ハクソク

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

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

2025年6月20日原文(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マシャリングを回避できるかも。ループ融合とか？他に面白い最適化ができるのかな？

Hacker Newsで議論の続きを見る