2026__MLSys2026__Optimizing PyTorch Inference with LLM-Based Multi-Agent Systems

# Optimizing PyTorch Inference with LLM-Based Multi-Agent Systems > [!info] Talk metadata > - **会議:** [[MLSys2026]] Day 2 (May 19 / Tue)、Research Track Oral "Agentic AI 1 & Multimodal/Generative Models" セッション（13:45 - 14:00 PDT、第4発表） > - **URL:** https://mlsys.org/virtual/2026/oral/3823 > - **登壇者:** Kirill Nagaitsev（Northwestern University, PhD candidate、DOE Computational Science Graduate Fellow）。研究テーマは LLM ベースコード生成・性能エンジニアリングの推論時（inference-time）技術。 > - **共著者:** Kirill Nagaitsev（Northwestern University）、Luka Grbcic・Samuel Williams・Costin Iancu（Lawrence Berkeley National Laboratory） > - **フレームワーク:** PIKE（**P**yTorch **I**nference **K**ernel **E**volution）。変種 PIKE-B（branching search）と PIKE-O（OpenEvolve ベース）。 > - **URL / 関連研究:** コード `https://github.com/pike-project/pike`、発表者ページ `https://kirn.io`、連絡先 `[email protected]` > [!abstract] 概要（論文 PDF アブストラクト・忠実日本語訳） > 利用可能な GPU ハードウェア上で性能を最大化することは、現代の AI 推論システムにとって継続的な課題である。従来手法にはカスタム GPU カーネルの記述や、高水準コードを特定の GPU ターゲット向けにチューニングする専用モデルコンパイラの利用が含まれる。近年の研究は、LLM ベースのマルチエージェントシステムがこうしたチューニングを効果的に実行でき、既存コンパイラをしばしば上回り、手動カーネル開発を不要にしうることを示している。しかし、このタスクにおけるマルチエージェントシステムの動的特性（dynamics）は未解明のままである。本研究では、マルチエージェント PyTorch 最適化システムを比較するための論理的フレームワークを提示する。評価により、exploit 重視の戦略はエラー修正エージェントと組み合わせたときに最良の性能を発揮し、性能は最適化ステップの粒度（granularity）と相関することを示す。最良の実装は、PyTorch における多様な機械学習アーキテクチャを網羅するベンチマーク KernelBench の多様なタスクにおいて、H100 GPU 上で PyTorch Eager 比平均 **2.88x**（torch.compile 比 1.85x）の高速化を達成する。コードは公開されている: `https://github.com/pike-project/pike` > [!note] 出典に関する注記 > 正式タイトルは論文 PDF 表紙の "Optimizing PyTorch Inference with LLM-Based Multi-Agent Systems"。フレームワーク名は PIKE（PyTorch Inference Kernel Evolution）であり、文字起こしの "Pike" は正しい音だが正式には大文字略称。発表者名は Kirill Nagaitsev（Northwestern）で、文字起こしの司会発言は名前を欠落させていた。 ## 問題: カーネル最適化の高コスト - ML ワークロードの性能最適化が現代の中心課題。GPU の進化に合わせ AI/ML ソフトウェアも継続的に適応する必要があるが、GPU カーネル最適化は性能エンジニアリングの最も技術的に難しい領域の一つで、並列性・メモリ階層・スケジューリングへの熟達を要する（論文 Introduction）。 - 典型的な流れ: モデル内の性能クリティカルなカーネルを特定し、ハードウェアに近い言語（Triton, CUDA、加えて CUTLASS/CuTe など）でカーネルを書く（スライド「Optimizing for Target Hardware」）。これは **開発工数・専門性を要するエンジニアリング集約的プロセス**であり、LLM を使う場合は **LLM 推論時コスト（inference-time cost）** を発生させる。 - 本研究は LLM のみでこれを行い、推論時コストの最小化に焦点を当てる。例として FlashAttention カーネル（Dao et al., 2022）は汎用 PyTorch 実装比 7.6x を達成したが、組み込み演算子化までに相当な時間と専門性を要した（論文）。 - **search ベース技術と agent ベース技術**は単独では性能最適化のような複雑タスクで有効と示されてきたが、両者の交差（overlap）は未解明。本研究は「性能最大化と推論時コスト最小化」の目的下で両技術の相互作用を分析する（スライド6「Dynamics largely unexplored - How to minimize cost here?」）。 ## 手法: search + agent を融合した進化的パイプライン最適化対象の例として、unoptimized な PyTorch（ReLU self-attention 変種、`att = (q @ k.transpose(-2,-1)) * s` → `masked_fill` → `F.relu` → `att @ v`）を H100 向けに最適化する（スライド7）。パイプライン（スライド8–11、論文 Algorithm 1 PIKE-B）: 1. **初期サンプリング**: LLM を N 回サンプリングし「等価かつ最適化されたコードを H100 向けに書け」と指示（"Write equivalent, optimized code targeting H100"）。 2. **エラー修正・正当性修正ループ**: LLM 誘導の反復エラー修正。正当性（correctness）は元の unoptimized PyTorch モデルに対するランダム入力での **差分テスト（differential testing）** で判定。各エラー修正ループには予算カットオフ（最大試行回数）がある。 3. **進化的アルゴリズム**: エラー修正を通過した working solution 群を進化的アルゴリズムに伝播。LLM を再びサンプリングし **mutation（突然変異）** と **crossover（交叉）** を駆動。 4. これを反復し、探索中に見つかった最良性能の **最適化済み Triton/CUDA コード**を出力。 ### フレームワーク PIKE と論理ステップ PIKE（multi-agent evolutionary framework）は最適化プロセスを論理ステップに分解し、各ステップに紐づくオプション/パラメータで多様なマルチエージェント戦略を比較可能にする（スライド13–15、論文 Figure 1）。論理ステップ（Logical Steps）と主なオプション: - **Library**（解ライブラリ。short-term / long-term、islands、explore/exploit） - **Seed Selection**（シード選択。explore/exploit 比 ε:(1-ε) でランダム選択 or elite 選択） - **Prompt Construction**（mutation / crossover / brainstorming） - **Evaluation**（コンパイル・実行・正当性検証・性能計測） - **Post-Processing**（解ライブラリへ追加、exit 条件で best 解を返す）主なエージェント（論文）: - **COA（Code Optimization Agent）**: 問題文・元モデル・シードから最適化版を生成。 - **EFA（Error Fixing Agent）**: コンパイル/正当性失敗時に LLM クエリでコード修正。最大 m 回試行。 - **IBA（Initial Brainstorming Agent）**: 入力コードから最適化アイデア n 件を生成（PIKE-B では IBA を含めず純粋に exploit-focused にする ablation あり）。本トークでは特に **islands** と **error-fixing** の 2 パラメータに焦点を当てる。 ## RQ1: error-fixing への予算配分 **問い**: 300 LLM クエリ予算のうち、明示的 error-fixing ステップに割くべきか、進化的最適化探索に全振りすべきか（スライド12, 20）。実験設計: PIKE 進化的フレームワーク内で、生成された各解に対し error-fixing を **最大5試行**（5 attempts max）か **0試行**（0 attempts max）に設定。0 試行の場合は壊れた解に一切修正を施さない（スライド20）。結果（スライド19, 21、Level 3-pike geomean speedup、PyTorch Eager 比、全バー 300 LLM クエリ予算）: - **PIKE-B（EFA あり、Gemini 2.5 Pro）: 2.88x** - **PIKE-B（cheap EFA、Gemini 2.5 Flash）: 2.59x** - **PIKE-B（no EFA、0 attempts）: 1.98x** - 比較対象: **METR 1.40x**、**torch.compile 1.64x**、**TensorRT 1.41x** - error-fixing は critical。初期に壊れた進化解も distinct な error-fixing で修正する価値がある（**5試行以内に解の約70–80%が修正される**）。予算トラジェクトリ（スライド22、論文 Figure 3）: - (a) クエリ予算 0–300/タスク: **100クエリ以内で torch.compile/eager を上回る**。 - (b) 金銭コスト（$）軸: **$25 予算では cheap error-fixing（Gemini 2.5 Flash）が勝つ**。弱い Flash を error-fixing に使っても Gemini 2.5 Pro の error-fixing と同等性能。 - 予算は **300 LLM クエリ/タスク = Gemini 2.5 Pro で約 $30–50/タスク**（スライド18）。Gemini 2.5 Pro は $1.25/$10.00 per 1M tokens、Flash は $0.30/$2.50 per 1M tokens（論文、Google 2025）。 ## RQ2: exploit vs explore（islands） **問い**: 300 LLM クエリ regime で、トップ性能コードに全注力（exploit）すべきか、広い探索基盤を維持（explore）すべきか（スライド12, 23）。戦略設計（スライド23）: - **Exploit-oriented evolution**: elite 解の exploit に集中、population diversity の維持は最小限。 - **Explore-oriented evolution**: 複数 islands（並列サブ集団で多様性維持）を使い、non-elite を含む大きなライブラリからサンプリング。結果（スライド24、Level 3-pike geomean、全て error-fixing 構成あり）: - **PIKE-B（1 island, top-3 sampling）: 2.88x** — 最良の exploit 重視戦略。 - **PIKE-O'（1 island, broader sampling）: 2.75x** — broad サンプリングへの移行でわずかに劣化。 - **PIKE-O'（3 islands）: 1.99x** — 3 islands（最も explore 重視）で大幅劣化。 - 結論: **数百クエリ予算 regime では多 island による diversity 維持は割に合わない**。 Level 5（SOTA vision/language ブロック）でも同様（スライド25、Level 5 geomean）。exploit 方向のパラメータ調整（islands を 3→1、elite をより重くサンプリング、ライブラリ履歴を短縮）が単調に高速化を向上: - **PIKE-B: 2.57x**、PIKE-O: 2.25x、PIKE-O (mut,npar,1isl): 2.33x、PIKE-O (mut,npar,1isl,EO): 2.41x、PIKE-O (mut,npar,1isl,EO,SL): 2.47x - 比較対象: **METR 1.50x**、**torch.compile 1.29x**、**TensorRT 1.25x** ## 実験セットアップ - **ベンチマーク KernelBench**（Ouyang et al., 2025）: 250 個の self-contained PyTorch タスク。Level 1（matmul/conv 等の単純演算）、Level 2（複合）、Level 3（大規模 ML アーキテクチャの critical component、例 AlexNet・MinGPT）。一部 100+ 行の PyTorch/Python に及ぶ（スライド16）。 - **METR refined 版**を採用し Level 1/2 を除外、Level 3 から LSTM/GRU を除外（計測ノイズ対策）。この filtered 版を **Level 3-pike**（30タスク、平均 85 LoC）と呼ぶ（スライド17、論文 Table 1）。 - **Level 5（新規）**: 2024-25 SOTA vision/language ブロック（DeepSeek-V3, Llama 3, RWKV, SD3, Mamba-2, S4, Hunyuan Video）、14タスク、平均 493 LoC。最大タスクは HunyuanTransformer（HunyuanVideo ブロック、1,268 LoC）。 - **OpenEvolve**（Sharma, 2025、AlphaEvolve に着想を得た OSS の LLM ベース進化フレームワーク）を拡張。OpenEvolve 単体には distinct な agent（error-fixing agent 等）が無いため、エージェントと CUDA/Triton 用カスタム evaluator を追加（スライド16）。 - **モデル**: 複数モデルをテストするが本トークは **Gemini 2.5 Pro** に焦点。error-fixing には cheap ablation として Gemini 2.5 Flash を使用。 - **予算**: 300 LLM クエリ/タスク。**ターゲットハードウェア: NVIDIA H100 GPU**（bare metal、80 GB HBM3、PCIe 接続、Intel Xeon Platinum 8480+ CPU、40+ CPU スレッド）。evaluator は Docker コンテナ化、20+ の CUDA/Triton コンパイルを並列実行。 ## 結論・takeaway **Key Takeaways**（スライド26）: 1. ML モデルコンパイラ（例 torch.compile）はマルチエージェント進化的手法に容易に上回られる。 2. 明示的な error-fixing ループへ予算を割くことが critical。 3. 数百クエリ（100s of queries）予算 regime では exploit 重視戦略が最良。性能は最適化ステップの粒度と相関し、aggressive かつ exploit 重視のステップが限られた予算内でより良い結果を生む（論文 contributions）。最良解は CUDA・Triton の双方で実装された。 ## Q&A - **Q（Sham, University of Washington、正当性保証）**: エージェントを使って最適化したコードが入力プログラムのセマンティクスを保つことをどう保証するか。テストケースはどう生成するか。 **A**: 元の PyTorch モデルに対するランダム入力での差分テスト（differential testing）を用いる。テスト入力は `torch.rand`/`torch.randn` で生成。 - **Q（mutation の意味）**: 最初のスライドで例を再生成するステップと mutate するステップがあったが、サンプルは単なるトークン列。数百の異なるトークンを mutate するとはどういう意味か。 **A**: AlphaEvolve のような研究に由来する考え方で、**プログラム自体が population 内の個体**であり、LLM にプログラムを mutate させる。 - **Q（Martin, AMD[?]、カーネル/言語）**: 生成されたカーネルを見たか。よく知られた最適化が主か、より興味深いものか。Triton が適切な言語か、Triton より柔軟な別言語が面白い可能性は。 **A**: 観察された最適化パターンは monolithic fused kernels、tiling パターン、intrinsics の利用。言語については Triton が非常に効果的で、理由の一つは組み込みの auto-tuning（エージェントが auto-tuning 機能を活用できた）。探索すべき空間は大きい。CUDA の解も生成され、問題依存（problem-dependent）。 - **Q（Jen、正当性基準）**: 最初の問いの correctness signal について。カーネルを最適化して比較すると数値差が生じうる。エージェントによる「ごまかし」を防ぐ correctness 基準をどう決めるか。 **A**: 先行研究と同様に元モデルに対する **誤差許容（error tolerance）** を許す。これは完璧な correctness チェックではなく、多くの documented なケースがあることは認識しており改善余地はあるが、本研究では error-tolerance 方式を採った。