Flash-MSA: Accelerating Million-Token Training with Sparse Attention Kernels

A developer released the first open-source training kernels for MiniMax Sparse Attention (MSA) on Hopper and Blackwell GPUs, enabling efficient million-token training with sparse attention. The kernels, built in CuTeDSL and referencing Flash Attention 4 and MSA inference code, support blockwise sparsity and GQA instead of MLA, making sparse attention accessible to models outside frontier labs.

Github https://github.com/nanduruganesh/flash-msa MiniMax Paper https://arxiv.org/abs/2606.13392 Trainer https://github.com/nanduruganesh/Megatron-LM Several frontier models 1 https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro , 2 https://huggingface.co/zai-org/GLM-5.2 , 3 https://huggingface.co/deepseek-ai/DeepSeek-V3.2 , 4 https://huggingface.co/meituan-longcat/LongCat-2.0 , 5 https://huggingface.co/MiniMaxAI/MiniMax-M3 use sparse attention to greatly speedup their inference, though no one has posted code to train it efficiently. Today I introduce the world's first performant open-source training kernels for Minimax Sparse Attention in CuTeDSL for Hopper and Blackwell GPUs. I did all of the dev work on Spheron https://www.spheron.network/ H100 and B200 rentals and with the help of referencing FA4 https://github.com/Dao-AILab/flash-attention/tree/main/flash attn , MSA inference https://github.com/MiniMax-AI/MSA , and Codex. Disclaimer: This is not an official implementation and I am not affiliated with MiniMax About MSA MSA is similar to Deepseek Sparse Attention https://arxiv.org/abs/2512.02556 , with some core changes Fig. 1 from the MSA Paper 1. Blockwise sparsity Instead of the proxy attention selecting individual KVs for the main attention, it selects them in blocks of 128 using max-pooling over the proxy scores. This introduces some nice caching properties for the kernels. 2. GQA instead of MLA for the main attention This one is particularly important because no western labs, to the best of my knowledge, have adopted MLA into their training, making the prevailing sparse attention formulation in frontier models like GLM-5.2, DSv4, which are fit to MLA, inaccessible to models here. 3. Group-wise specialization of proxy heads Replacing MLA with GQA introduces independent groups of queries within each layer, giving us the option to select a different subset of KVs with each proxy head instead of summing and scoring for the entire attention layer like DSA does. There is some evidence https://arxiv.org/abs/2410.10819 suggesting attention heads organically attend to different tokens, so this change should increase main attention expressivity. Kernel Design High level overview of kernel sequence To run MSA efficiently we need to repeat as little work as possible and not overload the registers / shared memory too much. In addition to regular flash registers Q tile, KV tile, O accumulator, LSE accumulator , we have to take into account streaming top-k accumulators in the forward. In the backward, we have to make space for a double-attention combined pass so we can calculate main attention grads and proxy attention grads, since proxy grads require access to both proxy and main attn probs. One nice benefit of block-sparsity is that since we only have to cache the indices of blocks instead of individual tokens like in DSA, it is feasible to store block indices all the way through until the backward - meaning in the entire train step, only the proxy forward is quadratic w.r.t. context length, everything else uses the cached sparse blocks from the proxy forward. Forward Order of operations is proxy attention - sparse main attention - send main attention output to next layer, save main LSE for backward. Proxy attention The proxy dot-product is a little different from regular Flash Attn as we don’t have to accumulate output anymore, but we do have to keep track of the top k attention scores and their respective indices as we stream across keys. Unlike Flash I also do not accumulate LSE for the backward and instead run a very cheap recompute of proxy dot-product over sparse activations to get the LSE during the backward. This was faster than fusing LSE+topk in the forward for me in practice. As each tile of QK^T is calculated, I grab the causal local max score for each chunk and do an insertion sort into the current top-k values for each query row, held in registers. I had to slice key blocks in half to make space for these top-k registers. Also, MSA dictates that the local block for each token must be unmasked sliding-window style, so I set the attention scores for the local KV block for each query to inf. Main attention The main attention is just a block-sparse flash attention forward. This has been done before in MoBA https://arxiv.org/abs/2502.13189 , so I copied their clever trick https://github.com/MoonshotAI/MoBA/blob/master/moba/moba efficient.py to re-parameterize block sparse attention into varlen flash. Backward To calculate gradients for the proxy heads, we need to fuse the proxy and main attention backwards because the proxy training signal requires access to both proxy attn probs and main attn probs at once. Since we saved block indices from the forward and only train both attentions on the sparse KV activations, the backward can be run in linear time. First we pull our cached block indices and invert the mapping of $B$ batch,proxy head,query,top k slot - key block into $B^{\ast}$ batch,proxy head,key block - queries that use this block . We use $B^{\ast}$ to schedule query chunks to optimize for reuse of shared sparse KV blocks. Then we run a quick sparse proxy attention pass over the selected blocks again using the MoBA varlen trick to get the proxy LSE, then we stream the fused proxy-main backward tasks, loading chunks of QKV, Q proxy, K proxy, and main lse. To account for loading so many heads into the registers we have to reduce the size of the Q chunks and KV chunks we use at a time. In each stream we calculate main and proxy attention probs, compute dQ,dK,dV, then compute proxy dQ,dK from the KL training term: KL Divergence Loss Recall the original KL loss term from DSA is $L^\iota = \sum t{D {KL} p t,s t \Vert Softmax I t, s t }$ Materializing both indexer and main attn prob. distributions to accumulate KL divergence would require many read/writes to shared mem and usage of additional registers that would significantly slow down training. Fortunately, there is a trick we can use to backprop atomically and still be mathematically equivalent to a full KL loss. With the proxy attn prob as $p {px}$ and main attn prob as $p$, expand the KL term: \ L^\iota = \sum t{D {KL} p t \Vert p {px,t} } = \sum t{p t log \frac{p t}{p {px,t}} }\ Use log rules to expand again: \ L^\iota = \sum t{p log p t - log p {px,t} } = \sum t{ p t log p t - p t log p {px,t} }\ We want to calculate grads into the next latent which is the pre-softmax attention score at position i not t , $z {px,i}$ \ \frac{\partial L^\iota}{\partial z {px,i}} = \frac{\partial}{\partial z {px,i}}\sum t{ p t log p t - p t log p {px,t} }\ Main probs $p t$ are detached from the proxy/KL-loss graph, so they are constant in this partial. \ \frac{\partial L^\iota}{\partial z {px,i}} = -\frac{\partial}{\partial z {px,t}} \sum t{p t log p {px,t} }\ We know the softmax logprob partial https://math.stackexchange.com/a/2340848 of a softmax output at position t $p {px,t}$ w.r.t. pre-softmax logit $z {px,i}$ is $\frac{\partial log p {px,t} }{\partial{z {px,i}}} = \delta {it}-p {px,i}$ Kronecker delta function $\delta {it}$ is only nonzero at i==t, so $\sum t{p t\delta {it}}=p i$. \ \frac{\partial L^\iota}{\partial z {px,i}} = -p i + \sum t{p tp {px,i}} = -p i + p {px,i}\sum t{p t}\ Because $p t$ is a probability distribution, $\sum t{p t}=1$. \ \frac{\partial L^\iota}{\partial z {px,i}} = -p i+p {px,i}\ In other words, gradient to proxy score from KL loss = proxy prob - main prob. This is the term I use in the kernel to calculate proxy gradients without needing to ever fully materialize KL. Warmup kernels In warmup mode, the main attention forward is dense and does not use block indices, so the proxy forward can be skipped entirely, we can train it fully in the backward. For the main attention warmup forward kernel I just call flash and save its returned output and lse, and return a placeholder KL. In the backward I call dense flash on the indexer just to get the LSE, then reuse the fused proxy+main attention backward from the sparse MSA kernels. Correctness To verify correctness of the the kernel forward and backward, I implemented MSA in eager Pytorch and swept for cosine similarity between both implementations' forward outputs and backward grads over several configs. The sweep is run in bf16 precision and the backwards include both target output loss and the internal KL loss. Typically the tolerance for precision at bf16 is 0.01. | Batch | Sequence | Q heads | Forward | Backward projection gradients | |||| |---|---|---|---|---|---|---|---|---| | Output | Q | K | V | Proxy Q | Proxy K | ||| | 1 | 4,096 | 8 | 0.9996 | 0.9996 | 0.9996 | 0.9999 | 1.0000 | 1.0000 | | 2 | 4,096 | 8 | 0.9996 | 0.9996 | 0.9996 | 0.9999 | 1.0000 | 1.0000 | | 4 | 4,096 | 8 | 0.9996 | 0.9996 | 0.9996 | 0.9999 | 1.0000 | 1.0000 | | 1 | 8,192 | 8 | 0.9985 | 0.9985 | 0.9985 | 0.9995 | 0.9999 | 1.0000 | | 2 | 8,192 | 8 | 0.9985 | 0.9985 | 0.9985 | 0.9995 | 1.0000 | 1.0000 | | 4 | 8,192 | 8 | 0.9985 | 0.9985 | 0.9985 | 0.9995 | 1.0000 | 1.0000 | | 1 | 4,096 | 16 | 0.9996 | 0.9995 | 0.9995 | 0.9999 | 0.9999 | 1.0000 | | 2 | 4,096 | 16 | 0.9996 | 0.9996 | 0.9996 | 0.9999 | 1.0000 | 1.0000 | | 4 | 4,096 | 16 | 0.9996 | 0.9996 | 0.9996 | 0.9999 | 1.0000 | 1.0000 | | 1 | 8,192 | 16 | 0.9985 | 0.9984 | 0.9984 | 0.9997 | 0.9999 | 1.0000 | | 2 | 8,192 | 16 | 0.9985 | 0.9984 | 0.9985 | 0.9997 | 1.0000 | 1.0000 | | 4 | 8,192 | 16 | 0.9985 | 0.9984 | 0.9985 | 0.9997 | 1.0000 | 1.0000 | | 1 | 4,096 | 32 | 0.9996 | 0.9996 | 0.9996 | 1.0000 | 0.9999 | 1.0000 | | 2 | 4,096 | 32 | 0.9996 | 0.9995 | 0.9996 | 1.0000 | 0.9999 | 1.0000 | | 4 | 4,096 | 32 | 0.9996 | 0.9995 | 0.9996 | 1.0000 | 1.0000 | 1.0000 | | 1 | 8,192 | 32 | 0.9985 | 0.9983 | 0.9984 | 0.9998 | 0.9999 | 1.0000 | | 2 | 8,192 | 32 | 0.9985 | 0.9984 | 0.9984 | 0.9998 | 0.9999 | 1.0000 | What’s next Increasing Fused Backward Parallelizability Currently the backward is bound by low tensor-pipe utilization and low occupancy from the fused backward's heavy register/shared memory requirements needed to run 2 attentions at once. For example, for the throughput sweep in this blog, Flash-MSA backward used 138 registers/thread, 105 KB shared memory/CTA, meanwhile H100s/B200s support a max of 255 reg/thread and 228 KB shared memory/CTA so I was bound to 1 CTA/SM. To increase eligible warps I tried different configs of tiling Q/KV more narrowly to achieve 2 CTAs/SM but the increased total CTAs from doing this created a net slowdown. The theoretical occupancy of the Flash-MSA backward is 12.5% compared to 18.75% in Flash-Attention. Router architectural speedups GLM has already proved it is stable and much faster to use IndexShare https://arxiv.org/abs/2603.12201 to share proxy heads across layers. Also, the indexer seems to always be served in low-precision during inference, so training it in low precision could help address train-inference matching and speed up training greatly if it was stable. Context parallelism To scale LLM training at long context, some form of CP is mandatory or the trainer will quickly run out of memory. There are a few options here 1. Headwise all-gather For clarity I will refer to an “MSA group” of attention heads as the subset of main attn query heads assigned to each proxy query head. Since MSA groups run independently of each other in both fwd/bwd, it is trivial to use TP-style CP folding for CP rank up to num proxy heads, where each device holds Seq len / CP rank tokens and its MSA group s , then does an A2A comm of the full sequence before calling MSA. This requires no changes to the kernels, in fact, you could probably do this right now in the current Megatron MSA fork https://github.com/nanduruganesh/Megatron-LM . 2. Ring Implementing Ring-style parallelism https://christianjmills.com/posts/cuda-mode-notes/lecture-013/ here is more difficult but probably the optimal way of doing CP for MSA as it allows for better overlap and higher CP rank than a2a. Ring would require something like overlap-exchange within the proxy forward with a way to stream the top-k values and indices across devices, then broadcast back to all devices, then a clever sparse-key host device lookup and retrieval for the main attention across devices. I do not know how to wire this into the kernel so I will have to contact my local CP expert before putting out this core feature, but if anything this will be the next blog post. Note about joint indexer-main attention training While the paper does include a formulation of MSA where a value and output proj are added to the proxy head then the proxy output is added to the main attention’s output to introduce grads from CE loss into the proxy weights and allegedly boost knowledge-type evals, implementing this will slow down training from adding new heads and accumulators to the fwd/bwd. MSA also states themselves in Table 6 of their paper that proper warmup can make up for not training the indexer with CE loss. Note about scheduler requiring proxy heads ≥ KV heads Since I schedule by mapping each proxy head to its respective main attention GQA group, the current kernels require that MSA groups = GQA groups. There is probably a simple way to invert the mapping for the scheduler/fused backward, but this would only be needed in cases where the model has more KV heads than proxy heads. However I expect stable training to require at least 4 MSA groups so to have KV heads proxy heads you'd need 4 GQA groups and in this current regime of transformers I have a strong contempt for models with 1024 KV cache/layer so I am not interested in implementing this path. Sweep across Top-k I thought it would be interesting to visualize the sparsity benefit as blocks decrease. I want to do some more sweeps across GQA and proxy MQA configurations when I get compute. I would also like to continued pre-train an existing GQA base model like Qwen3 using MSA to test the conversion, this will require some more compute. Top-k sweep done with same config as the MSA vs. FA sweep. - Sweeps taken with 16 Q heads, 2 KV heads, 4 proxy Q heads, 1 proxy K head, 128 headdim, B=1, dtype bf16 on an H100 using Pytorch 2.12.0 and CUDA 13.0, averaged over 5 train steps after 5 warmup steps at each sequence length. I was not able to run another sweep on B200, but my guess is FA4 will be much faster than it was on H100, but will still get gapped by Flash-MSA at long context. ↩ fnref:1