DiffusionGemma：vLLM原生支持的首个扩散LLM（dLLM）

vLLM 中的 DiffusionGemma

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

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

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

自回归 vs. 块扩散

自回归 vs. 块扩散解码。

DiffusionGemma 架构与采样循环

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

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

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

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

DiffusionGemma 块采样循环

DiffusionGemma 的每块采样循环。

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

熵界去噪

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

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

上下文中的块去噪

多个步骤上的熵界去噪。

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

自条件化

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

自条件化

自条件化反馈路径。

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

在 vLLM 中的实现

重用推测解码数据路径

vLLM 的引擎已经拥有一个非常成熟且稳定的推测解码路径。受 RFC #36155 的启发，我们重用了这条路径来实现 DiffusionGemma。在 vLLM 中为扩散 LLM 重用推测解码路径是自然而然的，因为在每个步骤中，当前的画布可以被视为一大组草稿 token，这些 token 要么被完全拒绝，要么被完全接受。这导致对核心 vLLM 组件（如调度器和 model runner）的更改非常少。一个显著的例外是，使用推测解码时，我们总是采样一个额外的 token（在推测解码文献中通常称为奖励 token），我们添加了对采样 0 个 token 的支持，并由 ModelState 控制。

具体来说，扩散按如下方式插入到现有堆栈中——调度器、model runner 和 Gemma4 骨干网络保持不变地重用，只有 ModelState 和采样器是扩散特定的：

DiffusionGemma 如何插入 vLLM 的推测解码堆栈

vLLM 软件抽象中的 DiffusionGemma。

ModelState 接口

在 ModelState 之前，向 V1 添加非自回归模型需要 fork model runner 并通过输入准备、注意力元数据和采样来传递扩散特定状态。ModelState 通过定义一组钩子来避免这种情况，model runner 在前向循环的每个阶段调用这些钩子：

钩子	DiffusionGemma 用它来...
prepare_inputs()	嵌入画布 token 并应用自条件化
prepare_attn()	设置每个请求的因果（编码器）vs. 双向（去噪）注意力
custom_sampler()	用 DiffusionSampler 替换默认采样器
add_request() / remove_request()	初始化和销毁每个请求的扩散状态（例如画布和自条件化概率）

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

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

整合：DiffusionGemmaModelState 和 DiffusionSampler

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

DiffusionSampler 取代了 vLLM 通常的 (Sampler, RejectionSampler) 对，并负责在阶段变化期间初始化和重置画布以及每个请求的扩散状态。每个步骤的工作是一个单一的 @torch.compile 函数 _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，请求在这些模式之间交替，因为 prompt 被预填充，然后画布被迭代去噪和接受。为了最小化延迟，vLLM 在每次前向传递期间将处于不同阶段的请求混合在批次中。因此，我们实现了 动态每序列因果注意力，它根据每个请求的因果性调整注意力掩码。下图描述了这种情况：这里，我们展示了一个包含三个请求的批次，每个请求处于不同阶段。

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

动态每序列因果注意力

动态每序列因果注意力。

我们在两个注意力后端中支持这种动态因果注意力：Triton Attention (TRITON_ATTN) 和 FlashAttention 4 (FLASH_ATTN)。在这两个后端中，单个布尔参数 causal 被替换为一个指示每个请求因果性的张量。掩码被适当更新，并且 tile 行为得以保留。

滑动窗口注意力

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

H100 和 H200 上的生成吞吐量：FP8 扩散 vs. 自回归基线

H100 和 H200 上的生成吞吐量——FP8 扩散 vs. 自回归基线。复现命令

致谢

感谢所有为将 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