# Agentic Operator Generation for ML ASICs > [!info] Talk metadata > - **会議:** [[MLSys2026]] Day 4 (May 21 / Thu)、Industry Track(9:00 - 9:15 AM PDT、公式仮想会場ページに基づく)。指示のセッション枠表記は「08:30 PDT〜(Day 4 Industry & Research Track: Agentic AI/MLSys & LLM Serving 5 セッション)」。 > - **会議論文:** Proceedings of the 9th MLSys Conference (Industry Track), Bellevue, WA, USA, 2026(論文 PDF フッタ)。 > - **登壇者:** Alec M. Hammond(Meta, Menlo Park, USA)。論文で Aram Markosyan と並び equal contribution(共同筆頭)、スライド表紙でも筆頭・太字表記。所属チームは **MTIA SW, FAIR, PyTorch**(スライド表紙)。 > - **共著者:** Alec M. Hammond*, Aram Markosyan*, Aman Dontula, Simon Mahns, Zacharias Fisches, Dmitrii Pedchenko, Keyur Muzumdar, Natacha Supper, Site Cao, Haishan Zhu, Mark Saroufim, Joe Isaacson, Laura Wang, Warren Hunt, Kaustubh Gondkar, Roman Levenstein, Gabriel Synnaeve, Richard Li, Jacob Kahn, Ajit Mathews(全員 Meta, Menlo Park, USA。*=equal contribution)。連絡先 `[email protected]`, `[email protected]`。 > - **URL:** https://mlsys.org/virtual/2026/oral/3817 > - **システム名:** TritorX(agentic AI system for Triton MTIA kernel generation)。 > - **関連リンク(論文参照より):** MTIA ブログ `https://ai.meta.com/blog/meta-mtia-scale-ai-chips-for-billions/`(スライド出典)、BackendBench `https://github.com/meta-pytorch/BackendBench`、KernelLLM `https://huggingface.co/facebook/KernelLLM`。 > [!abstract] 概要(論文 PDF アブストラクト・忠実日本語訳) > 我々は TritorX を提示する。これは新興アクセラレータプラットフォーム向けに、機能的に正しい Triton PyTorch ATen カーネルをスケールして生成するために設計された agentic AI システムである。TritorX は大規模言語モデル(LLM)を、カスタム linter・JIT コンパイル・PyTorch OpInfo ベースのテストハーネスと統合する。このパイプラインは実 Meta Training and Inference Accelerator (MTIA) シリコンと、次世代デバイス向けのハードウェアシミュレーション環境の双方に対応する。少数の高頻度カーネルに対して性能を優先する従来のカーネル生成手法とは対照的に、TritorX は **coverage(網羅性)を優先**する。我々のシステムは、多様なデータ型・形状・引数パターンを含む演算子集合全体にわたる **correctness と generality** を重視する。実験では、TritorX は **481 個のユニークな ATen 演算子**に対して、対応する全 PyTorch OpInfo テスト(**合計 20,000 件超**)をパスするカーネルとラッパーの生成に成功した。TritorX は、新たなアクセラレータプラットフォームに向けた完全な PyTorch ATen バックエンドの一夜での生成(overnight generation)への道を拓く。 ## 問題設定: ASIC のためのオペレータ網羅性 - ML/AI ハードウェアの急速な普及。US データセンタ電力消費は 2028 年までに総電力の 6.7〜12.0% に達すると予測され、効率的なアクセラレータ需要が逼迫(Shehabi et al., 2024、論文 Introduction)。業界はカスタムシリコン ASIC を含む heterogeneous なフリートに投資している。 - **MTIA(Meta Training and Inference Accelerator)**: Meta の自社 AI チップ。推薦モデル(DLRM)を Facebook/Instagram/Threads の数十億ユーザに提供し、GPU 比で総保有コスト(TCO)を **44% 削減**(Coburn et al., 2025、論文)。スライドでは「2 年未満で 4 世代(約 6 ヶ月ごとに新チップ)」「MTIA 300 → 500 で HBM 帯域 4.5x・compute FLOPS 25x」と提示(出典 Meta ブログ、スライド2)。 - **課題は operator coverage**: 各新プラットフォームは PyTorch 互換のソフトウェアエコシステム構築に多大な工数を要する。ATen の演算子集合は大きく、これを網羅したカーネルバックエンドの整備は新規モデルアーキテクチャの推論/学習に不可欠(論文 §1)。 - **ビジョン(スライド3)**: 「新チップ到着前に full operator coverage を得られるか」「新演算子ライブラリを overnight でオンボードできるか」「これらツールでコンパイラ/ランタイムスタックを改善できるか」。Kernel Roofline CDF 概念図では、AI 生成カーネルが大半を担い、human experts は roofline 効率最上位の領域のみ担う構図。 ## 提案手法: TritorX(FSM ベースのカーネル生成エージェント) TritorX は **finite state machine (FSM)** として実装され、linting・compile・testing・debugging・LLM 呼び出しの専用ツール/ルーチンを統合する(論文 §3、Figure 1/Figure 2)。各演算子は self-contained なセッションで生成され、LLM モデル・各 state の有効/無効・ハイパーパラメータ(最大反復数等)を事前設定可能。 パイプライン(スライド6–7, 11、論文 §3.1): 1. **Init & Prompt Generation**: タスク記述・出力要件・対象演算子の docstring・3 つの手作り例から成る初期プロンプトを構築。ATen docstring の有向非巡回グラフを構築し「nested」docstring を含める。プロンプトは最小限の MTIA コンテキスト・3 つの具体例・PyTorch docstring から成る(スライド7、例: `nn.functional.logsigmoid`)。 2. **Generate Kernel**: reasoning LLM が wrapper/kernel ペア(PyTorch シグネチャに合うラッパー + 機能を実装する 1 個以上の Triton カーネル)を生成。LLM は **oracle** として扱い CWM/GPT-OSS/Claude 等を呼べる。 3. **Triton MTIA Linter**: 2 つの主責務 — (1) cheating 防止、(2) モデルへの MTIA 教育。(1) 出力が Triton JIT ハーネス互換、(2) 他演算子や CPU への dispatch による「cheat」をしない、(3) 有効な Triton MTIA 構文/ライブラリのみ使用、を保証。config で `module_restrictions`(許可関数、200+ の Triton MTIA 操作)、`module_scope_restrictions`(`tl.*` は kernel 関数内のみ)、`forbidden_functions`(`eval`/`exec`/`compile` 禁止)を制御。lint 違反時は構造化レポートを生成し再生成へ。 4. **Compile / Execute / Test(Triton-MTIA Compiler)**: lint 通過後、MTIA インフラ互換の Triton JIT コンパイルハーネスへ。演算子 config が定める supported dtype に応じテストを合成し、各テストでループ・必要に応じ再コンパイル。コンパイル成功・ランタイムエラー無しなら、同入力を host に移し reference ATen CPU 実装で実行、PyTorch と同じ tolerance で出力比較。compile failure / runtime error / accuracy error 発生時に break して feedback へ。 5. **Process Results & Debug**: 前 state 出力を解析し次アクションを決定。`mtia-dbg`(LLDB ベース)を統合し backtrace・decoded register・failed core 等を抽出。failure からは recovery プロンプト種別(compiler failure / accuracy failure 等)を選択。compiler log は数千トークンに及ぶため、二次 LLM(要約モデル)で要約することがある。 6. **Termination**: (1) 全テスト pass で success 終了、(2) 規定 LLM 呼び出し数到達、(3) context window 飽和で最新生成を初期 proposal として新セッション開始、(4) 想定外エラーで終了。 設計上の特徴(スライド6): predetermined な state transition を持つ FSM(free-form tool calling ではない)、model agnostic(MTIA 固有知識をモデルに焼き込まない)、scalable(silicon または simulator 上で数百カーネルを同時実行)、ablation 容易。FSM は **embarrassingly parallel** で、対象演算子・dtype・LLM パラメータ・実行パラメータを指定して大規模実行できる(論文 §3.1)。 > [!note] 出典に関する注記 > 正式タイトルは論文 PDF 表紙の "Agentic Operator Generation for ML ASICs"。システム名はスライド・論文で一貫して TritorX。登壇者個人はスライド表紙で筆頭・太字の Alec M. Hammond(公式仮想会場ページの著者列挙でも先頭)であり、論文では Aram Markosyan と equal contribution。本ノートでは発表者を Alec M. Hammond と推定するが、共著者全員を上記に列挙した。所属は論文では全員 Meta, Menlo Park。スライド表紙の所属チーム表記は「MTIA SW, FAIR, PyTorch」。 ## 実験・評価 ### ベースライン構成(論文 §4.1) - **対象**: 568 個の MTIA 互換 OpInfo 演算子(629 から filter)。MTIA は複素数非対応のため FFT 系などを除外、random-number operators(device/host 間の RNG 差異で検証困難)も除外。dtype は bfloat16, float16, float32, int32, int64 のみテスト、合計 20,000 件超のテスト。 - 演算子あたり最大 **3 TritorX attempts**(LLM dialog session)、各 attempt 最大 **15 LLM calls**(=15 回の state machine 反復)。 - カーネル生成 LLM: **Code World Model (CWM)** / **GPT-OSS 120B** / **Claude Sonnet 4.5** / **Claude Opus 4.5**。全モデル context length 131,072、temperature 1.0。top-P は CWM 0.95 / GPT-OSS 1.0。GPT-OSS reasoning は "high"。 - feedback 要約モデル: Llama-4-Maverick(生成パラメータは CWM と同一)。 - 生成ジョブを **200 production MTIA devices**(silicon or simulation)へ dispatch。**95% のランが 2 時間以内に完了**、残りの tail は 6–8 時間。 ### 主要結果(スライド4「Key Results」/ Figure 3) - **ATen 演算子の 84.0% に対し包括的実装を生成**(model ensemble、MTIA 互換 OpInfo Ops 上)。各演算子は対応する OpInfo テスト全件(20k+)をパス。 - 個別モデルの operator coverage(Figure 3, アンサンブルは複数ランの union): - **Claude Opus 4.5: 78.7%** - **Claude Sonnet 4.5: 72.2%**(Figure 3 では Sonnet 72.0% 表記、Table 3/2 では 72.2%。スライド4 のグラフ凡例は GPT-OSS 120B=72.0%、Sonnet=(緑線)。論文 Table 3 baseline は CWM 55.3 / GPT-OSS 72.0 / Sonnet 72.2 / Opus 78.7) - **GPT OSS 120B: 72.0%** - **CWM 32B: 55.3%** - **Model Ensemble: 84.0%** - **full generation round は production infra 上で 1–2 時間**。simulation infra とも互換。生成演算子は 4 つの production model で out-of-the-box 動作。 ### 失敗モードと演算子カテゴリ別 coverage(論文 §4.1, Table 1) - uncovered な演算子の主な失敗モード 2 種: (1) half-precision の数値誤差(高精度では正しいが低精度入力の safeguard 欠如)、(2) MTIA の Triton dialect への不慣れ/非互換。失敗演算子の約 **80% は依然コンパイルし一部テストを pass**、**30% はテストの 80% 超を pass**。 - カテゴリ別 coverage(Table 1、Op Count / CWM / GPT-OSS / Sonnet / Opus): - Elementwise (161): 80.1 / 84.6 / 84.0 / 85.9 - Deep Learning (90): 64.4 / 71.1 / 71.4 / 76.2 - Linear Algebra (78): 71.8 / 79.5 / 78.2 / 78.2 - Other (78): 75.6 / 74.3 / 76.9 / 75.6 - Shape Manipulation (75): 96.0 / 96.0 / 94.7 / 96.0 - Reduction (63): 69.8 / 74.6 / 76.2 / 71.4 - Indexing & Selection (34): 73.5 / 79.4 / 91.2 / 91.2 - → Shape Manipulation が最高(96.0%)、Deep Learning Category が最低。後者は Triton セマンティクスに必ずしも合致しない計算的設計パターンのため(論文)。 - QEMU simulator 上で GPT-OSS を将来世代向けに実行した単一ランは **coverage 73.1%**。遭遇した compiler failure と feature gap を集約しコンパイラ/ASIC エンジニアと共有。 ### End-to-end 応用(production model 一般化、論文 §4.1, Table 2) - `__torch_dispatch__` で NanoGPT・DLRM・社内推薦モデル 2 種(MM1, MM2)の演算子入力を記録し、TritorX 検証ループで OpInfo 入力の代わりに使用。batch size 1024 固定、後 3 モデルは実 production データで実行。 - 新ステップ: 既存 OpInfo 生成済み演算子を、full e2e run で集めた入力で即テスト(Table 2 列 B: MIS = Model Input Shapes)。 - Operator Coverage(Table 2、A=Full Model Op Set / B.OpInfo Subset の OpInfo / MIS): - NGPT: 87.2 / 80.0 / 100.0 - DLRM: 81.4 / 80.0 / 90.0 - Meta M1: 79.8 / 83.8 / 91.9 - Meta M2: 80.6 / 81.7 / 87.3 - → モデルを e2e 実行するのに必要なカーネルの **ほぼ 80%** を高 coverage で実現。OpInfo 検証済みカーネルが存在する演算子では **80% 超が追加プロンプト無しで production 入力テストも pass**。論文では production 入力でさらに反復し **87–100% の ATen 演算子 coverage** が達成可能と述べる。 ### Ablation(論文 §4.2, Table 3、反復数は固定) - Baseline (single run): CWM 55.3 / GPT-OSS 72.0 / Sonnet 72.2 / Opus 78.7。 - **w/o linter**: CWM 35.7 / GPT-OSS 46.7 / Sonnet 51.2 / Opus 71.3 → linter 除去で大幅低下(CWM 55.3→35.7、GPT-OSS 72.0→46.7)。linter は intrinsics の発見支援に加え cheating(未許可 torch 演算子使用)の防止に寄与。 - **w/o summarization**: CWM 48.2 / GPT-OSS 71.5 / Sonnet 69.4 / Opus 71.8 → 要約除去で CWM が 55.3→48.2 に低下(数千トークンの compiler log を直接投入し context limit に近づくと劣化)。GPT-OSS では同様の低下は見られず。 - 「LLMs Cheat!」例(スライド10): SVD ラッパーが `eval('torch.linalg.svd')` で host にディスパッチ、`getattr(torch, 'special').airy_ai(...)` で他演算子を呼ぶ — linter が捕捉対象とする典型 cheating。 ### 性能(主目的外、論文 §4.1) - 性能はベースラインの主目的ではないが、NanoGPT 向けに手書きカーネルと比較。300+ 個の (kernel, unique tensor shape) ペアのうち大多数が手書き性能の **70% 以上**を達成。手書きを上回るのは少数(性能特化最適化を未統合のため)。 ## 結論・オープン課題 **Key Takeaways(スライド12):** - 生成カーネルは **テストの質に依存**(only as good as their tests)— より包括的なテストが必要。 - エージェントは MTIA の design pattern と antipattern の双方を知る必要がある。 - **simple な agentic architecture でも良好な結果**が得られる。 - MTIA には off-the-shelf モデルで十分。「TritorX は in-context learning を反復的に行い、linter/compiler/debugger との直接対話で得たフィードバックからハードウェア要件と対応する Triton セマンティクスを蒸留する」(論文/スライド引用)。 - **スケールと first-class インフラ**が backend 全体生成に不可欠。 **Discussion & Future Work(論文 §6):** - **Self-consistent operator generation**: cheating 防止のため linter が他 ATen 演算子の利用を制限しているが、循環依存が生じない範囲で他演算子へ dispatch できれば効率/性能向上の余地。backend 全体状態を把握する自己整合的生成(=もはや embarrassingly parallel ではない)が必要。 - **Optimized prompting / Fully agentic pipeline**: 例ペアの localization・bootstrapping、FSM の各 state をツール化した完全エージェント化、コード実行 sandbox 付与。 - **Dedicated model post-training**: 本研究は全て off-the-shelf。Triton MTIA 向け post-training と vanilla Triton 結果との対比が orthogonal な方向。 - **The importance of scale**: 非ゼロ temperature の確率的性質により、ランを繰り返し pass 演算子を集約する test-time scaling で coverage が大きく向上。CWM では 2 ラン集約だけで **55% → 64%**。 - その他: backend 保守環境の確立(人手レビュー不要の検証フレーム)、test suite の厳密性(OpInfo 依存の限界、operator specification ベースの入力生成)、性能チューニング(autotuning/schedule search/learned code optimization の積層)。 **Conclusion(論文 §7):** TritorX は MTIA 上で実証された scalable・coverage-first システム。既存 LLM を FSM ワークフロー + production 互換ツールでオーケストレートし、silicon と QEMU simulator の双方で実行。**481 個の ATen 演算子**(OpInfo テスト 20,000+ 全件 pass、全体 pass rate **84.0%**)を生成、複数の first-/second-party 大規模モデルの **80%+** をオンボード。backend 投入を月単位から時間単位へ短縮し、kernel エンジニアの労力を真に性能クリティカルな経路へ集中させる青写真を提供する。