[MoE Refactor] Migrate SM90 Cutlass W4A16 to MoeRunner (#26489)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>

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

@@ -0,0 +1,174 @@
+"""FlashInfer SM90 cutlass mixed-input W4A16 MXFP4 MoE fused func.
+
+Registered for ``("none", "flashinfer_mxfp4")``. Drives FlashInfer's
+``cutlass_fused_moe(use_w4_group_scaling=True)`` (PR #3084 in flashinfer,
+SM90 only). Quant methods build the quant_info each forward and call
+``MoeRunner.run(dispatch_output, quant_info)``.
+
+Two production call sites share this fused func:
+  - GPT-OSS via :class:`Mxfp4MoEMethod` (input pad/output trim + per-expert
+    SwiGLU scalars + per-expert bias)
+  - DSv4 via :class:`Mxfp4FlashinferCutlassMoEMethod` (no bias, optional
+    SwiGLU scalars, no padding)
+
+The SM100 trtllm-gen path also lives under ``MoeRunnerBackend.FLASHINFER_MXFP4``
+but is intentionally left in the legacy bypass path for now; migrating it is a
+follow-up.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.srt.distributed import get_tp_group
+from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+    use_symmetric_memory,
+)
+from sglang.srt.layers.dp_attention import is_allocation_symmetric
+from sglang.srt.layers.moe.moe_runner.base import (
+    MoeQuantInfo,
+    MoeRunnerConfig,
+    register_fused_func,
+)
+from sglang.srt.utils import is_flashinfer_available
+from sglang.srt.utils.common import next_power_of_2
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput
+    from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
+
+
+@dataclass
+class FlashInferMxfp4CutlassMoeQuantInfo(MoeQuantInfo):
+    """Quantization payload for the SM90 cutlass W4A16 MXFP4 MoE path.
+
+    Weights and scales are pre-interleaved at load time via
+    ``interleave_moe_{weights,scales}_for_sm90_mixed_gemm``; this dataclass
+    only carries references plus the per-call routing/topology fields.
+    """
+
+    # Pre-interleaved weights (uint8, packed FP4)
+    w13_weight: torch.Tensor  # [E, 2*N, K/2]
+    w2_weight: torch.Tensor  # [E, K, N/2]
+
+    # Pre-interleaved E8M0 block scales (uint8; viewed as int32 at call time)
+    w13_weight_scale: torch.Tensor  # [E, 2*N, K/32]
+    w2_weight_scale: torch.Tensor  # [E, K, N/32]
+
+    # Per-expert bias. GPT-OSS has both; DSv4 leaves both None.
+    w13_bias: Optional[torch.Tensor] = None  # bf16 [E, 2*N]
+    w2_bias: Optional[torch.Tensor] = None  # bf16 [E, K]
+
+    # Per-expert SwiGLU scalars (fp32 [E]). Either all three are present
+    # (clamped SwiGLU) or all three are None (kernel default SwiGLU).
+    swiglu_alpha: Optional[torch.Tensor] = None
+    swiglu_beta: Optional[torch.Tensor] = None
+    swiglu_limit: Optional[torch.Tensor] = None
+
+    # TP/EP topology (forwarded to the FlashInfer kernel)
+    moe_tp_size: int = 1
+    moe_tp_rank: int = 0
+    moe_ep_size: int = 1
+    moe_ep_rank: int = 0
+
+    # GPT-OSS pads its input hidden dim up to the (pre-padded) loaded weight
+    # width and trims the output back. DSv4 leaves this as ``None`` (no pad).
+    padded_hidden: Optional[int] = None
+
+
+def _flashinfer_cutlass_fused_moe():
+    """Lazy import — keeps non-flashinfer wheels importable."""
+    if not is_flashinfer_available():
+        raise RuntimeError(
+            "flashinfer_mxfp4 runner backend requires flashinfer to be installed."
+        )
+    from flashinfer.fused_moe import cutlass_fused_moe
+    from flashinfer.fused_moe.core import ActivationType
+
+    return cutlass_fused_moe, ActivationType
+
+
+@register_fused_func("none", "flashinfer_mxfp4")
+def fused_experts_none_to_flashinfer_mxfp4(
+    dispatch_output: "StandardDispatchOutput",
+    quant_info: MoeQuantInfo,
+    runner_config: MoeRunnerConfig,
+) -> "StandardCombineInput":
+    """SM90 W4A16 MXFP4 fused expert forward pass.
+
+    Mirrors the legacy ``Mxfp4MoEMethod._apply_sm90_cutlass`` and DSv4's
+    ``Mxfp4FlashinferCutlassMoEMethod.apply`` exactly; difference vs those is
+    that all per-layer state arrives via ``quant_info`` rather than via the
+    layer module, so this function is layer-agnostic.
+    """
+    from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
+    from sglang.srt.layers.moe.topk import TopKOutputChecker
+
+    assert isinstance(
+        quant_info, FlashInferMxfp4CutlassMoeQuantInfo
+    ), f"Unexpected quant_info type for flashinfer_mxfp4: {type(quant_info)}"
+
+    flashinfer_cutlass_fused_moe, ActivationType = _flashinfer_cutlass_fused_moe()
+
+    x = dispatch_output.hidden_states
+    topk_output = dispatch_output.topk_output
+
+    # Under ``--moe-runner-backend flashinfer_mxfp4`` topk may be in bypassed
+    # form (the SM100 trtllm-gen path does routing internally). The cutlass
+    # SM90 path needs explicit topk_ids / topk_weights; materialize here.
+    if TopKOutputChecker.format_is_bypassed(topk_output):
+        topk_output = topk_output.to_standard()
+    topk_ids = topk_output.topk_ids
+    topk_weights = topk_output.topk_weights
+
+    # GPT-OSS: pad input hidden dim up to the loaded weight width. DSv4
+    # leaves padded_hidden as None (or equal to origin_hidden), no pad.
+    origin_hidden = x.shape[-1]
+    padded_hidden = quant_info.padded_hidden
+    do_pad = padded_hidden is not None and padded_hidden != origin_hidden
+    if do_pad:
+        x = torch.nn.functional.pad(
+            x,
+            (0, padded_hidden - origin_hidden),
+            mode="constant",
+            value=0.0,
+        )
+
+    out_hidden = padded_hidden if do_pad else origin_hidden
+    output_dtype = torch.bfloat16
+    with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()):
+        out = torch.empty(x.shape[0], out_hidden, dtype=output_dtype, device=x.device)
+
+    flashinfer_cutlass_fused_moe(
+        input=x,
+        token_selected_experts=topk_ids.to(torch.int),
+        token_final_scales=topk_weights,
+        fc1_expert_weights=quant_info.w13_weight,
+        fc2_expert_weights=quant_info.w2_weight,
+        output_dtype=output_dtype,
+        quant_scales=[
+            quant_info.w13_weight_scale.view(torch.int32),
+            quant_info.w2_weight_scale.view(torch.int32),
+        ],
+        fc1_expert_biases=quant_info.w13_bias,
+        fc2_expert_biases=quant_info.w2_bias,
+        swiglu_alpha=quant_info.swiglu_alpha,
+        swiglu_beta=quant_info.swiglu_beta,
+        swiglu_limit=quant_info.swiglu_limit,
+        tp_size=quant_info.moe_tp_size,
+        tp_rank=quant_info.moe_tp_rank,
+        ep_size=quant_info.moe_ep_size,
+        ep_rank=quant_info.moe_ep_rank,
+        use_w4_group_scaling=True,
+        activation_type=ActivationType.Swiglu,
+        tune_max_num_tokens=next_power_of_2(x.shape[0]),
+        output=out,
+    )
+
+    if do_pad:
+        out = out[:, :origin_hidden].contiguous()
+
+    return StandardCombineInput(hidden_states=out)

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

@@ -61,6 +61,8 @@ class MoeRunner:
             self.runner_core = None  # FlashInfer TRT-LLM only supports fused path
         elif runner_backend.is_flashinfer_cutedsl():
             self.runner_core = None  # FlashInfer CuteDSL only supports fused path
+        elif runner_backend.is_flashinfer_mxfp4():
+            self.runner_core = None  # FlashInfer MXFP4 only supports fused path
         else:
             raise NotImplementedError(f"Unsupported runner backend: {runner_backend}")
 

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

@@ -77,9 +77,7 @@ if is_flashinfer_available():
         nvfp4_block_scale_interleave,
         trtllm_fp4_block_scale_moe,
     )
-    from flashinfer.fused_moe import cutlass_fused_moe as flashinfer_cutlass_fused_moe
     from flashinfer.fused_moe.core import (
-        ActivationType,
         get_w2_permute_indices_with_cache,
     )
 
@@ -1056,78 +1054,47 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
             or moe_runner_backend.is_triton()
             or moe_runner_backend.is_marlin()
         ):
+            self.runner = MoeRunner(moe_runner_backend, moe_runner_config)
+        elif (
+            moe_runner_backend.is_flashinfer_mxfp4()
+            and self._fi_kernel == "cutlass_sm90"
+        ):
+            # Register the fused func at runner construction so the FusedOpPool
+            # lookup at `MoeRunner.__init__` finds it.
+            import sglang.srt.layers.moe.moe_runner.flashinfer_mxfp4  # noqa: F401
+
             self.runner = MoeRunner(moe_runner_backend, moe_runner_config)
         else:
-            # TODO(cwan): refactor other backends
+            # Legacy bypass path (e.g. SM100 trtllm-gen under flashinfer_mxfp4)
+            # routes through `apply` without a MoeRunner. TODO(cwan): migrate.
             pass
 
-    def _apply_sm90_cutlass(self, layer, x, topk_output):
+    def _apply_sm90_cutlass(self, layer, dispatch_output):
         """SM90 (Hopper) MXFP4 x BF16 MoE via FlashInfer's cutlass mixed-input
-        path (PR #3084). The fused kernel does GEMM1 + SwiGLU + GEMM2 in one
-        call; weights/scales were pre-interleaved at load time."""
-        from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
-        from sglang.srt.layers.moe.topk import TopKOutputChecker
+        path (PR #3084). Routed through the unified ``MoeRunner`` -- this
+        helper only builds the quant_info; the actual kernel call lives in
+        :mod:`sglang.srt.layers.moe.moe_runner.flashinfer_mxfp4`."""
+        from sglang.srt.layers.moe.moe_runner.flashinfer_mxfp4 import (
+            FlashInferMxfp4CutlassMoeQuantInfo,
+        )
 
-        # Under ``--moe-runner-backend flashinfer_mxfp4`` the SGLang TopK layer
-        # emits BypassedTopKOutput by default (the SM100 trtllm-gen kernel does
-        # routing internally). The cutlass kernel needs explicit topk_ids /
-        # topk_weights, so materialize them here when bypassed.
-        if TopKOutputChecker.format_is_bypassed(topk_output):
-            topk_output = topk_output.to_standard()
-        topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids
-
-        # Pad input hidden dim to the (already-padded) loaded weight width.
-        origin_hidden = x.shape[-1]
-        padded_hidden = self._padded_hidden
-        if padded_hidden != origin_hidden:
-            x = torch.nn.functional.pad(
-                x,
-                (0, padded_hidden - origin_hidden),
-                mode="constant",
-                value=0.0,
-            )
-
-        output_dtype = torch.bfloat16
-        # Output is allocated at padded width (kernel writes padded_hidden
-        # columns), then trimmed back to origin_hidden before returning.
-        with use_symmetric_memory(
-            get_tp_group(), disabled=not is_allocation_symmetric()
-        ):
-            out_padded = torch.empty(
-                x.shape[0], padded_hidden, dtype=output_dtype, device=x.device
-            )
-
-        flashinfer_cutlass_fused_moe(
-            input=x,
-            token_selected_experts=topk_ids.to(torch.int),
-            token_final_scales=topk_weights,
-            fc1_expert_weights=layer.w13_weight,  # uint8 [E, 2*N, K/2] interleaved
-            fc2_expert_weights=layer.w2_weight,  # uint8 [E, K, N/2] interleaved
-            output_dtype=output_dtype,
-            quant_scales=[
-                layer.w13_weight_scale.view(torch.int32),
-                layer.w2_weight_scale.view(torch.int32),
-            ],
-            fc1_expert_biases=layer.w13_weight_bias,  # bf16 [E, 2*N]
-            fc2_expert_biases=layer.w2_weight_bias,  # bf16 [E, K]
+        quant_info = FlashInferMxfp4CutlassMoeQuantInfo(
+            w13_weight=layer.w13_weight,
+            w2_weight=layer.w2_weight,
+            w13_weight_scale=layer.w13_weight_scale,
+            w2_weight_scale=layer.w2_weight_scale,
+            w13_bias=layer.w13_weight_bias,
+            w2_bias=layer.w2_weight_bias,
             swiglu_alpha=layer.swiglu_alpha,
             swiglu_beta=layer.swiglu_beta,
             swiglu_limit=layer.swiglu_limit,
-            tp_size=layer.moe_tp_size,
-            tp_rank=layer.moe_tp_rank,
-            ep_size=layer.moe_ep_size,
-            ep_rank=layer.moe_ep_rank,
-            use_w4_group_scaling=True,
-            activation_type=ActivationType.Swiglu,
-            tune_max_num_tokens=next_power_of_2(x.shape[0]),
-            output=out_padded,
+            moe_tp_size=layer.moe_tp_size,
+            moe_tp_rank=layer.moe_tp_rank,
+            moe_ep_size=layer.moe_ep_size,
+            moe_ep_rank=layer.moe_ep_rank,
+            padded_hidden=self._padded_hidden,
         )
-
-        if padded_hidden != origin_hidden:
-            out = out_padded[:, :origin_hidden].contiguous()
-        else:
-            out = out_padded
-        return StandardCombineInput(hidden_states=out)
+        return self.runner.run(dispatch_output, quant_info)
 
     def apply(
         self,
@@ -1183,7 +1150,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
             return self.runner.run(dispatch_output, quant_info)
 
         if self._fi_kernel == "cutlass_sm90":
-            return self._apply_sm90_cutlass(layer, x, topk_output)
+            return self._apply_sm90_cutlass(layer, dispatch_output)
         if self.use_flashinfer:
             # When bf16 mode is enabled, we don't need to quantize the input,
             # TRT-LLM automatically handles quantization in the kernel implementation and pipelines it with GEMM operations,

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

@@ -26,15 +26,8 @@ import torch
 from torch.nn import Module
 from torch.nn.parameter import Parameter
 
-from sglang.srt.distributed import get_tp_group
-from sglang.srt.distributed.device_communicators.pynccl_allocator import (
-    use_symmetric_memory,
-)
-from sglang.srt.layers.dp_attention import is_allocation_symmetric
-from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
 from sglang.srt.layers.moe.topk import TopKOutputChecker
 from sglang.srt.utils import is_flashinfer_available, log_info_on_rank0
-from sglang.srt.utils.common import next_power_of_2
 
 # Silence the TRT-LLM cutlass autotune trace embedded inside FlashInfer's
 # cutlass_fused_moe. Its C++ logger reads TLLM_LOG_LEVEL on first kernel launch;
@@ -42,9 +35,6 @@ from sglang.srt.utils.common import next_power_of_2
 os.environ.setdefault("TLLM_LOG_LEVEL", "INFO")
 
 if is_flashinfer_available():
-    from flashinfer.fused_moe import cutlass_fused_moe as flashinfer_cutlass_fused_moe
-    from flashinfer.fused_moe.core import ActivationType
-
     try:
         from flashinfer.fused_moe import (
             interleave_moe_scales_for_sm90_mixed_gemm,
@@ -123,6 +113,9 @@ class Mxfp4FlashinferCutlassMoEMethod:
         )
 
     def create_moe_runner(self, layer: Module, moe_runner_config) -> None:
+        from sglang.srt.layers.moe.moe_runner.runner import MoeRunner
+        from sglang.srt.layers.moe.utils import MoeRunnerBackend
+
         self.moe_runner_config = moe_runner_config
 
         # DSv4 uses standard SwiGLU plus a config-driven activation clamp.
@@ -150,6 +143,12 @@ class Mxfp4FlashinferCutlassMoEMethod:
             self._swiglu_beta_tensor = None
             self._swiglu_limit_tensor = None
 
+        # Register the fused func at runner construction so the FusedOpPool
+        # lookup at `MoeRunner.__init__` finds it.
+        import sglang.srt.layers.moe.moe_runner.flashinfer_mxfp4  # noqa: F401
+
+        self.runner = MoeRunner(MoeRunnerBackend.FLASHINFER_MXFP4, moe_runner_config)
+
     def process_weights_after_loading(self, layer: Module) -> None:
         from sglang.srt.layers.quantization.utils import reorder_w1w3_to_w3w1
 
@@ -218,46 +217,30 @@ class Mxfp4FlashinferCutlassMoEMethod:
         layer: Module,
         dispatch_output: "DispatchOutput",
     ) -> "CombineInput":
+        from sglang.srt.layers.moe.moe_runner.flashinfer_mxfp4 import (
+            FlashInferMxfp4CutlassMoeQuantInfo,
+        )
+
+        # DSv4 always feeds StandardDispatchOutput; the fused func tolerates
+        # bypassed too but we keep the strict check here as a contract guard.
         topk_output = dispatch_output.topk_output
         if not TopKOutputChecker.format_is_standard(topk_output):
             raise ValueError(f"Unsupported topk output format: {topk_output.format}")
 
-        x = dispatch_output.hidden_states
-        topk_weights = topk_output.topk_weights
-        topk_ids = topk_output.topk_ids
-
-        output_dtype = torch.bfloat16
-        with use_symmetric_memory(
-            get_tp_group(), disabled=not is_allocation_symmetric()
-        ):
-            out = torch.empty(
-                x.shape[0], x.shape[-1], dtype=output_dtype, device=x.device
-            )
-
-        flashinfer_cutlass_fused_moe(
-            input=x,
-            token_selected_experts=topk_ids.to(torch.int),
-            token_final_scales=topk_weights,
-            fc1_expert_weights=layer.w13_weight,
-            fc2_expert_weights=layer.w2_weight,
-            output_dtype=output_dtype,
-            quant_scales=[
-                layer.w13_weight_scale_inv.view(torch.int32),
-                layer.w2_weight_scale_inv.view(torch.int32),
-            ],
-            fc1_expert_biases=None,  # DSv4 has no MoE expert bias.
-            fc2_expert_biases=None,
+        quant_info = FlashInferMxfp4CutlassMoeQuantInfo(
+            w13_weight=layer.w13_weight,
+            w2_weight=layer.w2_weight,
+            w13_weight_scale=layer.w13_weight_scale_inv,
+            w2_weight_scale=layer.w2_weight_scale_inv,
+            w13_bias=None,  # DSv4 has no MoE expert bias.
+            w2_bias=None,
             swiglu_alpha=self._swiglu_alpha_tensor,  # ones: standard SiLU gate
             swiglu_beta=self._swiglu_beta_tensor,  # zeros: standard up
             swiglu_limit=self._swiglu_limit_tensor,
-            tp_size=layer.moe_tp_size,
-            tp_rank=layer.moe_tp_rank,
-            ep_size=layer.moe_ep_size,
-            ep_rank=layer.moe_ep_rank,
-            use_w4_group_scaling=True,
-            activation_type=ActivationType.Swiglu,
-            tune_max_num_tokens=next_power_of_2(x.shape[0]),
-            output=out,
+            moe_tp_size=layer.moe_tp_size,
+            moe_tp_rank=layer.moe_tp_rank,
+            moe_ep_size=layer.moe_ep_size,
+            moe_ep_rank=layer.moe_ep_rank,
+            padded_hidden=None,  # DSv4 hidden_size is already a multiple of 128.
         )
-
-        return StandardCombineInput(hidden_states=out)
+        return self.runner.run(dispatch_output, quant_info)