iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Perf][DSV4.1] Pad shared experts for native MegaMoE fusion (#56568)

Browse Source 
Signed-off-by: Canlin <canlinguosdu@gmail.com>
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Canlin Guo
2026-09-16 17:08:42 +08:00
committed by GitHub
co-authored by Codex
parent 0384e72693
commit 03f67b3ad1
2 changed files with 128 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/models/test_deepseek_v4_mega_moe.py
+80 -18
View File
@@ -367,17 +367,21 @@ def test_deepseek_v4_mega_moe_padding_preserves_weights(monkeypatch, intermediat
[
pytest.param(128, 512, 128, True, False, id="aligned-block128"),
pytest.param(128, 512, 128, True, True, id="aligned-block128-mxfp8"),
pytest.param(5120, 2304, 32, False, False, id="deepseek-v41-flash"),
pytest.param(5120, 2304, 32, True, False, id="deepseek-v41-flash"),
pytest.param(128, 2304, 128, True, False, id="padded-block128"),
pytest.param(5120, 2304, 32, True, True, id="deepseek-v41-flash-mxfp8"),
pytest.param(128, 2304, 128, False, False, id="padded-packed-scale-fallback"),
],
)
def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights(
monkeypatch, hidden_size, intermediate_size, block_size, fused, mxfp8
):
"""V4.1-Flash's routed padding must preserve the serial shared-expert fallback."""
"""Shared fusion preserves checkpoint channels and zeros padded channels."""
class FakeDeepGemm:
transformed_dims: list[tuple[int, int]] = []
scale_inputs: list[tuple[int, ...]] = []
scale_values: list[torch.Tensor] = []
@staticmethod
def get_symm_buffer_for_mega_moe(*args, num_shared_experts=0, **kwargs):
@@ -402,6 +406,7 @@ def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights(
@classmethod
def transform_sf_into_required_layout(cls, sf, mn, k, *args, **kwargs):
cls.scale_inputs.append(tuple(sf.shape))
cls.scale_values.append(sf)
return torch.empty((sf.shape[0], mn, k // 128), dtype=torch.int32)
@classmethod
@@ -429,7 +434,10 @@ def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights(
def fp8_parameter(*shape):
return torch.nn.Parameter(
torch.empty(*shape, dtype=torch.float8_e4m3fn), requires_grad=False
torch.empty(*shape, dtype=torch.uint8)
.random_(1, 119)
.view(torch.float8_e4m3fn),
requires_grad=False,
)
def scale_parameter(*shape, dtype=torch.int32):
@@ -460,21 +468,34 @@ def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights(
linear.weight_block_size = (1, 32)
linear.weight_scale = torch.nn.Parameter(
linear.weight_scale_inv.view(torch.uint8)
.repeat_interleave(128, dim=0)
.repeat_interleave(4, dim=1),
.repeat_interleave(block_size, dim=0)
.repeat_interleave(block_size // 32, dim=1),
requires_grad=False,
)
del linear.weight_scale_inv
originals = []
for linear in (shared_experts.gate_up_proj, shared_experts.down_proj):
scale = linear.weight_scale if mxfp8 else linear.weight_scale_inv
if not fused:
# Already packed scales cannot be padded as checkpoint UE8M0 bytes.
linear.weight_scale_inv = scale_parameter(
linear.weight.shape[0], linear.weight.shape[1] // 128
)
continue
scale.data.view(torch.uint8).random_(124, 130)
block_m, block_k = linear.weight_block_size
scale_1x32 = (
experts._ue8m0_uint8_to_float(scale.view(torch.uint8))
.repeat_interleave(block_m, dim=0)
.repeat_interleave(block_k // 32, dim=1)
)
originals.append((linear.weight.view(torch.uint8).clone(), scale_1x32))
monkeypatch.setattr("vllm.utils.deep_gemm._import_deep_gemm", lambda: FakeDeepGemm)
original_gate_up_ptr = shared_experts.gate_up_proj.weight.data_ptr()
original_down_ptr = shared_experts.down_proj.weight.data_ptr()
experts.finalize_weights(shared_experts)
assert FakeDeepGemm.scale_inputs[-2:] == [
(1, 2 * intermediate_size, hidden_size // 32),
(1, hidden_size, intermediate_size // 32),
]
assert experts.has_fused_shared_experts is fused
if not fused:
assert experts.intermediate_size == 2560
@@ -490,15 +511,56 @@ def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights(
return
assert FakeDeepGemm.transformed_dims == [(3, 3), (2, 2)]
assert shared_experts.gate_up_proj.weight.data_ptr() != original_gate_up_ptr
assert (
experts._transformed_shared_l1_weights[0].data_ptr()
== shared_experts.gate_up_proj.weight.data_ptr()
)
assert (
experts._transformed_shared_l2_weights[0].data_ptr()
== shared_experts.down_proj.weight.data_ptr()
)
padded_size = experts.intermediate_size
assert FakeDeepGemm.scale_inputs[-2:] == [
(1, 2 * padded_size, hidden_size // 32),
(1, hidden_size, padded_size // 32),
]
for actual, original, fill in (
(
experts._transformed_shared_l1_weights[0].view(torch.uint8),
originals[0][0],
0,
),
(FakeDeepGemm.scale_values[-2][0], originals[0][1], 1),
):
actual = actual.unflatten(0, (2, padded_size))
assert torch.equal(
actual[:, :intermediate_size], original.unflatten(0, (2, intermediate_size))
)
assert torch.all(actual[:, intermediate_size:] == fill)
for actual, original, fill in (
(
experts._transformed_shared_l2_weights[0].view(torch.uint8),
originals[1][0],
0,
),
(FakeDeepGemm.scale_values[-1][0], originals[1][1], 1),
):
assert torch.equal(actual[:, : original.shape[1]], original)
assert torch.all(actual[:, original.shape[1] :] == fill)
if padded_size != intermediate_size:
# Generic linear post-load hooks still receive checkpoint-shaped weights.
for linear, (weight, _) in zip(
(shared_experts.gate_up_proj, shared_experts.down_proj), originals
):
assert torch.equal(linear.weight.view(torch.uint8), weight)
assert shared_experts.gate_up_proj.weight.data_ptr() == original_gate_up_ptr
assert shared_experts.down_proj.weight.data_ptr() == original_down_ptr
else:
assert shared_experts.gate_up_proj.weight.data_ptr() != original_gate_up_ptr
assert (
experts._transformed_shared_l1_weights[0].data_ptr()
== shared_experts.gate_up_proj.weight.data_ptr()
)
assert (
experts._transformed_shared_l2_weights[0].data_ptr()
== shared_experts.down_proj.weight.data_ptr()
)
transformed_l1 = experts._transformed_shared_l1_weights
experts.finalize_weights(shared_experts)
assert experts._transformed_shared_l1_weights is transformed_l1
@pytest.mark.parametrize("fused", [False, True])
vllm/models/deepseek_v4/nvidia/model.py
+48 -5
View File
@@ -204,6 +204,7 @@ class DeepseekV4MegaMoEExperts(nn.Module):
self.top_k = top_k
self.hidden_size = hidden_size
self.intermediate_size = intermediate_size
self.unpadded_intermediate_size = intermediate_size
self.num_shared_experts = num_shared_experts
self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
@@ -386,6 +387,17 @@ class DeepseekV4MegaMoEExperts(nn.Module):
# the generic linear post-load hook replaces the raw checkpoint scales
# with its 128x128 DeepGEMM layout.
checkpoint_scale_dtypes = (torch.float8_e8m0fnu, torch.uint8)
unpadded_size = self.unpadded_intermediate_size * self.num_shared_experts
padding = self.intermediate_size * self.num_shared_experts - unpadded_size
pad_weights = (
padding > 0
and gate_up_weight.dtype == torch.float8_e4m3fn
and down_weight.dtype == torch.float8_e4m3fn
and gate_up_weight.shape == (2 * unpadded_size, self.hidden_size)
and down_weight.shape == (self.hidden_size, unpadded_size)
and gate_up_scale.dtype in checkpoint_scale_dtypes
and down_scale.dtype in checkpoint_scale_dtypes
)
if (
gate_up_scale.dtype in checkpoint_scale_dtypes
and down_scale.dtype in checkpoint_scale_dtypes
@@ -396,6 +408,7 @@ class DeepseekV4MegaMoEExperts(nn.Module):
gate_up_scale,
gate_up_weight.shape[0],
gate_up_weight.shape[1],
padding=(padding, 0) if pad_weights else (0, 0),
)
down_scale = self._prepare_shared_expert_scale(
deep_gemm,
@@ -403,12 +416,27 @@ class DeepseekV4MegaMoEExperts(nn.Module):
down_scale,
down_weight.shape[0],
down_weight.shape[1],
padding=(0, padding) if pad_weights else (0, 0),
)
if gate_up_scale is None or down_scale is None:
self.num_shared_experts = 0
return
if pad_weights:
# Pad gate/up separately so the SwiGLU split stays at the midpoint.
gate_up_weight = (
torch.nn.functional.pad(
gate_up_weight.view(torch.uint8).unflatten(0, (2, unpadded_size)),
(0, 0, 0, padding),
)
.flatten(0, 1)
.view(gate_up_weight.dtype)
)
down_weight = torch.nn.functional.pad(
down_weight.view(torch.uint8), (0, padding)
).view(down_weight.dtype)
shared_intermediate_size = self.intermediate_size * self.num_shared_experts
expected_gate_up_shape = (
2 * shared_intermediate_size,
@@ -449,11 +477,11 @@ class DeepseekV4MegaMoEExperts(nn.Module):
# released instead of adding roughly 0.7 GiB per rank on DSV4-Flash.
# The generic linear post-load hook may still repack the serial scales,
# but this shared MLP is never called after native fusion is enabled.
gate_up.weight.data = transformed_l1[0]
self._transformed_shared_l1_weights = (
gate_up.weight.data,
transformed_l1[1],
)
# Padded weights need separate storage: the generic linear post-load
# hooks still expect the original checkpoint weight/scale shapes.
if not pad_weights:
gate_up.weight.data = transformed_l1[0]
self._transformed_shared_l1_weights = transformed_l1
self._transformed_shared_l2_weights = transformed_l2
def _prepare_shared_expert_scale(
@@ -463,6 +491,8 @@ class DeepseekV4MegaMoEExperts(nn.Module):
scale: torch.Tensor,
mn: int,
k: int,
*,
padding: tuple[int, int] = (0, 0),
) -> torch.Tensor | None:
block_size = getattr(linear, "weight_block_size", None)
if block_size is None or len(block_size) != 2:
@@ -497,6 +527,19 @@ class DeepseekV4MegaMoEExperts(nn.Module):
.repeat_interleave(block_k // 32, dim=1)[:mn, : k // 32]
.contiguous()
)
pad_m, pad_k = padding
if pad_m:
scale_1x32 = torch.nn.functional.pad(
scale_1x32.unflatten(0, (2, mn // 2)),
(0, 0, 0, pad_m),
value=1.0,
).flatten(0, 1)
mn += 2 * pad_m
if pad_k:
scale_1x32 = torch.nn.functional.pad(
scale_1x32, (0, pad_k // 32), value=1.0
)
k += pad_k
# The grouped API is used with a singleton dimension to request the
# MN-major, TMA-aligned packed-UE8M0 strides, then squeezed back to the
# 2D layout required for a shared expert.