Bs invar moe

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):

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

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()
 

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

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: