LLMをMegaKernelにコンパイルして低レイテンシ推論を実現する
(zhihaojia.medium.com)- LLM推論を単一のメガカーネルへ自動変換するコンパイラを開発
- MegaKernel（Persistentカーネル）方式は、LLM推論における計算と通信を完全に1つのGPUカーネルに統合し、非常に低いレイテンシを可能にする
- 既存のMLフレームワークやカーネルライブラリの分散構造により、パイプライン全体の単一カーネル化が非常に難しいという問題がある
- Mirage Persistent Kernel（MPK）は、コンパイラとランタイムシステムを通じて、マルチGPU LLM推論を自動的に高性能なmegakernelへ変換する
- MPKは演算グラフを細粒度のタスクグラフへ変換し、ソフトウェアパイプライニングと計算・通信のオーバーラップを最大化する
- MPKを適用すると既存システムと比べてトークン生成遅延が減少し、GPU数が増えるほど性能向上幅もさらに大きくなる
概要とMegaKernel方式の利点
- 大規模言語モデル（LLM）推論で遅延時間を減らす効果的な方法の1つは、すべての計算と通信の過程を単一の**megakernel（一貫型カーネル）**に融合する方式である
- この方式では、モデル全体のレイヤーごとの演算からGPU間通信まで、あらゆる処理を1つのGPUカーネルが途切れなく実行する
- 主な利点は次のとおり
- 繰り返しのカーネル呼び出しを省き、カーネル起動オーバーヘッドを除去
- レイヤー全体にわたってソフトウェアパイプライニングを実現可能
- 計算と通信を同時進行させることでレイテンシを隠蔽
従来の限界とMPKの登場
- 既存のPyTorch, Triton, TVMのようなMLフレームワークは、end-to-end megakernelの自動生成機能を本質的にはサポートしていない
- 実際のLLMシステムは、**NCCL/NVSHMEM（通信）、FlashInfer/FlashAttention（attention）、CUDA/Triton（カスタム演算）**など多様なカーネルライブラリの組み合わせで構成されており、単一カーネルへの統合が難しい
- こうした背景のもと、CMU、UW、Berkeley、NVIDIA、Tsinghuaの研究者らは**Mirage Persistent Kernel（MPK）**を開発した
- MPKはコンパイラとランタイムを組み合わせ、LLM推論の全パイプラインを自動的に高性能なmegakernelへ変換する
MPKの中核的価値
- MPKはカーネル起動オーバーヘッドを完全に除去し、層間の演算・データ読み込み・通信のオーバーラップを最大化することで、超低遅延のLLM推論環境を実現する
- 実際のテスト（39トークンのプロンプト、512トークン生成、speculative decoding不使用）では、
- NVIDIA A100 40GB GPU単体環境を基準に、vLLM/SGLangなど既存の最適化システムのトークンあたりデコード遅延（14.5ms）に対し、MPKは12.5msまで短縮した
- この数値は理論的下限（10ms）に近い（1.6TB/sメモリ帯域幅、16GBの重み読み込み基準）
- マルチGPU環境では、計算と通信を完全に統合することで、GPU数が増えるほどMPKの性能優位がより際立つ
MPKの動作構造の詳細
Part 1. コンパイラ – LLM演算グラフ → タスクグラフ変換
- 一般にLLMの演算は、各演算（例: 行列積、attention）または通信演算（例: all-reduce）をノード、データ依存性をエッジとする計算グラフで表される
- 従来設計では演算子ごとに別カーネルを実行する方式が一般的だが、これは実際の依存データ単位ではなくカーネル単位の依存性しか反映しないため、パイプライニングの機会が限られる
- 例: 行列積の後にall-reduceがある場合、all-reduceの実行は行列積全体が終わってからでないと始まらない。実際にはデータを分割し、部分実行や依存関係の活用が可能である
- MPKコンパイラは演算グラフを細分化し、実際のデータ単位に適したfine-grained task graphへ自動変換する
- 各タスク（四角形）は、個別のGPU SMに割り当てられる演算・通信単位である
- 各イベント（円）は、タスク間の同期ポイントである
- タスクおよびイベント間のエッジによって、効率的なデータ・制御依存性を表現する
- このタスクグラフにより、MPKでは計算と通信を部分的または並列に、より深くオーバーラップさせられる
- Mirage kernel superoptimizerにより、各タスクに適した高性能なCUDA実装も自動生成される
Part 2. ランタイム – メガカーネル内部でのタスクグラフ実行
- MPKランタイムは、タスクグラフをGPU上の1つのカーネル（メガカーネル）内部だけで完全に実行する方式である
- GPUのすべてのSM（Streaming Multiprocessors）を、静的にワーカーとスケジューラの役割へ分割する
ワーカー
- 各ワーカーはSM単位で動作し、専用のタスクキューを管理する
- ループ方式で
- 次のタスクをキューから取得
- 実行（例: matmul, attention, データ転送）
- 完了時にイベントへ通知
- 繰り返し処理
- これにより、各ワーカーのリソース利用を最適化し、非同期な階層演算を可能にする
スケジューラ
- 分散スケジューラが各SM内で単一warp単位で動作し、最大4つのスケジューラを同時実行できる
- 各スケジューラはアクティブ化されたイベントキューを管理し、条件を満たしたタスクをワーカーに割り当てる
- これにより、集中型の同期オーバーヘッドなしに大規模なタスク分散処理が可能になる
イベントベースの実行方式
- タスクが完了すると、特定のイベントカウンタを増加させる。カウンタがしきい値に達するとイベントが有効化され、スケジューラキューに挿入される
- スケジューラは、そのイベントに依存関係を持つ後続タスクを実行する
- これにより、細粒度のソフトウェアパイプライニングと計算・通信のオーバーラップが自然に実現される
- 例: あるレイヤーのmatmulと別のレイヤーのattentionを同時実行
- 部分的に完了したmatmulの結果が出た時点で、all-reduce通信を開始可能
- すべてのスケジューリングとタスク切り替えが単一カーネルコンテキスト内で行われるため、タスク間オーバーヘッドは1〜2マイクロ秒（μs）水準と非常に低い
今後の方向性
-
MPKの目標: 開発者が少量のPythonコード（数十行程度）を書くだけで、容易にLLMをmegakernelへコンパイルし、最大性能を発揮できるよう支援すること
-
主な発展方向
- 最新GPUアーキテクチャのサポート: たとえばNVIDIA Blackwell向け、warp単位に特化した方式など
- 動的workload処理: mixture-of-experts（MoE）など、動的な制御フローが必要なモデル向けのコンパイル戦略研究
- 高度なタスクスケジューリング: 優先度ベース、スループット最適化など、現代的なポリシーの研究と適用可能性の追求
-
MPKはGPUベースLLM推論のコンパイル・実行方式における根本的な転換点を示しており、コミュニティとの協力拡大を目指している
追加資料
- MPK（Mirage Persistent Kernel）のコードとドキュメント、最新の研究成果はGitHub（https://github.com/mirage-project/mirage）で確認できる
1件のコメント
Hacker Newsのコメント
著者へ。on-GPUインタプリタ方式が非常に有望な将来の方向性に見える点が興味深い。ほぼ同じアプローチを取る別の研究もあるので、関連投稿の参照を勧めたい。CUDAの根本的なプログラミングモデル（たとえばカーネル起動）が、細粒度のタスクベース並列化のために回避されているが、この方式がハードウェア活用率をより高めるのを実際に目にしている。CUDAがさまざまな面で私たちを縛っていたのではないかとも思う。著者の研究がPyTorchの実験的バックエンドに入る可能性にも期待している。それから、最初の部分の2つの段落がほぼ同一なので、軽微な誤植として指摘しておく。
vLLMとSGLangでしばらく密接に作業してきたが、このプロジェクトこそ後続プロジェクトの理想形だと確信している。演算依存グラフを分析し、演算を融合したり、より賢くタスクをスケジューリングしたりする点が印象的。チームに祝意を表したい。
記事とGitHub READMEに目を通したが、本当に素晴らしいプロジェクトだと思う。こうした最適化手法が、推論だけでなく学習段階にも適用できるのか気になる。特にbackward演算とgradient通信の融合が難題だとは認識している。現時点ではdynamic workload（たとえばMoE）をサポートしていないと理解しているが、最近ではMoEを単一カーネルで処理する論文 FlashDMoE: Fast Distributed MoE in a Single Kernel もある。
記事やREADMEまで読んでくれたことへの感謝。学習段階のサポートも可能ではあるが、一般に学習カーネルはより大きく、カーネル起動オーバーヘッドが大きな問題になりにくいため、推論（とくに低レイテンシ）がより大きな恩恵を受ける対象である。共有してもらったFlashDMoE論文も興味深く読んだし、MoEモデルのサポートも次の目標にしていることを強調する。
個人的には、gradientベース学習の最適化に時間を投じることにはやや懐疑的だ。実際、多くの学習タスクは離散値の性質を持っており、gradientベース学習ではうまく扱えないと考えている。
次の段階としては、直接Verilogにコンパイルして、AliExpressでLLMハードウェアを自分で買うのが夢だ。
Chiselなどのハードウェア技術を紹介する記事 を共有。AIやGPUが登場する以前には、こうしたソフトウェアからハードウェアへの直接変換というアイデアは有望なアプローチだった。CPUの進歩は停滞気味で、ソフトウェアとハードウェアの中間層をさらに最適化したいという願望は根強いが、GPUスタイルの並列計算が主流の加速方式であり続ける可能性が高い。汎用CPUは結局、GPUを管理する小さな頭脳の役割にとどまる見通しだ。ただし、ソフトウェアから直接ハードウェアへ移る方式が主流になるのは難しいだろうという見立てである。
5〜10年後にLLMの構造が安定化すれば、ハードウェアへ直接マッピングするのが現実的になるかもしれないという予想。現在の技術でも、数百億パラメータ級を1.5ビット近辺の超低精度ロジックゲートだけで単一ウェハに収められる可能性があるとの言及。精度が高くなるほどゲート数は指数的に増えるため、現状では重みメモリを保持しつつ計算ユニットを共有する方式のほうが効率的。将来的には超低精度LLMの開発が必須課題になる。
学習コストがすでに高いのに、さらにマスクのコストまで加われば状況はもっと厳しくなる、という冗談と、実際にはAIハードウェアスタートアップがこうした方向性の試みを長く続けてきたという冷静な評価。
LLM-in-a-box方式が実在するならかなり魅力的だという感想。近いうちにオフライン（air-gap）環境で働く機会がありそうで、そうしたソリューションは非常に有用だろうという期待。
実際にModalのGPU環境でコードを動かしてみたところ、研究で主張されている性能向上の数値が再現できた。mirageプロジェクトの結果コード を共有。Triton + FlashInferの組み合わせでは1トークンあたりレイテンシが19.2ms程度、MPKでは同条件で7.7msへと大幅に改善した。
以前、小さなCUDAコンテストに参加したことがある。画像やビジョン系の並列アルゴリズムだったが、賢く見せようとして中間結果をメモリにキャッシュした。コンテスト結果を見たら、他の人たちは自分よりずっと速いコードを提出していて驚いた。理由を見ると、中間結果などキャッシュせず、ひたすら再計算していた。メモリ往復より計算コストのほうがはるかに小さかったのだ。このプロジェクトもたぶん似た話だと思う。megakernelにコンパイルすることでlayer境界が消え、中間結果の共有は減って演算量は増えるが、全体としてはメモリ往復が減るので大きな利益になる。特に畳み込みネットワークではsweet spotがありそうだが、megakernelでこの部分をどう扱うのかは分からない。
今もLLMについて新しい比喩が次々に出てくる。もしかするとLLMをトランジスタのようなものと見なせるのではないかと思う。今はちょうど、パンチカードで掛け算しかできなかった部屋サイズのコンピュータの段階に似ていると想像している。100万件のo3-proクエリを同時に回せたら何が起こるのか、想像すると楽しい。
このプロジェクトはCMU（カーネギーメロン）出身。StanfordのHazy Researchでもmegakernelを扱ったブログ No Bubbles がある。この分野で競争が活発に進んでいる様子を見るのは印象的だ。（追記）「mirage」プロジェクトのより大きな全体像を扱う論文もあるが、megakernelアプローチ自体は扱っていない 論文リンク
投稿者本人からの返答。Stanfordとの研究が並行して進んでいることには同意する。主な違いは、自動化されたmegakernel生成コンパイラに注力している点だという。
Hazy ResearchのThunderKittensも非常にクールなライブラリだと言及。最近はNVIDIA GPUを最大限活用するために、形式化、パイプライニング、分割統治、効率最大化、そして専用コンパイラやDSL開発に大きな努力が注がれているという評価。
Qwen 8Bの性能値は、検証されるならかなり印象的だ。従来のmegakernel方式よりずっと実用的に感じられる。各SMごとに1つずつ常駐するこの種のカーネル方式は、かつてのLarrabeeを思い出させる。既存のCUDAではなく、伝統的なプロセス-スレッド-SIMD路線を進んでいたら、今の世界はどうなっていたのだろうと思う。
ソフトウェアベース推論ではなく、純粋なASIC方式で固定型LLMを作るというアイデア。コスト面の利点はあるのか。ソフトウェア側で追加処理や微調整ができる層を設ける余地はあるのか。実際、「十分に良い」水準にかなり近づいている以上、今後2〜4年のうちに専用チップへ固定して使う判断が出てきてもおかしくない。超特化ハードウェアの利点が、いったいどの時点で本当に光るのかが気になる。