Skip to content
LLM Inference
S13Decoding Acceleration·185 lines

Kernel Fusion & CUDA Graphs

At small batch sizes the GPU is idle waiting for the CPU. Fusing operations and replaying a captured graph removes hundreds of launches per token.

  • launch overhead
  • fusion
  • graph capture
  • static shapes
  • bucket padding

Motivation

The problem

Profile a batch-1 decode step on a modern GPU and you will find something absurd: the GPU sits idle for nearly half of it. Memory bandwidth is not the reason. The CPU cannot issue work fast enough.

A 32-layer model runs roughly eleven kernels per layer: two norms, three projections, RoPE, attention, the output projection, and three more for the FFN. That comes to about 350 kernel launches per token, each costing 5–10 µs of CPU time to build and submit. Two milliseconds of CPU work per token, while each kernel occupies the GPU for a few microseconds.

You have a machine capable of a quadrillion operations per second waiting on a Python loop.

Core idea

The solution

Two independent fixes, usually applied together.

  • Kernel fusion — combine adjacent operations into one kernel. Fewer launches, and the intermediate values stay in registers instead of round-tripping through HBM. RMSNorm fused with the residual add; Q, K and V as one packed matmul; RoPE folded into the attention prologue; SwiGLU's gate, multiply and activation in one pass.
  • CUDA graphs — record the entire sequence of launches once, then replay the whole recorded graph with a single CPU call. The GPU driver already knows the dependency structure, so there is nothing left to submit per kernel.
graph capture and replaypython
class GraphRunner:
    def __init__(self, model, batch_size):
        # Static buffers: the graph records ADDRESSES, not values.
        self.input_ids = torch.zeros(batch_size, dtype=torch.long, device="cuda")
        self.positions = torch.zeros(batch_size, dtype=torch.long, device="cuda")

        model(self.input_ids, self.positions)          # warm up, then capture
        self.graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(self.graph):
            self.output = model(self.input_ids, self.positions)

    def run(self, input_ids, positions):
        # Copy into the SAME buffers the graph was captured against.
        self.input_ids[: len(input_ids)].copy_(input_ids)
        self.positions[: len(positions)].copy_(positions)

        self.graph.replay()                            # ONE cpu call
        return self.output.clone()

A graph records addresses, not values

Capture writes down "run kernel A reading from pointer 0x7f… and writing to 0x7f…", and replay re-executes exactly that. Inputs must therefore be copied into the same buffers every time, and any tensor allocated fresh inside the captured region breaks the recording. Every constraint below follows from this.

Mechanics

How it works

Diagramfrom launch-bound to execution-bound
EAGER MODE — ONE LAUNCH PER KERNELCPUGPUDashed regions are bubbles: the GPU finished and is waiting for the CPU to submit the next kernel.A 32-layer model launches ~352 kernels per token. At 6 µs each that is 2.1 ms of CPU time alone.STEP 1 — FUSErms_norm+ residualq/k/v projone fused kernel3 launches → 1, and the intermediate never reaches HBMSTEP 2 — CAPTURE THE WHOLE STEP AS A GRAPHcapture oncerecord the DAGGRAPHreplay every stepone CPU callCPUGPUNo bubbles. The GPU runs back-to-back.THE CONSTRAINTA graph records fixedshapes and fixed addresses.Batch size changes everystep, so capture buckets.Engines capture graphs for batch sizes 1, 2, 4, 8, … and pad the real batch up to the next bucket.Prefill has too many distinct shapes to capture, so graphs are a decode-only optimisation.
Fusion removes launches; graphs remove the per-launch CPU cost of the ones that remain. The constraint panel explains why this optimisation applies to decode only.

Static shapes, and the bucketing workaround

A captured graph is valid for exactly the shapes it was captured with. Continuous batching, meanwhile, changes the batch size every single step; varying it is the entire point.

The standard resolution is bucketing: capture graphs for batch sizes 1, 2, 4, 8, 16, 24, 32, … and at run time pad the real batch up to the next captured size. Padding wastes a little compute; the launch saving is worth far more at small batches.

Prefill gets no graphs at all. Its sequence lengths are effectively unbounded, so there is no small set of buckets to capture. That costs nothing, because prefill is compute-bound and launch overhead is already hidden behind the arithmetic.

When this stops mattering

Launch overhead is a fixed cost per step, so its importance falls as the GPU work per step rises. At batch 1 it can be the majority of step time. At batch 256 the kernels take long enough that launches disappear behind them entirely.

CUDA graphs therefore matter most exactly where batching helps least: low-latency, low-concurrency serving. The two techniques cover for each other.

Explore

Try it

Simulatorlaunch overhead vs kernel time
kernels / step
288
step time
1.73 ms
bound by
CPU launch
vs unoptimised
1.00×
first 18 kernel launches
CPU
GPU
  • CPU: building and submitting a launch
  • GPU: executing the kernel
  • idle

The GPU lane has visible gaps: it finishes each kernel before the CPU has submitted the next. The accelerator is waiting on a Python-speed loop. Adding more GPU will not help.

where the time goes
CPU launch cost1.73 ms
GPU execution0.89 ms

Step time is the maximum of the two rather than the sum, because they overlap. Optimisation only helps if you shrink the larger one, so profile before you touch anything here.

GPU utilisation

52%

At batch 1 with 32 layers and no optimisation you are typically launch-bound: ~352 kernels at 6 µs each is 2.1 ms of pure CPU overhead per token, which can exceed the GPU work entirely. Turn on both toggles and watch it invert.

Start with batch 1 and both toggles off: the GPU lane is full of visible gaps and the readout says CPU-bound. Turn on fusion, then graphs, and watch the bottleneck move to the GPU. Now switch both toggles back off and drag batch size to 128. The gaps close on their own, because the kernels have grown long enough to hide the launches.

Build it

Implementation

The bucketing logic is short. Padding is the part people get wrong.

code/s13_cuda_graphs.py (excerpt)python
BUCKETS = [1, 2, 4, 8, 16, 24, 32, 48, 64, 96, 128, 192, 256]

class BucketedRunner:
    def __init__(self, model):
        self.graphs = {b: capture(model, b) for b in BUCKETS}

    def forward(self, batch):
        n = len(batch)
        bucket = next(b for b in BUCKETS if b >= n)

        padded = pad_to(batch, bucket)          # real requests + dummies
        out = self.graphs[bucket].run(padded)
        return out[:n]                          # discard the dummy rows
        # The dummies MUST be harmless: point them at a scratch KV block,
        # or they will write garbage into a real request's cache.

Padding rows must not touch real memory

The dummy rows in a padded batch run the same kernels as real requests, including the attention write into the KV cache. If a dummy row's block table points at a real block, it corrupts a real request's cache. Give padding its own scratch block and never reuse it. The resulting bug is intermittent, batch-size-dependent, and produces plausible output, which is about the worst combination available.
Run it locally
Models launch overhead against kernel execution time for a configurable model, showing where the crossover between CPU-bound and GPU-bound falls. Includes a fusion accounting model and a bucketing simulator that reports padding waste per bucket layout.
$ python code/s13_cuda_graphs.py
Expect: A batch-size sweep locating the launch-bound region, a fusion table showing kernels per layer falling from 11 to 4, and a bucket-layout comparison trading padding waste against pinned graph memory.

Only NumPy is required — setup instructions.

Production notes

In production

  • vLLM — captures decode graphs per batch bucket; enforce_eager=True disables it. Try that first when debugging a shape or memory bug.
  • TensorRT-LLM — takes this furthest: the entire model is compiled ahead of time into a fixed graph. Its speed comes from that, and so does the need to rebuild whenever shapes change.
  • torch.compile — automates fusion via Inductor and can capture graphs (mode="reduce-overhead"), recompiling per shape.
  • llama.cpp / picoLM — on CPU there is no launch overhead, but fusion still matters for a different reason: the fused dequant+dot pass halves memory traffic.

Exercises

  1. 1
    Compute the graph memory cost of your bucket layout, remembering that each captured graph pins its own static buffers. Find the layout that minimises padding waste under a fixed memory budget.
  2. 2
    Fuse RMSNorm with the residual add in the S03 implementation and confirm the output is unchanged. Count the eliminated intermediate tensors.
  3. 3
    Implement the padding-corruption bug deliberately, observe that output is subtly wrong only at certain batch sizes, then write the assertion that catches it.

Continue

What's next

Each forward pass is now as cheap as it can be. The remaining move is to stop doing one pass per token. S14 covers speculative decoding: it breaks the one-token-per-pass rule and leaves the output distribution untouched.

Check yourself

Questions

Answer before you look. A wrong answer tells you more than a right one, because the explanation points at the part of the chapter to reread.

Check yourselfQuestion 1 of 6

A CUDA graph records what, and why does that constrain how you use it?

Score 0/6