[KV Connector][Mooncake] Apply SWA lookup mask before hashing/key build (#47317)
Signed-off-by: Zhewen Li <zhewenli@inferact.ai> Co-authored-by: Zhewen Li <zhewenli@inferact.ai> Co-authored-by: Codex <codex@openai.com> Co-authored-by: Claude <noreply@anthropic.com>
Zhewen Li committed
Jul 9, 2026 at 12:51 UTC
2ded1b24e7ef48e21f71f51fc923a352deefbde2
2 files changed
+64
-7
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+56
@@ -1757,6 +1757,62 @@ def test_lookup_checks_all_potential_swa_hit_boundaries():
1757
]
1758
1759
1760
+def test_lookup_applies_swa_mask_before_accessing_hashes():
1761
+ """Lookup must apply the sparse SWA mask before touching group hashes, so
1762
+ false-mask chunks pay neither the hash access nor the key-construction cost.
1763
+ """
1764
+ from vllm.v1.kv_cache_interface import (
1765
+ FullAttentionSpec,
1766
+ KVCacheGroupSpec,
1767
+ SlidingWindowSpec,
1768
+ )
1769
+
1770
+ worker = _make_bare_worker(block_size=8)
1771
+ full = FullAttentionSpec(block_size=32, num_kv_heads=8, head_size=64, dtype=None)
1772
+ swa = SlidingWindowSpec(
1773
+ block_size=8, num_kv_heads=8, head_size=64, dtype=None, sliding_window=8
1774
+ )
1775
+ worker._kv_cache_groups = [
1776
+ KVCacheGroupSpec(["full"], full),
1777
+ KVCacheGroupSpec(["swa"], swa),
1778
+ ]
1779
+ worker.token_dbs = [
1780
+ ChunkedTokenDatabase(
1781
+ KeyMetadata("test-model", 0, 0, 0, 0, group_id=0),
1782
+ block_size=32,
1783
+ hash_block_size=8,
1784
+ ),
1785
+ ChunkedTokenDatabase(
1786
+ KeyMetadata("test-model", 0, 0, 0, 0, group_id=1),
1787
+ block_size=8,
1788
+ hash_block_size=8,
1789
+ ),
1790
+ ]
1791
+ worker.coord = mooncake_store_worker.MooncakeStoreCoordinator(
1792
+ worker._kv_cache_groups,
1793
+ scheduler_block_size=32,
1794
+ hash_block_size=8,
1795
+ retention_interval=0,
1796
+ )
1797
+ worker._init_lookup_key_prefixes()
1798
+
1799
+ block_hashes = _RecordingBlockHashes([f"h{i}".encode() for i in range(12)])
1800
+ accessed_before_rpc: list[int] = []
1801
+
1802
+ def exists(keys):
1803
+ accessed_before_rpc.extend(block_hashes.accessed)
1804
+ return [0] * len(keys)
1805
+
1806
+ worker.store.batch_is_exist.side_effect = exists
1807
+
1808
+ worker.lookup(96, block_hashes)
1809
+
1810
+ # Full-attention chunks (compact hashes at 3, 7, 11) plus only the reachable
1811
+ # SWA boundary tails (chunks 3, 7, 11) are hash-accessed. Every masked SWA
1812
+ # chunk in between is skipped before both the hash access and key build.
1813
+ assert accessed_before_rpc == [3, 7, 11, 3, 7, 11]
1814
+
1815
+
1816
# ---------------------------------------------------------------------------
1817
# register_kv_caches tests
1818
# ---------------------------------------------------------------------------
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+8
-7
@@ -59,6 +59,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.protocol import
59
RESP_OK,
60
)
61
from vllm.logger import init_logger
62
+from vllm.utils.math_utils import cdiv
63
from vllm.utils.network_utils import get_ip, make_zmq_socket
64
from vllm.v1.core.kv_cache_utils import (
65
BlockHash,
@@ -1473,14 +1474,14 @@ class MooncakeStoreWorker:
1474
group_hashes = self.coord.block_hashes_for_spec(
1475
block_hashes, self._kv_cache_groups[g_idx].kv_cache_spec
1476
)
1476
- for chunk_id, h in enumerate(group_hashes):
1477
- start_idx = chunk_id * spec_block_size
1478
- if start_idx >= token_len:
1479
- break
1480
- if lookup_mask is not None and (
1481
- chunk_id >= len(lookup_mask) or not lookup_mask[chunk_id]
1482
- ):
1477
+ max_chunks = min(len(group_hashes), cdiv(token_len, spec_block_size))
1478
+ mask_limit = (
1479
+ max_chunks if lookup_mask is None else min(max_chunks, len(lookup_mask))
1480
+ )
1481
+ for chunk_id in range(mask_limit):
1482
+ if lookup_mask is not None and not lookup_mask[chunk_id]:
1483
continue
1484
+ h = group_hashes[chunk_id]
1485
hash_hex = h.hex()
1486
for key_prefix in key_prefixes:
1487
candidate_keys.append(