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 2, 22:00 UTC
[KV Connector] Keep MooncakeStore full hits block-aligned (#43494)
Signed-off-by: Dao Le <daole@inferact.ai> Signed-off-by: Dao Le <Dao007forever@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
Dao007forever committed
May 23, 2026 at 23:15 UTC
0902d8e62fd905f7b02b6ca0b64900d105fb77e9
4 files changed
+74
-39
docs/features/mooncake_store_connector_usage.md
-1
@@ -204,7 +204,6 @@ the vLLM JSON config.
204
205
- `load_async` (bool): Enable asynchronous loading for better compute-I/O overlap. Default: `true`.
206
- `enable_cross_layers_blocks` (bool): Enable cross-layer block packing for reduced store operations. Default: `false`.
207
-- `discard_partial_chunks` (bool): Discard partial block chunks during store. Default: `true`.
207
- `lookup_rpc_port` (int): Custom port for the ZMQ lookup RPC socket. Default: `0`.
208
209
## Notes
tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py
+63
-1
@@ -18,7 +18,6 @@ def _make_bare_scheduler() -> MooncakeStoreScheduler:
18
scheduler.kv_role = "kv_both"
19
scheduler.original_block_size = 16
20
scheduler._block_size = 16
21
- scheduler._discard_partial_chunks = True
21
scheduler.load_specs = {}
22
scheduler._preempted_req_ids = set()
23
scheduler._unfinished_request_ids = {"req-0"}
@@ -344,3 +343,66 @@ def test_from_request_tracker_no_load_saves_normally():
343
assert req_meta.can_save is True
344
assert req_meta.load_spec is None
345
assert tracker.num_saved_tokens == 48
346
+
347
+
348
+class _StubLookupClient:
349
+ def __init__(self, hit_tokens: int) -> None:
350
+ self._hit_tokens = hit_tokens
351
+
352
+ def lookup(self, token_len: int, block_hashes: list[bytes]) -> int:
353
+ return self._hit_tokens
354
+
355
+
356
+def test_full_external_hit_keeps_kvpool_cached_tokens_block_aligned():
357
+ # When the external store hits the entire prompt, scheduler must leave at
358
+ # least one token uncomputed for sampling but stay on a block boundary.
359
+ # Otherwise the recv-side load mask floors token_len to
360
+ # (num_tokens-1)//block_size, the tail partial chunk is dropped, and -- if
361
+ # the local cache covers the aligned prefix -- key_list ends up empty
362
+ # (ZeroDivisionError in the recv thread's `tp_rank % len(key_list)`).
363
+ scheduler = _make_bare_scheduler()
364
+ scheduler.load_async = True
365
+ scheduler.client = _StubLookupClient(hit_tokens=48) # full hit on 48-token prompt
366
+
367
+ request = SimpleNamespace(
368
+ request_id="req-0",
369
+ num_tokens=48,
370
+ block_hashes=[b"h0", b"h1", b"h2"],
371
+ )
372
+
373
+ need_to_allocate, load_async = scheduler.get_num_new_matched_tokens(
374
+ request, num_computed_tokens=16
375
+ )
376
+
377
+ # 47 // 16 * 16 == 32 tokens left in external store after reserving the
378
+ # sub-block tail for sampling. 32 - 16 (local) == 16 to load.
379
+ assert need_to_allocate == 16
380
+ assert load_async is True
381
+ load_spec = scheduler.load_specs["req-0"]
382
+ assert load_spec.vllm_cached_tokens == 16
383
+ assert load_spec.kvpool_cached_tokens == 32
384
+ assert load_spec.kvpool_cached_tokens % 16 == 0
385
+
386
+
387
+def test_full_external_hit_with_full_local_hit_skips_load():
388
+ # When local prefix cache already covers the block-aligned external hit,
389
+ # there is nothing for the connector to load. The pre-fix behavior would
390
+ # have scheduled a 15-token load that the recv thread couldn't translate
391
+ # into any block-aligned key.
392
+ scheduler = _make_bare_scheduler()
393
+ scheduler.load_async = True
394
+ scheduler.client = _StubLookupClient(hit_tokens=48)
395
+
396
+ request = SimpleNamespace(
397
+ request_id="req-0",
398
+ num_tokens=48,
399
+ block_hashes=[b"h0", b"h1", b"h2"],
400
+ )
401
+
402
+ need_to_allocate, load_async = scheduler.get_num_new_matched_tokens(
403
+ request, num_computed_tokens=32
404
+ )
405
+
406
+ assert need_to_allocate == 0
407
+ assert load_async is False
408
+ assert "req-0" not in scheduler.load_specs
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py
+10
-29
@@ -68,12 +68,6 @@ class MooncakeStoreScheduler:
68
kv_cache_config, vllm_config
69
)
70
71
- self._discard_partial_chunks = (
72
- vllm_config.kv_transfer_config.get_from_extra_config(
73
- "discard_partial_chunks", True
74
- )
75
- )
76
-
71
# Per-request state
72
self.load_specs: dict[str, LoadSpec] = {} # to be loaded
73
self._request_trackers: dict[str, RequestTracker] = {} # scheduled new requests
@@ -88,18 +82,19 @@ class MooncakeStoreScheduler:
82
) -> tuple[int, bool]:
83
"""Check for external KV cache hit."""
84
# Look up against the full prefill range, not just the prompt.
91
- if self._discard_partial_chunks:
92
- token_len = request.num_tokens // self._block_size * self._block_size
93
- else:
94
- token_len = request.num_tokens
95
-
85
+ token_len = request.num_tokens // self._block_size * self._block_size
86
if token_len < self._block_size:
87
return 0, False
88
89
num_external_hit_tokens = self.client.lookup(token_len, request.block_hashes)
90
91
if num_external_hit_tokens == request.num_tokens:
102
- num_external_hit_tokens -= 1
92
+ # Leave a sub-block tail uncomputed for sampling, on a block
93
+ # boundary so the recv-side load mask covers every yielded chunk.
94
+ num_external_hit_tokens = max(
95
+ 0,
96
+ (request.num_tokens - 1) // self._block_size * self._block_size,
97
+ )
98
99
if num_external_hit_tokens < num_computed_tokens:
100
need_to_allocate = 0
@@ -214,9 +209,7 @@ class MooncakeStoreScheduler:
209
self._request_trackers[request.req_id] = request_tracker
210
211
last_chunk_tokens_num = (
217
- (len(prefill_tokens) // self._block_size * self._block_size)
218
- if self._discard_partial_chunks
219
- else len(prefill_tokens)
212
+ len(prefill_tokens) // self._block_size * self._block_size
213
)
214
215
req_meta = ReqMeta.from_request_tracker(
@@ -226,7 +219,6 @@ class MooncakeStoreScheduler:
219
skip_save=force_skip_save,
220
block_hashes=request_real.block_hashes,
221
is_last_chunk=(request_tracker.token_len >= last_chunk_tokens_num),
229
- discard_partial_chunks=self._discard_partial_chunks,
222
original_block_size=self.original_block_size,
223
)
224
if req_meta is not None:
@@ -269,9 +261,7 @@ class MooncakeStoreScheduler:
261
self._request_trackers[req_id] = request_tracker
262
263
last_chunk_tokens_num = (
272
- (len(prefill_tokens) // self._block_size * self._block_size)
273
- if self._discard_partial_chunks
274
- else len(prefill_tokens)
264
+ len(prefill_tokens) // self._block_size * self._block_size
265
)
266
req_meta = ReqMeta.from_request_tracker(
267
request_tracker,
@@ -282,7 +272,6 @@ class MooncakeStoreScheduler:
272
is_last_chunk=(
273
request_tracker.token_len >= last_chunk_tokens_num
274
),
285
- discard_partial_chunks=self._discard_partial_chunks,
275
original_block_size=self.original_block_size,
276
)
277
else:
@@ -310,9 +299,7 @@ class MooncakeStoreScheduler:
299
request_tracker.update(new_block_ids)
300
301
last_chunk_tokens_num = (
313
- (prefill_end // self._block_size * self._block_size)
314
- if self._discard_partial_chunks
315
- else prefill_end
302
+ prefill_end // self._block_size * self._block_size
303
)
304
req_meta = ReqMeta.from_request_tracker(
305
request_tracker,
@@ -323,7 +310,6 @@ class MooncakeStoreScheduler:
310
is_last_chunk=(
311
request_tracker.token_len >= last_chunk_tokens_num
312
),
326
- discard_partial_chunks=self._discard_partial_chunks,
313
original_block_size=self.original_block_size,
314
)
315
@@ -341,10 +327,6 @@ class MooncakeStoreScheduler:
327
if not load_spec:
328
continue
329
num_tokens_to_compute = load_spec.kvpool_cached_tokens
344
- if (num_tokens_to_compute % self._block_size != 0) and (
345
- num_tokens_to_compute == unfinished_req.num_tokens - 1
346
- ):
347
- num_tokens_to_compute = num_tokens_to_compute + 1
330
request_tracker = RequestTracker(
331
req_id=request_id,
332
token_len=num_tokens_to_compute,
@@ -358,7 +340,6 @@ class MooncakeStoreScheduler:
340
load_spec=load_spec,
341
skip_save=None,
342
block_hashes=unfinished_req.block_hashes,
361
- discard_partial_chunks=self._discard_partial_chunks,
343
)
344
if req_meta is not None:
345
meta.add_request(req_meta)
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+1
-8
@@ -1195,14 +1195,7 @@ class MooncakeStoreWorker:
1195
if load_spec is None or not load_spec.can_load:
1196
continue
1197
1198
- token_len = request.token_len_chunk
1199
- if (load_spec.kvpool_cached_tokens % self.block_size != 0) and (
1200
- load_spec.kvpool_cached_tokens == token_len - 1
1201
- ):
1202
- token_len = load_spec.kvpool_cached_tokens + 1
1203
- else:
1204
- token_len = load_spec.kvpool_cached_tokens
1205
- load_spec.token_len = token_len
1198
+ load_spec.token_len = load_spec.kvpool_cached_tokens
1199
1200
assert self.kv_recv_thread is not None
1201
self.kv_recv_thread.add_request(request)