diff --git a/python/sglang/srt/compilation/piecewise_context_manager.py b/python/sglang/srt/compilation/piecewise_context_manager.py index d8f6f5cbe..a87d31f3e 100644 --- a/python/sglang/srt/compilation/piecewise_context_manager.py +++ b/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: diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index 185f1bea3..ae70a66e2 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -957,16 +957,17 @@ 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( - topk_output - ), "Only standard topk output is supported for piecewise cuda graph" - return moe_forward_piecewise_cuda_graph_impl( - hidden_states, - topk_output.topk_weights, - topk_output.topk_ids, - topk_output.router_logits, - self.layer_id, - ) + if not TopKOutputChecker.format_is_standard(topk_output): + # Make sure there is torch lib op registration for the whole moe layer + return self.forward_impl(hidden_states, topk_output) + else: + return moe_forward_piecewise_cuda_graph_impl( + hidden_states, + topk_output.topk_weights, + topk_output.topk_ids, + topk_output.router_logits, + self.layer_id, + ) else: return self.forward_impl(hidden_states, topk_output) @@ -1129,16 +1130,17 @@ class FlashInferFusedMoE(FusedMoE): def forward(self, hidden_states: torch.Tensor, topk_output: TopKOutput): if is_in_piecewise_cuda_graph(): - assert TopKOutputChecker.format_is_standard( - topk_output - ), "Only standard topk output is supported for piecewise cuda graph" - return moe_forward_piecewise_cuda_graph_impl( - hidden_states, - topk_output.topk_weights, - topk_output.topk_ids, - topk_output.router_logits, - self.layer_id, - ) + if not TopKOutputChecker.format_is_standard(topk_output): + # Make sure there is torch lib op registration for the whole moe layer + return self.forward_impl(hidden_states, topk_output) + else: + return moe_forward_piecewise_cuda_graph_impl( + hidden_states, + topk_output.topk_weights, + topk_output.topk_ids, + topk_output.router_logits, + self.layer_id, + ) else: return self.forward_impl(hidden_states, topk_output) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 5eab72408..d23d8e9ad 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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 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 51e72aebc..0cb778c7b 100644 --- a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py +++ b/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( diff --git a/python/sglang/srt/models/gpt_oss.py b/python/sglang/srt/models/gpt_oss.py index 74ce2e3eb..2cf813bce 100644 --- a/python/sglang/srt/models/gpt_oss.py +++ b/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,10 +188,12 @@ class GptOssSparseMoeBlock(nn.Module): should_allreduce_fusion: bool = False, ) -> torch.Tensor: num_tokens, hidden_dim = hidden_states.shape - - router_logits, _ = self.router(hidden_states) - topk_output = self.topk(hidden_states, router_logits) - final_hidden_states = self.experts(hidden_states, topk_output) + 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) if self.tp_size > 1 and not should_allreduce_fusion: final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) @@ -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, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 8223d8683..ba82018f7 100644 --- a/python/sglang/srt/server_args.py +++ b/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: - self.moe_runner_backend = "auto" - logger.warning( - "Enable piecewise CUDA graph, enabling auto MOE kernel." - ) - elif is_blackwell_supported() and is_mxfp4_quant_format: + if 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."