Repository · Quick start · Docs · Benchmarks · Discussions · Discord
jax-linear-attention brings FLA-class linear-attention kernels to JAX. Mamba-2 and Mamba-3, Gated DeltaNet, KDA, RWKV-7, GLA and an FA3-style attention kernel, written in Pallas Mosaic GPU for NVIDIA Hopper, forward and backward, with a full LLM training stack around them. Every kernel is exact by default and measured against the latest upstream code.
pip install "jax-linear-attention[cuda]"Then call one op, jla.ops.linear_attention(...), or train a model from one TOML with jla train.
The quick start has both. It runs on
H100 and H200; on other devices every op falls back to plain JAX.
|
Kernels for the new layers Mamba-2/3, Gated DeltaNet and DeltaNet, KDA, RWKV-7, GLA, Mamba-1, mLSTM and Gated DeltaNet-2, plus FA3-style softmax attention with packed documents, windows, GQA and sinks. |
Exact by default Every recurrence ships three ways: a token-by-token reference, a chunked JAX op and the kernel. impl="auto" runs a kernel only where its GPU tests cover the call.
|
|
A training stack around them FSDP, data, tensor, expert and pipeline parallelism, FP8 training, dropless MoE, packed documents, exact resume, SFT and DPO, Hugging Face import and export, lm-evaluation-harness. |
Fast end to end 72.3% MFU on one H100 and 68.2% MFU per GPU on 8 H100s for a 1.3B transformer in FP8, and 1.04-1.67x flame (FLA's trainer) in bf16 on the same tokens. |
- Kernels: against the latest upstream code on one H100, forward + backward: Mamba-3 MIMO 2.70-5.37x the official recurrence kernels, the Mamba-2 scan 1.20-1.96x, KDA 1.47-1.78x FLA, GLA 1.52-1.77x FLA, Gated DeltaNet 1.71-2.68x FLA's Triton kernels.
- End to end: 1.3B models against flame on the same tokens: 1.04-1.67x in bf16 and 1.10-1.77x in FP8 against flame's FP8, and 1.47x flame per GPU for the transformer on 8 H100s.
- In the Gated DeltaNet-2 paper's setting: every 1.3B model of the paper's Figure 2 trains 1.14-1.41x faster on one H100 than the paper reports, with a 32,000-token vocabulary.
- Where we lose: FlashQLA beats our exact Gated DeltaNet split on long single sequences and few-head packed documents, and our FA3-style attention runs below FA3. Each case is listed with its numbers and the reason.
Every number was measured on H100s with no other jobs on the GPUs; the baseline commits and commands are in benchmarks/.
jax-linear-attention is free and open source under the Apache License 2.0. Contributions are welcome with a DCO sign-off (CONTRIBUTING.md). Parts of it follow or adapt flash-linear-attention, FlashAttention-3, Mamba, FlashQLA, FlashKDA, Tokamax and JAX, credited in NOTICE.
- Questions and ideas: Discussions
- Chat: Discord
- Bugs: open an issue
- Security problems: report them privately, never in a public issue
- Everything else: vikramlex@karlexai.com · @vikram_lex on X · karlex.ai
jax-linear-attention is made by 
KarLex AI.
