diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index 0d3c2b6583..9a5723d233 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -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]) diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index b1379cc9eb..4aba14cd03 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -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.