Moondream Photon: パイプライン化デコードによる GPU バブルの除去

Moondream の Photon 推論エンジンは、パイプライン化デコードという手法を実装することで、NVIDIA B200 GPU 上で最大 35% 高いデコードスループットを実現しています。このアプローチは、トークン生成ステップ間で CPU が必要なハウスキーピング作業を完了するのを待っている間に GPU がアイドル状態になる「GPU バブル」を排除します。

GPU バブルの問題

自己回帰型テキスト生成では、トークンが順次生成されます。各デコードループは通常、CPU と GPU の間の往復を必要とします。GPU はモデルのフォワードパスの重い演算を実行し、CPU は次のリクエストの選択、メタデータの設定、サンプリングされたトークンの記録といった重要なブックキーピングを担当します。

単一トークンに対する GPU の作業は比較的小さいため、CPU のブックキーピングの固定コストが大きなボトルネックになります。標準的なブロッキングループでは、GPU は次のトークンを開始する前に CPU がコミット・プラン・ランチサイクルを完了するのを待たなければならず、GPU バブルと呼ばれるアイドルギャップが生じます。

パイプライン化デコードの仕組み

Photon は CPU の作業と GPU 計算を重ね合わせることで GPU バブルを取り除きます。トークンが CPU にコミットされるのを待つ代わりに、Photon は前ステップの結果を CPU が処理している間に次の GPU フォワードパスを起動します。これは、次のフォワードパスがサンプリングされたトークンを CPU がデトークン化したりストリームしたりするのを待たずに、直接 GPU メモリから読むことができるため可能です。

安全に実装するために、Photon は以下の 3 つの主要メカニズムを使用します。

1. Ping-Pong スロット

CPU が結果を読み取る前に第 2 ステップが第 1 ステップの結果を上書きしないように、Photon は 2 つの作業バッファ（DecodeSlots）を「ピンポン」方式で交互に使用します。

• Compute Stream: 両スロットは単一のコンピュートストリームにフォワードパスをエンキューし、GPU 上での順次実行を保証します。
• Copy Stream: サンプリングされたトークンのデバイス‑ツー‑ホストコピーは別のコピーストリームに移動します。これにより、GPU が次のフォワードパスを実行している間に CPU がバックグラウンドで結果を読み取れます。
• Pinned Buffers: ページロックされたホストバッファを使用し、コピーがバックグラウンド DMA 転送として実行され、CPU のブロックを回避します。

2. Forward Now, Sample Later

空間タスク（座標やバウンディングボックスの返却など）で使用される制約付きデコードでは、モデルが生成できるトークンを制限するマスクが必要です。このステップ $t+1$ 用のマスクは、ステップ $t$ でサンプリングされたトークンに依存します。

Photon はこの依存関係を 3 つのフェーズに分割して解決します。

1. Launch: マスクを必要としないため、$t+1$ のフォワードパスを即座に起動します。
2. Commit: ステップ $t$ の結果をコミットし、$t+1$ のマスク決定に必要な状態を更新します。
3. Finalize Sampling: $t+1$ 用のマスクを構築し、トークンをサンプリングします。

この「コミット‑前‑ファイナライズ」順序により、GPU のフォワードパスは CPU がサンプリング制約を決定している間にバックグラウンドで走ります。

3. Zombie 管理

ステップ $t+1$ がステップ $t$ のコミット前に起動されるため、シーケンスがステップ $t$ で End‑of‑Sequence (EOS) トークンに達しても、物理的にはステップ $t+1$ のバッチに残っていることがあります。Photon はこれらを「ゾンビ」と呼びます。

複雑なキャンセルロジックを避けるため、Photon は参照カウント方式を使用します。

• finalized フラグ: シーケンスが EOS または長さ上限に達したときに true に設定されます。
• inflight_refs: まだインフライト中のステップがシーケンスを参照している数をカウントします。

コミット中にゾンビが検出された場合、コミットは単にスキップされます。シーケンスのリソース（KV ページや LoRA スロット）は inflight_refs がゼロになったときにのみ解放されます。

パフォーマンスへの影響とコストモデル

Photon の性能向上は、GPU が高速になるほど、またはモデルが小さくなるほど顕著です。CPU のブックキーピングコストは一定のままで、GPU フォワードパス時間が短縮されるためです。

ベンチマークされたスピードアップ

NVIDIA B200 で 32 ストリーム使用時、Photon はブロッキングデコードに比べてスループットが 35.4% 向上したことが観測されています。ハードウェアごとのスピードアップは以下の通りです。

ハードウェア	ストリーム数	Blocking (ms)	Pipelined (ms)	観測されたスピードアップ
RTX 3090	1	5.44	5.10	+6.5%
RTX 3090	32	11.74	10.52	+11.6%
B200	1	3.11	2.63	+17.6%
B200	32	5.55	3.98	+35.4%

ゾンビ税

先行実行には小さなペナルティ、いわゆる「ゾンビ税」があります。完了したシーケンスが不要なフォワードパスを 1 回実行することがあります。単一ストリームの場合、110 トークンのシーケンスで約 1% のオーバーヘッドになります。ただし、バッチ処理では GPU が重みのメモリ帯域に制限されているため、バッチに 1 行追加するコストはほぼ無視でき、影響はごくわずかです。

Prefill との統合

Photon はプレフィル（初期プロンプトと画像の処理）を 2 スロットパイプラインの別の起動として扱います。同じパイプラインを共有することで、新しいリクエストの高コストなプレフィルフォワードパスと、既存リクエストのデコードステップの CPU コミットを重ね合わせることができます。これは多数の短いリクエストがあるワークロードで特に効果的で、システムがデコードよりもプレフィルとアドミッションに多くの時間を費やす場合に有利です。

コミュニティの見解

技術実装は透明性が高く評価されていますが、最適化の効果はモデルサイズに大きく依存すると指摘する声もあります。ある批評家は、非常に大きなモデル（例: フォワードパスが 30‑40ms）では CPU‑GPU 同期オーバーヘッドが全体時間の割合が小さくなるため、MoE（Mixture of Experts）モデルにおけるカーネル最適化や通信スケジューリングの方が重要になると述べています。