From 278e3270a53291989a1776be443ab7065da2ba01 Mon Sep 17 00:00:00 2001 From: Tejas-Raj01 Date: Mon, 20 Jul 2026 19:55:24 +0530 Subject: [PATCH 1/3] fix: resolve silent request skipping in PRIORITY scheduling Signed-off-by: Tejas-Raj01 --- tests/v1/core/test_priority_preemption_bug.py | 37 +++++++++++++++++++ vllm/v1/core/sched/scheduler.py | 12 +++++- 2 files changed, 48 insertions(+), 1 deletion(-) create mode 100644 tests/v1/core/test_priority_preemption_bug.py diff --git a/tests/v1/core/test_priority_preemption_bug.py b/tests/v1/core/test_priority_preemption_bug.py new file mode 100644 index 000000000000..0aa29803aacd --- /dev/null +++ b/tests/v1/core/test_priority_preemption_bug.py @@ -0,0 +1,37 @@ +from vllm.v1.core.sched.scheduler import RequestStatus + + +def test_priority_scheduler_preempt_skipped_request(): + # Importing helper functions from the reproduction script + from repro import build_scheduler, make_request, mock_output + + scheduler = build_scheduler() + + # R1: Add request A (Worst priority) + A = make_request("A", num_tokens=200, priority=9, arrival_time=1.0) + scheduler.add_request(A) + out = scheduler.schedule() + scheduler.update_from_output(out, mock_output(out)) + + # R2: Add request B (Best priority) + B = make_request("B", num_tokens=15, priority=0, arrival_time=2.0) + scheduler.add_request(B) + out = scheduler.schedule() + scheduler.update_from_output(out, mock_output(out)) + + # R3: Add request C (Middle priority) + C = make_request("C", num_tokens=1, priority=1, arrival_time=3.0) + scheduler.add_request(C) + out = scheduler.schedule() + scheduler.update_from_output(out, mock_output(out)) + + # R4: Trigger preemption + out = scheduler.schedule(throttle_prefills=True) + + is_c_scheduled = "C" in out.num_scheduled_tokens + + assert is_c_scheduled, ( + "Bug present: Request C was silently skipped during scheduling " + "because the req_index was not decremented after preempting A." + ) + assert C.status == RequestStatus.RUNNING diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index a693061d5173..91b95a30f572 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -574,7 +574,18 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: self.running, key=lambda r: (r.priority, r.arrival_time), ) + + # Record the index of the preemption victim to + # maintain accurate loop state. + victim_index = self.running.index(preempted_req) self.running.remove(preempted_req) + + # Decrement the loop cursor if the removed request + # preceded the current iteration, preventing the + # silent omission of the subsequent request. + if victim_index < req_index: + req_index -= 1 + if preempted_req in scheduled_running_reqs: preempted_req_id = preempted_req.request_id scheduled_running_reqs.remove(preempted_req) @@ -592,7 +603,6 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: for i in preempted_encoder_inputs ) encoder_compute_budget += num_embeds_to_restore - req_index -= 1 else: preempted_req = self.running.pop() From 8e1ef7defb079459af80ddbdfb6c27c6596e9121 Mon Sep 17 00:00:00 2001 From: Tejas-Raj01 Date: Thu, 6 Aug 2026 05:21:23 +0530 Subject: [PATCH 2/3] fix: address reviewer feedback for victim removal Signed-off-by: Tejas-Raj01 --- vllm/v1/core/sched/scheduler.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 91b95a30f572..5d3a51f1ed2a 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -574,12 +574,10 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: self.running, key=lambda r: (r.priority, r.arrival_time), ) - # Record the index of the preemption victim to # maintain accurate loop state. victim_index = self.running.index(preempted_req) - self.running.remove(preempted_req) - + del self.running[victim_index] # Decrement the loop cursor if the removed request # preceded the current iteration, preventing the # silent omission of the subsequent request. From bd23a2ea2e84f9dd5f1b1e76d083a18749ee140b Mon Sep 17 00:00:00 2001 From: Tejas-Raj01 Date: Thu, 6 Aug 2026 05:38:49 +0530 Subject: [PATCH 3/3] fix: enhance priority scheduler with comprehensive request handling and initialization Signed-off-by: Tejas-Raj01 --- tests/v1/core/test_priority_preemption_bug.py | 120 +++++++++++++++++- 1 file changed, 116 insertions(+), 4 deletions(-) diff --git a/tests/v1/core/test_priority_preemption_bug.py b/tests/v1/core/test_priority_preemption_bug.py index 0aa29803aacd..da9e9f55d576 100644 --- a/tests/v1/core/test_priority_preemption_bug.py +++ b/tests/v1/core/test_priority_preemption_bug.py @@ -1,10 +1,122 @@ -from vllm.v1.core.sched.scheduler import RequestStatus +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch -def test_priority_scheduler_preempt_skipped_request(): - # Importing helper functions from the reproduction script - from repro import build_scheduler, make_request, mock_output +from vllm.config import ( + CacheConfig, + ModelConfig, + ParallelConfig, + SchedulerConfig, + VllmConfig, +) +from vllm.sampling_params import SamplingParams +from vllm.utils.hashing import sha256 +from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash +from vllm.v1.core.sched.scheduler import Scheduler +from vllm.v1.core.single_type_kv_cache_manager import register_all_kvcache_specs +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, +) +from vllm.v1.outputs import ModelRunnerOutput +from vllm.v1.request import Request, RequestStatus +from vllm.v1.structured_output import StructuredOutputManager + +BLOCK_SIZE = 16 +NUM_BLOCKS = 6 +MAX_MODEL_LEN = 1000 +LONG_PREFILL_THRESHOLD = 15 +EOS_TOKEN_ID = 50256 + + +def build_scheduler(): + model_config = ModelConfig( + model="facebook/opt-350m", + trust_remote_code=True, + dtype="float16", + seed=42, + skip_tokenizer_init=True, + max_model_len=MAX_MODEL_LEN, + ) + scheduler_config = SchedulerConfig( + max_num_seqs=5, + max_num_batched_tokens=200, + max_model_len=MAX_MODEL_LEN, + long_prefill_token_threshold=LONG_PREFILL_THRESHOLD, + enable_chunked_prefill=True, + is_encoder_decoder=model_config.is_encoder_decoder, + policy="priority", + watermark=0.0, + scheduler_reserve_full_isl=False, + ) + cache_config = CacheConfig( + block_size=BLOCK_SIZE, + gpu_memory_utilization=0.9, + cache_dtype="auto", + enable_prefix_caching=False, + ) + vllm_config = VllmConfig( + scheduler_config=scheduler_config, + model_config=model_config, + cache_config=cache_config, + parallel_config=ParallelConfig(), + ) + kv_cache_spec = FullAttentionSpec( + block_size=BLOCK_SIZE, num_kv_heads=1, head_size=1, dtype=torch.float32 + ) + kv_cache_config = KVCacheConfig( + num_blocks=NUM_BLOCKS, + kv_cache_tensors=[], + kv_cache_groups=[KVCacheGroupSpec(["layer"], kv_cache_spec)], + ) + cache_config.num_gpu_blocks = NUM_BLOCKS + register_all_kvcache_specs(vllm_config) + return Scheduler( + vllm_config=vllm_config, + kv_cache_config=kv_cache_config, + block_size=BLOCK_SIZE, + log_stats=True, + structured_output_manager=StructuredOutputManager(vllm_config), + ) + +_hash_init = False + + +def make_request(req_id, num_tokens, priority, arrival_time): + global _hash_init + if not _hash_init: + init_none_hash(sha256) + _hash_init = True + block_hasher = get_request_block_hasher(BLOCK_SIZE, sha256) + sp = SamplingParams(ignore_eos=True, max_tokens=50) + sp.update_from_generation_config({}, EOS_TOKEN_ID) + return Request( + request_id=req_id, + prompt_token_ids=[7] * num_tokens, + sampling_params=sp, + pooling_params=None, + priority=priority, + arrival_time=arrival_time, + block_hasher=block_hasher, + ) + + +def mock_output(out, tok=100): + ids = list(out.num_scheduled_tokens.keys()) + return ModelRunnerOutput( + req_ids=ids, + req_id_to_index={r: i for i, r in enumerate(ids)}, + sampled_token_ids=[[tok] for _ in ids], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ) + + +def test_priority_scheduler_preempt_skipped_request(): scheduler = build_scheduler() # R1: Add request A (Worst priority)