[Bugfix] Defer offload reads while transfers are pending (#46231)
Signed-off-by: test test <2260891073@qq.com>
Palaiologos1453 committed
Jun 21, 2026 at 19:30 UTC
d3ad8e8bcd1a015026981e479d5537549fba3e97
2 files changed
+106
-4
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+99
-4
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections.abc import Iterable
4
+from types import SimpleNamespace
5
from unittest.mock import MagicMock
6
7
import pytest
@@ -1278,11 +1279,11 @@ def test_reset_cache_finalizes_finished_request_with_pending_store(
1279
)
1280
1281
finalized: list[str] = []
1281
- runner.manager.on_request_finished.side_effect = (
1282
- lambda req_context: finalized.append(req_context.req_id)
1282
+ runner.manager.on_request_finished.side_effect = lambda req_context: (
1283
+ finalized.append(req_context.req_id)
1284
)
1284
- runner.manager.prepare_store.side_effect = (
1285
- lambda keys, req_context: generate_store_output(keys)
1285
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
1286
+ generate_store_output(keys)
1287
)
1288
1289
# Decode a couple of blocks and keep every transfer in flight, so the
@@ -1314,6 +1315,100 @@ def test_reset_cache_finalizes_finished_request_with_pending_store(
1315
assert req_id not in cs._req_status
1316
1317
1318
+def test_pending_transfer_defers_prefix_lookup():
1319
+ """A request with an in-flight store must not issue a load on re-admission.
1320
+
1321
+ With async scheduling, a preempted request's store can be flushed by the
1322
+ worker before the scheduler consumes its completion. If the request is
1323
+ re-admitted in that window, the connector should defer it instead of
1324
+ looking up offloaded blocks and later asserting when a load is queued while
1325
+ the store job is still tracked.
1326
+ """
1327
+ scheduler = object.__new__(OffloadingConnectorScheduler)
1328
+ scheduler.manager = MagicMock(spec=OffloadingManager)
1329
+
1330
+ request = SimpleNamespace(request_id="req-0")
1331
+ group_state = SimpleNamespace(block_ids=[1, 2, 3])
1332
+ req_status = SimpleNamespace(
1333
+ group_states=[group_state],
1334
+ transfer_jobs={123},
1335
+ )
1336
+ scheduler._req_status = {request.request_id: req_status}
1337
+
1338
+ matched_tokens, is_async = scheduler.get_num_new_matched_tokens(
1339
+ request,
1340
+ num_computed_tokens=0,
1341
+ )
1342
+
1343
+ assert matched_tokens is None
1344
+ assert is_async is False
1345
+ assert group_state.block_ids == []
1346
+ scheduler.manager.lookup.assert_not_called()
1347
+
1348
+
1349
+def test_async_preempt_readmit_before_transfer_output_is_deferred(request_runner):
1350
+ """A preempted request can be scheduled again before flush output is read.
1351
+
1352
+ EngineCore.step_with_batch_queue() may schedule a new batch while a prior
1353
+ preemption batch is still queued. The store completion from jobs_to_flush is
1354
+ only cleared when that queued output reaches update_from_output(), so the
1355
+ re-admission path must defer while the scheduler still tracks the store.
1356
+ """
1357
+ block_size = 4
1358
+ block_size_factor = 3
1359
+ offloaded_block_size = block_size * block_size_factor
1360
+
1361
+ runner = request_runner(
1362
+ block_size=block_size,
1363
+ num_gpu_blocks=100,
1364
+ async_scheduling=True,
1365
+ block_size_factor=block_size_factor,
1366
+ )
1367
+ free_block_queue = runner.scheduler.kv_cache_manager.block_pool.free_block_queue
1368
+ num_free_blocks_empty = free_block_queue.num_free_blocks
1369
+
1370
+ req_id = "0"
1371
+ runner.new_request(token_ids=[0] * offloaded_block_size * 2)
1372
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
1373
+ generate_store_output(keys)
1374
+ )
1375
+
1376
+ runner.run(decoded_tokens=[0], complete_transfers=False)
1377
+ runner.run(
1378
+ decoded_tokens=[0] * (2 * offloaded_block_size - block_size),
1379
+ complete_transfers=False,
1380
+ )
1381
+
1382
+ req_status = runner.connector_scheduler._req_status[req_id]
1383
+ pending_store_jobs = set(req_status.transfer_jobs)
1384
+ assert pending_store_jobs
1385
+ assert all(
1386
+ runner.connector_scheduler._jobs[jid].is_store for jid in pending_store_jobs
1387
+ )
1388
+
1389
+ free_block_queue.num_free_blocks = 0
1390
+ preempt_output = runner.scheduler.schedule()
1391
+ assert preempt_output.preempted_req_ids == {req_id}
1392
+ assert preempt_output.kv_connector_metadata is not None
1393
+ assert pending_store_jobs <= preempt_output.kv_connector_metadata.jobs_to_flush
1394
+ assert req_status.transfer_jobs == pending_store_jobs
1395
+
1396
+ # Simulate the async batch-queue window: schedule again before the
1397
+ # preemption batch's ModelRunnerOutput is consumed by update_from_output().
1398
+ free_block_queue.num_free_blocks = num_free_blocks_empty
1399
+ assert runner.scheduler.reset_prefix_cache()
1400
+ runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: len(
1401
+ key
1402
+ )
1403
+
1404
+ readmit_output = runner.scheduler.schedule()
1405
+
1406
+ assert readmit_output.num_scheduled_tokens == {}
1407
+ assert readmit_output.kv_connector_metadata is not None
1408
+ assert readmit_output.kv_connector_metadata.load_jobs == {}
1409
+ assert req_status.transfer_jobs == pending_store_jobs
1410
+
1411
+
1412
@pytest.mark.parametrize("async_scheduling", [True, False])
1413
def test_swa_alignment_skip(request_runner, async_scheduling: bool):
1414
"""SWA blocks unreachable by the load path are skipped during store.
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+7
@@ -647,6 +647,13 @@ class OffloadingConnectorScheduler:
647
for group_state in req_status.group_states:
648
group_state.block_ids.clear()
649
650
+ if req_status.transfer_jobs:
651
+ logger.debug(
652
+ "Delaying request %s since it still has in-flight transfers",
653
+ request.request_id,
654
+ )
655
+ return None, False
656
+
657
req_status.update_offload_keys()
658
req_status.num_locally_computed_tokens = num_computed_tokens
659