UniPrefill：基于块级动态稀疏化的通用长上下文预填充加速

# UniPrefill：通过分块动态稀疏化实现通用长上下文预填充加速

来源：https://arxiv.org/html/2605.06221  
报告编号：001  
GitHub：https://github.com/qhfan/UniPrefill.git  

**Qihang Fan**<sup>1,2,3,*</sup>, **Huaibo Huang**<sup>1,2,†</sup>, **Zhiying Wu**<sup>3</sup>, **Bingning Wang**<sup>3,‡</sup>, **Ran He**<sup>1,2</sup>  

<sup>1</sup> MAIS&NLPR, CASIA  
<sup>2</sup> UCAS  
<sup>3</sup> WeChat, Tencent  

###### 摘要

随着大型语言模型（LLMs）的快速进步，它们的能力不断增强，同时对上下文长度的需求也越来越长。为了提高长上下文处理的推理效率，近期提出了几种新颖的低复杂度混合架构，有效缓解了长上下文推理的计算负担。然而，现有的长上下文预填充（prefill）加速研究主要集中于稀疏注意力机制，这些机制仅在纯注意力模型上能达到最大加速效果。当转移到新兴架构——如线性/全注意力混合架构或滑动窗口/全注意力混合架构——时，这些预填充加速方法会出现显著的性能下降。此外，这类方法通常与连续批处理（continuous batching）不兼容，难以集成到 vLLM 等现代推理引擎中。为此，我们提出了 **UniPrefill**，这是一种适用于几乎所有模型架构的预填充加速框架，它在 token 级别直接加速模型计算。我们进一步将 UniPrefill 实现为连续批处理算子，并扩展了 vLLM 的调度策略以原生支持 UniPrefill 的预填充-解码协同处理和张量并行，使其能够无缝集成到 vLLM 中。UniPrefill 实现了高达 **2.1 倍**的首字延迟（TTFT）加速，且随着并发请求数量的增加，加速效果愈发显著。

††† 通讯作者。  
‡ 项目负责人。  
* 在微信实习期间的研究工作。

## 1 引言

大型语言模型（LLMs）的快速进步推动了其在日益多样化的现实应用中的部署，从文档理解、代码生成到多轮对话和检索增强生成 \[llama, llama2, qwen2.5-1m, qwen25technicalreport, qwen3technicalreport, qwentechnicalreport, glm2024chatglm\]。随着能力的扩展，LLM 预期处理的上下文长度也大幅增长——现代部署通常涉及数万 token 的序列，而对十万甚至百万 token 上下文的需求正变得司空见惯。这一趋势对推理效率造成了巨大压力，因为标准的 Softmax Self-Attention \[attention\] 机制随序列长度呈二次方扩展，在处理长上下文时会带来 prohibitive（令人望而却步）的计算成本。

为了解决二次复杂度瓶颈，新一代混合架构应运而生，将计算高效的层与全注意力层交错排列。两类代表性架构尤其受到关注：线性/全注意力混合架构，用线性递归机制替换部分注意力层 \[mamba, mamba2, yang2024gla, fan2024rect, fan2024breaking\]，将每层复杂度从 $O(N^2)$ 降至 $O(N)$；以及滑动窗口/全注意力混合架构，将大多数注意力层限制在固定的局部上下文窗口内，同时保留少量全局全注意力层以捕获长程依赖 \[gemmateam2025gemma3technicalreport, jiang2023mistral7b\]。这些混合设计大幅降低了长上下文推理的理论复杂度，并已被近期发布的众多生产级模型广泛采用。

![Figure 1: Prefill 吞吐量对比](https://arxiv.org/html/2605.06221#S1.F1)  
*图 1 标题：标准 Prefill 与 UniPrefill 在三种模型架构及不同批量大小下的预填充吞吐量对比（张量并行大小设为 8）。所有实验均在 vLLM 中进行，UniPrefill 深度集成于 vLLM 的连续批处理调度器。我们在 LLaMA-3.1-Instruct-8B \[llama3\]（全注意力）、Qwen3-Next-80B-A3B \[qwen3next_blog_2025\]（线性/全注意力混合）和 Gemma-3-12B \[gemmateam2025gemma3technicalreport\]（滑动窗口/全注意力混合）上评估了预填充吞吐量（K tokens/s），上下文长度从 4K 到 128K，批量大小为 1、4、16 和 64。实心条表示标准 Prefill，斜纹条表示 UniPrefill。UniPrefill 在所有三种架构上均实现了更高的吞吐量，且在更长的上下文长度和更大的批量大小下增益更为明显。*

尽管混合架构日益普及，但研究社区在预填充加速方面的努力仍高度集中在稀疏注意力 \[minference, mobamixtureblockattention, fan2026flashprefill\] 上。以 MInference \[minference\] 为代表的研究展示了令人印象深刻的预填充加速效果，在纯全注意力设置下对长序列实现了高达 10 倍的加速。然而，这种对稀疏注意力的关注带来了一个根本性限制：加速紧密耦合于全注意力操作本身。在全注意力仅占所有层一小部分的混合架构中，仅加速这些注意力层的边际效益大幅降低。例如，在比例为 3:1 的线性/全注意力混合模型中，现有稀疏注意力方法最多只能加速每四层中的一层，而主要的计算预算完全未被触及。这种架构上的不匹配使得现有的预填充加速方法在面对新一代混合模型时效果大打折扣。

现有预填充加速方法的第二个同样关键的局限性是与连续批处理的不兼容，而连续批处理是支撑 vLLM \[vLLM, zheng2024sglang\] 等现代高吞吐量推理引擎的调度范式。像 FlexPrefill \[flexprefill\] 这样的方法孤立地处理单个请求，并假设静态的批量组成，因此根本难以集成到请求动态进出的连续批处理调度器中。结果，这些方法大多停留在研究原型阶段，未能成功嵌入生产推理系统。

为了克服这两大局限，我们提出了 **UniPrefill**，这是一种通过利用以下关键洞察来实现架构无关加速的预填充加速框架：可以在全注意力层估计 token 的重要性，并将其传播到所有后续层。具体而言，UniPrefill 在每个全注意力层应用轻量级的分块评分标准，以识别并丢弃计算冗余的 token。一旦某个 token 被丢弃，它将从该块剩余层的所有下游计算中排除。这种级联效应意味着，在注意力层做出的单一 token 丢弃决策，会转化为整个层堆栈（而不仅仅是注意力子层）中计算量的成比例减少。因此，UniPrefill 同时显著减少了注意力 FLOPs 和 GEMM FLOPs，无论模型是纯全注意力 Transformer 还是混合架构，它都同样有效。

在系统设计之外，我们通过将 UniPrefill 实现为连续批处理算子 \[yu2022orca\] 并扩展 vLLM \[vLLM\] 的调度器以原生支持 UniPrefill token 丢弃机制下的预填充-解码协同处理，解决了系统集成挑战。这种紧密集成允许 UniPrefill 作为生产推理引擎中的透明加速层运行，无需更改模型权重或服务基础设施。

我们在 RULER \[hsieh2024ruler\] 基准上对多种模型架构评估了 UniPrefill。结果表明，UniPrefill 在几乎无精度损失的情况下，实现了高达 **2.1 倍** 的首字延迟（TTFT）加速，如图 1 所示。值得注意的是，加速效果随着并发请求数量的增加而有利扩展（见图 1），这使得 UniPrefill 特别适用于预填充成本为主要瓶颈的高并发生产服务场景。

我们的主要贡献总结如下：

*   我们提出了 **UniPrefill**，一种 token 级别的预填充加速框架，它在全注意力层丢弃 token 并将稀疏性传播到所有后续层，同时减少注意力和 GEMM FLOPs，从而在异构混合架构上实现一致的加速。
*   我们将 UniPrefill 实现为连续批处理算子，并通过支持预填充-解码协同处理和张量并行的扩展调度策略将其集成到 vLLM \[vLLM\] 中，实现了无缝的生产就绪部署。
*   在长上下文基准 RULER 上的大量实验表明，UniPrefill 实现了高达 **2.1 倍** 的 TTFT 加速，精度损失可忽略不计，且加速增益随请求并发度扩展。

## 2 相关工作

#### 混合 LLM 架构
为了克服 Softmax 注意力的二次复杂度，大量研究提出了高效的序列建模替代方案，包括状态空间模型、线性注意力变体和递归架构 \[mamba, mamba2, sun2023retentivenetworksuccessortransformer, yang2024gla, yang2024deltanet, fan2025sec, fan2024rect, minimax01scalingfoundationmodels, yang2024gdn, zhang2025kda\]。为了平衡效率和表达能力，混合架构应运而生，将全注意力与这些高效替代方案交错排列 \[qwen3next_blog_2025, lenz2025jamba, gemmateam2025gemma3technicalreport, xiao2026mimov2flash, jiang2023mistral7b\]，并被近期发布的众多生产模型广泛采用。然而，现有的预填充加速方法主要针对纯全注意力架构定制，限制了它们在这类新一代模型上的有效性。

#### 用于预填充加速的稀疏注意力
利用注意力得分矩阵中的固有稀疏性是加速预填充阶段的既定策略。一批研究识别出静态或动态稀疏模式——如垂直、斜向和分块稀疏结构——并跳过相应的注意力计算 \[minference, native-sparse-attention, mobamixtureblockattention, optimizingmixtureblockattention, flexprefill, chen2026vsprefill\]。这些方法在全注意力模型上展示了显著的加速效果 \[minference, flexprefill, xattention, wang2025proxyattn\]。然而，它们共有两个根本性限制：其加速紧密耦合于注意力操作本身，留给 FFN 和 GEMM 计算完全未加速；且它们通常与连续批处理 \[yu2022orca\] 不兼容，使得集成到 vLLM \[vLLM\] 等生产推理引擎变得非平凡。UniPrefill 通过在 token 级别操作并将稀疏性传播到所有层，解决了这两个限制。

## 3 方法

在本节中，我们介绍 UniPrefill，一种架构无关的预填充加速框架。整体流程如图 2 所示。

### 3.1 预备知识

考虑由包含 $B$ 个块（block）的混合 LLM 处理的输入序列 $\mathbf{x}=[x_1,\ldots,x_N]$。每个块 $b$ 包含一个全注意力层，后跟 $M_b$ 个子层（线性注意力、滑动窗口注意力、FFN 等）。令 $\mathbf{H}^{(b,0)} \in \mathbb{R}^{N \times d}$ 表示块输入。预填充的目标是计算用于下一个 token 预测的最终隐藏状态 $\mathbf{h}_N^{(L)}$：

$$
P(x_{N+1} \mid x_{1:N}) = \text{LMHead}\!\left(\mathbf{h}_N^{(L)}\right) \tag{1}
$$

标准预填充在每个全注意力层产生 $O(N^2 d_k)$ 的计算量，在每个 GEMM 子层产生 $O(Nd^2)$ 的计算量，每块总计 $O(N^2 d_k + M_b N d^2)$。

### 3.2 Token 重要性估计

由于下一个 token 预测仅依赖于 $\mathbf{h}_N^{(L)}$，token $i$ 对块 $b$ 中最终隐藏状态的贡献为：

$$
\mathbf{h}_N^{(b,1)} = \sum_{i=1}^N \mathbf{A}^{(b)}_{N,i} \cdot \mathbf{v}_i^{(b)} + \mathbf{h}_N^{(b,0)}, \tag{2}
$$

其中 $\mathbf{A}^{(b)}_{N,i} = \operatorname{softmax}_i\!\left(\mathbf{q}_N^{(b)} \{\mathbf{K}^{(b)}\}^\top / \sqrt{d_k}\right)$ 是全序列注意力权重。当 $\mathbf{A}^{(b)}_{N,i} \approx 0$ 时，token $i$ 对下一个 token 预测的贡献可忽略不计。为了减少估计方差，我们聚合最后 $n$ 个查询位置，而不是单个位置：

$$
s_i^{(b)} = \frac{1}{n} \sum_{j=N-n+1}^N \mathbf{A}^{(b)}_{j,i}, \tag{3}
$$

这需要 $n \times N$ 的注意力计算，成本为 $O(n N d_k)$，对于 $n \ll N$ 来说可以忽略不计。

在实践中，重要性估计和 token 选择在*块粒度*上操作。我们将输入序列划分为大小为 $G$ 的非重叠块：$\mathcal{B}_g = \{(g-1)G+1, \ldots, \min(gG, N)\}$，其中 $g=1, \ldots, \lceil N/G \rceil$。为了效率，首先计算部分 GEMM $\mathbf{S} = \mathbf{Q}_{[N-n:N]} \mathbf{K}^\top \in \mathbb{R}^{n \times N}$；然后*跨越完整序列维度*应用在线 softmax 以获得适当归一化的注意力权重，随后在每个块内缩减分数：

$$
\bar{s}_g^{(b)} = \frac{1}{G} \sum_{i \in \mathcal{B}_g} \frac{1}{n} \sum_{j=N-n+1}^N \mathbf{A}^{(b)}_{j,i}, \tag{4}
$$

其中 softmax 归一化在块缩减之前在完整的键序列上执行，确保 $\bar{s}_g^{(b)}$ 反映块 $g$ 捕获的真实注意力质量。这将选择决策的数量从 $N$ 减少到 $\lceil N/G \rceil$，同时保持重要性估计的准确性。

#### 与 SnapKV 的关系
我们的重要性估计与 SnapKV \[li2024snapkv\] 在表面相似，后者也使用观察窗口来识别重要 token。然而，这两种方法在目标和范围上存在根本差异。SnapKV 在所有层完成完整的 $N \times N$ 预填充后，应用其选择来压缩解码阶段的 KV 缓存——预填充 FLOPs 完全不受影响。UniPrefill 在预填充*期间*应用选择，将丢弃决策向前传播到所有后续层。形式上，SnapKV 每层最多节省 $O(N \cdot d_{kv})$ 解码时内存，而 UniPrefill 每块节省 $(1-\rho^{(b)}) \cdot M_b \cdot O(Nd^2)$ 预填充时 FLOPs，其中 $\rho^{(b)} = |\mathcal{S}^{(b)}|/N$ 是 token 保留率——这是一个随 $M_b$ 线性增长且在 SnapKV 中完全不存在的量。

### 3.3 前 p 百分比 Token 选择

令 $\pi$ 为按降序排列块级分数 $\{\bar{s}_g^{(b)}\}$ 的排列。我们保留最小的块集合：

$$
\mathcal{S}^{(b)} = \{\pi(1), \ldots, \pi(k^*)\}, \quad k^* = \min_k s.t. \sum_{j=1}^k \bar{s}_{\pi(j)}^{(b)} \geq \sum_g \bar{s}_g^{(b)} \cdot p
$$