[Core] Register persistent device tensors held outside modules

[aoshen02]

Overview

Adds a way to register every persistent device tensor as a parameter or buffer of the module that owns it, and a CI check that no new unregistered tensor appears. Sleep mode level 2, CUDA-graph address checks and weight reload only see parameters and buffers, so a tensor stored anywhere else comes back stale or freed after sleep/wake or a weight update. Context: weight reload RFC vllm-project#54477, tracker vllm-project#50079.

This is the first of a split series; follow-ups register the remaining sites per area and delete their entries from the CI allowlist. Follow-ups, each stacked only on this PR: #96 (Sentence-Transformers projector, sleep-mode bug fix), #97 (attention), #98 (linear kernels and schemes), #99 (MoE, NVFP4 sleep-mode bug fix), #100 (shared layers and model_executor models), #101 (DeepSeek-V4 / Kimi-K3), #102 (GLM5-next, Inkling, dots3).

Claims

• set_derived_buffer and register_constant_buffer give model code one way to keep weight-derived tensors and constants as non-persistent buffers, so state_dict() and checkpoints are unchanged.
• Non-module holders (quant methods by default) list their tensors in persistent_tensors(). register_held_tensors(model) registers them on the holding module after post-load and after layerwise reload.
• Layerwise reload refills derived buffers in place instead of leaving stale copies.
• test_initialization.py::can_initialize fails when an architecture keeps an unregistered device tensor, except for paths in KNOWN_UNREGISTERED. Today's remaining sites are listed there and are removed by the follow-ups.

Validation

GB200, image nightly-dev-arm64-cu130-ac9126e58aa7 overlaid with this change only.

Check	Result
test_initialization, all 396 registered architectures, dummy weights, oracle without the allowlist	363 pass with zero unregistered tensors. 29 report only paths now in KNOWN_UNREGISTERED. 4 fail the same way on main: DeepSeek-V4.1 (IndexError), LongcatFlashNgram (no cacheable KV group), Ultravox and HCXVisionV2 (transformers 5.18).
The 29 with the allowlist	29/29 pass
Unit tests: tests/model_executor/test_utils.py (set_derived_buffer, register_held_tensors through views, dtypes and wrapper subclasses, layerwise reload keeps derived state registered)	pass
pre-commit, mypy	pass

Gated checkpoints were swapped for ungated mirrors of the same architecture.

Details

• register_held_tensors registers each declared tensor as _held_<attr>_<name>, deduplicated per storage view (pointer, offset, shape, stride, dtype) and through traceable wrapper subclasses such as TorchAO. Rerunning it replaces every _held_* buffer; a holder still on meta keeps the imported buffer.
• Layerwise reload points derived buffers at their kernel storage before a layer is reprocessed, so set_derived_buffer refills them in place, and keeps a buffer that processing rebinds. The sharded RDT engine shares that preparation.
• Runtime state (KV caches, per-step scratch, workspaces, compile artifacts) is exempt by attribute name, with the reason in RUNTIME_STATE_ATTRS.
• The oracle only sees tensors reachable from module attributes; shape-keyed runtime caches (lru_cache, module globals) count as runtime state.

Not a duplicate. I searched open PRs for register buffer sleep mode, named_buffers sleep, unregistered tensors, persistent tensors buffer and derived buffer reload. @Ronald1995's Weight Reload series (vllm-project#53509, vllm-project#58991, vllm-project#58995, vllm-project#58999, vllm-project#59000, vllm-project#59002, vllm-project#59003, vllm-project#53510) and vllm-project#50521 / vllm-project#48539 make individual derived tensors reload-correct. This PR adds one registration contract and a CI check over every architecture; overlapping sites will rebase onto whichever lands first.

AI assistance. Written with Claude Code. Claude Fable 5.1, Codex (gpt-6-astra) and Kimi K3 audited the combined change in two rounds. The submitter reviewed every changed line and ran the checks 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