perf(vllm): batch MXFP8 MoE refit shuffles by seonjinn · Pull Request #3478 · NVIDIA-NeMo/RL

seonjinn · 2026-08-03T20:26:00Z

What does this PR do?

This PR extracts the MXFP8 MoE layout-transform optimization from #3294 into an independently reviewable change.

Replace the per-expert W13/W2 value and scale shuffle loop with four batched row gathers.
Cache the value-independent W13/W2 row permutations on each layer.
Keep batched gather outputs transient so CUDA tensor objects are not retained across vLLM sleep/wake boundaries.
Derive layout conversion from the selected MoE backend: FlashInfer TRTLLM always uses the batched path, the per-expert implementation remains a test oracle, and unsupported backends fail explicitly until they have a backend-specific conversion.

The optimization is transport-independent and does not modify weight-transfer metadata or precision negotiation. BF16-training/MXFP8-rollout support and the required vLLM 0.25.1 MXFP8 MoE kernel lifecycle are owned by #3477. The current branch temporarily retains the lifecycle block so it remains testable against main; it will be removed by rebasing after #3477 merges. The final #3478 diff will contain only the batched-shuffle optimization and its tests.

Why?

Each MoE layer has W13 and W2 value tensors plus their E8M0 scale tensors. The reference path launches the same four value-dependent row shuffles once per expert. Stacking experts preserves the exact row permutation while reducing Python dispatch and kernel-launch overhead.

The resulting layout is specific to Fp8MoeBackend.FLASHINFER_TRTLLM. Selecting another backend now fails before mutating weights rather than silently applying the FlashInfer layout.

Performance

Matched Qwen3-30B-A3B MXFP8, 8-node/32-GPU GB200, 20-step NCCL-Reshard A/B. Both arms used source 9706cdc, the same container, runtime, topology, rollout schedule, and instrumentation; only the shuffle implementation changed. Values are arithmetic means over Steps 3-9,11-19.

Metric	Per-expert reference	Batched	Change
Transfer + update (`n=15`)	5.899 s	0.894 s	-84.85%
E2E step time	208.136 s	204.397 s	-1.80%
E2E tokens/s/GPU	1228.204	1250.614	+1.82%

Step 11 lacks the transfer/update metric in both logs. The non-refit workload phases remained within 0.85%, and mean generation length was identical. Reference: job 5884885, W&B. Batched: job 5883555, W&B.

Correctness

CUDA tests compare the batched output bitwise with the per-expert reference for gated aligned shapes, gated shapes requiring scale padding, and non-gated MoE shapes. Tests also verify backend-derived dispatch, explicit rejection of unsupported backends, stable runtime parameter objects and storage across refits, and no caller-retained gather output buffers.

Signed-off-by: seonjinn <sna@nvidia.com>

copy-pr-bot · 2026-08-03T20:26:03Z

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 · 2026-08-03T20:27:29Z

/ok to test 5ee7704

seonjinn · 2026-08-03T22:04:11Z

/ok to test 5ee7704

Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn · 2026-08-05T00:04:08Z

/ok to test c795650

…-shuffle

copy-pr-bot · 2026-08-05T00:09:18Z

/ok to test 7812b5ead7c2f84ed6165ab0c322dd65e0ab15e9

@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 · 2026-08-05T00:09:31Z

/ok to test 7812b5e

Signed-off-by: seonjinn <sna@nvidia.com>

copy-pr-bot · 2026-08-05T00:14:54Z

/ok to test 3404e6f96c6120168850ae375c37d1d4eabf95c2

@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 · 2026-08-05T00:15:10Z

/ok to test 3404e6f

Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn · 2026-08-05T04:40:32Z

/ok to test 4ac5012

…-shuffle

seonjinn · 2026-08-05T05:24:00Z

/ok to test 3099836

Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn · 2026-08-05T07:48:30Z

/ok to test

copy-pr-bot · 2026-08-05T07:48:34Z

/ok to test

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

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

Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn · 2026-08-05T07:56:39Z

/ok to test

copy-pr-bot · 2026-08-05T07:56:42Z

/ok to test

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

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

seonjinn · 2026-08-05T07:59:13Z

/ok to test a74113d

Signed-off-by: seonjinn <sna@nvidia.com>

copy-pr-bot · 2026-08-05T08:27:07Z

/ok to test 6b5eeaa99771084bf6090516385300629767aa6a

@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 · 2026-08-05T08:27:16Z

/ok to test 6b5eeaa

Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn · 2026-08-05T08:45:54Z

/ok to test 52f30c7

seonjinn · 2026-08-05T12:34:30Z

Final validation for 52f30c7e8d99fbc7cb9138b2972546f4d7fed8ed:

Matched 20-step A/B

Qwen3-30B-A3B MXFP8 on 8 OCI-HSG GB200 nodes / 32 GPUs with NCCL Reshard. Both arms used source 9706cdc, the same immutable container, runtime, topology, rollout schedule, and instrumentation. Only refit_batched_moe_shuffle changed. Arithmetic means use Steps 3-9,11-19.

Metric	Per-expert reference	Batched	Change
Transfer + update (`n=15`)	5.899 s	0.894 s	-84.85%
E2E step time	208.136 s	204.397 s	-1.80%
E2E tokens/s/GPU	1228.204	1250.614	+1.82%

Step 11 lacks the transfer/update metric in both logs.

Reference: job 5884885, W&B
Batched: job 5883555, W&B

Both jobs recorded and synced all 20 steps. After training, both hit the same unrelated Ray Python-finalization assertion and exited 1:0; all requested comparison metrics were already complete.

Regression and CI evidence

Current CUDA test: job 5883915, 3 passed.
Parent implementation negative control: job 5884340, failed at the intended retained-output-buffer assertion.
Full CI run 30990382419: all 54 executed jobs passed; patch coverage is 100%.

youngeunkwon0405 · 2026-08-07T18:58:41Z

Should this behavior be derived from the selected backend rather than exposed as a user-facing knob? Both branches produce the FlashInfer TRTLLM layout, so setting refit_batched_moe_shuffle=false only selects the slower FI implementation; it does not make non-FlashInfer backends safe. Since the batched path is already verified bitwise against the reference, could we invoke it unconditionally when self.mxfp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM, keep the per-expert implementation only as a test oracle, and either dispatch the proper backend-specific refit conversion or fail loudly for other backends?

youngeunkwon0405 · 2026-08-07T22:34:19Z

For the other parts, looks reasonable to me.

Signed-off-by: seonjinn <sna@nvidia.com>

…e-split Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn · 2026-08-08T02:08:09Z

/ok to test 7436aa7

youngeunkwon0405

LGTM

seonjinn · 2026-08-11T01:09:30Z

@youngeunkwon0405 Thank you so much for your time and for reviewing this PR, Youngeun!

perf(vllm): batch MXFP8 MoE refit shuffles

5ee7704

Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn requested review from a team as code owners August 3, 2026 20:26

seonjinn mentioned this pull request Aug 3, 2026

perf: Reduce MXFP8 refit weight-transfer overheads #3294

Open

seonjinn added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label Aug 3, 2026

seonjinn requested a review from terrykong August 3, 2026 22:36

test(vllm): cover MXFP8 refit dispatch

c795650

Signed-off-by: seonjinn <sna@nvidia.com>

Merge remote-tracking branch 'origin/main' into sna/mxfp8-batched-moe…

7812b5e

…-shuffle

test(vllm): validate MXFP8 refit invariants

3404e6f

Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn added 5 commits August 4, 2026 21:39

test(vllm): require MXFP8 MoE kernel initialization

7125606

Signed-off-by: seonjinn <sna@nvidia.com>

fix(vllm): initialize MXFP8 MoE kernel after loading

c3ec212

Signed-off-by: seonjinn <sna@nvidia.com>

test(vllm): preserve MXFP8 kernel state across refits

494952d

Signed-off-by: seonjinn <sna@nvidia.com>

fix(vllm): preserve MXFP8 kernel across refits

17953ce

Signed-off-by: seonjinn <sna@nvidia.com>

test(vllm): isolate MXFP8 kernel lifecycle

4ac5012

Signed-off-by: seonjinn <sna@nvidia.com>

Merge remote-tracking branch 'origin/main' into sna/mxfp8-batched-moe…

3099836

…-shuffle

seonjinn mentioned this pull request Aug 5, 2026

feat(refit): support BF16 to MXFP8 NCCL reshard #3477

Merged

test(vllm): cover gated MXFP8 refit

2a2a35a

Signed-off-by: seonjinn <sna@nvidia.com>

test(vllm): compare gated shuffle storage

a74113d

Signed-off-by: seonjinn <sna@nvidia.com>

fix(vllm): release MXFP8 shuffle temporaries

6b5eeaa

Signed-off-by: seonjinn <sna@nvidia.com>

test(vllm): detect retained MXFP8 shuffle buffers

52f30c7

Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn requested a review from youngeunkwon0405 August 6, 2026 06:05

NVIDIA-NeMo deleted a comment from copy-pr-bot Bot Aug 7, 2026

seonjinn mentioned this pull request Aug 7, 2026

fix(vllm): support native BF16 FlashInfer TRTLLM refit #3545

Merged

seonjinn added 3 commits August 7, 2026 18:38

test(vllm): require backend-derived MXFP8 refit dispatch

dcf3b7b

Signed-off-by: seonjinn <sna@nvidia.com>

refactor(vllm): derive MXFP8 refit shuffle from backend

5f3a77c

Signed-off-by: seonjinn <sna@nvidia.com>

Merge remote-tracking branch 'upstream/main' into sna/pr3478-lifecycl…

7436aa7

…e-split Signed-off-by: seonjinn <sna@nvidia.com>

youngeunkwon0405 approved these changes Aug 10, 2026

View reviewed changes

terrykong merged commit d5fb8d0 into NVIDIA-NeMo:main Aug 11, 2026
83 of 84 checks passed

seonjinn mentioned this pull request Aug 24, 2026

perf(refit): batch MXFP8 expert prequantization #3804

Draft

Uh oh!

Conversation

seonjinn commented Aug 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Why?

Performance

Correctness

Uh oh!

copy-pr-bot Bot commented Aug 3, 2026

Uh oh!

seonjinn commented Aug 3, 2026

Uh oh!

seonjinn commented Aug 3, 2026

Uh oh!

seonjinn commented Aug 5, 2026

Uh oh!

copy-pr-bot Bot commented Aug 5, 2026

Uh oh!

seonjinn commented Aug 5, 2026

Uh oh!

copy-pr-bot Bot commented Aug 5, 2026

Uh oh!

seonjinn commented Aug 5, 2026

Uh oh!

seonjinn commented Aug 5, 2026

Uh oh!

seonjinn commented Aug 5, 2026

Uh oh!

seonjinn commented Aug 5, 2026

Uh oh!

copy-pr-bot Bot commented Aug 5, 2026

Uh oh!

seonjinn commented Aug 5, 2026

Uh oh!

copy-pr-bot Bot commented Aug 5, 2026

Uh oh!

seonjinn commented Aug 5, 2026

Uh oh!

copy-pr-bot Bot commented Aug 5, 2026

Uh oh!

seonjinn commented Aug 5, 2026

Uh oh!

seonjinn commented Aug 5, 2026

Uh oh!

seonjinn commented Aug 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Matched 20-step A/B

Regression and CI evidence

Uh oh!

youngeunkwon0405 commented Aug 7, 2026

Uh oh!

youngeunkwon0405 commented Aug 7, 2026

Uh oh!

seonjinn commented Aug 8, 2026

Uh oh!

youngeunkwon0405 left a comment

Choose a reason for hiding this comment

Uh oh!

seonjinn commented Aug 11, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

seonjinn commented Aug 3, 2026 •

edited

Loading

seonjinn commented Aug 5, 2026 •

edited

Loading