MISA: Mezcla de Atención Dispersa de Indexadores para Inferencia en Modelos de Lenguaje de Contexto Largo

MISA: Mixture of Indexer Sparse Attention for Long-Context LLM Inference

May 8, 2026

Autores: Ruijie Zhou, Fanxu Meng, Yufei Xu, Tongxuan Liu, Guangming Lu, Muhan Zhang, Wenjie Pei

cs.AI

Resumen

DeepSeek Sparse Attention (DSA) establece el estado del arte en atención dispersa de grano fino durante la inferencia al introducir un indexador aprendido a nivel de tokens que evalúa cada token de prefijo y selecciona los más relevantes para la atención principal. Para mantener su expresividad, el indexador utiliza múltiples cabezales de consulta (por ejemplo, 64 en DeepSeek-V3.2) que comparten el mismo conjunto de tokens seleccionados; este diseño multicabezal es precisamente lo que convierte al indexador en el principal costo en contextos largos. Proponemos MISA (Mixture of Indexer Sparse Attention), un reemplazo directo del indexador DSA que trata sus cabezales como un grupo de mezcla de expertos. Un enrutador ligero utiliza estadísticas económicas a nivel de bloques para seleccionar un subconjunto dependiente de la consulta con solo unos pocos cabezales activos, y solo esos cabezales ejecutan la costosa evaluación a nivel de tokens. Esto preserva la diversidad del grupo original de indexadores mientras reduce el costo por consulta, pasando de evaluar cada token de prefijo con todos los cabezales a hacerlo con solo un puñado de cabezales enrutados, más un término de enrutamiento insignificante calculado sobre un pequeño conjunto de claves agrupadas. Además, introducimos una variante jerárquica de MISA que utiliza el paso de enrutamiento para mantener un conjunto ampliado de candidatos y luego los reordena con el indexador DSA original, recuperando casi exactamente los tokens seleccionados finales. Con solo ocho cabezales activos y sin entrenamiento adicional, MISA iguala al indexador DSA denso en LongBench tanto para DeepSeek-V3.2 como para GLM-5, mientras opera con ocho y cuatro veces menos cabezales de indexación respectivamente, y supera en promedio a HISA. También preserva mapas de calor completamente verdes en pruebas Needle-in-a-Haystack hasta contextos de 128K tokens y recupera más del 92% de los tokens seleccionados por el indexador DSA por capa. Nuestro núcleo TileLang logra aproximadamente una aceleración de 3.82 veces sobre el núcleo original del indexador DSA en una sola GPU NVIDIA H200.

English

DeepSeek Sparse Attention (DSA) sets the state of the art for fine-grained inference-time sparse attention by introducing a learned token-wise indexer that scores every prefix token and selects the most relevant ones for the main attention. To remain expressive, the indexer uses many query heads (for example, 64 on DeepSeek-V3.2) that share the same selected token set; this multi-head design is precisely what makes the indexer the dominant cost on long contexts. We propose MISA (Mixture of Indexer Sparse Attention), a drop-in replacement for the DSA indexer that treats its indexer heads as a pool of mixture-of-experts. A lightweight router uses cheap block-level statistics to pick a query-dependent subset of only a few active heads, and only those heads run the heavy token-level scoring. This preserves the diversity of the original indexer pool while reducing the per-query cost from scoring every prefix token with every head to scoring it with only a handful of routed heads, plus a negligible router term computed on a small set of pooled keys. We further introduce a hierarchical variant of MISA that uses the routed pass to keep an enlarged candidate set and then re-ranks it with the original DSA indexer to recover the final selected tokens almost exactly. With only eight active heads and no additional training, MISA matches the dense DSA indexer on LongBench across DeepSeek-V3.2 and GLM-5 while running with eight and four times fewer indexer heads respectively, and outperforms HISA on average. It also preserves fully green Needle-in-a-Haystack heatmaps up to a 128K-token context and recovers more than 92% of the tokens selected by the DSA indexer per layer. Our TileLang kernel delivers roughly a 3.82 times speedup over DSA's original indexer kernel on a single NVIDIA H200 GPU.