HeadlinesBriefing favicon HeadlinesBriefing.com

Flash-MSA Open-Source Kernels Enable Minimax Sparse Attention Training

Hacker News •
×

A developer released Flash-MSA, the first performant open-source training kernels for Minimax Sparse Attention (MSA) written in CuTe DSL for Hopper (H100) and Blackwell (B200) GPUs. Development occurred on Spheron H100 and B200 rentals, referencing FlashAttention-4, MSA inference code, and Codex. The implementation is unofficial and unaffiliated with MiniMax.

MSA diverges from DeepSeek Sparse Attention (DSA) in three ways: blockwise sparsity selects key-value blocks of 128 via max-pooling over proxy scores; GQA replaces MLA for main attention, making sparse attention accessible to models not using MLA; and group-wise proxy head specialization lets each proxy head select different KV subsets, increasing expressivity.

The kernel fuses proxy and main attention backward passes to compute gradients in linear time. Proxy forward tracks top-k scores and indices in registers without accumulating LSE, recomputing it during backward over sparse activations. Main attention uses MoBA's varlen flash trick for block-sparse forward. The KL divergence gradient simplifies to proxy probability minus main probability, avoiding materialized distributions. Warmup mode runs dense main attention forward and skips proxy forward entirely.

Correctness was verified by sweeping cosine similarity between kernel and eager PyTorch implementations across configurations in bf16 precision, including both target output loss and internal KL loss gradients.