[Fix] Correct W4AFP8 DeepEP scaling and mode-specific dtypes (#33669)

python/sglang/srt/layers/moe/cutlass_w4a8_moe.py

@@ -425,6 +425,9 @@ def cutlass_w4a8_moe_deepep_normal(
         topk_weights,
         topk,
         c2.shape[1],
+        # DeepEP models apply routed_scaling_factor after the cross-rank
+        # combine, so this rank-local reduction must remain unscaled.
+        1.0,
         BLOCK_SIZE=512,
     )
 

python/sglang/srt/layers/moe/token_dispatcher/deepep.py

@@ -494,6 +494,8 @@ class _DeepEPDispatcherImplBase:
 
 
 class _DeepEPDispatcherImplNormal(_DeepEPDispatcherImplBase):
+    dispatch_mode = DeepEPMode.NORMAL
+
     def __init__(self, async_finish: bool, **kwargs):
         super().__init__(**kwargs)
 
@@ -654,6 +656,8 @@ class _DeepEPDispatcherImplNormal(_DeepEPDispatcherImplBase):
 
 
 class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
+    dispatch_mode = DeepEPMode.LOW_LATENCY
+
     def __init__(self, return_recv_hook: bool, **kwargs):
         super().__init__(**kwargs)
 

python/sglang/srt/layers/moe/utils.py

@@ -232,10 +232,11 @@ def get_deepep_output_dtype(self) -> DispatcherOutputDtype:
     0. Parse server argument.
     1. Parse deprecated environment variables.
     2. If quant_config contains input_global_scale → NVFP4 path.
-    3. Parse quant config
-    4. If flashinfer_cutedsl or is_cutlass backend is active → BF16 (it quantizes hidden_states internally).
-    5. Otherwise default for NPU → BF16 (the default for NPU).
-    6. Otherwise → FP8 (the default for most models like DeepSeek-V3).
+    3. Parse a mode-specific dtype from quant_config.
+    4. Parse a generic dtype from quant_config.
+    5. If flashinfer_cutedsl or is_cutlass backend is active → BF16 (it quantizes hidden_states internally).
+    6. Otherwise default for NPU → BF16 (the default for NPU).
+    7. Otherwise → FP8 (the default for most models like DeepSeek-V3).
     """
 
     # 0. Parse server argument.
@@ -258,12 +259,23 @@ def get_deepep_output_dtype(self) -> DispatcherOutputDtype:
         if input_global_scale is not None:
             return DispatcherOutputDtype.NVFP4
 
-        # 3. Parse quant config to determine the output dtype of dispatcher
+        # 3. Some MoE kernels require different wire formats for prefill and
+        # decode. Prefer a mode-specific override when the dispatcher exposes
+        # its concrete mode (normal or low_latency).
+        dispatch_mode = getattr(self, "dispatch_mode", None)
+        if dispatch_mode is not None:
+            mode_dispatcher_output_dtype = self.quant_config.get(
+                f"{dispatch_mode.value}_dispatcher_output_dtype", None
+            )
+            if mode_dispatcher_output_dtype is not None:
+                return DispatcherOutputDtype(mode_dispatcher_output_dtype)
+
+        # 4. Parse quant config to determine the output dtype of dispatcher
         dispatcher_output_dtype = self.quant_config.get("dispatcher_output_dtype", None)
         if dispatcher_output_dtype is not None:
             return DispatcherOutputDtype(dispatcher_output_dtype)
 
-    # 4. flashinfer_cutedsl / cutlass / humming expects BF16 dispatch
+    # 5. flashinfer_cutedsl / cutlass / humming expects BF16 dispatch
     if (
         get_moe_runner_backend().is_flashinfer_cutedsl()
         or get_moe_runner_backend().is_cutlass()
@@ -271,11 +283,11 @@ def get_deepep_output_dtype(self) -> DispatcherOutputDtype:
     ):
         return DispatcherOutputDtype.BF16
 
-    # 5. Default on NPU → BF16
+    # 6. Default on NPU → BF16
     if _is_npu:
         return DispatcherOutputDtype.BF16
 
-    # 6. Default → FP8
+    # 7. Default → FP8
     return DispatcherOutputDtype.FP8
 
 

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

@@ -283,6 +283,17 @@ class W4AFp8MoEMethod(FusedMoEMethodBase):
         )
         layer.w2_input_scale = Parameter(new_w2_input_scale, requires_grad=False)
 
+        if hasattr(layer, "dispatcher"):
+            # The normal kernel requantizes BF16 inputs with the checkpoint's
+            # static activation scale. The low-latency kernel instead consumes
+            # DeepEP's FP8 payload together with its per-token-group scales.
+            layer.dispatcher.set_quant_config(
+                {
+                    "normal_dispatcher_output_dtype": "bf16",
+                    "low_latency_dispatcher_output_dtype": "fp8",
+                }
+            )
+
     def create_moe_runner(
         self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
     ):
@@ -331,11 +342,18 @@ class W4AFp8MoEMethod(FusedMoEMethodBase):
         layer: DeepEPMoE,
         dispatch_output: DeepEPLLDispatchOutput,
     ) -> torch.Tensor:
-
-        from sglang.srt.layers.moe.cutlass_w4a8_moe import cutlass_w4a8_moe_deepep_ll
-
         hidden_states, hidden_scales, topk_ids, _, masked_m, _ = dispatch_output
 
+        if hidden_scales is None:
+            raise RuntimeError(
+                "W4AFP8 DeepEP low-latency requires FP8 dispatcher output "
+                "with per-token-group scales."
+            )
+
+        from sglang.srt.layers.moe.cutlass_w4a8_moe import (
+            cutlass_w4a8_moe_deepep_ll,
+        )
+
         output = cutlass_w4a8_moe_deepep_ll(
             hidden_states,
             hidden_scales,
@@ -367,10 +385,6 @@ class W4AFp8MoEMethod(FusedMoEMethodBase):
         layer: DeepEPMoE,
         dispatch_output: DeepEPNormalDispatchOutput,
     ) -> torch.Tensor:
-        from sglang.srt.layers.moe.cutlass_w4a8_moe import (
-            cutlass_w4a8_moe_deepep_normal,
-        )
-
         hidden_states, topk_idx, topk_weights = (
             dispatch_output.hidden_states,
             dispatch_output.topk_ids,
@@ -379,8 +393,18 @@ class W4AFp8MoEMethod(FusedMoEMethodBase):
         if isinstance(hidden_states, tuple):
             hidden_states = hidden_states[0]
 
+        if hidden_states.dtype != torch.bfloat16:
+            raise RuntimeError(
+                "W4AFP8 DeepEP normal requires BF16 dispatcher output, "
+                f"but got {hidden_states.dtype}."
+            )
+
         num_tokens = hidden_states.shape[0]
         if num_tokens > 0:
+            from sglang.srt.layers.moe.cutlass_w4a8_moe import (
+                cutlass_w4a8_moe_deepep_normal,
+            )
+
             return cutlass_w4a8_moe_deepep_normal(
                 hidden_states,
                 layer.w13_weight,