jax-linear-attention
github.com/jla-org/jax-linear-attention- Category
- Developer Tools
- Rank
- No. 1709Tools index
- Pricing
- Open Source
- Type
- TOOL
- GitHub
- 5 stars
- Latest release
- v0.1.0
- Date
About
jax-linear-attention provides Pallas Mosaic GPU kernels, forward and backward, for NVIDIA Hopper GPUs covering Mamba-2, Mamba-3, Gated DeltaNet, KDA, RWKV-7, GLA and an FA3-style attention kernel. It adds an LLM training stack around them for researchers training linear-attention and hybrid models in JAX. The README reports speedups over upstream Triton and CUDA kernels and lists cases where it loses.
Tags
jaxpallaslinear-attentionmambagpu-kernelsllm-training
Tech Stack
Python
Media
Comments (0)
No comments yet
Editorially curated, with community endorsements as a secondary signal. Corrections welcome.