[Ronald1995]

Overview

Attention sink writes bypassed parameter weight_loader callbacks, leaving stale sink values after layerwise reload during RL weight updates. Register checkpoint-backed sinks with a common TP loader and factory, and route model loading through the registered callbacks.

This consolidates and supersedes #57798, #57801, #57803, #57804, #57805, and #57806.

Claims

• DeepSeek V4/V4.1, HY-V4, MiMo-V2 (including MTP and Omni MTP), GPT-OSS, Laguna, and Qwen3-DFlash load sinks through their current parameter loader; GPT-OSS preserves MXFP4, Quark and other quantization dispatch.
• The loader partitions global checkpoint heads once, writes the local shard in place, and restores padded heads to -inf. The factory initializes sinks to -inf and registers the loader.
• Regression coverage checks changed sink values across two layerwise reloads, TP sharding, padding, storage preservation, quantization dispatch and pipeline-stage skips.

Validation

Current review revision: 7964a96f87

Merged upstream main at 61d9c777bf before committing the review fixes. Local Windows / Python 3.12 / PyTorch 2.13.0+cpu:

• Isolated CPU regression run: 35 passed, 1 skipped (CUDA case). Runs the exact repository sink test functions with real PyTorch, unchanged production reload modules, and production model loading methods/helpers extracted from source. It bypasses unrelated vLLM startup imports and model constructors, with stand-ins for unrelated configuration/attention/quantization infrastructure. This is isolated verification, not a standard-suite or full-model run.
• The six new DeepSeek V4/V4.1 and HY-V4 dispatch cases (TP ranks 0/1) detect an in-memory mutation that loads correct sink values while bypassing the parameter's current loader. Mutation harnesses remain outside the PR.
• Ruff checks/format, git diff --check, SPDX, forbidden-import and test-tethering checks passed.
• Standard pytest tests/model_executor/test_weight_utils.py tests/model_executor/model_loader/test_reload.py -k sink -q remains blocked during startup by missing uvloop on Windows. Full pre-commit could not initialize its environments because the configured local proxy is unreachable; commit hooks were bypassed for this revision after the checks above.
• No remote H20, GPU or full-model validation was performed for this review revision.

Historical H20 validation before this review revision

Previous revision

Environment: H20 Slurm cluster, node vllm-h20-02, 2 x NVIDIA H20-3e, driver 580.105.08; Python 3.12.14 / torch 2.13.0+cu130. Compiled extensions come from the per-commit precompiled cu130 wheel built for this PR's base commit (vllm-0.31.1rc1.dev295+gd4e92ccc3) with the branch's Python files overlaid, so Python and extensions are from the same commit without a rebuild.

Check	Result
pytest tests/model_executor/test_weight_utils.py tests/model_executor/model_loader/test_reload.py -k sink -q	41 passed (GPU node)
full test_weight_utils.py + test_reload.py, this branch	18 failed, 110 passed, 14 skipped, 2 xfailed, 7 errors
full same two files at upstream main (control)	18 failed, 70 passed, 13 skipped, 2 xfailed, 7 errors
failure-name diff, main -> branch	identical set (25 names); every additional pass is a sink test
run-precommit.sh (every hook the diff touches)	exit 0, no file rewritten
DCO over the branch range	all commits signed off, DCO pass

The failures are model-download / multi-GPU cases that fail identically before and after the change, so they carry no signal about this PR.

Additional out-of-tree verification (harness deliberately kept outside the PR, results in the review comments): each sink model's real attention module was built from tiny dimensions with dummy weights on H20, for TP=1 and TP=2, and driven through the real layerwise reload lifecycle, asserting the per-rank slice, the -inf padding and the preserved kernel storage; the pre-PR loading shape was reproduced as a control and never received the reloaded sinks.

These results were recorded before the current review changes and are not a rerun of the current head.

Details

A dedicated sink loader derives shard length from the checkpoint rather than the runtime parameter: DeepSeek pads its runtime sink, so sharded_weight_loader(0) would use the wrong shard length. Unpadded sinks use the same loader without reconstructing storage. Checkpoint heads must divide evenly across TP ranks.

The loader keeps the divisibility check and relies on tensor slicing and the default loader for incompatible shapes. Empty checkpoints are no longer explicitly rejected. Lifecycle rationale stays here rather than in long helper docstrings.

Tests in the existing CI-selected test_reload.py and test_weight_utils.py suites cover actual model dispatch for MiMo MTP/Omni MTP, Laguna, DFlash, GPT-OSS, DeepSeek V4/V4.1 and HY-V4. The DeepSeek/HY-V4 tests replace the current loader and verify the unsliced global checkpoint reaches it exactly once for each load. AST factory-name scanning and redundant empty/padding tests were removed.

No full-model accuracy evaluation, real-checkpoint end-to-end reload, TP>2 or multi-node run was performed. Pipeline parallelism is covered by the skip-path test. Backend-converted sink copies and GPT-OSS expert reload follow-ups are outside this PR; their behavior is unchanged.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill.
• 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: Adds focused regression coverage in existing CI-selected suites; validation limits are documented above.
• Code Quality and Style: Keeps code and comments concise and updates PR documentation.
• Pull Request Contents: Includes scope, validation evidence, root cause, tradeoffs and limitations.

[mergify]

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

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

[Ronald1995]

Verification on H20 (attention-sink series)

Verification of this PR used an integrated tree containing all six fix PRs, so the shared helper and every call site were exercised together.

Verified commit: 64c77ee60341e1aa51af3405a20934dfcd2ec8f2 (this PR's head, which is the shared helper plus the CUDA reload test)
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)
• The change is pure Python, so the verified revision was overlaid onto the prebuilt package with make-overlay.sh; no recompilation was needed.
• CPU run: H20 head node, same overlay, CPU torch only.

Results

Suite	Where	Result
tests/model_executor/model_loader/test_attention_sink_load_audit.py	head node (CPU)	39 passed
tests/model_executor/model_loader/test_attn_sink_reload.py	head node (CPU)	18 passed, 2 skipped (CUDA-only cases)
test_attn_sink_reload.py	H20 GPU	20 passed (the two CUDA cases run and pass)
GPU end-to-end sink reload script	H20 GPU	8/8 checks passed

GPU end-to-end output (real layerwise reload on device storage):

PASS  cold load is on device
PASS  cold load writes local heads
PASS  cold load keeps -inf padding
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
PASS  rank 1 keeps -inf padding
ALL GPU SINK CHECKS PASSED

Command: PYTHONPATH=<overlay> python -m pytest -q --import-mode=importlib <overlay>/tests/model_executor/model_loader/test_attn_sink_reload.py

Pre-commit

pre-commit run --files <changed files> on this branch, run twice to prove idempotence:

pass1 rc=0 dirty=0
pass2 rc=0 dirty=0

ruff check, ruff format, check-spdx-header, check-test-tethering, mypy (3.10-3.13) and the remaining hooks all pass; no hook rewrote a file. run-precommit.sh skips the pip-compile hooks because the sandbox cannot reach PyPI to build encodec/rouge-score/jsmin; no requirements file is touched by this PR.

Not verified

• No full-model checkpoint was loaded (no DeepSeek V4/V4.1/HY-V4/MiMo/GPT-OSS weights are available here), so this is a loader/contract-level verification plus a real GPU reload, not a full-model accuracy run.
• Run on one H20 node; TP>1 was exercised in-process by driving distinct rank head ranges, not as a multi-process TP run.

What changed in this revision

• The helper moved out of vllm/models/deepseek_v4/sink.py into vllm/model_executor/model_loader/attention_sink.py, next to default_weight_loader. It is part of the loading interface that HY-V4 loads through, so it does not belong under one model package.
• The reconstructed runtime tensor is built with torch.empty(param.shape, dtype=param.dtype, device=local_weight.device) instead of torch.full(..., -inf, dtype=...) plus an implicit dtype cast, so a checkpoint/parameter dtype mismatch is resolved before the load rather than relied on implicitly. Dtype coverage added.
• A head range wider than the padded parameter now raises ValueError instead of silently truncating; this is the one way the old shape could load the wrong heads. Coverage added.
• Added a CUDA-parametrized case to the reload suite (torch.cuda.is_available() skip guard) so CI exercises the device boundary, not just CPU. It passes on H20 (see above).

[Ronald1995]

/ci retry

[github-actions]

✅ Queued 1 failed job(s) for retry in Buildkite CI #94644.

[Ronald1995]

/ci retry

[github-actions]

✅ Queued 1 failed job(s) for retry in Buildkite CI #94644.

[Ronald1995]

/ci run

[github-actions]

✅ Triggered Buildkite CI #94710 for commit 15a4f28447c8.

[Ronald1995]

/ci run

[github-actions]

✅ Triggered Buildkite CI #94732 for commit 12d8e1425b80.

[Ronald1995]

/ci retry

[github-actions]

✅ Queued 1 failed job(s) for retry in Buildkite CI #94732.

[Ronald1995]

/ci retry

[github-actions]

✅ Queued 1 failed job(s) for retry in Buildkite CI #94732.