[hao-aaron]

Overview

MoE loaders record the padding they narrow away instead of writing it, and the padding is filled right before post-load processing, on every load, with the value each parameter was created with. A reload into storage that already held data (layerwise rematerialization, or an in-place reload into post-processed bytes) then post-processes exactly the bytes a cold start would. Replaces #50523 and aoshen02#179; it also batches the per-expert zeroing of aoshen02#179, which cost about 1.5 s per update on DeepSeek-V4.1 TP4 (~100k small memsets per rank), into one strided fill per region. Part of RFC #54477.

Claims

• Reload of padded MoE experts matches a cold start. gpt-oss-20b (MXFP4, experts padded 2880 → 3072): after a weight update, 96 padded expert parameters differ from a fresh server on main, and 0 with this PR (verifier NCCL sync); 280 → 0 for a disk reload at TP2.
• Cold start unchanged where main initialized the padding. Every parameter and buffer after a cold start is bitwise identical to main on gpt-oss-20b, Qwen3-30B-A3B (BF16 TP4, FP8 TP1) and Qwen3-30B-A3B-NVFP4 (TP4). Where main left padding uninitialized (e.g. unquantized MoE weights from torch.empty), it is now 0. Ones-initialized integer scales (Quark OCP_MX uint8 E8M0) now get 0 padding instead of 1; harmless, since their weight padding is zero. On DeepSeek-V4.1-Flash (TP4) every parameter is identical; only runtime buffers that change during generation (engram.staged_rows, topk_indices_buffer) differ.
• Covers MoE expert padding in every loader that goes through RoutedExperts.weight_loader (Inkling's experts bypass it and write their own padding, see [Model] Inkling, Muse Spark: write checkpoint tensors through the parameter's weight loader #64690). Hidden/intermediate rounding of weights, scales and biases, the gpt-oss combined, per-expert (streamed) and Quark paths, and block-aligned FP8 TP. The value is recorded right after create_weights: a padding_value attribute if set, else 1 for a floating-point parameter created with ones (its bytes equal the dtype's 1.0, e.g. torch.ones FP8 block scales), else 0. Uninitialized (torch.empty) storage never sets the value, even when it happens to hold a constant.
• Online quantization is unaffected. Layers post-processed during loading are not filled again: with online fp8_per_block (padded experts), cold start and reloads are bitwise identical to main.
• Cost: 4–7 ms per post-load processing pass on the models above, plus the marking in the loaders (gpt-oss-20b sync 0.41 → 0.45 s); nothing on the forward path.

Validation

All on GB200. "main" is the base commit; each side is compared with a fresh start of the same tree. The verifier, gpt-oss and NVFP4 runs used the current head 95bf429 against its merge base 15cc0a8; the unit-test A/B used 8f20479, and the Qwen3 BF16/FP8 rows earlier commits of this PR (a2080a4, and d36411e for online FP8). CI covers the unit tests on the current head.

Weight update through the verifier (vllm-state-transition-verifier run_bench.py --audit). gpt-oss-20b, trainer FSDP=2 → server TP1 over NCCL, one sync to a perturbed checkpoint (every tensor changed), every parameter, buffer and kernel-held tensor compared with a fresh server on that checkpoint:

	main	this PR
parameters ≠ fresh	96 (w13/w2 bias and weight scale)	0
non-parameter tensors moved / extra	72 / 96	72 / 96
sync time (one sync)	0.41 s	0.45 s

The 72/96 non-parameter tensors are moe_quant_config and TRT-LLM gemm1_* constants rebuilt by post-processing: a separate issue (#63090, aoshen02#180), unchanged here.

Disk reload (reload_weights, layerwise) on current main, reload to a perturbed checkpoint twice and back to the original:

Model, backend	main	this PR
gpt-oss-20b TP2, TRT-LLM MXFP4	280–288 expert tensors ≠ fresh; back-to-original differs from cold start in 480 tensors	0 expert tensors ≠ fresh; back-to-original: no parameter differs
Qwen3-30B-A3B BF16 TP4, TRT-LLM (intermediate 192 → 256)	pass; GSM8K 0.890 → 0.893	pass; GSM8K 0.892 → 0.892
Qwen3-30B-A3B BF16 TP4, online fp8_per_block (192 → 256)	no parameter ≠ fresh; cold start = this PR	no parameter ≠ fresh; 144 padded regions per rank recorded
Qwen3-30B-A3B-FP8 TP1 (no padding)	unchanged	unchanged
Qwen3-30B-A3B-NVFP4 TP4, TRT-LLM (no padding recorded)	960 side tensors extra / 384 moved	identical to main; cold start bitwise identical (8084 tensors)

Cold-start serving A/B (95bf429 vs 15cc0a8, same node, order main → PR → PR → main): GSM8K 5-shot greedy (1319 questions; gpt-oss via chat completions, reasoning effort low), vllm bench serve random 1024 in / 256 out at concurrency 64 (640 prompts), KV cache from the server log.

Model	GSM8K main → PR	Total throughput tok/s, main → PR	KV cache tokens, main → PR
gpt-oss-20b TP2, TRT-LLM MXFP4	0.861, 0.857 → 0.858, 0.864	70,881, 72,503 → 70,193, 73,090	10,970,211 → 10,971,579
Qwen3-30B-A3B BF16 TP4, TRT-LLM	0.894, 0.891 → 0.893, 0.892	48,643, 49,148 → 49,190, 49,289	6,374,288 → 6,374,288

The two main runs differ from each other as much as from the PR; the gpt-oss throughput gap follows run order (second run of each pair faster for both arms).

Unit tests (GB200, 8f20479 vs main side by side in one job): tests/kernels/moe/test_moe_weight_loading_padded.py, tests/kernels/moe/test_fp8_tp_loading.py, tests/quantization/test_online.py, tests/quantization/test_auto_gptq.py, tests/model_executor/model_loader/test_gpt_oss_weight_loading.py, tests/model_executor/test_utils.py, tests/model_executor/model_loader/test_reload.py.

	passed	failed
main	354	1 (test_kv_scale_reload)
this PR	359	1 (test_kv_scale_reload)

No test fails only with this PR.

New tests:

• test_moe_weight_loading_padded.py: loads into storage holding stale bytes; the loader leaves padding to post-processing, fill_padding writes each parameter's initial value (zeros, ones, torch.randn → 0, explicit padding_value) for every expert, a second fill after post-processing is a no-op, and a parameter whose layout changed is skipped.
• test_reload.py::test_padded_moe_reload_releases_each_layer: layerwise reload into rematerialized tensors holding other bytes; padding is 0 afterwards (fails without the fill in _layerwise_process).
• test_gpt_oss_weight_loading.py: the per-expert (streamed) loader records its padding.
• test_fp8_tp_loading.py: block-aligned FP8 TP padding is now written only by fill_padding (the test fails without it).
• test_direct_reload.py (g-padding): the toy model's loader writes a smaller checkpoint slice and records the rest with mark_padding; post-load processing re-lays out the whole tensor. Cold start and direct reload match a fresh cold start (fails without the fill in process_weights_after_loading).
• test_gpt_oss_weight_loading.py::test_quark_expert_padding_is_filled_before_post_processing: the Quark loader records its padding (per-expert and combined).
• test_moe_weight_loading_padded.py::test_initial_padding_value: a constant other than 1.0 (including all-0x01 bytes, as a freed bool buffer holds) gets 0.
• test_utils.py::test_padding_records_only_the_latest_load: a second load that splits a parameter differently is not overwritten by the first load's padding, and a parameter the second load wrote without padding is not filled.
• test_gpt_oss_weight_loading.py::test_streamed_mxfp4_expert_splits_like_the_checkpoint_loader: the per-expert loader splits the MXFP4 intermediate dim in whole blocks, as the full-checkpoint loader does, not by the padded per-rank size.

The last two fail on the previous head with the new tests copied in (slurm 42791) and pass on 624c71606 (slurm 42790): test_utils.py, test_gpt_oss_weight_loading.py, test_moe_weight_loading_padded.py, test_direct_reload.py, test_reload.py: 257 passed, 1 failed (test_kv_scale_reload, which also fails on main e0d00e4b8, slurm 42842).

FAILED test_utils.py::test_padding_records_only_the_latest_load
  assert layer.w.tolist() == [7] * 6 + [0] * 2   ->  At index 4 diff: 0.0 != 7
FAILED test_gpt_oss_weight_loading.py::test_streamed_mxfp4_expert_splits_like_the_checkpoint_loader
  rank 1 got 64 of its 128 rows (split by the padded 96 rows per rank)

Verifier command: run_bench.py --model-a <gpt-oss-20b> --update <perturbed copy> --backend nccl --layout fsdp=2 --server-gpus 0 --trainer-gpus 2,3 --warmup 0 --iters 1 --audit -- --max-model-len 4096, with the perturbed copy from validation/make_perturbed_checkpoint.py --variant 1.

Details

Root cause. Backends pad expert dims for kernel alignment, create_weights allocates the padding zero-filled, and loaders only write the checkpoint-backed slice. Post-processing then reads the whole padded tensor (TRT-LLM MXFP4 interleaves and shuffles over the padded size). A reload does not start from the allocation-time zeros: layerwise reload rematerializes with torch.empty_strided, in-place reload loads into the previous post-processed bytes, and the block-aligned FP8 TP path only refilled its padding because its loader happened to fill_ it.

Change.

• mark_padding(layer, name, region, value) records a padded region of layer.<name> (offset, shape, stride relative to the parameter) on the layer, which outlives the parameter tensors across reloads. A region of one expert is recorded once for all experts. Records hold only the latest load's regions: a parameter's first record after a fill starts its records over, and a fill drops parameters this load did not mark. Two loaders may split the same parameter differently (gpt-oss full-checkpoint vs per-expert), so earlier regions may cover real rows.
• fill_padding(module) writes the regions recorded since the last fill, and skips a parameter whose layout no longer matches the recorded one. It runs for every module before any post-processing at cold start, and per layer before post-processing in layerwise reload. A layer that online quantization post-processed during loading is therefore not filled again in kernel layout.
• RoutedExperts records instead of writing (_split_padding replaces _narrow_expert_data_for_padding), with the value recorded at construction. This covers hidden/intermediate padding, the gpt-oss combined, per-expert and Quark loaders (GptOssRoutedExperts._copy_to_expert keeps main's signature and finds its parameter from the storage, so other callers record too), and block-aligned FP8 TP, whose in-loader fill_ and pad_value plumbing are removed. The loader passes its parameter. The gpt-oss per-expert loader now splits the intermediate dim as the full-checkpoint loader does (whole MXFP4 blocks); it split by the padded per-rank size before, so the two loaders wrote different rows at TP >= 2 (on main too). recording is skipped on meta tensors (the layerwise dry run) and plain tensors, and a padded write into an nn.Parameter not registered on the layer warns once.
• The contract is documented on QuantizeMethodBase.process_weights_after_loading and fill_padding: code that runs post-processing itself must call fill_padding(layer) first.

Replaces #50523 and aoshen02#179.

• [Bugfix] Zero-fill rematerialized weights so reload preserves padding #50523 zero-fills every rematerialized tensor in materialize_meta_tensor. It fixes layerwise reload only, and only with zeros, so non-zero padding (unit scales) and in-place reload stay wrong.
• [MoE] Zero expert padding on every load, not only at allocation aoshen02/vllm#179 zeroes the padding inside the loader on every load, again only with zeros.
• This PR covers both reload paths, with each parameter's own initial value (e.g. unit scales), so neither is needed.

#62332 (refits into already-shuffled FlashInfer layouts) is orthogonal.

Limitations.

• Only MoE expert padding is recorded. Loaders that concatenate zeros onto the checkpoint tensor (DeepSeek-V4/V4.1 shared experts and drafters, Ling-3 shared experts, dots3 dense/shared MLP, GLM5-Next NoPE kv_a rows) are correct today, since the loader writes the zeros on every load; they need mark_padding once the load plan ([RFC]: Load plan — an engine-neutral API for sharded weight transfer #64418) stops running loaders.
• Code that runs post-processing itself must call fill_padding first (tests that call quant_method.process_weights_after_loading directly do so now). The sharded-RDT baked replay runs no loaders, so it fills nothing. It does bake padded unquantized MoE (weight_loader_numel excludes the padding) into uninitialized storage: intermediate padding is still correct there because unquantized post-processing re-zeroes it on the AITER and FlashInfer TRT-LLM backends (_zero_intermediate_padding, now commented as relied on), but hidden and bias padding of such a layer are not refilled on that path. The records persist on the layer, so a follow-up can replay them on that path.
• Padding is filled when loaders ran for that load. When transfer engines move to the load plan ([RFC]: Load plan — an engine-neutral API for sharded weight transfer #64418), RDT and nccl_m2n will no longer run loaders, so the plan needs an explicit fill_padding step before post-processing.
• A model built on meta (the weight-cache IPC loader) records padding value 0 for every parameter, including ones-initialized scales. Direct reload refuses pre-processed weights, but a layerwise reload of such a model would fill scale padding with 0 rather than the 1 its cold start wrote. Quant methods can set padding_value=1 on unit scales if that path needs it.
• A padded NVFP4 MoE was not available to test; Qwen3-30B-A3B-NVFP4 at TP4 records no padding.
• The evidence is tensor-level (bitwise against a fresh server), not accuracy. On every padded configuration tried, end-to-end output after a reload is also broken by the separate side-tensor issue above ([Model Loader] Keep side tensors on startup storage across layerwise reload #63090, [Model Loader] Write weight-derived tensors back in place on reload aoshen02/vllm#180), on main and with this PR alike: gpt-oss-20b GSM8K 0.33 → 0.00 after a layerwise reload on TRT-LLM MXFP4 (also on Triton MXFP4, which does not pad gpt-oss at TP1). That fix is needed too before reloaded gpt-oss matches a cold start end to end. Qwen3-30B-A3B BF16 already zeroes its intermediate padding in post-processing and keeps its accuracy on both sides.
• DeepSeek-V4.1 reload was not validated: layerwise reload runs out of memory on 4 GPUs on main as well.
• Out of scope: padding left unrecorded on purpose, because nothing reads it or every load rewrites it:
  • Vocab padding rows (VocabParallelEmbedding / ParallelLMHead): the loader zero-fills them, no token indexes them, and logits are sliced and masked.
  • DeepSeek-V4/V4.1 attention-sink padded heads: the loader writes -inf, and the outputs of padded heads are dropped.
  • Kimi-K3 KDA in_proj_qkvgfab padding rows and MiMo-V2 fused qkv pad rows: the forward splits them off.
  • Inkling gate padded expert rows and EP-alignment padding experts: never routed to.
  • Padding created by post-load processing itself (FlashInfer FP4/MXFP8 alignment, Marlin tiles, b12x, CUTLASS NVFP4, ModelOpt, online quantization): rebuilt from zeros on every run.
  • Stride-only padding (ROCm _maybe_pad_weight, gfx11 cache-aliasing padding): never read.
  • DBRX, Humming MoE (refused by direct reload), Cohere Compass, GGUF, bitsandbytes, HQQ, FP-Quant: not RL targets.

---

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 (Claude Code) was used to write the change and run the validation; the submitter reviewed every changed line.

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

[mergify]

Hi @hao-aaron, 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 #98145 for commit 55d1267e2283.

[aoshen02]

The batching makes sense to me: zeroing the padding in the loader per expert (aoshen02#179) cost about 1.5 s per update on DeepSeek-V4.1 TP4 (~100k small memsets per rank), and one strided fill per region removes that. A few things I'd like to see changed:

1. Pad values come from the parameter name. routed_experts.py:565 uses 1 when the name contains scale and the dtype is floating point, otherwise 0. That matches create_weights today by convention only: NVFP4 FP8 block scales get 1, uint8 MX scales get 0, and a backend that initializes differently or names a scale differently silently diverges from cold start. Cold-start identity was checked on gpt-oss, Qwen3 BF16/FP8 and DeepSeek-V4.1, not on a padded NVFP4 MoE. Could the value be declared where the parameter is created (e.g. a weight attr set in create_weights) instead of inferred?

2. A missed record is silent. When param is None or not found in self._parameters (:562), the padding is not recorded and stays stale with no error, which is exactly the bug this PR fixes. Skipping meta tensors is fine, but otherwise I'd assert, or at least warn once.

3. Every post-processing caller must call fill_padding first. That holds for the two loader paths, but sharded_rdt_engine.py:1396 already runs quant_method.process_weights_after_loading without it, and several tests had to be edited to add it. For post-processing that keeps shape, stride and dtype, the layout check is the only guard against filling processed bytes. Please state this contract where process_weights_after_loading is documented, so new callers do not miss it.

4. The block-aligned FP8 TP path writes its padding twice (:417 fill_ plus a record). No other path keeps an in-loader fallback, so I'd drop one of the two for consistency.

Posted by Claude Code on my behalf.

[hao-aaron]

/ci run

[github-actions]

✅ Triggered Buildkite CI #98858 for commit 1d3d59d316ae.

ℹ️ This PR is 1 commit behind upstream main. Merged upstream main into this branch as 1d3d59d316ae before starting CI. Pull the branch before pushing more changes.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @hao-aaron.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[hao-aaron]

/ci run

[github-actions]

✅ Triggered Buildkite CI #98975 for commit 292132564b26.

[hao-aaron]

/ci run

[github-actions]

✅ Triggered Buildkite CI #98981 for commit 877a17f46e73.

ℹ️ This PR is 1 commit behind upstream main. Merged upstream main into this branch as 877a17f46e73 before starting CI. Pull the branch before pushing more changes.

[hao-aaron]

/ci run

[github-actions]

✅ Triggered Buildkite CI #98984 for commit c777c75cefe8.

[hao-aaron]

/ci retry

[github-actions]

✅ Queued 4 failed job(s) for retry in Buildkite CI #98984.

[hao-aaron]

/ci retry

[github-actions]

✅ Queued 4 failed job(s) for retry in Buildkite CI #98984.

[hao-aaron]

/ci run

[github-actions]

✅ Triggered Buildkite CI #99031 for commit 26cc90fe7b00.

ℹ️ This PR is 11 commits behind upstream main. Merged upstream main into this branch as 26cc90fe7b00 before starting CI. Pull the branch before pushing more changes.

[hao-aaron]

Notes from a review pass, for reviewers. The fixes are in 624c71606 (see the description); what is left:

• Merge order with [Bugfix][Reload] Reload fixes, and a refit-equals-fresh harness #61287. [Bugfix][Reload] Reload fixes, and a refit-equals-fresh harness #61287 routes the Quark expert copies through the parameter's loader (gpt_oss_presliced_padded → self._copy_to_expert(data, loaded_weight)). _copy_to_expert now keeps main's signature and finds its parameter from the storage, so that call records padding, whichever PR merges first. The conflict at the Quark call sites resolves toward [Bugfix][Reload] Reload fixes, and a refit-equals-fresh harness #61287.
• Quark OCP_MX uint8 scales are created with torch.ones(uint8); _initial_padding_value returns 0 for integer dtypes, so their padding is now 0x00 where main had 0x01. Numerically harmless (the padded weights are zero), but not bit-identical to main on padded ROCm Quark MXFP4. Setting padding_value = 1 on them in quark_moe.py would restore it.
• Not tested: a layout-preserving in-place post-load step with value-1 padding (toy case g also passes with [Model Loader] Refill weight-derived tensors in place on every post-load rerun #64692's view.zero_() alone), and a RoutedExperts-level direct reload with EP local ids.
• Sharded-RDT replay fills no padding (in Limitations). Hidden and bias padding of unquantized MoE would hold uninitialized memory there; it is hard to reach today, since the only biased unquantized MoE (gpt-oss BF16) loads with param.copy_.
• Inkling writes its own AITER padding in its loaders ([Model] Inkling, Muse Spark: write checkpoint tensors through the parameter's weight loader #64690); moving it to mark_padding is [RFC]: Direct Weight Reload for RL Training #54477 F19.
• In RoutedExperts' gpt_oss_mxfp4 branch, hidden_dim=1, shard_dim=2 is right only because both trailing dims are narrowed (dim 1 is N for w13).

🤖 Generated with Claude Code