摘要
LLM 在编写 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 内核生成不同
LLM 在 GPU 内核生成方面取得了进展，但这一进展主要是在单 GPU 上衡量的。生产环境中的 AI 工作负载已不再适用这一框架：它们跨越多个 GPU，性能越来越受通信影响，而不仅仅是本地计算和内存。这种转变使得多 GPU 内核生成在三个方面成为一个不同的问题：
设计空间呈组合式扩展。 实践者组合使用张量并行、专家并行、数据并行、上下文并行和序列并行来适配硬件，每种组合都会产生不同的通信模式。性能模型发生了变化。 单 GPU 的屋顶模型围绕计算和内存带宽构建，而在多 GPU 代码中，瓶颈往往是互连。多 GPU 内核生成引入了一个关键的新设计选择： 如何在 GPU 之间移动数据——通过复制引擎、TMA、SM 加载/存储或 NVLS——以及是否将该移动与计算融合。
ParallelKernelBench
我们构建了 PKB 来测试模型是否能够超越纯 torch.dist，并实际编写生产级的多 GPU 内核。每个问题都从一个标准的 PyTorch + NCCL 实现和硬件拓扑描述开始。然后，模型必须将该参考替换为一个使用对称内存在 GPU 之间直接通信的 CUDA 内核。
为了确保这87个问题能够覆盖生产级并行类型的真实空间，我们基于分布式工作负载的分类法构建了它们。首先，我们确定了模型分片的主要方式——张量并行、上下文并行、数据并行、专家并行、序列并行以及FSDP/ZeRO——以及每种方式产生的通信模式。然后，我们从Megatron-LM、DeepSpeed、DeepEP、TensorRT-LLM、NeMo-RL等系统的代码库以及大量非LLM工作负载（如GNN路由、分布式FFT、高斯泼溅等）中选取了87个问题来覆盖该空间。另一个好处是，由于PKB的参考实现是用标准PyTorch + NCCL编写的，该基准测试并不依赖于任何特定的硬件代际。相反，它被设计为能够自然地随下一代硬件架构演进。
在评估模型之前，我们首先检查了PyTorch + NCCL基线是否留有真正的优化空间。一个考虑通信的屋顶线模型给出了肯定答案：大多数PKB问题受限于NVLink，且基线运行远低于硬件上限。因此，下一个问题很简单：模型能否缩小这一差距？
前沿模型在PKB上的表现
并不好。 在零样本设置下，最佳模型解决了87个问题中的28个，其中只有22个解决方案比PyTorch + NCCL基线更快。采样三次尝试后，最佳结果提升至36个正确解决方案和27个优于基线的解决方案，但fast1@3仍仅为31%。
成功案例集中在常见模式中：集合通信原语、张量并行GEMM以及Ulysses风格的上下文并行（通过全对全通信在注意力头之间分割序列维度）。这些是多GPU栈中最常在开源代码中出现的部分，因此模型可能对它们有更强的先验知识。
失败案例表明存在比CUDA语法更深层的问题。较弱的模型通常无法编译，但更强的推理模型经常生成能编译但返回错误结果的内核。难点在于协调秩、数据分区和集合通信顺序。
生成的内核还使用了非常狭窄的通信机制。大多数依赖于拷贝引擎或SM加载/存储指令，而更专门的机制如TMA和NVLS几乎不存在。在许多情况下，模型没有选择达到峰值性能所需的机制。我们将此归因于数据稀缺和硬件复杂性的结合：较新的原语如TMA和NVLS需要处理复杂且硬件特定的抽象，例如异步拷贝描述符或更新的NVLink拓扑，而模型缺乏强先验知识。这导致“可用”的分布式内核与真正优化的内核之间存在巨大差距（这是未来工作的明确目标！）
自然的下一步是给模型提供人类内核编写者会使用的相同反馈循环。我们将Gemini 3 Pro封装在一个智能体框架中，该框架可以访问代码仓库、终端、编译器输出、正确性测试、速度测量以及之前的尝试。模型不再只生成一个内核就停止，而是可以编译、运行基准测试、检查失败并修改。
这有所帮助，但实际收益有限。Gemini 3 Pro 从单次生成设置中的 24 个正确解决方案提升至 87 个中的 35 个，其中 26 个内核优于 PyTorch + NCCL 基线。改进来自修复语法错误、形状错误和简单的运行时错误。经过大约 20 轮优化后，性能趋于平稳。反馈有助于模型调试分布式内核，但剩余的失败凸显了一个更大的差距：无法推理秩协调、通信顺序以及 GPU 间传输机制的最优选择。
全新内核
除了标准集合通信的加速外，单次生成偶尔会为没有优化公共参考的工作负载生成真正新颖的高性能内核。这一能力展示了 AI 驱动优化的更广泛潜力；胜利不仅限于 Transformer，模型在状态空间模型、基因组流程和多模态 RL 循环中也能表现出色——这些领域中的专用内核与主流 LLM 堆栈相比仍大多未优化。
以下是三个示例，每个示例都用融合的对称内存内核替换了 NCCL 集合通信，这些内核直接在 NVLink 上传输数据。每个示例在 4 个 H100 GPU 上经过 100 次随机运行验证了正确性；速度测量遵循 ThunderKittens 2.0[5] 中的基准测试方法（按位相同输入、L2 感知输入组、500 次预热迭代和 100 次连续性能分析迭代）。
NeMo 词汇并行对数概率与 top-k/top-p 过滤（Gemini 3 Pro）
NVIDIA NeMo-RL 的 GRPO 训练循环中的核心步骤：在 top-k/top-p 过滤分布下计算词汇分片对数概率。PyTorch + NCCL 基线在过滤前跨秩收集完整词汇；生成的内核完全跳过这些集合通信，使用对称内存内联排列分片，同时将对数 softmax、令牌提取和目标收集融合为单个 warp-shuffle 归约。
M∈ {1024, 2048, 4096, 8192, 16384}。配置：4 个秩，B= 1，V= 32,000（每个秩 8,000 本地词汇），top-k= 10，top-p= 0.9。
Hyena 前向上下文并行（GPT-5.5）
Hyena 算子的上下文并行前向传播，其中 FFT 卷积需要序列全局上下文。参考实现通过重复的 all_to_all 调用在序列分片和通道分片布局之间交替；生成的内核将输入打包到一个对称分配中，并通过 NVLink 流式传输远程切片，在统一过程中计算门控和重新索引。值得注意的是，在较长序列长度下性能会下降。
l∈ {1024, 2048, 8192, 16384}，在 4 个 GPU 上。正确性和速度在 100 次试验中测量。
SAM 3 全收集掩码 IoU 抑制（GPT-5.5）
SAM 3 视频分割的跨 GPU 重复抑制：每帧后，秩通过交并比比较预测区域并清零重叠掩码。基线使用可变长度 all_gather 集合通信加上密集矩阵乘法；生成的解决方案将其压缩为对称内存内核的流水线，这些内核使用硬件 popcount 对掩码进行位打包并计算成对重叠。
进一步研究
PKB 目前有意限定在节点内 NVLink 范围内。自然的扩展方向是节点间网络——RoCE、InfiniBand——这些领域的设备端 API 生态尚不成熟，以及 TPU 等其他加速器和拓扑结构。我们也想了解更高层次的抽象是有益还是有害：PKB 已经接受 Triton 和 ParallelKittens 的解决方案，而 NCCL GIN 和 NVSHMEM 等新兴接口也值得作为目标进行研究。将支持扩展到这些范式将鼓励进一步研究 AI 智能体如何驾驭多样化的编程模型和硬件抽象。
更广泛的目标是为这一切背后更困难的问题提供一个具体目标：能够自主优化和管理大规模分布式基础设施的 LLM 系统。对于为语言模型训练和推理定制的基础设施，实现这种自主性最终可能弥合与能够处理自身端到端研究工程的 AI 智能体之间的差距。
我们发布 PKB 作为一个开放基准来推动这一目标。如果您想深入研究或贡献问题——尤其是节点间问题——我们期待您的反馈。请随时发送电子邮件至 [email protected] 或 [email protected]！
参考文献
- 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.