Robloxにおける3D制作のためのAI推論の高速化

• Robloxは、より応答性の高い反復処理を実現するため、3Dメッシュ生成を高速化するためにCUDA GraphsとKVキャッシュを導入しました。
• リリース当初、Cube 3Dモデルはトークンを7.8ミリ秒（従来は60.5ミリ秒）で生成でき、完全なオブジェクトを4秒（従来は31秒）で生成できるようになりました。 

今年初め、RobloxはCube 3D基盤モデルの最初の機能を公開しました。Cube 3Dを使用することで、クリエイターはテキストプロンプトから直接3Dモデルや環境を生成できます。 当初から、生成時間の遅延は本質的に反復的なプロセスを妨げるものであると認識し、レイテンシの最適化を最優先してきました。3月にCube 3Dがリリースされる前から、開発者とユーザーの両方にとって、推論ステップを7.8倍高速化し、応答性を向上させていました。 

リリース以来、いくつかの注目すべき体験（エクスペリエンス）において、57万8,000個以上のオブジェクトが生成されました。また、開発者からは、「猫」や「ハンバーガー」などのテキストプロンプトを使って、体験内でユーザーが3Dオブジェクトを生成できるようにしたいという要望も寄せられました。特に、ボイスチャットを使った人気の交流ゲーム「Mic Up」では、Cube 3Dを活用し、プレイヤーが楽しくインタラクティブにオブジェクトを生成できる仕組みを実現しました。 この実装では、プレイヤーは左側のメニューを開き、AIアイコンを含む追加機能にアクセスできます。そのアイコンをクリックすると、テキストプロンプトを入力して3Dオブジェクトを生成できます。ユーザーにとって、生成時間が長くなると操作の妨げとなり、自分のアイデアがリアルタイムで3Dに変換されるという魔法のような体験が損なわれてしまいます。 

私たちは、3D生成体験を「停止して待つ」ようなインタラクションから、レスポンスが良く自然な感覚のものへと変革し、迅速な実験を可能にしたいと考えました。シーンにオブジェクトを素早く追加できる機能は、開発者にとって極めて重要です。Cube 3Dを高速化するため、まず推論パイプラインのプロファイリングを行い、パフォーマンスのボトルネックを特定しました。高性能なGPUを使用しているにもかかわらず、操作の合間にかなりのアイドル時間が発生していることが判明しました。

現代のディープラーニングフレームワークは、GPU上での演算（またはカーネル）のスケジューリングと実行をCPUに依存しています。CPUは各演算を準備し、それをGPUに送信し、確認を待ってから次の演算の準備を行います。この待機時間がスケジューリングのボトルネックとなり、CPUが次のバッチを準備している間、GPUがアイドル状態になる可能性があります。理想的には、CPUがGPUよりも先行して動作し、演算を準備してキューに追加することで、GPUが常に処理すべき仕事を持っている状態を維持したいところです。

これは、Cube 3Dのようなトランスフォーマー型モデルの自己回帰デコーダーにとって特に問題となります。これらのモデルは、入力を処理し、トークンを順次生成する必要があるためです。これらのモデルでは、1回の生成に数千もの個別の演算が必要であり、シーケンスの各ステップごとに計算オーバーヘッドが蓄積されます。

「私たちは4次元的なインタラクションを可能にするものを構築したかったのです」と、エンジニアリング担当副社長のAnupam Singh氏は、Robloxが自己回帰型アプローチを選択した理由をこう説明しました。「単に車を作るだけでなく、車のドアを開けて中に入ることができるようにしたかったのです。」

各操作にかかる：

• 各カーネルの準備にかかるCPU時間
• カーネルの起動に伴うオーバーヘッド
• GPU実行時間（実際の計算）
• 完了確認時の同期オーバーヘッド

GPU上で高速に実行される小規模な演算の場合、このオーバーヘッドが推論時間の大部分を占める可能性があります。GPUが実際に計算を行っている時間は、推論時間全体のごく一部に過ぎない場合があります。

このボトルネックに対処するため、我々はCUDA Graphsを活用しました。これは、CPUの介入なしにGPU操作のシーケンスを記録・再生できる機能です。Cube 3Dのアーキテクチャにおける自己回帰デコーダコンポーネントは、テキストプロンプトを処理し、固定長ベクトルを通じて形状トークンを生成します。 

機能的には従来の大規模言語モデル（LLM）と似ていますが、当社のデュアルストリームデコーダアーキテクチャには重要な違いがあります。それは、2つの並列アテンションストリームを使用している点です。一方のストリームは条件トークン専用、もう一方は形状トークン専用です。市販のLLM推論エンジンは当社のニーズに適しておらず、特定のアーキテクチャに合わせたカスタム実装が必要でした。

CUDA Graphsは、GPU向けのマクロを記録するようなものだと考えてください。CPUが各コマンドを個別に発行する代わりに、GPU操作のシーケンス全体（グラフ）を記録し、単一のCPU命令でグラフ全体を起動します。このアプローチにより、推論中にCPUが各操作を個別にスケジューリングする必要がなくなるため、カーネル起動のオーバーヘッドが劇的に削減されます。グラフが起動されると、GPUはさらなる指示を待つことなく、シーケンス全体を自律的に実行します。

CUDA Graphsにはいくつかの制限があります。グラフ構造を事前に決定する必要があるため、固定のバッチサイズと入力次元が求められます。つまり、バッチサイズや入力形状ごとに個別のグラフを作成する必要があります。Cube 3Dでの当社のユースケースにおいては、一般的な入力次元に基づいて推論プロセスを標準化できたため、この制限は許容範囲内でした。

Cube 3DモデルにCUDAグラフを実装するには、アプローチを調整する必要がありました。従来のLLMでは、アテンション演算は常に同じシーケンス長で行われるため、処理対象となる形状は静的です。しかし、私たちのカスタムデュアルストリームアーキテクチャでは、シーケンス長のみを処理対象とするアテンション層もあれば、シーケンス長と条件長の組み合わせを処理対象とする層もあります。

こうした課題があったにもかかわらず、CUDA Graphsを実装した結果、目覚ましい成果が得られました。推論時の各トークンの生成時間を測定するために、出力トークンあたりの時間（TPOT）を使用しています。CUDA Graphsの実装後、TPOTは60.5ミリ秒から20.5ミリ秒へと改善し、2.9倍の高速化を達成しました。全体的な生成時間は31秒から10.5秒へと66%短縮されました。

レイテンシをさらに改善するため、業界全体で高い効果が実証されているLLM推論における標準的な手法であるKVキャッシュを実装しました。 

Cube 3Dのようなトランスフォーマーベースのモデルでは、各トークンの生成ごとに、それまでに生成されたすべてのトークンに基づいてキー（K）行列と値（V）行列を計算する必要があります。シーケンスが長くなるにつれ、トークンごとにこれらの行列を再計算することは、ますます非効率になっていきます。

KVキャッシュはこの問題を次のように解決します：

1. 以前に生成されたすべてのトークンに対する K 行列と V 行列を保存する
2. 新しいトークンに対してのみ K 行列と V 行列を計算する
3. これらの新しい行列をキャッシュされた値に追加する

このアプローチにより冗長な計算が排除され、新しいトークンごとに必要な処理負荷が軽減されます。これは、生成されるシーケンスが長くなるにつれて特に大きな効果を発揮します。

KVキャッシュをCUDA Graphの実装に統合する当社のアプローチは、従来のLLM推論と同様でした。KVキャッシュの導入により、TPOTはわずか7.8ミリ秒にまで短縮されました。全体的な生成時間は87%短縮され、当初の31秒からわずか4秒にまで短縮されました。この大幅な時間短縮により、このツールを利用するクリエイターにとって、その有効性が格段に高まりました。

最適化されたカーネル、さらに高速な推論を実現するモデルの量子化、ハードウェア固有の最適化、並列トークン生成など、レイテンシをさらに低減し、ユーザー体験を向上させる技術の研究を進めています。

この取り組みは、フルシーンの生成と理解へと拡大するとさらに複雑になります。そこでは、多くの3D要素がレイアウト内で互いに文脈に沿って連携する必要があります。また、私たちが作成する3Dオブジェクトやワールドには、ドアの開閉や車輪の回転など、完全な機能性を備えてほしいと考えています。それを実現するためには、シーン全体、完全に機能するオブジェクト、アバターへとスケールアップするための、迅速な生成と反復処理が必要です。 Cube 3D基盤モデルの拡張に伴い、さらなる改善点や新機能をご紹介できることを楽しみにしています。また、クリエイターコミュニティの皆様がこれらを活用して作り上げる没入感のある世界を見るのを楽しみにしています。