Attention Backends & Graph Optimization

The attention backend is where rubber meets road — it selects the right kernel for each batch composition and eliminates CPU overhead that would otherwise bottleneck decode throughput.

1. The CPU Overhead Problem

A decode step’s GPU compute time depends on model size and TP degree. For a 70B model with TP=8, each GPU loads ~17.5 GB of weights at 3,350 GB/s = ~5ms. For an 8B model on a single GPU: ~5ms. But launching CUDA kernels adds CPU overhead:

Operation	CPU time	GPU time
Python scheduler logic	~0.5 ms	0
CUDA kernel launch (per kernel)	~5-10 μs	—
Typical decode step (50+ kernels)	~0.5 ms launch overhead	~5 ms compute
Total		~5.5 ms (9% CPU overhead)

For small, fast models (1-3B, INT4 quantized) where GPU compute drops to <1ms, the CPU overhead becomes the dominant cost — kernels finish before the CPU can launch the next ones. This is where CUDA graphs become critical.

2. Pluggable Attention Backends

Modern engines support multiple attention implementations, selected based on the operation:

Backend	Best for	Key feature
FlashAttention-2	Prefill (long sequences)	Tiled, O(1) memory, no materialized attention matrix
FlashDecoding	Decode (single token, long KV)	Parallelizes across KV length, not just heads
FlashInfer	Both + PagedAttention native	Built-in support for paged KV layout
xFormers	General purpose	Wide hardware support

Prefill vs Decode Kernel Selection

Incoming batch:
  - 3 decode requests (1 token each)
  - 1 prefill chunk (512 tokens)

Engine selects:
  - FlashAttention-2 for the prefill chunk (long sequence, compute-heavy)
  - FlashDecoding for the decode requests (single token, memory-heavy)

The attention backend must handle mixed batches where prefill and decode coexist. Some implementations (FlashInfer) handle this natively; others require splitting the batch.

3. CUDA Graphs: Eliminating Launch Overhead

CUDA graphs capture a sequence of kernel launches into a single launchable graph — eliminating per-kernel CPU overhead:

# Without CUDA graph: launch 50 kernels one by one
for kernel in model_kernels:
    kernel.launch()  # each launch: ~10μs CPU overhead

# With CUDA graph: launch all 50 as a single graph
graph.launch()  # one launch: ~10μs total

Capture and Replay

# Capture phase (once, during warmup):
with torch.cuda.graph(graph):
    output = model(input_placeholder)  # record all kernel launches

# Replay phase (every step):
input_placeholder.copy_(real_input)  # update input in-place
graph.replay()  # replay all kernels — minimal CPU overhead

The Catch: Fixed Shapes

CUDA graphs require fixed tensor shapes. But batch sizes change every step (requests arrive/complete). Solutions:

Padding: Pad batch to fixed sizes (wastes compute on padding)
Multiple graphs: Capture graphs for common batch sizes (1, 2, 4, 8, 16, 32…)
Piecewise capture: Capture only the decode portion (fixed per-request shape), keep scheduling logic dynamic

4. torch.compile and Piecewise CUDA Graphs

PyTorch’s compiler can fuse operations and eliminate Python overhead:

# Compile the model forward pass
model.forward = torch.compile(model.forward, mode="reduce-overhead")

Benefits:

Operator fusion (combine multiple small kernels into one)
Memory planning (reduce intermediate tensor allocations)
Reduced Python/framework overhead

Piecewise CUDA Graphs

Piecewise CUDA Graph

The key challenge: full model CUDA graphs require fixed shapes, but the attention kernel’s shape changes with batch composition (different numbers of prefill/decode tokens). The solution is piecewise capture:

Identify static subgraphs: Parts of the model with fixed shapes regardless of batch composition (e.g., FFN layers — always [batch, hidden] → [batch, ffn_dim])
Capture each subgraph separately: Each gets its own CUDA graph
Dynamic parts stay ungraphed: Attention (variable KV length) runs dynamically
torch.compile handles this automatically: It identifies capturable regions and captures them

Full forward pass:
  [Embedding]  → static, captured as graph
  [Attention]  → dynamic (variable KV length), NOT captured
  [FFN]        → static, captured as graph
  [LayerNorm]  → static, captured as graph
  [LM Head]    → static, captured as graph

This gives ~80% of full CUDA graph benefit while handling variable batch compositions.

5. Practical Impact

Optimization	Decode step time (8B, single GPU)	When to use
No optimization	~5.5 ms (10% CPU overhead)	Never in production
CUDA graphs only	~5.1 ms (launch overhead eliminated)	Standard for decode
torch.compile only	~5.2 ms (fusion helps)	When shapes vary too much for graphs
CUDA graphs + compile	~5.0 ms (best of both)	Maximum throughput

For small models (1-3B INT4) where GPU time is <1ms, CUDA graphs reduce total step time from ~1.5ms to ~1.0ms

a 33% speedup. The smaller the model, the more CPU overhead matters.

# vLLM with CUDA graph and compile
vllm serve meta-llama/Llama-3.1-70B \
    --enforce-eager=false \       # enable CUDA graphs (default)
    --compilation-config '{"level": 3}'  # enable torch.compile

For the attention kernels themselves (FlashAttention algorithm), see Flash Attention v1. For tracing a full request through the engine, see Engine Walkthrough.

Self-Check

Explain why CPU overhead is a bottleneck for decode steps
Compare FlashAttention-2 vs FlashDecoding and when each is used
Explain CUDA graph capture and replay, and the fixed-shape constraint
Describe how engines handle mixed prefill/decode batches
Configure vLLM with CUDA graphs and torch.compile