[lmcafee-nvidia]

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

What does this PR do?

Makes the inference-optimized MoE experts (InferenceGroupedMLP, torch and vLLM grouped-GEMM backends) apply activation_func_clamp_value, glu_linear_offset, moe_mlp_glu_interleave_size and moe_apply_probs_on_input the way training does. Before this change, the experts dropped these options without an error.

The bug

InferenceGroupedMLP._resolve_mcore_activation_type maps every gated-SiLU configuration to plain SwiGLU. The torch and vllm (default) backends then receive only the activation type and the tanh clamp scale. TEGroupedMLP, the training path of the same module, applies four options that inference never read:

• activation_func_clamp_value: caps the gate at v and clamps the linear half to [-v, v].
• glu_linear_offset: silu(gate) * (linear + offset).
• moe_mlp_glu_interleave_size: FC1 rows alternate gate and linear blocks.
• moe_apply_probs_on_input: the expert input is scaled by the routing probability, and the output is left unweighted.

Config validation rejected these options only under inference_moe_nvfp4_recipe. A model trained with any of them and served with transformer_impl='inference_optimized' therefore produced wrong outputs. This affected BF16 and MXFP8 weights, and batch-invariant mode too. On H100 in the CI image, the expert outputs of main differ from training by a normwise relative error of 7.1 (clamp 0.5), 0.64 (offset 1.0) and 1.4 (interleave 32). A model-level matrix showed 0.8–1.7 nats of mean log-prob error.

Root cause and fix

The inference kernels implemented only plain SwiGLU, and nothing passed the options to them.

• GLU options (activations.py, batch_invariant.py): a frozen GLUOptions value holds the clamp, the linear offset and the interleave size. Two shared Triton helpers do the work. glu_columns finds the gate and linear columns of an output column in either FC1 layout. glu_inputs applies the clamp and the offset to the FP32 inputs. All four SwiGLU kernels use them: padded_swiglu (torch backend and MXFP8), bounded_silu_mul (vLLM), and the batch-invariant swiglu_with_probs and weighted_silu_mul_bounded.
  • The options are constexpr flags, so the default configuration computes exactly what it did before.
  • Like training's torch.clamp, the clamp keeps NaN inputs NaN.
  • Handling the interleave in the kernel keeps the per-expert parameters as views of the stacked weights, which refit relies on.
• moe_apply_probs_on_input: as in training, each routed input row is multiplied by its probability and rounded to BF16, and the expert output is not weighted again. The torch backend does this in the permute kernels (permute_tokens and the fused MXFP8 permute_and_quantize_mxfp8), which then write unit probabilities, so the activation and unpermute code is unchanged. The vLLM backend scales the A rows in its FC1 GEMM (MUL_ROUTED_WEIGHT_ON_INPUT). It then skips the probability in _moe_sum and, in batch-invariant mode, in the activation (weighted_silu_mul_bounded / squared_relu_with_probs accept None).
• Plumbing (fused_moe.py, vllm_fused_moe.py, experts.py): InferenceGroupedMLP builds GLUOptions from the config for gated units, which is the only case where training applies them. It passes the options and moe_apply_probs_on_input to mcore_fused_moe and vllm_fused_moe.
• Validation (transformer_config.py):
  • The FlashInfer MoE kernels apply routing probabilities only to the expert output. moe_apply_probs_on_input with inference_grouped_gemm_backend='flashinfer' is now rejected with a message that points to torch/vllm. The GLU options were already excluded on FlashInfer: gated units reach it only under the NVFP4 recipe, which rejects them.
  • The Transformer Engine GLU activations have no linear offset, so two training paths silently trained without glu_linear_offset while inference now applies it. Those paths are use_te_activation_func and the op-fuser ScaledSwiGLU (use_transformer_engine_op_fuser with SwiGLU experts). Both combinations are now rejected, so training and inference cannot disagree on the offset.

No new allocation, host sync or per-step setup, so the change is CUDA-graph safe.

In batch-invariant mode, the clamp keeps the inference kernels bit-identical to training's fused clamped weighted-SwiGLU (tested). For glu_linear_offset and the interleave, training always uses the unfused eager GLU. Its BF16 roundings are not reproduced; the same is true today for plain SwiGLU with bias_activation_fusion=False. Those cases now match training to BF16 tolerance but not bit for bit.

Tests

• tests/unit_tests/inference/test_moe_dispatching_and_routing.py::TestInferenceGroupedMLPMatchesTraining (1 GPU, TE): one expert module runs its inference path and its training path (TEGroupedMLP) on the same weights and hand-built routing.

  • It covers each option, both backends, and batch-invariant mode on and off. The batch-invariant torch case needs DeepGEMM and skips without it.
  • The routing is passed directly, without the router, so the top-1 case does not depend on router support for top-1.
  • It asserts a normwise error of at most 2e-2, and that each option moves the training output by at least 0.2. Plain SwiGLU is the control.

• test_config_rejects_flashinfer_probs_on_input (CPU), and tests/unit_tests/fusions/test_swiglu_fusion.py::test_swiglu_config_rejects_linear_offset_on_te_glu (CPU; fails without the new validation).

• Kernel reference tests:

  • permute_tokens and permute_and_quantize_mxfp8 with scale_by_probs against PyTorch, bit-exact; the MXFP8 one runs on H100.
  • The batch-invariant SwiGLU kernels with each GLU option, weighted and unweighted (test_vllm_fused_moe_batch_invariant.py).
  • The unweighted batch-invariant squared ReLU.
  • NaN propagation through the clamp in all four SwiGLU kernels (test_moe_permute.py).

• Determinism (test_inference_kernels.py, test_moe_kernels.py, manifest.py notes) adds:

  • replays of every changed kernel variant;
  • an InferenceGroupedMLP replay with the options on both backends;
  • test_batch_invariant_clamped_swiglu_matches_training, a bit-exact check of both batch-invariant SwiGLU kernels against clamped_weighted_swiglu.

  tools/check_kernel_determinism_coverage.py --base-ref origin/main passes.

On H100 in the CI image, at the first commit of this PR:

• The new tests fail on main in 13 cases: the 8 option cases of the parity test, the FlashInfer config test, and the 4 determinism tests that use the new kernel arguments. They pass with this PR: 39 passed, and 8 skipped. The skips are pre-existing Blackwell-only MXFP8/NVFP4 cases in test_inference_kernels.py, which also skip in the H100 CI bucket.
• The bit-exact batch-invariant clamp check passes.
• Neighbor tests pass: test_vllm_fused_moe.py (117), test_vllm_fused_moe_batch_invariant.py (16) and test_moe_permute.py (871).

On H100 in the CI image (rerun at this PR's head), the regression tests pass (126 passed; 11 Blackwell-only MXFP8/NVFP4 cases skipped, as in the H100 CI bucket), and 40 of them fail on main. The neighbor tests pass. The offset-validation tests pass on CPU and fail without the new check.

Issue tracking

Linked issue: none; this PR's regression tests reproduce the bug.

Code provenance

No code was copied from other repositories.

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

🤖 Generated with Claude Code

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.