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

Thanks for your detailed review and insightful advice. Below, we address each concern and outline planned revisions to strengthen the paper.

1. Sensitivity to Hyperparameter p

We appreciate the reviewer's valuable suggestion regarding the efficiency analysis of parameter p. In our current Figure 6, we have shown the impact of p on model performance (accuracy). The following table shows the impact on efficiency evaluated on the TrivialQA dataset. We will augment this analysis in the camera-ready version.

In summary, the p value directly controls pruning aggressiveness and affects the time of the Sparse Attention stage by the pruned tokens number. We will expand this analysis in the ablation study part of the camera-ready version.

Notably, although we introduce the hyperparameter p in order to get rid of k, p is a more reasonable and tunable hyperparameter. This is because p represents the accumulated probability, which is less impacted by the different distributions that may occur for different heads/layers/queries. In contrast, k is highly sensitive to different distributions, as illustrated in Figure 1 in the paper. This allows us to just tune p for a fixed model. We can use a small dataset to calibrate it.

2. Overhead of INT4 KV Cache

We fully acknowledge the overhead introduced by INT4 KV Cache in Twilight. As the reviewer correctly identified, this manifests in two key dimensions: (1) Storage overhead, as the reviewer mentioned, 12.5% extra overhead to GPU memory; (2) I/O overhead, which occurs during the attention forward process.

The wall-clock impact, as far as we understand, refers to the I/O overhead where the SpGEMV needs to read corresponding INT4 KV blocks to estimate the attention score. This can be found in the breakdown analysis of Figure 7, where “SpGEMV” part reflects this overhead. Note that SpGEMV is a memory-bound kernel like the Attention computation. Hence the total kernel time is roughly equal to the read time.

To deal with the Storage Overhead (on-device INT4 KV Cache), offloading may be a possible solution. But we find it is more suitable to offload the full precision KV instead of INT4 KV as the pruned tokens number (B1, follows notations in Section 4.3) is far less than the initial budget (B0). We discuss this potential scenario in Appendix D.

3. Improve Binary Search during Sampling

We sincerely appreciate the reviewer's valuable suggestion regarding alternative sampling approaches. After careful experimentation on the approach of CDF estimation, we find it is unsuitable for Twilight's architecture because our selector-pruner architecture requires higher precision. A concurrent work that also highlights the potential of top-p attention, Tactic (cited in our paper), makes this method work upon full attention.

For random methods like reservoir sampling, it is worth mentioning that MagicPIG, a work which is evaluated and compared with Twilight, uses another method called LSH sampling. Twilight tends to adopt non-randomized algorithms to maintain stability as an optimizer.

However, we fully acknowledge the optimization potential in our current binary search implementation and are actively exploring improvements. Notably, the FlashInfer community's recent breakthrough in sorting-free GPU sampling kernels presents a promising solution, as documented in their technical blog. Their novel Dual Pivot Rejection Sampling algorithm achieves an approximately 2x speedup over traditional binary search. We are trying to integrate it with Twilight and will incorporate this collaborative advancement in our camera-ready version.

4. Section of Limitations

We discuss the limitations in Appendix E, as required by the submission guidelines.

Thanks for your detailed review and insightful advice. Below, we address each concern and outline planned revisions to strengthen the paper.

1. Additional Computation Cost of Top-p Pruning

We sincerely appreciate the reviewer's insightful observation regarding the Pruner's computational overhead. As noted in Appendix E, we acknowledge that the current estimation process introduces non-trivial latency that can diminish speed benefits in certain edge cases, like the reviewer mentions, small batch sizes or scenarios with minimal token pruning. While we acknowledge the current computational overhead in the Pruner as noted by the reviewer, we maintain that Twilight's core contributions remain significant for the top-p sparse attention and hierarchical selector-pruner framework.

Thanks for your detailed review and insightful advice. Below, we address each concern and outline planned revisions to strengthen the paper.

1. Detailed Description of Efficient Kernel implementation

We fully agree with the reviewer's assessment regarding the importance of our kernel implementations to the overall contribution. While space constraints limited our presentation of these technical details in the current version, we will significantly expand the kernel implementation documentation in the camera-ready version, including warp-level reductions, better utilization of the memory bandwidth, fusion of dequantization and SpGEMV, details on head-wise load balancing, etc.

2. The Impact of p on Efficiency

We sincerely thank the reviewer for catching the incorrect x-axis label in Figure 6, which has now been corrected. Regarding the impact of the p value on computational efficiency, the following shows a basic statistics evaluated on the TrivialQA dataset:

In summary, the p value directly controls pruning aggressiveness and affects the time of the Sparse Attention stage by the pruned tokens number. We will expand this analysis in the ablation study part of the camera-ready version.

Notably, although we introduce the hyperparameter p in order to get rid of k, p is a more reasonable and tunable hyperparameter. This is because p represents the accumulated probability, which is less impacted by the different distributions that may occur for different heads/layers/queries. In contrast, k is highly sensitive to different distributions, as illustrated in Figure 1 in the paper. This allows us to just tune p for a fixed model. We can use a small dataset to calibrate it.

3. Question on Time Breakdown

Thanks for your interesting question! Your understanding is fully correct. And as Table 5 shows, Twilight can prune the original budget (~8192) to extremely low (100 to 400), which largely reduces the time of the Sparse Attention part.

Thanks for your detailed review and insightful advice. Below, we address each concern and outline planned revisions to strengthen the paper.

1. Comparing with Recent Sparse Attention Works

We sincerely appreciate the reviewer's valuable suggestions regarding prior work comparisons. We have already compared Twilight with two SOTA top-k decoding methods Quest/DoubleSparse and one SOTA non-top-k decoding method MagicPIG. But we acknowledge that our current presentation may have caused some confusion about our methodological focus. To clarify: Twilight specifically optimizes the decoding stage, which dominates the computations in emerging workloads like reasoning and conversation, while works [A-E] primarily target the prefill stage, addressing different optimization challenges. Our approach is fundamentally orthogonal to such prefill-stage optimizations. We will enhance the Related Work section to clearly delineate these two optimization domains and better position Twilight within the decode-stage optimization landscape.

2. Novelty

We respectfully but firmly disagree with the characterization of our work as "incremental." While our solution appears simple, this simplicity stems from our key novel insight (as listed below), while prior approaches to dynamic budget allocation in sparse attention (e.g., PyramidKV, AdaKV, DynamicKV) have overcomplicated the problem. Our core contributions are:

• Algorithmic Insight: We demonstrate that the essential dynamism in sparse attention actually lies in the “distribution”, and can be achieved straightforwardly through top-p sampling, bypassing the need for complex, over-parameterized controllers prevalent in prior work.
• Novel Hierarchical Architecture: We firstly introduce a hierarchical “selector-pruner” architecture, which successfully deploys top-p on existing top-k methods elegantly. System-Level Practicality: We show that top-p attention is not just theoretically sound but also hardware-friendly.

As evidenced by our experiments, this approach outperforms existing dynamic methods while being more deployable. We contend that the simplicity achieved through fundamental understanding should not be conflated with incremental progress.

3. The “Contradiction” of the Performance

We appreciate that the reviewer raises an insightful question: why top-p is better than top-k with less tokens while also being worse than full attention? We are glad to provide a clarification.

Why top-p is worse than full attention? The reason is that most training-free sparse attention designs are just numerical approximations to the full attention output, instead of actual knowledge distillation/retrieve. Training-free sparse attention is more like low-bit quantization, which inevitably causes accuracy loss. Trainable sparse attention (e.g. Native Sparse Attention from DeepSeek) has potential to surpass the full attention with better “focus”.

Why Twilight’s top-p is better than top-k with less tokens selected? Twilight employs a hierarchical optimization process with two distinct budget parameters:

• Initial Selection Budget (B0, follows the notation in Section 4.3): Our Token Selector (top-k) uses a conservative, relatively large B0 to maximize token coverage (Section 4.1). Crucially, this stage's computational cost is largely independent of B0's size so that we can use a conservative value of B0.

• Pruned Budget (B1): The subsequent pruning stage aggressively reduces the token count to B1 using INT4-approximated importance scores.

This two-stage approach provides key advantages over conventional top-k:

• Broader Context Awareness: B0 > typical top-k budgets. Therefore Twilight sees more candidate tokens.
• Precise Selection: INT4 scoring offers better importance estimation than typical top-k (which often reaches a sparsity of 1/16, as it is discussed in Section 4.

This "look more, prune better" strategy consistently outperforms standard top-k attention in both accuracy and efficiency metrics. Note that there is also a trade-off in this two-stage method, which is calculated in Section 4.3.

4. Limited Long-Context Performance

We respectfully disagree with the characterization of Twilight as "inferior" to MagicPIG on long-context benchmarks. Our analysis of Tables 2 and 3 suggests the two methods are statistically comparable, with performance differences falling within expected experimental variance (<1%). Notably, Twilight with DS slightly outperforms MagicPIG in both Tables.

We acknowledge MagicPIG as an excellent state-of-the-art sparse attention method that, like Twilight, demonstrates the advantages of moving beyond traditional top-k approaches. Both works provide independent validation of top-p sampling's superiority. Twilight has an advantage in respect of integration with existing algorithms due to its hierarchical architecture, while MagicPIG is a totally new algorithm-hardware co-design, with a unique sampling mechanism and attention computation workflow.

5. Average Post-Pruning Budget in Short-Context Evaluations

We appreciate the insightful advice from the reviewers. Below are the average post-pruning budget statistics in three short-context benchmarks. We will include them in our camera-ready version and we appreciate this opportunity to strengthen our experimental validation.

Thanks for your detailed review and insightful advice. Below, we address each concern and outline planned revisions to strengthen the paper.

1. Benefit-Overhead Tradeoff during Selection and Pruning

We appreciate the reviewer for mentioning the tradeoff between the extra overhead of Twilight and its benefits it brings. In Section 4.3, we discuss the overhead of Twilight and empirically analyze the time breakdown of Twilight’s attention operator in Figure 7, where the visualization makes this tradeoff particularly clear.

2. Analysis of Selected Token Distribution between Top-p and Top-k

We thank the reviewer for their insightful suggestion. In response, we will augment the experiment section with an in-depth analysis of token selection distributions (like Figure 3 in our paper, but more detailed) under both top-p and top-k sampling. This analysis will elucidate why top-p consistently achieves superior performance despite selecting fewer tokens, providing readers with a clearer understanding of its efficiency advantages.

3. Integration on Other Frameworks

We fully acknowledge the importance of system integration, especially given the practical focus of our method and its substantial system-level contributions. To clarify, Twilight's integration operates at multiple levels, which we categorize as follows:

• Kernel Library (FlashInfer). Currently Twilight leverages FlashInfer for high-performance implementations of key LLM operators (e.g., attention, FFN). While FlashInfer was chosen for its superior performance and flexibility, integration with alternative kernel libraries (e.g., Triton) is also possible. We plan to explore compiler-based approaches (e.g., Triton) in future work.
• Existing Sparse Algorithms (e.g., Quest). Here, "integration" refers to applying Twilight's optimizer to existing sparse algorithms. As detailed in Section 4, we demonstrate this through experiments with Quest, DS, and others.
• LLM Serving Frameworks (vLLM, SGLang). Serving frameworks typically build upon kernel libraries (e.g., FlashInfer is used by both vLLM and SGLang). Section 4.3 briefly discusses this integration layer, but we agree with the reviewer that concrete deployment examples would strengthen the work. We will add real-world implementation details in the camera-ready version.