Source-linked AI summary
HISA: Efficient Hierarchical Indexing for Fine-Grained Sparse Attention
Yufei Xu, Fanxu Meng, Fan Jiang, Yuxuan Wang, Ruijie Zhou, Zhaohui Wang, Jiexi Wu, Zhixin Pan, Xiaojuan Tang, Wenjie Pei, Tongxuan Liu, Di Yin, Xing Sun, Muhan Zhang
TL;DR
Token-level sparse attention makes downstream computation efficient, but its indexer still scans the full prefix for every query, creating a long-context bottleneck. HISA replaces that scan with block-level coarse filtering followed by token-level refinement, preserving the token-set interface and requiring no additional training. The paper reports substantial kernel speedups while maintaining nearly unchanged benchmark performance relative to DSA.
Problem
Token-level sparse attention still incurs full-prefix indexing for every query, leaving per-layer cost at O(L^2) as contexts grow.
Method
HISA hierarchically scores pooled block representatives, retains candidate blocks, and applies the original token indexer only within those blocks.
Results
HISA achieves up to 3.75× indexer speedup at 64K context and maintains nearly unchanged LongBench performance while outperforming the corresponding block-sparse baseline on NIAH.
Takeaways & Limitations
HISA is a training-free, plug-and-play indexer replacement that preserves the downstream Sparse MLA interface while reducing the search path.
Takeaways & Limitations
Mean-pooled block representations can lose information when blocks cross semantic boundaries, motivating overlapping or adaptive blocks and improved pooling.
Abstract
from arXiv · showhide
Token-level sparse attention mechanisms, exemplified by DeepSeek Sparse Attention (DSA), achieve fine-grained key selection by scoring every historical key for each query through a lightweight indexer, then computing attention only on the selected subset. While the downstream sparse attention itself scales favorably, the indexer must still scan the entire prefix for every query, introducing an per-layer bottleneck that grows prohibitively with context length. We propose HISA (Hierarchical Indexed Sparse Attention), a plug-and-play replacement for the indexer that rewrites the search path from a flat token scan into a two-stage hierarchical procedure: (1) a block-level coarse filtering stage that scores pooled block representations to discard irrelevant regions, followed by (2) a token-level refinement stage that applies the original indexer exclusively within the retained candidate blocks. HISA preserves the identical token-level top-sparse pattern consumed by the downstream Sparse MLA operator and requires no additional training. On kernel-level benchmarks, HISA achieves up to speedup at 64K context. On Needle-in-a-Haystack and LongBench, we directly replace the indexer in DeepSeek-V3.2 and GLM-5 with our HISA indexer, without any finetuning. HISA closely matches the original DSA in quality, while substantially outperforming block-sparse baselines.
1 Introduction
Long-context sparse attention reduces downstream computation but leaves token-level indexing as a full-prefix bottleneck. HISA addresses this by replacing flat scanning with hierarchical block filtering and token refinement while preserving compatibility with existing sparse attention.
- Long-context self-attention faces dominant prefill-latency and memory costs as context windows expand toward 1M tokens.
- Token-level sparse attention selects top-k historical keys, but its indexer still scores every prefix token for each query.
- O(L^2) per-layer indexing can become the dominant cost as context length grows, despite sparse downstream attention.
- HISA first scores pooled contiguous blocks, retains the top-m blocks, and then reranks at most mB tokens within those candidates.
- HISA preserves the selected token-set interface, reducing indexing complexity from O(L) to O(L/B + mB) per query without retraining or attention-architecture changes.
- The paper reports 2–4× kernel-level speedup at 64K contexts and performance comparable to DSA on Needle-in-a-Haystack and LongBench.
2 Related Work
Sparse-attention research spans hardware-friendly block selection, fine-grained token selection, and hierarchical coarse-to-fine designs. HISA is presented against these related strategies while retaining token-level outputs for downstream sparse attention.
- Block-sparse methods select fixed-size regions, improving GPU compatibility but forcing all tokens in a block to be retained or discarded together.
- HISA combines hierarchical coarse filtering with token-level refinement while producing the same per-query token-index structure as DSA.
- Token-sparse methods provide finer selection but face efficient importance-estimation challenges, including limited handling of layer- and query-specific variation.
- Recent hierarchical methods score block summaries globally before activating finer-grained attention within selected blocks.
3 Preliminary
DSA separates token-wise indexing from Sparse MLA: the indexer selects relevant token positions, and Sparse MLA attends only to their latent entries. The indexer remains quadratic across a layer because every query scans the prefix.
- DSA consists of a token-wise indexer and a downstream Sparse MLA operator.
- The indexer uses lightweight indexing keys and per-head gating weights to compute query-key relevance scores.
- The indexer selects top-k token indices and passes them to Sparse MLA.
- Because each query scores the full prefix, DSA indexing costs O(L^2) across all queries in a layer.
- Sparse MLA attends only to selected latent entries, reducing main attention cost from dense O(L^2) to sparse O(Lk).
- The selected token set is the interface HISA replaces while leaving the downstream Sparse MLA operator unchanged.
4 Method
HISA replaces DSA’s flat prefix scan with a coarse-to-fine hierarchy: pooled block scoring selects candidate blocks, then token-level refinement selects the final sparse tokens. The design preserves DSA’s downstream interface while reducing indexing cost and targeting efficiency gains in long-context settings.
- Compatibility and boundary behavior: HISA preserves the original per-query token-set output, so Sparse MLA, the attention architecture, and the KV-cache layout require no changes or retraining.This makes HISA a drop-in indexer replacement.
- Block partitioning and pooled keys: HISA partitions the causal prefix into contiguous blocks and mean-pools their indexing keys into representative keys for coarse scoring.The pooled representatives are used only as block-level proxies; token-level indexing keys and Sparse MLA KV states remain unchanged.
- Block-level coarse filtering: Stage 1 scores pooled block representatives and selects the top-m blocks, whose tokens form the candidate set for refinement.The selected candidate token set is the union of tokens in the retained blocks.
- Token-level refinement: Stage 2 applies the original DSA token-level scoring mechanism within the candidate set and selects the final top-k tokens.Sparse MLA then consumes the selected tokens using the same computation as in DSA.
- Boundary behavior: When t ≤ k, HISA selects all prefix tokens and is equivalent to dense attention.This is the short-prefix boundary regime.
- Boundary behavior: When k < t ≤ mB, HISA is equivalent to DSA; when t > mB, block pruning activates and its hierarchical advantage grows with sequence length.The third regime is the long-context setting targeted for efficiency gains.
- Complexity analysis: HISA’s per-query indexing cost is O(L/B + mB), and its per-layer cost is O(L^2/B + LmB), versus O(L^2) for the original indexer.The reduction is substantial when m ≪ M and B ≪ L, while performance degrades gracefully toward DSA as m approaches M.
5 Experiments
HISA is evaluated against DSA and Block-Sparse across kernel latency, retrieval, LongBench, attention visualization, and hyperparameter sensitivity. It closely preserves DSA quality while improving indexer efficiency and outperforming block-only selection.
- Kernel-Level Speedup: 2.16× speedup at 64K context with a 4:1 first-stage compression ratio, and up to 3.75× speedup under a fixed 8K budget.These measurements compare indexer overhead using TileLang kernels; end-to-end throughput also depends on Sparse MLA, KV-cache management, and other system components.
- Needle-in-a-Haystack: HISA closely matches DSA retrieval accuracy across context lengths and needle positions, with only marginal degradation at extreme lengths and depths.Block-Sparse shows noticeable degradation, especially when the needle is in the middle of the context.
- LongBench Evaluation: Across both models and all LongBench task categories, HISA performs very close to DSA and consistently surpasses DSA on Synthetic tasks.On GLM-5, HISA attains a higher average score, while Block-Sparse has a substantially larger gap and declines by 8.35% on GLM-5 Synthetic tasks.
- Visualization of Attention Scores: Attention weights tend to form contiguous, semantically coherent spans that persist across layers, while scores outside these spans are negligible.The visualization uses token positions on the x-axis and layer indices on the y-axis for a representative DeepSeek-V3.2 LongBench code sample.
- Hyperparameter Sensitivity: All three HISA configurations closely track DSA across five LongBench categories, while B=256 performs worse than the finer B=64 and B=128 settings.The configurations share mB = 8192 and k=2048 but vary block size and block-level top-m; Block-Sparse consistently underperforms HISA.
- Hyperparameter Sensitivity: Token-level refinement improves quality by pruning low-relevance tokens within selected blocks rather than retaining every token in those blocks.This distinction separates HISA from the Block-Sparse baseline, which performs block selection without token-level refinement.
6 Conclusion and Future Directions
HISA addresses the DSA indexer’s O(L^2) bottleneck with hierarchical filtering and achieves substantial kernel speedup while preserving long-context quality without additional training. Future work targets coarse-filter information loss, training-aware filtering, and end-to-end serving integration.
- 3.75× speedup over the DSA kernel is achieved at the kernel level with HISA’s hierarchical indexing approach.HISA directly replaces the token indexer in DeepSeek-V3.2 and GLM-5 without additional training.
- HISA maintains nearly unchanged LongBench performance and significantly outperforms the corresponding block-sparse baseline on NIAH.
- Future Directions: Mean-pooled block representations can lose information when blocks cross semantic boundaries and omit important tokens.
- Future Directions: Proposed mitigations include overlapping or adaptive blocks, max pooling, jointly training block scoring, and integrating HISA into realistic serving stacks.
A Algorithm Pseudocode
Algorithm 1 specifies HISA’s hierarchical indexer, including its inputs, block partitioning, coarse filtering, candidate expansion, token refinement, and final top-k selection.
- The indexer takes query and token indexing representations, block size B, block budget m, and token budget k as inputs.
- The prefix is partitioned into M = ⌈L/B⌉ contiguous blocks before block representations are computed.
- Stage 1: Block-level coarse filter: For each query, Stage 1 scores preceding blocks and selects the top m blocks together with the first and last blocks.
- Stage 2: Token-level refinement: Stage 2 applies token-level refinement over the union of tokens in the selected candidate blocks.
- The algorithm returns the k highest-scoring tokens from the refined candidate pool.
B Experimental Settings
The long-context evaluations use a zero-shot experimental setting.
- All long-context evaluations were conducted in a zero-shot setting.
B.1 Long-context Benchmarks
Long-context performance is evaluated on NIAH and LongBench using DeepSeek-V3.2 and GLM-5 under model- and task-specific serving configurations with within-task comparisons aligned.
- NIAH and LongBench are used to evaluate long-context performance on DeepSeek-V3.2 and GLM-5.
- Both models are deployed with the vLLM online serving framework using FP8 precision.
- NIAH Settings: NIAH uses a customized RULER2-derived evaluation codebase without chat templates to assess raw retrieval capabilities directly.
- LongBench Settings: LongBench uses lm-eval3 with configurations that vary according to model characteristics.
- LongBench Settings: DeepSeek-V3.2 uses its standard chat template, whereas GLM-5 is evaluated without one because template-based thinking exceeded generation limits and slowed inference.
- LongBench Settings: Concurrency defaults to 20 requests but is reduced for selected GLM-5 tasks because of out-of-memory issues.
- Fairness of Comparison: Methods are compared under strictly aligned settings within each model and task combination despite differences across models and tasks.