[Bugfix] Reset num_cached_tokens sentinel on request preemption

[xiaguan]

Description

When a request is preempted, _preempt_request resets num_computed_tokens = 0 but leaves num_cached_tokens at its previous value. On reschedule, the initialization guard:

if request.num_cached_tokens < 0:
    request.num_cached_tokens = num_computed_tokens

never fires again because num_cached_tokens is already ≥ 0. Meanwhile num_external_computed_tokens is re-queried from the KV connector (line 607) and can return a larger value than before (e.g. more blocks became available on a remote store). This breaks the invariant num_cached_tokens >= num_external_computed_tokens and produces a negative local_cache_hit in stats:

local_cache_hit = num_cached_tokens + recomputed - num_external_computed_tokens
                = 1500 + 0 - 2000 = -500  # BOOM

Root cause (_preempt_request, scheduler.py):

# Before fix
request.num_computed_tokens = 0
# num_cached_tokens NOT reset → stale value persists

Fix: reset num_cached_tokens to the sentinel value (-1) so the guard fires on the next schedule and both fields are initialised consistently from fresh values.

Test plan

Added test_preempt_resets_num_cached_tokens in tests/v1/core/test_scheduler.py:

1. Schedule a request with a MockKVConnector → num_cached_tokens is set.
2. Directly call _preempt_request → assert num_cached_tokens == -1.
3. Re-schedule → assert num_cached_tokens >= num_external_computed_tokens (invariant holds).

🤖 Generated with Claude Code

[orozery]

I think that num_cached_tokens is only used once in the request lifetime:

vllm/vllm/v1/metrics/stats.py

Line 344 in a40ee48

if is_prefilling:

Since is_prefilling will not be re-set to True upon preemption.
I actually think this is the desired behavior since we want to be able to characterize the kv cache hits stats per workload, which should be independent of whether preemptions occurred.

[gemini-code-assist]

Code Review

This pull request addresses a bug in the scheduler's preemption logic where num_cached_tokens was not reset, leading to stale values and potentially incorrect stats upon rescheduling. The fix correctly resets this field to its sentinel value of -1 in _preempt_request. The change is accompanied by a well-written and thorough unit test that validates the fix by simulating preemption and ensuring the relevant invariant holds on reschedule. The implementation is correct and effectively resolves the issue.

[xiaguan]

Thanks for the reply! That's a great point about the decode-phase preemption case — I agree is_prefilling correctly stays False there and the stats wouldn't be re-evaluated.

I guess the crash we're seeing is specifically with requests preempted during partial prefill — before they ever complete prefill and produce their first output. In that case, is_prefilling on the output_processor side is still True
(since line 621 req_state.is_prefilling = False was never reached). When the request is rescheduled and eventually completes prefill, PromptTokenStats.update_from_output() does get called, but num_cached_tokens still holds the
value from the original scheduling while num_external_computed_tokens has been refreshed by the connector re-query — which can result in the negative local_cache_hit.

[orozery]

I guess the crash we're seeing is specifically with requests preempted during partial prefill — before they ever complete prefill and produce their first output. In that case, is_prefilling on the output_processor side is still True (since line 621 req_state.is_prefilling = False was never reached). When the request is rescheduled and eventually completes prefill, PromptTokenStats.update_from_output() does get called, but num_cached_tokens still holds the value from the original scheduling while num_external_computed_tokens has been refreshed by the connector re-query — which can result in the negative local_cache_hit.

I see.
I think that num_external_computed_tokens should have the same "set-once" behavior as num_cached_tokens.

[markmc]

xref #34079

[markmc]

#37460 will resolve this, thank you