Pallas Mosaic GPU kernels for NVIDIA Hopper, forward and backward, for Mamba-2, Mamba-3, Gated DeltaNet, KDA, RWKV-7, GLA and more, with an FA3-style attention kernel and a full LLM training stack around them. Every kernel is exact by default and measured against the latest upstream code.
Open source · Apache-2.0 · Version 0.1.0
JLA installs from PyPI as jax-linear-attention. Call a kernel from your own JAX model with jla.ops.linear_attention(..., impl="auto"): it runs the Mosaic GPU kernel on an H100 and the chunked JAX op anywhere else. Or train a model from one TOML file with jla train.
Python 3.12 or newer and JAX 0.11.2 (the version is pinned). The kernels need an NVIDIA Hopper GPU (H100 or H200) and a CUDA 13 driver; on any other device every op falls back to its plain-JAX implementation. Quick start
pip install "jax-linear-attention[cuda]"CPU only (reference and chunked ops, small models):
pip install jax-linear-attentionThe layers behind newer language models (Mamba, Gated DeltaNet, KDA, RWKV-7, GLA) are recurrences, and they need dedicated kernels to train fast. PyTorch has them in Flash Linear Attention's Triton kernels; in JAX they have mostly run as plain JAX ops, which XLA cannot fuse into one efficient pass.
A token-by-token reference (the oracle), a chunk-parallel JAX op that runs on any JAX backend, and a Pallas Mosaic GPU kernel for Hopper. The tests hold each kernel to the reference on dense inputs, packed documents, initial and final states and every gradient.
impl="auto" runs a kernel only for the shapes, dtypes and options its GPU tests cover; every other call runs the chunked JAX op, so an unsupported call is slower, never wrong.
segment_ids reset the state at document boundaries inside every kernel, so a packed batch trains exactly as its documents would one by one.
Data, FSDP, tensor, expert and pipeline parallelism; FP8 with delayed scaling; dropless MoE on grouped-GEMM kernels; crash-consistent checkpoints with exact resume; SFT and DPO; Hugging Face import and export; perplexity and lm-evaluation-harness. Every option
On one H100, forward + backward: the Mamba-2 scan 1.20–1.96× the official Triton kernels and the whole Mamba-2 layer 1.17–1.41× the official fused training path, Mamba-3 SISO 1.58–2.32× and MIMO 2.70–5.37× the official recurrence kernels, GLA 1.52–1.77× FLA (dense), KDA 1.47–1.78× FLA, RWKV-7 1.06–1.46× FLA, and Gated DeltaNet 1.71–2.68× FLA's Triton kernels and 0.73–1.26× FlashQLA. Every shape and version
On one H100 and the same tokens, the transformer, sliding-window, Gated DeltaNet, KDA and Mamba-2 presets train 1.04–1.67× faster than flame (FLA's trainer, with torch.compile where it works) in bf16 and 1.10–1.77× in FP8, and Mamba-3 1.55× in bf16. Results
transformer-1p3b reaches 72.3% MFU in FP8 on one H100 and 68.2% per GPU on eight, and gdn-hd128-1p3b 55.0% and 52.3% against flame's 36.1% and 33.4%. MFU is model FLOPs over the 989 TFLOP/s bf16 spec; a power-capped H100 sustains about 681 TFLOP/s in bf16. Eight GPUs
The paper's Figure 2 times 1.3B hybrids on one H100 at 16,384 tokens per step. In that setting, every model of the figure trains 1.14–1.41× faster than the paper's numbers with a 32,000-token vocabulary, and 0.99–1.23× with Llama-3's 128,256. The Gated DeltaNet-2 hybrid runs at 52.1 and 49.1 thousand tokens/s at 2K × 8 and 16K × 1, against the paper's 38.0 and 36.1. The settings and the caveats
Forward + backward against the fastest upstream code that runs each shape: Mamba-2 and GLA lead at every length and head layout, Gated DeltaNet trails FlashQLA on long single rows of 8 heads of dim 128, and causal attention stays below FA3 (0.73–0.90× with the default backward). Attention results
The cases where JLA is slower than the alternatives, with their numbers. Every case and the reason
The forward is 0.90–0.92× FA3 on causal attention, short of the 0.95× target, and forward + backward 0.77–0.92× with the default backward (0.87–1.00× with the opt-in fused one); impl="auto" keeps cuDNN for dense causal and windowed calls.
On long single sequences (one row of 16K–64K tokens: 0.69–0.78× forward + backward at 8 heads, 0.68–0.73× at 16) and with packed documents at small B × H (one row of 8 heads with 16 documents: 0.74–0.77× from 16K tokens, 0.58–0.60× at 2K–8K). FlashQLA starts each sequence piece a few chunks early, from zero (an approximation); our split is exact, and its backward is two passes against FlashQLA's one fused kernel. At head dim 64, which FlashQLA does not support, packed documents in short rows also trail FLA (0.83–0.93× forward + backward at 2K–8K tokens).
At large B × H (0.57× at (8, 2048, 32)): FlashKDA keeps a bf16 state and approximate exp, and fuses its per-chunk preparation into one kernel.
The kernels need an H100 or H200 (everything else runs plain JAX), and JAX is pinned to 0.11.2 because Mosaic GPU's API changes with every JAX release. A few kernels (the Mamba-2 row kernels, the Gated DeltaNet piece scan, fp8_swiglu) are Pallas Triton, which JAX 0.11.2 deprecates; they move to Mosaic GPU before the pin does.
Our linear-attention kernels port the Triton kernels of Flash Linear Attention and follow its layer definitions, and its trainer, flame, is our end-to-end baseline. FlashAttention-3 shaped our attention kernel, our Mamba kernels port Mamba's, FlashQLA and FlashKDA showed the way for the Gated DeltaNet and KDA kernels, and Tokamax and JAX's Pallas examples are where our attention and GEMM kernels began. Acknowledgements
JLA is an independent project. It is not affiliated with or endorsed by Google or the JAX team; “JAX” names the framework it is built on.
JLA is free and open source under Apache-2.0. Version 0.1.0 is early development, and its APIs will change. Contributions are welcome, with a DCO sign-off on every commit.