From eae59b337e7ce689d0a51fcec72b6eb0333b1c71 Mon Sep 17 00:00:00 2001 From: Yuwei An Date: Fri, 14 Nov 2025 17:28:00 -0800 Subject: [PATCH] Piecewise Cuda Graph Support for gpt-oss model (#13045) Signed-off-by: Oasis-Git --- .../compilation/piecewise_context_manager.py | 16 ++++++++++++++ .../attention/flashinfer_mla_backend.py | 4 +--- .../piecewise_cuda_graph_runner.py | 21 ++++--------------- python/sglang/srt/models/deepseek_v2.py | 4 +--- python/sglang/srt/server_args.py | 7 ++++++- 5 files changed, 28 insertions(+), 24 deletions(-) diff --git a/python/sglang/srt/compilation/piecewise_context_manager.py b/python/sglang/srt/compilation/piecewise_context_manager.py index 711b4f487..a0d3c7dd2 100644 --- a/python/sglang/srt/compilation/piecewise_context_manager.py +++ b/python/sglang/srt/compilation/piecewise_context_manager.py @@ -4,6 +4,22 @@ from typing import Any, List, Optional from sglang.srt.model_executor.forward_batch_info import ForwardBatch +_in_piecewise_cuda_graph = False + + +def is_in_piecewise_cuda_graph(): + return _in_piecewise_cuda_graph + + +@contextmanager +def enable_piecewise_cuda_graph(): + global _in_piecewise_cuda_graph + _in_piecewise_cuda_graph = True + + yield + + _in_piecewise_cuda_graph = False + @dataclass class ForwardContext: diff --git a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py index 348b3d3f5..601a80cea 100644 --- a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py @@ -15,6 +15,7 @@ from typing import TYPE_CHECKING, Callable, Optional, Union import torch +from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph from sglang.srt.environ import envs from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.layers.attention.flashinfer_backend import ( @@ -22,9 +23,6 @@ from sglang.srt.layers.attention.flashinfer_backend import ( ) from sglang.srt.layers.dp_attention import get_attention_tp_size from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode -from sglang.srt.model_executor.piecewise_cuda_graph_runner import ( - is_in_piecewise_cuda_graph, -) from sglang.srt.server_args import get_global_server_args from sglang.srt.speculative.spec_info import SpecInput from sglang.srt.utils import ( 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 72f190f55..2322ebb22 100644 --- a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py @@ -26,7 +26,10 @@ import tqdm from sglang.srt.compilation.compilation_config import CompilationConfig from sglang.srt.compilation.compile import install_torch_compiled, set_compiled -from sglang.srt.compilation.piecewise_context_manager import set_forward_context +from sglang.srt.compilation.piecewise_context_manager import ( + enable_piecewise_cuda_graph, + set_forward_context, +) from sglang.srt.custom_op import CustomOp from sglang.srt.distributed import get_tensor_model_parallel_rank from sglang.srt.distributed.device_communicators.pynccl_allocator import ( @@ -55,22 +58,6 @@ logger = logging.getLogger(__name__) if TYPE_CHECKING: from sglang.srt.model_executor.model_runner import ModelRunner -_in_piecewise_cuda_graph = False - - -def is_in_piecewise_cuda_graph(): - return _in_piecewise_cuda_graph - - -@contextmanager -def enable_piecewise_cuda_graph(): - global _in_piecewise_cuda_graph - _in_piecewise_cuda_graph = True - - yield - - _in_piecewise_cuda_graph = False - @contextmanager def freeze_gc(enable_cudagraph_gc: bool): diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 486e88805..b174407c7 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -29,6 +29,7 @@ import tqdm from torch import nn from transformers import PretrainedConfig +from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph from sglang.srt.configs.model_config import ( get_nsa_index_head_dim, get_nsa_index_n_heads, @@ -109,9 +110,6 @@ from sglang.srt.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors -from sglang.srt.model_executor.piecewise_cuda_graph_runner import ( - is_in_piecewise_cuda_graph, -) from sglang.srt.model_loader.utils import maybe_executor_submit, should_async_load from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.server_args import get_global_server_args diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 88ec29b25..f5cfd59ac 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -998,7 +998,12 @@ class ServerArgs: self.dtype = "bfloat16" if self.moe_runner_backend == "auto": - if is_blackwell_supported() and is_mxfp4_quant_format: + if self.enable_piecewise_cuda_graph: + self.moe_runner_backend = "auto" + logger.warning( + "Enable piecewise CUDA graph, enabling auto MOE kernel." + ) + elif is_blackwell_supported() and is_mxfp4_quant_format: self.moe_runner_backend = "flashinfer_mxfp4" logger.warning( "Detected SM100 and MXFP4 quantization format for GPT-OSS model, enabling FlashInfer MXFP4 MOE kernel."