From 6988a1d77c570c17df50d7aa85e3bc2030d0c340 Mon Sep 17 00:00:00 2001 From: zack041 Date: Mon, 25 May 2026 23:40:38 -0700 Subject: [PATCH 1/3] accuracy fix mamba mtp prefix caching Signed-off-by: zack041 (cherry picked from commit 3812af4ed37258e8be9114b491d634b10c053aaa) (cherry picked from commit c48941dda114055ae6f02e374eda5d687dac46c8) --- vllm/v1/core/single_type_kv_cache_manager.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 478effcd7469..e054b07fd325 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -991,6 +991,12 @@ def find_longest_cache_hit( block_size = kv_cache_spec.block_size max_num_blocks = max_length // block_size + if use_eagle and max_num_blocks > 0: + # Full-attention cache hit lookup matches one extra block and then + # drops that final block as it's only partially accepted. Mamba/GDN state + # blocks are [null, ..., state] so popping after a match removes the state + # Instead, we can only search up to the boundary (not include final) + max_num_blocks -= 1 # Search from right to left and early stop when a match is found. for i in range(max_num_blocks - 1, -1, -1): if cached_block := block_pool.get_cached_block( From e27bada0b40b955fd10b43c665889f5e1cb1833e Mon Sep 17 00:00:00 2001 From: Alex Bilichenko Date: Fri, 12 Jun 2026 22:22:35 +0000 Subject: [PATCH 2/3] Fix Mamba MTP cache hit drop condition --- .../core/test_single_type_kv_cache_manager.py | 41 ++++++++++++++++++- vllm/v1/core/single_type_kv_cache_manager.py | 2 +- 2 files changed, 41 insertions(+), 2 deletions(-) diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py index 0e3e8879359a..94d315867c46 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/tests/v1/core/test_single_type_kv_cache_manager.py @@ -14,9 +14,14 @@ ) from vllm.v1.core.single_type_kv_cache_manager import ( ChunkedLocalAttentionManager, + MambaManager, SlidingWindowManager, ) -from vllm.v1.kv_cache_interface import ChunkedLocalAttentionSpec, SlidingWindowSpec +from vllm.v1.kv_cache_interface import ( + ChunkedLocalAttentionSpec, + MambaSpec, + SlidingWindowSpec, +) pytestmark = pytest.mark.cpu_test @@ -190,6 +195,40 @@ def run_one_case(block_is_cached, expect_length): ) +def test_mamba_possible_cached_prefix_with_eagle_drop(): + block_size = 2 + mamba_spec = MambaSpec( + block_size=block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + ) + + block_pool = BlockPool( + num_gpu_blocks=100, enable_caching=True, hash_block_size=block_size + ) + block_hash_list = [BlockHash(str(i).encode()) for i in range(3)] + + for i, block_hash in enumerate(block_hash_list): + block_pool.cached_block_hash_to_block.insert( + make_block_hash_with_group_id(block_hash, 0), + block_pool.blocks[i + 10], + ) + + computed_blocks = MambaManager.find_longest_cache_hit( + block_hashes=block_hash_list, + max_length=len(block_hash_list) * block_size, + kv_cache_group_ids=[0], + block_pool=block_pool, + kv_cache_spec=mamba_spec, + drop_eagle_block=True, + alignment_tokens=block_size, + )[0] + + assert len(computed_blocks) == 2 + assert computed_blocks[0] == block_pool.null_block + assert computed_blocks[1].block_id == 11 + + def test_chunked_local_attention_remove_skipped_blocks(): attention_spec = ChunkedLocalAttentionSpec( block_size=2, diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index e054b07fd325..0b6f1e2a547c 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -991,7 +991,7 @@ def find_longest_cache_hit( block_size = kv_cache_spec.block_size max_num_blocks = max_length // block_size - if use_eagle and max_num_blocks > 0: + if drop_eagle_block and max_num_blocks > 0: # Full-attention cache hit lookup matches one extra block and then # drops that final block as it's only partially accepted. Mamba/GDN state # blocks are [null, ..., state] so popping after a match removes the state From 9aae51a1aa4d20fdef913432994d963ef34142b7 Mon Sep 17 00:00:00 2001 From: Alex Bilichenko Date: Sun, 26 Jul 2026 18:46:14 +0000 Subject: [PATCH 3/3] Fix stale unpacking in Mamba eagle-drop cache-hit test `find_longest_cache_hit` returns `(computed_blocks, hit_length)` since #46384 ("[2/N][Core] support partial prefix cache hit for hybrid model"), where `computed_blocks` is a tuple with one block list per KV cache group. `test_mamba_possible_cached_prefix_with_eagle_drop` indexed only `[0]`, so it asserted against the tuple of groups rather than group 0's block list, and failed with `assert 1 == 2`. The other tests in this file already use `[0][0]`. Test-only change; the MambaManager logic is unaffected and its assertions (null_block at index 0, cached block 11 at index 1) pass once the correct list is selected. --- tests/v1/core/test_single_type_kv_cache_manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py index 94d315867c46..2c491e9bd52e 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/tests/v1/core/test_single_type_kv_cache_manager.py @@ -222,7 +222,7 @@ def test_mamba_possible_cached_prefix_with_eagle_drop(): kv_cache_spec=mamba_spec, drop_eagle_block=True, alignment_tokens=block_size, - )[0] + )[0][0] assert len(computed_blocks) == 2 assert computed_blocks[0] == block_pool.null_block