Fix Mamba prefix cache clamp for < 2 token chunk edge case (#6930)

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Author: santhnm2

megatron/core/inference/contexts/dynamic_context.py

@@ -3088,10 +3088,16 @@ def _find_mamba_match_count(
 
         mamba_map = self.mamba_slot_allocator.hash_to_block_id
         hashes = req.precomputed_block_hashes[start_block:end_block]
-        for i in range(len(hashes) - 1, -1, -1):
-            if hashes[i] in mamba_map:
-                return i + 1
-        return 0
+
+        # Mark the blocks in range whose hash the allocator still holds state
+        # for; the farthest such block is the match count. Intersecting against
+        # the range's hashes first keeps this bounded by the range rather than
+        # the size of the whole cache.
+        block_hashes = torch.tensor(hashes, dtype=torch.int64)
+        cached = mamba_map.keys() & set(hashes)
+        cached_hashes = torch.tensor(list(cached), dtype=torch.int64)
+        is_cached = torch.isin(block_hashes, cached_hashes)
+        return int(is_cached.nonzero()[-1].item()) + 1 if is_cached.any() else 0
 
     def _compute_prefix_match(
         self,
@@ -3148,7 +3154,9 @@ def _compute_prefix_match(
         # already had Mamba state restored during the first chunk.
         if self.is_hybrid_model and self.mamba_slot_allocator is not None and finished == 0:
             num_mamba_matched = self._find_mamba_match_count(
-                req, already_allocated_blocks, already_allocated_blocks + num_matched
+                req=req,
+                start_block=already_allocated_blocks,
+                end_block=already_allocated_blocks + num_matched,
             )
             if record_mamba_match:
                 req._mamba_num_matched_blocks = num_mamba_matched
@@ -3159,12 +3167,9 @@ def _compute_prefix_match(
                 raw_skip = num_mamba_matched * self.block_size_tokens
                 if raw_skip >= prefill_chunk_length:
                     # Back off to previous block with cached Mamba state
-                    mamba_map = self.mamba_slot_allocator.hash_to_block_id
-                    backed_off_blocks = 0
-                    for j in range(num_mamba_matched - 2, -1, -1):
-                        if req.precomputed_block_hashes[j] in mamba_map:
-                            backed_off_blocks = j + 1
-                            break
+                    backed_off_blocks = self._find_mamba_match_count(
+                        req=req, start_block=0, end_block=num_mamba_matched - 1
+                    )
                     prefix_skip_tokens = backed_off_blocks * self.block_size_tokens
                 else:
                     prefix_skip_tokens = raw_skip
@@ -3183,6 +3188,32 @@ def _compute_prefix_match(
             max_skip = prefill_chunk_length - 2
             prefix_skip_tokens = (max_skip // self.block_size_tokens) * self.block_size_tokens
 
+            # Rounding down can land on a block that has no cached Mamba state.
+            # add_request() restores from `prefix_skip_tokens // block_size - 1`
+            # unconditionally and, when `restore_to_live` misses, ZEROES the SSM
+            # state while still skipping the tokens -- the request then resumes
+            # mid-prompt from a zero state and produces a wrong (but internally
+            # coherent) distribution for its first generated token.
+            #
+            # Mamba boundaries are sparse: only the few positions selected in
+            # `compute_and_store_offsets` are cached, so the clamped boundary is
+            # frequently not one of them. A 5889-token prompt caches state only at
+            # block 22 (offset 5888), the clamp moves the skip to 5632, and the
+            # restore then targets block 21, which has none.
+            #
+            # Walk back to the nearest block that actually has cached state, the
+            # same way the `raw_skip >= prefill_chunk_length` branch above does.
+            if (
+                self.is_hybrid_model
+                and self.mamba_slot_allocator is not None
+                and finished == 0
+                and prefix_skip_tokens > 0
+            ):
+                usable = self._find_mamba_match_count(
+                    req=req, start_block=0, end_block=prefix_skip_tokens // self.block_size_tokens
+                )
+                prefix_skip_tokens = usable * self.block_size_tokens
+
         effective_prefill_chunk_length = prefill_chunk_length - prefix_skip_tokens
         num_blocks_from_pool = max(
             0, overall_required_blocks - already_allocated_blocks - num_matched

tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py

@@ -1046,6 +1046,37 @@ def test_mamba_prefill_skip_and_zero_prefill(self):
         ctx5.release_memory_blocks_from_request_indexes([0])
         assert not msa5.has_state(bid5) and bh5 not in msa5.hash_to_block_id
 
+    @pytest.mark.internal
+    def test_mamba_prefill_skip_clamp_lands_on_cached_block(self):
+        # The clamp that keeps effective_prefill_chunk_length >= 2 rounds the skip
+        # down to a block boundary, which can land on a block that has no cached
+        # Mamba state. The skip must walk back to the nearest block that does,
+        # otherwise add_request() zeroes the SSM state and resumes mid-prompt.
+        # One token past 3 full blocks: skipping all 3 leaves a 1-token chunk, so
+        # the clamp always fires and moves the boundary from block 3 to block 2.
+        ctx = self._mctx()
+        bs = ctx.block_size_tokens
+        prompt = self._prompt(bs * 3 + 1)
+        ctx.add_request(self._req(ctx, prompt.clone()))
+
+        # Only the last block has Mamba state, so the clamped boundary has none
+        # and there is no earlier cached block to fall back to: skip nothing.
+        self._mamba_allocate_and_register(ctx, self._block_ids(ctx, 0, 3)[2:])
+        req = self._req(ctx, prompt.clone(), request_id=2)
+        matched, _, _, _, prefix_skip, eff_chunk = ctx._compute_prefix_match(req, len(prompt))
+        assert len(matched) == 3 and prefix_skip == 0 and eff_chunk == len(prompt)
+
+        # Same clamp, but the first block is also cached: back off to it rather
+        # than all the way to zero.
+        ctx2 = self._mctx()
+        p2 = self._prompt(bs * 3 + 1)
+        ctx2.add_request(self._req(ctx2, p2.clone()))
+        blocks2 = self._block_ids(ctx2, 0, 3)
+        self._mamba_allocate_and_register(ctx2, [blocks2[0], blocks2[2]])
+        req2 = self._req(ctx2, p2.clone(), request_id=2)
+        m2, _, _, _, ps2, ec2 = ctx2._compute_prefix_match(req2, len(p2))
+        assert len(m2) == 3 and ps2 == bs and ec2 == len(p2) - bs
+
     @pytest.mark.internal
     def test_batch_invariant_mamba_chunked_prefill_scheduler_alignment(self):
         ctx = self._mctx(