Mamba prefix caching fixes

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

@@ -445,8 +445,11 @@ def _update_intermediate_metadata(
 
                 # Pad unused slots with safe defaults for CUDA graph replay:
                 # - chunk_indices=0: reads from chunk 0 (always exists), output ignored
-                # - abs_positions=d_conv: conv gather reads tokens [0..d_conv-1],
-                #   which are within bounds and produce a valid but unused state
+                # - abs_positions=d_conv: conv gather reads tokens [0..d_conv-1].
+                #   These are within bounds only when the prefill has at least
+                #   d_conv tokens; shorter sequences (e.g. small CUDA-graph warmup
+                #   buckets) would overrun the token axis, so _ssm_prefill clamps
+                #   the gather positions into range. The gathered state is unused.
                 if real_count < max_count:
                     self._intermediate_chunk_indices_buffer[real_count:max_count].fill_(0)
                     self._intermediate_abs_positions_buffer[real_count:max_count].fill_(self.d_conv)
@@ -464,7 +467,9 @@ def _update_intermediate_metadata(
             self.intermediate_abs_positions = self._intermediate_abs_positions_buffer[:max_count]
         else:
             # No extraction: fill with safe defaults for CUDA graph warmup
-            # (same rationale as padding comment above)
+            # (same rationale as padding comment above; abs_positions=d_conv may
+            # exceed a sub-d_conv warmup sequence, so _ssm_prefill clamps the
+            # gather positions into range and the gathered state is unused)
             self._intermediate_chunk_indices_buffer[:max_count] = 0
             self._intermediate_abs_positions_buffer[:max_count] = self.d_conv
             self.intermediate_count = 0

megatron/core/inference/contexts/dynamic_context.py

@@ -278,6 +278,13 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC
         # Prefix caching hit tracking (accumulated, reset by engine after logging).
         self.prefix_cache_hits = 0  # requests that matched at least one cached block
         self.prefix_cache_blocks_matched = 0  # total matched blocks across all requests
+        # Prefill compute accounting (drained into engine accumulators each step).
+        # computed = prompt tokens actually run through the model this step;
+        # skipped = prompt tokens whose prefill was skipped via a prefix-cache hit.
+        # A high skipped fraction confirms prefix caching is saving prefill compute
+        # (so any per-step latency growth is attention-over-context, not re-prefill).
+        self.prefix_cache_prefill_computed_tokens = 0
+        self.prefix_cache_prefill_skipped_tokens = 0
 
         # Engine step counter (used for logging, metrics, and event tracking)
         self.step_count = 0
@@ -2539,12 +2546,23 @@ def reset_tensors(self) -> None:
         self.token_to_block_idx.fill_(-1)
         self.token_to_local_position_within_kv_block.fill_(0)
 
-    def reset_metadata(self) -> None:
+    def reset_metadata(self, preserve_prefix_cache: bool = False) -> None:
         """Reset all bookkeeping state: counters, block allocator, attention/mamba state.
 
         This must be called after ``initialize_all_tensors()`` and after any
         suspend/resume cycle to bring the context back to a clean state.
+
+        Args:
+            preserve_prefix_cache: When True, keep the KV block allocator's prefix-cache
+                state (hash index, ref counts, cached blocks) intact. Used by the idle
+                ``dummy_forward`` path, which only needs to clear the transient one-token
+                step state -- wiping the allocator there would destroy cross-request prefix
+                reuse for any subsequent request (the engine idles between requests at low
+                concurrency, especially with EP > 1).
         """
+        # No cache to preserve when prefix caching is off: fall back to a full
+        # reset so the disabled path is byte-identical to the original behavior.
+        preserve_prefix_cache = preserve_prefix_cache and self.enable_prefix_caching
 
         # Reset request/token counts.
         self.total_request_count = 0
@@ -2567,7 +2585,8 @@ def reset_metadata(self) -> None:
         # Reset attention, mamba, and block allocator state.
         self.reset_attention_state()
         self.reset_mamba_state()
-        self.kv_block_allocator.reset()
+        if not preserve_prefix_cache:
+            self.kv_block_allocator.reset()
         self.request_to_kv_block_ids.fill_(-1)
 
         # Reset chunked prefill state
@@ -2579,7 +2598,7 @@ def reset_metadata(self) -> None:
             token_count=0, prefill_req_count=0, decode_req_count=0
         )
 
-    def reset(self) -> None:
+    def reset(self, preserve_prefix_cache: bool = False) -> None:
         """Reset entire context.
 
         This method does:
@@ -2590,18 +2609,31 @@ def reset(self) -> None:
         This method is useful after cuda graph warmup iterations, where the
         context's memory buffer is referenced by the cuda graph system and
         cannot be deallocated.
+
+        Args:
+            preserve_prefix_cache: When True, keep the KV and Mamba prefix-cache
+                state (hash indices, cached blocks/slots, LRU clock) intact. Used by
+                the idle ``dummy_forward`` path so an idle step between requests does
+                not destroy cross-request prefix reuse.
         """
+        # No cache to preserve when prefix caching is off: fall back to a full
+        # reset so the disabled path is byte-identical to the original behavior.
+        preserve_prefix_cache = preserve_prefix_cache and self.enable_prefix_caching
         self.reset_tensors()
-        self.reset_metadata()
+        self.reset_metadata(preserve_prefix_cache=preserve_prefix_cache)
 
         # Reset lifetime counters (not reset in reset_metadata, which is also
         # called during suspend/resume where these must persist).
-        self.step_count = 0
-        self.prefix_cache_lru_clock = 0
+        if not preserve_prefix_cache:
+            self.step_count = 0
+            self.prefix_cache_lru_clock = 0
 
-        # Reset Mamba cache state
-        if self.mamba_slot_allocator is not None:
-            self.mamba_slot_allocator.reset()
+            # Reset Mamba cache state
+            if self.mamba_slot_allocator is not None:
+                self.mamba_slot_allocator.reset()
+        # When preserving prefix cache (idle dummy_forward), keep step_count
+        # monotonic so the engine's periodic logging cadence
+        # (step_count % logging_step_interval) still fires for short requests.
 
     def current_input_and_position_ids(
         self, *, num_warmup_tokens: Optional[int] = None
@@ -3072,23 +3104,28 @@ def _register_range(start: int, end: int):
             else:
                 self._pending_mamba_zeros.append(mamba_idx)
 
-            # compute_and_store_offsets sets both CPU state (hash_to_block_id,
-            # _eos_cache_block_id_gpu) and GPU staging buffers.  Runs immediately
-            # because commit_intermediate_states() reads the CPU state after the
-            # forward pass.
-            if self.mamba_slot_allocator is not None:
-                self.mamba_slot_allocator.compute_and_store_offsets(
-                    req,
-                    current_id,
-                    prefix_skip_tokens,
-                    prefill_chunk_length,
-                    num_matched_blocks,
-                    matched_block_ids,
-                    overall_required_blocks,
-                )
+        # compute_and_store_offsets sets CPU state + GPU staging buffers that
+        # commit_intermediate_states() consumes after the forward pass. Run it for
+        # EVERY prefill chunk (not just the first): the last complete block of a
+        # multi-chunk prompt falls in a continuation chunk, and caching its Mamba
+        # state is precisely what lets a later turn skip prefill on a hybrid model.
+        # Mamba slot allocation / state restore above stays first-chunk-only.
+        if self.is_hybrid_model and self.mamba_slot_allocator is not None:
+            self.mamba_slot_allocator.compute_and_store_offsets(
+                req,
+                current_id,
+                prefix_skip_tokens,
+                prefill_chunk_length,
+                num_matched_blocks,
+                matched_block_ids,
+                overall_required_blocks,
+            )
 
         self.active_token_count += effective_prefill_chunk_length
         self.lifetime_prefill_token_count += effective_prefill_chunk_length
+        if self.enable_prefix_caching:
+            self.prefix_cache_prefill_computed_tokens += effective_prefill_chunk_length
+            self.prefix_cache_prefill_skipped_tokens += prefix_skip_tokens
         self.total_request_count += 1
         self.num_prefill_requests += 1
 

megatron/core/inference/contexts/mamba_slot_allocator.py

@@ -410,24 +410,36 @@ def compute_and_store_offsets(
             overall_required_blocks: Total blocks needed for this request.
         """
         ctx = self.context
+        bs = ctx.block_size_tokens
         prompt_len = len(req.prompt_tokens)
-        num_kv_matched = num_matched_blocks
-        kv_div_abs = num_kv_matched * ctx.block_size_tokens
-        last_aligned_abs = (prompt_len // ctx.block_size_tokens) * ctx.block_size_tokens
-        seq_len = prefill_chunk_length - skip_tokens  # effective prefill length
 
-        # Compute relative offsets (relative to prefill start after skip)
-        kv_div_rel = kv_div_abs - skip_tokens
-        last_aligned_rel = last_aligned_abs - skip_tokens
-        penultimate_abs = (overall_required_blocks - 1) * ctx.block_size_tokens
-        penultimate_rel = penultimate_abs - skip_tokens
+        # Absolute token position (from the prompt start) where THIS chunk's
+        # computed tokens begin. The first chunk computes from `skip_tokens` (the
+        # prefix that was skipped); continuation chunks compute from
+        # `finished_chunk_token_count` (with skip_tokens == 0). Framing the
+        # boundary offsets against this chunk start -- rather than assuming the
+        # first chunk -- lets us extract Mamba state at block boundaries that fall
+        # in ANY chunk. In particular the last complete block of a multi-chunk
+        # prompt lives in a continuation chunk; it was previously unreachable, so
+        # non-block-aligned prompts never cached a usable resume boundary and
+        # later turns could not skip prefill.
+        chunk_start = req.finished_chunk_token_count + skip_tokens
+        seq_len = prefill_chunk_length - skip_tokens  # tokens computed this chunk
+        is_last_chunk = req.finished_chunk_token_count + prefill_chunk_length >= prompt_len
+
+        # Candidate absolute block boundaries at which to cache Mamba state.
+        kv_div_abs = num_matched_blocks * bs
+        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
 
-        # Build offset list: include if > 0, < seq_len, and % mamba_chunk_size == 0
+        # Keep only boundaries that land inside this chunk's computed tokens and on
+        # a mamba-chunk boundary (required for mid-sequence state extraction).
         offsets_set = set()
-        for offset in [kv_div_rel, last_aligned_rel, penultimate_rel]:
+        for abs_pos in (kv_div_abs, last_aligned_abs, penultimate_abs):
+            offset = abs_pos - chunk_start
             if offset > 0 and offset < seq_len and offset % mamba_chunk_size == 0:
                 offsets_set.add(offset)
 
@@ -436,8 +448,8 @@ def compute_and_store_offsets(
 
         # CPU bookkeeping writes (no GPU kernel launches).
         if count > 0:
-            abs_tokens_cpu = torch.tensor([skip_tokens + o for o in offsets], dtype=torch.int64)
-            block_indices_cpu = abs_tokens_cpu // ctx.block_size_tokens - 1
+            abs_tokens_cpu = torch.tensor([chunk_start + o for o in offsets], dtype=torch.int64)
+            block_indices_cpu = abs_tokens_cpu // bs - 1
             bids_cpu = ctx.request_to_kv_block_ids[current_id][block_indices_cpu]
 
             self._intermediate_offsets_cpu[current_id, :count] = torch.tensor(
@@ -447,9 +459,13 @@ def compute_and_store_offsets(
             self._has_intermediates = True
         self._intermediate_counts_cpu[current_id] = count
 
-        # Block-aligned EOS: prompt_len is exactly block-aligned
-        if last_aligned_abs == prompt_len and prompt_len > 0:
-            last_block_idx = prompt_len // ctx.block_size_tokens - 1
+        # Block-aligned EOS: when the prompt length is exactly block-aligned, the
+        # request's live final state IS the last block boundary's state and can be
+        # cached directly. Only valid on the final chunk (otherwise the live state
+        # is mid-prompt). Non-block-aligned prompts cache their last complete block
+        # via the intermediate-extraction path above instead.
+        if is_last_chunk and last_aligned_abs == prompt_len and prompt_len > 0:
+            last_block_idx = prompt_len // bs - 1
             if last_block_idx >= 0:
                 self._eos_cache_block_id_cpu[current_id] = ctx.request_to_kv_block_ids[current_id][
                     last_block_idx