[Bugfix] Fix DeepSeek V4 mHC broadcast buffer for dummy load (#51368)
Signed-off-by: Hollow Man <hollowman@opensuse.org>
ℍ𝕠𝕝𝕝𝕠𝕨 𝕄𝕒𝕟 committed
Aug 19, 2026 at 03:02 UTC
a9f4afb66f77ba122d905611fa3f620abda98220
4 files changed
+46
-7
tests/models/test_deepseek_v4_mega_moe.py
+34
@@ -9,10 +9,13 @@ import torch
9
from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
10
bind_routed_experts_capturer,
11
)
12
+from vllm.models.deepseek_v4.nvidia.dspark import DSparkDeepseekV4ForCausalLM
13
from vllm.models.deepseek_v4.nvidia.model import (
14
+ DeepseekV4ForCausalLM,
15
DeepseekV4MegaMoEExperts,
16
make_deepseek_v4_expert_params_mapping,
17
)
18
+from vllm.models.deepseek_v4.nvidia.mtp import DeepSeekV4MTP
19
from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs
20
from vllm.platforms import current_platform
21
@@ -243,6 +246,37 @@ def test_deepseek_v4_mega_moe_fused_input_staging_is_bitwise_exact():
246
)
247
248
249
+def test_deepseek_v4_pwal_hook_finalizes_mega_moe_and_mhc_broadcast():
250
+ """The loader invokes the model-level PWAL hook for every load format,
251
+ so it must finalize megamoe + mhc broadcast weights to cover dummy
252
+ load, which skips load_weights()."""
253
+ calls = []
254
+ stub = SimpleNamespace(
255
+ model=SimpleNamespace(
256
+ finalize_mega_moe_weights=lambda: calls.append("mega_moe"),
257
+ finalize_mhc_broadcast_weights=lambda: calls.append("mhc"),
258
+ )
259
+ )
260
+
261
+ DeepseekV4ForCausalLM.process_weights_after_loading(stub)
262
+
263
+ assert calls == ["mega_moe", "mhc"]
264
+
265
+
266
+def test_deepseek_v4_drafter_pwal_hooks_finalize_mega_moe():
267
+ """MTP/DSpark drafters load as their own top-level models, so each needs
268
+ its own PWAL hook now that the megamoe forward no longer finalizes
269
+ weights lazily on first use."""
270
+ calls = []
271
+ mtp = SimpleNamespace(finalize_mega_moe_weights=lambda: calls.append("mtp"))
272
+ DeepSeekV4MTP.process_weights_after_loading(mtp)
273
+
274
+ dspark = SimpleNamespace(_finalize_moe=lambda: calls.append("dspark"))
275
+ DSparkDeepseekV4ForCausalLM.process_weights_after_loading(dspark)
276
+
277
+ assert calls == ["mtp", "dspark"]
278
+
279
+
280
@pytest.mark.skipif(
281
not torch.cuda.is_available(),
282
reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.",
vllm/models/deepseek_v4/nvidia/dspark.py
+4
-1
@@ -504,9 +504,9 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):
504
weight_loader(param, loaded_weight)
505
loaded_params.add(name)
506
507
- self._finalize_moe()
507
if self.model.confidence_head is not None and not loaded_confidence_head:
508
self.model.confidence_head = None
509
+ self.process_weights_after_loading()
510
logger.info_once("DSpark draft model loaded: %d params", len(loaded_params))
511
return loaded_params
512
@@ -514,6 +514,9 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):
514
for layer in self.model.layers:
515
layer.ffn.finalize_mega_moe_weights()
516
517
+ def process_weights_after_loading(self) -> None:
518
+ self._finalize_moe()
519
+
520
def _remap_dspark_name(self, name: str) -> str | None:
521
"""Map a checkpoint ``mtp.{i}.*`` name to this model's parameter path.
522
vllm/models/deepseek_v4/nvidia/model.py
+4
-5
@@ -504,10 +504,6 @@ class DeepseekV4MegaMoEExperts(nn.Module):
504
is_padding=is_padding,
505
)
506
507
- # This method must have been already called during the weight loading phase.
508
- # We call it again here to cover the dummy weight loading case.
509
- self.finalize_weights()
510
-
507
assert self._transformed_l1_weights is not None
508
assert self._transformed_l2_weights is not None
509
deep_gemm.fp8_fp4_mega_moe(
@@ -1543,9 +1539,12 @@ class DeepseekV4ForCausalLM(
1539
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1540
loader = AutoWeightsLoader(self, skip_substrs=["mtp."])
1541
loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
1542
+ self.process_weights_after_loading()
1543
+ return loaded_params
1544
+
1545
+ def process_weights_after_loading(self) -> None:
1546
self.model.finalize_mega_moe_weights()
1547
self.model.finalize_mhc_broadcast_weights()
1548
- return loaded_params
1548
1549
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1550
return self.model.get_expert_mapping()
vllm/models/deepseek_v4/nvidia/mtp.py
+4
-1
@@ -502,7 +502,7 @@ class DeepSeekV4MTP(nn.Module):
502
f"Use a checkpoint that includes MTP layer weights, "
503
f"or disable speculative decoding."
504
)
505
- self.finalize_mega_moe_weights()
505
+ self.process_weights_after_loading()
506
logger.info_once("MTP draft model loaded: %d params", len(loaded_params))
507
return loaded_params
508
@@ -510,6 +510,9 @@ class DeepSeekV4MTP(nn.Module):
510
for layer in self.model.layers.values():
511
layer.mtp_block.ffn.finalize_mega_moe_weights()
512
513
+ def process_weights_after_loading(self) -> None:
514
+ self.finalize_mega_moe_weights()
515
+
516
def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
517
"""
518
Rewrite the weight name to match the format of the original model.