YC Paper Club · Edition 03Talk 2 / 6 · Multi-GPU kernels · advanced
  1. 1Specialized Chips
  2. 2ParallelKittens
  3. 3Intelligence / Watt
  4. 4AI Writes Kernels
  5. 5Heterogeneous AI
  6. 6Madrona

ParallelKittens — Making Eight GPUs Feel Like One Machine.

Jul 30, 2026 · paperjuice · 11 min read · 2295 words advanced

ParallelKittens — Making Eight GPUs Feel Like One Machine.

Presented by Stuart Sul (Stanford / Cursor)

paperjuiceml-systemscudamulti-gpukernels

Imagine eight brilliant cooks sharing one kitchen. Each cook can chop faster every year, but the door between their workstations is barely getting wider. Soon the cooks spend more time passing bowls than cooking. That is the multi-GPU problem Stuart Sul brought to YC Paper Club: GPU arithmetic has accelerated faster than the links that move tensors between GPUs.

ParallelKittens asks a practical question: can a tiny vocabulary make communication-heavy GPU programs both readable and fast, without hiding the hardware that determines performance?

The problem before the framework

Large models do not fit neatly on one accelerator. Data parallelism copies a model and splits examples; tensor parallelism splits matrix operations; sequence parallelism splits tokens; expert parallelism routes tokens among mixture-of-experts workers. Every strategy creates a different traffic pattern. Libraries such as NCCL and NVSHMEM are essential, but a generic collective can lag new hardware features and makes it difficult to fuse communication with the surrounding compute.

The expensive mistake is to think of communication as a separate stage: compute, stop, transfer, stop, compute again. A fast kernel overlaps the two. While one tile travels across NVLink, another tile should occupy the tensor cores. The challenge is scheduling enough independent work without drowning the programmer in barriers, memory spaces, and device-specific instructions.

The kitten-sized idea

ParallelKittens extends ThunderKittens and treats remote GPU HBM as another level of the memory hierarchy. It supplies data structures for registers, shared memory, global memory, and peer global memory, then applies communication primitives suited to those locations.

Its program template has four worker roles—loader, consumer, communicator, and store—that support different schedules. Stuart reports roughly 50–100 lines of device code for kernels that match or beat hand-optimized implementations often hundreds or thousands of lines long.

Three knobs that explain most performance

  • Transfer mechanism: copy engines, TMA, and register load/store instructions have different message-size, SM-use, and in-network-compute trade-offs.
  • Resource schedule: intra-SM overlap assigns communication and compute to threads in one SM; inter-SM overlap dedicates some whole SMs to communication.
  • Design overhead: conveniences such as intermediate buffers add real data movement. Stuart reports up to an 80% all-reduce gain after removing that overhead in their fine-grained case.

A concrete story: ring attention

Suppose a long sequence is split across eight GPUs. Each GPU owns some queries and a block of keys and values. It computes attention against the local block, sends that block onward, and immediately computes against the next block that arrives. A naive implementation alternates network and math. A tiled pipeline keeps both busy: one tile is being consumed by tensor cores while the following tile crosses the fabric.

The same vocabulary stretches across all-reduce, all-gather plus matrix multiplication, ring attention, and expert routing. That reuse is the research contribution: not a single heroic kernel, but evidence that a small set of hardware-driven principles transfers across parallelism strategies.

What the numbers actually say

On Hopper and Blackwell systems, the paper reports up to 2.33× speedup for data- and tensor-parallel workloads, 4.08× for sequence-parallel workloads, and 1.22× for expert-parallel workloads. “Up to” matters: these are workload- and shape-dependent peaks, not a promise that every model becomes four times faster. The broader result is that compact kernels can match or beat mature implementations across several communication patterns.

Parallel patternWhat movesPK result reported
Data / tensorgradients or partial matrix resultsup to 2.33×
Sequencekey/value or sequence tilesup to 4.08×
Experttokens routed to expertsup to 1.22×

The honest boundaries

The presentation is specifically about tightly connected scale-up systems such as DGX and NVL72. Stuart also reports real adoption: Cursor uses ParallelKittens while training Composer on tens of thousands of Blackwell GPUs, and Together AI uses it to optimize inference workloads.

The takeaway

ParallelKittens does not make networking free. It makes the cost visible and schedulable. The story is the same as the kitchen: do not ask the cooks to wait at the door. Package work into tiles, pass them at the right moment, and keep every workstation busy.

Why networking becomes visible after single-GPU optimization

Stuart begins from a success story. FlashAttention and related work reduced waste inside one GPU. Kernel libraries improved tiling, memory movement, and fusion. As local computation becomes more efficient, time spent communicating among GPUs occupies a larger fraction of the total. This is an instance of Amdahl’s law: improving one component makes the remaining components determine the speed limit.

His cited example—networking consuming as much as half the runtime in a Llama prefill workload—is not a universal constant. It is evidence that communication can be first-order. The fraction changes with model shape, parallelism strategy, topology, and batch. The right conclusion is “measure communication,” not “networking is always 50%.”

From whole tensors to tiny transfers

Traditional distributed programs often treat communication as a coarse stage: finish a large computation, call a collective, wait, then continue. Modern kernels increasingly communicate tiles or tokens while other work is still running. The transfer may contain only a few kilobytes or less.

Fine granularity creates an opportunity and a problem. It enables overlap, so a GPU can compute on one tile while receiving another. But mechanisms designed for large messages may perform poorly when invoked repeatedly on tiny pieces. Startup cost, synchronization, intermediate buffers, and the number of SMs consumed by communication all matter.

The simplified GPU model

An SM, or streaming multiprocessor, is the unit that executes groups of GPU threads. Modern accelerators contain many SMs, backed by registers, shared memory, caches, and HBM. The exact counts differ by architecture; the programming principle does not. Expensive compute capacity is useful only when data arrives on time.

Within one GPU, a well-pipelined kernel fetches the next tile while computing the current tile. Across GPUs, remote HBM becomes another level in the memory hierarchy. NVLink provides a direct high-bandwidth path between GPUs, avoiding the need to route all peer traffic through shared PCIe and the host.

The kernel author’s goal is therefore not simply “send bytes quickly.” It is “arrange computation and communication so that the SMs rarely wait.” This requires knowing which mechanism moves the data, which execution resources perform the movement, and how completion is signaled.

Transfer mechanism 1: the copy engine

A copy engine is a DMA-style unit separate from the SMs. Its attraction is obvious: data can move while all SMs remain available for computation. The limitation in Stuart’s measurements is message size. The engine needs sufficiently large transfers to saturate NVLink; fine-grained kernels often send only a few kilobytes, where setup overhead prevents peak throughput.

This does not make copy engines bad. They remain appropriate for larger, less frequent transfers and cases where preserving every SM for computation matters. ParallelKittens turns the choice into an explicit workload decision instead of assuming the copy engine is always ideal.

Transfer mechanism 2: TMA

The Tensor Memory Accelerator can move multidimensional data asynchronously and maintain useful throughput at smaller message sizes. Stuart reports that on their Blackwell setup, roughly 15 of 148 SMs were enough for TMA-driven traffic to saturate NVLink. That leaves most SMs for compute while avoiding the small-message weakness of the copy engine.

TMA has a different limitation: it cannot use every in-network-compute capability required by some collective algorithms. Therefore “TMA is faster” is incomplete. It may be the right mechanism for all-gather plus GEMM in their evaluation, while a reduction needing network-side arithmetic may favor another path.

Transfer mechanism 3: register instructions

Ordinary load and store instructions give low-level control and can access features unavailable through TMA. The price is programming difficulty. Threads must issue accesses in patterns that coalesce efficiently. Register use can reduce occupancy, and the author must reason carefully about the relationship between communication instructions and useful arithmetic.

The three mechanisms form a trade space. Copy engines preserve SMs but prefer large messages. TMA handles smaller messages efficiently with some SM participation. Register instructions expose maximum control and in-network features but increase pressure and complexity.

Overlap strategy 1: intra-SM specialization

In intra-SM overlap, threads or warps within the same SM take different roles. Some move data while others compute. This is often called warp specialization. Communication and computation can coordinate through fast on-chip mechanisms, and the kernel does not reserve entire SMs exclusively for networking.

The constraint is resource alignment. Both activities share registers and shared memory. If they operate on unrelated data or require too much state, the SM runs out of fast local storage and occupancy falls. Intra-SM overlap works best when communication naturally feeds the computation running beside it.

Overlap strategy 2: inter-SM specialization

In inter-SM overlap, several SMs become communication workers and the rest compute. This separation is easier when network work and arithmetic have different schedules. It can also support prefetching remote data into local HBM before a later kernel consumes it.

The obvious cost is that communication SMs are not using their FLOPs for model arithmetic. Coordination crosses the memory hierarchy rather than remaining entirely within one SM. Nevertheless, Stuart shows that this can be the better schedule when local prefetching or loosely coupled activities dominate.

The remote-cache subtlety

One of the talk’s most useful details concerns repeated access to remote data. A fetch from peer HBM can pass through the local GPU’s memory path without leaving the data in the local cache hierarchy in the way a programmer might expect. Repeatedly reading a remote KV cache can therefore pay the remote cost again.

Prefetching the data into local HBM changes later accesses into local ones. That extra copy sounds wasteful, but it can win when the data is reused enough. Inter-SM communication workers are a natural way to stage this prefetch while other SMs continue computing.

Why different fused operators choose different schedules

Stuart names GEMM reduce-scatter as an intra-SM-friendly case and GEMM all-reduce as an inter-SM-friendly case. The point is not to memorize two labels. Reduce-scatter can align communication closely with the partial results being produced. All-reduce may benefit from communication workers that progress a more independent collective schedule.

A generic compiler cannot always infer these choices from a high-level graph, especially when hardware behavior is undocumented or rapidly changing. ParallelKittens makes the schedule visible to the programmer while keeping the code far smaller than hand-written low-level CUDA or PTX.

Design overhead: convenience can move bytes twice

Distributed libraries often use intermediate buffers because they simplify ownership, synchronization, and generality. Fine-grained workloads magnify that extra traffic. If every tiny exchange first lands in a staging buffer and is then copied to its destination, the system performs more memory movement than the mathematical operation requires.

Stuart reports an all-reduce improvement of up to 80% after stripping such overhead in their setting. The qualifier matters: this is a result for a particular implementation and regime. Its general lesson is that abstractions should expose controls that materially affect data movement rather than hiding them behind one universal interface.

What the ParallelKittens programming model contains

ParallelKittens extends ThunderKittens’ data-oriented approach. It represents values at several memory levels: register tiles, shared-memory tiles, global-memory layouts, and peer-global-memory layouts. Communication operations accept these structures and select mechanisms appropriate to the source, destination, and operation.

The four worker roles provide a readable pipeline. A loader brings local data into the fast hierarchy. A communicator exchanges data with peers. A consumer performs the central arithmetic. A store worker writes results. Real kernels can combine or specialize these roles, but the template makes the flow explicit.

How to read the code-size claim

Stuart reports roughly 50–100 lines of device code for kernels that match or outperform implementations hundreds or thousands of lines long. This is not a claim that multi-GPU programming is easy for a beginner. The framework packages hard-won knowledge about memory structures, transfer mechanisms, and worker schedules.

Small code matters for more than aesthetics. It is easier to inspect, modify for a new model shape, and present to an AI coding agent. It also reduces the surface where synchronization mistakes hide. But correctness and performance still require tests across shapes, devices, and concurrency.

Parallelism patterns covered by the evaluation

PatternWhy communication appearsTypical kernel opportunity
Data parallelworkers combine gradients or updatesfuse reduction with surrounding work
Tensor parallelone matrix operation is split across GPUsoverlap collectives with local GEMM
Sequence paralleltokens or attention state are distributedstream tiles as attention progresses
Expert paralleltokens move to selected expertsreduce routing and all-to-all bubbles

Production adoption and what it proves

Cursor uses ParallelKittens while training Composer on a large Blackwell fleet, and Together AI uses it for inference optimization. Adoption is evidence that the abstractions can survive outside a paper benchmark. It does not prove that every operator or topology should be rewritten in ParallelKittens.

Production introduces requirements beyond kernel speed: deterministic builds, integration with frameworks, monitoring, fallbacks, numerical validation, and maintenance across driver and architecture changes. A compact framework is valuable partly because those operational tasks become more manageable.

A method for designing a multi-GPU kernel

  1. Measure how much runtime is compute, local memory, and communication.
  2. Identify the message sizes and whether the operation needs in-network arithmetic.
  3. Compare copy-engine, TMA, and register-driven transfer behavior.
  4. Decide whether communication aligns with compute inside an SM or deserves separate SMs.
  5. Check whether remote data should be prefetched into local HBM.
  6. Remove unnecessary intermediate buffers and synchronization.
  7. Test correctness and speed across the full shape distribution.

The deeper lesson

ParallelKittens does not make networking disappear. It gives programmers a vocabulary for deciding how networking participates in a kernel. The lasting contribution of Stuart’s talk is the three-axis mental model: transfer mechanism, overlap schedule, and design overhead. When a multi-GPU workload stalls, those axes provide concrete questions instead of a vague instruction to “optimize communication.”

Sources and further reading

← prev: Specialized Chipsnext: Intelligence / Watt →
© cvam — written in plaintext, served warm