diff --git a/python/sglang/srt/compilation/compile.py b/python/sglang/srt/compilation/compile.py index b3feace52..46a9240fb 100644 --- a/python/sglang/srt/compilation/compile.py +++ b/python/sglang/srt/compilation/compile.py @@ -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) diff --git a/python/sglang/srt/compilation/piecewise_context_manager.py b/python/sglang/srt/compilation/piecewise_context_manager.py index e588b6335..20a08a997 100644 --- a/python/sglang/srt/compilation/piecewise_context_manager.py +++ b/python/sglang/srt/compilation/piecewise_context_manager.py @@ -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" +) diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 9e4bd54ab..a80e0e78f 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -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 diff --git a/python/sglang/srt/layers/attention/fla/layernorm_gated.py b/python/sglang/srt/layers/attention/fla/layernorm_gated.py index c9ec07974..b30cdacdb 100644 --- a/python/sglang/srt/layers/attention/fla/layernorm_gated.py +++ b/python/sglang/srt/layers/attention/fla/layernorm_gated.py @@ -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) diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index 841fbb791..39d44b1e6 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_backend.py @@ -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. " diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index 03e5516c2..052d051c9 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/python/sglang/srt/layers/communicator.py @@ -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: diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index aeb5e793c..2523ef851 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -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 diff --git a/python/sglang/srt/layers/quantization/awq.py b/python/sglang/srt/layers/quantization/awq.py index 42d3a6c4d..c05aa091a 100644 --- a/python/sglang/srt/layers/quantization/awq.py +++ b/python/sglang/srt/layers/quantization/awq.py @@ -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( diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index ef50634a6..66575c69e 100644 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -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: diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 3902a4caa..f4d2a63eb 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -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 diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index cee3dd2ea..585ab89f4 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -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 diff --git a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py index feccc5c56..88f840929 100644 --- a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py @@ -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" -) diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index c365d5c93..9eca43ce7 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py @@ -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 ( diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 04e97d952..f2dd98c6e 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -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: diff --git a/python/sglang/srt/models/minicpm3.py b/python/sglang/srt/models/minicpm3.py index 9755a6f6b..95dca19da 100644 --- a/python/sglang/srt/models/minicpm3.py +++ b/python/sglang/srt/models/minicpm3.py @@ -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): diff --git a/python/sglang/srt/models/minimax_m2.py b/python/sglang/srt/models/minimax_m2.py index c5ed94b47..7aad87c64 100644 --- a/python/sglang/srt/models/minimax_m2.py +++ b/python/sglang/srt/models/minimax_m2.py @@ -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: diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index bbb883a2d..4006de801 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -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: diff --git a/python/sglang/srt/models/qwen3.py b/python/sglang/srt/models/qwen3.py index b056317e4..95aaaa235 100644 --- a/python/sglang/srt/models/qwen3.py +++ b/python/sglang/srt/models/qwen3.py @@ -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") diff --git a/python/sglang/srt/models/qwen3_next.py b/python/sglang/srt/models/qwen3_next.py index e742dd543..42146e605 100644 --- a/python/sglang/srt/models/qwen3_next.py +++ b/python/sglang/srt/models/qwen3_next.py @@ -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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 0c13fea95..7aa58688c 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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", diff --git a/python/sglang/srt/utils/custom_op.py b/python/sglang/srt/utils/custom_op.py index 9b4bd90d6..cd708aa4c 100644 --- a/python/sglang/srt/utils/custom_op.py +++ b/python/sglang/srt/utils/custom_op.py @@ -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.``). + """ + 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) diff --git a/test/manual/nightly/test_vlms_piecewise_cuda_graph.py b/test/manual/nightly/test_vlms_piecewise_cuda_graph.py index 7a72dd3fa..8bacf3691 100644 --- a/test/manual/nightly/test_vlms_piecewise_cuda_graph.py +++ b/test/manual/nightly/test_vlms_piecewise_cuda_graph.py @@ -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", diff --git a/test/manual/nightly/test_vlms_vit_cuda_graph.py b/test/manual/nightly/test_vlms_vit_cuda_graph.py index d0f519848..4b2bc74d7 100644 --- a/test/manual/nightly/test_vlms_vit_cuda_graph.py +++ b/test/manual/nightly/test_vlms_vit_cuda_graph.py @@ -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", diff --git a/test/manual/piecewise_cudagraph/test_disaggregation_piecewise_cuda_graph.py b/test/manual/piecewise_cudagraph/test_disaggregation_piecewise_cuda_graph.py index 90e6f5547..086830a90 100644 --- a/test/manual/piecewise_cudagraph/test_disaggregation_piecewise_cuda_graph.py +++ b/test/manual/piecewise_cudagraph/test_disaggregation_piecewise_cuda_graph.py @@ -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( diff --git a/test/registered/mla/test_mla_deepseek_v3.py b/test/registered/mla/test_mla_deepseek_v3.py index c30d6827b..86c9fef79 100644 --- a/test/registered/mla/test_mla_deepseek_v3.py +++ b/test/registered/mla/test_mla_deepseek_v3.py @@ -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, diff --git a/test/registered/models/test_generation_models.py b/test/registered/models/test_generation_models.py index 2b4945d46..2606dbb5e 100644 --- a/test/registered/models/test_generation_models.py +++ b/test/registered/models/test_generation_models.py @@ -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), diff --git a/test/registered/models/test_kimi_linear_models_pcg.py b/test/registered/models/test_kimi_linear_models_pcg.py index 6c0094794..c464831dd 100644 --- a/test/registered/models/test_kimi_linear_models_pcg.py +++ b/test/registered/models/test_kimi_linear_models_pcg.py @@ -41,7 +41,6 @@ class TestKimiLinearPiecewiseCudaGraph(CustomTestCase): "--tp", "2", "--trust-remote-code", - "--enable-piecewise-cuda-graph", ], ) diff --git a/test/registered/models/test_qwen3_next_models_pcg.py b/test/registered/models/test_qwen3_next_models_pcg.py index bcba9aee9..7156ad042 100644 --- a/test/registered/models/test_qwen3_next_models_pcg.py +++ b/test/registered/models/test_qwen3_next_models_pcg.py @@ -40,9 +40,6 @@ class TestQwen3NextPiecewiseCudaGraph(CustomTestCase): other_args=[ "--tp", "4", - "--enable-piecewise-cuda-graph", - "--piecewise-cuda-graph-compiler", - "eager", ], ) diff --git a/test/registered/piecewise_cuda_graph/test_piecewise_cuda_graph_2_gpu.py b/test/registered/piecewise_cuda_graph/test_piecewise_cuda_graph_2_gpu.py index ac8176210..324e1eae1 100644 --- a/test/registered/piecewise_cuda_graph/test_piecewise_cuda_graph_2_gpu.py +++ b/test/registered/piecewise_cuda_graph/test_piecewise_cuda_graph_2_gpu.py @@ -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", diff --git a/test/registered/piecewise_cuda_graph/test_piecewise_cuda_graph_large_1_gpu.py b/test/registered/piecewise_cuda_graph/test_piecewise_cuda_graph_large_1_gpu.py index 8814cdcd3..6fa32efd4 100644 --- a/test/registered/piecewise_cuda_graph/test_piecewise_cuda_graph_large_1_gpu.py +++ b/test/registered/piecewise_cuda_graph/test_piecewise_cuda_graph_large_1_gpu.py @@ -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 diff --git a/test/registered/piecewise_cuda_graph/test_piecewise_cuda_graph_small_1_gpu.py b/test/registered/piecewise_cuda_graph/test_piecewise_cuda_graph_small_1_gpu.py index 1789b1080..60c14d37a 100644 --- a/test/registered/piecewise_cuda_graph/test_piecewise_cuda_graph_small_1_gpu.py +++ b/test/registered/piecewise_cuda_graph/test_piecewise_cuda_graph_small_1_gpu.py @@ -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" diff --git a/test/registered/quant/test_deepseek_v3_fp4_4gpu.py b/test/registered/quant/test_deepseek_v3_fp4_4gpu.py index 0ccf4b52e..6dd9ca13c 100644 --- a/test/registered/quant/test_deepseek_v3_fp4_4gpu.py +++ b/test/registered/quant/test_deepseek_v3_fp4_4gpu.py @@ -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", diff --git a/test/registered/radix_cache/test_swa_radix_cache_kl.py b/test/registered/radix_cache/test_swa_radix_cache_kl.py index e6222b0ad..f51d863d4 100644 --- a/test/registered/radix_cache/test_swa_radix_cache_kl.py +++ b/test/registered/radix_cache/test_swa_radix_cache_kl.py @@ -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", ], ) diff --git a/test/srt/ascend/test_ascend_piecewise_graph_prefill.py b/test/srt/ascend/test_ascend_piecewise_graph_prefill.py index 44122b716..939486582 100644 --- a/test/srt/ascend/test_ascend_piecewise_graph_prefill.py +++ b/test/srt/ascend/test_ascend_piecewise_graph_prefill.py @@ -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,