- Published on
Accelerating Million-Token Training with Sparse Attention Kernels
- Authors
- Name
- DeepHorus
A new open-source implementation of sparse attention training kernels, called Flash-MSA, promises to accelerate million-token training on Hopper and Blackwell GPUs. Developed independently by a researcher using Spheron H100 and B200 rentals, Flash-MSA is the first performant open-source kernel for MiniMax Sparse Attention (MSA) in CuTe DSL. This breakthrough addresses a critical gap in efficient training of sparse attention, which has primarily been available only for inference in frontier models such as GLM-5.2 and DeepSeek-V4.
Sparse attention mechanisms are widely adopted in large language models to reduce computational costs by focusing on a subset of key-value (KV) pairs. MiniMax's MSA variant differs from DeepSeek's approach in three key ways: blockwise sparsity, Grouped-Query Attention (GQA) instead of Multi-Head Latent Attention (MLA), and group-wise specialization of proxy heads. These changes introduce caching benefits and increase main attention expressivity by allowing independent KV selection per query group. However, until now, no open-source code existed to train MSA efficiently.
The Flash-MSA kernel design minimizes redundant computations and memory usage by streaming top-k accumulators during the forward pass and reusing cached block indices during the backward pass. Only the proxy attention forward is quadratic with respect to context length; all other operations leverage sparse blocks. The forward pass comprises proxy attention followed by sparse main attention, while the backward pass fuses proxy and main attention gradients to compute training signals. This approach significantly reduces the overall computational footprint for long-context training.
Key figures from the development highlight the efficiency gains: the proxy forward uses a cheap recompute instead of fusing log-sum-exp (LSE) with top-k selection, and blockwise sparsity reduces KV cache storage to only block indices. By re-parameterizing block sparse attention as variable-length flash attention, the kernel achieves further optimizations. These techniques collectively enable handling sequences up to millions of tokens in training, a feat previously challenging for open-source frameworks.
The implications for the AI community are substantial. Efficient training of sparse attention allows researchers to train models with much longer contextsâpotentially up to millions of tokensâwithout proportional increases in compute. This democratizes access to techniques used by labs like MiniMax and DeepSeek, which have kept their implementations proprietary. Moreover, the use of GQA rather than MLA makes the kernel compatible with models developed in Western labs, broadening its applicability.
Reactions from the AI research community indicate strong interest, as many had called for open-source training kernels to replicate frontier model capabilities. The developer notes that the work was made possible by referencing FlashAttention-4, the MSA inference implementation, and the CuTe DSL library. Licensing and integration with existing trainer frameworks are expected to be straightforward, though further validation on large-scale training runs is still needed.
Looking ahead, Flash-MSA could accelerate the development of next-generation large language models with ultra-long context windows. As frontier labs continue to push context lengths beyond one million tokens, open-source training kernels like this will be essential for enabling wider experimentation. The developer plans to release the full code on GitHub, inviting community contributions to refine and extend the implementation across additional GPU architectures.
