SGLang が拡散大規模モデルを強化：LLaDA 2.0 を即日サポート

TL;DR

SGLang 内での Diffusion Large Language Model (dLLM) フレームワークの設計と実装を興奮をもってご紹介します。既存の ChunkedPrefill メカニズムを活用することで、このシステムは以下を実現しました：

• シームレスな統合：SGLang エコシステムに組み込まれ、コアアーキテクチャの変更は不要。
• 性能の継承：既存の推論最適化技術の恩恵を受ける。
• 最大限の柔軟性：ユーザーが拡散デコードアルゴリズムを完全にカスタマイズ可能。

背景

動機

今年初め、LLaDA が初の Diffusion Large Language Model (dLLM) として登場し、学術界と産業界の注目を急速に集めました。中国人民大学とアントグループの共同研究成果であるこの技術は、dLLM の独自の実行パラダイムによる優れたデータ理解能力を示し、低レイテンシの小バッチシナリオでは推論速度が Auto-Regressive (AR) モデルを上回りました。

dLLM のパラメータ規模が拡大するにつれ、AR LLM と同様の scaling-law 効果を観察しました。より強力な dLLM を追求するため、私たちは 100B パラメータの LLaDA2.0-flash モデルを訓練しました。

しかし、LLaDA2.0-flash の訓練過程で、私たちは一連の AI インフラストラクチャエンジニアリングの課題に直面しました。特にモデル評価と RL 事後訓練の効率性と安定性が問題となりました。

課題

既存の dLLM 推論エンジンは大規模モデルの評価と RL 事後訓練のニーズを満たすことができません。例えば、Fast-dLLM は優れていますが、アルゴリズム研究により適しており、バッチ処理、スケジューリング、RL エコシステム統合、並列処理などの本番レベルのサービス能力が欠けています。

対照的に、SGLang は現在最も人気のある LLM 推論エンジンで、複数の利点があります：

1. 本番環境対応：数千社で導入済みで、成熟した信頼性の高いエンジニアリング能力を提供。
2. 技術的リーダーシップ：多数の先進的な推論最適化を統合し、コミュニティが継続的に新しい最適化を貢献。
3. 完全なエコシステム：RL 事後訓練エコシステムと高度に統合、特に分散ウェイトの GPU P2P 更新。

しかし SGLang は現在 Auto-Regressive 計算パラダイムのみをサポートしており、拡散計算には適応していません。したがって、課題は既存アーキテクチャを破壊することなく、dLLM のサポートを導入し、SGLang のすべての最適化の恩恵を受けられるようにすることです。

設計

重要な洞察

dLLM の現在の発展に基づき、いくつかの重要な洞察を抽出しました：

1. Bidirectional Attention Diffusion は計算コストが高く KV Cache 利用率が低いため、主流の dLLM は Block Diffusion アーキテクチャに移行している。
2. Block Diffusion の計算モードは SGLang の既存の Chunked-Prefill プロセスと高度に類似している。
3. AR モデルとは異なり、拡散言語モデルは多様なデコード戦略が必要で、専用インターフェースによる柔軟なカスタマイズをサポートする必要がある。

アーキテクチャ

私たちは SGLang の既存の Chunked-Prefill パイプラインを活用して、Block Diffusion LLM の計算サポートを実装しました。この手法は dLLM をシームレスに統合し、コアフレームワークを変更することなく、SGLang が蓄積した推論最適化から直接恩恵を受けます。

図に示すように、SGLang への修正は極めて控えめで、コアの周辺部分のみに触れています。元の generate request 実行フローは変更されず、主に Chunked Prefill を利用・修正し、prefill adder と chunked reqs の 2 つのコンポーネントに焦点を当てています。

SGLang の Chunked Prefill は GPU 利用率を最大化することを目的とし、単一チャンクのサイズは通常 2K-16K トークンです。しかし dLLM のデコードはブロックレベルで分割されます（例：LLaDA2.0 の 32 トークンブロック）。単一の大きなリクエストのロジックを踏襲すると GPU パフォーマンスを無駄にするため、効率的なバッチ処理を解決する必要があります。私たちは chunked reqs と prefill adder を修正し、単一の計算サイクル内で複数の Diffusion Block の処理をサポートしました。

さらに、デコード実行層では、TP Worker と Model Runner の間に拡散アルゴリズム抽象層を挿入しました：

• Worker が Diffusion モデルを識別すると、専用ブランチに入る。
• Diffusion アルゴリズムの run 関数を呼び出す。
• 内部で前方反復ループを通じて Model Runner を駆動し、完全な Block デコードまで続行。

Attention Mask

Block Diffusion と Chunk Prefill の単一前方伝播の最大の違いは attention mask の処理にあります：

• Block Diffusion は block-wise causal mask を使用。
• AR モデルの Chunk Prefill は従来の token-wise causal mask を使用。

Block Diffusion は Chunk Prefill の機能拡張と見なすことができます。単一の前方伝播は 2 つの部分の計算を含み、最終出力は連結されます：

1. Context Query：現在の Q_curr と既存の KV Cache の双方向注意力で、Chunk Prefill と完全に同じで、現在のブロックが履歴情報に注目することを保証。
2. Intra-Block Query：現在の Q_curr と自身の KV の計算。
   • Block Diffusion は双方向注意力を使用。
   • Chunk Prefill は causal Mask を使用。

可視化 Q_curr attention mask：

• Chunk Prefill (causal) は台形/三角マスク。
• Block Diffusion (双方向) は矩形マスク。

ストリーミング出力デモ

以下のアニメーションは LLaDA2.0-flash-CAP (100B / BF16) と gpt-oss-120B (117B / MXFP4) のストリーミング出力を比較しています。LLaDA2.0-flash-CAP は 8×H20 で SGLang dLLM TP8 を使用し、gpt-oss-120B は同じハードウェアで標準 AR プロセスを使用しています。

タスク：10 種類のプログラミング言語でクイックソートを実装 — 拡散 LLM が得意とするシナリオです。図に示すように、LLaDA2.0-flash-CAP のスループットは 935 tokens/s に達し、gpt-oss-120B の 263 tokens/s を大きく上回っています。

SGLang dLLM は AR モデルのようなストリーミング出力をサポートしますが、ブロック（例：32 トークン）単位で出力します。

使用方法

起動コマンド例

python3 -m sglang.launch_server \
  --model-path inclusionAI/LLaDA2.0-mini \ # サンプル HF/ローカルパス
  --dllm-algorithm LowConfidence \
  --dllm-algorithm-config ./config.yaml \ # オプション、アルゴリズムのデフォルト設定を使用
  --host 0.0.0.0 \
  --port 30000

注意：--dllm-algorithm-config の高度な設定で選択した --dllm-algorithm を使用します。この機能は設定とコードを分離し、ユーザーがカスタマイズしたアルゴリズムに統一されたエントリポイントでパラメータを渡すのに便利です。

クライアントコード例

他のサポートモデルと同様に、dLLM は REST API またはオフラインエンジン API を通じて使用できます。

Curl 生成リクエストの例：

curl -X POST "http://127.0.0.1:30000/generate" \
     -H "Content-Type: application/json" \
     -d '{
        "text": [
            "<role>SYSTEM</role>detailed thinking off<|role_end|><role>HUMAN</role>Write the number from 1 to 128<|role_end|><role>ASSISTANT</role>",
            "<role>SYSTEM<"