Add batch-invariant MoE and Mamba inference

Author: wdykas

megatron/core/inference/communication/torch_symm_triton/__init__.py

@@ -7,4 +7,5 @@
     multimem_all_gather_v,
     multimem_all_gatherv_3tensor,
     multimem_reduce_scatter_v,
+    ordered_reduce_scatter_v,
 )

megatron/core/inference/communication/torch_symm_triton/variable_collectives.py

@@ -354,6 +354,125 @@ def multimem_reduce_scatter_v(
     return output_tensor
 
 
+@triton.jit
+def _ordered_reduce_scatter_v_kernel(
+    local_ptr,
+    buffer_ptrs_dev,
+    signal_pad_ptrs,
+    local_tokens,
+    rank_token_offset_ptr,
+    ep_max_tokens_ptr,
+    input_byte_offset,
+    HIDDEN_SIZE: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    RANK: tl.constexpr,
+    WORLD_SIZE: tl.constexpr,
+):
+    """Variable-count reduce-scatter with explicit rank-order fp32 addition.
+
+    This is intentionally not a multimem.ld_reduce kernel. Each rank reads the
+    same token row from every peer symmetric buffer in rank order and accumulates
+    in fp32, which gives batch-invariant MoE a defined cross-rank reduction tree.
+    """
+    pid = tl.program_id(axis=0)
+
+    ep_max_tokens = tl.load(ep_max_tokens_ptr)
+    if pid >= ep_max_tokens:
+        return
+
+    symm_mem_sync(
+        signal_pad_ptrs,
+        None,
+        RANK,
+        WORLD_SIZE,
+        hasPreviousMemAccess=False,
+        hasSubsequentMemAccess=True,
+    )
+    sync_threads()
+
+    tid = tl.arange(0, BLOCK_SIZE)
+    rank_token_offset = tl.load(rank_token_offset_ptr)
+    buffer_ptrs = buffer_ptrs_dev.to(tl.pointer_type(tl.uint64))
+
+    for token_offset in range(pid, local_tokens, tl.num_programs(axis=0)):
+        global_token = rank_token_offset + token_offset
+
+        for channel_offset in range(0, HIDDEN_SIZE, BLOCK_SIZE):
+            offsets = channel_offset + tid
+            mask = offsets < HIDDEN_SIZE
+            acc = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
+
+            for src_rank in tl.range(0, WORLD_SIZE):
+                peer_base = tl.load(buffer_ptrs + src_rank).to(tl.pointer_type(tl.uint8))
+                peer_ptr = (peer_base + input_byte_offset).to(tl.pointer_type(tl.float32))
+                vals = tl.load(peer_ptr + global_token * HIDDEN_SIZE + offsets, mask=mask, other=0.0)
+                acc += vals
+
+            tl.store(local_ptr + token_offset * HIDDEN_SIZE + offsets, acc, mask=mask)
+
+
+def ordered_reduce_scatter_v(
+    output_tensor: torch.Tensor,
+    input_tensor: torch.Tensor,
+    symm_mem_hdl: _SymmetricMemory,
+    rank_token_offset: torch.Tensor,
+    ep_max_tokens: torch.Tensor,
+    per_rank_max_tokens: int,
+    input_byte_offset: int = 0,
+    **kwargs,
+) -> torch.Tensor:
+    """Variable-count reduce-scatter with fixed rank-order fp32 accumulation.
+
+    This is the batch-invariant alternative to multimem_reduce_scatter_v. It
+    uses symmetric memory for peer visibility, but performs no hardware FP
+    reduction; each output token is accumulated by explicitly loading peers in
+    rank order.
+    """
+    assert HAVE_TRITON, "Triton is required for ordered_reduce_scatter_v."
+    assert (
+        output_tensor.ndim == 2 and input_tensor.ndim == 2
+    ), "output_tensor and input_tensor must be 2-D [tokens, hidden_size]."
+    assert is_device_nvls_capable(
+        output_tensor.device
+    ), "ordered_reduce_scatter_v requires a Hopper+ GPU with NVLink (SM >= 9)."
+    assert (
+        output_tensor.dtype == input_tensor.dtype == torch.float32
+    ), "ordered_reduce_scatter_v requires fp32 input and output tensors."
+    assert (
+        rank_token_offset.numel() == 1
+        and rank_token_offset.dtype == torch.int32
+        and rank_token_offset.is_cuda
+    ), "rank_token_offset must be a scalar int32 CUDA tensor."
+
+    hidden_size = output_tensor.shape[1]
+    assert (
+        input_tensor.shape[1] == hidden_size
+    ), f"input and output hidden_size mismatch: {input_tensor.shape[1]} vs {hidden_size}"
+
+    MAX_NUM_BLOCKS = kwargs.get("max_num_blocks", 128)
+    MAX_BLOCK_SIZE = 1024
+    WARP_SIZE = 32
+    block_size = min(triton.next_power_of_2(hidden_size), MAX_BLOCK_SIZE)
+    num_warps = max(1, block_size // WARP_SIZE)
+    num_blocks = min(per_rank_max_tokens, MAX_NUM_BLOCKS)
+
+    _ordered_reduce_scatter_v_kernel[(num_blocks, 1, 1)](
+        output_tensor,
+        symm_mem_hdl.buffer_ptrs_dev,
+        symm_mem_hdl.signal_pad_ptrs_dev,
+        local_tokens=output_tensor.shape[0],
+        rank_token_offset_ptr=rank_token_offset,
+        ep_max_tokens_ptr=ep_max_tokens,
+        input_byte_offset=input_byte_offset,
+        HIDDEN_SIZE=hidden_size,
+        BLOCK_SIZE=block_size,
+        RANK=symm_mem_hdl.rank,
+        WORLD_SIZE=symm_mem_hdl.world_size,
+        num_warps=num_warps,
+    )
+    return output_tensor
+
+
 @triton.jit
 def _multimem_all_gatherv_3tensor_kernel(
     local_ptr_0,

megatron/core/inference/contexts/dynamic_context.py

@@ -303,6 +303,7 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC
         else:
             self.num_attention_heads_per_partition = 1
 
+        self.batch_invariant_mode = model_config.batch_invariant_mode
         self.num_speculative_tokens = inference_config.num_speculative_tokens
         assert self.num_speculative_tokens < inference_config.block_size_tokens, (
             f"num_speculative_tokens ({self.num_speculative_tokens}) must be < "
@@ -344,6 +345,12 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC
             self.mamba_ssm_states_dtype = mamba_inference_state_config.ssm_states_dtype
             self.mamba_chunk_size = mamba_inference_state_config.mamba_chunk_size
 
+            if self.batch_invariant_mode:
+                assert self.num_speculative_tokens == 0, (
+                    "batch_invariant_mode for Mamba dynamic inference only supports "
+                    "one-token decode; set num_speculative_tokens=0."
+                )
+
             # For hybrid models, the layer map converts the global layer index to the
             # corresponding attention layer index or Mamba layer index depending on the
             # layer type.
@@ -683,6 +690,13 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC
 
         # Deal with chunked prefill
         self.enable_chunked_prefill = inference_config.enable_chunked_prefill
+        if self.batch_invariant_mode and self.is_hybrid_model and self.enable_chunked_prefill:
+            # A chunk plus its final token must fit in one step; otherwise a prompt
+            # of that length can never advance without an invalid one-token tail.
+            assert self.max_tokens > self.mamba_chunk_size, (
+                "batch-invariant Mamba chunked prefill requires max_tokens > "
+                f"mamba_chunk_size ({self.mamba_chunk_size})."
+            )
 
         # FlashInfer.
         if inference_config.use_flashinfer_fused_rope is True:

megatron/core/inference/engines/dynamic_engine.py

@@ -1467,6 +1467,26 @@ def _find_mamba_match_count(self, req: DynamicInferenceRequest) -> int:
                 return i + 1
         return 0
 
+    def _mamba_batch_invariant_prefill_chunk_length(
+        self, req: DynamicInferenceRequest, capacity: int
+    ) -> int:
+        """Raw prefill length that computes an aligned chunk within ``capacity``.
+
+        Non-final calls must start and end at Mamba chunk boundaries. The final
+        prompt call may be shorter because it seeds the decode replay tail.
+        """
+        remaining = len(req.remaining_prompt_tokens)
+        if capacity >= remaining:
+            return remaining
+
+        chunk_size = self.context.mamba_chunk_size
+        computed_tokens = (capacity // chunk_size) * chunk_size
+        if remaining - computed_tokens == 1:
+            computed_tokens -= chunk_size
+        if computed_tokens <= 0:
+            return 0
+        return computed_tokens
+
     def schedule_waiting_requests(self):
         """Tries to schedule any requests in the waiting pool."""
         # Keep track of which requests get scheduled.
@@ -1571,6 +1591,9 @@ def schedule_chunked_prefill(self):
             # is_continuing_chunked_prefill is True if we are scheduling next
             # chunk of a existing chunked prefill request
             is_continuing_chunked_prefill = self.context.chunked_prefill_request_id >= 0
+            batch_invariant_mamba_prefill = (
+                self.context.batch_invariant_mode and self.context.is_hybrid_model
+            )
 
             # Check for conflicting block hashes.
             if prefix_caching_enabled and not is_continuing_chunked_prefill:
@@ -1630,29 +1653,47 @@ def schedule_chunked_prefill(self):
                                 not in self.context.kv_block_allocator.kv_hash_to_block_id
                             ):
                                 pending_block_hashes.add(block_hash)
-                    prefill_chunk_length = self.context.max_tokens - self.context.active_token_count
+                    available_prefill_length = (
+                        self.context.max_tokens - self.context.active_token_count
+                    )
+                    scheduled_prefill_length = available_prefill_length
+
+                    if batch_invariant_mamba_prefill:
+                        scheduled_prefill_length = self._mamba_batch_invariant_prefill_chunk_length(
+                            req, available_prefill_length
+                        )
+                        # No valid non-final Mamba chunk fits in the remaining token budget.
+                        if scheduled_prefill_length == 0:
+                            can_schedule = False
+                            break
 
                     # If this chunk would leave exactly 1 token for the final chunk, reduce
                     # this chunk by 1 or skip scheduling so the final chunk has 2 tokens.
                     # This avoids the edge case where max_seqlen_q=1 which results in a bug
                     # with the Flash Attention kernel.
                     # See https://github.com/Dao-AILab/flash-attention/issues/1537
-                    if remaining_len - prefill_chunk_length == 1:
-                        if prefill_chunk_length > 1:
-                            prefill_chunk_length -= 1
+                    # Batch-invariant Mamba handles this in its chunk-length helper.
+                    if (
+                        not batch_invariant_mamba_prefill
+                        and remaining_len - scheduled_prefill_length == 1
+                    ):
+                        if scheduled_prefill_length > 1:
+                            scheduled_prefill_length -= 1
                         else:
                             # We only have space for 1 token, but remaining is 2.
                             # Delay scheduling to avoid leaving exactly 1 token for the final chunk.
                             can_schedule = False
                             break
 
-                    self.context.add_request(req, prefill_chunk_length=prefill_chunk_length)
+                    self.context.add_request(req, prefill_chunk_length=scheduled_prefill_length)
                     self._loop.call_soon_threadsafe(
                         self._loop.create_task, self._notify_cond_for_new_request()
                     )
                     self.context.chunked_prefill_request_id = req.request_id
-                    req.remaining_prompt_tokens = req.remaining_prompt_tokens[prefill_chunk_length:]
-                    req.finished_chunk_token_count += prefill_chunk_length
+                    req.remaining_prompt_tokens = req.remaining_prompt_tokens[
+                        scheduled_prefill_length:
+                    ]
+                    req.finished_chunk_token_count += scheduled_prefill_length
                     # Still have tokens to prefill, so we break and keep the
                     # chunked prefill request at the head of the waiting queue
                     # Note that we do not need to continue check the queue, as the tokens are full

megatron/core/inference/moe/batch_invariant.py

@@ -0,0 +1,112 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+"""Batch-invariant inference MoE helpers."""
+
+from typing import Optional
+from unittest.mock import MagicMock
+
+import torch
+
+from megatron.core.transformer.custom_layers.batch_invariant_kernels import (
+    grouped_gemm_batch_invariant,
+    grouped_gemm_batch_invariant_alignment,
+    is_batch_invariant_mode_enabled,
+)
+from megatron.core.utils import null_decorator
+
+try:
+    import triton
+    import triton.language as tl
+
+    HAVE_TRITON = True
+except ImportError:
+    HAVE_TRITON = False
+
+if not HAVE_TRITON:
+    triton = MagicMock()
+    triton.jit = null_decorator
+    tl = MagicMock()
+
+
+def enabled() -> bool:
+    """Return whether global batch-invariant mode is active."""
+    return is_batch_invariant_mode_enabled()
+
+
+def grouped_mm(x_bf16: torch.Tensor, weight: torch.Tensor, offs: torch.Tensor) -> torch.Tensor:
+    """Batch-invariant BF16 grouped GEMM used by inference fused MoE."""
+    return grouped_gemm_batch_invariant(
+        x_bf16,
+        weight,
+        offs=offs.to(torch.int32),
+        m_total=x_bf16.shape[0],
+    )
+
+
+def grouped_mm_alignment() -> int:
+    """Per-expert row alignment required by the batch-invariant grouped GEMM."""
+    return grouped_gemm_batch_invariant_alignment()
+
+
+@triton.jit
+def _unpermute_tokens_in_expert_order_kernel(
+    expert_out_ptr,  # [output_size, hidden_dim] bf16 expert outputs
+    probs_ptr,  # [output_size] fp32 routing probabilities
+    inverse_map_ptr,  # [num_tokens, num_local_experts] permuted row or -1
+    valid_tokens_ptr,  # scalar int32 CUDA tensor: number of valid tokens
+    output_ptr,  # [num_tokens, hidden_dim] fp32 output buffer
+    hidden_dim,
+    num_local_experts: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+):
+    """Token-local batch-invariant unpermute.
+
+    Each program owns one output token and one hidden tile. Contributions are
+    accumulated in fp32 by increasing local expert id, avoiding atomic-add order.
+    """
+    tok = tl.program_id(0)
+    block_h = tl.program_id(1)
+    valid_tokens = tl.load(valid_tokens_ptr)
+    offsets = block_h * BLOCK_H + tl.arange(0, BLOCK_H)
+    mask_h = offsets < hidden_dim
+
+    acc = tl.zeros([BLOCK_H], dtype=tl.float32)
+    if tok < valid_tokens:
+        for lid in tl.range(0, num_local_experts):
+            pos = tl.load(inverse_map_ptr + tok * num_local_experts + lid)
+            if pos >= 0:
+                prob = tl.load(probs_ptr + pos)
+                vals = tl.load(expert_out_ptr + pos * hidden_dim + offsets, mask=mask_h).to(
+                    tl.float32
+                )
+                acc += vals * prob
+        tl.store(output_ptr + tok * hidden_dim + offsets, acc, mask=mask_h)
+
+
+def unpermute_tokens_in_expert_order(
+    expert_output: torch.Tensor,
+    permuted_probs: torch.Tensor,
+    inverse_map: torch.Tensor,
+    valid_tokens: torch.Tensor,
+    out: Optional[torch.Tensor],
+) -> torch.Tensor:
+    """Reduce local expert contributions token-by-token in fixed expert order."""
+    _, hidden_dim = expert_output.shape
+    num_tokens, num_local_experts = inverse_map.shape
+    if out is None:
+        out = torch.empty(
+            num_tokens, hidden_dim, dtype=torch.float32, device=expert_output.device
+        )
+
+    BLOCK_H = min(triton.next_power_of_2(hidden_dim), 1024)
+    grid = (num_tokens, triton.cdiv(hidden_dim, BLOCK_H))
+    _unpermute_tokens_in_expert_order_kernel[grid](
+        expert_output,
+        permuted_probs,
+        inverse_map,
+        valid_tokens,
+        out,
+        hidden_dim,
+        num_local_experts,
+        BLOCK_H=BLOCK_H,
+    )
+    return out