ParallelKittens: Communication and Compute in Multi-GPU Kernels
Multi-GPU execution adds a communication schedule to the model’s computation graph. A matrix multiplication can finish on one device while its consumer waits for a partial result from another. The time spent at that boundary depends on message size, topology, synchronization, and how much independent work is available during the transfer.
ParallelKittens approaches this as a kernel-design problem. It extends ThunderKittens with reusable primitives for combining computation and communication, so data movement can be scheduled alongside the operations that produce and consume it.1 The useful target is exposed communication: the portion of transfer time that remains on the critical path after overlap.
Communication follows the partition
The parallelism strategy determines what moves between devices.
Tensor parallelism partitions operations within a layer. Devices compute partial results and use collectives to combine or redistribute them. These boundaries can occur repeatedly during every generated token.
Expert parallelism places different experts on different devices. Tokens travel to the selected experts, then their outputs return for combination. The traffic depends on routing decisions as well as batch size.
Pipeline parallelism assigns consecutive stages to different devices. Activations move between stages, and the schedule determines whether a stage has useful work while it waits for the next microbatch.
Distributed attention partitions sequence or attention state. The kernel exchanges the shards needed to form the output, with communication shaped by the chosen attention algorithm.
These patterns create different opportunities. A regular reduction has a different schedule from a variable-sized exchange of expert tokens. The relevant network workload is the sequence of these operations, including their dependencies and sizes.
Collective operations
An all-reduce combines values across participants and returns the reduced result to each. An all-gather assembles distributed shards into a complete result on every participant. Reduce-scatter combines reduction with partitioned output. NCCL specifies these semantics independently of the algorithms used to execute them.2
The distinction is useful when reading a model trace:
| Boundary | Data that moves | Main scheduling dependency |
|---|---|---|
| Tensor-parallel layer | Partial activations or reductions | The next operation needs the combined result |
| Expert dispatch | Tokens and routing metadata | Each expert needs its assigned inputs |
| Expert combine | Weighted expert outputs | The following layer needs the combined activation |
| Pipeline stage | Activations | The next stage waits for a microbatch |
| Distributed attention | Attention-related shards | A tile waits for the required remote state |
The runtime selects a particular implementation for each boundary. That selection establishes the actual number of stages, temporary buffers, and synchronization points. The same high-level model can therefore produce different communication behavior under different serving engines.
Latency and bandwidth regimes
A first-order transfer model separates startup latency from transmission time:
transfer time ≈ startup latency + message bytes / effective bandwidth
Effective bandwidth is the rate achieved on the path being used. Collective algorithms add stages, and topology changes which links carry the traffic, but the model already explains why small and large transfers favor different optimizations.
For an illustrative 4 microsecond startup and 100 GB/s effective bandwidth, transmitting 4 KiB takes about 0.041 microseconds. Nearly all the elapsed time is startup. A 64 MiB transfer takes about 671 microseconds for transmission, so bandwidth dominates instead.
Doubling bandwidth barely changes the small-message case. Reducing launch overhead or combining messages can improve it substantially. Large transfers benefit more directly from higher sustained bandwidth and avoiding redundant copies.
Message fusion introduces a scheduling tradeoff. Combining several small transfers amortizes startup, but the combined transfer waits until its inputs are ready. If one result is available much earlier than the others, fusion can delay its consumer. The useful grouping follows readiness and dependencies as well as byte count.
The physical path
Within a server, communication may use PCIe or a dedicated GPU interconnect such as NVLink. Across servers, the path includes network adapters and a fabric such as InfiniBand or Ethernet with RDMA. Placement determines how often the execution crosses each boundary.
A frequently executed collective benefits from keeping its participants on a fast local path. A larger model may require crossing a slower boundary to fit, in which case the partition plan determines how much traffic pays that cost. Capacity and communication are coupled decisions.
RDMA can remove host-side copies from the transfer path. It does not remove queueing in the fabric, oversubscribed links, or uneven traffic among destinations. Those effects appear in the tail of distributed execution, especially when every participant waits for the slowest result.
Compute-communication overlap
For a serial stage, elapsed time is compute time plus communication time. With independent work and sufficient resources, the two can overlap:
serial stage: T = compute + communication
ideal overlap: T = max(compute, communication)
exposed communication = time remaining on the critical path
A stage with 6 ms of compute and 4 ms of communication takes 10 ms serially and 6 ms under ideal overlap, a 1.67x stage speedup. The bound assumes the transfer does not slow the computation through shared resource contention.
Its application-level effect depends on the fraction of total time in that stage. If it occupies 20 percent of request time, the idealized normalized total becomes approximately 0.8 + 0.2/1.67 = 0.92. That is about a 1.09x request speedup.
Tiling creates finer dependencies. A consumer may begin once the first tile arrives rather than waiting for the entire tensor. The producer, transfer, and consumer can then form a pipeline. Tile size controls reuse, synchronization frequency, and the amount of work exposed for overlap.
The pipeline also needs explicit ownership. A buffer cannot be overwritten while another device is still reading it, and a consumer must observe the producer’s completed writes before using the tile. Synchronization is part of the algorithm, with both a correctness cost and a scheduling cost.
The ParallelKittens template
ParallelKittens organizes kernels around loader, storer, consumer, and communicator roles. Loaders and storers move tiles, consumers perform computation, and communicators handle inter-GPU work. Its primitives and template encode recurring scheduling patterns across multi-GPU kernels.1
This makes transfer mechanism and resource assignment visible design choices. Copy-engine transfers and thread-level peer operations have different overheads. Dedicated communication resources can make progress independently, while using those resources also changes what remains available to computation.
The paper evaluates Hopper and Blackwell workloads and reports maxima of 2.33x for data- and tensor-parallel workloads, 4.08x for sequence-parallel workloads, and 1.22x for expert-parallel workloads.1 These results describe the tested kernels and baselines. Their significance is that a reusable design can capture gains across several communication patterns rather than requiring an unrelated implementation for every operator.
Reading a distributed trace
Nsight Systems separates CPU overhead, launch latency, and GPU execution time.3 Across multiple devices, those timelines explain whether a gap comes from delayed submission, a dependency, or data movement.
A collective that appears slow on most ranks can be waiting for one late participant. If seven GPUs arrive at the boundary early and the eighth finishes its preceding work late, the apparent collective duration includes that skew. The cause may be expert imbalance, input preparation, or a competing workload on the late device.
A trace therefore connects three intervals: production of the input, execution of the transfer, and the point where the consumer resumes. Communication that overlaps successfully may occupy a large interval without extending the critical path. A short transfer immediately before a dependent operation may matter more.
The shape sweep includes small and large messages, prefill and decode batches, and representative sequence lengths. Packing, unpacking, barriers, and launch overhead belong in the measured operation because they are part of moving a usable tensor to its consumer.
Kernel and system validation
A reference implementation establishes the expected output, with numerical tolerances appropriate to the operation. Partial tiles, changing rank counts, irregular expert assignments, and repeated execution exercise the boundaries most likely to expose ordering errors.
The performance comparison has two levels. A kernel benchmark measures whether the transfer and computation schedule improved. A serving benchmark measures whether that improvement reduces request latency, increases capacity, or lowers cost at the same output quality.
Custom fused kernels are most useful when the operation repeats often, accounts for a substantial share of runtime, and has a manageable shape space. A mature collective implementation remains a valuable baseline and fallback as model layouts and hardware change.
The connection between networking and kernel design is direct: the partition decides which data moves, the topology decides its path, and the schedule decides how much of the movement is visible to the next operation. Improving any of those can turn idle accelerator time into useful computation.
References
Footnotes
-
Sul, Arora, Spector, and Ré, ParallelKittens: Systematic and Practical Simplification of Multi-GPU AI Kernels, 2025. The reported 2.33x, 4.08x, and 1.22x maxima above come from this paper’s tested workloads. ↩ ↩2 ↩3
-
NVIDIA, NCCL collective operations. Definitions of all-reduce, all-gather, reduce-scatter, and related operations. ↩
-
NVIDIA, Understanding the Visualization of Overhead and Latency in Nsight Systems. Interpreting launch and execution timelines. ↩
Each rank starts with a local contribution to a vector.
- Reduce-scatter
Reduce matching elements across ranks and leave each rank with a distinct reduced shard.
- All-gather
Exchange the reduced shards until every rank holds the complete reduced vector.
This decomposition exposes where shard computation, transfer, and dependent work can be scheduled.
Toy independent work: compute takes 10 ms and communication takes 4 ms. Ignore launch, contention, and dependency overhead.
The best-case overlap bound is max(10, 4) = 10 ms, rather than their 14 ms serial sum.
An illustrative lower bound, not a ParallelKittens measurement. Dependencies or bandwidth contention can prevent complete overlap.