Flash-MSA: Open-Source Sparse Attention Kernels for Fast Training

Flash-MSA is a new open-source training kernel suite for MiniMax Sparse Attention (MSA), implemented in CuTeDSL for Hopper and Blackwell GPUs. While several frontier models rely on sparse attention for fast inference, until now no efficient open training implementation existed — this project fills that gap, developed on rented H100/B200 hardware and unaffiliated with MiniMax itself.

MSA builds on ideas from DeepSeek Sparse Attention but introduces block-level sparsity (128-token blocks selected via max-pooling), swaps MLA for GQA in the main attention path, and lets each proxy head independently select which KV blocks to attend to. The GQA choice matters because no major Western lab has adopted MLA in training, making DSA-style techniques used in models like DeepSeek-V4 hard to port directly.

The kernel design keeps compute linear wherever possible: block indices computed once in the proxy forward pass are cached and reused through the sparse main attention and the entire backward pass, so only the proxy forward remains quadratic in context length. The backward fuses proxy and main attention gradient computation, using an algebraic simplification of the KL-divergence loss (gradient = proxy probability minus main probability) to avoid ever materializing the full KL term — a significant savings in registers and memory. Correctness was validated against an eager PyTorch reference implementation via cosine similarity comparisons across configurations.

For engineers working on long-context training, this offers a rare look at how to implement sparse attention training efficiently at the kernel level, plus a working reference for teams wanting to adopt GQA-based sparse attention without MLA.