From 50f30a020b9b89f10e0184ce30fe10bdb0119808 Mon Sep 17 00:00:00 2001 From: lanqinghuan Date: Sat, 29 Aug 2026 12:40:35 +0800 Subject: [PATCH] [Prefix Cache] Support fine-grained SWA hits --- .../test_partial_prefix_cache_hits.py | 188 +++++++++++++++++- vllm/v1/core/kv_cache_coordinator.py | 36 +++- vllm/v1/core/single_type_kv_cache_manager.py | 160 ++++++++++++++- 3 files changed, 365 insertions(+), 19 deletions(-) diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index 11f34452254f..2ba5e40e1177 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -13,6 +13,7 @@ from tests.v1.core.test_prefix_caching import make_kv_cache_manager, make_request from vllm.utils.hashing import sha256 +from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_utils import ( KVCacheBlockCopy, get_block_hash, @@ -20,6 +21,7 @@ init_none_hash, ) from vllm.v1.core.sched.scheduler import Scheduler +from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, @@ -1594,7 +1596,7 @@ def test_hybrid_partial_hit_with_eagle_stays_within_group_blocks(): assert manager.allocate_slots(req1, 4, num_computed, computed_blocks) is not None -def test_hybrid_sliding_window_group_disables_partial_hash_hits(): +def test_hybrid_sliding_window_group_supports_partial_hash_hits(): hash_block_size = 2 sliding_window_block_size = 2 * hash_block_size mamba_block_size = 2 * sliding_window_block_size @@ -1629,7 +1631,7 @@ def test_hybrid_sliding_window_group_disables_partial_hash_hits(): dtype=torch.float32, sliding_window=sliding_window_block_size, ), - is_eagle_group=True, + is_eagle_group=False, ), ], ) @@ -1638,13 +1640,13 @@ def test_hybrid_sliding_window_group_disables_partial_hash_hits(): max_model_len=8192, enable_caching=True, hash_block_size=hash_block_size, - use_eagle=True, + use_eagle=False, ) - tokens = list(range(3 * sliding_window_block_size)) + tokens = list(range(mamba_block_size + hash_block_size)) request = make_request("0", tokens, hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(request) - assert not manager.coordinator.enable_partial_hash_hits + assert manager.coordinator.enable_partial_hash_hits assert ( manager.allocate_slots(request, mamba_block_size, num_computed, computed_blocks) is not None @@ -1661,8 +1663,182 @@ def test_hybrid_sliding_window_group_disables_partial_hash_hits(): ) computed_blocks, num_computed, _ = manager.get_computed_blocks(cached_request) - assert num_computed == mamba_block_size + assert num_computed == len(tokens) assert len(computed_blocks.blocks[0]) * hash_block_size == num_computed + swa_source = computed_blocks.blocks[2][-1] + assert ( + manager.allocate_slots(cached_request, 2, num_computed, computed_blocks) + is not None + ) + copies, retained = manager.take_kv_cache_block_copies() + swa_tail = manager.get_blocks(cached_request.request_id).blocks[2][-1] + assert KVCacheBlockCopy(swa_source.block_id, swa_tail.block_id) in copies + manager.block_pool.free_blocks(retained) + + +def test_hybrid_sliding_eagle_falls_back_to_coarse_mamba_checkpoint(): + """A fine EAGLE boundary can sit between restorable Mamba states. + + Reconciliation then repeatedly rewinds the SWA hit by one hash unit and + can collapse to zero. Enabling fine SWA lookup must preserve the usable + coarse Mamba checkpoint instead of regressing the pre-existing hit. + """ + hash_block_size = 2 + sliding_window_block_size = 2 * hash_block_size + mamba_block_size = 2 * sliding_window_block_size + kv_cache_config = KVCacheConfig( + num_blocks=128, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=hash_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=mamba_block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + KVCacheGroupSpec( + ["swa_draft"], + SlidingWindowSpec( + block_size=sliding_window_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=sliding_window_block_size, + ), + is_eagle_group=True, + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=hash_block_size, + use_eagle=True, + ) + + tokens = list(range(22)) + owner = make_request("owner", tokens, hash_block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(owner) + assert ( + manager.allocate_slots(owner, mamba_block_size, num_computed, computed_blocks) + is not None + ) + owner.num_computed_tokens = mamba_block_size + manager.new_step_starts() + assert manager.allocate_slots(owner, len(tokens) - mamba_block_size) is not None + owner.num_computed_tokens = len(tokens) + manager.free(owner) + manager.new_step_starts() + + replay = make_request("replay", tokens + [22, 23], hash_block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(replay) + + assert manager.coordinator.enable_partial_hash_hits + # The usable pre-existing checkpoint is one Mamba block (8 tokens). The + # fine-only convergence path reaches zero for this layout. + assert num_computed == mamba_block_size + assert manager.allocate_slots(replay, 2, num_computed, computed_blocks) is not None + + +def test_sliding_window_fine_hit_validates_physical_page_span(): + hash_block_size = 2 + block_size = 4 + spec = SlidingWindowSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=6, + ) + pool = BlockPool( + num_gpu_blocks=16, + enable_caching=True, + hash_block_size=hash_block_size, + ) + request = make_request("owner", list(range(12)), hash_block_size, sha256) + blocks = pool.get_new_blocks(3) + pool.cache_full_blocks(request, blocks, 0, 2, block_size, 0) + pool.cache_partial_block(request, blocks[2], 10, 0, block_size) + + def find(drop_eagle_block: bool): + return SlidingWindowManager.find_longest_cache_hit( + block_hashes=request.block_hashes, + max_length=10, + kv_cache_group_ids=[0], + block_pool=pool, + kv_cache_spec=spec, + drop_eagle_block=drop_eagle_block, + alignment_tokens=hash_block_size, + ) + + hit_blocks, hit_length = find(False) + assert hit_length == 10 + assert [block.block_id for block in hit_blocks[0]] == [ + pool.null_block.block_id, + blocks[1].block_id, + blocks[2].block_id, + ] + + hit_blocks, hit_length = find(True) + assert hit_length == 8 + assert [block.block_id for block in hit_blocks[0]] == [ + blocks[0].block_id, + blocks[1].block_id, + ] + + # A live endpoint is insufficient after an interior window page is evicted. + pool._maybe_evict_cached_block(blocks[1]) + _, hit_length = find(False) + assert hit_length == 4 + + +def test_sliding_window_fine_retention_pins_physical_page_span(): + hash_block_size = 2 + block_size = 4 + spec = SlidingWindowSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=6, + ) + + def retained(boundary: int, use_eagle: bool): + mask = SlidingWindowManager.reachable_block_mask( + start_block=0, + end_block=4, + alignment_tokens=hash_block_size, + kv_cache_spec=spec, + use_eagle=use_eagle, + retention_interval=0, + reachable_boundaries=(boundary,), + ) + assert mask is not None + return {idx for idx, keep in enumerate(mask) if keep} + + # The replay boundary for a 10-token prompt is token 9, which floors to + # the 8-token hash boundary. Its six-token SWA span intersects pages 0-1. + assert retained(9, use_eagle=False) == {0, 1} + # EAGLE proves the same 8-token hit using the 10-token endpoint, so the + # next physical page must remain cached even though it is not returned in + # the target block table. + assert retained(9, use_eagle=True) == {0, 1, 2} + # A shared-prefix junction can itself be a fine boundary. The 10-token + # hit's live window intersects physical pages 1-2. + assert retained(10, use_eagle=False) == {1, 2} @pytest.mark.parametrize("dcp_world_size", [1, 2, 4]) diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index c097542f8953..7b2f061b854c 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -758,6 +758,38 @@ def find_longest_cache_hit( self, block_hashes: list[BlockHash], max_cache_hit_length: int, + ) -> tuple[tuple[list[KVCacheBlock], ...], int, int]: + result = self._find_longest_cache_hit_once( + block_hashes, + max_cache_hit_length, + alignment_tokens=self._cache_hit_alignment_tokens, + enable_partial_hash_hits=self.enable_partial_hash_hits, + ) + + # A fine candidate can expose a cross-manager reconciliation gap: for + # example, an EAGLE/SWA group may accept a hash-unit boundary that an + # align-mode Mamba group cannot restore without an internal checkpoint. + # Do not let enabling fine lookup regress a usable coarse checkpoint. + # Only pay for the second lookup when the first pass observed such a + # gap; otherwise fine lookup is a superset of the coarse candidates. + if self.enable_partial_hash_hits and result[2] > 0: + coarse_result = self._find_longest_cache_hit_once( + block_hashes, + max_cache_hit_length, + alignment_tokens=self.scheduler_block_size, + enable_partial_hash_hits=False, + ) + if coarse_result[1] > result[1]: + return coarse_result + return result + + def _find_longest_cache_hit_once( + self, + block_hashes: list[BlockHash], + max_cache_hit_length: int, + *, + alignment_tokens: int, + enable_partial_hash_hits: bool, ) -> tuple[tuple[list[KVCacheBlock], ...], int, int]: """ Find the longest cache hit using an iterative fixed-point algorithm. @@ -827,7 +859,7 @@ def find_longest_cache_hit( if drop_eagle_block and not isinstance(spec, MambaSpec): eagle_margin = ( self.hash_block_size - if self.enable_partial_hash_hits + if enable_partial_hash_hits and manager_cls.supports_fine_grained_hash_lookup and group_block_size > self.hash_block_size else group_block_size @@ -842,7 +874,7 @@ def find_longest_cache_hit( block_pool=self.block_pool, kv_cache_spec=spec, drop_eagle_block=drop_eagle_block, - alignment_tokens=self._cache_hit_alignment_tokens, + alignment_tokens=alignment_tokens, dcp_world_size=( self.dcp_world_size if isinstance(spec, FullAttentionSpec) diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index b08f6f3c3875..4f3905226692 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -457,7 +457,7 @@ def cache_blocks( block_mask = self.reachable_block_mask( start_block=num_cached_blocks, end_block=num_full_blocks, - alignment_tokens=self.scheduler_block_size, + alignment_tokens=self._retention_alignment_tokens(), kv_cache_spec=self.kv_cache_spec, use_eagle=self.use_eagle, retention_interval=retention_interval, @@ -475,6 +475,9 @@ def cache_blocks( self.num_cached_block[request.request_id] = num_full_blocks + def _retention_alignment_tokens(self) -> int: + return self.scheduler_block_size + @classmethod def reachable_block_mask( cls, @@ -880,6 +883,8 @@ def remove_skipped_blocks( class SlidingWindowManager(SingleTypeKVCacheManager): + supports_fine_grained_hash_lookup: ClassVar[bool] = True + def __init__(self, kv_cache_spec: SlidingWindowSpec, **kwargs) -> None: super().__init__(kv_cache_spec, **kwargs) self.sliding_window = kv_cache_spec.sliding_window @@ -888,6 +893,12 @@ def __init__(self, kv_cache_spec: SlidingWindowSpec, **kwargs) -> None: # cached blocks. self.extra_retained_tokens = kv_cache_spec.extra_retained_tokens + def _retention_alignment_tokens(self) -> int: + hash_block_size = self.block_pool.hash_block_size + if hash_block_size < self.block_size: + return hash_block_size + return self.scheduler_block_size + @classmethod def _contiguous_blocks_for_hit( cls, window_size: int, block_size: int, use_eagle: bool @@ -919,10 +930,21 @@ def find_longest_cache_hit( ) assert dcp_world_size == 1, "DCP not support sliding window attn now." assert pcp_world_size == 1, "PCP not support sliding window attn now." - # Fine-grained partial hits are not supported for sliding window now - assert alignment_tokens % kv_cache_spec.block_size == 0, ( - "SlidingWindowManager does not support fine-grained (partial) cache hits" - ) + block_size = kv_cache_spec.block_size + if alignment_tokens < block_size: + assert block_size % alignment_tokens == 0 + assert alignment_tokens == block_pool.hash_block_size + return cls._find_fine_grained_cache_hit( + block_hashes=block_hashes, + max_length=max_length, + kv_cache_group_ids=kv_cache_group_ids, + block_pool=block_pool, + kv_cache_spec=kv_cache_spec, + drop_eagle_block=drop_eagle_block, + alignment_tokens=alignment_tokens, + ) + + assert alignment_tokens % block_size == 0 block_hashes = resolve_block_hashes( block_hashes, block_pool.hash_block_size, @@ -946,7 +968,6 @@ def find_longest_cache_hit( [block_pool.null_block] * max_num_blocks for _ in range(len(kv_cache_group_ids)) ) - block_size = kv_cache_spec.block_size num_contiguous_blocks = 0 match_found = False # Search from right to left and early stop when a match is found. @@ -1000,6 +1021,104 @@ def find_longest_cache_hit( hit_length = len(computed_blocks[0]) * block_size return computed_blocks, hit_length + @classmethod + def _find_fine_grained_cache_hit( + cls, + block_hashes: BlockHashList, + max_length: int, + kv_cache_group_ids: list[int], + block_pool: BlockPool, + kv_cache_spec: SlidingWindowSpec, + drop_eagle_block: bool, + alignment_tokens: int, + ) -> tuple[tuple[list[KVCacheBlock], ...], int]: + """Find a hash-granularity SWA hit backed by physical cache pages. + + Candidate endpoints are sparse partial-block entries published by + producers. An endpoint is usable only when every physical page that + intersects the post-EAGLE sliding window is still cached. The returned + lists remain in physical-page units, with nulls preserving global block + table indices before the live window. + """ + assert isinstance(block_hashes, Sequence) + block_size = kv_cache_spec.block_size + max_endpoint = ( + min( + max_length // alignment_tokens, + len(block_hashes), + ) + * alignment_tokens + ) + eagle_drop = alignment_tokens if drop_eagle_block else 0 + + for endpoint in range(max_endpoint, eagle_drop, -alignment_tokens): + endpoint_hash = block_hashes[endpoint // alignment_tokens - 1] + endpoint_blocks = block_pool.get_cached_block( + endpoint_hash, kv_cache_group_ids + ) + if endpoint_blocks is None: + continue + + hit_length = endpoint - eagle_drop + window_start = max(0, hit_length - (kv_cache_spec.sliding_window - 1)) + first_page = window_start // block_size + end_page = cdiv(hit_length, block_size) + endpoint_page = (endpoint - 1) // block_size + computed_blocks: tuple[list[KVCacheBlock], ...] = tuple( + [block_pool.null_block] * end_page + for _ in range(len(kv_cache_group_ids)) + ) + + complete = True + for page_idx in range(first_page, end_page): + cached: list[KVCacheBlock] | None + if page_idx == endpoint_page: + cached = endpoint_blocks + else: + page_end = (page_idx + 1) * block_size + page_hash = block_hashes[page_end // alignment_tokens - 1] + cached = block_pool.get_cached_block(page_hash, kv_cache_group_ids) + if cached is None: + complete = False + break + for group_blocks, block in zip(computed_blocks, cached): + group_blocks[page_idx] = block + + if complete: + return computed_blocks, hit_length + + return tuple([] for _ in kv_cache_group_ids), 0 + + def cache_blocks( + self, + request: Request, + num_tokens: int, + retention_interval: int | None = None, + ) -> None: + super().cache_blocks(request, num_tokens, retention_interval=retention_interval) + hash_block_size = self.block_pool.hash_block_size + if self.block_size == hash_block_size: + return + + boundary_tokens = request.num_prompt_tokens // hash_block_size * hash_block_size + if ( + boundary_tokens == 0 + or boundary_tokens > num_tokens + or boundary_tokens % self.block_size == 0 + ): + return + blocks = self.req_to_blocks[request.request_id] + block_idx = boundary_tokens // self.block_size + if block_idx >= len(blocks) or blocks[block_idx].is_null: + return + self.block_pool.cache_partial_block( + request=request, + block=blocks[block_idx], + num_tokens=boundary_tokens, + kv_cache_group_id=self.kv_cache_group_id, + block_size=self.block_size, + ) + @classmethod def reachable_block_mask( cls, @@ -1015,7 +1134,10 @@ def reachable_block_mask( if alignment_tokens is None: # Fast path: when the coordinator imposes no alignment constraint. return None - assert alignment_tokens % kv_cache_spec.block_size == 0 + assert ( + alignment_tokens % kv_cache_spec.block_size == 0 + or kv_cache_spec.block_size % alignment_tokens == 0 + ) block_size = kv_cache_spec.block_size # Contiguous blocks a hit needs at a boundary (incl. the EAGLE peek). @@ -1041,6 +1163,8 @@ def reachable_block_mask( else (None if retention_interval == 0 else retention_interval) ) if segment_tokens is not None: + if segment_tokens < block_size: + return None per_segment = segment_tokens // block_size if need >= per_segment: # Every block is reachable; cache them all. @@ -1055,10 +1179,24 @@ def reachable_block_mask( # the ``need``-block tail ending on each boundary explicitly. if retention_interval is not None: for boundary_tokens in reachable_boundaries: - aligned = boundary_tokens // alignment_tokens * alignment_tokens - end = aligned // block_size + shift - for j in range(max(start_block, end - need), min(end_block, end)): - mask[j - start_block] = True + if alignment_tokens < block_size: + hit = boundary_tokens // alignment_tokens * alignment_tokens + endpoint = hit + alignment_tokens if use_eagle else hit + first = ( + max(0, hit - (kv_cache_spec.sliding_window - 1)) // block_size + ) + end = cdiv(hit, block_size) + for j in range(max(start_block, first), min(end_block, end)): + mask[j - start_block] = True + if use_eagle and endpoint > 0: + endpoint_page = (endpoint - 1) // block_size + if start_block <= endpoint_page < end_block: + mask[endpoint_page - start_block] = True + else: + aligned = boundary_tokens // alignment_tokens * alignment_tokens + end = aligned // block_size + shift + for j in range(max(start_block, end - need), min(end_block, end)): + mask[j - start_block] = True return mask