ハクソク

世界を動かす技術を、日本語で。

LLMをメガカーネルにコンパイルする：低遅延推論への道

340日前原文(zhihaojia.medium.com)

概要

  • LLM推論の高速化 を目的とした 自動メガカーネル生成コンパイラMPK の紹介
  • 従来手法の課題 （カーネル起動回数・通信分断）を 統合カーネル化で解決
  • 単一・複数GPU環境両対応、最大6.7倍のレイテンシ短縮を実現
  • タスクグラフ変換とメガカーネル実行 による高効率化
  • 今後の展望 として最新GPU対応や動的ワークロードへの拡張を計画

LLM推論を高速化するメガカーネル自動生成コンパイラ「MPK」

  • LLM（大規模言語モデル）推論のボトルネック は、複数のGPUカーネル起動や外部通信による ハードウェア非効率
  • MPK は、これらの計算・通信を 単一メガカーネル に自動統合する コンパイラ兼ランタイムシステム
  • Python数十行でLLMをメガカーネルに変換 可能、開発者の負担を大幅軽減
  • カーネル起動オーバーヘッド排除、 計算・通信の重畳実行、 細粒度ソフトウェアパイプライニング を実現
  • 単一GPUでも既存最適化実装より低レイテンシ、 GPU数増加で効果拡大

メガカーネル化の意義とMPKの優位性

  • 単一カーネル起動 でモデル全体（計算・通信）を連続実行する「 メガカーネル設計」
  • カーネル起動オーバーヘッド排除、 多層パイプライニング、 計算・通信重畳 による 最大6.7倍高速化 （マルチGPU時）
  • 既存MLフレームワーク（PyTorch/Triton/TVM等）では統合カーネル生成不可
  • 多様な専用カーネル（NCCL/FlashInfer等）の統合困難性 を 自動化 で解決

MPKコンパイラ：LLM計算グラフの最適タスクグラフ化

  • LLM計算は計算グラフ（ノード＝演算・通信、エッジ＝依存関係）で表現
  • 従来は各演算ごとにGPUカーネルを個別起動 → パイプライニング機会損失
  • MPKは計算グラフを細粒度タスクグラフに自動変換
    • 各タスク＝GPU SM単位の計算・通信処理
    • イベント＝タスク間同期点
    • タスクとイベントの依存関係を明示し 最大限の並列・重畳実行 を引き出す
  • 部分的出力依存の通信（例：matmul→allreduce）を即時実行可能
  • Mirageカーネルスーパーオプティマイザ による各タスクの高効率CUDA実装自動生成

MPKランタイム：タスクグラフのメガカーネル内実行

  • GPU上で全タスクグラフを単一メガカーネル内で実行
  • SM（ストリーミングマルチプロセッサ）をワーカーとスケジューラに静的割当
    • ワーカー ：各SMでタスクキューを順次実行し、完了時にイベント通知
    • スケジューラ ：アクティベート済イベントを検出し、依存タスクを起動
    • 全体の同期・タスク遷移をカーネル内で完結 → タスク間オーバーヘッド1–2μs
  • イベント駆動型実行 で計算・通信の重畳や多層パイプライニングを実現
  • マルチGPUでも通信・計算の完全な重畳が可能

今後の展望と課題

  • 最新GPUアーキテクチャ（NVIDIA Blackwell等）対応
    • ワープ特化最適化 とメガカーネルモデルの統合が課題
  • 動的ワークロード（Mixture-of-Experts等）対応
    • 静的タスクグラフ から 動的制御フロー・条件分岐対応 への拡張を開発中
  • 高度なスケジューリング・タスク割当
    • 優先度・スループット最適化 等の応用（SLO対応推論やハイブリッドバッチ処理）

コミュニティへの呼びかけと情報源

  • MPKはLLM推論基盤の新たなパラダイムシフト
  • コード・ドキュメント・詳細情報 は公式GitHub（https://github.com/mirage-project/mirage）を参照
  • フィードバック・コントリビューション・共同研究を歓迎

Hackerたちの意見

次のステップは、Verilogに直接コンパイルして、AliExpressでいくつかのLLMを買うことだね。

└

https://riscv.org/blog/2021/02/hardware-description-language... これはAIやGPUが登場する前の有望なアイデアの一つだったね。CPUが停滞してる中で、人々は中間層のソフトウェアとハードウェアをさらに最適化したいと思うのは自然なこと。でも、GPUスタイルの並列計算が加速コンピューティングを支配するようになると思う。汎用CPUは、GPUを調整する小さな脳として残るんじゃないかな。ソフトウェアからハードウェアへの移行のアイデアは、主流にはならないかもしれないね。

└

そうだね… LLM-in-a-boxは結構面白そうだよ！これからエアギャップの作業があるから、そんなのがあったら便利だな。

└

だって、トレーニングコストが高すぎるのに、さらにマスクコストを追加するなんて。もっと真面目に言うと、これってAIハードウェアのスタートアップがずっとやってきたことじゃない？

└

5年から10年後、LLMが安定したら、ハードウェアに直接マッピングするのが理にかなうかもしれないね。今のプロセスだと、1000億パラメータが1.5ビット精度でロジックゲートに直接実装すれば、1枚のシリコンウエハーに収まるかも。もっと高い精度を使うとゲート数が指数的に増えるから、今はメモリに重みを保持して、計算ブロックを共有して使い回す方がいいと思う。ただ、将来的には超低精度のLLMをちゃんと動かせるようにしないとね。

これめっちゃクールだね。記事とGitHubのREADMEを見て楽しんだよ。これらの最適化が推論だけじゃなくて、トレーニングにも適用できるのか気になってた。ここでの課題は、逆計算と勾配通信を融合させることなんだろうね。それに、今のところ動的なワークロード、例えばMoEには対応してないみたい。最近、これにぴったりな論文を見つけたよ：「FlashDMoE: Fast Distributed MoE in a Single Kernel」 - https://arxiv.org/pdf/2506.04667

└

投稿とGitHubのREADMEを読んでくれてありがとう。トレーニングのサポートは確かに可能ですが、低遅延の推論ほどのメリットはないかもしれません。なぜなら、トレーニングは一般的にもっと大きなカーネルを含むからで、カーネルの起動オーバーヘッドがあまり重要でなくなるからです。FlashDMoEの取り組みを共有してくれてありがとう。次のステップはMoEモデルのサポートです。お楽しみに！

└

個人的には、勾配トレーニングの最適化に時間を投資するのはちょっと無駄だと思う。現実の多くのトレーニングタスクは、離散的な値を持っていて、勾配でトレーニングできないからね。

Qwen 8Bの数字、もし確認できたら、すごく印象的だね。前のメガカーネルよりずっと実用的だし。ただ、各SMに一つの持続カーネルがあるのはLarrabeeを思い出させるし、CUDAパスじゃなくて、伝統的なプロセススレッドSIMDパスをやったら世界はどうなるんだろうって考えちゃう。

vLLMとSGLangをここ数ヶ月密に使ってきたけど、これこそが私が想像していた後継プロジェクトの姿だよ。操作依存グラフを分析して、タスクを賢くスケジューリングするか、もしくは融合するっていうね。チームにおめでとう！

└

ポジティブなフィードバックをありがとう！私たちは、MPKが既存のLLMサービングシステムを強化できると信じています。特に低遅延のLLMサービングにおいてです。他の人たちと方向性についてコラボレーションできる機会にとてもワクワクしています。

これがCUDAグラフよりも大きな利点を提供する理由について、誰か直感的にわかる人いる？グラフのCPU起動コストは微小で、ほとんどの作業がGPUのスケジューラにオフロードされてることを示してる。メガカーネルを使えば、カーネルの境界でのI/Oマシャリングを回避できるかも。ループ融合とか？他に面白い最適化ができるのかな？

└

その通りだね。事前コンパイルされたCUDAグラフのCPU起動コストは本当に小さい。CUDAグラフは手動でカーネルを起動するよりも大きな進歩だけど、まだカーネルをモノリシックでブラックボックス的な操作として扱ってる。メガカーネルはその境界を消し去るんだ。記事の例のように、Matmul -> AllReduceがあると、AllReduceカーネルはMatmulカーネルが全部終わるまで始められない。依存関係はカーネルレベルにあるからね。メガカーネルでは、これらの操作を細かい「タスク」に分けてSM間でスケジュールするんだ。Matmulの最初のスライスからデータが必要なAllReduceタスクは、そのスライスがいくつかのSMによって計算されるとすぐに始められるけど、他のSMはまだMatmulの残りを処理してる。この細かいソフトウェアパイプラインと計算/通信のオーバーラップは、依存単位が全体のカーネルの時には実現できないんだよね。

└

グラフのCPU起動コストは微小です。 全くそうじゃないよ。カーネルの起動オーバーヘッドと同じくらいだよ。

このプロジェクトはCMUのものだよ。スタンフォードのHazy Researchもメガカーネルについて話してたよ：https://hazyresearch.stanford.edu/blog/2025-05-27-no-bubbles この分野での競争が見られて嬉しいね。（編集）：より大きな「ミラージュ」プロジェクトをカバーした関連論文だけど、「メガカーネル」アプローチは含まれてないよ：https://arxiv.org/abs/2405.05751

└

これはブログ記事の著者です。スタンフォードの取り組みが並行しているのはその通りです。主な違いは、私たちの焦点がコンパイルにあることです。つまり、メガカーネルを自動的に生成しやすくすることです。

└

Hazy ResearchにはThunderKittensっていう、かなりクールなライブラリもあるよ。NVIDIAのGPUモデルでGPUの効率を最大化するために、本当に形式化したり、パイプラインを作ったり、分割して攻略するために多くの努力がされているみたいだね。そして、いろんなもののためにコンパイラやDSLを書くことも。

スケールされたハードウェアでこれらのモデルを活用するための重要な発見だね。このアプローチは、LLMだけでなく他のタイプのニューラルネットワークにも応用できると思う。面白い分野になりそうだね。

└

フィードバックありがとう！はい、このアプローチは一般的で、他のMLワークロードにも適用できると信じています。

こんにちは、著者の方々。GPU上のインタープリタアプローチは有望な道に見えますね。この似たような同時進行の研究を見たことありますか？「https://news.ycombinator.com/item?id=44111673」CUDAプログラミングモデルの基本（例えばカーネルの起動）が、ハードウェアをより効果的に使うための細かいタスクベースの並列処理に置き換えられているのが興味深いです。CUDAが何かしらの面で私たちを足止めしているのかなって思います。あなたたちの研究がPyTorchの実験的なバックエンドに入る可能性はどれくらいですか？素晴らしい内容をシェアしてくれてありがとう！追伸、1部の最初の2段落がほぼ同じになってるよ。

└

素晴らしいフィードバックありがとう！スタンフォードのMegaKernelプロジェクトは似たような課題に取り組んでるけど、手動のCUDA実装に焦点を当ててるんだ。一方でMPKはコンパイラ主導のアプローチを取っていて、ユーザーがPyTorchレベルでLLMを表現すると、MPKが自動的に最適化されたメガカーネルにコンパイルしてくれるんだ。私たちの目標は、メガカーネルのプログラミングをもっとアクセスしやすくすること。CUDAが制限要因になることには完全に同意するよ、特にレイテンシに敏感なワークロードに関してはね。GPUが大きくなって速くなるにつれて、ハードウェアリソースを完全に活用するスタンドアロンカーネルを書くのがますます難しくなってる。特に小さなバッチサイズで低レイテンシを最適化する場合はね。 > あなたたちの研究がPyTorchの実験的なバックエンドに入る可能性はどれくらいですか？その方向には間違いなくワクワクしてるよ。MPKがPyTorchのメガカーネル生成をサポートできると信じてるし、実現方法を積極的に探ってるところだよ。お楽しみに！ > 追伸、1部の最初の2段落がほぼ同じになってるよ。指摘してくれてありがとう！投稿を最終化する時に重複した段落を削除するつもりだったんだ。

改善は本物だよ！しかも、多くの研究とは違って、コードがちゃんと動くんだ。Modal GPUを使って結果を再現できたよ。コードはここに置いておくね：「https://github.com/mirage-project/mirage/pull/327/files」Triton + FlashInfer: プロンプト長39、生成長264、トークンあたりのレイテンシ19.189573345762312 ms MPK: プロンプト長39、生成長334、トークンあたりのレイテンシ7.71875 ms

└

私たちの結果を再現してくれてありがとう！

著者の方々に質問です。彼らはこのスレッドにとても反応が良さそうだからね :)。1. 各タスクはどれくらい細かいの？例えば、従来の行列乗算カーネルでは、各スレッドブロックが結果行列の小さな出力タイルを担当してるよね。Mirageのメガカーネルでは、対応する小さな出力タイルごとにタスクがあるのかな？2. Mirageコンパイラはタスクグラフをどう形成するの？各オペレーターのデータフローのドメイン知識を個々の要素の粒度で持ってるのかな？再び行列乗算を例に取ると、特定の出力タイルはA行列の対応するM_BLOCK行を必要とするよね。もしA行列が以前の行列乗算の出力（+非線形性）だった場合、依存関係はそのM_BLOCK行に対応する出力タイルタスク全てになるのかな？

└
  1. MPKでは、各タスクは個々のSMにマッピングされるよ。タスクが処理する作業量は、従来のカーネルごとのオペレーターアプローチのスレッドブロックと似てる。2. TL;DR: MPKは、各タスクに関連する入力と出力のテンソルを追跡することで、タスク間の依存関係を自動的に分析するんだ。長いバージョン: MPKはimap、omap、fmap（Mirageの論文のセクション2を参照）を使って、各タスクの入力と出力のテンソルを決定するよ。タスクAがBが消費するテンソル要素を生成する場合、タスクAとタスクBの間に依存関係が生じるんだ。つまり、Aの出力がBの入力と重なる場合ね。 > 再び行列乗算を例に取ると、特定の出力タイルはA行列の対応するM_BLOCK行を必要とするよね。もしA行列が以前の行列乗算の出力（+非線形性）だった場合、依存関係はそのM_BLOCK行に対応する出力タイルタスク全てになるのかな？その通りだよ。この場合、AのそのM_BLOCK行を消費する全ての出力タイルタスクは、前のオペレーターでAの対応部分を生成する全てのタスクに依存することになるんだ。