Repository navigation
Conversation
…ssion cap The recycling-aware admission cap bounds the blocks a sliding window or chunked-local request holds at once, but it was compared with the request's block list, which pads the blocks out of the window with nulls. For a request whose computed prefix is long, the remaining need came out as zero, so the reservations of async KV loads ignored the blocks each prefill chunk still needs in those groups. Loads were admitted until the free pool was smaller than one chunk of the queue head, and with nothing running no block was ever freed. Count the blocks from the first one still in the window, capped, minus the real blocks the request already holds. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Boyang Han <yqszxx@gmail.com>
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment Once the PR is approved or has the If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
Chessing234
left a comment
There was a problem hiding this comment.
AI-assisted review by Codex. No blocking issue found in this static review.
Checked the capped admission arithmetic against the scheduler's in-flight reservation caller, the sliding-window/chunked-local skipped-token rules, prefix-hit eviction accounting, and the uncapped allocation path. Reviewed the added shared/evictable-prefix and remote-load regression cases. The full scheduler tests were not run locally.
|
Second model family for this bug: Gemma 4 31B on 0.31.0, standard What we sawGemma 4 31B (5 sliding-window groups, window 1024, block 16; 1 full-attention group, 32-token pages), Two returning requests with long prefixes in the CPU tier were admitted in one pass. Each took 1,006 blocks (686 full + 5 x 64 sliding). Usage hit 92% with 175 blocks free. Each then needed about 350 blocks for its next 1,024-token chunk. Neither fit. The waiting loop breaks on the first request it cannot allocate, and with Why (0.31.0 line numbers)In
CPU-only reproNo GPU, no model, no connector code: vLLM's own 0.31.0: 0.31.0 with this PR's diff applied to repro.py"""CPU-only repro: sliding-window blocks are not counted at admission for a long async (connector) prefix hit.
Uses only vLLM's KVCacheManager. No GPU, no model. Calls mirror what Scheduler.schedule() does for a waiting
request whose prefix is loaded asynchronously by a KV connector (e.g. OffloadingConnector):
allocate_slots(..., num_external_computed_tokens=hit, delay_cache_blocks=True,
full_sequence_must_fit=True, reserved_blocks=<sum of _request_remaining_blocks>)
KV layout matches Gemma 4 31B: 5 sliding-window groups (block 16, window 1024) and 1 full-attention group
(block 32), 2,188 blocks in the shared pool, max_model_len 32768, --max-num-batched-tokens 1024 with async
scheduling (2 batches in flight -> max_in_flight_tokens 2048).
Run: python repro.py
"""
import torch
from vllm import SamplingParams
from vllm.utils.hashing import sha256
from vllm.v1.core.kv_cache_manager import KVCacheManager
from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, SlidingWindowSpec
from vllm.v1.request import Request, RequestStatus
NUM_BLOCKS = 2188
PROMPT, HIT, CHUNK = 23000, 21952, 1024 # prompt tokens, tokens loaded by the connector, next prefill chunk
init_none_hash(sha256)
swa = [KVCacheGroupSpec([f"swa{g}"], SlidingWindowSpec(block_size=16, num_kv_heads=16, head_size=256,
dtype=torch.bfloat16, sliding_window=1024)) for g in range(5)]
full = KVCacheGroupSpec(["full"], FullAttentionSpec(block_size=32, num_kv_heads=4, head_size=512,
dtype=torch.bfloat16))
cfg = KVCacheConfig(num_blocks=NUM_BLOCKS, kv_cache_tensors=[], kv_cache_groups=swa + [full])
kv = KVCacheManager(cfg, max_model_len=32768, scheduler_block_size=32, hash_block_size=16,
max_in_flight_tokens=2048)
managers = kv.coordinator.single_type_managers
hasher = get_request_block_hasher(16, sha256)
def make_request(rid, seed):
r = Request(rid, [seed * 100000 + i for i in range(PROMPT)], SamplingParams(max_tokens=128), None,
block_hasher=hasher)
r.status = RequestStatus.WAITING
return r
def per_group_count(r, total_computed):
"""The full_sequence_must_fit count (kv_cache_manager.py:519-536), per KV cache group."""
return [m.get_num_blocks_to_allocate(r.request_id, r.num_tokens, [], total_computed, total_computed,
r.num_tokens, apply_admission_cap=True,
prefill_end=r.num_prompt_tokens) for m in managers]
def per_group_held(r):
return [sum(not b.is_null for b in m.req_to_blocks.get(r.request_id, ())) for m in managers]
def remaining(r):
"""Scheduler._request_remaining_blocks (scheduler.py:3003): the reservation for an in-flight async load."""
return kv.coordinator.get_num_blocks_to_allocate(
request_id=r.request_id, num_tokens=r.num_tokens, new_computed_blocks=kv.empty_kv_cache_blocks.blocks,
num_encoder_tokens=0, total_computed_tokens=r.num_computed_tokens,
num_local_computed_tokens=r.num_computed_tokens, num_tokens_main_model=r.num_tokens,
apply_admission_cap=True, prefill_end=r.num_prompt_tokens)
def admit_async_load(r, reserved):
"""The waiting loop for a request with an async external hit (scheduler.py:1194-1230)."""
out = kv.allocate_slots(r, 0, num_external_computed_tokens=HIT, delay_cache_blocks=True,
full_sequence_must_fit=True, reserved_blocks=reserved, has_scheduled_reqs=False)
if out is not None:
r.status, r.num_computed_tokens = RequestStatus.WAITING_FOR_REMOTE_KVS, HIT
return out is not None
groups = [f"swa{g}" for g in range(5)] + ["full"]
cap = managers[0].max_admission_blocks_per_request
print(f"pool {NUM_BLOCKS} blocks; sliding-window admission cap {cap} blocks/group")
print(f"request: {PROMPT} prompt tokens, {HIT} loaded asynchronously by the connector\n")
# (a) the admission count leaves out the sliding-window groups
A = make_request("A", 1)
counted = per_group_count(A, HIT)
free0 = kv.block_pool.get_num_free_blocks()
assert admit_async_load(A, reserved=0)
held = per_group_held(A)
print("(a) request A, admission check vs allocation")
print(f" {'group':6} {'counted':>8} {'allocated':>10}")
for g, c, h in zip(groups, counted, held):
print(f" {g:6} {c:8} {h:10}")
print(f" {'total':6} {sum(counted):8} {free0 - kv.block_pool.get_num_free_blocks():10}")
print(f" A's in-flight reservation (_request_remaining_blocks): {remaining(A)} blocks; "
f"sliding-window part: {sum(per_group_count(A, HIT)[:5])}\n")
# (b) a second async load is admitted next to A; afterwards neither can schedule its next chunk
B = make_request("B", 2)
reserved = remaining(A)
free_b = kv.block_pool.get_num_free_blocks()
ok = admit_async_load(B, reserved)
print(f"(b) request B: free {free_b}, reserved for A {reserved} -> {'ADMITTED' if ok else 'refused'}")
print(f" after both loads: usage {kv.usage:.1%}, free {kv.block_pool.get_num_free_blocks()} blocks")
loaded = [A, B] if ok else [A]
for r in loaded: # loads finish (_update_waiting_for_remote_kv); requests go back to WAITING
kv.cache_blocks(r, r.num_computed_tokens)
r.status = RequestStatus.WAITING
results = []
for r in loaded:
out = kv.allocate_slots(r, CHUNK, full_sequence_must_fit=True, has_scheduled_reqs=False)
results.append(out is not None)
print(f" {r.request_id}: next {CHUNK}-token chunk -> {'allocated' if out is not None else 'REFUSED'}")
if not any(results):
print(" STALL: no waiting request can schedule; with nothing running, nothing is preempted or freed.")Under loadThe stall above came from a Gemma 4 replay through our own async-load offload connector. Since we applied this diff to 0.31.0, we have run 8 replays of about an hour each (7 on Gemma 4 31B, 1 on Qwen3.8-27B; four KV-offload setups, 1x RTX PRO 6000). None stalled on admission. Checked A separate hardening, not needed for this fix: when nothing is running, the waiting loop could preempt or defer a waiting request that holds blocks from a finished async load, so an undercount anywhere cannot wedge the engine (#42568 is related). |
Purpose
Fixes a scheduler deadlock with asynchronous KV loads (KV connectors) on models that have sliding-window or chunked-local KV cache groups.
_request_remaining_blocks(#44560) reserves, for every request whose KV is still loading, the blocks it will need to resume prefill, so that further loads are not admitted into that space. For recycling-aware groups the reservation is capped bymax_admission_blocks_per_request, but the cap was compared withlen(req_to_blocks), and that list pads the positions outside the window with null blocks. For a request with a long computed prefix the length far exceeds the cap, the subtraction comes out negative, and the sliding-window reservation becomes zero, even though every prefill chunk still needs one slot per token in those groups. The non-fast path has the same problem throughnum_skipped_blocks.With the reservation undercounted, loads keep being admitted until the free pool is smaller than one prefill chunk of the queue head. The head cannot allocate,
schedule()breaks on it, and because nothing is running there is nothing to preempt: the engine sits atRunning: 0forever.Observed in serving. DeepSeek-V4-Flash (full-attention plus sliding-window KV groups with block sizes 4 and 8) served with pipeline parallelism, the CPU
OffloadingConnectorand DSpark speculative decoding, with contexts up to 1M tokens. While serving, the engine stopped atRunning: 0, Waiting: 20and never recovered. The scheduler state, dumped from the liveEngineCore:Fix. In the capped branch, count the blocks from the first one still inside the window, capped, and subtract the real blocks the request already holds (the non-null suffix of its block table) and the real blocks among the effective prefix hits. Execution allocation is unchanged; only the reservation and the admission estimate move. The uncapped path keeps its previous arithmetic.
Cost. The capped branch no longer takes the running-request fast path, but
apply_admission_cap=Trueis only used at admission time: by_request_remaining_blockswhen a new async load is admitted, and byallocate_slots(full_sequence_must_fit=True). Per-step allocation for running requests still uses the fast path. Inside the branch, the backward scan stops at the first null block, so it is bounded by the cap (the window plus in-flight tokens, in blocks).Not a duplicate
Searched open PRs and issues touching
max_admission_blocks_per_request,_request_remaining_blocks,get_num_blocks_to_allocateand the sliding-window admission path (2026-10-05):tests/v1/core/test_swa_inflight_window_free.py(out-of-window blocks must stay allocated until the step that reads them has been processed). This PR leaves allocation unchanged and only fixes the reservation.Test Plan
New tests:
tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py::test_async_load_reserves_sliding_window_blocks_for_inflight: full attention plus a sliding-window group (block 16, window 32), pool of 160 blocks, two 1280-token requests each async-loading 1024 tokens. Before the fix both loads are admitted and no tokens are ever scheduled after the loads complete; the test asserts that the scheduler makes progress.tests/v1/core/test_single_type_kv_cache_manager.py::test_admission_accounts_for_prefix_hits: for SWA and chunked-local, with shared hits, evictable hits and plain allocation, the capped estimate must equal the blocks the allocator actually pulls.Regression suites, run on a GPU node with and without the fix (
vllm/reverted, tests kept):End to end: the serving configuration above was restarted with the fix and the workload that triggered the deadlock was replayed for 30 minutes.
Test Result
test_admission_accounts_for_prefix_hits(its two uncapped cases pass on both sides, as the uncapped path is unchanged).AI assistance (Claude Code) was used to investigate, implement, test and draft this change. The submitter reviewed every changed line and ran the validation above.
Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.
🤖 Generated with Claude Code