Flash-MSA: Accelerating Million-Token Training With Sparse Attention Kernels | Ganesh Nanduru

[Github] [MiniMax Paper] [Trainer]

Flash-MSA vs Flash-Attention isolated train step.1

Several frontier models [1, 2, 3, 4, 5] 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 H100 and B200 rentals and with the help of referencing FA4, MSA inference, 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, 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 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.<br>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, so I copied their clever trick 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...