[YigongQin]

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

What does this PR do?

Add flashinfer_mega as an inference_grouped_gemm_backend, and as a backend of moe_inference_training_forward (from ). The backend runs FlashInfer's moe_ep megakernel, a single fused kernel that does the EP token transport and the SwiGLU expert MLP. With this PR, a training step's value-producing MoE forward can run the same megakernel that generation runs, while the backward still comes from the TE recompute pass. A training log-prob and a generation log-prob of the same token then come from the same kernel.

Precisions: bf16, mxfp8, nvfp4 and fp8_fp4 are available for generation. Only bf16 and mxfp8 can drive the training forward, because they are the ones whose kernel weights Megatron can rebuild from the parameters. Under mxfp8 the value is quantized and the gradient is that of the BF16 function (straight-through), as for the other backends of the mode.

Motivation: RL

Same goal as the base PR: a zero train/generation log-prob gap for on-policy rollouts. The base PR provides the generic mechanism, with the vllm and torch expert kernels. This PR adds a backend whose kernel fuses dispatch, expert compute and combine, so there is no separate dispatcher or grouped GEMM left to keep bitwise-aligned between training and generation.

1. Mega backend for generation

• megatron/core/inference/moe/mega/ (new):
  • adapter.py: MegatronMegaMoEAdapter wraps FlashInfer's MoEEpMegaLayer. Expert ids are sorted into a canonical order before the combine, because the kernel accumulates prob * expert_out in the order given and float addition is not associative. Generation's top-k order and the training router's mask order would otherwise differ in the last bit.
  • registry.py, _deps.py: the per-precision kernel registry (module-import guarded, so Megatron imports on builds without flashinfer.moe_ep) and the install check.
  • weights.py, training_weights.py: the kernel weight layout (gate/up interleaved in blocks of 32, K-major transpose), and, for mxfp8, quantization through FlashInfer's own mxfp8_quantize_per_block_32 and swizzle, so the bytes match what FlashInfer's preprocess_mega_weights produces.
• InferenceGroupedMLP (experts.py): _mega_forward, and a caller-owned kernel-weight buffer for bf16 and mxfp8. FlashInfer's own preprocessing snapshots the weights at construction, so a refit would not reach the kernel. Owning the buffer lets a refit rewrite it in place, with no teardown, EP collective or recompile.
• resharding/refit.py: the refit calls refresh_mega_weights() on every module that defines it, not only under an MXFP8 transform, because the megakernel keeps a transformed copy under every precision, bf16 included. The refit's final synchronize is scoped to the current stream.
• MegaLocalPassthroughDispatcher (token_dispatcher_inference.py): the megakernel owns EP transport, so every dispatch stage is a pass-through. DynamicInferenceContext skips the NCCL/NVLS per-step buffers for this backend.
• TransformerConfig: inference_mega_precision and inference_mega_max_tokens_per_rank, with construction-time checks that reject what the kernels cannot run: SwiGLU only (no squared-ReLU or other non-gated activations), no tanh-clamped activation, per-precision hidden_size / moe_ffn_hidden_size alignment, experts divisible by EP, no latent MoE, no Megatron-side MXFP8 expert weights. The backend composes with batch_invariant_mode and is exempt from that mode's NVLS requirement, because the megakernel does not use the NVLS dispatcher.
• pyproject.toml: flashinfer-python admitted up to <0.8.0 (flashinfer.moe_ep ships in 0.7.x).

2. Mega backend for the training value pass

• InferenceGroupedMLP._mega_training_forward_pass is the flashinfer_mega entry of the base PR's _inference_training_forward_pass. The kernel's weight layout is a permutation of the parameters, not a view, so unlike the vllm/torch BF16 path the layout cannot alias the parameters. It is repacked into a scratch buffer on every forward, so the kernel always reads the weights the optimizer last wrote.
• The scratch buffer and the kernel layer are shared by all MoE layers of the model, because the kernel's workspace would otherwise dominate memory with one per layer. Each layer repacks immediately before its own forward, and an ownership check after the kernel has read the scratch rejects a stale reader.
• moe_inference_training_forward accepts flashinfer_mega, restricted to bf16 and mxfp8. MoELayer selects the pass-through dispatcher for the value pass, so it never allocates the NVLS symmetric heap.

Tests

• tests/unit_tests/inference/test_flashinfer_mega_moe.py: config validation (backend, alignment, SwiGLU-only, tanh clamp, EP divisibility, precisions), the mega-specific restrictions on moe_inference_training_forward, and a smoke forward.
• tests/unit_tests/inference/test_mega_training_weights.py: the kernel-layout packers against FlashInfer's own preprocess_mega_weights, byte for byte, for bf16 and mxfp8; the shared scratch (stale-owner and mixed-geometry rejection); the MXFP8 requantization follows parameter updates.
• tests/unit_tests/inference/test_mega_training_forward.py (distributed, Blackwell, needs FlashInfer moe_ep):
  • TestValuePassRuns: fwd+bwd runs; the value pass reaches the megakernel exactly once; the same under eval().
  • TestTrainGenParity: bitwise train/generation parity for bf16 and mxfp8, in train and eval mode, and after a weight update.
  • TestTokenCountParity: under batch-invariant mode, a token's output does not depend on how many tokens shared the launch (generation wide vs chunked, and training wide vs generation chunked).
  • TestKernelDeterminism: repeated launches return identical bits.
  • TestAgainstStandardBf16: forward and gradients stay within tolerance of the TE bf16 path.
  • TestGenerationWeightOwnership: the owned repack equals FlashInfer's own preprocessing; a refit without refresh_mega_weights() is stale (negative control) and with it equals a freshly built layer; refit can find the hook; per-layer buffers are independent; unowned precisions refuse a refresh.
  • TestWholeTransformerLayer: the MoE layer inside a full TransformerLayer takes the megakernel and gives matching train/eval forwards.
• Determinism suite (tests/unit_tests/determinism/kernels/): test_mega_value_pass_replays (bf16, mxfp8; fwd+bwd with the shared-scratch repack under racing-stream contention); manifest notes updated for moe_experts and moe_token_dispatchers.

Test results: <FILL IN: unit + determinism results on Blackwell from the container (EP = world size), and the CI run link>

End-to-end validation: NeMo-RL GRPO

Tested on the RL side with both mega precisions:

• bf16 mega and mxfp8 mega GRPO runs (training forward on the megakernel, generation on the megakernel) match the original bf16 path (<FILL IN: what the original bf16 path is, e.g. the zero-KL BF16 baseline from Batch-invariance generation with vLLM fused-MoE backend at improved performance #6521, with vllm/TE experts>).

<FILL IN: model, cluster and parallelism (e.g. nodes x GPUs, EP, TP, colocated or not), number of steps, batch/sequence sizes, which metrics were compared and the observed numbers, job IDs / W&B links>

<FILL IN: figures, e.g. Generation KL Error and reward vs step for original bf16 / mega bf16 / mega mxfp8, with the axes labeled>

• Limits of this evidence: <FILL IN: number of runs, steps and sequence lengths; what "match" means here and to what logged precision; MXFP8 is straight-through, so its gradient is the BF16 function's by design>.

Scope and known limits

• Opt-in and additive: defaults are unchanged. flashinfer_mega is a new inference_grouped_gemm_backend value, and moe_inference_training_forward stays off by default.
• The training forward supports bf16 and mxfp8 only. nvfp4 and fp8_fp4 generation still lets FlashInfer preprocess and snapshot the weights, so refresh_mega_weights() raises NotImplementedError for them instead of serving pre-refit weights silently.
• SwiGLU experts only (the megakernel stacks gate and up into one weight and applies SwiGLU). No tanh-clamped activation, no latent MoE.
• Requires Blackwell (sm100 kernels) and flashinfer-python >= 0.7 with flashinfer.moe_ep. On other builds the backend raises a clear RuntimeError at construction, and the rest of Megatron imports unchanged.
• The repack runs every forward, and it is the cost of keeping the kernel in step with the optimizer. Throughput is not a goal: the mode exists for train/generation parity.
• uv.lock records the old flashinfer-python specifier: <FILL IN: refreshed in the container with uv lock, or left to the CI lockfile job>.

Issue tracking

Linked issue: none. This is a feature PR from an NVIDIA contributor; the design and the related work are tracked in the PRs below.

Related PRs

PR	What it provides	Relation to this PR
	Generic moe_inference_training_forward (value pass on the inference kernel, TE recompute backward) with vllm / torch backends	This PR is stacked on it and adds the flashinfer_mega backend.
#6521	Zero train/generation log-prob gap for BF16 MoE (Qwen3-style, SwiGLU)	The original BF16 path the mega runs are compared against.
#7302	MXFP8 rollout (generation) batch invariance, and TE grouped training	Already on main; the MXFP8 mega generation is checked against batch invariance in the same way.

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests (n/a: unit + determinism coverage; RL end-to-end validation lives in NeMo-RL)
• 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
• I have added proper typing to my code
• I have added relevant documentation (config-field docstrings only)
• 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.