[YigongQin]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

This PR adds one opt-in knob, moe_training_matches_inference, meant to be enabled only by RL training. It is off by default, and generation's forward is unchanged with it on or off.

With the knob on, the training forward of an inference-optimized MoE layer runs the same FlashInfer routed kernel that generation runs. A token's training log-prob then equals its generation log-prob bitwise, for FlashInfer routed MXFP8 and per-token NVFP4 MoE. The backward is unchanged: it comes from the Transformer Engine (TE) recompute pass. The requirement is that this adds no forward–backward inconsistency beyond what the default TE recipe already has; see Core metric.

Numbers below are for one MoE layer with Nano-v3 shapes, GB300 × 4, EP 4, eager; see Performance. The value pass runs through training's own router module and all-to-all dispatcher, so training needs none of generation's buffers.

Training keeps one derived copy of the expert weights for FlashInfer: BF16 plus NVFP4 for per-token NVFP4, and MCore MXFP8 plus FlashInfer's routed layout for MXFP8.

Per-token NVFP4: faster and smaller, no trade-off

Against TE's per-token NVFP4 recipe, the same quantization scheme:

• Bitwise train–gen parity for every token's log-prob.
• Log-prob forward 3.4–4.4× faster.
• Train step 2.3–3.1× faster than TE NVFP4 + recompute, and 1.05–1.21× faster than BF16 TE + recompute.
• Less memory: lower peak at every size (log-prob 165 vs 262 MiB at 1024 tokens per rank, 2633 vs 3580 MiB at 16384), and about 180 MiB less static memory per layer, because its cached kernel weights (about 180 MB per layer per rank) are smaller than TE NVFP4's cached quantized weights (about 350 MiB).
• Same gradient: the backward is TE's, and the forward is within 0.4% of TE's distance from BF16 on every seed and input (Core metric).

Checked by bitwise parity with generation (including after a refit), kernel batch invariance, bit-exact replay under contention, and the core metric; see What is tested.

MXFP8: faster, with extra memory during each call

Against TE MXFP8:

• Bitwise train–gen parity for every token's log-prob.
• Log-prob forward 1.4–1.7× faster.
• Train step 1.11–1.16× faster than TE MXFP8 + recompute, with the same peak, plus a kernel-weight rebuild of about 36 ms per MoE layer once per optimizer step (see Performance).
• Same gradient: the backward is TE's, and the forward is closer to BF16 than TE's on every seed and input.

Checked by the same parity and core-metric tests; the parity tests also cover a token's output across launch sizes.

The drawback is memory held between optimizer steps:

• Cached kernel weights: about 330 MB per MoE layer per rank, about 7.6 GB for 23 MoE layers. TE MXFP8 needs no extra copy, because its fp8_param weights already are its kernel's input. training_matches_inference.release_inference_value_pass_weights(model) frees it, for example before colocated generation.
• Small log-prob passes: each launch keeps at least the rows for the kernel's largest tile, so log-prob peak is 41 MiB higher than TE's at 1024 tokens per rank; it is lower from 2048.
• If memory is tight, generation can use the torch MXFP8 backend (inference_grouped_gemm_backend: torch), which already matches TE training bitwise under batch_invariant_mode, without this flag or the cache.

Usage

On top of generation's FlashInfer settings (inference_optimized, batch_invariant_mode, and per-token NVFP4 or MXFP8 with fp8_param), the RL trainer adds:

moe_training_matches_inference: true
moe_token_dispatcher_type: alltoall
recompute_granularity: selective
recompute_modules: [moe]

The value pass also runs under model.eval(), so the RL log-prob forward takes it too.

Training-side memory

• Activations: the value pass saves none. The cost is one extra MoE forward per layer.
• Kernel weights (sizes above): built on the first value pass after a weight change and held until the next optimizer step, state_dict() call, checkpoint load, in-place write to the parameters, or training_matches_inference.release_inference_value_pass_weights(model). The log-prob pass and the training microbatches of one step share a single build. Against TE without recompute, the MoE activations recompute frees outweigh the MXFP8 cache from about 4096 tokens per rank per microbatch.
• No generation buffers: the value pass uses training's all-to-all dispatcher, so it works the same with generation colocated, on other GPUs, or on another cluster. Training must use generation's EP size.

Background

Under batch_invariant_mode, the torch and vLLM inference MoE kernels already match TE training bitwise. FlashInfer, the production path for MXFP8 rollouts, cannot: #7302 reports the MXFP8 mismatch, and per-token NVFP4 from #8262 is batch invariant but does not match training either. FlashInfer fuses the activation and the FC2-input quantization into the FC1 epilogue, applies routing probabilities after FC2, and accumulates in its own order, so the training path cannot be aligned to it by hand. In RL, the gap shows up as a bias in the importance ratio.

FlashInfer owns only the per-rank expert compute and the sum of each token's routes on a rank (its finalize); the router, dispatch and cross-rank combine are Megatron's. Under batch_invariant_mode, training's router already picks generation's experts and training's all-to-all combine already adds the EP ranks as generation's does. The training forward therefore keeps training's router and dispatcher, and reproduces only FlashInfer's expert compute and its per-rank sum.

Scope

The forward and the backward come from different libraries, so the gradient is not the derivative of the forward. That only matters for recipes whose backward was the exact derivative of their forward to begin with, and those are out of scope:

• Consistent (out of scope): the backward reuses the forward's dequantized operands in high-precision GEMMs.
• Already inconsistent (in scope): TE MXFP8 quantizes separate row-wise and column-wise copies for the forward and backward, and quantizes the output gradient too, so its backward is a quantized estimate of the BF16 layer's gradient. The NVFP4 recompute runs in BF16, so it also computes the BF16 layer's gradient.

For the in-scope recipes, the backward stays exactly as it was. The forward moves from one quantized implementation to another that differs only in quantization details, accumulation order and where the routing probabilities are applied. FlashInfer's routed MoE has no backward, so an all-FlashInfer path is not an option.

Core metric: no extra forward–backward inconsistency

Claim. Using generation's forward does not make the gradient fit the loss worse than the default TE recipe already does.

Quantization has zero derivative almost everywhere, so a quantized backward computes the gradient of the BF16 layer, not of its quantized forward:

Recipe	Forward that produces the loss	Gradient the backward computes	Backward GEMMs
MXFP8, default TE	TE MXFP8	The BF16 layer's	MXFP8
MXFP8, this PR	FlashInfer MXFP8	The BF16 layer's (same TE backward)	MXFP8
NVFP4, default TE (per-token, backward_override='high_precision')	TE per-token NVFP4	The BF16 layer's	BF16
NVFP4, this PR	FlashInfer per-token NVFP4	The BF16 layer's (same TE backward)	BF16

The inconsistency is the gap between the quantized forward and the BF16 layer. Within each recipe the backward is identical, so the metric is each forward's relative distance from BF16 on the MoE output, with the same weights and tokens:

$$d_{\mathrm{recipe}} = \frac{\lVert y_{\mathrm{recipe}} - y_{\mathrm{BF16}}\rVert}{\lVert y_{\mathrm{BF16}}\rVert}$$

Pass condition, for every recipe, seed and input type:

$$d_{\mathrm{FI}} \le (1 + \alpha)\, d_{\mathrm{TE}}, \qquad \alpha = 1\%$$

Both forwards use the same quantization scheme, so either can land slightly closer to BF16; 1% is small next to the roughly 20% both are from BF16.

How it is measured: TestForwardBackwardMismatch in tests/unit_tests/inference/test_moe_training_matches_inference.py.

• Three layers share one set of weights: this PR's value pass, the default TE recipe, and TE BF16. For MXFP8, the BF16 layer gets the dequantized MXFP8 weights.
• The NVFP4 baseline is TE's per-token recipe, NVFP4BlockScaling(row_scaled_activation=True, disable_2d_quantization=True, disable_rht=True, disable_stochastic_rounding=True, backward_override='high_precision'), the same scheme as FlashInfer's, so the comparison measures the kernel rather than the scheme.
• Inputs are 512 tokens per rank with seeds 11–13, each as Gaussian tokens and as tokens with 5% scaled by 30. Only forwards run, so the comparison is deterministic.

uv run python -m torch.distributed.run --nproc-per-node 4 -m pytest -q -s \
  tests/unit_tests/inference/test_moe_training_matches_inference.py -k TestForwardBackwardMismatch

Results (EP = 4, GB200). Each cell is the range over seeds 11–13 (lower is better). The last number column is the largest amount by which this PR exceeds TE on any seed; negative means this PR is closer to BF16 on every seed.

Recipe	Input	This PR ($d_{\mathrm{FI}}$)	Default TE ($d_{\mathrm{TE}}$)	Largest excess over TE	Pass ($\le 1\%$)
MXFP8, squared ReLU	Gaussian	0.0413–0.0414	0.0415–0.0416	−0.17%	yes
MXFP8, squared ReLU	5% outliers	0.0400–0.0409	0.0405–0.0421	−1.2%	yes
NVFP4, squared ReLU	Gaussian	0.1957–0.1963	0.1958–0.1963	0.00%	yes
NVFP4, squared ReLU	5% outliers	0.1985–0.2021	0.1984–0.2027	+0.40%	yes
NVFP4, SwiGLU	Gaussian	0.2307–0.2325	0.2309–0.2325	+0.04%	yes
NVFP4, SwiGLU	5% outliers	0.2224–0.2315	0.2224–0.2321	+0.40%	yes

The test also prints the worst single token's distance, which is not part of the pass condition:

• NVFP4: this PR's worst token is closer to BF16 in 9 of 12 runs; the largest excess is 0.363 against 0.360.
• MXFP8: closer in 2 of 6 runs; the largest excess is 0.080 against 0.071. The average is still lower on every run, so the error is spread differently over tokens but is not larger overall.

Limits. This compares one layer's forward, not a training curve. A finite-difference gradient check does not work here, because the loss of a quantized forward is piecewise constant in the weights.

Performance

One MoE layer with Nano-v3 shapes (hidden 2688, 128 experts, top-6, squared ReLU, shared expert), GB300 × 4, EP 4, batch-invariant mode, eager, mean of 20 iterations. The baseline is TE in the same quantization: TE MXFP8, or TE's per-token NVFP4 recipe (FlashInfer's scheme, run outside batch-invariant mode). Peak is the extra memory allocated during the pass, on top of resident memory; the cached kernel weights are resident, so they appear under static memory, not peak. All numbers come from one run.

Train (forward + backward). The flag requires moe recompute, so the baseline is TE + recompute. The tables time microbatches after the first, where neither side converts weights.

Recipe	Tokens/rank	TE + recompute ms	This PR ms	Speedup	TE + recompute peak MiB	This PR peak MiB
MXFP8	1024	31.43	27.30	1.15×	733	732
MXFP8	4096	32.22	28.50	1.13×	1276	1275
MXFP8	16384	35.68	31.09	1.15×	4630	4632
NVFP4	1024	59.47	20.53	2.90×	750	721
NVFP4	4096	61.48	21.37	2.88×	1403	1334
NVFP4	16384	62.46	27.42	2.28×	4593	4416

Weight rebuild. The first value pass after each weight change also rebuilds the kernel weights: about 36 ms per MoE layer for MXFP8 (stack and FlashInfer shuffle of the TE MXFP8 parameters) and about 18 ms for NVFP4 (requantization from BF16), independent of the token count. The cost is shared across the step's microbatches and its log-prob pass. The rebuilt weights are allocated inside that first forward, so its saved memory includes them (about 326 MiB MXFP8, 200 MiB NVFP4).

• Both sides run the TE recompute forward and the TE backward; they differ in the first forward, where this PR's FlashInfer pass is faster than TE's. Peak is set by the recompute, so it is the same.
• The NVFP4 recompute in this PR runs TE BF16 (see What is changed), while TE NVFP4 + recompute runs the slow per-expert NVFP4 forward twice. Against BF16 TE + recompute (21.9–31.4 ms), this PR is 1.05–1.21× faster.
• Static memory per rank: MXFP8 1695 vs 1367 MiB (the MXFP8 kernel-weight cache), NVFP4 1543 vs 1721 MiB.

Log-prob (eval() forward under no_grad). Nothing is saved for backward and nothing is recomputed, so the baseline is plain TE.

Recipe	Tokens/rank	TE ms	This PR ms	Speedup	TE peak MiB	This PR peak MiB
MXFP8	1024	7.15	5.11	1.40×	235	276
MXFP8	4096	9.27	6.52	1.42×	898	763
MXFP8	16384	11.02	7.82	1.41×	3552	2698
NVFP4	1024	20.54	5.61	3.66×	262	165
NVFP4	4096	23.53	5.40	4.36×	921	660
NVFP4	16384	21.03	6.24	3.37×	3580	2633

• This PR's log-prob time is about 5–8 ms at every size. At small sizes it is bound by the all-to-all dispatcher's CPU work (kernel launches and the host sync for its split sizes); at 16384 tokens per rank the three NCCL all-to-alls take about 1.6 ms of it. These are the costs of training's own dispatcher, which the train step pays in its recompute anyway.
• TE NVFP4 takes 20–24 ms at every size because TE 2.18 runs row-scaled NVFP4 GEMMs one expert at a time in Python: per expert, a separate activation quantization, an FP32 GEMM, and a per-row rescale and copy. That is about 700 kernel launches per pass for about 2.6 ms of GPU work, so the pass is bound by CPU launch cost, and its timing varies between runs. BF16 TE takes 6.1–10.1 ms here. With CUDA graphs, the gap would shrink toward TE's GPU time.

Script: bench_match_generation_forward.py (not part of this PR). --weight-cache cached (the default) gives the tables above; --weight-cache rebuild resets the cache before every pass, which gives the rebuild cost.

What is changed

• TransformerConfig.moe_training_matches_inference (default off). Construction fails unless:
  • transformer_impl='inference_optimized' and inference_grouped_gemm_backend='flashinfer'. Torch and vLLM need no value pass.
  • MXFP8 or inference_moe_nvfp4_recipe is set. FlashInfer BF16 is not batch invariant.
  • batch_invariant_mode is on, moe_token_dispatcher_type='alltoall', and selective moe recompute is enabled.
  • Nothing makes the training router choose different experts than generation: no sinkhorn or quantile balancing, input jitter, forced load balancing, local CUDA graphs, or overlapped 1F1B schedule. Expert tensor parallelism is also rejected.
• MoELayer.forward hands off to training_matches_inference_forward when the flag is on and inference is not active; custom_forward is unchanged. Everything else is in the new moe/training_matches_inference.py, and the layer keeps no state:
  • With grad enabled, the layer runs under precision_aware_checkpoint, whose first call runs without grad and whose recompute runs with it. The grad mode picks the pass: no grad runs the value pass, grad runs the training custom_forward (the recompute, which supplies the gradient). A no_grad eval forward, the RL log-prob pass, runs the value pass directly.
  • Generation itself never enters the value pass, so its weights, buffers, CUDA graphs and kernels are untouched.
• The value pass runs the layer's own stages: router, preprocess, training's all-to-all dispatch, expert compute, all-to-all combine, and postprocess. Two of them differ from the training forward:
  • Routing calls the layer's router through its generation path (under InferenceMode.active(), for that call only), which returns each token's top-k routes in generation's order. Training's routing map does not keep that order, and FlashInfer sums a token's routes in it. This path also leaves the router's expert-bias token counts to the recompute.
  • Expert compute launches FlashInfer once over all dispatched rows, each row as a top-1 token with probability 1.0, so FlashInfer returns each route's expert output unscaled. The routed MXFP8 kernel has no tile config for a few rows per local expert, so MXFP8 launches are padded with unrouted rows to more than 128 rows per local expert (4097 rows for Nano-v3 at EP 4), where its tile is its largest; batch invariance makes the rows independent of the tile.
  • Unpermute uses a new BF16_RANK_PARTIALS mode of the batch-invariant unpermute kernel (moe/batch_invariant.py, unpermute_with_bf16_rank_partials). For each EP rank it adds the token's routes on that rank in top-k order as FP32 fused multiply-adds of probability and row, from zero, and quantizes the sum to BF16, as FlashInfer's finalize does. It then adds the ranks as the existing batch-invariant combine does, matching generation's ordered or multimem combine. The existing training path compiles the kernel with the mode off, unchanged.
  • The recompute pass runs the normal TE path and supplies the gradient.
• The value pass calls the same routed MXFP8 or NVFP4 FlashInfer kernel that generation calls, directly from training_matches_inference.py, on kernel weights built from the training parameters (MXFP8: the Major-K shuffle from TE's MXFP8 parameters; NVFP4: the requantization from BF16). InferenceGroupedMLP keeps no value-pass state and its forward has no value-pass branch; its only change is allowing the TE MXFP8 recompute with the flag on. Generation's _flashinfer_forward is unchanged.
• The kernel weights are cached in training_matches_inference.py, keyed weakly by experts module, with the address and version of every source tensor, so an in-place write to the parameters rebuilds them. TE writes into MXFP8 parameters bypass the version counter, so a process-wide optimizer step post-hook (torch.optim.optimizer.register_optimizer_step_post_hook) and state_dict and load_state_dict post-hooks on the experts, all registered there, also drop them. A reference-model swap that copies into state_dict() tensors, as NeMo-RL's does, is covered. training_matches_inference.release_inference_value_pass_weights(model) drops them on request, for example before colocated generation. The training pipeline (MegatronModule, the schedules, the optimizer) is unchanged.
• Reading the weights without running the expert linears, the value pass runs Megatron DDP's parameter all-gather pre-forward hooks on them; any other pre-forward hook, such as Megatron-FSDP's, and sharded (DTensor) expert parameters are rejected.
• NVFP4 refit: refresh_cache() requantizes the NVFP4 copy after a refit, replacing main's NotImplementedError.
• The NVFP4 recompute runs TE BF16, because the batch-invariant config rejects TE fp4 with inference-optimized MoE. The MXFP8 recompute runs TE MXFP8.

What is tested

The tests are in tests/unit_tests/inference/test_moe_training_matches_inference.py. GPU tests use EP = world size with the all-to-all dispatcher, the mxfp8, nvfp4-squared_relu and nvfp4-swiglu recipes, batch-invariant mode, and top-k 8:

1. Construction (CPU): the requirements above, and that the value pass keeps generation's top-k route order.
2. Forward and backward run, and the value pass hits the FlashInfer kernel once per step, including under eval().
3. The training forward equals generation bitwise, also after a weight update: the kernel weights are rebuilt after an optimizer step, an in-place copy into state_dict() tensors as NeMo-RL's reference swap does, or the release function, and reused otherwise. Training on 512 and 2048 tokens equals generation decoded 8 tokens at a time.
4. Generation buffers: with generation's NVLS buffers freed, the value pass still matches generation bitwise and allocates none of them.
5. MXFP8 launches are padded to the largest-tile floor and no further; other pre-forward hooks on the expert linears are rejected.
6. The core metric above.

tests/unit_tests/determinism/kernels/test_moe_kernels.py::test_inference_kernel_value_pass_replays[nvfp4-*]: the value pass, recompute and backward replay bit-exactly under contention (EP=1, all-to-all dispatcher).

tests/unit_tests/determinism/kernels/test_moe_kernels.py::test_mxfp8_training_stack_replays: the MXFP8 training-stack quantization replays bit-exactly.

tests/unit_tests/inference/test_flashinfer_nvfp4.py: a refit requantizes the NVFP4 kernel weights, and refit accepts built NVFP4 expert weights.

uv run python -m torch.distributed.run --nproc-per-node 4 -m pytest -q -s \
  tests/unit_tests/inference/test_moe_training_matches_inference.py
uv run python -m torch.distributed.run --nproc-per-node 4 -m pytest -q -s \
  tests/unit_tests/determinism/kernels/test_moe_kernels.py -k "value_pass or mxfp8_training_stack"
uv run python -m torch.distributed.run --nproc-per-node 4 -m pytest -q -s \
  tests/unit_tests/transformer/moe/test_moe_batch_invariant_unpermute.py
uv run python -m torch.distributed.run --nproc-per-node 4 -m pytest -q -s \
  tests/unit_tests/inference/test_flashinfer_nvfp4.py -k refit

Result (GB300 × 4): test_moe_training_matches_inference.py 52 passed; the determinism, batch-invariant unpermute and refit tests pass.

Limitations

• Same EP as generation: generation's combine is not EP-invariant, so training must use the same EP. An EP-invariant generation combine would remove this.
• All-to-all dispatcher only: the flex (DeepEP, HybridEP) and all-gather dispatchers have no batch-invariant combine that adds the EP ranks as generation does.
• ETP > 1 is rejected.

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• If this PR adds or changes a GPU kernel (Triton, jit_fuser/torch.compile, CUDA extension, TE or external-library dispatch, or a scatter/index accumulation), I have added or updated its bit-exact determinism test and registered it in tests/unit_tests/determinism/kernels/manifest.py (guide)
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

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

[YigongQin]

/review

[wdykas]

Do we have the flash infer with nvfp4 in the ci container?

[wdykas]

I think this cache can go stale in the main RL flow. NeMo-RL's use_reference_model swaps the reference weights in with param.copy_() (no optimizer step, no load_state_dict), so the reference logprob pass would reuse the policy's cached kernel weights and the KL would be computed against the policy. Nothing errors out.

Could we key the cache on the source params' _version (the underlying data tensors for MXFP8) instead of relying only on the optimizer/load hooks? That would also catch in-place copies, EMA swaps, etc.

[YigongQin]

Good catch, fixed. The cache is now keyed on the address and _version of every source tensor: the FP8 data and scales for TE MXFP8 parameters, and the parameter itself otherwise. That catches in-place copies and EMA swaps.

[wdykas]

For MXFP8 this is basically a second copy of every local expert weight (~330 MB/layer/rank on Nano), which gets big fast on larger MoEs where expert weights dominate memory. Also since it lives in a module-level dict rather than params/buffers, it stays on the GPU when the trainer offloads the model to CPU before colocated generation, unless someone calls release_inference_value_pass_weights, and nothing calls it today.

Could we at least drop it automatically on offload? Longer term it'd be nice to write FlashInfer's layout directly from the main-param update so there's only one copy.

[YigongQin]

Agreed that it's a second copy of the local expert weights in FlashInfer's layout: about 330 MB per layer per rank for MXFP8 at Nano-v3 shapes, and about 180 MB for NVFP4, which is smaller than TE NVFP4's own quantized copy. It's built on the first value pass after a weight change and reused until the next optimizer step, state_dict() call, checkpoint load or in-place write.

Megatron has no single offload point it could hook without guessing, so the cache is dropped explicitly: release_inference_value_pass_weights(model) frees it, and the next value pass rebuilds it from the parameters. For colocated generation, the trainer should call it before offloading. In NeMo-RL that's offload_before_refit / offload_after_refit in megatron_policy_worker.py. I'll send that as a NeMo-RL follow-up rather than reaching into the trainer from Megatron.

Writing FlashInfer's layout straight from the main-parameter update, so there's only one copy, is the right long-term fix. It needs a change in the optimizer's parameter-update path, so I'd like to keep it out of this PR.

[wdykas]

This hard-codes FlashInfer's internal tile selection (>128 routes per local expert -> 256 tile). If a FlashInfer release changes that heuristic, parity breaks silently. Could we pin the FlashInfer version for this feature, or check which tile actually got picked at runtime?

[YigongQin]

I removed the version check rather than pinning, because the floor can't silently break parity. In batch-invariant mode, the kernel's output rows don't depend on the tile it picks; the token-count parity tests check that. The floor only avoids the kernel's "No valid config" error on very small launches, so if FlashInfer changes its heuristic, the worst case is a loud failure, not a numeric difference.

[wdykas]

did you consider scoping the swap to just the experts instead of the whole MoE layer? Under batch_invariant_mode the training router already matches generation's, so as far as I can tell the real differences are FlashInfer's expert kernel and the combine order. Something like an expert-level op whose forward returns the FlashInfer output and whose backward recomputes the TE grouped MLP would let us keep training's own router and dispatcher. That would drop the InferenceMode toggling, borrowing generation's dispatcher/NVLS buffers, and the per-layer sync, and it would keep the training and inference paths separate like @zhongbozhu had suggested.

The combine order is the part I'm less sure about. Is that the main reason you went through generation's dispatcher? Curious if there's something else I'm missing.

[YigongQin]

did you consider scoping the swap to just the experts instead of the whole MoE layer? Under batch_invariant_mode the training router already matches generation's, so as far as I can tell the real differences are FlashInfer's expert kernel and the combine order. Something like an expert-level op whose forward returns the FlashInfer output and whose backward recomputes the TE grouped MLP would let us keep training's own router and dispatcher. That would drop the InferenceMode toggling, borrowing generation's dispatcher/NVLS buffers, and the per-layer sync, and it would keep the training and inference paths separate like @zhongbozhu had suggested.

The combine order is the part I'm less sure about. Is that the main reason you went through generation's dispatcher? Curious if there's something else I'm missing.

Done. The value pass now runs training's router module and all-to-all dispatcher, with no generation dispatcher, NVLS buffers or per-layer sync. Only the per-rank combine order needed matching: FlashInfer sums a token's routes on each rank in top-k order (fp32 FMA, then BF16). A new mode of the batch-invariant unpermute kernel does the same; with the mode off, the existing path is unchanged. FlashInfer runs each dispatched row as a top-1 token with probability 1.0.

[YigongQin]

/ok to test 6e77cf3