2.5 Attention backends and graph optimization

The first four sections were “macro” techniques—how to manage GPU memory, how to batch, how to reuse prefixes, how to chunk. This section zooms to the “micro” level: for the same batch of requests in one forward pass, how do you make it run faster? The answer is two things: pick an efficient Attention backend kernel, and use CUDA Graph + torch.compile to wipe out the invisible but costly CPU overhead in Decode. That is one of the important reasons vLLM V1 is faster than V0.

⏰ Currency note: default backend choice, CUDA Graph modes and capture sizes, and torch.compile behavior are all based on vLLM v0.32.0 (released 2026-10-18) and the matching official design docs.

📑 Contents

• 1. A counterintuitive bottleneck: the CPU holding the GPU back
• 2. Pluggable Attention backends
• 3. How backends handle mixed Prefill/Decode batches
• 4. CUDA Graph: pack a string of kernel launches into one
• 5. torch.compile and Piecewise CUDA Graph
• 6. Source walk-through: how vLLM splits and captures graphs
• 7. Switches and tradeoffs in practice
• Summary
• Self-check list
• References

---

1. A counterintuitive bottleneck: the CPU holding the GPU back

Start with a counterintuitive fact. The last chapter showed Decode is memory bound—the bottleneck is moving weights from GPU memory. But when the model is small and the batch is small too, a subtler bottleneck shows up: the cost of the CPU launching kernels.

The GPU does not start work on its own. Every operation (one matmul, one LayerNorm, one activation…) needs the CPU to tell it with a “launch” command. One Transformer forward has hundreds or thousands of such small ops, so hundreds or thousands of CPU→GPU launch calls. Each launch has fixed CPU cost (build arguments, driver scheduling, submit to the queue), on the order of microseconds.

Usually that cost barely matters—Prefill processes thousands of tokens in one step, so each GPU kernel runs for hundreds of microseconds or even milliseconds and swamps the launch cost. But Decode processes only 1 token (or a few) per step, so a kernel’s real compute may be just a few microseconds, and the CPU time to launch that kernel can exceed the GPU time to run it. The GPU works for a few microseconds, then sits waiting for the next command—compute idles again, this time not waiting for data but waiting for instructions.

The top half of the figure below is that case: yellow is the CPU launching kernels one by one, blue is the GPU actually computing, gray hatched cells are GPU idle time waiting for the next instruction. The bottom half is CUDA Graph from §4: the CPU submits once (green), eight kernels run back to back on the GPU, and orange dashed lines mark when each approach finishes the step.

📌 Key point: Decode has two kinds of “waiting”—waiting for data (memory bound; ease with batching/quantization) and waiting for instructions (CPU launch cost; ease with the graph optimizations in this section). On small models and small batches, the latter stands out. That is why graph optimization speeds Decode up so much.

---

2. Pluggable Attention backends

Attention is the heaviest, most carefully engineered op in a Transformer; its kernel implementation directly sets inference speed. vLLM’s design is not to lock to one implementation, but to offer pluggable Attention backends, picking the best one automatically or manually from hardware, model structure, and workload.

Common backends:

🔑 Core idea: Attention backends are different engineering implementations of the same math. Results are mathematically equivalent; they differ in memory access patterns, parallel strategies, and how well they fit specific hardware instructions and model structures. The right backend can change throughput a lot on the same card.

💡 Tip: most of the time you do not set this by hand; vLLM auto-picks a backend from GPU architecture, model type, and precision. In v0.32.0, Ampere and Hopper default to FLASH_ATTN, Blackwell to FLASHINFER; MLA models take a dedicated MLA backend. To intervene, use --attention-backend for the backend and --attention-config.flash_attn_version for FlashAttention version (2, 3, or 4). The support matrix moves quickly across releases; for a given card and model, follow the docs for your version.

---

3. How backends handle mixed Prefill/Decode batches

Here is a key link left from the previous section. vLLM V1’s unified scheduler packs Prefill chunks and Decode requests into the same step—so how does Attention run for that step? Prefill needs “attention among a whole stretch of new tokens + over history”; Decode needs only “1 new token over all history.” The compute shapes differ.

So the Attention backend must natively support mixed Prefill/Decode batches: one kernel call correctly handles a batch that mixes Prefill and Decode. Modern backends such as FlashAttention provide a varlen interface for exactly that—concatenate every request’s Query into one variable-length sequence, with start/end offsets per request, and finish in one kernel call without splitting Prefill and Decode into two calls.

📌 Key point: the “unified scheduler” (software erases Prefill/Decode) and “Attention backends that support mixed batches” (kernels erase Prefill/Decode) are a matched pair. Without the latter, V1’s elegant token-budget scheduling cannot land. That is why V1 requires capable Attention backends; early on, not every backend could support this.

---

4. CUDA Graph: pack a string of kernel launches into one

Back to the CPU launch cost in §1. CUDA Graph is NVIDIA’s remedy: capture a fixed sequence of GPU ops into a “graph,” then replay the whole graph with one command so the GPU runs every recorded kernel in order—hundreds or thousands of CPU launches collapse into one.

Analogy: without CUDA Graph, the teacher waits for a secretary to relay every sentence (each kernel needs its own CPU launch); with CUDA Graph, the teacher records the whole lecture in advance, hits play in class, and it runs straight through without the secretary relaying line by line. The bottom half of the §1 timeline is replay: after one submit the CPU is free to prepare the next scheduling step.

But CUDA Graph has a hard requirement: the recorded op sequence and tensor shapes must be fixed. That fits Decode perfectly—every step is “process a fixed-batch number of tokens,” shapes are stable, so you can capture a graph per batch size and pick the matching one at runtime.

⚠️ Note: because shapes must be fixed, vLLM only captures CUDA Graphs for a preset set of batch sizes. In v0.32.0 the default capture sizes are [1, 2, 4], plus every 8 from 8 to 248, plus every 16 from 256 to max_cudagraph_capture_size; max_cudagraph_capture_size defaults to 512 (1024 on datacenter-class Blackwell GPUs), and both are adjustable via compilation_config. At runtime the actual batch is padded up to the nearest captured size before replay—e.g. a 13-token batch runs the size-16 graph, with 3 empty slots; only batches above the max capture size fall back to per-kernel launch. So CUDA Graph gains show mostly on small, regular Decode batches; large Prefill chunks are harder to apply directly.

---

5. torch.compile and Piecewise CUDA Graph

CUDA Graph removes “launch overhead,” but there is another layer: can many small kernels themselves be fused into fewer, larger ones? That is what torch.compile does—it compiles the model’s forward into an optimized graph and cuts kernel count through fusion and similar passes. In vLLM V1, torch.compile is on by default, and all compilation finishes before serving starts, so requests do not trigger compile spikes.

There is a tension: CUDA Graph wants a fixed-shape stretch of ops, but Attention is the hardest part to fit under CUDA Graph—variable-length KV, complex masks, shapes that are not tidy. Giving up the whole graph’s CUDA Graph because Attention is awkward would be a shame.

vLLM V1’s fix is Piecewise CUDA Graph:

• Split the full forward graph at Attention ops (Attention is wrapped as an “opaque” custom op the compiler does not touch).
• Everything outside Attention (LayerNorm, GEMMs, activations, residuals… all regular token-wise ops) is captured and replayed with CUDA Graph.
• Attention itself stays eager, keeping flexibility for variable-length KV.

🔑 Core idea: Piecewise CUDA Graph = squeeze CPU overhead out of the regular parts with graph optimization, and leave Attention eager for flexibility. Each side takes what it is good at: the bulk of CPU launch cost (those hundreds or thousands of small ops) disappears under CUDA Graph, while Attention’s flexibility (mixed batches, variable-length KV) stays intact. That is one of the key engineering pieces behind V1’s Decode speedup over V0.

vLLM also supports Full CUDA Graph, recording Attention into the graph too, when the backend supports capture on that kind of batch. In v0.32.0, V1’s default mode is the combined FULL_AND_PIECEWISE: pure Decode batches have the most regular shapes and replay one FULL graph for the whole step; mixed batches that include Prefill fall back to PIECEWISE, with Attention outside the graph. Earlier releases defaulted to PIECEWISE; pooling models still default to PIECEWISE. Full mode set also includes NONE, PIECEWISE, FULL, FULL_DECODE_ONLY, set via cudagraph_mode in compilation_config.

The top half of the figure below is the split forward: blue are subgraphs from torch.compile, yellow is Attention left outside the graph, and intermediate layers share one subgraph kind. The bottom half is runtime choice under FULL_AND_PIECEWISE: green boxes are one FULL graph for the whole step; black boxes are one CUDA Graph per subgraph segment.

💡 Tip: whether a FULL graph can be used depends on how far the Attention backend supports CUDA Graph. Official docs split it into tiers: FA3 and Triton support any batch (ALWAYS); FA2 and FlashMLA only support batches with uniform shapes (UNIFORM_BATCH); FlashInfer only supports pure Decode batches with 1 token per request (UNIFORM_SINGLE_TOKEN_DECODE). Batches that do not meet the bar automatically fall back to PIECEWISE. Docs also note FULL graphs help Decode most on small models and MoE models.

---

6. Source walk-through: how vLLM splits and captures graphs

That was the Piecewise CUDA Graph idea; this section looks at how vLLM V1 does it in code. The core sits in the torch.compile integration (official design doc design/torch_compile). Code keeps moving; the description below follows main-branch docs, and concrete APIs follow your installed version.

6.1 Attention wrapped as an “opaque op”

The first step of splitting the graph is making the compiler “unable to see into” Attention. vLLM registers Attention as the custom op torch.ops.vllm.unified_attention_with_output. When PyTorch’s graph tracer (Dynamo) hits it, it treats it as a black box—no tracing inside, no trying to optimize it—but it can still capture a complete compute graph around it.

🔑 Core idea: wrapping Attention as an opaque op is what makes “splitting the graph” work. The compiler leaves Attention’s variable-length logic alone and only optimizes the regular-shaped layers around it—that is §5’s “split at Attention” in code.

6.2 One graph becomes three subgraph kinds

After Dynamo traces the model’s forward, it splits the graph on splitting_ops (usually Attention). Because a Transformer is a stack of isomorphic layers, the split yields only three distinct subgraphs, matching subgraphs ①②③ in the §5 figure:

• First layer: everything before the first Attention.
• Repeated middle layer: everything between two adjacent Attentions—this one subgraph is reused by every middle layer and compiled once.
• Last layer: everything after the last Attention through the output.

📌 Key point: because middle layers are isomorphic and subgraphs are reusable, compile cost stays low—even a model with dozens of layers only compiles three subgraphs. That is an engineering dividend from Transformer’s regular structure.

6.3 Capture and replay by batch size

The regular subgraphs from compilation are then captured with CUDA Graph. Because CUDA Graph needs fixed shapes, vLLM only captures for a preset set of batch sizes and pads the batch to the nearest capture size at runtime before replay:

from vllm import LLM
from vllm.config import CompilationConfig

llm = LLM(
    model="meta-llama/Llama-3.1-8B-Instruct",
    compilation_config=CompilationConfig(
        # Only capture CUDA Graphs for these batch sizes; pad up to the nearest size on replay.
        # Batches above the max (32 here) skip CUDA Graph
        cudagraph_capture_sizes=[1, 2, 4, 8, 16, 32],
        # Optional: Inductor auto-tune for specific static shapes (slow first time; off by default)
        compile_sizes=[1, 2, 4, 8],
    ),
)

• cudagraph_capture_sizes: which batch sizes get a captured graph. Unset uses the default list from §4; a longer list covers more finely and wastes less on padding, but capture time and memory grow.
• compile_sizes: lets Inductor auto-tune for listed static shapes and find faster Triton configs—faster at runtime, slower on first compile, so off by default.

💡 Tip: compile artifacts go into a compile cache, default ~/.cache/vllm/torch_compile_cache (the cache key includes all related config and traced source files). The cache directory can be copied across deployments to skip repeat compiles; to force a rebuild while debugging, set VLLM_DISABLE_COMPILE_CACHE=1.

---

7. Switches and tradeoffs in practice

For users, most of these optimizations are on by default and automatic; day to day you do not need to touch them. Knowing the tradeoffs still helps when debugging and tuning:

• ✅ Gains: Decode CPU overhead drops a lot; throughput and latency improve clearly on small models / small batches; torch.compile fusion cuts kernel count further.
• ❌ Costs:
  • Slower startup: serving start has to compile and capture CUDA Graphs, so first bring-up takes longer (vLLM softens this with the compile cache; the cache directory can be copied to speed later starts).
  • GPU memory: capturing CUDA Graphs for many batch sizes needs extra memory.
  • Harder debugging: a compiled graph is harder to debug than eager; hunting numeric issues may need temporarily turning these off.

⚠️ Note: for stubborn numeric anomalies or to check a kernel’s behavior, start temporarily with --enforce-eager; it turns off both torch.compile and CUDA Graph for a pure eager comparison. Remove it once you are done—production must keep these on, or you throw away Decode performance for nothing. Finer switches (e.g. set cudagraph_mode to NONE to disable only CUDA Graph) follow the official docs for your vLLM version.

---

📝 Summary

• Besides “waiting for data” (memory bound), Decode also “waits for instructions”—CPU per-kernel launch overhead, especially on small models / small batches.
• vLLM offers pluggable Attention backends (FlashAttention/FlashInfer/FlashMLA/Triton): same math, different engineering; usually auto-picked by GPU architecture (v0.32.0 defaults to FlashAttention on Ampere/Hopper, FlashInfer on Blackwell).
• Modern backends must natively support mixed Prefill/Decode batches, matching V1’s unified scheduler.
• CUDA Graph collapses hundreds or thousands of kernel launches into one replay and removes CPU overhead, but needs fixed shapes: batches are padded to the nearest capture size; batches above the max capture size skip the graph.
• Piecewise CUDA Graph splits at Attention: regular parts use CUDA Graph, Attention stays eager, combining performance and flexibility. V1’s default FULL_AND_PIECEWISE goes further: pure Decode batches use one FULL graph for the whole step.
• In source, Attention is wrapped as the opaque op unified_attention_with_output; the graph is split on splitting_ops into three reusable subgraphs, then captured per cudagraph_capture_sizes, with artifacts written to a copyable cache.
• These optimizations are on by default; the costs are slower startup, extra memory, and harder debugging. Use --enforce-eager temporarily when debugging; keep them on in production.

🎯 Self-check list

• Explain why CPU kernel-launch overhead becomes a bottleneck in Decode
• Name at least three vLLM Attention backends and when each fits
• Explain what “same math, different engineering” means for Attention backends
• Explain why the unified scheduler needs backends that support mixed Prefill/Decode batches
• Describe how CUDA Graph removes CPU overhead via “capture—replay” and its fixed-shape requirement
• Explain why Piecewise CUDA Graph splits at Attention, and how FULL_AND_PIECEWISE picks graphs by batch type
• Explain why the split yields only three subgraphs, and how cudagraph_capture_sizes relates to padding
• State the costs of these graph optimizations (startup, memory, debugging) and the production tradeoff

📚 References

• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• vLLM Design — torch.compile Integration
• vLLM Design — CUDA Graphs (cudagraph_mode and backend support levels)
• vLLM Design — Attention Backends (default backend selection)
• vLLM V1 Alpha Release (Piecewise CUDA Graph, FlashAttention 3)
• NVIDIA Developer Blog — Getting Started with CUDA Graphs
• vLLM — Engine Arguments (—attention-backend, —enforce-eager)