The Shift to Communication-Bound Workloads

As AI hardware evolves, the bottleneck for large-scale training and inference has shifted from raw compute to the interconnects between GPUs. Between the NVIDIA A100 (2020) and B200 (2024), BF16 tensor core throughput increased by 7.2x, while intra-node communication only improved by 3x and inter-node by 2x. Standard baselines like PyTorch + NCCL, which are designed for bulk transfers, frequently fall below 50% of the communication-aware roofline because they introduce synchronization overheads and fail to leverage fine-grained, direct NVLink transfers.

The Fundamentals of Multi-GPU Kernel Design

To optimize these workloads, developers must navigate three primary transfer mechanisms, each with distinct trade-offs:

  1. Copy Engine: Best for large messages; offloads work from the GPU processors but requires host/CPU initiation.
  2. Tensor Memory Acceleration (TMA): Device-initiated; saturates NVLink bandwidth with smaller messages, making it ideal for fine-grained communication.
  3. Register-Level Transfers: Necessary for leveraging in-network reductions via NVSwitch, though they consume precious register space.

Beyond transfer mechanisms, developers must choose between Intra-SM overlapping (specializing warps within a processor) and Inter-SM overlapping (dedicating entire processors to compute or communication). The choice depends on whether the compute and communication patterns align on the same data inputs.

ParallelKittens: A Practical Abstraction

Together AI developed ParallelKittens to simplify this complexity. It provides a set of minimal primitives that allow developers to inject multi-GPU communication logic into single-GPU kernels with roughly a dozen lines of code. This approach enables direct NVLink loads and stores, bypassing the staging overheads inherent in standard libraries like NCCL.

LLM Performance on ParallelKernelBench

To test if frontier models can reason through these trade-offs, the team created ParallelKernelBench, a suite of 87 real-world multi-GPU problems. The results were sobering:

  • Correctness vs. Speed: While models can generate correct code, they struggle to generate faster code. Correctness plateaus around 36/87 problems, but the number of solutions that actually outperform the baseline stalls near 31%.
  • The Reasoning Gap: Failures are rarely due to CUDA syntax. Instead, models fail on collective ordering, data partitioning, and selecting the correct transfer mechanism. Successes are largely limited to patterns heavily represented in public training data (e.g., standard tensor-parallel GEMMs).
  • Scaling Limits: Increasing test-time compute (sampling) improves correctness but does not significantly improve the ability to find optimal performance, suggesting that models are pattern-matching rather than reasoning from first principles about hardware topology.

Key Takeaways

  • Communication is the new compute: Optimize for the interconnect (NVLink/NVSwitch) rather than just the SMs.
  • Avoid bulk-transfer defaults: Standard libraries like NCCL are often too rigid for fine-grained, high-performance kernels.
  • Use specialized primitives: Abstractions like ParallelKittens allow for direct device-initiated transfers (TMA) that outperform CPU-initiated copy engines.
  • LLMs are not yet systems engineers: Models struggle with multi-GPU kernels because they lack a structural understanding of hardware topology and non-obvious performance trade-offs.
  • Prioritize topology awareness: When writing custom kernels, the choice between Intra-SM and Inter-SM scheduling is often the difference between peak performance and a bottlenecked system.

Notable Quotes

  • "Communication is increasingly consuming the majority of the runtime and yields low model flop utilization at scale."
  • "The design of NCCL really breaks down when you care about peak performance, fine-grained communication, and sort of non-trivial collectives that you want to fuse together."
  • "The success patterns here are really concentrated into familiar patterns... in other words, patterns that we see heavily represented on the internet rather than necessarily patterns that the model has used its reasoning abilities to think through."
  • "We found that there's deeper issues than CUDA syntax... models compile after a retry and then stall on collective ordering, data partitioning, and the choice between the copy engine, tensor memory acceleration, and register-level transfers."