MoE Refactor: Refactor modelopt_quant.py -> flashinfer_trllm.py

[b8zhong]

Motivation

Followup on #15151 (comment), and part of #8715

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[b8zhong]

/tag-and-rerun-ci one more time?

[ch-wan]

why not using runner?

[b8zhong]

Yeah, I think it would be good. The problem is (to my undersatnding), since we can only register 1 fused function for flashinfer_trtllm, while there might either be trtllm_fp4_block_scale_moe, or fused_experts_none_to_flashinfer_trtllm_fp8`.

Therefore to simplify the codes I didn't, other wise I suggest maybe the diff below to get the flashinfer_trtllm to register two different fused backends. But I am open to either.

git --no-pager diff
diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
index 74c56761a..c2b037bda 100644
--- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
+++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
@@ -207,7 +207,6 @@ class FlashInferTrtllmFp8MoeQuantInfo(MoeQuantInfo):
     use_routing_scales_on_input: bool = False
 
 
-@register_fused_func("none", "flashinfer_trtllm")
 def fused_experts_none_to_flashinfer_trtllm_fp8(
     dispatch_output: StandardDispatchOutput,
     quant_info: FlashInferTrtllmFp8MoeQuantInfo,
@@ -478,3 +477,21 @@ def fused_experts_none_to_flashinfer_trtllm_fp4(
     )[0]
 
     return StandardCombineInput(hidden_states=result)
+
+
+@register_fused_func("none", "flashinfer_trtllm")
+def fused_experts_none_to_flashinfer_trtllm(
+    dispatch_output: StandardDispatchOutput,
+    quant_info: MoeQuantInfo,
+    runner_config: MoeRunnerConfig,
+) -> StandardCombineInput:
+    """Dispatch to FP8 or FP4 FlashInfer TRT-LLM MoE based on quant_info type."""
+    if isinstance(quant_info, FlashInferTrtllmFp4MoeQuantInfo):
+        return fused_experts_none_to_flashinfer_trtllm_fp4(
+            dispatch_output, quant_info, runner_config
+        )
+    return fused_experts_none_to_flashinfer_trtllm_fp8(
+        dispatch_output,
+        cast(FlashInferTrtllmFp8MoeQuantInfo, quant_info),
+        runner_config,
+    )
diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py
index a71c0bc65..3b13ce408 100755
--- a/python/sglang/srt/layers/quantization/modelopt_quant.py
+++ b/python/sglang/srt/layers/quantization/modelopt_quant.py
@@ -1496,6 +1496,10 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
         self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
     ):
         self.moe_runner_config = moe_runner_config
+        if get_moe_runner_backend().is_flashinfer_trtllm():
+            self.runner = MoeRunner(
+                MoeRunnerBackend.FLASHINFER_TRTLLM, moe_runner_config
+            )
 
     def apply(
         self,
@@ -1514,11 +1518,11 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
         ), f"{activation=} missing from {ACT_STR_TO_TYPE_MAP.keys()=}"
         moe_runner_config = self.moe_runner_config
 
-        # FlashInfer TRTLLM FP4 path - check if layer has shuffled weights
+        # FlashInfer TRTLLM FP4 path - layer has shuffled weights only when
+        # backend is flashinfer_trtllm
         if hasattr(layer, "gemm1_weights_fp4_shuffled"):
             from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
                 FlashInferTrtllmFp4MoeQuantInfo,
-                fused_experts_none_to_flashinfer_trtllm_fp4,
             )
             from sglang.srt.layers.moe.utils import RoutingMethodType
 
@@ -1543,9 +1547,7 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
                 routing_method_type=routing_method_type,
             )
 
-            return fused_experts_none_to_flashinfer_trtllm_fp4(
-                dispatch_output, quant_info, moe_runner_config
-            )
+            return self.runner.run(dispatch_output, quant_info)
 
         if self.enable_flashinfer_cutlass_moe:
             assert (