Flash-MSA: Accelerating Million-Token Training with Sparse Attention Kernels
4 hours ago
- #Model Training
- #GPU Kernels
- #Sparse Attention
- 首个用于Hopper和Blackwell GPU的Minimax Sparse Attention (MSA)开源训练内核在CuTeDSL中发布。
- MSA采用块级稀疏设计，包含128个基于最大池化的稀疏块，使用GQA替代MLA，并对代理头进行分组专业化处理。
- 内核设计包括前向代理注意力（带top-k排序）与稀疏主注意力，以及利用缓存块索引实现高效融合的反向传播。
- 通过特定技巧避免完整矩阵计算，高效计算KL散度损失梯度，显著降低内存占用。
- 通过多种配置下内核与PyTorch实现的余弦相似度测试验证了正确性。
- 未来工作包括增强反向传播的并行性、探索路由器架构加速，以及实现上下文并行方法（如头间全收集或环形并行）。
- 备注提及联合索引器-主注意力训练的权衡，以及代理头相对于KV头的调度器需求。
- 通过top-k参数扫描可视化稀疏性优势，显示其性能优于Flash Attention，尤其在长上下文场景中。