Inference: Reduce mamba scratch space size by an order of magnitude.

megatron/core/inference/config.py

@@ -333,9 +333,13 @@ class InferenceConfig:
 
     This budget covers both buffers allocated by MambaSlotAllocator: the durable cache
     (ssm_states/conv_states, max_slots slots reused across requests) and the per-step
-    extraction scratch (intermediate_ssm_out/intermediate_conv_out, sized to the
-    worst-case 3 * max_requests slots). The scratch is reserved from this budget first,
-    so a larger max_requests leaves fewer durable slots."""
+    extraction scratch (intermediate_ssm_out/intermediate_conv_out). The scratch is
+    sized to the tighter of two per-step bounds,
+    ``min(ceil(max_tokens / block_size_tokens) + 1, 3 * max_requests)``, since a single
+    engine step can extract at most one state per block_size_tokens of its token budget
+    (and at most 3 per request). The scratch is reserved from this budget first, so a
+    smaller ``max_tokens`` (or ``max_requests``) shrinks the scratch and leaves more
+    durable cache slots."""
 
     # =================================
     # Logging config

megatron/core/inference/contexts/attention_context/mamba_metadata.py

@@ -14,14 +14,25 @@ class MambaMetadata:
     """Manages the metadata tensors required for Mamba layers during inference."""
 
     def __init__(
-        self, max_requests: int, max_tokens: int, mamba_chunk_size: int = 128, d_conv: int = 0
+        self,
+        max_requests: int,
+        max_tokens: int,
+        *,
+        max_intermediate_count: int,
+        mamba_chunk_size: int = 128,
+        d_conv: int = 0,
     ):
         """
         Initializes the Mamba slot allocator.
 
         Args:
             max_requests (int): The maximum number of concurrent requests.
             max_tokens (int): The maximum number of tokens.
+            max_intermediate_count (int): Per-step upper bound on Mamba
+                intermediate-state extractions; sizes the intermediate metadata
+                buffers. Computed once by DynamicInferenceContext (as
+                max_mamba_intermediate_states_per_step) and shared with
+                MambaSlotAllocator.
             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.
@@ -91,9 +102,9 @@ def __init__(
         )
         self.mamba_state_free_slot_count = self.max_requests
 
-        # Intermediate state extraction buffers (CUDA graph compatible)
-        # Each prefill request can produce up to 3 intermediate offsets
-        self.max_intermediate_count = MAX_INTERMEDIATE_OFFSETS_PER_REQUEST * max_requests
+        # Intermediate state extraction buffers (CUDA graph compatible). Sized by
+        # the per-step token-budget cap shared from DynamicInferenceContext.
+        self.max_intermediate_count = max_intermediate_count
         self._intermediate_chunk_indices_buffer = torch.zeros(
             self.max_intermediate_count, dtype=torch.int64, device=self.device
         )
@@ -381,13 +392,24 @@ def _update_intermediate_metadata(
             intermediate_counts_gpu: [real_prefill_count] int32 GPU tensor of
                 per-request offset counts (0-3), or None.
             real_prefill_count: Number of real (non-padding) prefill requests.
+            padded_prefill_count: Prefill request count after batch padding
+                (equals the captured graph bucket under CUDA graphs, or the
+                round-up-padded count in eager mode; always >= real_prefill_count).
+                Bounds the exposed/padded extent of the intermediate views via
+                ``max_count`` so CUDA graph replay always touches a fixed-size
+                region within the scratch buffers.
             cu_seqlens_gpu: GPU cu_seqlens tensor to read from. Defaults to
                 the legacy standalone ``_cu_seqlens_buffer`` used by
                 :meth:`update`; the coalesced production path passes the
                 shared ``ContextGPUView.mamba_cu_seqlens`` view.
         """
         chunk_size = self.mamba_chunk_size
-        max_count = padded_prefill_count * MAX_INTERMEDIATE_OFFSETS_PER_REQUEST
+        # Cap at the token-budget bound so the per-step views never exceed the
+        # buffers, even for high-prefill-count graph buckets where
+        # padded_prefill_count * MAX_INTERMEDIATE_OFFSETS_PER_REQUEST would.
+        max_count = min(
+            padded_prefill_count * MAX_INTERMEDIATE_OFFSETS_PER_REQUEST, self.max_intermediate_count
+        )
         if cu_seqlens_gpu is None:
             cu_seqlens_gpu = self._cu_seqlens_buffer
 
@@ -438,6 +460,13 @@ def _update_intermediate_metadata(
                 valid_abs_positions = abs_positions_2d[valid_mask]
 
                 real_count = valid_chunk_indices.numel()
+                # The token-budget bound guarantees this; fail loudly rather than
+                # silently overrun the scratch buffers if the candidate-offset
+                # logic in MambaSlotAllocator.compute_and_store_offsets changes.
+                assert real_count <= self.max_intermediate_count, (
+                    f"Mamba intermediate count {real_count} exceeds buffer size "
+                    f"{self.max_intermediate_count}"
+                )
                 self._intermediate_chunk_indices_buffer[:real_count] = valid_chunk_indices
                 self._intermediate_abs_positions_buffer[:real_count] = valid_abs_positions.to(
                     torch.int32

megatron/core/inference/contexts/dynamic_context.py

@@ -404,6 +404,19 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC
         # Block size tokens, bytes.
         kv_dtype_size_bytes = model_config.params_dtype.itemsize
         self.block_size_tokens = inference_config.block_size_tokens
+
+        # Mamba prefix caching extracts intermediate SSM states at KV block
+        # boundaries, but the chunked-scan kernel only materializes those states
+        # at multiples of mamba_chunk_size. A block boundary is therefore only a
+        # valid extraction point when block_size_tokens is a multiple of
+        # mamba_chunk_size; otherwise most/all boundaries are silently skipped.
+        if self.is_hybrid_model:
+            assert self.block_size_tokens % self.mamba_chunk_size == 0, (
+                f"block_size_tokens ({self.block_size_tokens}) must be a multiple "
+                f"of mamba_chunk_size ({self.mamba_chunk_size}) so that KV block "
+                "boundaries coincide with Mamba SSM chunk boundaries."
+            )
+
         if self.cache_mla_latent:
             #   one vector  c_t  (rank)  +  optional RoPE phase slice
             self.kv_reduced_dim = model_config.kv_lora_rank + model_config.qk_pos_emb_head_dim
@@ -598,6 +611,15 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC
 
         self.max_tokens = inference_config.max_tokens or self.DEFAULT_MAX_TOKENS
 
+        # Per-step upper bound on Mamba intermediate-state extractions, shared with
+        # MambaMetadata and MambaSlotAllocator so scratch/metadata buffers and the
+        # budget accounting agree. Bounded both by the token budget (one block
+        # boundary per block_size_tokens, +1 margin) and by the request budget
+        # (MAX_INTERMEDIATE_OFFSETS_PER_REQUEST per request); the tighter wins.
+        token_based_count = math.ceil(self.max_tokens / self.block_size_tokens) + 1
+        request_based_count = MAX_INTERMEDIATE_OFFSETS_PER_REQUEST * self.max_requests
+        self.max_mamba_intermediate_states_per_step = min(token_based_count, request_based_count)
+
         assert self.max_tokens >= self.max_requests, (
             f"max_tokens ({self.max_tokens}) must be >= "
             f"max_requests ({self.max_requests}), "
@@ -805,7 +827,7 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC
                 # budget first, then the rest sizes the "durable" cache
                 # (ssm_states/conv_states). mamba_bytes_per_req is the shared
                 # per-slot footprint of both.
-                scratch_slots = MAX_INTERMEDIATE_OFFSETS_PER_REQUEST * self.max_requests
+                scratch_slots = self.max_mamba_intermediate_states_per_step
                 scratch_bytes = scratch_slots * mamba_bytes_per_req
                 durable_slots = (prefix_cache_bytes - scratch_bytes) // mamba_bytes_per_req
                 durable_slots = max(durable_slots, 0)
@@ -864,6 +886,7 @@ def _allocate_mamba_states(self):
             self.mamba_metadata = MambaMetadata(
                 max_requests=self.max_requests,
                 max_tokens=self.max_tokens,
+                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],
             )
@@ -1700,26 +1723,24 @@ def _allocate_mamba_cache(self, mamba_gb: float) -> None:
         #                       `max_slots` slots (computed below).
         #   - "scratch" buffers: self.intermediate_ssm_out / self.intermediate_conv_out,
         #                       fixed CUDA-graph-safe staging for intermediate-state
-        #                       extraction, sized to the per-step worst case of
-        #                       `scratch_slots` = MAX_INTERMEDIATE_OFFSETS_PER_REQUEST
-        #                       * max_requests slots.
+        #                       extraction, sized to the per-step token-budget cap
+        #                       `scratch_slots` = max_mamba_intermediate_states_per_step.
         # The scratch is not part of the durable cache but consumes the same
         # per-slot bytes, so reserve it from the budget up front before sizing the
         # durable cache; otherwise total usage silently exceeds mamba_gb (and can
         # OOM) when scratch_slots > max_slots.
-        scratch_slots = MAX_INTERMEDIATE_OFFSETS_PER_REQUEST * self.max_requests
+        scratch_slots = self.max_mamba_intermediate_states_per_step
         scratch_bytes = scratch_slots * per_slot_bytes
         max_slots = (total_bytes - scratch_bytes) // per_slot_bytes  # durable slots
         if max_slots < 1:
             raise ValueError(
                 f"Mamba prefix cache budget (prefix_caching_mamba_gb={mamba_gb:.4g} GB) "
                 f"is too small. The CUDA-graph extraction scratch reserves "
-                f"{scratch_bytes / 1024**3:.4g} GB ({scratch_slots} slots = "
-                f"{MAX_INTERMEDIATE_OFFSETS_PER_REQUEST} offsets x {self.max_requests} "
-                f"requests x {per_slot_bytes / 1024:.1f} KB/slot), leaving room for "
+                f"{scratch_bytes / 1024**3:.4g} GB ({scratch_slots} slots x "
+                f"{per_slot_bytes / 1024:.1f} KB/slot), leaving room for "
                 f"fewer than one durable cache slot. Increase prefix_caching_mamba_gb "
                 f"to at least {(scratch_bytes + per_slot_bytes) / 1024**3:.4g} GB, or "
-                f"reduce max_requests."
+                f"reduce max_tokens."
             )
 
         self.mamba_slot_allocator = MambaSlotAllocator(

megatron/core/inference/contexts/mamba_slot_allocator.py

@@ -106,10 +106,10 @@ def __init__(
         # Pre-allocated "scratch" output buffers for CUDA graph compatible
         # extraction (GPU): per-step staging that the kernel writes intermediate
         # states into before commit copies them to the durable cache above. Sized
-        # to the per-step worst case (MAX_INTERMEDIATE_OFFSETS_PER_REQUEST *
-        # max_requests); the budget accounting in DynamicInferenceContext refers to
-        # these as the "scratch" buffers.
-        self.max_intermediate_count = MAX_INTERMEDIATE_OFFSETS_PER_REQUEST * context.max_requests
+        # by the per-step token budget computed once on the context; the budget
+        # accounting in DynamicInferenceContext refers to these as the "scratch"
+        # buffers.
+        self.max_intermediate_count = context.max_mamba_intermediate_states_per_step
         self.intermediate_ssm_out = torch.zeros(
             (num_mamba_layers, self.max_intermediate_count) + ssm_states_shape,
             dtype=ssm_states_dtype,
@@ -432,8 +432,10 @@ def compute_and_store_offsets(
         last_aligned_abs = (prompt_len // bs) * bs  # last complete block boundary
         penultimate_abs = (overall_required_blocks - 1) * bs
 
-        # Determine mamba_chunk_size from mamba config (128 is the standard SSM kernel chunk size)
-        mamba_chunk_size = 128
+        # SSM chunk size the mamba kernel actually runs with. States can only be
+        # extracted at multiples of this value, and it must match the value used
+        # in MambaMetadata (offset -> chunk-index conversion) to stay consistent.
+        mamba_chunk_size = ctx.mamba_chunk_size
 
         # Keep only boundaries that land inside this chunk's computed tokens and on
         # a mamba-chunk boundary (required for mid-sequence state extraction).

megatron/training/arguments.py

@@ -1978,9 +1978,10 @@ def _add_inference_args(parser):
                        'covers both the durable cache (the ssm_states/conv_states '
                        'slots reused across requests) and the per-step extraction '
                        'scratch (the intermediate_ssm_out/intermediate_conv_out '
-                       'buffers, sized to 3 * max_requests slots); the scratch is '
-                       'reserved first, so a larger max_requests leaves fewer durable '
-                       'slots.')
+                       'buffers, sized to min(ceil(max_tokens / block_size) + 1, '
+                       '3 * max_requests) slots); the scratch is reserved first, so a '
+                       'smaller max_tokens (or max_requests) shrinks the scratch and '
+                       'leaves more durable slots.')
     group.add_argument('--inference-dynamic-batching-cuda-graph-mixed-prefill-count',
                        type=int, default=16,
                        help='Number of mixed prefill requests to capture in a cuda graph.')

tests/unit_tests/inference/contexts/attention_metadata/test_mamba_metadata.py

@@ -14,7 +14,14 @@ def metadata_context(self):
         """Fixture to initialize MambaMetadata with standard constraints."""
         max_requests = 16
         max_tokens = 2048
-        metadata = MambaMetadata(max_requests=max_requests, max_tokens=max_tokens)
+        # Per-step intermediate-state cap (token budget / block_size + margin);
+        # value is irrelevant to these update() tests, which don't extract state.
+        max_intermediate_count = 17
+        metadata = MambaMetadata(
+            max_requests=max_requests,
+            max_tokens=max_tokens,
+            max_intermediate_count=max_intermediate_count,
+        )
 
         # Manually allocate some slots to simulate a running state.
         # We assume request_id i maps to mamba_slot i for simplicity in assertions.