[NVIDIA] Fix BF16 FlashInfer TRTLLM MoE RL weight update

[kaixih]

Summary

Fix BF16 FlashInfer TRTLLM MoE hot weight update.

With this change:

Pure inference is unchanged:

load_weights -> process_weights_after_loading -> packed params -> infer

RL rollout update changes from the broken flow:

packed params -> load canonical update weights into packed layout -> ERROR

to:
packed params -> restore canonical shape -> load canonical update weights -> process_weights_after_loading -> packed params

Changes

• Track canonical/packed layout metadata for BF16 FlashInfer TRTLLM MoE weights.
• Restore packed params to canonical shape before weight loading.
• Re-pack canonical weights in process_weights_after_loading().
• Handle fused w13 loading when destination halves have padding.
• Add a regression test for hot update after FlashInfer packing.

Test

python3 -m pytest -q test/registered/moe/test_flashinfer_trtllm_bf16_hot_update.py -s

<!-- pr-states:start -->
---
### CI States

Latest PR Test (Base): <!-- slot:pr-test:start -->:x: [Run #30841029629](https://github.com/sgl-project/sglang/actions/runs/30841029629)<!-- slot:pr-test:end -->
Latest PR Test (Extra): <!-- slot:pr-test-extra:start -->:x: [Run #30841030518](https://github.com/sgl-project/sglang/actions/runs/30841030518)<!-- slot:pr-test-extra:end -->
<!-- pr-states:end -->

[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!

[kaixih]

During the e2e testing, we find a more fundamental issue of the flashinfer_trtllm path. This issue actually affects the inference accuracy as well.

Basically, the destination param shard will be padded to 128 rows for the flashinfer_trtllm kenrel, while the checkpoint tensor is still unpadded. Previously, the loader used the padded destination shard size to slice the source checkpoint tensor, so higher TP ranks loaded the wrong rows.

Example:

• intermediate_size = 768
• TP = 8
• real source shard size = 768 / 8 = 96 rows
• FlashInfer destination shard size = round_up(96, 128) = 128 rows
• rank 3 should load source rows [288:384)
• old code used 128 * 3 and loaded source rows [384:512)

The new commit fix it. And in a pure SGLang fresh-load GSM8K-200 run with Qwen3-30B-A3B on B200

before padded TP fix: accuracy 0.775, invalid 0.000
after  padded TP fix: accuracy 0.930, invalid 0.000

[zianglih]

I cannot reproduce the inference accuracy drop with Qwen/Qwen3-30B-A3B-Instruct-2507 + tp8:

root@B200-141:/sgl-workspace/sglang# git log -1
commit 74c6294ba9afa2541468523e6e4830ea14f279f2 (HEAD -> main, origin/main, origin/HEAD)
Author: Zheng Wengang <zwg0606@gmail.com>
Date:   Thu May 21 11:07:21 2026 +0800

    [BugFix][EPD]Fix Qwen3VLMoe encoder-only AttributeError (#25759)
python -m sglang.launch_server --kv-cache-dtype bf16 --model Qwen/Qwen3-30B-A3B-Instruct-2507 --fp8-gemm-backend flashinfer_trtllm --moe-runner-backend flashinfer_trtllm --tp 8
python3 benchmark/gsm8k/bench_sglang.py --num-shots 8 --num-questions 1209 --parallel 1209 --platinum
Accuracy: 0.964
Invalid: 0.000
Latency: 8.201 s
Output throughput: 20902.976 token/s
Accuracy: 0.969
Invalid: 0.000
Latency: 10.292 s
Output throughput: 16598.841 token/s
Accuracy: 0.966
Invalid: 0.000
Latency: 9.686 s
Output throughput: 17633.333 token/s

Could you share the command you are using?

[kaixih]

So the clarification is: this is not an upstream-main issue. It is a PR-branch-without-7060cfa repro, and keeping 7060cfa makes the model accuracy normal again.

[kaixih]

Another data points:

Tried RL training with this fix. Compared to the triton backend (baseline), the accuracy matches and perf is >10% better:

Note: For the flashinfer_trtllm, we ran the test for 3 times with no --use-fault-tolerance so 3 lines of flashinfer above. Using --use-fault-tolerance may cause illegal memory access which might not be related to this PR and still being worked on.

@yueming-yuan to review.

[kaixih]

Finally, I figured out the issue. The previous IMA was not related to this PR. It was caused by my local experiment setup: I did not patch SGLang correctly inside the Miles container.

After fixing the setup, training runs cleanly. I tested 8xB200 Qwen3-30B RL training with this script in two configurations:

1. two rollout engines, TP4 each

2. one rollout engine, TP8

Both configurations run without issues. The raw reward curves match the Triton baseline, and FlashInfer shows about a 10% improvement in rollout_time.

@yueming-yuan can you take another look?

[kaixih]

cc. @nvpohanh

[kaixih]

I’m now pretty confident that the 2xTP4 setup has no issue. I ran another stress test by repeatedly launching the training job, and every run successfully passed step 40.

So the remaining issue appears to be specific to 1xTP8: for some mysterious reason, that configuration still triggers the IMA.

[b8zhong]

@kaixih Do we close this in favour of #33743