[PCG] GPT OSS Triton Kernel Support (#18405)

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

python/sglang/srt/compilation/piecewise_context_manager.py

@@ -56,6 +56,7 @@ class ForwardContext:
         self.attention_layer = None
         self.quant_config = None
         self.moe_layers = None
+        self.moe_fusions = None
 
     def set_forward_batch(self, forward_batch: ForwardBatch):
         self.forward_batch = forward_batch
@@ -69,6 +70,9 @@ class ForwardContext:
     def set_moe_layers(self, layers: List[Any]):
         self.moe_layers = layers
 
+    def set_moe_fusions(self, fusions: List[Any]):
+        self.moe_fusions = fusions
+
 
 _forward_context: Optional[ForwardContext] = None
 
@@ -85,6 +89,7 @@ def set_forward_context(
     attention_layers: List[Any],
     quant_config: Any,
     moe_layers: List[Any],
+    moe_fusions: List[Any],
 ):
     global _forward_context
     _forward_context = ForwardContext()
@@ -92,6 +97,7 @@ def set_forward_context(
     _forward_context.set_attention_layers(attention_layers)
     _forward_context.set_quant_config(quant_config)
     _forward_context.set_moe_layers(moe_layers)
+    _forward_context.set_moe_fusions(moe_fusions)
     try:
         yield
     finally:

python/sglang/srt/layers/moe/fused_moe_triton/layer.py

@@ -957,9 +957,10 @@ class FusedMoE(torch.nn.Module):
 
     def forward(self, hidden_states: torch.Tensor, topk_output: TopKOutput):
         if is_in_piecewise_cuda_graph():
-            assert TopKOutputChecker.format_is_standard(
+            if not TopKOutputChecker.format_is_standard(topk_output):
-                topk_output
+                # Make sure there is torch lib op registration for the whole moe layer
-            ), "Only standard topk output is supported for piecewise cuda graph"
+                return self.forward_impl(hidden_states, topk_output)
+            else:
                 return moe_forward_piecewise_cuda_graph_impl(
                     hidden_states,
                     topk_output.topk_weights,
@@ -1129,9 +1130,10 @@ class FlashInferFusedMoE(FusedMoE):
 
     def forward(self, hidden_states: torch.Tensor, topk_output: TopKOutput):
         if is_in_piecewise_cuda_graph():
-            assert TopKOutputChecker.format_is_standard(
+            if not TopKOutputChecker.format_is_standard(topk_output):
-                topk_output
+                # Make sure there is torch lib op registration for the whole moe layer
-            ), "Only standard topk output is supported for piecewise cuda graph"
+                return self.forward_impl(hidden_states, topk_output)
+            else:
                 return moe_forward_piecewise_cuda_graph_impl(
                     hidden_states,
                     topk_output.topk_weights,

python/sglang/srt/model_executor/model_runner.py

@@ -2126,6 +2126,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         language_model = getattr(self.model, "language_model", self.model)
         self.attention_layers = []
         self.moe_layers = []
+        self.moe_fusions = []
         for layer in language_model.model.layers:
             if hasattr(layer, "self_attn"):
                 if hasattr(layer.self_attn, "attn"):
@@ -2144,15 +2145,20 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                     self.attention_layers.append(layer.attention.attn)
 
             moe_block = None
+            moe_fusion = None
             if hasattr(layer, "mlp") and hasattr(layer.mlp, "experts"):
                 moe_block = layer.mlp.experts
+                moe_fusion = layer.mlp
             if hasattr(layer, "block_sparse_moe") and hasattr(
                 layer.block_sparse_moe, "experts"
             ):
                 moe_block = layer.block_sparse_moe.experts
+                moe_fusion = layer.block_sparse_moe
             if hasattr(layer, "moe") and hasattr(layer.moe, "experts"):
                 moe_block = layer.moe.experts
+                moe_fusion = layer.moe
             self.moe_layers.append(moe_block)
+            self.moe_fusions.append(moe_fusion)
 
         if len(self.attention_layers) < self.model_config.num_hidden_layers:
             # TODO(yuwei): support Non-Standard GQA

python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py

@@ -233,6 +233,7 @@ class PiecewiseCudaGraphRunner:
 
         self.attention_layers = self.model_runner.attention_layers
         self.moe_layers = self.model_runner.moe_layers
+        self.moe_fusions = self.model_runner.moe_fusions
 
         if get_global_graph_memory_pool() is None:
             set_global_graph_memory_pool(self.device_module.graph_pool_handle())
@@ -358,7 +359,11 @@ class PiecewiseCudaGraphRunner:
         set_dp_buffer_len(None, num_tokens, forward_batch.dp_padding_mode.is_max_len())
         set_is_extend_in_batch(False)
         with set_forward_context(
-            forward_batch, self.attention_layers, self.quant_config, self.moe_layers
+            forward_batch,
+            self.attention_layers,
+            self.quant_config,
+            self.moe_layers,
+            self.moe_fusions,
         ):
             _ = self.model_runner.model.forward(
                 forward_batch.input_ids,
@@ -520,7 +525,11 @@ class PiecewiseCudaGraphRunner:
 
             kwargs = {}
             with set_forward_context(
-                forward_batch, self.attention_layers, self.quant_config, self.moe_layers
+                forward_batch,
+                self.attention_layers,
+                self.quant_config,
+                self.moe_layers,
+                self.moe_fusions,
             ):
                 self.model_runner.model.forward(
                     forward_batch.input_ids,
@@ -684,6 +693,7 @@ class PiecewiseCudaGraphRunner:
                 self.attention_layers,
                 self.quant_config,
                 self.moe_layers,
+                self.moe_fusions,
             ):
                 with set_compiled(True):
                     output = self.model_runner.model.forward(

python/sglang/srt/models/gpt_oss.py

@@ -25,6 +25,10 @@ import torch
 from torch import nn
 from transformers import PretrainedConfig
 
+from sglang.srt.compilation.piecewise_context_manager import (
+    get_forward_context,
+    is_in_piecewise_cuda_graph,
+)
 from sglang.srt.distributed import (
     get_moe_expert_parallel_rank,
     get_moe_expert_parallel_world_size,
@@ -72,6 +76,7 @@ from sglang.srt.models.utils import (
 )
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import LazyValue, add_prefix, is_cuda, is_npu, make_layers
+from sglang.srt.utils.custom_op import register_custom_op
 
 _is_cuda = is_cuda()
 _is_npu = is_npu()
@@ -183,7 +188,9 @@ class GptOssSparseMoeBlock(nn.Module):
         should_allreduce_fusion: bool = False,
     ) -> torch.Tensor:
         num_tokens, hidden_dim = hidden_states.shape
-
+        if is_in_piecewise_cuda_graph():
+            final_hidden_states = moe_impl(self.layer_id, hidden_states)
+        else:
             router_logits, _ = self.router(hidden_states)
             topk_output = self.topk(hidden_states, router_logits)
             final_hidden_states = self.experts(hidden_states, topk_output)
@@ -195,6 +202,16 @@ class GptOssSparseMoeBlock(nn.Module):
         return ans
 
 
+@register_custom_op(out_shape="hidden_states")
+def moe_impl(layer_id: int, hidden_states: torch.Tensor) -> torch.Tensor:
+    forward_context = get_forward_context()
+    moe_fusion = forward_context.moe_fusions[layer_id]
+    router_logits, _ = moe_fusion.router(hidden_states)
+    topk_output = moe_fusion.topk(hidden_states, router_logits)
+    final_hidden_states = moe_fusion.experts(hidden_states, topk_output)
+    return final_hidden_states
+
+
 class GptOssAttention(nn.Module):
     def __init__(
         self,

python/sglang/srt/server_args.py

@@ -1363,12 +1363,7 @@ class ServerArgs:
                 self.dtype = "bfloat16"
 
             if self.moe_runner_backend == "auto":
-                if self.enable_piecewise_cuda_graph:
+                if is_blackwell_supported() and is_mxfp4_quant_format:
-                    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."