@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 9e88e969c0

[Perf][KVConnector][Mooncake] Parallelize KV load with a receive-thread pool (#45971)

Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Yifan Qiao committed Jun 24, 2026 at 18:25 UTC 9e88e969c08e11035793e939ea76e9b7e526bc0f

3 files changed +54 -27

tests/v1/kv_connector/unit/test_mooncake_store_worker.py

+5 -2

@@ -4,6 +4,7 @@

4 import json

5 import logging

6 import math

7 +import queue

8 import sys

9 import threading

10 import types

@@ -760,7 +761,7 @@ def test_store_worker_get_block_ids_with_load_errors_delegates_to_recv_thread():

761 recv_thread = MagicMock()

762 recv_thread.get_and_clear_block_ids_with_load_errors.return_value = {3, 4}

763 w = _make_bare_worker()

763 - w.kv_recv_thread = recv_thread

764 + w.kv_recv_threads = [recv_thread]

765

766 assert w.get_block_ids_with_load_errors() == {3, 4}

767 recv_thread.get_and_clear_block_ids_with_load_errors.assert_called_once_with()

@@ -1574,7 +1575,9 @@ def _make_bare_worker(

1575 worker.put_step = 1

1576 worker.enable_kv_events = False

1577 worker.kv_send_thread = None

1577 - worker.kv_recv_thread = None

1578 + worker.kv_recv_threads = []

1579 + worker.num_recv_threads = 1

1580 + worker.recv_request_queue = queue.Queue()

1581 worker.tp_size = 1

1582 worker.num_kv_head = 1

1583 worker.pp_size = 1

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

+40 -25

@@ -359,6 +359,7 @@ class KVTransferThread(threading.Thread):

359 ready_event: threading.Event,

360 name: str,

361 record_operation: Callable[..., None] | None = None,

362 + request_queue: queue.Queue[Any] | None = None,

363 ):

364 super().__init__(daemon=True, name=name)

365 self.store = store

@@ -368,7 +369,7 @@ class KVTransferThread(threading.Thread):

369 self.token_databases = token_databases

370 self._record_operation_cb = record_operation

371 self.done_task_lock = threading.Lock()

371 - self.request_queue: queue.Queue[Any] = queue.Queue()

372 + self.request_queue: queue.Queue[Any] = request_queue or queue.Queue()

373 self.finished_requests: set[str] = set()

374 self.kv_event_lock = threading.Lock()

375 self.kv_events: list[BlockStored] = []

@@ -749,6 +750,7 @@ class KVCacheStoreRecvingThread(KVTransferThread):

750 ready_event: threading.Event,

751 disk_offload_buffer_budget_bytes: int | None = None,

752 record_operation: Callable[..., None] | None = None,

753 + request_queue: queue.Queue[Any] | None = None,

754 ):

755 super().__init__(

756 store,

@@ -758,6 +760,7 @@ class KVCacheStoreRecvingThread(KVTransferThread):

760 ready_event,

761 name="KVCacheStoreRecvingThread",

762 record_operation=record_operation,

763 + request_queue=request_queue,

764 )

765 # _invalid_block_ids can be access by both the Worker and RecvingThread

766 self._invalid_block_ids_lock = threading.Lock()

@@ -1101,7 +1104,10 @@ class MooncakeStoreWorker:

1104 self.enable_kv_events = True

1105

1106 self.kv_send_thread: KVCacheStoreSendingThread | None = None

1104 - self.kv_recv_thread: KVCacheStoreRecvingThread | None = None

1107 + # Pool of load-receive threads

1108 + self.kv_recv_threads: list[KVCacheStoreRecvingThread] = []

1109 + self.num_recv_threads = max(1, envs.VLLM_MOONCAKE_LOAD_RECV_THREADS)

1110 + self.recv_request_queue: queue.Queue[ReqMeta] = queue.Queue()

1111 self.finished_store_req: set[str] = set()

1112 self._kv_connector_stats_lock = threading.Lock()

1113 self.kv_connector_stats = MooncakeStoreConnectorStats()

@@ -1260,19 +1266,30 @@ class MooncakeStoreWorker:

1266 )

1267 self.kv_send_thread.start()

1268

1263 - ready_event_recving = threading.Event()

1264 - self.kv_recv_thread = KVCacheStoreRecvingThread(

1265 - self.store,

1266 - self.coord,

1267 - self.token_dbs,

1268 - self.block_size,

1269 - self.tp_rank,

1270 - ready_event_recving,

1271 - disk_offload_buffer_budget_bytes=self.disk_offload_buffer_budget_bytes,

1272 - record_operation=self._record_kv_connector_operation,

1269 + self.kv_recv_threads = []

1270 + ready_events_recving = []

1271 + for i in range(self.num_recv_threads):

1272 + ready_event_recving = threading.Event()

1273 + recv_thread = KVCacheStoreRecvingThread(

1274 + self.store,

1275 + self.coord,

1276 + self.token_dbs,

1277 + self.block_size,

1278 + self.tp_rank,

1279 + ready_event_recving,

1280 + disk_offload_buffer_budget_bytes=self.disk_offload_buffer_budget_bytes,

1281 + record_operation=self._record_kv_connector_operation,

1282 + request_queue=self.recv_request_queue,

1283 + )

1284 + recv_thread.name = f"KVCacheStoreRecvingThread-{i}"

1285 + recv_thread.start()

1286 + self.kv_recv_threads.append(recv_thread)

1287 + ready_events_recving.append(ready_event_recving)

1288 + for ready_event_recving in ready_events_recving:

1289 + ready_event_recving.wait()

1290 + logger.info(

1291 + "Started %d Mooncake KV-load receive thread(s)", self.num_recv_threads

1292 )

1274 - self.kv_recv_thread.start()

1275 - ready_event_recving.wait()

1293

1294 def start_load_kv(

1295 self,

@@ -1306,9 +1323,7 @@ class MooncakeStoreWorker:

1323 continue

1324

1325 load_spec.token_len = load_spec.kvpool_cached_tokens

1309 -

1310 - assert self.kv_recv_thread is not None

1311 - self.kv_recv_thread.add_request(request)

1326 + self.recv_request_queue.put(request)

1327

1328 assert self.load_async, "load_async must be True for better performance."

1329 # Issue stores with CUDA event synchronization

@@ -1335,11 +1350,10 @@ class MooncakeStoreWorker:

1350 else set()

1351 )

1352

1338 - done_recving = (

1339 - self.kv_recv_thread.get_and_clear_finished_requests()

1340 - if self.load_async and self.kv_recv_thread is not None

1341 - else set()

1342 - )

1353 + done_recving: set[str] = set()

1354 + if self.load_async:

1355 + for recv_thread in self.kv_recv_threads:

1356 + done_recving |= recv_thread.get_and_clear_finished_requests()

1357

1358 logger.debug(

1359 "Completed send: %d, recv: %d, tp_rank: %d",

@@ -1350,9 +1364,10 @@ class MooncakeStoreWorker:

1364 return done_sending, done_recving

1365

1366 def get_block_ids_with_load_errors(self) -> set[int]:

1353 - if self.kv_recv_thread is None:

1354 - return set()

1355 - return self.kv_recv_thread.get_and_clear_block_ids_with_load_errors()

1367 + block_ids: set[int] = set()

1368 + for recv_thread in self.kv_recv_threads:

1369 + block_ids |= recv_thread.get_and_clear_block_ids_with_load_errors()

1370 + return block_ids

1371

1372 def _record_kv_connector_operation(

1373 self,

vllm/envs.py

@@ -200,6 +200,7 @@ if TYPE_CHECKING:

200 VLLM_NIXL_SIDE_CHANNEL_PORT: int = 5600

201 VLLM_MOONCAKE_BOOTSTRAP_PORT: int = 8998

202 VLLM_MOONCAKE_STORE_TIER_LOG: bool = False

203 + VLLM_MOONCAKE_LOAD_RECV_THREADS: int = 1

204 VLLM_MOONCAKE_DISK_STAGING_USABLE_RATIO: float = 0.9

205 MOONCAKE_PREFERRED_SEGMENT: str | None = None

206 MOONCAKE_REQUESTER_LOCAL_HOSTNAME: str | None = None

@@ -1531,6 +1532,14 @@ environment_variables: dict[str, Callable[[], Any]] = {

1532 "VLLM_MOONCAKE_STORE_TIER_LOG": lambda: (

1533 os.getenv("VLLM_MOONCAKE_STORE_TIER_LOG", "False").lower() in ("true", "1")

1534 ),

1535 + # Number of parallel KV-load receive threads per worker rank. Lets the

1536 + # per-request control overhead (Python prep + master key lookup) of one

1537 + # request overlap with the RDMA transfer of another, keeping the transfer

1538 + # engine's queue pairs busy. Helps when that overhead is significant or

1539 + # per-request batches are too small to saturate the link on their own.

1540 + "VLLM_MOONCAKE_LOAD_RECV_THREADS": lambda: int(

1541 + os.getenv("VLLM_MOONCAKE_LOAD_RECV_THREADS", "1")

1542 + ),

1543 # Fraction of the owner's DirectIO staging buffer to fill per GET batch.

1544 "VLLM_MOONCAKE_DISK_STAGING_USABLE_RATIO": lambda: float(

1545 os.getenv("VLLM_MOONCAKE_DISK_STAGING_USABLE_RATIO", "0.9")