diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index 852c75ae9cc..f8f43662f26 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -1757,6 +1757,62 @@ def test_lookup_checks_all_potential_swa_hit_boundaries(): ] +def test_lookup_applies_swa_mask_before_accessing_hashes(): + """Lookup must apply the sparse SWA mask before touching group hashes, so + false-mask chunks pay neither the hash access nor the key-construction cost. + """ + from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheGroupSpec, + SlidingWindowSpec, + ) + + worker = _make_bare_worker(block_size=8) + full = FullAttentionSpec(block_size=32, num_kv_heads=8, head_size=64, dtype=None) + swa = SlidingWindowSpec( + block_size=8, num_kv_heads=8, head_size=64, dtype=None, sliding_window=8 + ) + worker._kv_cache_groups = [ + KVCacheGroupSpec(["full"], full), + KVCacheGroupSpec(["swa"], swa), + ] + worker.token_dbs = [ + ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0, group_id=0), + block_size=32, + hash_block_size=8, + ), + ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0, group_id=1), + block_size=8, + hash_block_size=8, + ), + ] + worker.coord = mooncake_store_worker.MooncakeStoreCoordinator( + worker._kv_cache_groups, + scheduler_block_size=32, + hash_block_size=8, + retention_interval=0, + ) + worker._init_lookup_key_prefixes() + + block_hashes = _RecordingBlockHashes([f"h{i}".encode() for i in range(12)]) + accessed_before_rpc: list[int] = [] + + def exists(keys): + accessed_before_rpc.extend(block_hashes.accessed) + return [0] * len(keys) + + worker.store.batch_is_exist.side_effect = exists + + worker.lookup(96, block_hashes) + + # Full-attention chunks (compact hashes at 3, 7, 11) plus only the reachable + # SWA boundary tails (chunks 3, 7, 11) are hash-accessed. Every masked SWA + # chunk in between is skipped before both the hash access and key build. + assert accessed_before_rpc == [3, 7, 11, 3, 7, 11] + + # --------------------------------------------------------------------------- # register_kv_caches tests # --------------------------------------------------------------------------- diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index e60d2f47a4e..aa48d5123ef 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -59,6 +59,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.protocol import RESP_OK, ) from vllm.logger import init_logger +from vllm.utils.math_utils import cdiv from vllm.utils.network_utils import get_ip, make_zmq_socket from vllm.v1.core.kv_cache_utils import ( BlockHash, @@ -1473,14 +1474,14 @@ class MooncakeStoreWorker: group_hashes = self.coord.block_hashes_for_spec( block_hashes, self._kv_cache_groups[g_idx].kv_cache_spec ) - for chunk_id, h in enumerate(group_hashes): - start_idx = chunk_id * spec_block_size - if start_idx >= token_len: - break - if lookup_mask is not None and ( - chunk_id >= len(lookup_mask) or not lookup_mask[chunk_id] - ): + max_chunks = min(len(group_hashes), cdiv(token_len, spec_block_size)) + mask_limit = ( + max_chunks if lookup_mask is None else min(max_chunks, len(lookup_mask)) + ) + for chunk_id in range(mask_limit): + if lookup_mask is not None and not lookup_mask[chunk_id]: continue + h = group_hashes[chunk_id] hash_hex = h.hex() for key_prefix in key_prefixes: candidate_keys.append(