PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2605.07363v1cs.LGcs.AI

TL;DR

Long-context DSA indexers are costly because many heads score every prefix token. MISA routes each query to a few heads, matching DSA quality on LongBench while delivering substantial head and kernel-efficiency gains.

  • Problem

    DSA’s expressive multi-head indexer scores every prefix token with all heads, making head–token products the dominant long-context indexer cost.

  • Method

    MISA uses block-pooled routing to select a query-dependent subset of indexer heads, while MISA† re-ranks enlarged routed candidate sets with the full DSA indexer.

  • Results

    With eight active heads and no additional training, MISA matches dense DSA on LongBench within 0.5 average points, retains green 128K Needle-in-a-Haystack heatmaps, and achieves roughly 3.82× kernel speedup.

  • Takeaways & Limitations

    Head-level routing is a practical efficiency axis for fine-grained sparse attention that complements token-level sparsity schemes.

  • Takeaways & Limitations

    Speed experiments measure only TileLang kernel latency, and MISA does not reduce KV-cache memory-access volume or evaluate jointly trained routers.

Abstract

from arXiv · show

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.

1 Introduction

MISA addresses DSA’s dominant long-context indexer cost along the head axis by routing each query to only a few informative heads. It preserves DSA’s token-selection interface while achieving near-dense accuracy without additional training.

  • Design: MISA is a drop-in replacement that produces the same-size top-k token set and feeds the same Sparse Multi-Head Latent Attention operator as DSA.A hierarchical extension, MISA†, re-ranks routed candidates with the original DSA indexer to recover the dense top-k almost exactly.
  • Motivation: DSA uses H_I = 64 indexer heads because head diversity captures distinct relevance patterns, but scoring every prefix token with every head dominates long-context cost.DeepSeek-V3.2’s main attention has H = 128 query heads, while all heads share one selected token set.
  • MISA: MISA treats DSA’s indexer heads as MoE experts and uses cheap block-level statistics to route each query to h ≪ H_I active heads for token-level scoring.The routed heads retain fine-grained token localization while reducing computation along the head axis.
  • Results: Within 0.5 average points, MISA matches dense DSA on LongBench without additional training using h = 8 active heads, with 8× and 4× fewer indexer heads across DeepSeek-V3.2 and GLM-5, respectively.It also preserves full Needle-in-a-Haystack accuracy up to 128K context and recovers more than 92% of DSA-selected tokens per layer on LSHT.

2 Related work

Prior sparse-attention methods reduce long-context cost through static token patterns, cache eviction, block-level retrieval, learned gates, or sampling. Related MoE work motivates sparse routing, while prior attention-side MoE explores conditional computation within attention.

  • Sparse attention: Static sparse-attention methods use predefined windows, strides, or global tokens, decoupling selection from content.Examples include Sparse Transformer, Longformer, and BigBird.
  • Sparse attention: Cache-eviction methods remove tokens during decoding using attention-statistics heuristics, retaining attention sinks and recent tokens.The supplied passage names StreamingLLM and H2O as examples.
  • Sparse attention: Dynamic retrieval methods select content-relevant token blocks using pooled-key scores, trained selectors, locality-sensitive hashing, or learned sparse gates.The passage cites Quest, InfLLM, MoBA, MagicPIG, and SeerAttention.
  • Mixture of experts in language models: Mixture-of-experts language models use learned routers to activate sparse subsets of expert FFNs, while modern open-weight models commonly retain dense attention.The passage references Shazeer et al., GShard, Switch Transformer, Mixtral, and DeepSeek-MoE.
  • Mixture of experts in language models: Attention-side MoE has also been explored by assigning conditional computation to attention heads.The supplied passage introduces Mixture-of-Attention-Heads in this context.

3 Preliminaries

DSA combines a token-wise indexer that selects relevant tokens with Sparse MLA over those tokens, while HISA accelerates indexing through hierarchical coarse-to-fine search. The indexer is the main target for acceleration because its scan cost scales with prefix length, whereas Sparse MLA already reduces attention from O(L2) to O(Lk).

  • DSA overview: DSA has two components: a token-wise indexer selects a small relevant-token set, and Sparse MLA attends only over that set.The selected token set is the interface between the components.
  • Indexer in DSA: The DSA indexer selects top-k token indices from per-head relevance scores and passes them to Sparse MLA.It uses HI indexing heads, lightweight indexing keys, and per-head gating weights.
  • Indexer in DSA: O(HIL) is the per-query DSA indexer cost, while full-prefill cost grows as O(HIL2) per layer.The dominant work is the HI head–token products.
  • Sparse MLA in DSA: O(L2) becomes O(Lk) in Sparse MLA because each token stores one latent key–value entry shared across query heads and attention uses only selected entries.This leaves the indexer as the natural acceleration target without changing Sparse MLA.
  • Indexer in HISA: HISA partitions the prefix into M = ⌈L/B⌉ contiguous blocks, mean-pools each block, scores pooled keys, and selects top-m blocks before token-level rescoring.The first and last blocks are always included to retain the attention sink and local context; original DSA scoring then produces final top-k tokens for unmodified Sparse MLA.

4 Method

MISA replaces DSA’s costly all-head token scoring with query-dependent routing over a pool of indexer heads, using cheap block-level statistics to select only a few active experts. Its coarse-to-fine variant, MISA†, preserves token-level candidates and refines them with the full DSA indexer.

  • Motivation: DSA costs O(H_I L) per query because every prefix token is scored with all H_I heads, whose distinct relevance patterns cannot simply be collapsed.The heads specialize in different relevance patterns, and reducing their number measurably degrades retrieval.
  • MISA indexer: MISA treats the H_I indexer heads as experts and routes each query to h ≪ H_I active heads using cheap block-level statistics.Only the selected experts perform token-level scoring; routing operates on pooled block statistics.
  • MISA indexer: The router costs O(H_I M), while active-expert scoring reduces the indexer cost to O(H_I M + hL), with M = ⌈L/B⌉ ≪ L.MISA aggregates per-head affinities across blocks to select heads, unlike HISA, which aggregates across heads to select blocks.
  • MISA indexer: With h = 8, the dominant hL term is reduced 8× versus DSA for DeepSeek-V3.2 with H_I = 64 and 4× for GLM-5 with H_I = 32.All heads remain available in the pool, preserving indexer-pattern diversity while routing chooses which heads to consult per query.
  • MISA†: MISA† first uses a cheap MoE pass to retain an enlarged token-level candidate set, then applies the unmodified DSA indexer with all H_I heads to refine those candidates.Unlike HISA, MISA† does not use block-level top-m filtering; its candidate set need not be a union of full blocks.

5 Experimental results

MISA preserves DSA-quality retrieval and downstream performance while activating only a small subset of indexer heads at inference time. Across LongBench, NIAH, and kernel benchmarks, hierarchical reranking further recovers dense-indexer behavior and improves efficiency.

  • LongBench: MISA and MISA† activate 8 heads per query, using 1/8 of DeepSeek-V3.2’s and 1/4 of GLM-5’s indexer heads while keeping k = 2048 fixed.All sparse methods run at inference time without additional training.
  • LongBench: 50.85 vs. 51.05 on DeepSeek-V3.2 and 46.43 vs. 46.01 on GLM-5 show MISA matching or surpassing dense DSA on LongBench averages.MISA also outperforms Block-Sparse and HISA on average across both models.
  • LongBench: 50.95 vs. 51.05 reduces hierarchical MISA†’s DeepSeek-V3.2 average gap to DSA to 0.1 points, with every category within 0.4 points.Single-Document QA improves marginally under MISA†.
  • Needle-in-a-Haystack retrieval: Up to 128K context, both MISA variants reproduce DSA’s near-perfect NIAH retrieval grid, while MISA† is essentially indistinguishable from DSA.MISA† combines routed head selection with a token-level fine pass and operates with 1/8 of the per-token head–token products.
  • Kernel efficiency: 1-stage MISA is faster than DSA across all sequence lengths, while 2-stage MISA† is faster beyond 32K; memory traffic and load imbalance limit the theoretical speedup.The comparison is measured on a single NVIDIA H200 GPU.

6 Conclusion

MISA replaces the DSA indexer with routed sparse attention: a lightweight router activates only a query-dependent subset of heads for heavy token scoring while preserving the full indexer pool. With eight active heads and no additional training, it matches dense DSA on LongBench within 0.5 average points, retains fully green 128K Needle-in-a-Haystack heatmaps, and achieves roughly 3.82× kernel speedup over DSA.

  • 6 Conclusion: MISA uses a block-pooled router to select h ≪ H_I active heads, while only those heads perform the heavy token-level scan.Every indexer head remains available; routing chooses which heads to consult for each token.
  • 6 Conclusion: The routed design reduces dominant per-token cost from O(H_I L) to O(hL + H_I M) while preserving indexer-head diversity.The reduction comes from scoring prefix tokens with only the routed heads, alongside the router term over pooled keys.
  • 6 Conclusion: 0.5 average points is the maximum LongBench gap between MISA and the dense DSA indexer with h = 8 active heads and no additional training.This configuration corresponds to an 8× head reduction on DeepSeek-V3.2 and a 4× reduction on GLM-5, and MISA outperforms Block-Sparse and HISA on average.
  • 6 Conclusion: 128K context is supported by a fully green Needle-in-a-Haystack heatmap under the eight-active-head configuration.The result is reported for the two open-weight long-context models evaluated in the experiments.
  • 6 Conclusion: 3.82× wall-clock speedup is delivered by the TileLang kernel implementation over DSA's original indexer kernel.The speedup translates the routed indexer's computational savings into measured runtime improvement.

7 Limitation

The paper leaves end-to-end latency, KV-cache memory traffic, and router-training questions unresolved. Its speed measurements cover only the TileLang kernel, while experiments use pretrained DSA-based models without finetuning.

  • Speed experiments measure TileLang-kernel latency rather than full-model end-to-end latency.
  • MISA reduces DSA-stage computation but does not reduce KV-cache memory-access volume.
  • All reported results insert MISA into pretrained DSA-based models without finetuning.

A Per-layer agreement with the DSA top-k

MISA† is evaluated against the DSA top-2048 token set using per-layer IoU on the long LSHT subset of DeepSeek-V3.2. It consistently outperforms HISA across sequence lengths and layers, retaining more than 92% of DSA-selected tokens.

  • Per-layer agreement with the DSA top-k: MISA† uses per-layer IoU against the DSA top-2048 reference set on 10 long LSHT examples across all T layers of DeepSeek-V3.2.The reference uses all H_I = 64 DSA indexer heads.
  • Per-layer agreement with the DSA top-k: MISA† first generates k′ = 8192 candidates, then re-scores them with all H_I = 64 DSA heads to select the final k = 2048 tokens.This two-stage configuration enables comparison under an equivalent computational budget.
  • Per-layer agreement with the DSA top-k: More than 92% of the tokens selected by DSA are retained by MISA†, which consistently outperforms HISA across sequence lengths and at every layer.Figure 5 reports IoU against the DSA baseline as a function of token position and layer index.

B Ablation: active heads’ number of MISA†

The MISA† ablation varies the active-head count h in its coarse routed stage. It uses an enlarged candidate set followed by full DSA re-ranking, allowing aggressive head reduction while retaining relevant tokens for final selection.

  • Ablation: active heads’ number of MISA†: MISA† routes h active heads to select an enlarged candidate set of k′ = 4k = 8192 tokens before DSA refinement.The full HI = 64-head DSA indexer re-ranks the candidate set to extract the final k = 2048 tokens.
  • Ablation: active heads’ number of MISA†: Because the routed stage only needs to retain relevant tokens in the candidate set, DSA refinement can compensate for aggressive head reduction.This hierarchical design need not pinpoint the exact top-k during the routed pass.
  • Ablation: active heads’ number of MISA†: The ablation evaluates MISA† on DeepSeek-V3.2 with HI = 64, B = 1024, and 128K-context NIAH retrieval-accuracy heatmaps.The heatmaps vary context length on the x-axis and needle depth on the y-axis.

C Ablation: routing score

The routing score determines which heads are active in MISA, making its design critical. Among three candidates tested on DeepSeek-V3.2 at 128K, only the proposed block-pooled attention score consults prefix content and fully recovers DSA’s accuracy.

  • Routing-score candidates: At h = 8, B = 1024, and k = 2048, the study compares three routing scores using Needle-in-a-Haystack accuracy at 128K on DeepSeek-V3.2.The candidates are the indexer gating weight, the query-head ℓ2 norm, and the proposed block-pooled attention score.
  • Routing-score candidates: The gating-weight and query-norm proxies do not consult prefix content and both collapse on harder regions of the NIAH grid.The gating weight selects heads up-weighted by the DSA aggregator, while the query norm selects heads with more confident directions.
  • Routing-score candidates: The proposed block-pooled attention score is the only candidate that aggregates query-to-prefix evidence and fully recovers DSA’s accuracy.It averages attention over pooled prefix blocks before routing, directly incorporating the content that must be retrieved.

D Ablation: router block size

The router block-size sweep on DeepSeek-V3.2 shows that NIAH retrieval accuracy is largely insensitive to B across 128 to the full 131,072-token prefix, with only mild degradation at the largest blocks. Larger B reduces router cost by decreasing the number of pooled keys but lowers the spatial resolution of head-importance estimates.

  • Router block size: O(HIM): router cost depends on M = ⌈L/B⌉ pooled keys, so larger B trades lower router cost for coarser head-importance estimates.B controls the number of pooled keys and therefore the spatial resolution available to the router.
  • Router block size: 128 to 131,072: NIAH retrieval accuracy remains largely insensitive to router block size B across the full sweep.The sweep uses h = 8 and k = 2048 at 128K context.
  • Router block size: Only the largest block sizes mildly degrade NIAH accuracy, when routing is forced into a near-global pooled-key set.At the largest B, the router has near-global rather than spatially resolved pooling.

NeurIPS Paper Checklist

The checklist affirms that the paper’s claims, theoretical results, reproducibility information, and experimental details are adequately supported. It also directs readers to the paper’s discussion of limitations in Section 7.

  • Claims and scope: The abstract and introduction accurately reflect the paper’s theoretical and experimental contributions and scope.The justification points to Sections 4 and 5.
  • Limitations: The paper discusses its limitations in Section 7.The checklist identifies limitations as a dedicated discussion area and encourages addressing assumptions and robustness.
  • Theoretical analysis: All theoretical results have their assumptions and complete proofs provided in Section 4.The checklist answer is [Yes].
  • Reproducibility: The paper provides detailed experimental information across LongBench, NIAH, kernel speed, and head-count ablation studies to support reproducibility.The checklist answer is [Yes].
  • Experimental setting/details: The experimental setting specifies the models, datasets, hyperparameters, and hardware needed to understand and reproduce the results.The checklist answer is [Yes], with these details included in Section 5.
Loading 2605.07363v1…