MISA: Mischung von Indexer-Sparse-Attention für die LLM-Inferenz mit langem Kontext

Zusammenfassung

DeepSeek Sparse Attention (DSA) setzt den State-of-the-Art für feinkörnige Sparse-Attention zur Inferenzzeit, indem es einen gelernten tokenweisen Indexierer einführt, der jedes Präfix-Token bewertet und die relevantesten für die Hauptattention auswählt. Um ausdrucksstark zu bleiben, verwendet der Indexierer viele Query-Köpfe (z. B. 64 bei DeepSeek-V3.2), die denselben ausgewählten Token-Satz teilen; genau dieses Multi-Head-Design macht den Indexierer bei langen Kontexten zum dominierenden Kostenfaktor. Wir schlagen MISA (Mixture of Indexer Sparse Attention) vor, einen direkten Ersatz für den DSA-Indexierer, der dessen Indexierer-Köpfe als einen Pool von Mixture-of-Experts behandelt. Ein leichtgewichtiger Router verwendet kostengünstige Block-Level-Statistiken, um eine abfrageabhängige Teilmenge von nur wenigen aktiven Köpfen auszuwählen, und nur diese Köpfe führen das aufwändige Token-Level-Scoring durch. Dies bewahrt die Diversität des ursprünglichen Indexierer-Pools, während die Kosten pro Abfrage von der Bewertung jedes Präfix-Tokens mit jedem Kopf auf die Bewertung mit nur einer Handvoll gerouteter Köpfe reduziert werden, zuzüglich eines vernachlässigbaren Router-Terms, der auf einer kleinen Menge gepoolter Keys berechnet wird. Wir führen zudem eine hierarchische Variante von MISA ein, die den gerouteten Durchlauf nutzt, um einen erweiterten Kandidatensatz zu erhalten, und diesen dann mit dem ursprünglichen DSA-Indexierer neu bewertet, um die finalen ausgewählten Token nahezu exakt zu rekonstruieren. Mit nur acht aktiven Köpfen und ohne zusätzliches Training erreicht MISA auf LongBench sowohl mit DeepSeek-V3.2 als auch mit GLM-5 die Leistung des dichten DSA-Indexierers, während es mit acht bzw. vier Mal weniger Indexierer-Köpfen auskommt, und übertrifft HISA im Durchschnitt. Es bewahrt zudem vollständig grüne Needle-in-a-Haystack-Heatmaps bis zu einem Kontext von 128K Token und stellt pro Schicht mehr als 92 % der vom DSA-Indexierer ausgewählten Token wieder her. Unser TileLang-Kernel liefert auf einer einzelnen NVIDIA H200 GPU eine etwa 3,82-fache Beschleunigung gegenüber dem ursprünglichen Indexierer-Kernel von DSA.

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.