Summary (Overview)

  • MISA (Mixture of Indexer Sparse Attention) is a drop-in replacement for the DeepSeek Sparse Attention (DSA) indexer that treats its HIH^I indexer heads as a pool of mixture-of-experts (MoE) experts, routing only hHIh \ll H^I active heads per query.
  • A lightweight block-pooled router selects the query-dependent subset of heads using cheap statistics on M=L/BLM = \lceil L/B \rceil \ll L pooled keys, reducing per-query indexer cost from O(HIL)\mathcal{O}(H^I L) to O(hL+HIM)\mathcal{O}(hL + H^I M).
  • A hierarchical variant, MISA^{\dagger}, uses the MoE-routed pass to keep an enlarged candidate set (size k>kk' > k) and re-ranks it with the full DSA indexer to recover the final top-k almost exactly.
  • Without any additional training, MISA matches dense DSA on LongBench within 0.5 average points on both DeepSeek-V3.2 and GLM-5 while using only h=8h=8 active heads (8× and 4× fewer heads respectively), outperforms HISA on average, and preserves fully green Needle-in-a-Haystack heatmaps up to 128K context.
  • The TileLang kernel implementation delivers approximately a 3.82× speedup over DSA's original indexer kernel on a single NVIDIA H200 GPU.

Introduction and Theoretical Foundation

Frontier large language models now routinely process prefixes of hundreds of thousands to millions of tokens, making dense attention the dominant cost of both prefill and decode. Among sparse attention techniques, DeepSeek Sparse Attention (DSA) stands out as the best-performing fine-grained variant in production: it uses a lightweight learned indexer that scores every prefix token and feeds the top-k tokens into the main attention. This design carries into DeepSeek-V4's Compressed Sparse Attention (CSA), confirming that learned token-wise indexing remains the strongest building block.

The central problem: DSA's indexer is multi-head — although the main attention operates in MQA mode with a single shared KV entry per token, DSA uses HI=64H^I = 64 indexer heads. Each head specialises in a different relevance pattern (recency, syntactic role, lexical/semantic similarity), and the aggregated score benefits from this diversity. However, scoring each of the LL prefix tokens with all HIH^I heads makes the indexer the dominant cost on long contexts.

The key observation: While diversity across heads is essential when aggregated over a large pool, only a few heads are actually informative for any given query. The relevant set changes slowly along the prefix and can be identified from cheap block-level statistics.

The indexer relevance score in DSA is:

It,s=j=1HIwt,jIReLU(qt,jIksI)(1)I_{t,s} = \sum_{j=1}^{H^{I}} w_{t,j}^{I} \cdot \mathrm{ReLU}\big(\mathbf{q}_{t,j}^{I} \cdot \mathbf{k}_{s}^{I}\big)\tag{1}

with the top-k token selection:

Tt=TopK(It,:,k)(2)\mathcal{T}_{t} = \operatorname{TopK}(I_{t,:}, k)\tag{2}

Methodology

1. MISA: Mixture of Indexer Experts

MISA treats the HIH^I indexer heads as a pool of MoE experts. The method has two components:

Block-pooled router. The prefix is partitioned into M=L/BM = \lceil L/B \rceil blocks B1,,BM\mathcal{B}_1, \ldots, \mathcal{B}_M, each summarized by a pooled indexing key:

k~bI=Pool({ksIsBb})(4)\tilde{\mathbf{k}}_{b}^{I} = \operatorname{Pool}\bigl(\{\mathbf{k}_{s}^{I} \mid s \in \mathcal{B}_{b}\}\bigr)\tag{4}

For query position tt, the router computes per-head per-block affinities and aggregates across blocks:

At,j,b=wt,jIReLU(qt,jIk~bI),Et,j=1Mb=1MAt,j,b(7)A_{t,j,b} = w_{t,j}^{I} \cdot \mathrm{ReLU}\Big(\mathbf{q}_{t,j}^{I} \cdot \tilde{\mathbf{k}}_{b}^{I}\Big), \qquad E_{t,j} = \frac{1}{M} \sum_{b=1}^{M} |A_{t,j,b}|\tag{7}

The top-h heads are selected as the active expert set:

Ht=TopKj(Et,j,h)(8)\mathcal{H}_{t} = \mathrm{TopK}_{j}(E_{t,j}, h)\tag{8}

Sparse token scoring with active experts. Only active heads compute the token-level score:

I^t,s=jHtwt,jIReLU(qt,jIksI)(9)\hat{I}_{t,s} = \sum_{j \in \mathcal{H}_{t}} w_{t,j}^{I} \cdot \mathrm{ReLU}\big(\mathbf{q}_{t,j}^{I} \cdot \mathbf{k}_{s}^{I}\big)\tag{9}

with final selection:

Tt=TopKs(I^t,:,k)(10)\mathcal{T}_{t} = \operatorname{TopK}_{s}\left(\hat{I}_{t,:}, k\right)\tag{10}

Key distinction from HISA: Both methods compute the same per-head per-block affinities At,j,bA_{t,j,b}, but reduce them along orthogonal axes — HISA aggregates across heads to select top-m blocks (token-axis), while MISA aggregates across blocks to select top-h heads (head-axis). MISA uses a much coarser block partition (B=1024B=1024 vs. HISA's B=128B=128), keeping routing overhead negligible.

2. Hierarchical MISA (MISA^{\dagger})

A two-stage coarse-to-fine approach:

  • Coarse stage: MoE scoring selects an enlarged candidate set: Ωt=TopKs(I^t,:,k)\Omega_{t} = \operatorname{TopK}_{s}\left(\hat{I}_{t,:}, k^{\prime}\right) with k>kk' > k (Eq. 11)
  • Fine stage: Original DSA scoring with all HIH^I heads is applied within Ωt\Omega_t: Tt=TopKs({It,ssΩt},k)\mathcal{T}_{t} = \mathrm{TopK}_{s}(\{I_{t,s} \mid s \in \Omega_{t}\}, k) (Eq. 12)

Unlike HISA, MISA^{\dagger} keeps the coarse pass at full token granularity (head-level routing rather than block-level filtering), so tokens within the same block are still ranked individually.

Empirical Validation / Results

LongBench Results (Table 1)

All sparse methods applied at inference time without additional training. Token budget fixed at k=2048k = 2048 across all methods.

ModelIndexerHeadsSQAMQASumFSSynCodeAvg.
DeepSeek-V3.2DSA6450.8952.6622.1162.2469.8348.5651.05
Block6448.3649.7621.9059.4568.6749.0949.54
HISA6449.1751.9622.1361.6270.8348.9950.78
MISA850.8351.6222.3161.7770.0048.5450.85
MISA^{\dagger}850.9152.2722.0462.1069.8348.4550.95
GLM-5DSA3241.2327.8918.3963.2068.8456.5346.01
Block3238.3524.2916.9560.6460.4955.2942.67
HISA3242.4527.6217.9063.7869.3556.7946.32
MISA841.6429.1618.2263.5368.7657.2446.43
MISA^{\dagger}841.1628.6718.5663.5769.2657.8346.51

Task abbreviations: SQA = Single-Document QA, MQA = Multi-Document QA, Sum = Summarization, FS = Few-shot Learning, Syn = Synthetic Retrieval, Code = Code Completion.

Key findings:

  • MISA matches DSA within 0.20 points on DeepSeek-V3.2 (50.85 vs. 51.05) and surpasses it on GLM-5 (46.43 vs. 46.01)
  • MISA^{\dagger} closes the gap to 0.1 average points on DeepSeek-V3.2, with every per-category score within 0.4 points of DSA
  • Block-Sparse trails by 1.5–3.4 average points, confirming block-uniform selection is too coarse

Needle-in-a-Haystack (NIAH) Retrieval

Both MISA variants reproduce DSA's near-perfect green grid across the full depth–length plane up to 128K context. Block-Sparse shows visible accuracy holes at intermediate depths beyond ~32K; HISA shows minor degradations at deepest needle positions. MISA^{\dagger} is essentially indistinguishable from DSA.

Indexer Kernel Speed

  • 1-stage MISA: consistently faster than DSA across all sequence lengths
  • 2-stage MISA^{\dagger}: outperforms DSA when sequence length exceeds 32K
  • Achieved speedup: ~3.82× end-to-end over DSA's original indexer kernel on NVIDIA H200

Ablation: Number of Active Heads

  • h=1,2h=1, 2: too aggressive — visible accuracy holes in NIAH heatmap
  • h=4h=4: mitigates most deficiencies but suboptimal at 128K
  • h=8h=8: essentially indistinguishable from dense 64-head indexer
  • h=16h=16: no further gain despite twice the compute of h=8h=8

Default: h=8h=8 as the smallest setting that consistently matches DSA.

Theoretical and Practical Implications

Theoretical contributions:

  1. Head-axis routing as a new efficiency axis: MISA identifies the indexer's per-token head-token products as the dominant cost and introduces head-axis routing, complementary to token-axis hierarchies (HISA) and block-level methods (MoBA, Quest, InfLLM).
  2. Decoupling of routing from output: Unlike prior attention-side MoE methods (Mixture-of-Attention-Heads, MoH) that route the heads producing attention output, MISA routes only the heads producing the indexer score — the downstream attention remains dense over the chosen set. This decoupling enables very small expert counts (hHIh \ll H^I) without harming quality.
  3. Preservation of diversity: Every head remains available in the pool; routing only selects which ones to consult per query, preserving the full expressiveness of the indexer.

Practical implications:

  • Drop-in replacement: MISA requires no additional training and can be inserted into pretrained DSA-based models (DeepSeek-V3.2, GLM-5) with identical interface (same Sparse MLA operator, same token budget)
  • Hardware efficiency: Realized ~3.82× kernel speedup confirms head-axis routing translates to measurable savings
  • Scalability: The approach is complementary to token-level schemes and can be combined orthogonally with methods like HISA and IndexCache
  • Generalizability: The MoE-routed indexer concept could extend to future architectures like DeepSeek-V4's CSA, which applies DSA on compressed KV streams

Conclusion

MISA demonstrates that indexer-head-axis routing is a practical and complementary axis of efficiency for fine-grained sparse attention. The lightweight block-pooled router reduces the dominant per-token cost from O(HIL)\mathcal{O}(H^I L) to O(hL+HIM)\mathcal{O}(hL + H^I M) while preserving full indexer diversityaine. With h=8h=8 active heads and no additional training, MISA:

  • Matches dense DSA on LongBench within 0.5 average points on both models
  • Outperforms both Block-Sparse and HISA on average
  • Retains fully green NIAH heatmaps up to 128K context
  • Recovers more than 92% of DSA-selected tokens per layer
  • Delivers ~3.82× wall-clock kernel speedup on NVIDIA H200

Future directions (as noted in limitations):

  1. Measuring end-to-end model latency (not just indexer kernel)
  2. Reducing memory access volume to the KV cache in the indexer stage
  3. Jointly training the router with the indexer to further close residual quality gaps

The authors hope this work stimulates further exploration of head-level routing in sparse attention systems.

Related papers