[MoE] Fix flashinfer TRT-LLM BF16 expert weight reload on refit

[Kh4L]

Motivation

process_weights_after_loading rewrites BF16 MoE expert weights into the flashinfer TRT-LLM BlockMajorK layout whenever use_flashinfer_trtllm_moe is set, and that flag covers both flashinfer_trtllm and flashinfer_trtllm_routed. The inverse hook, however, was gated on is_flashinfer_trtllm_routed() alone. So with --moe-runner-backend flashinfer_trtllm the destination parameter stayed in block layout and the hot copy raised:

The size of tensor a (64) must match the size of tensor b (2048) at
non-singleton dimension 2

torch names the copy_ destination first, so the 64 is ours: block_k is 128 bytes and the conversion runs on a uint8 view, giving 128 / 2 bytes-per-bf16.

Nobody opts into this — flashinfer_trtllm is auto-selected on sm100 for BF16 MoE models when moe_runner_backend is left at "auto", so an RL weight refit on Blackwell hits it by default.

Fixes #27787

Modifications

1. Gate the inverse on the same flag that gates the transform (use_flashinfer_trtllm_moe), so the two cannot drift apart again.
2. The restore inverts the data, not just the shape. RL callers batch weights, so one refit is many update RPCs; with a shape-only restore an earlier bucket's re-derive would block expert slots that a later bucket has not written yet a second time.
3. New repack_weights_after_hot_update hook to re-derive the layout once the copies are done. update_weights_from_disk and the checkpoint-engine IPC path already re-run process_weights_after_loading; update_weights_from_tensor, update_weights_from_distributed and the bucketed variants call model.load_weights() directly and do not. Widening the gate on its own would therefore only trade the loud copy failure for a parameter left in canonical layout while the kernel reads BlockMajorK. The re-derive is called from a finally on all four update paths, so a mid-update exception cannot leave parameters canonical while the kernel expects block layout.

Files touched: layers/quantization/unquant.py, layers/quantization/base_config.py (no-op base hook), model_executor/model_runner_components/weight_updater.py, plus a new CPU unit test.

Note on the rebase onto current main: the only conflict was add/add with #33905 (XPU MoE row-stride padding), which inserts helpers at the same module-level position. Both blocks are kept and there is no interaction — the XPU padding applies under an XPU default device, the BlockMajorK rewrite under flashinfer on sm100, and #33905 only touches create_weights, which this PR does not.

Accuracy Tests

Reproduced against the real code path on CPU with Qwen3-30B-A3B geometry (hidden 2048, moe_intermediate 768, tp=2 -> 384 per partition):

before: postprocess (4,32,768,64) -> restore no-op   -> copy raises
after:  restore    -> (4,768,2048) -> copy ok -> repack -> (4,32,768,64)

test/registered/unit/layers/quantization/test_flashinfer_trtllm_bf16_moe_reload.py covers this end to end and is registered as a CPU test (11 cases): the block layout is applied on cold load and changes the shape; the restore runs for the non-routed backend and inverts the layout rather than reshaping it; a reload reproduces the cold-load layout exactly, including through the bucketed update path; the repack is a no-op when nothing was reverted and when the backend is inactive; non-bijective permutations and unexpected element counts are rejected; and the finally re-derive still runs when load_weights raises.

Speed Tests and Profiling

Not applicable to the forward path — this PR only changes the weight-update path. The re-derive runs once per update RPC, after the copies land; no kernel or forward-pass code is touched.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ✅ Run #31643837888
Latest PR Test (Extra): ❌ Run #31643837810

[nvpohanh]

@Kh4L could you fix the conflicts? thanks!

[nvpohanh]

@Kh4L could you rebase/merge with main again? github action was broken

[nvpohanh]

/tag-and-rerun-ci

[nvpohanh]

/rerun-failed-ci

[nvpohanh]

/rerun-failed-ci

[kpham-sgl]

@Kh4L can you help resolve the conflicts. Thanks!

[kaixih]

@Kh4L can you check if we are resolving the same issue in #25692? if yes, can you please check which one is better?