fix(inference): preserve batch invariance for Nemotron Nano

megatron/core/inference/moe/vllm_fused_moe.py

@@ -656,12 +656,18 @@ def vllm_fused_moe(
 
     topk_weights_flat = probs.reshape(-1).contiguous()
 
-    # FC1 + activation: [max_tokens, K] → [max_tokens*topk, N]
-    # SQUARED_RELU fuses into the GEMM epilogue (elementwise). SwiGLU pairs column c
-    # with c+N/2 across tiles and cannot, so FC1 runs unfused to the 2N-wide
-    # intermediate and gate/up is applied separately below.
+    # FC1 + activation: [max_tokens, K] → [max_tokens*topk, N]. Outside
+    # batch-invariant mode, SQUARED_RELU fuses into the GEMM epilogue while
+    # SwiGLU runs separately. Batch-invariant mode runs both activations
+    # separately so their probability application and BF16 rounding match the
+    # training TEGroupedMLP path exactly.
     assert activation_type in (ActivationType.SQUARED_RELU, ActivationType.SWIGLU)
     is_swiglu = activation_type == ActivationType.SWIGLU
+    fuse_squared_relu = not is_swiglu
+    if batch_invariant_mode:
+        # Apply the activation separately below to match training-side rounding.
+        fuse_squared_relu = False
+
     intermediate1 = torch.empty(
         num_valid, N, dtype=hidden_states.dtype, device=hidden_states.device
     )
@@ -677,10 +683,11 @@ def vllm_fused_moe(
         top_k=topk,
         config=config,
         grid_size=grid_size_fc1,
-        fuse_squared_relu=not is_swiglu,
+        fuse_squared_relu=fuse_squared_relu,
     )
-    if is_swiglu:
-        if batch_invariant_mode:
+    if batch_invariant_mode:
+        live_rows = (valid_tokens * topk).to(torch.int32)
+        if is_swiglu:
             # Match training: routing probabilities multiply at the activation
             # (before FC2), with the training kernel's exact rounding sequence
             # (single bf16 round of fp32 silu(gate)*up*prob). The reduction
@@ -691,14 +698,21 @@ def vllm_fused_moe(
                 intermediate1, topk_weights_flat, bound_elems
             )
         else:
-            # intermediate1 is [num_valid, 2N] (gate | up); reduce to [num_valid, N] via
-            # SiLU(gate) * up over the valid_tokens*topk live rows only.
-            n_rows = (valid_tokens * topk).to(torch.int32)
-            intermediate1 = bounded_silu_mul(intermediate1, n_rows)
-
-    # FC2: [max_tokens*topk, N] → [max_tokens*topk, K], without routing weights.
-    # Routing weights are applied in the reduction kernel to avoid an extra
-    # bf16 truncation of prob-scaled values before the topk summation.
+            # Nemotron-style weighted squared-ReLU rounds the squared BF16
+            # activation before its FP32 probability multiply, then rounds to
+            # BF16 again. Reuse the training-parity kernel with the flattened
+            # routing map as the live-row mask.
+            intermediate1 = batch_invariant.squared_relu_with_probs(
+                intermediate1, routing_map.reshape(-1), live_rows, topk_weights_flat
+            )
+    elif is_swiglu:
+        # intermediate1 is [num_valid, 2N] (gate | up); reduce to [num_valid, N] via
+        # SiLU(gate) * up over the valid_tokens*topk live rows only.
+        intermediate1 = bounded_silu_mul(intermediate1, (valid_tokens * topk).to(torch.int32))
+
+    # FC2: [max_tokens*topk, N] → [max_tokens*topk, K]. Batch-invariant mode
+    # already applied routing weights at the activation to match training;
+    # ordinary inference applies them in the reduction kernel.
     # Only local-expert blocks are processed; non-local positions are left
     # undefined and skipped by _moe_sum (which checks the routing map).
     intermediate3 = torch.empty(
@@ -722,6 +736,14 @@ def vllm_fused_moe(
     # Applies routing weights and accumulates in fp32, writes directly to
     # out (if provided), zeros rows beyond valid_tokens, and skips non-local
     # expert slots.
+    apply_routing_weights = True
+    accumulate_in_fp64 = False
+    if batch_invariant_mode:
+        # Probabilities were applied at the activation to mirror training. Use
+        # training's invariant within-rank accumulation for the unweighted sum.
+        apply_routing_weights = False
+        accumulate_in_fp64 = True
+
     return _moe_sum(
         intermediate3,
         probs,
@@ -733,9 +755,6 @@ def vllm_fused_moe(
         local_expert_start,
         num_local_experts,
         out=out,
-        # Batch-invariant mode: probs were already applied at the activation
-        # for SwiGLU (training parity), so the reduction uses unit weights;
-        # fp64 accumulation makes the topk sum order-independent by precision.
-        apply_weights=not (batch_invariant_mode and is_swiglu),
-        acc_fp64=batch_invariant_mode and is_swiglu,
+        apply_weights=apply_routing_weights,
+        acc_fp64=accumulate_in_fp64,
     )

megatron/core/ssm/ssm_inference.py

@@ -223,7 +223,18 @@ def ssm_dynamic_inference(
         if decode_req_count > 0:
             seq_len = 1 + context.num_speculative_tokens
             decode_token_count = decode_req_count * seq_len
-            zxBCdt_decode = zxBCdt[:decode_token_count] if prefill_req_count > 0 else zxBCdt
+            if context.batch_invariant_mode:
+                # Batch-invariant execution may include token-only rows to preserve
+                # model-wide M alignment. Those rows do not represent requests and
+                # must not be passed to the recurrent decode kernels.
+                assert decode_token_count <= zxBCdt.shape[0], (
+                    "Batch-invariant SSM metadata describes more decode tokens "
+                    f"({decode_token_count}) than the input projection contains "
+                    f"({zxBCdt.shape[0]})."
+                )
+                zxBCdt_decode = zxBCdt[:decode_token_count]
+            else:
+                zxBCdt_decode = zxBCdt[:decode_token_count] if prefill_req_count > 0 else zxBCdt
             # Reshape from [N*S, 1, d] to [N, S, d] for the decode kernels.
             zxBCdt_decode = zxBCdt_decode.squeeze(1).view(decode_req_count, seq_len, -1)
             y_decode = self.ssm_decode(
@@ -269,6 +280,17 @@ def ssm_dynamic_inference(
         else:
             raise RuntimeError("Dynamic inference called with 0 decode and 0 prefill requests")
 
+        if context.batch_invariant_mode:
+            # Restore the projection's token-only padding before the output projection.
+            # Its row count can be TP-local, unlike the context's global token count.
+            padding_token_count = zxBCdt.shape[0] - y.shape[0]
+            assert padding_token_count >= 0, (
+                "Batch-invariant SSM produced more token rows "
+                f"({y.shape[0]}) than the input projection contained ({zxBCdt.shape[0]})."
+            )
+            if padding_token_count > 0:
+                y = torch.cat((y, y.new_zeros(padding_token_count, *y.shape[1:])), dim=0)
+
         # Zero padding positions to avoid corrupting quantization amax calculations.
         if is_using_quantization_scales(self.config):
             y[context.padding_slice] = 0.0

megatron/core/transformer/attention.py

@@ -1159,7 +1159,20 @@ def flash_decode_and_prefill(
             # the number of tokens per request. Reshape to (B, S, H, D) so the
             # decode kernel sees batch=num_requests and seqlen_q=tokens_per_request.
             num_requests = seqlens_k.shape[0]
-            tokens_per_request = q.shape[0] // num_requests
+            if self.batch_invariant_mode:
+                # Batch-invariant CUDA-graph buckets can append token-only padding so
+                # model-wide M dimensions stay aligned. Those rows do not represent
+                # requests and must not be passed to attention.
+                input_token_count = q.shape[0]
+                tokens_per_request = int(max_seqlen_q)
+                metadata_token_count = num_requests * tokens_per_request
+                assert metadata_token_count <= input_token_count, (
+                    "Batch-invariant decode metadata describes more query tokens "
+                    f"({metadata_token_count}) than q contains ({input_token_count})."
+                )
+                q = q[:metadata_token_count]
+            else:
+                tokens_per_request = q.shape[0] // num_requests
             q = q.reshape(num_requests, tokens_per_request, q.shape[2], q.shape[3])
 
             # If using MLA we use the FlashMLA kernel
@@ -1273,6 +1286,20 @@ def flash_decode_and_prefill(
             output_total = output_total.reshape(
                 num_requests * tokens_per_request, 1, *output_total.shape[2:]
             )
+            if self.batch_invariant_mode:
+                padding_token_count = input_token_count - output_total.shape[0]
+                assert padding_token_count >= 0, (
+                    "Batch-invariant attention produced more query rows "
+                    f"({output_total.shape[0]}) than q contained ({input_token_count})."
+                )
+                if padding_token_count > 0:
+                    output_total = torch.cat(
+                        (
+                            output_total,
+                            output_total.new_zeros(padding_token_count, 1, *output_total.shape[2:]),
+                        ),
+                        dim=0,
+                    )
 
         return output_total
 

tests/unit_tests/inference/test_flash_decode.py

@@ -1,7 +1,14 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+from types import SimpleNamespace
+
+import pytest
 import torch
 
+import megatron.core.transformer.attention as attention_module
 from megatron.core.models.common.embeddings.rope_utils import apply_rotary_pos_emb_with_cos_sin
 from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding
+from megatron.core.transformer.attention import SelfAttention
 
 
 class TestRotaryEmbeddingWithPrecomputedCosSin:
@@ -29,3 +36,63 @@ def test_output_shapes_match(self):
         assert (
             output_flash_rotary.shape == expected_shape
         ), f"Outputs do not match: {output_flash_rotary.shape} != {expected_shape}"
+
+
+@pytest.mark.parametrize(
+    ("batch_invariant_mode", "num_requests", "tokens_per_request", "padded_token_count"),
+    [
+        (False, 728, 1, 728),
+        (True, 728, 1, 728),
+        (True, 728, 1, 768),
+        (False, 20, 3, 60),
+        (True, 20, 3, 64),
+    ],
+)
+def test_decode_attention_preserves_batch_invariant_token_padding(
+    monkeypatch, batch_invariant_mode, num_requests, tokens_per_request, padded_token_count
+):
+    """Only batch-invariant token-only rows bypass the attention kernel."""
+    attention = object.__new__(SelfAttention)
+    torch.nn.Module.__init__(attention)
+    attention.config = SimpleNamespace(window_size=None, window_attn_skip_freq=None)
+    attention.layer_number = 1
+    attention.batch_invariant_mode = batch_invariant_mode
+    attention.flash_attention_version = 4
+    attention.train(False)
+
+    kernel_queries = []
+
+    def fake_fa4_varlen(q, _k, _v, **_kwargs):
+        kernel_queries.append(q.clone())
+        return q.clone(), None
+
+    monkeypatch.setattr(attention_module, "HAVE_FA4", True)
+    monkeypatch.setattr(attention_module, "flash_attn4_varlen_func", fake_fa4_varlen, raising=False)
+
+    metadata_token_count = num_requests * tokens_per_request
+    num_heads = 2
+    head_dim = 4
+    q = torch.arange(padded_token_count * num_heads * head_dim, dtype=torch.float32).reshape(
+        padded_token_count, 1, num_heads, head_dim
+    )
+    cu_seqlens = torch.arange(0, metadata_token_count + 1, tokens_per_request, dtype=torch.int32)
+    seqlens_k = torch.ones(num_requests, dtype=torch.int32)
+
+    output = attention.flash_decode_and_prefill(
+        q=q,
+        k=torch.empty(0),
+        v=torch.empty(0),
+        max_seqlen_q=tokens_per_request,
+        max_seqlen_k=1,
+        cu_seqlens_q=cu_seqlens,
+        cu_seqlens_k=cu_seqlens,
+        seqlens_k=seqlens_k,
+        block_table=torch.zeros((num_requests, 1), dtype=torch.int32),
+        is_decode_only=True,
+    )
+
+    assert kernel_queries[0].shape == (metadata_token_count, num_heads, head_dim)
+    assert torch.equal(kernel_queries[0], q[:metadata_token_count, 0])
+    assert output.shape == q.shape
+    assert torch.equal(output[:metadata_token_count], q[:metadata_token_count])
+    assert torch.count_nonzero(output[metadata_token_count:]) == 0

tests/unit_tests/inference/test_moe_dispatching_and_routing.py

@@ -657,7 +657,8 @@ def _run_expert_and_combine():
         assert graph_output.shape == hidden_states.shape
         assert graph_output.dtype == torch.bfloat16
 
-    def test_batch_invariant_moe_matches_training(self):
+    @pytest.mark.parametrize("inference_grouped_gemm_backend", ["torch", "vllm"])
+    def test_batch_invariant_moe_matches_training(self, inference_grouped_gemm_backend):
         """The NVLS inference MoE path should exactly match training AllToAll."""
         from megatron.core.models.gpt.moe_module_specs import get_inference_optimized_moe_spec
         from megatron.core.parallel_state import get_expert_model_parallel_group
@@ -673,15 +674,15 @@ def test_batch_invariant_moe_matches_training(self):
             pytest.skip("Training-to-inference MoE parity requires expert parallelism.")
         if Utils.world_size & (Utils.world_size - 1):
             pytest.skip("NVLS Triton symmetric-memory barrier requires power-of-two EP size.")
-        if not HAVE_DEEPGEMM_BF16:
+        if inference_grouped_gemm_backend == "torch" and not HAVE_DEEPGEMM_BF16:
             pytest.skip("Batch-invariant torch MoE path requires DeepGEMM bf16 grouped kernels.")
 
         torch.manual_seed(2028)
         torch.cuda.manual_seed(2028)
 
         config = _make_base_config(
             expert_model_parallel_size=Utils.world_size,
-            inference_grouped_gemm_backend="torch",
+            inference_grouped_gemm_backend=inference_grouped_gemm_backend,
             inference_moe_token_dispatcher_type="nvls",
             batch_invariant_mode=True,
             attention_backend=AttnBackend.flash,
@@ -702,7 +703,8 @@ def test_batch_invariant_moe_matches_training(self):
             local_tokens, 1, config.hidden_size, device="cuda", dtype=torch.bfloat16
         )
 
-        with torch.no_grad(), set_batch_invariant_mode(True):
+        backend = "te_native" if inference_grouped_gemm_backend == "vllm" else None
+        with torch.no_grad(), set_batch_invariant_mode(True, backend=backend):
             training_output, _ = layer(hidden_states.clone())
             with InferenceMode.active():
                 inference_output, _ = layer(hidden_states.clone())

tests/unit_tests/ssm/test_gdp_dynamic_inference.py

@@ -48,6 +48,7 @@
 import torch
 
 from megatron.core import parallel_state
+from megatron.core.inference.batch_dimensions_utils import InferenceBatchDimensions
 from megatron.core.inference.config import InferenceConfig, MambaInferenceStateConfig
 from megatron.core.inference.contexts import StaticInferenceContext
 from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext
@@ -68,6 +69,7 @@
 from megatron.core.ssm.ops.gdp.fused_recurrent import fused_recurrent_gated_delta_rule_update
 from megatron.core.ssm.ops.gdp.metadata import build_gdp_chunk_descriptors, max_gdp_chunk_counts
 from megatron.core.ssm.packed_seq_helpers import check_fla_sequence_packing_support
+from megatron.core.ssm.ssm_inference import SSMDynamicInferenceMixin
 from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
 from megatron.core.transformer import TransformerConfig
 from megatron.core.transformer.cuda_graphs import delete_cuda_graphs
@@ -187,6 +189,81 @@ def _build_model(tp: int = 1) -> HybridModel:
     return model.cuda().eval()
 
 
+class _FakeSSM(SSMDynamicInferenceMixin):
+    def __init__(self, projected):
+        self.projected = projected
+        self.layer_number = 1
+        self.pp_layer_offset = 0
+        self.config = types.SimpleNamespace(fp8=False, fp4=False)
+        self.decode_inputs = []
+
+    def in_proj(self, _hidden_states):
+        return self.projected, None
+
+    def ssm_decode(
+        self,
+        zxBCdt,
+        _conv_state,
+        _ssm_state,
+        batch_indices,
+        intermediate_conv_state=None,
+        intermediate_ssm_state=None,
+    ):
+        self.decode_inputs.append((zxBCdt.clone(), batch_indices.clone()))
+        return zxBCdt[..., :2]
+
+    def out_proj(self, y):
+        return y, None
+
+
+@pytest.mark.parametrize(
+    ("batch_invariant_mode", "num_requests", "tokens_per_request", "padded_token_count"),
+    [
+        (False, 40, 1, 40),
+        (True, 40, 1, 40),
+        (True, 40, 1, 64),
+        (False, 20, 3, 60),
+        (True, 20, 3, 64),
+    ],
+)
+def test_decode_ssm_preserves_batch_invariant_token_padding(
+    batch_invariant_mode, num_requests, tokens_per_request, padded_token_count
+):
+    """Only batch-invariant token-only rows bypass SSM decode."""
+    metadata_token_count = num_requests * tokens_per_request
+    projected = torch.arange(padded_token_count * 4, dtype=torch.float32).reshape(
+        padded_token_count, 1, 4
+    )
+    mixer = _FakeSSM(projected)
+
+    batch_indices = torch.cat(
+        (torch.arange(4, dtype=torch.int32), torch.full((num_requests - 4,), -1, dtype=torch.int32))
+    )
+    context = types.SimpleNamespace(
+        batch_invariant_mode=batch_invariant_mode,
+        mamba_states_cache=lambda _layer, intermediate=False: (torch.empty(0), torch.empty(0)),
+        num_speculative_tokens=tokens_per_request - 1,
+        padded_batch_dimensions=InferenceBatchDimensions(
+            token_count=padded_token_count, prefill_req_count=0, decode_req_count=num_requests
+        ),
+        mamba_metadata=types.SimpleNamespace(batch_indices_decode=batch_indices),
+        padding_slice=slice(metadata_token_count, padded_token_count),
+    )
+
+    output, bias = mixer.ssm_dynamic_inference(torch.empty(0), context)
+
+    decode_input, received_indices = mixer.decode_inputs[0]
+    assert decode_input.shape == (num_requests, tokens_per_request, 4)
+    assert torch.equal(
+        decode_input, projected[:metadata_token_count, 0].view(num_requests, tokens_per_request, 4)
+    )
+    assert torch.equal(received_indices, batch_indices)
+    assert output.shape == (padded_token_count, 1, 2)
+    assert torch.equal(output[:metadata_token_count], projected[:metadata_token_count, :, :2])
+    assert torch.count_nonzero(output[metadata_token_count:]) == 0
+    assert bias is None
+
+
 class TestGDPDynamicInference:
     """Static/dynamic GDP inference equivalence against a full-sequence forward.