Article

LLM CUDA Graph Production Practices: Reducing Decode Launch Overhead with Shape Buckets, Piecewise Capture, and Eager Fallback

CPU kernel launch overhead in LLM decode is often overlooked. This deep dive explains CUDA Graph capture and replay, combined with shape bucket design, piecewise capture, warmup state machines, and eager fallback, delivering a production-ready solution to reduce tail latency and P99 jitter.

Background: GPU Is Busy, but CPU Can Also Slow Down Decode

LLM inference optimization often focuses on GPU utilization, batch size, and KV cache. However, in low-batch, short-sequence, or latency-sensitive scenarios, another bottleneck emerges: the CPU must repeatedly prepare and submit a large number of GPU kernels.

Autoregressive decoding requires a full model forward pass for each generated token. Each pass consists of many kernels, memory operations, and dependencies. If every step goes from Python and C++ runtime down to the CUDA driver, the cumulative launch cost of even fast kernels can amplify P95 and P99 latency jitter.

CUDA Graph doesn’t aim to make matrix multiplications faster. Instead, it captures a stable sequence of GPU work as a graph and replays it with a single graph launch, reducing the overhead of per-kernel parameter preparation and scheduling.

This optimization seems straightforward, but production deployment faces three challenges:

  1. LLM request batch sizes, sequence lengths, and execution paths are not completely fixed.
  2. Graph replay requires stable memory addresses and execution structures.
  3. Attention backends, MoE, sampling logic, and custom operators may not fully support capture.

Therefore, a production solution cannot rely on a single “enable CUDA Graph” switch. It must combine shape buckets, piecewise capture, warmup, and eager fallback.

Core Principle: Capture Once, Replay Stably

What CUDA Graph Records

PyTorch documentation describes CUDA Graph as a recording of work on a CUDA stream, including kernels, parameters, and dependencies. Capture does not execute the work immediately; it writes it into a graph. After capture, the same graph can be replayed multiple times.

The main benefit of graph replay is skipping repeated scheduling at the Python, C++, and CUDA driver levels, submitting an entire GPU workload with fewer host calls. It trades some dynamic flexibility for lower CPU overhead.

Two constraints are essential to understand:

  • Each replay typically executes the same kernel topology and parameter structure.
  • Input and output tensors must reuse the virtual memory addresses from capture. Business data should be copied into fixed buffers, not created with new addresses each time.

This is also where CUDA Graph differs from ordinary “code compilation.” Compilation addresses operator generation and optimization, while graph replay addresses repeated submission of already-determined GPU work. They can work together but are not the same thing.

Why Decode Benefits More

LLM inference is typically divided into Prefill and Decode:

PhaseCharacteristicsCPU Launch Overhead Ratio
PrefillProcesses a long prompt once, heavy matrix computationLow
DecodeProcesses few new tokens per step, many small kernelsHigh

Therefore, low-batch, interactive requests, and decode-heavy workloads are usually the best candidates for CUDA Graph. Conversely, if the service is already dominated by large-batch GEMM, inter-card communication, or queuing time, graph replay may bring only marginal improvement.

The Conflict Between Static Shapes and Dynamic Traffic

PyTorch documentation clearly states that standard CUDA Graph capture requires shapes and layouts to remain consistent during replay. LLM serving is inherently dynamic: the number of active sequences varies, Prefill and Decode mix, and some requests include LoRA, MoE, vision encoders, or custom logits processors.

Production systems typically don’t aim for a single graph covering all requests. Instead, they establish a limited set of shape buckets:

  • Batch sizes 1, 2, 4, 8, 16.
  • Separate pure Decode from Prefill/Mixed.
  • Optionally split by query length, speculative token count, or model branch.

At runtime, requests are mapped to the nearest available bucket, and the system decides whether to replay an existing graph or fall back to the eager path.

Engineering Implementation: From Traffic Profiling to Safe Fallback

Step 1: Prove Kernel Launch Is the Bottleneck

Don’t enable the feature first and then look for gains. Establish a baseline with a profiler:

  • Are there obvious CPU submission gaps between GPU kernels?
  • What is the CPU time ratio in a single decode step?
  • TPOT and P99 at batch sizes 1, 2, 4, 8.
  • Is the GPU already saturated by computation, memory bandwidth, or NCCL communication?
  • What is the traffic distribution across different prompt lengths and output lengths?

If most requests are high-batch or long-prefill, CUDA Graph may not be the top priority. If small-batch decode dominates and the timeline shows frequent kernel launch gaps, proceed to the next step.

Step 2: Design Shape Buckets Based on Real Distribution

More buckets are not always better. Each additional bucket typically means extra capture time, fixed input/output buffers, and graph cache management costs.

Start with a small set of high-coverage buckets:

cuda_graph_policy:
  decode_batch_buckets: [1, 2, 4, 8, 16]
  capture_prefill: false
  max_cached_graphs: 5
  fallback_mode: eager
  capture_on_startup: true

This configuration is an architectural example and does not correspond to a specific framework’s full parameters. Actual bucket sizes should come from online batch distribution, not mechanical powers of two.

The bucket selector needs to handle three cases explicitly:

def choose_execution_mode(batch_size: int, graph_cache: dict):
    bucket = smallest_bucket_not_less_than(batch_size)
    if bucket is None:
        return "eager", None
    if bucket in graph_cache:
        return "cuda_graph", bucket
    return "eager", None

If using ceiling buckets, ensure correctness via padding or valid token masks, and measure whether the extra computation offsets the graph gains.

Step 3: Prioritize Piecewise Capture Over Full Graph

Full capture requires the entire execution path to be CUDA Graph compatible. Production models often include attention, collectives, dynamic branches, and custom operators. A single incompatibility can cause capture failure.

A safer strategy is piecewise capture:

  • Stable, capturable operators go into the graph.
  • Attention or incompatible parts remain eager.
  • At runtime, choose Full, Piecewise, or None based on batch type and backend capability.

vLLM’s current documentation provides modes like NONE, PIECEWISE, FULL, FULL_DECODE_ONLY, and FULL_AND_PIECEWISE, using a dispatcher to select the execution path based on batch composition. Its default recommendation reflects an important principle: performance optimization must be evaluated together with backend compatibility and memory cost.

Example command:

vllm serve meta-llama/Llama-3.1-8B-Instruct \
  --compilation-config '{"cudagraph_mode":"FULL_AND_PIECEWISE"}'

When deploying, don’t copy mode names blindly. First, confirm the support matrix for your current vLLM version, attention backend, and model type.

Step 4: Integrate Warmup and Capture into the Startup State Machine

Graph capture is not a simple health check. It requires the model, CUDA context, memory pool, attention kernels, and communication paths to be in a stable state.

Consider splitting instance startup into four states:

StateMeaning
MODEL_LOADEDWeights loaded
EAGER_WARMEDSeveral eager warmup runs completed for typical shapes
GRAPH_CAPTUREDTarget buckets captured and validated
READYCan accept production traffic

PyTorch documentation requires eager warmup before capture and emphasizes that input/output tensors must be kept alive after capture. vLLM’s implementation also distinguishes warmup from actual capture to ensure the attention backend uses the correct kernel path.

If an instance becomes ready before graphs are prepared, the first user requests may bear the capture jitter or even trigger timeouts.

Step 5: Always Keep Eager Fallback

The graph path should not be a single point of failure. Fall back directly in these cases:

  • Batch or query shape misses a bucket.
  • Backend declares incompatibility with the current graph mode.
  • Capture, instantiation, or replay fails.
  • Dynamic features change the execution topology.
  • A new model version hasn’t completed graph compatibility validation.
  • Debugging requires a more complete error stack.

Fallback is not failure; it’s a compatibility strategy. The key is to log the reason. Otherwise, the system may appear normal while running in eager mode long-term, with the team mistakenly believing CUDA Graph is active.

Output these labels:

execution_mode=cuda_graph|piecewise|eager
bucket_id=decode_bs_4
fallback_reason=shape_miss|backend_unsupported|capture_error|disabled
model_revision=...
attention_backend=...

Step 6: Manage Graph Cache and Memory Budget

Each graph is not just a lightweight handle. Fixed tensors, private memory pools, workspace, and capture state can all increase resident memory.

Establish a dedicated budget for graphs:

  • Maximum number of graphs.
  • Memory increment per bucket.
  • Capture time.
  • Eviction policy for idle graphs.
  • Invalidation conditions after model or LoRA switching.
  • Traffic handling during graph reconstruction.

vLLM documentation notes that FULL_AND_PIECEWISE often performs better but requires more memory and has longer capture time. Production choices should be based on the ratio of P99 gain to memory cost, not defaulting to the most aggressive mode.

Applicable Scenarios

SuitableProceed with Caution
Low-batch interactive decode (online chat, code completion)Extremely discrete request shapes
Single model, fixed precision, stable execution pathHeavy dynamic control flow and incompatible custom operators
CPU scheduling overhead is significant in single decode stepMain bottleneck is large matrix computation, NCCL, or memory bandwidth
Shape distribution is concentrated, few buckets cover most requestsFrequent instance scaling and restart, capture cost cannot be amortized
More sensitive to TPOT and P99 jitter than peak throughputMemory near limit, no room for multiple graphs

Common Misconceptions

Misconception 1: CUDA Graph Equals Model Compilation

torch.compile, kernel fusion, and CUDA Graph solve different problems. Compilation optimizes the computation graph and generates kernels; CUDA Graph mainly reduces repeated submission overhead for stable GPU work sequences. They can be combined or used independently.

Misconception 2: More Graphs and Higher Coverage Are Always Better

More buckets may reduce shape misses but increase startup time, memory usage, and version management complexity. Production systems should focus on “how much core traffic is covered by a small number of buckets.”

Misconception 3: Only Look at Average Latency

The value of CUDA Graph is often seen in the stability of TPOT, P95, and P99. Comparing only average throughput may miss the benefit of reduced CPU scheduling jitter.

Misconception 4: Successful Capture Guarantees Correct Results

Graph replay reuses fixed addresses and execution structures. Buffer updates, masks, random number states, sampling parameters, and cross-stream dependencies can cause silent errors if mishandled. Consistency regression tests between eager and graph outputs are mandatory.

Misconception 5: Capture on the Fly When a Shape Miss Occurs

Synchronous capture triggered by an online request creates unpredictable latency spikes. A safer approach is to capture a whitelist of buckets during startup, route miss requests through eager, and use offline data to decide whether to add new buckets in the next version.

Deployment Checklist

  • Confirmed via profiler that CPU kernel launch is a visible bottleneck.
  • Buckets derived from real traffic distribution, with coverage recorded.
  • Eager, piecewise, and full paths all independently runnable.
  • Capture completed before instance is ready.
  • Output consistency tests between graph and eager cover sampling, masks, and edge shapes.
  • Attention backend and model version included in compatibility matrix.
  • Clear fallback for shape miss, backend incompatibility, and capture error.
  • Monitored graph hit rate, fallback rate, capture time, and memory increment.
  • Release supports toggling CUDA Graph off per instance, model, or tenant.
  • Model upgrades trigger recapture and compatibility gates.

FAQ

Does CUDA Graph always reduce TTFT?

Not necessarily. TTFT includes queuing, prefill, network, and scheduling. CUDA Graph more directly affects submission overhead for repeated GPU work, which is more visible in decode and small-compute-granularity scenarios. Whether TTFT improves depends on whether prefill is also captured and where the real TTFT bottleneck lies.

Should shape buckets be divided by batch size or sequence length?

First, determine which dimensions affect the execution path. Pure decode is often dominated by batch size and query length; prefill/mixed may also need to consider token count and attention metadata. Don’t build a multi-dimensional Cartesian product at once. Start with the dimension that most affects capture compatibility.

Can CUDA Graph update parameters without recapturing?

CUDA provides graph update and single-node parameter update capabilities, but updates have limitations on topology, context, and memory types. LLM frameworks typically encapsulate these details in their own dispatcher and graph cache layers. Business logic should not assume all dynamic changes can be updated in place.

References

  1. PyTorch CUDA Semantics — CUDA Graphs
  2. NVIDIA CUDA Programming Guide — CUDA Graphs
  3. vLLM Design Documents — CUDA Graphs
  4. PyGraph: Robust Compiler Support for CUDA Graphs in PyTorch
  5. Hybrid JIT-CUDA Graph Optimization for Low-Latency Large Language Model Inference

FAQ

Is CUDA Graph suitable for all LLM inference requests?
No. It's best for decode tasks with stable execution paths, shape-bucketable workloads, and significant CPU launch overhead. Highly dynamic requests or those with incompatible operators should keep eager fallback.
Does more shape buckets always mean higher hit rate?
Hit rate typically increases, but capture time, memory usage, and graph cache management costs also grow. Choose a small set of high-coverage buckets based on real traffic distribution.
Why might latency not improve after enabling CUDA Graph?
When requests are dominated by matrix computation, memory bandwidth, communication, or queuing time, kernel launch is not the main bottleneck, and graph replay gains are drowned out by other costs.
Does CUDA Graph always reduce TTFT?
Not necessarily. TTFT includes queuing, prefill, network, and scheduling. CUDA Graph more directly reduces submission overhead for repeated GPU work, which is more visible in decode and small-compute-granularity scenarios.