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 (清华大学, 麻省理工学院, 加州大学伯克利分校)

A1 主要贡献

本文深入探讨了在长上下文大语言模型（LLMs）中，利用注意力稀疏性进行加速的一个核心挑战：现有的大多数稀疏注意力算法采用固定的预算（即top-k方法中的k值）来选择计算中使用的token数量。这种静态决策在实际部署中存在严重问题，因为它无法适应真实场景的动态性，即在准确性和效率之间的最佳平衡点是不断变化的。

核心问题与研究目标：
- 问题：现有的top-k稀疏注意力方法依赖于一个固定的token预算，但这在不同的注意力头、层、查询甚至任务中都可能是次优的。如图1所示，某些注意力分布（称为“集中式注意力”）可能只需要很少的token就能累积足够的注意力权重，而固定的预算会导致“过度选择”，浪费计算资源。相反，对于其他较平坦的分布（称为“分散式注意力”），固定的预算可能导致“选择不足”，从而损害模型精度。此外，如图2所示，不同的稀疏算法为弥补其估计误差，需要不同程度的过度选择，这使得预算B的选择变得非常困难，需要针对每个算法进行离线校准。
- 目标：本文的目标是提出一种能够根据运行时注意力权重的动态分布，自适应地决定预算大小的方法，从而在不牺牲准确性的前提下，提升稀疏注意力的效率。

创新点与主要贡献：
1. 引入Top-p采样：本文揭示了一个关键洞见，即在稀疏注意力中利用top-p采样（也称核采样）的思想，可以实现高效和自适应的预算决策。与top-k选择固定数量的token不同，top-p选择累积概率超过阈值p的最小token集合，从而本质上适应了不同的注意力权重分布。
2. 提出Twilight框架：本文提出了Twilight，一个分层的KV缓存剪枝框架，它能够增强任何现有的稀疏注意力算法，赋予其自适应预算决策的能力，而无需对基础算法进行大规模修改。Twilight采用“先选择后剪枝”的架构：首先，让基础算法使用一个保守的、较大的预算选择一个候选token子集；然后，Twilight剪枝器再通过top-p进一步精炼这个子集。
3. 高效实现与评估：本文对Twilight进行了全面的准确性和效率评估。实验结果表明，Twilight可以在中长文场景下自适应地剪枝高达98%的token，且几乎没有精度损失。与最先进的稀疏注意力机制相比，Twilight在自注意力算子上实现了1.4倍的加速，在端到端解码中实现了1.35倍的加速。


(a) Top-k稀疏性（H2O, Quest等）固定的k-token预算导致过度选择或选择不足。
(b) Top-p稀疏性（本文方法）动态调整预算以累积足够的注意力权重。
图 1: top-k和top-p在注意力稀疏性方面的比较。近似注意力通常采用池化、通道剪枝和量化等技术来近似查询（$\tilde{Q}$）和键（$\tilde{K}$），并估计注意力权重。这些权重随后用于选择重要token以进行稀疏注意力计算。(a) 大多数现有设计使用的Top-k稀疏性依赖于固定的k-token预算，常常导致过度选择（$\sum \tilde{p}_i > 0.9$）或选择不足（$\sum \tilde{p}_i < 0.9$）。(b) 我们提出的top-p稀疏性动态调整预算，以累积恰好足够的注意力权重（$\sum \tilde{p}_i = 0.9$），从而实现更高效和自适应的稀疏注意力。


图 2: 在PG-19数据集上，不同top-k稀疏注意力方法中KV缓存预算与困惑度之间的关系。

A3 将Top-p采样引入稀疏注意力

问题形式化。我们首先对稀疏注意力机制进行形式化。在解码阶段，我们有查询向量$q \in \mathbb{R}^{1 \times d}$，以及KV缓存$K, V \in \mathbb{R}^{n \times d}$。这里，$d$表示头的维度，$n$表示上下文长度。
定义 3.1 (稀疏注意力)。设$\mathcal{I}$为所选索引的集合。稀疏注意力计算如下：

$$\hat{\mathbf{o}}=\operatorname{softmax}\left(\frac{\mathbf{q} \cdot \mathbf{K}^T}{\sqrt{d}}\right) \boldsymbol{\Lambda}_{\mathcal{I}} \mathbf{V}=\mathbf{W} \boldsymbol{\Lambda}_{\mathcal{I}} \mathbf{V} \in \mathbb{R}^{1 \times d}$$

其中$\Lambda_{\mathcal{I}} \in \mathbb{R}^{n \times n}$，$\Lambda_{\mathcal{I}}[i, j] = \begin{cases} 1 & \text{if } i=j \text{ and } i \in \mathcal{I} \ 0 & \text{otherwise} \end{cases}$。


图 3: 注意力权重中观察到的多样化分布。最左边的图像展示了一个“平坦”分布（分散式注意力），其中权重接近均匀分布。中间的图像描绘了一个“尖峰”分布（集中式注意力），其中权重集中在两侧的token上。当如最右边的图像所示叠加时，这些分布之间的差异变得显而易见。

设精确的注意力输出为$o = \text{softmax}(\frac{q \cdot K^T}{\sqrt{d}})V \in \mathbb{R}^{1 \times d}$。为了最小化误差$|o - \hat{o}|$，我们需要仔细选择在$\mathcal{I}$中使用的token子集。然而，在不加载完整KV缓存的情况下直接优化这个目标函数是具有挑战性的。根据Frobenius范数的次乘法性质，我们可以将误差界定为公式2。早期的研究表明，$V$的分布相对平滑【索引42，Atom: Low-bit quantization for efficient and accurate LLM serving+2024+MLSys】，这意味着$|V|_F$可以被视为一个常数。

$$\begin{aligned} \begin{aligned} \mathcal{L}=\|\mathbf{o}-\hat{\mathbf{o}}\| & =\|\mathbf{W}(\mathbf{\Lambda}_{\mathcal{I}}-\mathbf{1}^{n \times n}) \mathbf{V}\| \\ & \leq\|\mathbf{W}(\mathbf{\Lambda}_{\mathcal{I}}-\mathbf{1}^{n \times n})\| \cdot\|\mathbf{V}\|_F \end{aligned} \end{aligned}$$

因此，目标变为最小化$|W(\Lambda_{\mathcal{I}} - 1^{n \times n})| = 1 - \sum_{i \in \mathcal{I}} W[i]$，这意味着选择一个能最大化其注意力权重之和的token子集。如果我们固定这个子集的大小，即$|\mathcal{I}|$，那么我们就得到了oracle top-k注意力：
定义 3.2 (Oracle Top-k稀疏注意力)。给定预算$B$，

$$\mathcal{I}=\arg \max _{\mathcal{I}} \sum_{i \in \mathcal{I}} \mathbf{W}[i] \quad \text { s.t. }|\mathcal{I}|=B$$

这可作为当前top-k稀疏注意力方法的理论上界。

反思Top-k的问题。如前所述，top-k稀疏注意力的致命弱点在于难以确定一个普遍适用于所有场景的预算B。我们发现，这一困境与之前在LLM采样阶段遇到的问题非常相似，在采样阶段，模型从预测的概率分布中采样最终的输出token。为了解决top-k采样无法适应不同下一词分布的问题，研究者提出了核采样（Nucleus sampling）【索引43，The curious case of neural text degeneration+2019+arXiv】，也称为top-p采样。受此启发，我们更仔细地研究了注意力权重的分布。如公式2所示，输出误差受限于所选注意力权重之和。因此，目标应变为选择最小数量的token B来满足给定的输出误差要求。图3展示了图1中提到的几种真实LLM中的两种不同类型的注意力权重分布。很容易观察到，与尖峰分布相比，在平坦分布中必须选择更多的token才能达到相同的累积阈值。因此，我们认为预算动态性的核心原因是运行时注意力权重分布的动态性。因此，我们引入top-p稀疏注意力，通过直接对注意力权重之和应用一个阈值。
定义 3.3 (Oracle Top-p稀疏注意力)。给定阈值$p$，

$$\mathcal{I}=\arg\min_{\mathcal{I}}|\mathcal{I}|\quad\text{s.t.}\sum_{i\in\mathcal{I}}\mathbf{W}[i]\geq p$$

与top-k相比，top-p更有优势，因为它从公式2中提供了$(1 - p) \cdot |V|_F$的理论误差上界。在这种情况下，top-p将预算降低到尽可能低的水平，使其既高效又适应不同的分布。为了展示top-p如何减少预算，我们研究了一个注意力分数的真实分布，如图4所示。与两种固定预算策略B=16和B=1024（分别导致选择不足和过度选择）相比，p=0.8的点选择了一个非常小的预算B=97，以达到与B=1024相似的误差要求。


图 4: 在一个示例注意力头中，不同预算选择的累积注意力分数。

A2 方法细节

Twilight框架设计。借助高效且自适应的top-p稀疏注意力，我们的主要目标是利用它来赋予现有算法自适应的预算选择能力，而不是简单地发明另一种稀疏注意力设计。这主要有两个动机。一方面，尽管预算选择存在挑战，但现有的稀疏注意力算法凭借其有效的token选择策略，已在LLM服务系统【索引44，Efficient memory management for large language model serving with PagedAttention+2023+SOSP；索引45，Efficiently programming large language models using SGLang+2023+arXiv】中取得了显著成功。这些策略可以被我们的自适应稀疏性轻松复用和增强。另一方面，我们预计未来的稀疏注意力方法可能仍会采用top-k选择。通过开发一个通用的解决方案，我们的目标是自动为这些未来方法配备自适应注意力稀疏性，同时避免大规模的重新设计。因此，我们将我们的系统Twilight定位为现有算法的优化器。

应用Top-p的挑战。然而，将top-p应用于各种现有的稀疏注意力算法在算法和系统层面面临三个关键挑战。(C1) 并非所有算法都适用于top-p。Top-p对注意力权重的布局有严格的约束。例如，简单地在Quest【索引9，Quest: Query-aware sparsity for efficient long-context LLM inference+2024+arXiv】中用top-p替换top-k是行不通的，因为Quest对每页布局（每页16个token）的权重执行最大池化操作。此外，其他一些方法【索引46，TidalDecode: Fast and accurate LLM decoding with position persistent sparse attention+2024+arXiv；索引21，RetrievalAttention: Accelerating long-context LLM inference via vector retrieval+2024+arXiv】根本不使用注意力权重来选择关键token。(C2) 为top-p估计权重比为top-k更难。为了在不加载全部K数据的情况下找到关键token，通常使用K缓存的低精度表示【索引12，Post-training sparse attention with double sparsity+2024+arXiv；索引22，PQCache: Product quantization-based KVcache for long context LLM inference+2025+Proceedings of the ACM on Management of Data】。然而，top-p的精度要求高于top-k，因为前者需要一定程度的数值准确性，而后者仅要求顺序性。表1比较了top-k、top-p和全注意力。top-p注意力的精度要求介于其他两者之间，需要重新考虑K缓存的适当精度选择。(C3) 需要系统级优化。由于我们的工作是首次将top-p引入注意力权重，相关算法需要在GPU上高效实现，包括并行算法设计和核函数优化方面的努力。在4.1节中，我们通过提出一个统一的分层剪枝框架来解决C1。在4.2节中，我们通过高效的核函数实现（Top-p、SpGEMV、Attention）和K缓存的4位量化来减轻运行时开销，从而解决C2和C3。最后，在4.3节中，我们分析了Twilight的开销并讨论了一些额外问题。

表 1: 不同剪枝方法在注意力权重上的比较。“Normalization”指softmax。

4.1 采用“先选择后剪枝”架构的分层剪枝

提出两步式分层剪枝流程。为了统一支持各种稀疏注意力机制，我们提出了一个两步式的分层剪枝过程。我们首先将基础算法捕获为一个黑盒的“Token选择器”（Token Selector），只要它具有选择关键token子集的通用语义，而具体的选择算法细节无关紧要。我们让Token选择器使用一个保守的、相对较大的预算，例如1/4的稀疏度。然后，我们在此之后设置一个“Twilight剪枝器”（Twilight Pruner），通过仅保留top-p的token来进一步优化所选的索引，即注意力权重总和超过阈值p的最小子集。我们将此设计称为“先选择后剪枝”（Select-then-Prune）架构，如图5中间部分所示。最终的稀疏注意力核函数因此只在top-p的token上进行计算，从而实现了第3节中证明的效率和自适应性优势。


图 5: Twilight架构。Twilight整合了某个现有的基础稀疏注意力算法并对其进行进一步优化。它分三步计算自注意力。首先，Token选择器使用基础算法在一个保守的预算下选择关键token。然后，Twilight剪枝器通过top-p阈值化对所选的token子集进行剪枝。最后，将剪枝后的token索引传递给稀疏注意力核函数以执行注意力计算。

4.2 高效的核函数实现

Twilight架构细节概述。现在我们简要描述Twilight架构的细节，特别是剪枝器（Pruner）步骤。有关核函数实现的更多细节，请参阅附录B。

带有4位Key缓存量化的高效SpGEMV。剪枝器的起始部分与其他稀疏注意力算法类似，即估计token的重要性。正如我们在3.1节中形式化的，这可以通过估计q和K之间的相似度，即$q \cdot K$来完成。由于加载K是内存密集型的，我们通过将K量化到较低精度来减少内存访问成本。但是我们应该选择什么精度呢？表1显示top-p的精度要求介于top-k和全注意力之间。一些现有的top-k设计【索引12，Post-training sparse attention with double sparsity+2024+arXiv；索引22，PQCache: Product quantization-based KVcache for long context LLM inference+2025+Proceedings of the ACM on Management of Data】已将K缓存的压缩推向了1到2位的极低精度。对于全注意力，SageAttention【索引40，SageAttention: Accurate 8-bit attention for plug-and-play inference acceleration+2025+ICLR】已证明使用平滑K和逐块量化可达到8位精度。在本文中，我们通过经验发现，4位精度在top-p的准确性和效率之间取得了平衡，如图6所示。图中，2位量化的注意力权重总和显著下降，而4位和8位方法都保持了足够的稳定性。因此，我们基于FlashInfer【索引47，FlashInfer: Efficient and customizable attention engine for LLM inference serving+2025+arXiv】（一个用于LLM服务的高性能核函数库）实现了一个高效的稀疏GEMV（SpGEMV）核函数。这里的“稀疏”意味着量化的K缓存数据以分页方式【索引44，Efficient memory management for large language model serving with PagedAttention+2023+SOSP】存储/加载，以与原始KV缓存布局对齐。我们在GPU上维护这个额外的INT4非对称量化K缓存，如图5右侧所示。INT4的K向量在共享内存中解包和反量化，将从全局内存的数据访问减少到最多1/4，从而带来了可观的端到端加速。


图 6: 在p = 0.85时，不同量化精度下所选token的归一化注意力权重之和。

通过二分搜索实现高效Top-p。实现top-p采样的一种暴力方法是按降序对元素进行排序，然后累加它们直到总和达到阈值。这在像现代GPU这样的并行硬件中效率很低。由于我们的top-p方法受top-p采样的启发，我们也通过修改FlashInfer【索引47，FlashInfer: Efficient and customizable attention engine for LLM inference serving+2025+arXiv】中的top-p采样核函数来实现这个核函数。具体来说，我们的核函数采用了一种对并行友好的二分搜索方法，如算法1所示。请注意，像max、where和sum这样的元素级操作可以融合成一个单一的循环，该循环在GPU上进行了张量化处理。因此，我们不需要实例化像W0这样的中间变量。

算法 1 Top-p via Binary Search.
输入: 归一化的注意力权重 $W \in \mathbb{R}^{BS \times H \times N}$, top-p阈值 $p$, 超参数 $\epsilon$.
输出: 索引 $\mathcal{I}$, 掩码 $M \in {0, 1}^{BS \times H \times N}$.
$l = 0, r = \max(W), m = (l + r)/2;$
repeat
$W_0 = \text{where}(W < m, 0.0, W);$
$W_1 = \text{where}(W \leq l, \text{INF}, W);$
$W_2 = \text{where}(W > r, -\text{INF}, W);$
if $\text{sum}(W_0) \geq p$ then
$l = m;$
else
$r = m;$
end if
until $\max(W_2) - \min(W_1) \geq \epsilon$
选择索引 $\mathcal{I}$ 并设置掩码 $M$ 其中 $W \geq l;$
return $\mathcal{I}, M;$

感知头动态性的负载均衡。top-p剪枝器实现了头级别的动态预算，但也在注意力核函数中引发了负载不平衡问题。传统实现为所有头分配统一的计算资源。FlashInfer【索引47，FlashInfer: Efficient and customizable attention engine for LLM inference serving+2025+arXiv】深入研究了这种负载不平衡问题，但仅限于动态长度的请求。Twilight进一步重用FlashInfer中的负载均衡算法，通过扁平化头维度来解决头级别的动态性问题。

4.3 开销分析与讨论

执行时间。根据图5中的流水线，Twilight的执行时间由三部分组成：$T_{\text{TokenSel}} + T_{\text{Pruner}} + T_{\text{SparseAttn}}$。与没有Twilight的基线稀疏注意力相比，我们的方法引入了一个额外的延迟项$T_{\text{Pruner}}$，但减少了$T_{\text{SparseAttn}}$。我们的分层架构自然地与分层稀疏性相匹配，其中随着精度的提高，token的数量逐渐减少。假设Token选择器中的基础算法以1/16的稀疏度和/或精度降低来估计token重要性。那么理论上的加速可以表示为$N/16 + B_0 N/16 + B_0/4 + B_1$，其中$B_0 = |I_0|$是基础Token选择器的预算，而$B_1 = |I_1|$是经过Twilight用INT4剪枝后的预算。假设$B_0 = N/4$和$B_1 = N/64$，加速将约为2倍。这里我们忽略了top-p核函数的开销，因为当$B_0$在$N/8$到$N/4$左右时，SpGEMV主导了延迟。

内存开销。Twilight引入了一个额外的INT4量化K缓存，这带来了$1/2 \times 1/4 = 1/8$的额外KV缓存内存开销。然而，这种额外成本并非在所有情况下都会出现。首先，一些基础算法，如DS【索引12，Post-training sparse attention with double sparsity+2024+arXiv】，已经维护了一个INT4的K缓存。其次，最近的一些工作已经探索了INT4全注意力【索引41，SageAttention2: Efficient attention with thorough outlier smoothing and per-thread INT4 quantization+2024+arXiv】。这使我们能够直接在注意力计算中重用由INT4 K缓存计算出的估计注意力权重，而无需维护原始的FP16 K缓存。此外，如果GPU内存成为瓶颈，可以利用卸载和选择性量化（例如，仅为热门token保留额外的INT4 K缓存），我们将其作为未来工作。

与LLM服务系统的集成。我们的系统设计与PagedAttention【索引44，Efficient memory management for large language model serving with PagedAttention+2023+SOSP】天然对齐，因此Twilight可以无缝集成到流行的服务系统如vLLM【索引44，Efficient memory management for large language model serving with PagedAttention+2023+SOSP】和SGLang【索引45，Efficiently programming large language models using SGLang+2023+arXiv】中。其他常用技术，如前缀共享和多阶段注意力【索引48，Parrot: Efficient serving of LLM-based applications with semantic variable+2024+OSDI；索引45，Efficiently programming large language models using SGLang+2023+arXiv；索引49，RelayAttention for efficient large language model serving with long system prompts+2024+arXiv；索引50，ChunkAttention: Efficient self-attention with prefix-aware KV cache and two-phase partition+2024+arXiv；索引51，Cascade inference: Memory bandwidth efficient shared prefix batch decoding+2024+arXiv】，也与Twilight兼容，因为我们使用页级或token级的稀疏操作，并且可以实现灵活的计算流。

A4 实验环境

  • 模型架构：

    • Longchat-7B-v1.5-32k
    • LLaMA2-7B-Chat
    • LLaMA-3.1-8B-Instruct (上下文长度128k)
      这些模型覆盖了多头注意力（MHA）和分组查询注意力（GQA）两种主流实现。

  • 数据集与基准：

    • 长上下文基准：Longbench【索引1，LongBench: A bilingual, multitask benchmark for long context understanding+2024+ACL】和RULER【索引16，RULER: What’s the real context size of your long-context language models?+2024+arXiv】。
    • 中等上下文基准（500到2k tokens）：GSM8K【索引13，Training verifiers to solve math word problems+2021+arXiv】、COQA【索引14，CoQA: A conversational question answering challenge+2019+Transactions of the Association for Computational Linguistics】和PG-19【索引15，Compressive transformers for long-range sequence modelling+2019+arXiv】数据集上的困惑度测试。
    • 效率评估数据集：使用Longbench中的三个任务（Qasper用于问答，GovReport用于摘要，LCC用于编码），提示长度从10k到30k tokens不等。

  • 硬件配置：

    • 所有效率评估均在单张A100 GPU上进行。

  • 软件配置与实现细节：

    • 代码实现：使用CUDA和OpenAI Triton【索引60，Triton: an intermediate language and compiler for tiled neural network computations+2019+MAPL】实现。
    • 依赖库：集成了高性能核函数库FlashInfer【索引47，FlashInfer: Efficient and customizable attention engine for LLM inference serving+2025+arXiv】。对于GSM8K和COQA任务，使用了lm-harness框架【索引55，A framework for few-shot language model evaluation+2021+GitHub】进行评估。
    • 基线方法：
      • SOTA top-k稀疏注意力方法：Quest【索引9，Quest: Query-aware sparsity for efficient long-context LLM inference+2024+arXiv】和DS (Double Sparsity)【索引12，Post-training sparse attention with double sparsity+2024+arXiv】。
      • SOTA non-top-k方法：MagicPIG【索引30，MagicPIG: LSH sampling for efficient LLM generation+2024+arXiv】。
      • 全注意力/密集基线：PyTorch的scaled-dot-product-attention (SDPA)，后端为FlashAttention2 (FA2)【索引39，FlashAttention-2: Faster attention with better parallelism and work partitioning+2024+ICLR】和MemoryEfficient Attention【索引59，xFormers: A modular and hackable transformer modelling library+2022+GitHub】；以及FlashInfer【索引47，FlashInfer: Efficient and customizable attention engine for LLM inference serving+2025+arXiv】。

  • 超参数设置：Twilight的超参数p在LLaMA-2/3上设置为0.95，在Longchat上设置为0.85。

A4 实验结果

长上下文任务（Longbench）的准确性
如表2所示，在Longbench的12个任务上，Twilight显著提升了基线算法的性能。
- 在Longchat模型上，DS-Twilight和Quest-Twilight分别将平均得分提高了5.7%和2.5%，同时成功剪枝了基础算法过度选择的高达98%的冗余token。
- 在LLaMA-3.1-8B-Instruct模型上，Twilight以微小的预算增加为代价，实现了几乎为零的准确率损失（<1%）。

表 2: Longbench上12个不同任务的平均得分。我们报告了将Twilight与每个基础算法集成时的相对误差变化（改善或下降）。详细结果见附录C的表5。

RULER基准测试的准确性
在RULER基准上（表3），使用LLaMA-3.1-8B-Instruct模型，标准Quest表现不如非top-k方法，但DS表现出惊人的竞争力。
- 当与Twilight结合后，Quest-Twi的性能与SOTA的非top-k方法MagicPIG相当。
- DS-Twi则创造了新的记录，超越了所有现有方法。

表 3: RULER上的平均得分。

中等上下文任务的准确性
为了验证Twilight剪枝器本身不会对性能产生负面影响，实验在中等上下文任务上进行，所有基线使用与Twilight剪枝后相当的128预算。如表4所示，Twilight在GSM8K、COQA和PG-19困惑度测试中，性能均显著优于Quest和DS，与全注意力相比几乎没有损失。

表 4: 3个中等上下文基准测试的结果。

效率评估：自注意力加速
图7展示了在不同序列长度和批量大小下，自注意力算子的延迟和加速比。
- FlashInfer-Twi和Quest-Twi相较于FlashAttention2分别实现了高达6.5倍和15.8倍的加速。
- 它们也分别将各自的基础算法FlashInfer和Quest加速了2.4倍和1.4倍。


图 7: 不同序列长度和批量大小下的自注意力延迟和加速比。

效率评估：端到端解码加速
图8展示了在批量大小从32到256的不同服务场景下的端到端解码性能（每输出token时间，TPOT）。
- Quest-Twi相较于FlashInfer实现了高达3.9倍的加速。
- 相较于没有Twilight的Quest，Quest-Twi实现了1.35倍的加速。


图 8: 端到端服务场景中每输出token时间（TPOT）的改进。

消融研究
- 对阈值p的敏感性：如图9所示，实验分析了超参数p对准确性（PG-19困惑度）和效率（稀疏注意力延迟）的影响。结果表明，p是一个比k更合理且可调的超参数，因为它代表累积概率，受不同分布影响较小。对于Longchat-7B模型，p ≈ 0.85时在准确性和效率之间达到平衡。
- Twilight的时间分解：图10展示了Twilight三个组成部分（Token选择器、SpGEMV、Top P）的执行时间。结果与理论成本模型（4.3节）一致，表明Twilight通过引入较小的开销，显著减少了稀疏注意力核函数的时间。


图 9: 不同阈值p值下的PG-19困惑度（准确性）和稀疏注意力延迟（效率）。


图 10: 自注意力的时间分解。在批量大小为64时，Quest-Twi的性能大约是Quest的2倍。

A5 结论

本文首先指出现有的top-k稀疏注意力方法由于注意力权重分布的动态性，难以找到最优预算。随后，本文引入了Twilight，一个采用分层“先选择后剪枝”架构的框架，利用top-p采样来解决此问题。Twilight能够自适应地剪枝高达98%的token，为自注意力算子带来15.4倍的加速，并将端到端每token延迟降低3.9倍。与其所应用的基础稀疏注意力算法相比，Twilight提供了额外的1.4倍加速。我们的工作强调了自适应注意力稀疏性的重要性，并为未来稀疏注意力机制的研究铺平了一条有前景的道路。

A6 附录

A 不同层级的预算动态性

预算动态性分析。正如第1节所介绍的，存在着不同层级的预算动态性。我们提出并分析了四种不同层级的KV缓存预算动态性，如图11所示。它们分别是提示级（prompt-wise）、查询级（query-wise）、层级（layer-wise）和头级（head-wise）的动态性。


图 11: 在oracle top-p注意力中观察到的预算动态性。我们在四个维度上观察到动态性：不同的提示（任务）、同一提示内的不同查询、同一查询中的不同层，以及同一层中的不同头。

B 核函数实现细节

B.1 混合精度SpGEMV的实现

计算过程。如第4.2节所述，我们的实现需要一个GEMV核函数，用于计算FP16查询向量和INT4量化键矩阵的乘积（$q_{fp16} \cdot K_{int4}$），并支持分页索引。我们为此改编了FlashInfer【索引47，FlashInfer: Efficient and customizable attention engine for LLM inference serving+2025+arXiv】的注意力解码核函数。核函数的执行遵循两个主要步骤：（1）使用cp.async从全局内存异步加载量化的K缓存并将其反量化到共享内存；（2）计算点积。为减轻长内存延迟，我们采用了一个两阶段流水线，将后续块的数据加载与当前块的计算重叠。我们使用FP16来存储反量化后的K缓存，而不是FP32，以优化计算，因为这种精度权衡作为分数估计器是可以接受的。

反量化。我们遵循QServe【索引61，QServe: W4A8KV4 quantization and system co-design for efficient LLM serving+2024+arXiv】的设计，采用逐头、动态的KV量化，并使用与K缓存相同的分页内存布局存储每个头的FP16缩放因子和零点。K矩阵使用逐头的量化参数（缩放因子和零点）进行动态反量化。我们的反量化例程建立在【索引62，Who says elephants can’t run: Bringing large scale moe models into cloud scale production+2022+arXiv】的快速算法（在NVIDIA的FasterTransformer【索引63，FasterTransformer: Providing a script and recipe to run the highly optimized transformer-based encoder and decoder component+2023+GitHub】中实现）之上，该算法利用自定义PTX汇编指令在INT4和FP16之间进行高效的类型转换。

位打包。INT4的K元素被打包到一个uint8_t缓冲区中，每个8位字节存储两个4位元素——这与C++的字节寻址特性相符。为了简化反量化逻辑，我们首先为每个INT4元素加上一个+128的偏移量，将其转换为无符号值，然后再以交错方式打包。对这个打包缓冲区的地址计算被重新映射为以4位为粒度的步长；这是通过将有效字节偏移量减半来实现的【索引42，Atom: Low-bit quantization for efficient and accurate LLM serving+2024+MLSys】。我们进行了一项关于量化位数对效率影响的消融研究，如图12所示。我们对反量化和点积的优化使得该算子受内存带宽限制，因此能从量化中受益。


图 12: 不同量化位数的SpGEMV算子延迟。

B.2 处理分组查询注意力（Group Query Attention）

GQA的挑战与解决方案。分组查询注意力（GQA）【索引54，The Llama 3 herd of models+2024+arXiv】是近期模型架构（如LLaMA 3）广泛采用的一项技术，它将一组查询头映射到单个键值头。然而，这种结构与查询感知的稀疏注意力天然不兼容。不兼容的原因在于，查询感知的稀疏注意力依赖单个查询向量来识别重要token，但GQA在注意力头的粒度上造成了不匹配。一种暴力解决方案是为每个查询头独立加载token，但这会导致低效、重复的内存读取。Twilight通过在查询组的粒度上操作来解决此问题。具体来说，为给定查询组选择的token集合是该组内所有查询头识别出的token的并集【索引23，Native sparse attention: Hardware-aligned and natively trainable sparse attention+2025+arXiv】。

实现方式与权衡。如第4.2节所述，我们的top-p注意力机制原生支持头级动态性。然而，当与GQA集成时，这种头级动态性自然地转变为组级动态性，意味着同一组内的所有头共享一个共同的token预算。图13展示了我们采用扁平化分页KV缓存的注意力设计，它支持MHA的头级变长注意力（varlen attention）和GQA的组级变长注意力。这种设计是一个深思熟虑的权衡，平衡了实现效率与对现代注意力算法的兼容性。我们还在图13中比较了三种不同注意力实现的效率。


图 13: 左：Twilight中采用扁平化分页KV缓存的头级/组级变长注意力。右：在一个16k检索任务中，LLaMA-3.1-7B某层的真实预算分布上，三种注意力方法的比较。“Padded”指将所有头填充到最大预算长度；“Head Varlen”在头粒度加载KV，导致重复加载；“Group Varlen”在这两种方法之间取得了平衡。

C Longbench上的完整结果

请参阅表5。

表 5: Longbench上的完整结果。每个任务中的最高分（“Full”除外）已加粗显示。方法名称后显示了经过Twilight剪枝后的平均预算。我们还报告了将Twilight与每个基础算法集成时的相对误差变化（改善或下降）。

D 与Token丢弃方法的准确性比较

Token选择与Token丢弃的对比。如第2节所述，top-k稀疏注意力方法可大致分为两类：token丢弃和token选择。先前的研究【索引9，Quest: Query-aware sparsity for efficient long-context LLM inference+2024+arXiv】已证实，token选择通常优于token丢弃，因为后者不可避免地会造成不可逆的信息损失。为进一步验证这一观察，我们进行了Twilight与两种代表性的token丢弃方法——StreamingLLM【索引17，Efficient streaming language models with attention sinks+2023+arXiv】和SnapKV【索引18，SnapKV: LLM knows what you are looking for before generation+2024+arXiv】——的对比实验。如表6所示，DS-Twilight的性能显著优于这两种基线方法。

表 6: 在Longbench基准测试上，使用Longchat-7B-v1.5-32k模型对StreamingLLM、SnapKV和Twilight的比较。

E 在卸载场景下的效率评估

卸载场景下的性能增益。值得注意的是，在内存卸载场景中，每token的加载成本占主导地位，Twilight可以实现更显著的增益。这是因为Twilight以固定的估计成本减少了加载的token数量。表7显示，与Quest相比，Twilight可以实现高达16倍的加速。

表 7: 在卸载场景中，单个注意力算子的延迟（单位：微秒），其中KV缓存中相应的token从CPU内存加载。

F 局限性与未来工作

局限性分析。尽管Twilight有效地加速了现有的top-k稀疏注意力方法，但我们在图10中的分析揭示了不可忽略的估计开销。这使得Twilight在诸如大批量服务或卸载等场景中尤其具有优势，因为在这些场景中，从KV缓存加载token的成本占主导地位。第B.2节显示，头级动态性与GQA不友好，这给将Twilight与新模型架构集成带来了一些挑战。
未来工作展望。未来的研究可以专注于优化估计方法，以进一步提高端到端的延迟和吞吐量，以及如何将Twilight与其他模型架构（如多头潜在注意力（MLA）【索引64，Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model+2024+arXiv】）集成。