Subquadratic Attention with a Hierarchical Tree-Structured Indexer

Leon Lufkin; Cutter Dawes; Alexander Speicher; Kamesh Krishnamurthy; Tomás Figliolia; Beren Millidge

cs.CLOct 30, 2026

Subquadratic Attention with a Hierarchical Tree-Structured Indexer

Authors: Leon Lufkin, Cutter Dawes, Alexander Speicher, Kamesh Krishnamurthy, Tomás Figliolia, Beren Millidge

Organizations: Zyphra Research

Abstract

The quadratic scaling of attention is a fundamental bottleneck limiting the long-context capabilities of large language models (LLMs). A variety of approaches have attempted to circumvent this, but they either fail to match attention's expressiveness (e.g., sliding-window or linear attention variants with fixed-size state) or only achieve constant-factor speedups (e.g., sparse attention methods whose indexers themselves scale quadratically). In this paper, we derive an asymptotically subquadratic sparse attention method based on a hierarchical, tree-structured indexer, which we call Zyphra Hierarchical Attention (ZHA). ZHA constructs a tree of recursively compressed key summaries over the KV cache, which future queries traverse from the top down to retrieve a sparse subset of keys. Crucially, by iteratively growing the tree's depth, the total cost including routing and final attention grows sublinearly in sequence length, resulting in overall subquadratic (but superlinear) scaling for attention. We demonstrate that ZHA can be cheaply applied to existing pretrained models with only a short KL distillation training stage over a few billion tokens, and maintains performance on long-context retrieval and challenging reasoning, mathematical, and agentic coding evaluation tasks. Moreover, we show that our method is highly robust to tree geometry including very deep trees of up to 6 levels. Finally, we develop a detailed performance model and show that our approach scales subquadratically, dramatically outperforms competing methods beyond 1 million context, and ultimately has the potential to extend effective context windows to the hundreds of millions or billions of tokens.

Figures & tables

Figure 1: Dense top-

p

set size versus causal context length in Qwen3.6-35B-A3B on six fiction books (averaged across books, layers, and heads). We observe that at all

p<1

(shades of blue), the top-

p

set size increases sublinearly as a power law, indicating that longer sequences become increasingly compressible.

Figure 2: Fraction of KV entries of a 256K sequence covered by the union of top-

p

attention subsets as the number of future queries increases, for several values of

p

. We observe that for the higher

p

values we use in practice, a substantial fraction of the total KV cache is included in the union of the queries.

Figure 3: Overview of the ZHA architecture: (a) tree construction and (b) sparse routing. During tree construction, for each chunk of contiguous key nodes, we construct a compressed key node at the higher level by an attention-weighted sum of the keys, where the attention is computed with respect to the learned reference queries. In practice, we use multiple reference queries and store the key nodes for each in parallel. This process continues recursively with higher level key nodes being constructed from the summaries of lower levels. During routing, we start at the top of the tree and compute dot-product similarity scores between our query and the key nodes of that level. We select the set of nodes to expand using top-

p

, and then recursively descend the tree until we reach the leaf nodes corresponding to raw KV cache entries. The union of all the leaves we reach defines the sparse KV cache subset that we dispatch to the main attention.

Figure 4: Illustrative ZHA tree geometries of the three core tree parameters. We can trade off depth and width for a given level of compression in the tree, since compression scales as

C^{L}

. Secondly, at the cost of additional compute, we can increase the number of reference queries.

Figure 5: (a) Reasoning and (b) retrieval evaluation scores with dense attention, MSA, and deep (C4 L6) or wide (C256 L1) ZHA (

p=0.95

R=16

, converted from Qwen3.6-35B-A3B). Appendix B gives evaluation details and results across routing thresholds. We observe that ZHA with both wide and deep geometries maintains performance even at high sparsity levels, while MSA suffers performance degradation especially on retrieval tasks such as RULER and LongBench v2.

Figure 6: Geometry and sparsity have predictable effects on RULER accuracy (mean across tasks). Greater sparsity without degrading performance is achieved by: (a) decreasing compression

C^{L}

, (b) increasing depth at fixed compression (

C^{L}=4096

), or (c) making key summaries more expressive (with increasing reference-query number

R

) at fixed chunk size and depth. (The horizontal dotted lines in panels (a)–(c) show dense attention.) (d) For the simple geometry-to-performance model, predictions versus measurements across 64K, 128K, and 256K, holding out each length in turn (RMSE: 1.44 percentage points).

Figure 7: Per-level selection scales sublinearly in the ZHA router with predictable power-law exponent. (a)–(e): For ZHA geometries with depths

L=2

–

6

(

C=4

R=16

p=0.8

), per-level selection counts (normalized by first-bin count per-level) on AIME problems fitted by power-law scaling curves. (f): Per-level exponents fitted by cumulative top-down model, demonstrating that the across-depth variation in their exponents is well-described by a simple model that accumulates top-down contributions.

Figure 8: For a fixed top-p, and hence sparsity level, ZHA (here, $R=1$ , $C=8$ , $p=0.75$ ) scales subquadratically according to our decode roofline model, eventually dramatically outperforming both dense attention and the MSA sparse attention variant. Importantly, we have not empirically validated the scaling at these ultra-long context lengths due to the lack of existing LLMs which can operate at these contexts. Because of this, while we can obtain subquadratic latency, it is unclear to what extent reasoning and retrieval performance degrades at the level of depths required for these ultra-long contexts. Latencies are modeled per-token decode times for the attention subsystem of one layer at batch size 1.

Figure 9: Reference-query selection rates per tree level for the C4 L6 ZHA router on natural text. (a) Within each KV group and level, reference queries are independently ranked by selection frequency. Each point shows the level’s share of selections at that rank, averaged over layers and KV groups. Higher levels concentrate selections among fewer reference queries, while lower levels use them more uniformly. (b) Spearman correlation between levels’ selection rates. Higher levels have strongly aligned reference-query rankings. Agreement generally weakens as levels become further separated, with the chunk level showing the greatest divergence from the upper levels.

Figure 10: (a) Cosine similarity distribution of every pair of reference queries within a KV group; (b) norm distribution of every reference query. Reference queries within a KV group become slightly more similar and their norms more concentrated as the tree gets shallower and its chunks larger.

Appendix figures & tables34 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 11: Dense top- $p$ set size for individual query heads in Qwen3.6-35B-A3B, using the same six-book cohort and thresholds as Figure 1 . Rows identify full-attention layers and columns identify query heads, with zero-based indices. Each panel averages only across books. All panels share logarithmic axis limits; colors denote $p$ and the dotted black line gives the $p=1$ reference.

Table I : Percentage of individual rollout–layer–head fits for which the affine model beats the power law,

\rho=(1-R^{2}_{\mathrm{lin}})/(1-R^{2}_{\mathrm{pow}})<1

, evaluated in log-

y

space at the leaf level. Raw and Cum-max denote raw and cumulative-maximum targets. Each percentage uses 9,600 paired fits (30 problems

\times

16 rollouts

\times

10 layers

\times

2 KV heads), without averaging across rollouts before fitting.

Figure 12: Dense top- $p$ set size for Nemotron-3.5-Lightning-30B-A3B using the same six-book cohort as Figure 1 .

	Qwen ZHA		Nemotron ZHA		Qwen MSA
	Deep	Wide	Deep	Wide	Index
Trained module	ref. queries	ref. queries	ref. queries	ref. queries	index proj. + norms
Geometry	C4 F4 L6	C256 L1	C4 F4 L6	C256 L1	B128 top-16
Sparse-attention layers	10 / 40	10 / 40	6 / 52	6 / 52	10 / 40
Trainable params	81,920	81,920	24,576	24,576	7.87M
Frozen backbone params	34.66B	34.66B	31.58B	31.58B	34.66B

Appendix

Table II: Training setup and compute for the evaluated ZHA routers and MSA index. Backbones are frozen; ZHA trains only reference queries. Deep and wide ZHA use

R=16

Router	L0	L1	L2	L3	L4	L5
Qwen deep	0.021	0.097	0.180	0.253	0.302	0.319
Qwen wide	0.242	–	–	–	–	–
Nemotron deep	0.016	0.069	0.137	0.192	0.217	0.221
Nemotron wide	0.179	–	–	–	–	–

Appendix

Table III: Final per-level ZHA training KL, averaged over logged points in the last 100 steps.

Benchmark	Evaluated subset	Reported metric	Generation budget	Decoding / harness
AIME 2026 [ 32 ] 1	30 problems, 16 samples each	avg@16 (%)	81,920 output tokens	Thinking; temperature 1.0, top- $p$ 0.95, top- $k$ 20, presence penalty 1.5
SWE-bench Verified [ 22 , 5 ]	500 tasks, 1 run each	Resolved tasks (%)	32,768 tokens/turn; 204,800-token context; 250 steps	mini-swe-agent 2.4.6; bash text actions; thinking retained across turns; temperature 1.0, top- $p$ 0.95, top- $k$ 20
Terminal-Bench 2.1 [ 40 ] 2	All 89 tasks, 5 attempts each	Mean per-task success rate (%)	81,920 tokens/turn; 262,144-token context; 3 h per trial	Harbor 0.24 / Terminus-2; thinking retained across turns; temperature 1.0, top- $p$ 0.95, top- $k$ 20
MMLU-Pro [ 43 ]	Full test split: 12,032 questions; zero-shot CoT; 1 sample each	Accuracy (%)	81,920 output tokens	Thinking; temperature 1.0, top- $p$ 0.95, top- $k$ 20
GPQA Diamond [ 36 ]	All 198 questions, 8 samples each	avg@8 (%)	81,920 output tokens	Thinking; temperature 1.0, top- $p$ 0.95, top- $k$ 20
RULER-13 [ 16 ]	13 tasks $\times$ 100 examples at 256K; 64K and 128K in appendix	Mean task score (%)	Official per-task caps: 30–128 tokens	Greedy; thinking disabled

Appendix

Table IV : Evaluation datasets, scoring metrics, and inference protocols.

Figure 13: Qwen3.6-35B-A3B scores and decode key shares across routing thresholds for deep (C4 L6, blue) and wide (C256 L1, green) ZHA, with

R=16

. Darker shades indicate higher

p

. (a) Reasoning; (b) retrieval. Labels above bars show the benchmark scores and labels inside bars give key shares where available.

Figure 14: Qwen3.6-35B-A3B decode key shares for deep (C4 L6, blue) and wide (C256 L1, green) ZHA (

R=16

) at

p\in\{0.95,0.9,0.8\}

, compared with MSA (top-16 blocks of 128 tokens). (a) Reasoning; (b) retrieval. Darker shades indicate higher

p

; the dashed line marks dense attention (100%).

Figure 15: Dense attention vs. a ZHA router with one reference query per KV group, 8-token chunks, fan 8 and two levels (R1 C8 L2) on Qwen3.6-35B-A3B, with uniform routing top-

p

0.8, 0.75 and 0.7. (a) Reasoning (b) Retrieval.

Figure 16: Dense attention vs. the wide ZHA router (16 reference queries per KV group, 256-token chunks, one level; top-

p

0.95) on Qwen3.6-35B-A3B, with the router in BF16 and in NVFP4 (router operands, reference-key cache and scoring in FP4; backbone, KV cache and the sparse attention in BF 16). (a) Reasoning (b) Retrieval

Figure 17: Nemotron-3.5-Lightning-30B-A3B scores and decode key shares across routing thresholds for deep (C4 L6, blue) and wide (C256 L1, green) ZHA, with

R=16

. Darker shades indicate higher

p

. (a) Reasoning; (b) retrieval. Labels above bars show the benchmark scores and labels inside bars give key shares where available.

Figure 18: The 28 geometries used in the RULER fits. Each panel fixes the reference count

R

; dotted curves mark constant top-level comparession

C^{L}

. These geometries are evaluated at

p\in\{0.7,0.8,0.9,0.95,0.99\}

and sequence lengths of 64K, 128K, and 256K.

Figure 19: Per-task geometry, sparsity, and prediction accuracy for Qwen3.6-35B-A3B with reference-query-only routing; the final group shows the direct aggregate fit. Panels (a), (b), and (d) fix

R=16

. (a) Measured accuracy at 256K versus decode keys attended across 11 geometries, colored by

C^{L}

. (b) The matched

C^{L}=4096

geometries, colored by depth

L

. (c) The

C=8

L=2

slice at 256K, colored by reference count

R

. Curves sweep

p\in\{0.7,0.8,0.9,0.95\}

, with

p=0.99

included where available; markers denote depth in (a,b) and

R

in (c). Key shares use each task’s native decode counters (pooled counters for the aggregate), and the key-share axes are logarithmic. Dotted horizontal lines show dense attention. (d) All 150 leave-one-length-out predictions per response over 64K, 128K, and 256K, colored and marked by held-out length; inset values give RMSE in percentage points and dashed diagonals show equality. Each task has four independent parameters in equation 2 , fitted at the other two lengths across

R\in\{1,2,4,16\}

using the measured dense baseline at the held-out length. Accuracy uses 100 examples per task; vertical-axis ranges vary by task.

Figure 20: Measured RULER-13 accuracy at 256K versus (a) measured decode key share and (b) modeled decode-attention speedup over dense for Qwen3.6-35B-A3B (115 settings, 28 geometries). Speedups are H100 BF16 attention-layer projections at batch size one, not measured end-to-end timings; missing higher-level union counts use the slower modeled scenario. Colors denote

R

, markers

C

, and curves sweep

p

within each geometry. Black lines mark Pareto frontiers; dotted horizontal lines show dense accuracy, and the vertical guide marks

1\times

Variable	Definition
$B$	Batch size.
$T,\ Q(T)$	Context and query lengths; $Q(T)=1$ in decode and $Q(T)=T$ in full prefill.
$H_{q},\ H_{\mathrm{KV}},\ G$	Query and KV head counts, and GQA group size $G=H_{q}/H_{\mathrm{KV}}$ .
$d_{k},\ d_{v}$	Key and value dimensions per head.
$s_{\mathrm{QKV}},\ s_{\mathrm{accum}}$	Bytes per stored Q/K/V or output element, and per accumulator or softmax statistic.
$x$	Kernel stage: attention, split-KV partials, or merge.

Appendix

Table V : Variables for the dense-attention roofline model.

Mode	Storage	Median $e$	Median $\|e\|$	$90$ th percentile $\|e\|$
Causal prefill	BF16	$-0.9\%$	$15.2\%$	$57.2\%$
Causal prefill	FP16	$-1.0\%$	$16.5\%$	$54.9\%$
Decode	BF16	$12.9\%$	$12.9\%$	$54.1\%$
Decode	FP16	$13.2\%$	$13.3\%$	$53.8\%$

Appendix

Table 26

Panel of Figure 21 — Figure 21: Held-out dense-attention latency and timing divided by prediction. Within each subfigure, rows are MHA, GQA4, GQA8, and MQA; columns are causal prefill and decode. Colours distinguish $d\in\{64,128,256\}$ . Dots show all recorded timing samples and short ticks mark medians; solid lines in the latency views show worker-aware predictions. Each CUDA-event sample averages a replay batch, not an individual invocation. The lower views show $t_{\mathrm{Flash}}/t_{\mathrm{model}}$ ; the dashed $1\times$ line denotes exact prediction, and values below it indicate overprediction. Matching views retain the same layout, colours, marks, and axis limits across precisions; the ratio views share the same ratio scale.

Variable	Definition
$B$	Batch size.
$T,\ Q(T)$	Context and query lengths; $Q(T)=1$ in decode and $Q(T)=T$ in full prefill.
$H_{q},\ H_{\mathrm{KV}}$	Numbers of query heads and key/value heads, respectively.
$d_{k},\ d_{v}$	Query/key and value dimensions per head, respectively.
$G$	Query heads per KV head (GQA group size): $G=H_{q}/H_{\mathrm{KV}}$ .
$s_{\mathrm{QKV}}$	Bytes per stored Q/K/V, compressed-key, reference-bias, or routing-score element; BF16 baseline: $s_{\mathrm{QKV}}=2$ .

Appendix

Table VI : Variables for the ZHA roofline model.

Variable	Definition
$x$	Kernel label: routing, scan, attention, split partial/merge, fusion, or tree summarization/merge. Sparse stage quantities also carry $\ell$ when level-specific.
$\bm{\Phi}_{x}(T)$	Sparse-stage execution latency, excluding kernel launches.
$\bm{t}_{x,\mathrm{worker\ waves}}(T)$	Worker-wave contribution to sparse-stage execution latency.
$\mathbf{FLOPs}_{x}(T),\ \mathbf{bytes}_{x}(T)$	Sparse-stage operation count and memory traffic (bytes).
$\bm{N}_{\mathrm{tiles},x}(T)$	Product of launch-grid dimensions, not inner-loop tiles. In 1-D: $\lceil\bm{n}_{x}(T)/\kappa_{x}\rceil$ . Unsplit attention is a deterministic special case.
$\bm{n}_{x}(T)$	Work-item count for a generic 1-D partition.

Appendix

Table VII : Kernel and worker variables for the roofline model.

Quantity	A100 SXM4	H100 SXM5	MI300X	MI300A
Workers $N_{\mathrm{SM}}$	108 SMs	132 SMs	304 CUs	228 CUs
Hardware bound on $c_{x}$	$\leq 32$	$\leq 32$	$\leq 32$	$\leq 32$

Appendix

Table 30

Variable	Definition
$n,\ n_{\ell}(T)$	Full nodes per sequence; the balanced full-node count at level $\ell$ is $n_{\ell}(T)=\lfloor T/C^{\ell}\rfloor$ .
$N_{\mathrm{tiles},x}(n)$	Deterministic launch-grid size for a tree kernel processing $n$ full nodes per sequence.
$\mathrm{FLOPs}_{x}(n),\ \mathrm{bytes}_{x}(n)$	Deterministic operation count and memory traffic for that tree kernel.
$t_{\mathrm{summ}}(n),\ t_{\mathrm{merge}}(n)$	Leaf-summarization and parent-merge latencies, respectively, including launches.
$t_{\mathrm{tree}}^{\mathrm{prefill}}(T)$	Cost of constructing all full nodes in prefill.
$\overline{t}_{\mathrm{tree}}^{\mathrm{decode}}(T)$	Tree-maintenance cost per decoded token, amortized over full-node completions.

Appendix

Table VIII : Variables for tree construction and maintenance.

Variable	Definition
$C,\ k$	Tokens per block and maximum selected blocks per KV group, including the local block.
$M(T),\ k_{T}$	Visible block count and selected block count at context length $T$ .
$r(T),\ K(T)$	Tokens in the final, local block and total retained tokens per KV group.
$d_{I},\ d_{\mathrm{model}}$	Index-head width and model hidden width.
$W_{I}$	Combined index-projection width: $W_{I}=(H_{\mathrm{KV}}+1)d_{I}$ .
$\kappa_{\mathrm{proj}}$	Index-projection output channels per tile.

Appendix

Table IX : Additional variables for the MSA roofline model.

Stage $x$	$N_{x}$	$F_{x}$	$V_{x}$
Index projections	$B\lceil W_{I}/\kappa_{\mathrm{proj}}\rceil$	$2Bd_{\mathrm{model}}W_{I}$	$s[d_{\mathrm{model}}W_{I}+B(d_{\mathrm{model}}+W_{I})]$
Index + max-pool	$\nu_{I}\lceil T/\kappa_{I}\rceil$	$2BH_{\mathrm{KV}}d_{I}T$	$sd_{I}(\nu_{I}T+BH_{\mathrm{KV}})+s_{c}BH_{\mathrm{KV}}M$
Top- $k$	$BH_{\mathrm{KV}}$	—	$BH_{\mathrm{KV}}(s_{c}M+s_{\mathrm{id}}k_{T})$

Appendix

Table 33

Figure 22: Per-level selection scaling for

C=4

R=16

, and

p=0.7,0.9,0.95

, extending Figure 7 . The first five columns show depths 2–6; the sixth compares observed median exponents with the

r=1/2

profile. Bars show pointwise 95% whole-problem bootstrap intervals.

Figure 23: Observed top-level exponents versus total tree depth for

C=4

R=16

. Dotted lines fit constant (gray) and linear (blue) depth dependence separately at each

p

, using only the five top-level observations. Bars show pointwise 95% whole-problem bootstrap intervals.

Figure 24: Modeled BF16 decode latency for dense attention, MSA, and ZHA ( $R=1$ , $C=8$ ) at $p=0.7,0.75,0.8$ , with the H100 assumptions and nondecreasing latency-minimizing depth of Figure 8 . Latencies are modeled per-token decode times for the attention subsystem of one layer at batch size 1.

Figure 25: Retrieval traversals at 64K tokens for (a) single-needle and (b) multi-key 2 RULER tasks (Qwen3.6-35B-A3B, C16 L3, $R=16$ , $p=0.9$ ). Both show layer 35, query head 8, while emitting the first answer digit. Dark blue blocks are scored and expanded nodes; light blue blocks scored but pruned nodes; grey blocks are nodes which were never scored by the router. Dense-attention rows show chunk-level mass on a logarithmic scale. The green marker indicates the ground-truth needle location. Boxed regions are enlarged below to reveal finer traversal details. These examples illustrate task-dependent routing with top- $p$ : for single-needle, the router explores few branches before reaching the answer-containing chunk; for multi-key 2, whose context contains many distractor needles, it explores substantially more upper-level nodes. Both traversals ultimately attend to a single answer-containing chunk.

Figure 26: Agreement between tree levels in reference-query selection for Qwen3.6-35B-A3B ZHA ( $R=16$ ). Columns show C4 L6, C8 L4, C16 L3 and C64 L2; rows show RULER at 64K tokens and natural text at 32K tokens. For each layer/KV group, Spearman correlation is computed between the 16 reference-query selection shares at each pair of levels, then averaged over 20 groups (10 layers $\times$ 2 KV groups). Parenthesized axis labels give node spans in keys. Adjacent levels are similar in their reference selection rate while further apart levels have a different selection priority. This behavior holds across all geometries.

Figure 27: Typical selection behavior across levels: layer 19, KV group 0 of the C4 L6 ZHA router on RULER at 64K tokens. Adjacent levels have similar reference rankings, but agreement weakens over larger level separations; the chunk and top-level rankings are nearly unrelated (L1 versus L6: $\rho=0.12$ ).

Figure 28: Outlier selection across levels: layer 39, KV group 1 of the C4 L6 ZHA router on RULER at 64K tokens. Upper levels agree closely, while L1 reference rankings are negatively correlated with the top levels (L1 versus L6: $\rho=-0.35$ ).

Figure 29: Outlier selection across levels: Similar reference rankings across all levels: layer 3, KV group 1 of the C4 L6 ZHA router on RULER at 64K tokens. Chunk and top-level rankings remain strongly aligned ( $\rho=0.86$ ).

Figure 30: Reference-query cosine similarities across the four Qwen3.6-35B-A3B ZHA geometries. The C16 L3 references are ordered by average-linkage clustering with cosine distance; references in each other geometry are Hungarian-matched to them by cosine similarity. Matched references therefore occupy corresponding positions. Diagonal structure within the off-diagonal blocks indicates similar reference directions across geometries.

Figure 31: Reference-query cosine similarities with references independently ordered by top-level selection frequency within each geometry and KV group, most to least frequent. Frequencies were measured during routed prefill on 16 RULER prompts at 64K tokens and $p=0.9$ . No cross-geometry matching is applied; diagonal structure therefore indicates correspondence in both direction and selection rank.

Figure 32: Mean cosine similarities of centered reference keys across geometries, evaluated on the same 60 top-level nodes from four RULER prompts at 64K tokens. Each node covers 4096 keys. Within each node and KV group, the mean over the 16 reference keys is subtracted before computing cosine similarities, which are then averaged over nodes. C16 L3 references are clustered using reference-key cosine similarity; other geometries are Hungarian-matched using the same measure. The cross-geometry correspondence extends to the pooled key representations.

Efficient Attention

Efficient attention reduces the compute, memory traffic, and cache costs of Transformer attention while preserving useful long-range interactions.

See all 142 papers

Hierarchical Sparse Attention Done Right: Toward Infinite Context Modeling
Jul 3, 2026Xiang Hu, Xinyu Wei, Hao Gu +10Long-Context Language Model Inference Long-Context Language Modeling
Tencent HY Team · ShanghaiTech University · The Hong Kong University of Science and Technology +1
Latent Sparse Attention: Towards Subquadratic Ultra-Long Context Attention
Oct 30, 2026Ales Kubicek, Andrea Caciolai, Liangzu Peng +7Sparse Attention Long-Context Language Modeling
Prefrontal AI, London, United Kingdom · ETH Zurich, Zurich, Switzerland
HyScale: Scaling Long-Context Training for Learned Sparse Attention
Oct 23, 2026Yangshen Deng, Xinjing Huang, Yulong Wang +8Efficient Language Model Training Sparse Attention
University of Edinburgh · Tencent Hunyuan

Jul 3, 2026J/K move · Enter open · S save

Hierarchical Sparse Attention Done Right: Toward Infinite Context Modeling

Xiang Hu, Xinyu Wei, Hao Gu, Minshen Zhang+9

Tencent HY Team · ShanghaiTech University · The Hong Kong University of Science and Technology+1

Long-Context Language Model Inference Long-Context Language Modeling Long-Context Modeling Sparse Attention Hierarchical Attention

Open paper AskPDF arXiv

Scaling modern large language models (LLMs) to long contexts is limited by the quadratic computation cost, and poor length extrapolation of dense attention. Chunk-wise sparse attention offers a promising alternative, but all existing methods fall short of full attention because of their inaccurate chunk selection. We propose Hierarchical Landmark Sparse (HiLS) Attention, a chunk-wise sparse attention mechanism that learns chunk selection end-to-end under the language-modeling (LM) loss. HiLS factorizes attention hierarchically: each query performs attention independently with each retrieved chunk to extract chunk-specific information, and the resulting outputs are fused according to chunk retrieval scores. By incorporating retrieval scores into the forward attention computation, HiLS optimizes them directly with the LM loss, enabling end-to-end retrieval learning and native sparse training. Experimental results show that HiLS-Attention achieves performance comparable to, and in some cases better than, full attention at in-domain context lengths. Meanwhile, HiLS-Attention extrapolates more than

64\times

the training context length with 90% retrieval accuracy, far beyond full attention. Moreover, existing full-attention models can be converted to HiLS-Attention with lightweight continued pretraining, preserving in-domain performance while acquiring ultra-long-context extrapolation. Together with its sparse KV access and computation, HiLS-Attention breaks the usual efficiency-performance trade-off, enabling long-context LLMs that are both more efficient and more effective on general long-context tasks than their full-attention counterparts.

Abstract

Efficient Attention

Explore similar work

Hierarchical Sparse Attention Done Right: Toward Infinite Context Modeling

Latent Sparse Attention: Towards Subquadratic Ultra-Long Context Attention

HyScale: Scaling Long-Context Training for Learned Sparse Attention