ParallelKernelBench: Frontier LLMs can't write fast multi-GPU kernels (yet)

ParallelKernelBench: 前沿大语言模型尚无法快速编写多GPU内核（目前）

摘要

大语言模型在编写GPU内核方面取得了令人惊讶的成果 [1][2][3]，但目前几乎所有衡量这一进展的基准测试都是单GPU的。在实际生产中，通信通常是瓶颈：通信开销可能占推理延迟的20%以上 [4]，并且随着计算速度的提升远超互连带宽，这一差距还在不断扩大。

ParallelKernelBench（PKB）为多GPU内核生成提供了一个基准测试和评估框架，其中包括来自真实代码库的87个问题，任务是将PyTorch + NCCL替换为通过NVLink直接传输数据的CUDA内核。我们测试了前沿的编码模型，如GPT-5.5、Gemini 3 Pro、Opus 4.7等。评估结果显示了显著的性能差距：不到三分之一的问题被正确解决，且其中正确解决的问题中，不到四分之一的性能优于简单的基线。

我们将讨论它们为何失败，失败的模式是什么，以及一些模型出人意料地生成了比目前公开可用的任何内核都更快的内核的案例，其中包括一个针对NVIDIA NeMo-RL的GRPO训练循环的内核，该内核此前没有公开的优化参考。

为何多GPU内核生成与单GPU不同

大语言模型在GPU内核生成方面取得了进展，但这些进展大多是在单GPU上进行评估的。实际的AI工作负载不再适用这个框架：它们跨越多个GPU，性能越来越受到通信的影响，而不仅仅是本地计算和内存。这种转变使多GPU内核生成在三个方面成为一个不同的问题：

• 设计空间以组合方式扩展。从业者会将张量、专家、数据、上下文和序列并行性组合在一起以适应硬件，而每种组合都会产生不同的通信模式。

• 性能模型发生变化。单GPU的性能上限围绕计算和内存带宽构建。在多GPU代码中，瓶颈通常是互连。

• 多GPU内核生成引入了一个关键的新设计选择：如何在GPU之间移动数据——通过复制引擎、TMA、SM加载/存储或NVLS，以及是否将数据移动与计算融合。

ParallelKernelBench

我们构建了PKB来测试模型是否能够超越纯粹的torch.dist，真正编写生产级别的多GPU内核。每个问题都从标准的PyTorch + NCCL实现和硬件拓扑的描述开始。模型随后必须用一个通过对称内存直接在GPU之间通信的CUDA内核来替换该参考实现。

PKB评估流程。每个问题提供一个任务、硬件拓扑和PyTorch + NCCL参考实现；模型生成一个自定义的CUDA内核，并根据正确性、时钟速度提升和通信性能上限进行评估。

为了确保这87个问题能覆盖实际生产中并行类型的真实空间，我们根据分布式工作负载的分类构建了这些问题。首先，我们识别了模型分片的主要方式——张量、上下文、数据、专家、序列以及FSDP/ZeRO，以及每种方式所创建的通信模式。然后，我们从Megatron-LM、DeepSpeed、DeepEP、TensorRT-LLM、NeMo-RL等系统的代码库中选择了87个问题，以覆盖这些空间，同时还包括大量非LLM工作负载：图神经网络路由、分布式FFT、高斯点扩散等。另一个好处是，由于PKB引用是用标准PyTorch + NCCL编写的，因此基准测试不依赖于任何特定的硬件世代。相反，它被设计为能够自然地与下一代硬件架构共同发展。

对标准Transformer块并行化的分类。不同的分片策略在归一化、注意力和MLP之间创建了不同的通信模式，这里以一个代表性的Gemma3-27B层为例进行说明。

PKB问题在并行类型（左）和源代码库（右）上的覆盖范围，涵盖RL微调、LLM训练、内核库、视觉模型、图神经网络等。

在评估模型之前，我们首先检查了PyTorch + NCCL基线是否留有实际的提升空间。一个通信感知的屋顶线表明确实如此：大多数PKB问题受到NVLink的瓶颈限制，而基线运行远低于硬件上限。因此，下一个问题是简单的：模型是否能够弥补这一差距？

前沿模型在PKB上的表现

表现不佳。在零样本设置下，最佳模型解决了87个问题中的28个，其中只有22个解决方案比PyTorch + NCCL基线更快。采样三次尝试将最佳结果提高到36个正确解决方案和27个比基线更快的解决方案，但fast 1 @3仍然最高仅为31%。

类别

#

GPT-5.5

Claude Opus 4.7

Gemini 3 Pro

GLM-5.1

GLM-5.2

DeepSeek V4 Pro

pass @1→3

fast

1

@1→3

Collective Primitive

8

3→5

3→4

4→5

6→6

2→2

2→3

1→3

1→1

0→0

Tensor Parallel

17

1→2

3→3

0→1

Sequence Parallel

2

Context Parallel

12

7→8

5→6

7→7

7→9

5→7

Pipeline Parallel

Data Parallel

9

Expert Parallel

11

2→6

FSDP / ZeRO

Vocab Parallel

4

Graph Parallel

6

2→4

0→2

Geometric / Spatial

5

Other

Total

87

28→36

22→27

20→31

12→20

24→30

12→19

6→7

4→9

更多模型

→

单次尝试的LLM在PKB上表现不佳。pass@k报告了k次尝试中最佳的正确性；fast 1 @k统计了既正确又优于PyTorch + NCCL基线的解决方案。重复采样提高了这两个指标，但fast 1 @3仍然最高仅为31%（GPT-5.5），这表明存在超出采样噪声的基本限制。

成功集中在熟悉的模式上：集体原语、张量并行的GEMMs以及类似Ulysses的上下文并行，这种并行方式通过all-to-all通信将序列维度拆分到不同的注意力头中。这些是开源代码中最常见的多GPU堆栈部分，因此模型可能对这些部分有更强的先验知识。

失败表明的问题比CUDA语法更深层次。较弱的模型经常无法编译，但更强的推理模型经常生成可以编译但返回错误结果的内核。困难的部分在于推理关于秩协调、数据划分和集体顺序。

生成的内核也使用非常有限的通信机制。大多数依赖于复制引擎或SM的加载/存储指令，而更专业的机制如TMA和NVLS几乎不存在。在许多情况下，模型并未选择实现峰值性能所需的机制。我们将这种现象归因于数据稀缺性和硬件复杂性的结合：像TMA和NVLS这样的新原语需要处理复杂的、特定于硬件的抽象，例如异步复制描述符或新的NVLink拓扑结构，而模型缺乏明确的先验知识。这导致了“可用”的分布式内核与真正优化的内核之间存在巨大差距（这显然是未来工作的明确目标！）

fastp在前沿模型中的得分分布。即使是最优模型（GPT-5.5），随着加速阈值的增加，其表现迅速下降；在p = 1.0时，fast1的最高得分约为31%。

按模型划分的PKB失败模式。较弱的模型大多在编译时失败；更强的推理模型则更常生成能够编译但返回错误结果（输出不匹配）或死锁的内核。

GEMM + All-Gather的分析器跟踪：生成的CUDA内核（87.9 微秒）通过NVLink重叠计算和通信，其性能优于PyTorch + NCCL的参考实现（320.6 微秒），但仍高于理论上限（4.99 微秒）。

自然的下一步是让模型使用与人类内核编写者相同的反馈循环。我们为Gemini 3 Pro封装了一个代理式框架，使其能够访问仓库、终端、编译器输出、正确性测试、速度测量以及其之前的尝试。模型不再只生成一个内核并停止，而是可以编译、运行基准测试、检查失败情况并进行修订。

代理式评估循环：模型编写文件，运行正确性测试和基准测试，并在解决方案通过或步骤预算耗尽之前不断迭代。

这确实有所帮助，但实际收益较为有限。Gemini 3 Pro在单次生成设置中从24个正确解决方案提升到了87个中的35个，其中有26个内核的性能超过了PyTorch + NCCL的基准。这些提升来自于修复语法错误、形状错误和简单的运行时错误。经过大约20次优化步骤后，性能趋于稳定。反馈有助于模型调试分布式内核，但剩余的失败案例突显了一个更大的差距：模型无法推理关于秩协调、通信顺序以及GPU到GPU传输机制的最佳选择。

新的内核

用直接的NVLink加载和存储替换NCCL集合操作的示例：参考实现需要在NCCL集合操作旁边进行额外的重塑；生成的内核将计算操作与直接的对等内存读取融合。

除了在标准集合操作上的加速，单次生成有时还会为没有优化公开参考的负载生成真正的新高性能内核。这种能力预示了AI驱动优化的更广泛潜力；收益不仅限于Transformer模型，模型在状态空间模型、基因组流程和多模态强化学习循环等领域也能表现良好——这些领域中的专用内核相较于主流LLM堆栈仍大多未被优化。

以下是三个示例，每个示例都用融合的对称内存内核替换了 NCCL 集合操作，这些内核可以直接通过 NVLink 移动数据。每个示例都在 4 个 H100 GPU 上经过 100 次随机运行验证了正确性；速度测量遵循 ThunderKittens 2.0 [5] 中的基准测试方法（位相同输入、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 过滤后的分布下计算词汇分片的对数概率。PyTorch + NCCL 基线在过滤之前将完整的词汇表在各 rank 之间进行收集；生成的内核则完全跳过这些集合操作，使用对称内存在内部直接对分片进行排列，同时将 log-softmax、token 提取和目标收集融合为一个 warp-shuffle 减少操作。

参考实现

自定义内核

在序列长度 M ∈ {1024, 2048, 4096, 8192, 16384} 上相对于 PyTorch + NCCL 的加速。配置：4 个 rank，B = 1，V = 32,000（每个 rank 有 8,000 个本地词汇），top- k = 10，top- p = 0.9。

Hyena forward context parallelism (GPT-5.5)

Hyena 操作符的上下文并行前向传播，其中 FFT 卷积需要全局序列上下文。参考实现通过重复的 all_to_all 调用在序列分片和通道分片布局之间交替；生成的内核将输入打包到一个对称分配中，并通过 NVLink 流式传输远程切片，在统一的步骤中计算门控和重新索引。值得注意的是，性能在更长的序列长度上有所下降。

在 4 个 GPU 上，相对于 PyTorch + NCCL 在本地序列长度 l ∈ {1024, 2048, 8192, 16384} 上的加速。正确性和速度在 100 次试验中进行测量。

SAM 3 all-gathered mask IoU suppression (GPT-5.5)

SAM 3 视频分割的跨 GPU 重复抑制：在每一帧之后，各 rank 通过交并比比较预测区域，并将重叠的掩码置零。基线使用可变长度的 all_gather 集合操作加上密集矩阵乘法；生成的解决方案将其简化为一个对称内存内核的流水线，该内核将掩码位打包，并通过硬件 popcount 计算成对重叠。

在固定掩码分辨率 256 × 256，IoU 阈值 = 0.7 的情况下，相对于 PyTorch + NCCL 在跟踪对象数 N ∈ {16, 32, 64, 128, 256, 512} 上的加速。

进一步研究

目前 PKB 的范围仅限于节点内的 NVLink。自然的扩展是跨节点的互连 —— RoCE、InfiniBand —— 在这些互连中，设备端 API 的生态系统尚处于早期阶段，以及其他加速器和拓扑结构，如 TPUs。我们还想了解更高层次的抽象是否有助于或阻碍性能：PKB 已经接受 Triton 和 ParallelKittens 的解决方案，而像 NCCL GIN 和 NVSHMEM 这样的新兴接口也值得作为研究目标。扩展对这些范式的支持将鼓励进一步研究 AI 代理如何在多样化的编程模型和硬件抽象中导航。

更广泛的目标是所有这些背后更困难问题的一个具体目标：能够自主优化和管理大规模分布式基础设施的大型语言模型系统。对于专为语言模型训练和推理构建的基础设施，实现这种自主性最终可能弥合 AI 代理处理其端到端研究工程的能力之间的差距。

我们正在将 PKB 作为一项开放基准发布，以推动这一领域的发展。如果您想深入了解或贡献问题——尤其是涉及节点间交互的问题——我们非常希望听到您的声音。欢迎随时通过电子邮件 willychan2022@gmail.com 或 npaek@together.ai 与我们联系！

参考文献

• Standard Kernel. 在 PTX 层重新构想内核生成：一个从 DSL 学习并超越它们的 LLM 系统。Standard Kernel 博客，2026 年 4 月。

• Robert Tjarko Lange, Qi Sun, Aaditya Prasad, Maxence Faldor, Yujin Tang, 和 David Ha. 面向鲁棒的代理式 CUDA 内核基准测试、验证和优化。arXiv:2509.14279, 2025。

• Carlo Baronio, Pietro Marsella, Ben Pan, Simon Guo, 和 Silas Alberti. Kevin: 用于生成 CUDA 内核的多轮强化学习。arXiv:2507.11948, 2025。

• Raja Gond, Nipun Kwatra, 和 Ramachandran Ramjee. TokenWeave: 用于分布式 LLM 推理的高效计算-通信重叠。arXiv:2505.11329, 2026。

• Stuart Sul 和 Chris Ré. ThunderKittens 2.0: 为您的 GPU 提供更快的内核。Hazy Research, 2026 年 2 月。

@misc{chan2026parallelkernelbench, title = {ParallelKernelBench: LLM 能够编写快速的多 GPU 内核吗？}, author = {Willy Chan 和 Nathan Paek 和 Simon Guo 和 Simran Arora 和 Daniel Y. Fu}, year = {2026}, url = {https://nathanjpaek.github.io/parallel-kernel-bench/} }