SGLang and Meta Cut LLM Prefill Time 1.7x With Breakable CUDA Graph

SGLang and Meta rebuilt CUDA Graph support around a runner/backend split, delivering faster startup, broader kernel compatibility, and up to 1.93x prefill speedup.

·
·
SGLang and Meta Cut LLM Prefill Time 1.7x With Breakable CUDA Graph
  • SGLang and Meta refactored CUDA Graph support with a clean runner/backend split for reusable capture strategies.
  • Breakable CUDA Graph is now the default prefill backend, dropping torch.compile from the capture path entirely.
  • BCG builds prefill graphs 3.8-5.2x faster and uses roughly a quarter of the code versus TC piecewise.
  • Full CUDA Graph for prefill hits 1.93x speedup over eager on gpt-oss-120b via token bucket and request slot padding.
  • Memory reuse across shapes and segments keeps a 42-shape capture table at just 2.4 GB on GLM-5.2.
  • Capturing through chunked-prefill size eliminates the worst activation peak, making memory usage predictable.

Kernel launch overhead has been quietly taxing LLM inference for years. Every forward pass fires off hundreds of small GPU operations, and each launch from the CPU adds latency that compounds on latency-sensitive workloads. CUDA Graph fixes this by recording the whole sequence once and replaying it, but getting it to work inside a real serving engine is another story. The SGLang team, together with Meta, just published a deep dive on how they restructured CUDA Graph support and introduced several new capture techniques along the way.

Why one graph is never enough

An inference step is not a single kernel but a sequence of many GPU operations, and in modern LLM serving engines, repeatedly launching these operations from the CPU can introduce noticeable overhead, especially for latency-sensitive workloads. The catch is that CUDA Graph traditionally requires the entire captured region to be graph-compatible. Prefill attention often depends on runtime metadata, collectives may need host coordination, and features like LoRA or MoE routing add operations that resist capture. A single incompatible op can block CUDA Graph from covering most of the forward.

Before the refactor, SGLang had separate CUDA Graph runners for decode, prefill, and speculative decoding, each with its own overlapping logic. The new design splits responsibilities into two layers: a runner manages the execution-specific state needed for capture and replay (captured shapes, static input buffers, attention metadata, padding), while a backend determines how that execution is captured, whether as one full graph, a sequence of breakable segments, or compiler-generated pieces.

Breaking the graph on purpose

The headline technique is Breakable CUDA Graph (BCG), which SGLang claims as an originally proposed serving mechanism. Instead of tracing the full forward with a compiler and splitting the resulting graph, BCG inserts eager breaks directly during capture. Developers mark incompatible functions with an @eager_on_graph decorator, and the runtime does the rest.

Here is what actually happens under the hood: 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 tricky part is keeping tensor addresses stable across replays. The tensor returned by the eager function is retained as a persistent boundary buffer so its device address stays fixed, and on every replay the eager function runs normally and returns a fresh tensor which BCG copies into the retained buffer, so the next segment reads the updated value from the address it was originally captured with.

The result is a much lighter alternative to the torch.compile-based piecewise backend that used to fill this role. BCG reaches the same segmented execution as the torch.compile-based piecewise backend in roughly a quarter of the code (521 versus 1,771 lines), builds prefill graphs 3.8 to 5.2 times faster because no compilation is involved, and has broader coverage for complex functionality naturally.

The numbers that matter

Benchmarked on gpt-oss-120b (TP4 on 4x GB300) with prefill isolated, the three capture strategies stack up like this against eager execution:

BackendPrefill speedup vs eagerNotes
Full CUDA Graph1.93xExperimental, needs FA4 or FlashInfer
Breakable CUDA Graph1.70xNew default
TC piecewise (torch.compile)1.45xPays Dynamo guard checks per replay

BCG is 17% faster than the compiler-based backend at replay, not only at build time, because 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. Build-time savings are even more dramatic: torch.compile accounts for 78 to 86 percent of the time spent preparing prefill graphs, reaching 90 seconds on a 235B MoE and 158 seconds on GLM-5.2. BCG skips that phase entirely.

Making prefill fully static

Full CUDA Graph is easy for decode (one token per request, batch size is the only varying dimension) but hard for prefill, which varies in both total tokens and request count. SGLang's solution uses two kinds of padding:

  • Token padding: Pad the live batch to the nearest captured token bucket.
  • Request slot padding: Reserve a fixed number of request slots per graph. Unused ones become zero-length sentinels with zero sequence and extend lengths.
  • Metadata rewrite: Sentinel metadata is rewritten on every replay because the captured graph still reads the entire request table.

The two padding types have very different costs. Padded tokens become actual rows in the captured batch and therefore pass through dense projections as part of the same GEMMs, so they cost real compute. Empty request slots, on the other hand, contribute essentially no attention work because in FlashAttention's variable-length scheduler, work is derived from each sequence's actual length rather than assigning a fixed amount of computation to every request slot.

Memory that actually shrinks

Segmented capture could easily explode memory since each shape contains multiple segments with intermediates that stay alive. BCG avoids this through a shared CUDA Graph pool across segments, weak references at eager breaks (borrowed from a technique vLLM introduced so captured graphs could share output buffers), and a single maximum-sized output buffer sliced per shape. Even a large capture table stays modest: 42 shapes across a 78-layer MoE add 2.4 GB of graph memory on GLM-5.2.

The bigger insight is around the capture ceiling. If your capture ladder stops below the maximum prefill size, the largest prefill still runs eagerly and keeps its activation peak while you also pay for all the resident graphs below it. Capturing through the chunked-prefill size means the activation peak collapses (from 0.56 GB to 0.001 GB on gpt-oss-120b, and from 1.55 GB to 0.35 GB on GLM-5.2), turning a workload-dependent spike into a fixed allocation the engine can plan around.

What to reach for and when

BCG is now the default prefill backend in SGLang, and it also powers the diffusion stack, where Qwen-Image at 512x512 on a single B200 improves from 6.48s to 2.45s end-to-end latency, and Z-Image improves from 1.231s to 0.662s. Full CUDA Graph for prefill is still experimental and limited to FlashAttention (fa4) and FlashInfer, while TC piecewise remains available for platforms where BCG has not been validated. The broader lesson: BCG removes launch overhead but does not reduce model FLOPs or make compute-bound kernels cheaper, so its advantage is largest when exposed launch gaps are a meaningful fraction of execution time.

Trending
  • No trending articles

Comments

avatar

Next Reads