[KVConnector] MultiConnector: give every sub-connector the request's real blocks in `update_state_after_alloc` (#46865)
[KVConnector] MultiConnector: give every sub-connector the request's real blocks in `update_state_after_alloc` (#46865) Signed-off-by: deng451e <838677410@qq.com>
weishu committed
Jul 9, 2026 at 11:10 UTC
2285cfca464f9b8e3043df5c6517b4c77ccd0e0e
4 files changed
+37
-27
tests/v1/kv_connector/unit/test_multi_connector.py
+10
-10
@@ -272,7 +272,7 @@ def test_multi_example_connector_consistency():
272
"set_xfer_handshake_metadata_pp_aware",
273
"on_new_request",
274
"get_num_new_matched_tokens 0",
275
- "update_state_after_alloc num_blocks=[0] 0",
275
+ "update_state_after_alloc num_blocks=[7] 0",
276
"build_connector_meta",
277
]
278
# First three events are from initialization (register_kv_caches,
@@ -292,7 +292,7 @@ def test_multi_example_connector_consistency():
292
"set_xfer_handshake_metadata_pp_aware",
293
"on_new_request",
294
"get_num_new_matched_tokens 0",
295
- "update_state_after_alloc num_blocks=[0] 0",
295
+ "update_state_after_alloc num_blocks=[7] 0",
296
"build_connector_meta",
297
]
298
assert events["storage2-WORKER"][:8] == [
@@ -315,9 +315,9 @@ def test_multi_example_connector_consistency():
315
316
events = get_connector_events()
317
# get_num_new_matched_tokens will return new tokens from the first
318
- # connector so update_state_after_alloc will be with allocated blocks
319
- # on that one but with zero blocks for others (first nonzero match is
320
- # chosen).
318
+ # connector (first nonzero match is chosen), so update_state_after_alloc
319
+ # will report those external tokens on that one. Other connectors still
320
+ # receive the request's real blocks but with 0 external tokens.
321
storage1_scheduler_events = _ignore_event_collection(events["storage1-SCHEDULER"])
322
storage2_scheduler_events = _ignore_event_collection(events["storage2-SCHEDULER"])
323
assert storage1_scheduler_events[:4] == [
@@ -329,7 +329,7 @@ def test_multi_example_connector_consistency():
329
assert storage2_scheduler_events[:4] == [
330
"on_new_request",
331
"get_num_new_matched_tokens 0",
332
- "update_state_after_alloc num_blocks=[0] 0",
332
+ "update_state_after_alloc num_blocks=[7] 0",
333
"build_connector_meta",
334
]
335
@@ -345,15 +345,15 @@ def test_multi_example_connector_consistency():
345
346
events = get_connector_events()
347
# get_num_new_matched_tokens will be called for both connectors but will
348
- # return 0 from the first connector, but the second connector should have
349
- # a hit, so update_state_after_alloc will only be called with allocated
350
- # blocks for the second connector.
348
+ # return 0 from the first connector, while the second connector has a hit.
349
+ # Both connectors receive the request's real blocks, but only the chosen
350
+ # (second) connector reports external tokens.
351
storage1_scheduler_events = _ignore_event_collection(events["storage1-SCHEDULER"])
352
storage2_scheduler_events = _ignore_event_collection(events["storage2-SCHEDULER"])
353
assert storage1_scheduler_events[:4] == [
354
"on_new_request",
355
"get_num_new_matched_tokens 0",
356
- "update_state_after_alloc num_blocks=[0] 0",
356
+ "update_state_after_alloc num_blocks=[7] 0",
357
"build_connector_meta",
358
]
359
assert storage2_scheduler_events[:4] == [
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+7
-2
@@ -498,11 +498,16 @@ class KVConnectorBase_V1(ABC):
498
asynchronously loaded into, and second when any additional blocks
499
are allocated, after the load/transfer is complete.
500
501
+ Decide whether to load based on ``num_external_tokens``, not on
502
+ whether ``blocks`` is empty: ``blocks`` may be non-empty even when
503
+ ``num_external_tokens == 0`` (e.g. a non-chosen sub-connector of
504
+ MultiConnector still receives the request's real blocks).
505
+
506
Args:
507
request (Request): the request object.
508
blocks (KVCacheBlocks): the blocks allocated for the request.
504
- num_external_tokens (int): the number of tokens that will be
505
- loaded from the external KV cache.
509
+ num_external_tokens (int): the number of tokens to load from the
510
+ external KV cache. 0 means nothing should be loaded.
511
"""
512
pass
513
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py
+18
-12
@@ -550,17 +550,19 @@ class MoRIIOConnectorScheduler:
550
# remote_engine_id is returned by the prefill's request_finished.
551
# host/ports come from the request_id (parsed in add_new_req).
552
if "remote_engine_id" in params:
553
- # If remote_blocks and num_external_tokens = 0, we have
554
- # a full prefix cache hit on the D worker. We need to call
555
- # send_notify in _read_blocks to free the memory on the P.
556
-
557
- # Get unhashed blocks to pull from remote.
558
- local_block_ids = blocks.get_block_ids()[0]
559
- assert len(local_block_ids) <= len(remote_block_ids)
560
- if len(local_block_ids) == len(remote_block_ids):
561
- pass
553
+ if num_external_tokens > 0:
554
+ # Get unhashed blocks to pull from remote.
555
+ local_block_ids = blocks.get_block_ids()[0]
556
+ assert len(local_block_ids) <= len(remote_block_ids)
557
+ if len(local_block_ids) != len(remote_block_ids):
558
+ local_block_ids = remote_block_ids[
559
+ -len(local_block_ids) :
560
+ ]
561
else:
563
- local_block_ids = remote_block_ids[-len(local_block_ids) :]
562
+ # If remote_blocks and num_external_tokens = 0, we have
563
+ # a full prefix cache hit on the D worker. We need to call
564
+ # send_notify in _read_blocks to free the memory on the P.
565
+ local_block_ids = []
566
567
self._reqs_need_recv[request.request_id] = (
568
request,
@@ -593,7 +595,11 @@ class MoRIIOConnectorScheduler:
595
)
596
remote_notify_port = int(remote_notify_port)
597
596
- block_ids = blocks.get_block_ids()[0]
598
+ # num_external_tokens == 0: nothing to push, so don't tell the
599
+ # producer to write into these blocks.
600
+ block_notify_list = (
601
+ blocks.get_block_ids()[0] if num_external_tokens > 0 else []
602
+ )
603
604
for tp_index in range(self.tp_size):
605
target_port = remote_notify_port + get_port_offset(
@@ -603,7 +609,7 @@ class MoRIIOConnectorScheduler:
609
self.send_notify_block(
610
req_id=request.request_id,
611
transfer_id=request.kv_transfer_params["transfer_id"],
606
- block_notify_list=block_ids,
612
+ block_notify_list=block_notify_list,
613
host=remote_host,
614
port=target_port,
615
)
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+2
-3
@@ -403,14 +403,13 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
403
self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
404
):
405
chosen_connector = self._requests_to_connector.get(request.request_id, -1)
406
- empty_blocks = blocks.new_empty()
406
for i, c in enumerate(self._connectors):
407
if i == chosen_connector:
408
# Forward call to the chosen connector (if any).
409
c.update_state_after_alloc(request, blocks, num_external_tokens)
410
else:
412
- # Call with empty blocks for other connectors.
413
- c.update_state_after_alloc(request, empty_blocks, 0)
411
+ # Other connectors still receive the request's real blocks
412
+ c.update_state_after_alloc(request, blocks, 0)
413
414
def on_new_request(self, request: "Request") -> None:
415
for c in self._connectors: