CUDA Graphs and torch.compile for LLM Decode Latency: Killing Launch Overhead (2026)

CUDA Graphs and torch.compile for LLM Decode Latency: Killing Launch Overhead (2026)

CUDA Graphs and torch.compile for LLM Decode Latency: Killing Launch Overhead (2026)

A single decode step of a 7B to 8B parameter model on a modern data-center GPU should take a few milliseconds, because the hardware only has to stream the weights out of high-bandwidth memory once. In practice, a plain eager PyTorch loop often takes several times longer than that floor at batch size one. The GPU is not the problem. It is waiting for the CPU to tell it what to do next, one small kernel at a time.

CUDA graphs fix this by recording a whole sequence of kernels once and replaying it with a single launch call. Combined with torch.compile, which fuses many tiny operations into fewer kernels, the pair is the standard answer to LLM decode latency in 2026. It is also the reason vLLM, SGLang, TensorRT-LLM and gpt-fast all pad batches to a few fixed sizes and keep key-value (KV) cache tensors at fixed addresses.

The technique is simple to describe and surprisingly easy to get wrong. A graph captures pointers, not values. Anything dynamic, from variable batch sizes to growing caches to host-side control flow, either breaks the capture or silently multiplies memory use.

What this covers: why decode is launch-bound, how stream capture and graph replay work, what torch.compile adds on top (fusion and CUDAGraph Trees), how vLLM structures full and piecewise graphs, runnable code for manual and compiled capture, the failure modes that bite in production, and a decision checklist.

Context and Background

Autoregressive decoding generates one token per forward pass. Each pass reads every weight matrix once and the KV cache for every active sequence once, and does very little arithmetic per byte read. That makes decode memory-bandwidth-bound at small batch sizes. A back-of-envelope floor follows directly: an 8B-parameter model in bf16 is about 16 GB of weights; an H100 SXM exposes roughly 3.35 TB/s of HBM3 bandwidth according to NVIDIA’s published specification; so streaming the weights costs about 16 / 3350 seconds, or roughly 4.8 ms per token, before counting KV cache reads. Real systems reach a fraction of peak bandwidth, so treat 4.8 ms as a lower bound, not a prediction.

The surprise is that eager PyTorch frequently does not get near that bound at batch size one. A transformer layer in eager mode issues many separate operations: RMSNorm pieces, QKV projections, rotary embedding, attention, output projection, residual adds, the MLP’s gate, up and down projections, and an activation. Depending on implementation that is on the order of ten to thirty kernels per layer. For a 32-layer model, a single token can involve several hundred kernel launches. These counts are approximate and implementation-dependent; the reliable way to know yours is a profiler trace.

Each launch costs CPU time. The Python interpreter walks the module tree, the PyTorch dispatcher resolves the operator, allocates outputs, selects a kernel, and the CUDA driver enqueues the work. The exact cost varies by hardware, PyTorch version and operator, so I will not quote a single number; the mechanism is what matters. When the GPU finishes a kernel faster than the CPU can enqueue the next one, the stream runs dry and the device idles. Decode kernels at small batch are short, so this happens constantly.

This is not a new observation. The PyTorch team’s gpt-fast write-up, “Accelerating Generative AI with PyTorch II,” describes exactly this regime. Its baseline of 25.5 tokens per second for a Llama-style model at batch size one on an A100-80GB was “heavily CPU overhead bound.” Wrapping the decode function in torch.compile(..., mode="reduce-overhead", fullgraph=True) with a static KV cache lifted it to 107.0 tokens per second, a gain of more than 4x from overhead removal alone, before any quantization. Those are the authors’ reported numbers on their hardware; yours will differ, but the shape of the result is typical.

Two complementary tools attack the problem from different directions. CUDA graphs reduce the per-kernel CPU cost to nearly zero by moving launch work to a one-time capture. Kernel fusion, via torch.compile‘s Inductor backend and hand-written Triton or CUDA kernels, reduces the number of kernels and the number of round trips to HBM. For deeper background on the fusion side, see our walk-through of Triton GPU kernel programming and fused kernels, and for how attention itself avoids HBM traffic, see FlashAttention’s IO-aware tiling.

The engineering context in 2026 is that serving engines have absorbed most of this. vLLM, per its design documentation, treats CUDA graph capture as a first-class compilation mode with dispatch keys, capture sizes and per-backend support levels. You rarely hand-roll capture for production serving. But understanding the mechanism is what lets you diagnose a regression when a graph is skipped, when memory balloons after enabling capture, or when a custom kernel refuses to be captured. For how the engines compare end to end, see our inference benchmark of vLLM, TGI, SGLang and Triton. For the official reference on the CUDA side, the CUDA Programming Guide section on graphs is the primary source used throughout this article.

Why Decode Is Launch-Bound: The Reference Model

The direct answer: decode latency at small batch is the maximum of GPU time and CPU launch time, not their sum, because the two overlap through an asynchronous queue. When per-kernel GPU time is shorter than per-kernel CPU enqueue time, the CPU is the bottleneck and the GPU waits. CUDA graphs collapse hundreds of enqueues into one, making the GPU the bottleneck again.

CUDA graphs and kernel launch overhead in an eager LLM decode loop

Figure 1: The eager launch path. Every operator pays Python, dispatcher and driver cost; when the launch queue empties, the GPU idles.

Figure 1 shows the loop. CPU-side work is asynchronous with respect to the GPU: cudaLaunchKernel returns as soon as the work is queued, so the CPU can run ahead as long as it stays ahead. The condition for a healthy pipeline is simple. If the average CPU cost to enqueue one kernel is c and the average GPU duration of that kernel is g, the GPU stays busy only while c ≤ g on average. For large prefill matrix multiplications, g is hundreds of microseconds to milliseconds and the condition holds trivially. For decode at batch one, many kernels (norms, elementwise adds, small projections) have g in the single-digit-microsecond range, and c is larger.

The arithmetic of a launch-bound step

Consider an illustrative model, with all figures chosen for the sake of the arithmetic and not measured: 600 kernels per token, an average GPU duration of 8 microseconds per kernel, and an average CPU enqueue cost of 20 microseconds per kernel in eager mode. The GPU-only time is 600 x 8 = 4.8 ms. The CPU-only time is 600 x 20 = 12 ms. Because the CPU is the slower side, the step takes roughly 12 ms and the GPU sits idle for about 60 percent of it.

Now apply a graph. A graph replay costs one launch call plus the GPU’s own internal dispatch of graph nodes, which is much cheaper per node than a host-side launch because the driver has already validated and prepared the work. The step approaches the GPU-only figure. The speedup in this toy example is 12 / 4.8 = 2.5x. The real gain depends entirely on your ratio of c to g; the 4.2x observed in gpt-fast on an A100 means that setup was even more CPU-bound than my toy numbers.

This also explains why the benefit shrinks with bigger models and bigger batches. A 70B model sharded across eight GPUs has larger per-kernel GPU durations, so g rises toward c. A batch of 256 sequences makes every GEMM bigger. In both cases launch overhead becomes a smaller fraction of the step, and CUDA graphs matter less for throughput. They still matter for tail latency and for the minimum time-to-next-token, which is what interactive users feel.

Why overhead is a latency problem, not a throughput problem

Throughput-oriented batch serving can hide launch overhead by making kernels longer. Latency-oriented serving cannot. A chat assistant generating a single stream, a speculative-decoding draft model producing five tokens in a row, or a voice agent with a hard per-token deadline all operate in the regime where g is small. Speculative decoding is especially exposed: the draft model is small by design, so its kernels are short, and it runs sequentially several times per verification step. Graph capture of the draft loop is one of the highest-return optimizations in that architecture.

The CPU can also be the noisy neighbor

Launch overhead is not constant. It depends on interpreter speed, how many Python threads contend for the Global Interpreter Lock, CPU frequency scaling, and whether the process shares cores with a tokenizer, an HTTP server, or a metrics exporter. A serving process that looks fine in an isolated benchmark can show erratic decode latency under production load simply because the launching thread is descheduled. Graph replay is robust to this because the launching thread does almost nothing per step. That is a latency-variance benefit, not only a mean-latency benefit, and it is one reason p99 token latency often improves more than the mean.

How CUDA Graphs Work Under the Hood

A CUDA graph separates work submission into three stages, as NVIDIA’s programming guide describes: definition, instantiation, and execution. During definition you build a directed acyclic graph of nodes (kernels, memory copies, memsets, host callbacks, child graphs, and more) with dependency edges. Instantiation validates and prepares the graph, producing an executable graph. Execution launches it into a stream. The guide states the rationale directly: the per-launch setup overheads “can be paid once for the entire graph during instantiation.”

CUDA graphs lifecycle: warmup, capture, instantiate, replay, and tree branching

Figure 2: The lifecycle of a captured graph in PyTorch. Warmup, record, instantiate, then replay until a shape or address mismatch forces a new branch.

Stream capture

You rarely build graph nodes by hand. Instead you use stream capture: cudaStreamBeginCapture() puts a stream in capture mode, ordinary kernel launches into that stream are recorded as nodes instead of executed, and cudaStreamEndCapture() returns the graph. Cross-stream work is joined using cudaEventRecord() and cudaStreamWaitEvent(), and all forked streams must rejoin the origin stream before capture ends.

Capture imposes restrictions that explain most of the failures you will meet in practice. According to the programming guide:

  • The legacy default stream (cudaStreamLegacy) cannot be captured, though cudaStreamPerThread can.
  • Synchronizing or querying a capturing stream or a captured event is invalid.
  • Synchronous calls such as cudaMemcpy() are invalid while another stream in the context is capturing (unless that stream was created non-blocking).
  • Some APIs, such as cudaStreamAttachMemAsync(), are unsupported during capture.
  • Any invalid operation invalidates the capture, and cudaStreamEndCapture() then returns an error and a NULL graph.

The practical consequence for LLM code is blunt: no host-device synchronization inside the captured region. A .item() call, a tensor.cpu(), a data-dependent Python if on a GPU value, or torch.nonzero with its data-dependent output size will either fail capture or force the graph to break around it.

What a graph freezes

A captured kernel node stores the kernel function, its launch dimensions, and its argument values. Argument values include device pointers. On replay, the kernels read and write exactly the addresses recorded at capture time. This single fact generates the whole set of rules that follow. If you want a graph to process new data, you must write that data into the same memory the graph already points at. If a tensor was allocated at capture time, it must still be live at the same address at replay time, which is why PyTorch gives captured graphs a private memory pool.

Scalars are frozen too. If a kernel received a Python integer such as the current sequence length as a launch argument, the captured value is baked in. This is why naive growing-KV-cache code cannot be captured: the attention kernel’s sequence length argument and the cache tensor shape both change every step. The standard fix is a static cache preallocated to the maximum length with a mask or a device-side length tensor, so that the changing information lives in memory, not in kernel arguments.

Updating without re-recording

Re-capturing is not the only way to change a graph. The programming guide lists cudaGraphExecUpdate() for swapping in a topologically identical graph with the same dependency order, and per-node setters such as cudaGraphExecKernelNodeSetParams() for changing kernel arguments in an instantiated graph. There are also enable and disable calls, cudaGraphNodeSetEnabled() and cudaGraphNodeGetEnabled(), for kernel, memset and memcpy nodes. Topology changes still require re-instantiation. Inference engines use these update paths sparingly, since the common approach is simply to capture one graph per batch-size bucket and write inputs into static buffers.

Conditional nodes and device-side launch

The guide also documents more advanced features. Conditional nodes (IF with an optional else, WHILE, and SWITCH) allow control flow inside a graph, with the condition value set on the device through cudaGraphSetConditional() using handles created by cudaGraphConditionalHandleCreate(). This matters for LLM decode because the stop condition (“did we emit end-of-sequence?”) is device-resident data. In principle a WHILE node could run an entire decode loop without returning to the host. Device graph launch, enabled with cudaGraphInstantiateFlagDeviceLaunch, lets a kernel launch another graph, with documented modes such as fire-and-forget (limited to 120 launches per graph execution) and tail launch. Device-launchable graphs may contain only kernel, memcpy, memset and child-graph nodes. I did not verify which CUDA release introduced each feature from the page I read, so check the release notes for your toolkit before depending on them. Mainstream PyTorch inference does not use these yet; they are worth tracking as a way to take the host out of the loop entirely.

Memory cost of graphs

Graphs are not free in memory. The PyTorch documentation on CUDAGraph Trees notes that each kernel launch captured in a graph costs about 64 KB of device memory (stated for CUDA 12.4 and driver 550 or newer). A model with a few thousand captured kernels therefore consumes a few hundred megabytes per graph. Multiply by the number of batch-size buckets and the per-bucket cost becomes a real line item in your GPU memory budget. That is the reason serving engines capture a limited list of sizes and not every possible batch size.

What torch.compile Adds: Fusion and CUDAGraph Trees

torch.compile contributes two distinct things that are easy to conflate. The first is compilation: TorchDynamo traces Python into an FX graph, and the Inductor backend lowers it, fusing elementwise and reduction operations into generated Triton kernels and choosing between template kernels and vendor libraries for matrix multiplications. The second, enabled by mode="reduce-overhead", is CUDA graph management through CUDAGraph Trees.

torch.compile pipeline from Dynamo trace to Inductor fusion to a CUDA graph wrapper

Figure 3: The torch.compile pipeline for decode. Fusion reduces kernel count first; the CUDAGraph wrapper then removes the remaining launch cost.

Fusion shrinks the number of launches and bytes moved

Fusion helps launch overhead in two ways. It reduces the kernel count, so even an eager-style launch path has fewer launches to pay for. It also eliminates intermediate tensors, because an elementwise chain such as residual add, normalization scale and activation can run in a single pass over registers instead of writing and re-reading each intermediate from HBM. In a memory-bound regime, avoiding those intermediate round trips is a direct latency win, not just a launch-count win.

The gpt-fast article makes a second claim worth noting: torch.compile generates memory-bandwidth-bound matrix-vector and attention kernels that the authors report outperform cuBLAS and FlashAttention-2 in that decode regime. The post gives no figures for that comparison, so treat it as the authors’ observation, not a benchmark you can cite as a number. The point stands qualitatively: at batch size one a matvec is not the compute-bound GEMM that vendor libraries are most heavily tuned for.

The compile modes

The documented torch.compile modes relevant here are the default, reduce-overhead, max-autotune, and max-autotune-no-cudagraphs. As I read the PyTorch docs, reduce-overhead enables CUDA graphs to cut Python overhead at the price of some extra memory, and max-autotune additionally searches more kernel configurations (it also uses CUDA graphs on GPUs), with the no-cudagraphs variant available when you want autotuning without capture. I could not re-fetch the mode table for this revision, so confirm the exact semantics on the torch.compile page for your installed version; the documentation moves between releases.

How CUDAGraph Trees work

The PyTorch documentation describes CUDAGraph Trees as extending graph capture to torch.compile by recording a separate graph for each compiled region and each execution path. The recordings form a tree. If a function branches after a graph break, each branch records its own path, and the shared prefix is reused.

The lifecycle matches Figure 2:

  1. The first call warms up the graph, which runs Triton and cuBLAS autotuning without capture.
  2. The second call records the graph.
  3. Subsequent calls replay the recording.

All recordings share one private memory pool, so, according to the docs, there is no extra memory overhead compared with eager execution from the pool itself. When execution takes a new path, the pool is checkpointed to its state at the end of the shared prefix so that live tensors are not overwritten. This is what distinguishes trees from the older non-tree mode, which does not share memory between forward and backward graphs.

Static inputs: what gets copied

The docs state the input rules precisely. Inputs that are fresh eager tensors can change address between calls, so they are copied into static buffers before recording and replay. Parameters and buffers are assumed to keep the same addresses, and that assumption is checked at runtime. Outputs of other CUDAGraph Tree functions already live at fixed addresses and can feed later graphs without a copy; if that stability is violated, the system re-records. The key phrase is that values may change but addresses may not.

The consequence for decode: the per-step copy of the input token ids into a static buffer is cheap, but any large tensor passed fresh on every call (a big KV cache created outside the compiled function, for instance) would be copied each step. Keep large state as module buffers or static preallocated tensors that live at stable addresses.

Graph breaks multiply recordings

fullgraph=True tells Dynamo to raise an error on any graph break instead of silently splitting the function. For decode this is worth turning on during development. Every break means a separate compiled region, a separate recording, and Python executing between them, which reintroduces the overhead you set out to remove. The gpt-fast article uses fullgraph=True for exactly this reason: to verify there are no breaks.

Walk-through: Three Ways to Capture a Decode Step

The same idea appears at three levels of abstraction. Hand-written capture with torch.cuda.CUDAGraph shows the mechanism naked. torch.compile(mode="reduce-overhead") automates it. Serving engines such as vLLM wrap it in a dispatcher. Reading them in this order makes the engine-level behavior much easier to predict.

Level 1: manual capture with static buffers

The code below is a minimal, runnable pattern for a toy single-layer decode step. It follows the documented PyTorch API: warm up on a side stream, then capture with torch.cuda.graph, then replay by copying new data into static tensors. The model here is a stand-in; the same structure applies to a full transformer whose KV cache is preallocated.

import torch, time

torch.manual_seed(0)
dev = "cuda"
D, V = 4096, 32000

class TinyDecoder(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.emb = torch.nn.Embedding(V, D)
        self.ln = torch.nn.LayerNorm(D)
        self.mlp = torch.nn.Sequential(
            torch.nn.Linear(D, 4 * D), torch.nn.GELU(),
            torch.nn.Linear(4 * D, D))
        self.head = torch.nn.Linear(D, V, bias=False)

    def forward(self, tok):                       # tok: [B] int64
        h = self.emb(tok)
        h = h + self.mlp(self.ln(h))
        return self.head(h).argmax(dim=-1)        # next token ids, stays on GPU

model = TinyDecoder().to(dev).half().eval()
B = 1
static_tok = torch.zeros(B, dtype=torch.long, device=dev)

# 1) warmup on a side stream (required before capture)
s = torch.cuda.Stream()
s.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s), torch.no_grad():
    for _ in range(3):
        model(static_tok)
torch.cuda.current_stream().wait_stream(s)

# 2) capture
g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g), torch.no_grad():
    static_out = model(static_tok)

# 3) replay: write inputs in place, replay, read outputs
def step_graph(tok):
    static_tok.copy_(tok)
    g.replay()
    return static_out            # NOTE: overwritten on next replay; clone if kept

def step_eager(tok):
    with torch.no_grad():
        return model(tok)

def bench(fn, n=500):
    tok = torch.randint(0, V, (B,), device=dev)
    for _ in range(20): fn(tok)
    torch.cuda.synchronize(); t = time.perf_counter()
    for _ in range(n): fn(tok)
    torch.cuda.synchronize()
    return (time.perf_counter() - t) / n * 1e3

print(f"eager: {bench(step_eager):.3f} ms/step")
print(f"graph: {bench(step_graph):.3f} ms/step")

Three details deserve attention. The argmax stays on the GPU, so nothing synchronizes inside the region. The output tensor is static, which means a second replay overwrites it; clone it if you need to keep the previous token. And the benchmark times many steps between synchronizations, because timing a single asynchronous launch only measures how fast the CPU enqueued work. The numbers you get depend on your GPU and CPU, so I have not printed any; run it and compare. A single-layer toy will show a modest gain because there are few kernels; the gap grows with layer count.

Level 2: torch.compile with a static KV cache

For a real transformer, the pattern from gpt-fast is the reference. The structure is a static KV cache of maximum length, a position tensor that lives on the device, and a compiled decode function with mode="reduce-overhead". The snippet below shows the shape of the code, not a complete model.

import torch

# Pseudo-structure of a gpt-fast style decode step.
# model.setup_caches allocates K/V of shape [B, heads, MAX_LEN, head_dim] once.
model.setup_caches(max_batch_size=1, max_seq_length=MAX_LEN)

@torch.no_grad()
def decode_one_token(model, tok, input_pos):
    # tok: [1, 1] token ids, input_pos: [1] int tensor ON DEVICE
    logits = model(tok, input_pos)                # writes K/V at input_pos
    return logits[:, -1].argmax(dim=-1, keepdim=True)

# Prefill has variable prompt length -> keep it dynamic / uncaptured.
# Decode shapes are fixed -> compile with CUDA graphs.
decode_one_token = torch.compile(
    decode_one_token, mode="reduce-overhead", fullgraph=True)

tok = prompt_last_token                           # [1, 1]
pos = torch.tensor([prompt_len], device="cuda")
out = []
for _ in range(max_new_tokens):
    tok = decode_one_token(model, tok, pos).clone()   # clone: avoid stale output
    out.append(tok)
    pos += 1                                      # in-place on a device tensor

The .clone() is not decoration. Per the PyTorch docs, outputs from a CUDAGraph Tree can be overwritten by the next invocation, and accessing a stale one raises a RuntimeError. You can also call torch.compiler.cudagraph_mark_step_begin() to mark iteration boundaries, or configure torch._inductor.config.triton.cudagraph_trees_generation_cloning = "user_visible" to clone user-visible outputs automatically. Note that pos += 1 is an in-place update of a device tensor, so the compiled graph sees changing values at a stable address: values change, addresses do not.

The cache write inside the model is the subtle part. A static cache is updated with an index-put at input_pos and attention masks out the unused positions. That is wasteful in FLOPs (attention touches the masked region unless the kernel respects a length) but eliminates shape changes. This is the trade the gpt-fast authors chose, and they explicitly note that a growing cache defeats graph capture.

Hugging Face Transformers

Recent Transformers releases expose the same idea through a static cache option and compilation of the forward pass. The exact flag names and defaults have changed across releases, so I am not reproducing a snippet I have not verified against your installed version; consult the library’s generation documentation for “static cache” and torch.compile. The principle is unchanged: preallocate the cache, keep shapes fixed, and compile the decode forward.

Level 3: engine-level dispatch in vLLM

Production engines face a problem the single-stream examples do not: batch size and batch composition change every step. A graph captured for 17 sequences cannot serve 18. The standard answer is bucketing: capture graphs for a chosen list of sizes, pad the live batch up to the nearest captured size, replay, and discard the padded outputs.

vLLM-style CUDA graph dispatch with full graphs for decode and piecewise graphs for mixed batches

Figure 4: Dispatch in a serving engine. Uniform decode batches replay a full graph at a padded size; other batches use piecewise graphs with eager attention between pieces.

According to vLLM’s design documentation, the dispatch key is a BatchDescriptor holding num_tokens, num_reqs, a uniform flag and has_lora. A CudagraphDispatcher is the single source of truth for which keys have captured graphs, searching FULL first, then PIECEWISE, and falling back to NONE (eager) if nothing matches. A CUDAGraphWrapper captures or replays only when the runtime mode matches its own; otherwise it calls the model directly. Capture sizes come from cudagraph_capture_sizes in the compilation config, and the num_tokens in a key may be the padded length. I did not find the default capture size list or exact padding policy on the page I read, so check your version’s configuration reference.

The documented modes are:

Mode What is captured Typical fit
NONE Nothing; eager Debugging
PIECEWISE Everything except attention (and other incompatible ops); attention runs eagerly between pieces Maximum backend compatibility
FULL Full graphs for non-uniform batches; uniform decode reuses the same-size graph Small models, short prompts
FULL_DECODE_ONLY Full graphs for uniform decode only; prefill and mixed batches run eagerly Decode instances in disaggregated prefill/decode deployments
FULL_AND_PIECEWISE Full graphs for uniform decode, piecewise for everything else Default with piecewise compilation; usually fastest, most memory, longest capture

Why is attention excluded from piecewise capture? The docs call attention the main CUDA-graph-unsupported operation: backends often depend on batch metadata (sequence lengths, block tables) that varies in ways a frozen kernel argument cannot express. Piecewise compilation captures the rest of the model and leaves attention eager, which is compatible with every backend. Full graphs came later for backends that can capture attention with fixed-shape metadata. Each attention backend declares a cudagraph_support level (ALWAYS, UNIFORM_BATCH, UNIFORM_SINGLE_TOKEN_DECODE or NEVER), and vLLM downgrades the requested mode to the closest supported one. Hybrid models use the minimum capability across their backends.

You set the mode through the compilation config, for example on the command line with --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}' or in Python through compilation_config={"mode": 3, "cudagraph_mode": "FULL_AND_PIECEWISE"}. An experimental flag, use_inductor_graph_partition=True, moves graph partitioning into Inductor instead of after Dynamo and requires torch>=2.9. Treat this area as fast-moving: flag names and defaults shift between vLLM releases.

The padding tax

Padding is a deliberate trade. Replaying a graph sized for 32 sequences when 20 are live wastes compute and memory bandwidth on 12 phantom rows. At decode, however, the dominant cost is reading weights, which does not change with batch size, so the marginal cost of padded rows is small. The quantity worth checking is the KV read and the extra sampling work for padded rows, which can matter at long context. A denser list of capture sizes at small batch (where waste is proportionally largest) and a sparser list at large batch is the usual compromise, and it matches how memory per graph scales. Whatever your engine’s default is, measure the share of steps that fall into each bucket in production before tuning.

Choosing and Measuring: A Decision Matrix

The question is rarely “graphs or no graphs.” It is which layer of the stack should own capture, and what you give up. The matrix below summarizes the judgement calls. It reflects mechanisms described above and my own engineering opinion where marked; it is not a benchmark.

Approach Launch overhead Flexibility Memory cost Best when
Eager PyTorch Highest Complete Lowest Prototyping, prefill-heavy, large kernels
Manual torch.cuda.CUDAGraph Near zero Low; you manage static buffers Per-graph Small custom loops, draft models, research
torch.compile default Reduced by fusion High Low Moderate gains without static-shape constraints
torch.compile reduce-overhead Near zero after warmup Medium; shape and address rules Pool plus per-graph Fixed-shape decode with static KV cache
Engine-managed (vLLM and similar) Near zero for captured sizes Engine handles buckets Per bucket and mode Production serving with variable batches

How to measure whether you are launch-bound

Before changing anything, confirm the diagnosis. Capture a profiler trace of several decode steps with torch.profiler and look at the GPU timeline. In a launch-bound step, you will see small kernels separated by visible gaps, and the CPU-side track will show dense dispatcher activity running continuously. In a bandwidth-bound step, kernels are back to back, and the CPU track finishes early then waits at a synchronization point.

A cheap second test: run the same decode step at batch size 1 and batch size 8. If latency per step barely changes, the work is dominated by weight reads or by fixed overhead; combined with gaps in the trace, that points to launch cost. If latency scales with batch size, you are already compute or bandwidth-bound on activations and graphs will help less.

A third check is GPU utilization over time. Tools like nvidia-smi report the fraction of time any kernel was running, which is a coarse signal: a device that is “100 percent utilized” can still be launch-bound if tiny kernels keep it nominally busy. Prefer trace-level evidence over utilization percentages.

Setting expectations

Reported results are a better guide than hopes. The gpt-fast case moved from 25.5 to 107.0 tokens per second at batch size one on an A100-80GB, with the authors attributing the jump to removing CPU overhead. The same article then reports 72 percent memory bandwidth utilization at 107 tokens per second, which tells you the graph-based version was already close to the memory-bandwidth ceiling. That is the sign of success: once launch overhead is gone, the next improvements come from moving fewer bytes (quantization, speculative decoding), and indeed the article’s later steps (int8 at 157.4, int4 with GPTQ at 202.1 tokens per second) are bandwidth optimizations.

The sequence matters as an architectural principle. Remove launch overhead first, because it hides every other optimization: if the GPU is idle half the time, halving the bytes per token changes little. Then reduce bytes. Then, if you have headroom, raise batch size to improve throughput. This ordering is my interpretation of the evidence, not a published rule, but it is consistent with the gpt-fast progression.

Trade-offs, Gotchas, and What Goes Wrong

Graphs are a sharp tool. The failure modes cluster around the frozen-pointer model.

Dynamic shapes cause re-recording, and re-recording costs memory. PyTorch’s documentation says a new recording is made for every unique input shape, and each kernel in a graph costs roughly 64 KB of device memory in the stated configuration. A serving path that feeds raw, unpadded shapes will accumulate recordings until memory pressure or slow warmup dominates. The recommended mitigation is to pad inputs to a small set of fixed shapes. Setting torch._inductor.config.triton.cudagraph_skip_dynamic_graphs=True skips graphs with dynamic shapes entirely, trading speed for predictability.

Silent skipping hides regressions. Graphs are skipped for conditions including CPU operations, multi-device operations (NCCL is supported), free unbacked symbols, custom ops tagged cudagraph_unsafe, and operators such as aten._local_scalar_dense, aten._assert_scalar and aten.multinomial.default. The last is relevant to sampling: a sampling step using multinomial inside the compiled region can cause the graph to be skipped, which is why many engines keep sampling outside the captured region or implement it with graph-safe primitives. When mutated eager inputs trigger a skip, PyTorch logs “skipping due to mutated inputs.” Treat that log line as a performance alarm and grep for it in CI.

Mutation of inputs. Mutating eager inputs disables capture. Mutating parameters, buffers or prior tree outputs is supported only when torch._inductor.config.cudagraph_support_input_mutation = True. KV cache updates are exactly this kind of mutation, so verify your configuration instead of assuming it works.

Stale outputs. Because outputs live in the shared pool, a reference held across a replay can silently alias new data, or raise a RuntimeError in the compiled path. Clone what you retain. In a manual capture the failure is quieter: you read tokens from the previous step and nothing complains.

Host-device synchronization and data-dependent control flow. The docs list arbitrary control flow and .item() as unsupported; control flow expressed with torch.cond() can be captured. Early stopping logic that checks for the end-of-sequence token on the host each step therefore belongs outside the graph. A common pattern is to run several graphed steps, then check stop conditions once per chunk, accepting a few wasted tokens past the stop.

Capture-time memory spikes. Capturing many sizes at start-up, especially with FULL_AND_PIECEWISE, takes the longest and uses the most memory, per the vLLM docs. If you size the KV cache by free memory after loading the model, make sure graph memory is accounted for or you will see out-of-memory failures at the first large batch. Start-up time also grows, which matters for autoscaling cold starts.

Compile warmup cost. torch.compile pays Dynamo and Inductor time on the first call, plus Triton autotuning for max-autotune. The CUDAGraph Trees warmup call adds another pass before the first recording. Persist the compilation cache across restarts where your environment supports it, and do not measure latency in the first few iterations.

Custom kernels must be capture-safe. A custom CUDA or Triton kernel that allocates with cudaMalloc, synchronizes, or reads host-resident arguments will break capture. Allocation inside a captured region should go through the stream-ordered, graph-aware allocator, which PyTorch’s graph pool provides.

Multi-GPU adds its own constraints. NCCL collectives are supported in compiled CUDA graphs, per the docs, but tensor-parallel decode needs every rank to capture and replay the same sequence in lockstep. A rank that takes a different path (a skipped graph on one GPU, for instance) can stall the collective. Debug this by logging graph hit or miss per rank.

Anti-pattern: graphing the prefill. Prefill has variable lengths and large kernels. Capturing it multiplies recordings for little benefit, since g is large and launch overhead is already hidden. This is why gpt-fast, per its write-up, uses different compilation strategies for prefill and decode, and why vLLM’s FULL_DECODE_ONLY mode exists.

Practical Recommendations

Start by establishing that launch overhead is your problem. If a profile shows back-to-back kernels and a CPU that finishes early, graphs will not help; look at bytes moved instead. If it shows gaps, proceed.

For a single-model latency service or a draft model in speculative decoding, the lowest-risk path is torch.compile(mode="reduce-overhead", fullgraph=True) with a preallocated static KV cache and a device-resident position tensor, as in gpt-fast. Keep prefill in a separate, dynamic code path. Clone outputs you keep.

For multi-tenant production serving, let the engine own capture. Use the engine’s documented mode for your workload: a decode-focused instance in a prefill/decode disaggregated deployment is the natural fit for FULL_DECODE_ONLY, while a mixed workload usually wants the default FULL_AND_PIECEWISE if you can afford its memory. Tune the capture size list from observed batch-size histograms, not from guesses. Keep an eager fallback enabled so out-of-range batches degrade instead of fail.

For quantized, fused or custom attention kernels, check the backend’s declared graph support first. A kernel that is not capture-safe will silently demote you to a weaker mode.

A short checklist before shipping:

  • Profile a decode step and confirm gaps between kernels before adding graphs.
  • Preallocate the KV cache to its maximum size; no growing tensors in the captured region.
  • Keep position, sequence length and block-table data in device tensors updated in place, not in Python integers.
  • Run with fullgraph=True in development to catch graph breaks.
  • Pad batches to a small set of sizes; record how often each bucket is hit.
  • Grep logs for “skipping due to mutated inputs” and any CUDA graph skip messages; alert on them.
  • Budget memory: roughly 64 KB per captured kernel per graph in the documented configuration, times buckets, times modes.
  • Clone any output that outlives the next replay.
  • Keep sampling and stop-condition checks graph-safe or outside the graph.
  • Measure p50 and p99 per-token latency before and after, with warmup excluded.
  • Re-test after upgrading PyTorch, CUDA or the serving engine; defaults change.

Frequently Asked Questions

What are CUDA graphs and why do they speed up LLM decoding?

A CUDA graph records a sequence of GPU operations once and replays them with a single launch call. Decode runs hundreds of short kernels per token, and in eager mode the CPU can spend more time launching them than the GPU spends running them. Replaying a captured graph pays the setup cost once during instantiation, which removes most per-kernel CPU overhead and lets the GPU stay busy between kernels.

Does torch.compile use CUDA graphs automatically?

Not by default. Passing mode="reduce-overhead" enables CUDA graph management through CUDAGraph Trees, which record a graph per compiled region and execution path and replay them after a warmup call and a recording call. The default mode still fuses operations through Inductor but does not capture graphs. Check the documentation for your PyTorch version, because mode semantics and defaults evolve between releases.

Why do CUDA graphs require static shapes and fixed memory addresses?

A captured kernel stores its launch configuration and argument values, including device pointers. Replay reuses those exact values, so it reads and writes the same memory addresses. New data must be copied into the static buffers the graph already points at, and any change in tensor shape or scalar arguments requires a new recording. That is why serving engines pad batches to a few fixed sizes.

How much faster does a CUDA graph make decode?

It depends on how launch-bound your workload is. The gpt-fast write-up reports 25.5 tokens per second rising to 107.0 at batch size one on an A100-80GB after applying torch.compile with reduce-overhead and a static KV cache, a gain of more than 4x. Larger models and larger batches have longer kernels, so the relative gain shrinks. Measure on your own hardware.

What is the difference between full and piecewise CUDA graphs in vLLM?

Piecewise graphs capture everything except attention, which runs eagerly between pieces, and work with any attention backend. Full graphs capture the entire forward pass, including attention, for backends that support it. vLLM’s FULL_AND_PIECEWISE mode uses full graphs for uniform decode batches and piecewise graphs for others. It is usually the fastest but needs the most memory and capture time.

Can I use CUDA graphs with dynamic batch sizes?

Yes, through bucketing. Capture graphs for a chosen list of batch sizes, pad each live batch up to the nearest captured size, replay, and discard the outputs for padded rows. Batches larger than the biggest bucket fall back to eager execution. Each extra bucket costs device memory, so choose sizes based on your real batch-size distribution instead of capturing every value.

Further Reading

By Riju — about

Comments

No comments yet. Why don’t you start the discussion?

Leave a Reply

Your email address will not be published. Required fields are marked *