minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
EchO
2026-08-11 11:21:42 +08:00
committed by GitHub
parent 704808ed27
commit a58fa0388e
6 changed files with 327 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/cutlass_w4a8_moe.py
+3
View File
@@ -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
+4
View File
@@ -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
+20 -8
View File
@@ -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
+31 -7
View File
@@ -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,
test/registered/unit/layers/moe/test_w4afp8_deepep_dtype.py
+124
View File
@@ -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()
test/registered/unit/layers/moe/test_w4afp8_deepep_post_reorder.py
+145
View File
@@ -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()