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
[Core] Optimize scheduler request removal for single completions (#21917)
Signed-off-by: chiliu <chiliu@paypal.com> Signed-off-by: chiliu <cliu_whu@yeah.net> Co-authored-by: chiliu <chiliu@paypal.com>
633WHU committed
Aug 20, 2025 at 09:25 UTC
0167efe20d3d2280c3da6aea94a6f59afec5099c
2 files changed
+39
-8
vllm/v1/core/sched/scheduler.py
+6
-8
@@ -25,7 +25,7 @@ from vllm.v1.core.sched.output import (CachedRequestData, NewRequestData,
25
SchedulerOutput)
26
from vllm.v1.core.sched.request_queue import (SchedulingPolicy,
27
create_request_queue)
28
-from vllm.v1.core.sched.utils import check_stop
28
+from vllm.v1.core.sched.utils import check_stop, remove_all
29
from vllm.v1.engine import (EngineCoreEventType, EngineCoreOutput,
30
EngineCoreOutputs)
31
from vllm.v1.kv_cache_interface import KVCacheConfig
@@ -872,9 +872,7 @@ class Scheduler(SchedulerInterface):
872
873
# Remove the stopped requests from the running and waiting queues.
874
if stopped_running_reqs:
875
- self.running = [
876
- req for req in self.running if req not in stopped_running_reqs
877
- ]
875
+ self.running = remove_all(self.running, stopped_running_reqs)
876
if stopped_preempted_reqs:
877
# This is a rare case and unlikely to impact performance.
878
self.waiting.remove_requests(stopped_preempted_reqs)
@@ -1000,7 +998,7 @@ class Scheduler(SchedulerInterface):
998
else:
999
request_ids = set(request_ids)
1000
1003
- running_requests_to_remove = []
1001
+ running_requests_to_remove = set()
1002
waiting_requests_to_remove = []
1003
valid_requests = []
1004
@@ -1013,13 +1011,13 @@ class Scheduler(SchedulerInterface):
1011
1012
valid_requests.append(request)
1013
if request.status == RequestStatus.RUNNING:
1016
- running_requests_to_remove.append(request)
1014
+ running_requests_to_remove.add(request)
1015
else:
1016
waiting_requests_to_remove.append(request)
1017
1018
# Remove all requests from queues at once for better efficiency
1021
- for request in running_requests_to_remove:
1022
- self.running.remove(request)
1019
+ if running_requests_to_remove:
1020
+ self.running = remove_all(self.running, running_requests_to_remove)
1021
if waiting_requests_to_remove:
1022
self.waiting.remove_requests(waiting_requests_to_remove)
1023
vllm/v1/core/sched/utils.py
+33
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import contextlib
4
from typing import Optional
5
6
import torch
@@ -7,6 +8,38 @@ import torch
8
from vllm.v1.request import Request, RequestStatus
9
10
11
+def remove_all(lst: list, items_to_remove: set) -> list:
12
+ """Remove all items from a list that are in the items_to_remove set.
13
+
14
+ This method optimizes for the common case of removing a single item,
15
+ falling back to list comprehension for multiple items.
16
+
17
+ Args:
18
+ lst: The list to remove items from
19
+ items_to_remove: Set of items to remove
20
+
21
+ Returns:
22
+ Either the modified original list (for single item removal) or
23
+ a new list (for multiple item removal). Callers should use the
24
+ returned value.
25
+
26
+ Note:
27
+ For single item removal, this modifies the original list in-place
28
+ and returns it. For multiple items, it creates and returns a new list.
29
+ """
30
+ if not items_to_remove:
31
+ return lst
32
+
33
+ if len(items_to_remove) == 1:
34
+ # Fast path for single item removal (most common case)
35
+ item = next(iter(items_to_remove))
36
+ with contextlib.suppress(ValueError):
37
+ lst.remove(item)
38
+ return lst
39
+ # For multiple items, use list comprehension
40
+ return [item for item in lst if item not in items_to_remove]
41
+
42
+
43
def check_stop(request: Request,
44
max_model_len: int,
45
pooler_output: Optional[torch.Tensor] = None) -> bool: