GPU推論アクセラレーションの次のステップは、カーネル融合です。
編集：鄭佳美
編集：馬曉寧
スタンフォード大学のHazy Researchチームは、Llama-3.2-1Bオープンソースモデルのフォワード推論を「Megakernel」として統合し、低遅延推論能力を極限まで引き上げた画期的な最適化成果を発表しました。
対話型AIや人間が関与するインタラクティブなワークフローなど、リアルタイム性が極めて高い特定のアプリケーションでは、大規模言語モデルの応答速度は重要であるだけでなく、ユーザーエクスペリエンスの成否を決定することさえあります。
チームはLLM推論速度のボトルネックがメモリロードにあると考えています。彼らの研究によると、既存のオープンソース推論エンジン（vLLM、SGLangなど）は、H100のようなトップクラスのGPU上でも、超低遅延のシングルシーケンス生成タスクではメモリ帯域幅の50%未満しか利用できていませんでした。
これは主に、各Transformerモジュール層が数十から数百のCUDAカーネルに分解され、各カーネルが非常に小さな操作（RMS norm、Attention、MLP、Rotary Position Embeddingなど）を実行するため、それらの間に大量のコンテキストスイッチと待機が発生するためです。
さらに深刻なのは、これらのカーネルの起動と終了にかかるコストが、CUDA GraphやPDL（Programmatic Dependent Launch）などのメカニズムによって十分に隠されず、短時間のタスクではむしろ増幅されることです。
言い換えれば、GPUは「作業中」ではなく「作業を待つ」ことに多くの時間を費やしています。Hazyチームの研究は、まさにこの問題を中心に展開されています。
1 Megakernel：ゼロから設計された融合アプローチ
まず、実験結果から言うと、MegakernelはH100上での推論遅延を1ミリ秒未満にまで短縮し、メモリ帯域幅利用率は78%に達しました。これはvLLMと比較して2.5倍、SGLangと比較して1.5倍の改善です。さらに先進的なB200プラットフォームでは、遅延は600〜680マイクロ秒までさらに短縮され、理論上の限界に近づいています。
1回の完全な推論の時間配分を見ると、アクティベーションの保存、整合性待ち、データロードに250マイクロ秒、RMSNormとmatvec（matvecが95%を占める）に200マイクロ秒、重みロードはわずか30マイクロ秒で、パイプラインメカニズムは安定したパフォーマンスを示しています。ワープ間の同期とバリアは40マイクロ秒の遅延をもたらし、セットアップ、パラメータ渡し、ページステータスマーキングなどのその他のオーバーヘッドは合計で約80マイクロ秒です。
全体として、HazyチームのMegakernelは、綿密なスケジューリングの下で、現在のハードウェア性能をほぼ限界まで引き出しています。
これらの効果を得ることができたのは、Hazyチームが提案した、野心的だが効果的な設計思想、すなわち、フォワードプロパゲーション全体を単一のCUDAカーネル（いわゆるMegakernel）に統合することによるものです。
実験では、既存のThunderMLAアーキテクチャに基づき、GPU上で動作する軽量の「命令インタプリタ」システムを開発しました。このシステムは、各ストリーミングマルチプロセッサ（SM）に「実行計画」を事前に割り当てます。この計画には、順番に並んだ複数の命令が含まれており、各命令はTransformerモデル内の構造単位を表します。
これらの命令には以下が含まれます：
RMSNorm、QKVプロジェクション、RoPEを融合した複合命令；
Attention行列乗算と削減計算（長系列GQAをサポート）；
O-プロジェクションと残差加算；
MLPのRMSNorm、ゲートアクティベーション（SiLU）、アッププロジェクション；
ダウンプロジェクションと最終的な残差；
最終層のRMSNorm + 言語モデリングヘッド。
各命令は統一されたCUDAテンプレートに基づいて構築され、ロード、ストア、計算の標準化されたカプセル化を実現しています。命令間の依存関係は、インタプリタが実行前に静的に配置し、各SMは同じスケジュールを繰り返し再利用して複数のトークンを処理できます。
さらに、効率的なデータパスを確保するため、インタプリタはこれらの実行計画をモデル構造に応じて静的に編成し、スケジューリング時の動的な分岐を避け、スループットと並列実行能力を向上させています。
また、パイプライン計算を実現し、共有メモリの競合を防ぐため、チームはGPUの共有メモリをページング管理しました。例えば：
最初の213KBの共有メモリを13個の16KiBページに分割；
残りの部分は命令パラメータ、ページ割り当て情報などを格納；
各命令はロード前に明示的にページを要求し、完了後にインタプリタスケジューラに返却；
ページが解放されると、インタプリタは直ちに次の待機中の命令にそれを割り当てます。
このメカニズムにより、次の計算段階が可能な限り早く重みのプリロードを開始できることが保証され、帯域幅利用率が最大化され、「バブル」が排除されます。
しかし、Megakernelの構造は、従来のカーネル間の暗黙的な同期に依存できません。そのため、Hazyチームはカウンターシステムも使用しました。彼らはグローバルメモリに一連の整数を保持し、各命令が完了するたびに、対応するカウンターを+1します。ある命令が以前のステップの結果に依存する場合、その命令はカウンターが特定の値に達するまで待機してから実行されます。
例えば、MLPのダウンプロジェクション段階では、チームは中間状態を4つのチャンクに分割し、各チャンクは書き込み後すぐに後続の計算をトリガーすることで並列フローを実現します。さらに、チームは依存関係グラフを正確に設定することで、グローバルバリアを回避し、命令間の無駄な待機を大幅に削減し、カーネル全体の実行を理論上の並列性に可能な限り近づけました。
さらに、研究チームはCUDA非同期バリアの性能も測定し、バリアが「通過済み」の状態であっても、毎回60nsが必要であり、同期操作のコストは無視できないことを発見しました。実際の実行、特にmatrix-vector（行列とベクトルの乗算）のような重要な操作では、Hopperアーキテクチャ（H100など）では通常のCUDAコア（Tensor Coreではない）がより効果的である一方、BlackwellアーキテクチャではTensor Coreの性能が優れていることが判明しました。
これは、ハードウェアの世代が異なると、Megakernelの最適な実装パスもマイクロアーキテクチャの違いに適応すべきであり、単一のソリューションがすべてのプラットフォームで通用するわけではないことを示しています。
その他のコンテンツは、以下をクリックしてフォローしてください：
「AI科技評論」の許可なく、ウェブページ、フォーラム、コミュニティでのいかなる転載も固く禁じます！
公式アカウントでの転載は、「AI科技評論」のバックエンドで事前にメッセージを残し許可を得てください。転載時には出典を明記し、この公式アカウントの名刺を挿入してください。
// おすすめの読み物
UCL強化学習派：汪軍とその学生たち
大規模モデルの隠れたプレイヤーが登場：DeepSeekは左へ、面壁は右へ
基盤大規模モデル「六進二」：南のJiedian、北のZhipu