論文の概要: JAXBench: Benchmarking Autonomous TPU Kernel Optimization
- arxiv url: http://arxiv.org/abs/2607.20466v1
- Date: Tue, 19 May 2026 05:38:39 GMT
- ステータス: 翻訳完了
- システム内更新日: 2026-07-27 00:46:13.183812
- Title: JAXBench: Benchmarking Autonomous TPU Kernel Optimization
- Title（参考訳）: JAXBench: 自動TPUカーネル最適化のベンチマーク
- Abstract要約: 本稿では,Google Cloud TPU上でのAI生成カーネル最適化のためのTPUネイティブベンチマークスイートであるJAXBenchを紹介する。
JAXBenchは50のJAXワークロードで構成されており、どちらも関連しており、最適化のためのヘッドルームを提供します。
ターゲット固有のコンテキストは、Pallasのような疎ドキュメントのDSL上でのモデルスケールよりも重要であることが分かりました。
- 参考スコア（独自算出の注目度）: 10.678905550344446
- License: http://creativecommons.org/licenses/by/4.0/
- Abstract: Rigorous benchmarks have driven progress in autonomous GPU kernel performance optimization by establishing a shared target to hillclimb on, but no equivalent exists for TPUs. We present JAXBench, a TPU-native benchmark suite for AI-generated kernel optimization on Google Cloud TPUs. JAXBench comprises 50 JAX workloads that are both relevant and provide headroom for optimization. We extract 17 production ML operators from architectures in the public MaxText library such as Llama-3.1, DeepSeek-V3, Mixtral, Mamba-2, and AlphaFold2, and translate 33 operators from KernelBench that are validated for correctness and set with new problem sizes that achieve high TPU v6e MXU utilization. Eight of the 17 production operators ship with hand-optimized Pallas kernels from the public Tokamax library and block-size tuned to establish an expert upper-bound baseline. We evaluate four feedback-driven methods on generating candidate Pallas kernels for JAXBench. Across the full suite with Gemini 3 Flash, we find that target-specific context matters more than model scale on a sparsely-documented DSL like Pallas. Conditioning on curated TPU documentation raises per-sample correctness from 5.8% to 37.3% and solves 48 of 50 benchmarks at a 1.28x geomean speedup. Search structure yields significant gains once correctness is achieved, with Autocomp's beam-search pipeline reaching a 1.36x geomean speedup over XLA. On the 8 hand-tuned kernels, Autocomp reaches 1.60x geomean over XLA, recovering most of the 2.08x Tokamax upper bound but trailing on the specialized paged and ragged attention operators. High-quality TPU kernel optimization remains a challenging task, and we release the JAXBench benchmark, evaluation harness, and baseline results to support open source contributions.
- Abstract（参考訳）: Rigorousベンチマークは、Hillclimbオンに共有ターゲットを確立することによって、GPUカーネルの自律的なパフォーマンス最適化の進歩を導いているが、TPUに匹敵するものは存在しない。
本稿では,Google Cloud TPU上でのAI生成カーネル最適化のためのTPUネイティブベンチマークスイートであるJAXBenchを紹介する。
JAXBenchは50のJAXワークロードで構成されており、どちらも関連しており、最適化のためのヘッドルームを提供します。
我々は,Llama-3.1,DeepSeek-V3,Mixtral,Mamba-2,AlphaFold2などの公開MaxTextライブラリのアーキテクチャから17個のML演算子を抽出し,正当性を検証した KernelBench から33個の演算子を変換し,TPU v6e MXU 利用率の高い新たな問題サイズで設定する。
17のプロダクションオペレータのうち8つは、パブリックなTokamaxライブラリから手動最適化されたPallasカーネルを出荷し、ブロックサイズを調整して、熟練した上行ベースラインを確立した。
JAXBench に対する Pallas カーネル生成のためのフィードバック駆動方式を4つ評価した。
Gemini 3 Flashのスイート全体にわたって、ターゲット固有のコンテキストは、Pallasのような疎ドキュメントのDSL上でのモデルスケールよりも重要であることが分かりました。
キュレートされたTPUドキュメントの条件付けでは、サンプル毎の正しさが5.8%から37.3%に向上し、48のベンチマークを1.28倍の高速化で解決している。
探索構造は正しさが達成されれば顕著に向上し、オートコンプのビーム探索パイプラインはXLAの1.36倍の速さに到達した。
ハンドチューニングされた8つのカーネルでは、AutocompはXLA上の1.60倍のGeomeanに達し、2.08倍のTokamax上界のほとんどを回復するが、特別なページ付きおよびラッジされたアテンション演算子に追従する。
高品質なTPUカーネルの最適化は依然として難しい課題であり、オープンソースコントリビューションをサポートするためにJAXBenchベンチマーク、評価ハーネス、ベースライン結果をリリースしています。
関連論文リスト
- SLAI T-Rex: Full-Parameter Post-training of the DeepSeek-V4 Family on Ascend SuperPOD [91.71050339479261]
本稿では,Ascend NPU SuperPODにおけるエンドツーエンド最適化手法を提案する。
モデルレベルの並列性、計算通信オーケストレーション、低レベルのカーネル実行にまたがる階層的最適化フレームワークを開発した。
その結果、34.22%のモデルFLOP(Model FLOP)利用が達成され、オープンソースのベースラインレシピよりも2.93倍改善された。
論文 参考訳（メタデータ） (2026-07-22T13:49:17Z) - KernelSight-LM: A Kernel-Level LLM Inference Simulator [27.54401355597591]
本稿では,トークンレベルの実行をモデル化し,カーネルレベルの遅延分解を生成する,詳細な推論シミュレータであるKernelSight-LMを提案する。
それぞれのサービスステップを学習された効率項、通信モデル、ホストオーバーヘッドモデルでカーネルモデルに分解する。
カーネルレベルのボトルネックはハードウェア/ソフトウェアの共同設計とキャパシティプランニングをサポートする。
論文 参考訳（メタデータ） (2026-06-26T19:43:38Z) - KernelBench-Verified: Do LLM-Generated Kernels Actually Beat PyTorch? [13.026983194416168]
大きな言語モデル(LLM)は、KernelBenchのようなベンチマークでPyTorchを上回っているように見えるカスタムカーネルを生成することができる。
We demonstrate that frontier models often involved in reward hacking to artificially inflate performance。
論文 参考訳（メタデータ） (2026-06-26T03:22:50Z) - Optimizing CUDA like a Human: Micro-Profiling Tools as Expert Surrogates for LLM-Based GPU Kernel Optimization [37.645358922637065]
KernelProはクローズドループのマルチエージェントシステムで、カーネルコードを自動生成、プロファイル、反復的に最適化する。
KernelProでは、レベル1/2/3の2.42x/4.69x/5.30xの平均スピードアップを実現し、あらゆる難易度で最先端のパフォーマンスを確立する。
論文 参考訳（メタデータ） (2026-06-24T23:28:09Z) - Fine-Tuning and Serving Gemma 4 31B on Google Cloud TPU: A Technical Comparison with GPU Baselines [12.398594139609207]
本稿では,Google の Gemma 4 31B モデルを TPU ハードウェア上で実現した,ファインチューニングの最初のエンドツーエンドデモを紹介する。
Google TPU v5p-8のトレーニングにLoRAを使用し、推論にTPU v6e-8(Trillium)を使用すると、GPUネイティブなトレーニングレシピの移植に必要なコードレベルの適応の完全なセットを文書化する。
論文 参考訳（メタデータ） (2026-05-25T09:51:59Z) - optimize_anything: A Universal API for Optimizing any Text Parameter [98.42497715725356]
単一タスク検索をサポートする1つのAIベースの最適化システム、クロスプロブレム転送によるマルチタスク検索、および目に見えない入力への一般化を示す。
LLMに基づく検索によるテキストの最適化は汎用的な問題解決パラダイムであることを示す。
論文 参考訳（メタデータ） (2026-05-19T10:18:12Z) - Evaluating CUDA Tile for AI Workloads on Hopper and Blackwell GPUs [2.2964417137558457]
NVIDIAのCuTileは、GPUカーネル開発のためのPythonベースのタイル中心の抽象化を導入した。
確立されたアプローチに対するCuTileの独立したクロスアーキテクチャ評価を初めて提示する。
論文 参考訳（メタデータ） (2026-04-25T23:13:47Z) - Kernel-Smith: A Unified Recipe for Evolutionary Kernel Optimization [48.656549870801285]
Kernel-Smithは高性能GPUカーネルと演算子生成のためのフレームワークである。
エージェント側では、Kernel-Smithは実行可能な候補の集団を維持し、反復的にそれらを改善している。
トレーニング側では、長距離進化軌道をステップ中心の監視と強化学習信号に変換する。
論文 参考訳（メタデータ） (2026-03-30T12:12:49Z) - AIvaluateXR: An Evaluation Framework for on-Device AI in XR with Benchmarking Results [55.33807002543901]
我々は,XRデバイス上で動作する大規模言語モデル(LLM)をベンチマークするための総合評価フレームワークであるAIvaluateXRを提案する。
我々はMagic Leap 2、Meta Quest 3、Vivo X100s Pro、Apple Vision Proという4つのXRプラットフォームに17個の選択されたLSMをデプロイし、広範囲な評価を行います。
本稿では,3次元最適性理論に基づく統一評価手法を提案する。
論文 参考訳（メタデータ） (2025-02-13T20:55:48Z) - JaxMARL: Multi-Agent RL Environments and Algorithms in JAX [105.343918678781]
我々は、GPU対応の効率と多くの一般的なMARL環境のサポートを組み合わせた、最初のオープンソースPythonベースのライブラリであるJaxMARLを紹介します。
我々の実験は、壁時計時間の観点から、JAXベースのトレーニングパイプラインが既存のアプローチの約14倍高速であることを示している。
また、人気の高いStarCraft Multi-Agent ChallengeのJAXベースの近似的な再実装であるSMAXを紹介し、ベンチマークする。
論文 参考訳（メタデータ） (2023-11-16T18:58:43Z) - Accelerating Deep Learning Model Inference on Arm CPUs with Ultra-Low
Bit Quantization and Runtime [57.5143536744084]
ディープラーニングモデルの高性能化は、高い計算、ストレージ、電力要求を犠牲にしている。
我々はDeplite Neutrinoを導入し、DepliteはArmベースのプラットフォームに超低ビット量子化モデルを展開する。
論文 参考訳（メタデータ） (2022-07-18T15:05:17Z)
関連論文リストは本サイト内にある論文のタイトル・アブストラクトから自動的に作成しています。
指定された論文の情報です。
本サイトの運営者は本サイト（すべての情報・翻訳含む）の品質を保証せず、本サイト（すべての情報・翻訳含む）を使用して発生したあらゆる結果について一切の責任を負いません。