[Bugfix] Wait for offload keys before storing chunks (#52923)
Signed-off-by: chuanzhisongshu <49055103+982945902@users.noreply.github.com> Co-authored-by: lishuo121 <lishuo121@jd.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
chuanzhisongshu committed
Sep 3, 2026 at 16:07 UTC
092334c828e1114d6d8fb5828f7b0580e57e4a56
2 files changed
+42
-2
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+39
@@ -378,6 +378,45 @@ def test_abort_queued_request_does_not_build_store_job(
378
assert queued_req_id not in runner.connector_scheduler._req_status
379
380
381
+def test_store_jobs_wait_for_missing_offload_key(request_runner):
382
+ """Allocated chunks must wait until their block hash becomes available."""
383
+ block_size = 4
384
+ runner = request_runner(
385
+ block_size=block_size,
386
+ num_gpu_blocks=8,
387
+ async_scheduling=False,
388
+ )
389
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: None
390
+ runner.new_request(token_ids=[0] * (block_size * 2))
391
+ runner.run(decoded_tokens=[1])
392
+
393
+ req_id = str(runner.req_id)
394
+ req_status = runner.connector_scheduler._req_status[req_id]
395
+ req = req_status.req
396
+ group_state = req_status.group_states[0]
397
+ assert len(group_state.block_ids) >= 2
398
+ assert len(group_state.offload_keys) == 2
399
+
400
+ missing_hash = req.block_hashes.pop()
401
+ group_state.offload_keys.pop()
402
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
403
+ generate_store_output(keys)
404
+ )
405
+ scheduler_output = SimpleNamespace(
406
+ num_scheduled_tokens={req_id: 0}, finished_req_ids=set()
407
+ )
408
+
409
+ jobs = runner.connector_scheduler._build_store_jobs(scheduler_output)
410
+ assert len(jobs) == 1
411
+ assert group_state.next_stored_chunk_idx == 1
412
+
413
+ req.block_hashes.append(missing_hash)
414
+ req_status.update_offload_keys()
415
+ jobs = runner.connector_scheduler._build_store_jobs(scheduler_output)
416
+ assert len(jobs) == 1
417
+ assert group_state.next_stored_chunk_idx == 2
418
+
419
+
420
def test_scheduler_reports_lookup_sync_delay(request_runner):
421
runner = request_runner(
422
block_size=4,
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+3
-2
@@ -385,7 +385,7 @@ class RequestOffloadState:
385
group_state: RequestGroupState,
386
num_offloadable_tokens: int,
387
) -> int:
388
- """Number of allocated leading offloaded chunks eligible for store.
388
+ """Number of allocated and keyed leading chunks eligible for store.
389
390
For eagle/MTP groups the volatile trailing chunk of the offloadable
391
range is excluded while decoding: the draft-layer KV of the last
@@ -405,7 +405,8 @@ class RequestOffloadState:
405
num_allocated_chunks = (
406
len(group_state.block_ids) // self.config.blocks_per_chunk
407
)
408
- return min(num_chunks, num_allocated_chunks)
408
+ num_keyed_chunks = len(group_state.offload_keys)
409
+ return min(num_chunks, num_allocated_chunks, num_keyed_chunks)
410
411
def advance_stored_idx(self, num_offloadable_tokens: int) -> None:
412
# max(): at the prefill->decode transition of a chunk-aligned prompt,