Repository navigation
Conversation
…ss_weights_after_loading process_weights_after_loading no longer unregisters a loaded param or leaves it without its data, except where keeping the data would cost real memory: - Renames go through the new move_parameter helper, which registers the same param object under the runtime name too, so the loader writes through the checkpoint name and kernels read the runtime name of one storage. - Loaded params no kernel reads (consumed per-tensor scales, input scales of weight-only Marlin/Humming MoE backends) keep their data and are marked retired; the MoE quant configs get None for unused input scales instead of the layer attribute being set to None. - Large released sources (weights repacked into a different layout, KV-cache scale placeholders) stay registered with empty storage. kernel_parameters/get_live_param give the kernel view (live params, each object once under its runtime name) to the Humming conversion, and the Triton MXFP4 MoE wrapper moves to a plain <name>_triton attribute. Signed-off-by: aoshen02 <aoshen@inferact.ai> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- Marlin/CPU/zentorch/conch: a loaded zero-point param that the config does not use (symmetric GPTQ qzeros) is emptied in place again, as on main, instead of keeping its data (~1/group_size of the int4 weight). - Revert the unreachable deepseek_v4 XPU and dots3 vision changes. - Tests that asserted deleted or renamed params now assert the param is still registered (renamed ones alias the runtime name); drop the blanket non-empty check, since Triton MXFP4 release empties by design. Signed-off-by: aoshen02 <aoshen@inferact.ai> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
replace_parameter now copies a new tensor that is smaller than the param into the param's own storage (a prefix view) instead of rebinding .data to a new, smaller allocation, so a loaded param never ends up with fewer bytes than its checkpoint tensor (e.g. ModelOpt NVFP4 MoE w13_weight_scale_2 and w13_input_scale [E,2] -> [E], merged qkv per-tensor scales [3] -> scalar). Online quantization (uses_meta_device) still frees its high-precision source. The mixed-precision CUTLASS FP8 scale conversion and the K3 gfx942 MXFP4 repack now go through replace_parameter instead of rebinding or emptying .data directly. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: aoshen02 <aoshen@inferact.ai>
The Triton MXFP4 MoE path emptied w13/w2_weight and their scales after swizzling them into triton_kernels tensors. The params now hold the swizzled storage and the triton tensors point at it, so the loaded params keep their bytes without a second copy. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: aoshen02 <aoshen@inferact.ai>
A converted tensor under a new name is now registered as the leftover loaded param it came from (same object under both names, as move_parameter does) when it lives in that param's storage, or takes the place of a large leftover. This removes the second Parameter object that block-FP8 MoE w13/w2_weight_scale had over w13/w2_weight_scale_inv, and the emptied CT WNA16 *_weight_packed params. Other leftovers keep their data, retired. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: aoshen02 <aoshen@inferact.ai>
- Unused zero points (Marlin, Conch, CPU, Zentorch), dynamic 4-bit scales and INC ARK qzeros/scales keep their data, retired. - Zentorch WOQ/DA8W4, the DeepSeek-V4 CPU wo_a pack, CPU AWQ repack, CT W4A8 int8 MoE and CT WNA8O8 int weights put the repacked tensor into the loaded param (the same object is registered under the packed name where one exists). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: aoshen02 <aoshen@inferact.ai>
…put scales Humming block-FP8 MoE registered the runtime w13/w2_weight_scale as a second Parameter over the storage of the loaded w13/w2_weight_scale_inv, and the Humming NVFP4 linear did the same for input_scale_2 over input_global_scale. Both are now one object registered under both names (move_parameter). Signed-off-by: aoshen02 <aoshen@inferact.ai> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Aaron Hao <aaronhao@inferact.ai>
A result PWAL allocates during a direct reload is copied into the parameter's storage at the offset, shape and stride it had after cold-start PWAL, so results larger than the checkpoint tensor (e.g. DeepGEMM UE8M0 block scales) or re-laid-out at equal size (FlashInfer TRT-LLM FP8 MoE shuffle) stay at the address CUDA graphs captured. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Aaron Hao <aaronhao@inferact.ai>
MoE loaders now record the padding they narrow away (`mark_padding`) instead of writing it, and `fill_padding` writes it before every post-load processing pass: cold start, direct and layerwise reload, and the sharded-RDT replay. A reload into storage that holds the previous post-processed bytes, or a transfer that writes checkpoint bytes without running the loaders, then still feeds post-processing the same padding as a cold start. Covers hidden/intermediate padding of expert weights, scales and biases, the gpt-oss combined-expert path and block-aligned FP8 TP (unit scales). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Aaron Hao <aaronhao@inferact.ai>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
Some MoE backends round the expert hidden or intermediate size up, or pad a TP shard to whole weight blocks. Post-load processing (PWAL) reads the whole padded tensor: TRT-LLM MXFP4, for example, interleaves and shuffles over the padded size.
On a direct reload, the loaders write into storage that still holds the previous update's post-processed bytes. The padding is not checkpoint data, so nothing rewrites it, and PWAL processes stale kernel-layout bytes. Transfer engines that write checkpoint bytes without running the loaders (the load-plan direction for M2N and RDT) have the same problem even if the loaders were changed to zero the padding.
aoshen02#179 fixes the loader path by zeroing in the loader. This PR moves the fill to PWAL, so it also holds when loaders do not run.
Change
mark_padding(param, region, value): a loader records a padded region (offset, shape, stride relative to the parameter, and the fill value) instead of writing it. A region of one expert is recorded once for all experts.fill_padding(module): writes the recorded regions in checkpoint layout. It runs before every PWAL pass: cold start and direct finish (model_loader/utils.py), layerwise reload, and the sharded-RDT replay.RoutedExpertsrecords hidden/intermediate padding of weights, scales and biases, the gpt-oss combined-expert path, and block-aligned FP8 TP (unit scales, which the loader used tofill_).Behaviour change: hidden/intermediate padding used to keep whatever
create_weightsallocated (possibly uninitialized); it is now zeroed before every PWAL, cold start included.Tests
On GB200 (Slurm):
tests/kernels/moe/test_moe_weight_loading_padded.py: 87 passed. NewTestPaddingFilledByPostProcessingloads into dirty storage and checks that the loader leaves padding to PWAL andfill_paddingwrites it for every expert, including unit scales.tests/kernels/moe/test_fp8_tp_loading.py: 30 passed (FlashInfer cases included). The padding-leak test now callsfill_paddingthe way the PWAL drivers do.tests/model_executor/model_loader/test_direct_reload.py: 4 passed.tests/model_executor/model_loader/test_reload.py: the same 23 failures on this branch and on its base commit on one node; all are 401s fetching debug models that are not in the shared cache. Nothing fails only on this branch.End to end, gpt-oss-20b (MXFP4 experts padded 2880 → 3072,
FLASHINFER_TRTLLM_MXFP4_BF16), TP2,VLLM_WEIGHT_RELOAD_MODE=direct, two reloads from disk to a perturbed checkpoint, every parameter, buffer and off-registry tensor compared with a fresh start of the same tree:w13/w2bias and weight scale, every layer)The 144 moved tensors are
gemm1_alpha/beta/clamp_limiton the TRT-LLM MXFP4 experts object, rebuilt by every PWAL. That is the kernel-constant issue aoshen02#180 (keep_constants) covers, not padding; gpt-oss generation matches a fresh start only with both.AI assistance (Claude Code) was used to write this change and run the tests.
🤖 Generated with Claude Code