- 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.
- Correctness verified via cosine similarity tests between kernel and PyTorch implementations across various configurations.
- Future work includes increasing backward parallelizability, exploring router architectural speedups, and implementing context parallelism methods like headwise all-gather or ring parallelism.
- Notes mention joint indexer-main attention training trade-offs and scheduler requirements for proxy heads relative to KV heads.
- Sparsity benefits visualized via top-k sweeps, showing performance improvements over Flash Attention, especially at longer contexts.