From 13c190e0652091979f707c889b2837d966ded5e7 Mon Sep 17 00:00:00 2001 From: Siddharth Singh Date: Fri, 17 Jul 2026 11:30:59 -0700 Subject: [PATCH 01/10] reduce mamba scratch space size --- .../attention_context/mamba_metadata.py | 33 +++++++-- .../inference/contexts/dynamic_context.py | 28 ++++--- .../contexts/mamba_slot_allocator.py | 8 +- .../attention_metadata/test_mamba_metadata.py | 9 ++- .../contexts/test_dynamic_prefix_caching.py | 73 +++++++++++++++++-- 5 files changed, 127 insertions(+), 24 deletions(-) diff --git a/megatron/core/inference/contexts/attention_context/mamba_metadata.py b/megatron/core/inference/contexts/attention_context/mamba_metadata.py index 3d8c6d4f5b8..cec8ec071a4 100644 --- a/megatron/core/inference/contexts/attention_context/mamba_metadata.py +++ b/megatron/core/inference/contexts/attention_context/mamba_metadata.py @@ -14,7 +14,12 @@ 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. @@ -22,6 +27,11 @@ def __init__( 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,13 @@ 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 +454,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 diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index edb632fb74a..7549d8a40d8 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/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}), " @@ -802,7 +813,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) @@ -861,6 +872,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], ) @@ -1696,26 +1708,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( diff --git a/megatron/core/inference/contexts/mamba_slot_allocator.py b/megatron/core/inference/contexts/mamba_slot_allocator.py index 69f19d442bc..6212304af12 100644 --- a/megatron/core/inference/contexts/mamba_slot_allocator.py +++ b/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, 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 a7f579051ac..99bb046d97d 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 @@ -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. 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 66d3a963786..19987ebe42e 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py +++ b/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,68 @@ 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 + + # Tiny token budget: the per-request floor takes over. + ctx_small = self._mctx(block_size_tokens=256, max_tokens=256) + assert ( + ctx_small.max_mamba_intermediate_states_per_step + == MAX_INTERMEDIATE_OFFSETS_PER_REQUEST + ) + + @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): From f1ec8cb19a3a21bd4a05e7b45fa6f411cdddd030 Mon Sep 17 00:00:00 2001 From: Siddharth Singh Date: Fri, 17 Jul 2026 12:23:32 -0700 Subject: [PATCH 02/10] lint --- .../inference/contexts/attention_context/mamba_metadata.py | 3 +-- .../inference/contexts/test_dynamic_prefix_caching.py | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/megatron/core/inference/contexts/attention_context/mamba_metadata.py b/megatron/core/inference/contexts/attention_context/mamba_metadata.py index cec8ec071a4..11363b6c336 100644 --- a/megatron/core/inference/contexts/attention_context/mamba_metadata.py +++ b/megatron/core/inference/contexts/attention_context/mamba_metadata.py @@ -401,8 +401,7 @@ def _update_intermediate_metadata( # 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, + 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 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 19987ebe42e..36e1967d4e9 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py @@ -916,8 +916,7 @@ def test_max_intermediate_states_per_step_formula(self): # Tiny token budget: the per-request floor takes over. ctx_small = self._mctx(block_size_tokens=256, max_tokens=256) assert ( - ctx_small.max_mamba_intermediate_states_per_step - == MAX_INTERMEDIATE_OFFSETS_PER_REQUEST + ctx_small.max_mamba_intermediate_states_per_step == MAX_INTERMEDIATE_OFFSETS_PER_REQUEST ) @pytest.mark.internal From 49a553c7f54a4234aebef007b1ff47409ad9b28d Mon Sep 17 00:00:00 2001 From: Siddharth Singh Date: Fri, 17 Jul 2026 12:54:08 -0700 Subject: [PATCH 03/10] lint --- .../inference/contexts/test_dynamic_prefix_caching.py | 6 ------ 1 file changed, 6 deletions(-) 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 36e1967d4e9..62d0fb70695 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py @@ -913,12 +913,6 @@ def test_max_intermediate_states_per_step_formula(self): assert ctx.mamba_metadata.max_intermediate_count == expected assert ctx.mamba_slot_allocator.intermediate_ssm_out.shape[1] == expected - # Tiny token budget: the per-request floor takes over. - ctx_small = self._mctx(block_size_tokens=256, max_tokens=256) - assert ( - ctx_small.max_mamba_intermediate_states_per_step == MAX_INTERMEDIATE_OFFSETS_PER_REQUEST - ) - @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 From d25ec052fef332d2e3abf94879d8f0c301aff2b0 Mon Sep 17 00:00:00 2001 From: Siddharth Singh Date: Tue, 21 Jul 2026 02:11:05 -0700 Subject: [PATCH 04/10] address reviewer feedback --- megatron/core/inference/config.py | 10 ++- .../attention_context/mamba_metadata.py | 7 ++ .../inference/contexts/dynamic_context.py | 31 ++++--- .../contexts/mamba_slot_allocator.py | 6 +- .../contexts/test_dynamic_prefix_caching.py | 85 +++++++++++++++++-- 5 files changed, 116 insertions(+), 23 deletions(-) diff --git a/megatron/core/inference/config.py b/megatron/core/inference/config.py index 04af29e4a83..66f752ff551 100644 --- a/megatron/core/inference/config.py +++ b/megatron/core/inference/config.py @@ -330,9 +330,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 diff --git a/megatron/core/inference/contexts/attention_context/mamba_metadata.py b/megatron/core/inference/contexts/attention_context/mamba_metadata.py index 11363b6c336..953202a3c5b 100644 --- a/megatron/core/inference/contexts/attention_context/mamba_metadata.py +++ b/megatron/core/inference/contexts/attention_context/mamba_metadata.py @@ -17,6 +17,7 @@ def __init__( self, max_requests: int, max_tokens: int, + *, max_intermediate_count: int, mamba_chunk_size: int = 128, d_conv: int = 0, @@ -391,6 +392,12 @@ 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 diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 6993f40449f..9e03a1bcb75 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/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,16 +611,14 @@ 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, - ) + # 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 >= " diff --git a/megatron/core/inference/contexts/mamba_slot_allocator.py b/megatron/core/inference/contexts/mamba_slot_allocator.py index 6212304af12..21116b0cb7e 100644 --- a/megatron/core/inference/contexts/mamba_slot_allocator.py +++ b/megatron/core/inference/contexts/mamba_slot_allocator.py @@ -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). 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 62d0fb70695..f46cda264dc 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py @@ -56,6 +56,7 @@ def _ctx( rounder=64, enable_prefix_caching=True, max_tokens=None, + max_requests=None, prefix_caching_eviction_policy=PrefixCachingEvictionPolicy.LRU, mamba_config=None, prefix_caching_mamba_gb=None, @@ -80,6 +81,7 @@ def _ctx( paused_buffer_size_gb=0.2 * buffer_size_gb, block_size_tokens=block_size_tokens, max_tokens=max_tokens, + max_requests=max_requests, mamba_inference_state_config=mamba_config, use_flashinfer_fused_rope=None, unified_memory_level=0, @@ -892,27 +894,45 @@ def test_mamba_intermediate_offsets(self): @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. + # The extraction buffers are sized by the tighter of two per-step bounds: + # token-based: ceil(max_tokens / block_size) + 1 + # request-based: MAX_INTERMEDIATE_OFFSETS_PER_REQUEST * max_requests import math from megatron.core.inference.contexts.mamba_slot_allocator import ( MAX_INTERMEDIATE_OFFSETS_PER_REQUEST, ) - # Token budget dominates the floor. + def token_based(ctx): + return math.ceil(ctx.max_tokens / ctx.block_size_tokens) + 1 + + def request_based(ctx): + return MAX_INTERMEDIATE_OFFSETS_PER_REQUEST * ctx.max_requests + + # Token-limited regime: many requests, so the token budget is tighter. 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_requests >= 3 # ensure this regime is actually token-limited + expected = min(token_based(ctx), request_based(ctx)) + assert expected == token_based(ctx) # token bound wins here 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 + # Request-limited regime: few requests but a large token budget, so + # 3 * max_requests is the tighter bound. This is the case the token-only + # formula over-allocated for (e.g. 1 request + 16384 tokens once reserved + # 65 scratch slots a single request could never fill). + ctx2 = self._mctx(block_size_tokens=256, max_tokens=2048, max_requests=2) + expected2 = min(token_based(ctx2), request_based(ctx2)) + assert expected2 == request_based(ctx2) # request bound wins here + assert expected2 < token_based(ctx2) # ...and it is strictly tighter + assert ctx2.max_mamba_intermediate_states_per_step == expected2 + assert ctx2.mamba_slot_allocator.max_intermediate_count == expected2 + assert ctx2.mamba_metadata.max_intermediate_count == expected2 + assert ctx2.mamba_slot_allocator.intermediate_ssm_out.shape[1] == expected2 + @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 @@ -945,6 +965,55 @@ def test_intermediate_count_bounded_by_token_budget(self): assert md.intermediate_count <= budget assert md.intermediate_count == sum(md.per_request_intermediate_counts) + @pytest.mark.internal + def test_intermediate_count_fills_scratch_buffer(self): + # Reviewer follow-up: drive a single step that consumes nearly the whole + # scratch buffer (not just a few slots) and confirm it is never overrun. + # + # Realistic config: attention block_size=256, mamba_chunk_size=128. Since + # 256 is a multiple of 128, a 256-aligned block boundary is also a mamba + # chunk boundary (extractable); the mamba boundary at 128 is NOT a block + # boundary, so it is never a candidate. + # + # Each request is a fresh 257-token prompt (one token past a block): it + # crosses exactly one block boundary at token 256 -> exactly one + # intermediate offset, while consuming ~one block of tokens. Packing the + # token budget with these drives intermediate_count to max_tokens // 257, + # close to the token-budget bound -- filling nearly every scratch slot, + # unlike test_intermediate_count_bounded_by_token_budget (~1/3 of them). + import math + + bs = 256 + ctx = self._mctx(block_size_tokens=bs, max_tokens=2048, max_sequence_length=4096) + assert ctx.mamba_chunk_size == 128 # bs=256 is a multiple of mamba chunk 128 + budget = ctx.max_mamba_intermediate_states_per_step + + per_req = bs + 1 # 257: crosses the block boundary at bs=256 (a mamba-chunk multiple) + n = ctx.max_tokens // per_req + assert n >= 2 + assert n <= ctx.max_requests # all packed into a single step + for i in range(n): + ctx.add_request( + self._req(ctx, self._prompt(per_req, offset=i * 100000), request_id=i + 1) + ) + + ctx.initialize_attention_state() + ctx.transfer_bookkeeping_to_gpu() + + md = ctx.mamba_metadata + # Every request contributed exactly one intermediate offset... + assert md.intermediate_count == n + assert md.intermediate_count == sum(md.per_request_intermediate_counts) + # ...the step never overruns the scratch buffer... + assert md.intermediate_count <= budget + # ...and it fills all but a small, *derived* deficit: the block-spillover + # (per_req > bs, so fewer requests fit than there are blocks) plus the + # single +1 safety margin. n <= max_requests forces the token-based bound, + # so budget == ceil(max_tokens / bs) + 1. + assert budget == math.ceil(ctx.max_tokens / bs) + 1 + expected_unfilled = (math.ceil(ctx.max_tokens / bs) - n) + 1 + assert budget - md.intermediate_count == expected_unfilled + class TestMixedCachedAndFreshPrefill(PrefixCachingTestBase): From 2ee602742cff4a7e3ac62833119de3fb82e56914 Mon Sep 17 00:00:00 2001 From: Siddharth Singh Date: Tue, 21 Jul 2026 02:12:50 -0700 Subject: [PATCH 05/10] change arguments.py --- megatron/training/arguments.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 8174fbd5eb3..69f9d261c8a 100644 --- a/megatron/training/arguments.py +++ b/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.') From 15a7477eeadc98ab9535dfd85fbe705aad15e7b9 Mon Sep 17 00:00:00 2001 From: Siddharth Singh Date: Tue, 21 Jul 2026 11:18:50 -0700 Subject: [PATCH 06/10] remove +1 from the bound --- .../inference/contexts/dynamic_context.py | 19 +++---------------- 1 file changed, 3 insertions(+), 16 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index ba71d2641db..61aca503314 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -404,19 +404,6 @@ 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 @@ -614,9 +601,9 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC # 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 + # boundary per block_size_tokens) and by the request budget + # (MAX_INTERMEDIATE_OFFSETS_PER_REQUEST per request); + token_based_count = math.ceil(self.max_tokens / self.block_size_tokens) 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) From 992595bdc4dc26f882aca8f20e1da37bcf8e5dfb Mon Sep 17 00:00:00 2001 From: Siddharth Singh Date: Tue, 21 Jul 2026 11:43:05 -0700 Subject: [PATCH 07/10] add unit test with mamba chunk size != 128 --- .../contexts/test_dynamic_prefix_caching.py | 32 ++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) 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 f46cda264dc..19bda3f5a79 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py @@ -36,7 +36,7 @@ def teardown_class(cls): Utils.destroy_model_parallel() @staticmethod - def _mamba_config(): + def _mamba_config(mamba_chunk_size=128): from megatron.core.inference.config import MambaInferenceStateConfig return MambaInferenceStateConfig( @@ -45,6 +45,7 @@ def _mamba_config(): ssm_states_shape=(4, 16), conv_states_dtype=torch.float32, ssm_states_dtype=torch.float32, + mamba_chunk_size=mamba_chunk_size, ) def _ctx( @@ -1014,6 +1015,35 @@ def test_intermediate_count_fills_scratch_buffer(self): expected_unfilled = (math.ceil(ctx.max_tokens / bs) - n) + 1 assert budget - md.intermediate_count == expected_unfilled + @pytest.mark.internal + def test_intermediate_offsets_use_configured_mamba_chunk_size(self): + # Regression guard for the fix that reads mamba_chunk_size from the model + # config instead of hardcoding 128 in compute_and_store_offsets. + # + # With a 64-token mamba chunk and 64-token blocks, a 65-token prompt's + # block boundary at token 64 is a valid mamba-chunk multiple + # (64 % 64 == 0), so its state must be extracted and cached. The old + # hardcoded filter (64 % 128 != 0) would have wrongly skipped it, caching + # nothing and leaving no resume point for a later turn. + bs = 64 + ctx = self._mctx( + mamba_config=self._mamba_config(mamba_chunk_size=64), + block_size_tokens=bs, + max_sequence_length=512, + ) + assert ctx.mamba_chunk_size == 64 + msa = ctx.mamba_slot_allocator + + # Fresh 65-token prompt: crosses exactly the block boundary at token 64. + ctx.add_request(self._req(ctx, self._prompt(bs + 1))) + + count = msa._intermediate_counts_cpu[0].item() + # The boundary at token 64 was recorded -- would be 0 under the old + # hardcoded-128 filter, since 64 % 128 != 0. + assert count == 1 + offsets = msa._intermediate_offsets_cpu[0, :count].tolist() + assert offsets == [bs] + class TestMixedCachedAndFreshPrefill(PrefixCachingTestBase): From 27c8416849e1992464fc0c0984d2f406bdf5479e Mon Sep 17 00:00:00 2001 From: Siddharth Singh Date: Tue, 21 Jul 2026 13:37:41 -0700 Subject: [PATCH 08/10] lint --- megatron/core/inference/contexts/dynamic_context.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 61aca503314..57f0c4aaccd 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -603,7 +603,7 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC # budget accounting agree. Bounded both by the token budget (one block # boundary per block_size_tokens) and by the request budget # (MAX_INTERMEDIATE_OFFSETS_PER_REQUEST per request); - token_based_count = math.ceil(self.max_tokens / self.block_size_tokens) + token_based_count = math.ceil(self.max_tokens / self.block_size_tokens) 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) From 9930376f174f6d233ddc722692e0b77b9e64b026 Mon Sep 17 00:00:00 2001 From: Siddharth Singh Date: Wed, 22 Jul 2026 10:53:35 -0700 Subject: [PATCH 09/10] repair docstring --- megatron/core/inference/config.py | 2 +- megatron/training/arguments.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/megatron/core/inference/config.py b/megatron/core/inference/config.py index d07d7b5e6f8..1ee1f64f1db 100644 --- a/megatron/core/inference/config.py +++ b/megatron/core/inference/config.py @@ -335,7 +335,7 @@ class InferenceConfig: (ssm_states/conv_states, max_slots slots reused across requests) and the per-step 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 + ``min(ceil(max_tokens / block_size_tokens), 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 diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 5c416692abe..d6fcad28a98 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1978,7 +1978,7 @@ 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 min(ceil(max_tokens / block_size) + 1, ' + 'buffers, sized to min(ceil(max_tokens / block_size), ' '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.') From 9fb837b65d9b1bc6779d643cbab0caef0f67d134 Mon Sep 17 00:00:00 2001 From: Siddharth Singh Date: Wed, 22 Jul 2026 17:28:47 -0700 Subject: [PATCH 10/10] add tests --- .../inference/contexts/test_dynamic_prefix_caching.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) 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 19bda3f5a79..d4c61750991 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py @@ -905,7 +905,7 @@ def test_max_intermediate_states_per_step_formula(self): ) def token_based(ctx): - return math.ceil(ctx.max_tokens / ctx.block_size_tokens) + 1 + return math.ceil(ctx.max_tokens / ctx.block_size_tokens) def request_based(ctx): return MAX_INTERMEDIATE_OFFSETS_PER_REQUEST * ctx.max_requests @@ -1008,11 +1008,10 @@ def test_intermediate_count_fills_scratch_buffer(self): # ...the step never overruns the scratch buffer... assert md.intermediate_count <= budget # ...and it fills all but a small, *derived* deficit: the block-spillover - # (per_req > bs, so fewer requests fit than there are blocks) plus the - # single +1 safety margin. n <= max_requests forces the token-based bound, - # so budget == ceil(max_tokens / bs) + 1. - assert budget == math.ceil(ctx.max_tokens / bs) + 1 - expected_unfilled = (math.ceil(ctx.max_tokens / bs) - n) + 1 + # (per_req > bs, so fewer requests fit than there are blocks). n <= + # max_requests forces the token-based bound, so budget == ceil(max_tokens / bs). + assert budget == math.ceil(ctx.max_tokens / bs) + expected_unfilled = math.ceil(ctx.max_tokens / bs) - n assert budget - md.intermediate_count == expected_unfilled @pytest.mark.internal