A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 1, 18:45 UTC
[Bugfix][Mooncake] Fix per-group block_size/block_hash and group_idx in MooncakeStoreConnector KV events (#44103)
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Yifan Qiao committed
Jun 7, 2026 at 00:12 UTC
51ef6888315ad242e137fa02fb6500bcf643924c
5 files changed
+84
-20
tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py
-1
@@ -16,7 +16,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler impor
16
def _make_bare_scheduler() -> MooncakeStoreScheduler:
17
scheduler = object.__new__(MooncakeStoreScheduler)
18
scheduler.kv_role = "kv_both"
19
- scheduler.original_block_size = 16
19
scheduler._block_size = 16
20
scheduler.load_specs = {}
21
scheduler._preempted_req_ids = set()
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+78
-4
@@ -134,7 +134,6 @@ def _make_store_req(req_id: str, block_hashes: list[bytes]) -> ReqMeta:
134
block_ids=([0, 1],),
135
block_hashes=block_hashes,
136
can_save=True,
137
- original_block_size=16,
137
)
138
139
@@ -970,7 +969,6 @@ def test_store_sending_thread_clamps_token_len_to_lcm():
969
block_ids=([0, 1, 2],),
970
block_hashes=[b"a0", b"a1", b"a2"],
971
can_save=True,
973
- original_block_size=16,
972
)
973
)
974
@@ -1010,7 +1008,6 @@ def test_store_sending_thread_skips_when_token_len_below_lcm():
1008
block_ids=([0, 1],),
1009
block_hashes=[b"a0", b"a1"],
1010
can_save=True,
1013
- original_block_size=64,
1011
)
1012
)
1013
@@ -1088,7 +1085,6 @@ def test_store_sending_thread_only_stores_swa_blocks_in_window():
1085
block_ids=([0, 1], list(range(8))),
1086
block_hashes=hs,
1087
can_save=True,
1091
- original_block_size=32,
1088
)
1089
)
1090
@@ -1104,6 +1100,84 @@ def test_store_sending_thread_only_stores_swa_blocks_in_window():
1100
assert swa_hashes == {hs[3].hex(), hs[7].hex()}
1101
1102
1103
+def test_store_sending_thread_kv_events_use_group_chunk_metadata():
1104
+ from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash
1105
+ from vllm.v1.kv_cache_interface import (
1106
+ FullAttentionSpec,
1107
+ KVCacheGroupSpec,
1108
+ SlidingWindowSpec,
1109
+ )
1110
+
1111
+ store = MagicMock()
1112
+ store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
1113
+ store.batch_put_from_multi_buffers.return_value = [256, 256]
1114
+
1115
+ full_spec = FullAttentionSpec(
1116
+ block_size=32, num_kv_heads=8, head_size=64, dtype=None
1117
+ )
1118
+ swa_spec = SlidingWindowSpec(
1119
+ block_size=8,
1120
+ num_kv_heads=8,
1121
+ head_size=64,
1122
+ dtype=None,
1123
+ sliding_window=8,
1124
+ )
1125
+ coord = mooncake_store_worker.MooncakeStoreCoordinator(
1126
+ [KVCacheGroupSpec(["L0"], full_spec), KVCacheGroupSpec(["L1"], swa_spec)],
1127
+ scheduler_block_size=32,
1128
+ hash_block_size=8,
1129
+ )
1130
+
1131
+ db_full = ChunkedTokenDatabase(
1132
+ KeyMetadata("test-model", 0, 0, 0, 0, group_id=0),
1133
+ block_size=32,
1134
+ hash_block_size=8,
1135
+ )
1136
+ db_full.set_kv_caches_base_addr([0x1000])
1137
+ db_full.set_block_len([512])
1138
+ db_swa = ChunkedTokenDatabase(
1139
+ KeyMetadata("test-model", 0, 0, 0, 0, group_id=1),
1140
+ block_size=8,
1141
+ hash_block_size=8,
1142
+ )
1143
+ db_swa.set_kv_caches_base_addr([0x2000])
1144
+ db_swa.set_block_len([128])
1145
+
1146
+ thread = _make_store_sending_thread(
1147
+ store,
1148
+ coord=coord,
1149
+ token_databases=[db_full, db_swa],
1150
+ block_size=32,
1151
+ )
1152
+ thread.enable_kv_event = True
1153
+
1154
+ hs = [bytes([i + 1]) * 4 for i in range(4)]
1155
+ thread.add_stored_request("r0")
1156
+ thread._handle_request(
1157
+ ReqMeta(
1158
+ req_id="r0",
1159
+ token_len_chunk=32,
1160
+ block_ids=([0], list(range(4))),
1161
+ block_hashes=hs,
1162
+ can_save=True,
1163
+ token_ids=list(range(32)),
1164
+ )
1165
+ )
1166
+
1167
+ full_event, swa_event = thread.get_kv_events()
1168
+ assert full_event.group_idx == 0
1169
+ assert full_event.block_size == 32
1170
+ assert full_event.token_ids == list(range(32))
1171
+ assert full_event.block_hashes == [
1172
+ maybe_convert_block_hash(BlockHash(b"".join(hs)))
1173
+ ]
1174
+
1175
+ assert swa_event.group_idx == 1
1176
+ assert swa_event.block_size == 8
1177
+ assert swa_event.token_ids == list(range(24, 32))
1178
+ assert swa_event.block_hashes == [maybe_convert_block_hash(BlockHash(hs[3]))]
1179
+
1180
+
1181
def _auto_set_ready_event(*args, **kwargs):
1182
"""Side effect for mocked thread constructors that auto-sets ready_event."""
1183
for arg in args:
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py
-3
@@ -213,7 +213,6 @@ class ReqMeta:
213
current_event: torch.cuda.Event | None = None
214
215
token_ids: list[int] | None = None
216
- original_block_size: int | None = None
216
217
@staticmethod
218
def from_request_tracker(
@@ -223,7 +222,6 @@ class ReqMeta:
222
skip_save: bool | None = False,
223
block_hashes: list[BlockHash] | None = None,
224
is_last_chunk: bool | None = None,
226
- original_block_size: int | None = None,
225
) -> "ReqMeta | None":
226
"""Create ReqMeta from a RequestTracker."""
227
if block_hashes is None:
@@ -274,7 +272,6 @@ class ReqMeta:
272
block_hashes=block_hashes,
273
is_last_chunk=is_last_chunk,
274
token_ids=token_ids,
277
- original_block_size=original_block_size,
275
)
276
277
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py
+1
-8
@@ -59,11 +59,7 @@ class MooncakeStoreScheduler:
59
)
60
self.client = LookupKeyClient(vllm_config)
61
62
- self.pcp_size = vllm_config.parallel_config.prefill_context_parallel_size
63
- self.dcp_size = vllm_config.parallel_config.decode_context_parallel_size
64
- self.original_block_size = vllm_config.cache_config.block_size
65
- # LCM for multi-group HMA; bs * pcp * dcp for single-group. Matches
66
- # the engine's own scheduler block size by construction.
62
+ # Align with the engine's own scheduler_block_size and hash_block_size.
63
self._block_size, self._hash_block_size = resolve_kv_cache_block_sizes(
64
kv_cache_config, vllm_config
65
)
@@ -221,7 +217,6 @@ class MooncakeStoreScheduler:
217
skip_save=force_skip_save,
218
block_hashes=request_real.block_hashes,
219
is_last_chunk=(request_tracker.token_len >= last_chunk_tokens_num),
224
- original_block_size=self.original_block_size,
220
)
221
if req_meta is not None:
222
meta.add_request(req_meta)
@@ -274,7 +269,6 @@ class MooncakeStoreScheduler:
269
is_last_chunk=(
270
request_tracker.token_len >= last_chunk_tokens_num
271
),
277
- original_block_size=self.original_block_size,
272
)
273
else:
274
# Decode/chunked request
@@ -312,7 +306,6 @@ class MooncakeStoreScheduler:
306
is_last_chunk=(
307
request_tracker.token_len >= last_chunk_tokens_num
308
),
315
- original_block_size=self.original_block_size,
309
)
310
311
if req_meta is not None:
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+5
-4
@@ -507,7 +507,8 @@ class KVCacheStoreSendingThread(KVTransferThread):
507
return True
508
509
def _handle_request(self, req_meta: ReqMeta):
510
- # Cache hits are always a multiple of ``lcm_block_size`` tokens
510
+ # Cache hits are always a multiple of ``lcm_block_size`` tokens, which
511
+ # is also ``store_mask``'s precondition.
512
lcm_block_size = self.coord.lcm_block_size
513
token_len = req_meta.token_len_chunk // lcm_block_size * lcm_block_size
514
block_ids_per_group = req_meta.block_ids
@@ -550,7 +551,7 @@ class KVCacheStoreSendingThread(KVTransferThread):
551
starts.append(start)
552
ends.append(end)
553
keys.append(key.to_string())
553
- block_hashes.append(req_meta.block_hashes[chunk_idx])
554
+ block_hashes.append(BlockHash(bytes.fromhex(key.chunk_hash)))
555
group_indices.append(g_idx)
556
557
# Apply put_step striding for TP
@@ -627,10 +628,11 @@ class KVCacheStoreSendingThread(KVTransferThread):
628
block_hashes=[new_block_hashes[idx]],
629
parent_block_hash=prev_key_per_group.get(g_idx),
630
token_ids=token_ids,
630
- block_size=req_meta.original_block_size,
631
+ block_size=db.block_size,
632
lora_id=None,
633
medium="cpu",
634
lora_name=None,
635
+ group_idx=g_idx,
636
)
637
stored_events.append(stored_event)
638
prev_key_per_group[g_idx] = new_block_hashes[idx]
@@ -947,7 +949,6 @@ class MooncakeStoreWorker:
949
"load_async", True
950
)
951
self.cache_config = vllm_config.cache_config
950
- self.original_block_size = self.cache_config.block_size
952
self.block_size, self.hash_block_size = resolve_kv_cache_block_sizes(
953
kv_cache_config, vllm_config
954
)