GPUバブルを弾けさせる

← 全ブログに戻るMoondream EngineeringGPUバブルを弾けさせるMoondreamの推論エンジンであるPhotonは、VLM推論においてほぼリアルタイム（NVIDIA B200で約33ミリ秒）を実現しています。これは、GPUの動作を最適化することで、どのように最大35%高いデコードスループットを実現しているかを垣間見るものです。 2026年6月4日 AIモデルを可能な限り高速で実行するにはどうすればよいでしょうか？これはMoondream HQで私たちが常に頭を悩ませている問題です。GPUはモデル推論に関わるすべての計算を処理するため、一見すると、やるべきことは多くないように思えます。ただ、GPUに何をすべきかを指示し、答えを待つだけのように見えます。しかし、内部で実際にどのように機能しているかを見ると、GPUはしばしばアイドル状態になっていることがわかります。これは作業がないからではなく、CPUが次に何をすべきかをまだ指示していないためです。この現象はGPUバブルと呼ばれます。典型的なAIモデルがテキストを生成する場合、一度に1つのトークンを生成します（トークンとはテキストの塊で、おおよそ数文字です）。各トークンは前のトークンに依存します。これは自己回帰的（autoregressive）と呼ばれる特性で、そのため生成は順次行われます。2番目のトークンがなければ、3番目のトークンを計算することはできません。このデコードループには、CPUとGPU間の往復が含まれます。GPUは、次のトークンを生成するために数十億の算術演算を実行するなど、実際のモデルを実行するためのほとんどの重い処理を行います。しかし、CPUによっても驚くほど多くの作業が行われます。CPUは次に実行するリクエストを選択し、GPUが必要とするメタデータを設定し、モデルの出力から実際のトークンを選択して記録するなど、さまざまな処理を行います。課題は、1トークン分のGPU作業は小さいのに対し、CPUのハウスキーピングは各往復で発生する固定コストであることです。もしGPUが次のトークンを開始する前にそのハウスキーピングを待たなければならない場合、ループの一部でアイドル状態になります。これがGPUバブルが発生する理由です。この投稿では、Photonがパイプラインデコーディングと呼ばれる手法を使ってこれらのバブルをどのように隠しているかについて深く掘り下げていきます。この考え方は、2種類の作業をオーバーラップさせることです。つまり、CPUが最後のトークンを終えている間に、次のトークンのGPU作業を開始するのです。 バブル これが問題の形です。ブロッキングバージョン（上）では、すべてのステップがバトンパスです。CPUはフォワードを計画して起動し、GPUがそれを実行し、その後CPUが同期し、結果が着地するのを待ち、それらをコミットし、その後にようやく次のステップの計画を開始します。これは、計画が私たちが選択するトークンに依存するためです。例えば、モデルが応答を終えたことを示した場合、キューから新しい保留中のリクエストをスケジュールする必要があります。GPUは、CPUがコミット-計画-起動の作業を終えるのを待ってアイドル状態になります。 解決策は、ループをパイプライン化することです。現在のステップのトークンがまだ戻ってきてコミットされている間に、次のフォワードを起動します。これがパイプラインバージョン（下）です。フォワードは連続して実行され、CPUの作業はそれらの下にオーバーラップされます。私たちがこれを行える理由は、サンプリングしたばかりのトークンがGPUを離れる必要がないからです。次のフォワードは、その入力をGPUメモリから直接読み取ります。最終的には、トークンをデトークン化し、ストリーミングし、リクエストが完了したかどうかを判断するために、CPUにコピーが必要ですが、これは、次のフォワードがすでに実行されている間に、少し後でバックグラウンドで行うことができるハウスキーピングです。このコピーを待たないことが、バブルを取り除く動きです。これを安全に行うには、残りの部分で説明する3つのことが必要です。ステップバッファの衝突を防ぐ（ピンポン・スロット）、制約付きデコーディングでサンプリング順序を正しくする（すぐにフォワードし、後でサンプリングする）、リクエスト完了後のクリーンアップ（ゾンビ）です。 メカニズム1：ピンポン・スロット デコードステップを実行するために、GPUは一連の作業バッファを必要とします。入力（最後に生成されたトークンとそのシーケンス内の位置）をステージングする場所、モデルが出力（ロジット、語彙内の単語ごとに1つのスコア）を書き込む場所、サンプリングされたトークンを着地させる場所、およびアテンションカーネルが各シーケンスのキャッシュされたキーと値（KVキャッシュ）を見つけるために必要ないくつかの簿記情報です。私たちは両端にピン留めされた（ページロックされた）ホストバッファを保持しているため、GPUとの間のコピーはCPUをブロックするのではなく、バックグラウンドのDMA（ダイレクトメモリアクセス）転送として実行されます。これらのバッファは一度割り当てられ、すべてのステップで再利用されます。ランタイムにGPUメモリ割り当てを実行することを避けるために努力しています。なぜなら、それらはデバイスの同期を引き起こし、バブルを導入する可能性があるからです。固定されたバッファアドレスは、デコードステップを一度CUDAグラフとしてキャプチャし、それを再生するためにも必要であり、カーネル起動のオーバーヘッドを削減します。私たちはこのバンドルをDecodeSlotと呼んでいます。 これは機能しますが、パイプライン処理の阻害要因を導入します。バッファはステップが完了するまで使用されたままになるため、現在のステップが完了するまで次のステップを開始できません。2つのステップをオーバーラップさせるには、2番目のステップが独自の作業セットを必要とします。そうでなければ、CPUが読み取る前に最初のステップの結果を上書きしてしまう可能性があります。そのため、2つのスロットを保持し、ピンポン方式で交互に使用します。起動について注意すべき点が1つあります。CPUから起動を発行した瞬間にカーネルを実行するわけではありません。代わりに、それらをストリーム（GPUが順番に処理する順序付きキュー）にキューイングします。同じストリーム上の作業は順次実行されますが、別々のストリーム上の作業はオーバーラップできます。両方のスロットは同じ計算ストリームにフォワードを配置します。スロットはGPU並列処理のためではありません。CPUが1つのスロットの結果を処理している間に、GPUがもう1つのスロットのフォワードを実行できるようにするためにのみ存在します。フォワードはすべてその1つの計算ストリームを共有しますが、コピーは共有しません。各ステップのデバイスからホストへのコピー（簿記のためにサンプリングされたトークンを戻すもの）は別のコピー ストリームで行われるため、GPUが次のフォワードでビジーな間に実行できます。これが、私たちがそれを待たないことを可能にするものです。私たちは、ステップの出力が書き込まれた瞬間に記録されたイベントにコピーを固定します。そのため、そのステップの作業に正確に待ち、その後にキューに入れられたものには待ちません。スロットは、結果が読み取られた後、GPUが処理を終えた後だけでなく、結果が読み取られた後にのみ解放されます。そのピン留めされたホストバッファは、まだ転送中のコピーの着地場所であるため、スロットを新しいステップに早すぎに渡すと、転送途中のコピーを上書きしてしまい、デバッグが難しい破損バグを引き起こす可能性があります。そのため、スロットはそれを読み取るコミットを通じて予約されたままであり、そのコミットが完了した後にのみ解放されます。 メカニズム2：すぐにフォワードし、後でサンプリングする 次のフォワードは、CPUが最後のトークンで何をするかに依存しないため、先行して実行できます。しかし、次のステップに関する2つのことは、最後のステップのコミットされた結果に依存します。1つは、バッチ内にまだどのシーケンスが存在するかです。リクエストが完了した場合、次のフォワードには含まれるべきではありません。これは次のセクション（ゾンビ）で扱います。もう1つは、次のステップがサンプリングを許可されるトークンであり、これはこのセクションで扱います。これは制約付きデコーディングから来ています。Moondreamの空間スキルは、自由なテキストではなく、構造化された出力を返します。pointは座標を返し、detectはボックスを返し、segmentはアウトラインを返します。これらは、各ステップでモデルが生成できるトークンを制限することで、同じデコードループから得られます。許可されないトークンのスコア（ロジット）をサンプリング前に負の無限大に強制します。pointステップは座標を出力する必要があり、detectリクエストはx、y、サイズのサイクルをたどり、といった具合です。許可されるトークン（マスク）は、これまでに生成されたものに依存するため、ステップt+1のマスクは、tでサンプリングしたトークンに依存します。この依存関係はサンプリングにあり、フォワードにはありません。 各スケジューラティックは、起動、コミット、完了の3つのフェーズを通過します。 t+1のフォワードを起動します。これはマスクに依存しないため、すぐに実行されます。 ステップtをコミットします。進行中のコピーを待ち、リクエストのデコード状態を進めます。これはt+1のマスクを決定するために必要です。 t+1のサンプリングを完了します。状態が最新である状態で、マスクを構築し、サンプリングします。 t+1のサンプリングは、tのコミットがt+1のマスクを正しくするため、tのコミット後に実行されます。これを「コミット優先完了」順序と呼びます。GPUはステップ2と3を通じてt+1のフォワードを実行するため、コミットはクリティカルパスから消えます。プレーンテキストの場合、マスクはないため、フォワードとサンプリングの両方を1ステップ先行して実行できます。