祝贺 @GoogleDeepMind 发布 DiffusionGemma！一个基于 Gemma4 主干网络的 26B 参数扩散语言模型，也是 vLLM 原生支持的首个 dLLM。

它并行去噪 256 个 token 组成的块，而非逐个生成 token：在单张 H200（FP8）上，批大小为 1 时，输出速度可达 1200+ tok/s。

构建于模型运行器 v2 的 ModelState 之上，并结合现有的投机解码路径，仅需极少的调度器或运行器改动。FP8 和 NVFP4 检查点已在 @RedHat_AI 模型库上发布。感谢 @GoogleDeepMind、@RedHat_AI 和 @NVIDIAAI 团队！

http://vllm.ai/blog/2026-06-10-diffusion-gemma…

DiffusionGemma：vLLM 原生支持的首个扩散语言模型 (dLLM)

来源：https://vllm.ai/blog/2026-06-10-diffusion-gemma

vLLM 中的 DiffusionGemma

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

我们利用模型运行器 v2 (https://vllm.ai/blog/2026-03-24-mrv2) 的新 ModelState 抽象将 DiffusionGemma 集成到 vLLM 中。该抽象允许模型定义其自定义输入准备，并提供挂钩来管理每个请求的模型特定状态。结果在准确率上与原生的 Hugging Face 实现匹配，同时实现了高效的批处理服务。

与标准自回归 Transformer（从左到右逐个生成 token）不同，扩散语言模型通过对固定长度的画布进行迭代去噪来生成 token。这使得模型能够在多个去噪步骤中并行优化多个 token，从而有效地用额外计算换取内存带宽压力——这在批大小较小时尤其具有吸引力，因为此时空闲计算资源充足，而内存带宽是瓶颈。每次前向传播生成多个 token 可以带来极低的延迟响应。DiffusionGemma 特别地一次去噪一个包含 256 个 token 的画布。

自回归 vs. 块扩散解码。

DiffusionGemma 架构与采样循环

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

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

由于编码器使用普通的因果注意力，且提交的 KV 缓存与自回归模型完全相同，因此 vLLM 的自动前缀缓存可以开箱即用：共享的提示前缀可以在不同请求间重用，无需任何扩散特有的更改。

单个 256 token 块的循环工作方式如下：提示预填（编码器）后，画布被初始化为随机 token，然后其状态设置为去噪。每个去噪步骤在解码器模式下运行主干网络，覆盖整个画布，在每个位置采样一个候选 token，并决定保留哪些位置。一旦块停止变化，状态将设置回编码模式，并执行最后一次编码器传递来提交它——写入其 KV 缓存并输出 256 个 token——然后下一个块从新的随机画布开始。

DiffusionGemma 每块采样循环。

在一个块内，所有 256 个位置并行去噪；跨块时，生成仍是从左到右的，因为每个新块都依赖于之前所有已提交的 token。

熵界去噪

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

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

上下文中的块去噪。

一旦画布的“最佳猜测”（argmax）预测在连续几个步骤中停止变化并且其每个 token 的平均熵低于置信度阈值——或者达到硬性去噪步骤上限——则认为画布已收敛。此时提交的 token 是那个清晰的 argmax 预测，而不是步骤之间传递的带噪声的采样画布。

自条件化

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

自条件化反馈路径。

这为每个步骤提供了模型上一次信念的记忆，因此即使是重新添加了随机噪声的位置，也会从前一步携带信息，而不是从头开始。自条件化仅在解码器/去噪模式下激活——在编码器预填和提交传递中，该反馈被置零，因此这些传递看到的是普通的 token 嵌入。

vLLM 中的实现

复用投机解码数据路径

vLLM 引擎已经拥有非常成熟和稳定的投机解码路径。受 RFC #36155 (https://github.com/vllm-project/vllm/issues/36155) 的启发，我们复用此路径来实现 DiffusionGemma。在 vLLM 中为扩散 LLM 复用投机解码路径是一种自然的契合，因为每一步中，当前的画布可以被视为一大组草稿 token，这些 token 要么被完全拒绝，要么被完全接受。这导致对核心 vLLM 组件（如调度器和模型运行器）的改动极小。一个显著的例外是，在投机解码中我们总是采样一个额外的 token（通常称为投机解码文献中的奖励 token），因此添加了对采样 0 个 token 的支持，并由 ModelState 控制。

具体来说，扩散模型按如下方式插入现有堆栈——调度器、模型运行器和 Gemma4 主干网络被原样复用，只有 ModelState 和采样器是扩散特有的：

DiffusionGemma 在 vLLM 投机解码堆栈中的插入方式。

ModelState 接口

在 ModelState 之前，向 V1 添加非自回归模型将需要派生生模型运行器，并通过输入准备、注意力元数据和采样来线程化扩散特定状态。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.compile 函数 _compiled_sample_step，对所有正在运行中的解码请求进行向量化，涵盖三种情况：

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

在去噪期间，采样器报告 num_sampled = 0 和 num_rejected = query_len，因此 KV 缓存位置不移动；只有提交才推进它。将每个画布位置标记为拒绝告诉调度器将序列保持在其当前位置，并在下一步骤重新调度相同的块，这将整个去噪循环保持在现有投机解码记账中，而无需任何调度器更改。

动态每序列因果注意力

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

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

动态每序列因果注意力。

我们在两个注意力后端中支持这种动态因果注意力：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 个 key，将注意力沿对角线缩小为一个带状区域——而请求 1 的去噪画布使用对称窗口，关注两侧各 W 个 key，因此仅关注落在该窗口内的上下文 token。

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

量化检查点支持

DiffusionGemma 模型的量化检查点是使用 LLM Compressor (https://github.com/vllm-project/llm-compressor) 创建的，并以 compressed-tensors (https://github.com/vllm-project/compressed-tensors) 格式保存。其中包括一个具有量化权重和完全动态激活的 FP8 模型，以及一个将权重和激活都量化到 NVFP4 格式的 NVFP4 模型。

量化检查点可在 RedHatAI 模型库中找到：

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