[hao-aaron]

PR 6 of the direct weight reload stack (RFC vllm-project#54477), first part. Stacked on wu/stack/direct-g = vllm-project#57733 + vllm-project#62407 + vllm-project#58382.

Problem

GLM5-next runs its MLA projections (fused_qkv_a_proj = q_a + kv_a, q_b_proj, o_proj) in BF16, but its FP8 checkpoints (GLM-5.3-Flash) store them as block-FP8: weight + weight_scale_inv. The model's loader, and the MTP loader, held each weight until its scale arrived and dequantized it there.

That pairing buffer was a local of one load_weights call. An RL update that sends a weight and its scale in different calls (different buckets, or a different order) silently dropped both halves, and the projection kept its old values. A BF16 trainer works only because BF16 tensors skip that path.

Fix

Pair nothing in the loader. Register the checkpoint tensors as parameters under their checkpoint names, load each on its own, and combine them in process_weights_after_loading.

New Fp8DequantConfig / Fp8DequantLinearMethod (layers/quantization/fp8_dequant.py), for linear layers that a quantized checkpoint may store as block-FP8 but that run in BF16:

• weight is an FP8 view of the BF16 weight's own storage, with FP8 row i in the first half of the bytes of BF16 row i (row stride 2 * cols). There is no extra memory, and the TP sharding and fused-shard loading of the regular weight_loader_v2 path apply unchanged.
• weight_scale_inv is a regular block-scale parameter. It is retired after post-load processing ([Quantization] Keep deleted and renamed params registered after process_weights_after_loading vllm-project/vllm#62407) and stays registered, so a reload can load it again.
• A BF16 tensor under the weight's name (a BF16 checkpoint, or a trainer that sends BF16) is written into the same rows through the BF16 view, with no scale. Because both formats address the same rows, BF16 and FP8 parts of one fused weight can be mixed in one load.
• Post-load processing dequantizes the rows that hold FP8 bytes, in chunks, in place, so the weight keeps its startup address, which direct reload's CUDA graphs need. Afterwards every row holds runtime bytes, so a reload that leaves the weight out keeps it.
• Pre-processed weights (the weight cache) are kept as they are.

GLM5-next (attention.py) uses it for these projections whenever the checkpoint is quantized. The loaders (model.py, mtp.py) only adapt single tensors:

• Quark's weight_scale is loaded as weight_scale_inv.
• A NoPE kv_a_proj_with_mqa gets its zero RoPE rows on the weight, and its extra scale blocks on the scale, each tensor on its own.
• A BF16 weight counts its unneeded scale as loaded, so a BF16 checkpoint passes the default loader's completeness check.

_try_load_fp8_attn_proj and _dequant_fp8_block are removed.

Not in this PR (next parts of PR 6, same mechanism): the indexer wk FP8 dequant (DeepSeek-V2/V3.2 + MTP, HY-V4, dots3, GLM5-next) and MiMo-V2's fused FP8 qkv.

Tests

Unit tests on GB200, nightly dbcf568821 with this branch overlaid:

• tests/quantization/test_fp8_dequant.py, 13 passed (slurm 40414). It covers:

  • cold start equal to the dequantized checkpoint, at the startup address;
  • every tensor in its own load_weights call, with scales before weights;
  • BF16 and FP8 shards in one fused weight;
  • a BF16 checkpoint with no scales;
  • the TP row-parallel slice;
  • direct reload equal to cold start, with FP8 and with mixed updates, and parameter addresses unchanged;
  • a direct reload that leaves weights out keeps them;
  • layerwise reload (FP8 first, BF16 first, mixed);
  • pre-processed weights.

  Run on the first commit of this PR, before the review fixes, the layerwise FP8-first and pre-processed cases fail (slurm 40415).

• tests/quantization/test_quark.py -k glm5next: 4 passed (40409). The two loader tests are rewritten for the per-tensor adaptation.

• tests/model_executor/model_loader/test_direct_reload.py: 11 passed. test_model_finalize.py: 8 passed.

End to end on GB200 (slurm 40417), with the sharded-trainer benchmark (vllm-state-transition-verifier sharded-trainer-bench):

• Model: GLM-5.3-Flash truncated to 4 decoder layers (3 KDA, 1 MLA with MoE) plus the vision tower.
• Server: TP=2, with PR 5 ([Model Loader] Write weight-derived tensors back in place on reload aoshen02/vllm#180) also merged for the GLM-5 caches.
• Trainer: fsdp=2, NCCL engine, 1 warmup + 2 measured syncs alternating two perturbed checkpoints. Every tensor changes, and scales are scaled by powers of two.
• The audit compares every device tensor with a fresh server on the last checkpoint sent.
• The benchmark's WeightSource subclasses were adapted to the [RL] Declare weight ownership and layout on ParamMeta vllm-project/vllm#63256 contract for this run: metadata as a cached_property, iter_full, and no m2n_source import.

Run	Compared	Mismatched	Moved	Memory growth
this PR, direct, FP8 updates	218	2 (_workspace_buffer)	0	0
this PR, direct, BF16 update (MLA projections dequantized, scales dropped)	218	2 + 6 retired weight_scale_inv	0	0
base (direct-g + PR 5), direct, FP8 updates	212	2 (_workspace_buffer)	0	0
this PR, layerwise, FP8 updates	218	2 (_workspace_buffer)	0	0

• _workspace_buffer is the MLA backend's scratch space. It differs from the fresh server in every run, including the base, so it is not caused by this PR.
• In the BF16 run, the six retired scales are the cold-start FP8 scales, which the BF16 update does not resend. The fresh BF16 server never loaded any. No kernel reads them.
• All MLA projection weights match on both ranks. That includes q_b_proj and o_proj, which are sharded across the two ranks.
• With this engine and bucket size the base did not split a weight from its scale, so the failure this PR fixes is shown by the unit tests (separate calls, any order), not by this run.
• The MTP layer is covered by the unit-level loader changes only. The run had no speculative decoding.

🤖 Generated with Claude Code

https://claude.ai/code/session_01SeFSSrpw7yTpLz2C2sHDKA