fix(vllm): support BF16 TRTLLM NCCL reshard refit

[seonjinn]

What this PR does

Makes weight refit work for BF16 MoE models running on the FlashInfer TRTLLM backend when using refit_transport=nccl_reshard.

Without this fix, a refit silently corrupts the MoE expert weights and the model generates garbage. With this fix, expert weights load correctly, and any configuration we cannot handle is rejected with a clear error instead of corrupting anything.

#3545 already fixed this for the ipc and collective transports. This PR extends the fix to nccl_reshard, which needs a different approach.

The problem in one picture

What the trainer sends:            What vLLM actually stores (TRTLLM):
┌──────────────────────┐           ┌──────────────────────────────┐
│ expert weights in    │  ──X──>   │ same weights, repacked into  │
│ checkpoint layout    │           │ a kernel-private layout      │
└──────────────────────┘           └──────────────────────────────┘

nccl_reshard streams every received tensor directly into vLLM's live parameter buffers. That is fast, and it is correct as long as the buffer stores the weight in the same layout as the checkpoint.

The FlashInfer TRTLLM MoE backend breaks that assumption: at model load time, vLLM repacks the grouped expert weights (gate/up/down projections) into a private layout its kernels need. Streaming checkpoint-layout bytes into a repacked buffer writes the right numbers into the wrong places — no error, just a corrupted model.

Why we can't just reuse the #3545 fix

#3545 fixes ipc/collective by reloading the whole model through vLLM's normal weight loader (which knows how to repack). But nccl_reshard receives most weights directly into live buffers. A whole-model reload would reset all buffers back to checkpoint layout — including the dense/attention buffers that direct receive depends on and that were never broken. The fix must be surgical: touch only the expert weights.

How the fix works

Split the weights into two groups and treat them differently:

Weight kind	Path	Why
Dense, attention, embeddings	Direct receive into live buffers (unchanged)	Their layout matches the checkpoint — always was correct
Grouped MoE expert weights	Staging tensor → vLLM's own weight loader	Only vLLM knows how to repack them

For the expert weights, concretely:

1. Receive into a staging tensor, not the live buffer. Each rank allocates a fresh BF16 tensor shaped as its expert-parallel slice of the global [num_experts, ...] weight.
2. Replay through vLLM's loader. After the transfer, each expert in the staging tensor is handed to vLLM's normal weight loader under its global name (experts.<expert_start + local_idx>.<proj>.weight). vLLM repacks it exactly like it does at model load.
3. Reload lifecycle covers only the MoE modules. Every other buffer stays live, so the direct-receive fast path keeps working for the rest of the model. One finalize() at the end rebuilds the kernel layout.

When we refuse to refit (on purpose)

Two layouts cannot be mapped to "one named load per expert", so we reject them loudly at setup instead of guessing:

• Expert weights sharded on a non-expert dim (e.g. a TP slice through the middle of one expert's matrix) → ValueError when the refit map is built.
• Non-linear expert placement (expert_placement_strategy != "linear", e.g. round-robin) → RuntimeError, because "global expert ID = expert_start + local_idx" would be wrong.

And if anything fails after the expert buffers have been invalidated mid-refit, the worker is marked permanently unusable rather than serving a half-updated model.

What is NOT affected

Everything else is untouched: other MoE backends, other transports, quantized (FP8/MXFP8) paths, dense models.

vLLM version compatibility

The backend detection touches three vLLM internals, verified byte-identical at v0.25.1 (our current pin) and v0.28.0: UnquantizedMoeBackend.FLASHINFER_TRTLLM, UnquantizedFusedMoEMethod.unquantized_backend, and ExpertMapManager.placement_strategy. If a future vLLM moves them, the guarded imports degrade to "backend not detected" and the placement check raises — there is no silent-corruption path.

How it was validated

• Unit tests (CPU): staging shape/dtype, global expert-ID naming, reload lifecycle order, and both rejected layouts.
• End to end: 1-step Qwen3-30B-A3B BF16 GRPO on 16 GB200 — 54.00 GiB FFN reshard completed, Policy generation refit completed successfully, moe_backend=flashinfer_trtllm.
• pre-commit passes.

Dependency

#3545 has merged and this branch is synced with latest main, so the diff here is only the nccl_reshard extension.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[seonjinn]

/ok to test 9d5a10a

[seonjinn]

/ok to test 2844542

[seonjinn]

Self-review completed against the latest main branch.

• Checked backend gating, expert-local shard selection, global expert IDs, staging dtype, native reload lifecycle, transport failure behavior, and unchanged dense/attention paths.
• Verified the vLLM native reload API through targeted unit tests and an end-to-end BF16 TRTLLM refit run.
• Found one missing negative test during review: tensor-sharded expert destinations were rejected in code but not covered. Added test_build_hf_to_local_param_map_rejects_trtllm_tensor_sharding.
• Validation: 52 targeted vLLM tests passed, 8 xferdtensor tests passed, and pre-commit passed on all changed files.

No unresolved correctness findings remain in the NCCL reshard extension. The PR remains draft because it includes and depends on #3545.

[seonjinn]

Latest-main end-to-end validation completed at 28445421d502a31cc31975affafbc629be0745a1.

• Hardware: 16 GB200 GPUs (4 nodes x 4 GPUs)
• Model: Qwen3-30B-A3B
• Configuration: BF16 training, BF16 rollout, moe_backend=flashinfer_trtllm, refit_transport=nccl_reshard
• Scope: one full async GRPO step
• Result: completed successfully with exit code 0
• Transfer split: 54.00 GiB through NCCL reshard and 2.87 GiB through the misc broadcast path (95.0% through reshard)
• Runtime confirmation: Policy generation refit completed successfully on all generation ranks

This verifies the implementation on the current main-based branch, including generation startup, refit, rollout collection, logprob/advantage processing, and the training step.

[seonjinn]

Mixed-scope validation update on 5b14cad6:

• Qwen and Nemotron-style module names were tested with (first N, last M) values (0,0), (5,0), (0,4), (1,1), (2,6), (3,5), (1,3), (7,3), (26,26), (30,30), (40,0), and (0,40).
• Boundary-layer QKVO and routed experts stay BF16; middle-layer QKVO and routed experts use MXFP8; gates and shared experts stay BF16.
• Invalid negative, non-integer, and out-of-range boundaries fail before vLLM patching.
• Focused scope result: 35 passed, 42 deselected. GitHub checks pass.

The combined mixed-refit stack also completed two colocated sync CUDA IPC runs with CUDA Graph on:

Model	BF16 boundary	Result	gen_kl_error mean	Run
Nemotron 3.5 Lightning 30B-A3B	first 1 / last 3	20/20 steps	0.004403	W&B
Qwen3.5 35B-A3B	first 3 / last 5	20/20 steps	0.001732	W&B

These end-to-end runs include the dependent sync-refit changes and are compatibility evidence, not PR #3659-only performance results.

Nano caveat: its forward gen_kl_error meets the existing MXFP8 functional threshold, but reverse policy_kl_error has large tail outliers. A matched first-1/last-3 sync MoE-only control also completed 20/20 steps. Compared with MoE-only, QKVO+MoE reduced E2E throughput by 4.67%, reduced generation throughput by 4.28%, doubled mean generation KL, and increased reverse-KL median from 35.86 to 180.89. The refit mechanism works, but Nano should keep QKVO in BF16 and use MoE-only MXFP8 until the remaining reverse-KL tail is understood.

[seonjinn]

/ok to test 96101df

[copy-pr-bot]

/ok to test 8eddc45aa060f27f96e4a89b370051bda08e7d34

@seonjinn, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

[seonjinn]

/ok to test 8eddc45

[seonjinn]

/ok to test f20709b

[seonjinn]

/ok to test eb7529f

[sharonyu-115]

/ok to test 014a305