- First open-source training kernels for Minimax Sparse Attention (MSA) in CuTeDSL for Hopper and Blackwell GPUs are introduced.
- MSA features blockwise sparsity with 128 blocks using max-pooling, GQA instead of MLA, and group-wise specialization of proxy heads.
- Kernel design includes forward proxy attention with top-k sorting and sparse main attention, and a fused backward pass using cached block indices for efficiency.
- KL divergence loss gradients are efficiently computed using a trick to avoid full materialization, reducing memory usage.