@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1fc2cee50a
Fork 0

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, 18:15 UTC

[KVConnector][Mooncake] Wire reset_cache cascade end-to-end (#42694)

Signed-off-by: aoshen524 <aoshen524@gmail.com> Signed-off-by: Ao Shen <aoshen@inferact.ai> Co-authored-by: aoshen524 <aoshen524@gmail.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

aoshen02 committed May 27, 2026 at 11:52 UTC 1fc2cee50a09a094b9f2bbdfcb0ab0cadb536712
8 files changed +517 -14
tests/v1/core/test_scheduler.py
+21
@@ -774,6 +774,27 @@ def test_scheduler_reset_prefix_cache():
774 assert scheduler.waiting[i] == request
775
776
777 +def test_reset_connector_cache_no_connector_is_no_op_success():
778 + """``reset_connector_cache`` must return True when no connector is
779 + configured.
780 +
781 + Without this, ``reset_prefix_cache(reset_connector=True)`` returns
782 + ``False`` on every engine that doesn't have a KV connector configured —
783 + even when the local prefix cache reset succeeded — and any caller that
784 + interprets the return value as "did the reset I asked for succeed?"
785 + sees a spurious failure.
786 + """
787 + scheduler = create_scheduler(enable_prefix_caching=True)
788 + assert scheduler.connector is None
789 +
790 + # No-connector reset is treated as success.
791 + assert scheduler.reset_connector_cache() is True
792 +
793 + # End-to-end: reset_prefix_cache(reset_connector=True) on an idle
794 + # scheduler succeeds with or without a connector.
795 + assert scheduler.reset_prefix_cache(reset_connector=True) is True
796 +
797 +
798 # Note - these test cases mirror some of those in test_rejection_sampler.py
799 @pytest.mark.parametrize(
800 "spec_tokens,output_tokens,expected",
tests/v1/kv_connector/unit/test_mooncake_store_connector.py
+330
@@ -11,6 +11,11 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
11 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
12 connector as mooncake_store_connector,
13 )
14 +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
15 + protocol,
16 + scheduler,
17 + worker,
18 +)
19 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
20 MooncakeStoreConnectorMetadata,
21 )
@@ -284,3 +289,328 @@ def test_update_connector_output_and_take_events():
289 assert connector._kv_cache_events is kv_events
290 assert list(connector.take_events()) == [event]
291 assert connector._kv_cache_events is None
292 +
293 +
294 +# ============================================================
295 +# reset_cache() — RL hard-reset path via typed LookupKey protocol
296 +# ============================================================
297 +
298 +
299 +def test_reset_cache_scheduler_role_delegates_to_reset_store():
300 + """SCHEDULER role reset_cache() routes to scheduler.reset_store()."""
301 + vllm_config = _make_vllm_config()
302 + kv_cache_config = _make_kv_cache_config()
303 +
304 + with (
305 + set_current_vllm_config(vllm_config),
306 + patch(
307 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
308 + "connector.MooncakeStoreScheduler"
309 + ) as mock_scheduler_cls,
310 + ):
311 + conn = mooncake_store_connector.MooncakeStoreConnector(
312 + vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
313 + )
314 +
315 + mock_scheduler_cls.return_value.reset_store.return_value = True
316 + assert conn.reset_cache() is True
317 + mock_scheduler_cls.return_value.reset_store.assert_called_once_with()
318 +
319 +
320 +def test_reset_cache_scheduler_role_propagates_failure():
321 + """SCHEDULER role surfaces False when scheduler.reset_store() fails."""
322 + vllm_config = _make_vllm_config()
323 + kv_cache_config = _make_kv_cache_config()
324 +
325 + with (
326 + set_current_vllm_config(vllm_config),
327 + patch(
328 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
329 + "connector.MooncakeStoreScheduler"
330 + ) as mock_scheduler_cls,
331 + ):
332 + conn = mooncake_store_connector.MooncakeStoreConnector(
333 + vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
334 + )
335 +
336 + mock_scheduler_cls.return_value.reset_store.return_value = False
337 + assert conn.reset_cache() is False
338 +
339 +
340 +def test_reset_cache_worker_role_returns_none():
341 + """WORKER role reset_cache() is a no-op; reset is driven via ZMQ admin."""
342 + vllm_config = _make_vllm_config()
343 + kv_cache_config = _make_kv_cache_config()
344 +
345 + with (
346 + set_current_vllm_config(vllm_config),
347 + patch(
348 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
349 + "connector.MooncakeStoreWorker"
350 + ),
351 + ):
352 + conn = mooncake_store_connector.MooncakeStoreConnector(
353 + vllm_config, KVConnectorRole.WORKER, kv_cache_config
354 + )
355 +
356 + assert conn.reset_cache() is None
357 +
358 +
359 +def test_scheduler_reset_store_returns_client_reset_result():
360 + """MooncakeStoreScheduler.reset_store() returns LookupKeyClient.reset()."""
361 + vllm_config = _make_vllm_config()
362 + kv_cache_config = _make_kv_cache_config()
363 +
364 + with (
365 + set_current_vllm_config(vllm_config),
366 + patch(
367 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
368 + "scheduler.LookupKeyClient"
369 + ) as mock_client_cls,
370 + ):
371 + sched = scheduler.MooncakeStoreScheduler(vllm_config, kv_cache_config)
372 +
373 + mock_client_cls.return_value.reset.return_value = True
374 + assert sched.reset_store() is True
375 + mock_client_cls.return_value.reset.assert_called_once_with()
376 +
377 +
378 +def test_scheduler_reset_store_handles_rpc_exception():
379 + """Exceptions from the ZMQ reset RPC convert to False, not raise."""
380 + vllm_config = _make_vllm_config()
381 + kv_cache_config = _make_kv_cache_config()
382 +
383 + with (
384 + set_current_vllm_config(vllm_config),
385 + patch(
386 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
387 + "scheduler.LookupKeyClient"
388 + ) as mock_client_cls,
389 + ):
390 + sched = scheduler.MooncakeStoreScheduler(vllm_config, kv_cache_config)
391 +
392 + mock_client_cls.return_value.reset.side_effect = RuntimeError("rpc timed out")
393 + assert sched.reset_store() is False
394 +
395 +
396 +def test_lookup_key_client_lookup_prepends_typed_tag():
397 + """LookupKeyClient.lookup() puts LOOKUP_MSG tag at frame 0."""
398 + vllm_config = _make_vllm_config()
399 +
400 + with patch(
401 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
402 + "worker.make_zmq_socket"
403 + ) as mock_make_socket:
404 + client = worker.LookupKeyClient(vllm_config)
405 +
406 + fake_socket = mock_make_socket.return_value
407 + fake_socket.recv.return_value = (5).to_bytes(4, "big")
408 +
409 + assert client.lookup(token_len=128, block_hashes=[]) == 5
410 +
411 + sent_frames = fake_socket.send_multipart.call_args[0][0]
412 + assert sent_frames[0] == protocol.LOOKUP_MSG
413 + assert int.from_bytes(sent_frames[1], "big") == 128
414 +
415 +
416 +def test_lookup_key_client_reset_uses_typed_protocol():
417 + """LookupKeyClient.reset() sends RESET_MSG and parses RESP_OK / RESP_ERR."""
418 + vllm_config = _make_vllm_config()
419 +
420 + with patch(
421 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
422 + "worker.make_zmq_socket"
423 + ) as mock_make_socket:
424 + client = worker.LookupKeyClient(vllm_config)
425 +
426 + fake_socket = mock_make_socket.return_value
427 +
428 + # ACK path: server returns RESP_OK -> client returns True.
429 + fake_socket.recv.return_value = protocol.RESP_OK
430 + assert client.reset() is True
431 + assert fake_socket.send.call_args[0][0] == protocol.RESET_MSG
432 +
433 + # NACK path: server returns RESP_ERR -> client returns False.
434 + fake_socket.recv.return_value = protocol.RESP_ERR
435 + assert client.reset() is False
436 +
437 +
438 +def test_protocol_tags_are_distinct_and_non_empty():
439 + """Protocol tags must be unique and non-empty to avoid collision."""
440 + tags = {protocol.LOOKUP_MSG, protocol.RESET_MSG}
441 + assert len(tags) == 2
442 + for tag in tags:
443 + assert isinstance(tag, bytes)
444 + assert len(tag) > 0
445 + assert protocol.RESP_OK != protocol.RESP_ERR
446 +
447 +
448 +def test_scheduler_reset_connector_cache_invokes_connector_reset():
449 + """Cascade test: Scheduler.reset_prefix_cache(reset_connector=True)
450 + cascades into MooncakeStoreConnector.reset_cache without dragging in
451 + the heavy KVCacheManager fixtures.
452 + """
453 + vllm_config = _make_vllm_config()
454 + kv_cache_config = _make_kv_cache_config()
455 +
456 + with (
457 + set_current_vllm_config(vllm_config),
458 + patch(
459 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
460 + "connector.MooncakeStoreScheduler"
461 + ) as mock_scheduler_cls,
462 + ):
463 + conn = mooncake_store_connector.MooncakeStoreConnector(
464 + vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
465 + )
466 +
467 + mock_scheduler_cls.return_value.reset_store.return_value = True
468 +
469 + class _StubScheduler:
470 + def __init__(self, c):
471 + self.connector = c
472 +
473 + def reset_connector_cache(self):
474 + return self.connector.reset_cache() is not False
475 +
476 + sched = _StubScheduler(conn)
477 + assert sched.reset_connector_cache() is True
478 + mock_scheduler_cls.return_value.reset_store.assert_called_once_with()
479 +
480 + mock_scheduler_cls.return_value.reset_store.reset_mock()
481 + mock_scheduler_cls.return_value.reset_store.return_value = False
482 + assert sched.reset_connector_cache() is False
483 +
484 +
485 +def test_reset_cache_scheduler_role_clears_local_state():
486 + """SCHEDULER reset_cache() must clear scheduler-side state that points
487 + at master keys we're about to wipe -- pending load_specs and
488 + accumulated _kv_cache_events both reference keys whose blobs are
489 + about to be remove_all'd, so reading them after reset would surface
490 + stale references to wiped keys.
491 + """
492 + from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( # noqa: E501
493 + LoadSpec,
494 + )
495 +
496 + vllm_config = _make_vllm_config()
497 + kv_cache_config = _make_kv_cache_config()
498 +
499 + with (
500 + set_current_vllm_config(vllm_config),
501 + patch(
502 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
503 + "connector.MooncakeStoreScheduler"
504 + ) as mock_scheduler_cls,
505 + ):
506 + conn = mooncake_store_connector.MooncakeStoreConnector(
507 + vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
508 + )
509 +
510 + # Seed both sentinel pieces of stale-reference state.
511 + sched_inst = mock_scheduler_cls.return_value
512 + sched_inst.load_specs = {
513 + "req-A": LoadSpec(vllm_cached_tokens=0, kvpool_cached_tokens=128, can_load=True)
514 + }
515 + conn._kv_cache_events = mooncake_store_connector.MooncakeStoreKVEvents(
516 + num_workers=1
517 + )
518 + sched_inst.reset_store.return_value = True
519 +
520 + assert conn.reset_cache() is True
521 +
522 + # Both stale references must be cleared by the time reset_store is
523 + # invoked downstream (load_specs flushed dict, events nulled).
524 + assert sched_inst.load_specs == {}
525 + assert conn._kv_cache_events is None
526 +
527 +
528 +def test_lookup_key_server_reset_drains_send_queue_before_remove_all():
529 + """LookupKeyServer RESET handler must drain the send thread's
530 + request_queue BEFORE calling store.remove_all -- otherwise stale
531 + puts that were already in flight when the caller paused generation
532 + can land on the master AFTER remove_all and silently repopulate it
533 + with KV hashed against the previous-policy weights.
534 + """
535 + # Exercise the handler logic directly with mocks for the send thread
536 + # and store. We assert (a) join() is called, (b) remove_all is called,
537 + # and (c) join() comes BEFORE remove_all in the call order. The full
538 + # LookupKeyServer is heavy (binds a real ZMQ REP socket), so we drive
539 + # just the dispatch branch here via a stub equivalent.
540 + from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
541 + protocol,
542 + )
543 +
544 + call_order: list[str] = []
545 +
546 + fake_send_queue = MagicMock()
547 + fake_send_queue.join.side_effect = lambda: call_order.append("join")
548 +
549 + fake_store = MagicMock()
550 + fake_store.remove_all.side_effect = lambda force: call_order.append(
551 + f"remove_all(force={force})"
552 + )
553 +
554 + fake_send_thread = MagicMock()
555 + fake_send_thread.request_queue = fake_send_queue
556 +
557 + fake_store_worker = MagicMock()
558 + fake_store_worker.kv_send_thread = fake_send_thread
559 + fake_store_worker.store = fake_store
560 +
561 + fake_socket = MagicMock()
562 + sent: list[bytes] = []
563 + fake_socket.send.side_effect = lambda frame: sent.append(frame)
564 +
565 + # Mirror the body of LookupKeyServer.process_request RESET_MSG branch.
566 + # Keeping this inline (instead of importing the closure) keeps the
567 + # test independent of the live thread lifecycle.
568 + msg_type = protocol.RESET_MSG
569 + if msg_type == protocol.RESET_MSG:
570 + try:
571 + if fake_store_worker.kv_send_thread is not None:
572 + fake_store_worker.kv_send_thread.request_queue.join()
573 + fake_store_worker.store.remove_all(force=True)
574 + fake_socket.send(protocol.RESP_OK)
575 + except Exception:
576 + fake_socket.send(protocol.RESP_ERR)
577 +
578 + # Drain must happen before remove_all.
579 + assert call_order == ["join", "remove_all(force=True)"]
580 + # Worker reported success.
581 + assert sent == [protocol.RESP_OK]
582 +
583 +
584 +def test_lookup_key_server_reset_skips_drain_when_no_send_thread():
585 + """When the worker has no send thread (e.g. consumer-only role
586 + configurations), the RESET handler must still call remove_all
587 + instead of dereferencing a None send thread.
588 + """
589 + from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
590 + protocol,
591 + )
592 +
593 + call_order: list[str] = []
594 + fake_store = MagicMock()
595 + fake_store.remove_all.side_effect = lambda force: call_order.append("remove_all")
596 +
597 + fake_store_worker = MagicMock()
598 + fake_store_worker.kv_send_thread = None
599 + fake_store_worker.store = fake_store
600 +
601 + fake_socket = MagicMock()
602 + sent: list[bytes] = []
603 + fake_socket.send.side_effect = lambda frame: sent.append(frame)
604 +
605 + msg_type = protocol.RESET_MSG
606 + if msg_type == protocol.RESET_MSG:
607 + try:
608 + if fake_store_worker.kv_send_thread is not None:
609 + fake_store_worker.kv_send_thread.request_queue.join()
610 + fake_store_worker.store.remove_all(force=True)
611 + fake_socket.send(protocol.RESP_OK)
612 + except Exception:
613 + fake_socket.send(protocol.RESP_ERR)
614 +
615 + assert call_order == ["remove_all"]
616 + assert sent == [protocol.RESP_OK]
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py
+17
@@ -200,6 +200,23 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):
200 assert self.connector_scheduler is not None
201 return self.connector_scheduler.request_finished(request, block_ids)
202
203 + def reset_cache(self) -> bool | None:
204 + """Reset the external Mooncake store on prefix-cache reset.
205 +
206 + Drains the worker send queue, then runs ``remove_all`` on the
207 + Mooncake master. Caller must first pause generation (e.g.
208 + ``pause_generation``) so no new puts are enqueued during drain.
209 +
210 + Returns True on ack, False on failure, None for the worker role.
211 + """
212 + if self.role == KVConnectorRole.SCHEDULER:
213 + assert self.connector_scheduler is not None
214 + # Clear local references to keys we're about to wipe.
215 + self.connector_scheduler.load_specs.clear()
216 + self._kv_cache_events = None
217 + return self.connector_scheduler.reset_store()
218 + return None
219 +
220 def update_connector_output(self, connector_output: KVConnectorOutput):
221 kv_cache_events = connector_output.kv_cache_events
222 if not kv_cache_events or not isinstance(
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/protocol.py new
+36
@@ -0,0 +1,36 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Wire-format constants for the LookupKey ZMQ admin channel.
4 +
5 +This is the single source of truth shared by ``LookupKeyClient`` and
6 +``LookupKeyServer`` on the scheduler<->worker rank-0 admin channel.
7 +
8 +Wire format (REQ/REP over IPC):
9 +
10 + Request: [msg_type: bytes] [payload_frames...]
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
15 + Response: [hit_count: u32 big-endian, 4 bytes]
16 +
17 + msg_type == RESET_MSG:
18 + (no payload frames)
19 + Response: [RESP_OK] or [RESP_ERR]
20 +
21 +The first frame of every request is a named bytes tag (not a numeric
22 +sentinel that aliases the data field) so the protocol stays
23 +self-describing and extensible: adding new admin commands requires
24 +only a new tag and a new dispatch branch.
25 +
26 +Mirrors the named-tag convention used by the NIXL connector (see
27 +``vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py``).
28 +"""
29 +
30 +# Request message-type tags. Frame 0 of every request.
31 +LOOKUP_MSG: bytes = b"lookup"
32 +RESET_MSG: bytes = b"reset"
33 +
34 +# Single-byte response status codes for admin commands.
35 +RESP_OK: bytes = b"\x01"
36 +RESP_ERR: bytes = b"\x00"
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py
+27
@@ -371,3 +371,30 @@ class MooncakeStoreScheduler:
371 request.request_id,
372 )
373 return delay_free_blocks, None
374 +
375 + def reset_store(self) -> bool:
376 + """Trigger a global ``remove_all(force=True)`` on the Mooncake master.
377 +
378 + Routes through the existing LookupKey ZMQ admin channel to worker
379 + rank 0, which owns the ``MooncakeDistributedStore`` handle.
380 +
381 + Ordering assumption: caller (typically
382 + ``Scheduler.reset_connector_cache``, invoked via
383 + ``reset_prefix_cache(reset_connector=True)``) MUST ensure no
384 + in-flight Mooncake lookups or transfers. For RL workflows this is
385 + satisfied at the step boundary after weight updates and rollout
386 + drain. Violating this can allow stale KV to be served on the next
387 + request, defeating the hard-reset guarantee.
388 +
389 + Returns True on ACK from worker, False on NACK or RPC error.
390 + """
391 + try:
392 + ok = self.client.reset()
393 + if ok:
394 + logger.info("Mooncake store reset via remove_all succeeded.")
395 + else:
396 + logger.warning("Mooncake store reset returned NACK from worker.")
397 + return ok
398 + except Exception as e:
399 + logger.error("Mooncake reset_store RPC failed: %s", e)
400 + return False
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+64 -10
@@ -50,6 +50,12 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
50 PoolKey,
51 ReqMeta,
52 )
53 +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.protocol import ( # noqa: E501
54 + LOOKUP_MSG,
55 + RESET_MSG,
56 + RESP_ERR,
57 + RESP_OK,
58 +)
59 from vllm.logger import init_logger
60 from vllm.utils.network_utils import get_ip, make_zmq_socket
61 from vllm.v1.core.kv_cache_utils import (
@@ -1424,7 +1430,14 @@ class MooncakeStoreWorker:
1430
1431
1432 class LookupKeyServer:
1427 - """ZMQ server on worker rank 0 for handling prefix lookup queries."""
1433 + """ZMQ server on worker rank 0 for the LookupKey admin channel.
1434 +
1435 + Handles two request types, tagged at frame 0:
1436 + - ``LOOKUP_MSG``: prefix-cache hit query, returns hit count.
1437 + - ``RESET_MSG``: drains the send thread queue, then runs
1438 + ``store.remove_all(force=True)``. Caller must have paused the
1439 + scheduler first.
1440 + """
1441
1442 def __init__(
1443 self,
@@ -1450,13 +1463,37 @@ class LookupKeyServer:
1463 def process_request():
1464 while self.running:
1465 all_frames = self.socket.recv_multipart(copy=False)
1453 - token_len = int.from_bytes(all_frames[0], byteorder="big")
1454 - hash_frames = all_frames[1:]
1455 - hashes_str = self.decoder.decode(hash_frames)
1456 - block_hashes = [BlockHash(bytes.fromhex(s)) for s in hashes_str]
1457 - result = self.store_worker.lookup(token_len, block_hashes)
1458 - response = result.to_bytes(4, "big")
1459 - self.socket.send(response)
1466 + msg_type = bytes(all_frames[0])
1467 +
1468 + if msg_type == LOOKUP_MSG:
1469 + token_len = int.from_bytes(all_frames[1], byteorder="big")
1470 + hash_frames = all_frames[2:]
1471 + hashes_str = self.decoder.decode(hash_frames)
1472 + block_hashes = [BlockHash(bytes.fromhex(s)) for s in hashes_str]
1473 + result = self.store_worker.lookup(token_len, block_hashes)
1474 + self.socket.send(result.to_bytes(4, "big"))
1475 +
1476 + elif msg_type == RESET_MSG:
1477 + try:
1478 + # Drain in-flight puts before wiping the master;
1479 + # otherwise stale puts can repopulate it post-reset.
1480 + # Safe across HMA: store.remove_all wipes the underlying
1481 + # flat key space, clearing every (group_id, hash) entry.
1482 + if self.store_worker.kv_send_thread is not None:
1483 + self.store_worker.kv_send_thread.request_queue.join()
1484 + self.store_worker.store.remove_all(force=True)
1485 + logger.info("Mooncake store reset via remove_all succeeded.")
1486 + self.socket.send(RESP_OK)
1487 + except Exception as e:
1488 + logger.error("Mooncake remove_all failed: %s", e)
1489 + self.socket.send(RESP_ERR)
1490 +
1491 + else:
1492 + logger.warning(
1493 + "LookupKeyServer received unknown msg_type: %r",
1494 + msg_type,
1495 + )
1496 + self.socket.send(RESP_ERR)
1497
1498 self.thread = threading.Thread(target=process_request, daemon=True)
1499 self.thread.start()
@@ -1473,7 +1510,12 @@ class LookupKeyServer:
1510
1511
1512 class LookupKeyClient:
1476 - """ZMQ client for querying prefix cache hits from worker."""
1513 + """ZMQ client for the LookupKey admin channel.
1514 +
1515 + Routes both prefix-cache lookups and admin commands (currently:
1516 + ``reset``) to ``LookupKeyServer`` on worker rank 0. The first frame
1517 + of every request is a named tag from ``protocol.py``.
1518 + """
1519
1520 def __init__(self, vllm_config: VllmConfig):
1521 self.encoder = MsgpackEncoder()
@@ -1490,12 +1532,24 @@ class LookupKeyClient:
1532 hash_strs = [h.hex() for h in block_hashes]
1533 hash_frames = self.encoder.encode(hash_strs)
1534 token_len_bytes = token_len.to_bytes(4, byteorder="big")
1493 - all_frames = [token_len_bytes] + list(hash_frames)
1535 + all_frames = [LOOKUP_MSG, token_len_bytes] + list(hash_frames)
1536 self.socket.send_multipart(all_frames, copy=False)
1537 resp = self.socket.recv()
1538 result = int.from_bytes(resp, "big")
1539 return result
1540
1541 + def reset(self) -> bool:
1542 + """Trigger ``store.remove_all(force=True)`` on worker rank 0.
1543 +
1544 + Ordering assumption: caller MUST ensure no in-flight Mooncake
1545 + lookups or transfers when invoking reset. In RL workflows this
1546 + holds naturally at the step boundary after weight updates and
1547 + rollout drain. Returns True on ACK, False on NACK.
1548 + """
1549 + self.socket.send(RESET_MSG)
1550 + resp = self.socket.recv()
1551 + return bytes(resp) == RESP_OK
1552 +
1553 def close(self):
1554 self.socket.close(linger=0)
1555
vllm/v1/core/sched/scheduler.py
+10 -2
@@ -1945,8 +1945,16 @@ class Scheduler(SchedulerInterface):
1945
1946 def reset_connector_cache(self) -> bool:
1947 if self.connector is None:
1948 - logger.warning("reset_connector called but no KV connector is configured.")
1949 - return False
1948 + # No connector attached -> nothing to reset, treat as success so
1949 + # callers that unconditionally request a connector reset (e.g. as
1950 + # part of a cache-clearing cascade after a weight update) don't
1951 + # see reset_prefix_cache() flip to False purely because they
1952 + # didn't configure a connector.
1953 + logger.debug(
1954 + "reset_connector requested but no KV connector is configured; "
1955 + "treating as no-op success."
1956 + )
1957 + return True
1958
1959 if self.connector.reset_cache() is False:
1960 return False
vllm/v1/engine/core.py
+12 -2
@@ -655,8 +655,18 @@ class EngineCore:
655 # Reset the GPU model runner's encoder cache (physical storage)
656 self.model_executor.reset_encoder_cache()
657
658 - def _reset_caches(self, reset_running_requests=True) -> None:
659 - self.reset_prefix_cache(reset_running_requests=reset_running_requests)
658 + def _reset_caches(
659 + self,
660 + reset_running_requests: bool = True,
661 + reset_connector: bool = True,
662 + ) -> None:
663 + # reset_connector=True so external connectors clear alongside
664 + # local caches, matching the pause_generation(clear_cache=True)
665 + # contract. No-op when no connector is configured.
666 + self.reset_prefix_cache(
667 + reset_running_requests=reset_running_requests,
668 + reset_connector=reset_connector,
669 + )
670 self.reset_mm_cache()
671 self.reset_encoder_cache()
672