fix(inference): preserve batch invariance for Nemotron Nano

[wdykas]

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

Summary

This is a targeted follow-up to #6521, which introduced batch-invariant CUDA-graph token alignment and enabled the vLLM
fused-MoE backend under batch-invariant mode.

That MR intentionally pads token counts to 64-token boundaries while leaving request metadata unchanged. For Nemotron Nano, this exposed two missing cases:

• Decode attention and SSM treated token-only padding rows as real request tokens, causing invalid reshape/input dimensions.
• The vLLM squared-ReLU MoE path did not reproduce the activation, routing-probability, and rounding sequence used by the training TEGroupedMLP path, breaking bitwise
  generation/training parity.

we could get rid of this and just us the version where I overwrite the Rmsnorms with our triton kernel but this is much faster. The rmsnorm is 2x slower in triton so I would like to keep this path.

Changes

• Exclude batch-invariant token-only padding before decode attention and SSM kernels, then restore zero padding before their output projections.
• Pin the vLLM MoE K-reduction recipe under batch-invariant mode.
• Match training-side routing-probability application and BF16 rounding for both squared-ReLU and SwiGLU.
• Match the invariant training unpermute with an unweighted FP64 within-rank reduction.
• Keep every behavioral change explicitly gated by batch-invariant mode; ordinary inference behavior is unchanged.

The token alignment introduced by #6521 remains intact—the consumers now correctly distinguish aligned token rows from actual requests.

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

[wdykas]

/ok to test 5eb83f5

[wdykas]

@utkarsh530 In case you need for your work.

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

[santhnm2]

Is it possible to consolidate the new unit tests into existing files?

[wdykas]

/ok to test 1961b5e

[wdykas]

/claude review

[wdykas]

/ok to test a01cc0d

[wdykas]

/ok to test c837c0a

[utkarsh530]

@wdykas Thanks! This is quite helpful

[nemo-automation-bot]

🔄 Merge queue validation started!

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

[nemo-automation-bot]

🔄 Merge queue validation started!

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