Flash-MSA es una implementación open source, escrita en CuTeDSL, de los kernels de entrenamiento para Minimax Sparse Attention (MSA), un mecanismo de atención dispersa similar a DeepSeek Sparse Attention pero con tres diferencias clave: sparsity por bloques de 128 tokens seleccionados vía max-pooling, uso de Grouped-Query Attention (GQA) en lugar de Multi-Latent Attention para la atención principal, y especialización por grupos de las cabezas proxy. El proyecto, desarrollado por Nanduru Ganesh sobre GPUs H100 y B200 alquiladas en Spheron, ofrece por primera vez código público y eficiente para entrenar este tipo de atención en hardware Hopper y Blackwell. Entre sus aportes técnicos destacan: un forward que encadena atención proxy seguida de atención principal dispersa, almacenando en caché los índices de bloque para reutilizarlos en el backward y mantener el entrenamiento lineal en la longitud de contexto; un backward fusionado que combina las pasadas de la atención proxy y principal para calcular gradientes mediante un atajo algebraico sobre la divergencia KL, evitando materializar ambas distribuciones en memoria compartida; y un modo warmup que ejecuta atención principal densa durante la fase inicial del entrenamiento. La corrección se valida comparando las salidas y gradientes con una implementación eager en PyTorch, alcanzando similitud coseno de 0,9996 a 0,9985 en precisión bf16 para secuencias de 4.096 y 8.192 tokens. El repositorio incluye integración con Megatron-LM. Su utilidad principal es permitir a investigadores y laboratorios sin acceso a MLA entrenar modelos frontera que aplican atención dispersa estilo GLM-5.2 o DeepSeek-V4 a contextos largos.
Noticias agregadas con IA