perf: Reduce MXFP8 refit weight-transfer overheads

[seonjinn]

PR decomposition

This PR retains the original integrated MXFP8 refit optimization branch and its review history.

• The batched MXFP8 MoE shuffle has been extracted into perf(vllm): batch MXFP8 MoE refit shuffles #3478 for independent review.
• BF16 trainer storage to MXFP8 vLLM with NCCL Reshard is implemented separately in feat(refit): support BF16 to MXFP8 NCCL reshard #3477.
• The extracted shuffle remains present in this historical integrated branch for now; it will be dropped or rebased onto perf(vllm): batch MXFP8 MoE refit shuffles #3478 before merge to avoid landing duplicate code.

What does this PR do?

This PR removes repeated transfer and receiver-side work from the MXFP8 rollout refit path.

• Prequantize eligible weights on the trainer and transfer E4M3 values plus E8M0 scales instead of full BF16 payloads.
• Reuse persistent CUDA-IPC staging buffers instead of allocating them on every refit.
• Batch the FlashInfer TRTLLM MoE W13/W2 value and scale layout transformations.
• Cache topology-dependent metadata and validated vLLM loader routes.
• Support slim post-refit offload and pinned-reference swap as opt-in runtime controls.
• Preserve the baseline paths behind configuration fallbacks.

Performance

Matched 20-step B200 runs, arithmetic mean over steady-state steps:

Model	Transfer + update	Total refit
Qwen3-30B-A3B	9.67 → 2.98 s (-69.2%)	17.2 → 13.1 s (-23.8%)
Qwen3-235B-A22B	18.40 → 7.26 s (-60.5%)	40.2 → 30.4 s (-24.4%)
Nemotron Nano 3	3.56 → 1.69 s (-52.6%)	13.2 → 11.6 s (-12.0%)

On Qwen3-235B-A22B, E2E step time decreased from 314.3 s to 305.0 s and throughput increased from 148.4 to 153.1 tokens/s/GPU, with identical mean rollout reward of 0.57.

Correctness

• MXFP8 prequantization is checked against the receiver quantization path.
• Batched MoE shuffles are compared bitwise with the per-expert reference path.
• Runtime options default to the existing behavior unless explicitly enabled.
• Matched training runs completed without reward regression in the reported validation windows.

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

Isolated-branch reproduction (follow-up promised in the PR body).

Setup: Qwen3-30B-A3B GRPO (OpenMathInstruct-2), 2 nodes x 8 B200, grpo-qwen3-30ba3b-4n4g-mxfp8-rollout.yaml, 20 steps, means over steps 3-20. Control = current main behavior (all flags off); treatment = this branch with refit_prequantize, persistent IPC buffers (4 GB), slim offload, cached loaders, pinned reference swap (batched shuffle is default-on). Same cluster, back-to-back runs.

steps 3-20 mean	flags off (main)	this PR	delta
transfer_and_update_weights (s)	9.67	2.98	-69%
prepare_for_generation total (s)	17.2	13.1	-4.1 s
generation / logprob / training (s)	49.8 / 84.6 / 91.0	51.8 / 82.0 / 90.6	unchanged (noise)
total step (s)	256.3	250.9	-2.1%
train reward	0.528	0.529	=

The saving is confined to the refit phase, exactly as claimed; all compute phases and reward are untouched. NRL_MXFP8_SHUFFLE_VERIFY=1 bit-exact assert passed on first refit.

[seonjinn]

/ok to test 740c825

[seonjinn]

/ok to test 3f60f5d

[seonjinn]

Isolated reproduction at 235B scale (follow-up to the 30B numbers above).

Setup: Qwen3-235B-A22B GRPO, 8 nodes x 8 B200 (64 GPUs), grpo-qwen3-235b-16n4g-mxfp8-rollout.yaml, vLLM TP 8, 20 steps, means over steps 3-20. Control = same clean branch with all refit flags off; treatment = refit_prequantize + persistent IPC buffers + slim offload + cached loaders + pinned reference swap (batched shuffle default-on).

steps 3-20 mean	flags off (main)	this PR	delta
transfer_and_update_weights (s)	21.8	7.4	-66%
prepare_for_generation total (s)	42.2	28.8	-13.4 s

The refit-transfer saving grows with model size (30B: 9.7 -> 3.0 s; 235B: 21.8 -> 7.4 s), because larger weights spend proportionally more of the refit in transport - exactly the phase this PR targets. Batched-shuffle bit-exact verify passed on first refit.

[seonjinn]

/ok to test 04af034

[seonjinn]

/ok to test e06f7a8

[seonjinn]

/ok to test 3ebdf9d

[seonjinn]

/ok to test b88aa2c

[seonjinn]

/ok to test 438e20d

[seonjinn]

/ok to test 96ba299

[seonjinn]

/ok to test 96ba299

[seonjinn]

/ok to test 2b80a85

[seonjinn]

/ok to test 0b6b5cc

[seonjinn]

/ok to test cd7732c

[seonjinn]

/ok to test 5ee7704

[seonjinn]

PR decomposition update: I restored the original integrated #3294 branch. The batched MXFP8 MoE shuffle has been extracted unchanged into #3478 for independent review. BF16 trainer storage to MXFP8 vLLM NCCL-Reshard support is tracked separately in #3477. Before merging #3294, I will drop or rebase the extracted shuffle portion onto #3478 so duplicate code does not land.

[seonjinn]

/ok to test cd7732c

[terrykong]

@youngeunkwon0405 could you help review first, i'll review afterwards

[youngeunkwon0405]

PR decomposition update: I restored the original integrated #3294 branch. The batched MXFP8 MoE shuffle has been extracted unchanged into #3478 for independent review. BF16 trainer storage to MXFP8 vLLM NCCL-Reshard support is tracked separately in #3477. Before merging #3294, I will drop or rebase the extracted shuffle portion onto #3478 so duplicate code does not land.

Looks like you have the PR decomposition plan. Then, this PR will be just for a reference? And the decomposed PR will be the actual target for merge?