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] V1 Fix the cursor leakage issue during request scheduling. (#21173)
Signed-off-by: CLFutureX <775523362@qq.com>
PiteXChen committed
Aug 5, 2025 at 09:34 UTC
2dffac464c82ac7c509c78f7d12a7c72ea765a63
2 files changed
+100
-3
tests/v1/core/test_scheduler.py
+95
-2
@@ -1307,13 +1307,18 @@ def create_requests_with_priority(
1307
mm_positions: Optional[list[PlaceholderRange]] = None,
1308
max_tokens: int = 16,
1309
stop_token_ids: Optional[list[int]] = None,
1310
- prompt_logprobs: Optional[int] = None):
1310
+ prompt_logprobs: Optional[int] = None,
1311
+ request_ids: Optional[list[str]] = None):
1312
"""Create requests with specified priorities and arrival times."""
1313
assert len(priorities) == num_requests
1314
if arrival_times is not None:
1315
assert len(arrival_times) == num_requests
1316
else:
1317
arrival_times = [float(i) for i in range(num_requests)]
1318
+ if request_ids is not None:
1319
+ assert len(request_ids) == num_requests
1320
+ else:
1321
+ request_ids = [f"{i}" for i in range(num_requests)]
1322
1323
sampling_params = SamplingParams(ignore_eos=False,
1324
max_tokens=max_tokens,
@@ -1328,7 +1333,7 @@ def create_requests_with_priority(
1333
mm_position = None
1334
mm_inputs = None
1335
request = Request(
1331
- request_id=f"{i}",
1336
+ request_id=request_ids[i],
1337
prompt_token_ids=[i] * num_tokens,
1338
sampling_params=sampling_params,
1339
pooling_params=None,
@@ -1829,3 +1834,91 @@ def test_schedule_skip_tokenizer_init_structured_output_request():
1834
assert len(output.scheduled_new_reqs) == 0
1835
assert len(scheduler.running) == 0
1836
assert len(scheduler.waiting) == 1
1837
+
1838
+
1839
+def test_priority_scheduling_preemption_victim_iterator_order():
1840
+ """Test that the scheduling order is maintained after
1841
+ preempting lower-priority requests."""
1842
+ scheduler = create_scheduler_with_priority(
1843
+ max_num_batched_tokens=200,
1844
+ num_blocks=9,
1845
+ )
1846
+ # Add three priority requests first.
1847
+ priority_requests = create_requests_with_priority(
1848
+ num_requests=3,
1849
+ priorities=[3, 4, 5],
1850
+ arrival_times=[1.0, 2.0, 3.0],
1851
+ num_tokens=15,
1852
+ request_ids=["1", "2", "3"],
1853
+ )
1854
+
1855
+ for request in priority_requests:
1856
+ scheduler.add_request(request)
1857
+ # After scheduling, transfer from the waiting queue to the running queue.
1858
+ # At this time, 3 blocks have been allocated, and 5 available blocks remain.
1859
+ output = scheduler.schedule()
1860
+
1861
+ model_output = ModelRunnerOutput(
1862
+ req_ids=[req.request_id for req in priority_requests],
1863
+ req_id_to_index={
1864
+ req.request_id: i
1865
+ for i, req in enumerate(priority_requests)
1866
+ },
1867
+ sampled_token_ids=[[15] for _ in priority_requests],
1868
+ spec_token_ids=None,
1869
+ logprobs=None,
1870
+ prompt_logprobs_dict={},
1871
+ pooler_output=[],
1872
+ )
1873
+ scheduler.update_from_output(output, model_output)
1874
+
1875
+ # Add tow high priority requests.
1876
+ high_priority_requests = create_requests_with_priority(
1877
+ num_requests=2,
1878
+ priorities=[1, 2],
1879
+ arrival_times=[4.0, 5.0],
1880
+ num_tokens=16,
1881
+ request_ids=["4", "5"],
1882
+ )
1883
+ for request in high_priority_requests:
1884
+ scheduler.add_request(request)
1885
+
1886
+ # After scheduling, transfer the two high-priority requests from
1887
+ # the waiting queue to the running queue.
1888
+ # the IDs of the requests in the running queue are: 1, 2, 3, 4, 5.
1889
+ # At this time, 3+2 blocks have been allocated,
1890
+ # and 3 available blocks remain.
1891
+ output = scheduler.schedule()
1892
+
1893
+ merge_requests = priority_requests + high_priority_requests
1894
+
1895
+ model_output = ModelRunnerOutput(
1896
+ req_ids=[req.request_id for req in merge_requests],
1897
+ req_id_to_index={
1898
+ req.request_id: i
1899
+ for i, req in enumerate(merge_requests)
1900
+ },
1901
+ sampled_token_ids=[[1] for _ in merge_requests],
1902
+ spec_token_ids=None,
1903
+ logprobs=None,
1904
+ prompt_logprobs_dict={},
1905
+ pooler_output=[],
1906
+ )
1907
+ scheduler.update_from_output(output, model_output)
1908
+
1909
+ # At this time, the request with the lowest priority
1910
+ # (request.id = 2) will be preempted, freeing up 2 blocks,
1911
+ # which exactly meets the resource allocation requirements
1912
+ # for request.id = 4 and request.id = 5.
1913
+ output = scheduler.schedule()
1914
+
1915
+ # Should schedule the new request without preemption.
1916
+ assert len(scheduler.running) == 4 #
1917
+ assert len(scheduler.waiting) == 1 #
1918
+
1919
+ running_priorities = [req.priority for req in scheduler.running]
1920
+ running_req_ids = [req.request_id for req in scheduler.running]
1921
+
1922
+ assert running_priorities == [3, 4, 1, 2]
1923
+ assert running_req_ids == ["1", "2", "4", "5"]
1924
+ assert scheduler.waiting.peek_request().priority == 5
vllm/v1/core/sched/scheduler.py
+5
-1
@@ -257,7 +257,11 @@ class Scheduler(SchedulerInterface):
257
self.running,
258
key=lambda r: (r.priority, r.arrival_time),
259
)
260
- self.running.remove(preempted_req)
260
+ preempted_index = self.running.index(preempted_req)
261
+ if preempted_index <= req_index:
262
+ req_index -= 1
263
+ scheduled_running_reqs.remove(preempted_req)
264
+ self.running.pop(preempted_index)
265
else:
266
preempted_req = self.running.pop()
267