diff --git a/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py b/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py index 613d4a0eb..3cc372057 100644 --- a/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py +++ b/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, ) diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py index f33c74f77..5ec2e777d 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py +++ b/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) diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py index 639e32831..8faaa984b 100644 --- a/python/sglang/srt/layers/moe/utils.py +++ b/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 diff --git a/python/sglang/srt/layers/quantization/w4afp8.py b/python/sglang/srt/layers/quantization/w4afp8.py index 42b1dc24c..cdf5b1e52 100644 --- a/python/sglang/srt/layers/quantization/w4afp8.py +++ b/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, diff --git a/test/registered/unit/layers/moe/test_w4afp8_deepep_dtype.py b/test/registered/unit/layers/moe/test_w4afp8_deepep_dtype.py new file mode 100644 index 000000000..d9b05c7b4 --- /dev/null +++ b/test/registered/unit/layers/moe/test_w4afp8_deepep_dtype.py @@ -0,0 +1,124 @@ +"""CPU regressions for W4AFP8 DeepEP dispatcher dtypes.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import torch + +from sglang.srt.layers.moe import utils as moe_utils +from sglang.srt.layers.moe.token_dispatcher import deepep +from sglang.srt.layers.moe.utils import ( + DeepEPMode, + DispatcherOutputDtype, + MoeRunnerBackend, +) +from sglang.srt.layers.quantization import w4afp8 +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +class TestW4AFP8DeepEPDispatcherDtype(CustomTestCase): + def test_w4afp8_sets_mode_specific_dispatcher_dtypes(self): + dispatcher = Mock() + layer = SimpleNamespace( + dispatcher=dispatcher, + w2_weight=torch.empty(0), + w13_weight_scale_inv=torch.ones((1, 1, 4)), + w2_weight_scale_inv=torch.ones((1, 1, 4)), + w13_input_scale=torch.ones(1), + w2_input_scale=torch.ones(1), + ) + + w4afp8.W4AFp8MoEMethod(SimpleNamespace()).process_weights_after_loading(layer) + + dispatcher.set_quant_config.assert_called_once_with( + { + "normal_dispatcher_output_dtype": "bf16", + "low_latency_dispatcher_output_dtype": "fp8", + } + ) + + def test_mode_specific_dtype_selection(self): + quant_config = { + "normal_dispatcher_output_dtype": "bf16", + "low_latency_dispatcher_output_dtype": "fp8", + } + + with ( + patch.object(moe_utils, "get_server_args", return_value=None), + patch.object( + moe_utils.envs.SGLANG_DEEPEP_BF16_DISPATCH, + "get", + return_value=False, + ), + patch.object( + moe_utils, + "get_moe_runner_backend", + return_value=MoeRunnerBackend.AUTO, + ), + ): + normal_dtype = moe_utils.get_deepep_output_dtype( + SimpleNamespace( + quant_config=quant_config, + dispatch_mode=DeepEPMode.NORMAL, + ) + ) + low_latency_dtype = moe_utils.get_deepep_output_dtype( + SimpleNamespace( + quant_config=quant_config, + dispatch_mode=DeepEPMode.LOW_LATENCY, + ) + ) + + self.assertEqual( + deepep._DeepEPDispatcherImplNormal.dispatch_mode, DeepEPMode.NORMAL + ) + self.assertEqual( + deepep._DeepEPDispatcherImplLowLatency.dispatch_mode, + DeepEPMode.LOW_LATENCY, + ) + self.assertEqual(normal_dtype, DispatcherOutputDtype.BF16) + self.assertEqual(low_latency_dtype, DispatcherOutputDtype.FP8) + + def test_normal_rejects_fp8_and_preserves_empty_bf16(self): + method = w4afp8.W4AFp8MoEMethod(SimpleNamespace()) + empty_topk_ids = torch.empty((0, 1), dtype=torch.int64) + empty_topk_weights = torch.empty((0, 1), dtype=torch.float32) + + fp8_dispatch_output = SimpleNamespace( + hidden_states=torch.empty((0, 128), dtype=torch.float8_e4m3fn), + topk_ids=empty_topk_ids, + topk_weights=empty_topk_weights, + ) + with self.assertRaisesRegex(RuntimeError, "requires BF16"): + method.apply_deepep_normal(SimpleNamespace(), fp8_dispatch_output) + + bf16_dispatch_output = SimpleNamespace( + hidden_states=torch.empty((0, 128), dtype=torch.bfloat16), + topk_ids=empty_topk_ids, + topk_weights=empty_topk_weights, + ) + output = method.apply_deepep_normal(SimpleNamespace(), bf16_dispatch_output) + self.assertEqual(output.dtype, torch.bfloat16) + self.assertEqual(output.shape, (0, 128)) + + def test_low_latency_requires_fp8_scales(self): + method = w4afp8.W4AFp8MoEMethod(SimpleNamespace()) + dispatch_output = ( + torch.empty((1, 1, 128), dtype=torch.bfloat16), + None, + torch.empty((0, 1), dtype=torch.int64), + torch.empty((0, 1), dtype=torch.float32), + torch.zeros(1, dtype=torch.int32), + 0, + ) + + with self.assertRaisesRegex(RuntimeError, "requires FP8"): + method.apply_deepep_ll(SimpleNamespace(), dispatch_output) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/layers/moe/test_w4afp8_deepep_post_reorder.py b/test/registered/unit/layers/moe/test_w4afp8_deepep_post_reorder.py new file mode 100644 index 000000000..a504d4c04 --- /dev/null +++ b/test/registered/unit/layers/moe/test_w4afp8_deepep_post_reorder.py @@ -0,0 +1,145 @@ +"""Regression test for W4AFP8 DeepEP-normal post-reorder scaling.""" + +import sys +import unittest +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock, patch + +import torch + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +# The function under test is a GPU implementation, but this test replaces every +# launched kernel and only verifies the host-side call contract. Stub the +# extension symbols so importing the module remains valid on CPU CI runners. +_sgl_kernel_stub = ModuleType("sgl_kernel") +_sgl_kernel_stub.cutlass_w4a8_moe_mm = Mock() +_sgl_kernel_stub.get_cutlass_w4a8_moe_mm_data = Mock() +_sgl_kernel_stub.silu_and_mul = Mock() +with patch.dict(sys.modules, {"sgl_kernel": _sgl_kernel_stub}): + from sglang.srt.layers.moe import cutlass_w4a8_moe as w4a8_moe + + +class _KernelLauncher: + def __init__(self, fn): + self.fn = fn + + def __getitem__(self, _grid): + return self.fn + + +class TestW4AFP8DeepEPNormalPostReorder(CustomTestCase): + def test_post_reorder_receives_neutral_routed_scale(self): + """The local reduction is unscaled; DeepEP scales after rank combine.""" + + num_tokens, hidden_size, intermediate_size = 2, 8, 4 + num_experts, topk = 2, 2 + topk_ids = torch.tensor([[0, 1], [1, 0]], dtype=torch.int64) + topk_weights = torch.full((num_tokens, topk), 0.5, dtype=torch.float32) + src2dst = torch.arange(num_tokens * topk, dtype=torch.int64) + + def fake_post_reorder( + _down_output, + output, + _src2dst, + _topk_ids, + _topk_weights, + _topk, + _hidden_size, + routed_scaling_factor, + *, + BLOCK_SIZE, + ): + self.assertEqual(routed_scaling_factor, 1.0) + self.assertEqual(BLOCK_SIZE, 512) + output.zero_() + + noop_launcher = _KernelLauncher(lambda *args, **kwargs: None) + post_reorder_launcher = _KernelLauncher(fake_post_reorder) + preprocess_result = ( + torch.arange(num_tokens * topk), + src2dst, + torch.empty(0), + ) + + strides = torch.zeros((num_experts, 3), dtype=torch.int64) + expert_offsets = torch.zeros(num_experts + 1, dtype=torch.int32) + problem_sizes = torch.zeros((num_experts, 3), dtype=torch.int32) + layer = SimpleNamespace( + w13_weight=torch.zeros( + (num_experts, intermediate_size * 2, hidden_size // 2), + dtype=torch.int8, + ), + w2_weight=torch.zeros( + (num_experts, hidden_size, intermediate_size // 2), + dtype=torch.int8, + ), + w13_weight_scale_inv=torch.ones((num_experts, 1, 1)), + w2_weight_scale_inv=torch.ones((num_experts, 1, 1)), + w13_input_scale=torch.ones(1), + w2_input_scale=torch.ones(1), + ) + + with ( + patch.object( + w4a8_moe, + "deepep_run_moe_deep_preprocess", + return_value=preprocess_result, + ), + patch.object(w4a8_moe, "deepep_permute_triton_kernel", noop_launcher), + patch.object( + w4a8_moe, + "deepep_post_reorder_triton_kernel", + post_reorder_launcher, + ), + patch.object( + w4a8_moe, + "get_cutlass_w4a8_moe_mm_data", + new=lambda *args, **kwargs: None, + create=True, + ), + patch.object( + w4a8_moe, + "cutlass_w4a8_moe_mm", + new=lambda *args, **kwargs: None, + create=True, + ), + patch.object( + w4a8_moe, + "per_tensor_quant_fp8", + new=lambda *args, **kwargs: None, + ), + patch.object(w4a8_moe, "silu_and_mul", new=lambda *args, **kwargs: None), + ): + output = w4a8_moe.cutlass_w4a8_moe_deepep_normal( + torch.ones((num_tokens, hidden_size), dtype=torch.bfloat16), + layer.w13_weight, + layer.w2_weight, + layer.w13_weight_scale_inv, + layer.w2_weight_scale_inv, + topk_weights, + topk_ids, + strides, + strides, + strides, + strides, + strides, + strides, + strides, + strides, + expert_offsets, + problem_sizes, + problem_sizes, + layer.w13_input_scale, + layer.w2_input_scale, + ) + + self.assertEqual(output.shape, (num_tokens, hidden_size)) + self.assertEqual(output.dtype, torch.bfloat16) + + +if __name__ == "__main__": + unittest.main()