[MoE Refactor] Migrate flashinfer_cutedsl + DeepEP to MoeRunner (#25525)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

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

@@ -105,6 +105,12 @@ class DeepEPMoE(FusedMoE):
             and envs.SGLANG_DEEPEP_BF16_DISPATCH.get()
         ):
             self.deprecate_flag = True
+        elif (
+            get_moe_runner_backend().is_flashinfer_cutedsl()
+            and quant_config is not None
+            and quant_config.get_name() == "modelopt_fp4"
+        ):
+            self.deprecate_flag = True
         else:
             self.deprecate_flag = False
 
@@ -134,15 +140,9 @@ class DeepEPMoE(FusedMoE):
             self.deepep_mode.enable_low_latency()
             and not _is_npu
             and not _is_hip
-            and not (
-                get_moe_runner_backend().is_flashinfer_cutedsl()
-                and self.quant_config is not None
-                and self.quant_config.get_name() == "modelopt_fp4"
-            )
             and quant_config is not None
         ):
-            # AMD HIP, NPU supports low_latency deepep without deepgemm
+            # AMD HIP and NPU support low_latency DeepEP without DeepGEMM.
-            # NV FP4 quantization with flashinfer_cutedsl also supports low_latency deepep without deepgemm
             # Unquantized draft MoE uses BF16 DeepEP dispatch and a local fallback.
             assert (
                 deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
@@ -220,12 +220,6 @@ class DeepEPMoE(FusedMoE):
         elif DispatchOutputChecker.format_is_deepep_ll(dispatch_output):
             if self.quant_config is None:
                 output = self.forward_unquantized_deepep_ll(dispatch_output)
-            elif (
-                get_moe_runner_backend().is_flashinfer_cutedsl()
-                and self.quant_config is not None
-                and self.quant_config.get_name() == "modelopt_fp4"
-            ):
-                output = self.forward_flashinfer_cutedsl(dispatch_output)
             elif self.use_w4afp8:
                 output = self.forward_cutlass_w4afp8_masked(dispatch_output)
             else:
@@ -288,22 +282,6 @@ class DeepEPMoE(FusedMoE):
             output = output + w2_bias.unsqueeze(1)
         return output.masked_fill(~valid_mask, 0)
 
-    def forward_flashinfer_cutedsl(
-        self,
-        dispatch_output: DeepEPLLDispatchOutput,
-    ):
-        hidden_states, hidden_states_scale, _, _, masked_m, _ = dispatch_output
-        assert self.quant_method is not None
-        assert self.moe_runner_config.activation == "silu"
-
-        output = self.quant_method.apply_without_routing_weights(
-            layer=self,
-            x=(hidden_states, hidden_states_scale),
-            masked_m=masked_m,
-            moe_runner_config=self.moe_runner_config,
-        )
-        return output
-
     def forward_cutlass_w4afp8(
         self,
         dispatch_output: DeepEPNormalDispatchOutput,

python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import logging
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Optional
 
 import torch
 
@@ -14,7 +14,10 @@ from sglang.srt.layers.moe.moe_runner.base import (
 from sglang.srt.utils.common import log_info_on_rank0, print_warning_once
 
 if TYPE_CHECKING:
+    from sglang.srt.batch_overlap.single_batch_overlap import DownGemmOverlapArgs
     from sglang.srt.layers.moe.token_dispatcher import (
+        DeepEPLLCombineInput,
+        DeepEPLLDispatchOutput,
         StandardCombineInput,
         StandardDispatchOutput,
     )
@@ -283,28 +286,49 @@ def ensure_cutedsl_wrapper(layer: torch.nn.Module) -> None:
 
 @dataclass
 class CuteDslFp4MoeQuantInfo(MoeQuantInfo):
-    """Quantization payload consumed by FlashInfer CuteDSL FP4 MoE kernels."""
+    """Quantization payload for FlashInfer CuteDSL FP4 MoE kernels.
 
-    # Lazily-created CuteDslMoEWrapper (stashed on layer)
+    Shared by the two CuteDSL runner entries:
-    wrapper: Any
 
-    # Weights (uint8 FP4 packed)
+    * "v2" standard path (a2a=``none``/``flashinfer``): consumed by the
+      ``@register_fused_func("none", "flashinfer_cutedsl")`` entry, which
+      drives ``CuteDslMoEWrapper.run``. Weights are ``[Up, Gate]``
+      interleaved with MMA-layout blockscales. ``wrapper`` is set;
+      ``w*_scale`` are scalarized.
+
+    * "v1" DeepEP low-latency path (a2a=``deepep``): consumed by the
+      ``@register_fused_func("deepep", "flashinfer_cutedsl")`` entry,
+      which drives ``flashinfer_cutedsl_moe_masked``. Weights are
+      ``[Gate, Up]`` non-interleaved with swizzled blockscales.
+      ``wrapper`` is ``None``; ``w*_scale`` are per-expert.
+    """
+
+    # FP4 packed weights (uint8)
     w13_weight: torch.Tensor
     w2_weight: torch.Tensor
 
-    # Block-scale factors
+    # Block-scale factors (MMA layout for v2, swizzled for v1)
     w13_weight_sf: torch.Tensor
     w2_weight_sf: torch.Tensor
 
-    # Per-expert GEMM scales
+    # Per-expert GEMM dequant alphas (scalarized for v2, per-expert for v1)
     w1_alpha: torch.Tensor
     w2_alpha: torch.Tensor
 
-    # Intermediate quantization scale (fc2 input)
+    # Activation quant scales (1 / raw_input_scale).
-    fc2_input_scale: torch.Tensor
+    #   - a1_scale: quantizes hidden_states before GEMM1
+    #   - a2_scale: quantizes GEMM1 output before GEMM2 (a.k.a. fc2 input)
+    a1_scale: torch.Tensor
+    a2_scale: torch.Tensor
 
-    # Activation quantization scale (scalarized)
+    # v2 only: lazily-created CuteDslMoEWrapper (``None`` on the v1 path).
-    input_scale: torch.Tensor
+    wrapper: Optional[Any] = None
+
+    # v1 only: ``True`` when DeepEP pre-quantizes activations to NVFP4.
+    use_nvfp4_dispatch: bool = False
+
+    # v1 only: SBO down-GEMM overlap args.
+    down_gemm_overlap_args: Optional["DownGemmOverlapArgs"] = None
 
 
 @register_fused_func("none", "flashinfer_cutedsl")
@@ -318,6 +342,7 @@ def fused_experts_none_to_flashinfer_cutedsl_fp4(
     from sglang.srt.layers.quantization.fp4_utils import fp4_quantize
 
     assert runner_config.activation == "silu", "Only silu is supported for CuteDSL MoE."
+    assert quant_info.wrapper is not None, "CuteDSL v2 path requires CuteDslMoEWrapper."
 
     hidden_states = dispatch_output.hidden_states
     topk_output = dispatch_output.topk_output
@@ -330,7 +355,7 @@ def fused_experts_none_to_flashinfer_cutedsl_fp4(
 
     x_fp4, x_sf = fp4_quantize(
         hidden_states,
-        quant_info.input_scale,
+        quant_info.a1_scale,
         sf_vec_size=_FP4_SF_VEC_SIZE,
         is_sf_swizzled_layout=False,
     )
@@ -343,10 +368,75 @@ def fused_experts_none_to_flashinfer_cutedsl_fp4(
         w1_weight=quant_info.w13_weight,
         w1_weight_sf=quant_info.w13_weight_sf,
         w1_alpha=quant_info.w1_alpha,
-        fc2_input_scale=quant_info.fc2_input_scale,
+        fc2_input_scale=quant_info.a2_scale,
         w2_weight=quant_info.w2_weight,
         w2_weight_sf=quant_info.w2_weight_sf,
         w2_alpha=quant_info.w2_alpha,
     )
 
     return StandardCombineInput(hidden_states=output)
+
+
+@register_fused_func("deepep", "flashinfer_cutedsl")
+def fused_experts_deepep_to_flashinfer_cutedsl_fp4(
+    dispatch_output: DeepEPLLDispatchOutput,
+    quant_info: CuteDslFp4MoeQuantInfo,
+    runner_config: MoeRunnerConfig,
+) -> DeepEPLLCombineInput:
+    from sglang.srt.layers.moe.flashinfer_cutedsl_moe import (
+        flashinfer_cutedsl_moe_masked,
+    )
+    from sglang.srt.layers.moe.token_dispatcher.deepep import DeepEPLLCombineInput
+
+    assert runner_config.activation == "silu", "Only silu is supported for CuteDSL MoE."
+    assert (
+        not runner_config.apply_router_weight_on_input
+    ), "apply_router_weight_on_input is not supported for Flashinfer"
+
+    hidden_states, hidden_states_scale, _, _, masked_m, _ = dispatch_output
+
+    # flashinfer_cutedsl_moe_masked reinterprets scales as float8_e4m3fn.
+    # Same-dtype .view is a no-op; only wider dtypes (e.g. int32-packed
+    # UE8M0) need stride(-1)==1.
+    if (
+        quant_info.use_nvfp4_dispatch
+        and hidden_states_scale is not None
+        and hidden_states_scale.element_size() != 1
+        and hidden_states_scale.stride(-1) != 1
+    ):
+        raise AssertionError(
+            f"NVFP4 dispatch scale has stride(-1)={hidden_states_scale.stride(-1)}, "
+            f"dtype={hidden_states_scale.dtype}; .view(float8_e4m3fn) requires stride(-1)==1. "
+            "Try SGLANG_MOE_NVFP4_DISPATCH=0 or check DeepEP version."
+        )
+
+    overlap = quant_info.down_gemm_overlap_args
+    output = flashinfer_cutedsl_moe_masked(
+        hidden_states=(hidden_states, hidden_states_scale),
+        input_global_scale=(
+            None if quant_info.use_nvfp4_dispatch else quant_info.a1_scale
+        ),
+        w1=quant_info.w13_weight,
+        w1_blockscale=quant_info.w13_weight_sf,
+        w1_alpha=quant_info.w1_alpha,
+        w2=quant_info.w2_weight,
+        a2_global_scale=quant_info.a2_scale,
+        w2_blockscale=quant_info.w2_weight_sf,
+        w2_alpha=quant_info.w2_alpha,
+        masked_m=masked_m,
+        **(
+            dict(
+                down_sm_count=overlap.num_sms,
+                down_signals=overlap.signal,
+                down_start_event=overlap.start_event,
+            )
+            if overlap is not None
+            else {}
+        ),
+    )
+
+    return DeepEPLLCombineInput(
+        hidden_states=output,
+        topk_ids=dispatch_output.topk_ids,
+        topk_weights=dispatch_output.topk_weights,
+    )

python/sglang/srt/layers/moe/moe_runner/runner.py

@@ -163,11 +163,9 @@ class MoeRunner:
     def set_overlap_args(
         self, down_gemm_overlap_args: DownGemmOverlapArgs, meta_overlap_args: dict
     ):
-        assert self.fused_func is None, "Fused func is not supported for overlap args"
         self.down_gemm_overlap_args = down_gemm_overlap_args
         self.meta_overlap_args = meta_overlap_args
 
     def clear_overlap_args(self) -> None:
-        assert self.fused_func is None, "Fused func is not supported for overlap args"
         self.down_gemm_overlap_args = None
         self.meta_overlap_args = None

python/sglang/srt/layers/quantization/modelopt_quant.py

@@ -67,7 +67,6 @@ from sglang.srt.utils.custom_op import register_custom_op
 from sglang.srt.utils.patch_torch import register_fake_if_exists
 
 if TYPE_CHECKING:
-    from sglang.srt.batch_overlap.single_batch_overlap import DownGemmOverlapArgs
     from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
     from sglang.srt.layers.moe.token_dispatcher import (
         CombineInput,
@@ -1551,24 +1550,24 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
     # ----- CuteDSL v1 vs v2 path helpers -----
     #
     # "v1": cutedsl + deepep low-latency.
-    #   - Bypasses MoeRunner entirely; calls apply_without_routing_weights ->
+    #   - MoeRunner fused func calls flashinfer_cutedsl_moe_masked
-    #     flashinfer_cutedsl_moe_masked (grouped_gemm_nt_masked).
+    #     (grouped_gemm_nt_masked).
     #   - Expects W13 in default [Gate, Up] order, NOT interleaved.
     #   - Uses swizzled blockscales directly (w13_blockscale_swizzled).
     #
     # "v2" (standard): cutedsl + none/flashinfer a2a.
-    #   - Uses MoeRunner with @register_fused_func CuteDslMoEWrapper kernels.
+    #   - MoeRunner fused func calls CuteDslMoEWrapper kernels.
     #   - Expects W13 in [Up, Gate] order, interleaved in 64-row chunks.
     #   - Uses MMA-layout blockscales (w13_blockscale_mma).
 
     @property
     def _is_cutedsl_v1_deepep(self) -> bool:
-        """CuteDSL v1 + DeepEP low-latency path (no MoeRunner)."""
+        """CuteDSL v1 + DeepEP low-latency path (masked grouped GEMM)."""
         return is_flashinfer_cutedsl_v1_path()
 
     @property
     def _is_cutedsl_v2_standard(self) -> bool:
-        """New CuteDSL standard path (a2a=none or flashinfer, uses MoeRunner)."""
+        """CuteDSL v2 standard path (a2a=none or flashinfer, uses CuteDslMoEWrapper)."""
         return self.enable_flashinfer_cutedsl_moe and not self._is_cutedsl_v1_deepep
 
     def create_weights(
@@ -1998,11 +1997,6 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
         if moe_runner_backend.is_flashinfer_cutedsl():
             import sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl  # noqa: F401 – triggers @register_fused_func
 
-            # CuteDSL v1 (deepep) uses the apply_without_routing_weights
-            # path (flashinfer_cutedsl_moe_masked) and does not need a MoeRunner.
-            if self._is_cutedsl_v1_deepep:
-                return
-
         if not moe_runner_backend.is_flashinfer_cutlass():
             self.runner = MoeRunner(moe_runner_backend, moe_runner_config)
 
@@ -2013,10 +2007,10 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
     ) -> CombineInput:
         from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
 
-        x = dispatch_output.hidden_states
+        # Note: dispatch_output may be a DeepEPLLDispatchOutput (no topk_output
-        x_sf = dispatch_output.hidden_states_scale
+        # attribute -- topk_ids/topk_weights live directly on the dispatch
-        topk_output = dispatch_output.topk_output
+        # tuple). Defer per-attribute access to the branches that actually
-
+        # consume them.
         activation = self.moe_runner_config.activation
 
         assert (
@@ -2054,33 +2048,49 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
 
             return self.runner.run(dispatch_output, quant_info)
 
-        # CuteDSL v2 standard path (a2a=none/flashinfer).
-        # The v1 (deepep) path never reaches apply(); it goes through
-        # apply_without_routing_weights instead.
         if self.enable_flashinfer_cutedsl_moe:
             from sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl import (
                 CuteDslFp4MoeQuantInfo,
                 ensure_cutedsl_wrapper,
             )
 
+            if self._is_cutedsl_v1_deepep:
+                # v1 path: DeepEP low-latency + flashinfer_cutedsl_moe_masked.
+                # Weights are [Gate, Up] (non-interleaved) with swizzled blockscales.
+                quant_info = CuteDslFp4MoeQuantInfo(
+                    w13_weight=layer.w13_weight,
+                    w2_weight=layer.w2_weight,
+                    w13_weight_sf=layer.w13_blockscale_swizzled,
+                    w2_weight_sf=layer.w2_blockscale_swizzled,
+                    w1_alpha=layer.g1_alphas,
+                    w2_alpha=layer.g2_alphas,
+                    a1_scale=layer.w13_input_scale_quant,
+                    a2_scale=layer.w2_input_scale_quant,
+                    use_nvfp4_dispatch=MOE_NVFP4_DISPATCH,
+                    down_gemm_overlap_args=getattr(
+                        self.runner, "down_gemm_overlap_args", None
+                    ),
+                )
+                return self.runner.run(dispatch_output, quant_info)
+
+            # v2 standard path (a2a=none/flashinfer): uses CuteDslMoEWrapper
+            # with [Up, Gate] interleaved weights and MMA blockscales.
             ensure_cutedsl_wrapper(layer)
             w1_alpha, fc2_input_scale, w2_alpha = layer._cutedsl_scales
-            w1_weight_sf = getattr(
-                layer, "w13_blockscale_mma", layer.w13_blockscale_swizzled
-            )
-            w2_weight_sf = getattr(
-                layer, "w2_blockscale_mma", layer.w2_blockscale_swizzled
-            )
             quant_info = CuteDslFp4MoeQuantInfo(
-                wrapper=layer._cutedsl_wrapper,
                 w13_weight=layer.w13_weight,
                 w2_weight=layer.w2_weight,
-                w13_weight_sf=w1_weight_sf,
+                w13_weight_sf=getattr(
-                w2_weight_sf=w2_weight_sf,
+                    layer, "w13_blockscale_mma", layer.w13_blockscale_swizzled
+                ),
+                w2_weight_sf=getattr(
+                    layer, "w2_blockscale_mma", layer.w2_blockscale_swizzled
+                ),
                 w1_alpha=w1_alpha,
                 w2_alpha=w2_alpha,
-                fc2_input_scale=fc2_input_scale,
+                a1_scale=layer._cutedsl_input_scale,
-                input_scale=layer._cutedsl_input_scale,
+                a2_scale=fc2_input_scale,
+                wrapper=layer._cutedsl_wrapper,
             )
             return self.runner.run(dispatch_output, quant_info)
 
@@ -2092,6 +2102,9 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
             ), "apply_router_weight_on_input is not supported for Flashinfer"
             # TRTLLM Cutlass moe takes in activations in BF16/Half/nvfp4 precision
             # and fp4 quantized weights loaded from the checkpoint
+            x = dispatch_output.hidden_states
+            x_sf = dispatch_output.hidden_states_scale
+            topk_output = dispatch_output.topk_output
             topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids
 
             output_dtype = torch.bfloat16
@@ -2145,6 +2158,8 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
 
         from sglang.srt.layers.moe.cutlass_moe import cutlass_moe_fp4
 
+        x = dispatch_output.hidden_states
+        topk_output = dispatch_output.topk_output
         topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids
         output = cutlass_moe_fp4(
             a=x,
@@ -2163,75 +2178,3 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
         ).to(x.dtype)
         # Scale by routed_scaling_factor is fused into select_experts.
         return StandardCombineInput(hidden_states=output)
-
-    def apply_without_routing_weights(
-        self,
-        layer: FusedMoE,
-        x: tuple[torch.Tensor, Optional[torch.Tensor]],
-        masked_m: torch.Tensor,
-        moe_runner_config: MoeRunnerConfig,
-    ) -> torch.Tensor:
-        """CuteDSL v1 (deepep low-latency) path.
-
-        Called by the DeepEP dispatcher instead of apply().  Uses
-        flashinfer_cutedsl_moe_masked (grouped_gemm_nt_masked) directly,
-        bypassing MoeRunner.  Weights must be in default [Gate, Up] order
-        and NOT interleaved -- see _is_cutedsl_v1_deepep guards in
-        process_weights_after_loading and load_up_proj_weight_first.
-        """
-        assert (
-            moe_runner_config.activation == "silu"
-        ), "Only SiLU activation is supported."
-
-        assert self.enable_flashinfer_cutedsl_moe, "only support flashinfer cutedsl moe"
-        assert (
-            not moe_runner_config.apply_router_weight_on_input
-        ), "apply_router_weight_on_input is not supported for Flashinfer"
-
-        from sglang.srt.layers.moe.flashinfer_cutedsl_moe import (
-            flashinfer_cutedsl_moe_masked,
-        )
-
-        # flashinfer_cutedsl_moe_masked reinterprets scales as float8_e4m3fn.
-        # Same-dtype .view is a no-op; only wider dtypes (e.g. int32-packed
-        # UE8M0) need stride(-1)==1.
-        if (
-            MOE_NVFP4_DISPATCH
-            and x[1] is not None
-            and x[1].element_size() != 1
-            and x[1].stride(-1) != 1
-        ):
-            raise AssertionError(
-                f"NVFP4 dispatch scale has stride(-1)={x[1].stride(-1)}, "
-                f"dtype={x[1].dtype}; .view(float8_e4m3fn) requires stride(-1)==1. "
-                "Try SGLANG_MOE_NVFP4_DISPATCH=0 or check DeepEP version."
-            )
-
-        down_gemm_overlap_args: Optional[DownGemmOverlapArgs] = getattr(
-            layer, "down_gemm_overlap_args", None
-        )
-
-        out = flashinfer_cutedsl_moe_masked(
-            hidden_states=x,
-            input_global_scale=(
-                None if MOE_NVFP4_DISPATCH else layer.w13_input_scale_quant
-            ),
-            w1=layer.w13_weight,
-            w1_blockscale=layer.w13_blockscale_swizzled,
-            w1_alpha=layer.g1_alphas,
-            w2=layer.w2_weight,
-            a2_global_scale=layer.w2_input_scale_quant,
-            w2_blockscale=layer.w2_blockscale_swizzled,
-            w2_alpha=layer.g2_alphas,
-            masked_m=masked_m,
-            **(
-                dict(
-                    down_sm_count=down_gemm_overlap_args.num_sms,
-                    down_signals=down_gemm_overlap_args.signal,
-                    down_start_event=down_gemm_overlap_args.start_event,
-                )
-                if down_gemm_overlap_args is not None
-                else {}
-            ),
-        )
-        return out

test/registered/moe/test_cutedsl_moe.py

@@ -780,7 +780,8 @@ class TestCuteDslV2(unittest.TestCase):
 class TestCuteDslV1(unittest.TestCase):
     """Correctness tests for the CuteDSL v1 (deepep) path.
 
-    The v1 path (apply_without_routing_weights -> flashinfer_cutedsl_moe_masked)
+    The v1 path (flashinfer_cutedsl_moe_masked, dispatched via the
+    @register_fused_func("deepep", "flashinfer_cutedsl") MoeRunner entry)
     is used when --moe-runner-backend flashinfer_cutedsl and --moe-a2a-backend
     deepep are combined.  It expects:
       - W13 in default [Gate, Up] order (load_up_proj_weight_first = False)