62 lines
2.6 KiB
Diff
62 lines
2.6 KiB
Diff
From f4efb5ef63426dad77a45b5a9787fd05fbac5638 Mon Sep 17 00:00:00 2001
|
|
From: halt95 <halt95@users.noreply.github.com>
|
|
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
|
|
|