扩散大语言模型的词缀缓存

# 扩散大语言模型的词缀缓存
来源：https://arxiv.org/html/2608.26140
Kaihua Liang1 An Zhong1 Xin Tan2 Zafar Ayyub Qazi1,3 Hong Xu2 Jian Weng1 Marco Canini1
1 KAUST 2 CUHK 3 LUMS
\{kaihua\.liang, marco\}@kaust\.edu\.sa
![未命名图片](https://arxiv.org/html/2608.26140v1/x1.png)
https://github.com/sands-lab/ACache
###### 摘要
扩散大语言模型（DLLMs）实现了非自回归解码和双向上下文建模，但高效推理仍然具有挑战性。与自回归系统的键值（KV）缓存可以重用于共享前缀不同，DLLMs通过双向注意力将共享上下文词元的键值状态与不断演化的生成词元耦合在一起，这使得朴素的缓存重用会过时，而完整的重新计算又代价高昂。我们提出了ACache，一种面向词缀的缓存重用机制，用于DLLMs中共享文本片段（超越前缀）。ACache通过测量词缀词元对掩码生成词元的影响力，识别出一小部分请求特定的关键词缀子集，称为锚点词元，并选择性地仅重新计算这些词元的键值状态，同时重用其余的词缀缓存。ACache构建在Fast-dLLM之上，在不同设置下，仅重新计算约20%的词缀词元即可恢复因直接重用词缀缓存而导致的精度损失。我们还在Nano-vLLM引擎之上构建了一个共享前缀原型，表明ACache将重新计算延迟降低了高达55.7%，并将端到端吞吐量提高了高达1.68倍。

**扩散大语言模型的词缀缓存**
Kaihua Liang1 An Zhong1 Xin Tan2 Zafar Ayyub Qazi1,3 Hong Xu2 Jian Weng1 Marco Canini1
1 KAUST 2 CUHK 3 LUMS
\{kaihua\.liang, marco\}@kaust\.edu\.sa
![未命名图片](https://arxiv.org/html/2608.26140v1/x2.png)
https://github.com/sands-lab/ACache

## 1 引言
大语言模型（LLMs）已迅速成为通用智能系统的基础，推动了诸如与用户意图对齐的对话助手（Ouyanget al., 2022）、代码生成与编程副驾驶（Chenet al., 2021），以及能够推理、与环境交互并调用外部API或工具的智能体系统（Yaoet al., 2023；Patilet al., 2024）等应用的进展。最近，扩散大语言模型（DLLMs）（Nieet al., 2025；Yeet al., 2025）作为传统自回归（AR）模型的有前景的替代方案出现，利用双向依赖性进行文本建模。与AR模型严格顺序生成词元不同，DLLMs执行非自回归词元预测的迭代去噪，实现低延迟并行解码（Israelet al., 2025；Wuet al., 2026a）和更丰富的双向上下文建模（Niet al., 2025；Heet al., 2026）。对于AR模型，前缀缓存已成为高效推理的成熟优化技术，可重用系统提示和提示模板等共享前缀的键值（KV）缓存（Gimet al., 2024；Zhenget al., 2024）。相比之下，DLLMs的全局双向注意力将共享前缀词元的键值状态与序列中动态演化的生成词元耦合在一起（Maet al., 2025；Liuet al., 2025）。如图1（左和中）所示，AR推理支持直接前缀重用，这在DLLMs的双向注意力下会变得过时。因此，大多数实用的DLLM推理框架，如dLLM-Cache（Liuet al., 2025）、dKV-Cache（Maet al., 2025）和Fast-dLLM（Wuet al., 2026a），周期性地重新计算键值状态以保持上下文一致性，而朴素地重用公共前缀缓存会引入过时的上下文并显著降低生成质量。

**图1：** AR LLM和DLLMs中的缓存重用。**左：** 因果AR推理支持直接重用共享前缀。**中：** 在DLLMs中，双向注意力将共享上下文的键值状态与演化的生成词元耦合，阻止了朴素重用。**右：** ACache重用共享词缀，同时选择性重新计算锚点词元。

为了解决这一局限性，我们提出了ACache（词缀缓存），一种专门为DLLMs设计的细粒度缓存重用机制。我们研究词缀（即超越前缀的共享连续文本片段），因为即使这些片段出现在上下文的开头、中间或结尾，双向注意力也能实现键值缓存重用，而传统的AR键值缓存重用主要以词缀为中心。如图1（右）所示，ACache在重用共享词缀的同时，仅选择性重新计算一小部分关键的词元。我们将这些关键的词元（Zhanget al., 2023；Xiaoet al., 2024）称为锚点词元。ACache通过测量掩码生成词元对词缀词元的跨注意力重要性来识别锚点词元，并仅重新计算这些词元的键值状态，同时重用其余的词缀缓存。我们在最先进的Fast-dLLM框架上实现了ACache，用选择性的锚点词元重新计算替代其周期性的全缓存重新计算。在多个基准测试中，ACache仅重新计算一小部分词缀词元即可恢复因直接重用词缀缓存而损失的精度。更广泛地说，ACache表明词缀重用是DLLMs中双向注意力所开启的一种系统优化机会，激发了超越传统词缀中心推理（Srivatsaet al., 2025；Panet al., 2025；Yuanet al., 2026）的缓存重用设计。我们还在Nano-vLLM（一个轻量级推理引擎）上构建了一个共享前缀的ACache原型，以验证其在具有分页内存管理和连续批处理（Yuet al., 2022）的现实推理堆栈中的计算和内存效率收益。

总而言之，我们的贡献如下：
- • 我们提出了ACache，一种面向DLLMs的细粒度词缀缓存机制，针对标准前缀之外的共享文本片段，并引入了锚点词元，以实现仅对一小部分请求特定的词缀词元进行选择性键值重计算。
- • 我们在Fast-dLLM上实现了ACache，并在多个基准上进行了评估，结果表明在重新计算约20%的词缀词元时，ACache在不同设置下恢复了因直接重用词缀缓存而导致的精度损失。
- • 我们在Nano-vLLM引擎之上构建了一个共享前缀的系统原型，并表明在具有分页内存管理和连续批处理的现实推理堆栈下，ACache将重新计算延迟降低了高达55.7%，并将端到端吞吐量提高了高达1.68倍。

## 2 背景
### 2.1 基于扩散的语言模型
基于扩散的语言模型通过迭代地对部分掩码序列进行去噪来生成文本，使预测能够利用双向上下文，而不仅仅是左侧前缀。这种非因果公式出现在连续嵌入模型如Diffusion-LM（Liet al., 2022）和离散状态模型如D3PM（Austinet al., 2021a）中。最近的离散扩散模型加强了这一方向：SEDD（Louet al., 2024）引入了离散数据分数熵目标，MDLM（Sahooet al., 2024）表明掩码吸收态扩散与现代训练方法结合，可以在保持并行去噪和填充能力的同时，接近自回归困惑度。这些思想现已扩展到十亿参数模型。LLaDA（Nieet al., 2025）表明扩散模型可以遵循标准的预训练和微调流程，并在通用、数学和编码任务上保持竞争力，而Dream（Yeet al., 2025）引入了基于AR的初始化和自适应噪声重调度，并强调了DLLM的规划和灵活生成能力。在解码过程中，DLLMs通过不同调度策略（Nieet al., 2025；Ben-Hamuet al., 2025；Wuet al., 2026a）迭代更新掩码响应，进行并行预测和部分承诺。这种全序列去噪使得任意顺序的解码和填充成为可能，但也造成了步骤间的显著冗余，使得跨步骤的缓存重用成为一个自然的效率优化目标。

### 2.2 DLLMs中的缓存技术
与自回归LLMs不同，DLLMs不能直接采用标准AR键值缓存机制用于因果解码。因为每个解码步骤都使用全局双向注意力，词元表示（包括提示词元）依赖于演化的掩码响应词元，所以键值缓存不是仅追加的，并且很快会过时。因此，实用的DLLM推理仍然受限于每个解码步骤重新计算全序列键值缓存的代价。因此，最近的工作探索了近似的缓存策略。dLLM-Cache（Liuet al., 2025）使重用策略更加不对称：它以长间隔重新计算提示特征，更频繁地更新响应特征，并使用特征相似性来仅自适应地重新计算一部分响应词元。dKV-Cache（Maet al., 2025）专注于已解码的词元，重用它们跨解码步骤的键值状态，并在固定大小的块内延迟重新计算，以减少对已稳定词元的重复计算。Fast-dLLM（Wuet al., 2026a）利用块级时间局部性，在解码块内重用键值状态，并周期性地重新计算全缓存，同时将这种模式与并行解码相结合。然而，这些近似技术仍然需要周期性地重新计算全缓存，并且无法为DLLMs中的共享输入启用跨请求的缓存重用。据我们所知，ACache是第一个明确解决这一局限性的工作。

## 3 ACache设计
### 3.1 预备知识
我们形式化DLLM的推理循环。给定长度为p的提示x₁:p和n个词元的响应区域，令T表示总解码步数，令z⁽ᵗ⁾表示第t步的完整序列，令Mₜ表示当前掩码位置的集合。生成从完全掩码的响应开始：z⁽ᵀ⁾ = [x₁:p, [MASK]ⁿ] (1)
在每个解码步骤t，模型并行预测所有当前掩码位置：ŷᵢ⁽ᵗ⁾ ~ pθ(· | z⁽ᵗ⁾), i ∈ Mₜ (2)
仅承诺一部分置信的预测。令Cₜ ⊆ Mₜ表示在步骤t选择承诺的位置。下一状态更新为：zᵢ⁽ᵗ⁻¹⁾ = { ŷᵢ⁽ᵗ⁾, i ∈ Cₜ; [MASK], i ∈ Mₜ \ Cₜ } (3)
不同的解码调度定义Cₜ的方式不同，但它们都遵循这种迭代的“预测然后承诺”模式。承诺通常由预测词元的置信度阈值控制（Wuet al., 2026a；Xiaoet al., 2026；Luet al., 2026；Luoet al., 2026a, b）。虽然这个基础循环并行预测所有掩码位置，但DLLMs的实际推理方案通常施加更有结构的块级调度（Nieet al., 2025；Wuet al., 2026a；Luoet al., 2026a）。ACache在此公共推理循环之上运行。

现在考虑一个共享词缀，它占据z⁽ᵗ⁾中一组连续的非掩码位置F = {l, l+1, ..., r}，长度L = |F|。在代表性的DLLM推理框架中（Maet al., 2025；Liuet al., 2025；Wuet al., 2026a），键值缓存在相邻步骤之间重用，并周期性地重新计算以保持上下文一致性。在每个这样的重新计算点，F中的所有词缀词元都会被重新计算，尽管底层的词缀文本在请求之间是共享的。ACache旨在通过识别一小部分关键的词缀词元子集来解决这种低效性，这些词元的键值状态必须在当前请求特定上下文中更新。

**图2：** 关于锚点词元重要性的经验观察。在第5节设置下，在三个随机选择的GSM8K样本（num\_fewshot=2）上，掩码到词缀的注意力重要性集中在一小部分词缀词元上，这些词元的位置在解码步骤之间保持稳定。

### 3.2 锚点选择
ACache仅在每个请求开始时（在解码循环开始之前）选择一次锚点词元，然后在该请求的整个过程中重用相同的锚点集合。为了使此选择依赖于请求同时保留词缀重用，我们遵循MaskKV（Huanget al., 2025）的见解。具体而言，ACache在当前请求的非词缀位置上运行一次性探测，同时使用预计算的词缀缓存作为固定的过去键值状态；在选择过程中，词缀状态不会被转发或更新。然后，它通过掩码生成词元给予每个词缀词元的注意力来衡量其重要性。令Qmask⁽ⁿʰ⁾表示第n层第h个头的掩码位置的查询，令Kaffix⁽ⁿʰ⁾表示预计算的词缀词元的键。对于词缀词元j∈F和掩码位置k∈Mₜ，我们用a_{j,k}⁽ⁿʰ⁾表示从掩码查询位置k到词缀键位置j的注意力权重：a_{j,k}⁽ⁿʰ⁾ = [Softmax(Qmask⁽ⁿʰ⁾(Kaffix⁽ⁿʰ⁾)ᵀ / √dₕ)]_{k,j}。