ParallelKernelBench: Can LLMs Write Fast Multi-GPU Kernels?

发表时间: 2026-06 · alphaXiv:2606.parallel-kernel-bench (Stanford / Caltech / UCSD / Together AI)

作者及机构:Willy Chan (Stanford University, Together AI), Nathan Paek (Stanford University, Together AI), Simon Guo (Stanford University), Simran Arora (California Institute of Technology, Together AI), Daniel Y. Fu (University of California, San Diego, Together AI)


A1 主要贡献(总结)


A3 背景知识/关键Observation/设计原则(缩写)

GPU架构与内存分级。现代数据中心级NVIDIA GPU在成千上万个并行硬件线程上执行被称为内核(Kernel)的程序,这些线程组织在三个层次中:线程是基本执行单元;每32个线程组成一个线程束(Warp),在单个流式多处理器(SM)上执行;线程束进一步组织为线程块(Block),共享片上快速的共享内存(Shared Memory)。GPU还使用高带宽内存(HBM)作为主要的全局工作内存。在多GPU系统中,GPU之间通过NVLink连接,允许一个GPU上执行的线程直接读写对端GPU的HBM。

互联硬件层级与网络计算。现代GPU集群暴露了多级互联结构。PCIe(64 GB/s)负责CPU到GPU的传输,而NVLink(单向每GPU 450 GB/s)提供节点内高带宽的GPU到GPU点对点连接。NVSwitch将节点内的所有GPU连接成全互连结构,并通过NVLink SHARP(NVLS)支持网络内计算(In-Network Compute),通过多内存(multimem)指令直接加速多播和归约操作,无需将数据路由通过GPU的SM。

三种数据传输机制的权衡。根据ParallelKittens【[48],Sul等人,Parallelkittens: Systematic and practical simplification of multi-gpu ai kernels, 2025】的分类,GPU间的数据传输主要通过三种机制进行。复制引擎(Copy Engine, CE)由主机端发起,对于大消息能饱和NVLink带宽,但无法在内核内部调用。张量内存协处理器(Tensor Memory Accelerator, TMA)能够异步移动数据而不消耗SM资源,适合将通信与计算重叠。最后,SM加载/存储(Load/Store, LSA)指令直接利用SM进行传输;尽管这增加了寄存器压力,但它支持通过NVSwitch进行NVLS网络内计算,而其他两种机制不支持。

通信软件栈与对称内存。NCCL是GPU集体通信的标准库,被 torch.distributed 使用,它将all-reduce和reduce-scatter等操作暴露为在CUDA流上排队的异步内核【[22],Hu等人,Demystifying nccl: An in-depth analysis of gpu communication protocols and algorithms, 2026】。虽然NCCL使用方便,但它通常在粗粒度的集体通信层级运行,难以灵活重叠计算与通信。为了进行更底层的控制,程序员可以使用对称堆(Symmetric Heap)——一个在所有参与GPU上注册的内存区域——在内核内部执行单边put和get操作。PyTorch通过其对称内存(Symmetric Memory)抽象暴露了这一功能,可在CUDA中通过统一虚拟寻址(UVA)访问,将远程GPU内存映射到单个虚拟地址空间中,使内核能够使用标准load/store指令直接读写对端GPU。

现有基准测试的局限性。如Table 1所示,现有的内核生成基准测试(如KernelBench [41], FlashInfer-Bench [54], SOL-ExecBench [32])均只关注单GPU场景。在线代码源中完全不包含分布式问题,且未覆盖完整的AI并行工作负载。这凸显了多GPU内核生成基准测试的空白。


A2 方法细节(缩写)

PKB任务格式定义。每个PKB任务向大语言模型提供一个用PyTorch和 torch.distributed NCCL操作编写的未优化参考实现,以及一个指定GPU数量和节点内硬件配置的系统拓扑。模型的任务是将此参考实现重写为使用统一虚拟寻址(UVA)的高性能设备端CUDA内核。该过程使用PyTorch的对称内存API(Symmetric Memory API)来实现低延迟的GPU点对点(P2P)通信。任务的输入输出和评估流程如图1所示。

Figure 1 ParallelKernelBench评估语言模型在多GPU内核优化上的表现
Figure 1 ParallelKernelBench评估语言模型在多GPU内核优化上的表现

输入准备与环境预设。每个任务在通用的 create_input_tensor() 函数中初始化每个GPU rank的输入张量和大小,并配置好所有的通信组。在调用模型的 solution() 函数之前,分片方案(Sharding Scheme)已经完全确定,该函数在每个rank上以相同的方式运行。系统向模型提供:(i) 使用未优化的 torch.distributed 和NCCL集体通信的参考实现;(ii) 系统提示词和上下文示例(Listing 1);(iii) 描述硬件能力(如NVLink P2P访问和网络内归约支持)的自然语言系统拓扑说明。

模型输出与内联编译。模型需要输出一个重写的 solution 函数,用细粒度的、重叠计算与通信的内核替换未优化的集体通信。评估系统使用 PyTorch 的 load_inline 工具动态编译CUDA代码。对称堆指针在多次运行之间会被缓存,从而在测量延迟时排除设置(Setup)的开销。

问题选择原则与分类覆盖。PKB的问题来源于两个方面:一是从GitHub开源仓库中抓取的高频真实部署操作,二是从优化后的库实现中提取的内核,以确保工作负载的实用性。如图2所示,PKB将标准的Transformer层(Norm -> Attention -> Norm -> MLP)分解为可并行的操作,并标注对应的NCCL集体通信,指导构建覆盖张量并行、专家并行、上下文并行、序列并行和数据并行的87个问题(见Figure 3)。

Figure 2 标准Transformer工作负载的并行策略分类
Figure 2 标准Transformer工作负载的并行策略分类
Figure 3 ParallelKernelBench覆盖的分布式工作负载分布
Figure 3 ParallelKernelBench覆盖的分布式工作负载分布

评估指标设计与正确性验证。为了对比不同模型,PKB采用了KernelBench【[41],Ouyang等人,Kernelbench: Can llms write efficient gpu kernels?, 2025】引入的 $fast_p$ 指标,即模型生成的内核既正确又比基线快 $p$ 倍的环境所占的比例。正确性验证通过在5次独立尝试中,使用随机的每rank PyTorch输入,验证生成的 solution() 的输出与参考实留在数值容差内一致。性能测量则在500次预热(Warmup)迭代后,对100次连续启动且无中间同步的分析迭代进行墙上时间(Wall-clock time)的平均。

通信感知屋顶线模型公式。为了衡量生成的内核是否充分利用了硬件,PKB引入了 $roof_u$ 指标,即生成内核正确且达到至少 $u$ 比例屋顶线极限的任务比例。每个任务的理论屋顶线性能通过计算FLOP数和传输字节数来估计。对于划分为 $k$ 个阶段的问题 $i$,阶段感知的屋顶线时间 $T_{roof,i}$ 计算公式为:

$$T_{roof,i} = \sum_{k} \max \left( \frac{W_k}{P_{peak}}, \frac{Q_{mem,k}}{B_{mem}}, \frac{Q_{comm,k}}{B_{comm}} \right)$$

其中,$W_k$ 为第 $k$ 阶段的 FLOP 数,$Q_{mem,k}$ 为该阶段跨越 L2-HBM 边界的字节数,$Q_{comm,k}$ 为该阶段的通信量(以字节为单位)。$P_{peak}$、$B_{mem}$ 和 $B_{comm}$ 分别代表峰值计算性能、HBM 带宽和 NVLink 带宽。

屋顶线模型的参数绑定。对于H100 SXM5 GPU,其计算峰值 $P_{peak}$ 根据数据精度和执行路径设置:BF16/FP16 张量核心为 989 TFLOP/s,FP8 张量核心为 1979 TFLOP/s,TF32 张量核心为 494 TFLOP/s,FP32 单精度为 67 TFLOP/s。HBM 带宽 $B_{mem}$ 设为 3.35 TB/s。NVLink 单向带宽 $B_{comm}$ 设为 450 GB/s。当一个阶段跨越多个进程组轴(如张量并行TP和上下文并行CP)时,每个阶段在各自的轴上计费,共享网络资源的时间在取最大值前进行求和。

基线硬件利用率分析。通过对所有PKB问题进行分析,确定其绑定资源(即在公式中贡献最大的资源)。结果显示,75%的任务受NVLink带宽限制,18%受HBM带宽限制,7%受计算限制。如Figure 4所示,大多数PyTorch基线在通信限制任务中的利用率低于其理论屋顶线的50%。

Figure 4 PyTorch基线远低于屋顶线且LLM优化方案可以缩小该差距
Figure 4 PyTorch基线远低于屋顶线且LLM优化方案可以缩小该差距

A4 实验环境(总结)


A4 实验结果(总结)

Figure 5 各模型在PKB上的失败原因分布及通信机制使用情况
Figure 5 各模型在PKB上的失败原因分布及通信机制使用情况
Figure 6 单次尝试生成的内核很难超越 unoverlapped 基线
Figure 6 单次尝试生成的内核很难超越 unoverlapped 基线
Figure 7 智能体反馈机制下 Gemini 3 Pro 在PKB上的性能爬升
Figure 7 智能体反馈机制下 Gemini 3 Pro 在PKB上的性能爬升
Figure 7b 智能体迭代过程中的加速曲线
Figure 7b 智能体迭代过程中的加速曲线

A5 结论(总结)

论文提出了 ParallelKernelBench (PKB),这是首个用于多GPU内核生成的基准测试和评估框架。研究表明,尽管前沿LLM在单GPU内核生成上取得了显著进展,但在编写高性能多GPU内核方面仍面临严峻挑战。在单次尝试下,即使是最好的模型也只能正确解决约32%的问题,且仅有25%的生成内核能超越未重叠的 PyTorch + NCCL 基线。错误分析表明,模型主要受阻于复杂的分布式内存协调、Rank之间的同步以及对高级硬件通信机制(如TMA、NVLS)的欠利用。

不过,通过智能体迭代反馈,LLM的性能能得到明显改善。此外,LLM在一些非Transformer的、未被人类专家深度优化的小众领域(如Hyena算子、SAM3掩码抑制、NeMo-RL中的词表并行对数概率计算)中,成功生成了能够替代现有参考实现的高性能全新内核。

未来工作展望

  1. 扩展PKB以支持节点间的通信场景(如通过 RoCE 或 InfiniBand),并支持更多的硬件拓扑和加速器(如 Google TPU)。
  2. 在内核生成中探索更高级的编程抽象(如 ParallelKittens、Triton-Distributed)以及新兴的底层接口(如 cuTile 互联功能、NCCL GIN 和 NVSHMEM)。
  3. 探索更先进的智能体推理机制和系统,以提升LLM在复杂分布式系统优化任务中的代码生成质量。

A6 附录(缩写)

Appendix A: 提示词与上下文示例

向量对称相加示例代码分析。如Listing 1所示,PKB提供了一个极简的上下文示例:一个跨越两个GPU rank的元素级 float32 向量相加内核。该示例演示了两个核心原语:(1) 使用 PyTorch 的对称内存 API (torch.distributed._symmetric_memory) 在对称堆上分配对端可直接寻址的缓冲区;(2) 在 CUDA 内核中,通过将远程地址强制转换为指针(使用 reinterpret_cast),利用统一虚拟寻址(UVA)和NVLink点对点(P2P)直接读取对端内存。
系统提示词结构设计。Listing 2展示了系统提示词的结构,该提示词会动态注入当前任务的参考代码、硬件拓扑说明(如 NVIDIA H100 SXM、4个GPU互联、NVLink支持等)以及优化原则。优化原则强调:(1) 最小化热路径上的原生 PyTorch 开销;(2) 优先采用设备端通信(如对称内存、UVA/P2P);(3) 最大化计算与通信的重叠(如双缓冲、多流、流水线分块)。

Appendix B: GPU通信框架调查

设备端融合通信优势。传统的多GPU程序将计算与通信分离,而PKB的核心设计围绕设备端融合通信展开。通过在内核内部使用对称堆API、UVA指针和NVSwitch的 multimem 指令,可以消除主机端启动网络集体的开销。
跨节点与新兴API生态。虽然PKB目前局限于节点内NVLink域,但附录指出了跨节点设备端通信的现状。NVSHMEM 提供了跨节点的一边内存访问(RMA)和集体操作,但其安全机制可能带来额外开销【[48],Sul等人,Parallelkittens, 2025】。Triton Distributed 主要针对 H800 设计。最近引入的 NCCL GIN(NCCL 2.28+)提供了设备端API,允许GPU在无需CPU参与的情况下直接发起RDMA传输【[22],Hu等人,Demystifying nccl, 2026】。

Appendix C: 问题来源与分类

数据源与覆盖范围。PKB的87个问题来自28个开源仓库(如Megatron-LM [37], DeepSpeed [25], DeepEP [59]等)。工作负载主要分为六个抽象层级(Table 4):层级1为纯集体通信,层级2至6将计算与通信进行不同程度的融合(如TP-MLP、MoE路由、环形注意力等)。Table 5总结了22个非LLM问题(来自GNN、三维高斯泼溅、分子动力学等领域)所带来的独特通信特征,例如不规则的稀疏点对点通信、多对多动态大小传输等。

Appendix D: 通信感知屋顶线模型细节与实例

阶段划分与依赖边界。为了避免简单取最大值低估了无法重叠的硬同步时间,PKB采用阶段感知(Phase-aware)屋顶线。在可以进行瓦片流水线(Tile-pipelining)的阶段内,计算、内存和通信完全重叠;而在需要完整中间结果的步骤之间,则将各阶段的屋顶线时间相加。
Hyena Conv1D边界交换算子分析。以问题72(Hyena Conv1D 边界交换)为例,该问题包含两个阶段:(1) Halo P2P阶段:通信量为 $2BH(K-1)$ 字节,无计算和内存流量;(2) 深度一维卷积阶段:计算量为 $4BHSK$ FLOP,内存流量为 $(4BHS + HK)$ 字节。在 batch size $B=1$,通道数 $H=1024$,卷积核大小 $K=7$,本地序列长度 $S=1024$,GPU数 $W=4$,精度为2字节(BF16)的配置下,计算时间为 438 ns,内存时间为 2508 ns,通信时间为 55 ns,因此理论屋顶线时间 $T_{roof,72} \approx 2.51 \, \mu\text{s}$,受内存带宽限制。

Appendix E: 智能体脚手架(Agentic Harness)

Mini-SWE-Agent 工作流。智能体运行在基于 mini-swe-agent【[55],Yang等人,SWE-agent: Agent-computer interfaces enable automated software engineering, 2024】的闭环环境中。如图9所示,智能体被允许通过终端命令修改本地代码、编译运行,并使用专门的评测脚本直接在真实H100硬件上测试正确性和运行耗时。系统设置了最多20次迭代上限,智能体可利用编译器报错和性能反馈不断调整内核参数直至满意。

Appendix F: 重点内核个案分析

NCCL集体通信的向量化复制优化。Gemini 3 Pro 在优化 AllGather 和 Scatter 时,通过动态选择最宽的向量化内存指令(如在16字节对齐时使用 128位的 uint4 读写,在8字节对齐时回退到 64位的 uint2)来最大化NVLink的吞吐量,如图11所示。这使得它在小数据量下超越了NCCL本身(Figure 10与Figure 10b)。

Figure 10 Gemini 3 Pro在AllGather上的性能表现
Figure 10 Gemini 3 Pro在AllGather上的性能表现
Figure 10b Gemini 3 Pro在Scatter上的性能表现
Figure 10b Gemini 3 Pro在Scatter上的性能表现
Figure 11 Gemini 3 Pro的向量化内存拷贝内核优化
Figure 11 Gemini 3 Pro的向量化内存拷贝内核优化

Ulysses序列并行拉取内核。在DeepSpeed-Ulysses Attention中,Gemini 3 Pro 将复杂的 sequence-to-head 转换重写为直接的 NVLink 拉取(Pull)内核(Figure 13)。每个线程根据自身的全局索引,计算出对端GPU的Rank和对应偏移,通过UVA指针直接跨NVLink拉取所需数据。在本地序列长度 $T=1024$ 时,该方案相比 PyTorch+NCCL 实现了 ∼3.73倍的加速(Figure 12)。

Figure 12 Ulysses QKV All-to-All内核相对于参考实现的加速比
Figure 12 Ulysses QKV All-to-All内核相对于参考实现的加速比
Figure 13 Gemini 3 Pro生成的NVLink拉取内核核心逻辑
Figure 13 Gemini 3 Pro生成的NVLink拉取内核核心逻辑

NeMo词表并行对数概率计算融合。针对NeMo-RL中的 vocab-parallel log-probability 计算,Gemini 3 Pro 生成了融合内核(Figure 15),利用 128位的 uint4 写入(Push)操作将 Logits 直接推送到 peer 的对称内存,并使用共享内存中的 warp-shuffle 归约完成 log-softmax 和目标词提取。该内核在 $l=1024$ 时取得了 1.28倍的加速(Figure 14)。

Figure 14 VocabP LogProb TopK内核加速比变化
Figure 14 VocabP LogProb TopK内核加速比变化
Figure 15 Gemini 3 Pro生成的融合对数概率与归约内核
Figure 15 Gemini 3 Pro生成的融合对数概率与归约内核

Hyena前向CP内核的带宽饱和瓶颈。对于Hyena的上下文并行前向传播,GPT-5.5 设计的内核(Figure 17)将三次 all_to_all 的输入打包进单个对称分配中,使用单个 UVA 读取内核(gather_x1_and_u)在读取的同时在线计算 $x_2 \odot v$ 并执行逆向 zigzag 重新索引。如Figure 16所示,在小数据量下该内核取得了 ∼3.55倍的加速;但在大序列长度($l \ge 8192$)下,由于细粒度的线程级UVA访问无法像NCCL大块传输那样饱和NVLink带宽,其性能下降至基线的 0.70倍。

Figure 16 Hyena CP内核在不同本地序列长度下的加速比
Figure 16 Hyena CP内核在不同本地序列长度下的加速比
Figure 17 GPT-5.5在Hyena CP上消除冗余All-to-All的融合方案
Figure 17 GPT-5.5在Hyena CP上消除冗余All-to-All的融合方案

SAM3 IoU抑制的空间压缩优化。在SAM3的 IoU 抑制任务中,GPT-5.5 设计的内核(Figure 19)在跨UVA读取对端掩码的同时,在线将其压缩为 32位字(Bitpacking),并使用硬件快速位计数指令 __popc 计算交集,避免了基线中将掩码转为 float 矩阵再进行密集乘法的高额开销。在16个对象时该内核加速比为 1.38倍,但随着对象数增加,二次方复杂度的计算占主导,加速比收敛至 1.03倍(Figure 18)。

Figure 18 SAM3 IoU抑制内核随目标物体数量增加的性能加速趋势
Figure 18 SAM3 IoU抑制内核随目标物体数量增加的性能加速趋势
Figure 19 GPT-5.5为SAM3掩码抑制设计的位打包对称内存内核
Figure 19 GPT-5.5为SAM3掩码抑制设计的位打包对称内存内核

Appendix G: Blackwell GPU Performance

Blackwell平台性能衰减分析。在 Blackwell (B200) 平台上,使用 B200 节点对 Gemini 3 Pro、Claude Opus 4.7 和 GPT-5.5 进行的测试表明,正确率仍低于 30%,且加速比随着要求提升而急剧衰减(Figure 20)。这印证了LLM在分布式并行内核生成上的主要瓶颈是逻辑推理能力,而非硬件平台的限制。

Figure 20 三种前沿模型在Blackwell平台上的性能加速比曲线
Figure 20 三种前沿模型在Blackwell平台上的性能加速比曲线

方法细节中引用的参考文献汇总

  1. 参考文献【22】:Zhiyi Hu et al. "Demystifying nccl: An in-depth analysis of gpu communication protocols and algorithms", 2026.

    • 引用段落:Section 2.1 "Communication Software" 与 Appendix B
    • 原文描述:NCCL是GPU集体通信的标准库,它将 all-reduce 和 reduce-scatter 等操作暴露为在 CUDA 流上排队的异步内核。
  2. 参考文献【41】:Anne Ouyang et al. "Kernelbench: Can llms write efficient gpu kernels?", 2025.

    • 引用段落:Section 1 "Introduction" 与 Section 4.1 "Metric Design"
    • 原文描述:KernelBench 建立了单 GPU 内核生成的标准评估框架,并引入了 $fast_p$ 指标,用于衡量生成内核在正确且达到特定加速比时的比例。
  3. 参考文献【48】:Stuart H. Sul et al. "Parallelkittens: Systematic and practical simplification of multi-gpu ai kernels", 2025.

    • 引用段落:Section 2.1 "Transfer Mechanisms" 与 Appendix B
    • 原文描述:提出了 GPU 间数据传输的三种主要机制分类(Copy Engine, TMA, SM load/store),以及基于对称内存编程的简化多 GPU 内核开发框架。
  4. 参考文献【52】:Samuel Williams et al. "Roofline: an insightful visual performance model for multicore architectures", 2009.

    • 引用段落:Section 1 "Introduction" 与 Section 4.2 "Understanding the Communication Bottleneck"
    • 原文描述:经典的 Roofline 模型通过单节点的峰值计算性能和 HBM 带宽来限制内核的运行时间。
  5. 参考文献【55】:John Yang et al. "SWE-agent: Agent-computer interfaces enable automated software engineering", 2024.

    • 引用段落:Section 5.3 "Impact of Agentic Harness & Iterative Feedback" 与 Appendix E
    • 原文描述:Mini-SWE-Agent 提供了一个允许 LLM 运行终端命令、观察环境反馈并进行闭环迭代的代码维护与软件工程框架。
  6. 参考文献【25】:Sam Ade Jacobs et al. "Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models", 2023.

    • 引用段落:Appendix F.2 "Context Parallel: Ulysses (Gemini 3 Pro)"
    • 原文描述:Ulysses Attention 通过在注意力计算前将数据从序列并行转换为头并行布局,计算后再转换回来,从而支持极长序列的分布式训练。
  7. 参考文献【43】:Michael Poli et al. "Hyena hierarchy: Towards larger convolutional language models", 2023.

    • 引用段落:Section 6.2 "Net New Kernels" 与 Appendix F.4 "Hyena Forward CP (GPT-5.5)"
    • 原文描述:Hyena 算子通过全局的一维卷积替代注意力机制,其上下文并行(CP)前向传播需要在序列并行和通道并行布局之间反复转换。
  8. 参考文献【10】:Nicolas Carion et al. "Sam 3: Segment anything with concepts", 2025.

    • 引用段落:Section 6.2 "Net New Kernels" 与 Appendix F.5 "SAM3 all-gathered mask IoU suppression (GPT-5.5)"
    • 原文描述:SAM3 在分布式视频对象分割中,需要跨 GPU 收集预测掩码并进行全局相交比(IoU)计算以抑制重复预测。