First Open-Source Training Kernels for Sparse Attention Released, Enabling Million-Token LLM Training
Key Takeaways
- ▸First performant open-source training kernels for sparse attention, enabling 1M+ token training at production scale
- ▸Block-wise sparsity (128-token blocks) with GQA compatibility makes it accessible to existing frameworks, unlike MLA-dependent alternatives
- ▸Optimized for NVIDIA H100 and B200 GPUs using CuTeDSL, with only proxy attention forward pass remaining quadratic in context length
Summary
An independent developer has released Flash-MSA, the first open-source training kernels for MiniMax Sparse Attention (MSA), a technique used by frontier models to accelerate inference. Previously, sparse attention was only efficiently implemented for inference; Flash-MSA brings performant training support to the open-source community, with optimized kernels for NVIDIA's Hopper and Blackwell GPUs.
The implementation uses block-wise sparsity (128-token blocks) combined with Group Query Attention (GQA), making it compatible with existing training frameworks unlike alternative sparse attention approaches that require Multihead Latent Attention (MLA). The kernels achieve this by carefully managing register and memory pressure, using top-k selection during the proxy attention forward pass, and recomputing sparse attention during backpropagation to minimize memory overhead.
Developed on rental GPUs (H100 and B200), Flash-MSA enables training of models with massive context windows without quadratic memory scaling. The only quadratic component is the proxy attention forward pass; all other operations leverage cached sparse block indices. This breakthrough addresses a critical gap in the open-source LLM ecosystem, where sparse attention has remained confined to proprietary implementations.
- Group-wise specialization of proxy heads increases expressivity by allowing different KV subsets per attention head
- Eliminates the barrier between frontier model inference speedups and open-source training, democratizing sparse attention techniques
Editorial Opinion
This is a significant democratization moment for sparse attention training. Frontier models (GLM-5.2, DeepSeek-v4, MiniMax) have long enjoyed sparse attention's efficiency gains at inference time, but the gap between closed-source inference and open training has prevented the community from experimenting with these techniques at scale. Flash-MSA bridges that gap by choosing GQA over MLA—a pragmatic decision that prioritizes compatibility over strict adherence to proprietary architectures. For the first time, researchers can now train sparse-attention models on consumer/rented hardware without proprietary frameworks, which should accelerate open-model development.
