🤖 machine learning

AutoMegaKernel: A Statically-Checked Agent Harness for Self-Retargeting Megakernel Synthesis

AutoMegaKernelは、HuggingFaceのLlamaファミリーモデル向けに単一の永続的なCUDAメガカーネルを自動合成する静的チェック済みエージェントハーネスであり、自己改善型エージェントループを通じて、多様なNVIDIAアーキテクチャ間でのデッドロックのない実行を保証しつつ、推論クラスのGPUにおいてベースラインに対して最大1.72倍の高速化を実現します。

原著者： Jaber Jaber, Osama Jaber

公開日 2026-06-09

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Jaber Jaber, Osama Jaber

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

巨大で複雑な工場（大規模言語モデル）を運営して、たった一つの製品（テキストの単語一つ）を製造しようとしている場面を想像してください。

旧来の方法（標準的なAI）：
現在、この工場を運営することは、作業員を倉庫に送り、特定の道具を取りに行かせ、それを組み立てラインに戻して、わずかな作業を行い、また次の道具を取りに倉庫へ送り返す、という作業を単語ごとに数十回繰り返すようなものです。

問題点： 作業員は、組み立てる時間よりも、倉庫へ道具を取りに往復している時間の方が長くなってしまいます。これは「帯域幅制限（bandwidth-bound）」と呼ばれます。工場は、作業員の組み立て速度ではなく、作業員の歩く速さによって制限されてしまうのです。
コスト： 道具を取りに立ち寄るたびに、わずかな遅延（起動レイテンシ）が発生します。これを一歩ごとに繰り返すと、膨大な時間が積み重なります。

新しい方法（AutoMegaKernel）：
著者たちは、AutoMegaKernel (AMK) と呼ばれるシステムを構築しました。AMKは、ルールの概念を根底から変える革命的な新しい工場マネージャーだと考えてください。

1. 「ワンショット」工場

作業員を何度も往復させる代わりに、AMKは**一度の起動（シングルローンチ）**で全工程を完了するように工場を構成します。

比喩： 作業員が一つずつ道具を取りに行くのではなく、あらゆる機械に作業員が配置されている状態を想像してください。彼らはすべて直接的な内部パイプラインでつながっています。「スタート」ボタンが押されると、誰一人として床を離れることなく、組み立てライン全体が一連の連続した動きとして、最初から最後まで走り抜けます。
結果： これにより、「歩行時間」と「停止・再開」による遅延が取り除かれ、理論上、プロセスが大幅に高速化されます。

2. 「安全検査官」（静的検証エージェント）

全員が同時に働く工場を作ることは危険を伴います。もし作業員が完璧に連携できていなければ、衝突したり（レースコンディション）、道具が届くのを永遠に待ち続けたり（デッドロック）する可能性があります。

イノベーション： 通常、これを実現するには人間の専門家が手動でコードを書く必要がありますが、これは困難でミスが起きやすい作業です。AMKは、工場のレイアウトを自動的に設計するAIエージェントを使用します。
セーフティネット： 工場が稼働する前に、**凍結された変更不可能な安全検査官（バリデータ）**がAIの設計図をチェックします。そして、以下のことを数学的に証明します：
- 誰も永遠に待ち続けることがないこと。
- 二人の作業員が同時に同じ道具を使おうとしないこと。
- 操作の順序が完璧であること。
主張： 著者らは、この検査官を7,160個の「トリッキーな」あるいは「壊れた」設計図に対してテストしました。検査官は、すべての不安全な設計図を検知し、工場が始まる前に拒否しました。危険な計画を一度も通すことはありませんでした。

3. 「ユニバーサル・アダプター」（セルフ・リターゲティング）

通常、特定の建物（特定のコンピュータチップ）向けに設計された工場は、別の建物では機能しません。

魔法： AMKは「セルフ・リターゲティング（自己再標的化）」が可能です。設計図を一度書けば、システムは人間が各チップ用にコードを書き直すことなく、異なる種類のコンピュータチップ（RTX 5090、A100、H100など）で完璧に動作するように自動的に適応します。

4. 「軽量版」（量子化）

このシステムは、「より軽い」道具を使って工場を運営する方法も編み出しました。

比喩： 作業員は重いフルサイズの道具（高精度な数学）の代わりに、軽量でコンパクトな道具（int8またはint4精度）を使用します。
メリット： 道具が軽いため、一度に多くの道具を運ぶことができ、作業員の動きも速くなります。
- Int8： 重い道具と同等の性能を維持しながら（ロスレス）、12%高速化します。
- Int4： さらに高速に動き（「歩行時間」を2倍以上削減）、最終的な製品の完璧さはわずかに低下する可能性がありますが（ロスの発生）、読み取り可能なレベルは維持されます。

5. 正直な結果（「本音」の話）

著者たちは、このシステムがどこで勝ち、どこで負けるのかについて非常に透明です。

勝利する場面： 「インファレンス（推論）クラス」のチップ（L4、L40S、A10G、およびコンシューマー向けのRTX 5090など）において、AMKシステムは、単語を一つずつ生成する際に現在の業界標準（cuBLAS）よりも高速です。これによって、競合を1.1倍から1.3倍上回ります。
- 理由： 他のシステムを悩ませている「往復の歩行」による遅延を、見事に排除することに成功したためです。
敗北する場面： 巨大で高速な「トレーニング（学習）クラス」のチップ（A100やH100など）では、現時点のバージョンは標準的なシステムよりも低速です。
- 理由： 作業員同士の「安全な連携（同期）」に、わずかな時間がかかるためです。最も高速なチップにおいては、この微小な遅延がボトルネックとなります。著者らはこれを正直に認めています。「私たちはボトルネックを特定しており、それについて誠実に対処しています。」

まとめ

**Autoメガカーネル（AutoMegaKernel）**は、AIモデルを実行するための「ワンショット」工場を設計するために、AIエージェントを使用するシステムです。これには、設計が決してクラッシュしないことを保証する厳格な安全検査官が含まれています。また、さまざまなコンピュータチップ上で自動的に動作することにも成功しています。

大きな勝利： 現在、多くの一般的なコンシューマー向けおよびデータセンター向けチップにおいて、テキストを一度に一単語ずつ生成する最も速い方法であり、不要な遅延を取り除くことで標準的なツールを打ち負かしています。
課題： まだ完璧ではありません。絶対的な最高速のスーパーコンピュータ用チップにおいては、調整のオーバーヘッドが依然としてボトルネックとなり、既存の巨人をわずかに下回りますが、著者らはこのシステムが安全で、正確であり、自己改善が可能であることを証明しました。

技術要約: AutoMegaKernel (AMK)

問題提起

バッチサイズ1におけるシングルストリームの自己回帰的デコーディングは、根本的に帯域幅に制約されます。レイテンシの底（レイテンシフロア）は、総ウェイト・バイト数とHBM帯域幅の比率によって決定されます（ $t_{min} = \text{weight\_bytes} / \text{HBM\_bandwidth}$ ）。従来の実行フレームワーク（例：PyTorch, cuBLAS）は、オペレータごとに個別のカーネルを起動するため、重大なCPU起動レイテンシが発生し、オペレーション間でアクティベーションがHBMを経由してラウンドトリップすることを強いるため、この底に到達できていません。CUDA Graphsは起動オーバーヘッドを軽減しますが、オペレータ間の境界とHBMへのラウンドトリップという問題は残ります。

既存の「メガカーネル」アプローチ（例：Mirage Persistent Kernel）や手動最適化カーネルは、フォワードパス全体を単一のパーシステントカーネルに融合することで、レイテンシの底に対処します。しかし、これらのシステムは信頼性とポータビリティに欠けています。これらはしばしば静的な安全ゲートを省略しており（デッドロックやレースコンディションのリスクがある）、特定のモデルやアーキテクチャを個別にターゲットとする必要があり、自律的なエージェントがシステムを安全に最適化するための構造化されたインターフェースを提供していません。

手法

AutoMegaKernel (AMK) は、HuggingFaceのLlamaファミリーモデルを、モデルごとの手書きCUDAを用いることなく、単一のパーシステントな協調カーネルへとコンパイルする、自律的なエージェント駆動型システムです。本システムは4つのレイヤーと2つのオートリサーチループで構成されています。

レイヤー 0 (VM): 手書きの、固定された、アーキテクチャごとのパーシステント・メガカーネル（vm/scheduler.cu）。これは、1つのStreaming Multiprocessor (SM) ごとに1つの協調スレッドブロックを起動します。各ブロックはローカルキューを所有し、wait → dispatch → signal 命令を実行しながら、タスクのグローバルなトポロジカル順序を辿ります。同期にはロックを避け、単調増加カウンタに依存します。
レイヤー 1 (Instructions): 個別に正当性がチェックされた、ABI準拠のマイクロカーネル（GEMV、アテンション、RMSNormなど）。
レイヤー 2 (Scheduler): エージェント向けの編集サーフェス。モデルをインポートし、型付きの Schedule IR（単調増加カウンタによって同期されたタスクDAG）へと低レベル化し、ScheduleConfig オブジェクトを介して構造的な構成（タイリング、融合、SM割り当て、パイプラインの深さ）を提案します。
レイヤー 3 (Roadmap): 将来的なダイナミズム（例：継続的バッチング）のためのプレースホルダー。

安全ゲート (The Safety Gate):
核心となる革新は、静的にチェックされたバリデータ（schedule/ir.py 内の validate()）です。あらゆるGPU起動の前に、バリデータはSchedule IRに対して一連の静的グラフチェックを実行し、以下を証明します。

デッドロックフリー: プロデューサー・コンシューマー間のDAGが非巡回であること、待機閾値が充足可能であること（ $1 \le t \le \#producers$ ）、および各SMのキューの順序が有効であること。
レースフリー: 共有カウンタが正しく結合されていること（マルチプロデューサーのカウンタを待機する前に、すべてのプロデューサーが完了している必要がある）、およびメモリ・アクセスが「happens-before」関係を遵守していること。
正当性: 出力の到達可能性と参照整合性の検証。

スケジュールがこれらのチェックに失敗した場合、即座に拒否されます。これにより、エージェントが提案した不安全なスケジュールがGPUをハングさせることを防ぎます。

オートリサーチループ:
AMKには、ScheduleConfig の変更を反復的に提案し、それらを検証し、CPUリファレンスに対して正当性ゲートを実行し、レイテンシを測定し、変更を保持または破棄する、エージェント駆動型の unattended ループが含まれています。このループは、自身のベースラインに対してメガカーネルを自律的に改善します。

主な貢献

静的にチェックされた Schedule IR: 信頼できるベースラインに基づき、静的グラフチェック（機械的な証明ではなく）を通じてデッドロックおよびレースフリーを証明する、固定されたバリデータ。これは、7,160個の敵対的に構築されたスケジュール（うち6,091個は不安全）に対し、360個の実際のモデル低レベル化すべてを受け入れつつ、ゼロの誤受理（false-accepts）を達成しました。
自動生成のカバレッジ: AMKは、手書きのCUDAなしで、10個中10個のサポート対象モデル（SmolLM2-135MやTinyLlama-1.1Bなどの実際のチェックポイントを含む）に対して正しいメガカーネルを自動生成します。また、サポートされていないバリアント（例：バイアス付きプロジェクション）を正しく拒否するか、あるいはその盲点を正直に記録します。
セルフ・リターゲティング: 単一のソースコードベースから、sm_80 (A100)、sm_90 (H100)、および sm_120 (RTX 5090) 用の正しいメガカーネルを生成し、教師あり実行（eager execution）とビット単位で一致、またはタイトなfp32許容誤差（ $\le 3.8 \times 10^{-5}$ ）内に収めます。
量子化メガカーネル: 同一のIRパスから、int8（貪欲法によるロスレス、1.12倍高速）および int4（ロスあり、ウェイト・トラフィックを2.42倍削減）のウェイト・オンリー量子化カーネルを自動生成します。
エージェント駆動型コンパイラ: 正当性を構造的な不変量として扱うことで、エージェントがランタイムハングのリスクなく、安全にスケジュールの探索空間を探索できるようにしています。
誠実なパフォーマンス研究: 本論文は、推論クラスおよびトレーニングクラスのGPUにおける結果を報告しており、システムが勝利する箇所と、確立されたベースラインに劣る箇所を明示的に開示しています。

実験結果

正当性とポータビリティ:

RTX 5090、A100、および H100 において、AMKはHuggingFaceの貪欲デコードをトークン単位で完全に再現し、SmolLM2-135M チェックポイントにおいて教師強制（teacher-forced）パープレキシティを $2.5 \times 10^{-7}$ まで一致させました。
バリデータは、ストレス・テストにおける6,091個の不安全なスケジュールをすべて拒否し、0%の誤受理率を実証しました。

推論クラスGPUにおけるパフォーマンス (勝利):

RTX 5090 (コンシューマ向け): オートチューニングされた int8 ウェイト・オンリー (W8A16) メガカーネルは、バッチ1デコードにおいて、CUDAグラフ化された cuBLAS (bf16) を 1.19–1.23× 上回りました。これは精度非対称な比較（int8 対 bf16）であり、この勝利は、より速いカーネルによるものではなく、流れるウェイト・バイト数を削減したこと（0.61倍）によるものです。
データセンター推論フリート: int8 メガカーネルは、推論クラスのフリート全体で CUDA グラフ化された cuBLAS bf16 を上回りました。
- NVIDIA L4: 最大 1.33× 高速（モデルサイズとともに向上）。
- L40S: 1.25–1.27× 高速。
- A10G: スケールに応じて最大 1.08× 高速。
自己改善: オートリサーチ・ループは、RTX 5090 において、デフォルトの構成に対して 1.25–1.72× の自動改善を実現しました。

トレーニングクラスGPUにおけるパフォーマンス (敗北):

高帯域幅のトレーニングクラスGPU (A100, H100) では、自動生成されたメガカーネルは CUDA グラフ化された cuBLAS および vLLM に及びません。
- H100 では、AMK は vLLM のデフォルトの CUDA グラフ・パスよりも 1.65× 遅い 結果となりました。
- A100 では、最適化された bf16 GEMV は、測定されたHBMピークの 約63% にしか達しませんでした（cuBLAS の約90%に対し）。
診断: ボトルネックは、タイルごとのクロスSM同期コスト（cuBLASが支払わない、メガカーネルが支払う固定コスト）です。高帯域幅のシリコンでは、この同期オーバーヘッドが支配的になりますが、推論クラスのGPUでは、融合と量子化による帯域幅の節約がこのコストを相殺します。

重要性と主張

本論文は、AutoMegaKernel が「生の速度」から「体系的な安全性とポータビリティ」への転換を象徴していると主張しています。その主要な貢献は、自律的なエージェントがモデル全体のスケジュール空間を安全に探索することを可能にする、静的にチェックされたコンパイラ・アーキテクチャです。

信頼 (Trust): GPUのハングのリスクを、実行前に不安全なスケジュールを拒否する静的安全ゲートに置き換えます。
汎用性 (Generality): 手動のCUDAチューニングなしに、アーキテクチャ（sm_80/90/120）やモデルファミリーを越えたセルフ・リターゲティングを実現します。
誠実さ (Honesty): 著者らは、パフォーマンスの結果がレジーム依存であることを明示しています。システムは、演算の融合と量子化を利用してメモリ・トラフィックを削減することで、推論クラスのGPUで勝利しますが、現在のところ、より粗い同期（将来の課題）を必要とする同期オーバーヘッドのために、ハイエンドのトレーニング・ハードウェアでは負けています。

結論として、現在のカーネル品質は、高性能なトレーニング・ハードウェアにおける手動最適化されたベースラインにはまだ及びませんが、静的にチェックされたハーネスは、そのギャップを埋めるために必要な（より粗い同期や、より深いパイプライン化などの）将来の改善を安全に自動化するための、不可欠な基盤を提供します。