Stratified Sparse Attention acelera la decodificación 2,96 veces con un contexto de 1M

Un equipo de investigadores ha propuesto Stratified Sparse Attention (SSA), una forma de acceder a una pequeña fracción del KV-cache y perder bastante menos calidad que los esquemas dispersos anteriores. En RULER-HARD con Ministral-3-8B y una dispersión de la atención de 50x, SSA eleva la precisión del 81,8% al 84,2% con un contexto de 32K y del 63,8% al 69,3% con 128K. Eso supone el 94% y el 85% del resultado de la atención densa completa. Con un contexto de un millón de tokens, los kernels CUDA optimizados aceleran la decodificación hasta 2,96 veces respecto a FlashInfer.

Dónde pierde calidad la atención dispersa

Con contextos largos, el cuello de botella de la inferencia no es tanto el cálculo como la lectura del KV-cache, que crece en cada paso de decodificación. La atención dispersa (sparse attention) reduce el coste porque accede solo a una parte del caché. Aquí hay una bifurcación.

El primer camino es la selección determinista. Top-k o top-p toman los tokens con las puntuaciones más altas. Así se captura de forma fiable lo principal, pero se descarta sin más toda la atención restante, que token a token es pequeña y, sumada, puede ser significativa. El segundo camino es el muestreo. Cubre mejor la cola de la distribución, pero paga con varianza: el resultado varía más de una ejecución a otra.

Los enfoques híbridos combinan top-k con muestreo uniforme. Los autores muestran que, con una dispersión agresiva, la brecha respecto a la atención densa sigue siendo grande. El muestreo uniforme es ciego: gasta presupuesto en tokens que apenas influyen en la respuesta.

Cómo SSA divide el caché en estratos

SSA construye un esquema unificado con dos partes. El modelo lee de forma determinista los tokens con puntuaciones altas. El resto del caché no se descarta ni se elige al azar: sus tokens se muestrean con probabilidades proporcionales a las puntuaciones proxy de atención. Para que la contribución de los tokens elegidos no distorsione el resultado, SSA introduce una corrección por la probabilidad de que cada uno entre en la muestra. Así se recupera la masa de atención que queda fuera del conjunto seleccionado.

La idea se parece al muestreo estratificado clásico: la zona importante se cubre por completo y el resto se estima a partir de una submuestra con pesos. Los autores subrayan que SSA es una generalización. Top-k, top-p y el híbrido de top-k con muestreo resultan ser casos particulares suyos, de modo que el método puede verse como un marco común para toda la familia.

Cifras y aceleración

La evaluación se realizó en cuatro benchmarks (RULER, LongBench, LOFT y AIME) y con cinco modelos de 3B a 27B parámetros. Con la misma dispersión, SSA supera de forma consistente a los métodos de referencia. Los resultados principales:

• RULER-HARD, Ministral-3-8B, 32K: 81,8% → 84,2% frente al mejor baseline práctico, 94% de la atención densa;
• RULER-HARD, Ministral-3-8B, 128K: 63,8% → 69,3%, 85% de la atención densa;
• régimen: dispersión de la atención de 50x;
• velocidad: hasta 2,96x en decodificación respecto a FlashInfer con un contexto de 1M;
• código abierto: github.com/skylight-org/SSA-Paper-Code.

La segunda parte del resultado también importa. Un método más preciso suele requerir operaciones más complejas, y la aceleración puede desaparecer. Con kernels CUDA optimizados, los autores muestran que la mejora de calidad no se pierde en sobrecostes. La ganancia de precisión es mayor con 128K que con 32K (5,5 p. p. frente a 2,4 p. p.), es decir, el método ayuda más donde el contexto es más largo y la dispersión resulta más dañina.

Qué obtienen los profesionales

Para los equipos que sirven contextos largos, el KV-cache es una partida directa de gasto en memoria de GPU y ancho de banda. Si una reducción de 50x en los accesos al caché conserva entre el 85% y el 94% de la calidad de la atención densa en tareas difíciles, los ingenieros disponen de otra palanca entre coste y precisión. Resulta especialmente interesante para escenarios agénticos y para el trabajo con grandes bases de código, donde el contexto crece rápido.

También se ven las limitaciones. Parte de las cifras corresponde a un solo modelo, Ministral-3-8B, y con 128K sigue habiendo un retraso notable respecto al modo denso: el 85% es un buen resultado, pero no equivale a paridad. Para producción habrá que probar el método con tareas propias y comprobar la compatibilidad de los kernels con el stack en uso. El código abierto rebaja la barrera: se pueden reproducir los experimentos y comparar SSA con los esquemas top-k actuales sin esperar a que se implemente en los frameworks populares.