[Bugfix] Stop mamba align prefill at the replay boundary

tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py

@@ -75,6 +75,34 @@ def test_mamba_align_split_partial_tail_schedule():
     assert split(self=mock, request=req2, num_new_tokens=1000) == 512
 
 
+def test_mamba_align_split_stops_at_replay_boundary():
+    """A prompt whose length is an exact multiple of block_size still gets a
+    chunk end one block below it. Otherwise the only cached Mamba state sits
+    at num_tokens, one token above the cache-hit cap (num_tokens - 1), and a
+    later request replaying the prompt misses the prefix cache entirely."""
+    block_size = 512
+    hash_block_size = 32
+    mock = SimpleNamespace(
+        cache_config=SimpleNamespace(block_size=block_size),
+        use_eagle=False,
+        hash_block_size=hash_block_size,
+        mamba_partial_cache_hit=True,
+    )
+    split = Scheduler._mamba_block_aligned_split
+
+    req = make_request("0", [0] * 4096, hash_block_size, sha256)
+    req.num_computed_tokens = 0
+    assert split(self=mock, request=req, num_new_tokens=4096) == 3584
+    req.num_computed_tokens = 3584
+    assert split(self=mock, request=req, num_new_tokens=512) == 512
+
+    # One token more: the replay boundary is the last block boundary already,
+    # so the chunking is unchanged.
+    req2 = make_request("1", [0] * 4097, hash_block_size, sha256)
+    req2.num_computed_tokens = 0
+    assert split(self=mock, request=req2, num_new_tokens=4097) == 4096
+
+
 def test_hybrid_mamba_align_partial_hash_hit():
     hash_block_size = 2
     mamba_block_size = 2 * hash_block_size

vllm/v1/core/sched/scheduler.py

@@ -394,6 +394,14 @@ def _mamba_block_aligned_split(
             end = end // block_size * block_size
 
         next_block_boundary = (start // block_size + 1) * block_size
+        # A later request replaying this prompt resumes at `num_tokens - 1` at
+        # the latest (`get_computed_blocks` caps the hit there), so the block
+        # boundary at or below that must be a chunk end too: when `num_tokens`
+        # is an exact multiple of `block_size` the only state cached below it
+        # would otherwise sit one token above the cap, missing entirely.
+        replay_boundary = min(
+            (request.num_tokens - 1) // block_size * block_size, last_cache_position
+        )
         tail_boundary = (
             request.num_prompt_tokens // self.hash_block_size * self.hash_block_size
             if self.mamba_partial_cache_hit
@@ -408,6 +416,8 @@ def _mamba_block_aligned_split(
             else 0,
             # Never run past the last cacheable block boundary mid-chunk.
             last_cache_position,
+            # Keep the boundary a replay of this prompt can resume from.
+            replay_boundary,
             # Fine-grained hits: the prompt's partial-tail entry can only be
             # registered by a chunk ending exactly at its last hash boundary.
             tail_boundary