Show HN: We ported FLA's Triton kernels to Jax Pallas; they're faster on H100

Category: ai-ml

Tags: jax, linear-attention, kernel, gpu, mamba, pallas

Score: 7.5/10 (Innovation: 6, Technical: 8, Documentation: 9, Utility: 7)

jax-linear-attention ports Flash Linear Attention's Triton kernels to JAX Pallas, providing fast forward/backward kernels for linear attention and SSM architectures like Mamba-2, Gated DeltaNet, KDA, RWKV-7, and GLA on NVIDIA Hopper GPUs. It includes a full LLM training stack and benchmarks showing up to 5.37x kernel speedups and 72.3% MFU on H100, filling a gap for efficient linear attention training in the JAX ecosystem.

Target audience: machine learning researchers and engineers training linear-attention and state-space models in JAX

Repository: https://github.com/jla-org/jax-linear-attention · Python · Apache-2.0 · 4 stars

View on Hacker News