BetterSWE DigestWhat is BetterSWE?

Deep Dive · Issue 24

Why CUDA Graphs Break on Every LLM Token

CUDA Graphs freeze GPU kernels for instant replay. Every LLM token changes the shapes the graph was built for. The conflict is built in.

GPU receiving kernel launches from CPU with a shape mismatch blocking one launch

The setup

A Replay That Can't Keep Up

GPU receiving kernel launches from CPU with a shape mismatch blocking one launch

CUDA Graphs record a sequence of GPU kernel launches once and replay the entire batch in a single CPU call. For LLM serving, where every microsecond of CPU overhead between kernels shows up as added per-token latency, this is a meaningful win.

The catch is that a captured graph is a frozen recording. Every tensor shape, stride, and kernel argument is fixed at capture time, but each LLM decoding step extends the sequence by one position, changing the attention mask shape and KV cache length. A graph captured for sequence length 128 cannot replay at 129.

One token later, the captured graph is already stale.

The mechanism

What Capture Actually Freezes

A captured graph as a sealed recording with a shape mismatch crack

During capture, PyTorch runs the forward pass once while CUDA records every kernel, its arguments, and the memory addresses of every input and output tensor. The result is a graph object you can launch with a single cudaGraphLaunch call. The CPU does no Python evaluation, no shape computation, no kernel selection on replay. It fires the recording.

Pointers, not shapes

What gets frozen is more than the kernel list. The graph stores pointer values to specific memory addresses. If a tensor is reallocated because its shape changed, the graph still points at the old address. The graph also encodes kernel launch configurations like grid and block dimensions, which are often derived from tensor shapes at capture time. Change the shape and the kernel configuration is wrong.

Captured graph

fixed shapes, fixed addresses, one replay call

Live execution

shapes computed per step, kernels selected at runtime

Why you can't just update the pointers

You might assume the fix is simple: keep the same kernels, just update the tensor pointers. Some graph APIs do support pointer updates via cudaGraphExecKernelNodeSetParams. But the problem goes deeper than addresses. Kernel grid dimensions, shared memory sizes, and even the selection of which kernel to call can depend on the input shape. A short sequence might use a fused attention kernel; a long one might need a different tiling strategy. The graph captures one code path, and the right path changes with the shape.

The cost

Why the CPU Is the Bottleneck

CPU-GPU timeline with launch overhead gaps eliminated by graph replay

Without graphs, every kernel launch costs CPU time. Each cudaLaunchKernel call takes roughly 3 to 10 microseconds of CPU overhead, and an LLM decoding step for a single token can involve hundreds of kernels. At 200 kernels and 5 μs each, that is 1 ms of pure CPU overhead per token, spent on bookkeeping the GPU does not need.

toverhead=nkernels×tlaunch=200×5 μs=1 ms per tokent_{\text{overhead}} = n_{\text{kernels}} \times t_{\text{launch}} = 200 \times 5\,\mu\text{s} = 1\,\text{ms per token}
CPU launch overhead per decoding step at typical kernel counts

For small models or short sequences, the GPU finishes the actual computation in a few milliseconds, so 1 ms of launch overhead is a large fraction of total latency. CUDA Graph replay collapses that overhead to a single CPU call, which is why the technique has become standard in production LLM serving.

A graph is a recording, and a recording can't improvise.

The workarounds

Bucketing, Padding, and Re-Capture

Shape buckets with padding and a JIT compiler adding new graphs to a cache

The most common workaround is shape bucketing. You capture several graphs, each covering a range of sequence lengths, and pick the right one at inference time. A request with sequence length 130 rounds up to the 128-256 bucket, and the GPU pads the gap.

Bucketing trades wasted GPU work for graph reuse.

A newer approach combines JIT compilation with graph capture. When a shape falls outside existing buckets, a JIT compiler generates and captures a new graph on the fly and caches it for reuse, avoiding both padding waste and re-capture cost.

PyGraph goes further by making PyTorch itself robust to dynamic patterns, handling control flow, shape changes, and memory reuse inside the graph so the framework can capture more without manual bucketing.

Wrap-up

The Tension That Won't Go Away

Why it matters

The tension between static capture and dynamic shapes is the fundamental tradeoff of recording GPU work for replay: you gain a CPU-free execution path and lose the ability to adapt. Every LLM serving framework that uses graphs pays this tax through padding waste, re-capture latency, or the engineering cost of maintaining a graph cache. Hybrid JIT approaches now appearing in research suggest the field is converging on a compromise: capture what you can, compile what you can't, cache everything.

The short version

  • CUDA Graphs record kernel launches for one-shot replay, eliminating per-kernel CPU overhead that can reach **1 ms per token**.
  • Every LLM decoding step changes sequence length, attention mask shape, and KV cache size, invalidating the captured graph.
  • Frozen memory addresses and kernel configurations mean even a one-position shape change breaks replay.
  • Shape bucketing with padding is the common workaround; hybrid JIT compilation and capture is the emerging fix.
  • The conflict is structural: graphs want static shapes, and LLM inference is dynamic by definition.

Sources

  1. 1.Hybrid JIT–CUDA Graph Optimization for Low-Latency Large Language Model Inference
  2. 2.CUDA Graphs and Dynamic Shapes: Performance Optimization for Variable-Length Sequences | ClusterBid · ClusterBid
  3. 5.Speed, Python: Pick Two. How CUDA Graphs Enable Fast Python Code for Deep Learning
  4. 6.Handling Dynamic Patterns — CUDA Graph Best Practice for PyTorch
  5. 8.PyTorch CUDA Graph Capture - Lei Mao's Log Book
  6. 9.PyGraph: Robust Compiler Support for CUDA Graphs in PyTorch
  7. 11.CUDA Graphs in LLM Inference: Deep Dive - DEV Community