SGL-Diffusion における AR+DiT フルスタックパフォーマンス最適化

TL;DR

HFバックエンドをSRTに置き換えてARモデリングを高速化し、並列処理の競合を解消。ARには専用TPを、DiTにはSPを採用する。動態バッチ処理によりハードウェア使用率を向上させ、処理済み画像の早期返却をサポート。デバイスごとに1つのdenoiserによるDiTの並列実行を実現し、AR結果をバッファリングすることでARとDiTのワークフローをオーバーラップさせる。

1. 背景

ハイブリッド自己回帰-拡散（AR+DiT）生成フレームワークは、ARによるグローバルコンテキストのモデリングとDiTによるローカル詳細の精緻化を組み合わせたものである。GLM-Imageはその典型的な事例であり、まず9Bの視覚言語モデルが自己回帰的にセマンティックな事前トークンを生成し、次に7BのDiTが30〜50ステップのデノイジングにより高解像度画像を生成する。

ネイティブなデプロイでは、ARエンコーダ・DiTデノイザ・VAEデコーダが単一プロセス上でチェーン状に動作するため、アーキテクチャの密結合、並行処理時の低いリソース使用率、リソース割り当てのミスマッチという3つの課題が生じる。

2. ARバックエンドのSRT化（PR #25381）

ARフェーズを拡散ワーカープロセスから切り離し、独立したSRTサービスとして展開することで、独立した重みのロードとスケジューリングをサポートする。ARサーバーはTPを個別に設定できるようになり、DiTのSP戦略による制約を受けなくなる。

パフォーマンスの向上は顕著であり、シングルカードのエンドツーエンドレイテンシは154.6秒から78.3秒へ（-49.4%）短縮され、4カード異種構成ではさらに35.2秒（-77.2%）まで低下した。

3. 動的バッチ処理と早期返却のサポート（PR #30683）

動的バッチ処理をGLM-Imageに拡張し、ARフェーズのみにバッチ処理を適用する。また、バッチ全体の完了を待たずに済むよう早期返却（early return）をサポートする。スループットはバッチサイズに応じて向上し、BS1の0.0388 img/sからBS16の0.1368 img/sへと増加する。

4. 分離型AR-to-DiTファンアウトアーキテクチャ（PR #31320）

2つのフェーズを完全に分離し、ARでは大バッチ+TPを採用、DiTはNPU単体でbatch=1が最適となる構成で動作させることにより、異種並列処理とワークフローのオーバーラップを実現する。