@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 3b3d5287fa

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 1, 18:45 UTC

[BugFix] Resolve multiple async kv load deadlock (#44560)

Signed-off-by: Nick Hill <nickhill123@gmail.com>

Nick Hill committed Jun 6, 2026 at 16:05 UTC 3b3d5287faee2d4cf84a93419fe28a1f65c93749

6 files changed +214 -1

tests/v1/core/test_scheduler.py

+78

@@ -4474,3 +4474,81 @@ def test_ec_connector_pending_prefetch_only_checks_future_mm_features():

4474 f"Expected only {HASH_FUTURE!r} from future mm feature filtering, "

4475 f"got {future_hashes!r}. Past/boundary features must be filtered out."

4476 )

4477 +

4478 +

4479 +def test_async_load_reservation_prevents_wedge_e2e():

4480 + """Same wedge scenario as PR #40968's lateral-preemption e2e test, but

4481 + resolved by reservation-based admission control instead of preemption.

4482 +

4483 + A (8 blocks) and B (5 blocks) both want an async KV load, sharing a 4-block

4484 + prefix, in a 10-block pool (9 usable). Admitting both loads would wedge:

4485 + once their recvs finish neither can complete its local prefill (8+5 > 9).

4486 +

4487 + Here the reservation gate refuses to admit B's load while A's full sequence

4488 + is still reserved, so B never holds blocks and A is free to complete - no

4489 + deadlock, and (unlike lateral preemption) B is never preempted.

4490 + """

4491 + BLOCK_SIZE = 16

4492 + A_TOKENS = BLOCK_SIZE * 8 # bigger request

4493 + B_TOKENS = BLOCK_SIZE * 5 # smaller request

4494 + MATCHED_TOKENS = BLOCK_SIZE * 4 # 4-block prefix loaded for both

4495 + NUM_BLOCKS = 10 # 9 usable; both prefixes fit, but not both full sequences

4496 +

4497 + scheduler = create_scheduler(

4498 + block_size=BLOCK_SIZE,

4499 + num_blocks=NUM_BLOCKS,

4500 + max_num_seqs=4,

4501 + max_num_batched_tokens=A_TOKENS * 2,

4502 + use_kv_connector=mock_kv(matched_tokens=MATCHED_TOKENS, is_async=True),

4503 + )

4504 +

4505 + [a] = create_requests(

4506 + num_requests=1, num_tokens=A_TOKENS, block_size=BLOCK_SIZE, req_ids=["a"]

4507 + )

4508 + [b] = create_requests(

4509 + num_requests=1, num_tokens=B_TOKENS, block_size=BLOCK_SIZE, req_ids=["b"]

4510 + )

4511 + scheduler.add_request(a)

4512 + scheduler.add_request(b)

4513 +

4514 + EMPTY_OUTPUT = ModelRunnerOutput(

4515 + req_ids=[],

4516 + req_id_to_index={},

4517 + sampled_token_ids=[],

4518 + logprobs=None,

4519 + prompt_logprobs_dict={},

4520 + pooler_output=[],

4521 + )

4522 +

4523 + req_to_blocks = scheduler.kv_cache_manager.coordinator.single_type_managers[

4524 + 0

4525 + ].req_to_blocks

4526 +

4527 + # Step 1: A's load is admitted; B's is held back by the reservation (B never

4528 + # holds blocks, so the wedge precondition - both holding prefixes - is gone).

4529 + out1 = scheduler.schedule()

4530 + assert a.status == RequestStatus.WAITING_FOR_REMOTE_KVS

4531 + assert a.num_computed_tokens == MATCHED_TOKENS

4532 + assert b.status == RequestStatus.WAITING

4533 + assert b.request_id not in req_to_blocks

4534 + assert len(scheduler.running) == 0

4535 + scheduler.update_from_output(out1, EMPTY_OUTPUT)

4536 +

4537 + # Step 2: nothing changes until A's recv lands.

4538 + out2 = scheduler.schedule()

4539 + assert len(scheduler.running) == 0

4540 + a_finished = dataclasses.replace(

4541 + EMPTY_OUTPUT,

4542 + kv_connector_output=KVConnectorOutput(finished_recving=[a.request_id]),

4543 + )

4544 + scheduler.update_from_output(out2, a_finished)

4545 +

4546 + # Step 3: A makes forward progress straight to RUNNING - no preemption was

4547 + # needed because B never wedged it.

4548 + out3 = scheduler.schedule()

4549 + assert a.status == RequestStatus.RUNNING

4550 + assert a in scheduler.running

4551 + assert a.request_id in {req.req_id for req in out3.scheduled_new_reqs}

4552 + assert b.status == RequestStatus.WAITING

4553 + assert b.num_preemptions == 0

4554 + assert b.request_id not in req_to_blocks

tests/v1/e2e/general/test_mamba_prefix_cache.py

@@ -180,6 +180,7 @@ def get_fake_allocate_slots_fn(original_allocate_slots_fn: Callable):

180 delay_cache_blocks: bool = False,

181 num_encoder_tokens: int = 0,

182 full_sequence_must_fit: bool = False,

183 + reserved_blocks: int = 0,

184 ):

185 ret = original_allocate_slots_fn(

186 self,

@@ -192,6 +193,7 @@ def get_fake_allocate_slots_fn(original_allocate_slots_fn: Callable):

193 delay_cache_blocks,

194 num_encoder_tokens,

195 full_sequence_must_fit,

196 + reserved_blocks,

197 )

198 if cur_step_action is not None:

199 cur_block_ids = self.coordinator.single_type_managers[0].req_to_blocks[

tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py

+79

@@ -655,3 +655,82 @@ def test_p_side_chunked_prefill_mamba(mock_platform):

655 outputs = engine_core_outputs[0].outputs

656 assert len(outputs) == 1

657 assert outputs[0].finish_reason == FinishReason.LENGTH

658 +

659 +

660 +def test_async_load_reserves_blocks_for_inflight():

661 + """A second async KV-connector load is not admitted if its initial

662 + allocation would consume blocks reserved for an already in-flight sequence.

663 +

664 + req_a gets a 1-block prefix (full sequence = 4 blocks), reserving 3 more.

665 + req_b would need a 4-block initial allocation, but only

666 + (free - req_a's 3-block reservation) = 3 blocks are available to it, so it is

667 + held back in WAITING (holding no blocks) rather than wedging req_a.

668 + """

669 + vllm_config = create_vllm_config()

670 + BLOCK_SIZE = vllm_config.cache_config.block_size

671 + scheduler = create_scheduler(vllm_config, num_blocks=8) # usable = 7

672 +

673 + req_a = create_request(

674 + request_id=1,

675 + block_size=BLOCK_SIZE,

676 + num_tokens=BLOCK_SIZE * 4,

677 + do_remote_prefill=True,

678 + num_remote_blocks=1,

679 + )

680 + req_b = create_request(

681 + request_id=2,

682 + block_size=BLOCK_SIZE,

683 + num_tokens=BLOCK_SIZE * 5,

684 + do_remote_prefill=True,

685 + num_remote_blocks=1,

686 + )

687 + scheduler.add_request(req_a)

688 + scheduler.add_request(req_b)

689 +

690 + # Partial external matches: req_a loads 1 block, req_b loads 4 blocks.

691 + with patch.object(

692 + scheduler.connector,

693 + "get_num_new_matched_tokens",

694 + side_effect=[(BLOCK_SIZE, True), (BLOCK_SIZE * 4, True)],

695 + ):

696 + scheduler.schedule()

697 +

698 + assert req_a.status == RequestStatus.WAITING_FOR_REMOTE_KVS

699 + assert req_b.status == RequestStatus.WAITING

700 +

701 + req_to_blocks = scheduler.kv_cache_manager.coordinator.single_type_managers[

702 + 0

703 + ].req_to_blocks

704 + assert req_a.request_id in req_to_blocks

705 + assert req_b.request_id not in req_to_blocks

706 +

707 +

708 +def test_async_loads_both_admitted_when_pool_fits():

709 + """Sanity: with a pool large enough, the reservation gate admits both async

710 + loads (it is not over-conservative)."""

711 + vllm_config = create_vllm_config()

712 + BLOCK_SIZE = vllm_config.cache_config.block_size

713 + scheduler = create_scheduler(vllm_config, num_blocks=64)

714 +

715 + reqs = [

716 + create_request(

717 + request_id=i,

718 + block_size=BLOCK_SIZE,

719 + num_tokens=BLOCK_SIZE * 5,

720 + do_remote_prefill=True,

721 + num_remote_blocks=1,

722 + )

723 + for i in (1, 2)

724 + ]

725 + for req in reqs:

726 + scheduler.add_request(req)

727 +

728 + with patch.object(

729 + scheduler.connector,

730 + "get_num_new_matched_tokens",

731 + side_effect=[(BLOCK_SIZE, True), (BLOCK_SIZE, True)],

732 + ):

733 + scheduler.schedule()

734 +

735 + for req in reqs:

736 + assert req.status == RequestStatus.WAITING_FOR_REMOTE_KVS

tests/v1/kv_connector/unit/utils.py

@@ -56,6 +56,7 @@ def assert_scheduler_empty(scheduler: Scheduler):

56 assert len(scheduler.running) == 0

57 assert len(scheduler.finished_req_ids) == 0

58 assert len(scheduler.finished_recving_kv_req_ids) == 0

59 + assert len(scheduler._inflight_prefills) == 0

61 # EncoderCacheManager.

62 assert len(scheduler.encoder_cache_manager.freed) == 0

vllm/v1/core/kv_cache_manager.py

+8 -1

@@ -246,6 +246,7 @@ class KVCacheManager:

246 delay_cache_blocks: bool = False,

247 num_encoder_tokens: int = 0,

248 full_sequence_must_fit: bool = False,

249 + reserved_blocks: int = 0,

250 ) -> KVCacheBlocks | None:

251 """Add slots for a request with new tokens to append.

252

@@ -271,6 +272,11 @@ class KVCacheManager:

272 free blocks to hold the full sequence, accounting for prefix cache hits

273 and sliding window. Used as an admission gate to prevent over-admitting

274 requests when chunked prefill would otherwise only check the first chunk

275 + reserved_blocks: Number of free blocks that must be left available for

276 + other in-flight sequences to complete. The actual allocation is only

277 + made if it fits within (free blocks - reserved_blocks). Used to gate

278 + async KV-connector loads so their initial allocation cannot consume

279 + blocks an already in-flight (prefilling) sequence is relying on.

280

281 Blocks layout:

282 ```

@@ -386,7 +392,8 @@ class KVCacheManager:

392 num_tokens_main_model=num_tokens_main_model,

393 )

394

389 - if num_blocks_to_allocate > self.block_pool.get_num_free_blocks():

395 + available_blocks = self.block_pool.get_num_free_blocks() - reserved_blocks

396 + if num_blocks_to_allocate > available_blocks:

397 # Cannot allocate new blocks

398 return None

399

vllm/v1/core/sched/scheduler.py

+46

@@ -286,6 +286,10 @@ class Scheduler(SchedulerInterface):

286

287 self._pause_state: PauseState = PauseState.UNPAUSED

288

289 + # In-flight requests still prefilling (prefill chunks + in-progress

290 + # async KV loads). Their remaining-block reservation gates async loads.

291 + self._inflight_prefills: set[Request] = set()

292 +

293 def _mamba_block_aligned_split(

294 self,

295 request: Request,

@@ -746,6 +750,14 @@ class Scheduler(SchedulerInterface):

750 for i in encoder_inputs_to_schedule

751 )

752

753 + reserved_blocks = 0

754 + if load_kv_async:

755 + # An async load holds its blocks for the whole transfer with

756 + # no forward progress and isn't preemptible here. Admit it

757 + # only if it fits in (free - other in-flight reservations), to

758 + # avoid deadlock and predictable preemptions.

759 + reserved_blocks = self._inflight_prefill_reserved_blocks()

760 +

761 new_blocks = self.kv_cache_manager.allocate_slots(

762 request,

763 num_new_tokens,

@@ -756,6 +768,7 @@ class Scheduler(SchedulerInterface):

768 delay_cache_blocks=load_kv_async,

769 num_encoder_tokens=num_encoder_tokens,

770 full_sequence_must_fit=self.scheduler_reserve_full_isl,

771 + reserved_blocks=reserved_blocks,

772 )

773

774 if new_blocks is None:

@@ -807,6 +820,7 @@ class Scheduler(SchedulerInterface):

820 # _update_waiting_for_remote_kv will then cache

821 # only the successfully loaded tokens.

822 request.num_computed_tokens = num_computed_tokens

823 + self._inflight_prefills.add(request)

824 continue

825

826 self.running.append(request)

@@ -830,6 +844,9 @@ class Scheduler(SchedulerInterface):

844 token_budget -= num_new_tokens

845 request.status = RequestStatus.RUNNING

846 request.num_computed_tokens = num_computed_tokens

847 + # Only track requests that will still be prefilling after this chunk.

848 + if num_computed_tokens + num_new_tokens < request.num_tokens:

849 + self._inflight_prefills.add(request)

850 # Encoder-related.

851 if encoder_inputs_to_schedule:

852 scheduled_encoder_inputs[request_id] = encoder_inputs_to_schedule

@@ -965,6 +982,7 @@ class Scheduler(SchedulerInterface):

982 )

983 self.kv_cache_manager.free(request)

984 self.encoder_cache_manager.free(request)

985 + self._inflight_prefills.discard(request)

986 request.status = RequestStatus.PREEMPTED

987 request.num_computed_tokens = 0

988 if request.spec_token_ids:

@@ -996,6 +1014,9 @@ class Scheduler(SchedulerInterface):

1014 scheduler_output.has_structured_output_requests |= (

1015 request.use_structured_output and not request.is_prefill_chunk

1016 )

1017 + # Drop from the in-flight-prefill set once it's no longer prefilling.

1018 + if not request.is_prefill_chunk:

1019 + self._inflight_prefills.discard(request)

1020

1021 # Snapshot block IDs for routed experts before forward starts.

1022 # A concurrent schedule() may preempt requests and free blocks

@@ -1869,6 +1890,7 @@ class Scheduler(SchedulerInterface):

1890 ) -> dict[str, Any] | None:

1891 assert request.is_finished()

1892

1893 + self._inflight_prefills.discard(request)

1894 connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request)

1895 self.encoder_cache_manager.free(request)

1896 request_id = request.request_id

@@ -2105,6 +2127,30 @@ class Scheduler(SchedulerInterface):

2127

2128 return self.connector.request_finished_all_groups(request, block_ids)

2129

2130 + def _request_remaining_blocks(self, request: Request) -> int:

2131 + """Blocks `request` still needs to allocate to hold its full sequence."""

2132 + full_num_tokens = min(request.num_tokens, self.max_model_len)

2133 + return self.kv_cache_manager.coordinator.get_num_blocks_to_allocate(

2134 + request_id=request.request_id,

2135 + num_tokens=full_num_tokens,

2136 + new_computed_blocks=self.kv_cache_manager.empty_kv_cache_blocks.blocks,

2137 + num_encoder_tokens=0,

2138 + total_computed_tokens=request.num_computed_tokens,

2139 + num_tokens_main_model=full_num_tokens,

2140 + apply_admission_cap=True,

2141 + )

2142 +

2143 + def _inflight_prefill_reserved_blocks(self) -> int:

2144 + """Blocks in-flight prefills still need to finish (their reservation).

2145 +

2146 + Sums remaining full-ISL blocks over `self._inflight_prefills` (running

2147 + prefills + in-progress async loads). The candidate async load isn't yet

2148 + in the set, so it's naturally excluded.

2149 + """

2150 + return sum(

2151 + self._request_remaining_blocks(req) for req in self._inflight_prefills

2152 + )

2153 +

2154 def _update_waiting_for_remote_kv(self, request: Request) -> None:

2155 """

2156 KV Connector: update request state after async recv is finished.