ParallelKernelBench: Can LLMs Write Fast Multi-GPU Kernels?
这是什么
现有 kernel 生成基准(KernelBench、TritonBench 等)几乎全是单 GPU 问题,而真实训练/推理系统里通信开销可占推理延迟 20% 以上,计算-通信重叠、集合通信融合才是分布式性能工程的主战场。ParallelKernelBench(PKB)由 Together AI(Willy Chan、Simran Arora、Dan Fu 等,含 Stanford KernelBench 原班人马 Simon Guo)于 2026-06 发布,把任务定义为:给定一段 PyTorch + NCCL 的参考实现,让模型改写成利用 symmetric memory 直接走 NVLink 的细粒度 CUDA(或 Triton / ParallelKittens DSL)kernel。
87 个问题全部取自真实代码库(Megatron-LM、DeepSpeed、DeepEP、TensorRT-LLM、NeMo-RL,外加 GNN、分布式 FFT、Gaussian splatting 等非 LLM 负载),按并行模式分 12 类:Tensor Parallel 17、Context Parallel 12、Expert Parallel 11、Data Parallel 9、FSDP/ZeRO 9、集合通信原语 8 等,Pipeline Parallel 只有 1 题。评测在 4x H100(NVLink/NVSwitch)上做,逐 rank 比对输出、100 次随机输入验证正确性,性能按 ThunderKittens 2 的规范计时(500 warmup + 100 timed)。
机制与做法
任务与评测协议
每个问题在 reference/ 下有一份 PyTorch+NCCL 参考实现;模型拿到任务说明 + 网络拓扑/硬件能力描述(如 4x H100、NVSwitch、NVLS 开启)+ 参考代码,生成单文件解答放进 solutions_<backend>/。eval 模式用 torchrun 起多进程,对每次随机输入先跑 reference 再跑候选,逐 rank 比对 rank_*.pt 张量(atol/rtol 容差),首个不匹配即判错;性能报告相对 reference 的 speedup,另有近似 roofline 估计通信利用率。
生成侧支持 zero-shot(kernelgen/generate_kernel.py,prompt 模板在 prompts.toml)和 agentic 两种:后者接 mini-swe-agent,允许模型拿编译/测试反馈迭代。评测可本地跑,也可用 Together Containers(jig)云端跑——作者自嘲'不是每个人都能在 Costco 买到 DGX 8xH100'。
主要结果与失败模式
zero-shot:GPT-5.5 最好,28/87 正确、22 个快于 baseline;3 采样下 36 正确、27 更快,fast1@3 峰值 31%。Gemini 3 Pro pass@1→3 为 24→30,Claude Opus 4.7 为 20→31,DeepSeek V4 Pro 只有 2→3。Pipeline Parallel 类没有任何模型解出。agentic 模式(Gemini 3 Pro + 编译/测试反馈)从 24 提升到 35/87,但约 20 步 refinement 后就平台化。
失败模式分层明显:弱模型死在编译期;强推理模型能编译但输出错误或直接死锁——难点在 rank 间协调、数据切分和集合通信的顺序推理。生成的 kernel 几乎只用 copy engine 或 SM load/store,TMA 和 NVLS 这类新硬件特性基本缺席。也有亮点:Gemini 3 Pro 给 NeMo-RL 的 vocab-parallel log-prob 写出了没有公开优化版可比的融合 kernel,GPT-5.5 在 SAM 3 mask IoU suppression 上用上了 bitpacked mask + 硬件 popcount。
关键结果
- 87 个问题、12 个并行类别,全部来自真实代码库(Megatron-LM、DeepSpeed、DeepEP、TensorRT-LLM、NeMo-RL 及 GNN/FFT/3DGS 等),在 4x H100 上评测。
- 最强模型 GPT-5.5 zero-shot 仅 28/87 正确、22 个跑赢 PyTorch+NCCL baseline;fast1@3 峰值 31%——即'不到三分之一'。
- Pipeline Parallel 类 0 解;DeepSeek V4 Pro 几乎全挂(pass@1 只有 2/87)。
- agentic 反馈循环有效但有限:Gemini 3 Pro 从 24 提到 35/87 正确,约 20 步后收益平台化。
- 失败模式:强模型的 kernel 能编译但输出错或死锁,核心短板是 rank 协调与集合通信顺序推理;TMA、NVLS 等硬件特性几乎无人使用。
- 个别生成 kernel 超越公开实现(如 GEMM+All-Gather 87.9μs vs baseline 320.6μs),但离 roofline(4.99μs)仍差一个数量级。
实证核查
与我们方向的关系
对 ai4ai 方向这是 KernelBench 谱系的自然延伸,也是目前少有的多 GPU/通信感知 kernel 生成基准:它把'LLM 自动化 GPU 编程'的短板从算子级(单卡已被 RL 训出的模型刷得很高)推到系统级(rank 协调、通信-计算重叠),后者恰是分布式训练自动化真正值钱的部分。31% 的 fast1@3 上限和 Pipeline Parallel 零解说明这块还远未饱和,适合作为 agent/RL for kernel generation 工作的下一个靶子。
可借鉴的点:(1) 逐 rank 输出比对 + 随机多 trial 的分布式正确性验证协议,比单卡 allclose 严格得多;(2) prompt 里显式给网络拓扑与硬件能力(NVSwitch、NVLS 开关)的做法;(3) agentic 反馈约 20 步即平台化的观察——单靠编译/测试信号迭代对通信 bug(死锁、错序)帮助有限,可能需要 profiler/trace 级反馈。注意目前仅限单机 NVLink,internode(NVSHMEM、prefill-decode 分离)在 roadmap 上还没做。