[Bugfix] Free out-of-window KV blocks on the processed-token basis

[ivanium]

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 vllm-project/vllm#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: vllm-project#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 vllm-project#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).
• Not yet run: GPU e2e garble repro (async + SWA + SimpleCPUOffload).

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

🤖 Generated with Claude Code