Moondreamのブログ「Popping the GPU Bubble」は、AI推論を速くする話です。
ただし、単に「GPUを強く使えば速くなる」という雑な話ではありません。むしろ逆で、GPUがサボってしまう“すき間時間”をどう消すかに真正面から向き合っています。ここが面白い。
AIの推論では、GPUが大量の計算をして次のトークン（ざっくり言うと、文章のかたまり）を出します。ところが実際には、GPUが計算している時間よりも、CPUが「次は何をやるか」を準備している時間のせいで、GPUが待たされることがあります。これが記事のいう GPU bubble です。風船みたいに膨らんだ無駄な待ち時間、という比喩ですね。
Moondreamはこの待ち時間を、pipelined decoding という方法で減らしています。要するに、
「前のトークンの後片付けをCPUがしているあいだに、GPUは次のトークンの計算をもう始める」
という並行処理です。これだけ聞くと当たり前に見えますが、実際に安全にやるのがかなり難しい。そこでこの記事では、その工夫をかなり丁寧に説明しています。
AIモデルの文章生成は、1トークンずつ順番に進みます。
「3つ目の単語」は「2つ目」がないと決められないので、先読みしすぎることができません。これを autoregressive と呼びます。簡単に言うと、前の結果を見て次を決める方式です。
この方式だと、毎回「GPUが計算 → CPUが結果を受け取る → CPUが次の準備 → GPUがまた計算」という往復が起きます。
GPUの計算そのものは重いのですが、1トークン分の計算は案外小さい。すると、CPU側の事務作業が相対的に重く見えてきます。これが困る。せっかく高級なGPUを積んでいるのに、肝心のGPUが待ちぼうけになるわけです。もったいないにもほどがあります。
Moondreamの答えは、処理の順番を少しずらすことです。
従来は、CPUが前の処理を完全に終えてから次のGPU計算を始めます。これだと直列です。一直線。わかりやすいけれど遅い。
それに対して pipelined decoding では、前のトークンの後処理をしている最中に、次の forward をGPUへ投げる。
こうすると、CPUの雑務がGPUの計算の裏に隠れます。表から見ると待ち時間が減ったように見える。実際にGPUの利用効率が上がる、というわけです。
記事によると Photon はこれで decode throughput を最大35%向上させているとのこと。
この数字はかなり大きいです。単純なアルゴリズムの改善というより、GPUの使い方そのものを磨いた結果の勝利だと思います。
ここが技術記事としていちばんおもしろいところです。
「じゃあ次の計算を先に投げておけばいいじゃん」で終わるほど、実装は甘くありません。
GPUで1ステップを回すには、入力や出力を置くための作業領域が要ります。
Moondreamはこれを DecodeSlot と呼んでいます。しかも、毎回その場で確保すると遅いし、GPUの処理を止めかねないので、あらかじめ固定で確保して使い回します。CUDA graph を使う都合でも、アドレスが毎回変わらないほうが都合がいい。
ただし、1セットしかないと次の処理が前の結果を上書きしてしまう。そこで 2つのslotを交互に使う ping-pong にする。
片方をGPUが使っているあいだ、もう片方をCPUが読む。こうすると衝突しません。
この「2つに分ける」という発想は、地味だけどかなり好きです。派手な魔法ではなく、事故を防ぐための堅実な工夫だからです。現場っぽさがある。
Moondreamは、画像から点を返したり、箱を返したりするような constrained decoding を使います。
これは「モデルが何でも自由にしゃべる」のではなく、「出していいトークンを制限する」やり方です。たとえば座標を返す場面なら、数字や記号の並びだけを許す、といった具合です。
ここで厄介なのは、次に何を出してよいかのルールが、直前の結果に依存することです。
つまり、GPUの計算だけ先に回せても、最終的なサンプリングは前のトークンの確定を待たないと正しくできない。
そこで記事では、
これを commit-before-finalize と呼んでいます。
要するに、GPUが「候補を出す」ところまでは先にやっておき、何を採用するかの最終判断だけ少し後ろにずらすわけです。こうすると、GPUの重い仕事は前に押し出され、CPUの都合は裏に隠れます。
これも実装者が頭を抱えそうな話です。
次のステップをGPUに投げる時点では、まだ前のステップの結果が完全には確定していないことがあります。ところが、その直前で「このリクエストは終了です」と判明する場合がある。すると、もう始めてしまったGPUの仕事を途中で止めるのは大変です。
Moondreamは、こういう半端な状態のリクエストを zombies と呼んでいます。名前の付け方がちょっと好きです。確かに、もう終わったのにワークロードの中に残っている感じはゾンビっぽい。
ここでは、
結果として、先に終了扱いにして、あとで安全に解放する形になります。
これ、地味ですがかなり賢い。途中キャンセルの例外処理を山ほど書くより、参照カウントで自然に片づくようにしたほうが、たしかに強いです。
記事は decode だけで終わりません。
実運用では、まず入力プロンプトや画像をまとめて処理する prefill があり、そのあとに1トークンずつ生成する decode が続きます。prefill は一回が重いので、これもボトルネックになりやすい。
Moondreamは、この prefill を別物として扱わず、同じ2-slot pipeline に流し込んでいるのが面白いです。
つまり、GPUが decode を処理している裏で prefill を走らせたり、その逆をやったりできる。短いリクエストが多い場合、実はほとんどの時間が prefill と受付処理に消えるので、ここをうまく重ねる効果は大きいはずです。
個人的には、ここがかなり実務的で好きです。
「理論的にきれい」より、「本番でちゃんと効く」ことを優先している匂いがするからです。
このブログは、単に「Moondreamは速いです」という宣伝ではありません。
むしろ、速さはGPUの性能だけでは決まらない、という話です。
GPUは強い。でも、CPUとの連携が雑だと、その強さはすぐ目減りする。
逆に言えば、GPUの計算を少しでも止めないように設計すれば、同じハードでもかなり伸びる。Photon はその代表例だと思います。
しかも面白いのは、ここでやっているのが「新しいモデル」ではなく、実行系の設計だということです。
AIの世界はどうしてもモデルの大きさやベンチマークの数字が目立ちますが、実際にサービスとして速く・安定して動かすには、こういう泥臭い工夫が効く。私はこの手の話がかなり好きです。派手さはないけれど、最後に勝つのはこういう地味な最適化だったりします。
