Flash-MSA: núcleos abiertos para entrenar atención dispersa a escala de millones de tokens

Flash-MSA: Accelerating Million-Token Training with Sparse Attention Kernels

Flash-MSA presenta los primeros kernels de entrenamiento de código abierto para MiniMax Sparse Attention (MSA), optimizados para GPUs Hopper y Blackwell. A diferencia de la inferencia, el entrenamiento eficiente de atención dispersa no tenía implementación pública. Los kernels, escritos en CuTeDSL, logran un rendimiento comparable a FlashAttention en pasos de entrenamiento aislados, con una similitud coseno superior a 0.998 en precisión bf16. El diseño aprovecha la dispersión por bloques para almacenar índices de bloques y evitar recalcular la atención principal en el backward, reduciendo la complejidad a lineal. Incluye una fusión del backward de atención proxy y principal, y un truco matemático para calcular el gradiente de la pérdida KL sin materializar las distribuciones completas.