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

MISA is a Mixture-of-Experts mechanism for indexer sparse attention in long-context LLM inference. Its core innovation is treating the index heads of DeepSeek Sparse Attention as an expert pool, dynamically selecting a small number of active heads (only 8 in experiments) via a lightweight router for token-level scoring. Without additional training, its performance is comparable to the original 64-head indexer, while achieving a 3.82x kernel speedup.

As LLM text processing length expands, the O(n²) complexity of standard self-attention becomes a bottleneck. Sparse attention reduces cost by selecting important token pairs. DeepSeek Sparse Attention (DSA) introduces a learnable token-level indexer to implement token scoring, dynamic selection, and multi-head sharing. However, the indexer uses a large number of query heads (e.g., 64), leading to computational burden—under long contexts, the indexer's cost even exceeds that of the main attention.

MISA optimizes via the Mixture-of-Experts mechanism:

1. Core Architecture: A lightweight router first performs block-level statistics to capture coarse-grained query patterns, then dynamically selects 8 active heads. Only active heads execute token-level scoring, reducing computation.
2. Hierarchical Variant: First route to expand the candidate set, then reorder using the original DSA to balance efficiency and quality, recovering over 92% of the originally selected tokens.

Experimental results:

• LongBench Benchmark: DeepSeek-V3.2 with 8 heads achieves performance equivalent to the 64-head version; GLM-5 with 8 heads matches the 32-head version.
• Needle-in-a-Haystack: Maintains a fully green heatmap under 128K context, with no key information missing.
• Comparison with HISA: Outperforms on average with higher efficiency.
• Kernel Speedup: Achieves a 3.82x speedup on NVIDIA H200, thanks to memory optimization and parallelism improvement.

MISA’s prominent advantage is zero-training plug-and-play: No fine-tuning or retraining of pre-trained models is required; existing DSA models can be directly replaced, avoiding training performance degradation, eliminating the need for expensive computing resources, and enabling rapid deployment.

Technical Insights:

• Expert redundancy is widespread; 8 out of 64 heads can achieve similar performance.
• Dynamic computation (conditional computation) is worth exploring.
• Hierarchical design (coarse screening + fine ranking) is universal.
• Algorithm-system synergy (e.g., TileLang kernel) unlocks potential. Future directions can further explore the application of dynamic computation and hierarchical design in more scenarios.

MISA optimizes the sparse attention indexer from a Mixture-of-Experts perspective, improving computational efficiency without sacrificing quality. It not only enhances long-context LLM inference performance but also demonstrates the potential of dynamic routing and hierarchical design. As context length grows, such efficient sparse attention technologies will become increasingly important.