SGLang-Diffusion：プロダクショングレードのビデオ生成のための高度な最適化

2ヶ月の進捗更新に続き、私たちはSGLang-Diffusionをプロダクショングレードのビデオ生成フレームワークとして深層最適化したことを喜んでお伝えします。これらの改善はスケーラビリティ、効率性、安定性を重視しており、大規模な拡散モデルのデプロイメントにとって重要です。

概要

ビデオ生成モデルの複雑性が増す中、私たちは推論パイプライン全体における重要なボトルネックを特定し、対処しました：

• インテリジェント並列化：トークンレベルのシーケンス分割と並列フォールディングにより、最適なリソース利用を実現
• 分散VAE：並列エンコード/デコードにより、高解像度ビデオのメモリボトルネックを解消
• プロダクショングレードサービング：Fixed Cache-DiT統合のバグを修正し、安定したマルチリクエストサービングを確保
• 最適化されたI/O：ビデオ保存を高速化し、不要なシリアライゼーションを排除
• 融合カーネル：LayerNormバリアント用のカスタムJITカーネルにより、GPUのアイドル時間を削減

以下、技術的詳細について深く掘り下げます。

主要な改善点

1. SP-シャーディングの改善：フレームレベルからトークンレベルへ

Video DiTモデルの入力テンソルの形状は通常B, T, H, W, Cで、一般的な設定num_frames=81では1, 21, 90, 160, 3となります。

8×H100構成で、Ulysses Sequence Parallel (N=8)を使用する場合、フレームワークは非アテンション操作でシーケンス次元に沿って分割し、その後all-to-all通信でアテンションヘッドの次元分割に切り替える必要があります。

従来の手法：フレームレベル分割

初期の実装ではT（時間）次元に沿って直接分割していました。しかし21フレームは8 GPUで均等に分割できないため、2つの最適でない方法が生じました：

1. フレーム数の調整：前処理時にnum_framesをNで割り切れるように修正
2. トークンパディング：時間次元をNの倍数にパディング（21 → 24）

フレームレベルのパディングは大きなオーバーヘッドを導入します：各パディングトークンには追加のH × W × C計算が必要です。

新しい手法：トークンレベル分割

パディングのオーバーヘッドを最小化するため、現在はT × H × Wを単一のシーケンス次元に平坦化してから分割しています。主な利点：

• パディングの削減またはゼロ化：一般的な解像度とVAE設定では、H × Wは8で割り切れることが多く、パディングが完全に不要
• 通信量の削減：パディングが必要な場合でも、オーバーヘッドはフレームレベルよりはるかに小さい

比較：形状と通信量の分析

この最適化により、より高速な通信とより小さなメモリ占有を実現し、特にビデオモデルに適しています。詳細は関連PRを参照してください。

2. Parallel Folding：テキストエンコーダーとDiTの並列化の分離

元の実装では、Text EncoderとDiTは同じTensor Parallel (TP)グループを共有していました。DiTがSequence Parallel (SP)のみを使用する場合、Text EncoderはTP=1で実行され、各GPUが完全なモデルのコピーを保持し、メモリと計算を無駄にしていました。

両者の計算が完全に分離されていることを考慮し、Parallel Foldingを導入しました：Text EncoderはDiTのSPグループをそのTPグループとして使用します。

実際の効果：

• テキストエンコーダー：SPグループ全体にTPを適用し、速度を最大化しメモリを削減
• デノイザー：SPを適用してシーケンス処理のスループットとメモリを最適化

両者が互いに干渉しないことを確保し、最適な並列化戦略を使用して、全体的な効率を向上させます。詳細は関連PRを参照してください。

3. Parallel VAE：分散エンコード/デコード

VAEのエンコード/デコードには密な3D畳み込みが含まれます。高解像度ビデオでは、単一GPUの実装は遅く、OOMになりやすいです。

一般的な緩和方法：

1. チャンク化：特徴マップのチャンクを順次処理し、ピークメモリを削減するがレイテンシが増加
2. 並列化：GPU間で並行してチャンクを処理し、メモリとレイテンシの両方を削減

Wan-VAEのためにParallel VAEを実装しました：

• 高さ分割：高さ次元に沿ってRank間で特徴マップを分割
• Conv操作：halo_exchange（P2P）を使用して境界ピクセルを共有し、数学的等価性を確保
• アテンション操作：必要に応じてall_gatherでグローバルコンテキストを取得
• 結果の集約：エンコード/デコードの最後でall_gatherして全高さを再構築

VAEのボトルネックを解消し、OOMなしでより高い解像度とより長いシーケンスをサポートします。

4. Cache-DiTサービング：マルチリクエスト安定性の修正

Cache-DiTは残差をキャッシュして冗長な計算をスキップすることで推論を高速化しますが、SCMの正しいnum_inference_steps設定に依存します。

問題：Wan2.2のデュアルTransformerアーキテクチャでは、transformerとtransformer_2がそれぞれnum_high_noise_stepsとnum_low_noise_stepsを実行します。初期実装には2つのバグがありました：

1. 両方とも誤って総num_inference_stepsでキャッシュコンテキストを設定
2. サービスモードでは、キャッシュがリクエスト間で永続化され、異なるnum_inference_stepsにより形状の不一致が発生しサーバーがクラッシュ

解決策：

1. num_high_noise_stepsとnum_low_noise_stepsを使用して独立したコンテキストを設定
2. 新しいリクエストごとにタイムステップ分割を再計算し、Cache-DiT APIを利用してコンテキストを更新し、完全に分離

安定したプロダクショングレードのCache-DiT高速化サービングを確保します。

5. ビデオ保存の最適化：シリアライゼーションのオーバーヘッドの削除

サービスアーキテクチャでは、scheduler_clientとgpu_workerがZMQを介して通信します。

以前：gpu_workerは推論後にテンソルをシリアライズし、ZMQで送信し、scheduler_clientがデシリアライズして保存し、シリアライゼーション/デシリアライゼーションとメモリコピーのオーバーヘッドを導入していました。

新しい方式：gpu_workerが出力テンソルを直接処理してビデオを保存し、scheduler_clientにはファイルパスのみを返します。

シリアライゼーションのオーバーヘッドを排除し、重複コピーを回避します。

6. WanVideo LayerNorm融合：CuTeDSL JITカーネル

WanVideoは2つのLayerNormモードを導入しました：

1. LayerNormScaleShift：y = LN(x) * (1 + scale) + shift
2. ScaleResidualLayerNormScaleShift：residual_out = residual + gate * x；y = LN(residual_out) * (1 + scale) + shift

個別のカーネル実装は、複数の起動と中間メモリトラフィックを引き起こし、GPUのアイドル時間を生成します。

解決策：CuTeDSLを使用して融合JITカーネルを実装（sglang/jit_kernel/diffusion/cutedslに配置）し、単一カーネルで操作を融合します。

利点：

• カーネル起動のオーバーヘッドを削減
• メモリトラフィックを削減
• GPU利用率とスループットを向上

多層WanVideoでのマイクロ最適化が蓄積されて顕著な効果をもたらします。

パフォーマンス結果

SGLang-DiffusionとLightX2VのWan2.2 T2Vでの異なる設定下での比較。

今後の展望

私たちは拡散モデルサービングの境界を継続的に押し広げています。詳細はSGLang-Diffusion 2026Q1ロードマップをご覧ください。

プロダクションデプロイメントの最適化に関する更なる更新にご期待ください。

謝辞

• 貢献者への感謝：Skywork.ai、Song Rui、SGLang-Diffusionチーム
• 計算パートナーのサポートに特別な感謝

SGLang-Diffusionを活用した拡散生成を試す：APIFree

詳細情報

• Slackチャンネル：#diffusion（参加：slack.sglang.io）
• SGLang-Diffusion Cookbook
• SGLang-Diffusionドキュメント
• 前回の更新：2ヶ月の進捗