GoKawiil[Github] [MiniMax Paper] [Trainer]
Flash-MSA vs Flash-Attention isolated train step.
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
... continue reading