Source-linked AI summary
Double-P: Hierarchical Top-P Sparse Attention for Long-Context LLMs
Wentao Ni, Kangqi Zhang, Zhongming Yu, Oren Nelson, Mingu Lee, Hong Cai, Fatih Porikli, Jongryool Kim, Zhijian Liu, Jishen Zhao
TL;DR
Fixed-budget top-k attention cannot adapt to heterogeneous attention distributions, while existing top-p methods do not jointly optimize estimation, selection overhead, and sparse-attention cost. Double-P uses cluster-level estimation followed by adaptive token-level refinement. It achieves near-zero accuracy gap to full attention and up to 1.26× end-to-end decoding speedup over existing methods.
Problem
Method
Double-P estimates attention mass at the cluster level and adaptively refines token-level computation through a second top-p stage.
Source: §4 · p. 4, §4.2 · p. 5, §4.3 · p. 6Results
Takeaways & Limitations
Double-P combines probability-based pruning with hierarchical sparse attention while preserving top-p attention guarantees.
Source: §1 · p. 2, §6 · p. 8Token-level top-p estimation uses fixed budgets that may fail to recover target attention mass and incurs selection overhead that scales linearly with context length.
Source: §1 · p. 2
Abstract
from arXiv · showhide
As long-context inference becomes central to large language models (LLMs), attention over growing key-value caches emerges as a dominant decoding bottleneck, motivating sparse attention for scalable inference. Fixed-budget top-k sparse attention cannot adapt to heterogeneous attention distributions across heads and layers, whereas top-p sparse attention directly preserves attention mass and provides stronger accuracy guarantees. Existing top-p methods, however, fail to jointly optimize top-p accuracy, selection overhead, and sparse attention cost, which limits their overall efficiency. We present Double-P, a hierarchical sparse attention framework that optimizes all three stages. Double-P first performs coarse-grained top-p estimation at the cluster level using size-weighted centroids, then adaptively refines computation through a second top-p stage that allocates token-level attention only when needed. Across long-context benchmarks, Double-P consistently achieves near-zero accuracy drop, reducing attention computation overhead by up to 1.8x and delivers up to 1.3x end-to-end decoding speedup over state-of-the-art fixed-budget sparse attention methods.
1. Introduction
Long-context decoding makes attention a major bottleneck, while fixed token budgets cannot accommodate varying attention distributions across heads, layers, and decoding steps. Double-P combines cluster-level top-p estimation with adaptive token-level refinement to improve accuracy–efficiency trade-offs.
1. Introduction
- Long-context inference makes attention a major decoding bottleneck because computation and memory overhead grow with the key–value cache.Sparse attention addresses this by restricting computation to a subset of tokens estimated to be most relevant to the query.Source: §1 · p. 1
- Top-p selection preserves a target attention mass rather than selecting a fixed number of tokens, but efficient deployment must jointly optimize estimation, selection overhead, and sparse attention cost.The introduction contrasts this with cluster-based approaches that use fixed budgets and thus behave like top-k methods without explicit mass control or top-p accuracy guarantees.Source: §1 · p. 2
- Fixed token budgets cannot reliably satisfy top-p constraints because attention distributions vary across heads and layers.This variability makes a budget adequate for some heads overly conservative or insufficient for others.Source: §1 · pp. 1–2
- Double-P achieves up to 2.27× self-attention speedup and 1.26× end-to-end decoding speedup, forming a superior accuracy–efficiency Pareto frontier over existing methods.The comparison is on Ruler at 32k context length.Source: Figure 1 · p. 1, §1 · p. 2
- Double-P estimates attention mass at the cluster level and then adaptively refines token-level computation through a second top-p stage.It is presented as preserving attention-mass guarantees while reducing estimation and sparse-attention costs.Source: §1 · p. 2
2. Background
During autoregressive decoding, each generated token attends over cached representations, making full attention costly as context grows. Sparse attention reduces computation by selecting tokens or representations, alongside other inference optimizations that target memory, prompt processing, or sequence complexity.
2.1. LLM and Attention Mechanism
- During decoding, each output token attends over previously cached keys and values, and full attention becomes a latency bottleneck as sequence length grows.The background states that computation and memory-bandwidth requirements scale quadratically with sequence length.Source: §2.1 · pp. 2–3
2.2. Sparse Attention
- Sparse attention approximates full attention by selecting a subset of tokens, with prior approaches including static selection, eviction, retrieval, prefill acceleration, trainable routing, and top-p selection.Retrieval approaches operate at page, cluster, or token granularity; their retrieval quality can be sensitive to approximation error and fixed budgets.Source: §2.2 · p. 3
2.3. Other Inference Optimizations
- Other inference optimizations allocate KV-cache budgets across layers or heads, quantize model weights or cache values, or replace softmax attention with linear-time formulations.The cited approaches aim to reduce memory, computation, or explicit KV-cache storage.Source: §2.3 · p. 3
3. Motivation
Top-p attention adapts the number of selected tokens to preserve a target attention mass, but current approaches face estimation-cost and fixed-budget limitations. These issues motivate jointly optimizing estimation, selection, and sparse computation through a hierarchical approach.
3. Motivation
- 91.9% of LongBench samples fail to reach target attention mass p = 0.95 under the commonly recommended fixed budget for token-level top-p estimation.The same budget can recover nearly all attention mass for some samples and severely underestimate it for others, leading to conservative over-provisioning.Source: §3 · pp. 3–4, Figure 3 · p. 4
- Token-level attention score estimation and top-p selection together account for a large fraction of decoding time and scale linearly with the number of tokens considered.This overhead persists even when many tokens are pruned from the final sparse-attention computation.Source: §3 · p. 4, Figure 4 · p. 4
- Cluster-based methods reduce sparse-attention cost but use fixed cluster or token budgets, leaving preserved attention mass without explicit control or top-p-aligned accuracy guarantees.This contrasts with top-p methods, which retain the minimal token set whose cumulative attention exceeds a target threshold.Source: §3 · pp. 3–4
- Existing methods either provide top-p selection at high estimation cost or gain efficiency through fixed-budget top-k behavior, motivating joint optimization of estimation, selection, and sparse computation.The motivation identifies a hierarchical approach as a way to align probability-based pruning with sparse-attention structure.Source: §3 · p. 4
4. Double-P Sparse Attention
Double-P jointly targets top-p estimation accuracy, selection overhead, and sparse-attention cost through cluster-level estimation followed by adaptive token-level refinement.
4. Double-P Sparse Attention
- Double-P combines cluster-level top-p estimation with adaptive token-level refinement to jointly optimize accuracy, selection overhead, and sparse-attention cost.The design replaces fixed token budgets with probability-driven pruning.Source: §4 · p. 4
4.1. Prefilling
- During prefilling, k-means groups similar keys into clusters whose centroids, sizes, and value aggregates compactly summarize the KV cache.These representations support low-cost cluster-level attention-mass estimation and later mixed-granularity computation.Source: §4.1 · pp. 4–5
4.2. Stage 1: Cluster-Level Attention Score Estimation
- Cluster mass is approximated by scaling the centroid logit by cluster size, then normalized and sorted for top-p selection of the smallest set exceeding p1.The approximate masses provide a lower-cost estimate than evaluating all token-level logits.Source: §4.2 · p. 5
4.3. Stage 2: Adaptive Token Budget Allocation for Sparse Attention
- The second top-p stage assigns exact token attention to high-impact clusters and centroid approximations to the rest, with selected cluster counts closely tracking the minimum needed for bounded error.Approximation error is concentrated in a few clusters and varies across layers and decoding steps.Source: §4.3 · pp. 5–6, Figure 7 · p. 6, Figure 6 · p. 6
4.4. Efficient kernel implementation
- Custom top-p kernels use prefix sums and early stopping, fused gathering creates contiguous buffers, and weighted FlashAttention processes exact tokens and centroid summaries together.The fused attention design avoids separate kernels for exact and approximate contributions.Source: §4.4 · p. 6
5. Evaluation
The evaluation measures long-context accuracy and decoding efficiency. Double-P leads sparse attention baselines in accuracy while reducing attention and end-to-end latency.
5.1. Experimental Setup
- Experiments compare Quest, RetroInfer, and Twilight using LLaMA3.1-8B up to 128K tokens, with accuracy measured on RULER and LongBench.The evaluation covers downstream task accuracy and end-to-end decoding latency.Source: §5 · p. 6, §5.1 · p. 7
5.2. Accuracy
- Double-P leads sparse methods on RULER (+1.26 average points) and LongBench, with near-zero accuracy drop versus full attention.RULER gains over the strongest baseline are +1.31, +1.79, and +0.78 absolute points at 16K, 32K, and 64K, respectively.Source: §5.2 · p. 7, Table 2 · p. 7
5.3. Efficiency
- Attention speedups reach 1.74× over Quest-Twilight and 1.78× over RetroInfer; cluster approximation reduces estimation overhead, while adaptive budgets cut final attention cost.Token-level estimation accounts for over 60% of existing top-p attention latency.Source: §5.3 · pp. 7–8, Figure 9 · p. 8
- Up to 1.11× versus Quest and 1.26× versus RetroInfer, Double-P speeds end-to-end decoding while maintaining near-zero accuracy drop; it reaches 2.23× versus full attention.Source: §5.3 · p. 8
5.4. Ablation Study
- Larger cluster-level and token-level thresholds preserve more attention mass and improve accuracy, while smaller thresholds reduce computation; LLaMA-3.1-8B uses (p1, p2) = (0.95, 0.7).The paper also evaluates Qwen-3-8B on RULER and LongBench using (p1, p2) = (0.99, 0.8).Source: §5.4 · p. 8, Table 3 · p. 8
6. Conclusion
Double-P hierarchically optimizes top-p mass estimation, candidate selection, and sparse attention for long-context LLM inference. It reports higher accuracy than prior sparse methods alongside faster attention and decoding.
6. Conclusion
- Double-P combines adaptive probability-driven computation across multiple levels to preserve top-p attention guarantees without fixed-budget constraints.Source: §6 · p. 8
- Double-P remains within a near-zero accuracy gap to full attention and achieves higher accuracy than prior sparse attention methods.Source: §6 · p. 8
- Speedups reach 1.78× over token-level top-p methods and 1.26× over strong sparse baselines end-to-end, and 2.23× versus full attention.Source: §6 · p. 8
Impact Statement
The paper presents an algorithmic and system-level framework for improving long-context inference efficiency and accuracy. It describes no direct negative societal or ethical consequences unique to the work.
Impact Statement
- The work targets computational and memory overhead while preserving attention quality in long-context LLM inference.Source: Impact Statement · p. 9
- The techniques are algorithmic and system-level improvements to existing attention mechanisms, not new capabilities that fundamentally alter LLMs’ intended use.Source: Impact Statement · p. 9
- The authors foresee no direct negative societal or ethical consequences unique to this work beyond those associated with LLMs more broadly.Source: Impact Statement · p. 9
A.1. Ruler
Appendix tables report RULER results across tasks for LLaMA-3.1-8B and Qwen3-8B at 16K, 32K, and 64K context lengths.
A.1. Ruler
- Table 4 reports RULER results across tasks for LLaMA-3.1-8B and Qwen3-8B at 16K context length.Source: Table 4 · p. 12
- Table 5 reports RULER results across tasks for LLaMA-3.1-8B and Qwen3-8B at 32K context length.Source: Table 5 · p. 12
- Table 6 reports RULER results across tasks for LLaMA-3.1-8B and Qwen3-8B at 64K context length.Source: Table 6 · p. 12
A.2. LongBench
Table 7 presents LongBench results across tasks for LLaMA-3.1-8B and Qwen3-8B.
A.2. LongBench
- Table 7 reports LongBench results across tasks.Source: Table 7 · p. 13
- The reported results include LLaMA-3.1-8B.Source: Table 7 · p. 13
- The reported results also include Qwen3-8B.Source: Table 7 · p. 13