Breakable CUDA Graph in SGLang: 5× faster graph builds, 1.93× faster prefill

Background

To start, the basic idea behind CUDA Graph: an inference step is not a single kernel but a sequence of many GPU operations. CUDA Graph records the GPU workload once and then replays all the kernels at once, which reduces the CPU overhead of launching small kernels one by one.

While maintaining SGLang, though, we found that using CUDA Graph well takes much more than recording every operation once. Plenty of operations are not graph-compatible. These operations sit in the middle of the forward, and a single incompatible one is enough to keep a large stretch of it off CUDA Graph.

The natural idea, then, is to cut the CUDA Graph before those incompatible operations and resume afterward. The diagram below compares three ways of executing the same forward:

The intuition of BCG is simple; the engineering is not. Our first answer went through torch.compile, referred to as tc_piecewise in the rest of this article: let the compiler trace the full forward, then split the graph at registered points. It got the job done, but compilation time grew with model complexity, and it was awkward to adapt to our own custom optimized kernels. Here is that route in full:

Since capture already executes operators in order, why should a compiler be the one to find the break points? Following that idea, we built Breakable CUDA Graph: eager breaks inserted directly while capture is in progress, with no compiler needed to understand the forward. It reaches the same segmented execution in roughly a quarter of the code (521 versus 1,771 lines); it builds prefill graphs 3.8–5.2× faster because no compilation is involved; and measured on prefill alone it runs 1.70× over eager. BCG first landed on the decode path (#19102) and was later extended to prefill (#22218), where it is now the default.

The rest of this post walks through how SGLang implements BCG and full CUDA Graph for prefill.

Breakable CUDA Graph

Design and mechanism

BCG lets developers mark the incompatible region directly with @eager_on_graph. During capture, the current graph segment is closed when execution reaches the marked function, the function runs eagerly, and capture resumes afterward in a new segment. The forward becomes a sequence of CUDA Graph segments separated by eager regions.

During capture the marked function runs once between the two segments, and the tensor it returns is retained as a persistent boundary buffer whose device address stays fixed; the following segment is then captured against that address. On replay the eager function runs normally and returns a fresh tensor, which BCG copies into the retained buffer. The diagram below walks through one capture and two replays:

Benefits

Faster startup. Compared with the traditional compiler-based piecewise graph, setup is dominated by compilation rather than capture: torch.compile accounts for 78–86% of the time spent preparing prefill graphs, and it grows with model complexity. BCG eliminates that 78–86% share entirely: prefill graphs build 3.8–5.2× faster (42 captured shapes per configuration, TP4, 4×GB300).

Broader compatibility. SGLang relies heavily on custom CUDA, Triton, and JIT-compiled kernels, which are not native PyTorch operators. Making them visible to torch.compile usually means wrapping them through torch.library and providing a fake implementation for tracing. More importantly, where the breaks can go is constrained as well: inputs and outputs crossing a boundary have to be types the compiler can represent, and when the natural cut involves unusual runtime state or return types, you end up moving the cut or widening the eager region just to hand the compiler an interface it understands. Keeping CUDA Graph working increasingly felt like maintaining a torch.compile integration. BCG lifts that constraint at eager breaks, and there is no separate compiler to maintain. The same segmented execution takes 521 lines in BCG, versus 1,771 in the old approach.

Easier debugging. Once captured, a CUDA Graph is a black box at replay: ordinary Python does not execute inside it, which makes prints, assertions, and step-by-step inspection difficult. BCG naturally keeps eager regions where ordinary Python still runs on every replay. SGLang extends the idea with --debug-cuda-graph, which wraps the whole decode forward in a single eager break so the model executes eagerly. That gives a useful debugging boundary: if the problem persists, it most likely lives in the model or ouside the graph; if it disappears, capture itself is the first suspect.

BCG has also been adopted by SGLang's diffusion stack (#27436): on a single B200, Qwen-Image at 512×512 drops from 6.48 s to 2.45 s end-to-end, and Z-Image from 1.231 s to 0.662 s.

Going further: full CUDA Graph for prefill

As covered earlier, BCG buys compatibility by breaking the CUDA Graph: incompatible operations stay outside the graphs, at the cost of a short eager stretch at every break. Full CUDA Graph, by contrast, captures the whole forward as one graph, with no eager regions and the fewest replay-time launches. This section pushes the same capture onto prefill.

Making prefill static

Inherited technical challenges. Unlike a decode AR step, whose size is fixed, prefill is hard to capture because two things vary: the length of a prefill, and the number of requests its tokens belong to. CUDA Graph requires both dimensions to stay fixed. On top of that, some attention backends depend on runtime metadata. Full CUDA Graph was therefore long difficult to apply to the prefill stage, which is one of the main reasons the prefill path adopted BCG. More recently, #27988 redesigned how request slots and attention metadata are represented, so supported attention backends no longer have to run outside the graph.

Solution. SGLang does two things.

1. It fixes the token dimension with token buckets: a live batch is padded to the nearest captured token count.
2. It handles the request dimension separately: each graph reserves a fixed number of request slots, live requests take the first ones, and the rest are rewritten as zero-length sentinels.

A batch carrying more requests than there are reserved slots falls back to eager. Sentinel and attention metadata both have to be rebuilt outside the graph on every replay, so full prefill capture currently works only on FlashAttention (fa4) and FlashInfer, the two backends that prepare extend-mode metadata this way. Full prefill capture is also still experimental and has to be enabled explicitly; the engine warns that full is experimental and points to BCG or tc_piecewise for production workloads. Broadening backend support and tuning the bucket and slot choices are still ahead of us.

As shown below, the diagram covers how we pad both the token and the request dimension in different scenarios:

Prefill benchmark

We measured prefill on its own: a fixed input length with a single output token, one request at a time, decode graphs disabled in every arm, on gpt-oss-120b (TP4, 4×GB300), where all four paths run. Full capture is 1.93× faster than eager, BCG 1.70×, and tc_piecewise 1.45×, so BCG is also 17% faster than the compiler-based backend at replay, not only at build time. The gap comes from what each one does per forward: BCG replays its recorded segments directly, while tc_piecewise calls back into the compiled callable every time, paying Torch Dynamo's guard checks and dispatch before its own captured pieces run. On GLM-5.2 only BCG can capture at all — tc_piecewise cannot trace the forward and full capture has no path for its sparse attention — and it is 1.60× over eager there. Every curve is flat across a 32× range in prompt length, which is the signature of launch overhead rather than compute:

Memory footprint of CUDA Graphs

Both of the things above affect memory: BCG cuts a graph capture into segments, and full capture brings larger prefill shapes in. CUDA Graph memory is resident: replay requires intermediate tensors to keep their addresses, so whatever is allocated at capture time stays until the process exits. That leaves two questions.

1. With BCG cutting the forward into segments, resident memory must not grow with the number of segments.
2. For resident graph memory to stand in for the worst-case activation peak, capture has to cover the largest prefill shape.

Reuse inside a segmented capture

Imagine one prefill forward that BCG breaks into N segments. If every segment's intermediates occupy and retain their own memory, graph memory could easily multiply (N segments × the number of captured shapes). The fix is simple: a CUDA Graph replays in order and only one segment runs at a time, so there is no need to keep every segment's intermediates alive at once. A later segment can overwrite what an earlier one left behind and share the same memory. BCG avoids that growth through three forms of reuse.

• Segments share one memory pool. Every segment of a captured shape uses the same CUDA Graph pool, so intermediate storage can be reused rather than pinned separately by each segment.
• Weak references at eager breaks. When the storage of a tensor passed into a break is already owned by the graph pool, the tensor is held through a weak reference, avoiding unnecessary Python references that would extend its lifetime. The weak-reference technique comes from vLLM PR #9724, which introduced it so that captured graphs could share an output buffer instead of each graph pinning one of its own.
• Capture sizes share one output buffer. All capture sizes share a single maximum-sized output buffer, with each shape slicing only the rows it needs, instead of allocating an output buffer per shape.

With these reuse mechanisms, even a large capture table remains modest: 42 shapes across a 78-layer MoE add 2.4 GB of graph memory on GLM-5.2.

Capture through the chunked-prefill size

Capturing a shape swaps its transient peak for resident memory, and the swap only pays off for shapes that actually replay. If capture stops at some size n, prefills larger than n still fall back to eager with the peak intact, while the resident memory for the smaller shapes has already been spent, and total memory ends up above the no-graph baseline. What matters is always the capture ceiling, not how many shapes are captured.

chunked_prefill_size is the upper bound on a single prefill forward, so capturing all the way up to it lets the largest prefill replay a graph and removes the worst-case activation peak. On gpt-oss-120b that peak drops from 0.56 GB to 0.001 GB, and on GLM-5.2 from 1.55 GB to 0.35 GB, putting total memory 0.51 GB and 1.10 GB below the no-graph baseline. As the figure shows, the amount is small against a total footprint of a few hundred gigabytes, but it means CUDA Graph is saving memory rather than costing it. The other benefit is predictable memory: a peak that used to move with the workload becomes a fixed allocation established at capture time, which the engine can account for up front instead of reserving headroom for large prefills:

Acknowledgments

This work was a collaboration between the SGLang team and the Meta team. SGLang: Yuwei An*, Cheng Wan, Xiaoyu Zhang, Mick Qian, Baizhou Zhang, Yusheng Su, Ke Bao; Meta: Shiyang Chen*, Lianmin Zheng (* equal contribution). We also thank the NVIDIA, AMD, Thinking Machines Lab, and Meta PyTorch teams for their help along the way.

Further reading

• LMSYS Blog: Advanced CUDA Graph Techniques in SGLang
• sglang#19102: first BCG implementation
• sglang#22218: BCG for prefill
• sglang#23906: runner/backend refactor
• sglang#27988: full CUDA Graph for prefill
• NVIDIA: Getting Started with CUDA Graphs

The visualized tutorial was written by Yichi Zhang and the SGLang Team.