chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Piecewise Cuda Graph set default (#16331)

Browse Source
This commit is contained in:
Yuwei An
2026-03-02 23:18:07 +08:00
committed by GitHub
parent 468e3dc56b
commit c64274c746
34 changed files with 505 additions and 143 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/compilation/compile.py
+2 -14
View File
@@ -1,31 +1,19 @@
import contextvars
import inspect
import logging
import os
import sys
import types
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Any, Callable, Optional, Union
import torch
from sglang.srt.compilation.compilation_config import CompilationConfig
from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
from sglang.srt.utils.common import rank0_log
logger = logging.getLogger(__name__)
_COMPILE_ENABLED = contextvars.ContextVar("_COMPILE_ENABLED", default=False)
@contextmanager
def set_compiled(enabled: bool = True):
token = _COMPILE_ENABLED.set(enabled)
try:
yield
finally:
_COMPILE_ENABLED.reset(token)
@dataclass
class IntermediateTensors:
@@ -200,7 +188,7 @@ def install_torch_compiled(
state["compiled_callable"] = compiled_callable
def trampoline(self, *args, **kwargs):
use_compiled = _COMPILE_ENABLED.get()
use_compiled = is_in_piecewise_cuda_graph()
if use_compiled:
if not state["compiled"]:
_ensure_compiled(self, *args, **kwargs)
python/sglang/srt/compilation/piecewise_context_manager.py
+23 -5
View File
@@ -1,11 +1,15 @@
from __future__ import annotations
import logging
from contextlib import contextmanager
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, List, Optional
import torch
logger = logging.getLogger(__name__)
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@@ -38,10 +42,17 @@ def enable_piecewise_cuda_graph_compile():
def enable_piecewise_cuda_graph():
global _in_piecewise_cuda_graph
_in_piecewise_cuda_graph = True
yield
_in_piecewise_cuda_graph = False
try:
yield
except Exception as e:
logger.error(
"Piecewise CUDA Graph failed with error: %s\n%s",
e,
PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG,
)
raise
finally:
_in_piecewise_cuda_graph = False
@contextmanager
@@ -56,7 +67,7 @@ def set_pcg_capture_stream(stream: torch.cuda.Stream):
class ForwardContext:
def __init__(self):
self.forward_batch = None
self.attention_layer = None
self.attention_layers = None
self.quant_config = None
self.moe_layers = None
self.moe_fusions = None
@@ -105,3 +116,10 @@ def set_forward_context(
yield
finally:
_forward_context = None
PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG = (
"Piecewise CUDA Graph is enabled by default as an experimental feature.\n"
"To work around this error, add --disable-piecewise-cuda-graph to your launch command.\n"
"Please report this issue at https://github.com/sgl-project/sglang/issues/new/choose"
)
python/sglang/srt/configs/model_config.py
+19
View File
@@ -193,6 +193,10 @@ class ModelConfig:
self.is_local_attention_model = is_local_attention_model(
self.hf_config.architectures
)
self.is_piecewise_cuda_graph_disabled_model = (
is_piecewise_cuda_graph_disabled_model(self.hf_config.architectures)
or is_deepseek_nsa(self.hf_text_config)
)
self.dtype = _get_and_verify_dtype(self.hf_text_config, dtype)
# Derive context length and model shapes
@@ -1274,6 +1278,14 @@ multimodal_model_archs = [
"KimiK25ForConditionalGeneration",
]
piecewise_cuda_graph_disabled_model_archs = [
"DeepseekV32ForCausalLM",
"Qwen3NextForCausalLM",
"GlmMoeDsaForCausalLM",
"BailingMoeV2_5ForCausalLM",
"LLaDAModelLM",
]
if external_mm_model_arch := envs.SGLANG_EXTERNAL_MM_MODEL_ARCH.get():
multimodal_model_archs.append(external_mm_model_arch)
@@ -1329,6 +1341,13 @@ def is_multimodal_chunked_prefill_supported(model_architectures: List[str]):
return True
def is_piecewise_cuda_graph_disabled_model(model_architectures: List[str]):
return any(
arch in piecewise_cuda_graph_disabled_model_archs
for arch in model_architectures
)
def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float:
if scale <= 1:
return 1.0
python/sglang/srt/layers/attention/fla/layernorm_gated.py
+1 -1
View File
@@ -180,7 +180,7 @@ def calc_rows_per_block(M: int, device: torch.device) -> int:
# When piecewise cuda graph is enabled, use a constant value to avoid
# torch.compile creating guards on the dynamic batch dimension.
try:
if get_global_server_args().enable_piecewise_cuda_graph:
if not get_global_server_args().disable_piecewise_cuda_graph:
return MAX_ROWS_PER_BLOCK
except ValueError:
# Global server args not initialized (e.g., in unit tests)
python/sglang/srt/layers/attention/flashinfer_backend.py
+1 -1
View File
@@ -243,7 +243,7 @@ class FlashInferAttnBackend(AttentionBackend):
if is_sm100_supported():
# Disable CUTLASS backend when piecewise cuda graph is enabled
# due to TMA descriptor initialization issues on B200
if model_runner.server_args.enable_piecewise_cuda_graph:
if not model_runner.server_args.disable_piecewise_cuda_graph:
logger.warning(
"CUTLASS backend is disabled when piecewise cuda graph is enabled "
"due to TMA descriptor initialization issues on B200. "
python/sglang/srt/layers/communicator.py
+1 -1
View File
@@ -198,7 +198,7 @@ class AttnTpContext:
and not is_dp_attention_enabled()
and get_moe_a2a_backend().is_none()
and not enable_moe_dense_fully_dp()
and not get_global_server_args().enable_piecewise_cuda_graph
and get_global_server_args().disable_piecewise_cuda_graph
and get_global_server_args().speculative_algorithm != "EAGLE3"
)
if get_global_server_args().enable_attn_tp_input_scattered:
python/sglang/srt/layers/moe/topk.py
+31 -1
View File
@@ -79,7 +79,37 @@ if _is_cuda:
from sgl_kernel import moe_fused_gate
try:
from flashinfer.fused_moe import fused_topk_deepseek
from flashinfer.fused_moe import fused_topk_deepseek as _fused_topk_deepseek
from sglang.srt.utils.custom_op import register_custom_op
@register_custom_op(
op_name="fused_topk_deepseek",
mutates_args=["topk_weights", "topk_ids"],
)
def fused_topk_deepseek(
gating_output: torch.Tensor,
correction_bias: torch.Tensor,
num_expert_group: int,
topk_group: int,
topk: int,
scaling_factor: float,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
renormalize: bool,
) -> None:
_fused_topk_deepseek(
gating_output,
correction_bias,
num_expert_group,
topk_group,
topk,
scaling_factor,
topk_weights,
topk_ids,
renormalize,
)
except ImportError:
fused_topk_deepseek = None
python/sglang/srt/layers/quantization/awq.py
+7 -12
View File
@@ -65,7 +65,14 @@ if _is_cuda:
awq_marlin_moe_repack,
awq_marlin_repack,
)
from sglang.srt.utils.custom_op import register_custom_op_from_extern
awq_dequantize = register_custom_op_from_extern(
awq_dequantize,
fake_impl=lambda qweight, scales, qzeros: qweight.new_empty(
qweight.shape[:-1] + (qweight.shape[-1] * 8,), dtype=scales.dtype
),
)
elif _is_hip:
from sglang.srt.layers.quantization.awq_triton import (
@@ -952,18 +959,6 @@ class AWQMoEAscendMethod(AWQMoEMethod):
# Register fake implementations for torch.compile support
if _is_cuda:
@register_fake_if_exists("sgl_kernel::awq_dequantize")
def _(
qweight,
scales,
qzeros,
ch_axis,
group_size,
num_bits,
):
out_shape = qweight.shape[:-1] + (qweight.shape[-1] * 32 // num_bits,)
return qweight.new_empty(out_shape, dtype=scales.dtype)
@register_fake_if_exists("sgl_kernel::awq_marlin_repack")
def _(b_q_weight, size_k, size_n, num_bits):
return b_q_weight.new_empty(
python/sglang/srt/layers/quantization/fp8_utils.py
+29 -2
View File
@@ -181,7 +181,35 @@ def _check_cutlass_block_fp8_hardware_support() -> bool:
if is_blackwell_supported() and is_flashinfer_available():
from flashinfer.gemm import gemm_fp8_nt_groupwise
from flashinfer.gemm import gemm_fp8_nt_groupwise as _raw_gemm_fp8_nt_groupwise
from sglang.srt.utils.custom_op import register_custom_op
# Wrap gemm_fp8_nt_groupwise as a custom op so torch.compile does not trace
# into flashinfer's JIT compilation code (pathlib/cubin_loader ops).
@register_custom_op(
op_name="flashinfer_gemm_fp8_nt_groupwise",
mutates_args=[],
fake_impl=lambda q_input, weight, x_scale, weight_scale, out_dtype: (
q_input.new_empty((q_input.shape[0], weight.shape[0]), dtype=out_dtype)
),
)
def gemm_fp8_nt_groupwise(
q_input: torch.Tensor,
weight: torch.Tensor,
x_scale: torch.Tensor,
weight_scale: torch.Tensor,
out_dtype: torch.dtype,
) -> torch.Tensor:
return _raw_gemm_fp8_nt_groupwise(
q_input,
weight,
x_scale,
weight_scale,
out_dtype=out_dtype,
backend="trtllm",
)
if is_sm90_supported() and is_flashinfer_available():
# FlashInfer SM90 DeepGEMM with automatic swapAB optimization for small M
@@ -350,7 +378,6 @@ def flashinfer_gemm_w8a8_block_fp8_linear_with_fallback(
x_scale,
weight_scale,
out_dtype=input_2d.dtype,
backend="trtllm",
)
if bias is not None:
python/sglang/srt/managers/schedule_batch.py
+1 -1
View File
@@ -2270,7 +2270,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
if (
self.forward_mode.is_decode()
and server_args.enable_piecewise_cuda_graph
and not server_args.disable_piecewise_cuda_graph
and not self.tree_cache.is_chunk_cache()
):
return
python/sglang/srt/model_executor/model_runner.py
+18 -31
View File
@@ -107,7 +107,6 @@ from sglang.srt.layers.moe.routed_experts_capturer import (
get_global_experts_capturer,
set_global_experts_capturer,
)
from sglang.srt.layers.moe.utils import get_moe_a2a_backend
from sglang.srt.layers.pooler import EmbeddingPoolerOutput
from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype
from sglang.srt.layers.sampler import create_sampler
@@ -1656,32 +1655,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
or self.hybrid_lightning_config
)
def can_run_piecewise_cuda_graph(self):
if self.is_draft_worker:
return False
if self.server_args.enable_torch_compile:
log_info_on_rank0(
logger,
"Disable piecewise CUDA graph because piecewise_cuda_graph has conflict with torch compile",
)
return False
if self.pp_size > 1:
# TODO(yuwei): support PP
log_info_on_rank0(
logger,
"Disable piecewise CUDA graph because piecewise_cuda_graph does not support PP",
)
return False
if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake():
# TODO(yuwei): fix the compilation errors for MOE A2A backend
log_info_on_rank0(
logger,
"Disable piecewise CUDA graph due to existing compilation errors",
)
return False
return True
def configure_kv_cache_dtype(self):
if self.server_args.kv_cache_dtype == "auto":
quant_config = getattr(self.model, "quant_config", None)
@@ -2183,10 +2156,24 @@ class ModelRunner(ModelRunnerKVCacheMixin):
"""Initialize piecewise CUDA graph runner."""
self.piecewise_cuda_graph_runner = None
if (
not self.server_args.enable_piecewise_cuda_graph
or not self.can_run_piecewise_cuda_graph()
):
if self.server_args.disable_piecewise_cuda_graph:
logger.info(
"Disable piecewise CUDA graph because --disable-piecewise-cuda-graph is set"
)
return
# Disable piecewise CUDA graph for non-language models
if not hasattr(self.model, "model"):
logger.warning(
"Disable piecewise CUDA graph because the model is not a language model"
)
return
# Disable piecewise CUDA graph for non capture size
if not self.server_args.piecewise_cuda_graph_tokens:
logger.warning(
"Disable piecewise CUDA graph because the capture size is not set"
)
return
# Collect attention layers and moe layers from the model
python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py
+24 -29
View File
@@ -27,7 +27,7 @@ import tqdm
from sglang.srt.batch_overlap.two_batch_overlap import TboCudaGraphRunnerPlugin
from sglang.srt.compilation.compilation_config import CompilationConfig
from sglang.srt.compilation.compile import install_torch_compiled, set_compiled
from sglang.srt.compilation.compile import install_torch_compiled
from sglang.srt.compilation.piecewise_context_manager import (
enable_piecewise_cuda_graph,
enable_piecewise_cuda_graph_compile,
@@ -196,7 +196,9 @@ class PiecewiseCudaGraphRunner:
if model_runner.server_args.enable_return_hidden_states:
self.capture_hidden_mode = CaptureHiddenMode.FULL
self.max_num_tokens = max(self.capture_num_tokens)
self.max_num_tokens = (
max(self.capture_num_tokens) if self.capture_num_tokens else 8192
)
self.max_bs = model_runner.req_to_token_pool.size
self.is_multimodal = model_runner.is_multimodal
@@ -278,6 +280,10 @@ class PiecewiseCudaGraphRunner:
with patch_model(
language_model.model, self.compile_config.compiler
) as patched_model:
# Dummy warmup for jit kernel
self.warmup_compile(num_tokens=self.capture_num_tokens[0])
install_torch_compiled(
patched_model,
fullgraph=True,
@@ -286,7 +292,7 @@ class PiecewiseCudaGraphRunner:
graph_pool=get_global_graph_memory_pool(),
)
with set_compiled(True), enable_piecewise_cuda_graph_compile():
with enable_piecewise_cuda_graph_compile():
compile_range = (
tqdm.tqdm(list(reversed(self.capture_num_tokens)))
if get_tensor_model_parallel_rank() == 0
@@ -297,7 +303,7 @@ class PiecewiseCudaGraphRunner:
compile_range.set_description(
f"Compiling num tokens ({num_tokens=})"
)
self.warmup_torch_compile(num_tokens=num_tokens)
self.warmup_compile(num_tokens=num_tokens)
set_global_graph_memory_pool(self.device_module.graph_pool_handle())
set_graph_pool_id(get_global_graph_memory_pool())
@@ -305,16 +311,11 @@ class PiecewiseCudaGraphRunner:
self.device_module.synchronize()
self.model_runner.tp_group.barrier()
# Capture
try:
self.capture()
except RuntimeError as e:
raise Exception(
f"Capture cuda graph failed: {e}\n{PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG}"
)
self.capture()
self.raw_num_tokens = 0
def warmup_torch_compile(self, num_tokens: int):
def warmup_compile(self, num_tokens: int):
"""Warmup the model with a simple forward pass before CUDA graph capture."""
buffers = self.buffers
input_ids = buffers.input_ids[:num_tokens]
@@ -409,6 +410,10 @@ class PiecewiseCudaGraphRunner:
return torch.int64 if not is_npu() else torch.int32
def can_run(self, forward_batch: ForwardBatch):
# Disable piecewise cuda graph for input embeddings
# TODO(yuwei): fix it
if forward_batch.input_embeds is not None:
return False
num_tokens = len(forward_batch.input_ids)
if forward_batch.return_logprob:
for start_len, seq_len in zip(
@@ -452,8 +457,7 @@ class PiecewiseCudaGraphRunner:
f"Capturing num tokens ({num_tokens=} {avail_mem=:.2f} GB)"
)
with set_compiled(True):
self.capture_one_batch_size(num_tokens)
self.capture_one_batch_size(num_tokens)
def capture_one_batch_size(self, num_tokens: int):
buffers = self.buffers
@@ -722,6 +726,7 @@ class PiecewiseCudaGraphRunner:
temperature=forward_batch.temperature,
top_p_normalized_logprobs=forward_batch.top_p_normalized_logprobs,
top_p=forward_batch.top_p,
dimensions=forward_batch.dimensions,
)
return static_forward_batch
@@ -743,13 +748,12 @@ class PiecewiseCudaGraphRunner:
self.moe_layers,
self.moe_fusions,
):
with set_compiled(True):
output = self.model_runner.model.forward(
static_forward_batch.input_ids,
static_forward_batch.positions,
static_forward_batch,
**kwargs,
)
output = self.model_runner.model.forward(
static_forward_batch.input_ids,
static_forward_batch.positions,
static_forward_batch,
**kwargs,
)
if isinstance(output, LogitsProcessorOutput):
return LogitsProcessorOutput(
next_token_logits=output.next_token_logits[
@@ -798,12 +802,3 @@ class PiecewiseCudaGraphRunner:
)
return spec_info
PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG = (
"Possible solutions:\n"
"1. set --mem-fraction-static to a smaller value (e.g., 0.8 or 0.7)\n"
"2. set --piecewise-cuda-graph-max-tokens to a smaller value (e.g., 512)\n"
"3. disable Piecewise CUDA graph by unset --enable-piecewise-cuda-graph\n"
"Open an issue on GitHub https://github.com/sgl-project/sglang/issues/new/choose \n"
)
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py
+27 -1
View File
@@ -30,7 +30,33 @@ if TYPE_CHECKING:
from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA
if _is_cuda:
from sgl_kernel import bmm_fp8
from sgl_kernel import bmm_fp8 as _raw_bmm_fp8
from sglang.srt.utils.custom_op import register_custom_op
# TODO(yuwei): remove this wrapper after sgl-kernel registers its own fake/meta impl
# Wrap bmm_fp8 as a custom op so torch.compile does not trace into
# torch.cuda.current_blas_handle() (which returns a non-Tensor).
@register_custom_op(mutates_args=["out"])
def _bmm_fp8_op(
A: torch.Tensor,
B: torch.Tensor,
out: torch.Tensor,
A_scale: torch.Tensor,
B_scale: torch.Tensor,
) -> None:
_raw_bmm_fp8(A, B, A_scale, B_scale, out.dtype, out)
def bmm_fp8(A, B, A_scale, B_scale, dtype, out=None):
if out is None:
out = torch.empty(
(A.shape[0], A.shape[1], B.shape[2]),
device=A.device,
dtype=dtype,
)
_bmm_fp8_op(A, B, out, A_scale, B_scale)
return out
if _use_aiter:
from aiter.ops.triton.batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant import (
python/sglang/srt/models/deepseek_v2.py
+1 -1
View File
@@ -1967,7 +1967,7 @@ class DeepseekV2Model(nn.Module):
# NOTE: torch dynamo does not support graph break in context manager
ctx = (
nullcontext()
if get_global_server_args().enable_piecewise_cuda_graph
if not get_global_server_args().disable_piecewise_cuda_graph
else get_global_expert_distribution_recorder().with_current_layer(i)
)
with ctx:
python/sglang/srt/models/minicpm3.py
+26 -1
View File
@@ -42,7 +42,32 @@ from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.utils import add_prefix, is_cuda
if is_cuda():
from sgl_kernel import bmm_fp8
from sgl_kernel import bmm_fp8 as _raw_bmm_fp8
from sglang.srt.utils.custom_op import register_custom_op
# TODO(yuwei): remove this wrapper after sgl-kernel registers its own fake/meta impl
# Wrap bmm_fp8 as a custom op so torch.compile does not trace into
# torch.cuda.current_blas_handle() (which returns a non-Tensor).
@register_custom_op(mutates_args=["out"])
def _bmm_fp8_op(
A: torch.Tensor,
B: torch.Tensor,
out: torch.Tensor,
A_scale: torch.Tensor,
B_scale: torch.Tensor,
) -> None:
_raw_bmm_fp8(A, B, A_scale, B_scale, out.dtype, out)
def bmm_fp8(A, B, A_scale, B_scale, dtype, out=None):
if out is None:
out = torch.empty(
(A.shape[0], A.shape[1], B.shape[2]),
device=A.device,
dtype=dtype,
)
_bmm_fp8_op(A, B, out, A_scale, B_scale)
return out
class MiniCPM3MLP(nn.Module):
python/sglang/srt/models/minimax_m2.py
+3 -3
View File
@@ -445,7 +445,7 @@ class MiniMaxM2MoE(nn.Module):
if router_logits is not None:
ctx = (
nullcontext()
if get_global_server_args().enable_piecewise_cuda_graph
if not get_global_server_args().disable_piecewise_cuda_graph
else get_global_expert_distribution_recorder().with_current_layer(
self.layer_id
)
@@ -483,7 +483,7 @@ class MiniMaxM2MoE(nn.Module):
if self.ep_size > 1:
ctx = (
nullcontext()
if get_global_server_args().enable_piecewise_cuda_graph
if not get_global_server_args().disable_piecewise_cuda_graph
else get_global_expert_distribution_recorder().with_current_layer(
self.layer_id
)
@@ -909,7 +909,7 @@ class MiniMaxM2Model(nn.Module):
for i in range(self.start_layer, self.end_layer):
ctx = (
nullcontext()
if get_global_server_args().enable_piecewise_cuda_graph
if not get_global_server_args().disable_piecewise_cuda_graph
else get_global_expert_distribution_recorder().with_current_layer(i)
)
with ctx:
python/sglang/srt/models/qwen2_moe.py
+1 -1
View File
@@ -642,7 +642,7 @@ class Qwen2MoeModel(nn.Module):
for i in range(self.start_layer, self.end_layer):
ctx = (
nullcontext()
if get_global_server_args().enable_piecewise_cuda_graph
if not get_global_server_args().disable_piecewise_cuda_graph
else get_global_expert_distribution_recorder().with_current_layer(i)
)
with ctx:
python/sglang/srt/models/qwen3.py
+1 -1
View File
@@ -300,7 +300,7 @@ class Qwen3DecoderLayer(nn.Module):
cache=(
[self.mlp.gate_up_proj.weight, self.mlp.down_proj.weight]
if _is_npu
and not get_global_server_args().enable_piecewise_cuda_graph
and not get_global_server_args().disable_piecewise_cuda_graph
and (
hasattr(self.mlp.gate_up_proj, "weight")
and hasattr(self.mlp.down_proj, "weight")
python/sglang/srt/models/qwen3_next.py
+6 -2
View File
@@ -303,7 +303,7 @@ class Qwen3GatedDeltaNet(nn.Module):
device=torch.get_device_module().current_device(),
dtype=config.torch_dtype,
)
if get_global_server_args().enable_piecewise_cuda_graph
if not get_global_server_args().disable_piecewise_cuda_graph
else FusedRMSNormGated(
self.head_v_dim,
eps=self.layer_norm_epsilon,
@@ -387,7 +387,11 @@ class Qwen3GatedDeltaNet(nn.Module):
return query, key, value, z, b, a
def _forward_input_proj(self, hidden_states: torch.Tensor):
if _is_cpu or _is_npu or get_global_server_args().enable_piecewise_cuda_graph:
if (
_is_cpu
or _is_npu
or not get_global_server_args().disable_piecewise_cuda_graph
):
DUAL_STREAM_TOKEN_THRESHOLD = 0
else:
DUAL_STREAM_TOKEN_THRESHOLD = 1024
python/sglang/srt/server_args.py
+115 -8
View File
@@ -608,7 +608,8 @@ class ServerArgs:
enable_single_batch_overlap: bool = False
tbo_token_distribution_threshold: float = 0.48
enable_torch_compile: bool = False
enable_piecewise_cuda_graph: bool = False
disable_piecewise_cuda_graph: bool = False
enforce_piecewise_cuda_graph: bool = False
enable_torch_compile_debug_mode: bool = False
torch_compile_max_bs: int = 32
piecewise_cuda_graph_max_tokens: Optional[int] = None
@@ -731,6 +732,9 @@ class ServerArgs:
self._handle_cpu_backends()
self._handle_npu_backends()
# Handle piecewise CUDA graph.
self._handle_piecewise_cuda_graph()
# Get GPU memory capacity, which is a common dependency for several configuration steps.
gpu_mem = get_device_memory_capacity(self.device)
@@ -901,6 +905,66 @@ class ServerArgs:
)
self.piecewise_cuda_graph_compiler = "eager"
def _handle_piecewise_cuda_graph(self):
# Skip auto-disable when enforce flag is set (for testing)
if self.enforce_piecewise_cuda_graph:
self.disable_piecewise_cuda_graph = False
return
# Disable piecewise cuda graph with following conditions:
# 1. Disable Model Arch
if self.get_model_config().is_piecewise_cuda_graph_disabled_model:
self.disable_piecewise_cuda_graph = True
# 2. Speculative decoding
if self.speculative_algorithm is not None:
self.disable_piecewise_cuda_graph = True
# 3. DP attention
if self.enable_dp_attention:
self.disable_piecewise_cuda_graph = True
# 4. Torch compile
if self.enable_torch_compile:
self.disable_piecewise_cuda_graph = True
# 5. Pipeline parallelism
if self.pp_size > 1:
self.disable_piecewise_cuda_graph = True
# 6. Non-CUDA hardware (AMD, NPU, etc.)
if is_hip() or is_npu():
self.disable_piecewise_cuda_graph = True
# 7. MoE A2A backend
if self.moe_a2a_backend != "none":
self.disable_piecewise_cuda_graph = True
# 8. LoRA
if self.lora_paths or self.enable_lora:
self.disable_piecewise_cuda_graph = True
# 9. Multimodal / VLM models
if self.get_model_config().is_multimodal:
self.disable_piecewise_cuda_graph = True
# 10. GGUF quantized models (custom dequant ops unsupported by torch.compile)
if (
self.load_format == "gguf"
or self.quantization == "gguf"
or check_gguf_file(self.model_path)
):
self.disable_piecewise_cuda_graph = True
# 11. DLLM (diffusion LLM) models (context manager in forward breaks dynamo)
if self.dllm_algorithm is not None:
self.disable_piecewise_cuda_graph = True
# 12. CPU offload (breaks dynamo)
if self.cpu_offload_gb > 0 or self.enable_hierarchical_cache:
self.disable_piecewise_cuda_graph = True
# 13. Deterministic inference
if self.enable_deterministic_inference:
self.disable_piecewise_cuda_graph = True
# 14. PD disaggregation
if self.disaggregation_mode != "null":
self.disable_piecewise_cuda_graph = True
# 15. Symmetric memory (torch.cuda.use_mem_pool is untraceable by dynamo)
if self.enable_symm_mem:
self.disable_piecewise_cuda_graph = True
# 16. Expert distribution recorder
if self.enable_eplb or self.expert_distribution_recorder_mode is not None:
self.disable_piecewise_cuda_graph = True
def _handle_gpu_memory_settings(self, gpu_mem):
"""
Configure GPU memory-dependent settings including
@@ -1005,6 +1069,19 @@ class ServerArgs:
else:
self.piecewise_cuda_graph_max_tokens = 2048
# If max_total_tokens is set, cap pcg tokens to not exceed max_total_tokens
if self.max_total_tokens is not None:
self.piecewise_cuda_graph_max_tokens = min(
self.piecewise_cuda_graph_max_tokens, self.max_total_tokens
)
# For Llama2 series models, the max tokens is limited to 4096
# TODO(yuwei): remove this after the issue is fixed
if "llama-2" in self.model_path.lower():
self.piecewise_cuda_graph_max_tokens = min(
self.piecewise_cuda_graph_max_tokens, 4096
)
if self.piecewise_cuda_graph_tokens is None:
self.piecewise_cuda_graph_tokens = (
self._generate_piecewise_cuda_graph_tokens()
@@ -1034,9 +1111,13 @@ class ServerArgs:
reserved_mem += self.cuda_graph_max_bs * self.dp_size * 1.5
# For piecewise cuda graphs
if self.enable_piecewise_cuda_graph:
# Only calculate the memory overhead for Non-Torch Memory use since the Torch Memory can be reused with Cuda Graph Capture
reserved_mem += len(self.piecewise_cuda_graph_tokens) * 8
if not self.disable_piecewise_cuda_graph:
if not self.use_mla_backend():
# Only calculate the memory overhead for Non-Torch Memory use since the Torch Memory can be reused with Cuda Graph Capture
reserved_mem += len(self.piecewise_cuda_graph_tokens) * 8
else:
# For MLA backend the memory overhead is much higher than expected with fa3
reserved_mem += 1.5 * 1024
if gpu_mem is not None and gpu_mem > 60 * 1024:
reserved_mem = max(reserved_mem, 10 * 1024)
@@ -1267,7 +1348,7 @@ class ServerArgs:
else:
# DeepSeek V3/R1/V3.1
if self.enable_piecewise_cuda_graph:
if not self.disable_piecewise_cuda_graph:
logger.info("Piecewise CUDA graph is enabled, use MLA for prefill.")
if is_sm100_supported():
@@ -2239,6 +2320,22 @@ class ServerArgs:
self.ep_size == 1
), "FP8/MXFP8 Cutlass MoE is only supported with ep_size == 1"
# TODO(yuwei): Fix piecewise cuda graph support for bypassed topk MoE backends.
# Exception: GptOssForCausalLM wraps the entire MoE block in its own
# custom op (moe_impl), so bypassed topk is handled inside the op body.
if (
not self.enforce_piecewise_cuda_graph
and self.moe_runner_backend in ("flashinfer_trtllm", "flashinfer_mxfp4")
and self.get_model_config().hf_config.architectures[0]
!= "GptOssForCausalLM"
):
self.disable_piecewise_cuda_graph = True
logger.info(
f"Piecewise cuda graph is disabled for MoE runner backend "
f"'{self.moe_runner_backend}' (bypassed topk is incompatible "
f"with torch.compile)."
)
def _handle_a2a_moe(self):
if self.moe_a2a_backend == "deepep":
if self.deepep_mode == "normal":
@@ -2624,7 +2721,7 @@ class ServerArgs:
self.disaggregation_transfer_backend != "fake"
), "Prefill server does not support 'fake' as the transfer backend"
if not self.enable_piecewise_cuda_graph:
if self.disable_piecewise_cuda_graph:
self.disable_cuda_graph = True
logger.warning(
"Cuda graph is disabled for prefill server when piecewise cuda graph is not enabled."
@@ -4712,9 +4809,19 @@ class ServerArgs:
help="Enable debug mode for torch compile",
)
parser.add_argument(
"--enable-piecewise-cuda-graph",
"--disable-piecewise-cuda-graph",
action="store_true",
help="Optimize the model with piecewise cuda graph for extend/prefill only. Experimental feature.",
help="Disable piecewise cuda graph for extend/prefill.",
)
parser.add_argument(
"--enable-piecewise-cuda-graph",
action=DeprecatedAction,
help="Deprecated: Piecewise cuda graph is enabled by default. Use --enforce-piecewise-cuda-graph to skip auto-disable conditions.",
)
parser.add_argument(
"--enforce-piecewise-cuda-graph",
action="store_true",
help="Enforce piecewise cuda graph, skipping all auto-disable conditions. Used for testing.",
)
parser.add_argument(
"--piecewise-cuda-graph-tokens",
python/sglang/srt/utils/custom_op.py
+144
View File
@@ -4,6 +4,7 @@ import inspect
from typing import Any, Callable, List, Optional, TypeVar, Union, overload
import torch
import torch.library
F = TypeVar("F", bound=Callable)
@@ -189,3 +190,146 @@ class CustomOpWrapper:
)
return fake_impl
def register_custom_op_from_extern(
fn: Callable,
*,
op_name: Optional[str] = None,
mutates_args: Optional[List[str]] = None,
out_shape: Optional[Union[int, str]] = None,
out_dtype: Optional[torch.dtype] = None,
fake_impl: Optional[Callable] = None,
computed_args: Optional[dict] = None,
) -> Callable:
"""Wrap an external library function as a custom op for torch.compile compatibility.
Use this to wrap functions from external libraries (e.g. flashinfer kernels) that
perform operations incompatible with torch.compile/dynamo tracing, such as JIT
compilation, file I/O, or dynamic module loading.
The wrapped function becomes an opaque node in the compiled graph. Dynamo will
not trace inside it, avoiding tracing failures. A fake implementation is used
for shape/dtype propagation during compilation.
The external function must have type annotations compatible with
``torch.library.infer_schema`` (``torch.Tensor``, ``int``, ``float``, ``bool``,
``Optional[torch.Tensor]``, etc.).
This function is idempotent: calling it multiple times with the same ``op_name``
(or ``fn.__name__``) safely skips re-registration.
Example usage::
from flashinfer.fused_moe import trtllm_fp8_block_scale_moe
trtllm_fp8_block_scale_moe = register_custom_op_from_extern(
trtllm_fp8_block_scale_moe,
out_shape="hidden_states",
out_dtype=torch.bfloat16,
computed_args={
"tune_max_num_tokens": lambda hidden_states, **kw: next_power_of_2(
hidden_states.shape[0]
),
},
)
:param fn: The external function to wrap.
:param op_name: The name of the custom operator.
Defaults to ``fn.__name__``.
:param mutates_args: A list of argument names that are mutated in-place.
Defaults to ``[]``.
:param out_shape: The position (int) or name (str) of the argument whose shape
matches the output tensor. Used to auto-generate a fake
implementation. Set to ``None`` for inplace-only operators.
:param out_dtype: Override the output dtype in the fake implementation.
If ``None``, ``torch.empty_like`` is used (same dtype as the
reference tensor). Useful when the output dtype differs from
the input (e.g. fp8 input -> bf16 output).
:param fake_impl: A custom fake implementation for shape/dtype propagation.
Only one of ``out_shape`` or ``fake_impl`` should be provided.
:param computed_args: A dict mapping argument names to callables. These arguments
are excluded from the custom op schema and computed inside
the op body at runtime. Each callable receives the other
arguments as keyword args and returns the computed value.
Use this for arguments that vary dynamically (e.g.
``tune_max_num_tokens``) to avoid torch.compile recompilation.
:return: The registered custom op callable (``torch.ops.sglang.<op_name>``).
"""
name = op_name or fn.__name__
computed_args = computed_args or {}
assert not (
out_shape is not None and fake_impl is not None
), "Only one of `out_shape` or `fake_impl` should be provided."
# If computed_args specified, create a wrapper with a reduced signature
# that computes the excluded args inside the op body.
if computed_args:
original_fn = fn
original_sig = inspect.signature(fn)
# Build new signature excluding computed args
new_params = [
p
for param_name, p in original_sig.parameters.items()
if param_name not in computed_args
]
new_sig = original_sig.replace(parameters=new_params)
def wrapper(*args, **kwargs):
bound = new_sig.bind(*args, **kwargs)
bound.apply_defaults()
# Compute excluded args from the bound arguments
for arg_name, compute_fn in computed_args.items():
bound.arguments[arg_name] = compute_fn(**bound.arguments)
return original_fn(**bound.arguments)
wrapper.__name__ = fn.__name__
wrapper.__qualname__ = fn.__qualname__
wrapper.__module__ = fn.__module__
wrapper.__signature__ = new_sig
# Build annotations without computed args, preserving return type
wrapper.__annotations__ = {
k: v
for k, v in getattr(fn, "__annotations__", {}).items()
if k not in computed_args
}
fn = wrapper
# Generate fake_impl from out_shape if needed
fake_sig = inspect.signature(fn)
if fake_impl is None and out_shape is not None:
def _fake_impl(*args, **kwargs):
bound = fake_sig.bind(*args, **kwargs)
bound.apply_defaults()
try:
ref = (
bound.args[out_shape]
if isinstance(out_shape, int)
else bound.arguments[out_shape]
)
except (IndexError, KeyError):
raise RuntimeError(
f"Cannot find output argument at position `{out_shape}` for "
f"external function `{name}` with signature `{fake_sig}`."
)
if out_dtype is not None:
return torch.empty(ref.shape, dtype=out_dtype, device=ref.device)
return torch.empty_like(ref)
fake_impl = _fake_impl
elif fake_impl is None:
fake_impl = lambda *args, **kwargs: None
from sglang.srt.utils.common import direct_register_custom_op
direct_register_custom_op(
op_name=name,
op_func=fn,
mutates_args=mutates_args or [],
fake_impl=fake_impl,
)
return getattr(torch.ops.sglang, name)
test/manual/nightly/test_vlms_piecewise_cuda_graph.py
+1 -1
View File
@@ -137,7 +137,7 @@ class TestVLMPiecewiseCudaGraph(CustomTestCase):
"--trust-remote-code",
"--piecewise-cuda-graph-max-tokens",
"8192",
"--enable-piecewise-cuda-graph",
"--enforce-piecewise-cuda-graph",
"--tp=8",
"--piecewise-cuda-graph-compiler=eager",
"--disable-radix-cache",
test/manual/nightly/test_vlms_vit_cuda_graph.py
+1 -1
View File
@@ -140,7 +140,7 @@ class TestVLMViTCudaGraph(CustomTestCase):
other_args=[
"--mm-attention-backend",
"fa3",
"--enable-piecewise-cuda-graph",
"--enforce-piecewise-cuda-graph",
"--piecewise-cuda-graph-max-tokens",
"8192",
"--chunked-prefill-size",
test/manual/piecewise_cudagraph/test_disaggregation_piecewise_cuda_graph.py
+1 -1
View File
@@ -38,7 +38,7 @@ class TestDisaggregationPiecewiseCudaGraph(PDDisaggregationServerBase):
"prefill",
"--tp",
"1",
"--enable-piecewise-cuda-graph",
"--enforce-piecewise-cuda-graph",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
test/registered/mla/test_mla_deepseek_v3.py
+8 -1
View File
@@ -111,7 +111,14 @@ class TestMLADeepseekV3Fa3Fp8Kvcache(CustomTestCase):
]
if is_cuda():
other_args.extend(
["--attention-backend", "fa3", "--cuda-graph-max-bs", "2"]
[
"--attention-backend",
"fa3",
"--mem-fraction-static",
"0.8",
"--cuda-graph-max-bs",
"2",
]
)
cls.process = popen_launch_server(
cls.model,
test/registered/models/test_generation_models.py
+1 -1
View File
@@ -66,7 +66,7 @@ CI_MODELS = [
# the complete set of models to test sglang's generation model
ALL_MODELS = [
*CI_MODELS,
ModelCase("Qwen/Qwen2-1.5B"),
ModelCase("Qwen/Qwen2-1.5B", decode_tolerance=7e-2),
ModelCase("Qwen/Qwen2.5-14B-Instruct"),
ModelCase("HuggingFaceTB/SmolLM-135M-Instruct", skip_long_prompt=True),
ModelCase("allenai/OLMo-1B-0724-hf", decode_tolerance=8e-2, skip_long_prompt=True),
test/registered/models/test_kimi_linear_models_pcg.py
-1
View File
@@ -41,7 +41,6 @@ class TestKimiLinearPiecewiseCudaGraph(CustomTestCase):
"--tp",
"2",
"--trust-remote-code",
"--enable-piecewise-cuda-graph",
],
)
test/registered/models/test_qwen3_next_models_pcg.py
-3
View File
@@ -40,9 +40,6 @@ class TestQwen3NextPiecewiseCudaGraph(CustomTestCase):
other_args=[
"--tp",
"4",
"--enable-piecewise-cuda-graph",
"--piecewise-cuda-graph-compiler",
"eager",
],
)
test/registered/piecewise_cuda_graph/test_piecewise_cuda_graph_2_gpu.py
-1
View File
@@ -27,7 +27,6 @@ class TestPiecewiseCudaGraphTP(CustomTestCase):
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--enable-piecewise-cuda-graph",
"--piecewise-cuda-graph-compiler",
"eager",
"--tp",
test/registered/piecewise_cuda_graph/test_piecewise_cuda_graph_large_1_gpu.py
+2 -3
View File
@@ -27,7 +27,6 @@ class TestPiecewiseCudaGraphQwen3MoE(CustomTestCase):
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--enable-piecewise-cuda-graph",
"--piecewise-cuda-graph-compiler",
"eager",
],
@@ -65,7 +64,7 @@ class TestPiecewiseCudaGraphGPTQ(CustomTestCase):
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=["--enable-piecewise-cuda-graph"],
other_args=[],
)
@classmethod
@@ -101,7 +100,7 @@ class TestPiecewiseCudaGraphAWQ(CustomTestCase):
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=["--enable-piecewise-cuda-graph"],
other_args=[],
)
@classmethod
test/registered/piecewise_cuda_graph/test_piecewise_cuda_graph_small_1_gpu.py
+3 -10
View File
@@ -33,7 +33,7 @@ class TestPiecewiseCudaGraphCorrectness(CustomTestCase):
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=["--enable-piecewise-cuda-graph"],
other_args=[],
)
@classmethod
@@ -57,8 +57,7 @@ class TestPiecewiseCudaGraphBenchmark(CustomTestCase):
def test_latency(self):
prefill_latency, _, _ = run_bench_one_batch(
DEFAULT_MODEL_NAME_FOR_TEST,
other_args=["--enable-piecewise-cuda-graph"],
DEFAULT_MODEL_NAME_FOR_TEST, other_args=[]
)
self.assertLess(prefill_latency, 0.015)
@@ -76,7 +75,6 @@ class TestPiecewiseCudaGraphLlama31FP4(CustomTestCase):
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--enable-piecewise-cuda-graph",
"--quantization",
"modelopt_fp4",
"--mem-fraction-static",
@@ -112,7 +110,6 @@ class TestPiecewiseCudaGraphDeepSeek(CustomTestCase):
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--enable-piecewise-cuda-graph",
"--piecewise-cuda-graph-compiler",
"eager",
"--piecewise-cuda-graph-max-tokens",
@@ -152,7 +149,6 @@ class TestPiecewiseCudaGraphFP8(CustomTestCase):
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--enable-piecewise-cuda-graph",
"--quantization",
"modelopt_fp8",
"--kv-cache-dtype",
@@ -191,7 +187,6 @@ class TestPiecewiseCudaGraphQwen25VL(CustomTestCase):
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--enable-piecewise-cuda-graph",
"--piecewise-cuda-graph-compiler",
"eager",
"--disable-radix-cache",
@@ -232,7 +227,6 @@ class TestPiecewiseCudaGraphInternVL25(CustomTestCase):
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--enable-piecewise-cuda-graph",
"--piecewise-cuda-graph-compiler",
"eager",
"--disable-radix-cache",
@@ -273,7 +267,6 @@ class TestPiecewiseCudaGraphQwen25VLEmbedding(CustomTestCase):
model_path=model_path,
enable_multimodal=True,
is_embedding=True,
enable_piecewise_cuda_graph=True,
piecewise_cuda_graph_compiler="eager",
)
out = engine.encode([text], image_data=[DEFAULT_IMAGE_URL])[0]["embedding"]
@@ -284,7 +277,7 @@ class TestPiecewiseCudaGraphQwen25VLEmbedding(CustomTestCase):
model_path=model_path,
enable_multimodal=True,
is_embedding=True,
enable_piecewise_cuda_graph=False,
disable_piecewise_cuda_graph=True,
)
out_without_pcg = engine.encode([text], image_data=[DEFAULT_IMAGE_URL])[0][
"embedding"
test/registered/quant/test_deepseek_v3_fp4_4gpu.py
-1
View File
@@ -100,7 +100,6 @@ class TestDeepseekV3FP4PiecewiseCudaGraph(CustomTestCase):
"flashinfer_trtllm",
"--quantization",
"modelopt_fp4",
"--enable-piecewise-cuda-graph",
"--kv-cache-dtype",
"fp8_e4m3",
"--model-loader-extra-config",
test/registered/radix_cache/test_swa_radix_cache_kl.py
+5 -1
View File
@@ -31,11 +31,15 @@ class TestSWARadixCacheKL(CustomTestCase):
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
# Use a lower mem-fraction-static to avoid OOM during input logprobs
# gathering. With PCG enabled, more memory is reserved for CUDA graph
# captures, so the static fraction should be lower.
other_args=[
"--tp-size",
"1",
"--mem-fraction-static",
"0.75",
"0.70",
"--disable-piecewise-cuda-graph",
],
)
test/srt/ascend/test_ascend_piecewise_graph_prefill.py
+2 -2
View File
@@ -36,7 +36,7 @@ class TestPiecewiseGraphPrefillCorrectness(CustomTestCase):
"ascend",
"--cuda-graph-bs",
128,
"--enable-piecewise-cuda-graph",
"--enforce-piecewise-cuda-graph",
"--piecewise-cuda-graph-tokens",
*TOKENS_TO_CAPTURE,
],
@@ -77,7 +77,7 @@ class TestPiecewiseGraphPrefillBenchmark(CustomTestCase):
0.8,
"--attention-backend",
"ascend",
"--enable-piecewise-cuda-graph",
"--enforce-piecewise-cuda-graph",
"--piecewise-cuda-graph-tokens",
]
+ TOKENS_TO_CAPTURE,