[Bugfix][V1] Free out-of-window blocks on the processed-token basis under async scheduling

[Saddss]

Attribution. Core fix is the main contribution of @ivanium (Yifan Qiao), from ivanium/vllm#50 (original commit authorship preserved). The problem/fix write-up below is transcribed from that PR. @Saddss added the empirical verification, the regression analysis, and additional test cases (second commit).

Problem

With async scheduling (now default) or PP, the scheduler runs ahead of the GPU: request.num_computed_tokens is advanced at schedule time. remove_skipped_blocks (SWA / chunked-local / mamba) used this optimistic count as the free boundary, causing two bugs:

1. [Bug] SimpleCPUOffload: residual output garbling under high concurrency — load path lacks cross-stream sync with compute (WAR), symmetric gap left by #46278 #47282 (load-WAR): blocks still covered by an in-flight step's attention window are returned to the pool; a KV-consumer connector reallocates them as load destinations and overwrites them on an unordered stream mid-read.
2. Spec-rejection misfree (silent): rejected spec tokens roll the boundary back below already-nulled blocks, so attention reads null KV. Upstream only patched this locally for Mamba (num_speculative_tokens subtraction); SWA / chunked-local were exposed.

Fix

One mechanism fixes both: track Request.num_in_flight_tokens (scheduled minus processed, structurally symmetric across preemption / spec rejection / failed-KV-load) and free on the settled basis num_computed_tokens - num_in_flight_tokens.

Invariant: a block is returned to the pool only when it is outside the attention window of every unprocessed step and no rollback can make it needed again. Sync scheduling has zero in-flight tokens — behavior unchanged, no gating flag.

Consequential (not an independent fix): the SWA / chunked-local per-request admission bound grows to W - 1 + max_concurrent_batches * max_num_batched_tokens, applied in the single source of truth used by both startup pool sizing and the runtime gate. MambaManager's local subtraction is removed (subsumed exactly, now also covering prefill chunks and PP depth).

Alternatives considered: #47653 (BlockPool fence) and a scheduler-callback defer both fence the free downstream; this PR corrects the lying input instead — no new queues, no BlockPool state, no callbacks. Request-level finish/preempt frees remain behind the #45357 fence (write-after-free, unreachable by a lagged basis).

Tests

• New: tests/v1/core/test_swa_inflight_window_free.py — counter symmetry, out-of-window blocks held until the reader step is processed, sync unchanged, admission cap.
• Regression (all pass): test_deferred_block_free, test_single_type_kv_cache_manager, test_kv_cache_utils, test_async_scheduler, test_prefix_caching, offloading-connector scheduler suite (82), test_scheduler.py failure set byte-identical to clean main (35 pre-existing env failures).

AI assistance was used for this change; every line reviewed by the submitter.

---

Added by @Saddss

Additional test coverage

Extends tests/v1/core/test_swa_inflight_window_free.py (7 tests total, all pass on latest nightly):

• test_chunked_local_free_waits_for_in_flight_step / test_chunked_local_admission_cap_accounts_for_overlapping_batches — chunked-local analogues of the SWA in-flight-free and admission tests (it prunes whole chunks and is exposed to the same WAR).
• test_connector_finish_frees_on_settled_basis — covers the second prune site, Scheduler._connector_finished: the out-of-window prune done before the block table is handed to a KV connector (the simple CPU offload / NIXL store, i.e. the exact [Bug] SimpleCPUOffload: residual output garbling under high concurrency — load path lacks cross-stream sync with compute (WAR), symmetric gap left by #46278 #47282 load-WAR site) must also use the settled basis. Control-verified: reverting _connector_finished to raw num_computed_tokens makes this test fail.

GPU e2e garble repro (the item left "not yet run" above)

Gemma-3-27B-class model, single GPU, simple CPU offload (100 GB), async scheduling:

	stock nightly	with fix
Garbled responses	4.5% (27/600)	0.0% (0/600)

Performance (no offload, KV-bound)

• Normal length (~3.7k tok/req): throughput/latency identical to stock — the reduced reported KV size / max-concurrency from the wider admission bound is cosmetic for this shape.
• Near-max length (~15.5k tok/req): ~34% lower total throughput, +51% duration, +125% mean TTFT — the cost of the (intentional) more conservative admission cap for SWA / chunked-local when overlapping batches are possible. Known trade-off, to be revisited via admission-logic optimization for KV-constrained cases.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[ivanium]

@claude review

[ivanium]

cc @njhill and @tdoublep . Hopefully this makes KV cache manager slightly easier to reason about

[ivanium]

+viz @ZJY0516 and @zhewenl as well

[njhill]

Thanks @ivanium @Saddss, nice fix/simplification, this is an unfortunate bug/omission.

I agree with @ZJY0516 that it doesn't make sense to pass max_concurrent_batches. Within the kv cache manager we only care about the derived "max inflight tokens" so I think we can simplify and in particular not pass through two separate values.

I have pushed a commit here to do that, based on this branch, PTAL: njhill@00c404d. Feel free to cherry-pick it into here if it looks good.

I also renamed the remove_skipped_blocks() num_computed_tokens arg to processed_computed_tokens for clarity.

[njhill]

I think we can remove this comment here altogether since the semantics are already covered in the arg definiton (at least with the changes in my additional commit).

[Saddss]

Thanks @ZJY0516 and @njhill for the review and the suggestion.

@ivanium since this is part of your change, could you take a look at njhill/vllm@00c404d and see what you think? I'm happy to help with the cherry-pick and testing.

[ivanium]

Thanks @ZJY0516 and @njhill for the review and the suggestion.

@ivanium since this is part of your change, could you take a look at njhill/vllm@00c404d and see what you think? I'm happy to help with the cherry-pick and testing.

Thanks! I think the comments and @njhill are great and I cherry picked it. I pushed it directly for simplicity, but feel free to double check and modify/revert them if you find any issues.

[Saddss]

Thanks @ZJY0516 and @njhill for the review and the suggestion.
@ivanium since this is part of your change, could you take a look at njhill/vllm@00c404d and see what you think? I'm happy to help with the cherry-pick and testing.

Thanks! I think the comments and @njhill are great and I cherry picked it. I pushed it directly for simplicity, but feel free to double check and modify/revert them if you find any issues.

Double-checked on my side — re-ran the offload garble repro on the updated branch (with @njhill's max_in_flight_tokens change): 100 GB simple CPU offload + async scheduling, 600 requests → 0 garbled (0/600), 0 errors (same setup previously reproduced it on stock). The full test_swa_inflight_window_free.py suite (7 tests, incl. the connector-finish case) also passes on latest nightly. LGTM on my side — thanks @ZJY0516 @njhill @ivanium!

[ZJY0516]

@Saddss Hi, the failing test e2e-core-1-gpu doesn't fail on main, PTAL

[Saddss]

Hi @ZJY0516, thanks for the catch! Confirmed it's not flaky — fixed and pushed.
It only fails on this PR (not main) precisely because this PR's fix changes when out-of-window blocks get nulled under async scheduling (a block is held one step longer). test_mamba_prefix_cache checks a hard-coded golden of which blocks are held/freed at each step, so once the behavior changes the golden no longer matches — hence [1,1,1,1,1] != [0,1,1,1,1]. The old golden was actually the "async + pre-fix" values; sync/async weren't distinguished before.
What I did: pinned the existing mrv1/mrv2 to async_scheduling=False and added mrv1_async/mrv2_async for the async path, with goldens re-recorded from actual runs per mode.
@njhill Thanks for merging it so quickly! That said, it landed before I could post this, so it'd be great if you could still give the test change a second look when you get a chance. PTAL, thanks! cc @ivanium