diff --git a/megatron/core/inference/moe/vllm_fused_moe.py b/megatron/core/inference/moe/vllm_fused_moe.py index a14a40261f3..1b9a0323c43 100644 --- a/megatron/core/inference/moe/vllm_fused_moe.py +++ b/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, ) diff --git a/megatron/core/ssm/ssm_inference.py b/megatron/core/ssm/ssm_inference.py index 16fd356cc96..d8a249a9604 100644 --- a/megatron/core/ssm/ssm_inference.py +++ b/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 diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index 20a586204ed..ab7c9614cf5 100644 --- a/megatron/core/transformer/attention.py +++ b/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 diff --git a/tests/unit_tests/inference/test_flash_decode.py b/tests/unit_tests/inference/test_flash_decode.py index 77ac08c0618..91cfd3f079c 100644 --- a/tests/unit_tests/inference/test_flash_decode.py +++ b/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 diff --git a/tests/unit_tests/inference/test_moe_dispatching_and_routing.py b/tests/unit_tests/inference/test_moe_dispatching_and_routing.py index 79876b7d4fa..12f1200d8b4 100644 --- a/tests/unit_tests/inference/test_moe_dispatching_and_routing.py +++ b/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,7 +674,7 @@ 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) @@ -681,7 +682,7 @@ def test_batch_invariant_moe_matches_training(self): 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()) diff --git a/tests/unit_tests/ssm/test_gdp_dynamic_inference.py b/tests/unit_tests/ssm/test_gdp_dynamic_inference.py index c8135a8a9d7..1880f74e4ad 100644 --- a/tests/unit_tests/ssm/test_gdp_dynamic_inference.py +++ b/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.