[Bugfix][Mamba] Restore prompt-tail prefix-cache hits with MTP (#58368)
Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Co-authored-by: Benjamin Chislett <bchislett@nvidia.com>
Netanel Haber committed
Sep 25, 2026 at 02:46 UTC
d5051abaf19a1cce19159756a79d27ccb0b2fccf
2 files changed
+23
-1
tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py
+22
@@ -196,6 +196,28 @@ def test_sibling_resumes_below_the_block_grid_when_the_prefix_ends_early():
196
assert hit == resume, f"expected the resume point at {resume}, got {hit}"
197
198
199
+@pytest.mark.parametrize("mtp_draft_group_identified", [False, True])
200
+def test_mamba_prefix_cache_drops_hash_block(mtp_draft_group_identified):
201
+ block_size, hash_block_size = 1024, 64
202
+ manager = _manager(
203
+ block_size,
204
+ hash_block_size,
205
+ eagle_group=0 if mtp_draft_group_identified else None,
206
+ )
207
+ prompt_a = PREFIX[:1600]
208
+ hit_after_mtp_block_drop = len(prompt_a) - hash_block_size
209
+
210
+ _prefill(
211
+ manager,
212
+ _stub(manager, block_size, hash_block_size),
213
+ make_request("A", prompt_a, hash_block_size, sha256),
214
+ )
215
+
216
+ tokens_after_a = [-1] * 500
217
+ hit = _sibling_hit(manager, len(prompt_a), tokens_after_a, hash_block_size)
218
+ assert hit == hit_after_mtp_block_drop
219
+
220
+
221
# --------------------------------------------------------------------------
222
# Invariants the two clauses have to keep
223
vllm/v1/core/single_type_kv_cache_manager.py
+1
-1
@@ -2078,7 +2078,7 @@ class MambaManager(SingleTypeKVCacheManager):
2078
latest_prompt_hash_boundary = (
2079
request.num_prompt_tokens // hash_block_size
2080
) * hash_block_size
2081
- if self.use_eagle:
2081
+ if self.drop_eagle_checkpoint_block:
2082
# Eagle groups match one hash unit past the candidate and drop it,
2083
# so register the tail one unit lower.
2084
latest_prompt_hash_boundary = max(