From f4efb5ef63426dad77a45b5a9787fd05fbac5638 Mon Sep 17 00:00:00 2001 From: halt95 Date: Fri, 25 Sep 2026 02:46:42 +0300 Subject: [PATCH 063/103] mamba: restore prompt-tail prefix-cache hits with MTP (port of upstream #58368) Port of the upstream vLLM change #58368 (Apache-2.0). With MTP enabled, a prompt's tail block can again be served from the prefix cache. --- .../test_mamba_eagle_resume_checkpoint.py | 22 +++++++++++++++++++ vllm/v1/core/single_type_kv_cache_manager.py | 2 +- 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py b/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py index 059b400..87a056a 100644 --- a/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py +++ b/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py @@ -196,6 +196,28 @@ def test_sibling_resumes_below_the_block_grid_when_the_prefix_ends_early(): assert hit == resume, f"expected the resume point at {resume}, got {hit}" +@pytest.mark.parametrize("mtp_draft_group_identified", [False, True]) +def test_mamba_prefix_cache_drops_hash_block(mtp_draft_group_identified): + block_size, hash_block_size = 1024, 64 + manager = _manager( + block_size, + hash_block_size, + eagle_group=0 if mtp_draft_group_identified else None, + ) + prompt_a = PREFIX[:1600] + hit_after_mtp_block_drop = len(prompt_a) - hash_block_size + + _prefill( + manager, + _stub(manager, block_size, hash_block_size), + make_request("A", prompt_a, hash_block_size, sha256), + ) + + tokens_after_a = [-1] * 500 + hit = _sibling_hit(manager, len(prompt_a), tokens_after_a, hash_block_size) + assert hit == hit_after_mtp_block_drop + + # -------------------------------------------------------------------------- # Invariants the two clauses have to keep diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 8242e25..76ba691 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -2126,7 +2126,7 @@ class MambaManager(SingleTypeKVCacheManager): latest_prompt_hash_boundary = ( request.num_prompt_tokens // hash_block_size ) * hash_block_size - if self.use_eagle: + if self.drop_eagle_checkpoint_block: # Eagle groups match one hash unit past the candidate and drop it, # so register the tail one unit lower. latest_prompt_hash_boundary = max( -- 2.43.0