DiffusionGemma: vLLMでネイティブにサポートされた初のDiffusion LLM

vLLMは、Gemma4バックボーンに基づいて構築された26Bパラメータの離散拡散言語モデルであるDiffusionGemmaを統合しました。この統合は、vLLMにおける初の拡散LLM（dLLM）のネイティブサポートを意味し、逐次的なトークン生成から反復的なブロック精緻化への移行を可能にし、低バッチサイズにおけるレイテンシを大幅に削減します。

DiffusionGemmaのアーキテクチャとデコーディング

DiffusionGemmaは、トークンを左から右へ一つずつ生成する標準的な自己回帰デコーディングを、256トークンの固定長キャンバスを反復的にデノイズするプロセスに置き換えます。このアプローチは、メモリ帯域幅の負荷を計算量とトレードオフにすることで、モデルが複数のデノイジングステップにわたって複数のトークンを並列に精緻化することを可能にします。

デュアルモード動作

DiffusionGemmaは、2つの異なるモードで動作する単一の重みセットを使用します：

• Encoder Mode: 因果的アテンション（causal attention）を使用してプロンプトをプリフィルし、完了したブロックをKVキャッシュに「コミット」します。
• Decoder Mode: 双方向アテンション（bidirectional attention）を採用し、キャンバス内のすべての位置が他のすべての位置にアテンションを払うことで、同時並行的な精緻化を可能にします。

エンコーダーモードは標準的な因果的アテンションを使用するため、vLLMの自動プリフィックスキャッシュ（automatic prefix caching）は互換性が維持され、変更なしで使用できます。

サンプリングループとエントロピー境界デノイジング

生成は256トークンのブロック単位で行われます。プロンプトのプリフィル後、キャンバスはランダムなトークンで初期化されます。その後、モデルはすべての位置で候補トークンをサンプリングするデノイジングループに入ります。

どのトークンを保持するかを決定するために、DiffusionGemmaは**エントロピー境界（entropy-bound）ルールを使用します。これは、累積エントロピーが固定の予算を超えるまで、確信度の高いものから低いものへとトークンを受け入れる仕組みです。キャンバスが収束（converged）**すると（argmax予測が数ステップ安定し、トークンあたりの平均エントロピーが閾値を下回る、またはステップ制限に達した状態と定義）、トークンはエンコーダーパスを介してコミットされ、プロセスは次のブロックに対して繰り返されます。

安定性のための自己条件付け

収束を加速させ安定性を高めるために、モデルは**自己条件付け（self-conditioning）**を使用します。デノイジングステップの間、モデルは前の予測の完全なsoftmax分布に基づいて条件付けられます。この分布は、トークン埋め込みの確率重み付き平均に変換され、ゲート付きMLPを介してキャンバスの埋め込みに追加され、モデルに以前の信念のメモリを提供します。

vLLMの実装詳細

DiffusionGemmaの統合には、非自己回帰的なデコーディングパターンのサポートが必要でした。vLLMは、いくつかのアーキテクチャ抽象化を活用することでこれを実現しました：

Speculative Decoding データパス

vLLMは既存の投機的デコーディング（speculative decoding）パスを再利用し、現在のキャンバスを、完全に拒否されるか完全に受け入れられるドラフトトークンのセットとして扱います。これにより、チームはコアとなるスケジューラとモデルランナーを最小限の変更で維持することができました。

ModelState インターフェース

モデルランナーのフォークを避けるため、vLLMはModelState抽象化を利用しました。これにより、DiffusionGemmaは汎用ランナーを変更することなく、カスタム動作を定義するためのフックを提供できます：

フック	DiffusionGemmaにおける目的
prepare_inputs()	キャンバスの埋め込みと自己条件付けの適用を処理する。
prepare_attn()	リクエストごとに因果的（エンコーダー）アテンションと双方向（デノイズ）アテンションを切り替える。
custom_sampler()	デフォルトのサンプラーを置き換えるために DiffusionSampler をインストールする。
add_request() / remove_request()	リクエストごとの拡散状態（キャンバス、確率）のライフサイクルを管理する。

動的なシーケンスごとの因果的アテンション

単一のバッチに異なる段階（プリフィル、デノイズ、コミット）のリクエストが含まれる可能性があるため、vLLMは**動的なシーケンスごとの因果的アテンション（dynamic per-sequence causal attention）**を実装しました。これにより、アテンションマスクが各リクエストの特定の因果関係の要件に適応できるようになります。この機能は、Triton Attention (TRITON_ATTN) と FlashAttention 4 (FLASH_ATTN) の両方のバックエンドでサポートされています。

対称スライディングウィンドウアテンション

スライディングウィンドウアテンションを使用するレイヤーの場合、DiffusionGemmaはデノイジングフェーズ中に対称ウィンドウを必要とします。因果的リクエストは片側ウィンドウ（前の $W$ トークンにアテンションを払う）を使用しますが、双方向リクエストは両側の $W$ トークンにアテンションを払い、合計ウィンドウサイズは $2W + 1$ となります。

パフォーマンスと量子化

スループット結果

DiffusionGemmaのアーキテクチャは、極めて低レイテンシの推論を可能にします。単一GPUのバッチサイズ1で実施されたベンチマークでは、自己回帰ベースラインと比較して大幅なスループットの向上が示されています：

• H200 (FP8): 1,288 generation tokens/second (6$\times$ 標準的な自己回帰ベースライン、3$\times$ マルチトークン予測)
• H100 (FP8): 1,008 generation tokens/second (5$\times$ 標準的な自己回帰ベースライン、2.6$\times$ マルチトークン予測)

量子化チェックポイント

モデルの量子化バージョンはLLM Compressorを使用して作成され、compressed-tensors 形式で利用可能です。サポートされている形式は以下の通りです：

• FP8: 完全動的アクティベーションを備えた量子化重み。
• NVFP4: 重みとアクティベーションの両方がNVFP4形式に量子化。