Hierarchical Sparse Attention Done Right: Toward Infinite Context Modeling

Xiang Hu, Xinyu Wei, Hao Gu, Minshen Zhang, Tian Liang, Huayang Li, Lei Zhu, Yan Wang, Sirui Han, Yushi Bai, Kewei Tu, Haitao Mi, Leo Liang

HiLS-Attention enables end-to-end learnable sparse retrieval, achieving ultra-long context extrapolation.

How can we replace non-parametric chunk selection in sparse attention with a differentiable, learnable mechanism to enable efficient, long-context LLM scaling?

Scaling Large Language Models (LLMs) to long contexts is bottlenecked by the quadratic cost of full attention and the inaccurate chunk selection of existing sparse methods. HiLS-Attention replaces non-parametric chunk summaries with learnable landmark-based surrogates, allowing the model to optimize chunk selection end-to-end via the language-modeling loss. This approach achieves performance comparable to full attention on in-domain tasks while enabling 64× to 512× length extrapolation with high retrieval accuracy.

Paper Primer

Existing sparse attention methods rely on mean-pooled chunk summaries that fail to capture the concentrated attention mass required for precise retrieval. HiLS-Attention (Hierarchical Landmark Sparse Attention) solves this by using a first-order Taylor expansion of the full-attention chunk mass to derive a learnable, entropy-calibrated surrogate score.

The mechanism functions like a hierarchical router: it first computes a surrogate score for each chunk using a landmark token, selects the top-K chunks, and then fuses the intra-chunk token-level information with the inter-chunk retrieval scores. This allows the model to suppress irrelevant chunks and promote those useful for next-token prediction during the forward pass.

HiLS-Attention achieves superior long-context extrapolation compared to full attention.

In needle-in-a-haystack retrieval tasks, the model maintains over 90% accuracy at 4M context length despite being trained on only 8K context.

HiLS-Attention preserves short-context performance while reducing computational overhead.

Full-attention models converted to HiLS-Attention via lightweight continued pretraining match or exceed baseline performance on LongBench and general downstream tasks.

Why is this approach more effective than simply using mean-pooled chunk summaries?

Mean-pooled summaries are only accurate when attention logits are uniform; they fail when specific tokens dominate, which is the typical case for precise in-context retrieval. HiLS-Attention’s surrogate incorporates an entropy-based bias term that adaptively accounts for these concentrated distributions.

Does this method require training a model from scratch?

No. The authors demonstrate that existing full-attention models can be converted to HiLS-Attention using lightweight continued pretraining, where only the landmark token embeddings and low-rank query adapters are updated.

HiLS-Attention breaks the efficiency-performance trade-off by enabling native sparse training that is both more efficient during inference and more effective at long-range retrieval than full-attention baselines.

HiLS-Attention provides a robust, end-to-end learnable framework for sparse attention, making ultra-long or "infinite" context modeling a practical reality for parameter-constrained LLMs.

Introduction and Motivation

We expose why heuristic sparse attention fails and introduce HiLS‑Attention as a learnable, high‑fidelity sparse alternative.

Standard sparse attention methods select chunks with non‑differentiable heuristics, which leads to poor information fidelity and hurts long‑context performance. The naive Block Sparse Attention (BSA) computes full attention, aggregates token‑level attention mass into each chunk (the “Chunk Mass” $Z_c$), and then picks the top‑K chunks—an accurate but computationally prohibitive pattern. Our central motivation is that this heuristic selection is the bottleneck preventing sparse models from matching full‑attention quality.

**Figure 1.** After only 50B continued-training tokens, HiLS-Attention inherits the capability of full attention while bringing two key advantages: *strong ultra-long context extrapolation* beyond the YaRN-extended 4x length (Fig. 1a) and *faster inference* (Fig. 1b). Meanwhile, it preserves comparable performance for short- and medium-context tasks, within both the original training length and the YaRN-extrapolated range (Fig. 1c & 1d).

HiLS‑Attention enables long‑context scaling with sparse compute.

Limitations of Naive Sparse Attention

Naive block‑sparse attention is exact but computationally costly, and common approximations break fidelity.

Naive Block Sparse Attention (BSA) can select exactly the most informative chunks, but it must evaluate every token‑to‑token logit inside each candidate chunk, which erodes the computational savings that sparsity promises.

BSA partitions the sequence into equal‑sized chunks, lets each query attend to a local sliding window and to the top‑$K$ distant chunks chosen by their total attention mass.

Existing sparse‑attention schemes replace the exact chunk mass with cheap proxies such as mean‑pooled keys, but those proxies are accurate only under restrictive logit distributions, leading to systematic ranking errors.

Mean‑pooled keys capture the average logit in a chunk, which is faithful when all logits are similar, but they collapse to the maximum when a single token dominates, missing the true LogSumExp behavior.

The HiLS-Attention Mechanism

We introduce HiLS-Attention, a differentiable sparse attention that learns chunk summaries.

The analysis of naive block sparse attention revealed that heuristic chunk selection discards useful information. HiLS‑Attention replaces the non‑parametric chunk mass $Z_{i,c}$ with a learnable surrogate $\hat Z_{i,c}$ to enable fully differentiable sparse routing.

HiLS‑Attention learns a compact summary for each chunk, letting the model assign attention mass via a differentiable surrogate instead of a hard heuristic.

Chunk $c_1$ contains tokens $t_1,t_2$. Using $q'_c$, we obtain scores $s_1=0.2$, $s_2=0.8$, yielding $p_1=0.27$, $p_2=0.73$.

Weighted summary $k'_{c_1}=0.27\,k_1+0.73\,k_2$ and bias $b'_{c_1}=-(0.27\log0.27+0.73\log0.73)=0.61$.

Surrogate score $\hat s_{i,c_1}=q_i^{\top}k'_{c_1}/\sqrt{d}+0.61$ is computed for each query $q_i$; similarly for $c_2,c_3,c_4$.

Inter‑chunk softmax ranks the four $\hat s_{i,c}$ values; the top‑$K=2$ chunks (e.g., $c_2$ and $c_4$) are kept.

Within each selected chunk, an intra‑chunk softmax distributes the mass among its two tokens, completing the attention output for token $i$.

The surrogate captures both the average relevance of a chunk (via $k'_c$) and the distribution’s certainty (via $b'_c$), enabling the model to learn which chunks are truly informative.

How does HiLS‑Attention differ from naive Block Sparse Attention that selects chunks by exact mass $Z_{i,c}$?

Naive BSA computes the exact chunk mass $Z_{i,c}$ with a full $QK$ pass, which is non‑differentiable and costly. HiLS‑Attention replaces $Z_{i,c}$ by a learned surrogate $(k'_c,b'_c)$ that is cheap to compute, differentiable, and can be optimized jointly with the language‑modeling loss.

HoPE augments each token with positional embeddings at multiple granularities, allowing the model to reason about both fine‑grained token order and coarse‑grained chunk locations.

Token $t_3$’s content embedding $e_3$ is summed with $p_3=23$ to form $e'_3$.

Query $q_3$ is computed from $e'_3$, influencing both the intra‑chunk attention scores and the landmark query $q'_{c_2}$.

The chunk‑level embedding $p^{(1)}_{c_2}=20$ is shared by $t_3$ and $t_4$, biasing the surrogate $k'_{c_2}$ toward the chunk’s overall position.

Hierarchical positions let the model encode coarse location information without inflating the per‑token embedding size, improving long‑range routing.

Why not use a single flat positional embedding instead of the hierarchical scheme?

A flat embedding would require each token to carry the full absolute position, which grows linearly with context length and offers no explicit chunk‑level bias. HoPE’s hierarchy provides a shared chunk signal that scales gracefully and directly informs the landmark query used in HiLS‑Attention.

**Figure 3.** An overview of HiLS-Attention. We omit the scaling factor $\frac{1}{\sqrt{d}}$ for simplicity. Naive block sparse attention selects the top-$K$ chunks by their exact mass $Z_c$, e.g., chunks 1 and 3 when $K = 2$, but computing all $Z_c$ requires a full QK computation. HiLS-Attention instead uses compressed chunk keys $k'_c$ to efficiently estimate a chunk-mass surrogate $Z'_c \propto \exp(q^\top k'_c)$. It factorizes attention into two stages: an **inter-chunk softmax**, which specifies the total attention mass assigned to each chunk, and an **intra-chunk softmax**, which distributes each chunk's attention mass among its tokens. Since $Z'_c$ parameterizes the forward attention weights, gradients from the next-token prediction loss can be directly backpropagated to the compressed key $k'_c$, enabling end-to-end learning.

HiLS‑Attention routing algorithm (per query token $i$)

By learning chunk summaries and hierarchical positions, HiLS‑Attention achieves differentiable sparse routing while preserving information fidelity. The next section details how we implement and train the model at scale.

Implementation and Training Recipes

Practical design choices that make HiLS‑Attention fast and trainable.

Implementing HiLS‑Attention requires a handful of architectural and kernel‑level tweaks. Below we describe the choices that let the method stay differentiable, memory‑efficient, and GPU‑friendly.

Instead of rotating every dimension, HoPE keeps only those rotation periods that fit the training context length and replaces the rest with a non‑rotating embedding (NoPE).

How does HoPE differ from the standard RoPE that the reader might assume?

RoPE applies a sinusoidal rotation to every model dimension, tying the rotation period to the absolute position. HoPE selectively disables rotation for dimensions whose period would exceed the maximum training length, so those dimensions act like a plain embedding (NoPE) instead of a rotating one.

Because a chunk key $k'_c$ summarizes many tokens, the original token query $q_i$ is a poor estimator of the chunk’s relevance; Q‑Cal adds a cheap low‑rank correction to the query before scoring.

Compute $\Delta q_i = W_{\text{up}}W_{\text{down}}h_i = W_{\text{up}}(1,0)^{\top} = (0.1,0.2,0.1,0.2)$.

Form calibrated query $\hat{q}_i = q_i + \Delta q_i = (0.6, -0.3, 0.6, -0.3)$.

Dot‑product $\hat{q}_i^{\top}k'_c = 0.6\cdot0.2 + (-0.3)\cdot0.2 + 0.6\cdot0.2 + (-0.3)\cdot0.2 = 0.12$.

Scale by $\sqrt{d}=2$: $\frac{0.12}{2}=0.06$.

Add bias $b'_c=0$ (for simplicity) → final score $\hat{s}_{i,c}=0.06$.

Even a tiny low‑rank correction can shift the query enough to change the relative ordering of chunk scores, which explains the observed perplexity gains.

Why not use a full‑rank linear layer instead of the low‑rank $W_{\text{down}}$/$W_{\text{up}}$ pair?

A full‑rank layer would add $d_{\text{model}}\times d$ parameters per head, dramatically increasing memory and compute. The low‑rank factorization keeps the parameter count to $d_{\text{model}}\times r + r\times d$, where $r\ll d_{\text{model}}$, preserving efficiency while still providing a useful query adjustment.

When several query heads share a KV head, we let each head compute its own chunk scores, then take the maximum across heads so that any head that finds a chunk important forces its inclusion.

How does taking the maximum differ from averaging the head scores?

Maximum ensures that if any head deems a chunk valuable, the chunk is kept, preserving head‑level flexibility. Averaging could dilute a strong signal from a single head, causing the chunk to be dropped even though it is crucial for that head.

To avoid the engineering overhead of per‑chunk landmark tokens, we instead use a single shared learnable query per layer; this matches in‑domain performance but harms extrapolation.

**Figure 4.** Kernel design of NSA and HiLS-Attention. Kernel design of NSA and HiLS-Attention. (a) NSA handles one query token per tile and computes attention over its selected chunks. Each Tensor Core operation has shape $(G, d) \times (d, S)$, where $G$ is the GQA group size and $S$ is the chunk size. (b) HiLS-Attention packs $M$ adjacent query tokens, attends to the union of their selected chunks, and enlarges the Tensor Core operation to $(M \times G, d) \times (d, S)$. This packing reuses overlapping K/V chunks across adjacent tokens.

Sparse attention kernels suffer from under‑utilized Tensor Cores because each token often attends to a different chunk set. HiLS‑Attention’s packing strategy aligns multiple queries so that the same K/V chunks are reused, enabling larger matrix tiles.

We batch $M$ adjacent query tokens and their $G$ heads, then attend to the union of their selected chunks, turning many small matrix multiplies into a single larger one that fits Tensor Core tile requirements.

Stack the $M\times G = 8$ query vectors into a matrix $Q_{\text{pack}}$ of shape $(8,8)$.

Load the four selected chunk key/value matrices $K_{\text{union}}$, $V_{\text{union}}$ of shape $(8,4)$ each.

Perform GEMM $A = Q_{\text{pack}}\,K_{\text{union}}^{\top}$ → shape $(8,4)$.

Apply softmax row‑wise on $A$ to obtain attention weights.

Compute output $O = \text{softmax}(A)\,V_{\text{union}}$ → shape $(8,8)$.

Unpack $O$ into two separate query outputs by slicing the first 4 rows for query 1 and the next 4 rows for query 2.

Because the two queries share three of four chunks, the packed GEMM reuses the same $K$/$V$ data, cutting memory traffic by roughly 50 % while still delivering the exact same per‑query results.

Does packing change the attention distribution compared to processing queries separately?

No. Each query’s attention weights are computed using the same softmax over the union of chunks, but because the union contains all chunks any individual query would have seen, the resulting per‑query distribution is identical to the separate‑query case.

This packed kernel design is the primary reason HiLS‑Attention achieves real‑world speedups: it aligns the computation with Tensor Core tile sizes, eliminates redundant K/V loads, and removes the hard $G\ge16$ requirement of prior NSA kernels.

We evaluate two continued‑training (CPT) recipes to integrate HiLS‑Attention without losing the base model’s capabilities.

Freeze all original parameters and train only the newly introduced landmark embeddings and the Q‑Cal projection matrices, which together constitute < 1 % of the model.

Why does training on just 5 B tokens recover the base model’s quality?

The trainable parameters are a tiny fraction of the whole model, so a modest amount of data is enough to adjust the new query‑calibration and landmark embeddings to the distribution the base model already knows.

Re‑initialize landmark embeddings and Q‑Cal adapters, then jointly fine‑tune every parameter, which is especially beneficial after swapping RoPE for HoPE.

By packing adjacent queries, using a low‑rank query adapter, and adopting HoPE, HiLS‑Attention becomes both GPU‑efficient and easy to fine‑tune.

Small-Scale Empirical Evaluation

Small models show HiLS‑Attn’s strong perplexity and retrieval gains.

Standard sparse attention relies on non‑differentiable chunk selection, which harms fidelity. HiLS‑Attn replaces those heuristics with learnable summaries, enabling fully differentiable, high‑fidelity sparse attention.

HiLS‑Attn‑HoPE attains the lowest perplexity across all context lengths, reaching $5.01$ at $512\text{K}$ tokens.

Table 1 reports a PPL of $5.01$ for HiLS‑Attn‑HoPE, beating the next best $8.47$ (Full‑Attn RoPE).

A synthetic long‑context task that hides short “needle” strings in a long document and asks the model to retrieve them at the end, measuring pure in‑context retrieval fidelity.

**Figure 2.** In-context retrieval results.

Scaling to Ultra-Long Contexts

HiLS‑Attention scales to 256 K context with lower perplexity and faster inference.

HiLS‑Attention achieves up to 15.7× faster decode latency than full attention at a 512 K context length.

Figure 6 shows a 15.7× decode speedup for HiLS‑Attention at 512 K.

**Table 3.** Perplexity for 345M models continue trained with a 256K context length on 10B

**Figure 5.** Perplexity (a) and RULER accuracy (b) of the 1.4B model at different training steps. Left: Full-Attention with RoPE; right: HiLS-Attention with HoPE. The annotated values on the curves correspond to the final checkpoint (143k steps), which is highlighted with star markers and thicker lines. The detailed per-step results are deferred to Appendix H (Tab. 15 &amp; Tab. 16).

**Figure 6.** Inference latency of HiLS-Attention vs. full attention at the 345M scale on a single NVIDIA H800 (batch size 1, bf16, chunk size 64, top-k = 32, 512-token sliding window). Prefill is the warm-run latency for the full input; decode is the median per-token latency with CUDA graphs. Speedup is full attention / HiLS-Attention (> 1 means HiLS-Attention is faster). HiLS-Attention reaches parity at ~16K and is 13.5x/15.7x faster (prefill/decode) at 512K.

Ablation Study Overview

Isolating each HiLS‑Attention component reveals its impact on long‑context performance.

We evaluate the ablation configurations listed in Table 5, comparing each against the default HiLS‑Attention setting.

Low‑rank query calibration (r = 64) adds less than 1 % parameter overhead across all context lengths.

Table 6 shows overheads of 0.6 % (64), 1.2 % (128), 4.9 % (512) and 0 % for longer contexts.

Removing the $\Delta$q term (w/o Q‑Cal) causes a dramatic increase in perplexity, indicating that the low‑rank path is essential for decoupling token‑level attention from chunk‑level mass.

Increasing the rank to r = 128 (Table 7) harms length extrapolation, with perplexity rising sharply for contexts beyond 32 K, suggesting that over‑parameterizing the query projection destabilizes scaling.

Landmark tokens are indispensable: the landmark‑free variant (w/o lmk, shared qc) quickly loses performance on ultra‑long sequences, confirming that learned chunk queries $q'_c$ require dedicated landmark embeddings.

**Figure 7.** Chunk-id overlap among adjacent query tokens. Left: loaded union size for the final $M = 16$ queries versus the visible historical chunks; percentages denote loaded fractions. Right: normalized overlap as group size $M$ increases, with the dashed line showing inter-block reuse at $M = 16$. Error bars show standard deviation across HiLS layers and heads, and shaded regions show the layer-wise min-max range.

Both HoPE positional encoding and query augmentation ($\Delta$q) consistently improve scores across the RULER benchmark (Table 5), indicating that richer positional signals and calibrated queries jointly boost fidelity.

Ablation Results: Performance Metrics

Key ablations reveal which HiLS‑Attention components most affect RULER performance.

This ablation suite isolates each design choice in HiLS‑Attention and measures its impact on the RULER benchmark (345 M model, 8 K context).

Removing Prop. 3.1 from HiLS‑Attn‑HoPE drops the RULER score from 100 % to 96 %.

Table 7 reports 100 % with the full component and 96 % when Prop. 3.1 is omitted.

Overall, the ablations demonstrate three clear trends: (1) Prop. 3.1 is indispensable for HiLS‑Attention; (2) HoPE’s hybrid positional encoding consistently outperforms pure RoPE or NoPE when combined with chunk compression; and (3) landmark‑token queries, generated by the full Transformer stack, are crucial for robust long‑context extrapolation.

Ablation Results: LongBench Evaluation

Native sparse training ablations reveal which components drive performance gains.

We now ask the canonical ablation question: does removing a HiLS‑Attention component hurt performance?

LongBench‑v1 aggregates a suite of long‑context language tasks, measuring how well a model retains information as context length grows.

**Table 8.** Downstream task evaluation results for 1.4B model after 300B-token training.

Removing the HoPE embedding (RoPE variant) drops LAMBADA by 1.5 points.

Table 8 reports 57.13 for HoPE vs 55.64 for RoPE.

Switching from HoPE to NoPE collapses performance on HellaSwag by ≈ 32 points.

HellaSwag scores: 57.05 (HoPE) versus 25.00 (NoPE) in Table 8.

**Table 1.** Comparison of model performance across various benchmarks.

**Table 10.** Validation perplexity on the Olmo3 pretraining corpus at different context lengths. Lower is better. Models are continued-pretrained at 8K length. Perplexity values are reported to three decimal places to better highlight small differences across models.

**Table 11.** LongBench-v1 scores grouped by context length. Overall is weighted by the number of samples in each dataset.

HiLS‑Attn‑HoPE’s LongBench‑v1 score at 128 K is 5.8 points lower than RoPE, indicating superior long‑range fidelity.

Table 11: HoPE = 5.8 vs RoPE = 4.1 (lower is better).

Related Work

We position HiLS‑Attention among prior sparse‑attention approaches, highlighting key design differences.

Sparse attention methods fall into two families: token‑wise approaches that attend individually to each token, and block‑wise approaches that operate on groups of tokens. This work concentrates on block‑wise sparse attention because it can enable native sparse training, a prerequisite for truly infinite‑length contexts.

Block‑wise methods differ in how they summarize each chunk. Non‑parametric summaries use fixed, heuristic statistics such as mean‑pooled representations; examples include NSA, MoBA, DashAttention, and SeerAttention, all of which struggle to retrieve the exact needle even at training length, indicating poor chunk relevance estimation.

Parametric summaries learn a representation for each chunk. Landmark Attention (LMK‑Attn) introduces landmark tokens but still requires dense attention during training, while the HSA series designs specialized kernels and combines NoPE‑based HSA with RoPE‑based SWA, incurring substantial parameter overhead and later perplexity dominated by SWA.

Compared with LMK‑Attn, HiLS‑Attention provides native sparse training and a more accurate chunk‑mass estimate, yielding markedly lower perplexity and better extrapolation. Inspired by HSA’s NoPE component, we adopt HoPE to retain positional information while supporting ultra‑long contexts.

We hypothesize that HiLS‑Attention’s strong length extrapolation stems from higher‑quality in‑context retrieval, enabled by its compression‑and‑retrieval bias: token‑level noise is attenuated while shared semantics survive. The algorithm compresses each chunk to a lightweight summary, reducing per‑token retrieval cost to $O(L/S)$ and overall attention cost to $O(LKS)$, with retrieval further reducible to $O(L\log(L/S))$ via FAISS.

The table presents performance metrics for various methods across two categories of LongBench-v1: those with context lengths less than 8K and those greater than 8K, along with an overall percentage score.

**Table 12.** Comparison of sparse-attention variants in terms of in-context retrieval capability, require full QK computation during training, extrapolation ability, training strategy, and compatibility with CPT. The weaknesses of each method are highlighted in red.

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers