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

I introduce Flash-MSA, the first performant open-source training kernels for MiniMax Sparse Attention designed for Hopper and Blackwell GPUs. Developed on Spheron rentals, this implementation uses blockwise sparsity and GQA to overcome limitations in current frontier models. By optimizing memory usage and fusing backward passes, I enable efficient training of million-token contexts while maintaining high precision compared to eager PyTorch implementations.

"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."