MISA: Mengsel van Indexer Schaarse Aandacht voor Lange-Context LLM Inferentie

Samenvatting

DeepSeek Sparse Attention (DSA) stelt de nieuwste standaard voor fijnmazige sparse aandacht tijdens inferentie door een aangeleerde token-gewijze indexeerder te introduceren die elke prefix-token scoort en de meest relevante selecteert voor de hoofdaandacht. Om expressief te blijven, gebruikt de indexeerder meerdere query-koppen (bijvoorbeeld 64 op DeepSeek-V3.2) die dezelfde geselecteerde tokenset delen; dit multi-head ontwerp is precies wat de indexeerder de dominante kostenpost maakt bij lange contexten. Wij stellen MISA (Mixture of Indexer Sparse Attention) voor, een directe vervanging voor de DSA-indexeerder die de indexeerder-koppen behandelt als een pool van Mixture-of-Experts. Een lichtgewicht router gebruikt goedkope blokniveau-statistieken om een query-afhankelijke subset van slechts enkele actieve koppen te kiezen, en alleen die koppen voeren de zware token-niveau scoring uit. Dit behoudt de diversiteit van de oorspronkelijke indexeerder-pool terwijl de kosten per query worden verminderd van het scoren van elke prefix-token met elke kop tot het scoren ervan met slechts een handvol gerouteerde koppen, plus een verwaarloosbare routerterm die wordt berekend op een kleine set samengevoegde sleutels. We introduceren verder een hiërarchische variant van MISA die de gerouteerde pas gebruikt om een vergrote kandidaatset te behouden en deze vervolgens herrangschikt met de originele DSA-indexeerder om de uiteindelijk geselecteerde tokens vrijwel exact te herwinnen. Met slechts acht actieve koppen en zonder extra training evenaart MISA de dichte DSA-indexeerder op LongBench voor DeepSeek-V3.2 en GLM-5, terwijl het respectievelijk acht en vier keer minder indexeerder-koppen gebruikt, en presteert het gemiddeld beter dan HISA. Het behoudt ook volledig groene 'Needle-in-a-Haystack'-heatmaps voor contexten tot 128K tokens en herwint meer dan 92% van de door de DSA-indexeerder per laag geselecteerde tokens. Onze TileLang-kernel levert een ruwe versnelling van 3,82 keer ten opzichte van DSA's originele indexeerder-kernel op een enkele NVIDIA H200 GPU.

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.