diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index 9a5723d233..67029b50a9 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -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) diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 4aba14cd03..7fef691bd1 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -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,