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

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

@@ -14,14 +14,24 @@ 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 +101,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
         )
@@ -387,7 +397,12 @@ def _update_intermediate_metadata(
                 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 +453,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

@@ -598,6 +598,17 @@ 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, computed
+        # once here and shared with MambaMetadata and MambaSlotAllocator so the
+        # extraction scratch/metadata buffers and the budget accounting all agree.
+        # A step produces at most one block boundary per block_size_tokens of its
+        # token budget; the floor keeps buffers usable for tiny max_tokens and the
+        # +1 is a small safety margin.
+        self.max_mamba_intermediate_states_per_step = max(
+            MAX_INTERMEDIATE_OFFSETS_PER_REQUEST,
+            math.ceil(self.max_tokens / self.block_size_tokens) + 1,
+        )
+
         assert self.max_tokens >= self.max_requests, (
             f"max_tokens ({self.max_tokens}) must be >= "
             f"max_requests ({self.max_requests}), "
@@ -805,7 +816,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 +875,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],
             )
@@ -1699,26 +1711,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,

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.

tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py

@@ -735,11 +735,12 @@ def test_hybrid_prefix_caching_without_mamba_budget_warns(self, caplog):
 
     @pytest.mark.internal
     def test_mamba_cache_budget_too_small_raises(self):
-        # The CUDA-graph extraction scratch (MAX_INTERMEDIATE_OFFSETS_PER_REQUEST *
-        # max_requests slots) is reserved from prefix_caching_mamba_gb before the
-        # durable cache is sized. A budget too small to fit the scratch plus at
-        # least one durable slot is a hard configuration error, not a silent
-        # over-allocation (which previously could OOM at startup).
+        # The CUDA-graph extraction scratch (sized to the per-step token-budget
+        # cap, max_mamba_intermediate_states_per_step) is reserved from
+        # prefix_caching_mamba_gb before the durable cache is sized. A budget too
+        # small to fit the scratch plus at least one durable slot is a hard
+        # configuration error, not a silent over-allocation (which previously
+        # could OOM at startup).
         with pytest.raises(ValueError, match="prefix cache budget"):
             self._mctx(prefix_caching_mamba_gb=1e-5)
 
@@ -889,6 +890,61 @@ def test_mamba_intermediate_offsets(self):
                 torch.full_like(ctx5.mamba_slot_allocator.conv_states[layer, slot5], layer + 1.0),
             )
 
+    @pytest.mark.internal
+    def test_max_intermediate_states_per_step_formula(self):
+        # The extraction buffers are sized by the per-step token budget:
+        # ceil(max_tokens / block_size) + 1, floored at
+        # MAX_INTERMEDIATE_OFFSETS_PER_REQUEST -- independent of max_requests.
+        import math
+
+        from megatron.core.inference.contexts.mamba_slot_allocator import (
+            MAX_INTERMEDIATE_OFFSETS_PER_REQUEST,
+        )
+
+        # Token budget dominates the floor.
+        ctx = self._mctx(block_size_tokens=256, max_tokens=2048)
+        expected = max(
+            MAX_INTERMEDIATE_OFFSETS_PER_REQUEST,
+            math.ceil(ctx.max_tokens / ctx.block_size_tokens) + 1,
+        )
+        assert ctx.max_mamba_intermediate_states_per_step == expected
+        # The single value is shared everywhere it's consumed.
+        assert ctx.mamba_slot_allocator.max_intermediate_count == expected
+        assert ctx.mamba_metadata.max_intermediate_count == expected
+        assert ctx.mamba_slot_allocator.intermediate_ssm_out.shape[1] == expected
+
+    @pytest.mark.internal
+    def test_intermediate_count_bounded_by_token_budget(self):
+        # Claim: a single engine step emits at most max_tokens / block_size Mamba
+        # intermediate states, regardless of how many prefill requests it packs.
+        # Fill the token budget with fresh multi-block prefills and confirm the
+        # extracted count never exceeds the scratch buffer.
+        bs = 256
+        ctx = self._mctx(block_size_tokens=bs, max_tokens=2048, max_sequence_length=4096)
+        budget = ctx.max_mamba_intermediate_states_per_step
+
+        # Non-block-aligned 2.5-block prefills (each crosses a block boundary on a
+        # mamba-chunk multiple -> one intermediate offset). Distinct content so
+        # they never prefix-match one another.
+        per_req = bs * 2 + bs // 2  # 640 tokens
+        n = ctx.max_tokens // per_req
+        assert n >= 2
+        for i in range(n):
+            ctx.add_request(
+                self._req(ctx, self._prompt(per_req, offset=i * 100000), request_id=i + 1)
+            )
+
+        # Drive the step's metadata computation (populates intermediate_count).
+        ctx.initialize_attention_state()
+        ctx.transfer_bookkeeping_to_gpu()
+
+        md = ctx.mamba_metadata
+        # Extraction actually fired (guards against a silent no-op test)...
+        assert md.intermediate_count > 0
+        # ...and the packed step never exceeds the token-budget bound.
+        assert md.intermediate_count <= budget
+        assert md.intermediate_count == sum(md.per_request_intermediate_counts)
+
 
 class TestMixedCachedAndFreshPrefill(PrefixCachingTestBase):