EP Support for Piecewise Cuda Graph (#14164)

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

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
 

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)

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:

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:

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

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,
+)

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

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
+        # Collect attention layers and moe layers from the model
-        self.attention_layers = []
         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(

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 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_ids=input_ids,
-                input_embeds=(
+                input_embeds=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
-                ),
                 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=(
+                mrope_positions=mrope_positions,
-                    self.mrope_positions[:, :num_tokens] if self.is_multimodal else None
-                ),
                 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(

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": [

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