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

[Perf][DSV4.1] Remove MegaMoE padding and shared padding workaround (#57204)

Browse Source 
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
This commit is contained in:
Woosuk Kwon
2026-09-16 19:09:30 +00:00
committed by GitHub
parent bc0f47cd03
commit 30b847c1fa
2 changed files with 42 additions and 157 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/models/test_deepseek_v4_mega_moe.py
+37 -92
View File
@@ -300,8 +300,10 @@ def test_deepseek_v4_mega_moe_weight_loader_uses_ep_expert_ownership():
@pytest.mark.parametrize("intermediate_size", [512, 2304])
def test_deepseek_v4_mega_moe_padding_preserves_weights(monkeypatch, intermediate_size):
"""Pad gate/up separately and zero every added down-projection column."""
def test_deepseek_v4_mega_moe_preserves_checkpoint_dimensions(
monkeypatch, intermediate_size
):
"""Keep native widths so V4.1 routed and shared experts can fuse."""
experts = DeepseekV4MegaMoEExperts(
SimpleNamespace(
scheduler_config=SimpleNamespace(max_num_batched_tokens=4),
@@ -333,29 +335,15 @@ def test_deepseek_v4_mega_moe_padding_preserves_weights(monkeypatch, intermediat
)
experts.finalize_weights()
padded_size = experts.intermediate_size
assert padded_size == (intermediate_size + 511) // 512 * 512
for actual, original in zip(experts._transformed_l1_weights, originals[:2]):
assert experts.intermediate_size == intermediate_size
transformed = (*experts._transformed_l1_weights, *experts._transformed_l2_weights)
for actual, original in zip(transformed, originals):
expected = (
original
if actual.dtype == torch.int8
else (experts._ue8m0_uint8_to_float(original))
else experts._ue8m0_uint8_to_float(original)
)
actual = actual.view(expected.dtype).unflatten(1, (2, padded_size))
assert torch.equal(
actual[:, :, :intermediate_size],
expected.unflatten(1, (2, intermediate_size)),
)
assert torch.count_nonzero(actual[:, :, intermediate_size:]) == 0
for actual, original in zip(experts._transformed_l2_weights, originals[2:]):
expected = (
original
if actual.dtype == torch.int8
else (experts._ue8m0_uint8_to_float(original))
)
actual = actual.view(expected.dtype)
assert torch.equal(actual[..., : original.shape[-1]], expected)
assert torch.count_nonzero(actual[..., original.shape[-1] :]) == 0
assert torch.equal(actual.view(expected.dtype), expected)
transformed_l1 = experts._transformed_l1_weights
experts.finalize_weights()
@@ -363,20 +351,19 @@ def test_deepseek_v4_mega_moe_padding_preserves_weights(monkeypatch, intermediat
@pytest.mark.parametrize(
"hidden_size,intermediate_size,block_size,fused,mxfp8",
"hidden_size,intermediate_size,block_size,mxfp8",
[
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, 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"),
pytest.param(128, 512, 128, False, id="aligned-block128"),
pytest.param(128, 512, 128, True, id="aligned-block128-mxfp8"),
pytest.param(5120, 2304, 32, False, id="deepseek-v41-flash"),
pytest.param(128, 2304, 128, False, id="native-block128"),
pytest.param(5120, 2304, 32, True, id="deepseek-v41-flash-mxfp8"),
],
)
def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights(
monkeypatch, hidden_size, intermediate_size, block_size, fused, mxfp8
monkeypatch, hidden_size, intermediate_size, block_size, mxfp8
):
"""Shared fusion preserves checkpoint channels and zeros padded channels."""
"""Shared fusion preserves checkpoint weights and scales at native widths."""
class FakeDeepGemm:
transformed_dims: list[tuple[int, int]] = []
@@ -476,12 +463,6 @@ def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights(
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 = (
@@ -493,70 +474,34 @@ def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights(
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 experts.has_fused_shared_experts is fused
if not fused:
assert experts.intermediate_size == 2560
assert FakeDeepGemm.transformed_dims == [(3, 3)]
assert experts.num_shared_experts == 0
assert shared_experts.gate_up_proj.weight.data_ptr() == original_gate_up_ptr
assert shared_experts.down_proj.weight.data_ptr() == original_down_ptr
assert shared_experts.gate_up_proj.weight.shape == (
2 * intermediate_size,
hidden_size,
)
assert shared_experts.down_proj.weight.shape == (hidden_size, intermediate_size)
return
assert experts.has_fused_shared_experts
assert experts.intermediate_size == intermediate_size
assert FakeDeepGemm.transformed_dims == [(3, 3), (2, 2)]
padded_size = experts.intermediate_size
assert FakeDeepGemm.scale_inputs[-2:] == [
(1, 2 * padded_size, hidden_size // 32),
(1, hidden_size, padded_size // 32),
(1, 2 * intermediate_size, hidden_size // 32),
(1, hidden_size, intermediate_size // 32),
]
for actual, original, fill in (
for transformed, actual_scale, (weight, scale) in zip(
(
experts._transformed_shared_l1_weights[0].view(torch.uint8),
originals[0][0],
0,
experts._transformed_shared_l1_weights,
experts._transformed_shared_l2_weights,
),
(FakeDeepGemm.scale_values[-2][0], originals[0][1], 1),
FakeDeepGemm.scale_values[-2:],
originals,
):
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()
)
assert torch.equal(transformed[0].view(torch.uint8), weight)
assert torch.equal(actual_scale[0], scale)
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)
vllm/models/deepseek_v4/nvidia/model.py
+5 -65
View File
@@ -204,7 +204,6 @@ 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
@@ -387,17 +386,6 @@ 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
@@ -408,7 +396,6 @@ 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,
@@ -416,27 +403,12 @@ 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,
@@ -477,11 +449,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.
# 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
gate_up.weight.data = transformed_l1[0]
self._transformed_shared_l1_weights = (
gate_up.weight.data,
transformed_l1[1],
)
self._transformed_shared_l2_weights = transformed_l2
def _prepare_shared_expert_scale(
@@ -491,8 +463,6 @@ 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:
@@ -527,19 +497,6 @@ 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.
@@ -558,23 +515,6 @@ class DeepseekV4MegaMoEExperts(nn.Module):
if self._transformed_l1_weights is None:
self._check_runtime_supported()
# MegaMoE's 1x32 activation scales need 16-byte TMA rows, so
# pad each gate/up half and the down projection to a multiple of 512.
padded_size = (self.intermediate_size + 511) // 512 * 512
padding = padded_size - self.intermediate_size
if padding:
for param in (self.w13_weight, self.w13_weight_scale):
gate_up = param.data.unflatten(1, (2, self.intermediate_size))
param.data = torch.nn.functional.pad(
gate_up, (0, 0, 0, padding)
).flatten(1, 2)
self.w2_weight.data = torch.nn.functional.pad(
self.w2_weight.data, (0, padding // 2)
)
self.w2_weight_scale.data = torch.nn.functional.pad(
self.w2_weight_scale.data, (0, padding // 32)
)
self.intermediate_size = padded_size
w13_scale = deep_gemm.transform_sf_into_required_layout(
self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(),
2 * self.intermediate_size,