[Model Loader] Run model-level post-load finalization from the loader hook

[aoshen02]

Stacked on #59927 (DeepSeek-V4 MegaMoE shared-expert finalize independent of linear post-load order). This branch contains that commit first; review only the top commit, [Model Loader] Run model-level post-load finalization from the loader hook.

Overview

Several models finalized weight-derived state from inside load_weights, lazily from forward() / EPLB get_expert_weights(), or not at all. BaseModelLoader.load_model already dispatches the model's process_weights_after_loading once per cold load, including load_format="dummy", which never calls load_weights. This PR moves that work to the hook, so every cold-load path runs it once after all weights are in place. It also routes raw (non-vLLM-serialized) tensorizer checkpoints through BaseModelLoader.load_model, the same change as #55318.

Context: weight reload (RFC #54477, tracker #50079) needs load_weights to only load. RFC #59502 assumes the model hook is cold-start-only, as it is here.

What changes

model / path	before	after
Kimi-K3 MTP (MegaMoE)	finalize_weights in load_weights; also lazily from forward()	KimiK3MTP.process_weights_after_loading; forward() asserts
Kimi-K3 target (MegaMoE)	hook, plus a lazy call in forward()	hook only; forward() asserts
DeepSeek-V4 NVIDIA target, MTP, DSpark	finalize_mega_moe_weights in load_weights	hook only
DeepSeek-V4 XPU target, MTP, DSpark	finalize_mega_moe_weights in load_weights, plus a lazy call in forward()	hook only; forward() asserts
DeepSeek-V4 VL wrapper	sorted weights and skipped its hook for a child that finalizes in load_weights (finalizes_weights_during_load)	always delegates the hook; the flag is gone
DeepSeek-V4 NVIDIA get_expert_weights() (EPLB)	lazily finalized	asserts
MiniCPM-V / MiniCPM-O resampler device	moved in load_weights	hook
OpenPangu text root, MTP	text: sweep in OpenPanguModel.load_weights; MTP: never refreshed sink keys	both roots also sweep post_weight_load from the hook; VL root delegates to its language model
Tensorizer, raw checkpoint	private helper: load_weights, then return (no per-layer quant processing, no model hook)	BaseModelLoader.load_model
Medusa	token_map.to(...) with the result discarded	removed (no-op)

Moving DeepSeek-V4 NVIDIA to the hook means the MegaMoE shared-expert fusion always runs after the linear hook has packed its scales; #59927 (the first commit here) makes that correct.

Not changed: DeepSeek-V4.1 still finalizes inside load_weights. Its mega-attention finalize permutes the wq_b / wo_a weight scales in place, and running it after the linear hook is not verified yet.

Evidence

A/B on GB200, image vllm/vllm-openai:nightly-dev-arm64-cu130-ac9126e58aa7 (sha256:8bafba97…, built from ac9126e58aa7).

• A = that commit + [Bugfix][Spec Decode] Scope MTP completeness checks for Kimi K3 #50080.
• B = A + [Model][DeepSeek-V4] Make MegaMoE shared-expert finalize independent of linear post-load order #59927 + this PR (same patch-ids as the pushed commits). The Vision-Exp and GSM8K rows were run after the DeepSeek-V4 NVIDIA move to the hook.
• Python files are overlaid per file, with SHA-256 gates against the image commit.
• Sink-attention models need draft [Bugfix][Core] Accept needs_kv_cache_zeroing in SinkFullAttentionManager #59512 to start at all, on that commit and on main. Those rows apply [Bugfix][Core] Accept needs_kv_cache_zeroing in SinkFullAttentionManager #59512 to both arms.
• Greedy decoding; token ids and top-5 logprobs are compared as raw floats, exactly. A and B run in separate processes.

Fixed (fails before, works after):

case	before (A)	after (B)
Qwen3-0.6B-FP8, Qwen3-4B-FP8 from a raw tensorizer file	crash (DeepGEMM fp8_gemm_nt, first forward)	identical to a default load
DeepSeek-V4-Flash 4-layer config (FP4 experts), MegaMoE + EP, dummy	startup crash, DeepGEMM layout.hpp:114	starts with the shared expert fused (fixed by #59927, see there)
MiniCPM-V-2_6, dummy	startup crash: indices should be either on cpu or on the same device as the indexed tensor (cpu)	starts; pos_embed on cuda:0
MiniCPM-o-2_6, dummy	same crash	starts, generates
openPangu-R-7B-2512 (sinks), dummy	27/27 layers: cached sink key != k_layernorm(param_sink_key) (max abs diff 1.52)	0/27
Kimi-K3 4-layer + 1 MTP config (56 experts), MegaMoE, MTP, dummy	MTP experts finalized from forward(); 32 more finalize_weights calls during generation	finalized once by KimiK3MTP.process_weights_after_loading; 0 calls during generation

Unchanged (bit-identical before/after, real weights):

case	result
Qwen3-0.6B-FP8, Qwen3-4B-FP8, default load	identical
Qwen3-0.6B + online fp8, default and raw tensorizer	identical
MiniCPM-V-2_6, image + text	identical
MiniCPM-o-2_6, 4 text + 1 image	identical
DeepSeek-V4-Flash-Vision-Exp, DP4 / EP4, MegaMoE, DSpark (5 speculative tokens), vllm serve	identical; DSpark totals equal (290 draft tokens, 139 accepted)
DeepSeek-V4-Flash, TP2 / EP, MegaMoE, MTP (2 tokens), GSM8K 5-shot, 1319 questions (CI config DeepSeek-V4-Flash-deepgemm-mega-moe.yaml)	A: 0.953 and 0.947 over two runs; this PR: 0.951. Shared expert fused in both (no fallback warning). GSM8K differs from run to run on main as well, so this row checks accuracy, not bit-identity.

Reproduction

Scripts are in the PR author's validation directory; each records token ids and raw logprobs to JSON, compared with cmp_ab.py.

• g3_suite.sh: every "fixed" row.
• g1_ab.py / g1_serve.py: the "unchanged" rows. --probe sink compares attn.sink_key with k_layernorm(param_sink_key) per layer.
• The DeepSeek-V4 4-layer config is DeepSeek-V4-Flash with 4 layers, FP4 experts and compress_ratios [0, 0, 4, 128, 0]. The existing 4-layer Base proxy has FP8 experts, which MegaMoE rejects.

Unit tests (local, humming-kernels==0.1.16 as pinned):

pytest tests/models/test_deepseek_v4_vl.py tests/models/test_openpangu_post_load.py \
  "tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py::test_deserialized_hf_weights_are_ready_for_forward" \
  tests/models/kimi_k3/test_weight_loading.py \
  tests/model_executor/model_loader/test_reload.py -m "not slow_test" \
  -k "not test_reload_weights and not test_kv_scale_reload and not test_online_quantize_reload"

The two new behaviour tests fail on main and pass here:

• The tensorizer test fails with 'Linear' object has no attribute 'cpu_linear'. It is adapted from [Bugfix][Model Loader] Run post-load processing for raw Tensorizer weights #55318's test.
• The OpenPangu MTP test has no hook to call.

pre-commit passes, except mypy-3.10 on vllm/models/deepseek_v4/xpu/model.py, which reports the same five deep_gemm attr-defined errors on unmodified main. The new and moved test files are wired into the "Basic Models (Other)" Buildkite job, NVIDIA and AMD.

Still running (results will be added here):

• Full Kimi-K3 (1.5T MXFP4) on 4 GB200 nodes, TP4 x DP4, EP16, MegaMoE: before/after cold-start comparison, then a weight update of a perturbed checkpoint compared against a fresh engine.
• DeepSeek-V4.1-Flash, DP4 / EP4, MegaMoE.

Not validated:

• XPU. The DeepSeek-V4 XPU move to the hook is untested.
• openPangu-VL-7B. Its required remote code fails to import with this image's transformers.
• OpenPangu sinks on real weights. No public sink checkpoint loads in vLLM today: R-7B-2512 has a merge_conv module; R-72B-2512's config lacks first_k_dense_replace.

Limitations (found, not fixed here)

• finalize_layerwise_reload does not dispatch the model hook. Derived state that only the hook rebuilds is not refreshed by a layerwise weight reload; reload-time refresh is RFC [RFC]: Direct Weight Reload for RL Training #54477's scope. OpenPangu's text model also sweeps in load_weights; its MTP drafter does not.
• The vLLM-serialized tensorizer branch still skips post-load by design: it restores an already-processed snapshot. MegaMoE state is not in that snapshot.

Overlap. The raw-tensorizer change is the same as draft #55318 (nicole-lihui), and its test is adapted from there. This PR needs it: it moves MiniCPM's device fix-up out of load_weights, which was the only finalization that path got. Whichever lands first, the other drops the hunk.

Not a duplicate. I searched for finalize_mega_moe_weights, process_weights_after_loading load_weights, tensorizer process_weights_after_loading and post-load finalization loader. Related open PRs: #55318 (above), #58889 (runtime state for the vLLM-tensorized branch) and #42823 (makes load_weights re-entrant via layerwise reload; complementary).

AI assistance. Written with Claude Code. Audited adversarially by Claude Fable 5.1, Codex (gpt-6-astra) and Kimi K3 over three rounds. Round 1 found two problems:

• A follow-up commit moving the DFlash/DSpark build into the hook broke CPU and quantized drafters, so it was dropped.
• The previous description claimed changes the diff did not contain.

The submitter reviewed every changed line and ran the tests and A/B above.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

🤖 Generated with Claude Code

[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.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @aoshen02.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[vllm-agent]

CI selector (shadow): 45 test steps (93 jobs) instead of 66 (99 jobs)

Shadow mode: this changes nothing about what CI runs. It shows what the evidence-based selector would pick for this PR, next to today's rules. How it works.

Feedback welcome: reply here if it would skip a step this change needs, or runs something unrelated.

steps (jobs)	Today's rules	Selector	Would skip	Would add
NVIDIA, CPU and others	66 (99)	45 (93)	40 (45)	19 (39)
AMD mirrors	75 (105)	21 (48)	58 (67)	4 (10)

Selector would run (45)

• arm-cpu-test ×3
• basic-correctness-cumem
• basic-correctness-sleep-mode
• basic-models-test-other-cpu
• basic-models-tests-extra-initialization ×14
• basic-models-tests-initialization
• basic-models-tests-other
• cpu-kernel-tests ×2
• cpu-language-generation-and-pooling-model-tests ×3
• cpu-multi-modal-model-tests-n ×4
• cpu-multimodal-config
• cpu-params-env-tokenizers-parser
• cpu-reasoning-renderers
• cpu-spec-decode-tests
• deepseek-v4-kernel-test-b200
• deepseek-v4-kernel-test-h100
• distributed-dp-tests-4-gpus
• distributed-model-tests-2-gpus ×3
• distributed-torchrun-examples-4-gpus
• engine
• entrypoints-integration-api-server ×4
• examples
• glm5next-unit-tests
• inkling-unit-tests-b200
• kernels-attention-test ×7
• kernels-flashmla-test-h100
• kernels-mhc-test-b200
• kernels-root-misc-test-b200
• kimi-k3-prefix-cache-4xb200 ×2
• kimi-k3-unit-tests-b200
• kv-offload-large
• language-models-tests-extra-standard ×2
• language-models-tests-standard
• lora-tp-distributed ×4
• model-executor
• multi-modal-models-standard-4-other-whisper
• multi-modal-processor ×4
• multi-modal-processor-cpu ×4
• quantization ×4
• qwen4-exp-unit-tests
• qwen4-exp-unit-tests-cpu
• v1-attention-b200 ×2
• v1-attention-h100-mi300 ×2
• v1-executor-worker
• v1-spec-decode

Would skip (today's rules run them) (40)

• amd-lm-eval-small-models-harness
• amd-qwen3-next-mtp-async-eplb-accuracy
• ascend-npu-test
• async-engine-inputs-utils-worker
• basic-correctness ×2
• basic-correctness-cpu-offload
• basic-correctness-prefetch-offload
• benchmarks-cli-test
• entrypoints-integration-api-server-generate
• entrypoints-integration-api-server-openai-chat_completion
• entrypoints-integration-api-server-openai-completion
• entrypoints-integration-llm
• entrypoints-integration-multimodal
• entrypoints-integration-pooling
• entrypoints-integration-responses-api
• entrypoints-integration-speech_to_text
• fusion-e2e-quick-h100
• fusion-e2e-tp2-b200
• fusion-e2e-tp2-quick-h100
• kernels-fla-ops-test-b200
• language-models-tests-granite-l4-compatibility
• language-models-tests-hybrid ×2
• metrics-tracing-2-gpus
• multi-modal-models-standard-1-qwen2
• multi-modal-models-standard-2-qwen3-gemma
• multi-modal-models-standard-3-llava-qwen2-vl
• pipeline-context-parallelism-4-gpus
• pytorch-compilation-dynamic-shapes
• pytorch-compilation-passes-unit-tests
• pytorch-compilation-unit-tests
• pytorch-compilation-unit-tests-h100
• pytorch-fullgraph-cudagraph-l4-compatibility
• pytorch-fullgraph-test
• regression
• v1-core
• v1-kv-connectors ×4
• v1-kv-offload
• v1-logits-oracle
• v1-metrics-lmeval
• v1-sample

Would add (today's rules do not run them) (19)

• arm-cpu-test ×3 (code map)
• cpu-kernel-tests ×2 (code map)
• cpu-multi-modal-model-tests-n ×4 (code map)
• cpu-multimodal-config (code map)
• cpu-params-env-tokenizers-parser (code map)
• cpu-reasoning-renderers (code map)
• cpu-spec-decode-tests (code map)
• deepseek-v4-kernel-test-h100 (code map)
• distributed-dp-tests-4-gpus (code map)
• distributed-torchrun-examples-4-gpus (Python record)
• engine (code map)
• examples (code map)
• kernels-attention-test ×7 (code map)
• kernels-flashmla-test-h100 (code map)
• kv-offload-large (code map)
• lora-tp-distributed ×4 (Python record)
• quantization ×4 (code map)
• v1-attention-b200 ×2 (code map)
• v1-attention-h100-mi300 ×2 (code map)

AMD mirrors: would skip (58)

• async-engine-inputs-utils-worker
• basic-correctness ×2
• basic-correctness-cpu-offload
• basic-correctness-prefetch-offload
• benchmarks-cli-test
• entrypoints-integration-api-server-generate
• entrypoints-integration-api-server-openai-chat_completion
• entrypoints-integration-api-server-openai-completion
• entrypoints-integration-llm
• entrypoints-integration-multimodal
• entrypoints-integration-pooling
• entrypoints-integration-responses-api
• entrypoints-integration-speech_to_text
• extract-hidden-states-integration-2-gpus
• fusion-e2e-quick-h100
• fusion-e2e-tp2-ar-rms-config-sweep-h100
• fusion-e2e-tp2-b200
• fusion-e2e-tp2-quick-h100
• kernels-fla-ops-test-b200
• kernels-root-misc-test-b200
• language-models-tests-extra-standard ×2
• language-models-tests-granite-l4-compatibility
• language-models-tests-hybrid ×2
• language-models-tests-standard
• lm-eval-small-models
• metrics-tracing-2-gpus
• mrcr-eval-small-models
• multi-modal-accuracy-eval-small-models
• multi-modal-models-standard-1-qwen2
• multi-modal-models-standard-2-qwen3-gemma
• multi-modal-models-standard-3-llava-qwen2-vl
• openai-api-correctness
• pipeline-context-parallelism-4-gpus
• platform-tests
• pytorch-compilation-dynamic-shapes
• pytorch-compilation-passes-unit-tests
• pytorch-compilation-unit-tests
• pytorch-compilation-unit-tests-h100
• pytorch-fullgraph-cudagraph-l4-compatibility
• pytorch-fullgraph-test
• quantized-models-test
• quantized-moe-test-b200
• regression
• spec-decode-draft-model ×4
• spec-decode-eagle-1-deepseek-qwen
• spec-decode-eagle-2-llama3-qwen-vl-other
• spec-decode-mtp-deepseek-mimo
• spec-decode-mtp-gemma4
• spec-decode-mtp-qwen3-5
• spec-decode-ngram-suffix
• spec-decode-speculators
• v1-core
• v1-kv-connectors ×4
• v1-kv-offload
• v1-logits-oracle
• v1-metrics-lmeval
• v1-sample
• v1-spec-decode

AMD mirrors: would add (4)

• distributed-torchrun-examples-4-gpus (Python record)
• examples (code map)
• lora-tp-distributed ×4 (Python record)
• quantization ×4 (code map)

_{22 changed files · base 0872ddf4b9 · head 38e1e36fa2 · Python record: build 92785 at 18f8f96025 · kernel record: table 18f8f96 (build 92785), map 18f8f96 · not counted: 10 build steps, 5 A100 steps the generator no longer emits, 21 optional steps the selector would also run}

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @aoshen02.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork