ParallelKernelBench：前沿LLM尚无法编写快速多GPU内核
ParallelKernelBench: Frontier LLMs can't write fast multi-GPU kernels (yet)
ParallelKernelBench (PKB) 是一个用于多 GPU kernel 生成的 benchmark 和评估框架，包含来自 Megatron-LM、DeepSpeed、TensorRT-LLM 等真实代码库的 87 个问题，任务是用通过 NVLink 直接传输数据的 CUDA kernel 替换 PyTorch + NCCL 实现。测试 GPT-5.5、Gemini 3 Pro、Claude Opus 4.7 等模型后，zero-shot 下正确解决问题不到三分之一，其中快于 baseline 的不到四分之一。agentic 循环将 Gemini 3 Pro 的正确率从 24 提升至 35，但性能在约 20 步后趋于平稳。模型在 NeMo-RL GRPO 训练、Hyena 上下文并行和 SAM 3 视频分割等任务中生成了比公开实现更快的 kernel。
LLM 在编写 GPU kernel 方面的表现已经出奇地好**[1][2][3]，但几乎所有衡量这一进展的现有 benchmark 都是单 GPU 的。在生产环境中，通信往往是瓶颈：通信开销可能占推理延迟的 20% 以上[4]**，并且随着计算能力增长快于互连带宽，这一差距还在持续扩大。
(PKB) 提供了一个用于多 GPU kernel 生成的 benchmark 和评估框架，包含来自真实代码库的 87 个问题，其任务是用一个通过 NVLink 直接传输数据的 CUDA kernel 替换 PyTorch + NCCL。我们测试了前沿的编码模型，如 GPT-5.5、Gemini 3 Pro、Opus 4.7 等。评估结果显示，各模型普遍存在显著的性能差距：正确解决的问题不到三分之一，而其中能超越朴素 baseline 的更是不到四分之一。
我们将探讨它们失败的原因、失败的模式，以及一些模型出人意料地生成了比任何公开实现都更快的 kernel 的案例，其中包括一个用于 NVIDIA NeMo-RL 的 GRPO 训练循环的 kernel，该循环此前没有经过优化的公开参考实现。
为什么多 GPU 与单 GPU kernel 生成不同
LLM 在 GPU kernel 生成方面取得了进展，但这种进展主要是在单 GPU 上衡量的。生产级 AI 工作负载已不再适用这一框架：它们横跨多个 GPU，性能越来越多地由通信决定，而不仅仅是本地计算和内存。这种转变使得多 GPU kernel 生成在三个方面成为一个不同的问题：
- 设计空间呈组合式爆炸增长。 从业者组合使用张量并行、专家并行、数据并行、上下文并行和序列并行来适配硬件，每种组合都会产生不同的通信模式。
- 性能模型发生了变化。 单 GPU 的 roofline 模型围绕计算和内存带宽构建。在多 GPU 代码中，瓶颈往往是互连。
- 多 GPU kernel 生成引入了一个关键的新设计选择： 如何在 GPU 之间移动数据——通过 copy engine、TMA、SM load/store 或 NVLS——以及是否将该移动与计算融合。
我们构建 PKB 是为了测试模型能否超越纯粹的 torch.dist，真正编写生产级的多 GPU kernel。每个问题都从一个标准的 PyTorch + NCCL 实现和硬件拓扑描述开始。然后，模型必须用使用对称内存（symmetric memory）在 GPU 之间直接通信的 CUDA kernel 替换该参考实现。
PKB 评估流程。每个问题提供一个任务、硬件拓扑和 PyTorch + NCCL 参考实现；模型生成一个自定义 CUDA kernel，并对其正确性、挂钟加速比和通信 roofline 进行评估。
为了确保这 87 个问题覆盖生产级并行类型的真实空间，我们根据分布式工作负载的分类法构建了它们。首先，我们确定了模型被分片的主要方式——张量并行、上下文并行、数据并行、专家并行、序列并行和 FSDP/ZeRO——以及每种方式产生的通信模式。然后，我们选择了 87 个问题来覆盖该空间，这些问题取自 Megatron-LM、DeepSpeed、DeepEP、TensorRT-LLM、NeMo-RL 等系统的代码库，以及大量非 LLM 工作负载：GNN 路由、分布式 FFT、高斯泼溅等。另一个好处是，由于 PKB 的参考实现是用标准 PyTorch + NCCL 编写的，因此该 benchmark 不依赖于任何特定的单一代硬件。相反，它被设计为能够自然地随着下一代硬件架构的发展而演进。
标准 Transformer 模块的并行化分类法。不同的分片策略在归一化、attention 和 MLP 上产生不同的通信模式，此处以代表性的 Gemma3-27B 层为例进行说明。
PKB 问题在并行类型（左）和源代码库（右）上的覆盖范围，涵盖 RL 后训练、LLM 训练、kernel 库、视觉模型、GNN 等。
在评估模型之前，我们首先检查了 PyTorch + NCCL baseline 是否留下了真正的优化空间。一个考虑通信的 roofline 模型给出了肯定的答案：大多数 PKB 问题的瓶颈在于 NVLink，而 baseline 的运行速度远低于硬件上限。那么下一个问题很简单：模型能否缩小这一差距？
前沿模型在 PKB 上的表现
并不好。 在 zero-shot 设置下，最好的模型解决了 87 个问题中的 28 个，其中只有 22 个解决方案比 PyTorch + NCCL baseline 更快。采样三次尝试后，最佳结果提升到 36 个正确解决方案和 27 个快于 baseline 的解决方案，但 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 | pass @1→3 | fast 1 @1→3 | pass @1→3 | fast 1 @1→3 | pass @1→3 | fast 1 @1→3 |
| 集合通信原语 | 8 | 3→5 | 3→4 | 4→5 | 3→4 | 6→6 | 2→2 |
| 张量并行 | 17 | 2→2 | 1→2 | 1→3 | 1→3 | 3→3 | 1→1 |
| 序列并行 | 2 | 1→1 | 1→1 | 0→0 | 0→0 | 0→0 | 0→0 |
| 上下文并行 | 12 | 7→8 | 5→6 | 7→7 | 3→4 | 7→9 | 5→7 |
| 流水线并行 | 1 | 0→0 | 0→0 | 0→0 | 0→0 | 0→0 | 0→0 |
| 数据并行 | 9 | 1→2 | 1→2 | 1→1 | 1→1 | 0→1 | 0→0 |
| 专家并行 | 11 | 3→3 | 2→3 | 2→6 | 0→1 | 3→4 | 0→1 |
| FSDP / ZeRO | 9 | 3→4 | 3→4 | 3→3 | 3→3 | 1→1 | 0→1 |
| 词表并行 | 4 | 2→2 | 1→1 | 1→2 | 1→2 | 2→2 | 2→2 |
| 图并行 | 6 | 2→4 | 2→3 | 0→2 | 0→2 | 1→1 | 1→1 |
| 几何/空间 | 5 | 3→3 | 3→3 | 1→2 | 0→0 | 0→1 | 0→1 |
| 其他 | 2 | 1→2 | 0→1 | 0→0 | 0→0 | 1→2 | 1→2 |
| 总计 | 87 | 28→36 | 22→27 | 20→31 | 12→20 | 24→30 | 12→19 |
更多模型 →
单次 LLM 在 PKB 上表现挣扎。 pass@k 报告了 k 次尝试中最佳结果的正确性；fast 1@k 统计了既正确又优于 PyTorch + NCCL baseline 的解决方案数量。重复采样改善了两个指标，但 fast 1@3 仍然只有 31%（GPT-5.5），这表明存在超出采样噪声的根本性限制。
成功案例集中在熟悉的模式上：集合通信原语、张量并行 GEMM 和 Ulysses 风格的上下文并行（通过 all-to-all 通信在 attention head 之间分割序列维度）。这些是多 GPU 栈中在开源代码里最可见的部分，因此模型可能对它们有更强的先验知识。
失败案例表明存在比 CUDA 语法更深层的问题。较弱的模型通常无法编译，但更强的推理模型经常生成能编译但返回错误结果的 kernel。难点在于推理 rank 协调、数据分区和集合通信顺序。
生成的 kernel 也使用了非常狭窄的通信机制集。大多数依赖 copy engine 或 SM load/store 指令，而更专门的机制如 TMA 和 NVLS 几乎不存在。在许多情况下，模型没有选择达到峰值性能所需的机制。我们将其归因于数据稀缺和硬件复杂性的结合：较新的原语如 TMA 和 NVLS 需要驾驭复杂、特定于硬件的抽象，例如异步拷贝描述符或更新的 NVLink 拓扑，而模型在这些方面缺乏强先验知识。这在“能工作”的分布式 kernel 和真正优化过的 kernel 之间留下了巨大差距（这是未来工作的明确目标！）
前沿模型的 fastp 分数分布。即使是最好的模型（GPT-5.5），随着加速比阈值的增加，分数也会迅速下降；在 p = 1.0 时，fast1 最高约为 31%。
按模型划分的 PKB 失败模式。较弱的模型主要在编译时失败；更强的推理模型更常产生能编译但返回错误结果（输出不匹配）或死锁的 kernel。
GEMM + All-Gather 的 profiler 跟踪：生成的 CUDA kernel（87.9 μs）在 NVLink 上重叠计算和通信，优于 PyTorch + NCCL 参考实现（320.6 μs），但仍高于理论 roofline（4.99 μs）。
自然的下一步是给模型提供人类 kernel 编写者会使用的相同反馈循环。我们将 Gemini 3 Pro 封装在一个 agentic harness 中，使其能够访问仓库、终端、编译器输出、正确性测试、速度测量以及之前的尝试。模型不是生成一个 kernel 就停止，而是可以编译、运行 benchmark、检查失败并修改。
Agentic 评估循环：模型编写文件，运行正确性测试和 benchmark，并迭代直到解决方案通过或步骤预算耗尽。
这有所帮助，但实际收益有限。Gemini 3 Pro 从单次设置下的 24 个正确解决方案提高到 87 个中的 35 个，其中 26 个 kernel 优于 PyTorch + NCCL baseline。收益来自于修复语法错误、形状错误和简单的运行时错误。在大约 20 次优化步骤后，性能趋于平稳。反馈有助于模型调试分布式 kernel，但剩余的失败突显了一个更大的差距：无法推理 rank 协调、通信顺序以及 GPU 间传输机制的最优选择。
全新的 kernel
用直接 NVLink load/store 替换 NCCL 集合通信的示例：参考实现需要额外的 reshape 和 NCCL 集合通信；生成的 kernel 将计算操作与直接的 peer memory 读取融合在一起。
除了在标准集合通信上实现加速外，单次生成偶尔会为没有优化过的公开参考实现的工作负载，产生真正新颖的高性能 kernel。这种能力预示了 AI 驱动优化的更广泛潜力；胜利不仅限于 Transformer，模型在状态空间模型、基因组学流程和多模态 RL 循环等领域也能表现出色——与主流的 LLM 栈相比，这些领域的专用 kernel 在很大程度上仍未得到优化。
以下是三个示例，每个都用融合的对称内存 kernel 替换了 NCCL 集合通信，该 kernel 通过 NVLink 直接传输数据。每个示例都在 4 个 H100 GPU 上经过 100 次随机运行验证了正确性；速度测量遵循 ThunderKittens 2.0[5] 中的 benchmark 方法论（按位相同输入、L2 感知输入组、500 次预热迭代和 100 次背靠背性能分析迭代）。
NeMo 词表并行 log-prob 与 top-k/top-p 过滤（Gemini 3 Pro）
NVIDIA NeMo-RL 的 GRPO 训练循环中的一个核心步骤：在 top-k/top-p 过滤分布下计算词表分片的 log-probability。PyTorch + NCCL baseline 在过滤前跨 rank 收集完整词表；生成的 kernel 完全跳过了这些集合通信，使用对称内存就地排列分片，同时将 log-softmax、token 提取和目标 gather 融合到单个 warp-shuffle reduction 中。
与 PyTorch + NCCL 相比，在序列长度 M ∈ {1024, 2048, 4096, 8192, 16384} 上的加速比。配置：4 个 rank，B = 1，V = 32,000（每个 rank 本地词表 8,000），top-k = 10，top-p = 0.9。
Hyena 前向上下文并行（GPT-5.5）
Hyena 算子的上下文并行前向传播，其中 FFT 卷积需要序列全局上下文。参考实现通过重复的 all_to_all 调用在序列分片和通道分片布局之间交替；生成的 kernel 将输入打包到一个对称内存分配中，并通过 NVLink 流式传输远程切片，在统一的过程中计算门控和重索引。值得注意的是，在较长的序列长度下性能会下降。
与 PyTorch + NCCL 相比，在 4 个 GPU 上，本地序列长度 l ∈ {1024, 2048, 8192, 16384} 的加速比。正确性和速度在 100 次试验中测量。
SAM 3 全收集 mask IoU 抑制（GPT-5.5）
SAM 3 视频分割的跨 GPU 重复抑制：在每帧之后，rank 通过交并比比较预测区域，并将重叠的 mask 置零。baseline 使用可变长度的 all_gather 集合通信加上一个密集 matmul；生成的解决方案将其压缩为一个对称内存 kernel 的流水线，该流水线对 mask 进行位打包，并使用硬件 popcount 计算成对重叠。
与 PyTorch + NCCL 相比，在固定 mask 分辨率 256 × 256、IoU 阈值 = 0.7 下，跟踪对象 N ∈ {16, 32, 64, 128, 256, 512} 的加速比。
进一步研究
PKB 目前有意限定在节点内 NVLink。自然的扩展是节点间网络——RoCE、InfiniBand——在这些网络上，设备端 API 生态较新，以及其他加速器和拓扑如 TPU。我们也想知道更高级的抽象是有帮助还是有害：PKB 已经接受 Triton 和 ParallelKittens 解决方案，而新兴接口如 NCCL GIN 和 NVSHMEM 也值得作为目标进行研究。将支持扩展到这些范式将鼓励进一步研究 AI agent 如何驾驭多样化的编程模型和硬件抽象。
更广泛的目标是为这一切背后更困难的问题提供一个具体目标：能够自主优化和管理大规模分布式基础设施的 LLM 系统。对于为语言模型训练和推理定制的基础设施，实现这种自主性最终可以弥合与能够处理自身端到端研究工程的 AI agent 之间的差距。
我们正在将 PKB 作为一个开放的 benchmark 发布，以推动这一目标。如果您想深入研究或贡献问题——尤其是节点间问题——我们很乐意听取您的意见。请随时发送电子邮件至 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/} }