Piecewise CUDA Graph

​

Motivation

​

Usage

python3 -m sglang.launch_server \
    --model-path meta-llama/Llama-3.1-8B-Instruct

​

Disable PCG

python3 -m sglang.launch_server \
    --model-path meta-llama/Llama-3.1-8B-Instruct \
    --disable-piecewise-cuda-graph

​

Custom capture sizes

python3 -m sglang.launch_server \
    --model-path meta-llama/Llama-3.1-8B-Instruct \
    --piecewise-cuda-graph-max-tokens 2048

​

Server Args

​

Bug Report

​

For Users

Piecewise CUDA Graph is enabled by default as an experimental feature.
To work around this error, add --disable-piecewise-cuda-graph to your launch command.
Please report this issue at https://github.com/sgl-project/sglang/issues/new/choose

1. The full error traceback
2. Model name and quantization method
3. Launch command with all arguments
4. GPU type and driver version

​

For Developers

from sglang.srt.utils.custom_op import register_custom_op

# Inplace operator (no return value)
@register_custom_op(mutates_args=["output_q", "output_s"])
def per_token_group_quant_8bit(
    input: torch.Tensor,
    output_q: torch.Tensor,
    output_s: torch.Tensor,
) -> None:
    # kernel implementation ...

# out_shape indicates which argument has the same shape as the output
@register_custom_op(mutates_args=["x"], out_shape=0)
def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    return x.add_(y)

​

How it works

​

Torch compile backend

model.forward wrapper
→ torch.compile(..., backend=SGLangBackend)
→ FX graph
→ split_graph() at registered split ops
→ split_gm (top-level graph that chains the pieces)
→ replace capturable submodules with CUDAPiecewiseBackend
→ runtime dispatch: eager split ops + per-piece capture/replay

• Install: install_torch_compiled() replaces model.forward with a wrapper function. When is_in_piecewise_cuda_graph() returns True, the wrapper dispatches to the compiled callable; otherwise it falls back to the original forward. The first invocation through this path triggers Dynamo tracing and graph compilation — CUDA graph replay only happens after the capture phase completes.
• Split: When torch.compile traces the model, SGLangBackend receives the FX graph and calls split_graph(). Ops listed in CompilationConfig.split_ops are treated as split points, so the graph is cut at each one. These split-op submodules are left to run eagerly at runtime, while the surrounding submodules are compiled and wrapped by CUDAPiecewiseBackend. The result is a top-level “stitching graph” (split_gm) with children such as submod_0, submod_1, … interleaving capturable subgraphs and eager split-op submodules.
• Replace: PiecewiseCompileInterpreter iterates over each capturable submodule in split_gm, compiles it for general (dynamic) shapes, and replaces it in-place with a CUDAPiecewiseBackend instance. Split-op submodules (e.g., attention, all-reduce) are left as-is and run eagerly at runtime.
• Dispatch: At runtime, calling split_gm executes the stitching graph, which calls each submodule in order. Split-op submodules run eagerly. Each CUDAPiecewiseBackend submodule goes through three phases:
  • Compile warmup — runs the general-shape compiled path.
  • Capture — for each capture size, runs one warmup pass then records a CUDA graph.
  • Steady-state replay — replays the captured CUDA graph for each forward pass.

​

Piecewise cuda graph runner

• Compile — Warms up JIT kernels with a dummy forward pass, then wraps the model with torch.compile, triggering Dynamo tracing to split the FX graph and create CUDAPiecewiseBackend instances for each subgraph piece.
• Capture — Iterates over capture sizes in reverse order (largest first). For each size, runs the forward pass twice (one warmup, one CUDA graph capture).
• Replay — At runtime, finds the smallest captured size >= actual token count via binary search, copies inputs into static buffers with zero-padding, replays the captured CUDA graphs, and slices outputs back to the actual token count.

​

Memory optimization

​

Shape configuration

​

Compatibility

• Disabled model architectures (e.g., DeepseekV32ForCausalLM)
• Speculative decoding
• DP attention
• Pipeline parallelism (pp_size > 1)
• Non-CUDA hardware (AMD ROCm, Ascend NPU)
• MoE A2A backend
• LoRA
• Multimodal / VLM models
• DLLM (diffusion LLM)
• Deterministic inference
• PD disaggregation
• Expert distribution recorder / EPLB

​

Code Reference