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..2c491e9bd52e 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][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 478effcd7469..0b6f1e2a547c 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 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 + # 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(