[aoshen02]

Overview

Part 3 of 3 of keeping the loaded parameter objects through process_weights_after_loading (PWAL). After #62326 and #62405 (both merged), a loaded param can still disappear when PWAL deletes or renames it. This PR keeps those params registered, so after PWAL every loader-visible param is still registered under the same name, as the same object, and keeps its loaded data unless that would cost real memory.

Claims

• Renames alias one param under both names. The new move_parameter(layer, src, dst) registers the src param object under dst as well and records dst as its runtime name. The loader writes through the checkpoint name and kernels read the runtime name, both pointing at one storage, so a later replace_parameter on either name updates both and no extra memory is used. This covers:
  • compressed-tensors *_weight_packed (linear, NVFP4/MXFP4/W4A8 MoE);
  • INC MXFP4;
  • the NVFP4 MoE oracle;
  • Humming NVFP4 and Humming's checkpoint-to-standard renames;
  • compressed-tensors W8A16 block FP8, which now aliases weight_scale_inv to weight_scale.
• Loaded params no kernel reads keep their data and are marked retired (retire_parameter). This covers:
  • consumed NVFP4 per-tensor scales (ModelOpt, Quark);
  • small Humming leftovers (at most 1/256 of the largest converted tensor);
  • the input scales of ModelOpt and Quark NVFP4 MoE when a weight-only backend (Marlin, Humming) is selected. These used to be set to None, which unregistered them. The method now leaves them registered and passes None to the MoE quant config.
• Large released sources stay registered with empty storage (replace_parameter(layer, name, t.new_empty(0))). Their data was converted into a different-layout kernel tensor, so keeping it would double that memory. This covers:
  • unused zero points, e.g. symmetric GPTQ qzeros (Conch, CPU, Marlin, Zentorch via MPLinearKernel._release_unused_zp; int8 scaled-mm; Quark int8 linear and MoE);
  • KV-cache scale placeholders (kv_cache.py, compressed-tensors KV);
  • released weights (XPU W4A8, wNa8o8, dynamic 4-bit, CPU int8, Triton MXFP4 MoE, INC ARK);
  • Humming leftovers larger than 1/256 of the largest converted tensor.
• Not changed: dispatch_cpu_unquantized_gemm keeps main's fresh Parameter(empty(0)) on release, since the released weight can be a tied lm_head shared with embed_tokens.
• Readers see the kernel view.
  • kernel_parameters(layer) returns the live params, each object once under its runtime name. get_live_param treats retired or empty params as absent. Humming uses both helpers.
  • The Triton MXFP4 MoE wrapper is stored as a plain <name>_triton attribute read via triton_weight(), so it no longer takes over the param name.
• Coverage (whole-tree enumeration of PWAL-reachable sites on main, 1786 rows split across the three PRs): 73 delete/rename sites fixed here, plus 36 sites converted together with them (Humming's delete-then-recreate pattern), plus 7 n/a rows inside the changed hunks.

Validation

• pre-commit passes.

• On GB200 (rebased on main), the tests of every changed file: tests/quantization/ (all), tests/model_executor/test_utils.py, model_loader/test_reload.py, and the kernel/MoE/LoRA/model tests that import changed modules: 1609 passed.

  • Every remaining failure fails identically on main in the same image: Int8 cutlass on SM100, Hopper-only Machete/W4A8, gated repos, DeepGEMM UE8M0 int32 scales on SM100 (fp8_block_enabled[auto], fp8_pb_wo_checkpoint_setup), qutlass forward, int4 per-token KV kernel, AutoRound MXFP8 MoE NotImplementedError, and test_kv_scale_reload.
  • Existing tests that asserted a deleted or renamed param now assert it is still registered (renames alias the runtime name; Humming leaves the large unused checkpoint weights empty).
  • test_keep_param_objects.py runs ModelOpt and compressed-tensors NVFP4 linear and MoE with Marlin and Humming forced; test_utils.py covers aliasing, retiring, and the runtime name.

• End-to-end on GB200 (before the rebase), main vs the stack with this PR. Each run is a cold start with GSM8K (1319 questions, 5-shot), greedy decoding, a byte digest of every post-PWAL tensor, and an object census of loader-visible params:

  Config	GSM8K main / stack	Post-PWAL tensors	Loader params registered (main → stack)	Kept with empty storage
  ModelOpt NVFP4 (Qwen3-30B-A3B)	0.8878 / 0.8886	identical	69% → 100%	0
  ... forced Marlin (3 runs each)	0.8749, 0.8901, 0.8863 / 0.8779, 0.8893, 0.8855	identical	62% → 100%	0
  ... forced Humming (3 runs each)	0.8802, 0.8901, 0.8870 / 0.8734, 0.8863, 0.8961	identical	69% → 100%	0
  compressed-tensors NVFP4, forced Humming	0.8931 / 0.8946	identical	36% → 100%	0
  compressed-tensors MXFP4 MoE (CUTLASS, *_weight_packed → *_weight)	0.8582 / 0.8582 (1314/1319 identical text)	identical	46% → 100%	0

  • Marlin and Humming are not run-to-run deterministic. Two main runs agree on about 390 of 1319 GSM8K outputs, so their rows show three runs per arm.
  • The extra tensors in the stack are aliases, with no extra memory, or kept scalars and per-tensor scales. Kept data adds at most 0.15 MiB per model (under 0.001%).

• Direct load of W1 into the live params (model.load_weights, no PWAL): every checkpoint name now resolves.

  • On main, ModelOpt NVFP4 with Marlin or Humming fails with has no parameter 'w2_input_scale', and compressed-tensors MXFP4 MoE fails with has no parameter 'w2_weight_packed'.
  • With this stack, compressed-tensors MXFP4 MoE loads all 675 tensors with no param moved. ModelOpt NVFP4 then stops at a shape mismatch, because Marlin/Humming repack to a different layout.
  • Direct load is only exact for layout-preserving PWAL, as on main; layerwise reload is the supported path.

• Layerwise reload W0 → W1:

  • ModelOpt NVFP4: all tensors are byte-identical to a cold start on W1, and greedy matches main 8/8.
  • compressed-tensors MXFP4 MoE: all tensors are byte-identical to a cold start on W1, and greedy matches 8/8.

• Earlier runs of the stack ([Quantization] replace_parameter keeps the parameter object by default #62326) gave tests/model_executor/model_loader/test_reload.py: 74 passed, 1 failed (test_kv_scale_reload), the same as on main.

Details

• Reload: layerwise reload restores checkpoint-format params before loading. It goes through layer._parameters, so both alias names are restored, and a param kept empty is never handed to a weight_loader.
• Humming: used to delete every param and re-create the set it needs. It now writes through replace_parameter. Leftovers are skipped if they alias a replaced param, retired if small, and emptied if large.
• Weight-only NVFP4 MoE input scales: Humming's MoE quant config reads the input scales from the layer, so they are retired rather than left live. A live loaded scale would switch Humming to the static checkpoint scale, while main uses none.
• Known limitation: an aliased param appears under both names in state_dict(), so the UVA offloader and elastic-EP scale-up copy it twice (no data change).
• Not kept (exceptions, listed in [Quantization] replace_parameter keeps the parameter object by default #62326): the online FP8 requantization-source weight_scale, AutoGPTQ placeholder MoE biases, and torchao tensor subclasses.

---

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.

AI assistance: prepared with Claude Code; review passes by Codex, Kimi K3 and Claude Fable.

🤖 Generated with Claude Code

Update: loaded storage is never shrunk, emptied or split (0dcd74c..682db58)

A loader-visible param now keeps the checkpoint bytes it had at create_weights after PWAL. Its storage is never emptied. When two names refer to the same bytes, they are registered as one Parameter object. This supersedes the "kept with empty storage" item above for every site that has a plain tensor to back.

• Shrink. replace_parameter copies a smaller result into a prefix view of the param's own storage instead of rebinding it to a new, smaller allocation. Results of equal or larger size are still rebound. Online quantization (uses_meta_device) still frees its bf16 source. This covers:
  • ModelOpt/Quark NVFP4 MoE w13_weight_scale_2 / w13_input_scale ([E,2] → [E]);
  • merged per-tensor scales ([3] → [1] or scalar: Humming NVFP4 weight_scale_2, CT NVFP4 weight_global_scale / input_global_scale, FP8 input_scale.max());
  • the CUTLASS W4A8 FP8 scales and the CPU AWQ repack, which now go through replace_parameter.
• Triton MXFP4 MoE. The weights and scales are no longer emptied. Each param holds the swizzled storage, and its triton_kernels tensor points at it, so no second copy is kept.
• Humming. A converted tensor under a new name is registered as the leftover loaded param it came from (move_parameter). This applies when the tensor lives in that param's storage or replaces a large leftover; other leftovers are retired with their data.
  • This removes the second Parameter object that block-FP8 MoE w13/w2_weight_scale and NVFP4 input_scale_2 had over *_weight_scale_inv / input_global_scale.
  • It also removes the emptied CT WNA16 *_weight_packed.
• No more emptied loaded params.
  • Unused zero points (Marlin, Conch, CPU, Zentorch), dynamic 4-bit scales and INC ARK qzeros/scales are retired with their data.
  • Zentorch, the DeepSeek-V4 CPU wo_a pack, CT W4A8 int8 MoE and CT WNA8O8 put the repacked tensor into the loaded param's storage. The K3 gfx942 MXFP4 repack no longer empties its source.
• Not changed (reported): sites whose kernel copy is opaque or split across several tensors still release the loaded weight. Keeping it would double weight memory:
  • CPU oneDNN int8 (scaled_mm/cpu.py) and dispatch_cpu_unquantized_gemm;
  • AMX MLA kv_b_proj;
  • the SM12x b12x linear/MoE/Q8_0 packs.
• Coverage: whole-tree AST call-graph enumeration from every PWAL root (1715 site rows), cross-checked with a git grep for release and alias patterns.
  • 1030 rows are not applicable: creation time, or not a loader param.
  • 409 replace_parameter rows are covered by the shrink rule.
  • The remaining rows were reviewed by hand. The .data rebinds there keep the same or a larger size, except the two converted above.

Census on GB200 (verify-generic-restore census: C1 = storage ≥ checkpoint bytes, C2 = another param or buffer overlapping the bytes). Before = a693c5e; after = 682db58. 4-layer slices; gpt-oss-20b full.

config	C1 before → after	C2 before → after
ModelOpt NVFP4 MoE default / FI-CUTLASS / CuteDSL / vLLM CUTLASS	8 → 0 each	0 → 0
ModelOpt NVFP4 Marlin / Humming	4 / 8 → 0	0 → 0
CT NVFP4 default / Marlin	8 → 0	12 → 12 (KV-scale buffer views, harmless)
CT NVFP4 Humming (new)	– → 0	– → 12 (same KV-scale views)
gpt-oss Triton MXFP4	96 → 0	0 → 0
Qwen3-30B block-FP8 Humming	0 → 0	8 → 0
controls: q30-fp8 default/Marlin, gpt-oss default	0 → 0	0 → 0

In every config, all loader params are still registered as the same object (missing=0), and generation works.

• Tests (GB200): model_executor/test_utils.py, quantization/test_keep_param_objects.py (now also asserts storage ≥ loaded bytes and no two param objects on one storage), test_moe_wna16, test_zen_da8w4, test_gpt_oss_triton_kernels, test_gfx942_int4, test_auto_round, test_cpu_wna16, test_triton_w4a16, test_cutlass_w4a16, test_reload, and the MXFP4 / scaled-mm kernel selection tests.
  • 299 passed.
  • 10 failed, and each fails identically on main: Hopper-only Machete, AutoRound MXFP8 MoE, and test_kv_scale_reload.
  • pre-commit passes.
• Greedy GSM8K, 5-shot completion:
  • ModelOpt NVFP4 Qwen3-30B-A3B, default backend, 300 questions: this change vs a32836e agrees on 245–256/300 outputs, the same as run-to-run (262/300). Accuracy is 0.927 / 0.930 vs 0.923.
  • gpt-oss-20b on Triton MXFP4 (chat, 200 questions): byte-identical to main (200/200 vs a main repeat), accuracy 0.675.

🤖 Generated with Claude Code

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

[hao-aaron]

/ci run

[github-actions]

✅ Triggered Buildkite CI #97799 for commit 81e789157847.

[hao-aaron]

/ci run

[github-actions]

✅ Triggered Buildkite CI #97876 for commit e61d991ffb53.

[hao-aaron]

/ci run

[github-actions]

✅ Triggered Buildkite CI #98031 for commit 792089aa89ef.

[mergify]

Hi @aoshen02, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[hao-aaron]

/ci run

[github-actions]

✅ Triggered Buildkite CI #98047 for commit 9055550af931.

[hao-aaron]

/ci run

[github-actions]

✅ Triggered Buildkite CI #98058 for commit 93bbe8ff662d.