跳到正文
Together AI Blog·· 2026-06-23精选AI 评分61

Together AI 发布 ParallelKernelBench：前沿 LLM 还写不出快的多 GPU 内核

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

AI 导读

Together AI 发布 ParallelKernelBench（PKB），一个面向多 GPU 内核生成的基准与评测框架，包含 87 道取自 Megatron-LM、DeepSpeed、TensorRT-LLM、NeMo-RL 等真实代码库的题目，要求模型用 CUDA 内核替换 PyTorch + NCCL 并直接通过 NVLink 通信。

推荐理由

Together AI 用 87 道真实多 GPU 内核题测前沿模型，给出通过率与失败模式，可据此判断分布式内核生成的当前边界。

正文 · AI 翻译

概述

LLM 在编写 GPU 内核方面已经变得出奇地出色[1][2][3]，但几乎所有衡量这一进展的现有基准都是单 GPU 的。在生产环境中，通信往往是瓶颈：通信开销可占推理延迟的 20% 以上[4]，而且随着计算扩展速度快于互连带宽，这一差距还在不断扩大。

ParallelKernelBench（PKB）提供了一个用于多 GPU 内核生成的基准和评估框架，包含来自真实代码库的 87 个问题，其任务是用直接在 NVLink 上传输数据的 CUDA 内核替换 PyTorch + NCCL。我们测试了 GPT-5.5、Gemini 3 Pro、Opus 4.7 等前沿编码模型。评估揭示了普遍存在的显著性能差距：不到三分之一的问题被正确解决，而其中又只有不到四分之一优于朴素基线。

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

为什么多 GPU 内核生成不同于单 GPU 内核生成

LLM 在 GPU 内核生成方面已取得进展，但这种进展大多是在单 GPU 上衡量的。生产级 AI 工作负载已不再符合这一框架：它们跨多个 GPU，性能越来越取决于通信，而不仅仅是本地计算和内存。这一转变使多 GPU 内核生成在三个方面成为一个不同的问题：

  • 设计空间呈组合式扩展。从业者组合张量并行、专家并行、数据并行、上下文并行和序列并行以适配硬件，而每种组合都会产生不同的通信模式。
  • 性能模型发生变化。单 GPU 的 roofline 围绕计算和内存带宽构建。在多 GPU 代码中，瓶颈往往是互连。
  • 多 GPU 内核生成引入了一个关键的新设计选择：如何在 GPU 之间移动数据——通过复制引擎、TMA、SM load/store 还是 NVLS——以及是否将该数据移动与计算融合。

ParallelKernelBench

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

PKB 评估流水线。每个问题提供任务、硬件拓扑和 PyTorch + NCCL 参考实现；模型生成自定义 CUDA 内核，并对其正确性、实际运行加速比和通信 roofline 进行评估。

为了确保这 87 个问题覆盖生产环境中并行类型的真实空间，我们基于分布式工作负载的分类体系构建了它们。首先，我们确定了模型分片的主要方式——张量、上下文、数据、专家、序列以及 FSDP/ZeRO——以及每种方式所产生的通信模式。然后，我们从 Megatron-LM、DeepSpeed、DeepEP、TensorRT-LLM、NeMo-RL 等系统的代码库中选取了 87 个问题来覆盖这一空间，此外还包括一长串非 LLM 工作负载：GNN 路由、分布式 FFT、高斯泼溅等。另一个好处是，由于 PKB 参考实现是用标准 PyTorch + NCCL 编写的，该基准测试不绑定于任何单一特定的硬件代际。相反，它被设计为能够随下一代硬件架构自然演进。

并行化标准 transformer 块的分类体系。不同的分片策略在归一化、注意力和 MLP 中产生不同的通信模式，此处以一个具有代表性的 Gemma3-27B 层为例进行说明。

PKB 问题在并行类型（左）和源代码库（右）上的覆盖情况，涵盖 RL 后训练、LLM 训练、内核库、视觉模型、GNN 等。

在评估模型之前，我们首先检查 PyTorch + NCCL 基线是否留有真正的提升空间。一个通信感知的 roofline 模型表明确实如此：大多数 PKB 问题受限于 NVLink，而基线运行远低于硬件上限。所以下一个问题很简单：模型能弥合这一差距吗？

前沿模型在 PKB 上的表现如何

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

类别 # 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
集合原语 8 3→53→4 4→53→4 6→62→2 2→32→3 1→31→1 0→00→0
张量并行 17 2→21→2 1→31→3 3→31→1 1→10→0 0→10→0 0→00→0
序列并行 2 1→11→1 0→00→0 0→00→0 0→00→0 0→00→0 0→00→0
上下文并行 12 7→85→6 7→73→4 7→95→7 2→20→0 2→20→0 2→30→1
流水线并行 1 0→00→0 0→00→0 0→00→0 0→00→0 0→00→0 0→00→0
数据并行 9 1→21→2 1→11→1 0→10→0 0→00→0 0→00→0 0→00→0
专家并行 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
词表并行 4 2→21→1 1→21→2 2→22→2 0→00→0 0→10→1 0→00→0
图并行 6 2→42→3 0→20→2 1→11→1 0→00→0 0→00→0 0→00→0
几何 / 空间 5 3→33→3 1→20→0 0→10→1 0→00→0 0→00→0 0→00→0
其他 2 1→20→1 0→00→0 1→21→2 0→00→0 0→10→0 0→00→0
总计 87 28→3622→27 20→3112→20 24→3012→19 6→72→3 4→91→2 2→30→1

更多模型 →

单次尝试的 LLM 在 PKB 上表现挣扎。pass@k 报告 k 次尝试中最佳结果的正确性；fast1@k 统计既正确又优于 PyTorch + NCCL 基线的解决方案。重复采样可改善这两项指标，但 fast1@3 仍然最高仅达 31%（GPT-5.5），这表明存在超越采样噪声的根本性局限。

成功案例集中在熟悉的模式上：集合原语、张量并行 GEMM，以及 Ulysses 风格的上下文并行——后者利用全对全通信将序列维度拆分到注意力头上。这些是多 GPU 技术栈中在开源代码里最常出现的部分，因此模型可能对它们有更强的先验知识。

失败案例表明问题比 CUDA 语法更深层。较弱的模型往往无法编译，但更强的推理模型经常生成能够编译却返回错误结果的内核。真正的难点在于对 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 就停下，而是可以编译、运行基准测试、检查失败并进行修改。

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

这有所帮助，但实际收益有限。Gemini 3 Pro 从单次生成设置下的 24 个正确解提升到 87 个中的 35 个，其中 26 个 kernel 优于 PyTorch + NCCL 基线。收益来自修复语法错误、形状错误和简单的运行时 bug。经过大约 20 轮改进后，性能趋于停滞。反馈能帮助模型调试分布式 kernel，但剩余的失败凸显了一个更大的差距：无法对 rank 协调、通信顺序以及 GPU 间传输机制的最优选择进行推理。

净新增 kernel

用直接 NVLink load 和 store 替换 NCCL 集合通信的示例：参考实现需要在 NCCL 集合通信之外进行额外的 reshape；生成的 kernel 将计算操作与直接的对等内存读取融合在一起。

除了在标准集合通信上的加速之外，单次生成偶尔还能为没有优化公开参考实现的工作负载带来真正全新的高性能 kernel。这一能力让我们得以一窥 AI 驱动优化的更广泛潜力；收益并不局限于 Transformer，因为模型在状态空间模型、基因组流水线和多模态 RL 循环上也能表现出色——在这些领域，专用 kernel 相比主流 LLM 技术栈仍基本未得到优化。

下面是三个示例，每个示例都用融合的对称内存内核替换 NCCL 集合通信，直接在 NVLink 上移动数据。每个示例都在 4 块 H100 GPU 上经过 100 次随机运行验证了正确性；速度测量遵循 ThunderKittens 2.0[5] 中的基准测试方法（按位相同的输入、L2 感知的输入分组、500 次预热迭代，以及 100 次连续性能分析迭代）。

NeMo 词表并行 log-prob，带 top-k/top-p 过滤（Gemini 3 Pro）

NVIDIA NeMo-RL 的 GRPO 训练循环中的一个核心步骤：在 top-k/top-p 过滤分布下计算按词表分片的 log-probabilities。PyTorch + NCCL 基线在过滤前跨 rank 收集完整词表；生成的内核完全跳过这些集合通信，使用对称内存就地置换分片，同时将 log-softmax、token 提取和目标 gather 融合为单次 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 前向上下文并行（GPT-5.5）

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

参考实现

自定义内核

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

SAM 3 全收集掩码 IoU 抑制（GPT-5.5）

SAM 3 视频分割的跨 GPU 重复抑制：每帧之后，各 rank 通过交并比比较预测区域，并将重叠的掩码置零。基线使用可变长度的 all_gather 集合通信加上稠密 matmul；生成的解决方案将其折叠为一系列对称内存内核，对掩码进行位打包，并用硬件 popcount 计算成对重叠。

参考实现

自定义内核

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

进一步研究

PKB 目前有意限定于节点内 NVLink。自然的扩展是节点间互连——RoCE、InfiniBand——在这些领域设备端 API 生态更年轻，以及其他加速器和拓扑，如 TPU。我们还想知道更高层抽象是有帮助还是有损害：PKB 已经接受 Triton 和 ParallelKittens 解决方案，而 NCCL GIN 和 NVSHMEM 等新兴接口作为目标值得研究。将支持扩展到这些范式将鼓励进一步研究 AI 智能体如何驾驭多样化的编程模型和硬件抽象。

更宏大的目标是为这一切背后更困难的问题设定一个具体目标：能够自主优化和管理大规模分布式基础设施的 LLM 系统。对于为语言模型训练和推理而定制构建的基础设施而言，实现这种自主性最终可能弥合差距，通向能够处理自身端到端研究工程的 AI 智能体。

我们发布 PKB 作为一个开放基准来推动这一方向。如果你想深入了解或贡献问题——尤其是节点间的问题——我们很乐意听到你的声音。欢迎发送邮件至 [email protected] 或 [email protected]！

参考文献

  1. Standard Kernel. 在 PTX 层重新构想内核生成：一个从 DSL 中学习并超越它们的 LLM 系统. Standard Kernel 博客，2026 年 4 月。
  2. Robert Tjarko Lange, Qi Sun, Aaditya Prasad, Maxence Faldor, Yujin Tang, and David Ha. 迈向稳健的智能体式 CUDA 内核基准测试、验证与优化. arXiv:2509.14279, 2025.
  3. Carlo Baronio, Pietro Marsella, Ben Pan, Simon Guo, and Silas Alberti. Kevin：用于生成 CUDA 内核的多轮强化学习. arXiv:2507.11948, 2025.
  4. Raja Gond, Nipun Kwatra, and Ramachandran Ramjee. TokenWeave：面向分布式 LLM 推理的高效计算-通信重叠. arXiv:2505.11329, 2026.
  5. Stuart Sul and Chris Ré. ThunderKittens 2.0：为你的 GPU 打造更快的内核. Hazy Research，2026 年 2 月。

@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/} }

来源：Together AI Blog · together.ai

开源生态部署工程论文研究评测基准#论文/研究#评测/基准#部署/工程#开源/仓库