From 923e33d855b6456b285b9ef87770ea6df3ef85d6 Mon Sep 17 00:00:00 2001 From: William Dykas Date: Wed, 22 Jul 2026 12:09:54 -0700 Subject: [PATCH 01/14] Add batch-invariant MoE and Mamba inference Signed-off-by: William Dykas --- .../torch_symm_triton/__init__.py | 1 + .../torch_symm_triton/variable_collectives.py | 119 +++ .../inference/contexts/dynamic_context.py | 14 + .../core/inference/engines/dynamic_engine.py | 55 +- .../core/inference/moe/batch_invariant.py | 112 +++ megatron/core/inference/moe/fused_moe.py | 30 +- megatron/core/inference/moe/permute.py | 58 +- megatron/core/ssm/mamba_mixer.py | 67 +- .../core/ssm/ops/batch_invariant_decode.py | 247 ++++++ megatron/core/ssm/ops/ssd_bmm.py | 49 +- megatron/core/ssm/ops/ssd_chunk_scan.py | 90 +- megatron/core/ssm/ops/ssd_chunk_state.py | 62 +- megatron/core/ssm/ops/ssd_combined.py | 104 +++ megatron/core/ssm/ops/ssd_state_passing.py | 87 +- .../core/tensor_parallel/inference_layers.py | 18 +- .../custom_layers/batch_invariant_kernels.py | 788 +++++++++++++++++- .../core/transformer/moe/batch_invariant.py | 92 ++ megatron/core/transformer/moe/moe_utils.py | 54 +- .../core/transformer/moe/token_dispatcher.py | 7 +- .../moe/token_dispatcher_inference.py | 19 +- .../core/transformer/transformer_config.py | 48 +- pyproject.toml | 5 +- .../contexts/test_dynamic_prefix_caching.py | 42 + .../test_moe_dispatching_and_routing.py | 219 +++++ .../unit_tests/inference/test_moe_permute.py | 44 + .../unit_tests/rl/test_rl_batch_invariant.py | 83 ++ .../ssm/ops/test_batch_invariant_decode.py | 646 ++++++++++++++ .../moe/test_moe_batch_invariant.py | 259 ++++++ 28 files changed, 3333 insertions(+), 86 deletions(-) create mode 100644 megatron/core/inference/moe/batch_invariant.py create mode 100644 megatron/core/ssm/ops/batch_invariant_decode.py create mode 100644 megatron/core/transformer/moe/batch_invariant.py create mode 100644 tests/unit_tests/ssm/ops/test_batch_invariant_decode.py create mode 100644 tests/unit_tests/transformer/moe/test_moe_batch_invariant.py diff --git a/megatron/core/inference/communication/torch_symm_triton/__init__.py b/megatron/core/inference/communication/torch_symm_triton/__init__.py index 75da02eaf4b..53523357567 100644 --- a/megatron/core/inference/communication/torch_symm_triton/__init__.py +++ b/megatron/core/inference/communication/torch_symm_triton/__init__.py @@ -7,4 +7,5 @@ multimem_all_gather_v, multimem_all_gatherv_3tensor, multimem_reduce_scatter_v, + ordered_reduce_scatter_v, ) diff --git a/megatron/core/inference/communication/torch_symm_triton/variable_collectives.py b/megatron/core/inference/communication/torch_symm_triton/variable_collectives.py index a32b20b9a14..e1b39716f10 100644 --- a/megatron/core/inference/communication/torch_symm_triton/variable_collectives.py +++ b/megatron/core/inference/communication/torch_symm_triton/variable_collectives.py @@ -354,6 +354,125 @@ def multimem_reduce_scatter_v( return output_tensor +@triton.jit +def _ordered_reduce_scatter_v_kernel( + local_ptr, + buffer_ptrs_dev, + signal_pad_ptrs, + local_tokens, + rank_token_offset_ptr, + ep_max_tokens_ptr, + input_byte_offset, + HIDDEN_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + RANK: tl.constexpr, + WORLD_SIZE: tl.constexpr, +): + """Variable-count reduce-scatter with explicit rank-order fp32 addition. + + This is intentionally not a multimem.ld_reduce kernel. Each rank reads the + same token row from every peer symmetric buffer in rank order and accumulates + in fp32, which gives batch-invariant MoE a defined cross-rank reduction tree. + """ + pid = tl.program_id(axis=0) + + ep_max_tokens = tl.load(ep_max_tokens_ptr) + if pid >= ep_max_tokens: + return + + symm_mem_sync( + signal_pad_ptrs, + None, + RANK, + WORLD_SIZE, + hasPreviousMemAccess=False, + hasSubsequentMemAccess=True, + ) + sync_threads() + + tid = tl.arange(0, BLOCK_SIZE) + rank_token_offset = tl.load(rank_token_offset_ptr) + buffer_ptrs = buffer_ptrs_dev.to(tl.pointer_type(tl.uint64)) + + for token_offset in range(pid, local_tokens, tl.num_programs(axis=0)): + global_token = rank_token_offset + token_offset + + for channel_offset in range(0, HIDDEN_SIZE, BLOCK_SIZE): + offsets = channel_offset + tid + mask = offsets < HIDDEN_SIZE + acc = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + + for src_rank in tl.range(0, WORLD_SIZE): + peer_base = tl.load(buffer_ptrs + src_rank).to(tl.pointer_type(tl.uint8)) + peer_ptr = (peer_base + input_byte_offset).to(tl.pointer_type(tl.float32)) + vals = tl.load(peer_ptr + global_token * HIDDEN_SIZE + offsets, mask=mask, other=0.0) + acc += vals + + tl.store(local_ptr + token_offset * HIDDEN_SIZE + offsets, acc, mask=mask) + + +def ordered_reduce_scatter_v( + output_tensor: torch.Tensor, + input_tensor: torch.Tensor, + symm_mem_hdl: _SymmetricMemory, + rank_token_offset: torch.Tensor, + ep_max_tokens: torch.Tensor, + per_rank_max_tokens: int, + input_byte_offset: int = 0, + **kwargs, +) -> torch.Tensor: + """Variable-count reduce-scatter with fixed rank-order fp32 accumulation. + + This is the batch-invariant alternative to multimem_reduce_scatter_v. It + uses symmetric memory for peer visibility, but performs no hardware FP + reduction; each output token is accumulated by explicitly loading peers in + rank order. + """ + assert HAVE_TRITON, "Triton is required for ordered_reduce_scatter_v." + assert ( + output_tensor.ndim == 2 and input_tensor.ndim == 2 + ), "output_tensor and input_tensor must be 2-D [tokens, hidden_size]." + assert is_device_nvls_capable( + output_tensor.device + ), "ordered_reduce_scatter_v requires a Hopper+ GPU with NVLink (SM >= 9)." + assert ( + output_tensor.dtype == input_tensor.dtype == torch.float32 + ), "ordered_reduce_scatter_v requires fp32 input and output tensors." + assert ( + rank_token_offset.numel() == 1 + and rank_token_offset.dtype == torch.int32 + and rank_token_offset.is_cuda + ), "rank_token_offset must be a scalar int32 CUDA tensor." + + hidden_size = output_tensor.shape[1] + assert ( + input_tensor.shape[1] == hidden_size + ), f"input and output hidden_size mismatch: {input_tensor.shape[1]} vs {hidden_size}" + + MAX_NUM_BLOCKS = kwargs.get("max_num_blocks", 128) + MAX_BLOCK_SIZE = 1024 + WARP_SIZE = 32 + block_size = min(triton.next_power_of_2(hidden_size), MAX_BLOCK_SIZE) + num_warps = max(1, block_size // WARP_SIZE) + num_blocks = min(per_rank_max_tokens, MAX_NUM_BLOCKS) + + _ordered_reduce_scatter_v_kernel[(num_blocks, 1, 1)]( + output_tensor, + symm_mem_hdl.buffer_ptrs_dev, + symm_mem_hdl.signal_pad_ptrs_dev, + local_tokens=output_tensor.shape[0], + rank_token_offset_ptr=rank_token_offset, + ep_max_tokens_ptr=ep_max_tokens, + input_byte_offset=input_byte_offset, + HIDDEN_SIZE=hidden_size, + BLOCK_SIZE=block_size, + RANK=symm_mem_hdl.rank, + WORLD_SIZE=symm_mem_hdl.world_size, + num_warps=num_warps, + ) + return output_tensor + + @triton.jit def _multimem_all_gatherv_3tensor_kernel( local_ptr_0, diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 4a0d0cba518..f00e52bb4c9 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -303,6 +303,7 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC else: self.num_attention_heads_per_partition = 1 + self.batch_invariant_mode = model_config.batch_invariant_mode self.num_speculative_tokens = inference_config.num_speculative_tokens assert self.num_speculative_tokens < inference_config.block_size_tokens, ( f"num_speculative_tokens ({self.num_speculative_tokens}) must be < " @@ -344,6 +345,12 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC self.mamba_ssm_states_dtype = mamba_inference_state_config.ssm_states_dtype self.mamba_chunk_size = mamba_inference_state_config.mamba_chunk_size + if self.batch_invariant_mode: + assert self.num_speculative_tokens == 0, ( + "batch_invariant_mode for Mamba dynamic inference only supports " + "one-token decode; set num_speculative_tokens=0." + ) + # For hybrid models, the layer map converts the global layer index to the # corresponding attention layer index or Mamba layer index depending on the # layer type. @@ -683,6 +690,13 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC # Deal with chunked prefill self.enable_chunked_prefill = inference_config.enable_chunked_prefill + if self.batch_invariant_mode and self.is_hybrid_model and self.enable_chunked_prefill: + # A chunk plus its final token must fit in one step; otherwise a prompt + # of that length can never advance without an invalid one-token tail. + assert self.max_tokens > self.mamba_chunk_size, ( + "batch-invariant Mamba chunked prefill requires max_tokens > " + f"mamba_chunk_size ({self.mamba_chunk_size})." + ) # FlashInfer. if inference_config.use_flashinfer_fused_rope is True: diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 92efff36073..317e6ca1edf 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -1467,6 +1467,26 @@ def _find_mamba_match_count(self, req: DynamicInferenceRequest) -> int: return i + 1 return 0 + def _mamba_batch_invariant_prefill_chunk_length( + self, req: DynamicInferenceRequest, capacity: int + ) -> int: + """Raw prefill length that computes an aligned chunk within ``capacity``. + + Non-final calls must start and end at Mamba chunk boundaries. The final + prompt call may be shorter because it seeds the decode replay tail. + """ + remaining = len(req.remaining_prompt_tokens) + if capacity >= remaining: + return remaining + + chunk_size = self.context.mamba_chunk_size + computed_tokens = (capacity // chunk_size) * chunk_size + if remaining - computed_tokens == 1: + computed_tokens -= chunk_size + if computed_tokens <= 0: + return 0 + return computed_tokens + def schedule_waiting_requests(self): """Tries to schedule any requests in the waiting pool.""" # Keep track of which requests get scheduled. @@ -1571,6 +1591,9 @@ def schedule_chunked_prefill(self): # is_continuing_chunked_prefill is True if we are scheduling next # chunk of a existing chunked prefill request is_continuing_chunked_prefill = self.context.chunked_prefill_request_id >= 0 + batch_invariant_mamba_prefill = ( + self.context.batch_invariant_mode and self.context.is_hybrid_model + ) # Check for conflicting block hashes. if prefix_caching_enabled and not is_continuing_chunked_prefill: @@ -1630,29 +1653,47 @@ def schedule_chunked_prefill(self): not in self.context.kv_block_allocator.kv_hash_to_block_id ): pending_block_hashes.add(block_hash) - prefill_chunk_length = self.context.max_tokens - self.context.active_token_count + available_prefill_length = ( + self.context.max_tokens - self.context.active_token_count + ) + scheduled_prefill_length = available_prefill_length + + if batch_invariant_mamba_prefill: + scheduled_prefill_length = self._mamba_batch_invariant_prefill_chunk_length( + req, available_prefill_length + ) + # No valid non-final Mamba chunk fits in the remaining token budget. + if scheduled_prefill_length == 0: + can_schedule = False + break # If this chunk would leave exactly 1 token for the final chunk, reduce # this chunk by 1 or skip scheduling so the final chunk has 2 tokens. # This avoids the edge case where max_seqlen_q=1 which results in a bug # with the Flash Attention kernel. # See https://github.com/Dao-AILab/flash-attention/issues/1537 - if remaining_len - prefill_chunk_length == 1: - if prefill_chunk_length > 1: - prefill_chunk_length -= 1 + # Batch-invariant Mamba handles this in its chunk-length helper. + if ( + not batch_invariant_mamba_prefill + and remaining_len - scheduled_prefill_length == 1 + ): + if scheduled_prefill_length > 1: + scheduled_prefill_length -= 1 else: # We only have space for 1 token, but remaining is 2. # Delay scheduling to avoid leaving exactly 1 token for the final chunk. can_schedule = False break - self.context.add_request(req, prefill_chunk_length=prefill_chunk_length) + self.context.add_request(req, prefill_chunk_length=scheduled_prefill_length) self._loop.call_soon_threadsafe( self._loop.create_task, self._notify_cond_for_new_request() ) self.context.chunked_prefill_request_id = req.request_id - req.remaining_prompt_tokens = req.remaining_prompt_tokens[prefill_chunk_length:] - req.finished_chunk_token_count += prefill_chunk_length + req.remaining_prompt_tokens = req.remaining_prompt_tokens[ + scheduled_prefill_length: + ] + req.finished_chunk_token_count += scheduled_prefill_length # Still have tokens to prefill, so we break and keep the # chunked prefill request at the head of the waiting queue # Note that we do not need to continue check the queue, as the tokens are full diff --git a/megatron/core/inference/moe/batch_invariant.py b/megatron/core/inference/moe/batch_invariant.py new file mode 100644 index 00000000000..8e7e2869a29 --- /dev/null +++ b/megatron/core/inference/moe/batch_invariant.py @@ -0,0 +1,112 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Batch-invariant inference MoE helpers.""" + +from typing import Optional +from unittest.mock import MagicMock + +import torch + +from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( + grouped_gemm_batch_invariant, + grouped_gemm_batch_invariant_alignment, + is_batch_invariant_mode_enabled, +) +from megatron.core.utils import null_decorator + +try: + import triton + import triton.language as tl + + HAVE_TRITON = True +except ImportError: + HAVE_TRITON = False + +if not HAVE_TRITON: + triton = MagicMock() + triton.jit = null_decorator + tl = MagicMock() + + +def enabled() -> bool: + """Return whether global batch-invariant mode is active.""" + return is_batch_invariant_mode_enabled() + + +def grouped_mm(x_bf16: torch.Tensor, weight: torch.Tensor, offs: torch.Tensor) -> torch.Tensor: + """Batch-invariant BF16 grouped GEMM used by inference fused MoE.""" + return grouped_gemm_batch_invariant( + x_bf16, + weight, + offs=offs.to(torch.int32), + m_total=x_bf16.shape[0], + ) + + +def grouped_mm_alignment() -> int: + """Per-expert row alignment required by the batch-invariant grouped GEMM.""" + return grouped_gemm_batch_invariant_alignment() + + +@triton.jit +def _unpermute_tokens_in_expert_order_kernel( + expert_out_ptr, # [output_size, hidden_dim] bf16 expert outputs + probs_ptr, # [output_size] fp32 routing probabilities + inverse_map_ptr, # [num_tokens, num_local_experts] permuted row or -1 + valid_tokens_ptr, # scalar int32 CUDA tensor: number of valid tokens + output_ptr, # [num_tokens, hidden_dim] fp32 output buffer + hidden_dim, + num_local_experts: tl.constexpr, + BLOCK_H: tl.constexpr, +): + """Token-local batch-invariant unpermute. + + Each program owns one output token and one hidden tile. Contributions are + accumulated in fp32 by increasing local expert id, avoiding atomic-add order. + """ + tok = tl.program_id(0) + block_h = tl.program_id(1) + valid_tokens = tl.load(valid_tokens_ptr) + offsets = block_h * BLOCK_H + tl.arange(0, BLOCK_H) + mask_h = offsets < hidden_dim + + acc = tl.zeros([BLOCK_H], dtype=tl.float32) + if tok < valid_tokens: + for lid in tl.range(0, num_local_experts): + pos = tl.load(inverse_map_ptr + tok * num_local_experts + lid) + if pos >= 0: + prob = tl.load(probs_ptr + pos) + vals = tl.load(expert_out_ptr + pos * hidden_dim + offsets, mask=mask_h).to( + tl.float32 + ) + acc += vals * prob + tl.store(output_ptr + tok * hidden_dim + offsets, acc, mask=mask_h) + + +def unpermute_tokens_in_expert_order( + expert_output: torch.Tensor, + permuted_probs: torch.Tensor, + inverse_map: torch.Tensor, + valid_tokens: torch.Tensor, + out: Optional[torch.Tensor], +) -> torch.Tensor: + """Reduce local expert contributions token-by-token in fixed expert order.""" + _, hidden_dim = expert_output.shape + num_tokens, num_local_experts = inverse_map.shape + if out is None: + out = torch.empty( + num_tokens, hidden_dim, dtype=torch.float32, device=expert_output.device + ) + + BLOCK_H = min(triton.next_power_of_2(hidden_dim), 1024) + grid = (num_tokens, triton.cdiv(hidden_dim, BLOCK_H)) + _unpermute_tokens_in_expert_order_kernel[grid]( + expert_output, + permuted_probs, + inverse_map, + valid_tokens, + out, + hidden_dim, + num_local_experts, + BLOCK_H=BLOCK_H, + ) + return out diff --git a/megatron/core/inference/moe/fused_moe.py b/megatron/core/inference/moe/fused_moe.py index f6c0af4e94e..70cb3da7ea3 100644 --- a/megatron/core/inference/moe/fused_moe.py +++ b/megatron/core/inference/moe/fused_moe.py @@ -21,6 +21,8 @@ ) from megatron.core.inference.quantization.mxfp8_tensor import MXFP8Tensor +from . import batch_invariant + try: from torch.nn.functional import grouped_mm @@ -130,8 +132,17 @@ def mcore_fused_moe( use_mxfp8 = isinstance(fc1_weight, MXFP8Tensor) # Fused quant kernels only apply to MXFP8 path use_fused_quant = use_mxfp8 and not disable_fused_quant_kernels + batch_invariant_mode = batch_invariant.enabled() - if use_mxfp8: + if batch_invariant_mode: + # The MXFP8 path uses scaled_grouped_mm and is not batch invariant. + assert not use_mxfp8, ( + "batch_invariant_mode requires the bf16 grouped GEMM path; got " + "MXFP8 weights. Disable mxfp8 or batch_invariant_mode." + ) + mm_fn = batch_invariant.grouped_mm + expert_alignment = batch_invariant.grouped_mm_alignment() + elif use_mxfp8: assert ( HAVE_SCALED_GMM ), "torch.nn.functional.scaled_grouped_mm not available. Install PyTorch 2.10+." @@ -152,6 +163,7 @@ def mcore_fused_moe( # --- Pre-processing: permute --- if use_fused_quant: # Fused permute + MXFP8 quantize: single kernel produces MXFP8Tensor + batch_invariant_inverse_map = None hidden_states, permuted_probs, permutation_map, offs = permute_and_quantize_mxfp8( hidden_states, probs, @@ -162,7 +174,7 @@ def mcore_fused_moe( alignment=expert_alignment, ) else: - hidden_states, permuted_probs, permutation_map, offs = permute_tokens( + permuted = permute_tokens( hidden_states, probs, routing_map, @@ -170,7 +182,12 @@ def mcore_fused_moe( num_local_experts, valid_tokens, alignment=expert_alignment, + return_batch_invariant_inverse_map=batch_invariant_mode, ) + hidden_states, permuted_probs, permutation_map, offs = permuted[:4] + # Maps each (token, local expert) pair to its row in the expert-grouped buffer, + # allowing batch-invariant unpermute to read contributions in fixed expert order. + batch_invariant_inverse_map = permuted[4] if batch_invariant_mode else None # --- FC1 -> activation -> FC2 --- # Quantize if MXFP8 path and hidden_states not already quantized (fused permute+quant @@ -191,5 +208,12 @@ def mcore_fused_moe( # --- Post-processing: unpermute --- return unpermute_tokens( - fc2_output, permuted_probs, permutation_map, max_tokens, n_used, valid_tokens, out=out + fc2_output, + permuted_probs, + permutation_map, + max_tokens, + n_used, + valid_tokens, + out=out, + batch_invariant_inverse_map=batch_invariant_inverse_map, ) diff --git a/megatron/core/inference/moe/permute.py b/megatron/core/inference/moe/permute.py index 6906c877061..70377fa17bf 100644 --- a/megatron/core/inference/moe/permute.py +++ b/megatron/core/inference/moe/permute.py @@ -15,6 +15,8 @@ from megatron.core.utils import null_decorator +from . import batch_invariant + try: import triton import triton.language as tl @@ -239,6 +241,7 @@ def _permute_tokens_kernel( out_hidden_ptr, # [output_size, hidden_dim] output: permuted hidden states out_probs_ptr, # [output_size] output: permuted probabilities out_src_idx_ptr, # [output_size] output: permutation_map (original token index, -1 for padding) + inverse_map_ptr, # [max_tokens, num_local_experts] token/local-expert -> permuted row counters_ptr, # [num_local_experts] exclusive offsets, atomically incremented valid_tokens_ptr, # scalar int32 CUDA tensor: number of valid tokens this iteration hidden_dim, # hidden dimension @@ -248,6 +251,7 @@ def _permute_tokens_kernel( num_local_experts: tl.constexpr, # number of experts on this rank BLOCK_H: tl.constexpr, # tile size for copying hidden_dim NUM_BLOCKS: tl.constexpr, # grid size (fixed for CG) + HAS_INVERSE: tl.constexpr, # whether to write inverse_map_ptr ): """Permute tokens into expert-grouped order. @@ -282,6 +286,8 @@ def _permute_tokens_kernel( tl.store(out_probs_ptr + pos, tl.load(probs_ptr + tok * topk + k)) # Record source token index for unpermute tl.store(out_src_idx_ptr + pos, tok) + if HAS_INVERSE: + tl.store(inverse_map_ptr + tok * num_local_experts + lid, pos) def permute_tokens( @@ -292,6 +298,7 @@ def permute_tokens( num_local_experts: int, valid_tokens: torch.Tensor, alignment: int = 1, + return_batch_invariant_inverse_map: bool = False, ) -> tuple: """Permute tokens into expert-grouped order. @@ -308,15 +315,22 @@ def permute_tokens( valid_tokens: scalar int32 CUDA tensor with the number of valid tokens this iteration. Fixed address; value updated each step before graph replay. alignment: per-expert token alignment (default 1). + return_batch_invariant_inverse_map: if True, also return the map used by + batch-invariant unpermute. Returns: - (permuted_hidden, permuted_probs, permutation_map, inclusive_offsets) + By default, returns the original 4-tuple: + (permuted_hidden, permuted_probs, permutation_map, inclusive_offsets). + If return_batch_invariant_inverse_map=True, appends the inverse map as a + fifth return value. - permuted_hidden: [output_size, hidden_size] - permuted_probs: [output_size] - permutation_map: [output_size] int32, maps each permuted row back to its original token index. Used by unpermute_tokens to scatter expert outputs back and by activation kernels to skip padding rows (-1). - inclusive_offsets: [num_local_experts] int32 cumulative offsets for grouped_mm + - inverse map: [max_tokens, num_local_experts] int32 map from token/local-expert + to permuted row, only present when requested. """ max_tokens, hidden_dim = hidden_states.shape topk = probs.shape[1] @@ -342,12 +356,24 @@ def permute_tokens( ) permuted_probs = torch.empty(output_size, dtype=probs.dtype, device=probs.device) permutation_map = torch.empty(output_size, dtype=torch.int32, device=probs.device) + batch_invariant_inverse_map = None + if return_batch_invariant_inverse_map: + batch_invariant_inverse_map = torch.full( + (max_tokens, num_local_experts), -1, dtype=torch.int32, device=probs.device + ) # Only initialize [0, n_used) to -1; activation and unpermute kernels are gated # by the same inclusive_expert_offsets[-1] pointer so they never read beyond n_used. init_permutation_map(permutation_map, inclusive_expert_offsets[-1:]) BLOCK_H = min(triton.next_power_of_2(hidden_dim), 1024) max_pairs = max_tokens * topk NUM_BLOCKS = min(max_pairs, 512) + # The inverse-map pointer is unused when HAS_INVERSE=False. Reuse an existing + # int32 device tensor instead of allocating a dummy buffer for that kernel variant. + inverse_map_ptr = ( + batch_invariant_inverse_map + if batch_invariant_inverse_map is not None + else permutation_map + ) _permute_tokens_kernel[(NUM_BLOCKS,)]( hidden_states, probs, @@ -355,6 +381,7 @@ def permute_tokens( permuted_hidden, permuted_probs, permutation_map, + inverse_map_ptr, exclusive_expert_offsets, valid_tokens, hidden_dim, @@ -364,7 +391,16 @@ def permute_tokens( num_local_experts, BLOCK_H=BLOCK_H, NUM_BLOCKS=NUM_BLOCKS, + HAS_INVERSE=batch_invariant_inverse_map is not None, ) + if return_batch_invariant_inverse_map: + return ( + permuted_hidden, + permuted_probs, + permutation_map, + inclusive_expert_offsets, + batch_invariant_inverse_map, + ) return permuted_hidden, permuted_probs, permutation_map, inclusive_expert_offsets @@ -440,6 +476,7 @@ def unpermute_tokens( n_used: torch.Tensor, valid_tokens: torch.Tensor, out: torch.Tensor = None, + batch_invariant_inverse_map: torch.Tensor = None, ) -> torch.Tensor: """Unpermute expert outputs back to original token order. @@ -464,6 +501,25 @@ def unpermute_tokens( permuted_probs.dtype == torch.float32 ), f"permuted_probs must be fp32, got {permuted_probs.dtype}" output_size, hidden_dim = expert_output.shape + + # Triton kernel below uses tl.atomic_add (non-deterministic). Batch-invariant + # MoE instead reduces each token independently in fixed local-expert order, + # so unrelated tokens cannot affect the accumulation tree. + if batch_invariant.enabled(): + assert batch_invariant_inverse_map is not None, ( + "batch-invariant MoE unpermute requires its inverse map" + ) + # The expert-order kernel stores every row tok < valid_tokens, including zero + # rows for tokens with no local expert contribution. Rows beyond + # valid_tokens are not read by the graphed RSV combine. + return batch_invariant.unpermute_tokens_in_expert_order( + expert_output, + permuted_probs, + batch_invariant_inverse_map, + valid_tokens, + out, + ) + BLOCK_H = min(triton.next_power_of_2(hidden_dim), 1024) if out is None: out = torch.empty(num_tokens, hidden_dim, dtype=torch.float32, device=expert_output.device) diff --git a/megatron/core/ssm/mamba_mixer.py b/megatron/core/ssm/mamba_mixer.py index 58b776538c8..d2091952eb8 100644 --- a/megatron/core/ssm/mamba_mixer.py +++ b/megatron/core/ssm/mamba_mixer.py @@ -26,6 +26,7 @@ from megatron.core.inference.utils import InferenceMode from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.ssm.ops.batch_invariant_decode import MambaBatchInvariantDecode from megatron.core.ssm.ops.causal_conv1d_triton import causal_conv1d_update from megatron.core.ssm.ops.mamba_ssm import selective_state_update from megatron.core.tensor_parallel import get_cuda_rng_tracker @@ -438,6 +439,10 @@ def forward( return self._dynamic_inference(hidden_states, inference_context) else: assert inference_context.is_static_batching() + assert not self.config.batch_invariant_mode, ( + "batch_invariant_mode for Mamba inference is only supported with " + "DynamicInferenceContext." + ) assert not self.config.sequence_parallel conv_state, ssm_state = self._get_states_from_cache(inference_context, batch) if inference_context.seqlen_offset > 0: @@ -936,7 +941,20 @@ def _ssm_prefill( chunk_starts = cu_chunk_seqlens[:-1] seq_idx_for_varlen = seq_idx[0, chunk_starts].contiguous() - ssm_varlen_result = mamba_chunk_scan_combined_varlen( + # Batch-invariant decode replays the partial prefill tail, so keep + # the cached SSM state at the last complete chunk boundary. + boundary_chunk_indices = None + has_boundary = None + if self.config.batch_invariant_mode: + prefill_lens = (cu_seqlens[1:] - cu_seqlens[:-1]).to(torch.long) + tail_lens = prefill_lens % self.chunk_size + has_boundary = prefill_lens >= self.chunk_size + # A partial tail uses the preceding full chunk's state. + boundary_chunk_indices = ( + last_chunk_indices.to(torch.long) - (tail_lens > 0).to(torch.long) + ).clamp(min=0) + + scan_result = mamba_chunk_scan_combined_varlen( x=x, dt=dt, A=A, @@ -955,23 +973,40 @@ def _ssm_prefill( z=z if not self.rmsnorm else None, dt_bias=self.cp.get_dt_bias().float(), initial_states=initial_ssm_state, - return_intermediate_states=False, - intermediate_chunk_indices=intermediate_chunk_indices, + return_intermediate_states=self.config.batch_invariant_mode, + intermediate_chunk_indices=( + None if self.config.batch_invariant_mode else intermediate_chunk_indices + ), dt_softplus=True, dt_limit=(0.0, float("inf")), state_dtype=ssm_state.dtype, ) - if intermediate_chunk_indices is not None: - ssm_varlen_states, intermediate_ssm_states = ssm_varlen_result - else: - ssm_varlen_states = ssm_varlen_result - intermediate_ssm_states = None - y = y.unsqueeze(0) z = z.unsqueeze(0) - tensor_masked_update(ssm_state, batch_indices, ssm_varlen_states) + if self.config.batch_invariant_mode: + scan_states = scan_result + boundary_mask = has_boundary.view(-1, 1, 1, 1) + cache_states = torch.where( + boundary_mask, scan_states[boundary_chunk_indices], initial_ssm_state + ) + intermediate_ssm_states = ( + scan_states[intermediate_chunk_indices] + if intermediate_chunk_indices is not None + else None + ) + elif intermediate_chunk_indices is not None: + cache_states, intermediate_ssm_states = scan_result + else: + cache_states = scan_result + intermediate_ssm_states = None + + tensor_masked_update(ssm_state, batch_indices, cache_states) + if self.config.batch_invariant_mode: + self._batch_invariant_decode().seed( + x, dt, B, C, cu_seqlens, batch_indices, max_batch=ssm_state.shape[0] + ) # Write intermediate states to pre-allocated output buffers # All tensor ops, no Python loops, fully CUDA graph compatible. @@ -1046,6 +1081,12 @@ def _get_decode_A_neg_exp(self) -> torch.Tensor: self._A_neg_exp_cache_stale = False return self._A_neg_exp_cache.view(-1, 1, 1).expand(-1, self.headdim, self.d_state) + def _batch_invariant_decode(self) -> MambaBatchInvariantDecode: + """Batch-invariant decode adapter, created on first use.""" + if not hasattr(self, "_batch_invariant_decoder"): + self._batch_invariant_decoder = MambaBatchInvariantDecode(self) + return self._batch_invariant_decoder + def train(self, mode: bool = True): """Mark the decode cache stale; weights may have updated.""" if mode: @@ -1192,6 +1233,12 @@ def _ssm_decode( y = y * self.act(z) # (B D) y = y.unsqueeze(1) # Restore seq dimension + elif self.config.batch_invariant_mode: + assert batch_indices is not None, ( + "batch_invariant_mode for Mamba decode requires dynamic batching " + "batch_indices." + ) + y = self._batch_invariant_decode().step(x, dt, B, C, batch_indices, ssm_state) else: A = self._get_decode_A_neg_exp() diff --git a/megatron/core/ssm/ops/batch_invariant_decode.py b/megatron/core/ssm/ops/batch_invariant_decode.py new file mode 100644 index 00000000000..09fe5ff5e6d --- /dev/null +++ b/megatron/core/ssm/ops/batch_invariant_decode.py @@ -0,0 +1,247 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +"""Batch-invariant Mamba decode using buffered chunk replay.""" + +from dataclasses import dataclass + +import torch + +from megatron.core.ssm.ops.ssd_combined import mamba_chunk_scan_decode_rows + + +@dataclass +class BatchInvariantDecodeBuffers: + """Per-slot persistent state for the buffered decode scan.""" + + x: torch.Tensor # (max_batch + 1, chunk_size, nheads, headdim) + dt: torch.Tensor # (max_batch + 1, chunk_size, nheads) + B: torch.Tensor # (max_batch + 1, chunk_size, ngroups, dstate) + C: torch.Tensor # (max_batch + 1, chunk_size, ngroups, dstate) + # Tokens buffered since the slot's last chunk boundary; doubles as the + # write cursor for the next token. + num_buffered: torch.Tensor # (max_batch + 1,) int32 + # Per-entry target-row output, allocated once and sliced per step. + out: torch.Tensor # (max_batch + 1, nheads, headdim) + + @classmethod + def allocate( + cls, + max_batch: int, + chunk_size: int, + nheads: int, + headdim: int, + ngroups: int, + dstate: int, + device: torch.device, + dtype: torch.dtype, + ) -> "BatchInvariantDecodeBuffers": + """Allocate the per-slot decode buffers.""" + # Padding entries use batch index -1. Map them to an extra row so + # fixed-shape graph code can write without touching a live request. + rows = max_batch + 1 + return cls( + x=torch.zeros(rows, chunk_size, nheads, headdim, device=device, dtype=dtype), + dt=torch.zeros(rows, chunk_size, nheads, device=device, dtype=dtype), + B=torch.zeros(rows, chunk_size, ngroups, dstate, device=device, dtype=dtype), + C=torch.zeros(rows, chunk_size, ngroups, dstate, device=device, dtype=dtype), + num_buffered=torch.zeros(rows, device=device, dtype=torch.int32), + out=torch.empty(rows, nheads, headdim, device=device, dtype=dtype), + ) + + @property + def trash_row(self) -> int: + """Write sink for padding entries (the buffers' extra last row).""" + return self.num_buffered.shape[0] - 1 + + def map_slots(self, batch_indices: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Map padding entries to the trash row.""" + slots = batch_indices.to(torch.long) + is_active = slots >= 0 + return slots.masked_fill(~is_active, self.trash_row), is_active + + def seed( + self, + x: torch.Tensor, + dt: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + cu_seqlens: torch.Tensor, + batch_indices: torch.Tensor, + ) -> None: + """Store each prefill's unfinished chunk for decode replay.""" + chunk_size = self.x.shape[1] + num_seqs = cu_seqlens.numel() - 1 + + seq_starts = cu_seqlens[:-1].to(torch.long) + seq_ends = cu_seqlens[1:].to(torch.long) + prefill_lens = seq_ends - seq_starts + # Covers every case: prefill_len < chunk_size gives prefill_len, + # boundary-aligned gives 0. + tail_lens = prefill_lens % chunk_size + + # Redirect padding entries (batch_indices < 0) to the trash row so all + # writes below are unconditional. + slots, is_active = self.map_slots(batch_indices[:num_seqs]) + + # Fill unused rows with a valid token from the same sequence. The row-gated + # kernel evaluates a full M-block, so finite padding prevents masked NaNs + # from reaching the target row through tensor-core operations. + offsets = torch.arange(chunk_size, device=x.device, dtype=torch.long) + safe_tail_lens = torch.clamp(tail_lens, min=1) + safe_tail_offsets = torch.minimum( + offsets.unsqueeze(0), (safe_tail_lens - 1).unsqueeze(1) + ) + safe_tail_starts = torch.where( + tail_lens > 0, + seq_ends - tail_lens, + torch.clamp(seq_ends - 1, min=0), + ) + tail_token_idx = (safe_tail_starts.unsqueeze(1) + safe_tail_offsets).clamp( + max=x.shape[0] - 1 + ) + + self.x[slots] = x[tail_token_idx] + self.dt[slots] = dt[tail_token_idx] + self.B[slots] = B[tail_token_idx] + self.C[slots] = C[tail_token_idx] + + # Keep the trash row's count pinned at 0 so its buffer writes stay in + # bounds. + self.num_buffered[slots] = torch.where( + is_active, tail_lens, torch.zeros_like(tail_lens) + ).to(torch.int32) + + +def batch_invariant_decode_buffered_scan( + buffers: BatchInvariantDecodeBuffers, + x: torch.Tensor, # (decode_batch_size, 1, nheads, headdim) + dt: torch.Tensor, # (decode_batch_size, 1, nheads) + B: torch.Tensor, # (decode_batch_size, 1, ngroups, dstate) + C: torch.Tensor, # (decode_batch_size, 1, ngroups, dstate) + A: torch.Tensor, + D: torch.Tensor, + dt_bias: torch.Tensor, + batch_indices: torch.Tensor, + ssm_state: torch.Tensor, +) -> torch.Tensor: + """Run one decode token with full chunk-scan arithmetic. + + Mutates the replay buffers and commits ``ssm_state`` when a chunk fills. + """ + decode_batch_size, tokens_per_entry, nheads, headdim = x.shape + dstate = B.shape[-1] + chunk_size = buffers.x.shape[1] + assert tokens_per_entry == 1, ( + "batch-invariant Mamba decode assumes one new token per request " + "per call (no speculative decoding)." + ) + output_capacity = buffers.out.shape[0] + assert decode_batch_size <= output_capacity, ( + f"decode batch size {decode_batch_size} exceeds the output buffer capacity " + f"({output_capacity}); increase max_batch." + ) + + # Redirect padding entries (batch_indices < 0) to the trash row so the + # buffer writes below are unconditional. + slots, is_active = buffers.map_slots(batch_indices) + # ssm_state is engine-owned and has no trash row: clamp for reads. Its + # only writes happen in-kernel for crossing slots, which never alias. + state_slots = slots.clamp(max=buffers.trash_row - 1) + + # Write the new token at each slot's cursor; write_pos is also the + # token's intra-chunk row, the one row the scan must produce. + write_pos = buffers.num_buffered[slots].to(torch.long) + buffers.x[slots, write_pos] = x[:, 0] + buffers.dt[slots, write_pos] = dt[:, 0] + buffers.B[slots, write_pos] = B[:, 0] + buffers.C[slots, write_pos] = C[:, 0] + + # A slot crosses its chunk boundary when this token fills the buffer. + crossed = (write_pos + 1 == chunk_size) & is_active + out = buffers.out[:decode_batch_size] + + # Run the gated pipeline over the buffers and ssm_state in place. State + # passing writes crossing slots' boundary states straight into + # ssm_state, so no scatter is needed afterwards. + mamba_chunk_scan_decode_rows( + buffers.x.view(-1, nheads, headdim), + buffers.dt.view(-1, nheads), + A, + buffers.B.view(-1, buffers.B.shape[-2], dstate), + buffers.C.view(-1, buffers.C.shape[-2], dstate), + chunk_size, + chunk_starts=(slots * chunk_size).to(torch.int32), + slots=state_slots.to(torch.int32), + target_rows=write_pos.to(torch.int32), + chunk_flags=crossed.to(torch.int32), + initial_states=ssm_state, + out=out, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + ) + + # The scan stored each entry's target row at out[i]; padding entries + # return zeros. + y = torch.where(is_active.view(-1, 1, 1), out, torch.zeros_like(out)).unsqueeze(1) + + # Crossed slots restart their buffer; the rest advance. Padding entries + # write 0 to the trash row, keeping its cursor pinned in bounds. + buffers.num_buffered[slots] = torch.where( + crossed | ~is_active, torch.zeros_like(write_pos), write_pos + 1 + ).to(torch.int32) + + return y + + +class MambaBatchInvariantDecode: + """Adapter between a MambaMixer and the buffered decode.""" + + def __init__(self, mixer): + # The gate is applied outside the scan (RMSNormGated), so the + # buffers carry no z. Enforced here because the decode path would + # otherwise silently drop it. + assert mixer.rmsnorm, "batch_invariant_mode requires rmsnorm=True" + self.mixer = mixer + self.buffers: BatchInvariantDecodeBuffers | None = None + + def _get_buffers(self, max_batch, x, B) -> BatchInvariantDecodeBuffers: + if self.buffers is None: + nheads, headdim = x.shape[-2:] + ngroups, dstate = B.shape[-2:] + self.buffers = BatchInvariantDecodeBuffers.allocate( + max_batch, + self.mixer.chunk_size, + nheads, + headdim, + ngroups, + dstate, + x.device, + x.dtype, + ) + return self.buffers + + def seed(self, x, dt, B, C, cu_seqlens, batch_indices, max_batch) -> None: + """Seed replay buffers from the prefill tail.""" + buffers = self._get_buffers(max_batch, x, B) + buffers.seed(x, dt, B, C, cu_seqlens, batch_indices) + + def step(self, x, dt, B, C, batch_indices, ssm_state) -> torch.Tensor: + """Run one decode step using the mixer's flattened layouts.""" + mixer = self.mixer + batch = x.shape[0] + x = x.view(batch, 1, -1, mixer.headdim) + B = B.view(batch, 1, mixer.ngroups_local_tp, -1) + C = C.view(batch, 1, mixer.ngroups_local_tp, -1) + + A = -torch.exp(mixer.cp.get_A_log().float()) + D = mixer.cp.get_D() + if mixer.D_has_hdim: + D = D.float().view(-1, mixer.headdim) + dt_bias = mixer.cp.get_dt_bias().float() + + buffers = self._get_buffers(ssm_state.shape[0], x, B) + + y = batch_invariant_decode_buffered_scan( + buffers, x, dt, B, C, A, D, dt_bias, batch_indices, ssm_state + ) + return y.reshape(batch, 1, -1) diff --git a/megatron/core/ssm/ops/ssd_bmm.py b/megatron/core/ssm/ops/ssd_bmm.py index 0cbb07fdbf5..ba07563a1f7 100644 --- a/megatron/core/ssm/ops/ssd_bmm.py +++ b/megatron/core/ssm/ops/ssd_bmm.py @@ -67,7 +67,8 @@ def _bmm_chunk_fwd_kernel( a_ptr, b_ptr, out_ptr, - cu_chunk_seqlens_ptr, + chunk_offsets_ptr, + target_rows_ptr, # Matrix dimensions seqlen, chunk_size: tl.constexpr, @@ -85,6 +86,7 @@ def _bmm_chunk_fwd_kernel( stride_outn: tl.constexpr, # Meta-parameters IS_CAUSAL: tl.constexpr, + HAS_TARGET_ROWS: tl.constexpr, dot_dtype: tl.constexpr, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, @@ -99,9 +101,20 @@ def _bmm_chunk_fwd_kernel( if IS_CAUSAL: if pid_n * BLOCK_SIZE_N >= (pid_m + 1) * BLOCK_SIZE_M: return + if HAS_TARGET_ROWS: + # Keep the target row's tile and its causal column tiles. + tr = tl.load(target_rows_ptr + pid_c) + if pid_m != tr // BLOCK_SIZE_M: + return + if pid_n * BLOCK_SIZE_N > tr: + return - chunk_seqlen_start = tl.load(cu_chunk_seqlens_ptr + pid_c) - chunk_seqlen_end = tl.load(cu_chunk_seqlens_ptr + pid_c + 1) + chunk_seqlen_start = tl.load(chunk_offsets_ptr + pid_c) + if HAS_TARGET_ROWS: + # Fixed windows need only a start offset. + chunk_seqlen_end = chunk_seqlen_start + chunk_size + else: + chunk_seqlen_end = tl.load(chunk_offsets_ptr + pid_c + 1) a_ptr += chunk_seqlen_start * stride_a_seqlen + pid_h * stride_a_head b_ptr += chunk_seqlen_start * stride_b_seqlen + pid_h * stride_b_head @@ -140,7 +153,16 @@ def _bmm_chunk_fwd_kernel( tl.store(out_ptrs, out, mask=(offs_m[:, None] < chunk_size) & (offs_n[None, :] < chunk_size)) -def _bmm_chunk_fwd(a, b, chunk_size, cu_chunk_seqlens, causal=False, output_dtype=None): +def _bmm_chunk_fwd( + a, + b, + chunk_size, + cu_chunk_seqlens, + causal=False, + output_dtype=None, + target_rows=None, + chunk_starts=None, +): """ Argument: a: (seqlen, ngroups, k) @@ -149,9 +171,23 @@ def _bmm_chunk_fwd(a, b, chunk_size, cu_chunk_seqlens, causal=False, output_dtyp cu_chunk_seq_lens: (nchunks+1,) causal: if True, then out[i, j] for i > j will be arbitrary, only out[i, j] for i <= j are guaranteed to be correct. + target_rows: optional (nchunks,) int32. Decode mode: compute only the + M-block containing target_rows[c] (and N-blocks up to it); other + output entries are left uninitialized. Return: out: (nchunks, ngroups, chunk_size, chunk_size) """ + has_target_rows = target_rows is not None + assert (chunk_starts is not None) == has_target_rows, ( + "target_rows and chunk_starts must be provided together" + ) + if has_target_rows: + # chunk_starts has one fixed-window start per chunk. + chunk_offsets = chunk_starts + nchunks = len(chunk_starts) + else: + chunk_offsets = cu_chunk_seqlens + nchunks = len(cu_chunk_seqlens) - 1 seqlen, ngroups, k = a.shape assert b.shape == a.shape if a.stride(-1) != 1 and a.stride(0) != 1: @@ -159,7 +195,6 @@ def _bmm_chunk_fwd(a, b, chunk_size, cu_chunk_seqlens, causal=False, output_dtyp if b.stride(-1) != 1 and b.stride(0) != 1: b = b.contiguous() - nchunks = len(cu_chunk_seqlens) - 1 # Allocates output. out_dtype = a.dtype if output_dtype is None else output_dtype out = torch.empty((nchunks, ngroups, chunk_size, chunk_size), device=a.device, dtype=out_dtype) @@ -178,7 +213,8 @@ def _bmm_chunk_fwd(a, b, chunk_size, cu_chunk_seqlens, causal=False, output_dtyp a_ptr=a, b_ptr=b, out_ptr=out, - cu_chunk_seqlens_ptr=cu_chunk_seqlens, + chunk_offsets_ptr=chunk_offsets, + target_rows_ptr=target_rows, seqlen=seqlen, chunk_size=chunk_size, K=k, @@ -194,6 +230,7 @@ def _bmm_chunk_fwd(a, b, chunk_size, cu_chunk_seqlens, causal=False, output_dtyp stride_outm=out.stride(-2), stride_outn=out.stride(-1), IS_CAUSAL=causal, + HAS_TARGET_ROWS=has_target_rows, dot_dtype=dot_dtype, ) return out diff --git a/megatron/core/ssm/ops/ssd_chunk_scan.py b/megatron/core/ssm/ops/ssd_chunk_scan.py index 521a294db5d..88dcf10bdec 100644 --- a/megatron/core/ssm/ops/ssd_chunk_scan.py +++ b/megatron/core/ssm/ops/ssd_chunk_scan.py @@ -88,7 +88,8 @@ def _chunk_scan_fwd_kernel( states_ptr, D_ptr, initstates_ptr, - cu_chunk_seqlens_ptr, + chunk_offsets_ptr, + target_rows_ptr, # Matrix dimensions chunk_size: tl.constexpr, hdim: tl.constexpr, @@ -133,6 +134,7 @@ def _chunk_scan_fwd_kernel( HAS_D: tl.constexpr, D_HAS_HDIM: tl.constexpr, HAS_Z: tl.constexpr, + HAS_TARGET_ROWS: tl.constexpr, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, @@ -145,9 +147,17 @@ def _chunk_scan_fwd_kernel( num_pid_n = tl.cdiv(hdim, BLOCK_SIZE_N) pid_m = tl.program_id(axis=0) // num_pid_n pid_n = tl.program_id(axis=0) % num_pid_n + if HAS_TARGET_ROWS: + # Keep the tile containing the only output row consumed. + if pid_m != tl.load(target_rows_ptr + pid_c) // BLOCK_SIZE_M: + return cb_ptr += pid_c * stride_cb_chunk + (pid_h // nheads_ngroups_ratio) * stride_cb_head - chunk_seqlen_start = tl.load(cu_chunk_seqlens_ptr + pid_c) - chunk_seqlen_end = tl.load(cu_chunk_seqlens_ptr + pid_c + 1) + chunk_seqlen_start = tl.load(chunk_offsets_ptr + pid_c) + if HAS_TARGET_ROWS: + # Fixed windows need only a start offset. + chunk_seqlen_end = chunk_seqlen_start + chunk_size + else: + chunk_seqlen_end = tl.load(chunk_offsets_ptr + pid_c + 1) x_ptr += chunk_seqlen_start * stride_x_seqlen + pid_h * stride_x_head dt_ptr += pid_c * stride_dt_chunk + pid_h * stride_dt_head dA_cumsum_ptr += pid_c * stride_dA_cs_chunk + pid_h * stride_dA_cs_head @@ -159,20 +169,31 @@ def _chunk_scan_fwd_kernel( seq_idx_ptr += pid_c * stride_seq_idx_chunk seq_idx = tl.load(seq_idx_ptr) - seq_idx_prev = tl.load(seq_idx_ptr - stride_seq_idx_chunk, mask=pid_c >= 1, other=-1) - - if HAS_INITSTATES and (seq_idx != seq_idx_prev): + if HAS_TARGET_ROWS: + # Each fixed window starts from its indexed cached state. + seq_idx_prev = -1 prev_states_ptr = ( initstates_ptr + seq_idx * stride_init_states_batch + pid_h * stride_init_states_head ) prev_states_hdim = stride_init_states_hdim prev_states_dstate = stride_init_states_dstate else: - prev_states_ptr = ( - states_ptr + (pid_c - 1) * stride_states_chunk + pid_h * stride_states_head - ) - prev_states_hdim = stride_states_hdim - prev_states_dstate = stride_states_dstate + seq_idx_prev = tl.load(seq_idx_ptr - stride_seq_idx_chunk, mask=pid_c >= 1, other=-1) + + if HAS_INITSTATES and (seq_idx != seq_idx_prev): + prev_states_ptr = ( + initstates_ptr + + seq_idx * stride_init_states_batch + + pid_h * stride_init_states_head + ) + prev_states_hdim = stride_init_states_hdim + prev_states_dstate = stride_init_states_dstate + else: + prev_states_ptr = ( + states_ptr + (pid_c - 1) * stride_states_chunk + pid_h * stride_states_head + ) + prev_states_hdim = stride_states_hdim + prev_states_dstate = stride_states_dstate chunk_size_limit = chunk_seqlen_end - chunk_seqlen_start @@ -305,13 +326,32 @@ def _chunk_scan_fwd_kernel( ).to(tl.float32) acc *= z * tl.sigmoid(z) - out_ptr += chunk_seqlen_start * stride_out_seqlen + pid_h * stride_out_head - out_ptrs = out_ptr + ( - stride_out_seqlen * offs_out_m[:, None] + offs_out_n[None, :] * stride_out_hdim - ) - tl.store( - out_ptrs, acc, mask=(offs_out_m[:, None] < chunk_size_limit) & (offs_out_n[None, :] < hdim) - ) + if HAS_TARGET_ROWS: + # Store just the target row to a compact (nchunks, nheads, hdim) + # output; nothing else is consumed downstream. Same acc values as + # the full store, only the mask is narrower. + tr = tl.load(target_rows_ptr + pid_c) + out_ptr += pid_c * stride_out_seqlen + pid_h * stride_out_head + # All M-lanes alias the same output row (row stride 0); the mask + # keeps only lane tr, so one lane stores per column. + out_ptrs = out_ptr + ( + offs_out_m[:, None] * 0 + offs_out_n[None, :] * stride_out_hdim + ) + tl.store( + out_ptrs, + acc, + mask=(offs_out_m[:, None] == tr) & (offs_out_n[None, :] < hdim), + ) + else: + out_ptr += chunk_seqlen_start * stride_out_seqlen + pid_h * stride_out_head + out_ptrs = out_ptr + ( + stride_out_seqlen * offs_out_m[:, None] + offs_out_n[None, :] * stride_out_hdim + ) + tl.store( + out_ptrs, + acc, + mask=(offs_out_m[:, None] < chunk_size_limit) & (offs_out_n[None, :] < hdim), + ) def _chunk_scan_fwd( @@ -327,8 +367,18 @@ def _chunk_scan_fwd( D=None, z=None, initial_states=None, + target_rows=None, + chunk_starts=None, ): assert seq_idx is not None, "this implementation requires seq_idx" + has_target_rows = target_rows is not None + assert (chunk_starts is not None) == has_target_rows, ( + "target_rows and chunk_starts must be provided together" + ) + if has_target_rows: + chunk_offsets = chunk_starts + else: + chunk_offsets = cu_chunk_seqlens seqlen, nheads, headdim = x.shape _, nchunks, chunk_size = dt.shape @@ -375,7 +425,8 @@ def _chunk_scan_fwd( states_ptr=states, D_ptr=D, initstates_ptr=initial_states, - cu_chunk_seqlens_ptr=cu_chunk_seqlens, + chunk_offsets_ptr=chunk_offsets, + target_rows_ptr=target_rows, chunk_size=chunk_size, hdim=headdim, dstate=dstate, @@ -417,6 +468,7 @@ def _chunk_scan_fwd( HAS_D=D is not None, D_HAS_HDIM=D.dim() == 2 if D is not None else True, HAS_Z=z is not None, + HAS_TARGET_ROWS=has_target_rows, BLOCK_SIZE_DSTATE=max(triton.next_power_of_2(dstate), 16), IS_TRITON_22=TRITON_22, HAS_INITSTATES=initial_states is not None, diff --git a/megatron/core/ssm/ops/ssd_chunk_state.py b/megatron/core/ssm/ops/ssd_chunk_state.py index 473af1491aa..4a698630268 100644 --- a/megatron/core/ssm/ops/ssd_chunk_state.py +++ b/megatron/core/ssm/ops/ssd_chunk_state.py @@ -51,11 +51,12 @@ def _chunk_cumsum_fwd_kernel( dt_bias_ptr, dt_out_ptr, dA_cumsum_ptr, - cu_chunk_seqlens_ptr, + chunk_offsets_ptr, # Matrix dimension seqlen, nheads: tl.constexpr, chunk_size: tl.constexpr, + HAS_CHUNK_STARTS: tl.constexpr, dt_min: tl.constexpr, dt_max: tl.constexpr, # Strides @@ -80,8 +81,12 @@ def _chunk_cumsum_fwd_kernel( pid_c = tl.program_id(axis=0).to(tl.int64) pid_h = tl.program_id(axis=1) - chunk_seqlen_start = tl.load(cu_chunk_seqlens_ptr + pid_c) - chunk_seqlen_end = tl.load(cu_chunk_seqlens_ptr + pid_c + 1) + chunk_seqlen_start = tl.load(chunk_offsets_ptr + pid_c) + if HAS_CHUNK_STARTS: + # Fixed windows need only a start offset. + chunk_seqlen_end = chunk_seqlen_start + chunk_size + else: + chunk_seqlen_end = tl.load(chunk_offsets_ptr + pid_c + 1) dt_ptr += chunk_seqlen_start * stride_dt_seqlen dt_out_ptr += pid_c * stride_dt_out_chunk @@ -179,7 +184,8 @@ def _chunk_state_fwd_kernel( states_ptr, dt_ptr, dA_cumsum_ptr, - cu_chunk_seqlens_ptr, + chunk_offsets_ptr, + chunk_flags_ptr, # Matrix dimensions hdim: tl.constexpr, dstate: tl.constexpr, @@ -204,6 +210,7 @@ def _chunk_state_fwd_kernel( stride_dA_cs_chunk: tl.int64, stride_dA_cs_csize: tl.constexpr, # Meta-parameters + HAS_CHUNK_FLAGS: tl.constexpr, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, @@ -213,8 +220,16 @@ def _chunk_state_fwd_kernel( num_pid_n = tl.cdiv(dstate, BLOCK_SIZE_N) pid_m = tl.program_id(axis=0) // num_pid_n pid_n = tl.program_id(axis=0) % num_pid_n - chunk_seqlen_start = tl.load(cu_chunk_seqlens_ptr + pid_c) - chunk_seqlen_end = tl.load(cu_chunk_seqlens_ptr + pid_c + 1) + if HAS_CHUNK_FLAGS: + # Only completed chunks need a boundary state. + if tl.load(chunk_flags_ptr + pid_c) == 0: + return + chunk_seqlen_start = tl.load(chunk_offsets_ptr + pid_c) + if HAS_CHUNK_FLAGS: + # Chunk flags are paired with fixed-window starts. + chunk_seqlen_end = chunk_seqlen_start + chunk_size + else: + chunk_seqlen_end = tl.load(chunk_offsets_ptr + pid_c + 1) b_ptr += chunk_seqlen_start * stride_b_seqlen + (pid_h // nheads_ngroups_ratio) * stride_b_head x_ptr += chunk_seqlen_start * stride_x_seqlen + pid_h * stride_x_head dt_ptr += pid_c * stride_dt_chunk + pid_h * stride_dt_head @@ -277,12 +292,18 @@ def _chunk_cumsum_fwd( dt_bias=None, dt_softplus=False, dt_limit=(0.0, float("inf")), + chunk_starts=None, ): seqlen, nheads = dt.shape assert A.shape == (nheads,) if dt_bias is not None: assert dt_bias.shape == (nheads,) - nchunks = cu_chunk_seqlens.shape[0] - 1 + if chunk_starts is not None: + chunk_offsets = chunk_starts + nchunks = chunk_starts.shape[0] + else: + chunk_offsets = cu_chunk_seqlens + nchunks = cu_chunk_seqlens.shape[0] - 1 dt_out = torch.empty(nheads, nchunks, chunk_size, device=dt.device, dtype=torch.float32) dA_cumsum = torch.empty(nheads, nchunks, chunk_size, device=dt.device, dtype=torch.float32) grid_chunk_cs = lambda META: (nchunks, triton.cdiv(nheads, META["BLOCK_SIZE_H"])) @@ -293,7 +314,7 @@ def _chunk_cumsum_fwd( dt_bias_ptr=dt_bias, dt_out_ptr=dt_out, dA_cumsum_ptr=dA_cumsum, - cu_chunk_seqlens_ptr=cu_chunk_seqlens, + chunk_offsets_ptr=chunk_offsets, seqlen=seqlen, nheads=nheads, chunk_size=chunk_size, @@ -309,6 +330,7 @@ def _chunk_cumsum_fwd( stride_dA_cs_head=dA_cumsum.stride(0), stride_dA_cs_chunk=dA_cumsum.stride(1), stride_dA_cs_csize=dA_cumsum.stride(2), + HAS_CHUNK_STARTS=chunk_starts is not None, DT_SOFTPLUS=dt_softplus, HAS_DT_BIAS=dt_bias is not None, BLOCK_SIZE_CHUNK=triton.next_power_of_2(chunk_size), @@ -316,7 +338,25 @@ def _chunk_cumsum_fwd( return dA_cumsum, dt_out -def _chunk_state_fwd(B, x, dt, dA_cumsum, cu_chunk_seqlens, states=None, states_in_fp32=True): +def _chunk_state_fwd( + B, + x, + dt, + dA_cumsum, + cu_chunk_seqlens, + states=None, + states_in_fp32=True, + chunk_flags=None, + chunk_starts=None, +): + has_chunk_flags = chunk_flags is not None + assert (chunk_starts is not None) == has_chunk_flags, ( + "chunk_flags and chunk_starts must be provided together" + ) + if has_chunk_flags: + chunk_offsets = chunk_starts + else: + chunk_offsets = cu_chunk_seqlens seqlen, nheads, headdim = x.shape _, nchunks, chunk_size = dt.shape _, ngroups, dstate = B.shape @@ -345,7 +385,8 @@ def _chunk_state_fwd(B, x, dt, dA_cumsum, cu_chunk_seqlens, states=None, states_ states_ptr=states, dt_ptr=dt, dA_cumsum_ptr=dA_cumsum, - cu_chunk_seqlens_ptr=cu_chunk_seqlens, + chunk_offsets_ptr=chunk_offsets, + chunk_flags_ptr=chunk_flags, hdim=headdim, dstate=dstate, chunk_size=chunk_size, @@ -367,6 +408,7 @@ def _chunk_state_fwd(B, x, dt, dA_cumsum, cu_chunk_seqlens, states=None, states_ stride_dA_cs_head=dA_cumsum.stride(0), stride_dA_cs_chunk=dA_cumsum.stride(1), stride_dA_cs_csize=dA_cumsum.stride(2), + HAS_CHUNK_FLAGS=has_chunk_flags, ) return states diff --git a/megatron/core/ssm/ops/ssd_combined.py b/megatron/core/ssm/ops/ssd_combined.py index 4fcee98b13e..bc227ccca7c 100644 --- a/megatron/core/ssm/ops/ssd_combined.py +++ b/megatron/core/ssm/ops/ssd_combined.py @@ -160,6 +160,110 @@ def _mamba_chunk_scan_combined_fwd( return final_states +def mamba_chunk_scan_decode_rows( + x, + dt, + A, + B, + C, + chunk_size, + chunk_starts, + slots, + target_rows, + chunk_flags, + initial_states, + out, + D=None, + dt_bias=None, + dt_softplus=False, + dt_limit=(0.0, float("inf")), +): + """Row-gated chunk scan for batch-invariant single-token decode. + + Same 5-kernel pipeline as the full varlen scan, run directly over the + persistent per-slot buffers: chunk c is the fixed chunk_size window at + chunk_starts[c], and every chunk is its own sequence starting from + initial_states[slots[c]]. The kernels are gated to what a decode step + actually consumes: bmm and the scan compute only the block containing + target_rows[c], and the chunk-state matmul runs only where chunk_flags + is set (the slot crosses its boundary, the one step its state is read). + The blocks that do run execute the same instructions as the ungated + kernels, so the outputs match a full scan bitwise. + + Args: + x/dt/B/C: flattened persistent buffers, (num_rows * chunk_size, ...). + chunk_starts: (nseq,) int32, window offset per chunk + (slot * chunk_size for per-slot buffers). + slots: (nseq,) int32, live-cache row containing each chunk's incoming + state. May repeat for padding entries. + target_rows: (nseq,) int32, the only output row read per chunk. + chunk_flags: (nseq,), nonzero where the slot crosses its boundary. + initial_states: (num_states, nheads, headdim, dstate), the engine's + live SSM cache. Crossing chunks update it in place. + out: (nseq, nheads, headdim), receives each chunk's target row. + """ + dA_cumsum, dt = _chunk_cumsum_fwd( + dt, + A, + chunk_size, + None, + dt_bias=dt_bias, + dt_softplus=dt_softplus, + dt_limit=dt_limit, + chunk_starts=chunk_starts, + ) + # Only boundary-crossing chunks produce a state; state passing masks the rest. + states = _chunk_state_fwd( + B, + x, + dt, + dA_cumsum, + None, + states_in_fp32=True, + chunk_flags=chunk_flags, + chunk_starts=chunk_starts, + ) + CB = _bmm_chunk_fwd( + C, + B, + chunk_size, + None, + output_dtype=torch.float32, + target_rows=target_rows, + chunk_starts=chunk_starts, + ) + # The scan must run before state passing: the snapshot below overwrites + # crossing slots' rows in ssm_state, and the scan reads initial_states + # from that same cache. The scan never reads state passing's output in + # decode mode (every chunk is its own sequence), so `states` is just a + # shape-valid placeholder for the unused carried-state pointer. + _chunk_scan_fwd( + CB, + x, + dt, + dA_cumsum, + C, + states, + None, + out, + slots, + D=D, + initial_states=initial_states, + target_rows=target_rows, + chunk_starts=chunk_starts, + ) + _state_passing_fwd( + states.flatten(-2), + dA_cumsum, + None, + initial_states=initial_states.flatten(-2), + seq_idx=slots, + dst_states=initial_states.flatten(-2), + dst_indices=slots, + dst_flags=chunk_flags, + ) + + def mamba_chunk_scan_combined_varlen( x, dt, diff --git a/megatron/core/ssm/ops/ssd_state_passing.py b/megatron/core/ssm/ops/ssd_state_passing.py index 65b81a0ec31..bada474136c 100644 --- a/megatron/core/ssm/ops/ssd_state_passing.py +++ b/megatron/core/ssm/ops/ssd_state_passing.py @@ -32,6 +32,9 @@ def _state_passing_fwd_kernel( initstates_ptr, seq_idx_ptr, cu_chunk_seqlens_ptr, + dst_states_ptr, + dst_indices_ptr, + dst_flags_ptr, # Matrix dimensions dim: tl.constexpr, nchunks, @@ -51,8 +54,12 @@ def _state_passing_fwd_kernel( stride_initstates_head: tl.int64, stride_initstates_dim: tl.constexpr, stride_seq_idx_chunk: tl.constexpr, + stride_dst_batch: tl.int64, + stride_dst_head: tl.int64, + stride_dst_dim: tl.constexpr, # Meta-parameters HAS_INITSTATES: tl.constexpr, + HAS_DST_STATES: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): pid_h = tl.program_id(axis=1) @@ -77,11 +84,23 @@ def _state_passing_fwd_kernel( prev_seq_idx = 0 for c in range(nchunks): - new_states = tl.load(states_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) + if HAS_DST_STATES: + dst_flag = tl.load(dst_flags_ptr + c) != 0 + else: + dst_flag = True + # Unflagged destination chunks have no chunk state. + new_states = tl.load(states_ptrs, mask=(offs_m < dim) & dst_flag, other=0.0).to( + tl.float32 + ) dA_cs = tl.load(dA_cs_ptr).to(tl.float32) seq_idx = tl.load(seq_idx_ptr + c * stride_seq_idx_chunk) + if HAS_DST_STATES: + # Destination chunks start from their indexed initial state. + is_new_seq = True + else: + is_new_seq = prev_seq_idx != seq_idx # we have started a new sequence - if prev_seq_idx != seq_idx: + if is_new_seq: if HAS_INITSTATES: initstates_ptrs = ( initstates_ptr @@ -95,7 +114,20 @@ def _state_passing_fwd_kernel( prev_seq_idx = seq_idx states = tl.exp(dA_cs) * states + new_states - tl.store(out_ptrs, states, mask=offs_m < dim) + if not HAS_DST_STATES: + tl.store(out_ptrs, states, mask=offs_m < dim) + + if HAS_DST_STATES: + # Commit completed chunks directly to the state cache. + if dst_flag: + dst_idx = tl.load(dst_indices_ptr + c).to(tl.int64) + dst_ptrs = ( + dst_states_ptr + + dst_idx * stride_dst_batch + + pid_h * stride_dst_head + + offs_m * stride_dst_dim + ) + tl.store(dst_ptrs, states, mask=offs_m < dim) states_ptrs += stride_states_chunk dA_cs_ptr += stride_dA_cs_chunk @@ -103,20 +135,48 @@ def _state_passing_fwd_kernel( def _state_passing_fwd( - states, dA_cumsum, cu_chunk_seqlens, seq_idx, initial_states=None, out_dtype=None + states, + dA_cumsum, + cu_chunk_seqlens, + seq_idx, + initial_states=None, + out_dtype=None, + dst_states=None, + dst_indices=None, + dst_flags=None, ): + """ + dst_states/dst_indices/dst_flags write flagged boundary states directly to + dst_states without allocating an output tensor. + """ nchunks, nheads, dim = states.shape chunk_size = dA_cumsum.shape[-1] assert dA_cumsum.shape == (nheads, nchunks, chunk_size) seqlen = seq_idx.shape[-1] - out_dtype = states.dtype if out_dtype is None else out_dtype - out = torch.empty((nchunks, nheads, dim), device=states.device, dtype=out_dtype) + has_dst = dst_states is not None + assert (dst_indices is not None) == has_dst and (dst_flags is not None) == has_dst, ( + "dst_states, dst_indices, and dst_flags must be provided together" + ) + if not has_dst: + out_dtype = states.dtype if out_dtype is None else out_dtype + out = torch.empty((nchunks, nheads, dim), device=states.device, dtype=out_dtype) + out_strides = out.stride() + else: + out = states + out_strides = (0, 0, 0) initial_states_strides = ( (initial_states.stride(0), initial_states.stride(1), initial_states.stride(2)) if initial_states is not None else (0, 0, 0) ) + if has_dst: + assert dst_states.shape[1] == nheads and dst_states.shape[2] == dim + dst_strides = ( + (dst_states.stride(0), dst_states.stride(1), dst_states.stride(2)) + if has_dst + else (0, 0, 0) + ) grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE"]), nheads) with torch.cuda.device(states.device.index): @@ -127,6 +187,9 @@ def _state_passing_fwd( initstates_ptr=initial_states, seq_idx_ptr=seq_idx, cu_chunk_seqlens_ptr=cu_chunk_seqlens, + dst_states_ptr=dst_states, + dst_indices_ptr=dst_indices, + dst_flags_ptr=dst_flags, dim=dim, nchunks=nchunks, seqlen=seqlen if seq_idx is not None else 0, @@ -134,9 +197,9 @@ def _state_passing_fwd( stride_states_chunk=states.stride(0), stride_states_head=states.stride(1), stride_states_dim=states.stride(2), - stride_out_chunk=out.stride(0), - stride_out_head=out.stride(1), - stride_out_dim=out.stride(2), + stride_out_chunk=out_strides[0], + stride_out_head=out_strides[1], + stride_out_dim=out_strides[2], stride_dA_cs_head=dA_cumsum.stride(0), stride_dA_cs_chunk=dA_cumsum.stride(1), stride_dA_cs_csize=dA_cumsum.stride(2), @@ -144,6 +207,10 @@ def _state_passing_fwd( stride_initstates_head=initial_states_strides[1], stride_initstates_dim=initial_states_strides[2], stride_seq_idx_chunk=seq_idx.stride(0), + stride_dst_batch=dst_strides[0], + stride_dst_head=dst_strides[1], + stride_dst_dim=dst_strides[2], HAS_INITSTATES=initial_states is not None, + HAS_DST_STATES=has_dst, ) - return out + return None if has_dst else out diff --git a/megatron/core/tensor_parallel/inference_layers.py b/megatron/core/tensor_parallel/inference_layers.py index 14ac28fbefa..68ab2bf382b 100644 --- a/megatron/core/tensor_parallel/inference_layers.py +++ b/megatron/core/tensor_parallel/inference_layers.py @@ -24,6 +24,10 @@ gather_from_tensor_model_parallel_region, reduce_scatter_to_sequence_parallel_region, ) +from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( + is_batch_invariant_mode_enabled, + rmsnorm_batch_invariant, +) from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.utils import get_tensor_model_parallel_group_if_none @@ -41,6 +45,9 @@ def _te_rms_norm_kernel(x: torch.Tensor, weight: torch.Tensor, eps: float): + # Use the same RMSNorm kernel as the training recompute. + if is_batch_invariant_mode_enabled(): + return rmsnorm_batch_invariant(x, weight, eps).to(x.dtype) x_shape = x.shape x = x.view(-1, x.size(-1)) out, _, _ = tex.rmsnorm_fwd( @@ -405,11 +412,14 @@ def _matmul_reduce_scatter(self, x, residual=None): # RS requires bf16 (hardware multimem reduce is bf16-only). # Check the matmul output shape: if it is NVLS-eligible, the RS output # (world_size times smaller on dim 0) is too. + # TP sequence-parallel RS: use NCCL in batch-invariant mode to match + # the training reduction path. This does not affect MoE EP NVLS. can_use_nvls = ( self.triton_nvls_kernels_allowed and x.dtype == torch.bfloat16 and are_tensors_nvls_eligible(x) and symm_mem_buffer["handle"] is not None + and not is_batch_invariant_mode_enabled() ) if can_use_nvls: @@ -532,7 +542,13 @@ def inference_reduce_scatter_to_sequence_parallel_region( config, 'inference_disable_triton_nvls_kernels', False ) - if triton_nvls_kernels_allowed and SymmetricMemoryManager.is_initialized("tp"): + # TP sequence-parallel RS: use NCCL in batch-invariant mode to match + # training. This does not affect MoE EP NVLS. + if ( + triton_nvls_kernels_allowed + and SymmetricMemoryManager.is_initialized("tp") + and not is_batch_invariant_mode_enabled() + ): buf = SymmetricMemoryManager.get_buffer("tp", process_group=tp_group) symm_mem_buffer = buf.maybe_get_tensor(list(x.size()), dtype=x.dtype) diff --git a/megatron/core/transformer/custom_layers/batch_invariant_kernels.py b/megatron/core/transformer/custom_layers/batch_invariant_kernels.py index 6b4311fe540..a14489c5efb 100644 --- a/megatron/core/transformer/custom_layers/batch_invariant_kernels.py +++ b/megatron/core/transformer/custom_layers/batch_invariant_kernels.py @@ -6,6 +6,7 @@ import contextlib import importlib import importlib.util +import inspect import logging from collections import namedtuple from collections.abc import Callable @@ -28,11 +29,29 @@ tl = MagicMock() HAVE_TRITON = False +try: + import deep_gemm + + HAVE_DEEPGEMM_BF16 = all( + hasattr(deep_gemm, name) + for name in ( + "m_grouped_bf16_gemm_nt_contiguous", + "k_grouped_bf16_gemm_tn_contiguous", + "bf16_gemm_nn", + ) + ) +except ImportError: + deep_gemm = None + HAVE_DEEPGEMM_BF16 = False + __all__ = [ "set_batch_invariant_mode", "is_batch_invariant_mode_enabled", "disable_batch_invariant_mode", "enable_batch_invariant_mode", + "grouped_gemm_batch_invariant", + "grouped_gemm_batch_invariant_alignment", + "HAVE_DEEPGEMM_BF16", ] @@ -478,13 +497,48 @@ def mean_dim( return output +# Kernel backend for mm / addmm. Production uses DeepGEMM; the Triton option +# remains available to tests that exercise non-bf16 operators. +# "deepgemm" (default): DeepGEMM `bf16_gemm_nn` — bitwise-identical to +# `torch.mm`. Requires bf16 CUDA inputs on Hopper/Blackwell. +# "triton": batch-invariant Triton `matmul_persistent` — works on any CUDA +# device with bf16/fp16/fp32. Has small rounding drift vs `torch.mm`. +_BATCH_INVARIANT_BACKENDS = ("deepgemm", "triton") +_BATCH_INVARIANT_BACKEND: str = "deepgemm" + + +def _mm_deepgemm(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """`a @ b` via DeepGEMM `bf16_gemm_nn`. Both inputs are row-major. + + Bitwise-identical to `torch.mm` on Hopper/Blackwell, deterministic across + runs, batch-invariant. + """ + if a.dtype != torch.bfloat16: + raise RuntimeError( + f"The DeepGEMM batch-invariant backend requires bf16 inputs " + f"(got {a.dtype}); use backend='triton' for fp16/fp32." + ) + M = a.shape[0] + N = b.shape[1] + d = torch.empty(M, N, device=a.device, dtype=a.dtype) + deep_gemm.bf16_gemm_nn(a, b, d) + return d + + def mm_batch_invariant(a, b): - """Batch-invariant replacement for `aten::mm` using a persistent matmul kernel.""" + """Batch-invariant replacement for `aten::mm`.""" + if _BATCH_INVARIANT_BACKEND == "deepgemm": + return _mm_deepgemm(a, b) return matmul_persistent(a, b) def addmm_batch_invariant(bias, a, b): - """Batch-invariant replacement for `aten::addmm` using a persistent matmul kernel.""" + """Batch-invariant replacement for `aten::addmm`.""" + if _BATCH_INVARIANT_BACKEND == "deepgemm": + out = _mm_deepgemm(a, b) + if bias is not None: + out = out + bias + return out return matmul_persistent(a, b, bias=bias) @@ -525,6 +579,7 @@ def get_batch_invariant_attention_block_size() -> AttentionBlockSize: _MEG_TE_GENERAL_GEMM_ORIG = None _TE_RMSNORM_FUNC_ORIGS: Dict[str, Any] = {} _TE_GEMM_FUNC_ORIGS: Dict[str, Any] = {} +_TE_GROUPED_GEMM_FUNC_ORIGS: Dict[str, Any] = {} def _import_module_if_available(name: str): @@ -617,6 +672,335 @@ def _patched(*args, **kwargs): _TE_RMSNORM_FUNC_ORIGS[name] = orig setattr(te_layernorm_mod, name, _make_rmsnorm_patched(orig)) + # Patch TE.general_grouped_gemm at every known import site so that + # TEGroupedMLP (forward + dgrad + wgrad) goes through DeepGEMM in bf16. + _te_patch_general_grouped_gemm() + + +def _te_patch_general_grouped_gemm() -> None: + """Replace TE.general_grouped_gemm with a batch-invariant dispatcher. + + Patches the symbol at three import sites — the consumer module + (transformer_engine.pytorch.module.grouped_linear), the package-level + re-export (transformer_engine.pytorch.cpp_extensions), and the source + module (transformer_engine.pytorch.cpp_extensions.gemm). Stores originals + in _TE_GROUPED_GEMM_FUNC_ORIGS so the unpatch can restore them. + """ + te_grouped_linear_mod = _import_module_if_available( + "transformer_engine.pytorch.module.grouped_linear" + ) + if te_grouped_linear_mod is not None and hasattr(te_grouped_linear_mod, "general_grouped_gemm"): + key = "module.grouped_linear.general_grouped_gemm" + if key not in _TE_GROUPED_GEMM_FUNC_ORIGS: + _TE_GROUPED_GEMM_FUNC_ORIGS[key] = te_grouped_linear_mod.general_grouped_gemm + te_grouped_linear_mod.general_grouped_gemm = _te_general_grouped_gemm_patched + + te_cpp = _import_module_if_available("transformer_engine.pytorch.cpp_extensions") + if te_cpp is not None and hasattr(te_cpp, "general_grouped_gemm"): + key = "cpp_extensions.general_grouped_gemm" + if key not in _TE_GROUPED_GEMM_FUNC_ORIGS: + _TE_GROUPED_GEMM_FUNC_ORIGS[key] = te_cpp.general_grouped_gemm + te_cpp.general_grouped_gemm = _te_general_grouped_gemm_patched + + te_cpp_gemm = _import_module_if_available("transformer_engine.pytorch.cpp_extensions.gemm") + if te_cpp_gemm is not None and hasattr(te_cpp_gemm, "general_grouped_gemm"): + key = "cpp_extensions.gemm.general_grouped_gemm" + if key not in _TE_GROUPED_GEMM_FUNC_ORIGS: + _TE_GROUPED_GEMM_FUNC_ORIGS[key] = te_cpp_gemm.general_grouped_gemm + te_cpp_gemm.general_grouped_gemm = _te_general_grouped_gemm_patched + + +def _te_unpatch_general_grouped_gemm() -> None: + """Restore the originals captured by _te_patch_general_grouped_gemm.""" + module_paths = { + "module.grouped_linear.general_grouped_gemm": ( + "transformer_engine.pytorch.module.grouped_linear", + "general_grouped_gemm", + ), + "cpp_extensions.general_grouped_gemm": ( + "transformer_engine.pytorch.cpp_extensions", + "general_grouped_gemm", + ), + "cpp_extensions.gemm.general_grouped_gemm": ( + "transformer_engine.pytorch.cpp_extensions.gemm", + "general_grouped_gemm", + ), + } + for key, (mod_name, attr) in module_paths.items(): + if key not in _TE_GROUPED_GEMM_FUNC_ORIGS: + continue + mod = _import_module_if_available(mod_name) + if mod is not None and hasattr(mod, attr): + setattr(mod, attr, _TE_GROUPED_GEMM_FUNC_ORIGS[key]) + _TE_GROUPED_GEMM_FUNC_ORIGS.pop(key, None) + + +def _get_original_te_grouped_gemm(): + for key in ( + "module.grouped_linear.general_grouped_gemm", + "cpp_extensions.general_grouped_gemm", + "cpp_extensions.gemm.general_grouped_gemm", + ): + orig = _TE_GROUPED_GEMM_FUNC_ORIGS.get(key) + if orig is not None: + return orig + return None + + +def _original_te_grouped_gemm_has_quantization_params(orig) -> bool: + try: + return "quantization_params" in inspect.signature(orig).parameters + except (TypeError, ValueError): + return False + + +def _call_original_te_grouped_gemm( + orig, + A, + B, + out, + quantization_params, + out_dtype, + *, + layout, + m_splits, + gelu, + grad, + accumulate, + bias, + use_bias, + use_split_accumulator, + D_dtype, + single_output, +): + kwargs = dict( + layout=layout, + m_splits=m_splits, + gelu=gelu, + grad=grad, + accumulate=accumulate, + bias=bias, + use_bias=use_bias, + use_split_accumulator=use_split_accumulator, + D_dtype=D_dtype, + single_output=single_output, + ) + if _original_te_grouped_gemm_has_quantization_params(orig): + return orig(A, B, out, quantization_params, out_dtype, **kwargs) + return orig(A, B, out, out_dtype, **kwargs) + + +def _is_bf16_grouped_path(A, B, quantization_params, gelu: bool) -> bool: + """Decide if TE's general_grouped_gemm call can be served by DeepGEMM bf16.""" + if gelu: + return False + if not HAVE_DEEPGEMM_BF16: + return False + if not (isinstance(A, list) and isinstance(B, list)): + return False + if len(A) != len(B) or len(A) == 0: + return False + if quantization_params is not None and any(q is not None for q in quantization_params): + return False + for t in (*A, *B): + if not isinstance(t, torch.Tensor): + return False + if t.dtype != torch.bfloat16: + return False + return True + + +def _te_general_grouped_gemm_patched( + A, + B, + out, + quantization_params=None, + out_dtype=None, + layout: str = "TN", + m_splits=None, + gelu: bool = False, + grad: bool = False, + accumulate: bool = False, + bias=None, + use_bias: bool = False, + use_split_accumulator: bool = False, + D_dtype=None, + single_output: bool = False, +): + """Batch-invariant replacement for TE general_grouped_gemm. + + Dispatches by (layout, single_output, grad) to forward / dgrad / wgrad + implementations backed by DeepGEMM. Falls back to TE's original for any + case we cannot guarantee batch-invariant: quantized inputs, gelu fusion, + non-bf16 dtypes, or unsupported (layout, mode) combinations. + """ + # TE versions differ here: + # old: general_grouped_gemm(A, B, out, out_dtype, ...) + # new: general_grouped_gemm(A, B, out, quantization_params, out_dtype, ...) + if out_dtype is None and isinstance(quantization_params, torch.dtype): + out_dtype = quantization_params + quantization_params = None + + if not _is_bf16_grouped_path(A, B, quantization_params, gelu): + orig = _get_original_te_grouped_gemm() + if orig is None: + raise RuntimeError( + "Batch-invariant grouped GEMM patch was invoked but no original " + "TE general_grouped_gemm was captured; patching order issue." + ) + return _call_original_te_grouped_gemm( + orig, + A, + B, + out, + quantization_params, + out_dtype, + layout=layout, + m_splits=m_splits, + gelu=gelu, + grad=grad, + accumulate=accumulate, + bias=bias, + use_bias=use_bias, + use_split_accumulator=use_split_accumulator, + D_dtype=D_dtype, + single_output=single_output, + ) + + # Dispatch by TE's call convention. + # In TE _GroupedLinear: + # forward -> layout="TN", single_output=True, grad=False (A=weights, B=inputmats) + # dgrad -> layout="NN", single_output=True, grad=True (A=weights, B=grad_y) + # wgrad -> layout="NT", single_output=False, grad=True (A=inputmats, B=grad_y) + if single_output and layout == "TN" and not grad: + return _batch_invariant_te_grouped_forward(A, B, out, m_splits, bias, use_bias, accumulate) + if single_output and layout == "NN" and grad: + return _batch_invariant_te_grouped_dgrad(A, B, out, m_splits, accumulate) + if (not single_output) and layout == "NT" and grad: + return _batch_invariant_te_grouped_wgrad(A, B, out, m_splits, use_bias, accumulate) + # Unknown TE call shape — defer to the original. + orig = _get_original_te_grouped_gemm() + if orig is None: + raise RuntimeError( + "Batch-invariant grouped GEMM patch was invoked but no original " + "TE general_grouped_gemm was captured; patching order issue." + ) + return _call_original_te_grouped_gemm( + orig, + A, + B, + out, + quantization_params, + out_dtype, + layout=layout, + m_splits=m_splits, + gelu=gelu, + grad=grad, + accumulate=accumulate, + bias=bias, + use_bias=use_bias, + use_split_accumulator=use_split_accumulator, + D_dtype=D_dtype, + single_output=single_output, + ) + + +def _stack_weights_for_deepgemm(weights: List[torch.Tensor]) -> torch.Tensor: + """Stack a per-expert weight list into a contiguous [E, N, K] buffer.""" + if not weights: + return torch.empty(0) + return torch.stack([w.contiguous() for w in weights], dim=0) + + +def _batch_invariant_te_grouped_forward(A, B, out, m_splits, bias, use_bias, accumulate): + """TE forward: Y = X @ W^T per expert, then optional bias. + + A = weights: List[Tensor[N, K]] + B = inputmats: List[Tensor[m_i, K]] + out: [single Tensor[M_total, N]] (single_output=True) + """ + assert not accumulate, "Forward never accumulates" + assert len(out) == 1, "single_output=True forward expects a single out tensor" + out_buf = out[0] + w_stack = _stack_weights_for_deepgemm(A) + x_cat = torch.cat([b.contiguous() for b in B], dim=0) + m_total = x_cat.shape[0] + m_indices = _m_splits_to_m_indices(m_splits, x_cat.device, m_total) + + y = _bf16_grouped_gemm_contiguous(x_cat, w_stack, m_indices, m_splits) + if use_bias and bias is not None: + offset = 0 + for i, m in enumerate(m_splits): + if m == 0: + continue + b_i = bias[i] if i < len(bias) else None + if b_i is not None and b_i.numel() > 0: + y[offset : offset + m] = y[offset : offset + m] + b_i.to(y.dtype) + offset += m + + if y.dtype != out_buf.dtype: + y = y.to(out_buf.dtype) + out_buf.copy_(y) + # TE's contract: (out_list, bias_or_grad_bias, gelu_input) + return out, bias if use_bias else [None] * len(A), None + + +def _batch_invariant_te_grouped_dgrad(A, B, out, m_splits, accumulate): + """TE dgrad: dX = dY @ W per expert. + + A = weights: List[Tensor[N, K]] + B = grad_y_per_expert: List[Tensor[m_i, N]] + out: [single Tensor[M_total, K]] (single_output=True) + """ + assert not accumulate, "Dgrad never accumulates" + assert len(out) == 1 + out_buf = out[0] + w_stack = _stack_weights_for_deepgemm(A) + dy_cat = torch.cat([b.contiguous() for b in B], dim=0) + m_total = dy_cat.shape[0] + m_indices = _m_splits_to_m_indices(m_splits, dy_cat.device, m_total) + # NT call interprets B as [E, out_dim, in_dim]; for dgrad we need W as [E, K, N] + w_kn = w_stack.transpose(1, 2).contiguous() + dx = _bf16_grouped_gemm_contiguous(dy_cat, w_kn, m_indices, m_splits) + if dx.dtype != out_buf.dtype: + dx = dx.to(out_buf.dtype) + out_buf.copy_(dx) + return out, [None] * len(A), None + + +def _batch_invariant_te_grouped_wgrad(A, B, out, m_splits, use_bias, accumulate): + """TE wgrad: dW[g] = dY[g]^T @ X[g], plus optional dbias[g] = sum(dY[g], dim=0). + + A = inputmats: List[Tensor[m_i, K]] + B = grad_y: List[Tensor[m_i, N]] + out: List[Tensor[N, K]] per expert (single_output=False) + """ + E = len(m_splits) + x_cat = torch.cat([a.contiguous() for a in A], dim=0) + dy_cat = torch.cat([b.contiguous() for b in B], dim=0) + m_total = x_cat.shape[0] + assert sum(m_splits) == m_total + dw_stack = _bf16_grouped_gemm_wgrad_contiguous(dy_cat, x_cat, m_splits) + + grad_bias = [None] * E + if use_bias: + offset = 0 + for i, m in enumerate(m_splits): + if m > 0: + grad_bias[i] = dy_cat[offset : offset + m].sum(dim=0) + offset += m + + for i in range(E): + target = out[i] + contrib = dw_stack[i] + if contrib.dtype != target.dtype: + contrib = contrib.to(target.dtype) + if accumulate: + target.add_(contrib) + else: + target.copy_(contrib) + return out, grad_bias, None + def _te_unpatch_for_batch_invariant(): """Restore original Transformer Engine functions if they were patched.""" @@ -690,6 +1074,9 @@ def _te_unpatch_for_batch_invariant(): else: _TE_GEMM_FUNC_ORIGS.pop(key, None) + # Restore TE general_grouped_gemm at every patched import site. + _te_unpatch_general_grouped_gemm() + def _extract_te_gemm_args(args: tuple, kwargs: Dict[str, Any]): """Utility to parse TE general_gemm flexible signature. @@ -739,7 +1126,7 @@ def forward( opA = opA.reshape(-1, opA.shape[-1]) elif opA.dim() < 2: raise ValueError(f"opA has insufficient dimensions: {opA.shape}") - assert opA.dim() == 2, f"opA must be 2D for matmul_persistent, got shape {opA.shape}" + assert opA.dim() == 2, f"opA must be 2D, got shape {opA.shape}" # Flatten all leading dims of opB except the last feature dim to match TE behavior if opB.dim() >= 2: @@ -751,7 +1138,7 @@ def forward( opB_2d = opB # Perform GEMM: (N_total, K) @ (K, O) -> (N_total, O) - base_2d = matmul_persistent(opB_2d, opA, bias=None) + base_2d = mm_batch_invariant(opB_2d, opA) # Reshape back to original leading dims with output features at the end out = base_2d.reshape(*leading_shape, base_2d.shape[-1]) @@ -945,6 +1332,256 @@ def rmsnorm_batch_invariant(x: torch.Tensor, weight: torch.Tensor, eps: float) - return BatchInvariantRMSNormFn.apply(x, weight, eps, False) +# --------------------------------------------------------------------------- +# Batch-invariant grouped GEMM (DeepGEMM-backed). Used by MoE so that training +# (TEGroupedMLP via patched TE.general_grouped_gemm) and inference +# (InferenceGroupedMLP via patched _bf16_grouped_mm) produce bitwise-identical +# outputs for the same inputs. This is what gives RL rollout==train log-prob +# parity for MoE models. +# --------------------------------------------------------------------------- + + +def _require_deepgemm_bf16(op: str) -> None: + """Raise a clear error if DeepGEMM bf16 grouped bindings are unavailable.""" + if not HAVE_DEEPGEMM_BF16: + raise RuntimeError( + f"Batch-invariant grouped GEMM ({op}) requires DeepGEMM with bf16 bindings. " + "Install via `uv pip install -e .[batch_invariant]` (pins a DeepGEMM commit " + "that exposes m_grouped_bf16_gemm_nt_contiguous), or disable " + "transformer_config.batch_invariant_mode for MoE models." + ) + + +def _offs_to_m_indices(offs: torch.Tensor, m_total: int) -> torch.Tensor: + """Convert inclusive cumulative per-expert offsets to per-row expert ids. + + offs: int32 [num_experts] inclusive offsets — offs[i] is the (exclusive) end + row of expert i in the contiguous M dimension. Equivalently, offs[i] is + the start of expert i+1. + Returns: int32 [m_total] m_indices[r] = expert id for row r. Rows past offs[-1] + (post-padding tail when m_total > offs[-1]) get -1; DeepGEMM skips + those rows. + """ + rows = torch.arange(m_total, device=offs.device, dtype=torch.int32) + # For row r, expert id = bisect_right(offs, r). torch.searchsorted is deterministic. + m_indices = torch.searchsorted(offs, rows, right=True).to(torch.int32) + n_used = offs[-1].to(torch.int32) + m_indices = torch.where(rows < n_used, m_indices, torch.full_like(m_indices, -1)) + return m_indices + + +def _m_splits_to_m_indices(m_splits: List[int], device: torch.device, m_total: int) -> torch.Tensor: + """Convert TE per-expert token counts (List[int]) to int32 [m_total] m_indices. + + No padding rows in TE training path — sum(m_splits) == m_total exactly. + """ + assert sum(m_splits) == m_total, f"m_splits sum ({sum(m_splits)}) != m_total ({m_total})" + parts = [ + torch.full((n,), i, device=device, dtype=torch.int32) + for i, n in enumerate(m_splits) + if n > 0 + ] + if not parts: + return torch.empty(0, device=device, dtype=torch.int32) + return torch.cat(parts, dim=0) + + +# DeepGEMM's contiguous M-grouped and K-grouped bf16 GEMMs require each +# per-expert block on the grouped axis to be a multiple of this alignment +# (typically 128 on SM90/SM100). We pad inputs to satisfy this, then strip the +# padding from the output. Padding rows are zeros (correct identity for the +# reduction sum) and tagged with m_indices=-1 for the M-grouped case so the +# kernel can skip them in store. +_DEEPGEMM_M_ALIGNMENT: Optional[int] = None + + +def _deepgemm_m_alignment() -> int: + """Lazily fetch DeepGEMM's required per-expert block alignment.""" + global _DEEPGEMM_M_ALIGNMENT + if _DEEPGEMM_M_ALIGNMENT is None: + _DEEPGEMM_M_ALIGNMENT = int(deep_gemm.get_m_alignment_for_contiguous_layout()) + return _DEEPGEMM_M_ALIGNMENT + + +def grouped_gemm_batch_invariant_alignment() -> int: + """Return the M alignment required by the DeepGEMM grouped-GEMM backend.""" + _require_deepgemm_bf16("get_m_alignment_for_contiguous_layout") + return _deepgemm_m_alignment() + + +def _pad_for_m_grouped(a: torch.Tensor, counts: List[int]) -> tuple: + """Pad an M-grouped contiguous input to satisfy DeepGEMM's per-expert M alignment. + + Returns the padded input, row-to-expert map, and padded counts. + The padded layout groups expert i's true rows contiguously at the start of + its 128-aligned block; remaining rows in the block are zero with m_indices=-1. + """ + alignment = _deepgemm_m_alignment() + padded_counts = [((count + alignment - 1) // alignment) * alignment for count in counts] + M_pad = sum(padded_counts) + if M_pad == 0: + return ( + torch.empty(0, a.shape[1], device=a.device, dtype=a.dtype), + torch.empty(0, device=a.device, dtype=torch.int32), + padded_counts, + ) + + a_padded = torch.zeros(M_pad, a.shape[1], device=a.device, dtype=a.dtype) + m_indices_padded = torch.full((M_pad,), -1, device=a.device, dtype=torch.int32) + src = 0 + dst = 0 + for i, (count, padded_count) in enumerate(zip(counts, padded_counts)): + if count > 0: + a_padded[dst : dst + count] = a[src : src + count] + m_indices_padded[dst : dst + count] = i + src += count + dst += padded_count + return a_padded, m_indices_padded, padded_counts + + +def _bf16_grouped_gemm_contiguous( + a: torch.Tensor, b: torch.Tensor, m_indices: torch.Tensor, counts: List[int] +) -> torch.Tensor: + """bf16 M-grouped GEMM via DeepGEMM. Deterministic / batch-invariant. + + a: [M_total, K] bf16, contiguous, expert-grouped (rows of expert i + are contiguous; m_indices is sorted). + b: [E, N, K] bf16, contiguous (per-expert weights, NT layout — + DeepGEMM transposes B internally). + m_indices: [M_total] int32, expert id per row (-1 to skip). + Returns: [M_total, N] bf16 with rows in the same order as `a`. + + Handles DeepGEMM's per-expert M alignment requirement by padding/unpadding + internally; the caller does not need pre-padded inputs. + """ + _require_deepgemm_bf16("m_grouped_bf16_gemm_nt_contiguous") + assert ( + a.dtype == torch.bfloat16 and b.dtype == torch.bfloat16 + ), f"bf16 grouped GEMM requires bf16; got a.dtype={a.dtype}, b.dtype={b.dtype}" + assert a.is_contiguous() and b.is_contiguous(), "a, b must be contiguous" + assert ( + m_indices.dtype == torch.int32 and m_indices.is_contiguous() + ), "m_indices must be int32 contiguous" + M_total, K = a.shape + E, N, K_b = b.shape + assert K == K_b, f"K mismatch between a ({K}) and b ({K_b})" + assert ( + m_indices.shape[0] == M_total + ), f"m_indices length {m_indices.shape[0]} != M_total {M_total}" + + assert len(counts) == E and sum(counts) == M_total + a_padded, m_indices_padded, padded_counts = _pad_for_m_grouped(a, counts) + M_pad = a_padded.shape[0] + if M_pad == 0: + return torch.zeros(M_total, N, device=a.device, dtype=torch.bfloat16) + + d_padded = torch.empty(M_pad, N, device=a.device, dtype=torch.bfloat16) + deep_gemm.m_grouped_bf16_gemm_nt_contiguous(a_padded, b, d_padded, m_indices_padded) + + # Strip padding: copy each expert's true rows back to a [M_total, N] tensor. + d = torch.empty(M_total, N, device=a.device, dtype=torch.bfloat16) + src = 0 + dst = 0 + for count, padded_count in zip(counts, padded_counts): + if count > 0: + d[src : src + count] = d_padded[dst : dst + count] + src += count + dst += padded_count + return d + + +def _bf16_grouped_gemm_aligned_contiguous( + a: torch.Tensor, b: torch.Tensor, m_indices: torch.Tensor +) -> torch.Tensor: + """DeepGEMM M-grouped GEMM for already aligned expert blocks. + + This path is used by inference CUDA graphs. The caller is responsible for + using `grouped_gemm_batch_invariant_alignment()` when building expert + offsets, so no device-to-host count extraction or dynamic padding is needed + on the captured path. + """ + _require_deepgemm_bf16("m_grouped_bf16_gemm_nt_contiguous") + assert ( + a.dtype == torch.bfloat16 and b.dtype == torch.bfloat16 + ), f"bf16 grouped GEMM requires bf16; got a.dtype={a.dtype}, b.dtype={b.dtype}" + assert a.is_contiguous() and b.is_contiguous(), "a, b must be contiguous" + assert ( + m_indices.dtype == torch.int32 and m_indices.is_contiguous() + ), "m_indices must be int32 contiguous" + M_total, K = a.shape + E, N, K_b = b.shape + assert K == K_b, f"K mismatch between a ({K}) and b ({K_b})" + assert ( + m_indices.shape[0] == M_total + ), f"m_indices length {m_indices.shape[0]} != M_total {M_total}" + + d = torch.empty(M_total, N, device=a.device, dtype=torch.bfloat16) + if M_total == 0: + return d + deep_gemm.m_grouped_bf16_gemm_nt_contiguous(a, b, d, m_indices) + return d + + +def _bf16_grouped_gemm_wgrad_contiguous( + grad_y: torch.Tensor, x: torch.Tensor, counts: List[int] +) -> torch.Tensor: + """K-grouped TN GEMM producing per-expert weight gradients via DeepGEMM. + + grad_y: [M_total, N] bf16, contiguous, expert-grouped. + x: [M_total, K] bf16, contiguous, expert-grouped (same row ordering). + Returns: [E, N, K] bf16 stacked per-expert wgrad. + + DeepGEMM's k_grouped_bf16 kernel computes in fp32 and requires fp32 d/c + accumulators; we cast the result back to bf16. Per-expert K alignment is + handled by padding internally. + """ + _require_deepgemm_bf16("k_grouped_bf16_gemm_tn_contiguous") + assert grad_y.dtype == torch.bfloat16 and x.dtype == torch.bfloat16 + assert grad_y.is_contiguous() and x.is_contiguous() + M_total, N = grad_y.shape + M_total_b, K = x.shape + assert M_total == M_total_b + + num_experts = len(counts) + assert sum(counts) == M_total + alignment = _deepgemm_m_alignment() + padded_counts = [((count + alignment - 1) // alignment) * alignment for count in counts] + M_pad = sum(padded_counts) + if M_pad == 0: + return torch.zeros(num_experts, N, K, device=grad_y.device, dtype=torch.bfloat16) + + grad_y_pad = torch.zeros(M_pad, N, device=grad_y.device, dtype=torch.bfloat16) + x_pad = torch.zeros(M_pad, K, device=x.device, dtype=torch.bfloat16) + src = 0 + dst = 0 + for c, cp in zip(counts, padded_counts): + if c > 0: + grad_y_pad[dst : dst + c] = grad_y[src : src + c] + x_pad[dst : dst + c] = x[src : src + c] + src += c + dst += cp + + ks_tensor = torch.tensor(padded_counts, dtype=torch.int32, device=grad_y.device) + d_fp32 = torch.zeros(num_experts, N, K, device=grad_y.device, dtype=torch.float32) + c_zero = torch.zeros(num_experts, N, K, device=grad_y.device, dtype=torch.float32) + deep_gemm.k_grouped_bf16_gemm_tn_contiguous( + grad_y_pad, x_pad, d_fp32, padded_counts, ks_tensor, c_zero + ) + return d_fp32.to(torch.bfloat16) + + +def grouped_gemm_batch_invariant( + a: torch.Tensor, + b: torch.Tensor, + *, + offs: torch.Tensor, + m_total: int, +) -> torch.Tensor: + """Run the graph-safe grouped GEMM over pre-aligned inference expert blocks.""" + m_indices = _offs_to_m_indices(offs, m_total).contiguous() + return _bf16_grouped_gemm_aligned_contiguous(a.contiguous(), b.contiguous(), m_indices) + + def _te_rmsnorm_forward_patched(self, x: torch.Tensor) -> torch.Tensor: """Patched TE RMSNorm.forward that routes to batch-invariant implementation with autograd support. @@ -962,11 +1599,30 @@ def is_batch_invariant_mode_enabled(): return _batch_invariant_MODE -def enable_batch_invariant_mode(): - """Enable global batch-invariant mode and patch Aten/TE kernels.""" - global _batch_invariant_MODE, _batch_invariant_LIB +def enable_batch_invariant_mode(backend: str = "deepgemm"): + """Enable global batch-invariant mode and patch Aten/TE kernels. + + Args: + backend: which kernel to dispatch `aten::mm`/`aten::addmm` through. + "deepgemm" (default) routes bf16 CUDA inputs through DeepGEMM + `bf16_gemm_nn`. "triton" routes through the batch-invariant + Triton `matmul_persistent` kernel (works for bf16/fp16/fp32 and + on any CUDA device). Grouped GEMM always uses DeepGEMM regardless. + """ + global _batch_invariant_MODE, _batch_invariant_LIB, _BATCH_INVARIANT_BACKEND if _batch_invariant_MODE: return + if backend not in _BATCH_INVARIANT_BACKENDS: + raise ValueError( + f"Unknown batch-invariant backend {backend!r}; " + f"expected one of {_BATCH_INVARIANT_BACKENDS}." + ) + if backend == "deepgemm" and not HAVE_DEEPGEMM_BF16: + raise RuntimeError( + "The DeepGEMM batch-invariant backend requires DeepGEMM with " + "bf16 bindings. Install DeepGEMM or use backend='triton'." + ) + _BATCH_INVARIANT_BACKEND = backend dispatch_key = getattr(torch.accelerator.current_accelerator(), "type", "cpu").upper() _batch_invariant_MODE = True _batch_invariant_LIB = torch.library.Library("aten", "IMPL") @@ -976,6 +1632,10 @@ def enable_batch_invariant_mode(): _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, dispatch_key) # Also patch Transformer Engine kernels when available _te_patch_for_batch_invariant() + # Pin the Mamba autotuners so rollout and training processes can't end + # up on different tile configs (and therefore different fp32 reduction + # orders) through autotune timing noise. + _pin_mamba_autotuners() def disable_batch_invariant_mode(): @@ -987,33 +1647,137 @@ def disable_batch_invariant_mode(): _batch_invariant_LIB = None # Restore Transformer Engine kernels if previously patched _te_unpatch_for_batch_invariant() + _unpin_mamba_autotuners() + + +# (autotuner, original configs list) pairs saved by _pin_mamba_autotuners. +_PINNED_AUTOTUNERS: list = [] + +# Rollout uses the repo kernels while training uses mamba_ssm. Pin a config +# present in both copies so autotune timing cannot change the reduction order. +_PINNED_MAMBA_CONFIGS = { + "_bmm_chunk_fwd_kernel": {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32}, + "_chunk_scan_fwd_kernel": {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, + "_chunk_state_fwd_kernel": {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32}, + "_chunk_cumsum_fwd_kernel": {"BLOCK_SIZE_H": 8}, + "_state_passing_fwd_kernel": {"BLOCK_SIZE": 1024}, +} + + +def _pin_mamba_autotuners(): + """Pin the Mamba chunked-scan forward kernels to fixed tile configs. + + BLOCK sizes determine the fp32 reduction grouping inside tl.dot loops, + so rollout/train parity needs the inference process (repo ssd_* kernels) + and the training process (mamba_ssm package kernels) to pick the same + config. Triton's autotuner re-benchmarks per process and timing noise + can flip the winner; we've seen that break parity in practice. Pinning + both sides to the same config removes the benchmark from the loop. + + Only the five forward kernels matter for parity; backward kernels only + affect gradients. + """ + global _PINNED_AUTOTUNERS + try: + from triton.runtime.autotuner import Autotuner + except ImportError: + return + + kernels = [] + try: + from megatron.core.ssm.ops import ( + ssd_bmm as r_bmm, + ssd_chunk_scan as r_scan, + ssd_chunk_state as r_state, + ssd_state_passing as r_pass, + ) + + kernels += [ + r_bmm._bmm_chunk_fwd_kernel, + r_scan._chunk_scan_fwd_kernel, + r_state._chunk_cumsum_fwd_kernel, + r_state._chunk_state_fwd_kernel, + r_pass._state_passing_fwd_kernel, + ] + except ImportError: + pass + try: + from mamba_ssm.ops.triton import ( + ssd_bmm as p_bmm, + ssd_chunk_scan as p_scan, + ssd_chunk_state as p_state, + ssd_state_passing as p_pass, + ) + + kernels += [ + p_bmm._bmm_chunk_fwd_kernel, + p_scan._chunk_scan_fwd_kernel, + p_state._chunk_cumsum_fwd_kernel, + p_state._chunk_state_fwd_kernel, + p_pass._state_passing_fwd_kernel, + ] + except (ImportError, AttributeError): + pass + + for kernel in kernels: + if not isinstance(kernel, Autotuner) or len(kernel.configs) <= 1: + continue + name = getattr(getattr(kernel, "fn", None), "__name__", "") + expected = _PINNED_MAMBA_CONFIGS[name] + chosen = next( + cfg + for cfg in kernel.configs + if all(cfg.kwargs.get(key) == value for key, value in expected.items()) + ) + _PINNED_AUTOTUNERS.append((kernel, kernel.configs)) + kernel.configs = [chosen] + if hasattr(kernel, "cache"): + kernel.cache.clear() + + +def _unpin_mamba_autotuners(): + """Restore the original autotune config lists saved by _pin_mamba_autotuners.""" + global _PINNED_AUTOTUNERS + for kernel, original in _PINNED_AUTOTUNERS: + kernel.configs = original + if hasattr(kernel, "cache"): + kernel.cache.clear() + _PINNED_AUTOTUNERS = [] @contextlib.contextmanager -def set_batch_invariant_mode(enabled: bool = True): +def set_batch_invariant_mode(enabled: bool = True, backend: Optional[str] = None): """Context manager to toggle global batch-invariant mode. When `enabled` is True, batch-invariant kernels are enabled for the duration of the context; when False, they are disabled for the duration. This implementation is re-entrant and correctly restores the previous state even under nesting. + The helper default remains "triton" for tests that exercise non-bf16 operators. """ global _batch_invariant_MODE, _batch_invariant_LIB # Save the previous on/off state so we can correctly restore it, even under # nested usage or when toggling from True->False inside an outer True scope. prev_enabled = _batch_invariant_MODE + prev_backend = _BATCH_INVARIANT_BACKEND # Apply the requested state only if it differs from the current one. if enabled and not prev_enabled: - enable_batch_invariant_mode() + enable_batch_invariant_mode(backend=backend or "triton") + elif enabled and prev_enabled and backend is not None and backend != prev_backend: + raise RuntimeError( + "Cannot switch batch-invariant backend inside an active context " + f"(active={prev_backend!r}, requested={backend!r})." + ) elif not enabled and prev_enabled: disable_batch_invariant_mode() try: yield finally: - # Restore the previous state. If we turned BIK on at entry, turn it off here. - # If we turned it off at entry (inside an outer True scope), turn it back on. + # Restore the previous state. If we turned batch-invariant mode on at + # entry, turn it off here. If we turned it off at entry (inside an + # outer True scope), turn it back on. if enabled and not prev_enabled: disable_batch_invariant_mode() elif not enabled and prev_enabled: - enable_batch_invariant_mode() + enable_batch_invariant_mode(backend=prev_backend) diff --git a/megatron/core/transformer/moe/batch_invariant.py b/megatron/core/transformer/moe/batch_invariant.py new file mode 100644 index 00000000000..79a0611bf65 --- /dev/null +++ b/megatron/core/transformer/moe/batch_invariant.py @@ -0,0 +1,92 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Batch-invariant MoE permutation helpers.""" + +from typing import Optional + +import torch + +from megatron.core import parallel_state + + +def build_inverse_permutation_map( + routing_map: torch.Tensor, + flat_sorted: torch.Tensor, + sorted_indices: torch.Tensor, + num_out_tokens: int, +) -> torch.Tensor: + """Build token/top-k -> permuted-row and expert-id map for batch-invariant unpermute. + + The regular permutation map is row -> token. Batch-invariant unpermute needs + the inverse ownership model so each output token can read its routed rows and + add them in a fixed order. + """ + num_tokens = routing_map.size(0) + assert isinstance( + num_out_tokens, int + ), "batch-invariant graph unpermute requires static num_out_tokens" + assert num_out_tokens % num_tokens == 0, ( + "batch-invariant graph unpermute expects fixed top-k per token" + ) + + topk = num_out_tokens // num_tokens + row_ids = torch.arange(num_out_tokens, device=routing_map.device, dtype=torch.long) + expert_ids = torch.div(flat_sorted, num_tokens, rounding_mode='floor').to(torch.long) + token_ids = sorted_indices.to(torch.long) + + slots_by_token_expert = routing_map.bool().to(torch.long).cumsum(dim=1) - 1 + row_slots = slots_by_token_expert[token_ids, expert_ids] + linear_slots = token_ids * topk + row_slots + + inverse_rows = torch.full((num_tokens, topk), -1, device=routing_map.device, dtype=torch.long) + inverse_experts = torch.full((num_tokens, topk), -1, device=routing_map.device, dtype=torch.long) + inverse_rows.view(-1).scatter_(0, linear_slots, row_ids) + inverse_experts.view(-1).scatter_(0, linear_slots, expert_ids) + return torch.stack((inverse_rows, inverse_experts), dim=0) + + +def unpermute( + permuted_tokens: torch.Tensor, + restore_shape: torch.Size, + *, + probs: Optional[torch.Tensor], + num_experts: int, + inverse_map: torch.Tensor, +) -> torch.Tensor: + """Batch-invariant MoE unpermute. + + Accumulation is token-owned. The AllToAll inverse map avoids data-dependent + shapes and adds contributions by EP rank then top-k slot, matching the + inference NVLS rank-ordered combine. + """ + input_dtype = permuted_tokens.dtype + output_tokens = torch.zeros(restore_shape, dtype=torch.float32, device=permuted_tokens.device) + ep_size = parallel_state.get_expert_model_parallel_world_size() or 1 + assert num_experts % ep_size == 0, "batch-invariant MoE expects contiguous EP shards" + experts_per_rank = num_experts // ep_size + inverse_rows = inverse_map[0] + inverse_experts = inverse_map[1] + topk = inverse_rows.size(1) + + for ep_rank in range(ep_size): + rank_partial = torch.zeros_like(output_tokens) + start_expert = ep_rank * experts_per_rank + end_expert = start_expert + experts_per_rank + + for k in range(topk): + row_ids = inverse_rows[:, k] + expert_ids = inverse_experts[:, k] + valid_mask = ( + (row_ids >= 0) & (expert_ids >= start_expert) & (expert_ids < end_expert) + ) + + safe_rows = row_ids.clamp_min(0) + chunk = permuted_tokens.index_select(0, safe_rows).to(torch.float32) + if probs is not None: + safe_experts = expert_ids.clamp_min(0) + chunk = chunk * probs.gather(1, safe_experts.unsqueeze(1)).to(torch.float32) + chunk.masked_fill_(~valid_mask.unsqueeze(-1), 0.0) + rank_partial += chunk + + output_tokens += rank_partial + + return output_tokens.to(dtype=input_dtype) diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py index fbadcb7d3da..aae7a731c84 100644 --- a/megatron/core/transformer/moe/moe_utils.py +++ b/megatron/core/transformer/moe/moe_utils.py @@ -19,7 +19,14 @@ ) from megatron.core.tensor_parallel.mappings import reduce_from_tensor_model_parallel_region from megatron.core.transformer.cuda_graphs import is_graph_capturing +from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( + is_batch_invariant_mode_enabled, +) from megatron.core.transformer.enums import CudaGraphModule +from megatron.core.transformer.moe.batch_invariant import ( + build_inverse_permutation_map as build_batch_invariant_inverse_permutation_map, + unpermute as batch_invariant_unpermute, +) from megatron.core.transformer.moe.router_replay import RouterReplay from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.utils import internal_api, is_te_min_version @@ -308,6 +315,7 @@ def permute( drop_and_pad: bool = False, tokens_per_expert: Optional[torch.Tensor] = None, align_size: int = 0, + return_batch_invariant_inverse_map: bool = False, ) -> Tuple[ torch.Tensor, Optional[torch.Tensor], @@ -340,6 +348,8 @@ def permute( tokens_per_expert (torch.Tensor, optional): Tensor of shape `[num_experts]` containing actual token counts per expert. align_size (int, optional): The alignment size for the input tensor for fp8 or fp4. + return_batch_invariant_inverse_map (bool, optional): Return a fixed-shape + batch-invariant inverse map in the `pad_offsets` slot for graph-safe unpermute. Returns: Tuple[ @@ -352,6 +362,10 @@ def permute( The permuted tokens, (optional) permuted probs, sorted indices, (optional) pad_offsets, (optional) padded_tokens_per_expert. """ + if return_batch_invariant_inverse_map: + assert not fused, "batch-invariant MoE permute requires the unfused path" + assert not drop_and_pad, "batch-invariant MoE supports dynamic dropless routing only" + if fused and probs is None: if not HAVE_TE or fused_permute is None: raise ValueError("fused_permute is not available. Please install TE >= 2.1.0.") @@ -386,6 +400,7 @@ def permute( num_tokens, hidden = tokens.shape num_experts = routing_map.shape[1] permuted_probs = None + batch_invariant_inverse_map = None if drop_and_pad and not (num_out_tokens is None): capacity = num_out_tokens // num_experts assert not routing_map.requires_grad @@ -412,6 +427,7 @@ def permute( assert ( num_out_tokens is not None ), "num_out_tokens is required for the argsort-based permute" + routing_map_for_inverse = routing_map # mask [num_tokens, num_experts] -> [num_experts, num_tokens] routing_map = routing_map.bool().T.contiguous() @@ -426,10 +442,24 @@ def permute( if probs is not None: permuted_probs = probs.T.contiguous().reshape(-1)[flat_sorted] + if return_batch_invariant_inverse_map: + batch_invariant_inverse_map = build_batch_invariant_inverse_permutation_map( + routing_map_for_inverse, + flat_sorted, + sorted_indices, + num_out_tokens, + ) + # use the mapping to permute the tokens permuted_input = tokens.index_select(0, sorted_indices) - return permuted_input, permuted_probs, sorted_indices, None, tokens_per_expert + return ( + permuted_input, + permuted_probs, + sorted_indices, + batch_invariant_inverse_map, + tokens_per_expert, + ) def unpermute( @@ -441,6 +471,7 @@ def unpermute( fused: bool = False, drop_and_pad: bool = False, pad_offsets: Optional[torch.Tensor] = None, + batch_invariant_inverse_map: Optional[torch.Tensor] = None, ) -> torch.Tensor: """ Restore the original order of tokens after permutation. If probs are provided, it @@ -466,10 +497,18 @@ def unpermute( Tensor of per-expert cumulative padding offsets used to remove padding added during permutation. This is the fourth output of `moe_permute_and_pad_with_probs` and is required when unpermuting padded outputs. Defaults to None. + batch_invariant_inverse_map (torch.Tensor, optional): Fixed-shape + `[2, num_tokens, topk]` map from token/top-k slot to permuted row and + global expert id. Used by batch-invariant CUDA graph paths. Returns: torch.Tensor: The tokens restored to their original order. """ + batch_invariant_mode = is_batch_invariant_mode_enabled() + if batch_invariant_mode: + assert not fused, "batch-invariant MoE unpermute requires the unfused path" + assert not drop_and_pad, "batch-invariant MoE supports dynamic dropless routing only" + if fused: if not HAVE_TE or fused_unpermute is None: raise ValueError("fused_unpermute is not available. Please install TE >= 2.1.0.") @@ -484,6 +523,19 @@ def unpermute( **extra_kwargs, ) + if batch_invariant_mode: + assert routing_map is not None, "batch-invariant MoE unpermute requires routing_map" + assert batch_invariant_inverse_map is not None, ( + "batch-invariant MoE unpermute requires the AllToAll inverse map" + ) + return batch_invariant_unpermute( + permuted_tokens, + restore_shape, + probs=probs, + num_experts=routing_map.size(1), + inverse_map=batch_invariant_inverse_map, + ) + _, hidden = restore_shape input_dtype = permuted_tokens.dtype diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py index 9e012dcbd88..69aaa79187e 100644 --- a/megatron/core/transformer/moe/token_dispatcher.py +++ b/megatron/core/transformer/moe/token_dispatcher.py @@ -461,6 +461,9 @@ def __init__( 'reversed_local_input_permutation_mapping', 'routing_map', ] + self.batch_invariant_inverse_permutation_mapping = None + if self.config.batch_invariant_mode: + self.cudagraph_attrs.append('batch_invariant_inverse_permutation_mapping') self.shared_experts = None @@ -639,7 +642,7 @@ def dispatch_preprocess( permutated_local_input_tokens, permuted_probs, self.reversed_local_input_permutation_mapping, - _, + self.batch_invariant_inverse_permutation_mapping, _, ) = permute( hidden_states, @@ -648,6 +651,7 @@ def dispatch_preprocess( num_out_tokens=self.num_out_tokens, fused=self.config.moe_permute_fusion, drop_and_pad=self.drop_and_pad, + return_batch_invariant_inverse_map=self.config.batch_invariant_mode, ) return permutated_local_input_tokens, permuted_probs @@ -867,6 +871,7 @@ def combine_postprocess(self, permutated_local_input_tokens): routing_map=self.routing_map, fused=self.config.moe_permute_fusion, drop_and_pad=self.drop_and_pad, + batch_invariant_inverse_map=self.batch_invariant_inverse_permutation_mapping, ) # Reshape the output tensor diff --git a/megatron/core/transformer/moe/token_dispatcher_inference.py b/megatron/core/transformer/moe/token_dispatcher_inference.py index e85115528b8..84fb4cda68f 100644 --- a/megatron/core/transformer/moe/token_dispatcher_inference.py +++ b/megatron/core/transformer/moe/token_dispatcher_inference.py @@ -31,6 +31,7 @@ from megatron.core.inference.communication.torch_symm_triton import ( multimem_all_gatherv_3tensor, multimem_reduce_scatter_v, + ordered_reduce_scatter_v, ) from megatron.core.inference.moe import InferenceGroupedGemmBackend from megatron.core.inference.moe.metadata import fused_metadata_update @@ -43,6 +44,9 @@ from megatron.core.transformer.moe.shared_experts import SharedExpertMLP from megatron.core.transformer.moe.token_dispatcher import MoEAllGatherTokenDispatcher from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( + is_batch_invariant_mode_enabled, +) from megatron.core.typed_torch import apply_module from megatron.core.utils import get_pg_rank, get_pg_size @@ -554,9 +558,13 @@ def combine_preprocess(self, expert_output): def token_combine(self, hidden_states): """ReduceScatter-V: sum expert outputs across EP ranks, scatter to local tokens. + In batch-invariant mode, the symmetric RSV buffer is still used for data + visibility, but the rank reduction is an explicit fp32 rank-order loop + rather than a hardware multimem reduction. + Args: - hidden_states: [global_max, hidden_size] expert outputs (fp32 when - written directly to the RSV buffer, bf16 otherwise). + hidden_states: [global_max, hidden_size] expert outputs (fp32 + when written directly to the RSV buffer, bf16 otherwise). Returns: [local_tokens, hidden_size] bf16 local token outputs. @@ -574,7 +582,12 @@ def token_combine(self, hidden_states): dtype=rsv["tensor"].dtype, device=hidden_states.device, ) - multimem_reduce_scatter_v( + reduce_scatter_v = ( + ordered_reduce_scatter_v + if is_batch_invariant_mode_enabled() + else multimem_reduce_scatter_v + ) + reduce_scatter_v( output, rsv["tensor"], rsv["handle"], diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 4e45fa9890d..61170f06911 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -992,7 +992,7 @@ class TransformerConfig(ModelParallelConfig): batch_invariant_mode: bool = False """If true, uses batch-invariant kernels that provide deterministic forward execution regardless of batch size. This ensures bitwise identical results when the same inputs are processed - in different batch configurations. This will significantly affect speed of + in different batch configurations. This will significantly affect speed of training and inference as the kernels are not full optimized. Defaults to False.""" @@ -1337,6 +1337,21 @@ def __post_init__(self): "Set inference_grouped_gemm_backend to 'torch' for MXFP8." ) + if self.batch_invariant_mode: + if self.inference_grouped_gemm_backend != InferenceGroupedGemmBackend.TORCH: + raise ValueError( + "batch_invariant_mode requires " + "inference_grouped_gemm_backend='torch'." + ) + if ( + self.expert_model_parallel_size > 1 + and self.inference_moe_token_dispatcher_type != "nvls" + ): + raise ValueError( + "batch_invariant_mode with inference-optimized MoE and expert " + "parallelism requires inference_moe_token_dispatcher_type='nvls'." + ) + if self.num_moe_experts is not None and self.num_moe_experts <= 0: raise ValueError("num_moe_experts must be non-negative.") @@ -2470,6 +2485,37 @@ def _scope_to_str(s): assert ( self.attention_backend == AttnBackend.flash ), "Batch invariant mode only supports FlashAttention" + if (self.num_moe_experts or 0) > 0: + from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( + HAVE_DEEPGEMM_BF16, + ) + + if self.transformer_impl != "inference_optimized": + assert self.moe_token_dispatcher_type == "alltoall", ( + "Batch-invariant MoE training requires " + "moe_token_dispatcher_type='alltoall'." + ) + assert HAVE_DEEPGEMM_BF16, ( + "batch_invariant_mode=True with MoE requires DeepGEMM with bf16 " + "grouped-GEMM bindings (m_grouped_bf16_gemm_nt_contiguous). " + "Install via `uv pip install -e .[batch_invariant]`." + ) + assert not ( + self.fp8 or self.fp4 + ), "Batch-invariant MoE is bf16-only. Disable fp8/fp4 to use it." + assert not ( + self.moe_permute_fusion or self.moe_permute_fusion_into_hybridep + ), ( + "Batch-invariant MoE requires the unfused permute/unpermute path so " + "top-k reductions use the fixed batch-invariant add tree." + ) + assert not ( + self.moe_pad_expert_input_to_capacity + or self.moe_pad_experts_for_cuda_graph_inference + ), ( + "Batch-invariant MoE supports dynamic dropless routing only. " + "Disable MoE capacity/expert padding." + ) @dataclass diff --git a/pyproject.toml b/pyproject.toml index 9f376260c0d..0efbcf5e583 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -179,7 +179,8 @@ linting = [ "pylint==3.2.6", ] ci = ["python-gitlab", "slack-sdk", "pandas"] -no_pypi_wheels = ["flash_mla", "emerging_optimizers"] +batch_invariant = ["deep_gemm"] +no_pypi_wheels = ["flash_mla", "emerging_optimizers", "deep_gemm"] [tool.uv] default-groups = ["linting", "build", "test"] @@ -189,6 +190,7 @@ no-build-isolation-package = [ "mamba-ssm", "transformer-engine", "transformer-engine-torch", + "deep_gemm", ] link-mode = "copy" conflicts = [[{ extra = "lts" }, { extra = "dev" }]] @@ -206,6 +208,7 @@ override-dependencies = [ flash_mla = [ { git = "https://github.com/deepseek-ai/FlashMLA", rev = "9edee0c022cd0938148a18e334203b0aab43aa19" }, ] +deep_gemm = { git = "https://github.com/deepseek-ai/DeepGEMM.git", rev = "714dd1a4a980f7937a74343d19a8eba4fe321480" } transformer-engine = { git = "https://github.com/NVIDIA/TransformerEngine.git", rev = "42b840051647eef89761a16dfdff87e82bb253ab" } nemo-run = { git = "https://github.com/NVIDIA-NeMo/Run.git", rev = "17ae86b64d7f75653351664f5d8c9e466faede00" } emerging_optimizers = { git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git", rev = "v0.2.0" } diff --git a/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py b/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py index 84898db60d8..cf2ed7c5293 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py @@ -18,6 +18,7 @@ ) from megatron.core.inference.sampling_params import SamplingParams from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.enums import AttnBackend from megatron.core.transformer.transformer_config import TransformerConfig from tests.unit_tests.test_utilities import Utils @@ -54,11 +55,14 @@ def _ctx( block_size_tokens=32, max_sequence_length=512, rounder=64, + max_requests=None, enable_prefix_caching=True, max_tokens=None, prefix_caching_eviction_policy=PrefixCachingEvictionPolicy.LRU, mamba_config=None, prefix_caching_mamba_gb=None, + batch_invariant_mode=False, + enable_chunked_prefill=False, ): DynamicInferenceContext.ROUNDER = rounder DynamicInferenceContext.TOKEN_ROUNDER = rounder @@ -73,9 +77,12 @@ def _ctx( tensor_model_parallel_size=1, pipeline_model_parallel_size=1, use_cpu_initialization=True, + batch_invariant_mode=batch_invariant_mode, + attention_backend=AttnBackend.flash if batch_invariant_mode else AttnBackend.auto, ) inference_config = InferenceConfig( max_sequence_length=max_sequence_length, + max_requests=max_requests, buffer_size_gb=buffer_size_gb, paused_buffer_size_gb=0.2 * buffer_size_gb, block_size_tokens=block_size_tokens, @@ -84,6 +91,7 @@ def _ctx( use_flashinfer_fused_rope=None, unified_memory_level=0, enable_prefix_caching=enable_prefix_caching, + enable_chunked_prefill=enable_chunked_prefill, prefix_caching_eviction_policy=prefix_caching_eviction_policy, prefix_caching_mamba_gb=prefix_caching_mamba_gb, ) @@ -778,6 +786,40 @@ def test_mamba_prefill_skip_and_zero_prefill(self): ctx5.release_memory_blocks_from_request_indexes([0]) assert not msa5.has_state(bid5) and bh5 not in msa5.hash_to_block_id + @pytest.mark.internal + def test_batch_invariant_mamba_chunked_prefill_scheduler_alignment(self): + ctx = self._mctx(block_size_tokens=32, batch_invariant_mode=True) + engine = _StubEngine(ctx, enable_chunked_prefill=True) + req = self._req(ctx, self._prompt(500)) + + assert engine._mamba_batch_invariant_prefill_chunk_length(req, 300) == 256 + assert engine._mamba_batch_invariant_prefill_chunk_length(req, 100) == 0 + + short_req = self._req(ctx, self._prompt(200), request_id=2) + assert engine._mamba_batch_invariant_prefill_chunk_length(short_req, 300) == 200 + + one_left_req = self._req(ctx, self._prompt(ctx.mamba_chunk_size + 1), request_id=3) + assert ( + engine._mamba_batch_invariant_prefill_chunk_length( + one_left_req, ctx.mamba_chunk_size + ) + == 0 + ) + assert ( + engine._mamba_batch_invariant_prefill_chunk_length( + one_left_req, ctx.mamba_chunk_size + 1 + ) + == ctx.mamba_chunk_size + 1 + ) + + with pytest.raises(AssertionError, match="max_tokens > mamba_chunk_size"): + self._mctx( + batch_invariant_mode=True, + enable_chunked_prefill=True, + max_tokens=ctx.mamba_chunk_size, + max_requests=64, + ) + @pytest.mark.internal def test_mamba_intermediate_offsets(self): bs = 256 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 49b5df613f7..0e941ea8568 100644 --- a/tests/unit_tests/inference/test_moe_dispatching_and_routing.py +++ b/tests/unit_tests/inference/test_moe_dispatching_and_routing.py @@ -230,6 +230,21 @@ def test_init(self): assert dispatcher.topk == NANOV3_BASE["moe_router_topk"] assert dispatcher.ep_size == Utils.world_size + def test_init_rejects_batch_invariant_ep(self): + """Batch-invariant MoE on the inference EP path is NVLS-only on this branch.""" + if Utils.world_size == 1: + pytest.skip( + "NCCL batch-invariant rejection is only relevant with expert parallelism." + ) + + with pytest.raises(ValueError, match="requires inference_moe_token_dispatcher_type"): + self._make_dispatcher( + batch_invariant_mode=True, + attention_backend=AttnBackend.flash, + inference_grouped_gemm_backend="torch", + inference_moe_token_dispatcher_type="nccl", + ) + @pytest.mark.parametrize("use_allgather_v", [False, True]) def test_dispatch_combine(self, use_allgather_v): """Dispatch+combine correctness for both CG (equal-count) and prefill (variable-count) paths. @@ -326,6 +341,9 @@ def _make_dispatcher(self): NVLSAllGatherVDispatcher, ) + if Utils.world_size <= 0 or Utils.world_size & (Utils.world_size - 1): + pytest.skip("NVLS Triton symmetric-memory barrier requires power-of-two EP size.") + config = _make_base_config(expert_model_parallel_size=Utils.world_size) num_local_experts = config.num_moe_experts // Utils.world_size ep_rank = torch.distributed.get_rank() if Utils.world_size > 1 else 0 @@ -439,3 +457,204 @@ def test_cuda_graph_dispatch_combine(self, max_rank_tokens, seed): expected_combined = (global_hidden[start:end].float() * ep_size).bfloat16() torch.testing.assert_close(graph_combined, expected_combined, atol=0, rtol=0) + + def test_cuda_graph_batch_invariant_combine_uses_ordered_symmetric_memory(self, monkeypatch): + """Batch-invariant mode should use ordered peer loads on NVLS dispatcher. + + The graph path still writes local partials into the symmetric RSV buffer, + but the combine must not use multimem.ld_reduce in batch-invariant mode. + """ + from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( + set_batch_invariant_mode, + ) + from megatron.core.transformer.moe import token_dispatcher_inference + + if Utils.world_size < 2: + pytest.skip("Ordered RSV combine requires expert-parallel world_size > 1.") + + torch.manual_seed(2026) + torch.cuda.manual_seed(2026) + + dispatcher = self._make_dispatcher() + ep_size = dispatcher.ep_size + hidden_size = NANOV3_BASE["hidden_size"] + topk = NANOV3_BASE["moe_router_topk"] + num_experts = NANOV3_BASE["num_moe_experts"] + rank = torch.distributed.get_rank() if ep_size > 1 else 0 + + max_rank_tokens = 24 + tokens_per_rank = [max(1, max_rank_tokens + r - (ep_size - 1)) for r in range(ep_size)] + local_tokens = tokens_per_rank[rank] + total_tokens = sum(tokens_per_rank) + global_max = _NVLS_ENGINE_MAX_TOKENS * ep_size + + global_hidden = torch.randn(total_tokens, hidden_size, device="cuda", dtype=torch.bfloat16) + global_probs = torch.randn(total_tokens, topk, device="cuda", dtype=torch.float32) + global_routing_map = torch.randint(0, num_experts, (total_tokens, topk), device="cuda") + if ep_size > 1: + torch.distributed.broadcast(global_hidden, src=0) + torch.distributed.broadcast(global_probs, src=0) + torch.distributed.broadcast(global_routing_map, src=0) + + start = sum(tokens_per_rank[:rank]) + end = start + local_tokens + static_hidden = global_hidden[start:end].contiguous() + static_probs = global_probs[start:end].contiguous() + static_routing_map = global_routing_map[start:end].contiguous() + + ordered_calls = {"value": 0} + orig_ordered_reduce_scatter_v = token_dispatcher_inference.ordered_reduce_scatter_v + + def _tracked_ordered_reduce_scatter_v(*args, **kwargs): + ordered_calls["value"] += 1 + return orig_ordered_reduce_scatter_v(*args, **kwargs) + + monkeypatch.setattr( + token_dispatcher_inference, + "ordered_reduce_scatter_v", + _tracked_ordered_reduce_scatter_v, + ) + + with torch.no_grad(), set_batch_invariant_mode(True): + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + dispatcher.routing_map = static_routing_map + dispatcher._local_tokens = local_tokens + d_hidden, _ = dispatcher.token_dispatch(static_hidden, static_probs) + dispatcher.token_combine(d_hidden.clone()) + torch.cuda.current_stream().wait_stream(s) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + dispatcher.routing_map = static_routing_map + dispatcher._local_tokens = local_tokens + d_hidden, _ = dispatcher.token_dispatch(static_hidden, static_probs) + assert d_hidden.shape[0] == global_max + graph_combined = dispatcher.token_combine(d_hidden.clone()) + + graph.replay() + + assert ordered_calls["value"] > 0 + assert graph_combined.shape == (local_tokens, hidden_size) + expected_combined = (global_hidden[start:end].float() * ep_size).bfloat16() + torch.testing.assert_close(graph_combined, expected_combined, atol=0, rtol=0) + + def test_cuda_graph_batch_invariant_moe_layer_uses_ordered_rsv(self, monkeypatch): + """A real inference MoE layer should use ordered RSV combine in batch-invariant mode. + + This catches the production branch in InferenceGroupedMLP: mcore_fused_moe + writes deterministic local partials into the symmetric RSV buffer, then + token_combine uses explicit rank-order fp32 loads. + """ + 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 + from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( + HAVE_DEEPGEMM_BF16, + set_batch_invariant_mode, + ) + from megatron.core.transformer.moe.token_dispatcher_inference import ( + NVLSAllGatherVDispatcher, + ) + from megatron.core.transformer.moe import token_dispatcher_inference + + if Utils.world_size < 2: + pytest.skip("NVLS RSV branch test requires expert-parallel world_size > 1.") + 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: + pytest.skip("Batch-invariant torch MoE path requires DeepGEMM bf16 grouped kernels.") + + torch.manual_seed(2027) + torch.cuda.manual_seed(2027) + + config = _make_base_config( + expert_model_parallel_size=Utils.world_size, + inference_grouped_gemm_backend="torch", + inference_moe_token_dispatcher_type="nvls", + batch_invariant_mode=True, + attention_backend=AttnBackend.flash, + moe_shared_expert_intermediate_size=None, + ) + ep_group = get_expert_model_parallel_group() + NVLSAllGatherVDispatcher.allocate_buffers( + per_rank_worst_case_token_count=_NVLS_ENGINE_MAX_TOKENS, + topk=config.moe_router_topk, + hidden_size=config.hidden_size, + ep_group=ep_group, + ) + + layer = get_inference_optimized_moe_spec()(config=config).cuda().eval() + assert isinstance(layer._inference_token_dispatcher, NVLSAllGatherVDispatcher) + layer.token_dispatcher = layer._inference_token_dispatcher + layer.shared_expert_overlap = layer._inference_token_dispatcher.shared_experts is not None + assert not hasattr(layer.experts, "_batch_invariant_global_unpermute") + + used_rsv = {"value": False} + orig_get_rsv_tensor = NVLSAllGatherVDispatcher._get_rsv_tensor.__func__ + + def _tracked_get_rsv_tensor(cls): + tensor = orig_get_rsv_tensor(cls) + used_rsv["value"] = used_rsv["value"] or tensor is not None + return tensor + + monkeypatch.setattr( + NVLSAllGatherVDispatcher, "_get_rsv_tensor", classmethod(_tracked_get_rsv_tensor) + ) + ordered_calls = {"value": 0} + orig_ordered_reduce_scatter_v = token_dispatcher_inference.ordered_reduce_scatter_v + + def _tracked_ordered_reduce_scatter_v(*args, **kwargs): + ordered_calls["value"] += 1 + return orig_ordered_reduce_scatter_v(*args, **kwargs) + + monkeypatch.setattr( + token_dispatcher_inference, + "ordered_reduce_scatter_v", + _tracked_ordered_reduce_scatter_v, + ) + + local_tokens = 16 + hidden_states = torch.randn( + local_tokens, 1, config.hidden_size, device="cuda", dtype=torch.bfloat16 + ) + probs = torch.randn( + local_tokens, config.moe_router_topk, device="cuda", dtype=torch.float32 + ) + routing_map = ( + torch.arange(local_tokens * config.moe_router_topk, device="cuda") + .reshape(local_tokens, config.moe_router_topk) + .remainder(config.num_moe_experts) + .to(torch.int64) + ) + + def _run_expert_and_combine(): + preprocessed_hidden, preprocessed_probs = layer.preprocess( + hidden_states, probs, routing_map + ) + dispatched_hidden, dispatched_probs = layer.dispatch( + preprocessed_hidden, preprocessed_probs + ) + output, _ = layer.routed_experts_compute(dispatched_hidden, dispatched_probs) + output = layer.combine(output) + return layer.postprocess(output, None) + + with torch.no_grad(), InferenceMode.active(), set_batch_invariant_mode(True): + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + _run_expert_and_combine() + torch.cuda.current_stream().wait_stream(s) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_output = _run_expert_and_combine() + + graph.replay() + + assert used_rsv["value"] + assert ordered_calls["value"] > 0 + assert graph_output.shape == hidden_states.shape + assert graph_output.dtype == torch.bfloat16 diff --git a/tests/unit_tests/inference/test_moe_permute.py b/tests/unit_tests/inference/test_moe_permute.py index 6bddf515b14..e5a67d2d445 100644 --- a/tests/unit_tests/inference/test_moe_permute.py +++ b/tests/unit_tests/inference/test_moe_permute.py @@ -393,6 +393,50 @@ def test_multiple_topk_accumulation(self, topk): result[0], torch.full((hidden_dim,), expected_val, device="cuda"), atol=1e-4, rtol=1e-4 ) + def test_batch_invariant_unpermute_is_token_local(self): + """Unrelated earlier tokens must not affect another token's top-k sum.""" + from megatron.core.inference.moe.permute import unpermute_tokens + from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( + set_batch_invariant_mode, + ) + + hidden_dim = 8 + permutation_map = torch.empty(3, dtype=torch.int32, device="cuda") + probs = torch.ones(3, device="cuda", dtype=torch.float32) + + # Token 1 has two unit contributions in both layouts. Layout B adds a + # huge unrelated token-0 contribution before it; token 1 must not drift. + expert_output_a = torch.ones(2, hidden_dim, device="cuda", dtype=torch.bfloat16) + inverse_a = torch.tensor([[-1, -1], [0, 1]], dtype=torch.int32, device="cuda") + + expert_output_b = torch.ones(3, hidden_dim, device="cuda", dtype=torch.bfloat16) + expert_output_b[0] = 1e20 + inverse_b = torch.tensor([[0, -1], [1, 2]], dtype=torch.int32, device="cuda") + + with set_batch_invariant_mode(True): + out_a = unpermute_tokens( + expert_output_a, + probs[:2], + permutation_map[:2], + 2, + _vt(2), + _vt(2), + batch_invariant_inverse_map=inverse_a, + ) + out_b = unpermute_tokens( + expert_output_b, + probs, + permutation_map, + 2, + _vt(3), + _vt(2), + batch_invariant_inverse_map=inverse_b, + ) + + expected = torch.full((hidden_dim,), 2.0, device="cuda") + torch.testing.assert_close(out_a[1], expected, rtol=0.0, atol=0.0) + torch.testing.assert_close(out_b[1], expected, rtol=0.0, atol=0.0) + @pytest.mark.internal class TestPermuteUnpermuteRoundtrip: diff --git a/tests/unit_tests/rl/test_rl_batch_invariant.py b/tests/unit_tests/rl/test_rl_batch_invariant.py index ab339755307..2e834d5bdaf 100644 --- a/tests/unit_tests/rl/test_rl_batch_invariant.py +++ b/tests/unit_tests/rl/test_rl_batch_invariant.py @@ -31,3 +31,86 @@ def test_selective_log_softmax_batch_invariant(): # If the kernel is batch invariant, each example's output should not depend # on its position in the batch. assert torch.equal(bik_logps, bik_logps_perm[perm.argsort()]) + + +def test_moe_unpermute_batch_invariant_inverse_map_rank_tree(): + from megatron.core import parallel_state + from megatron.core.transformer.moe.moe_utils import unpermute + + hidden = 4 + tokens = torch.tensor( + [[1e20], [1.0], [-1e20], [1.0]], device="cuda", dtype=torch.float32 + ).expand(4, hidden) + sorted_indices = torch.zeros(4, device="cuda", dtype=torch.int64) + routing_map = torch.ones(1, 4, device="cuda", dtype=torch.bool) + inverse_map = torch.tensor( + [[[0, 1, 2, 3]], [[0, 1, 2, 3]]], device="cuda", dtype=torch.int64 + ) + + parallel_state.set_expert_model_parallel_world_size(2) + try: + with set_batch_invariant_mode(True): + out = unpermute( + tokens, + sorted_indices, + (1, hidden), + routing_map=routing_map, + batch_invariant_inverse_map=inverse_map, + ) + finally: + parallel_state.set_expert_model_parallel_world_size(None) + + torch.testing.assert_close(out[0], torch.zeros(hidden, device="cuda"), rtol=0.0, atol=0.0) + + +def test_moe_batch_invariant_permute_unpermute_cuda_graph_non_padded(): + from megatron.core import parallel_state + from megatron.core.transformer.moe.moe_utils import permute, unpermute + + torch.manual_seed(123) + num_tokens, hidden, num_experts, topk = 6, 8, 4, 2 + tokens = torch.randn(num_tokens, hidden, device="cuda", dtype=torch.bfloat16) + routing_map = torch.zeros(num_tokens, num_experts, device="cuda", dtype=torch.bool) + routing_map[:, 0] = True + routing_map[:, 2] = True + probs = torch.rand(num_tokens, num_experts, device="cuda", dtype=torch.float32) + + def _run(): + permuted, _, sorted_indices, inverse_map, _ = permute( + tokens, + routing_map, + probs=probs, + num_out_tokens=num_tokens * topk, + return_batch_invariant_inverse_map=True, + ) + return unpermute( + permuted, + sorted_indices, + tokens.shape, + probs=probs, + routing_map=routing_map, + batch_invariant_inverse_map=inverse_map, + ) + + parallel_state.set_expert_model_parallel_world_size(2) + try: + with torch.no_grad(), set_batch_invariant_mode(True): + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + expected = _run() + torch.cuda.current_stream().wait_stream(stream) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_out = _run() + graph.replay() + finally: + parallel_state.set_expert_model_parallel_world_size(None) + + torch.testing.assert_close(graph_out, expected, rtol=0.0, atol=0.0) + reference = ( + tokens.float() * probs[:, 0, None] + tokens.float() * probs[:, 2, None] + ).to(tokens.dtype) + torch.testing.assert_close(graph_out, reference, rtol=0.0, atol=0.0) diff --git a/tests/unit_tests/ssm/ops/test_batch_invariant_decode.py b/tests/unit_tests/ssm/ops/test_batch_invariant_decode.py new file mode 100644 index 00000000000..c99737cae46 --- /dev/null +++ b/tests/unit_tests/ssm/ops/test_batch_invariant_decode.py @@ -0,0 +1,646 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +"""Tests for batch-invariant Mamba decode.""" + +import unittest + +import torch + +try: + from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined + + from megatron.core.ssm.ops.batch_invariant_decode import ( + BatchInvariantDecodeBuffers, + batch_invariant_decode_buffered_scan, + ) + from megatron.core.ssm.ops.ssd_combined import mamba_chunk_scan_combined_varlen + + HAVE_BATCH_INVARIANT_DECODE = True +except ImportError: + HAVE_BATCH_INVARIANT_DECODE = False + + +def _full_scan(x, dt, A, B, C, D, dt_bias, chunk_size, initial_states=None): + """Reference: a single `mamba_chunk_scan_combined` over the whole sequence.""" + y, final = mamba_chunk_scan_combined( + x, dt, A, B, C, chunk_size, + D=D, z=None, dt_bias=dt_bias, dt_softplus=True, + initial_states=initial_states, return_final_states=True, + ) + return y, final + + +@unittest.skipIf(not HAVE_BATCH_INVARIANT_DECODE, "mamba_ssm / batch_invariant_decode unavailable") +@unittest.skipIf(not torch.cuda.is_available(), "CUDA required") +class TestBatchInvariantDecodeBufferedScan(unittest.TestCase): + """Verify the batch-invariant decode scan matches a full-sequence scan bitwise.""" + + @classmethod + def setUpClass(cls): + # Pin the Mamba autotuners exactly like enable_batch_invariant_mode + # does in production: without pinning, autotune timing noise can pick + # different tile configs per process and flake the bitwise asserts. + from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( + _pin_mamba_autotuners, + ) + + _pin_mamba_autotuners() + + def setUp(self): + torch.manual_seed(0) + # No global flags: batch_invariant_decode_buffered_scan is a pure tensor-ops function + # and batch-invariant mode by design does not require + # torch.use_deterministic_algorithms. + self.device = torch.device("cuda") + self.dtype = torch.bfloat16 + # Small but non-trivial mamba dims. + self.nh = 8 + self.headdim = 32 + self.ngroups = 1 + self.dstate = 16 + self.chunk_size = 32 + self.A = -torch.exp( + torch.randn(self.nh, device=self.device, dtype=torch.float32).abs() + ) + self.D = torch.randn(self.nh, device=self.device, dtype=torch.float32) * 0.1 + self.dt_bias = ( + torch.randn(self.nh, device=self.device, dtype=torch.float32) * 0.01 + ) + + def _make_seq(self, total_len): + """Generate a (1, total_len, ...) random mamba input sequence.""" + nh, p, ng, n = self.nh, self.headdim, self.ngroups, self.dstate + return ( + torch.randn(1, total_len, nh, p, device=self.device, dtype=self.dtype) * 0.1, + torch.randn(1, total_len, nh, device=self.device, dtype=self.dtype).abs() * 0.1, + torch.randn(1, total_len, ng, n, device=self.device, dtype=self.dtype) * 0.1, + torch.randn(1, total_len, ng, n, device=self.device, dtype=self.dtype) * 0.1, + ) + + def _make_bufs(self, max_batch): + return BatchInvariantDecodeBuffers.allocate( + max_batch, self.chunk_size, + self.nh, self.headdim, self.ngroups, self.dstate, + self.device, self.dtype, + ) + + def _seed_from_prefill(self, bufs, x, dt, B, C, prefill_len, slot, max_batch): + """Run the prefill through the reference scan, store its ssm_state at + the slot, and seed the batch-invariant buffer with the partial-chunk tail.""" + # Production batch-invariant prefill keeps ssm_state at a full Mamba chunk + # boundary. Short prefills therefore keep the zero initial boundary; + # longer prefills store the largest chunk-aligned prefix state. + ssm_state = torch.zeros( + max_batch, self.nh, self.headdim, self.dstate, + device=self.device, dtype=self.dtype, + ) + if prefill_len >= self.chunk_size: + # Prefill on the largest chunk-aligned prefix; the tail goes in the buffer. + aligned = (prefill_len // self.chunk_size) * self.chunk_size + _, final = _full_scan( + x[:, :aligned], dt[:, :aligned], self.A, + B[:, :aligned], C[:, :aligned], self.D, self.dt_bias, self.chunk_size, + initial_states=None, + ) + ssm_state[slot] = final[0].to(self.dtype) + + cu = torch.tensor([0, prefill_len], dtype=torch.int32, device=self.device) + batch_indices = torch.tensor([slot], dtype=torch.int32, device=self.device) + # Buffer seeding expects the flat layout used by the mixer's prefill path. + # Here total == prefill_len since we have 1 sequence. + bufs.seed( + x[0, :prefill_len], + dt[0, :prefill_len], + B[0, :prefill_len], + C[0, :prefill_len], + cu, batch_indices, + ) + return ssm_state + + def _decode_one_step(self, bufs, x, dt, B, C, pos, slot, ssm_state): + """Call batch_invariant_decode_buffered_scan for the single token at index `pos`.""" + batch_indices = torch.tensor([slot], dtype=torch.int32, device=self.device) + return batch_invariant_decode_buffered_scan( + bufs, + x[:, pos : pos + 1], + dt[:, pos : pos + 1], + B[:, pos : pos + 1], + C[:, pos : pos + 1], + self.A, self.D, self.dt_bias, + batch_indices, + ssm_state, + ) + + def _varlen_boundary_state_from_prefill( + self, x, dt, B, C, prefill_len, initial_states=None + ): + """Production-shaped varlen prefill returning the last full chunk boundary state.""" + chunk_boundaries = [0] + pos = self.chunk_size + while pos < prefill_len: + chunk_boundaries.append(pos) + pos += self.chunk_size + chunk_boundaries.append(prefill_len) + + cu_chunk_seqlens = torch.tensor( + chunk_boundaries, dtype=torch.int32, device=self.device + ) + last_chunk_indices = torch.tensor( + [len(chunk_boundaries) - 2], dtype=torch.int32, device=self.device + ) + tail_len = prefill_len % self.chunk_size + has_boundary = prefill_len >= self.chunk_size + boundary_idx = last_chunk_indices.to(torch.long) + if tail_len != 0: + boundary_idx = boundary_idx - 1 + boundary_idx = boundary_idx.clamp(min=0) + + out = torch.zeros_like(x[0, :prefill_len]) + seq_idx = torch.zeros( + len(chunk_boundaries) - 1, dtype=torch.int32, device=self.device + ) + chunk_states = mamba_chunk_scan_combined_varlen( + x=x[0, :prefill_len], + dt=dt[0, :prefill_len], + A=self.A, + B=B[0, :prefill_len], + C=C[0, :prefill_len], + chunk_size=self.chunk_size, + cu_chunk_seqlens=cu_chunk_seqlens, + last_chunk_indices=last_chunk_indices, + seq_idx=seq_idx, + out=out, + D=self.D, + z=None, + dt_bias=self.dt_bias, + initial_states=initial_states, + return_intermediate_states=True, + dt_softplus=True, + dt_limit=(0.0, float("inf")), + state_dtype=self.dtype, + ) + final_state = chunk_states[last_chunk_indices] + boundary_state = chunk_states[boundary_idx] + if not has_boundary: + boundary_state = ( + torch.zeros_like(boundary_state) + if initial_states is None + else initial_states + ) + return final_state, boundary_state + + def _assert_bitwise(self, a, b, msg): + # bf16 outputs — bitwise-equal is the actual batch-invariant claim. + diff = (a.float() - b.float()).abs().max().item() + self.assertEqual(diff, 0.0, f"{msg}: max_abs_diff={diff:.3e}") + + def test_single_decode_matches_full_scan(self): + """Default case: prefill > chunk_size, single decode token, partial tail.""" + max_batch, slot = 4, 1 + for prefill_len in [33, 50, 95, 128]: + with self.subTest(prefill_len=prefill_len): + total = prefill_len + 1 + x, dt, B, C = self._make_seq(total) + # Reference: full scan over the whole (prefill + 1) sequence. + y_full, _ = _full_scan( + x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size, + ) + # batch-invariant: seed from prefill, then one decode step. + bufs = self._make_bufs(max_batch) + ssm_state = self._seed_from_prefill( + bufs, x, dt, B, C, prefill_len, slot, max_batch, + ) + y_batch_invariant = self._decode_one_step( + bufs, x, dt, B, C, prefill_len, slot, ssm_state, + ) + self._assert_bitwise( + y_batch_invariant[0, 0], y_full[0, prefill_len], + f"prefill_len={prefill_len}", + ) + + def test_dynamic_prefill_uses_boundary_state_not_prompt_end_state(self): + """Production prefill returns the prompt-end state too, but batch-invariant decode + must keep the cache at the last full chunk boundary and put the tail in + the replay buffer.""" + max_batch, slot = 4, 1 + for prefill_len in [31, 33, 50, 95, 128]: + with self.subTest(prefill_len=prefill_len): + total = prefill_len + 1 + x, dt, B, C = self._make_seq(total) + y_full, _ = _full_scan( + x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size, + ) + + _, boundary_state = self._varlen_boundary_state_from_prefill( + x, dt, B, C, prefill_len + ) + ssm_state = torch.randn( + max_batch, self.nh, self.headdim, self.dstate, + device=self.device, dtype=self.dtype, + ) + ssm_state[slot] = boundary_state[0].to(self.dtype) + + bufs = self._make_bufs(max_batch) + cu = torch.tensor([0, prefill_len], dtype=torch.int32, device=self.device) + batch_indices = torch.tensor([slot], dtype=torch.int32, device=self.device) + bufs.seed( + x[0, :prefill_len], + dt[0, :prefill_len], + B[0, :prefill_len], + C[0, :prefill_len], + cu, + batch_indices, + ) + y_batch_invariant = self._decode_one_step( + bufs, x, dt, B, C, prefill_len, slot, ssm_state + ) + self._assert_bitwise( + y_batch_invariant[0, 0], y_full[0, prefill_len], + f"dynamic prefill boundary state prefill_len={prefill_len}", + ) + + def test_chunked_prefill_handoff_matches_full_scan(self): + """Splitting prefill at a Mamba boundary preserves exact decode output.""" + max_batch, slot = 2, 0 + first_chunk_len = 2 * self.chunk_size + + for final_chunk_len in [20, self.chunk_size + 13]: + with self.subTest(final_chunk_len=final_chunk_len): + prefill_len = first_chunk_len + final_chunk_len + x, dt, B, C = self._make_seq(prefill_len + 1) + y_full, _ = _full_scan( + x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size, + ) + + _, first_boundary = self._varlen_boundary_state_from_prefill( + x[:, :first_chunk_len], + dt[:, :first_chunk_len], + B[:, :first_chunk_len], + C[:, :first_chunk_len], + first_chunk_len, + ) + _, final_boundary = self._varlen_boundary_state_from_prefill( + x[:, first_chunk_len:prefill_len], + dt[:, first_chunk_len:prefill_len], + B[:, first_chunk_len:prefill_len], + C[:, first_chunk_len:prefill_len], + final_chunk_len, + initial_states=first_boundary, + ) + + ssm_state = torch.zeros( + max_batch, self.nh, self.headdim, self.dstate, + device=self.device, dtype=self.dtype, + ) + ssm_state[slot] = final_boundary[0] + bufs = self._make_bufs(max_batch) + cu = torch.tensor( + [0, final_chunk_len], dtype=torch.int32, device=self.device + ) + bufs.seed( + x[0, first_chunk_len:prefill_len], + dt[0, first_chunk_len:prefill_len], + B[0, first_chunk_len:prefill_len], + C[0, first_chunk_len:prefill_len], + cu, + torch.tensor([slot], dtype=torch.int32, device=self.device), + ) + y_batch_invariant = self._decode_one_step( + bufs, x, dt, B, C, prefill_len, slot, ssm_state + ) + self._assert_bitwise( + y_batch_invariant[0, 0], + y_full[0, prefill_len], + f"chunked prefill final_chunk_len={final_chunk_len}", + ) + + def test_seed_ignores_nonfinite_physical_padding_rows(self): + """Dynamic prefill can carry padded physical token rows after the real + prefix. Seed must duplicate a valid per-sequence tail token into unused + replay-buffer rows; otherwise masked future rows can still poison the + row-gated Triton dot as 0 * NaN.""" + max_batch, slot = 4, 0 + prefill_len = self.chunk_size + 1 + total = prefill_len + 1 + x, dt, B, C = self._make_seq(total) + y_full, _ = _full_scan( + x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size, + ) + + _, boundary_state = self._varlen_boundary_state_from_prefill(x, dt, B, C, prefill_len) + ssm_state = torch.zeros( + max_batch, self.nh, self.headdim, self.dstate, + device=self.device, dtype=self.dtype, + ) + ssm_state[slot] = boundary_state[0].to(self.dtype) + + nan_x = torch.full_like(x[0, :1], float("nan")) + nan_dt = torch.full_like(dt[0, :1], float("nan")) + nan_B = torch.full_like(B[0, :1], float("nan")) + nan_C = torch.full_like(C[0, :1], float("nan")) + + bufs = self._make_bufs(max_batch) + cu = torch.tensor([0, prefill_len], dtype=torch.int32, device=self.device) + batch_indices = torch.tensor([slot], dtype=torch.int32, device=self.device) + bufs.seed( + torch.cat([x[0, :prefill_len], nan_x], dim=0), + torch.cat([dt[0, :prefill_len], nan_dt], dim=0), + torch.cat([B[0, :prefill_len], nan_B], dim=0), + torch.cat([C[0, :prefill_len], nan_C], dim=0), + cu, + batch_indices, + ) + self.assertTrue(torch.isfinite(bufs.x[slot]).all()) + self.assertTrue(torch.isfinite(bufs.dt[slot]).all()) + self.assertTrue(torch.isfinite(bufs.B[slot]).all()) + self.assertTrue(torch.isfinite(bufs.C[slot]).all()) + + y_batch_invariant = self._decode_one_step( + bufs, x, dt, B, C, prefill_len, slot, ssm_state + ) + self._assert_bitwise( + y_batch_invariant[0, 0], y_full[0, prefill_len], + "nonfinite physical padding rows", + ) + + def test_short_prefill_uses_zero_boundary_state(self): + """prefill_len < chunk_size: decode replays from the zero boundary.""" + max_batch, slot = 2, 0 + for prefill_len in [1, 7, 16, 31]: + with self.subTest(prefill_len=prefill_len): + total = prefill_len + 1 + x, dt, B, C = self._make_seq(total) + y_full, _ = _full_scan( + x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size, + ) + bufs = self._make_bufs(max_batch) + ssm_state = self._seed_from_prefill( + bufs, x, dt, B, C, prefill_len, slot, max_batch, + ) + y_batch_invariant = self._decode_one_step( + bufs, x, dt, B, C, prefill_len, slot, ssm_state, + ) + self._assert_bitwise( + y_batch_invariant[0, 0], y_full[0, prefill_len], + f"prefill_len={prefill_len}", + ) + + def test_multi_step_decode_across_chunk_boundary(self): + """Step decode several times so the per-slot buffer fills, crosses + a chunk boundary, and resets. Each step must match the full scan.""" + max_batch, slot = 2, 0 + prefill_len = 20 # < chunk_size, so first decode step will keep growing buf + n_decode = self.chunk_size + 5 # enough to cross at least one boundary + total = prefill_len + n_decode + x, dt, B, C = self._make_seq(total) + y_full, _ = _full_scan( + x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size, + ) + + bufs = self._make_bufs(max_batch) + ssm_state = self._seed_from_prefill( + bufs, x, dt, B, C, prefill_len, slot, max_batch, + ) + + for k in range(n_decode): + pos = prefill_len + k + y_batch_invariant = self._decode_one_step(bufs, x, dt, B, C, pos, slot, ssm_state) + self._assert_bitwise( + y_batch_invariant[0, 0], y_full[0, pos], + f"step k={k} (pos={pos}, num_buffered_before={bufs.num_buffered[slot].item()})", + ) + + def test_multi_slot_independent_streams(self): + """Two slots with different prefill lengths decoded in the same call — + each slot's output must match its own full scan.""" + max_batch = 4 + slots = [0, 2] + prefill_lens = [25, 70] # one short, one long with a boundary state + x_per_slot, dt_per_slot, B_per_slot, C_per_slot = [], [], [], [] + y_refs = [] + for plen in prefill_lens: + x, dt, B, C = self._make_seq(plen + 1) + x_per_slot.append(x); dt_per_slot.append(dt) + B_per_slot.append(B); C_per_slot.append(C) + y_full, _ = _full_scan( + x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size, + ) + y_refs.append(y_full[0, plen]) + + bufs = self._make_bufs(max_batch) + # Per-slot seeding (each slot's prefill done independently). + ssm_state = torch.zeros( + max_batch, self.nh, self.headdim, self.dstate, + device=self.device, dtype=self.dtype, + ) + for slot, plen, x, dt, B, C in zip( + slots, prefill_lens, x_per_slot, dt_per_slot, B_per_slot, C_per_slot, + ): + partial = self._seed_from_prefill( + bufs, x, dt, B, C, plen, slot, max_batch, + ) + ssm_state[slot] = partial[slot] + + # Both slots step at once. + x_step = torch.cat( + [x_per_slot[i][:, prefill_lens[i] : prefill_lens[i] + 1] for i in range(2)], + dim=0, + ) + dt_step = torch.cat( + [dt_per_slot[i][:, prefill_lens[i] : prefill_lens[i] + 1] for i in range(2)], + dim=0, + ) + B_step = torch.cat( + [B_per_slot[i][:, prefill_lens[i] : prefill_lens[i] + 1] for i in range(2)], + dim=0, + ) + C_step = torch.cat( + [C_per_slot[i][:, prefill_lens[i] : prefill_lens[i] + 1] for i in range(2)], + dim=0, + ) + batch_indices = torch.tensor(slots, dtype=torch.int32, device=self.device) + y_batch_invariant = batch_invariant_decode_buffered_scan( + bufs, x_step, dt_step, B_step, C_step, + self.A, self.D, self.dt_bias, batch_indices, ssm_state, + ) + for i, plen in enumerate(prefill_lens): + self._assert_bitwise( + y_batch_invariant[i, 0], y_refs[i], + f"multi-slot slot={slots[i]} prefill_len={plen}", + ) + + def test_inactive_padding_entries(self): + """batch_indices mixing -1 padding entries with active slot 0 (the CUDA- + graph padding pattern). Padding entries must not perturb slot 0's buffer + or output — they are redirected to the buffers' trash row.""" + max_batch, slot = 2, 0 + prefill_len = 50 + n_decode = 8 + total = prefill_len + n_decode + x, dt, B, C = self._make_seq(total) + y_full, _ = _full_scan( + x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size, + ) + + bufs = self._make_bufs(max_batch) + ssm_state = self._seed_from_prefill( + bufs, x, dt, B, C, prefill_len, slot, max_batch, + ) + batch_indices = torch.tensor([slot, -1, -1], dtype=torch.int32, device=self.device) + for k in range(n_decode): + pos = prefill_len + k + # Entry 0 carries the real token; padding entries carry garbage. + def pad3(t): + junk = torch.randn( + 2, *t.shape[1:], device=t.device, dtype=t.dtype + ) + return torch.cat([t, junk], dim=0) + + y_batch_invariant = batch_invariant_decode_buffered_scan( + bufs, + pad3(x[:, pos : pos + 1]), + pad3(dt[:, pos : pos + 1]), + pad3(B[:, pos : pos + 1]), + pad3(C[:, pos : pos + 1]), + self.A, self.D, self.dt_bias, + batch_indices, + ssm_state, + ) + self._assert_bitwise( + y_batch_invariant[0, 0], y_full[0, pos], f"padded step k={k}" + ) + # Padding entries must return zeros. + self.assertEqual(y_batch_invariant[1:].abs().max().item(), 0.0) + + def test_cuda_graph_replay_matches_full_scan(self): + """A captured decode step advances persistent state exactly across replays.""" + max_batch, slot = 2, 0 + prefill_len = 20 + x, dt, B, C = self._make_seq(prefill_len + 2) + y_full, _ = _full_scan( + x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size, + ) + + # Compile Triton before capture without touching the graph's buffers. + warmup_bufs = self._make_bufs(max_batch) + warmup_state = self._seed_from_prefill( + warmup_bufs, x, dt, B, C, prefill_len, slot, max_batch, + ) + self._decode_one_step( + warmup_bufs, x, dt, B, C, prefill_len, slot, warmup_state + ) + + bufs = self._make_bufs(max_batch) + ssm_state = self._seed_from_prefill( + bufs, x, dt, B, C, prefill_len, slot, max_batch, + ) + batch_indices = torch.tensor([slot], dtype=torch.int32, device=self.device) + static_x = x[:, prefill_len : prefill_len + 1].clone() + static_dt = dt[:, prefill_len : prefill_len + 1].clone() + static_B = B[:, prefill_len : prefill_len + 1].clone() + static_C = C[:, prefill_len : prefill_len + 1].clone() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_output = batch_invariant_decode_buffered_scan( + bufs, + static_x, + static_dt, + static_B, + static_C, + self.A, + self.D, + self.dt_bias, + batch_indices, + ssm_state, + ) + + # Capture executes once, so restore the replay cursor before the first replay. + cu = torch.tensor([0, prefill_len], dtype=torch.int32, device=self.device) + bufs.seed( + x[0, :prefill_len], + dt[0, :prefill_len], + B[0, :prefill_len], + C[0, :prefill_len], + cu, + batch_indices, + ) + graph.replay() + self._assert_bitwise( + graph_output[0, 0], y_full[0, prefill_len], "CUDA graph replay step 0" + ) + + static_x.copy_(x[:, prefill_len + 1 : prefill_len + 2]) + static_dt.copy_(dt[:, prefill_len + 1 : prefill_len + 2]) + static_B.copy_(B[:, prefill_len + 1 : prefill_len + 2]) + static_C.copy_(C[:, prefill_len + 1 : prefill_len + 2]) + graph.replay() + self._assert_bitwise( + graph_output[0, 0], y_full[0, prefill_len + 1], "CUDA graph replay step 1" + ) + + def test_crossing_with_dominant_carried_state(self): + """Boundary crossing where the carried state dominates the output + (weak decay: A ~ -0.01 → exp(dA_cs) ≈ 1). Guards the pipeline + ordering: the scan must consume the PRE-step boundary state, not the + one the fused snapshot writes during the same call — with strong + decay that corruption can round away in bf16 and hide.""" + max_batch, slot = 2, 0 + prefill_len = 20 + n_decode = self.chunk_size + 5 + total = prefill_len + n_decode + x, dt, B, C = self._make_seq(total) + + weak_A = self.A * 0.01 + y_full, _ = mamba_chunk_scan_combined( + x, dt, weak_A, B, C, self.chunk_size, + D=self.D, z=None, dt_bias=self.dt_bias, dt_softplus=True, + initial_states=None, return_final_states=True, + ) + + bufs = self._make_bufs(max_batch) + ssm_state = torch.zeros( + max_batch, self.nh, self.headdim, self.dstate, + device=self.device, dtype=self.dtype, + ) + cu = torch.tensor([0, prefill_len], dtype=torch.int32, device=self.device) + batch_indices = torch.tensor([slot], dtype=torch.int32, device=self.device) + bufs.seed( + x[0, :prefill_len], dt[0, :prefill_len], + B[0, :prefill_len], C[0, :prefill_len], cu, batch_indices, + ) + for k in range(n_decode): + pos = prefill_len + k + y_batch_invariant = batch_invariant_decode_buffered_scan( + bufs, + x[:, pos : pos + 1], dt[:, pos : pos + 1], + B[:, pos : pos + 1], C[:, pos : pos + 1], + weak_A, self.D, self.dt_bias, + batch_indices, ssm_state, + ) + self._assert_bitwise( + y_batch_invariant[0, 0], y_full[0, pos], f"weak-decay step k={k}" + ) + + def test_deterministic_across_calls(self): + """Same inputs → bitwise-identical output across repeated invocations.""" + max_batch, slot = 2, 0 + prefill_len = 50 + total = prefill_len + 1 + x, dt, B, C = self._make_seq(total) + + outs = [] + for _ in range(3): + bufs = self._make_bufs(max_batch) + ssm_state = self._seed_from_prefill( + bufs, x, dt, B, C, prefill_len, slot, max_batch, + ) + outs.append( + self._decode_one_step(bufs, x, dt, B, C, prefill_len, slot, ssm_state) + ) + for i in range(1, len(outs)): + self.assertTrue(torch.equal(outs[0], outs[i]), + f"determinism: run {i} differs from run 0") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit_tests/transformer/moe/test_moe_batch_invariant.py b/tests/unit_tests/transformer/moe/test_moe_batch_invariant.py new file mode 100644 index 00000000000..b9a848068ca --- /dev/null +++ b/tests/unit_tests/transformer/moe/test_moe_batch_invariant.py @@ -0,0 +1,259 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Tests for batch-invariant MoE grouped GEMM.""" +import pytest +import torch +import torch.nn.functional as F + +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( + HAVE_DEEPGEMM_BF16, + _bf16_grouped_gemm_contiguous, + _m_splits_to_m_indices, + _offs_to_m_indices, + set_batch_invariant_mode, +) + + +def _hopper_or_newer() -> bool: + if not torch.cuda.is_available(): + return False + major, _ = torch.cuda.get_device_capability() + return major >= 9 + + +pytestmark = [ + pytest.mark.skipif( + not HAVE_DEEPGEMM_BF16, + reason="DeepGEMM with bf16 grouped bindings is required for MoE batch-invariant tests.", + ), + pytest.mark.skipif( + not _hopper_or_newer(), reason="DeepGEMM bf16 grouped kernels require Hopper (sm_90+)." + ), +] + + +# --------------------------------------------------------------------------- +# Index-conversion helpers +# --------------------------------------------------------------------------- + + +def test_m_splits_to_m_indices_basic(): + m_splits = [3, 0, 5, 2] + m_total = sum(m_splits) + out = _m_splits_to_m_indices(m_splits, torch.device("cuda"), m_total) + expected = torch.tensor([0, 0, 0, 2, 2, 2, 2, 2, 3, 3], dtype=torch.int32, device="cuda") + assert torch.equal(out, expected) + + +def test_offs_to_m_indices_basic(): + # Three experts with 4/2/3 tokens, plus 1 row of post-padding (-1). + offs = torch.tensor([4, 6, 9], dtype=torch.int32, device="cuda") + m_total = 10 # one trailing pad row past offs[-1]=9 + out = _offs_to_m_indices(offs, m_total) + expected = torch.tensor([0, 0, 0, 0, 1, 1, 2, 2, 2, -1], dtype=torch.int32, device="cuda") + assert torch.equal(out, expected) + + +# --------------------------------------------------------------------------- +# Kernel-level invariance +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("E", [2, 4, 8]) +def test_grouped_gemm_split_invariance(E): + """Splitting the M dimension at expert boundaries must give bitwise-identical + output to the full call.""" + torch.manual_seed(0) + K, N = 128, 96 + # Build expert-grouped tokens: 32 tokens per expert + per_expert = 32 + M = per_expert * E + x = torch.randn(M, K, device="cuda", dtype=torch.bfloat16) + w = torch.randn(E, N, K, device="cuda", dtype=torch.bfloat16) + m_indices = torch.repeat_interleave( + torch.arange(E, device="cuda", dtype=torch.int32), + torch.tensor([per_expert] * E, device="cuda", dtype=torch.int32), + ).contiguous() + + counts = [per_expert] * E + y_full = _bf16_grouped_gemm_contiguous(x, w, m_indices, counts) + + # Split into halves at expert boundaries (mid-expert split is not legal for + # contiguous-layout DeepGEMM — must split on a boundary). + half = (E // 2) * per_expert + y0 = _bf16_grouped_gemm_contiguous( + x[:half].contiguous(), + w, + m_indices[:half].contiguous(), + counts[: E // 2] + [0] * (E // 2), + ) + y1 = _bf16_grouped_gemm_contiguous( + x[half:].contiguous(), + w, + m_indices[half:].contiguous(), + [0] * (E // 2) + counts[E // 2 :], + ) + y_cat = torch.cat([y0, y1], dim=0) + assert torch.equal( + y_full, y_cat + ), f"max abs diff: {(y_full.float() - y_cat.float()).abs().max().item()}" + + +def test_grouped_gemm_per_expert_token_count_invariance(): + """For a fixed expert id, the per-row output must be identical regardless of + how many *other* expert rows surround it in the batch.""" + torch.manual_seed(1) + E, K, N = 4, 64, 48 + w = torch.randn(E, N, K, device="cuda", dtype=torch.bfloat16) + x_target = torch.randn(8, K, device="cuda", dtype=torch.bfloat16) + + # Layout A: just expert 1's tokens. + m_indices_A = torch.full((8,), 1, dtype=torch.int32, device="cuda") + y_A = _bf16_grouped_gemm_contiguous(x_target, w, m_indices_A, [0, 8, 0, 0]) + + # Layout B: expert 0 (16 rows), then expert 1 (8 rows, same x_target), + # then expert 3 (12 rows). + x_pad0 = torch.randn(16, K, device="cuda", dtype=torch.bfloat16) + x_pad3 = torch.randn(12, K, device="cuda", dtype=torch.bfloat16) + x_B = torch.cat([x_pad0, x_target, x_pad3], dim=0).contiguous() + m_indices_B = torch.cat( + [ + torch.full((16,), 0, dtype=torch.int32, device="cuda"), + torch.full((8,), 1, dtype=torch.int32, device="cuda"), + torch.full((12,), 3, dtype=torch.int32, device="cuda"), + ], + dim=0, + ).contiguous() + y_B = _bf16_grouped_gemm_contiguous(x_B, w, m_indices_B, [16, 8, 0, 12]) + + # The 8 rows assigned to expert 1 inside y_B must match y_A bitwise. + y_B_target = y_B[16 : 16 + 8] + assert torch.equal(y_A, y_B_target) + + +# --------------------------------------------------------------------------- +# End-to-end: TEGroupedMLP batch-invariance +# --------------------------------------------------------------------------- + + +@pytest.fixture +def _moe_env(): + """Spin up a tiny model-parallel env for MoELayer construction.""" + from megatron.core.utils import is_te_min_version + from tests.unit_tests.test_utilities import Utils + + if not is_te_min_version("1.9.0.dev0"): + pytest.skip("TE GroupedLinear requires TE >= 1.9.0.dev0") + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + try: + yield + finally: + Utils.destroy_model_parallel() + + +def _build_moe_layer(hidden_size=64, ffn=128, num_experts=4, topk=1): + from megatron.core.models.gpt.gpt_layer_specs import ( + get_gpt_layer_with_transformer_engine_submodules, + ) + from megatron.core.transformer.enums import AttnBackend + from megatron.core.transformer.moe.moe_layer import MoELayer + from megatron.core.transformer.spec_utils import get_submodules + from megatron.core.transformer.transformer_config import TransformerConfig + + cfg = TransformerConfig( + num_layers=1, + hidden_size=hidden_size, + num_attention_heads=4, + num_moe_experts=num_experts, + moe_ffn_hidden_size=ffn, + moe_grouped_gemm=True, + moe_router_topk=topk, + moe_token_dispatcher_type="alltoall", + moe_router_load_balancing_type="sinkhorn", + gated_linear_unit=False, + activation_func=F.gelu, + add_bias_linear=False, + params_dtype=torch.bfloat16, + bf16=True, + attention_backend=AttnBackend.flash, + batch_invariant_mode=True, + ) + submodules = get_submodules( + get_gpt_layer_with_transformer_engine_submodules( + cfg.num_moe_experts, moe_grouped_gemm=True + ).mlp + ) + return MoELayer(cfg, submodules).cuda().eval(), cfg + + +def test_tegroupedmlp_batch_invariant_split(_moe_env): + """Splitting the batch and concatenating outputs must give bitwise-identical + results to the full batch — the basic batch-invariance contract.""" + from megatron.core.transformer.moe.experts import TEGroupedMLP + + layer, cfg = _build_moe_layer() + assert isinstance(layer.experts, TEGroupedMLP) + + torch.manual_seed(0) + M = 48 + x = torch.randn(M, 1, cfg.hidden_size, device="cuda", dtype=torch.bfloat16) + + with torch.no_grad(), set_batch_invariant_mode(True): + y_full, _ = layer(x) + y0, _ = layer(x[: M // 2]) + y1, _ = layer(x[M // 2 :]) + y_cat = torch.cat([y0, y1], dim=0) + assert torch.equal(y_full, y_cat), ( + f"TEGroupedMLP not batch-invariant under halving; max abs diff: " + f"{(y_full.float() - y_cat.float()).abs().max().item()}" + ) + + +def test_tegroupedmlp_per_token_invariance_across_batch_sizes(_moe_env): + """The strongest batch-invariance check: a fixed set of "target" tokens must + produce the *exact same output* regardless of what other tokens surround + them in the batch. This is what RL log-prob parity needs.""" + layer, cfg = _build_moe_layer() + torch.manual_seed(1) + + # 8 target tokens whose outputs we lock in by running them alone. + target = torch.randn(8, 1, cfg.hidden_size, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(), set_batch_invariant_mode(True): + y_target_alone, _ = layer(target) + + # Now embed those same 8 tokens at different positions inside larger batches + # of varying sizes, with random surrounding tokens. + for pad_left, pad_right in [(0, 16), (40, 0), (24, 24), (5, 13), (1, 1)]: + left = torch.randn(pad_left, 1, cfg.hidden_size, device="cuda", dtype=torch.bfloat16) + right = torch.randn(pad_right, 1, cfg.hidden_size, device="cuda", dtype=torch.bfloat16) + big = torch.cat([left, target, right], dim=0) + with torch.no_grad(), set_batch_invariant_mode(True): + y_big, _ = layer(big) + y_target_in_big = y_big[pad_left : pad_left + 8] + assert torch.equal(y_target_alone, y_target_in_big), ( + f"Per-token output drifted when batch shape changed " + f"(pad_left={pad_left}, pad_right={pad_right}); " + f"max abs diff: " + f"{(y_target_alone.float() - y_target_in_big.float()).abs().max().item()}" + ) + + +def test_tegroupedmlp_invariance_under_permutation(_moe_env): + """Permuting the input batch and undoing the permutation in the output + yields bitwise-identical results. Different routing distribution per + micro-batch position, same kernel output.""" + layer, cfg = _build_moe_layer() + torch.manual_seed(2) + M = 32 + x = torch.randn(M, 1, cfg.hidden_size, device="cuda", dtype=torch.bfloat16) + + perm = torch.randperm(M, device="cuda") + with torch.no_grad(), set_batch_invariant_mode(True): + y_ref, _ = layer(x) + y_perm, _ = layer(x[perm]) + y_unperm = y_perm[perm.argsort()] + assert torch.equal(y_ref, y_unperm), ( + f"MoE output not invariant to batch permutation; max abs diff: " + f"{(y_ref.float() - y_unperm.float()).abs().max().item()}" + ) From 11c78df45fbc6a2ef61900a6d17857b5e4ede3e2 Mon Sep 17 00:00:00 2001 From: root Date: Thu, 23 Jul 2026 06:48:38 -0700 Subject: [PATCH 02/14] Preserve FP32 Mamba state across decode chunks Signed-off-by: root --- megatron/core/inference/config.py | 13 ++- .../inference/contexts/dynamic_context.py | 4 + .../core/ssm/ops/batch_invariant_decode.py | 4 + .../ssm/ops/test_batch_invariant_decode.py | 109 ++++++++++++------ 4 files changed, 95 insertions(+), 35 deletions(-) diff --git a/megatron/core/inference/config.py b/megatron/core/inference/config.py index e8769f3d6e7..d02e358e569 100644 --- a/megatron/core/inference/config.py +++ b/megatron/core/inference/config.py @@ -37,7 +37,7 @@ class MambaInferenceStateConfig: """The dtype to use for the Mamba conv state tensor. Defaults to the model dtype.""" ssm_states_dtype: torch.dtype - """The dtype to use for the Mamba SSM state tensor. Defaults to the model dtype.""" + """The dtype to use for Mamba SSM state. Batch-invariant mode requires FP32.""" mamba_chunk_size: int = 128 """The chunk size used by the Mamba SSM Triton kernels.""" @@ -60,7 +60,16 @@ def from_model( ) if conv_states_dtype is None: conv_states_dtype = model.config.params_dtype - if ssm_states_dtype is None: + if model.config.batch_invariant_mode: + if ssm_states_dtype not in (None, torch.float32): + raise ValueError( + "batch_invariant_mode requires FP32 Mamba SSM states; " + f"got {ssm_states_dtype}." + ) + # State passing carries an unrounded FP32 boundary value across + # chunks. Rounding the cache to BF16 changes the next transition. + ssm_states_dtype = torch.float32 + elif ssm_states_dtype is None: ssm_states_dtype = model.config.params_dtype mamba_chunk_size = 128 for layer_type, layer in zip(decoder.layer_type_list, decoder.layers): diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index f00e52bb4c9..eef448852a0 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -350,6 +350,10 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC "batch_invariant_mode for Mamba dynamic inference only supports " "one-token decode; set num_speculative_tokens=0." ) + assert self.mamba_ssm_states_dtype == torch.float32, ( + "batch_invariant_mode requires FP32 Mamba SSM states so state-passing " + "boundaries are not rounded between decode chunks." + ) # For hybrid models, the layer map converts the global layer index to the # corresponding attention layer index or Mamba layer index depending on the diff --git a/megatron/core/ssm/ops/batch_invariant_decode.py b/megatron/core/ssm/ops/batch_invariant_decode.py index 09fe5ff5e6d..3244cf4874b 100644 --- a/megatron/core/ssm/ops/batch_invariant_decode.py +++ b/megatron/core/ssm/ops/batch_invariant_decode.py @@ -134,6 +134,10 @@ def batch_invariant_decode_buffered_scan( "batch-invariant Mamba decode assumes one new token per request " "per call (no speculative decoding)." ) + assert ssm_state.dtype == torch.float32, ( + "batch-invariant Mamba decode requires an FP32 SSM state cache to preserve " + "the state-passing carry across chunk boundaries." + ) output_capacity = buffers.out.shape[0] assert decode_batch_size <= output_capacity, ( f"decode batch size {decode_batch_size} exceeds the output buffer capacity " diff --git a/tests/unit_tests/ssm/ops/test_batch_invariant_decode.py b/tests/unit_tests/ssm/ops/test_batch_invariant_decode.py index c99737cae46..9e9191ff541 100644 --- a/tests/unit_tests/ssm/ops/test_batch_invariant_decode.py +++ b/tests/unit_tests/ssm/ops/test_batch_invariant_decode.py @@ -83,16 +83,24 @@ def _make_bufs(self, max_batch): self.device, self.dtype, ) + def _make_ssm_state(self, max_batch): + """Production BIK state cache: FP32 carry across Mamba chunks.""" + return torch.zeros( + max_batch, + self.nh, + self.headdim, + self.dstate, + device=self.device, + dtype=torch.float32, + ) + def _seed_from_prefill(self, bufs, x, dt, B, C, prefill_len, slot, max_batch): """Run the prefill through the reference scan, store its ssm_state at the slot, and seed the batch-invariant buffer with the partial-chunk tail.""" # Production batch-invariant prefill keeps ssm_state at a full Mamba chunk # boundary. Short prefills therefore keep the zero initial boundary; # longer prefills store the largest chunk-aligned prefix state. - ssm_state = torch.zeros( - max_batch, self.nh, self.headdim, self.dstate, - device=self.device, dtype=self.dtype, - ) + ssm_state = self._make_ssm_state(max_batch) if prefill_len >= self.chunk_size: # Prefill on the largest chunk-aligned prefix; the tail goes in the buffer. aligned = (prefill_len // self.chunk_size) * self.chunk_size @@ -101,7 +109,7 @@ def _seed_from_prefill(self, bufs, x, dt, B, C, prefill_len, slot, max_batch): B[:, :aligned], C[:, :aligned], self.D, self.dt_bias, self.chunk_size, initial_states=None, ) - ssm_state[slot] = final[0].to(self.dtype) + ssm_state[slot] = final[0] cu = torch.tensor([0, prefill_len], dtype=torch.int32, device=self.device) batch_indices = torch.tensor([slot], dtype=torch.int32, device=self.device) @@ -176,7 +184,7 @@ def _varlen_boundary_state_from_prefill( return_intermediate_states=True, dt_softplus=True, dt_limit=(0.0, float("inf")), - state_dtype=self.dtype, + state_dtype=torch.float32, ) final_state = chunk_states[last_chunk_indices] boundary_state = chunk_states[boundary_idx] @@ -217,6 +225,55 @@ def test_single_decode_matches_full_scan(self): f"prefill_len={prefill_len}", ) + def test_rejects_bf16_state_cache(self): + """A rounded state cache cannot preserve carry across multiple chunks.""" + x, dt, B, C = self._make_seq(1) + bufs = self._make_bufs(max_batch=2) + ssm_state = torch.zeros( + 2, + self.nh, + self.headdim, + self.dstate, + device=self.device, + dtype=torch.bfloat16, + ) + with self.assertRaisesRegex(AssertionError, "requires an FP32 SSM state cache"): + self._decode_one_step(bufs, x, dt, B, C, pos=0, slot=0, ssm_state=ssm_state) + + def test_inference_config_uses_fp32_state_cache(self): + """BIK model-derived inference config cannot select a rounded state dtype.""" + from types import SimpleNamespace + from unittest.mock import patch + + from megatron.core.inference.config import MambaInferenceStateConfig + from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols + + model = SimpleNamespace( + config=SimpleNamespace( + batch_invariant_mode=True, + params_dtype=torch.bfloat16, + ) + ) + decoder = SimpleNamespace( + layer_type_list=[Symbols.MAMBA], + layers=[SimpleNamespace(mixer=SimpleNamespace(chunk_size=self.chunk_size))], + mamba_state_shapes_per_request=lambda: ((4, 8), (8, 32, 16)), + ) + with patch( + "megatron.core.inference.config.get_attr_wrapped_model", + return_value=decoder, + ): + config = MambaInferenceStateConfig.from_model(model) + self.assertEqual(config.ssm_states_dtype, torch.float32) + with self.assertRaisesRegex(ValueError, "requires FP32 Mamba SSM states"): + MambaInferenceStateConfig.from_model( + model, + ssm_states_dtype=torch.bfloat16, + ) + model.config.batch_invariant_mode = False + config = MambaInferenceStateConfig.from_model(model) + self.assertEqual(config.ssm_states_dtype, torch.bfloat16) + def test_dynamic_prefill_uses_boundary_state_not_prompt_end_state(self): """Production prefill returns the prompt-end state too, but batch-invariant decode must keep the cache at the last full chunk boundary and put the tail in @@ -233,11 +290,8 @@ def test_dynamic_prefill_uses_boundary_state_not_prompt_end_state(self): _, boundary_state = self._varlen_boundary_state_from_prefill( x, dt, B, C, prefill_len ) - ssm_state = torch.randn( - max_batch, self.nh, self.headdim, self.dstate, - device=self.device, dtype=self.dtype, - ) - ssm_state[slot] = boundary_state[0].to(self.dtype) + ssm_state = torch.randn_like(self._make_ssm_state(max_batch)) + ssm_state[slot] = boundary_state[0] bufs = self._make_bufs(max_batch) cu = torch.tensor([0, prefill_len], dtype=torch.int32, device=self.device) @@ -287,10 +341,7 @@ def test_chunked_prefill_handoff_matches_full_scan(self): initial_states=first_boundary, ) - ssm_state = torch.zeros( - max_batch, self.nh, self.headdim, self.dstate, - device=self.device, dtype=self.dtype, - ) + ssm_state = self._make_ssm_state(max_batch) ssm_state[slot] = final_boundary[0] bufs = self._make_bufs(max_batch) cu = torch.tensor( @@ -327,11 +378,8 @@ def test_seed_ignores_nonfinite_physical_padding_rows(self): ) _, boundary_state = self._varlen_boundary_state_from_prefill(x, dt, B, C, prefill_len) - ssm_state = torch.zeros( - max_batch, self.nh, self.headdim, self.dstate, - device=self.device, dtype=self.dtype, - ) - ssm_state[slot] = boundary_state[0].to(self.dtype) + ssm_state = self._make_ssm_state(max_batch) + ssm_state[slot] = boundary_state[0] nan_x = torch.full_like(x[0, :1], float("nan")) nan_dt = torch.full_like(dt[0, :1], float("nan")) @@ -428,10 +476,7 @@ def test_multi_slot_independent_streams(self): bufs = self._make_bufs(max_batch) # Per-slot seeding (each slot's prefill done independently). - ssm_state = torch.zeros( - max_batch, self.nh, self.headdim, self.dstate, - device=self.device, dtype=self.dtype, - ) + ssm_state = self._make_ssm_state(max_batch) for slot, plen, x, dt, B, C in zip( slots, prefill_lens, x_per_slot, dt_per_slot, B_per_slot, C_per_slot, ): @@ -579,14 +624,15 @@ def test_cuda_graph_replay_matches_full_scan(self): ) def test_crossing_with_dominant_carried_state(self): - """Boundary crossing where the carried state dominates the output + """Repeated boundary crossings where the carried state dominates the output (weak decay: A ~ -0.01 → exp(dA_cs) ≈ 1). Guards the pipeline - ordering: the scan must consume the PRE-step boundary state, not the - one the fused snapshot writes during the same call — with strong - decay that corruption can round away in bf16 and hide.""" + ordering and FP32 state-passing carry. With strong decay, either + corruption can round away in BF16 and hide.""" max_batch, slot = 2, 0 prefill_len = 20 - n_decode = self.chunk_size + 5 + # Cross twice: the second transition detects an accidental BF16 + # store/reload of state passing's FP32 carry. + n_decode = 2 * self.chunk_size + 5 total = prefill_len + n_decode x, dt, B, C = self._make_seq(total) @@ -598,10 +644,7 @@ def test_crossing_with_dominant_carried_state(self): ) bufs = self._make_bufs(max_batch) - ssm_state = torch.zeros( - max_batch, self.nh, self.headdim, self.dstate, - device=self.device, dtype=self.dtype, - ) + ssm_state = self._make_ssm_state(max_batch) cu = torch.tensor([0, prefill_len], dtype=torch.int32, device=self.device) batch_indices = torch.tensor([slot], dtype=torch.int32, device=self.device) bufs.seed( From 79df664a321a35bfb0854a55a67c82941ba1bd84 Mon Sep 17 00:00:00 2001 From: root Date: Thu, 23 Jul 2026 07:31:30 -0700 Subject: [PATCH 03/14] Clarify batch-invariant decoder accessor Signed-off-by: root --- .../torch_symm_triton/variable_collectives.py | 4 +- .../core/inference/moe/batch_invariant.py | 9 +- megatron/core/inference/moe/permute.py | 16 +- megatron/core/ssm/mamba_mixer.py | 9 +- .../core/ssm/ops/batch_invariant_decode.py | 26 +- megatron/core/ssm/ops/ssd_bmm.py | 6 +- megatron/core/ssm/ops/ssd_chunk_scan.py | 16 +- megatron/core/ssm/ops/ssd_chunk_state.py | 6 +- megatron/core/ssm/ops/ssd_state_passing.py | 14 +- .../custom_layers/batch_invariant_kernels.py | 26 +- .../core/transformer/moe/batch_invariant.py | 14 +- megatron/core/transformer/moe/moe_utils.py | 13 +- .../moe/token_dispatcher_inference.py | 6 +- .../core/transformer/transformer_config.py | 7 +- .../contexts/test_dynamic_prefix_caching.py | 4 +- .../test_moe_dispatching_and_routing.py | 6 +- .../unit_tests/rl/test_rl_batch_invariant.py | 10 +- .../ssm/ops/test_batch_invariant_decode.py | 294 ++++++++---------- .../moe/test_moe_batch_invariant.py | 10 +- 19 files changed, 208 insertions(+), 288 deletions(-) diff --git a/megatron/core/inference/communication/torch_symm_triton/variable_collectives.py b/megatron/core/inference/communication/torch_symm_triton/variable_collectives.py index e1b39716f10..dcf8ec4924d 100644 --- a/megatron/core/inference/communication/torch_symm_triton/variable_collectives.py +++ b/megatron/core/inference/communication/torch_symm_triton/variable_collectives.py @@ -405,7 +405,9 @@ def _ordered_reduce_scatter_v_kernel( for src_rank in tl.range(0, WORLD_SIZE): peer_base = tl.load(buffer_ptrs + src_rank).to(tl.pointer_type(tl.uint8)) peer_ptr = (peer_base + input_byte_offset).to(tl.pointer_type(tl.float32)) - vals = tl.load(peer_ptr + global_token * HIDDEN_SIZE + offsets, mask=mask, other=0.0) + vals = tl.load( + peer_ptr + global_token * HIDDEN_SIZE + offsets, mask=mask, other=0.0 + ) acc += vals tl.store(local_ptr + token_offset * HIDDEN_SIZE + offsets, acc, mask=mask) diff --git a/megatron/core/inference/moe/batch_invariant.py b/megatron/core/inference/moe/batch_invariant.py index 8e7e2869a29..270bbf6d445 100644 --- a/megatron/core/inference/moe/batch_invariant.py +++ b/megatron/core/inference/moe/batch_invariant.py @@ -35,10 +35,7 @@ def enabled() -> bool: def grouped_mm(x_bf16: torch.Tensor, weight: torch.Tensor, offs: torch.Tensor) -> torch.Tensor: """Batch-invariant BF16 grouped GEMM used by inference fused MoE.""" return grouped_gemm_batch_invariant( - x_bf16, - weight, - offs=offs.to(torch.int32), - m_total=x_bf16.shape[0], + x_bf16, weight, offs=offs.to(torch.int32), m_total=x_bf16.shape[0] ) @@ -93,9 +90,7 @@ def unpermute_tokens_in_expert_order( _, hidden_dim = expert_output.shape num_tokens, num_local_experts = inverse_map.shape if out is None: - out = torch.empty( - num_tokens, hidden_dim, dtype=torch.float32, device=expert_output.device - ) + out = torch.empty(num_tokens, hidden_dim, dtype=torch.float32, device=expert_output.device) BLOCK_H = min(triton.next_power_of_2(hidden_dim), 1024) grid = (num_tokens, triton.cdiv(hidden_dim, BLOCK_H)) diff --git a/megatron/core/inference/moe/permute.py b/megatron/core/inference/moe/permute.py index 70377fa17bf..4046877ec41 100644 --- a/megatron/core/inference/moe/permute.py +++ b/megatron/core/inference/moe/permute.py @@ -370,9 +370,7 @@ def permute_tokens( # The inverse-map pointer is unused when HAS_INVERSE=False. Reuse an existing # int32 device tensor instead of allocating a dummy buffer for that kernel variant. inverse_map_ptr = ( - batch_invariant_inverse_map - if batch_invariant_inverse_map is not None - else permutation_map + batch_invariant_inverse_map if batch_invariant_inverse_map is not None else permutation_map ) _permute_tokens_kernel[(NUM_BLOCKS,)]( hidden_states, @@ -506,18 +504,14 @@ def unpermute_tokens( # MoE instead reduces each token independently in fixed local-expert order, # so unrelated tokens cannot affect the accumulation tree. if batch_invariant.enabled(): - assert batch_invariant_inverse_map is not None, ( - "batch-invariant MoE unpermute requires its inverse map" - ) + assert ( + batch_invariant_inverse_map is not None + ), "batch-invariant MoE unpermute requires its inverse map" # The expert-order kernel stores every row tok < valid_tokens, including zero # rows for tokens with no local expert contribution. Rows beyond # valid_tokens are not read by the graphed RSV combine. return batch_invariant.unpermute_tokens_in_expert_order( - expert_output, - permuted_probs, - batch_invariant_inverse_map, - valid_tokens, - out, + expert_output, permuted_probs, batch_invariant_inverse_map, valid_tokens, out ) BLOCK_H = min(triton.next_power_of_2(hidden_dim), 1024) diff --git a/megatron/core/ssm/mamba_mixer.py b/megatron/core/ssm/mamba_mixer.py index d2091952eb8..152e9c4573f 100644 --- a/megatron/core/ssm/mamba_mixer.py +++ b/megatron/core/ssm/mamba_mixer.py @@ -1004,7 +1004,7 @@ def _ssm_prefill( tensor_masked_update(ssm_state, batch_indices, cache_states) if self.config.batch_invariant_mode: - self._batch_invariant_decode().seed( + self._get_batch_invariant_decoder().seed( x, dt, B, C, cu_seqlens, batch_indices, max_batch=ssm_state.shape[0] ) @@ -1081,7 +1081,7 @@ def _get_decode_A_neg_exp(self) -> torch.Tensor: self._A_neg_exp_cache_stale = False return self._A_neg_exp_cache.view(-1, 1, 1).expand(-1, self.headdim, self.d_state) - def _batch_invariant_decode(self) -> MambaBatchInvariantDecode: + def _get_batch_invariant_decoder(self) -> MambaBatchInvariantDecode: """Batch-invariant decode adapter, created on first use.""" if not hasattr(self, "_batch_invariant_decoder"): self._batch_invariant_decoder = MambaBatchInvariantDecode(self) @@ -1235,10 +1235,9 @@ def _ssm_decode( y = y.unsqueeze(1) # Restore seq dimension elif self.config.batch_invariant_mode: assert batch_indices is not None, ( - "batch_invariant_mode for Mamba decode requires dynamic batching " - "batch_indices." + "batch_invariant_mode for Mamba decode requires dynamic batching " "batch_indices." ) - y = self._batch_invariant_decode().step(x, dt, B, C, batch_indices, ssm_state) + y = self._get_batch_invariant_decoder().step(x, dt, B, C, batch_indices, ssm_state) else: A = self._get_decode_A_neg_exp() diff --git a/megatron/core/ssm/ops/batch_invariant_decode.py b/megatron/core/ssm/ops/batch_invariant_decode.py index 3244cf4874b..b2e41cc87dd 100644 --- a/megatron/core/ssm/ops/batch_invariant_decode.py +++ b/megatron/core/ssm/ops/batch_invariant_decode.py @@ -12,15 +12,15 @@ class BatchInvariantDecodeBuffers: """Per-slot persistent state for the buffered decode scan.""" - x: torch.Tensor # (max_batch + 1, chunk_size, nheads, headdim) - dt: torch.Tensor # (max_batch + 1, chunk_size, nheads) - B: torch.Tensor # (max_batch + 1, chunk_size, ngroups, dstate) - C: torch.Tensor # (max_batch + 1, chunk_size, ngroups, dstate) + x: torch.Tensor # (max_batch + 1, chunk_size, nheads, headdim) + dt: torch.Tensor # (max_batch + 1, chunk_size, nheads) + B: torch.Tensor # (max_batch + 1, chunk_size, ngroups, dstate) + C: torch.Tensor # (max_batch + 1, chunk_size, ngroups, dstate) # Tokens buffered since the slot's last chunk boundary; doubles as the # write cursor for the next token. num_buffered: torch.Tensor # (max_batch + 1,) int32 # Per-entry target-row output, allocated once and sliced per step. - out: torch.Tensor # (max_batch + 1, nheads, headdim) + out: torch.Tensor # (max_batch + 1, nheads, headdim) @classmethod def allocate( @@ -87,13 +87,9 @@ def seed( # from reaching the target row through tensor-core operations. offsets = torch.arange(chunk_size, device=x.device, dtype=torch.long) safe_tail_lens = torch.clamp(tail_lens, min=1) - safe_tail_offsets = torch.minimum( - offsets.unsqueeze(0), (safe_tail_lens - 1).unsqueeze(1) - ) + safe_tail_offsets = torch.minimum(offsets.unsqueeze(0), (safe_tail_lens - 1).unsqueeze(1)) safe_tail_starts = torch.where( - tail_lens > 0, - seq_ends - tail_lens, - torch.clamp(seq_ends - 1, min=0), + tail_lens > 0, seq_ends - tail_lens, torch.clamp(seq_ends - 1, min=0) ) tail_token_idx = (safe_tail_starts.unsqueeze(1) + safe_tail_offsets).clamp( max=x.shape[0] - 1 @@ -113,10 +109,10 @@ def seed( def batch_invariant_decode_buffered_scan( buffers: BatchInvariantDecodeBuffers, - x: torch.Tensor, # (decode_batch_size, 1, nheads, headdim) - dt: torch.Tensor, # (decode_batch_size, 1, nheads) - B: torch.Tensor, # (decode_batch_size, 1, ngroups, dstate) - C: torch.Tensor, # (decode_batch_size, 1, ngroups, dstate) + x: torch.Tensor, # (decode_batch_size, 1, nheads, headdim) + dt: torch.Tensor, # (decode_batch_size, 1, nheads) + B: torch.Tensor, # (decode_batch_size, 1, ngroups, dstate) + C: torch.Tensor, # (decode_batch_size, 1, ngroups, dstate) A: torch.Tensor, D: torch.Tensor, dt_bias: torch.Tensor, diff --git a/megatron/core/ssm/ops/ssd_bmm.py b/megatron/core/ssm/ops/ssd_bmm.py index ba07563a1f7..65d89547f7a 100644 --- a/megatron/core/ssm/ops/ssd_bmm.py +++ b/megatron/core/ssm/ops/ssd_bmm.py @@ -178,9 +178,9 @@ def _bmm_chunk_fwd( out: (nchunks, ngroups, chunk_size, chunk_size) """ has_target_rows = target_rows is not None - assert (chunk_starts is not None) == has_target_rows, ( - "target_rows and chunk_starts must be provided together" - ) + assert ( + chunk_starts is not None + ) == has_target_rows, "target_rows and chunk_starts must be provided together" if has_target_rows: # chunk_starts has one fixed-window start per chunk. chunk_offsets = chunk_starts diff --git a/megatron/core/ssm/ops/ssd_chunk_scan.py b/megatron/core/ssm/ops/ssd_chunk_scan.py index 88dcf10bdec..aa92a4466d4 100644 --- a/megatron/core/ssm/ops/ssd_chunk_scan.py +++ b/megatron/core/ssm/ops/ssd_chunk_scan.py @@ -334,14 +334,8 @@ def _chunk_scan_fwd_kernel( out_ptr += pid_c * stride_out_seqlen + pid_h * stride_out_head # All M-lanes alias the same output row (row stride 0); the mask # keeps only lane tr, so one lane stores per column. - out_ptrs = out_ptr + ( - offs_out_m[:, None] * 0 + offs_out_n[None, :] * stride_out_hdim - ) - tl.store( - out_ptrs, - acc, - mask=(offs_out_m[:, None] == tr) & (offs_out_n[None, :] < hdim), - ) + out_ptrs = out_ptr + (offs_out_m[:, None] * 0 + offs_out_n[None, :] * stride_out_hdim) + tl.store(out_ptrs, acc, mask=(offs_out_m[:, None] == tr) & (offs_out_n[None, :] < hdim)) else: out_ptr += chunk_seqlen_start * stride_out_seqlen + pid_h * stride_out_head out_ptrs = out_ptr + ( @@ -372,9 +366,9 @@ def _chunk_scan_fwd( ): assert seq_idx is not None, "this implementation requires seq_idx" has_target_rows = target_rows is not None - assert (chunk_starts is not None) == has_target_rows, ( - "target_rows and chunk_starts must be provided together" - ) + assert ( + chunk_starts is not None + ) == has_target_rows, "target_rows and chunk_starts must be provided together" if has_target_rows: chunk_offsets = chunk_starts else: diff --git a/megatron/core/ssm/ops/ssd_chunk_state.py b/megatron/core/ssm/ops/ssd_chunk_state.py index 4a698630268..70d2c599d5d 100644 --- a/megatron/core/ssm/ops/ssd_chunk_state.py +++ b/megatron/core/ssm/ops/ssd_chunk_state.py @@ -350,9 +350,9 @@ def _chunk_state_fwd( chunk_starts=None, ): has_chunk_flags = chunk_flags is not None - assert (chunk_starts is not None) == has_chunk_flags, ( - "chunk_flags and chunk_starts must be provided together" - ) + assert ( + chunk_starts is not None + ) == has_chunk_flags, "chunk_flags and chunk_starts must be provided together" if has_chunk_flags: chunk_offsets = chunk_starts else: diff --git a/megatron/core/ssm/ops/ssd_state_passing.py b/megatron/core/ssm/ops/ssd_state_passing.py index bada474136c..9f3a4e0a551 100644 --- a/megatron/core/ssm/ops/ssd_state_passing.py +++ b/megatron/core/ssm/ops/ssd_state_passing.py @@ -89,9 +89,7 @@ def _state_passing_fwd_kernel( else: dst_flag = True # Unflagged destination chunks have no chunk state. - new_states = tl.load(states_ptrs, mask=(offs_m < dim) & dst_flag, other=0.0).to( - tl.float32 - ) + new_states = tl.load(states_ptrs, mask=(offs_m < dim) & dst_flag, other=0.0).to(tl.float32) dA_cs = tl.load(dA_cs_ptr).to(tl.float32) seq_idx = tl.load(seq_idx_ptr + c * stride_seq_idx_chunk) if HAS_DST_STATES: @@ -154,9 +152,9 @@ def _state_passing_fwd( assert dA_cumsum.shape == (nheads, nchunks, chunk_size) seqlen = seq_idx.shape[-1] has_dst = dst_states is not None - assert (dst_indices is not None) == has_dst and (dst_flags is not None) == has_dst, ( - "dst_states, dst_indices, and dst_flags must be provided together" - ) + assert (dst_indices is not None) == has_dst and ( + dst_flags is not None + ) == has_dst, "dst_states, dst_indices, and dst_flags must be provided together" if not has_dst: out_dtype = states.dtype if out_dtype is None else out_dtype out = torch.empty((nchunks, nheads, dim), device=states.device, dtype=out_dtype) @@ -173,9 +171,7 @@ def _state_passing_fwd( if has_dst: assert dst_states.shape[1] == nheads and dst_states.shape[2] == dim dst_strides = ( - (dst_states.stride(0), dst_states.stride(1), dst_states.stride(2)) - if has_dst - else (0, 0, 0) + (dst_states.stride(0), dst_states.stride(1), dst_states.stride(2)) if has_dst else (0, 0, 0) ) grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE"]), nheads) diff --git a/megatron/core/transformer/custom_layers/batch_invariant_kernels.py b/megatron/core/transformer/custom_layers/batch_invariant_kernels.py index a14489c5efb..ead2b53d66c 100644 --- a/megatron/core/transformer/custom_layers/batch_invariant_kernels.py +++ b/megatron/core/transformer/custom_layers/batch_invariant_kernels.py @@ -1571,11 +1571,7 @@ def _bf16_grouped_gemm_wgrad_contiguous( def grouped_gemm_batch_invariant( - a: torch.Tensor, - b: torch.Tensor, - *, - offs: torch.Tensor, - m_total: int, + a: torch.Tensor, b: torch.Tensor, *, offs: torch.Tensor, m_total: int ) -> torch.Tensor: """Run the graph-safe grouped GEMM over pre-aligned inference expert blocks.""" m_indices = _offs_to_m_indices(offs, m_total).contiguous() @@ -1685,12 +1681,10 @@ def _pin_mamba_autotuners(): kernels = [] try: - from megatron.core.ssm.ops import ( - ssd_bmm as r_bmm, - ssd_chunk_scan as r_scan, - ssd_chunk_state as r_state, - ssd_state_passing as r_pass, - ) + from megatron.core.ssm.ops import ssd_bmm as r_bmm + from megatron.core.ssm.ops import ssd_chunk_scan as r_scan + from megatron.core.ssm.ops import ssd_chunk_state as r_state + from megatron.core.ssm.ops import ssd_state_passing as r_pass kernels += [ r_bmm._bmm_chunk_fwd_kernel, @@ -1702,12 +1696,10 @@ def _pin_mamba_autotuners(): except ImportError: pass try: - from mamba_ssm.ops.triton import ( - ssd_bmm as p_bmm, - ssd_chunk_scan as p_scan, - ssd_chunk_state as p_state, - ssd_state_passing as p_pass, - ) + from mamba_ssm.ops.triton import ssd_bmm as p_bmm + from mamba_ssm.ops.triton import ssd_chunk_scan as p_scan + from mamba_ssm.ops.triton import ssd_chunk_state as p_state + from mamba_ssm.ops.triton import ssd_state_passing as p_pass kernels += [ p_bmm._bmm_chunk_fwd_kernel, diff --git a/megatron/core/transformer/moe/batch_invariant.py b/megatron/core/transformer/moe/batch_invariant.py index 79a0611bf65..9baf66dab65 100644 --- a/megatron/core/transformer/moe/batch_invariant.py +++ b/megatron/core/transformer/moe/batch_invariant.py @@ -24,9 +24,9 @@ def build_inverse_permutation_map( assert isinstance( num_out_tokens, int ), "batch-invariant graph unpermute requires static num_out_tokens" - assert num_out_tokens % num_tokens == 0, ( - "batch-invariant graph unpermute expects fixed top-k per token" - ) + assert ( + num_out_tokens % num_tokens == 0 + ), "batch-invariant graph unpermute expects fixed top-k per token" topk = num_out_tokens // num_tokens row_ids = torch.arange(num_out_tokens, device=routing_map.device, dtype=torch.long) @@ -38,7 +38,9 @@ def build_inverse_permutation_map( linear_slots = token_ids * topk + row_slots inverse_rows = torch.full((num_tokens, topk), -1, device=routing_map.device, dtype=torch.long) - inverse_experts = torch.full((num_tokens, topk), -1, device=routing_map.device, dtype=torch.long) + inverse_experts = torch.full( + (num_tokens, topk), -1, device=routing_map.device, dtype=torch.long + ) inverse_rows.view(-1).scatter_(0, linear_slots, row_ids) inverse_experts.view(-1).scatter_(0, linear_slots, expert_ids) return torch.stack((inverse_rows, inverse_experts), dim=0) @@ -75,9 +77,7 @@ def unpermute( for k in range(topk): row_ids = inverse_rows[:, k] expert_ids = inverse_experts[:, k] - valid_mask = ( - (row_ids >= 0) & (expert_ids >= start_expert) & (expert_ids < end_expert) - ) + valid_mask = (row_ids >= 0) & (expert_ids >= start_expert) & (expert_ids < end_expert) safe_rows = row_ids.clamp_min(0) chunk = permuted_tokens.index_select(0, safe_rows).to(torch.float32) diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py index aae7a731c84..923c730fa53 100644 --- a/megatron/core/transformer/moe/moe_utils.py +++ b/megatron/core/transformer/moe/moe_utils.py @@ -25,8 +25,8 @@ from megatron.core.transformer.enums import CudaGraphModule from megatron.core.transformer.moe.batch_invariant import ( build_inverse_permutation_map as build_batch_invariant_inverse_permutation_map, - unpermute as batch_invariant_unpermute, ) +from megatron.core.transformer.moe.batch_invariant import unpermute as batch_invariant_unpermute from megatron.core.transformer.moe.router_replay import RouterReplay from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.utils import internal_api, is_te_min_version @@ -444,10 +444,7 @@ def permute( if return_batch_invariant_inverse_map: batch_invariant_inverse_map = build_batch_invariant_inverse_permutation_map( - routing_map_for_inverse, - flat_sorted, - sorted_indices, - num_out_tokens, + routing_map_for_inverse, flat_sorted, sorted_indices, num_out_tokens ) # use the mapping to permute the tokens @@ -525,9 +522,9 @@ def unpermute( if batch_invariant_mode: assert routing_map is not None, "batch-invariant MoE unpermute requires routing_map" - assert batch_invariant_inverse_map is not None, ( - "batch-invariant MoE unpermute requires the AllToAll inverse map" - ) + assert ( + batch_invariant_inverse_map is not None + ), "batch-invariant MoE unpermute requires the AllToAll inverse map" return batch_invariant_unpermute( permuted_tokens, restore_shape, diff --git a/megatron/core/transformer/moe/token_dispatcher_inference.py b/megatron/core/transformer/moe/token_dispatcher_inference.py index 84fb4cda68f..0d1b133d5ca 100644 --- a/megatron/core/transformer/moe/token_dispatcher_inference.py +++ b/megatron/core/transformer/moe/token_dispatcher_inference.py @@ -41,12 +41,12 @@ gather_from_sequence_parallel_region, reduce_scatter_to_sequence_parallel_region, ) -from megatron.core.transformer.moe.shared_experts import SharedExpertMLP -from megatron.core.transformer.moe.token_dispatcher import MoEAllGatherTokenDispatcher -from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( is_batch_invariant_mode_enabled, ) +from megatron.core.transformer.moe.shared_experts import SharedExpertMLP +from megatron.core.transformer.moe.token_dispatcher import MoEAllGatherTokenDispatcher +from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.typed_torch import apply_module from megatron.core.utils import get_pg_rank, get_pg_size diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 61170f06911..7c2bffd456c 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1340,8 +1340,7 @@ def __post_init__(self): if self.batch_invariant_mode: if self.inference_grouped_gemm_backend != InferenceGroupedGemmBackend.TORCH: raise ValueError( - "batch_invariant_mode requires " - "inference_grouped_gemm_backend='torch'." + "batch_invariant_mode requires " "inference_grouped_gemm_backend='torch'." ) if ( self.expert_model_parallel_size > 1 @@ -2503,9 +2502,7 @@ def _scope_to_str(s): assert not ( self.fp8 or self.fp4 ), "Batch-invariant MoE is bf16-only. Disable fp8/fp4 to use it." - assert not ( - self.moe_permute_fusion or self.moe_permute_fusion_into_hybridep - ), ( + assert not (self.moe_permute_fusion or self.moe_permute_fusion_into_hybridep), ( "Batch-invariant MoE requires the unfused permute/unpermute path so " "top-k reductions use the fixed batch-invariant add tree." ) diff --git a/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py b/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py index cf2ed7c5293..2512cdc98fe 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py @@ -800,9 +800,7 @@ def test_batch_invariant_mamba_chunked_prefill_scheduler_alignment(self): one_left_req = self._req(ctx, self._prompt(ctx.mamba_chunk_size + 1), request_id=3) assert ( - engine._mamba_batch_invariant_prefill_chunk_length( - one_left_req, ctx.mamba_chunk_size - ) + engine._mamba_batch_invariant_prefill_chunk_length(one_left_req, ctx.mamba_chunk_size) == 0 ) assert ( 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 0e941ea8568..a1c858031be 100644 --- a/tests/unit_tests/inference/test_moe_dispatching_and_routing.py +++ b/tests/unit_tests/inference/test_moe_dispatching_and_routing.py @@ -233,9 +233,7 @@ def test_init(self): def test_init_rejects_batch_invariant_ep(self): """Batch-invariant MoE on the inference EP path is NVLS-only on this branch.""" if Utils.world_size == 1: - pytest.skip( - "NCCL batch-invariant rejection is only relevant with expert parallelism." - ) + pytest.skip("NCCL batch-invariant rejection is only relevant with expert parallelism.") with pytest.raises(ValueError, match="requires inference_moe_token_dispatcher_type"): self._make_dispatcher( @@ -554,10 +552,10 @@ def test_cuda_graph_batch_invariant_moe_layer_uses_ordered_rsv(self, monkeypatch HAVE_DEEPGEMM_BF16, set_batch_invariant_mode, ) + from megatron.core.transformer.moe import token_dispatcher_inference from megatron.core.transformer.moe.token_dispatcher_inference import ( NVLSAllGatherVDispatcher, ) - from megatron.core.transformer.moe import token_dispatcher_inference if Utils.world_size < 2: pytest.skip("NVLS RSV branch test requires expert-parallel world_size > 1.") diff --git a/tests/unit_tests/rl/test_rl_batch_invariant.py b/tests/unit_tests/rl/test_rl_batch_invariant.py index 2e834d5bdaf..7e62022d6b8 100644 --- a/tests/unit_tests/rl/test_rl_batch_invariant.py +++ b/tests/unit_tests/rl/test_rl_batch_invariant.py @@ -43,9 +43,7 @@ def test_moe_unpermute_batch_invariant_inverse_map_rank_tree(): ).expand(4, hidden) sorted_indices = torch.zeros(4, device="cuda", dtype=torch.int64) routing_map = torch.ones(1, 4, device="cuda", dtype=torch.bool) - inverse_map = torch.tensor( - [[[0, 1, 2, 3]], [[0, 1, 2, 3]]], device="cuda", dtype=torch.int64 - ) + inverse_map = torch.tensor([[[0, 1, 2, 3]], [[0, 1, 2, 3]]], device="cuda", dtype=torch.int64) parallel_state.set_expert_model_parallel_world_size(2) try: @@ -110,7 +108,7 @@ def _run(): parallel_state.set_expert_model_parallel_world_size(None) torch.testing.assert_close(graph_out, expected, rtol=0.0, atol=0.0) - reference = ( - tokens.float() * probs[:, 0, None] + tokens.float() * probs[:, 2, None] - ).to(tokens.dtype) + reference = (tokens.float() * probs[:, 0, None] + tokens.float() * probs[:, 2, None]).to( + tokens.dtype + ) torch.testing.assert_close(graph_out, reference, rtol=0.0, atol=0.0) diff --git a/tests/unit_tests/ssm/ops/test_batch_invariant_decode.py b/tests/unit_tests/ssm/ops/test_batch_invariant_decode.py index 9e9191ff541..aa944de0ff1 100644 --- a/tests/unit_tests/ssm/ops/test_batch_invariant_decode.py +++ b/tests/unit_tests/ssm/ops/test_batch_invariant_decode.py @@ -22,9 +22,18 @@ def _full_scan(x, dt, A, B, C, D, dt_bias, chunk_size, initial_states=None): """Reference: a single `mamba_chunk_scan_combined` over the whole sequence.""" y, final = mamba_chunk_scan_combined( - x, dt, A, B, C, chunk_size, - D=D, z=None, dt_bias=dt_bias, dt_softplus=True, - initial_states=initial_states, return_final_states=True, + x, + dt, + A, + B, + C, + chunk_size, + D=D, + z=None, + dt_bias=dt_bias, + dt_softplus=True, + initial_states=initial_states, + return_final_states=True, ) return y, final @@ -58,13 +67,9 @@ def setUp(self): self.ngroups = 1 self.dstate = 16 self.chunk_size = 32 - self.A = -torch.exp( - torch.randn(self.nh, device=self.device, dtype=torch.float32).abs() - ) + self.A = -torch.exp(torch.randn(self.nh, device=self.device, dtype=torch.float32).abs()) self.D = torch.randn(self.nh, device=self.device, dtype=torch.float32) * 0.1 - self.dt_bias = ( - torch.randn(self.nh, device=self.device, dtype=torch.float32) * 0.01 - ) + self.dt_bias = torch.randn(self.nh, device=self.device, dtype=torch.float32) * 0.01 def _make_seq(self, total_len): """Generate a (1, total_len, ...) random mamba input sequence.""" @@ -78,20 +83,20 @@ def _make_seq(self, total_len): def _make_bufs(self, max_batch): return BatchInvariantDecodeBuffers.allocate( - max_batch, self.chunk_size, - self.nh, self.headdim, self.ngroups, self.dstate, - self.device, self.dtype, + max_batch, + self.chunk_size, + self.nh, + self.headdim, + self.ngroups, + self.dstate, + self.device, + self.dtype, ) def _make_ssm_state(self, max_batch): """Production BIK state cache: FP32 carry across Mamba chunks.""" return torch.zeros( - max_batch, - self.nh, - self.headdim, - self.dstate, - device=self.device, - dtype=torch.float32, + max_batch, self.nh, self.headdim, self.dstate, device=self.device, dtype=torch.float32 ) def _seed_from_prefill(self, bufs, x, dt, B, C, prefill_len, slot, max_batch): @@ -105,8 +110,14 @@ def _seed_from_prefill(self, bufs, x, dt, B, C, prefill_len, slot, max_batch): # Prefill on the largest chunk-aligned prefix; the tail goes in the buffer. aligned = (prefill_len // self.chunk_size) * self.chunk_size _, final = _full_scan( - x[:, :aligned], dt[:, :aligned], self.A, - B[:, :aligned], C[:, :aligned], self.D, self.dt_bias, self.chunk_size, + x[:, :aligned], + dt[:, :aligned], + self.A, + B[:, :aligned], + C[:, :aligned], + self.D, + self.dt_bias, + self.chunk_size, initial_states=None, ) ssm_state[slot] = final[0] @@ -120,7 +131,8 @@ def _seed_from_prefill(self, bufs, x, dt, B, C, prefill_len, slot, max_batch): dt[0, :prefill_len], B[0, :prefill_len], C[0, :prefill_len], - cu, batch_indices, + cu, + batch_indices, ) return ssm_state @@ -133,14 +145,14 @@ def _decode_one_step(self, bufs, x, dt, B, C, pos, slot, ssm_state): dt[:, pos : pos + 1], B[:, pos : pos + 1], C[:, pos : pos + 1], - self.A, self.D, self.dt_bias, + self.A, + self.D, + self.dt_bias, batch_indices, ssm_state, ) - def _varlen_boundary_state_from_prefill( - self, x, dt, B, C, prefill_len, initial_states=None - ): + def _varlen_boundary_state_from_prefill(self, x, dt, B, C, prefill_len, initial_states=None): """Production-shaped varlen prefill returning the last full chunk boundary state.""" chunk_boundaries = [0] pos = self.chunk_size @@ -149,9 +161,7 @@ def _varlen_boundary_state_from_prefill( pos += self.chunk_size chunk_boundaries.append(prefill_len) - cu_chunk_seqlens = torch.tensor( - chunk_boundaries, dtype=torch.int32, device=self.device - ) + cu_chunk_seqlens = torch.tensor(chunk_boundaries, dtype=torch.int32, device=self.device) last_chunk_indices = torch.tensor( [len(chunk_boundaries) - 2], dtype=torch.int32, device=self.device ) @@ -163,9 +173,7 @@ def _varlen_boundary_state_from_prefill( boundary_idx = boundary_idx.clamp(min=0) out = torch.zeros_like(x[0, :prefill_len]) - seq_idx = torch.zeros( - len(chunk_boundaries) - 1, dtype=torch.int32, device=self.device - ) + seq_idx = torch.zeros(len(chunk_boundaries) - 1, dtype=torch.int32, device=self.device) chunk_states = mamba_chunk_scan_combined_varlen( x=x[0, :prefill_len], dt=dt[0, :prefill_len], @@ -190,9 +198,7 @@ def _varlen_boundary_state_from_prefill( boundary_state = chunk_states[boundary_idx] if not has_boundary: boundary_state = ( - torch.zeros_like(boundary_state) - if initial_states is None - else initial_states + torch.zeros_like(boundary_state) if initial_states is None else initial_states ) return final_state, boundary_state @@ -209,20 +215,15 @@ def test_single_decode_matches_full_scan(self): total = prefill_len + 1 x, dt, B, C = self._make_seq(total) # Reference: full scan over the whole (prefill + 1) sequence. - y_full, _ = _full_scan( - x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size, - ) + y_full, _ = _full_scan(x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size) # batch-invariant: seed from prefill, then one decode step. bufs = self._make_bufs(max_batch) - ssm_state = self._seed_from_prefill( - bufs, x, dt, B, C, prefill_len, slot, max_batch, - ) + ssm_state = self._seed_from_prefill(bufs, x, dt, B, C, prefill_len, slot, max_batch) y_batch_invariant = self._decode_one_step( - bufs, x, dt, B, C, prefill_len, slot, ssm_state, + bufs, x, dt, B, C, prefill_len, slot, ssm_state ) self._assert_bitwise( - y_batch_invariant[0, 0], y_full[0, prefill_len], - f"prefill_len={prefill_len}", + y_batch_invariant[0, 0], y_full[0, prefill_len], f"prefill_len={prefill_len}" ) def test_rejects_bf16_state_cache(self): @@ -230,12 +231,7 @@ def test_rejects_bf16_state_cache(self): x, dt, B, C = self._make_seq(1) bufs = self._make_bufs(max_batch=2) ssm_state = torch.zeros( - 2, - self.nh, - self.headdim, - self.dstate, - device=self.device, - dtype=torch.bfloat16, + 2, self.nh, self.headdim, self.dstate, device=self.device, dtype=torch.bfloat16 ) with self.assertRaisesRegex(AssertionError, "requires an FP32 SSM state cache"): self._decode_one_step(bufs, x, dt, B, C, pos=0, slot=0, ssm_state=ssm_state) @@ -249,27 +245,18 @@ def test_inference_config_uses_fp32_state_cache(self): from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols model = SimpleNamespace( - config=SimpleNamespace( - batch_invariant_mode=True, - params_dtype=torch.bfloat16, - ) + config=SimpleNamespace(batch_invariant_mode=True, params_dtype=torch.bfloat16) ) decoder = SimpleNamespace( layer_type_list=[Symbols.MAMBA], layers=[SimpleNamespace(mixer=SimpleNamespace(chunk_size=self.chunk_size))], mamba_state_shapes_per_request=lambda: ((4, 8), (8, 32, 16)), ) - with patch( - "megatron.core.inference.config.get_attr_wrapped_model", - return_value=decoder, - ): + with patch("megatron.core.inference.config.get_attr_wrapped_model", return_value=decoder): config = MambaInferenceStateConfig.from_model(model) self.assertEqual(config.ssm_states_dtype, torch.float32) with self.assertRaisesRegex(ValueError, "requires FP32 Mamba SSM states"): - MambaInferenceStateConfig.from_model( - model, - ssm_states_dtype=torch.bfloat16, - ) + MambaInferenceStateConfig.from_model(model, ssm_states_dtype=torch.bfloat16) model.config.batch_invariant_mode = False config = MambaInferenceStateConfig.from_model(model) self.assertEqual(config.ssm_states_dtype, torch.bfloat16) @@ -283,9 +270,7 @@ def test_dynamic_prefill_uses_boundary_state_not_prompt_end_state(self): with self.subTest(prefill_len=prefill_len): total = prefill_len + 1 x, dt, B, C = self._make_seq(total) - y_full, _ = _full_scan( - x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size, - ) + y_full, _ = _full_scan(x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size) _, boundary_state = self._varlen_boundary_state_from_prefill( x, dt, B, C, prefill_len @@ -308,7 +293,8 @@ def test_dynamic_prefill_uses_boundary_state_not_prompt_end_state(self): bufs, x, dt, B, C, prefill_len, slot, ssm_state ) self._assert_bitwise( - y_batch_invariant[0, 0], y_full[0, prefill_len], + y_batch_invariant[0, 0], + y_full[0, prefill_len], f"dynamic prefill boundary state prefill_len={prefill_len}", ) @@ -321,9 +307,7 @@ def test_chunked_prefill_handoff_matches_full_scan(self): with self.subTest(final_chunk_len=final_chunk_len): prefill_len = first_chunk_len + final_chunk_len x, dt, B, C = self._make_seq(prefill_len + 1) - y_full, _ = _full_scan( - x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size, - ) + y_full, _ = _full_scan(x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size) _, first_boundary = self._varlen_boundary_state_from_prefill( x[:, :first_chunk_len], @@ -344,9 +328,7 @@ def test_chunked_prefill_handoff_matches_full_scan(self): ssm_state = self._make_ssm_state(max_batch) ssm_state[slot] = final_boundary[0] bufs = self._make_bufs(max_batch) - cu = torch.tensor( - [0, final_chunk_len], dtype=torch.int32, device=self.device - ) + cu = torch.tensor([0, final_chunk_len], dtype=torch.int32, device=self.device) bufs.seed( x[0, first_chunk_len:prefill_len], dt[0, first_chunk_len:prefill_len], @@ -373,9 +355,7 @@ def test_seed_ignores_nonfinite_physical_padding_rows(self): prefill_len = self.chunk_size + 1 total = prefill_len + 1 x, dt, B, C = self._make_seq(total) - y_full, _ = _full_scan( - x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size, - ) + y_full, _ = _full_scan(x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size) _, boundary_state = self._varlen_boundary_state_from_prefill(x, dt, B, C, prefill_len) ssm_state = self._make_ssm_state(max_batch) @@ -402,12 +382,9 @@ def test_seed_ignores_nonfinite_physical_padding_rows(self): self.assertTrue(torch.isfinite(bufs.B[slot]).all()) self.assertTrue(torch.isfinite(bufs.C[slot]).all()) - y_batch_invariant = self._decode_one_step( - bufs, x, dt, B, C, prefill_len, slot, ssm_state - ) + y_batch_invariant = self._decode_one_step(bufs, x, dt, B, C, prefill_len, slot, ssm_state) self._assert_bitwise( - y_batch_invariant[0, 0], y_full[0, prefill_len], - "nonfinite physical padding rows", + y_batch_invariant[0, 0], y_full[0, prefill_len], "nonfinite physical padding rows" ) def test_short_prefill_uses_zero_boundary_state(self): @@ -417,19 +394,14 @@ def test_short_prefill_uses_zero_boundary_state(self): with self.subTest(prefill_len=prefill_len): total = prefill_len + 1 x, dt, B, C = self._make_seq(total) - y_full, _ = _full_scan( - x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size, - ) + y_full, _ = _full_scan(x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size) bufs = self._make_bufs(max_batch) - ssm_state = self._seed_from_prefill( - bufs, x, dt, B, C, prefill_len, slot, max_batch, - ) + ssm_state = self._seed_from_prefill(bufs, x, dt, B, C, prefill_len, slot, max_batch) y_batch_invariant = self._decode_one_step( - bufs, x, dt, B, C, prefill_len, slot, ssm_state, + bufs, x, dt, B, C, prefill_len, slot, ssm_state ) self._assert_bitwise( - y_batch_invariant[0, 0], y_full[0, prefill_len], - f"prefill_len={prefill_len}", + y_batch_invariant[0, 0], y_full[0, prefill_len], f"prefill_len={prefill_len}" ) def test_multi_step_decode_across_chunk_boundary(self): @@ -440,20 +412,17 @@ def test_multi_step_decode_across_chunk_boundary(self): n_decode = self.chunk_size + 5 # enough to cross at least one boundary total = prefill_len + n_decode x, dt, B, C = self._make_seq(total) - y_full, _ = _full_scan( - x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size, - ) + y_full, _ = _full_scan(x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size) bufs = self._make_bufs(max_batch) - ssm_state = self._seed_from_prefill( - bufs, x, dt, B, C, prefill_len, slot, max_batch, - ) + ssm_state = self._seed_from_prefill(bufs, x, dt, B, C, prefill_len, slot, max_batch) for k in range(n_decode): pos = prefill_len + k y_batch_invariant = self._decode_one_step(bufs, x, dt, B, C, pos, slot, ssm_state) self._assert_bitwise( - y_batch_invariant[0, 0], y_full[0, pos], + y_batch_invariant[0, 0], + y_full[0, pos], f"step k={k} (pos={pos}, num_buffered_before={bufs.num_buffered[slot].item()})", ) @@ -467,50 +436,51 @@ def test_multi_slot_independent_streams(self): y_refs = [] for plen in prefill_lens: x, dt, B, C = self._make_seq(plen + 1) - x_per_slot.append(x); dt_per_slot.append(dt) - B_per_slot.append(B); C_per_slot.append(C) - y_full, _ = _full_scan( - x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size, - ) + x_per_slot.append(x) + dt_per_slot.append(dt) + B_per_slot.append(B) + C_per_slot.append(C) + y_full, _ = _full_scan(x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size) y_refs.append(y_full[0, plen]) bufs = self._make_bufs(max_batch) # Per-slot seeding (each slot's prefill done independently). ssm_state = self._make_ssm_state(max_batch) for slot, plen, x, dt, B, C in zip( - slots, prefill_lens, x_per_slot, dt_per_slot, B_per_slot, C_per_slot, + slots, prefill_lens, x_per_slot, dt_per_slot, B_per_slot, C_per_slot ): - partial = self._seed_from_prefill( - bufs, x, dt, B, C, plen, slot, max_batch, - ) + partial = self._seed_from_prefill(bufs, x, dt, B, C, plen, slot, max_batch) ssm_state[slot] = partial[slot] # Both slots step at once. x_step = torch.cat( - [x_per_slot[i][:, prefill_lens[i] : prefill_lens[i] + 1] for i in range(2)], - dim=0, + [x_per_slot[i][:, prefill_lens[i] : prefill_lens[i] + 1] for i in range(2)], dim=0 ) dt_step = torch.cat( - [dt_per_slot[i][:, prefill_lens[i] : prefill_lens[i] + 1] for i in range(2)], - dim=0, + [dt_per_slot[i][:, prefill_lens[i] : prefill_lens[i] + 1] for i in range(2)], dim=0 ) B_step = torch.cat( - [B_per_slot[i][:, prefill_lens[i] : prefill_lens[i] + 1] for i in range(2)], - dim=0, + [B_per_slot[i][:, prefill_lens[i] : prefill_lens[i] + 1] for i in range(2)], dim=0 ) C_step = torch.cat( - [C_per_slot[i][:, prefill_lens[i] : prefill_lens[i] + 1] for i in range(2)], - dim=0, + [C_per_slot[i][:, prefill_lens[i] : prefill_lens[i] + 1] for i in range(2)], dim=0 ) batch_indices = torch.tensor(slots, dtype=torch.int32, device=self.device) y_batch_invariant = batch_invariant_decode_buffered_scan( - bufs, x_step, dt_step, B_step, C_step, - self.A, self.D, self.dt_bias, batch_indices, ssm_state, + bufs, + x_step, + dt_step, + B_step, + C_step, + self.A, + self.D, + self.dt_bias, + batch_indices, + ssm_state, ) for i, plen in enumerate(prefill_lens): self._assert_bitwise( - y_batch_invariant[i, 0], y_refs[i], - f"multi-slot slot={slots[i]} prefill_len={plen}", + y_batch_invariant[i, 0], y_refs[i], f"multi-slot slot={slots[i]} prefill_len={plen}" ) def test_inactive_padding_entries(self): @@ -522,22 +492,17 @@ def test_inactive_padding_entries(self): n_decode = 8 total = prefill_len + n_decode x, dt, B, C = self._make_seq(total) - y_full, _ = _full_scan( - x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size, - ) + y_full, _ = _full_scan(x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size) bufs = self._make_bufs(max_batch) - ssm_state = self._seed_from_prefill( - bufs, x, dt, B, C, prefill_len, slot, max_batch, - ) + ssm_state = self._seed_from_prefill(bufs, x, dt, B, C, prefill_len, slot, max_batch) batch_indices = torch.tensor([slot, -1, -1], dtype=torch.int32, device=self.device) for k in range(n_decode): pos = prefill_len + k + # Entry 0 carries the real token; padding entries carry garbage. def pad3(t): - junk = torch.randn( - 2, *t.shape[1:], device=t.device, dtype=t.dtype - ) + junk = torch.randn(2, *t.shape[1:], device=t.device, dtype=t.dtype) return torch.cat([t, junk], dim=0) y_batch_invariant = batch_invariant_decode_buffered_scan( @@ -546,13 +511,13 @@ def pad3(t): pad3(dt[:, pos : pos + 1]), pad3(B[:, pos : pos + 1]), pad3(C[:, pos : pos + 1]), - self.A, self.D, self.dt_bias, + self.A, + self.D, + self.dt_bias, batch_indices, ssm_state, ) - self._assert_bitwise( - y_batch_invariant[0, 0], y_full[0, pos], f"padded step k={k}" - ) + self._assert_bitwise(y_batch_invariant[0, 0], y_full[0, pos], f"padded step k={k}") # Padding entries must return zeros. self.assertEqual(y_batch_invariant[1:].abs().max().item(), 0.0) @@ -561,23 +526,17 @@ def test_cuda_graph_replay_matches_full_scan(self): max_batch, slot = 2, 0 prefill_len = 20 x, dt, B, C = self._make_seq(prefill_len + 2) - y_full, _ = _full_scan( - x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size, - ) + y_full, _ = _full_scan(x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size) # Compile Triton before capture without touching the graph's buffers. warmup_bufs = self._make_bufs(max_batch) warmup_state = self._seed_from_prefill( - warmup_bufs, x, dt, B, C, prefill_len, slot, max_batch, - ) - self._decode_one_step( - warmup_bufs, x, dt, B, C, prefill_len, slot, warmup_state + warmup_bufs, x, dt, B, C, prefill_len, slot, max_batch ) + self._decode_one_step(warmup_bufs, x, dt, B, C, prefill_len, slot, warmup_state) bufs = self._make_bufs(max_batch) - ssm_state = self._seed_from_prefill( - bufs, x, dt, B, C, prefill_len, slot, max_batch, - ) + ssm_state = self._seed_from_prefill(bufs, x, dt, B, C, prefill_len, slot, max_batch) batch_indices = torch.tensor([slot], dtype=torch.int32, device=self.device) static_x = x[:, prefill_len : prefill_len + 1].clone() static_dt = dt[:, prefill_len : prefill_len + 1].clone() @@ -610,9 +569,7 @@ def test_cuda_graph_replay_matches_full_scan(self): batch_indices, ) graph.replay() - self._assert_bitwise( - graph_output[0, 0], y_full[0, prefill_len], "CUDA graph replay step 0" - ) + self._assert_bitwise(graph_output[0, 0], y_full[0, prefill_len], "CUDA graph replay step 0") static_x.copy_(x[:, prefill_len + 1 : prefill_len + 2]) static_dt.copy_(dt[:, prefill_len + 1 : prefill_len + 2]) @@ -638,9 +595,18 @@ def test_crossing_with_dominant_carried_state(self): weak_A = self.A * 0.01 y_full, _ = mamba_chunk_scan_combined( - x, dt, weak_A, B, C, self.chunk_size, - D=self.D, z=None, dt_bias=self.dt_bias, dt_softplus=True, - initial_states=None, return_final_states=True, + x, + dt, + weak_A, + B, + C, + self.chunk_size, + D=self.D, + z=None, + dt_bias=self.dt_bias, + dt_softplus=True, + initial_states=None, + return_final_states=True, ) bufs = self._make_bufs(max_batch) @@ -648,21 +614,28 @@ def test_crossing_with_dominant_carried_state(self): cu = torch.tensor([0, prefill_len], dtype=torch.int32, device=self.device) batch_indices = torch.tensor([slot], dtype=torch.int32, device=self.device) bufs.seed( - x[0, :prefill_len], dt[0, :prefill_len], - B[0, :prefill_len], C[0, :prefill_len], cu, batch_indices, + x[0, :prefill_len], + dt[0, :prefill_len], + B[0, :prefill_len], + C[0, :prefill_len], + cu, + batch_indices, ) for k in range(n_decode): pos = prefill_len + k y_batch_invariant = batch_invariant_decode_buffered_scan( bufs, - x[:, pos : pos + 1], dt[:, pos : pos + 1], - B[:, pos : pos + 1], C[:, pos : pos + 1], - weak_A, self.D, self.dt_bias, - batch_indices, ssm_state, - ) - self._assert_bitwise( - y_batch_invariant[0, 0], y_full[0, pos], f"weak-decay step k={k}" + x[:, pos : pos + 1], + dt[:, pos : pos + 1], + B[:, pos : pos + 1], + C[:, pos : pos + 1], + weak_A, + self.D, + self.dt_bias, + batch_indices, + ssm_state, ) + self._assert_bitwise(y_batch_invariant[0, 0], y_full[0, pos], f"weak-decay step k={k}") def test_deterministic_across_calls(self): """Same inputs → bitwise-identical output across repeated invocations.""" @@ -674,15 +647,12 @@ def test_deterministic_across_calls(self): outs = [] for _ in range(3): bufs = self._make_bufs(max_batch) - ssm_state = self._seed_from_prefill( - bufs, x, dt, B, C, prefill_len, slot, max_batch, - ) - outs.append( - self._decode_one_step(bufs, x, dt, B, C, prefill_len, slot, ssm_state) - ) + ssm_state = self._seed_from_prefill(bufs, x, dt, B, C, prefill_len, slot, max_batch) + outs.append(self._decode_one_step(bufs, x, dt, B, C, prefill_len, slot, ssm_state)) for i in range(1, len(outs)): - self.assertTrue(torch.equal(outs[0], outs[i]), - f"determinism: run {i} differs from run 0") + self.assertTrue( + torch.equal(outs[0], outs[i]), f"determinism: run {i} differs from run 0" + ) if __name__ == "__main__": diff --git a/tests/unit_tests/transformer/moe/test_moe_batch_invariant.py b/tests/unit_tests/transformer/moe/test_moe_batch_invariant.py index b9a848068ca..3659dbd6949 100644 --- a/tests/unit_tests/transformer/moe/test_moe_batch_invariant.py +++ b/tests/unit_tests/transformer/moe/test_moe_batch_invariant.py @@ -82,16 +82,10 @@ def test_grouped_gemm_split_invariance(E): # contiguous-layout DeepGEMM — must split on a boundary). half = (E // 2) * per_expert y0 = _bf16_grouped_gemm_contiguous( - x[:half].contiguous(), - w, - m_indices[:half].contiguous(), - counts[: E // 2] + [0] * (E // 2), + x[:half].contiguous(), w, m_indices[:half].contiguous(), counts[: E // 2] + [0] * (E // 2) ) y1 = _bf16_grouped_gemm_contiguous( - x[half:].contiguous(), - w, - m_indices[half:].contiguous(), - [0] * (E // 2) + counts[E // 2 :], + x[half:].contiguous(), w, m_indices[half:].contiguous(), [0] * (E // 2) + counts[E // 2 :] ) y_cat = torch.cat([y0, y1], dim=0) assert torch.equal( From 3892c397b2bf9b83e5306fc8f1f1c4a58c65cc55 Mon Sep 17 00:00:00 2001 From: root Date: Tue, 28 Jul 2026 08:42:06 -0700 Subject: [PATCH 04/14] Fix batch-invariant RL verification issues Signed-off-by: root --- .../attention_context/triton/tensor_ops.py | 28 +++--- .../unit_tests/inference/test_moe_permute.py | 14 ++- .../ssm/ops/test_batch_invariant_decode.py | 95 ++++++++++--------- 3 files changed, 75 insertions(+), 62 deletions(-) diff --git a/megatron/core/inference/contexts/attention_context/triton/tensor_ops.py b/megatron/core/inference/contexts/attention_context/triton/tensor_ops.py index 88efa70994a..6ac68053195 100644 --- a/megatron/core/inference/contexts/attention_context/triton/tensor_ops.py +++ b/megatron/core/inference/contexts/attention_context/triton/tensor_ops.py @@ -128,9 +128,7 @@ def _tensor_masked_update_kernel_2d( + (row_offsets.to(tl.int64) * stride_state_d0) ) src_ptr = ( - NEW_STATES_PTR - + (pid_batch * stride_new_b.to(tl.int64)) - + (row_offsets.to(tl.int64) * stride_new_d0) + NEW_STATES_PTR + (pid_batch * stride_new_b) + (row_offsets.to(tl.int64) * stride_new_d0) ) val = tl.load(src_ptr, mask=mask) @@ -170,21 +168,17 @@ def _tensor_masked_update_kernel_3d( # Given shape (batch, D0, D1) # idx_d1 = flat_idx % D1 # idx_d0 = flat_idx // D1 - idx_d1 = flat_offsets % SIZE_D1.to(tl.int64) - idx_d0 = flat_offsets // SIZE_D1.to(tl.int64) + idx_d1 = flat_offsets % SIZE_D1 + idx_d0 = flat_offsets // SIZE_D1 # Calculate pointers using specific strides dst_offset = ( - (target_idx.to(tl.int64) * stride_state_b.to(tl.int64)) + (target_idx.to(tl.int64) * stride_state_b) + (idx_d0 * stride_state_d0) + (idx_d1 * stride_state_d1) ) - src_offset = ( - (pid_batch * stride_new_b.to(tl.int64)) - + (idx_d0 * stride_new_d0) - + (idx_d1 * stride_new_d1) - ) + src_offset = (pid_batch * stride_new_b) + (idx_d0 * stride_new_d0) + (idx_d1 * stride_new_d1) dst_ptr = STATES_PTR + dst_offset src_ptr = NEW_STATES_PTR + src_offset @@ -232,21 +226,21 @@ def _tensor_masked_update_kernel_4d( # idx_d1 = temp % D1 # idx_d0 = temp // D1 - idx_d2 = flat_offsets % SIZE_D2.to(tl.int64) - temp = flat_offsets // SIZE_D2.to(tl.int64) - idx_d1 = temp % SIZE_D1.to(tl.int64) - idx_d0 = temp // SIZE_D1.to(tl.int64) + idx_d2 = flat_offsets % SIZE_D2 + temp = flat_offsets // SIZE_D2 + idx_d1 = temp % SIZE_D1 + idx_d0 = temp // SIZE_D1 # Calculate pointers using specific strides dst_offset = ( - (target_idx.to(tl.int64) * stride_state_b.to(tl.int64)) + (target_idx.to(tl.int64) * stride_state_b) + (idx_d0 * stride_state_d0) + (idx_d1 * stride_state_d1) + (idx_d2 * stride_state_d2) ) src_offset = ( - (pid_batch * stride_new_b.to(tl.int64)) + (pid_batch * stride_new_b) + (idx_d0 * stride_new_d0) + (idx_d1 * stride_new_d1) + (idx_d2 * stride_new_d2) diff --git a/tests/unit_tests/inference/test_moe_permute.py b/tests/unit_tests/inference/test_moe_permute.py index cdcf96cc201..3d90916ad4b 100644 --- a/tests/unit_tests/inference/test_moe_permute.py +++ b/tests/unit_tests/inference/test_moe_permute.py @@ -50,14 +50,15 @@ def _make_inputs(num_tokens, hidden_dim, topk, num_experts, seed=42): def test_batch_invariant_squared_relu_applies_probs_before_fc2(): - """Match training's BF16 activation and probability rounding order.""" + """Match training's probability placement and BF16 rounding before FC2.""" from megatron.core.activations import squared_relu from megatron.core.inference.moe.activations import padded_squared_relu torch.manual_seed(17) - rows, hidden = 37, 1856 + rows, hidden, output_size = 37, 1856, 512 x = torch.randn(rows, hidden, device="cuda", dtype=torch.bfloat16) probs = torch.rand(rows, device="cuda", dtype=torch.float32) + fc2_weight = torch.randn(output_size, hidden, device="cuda", dtype=torch.bfloat16) permutation_map = torch.arange(rows, device="cuda", dtype=torch.int32) unweighted = padded_squared_relu(x, permutation_map, _vt(rows)) @@ -68,6 +69,15 @@ def test_batch_invariant_squared_relu_applies_probs_before_fc2(): assert torch.equal(unweighted, expected_unweighted) assert torch.equal(actual, expected) + training_output = expected @ fc2_weight.T + inference_output = actual @ fc2_weight.T + old_inference_output = ((unweighted @ fc2_weight.T).float() * probs.unsqueeze(1)).to( + torch.bfloat16 + ) + + assert torch.equal(inference_output, training_output) + assert not torch.equal(old_inference_output, training_output) + @pytest.mark.internal class TestComputeLocalTokensPerExpert: diff --git a/tests/unit_tests/ssm/ops/test_batch_invariant_decode.py b/tests/unit_tests/ssm/ops/test_batch_invariant_decode.py index 994e7433ec9..901cc08894a 100644 --- a/tests/unit_tests/ssm/ops/test_batch_invariant_decode.py +++ b/tests/unit_tests/ssm/ops/test_batch_invariant_decode.py @@ -81,9 +81,9 @@ def _make_seq(self, total_len): torch.randn(1, total_len, ng, n, device=self.device, dtype=self.dtype) * 0.1, ) - def _make_bufs(self, max_batch): + def _make_bufs(self, max_requests): return BatchInvariantDecodeBuffers.allocate( - max_batch, + max_requests, self.chunk_size, self.nh, self.headdim, @@ -93,19 +93,24 @@ def _make_bufs(self, max_batch): self.dtype, ) - def _make_ssm_state(self, max_batch): + def _make_ssm_state(self, max_requests): """Production BIK state cache: FP32 carry across Mamba chunks.""" return torch.zeros( - max_batch, self.nh, self.headdim, self.dstate, device=self.device, dtype=torch.float32 + max_requests, + self.nh, + self.headdim, + self.dstate, + device=self.device, + dtype=torch.float32, ) - def _seed_from_prefill(self, bufs, x, dt, B, C, prefill_len, slot, max_batch): + def _seed_from_prefill(self, bufs, x, dt, B, C, prefill_len, slot, max_requests): """Run the prefill through the reference scan, store its ssm_state at the slot, and seed the batch-invariant buffer with the partial-chunk tail.""" # Production batch-invariant prefill keeps ssm_state at a full Mamba chunk # boundary. Short prefills therefore keep the zero initial boundary; # longer prefills store the largest chunk-aligned prefix state. - ssm_state = self._make_ssm_state(max_batch) + ssm_state = self._make_ssm_state(max_requests) if prefill_len >= self.chunk_size: # Prefill on the largest chunk-aligned prefix; the tail goes in the buffer. aligned = (prefill_len // self.chunk_size) * self.chunk_size @@ -211,7 +216,7 @@ def _assert_bitwise(self, a, b, msg): def test_single_decode_matches_full_scan(self): """Default case: prefill > chunk_size, single decode token, partial tail.""" - max_batch, slot = 4, 1 + max_requests, slot = 4, 1 for prefill_len in [33, 50, 95, 128]: with self.subTest(prefill_len=prefill_len): total = prefill_len + 1 @@ -219,8 +224,10 @@ def test_single_decode_matches_full_scan(self): # Reference: full scan over the whole (prefill + 1) sequence. y_full, _ = _full_scan(x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size) # batch-invariant: seed from prefill, then one decode step. - bufs = self._make_bufs(max_batch) - ssm_state = self._seed_from_prefill(bufs, x, dt, B, C, prefill_len, slot, max_batch) + bufs = self._make_bufs(max_requests) + ssm_state = self._seed_from_prefill( + bufs, x, dt, B, C, prefill_len, slot, max_requests + ) y_batch_invariant = self._decode_one_step( bufs, x, dt, B, C, prefill_len, slot, ssm_state ) @@ -231,7 +238,7 @@ def test_single_decode_matches_full_scan(self): def test_rejects_bf16_state_cache(self): """A rounded state cache cannot preserve carry across multiple chunks.""" x, dt, B, C = self._make_seq(1) - bufs = self._make_bufs(max_batch=2) + bufs = self._make_bufs(max_requests=2) ssm_state = torch.zeros( 2, self.nh, self.headdim, self.dstate, device=self.device, dtype=torch.bfloat16 ) @@ -267,7 +274,7 @@ def test_dynamic_prefill_uses_boundary_state_not_prompt_end_state(self): """Production prefill returns the prompt-end state too, but batch-invariant decode must keep the cache at the last full chunk boundary and put the tail in the replay buffer.""" - max_batch, slot = 4, 1 + max_requests, slot = 4, 1 for prefill_len in [31, 33, 50, 95, 128]: with self.subTest(prefill_len=prefill_len): total = prefill_len + 1 @@ -277,10 +284,10 @@ def test_dynamic_prefill_uses_boundary_state_not_prompt_end_state(self): _, boundary_state = self._varlen_boundary_state_from_prefill( x, dt, B, C, prefill_len ) - ssm_state = torch.randn_like(self._make_ssm_state(max_batch)) + ssm_state = torch.randn_like(self._make_ssm_state(max_requests)) ssm_state[slot] = boundary_state[0] - bufs = self._make_bufs(max_batch) + bufs = self._make_bufs(max_requests) cu = torch.tensor([0, prefill_len], dtype=torch.int32, device=self.device) batch_indices = torch.tensor([slot], dtype=torch.int32, device=self.device) bufs.seed( @@ -303,7 +310,7 @@ def test_dynamic_prefill_uses_boundary_state_not_prompt_end_state(self): def test_chunked_prefill_handoff_matches_full_scan(self): """Splitting prefill at a Mamba boundary preserves exact decode output.""" - max_batch, slot = 2, 0 + max_requests, slot = 2, 0 first_chunk_len = 2 * self.chunk_size for final_chunk_len in [20, self.chunk_size + 13]: @@ -328,9 +335,9 @@ def test_chunked_prefill_handoff_matches_full_scan(self): initial_states=first_boundary, ) - ssm_state = self._make_ssm_state(max_batch) + ssm_state = self._make_ssm_state(max_requests) ssm_state[slot] = final_boundary[0] - bufs = self._make_bufs(max_batch) + bufs = self._make_bufs(max_requests) cu = torch.tensor([0, final_chunk_len], dtype=torch.int32, device=self.device) bufs.seed( x[0, first_chunk_len:prefill_len], @@ -355,14 +362,14 @@ def test_seed_ignores_nonfinite_physical_padding_rows(self): prefix. Seed must duplicate a valid per-sequence tail token into unused replay-buffer rows; otherwise masked future rows can still poison the row-gated Triton dot as 0 * NaN.""" - max_batch, slot = 4, 0 + max_requests, slot = 4, 0 prefill_len = self.chunk_size + 1 total = prefill_len + 1 x, dt, B, C = self._make_seq(total) y_full, _ = _full_scan(x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size) _, boundary_state = self._varlen_boundary_state_from_prefill(x, dt, B, C, prefill_len) - ssm_state = self._make_ssm_state(max_batch) + ssm_state = self._make_ssm_state(max_requests) ssm_state[slot] = boundary_state[0] nan_x = torch.full_like(x[0, :1], float("nan")) @@ -370,7 +377,7 @@ def test_seed_ignores_nonfinite_physical_padding_rows(self): nan_B = torch.full_like(B[0, :1], float("nan")) nan_C = torch.full_like(C[0, :1], float("nan")) - bufs = self._make_bufs(max_batch) + bufs = self._make_bufs(max_requests) cu = torch.tensor([0, prefill_len], dtype=torch.int32, device=self.device) batch_indices = torch.tensor([slot], dtype=torch.int32, device=self.device) bufs.seed( @@ -395,14 +402,16 @@ def test_seed_ignores_nonfinite_physical_padding_rows(self): def test_short_prefill_uses_zero_boundary_state(self): """prefill_len < chunk_size: decode replays from the zero boundary.""" - max_batch, slot = 2, 0 + max_requests, slot = 2, 0 for prefill_len in [1, 7, 16, 31]: with self.subTest(prefill_len=prefill_len): total = prefill_len + 1 x, dt, B, C = self._make_seq(total) y_full, _ = _full_scan(x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size) - bufs = self._make_bufs(max_batch) - ssm_state = self._seed_from_prefill(bufs, x, dt, B, C, prefill_len, slot, max_batch) + bufs = self._make_bufs(max_requests) + ssm_state = self._seed_from_prefill( + bufs, x, dt, B, C, prefill_len, slot, max_requests + ) y_batch_invariant = self._decode_one_step( bufs, x, dt, B, C, prefill_len, slot, ssm_state ) @@ -413,15 +422,15 @@ def test_short_prefill_uses_zero_boundary_state(self): def test_multi_step_decode_across_chunk_boundary(self): """Step decode several times so the per-slot buffer fills, crosses a chunk boundary, and resets. Each step must match the full scan.""" - max_batch, slot = 2, 0 + max_requests, slot = 2, 0 prefill_len = 20 # < chunk_size, so first decode step will keep growing buf n_decode = self.chunk_size + 5 # enough to cross at least one boundary total = prefill_len + n_decode x, dt, B, C = self._make_seq(total) y_full, _ = _full_scan(x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size) - bufs = self._make_bufs(max_batch) - ssm_state = self._seed_from_prefill(bufs, x, dt, B, C, prefill_len, slot, max_batch) + bufs = self._make_bufs(max_requests) + ssm_state = self._seed_from_prefill(bufs, x, dt, B, C, prefill_len, slot, max_requests) for k in range(n_decode): pos = prefill_len + k @@ -435,7 +444,7 @@ def test_multi_step_decode_across_chunk_boundary(self): def test_multi_slot_independent_streams(self): """Two slots with different prefill lengths decoded in the same call — each slot's output must match its own full scan.""" - max_batch = 4 + max_requests = 4 slots = [0, 2] prefill_lens = [25, 70] # one short, one long with a boundary state x_per_slot, dt_per_slot, B_per_slot, C_per_slot = [], [], [], [] @@ -449,13 +458,13 @@ def test_multi_slot_independent_streams(self): y_full, _ = _full_scan(x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size) y_refs.append(y_full[0, plen]) - bufs = self._make_bufs(max_batch) + bufs = self._make_bufs(max_requests) # Per-slot seeding (each slot's prefill done independently). - ssm_state = self._make_ssm_state(max_batch) + ssm_state = self._make_ssm_state(max_requests) for slot, plen, x, dt, B, C in zip( slots, prefill_lens, x_per_slot, dt_per_slot, B_per_slot, C_per_slot ): - partial = self._seed_from_prefill(bufs, x, dt, B, C, plen, slot, max_batch) + partial = self._seed_from_prefill(bufs, x, dt, B, C, plen, slot, max_requests) ssm_state[slot] = partial[slot] # Both slots step at once. @@ -494,15 +503,15 @@ def test_inactive_padding_entries(self): """batch_indices mixing -1 padding entries with active slot 0 (the CUDA- graph padding pattern). Padding entries must not write replay buffers, perturb slot 0, or produce nonzero output.""" - max_batch, slot = 3, 0 + max_requests, slot = 3, 0 prefill_len = 50 n_decode = 8 total = prefill_len + n_decode x, dt, B, C = self._make_seq(total) y_full, _ = _full_scan(x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size) - bufs = self._make_bufs(max_batch) - ssm_state = self._seed_from_prefill(bufs, x, dt, B, C, prefill_len, slot, max_batch) + bufs = self._make_bufs(max_requests) + ssm_state = self._seed_from_prefill(bufs, x, dt, B, C, prefill_len, slot, max_requests) batch_indices = torch.tensor([slot, -1, -1], dtype=torch.int32, device=self.device) inactive_counts = bufs.num_buffered[1:].clone() for k in range(n_decode): @@ -564,20 +573,20 @@ def test_seed_skips_padding_entries(self): def test_cuda_graph_replay_matches_full_scan(self): """A captured decode step advances persistent state exactly across replays.""" - max_batch, slot = 2, 0 + max_requests, slot = 2, 0 prefill_len = 20 x, dt, B, C = self._make_seq(prefill_len + 2) y_full, _ = _full_scan(x, dt, self.A, B, C, self.D, self.dt_bias, self.chunk_size) # Compile Triton before capture without touching the graph's buffers. - warmup_bufs = self._make_bufs(max_batch) + warmup_bufs = self._make_bufs(max_requests) warmup_state = self._seed_from_prefill( - warmup_bufs, x, dt, B, C, prefill_len, slot, max_batch + warmup_bufs, x, dt, B, C, prefill_len, slot, max_requests ) self._decode_one_step(warmup_bufs, x, dt, B, C, prefill_len, slot, warmup_state) - bufs = self._make_bufs(max_batch) - ssm_state = self._seed_from_prefill(bufs, x, dt, B, C, prefill_len, slot, max_batch) + bufs = self._make_bufs(max_requests) + ssm_state = self._seed_from_prefill(bufs, x, dt, B, C, prefill_len, slot, max_requests) batch_indices = torch.tensor([slot], dtype=torch.int32, device=self.device) static_x = x[:, prefill_len : prefill_len + 1].clone() static_z = static_x.clone() @@ -630,7 +639,7 @@ def test_crossing_with_dominant_carried_state(self): (weak decay: A ~ -0.01 → exp(dA_cs) ≈ 1). Guards the pipeline ordering and FP32 state-passing carry. With strong decay, either corruption can round away in BF16 and hide.""" - max_batch, slot = 2, 0 + max_requests, slot = 2, 0 prefill_len = 20 # Cross twice: the second transition detects an accidental BF16 # store/reload of state passing's FP32 carry. @@ -654,8 +663,8 @@ def test_crossing_with_dominant_carried_state(self): return_final_states=True, ) - bufs = self._make_bufs(max_batch) - ssm_state = self._make_ssm_state(max_batch) + bufs = self._make_bufs(max_requests) + ssm_state = self._make_ssm_state(max_requests) cu = torch.tensor([0, prefill_len], dtype=torch.int32, device=self.device) batch_indices = torch.tensor([slot], dtype=torch.int32, device=self.device) bufs.seed( @@ -686,15 +695,15 @@ def test_crossing_with_dominant_carried_state(self): def test_deterministic_across_calls(self): """Same inputs → bitwise-identical output across repeated invocations.""" - max_batch, slot = 2, 0 + max_requests, slot = 2, 0 prefill_len = 50 total = prefill_len + 1 x, dt, B, C = self._make_seq(total) outs = [] for _ in range(3): - bufs = self._make_bufs(max_batch) - ssm_state = self._seed_from_prefill(bufs, x, dt, B, C, prefill_len, slot, max_batch) + bufs = self._make_bufs(max_requests) + ssm_state = self._seed_from_prefill(bufs, x, dt, B, C, prefill_len, slot, max_requests) outs.append(self._decode_one_step(bufs, x, dt, B, C, prefill_len, slot, ssm_state)) for i in range(1, len(outs)): self.assertTrue( From 52858063d62897e40e3dc88e096f2531372b19d9 Mon Sep 17 00:00:00 2001 From: root Date: Tue, 28 Jul 2026 15:25:52 -0700 Subject: [PATCH 05/14] Finalize exact batch-invariant Mamba and MoE execution Signed-off-by: root --- .../attention_context/triton/tensor_ops.py | 28 ++-- .../inference/contexts/dynamic_context.py | 20 ++- .../contexts/mamba_slot_allocator.py | 22 +++- megatron/core/ssm/mamba_mixer.py | 15 +-- .../core/ssm/ops/batch_invariant_decode.py | 78 ++++++++--- megatron/core/ssm/ops/ssd_chunk_scan.py | 1 - megatron/core/ssm/ops/ssd_chunk_state.py | 4 +- megatron/core/ssm/ops/ssd_state_passing.py | 8 +- .../custom_layers/batch_invariant_kernels.py | 124 +++--------------- .../core/transformer/transformer_config.py | 4 + megatron/rl/rl_utils.py | 12 +- .../contexts/test_dynamic_prefix_caching.py | 30 ++++- .../unit_tests/rl/test_rl_batch_invariant.py | 3 +- 13 files changed, 183 insertions(+), 166 deletions(-) diff --git a/megatron/core/inference/contexts/attention_context/triton/tensor_ops.py b/megatron/core/inference/contexts/attention_context/triton/tensor_ops.py index 6ac68053195..88efa70994a 100644 --- a/megatron/core/inference/contexts/attention_context/triton/tensor_ops.py +++ b/megatron/core/inference/contexts/attention_context/triton/tensor_ops.py @@ -128,7 +128,9 @@ def _tensor_masked_update_kernel_2d( + (row_offsets.to(tl.int64) * stride_state_d0) ) src_ptr = ( - NEW_STATES_PTR + (pid_batch * stride_new_b) + (row_offsets.to(tl.int64) * stride_new_d0) + NEW_STATES_PTR + + (pid_batch * stride_new_b.to(tl.int64)) + + (row_offsets.to(tl.int64) * stride_new_d0) ) val = tl.load(src_ptr, mask=mask) @@ -168,17 +170,21 @@ def _tensor_masked_update_kernel_3d( # Given shape (batch, D0, D1) # idx_d1 = flat_idx % D1 # idx_d0 = flat_idx // D1 - idx_d1 = flat_offsets % SIZE_D1 - idx_d0 = flat_offsets // SIZE_D1 + idx_d1 = flat_offsets % SIZE_D1.to(tl.int64) + idx_d0 = flat_offsets // SIZE_D1.to(tl.int64) # Calculate pointers using specific strides dst_offset = ( - (target_idx.to(tl.int64) * stride_state_b) + (target_idx.to(tl.int64) * stride_state_b.to(tl.int64)) + (idx_d0 * stride_state_d0) + (idx_d1 * stride_state_d1) ) - src_offset = (pid_batch * stride_new_b) + (idx_d0 * stride_new_d0) + (idx_d1 * stride_new_d1) + src_offset = ( + (pid_batch * stride_new_b.to(tl.int64)) + + (idx_d0 * stride_new_d0) + + (idx_d1 * stride_new_d1) + ) dst_ptr = STATES_PTR + dst_offset src_ptr = NEW_STATES_PTR + src_offset @@ -226,21 +232,21 @@ def _tensor_masked_update_kernel_4d( # idx_d1 = temp % D1 # idx_d0 = temp // D1 - idx_d2 = flat_offsets % SIZE_D2 - temp = flat_offsets // SIZE_D2 - idx_d1 = temp % SIZE_D1 - idx_d0 = temp // SIZE_D1 + idx_d2 = flat_offsets % SIZE_D2.to(tl.int64) + temp = flat_offsets // SIZE_D2.to(tl.int64) + idx_d1 = temp % SIZE_D1.to(tl.int64) + idx_d0 = temp // SIZE_D1.to(tl.int64) # Calculate pointers using specific strides dst_offset = ( - (target_idx.to(tl.int64) * stride_state_b) + (target_idx.to(tl.int64) * stride_state_b.to(tl.int64)) + (idx_d0 * stride_state_d0) + (idx_d1 * stride_state_d1) + (idx_d2 * stride_state_d2) ) src_offset = ( - (pid_batch * stride_new_b) + (pid_batch * stride_new_b.to(tl.int64)) + (idx_d0 * stride_new_d0) + (idx_d1 * stride_new_d1) + (idx_d2 * stride_new_d2) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index b542bd2cc06..0d45e34967f 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -1094,7 +1094,10 @@ def initialize_all_tensors(self) -> None: + _mha_cu_kv_seq_lengths_bytes + _mha_block_table_bytes ) - # Mamba section (hybrid models only). Must match MambaMetadata and ContextGPUView. + # Mamba section (hybrid models only). Must match the MambaMetadata + # shapes (mirrors the layout documented in ContextGPUView). + # batch_indices_decode is int32 in batch-invariant mode and int64 otherwise; + # all other fields are int32. if self.is_hybrid_model: self._mamba_decode_indices_dtype = ( torch.int32 if self.batch_invariant_mode else torch.int64 @@ -2844,6 +2847,15 @@ def _find_mamba_match_count( mamba_map = self.mamba_slot_allocator.hash_to_block_id hashes = req.precomputed_block_hashes[start_block:end_block] for i in range(len(hashes) - 1, -1, -1): + block_count = start_block + i + 1 + if ( + self.batch_invariant_mode + and (block_count * self.block_size_tokens) % self.mamba_chunk_size + ): + # Restarting between Mamba chunk boundaries changes the + # subsequent reduction grouping. Recompute from the latest + # aligned cached state instead. + continue if hashes[i] in mamba_map: return i + 1 return 0 @@ -4248,10 +4260,8 @@ def calculate_log_probs( log_probs (Tensor): Used to compute top n logprobs later if required. """ - # Keep the same log-softmax input dtype as training in batch-invariant mode. - logits_squeezed = logits.squeeze(0) - if not self.batch_invariant_mode: - logits_squeezed = logits_squeezed.float() + # Calculate log_probs (sequence_length x vocab_size) + logits_squeezed = logits.squeeze(0).float() n_active = self.total_request_count - self.paused_request_count if only_last_token_logits or self.is_decode_only(): diff --git a/megatron/core/inference/contexts/mamba_slot_allocator.py b/megatron/core/inference/contexts/mamba_slot_allocator.py index 21116b0cb7e..ab8d1d67478 100644 --- a/megatron/core/inference/contexts/mamba_slot_allocator.py +++ b/megatron/core/inference/contexts/mamba_slot_allocator.py @@ -1,5 +1,6 @@ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +import math from typing import TYPE_CHECKING, Dict import torch @@ -437,10 +438,19 @@ def compute_and_store_offsets( # in MambaMetadata (offset -> chunk-index conversion) to stay consistent. mamba_chunk_size = ctx.mamba_chunk_size + candidates = (kv_div_abs, last_aligned_abs, penultimate_abs) + if ctx.batch_invariant_mode: + # A reusable BIK state must be on both the KV-block and Mamba-chunk + # grids. Round candidates back so a short suffix is recomputed. + restore_stride = math.lcm(bs, mamba_chunk_size) + candidates = tuple( + position // restore_stride * restore_stride for position in candidates + ) + # Keep only boundaries that land inside this chunk's computed tokens and on # a mamba-chunk boundary (required for mid-sequence state extraction). offsets_set = set() - for abs_pos in (kv_div_abs, last_aligned_abs, penultimate_abs): + for abs_pos in candidates: offset = abs_pos - chunk_start if offset > 0 and offset < seq_len and offset % mamba_chunk_size == 0: offsets_set.add(offset) @@ -466,7 +476,15 @@ def compute_and_store_offsets( # cached directly. Only valid on the final chunk (otherwise the live state # is mid-prompt). Non-block-aligned prompts cache their last complete block # via the intermediate-extraction path above instead. - if is_last_chunk and last_aligned_abs == prompt_len and prompt_len > 0: + is_reusable_mamba_boundary = ( + not ctx.batch_invariant_mode or prompt_len % mamba_chunk_size == 0 + ) + if ( + is_last_chunk + and last_aligned_abs == prompt_len + and prompt_len > 0 + and is_reusable_mamba_boundary + ): last_block_idx = prompt_len // bs - 1 if last_block_idx >= 0: self._eos_cache_block_id_cpu[current_id] = ctx.request_to_kv_block_ids[current_id][ diff --git a/megatron/core/ssm/mamba_mixer.py b/megatron/core/ssm/mamba_mixer.py index dc776d1d0b5..31b30748e9e 100644 --- a/megatron/core/ssm/mamba_mixer.py +++ b/megatron/core/ssm/mamba_mixer.py @@ -986,8 +986,6 @@ def _ssm_prefill( # Batch-invariant decode replays the partial prefill tail, so keep # the cached SSM state at the last complete chunk boundary. - boundary_chunk_indices = None - has_boundary = None if self.config.batch_invariant_mode: prefill_lens = (cu_seqlens[1:] - cu_seqlens[:-1]).to(torch.long) tail_lens = prefill_lens % self.chunk_size @@ -1263,7 +1261,12 @@ def _ssm_decode( dim=-1, ) # SSM step - if selective_state_update is None: + if self.config.batch_invariant_mode: + assert ( + batch_indices is not None + ), "batch_invariant_mode for Mamba decode requires batch_indices from dynamic batching." + y = self._get_batch_invariant_decoder().step(x, z, dt, B, C, batch_indices, ssm_state) + elif selective_state_update is None: # Fallback uses 1D A; the decode cache is pre-expanded for Triton. A = -torch.exp(self.A_log.float()) # TODO(ksanthanam): Consider deprecating this path @@ -1320,13 +1323,7 @@ def _ssm_decode( y = rearrange(y, "b h p -> b (h p)") if not self.rmsnorm: y = y * self.act(z) # (B D) - y = y.unsqueeze(1) # Restore seq dimension - elif self.config.batch_invariant_mode: - assert ( - batch_indices is not None - ), "batch_invariant_mode for Mamba decode requires batch_indices from dynamic batching." - y = self._get_batch_invariant_decoder().step(x, z, dt, B, C, batch_indices, ssm_state) else: A = self._get_decode_A_neg_exp() diff --git a/megatron/core/ssm/ops/batch_invariant_decode.py b/megatron/core/ssm/ops/batch_invariant_decode.py index c3990fef20d..ba3e866db7c 100644 --- a/megatron/core/ssm/ops/batch_invariant_decode.py +++ b/megatron/core/ssm/ops/batch_invariant_decode.py @@ -4,13 +4,55 @@ from dataclasses import dataclass import torch +import triton +import triton.language as tl -from megatron.core.inference.contexts.attention_context.triton.tensor_ops import ( - tensor_masked_update, -) from megatron.core.ssm.ops.ssd_combined import mamba_chunk_scan_decode_rows +@triton.jit +def _masked_update_rows_kernel( + states_ptr, + indices_ptr, + values_ptr, + state_row_stride, + value_row_stride, + ROW_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """Copy contiguous rows, skipping entries whose destination index is -1.""" + src_row = tl.program_id(0) + dst_row = tl.load(indices_ptr + src_row) + if dst_row < 0: + return + + offsets = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < ROW_SIZE + values = tl.load(values_ptr + src_row * value_row_stride + offsets, mask=mask) + tl.store(states_ptr + dst_row * state_row_stride + offsets, values, mask=mask) + + +def _masked_update_rows(states: torch.Tensor, indices: torch.Tensor, values: torch.Tensor) -> None: + """Copy rows into persistent BIK buffers without touching inactive graph lanes.""" + assert states.ndim == values.ndim == 2 + assert states.stride(1) == values.stride(1) == 1 + assert indices.dtype == torch.int32 and indices.numel() == values.shape[0] + row_size = states.shape[1] + assert values.shape[1] == row_size + + block_size = min(triton.next_power_of_2(row_size), 1024) + grid = (values.shape[0], triton.cdiv(row_size, block_size)) + _masked_update_rows_kernel[grid]( + states, + indices, + values, + states.stride(0), + values.stride(0), + ROW_SIZE=row_size, + BLOCK_SIZE=block_size, + ) + + @dataclass class BatchInvariantDecodeBuffers: """Per-slot persistent state for the buffered decode scan.""" @@ -88,12 +130,12 @@ def seed( ) slots = batch_indices[:num_seqs] - tensor_masked_update(self.x.flatten(1), slots, x[tail_token_idx].flatten(1)) - tensor_masked_update(self.z.flatten(1), slots, z[tail_token_idx].flatten(1)) - tensor_masked_update(self.dt.flatten(1), slots, dt[tail_token_idx].flatten(1)) - tensor_masked_update(self.B.flatten(1), slots, B[tail_token_idx].flatten(1)) - tensor_masked_update(self.C.flatten(1), slots, C[tail_token_idx].flatten(1)) - tensor_masked_update( + _masked_update_rows(self.x.flatten(1), slots, x[tail_token_idx].flatten(1)) + _masked_update_rows(self.z.flatten(1), slots, z[tail_token_idx].flatten(1)) + _masked_update_rows(self.dt.flatten(1), slots, dt[tail_token_idx].flatten(1)) + _masked_update_rows(self.B.flatten(1), slots, B[tail_token_idx].flatten(1)) + _masked_update_rows(self.C.flatten(1), slots, C[tail_token_idx].flatten(1)) + _masked_update_rows( self.num_buffered.unsqueeze(1), slots, tail_lens.to(torch.int32).unsqueeze(1) ) @@ -139,13 +181,17 @@ def batch_invariant_decode_buffered_scan( active = batch_indices >= 0 safe_slots = batch_indices.clamp_min(0) write_pos = buffers.num_buffered[safe_slots].to(torch.long) - buffer_rows = torch.where(active, safe_slots * chunk_size + write_pos, -1) + buffer_rows = torch.where(active, safe_slots * chunk_size + write_pos, -1).to(torch.int32) - tensor_masked_update(buffers.x.view(-1, nheads, headdim), buffer_rows, x[:, 0]) - tensor_masked_update(buffers.z.view(-1, nheads, headdim), buffer_rows, z[:, 0]) - tensor_masked_update(buffers.dt.view(-1, nheads), buffer_rows, dt[:, 0]) - tensor_masked_update(buffers.B.view(-1, buffers.B.shape[-2], dstate), buffer_rows, B[:, 0]) - tensor_masked_update(buffers.C.view(-1, buffers.C.shape[-2], dstate), buffer_rows, C[:, 0]) + _masked_update_rows(buffers.x.view(-1, nheads * headdim), buffer_rows, x[:, 0].flatten(1)) + _masked_update_rows(buffers.z.view(-1, nheads * headdim), buffer_rows, z[:, 0].flatten(1)) + _masked_update_rows(buffers.dt.view(-1, nheads), buffer_rows, dt[:, 0]) + _masked_update_rows( + buffers.B.view(-1, buffers.B.shape[-2] * dstate), buffer_rows, B[:, 0].flatten(1) + ) + _masked_update_rows( + buffers.C.view(-1, buffers.C.shape[-2] * dstate), buffer_rows, C[:, 0].flatten(1) + ) crossed = active & (write_pos + 1 == chunk_size) target_rows.copy_(torch.where(active, write_pos, -1).to(torch.int32)) @@ -175,7 +221,7 @@ def batch_invariant_decode_buffered_scan( ) next_write_pos = torch.where(crossed, 0, write_pos + 1).to(torch.int32) - tensor_masked_update( + _masked_update_rows( buffers.num_buffered.unsqueeze(1), batch_indices, next_write_pos.unsqueeze(1) ) diff --git a/megatron/core/ssm/ops/ssd_chunk_scan.py b/megatron/core/ssm/ops/ssd_chunk_scan.py index 6e6b12d1e7f..8066965b12b 100644 --- a/megatron/core/ssm/ops/ssd_chunk_scan.py +++ b/megatron/core/ssm/ops/ssd_chunk_scan.py @@ -174,7 +174,6 @@ def _chunk_scan_fwd_kernel( seq_idx = tl.load(seq_idx_ptr) if HAS_TARGET_ROWS: # Each fixed window starts from its indexed cached state. - seq_idx_prev = -1 prev_states_ptr = ( initstates_ptr + seq_idx * stride_init_states_batch + pid_h * stride_init_states_head ) diff --git a/megatron/core/ssm/ops/ssd_chunk_state.py b/megatron/core/ssm/ops/ssd_chunk_state.py index 9f70c9ccb3d..8a21de1bef7 100644 --- a/megatron/core/ssm/ops/ssd_chunk_state.py +++ b/megatron/core/ssm/ops/ssd_chunk_state.py @@ -57,7 +57,6 @@ def _chunk_cumsum_fwd_kernel( seqlen, nheads: tl.constexpr, chunk_size: tl.constexpr, - HAS_CHUNK_STARTS: tl.constexpr, HAS_TARGET_ROWS: tl.constexpr, dt_min: tl.constexpr, dt_max: tl.constexpr, @@ -87,7 +86,7 @@ def _chunk_cumsum_fwd_kernel( return chunk_seqlen_start = tl.load(chunk_offsets_ptr + pid_c) - if HAS_CHUNK_STARTS: + if HAS_TARGET_ROWS: # Fixed windows need only a start offset. chunk_seqlen_end = chunk_seqlen_start + chunk_size else: @@ -341,7 +340,6 @@ def _chunk_cumsum_fwd( stride_dA_cs_head=dA_cumsum.stride(0), stride_dA_cs_chunk=dA_cumsum.stride(1), stride_dA_cs_csize=dA_cumsum.stride(2), - HAS_CHUNK_STARTS=chunk_starts is not None, HAS_TARGET_ROWS=has_target_rows, DT_SOFTPLUS=dt_softplus, HAS_DT_BIAS=dt_bias is not None, diff --git a/megatron/core/ssm/ops/ssd_state_passing.py b/megatron/core/ssm/ops/ssd_state_passing.py index 07033067641..9c12537d7f1 100644 --- a/megatron/core/ssm/ops/ssd_state_passing.py +++ b/megatron/core/ssm/ops/ssd_state_passing.py @@ -154,10 +154,8 @@ def _state_passing_fwd( if not has_dst: out_dtype = states.dtype if out_dtype is None else out_dtype out = torch.empty((nchunks, nheads, dim), device=states.device, dtype=out_dtype) - out_strides = out.stride() else: out = states - out_strides = (0, 0, 0) initial_states_strides = ( (initial_states.stride(0), initial_states.stride(1), initial_states.stride(2)) @@ -189,9 +187,9 @@ def _state_passing_fwd( stride_states_chunk=states.stride(0), stride_states_head=states.stride(1), stride_states_dim=states.stride(2), - stride_out_chunk=out_strides[0], - stride_out_head=out_strides[1], - stride_out_dim=out_strides[2], + stride_out_chunk=out.stride(0), + stride_out_head=out.stride(1), + stride_out_dim=out.stride(2), stride_dA_cs_head=dA_cumsum.stride(0), stride_dA_cs_chunk=dA_cumsum.stride(1), stride_dA_cs_csize=dA_cumsum.stride(2), diff --git a/megatron/core/transformer/custom_layers/batch_invariant_kernels.py b/megatron/core/transformer/custom_layers/batch_invariant_kernels.py index b6f058a1669..abb21f36edb 100644 --- a/megatron/core/transformer/custom_layers/batch_invariant_kernels.py +++ b/megatron/core/transformer/custom_layers/batch_invariant_kernels.py @@ -6,7 +6,6 @@ import contextlib import importlib import importlib.util -import inspect import logging from collections import namedtuple from collections.abc import Callable @@ -713,8 +712,12 @@ def _patched(*args, **kwargs): # TEGroupedMLP (forward + dgrad + wgrad) goes through DeepGEMM in bf16. _te_patch_general_grouped_gemm() - # Fused LayerNormLinear and LayerNormMLP call apply_normalization - # instead of RMSNorm.forward. + # Patch the fused-module normalization entry (`apply_normalization`). TE's + # fused LayerNormLinear / LayerNormMLP call this instead of RMSNorm.forward, + # so without this patch their internal RMSNorm runs TE's tex kernel, whose + # within-row reduction strategy depends on the total row count — i.e. it is + # NOT batch-invariant (observed: same rows, different output at 928 vs 2274 + # rows on GB200, 1 bf16 ulp per layer, amplifying across depth). import transformer_engine.pytorch.module._common as te_common for mod_name, mod in ( @@ -796,68 +799,13 @@ def _te_unpatch_general_grouped_gemm() -> None: _TE_GROUPED_GEMM_FUNC_ORIGS.pop(key, None) -def _get_original_te_grouped_gemm(): - for key in ( - "module.grouped_linear.general_grouped_gemm", - "cpp_extensions.general_grouped_gemm", - "cpp_extensions.gemm.general_grouped_gemm", - ): - orig = _TE_GROUPED_GEMM_FUNC_ORIGS.get(key) - if orig is not None: - return orig - return None - - -def _original_te_grouped_gemm_has_quantization_params(orig) -> bool: - try: - return "quantization_params" in inspect.signature(orig).parameters - except (TypeError, ValueError): - return False - - -def _call_original_te_grouped_gemm( - orig, - A, - B, - out, - quantization_params, - out_dtype, - *, - layout, - m_splits, - gelu, - grad, - accumulate, - bias, - use_bias, - use_split_accumulator, - D_dtype, - single_output, -): - kwargs = dict( - layout=layout, - m_splits=m_splits, - gelu=gelu, - grad=grad, - accumulate=accumulate, - bias=bias, - use_bias=use_bias, - use_split_accumulator=use_split_accumulator, - D_dtype=D_dtype, - single_output=single_output, - ) - if _original_te_grouped_gemm_has_quantization_params(orig): - return orig(A, B, out, quantization_params, out_dtype, **kwargs) - return orig(A, B, out, out_dtype, **kwargs) - - def _is_bf16_grouped_path(A, B, quantization_params, gelu: bool) -> bool: """Decide if TE's general_grouped_gemm call can be served by DeepGEMM bf16.""" if gelu: return False if not HAVE_DEEPGEMM_BF16: return False - if not (isinstance(A, list) and isinstance(B, list)): + if not isinstance(A, (list, tuple)) or not isinstance(B, (list, tuple)): return False if len(A) != len(B) or len(A) == 0: return False @@ -891,9 +839,8 @@ def _te_general_grouped_gemm_patched( """Batch-invariant replacement for TE general_grouped_gemm. Dispatches by (layout, single_output, grad) to forward / dgrad / wgrad - implementations backed by DeepGEMM. Falls back to TE's original for any - case we cannot guarantee batch-invariant: quantized inputs, gelu fusion, - non-bf16 dtypes, or unsupported (layout, mode) combinations. + implementations backed by DeepGEMM. Unsupported calls fail rather than + silently using a kernel that is not guaranteed batch invariant. """ # TE versions differ here: # old: general_grouped_gemm(A, B, out, out_dtype, ...) @@ -903,29 +850,9 @@ def _te_general_grouped_gemm_patched( quantization_params = None if not _is_bf16_grouped_path(A, B, quantization_params, gelu): - orig = _get_original_te_grouped_gemm() - if orig is None: - raise RuntimeError( - "Batch-invariant grouped GEMM patch was invoked but no original " - "TE general_grouped_gemm was captured; patching order issue." - ) - return _call_original_te_grouped_gemm( - orig, - A, - B, - out, - quantization_params, - out_dtype, - layout=layout, - m_splits=m_splits, - gelu=gelu, - grad=grad, - accumulate=accumulate, - bias=bias, - use_bias=use_bias, - use_split_accumulator=use_split_accumulator, - D_dtype=D_dtype, - single_output=single_output, + raise RuntimeError( + "Batch-invariant grouped GEMM requires unquantized BF16 tensor sequences " + "with GELU fusion disabled." ) # Dispatch by TE's call convention. @@ -939,30 +866,9 @@ def _te_general_grouped_gemm_patched( return _batch_invariant_te_grouped_dgrad(A, B, out, m_splits, accumulate) if (not single_output) and layout == "NT" and grad: return _batch_invariant_te_grouped_wgrad(A, B, out, m_splits, use_bias, accumulate) - # Unknown TE call shape — defer to the original. - orig = _get_original_te_grouped_gemm() - if orig is None: - raise RuntimeError( - "Batch-invariant grouped GEMM patch was invoked but no original " - "TE general_grouped_gemm was captured; patching order issue." - ) - return _call_original_te_grouped_gemm( - orig, - A, - B, - out, - quantization_params, - out_dtype, - layout=layout, - m_splits=m_splits, - gelu=gelu, - grad=grad, - accumulate=accumulate, - bias=bias, - use_bias=use_bias, - use_split_accumulator=use_split_accumulator, - D_dtype=D_dtype, - single_output=single_output, + raise RuntimeError( + "Unsupported batch-invariant grouped GEMM call: " + f"layout={layout!r}, single_output={single_output}, grad={grad}." ) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 0d3c5cf2668..99b2a232b47 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -2864,6 +2864,10 @@ def _scope_to_str(s): ) if self.batch_invariant_mode: + assert self.params_dtype == torch.bfloat16, ( + "Batch invariant mode supports BF16 model parameters only; " + f"got {self.params_dtype}." + ) assert ( self.attention_backend == AttnBackend.flash ), "Batch invariant mode only supports FlashAttention (--attention-backend flash)" diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index a8922e69709..2f9a7360e44 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -737,6 +737,9 @@ def selective_log_softmax(logits, index): # logsumexp approach is unstable with bfloat16, fall back to slightly less efficent approach per_token_logps = [] for row_logits, row_labels in zip(logits, index): # loop to reduce peak mem consumption + # Match inference by running batch-invariant log-softmax in FP32. + if use_bik_logsoftmax: + row_logits = row_logits.float() row_logps = torch.nn.functional.log_softmax(row_logits, dim=-1) row_per_token_logps = row_logps.gather(dim=-1, index=row_labels.unsqueeze(-1)).squeeze( -1 @@ -1627,9 +1630,12 @@ def prepare_data_for_update( use_single_mempool=args.cuda_graph_use_single_mempool, ) - dtype = ( - torch.bfloat16 if args.bf16 else (torch.float16 if args.fp16 else torch.float32) - ) + if is_batch_invariant_mode_enabled(): + dtype = torch.float32 + else: + dtype = ( + torch.bfloat16 if args.bf16 else (torch.float16 if args.fp16 else torch.float32) + ) pg_collection = get_attr_wrapped_model(model, "pg_collection") pp_group = pg_collection.pp diff --git a/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py b/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py index 81f55da2a17..78143a195a7 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py @@ -70,7 +70,7 @@ def _ctx( DynamicInferenceContext.REQUEST_ROUNDER = rounder transformer_config = TransformerConfig( - params_dtype=torch.float32, + params_dtype=torch.bfloat16 if batch_invariant_mode else torch.float32, num_layers=4, kv_channels=8, num_attention_heads=2, @@ -83,6 +83,9 @@ def _ctx( flash_attention_version=4 if batch_invariant_mode else None, attention_dropout=0.0 if batch_invariant_mode else 0.1, ) + if batch_invariant_mode: + max_tokens = 512 if max_tokens is None else max_tokens + max_requests = 64 if max_requests is None else max_requests inference_config = InferenceConfig( max_sequence_length=max_sequence_length, buffer_size_gb=buffer_size_gb, @@ -809,6 +812,20 @@ def test_mamba_cache_lifecycle(self): self._mamba_allocate_and_register(ctx6, self._block_ids(ctx6, 0, 4)[:2]) req6 = self._req(ctx6, p6.clone(), request_id=2) assert ctx6._find_mamba_match_count(req6, 0, len(req6.precomputed_block_hashes)) == 2 + + # BIK restores only at Mamba chunk boundaries. A later unaligned + # cached state would change the reduction grouping after the restore. + ctx_bik = self._mctx(block_size_tokens=32, batch_invariant_mode=True) + p_bik = self._prompt(32 * 8) + ctx_bik.add_request(self._req(ctx_bik, p_bik.clone())) + self._mamba_allocate_and_register(ctx_bik, self._block_ids(ctx_bik, 0, 8)[:5]) + req_bik = self._req(ctx_bik, p_bik.clone(), request_id=2) + assert ( + ctx_bik._find_mamba_match_count(req_bik, 0, len(req_bik.precomputed_block_hashes)) == 4 + ) + *_, prefix_skip, effective_prefill = ctx_bik._compute_prefix_match(req_bik, len(p_bik)) + assert prefix_skip == 128 and effective_prefill == 128 + # no match when no mamba hashes registered ctx7 = self._mctx() ctx7.add_request(self._req(ctx7, self._prompt(bs * 3))) @@ -1020,6 +1037,17 @@ def test_mamba_intermediate_offsets(self): ctx3.transfer_bookkeeping_to_gpu() assert ctx3.mamba_slot_allocator._eos_cache_block_id_cpu[1].item() >= 0 + # BIK does not spend durable slots on block boundaries that cannot be + # restored without changing Mamba's chunk reduction grouping. + ctx_bik_unaligned = self._mctx(block_size_tokens=32, batch_invariant_mode=True) + ctx_bik_unaligned.add_request(self._req(ctx_bik_unaligned, self._prompt(160).clone())) + assert ctx_bik_unaligned.mamba_slot_allocator._eos_cache_block_id_cpu[0].item() < 0 + assert ctx_bik_unaligned.mamba_slot_allocator._intermediate_offsets_cpu[0, 0].item() == 128 + + ctx_bik_aligned = self._mctx(block_size_tokens=32, batch_invariant_mode=True) + ctx_bik_aligned.add_request(self._req(ctx_bik_aligned, self._prompt(128).clone())) + assert ctx_bik_aligned.mamba_slot_allocator._eos_cache_block_id_cpu[0].item() >= 0 + # intermediate output buffers are pre-allocated ctx4 = self._mctx() msa4 = ctx4.mamba_slot_allocator diff --git a/tests/unit_tests/rl/test_rl_batch_invariant.py b/tests/unit_tests/rl/test_rl_batch_invariant.py index 7e62022d6b8..53545093e72 100644 --- a/tests/unit_tests/rl/test_rl_batch_invariant.py +++ b/tests/unit_tests/rl/test_rl_batch_invariant.py @@ -14,7 +14,7 @@ def test_selective_log_softmax_batch_invariant(): B, S, V = 4, 7, 16 device = torch.device("cuda") - logits = torch.randn(B, S, V, dtype=torch.float32, device=device) + logits = torch.randn(B, S, V, dtype=torch.bfloat16, device=device) labels = torch.randint(low=0, high=V, size=(B, S), device=device) # Randomly permute the batch dimension; a batch-invariant implementation should @@ -30,6 +30,7 @@ def test_selective_log_softmax_batch_invariant(): # Undo the permutation on the permuted outputs and compare elementwise. # If the kernel is batch invariant, each example's output should not depend # on its position in the batch. + assert bik_logps.dtype == torch.float32 assert torch.equal(bik_logps, bik_logps_perm[perm.argsort()]) From 44533eebfe98ea36cefd858fbfc3eeef7d90205a Mon Sep 17 00:00:00 2001 From: root Date: Tue, 28 Jul 2026 15:36:50 -0700 Subject: [PATCH 06/14] Drop batch-invariant Mamba prefix caching Signed-off-by: root --- .../inference/contexts/dynamic_context.py | 13 ++---- .../contexts/mamba_slot_allocator.py | 22 +-------- .../core/inference/engines/dynamic_engine.py | 16 +++---- .../contexts/test_dynamic_prefix_caching.py | 45 +++++-------------- 4 files changed, 26 insertions(+), 70 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 0d45e34967f..4824525b390 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -358,6 +358,10 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC self.mamba_chunk_size = mamba_inference_state_config.mamba_chunk_size if self.batch_invariant_mode: + assert not self.enable_prefix_caching, ( + "batch_invariant_mode does not support Mamba prefix caching; " + "set enable_prefix_caching=False." + ) assert self.num_speculative_tokens == 0, ( "batch_invariant_mode for Mamba dynamic inference only supports " "one-token decode; set num_speculative_tokens=0." @@ -2847,15 +2851,6 @@ def _find_mamba_match_count( mamba_map = self.mamba_slot_allocator.hash_to_block_id hashes = req.precomputed_block_hashes[start_block:end_block] for i in range(len(hashes) - 1, -1, -1): - block_count = start_block + i + 1 - if ( - self.batch_invariant_mode - and (block_count * self.block_size_tokens) % self.mamba_chunk_size - ): - # Restarting between Mamba chunk boundaries changes the - # subsequent reduction grouping. Recompute from the latest - # aligned cached state instead. - continue if hashes[i] in mamba_map: return i + 1 return 0 diff --git a/megatron/core/inference/contexts/mamba_slot_allocator.py b/megatron/core/inference/contexts/mamba_slot_allocator.py index ab8d1d67478..21116b0cb7e 100644 --- a/megatron/core/inference/contexts/mamba_slot_allocator.py +++ b/megatron/core/inference/contexts/mamba_slot_allocator.py @@ -1,6 +1,5 @@ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -import math from typing import TYPE_CHECKING, Dict import torch @@ -438,19 +437,10 @@ def compute_and_store_offsets( # in MambaMetadata (offset -> chunk-index conversion) to stay consistent. mamba_chunk_size = ctx.mamba_chunk_size - candidates = (kv_div_abs, last_aligned_abs, penultimate_abs) - if ctx.batch_invariant_mode: - # A reusable BIK state must be on both the KV-block and Mamba-chunk - # grids. Round candidates back so a short suffix is recomputed. - restore_stride = math.lcm(bs, mamba_chunk_size) - candidates = tuple( - position // restore_stride * restore_stride for position in candidates - ) - # Keep only boundaries that land inside this chunk's computed tokens and on # a mamba-chunk boundary (required for mid-sequence state extraction). offsets_set = set() - for abs_pos in candidates: + for abs_pos in (kv_div_abs, last_aligned_abs, penultimate_abs): offset = abs_pos - chunk_start if offset > 0 and offset < seq_len and offset % mamba_chunk_size == 0: offsets_set.add(offset) @@ -476,15 +466,7 @@ def compute_and_store_offsets( # cached directly. Only valid on the final chunk (otherwise the live state # is mid-prompt). Non-block-aligned prompts cache their last complete block # via the intermediate-extraction path above instead. - is_reusable_mamba_boundary = ( - not ctx.batch_invariant_mode or prompt_len % mamba_chunk_size == 0 - ) - if ( - is_last_chunk - and last_aligned_abs == prompt_len - and prompt_len > 0 - and is_reusable_mamba_boundary - ): + if is_last_chunk and last_aligned_abs == prompt_len and prompt_len > 0: last_block_idx = prompt_len // bs - 1 if last_block_idx >= 0: self._eos_cache_block_id_cpu[current_id] = ctx.request_to_kv_block_ids[current_id][ diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index d8fd6522b6a..98787ee913e 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -1608,24 +1608,24 @@ def get_prefix_coordination_metrics(self) -> dict: return {"waits": self._prefix_coordination_waits} def _mamba_batch_invariant_prefill_chunk_length( - self, req: DynamicInferenceRequest, capacity: int, prefix_skip: int = 0 + self, req: DynamicInferenceRequest, capacity: int ) -> int: - """Prefill span that computes an aligned Mamba chunk within ``capacity``. + """Raw prefill length that computes an aligned chunk within ``capacity``. Non-final calls must start and end at Mamba chunk boundaries. The final prompt call may be shorter because it seeds the decode replay tail. """ - remaining = len(req.remaining_prompt_tokens) - prefix_skip + remaining = len(req.remaining_prompt_tokens) if capacity >= remaining: - return prefix_skip + remaining + return remaining chunk_size = self.context.mamba_chunk_size computed_tokens = (capacity // chunk_size) * chunk_size if remaining - computed_tokens == 1: computed_tokens -= chunk_size if computed_tokens <= 0: - return prefix_skip - return prefix_skip + computed_tokens + return 0 + return computed_tokens def schedule_waiting_requests(self): """Tries to schedule any requests in the waiting pool.""" @@ -1875,9 +1875,9 @@ def schedule_chunked_prefill(self): if batch_invariant_mamba_prefill: prefill_chunk_length = self._mamba_batch_invariant_prefill_chunk_length( - req, computed_chunk, prefix_skip + req, computed_chunk ) - if prefill_chunk_length == prefix_skip: + if prefill_chunk_length == 0: can_schedule = False break else: diff --git a/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py b/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py index 78143a195a7..885779c0757 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py @@ -813,19 +813,6 @@ def test_mamba_cache_lifecycle(self): req6 = self._req(ctx6, p6.clone(), request_id=2) assert ctx6._find_mamba_match_count(req6, 0, len(req6.precomputed_block_hashes)) == 2 - # BIK restores only at Mamba chunk boundaries. A later unaligned - # cached state would change the reduction grouping after the restore. - ctx_bik = self._mctx(block_size_tokens=32, batch_invariant_mode=True) - p_bik = self._prompt(32 * 8) - ctx_bik.add_request(self._req(ctx_bik, p_bik.clone())) - self._mamba_allocate_and_register(ctx_bik, self._block_ids(ctx_bik, 0, 8)[:5]) - req_bik = self._req(ctx_bik, p_bik.clone(), request_id=2) - assert ( - ctx_bik._find_mamba_match_count(req_bik, 0, len(req_bik.precomputed_block_hashes)) == 4 - ) - *_, prefix_skip, effective_prefill = ctx_bik._compute_prefix_match(req_bik, len(p_bik)) - assert prefix_skip == 128 and effective_prefill == 128 - # no match when no mamba hashes registered ctx7 = self._mctx() ctx7.add_request(self._req(ctx7, self._prompt(bs * 3))) @@ -941,22 +928,21 @@ def test_mamba_prefill_skip_and_zero_prefill(self): @pytest.mark.internal def test_batch_invariant_mamba_chunked_prefill_scheduler_alignment(self): - ctx = self._mctx(block_size_tokens=32, batch_invariant_mode=True) + ctx = self._mctx( + block_size_tokens=32, batch_invariant_mode=True, enable_prefix_caching=False + ) engine = _StubEngine(ctx, enable_chunked_prefill=True) - req = self._req(ctx, self._prompt(500)) + req = self._req(ctx, self._prompt(500), enable_prefix_caching=False) assert engine._mamba_batch_invariant_prefill_chunk_length(req, 300) == 256 assert engine._mamba_batch_invariant_prefill_chunk_length(req, 100) == 0 - short_req = self._req(ctx, self._prompt(200), request_id=2) + short_req = self._req(ctx, self._prompt(200), request_id=2, enable_prefix_caching=False) assert engine._mamba_batch_invariant_prefill_chunk_length(short_req, 300) == 200 - assert engine._mamba_batch_invariant_prefill_chunk_length(req, 200, prefix_skip=128) == 256 - assert ( - engine._mamba_batch_invariant_prefill_chunk_length(short_req, 72, prefix_skip=128) - == 200 - ) - one_left_req = self._req(ctx, self._prompt(ctx.mamba_chunk_size + 1), request_id=3) + one_left_req = self._req( + ctx, self._prompt(ctx.mamba_chunk_size + 1), request_id=3, enable_prefix_caching=False + ) assert ( engine._mamba_batch_invariant_prefill_chunk_length(one_left_req, ctx.mamba_chunk_size) == 0 @@ -971,11 +957,15 @@ def test_batch_invariant_mamba_chunked_prefill_scheduler_alignment(self): with pytest.raises(AssertionError, match="max_tokens > mamba_chunk_size"): self._mctx( batch_invariant_mode=True, + enable_prefix_caching=False, enable_chunked_prefill=True, max_tokens=ctx.mamba_chunk_size, max_requests=64, ) + with pytest.raises(AssertionError, match="does not support Mamba prefix caching"): + self._mctx(batch_invariant_mode=True) + @pytest.mark.internal def test_mamba_intermediate_offsets(self): bs = 256 @@ -1037,17 +1027,6 @@ def test_mamba_intermediate_offsets(self): ctx3.transfer_bookkeeping_to_gpu() assert ctx3.mamba_slot_allocator._eos_cache_block_id_cpu[1].item() >= 0 - # BIK does not spend durable slots on block boundaries that cannot be - # restored without changing Mamba's chunk reduction grouping. - ctx_bik_unaligned = self._mctx(block_size_tokens=32, batch_invariant_mode=True) - ctx_bik_unaligned.add_request(self._req(ctx_bik_unaligned, self._prompt(160).clone())) - assert ctx_bik_unaligned.mamba_slot_allocator._eos_cache_block_id_cpu[0].item() < 0 - assert ctx_bik_unaligned.mamba_slot_allocator._intermediate_offsets_cpu[0, 0].item() == 128 - - ctx_bik_aligned = self._mctx(block_size_tokens=32, batch_invariant_mode=True) - ctx_bik_aligned.add_request(self._req(ctx_bik_aligned, self._prompt(128).clone())) - assert ctx_bik_aligned.mamba_slot_allocator._eos_cache_block_id_cpu[0].item() >= 0 - # intermediate output buffers are pre-allocated ctx4 = self._mctx() msa4 = ctx4.mamba_slot_allocator From 365547fad89940ae100aecf268b7b6fe18b03459 Mon Sep 17 00:00:00 2001 From: root Date: Tue, 28 Jul 2026 16:09:47 -0700 Subject: [PATCH 07/14] Isolate batch-invariant inference kernels Signed-off-by: root --- .../torch_symm_triton/__init__.py | 1 - .../torch_symm_triton/variable_collectives.py | 121 ------------ megatron/core/inference/moe/activations.py | 27 +-- .../core/inference/moe/batch_invariant.py | 172 ++++++++++++++++++ megatron/core/inference/moe/fused_moe.py | 4 +- megatron/core/ssm/mamba_mixer.py | 68 ++++--- .../core/ssm/ops/batch_invariant_decode.py | 91 ++++++++- megatron/core/ssm/ops/ssd_combined.py | 107 ----------- .../moe/token_dispatcher_inference.py | 10 +- .../test_moe_dispatching_and_routing.py | 16 +- .../unit_tests/inference/test_moe_permute.py | 3 +- 11 files changed, 311 insertions(+), 309 deletions(-) diff --git a/megatron/core/inference/communication/torch_symm_triton/__init__.py b/megatron/core/inference/communication/torch_symm_triton/__init__.py index 53523357567..75da02eaf4b 100644 --- a/megatron/core/inference/communication/torch_symm_triton/__init__.py +++ b/megatron/core/inference/communication/torch_symm_triton/__init__.py @@ -7,5 +7,4 @@ multimem_all_gather_v, multimem_all_gatherv_3tensor, multimem_reduce_scatter_v, - ordered_reduce_scatter_v, ) diff --git a/megatron/core/inference/communication/torch_symm_triton/variable_collectives.py b/megatron/core/inference/communication/torch_symm_triton/variable_collectives.py index dcf8ec4924d..a32b20b9a14 100644 --- a/megatron/core/inference/communication/torch_symm_triton/variable_collectives.py +++ b/megatron/core/inference/communication/torch_symm_triton/variable_collectives.py @@ -354,127 +354,6 @@ def multimem_reduce_scatter_v( return output_tensor -@triton.jit -def _ordered_reduce_scatter_v_kernel( - local_ptr, - buffer_ptrs_dev, - signal_pad_ptrs, - local_tokens, - rank_token_offset_ptr, - ep_max_tokens_ptr, - input_byte_offset, - HIDDEN_SIZE: tl.constexpr, - BLOCK_SIZE: tl.constexpr, - RANK: tl.constexpr, - WORLD_SIZE: tl.constexpr, -): - """Variable-count reduce-scatter with explicit rank-order fp32 addition. - - This is intentionally not a multimem.ld_reduce kernel. Each rank reads the - same token row from every peer symmetric buffer in rank order and accumulates - in fp32, which gives batch-invariant MoE a defined cross-rank reduction tree. - """ - pid = tl.program_id(axis=0) - - ep_max_tokens = tl.load(ep_max_tokens_ptr) - if pid >= ep_max_tokens: - return - - symm_mem_sync( - signal_pad_ptrs, - None, - RANK, - WORLD_SIZE, - hasPreviousMemAccess=False, - hasSubsequentMemAccess=True, - ) - sync_threads() - - tid = tl.arange(0, BLOCK_SIZE) - rank_token_offset = tl.load(rank_token_offset_ptr) - buffer_ptrs = buffer_ptrs_dev.to(tl.pointer_type(tl.uint64)) - - for token_offset in range(pid, local_tokens, tl.num_programs(axis=0)): - global_token = rank_token_offset + token_offset - - for channel_offset in range(0, HIDDEN_SIZE, BLOCK_SIZE): - offsets = channel_offset + tid - mask = offsets < HIDDEN_SIZE - acc = tl.zeros([BLOCK_SIZE], dtype=tl.float32) - - for src_rank in tl.range(0, WORLD_SIZE): - peer_base = tl.load(buffer_ptrs + src_rank).to(tl.pointer_type(tl.uint8)) - peer_ptr = (peer_base + input_byte_offset).to(tl.pointer_type(tl.float32)) - vals = tl.load( - peer_ptr + global_token * HIDDEN_SIZE + offsets, mask=mask, other=0.0 - ) - acc += vals - - tl.store(local_ptr + token_offset * HIDDEN_SIZE + offsets, acc, mask=mask) - - -def ordered_reduce_scatter_v( - output_tensor: torch.Tensor, - input_tensor: torch.Tensor, - symm_mem_hdl: _SymmetricMemory, - rank_token_offset: torch.Tensor, - ep_max_tokens: torch.Tensor, - per_rank_max_tokens: int, - input_byte_offset: int = 0, - **kwargs, -) -> torch.Tensor: - """Variable-count reduce-scatter with fixed rank-order fp32 accumulation. - - This is the batch-invariant alternative to multimem_reduce_scatter_v. It - uses symmetric memory for peer visibility, but performs no hardware FP - reduction; each output token is accumulated by explicitly loading peers in - rank order. - """ - assert HAVE_TRITON, "Triton is required for ordered_reduce_scatter_v." - assert ( - output_tensor.ndim == 2 and input_tensor.ndim == 2 - ), "output_tensor and input_tensor must be 2-D [tokens, hidden_size]." - assert is_device_nvls_capable( - output_tensor.device - ), "ordered_reduce_scatter_v requires a Hopper+ GPU with NVLink (SM >= 9)." - assert ( - output_tensor.dtype == input_tensor.dtype == torch.float32 - ), "ordered_reduce_scatter_v requires fp32 input and output tensors." - assert ( - rank_token_offset.numel() == 1 - and rank_token_offset.dtype == torch.int32 - and rank_token_offset.is_cuda - ), "rank_token_offset must be a scalar int32 CUDA tensor." - - hidden_size = output_tensor.shape[1] - assert ( - input_tensor.shape[1] == hidden_size - ), f"input and output hidden_size mismatch: {input_tensor.shape[1]} vs {hidden_size}" - - MAX_NUM_BLOCKS = kwargs.get("max_num_blocks", 128) - MAX_BLOCK_SIZE = 1024 - WARP_SIZE = 32 - block_size = min(triton.next_power_of_2(hidden_size), MAX_BLOCK_SIZE) - num_warps = max(1, block_size // WARP_SIZE) - num_blocks = min(per_rank_max_tokens, MAX_NUM_BLOCKS) - - _ordered_reduce_scatter_v_kernel[(num_blocks, 1, 1)]( - output_tensor, - symm_mem_hdl.buffer_ptrs_dev, - symm_mem_hdl.signal_pad_ptrs_dev, - local_tokens=output_tensor.shape[0], - rank_token_offset_ptr=rank_token_offset, - ep_max_tokens_ptr=ep_max_tokens, - input_byte_offset=input_byte_offset, - HIDDEN_SIZE=hidden_size, - BLOCK_SIZE=block_size, - RANK=symm_mem_hdl.rank, - WORLD_SIZE=symm_mem_hdl.world_size, - num_warps=num_warps, - ) - return output_tensor - - @triton.jit def _multimem_all_gatherv_3tensor_kernel( local_ptr_0, diff --git a/megatron/core/inference/moe/activations.py b/megatron/core/inference/moe/activations.py index d856586004d..ae5e4560ce3 100644 --- a/megatron/core/inference/moe/activations.py +++ b/megatron/core/inference/moe/activations.py @@ -35,12 +35,10 @@ def _squared_relu_kernel( output_ptr, src_idx_ptr, n_used_ptr, - probs_ptr, N, max_rows, # output_size (fixed for CG) BLOCK_N: tl.constexpr, NUM_BLOCKS: tl.constexpr, # grid size (fixed for CG) - APPLY_PROBS: tl.constexpr, ): """Squared ReLU that skips rows beyond n_used and alignment-padding rows (perm_map == -1). @@ -59,18 +57,11 @@ def _squared_relu_kernel( m = o < N x = tl.load(input_ptr + row * N + o, mask=m).to(tl.float32) r = tl.maximum(x, 0.0) - activated = (r * r).to(tl.bfloat16) - if APPLY_PROBS: - prob = tl.load(probs_ptr + row) - activated = (activated.to(tl.float32) * prob).to(tl.bfloat16) - tl.store(output_ptr + row * N + o, activated, mask=m) + tl.store(output_ptr + row * N + o, (r * r).to(tl.bfloat16), mask=m) def padded_squared_relu( - x: torch.Tensor, - permutation_map: torch.Tensor, - n_used: torch.Tensor, - probs: torch.Tensor | None = None, + x: torch.Tensor, permutation_map: torch.Tensor, n_used: torch.Tensor ) -> torch.Tensor: """Squared ReLU activation that skips rows beyond n_used and alignment-padding rows. @@ -78,25 +69,13 @@ def padded_squared_relu( x: [output_size, ffn_hidden] BF16 FC1 output. permutation_map: [output_size] int32, original token index or -1 for padding. n_used: scalar int32 CUDA tensor = inclusive_expert_offsets[-1]. - probs: optional FP32 router probabilities applied after the BF16 - activation rounding, matching the training path before FC2. """ M, N = x.shape out = torch.empty(M, N, dtype=x.dtype, device=x.device) BLOCK_N = min(triton.next_power_of_2(N), 1024) NUM_BLOCKS = min(M, 512) - probs_ptr = probs if probs is not None else permutation_map _squared_relu_kernel[(NUM_BLOCKS,)]( - x, - out, - permutation_map, - n_used, - probs_ptr, - N, - M, - BLOCK_N=BLOCK_N, - NUM_BLOCKS=NUM_BLOCKS, - APPLY_PROBS=probs is not None, + x, out, permutation_map, n_used, N, M, BLOCK_N=BLOCK_N, NUM_BLOCKS=NUM_BLOCKS ) return out diff --git a/megatron/core/inference/moe/batch_invariant.py b/megatron/core/inference/moe/batch_invariant.py index ca6514d755a..06dfddc2869 100644 --- a/megatron/core/inference/moe/batch_invariant.py +++ b/megatron/core/inference/moe/batch_invariant.py @@ -6,6 +6,11 @@ import torch +from megatron.core.inference.communication.torch_symm_triton.barrier import symm_mem_sync +from megatron.core.inference.communication.torch_symm_triton.utils import ( + is_device_nvls_capable, + sync_threads, +) from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( grouped_gemm_batch_invariant, grouped_gemm_batch_invariant_alignment, @@ -26,6 +31,11 @@ triton.jit = null_decorator tl = MagicMock() +try: + from torch._C._distributed_c10d import _SymmetricMemory +except ImportError: + _SymmetricMemory = MagicMock() + def enabled() -> bool: """Return whether global batch-invariant mode is active.""" @@ -44,6 +54,168 @@ def grouped_mm_alignment() -> int: return grouped_gemm_batch_invariant_alignment() +@triton.jit +def _squared_relu_with_probs_kernel( + input_ptr, + output_ptr, + permutation_map_ptr, + n_used_ptr, + probs_ptr, + hidden_size, + max_rows, + BLOCK_SIZE: tl.constexpr, + NUM_BLOCKS: tl.constexpr, +): + """Apply squared ReLU and router probabilities in training order.""" + pid = tl.program_id(0) + n_used = tl.load(n_used_ptr) + if pid >= n_used: + return + + for row in tl.range(pid, max_rows, NUM_BLOCKS): + if row < n_used: + if tl.load(permutation_map_ptr + row) >= 0: + prob = tl.load(probs_ptr + row) + for offset in tl.range(0, hidden_size, BLOCK_SIZE): + cols = offset + tl.arange(0, BLOCK_SIZE) + mask = cols < hidden_size + value = tl.load(input_ptr + row * hidden_size + cols, mask=mask).to(tl.float32) + value = tl.maximum(value, 0.0) + value = (value * value).to(tl.bfloat16) + value = (value.to(tl.float32) * prob).to(tl.bfloat16) + tl.store(output_ptr + row * hidden_size + cols, value, mask=mask) + + +def squared_relu_with_probs( + x: torch.Tensor, permutation_map: torch.Tensor, n_used: torch.Tensor, probs: torch.Tensor +) -> torch.Tensor: + """Match training's BF16 squared-ReLU rounding before the FP32 probability multiply.""" + num_rows, hidden_size = x.shape + out = torch.empty_like(x) + block_size = min(triton.next_power_of_2(hidden_size), 1024) + num_blocks = min(num_rows, 512) + _squared_relu_with_probs_kernel[(num_blocks,)]( + x, + out, + permutation_map, + n_used, + probs, + hidden_size, + num_rows, + BLOCK_SIZE=block_size, + NUM_BLOCKS=num_blocks, + ) + return out + + +@triton.jit +def _ordered_reduce_scatter_v_kernel( + local_ptr, + buffer_ptrs_dev, + signal_pad_ptrs, + local_tokens, + rank_token_offset_ptr, + ep_max_tokens_ptr, + input_byte_offset, + HIDDEN_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + RANK: tl.constexpr, + WORLD_SIZE: tl.constexpr, +): + """Reduce peer rows with an explicit rank-order FP32 sum.""" + pid = tl.program_id(axis=0) + + ep_max_tokens = tl.load(ep_max_tokens_ptr) + if pid >= ep_max_tokens: + return + + symm_mem_sync( + signal_pad_ptrs, + None, + RANK, + WORLD_SIZE, + hasPreviousMemAccess=False, + hasSubsequentMemAccess=True, + ) + sync_threads() + + tid = tl.arange(0, BLOCK_SIZE) + rank_token_offset = tl.load(rank_token_offset_ptr) + buffer_ptrs = buffer_ptrs_dev.to(tl.pointer_type(tl.uint64)) + + for token_offset in range(pid, local_tokens, tl.num_programs(axis=0)): + global_token = rank_token_offset + token_offset + + for channel_offset in range(0, HIDDEN_SIZE, BLOCK_SIZE): + offsets = channel_offset + tid + mask = offsets < HIDDEN_SIZE + acc = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + + for src_rank in tl.range(0, WORLD_SIZE): + peer_base = tl.load(buffer_ptrs + src_rank).to(tl.pointer_type(tl.uint8)) + peer_ptr = (peer_base + input_byte_offset).to(tl.pointer_type(tl.float32)) + values = tl.load( + peer_ptr + global_token * HIDDEN_SIZE + offsets, mask=mask, other=0.0 + ) + acc += values + + tl.store(local_ptr + token_offset * HIDDEN_SIZE + offsets, acc, mask=mask) + + +def ordered_reduce_scatter_v( + output_tensor: torch.Tensor, + input_tensor: torch.Tensor, + symm_mem_hdl: _SymmetricMemory, + rank_token_offset: torch.Tensor, + ep_max_tokens: torch.Tensor, + per_rank_max_tokens: int, + input_byte_offset: int = 0, + **kwargs, +) -> torch.Tensor: + """Reduce-scatter variable token rows with a fixed FP32 rank order.""" + assert HAVE_TRITON, "Triton is required for ordered_reduce_scatter_v." + assert ( + output_tensor.ndim == 2 and input_tensor.ndim == 2 + ), "output_tensor and input_tensor must be 2-D [tokens, hidden_size]." + assert is_device_nvls_capable( + output_tensor.device + ), "ordered_reduce_scatter_v requires a Hopper+ GPU with NVLink (SM >= 9)." + assert ( + output_tensor.dtype == input_tensor.dtype == torch.float32 + ), "ordered_reduce_scatter_v requires fp32 input and output tensors." + assert ( + rank_token_offset.numel() == 1 + and rank_token_offset.dtype == torch.int32 + and rank_token_offset.is_cuda + ), "rank_token_offset must be a scalar int32 CUDA tensor." + + hidden_size = output_tensor.shape[1] + assert ( + input_tensor.shape[1] == hidden_size + ), f"input and output hidden_size mismatch: {input_tensor.shape[1]} vs {hidden_size}" + + max_num_blocks = kwargs.get("max_num_blocks", 128) + block_size = min(triton.next_power_of_2(hidden_size), 1024) + num_warps = max(1, block_size // 32) + num_blocks = min(per_rank_max_tokens, max_num_blocks) + + _ordered_reduce_scatter_v_kernel[(num_blocks, 1, 1)]( + output_tensor, + symm_mem_hdl.buffer_ptrs_dev, + symm_mem_hdl.signal_pad_ptrs_dev, + local_tokens=output_tensor.shape[0], + rank_token_offset_ptr=rank_token_offset, + ep_max_tokens_ptr=ep_max_tokens, + input_byte_offset=input_byte_offset, + HIDDEN_SIZE=hidden_size, + BLOCK_SIZE=block_size, + RANK=symm_mem_hdl.rank, + WORLD_SIZE=symm_mem_hdl.world_size, + num_warps=num_warps, + ) + return output_tensor + + @triton.jit def _unpermute_tokens_in_expert_order_kernel( expert_out_ptr, # [output_size, hidden_dim] bf16 expert outputs diff --git a/megatron/core/inference/moe/fused_moe.py b/megatron/core/inference/moe/fused_moe.py index 17dd5ce69a1..3f49ba60e43 100644 --- a/megatron/core/inference/moe/fused_moe.py +++ b/megatron/core/inference/moe/fused_moe.py @@ -202,7 +202,9 @@ def mcore_fused_moe( n_used = offs[-1:] if batch_invariant_mode: # Match training: BF16 activation, FP32 probability multiply, then BF16 before FC2. - activation_out = activation_func(fc1_output, permutation_map, n_used, probs=permuted_probs) + activation_out = batch_invariant.squared_relu_with_probs( + fc1_output, permutation_map, n_used, permuted_probs + ) else: activation_out = activation_func(fc1_output, permutation_map, n_used) # Fused activation+quant returns MXFP8Tensor; otherwise quantize separately. diff --git a/megatron/core/ssm/mamba_mixer.py b/megatron/core/ssm/mamba_mixer.py index 31b30748e9e..773b26502da 100644 --- a/megatron/core/ssm/mamba_mixer.py +++ b/megatron/core/ssm/mamba_mixer.py @@ -1200,7 +1200,30 @@ def _ssm_decode( ) # Conv step - if causal_conv1d_update_triton is None: + if self.config.batch_invariant_mode: + # Match the causal-conv1d arithmetic used by the training forward. + assert ( + causal_conv1d_update_cuda is not None + ), "Batch-invariant Mamba decode requires causal-conv1d" + assert seq_len == 1, "Batch-invariant Mamba decode supports one token per request" + assert ( + intermediate_conv_state is None + ), "Batch-invariant Mamba decode does not support speculative decoding" + assert ( + batch_indices is not None and batch_indices.dtype == torch.int32 + ), "Batch-invariant Mamba decode requires int32 dynamic-batching indices" + + xBC_dtype = xBC.dtype + xBC = causal_conv1d_update_cuda( + xBC.to(conv_state.dtype).squeeze(1), + conv_state, + rearrange(self.conv1d_weight, "d 1 w -> d w").to(conv_state.dtype), + self.conv1d_bias.to(conv_state.dtype), + self.activation, + conv_state_indices=batch_indices, + ).unsqueeze(1) + xBC = xBC.to(xBC_dtype) + elif causal_conv1d_update_triton is None: # TODO(ksanthanam): Consider deprecating this path assert seq_len == 1, "Native PyTorch fallback only supports 1 token at a time" xBC_squeeze = xBC.squeeze(1) @@ -1216,40 +1239,15 @@ def _ssm_decode( # tensors to the conv state dtype for causal_conv1d_update and then cast xBC # back to the original dtype xBC_dtype = xBC.dtype - xBC = xBC.to(conv_state.dtype) - weight = rearrange(self.conv1d_weight, "d 1 w -> d w").to(conv_state.dtype) - bias = self.conv1d_bias.to(conv_state.dtype) - if self.config.batch_invariant_mode: - # Match the causal-conv1d arithmetic used by the training forward. - assert ( - causal_conv1d_update_cuda is not None - ), "Batch-invariant Mamba decode requires causal-conv1d" - assert seq_len == 1, "Batch-invariant Mamba decode supports one token per request" - assert ( - intermediate_conv_state is None - ), "Batch-invariant Mamba decode does not support speculative decoding" - assert ( - batch_indices is not None and batch_indices.dtype == torch.int32 - ), "Batch-invariant Mamba decode requires int32 dynamic-batching indices" - xBC = causal_conv1d_update_cuda( - xBC.squeeze(1), - conv_state, - weight, - bias, - self.activation, - conv_state_indices=batch_indices, - ).unsqueeze(1) - else: - xBC = causal_conv1d_update_triton( - xBC, - conv_state, - weight, - bias, - self.activation, - conv_state_indices=batch_indices, - intermediate_conv_states=intermediate_conv_state, - ) - xBC = xBC.to(xBC_dtype) + xBC = causal_conv1d_update_triton( + xBC.to(conv_state.dtype), + conv_state, + rearrange(self.conv1d_weight, "d 1 w -> d w").to(conv_state.dtype), + self.conv1d_bias.to(conv_state.dtype), + self.activation, + conv_state_indices=batch_indices, + intermediate_conv_states=intermediate_conv_state, + ).to(xBC_dtype) x, B, C = torch.split( xBC, diff --git a/megatron/core/ssm/ops/batch_invariant_decode.py b/megatron/core/ssm/ops/batch_invariant_decode.py index ba3e866db7c..fd532849d19 100644 --- a/megatron/core/ssm/ops/batch_invariant_decode.py +++ b/megatron/core/ssm/ops/batch_invariant_decode.py @@ -7,7 +7,10 @@ import triton import triton.language as tl -from megatron.core.ssm.ops.ssd_combined import mamba_chunk_scan_decode_rows +from megatron.core.ssm.ops.ssd_bmm import _bmm_chunk_fwd +from megatron.core.ssm.ops.ssd_chunk_scan import _chunk_scan_fwd +from megatron.core.ssm.ops.ssd_chunk_state import _chunk_cumsum_fwd, _chunk_state_fwd +from megatron.core.ssm.ops.ssd_state_passing import _state_passing_fwd @triton.jit @@ -53,6 +56,90 @@ def _masked_update_rows(states: torch.Tensor, indices: torch.Tensor, values: tor ) +def _mamba_chunk_scan_decode_rows( + x, + z, + dt, + A, + B, + C, + chunk_size, + chunk_starts, + slots, + target_rows, + chunk_flags, + initial_states, + out, + D=None, + dt_bias=None, + dt_softplus=False, + dt_limit=(0.0, float("inf")), +): + """Run the training scan pipeline over buffered decode chunks. + + Each kernel computes only the row or boundary consumed by this decode step, + while preserving the training kernel's arithmetic for that result. + """ + dA_cumsum, dt = _chunk_cumsum_fwd( + dt, + A, + chunk_size, + None, + dt_bias=dt_bias, + dt_softplus=dt_softplus, + dt_limit=dt_limit, + chunk_starts=chunk_starts, + target_rows=target_rows, + ) + states = _chunk_state_fwd( + B, + x, + dt, + dA_cumsum, + None, + states_in_fp32=True, + chunk_flags=chunk_flags, + chunk_starts=chunk_starts, + ) + CB = _bmm_chunk_fwd( + C, + B, + chunk_size, + None, + output_dtype=torch.float32, + target_rows=target_rows, + chunk_starts=chunk_starts, + ) + # Scan before state passing because both read the incoming live state and + # state passing overwrites crossing slots with the outgoing boundary state. + _chunk_scan_fwd( + CB, + x, + dt, + dA_cumsum, + C, + states, + None, + out, + slots, + D=D, + z=z, + initial_states=initial_states, + target_rows=target_rows, + chunk_starts=chunk_starts, + ) + _state_passing_fwd( + states.flatten(-2), + dA_cumsum, + None, + initial_states=initial_states.flatten(-2), + seq_idx=slots, + dst_states=initial_states.flatten(-2), + dst_indices=slots, + dst_flags=chunk_flags, + ) + + @dataclass class BatchInvariantDecodeBuffers: """Per-slot persistent state for the buffered decode scan.""" @@ -201,7 +288,7 @@ def batch_invariant_decode_buffered_scan( # Run the gated pipeline over the buffers and ssm_state in place. State # passing writes crossing slots' boundary states straight into # ssm_state, so no scatter is needed afterwards. - mamba_chunk_scan_decode_rows( + _mamba_chunk_scan_decode_rows( buffers.x.view(-1, nheads, headdim), buffers.z.view(-1, nheads, headdim), buffers.dt.view(-1, nheads), diff --git a/megatron/core/ssm/ops/ssd_combined.py b/megatron/core/ssm/ops/ssd_combined.py index a685d79a4fc..4fcee98b13e 100644 --- a/megatron/core/ssm/ops/ssd_combined.py +++ b/megatron/core/ssm/ops/ssd_combined.py @@ -160,113 +160,6 @@ def _mamba_chunk_scan_combined_fwd( return final_states -def mamba_chunk_scan_decode_rows( - x, - z, - dt, - A, - B, - C, - chunk_size, - chunk_starts, - slots, - target_rows, - chunk_flags, - initial_states, - out, - D=None, - dt_bias=None, - dt_softplus=False, - dt_limit=(0.0, float("inf")), -): - """Row-gated chunk scan for batch-invariant single-token decode. - - Same 5-kernel pipeline as the full varlen scan, run directly over the - persistent per-slot buffers: chunk c is the fixed chunk_size window at - chunk_starts[c], and every chunk is its own sequence starting from - initial_states[slots[c]]. The kernels are gated to what a decode step - actually consumes: bmm and the scan compute only the block containing - target_rows[c], and the chunk-state matmul runs only where chunk_flags - is set (the slot crosses its boundary, the one step its state is read). - The blocks that do run execute the same instructions as the ungated - kernels, so the outputs match a full scan bitwise. - - Args: - x/z/dt/B/C: flattened persistent buffers, (num_rows * chunk_size, ...). - chunk_starts: (nseq,) int32, window offset per chunk - (slot * chunk_size for per-slot buffers). - slots: (nseq,) int32, live-cache row containing each chunk's incoming - state. May repeat for padding entries. - target_rows: (nseq,) int32, the only output row read per chunk. - chunk_flags: (nseq,), nonzero where the slot crosses its boundary. - initial_states: (num_states, nheads, headdim, dstate), the engine's - live SSM cache. Crossing chunks update it in place. - out: (nseq, nheads, headdim), receives each chunk's target row. - """ - dA_cumsum, dt = _chunk_cumsum_fwd( - dt, - A, - chunk_size, - None, - dt_bias=dt_bias, - dt_softplus=dt_softplus, - dt_limit=dt_limit, - chunk_starts=chunk_starts, - target_rows=target_rows, - ) - # Only boundary-crossing chunks produce a state; state passing masks the rest. - states = _chunk_state_fwd( - B, - x, - dt, - dA_cumsum, - None, - states_in_fp32=True, - chunk_flags=chunk_flags, - chunk_starts=chunk_starts, - ) - CB = _bmm_chunk_fwd( - C, - B, - chunk_size, - None, - output_dtype=torch.float32, - target_rows=target_rows, - chunk_starts=chunk_starts, - ) - # The scan must run before state passing: the snapshot below overwrites - # crossing slots' rows in ssm_state, and the scan reads initial_states - # from that same cache. The scan never reads state passing's output in - # decode mode (every chunk is its own sequence), so `states` is just a - # shape-valid placeholder for the unused carried-state pointer. - _chunk_scan_fwd( - CB, - x, - dt, - dA_cumsum, - C, - states, - None, - out, - slots, - D=D, - z=z, - initial_states=initial_states, - target_rows=target_rows, - chunk_starts=chunk_starts, - ) - _state_passing_fwd( - states.flatten(-2), - dA_cumsum, - None, - initial_states=initial_states.flatten(-2), - seq_idx=slots, - dst_states=initial_states.flatten(-2), - dst_indices=slots, - dst_flags=chunk_flags, - ) - - def mamba_chunk_scan_combined_varlen( x, dt, diff --git a/megatron/core/transformer/moe/token_dispatcher_inference.py b/megatron/core/transformer/moe/token_dispatcher_inference.py index d3b9e99425e..1a2ae10b24a 100644 --- a/megatron/core/transformer/moe/token_dispatcher_inference.py +++ b/megatron/core/transformer/moe/token_dispatcher_inference.py @@ -31,9 +31,8 @@ from megatron.core.inference.communication.torch_symm_triton import ( multimem_all_gatherv_3tensor, multimem_reduce_scatter_v, - ordered_reduce_scatter_v, ) -from megatron.core.inference.moe import InferenceGroupedGemmBackend +from megatron.core.inference.moe import InferenceGroupedGemmBackend, batch_invariant from megatron.core.inference.moe.metadata import fused_metadata_update from megatron.core.inference.symmetric_memory import SymmetricMemoryManager from megatron.core.process_groups_config import ProcessGroupCollection @@ -41,9 +40,6 @@ gather_from_sequence_parallel_region, reduce_scatter_to_sequence_parallel_region, ) -from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( - is_batch_invariant_mode_enabled, -) from megatron.core.transformer.moe.inference_routing_mask_kernel import mask_routing_padding from megatron.core.transformer.moe.shared_experts import SharedExpertMLP from megatron.core.transformer.moe.token_dispatcher import MoEAllGatherTokenDispatcher @@ -646,8 +642,8 @@ def token_combine(self, hidden_states): device=hidden_states.device, ) reduce_scatter_v = ( - ordered_reduce_scatter_v - if is_batch_invariant_mode_enabled() + batch_invariant.ordered_reduce_scatter_v + if batch_invariant.enabled() else multimem_reduce_scatter_v ) reduce_scatter_v( 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 43ecb16971c..3d5353ab000 100644 --- a/tests/unit_tests/inference/test_moe_dispatching_and_routing.py +++ b/tests/unit_tests/inference/test_moe_dispatching_and_routing.py @@ -464,10 +464,10 @@ def test_cuda_graph_batch_invariant_combine_uses_ordered_symmetric_memory(self, The graph path still writes local partials into the symmetric RSV buffer, but the combine must not use multimem.ld_reduce in batch-invariant mode. """ + from megatron.core.inference.moe import batch_invariant from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( set_batch_invariant_mode, ) - from megatron.core.transformer.moe import token_dispatcher_inference if Utils.world_size < 2: pytest.skip("Ordered RSV combine requires expert-parallel world_size > 1.") @@ -503,16 +503,14 @@ def test_cuda_graph_batch_invariant_combine_uses_ordered_symmetric_memory(self, static_routing_map = global_routing_map[start:end].contiguous() ordered_calls = {"value": 0} - orig_ordered_reduce_scatter_v = token_dispatcher_inference.ordered_reduce_scatter_v + orig_ordered_reduce_scatter_v = batch_invariant.ordered_reduce_scatter_v def _tracked_ordered_reduce_scatter_v(*args, **kwargs): ordered_calls["value"] += 1 return orig_ordered_reduce_scatter_v(*args, **kwargs) monkeypatch.setattr( - token_dispatcher_inference, - "ordered_reduce_scatter_v", - _tracked_ordered_reduce_scatter_v, + batch_invariant, "ordered_reduce_scatter_v", _tracked_ordered_reduce_scatter_v ) with torch.no_grad(), set_batch_invariant_mode(True): @@ -548,13 +546,13 @@ def test_cuda_graph_batch_invariant_moe_layer_uses_ordered_rsv(self, monkeypatch writes deterministic local partials into the symmetric RSV buffer, then token_combine uses explicit rank-order fp32 loads. """ + from megatron.core.inference.moe import batch_invariant 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 from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( HAVE_DEEPGEMM_BF16, set_batch_invariant_mode, ) - from megatron.core.transformer.moe import token_dispatcher_inference from megatron.core.transformer.moe.token_dispatcher_inference import ( NVLSAllGatherVDispatcher, ) @@ -605,16 +603,14 @@ def _tracked_get_rsv_tensor(cls): NVLSAllGatherVDispatcher, "_get_rsv_tensor", classmethod(_tracked_get_rsv_tensor) ) ordered_calls = {"value": 0} - orig_ordered_reduce_scatter_v = token_dispatcher_inference.ordered_reduce_scatter_v + orig_ordered_reduce_scatter_v = batch_invariant.ordered_reduce_scatter_v def _tracked_ordered_reduce_scatter_v(*args, **kwargs): ordered_calls["value"] += 1 return orig_ordered_reduce_scatter_v(*args, **kwargs) monkeypatch.setattr( - token_dispatcher_inference, - "ordered_reduce_scatter_v", - _tracked_ordered_reduce_scatter_v, + batch_invariant, "ordered_reduce_scatter_v", _tracked_ordered_reduce_scatter_v ) local_tokens = 16 diff --git a/tests/unit_tests/inference/test_moe_permute.py b/tests/unit_tests/inference/test_moe_permute.py index 3d90916ad4b..8be6aec7f59 100644 --- a/tests/unit_tests/inference/test_moe_permute.py +++ b/tests/unit_tests/inference/test_moe_permute.py @@ -53,6 +53,7 @@ def test_batch_invariant_squared_relu_applies_probs_before_fc2(): """Match training's probability placement and BF16 rounding before FC2.""" from megatron.core.activations import squared_relu from megatron.core.inference.moe.activations import padded_squared_relu + from megatron.core.inference.moe.batch_invariant import squared_relu_with_probs torch.manual_seed(17) rows, hidden, output_size = 37, 1856, 512 @@ -62,7 +63,7 @@ def test_batch_invariant_squared_relu_applies_probs_before_fc2(): permutation_map = torch.arange(rows, device="cuda", dtype=torch.int32) unweighted = padded_squared_relu(x, permutation_map, _vt(rows)) - actual = padded_squared_relu(x, permutation_map, _vt(rows), probs=probs) + actual = squared_relu_with_probs(x, permutation_map, _vt(rows), probs) expected_unweighted = squared_relu(x) expected = (squared_relu(x) * probs.unsqueeze(1)).to(torch.bfloat16) From 75230ec5ef2cd6c1d0379a4c5c4cad3c6198d480 Mon Sep 17 00:00:00 2001 From: root Date: Wed, 29 Jul 2026 05:10:52 -0700 Subject: [PATCH 08/14] Adapt Mamba BIK tests to raw-state scan API Signed-off-by: root --- tests/unit_tests/ssm/ops/test_batch_invariant_decode.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit_tests/ssm/ops/test_batch_invariant_decode.py b/tests/unit_tests/ssm/ops/test_batch_invariant_decode.py index 901cc08894a..904a6ab9e22 100644 --- a/tests/unit_tests/ssm/ops/test_batch_invariant_decode.py +++ b/tests/unit_tests/ssm/ops/test_batch_invariant_decode.py @@ -181,7 +181,7 @@ def _varlen_boundary_state_from_prefill(self, x, dt, B, C, prefill_len, initial_ out = torch.zeros_like(x[0, :prefill_len]) seq_idx = torch.zeros(len(chunk_boundaries) - 1, dtype=torch.int32, device=self.device) - chunk_states = mamba_chunk_scan_combined_varlen( + _, chunk_states = mamba_chunk_scan_combined_varlen( x=x[0, :prefill_len], dt=dt[0, :prefill_len], A=self.A, @@ -196,7 +196,7 @@ def _varlen_boundary_state_from_prefill(self, x, dt, B, C, prefill_len, initial_ z=None, dt_bias=self.dt_bias, initial_states=initial_states, - return_intermediate_states=True, + return_raw_states=True, dt_softplus=True, dt_limit=(0.0, float("inf")), state_dtype=torch.float32, From 356eaea1b82e8f20c08135ca68bceb757a499e2f Mon Sep 17 00:00:00 2001 From: root Date: Wed, 29 Jul 2026 06:19:15 -0700 Subject: [PATCH 09/14] lint Signed-off-by: root --- tests/unit_tests/models/test_gpt_model_batch_invariant.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit_tests/models/test_gpt_model_batch_invariant.py b/tests/unit_tests/models/test_gpt_model_batch_invariant.py index 1e93687bcd7..d3a97a2c4db 100644 --- a/tests/unit_tests/models/test_gpt_model_batch_invariant.py +++ b/tests/unit_tests/models/test_gpt_model_batch_invariant.py @@ -30,7 +30,9 @@ from tests.unit_tests.test_utilities import Utils try: - from flash_attn_3.flash_attn_interface import _flash_attn_forward + from flash_attn_3.flash_attn_interface import ( + _flash_attn_forward, + ) from flash_attn_3.flash_attn_interface import ( flash_attn_with_kvcache as flash_attn3_with_kvcache, ) From 0ed0d9a6fc70bcbbbb8173cc9d61205a7b81ba3a Mon Sep 17 00:00:00 2001 From: root Date: Thu, 30 Jul 2026 05:51:03 -0700 Subject: [PATCH 10/14] Update lockfile for batch-invariant dependencies Signed-off-by: root --- uv.lock | 29 ++++++++++++++++++++--------- 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/uv.lock b/uv.lock index 15036ceb4ce..61feedc44e4 100644 --- a/uv.lock +++ b/uv.lock @@ -1012,6 +1012,11 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/05/7f/798705f5296a58ca505d600456748d1be48078eac8a7050d8a98bc9edb89/decorator-5.3.1-py3-none-any.whl", hash = "sha256:f47fe6fdbd2edd623ecfe36875d37aba411624e2670dd395dddae1358689bb3c", size = 10365, upload-time = "2026-05-18T06:03:26.517Z" }, ] +[[package]] +name = "deep-gemm" +version = "2.5.0+local" +source = { git = "https://github.com/deepseek-ai/DeepGEMM.git?rev=714dd1a4a980f7937a74343d19a8eba4fe321480#714dd1a4a980f7937a74343d19a8eba4fe321480" } + [[package]] name = "defusedxml" version = "0.7.1" @@ -2220,6 +2225,9 @@ training = [ ] [package.dev-dependencies] +batch-invariant = [ + { name = "deep-gemm" }, +] build = [ { name = "cython" }, { name = "hatchling" }, @@ -2251,6 +2259,7 @@ linting = [ { name = "ruff" }, ] no-pypi-wheels = [ + { name = "deep-gemm" }, { name = "emerging-optimizers" }, { name = "flash-mla" }, ] @@ -2320,6 +2329,7 @@ requires-dist = [ provides-extras = ["training", "mlm", "dev", "lts", "te", "ssm"] [package.metadata.requires-dev] +batch-invariant = [{ name = "deep-gemm", git = "https://github.com/deepseek-ai/DeepGEMM.git?rev=714dd1a4a980f7937a74343d19a8eba4fe321480" }] build = [ { name = "cython", specifier = ">=3.0.0" }, { name = "hatchling" }, @@ -2351,6 +2361,7 @@ linting = [ { name = "ruff", specifier = "~=0.9.0" }, ] no-pypi-wheels = [ + { name = "deep-gemm", git = "https://github.com/deepseek-ai/DeepGEMM.git?rev=714dd1a4a980f7937a74343d19a8eba4fe321480" }, { name = "emerging-optimizers", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.2.0" }, { name = "flash-mla", git = "https://github.com/deepseek-ai/FlashMLA?rev=nv_dev" }, ] @@ -5099,15 +5110,15 @@ name = "torch" version = "2.13.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cuda-bindings", marker = "python_version < '0'" }, - { name = "filelock" }, - { name = "fsspec" }, - { name = "jinja2" }, - { name = "networkx" }, - { name = "setuptools" }, - { name = "sympy" }, - { name = "triton" }, - { name = "typing-extensions" }, + { name = "cuda-bindings", marker = "python_full_version < '3.15' and sys_platform == 'linux'" }, + { name = "filelock", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "fsspec", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "jinja2", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "networkx", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "setuptools", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "sympy", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "triton", marker = "sys_platform == 'never'" }, + { name = "typing-extensions", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] [[package]] From 7aee3ebddcae990dec09581b4a08789560b03041 Mon Sep 17 00:00:00 2001 From: root Date: Thu, 30 Jul 2026 05:54:41 -0700 Subject: [PATCH 11/14] Fix batch-invariant test formatting Signed-off-by: root --- tests/unit_tests/models/test_gpt_model_batch_invariant.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/unit_tests/models/test_gpt_model_batch_invariant.py b/tests/unit_tests/models/test_gpt_model_batch_invariant.py index d3a97a2c4db..1e93687bcd7 100644 --- a/tests/unit_tests/models/test_gpt_model_batch_invariant.py +++ b/tests/unit_tests/models/test_gpt_model_batch_invariant.py @@ -30,9 +30,7 @@ from tests.unit_tests.test_utilities import Utils try: - from flash_attn_3.flash_attn_interface import ( - _flash_attn_forward, - ) + from flash_attn_3.flash_attn_interface import _flash_attn_forward from flash_attn_3.flash_attn_interface import ( flash_attn_with_kvcache as flash_attn3_with_kvcache, ) From 8e5e1f584e669b3726920bd2f8bbd6cd1ff67a91 Mon Sep 17 00:00:00 2001 From: root Date: Thu, 30 Jul 2026 12:56:44 -0700 Subject: [PATCH 12/14] Fix hybrid MoE batch-invariant test config Signed-off-by: root --- tests/unit_tests/inference/test_hybrid_moe.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/unit_tests/inference/test_hybrid_moe.py b/tests/unit_tests/inference/test_hybrid_moe.py index ac14ff0c61e..b23cc0f9660 100644 --- a/tests/unit_tests/inference/test_hybrid_moe.py +++ b/tests/unit_tests/inference/test_hybrid_moe.py @@ -31,10 +31,12 @@ from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.core.ssm.mamba_mixer import _check_mamba_sequence_packing_support from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.attention import HAVE_FA4 from megatron.core.transformer.cuda_graphs import _CudagraphGlobalRecord, delete_cuda_graphs from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( te_supports_batch_invariant_attention, ) +from megatron.core.transformer.enums import AttnBackend from megatron.core.transformer.moe.token_dispatcher_inference import NVLSAllGatherVDispatcher from megatron.core.utils import is_fa_min_version from tests.unit_tests.inference.test_moe_dispatching_and_routing import ( @@ -282,6 +284,9 @@ def test_batch_invariant_prefill_matches_full_forward(self): config = _make_base_config( num_layers=3, batch_invariant_mode=True, + attention_backend=AttnBackend.flash, + attention_dropout=0.0, + flash_attention_version=4 if HAVE_FA4 else 3, inference_grouped_gemm_backend="torch", inference_moe_token_dispatcher_type="nvls", ) From 1a48fcb2069cd4987a78c37f7741ac5da4484c9f Mon Sep 17 00:00:00 2001 From: root Date: Thu, 30 Jul 2026 14:26:09 -0700 Subject: [PATCH 13/14] Fix hybrid MoE FlashAttention test setup Signed-off-by: root --- tests/unit_tests/inference/test_hybrid_moe.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/unit_tests/inference/test_hybrid_moe.py b/tests/unit_tests/inference/test_hybrid_moe.py index b23cc0f9660..90335933ed2 100644 --- a/tests/unit_tests/inference/test_hybrid_moe.py +++ b/tests/unit_tests/inference/test_hybrid_moe.py @@ -43,7 +43,7 @@ NANOV3_BASE, _make_base_config, ) -from tests.unit_tests.test_utilities import Utils +from tests.unit_tests.test_utilities import Utils, clear_nvte_env_vars # Request state constants for parametrized tests. NONE = "none" # 0 requests (dummy rank) @@ -281,6 +281,7 @@ def test_batch_invariant_prefill_matches_full_forward(self): ) model_parallel_cuda_manual_seed(123, inference_rng_tracker=True, force_reset_rng=True) + clear_nvte_env_vars() config = _make_base_config( num_layers=3, batch_invariant_mode=True, From 641d239aed94e843578aaf408dc73c936422e392 Mon Sep 17 00:00:00 2001 From: root Date: Fri, 31 Jul 2026 07:14:40 -0700 Subject: [PATCH 14/14] Skip BIK attention test without FA3 or FA4 Signed-off-by: root --- tests/unit_tests/inference/test_hybrid_moe.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/unit_tests/inference/test_hybrid_moe.py b/tests/unit_tests/inference/test_hybrid_moe.py index 90335933ed2..9587cc2d285 100644 --- a/tests/unit_tests/inference/test_hybrid_moe.py +++ b/tests/unit_tests/inference/test_hybrid_moe.py @@ -31,7 +31,7 @@ from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.core.ssm.mamba_mixer import _check_mamba_sequence_packing_support from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed -from megatron.core.transformer.attention import HAVE_FA4 +from megatron.core.transformer.attention import HAVE_FA3, HAVE_FA4 from megatron.core.transformer.cuda_graphs import _CudagraphGlobalRecord, delete_cuda_graphs from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( te_supports_batch_invariant_attention, @@ -63,8 +63,11 @@ # independently. _EP_SIZE = 4 requires_te_batch_invariant_attention = pytest.mark.skipif( - not te_supports_batch_invariant_attention(), - reason="Batch-invariant attention requires TransformerEngine PR #3204 or >= 2.18.", + not te_supports_batch_invariant_attention() or not (HAVE_FA3 or HAVE_FA4), + reason=( + "Batch-invariant attention requires TransformerEngine PR #3204 or >= 2.18 " + "and FlashAttention-3 or -4." + ), ) # Combinatorial sweep: unordered combinations with repetition of ALL_STATES