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 15, 17:15 UTC
[Misc][V1] Misc code streamlining (#15723)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Mar 28, 2025 at 20:59 UTC
6d531ad7b810522fde902cb1cbf95f52bfc77860
5 files changed
+33
-39
vllm/distributed/utils.py
+1
-4
@@ -207,10 +207,7 @@ class StatelessProcessGroup:
207
def barrier(self):
208
"""A barrier to synchronize all ranks."""
209
for i in range(self.world_size):
210
- if i == self.rank:
211
- self.broadcast_obj(None, src=self.rank)
212
- else:
213
- self.broadcast_obj(None, src=i)
210
+ self.broadcast_obj(None, src=i)
211
212
@staticmethod
213
def create(
vllm/v1/core/sched/scheduler.py
+25
-30
@@ -269,29 +269,26 @@ class Scheduler(SchedulerInterface):
269
270
request = self.waiting[0]
271
272
- # Waiting request skipping logic
273
- is_skipped = False
272
# Skip request if the structured output request is still waiting
275
- # for FSM.
276
- if (not is_skipped
277
- and request.status == RequestStatus.WAITING_FOR_FSM):
273
+ # for FSM compilation.
274
+ if request.status == RequestStatus.WAITING_FOR_FSM:
275
structured_output_req = request.structured_output_request
279
- is_skipped = (not structured_output_req
280
- or not structured_output_req.grammar)
281
- if not is_skipped:
276
+ if structured_output_req and structured_output_req.grammar:
277
request.status = RequestStatus.WAITING
283
-
284
- # Skip request if max_loras can't be honored.
285
- if (not is_skipped and self.lora_config
286
- and request.lora_request):
287
- req_lora_id = request.lora_request.lora_int_id
288
- is_skipped = (len(scheduled_loras)
289
- == self.lora_config.max_loras
290
- and (req_lora_id not in scheduled_loras))
291
-
292
- if is_skipped:
293
- skipped_waiting_requests.appendleft(request)
278
+ else:
279
+ self.waiting.popleft()
280
+ skipped_waiting_requests.appendleft(request)
281
+ continue
282
+
283
+ # Check that adding the request still respects the max_loras
284
+ # constraint.
285
+ if self.lora_config and request.lora_request and (
286
+ len(scheduled_loras) == self.lora_config.max_loras
287
+ and request.lora_request.lora_int_id
288
+ not in scheduled_loras):
289
+ # Scheduling would exceed max_loras, skip.
290
self.waiting.popleft()
291
+ skipped_waiting_requests.appendleft(request)
292
continue
293
294
# Get already-cached tokens.
@@ -602,8 +599,9 @@ class Scheduler(SchedulerInterface):
599
# OPTIMIZATION: Avoid list(set) if the set is empty.
600
if cached_encoder_input_ids:
601
for input_id in list(cached_encoder_input_ids):
605
- start_pos = request.mm_positions[input_id]["offset"]
606
- num_tokens = request.mm_positions[input_id]["length"]
602
+ mm_positions = request.mm_positions[input_id]
603
+ start_pos = mm_positions["offset"]
604
+ num_tokens = mm_positions["length"]
605
if start_pos + num_tokens <= request.num_computed_tokens:
606
# The encoder output is already processed and stored
607
# in the decoder's KV cache.
@@ -616,25 +614,24 @@ class Scheduler(SchedulerInterface):
614
615
stopped = False
616
new_logprobs = None
619
- new_token_ids: list[int] = []
617
+ new_token_ids = generated_token_ids
618
619
# Append generated tokens and check for stop. Note that if
620
# a request is still being prefilled, we expect the model runner
621
# to return empty token ids for the request.
624
- for output_token_id in generated_token_ids:
622
+ for num_new, output_token_id in enumerate(new_token_ids, 1):
623
request.append_output_token_ids(output_token_id)
626
- new_token_ids.append(output_token_id)
624
625
# Check for stop and update request state.
626
# This must be called before we make the EngineCoreOutput.
627
stopped = check_stop(request, self.max_model_len)
628
if stopped:
629
self._free_request(request)
630
+ del new_token_ids[num_new:] # Trim new tokens if needed.
631
break
632
633
# Extract sample logprobs if needed.
636
- if (request.sampling_params.logprobs is not None
637
- and logprobs is not None):
634
+ if request.sampling_params.logprobs is not None and logprobs:
635
# NOTE: once we support N tokens per step (spec decode),
636
# the outer lists can be of length > 1.
637
new_logprobs = logprobs.slice(req_index, req_index + 1)
@@ -644,9 +641,7 @@ class Scheduler(SchedulerInterface):
641
# should not be None if use_structured_output, we have
642
# check above, so safe to ignore type warning
643
request.structured_output_request.grammar.accept_tokens( # type: ignore[union-attr]
647
- request.request_id,
648
- new_token_ids,
649
- )
644
+ req_id, new_token_ids)
645
646
# Get prompt logprobs for this request.
647
prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
@@ -665,7 +660,7 @@ class Scheduler(SchedulerInterface):
660
# Invariant: EngineCore returns no partial prefill outputs.
661
assert not prompt_logprobs_tensors
662
668
- self.scheduled_req_ids.remove(request.request_id)
663
+ self.scheduled_req_ids.remove(req_id)
664
if not stopped:
665
new_running.append(request)
666
vllm/v1/engine/core_client.py
+1
-1
@@ -416,9 +416,9 @@ class SyncMPClient(MPClient):
416
417
def process_outputs_socket():
418
shutdown_socket = ctx.socket(zmq.PAIR)
419
- shutdown_socket.bind(shutdown_path)
419
out_socket = make_zmq_socket(ctx, output_path, zmq.constants.PULL)
420
try:
421
+ shutdown_socket.bind(shutdown_path)
422
poller = zmq.Poller()
423
poller.register(shutdown_socket)
424
poller.register(out_socket)
vllm/v1/engine/output_processor.py
+1
-1
@@ -328,7 +328,7 @@ class OutputProcessor:
328
# 2) Detokenize the token ids into text and perform stop checks.
329
stop_string = req_state.detokenizer.update(
330
new_token_ids, finish_reason == FinishReason.STOP)
331
- if stop_string and finish_reason != FinishReason.STOP:
331
+ if stop_string:
332
finish_reason = FinishReason.STOP
333
stop_reason = stop_string
334
vllm/v1/request.py
+5
-3
@@ -93,9 +93,11 @@ class Request:
93
token_ids: Union[int, list[int]],
94
) -> None:
95
if isinstance(token_ids, int):
96
- token_ids = [token_ids]
97
- self._output_token_ids.extend(token_ids)
98
- self._all_token_ids.extend(token_ids)
96
+ self._output_token_ids.append(token_ids)
97
+ self._all_token_ids.append(token_ids)
98
+ else:
99
+ self._output_token_ids.extend(token_ids)
100
+ self._all_token_ids.extend(token_ids)
101
102
@property
103
def num_tokens(self) -> int: