为扩散语言模型启用共享前缀的KV缓存

# 为扩散语言模型启用共享前缀的 KV 缓存
**来源:** https://arxiv.org/html/2606.07571

Younghun Go††感谢：贡献相等。 Jaehoon Han1Changyong Shin Chuk Yoo††感谢：通讯作者。 Gyeongsik Yang2
韩国大学，韩国

###### 摘要

键值（KV）缓存在共享前缀场景下对于高吞吐量的大语言模型（LLM）服务至关重要，但在新兴的扩散语言模型（DLM）中面临严峻挑战。在DLM中，双向注意力意味着更新任意一个token都会动态改变整个上下文及其对应的KV。因此，为LLM设计的现有缓存技术假设KV一旦计算即保持不变，会破坏共享前缀的KV。我们的实验表明，将这些技术应用于DLM会导致模型准确率几乎降为零。为了解锁DLM的高吞吐量服务，我们提出了双向前缀缓存BiCache，据我们所知，这是首个针对DLM中共享前缀的KV缓存技术。BiCache的设计基于我们全面分析中的关键观察：共享前缀的KV在浅层中保持稳定且可复用，而浅层的深度取决于每个请求中共享前缀token的比例。因此，BiCache动态地识别出可安全复用共享前缀KV的层深度，并消除冗余计算。评估表明，与现有技术相比，BiCache在不降低准确率（仅有0-1.8%的差异）的情况下，将服务吞吐量显著提高了36.3%–98.3%。

# 为扩散语言模型启用共享前缀的 KV 缓存

Younghun Go††感谢：贡献相等。 Jaehoon Han1Changyong Shin Chuk Yoo††感谢：通讯作者。 Gyeongsik Yang2
韩国大学，韩国

## 1 引言

键值（KV）缓存是高效大语言模型（LLM）服务的基本优化手段。在注意力层中，token隐藏状态被投影为键张量和值张量，存储这些张量使得服务系统能够在后续生成步骤中复用它们，而不是重新计算，从而改善服务延迟和吞吐量。KV缓存在共享前缀场景下尤为重要，即相同的token序列（例如系统提示）出现在多个请求的开头，其KV可以被缓存并在这些请求之间复用。在实际工作负载中，共享前缀占用了85-97%的提示token[37]。大多数现有LLM基于自回归模型（ARM），其顺序生成token[41]。在ARM中，先前生成token的KV一旦计算即保持不变。因此，现代服务系统采用“共享前缀缓存”，即存储共享前缀的KV一次，并在请求之间复用[38]，这使得共享前缀缓存成为高吞吐量服务的关键优化。

最近，扩散语言模型（DLM）作为ARM的一个有前景的替代方案出现[39]。ARM顺序生成token，导致延迟随输出长度线性增长[40]。此外，其因果注意力可能导致逆向诅咒，即模型可能正确回答从左到右的关系，但无法推断其逆关系[27, 26]。相比之下，DLM通过双向注意力在迭代去噪步骤中在任意位置生成token，允许所有位置在每个步骤中交互，从而缓解了ARM严格的从左到右生成瓶颈。然而，DLM的双向生成过程使得共享前缀缓存变得更加困难。在双向注意力下，更新一个token会改变所有其他token的KV，包括共享前缀中的token。因此，共享前缀KV在不同步骤或不同请求之间并非不变，这是与ARM的关键区别。直接复用ARM中的共享前缀KV可能会破坏注意力上下文，并严重降低模型准确率。我们的动机实验表明，将针对ARM设计的共享前缀缓存直接应用于DLM，会导致在多个基准测试上准确率几乎降为零（§3）。近期的研究[36, 33, 35]引入了针对DLM的KV缓存技术，但它们专注于在单个请求内复用KV，并未解决跨请求的共享前缀缓存。由于共享前缀在实际工作负载中占主导地位且对服务吞吐量至关重要，因此为DLM启用共享前缀缓存仍是一个重大的开放挑战。

为了应对这一挑战，我们提出了双向前缀缓存BiCache，据我们所知，这是首个针对DLM的共享前缀缓存技术。BiCache的核心思想是根据层深度和请求的共享前缀比例选择性复用KV。通过对共享前缀KV在请求、层和去噪步骤之间如何变化的深入分析，我们得到了三个重要观察（§4）。第一，共享前缀KV在浅层中跨请求保持高度相似，这使得在这些层中实现共享前缀缓存成为可能。第二，可以安全复用缓存KV的浅层深度与共享前缀比例（即每个请求中共享前缀token所占比例）密切相关。第三，虽然剩余的深层不能直接跨请求复用共享前缀KV，但它们可以通过周期性刷新技术在单个请求内减少冗余计算（§5）。基于这些观察，BiCache引入了两种实用机制：共享前缀分析和分层分区缓存。共享前缀分析确定到哪一层深度可以安全复用共享前缀KV而不损害准确率。在服务过程中，分层分区缓存在浅层直接复用缓存好的共享前缀KV，并在深层应用周期性刷新以消除冗余KV计算。

本研究的主要贡献如下：
-   • 刻画了DLM中共享前缀KV在请求、层和步骤之间的变化规律。
-   • 提出了BiCache，据我们所知，这是首个针对DLM的共享前缀缓存技术。
-   • 在准确率不降低的情况下（仅有0-1.8%的差异），实现了36.3%–98.3%的吞吐量显著提升。

## 2 背景

### 2.1 DLM LLaDA

LLaDA[39]是一个基于Transformer的DLM，也是后来许多DLM的基础。图1展示了其生成过程。它有三个关键超参数：Transformer层数LL、输出长度gg和去噪步数ss。在本文中，我们使用“步”指代一个去噪步骤。通过ss步，模型生成gg个token，每步预测⌊g/s⌋或⌊g/s⌋+1个token（例如，g=8，s=3给出每步{3,3,2}个token）。当新请求到达时，LLaDA通过在提示token后追加gg个[MASK] token来形成输入字符串（1）。在每个步骤中，输入字符串通过LL个Transformer层的前向传播进行处理（2）。在每个步骤中，DLM应用双向注意力并对指定数量的[MASK] token进行去噪，从而填入预测词（token）。在图1中，步骤1去噪三个token（如上所述），因此第一步去噪“explains”、“prefix”和“caching”（3）。包含三个去噪后token的输入字符串传递到下一步，在此步骤中，使用LL个Transformer层再次进行去噪。上一步的去噪token保持不变，模型去噪指定数量的掩码token（本例中为三个），生成“clearly”、“why”和“impossible”。此迭代过程持续s步，直到所有[MASK] token都被去噪。

图1：DLM中token生成的工作流程。

### 2.2 键值与双向注意力

在这里，我们详细描述DLM层。我们解释隐藏状态如何被投影为KV，以及双向注意力如何聚合它们以更新token表示。最后，我们解释这种双向特性如何导致KV随步骤变化。

**KV计算。** 在Transformer层l∈{1,...,L}和步t∈{1,...,s}，token的隐藏状态表示为hl,t∈Rd，其中dd是隐藏维度（每个token表示的特征大小）。该层使用其投影矩阵WlK∈Rdk×d和WlV∈Rdv×d计算键向量和值向量。这里，dk和dv分别表示键向量和值向量的维度；键向量（kl,t∈Rdk）用于计算注意力权重，值向量（vl,t∈Rdv）根据这些权重进行聚合。计算方式为：

kl,t = WlK hl,t，vl,t = WlV hl,t。   (1)

为简洁起见，全文将键值对表示为KVl,t ≜ (kl,t, vl,t)。

**双向注意力。** 使用键和值，层ll计算注意力输出ol,t以从输入字符串中收集上下文。具体来说，它使用查询向量ql,t（也通过WlQ从hl,t投影得到）与输入字符串中所有token的键进行比较，其中n表示输入字符串中的token数量：

ol,t = ∑_{j=1}^{n} softmax_j( (ql,t · kl,t^(j)) / √dk ) vl,t^(j)。   (2)

此处，ol,t是值向量vl,t^(j)的加权和，代表聚合后的全局上下文。权重由经过softmax归一化的缩放点积给出，衡量第jj个token在当前步tt与当前token的相关性。为了生成每个token的ol,t，查询会关注所有nn个token。这个输出ol,t会更新下一层的隐藏状态hl+1,t。经过LL层后，最终的隐藏状态被转换为词汇logits，以预测最可能的单词（例如图1中的“explains”）。

重要的是，在双向注意力下，在一步中解掩一个token会改变其他token的ol,t。这会传播到它们的隐藏状态，并通过公式(1)修改其键和值向量。例如，在图1中，即使“explains”这个token已经被解掩，其KV在后续步骤中仍会随着周围token的变化而持续更新。

### 2.3 共享前缀及其缓存

共享前缀在实践中很常见。例如，许多生产级LLM服务系统在每个用户请求前附加一个固定的系统提示，以提供格式指南、安全和伦理约束以及角色说明[18]。真实世界共享前缀的示例见附录A。此外，在多轮对话中，共享前缀自然产生。当用户在同一个会话中提交新请求时，模型的输入字符串通常包含之前的用户请求和响应作为共享前缀，以便模型能够维护对话上下文[19]。因此，提示token中有很大一部分在请求之间重复，在实际部署中最高可达提示token总长度的97%[37]。

共享前缀缓存通过以下方式避免冗余KV计算：1） 一次性计算共享前缀的KV，然后2） 为后续输入包含相同前缀的请求复用这些KV。在ARM[38]中，一旦共享前缀的KV被缓存，它们可以直接用于其他请求的所有层和所有步骤。这是因为在ARM的因果注意力下，即使添加了额外的后缀token，前缀的KV值也保持不变。先前的工作报告称，ARM中的共享前缀缓存可以提供约10倍的延迟改善，使其成为实际LLM服务系统的关键优化[20]。

## 3 动机：DLM需要新的共享前缀缓存技术

在本节中，我们将为ARM设计的现有共享前缀缓存应用于DLM，并展示它会导致高准确率下降，这证明需要新技术。

### 3.1 实验设置

所有实验均在配备180 GB内存的NVIDIA B200 GPU上进行。我们评估具有32个Transformer层的LLaDA[39]。设置g=256，s=128；因此每一步解掩两个token。我们使用batch size 1以将模型放入GPU内存，这与其它研究中的配置一致[6, 5]。LLaDA在广泛使用的基准测试上进行评估：1) ARC-Challenge-Chat（科学选择题问答），2) GPQA（研究生水平问答），以及3) MATH-500（竞赛风格数学）。类似现有的LLM服务系统和研究[17, 16]，我们在每个请求前附加两个系统提示。首先，我们附加xAI的Grok-4模型[25]的真实系统提示（详情见附录A），提供角色和安全指令。其次，我们附加另一个系统提示，强制要求特定的输出格式（即在最终答案前附加##），以满足基准测试中的严格答案匹配要求（如下所述）。跨基准测试，共享前缀平均占请求的85%（详情见附录§D.5），这与实际工作负载一致[37]。我们评估两个指标：1) 吞吐量和2) 准确率。吞吐量以每秒生成的token数（tokens/s）衡量，