PaperlayerView exact revision on arXiv

Source-linked AI summary

Twilight: Adaptive Attention Sparsity with Hierarchical Top-$p$ Pruning

Chaofan Lin, Jiaming Tang, Shuo Yang, Hanshuo Wang, Tian Tang, Boyu Tian, Ion Stoica, Song Han, Mingyu Gao

arXiv:2502.02770v5cs.LGcs.CL

TL;DR

Long-context decoding faces rising KV-cache memory and data-loading costs, while fixed top-k budgets do not adapt to changing attention distributions. Twilight combines a base token selector with top-p pruning, and the paper reports up to 98% token pruning alongside faster attention and end-to-end decoding.

  • Problem

    Long-context decoding incurs growing KV-cache memory and repeated memory-loading costs, while fixed top-k budgets struggle to match runtime-varying attention distributions.

    Source: §1 · p. 1

  • Method

    Twilight uses a hierarchical select-then-prune framework: a base sparse-attention algorithm selects a conservative subset, then top-p pruning adaptively retains tokens.

    Source: §1 · p. 2

  • Results

    Twilight adaptively prunes up to 98% of tokens, with a 15.4× self-attention speedup and a 3.9× reduction in end-to-end per-token latency.

    Source: §6 · p. 10

  • Takeaways & Limitations

    The accumulated-probability threshold p is less influenced by distribution differences across heads, layers, and queries than k, making it a more reasonable tunable hyperparameter.

    Source: §5.3 · p. 10

    Twilight adds 1/8 extra KV-cache memory overhead, though some existing INT4 caches can be reused instead.

    Source: §4.3 · p. 7

Abstract

from arXiv · show

Leveraging attention sparsity to accelerate long-context large language models (LLMs) has been a hot research topic. However, current algorithms such as sparse attention or key-value (KV) cache compression tend to use a fixed budget, which presents a significant challenge during deployment because it fails to account for the dynamic nature of real-world scenarios, where the optimal balance between accuracy and efficiency can vary greatly. In this paper, we find that borrowing top-$p$ sampling (nucleus sampling) to sparse attention can surprisingly achieve adaptive budgeting. Based on this, we propose Twilight, a framework to bring adaptive sparsity to any existing sparse attention algorithm without sacrificing their accuracy. Empirical results show that Twilight can adaptively prune at most 98% of redundant tokens, leading to $15.4\times$ acceleration in self-attention operations and $3.9\times$ acceleration in end-to-end per token latency in long context LLM decoding.

1 Introduction

Long-context decoding faces growing KV-cache costs, while fixed top-k budgets struggle to fit changing attention distributions. Twilight applies hierarchical top-p pruning to adapt token selection; the paper reports faster sparse attention with nearly no accuracy loss.

1 Introduction

  • Long-context capabilities support applications including retrieval, summarization, code generation, video processing, and chain-of-thought reasoning.The introduction notes context windows reaching 1M to 10M tokens.Source: §1 · p. 1
  • Long-context decoding incurs growing computational and memory costs because the KV cache grows with sequence length and must be repeatedly loaded.The resulting memory use also challenges continued scaling of long-context LLMs.Source: §1 · p. 1
  • Fixed top-k budgets can over-select focused attention or under-select diffuse attention, while the best budget also varies by algorithm.Different algorithms may need offline budget calibration to compensate for estimation inaccuracy.Source: §1 · p. 2, Figure 1 · p. 2, Figure 2 · p. 2
  • Twilight first selects a conservative subset with the base algorithm, then retains only top-p tokens to adapt the budget dynamically.Top-p selects enough tokens to accumulate the target attention weight.Source: §1 · p. 2
  • 1.4× speedup over existing sparse attention methods with nearly no accuracy loss; self-attention and end-to-end decoding speed up 1.4× and 1.35×, respectively.The evaluation covers medium-context and long-context benchmarks.Source: §1 · pp. 2–3

2 Related Work

Related work spans top-k sparse attention, dynamic-budget studies, alternatives to top-k, and other KV-cache optimizations. Twilight is positioned as addressing budget selection and as orthogonal to several cache optimizations.

2 Related Work

  • Top-k sparse-attention methods differ in token eviction, criticality estimation, and trainability, but require advance budget selection and can over- or under-select.The cited methods include KV-cache compression, token-selection methods, and trainable sparse attention.Source: §2 · p. 3
  • Prior dynamic-budget studies report variation across layers, heads, and prompts; this paper attributes that dynamism to differences in attention-weight distributions.The cited studies tend to focus on only one aspect of the variation.Source: §2 · p. 3
  • Other non-top-k approaches include LSH sampling, prefill-stage adaptive sparsity, and Tactic's top-p sparsity with function-fitted weight estimates.The passage says Tactic usually overestimates the budget and that MagicPIG requires algorithm-system co-design.Source: §2 · p. 3
  • Twilight is described as orthogonal to KV-cache quantization, linear attention, and memory-efficient attention, and can be combined with them.Source: §2 · p. 3

3 Bringing Top-p Sampling to Sparse Attention

This section frames sparse attention as selecting token subsets to approximate full attention, then argues that runtime-varying attention distributions make a universal top-k budget difficult and motivates top-p selection by cumulative weight.

3.1 Problem Formulation

  • Oracle top-k selects a fixed-size subset maximizing summed attention weights; this formulation is a theoretical upper bound for current top-k sparse attention methods.Source: §3.1 · p. 4

3.2 Rethinking the Problem of Top-k

  • Top-p selects the minimum number of tokens needed to meet a cumulative attention threshold, adapting the budget to runtime-varying distributions.The passage explains that flat distributions need more tokens than peaked distributions to reach the same threshold.Source: §3.2 · p. 4
  • At p = 0.8, top-p selects B = 97 for an error requirement similar to B = 1024, while fixed B = 16 under-selects.The passage contrasts fixed B = 16 and B = 1024 as under-selection and over-selection, respectively.Source: §3.2 · p. 5

4 Twilight

Twilight adapts existing sparse attention through hierarchical top-p pruning while addressing compatibility, precision, and GPU implementation challenges. Its design combines a conservative base selector with top-p pruning and kernel optimizations, with explicit runtime and memory tradeoffs.

4 Twilight

  • Twilight targets three hurdles in adapting top-p to existing sparse attention: incompatible selection layouts, higher weight-precision needs, and GPU kernel and load-balancing work.The framework addresses these challenges with hierarchical pruning, efficient kernels, and 4-bit K-cache quantization.Source: §4 · p. 5

4.1 Hierarchical Pruning with a Select-then-Prune Architecture

  • The Token Selector treats the base algorithm as a black box with a conservative budget, then the pruner retains the minimum top-p subset before sparse attention runs.The final sparse-attention kernel computes only on the retained top-p tokens.Source: §4.1 · pp. 5–6, Figure 5 · p. 6

4.2 Efficient Kernel Implementations

  • At p = 0.85, 2-bit quantization lowers selected-token attention-weight sums, while 4-bit and 8-bit methods maintain stability.The paper reports that 4-bit precision balances accuracy and efficiency for top-p.Source: §4.2 · p. 6, Figure 6 · p. 6
  • INT4 K vectors are dequantized in shared memory, limiting global-memory accesses to at most one-quarter; binary search fuses top-p operations, while flattened heads reuse FlashInfer load balancing.The sparse GEMV kernel stores quantized K cache in paged form.Source: §4.2 · pp. 6–7

4.3 Overhead Analysis and Discussion

  • For B0 = N/4 and B1 = N/64, modeled speedup is approximately 2×; top-p overhead is omitted because SpGEMV dominates when B0 is around N/8 to N/4.Execution time comprises Token Selector, Pruner, and sparse-attention stages.Source: §4.3 · p. 7
  • Twilight adds 1/8 extra KV-cache memory for INT4 K, though some base algorithms already maintain it; the design aligns with PagedAttention and common serving techniques.The paper suggests offloading or selective quantization if GPU memory becomes a bottleneck.Source: §4.3 · p. 7

5 Evaluation

The evaluation compares Twilight-enhanced sparse-attention methods with established baselines across long- and medium-context benchmarks, then examines efficiency and parameter sensitivity. Results cover accuracy, latency, and execution-time trade-offs.

5.1 Accuracy Evaluation

  • The evaluation used Longbench, RULER, GSM8K, COQA, and PG-19, comparing Quest, DS, and MagicPIG where supported.Experiments used three models and included both MHA and GQA implementations; MagicPIG was omitted for LLaMA-2 because official support was unavailable.Source: §5.1 · p. 8
  • Up to 5.7% score improvement over the original method and pruning of up to 98% of redundant tokens were reported on Longbench; LLaMA-3.1 had less than 1% accuracy loss.Table 2 reports average scores across 12 Longbench tasks; the LLaMA-3.1 result came with a slight increase in budget usage.Source: Table 2 · p. 8, §5.1 · p. 8
  • On RULER, Quest-Twi reached performance comparable to MagicPIG, while DS-Twi surpassed all existing methods.The evaluation used LLaMA-3.1-8B-Instruct and included specialized non-retrieval tests.Source: §5.1 · p. 8, Table 3 · p. 9
  • On three medium-context benchmarks, Twilight outperformed Quest and DS by significant margins, with nearly zero loss versus full attention.The baselines used a budget of 128, comparable to the budget after Twilight pruning.Source: §5.1 · p. 8, Table 4 · p. 9

5.2 Efficiency Evaluation

  • At batch size 64, Quest-Twi outperformed Quest by about 2× in self-attention time; Figures 7–8 report latency and TPOT results across serving settings.The efficiency evaluation also presents latency and speedups across sequence lengths and batch sizes.Source: Figure 7 · p. 9, Figure 8 · p. 9, Figure 10 · p. 10

5.3 Ablation Study

  • At p ≈ 0.85, accuracy and efficiency balanced on the tested PG-19 perplexity and TrivialQA kernel-speed measures, motivating that setting for Longchat.For a 32k retrieval task, Twilight pruned Quest's budget from 8192 to 256, reducing sparse-attention kernel time with minor overhead.Source: §5.3 · p. 10

6 Conclusion

The conclusion presents Twilight as a select-then-prune framework for budget variation in sparse attention, reporting token pruning and speedups. The appendix identifies four levels of budget dynamism and describes the full Longbench results table.

6 Conclusion

  • Up to 98% token pruning yielded a 15.4× self-attention speedup and a 3.9× reduction in end-to-end per-token latency; Twilight added 1.4× speedup over its base sparse-attention algorithm.The conclusion attributes the budget challenge to dynamic attention-weight distributions and describes Twilight's hierarchical select-then-prune design.Source: §6 · p. 10

A Budget Dynamism at Different Levels

B.1 Implementation of Mixed-Precision SpGEMV

The implementation uses a quantized-key SpGEMV kernel and adapts token selection to GQA's grouped query heads. Its design combines on-the-fly dequantization with group-level token selection and variable-length attention.

B.1 Implementation of Mixed-Precision SpGEMV

  • The SpGEMV kernel computes FP16-query and INT4-key products, loading and dequantizing keys asynchronously while overlapping later-block loads with current-block computation.FP16 stores the dequantized key cache as a tolerable score-estimation trade-off.Source: §B.1 · p. 20
  • Per-head dynamic quantization stores FP16 scales and zero points; INT4 key values are packed two per byte and dequantized on the fly.The dequantization routine uses custom PTX conversions between INT4 and FP16.Source: §B.1 · p. 20
  • The quantization-bit ablation reports that dequantization and dot-product optimizations make the operator memory-bound, so it benefits from quantization.Figure 12 shows SpGEMV latency at different quantization bit widths.Source: §B.1 · p. 20, Figure 12 · p. 20

B.2 Tackling Group Query Attention

  • For GQA, Twilight selects the union of tokens identified by query heads in each group, avoiding the repeated reads of loading tokens independently per head.The flattened paged KV cache supports head-wise variable-length attention for MHA and group-wise variable-length attention for GQA.Source: §B.2 · pp. 20–21, Figure 13 · p. 20

C Full Results on Longbench

The supplied passage directs readers to Table 5 for the full Longbench results.

C Full Results on Longbench

D Accuracy Comparison with Token Dropping Methods

DS-Twilight outperforms StreamingLLM and SnapKV on Longbench with the Longchat-7B-v1.5-32k model. The section compares Twilight with token-dropping methods.

D Accuracy Comparison with Token Dropping Methods

E Efficiency Evaluation in Offloading Scenarios

In offloading scenarios, Twilight reduces the number of loaded tokens with a fixed estimation cost, and can achieve up to 16× speedups compared with Quest.

E Efficiency Evaluation in Offloading Scenarios

F Limitations and Future Work

Twilight has non-negligible estimation overheads, making it particularly advantageous when KV-cache loading dominates; head-wise dynamism also creates integration challenges with GQA.

F Limitations and Future Work

  • Twilight has non-negligible estimation overheads, so it is particularly advantageous when KV-cache loading dominates, such as large-batch serving or offloading.Future work could optimize estimation to improve end-to-end latency and throughput.Source: F Limitations and Future Work · p. 22
  • Head-wise dynamism is unfriendly with GQA and creates challenges integrating Twilight with new model architectures, including MLA.The paper identifies architecture integration and estimation optimization as future research directions.Source: F Limitations and Future Work · p. 22
Loading 2502.02770v5…