Skip to content
mini-sglang from Scratch Stage 4 · Faster and bigger

CUDA Graph

One decode step of a 28-layer model launches hundreds of kernels. At a small batch size each kernel computes for only a few microseconds, while the CPU takes a few microseconds to launch one, so the GPU spends most of its time waiting for the CPU. CUDA Graph records a whole forward pass so one replay launches every kernel. The price is that every address the recorded kernels read and write is now fixed, so every input has to be copied into fixed buffers first. This chapter implements mini-sglang's GraphRunner and verifies the "every input goes through a fixed buffer" contract with a CPU emulation.

Self-test: if you can answer these, skip the chapter

  1. Why is CUDA Graph used for decode and not for prefill?
  2. The recorded batch size is 4 and the real batch has 3 requests. What now? Where does the padded request's output go?
  3. The attention metadata (each request's KV length, the page table) differs every step. How does the graph see the new values?
  4. Why is the largest batch size recorded first?
Answers (try it yourself first, then expand)
  1. A decode step has very few tokens but hundreds of small kernels, so the CPU's launch cost dominates and a graph removes it; a prefill's shape (its token count) differs every time across a wide range, which is impossible to record for, and prefill does enough compute that launch overhead never mattered.
  2. Pad the batch to 4 with a dummy request: its inputs and metadata live in the fixed buffers, its KV is written to the extra page allocated for it in the KV pool, and its result is discarded.
  3. All the metadata lives in buffers fixed at capture time, which are the only addresses the graph knows; before each replay, this step's new values are copied into those buffers.
  4. Recording the largest first makes it allocate temporary memory from the pool at its largest requirement, after which the smaller graphs reuse the same memory pool and add nothing.

Files you will write: engine/graph.py, plus init_capture_graph, prepare_for_capture and prepare_for_replay in the attention backends.

Files so far: 75 (1 new, 1 modified, 0 brought in early)
minisgl/
├── __init__.py
├── core.py
├── env.py
├── __main__.py
├── shell.py
├── attention/
│   ├── __init__.py
│   ├── base.py
│   ├── utils.py
│   ├── torch_backend.py  ← modified in this chapter
│   ├── fi.py
│   ├── fa.py
├── distributed/  (3 files, unchanged)
├── engine/
│   ├── __init__.py
│   ├── config.py
│   ├── engine.py
│   ├── sample.py
│   ├── graph.py  ← new in this chapter
├── kernel/  (2 files, unchanged)
├── kvcache/  (5 files, unchanged)
├── layers/  (8 files, unchanged)
├── llm/  (2 files, unchanged)
├── message/  (5 files, unchanged)
├── models/  (7 files, unchanged)
├── moe/  (2 files, unchanged)
├── scheduler/  (9 files, unchanged)
├── server/  (4 files, unchanged)
├── tokenizer/  (4 files, unchanged)
├── utils/  (8 files, unchanged)
This step's main: examples/ch18_cuda_graph.py — it uses only the files above; python tools/steps.py check rebuilds this tree chapter by chapter and runs it.

Animation: capturing, padding and replaying a CUDA Graph (about 1.5 minutes, Chinese narration and subtitles)

Capture and replay

A CUDA Graph records "kernels plus arguments", and the pointers among those arguments are fixed at capture time. So:

  1. inputs live in fixed buffers. Before capture, the batch's input_ids, positions and out_loc point at slices of the buffers; before replay, the current batch's values are copied in;
  2. the attention metadata lives in fixed buffers too, which the attention backend takes care of: at capture the metadata points at the backend's own buffers, and before replay the new metadata is copied in;
  3. one graph per batch size, with a real batch padded up to the nearest recorded size.
capture (once per batch size) dummy batch bs dummy requests fixed buffers input_ids · positions · out_loc · logits the attention backend's captured buffers seq_lens · page_table (or the FlashInfer wrapper) points to torch.cuda.graph: records every kernel of model.forward() (pointers fixed) replay (every decode step) current batch 3 requests → padded to 4 copy_from copy the inputs into the fixed buffers prepare_for_replay copy the metadata into the captured buffers graph.replay() Miss one input and the graph silently reads the old value. The CPU emulation EmulatedGraph exists to catch exactly this.
capture and replay: every input goes through a fixed buffer

minisgl/engine/graph.py · GraphCaptureBuffer
@dataclass
class GraphCaptureBuffer:
    input_ids: torch.Tensor
    out_loc: torch.Tensor
    positions: torch.Tensor
    logits: torch.Tensor

    @classmethod
    def init(cls, bs: int, vocab_size: int, device: torch.device) -> GraphCaptureBuffer:
        return GraphCaptureBuffer(
            input_ids=torch.zeros(bs, dtype=torch.int32, device=device),
            out_loc=torch.zeros(bs, dtype=torch.int32, device=device),
            positions=torch.zeros(bs, dtype=torch.int32, device=device),
            logits=torch.empty(bs, vocab_size, dtype=torch.float32, device=device),
        )

    def set_batch(self, batch: Batch) -> None:
        """At capture: point the batch's input fields at the buffers."""
        s = slice(batch.padded_size)
        batch.input_ids, batch.out_loc, batch.positions = (
            self.input_ids[s], self.out_loc[s], self.positions[s])

    def copy_from(self, batch: Batch) -> None:
        """Before replay: copy the current batch's inputs into the buffers."""
        s = slice(batch.padded_size)
        self.input_ids[s] = batch.input_ids
        self.out_loc[s] = batch.out_loc
        self.positions[s] = batch.positions
minisgl/engine/graph.py · GraphRunner._capture_graphs
def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel) -> None:
    self.attn_backend.init_capture_graph(max_seq_len=max_seq_len, bs_list=self.graph_bs_list)
    synchronize(self.device)
    logger.info_rank0(f"Capturing graphs for batch sizes {self.graph_bs_list}")
    self.buffer = GraphCaptureBuffer.init(self.max_graph_bs, vocab_size, self.device)
    pool = None
    for bs in sorted(self.graph_bs_list, reverse=True):  # record the largest first so the smaller ones reuse its memory pool
        batch = Batch(reqs=[self.dummy_req] * bs, phase="decode")
        batch.padded_reqs = batch.reqs
        self.attn_backend.prepare_for_capture(batch)
        self.buffer.set_batch(batch)
        logits = self.buffer.logits[:bs]
        ctx = get_global_ctx()
        with ctx.forward_batch(batch):
            logits.copy_(model.forward())  # one warm-up: trigger the lazy initializations so they are not recorded into the graph
            if is_cuda(self.device):
                graph = torch.cuda.CUDAGraph()
                with torch.cuda.graph(graph, pool=pool, stream=self.stream):
                    logits.copy_(model.forward())
                if pool is None:
                    pool = graph.pool()
            else:
                graph = EmulatedGraph(lambda m=model, lg=logits: lg.copy_(m.forward()), batch)
        self.graph_map[bs] = graph
    logger.info_rank0(f"Free memory after capture: {get_free_memory(self.device) / 2**30:.2f} GiB")
The batch used for capture is all dummy requests (chapter 6: a dummy has its own row in the page table, pointing at its own page in the KV pool). Each batch size runs once normally first (a warm-up that triggers the libraries' lazy initialization so it does not get recorded into the graph), then runs again inside a torch.cuda.graph context to capture. Recording starts from the largest batch size, and the smaller graphs then reuse the same memory pool (pool), so memory use is set by the largest.

The batch sizes recorded by default are 1, 2, 4 and every multiple of 8 from 8 up to a ceiling set by memory (256 on cards above 80 GB, 160 otherwise).

Padding and replay

minisgl/engine/graph.py · GraphRunner.pad_batch
def pad_batch(self, batch: Batch) -> None:
    """Pad the batch up to the nearest recorded batch size with dummy requests, whose results are discarded."""
    padded = (next(bs for bs in self.graph_bs_list if bs >= batch.size)
              if self.can_use_cuda_graph(batch) else batch.size)
    batch.padded_reqs = batch.reqs + [self.dummy_req] * (padded - batch.size)
pad_batch is the first thing the scheduler calls when preparing a batch (chapter 7): when a graph can be used, dummy requests pad the batch up to the nearest recorded size. The padded dummies get positions, inputs and metadata too, their KV goes into the dummy page, and their logits are discarded (logits[:batch.size]).

minisgl/engine/graph.py · GraphRunner.replay
def replay(self, batch: Batch) -> torch.Tensor:
    assert self.can_use_cuda_graph(batch)
    self.buffer.copy_from(batch)
    self.attn_backend.prepare_for_replay(batch)
    self.graph_map[batch.padded_size].replay()  # type: ignore[attr-defined]
    return self.buffer.logits[: batch.size]
The attention backend's three hooks, with the reference backend as the example (the same as the official FlashAttention backend's approach):

minisgl/attention/torch_backend.py · TorchAttnBackend.prepare_for_capture
def prepare_for_capture(self, batch: Batch) -> None:
    bs = batch.size
    assert bs in self.capture_bs and self.capture is not None
    cap = self.capture
    batch.attn_metadata = TorchAttnMetadata(
        cu_seqlens_q=cap.cu_seqlens_q[: bs + 1],  # always [0, 1, ..., bs] during decode
        cache_seqlens=cap.seq_lens[:bs],
        page_table=cap.page_table[:bs, :],
        max_seqlen_q=1,
    )
minisgl/attention/torch_backend.py · TorchAttnBackend.prepare_for_replay
def prepare_for_replay(self, batch: Batch) -> None:
    metadata, bs = batch.attn_metadata, batch.padded_size
    assert isinstance(metadata, TorchAttnMetadata) and self.capture is not None
    table_len = metadata.page_table.size(1)
    self.capture.seq_lens[:bs].copy_(metadata.cache_seqlens)
    self.capture.page_table[:bs, :table_len].copy_(metadata.page_table)
The FlashInfer backend works differently: it offers wrappers made for CUDA Graph that take fixed-address buffers at construction, and plan copies the new metadata into them, so planning again before replay is all it takes (FlashInferBackend.prepare_for_replay, chapter 17).

The CPU emulation

There is no CUDA Graph on a CPU, but we still want to check that the "every input goes through a fixed buffer" contract holds, because that is the most common source of CUDA Graph bugs: one input is not copied in, the graph reads the old value from capture time (or from the previous step), and nothing reports an error.

minisgl/engine/graph.py · EmulatedGraph
class EmulatedGraph:
    """The CPU stand-in: the same interface as torch.cuda.CUDAGraph's replay."""

    def __init__(self, fn: Callable[[], None], batch: Batch) -> None:
        self.fn, self.batch = fn, batch

    def replay(self) -> None:
        ctx = get_global_ctx()
        saved, ctx._batch = ctx._batch, self.batch  # a real graph reads no Python objects, which this imitates
        try:
            self.fn()
        finally:
            ctx._batch = saved
During "capture" EmulatedGraph records the function to execute and the batch as it was at capture time. In that batch the input fields point at the buffers and the attention metadata points at the backend's buffers; at replay it runs the forward pass again with that batch rather than the current one. So it behaves like a real CUDA Graph: it can only see what is in the buffers.

examples/ch18_cuda_graph.py
"""Chapter 18: the CUDA Graph emulation on a CPU, with padding, replay, and what happens when one input is not copied."""

import torch
from minisgl.core import SamplingParams
from minisgl.engine.graph import GraphCaptureBuffer
from minisgl.env import ENV
from minisgl.llm import LLM

ENV.DISABLE_OVERLAP_SCHEDULING.value = True
PROMPTS = ["The capital of France is", "List three prime numbers:", "def fibonacci(n):"]
KW = dict(dtype=torch.float32, max_running_req=8, num_page_override=512, max_seq_len_override=128)


def run(cuda_graph_max_bs):
    llm = LLM("models/Qwen3-0.6B", cuda_graph_max_bs=cuda_graph_max_bs, **KW)
    runner = llm.engine.graph_runner
    log = []
    orig = runner.replay
    runner.replay = lambda b: (log.append(f"{b.size}->{b.padded_size}"), orig(b))[1]
    out = llm.generate(PROMPTS, SamplingParams(max_tokens=6, ignore_eos=True))
    sizes = sorted(runner.graph_map)
    llm.shutdown()
    return [o["token_ids"] for o in out], sizes, log


eager, _, _ = run(0)
graph, sizes, log = run(4)
print("recorded batch sizes:", sizes)
print("the replay of each decode step (real batch size -> padded):", log)
print("matches the output without a graph:", graph == eager)

orig_copy = GraphCaptureBuffer.copy_from


def buggy_copy_from(self, batch):  # positions was left out
    s = slice(batch.padded_size)
    self.input_ids[s] = batch.input_ids
    self.out_loc[s] = batch.out_loc


GraphCaptureBuffer.copy_from = buggy_copy_from
wrong, _, _ = run(4)
GraphCaptureBuffer.copy_from = orig_copy
print("matches the correct output after positions was not copied:", wrong == eager)
print("  correct:", eager[0])
print("  wrong:", wrong[0])
What it prints
recorded batch sizes: [1, 2, 4]
the replay of each decode step (real batch size -> padded): ['3->4', '3->4', '3->4', '3->4', '3->4']
matches the output without a graph: True
matches the correct output after positions was not copied: False
  correct: [12095, 13, 576, 6722, 315, 15344]
  wrong: [12095, 1103, 9736, 6686, 374, 264]
Three requests decode with every step padded to 4 and replayed, and the output is identical to running without a graph. Then a bug is introduced on purpose: positions is not copied before replay. RoPE inside the graph reads the positions left in the buffer from capture time (all zeros) and the output is wrong from the second token on, with no error anywhere. It would look exactly the same on a real GPU.

Difference from upstream

EmulatedGraph is ours: on a CPU with cuda_graph_max_bs set explicitly, it stands in for torch.cuda.CUDAGraph (graphs are off by default on a CPU). The capture loop, the padding and the replay logic are the same as upstream's.

The reference backend cannot be captured on a real GPU

The torch reference backend's forward first .tolist()s cu_seqlens_q and each request's KV length back to the CPU, then slices by those numbers. A device → host copy is forbidden during capture (Cannot copy between CPU and CUDA tensors during CUDA graph capture), and even if you got around it, the recorded lengths would be frozen at their capture-time values and replay would be wrong. So the engine turns CUDA Graph off automatically when the torch backend is used on CUDA — that backend exists to check answers, not to be fast. On a CPU it runs through EmulatedGraph (which just re-runs the function), which is not subject to this, so this chapter's verification works as before.

The official implementation

Tests

tests/test_ch18_cuda_graph.py · test_forgetting_to_copy_an_input_breaks_replay
def test_forgetting_to_copy_an_input_breaks_replay(monkeypatch):
    """The counter-example: positions is not copied before replay (RoPE in the graph still reads the capture-time positions), so the result is wrong."""
    ENV.DISABLE_OVERLAP_SCHEDULING.value = True

    def buggy_copy_from(self, batch: Batch) -> None:
        s = slice(batch.padded_size)
        self.input_ids[s] = batch.input_ids
        self.out_loc[s] = batch.out_loc  # positions was left out

    monkeypatch.setattr(GraphCaptureBuffer, "copy_from", buggy_copy_from)
    llm = build_llm(QWEN3, cuda_graph_max_bs=4)
    out = llm.generate(PROMPTS[:3], greedy(6))
    assert [o["token_ids"] for o in out] != hf_greedy(QWEN3, PROMPTS[:3], 6)
    llm.shutdown()
The same file also has end-to-end tests for the reference, FlashInfer and FlashAttention backends under the emulated graph (every decode step replays, and the output matches Hugging Face), plus the batch-size list computation.

How to explain it

On CUDA Graph: a decode step is hundreds of small kernels where the CPU's launch cost exceeds the GPU's compute, so CUDA Graph records a whole decode forward pass and replays it in one go; a prefill's token count varies widely and it is compute-heavy anyway, so it neither needs nor suits recording. What is recorded are fixed addresses, so every input and all the attention metadata must live in fixed buffers with the new values copied in before replay (miss one and it fails silently). One graph is recorded per batch size and a real batch is padded with dummy requests to the nearest one, with the dummies' KV written to a page reserved for them. The largest batch size is recorded first so the memory pool it allocates is reused by the later graphs. The gain is largest with small models and small batches.

Exercises

  1. With 5 requests in the batch and [1, 2, 4, 8] recorded, what is it padded to? How much extra compute does that cost? How would you choose the list of batch sizes to balance "number of graphs (memory)" against "padding waste"?
  2. Could sampling be recorded into the graph too? Why does mini-sglang not do so?
  3. Under tensor parallelism, can the all-reduce inside the graph be recorded? What has to be watched?
Answers
  1. Padded to 8, computing 3 extra dummy requests, about 60% more. At small batch sizes every size is common, so the steps should be small (1, 2, 4, 8); at large batch sizes the kernels are already well filled and the padding waste is a smaller fraction, so the steps can be wide (every 8 or 16). That is the thinking behind the official default list.
  2. It could, but the sampling parameters (temperature, top-k, top-p) differ per request and the random-number state is needed too, all of which would have to go into fixed buffers; and greedy sampling is just one argmax, which costs almost nothing.
  3. Yes, NCCL supports capturing collectives in a CUDA Graph. Every rank has to capture and replay the same graphs in the same order (chapter 14's "every rank identical" again), and the graphs must be destroyed before the NCCL communicator.

Summary

  • CUDA Graph records a whole decode forward pass so one replay launches every kernel, removing the CPU's launch cost at small batch sizes.
  • Every input and all the attention metadata must live in fixed buffers and be copied in before replay; there is one graph per batch size and a real batch is padded with dummy requests.
  • EmulatedGraph emulates the "can only see the buffers" behaviour on a CPU: miss one input copy and the output is wrong, silently.

Exercises for this chapter (1)

Write the code and run the tests in your browser, or judge locally with python practice/judge.py. The exercises are in Chinese for now.

Found a mistake or a better way to say it? Edit this page on GitHubReport an issue