fix(nvfp4): make process_weights_after_loading hot-reload-safe via alias-when-same-shape

[ch-wan]

Motivation

PR #25107 freed unused NVFP4 source scales (input_scale / weight_scale / weight_scale_2 and their MoE counterparts) inside process_weights_after_loading to reclaim ~15 GiB/rank on Kimi-K2.5 NVFP4 TP=4. That broke test_update_weights_from_disk_blackwell.py::TestServerUpdateWeightsFromDiskNVFP4:

• 1st invocation: derive alpha / input_scale_inv / weight_scale_interleaved, mutate layer.weight in place (TRTLLM shuffle_matrix_a permutes rows, CUTLASS pad_nvfp4_weight pads N/K), then del the source scales.
• /update_weights_from_disk → model.load_weights(...):
  • Refills layer.weight from safetensors with raw bytes (overwriting the previously-shuffled/padded representation).
  • Can't refill the deleted source-scale slots.
• 2nd process_weights_after_loading call: layer.input_scale.max() → AttributeError. Even with a guard, the new raw layer.weight mismatches the cached shuffled-layout weight_scale_interleaved / alpha, so the GEMM produces garbage tokens.

Neither a full revert (loses ~15 GiB/rank) nor a _weights_postprocessed skip-on-second-pass flag (silently broken outputs after hot reload) is acceptable.

Approach: alias-when-same-shape

Drop the dels. Let process_weights_after_loading run on every call. Introduce alias_or_bind_derived_param in python/sglang/srt/layers/utils/common.py: when the derived tensor has the same shape & dtype as the source Parameter, write the derived bytes into the source's storage in place and register the derived attribute name as an alias of the source Parameter. Two attribute names then share one underlying buffer.

For each call:

• apply() reads via the derived name (weight_scale_interleaved, *_blockscale_swizzled) and sees the post-processed bytes.
• update_weights_from_disk refills the source slot (still registered under its original name) — the next process_weights_after_loading overwrites the same buffer with the re-derived bytes.
• Peak GPU memory is the source size, not source + derived.

When shapes diverge (genuine padding required for non-aligned dims), fall back to allocating a separate Parameter via copy_or_rebind_param — no memory savings in that case, but correctness preserved.

For Kimi-K2.5 NVFP4 dims (H=7168, moe_intermediate=2048, group_size=16, etc.), the relevant pairs are all byte-for-byte same-size:

• Linear: weight_scale ↔ weight_scale_interleaved (CUTLASS + TRTLLM paths)
• MoE: w13_weight_scale ↔ w13_blockscale_swizzled, w2_weight_scale ↔ w2_blockscale_swizzled

layer.weight is left untouched semantically (copy_or_rebind_param continues to do in-place copy when shape matches), so the in-place shuffle/pad runs on every hot reload too.

The TRTLLM-branch MoE additionally aliases the unused *_blockscale_swizzled placeholders to the in-place-shuffled w*_weight_scale after align_fp4_moe_weights_for_flashinfer_trtllm, freeing the placeholders allocated in create_weights.

What stayed scalar (not aliased)

alpha and input_scale_inv (Linear) and g{1,2}_alphas / w{13,2}_input_scale_quant (MoE) remain plain copy_or_rebind_param Parameters. An earlier attempt to alias them into the [N_partitions] source slot broke fused-QKV linears: a downstream call tries to view the scale as [1] and hits RuntimeError: shape '[1]' is invalid for input of size 3. Per-linear savings would have been negligible anyway (~bytes); reverted in commit bf0b436.

Modifications

python/sglang/srt/layers/utils/common.py: add alias_or_bind_derived_param(module, source_name, derived_name, derived_value).

python/sglang/srt/layers/quantization/modelopt_quant.py:

• ModelOptFp4LinearMethod.process_weights_after_loading: replace copy_or_rebind_param(layer, 'weight_scale_interleaved', ...) + del layer.weight_scale with alias_or_bind_derived_param(layer, 'weight_scale', 'weight_scale_interleaved', ...) on both the TRTLLM and CUTLASS branches. Remove the del layer.input_scale, layer.weight_scale_2 line.
• ModelOptNvFp4FusedMoEMethod.process_weights_after_loading:
  • Remove del layer.w13_input_scale, layer.w2_input_scale.
  • TRTLLM branch: replace del layer.w13_blockscale_swizzled, layer.w2_blockscale_swizzled with two lines aliasing the swizzled-name slots to w13_weight_scale / w2_weight_scale.
  • non-TRTLLM (CUTLASS / CuteDSL) branch: replace the two copy_or_rebind_param(layer, '*_blockscale_swizzled', ...) + trailing del layer.w13_weight_scale, layer.w2_weight_scale with alias_or_bind_derived_param(layer, 'w13_weight_scale', 'w13_blockscale_swizzled', ...) / same for w2.

ModelOptFp8LinearMethod and ModelOptFp8MoEMethod are unchanged.

Accuracy / hot-reload tests

test/registered/rl/test_update_weights_from_disk_blackwell.py::TestServerUpdateWeightsFromDiskNVFP4:

• Reloads nvidia/Qwen3-30B-A3B-NVFP4 from disk twice with TP=4 + FlashInfer-TRTLLM FP4 GEMM + FlashInfer-TRTLLM-routed MoE.
• Asserts decode logprobs unchanged within atol=1e-4.

Verified on a 4× GB300 node:

Ran 1 test in 290.006s

OK

Speed / memory

No latency impact — the process_weights_after_loading path is one-shot at startup (and once per update_weights_from_disk call). The same ~15 GiB/rank peak-memory win from PR #25107 is preserved on Kimi-K2.5 NVFP4 TP=4 (all relevant pairs hit the alias path; no separate buffer is allocated for the derived view).

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests. (Reusing existing test_update_weights_from_disk_blackwell.py.)
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed. (Same-checkpoint hot-reload logprob match; no perf-path change.)
• Follow the SGLang code style guidance.

🤖 Generated with Claude Code

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[ch-wan]

/tag-and-rerun-ci

[alisonshao]

/rerun-test test_update_weights_from_disk_blackwell.py

[github-actions]

🚀 4-gpu-b200 (1 test): ❌ View workflow run

cd test/ && python3 registered/rl/test_update_weights_from_disk_blackwell.py