[Model][Core] Register every persistent device tensor as a parameter or buffer

[aoshen02]

Purpose

Sleep mode level 2 saves and restores only named_buffers(), and CUDA-graph address checks and weight reload only see parameters and buffers. A persistent device tensor stored any other way comes back stale or freed after sleep/wake or a weight update. That covers plain attributes, tensors held by quant methods, kernels, MoE quant configs or attention impls, and module-level caches.

This PR makes it a rule that every persistent device tensor is a parameter or a buffer, and enforces it in CI. Earlier versions of this PR registered weight-derived tensors only; this version covers every persistent device tensor, constants included, under the contract below.

Contract

• Model code registers state on the module that owns it, as a non-persistent buffer, so state_dict() and checkpoints are unchanged.
• Derived tensors are registered as None in __init__ and filled with set_derived_buffer. Later values are copied in place, and a shape or dtype change raises.
• Constants no checkpoint provides use register_constant_buffer, so layerwise reload does not wait for them.
• Objects that are not modules list their tensors in persistent_tensors(). This covers quant methods, linear kernels, fused-MoE experts / kernels / prepare-finalize / quant configs, attention impls and schemes. register_held_tensors(model) registers each of them on the module that holds the object, as _held_<attr>_<name>. It runs after process_weights_after_loading and after layerwise reload, and deduplicates per storage view, through traceable wrapper subclasses such as TorchAO.
• Persistent tensors are built in __init__, create_weights or post-load, never on the first forward.
• Runtime state is exempt by attribute name, with the reason listed in tests/models/utils.py::RUNTIME_STATE_ATTRS: KV caches, per-step scratch, workspaces, compile artifacts.

Changes

• Framework: vllm/model_executor/utils.py (set_derived_buffer, register_constant_buffer, register_held_tensors, held_tensors, dataclass_tensors, storage_views), called from model_loader/utils.py::process_weights_after_loading.
• Reload: reload/layerwise.py points derived buffers at their kernel storage before a layer is reprocessed, so values are refilled in place. A buffer that processing rebinds is kept instead of the stale one. The MoE runner hook and held-tensor registration also run after reload. The sharded RDT engine shares the same preparation.
• MoE runner: the fused router and shared-expert gate is built at post-load under device_loading_context, instead of lazily on the first forward.
• DeepGEMM warmup: it read the MoE quant config by building a second one. On FlashInfer CUTLASS per-tensor FP8 MoE that re-registered g1_alphas/g2_alphas and orphaned the tensors the kernel holds. It now reads the layer's existing config.
• Sites:
  • attention: sinks, ALiBi slopes, MLA absorbed weights, static sink K/V, TurboQuant constants
  • MoE and quantization: FP8 / NVFP4 / MXFP4 / W4A8 scales, CUTLASS strides, TRT-LLM constants, triton_kernels storage, b12x packed weights, Humming locks, Hadamard shared weights
  • layers: RMSNorm without weight, linear-attention slopes
  • models: DFlash / DSpark / LiLiCorr fused buffers, Gemma4 MTP, GLM5-next, Inkling, Kimi-K3, DeepSeek-V4 / V4.1 (NVIDIA, ROCm, XPU, CPU), dots3 audio STFT constants, pixtral and OpenPangu-VL rotary, Granite Speech, Sentence-Transformers projector weights under DispatchPooler
• CI: test_initialization.py::can_initialize asserts that find_unregistered_tensors finds nothing. Unit tests cover set_derived_buffer, register_held_tensors and the layerwise reload path.

Not covered (known, out of scope)

• Model-level hooks on layerwise reload. Layerwise reload does not re-run model-level post-load hooks, so DeepSeek-V4 / Kimi-K3 MegaMoE and GLM5-next derived weights are not rebuilt by a reload. This was already true on main. Moving these finalizers to the loader hook is [Model Loader] Run model-level post-load finalization from the loader hook vllm-project/vllm#57733.
• CUDA graph addresses after reload. A tensor that a reload rebuilds as a new object, such as an FP8 MoE quant config, is registered as the new tensor. A graph captured earlier still holds the old address, as on main.
• gpt-oss sinks on layerwise reload. gpt-oss layerwise reload returns NaN on both main and this PR because the OAIAttention sinks are not reloaded. This is fixed separately.
• Coverage of the CI oracle. It only sees tensors reachable from module attributes. Runtime shape-keyed caches (lru_cache, module globals) are runtime state.

Test Plan / Results

All GPU checks ran on GB200. The image was nightly-dev-arm64-cu130-ac9126e58aa7, overlaid with this change file by file, with SHA-256 gates. Base arm I is the same image with no overlay. Determinism checks ran with kernel_config.enable_flashinfer_autotune=false, because autotune alone makes several models differ from run to run.

Unregistered-tensor oracle (find_unregistered_tensors):

Check	Result
test_initialization over all 396 registered architectures, dummy weights	392 pass with zero unregistered tensors. 4 fail the same way without this PR: DeepSeek-V4.1 (IndexError), LongcatFlashNgram (no cacheable KV group), Ultravox and HCXVisionV2 (transformers 5.18 incompatibilities). Gated checkpoints were swapped for ungated mirrors of the same architecture.
Quantized checkpoints, 46 configs: FP8 / block FP8 / NVFP4 / MXFP4 / MXFP8 / AWQ / GPTQ / W4A16 / AutoRound / QuIP / SpinQuant / Quark / FP8 KV, dense and MoE	All 35 runnable configs: zero. 11 cannot run on this node: int8 CUTLASS on SM100, AMD-only, bitsandbytes not installed, GPTQ desc_act, W4A8 scheme.
Real weights: Kimi-K3 6-layer slice (TP4, EP, MegaMoE), DeepSeek-V4-Flash 4-layer slice (EP, MegaMoE), gpt-oss-20b, Qwen3-30B-A3B FP8 and NVFP4, DeepSeek-Coder-V2-Lite-FP8	Zero

On main the oracle finds, for example: FlashInfer CUTLASS FP8 MoE scales computed in the quant config, MXFP4 TRT-LLM constants, CUTLASS strides, attention sinks, Sentence-Transformers projector weights and MoE runner fused gates.

Sleep level 2 round trip. Generate, then sleep(2), wake_up(), reload_weights(), then generate again. The second run must be bit-identical (48 greedy tokens and top-5 logprobs).

Model	main	this PR
Qwen3-30B-A3B-NVFP4	differs	bit-identical
embeddinggemma-300m (ST projector)	max abs diff 0.153	<= 1.5e-8 over 4 runs (fp32 noise)
Qwen3-0.6B-FP8, Qwen3-30B-A3B-FP8 (also reload only), Qwen1.5-MoE, Qwen3-Next-80B-FP8, Qwen2.5-VL-3B, Gemma-3-4B, Gemma-4-E2B, Ministral-3-3B	bit-identical	bit-identical

DeepSeek-V4, Kimi-K3 and DeepSeek-Coder-V2-Lite already differ between two generations in one process without sleeping, on main and on this PR alike, so this check cannot decide them.

Outputs unchanged (main vs this PR, warm caches, same GPU). Greedy tokens and top-5 logprobs are bit-identical for every model that is deterministic from run to run: gpt-oss-20b, Qwen3-30B-A3B FP8 and NVFP4, Qwen3-0.6B (bf16 and FP8), Qwen3-Next-80B-FP8, Ministral-3-3B, Gemma-3-4B, DeepSeek-Coder-V2-Lite-FP8, DeepSeek-V4 4L and Kimi-K3 6L real weights.

Throughput is unchanged (vllm bench throughput, 1024 in / 256 out, 500 prompts, 3 reps): gpt-oss-20b 106.05 vs 106.39 req/s; Qwen3-30B-A3B-FP8 75.71 vs 75.77 req/s.

Unit tests: tests/model_executor/test_utils.py, tests/quantization/test_held_tensors.py, tests/model_executor/model_loader/, tests/models/test_deepseek_v4_mega_moe.py and tests/kernels/core/test_layernorm.py show the same failures as on main, plus the new tests. The failures on both are local _C and flashinfer build mismatches. pre-commit passes.

Not a duplicate

I searched open PRs for register buffer sleep mode, named_buffers sleep, derived weight buffer, unregistered tensors and persistent tensors buffer. Related work:

• The Weight Reload series by @Ronald1995 refreshes derived state in stable storage on reload, one model or backend at a time: [Weight Reload] Add module-owned runtime buffers and staged reload lifecycle vllm-project/vllm#53509 (module-owned buffers and staged reload lifecycle), [Weight Reload][GLM5-next] Refresh derived caches in stable storage vllm-project/vllm#58991, [Weight Reload][DeepSeek] Refresh hyperconnection broadcast buffers vllm-project/vllm#58995, [Weight Reload][ROCm] Refresh DeepSeek preshuffled weight scales vllm-project/vllm#58999, [Weight Reload][Gemma4] Refresh MTP gathered LM head in stable storage vllm-project/vllm#59000, [Sleep Mode][FlashInfer] Make DeepSeek epilogue alphas module-owned vllm-project/vllm#59002, [Weight Reload][MegaMoE] Rebuild module-owned packed expert buffers vllm-project/vllm#59003. [Sleep Mode] Restore Humming lock workspaces after wake and reload vllm-project/vllm#53510 restores Humming lock workspaces after wake.
• [Bugfix] Preserve FlashInfer CUTLASS MoE constant storage across weight reload vllm-project/vllm#50521 and [Bugfix] Preserve Machete act-order permutation storage across weight reload vllm-project/vllm#48539 preserve FlashInfer CUTLASS constants and Machete act-order permutations across reload.

These PRs make individual derived tensors reload-correct. This PR defines one registration contract for every persistent device tensor and enforces it with a CI oracle over every architecture. Some sites overlap: the Gemma4 MTP LM head, MegaMoE buffers, DeepSeek hyper-connection broadcast, FlashInfer epilogue alphas and Humming locks. I will rebase onto whichever of those lands first.

AI assistance

Written with Claude Code. Fable 5.1, Codex (gpt-6-astra) and Kimi K3 each audited the change in two rounds, and every finding was fixed or is listed above. The submitter reviewed every changed line and ran the tests and GPU checks above.

🤖 Generated with Claude Code