[Perf][KVConnector][Mooncake] Compact chunk-hash keys and zero-copy lookup wire format (#45969)
Yifan Qiao committed
Jun 20, 2026 at 15:00 UTC
ab7fcbdd5dbcb457c61f722e0a854de29491cf4d
7 files changed
+164
-53
tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py
+5
-2
@@ -7,7 +7,10 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator imp
7
ExternalCachedBlockPool,
8
MooncakeStoreCoordinator,
9
)
10
-from vllm.v1.core.kv_cache_utils import BlockHash, BlockHashListWithBlockSize
10
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
11
+ chunk_hashes_for_block_size,
12
+)
13
+from vllm.v1.core.kv_cache_utils import BlockHash
14
from vllm.v1.kv_cache_interface import (
15
FullAttentionSpec,
16
KVCacheGroupSpec,
@@ -182,7 +185,7 @@ def test_coordinator_group_block_size_double_hash():
185
]
186
coord = _make_coord(groups, hash_block_size=16)
187
hs = _hashes(4)
185
- big_hashes = list(BlockHashListWithBlockSize(hs, 16, 32))
188
+ big_hashes = list(chunk_hashes_for_block_size(hs, 16, 32))
189
exists = {(0, bytes(h)) for h in hs}
190
exists |= {(1, bytes(bh)) for bh in big_hashes}
191
cmap = ExternalCachedBlockPool(exists)
tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py
+6
-7
@@ -323,8 +323,8 @@ def test_recv_skips_swa_blocks_before_window():
323
324
def test_chunked_token_database_hash_block_size_smaller_than_block_size():
325
"""DSv4-style: hash_block_size=4, group block_size=16 — process_tokens
326
- must merge every 4 fine hashes into one chunk hash via
327
- BlockHashListWithBlockSize."""
326
+ keys each 16-token chunk by its last fine hash, keeping the Mooncake key
327
+ at one digest instead of concatenating all 4 fine hashes."""
328
md = KeyMetadata("m", 0, 0, 0, 0, group_id=3)
329
db = ChunkedTokenDatabase(md, block_size=16, hash_block_size=4)
330
db.set_kv_caches_base_addr([0])
@@ -335,8 +335,7 @@ def test_chunked_token_database_hash_block_size_smaller_than_block_size():
335
assert len(out) == 2
336
assert out[0][0] == 0 and out[0][1] == 16
337
assert out[1][0] == 16 and out[1][1] == 32
338
- # Each chunk's hash is the concatenation of 4 fine hashes.
339
- expected0 = b"".join(fine_hashes[0:4]).hex()
340
- expected1 = b"".join(fine_hashes[4:8]).hex()
341
- assert out[0][2].chunk_hash == expected0
342
- assert out[1][2].chunk_hash == expected1
338
+ # Each chunk's hash is its last (4th) fine hash, which already chains the
339
+ # prior three.
340
+ assert out[0][2].chunk_hash == fine_hashes[3].hex()
341
+ assert out[1][2].chunk_hash == fine_hashes[7].hex()
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+35
-3
@@ -23,6 +23,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
23
worker as mooncake_store_worker,
24
)
25
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
26
+ BlobBlockHashes,
27
ChunkedTokenDatabase,
28
KeyMetadata,
29
LoadSpec,
@@ -32,6 +33,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
33
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.metrics import (
34
MooncakeStoreConnectorStats,
35
)
36
+from vllm.v1.core.kv_cache_utils import BlockHash
37
38
39
def _default_send_coord() -> mooncake_store_worker.MooncakeStoreCoordinator:
@@ -1179,9 +1181,9 @@ def test_store_sending_thread_kv_events_use_group_chunk_metadata():
1181
assert full_event.group_idx == 0
1182
assert full_event.block_size == 32
1183
assert full_event.token_ids == list(range(32))
1182
- assert full_event.block_hashes == [
1183
- maybe_convert_block_hash(BlockHash(b"".join(hs)))
1184
- ]
1184
+ # block_size=32 over hash_block_size=8 (scale 4): the chunk is keyed by its
1185
+ # last sub-hash, not the concatenation of all four.
1186
+ assert full_event.block_hashes == [maybe_convert_block_hash(BlockHash(hs[3]))]
1187
1188
assert swa_event.group_idx == 1
1189
assert swa_event.block_size == 8
@@ -1749,3 +1751,33 @@ def test_store_worker_close_swallows_store_errors():
1751
worker.close()
1752
1753
assert worker.store is None
1754
+
1755
+
1756
+def test_blob_block_hashes_wire_roundtrip():
1757
+ """The lookup wire format sends a ``hash_len`` frame plus the raw hashes
1758
+ concatenated back-to-back; the server rebuilds them through a zero-copy
1759
+ ``BlobBlockHashes`` view over the frame buffer."""
1760
+ hashes = [BlockHash(bytes([i]) * 16) for i in range(5)]
1761
+ hash_len = len(hashes[0])
1762
+
1763
+ # Client side (LookupKeyClient._lookup): flat payload frame.
1764
+ blob = b"".join(hashes)
1765
+
1766
+ # Server side (LookupKeyServer): view over the frame buffer (a memoryview),
1767
+ # never materializing the full hash list upfront.
1768
+ view = BlobBlockHashes(memoryview(blob), hash_len)
1769
+
1770
+ assert len(view) == 5
1771
+ assert list(view) == hashes # default Sequence iter terminates via IndexError
1772
+ assert [bytes(h) for h in view] == hashes
1773
+ assert bytes(view[-1]) == hashes[-1]
1774
+ assert [bytes(h) for h in view[1:3]] == hashes[1:3]
1775
+ with pytest.raises(IndexError):
1776
+ _ = view[5]
1777
+
1778
+
1779
+def test_blob_block_hashes_empty():
1780
+ """Empty lookups send hash_len=0 and an empty payload."""
1781
+ view = BlobBlockHashes(memoryview(b""), 0)
1782
+ assert len(view) == 0
1783
+ assert list(view) == []
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py
+12
-12
@@ -2,13 +2,15 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""External-store cache-hit coordinator for MooncakeStoreConnector."""
4
5
+from collections.abc import Sequence
6
from typing import cast
7
8
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
9
+ chunk_hashes_for_block_size,
10
+)
11
from vllm.v1.core.block_pool import BlockPool
12
from vllm.v1.core.kv_cache_utils import (
13
BlockHash,
10
- BlockHashList,
11
- BlockHashListWithBlockSize,
14
KVCacheBlock,
15
)
16
from vllm.v1.core.single_type_kv_cache_manager import (
@@ -120,7 +122,7 @@ class MooncakeStoreCoordinator:
122
123
def find_longest_cache_hit(
124
self,
123
- block_hashes: list[BlockHash],
125
+ block_hashes: Sequence[BlockHash],
126
max_length: int,
127
cached_block_pool: ExternalCachedBlockPool,
128
*,
@@ -147,7 +149,7 @@ class MooncakeStoreCoordinator:
149
150
def load_mask(
151
self,
150
- block_hashes: list[BlockHash],
152
+ block_hashes: Sequence[BlockHash],
153
token_len: int,
154
) -> tuple[list[bool], ...]:
155
"""Per-group load masks: ``mask[g][i]`` is True iff group ``g``'s
@@ -236,17 +238,15 @@ class MooncakeStoreCoordinator:
238
return tuple(masks)
239
240
def block_hashes_for_spec(
239
- self, block_hashes: list[BlockHash], spec: KVCacheSpec
240
- ) -> BlockHashList:
241
- if spec.block_size == self.hash_block_size:
242
- return block_hashes
243
- return BlockHashListWithBlockSize(
241
+ self, block_hashes: Sequence[BlockHash], spec: KVCacheSpec
242
+ ) -> Sequence[BlockHash]:
243
+ return chunk_hashes_for_block_size(
244
block_hashes, self.hash_block_size, spec.block_size
245
)
246
247
def _find_hit_blocks(
248
self,
249
- block_hashes: list[BlockHash],
249
+ block_hashes: Sequence[BlockHash],
250
max_length: int,
251
cached_block_pool: ExternalCachedBlockPool,
252
*,
@@ -264,7 +264,7 @@ class MooncakeStoreCoordinator:
264
spec, group_ids, manager_cls = self.attention_groups[0]
265
hashes = self.block_hashes_for_spec(block_hashes, spec)
266
hit_blocks = manager_cls.find_longest_cache_hit(
267
- block_hashes=hashes,
267
+ block_hashes=hashes, # type: ignore[arg-type]
268
max_length=max_length,
269
kv_cache_group_ids=group_ids,
270
block_pool=cast(BlockPool, cached_block_pool),
@@ -304,7 +304,7 @@ class MooncakeStoreCoordinator:
304
_max_length = min(curr_hit_length + spec.block_size, max_length)
305
hashes = self.block_hashes_for_spec(block_hashes, spec)
306
hit_blocks = manager_cls.find_longest_cache_hit(
307
- block_hashes=hashes,
307
+ block_hashes=hashes, # type: ignore[arg-type]
308
max_length=_max_length,
309
kv_cache_group_ids=group_ids,
310
block_pool=cast(BlockPool, cached_block_pool),
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py
+78
-9
@@ -5,8 +5,9 @@
5
# (vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/).
6
"""Data classes for MooncakeStoreConnector."""
7
8
-from collections.abc import Iterable
8
+from collections.abc import Iterable, Sequence
9
from dataclasses import dataclass
10
+from typing import cast
11
12
import torch
13
@@ -23,6 +24,77 @@ from vllm.v1.core.kv_cache_utils import (
24
logger = init_logger(__name__)
25
26
27
+class BlobBlockHashes(Sequence[BlockHash]):
28
+ """Lazy view over a flat buffer of fixed-size block hashes to avoid the overhead
29
+ of materializing all hashes upfront.
30
+ """
31
+
32
+ def __init__(self, blob: memoryview, hash_len: int):
33
+ self._blob = blob
34
+ self._hash_len = hash_len
35
+ self._n = len(blob) // hash_len if hash_len else 0
36
+
37
+ def __len__(self) -> int:
38
+ return self._n
39
+
40
+ def __getitem__(self, idx):
41
+ if isinstance(idx, slice):
42
+ return [self[i] for i in range(*idx.indices(self._n))]
43
+ if idx < 0:
44
+ idx += self._n
45
+ if not 0 <= idx < self._n:
46
+ raise IndexError(idx)
47
+ off = idx * self._hash_len
48
+ return BlockHash(self._blob[off : off + self._hash_len])
49
+
50
+
51
+class _CompactChunkHashList(BlockHashListWithBlockSize):
52
+ """View that keys each ``block_size`` chunk by the last constituent
53
+ ``hash_block_size`` hash instead of concatenating all of them.
54
+
55
+ The engine chains block hashes (each hash folds in the previous one), so the
56
+ final sub-block hash of a chunk already uniquely identifies the whole chunk
57
+ and its prefix. Using it keeps a Mooncake key at a single hash digest
58
+ regardless of the ``block_size`` / ``hash_block_size`` ratio, instead of
59
+ growing the key linearly with it (e.g. 64x for ``block_size=256``,
60
+ ``hash_block_size=4``).
61
+ """
62
+
63
+ def __init__(
64
+ self,
65
+ block_hashes: Sequence[BlockHash],
66
+ hash_block_size: int,
67
+ target_block_size: int,
68
+ ):
69
+ # Accept any indexable sequence (e.g. the lazy ``BlobBlockHashes``), not
70
+ # just ``list``; the base only indexes/sizes it.
71
+ assert target_block_size % hash_block_size == 0
72
+ self.block_hashes = block_hashes # type: ignore[assignment]
73
+ self.scale_factor = target_block_size // hash_block_size
74
+
75
+ def _get_value_at(self, idx: int) -> BlockHash:
76
+ return self.block_hashes[idx * self.scale_factor + self.scale_factor - 1]
77
+
78
+
79
+def chunk_hashes_for_block_size(
80
+ block_hashes: Sequence[BlockHash],
81
+ hash_block_size: int,
82
+ block_size: int,
83
+) -> Sequence[BlockHash]:
84
+ """Map ``hash_block_size``-granular block hashes to one compact hash per
85
+ ``block_size`` chunk (the chunk's last sub-hash). Returns ``block_hashes``
86
+ unchanged when the two sizes are equal.
87
+ """
88
+ if block_size == hash_block_size:
89
+ return block_hashes
90
+ # Structurally a Sequence[BlockHash] (indexable + sized); the base class
91
+ # just isn't declared as one.
92
+ return cast(
93
+ "Sequence[BlockHash]",
94
+ _CompactChunkHashList(block_hashes, hash_block_size, block_size),
95
+ )
96
+
97
+
98
@dataclass
99
class KeyMetadata:
100
"""Metadata for constructing pool keys."""
@@ -138,18 +210,15 @@ class ChunkedTokenDatabase:
210
Args:
211
token_len: Total number of tokens.
212
block_hashes: Block hashes computed at ``hash_block_size`` granularity.
141
- When ``block_size > hash_block_size`` consecutive hashes are merged
142
- up to the group's ``block_size`` via ``BlockHashListWithBlockSize``.
213
+ When ``block_size > hash_block_size`` each group's ``block_size`` chunk
214
+ is keyed by its last sub-hash via ``chunk_hashes_for_block_size``.
215
mask_num: Number of tokens to skip from the beginning.
216
"""
217
if not block_hashes:
218
return
147
- if self.block_size == self.hash_block_size:
148
- chunk_hashes: Iterable[BlockHash] = block_hashes
149
- else:
150
- chunk_hashes = BlockHashListWithBlockSize(
151
- block_hashes, self.hash_block_size, self.block_size
152
- )
219
+ chunk_hashes: Iterable[BlockHash] = chunk_hashes_for_block_size(
220
+ block_hashes, self.hash_block_size, self.block_size
221
+ )
222
for chunk_id, h in enumerate(chunk_hashes):
223
start_idx = chunk_id * self.block_size
224
if start_idx >= token_len:
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/protocol.py
+4
-1
@@ -11,7 +11,10 @@ Wire format (REQ/REP over IPC):
11
12
msg_type == LOOKUP_MSG:
13
frame 1: token_len (u32 big-endian, 4 bytes)
14
- frame 2..n: msgpack-encoded list[str] of block-hash hex digests
14
+ frame 2: hash_len (u16 big-endian, 2 bytes) — byte length of each
15
+ fixed-size block hash (0 when there are no hashes)
16
+ frame 3: raw block hashes concatenated back-to-back (each hash_len
17
+ bytes); the server splits on hash_len
18
Response: [hit_count: u32 big-endian, 4 bytes]
19
20
msg_type == RESET_MSG:
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+24
-19
@@ -18,7 +18,7 @@ import socket
18
import threading
19
import time
20
from collections import defaultdict
21
-from collections.abc import Callable
21
+from collections.abc import Callable, Sequence
22
from concurrent.futures import Future, ThreadPoolExecutor
23
from dataclasses import dataclass
24
from typing import Any, Literal, TypeVar
@@ -45,6 +45,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator imp
45
MooncakeStoreCoordinator,
46
)
47
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( # noqa: E501
48
+ BlobBlockHashes,
49
ChunkedTokenDatabase,
50
KeyMetadata,
51
MooncakeStoreConnectorMetadata,
@@ -65,7 +66,6 @@ from vllm.v1.core.kv_cache_utils import (
66
resolve_kv_cache_block_sizes,
67
)
68
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec
68
-from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
69
70
from .metrics import MooncakeStoreConnectorStats
71
@@ -1372,7 +1372,7 @@ class MooncakeStoreWorker:
1372
1373
return finished_sending
1374
1375
- def lookup(self, token_len: int, block_hashes: list[BlockHash]) -> int:
1375
+ def lookup(self, token_len: int, block_hashes: Sequence[BlockHash]) -> int:
1376
"""Check how many prefix tokens exist in the store.
1377
1378
Checks across all TP ranks and PP ranks.
@@ -1392,6 +1392,11 @@ class MooncakeStoreWorker:
1392
group_hashes = self.coord.block_hashes_for_spec(
1393
block_hashes, self._kv_cache_groups[g_idx].kv_cache_spec
1394
)
1395
+ metadata_templates = [
1396
+ dataclasses.replace(db.metadata, tp_rank=tp, pp_rank=pp)
1397
+ for tp in range(tp_count)
1398
+ for pp in range(self.pp_size)
1399
+ ]
1400
for chunk_id, h in enumerate(group_hashes):
1401
start_idx = chunk_id * spec_block_size
1402
if start_idx >= token_len:
@@ -1400,11 +1405,11 @@ class MooncakeStoreWorker:
1405
chunk_id >= len(lookup_mask) or not lookup_mask[chunk_id]
1406
):
1407
continue
1403
- for tp in range(tp_count):
1404
- for pp in range(self.pp_size):
1405
- md = dataclasses.replace(db.metadata, tp_rank=tp, pp_rank=pp)
1406
- candidate_keys.append(PoolKey(md, h.hex()).to_string())
1407
- candidate_meta.append((g_idx, bytes(h)))
1408
+ h_hex = h.hex()
1409
+ h_bytes = bytes(h)
1410
+ for md in metadata_templates:
1411
+ candidate_keys.append(PoolKey(md, h_hex).to_string())
1412
+ candidate_meta.append((g_idx, h_bytes))
1413
1414
if not candidate_keys:
1415
return 0
@@ -1483,7 +1488,6 @@ class LookupKeyServer:
1488
store_worker: MooncakeStoreWorker,
1489
vllm_config: VllmConfig,
1490
):
1486
- self.decoder = MsgpackDecoder()
1491
self.ctx = zmq.Context() # type: ignore[attr-defined]
1492
socket_path = get_zmq_rpc_path_lookup(vllm_config)
1493
self._ipc_path = socket_path.removeprefix("ipc://")
@@ -1506,9 +1510,9 @@ class LookupKeyServer:
1510
1511
if msg_type == LOOKUP_MSG:
1512
token_len = int.from_bytes(all_frames[1], byteorder="big")
1509
- hash_frames = all_frames[2:]
1510
- hashes_str = self.decoder.decode(hash_frames)
1511
- block_hashes = [BlockHash(bytes.fromhex(s)) for s in hashes_str]
1513
+ hash_len = int.from_bytes(all_frames[2], byteorder="big")
1514
+ blob = all_frames[3].buffer
1515
+ block_hashes = BlobBlockHashes(blob, hash_len)
1516
result = self.store_worker.lookup(token_len, block_hashes)
1517
self.socket.send(result.to_bytes(4, "big"))
1518
@@ -1557,7 +1561,6 @@ class LookupKeyClient:
1561
"""
1562
1563
def __init__(self, vllm_config: VllmConfig):
1560
- self.encoder = MsgpackEncoder()
1564
self.ctx = zmq.Context() # type: ignore[attr-defined]
1565
socket_path = get_zmq_rpc_path_lookup(vllm_config)
1566
self.socket = make_zmq_socket(
@@ -1574,14 +1577,16 @@ class LookupKeyClient:
1577
self.futures: dict[str, Future[int]] = {}
1578
1579
def _lookup(self, token_len: int, block_hashes: list[BlockHash]) -> int:
1577
- hash_strs = [h.hex() for h in block_hashes]
1578
- hash_frames = self.encoder.encode(hash_strs)
1579
- token_len_bytes = token_len.to_bytes(4, byteorder="big")
1580
- all_frames = [LOOKUP_MSG, token_len_bytes] + list(hash_frames)
1580
+ hash_len = len(block_hashes[0]) if block_hashes else 0
1581
+ all_frames = (
1582
+ LOOKUP_MSG,
1583
+ token_len.to_bytes(4, byteorder="big"),
1584
+ hash_len.to_bytes(2, byteorder="big"),
1585
+ b"".join(block_hashes),
1586
+ )
1587
self.socket.send_multipart(all_frames, copy=False)
1588
resp = self.socket.recv()
1583
- result = int.from_bytes(resp, "big")
1584
- return result
1589
+ return int.from_bytes(resp, "big")
1590
1591
def lookup(
1592
self,