fix: resolve silent request skipping in PRIORITY scheduling (#49206)
Signed-off-by: Tejas-Raj01 <rajtejas.xyz@gmail.com>
Tejas committed
Aug 7, 2026 at 03:51 UTC
4d341ca829d7fbad351b3a5a17d1405e63dc5bf2
2 files changed
+159
-2
tests/v1/core/test_priority_preemption_bug.py
new
+149
@@ -0,0 +1,149 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+from vllm.config import (
7
+ CacheConfig,
8
+ ModelConfig,
9
+ ParallelConfig,
10
+ SchedulerConfig,
11
+ VllmConfig,
12
+)
13
+from vllm.sampling_params import SamplingParams
14
+from vllm.utils.hashing import sha256
15
+from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
16
+from vllm.v1.core.sched.scheduler import Scheduler
17
+from vllm.v1.core.single_type_kv_cache_manager import register_all_kvcache_specs
18
+from vllm.v1.kv_cache_interface import (
19
+ FullAttentionSpec,
20
+ KVCacheConfig,
21
+ KVCacheGroupSpec,
22
+)
23
+from vllm.v1.outputs import ModelRunnerOutput
24
+from vllm.v1.request import Request, RequestStatus
25
+from vllm.v1.structured_output import StructuredOutputManager
26
+
27
+BLOCK_SIZE = 16
28
+NUM_BLOCKS = 6
29
+MAX_MODEL_LEN = 1000
30
+LONG_PREFILL_THRESHOLD = 15
31
+EOS_TOKEN_ID = 50256
32
+
33
+
34
+def build_scheduler():
35
+ model_config = ModelConfig(
36
+ model="facebook/opt-350m",
37
+ trust_remote_code=True,
38
+ dtype="float16",
39
+ seed=42,
40
+ skip_tokenizer_init=True,
41
+ max_model_len=MAX_MODEL_LEN,
42
+ )
43
+ scheduler_config = SchedulerConfig(
44
+ max_num_seqs=5,
45
+ max_num_batched_tokens=200,
46
+ max_model_len=MAX_MODEL_LEN,
47
+ long_prefill_token_threshold=LONG_PREFILL_THRESHOLD,
48
+ enable_chunked_prefill=True,
49
+ is_encoder_decoder=model_config.is_encoder_decoder,
50
+ policy="priority",
51
+ watermark=0.0,
52
+ scheduler_reserve_full_isl=False,
53
+ )
54
+ cache_config = CacheConfig(
55
+ block_size=BLOCK_SIZE,
56
+ gpu_memory_utilization=0.9,
57
+ cache_dtype="auto",
58
+ enable_prefix_caching=False,
59
+ )
60
+ vllm_config = VllmConfig(
61
+ scheduler_config=scheduler_config,
62
+ model_config=model_config,
63
+ cache_config=cache_config,
64
+ parallel_config=ParallelConfig(),
65
+ )
66
+ kv_cache_spec = FullAttentionSpec(
67
+ block_size=BLOCK_SIZE, num_kv_heads=1, head_size=1, dtype=torch.float32
68
+ )
69
+ kv_cache_config = KVCacheConfig(
70
+ num_blocks=NUM_BLOCKS,
71
+ kv_cache_tensors=[],
72
+ kv_cache_groups=[KVCacheGroupSpec(["layer"], kv_cache_spec)],
73
+ )
74
+ cache_config.num_gpu_blocks = NUM_BLOCKS
75
+ register_all_kvcache_specs(vllm_config)
76
+ return Scheduler(
77
+ vllm_config=vllm_config,
78
+ kv_cache_config=kv_cache_config,
79
+ block_size=BLOCK_SIZE,
80
+ log_stats=True,
81
+ structured_output_manager=StructuredOutputManager(vllm_config),
82
+ )
83
+
84
+
85
+_hash_init = False
86
+
87
+
88
+def make_request(req_id, num_tokens, priority, arrival_time):
89
+ global _hash_init
90
+ if not _hash_init:
91
+ init_none_hash(sha256)
92
+ _hash_init = True
93
+ block_hasher = get_request_block_hasher(BLOCK_SIZE, sha256)
94
+ sp = SamplingParams(ignore_eos=True, max_tokens=50)
95
+ sp.update_from_generation_config({}, EOS_TOKEN_ID)
96
+ return Request(
97
+ request_id=req_id,
98
+ prompt_token_ids=[7] * num_tokens,
99
+ sampling_params=sp,
100
+ pooling_params=None,
101
+ priority=priority,
102
+ arrival_time=arrival_time,
103
+ block_hasher=block_hasher,
104
+ )
105
+
106
+
107
+def mock_output(out, tok=100):
108
+ ids = list(out.num_scheduled_tokens.keys())
109
+ return ModelRunnerOutput(
110
+ req_ids=ids,
111
+ req_id_to_index={r: i for i, r in enumerate(ids)},
112
+ sampled_token_ids=[[tok] for _ in ids],
113
+ logprobs=None,
114
+ prompt_logprobs_dict={},
115
+ pooler_output=[],
116
+ )
117
+
118
+
119
+def test_priority_scheduler_preempt_skipped_request():
120
+ scheduler = build_scheduler()
121
+
122
+ # R1: Add request A (Worst priority)
123
+ A = make_request("A", num_tokens=200, priority=9, arrival_time=1.0)
124
+ scheduler.add_request(A)
125
+ out = scheduler.schedule()
126
+ scheduler.update_from_output(out, mock_output(out))
127
+
128
+ # R2: Add request B (Best priority)
129
+ B = make_request("B", num_tokens=15, priority=0, arrival_time=2.0)
130
+ scheduler.add_request(B)
131
+ out = scheduler.schedule()
132
+ scheduler.update_from_output(out, mock_output(out))
133
+
134
+ # R3: Add request C (Middle priority)
135
+ C = make_request("C", num_tokens=1, priority=1, arrival_time=3.0)
136
+ scheduler.add_request(C)
137
+ out = scheduler.schedule()
138
+ scheduler.update_from_output(out, mock_output(out))
139
+
140
+ # R4: Trigger preemption
141
+ out = scheduler.schedule(throttle_prefills=True)
142
+
143
+ is_c_scheduled = "C" in out.num_scheduled_tokens
144
+
145
+ assert is_c_scheduled, (
146
+ "Bug present: Request C was silently skipped during scheduling "
147
+ "because the req_index was not decremented after preempting A."
148
+ )
149
+ assert C.status == RequestStatus.RUNNING
vllm/v1/core/sched/scheduler.py
+10
-2
@@ -593,7 +593,16 @@ class Scheduler(SchedulerInterface):
593
self.running,
594
key=lambda r: (r.priority, r.arrival_time),
595
)
596
- self.running.remove(preempted_req)
596
+ # Record the index of the preemption victim to
597
+ # maintain accurate loop state.
598
+ victim_index = self.running.index(preempted_req)
599
+ del self.running[victim_index]
600
+ # Decrement the loop cursor if the removed request
601
+ # preceded the current iteration, preventing the
602
+ # silent omission of the subsequent request.
603
+ if victim_index < req_index:
604
+ req_index -= 1
605
+
606
if preempted_req in scheduled_running_reqs:
607
preempted_req_id = preempted_req.request_id
608
scheduled_running_reqs.remove(preempted_req)
@@ -611,7 +620,6 @@ class Scheduler(SchedulerInterface):
620
for i in preempted_encoder_inputs
621
)
622
encoder_compute_budget += num_embeds_to_restore
614
- req_index -= 1
623
else:
624
preempted_req = self.running.pop()
625