@tduncan / Nvidia-Optimization-Stack-VLLM / commits / dc02934a07

[Bugfix][KV Offload] Stop offloading the final sampled token's KV slot (#54288)

Signed-off-by: almogtavor <almogtavor@gmail.com> Signed-off-by: Almog Tavor <70065337+almogtavor@users.noreply.github.com> Co-authored-by: Or Ozeri <or@ozery.com>

Almog Tavor committed Sep 6, 2026 at 14:22 UTC dc02934a07672eb778a8d819f4e35d0f93ee427b

2 files changed +40 -31

tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py

+36 -30

@@ -371,18 +371,20 @@ def test_scheduler_reports_allocation_failure(request_runner):

371

372 @pytest.mark.parametrize("async_scheduling", [True, False])

373 @pytest.mark.parametrize("prompt_offset", [-1, -2])

374 -def test_last_block_offloaded_at_request_finish(

374 +def test_final_sampled_token_does_not_complete_an_offloaded_block(

375 request_runner, async_scheduling: bool, prompt_offset: int

376 ):

377 - """EOS fills the last block at request finish - verify the final block is stored.

377 + """A block whose last slot holds the final sampled token is not stored.

378

379 - prompt = block_size + prompt_offset tokens → not a full block at schedule time,

380 - so _build_store_jobs creates no store job. After EOS, request_finished

381 - keeps req_status alive so _build_store_jobs can process it on the next step.

379 + That token came out of the previous position's forward pass; its own KV

380 + slot is never written, and under spec decode it holds the first rejected

381 + draft's KV. Storing the block would publish it under a content hash that

382 + the GPU prefix cache itself declines to commit, so a later request with

383 + the same tokens would load the unwritten slot.

384

383 - prompt_offset=-1: EOS fills the block → store job created on next step.

384 - prompt_offset=-2: block remains partial → no store job, cleanup in

385 - _build_store_jobs deletes req_status.

385 + prompt = 2 * block_size + prompt_offset tokens. Block 0 is filled by

386 + prefill and stored either way. At prompt_offset=-1 the EOS fills block 1,

387 + which the finishing step must still decline; at -2 block 1 stays partial.

388 """

389 block_size = 4

390 runner = request_runner(

@@ -390,18 +392,11 @@ def test_last_block_offloaded_at_request_finish(

392 num_gpu_blocks=10,

393 async_scheduling=async_scheduling,

394 )

393 - # prompt = block_size + prompt_offset tokens

394 - runner.new_request(token_ids=[0] * (block_size + prompt_offset))

395 + runner.new_request(token_ids=[0] * (2 * block_size + prompt_offset))

396 runner.manager.prepare_store.side_effect = lambda keys, req_context: (

397 generate_store_output(list(keys))

398 )

398 -

399 - if prompt_offset == -1:

400 - # EOS fills the block, so a store job is created for block 0.

401 - runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=(0,))

402 - else:

403 - # Block remains partial, so no store job is created.

404 - runner.run(decoded_tokens=[EOS_TOKEN_ID])

399 + runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=(0,))

400

401 cs = runner.connector_scheduler

402 # After the full run completes, req_status is cleaned up.

@@ -933,10 +928,13 @@ def test_on_request_finished_not_deferred_until_store_completion(

928 def test_on_request_finished_fires_after_final_block_store(

929 request_runner, async_scheduling: bool

930 ):

936 - """on_request_finished fires after the final-block prepare_store at EOS.

931 + """on_request_finished fires after the final prepare_store of the request.

932

938 - When EOS fills a partial block, request_finished() keeps req_status alive

939 - so _build_store_jobs can create a store job for it on the next step.

933 + request_finished() keeps req_status alive so _build_store_jobs can retry a

934 + block on the next step; the retry must be issued before the manager is told

935 + the request is over. Every attempt made before the request finishes is

936 + declined here, so the store that succeeds is the one from the finishing

937 + step and the ordering has something to order.

938 """

939 block_size = 4

940 runner = request_runner(

@@ -950,22 +948,31 @@ def test_on_request_finished_fires_after_final_block_store(

948 ("on_request_finished", req_context.req_id)

949 )

950

951 + runner.new_request(token_ids=[0] * block_size)

952 + request = runner.scheduler.requests[str(runner.req_id)]

953 +

954 def prepare_store(keys, req_context):

954 - calls.append(("prepare_store", req_context.req_id))

955 - return generate_store_output(keys)

955 + finished = request.is_finished()

956 + calls.append(

957 + ("stored_after_finish" if finished else "prepare_store", req_context.req_id)

958 + )

959 + return generate_store_output(keys) if finished else None

960

961 runner.manager.prepare_store.side_effect = prepare_store

962

959 - runner.new_request(token_ids=[0] * (block_size - 1))

963 runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=(0,))

964

965 req_id = str(runner.req_id)

966 assert calls.count(("on_request_finished", req_id)) == 1, calls

967 + assert calls.count(("stored_after_finish", req_id)) == 1, calls

968

969 finished_idx = calls.index(("on_request_finished", req_id))

966 - prepare_indices = [i for i, c in enumerate(calls) if c == ("prepare_store", req_id)]

967 - assert prepare_indices, calls

968 - assert finished_idx > max(prepare_indices), calls

970 + store_indices = [

971 + i

972 + for i, c in enumerate(calls)

973 + if c in (("prepare_store", req_id), ("stored_after_finish", req_id))

974 + ]

975 + assert finished_idx > max(store_indices), calls

976

977

978 @pytest.mark.parametrize("async_scheduling", [True, False])

@@ -1156,10 +1163,9 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo

1163 touch_calls = runner.manager.touch.call_args_list

1164 assert len(touch_calls) == 6

1165

1159 - # EOS fills the 7th block (offset 6). The extra schedule step processes

1160 - # finished_req_ids and stores block 6 for both groups before the request's

1161 - # GPU blocks are freed.

1162 - runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=(6,))

1166 + # EOS lands in the last slot of the 7th block (offset 6). No forward pass

1167 + # writes that slot, so the finishing step declines the block.

1168 + runner.run(decoded_tokens=[EOS_TOKEN_ID])

1169

1170 runner.scheduler.reset_prefix_cache()

1171

vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py

+4 -1

@@ -1384,7 +1384,10 @@ class OffloadingConnectorScheduler:

1384 if req.status is RequestStatus.FINISHED_ABORTED:

1385 num_tokens_after_batch = req.num_computed_tokens

1386 elif req.is_finished():

1387 - num_tokens_after_batch = req.num_tokens

1387 + # Clamp to the GPU prefix cache's commit point. The final sampled

1388 + # token's slot is never committed (under spec decode it holds a

1389 + # rejected draft's KV), so a block ending there must not be stored.

1390 + num_tokens_after_batch = max(req.num_prompt_tokens, req.num_tokens - 1)

1391 else:

1392 num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]

1393 num_tokens_after_batch = req.num_computed_tokens + num_scheduled_tokens