[Core][KV Connector] Resolve connector block tables for every scheduled request (#54853)
Signed-off-by: Zhewen Li <zhewenli@inferact.ai> Co-authored-by: Zhewen Li <zhewenli@inferact.ai> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: Codex <noreply@openai.com>
Zhewen Li committed
Sep 8, 2026 at 19:15 UTC
0b066293f3c738a0cbd3a087bf893f2f4dcd61f2
6 files changed
+105
-31
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+8
-4
@@ -226,7 +226,8 @@ def test_partial_tail_store_uses_attention_and_recurrent_cow_sources():
226
227
output = SimpleNamespace(
228
kv_connector_block_state=KVConnectorBlockState(
229
- block_ids={},
229
+ req_ids=set(),
230
+ resolve_block_ids={}.__getitem__,
231
boundary_state_offloads={"req": [(1, 99, 28)]},
232
)
233
)
@@ -285,7 +286,8 @@ def test_aligned_boundary_store_uses_exact_source_with_partial_tail():
286
287
output = SimpleNamespace(
288
kv_connector_block_state=KVConnectorBlockState(
288
- block_ids={},
289
+ req_ids=set(),
290
+ resolve_block_ids={}.__getitem__,
291
boundary_state_offloads={"req": [(1, 98, 16), (1, 99, 28)]},
292
)
293
)
@@ -310,7 +312,8 @@ def test_aligned_boundary_store_flushes_before_cow_destination_reuse():
312
313
output = SchedulerOutput.make_empty()
314
output.kv_connector_block_state = KVConnectorBlockState(
313
- block_ids={},
315
+ req_ids=set(),
316
+ resolve_block_ids={}.__getitem__,
317
boundary_state_offloads={"req": [(1, 99, 16)]},
318
)
319
meta = scheduler.build_connector_meta(output)
@@ -4495,7 +4498,8 @@ class TestMambaHybridOffloadServing:
4498
num_scheduled_tokens={"A": self.PROMPT_TOKENS},
4499
finished_req_ids=set(),
4500
kv_connector_block_state=KVConnectorBlockState(
4498
- block_ids={},
4501
+ req_ids=set(),
4502
+ resolve_block_ids={}.__getitem__,
4503
boundary_state_offloads={"A": [(1, 101, self.MAMBA_BLOCK)]},
4504
),
4505
)
tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py
+3
-1
@@ -57,8 +57,10 @@ def _make_connector_block_state(
57
block_ids: tuple[list[int], ...] | None = None,
58
offloads: list[tuple[int, int, int]] | None = None,
59
) -> KVConnectorBlockState:
60
+ tables = {} if block_ids is None else {"req-0": block_ids}
61
return KVConnectorBlockState(
61
- block_ids={} if block_ids is None else {"req-0": block_ids},
62
+ req_ids=set(tables),
63
+ resolve_block_ids=tables.__getitem__,
64
boundary_state_offloads=({} if offloads is None else {"req-0": offloads}),
65
)
66
tests/v1/kv_connector/unit/test_scheduler_kv_connector_override.py
+76
-8
@@ -6,6 +6,7 @@ import pytest
6
7
import vllm.plugins as plugins_module
8
from tests.v1.core.utils import create_requests, create_scheduler
9
+from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
10
from vllm.distributed.kv_transfer.kv_connector.factory import (
11
KVConnectorFactory,
12
)
@@ -15,9 +16,10 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
16
)
17
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
18
from vllm.v1.core.kv_cache_utils import BlockHash
18
-from vllm.v1.core.sched.output import SchedulerOutput
19
+from vllm.v1.core.sched.output import KVConnectorBlockState, SchedulerOutput
20
from vllm.v1.core.sched.scheduler import Scheduler
21
from vllm.v1.kv_cache_interface import KVCacheConfig
22
+from vllm.v1.outputs import ModelRunnerOutput
23
from vllm.v1.request import Request
24
25
@@ -54,7 +56,11 @@ class DummyKVConnector(KVConnectorBase_V1):
56
)
57
block_state = scheduler_output.kv_connector_block_state
58
assert block_state is not None
57
- block_ids_by_req = block_state.block_ids
59
+ block_ids_by_req = {}
60
+ for req_id in block_state.req_ids:
61
+ block_ids = block_state.get_block_ids(req_id)
62
+ assert block_ids is not None
63
+ block_ids_by_req[req_id] = block_ids
64
return DummyConnectorMetadata(
65
block_hashes_by_req=block_hashes_by_req,
66
block_ids_by_req=block_ids_by_req,
@@ -75,11 +81,12 @@ class DummyKVConnector(KVConnectorBase_V1):
81
82
def _my_plugin():
83
"""Registers the dummy KV connector and overrides _build_kv_connector_meta"""
78
- KVConnectorFactory.register_connector(
79
- "DummyKVConnector",
80
- __name__,
81
- DummyKVConnector.__name__,
82
- )
84
+ if "DummyKVConnector" not in KVConnectorFactory._registry:
85
+ KVConnectorFactory.register_connector(
86
+ "DummyKVConnector",
87
+ __name__,
88
+ DummyKVConnector.__name__,
89
+ )
90
91
def _custom_build_kv_connector_meta(
92
self, connector: KVConnectorBase_V1, scheduler_output: SchedulerOutput
@@ -126,7 +133,8 @@ def test_connector_receives_block_hashes(_load_plugin):
133
134
output = scheduler.schedule()
135
129
- # Verify the connector metadata was built with block hashes.
136
+ # Exercise worker serialization without opening sockets or shared memory.
137
+ MessageQueue(n_reader=0, n_local_reader=0).enqueue(output)
138
meta = output.kv_connector_metadata
139
assert isinstance(meta, DummyConnectorMetadata)
140
assert len(meta.block_hashes_by_req) == 3
@@ -144,3 +152,63 @@ def test_connector_receives_block_hashes(_load_plugin):
152
}
153
assert meta.block_ids_by_req == expected_block_ids
154
assert output.kv_connector_block_state is None
155
+
156
+
157
+def _make_model_runner_output(scheduler: Scheduler) -> ModelRunnerOutput:
158
+ return ModelRunnerOutput(
159
+ req_ids=[req.request_id for req in scheduler.running],
160
+ req_id_to_index={req.request_id: i for i, req in enumerate(scheduler.running)},
161
+ sampled_token_ids=[[1000]] * len(scheduler.running),
162
+ logprobs=None,
163
+ prompt_logprobs_dict={},
164
+ pooler_output=[],
165
+ )
166
+
167
+
168
+def test_connector_block_state_covers_steps_without_new_blocks(_load_plugin):
169
+ """A store save lands on the step that fills a block, which allocates none.
170
+
171
+ The connector must still be able to look up that request's current table.
172
+ """
173
+ block_size = 16
174
+ scheduler = create_scheduler(
175
+ use_kv_connector="DummyKVConnector", block_size=block_size
176
+ )
177
+ (request,) = create_requests(
178
+ num_requests=1, num_tokens=3 * block_size - 1, block_size=block_size
179
+ )
180
+ scheduler.add_request(request)
181
+
182
+ output = scheduler.schedule()
183
+ scheduler.update_from_output(output, _make_model_runner_output(scheduler))
184
+
185
+ # The first decode token fills the third block: no allocation this step.
186
+ output = scheduler.schedule()
187
+ assert output.scheduled_cached_reqs.req_ids == [request.request_id]
188
+ [new_block_ids] = output.scheduled_cached_reqs.new_block_ids
189
+ assert not new_block_ids
190
+
191
+ meta = output.kv_connector_metadata
192
+ assert isinstance(meta, DummyConnectorMetadata)
193
+ block_ids = meta.block_ids_by_req.get(request.request_id)
194
+ assert block_ids is not None
195
+ assert block_ids == scheduler.kv_cache_manager.get_block_ids(request.request_id)
196
+ assert len(block_ids[0]) == 3
197
+
198
+
199
+def test_connector_block_state_resolves_only_offered_requests():
200
+ tables = {"a": ([1, 2],), "b": ([3],)}
201
+ resolve_block_ids = MagicMock(side_effect=tables.__getitem__)
202
+ block_state = KVConnectorBlockState(
203
+ req_ids={"a"},
204
+ resolve_block_ids=resolve_block_ids,
205
+ boundary_state_offloads={},
206
+ )
207
+
208
+ assert block_state.get_block_ids("b") is None
209
+ resolve_block_ids.assert_not_called()
210
+ assert block_state.get_block_ids("a") == ([1, 2],)
211
+ resolve_block_ids.assert_called_once_with("a")
212
+
213
+ tables["a"] = ([4, 5],)
214
+ assert block_state.get_block_ids("a") == ([4, 5],)
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py
+1
-1
@@ -434,7 +434,7 @@ class MooncakeStoreScheduler:
434
"Current block tables are required for Mooncake store jobs"
435
)
436
for req_meta in save_metas:
437
- block_ids = block_state.block_ids.get(req_meta.req_id)
437
+ block_ids = block_state.get_block_ids(req_meta.req_id)
438
assert block_ids is not None, (
439
f"Missing current block table for store request {req_meta.req_id}"
440
)
vllm/v1/core/sched/output.py
+10
-2
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from collections.abc import Callable
5
from dataclasses import dataclass
6
from functools import cached_property
7
from typing import TYPE_CHECKING
@@ -209,11 +210,18 @@ class ScheduledEncoderInputStats:
210
class KVConnectorBlockState:
211
"""Scheduler-local block state offered to a producer-side KV connector."""
212
212
- # Authoritative current block-table snapshots.
213
- block_ids: dict[str, tuple[list[int], ...]]
213
+ # Requests scheduled this step and requests with a boundary-state hand-off.
214
+ req_ids: set[str]
215
+ # Resolve on access to avoid copying tables the connector never reads.
216
+ resolve_block_ids: Callable[[str], tuple[list[int], ...]]
217
# Exact Mamba "align" boundary-state hand-offs.
218
boundary_state_offloads: dict[str, list[tuple[int, int, int]]]
219
220
+ def get_block_ids(self, req_id: str) -> tuple[list[int], ...] | None:
221
+ if req_id not in self.req_ids:
222
+ return None
223
+ return self.resolve_block_ids(req_id)
224
+
225
226
@dataclass
227
class SchedulerOutput:
vllm/v1/core/sched/scheduler.py
+7
-15
@@ -1359,24 +1359,16 @@ class Scheduler(SchedulerInterface):
1359
1360
kv_connector_block_state = None
1361
if self.connector is not None:
1362
- snapshot_req_ids = {req.req_id for req in new_reqs_data}
1363
- snapshot_req_ids.update(
1364
- req_id
1365
- for req_id, block_ids in zip(
1366
- cached_reqs_data.req_ids,
1367
- cached_reqs_data.new_block_ids,
1368
- strict=True,
1369
- )
1370
- if block_ids
1371
- )
1372
- snapshot_req_ids.update(
1362
+ # Any request scheduled this step can become a connector job now,
1363
+ # not only the ones that were allocated blocks: a store save lands
1364
+ # on the step that fills a block, which allocated none.
1365
+ block_state_req_ids = set(num_scheduled_tokens)
1366
+ block_state_req_ids.update(
1367
req_id for req_id in boundary_state_offloads if req_id in self.requests
1368
)
1369
kv_connector_block_state = KVConnectorBlockState(
1376
- block_ids={
1377
- req_id: self.kv_cache_manager.get_block_ids(req_id)
1378
- for req_id in snapshot_req_ids
1379
- },
1370
+ req_ids=block_state_req_ids,
1371
+ resolve_block_ids=self.kv_cache_manager.get_block_ids,
1372
boundary_state_offloads=boundary_state_offloads,
1373
)
1374