[Bugfix] Cap SWA/chunked-local runtime admission to startup pool-sizing bound

[Dao007forever]

Purpose

Fixes #39734 (and the same root cause that motivates #39866 / #40027): on hybrid full + sliding-window models with the pool sized at the startup minimum, runtime admission rejects long-but-fittable prompts because the SWA admission gate budgets the full sequence even though SlidingWindowManager.remove_skipped_blocks recycles out-of-window blocks each chunk. The startup pool sizer already accounts for this; the runtime gate did not. The two formulas drifted, producing scheduler stalls and head-of-line blocking on Gemma-4 31B-class workloads.

This PR makes startup pool sizing and runtime admission use the same recycling-aware bound by introducing a max_admission_blocks_per_request method on SlidingWindowSpec and ChunkedLocalAttentionSpec. Both max_memory_usage_bytes (startup sizer) and get_num_blocks_to_allocate (runtime gate) call it. They cannot drift.

Relationship to existing PRs (per AGENTS.md duplicate-work check)

• [Scheduler] Cap SWA admission budget at sliding_window + chunk_size #39866 (open) — Same observed symptom; different design. [Scheduler] Cap SWA admission budget at sliding_window + chunk_size #39866 caps in the coordinator by adding a separate admission method (get_num_blocks_needed_for_admission) alongside the existing get_num_blocks_to_allocate. The token-cap formula min(num_tokens, sliding_window + max_num_batched_tokens) is duplicated separately from the startup pool-sizing formula in max_memory_usage_bytes, so the two can drift. SWA only.
• Handle capped SWA admission in allocator path #40027 (open, follow-up to [Scheduler] Cap SWA admission budget at sliding_window + chunk_size #39866) — Patches the admission/allocation gap that [Scheduler] Cap SWA admission budget at sliding_window + chunk_size #39866 introduces, where admission succeeds with the capped count but the allocator still requires the uncapped count and raises ValueError: Cannot get 1 free blocks from the pool.

This PR's design avoids that gap by construction:

Aspect	#39866 + #40027	This PR
Where the cap lives	Coordinator wraps the manager call	Inside SlidingWindowManager.get_num_blocks_to_allocate
Admission vs. allocation	Two paths (capped vs. uncapped) → needs #40027 guard	Same path → no gap
Source of truth	Cap formula duplicated in coordinator	max_admission_blocks_per_request on the spec, called from both max_memory_usage_bytes and the manager
Block-level math	Token cap (sliding_window + max_num_batched_tokens)	Block cap matching startup formula (cdiv(sliding_window-1+max_num_batched_tokens, block_size) + 1, capped at max_model_len)
ChunkedLocalAttention	Not covered	Covered by the same mechanism

I'd be happy to fold useful pieces (e.g., the head-of-line repro in #39866's bench numbers) into the discussion, or close in favor of either if the maintainers prefer the other shape.

Changes

• vllm/v1/kv_cache_interface.py — add SlidingWindowSpec.max_admission_blocks_per_request and ChunkedLocalAttentionSpec.max_admission_blocks_per_request. Refactor each spec's max_memory_usage_bytes to call its own admission method, locking startup and runtime to one formula.
• vllm/v1/core/single_type_kv_cache_manager.py — SlidingWindowManager and ChunkedLocalAttentionManager accept max_admission_blocks_per_request and cap num_tokens inside get_num_blocks_to_allocate to max_admission_blocks_per_request * block_size (recycling-aware reservation, not lifetime touches).
• vllm/v1/core/kv_cache_coordinator.py, vllm/v1/core/kv_cache_manager.py, vllm/v1/core/sched/scheduler.py, vllm/v1/simple_kv_offload/manager.py — plumb max_num_batched_tokens through so the spec method can be called at manager-construction time. Defaults to max_model_len when not supplied (preserves the legacy uncapped behavior for direct callers).

Safety argument

Per-request peak real-held blocks ≤ reservation, because SlidingWindowManager.remove_skipped_blocks (and the chunked-local equivalent) is invoked from allocate_slots before each chunk's get_num_blocks_to_allocate, swapping out-of-window blocks for null_block. The admission gate enforces sum(reservations) ≤ pool. Therefore sum(peak_real_held) ≤ pool. Drift between the two formulas would re-introduce #39734 or, worse, mid-prefill OOM — they MUST stay in sync, which is exactly what the single-source-of-truth design guarantees.

Test plan

.venv/bin/python -m pytest tests/v1/core/test_prefix_caching.py tests/v1/core/test_single_type_kv_cache_manager.py -v

New tests:

• test_can_fit_full_sequence_swa_cap_admits_long_prompt — hybrid full+SWA, pool sized at startup minimum, prompt longer than SWA window. Without the cap: 32 (full) + 32 (SWA) = 64 blocks demanded → reject. With the cap: 32 (full) + 13 (SWA) = 45 ≤ pool. Asserts admission succeeds.
• test_can_fit_full_sequence_full_attention_still_gates_oversized — confirms the cap does not weaken the genuine-too-big rejection path on full-attention groups.

Reproduce head-of-line blocking on Gemma3:

.venv/bin/vllm serve google/gemma-3-12b-it \
    --tensor-parallel-size 1 \
    --max-num-seqs 8 \
    --max-num-batched-tokens 8192 \
    --max-model-len 131072 \
    --gpu-memory-utilization 0.25 \
    --enable-prefix-caching \
    --async-scheduling \
    --host 0.0.0.0 --port 8000

GPU KV cache size: 51,120 tokens
Maximum concurrency for 131,072 tokens per request: 1.73x

Send 4 requests of size: [100_000, 5_000, 13_000, 14_000] ("hello" * n).

Before the fix:

Running: 0 reqs, Waiting: 4 reqs

After the fix all 4 went through.

Test results

============================ 65 passed in 12.25s ============================

AI-assisted contribution disclosure

This change was developed with AI assistance (Claude Code). The submitter has reviewed every changed line, runs the test suite locally, and stands behind the design choice and the safety argument above.

---

Signed-off-by: Dao Le Dao007forever@gmail.com

Co-authored-by: Claude

[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.

[gemini-code-assist]

Code Review

This pull request implements a recycling-aware admission cap for Sliding Window Attention (SWA) and Chunked Local Attention to ensure consistency between startup pool sizing and runtime admission gating. By introducing max_admission_blocks_per_request in the attention specifications and overriding get_num_blocks_to_allocate in the corresponding managers, the system now correctly accounts for block recycling during chunked prefill, preventing unnecessary request rejections for long prompts. The changes also include propagating max_num_batched_tokens through the KV cache coordination layers and adding comprehensive tests to verify the new admission logic. I have no feedback to provide.

[njhill]

Thanks @Dao007forever! I just have some minor simplification suggestions, see: njhill@f0c7ac6

[jhaotingc]

Hey @Dao007forever do we need #39866 with your PR? Seems that we don't.
Is there a perf data in for Gemma-4? If there's none it's fine.
TY

[Dao007forever]

Hey @jhaotingc, thanks for reaching out! I don't think we need #39866 since they are doing the same thing, just at different places (coordinator vs manager).

I don't have perf data for Gemma-4, I only have GB200 and don't think it'd be interesting :(.

[besquared]

Thanks @Dao007forever!