[Fix] Correct W4AFP8 DeepEP scaling and mode-specific dtypes (#33669)
This commit is contained in:
EchO
@@ -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()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user