DashAttention alcanza 75% de dispersidad manteniendo precisión de atención completa

Un equipo de investigación de Edimburgo, Cambridge, Instituto Superior Técnico y Tsinghua publicó DashAttention, un reemplazo directo para la atención jerárquica top-k utilizada en NSA e InfLLMv2 que mantiene la etapa de selección KV completamente diferenciable. El pipeline completo de dos etapas ahora puede optimizarse extremo a extremo. El artículo (arXiv 2605.18753, publicado el 18 de mayo) incluye un kernel GPU Triton y cuatro puntos de control MiniCPM-4-8B ajustados finamente lanzados en Hugging Face para comparación directa.

El problema central: NSA e InfLLMv2 aplican un paso de selección de bloque top-k coarse, luego ejecutan atención softmax estándar solo sobre los bloques seleccionados. Top-k asume que cada query necesita el mismo número de bloques de tokens relevantes. La operación de selección discreta corta el flujo de gradientes entre las etapas coarse y fine, por lo que el puntuador coarse no puede aprender de la pérdida de atención downstream. Las dos etapas se entrenan con objetivos desalineados.

DashAttention reemplaza top-k con α-entmax, una transformación dispersa diferenciable que generaliza softmax. Los tokens cuyas puntuaciones caen por debajo de un umbral adaptativo por query reciben exactamente cero probabilidad. El umbral se determina por la distribución de entrada en lugar de una k fija. Una query con un bloque de contexto dominante atiende a fewer bloques que el máximo configurado; una query con muchos chunks relevantes atiende a más. La salida de α-entmax actúa como una ponderación previa para la atención softmax de segunda etapa. Porque α-entmax es diferenciable, los gradientes fluyen desde la etapa softmax completamente a través del selector de bloques. Los autores llaman a esta propiedad "no dispersiva" — la masa de atención se concentra en bloques relevantes en lugar de dispersarse entre bloques irrelevantes.

DashAttention iguala la calidad de atención completa con 75% de dispersidad y logra una frontera Pareto mejor que NSA e InfLLMv2. La brecha se amplía en regímenes de alta dispersidad. La precisión de los competidores se degrada más rápidamente por encima de rangos de dispersidad cómodos, mientras que la selección adaptativa de DashAttention retiene más precisión a escala. Los benchmarks se ejecutaron en MiniCPM-4-8B base. Las cuatro variantes — atención completa, InfLLMv2, NSA y DashAttention — están disponibles en Hugging Face bajo la organización fasa-org. Los autores reportan un speedup sobre FlashAttention-3 en tiempo de inferencia; el multiplicador exacto requiere la Sección 5 del artículo.

La implementación Triton es instalable vía pip install -e . desde github.com/fasa-org/dash-attention. La interfaz envuelve queries, keys, values y un vector de clasificación por cabeza (head_cls) para puntuación coarse. Opcionalmente retorna una máscara active_blocks indicando qué bloques fueron seleccionados por query, útil para perfilar dispersidad real en tiempo de ejecución. Los flags chunk_size y estimate_diagonal controlan granularidad de bloque y corrección diagonal para normalización α-entmax. GQA es soportado vía enable_gqa=True.

Antes del despliegue en producción, quedan dos brechas. El artículo está en pre-revisión por pares y las afirmaciones de eficiencia necesitan verificación más allá del resumen. Los datos de TTFT y throughput de generación de tokens no se divulgan en materiales públicos. El paso de selección de bloque α-entmax añade una llamada kernel adicional comparado con métodos top-k más simples; el overhead en longitudes de contexto cortas (menores de 16k) no está caracterizado. El vector de clasificación head_cls es una adición arquitectónica que no existe en puntos de control Llama-family estándar. Adoptar DashAttention requiere ajuste fino — no puede ser intercambiado como una optimización en tiempo de servicio sin reentrenamiento. Los cuatro puntos de control 8B lanzados proporcionan un punto de partida reproducible, pero los costos de ajuste fino a escala de producción no se divulgan.

Si está ajustando finamente para tareas de contexto largo por encima de 32k y NSA o InfLLMv2 está en su lista corta, la selección diferenciable y dispersidad adaptativa justifican una comparación directa en su suite de evaluación. Espere los números exactos de eficiencia antes de dimensionar presupuesto de GPU.