[Ronald1995]

Fix the MiMo-V2 target model's attention sink bias load. load_weights slices the checkpoint tensor for the local TP rank, then uses the same current parameter loader path as other weights.

The MTP model already follows this contract and is unchanged. This PR changes only mimo_v2.py.

Validation

• ruff check, ruff format --check, and git diff --check pass.

• The integrated static audit finds no direct checkpoint sink copy. The new runtime tests have not been executed, as requested.

AI assistance

AI assistance was used.

[mergify]

Documentation preview: https://vllm--57804.org.readthedocs.build/en/57804/

[Ronald1995]

Verification on H20

Verified commit: 50ffbf6b74a8b888e72b90ef07247093f80e2bb1 (this PR's head; it covers MiMo-V2 target)
Integration commit: c623df9e4bfe1133a3205b94e2d771ba43876ec7 (attn-sink-integration = #57797 + #57798 + #57801 + #57803 + #57804 + #57805 + #57806)

Environment

• Cluster: H20 Slurm, node vllm-h20-01, 1x NVIDIA H20-3e (--partition=h20 --gres=gpu:h20:1)
• Python: $HOME/vllm-shared/prebuilt-venv/bin/python (torch 2.13.0+cu130 + compiled vllm from the shared prebuilt baseline)
• Pure-Python change, so the verified revision was overlaid onto the prebuilt package with make-overlay.sh; nothing was recompiled.

Results

Suite	Where	Result
test_attention_sink_load_audit.py (from #57806)	head node + GPU node	39 passed
test_attn_sink_reload.py (from #57797)	head node (CPU)	18 passed, 2 skipped (CUDA-only)
test_attn_sink_reload.py	H20 GPU	20 passed
GPU end-to-end sink reload script	H20 GPU	8/8 checks passed

The GPU script drives a real layerwise reload (record_metadata_for_reloading -> initialize_layerwise_reload -> load_weights -> finalize_layerwise_reload) on CUDA tensors and asserts the live update lands in the kernel tensor storage with the -inf padding tail re-established and the storage pointer unchanged:

PASS  live update reaches kernel tensor
PASS  live update re-establishes -inf padding
PASS  kernel tensor storage is stable
PASS  rank 1 live update loads its own heads
ALL GPU SINK CHECKS PASSED

Pre-commit

pre-commit run --files <changed files>
pass1 rc=0 dirty=0
pass2 rc=0 dirty=0

No hook rewrote a file. run-precommit.sh skips the pip-compile hooks because the sandbox cannot reach PyPI; no requirements file is touched here.

Not verified

• No full-model checkpoint was loaded (no matching weights available), so this is loader/contract-level plus a real GPU reload, not a full-model accuracy run.
• Single GPU; TP>1 was exercised in-process via distinct rank head ranges, not as a multi-process TP run.

What changed in this revision

The earlier revision reindented the whole tail of load_weights to drop an else, which also moved maybe_remap_kv_scale_name in front of the sink handling. That is unrelated to the sink write and changes behavior for a non-sink path. This revision keeps the sink branch self-contained: narrow, load through the current loader, loaded_params.add, continue. The diff is now the sink branch only.

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