Otevřená knihovna jax-linear-attention přináší do JAX výpočetní jádra pro lineární pozornost a modely stavového prostoru na GPU Nvidia Hopper. Používá Pallas Mosaic a zahrnuje dopředný i zpětný průchod. Dokumentace uvádí podporu rodin Mamba, Gated DeltaNet, KDA, RWKV a dalších, doplněnou trénovací sestavou jazykových modelů.
Autoři zveřejňují měření na H100 a srovnání s příslušnými implementacemi včetně případů, kde jejich řešení prohrává. Zrychlení konkrétního jádra nelze zaměnit za stejné zrychlení celého tréninku nebo za výsledek na jiném hardwaru. Projekt míří na výzkumníky a vývojáře pracující v JAX; pro praktické použití je nutné prověřit podporovanou architekturu, přesnost výpočtu a podmínky měření.
