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 14, 05:30 UTC
[Bugfix] Fix for V1 priority scheduling crashes at preemption (#23713)
Signed-off-by: Hanchenli <lihanc2002@gmail.com>
Hanchenli committed
Aug 27, 2025 at 17:44 UTC
5da4f5d857933329aaca779e3a81f1385c84e34a
2 files changed
+90
-3
tests/v1/core/test_scheduler.py
+88
-3
@@ -1293,7 +1293,8 @@ def create_requests_with_priority(
1293
mm_positions: Optional[list[list[PlaceholderRange]]] = None,
1294
max_tokens: int = 16,
1295
stop_token_ids: Optional[list[int]] = None,
1296
- prompt_logprobs: Optional[int] = None):
1296
+ prompt_logprobs: Optional[int] = None,
1297
+ starting_idx: int = 0):
1298
"""Create requests with specified priorities and arrival times."""
1299
assert len(priorities) == num_requests
1300
if arrival_times is not None:
@@ -1315,8 +1316,8 @@ def create_requests_with_priority(
1316
mm_position = None
1317
mm_kwargs = None
1318
request = Request(
1318
- request_id=f"{i}",
1319
- prompt_token_ids=[i] * num_tokens,
1319
+ request_id=f"{i + starting_idx}",
1320
+ prompt_token_ids=[i + starting_idx] * num_tokens,
1321
sampling_params=sampling_params,
1322
pooling_params=None,
1323
multi_modal_kwargs=mm_kwargs,
@@ -1813,3 +1814,87 @@ def test_schedule_skip_tokenizer_init_structured_output_request():
1814
assert len(output.scheduled_new_reqs) == 0
1815
assert len(scheduler.running) == 0
1816
assert len(scheduler.waiting) == 1
1817
+
1818
+
1819
+def test_priority_scheduling_preemption_when_out_of_kv():
1820
+ """Test that priority scheduling preempts lower priority requests
1821
+ when out of KV cache space."""
1822
+ # Create scheduler with very limited memory to force preemption
1823
+ scheduler = create_scheduler_with_priority(
1824
+ max_num_seqs=2, # Allow multiple requests
1825
+ max_num_batched_tokens=200,
1826
+ num_blocks=5, # Can hold 64 tokens (first block is null)
1827
+ block_size=16, # Standard block size
1828
+ )
1829
+
1830
+ # Create a request and schedule it
1831
+ request_low = create_requests_with_priority(
1832
+ num_requests=1,
1833
+ priorities=[1],
1834
+ arrival_times=[0.0],
1835
+ num_tokens=30,
1836
+ starting_idx=0,
1837
+ )[0]
1838
+ scheduler.add_request(request_low)
1839
+ output = scheduler.schedule()
1840
+ assert len(output.scheduled_new_reqs) == 1
1841
+ assert len(scheduler.waiting) == 0
1842
+ assert len(scheduler.running) == 1
1843
+
1844
+ # Simulate model execution
1845
+ model_output = ModelRunnerOutput(
1846
+ req_ids=[request_low.request_id],
1847
+ req_id_to_index={request_low.request_id: 0},
1848
+ sampled_token_ids=[[100]],
1849
+ # spec_token_ids=None,
1850
+ logprobs=None,
1851
+ prompt_logprobs_dict={},
1852
+ pooler_output=[],
1853
+ )
1854
+ scheduler.update_from_output(output, model_output)
1855
+
1856
+ # Create a high priority request and schedule it
1857
+ request_high = create_requests_with_priority(
1858
+ num_requests=1,
1859
+ priorities=[0],
1860
+ arrival_times=[1.0],
1861
+ num_tokens=32,
1862
+ starting_idx=1,
1863
+ )[0]
1864
+ scheduler.add_request(request_high)
1865
+ output = scheduler.schedule()
1866
+ # KV cache should be full at this point
1867
+ assert scheduler.kv_cache_manager.block_pool.get_num_free_blocks() == 0
1868
+ assert len(output.scheduled_new_reqs) == 1
1869
+ assert output.scheduled_cached_reqs.num_reqs == 1
1870
+ assert len(scheduler.waiting) == 0
1871
+ assert len(scheduler.running) == 2
1872
+
1873
+ # Simulate model execution
1874
+ requests = [request_low, request_high]
1875
+ model_output = ModelRunnerOutput(
1876
+ req_ids=[req.request_id for req in requests],
1877
+ req_id_to_index={
1878
+ req.request_id: i
1879
+ for i, req in enumerate(requests)
1880
+ },
1881
+ sampled_token_ids=[[100] for _ in requests],
1882
+ # spec_token_ids=None,
1883
+ logprobs=None,
1884
+ prompt_logprobs_dict={},
1885
+ pooler_output=[],
1886
+ )
1887
+ scheduler.update_from_output(output, model_output)
1888
+
1889
+ # Schedule again - this should trigger preemption
1890
+ # req_low needs 32 tokens = 2 blocks
1891
+ # req_high needs 33 tokens = 3 blocks
1892
+ # so doesn't fit in 4 blocks.
1893
+ output = scheduler.schedule()
1894
+
1895
+ # Should have preempted req_low
1896
+ assert len(output.scheduled_new_reqs) == 0
1897
+ assert output.scheduled_cached_reqs.num_reqs == 1
1898
+ assert output.scheduled_cached_reqs.req_ids[0] == request_high.request_id
1899
+ assert len(scheduler.waiting) == 1
1900
+ assert len(scheduler.running) == 1
\ No newline at end of file
vllm/v1/core/sched/scheduler.py
+2
@@ -253,6 +253,8 @@ class Scheduler(SchedulerInterface):
253
key=lambda r: (r.priority, r.arrival_time),
254
)
255
self.running.remove(preempted_req)
256
+ if preempted_req in scheduled_running_reqs:
257
+ scheduled_running_reqs.remove(preempted_req)
258
else:
259
preempted_req = self.running.pop()
260