Loading post…
Together AIBenchmark
ParallelKernelBench: Frontier LLMs can't write fast multi-GPU kernels (yet)
Summary
ParallelKernelBench (PKB) measures whether language models can replace PyTorch + NCCL implementations with fast CUDA kernels for multi-GPU workloads, where communication can account for over 20% of inference latency. Its 87 problems come from real codebases and span tensor, context, data, expert, sequence, and FSDP/ZeRO parallelism, with hardware topology supplied and symmetric memory used for direct GPU communication over NVLink. In zero-shot evaluation, GPT-5.5 solved 28 problems and produced 22 faster-than-baseline solutions; sampling three attempts raised those figures to 36 and 27, but the best fast 1@3 result was only 31%. Models performed best on familiar collective primitives, tensor-parallel GEMMs, and Ulysses-style context parallelism, while failures often involved rank coordination, data partitioning, collective ordering, and underuse of TMA or NVLS.
Context
Multi-GPU kernel generation differs from single-GPU generation because parallelism compositions create varied communication patterns, interconnect bandwidth can dominate performance, and kernels must choose among copy engines, TMA, SM load/store, and NVLS while deciding whether to fuse communication with computation. PKB’s communication-aware roofline found that most problems are NVLink-bottlenecked and that PyTorch + NCCL baselines leave performance headroom.
Approach / What changed
PKB gives each model a task, hardware topology, and PyTorch + NCCL reference implementation. The model must produce a custom CUDA kernel using symmetric memory for direct multi-GPU communication, and the result is evaluated for correctness, wall-clock speedup, and communication roofline. The 87 tasks were selected from a taxonomy covering distributed workload types and codebases including Megatron-LM, DeepSpeed, DeepEP, TensorRT-LLM, and NeMo-RL.
Takeaways
PKB contains 87 real-world problems spanning tensor, context, data, expert, sequence, FSDP/ZeRO, and other parallelism patterns.
GPT-5.5 improved from 28 correct and 22 faster-than-baseline solutions at one attempt to 36 and 27 with three attempts, but fast 1@3 reached only 31%.
Generated kernels commonly used copy engines or SM load/store instructions; TMA and NVLS were almost absent, and difficult failures involved rank coordination, partitioning, and collective ordering.
FAQ
What does PKB require a model to generate?
Each task starts with a PyTorch + NCCL implementation and a hardware topology. The model must replace it with a CUDA kernel that uses symmetric memory to communicate directly across GPUs over NVLink.
Which multi-GPU patterns were models best at handling?
Successes were concentrated in collective primitives, tensor-parallel GEMMs, and Ulysses-style context parallelism. These patterns are more visible in open-source code, giving models stronger prior examples.
Why can a compiling kernel still be incorrect or slow?
The difficult parts extend beyond CUDA syntax to rank coordination, data partitioning, and collective ordering. Models also used a narrow set of communication mechanisms and rarely selected specialized mechanisms such as TMA or NVLS.
Source and freshness
- Source
- together.ai
- Published
- Added to Yomu
Summaries are generated from the original by a language model. Check the original before relying on it.
Keep reading
Related reading
The Good Stuff
Paid picks we use ourselves. Links carry a referral code — a subscription through one earns us a commission at no extra cost to you.
The best engineering reads, once a week
One useful email. Unsubscribe anytime.