DiffusionGemma 成为 vLLM 原生支持的首个扩散语言模型

提示：想要部署 DiffusionGemma？请参阅 vLLM 配方 获取部署说明。

Google 的 DiffusionGemma 是一个基于 Gemma4 骨干构建的 26B 参数离散扩散语言模型，也是 vLLM 中支持的第一个 dLLM。将 DiffusionGemma 集成到 vLLM 中需要支持一种根本不同的解码模式。dLLM 无法干净地适配标准的自回归服务路径：它们需要双向注意力、迭代细化、基于块的生成，以及每个去噪步骤中的自定义采样行为。

我们使用 model runner v2 的 新 ModelState 抽象将 DiffusionGemma 集成到 vLLM 中，该抽象允许模型定义其自定义输入准备，并提供用于管理每个请求的模型特定状态的钩子。结果与 Hugging Face 参考实现的准确性相匹配，同时实现了高效的批处理服务。

与标准自回归 Transformer 从左到右一次生成一个 token 不同，扩散语言模型通过迭代去噪固定长度的画布来生成 token。这允许模型在多个去噪步骤中并行细化多个 token，有效地用内存带宽压力换取额外的计算——在低批量大小下，这是一个特别有吸引力的权衡，因为此时备用计算充足而内存带宽是瓶颈。每次前向传递生成许多 token 可以转化为非常低延迟的响应。DiffusionGemma 具体一次去噪 256 个 token 的画布。

DiffusionGemma 架构与采样循环

DiffusionGemma 构建在标准 Gemma4 骨干之上，但以两种模式运行，共享相同的权重——一组层，以两种方式使用：

• 编码器模式使用因果注意力并写入 KV 缓存。每个块运行两次：一次用于预填充提示，一次用于“提交”完成的块。
• 解码器模式使用双向注意力，仅读取 KV 缓存。这是去噪模式——画布中的每个位置都可以关注其他每个位置，这正是让模型能够一次性细化整个块的原因。

由于编码器使用普通的因果注意力，并且提交的 KV 完全按照自回归模型的方式写入，vLLM 的自动前缀缓存开箱即用：共享的提示前缀在请求之间被重用，无需任何扩散特定的更改。

单个 256-token 块的循环如下。在提示被预填充（编码器）后，画布被初始化为随机 token，然后其状态被设置为去噪。每个去噪步骤在解码器模式下对整个画布运行骨干，在每个位置采样一个候选 token，并决定保留哪些位置。一旦块停止变化，状态被设置回编码，最后的编码器传递提交它——写入其 KV 并发出 256 个 token——下一个块从新的随机画布开始。

在一个块内，所有 256 个位置并行去噪；跨块时，生成仍然是左到右的，因为每个新块都以所有先前提交的 token 为条件。

熵约束去噪

每个去噪步骤都会重新采样画布上的所有位置，但只有模型确信的位置会被保留；其余位置被丢弃，并在下一步替换为全新的随机 token。置信度通过每个位置预测分布的熵来衡量——低熵意味着模型基本已经拿定主意。

DiffusionGemma 使用熵约束规则来决定接受多少个位置：它从最确信到最不确信依次遍历位置，接受 token，直到其累积熵超过固定预算。在早期，模型几乎对所有事情都不确定，因此只有少数位置被锁定。随着这些锚点将上下文传播到相邻位置，分布变得更加尖锐，更多位置落入预算之内，整个块在若干步骤内逐渐清晰聚焦。

当画布的最佳猜测（argmax）预测连续若干步骤不再变化，并且其平均每 token 熵低于置信度阈值时，该画布被视为已收敛——或者它达到硬性去噪步骤上限。此时，已提交的 token 就是那个干净的 argmax 预测，而不是步骤之间携带的带噪采样画布。

自条件化

为了使去噪循环更稳定并更快收敛，DiffusionGemma 使用了自条件化：在步骤之间，模型以自己上一步的预测为条件。它不反馈硬 token，而是反馈上一步完整的 softmax 分布，将其转换为 token 嵌入的概率加权平均值，并通过一个小型门控 MLP 将其加到画布嵌入上，然后再进行下一轮。

这使每一步都能记住模型上次的信念，因此即使是被重新加噪为随机 token 的位置，也会携带上一步的信息，而不必从头开始。自条件化仅在解码器/去噪模式下激活——在编码器预填充和提交阶段，反馈被置零，因此这些阶段看到的是普通的 token 嵌入。

在 vLLM 中的实现

复用投机解码数据路径

vLLM 的引擎已经拥有一条非常成熟且稳定的投机解码路径。受 RFC #36155 启发，我们复用这条路径来实现 DiffusionGemma。在 vLLM 中为扩散 LLM 复用投机解码路径是一种自然的契合，因为在每一步，当前画布都可以被视为一大组草稿 token，它们要么被完全拒绝，要么被完全接受。这使得对 vLLM 核心组件（如调度器和模型运行器）的改动非常小。值得注意的例外是，在投机解码中我们总是额外采样一个 token（在投机解码文献中通常称为奖励 token），因此我们增加了对采样 0 个 token 的支持，并由 ModelState 控制。

具体来说，扩散按如下方式接入现有技术栈——调度器、模型运行器和 Gemma4 主干均原样复用，只有 ModelState 和采样器是扩散专用的：

ModelState 接口

在 ModelState 出现之前，向 V1 添加一个非自回归模型需要 fork 模型运行器，并在输入准备、注意力元数据和采样过程中贯穿扩散特定的状态。ModelState 通过定义一组钩子来避免这一点，运行器在前向循环的每个阶段都会调用这些钩子：

模型通过在模型类上定义 get_model_state_cls() 来自我注册其 ModelState。模型运行器保持通用。在每一步，它调用 prepare_attn(...) 来构建元数据，将 prepare_inputs(...) 合并到前向 kwargs 中，并将采样委托给 custom_sampler()->DiffusionSampler 安装的任何采样器。

这意味着添加一个新的块扩散模型需要实现一个 ModelState 并在模型类上进行一行注册，而无需更改运行器、调度器或任何共享基础设施。我们相信这可以作为未来向 vLLM 干净地添加扩散语言模型的蓝图。

整合起来：DiffusionGemmaModelState 和 DiffusionSampler

DiffusionGemmaModelState 是 DiffusionGemma 的 ModelState 实现。它保存每个请求的状态（主要与扩散循环相关）：一个阶段标志，表示请求是处于提交还是去噪阶段，当前的 canvas，用于收敛检查的历史记录，自条件概率等。此状态存在于预分配的 GPU 张量中，并就地更新。DiffusionGemmaModelState.prepare_inputs() 嵌入画布 token 并应用自条件：它从上一步去噪步骤中获取 softmax 分布（来自内部每个请求的状态），计算 token 嵌入的概率加权平均值，并将其通过门控 MLP 馈送，以便模型可以看到自己之前的预测。prepare_attn() 构建注意力元数据，使用阶段标志来决定注意力应该是因果的（提交阶段/编码器）还是双向的（去噪阶段/解码器）。由于单个批次可以混合包含预填充、去噪和提交请求，并且每个请求的因果标志是在 GPU 上异步设置的，我们不得不对注意力内核进行一些修改，我们将在后面的章节中讨论。

DiffusionSampler 取代了 vLLM 通常的 (Sampler, RejectionSampler) 对，负责在阶段变化期间初始化和重置画布以及每个请求的扩散状态。每步工作是一个单一的 @torch.compiled 函数 _compiled_sample_step，在所有进行中的解码请求上向量化，涵盖三种情况：

• 预填充：将画布初始化为随机 token 并返回 num_sampled = 0。
• 去噪：对 logits 进行温度缩放，使用 Gumbel-max 技巧（argmax(logits/T + gumbel_noise)）在每个画布位置抽取一个候选 token，接受最自信的位置直到熵界限，并将其余部分重新噪声化为随机 token。该步骤还记录 argmax 画布并检查收敛：argmax 画布在配置的步数内保持稳定且平均熵低于阈值，或达到步数上限。
• 提交：发出干净的 argmax_canvas（num_sampled = 256），为下一个块重新初始化画布，并重置每个请求的状态。

在去噪过程中，采样器报告 num_sampled = 0 和 num_rejected = query_len，因此 KV 缓存位置不会移动；只有提交才会推进它。将每个画布位置标记为已拒绝会告诉调度器保持序列位置不变，并在下一步重新调度同一个块，从而使整个去噪循环保持在现有的推测解码核算之内，无需对调度器做任何更改。

动态逐序列因果注意力

如上所述，DiffusionGemma 以两种模式运行：使用因果注意力的 编码器 模式和使用双向注意力的 解码器 模式。到目前为止，因果性是一个批次范围的属性——一次前向传播中的每个请求共享相同的掩码类型。典型的解码器模型仅使用因果注意力，而像 Whisper 这样的编码器-解码器模型在其编码器层中仅使用双向注意力。然而对于 DiffusionGemma，随着提示被预填充，然后画布被迭代去噪并接受，请求会在这些模式之间交替。为了最小化延迟，vLLM 在每次前向传播中将处于不同阶段的请求混合在同一批次中。因此，我们实现了 动态逐序列因果注意力，它会根据每个请求的因果性调整注意力掩码。下图描绘了这种情况：这里，我们展示了一个包含三个请求的批次，每个请求处于不同的阶段。

• 请求 0 是长度为 6 的预填充，因此它使用因果注意力（“编码器”传递），其中对角线以上的条目被掩蔽——每个查询 token 只关注来自其自身及之前 token 的键。我们还注意到，注意力是按块计算的（在此示例中形状为 2x2，但在实践中这些块要大得多，并且具有依赖于硬件的调优），仅包含被掩蔽条目的块会被完全跳过，从而节省计算量以及从 HBM 加载其 K/V 块的内存带宽。
• 请求 1 已经完成了长度为 6 的预填充，现在正在以解码器模式生成新 token。在大小为 4 的画布内，所有查询使用双向注意力关注画布中的所有键。它们还关注上下文中的所有键。没有条目被掩蔽，也没有块被跳过。
• 最后，请求 2 已完成其去噪步骤，其画布已准备好被接受。我们最后一次运行编码器传递，使用因果注意力，并用新接受的 token 的条目填充 KV 缓存。同样，所有查询也关注缓存的键。

我们在两个注意力后端中支持这种动态因果注意力：Triton Attention（TRITON_ATTN）和 FlashAttention 4（FLASH_ATTN）。在这两个后端中，单个布尔参数 causal 被替换为一个张量，用于指示每个请求的因果性。掩码会相应更新，分块行为得以保留。

滑动窗口注意力

最后，DiffusionGemma 的某些层使用滑动窗口注意力。对于画布中的 token，滑动窗口注意力也必须变为对称的：对于窗口大小 W，画布 token 不再只关注自身及其之前的 W 个 token，而是还关注其之后的 W 个 token，总窗口大小为 2*W + 1。我们在下面描绘了这一点：

与之前一样，同样的三个请求显示在一个带有 W=2 的滑动窗口层上。请求 0 和 2（预填充和接受）保持单侧因果窗口——每个查询关注自身及其之前的 W 个键，将注意力收窄到沿对角线的一条带内——而请求 1 的去噪画布使用对称窗口，关注两侧各 W 个键，因此只关注落在其中的上下文 token。

要在两个后端中支持这一点，只需修改双向请求的窗口右边界：因果请求保持仅左侧窗口，而双向请求使用每侧 W 的对称窗口。

量化检查点支持

DiffusionGemma 模型的量化检查点是使用 LLM Compressor 创建的，并以 compressed-tensors 格式保存。其中包括一个权重已量化且激活完全动态的 FP8 模型，以及一个权重和激活均量化为 NVFP4 格式的 NVFP4 模型。

量化检查点可在 RedHatAI hub 上找到：

1. https://huggingface.co/RedHatAI/diffusiongemma-26B-A4B-it-NVFP4
2. https://huggingface.co/RedHatAI/diffusiongemma-26B-A4B-it-FP8-dynamic

为验证模型的准确性，我们在启用和未启用思考的情况下，使用 vLLM 在 AIME 2025、GPQA Diamond 和 GSM8k 基准上进行了初步评估。评估和恢复分数见模型卡。

结果

DiffusionGemma 的架构可实现极低延迟推理，非常适合交互式应用。为评估我们的实现在此场景下的性能，我们使用内置的 vllm bench serve 在单块 H100 和 H200 上以批大小 1 对 vLLM 进行了基准测试。FP8 扩散模型在 H200 上达到每秒 1,288 个生成 token（约为标准自回归基线的 6 倍，约为使用多 token 预测的基线的 3 倍），在 H100 上达到每秒 1,008 个 token（分别约为 5 倍和 2.6 倍）。

致谢

感谢所有为将 DiffusionGemma 引入 vLLM 做出贡献的人。这是 Google DeepMind 与 vLLM 团队之间的紧密合作。

• Google DeepMind： Martin Kukla、João Gante、Luciano Martins
• vLLM： Lucas Wilkinson、Matthew Bonanni、Nicolò Lucchesi、Dipika Sikka、Doug Smith、Edward Arthur Quarm Jnr、Alon Kellner（Red Hat）、Nick Hill（Inferact）
• NVIDIA： Dimitrios Bariamis、Alec Kohlhoff、Porras Huang、Eugene Rakhmatulin