Piecewise Cuda Graph Support for gpt-oss model (#13045)

Signed-off-by: Oasis-Git <ayw.sirius19@gmail.com>

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:

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 (

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):

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

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."