Support MXFP8 and deferred route weighting in DeepEP v2 (#40030)
Co-authored-by: metamergebot <324680979+metamergebot@users.noreply.github.com> Co-authored-by: Xingyu Liu <38244988+charlotte12l@users.noreply.github.com> Co-authored-by: pranjalssh <14260275+pranjalssh@users.noreply.github.com>
This commit is contained in:
metamergebot
co-authored by
metamergebot
Xingyu Liu
pranjalssh
metamergebot
Xingyu Liu
pranjalssh
parent
6cc9090d1f
commit
0e5347db82
@@ -144,6 +144,7 @@ class TestDeepEPv2BufferLifecycle(CustomTestCase):
|
||||
impl = object.__new__(deepep_v2._DeepEPv2Impl)
|
||||
impl.num_max_dispatch_tokens_per_rank = 4
|
||||
impl.hidden_size = 128
|
||||
impl.activation_scale_block_size = 128
|
||||
impl.router_topk = 2
|
||||
impl._validate_common(torch.empty(4, 128), torch.zeros(4, 2))
|
||||
with self.assertRaisesRegex(ValueError, "per-rank buffer capacity"):
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
"""Tests for the DeepEP v2 expanded/masked repack kernels."""
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
|
||||
@@ -113,6 +114,80 @@ class TestDeepEPv2MaskedSlab(CustomTestCase):
|
||||
def test_empty_experts(self):
|
||||
self._check_expand_roundtrip([0, 0, 0, 0], torch.bfloat16, with_scale=False)
|
||||
|
||||
def test_runner_defers_expanded_route_weighting(self):
|
||||
from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
|
||||
from sglang.srt.layers.moe.moe_runner.deep_gemm import (
|
||||
DeepGemmRunnerOutput,
|
||||
post_permute_deep_gemm_to_deepep_v2,
|
||||
)
|
||||
from sglang.srt.layers.moe.token_dispatcher.base import RoutewiseLayout
|
||||
|
||||
counts = [3, 0, 2]
|
||||
recv_x, _, psum, starts, total = _build_layout(
|
||||
counts, self.ALIGN, self.HIDDEN, torch.bfloat16
|
||||
)
|
||||
masked_x, _, _ = expand_to_masked_slab(
|
||||
recv_x, None, psum, len(counts), self.MAX_M, self.ALIGN
|
||||
)
|
||||
weights = torch.full((total,), 0.25, device=DEVICE)
|
||||
state = {
|
||||
"deepep_v2_expanded": True,
|
||||
"deepep_v2_masked": True,
|
||||
"deepep_v2_psum": psum,
|
||||
"deepep_v2_total_expanded": total,
|
||||
"deepep_v2_expert_alignment": self.ALIGN,
|
||||
"topk_weights": weights,
|
||||
}
|
||||
rows = _real_rows(starts, counts)
|
||||
for no_combine in (False, True):
|
||||
with self.subTest(no_combine=no_combine):
|
||||
output = post_permute_deep_gemm_to_deepep_v2(
|
||||
DeepGemmRunnerOutput(masked_x),
|
||||
None,
|
||||
MoeRunnerConfig(no_combine=no_combine),
|
||||
state,
|
||||
)
|
||||
expected = recv_x[rows] if no_combine else recv_x[rows] * 0.25
|
||||
self.assertTrue(torch.equal(output.hidden_states[rows], expected))
|
||||
self.assertEqual(
|
||||
output.routewise_layout,
|
||||
RoutewiseLayout.EXPANDED if no_combine else None,
|
||||
)
|
||||
if no_combine:
|
||||
self.assertTrue(torch.equal(output.topk_weights, weights))
|
||||
|
||||
def test_runner_restores_token_topk_routes_and_masks_nonlocal_slots(self):
|
||||
from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
|
||||
from sglang.srt.layers.moe.moe_runner.deep_gemm import (
|
||||
DeepGemmRunnerOutput,
|
||||
post_permute_deep_gemm_to_deepep_v2,
|
||||
)
|
||||
from sglang.srt.layers.moe.token_dispatcher.base import RoutewiseLayout
|
||||
|
||||
hidden = torch.tensor([[3.0], [5.0], [7.0]], device=DEVICE)
|
||||
weights = torch.tensor([[0.25, 0.0], [0.5, 0.75]], device=DEVICE)
|
||||
state = {
|
||||
"topk_ids": torch.tensor([[0, -1], [1, 0]], device=DEVICE),
|
||||
"topk_weights": weights,
|
||||
"output_index": torch.tensor([[1, -1], [0, 2]], device=DEVICE),
|
||||
}
|
||||
for empty in (False, True):
|
||||
with self.subTest(empty=empty):
|
||||
if empty:
|
||||
state["output_index"] = torch.full((2, 2), -1, device=DEVICE)
|
||||
output = post_permute_deep_gemm_to_deepep_v2(
|
||||
DeepGemmRunnerOutput(hidden[:0] if empty else hidden),
|
||||
None,
|
||||
MoeRunnerConfig(no_combine=True),
|
||||
state,
|
||||
)
|
||||
expected = torch.tensor([[[5.0], [0.0]], [[3.0], [7.0]]], device=DEVICE)
|
||||
if empty:
|
||||
expected.zero_()
|
||||
self.assertTrue(torch.equal(output.hidden_states, expected))
|
||||
self.assertIs(output.topk_weights, weights)
|
||||
self.assertEqual(output.routewise_layout, RoutewiseLayout.TOKEN_TOPK)
|
||||
|
||||
def test_single_hot_expert(self):
|
||||
self._check_expand_roundtrip(
|
||||
[0, self.MAX_M, 0, 0], torch.bfloat16, with_scale=False, topk=True
|
||||
@@ -133,7 +208,7 @@ class TestDeepEPv2MaskedSlab(CustomTestCase):
|
||||
recv_x, None, psum, len(counts), self.MAX_M, self.ALIGN
|
||||
)
|
||||
|
||||
def _production_packed_ue8m0_layout(self, counts):
|
||||
def _production_packed_ue8m0_layout(self, counts, group_size):
|
||||
"""Build expanded rows with the production packed UE8M0 quantizer."""
|
||||
from sglang.kernels.ops.quantization.fp8_kernel import (
|
||||
sglang_per_token_group_quant_fp8,
|
||||
@@ -146,7 +221,7 @@ class TestDeepEPv2MaskedSlab(CustomTestCase):
|
||||
)
|
||||
recv_x, recv_x_scale = sglang_per_token_group_quant_fp8(
|
||||
raw,
|
||||
128,
|
||||
group_size,
|
||||
column_major_scales=True,
|
||||
scale_tma_aligned=True,
|
||||
scale_ue8m0=True,
|
||||
@@ -157,9 +232,14 @@ class TestDeepEPv2MaskedSlab(CustomTestCase):
|
||||
return recv_x, recv_x_scale, psum, starts, total, hidden
|
||||
|
||||
def test_fp8_packed_ue8m0_scale_from_production_quantizer(self):
|
||||
for group_size in (32, 128):
|
||||
with self.subTest(group_size=group_size):
|
||||
self._check_packed_scale(group_size)
|
||||
|
||||
def _check_packed_scale(self, group_size):
|
||||
counts = [3, 1, 6, 2]
|
||||
recv_x, recv_x_scale, psum, starts, _, hidden = (
|
||||
self._production_packed_ue8m0_layout(counts)
|
||||
self._production_packed_ue8m0_layout(counts, group_size)
|
||||
)
|
||||
E = len(counts)
|
||||
masked_x, masked_x_scale, masked_m = expand_to_masked_slab(
|
||||
@@ -175,10 +255,15 @@ class TestDeepEPv2MaskedSlab(CustomTestCase):
|
||||
torch.testing.assert_close(masked_x_scale[e, j], recv_x_scale[s + j])
|
||||
|
||||
def test_expand_under_cuda_graph_capture(self):
|
||||
for group_size in (32, 128):
|
||||
with self.subTest(group_size=group_size):
|
||||
self._check_graph_capture(group_size)
|
||||
|
||||
def _check_graph_capture(self, group_size):
|
||||
# Exercise replay with the production packed scale layout.
|
||||
counts = [3, 1, 6, 2]
|
||||
recv_x, recv_x_scale, psum, starts, _, _ = self._production_packed_ue8m0_layout(
|
||||
counts
|
||||
counts, group_size
|
||||
)
|
||||
E = len(counts)
|
||||
warm = torch.cuda.Stream()
|
||||
@@ -235,10 +320,34 @@ class TestDeepEPv2HandleLifecycle(CustomTestCase):
|
||||
|
||||
impl._get_buffer = _boom
|
||||
with self.assertRaisesRegex(RuntimeError, "boom"):
|
||||
impl.combine(None)
|
||||
impl.combine(SimpleNamespace(routewise_layout=None))
|
||||
self.assertIsNone(impl._handle)
|
||||
self.assertFalse(impl._pad_empty_combine)
|
||||
|
||||
def test_unfinalized_routes_are_rejected_and_release_handle(self):
|
||||
from sglang.srt.layers.moe.token_dispatcher.base import (
|
||||
CombineInputChecker,
|
||||
RoutewiseLayout,
|
||||
)
|
||||
from sglang.srt.layers.moe.token_dispatcher.deepep_v2 import (
|
||||
DeepEPv2CombineInput,
|
||||
)
|
||||
|
||||
impl = self._bare_impl()
|
||||
impl._handle = object()
|
||||
output = DeepEPv2CombineInput(
|
||||
torch.empty(2, 8), torch.ones(2), RoutewiseLayout.EXPANDED
|
||||
)
|
||||
self.assertTrue(CombineInputChecker.needs_model_route_finalization(output))
|
||||
with self.assertRaisesRegex(ValueError, "model route finalization"):
|
||||
impl.combine(output)
|
||||
self.assertIsNone(impl._handle)
|
||||
self.assertFalse(
|
||||
CombineInputChecker.needs_model_route_finalization(
|
||||
DeepEPv2CombineInput(torch.empty(2, 8), None)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -52,6 +52,13 @@ class _FakeBuffer:
|
||||
num_recv = (x[0] if isinstance(x, tuple) else x).shape[0]
|
||||
topk_idx = kwargs["topk_idx"]
|
||||
topk_weights = kwargs["topk_weights"]
|
||||
if kwargs["do_expand"]:
|
||||
if isinstance(x, tuple):
|
||||
x = tuple(t.repeat_interleave(TOPK, dim=0) for t in x)
|
||||
else:
|
||||
x = x.repeat_interleave(TOPK, dim=0)
|
||||
topk_idx = None
|
||||
topk_weights = topk_weights.flatten()
|
||||
event = SimpleNamespace(event=None, current_stream_wait=lambda: None)
|
||||
return x, topk_idx, topk_weights, _FakeHandle(num_recv), event
|
||||
|
||||
@@ -93,7 +100,9 @@ class _DeepEPv2WireDtypeBase(CustomTestCase):
|
||||
for item in reversed(self._patches):
|
||||
item.stop()
|
||||
|
||||
def _dispatch(self, use_fp8_dispatch, num_tokens=8, is_extend_in_batch=True):
|
||||
def _dispatch(
|
||||
self, use_fp8_dispatch, num_tokens=8, is_extend_in_batch=True, group_size=128
|
||||
):
|
||||
dispatcher = deepep_v2.DeepEPv2Dispatcher(
|
||||
group=_FakeGroup(),
|
||||
router_topk=TOPK,
|
||||
@@ -102,6 +111,7 @@ class _DeepEPv2WireDtypeBase(CustomTestCase):
|
||||
hidden_size=HIDDEN,
|
||||
params_dtype=torch.bfloat16,
|
||||
use_fp8_dispatch=use_fp8_dispatch,
|
||||
activation_scale_block_size=group_size,
|
||||
)
|
||||
dispatcher._impl.num_max_dispatch_tokens_per_rank = NUM_MAX_TOKENS
|
||||
hidden_states = torch.randn((num_tokens, HIDDEN), dtype=torch.bfloat16)
|
||||
@@ -116,6 +126,18 @@ class _DeepEPv2WireDtypeBase(CustomTestCase):
|
||||
|
||||
|
||||
class TestDeepEPv2WireDtype(_DeepEPv2WireDtypeBase):
|
||||
def test_mxfp8_scale_group_survives_both_dispatch_layouts(self):
|
||||
for is_extend in (True, False):
|
||||
with self.subTest(is_extend=is_extend):
|
||||
_, out = self._dispatch(
|
||||
use_fp8_dispatch=True,
|
||||
is_extend_in_batch=is_extend,
|
||||
group_size=32,
|
||||
)
|
||||
self.assertEqual(out.activation_scale_block_size, 32)
|
||||
self.assertEqual(out.hidden_states_scale.shape[-1], HIDDEN // 32)
|
||||
self.assertEqual(out.is_expanded, not is_extend)
|
||||
|
||||
def test_bf16_dispatch_sends_unquantized_activations(self):
|
||||
hidden_states, out = self._dispatch(use_fp8_dispatch=False)
|
||||
self.assertIs(_FakeBuffer.last.dispatch_x, hidden_states)
|
||||
|
||||
@@ -76,13 +76,18 @@ def _fp8_method(**overrides):
|
||||
return method
|
||||
|
||||
|
||||
def test_deepep_v2_quant_contract_accepts_blockwise_fp8(_moe_flags):
|
||||
@pytest.mark.parametrize("use_mxfp8,block_size", [(False, [128, 128]), (True, [1, 32])])
|
||||
def test_deepep_v2_quant_contract_accepts_blockwise_fp8(
|
||||
_moe_flags, use_mxfp8, block_size
|
||||
):
|
||||
from sglang.srt.layers.moe.fused_moe_triton.layer import (
|
||||
_validate_deepep_v2_quant_method,
|
||||
)
|
||||
|
||||
_moe_flags.a2a_backend = MoeA2ABackend.DEEPEP_V2
|
||||
_validate_deepep_v2_quant_method(_fp8_method(weight_block_size=[128, 128]))
|
||||
_validate_deepep_v2_quant_method(
|
||||
_fp8_method(weight_block_size=block_size, use_mxfp8=use_mxfp8)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -90,7 +95,7 @@ def test_deepep_v2_quant_contract_accepts_blockwise_fp8(_moe_flags):
|
||||
[
|
||||
({"activation_scheme": "static"}, "activation_scheme"),
|
||||
({"weight_block_size": None}, "weight_block_size"),
|
||||
({"weight_block_size": (1, 32), "use_mxfp8": True}, "MXFP8"),
|
||||
({"weight_block_size": (128, 128), "use_mxfp8": True}, "MXFP8"),
|
||||
({"is_fp4_expert": True}, "FP4 experts"),
|
||||
],
|
||||
)
|
||||
@@ -149,5 +154,46 @@ def test_deepep_v2_runner_backstop(_moe_flags):
|
||||
assert MoeRunner(MoeRunnerBackend.DEEP_GEMM, MoeRunnerConfig()).runner_core
|
||||
|
||||
|
||||
def test_deepep_v2_registration_uses_primary_architecture_and_rejects_conflicts():
|
||||
from sglang.srt.configs.moe_model_registry import (
|
||||
model_requires_fp32_silu_mul,
|
||||
model_supports_deepep_v2,
|
||||
register_deepep_v2_model,
|
||||
)
|
||||
|
||||
register_deepep_v2_model("TestRoutewiseMoe", silu_mul_keep_fp32=True)
|
||||
config = SimpleNamespace(architectures=["TestRoutewiseMoe"])
|
||||
assert model_supports_deepep_v2(config)
|
||||
assert model_requires_fp32_silu_mul(config)
|
||||
config.architectures = ["UnsupportedMoe", "TestRoutewiseMoe"]
|
||||
assert not model_supports_deepep_v2(config)
|
||||
assert not model_requires_fp32_silu_mul(config)
|
||||
with pytest.raises(ValueError, match="Conflicting"):
|
||||
register_deepep_v2_model("TestRoutewiseMoe", silu_mul_keep_fp32=False)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("block_size,width", [(32, 16), (128, 4)])
|
||||
def test_mxfp8_recipes_keep_activation_and_weight_groups_separate(block_size, width):
|
||||
# The packed format is a layout contract, independent of the host GPU.
|
||||
from unittest.mock import patch
|
||||
|
||||
from sglang.srt.layers import deep_gemm_wrapper
|
||||
from sglang.srt.layers.moe.moe_runner.deep_gemm import DeepGemmMoeQuantInfo
|
||||
|
||||
with patch.object(deep_gemm_wrapper, "DEEPGEMM_SCALE_UE8M0", True):
|
||||
quant = DeepGemmMoeQuantInfo(
|
||||
None, None, True, block_shape=[1, 32], use_mxfp8=True
|
||||
)
|
||||
assert quant.scale_recipes(
|
||||
activation_block_size=block_size, hidden_size=2048, activation_scale_width=width
|
||||
) == ((1, block_size), (1, 32))
|
||||
with pytest.raises(AssertionError, match="activation scale mismatch"):
|
||||
quant.scale_recipes(
|
||||
activation_block_size=block_size,
|
||||
hidden_size=2048,
|
||||
activation_scale_width=width + 1,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
|
||||
Reference in New Issue
Block a user