diff --git a/megatron/core/inference/config.py b/megatron/core/inference/config.py index 46d2dbca1b0..6b3e715d531 100644 --- a/megatron/core/inference/config.py +++ b/megatron/core/inference/config.py @@ -38,7 +38,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.""" @@ -61,7 +61,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/attention_context/mamba_metadata.py b/megatron/core/inference/contexts/attention_context/mamba_metadata.py index 9984b2dd71a..045ede4b502 100644 --- a/megatron/core/inference/contexts/attention_context/mamba_metadata.py +++ b/megatron/core/inference/contexts/attention_context/mamba_metadata.py @@ -21,6 +21,7 @@ def __init__( max_intermediate_count: int, mamba_chunk_size: int = 128, d_conv: int = 0, + decode_indices_dtype: torch.dtype = torch.int64, ): """ Initializes the Mamba slot allocator. @@ -36,12 +37,15 @@ def __init__( mamba_chunk_size (int): The chunk size used by the Mamba SSM Triton kernels. d_conv (int): Convolution window size (from mamba_conv_states_shape[-1]). Used for vectorized conv state extraction at intermediate offsets. + decode_indices_dtype (torch.dtype): Dtype for decode state-slot indices. """ self.max_requests = max_requests self.max_tokens = max_tokens self.mamba_chunk_size = mamba_chunk_size self.d_conv = d_conv self.device = torch.cuda.current_device() + assert decode_indices_dtype in (torch.int32, torch.int64) + self.decode_indices_dtype = decode_indices_dtype # Maximum possible chunks across all batch configurations self.max_chunks = max_tokens // mamba_chunk_size + max_requests @@ -52,9 +56,10 @@ def __init__( ) # Map from requests to slots in the static Mamba state buffer for active decode requests. - # int64 so selective_state_update can index directly without a per-layer upcast kernel; + # Non-BIK decode uses int64 for selective_state_update; BIK uses int32 + # for the exact causal-conv1d update kernel. self._batch_indices_decode_buffer = torch.full( - (self.max_requests,), -1, dtype=torch.int64, device=self.device + (self.max_requests,), -1, dtype=self.decode_indices_dtype, device=self.device ) # Map from requests to slots in the static Mamba state buffer for active prefill requests diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 102a5e5e55b..56f935e6cc0 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -315,6 +315,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 < " @@ -357,6 +358,20 @@ 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 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." + ) + 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 # layer type. @@ -722,6 +737,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: @@ -877,6 +899,7 @@ def _allocate_mamba_states(self): max_intermediate_count=self.max_mamba_intermediate_states_per_step, mamba_chunk_size=self.mamba_chunk_size, d_conv=self.mamba_conv_states_shape[-1], + decode_indices_dtype=self._mamba_decode_indices_dtype, ) # Bind the unified CPU/GPU buffers so the per-step Mamba metadata # fields ride along with the single coalesced H2D in @@ -1081,12 +1104,18 @@ def initialize_all_tensors(self) -> None: ) # Mamba section (hybrid models only). Must match the MambaMetadata # shapes (mirrors the layout documented in ContextGPUView). - # batch_indices_decode is int64; all other fields are int32. + # batch_indices_decode is int32 in batch-invariant mode and int64 otherwise; + # all other fields are int32. if self.is_hybrid_model: - # mamba_batch_indices_decode is int64; pad to 8-byte alignment. - _mamba_align_pad = (8 - _pre_mamba_bytes % 8) % 8 + self._mamba_decode_indices_dtype = ( + torch.int32 if self.batch_invariant_mode else torch.int64 + ) + _decode_index_bytes = 4 if self.batch_invariant_mode else 8 + _mamba_align_pad = ( + _decode_index_bytes - _pre_mamba_bytes % _decode_index_bytes + ) % _decode_index_bytes self._max_mamba_chunks = self.max_tokens // self.mamba_chunk_size + self.max_requests - _mamba_batch_indices_decode_bytes = self.max_requests * 8 + _mamba_batch_indices_decode_bytes = self.max_requests * _decode_index_bytes _mamba_batch_indices_prefill_bytes = self.max_requests * 4 _mamba_seq_idx_bytes = self.max_tokens * 4 _mamba_cu_seqlens_bytes = (self.max_requests + 1) * 4 @@ -1250,7 +1279,7 @@ def initialize_all_tensors(self) -> None: _off += _mamba_align_pad self._cpu_mamba_batch_indices_decode = self._cpu_bookkeeping_buf[ _off : _off + _mamba_batch_indices_decode_bytes - ].view(torch.int64) + ].view(self._mamba_decode_indices_dtype) _off += _mamba_batch_indices_decode_bytes self._cpu_mamba_batch_indices_prefill = self._cpu_bookkeeping_buf[ _off : _off + _mamba_batch_indices_prefill_bytes @@ -1297,6 +1326,9 @@ def initialize_all_tensors(self) -> None: max_kv_blocks=self.max_kv_block_count, device=torch.cuda.current_device(), max_mamba_chunks=self._max_mamba_chunks, + mamba_decode_indices_dtype=( + self._mamba_decode_indices_dtype if self.is_hybrid_model else torch.int64 + ), ) self._bookkeeping_h2d_done_event = torch.cuda.Event() diff --git a/megatron/core/inference/contexts/gpu_view.py b/megatron/core/inference/contexts/gpu_view.py index 2066375d19e..d92205f137f 100644 --- a/megatron/core/inference/contexts/gpu_view.py +++ b/megatron/core/inference/contexts/gpu_view.py @@ -31,7 +31,9 @@ def __init__( max_kv_blocks: int, device: torch.device, max_mamba_chunks: int = 0, + mamba_decode_indices_dtype: torch.dtype = torch.int64, ): + assert mamba_decode_indices_dtype in (torch.int32, torch.int64) # Field layout (must match DynamicInferenceContext's CPU buffer layout): # int64 token fields first (auto 8-byte alignment), then int32 token # fields, then int32 request fields, then int32 MHA fields, then @@ -63,7 +65,7 @@ def __init__( mha_block_table_bytes = max_bs * max_kv_blocks * 4 # Mamba section, only present for hybrid models. - # mamba_batch_indices_decode int64 (max_bs,) + # mamba_batch_indices_decode int32 or int64 (max_bs,) # mamba_batch_indices_prefill int32 (max_bs,) # mamba_seq_idx int32 (1, max_tokens) # mamba_cu_seqlens int32 (max_bs + 1,) @@ -86,9 +88,11 @@ def __init__( ) if max_mamba_chunks > 0: - # mamba_batch_indices_decode is int64; pad to 8-byte alignment. - mamba_align_pad = (8 - pre_mamba_bytes % 8) % 8 - mamba_batch_indices_decode_bytes = max_bs * 8 + decode_index_bytes = 4 if mamba_decode_indices_dtype == torch.int32 else 8 + mamba_align_pad = ( + decode_index_bytes - pre_mamba_bytes % decode_index_bytes + ) % decode_index_bytes + mamba_batch_indices_decode_bytes = max_bs * decode_index_bytes mamba_batch_indices_prefill_bytes = max_bs * 4 mamba_seq_idx_bytes = max_tokens * 4 mamba_cu_seqlens_bytes = (max_bs + 1) * 4 @@ -202,7 +206,7 @@ def __init__( off += mamba_align_pad self.mamba_batch_indices_decode = self._buf[ off : off + mamba_batch_indices_decode_bytes - ].view(torch.int64) + ].view(mamba_decode_indices_dtype) off += mamba_batch_indices_decode_bytes self.mamba_batch_indices_prefill = self._buf[ off : off + mamba_batch_indices_prefill_bytes diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 833b65dd15f..770613fdd85 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -1631,6 +1631,26 @@ 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 + ) -> 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) -> None: """Try to schedule requests from the waiting pool.""" # Keep track of which requests get scheduled. @@ -1882,6 +1902,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: @@ -1934,7 +1957,15 @@ def schedule_chunked_prefill(self): else: computed_chunk = computed_budget - prefill_chunk_length = prefix_skip + computed_chunk + if batch_invariant_mamba_prefill: + prefill_chunk_length = self._mamba_batch_invariant_prefill_chunk_length( + req, computed_chunk + ) + if prefill_chunk_length == 0: + can_schedule = False + break + else: + prefill_chunk_length = prefix_skip + computed_chunk # Mamba prefix caching: keep chunk boundaries block-aligned. # compute_and_store_offsets() records a recurrent-state snapshot at a @@ -1970,7 +2001,7 @@ def schedule_chunked_prefill(self): # See https://github.com/Dao-AILab/flash-attention/issues/1537 # The -1 is safe after CG snapping: is_applicable_for_batch_dim matches on # cg.token_count >= real.token_count, so the snapped CG still covers token_count-1. - if remaining_len - prefill_chunk_length == 1: + if not batch_invariant_mamba_prefill and remaining_len - prefill_chunk_length == 1: if computed_chunk > 1: prefill_chunk_length -= 1 else: diff --git a/megatron/core/inference/moe/batch_invariant.py b/megatron/core/inference/moe/batch_invariant.py new file mode 100644 index 00000000000..06dfddc2869 --- /dev/null +++ b/megatron/core/inference/moe/batch_invariant.py @@ -0,0 +1,275 @@ +# 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.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, + 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() + +try: + from torch._C._distributed_c10d import _SymmetricMemory +except ImportError: + _SymmetricMemory = 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 _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 + 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: + vals = tl.load(expert_out_ptr + pos * hidden_dim + offsets, mask=mask_h).to( + tl.float32 + ) + acc += vals + tl.store(output_ptr + tok * hidden_dim + offsets, acc, mask=mask_h) + + +def unpermute_tokens_in_expert_order( + expert_output: 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, + 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..3f49ba60e43 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 @@ -183,7 +200,13 @@ def mcore_fused_moe( # number of rows actually used by experts this iteration (valid tokens + alignment # padding within expert blocks). Passed to activation and unpermute to skip unused rows. n_used = offs[-1:] - activation_out = activation_func(fc1_output, permutation_map, n_used) + if batch_invariant_mode: + # Match training: BF16 activation, FP32 probability multiply, then BF16 before FC2. + 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. if use_mxfp8 and not isinstance(activation_out, MXFP8Tensor): activation_out = MXFP8Tensor.from_bf16(activation_out, backend="triton") @@ -191,5 +214,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, + None if batch_invariant_mode else 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..f65ac5c4200 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,22 @@ 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 +379,7 @@ def permute_tokens( permuted_hidden, permuted_probs, permutation_map, + inverse_map_ptr, exclusive_expert_offsets, valid_tokens, hidden_dim, @@ -364,7 +389,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 @@ -434,12 +468,13 @@ def _unpermute_tokens_kernel( def unpermute_tokens( expert_output: torch.Tensor, - permuted_probs: torch.Tensor, + permuted_probs: Optional[torch.Tensor], permutation_map: torch.Tensor, num_tokens: int, 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. @@ -448,7 +483,8 @@ def unpermute_tokens( Args: expert_output: [output_size, hidden_dim] expert outputs in permuted order. - permuted_probs: [output_size] fp32 routing probabilities. + permuted_probs: [output_size] fp32 routing probabilities, or None when + batch-invariant inference applied them before FC2. permutation_map: [output_size] int32, original token index or -1 for padding. num_tokens: max token count (output buffer height); always fixed for CG. n_used: scalar int32 CUDA tensor = inclusive_expert_offsets[-1]. Rows @@ -460,10 +496,25 @@ def unpermute_tokens( Pass a symmetric memory tensor to scatter directly into it, avoiding a separate copy before RSV. If None, a local buffer is allocated. """ - assert ( - 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, batch_invariant_inverse_map, valid_tokens, out + ) + + assert ( + permuted_probs is not None and permuted_probs.dtype == torch.float32 + ), "permuted_probs must be fp32" 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/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py index 9fc94365045..53522dd8b2b 100644 --- a/megatron/core/models/common/language_module/language_module.py +++ b/megatron/core/models/common/language_module/language_module.py @@ -110,6 +110,12 @@ def _set_attention_backend(self): Transformer engine works based on optout. By default all three attention backend flags are set to 1. So if the user choses a particular attention backend we set the other two to 0. If the user choses local, we set all 3 TE env variables to 0. """ + if self.config.batch_invariant_mode: + from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( + assert_te_supports_batch_invariant_attention, + ) + + assert_te_supports_batch_invariant_attention() def check_and_set_env_variable( env_variable_name: str, expected_value: int, attn_type: AttnBackend diff --git a/megatron/core/ssm/mamba_mixer.py b/megatron/core/ssm/mamba_mixer.py index 7a44c8a493a..f6ae07dd230 100644 --- a/megatron/core/ssm/mamba_mixer.py +++ b/megatron/core/ssm/mamba_mixer.py @@ -25,6 +25,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.intermediate_extraction import ( scatter_intermediate_conv, @@ -60,10 +61,12 @@ try: from causal_conv1d import causal_conv1d_fn + from causal_conv1d import causal_conv1d_update as causal_conv1d_update_cuda from causal_conv1d.causal_conv1d_varlen import causal_conv1d_varlen_states except ImportError: causal_conv1d_fn = None + causal_conv1d_update_cuda = None try: from mamba_ssm.ops.triton.layernorm_gated import RMSNorm as RMSNormGated @@ -488,6 +491,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: @@ -988,12 +995,25 @@ def _ssm_prefill( chunk_starts = cu_chunk_seqlens[:-1] seq_idx_for_varlen = seq_idx[0, chunk_starts].contiguous() + # Batch-invariant decode replays the partial prefill tail, so keep + # the cached SSM state at the last complete chunk boundary. + 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) + # Extraction is enabled when the slot allocator wired buffers in via # the caller. When enabled, the chunk scan returns its raw states so # our Triton kernels do a fused gather+conditional-scatter directly, # skipping the dense intermediate tensor and the padded-slot writes. extract_intermediates = ( - intermediate_chunk_indices is not None and intermediate_ssm_out is not None + not self.config.batch_invariant_mode + and intermediate_chunk_indices is not None + and intermediate_ssm_out is not None ) ssm_varlen_result = mamba_chunk_scan_combined_varlen( x=x, @@ -1011,16 +1031,16 @@ def _ssm_prefill( if self.D_has_hdim else self.cp.get_D() ), - z=z if not self.rmsnorm else None, + z=z if (self.config.batch_invariant_mode or not self.rmsnorm) else None, dt_bias=self.cp.get_dt_bias().float(), initial_states=initial_ssm_state, - return_raw_states=extract_intermediates, + return_raw_states=self.config.batch_invariant_mode or extract_intermediates, dt_softplus=True, dt_limit=(0.0, float("inf")), state_dtype=ssm_state.dtype, ) - if extract_intermediates: + if self.config.batch_invariant_mode or extract_intermediates: ssm_varlen_states, raw_ssm_states = ssm_varlen_result else: ssm_varlen_states = ssm_varlen_result @@ -1029,7 +1049,26 @@ def _ssm_prefill( y = y.unsqueeze(0) z = z.unsqueeze(0) - tensor_masked_update(ssm_state, batch_indices, ssm_varlen_states) + if self.config.batch_invariant_mode: + boundary_mask = has_boundary.view(-1, 1, 1, 1) + cache_states = torch.where( + boundary_mask, raw_ssm_states[boundary_chunk_indices], initial_ssm_state + ) + else: + cache_states = ssm_varlen_states + + tensor_masked_update(ssm_state, batch_indices, cache_states) + if self.config.batch_invariant_mode: + self._get_batch_invariant_decoder().seed( + x, + z.squeeze(0), + dt, + B, + C, + cu_seqlens, + batch_indices, + max_requests=ssm_state.shape[0], + ) if extract_intermediates: # Fused gather+conditional-scatter for SSM: read row @@ -1090,7 +1129,7 @@ def _ssm_prefill( if self.rmsnorm: z = rearrange(z, "b l h p -> l b (h p)").contiguous() z = self.cp.post_conv_ssm(z) - y = self.norm(y, z) + y = self.norm(y, None if self.config.batch_invariant_mode else z) return y @@ -1112,6 +1151,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 _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) + return self._batch_invariant_decoder + def train(self, mode: bool = True): """Mark the decode cache stale; weights may have updated.""" if mode: @@ -1161,7 +1206,30 @@ def _ssm_decode( ) # Conv step - if causal_conv1d_update 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 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) @@ -1197,7 +1265,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 @@ -1286,7 +1359,7 @@ def _ssm_decode( y = rearrange(y, "b s h p -> b s (h p)") if self.rmsnorm: - y = self.norm(y, z) + y = self.norm(y, None if self.config.batch_invariant_mode else z) return y 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..fd532849d19 --- /dev/null +++ b/megatron/core/ssm/ops/batch_invariant_decode.py @@ -0,0 +1,369 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +"""Batch-invariant Mamba decode using buffered chunk replay.""" + +from dataclasses import dataclass + +import torch +import triton +import triton.language as tl + +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 +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, + ) + + +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.""" + + x: torch.Tensor # (max_requests, chunk_size, nheads, headdim) + z: torch.Tensor # (max_requests, chunk_size, nheads, headdim) + dt: torch.Tensor # (max_requests, chunk_size, nheads) + B: torch.Tensor # (max_requests, chunk_size, ngroups, dstate) + C: torch.Tensor # (max_requests, 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_requests,) int32 + # Per-entry target-row output, allocated once and sliced per step. + out: torch.Tensor # (max_requests, nheads, headdim) + target_rows: torch.Tensor # (max_requests,) int32 + chunk_flags: torch.Tensor # (max_requests,) int32 + + @classmethod + def allocate( + cls, + max_requests: 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.""" + return cls( + x=torch.zeros(max_requests, chunk_size, nheads, headdim, device=device, dtype=dtype), + z=torch.zeros(max_requests, chunk_size, nheads, headdim, device=device, dtype=dtype), + dt=torch.zeros(max_requests, chunk_size, nheads, device=device, dtype=dtype), + B=torch.zeros(max_requests, chunk_size, ngroups, dstate, device=device, dtype=dtype), + C=torch.zeros(max_requests, chunk_size, ngroups, dstate, device=device, dtype=dtype), + num_buffered=torch.zeros(max_requests, device=device, dtype=torch.int32), + out=torch.empty(max_requests, nheads, headdim, device=device, dtype=dtype), + target_rows=torch.empty(max_requests, device=device, dtype=torch.int32), + chunk_flags=torch.empty(max_requests, device=device, dtype=torch.int32), + ) + + def seed( + self, + x: torch.Tensor, + z: 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 + + # 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 + ) + + slots = batch_indices[:num_seqs] + _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) + ) + + +def batch_invariant_decode_buffered_scan( + buffers: BatchInvariantDecodeBuffers, + x: torch.Tensor, # (decode_batch_size, 1, nheads, headdim) + z: 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)." + ) + 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 " + f"({output_capacity}); increase max_requests." + ) + + out = buffers.out[:decode_batch_size] + target_rows = buffers.target_rows[:decode_batch_size] + chunk_flags = buffers.chunk_flags[:decode_batch_size] + + 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).to(torch.int32) + + _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)) + chunk_flags.copy_(crossed.to(torch.int32)) + out.zero_() + + # 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.z.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=batch_indices * chunk_size, + slots=batch_indices, + target_rows=target_rows, + chunk_flags=chunk_flags, + initial_states=ssm_state, + out=out, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + ) + + next_write_pos = torch.where(crossed, 0, write_pos + 1).to(torch.int32) + _masked_update_rows( + buffers.num_buffered.unsqueeze(1), batch_indices, next_write_pos.unsqueeze(1) + ) + + return out.unsqueeze(1) + + +class MambaBatchInvariantDecode: + """Adapter between a MambaMixer and the buffered decode.""" + + def __init__(self, mixer): + # Training applies z inside the chunk scan before RMSNormGated, so + # decode buffers and replays z through that same kernel path. + assert mixer.rmsnorm, "batch_invariant_mode requires rmsnorm=True" + self.mixer = mixer + self.buffers: BatchInvariantDecodeBuffers | None = None + + def _get_buffers(self, max_requests, x, B) -> BatchInvariantDecodeBuffers: + if self.buffers is None: + nheads, headdim = x.shape[-2:] + ngroups, dstate = B.shape[-2:] + self.buffers = BatchInvariantDecodeBuffers.allocate( + max_requests, + self.mixer.chunk_size, + nheads, + headdim, + ngroups, + dstate, + x.device, + x.dtype, + ) + return self.buffers + + def seed(self, x, z, dt, B, C, cu_seqlens, batch_indices, max_requests) -> None: + """Seed replay buffers from the prefill tail.""" + buffers = self._get_buffers(max_requests, x, B) + buffers.seed(x, z, dt, B, C, cu_seqlens, batch_indices) + + def step(self, x, z, 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) + z = z.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, z, 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..1b3a819ee14 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,22 @@ 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 tr < 0: + return + 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 +155,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 +173,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 +197,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 +215,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 +232,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..8066965b12b 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,20 @@ 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. + target_row = tl.load(target_rows_ptr + pid_c) + if target_row < 0: + return + if pid_m != target_row // 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 +172,30 @@ 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. 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 +328,26 @@ 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 +363,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 +421,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 +464,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..8a21de1bef7 100644 --- a/megatron/core/ssm/ops/ssd_chunk_state.py +++ b/megatron/core/ssm/ops/ssd_chunk_state.py @@ -51,11 +51,13 @@ def _chunk_cumsum_fwd_kernel( dt_bias_ptr, dt_out_ptr, dA_cumsum_ptr, - cu_chunk_seqlens_ptr, + chunk_offsets_ptr, + target_rows_ptr, # Matrix dimension seqlen, nheads: tl.constexpr, chunk_size: tl.constexpr, + HAS_TARGET_ROWS: tl.constexpr, dt_min: tl.constexpr, dt_max: tl.constexpr, # Strides @@ -79,9 +81,16 @@ def _chunk_cumsum_fwd_kernel( # https://github.com/triton-lang/triton/issues/1058 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) + if HAS_TARGET_ROWS: + if tl.load(target_rows_ptr + pid_c) < 0: + return + + 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) dt_ptr += chunk_seqlen_start * stride_dt_seqlen dt_out_ptr += pid_c * stride_dt_out_chunk @@ -179,7 +188,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 +214,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 +224,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 +296,23 @@ def _chunk_cumsum_fwd( dt_bias=None, dt_softplus=False, dt_limit=(0.0, float("inf")), + chunk_starts=None, + target_rows=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 + 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 + 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 +323,8 @@ 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, + target_rows_ptr=target_rows, seqlen=seqlen, nheads=nheads, chunk_size=chunk_size, @@ -309,6 +340,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_TARGET_ROWS=has_target_rows, DT_SOFTPLUS=dt_softplus, HAS_DT_BIAS=dt_bias is not None, BLOCK_SIZE_CHUNK=triton.next_power_of_2(chunk_size), @@ -316,7 +348,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 +395,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 +418,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_state_passing.py b/megatron/core/ssm/ops/ssd_state_passing.py index 65b81a0ec31..9c12537d7f1 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) @@ -66,7 +73,7 @@ def _state_passing_fwd_kernel( states_ptrs = states_ptr + offs_m * stride_states_dim out_ptrs = out_ptr + offs_m * stride_out_dim - if HAS_INITSTATES: + if HAS_INITSTATES and not HAS_DST_STATES: initstates_ptrs = ( initstates_ptr + pid_h * stride_initstates_head + offs_m * stride_initstates_dim ) @@ -75,27 +82,46 @@ def _state_passing_fwd_kernel( else: states = tl.zeros((BLOCK_SIZE,), dtype=tl.float32) - prev_seq_idx = 0 + prev_seq_idx = tl.full((), 0, tl.int64) for c in range(nchunks): - new_states = tl.load(states_ptrs, mask=offs_m < dim, 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) - # we have started a new sequence - if prev_seq_idx != seq_idx: - if HAS_INITSTATES: - initstates_ptrs = ( - initstates_ptr - + seq_idx * stride_initstates_batch - + pid_h * stride_initstates_head - + offs_m * stride_initstates_dim - ) - states = tl.load(initstates_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 + if dst_flag: + new_states = tl.load(states_ptrs, mask=offs_m < dim, 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).to(tl.int64) + if HAS_DST_STATES: + # Destination chunks start from their indexed initial state. + is_new_seq = True else: - states = tl.zeros((BLOCK_SIZE,), dtype=tl.float32) + is_new_seq = prev_seq_idx != seq_idx + if is_new_seq: + if HAS_INITSTATES: + initstates_ptrs = ( + initstates_ptr + + seq_idx * stride_initstates_batch + + pid_h * stride_initstates_head + + offs_m * stride_initstates_dim + ) + states = tl.load(initstates_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) + else: + states = tl.zeros((BLOCK_SIZE,), dtype=tl.float32) - prev_seq_idx = seq_idx - states = tl.exp(dA_cs) * states + new_states - tl.store(out_ptrs, states, mask=offs_m < dim) + prev_seq_idx = seq_idx + states = tl.exp(dA_cs) * states + new_states + if HAS_DST_STATES: + 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) + else: + tl.store(out_ptrs, states, mask=offs_m < dim) states_ptrs += stride_states_chunk dA_cs_ptr += stride_dA_cs_chunk @@ -103,20 +129,44 @@ 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) + else: + out = states 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 +177,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, @@ -144,6 +197,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 2adefc58634..87ba3023d2a 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( @@ -413,11 +420,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: @@ -540,7 +550,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 a83e298eee2..abb21f36edb 100644 --- a/megatron/core/transformer/custom_layers/batch_invariant_kernels.py +++ b/megatron/core/transformer/custom_layers/batch_invariant_kernels.py @@ -12,6 +12,7 @@ from typing import Any, Dict, List, Optional import torch +from packaging.version import Version try: import triton @@ -28,15 +29,60 @@ 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", + "assert_te_supports_batch_invariant_attention", + "te_supports_batch_invariant_attention", + "HAVE_DEEPGEMM_BF16", ] _LOGGER = logging.getLogger(__name__) +_TE_BATCH_INVARIANT_COMMIT = "cb4a45fd" +_TE_BATCH_INVARIANT_MIN_VERSION = Version("2.18") + + +def te_supports_batch_invariant_attention() -> bool: + """Return whether TE supports explicit FlashAttention version selection.""" + import transformer_engine + + te_version = Version(transformer_engine.__version__) + te_revision = te_version.local or "" + return te_version >= _TE_BATCH_INVARIANT_MIN_VERSION or te_revision.startswith( + _TE_BATCH_INVARIANT_COMMIT + ) + + +def assert_te_supports_batch_invariant_attention() -> None: + """Require TE's explicit FlashAttention version selection.""" + import transformer_engine + + te_version = Version(transformer_engine.__version__) + assert te_supports_batch_invariant_attention(), ( + "Batch-invariant attention requires TransformerEngine PR #3204 " + f"({_TE_BATCH_INVARIANT_COMMIT}) or TransformerEngine >= " + f"{_TE_BATCH_INVARIANT_MIN_VERSION}; found {te_version}." + ) def _matmul_launch_metadata( @@ -312,7 +358,7 @@ def log_softmax(input: torch.Tensor, dim: int = -1) -> torch.Tensor: Args: input: Input tensor dim: Dimension along which to compute log_softmax (only -1 or last dim supported) - >> Stashed changes + Returns: Tensor with log_softmax applied along the specified dimension """ @@ -478,13 +524,52 @@ def mean_dim( return output +# Production uses DeepGEMM for bf16 and the deterministic Triton kernel for +# intentional higher-precision operations such as the fp32 MoE router. +_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" + and a.dtype == torch.bfloat16 + and b.dtype == torch.bfloat16 + ): + 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" + and a.dtype == torch.bfloat16 + and b.dtype == torch.bfloat16 + ): + out = _mm_deepgemm(a, b) + if bias is not None: + out = out + bias + return out return matmul_persistent(a, b, bias=bias) @@ -525,6 +610,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] = {} _TE_APPLY_NORM_ORIGS: Dict[str, Any] = {} @@ -622,6 +708,10 @@ 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() + # 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 @@ -651,6 +741,234 @@ def _patched(*args, **kwargs): mod.apply_normalization = _te_apply_normalization_patched +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 _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, tuple)) or not isinstance(B, (list, tuple)): + 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. 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, ...) + # 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): + raise RuntimeError( + "Batch-invariant grouped GEMM requires unquantized BF16 tensor sequences " + "with GELU fusion disabled." + ) + + # 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) + raise RuntimeError( + "Unsupported batch-invariant grouped GEMM call: " + f"layout={layout!r}, single_output={single_output}, grad={grad}." + ) + + +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.""" global _TE_GENERAL_GEMM_ORIG, _TE_RMSNORM_ORIG_FWD, _MEG_TE_GENERAL_GEMM_ORIG @@ -731,6 +1049,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. @@ -780,7 +1101,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: @@ -792,7 +1113,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]) @@ -1052,6 +1373,252 @@ 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. @@ -1069,11 +1636,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") @@ -1083,6 +1669,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(): @@ -1094,33 +1684,133 @@ 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 + 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, + 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 + 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, + 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..9baf66dab65 --- /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 9b7cf177c79..dfdb9a14460 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, +) +from megatron.core.transformer.moe.batch_invariant import unpermute as batch_invariant_unpermute from megatron.core.transformer.moe.moe_logging import get_moe_metrics_tracker from megatron.core.transformer.moe.router_replay import RouterReplay from megatron.core.transformer.transformer_config import TransformerConfig @@ -343,6 +350,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], @@ -375,6 +383,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[ @@ -387,6 +397,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.") @@ -421,6 +435,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 @@ -447,6 +462,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() @@ -461,10 +477,21 @@ 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( @@ -476,6 +503,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 @@ -501,10 +529,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.") @@ -519,6 +555,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 @@ -820,7 +869,12 @@ def _compute_topk( ) else: # Sorting top-k turned off during inference - return torch.topk(scores, k=topk, dim=1, sorted=torch.is_grad_enabled()) + return torch.topk( + scores, + k=topk, + dim=1, + sorted=torch.is_grad_enabled() or is_batch_invariant_mode_enabled(), + ) def compute_topk(scores, topk, num_groups=None, group_topk=None): # Default behavior if no replay is active diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py index e4591ce3acf..55bbe96f5f3 100644 --- a/megatron/core/transformer/moe/router.py +++ b/megatron/core/transformer/moe/router.py @@ -7,6 +7,9 @@ from megatron.core.inference.utils import InferenceMode from megatron.core.jit import jit_fuser +from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( + is_batch_invariant_mode_enabled, +) from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.moe.moe_logging import get_moe_metrics_tracker from megatron.core.transformer.moe.moe_utils import ( @@ -953,7 +956,12 @@ def _forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = N if self.qb_beta is not None: precomputed_indices = (logits - self.qb_beta).topk(self.topk, dim=1).indices - probs, top_indices = self._compiled_topk_routing( + routing = ( + topk_routing_with_score_function + if is_batch_invariant_mode_enabled() + else self._compiled_topk_routing + ) + probs, top_indices = routing( logits, self.topk, use_pre_softmax=self.config.moe_router_pre_softmax, diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py index 5743a047960..256f38cc6fb 100644 --- a/megatron/core/transformer/moe/token_dispatcher.py +++ b/megatron/core/transformer/moe/token_dispatcher.py @@ -482,6 +482,9 @@ def __init__( 'hidden_shape', 'probs', ] + 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 @@ -660,7 +663,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, @@ -669,6 +672,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 @@ -888,6 +892,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 b08d88f2641..1a2ae10b24a 100644 --- a/megatron/core/transformer/moe/token_dispatcher_inference.py +++ b/megatron/core/transformer/moe/token_dispatcher_inference.py @@ -32,7 +32,7 @@ multimem_all_gatherv_3tensor, multimem_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 @@ -617,9 +617,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. @@ -637,7 +641,12 @@ def token_combine(self, hidden_states): dtype=rsv["tensor"].dtype, device=hidden_states.device, ) - multimem_reduce_scatter_v( + reduce_scatter_v = ( + batch_invariant.ordered_reduce_scatter_v + if batch_invariant.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 4f9de546161..4a8a6fe1431 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1528,6 +1528,20 @@ 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.") @@ -2873,6 +2887,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)" @@ -2893,6 +2911,35 @@ def _scope_to_str(s): assert ( self.attention_dropout == 0.0 ), "Batch invariant mode does not support attention dropout" + 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/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index d4424d17f74..387336472f7 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -739,6 +739,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 @@ -1690,9 +1693,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/pyproject.toml b/pyproject.toml index e19fe870334..849259f4d38 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -174,7 +174,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] managed = true @@ -185,6 +186,7 @@ no-build-isolation-package = [ "mamba-ssm", "transformer-engine", "transformer-engine-torch", + "deep_gemm", "fast-hadamard-transform", ] link-mode = "copy" @@ -229,6 +231,7 @@ requires-dist = ["torch", "packaging", "ninja"] flash_mla = [ { git = "https://github.com/deepseek-ai/FlashMLA", rev = "nv_dev" }, ] +deep_gemm = { git = "https://github.com/deepseek-ai/DeepGEMM.git", rev = "714dd1a4a980f7937a74343d19a8eba4fe321480" } transformer-engine = { git = "https://github.com/NVIDIA/TransformerEngine.git", rev = "e7c550c5f80636cf841a8204b1d6f85a5f3f28b7" } nemo-run = { git = "https://github.com/NVIDIA-NeMo/Run.git", rev = "ddd40a8f24847f5c919f911d0240bd622653612f" } emerging_optimizers = { git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git", rev = "v0.2.0" } diff --git a/tests/unit_tests/inference/contexts/attention_metadata/test_mamba_metadata.py b/tests/unit_tests/inference/contexts/attention_metadata/test_mamba_metadata.py index 99bb046d97d..43a7d57d8ba 100644 --- a/tests/unit_tests/inference/contexts/attention_metadata/test_mamba_metadata.py +++ b/tests/unit_tests/inference/contexts/attention_metadata/test_mamba_metadata.py @@ -31,6 +31,15 @@ def metadata_context(self): yield metadata metadata.reset() + @pytest.mark.internal + @pytest.mark.parametrize("dtype", [torch.int32, torch.int64]) + def test_decode_indices_dtype(self, dtype): + metadata = MambaMetadata( + max_requests=4, max_tokens=16, max_intermediate_count=1, decode_indices_dtype=dtype + ) + + assert metadata._batch_indices_decode_buffer.dtype == dtype + def _run_update_test( self, metadata: MambaMetadata, 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 fa166af9669..1c1d48ff3fb 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 @@ -61,13 +62,15 @@ def _ctx( 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 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, @@ -75,7 +78,14 @@ 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, + 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, @@ -87,6 +97,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, ) @@ -914,6 +925,46 @@ 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, enable_prefix_caching=False + ) + engine = _StubEngine(ctx, enable_chunked_prefill=True) + 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, enable_prefix_caching=False) + 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, enable_prefix_caching=False + ) + 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_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 diff --git a/tests/unit_tests/inference/contexts/test_gpu_view.py b/tests/unit_tests/inference/contexts/test_gpu_view.py index 63a838f563e..1bd48aeaa21 100644 --- a/tests/unit_tests/inference/contexts/test_gpu_view.py +++ b/tests/unit_tests/inference/contexts/test_gpu_view.py @@ -93,3 +93,20 @@ def test_layout_with_and_without_mamba(self, max_mamba_chunks): for name in MAMBA_VIEWS_INT32: assert getattr(v, name) is not None assert getattr(v, name).dtype == torch.int32 + + @pytest.mark.parametrize("dtype", [torch.int32, torch.int64]) + def test_mamba_decode_indices_dtype(self, dtype): + """The runtime-selected decode dtype must not change the remaining layout.""" + v = ContextGPUView( + max_requests=MAX_REQUESTS, + max_tokens=MAX_TOKENS, + max_kv_blocks=MAX_KV_BLOCKS, + device=torch.device("cuda"), + max_mamba_chunks=MAX_MAMBA_CHUNKS, + mamba_decode_indices_dtype=dtype, + ) + + assert v.mamba_batch_indices_decode.dtype == dtype + assert v.mamba_batch_indices_decode.shape == (MAX_REQUESTS,) + assert v.mamba_batch_indices_prefill.dtype == torch.int32 + assert v.mamba_conv_seq_start.shape == (MAX_TOKENS,) diff --git a/tests/unit_tests/inference/test_hybrid_moe.py b/tests/unit_tests/inference/test_hybrid_moe.py index a8cc9b743e6..9587cc2d285 100644 --- a/tests/unit_tests/inference/test_hybrid_moe.py +++ b/tests/unit_tests/inference/test_hybrid_moe.py @@ -31,14 +31,19 @@ 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_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, +) +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 ( 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) @@ -57,6 +62,13 @@ # ranks form data-parallel replicas, each running the same EP combo # independently. _EP_SIZE = 4 +requires_te_batch_invariant_attention = pytest.mark.skipif( + 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 # across the EP ranks. Since rank assignment is symmetric (shuffling ranks @@ -261,6 +273,72 @@ def _assert_cuda_graphs_were_replayed(expect_replayed, rank, label): class TestDynamicInferenceNVLS(_TestDynamicInferenceBase): """NVLS dispatcher: combinatorial sweep of EP request states.""" + @requires_te_batch_invariant_attention + @torch.inference_mode() + def test_batch_invariant_prefill_matches_full_forward(self): + """Dynamic prefill should exactly match the full-sequence forward.""" + from megatron.core.inference.inference_request import DynamicInferenceRequest + from megatron.core.inference.sampling_params import SamplingParams + from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( + set_batch_invariant_mode, + ) + + 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, + 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", + ) + model = HybridModel( + config=config, + hybrid_stack_spec=hybrid_inference_stack_spec, + vocab_size=self.VOCAB_SIZE, + max_sequence_length=self.MAX_SEQ_LEN, + hybrid_layer_pattern="ME*", + ).cuda() + model.eval() + + input_ids = torch.arange(64, device="cuda", dtype=torch.long).unsqueeze(0) + with set_batch_invariant_mode(True): + InferenceMode.unset_active() + full_logits = model( + input_ids=input_ids, + position_ids=None, + attention_mask=None, + runtime_gather_output=True, + ) + + ctx = self._build_context( + model, + num_cuda_graphs=0, + use_cuda_graphs_for_non_decode_steps=False, + max_requests=4, + max_tokens=128, + ) + request = DynamicInferenceRequest( + request_id=0, + prompt_tokens=input_ids.cpu().squeeze(0), + sampling_params=SamplingParams(num_tokens_to_generate=1, termination_id=-1), + ) + ctx.add_request(request) + ctx.initialize_attention_state() + + InferenceMode.set_active() + inference_logits = model( + input_ids=input_ids, + position_ids=None, + attention_mask=None, + inference_context=ctx, + runtime_gather_output=True, + ) + + torch.testing.assert_close(inference_logits, full_logits, atol=0, rtol=0) + # ------------------------------------------------------------------ # test_ep_state_cross_product: combinatorial sweep with mixed CUDA graphs # ------------------------------------------------------------------ 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 5b21ab4c364..3d5353ab000 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, + flash_attention_version=4, + attention_dropout=0.0, + 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 @@ -440,6 +458,257 @@ 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.inference.moe import batch_invariant + from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( + set_batch_invariant_mode, + ) + + 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 = 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( + batch_invariant, "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.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.token_dispatcher_inference import ( + NVLSAllGatherVDispatcher, + ) + + 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, + flash_attention_version=4, + attention_dropout=0.0, + 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 = 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( + batch_invariant, "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 + + def test_batch_invariant_moe_matches_training(self): + """The NVLS inference MoE path should exactly match training AllToAll.""" + from megatron.core.models.gpt.moe_module_specs import get_inference_optimized_moe_spec + from megatron.core.parallel_state import get_expert_model_parallel_group + 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, + ) + + if Utils.world_size < 2: + pytest.skip("Training-to-inference MoE parity requires expert parallelism.") + if Utils.world_size & (Utils.world_size - 1): + pytest.skip("NVLS Triton symmetric-memory barrier requires power-of-two EP size.") + if not HAVE_DEEPGEMM_BF16: + pytest.skip("Batch-invariant torch MoE path requires DeepGEMM bf16 grouped kernels.") + + torch.manual_seed(2028) + torch.cuda.manual_seed(2028) + + config = _make_base_config( + expert_model_parallel_size=Utils.world_size, + inference_grouped_gemm_backend="torch", + inference_moe_token_dispatcher_type="nvls", + batch_invariant_mode=True, + attention_backend=AttnBackend.flash, + flash_attention_version=4, + attention_dropout=0.0, + moe_shared_expert_intermediate_size=None, + ) + 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=get_expert_model_parallel_group(), + ) + + layer = get_inference_optimized_moe_spec()(config=config).cuda().eval() + local_tokens = 17 + torch.distributed.get_rank() + hidden_states = torch.randn( + local_tokens, 1, config.hidden_size, device="cuda", dtype=torch.bfloat16 + ) + + with torch.no_grad(), set_batch_invariant_mode(True): + training_output, _ = layer(hidden_states.clone()) + with InferenceMode.active(): + inference_output, _ = layer(hidden_states.clone()) + + torch.testing.assert_close(inference_output, training_output, atol=0, rtol=0) + # ────────────────────────────────────────────────────────────────────── # mask_routing_padding kernel diff --git a/tests/unit_tests/inference/test_moe_permute.py b/tests/unit_tests/inference/test_moe_permute.py index 6bddf515b14..8be6aec7f59 100644 --- a/tests/unit_tests/inference/test_moe_permute.py +++ b/tests/unit_tests/inference/test_moe_permute.py @@ -49,6 +49,37 @@ def _make_inputs(num_tokens, hidden_dim, topk, num_experts, seed=42): return hidden, probs, routing_map +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 + 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)) + 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) + + 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: @@ -393,6 +424,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/models/test_gpt_model_batch_invariant.py b/tests/unit_tests/models/test_gpt_model_batch_invariant.py index b52fd64f592..1e93687bcd7 100644 --- a/tests/unit_tests/models/test_gpt_model_batch_invariant.py +++ b/tests/unit_tests/models/test_gpt_model_batch_invariant.py @@ -18,7 +18,10 @@ from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec from megatron.core.models.gpt.gpt_model import GPTModel from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed -from megatron.core.transformer.custom_layers.batch_invariant_kernels import set_batch_invariant_mode +from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( + set_batch_invariant_mode, + te_supports_batch_invariant_attention, +) from megatron.core.transformer.enums import AttnBackend from megatron.core.transformer.module import Float16Module from megatron.core.transformer.transformer_config import TransformerConfig @@ -48,6 +51,10 @@ # Batch-invariant mode requires an explicit FlashAttention version; pick the newest # one available so training and inference run the same kernel. _BIK_FA_VERSION = 4 if HAVE_FA4 else 3 +pytestmark = pytest.mark.skipif( + not te_supports_batch_invariant_attention(), + reason="Batch-invariant attention requires TransformerEngine PR #3204 or >= 2.18.", +) class DummyTokenizer: diff --git a/tests/unit_tests/rl/test_rl_batch_invariant.py b/tests/unit_tests/rl/test_rl_batch_invariant.py index ab339755307..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,4 +30,86 @@ 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()]) + + +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..904a6ab9e22 --- /dev/null +++ b/tests/unit_tests/ssm/ops/test_batch_invariant_decode.py @@ -0,0 +1,715 @@ +# 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=x, + 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_requests): + return BatchInvariantDecodeBuffers.allocate( + max_requests, + self.chunk_size, + self.nh, + self.headdim, + self.ngroups, + self.dstate, + self.device, + self.dtype, + ) + + def _make_ssm_state(self, max_requests): + """Production BIK state cache: FP32 carry across Mamba chunks.""" + return torch.zeros( + 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_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_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 + _, 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] + + 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], + 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], + 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_raw_states=True, + dt_softplus=True, + dt_limit=(0.0, float("inf")), + state_dtype=torch.float32, + ) + 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_requests, 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_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 + ) + self._assert_bitwise( + y_batch_invariant[0, 0], y_full[0, prefill_len], 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_requests=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 + the replay buffer.""" + 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 + 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_like(self._make_ssm_state(max_requests)) + ssm_state[slot] = boundary_state[0] + + 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( + x[0, :prefill_len], + 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_requests, 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 = self._make_ssm_state(max_requests) + ssm_state[slot] = final_boundary[0] + 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], + 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_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_requests) + 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")) + 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_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( + torch.cat([x[0, :prefill_len], nan_x], dim=0), + 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.z[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_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_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 + ) + 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_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_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 + 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_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 = [], [], [], [] + 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_requests) + # Per-slot seeding (each slot's prefill done independently). + 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_requests) + 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, + 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 write replay buffers, + perturb slot 0, or produce nonzero output.""" + 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_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): + 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(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) + torch.testing.assert_close(bufs.num_buffered[1:], inactive_counts) + + def test_seed_skips_padding_entries(self): + """Padded prefill entries must not write any persistent replay state.""" + prefill_len = 20 + x, dt, B, C = self._make_seq(prefill_len) + bufs = self._make_bufs(3) + bufs.x[1].normal_() + bufs.z[1].normal_() + bufs.dt[1].normal_() + bufs.B[1].normal_() + bufs.C[1].normal_() + bufs.num_buffered[1] = 7 + before = tuple( + tensor[1].clone() + for tensor in (bufs.x, bufs.z, bufs.dt, bufs.B, bufs.C, bufs.num_buffered) + ) + + bufs.seed( + x[0], + x[0], + dt[0], + B[0], + C[0], + torch.tensor([0, prefill_len, prefill_len], dtype=torch.int32, device=self.device), + torch.tensor([0, -1], dtype=torch.int32, device=self.device), + ) + + for tensor, expected in zip( + (bufs.x, bufs.z, bufs.dt, bufs.B, bufs.C, bufs.num_buffered), before + ): + torch.testing.assert_close(tensor[1], expected) + + def test_cuda_graph_replay_matches_full_scan(self): + """A captured decode step advances persistent state exactly across replays.""" + 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_requests) + warmup_state = self._seed_from_prefill( + 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_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() + 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_z, + 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], + 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_z.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): + """Repeated boundary crossings where the carried state dominates the output + (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_requests, slot = 2, 0 + prefill_len = 20 + # 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) + + 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=x, + dt_bias=self.dt_bias, + dt_softplus=True, + initial_states=None, + return_final_states=True, + ) + + 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( + x[0, :prefill_len], + 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], + 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_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_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( + 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/ssm/ops/test_ssm_kernel.py b/tests/unit_tests/ssm/ops/test_ssm_kernel.py index 6ef5610f819..62c25f43f0f 100644 --- a/tests/unit_tests/ssm/ops/test_ssm_kernel.py +++ b/tests/unit_tests/ssm/ops/test_ssm_kernel.py @@ -2,6 +2,7 @@ import math import unittest +from types import SimpleNamespace from unittest.mock import MagicMock import torch @@ -82,6 +83,7 @@ def setUp(self): # Create the Mixer instance directly self.mixer = MagicMock(spec=MambaMixer) + self.mixer.config = SimpleNamespace(batch_invariant_mode=False) self.mixer.d_state = self.d_state self.mixer.d_conv = self.d_conv self.mixer.headdim = self.headdim 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..a163d5ea09d --- /dev/null +++ b/tests/unit_tests/transformer/moe/test_moe_batch_invariant.py @@ -0,0 +1,255 @@ +# 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, + flash_attention_version=4, + attention_dropout=0.0, + 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()}" + ) diff --git a/tests/unit_tests/transformer/test_te_layers_batch_invariant.py b/tests/unit_tests/transformer/test_te_layers_batch_invariant.py index 685e9332025..9431db4afd7 100644 --- a/tests/unit_tests/transformer/test_te_layers_batch_invariant.py +++ b/tests/unit_tests/transformer/test_te_layers_batch_invariant.py @@ -16,7 +16,12 @@ ) from megatron.core.tensor_parallel.layers import ColumnParallelLinear from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed -from megatron.core.transformer.custom_layers.batch_invariant_kernels import set_batch_invariant_mode +from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( + HAVE_DEEPGEMM_BF16, + assert_te_supports_batch_invariant_attention, + set_batch_invariant_mode, + te_supports_batch_invariant_attention, +) from megatron.core.transformer.enums import AttnBackend, AttnMaskType from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.utils import init_method_normal, is_te_min_version @@ -38,6 +43,26 @@ # Batch-invariant mode requires an explicit FlashAttention version. _BIK_FA_VERSION = 4 if HAVE_FA4 else 3 +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.", +) + + +@pytest.mark.parametrize("te_version", ("2.17.0+cb4a45fd", "2.18.0", "2.19.0.dev0")) +def test_batch_invariant_accepts_compatible_te(monkeypatch, te_version): + import transformer_engine + + monkeypatch.setattr(transformer_engine, "__version__", te_version) + assert_te_supports_batch_invariant_attention() + + +def test_batch_invariant_rejects_incompatible_te(monkeypatch): + import transformer_engine + + monkeypatch.setattr(transformer_engine, "__version__", "2.17.0+deadbeef") + with pytest.raises(AssertionError, match="TransformerEngine PR #3204"): + assert_te_supports_batch_invariant_attention() # ============================================================================ @@ -330,6 +355,7 @@ def test_column_parallel_linear_batch_invariant_randomized(): not (is_te_min_version("2.10.0") and HAVE_FA3), reason="TE attention BIK tests require TE >= 2.10.0 and FlashAttention-3", ) +@requires_te_batch_invariant_attention def test_te_attention_layer_batch_invariant_randomized(): torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False @@ -754,3 +780,25 @@ def test_bik_te_general_gemm_numerical_parity(dtype): C_bik = _te_general_gemm(A, B, out_dtype=dtype, layout="TN")[0] torch.testing.assert_close(C_bik, C_ref, **_tols(dtype)) + + +@pytest.mark.skipif(not HAVE_DEEPGEMM_BF16, reason="DeepGEMM bf16 bindings are unavailable") +def test_bik_te_general_gemm_deepgemm_backend_supports_fp32_router(): + torch.manual_seed(123) + M1, M2, K, N = 37, 23, 128, 128 + A1 = torch.randn(M1, K, **_device(torch.float32)) + A2 = torch.randn(M2, K, **_device(torch.float32)) + A = torch.cat([A1, A2], dim=0) + B = torch.randn(K, N, **_device(torch.float32)) + + with set_batch_invariant_mode(True, backend="deepgemm"): + full = _te_general_gemm(A, B, out_dtype=torch.float32, layout="TN")[0] + chunks = torch.cat( + [ + _te_general_gemm(A1, B, out_dtype=torch.float32, layout="TN")[0], + _te_general_gemm(A2, B, out_dtype=torch.float32, layout="TN")[0], + ], + dim=1, + ) + + assert torch.equal(full, chunks) 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]]