PidokuInfra
07CUDA Graphs
On this page

CUDA Graphs

Intermediate Advanced 1h Difficulty 3/5 Topic 07 of 12

Prerequisites 05, 06, IV.08

1. What is it?#

Recording a sequence of GPU operations once, then replaying the whole sequence with a single launch.

WITHOUT:  launch k1, launch k2, ..., launch k350      350 × ~5 µs CPU = 1.75 ms
WITH:     graph.replay()                              1 × ~5 µs CPU = 5 µs

For LLM decode at small batch, this is typically a 20-45% end-to-end speedup for a config flag.

2. Why does it exist?#

Because the launch path (file 05) costs microseconds per kernel, and decode issues hundreds of kernels per token, each doing microseconds of work. The CPU becomes the critical path.

A graph moves all the per-launch work to a one-time setup, leaving only “execute this recorded DAG.”

3. Simple analogy#

A pre-programmed CNC machine versus manual operation.

Manual: for each of 350 operations, the operator reads the instruction, positions the tool, executes. The reading and positioning takes as long as the cutting.

Programmed: load the program once, press start. The machine executes all 350 operations back-to-back with no operator involvement.

The constraint is the same in both cases: the program is fixed. You can change the stock (input data) but not the operations (shapes, addresses).

4. Tiny example#

Python
import torch, time

model = ...  # your model
device = 'cuda'

# --- 1. Static input/output buffers (addresses are baked into the graph) ---
static_input = torch.zeros(1, 1, dtype=torch.long, device=device)
static_cache = ...   # pre-allocated KV cache

# --- 2. Warmup on a side stream (REQUIRED before capture) ---
s = torch.cuda.Stream()
s.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s):
    for _ in range(3):
        static_output = model(static_input, static_cache)
torch.cuda.current_stream().wait_stream(s)

# --- 3. Capture ---
g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
    static_output = model(static_input, static_cache)

# --- 4. Replay: copy new data INTO the static buffer, then replay ---
def step(token_id):
    static_input.copy_(token_id)      # write into the SAME memory
    g.replay()                        # one launch
    return static_output.clone()      # read the result before the next replay

Three rules visible here:

  1. Warmup before capture (allocates workspaces, resolves lazy initialization).
  2. Inputs and outputs live in fixed buffers; you copy into and out of them.
  3. Replay does not take arguments — the addresses are baked in.

Measure it:

Python
def bench(fn, n=200):
    for _ in range(10): fn()
    torch.cuda.synchronize(); t0 = time.perf_counter()
    for _ in range(n): fn()
    torch.cuda.synchronize()
    return (time.perf_counter()-t0)/n

print(f"eager: {bench(lambda: model(static_input, static_cache))*1e3:.2f} ms")
print(f"graph: {bench(g.replay)*1e3:.2f} ms")

5. Technical explanation#

What gets captured#

✓ Kernel launches
✓ Memory copies (device-to-device, and host copies with fixed addresses)
✓ Memsets
✓ Dependencies between them (inferred from stream ordering)

✗ Host code (any Python between kernels is NOT captured)
✗ Dynamic memory allocation (cudaMalloc during capture fails, unless using
    the graph memory pool)
✗ Anything shape-dependent decided on the host
✗ Synchronization (cudaStreamSynchronize during capture is an error)

The constraints, and how engines work around them#

CONSTRAINT 1: static shapes
  Decode has shape (batch, 1, d). Batch varies.
  → Capture graphs for a set of batch sizes: {1,2,4,8,16,24,32,48,64,...}
    and pad the actual batch up to the nearest captured size.

CONSTRAINT 2: static addresses
  All tensors must live at fixed addresses.
  → Pre-allocate everything: input buffer, KV cache pool, output buffer,
    block tables, sampling metadata. Copy new values in each step.

CONSTRAINT 3: no host control flow
  "if this sequence finished, stop" is a host decision.
  → Make control flow data-driven on the GPU, or handle it outside the graph.

CONSTRAINT 4: variable context length
  Attention reads a variable number of KV blocks.
  → The block table and sequence lengths are DATA in fixed-size buffers.
    The kernel loops based on the data. No recapture needed.

That last one is elegant and worth noting: paged attention makes context length a runtime value rather than a shape, which is exactly what makes graph capture feasible for LLM decode. PagedAttention solved a compilation problem as a side effect of solving a memory problem.

The padding cost#

Captured sizes: {1, 2, 4, 8, 16, 24, 32, 48, 64, 96, 128, 192, 256}
Actual batch 37 → run as 48

Wasted work: the padded rows do real computation on garbage.
  GEMM: 48 rows instead of 37 → 30% more GEMM work
        (but GEMMs round to tile boundaries anyway, so often ~0 real cost)
  Attention: 48 sequences' KV read instead of 37 → 30% more KV bandwidth  ← REAL
  
Net: roughly (padded/actual - 1) × KV_fraction_of_step_time

With a good capture list the average padding is 5-15%, and the launch-overhead saving is 25-45%. Net win, clearly — but tune the list if your batch distribution is unusual.

Memory cost#

Each captured graph holds:
  - its own static buffers for intermediates (via the graph memory pool)
  - the graph structure itself

For a 70B model with 15 captured batch sizes: 1-3 GB total.
That's KV cache you don't have — factor it into capacity planning.

Engines let you limit this (--cuda-graph-max-bs and similar) or disable it (--enforce-eager).

Graph update#

Python
# Instead of recapturing, you can update parameters of an instantiated graph
g.replay()   # PyTorch handles this; raw CUDA has cudaGraphExecKernelNodeSetParams

Raw CUDA allows updating kernel parameters in an instantiated graph without rebuilding it — useful for changing a pointer or a scalar. PyTorch’s abstraction mostly hides this.

6. Under the hood#

What instantiation does:

1. Topologically sort the recorded nodes
2. Resolve all dependencies into a hardware-executable form
3. Pre-allocate any needed internal state
4. Produce a cudaGraphExec_t

Replay then just tells the GPU “execute this pre-built structure.” The GPU’s front-end walks the DAG directly, with no host round trip per node — which is why the GPU-side gaps shrink too, not just the CPU time.

7. Performance implications#

Configuration                  Speedup from graphs
70B, batch 1                   1.35-1.45x
70B, batch 8                   1.20-1.30x
70B, batch 64                  1.08-1.15x
70B, batch 256                 1.02-1.06x
7B, batch 1                    1.40-1.60x   ← more kernels per unit of work
Prefill                        ~1.00x (not captured)

The benefit is inversely proportional to work per kernel. Small models and small batches benefit most.

8. Production implications#

  • Leave graphs enabled. They’re on by default in vLLM and TensorRT-LLM (as “in-flight batching with graph capture”).
  • Budget the startup time and memory. Capture takes 20-60 seconds and 1-3 GB.
  • Tune the capture list to your batch distribution if it’s unusual.
  • --enforce-eager for debugging only. Graphs make stack traces less informative; disable temporarily when chasing a bug, then re-enable.
  • Check feature compatibility. Multi-LoRA, some speculative decoding modes, and certain sampling features can be incompatible with capture in some engine versions. If enabling a feature silently disables graphs, you’ve just lost 25%.
  • Watch for capture failures at startup. They’re often logged as warnings and fall back to eager — a silent regression.

9. Common mistakes#

Skipping the warmup before capture. Capture fails or records lazy-initialization work.

Allocating during capture. Fails. Pre-allocate everything.

Reading the output after the next replay. The output buffer is overwritten. Clone or consume before replaying.

Capturing too many batch sizes. Long startup, high memory.

Capturing too few. Excessive padding.

Leaving --enforce-eager on in production.

Assuming graphs help prefill. They don’t.

10. Hands-on exercise#

A. Capture a decode step. Take a small model, pre-allocate a KV cache, and capture one decode step. Verify the output matches eager execution. Measure both.

B. Batch size sweep. Capture graphs for batch sizes {1,2,4,8,16,32,64}. Measure the speedup at each. Plot speedup vs batch size. Confirm the trend in section 7.

C. Padding cost. For each captured size, measure the time. Then compute the expected cost of running batch 37 as batch 48. Verify empirically.

D. Memory cost. Measure GPU memory before and after capturing 15 graph sizes for a real model. How much KV cache did you give up?

E. Break it. Try to capture a region containing (i) a .item() call, (ii) a cudaMalloc, (iii) a data-dependent branch. Observe the failure modes.

11. Interview questions#

  1. What are CUDA graphs and what problem do they solve?
  2. What are the constraints on graph capture?
  3. How do engines handle variable batch size with graphs? What does that cost?
  4. Why does paged attention make graph capture feasible for LLM decode?
  5. Why do graphs help small batch more than large batch?
  6. What is the memory and startup cost of graph capture?
  7. Name three things that would make graph capture fail.

12. Further reading#

  • [REFERENCE] CUDA C++ Programming Guide, “CUDA Graphs”
  • [REFERENCE] PyTorch torch.cuda.CUDAGraph and make_graphed_callables
  • [REFERENCE] vLLM’s graph capture implementation (model_runner.py)
  • Next: 08 — CUDA programming fundamentals
esc

↑↓ navigate↵ openesc close