From 9d0347b33affec03154667fc8cf666113f195cb0 Mon Sep 17 00:00:00 2001 From: Yuwei An Date: Fri, 19 Dec 2025 09:59:27 -0800 Subject: [PATCH] EP Support for Piecewise Cuda Graph (#14164) Signed-off-by: Oasis-Git --- python/sglang/srt/compilation/backend.py | 12 +-- .../srt/compilation/compilation_config.py | 7 ++ .../compilation/piecewise_context_manager.py | 11 ++- python/sglang/srt/layers/communicator.py | 1 + python/sglang/srt/layers/moe/ep_moe/layer.py | 25 +++++- .../srt/layers/moe/fused_moe_triton/layer.py | 86 ++++++++++++++++++- python/sglang/srt/layers/moe/topk.py | 19 ++-- .../sglang/srt/model_executor/model_runner.py | 24 +++++- .../piecewise_cuda_graph_runner.py | 49 ++++++----- test/srt/run_suite.py | 2 +- test/srt/test_piecewise_cuda_graph_2_gpu.py | 44 ++++++++++ 11 files changed, 224 insertions(+), 56 deletions(-) diff --git a/python/sglang/srt/compilation/backend.py b/python/sglang/srt/compilation/backend.py index 08c181415..8af025707 100644 --- a/python/sglang/srt/compilation/backend.py +++ b/python/sglang/srt/compilation/backend.py @@ -26,16 +26,6 @@ from sglang.srt.utils.common import is_npu, rank0_log logger = logging.getLogger(__name__) -SPLIT_OPS = [ - "sglang.unified_attention_with_output", - "sglang.gdn_with_output", -] - - -def add_split_ops(ops): - SPLIT_OPS.extend(ops) - - def make_compiler(config: CompilationConfig): if config.compiler == "eager": return EagerAdapter() @@ -433,7 +423,7 @@ class SGLangBackend: self.split_gm, self.piecewise_graphs = split_graph( graph, - SPLIT_OPS, + self.compile_config.split_ops, ) from torch._dynamo.utils import lazy_format_graph_code diff --git a/python/sglang/srt/compilation/compilation_config.py b/python/sglang/srt/compilation/compilation_config.py index 5ddafe8af..9da12bc91 100644 --- a/python/sglang/srt/compilation/compilation_config.py +++ b/python/sglang/srt/compilation/compilation_config.py @@ -15,6 +15,13 @@ class CompilationConfig: self.capture_sizes = capture_sizes self.compiler = compiler self.enable_debug_mode = enable_debug_mode + self.split_ops = [ + "sglang.unified_attention_with_output", + "sglang.gdn_with_output", + ] + + def add_split_op(self, op: str): + self.split_ops.append(op) def add_traced_file(self, file_path: str): self.traced_files.add(file_path) diff --git a/python/sglang/srt/compilation/piecewise_context_manager.py b/python/sglang/srt/compilation/piecewise_context_manager.py index a0d3c7dd2..bafb5645d 100644 --- a/python/sglang/srt/compilation/piecewise_context_manager.py +++ b/python/sglang/srt/compilation/piecewise_context_manager.py @@ -26,6 +26,8 @@ class ForwardContext: def __init__(self): self.forward_batch = None self.attention_layer = None + self.quant_config = None + self.moe_layers = None def set_forward_batch(self, forward_batch: ForwardBatch): self.forward_batch = forward_batch @@ -36,6 +38,9 @@ class ForwardContext: def set_quant_config(self, quant_config: Any): self.quant_config = quant_config + def set_moe_layers(self, layers: List[Any]): + self.moe_layers = layers + _forward_context: Optional[ForwardContext] = None @@ -48,13 +53,17 @@ def get_forward_context() -> Optional[ForwardContext]: @contextmanager def set_forward_context( - forward_batch: ForwardBatch, attention_layers: List[Any], quant_config: Any + forward_batch: ForwardBatch, + attention_layers: List[Any], + quant_config: Any, + moe_layers: List[Any], ): global _forward_context _forward_context = ForwardContext() _forward_context.set_forward_batch(forward_batch) _forward_context.set_attention_layers(attention_layers) _forward_context.set_quant_config(quant_config) + _forward_context.set_moe_layers(moe_layers) try: yield finally: diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index 33ce75364..a39dcef47 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/python/sglang/srt/layers/communicator.py @@ -545,6 +545,7 @@ class LayerCommunicator: return True return False + # NOTE: This function will cause torch recompilation def should_fuse_mlp_allreduce_with_next_layer( self, forward_batch: ForwardBatch ) -> bool: diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py index 4f22ba798..a525637ec 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/python/sglang/srt/layers/moe/ep_moe/layer.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any, Dict, 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.hardware_backend.npu.quantization.fused_moe_method_npu import ( NPUW4A16Int4DynamicMoEMethod, @@ -20,7 +21,7 @@ from sglang.srt.layers.moe.token_dispatcher.deepep import ( DeepEPLLCombineInput, DeepEPNormalCombineInput, ) -from sglang.srt.layers.moe.topk import TopKOutput +from sglang.srt.layers.moe.topk import TopKOutput, TopKOutputChecker from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.quantization.fp8 import Fp8Config from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz @@ -90,7 +91,6 @@ class DeepEPMoE(FusedMoE): routed_scaling_factor=routed_scaling_factor, **kwargs, ) - if _use_aiter or _is_npu: self.deprecate_flag = False elif deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM and isinstance( @@ -152,9 +152,28 @@ class DeepEPMoE(FusedMoE): 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 torch.ops.sglang.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) + + def forward_impl( + self, + hidden_states: torch.Tensor, + topk_output: TopKOutput, + ): if self.deprecate_flag: - return super().forward( + return super().forward_impl( hidden_states, topk_output, ) 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 73fe410b8..c72e11e06 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -8,6 +8,10 @@ import torch from sglang.srt.batch_overlap.single_batch_overlap import DownGemmOverlapArgs from sglang.srt.batch_overlap.two_batch_overlap import MaybeTboDeepEPDispatcher +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, @@ -37,7 +41,7 @@ from sglang.srt.layers.moe.token_dispatcher.standard import ( StandardDispatcher, StandardDispatchOutput, ) -from sglang.srt.layers.moe.topk import TopKOutput, TopKOutputChecker +from sglang.srt.layers.moe.topk import StandardTopKOutput, TopKOutput, TopKOutputChecker from sglang.srt.layers.moe.utils import RoutingMethodType from sglang.srt.layers.quantization.base_config import ( FusedMoEMethodBase, @@ -57,6 +61,7 @@ from sglang.srt.utils import ( next_power_of_2, round_up, ) +from sglang.srt.utils.common import direct_register_custom_op if is_flashinfer_available(): from flashinfer import fp4_quantize @@ -882,6 +887,21 @@ 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 torch.ops.sglang.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) + + def forward_impl(self, hidden_states: torch.Tensor, topk_output: TopKOutput): origin_hidden_states_dim = hidden_states.shape[-1] assert self.quant_method is not None @@ -1039,6 +1059,21 @@ class FlashInferFusedMoE(FusedMoE): super().__init__(*args, **kwargs) 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 torch.ops.sglang.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) + + def forward_impl(self, hidden_states: torch.Tensor, topk_output: TopKOutput): assert ( self.moe_runner_config.activation == "silu" ), "Only silu is supported for flashinfer trtllm moe" @@ -1150,6 +1185,21 @@ class FlashInferFP4MoE(FusedMoE): return hs_fp4, hs_sf 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 torch.ops.sglang.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) + + def forward_impl(self, hidden_states: torch.Tensor, topk_output: TopKOutput): """Forward pass using FP4 TRTLLM kernel. Args: @@ -1234,3 +1284,37 @@ class FlashInferFP4MoE(FusedMoE): )[0] return result + + +def moe_forward_piecewise_cuda_graph_impl( + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + router_logits: torch.Tensor, + layer_id: int, +) -> torch.Tensor: + # only standard topk output is supported for piecewise cuda graph + topk_output = StandardTopKOutput( + topk_weights=topk_weights, topk_ids=topk_ids, router_logits=router_logits + ) + forward_context = get_forward_context() + moe_layer = forward_context.moe_layers[layer_id] + return moe_layer.forward_impl(hidden_states, topk_output) + + +def moe_forward_piecewise_cuda_graph_impl_fake( + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + router_logits: torch.Tensor, + layer_id: int, +) -> torch.Tensor: + return torch.empty_like(hidden_states) + + +direct_register_custom_op( + op_name="moe_forward_piecewise_cuda_graph_impl", + op_func=moe_forward_piecewise_cuda_graph_impl, + mutates_args=[], + fake_impl=moe_forward_piecewise_cuda_graph_impl_fake, +) diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index a802647e8..a43fa350d 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -17,7 +17,7 @@ from __future__ import annotations import logging import math from dataclasses import dataclass -from enum import Enum, auto +from enum import IntEnum, auto from typing import ( TYPE_CHECKING, Callable, @@ -120,33 +120,24 @@ class TopKOutputChecker: @staticmethod def format_is_standard(topk_output: TopKOutput) -> TypeGuard[StandardTopKOutput]: - return topk_output.format.is_standard() + return isinstance(topk_output, StandardTopKOutput) @staticmethod def format_is_triton_kernels( topk_output: TopKOutput, ) -> TypeGuard[TritonKernelTopKOutput]: - return topk_output.format.is_triton_kernels() + return isinstance(topk_output, TritonKernelTopKOutput) @staticmethod def format_is_bypassed(topk_output: TopKOutput) -> TypeGuard[BypassedTopKOutput]: - return topk_output.format.is_bypassed() + return isinstance(topk_output, BypassedTopKOutput) -class TopKOutputFormat(Enum): +class TopKOutputFormat(IntEnum): STANDARD = auto() TRITON_KERNEL = auto() BYPASSED = auto() - def is_standard(self) -> bool: - return self == TopKOutputFormat.STANDARD - - def is_triton_kernels(self) -> bool: - return self == TopKOutputFormat.TRITON_KERNEL - - def is_bypassed(self) -> bool: - return self == TopKOutputFormat.BYPASSED - @runtime_checkable class TopKOutput(Protocol): diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 35a9080ec..26d61239f 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -97,6 +97,7 @@ from sglang.srt.layers.dp_attention import ( set_is_extend_in_batch, ) from sglang.srt.layers.logits_processor import LogitsProcessorOutput +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 @@ -1723,6 +1724,13 @@ class ModelRunner: "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 init_memory_pool( @@ -2623,9 +2631,10 @@ class ModelRunner: ): return - # Collect attention layers from the model - self.attention_layers = [] + # Collect attention layers and moe layers from the model self.model.model = resolve_language_model(self.model) + self.attention_layers = [] + self.moe_layers = [] for layer in self.model.model.layers: if hasattr(layer, "self_attn"): if hasattr(layer.self_attn, "attn"): @@ -2643,6 +2652,17 @@ class ModelRunner: if hasattr(layer.attention, "attn"): self.attention_layers.append(layer.attention.attn) + moe_block = None + if hasattr(layer, "mlp") and hasattr(layer.mlp, "experts"): + moe_block = layer.mlp.experts + if hasattr(layer, "block_sparse_moe") and hasattr( + layer.block_sparse_moe, "experts" + ): + moe_block = layer.block_sparse_moe.experts + if hasattr(layer, "moe") and hasattr(layer.moe, "experts"): + moe_block = layer.moe.experts + self.moe_layers.append(moe_block) + if len(self.attention_layers) < self.model_config.num_hidden_layers: # TODO(yuwei): support Non-Standard GQA log_info_on_rank0( 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 3d54b0c7f..94a16228b 100644 --- a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py @@ -44,6 +44,7 @@ from sglang.srt.layers.dp_attention import ( set_is_extend_in_batch, ) from sglang.srt.layers.logits_processor import LogitsProcessorOutput +from sglang.srt.layers.moe.utils import get_moe_a2a_backend from sglang.srt.layers.pooler import EmbeddingPoolerOutput from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, @@ -195,6 +196,11 @@ class PiecewiseCudaGraphRunner: self.model_runner.server_args.piecewise_cuda_graph_compiler, self.model_runner.server_args.enable_torch_compile_debug_mode, ) + if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake(): + self.compile_config.add_split_op( + "sglang.moe_forward_piecewise_cuda_graph_impl" + ) + self.quant_config = getattr(self.model_runner.model, "quant_config", None) # Batch sizes to capture @@ -243,6 +249,7 @@ class PiecewiseCudaGraphRunner: ) self.attention_layers = self.model_runner.attention_layers + self.moe_layers = self.model_runner.moe_layers if get_global_graph_memory_pool() is None: set_global_graph_memory_pool(self.device_module.graph_pool_handle()) @@ -260,10 +267,9 @@ class PiecewiseCudaGraphRunner: compile_config=self.compile_config, graph_pool=get_global_graph_memory_pool(), ) - - with set_compiled(True): - self.warmup_torch_compile() - + with freeze_gc(self.model_runner.server_args.enable_cudagraph_gc): + with set_compiled(True): + self.warmup_torch_compile() # Capture try: self.capture() @@ -277,28 +283,24 @@ class PiecewiseCudaGraphRunner: def warmup_torch_compile(self): """Warmup the model with a simple forward pass before CUDA graph capture.""" num_tokens = 2 - + input_ids = self.input_ids[:num_tokens] + input_embeds = self.input_embeds[:num_tokens] if self.is_multimodal else None out_cache_loc = self.out_cache_loc[:num_tokens] out_cache_loc_swa = ( self.out_cache_loc_swa[:num_tokens] if self.out_cache_loc_swa is not None else None ) + positions = self.positions[:num_tokens] + mrope_positions = ( + self.mrope_positions[:, :num_tokens] if self.is_multimodal else None + ) with torch.device(self.device): forward_batch = ForwardBatch( forward_mode=ForwardMode.EXTEND, batch_size=1, - input_ids=(torch.randint(0, 100, (num_tokens,), device=self.device)), - input_embeds=( - torch.randn( - num_tokens, - self.model_runner.model_config.hidden_size, - dtype=self.model_runner.dtype, - device=self.device, - ) - if self.is_multimodal - else None - ), + input_ids=input_ids, + input_embeds=input_embeds, req_pool_indices=torch.arange(1, device=self.device), seq_lens=torch.tensor([num_tokens], device=self.device), next_token_logits_buffer=None, @@ -319,14 +321,12 @@ class PiecewiseCudaGraphRunner: extend_prefix_lens_cpu=torch.tensor([num_tokens], device="cpu"), extend_seq_lens_cpu=torch.tensor([num_tokens], device="cpu"), extend_logprob_start_lens_cpu=torch.tensor([num_tokens], device="cpu"), - positions=torch.arange(num_tokens, device=self.device), + positions=positions, global_num_tokens_gpu=None, global_num_tokens_for_logprob_gpu=None, dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(), global_dp_buffer_len=None, - mrope_positions=( - self.mrope_positions[:, :num_tokens] if self.is_multimodal else None - ), + mrope_positions=mrope_positions, spec_algorithm=None, spec_info=None, capture_hidden_mode=CaptureHiddenMode.NULL, @@ -338,7 +338,7 @@ class PiecewiseCudaGraphRunner: # Attention backend self.model_runner.attn_backend.init_forward_metadata(forward_batch) with set_forward_context( - forward_batch, self.attention_layers, self.quant_config + forward_batch, self.attention_layers, self.quant_config, self.moe_layers ), disable_ca_comm(self.model_runner.tp_group): _ = self.model_runner.model.forward( forward_batch.input_ids, @@ -483,7 +483,7 @@ class PiecewiseCudaGraphRunner: kwargs = {} with set_forward_context( - forward_batch, self.attention_layers, self.quant_config + forward_batch, self.attention_layers, self.quant_config, self.moe_layers ): self.model_runner.model.forward( forward_batch.input_ids, @@ -608,7 +608,10 @@ class PiecewiseCudaGraphRunner: static_forward_batch = self.replay_prepare(forward_batch, **kwargs) # Replay with set_forward_context( - static_forward_batch, self.attention_layers, self.quant_config + static_forward_batch, + self.attention_layers, + self.quant_config, + self.moe_layers, ): with set_compiled(True): output = self.model_runner.model.forward( diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index 37c6cb0d1..782c4ff23 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -140,7 +140,7 @@ suites = { TestFile("test_dp_attention.py", 350), TestFile("test_load_weights_from_remote_instance.py", 72), TestFile("test_patch_torch.py", 19), - TestFile("test_piecewise_cuda_graph_2_gpu.py", 200), + TestFile("test_piecewise_cuda_graph_2_gpu.py", 400), TestFile("test_eagle_dp_attention.py", 200), ], "per-commit-4-gpu": [ diff --git a/test/srt/test_piecewise_cuda_graph_2_gpu.py b/test/srt/test_piecewise_cuda_graph_2_gpu.py index 5b6c40f42..100edae65 100644 --- a/test/srt/test_piecewise_cuda_graph_2_gpu.py +++ b/test/srt/test_piecewise_cuda_graph_2_gpu.py @@ -53,5 +53,49 @@ class TestPiecewiseCudaGraphQwen3OmniMOE(CustomTestCase): self.assertGreaterEqual(metrics["score"], 0.70) +class TestPiecewiseCudaGraphFusedMoE(CustomTestCase): + """Test piecewise CUDA graph with FusedMoE Backend""" + + @classmethod + def setUpClass(cls): + cls.model = "Qwen/Qwen3-Coder-30B-A3B-Instruct" + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--enable-piecewise-cuda-graph", + "--piecewise-cuda-graph-compiler", + "eager", + "--tp", + "2", + "--ep-size", + "2", + ], + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_gsm8k_accuracy(self): + """Test GSM8K accuracy with 8-shot setting""" + num_examples = 2000 + + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="mgsm_en", + num_examples=num_examples, + num_threads=min(num_examples, 1024), + ) + + metrics = run_eval(args) + print(f"GSM8K Accuracy: {metrics['score']:.3f}") + + self.assertGreaterEqual(metrics["score"], 0.90) + + if __name__ == "__main__": unittest.main()