@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 35e4dd4a69

[KV Connector][Mooncake] Async lookup to reduce scheduler overhead (#45659)

Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>

Yifan Qiao committed Jun 18, 2026 at 14:44 UTC 35e4dd4a69b6b95feb74866341daa46c3836aed0

6 files changed +197 -11

docs/features/mooncake_store_connector_usage.md

@@ -203,6 +203,7 @@ the vLLM JSON config.

203 ### kv_connector_extra_config

204

205 - `load_async` (bool): Enable asynchronous loading for better compute-I/O overlap. Default: `true`.

206 +- `lookup_async` (bool): Run the external prefix-cache lookup on a background thread so it never blocks the scheduler step. The request is held until the in-flight lookup completes, then resumed on a later step. Default: `false`.

207 - `enable_cross_layers_blocks` (bool): Enable cross-layer block packing for reduced store operations. Default: `false`.

208 - `lookup_rpc_port` (int): Custom port for the ZMQ lookup RPC socket. Default: `0`.

209 - `cache_prefix` (str): Namespace prepended to every store key. Lets separate deployments share one Mooncake master without polluting each other — instances configured with different prefixes never see each other's cached blocks, even for identical prompts. All instances that should share a prefix cache must use the same value. Default: `""` (no prefix; keys are byte-identical to the unprefixed format).

tests/v1/kv_connector/unit/test_mooncake_store_connector.py

+126 -1

@@ -1,6 +1,8 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 +import threading

5 +import time

6 from unittest.mock import MagicMock, patch

8 from vllm.config import set_current_vllm_config

@@ -406,7 +408,9 @@ def test_lookup_key_client_lookup_prepends_typed_tag():

408 fake_socket = mock_make_socket.return_value

409 fake_socket.recv.return_value = (5).to_bytes(4, "big")

410

409 - assert client.lookup(token_len=128, block_hashes=[]) == 5

411 + # Blocking lookup (non_block defaults to False) runs on the executor and

412 + # returns the resolved hit length.

413 + assert client.lookup("req0", token_len=128, block_hashes=[]) == 5

414

415 sent_frames = fake_socket.send_multipart.call_args[0][0]

416 assert sent_frames[0] == protocol.LOOKUP_MSG

@@ -435,6 +439,127 @@ def test_lookup_key_client_reset_uses_typed_protocol():

439 assert client.reset() is False

440

441

442 +def _poll_lookup(client, req_id, token_len=128, block_hashes=(), timeout=5.0):

443 + """Drive non-blocking lookup until the executor completes it."""

444 + deadline = time.monotonic() + timeout

445 + while time.monotonic() < deadline:

446 + result = client.lookup(req_id, token_len, list(block_hashes), non_block=True)

447 + if result is not None:

448 + return result

449 + time.sleep(0.005)

450 + return None

451 +

452 +

453 +def _gated_recv(gate: threading.Event, value: int):

454 + """Mock recv side-effect that blocks until ``gate`` is set, so the

455 + executor's lookup can be held pending deterministically."""

456 +

457 + def recv():

458 + gate.wait()

459 + return value.to_bytes(4, "big")

460 +

461 + return recv

462 +

463 +

464 +def test_lookup_key_client_non_block_lookup_async():

465 + """Non-blocking lookup defers to the executor: None first, hit once the

466 + Future resolves."""

467 + vllm_config = _make_vllm_config()

468 +

469 + with patch(

470 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."

471 + "worker.make_zmq_socket"

472 + ) as mock_make_socket:

473 + client = worker.LookupKeyClient(vllm_config)

474 +

475 + fake_socket = mock_make_socket.return_value

476 + # Hold the executor's lookup pending until we release the gate.

477 + gate = threading.Event()

478 + fake_socket.recv.side_effect = _gated_recv(gate, 7)

479 +

480 + # First query submits the lookup and returns None while it is in flight.

481 + assert client.lookup("req1", 128, [], non_block=True) is None

482 + # Release the executor; a later poll returns the hit length.

483 + gate.set()

484 + assert _poll_lookup(client, "req1") == 7

485 + # Future is consumed (popped) on read.

486 + assert "req1" not in client.futures

487 +

488 +

489 +def test_lookup_key_client_discard_clears_state():

490 + """discard() drops a completed lookup Future so it is not served stale."""

491 + vllm_config = _make_vllm_config()

492 +

493 + with patch(

494 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."

495 + "worker.make_zmq_socket"

496 + ) as mock_make_socket:

497 + client = worker.LookupKeyClient(vllm_config)

498 +

499 + fake_socket = mock_make_socket.return_value

500 + gate = threading.Event()

501 + fake_socket.recv.side_effect = _gated_recv(gate, 9)

502 +

503 + # Submit while gated so the call returns None and the Future stays in

504 + # `futures` (unconsumed) once it resolves.

505 + assert client.lookup("req2", 128, [], non_block=True) is None

506 + gate.set()

507 + deadline = time.monotonic() + 5.0

508 + while time.monotonic() < deadline:

509 + if client.futures["req2"].done():

510 + break

511 + time.sleep(0.005)

512 + # discard() drops the completed result before any lookup consumes it.

513 + client.discard("req2")

514 + assert "req2" not in client.futures

515 + # A fresh query re-submits rather than returning a stale value: hold the

516 + # gate so the resubmitted lookup stays in flight.

517 + gate.clear()

518 + assert client.lookup("req2", 128, [], non_block=True) is None

519 + gate.set() # release the executor so the worker thread can drain

520 +

521 +

522 +def test_get_num_new_matched_tokens_async_defers_then_reports():

523 + """Async lookup returns (None, False) until ready, then the hit count."""

524 + vllm_config = create_vllm_config(

525 + kv_connector="MooncakeStoreConnector",

526 + kv_role="kv_both",

527 + kv_connector_extra_config={"lookup_async": True},

528 + )

529 + kv_cache_config = _make_kv_cache_config()

530 +

531 + with (

532 + set_current_vllm_config(vllm_config),

533 + patch(

534 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."

535 + "scheduler.LookupKeyClient"

536 + ) as mock_client_cls,

537 + ):

538 + sched = scheduler.MooncakeStoreScheduler(vllm_config, kv_cache_config)

539 +

540 + assert sched.lookup_async is True

541 + mock_client = mock_client_cls.return_value

542 +

543 + block_size = sched._block_size

544 + request = MagicMock()

545 + request.request_id = "r1"

546 + request.num_tokens = 4 * block_size

547 + request.block_hashes = []

548 +

549 + # Lookup not ready -> defer.

550 + mock_client.lookup.return_value = None

551 + assert sched.get_num_new_matched_tokens(request, 0) == (None, False)

552 + assert "r1" not in sched.load_specs

553 +

554 + # Lookup ready with a hit -> report need_to_allocate + async-load flag.

555 + hit = 3 * block_size

556 + mock_client.lookup.return_value = hit

557 + need, load_async = sched.get_num_new_matched_tokens(request, 0)

558 + assert need == hit

559 + assert load_async == sched.load_async

560 + assert sched.load_specs["r1"].kvpool_cached_tokens == hit

561 +

562 +

563 def test_protocol_tags_are_distinct_and_non_empty():

564 """Protocol tags must be unique and non-empty to avoid collision."""

565 tags = {protocol.LOOKUP_MSG, protocol.RESET_MSG}

tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py

+8 -1

@@ -16,6 +16,7 @@ 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.lookup_async = False

20 scheduler._block_size = 16

21 scheduler.load_specs = {}

22 scheduler._preempted_req_ids = set()

@@ -405,7 +406,13 @@ class _StubLookupClient:

406 def __init__(self, hit_tokens: int) -> None:

407 self._hit_tokens = hit_tokens

408

408 - def lookup(self, token_len: int, block_hashes: list[bytes]) -> int:

409 + def lookup(

410 + self,

411 + req_id: str,

412 + token_len: int,

413 + block_hashes: list[bytes],

414 + non_block: bool = False,

415 + ) -> int:

416 return self._hit_tokens

417

418

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py

+1 -1

@@ -176,7 +176,7 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):

176 self,

177 request: Request,

178 num_computed_tokens: int,

179 - ) -> tuple[int, bool]:

179 + ) -> tuple[int | None, bool]:

180 assert self.connector_scheduler is not None

181 return self.connector_scheduler.get_num_new_matched_tokens(

182 request, num_computed_tokens

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py

+19 -6

@@ -54,9 +54,9 @@ class MooncakeStoreScheduler:

54 ):

55 assert vllm_config.kv_transfer_config is not None

56 self.kv_role = vllm_config.kv_transfer_config.kv_role

57 - self.load_async = vllm_config.kv_transfer_config.kv_connector_extra_config.get(

58 - "load_async", True

59 - )

57 + kvc_extra_config = vllm_config.kv_transfer_config.kv_connector_extra_config

58 + self.load_async = kvc_extra_config.get("load_async", True)

59 + self.lookup_async = kvc_extra_config.get("lookup_async", False)

60 self.client = LookupKeyClient(vllm_config)

62 # Align with the engine's own scheduler_block_size and hash_block_size.

@@ -75,14 +75,26 @@ class MooncakeStoreScheduler:

75 self,

76 request: Request,

77 num_computed_tokens: int,

78 - ) -> tuple[int, bool]:

79 - """Check for external KV cache hit."""

78 + ) -> tuple[int | None, bool]:

79 + """Check for external KV cache hit.

80 +

81 + Returns ``(None, False)`` when an async lookup is still in flight,

82 + signaling the scheduler to retry this request on a later step.

83 + """

84 # Look up against the full prefill range, not just the prompt.

85 token_len = request.num_tokens // self._block_size * self._block_size

86 if token_len < self._block_size:

87 return 0, False

85 - num_external_hit_tokens = self.client.lookup(token_len, request.block_hashes)

89 + num_external_hit_tokens = self.client.lookup(

90 + request.request_id,

91 + token_len,

92 + request.block_hashes,

93 + non_block=self.lookup_async,

94 + )

95 + if num_external_hit_tokens is None:

96 + # Lookup not ready yet; scheduler will retry on a later step.

97 + return None, False

99 if num_external_hit_tokens == request.num_tokens:

100 # Leave a sub-block tail uncomputed for sampling, on a block

@@ -158,6 +170,7 @@ class MooncakeStoreScheduler:

170 force_skip_save = self.kv_role == "kv_consumer"

171

172 for finished_req_id in scheduler_output.finished_req_ids:

173 + self.client.discard(finished_req_id)

174 self.load_specs.pop(finished_req_id, None)

175 self._request_trackers.pop(finished_req_id, None)

176 self._unfinished_requests.pop(finished_req_id, None)

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py

+42 -2

@@ -19,6 +19,7 @@ import threading

19 import time

20 from collections import defaultdict

21 from collections.abc import Callable

22 +from concurrent.futures import Future, ThreadPoolExecutor

23 from dataclasses import dataclass

24 from typing import Any, Literal, TypeVar

@@ -1560,7 +1561,13 @@ class LookupKeyClient:

1561 bind=False,

1562 )

1563

1563 - def lookup(self, token_len: int, block_hashes: list[BlockHash]) -> int:

1564 + # Async lookup support

1565 + self.executor = ThreadPoolExecutor(

1566 + max_workers=1, thread_name_prefix="MooncakeLookupClient"

1567 + )

1568 + self.futures: dict[str, Future[int]] = {}

1569 +

1570 + def _lookup(self, token_len: int, block_hashes: list[BlockHash]) -> int:

1571 hash_strs = [h.hex() for h in block_hashes]

1572 hash_frames = self.encoder.encode(hash_strs)

1573 token_len_bytes = token_len.to_bytes(4, byteorder="big")

@@ -1570,7 +1577,36 @@ class LookupKeyClient:

1577 result = int.from_bytes(resp, "big")

1578 return result

1579

1573 - def reset(self) -> bool:

1580 + def lookup(

1581 + self,

1582 + req_id: str,

1583 + token_len: int,

1584 + block_hashes: list[BlockHash],

1585 + non_block: bool = False,

1586 + ) -> int | None:

1587 + """If non_block is True, will return None until the result is ready,

1588 + so the caller retries on a later step."""

1589 + future = self.futures.get(req_id)

1590 + if future is None:

1591 + future = self.executor.submit(self._lookup, token_len, list(block_hashes))

1592 + self.futures[req_id] = future

1593 + if non_block and not future.done():

1594 + return None

1595 + try:

1596 + return future.result()

1597 + except Exception as e:

1598 + logger.error("Async Mooncake lookup failed for %s: %s", req_id, e)

1599 + return 0

1600 + finally:

1601 + del self.futures[req_id]

1602 +

1603 + def discard(self, req_id: str) -> None:

1604 + """Drop any cached/in-flight lookup for ``req_id`` (e.g. on abort)."""

1605 + future = self.futures.pop(req_id, None)

1606 + if future is not None:

1607 + future.cancel()

1608 +

1609 + def _reset(self) -> bool:

1610 """Trigger ``store.remove_all(force=True)`` on worker rank 0.

1611

1612 Ordering assumption: caller MUST ensure no in-flight Mooncake

@@ -1582,7 +1618,11 @@ class LookupKeyClient:

1618 resp = self.socket.recv()

1619 return bytes(resp) == RESP_OK

1620

1621 + def reset(self) -> bool:

1622 + return self.executor.submit(self._reset).result()

1623 +

1624 def close(self):

1625 + self.executor.shutdown(wait=False, cancel_futures=True)

1626 self.socket.close(linger=0)

1627

1628