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

Researchers have developed Flash-MSA, the first open-source training kernels for MiniMax Sparse Attention that enable efficient million-token model training on Hopper and Blackwell GPUs. The implementation uses blockwise sparsity (selecting KV pairs in blocks of 128), GQA instead of MLA for main attention, and group-wise specialization of proxy heads to improve expressivity while reducing computational overhead. By caching sparse block indices throughout the training step, only the proxy attention forward pass remains quadratic with context length, while other operations maintain linear complexity.