Flash-MSA es el primer conjunto de kernels de código abierto que permite entrenar de forma eficiente modelos con MiniMax Sparse Attention (MSA), un mecanismo de atención dispersa inspirado en DeepSeek Sparse Attention que varios modelos frontera —DeepSeek V3.2, DeepSeek V4-Pro, GLM-5.2, LongCat-2.0 y MiniMax-M3— ya usan en inferencia. Hasta ahora no existía código público capaz de entrenar este tipo de arquitecturas a escala.
El proyecto, desarrollado por Nanduru Ganesh y publicado en GitHub, está implementado en CuTeDSL para GPUs Hopper y Blackwell y se integra con Megatron-LM. Frente a la formulación basada en MLA de DeepSeek, MSA introduce tres cambios clave: dispersa por bloques de 128 tokens seleccionados mediante max-pooling, atención principal con GQA en lugar de MLA, y especialización por grupos de las cabezas proxy para aumentar la expresividad.
El diseño de los kernels minimiza el trabajo repetido y aprovecha que la dispersión por bloques permite cachear los índices seleccionados durante todo el paso de entrenamiento: solo el forward proxy es cuadrático respecto a la longitud de contexto; el resto opera linealmente. El backward fusiona los pases proxy y principal para calcular gradientes sin materializar la divergencia KL completa. El autor valida la corrección comparando frente a una implementación eager en PyTorch en precisión bf16 con tolerancia de 0,01 y obtiene similitud coseno ≥0,9985 en secuencias de 4.096 y 8.192 tokens. El autor aclara que no es una implementación oficial ni está afiliado a MiniMax.