Batch-invariance generation with vLLM fused-MoE backend at improved performance

[utkarsh530]

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

Batch-invariant mode: vLLM fused-MoE backend, te_native GEMM backend, and gated-model support

Summary

batch_invariant_mode currently forces the inference-optimized MoE onto the torch/DeepGEMM
backend, which costs ~6× generation throughput on the same engine, and it cannot run
gated-activation MoE models (e.g. Qwen3-MoE — the models #5700 enabled) at all. This PR makes
the vLLM Triton fused-MoE backend batch-invariant, adds a zero-overhead te_native GEMM
backend, and fixes gated-SwiGLU support under the flag — bringing batch-invariant generation
to parity with the engine's non-invariant throughput.

Measured on Qwen3-30B-A3B, EP8/TP1, 8×B200, full CUDA graphs, BS256/OSL1024:

configuration	tok/s	invariance overhead
BI off, vLLM backend (baseline)	~25,400	—
BI on, as shipped today (torch + DeepGEMM)	~4,070	6.2×
BI on, this PR (vLLM backend + te_native)	~25,100	~0.99×

What's included

• 64-multiple floor for CUDA-graph token buckets — bucket auto-sizing injects 1- and
  2-token decode buckets whose graphed norms execute in a different M%32 reduction bit-class
  than eager steps (TOKEN_ROUNDER=64), breaking cross-batch bit-equality. All sizing
  distributions floored under BI mode; request counts untouched.
• Gated SwiGLU under batch-invariant mode — the BI MoE activation was hard-wired to
  squared-ReLU; gated models crashed with a grouped-GEMM K mismatch. Adds swiglu_with_probs
  (graph-safe, training-parity rounding).
• Batch-invariant vLLM Triton fused-MoE backend — pins only the K-reduction recipe
  (BLOCK_SIZE_K; the kernel is fp32-accumulate with no split-K, so M/N tiling and pipeline
  depth cannot affect bits and stay hint-adaptive), applies routing probabilities at the
  activation with the training kernel's exact rounding (device-bounded, CUDA-graph-safe
  kernel), and reduces top-k with unit weights in fp64.
• te_native GEMM backend — keeps the native cuBLASLt kernels for every dense GEMM
  (aten and TE) and obtains batch invariance via cuBLASLt workspace starvation: split-K
  variants require workspace, so a ~1KB workspace pins every M to the same serial-K reduction
  recipe at native speed. Includes a workaround for TE (≤2.15 verified) hardcoding a 32MiB
  workspace and ignoring CUBLASLT_WORKSPACE_SIZE. Under te_native, native TE RMSNorm is
  also kept (the 64-multiple alignment discipline holds its M%32 bit-class constant).
• Backend selection — TransformerConfig.batch_invariant_backend
  ("deepgemm" | "triton" | "te_native") plumbed through --batch-invariant-backend.
• Tests — tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py: bucket-floor
  properties across all sizing distributions, kernel repeat-determinism / row-locality /
  device-bound soundness under NaN-poisoned tails, _moe_sum option coverage, an end-to-end
  bitwise invariance test (same tokens across different co-batch sizes and launch-config
  classes → torch.equal), and te_native enable/disable round-trip. 11/11 passing;
  the pre-existing BI/inference suites show zero regressions vs stock main
  (identical pass/skip/fail sets in the same environment).

Usage

--batch-invariant-mode --batch-invariant-backend te_native \
--transformer-impl inference_optimized --inference-grouped-gemm-backend vllm \
--inference-moe-token-dispatcher-type nvls \
--attention-backend flash --flash-attention-version 4 --attention-dropout 0.0

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact @NVIDIA/mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

• New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
• Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your 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.

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[wdykas]

/claude strict-review

[wdykas]

I think this is un-needed for the old BiK paths if we are actually overriding the norms but I dont think this will make much of a difference for perf regardless to I am ok activating every time

[utkarsh530]

Yeah you are right. The floor is only load-bearing for the te_native backend (native rmsnorm's M%32 bit-class); with the deepgemm/triton backends the substituted norm makes it redundant. I kept it unconditional so bucket geometry stays uniform across backends (and the collapsed 1/2/4-token buckets slightly reduce graph count/capture time), but happy to gate it on te_native if you'd prefer.

[wdykas]

nope this is fine

[wdykas]

Does this work on all gpu types(hopper, Blackwell, Rubin)?

[utkarsh530]

I tested right now on Blackwell, but vLLM similarly guards this only for Hopper and Blackwell: https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/batch_invariant.py#L915-L927

[wdykas]

I dont think we need to care about ampere but I guess we could guard for when we use Rubin until we test

[wdykas]

For the most part looks good to me. Only question is around do we know this path will work on hopper Blackwell and Rubin?

[claude]

[CRITICAL Correctness] Rounding token_count up here silently deletes decode buckets instead of aligning them, whenever the request count is not a multiple of 64.

_batch_invariant_token_floor rounds up to the next 64-multiple, but InferenceBatchDimensions.is_valid rejects any dims where

token_count > prefill_req_count * max_sequence_length
              + decode_req_count * (num_speculative_tokens + 1)

For the decode-only call sites (prefill_req_count=0, lines ~562-568 and ~629-635) the right-hand side is exactly decode_req_count * (spec + 1), which is the pre-rounding token_count. So as soon as the floor bumps the value, the bucket fails validation and is dropped by the if batch_dim.is_valid(...) guard — no warning, no fallback bucket.

Concretely, with max_requests=100, num_speculative_tokens=0, tp_size=1: decode_req_count = 100, token_count = 100 → floored to 128 → 128 > 100 * 1 → the largest decode graph is never captured. Every decode step at high concurrency then falls back to eager, which both costs the throughput the PR is claiming and — more importantly — puts those steps on a different code path from the graphed ones, which is the opposite of what batch-invariant mode is for. The same happens for any auto-computed bucket size that isn't a 64-multiple.

TestCudaGraphBucket64Floor can't catch this: it only asserts token_count % 64 == 0 over surviving buckets and uses max_requests=512, itself a 64-multiple, so nothing is ever dropped.

Fix: raise the request budget along with the token count so the invariant still holds, and clamp to max_requests so we degrade to a smaller aligned bucket rather than to nothing:

def add_if_valid(token_count: int, prefill_req_count: int, decode_req_count: int) -> None:
    """Helper to create and append batch dimension to list only if it's valid."""
    if _batch_invariant_mode_enabled():
        aligned = _batch_invariant_token_align(token_count)
        if prefill_req_count == 0 and num_speculative_tokens == 0:
            # Decode-only dims are validated against decode_req_count, so the
            # request budget has to grow with the aligned token count or the
            # bucket is silently rejected below.
            decode_req_count = min(aligned, max_requests)
            aligned = min(aligned, decode_req_count)
        token_count = aligned
    ...

and extend the test with a non-64-multiple max_requests (e.g. 100) asserting the largest decode bucket still exists.

[utkarsh530]

Fixed in 8d2f97d. Went with validate-then-align rather than raising the request budget: the alignment is padding (same semantics as the eager TOKEN_ROUNDER), so validity is checked on the unpadded dims and the request count stays truthful; note the proposed min(aligned, decode_req_count) clamp would un-align the tokens again. Added the max_requests=100 regression test asserting the largest decode bucket survives padded.

[claude]

Strict review summary

Findings: CRITICAL: 2 · IMPORTANT: 3 · SUGGESTION: 5

The core idea of this PR is sound, and the parts that matter most for bit-equality — the Inductor-verbatim fp32 SwiGLU sequence with a single final bf16 round, the gate|up half-split convention matching activations.py, the fixed for t in range(topk) reduction order, keeping bound_elems as a device tensor so no host sync breaks CUDA graphs — all check out against the training references in megatron/core/fusions/fused_bias_swiglu.py. No new direct parallel_state.get_*_group() reads anywhere in the diff. No checkpoint-format, architecture, or default-value changes; the inference_grouped_gemm_backend validation change is a backward-compatible widening, and both new _te_patch_for_batch_invariant kwargs default to False, so existing callers are untouched. Every newly added identifier has a real runtime use path.

Most impactful findings

1. Decode CUDA-graph buckets are silently dropped, not aligned (batch_dimensions_utils.py:463-473) — CRITICAL.
add_if_valid rounds token_count up to a 64-multiple, but InferenceBatchDimensions.is_valid rejects decode-only dims where token_count > decode_req_count * (spec + 1). Any max_requests that is not a 64-multiple therefore loses its largest decode graph entirely — e.g. max_requests=100 gives 100 tokens, floored to 128, and 128 > 100 so the bucket is dropped with no warning. Those decode steps fall back to eager, which costs the throughput this PR is claiming and puts them on a different code path from the graphed ones, which is the opposite of the feature purpose. TestCudaGraphBucket64Floor cannot catch it: it uses max_requests=512 and only asserts alignment over the buckets that survive.

2. apply_weights=False is a silent no-op unless acc_fp64=True (vllm_fused_moe.py:495-503) — CRITICAL.
The kernel nests APPLY_WEIGHTS inside the ACC_FP64 branch, so the fp32 path always applies routing probabilities regardless of the flag. The one production call site happens to couple the flags, so this is latent rather than live — but the documented combination double-applies probabilities with no error, and the new tests cover (False, True) and (True, False) while missing exactly (False, False).

3. The te_native workspace patch leaks permanently (batch_invariant_kernels.py:1665-1679) — IMPORTANT.
disable_batch_invariant_mode() never restores get_cublas_workspace_size_bytes, so after one set_batch_invariant_mode(...) block every subsequent TE GEMM in the process runs with a 1 KB cuBLAS workspace — a large silent throughput regression, most visible in a test file that enables BI and then runs unrelated GEMM tests. Separately, os.environ.setdefault makes the env pin a no-op wherever CUBLAS_WORKSPACE_CONFIG is already set — which is about 20 of our own model_config.yaml files (:4096:8) plus one inference example script. In those environments te_native quietly stops being batch-invariant, and the symptom is non-identical output rather than an error.

Also flagged: acc_fp64=batch_invariant_mode pays fp64 on the squared-ReLU path where the reduction is already order-deterministic (suggest and is_swiglu); batch_invariant_backend: str should be Literal[...] so the auto-generated flag gets argparse choices and __post_init__ validation like every other BI setting; and _batch_invariant_token_floor rounds up, a name/behaviour inversion that is exactly what hides finding 1.

Risk assessment

Medium-high, concentrated in the CUDA-graph sizing change. Everything gated behind batch_invariant_mode is opt-in and cannot affect default training or inference, which bounds the blast radius considerably. But finding 1 fires on any non-64-multiple max_requests — an ordinary configuration, not a corner case — and degrades silently into the eager path, so it will present as "batch-invariant mode is slower than advertised" rather than as a bug. The finding 3 workspace leak escapes the opt-in boundary altogether and will affect co-resident non-BI work in the same process. Findings 1 and 2 both want a test case added alongside the fix; finding 3 wants the round-trip assertion in TestTeNativeBackend tightened to check restoration. The remaining items are low-risk cleanups that can land in the same pass.

[wdykas]

/ok to test 8a25689

[wdykas]

/ok to test a5aa088

[nemo-automation-bot]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/32581507186