[Bugfix] Fix DeepSeek V4 mHC broadcast buffer for dummy load

[HollowMan6]

Purpose

With --load-format dummy, load_weights() is skipped, so
hc_attn_fn_broadcast stays None and the first decoder layer's 2-D
broadcast path fails at assert self.hc_attn_fn_broadcast is not None.

Fix: define a model-level process_weights_after_loading() on
DeepseekV4ForCausalLM. The loader invokes this hook for every load
format — including dummy load — so the broadcast buffer is always
populated before the first forward.

The hook also finalizes MegaMoE weights, replacing
the lazy finalize_weights() call in the MegaMoE forward that existed only
to cover dummy load. This takes the branch out of the forward hot path and
moves the heavy weight repack (which also frees the original params) out of
the first forward, so profiling sees steady-state memory. The EPLB-side
guard call in get_expert_mapping is unchanged.

Because the MTP and DSpark drafters load as their own top-level models
(registry → get_model → loader → PWAL), each gets a matching hook —
the removed forward-time call was their only finalize path under dummy
load.

The companion fix for weight refit (in-place copy_() refresh so the
buffer address stays stable for captured CUDA graphs) is split into a
separate PR.

Test Plan

Purpose

With --load-format dummy, load_weights() is skipped, so
hc_attn_fn_broadcast stays None and the first decoder layer's 2-D
broadcast path fails at assert self.hc_attn_fn_broadcast is not None.

Fix: define a model-level process_weights_after_loading() on
DeepseekV4ForCausalLM. The loader invokes this hook for every load
format — including dummy load — so the broadcast buffer is always
populated before the first forward.

Per review (@aoshen02), the hook also finalizes MegaMoE weights, replacing
the lazy finalize_weights() call in the MegaMoE forward that existed only
to cover dummy load. This takes the branch out of the forward hot path and
moves the heavy weight repack (which also frees the original params) out of
the first forward, so profiling sees steady-state memory. The EPLB-side
guard call in get_expert_mapping is unchanged.

Because the MTP and DSpark drafters load as their own top-level models
(registry → get_model → loader → PWAL), each gets a matching hook —
the removed forward-time call was their only finalize path under dummy
load.

The companion fix for weight refit (in-place copy_() refresh so the
buffer address stays stable for captured CUDA graphs) is split into a
separate PR.

Not a duplicate

Searched open PRs for this area: #50645 guards the broadcast kernel on
DeepGEMM support, #49707 fixes mHC warmup coverage (consumes the buffer,
doesn't populate it), #47807 streamlines warmup on the kernel side. None
touch load-time weight finalization or MegaMoE finalize placement.

Test Plan

Two unit tests added to tests/models/test_deepseek_v4_mega_moe.py:

• test_deepseek_v4_pwal_hook_finalizes_mega_moe_and_mhc_broadcast — the
  target model's hook runs both finalizes.
• test_deepseek_v4_drafter_pwal_hooks_finalize_mega_moe — the MTP and
  DSpark hooks delegate to their MegaMoE finalize.

End to end tests on verl side with DSV4 training

Test Result

• New tests: 2 passed; full file: 9 passed (GPU).
• Mutation check: with the three model files reverted to base, both new
  tests fail — they pin the hooks, not incidentals.
• ruff check + ruff format --check clean.

End to end tests on verl side with DSV4 training returned to normal training inference mismatch.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[Copilot]

Pull request overview

Fixes DeepSeek V4 mHC broadcast-buffer lifecycle so hc_attn_fn_broadcast is always populated for dummy-load initialization and remains pointer-stable across weight refits (important for CUDA graph replay correctness).

Changes:

• Call finalize_mhc_broadcast_weights() during DeepseekV4Model.__init__ to avoid dummy-load crashes when load_weights() is skipped.
• Update finalize_mhc_broadcast_weights() to refresh the existing broadcast tensor in-place (vs rebinding) to keep CUDA-graph-captured pointers valid across refits.
• Add unit tests asserting (1) correct stream-summed initialization and (2) in-place refresh behavior on refit.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
vllm/models/deepseek_v4/nvidia/model.py	Ensures broadcast buffer is initialized for dummy load and refreshed in-place on subsequent finalization to keep tensor address stable.
tests/kernels/test_mhc_kernels.py	Adds regression tests covering initial broadcast finalization and in-place refresh contract across refit.

Suppressed comments (1)

tests/kernels/test_mhc_kernels.py:392

• The new DeepSeek V4 broadcast-buffer tests should be skipped when the DeepSeek V4 NVIDIA model import isn't available (see HAS_DEEPSEEK_V4_NVIDIA from the guarded import), otherwise they will error at runtime (e.g., calling DeepseekV4Model.finalize_mhc_broadcast_weights). Add a skip marker to these tests so other kernel tests remain runnable.

def test_deepseek_v4_mhc_broadcast_finalize_sums_hc_streams(monkeypatch):
    """First finalize (run at __init__ so dummy load has a valid buffer)
    allocates hc_attn_fn_broadcast as hc_attn_fn summed over hc streams."""
    _patch_first_rank_pp_group(monkeypatch)
    layer = _make_mhc_decoder_layer(hc_mult=2, hidden_size=8)

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[github-actions]

✅ @HollowMan6, CI is now available for this PR.

• /ci run starts a CI build.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /ci cancel cancels scheduled or running CI builds for this PR branch.

[aoshen02]

LGTM, @WoosukKwon could you take a look?
@HollowMan6 Please add a gsm8k test result before and after this PR and make sure that init with dummy weight loader and then reload, the gsm8k is the same as normal weight loader.

[aoshen02]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84397 for commit 177101df1ff3.

[github-actions]

✅ Triggered Buildkite CI #84427 for commit e48860c91421.

[HollowMan6]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84449 for commit 5fd174ec98f6.

[HollowMan6]

@aoshen02 Now I verified this with verl e2e and found that removing finalize* from load_weights 177101d will actually break the sync, so I reverted the change in the last commit, and it should still be fine: #51368 (comment).

After the fix, I checked with verl e2e again by init with dummy weight loader and then reload with the normal weight load, and the training inference mismatch (log probs diff) are consistent and reasonably low at step 1, so I think we are good to go now:

rollout_actor_probs_pearson_corr:0.9952459931373596 rollout_probs_diff_mean:0.0063540199771523476