Flash-MSA: Accelerating Million-Token Training with Sparse Attention Kernels
Summary
Flash-MSA presents open-source training kernels for Minimax Sparse Attention (MSA), enabling efficient million-token training on specialized GPUs. It covers blockwise sparsity, GQA-based main attention, proxy-head grouping, and a KL-loss trick to avoid full materialization, plus forward/backward kernel design and planned future optimizations.