DiffusionGemma: 首個原生支援於 vLLM 的擴散 LLM

vLLM 已整合 DiffusionGemma，這是一個建立在 Gemma4 骨幹上的 26B 參數離散擴散語言模型。此整合標誌著 vLLM 首次原生支援擴散 LLM (dLLM)，使得從序列式 token 生成轉變為迭代區塊精煉，顯著降低低批次大小時的延遲。

DiffusionGemma 架構與解碼

DiffusionGemma 以一個迭代去噪固定長度為 256 個 token 的畫布的過程，取代標準的自回歸解碼（即從左到右逐個生成 token）。此方法用額外的計算換取記憶體頻寬壓力，使模型能夠在多個去噪步驟中並行精煉多個 token。

雙模式操作

DiffusionGemma 使用一組權重在兩種不同模式下運作：

• Encoder Mode: 使用因果注意力來預填提示詞，並將已完成的區塊「提交」到 KV 快取。
• Decoder Mode: 使用雙向注意力，使畫布上的每個位置都能關注其他所有位置，以進行同時精煉。

因為編碼器模式使用標準的因果注意力，vLLM 的自動前緩存保持相容且無需修改即可運作。

取樣循環與熵限制去噪

生成以 256-token 為一個區塊進行。在提示詞預填後，畫布會被初始化為隨機 token。接著，模型進入去噪循環，在每個位置採樣候選 token。

為決定保留哪些 token，DiffusionGemma 使用一個 entropy-bound 規則：它會從最有信心到最不有信心的 token 接受，直到累積熵超過固定預算。當畫布達到 converged 狀態時——定義為 argmax 預測在數個步驟中保持穩定且每個 token 的平均熵低於某個閾值（或達到步驟上限）——則透過編碼器通過將 token 提交，並對下一個區塊重複此過程。

為穩定性的自我條件化

為加速收斂並提高穩定性，模型使用 self-conditioning。在去噪步驟之間，模型會根據其先前預測的完整 softmax 分布進行條件化。該分布會被轉換為 token 嵌入的機率加權平均，並通過一個 gated MLP 加到畫布嵌入中，為模型提供其先前信念的記憶。

vLLM 實作細節

整合 DiffusionGemma 需要支援非自回歸解碼模式。vLLM 透過利用幾種架構抽象來實現此目標：

投機解碼解碼路徑

vLLM 重複使用其現有的投機解碼路徑，將當前畫布視為一組要么完全碼資料路徑

vLLM 重複使用其現有的投機解碼路徑，將當前畫布視為一組要么完全被拒絕要么完全被接受的草稿 token。這使得團隊能以最小的變動維護核心排程器和模型運行器。

ModelState 介面

為避免分叉模型運行器，vLLM 使用了 ModelState 抽象。這提供了鉤子，使 DiffusionGemma 能在不更改通用運行器的情況下定義自訂行為：

Hook	在 DiffusionGemma 中的用途
prepare_inputs()	處理畫布嵌入和自我條件化的應用。
prepare_attn()	每個請求在因果（編碼器）和雙向（去噪）注意力之間切換。
custom_sampler()	安裝 DiffusionSampler 以取代預設採樣器。
add_request() / remove_request()	管理每個請求的擴散狀態（畫布、機率）的生命週期。

動態每序列因果注意力

因為單個批次可能包含處於不同階段（預填、去噪、提交）的請求，vLLM 實作了 動態每序列因果注意力。這使得注意力遮罩能夠適應每個請求特定的因果需求。此功能同時支援 Triton Attention (TRITON_ATTN) 和 FlashAttention 4 (FLASH_ATTN) 後端。

對稱滑動窗口注意力

對於使用滑動窗口注意力的層，DiffusionGemma 在去噪階段需要對稱的窗口。而因果請求使用單側窗口（僅關注前 $W$ 個 token），雙向請求則在兩側各關注 $W$ 個 token，總窗口大小為 $2W + 1$。

效能與量化

吞吐量結果

DiffusionGemma 的架構使得極低延遲的推理成為可能。在單個 GPU 上，批次大小為 1 的基準測試顯示，相較於自回歸基線，吞吐量有顯著提升：

• H200 (FP8): 每秒 1,288 個生成 token（約為標準自回歸基線的 6 倍，約為多 token 預測的 3 倍）。
• H100 (FP8): 每秒 1,008 個生成 token（約為標準自回歸基線的 5 倍，約為多 token 預測的 2.6 倍）。

量化檢查點

模型的量化版本是透過 LLM Compressor 建立的，並以 compressed-tensors 格式提供。支援的格式包括：

• FP8: 權重進行量化，激活函數為完全動態。* NVFP4: 權重和激活函數均量化為 NVFP4 格式。