[BugFix][KV Connector] Fix Deadlock with KVConnector + MTP under KV Pressure (#57104)
Signed-off-by: Robert Shaw <robertgshaw2@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
Robert Shaw committed
Sep 16, 2026 at 18:22 UTC
6c3d9816afbb3e5499a573248d1cf6fcb00c05cd
3 files changed
+92
-5
tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py
+60
@@ -731,3 +731,63 @@ def test_async_loads_both_admitted_when_pool_fits():
731
732
for req in reqs:
733
assert req.status == RequestStatus.WAITING_FOR_REMOTE_KVS
734
+
735
+
736
+def test_async_load_reserves_blocks_for_promotion_margin():
737
+ """An async load is not admitted unless the blocks its own promotion will
738
+ need are still free.
739
+
740
+ A parked load is allocated without lookahead slots, but promotion pads it
741
+ to ``1 + num_spec_tokens`` and asks for the lookahead margin on top. Without
742
+ reserving that margin, two loads can be admitted that together consume the
743
+ whole pool; the head then fails ``allocate_slots`` forever while the load
744
+ behind it is never reached, and with nothing running no block is ever freed.
745
+
746
+ req_a (4 blocks) and req_b (3 blocks) exactly fill the 7 usable blocks, so
747
+ req_b must be held back in WAITING holding no blocks, leaving req_a room to
748
+ be promoted and run.
749
+ """
750
+ vllm_config = create_vllm_config(num_speculative_tokens=3)
751
+ BLOCK_SIZE = vllm_config.cache_config.block_size
752
+ scheduler = create_scheduler(vllm_config, num_blocks=8) # usable = 7
753
+
754
+ req_a = create_request(
755
+ request_id=1,
756
+ block_size=BLOCK_SIZE,
757
+ num_tokens=BLOCK_SIZE * 4,
758
+ do_remote_prefill=True,
759
+ max_tokens=1,
760
+ )
761
+ req_b = create_request(
762
+ request_id=2,
763
+ block_size=BLOCK_SIZE,
764
+ num_tokens=BLOCK_SIZE * 3,
765
+ do_remote_prefill=True,
766
+ max_tokens=1,
767
+ )
768
+ scheduler.add_request(req_a)
769
+ scheduler.add_request(req_b)
770
+
771
+ # Both get a full external hit, so each would hold its whole prompt.
772
+ with patch.object(
773
+ scheduler.connector,
774
+ "get_num_new_matched_tokens",
775
+ side_effect=[(BLOCK_SIZE * 4, True), (BLOCK_SIZE * 3, True)],
776
+ ):
777
+ scheduler.schedule()
778
+
779
+ assert req_a.status == RequestStatus.WAITING_FOR_REMOTE_KVS
780
+ assert req_b.status == RequestStatus.WAITING
781
+ req_to_blocks = scheduler.kv_cache_manager.coordinator.single_type_managers[
782
+ 0
783
+ ].req_to_blocks
784
+ assert req_b.request_id not in req_to_blocks
785
+
786
+ # req_a's load lands: it must be promotable and actually get scheduled.
787
+ scheduler.update_from_output(
788
+ scheduler.schedule(),
789
+ create_model_runner_output([], finished_recving={req_a.request_id}),
790
+ )
791
+ scheduler_output = scheduler.schedule()
792
+ assert req_a.status == RequestStatus.RUNNING
793
+ assert scheduler_output.num_scheduled_tokens[req_a.request_id] > 0
tests/v1/kv_connector/unit/utils.py
+8
@@ -17,6 +17,7 @@ from vllm.config import (
17
KVTransferConfig,
18
ModelConfig,
19
SchedulerConfig,
20
+ SpeculativeConfig,
21
VllmConfig,
22
)
23
from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
@@ -106,6 +107,7 @@ def create_vllm_config(
107
kv_connector_module_path: str | None = None,
108
kv_role: str = "kv_consumer",
109
disable_hybrid_kv_cache_manager: bool | None = None,
110
+ num_speculative_tokens: int | None = None,
111
) -> VllmConfig:
112
"""Initialize VllmConfig For Testing."""
113
model_config = ModelConfig(
@@ -141,6 +143,11 @@ def create_vllm_config(
143
kv_load_failure_policy=kv_load_failure_policy,
144
)
145
attention_config = AttentionConfig(backend=attention_backend)
146
+ speculative_config = (
147
+ SpeculativeConfig(model="ngram", num_speculative_tokens=num_speculative_tokens)
148
+ if num_speculative_tokens is not None
149
+ else None
150
+ )
151
return VllmConfig(
152
scheduler_config=scheduler_config,
153
model_config=model_config,
@@ -148,6 +155,7 @@ def create_vllm_config(
155
kv_transfer_config=kv_transfer_config,
156
device_config=DeviceConfig("cpu"),
157
attention_config=attention_config,
158
+ speculative_config=speculative_config,
159
)
160
161
vllm/v1/core/sched/scheduler.py
+24
-5
@@ -31,6 +31,7 @@ from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
31
from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
32
from vllm.multimodal.encoder_budget import MultiModalBudget
33
from vllm.multimodal.utils import get_mm_features_in_window
34
+from vllm.utils.math_utils import cdiv
35
from vllm.v1.core.encoder_cache_manager import (
36
EncoderCacheManager,
37
EncoderDecoderCacheManager,
@@ -1140,9 +1141,13 @@ class Scheduler(SchedulerInterface):
1141
if load_kv_async:
1142
# An async load holds its blocks for the whole transfer with
1143
# no forward progress and isn't preemptible here. Admit it
1143
- # only if it fits in (free - other in-flight reservations), to
1144
- # avoid deadlock and predictable preemptions.
1145
- reserved_blocks = self._inflight_prefill_reserved_blocks()
1144
+ # only if it fits in (free - other in-flight reservations)
1145
+ # plus its own spec decode step blocks, to avoid deadlock and
1146
+ # predictable preemptions.
1147
+ reserved_blocks = (
1148
+ self._inflight_prefill_reserved_blocks()
1149
+ + self._spec_decode_step_blocks()
1150
+ )
1151
1152
new_blocks = self.kv_cache_manager.allocate_slots(
1153
request,
@@ -2892,9 +2897,9 @@ class Scheduler(SchedulerInterface):
2897
return delay_free or partial_tail_delay, kv_xfer_params
2898
2899
def _request_remaining_blocks(self, request: Request) -> int:
2895
- """Blocks `request` still needs to allocate to hold its full sequence."""
2900
+ """Blocks `request` still needs to hold its full sequence and be promoted."""
2901
full_num_tokens = min(request.num_tokens, self.max_model_len)
2897
- return self.kv_cache_manager.coordinator.get_num_blocks_to_allocate(
2902
+ num_blocks = self.kv_cache_manager.coordinator.get_num_blocks_to_allocate(
2903
request_id=request.request_id,
2904
num_tokens=full_num_tokens,
2905
new_computed_blocks=self.kv_cache_manager.empty_kv_cache_blocks.blocks,
@@ -2904,6 +2909,20 @@ class Scheduler(SchedulerInterface):
2909
num_tokens_main_model=full_num_tokens,
2910
apply_admission_cap=True,
2911
)
2912
+ return num_blocks + self._spec_decode_step_blocks()
2913
+
2914
+ def _spec_decode_step_blocks(self) -> int:
2915
+ """Number of blocks for the extra KV slots a spec decode step needs.
2916
+
2917
+ When using async kv load, scheduler must reserve enough blocks for
2918
+ full sequence + the spec decode step, otherwise request cannot be
2919
+ able to run after the async_load if we are out of kv blocks.
2920
+ """
2921
+ if not self.num_spec_tokens:
2922
+ return 0
2923
+ return cdiv(
2924
+ 1 + self.num_spec_tokens + self.num_lookahead_tokens, self.block_size
2925
+ )
2926
2927
def _inflight_prefill_reserved_blocks(self) -> int:
2928
"""Num blocks in-flight prefills still need to finish (their reservation)."""