ParallelKernelBench:最先端のLLMはまだ高速なマルチGPUカーネルを記述できない(現時点では)
Together AI が公開したベンチマーク「ParallelKernelBench」は、現在の最先端大規模言語モデルが依然として効率的なマルチ GPU 用カーネルコードを生成する能力に重大な欠陥を抱えていることを示した。
キーポイント
LLM のハードウェア最適化能力の限界
最先端の LLM でも、複雑なマルチ GPU 環境向けの効率的なカーネルコードを自動生成する際に依然として困難を抱えていることが実証された。
専用ベンチマーク「ParallelKernelBench」の導入
Together AI が、LLM のハードウェア最適化能力を厳密に評価するために新たなベンチマークを公開し、業界標準となる可能性を示した。
生成コードの実用性の課題
モデルが生成するコードは構文上正しくても、実際のマルチ GPU 環境での実行効率やスケーラビリティにおいて不十分なケースが多いことが判明した。
影響分析・編集コメントを表示
影響分析
この分析は、LLM がハードウェアレベルの最適化まで自律的に処理できるという過度な期待に水を差し、開発現場における人間の監督と専門家の介入の重要性を再認識させるものです。今後は、LLM を活用したコード生成ツールが「効率的な実装」を支援する補助的な役割に留まるか、あるいは新しいアーキテクチャやトレーニング手法によってこの壁を突破できるかが注目されます。
編集コメント
LLM のコード生成能力が「動く」レベルから「高速で効率的に動作する」レベルへ移行するには、まだ多くのハードルが残っていることが浮き彫りになりました。
最前線のモデルは、87 の実世界の問題のうち 3 つに満たない割合しか解決できませんが、生成されたカーネルの中には公開されているどのものよりも優れたものが存在します。
image
image要約
大規模言語モデル(LLM)は GPU カーネル[1][2][3]の生成において驚くほど優秀になってきましたが、その進捗を測定する現在のベンチマークのほとんどは単一 GPU 向けです。実際の運用環境では、通信がボトルネックとなることが多く、通信オーバーヘッドは推論レイテンシの 20% を超える割合を占めることがあり[4]、計算能力の拡大速度がインターコネクト帯域幅の拡大速度を上回るにつれて、この格差はさらに広がっています。
ParallelKernelBench(PKB)は、マルチ GPU カーネル生成のためのベンチマークおよび評価フレームワークを提供しており、PyTorch と NCCL を置き換え、データを NVLink 上で直接転送する CUDA カーネルへ書き換えるというタスクを含む、実コードベースから抽出された 87 の問題を含んでいます。GPT-5.5、Gemini 3 Pro、Opus 4.7 など、最前線のコーディングモデルをテストしました。その評価結果は、全体的に顕著な性能の格差を示すものでした:正しく解決された問題は全体の 3 分の 1 に満たず、そのうちでも単純なベースラインを上回ったものは 4 分の 1 にも満たないものでした。
なぜ失敗するのか、どのようなパターンがあるのか、そして公に利用可能などの手法よりも驚くほど高速なカーネルを生成したいくつかのケースについて解説します。その中には、事前の最適化された公開リファレンスが存在しないNVIDIA NeMo-RL の GRPO 学習ループ向けのものも含まれています。
なぜマルチ GPU はシングル GPU カーネル生成とは異なるのか
LLM は GPU カーネル生成において進歩を遂げましたが、その進歩は主にシングル GPU の文脈で測定されてきました。しかし、実際の AI ワークロードはもはやそのような枠組みには収まりません。これらは複数の GPU にまたがり、パフォーマンスは単なるローカルな計算やメモリではなく、通信によってますます決定されるようになっています。この変化により、マルチ GPU カーネル生成は以下の 3 つの点で異なる問題となります。
- 設計空間が組み合わせ的に拡大します。実践者は、ハードウェアに適合させるためにテンソル並列化、エキスパート並列化、データ並列化、コンテキスト並列化、シーケンス並列化を組み合わせて使用しますが、それぞれの組み合わせが異なる通信パターンを生み出します。
- パフォーマンスモデルが変化します。シングル GPU の屋根線(roofline)は計算とメモリ帯域幅を中心に構築されていますが、マルチ GPU コードではボトルネックはしばしば相互接続(interconnect)となります。
- マルチ GPU カーネル生成には、GPU 間でのデータ転送方法(コピーエンジン、TMA、SM ロード/ストア、または NVLS を通じて行うか)や、その転送を計算と融合させるかどうかという、決定的に新しい設計上の選択が新たに生じます。
ParallelKernelBench
PKB は、モデルが純粋な torch.dist の使用を超えて、実際に本番環境向けのマルチ GPU カーネルを記述できるかをテストするために構築されました。各問題は、標準的な PyTorch + NCCL 実装とハードウェアトポロジの説明から始まります。その後、モデルは対称メモリを使用して GPU 間で直接通信する CUDA カーネルにその参照実装を置き換える必要があります。
PKB 評価パイプライン。各問題にはタスク、ハードウェアトポロジ、および PyTorch + NCCL の参照実装が提供され、モデルはカスタム CUDA カーネルを生成します。このカーネルは、正しさ、壁際速度向上率、通信の屋根線(communication roofline)に基づいて評価されます。
87 の問題が生産環境における並列処理の実際の範囲を網羅していることを保証するため、分散ワークロードの分類体系から問題を構築しました。まず、モデルがシャードされる主要な方法(テンソル、コンテキスト、データ、エキスパート、シーケンス、FSDP/ZeRO)と、それぞれが生み出す通信パターンを特定しました。次に、Megatron-LM、DeepSpeed、DeepEP、TensorRT-LLM、NeMo-RL などのシステムコードベースから、さらに GNN ルーティング、分散 FFT(高速フーリエ変換)、ガウシアンスプラッティングなど、LLM 以外のワークロードの長尾部分も含めて、その範囲をカバーする 87 の問題を選択しました。もう一つの利点は、PKB の参照実装が標準的な PyTorch + NCCL で記述されているため、このベンチマークが特定のハードウェア世代に縛られないことです。むしろ、次世代のハードウェアアーキテクチャとともに自然に進化するように設計されています。
標準的なトランスフォーマーブロックの並列化のための分類体系。異なるシャード戦略は、正規化、アテンション、MLP 全体にわたって固有の通信パターンを生み出し、ここでは代表的な Gemma3-27B のレイヤーを例に示しています。image
並列化の種類(左)とソースコードベース(右)にわたる PKB の問題カバレッジは、RL 後のトレーニング、LLM トレーニング、カーネルライブラリ、ビジョンモデル、GNN などを含んでいます。モデルを評価する前に、まず PyTorch + NCCL ベースラインが実際に性能向上の余地を残しているかを確認しました。通信を考慮したルーフライン分析によると、答えはイエスです:PKB の問題の多くは NVLink によってボトルネック化されており、ベースラインはハードウェアの上限から遠く離れています。したがって、次の質問は単純になります:モデルはこのギャップを埋められるでしょうか?
PKB における最先端モデルのパフォーマンス
あまり良くありません。 ゼロショット設定では、最良のモデルが 87 の問題のうち 28 を解決しましたが、そのうち PyTorch + NCCL ベースラインより高速なものはわずか 22 件でした。3 回の試行をサンプリングすることで最良の結果は改善され、正解数は 36 件、ベースラインより高速な解決策は 27 件に達しますが、fast1@3 は依然として 31% で頭打ちになります。
Category # GPT-5.5 Claude Opus 4.7 Gemini 3 Pro GLM-5.1 GLM-5.2 DeepSeek V4 Pro
@1→3 のパス
fast1 @1→3
@1→3 のパス
fast1 @1→3
@1→3 のパス
fast1 @1→3
@1→3 のパス
fast1 @1→3
@1→3 のパス
fast1 @1→3
@1→3 のパス
fast1 @1→3
@1→3 のパス
fast1 @1→3
Collective Primitive (集合プリミティブ)
8
3→53→4
4→53→4
6→62→2
2→32→3
1→31→1
0→00→0
Tensor Parallel (テンソル並列)
17
2→21→2
1→31→3
3→31→1
1→10→0
0→10→0
0→00→0
Sequence Parallel (シーケンス並列)
2
1→11→1
0→00→0
0→00→0
0→00→0
0→00→0
0→00→0
Context Parallel (コンテキスト並列)
12
7→85→6
7→73→4
7→95→7
2→20→0
2→20→0
2→30→1
Pipeline Parallel (パイプライン並列)
1
0→00→0
0→00→0
0→00→0
0→00→0
0→00→0
0→00→0
Data Parallel (データ並列)
9
1→21→2
1→11→1
0→10→0
0→00→0
0→00→0
0→00→0
Expert Parallel (エキスパート並列)
11
3→32→3
2→60→1
3→40→1
0→00→0
1→10→0
0→00→0
FSDP / ZeRO
9
3→43→4
3→33→3
1→10→1
1→10→0
0→00→0
0→00→0
Vocab Parallel
4
2→21→1
1→21→2
2→22→2
0→00→0
0→10→1
0→00→0
Graph Parallel
6
2→42→3
0→20→2
1→11→1
0→00→0
0→00→0
0→00→0
Geometric / Spatial
5
3→33→3
1→20→0
0→10→1
0→00→0
0→00→0
0→00→0
Other
2
1→20→1
0→00→0
1→21→2
0→00→0
0→10→0
0→00→0
More models →
Single-shot LLMs struggle on PKB. pass@k reports correctness on the best of k attempts; fast1@k counts solutions that are both correct and outperform the PyTorch + NCCL baseline. Repeated sampling improves both metrics, but fast1@3 still peaks at just 31% (GPT-5.5), suggesting fundamental limitations beyond sampling noise.
成功事例は、集約プリミティブ、テンソル並列 GEMM(行列乗算)、およびアテンションヘッド全体にシーケンス次元を分割するオール・トゥー・オール通信を用いる Ulysses スタイルのコンテキスト並列化といった、よく知られたパターンに集中しています。これらはオープンソースコード上で最も目立つマルチ GPU スタックの一部であり、モデルはこれらの分野に対してより強い事前知識(プライア)を持っている可能性があります。
失敗事例は CUDA 構文よりも深い問題を示唆しています。 weaker なモデルはコンパイル自体に失敗することが多いですが、推論能力が強いモデルであっても、コンパイルは通るものの誤った結果を返すカーネルを生成することが頻繁にあります。難しいのは、ランク間の調整、データ分割、および集約操作の順序付けに関する推論です。
生成されたカーネルも、非常に限られた通信メカニズムセットしか使用していません。大半はコピーエンジンや SM のロード/ストア命令に依存しており、TMA や NVLS といったより専門的なメカニズムはほとんど見当たりません。多くの場合、モデルはピークパフォーマンスに必要なメカニズムを選択していません。これはデータ不足とハードウェアの複雑さの組み合わせによるものと考えられます。新しいプリミティブである TMA(Tensor Memory Accelerator)や NVLS(NVIDIA Link System)を利用するには、非同期コピー記述子や最新の NVLink トポロジなど、細部まで入り組んだハードウェア固有の抽象化を扱う必要があり、モデルにはこれに対する強い事前知識がありません。その結果、「動作する」分散カーネルと「真に最適化された」カーネルの間には大きな隔たりが生じています(これは今後の研究における明確な目標です!)
frontier モデル全体における fastp スコアの分布。最良のモデル(GPT-5.5)でさえ、速度向上の閾値が増加するにつれて急速に性能が低下します;p = 1.0 の場合、fast1 は約 31% で頭打ちになりますimage
モデル別の PKB(ParallelKernelBench)の失敗モード。 weaker モデルは主にコンパイル時に失敗しますが、推論能力が強いモデルでは、コンパイルは成功するものの結果が誤り(出力不一致)となるか、デッドロックが発生するカーネルを生成することがより頻繁に見られますimage
GEMM + All-Gather のプロファイラトレース:生成された CUDA カーネル(87.9 μs)は、NVLink を介して計算と通信をオーバーラップさせ、PyTorch + NCCL のリファレンス(320.6 μs)を上回る性能を発揮しますが、依然として理論上の roofline(4.99 μs)には及びません。自然な次のステップは、人間がカーネルライターとして使用するのと同じフィードバックループをモデルに提供することです。私たちは Gemini 3 Pro を agentic harness にラップし、リポジトリへのアクセス権、ターミナル、コンパイラの出力、正しさのテスト、速度測定、そして過去の試行結果を利用できるようにしました。一度カーネルを生成して終了するのではなく、モデルはコンパイルを実行し、ベンチマークを走らせ、失敗を検証し、修正を加えることができます。
エージェント型評価ループ:モデルはファイルを作成し、正誤テストとベンチマークを実行し、解決策が合格するかステップ予算が尽きるまで反復します。これは役立ちましたが、実用的な効果は限定的でした。Gemini 3 Pro は、単一ショット設定で 24 件だった正解数が 87 件中 35 件に改善され、そのうち 26 のカーネルが PyTorch + NCCL ベースラインを上回りました。この向上は、構文エラー、形状の誤り、単純なランタイムバグの修正によるものです。約 20 回の改良ステップの後、性能は頭打ちとなりました。フィードバックはモデルが分散カーネルをデバッグするのを助けますが、残りの失敗はより大きなギャップを浮き彫りにしています:ランク間の調整、通信順序、および GPU から GPU への転送メカニズムの最適な選択に関する推論能力の欠如です。
新しいカーネル
NCCL コレクティブを直接の NVLink 読み書きに置き換える例:参照実装では NCCL コレクティブに加えて追加的な形状変更が必要ですが、生成されたカーネルは計算操作と直接のピアメモリ読み込みを融合させています。標準的なコレクティブにおける速度向上に加え、単一の生成試行で、最適化された公開参照が存在しないワークロードに対して真に新しい高性能なカーネルが偶然発見されることもあります。この機能は、AI 駆動型最適化のより広範な可能性を示す予兆を提供するものであり、その成功はトランスフォーマーに限定されません。モデルは状態空間モデル、ゲノムパイプライン、マルチモーダル RL ループといった領域でも良好な性能を発揮できます。これらは、主流の LLM スタックと比較して専門化されたカーネルが依然としてほとんど最適化されていない分野です。
以下に 3 つの例を示します。いずれも NCCL コレクティブを融合対称メモリカーネルに置き換え、データを直接 NVLink を経由で転送するものです。各ケースは 4 枚の H100 GPU および 100 回のランダム化された実行において正しさが検証されており、速度測定は ThunderKittens 2.0[[5]](https://nathanjpaek.github.io/parallel-kernel-bench/#ref-thunderkittens) のベンチマーク手法に従って行われました(ビット単位の同一入力、L2 対応入力グループ、500 回のウォームアップ反復、および 100 回の連続プロファイリング反復)。
NeMo vocab-parallel log-prob with top-k/top-p filtering (Gemini 3 Pro)
NVIDIA NeMo-RL の GRPO 学習ループにおける中核的なステップ:top-*k*/top-*p* でフィルタリングされた分布の下で、語彙をランク間で分割した対数尤度(log-probabilities)を計算する。PyTorch + NCCL ベースラインでは、フィルタリング前に全ランクにわたって完全な語彙を集約しますが、生成されたカーネルはそれらの通信操作を完全に省略し、対称メモリを使用してシャードをインラインで並べ替えながら、log-softmax、トークン抽出、ターゲット集約を単一の warp-shuffle 還元演算に融合します。
シーケンス長 *M* ∈ {1024, 2048, 4096, 8192, 16384} にわたる PyTorch + NCCL に対する速度向上比。設定:ランク数 4、*B* = 1、*V* = 32,000(各ランクのローカル語彙サイズは 8,000)、top-*k* = 10、top-*p* = 0.9。
Hyena forward context parallelism (GPT-5.5)
Hyena オペレータにおける文脈並列な順伝播パス。ここで FFT 畳み込み(FFT convolution)にはシーケンス全体のコンテキストが必要です。参照実装では、繰り返し行われる all_to_all 呼び出しを通じてシーケンス分割レイアウトとチャネル分割レイアウトを交互に切り替えますが、生成されたカーネルは入力を単一の対称アロケーションにパッキングし、NVLink を経由してリモートスライスをストリーミングしながら、ゲート計算と再インデックス化を統一されたパスで実行します。特筆すべきは、より長いシーケンス長において性能が低下することです。
ローカルシーケンス長 l ∈ {1024, 2048, 8192, 16384} における 4 GPU 環境での PyTorch + NCCL に対する速度向上率。正解性と速度は 100 回の試行で測定されました。
SAM 3 の全結合マスク IoU サプレッション (GPT-5.5)
SAM 3 動画セグメンテーションにおける GPU 間重複抑制:各フレーム処理後に、ランク間で予測領域を交差率(IoU: Intersection-over-Union)で比較し、重なるマスクをゼロにします。ベースラインでは可変長の all_gather コレクティブと密行列乗算を使用しますが、生成されたソリューションはこれを、マスクをビットパックしてハードウェアの popcount 機能を用いて対向重複を計算する一連の対称メモリカーネルパイプラインへと圧縮しています。
固定マスク解像度 256 × 256、IoU しきい値 = 0.7 の条件下で、追跡対象数 N ∈ {16, 32, 64, 128, 256, 512} における PyTorch + NCCL に対する速度向上率。
さらに研究を進めるために
PKB は現在、意図的にノード内 NVLink にスコープを限定しています。自然な拡張先は、デバイス側 API の状況がより若く、他のアクセラレータやトポロジ(TPU など)が存在するノード間ファブリック(RoCE, InfiniBand)です。また、より高レベルの抽象化が役立つか、あるいは害になるかについても知りたいと考えています。PKB はすでに Triton や ParallelKittens のソリューションを受け付けており、NCCL GIN や NVSHMEM といった新興インターフェースも調査対象として価値があります。これらのパラダイムへのサポート拡大は、AI エージェントが多様なプログラミングモデルやハードウェア抽象化をどのようにナビゲートするかという点におけるさらなる研究を促すことになります。
より広範な目標は、このすべてに背後にある困難な問題に対する具体的なターゲットを設定することです。つまり、大規模分散インフラストラクチャを自律的に最適化・管理できる LLM システムです。言語モデルのトレーニングや推論のために特別に構築されたインフラストラクチャにおいて、この自律性を達成することは、最終的に自分たちのエンドツーエンドの研究エンジニアリングを処理できる AI エージェントへの格差を埋める可能性があります。
私たちはそれを推進するために PKB をオープンベンチマークとして公開します。深く掘り下げたい場合や問題を提供したい場合(特にノード間に関するもの)、ぜひご連絡ください。willychan2022@gmail.com または npaek@together.ai までメールをお送りください!
参考文献
- Standard Kernel. Reimagining Kernel Generation at the PTX Layer: An LLM System Learning from DSLs to Outperform Them. Standard Kernel Blog, April 2026.
- Robert Tjarko Lange, Qi Sun, Aaditya Prasad, Maxence Faldor, Yujin Tang, and David Ha. Towards Robust Agentic CUDA Kernel Benchmarking, Verification, and Optimization. arXiv:2509.14279, 2025.
- Carlo Baronio, Pietro Marsella, Ben Pan, Simon Guo, and Silas Alberti. Kevin: Multi-Turn RL for Generating CUDA Kernels. arXiv:2507.11948, 2025.
- Raja Gond, Nipun Kwatra, and Ramachandran Ramjee. TokenWeave: Efficient Compute-Communication Overlap for Distributed LLM Inference. arXiv:2505.11329, 2026.
- Stuart Sul and Chris Ré. ThunderKittens 2.0: Even Faster Kernels for Your GPUs. Hazy Research, February 2026.
@misc{chan2026parallelkernelbench,
title = {ParallelKernelBench: Can LLMs Write Fast Multi-GPU Kernels?},
author = {Willy Chan and Nathan Paek and Simon Guo and Simran Arora and Daniel Y. Fu},
year = {2026},
url = {https://nathanjpaek.github.io/parallel-kernel-bench/}
}
原文を表示
The best frontier model solves under a third of 87 real-world problems — but a few generated kernels beat anything publicly available.


Summary
LLMs have gotten surprisingly good at writing GPU kernels[1][2][3], but almost all current benchmarks measuring that progress are single-GPU. In production, communication is often the bottleneck: communication overhead can account for over 20% of inference latency[4], and that gap keeps widening as compute scales faster than interconnect bandwidth.
ParallelKernelBench (PKB) offers a benchmark and evaluation framework for multi-GPU kernel generation and includes 87 problems from real codebases where the task is replacing PyTorch + NCCL with a CUDA kernel that moves data directly over NVLink. We tested frontier coding models such as GPT-5.5, Gemini 3 Pro, Opus 4.7, and others. The evaluation revealed significant performance gaps across the board: under a third of problems were solved correctly, and fewer than a quarter of those beat the naive baseline.
We'll cover why they fail, what the patterns look like, and a few cases where models surprisingly produced kernels faster than anything publicly available, including one for NVIDIA NeMo-RL's GRPO training loop, which has no prior optimized public reference.
Why multi-GPU is different from single-GPU kernel generation
LLMs have made progress on GPU kernel generation, but that progress has mostly been measured on a single GPU. Production AI workloads no longer fit that frame: they span multiple GPUs, and performance is increasingly shaped by communication rather than just local compute and memory. That shift makes multi-GPU kernel generation a different problem in three ways:
- The design space expands combinatorially. Practitioners compose tensor, expert, data, context, and sequence parallelism to fit the hardware, and each composition creates a different communication pattern.
- The performance model changes. A single-GPU roofline is built around compute and memory bandwidth. In multi-GPU code, the bottleneck is often the interconnect.
- Multi-GPU kernel generation introduces a critical new design choice: how to move data between GPUs — through the copy engine, TMA, SM load/store, or NVLS — and whether to fuse that movement with compute.
ParallelKernelBench
We built PKB to test whether models can move beyond pure torch.dist and actually write production multi-GPU kernels. Each problem starts from a standard PyTorch + NCCL implementation and a description of the hardware topology. The model then has to replace that reference with a CUDA kernel that communicates directly across GPUs using symmetric memory.
To make sure the 87 problems cover the real space of production parallelism types, we built them from a taxonomy of distributed workloads. First, we identified the major ways models get sharded — tensor, context, data, expert, sequence, and FSDP/ZeRO — along with the communication patterns each one creates. Then we chose 87 problems to cover that space taken from the codebases of systems like Megatron-LM, DeepSpeed, DeepEP, TensorRT-LLM, NeMo-RL, as well as a long tail of non-LLM workloads: GNN routing, distributed FFTs, Gaussian splatting, etc. Another benefit is that because PKB references are written in standard PyTorch + NCCL, the benchmark is not tied to any single, particular hardware generation. Instead, it is designed to naturally evolve alongside next-generation hardware architectures.
Before evaluating models, we first checked whether the PyTorch + NCCL baselines leave real headroom. A communication-aware roofline says yes: most PKB problems are bottlenecked by NVLink, and the baselines run far below the hardware ceiling. So the next question is simple: can models close that gap?
How frontier models do on PKB
Not well. In the zero-shot setting, the best model solves 28 of 87 problems, and only 22 of those solutions are faster than the PyTorch + NCCL baseline. Sampling three attempts improves the best result to 36 correct solutions and 27 faster-than-baseline solutions, but fast1@3 still tops out at 31%.
Category
#
GPT-5.5
Claude Opus 4.7
Gemini 3 Pro
GLM-5.1
GLM-5.2
DeepSeek V4 Pro
pass @1→3
fast1 @1→3
pass @1→3
fast1 @1→3
pass @1→3
fast1 @1→3
pass @1→3
fast1 @1→3
pass @1→3
fast1 @1→3
pass @1→3
fast1 @1→3
Collective Primitive
8
3→53→4
4→53→4
6→62→2
2→32→3
1→31→1
0→00→0
Tensor Parallel
17
2→21→2
1→31→3
3→31→1
1→10→0
0→10→0
0→00→0
Sequence Parallel
2
1→11→1
0→00→0
0→00→0
0→00→0
0→00→0
0→00→0
Context Parallel
12
7→85→6
7→73→4
7→95→7
2→20→0
2→20→0
2→30→1
Pipeline Parallel
1
0→00→0
0→00→0
0→00→0
0→00→0
0→00→0
0→00→0
Data Parallel
9
1→21→2
1→11→1
0→10→0
0→00→0
0→00→0
0→00→0
Expert Parallel
11
3→32→3
2→60→1
3→40→1
0→00→0
1→10→0
0→00→0
FSDP / ZeRO
9
3→43→4
3→33→3
1→10→1
1→10→0
0→00→0
0→00→0
Vocab Parallel
4
2→21→1
1→21→2
2→22→2
0→00→0
0→10→1
0→00→0
Graph Parallel
6
2→42→3
0→20→2
1→11→1
0→00→0
0→00→0
0→00→0
Geometric / Spatial
5
3→33→3
1→20→0
0→10→1
0→00→0
0→00→0
0→00→0
Other
2
1→20→1
0→00→0
1→21→2
0→00→0
0→10→0
0→00→0
Total
87
28→3622→27
20→3112→20
24→3012→19
6→72→3
4→91→2
2→30→1
More models →
Single-shot LLMs struggle on PKB. pass@k reports correctness on the best of k attempts; fast1@k counts solutions that are both correct and outperform the PyTorch + NCCL baseline. Repeated sampling improves both metrics, but fast1@3 still peaks at just 31% (GPT-5.5), suggesting fundamental limitations beyond sampling noise.
The successes are concentrated in familiar patterns: collective primitives, tensor-parallel GEMMs, and Ulysses-style context parallelism, which splits the sequence dimension across attention heads using all-to-all communication. These are the parts of the multi-GPU stack most visible in open-source code, so models likely have stronger priors for them.
The failures suggest a deeper issue than CUDA syntax. Weaker models often fail to compile, but stronger reasoning models frequently produce kernels that compile and return incorrect results. The hard part is reasoning about rank coordination, data partitioning, and collective ordering.
Generated kernels also use a very narrow set of communication mechanisms. Most rely on copy engines or SM load/store instructions, while more specialized mechanisms such as TMA and NVLS are almost absent. In many cases, models do not choose the mechanism needed for peak performance. We attribute this to a combination of data scarcity and hardware complexity: newer primitives like TMA and NVLS require navigating intricate, hardware-specific abstractions, such as asynchronous copy descriptors or newer NVLink topologies, where models lack strong priors. That leaves a large gap between “working” distributed kernels and genuinely optimized ones (a clear target for future work!)
The natural next step is to give the model the same feedback loop a human kernel writer would use. We wrapped Gemini 3 Pro in an agentic harness with access to the repository, a terminal, compiler output, correctness tests, speed measurements, and its previous attempts. Instead of producing one kernel and stopping, the model could compile, run the benchmark, inspect failures, and revise.
This helped, but the practical gains are modest. Gemini 3 Pro improved from 24 correct solutions in the single-shot setting to 35 out of 87, with 26 kernels beating the PyTorch + NCCL baseline. The gains came from fixing syntax errors, shape mistakes, and simple runtime bugs. After roughly 20 refinement steps, performance plateaued. Feedback helps models debug distributed kernels, but the remaining failures highlight a much bigger gap: an inability to reason about rank coordination, communication ordering, and the optimal choice of GPU-to-GPU transfer mechanisms.
Net new kernels
Beyond speedups on standard collectives, single-shot generation occasionally surfaces genuinely new high-performance kernels for workloads with no optimized public reference. This capability offers a preview of the broader potential for AI-driven optimization; wins are not limited to Transformers, as models can do well on state-space models, genomic pipelines, and multimodal RL loops — domains where specialized kernels remain largely unoptimized compared to mainstream LLM stacks.
Below are three examples, each replacing NCCL collectives with fused symmetric-memory kernels that move data directly over NVLink. Each was verified for correctness over 4 H100 GPUs and 100 randomized runs; speed measurements follow the benchmarking methodology in ThunderKittens 2.0[[5]](https://nathanjpaek.github.io/parallel-kernel-bench/#ref-thunderkittens)(bitwise-identical inputs, L2-aware input groups, 500 warmup iterations, and 100 back-to-back profiling iterations).
NeMo vocab-parallel log-prob with top-k/top-p filtering (Gemini 3 Pro)
A core step in NVIDIA NeMo-RL's GRPO training loop: compute vocabulary-sharded log-probabilities under a top-*k*/top-*p* filtered distribution. The PyTorch + NCCL baseline gathers the full vocabulary across ranks before filtering; the generated kernel skips those collectives entirely, using symmetric memory to permute shards inline while fusing log-softmax, token extraction, and target gather into a single warp-shuffle reduction.
Hyena forward context parallelism (GPT-5.5)
The context-parallel forward pass for the Hyena operator, where FFT convolution needs sequence-global context. The reference alternates between sequence- and channel-sharded layouts via repeated all_to_all calls; the generated kernel packs inputs into one symmetric allocation and streams remote slices over NVLink, computing gating and reindexing in a unified pass. Notably, performance degrades at longer sequence lengths.
SAM 3 all-gathered mask IoU suppression (GPT-5.5)
Cross-GPU duplicate suppression for SAM 3 video segmentation: after each frame, ranks compare predicted regions via intersection-over-union and zero out overlapping masks. The baseline uses variable-length all_gather collectives plus a dense matmul; the generated solution collapses this into a pipeline of symmetric-memory kernels that bitpack masks and compute pairwise overlap with hardware popcount.
Further research
PKB is intentionally scoped to intra-node NVLink today. The natural extensions are inter-node fabrics — RoCE, InfiniBand — where the device-side API landscape is younger, and other accelerators and topologies like TPUs. We'd also like to know whether higher-level abstractions help or hurt: PKB already accepts Triton and ParallelKittens solutions, and emerging interfaces like NCCL GIN and NVSHMEM are worth studying as targets. Expanding support to these paradigms will encourage further research into how AI agents navigate diverse programming models and hardware abstractions.
The broader goal is a concrete target for the harder problem behind all of this: LLM systems that can autonomously optimize and manage large-scale distributed infrastructure. For infrastructure custom-built for language model training and inference, achieving this autonomy could ultimately bridge the gap to AI agents capable of handling their own end-to-end research engineering.
We're releasing PKB as an open benchmark to push on that. If you want to dig in or contribute problems — especially inter-node ones — we'd love to hear from you. Feel free to email willychan2022@gmail.com or npaek@together.ai!
References
- Standard Kernel. Reimagining Kernel Generation at the PTX Layer: An LLM System Learning from DSLs to Outperform Them. Standard Kernel Blog, April 2026.
- Robert Tjarko Lange, Qi Sun, Aaditya Prasad, Maxence Faldor, Yujin Tang, and David Ha. Towards Robust Agentic CUDA Kernel Benchmarking, Verification, and Optimization. arXiv:2509.14279, 2025.
- Carlo Baronio, Pietro Marsella, Ben Pan, Simon Guo, and Silas Alberti. Kevin: Multi-Turn RL for Generating CUDA Kernels. arXiv:2507.11948, 2025.
- Raja Gond, Nipun Kwatra, and Ramachandran Ramjee. TokenWeave: Efficient Compute-Communication Overlap for Distributed LLM Inference. arXiv:2505.11329, 2026.
- Stuart Sul and Chris Ré. ThunderKittens 2.0: Even Faster Kernels for Your GPUs. Hazy Research, February 2026.
@misc{chan2026parallelkernelbench,
title = {ParallelKernelBench: Can LLMs Write Fast Multi-GPU Kernels?},
author = {Willy Chan and Nathan Paek and Simon Guo and Simran Arora and Daniel Y. Fu},
year = {2026},
url = {https://nathanjpaek.github.io/parallel-kernel-bench/}
}
関連記事
ポッドキャスト:AI に自我があるなら『帝国時代 II』にもあるという論文について
Matthew が、大規模言語モデルに自我があると仮定した場合、古典的ゲーム『帝国時代 II』も同様に自我を持つと主張する興味深い論文を紹介した。
トークン終末が到来:企業、AI への支出抑制に躍起
コンサルティング大手のアクセンチュアは、非技術職による PDF からスライド作成などの些細なタスクでの AI トークン予算の浪費を防ぐため、業界全体で急激に増加するトークン支出を抑制しようとしている。
2026 年にローカルで実行可能なトップ 7 つのコーディングモデル
KDnuggets が選定した、2026 年版のローカル環境で動作する主要な 7 つのコード生成 AI モデルを紹介している。
今日のまとめ
AI日報で今日の重要ニュースをまとめ読み