LaCache: 扩散大语言模型的精确缓存与精度自适应推理

# LaCache：扩散大规模语言模型的精确缓存与精度自适应推理
来源：https://arxiv.org/html/2607.16339

陈行如¹，梁泽朗¹††，马永嘉¹††，詹吉庆¹，杨淑玲¹，温联¹，战坤¹††

¹理想汽车股份有限公司。通讯作者。

######  摘要

基于扩散的大规模语言模型（DLLM）通过文本生成中的半自回归（SAR）解码实现并行生成。然而，现有方法存在严重的算子级冗余：它们在去噪步骤中对整个序列进行重新计算，忽略了前缀和掩码后缀在一个块内保持不变的事实。我们提出LaCache，一种无需训练的加速框架，通过无损缓存和混合精度来缓解这种冗余。具体而言，LaCache采用**无损状态记忆化（LSM）**，通过缓存三种类型的中间结果：（i）**EmbedCache**用于嵌入输出，（ii）**RoPECache**用于逐令牌预注意力状态，（iii）**FACache**用于FlashAttention中的在线softmax统计量。这些缓存使模型能够跳过对不变令牌的冗余计算，同时不改变输出。为了进一步缓解内存带宽瓶颈，LaCache集成了针对FFN层的逐组FP8量化策略，该策略适应扩散过程中随步骤变化的激活分布。实验表明，单独的LaCache相比原始DLLM实现了约1.3×端到端加速。当与现有加速方法结合时，LaCache达到高达40.2×端到端加速，同时保持相当的任务准确度。

LaCache：扩散大规模语言模型的无损缓存与精度自适应推理
陈行如¹††，梁泽朗¹††，马永嘉¹††，詹吉庆¹，杨淑玲¹，温联¹，战坤¹††
¹理想汽车股份有限公司。††主要联系[email protected]

## 1 引言

大规模语言模型（LLMs）(Yang et al., 2025a (https://arxiv.org/html/2607.16339#bib.bib36); Dubey et al., 2024 (https://arxiv.org/html/2607.16339#bib.bib37); Brown et al., 2020 (https://arxiv.org/html/2607.16339#bib.bib19); Ma et al., 2026 (https://arxiv.org/html/2607.16339#bib.bib46))广泛应用于聊天(Mao et al., 2024 (https://arxiv.org/html/2607.16339#bib.bib40))、代码生成(Jiang et al., 2024 (https://arxiv.org/html/2607.16339#bib.bib38))和复杂推理(Wan et al., 2024 (https://arxiv.org/html/2607.16339#bib.bib39); Liu et al., 2025a (https://arxiv.org/html/2607.16339#bib.bib26); Feng et al., 2026 (https://arxiv.org/html/2607.16339#bib.bib52))。在部署中，推理延迟和计算成本通常是主要瓶颈(Lou et al., 2023 (https://arxiv.org/html/2607.16339#bib.bib32))。自回归解码逐个生成令牌，限制了效率和硬件并行性(Dong et al., 2025 (https://arxiv.org/html/2607.16339#bib.bib29); Feng et al., 2025 (https://arxiv.org/html/2607.16339#bib.bib44); Yang et al., 2025c (https://arxiv.org/html/2607.16339#bib.bib45))。这促使了并行解码和非自回归生成的发展。

参见图注
图1：**左图：** DLLMs将总生成令牌划分为连续块，并在每个块内的多个步骤中进行并行解码。在一个块生成期间，块外的令牌保持不变，导致重复计算。**右图：** LLaDA-Instruct中各模块的延迟主要集中在Transformer层和最终采样器。使用我们的方法LaCache后，各模块的延迟显著降低。

扩散大规模语言模型（DLLMs）(Nie et al., 2025 (https://arxiv.org/html/2607.16339#bib.bib1); Ye et al., 2025 (https://arxiv.org/html/2607.16339#bib.bib3); Lyu et al., 2019 (https://arxiv.org/html/2607.16339#bib.bib33); Yang et al., 2025b (https://arxiv.org/html/2607.16339#bib.bib51); Ma et al., 2025 (https://arxiv.org/html/2607.16339#bib.bib43); Khanna et al., 2025 (https://arxiv.org/html/2607.16339#bib.bib21))将文本生成建模为迭代去噪，使得多个令牌位置可以在一步中更新，并采用双向注意力以利用全部上下文。尽管有这种潜力，开源DLLMs存在显著的计算开销(Nie et al., 2025 (https://arxiv.org/html/2607.16339#bib.bib1); Luo et al., 2024 (https://arxiv.org/html/2607.16339#bib.bib49))。我们将这种低效率归因于两个主要限制：首先，双向注意力阻碍了标准KV缓存机制的直接重用。其次，并行解码可能因违反令牌依赖性而降低生成质量(Wu et al., 2025 (https://arxiv.org/html/2607.16339#bib.bib4); Xie et al., 2025 (https://arxiv.org/html/2607.16339#bib.bib50); Ma et al., 2024 (https://arxiv.org/html/2607.16339#bib.bib48))。

参见图注
图2：一次推理步骤的完整过程。在第一步更新缓存后，后续步骤中输入只需包含当前块的令牌，从而节省了Transformer第一层的大部分计算工作，包括嵌入层、归一化、QKV线性变换、旋转位置编码和FlashAttention。此外，混合精度加速方法应用于除第一层外的所有层：gateup线性层、ffout线性层。FP8量化操作与ff Norm和SwiGLU融合，进一步节省了开销。此外，Transformer层的输出将被再次截断，仅保留当前块的令牌用于后续采样操作，从而降低float Softmax操作的延迟。

大多数开源DLLMs采用半自回归（SAR）(Nie et al., 2024 (https://arxiv.org/html/2607.16339#bib.bib41))、块状推理方案。模型用一个固定数量的[MASK]令牌初始化未来位置，并将其划分为块。然后在每个块内执行多次去噪步骤。在这些步骤中，只有当前块令牌发生变化，而提示和其他块保持相同的令牌ID。尽管如此，模型在每一步都执行完整的前向传播，导致大量的冗余计算。如图1（左）所示，一个块内的高输入重叠导致重复计算(Wei et al., 2025b (https://arxiv.org/html/2607.16339#bib.bib28))。这出现在逐令牌算子中，如嵌入、QKV投影和RoPE。也出现在注意力计算内部，特别是第一层FlashAttention的分块计算中。图1（右）的性能分析结果显示，延迟主要由Transformer层和最终采样器主导。这一观察表明，消除对不变令牌的重复算子级计算可以带来稳定且可组合的端到端加速。

现有的推理加速技术主要通过减少去噪步骤数或有效上下文长度来工作。Fast dLLM(Wu et al., 2025 (https://arxiv.org/html/2607.16339#bib.bib4))采用置信度阈值解码和近似缓存来提高吞吐量。DPad(Chen et al., 2025 (https://arxiv.org/html/2607.16339#bib.bib5))应用后缀丢弃，包括滑动窗口和距离衰减，以减少注意力中使用的令牌数量。其他工作探索指导令牌或自适应解码调度(Wei et al., 2025a (https://arxiv.org/html/2607.16339#bib.bib7); Sun et al., 2025 (https://arxiv.org/html/2607.16339#bib.bib6))。虽然这些方法有效且与我们的方法互补，但它们并没有直接消除SAR块内由不变令牌引起的算子级冗余，从而留下了进一步加速的空间。

我们提出LaCache，一种针对DLLMs中SAR推理的无需训练加速框架。LaCache通过重用在每个SAR块内去噪步骤之间不变令牌的中间状态来减少算子级冗余。它结合了**无损状态记忆化（LSM）**（重用输出相同的计算）和一种提高硬件效率的混合精度策略。首先，**EmbedCache**存储不变令牌的嵌入输出。其次，**RoPECache**存储第一层逐令牌预注意力结果，包括RoPE。第三，**FACache**存储第一层FlashAttention中的在线softmax状态，包括运行最大值、归一化器和未归一化输出累加器，这使我们能够跳过不涉及当前块的注意力分块。受不同层和去噪步骤间激活分布多样性的启发，LaCache进一步在FFN层上应用细粒度逐组FP8量化，该量化根据局部激活模式调整缩放比例，并在与周围算子融合的FP8 Tensor Core内核上运行，以缓解内存带宽瓶颈，同时保持模型精度。

实验表明，LaCache相比原始DLLM实现了约1.3×端到端推理加速。当与现有加速方案结合时，它达到高达40.2×端到端加速，同时保持相当的质量。我们的主要贡献如下：

- • 我们识别了DLLMs中SAR推理的算子级冗余，并引入了三个无损缓存组件，覆盖逐令牌算子和第一层FlashAttention的在线softmax状态，减轻了每个SAR块内的冗余计算。
- • 我们集成了一种面向FFN层的逐组FP8量化，适应不同的激活分布和FP8 Tensor Core内核，在保持精度的同时实现加速。
- • 我们在多个模型和基准测试上展示了一致的加速效果，表明LaCache与现有加速方法具有良好的组合性。

## 2 相关工作

**加速扩散语言模型。** 最近的研究从多个角度优化dLLM推理效率。一些工作(Jiang et al., 2025 (https://arxiv.org/html/2607.16339#bib.bib27); Liu et al., 2025b (https://arxiv.org/html/2607.16339#bib.bib42); Wu et al., 2025 (https://arxiv.org/html/2607.16339#bib.bib4))引入近似KV缓存机制以在去噪步骤间重用隐藏状态。具体来说，Fast-dLLM(Wu et al., 2025 (https://arxiv.org/html/2607.16339#bib.bib4))通过置信度感知并行解码提高吞吐量，用基于阈值的策略替代top-k采样。Dllm-cache(Liu et al., 2025b (https://arxiv.org/html/2607.16339#bib.bib42))存储和检索Transformer推理过程中生成的键值（KV）状态。从互补的角度，DPad(Chen et al., 2025 (https://arxiv.org/html/2607.16339#bib.bib5))通过修剪上下文长度来降低注意力的计算成本，采用带滑动窗口的后缀丢弃。其他工作关注自适应解码调度，例如插入指导令牌(Wei et al., 2025a (https://arxiv.org/html/2607.16339#bib.bib7))或监控令牌稳定性以动态调整生成阈值(Sun et al., 2025 (https://arxiv.org/html/2607.16339#bib.bib6))。虽然这些方法有效地优化了解码调度或上下文图，但它们并没有直接消除一个块内不变令牌上的算子级冗余，这正是我们的框架旨在弥补的差距。

**低精度推理。** 量化已被广泛采用以加速LLM推理。SmoothQuant(Xiao et al., 2023 (https://arxiv.org/html/2607.16339#bib.bib10))提出逐通道量化，通过跨通道平滑激活异常值以保持精度。COAT(Xie et al., 2024 (https://arxiv.org/html/2607.16339#bib.bib11))引入混合粒度量化，结合逐张量和逐组方案以平衡效率和精度。类似地，DeepSeek-V3(Liu et al., 2024 (https://arxiv.org/html/2607.16339#bib.bib9))对线性层采用逐组量化。MOSS(Zhang et al., 2025 (https://arxiv.org/html/2607.16339#bib.bib12))通过两级缩放策略进一步细化，以减少训练过程中的精度损失。基于这些见解，LaCache将逐组FP8量化应用于dLLM中的FFN线性层，利用了SAR推理中观察到的激活分布特性。

## 3 方法

参见图注
图3：FACache的更新和重用方案。**左图：** 在第一步的第一个Flash-attention中更新缓存。当qkv分块与当前生成块无重叠时，缓存按照相同的计算过程收集中间结果。**右图：** 在后续步骤的第一个Flash-attention中重用缓存。当q分块与当前生成块无重叠时，寄存器将使用缓存中的数据初始化。如果kv分块也与当前生成块无重叠，则可以直接跳过计算。

在本节中，我们首先分析DLLMs中的冗余计算，这种冗余由输入中的重复引起，以及推理过程中适用于量化的稳定激活动态。然后我们介绍无损缓存策略，该策略包含推理过程中中间结果的三个部分。最后，我们提出一种细粒度混合精度策略，以几乎无精度损失的方式加速推理。

### 3.1 概述

LaCache通过集成三个专门组件提出了一种无损状态记忆化（LSM）策略：**EmbedCache**、**RoPECache**和**FACache**。该机制确保数学上等价的输出，同时严格跳过不变令牌的冗余计算。为了进一步提高吞吐量，我们为计算密集的线性层采用了细粒度FP8混合精度策略。整体架构如图2所示。

### 3.2 DLLMs中的加速空间

#### 3.2.1 DLLM推理中的冗余计算

在一个SAR块的多个生成步骤中，只有当前块中的令牌被更新，而块外的令牌ID保持不变。因此，冗余计算出现在逐令牌算子（嵌入、注意力归一化、QKV投影、RoPE）以及第一层FlashAttention分块过程的部分中。通过适当的缓存策略，可以在不改变模型输出的情况下减少这些冗余开销。

DLLMs的结构仍然是基于Transformer的，线性层主导运行时。加速线性层的常用方法是低精度计算(Kalamkar et al., 2019 (https://arxiv.org/html/2607.16339#bib.bib23); Lee et al., 2009 (https://arxiv.org/html/2607.16339#bib.bib25); Micikevicius et al., 2017 (https://arxiv.org/html/2607.16339#bib.bib24); Liu et al., 2024 (https://arxiv.org/html/2607.16339#bib.bib9))。如图4所示，多个步骤中第一层激活的峰值集中在几个通道上，并且随着未来令牌被解掩码，峰值密度降低。在中间层，激活值相对平坦，只有少数通道有零星尖峰。在较深层，同一令牌的大多数通道的值基本相同。这表明细粒度方案可以更好地平衡效率和精度，特别是在包含许多掩码令牌的SAR输入下。

### 3.3 无损状态记忆化

为了消除算子级冗余，我们引入了无损状态记忆化（LSM）。它严格记忆不变区域的逐令牌算子的中间结果以及第一层FlashAttention的内部在线softmax统计量。记忆的状态在每个SAR块的第一个去噪步骤中填充，并在后续步骤中重用。