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

Recent frontier models such as DeepSeek-V4 and GLM-5.2 use Sparse Attention to improve inference speed, but code for training it efficiently has been lacking. Flash-MSA addresses this by providing high-performance open-source training kernels that support Hopper and Blackwell GPUs.

Key features of MSA (MiniMax Sparse Attention):

• Blockwise Sparsity: Instead of individual KV tokens, blocks of 128 are selected, improving caching efficiency.
• GQA (Grouped-Query Attention) applied: Uses GQA instead of the existing MLA approach, making it accessible across diverse model architectures while increasing the representational power of attention heads.
• Group-wise Specialization: Allows each proxy head to select a different KV subset, ensuring diversity in attention.

Key points of kernel design:

• Forward Pass: Performs proxy attention followed by Sparse Main Attention, caching only the top-k indices during proxy attention to reduce memory overhead.
• Backward Pass: Performs a Double-attention combined pass to compute gradients for both main attention and proxy attention, maximizing efficiency.
• Memory Optimization: Thanks to block-sparsity, only the block indices from the proxy forward results need to be stored throughout the training step, significantly reducing memory usage.