[aoshen02]

Problem

process_weights_after_loading reruns on every weight reload (layerwise and direct). Many quant methods, kernels and models build tensors derived from the weights in that step: kernel constants, activation global scales, decode-path copies, fused drafter buffers. They allocate them as new tensors each time. Captured CUDA graphs and kernels keep reading the tensors from startup, so after a reload they read freed memory.

The weights themselves are reloaded correctly (diff 0 against cold start), yet accuracy drops or collapses. The symptom is model-specific and silent: no error, only wrong output.

Fix

Write weight-derived tensors back into their startup storage on every rerun instead of allocating new ones.

• Parameters go through replace_parameter, which already refills in place; it now returns the kept parameter.
• Buffers and plain attributes go through the new refill_derived(obj, name, tensor): copy in place when shape, dtype and device match, else replace. It returns the held tensor.
• MoE experts objects, which quant methods rebuild on every rerun, name their constants in kept_constants. keep_constants(layer) keeps them on the layer through refill_derived and hands them back to the new object.

Sites:

• MoE kernel constants:
  • TRTLLM MXFP4 SwiGLU gemm1_alpha/beta/clamp_limit.
  • TRTLLM FP8 _g1_alphas/_g2_alphas/_g1_scale_c and SwiGLU gemm1_alpha/beta/clamp_limit.
  • FlashInfer CUTLASS (FP8/NVFP4/MXFP4) SwiGLU constants and MXFP8 fake_input_scale.
  • CuteDSL NVFP4 per_token_global_scale.
  • FP8 per-tensor FlashInfer CUTLASS g1/g2_alphas.
  • NVFP4 MoE a1/a2_gscale.
  • TRTLLM NVFP4 g1_scale_c and gemm1_* (registered by the experts' post-load step; kept from the constructor for Quark NVFP4, which skips that step).
  • CUTLASS FP8 ab/c_strides and CUTLASS W4A8 a/b/c/s_strides (online fp8_per_channel MoE on Qwen3.5-35B-A3B: illegal memory access after a reload before, GSM8K 0.644 → 0.623 with oracle unchanged after).
  • A runtime-dispatch wrapper (FallbackExperts: Triton-or-CUTLASS, Triton-or-DeepGEMM) passes keep_constants to both experts it holds.
• Kimi-K3 specialized decode: the packed MoE front, the padded TP16 dense down projection, and the KDA conv weight and FP32 norm. The MLA reciprocal q/k (inverse) scales are refilled in their buffers.
• GLM-5-next: indexer gate and KDA merged conv weight.
• Drafters:
  • DFlash/DSpark fused K/V and norm stacks (Qwen3, Laguna, Gemma4, Muse).
  • Gemma4 MTP gathered lm_head.
  • LiLiCorr inference buffers.
• Muse Spark:
  • The residual alpha/beta buffers.
  • The NVIDIA router output multiplier, now folded by the router weight's loader (once per load) instead of once per post-load run.
  • The NVIDIA NVFP4 MoE constants (g1_alphas, g1_scale_c, g2_alphas, a1/a2_gscale).

Tests

Unit tests:

• tests/model_executor/test_utils.py: refill_derived refills buffers and attributes in place and never makes a parameter.
• tests/quantization/test_keep_param_objects.py: a rebuilt experts object keeps the first object's constant addresses, for TRTLLM MXFP4/FP8/NVFP4, FlashInfer CUTLASS, CuteDSL, CUTLASS FP8/W4A8 and a FallbackExperts wrapper. Also covers the FP8 FlashInfer CUTLASS scales.
• tests/models/test_muse_spark_logits_local.py: the router fold happens once per load; the NVFP4 MoE constants keep their addresses across reruns.

All pass on GB200 (slurm 39929, 40053, 40106).

E2E on GB200: same-weight NUPD≥2 reloads. GSM8K is compared before and after, and every tensor is fingerprinted against cold start.

Reload in place (direct mode), stacked on vllm-project#57733 + vllm-project#62407 + vllm-project#58382 (+ fork #179):

Model	Job	GSM8K pre → post	Tensors changed
Muse Spark NVFP4 (NUPD=3)	40042	0.8385 → 0.8378 (oracle PASS)	0/3011
Muse Spark NVFP4, this head (NUPD=3)	40107	0.815 → 0.839 (oracle PASS; 98 lost / 130 gained of 1319, sampling noise)	0/3311
Muse Spark NVFP4, earlier head	40090	0.836 → 0.823 (oracle PASS; 125 lost / 107 gained of 1319, within sampling noise)	0/3311
DeepSeek-V4.1 (NUPD=3)	40010	0.909 → 0.904	0/1709
MiniMax-M3.1 vendor image (NUPD=3)	40013	0.9507 → 0.9507	0/3148
gpt-oss-120b TRTLLM MXFP4	39950	0.938 → 0.936	0/796

• Muse before this PR:
  • Without the router fold fix: 0.848 → 0.690, with 62 router weights differing from cold start (39889).
  • With the router fold fix but without the MoE constants fix: 0.833 → 0.751, with 420 tensors at new addresses (40036).
• Address check: with a moved-address check on, Muse 40044 completes NUPD=3 with nothing moved.

Layerwise reload, CuteDSL NVFP4 MoE, Qwen3-30B-A3B NVFP4:

	Job	GSM8K pre → post
Base	39882	0.832 → 0.039
This PR	39883	0.840 → 0.844

Dependencies

Stacked on vllm-project#62407 (keep param objects across post-load processing) and vllm-project#57733 (post-load placement). The base branch of this PR is their merge, 5e23af126d. Follow-up to vllm-project#54477.

🤖 Generated with Claude Code

[hao-aaron]

Gap to cover later: DeepSeek-V4 ROCm compressor GEMM fusion (found while auditing storage sharing for vllm-project#64336)

vllm/models/deepseek_v4/amd/rocm.py:1217 prepare_compressor_gemm_fusion (run from DeepseekV4ForCausalLM.process_weights_after_loading, amd/model.py:1578) does

fused_weight = torch.cat((main_weight, indexer_weight), dim=0)
main_weight.set_(fused_weight[:main_size])
indexer_weight.set_(fused_weight[main_size:])

so two checkpoint params end up on one storage, and the indexer weight moves to offset main_size * K while its recorded checkpoint offset is 0. On a direct reload:

• begin (_checkpoint_views) re-views the indexer weight at offset 0 of the fused storage, so the loader writes the indexer's checkpoint bytes over the main compressor rows;
• finish returns early (_fused_compressor_weight is not None), so the fused weight is never rebuilt. The early return also leaves it stale after a layerwise reload.

The forward (_run_parallel_input_projections, rocm.py:1288) only reads _fused_compressor_weight, so this fits this PR's refill_derived: keep the params on their own storage and refill the fused copy in place.

fused = torch.cat((main_weight, indexer_weight), dim=0)
refill_derived(self, "_fused_compressor_weight", fused)  # in place on reload
self._fused_compressor_split_sizes = (main_size, indexer_size)
# drop the set_() calls and the early return

Cost: one extra copy of the C4 compressor weights per layer (the set_ was saving that). Scope: DeepSeek-V4 (not V4.1, whose version returns False) on ROCm AITER, C4 layers with an indexer, NoopOffloader. Not a tracked family in vllm-project#54477 and not validated; untested here (no AMD hardware).

Posted by Claude Code on behalf of Aaron Hao.