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, 21:00 UTC
[V1] Refactor parallel sampling support (#13774)
Signed-off-by: Mark McLoughlin <markmc@redhat.com>
Mark McLoughlin committed
Mar 3, 2025 at 16:15 UTC
4167252eafa214d6aa4ecfb627d7b16a31095f08
5 files changed
+201
-464
vllm/v1/engine/async_llm.py
+21
-40
@@ -25,7 +25,7 @@ from vllm.usage.usage_lib import UsageContext
25
from vllm.utils import cdiv, kill_process_tree
26
from vllm.v1.engine.core_client import EngineCoreClient
27
from vllm.v1.engine.output_processor import OutputProcessor
28
-from vllm.v1.engine.parallel_sampling import generate_parallel_sampling_async
28
+from vllm.v1.engine.parallel_sampling import ParentRequest
29
from vllm.v1.engine.processor import Processor
30
from vllm.v1.executor.abstract import Executor
31
from vllm.v1.metrics.loggers import (LoggingStatLogger, PrometheusStatLogger,
@@ -145,25 +145,30 @@ class AsyncLLM(EngineClient):
145
"""Add new request to the AsyncLLM."""
146
147
# 1) Create a new output queue for the request.
148
- if self.output_processor.is_request_active(request_id):
149
- raise ValueError(f"Request id {request_id} already running.")
148
queue: asyncio.Queue[RequestOutput] = asyncio.Queue()
149
152
- # 2) Convert Input --> Request.
153
- request = self.processor.process_inputs(request_id, prompt, params,
154
- arrival_time, lora_request,
155
- trace_headers,
156
- prompt_adapter_request,
157
- priority)
150
+ # 2) Fan out child requests (for n>1)
151
+ parent_req = ParentRequest.from_params(request_id, params)
152
+ n = params.n if isinstance(params, SamplingParams) else 1
153
+ for idx in range(n):
154
+ if parent_req is not None:
155
+ request_id, params = parent_req.get_child_info(idx)
156
159
- # 3) Add the request to OutputProcessor (this process).
160
- self.output_processor.add_request(request, queue)
157
+ # 3) Convert Input --> Request.
158
+ request = self.processor.process_inputs(request_id, prompt, params,
159
+ arrival_time, lora_request,
160
+ trace_headers,
161
+ prompt_adapter_request,
162
+ priority)
163
162
- # 4) Add the EngineCoreRequest to EngineCore (separate process).
163
- await self.engine_core.add_request_async(request)
164
+ # 4) Add the request to OutputProcessor (this process).
165
+ self.output_processor.add_request(request, parent_req, idx, queue)
166
165
- if self.log_requests:
166
- logger.info("Added request %s.", request_id)
167
+ # 5) Add the EngineCoreRequest to EngineCore (separate process).
168
+ await self.engine_core.add_request_async(request)
169
+
170
+ if self.log_requests:
171
+ logger.info("Added request %s.", request_id)
172
173
return queue
174
@@ -172,7 +177,7 @@ class AsyncLLM(EngineClient):
177
# requests we don't need to send multiple messages to core proc,
178
# and so we don't need multiple streams which then get
179
# re-multiplexed in the API server anyhow.
175
- async def _generate(
180
+ async def generate(
181
self,
182
prompt: PromptType,
183
sampling_params: SamplingParams,
@@ -243,30 +248,6 @@ class AsyncLLM(EngineClient):
248
await self.abort(request_id)
249
raise
250
246
- def generate(
247
- self,
248
- prompt: PromptType,
249
- sampling_params: SamplingParams,
250
- request_id: str,
251
- lora_request: Optional[LoRARequest] = None,
252
- trace_headers: Optional[Mapping[str, str]] = None,
253
- prompt_adapter_request: Optional[PromptAdapterRequest] = None,
254
- priority: int = 0,
255
- ) -> AsyncGenerator[RequestOutput, None]:
256
- kwargs = dict(prompt=prompt,
257
- sampling_params=sampling_params,
258
- request_id=request_id,
259
- lora_request=lora_request,
260
- trace_headers=trace_headers,
261
- prompt_adapter_request=prompt_adapter_request,
262
- priority=priority)
263
- if sampling_params.n is None or sampling_params.n == 1:
264
- return self._generate(**kwargs)
265
- else:
266
- # Special handling for parallel sampling requests
267
- return generate_parallel_sampling_async(generate=self._generate,
268
- **kwargs)
269
-
251
async def _run_output_handler(self):
252
"""Background loop: pulls from EngineCore and pushes to AsyncStreams."""
253
vllm/v1/engine/llm_engine.py
+22
-52
@@ -22,7 +22,7 @@ from vllm.transformers_utils.tokenizer_group import (
22
from vllm.usage.usage_lib import UsageContext
23
from vllm.v1.engine.core_client import EngineCoreClient
24
from vllm.v1.engine.output_processor import OutputProcessor
25
-from vllm.v1.engine.parallel_sampling import SyncParallelSamplingManager
25
+from vllm.v1.engine.parallel_sampling import ParentRequest
26
from vllm.v1.engine.processor import Processor
27
from vllm.v1.executor.abstract import Executor
28
@@ -50,9 +50,6 @@ class LLMEngine:
50
self.model_config = vllm_config.model_config
51
self.cache_config = vllm_config.cache_config
52
53
- # Bookkeeping for parallel sampling requests
54
- self.parallel_manager = SyncParallelSamplingManager()
55
-
53
# important: init dp group before init the engine_core
54
self.parallel_config = vllm_config.parallel_config
55
self.dp_enabled = self.parallel_config.data_parallel_size > 1 # noqa
@@ -120,8 +117,7 @@ class LLMEngine:
117
multiprocess_mode=enable_multiprocessing)
118
119
def get_num_unfinished_requests(self) -> int:
123
- return self.parallel_manager.get_num_unfinished_requests(
124
- self.output_processor.get_num_unfinished_requests())
120
+ return self.output_processor.get_num_unfinished_requests()
121
122
def has_unfinished_requests(self) -> bool:
123
has_unfinished = self.output_processor.has_unfinished_requests()
@@ -157,48 +153,25 @@ class LLMEngine:
153
prompt_adapter_request: Optional[PromptAdapterRequest] = None,
154
priority: int = 0,
155
) -> None:
160
- """Add request."""
161
- kwargs = dict(request_id=request_id,
162
- prompt=prompt,
163
- params=params,
164
- arrival_time=arrival_time,
165
- lora_request=lora_request,
166
- trace_headers=trace_headers,
167
- prompt_adapter_request=prompt_adapter_request,
168
- priority=priority)
169
- # Handle parallel sampling requests differently.
170
- if params is None or isinstance(params,
171
- PoolingParams) or params.n == 1:
172
- self._add_request(**kwargs)
173
- else:
174
- # Special handling for parallel sampling requests
175
- self.parallel_manager.add_request_parallel_sampling(
176
- add_request=self._add_request, **kwargs)
177
-
178
- def _add_request(
179
- self,
180
- request_id: str,
181
- prompt: PromptType,
182
- params: Union[SamplingParams, PoolingParams],
183
- arrival_time: Optional[float] = None,
184
- lora_request: Optional[LoRARequest] = None,
185
- trace_headers: Optional[Mapping[str, str]] = None,
186
- prompt_adapter_request: Optional[PromptAdapterRequest] = None,
187
- priority: int = 0,
188
- ) -> None:
189
- """Add request, `n=1`"""
190
- # 1) Process raw inputs into the request.
191
- request = self.processor.process_inputs(request_id, prompt, params,
192
- arrival_time, lora_request,
193
- trace_headers,
194
- prompt_adapter_request,
195
- priority)
196
-
197
- # 2) Make a new RequestState and queue.
198
- self.output_processor.add_request(request)
199
-
200
- # 3) Add the request to EngineCore.
201
- self.engine_core.add_request(request)
156
+ # 1) Fan out child requests (for n>1)
157
+ parent_req = ParentRequest.from_params(request_id, params)
158
+ n = params.n if isinstance(params, SamplingParams) else 1
159
+ for idx in range(n):
160
+ if parent_req is not None:
161
+ request_id, params = parent_req.get_child_info(idx)
162
+
163
+ # 2) Process raw inputs into the request.
164
+ request = self.processor.process_inputs(request_id, prompt, params,
165
+ arrival_time, lora_request,
166
+ trace_headers,
167
+ prompt_adapter_request,
168
+ priority)
169
+
170
+ # 3) Make a new RequestState and queue.
171
+ self.output_processor.add_request(request, parent_req, idx)
172
+
173
+ # 3) Add the request to EngineCore.
174
+ self.engine_core.add_request(request)
175
176
def step(self) -> list[RequestOutput]:
177
@@ -217,10 +190,7 @@ class LLMEngine:
190
# 3) Abort any reqs that finished due to stop strings.
191
self.engine_core.abort_requests(processed_outputs.reqs_to_abort)
192
220
- request_outputs = processed_outputs.request_outputs
221
-
222
- # 4) Process unfinished parallel sampling requests
223
- return self.parallel_manager.step(request_outputs)
193
+ return processed_outputs.request_outputs
194
195
def get_model_config(self):
196
return self.model_config
vllm/v1/engine/output_processor.py
+112
-69
@@ -4,13 +4,14 @@ import asyncio
4
from dataclasses import dataclass
5
from typing import Optional, Union
6
7
-from vllm.outputs import RequestOutput
7
+from vllm.outputs import CompletionOutput, RequestOutput
8
from vllm.sampling_params import RequestOutputKind
9
from vllm.transformers_utils.tokenizer import AnyTokenizer
10
from vllm.transformers_utils.tokenizer_group import BaseTokenizerGroup
11
from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest, FinishReason
12
from vllm.v1.engine.detokenizer import IncrementalDetokenizer
13
from vllm.v1.engine.logprobs import LogprobsProcessor
14
+from vllm.v1.engine.parallel_sampling import ParentRequest
15
from vllm.v1.metrics.stats import (IterationStats, LoRARequestStates,
16
RequestStateStats)
17
@@ -27,6 +28,8 @@ class RequestState:
28
def __init__(
29
self,
30
request_id: str,
31
+ parent_req: Optional[ParentRequest],
32
+ request_index: int,
33
lora_name: Optional[str],
34
output_kind: RequestOutputKind,
35
prompt: Optional[str],
@@ -38,6 +41,8 @@ class RequestState:
41
log_stats: bool,
42
):
43
self.request_id = request_id
44
+ self.parent_req = parent_req
45
+ self.request_index = request_index
46
self.lora_name = lora_name
47
self.output_kind = output_kind
48
self.prompt = prompt
@@ -56,11 +61,15 @@ class RequestState:
61
cls,
62
tokenizer: AnyTokenizer,
63
request: EngineCoreRequest,
64
+ parent_req: Optional[ParentRequest],
65
+ request_index: int,
66
queue: Optional[asyncio.Queue[RequestOutput]],
67
log_stats: bool,
68
) -> "RequestState":
69
return cls(
70
request_id=request.request_id,
71
+ parent_req=parent_req,
72
+ request_index=request_index,
73
lora_name=(request.lora_request.name
74
if request.lora_request is not None else None),
75
output_kind=request.sampling_params.output_kind,
@@ -79,6 +88,88 @@ class RequestState:
88
log_stats=log_stats,
89
)
90
91
+ def make_request_output(
92
+ self,
93
+ new_token_ids: list[int],
94
+ finish_reason: Optional[FinishReason],
95
+ stop_reason: Union[int, str, None],
96
+ ) -> Optional[RequestOutput]:
97
+
98
+ finished = finish_reason is not None
99
+ output_kind = self.output_kind
100
+ final_only = output_kind == RequestOutputKind.FINAL_ONLY
101
+
102
+ # In follow up, we will switch to invariant where EngineCore
103
+ # does not stream partial prefills.
104
+ if not finished and (self.is_prefilling or final_only):
105
+ # Only the final output is required in FINAL_ONLY mode.
106
+ return None
107
+
108
+ def new_request_output(request_id: str) -> RequestOutput:
109
+ return self._new_request_output(request_id, finished)
110
+
111
+ completion_output = self._new_completion_output(
112
+ new_token_ids, finish_reason, stop_reason)
113
+
114
+ if self.parent_req is not None:
115
+ return self.parent_req.make_request_output(final_only,
116
+ completion_output,
117
+ new_request_output)
118
+
119
+ request_output = new_request_output(self.request_id)
120
+ request_output.outputs.append(completion_output)
121
+ return request_output
122
+
123
+ def _new_request_output(
124
+ self,
125
+ request_id: str,
126
+ finished: bool,
127
+ ) -> RequestOutput:
128
+
129
+ if self.output_kind == RequestOutputKind.DELTA:
130
+ # Side effect: logprobs processor forgets prompt logprobs
131
+ prompt_logprobs = self.logprobs_processor.pop_prompt_logprobs()
132
+ else:
133
+ prompt_logprobs = self.logprobs_processor.prompt_logprobs
134
+
135
+ return RequestOutput(
136
+ request_id=request_id,
137
+ prompt=self.prompt,
138
+ prompt_token_ids=self.prompt_token_ids,
139
+ prompt_logprobs=prompt_logprobs,
140
+ outputs=[],
141
+ finished=finished,
142
+ )
143
+
144
+ def _new_completion_output(
145
+ self,
146
+ token_ids: list[int],
147
+ finish_reason: Optional[FinishReason],
148
+ stop_reason: Union[int, str, None],
149
+ ) -> CompletionOutput:
150
+
151
+ finished = finish_reason is not None
152
+ delta = self.output_kind == RequestOutputKind.DELTA
153
+
154
+ # Prepare text and token_ids, based on delta mode
155
+ text = self.detokenizer.get_next_output_text(finished, delta)
156
+ if not delta:
157
+ token_ids = self.detokenizer.output_token_ids
158
+
159
+ # Prepare logprobs, based on delta mode
160
+ logprobs = self.logprobs_processor.logprobs
161
+ if delta and logprobs:
162
+ logprobs = logprobs[-len(token_ids):]
163
+
164
+ return CompletionOutput(
165
+ index=self.request_index,
166
+ text=text,
167
+ token_ids=token_ids,
168
+ logprobs=logprobs,
169
+ cumulative_logprob=self.logprobs_processor.cumulative_logprob,
170
+ finish_reason=str(finish_reason) if finished else None,
171
+ stop_reason=stop_reason if finished else None)
172
+
173
174
class OutputProcessor:
175
"""Process EngineCoreOutputs into RequestOutputs."""
@@ -93,9 +184,6 @@ class OutputProcessor:
184
self.request_states: dict[str, RequestState] = {}
185
self.lora_states = LoRARequestStates()
186
96
- def is_request_active(self, request_id: str) -> bool:
97
- return request_id in self.request_states
98
-
187
def get_num_unfinished_requests(self):
188
return len(self.request_states)
189
@@ -114,6 +202,8 @@ class OutputProcessor:
202
def add_request(
203
self,
204
request: EngineCoreRequest,
205
+ parent_req: Optional[ParentRequest] = None,
206
+ request_index: int = 0,
207
queue: Optional[asyncio.Queue[RequestOutput]] = None,
208
) -> None:
209
request_id = request.request_id
@@ -123,6 +213,8 @@ class OutputProcessor:
213
req_state = RequestState.from_new_request(
214
tokenizer=self.tokenizer.get_lora_tokenizer(request.lora_request),
215
request=request,
216
+ parent_req=parent_req,
217
+ request_index=request_index,
218
queue=queue,
219
log_stats=self.log_stats)
220
self.request_states[request_id] = req_state
@@ -202,8 +294,8 @@ class OutputProcessor:
294
req_state.logprobs_processor.update_from_output(engine_core_output)
295
296
# 4) Create and handle RequestOutput objects.
205
- if request_output := self._make_request_output(
206
- req_state, new_token_ids, finish_reason, stop_reason):
297
+ if request_output := req_state.make_request_output(
298
+ new_token_ids, finish_reason, stop_reason):
299
if req_state.queue is not None:
300
# AsyncLLM: put into queue for handling by generate().
301
req_state.queue.put_nowait(request_output)
@@ -211,18 +303,17 @@ class OutputProcessor:
303
# LLMEngine: return list of RequestOutputs.
304
request_outputs.append(request_output)
305
214
- # Free completed requests.
215
- if request_output.finished:
216
- self.request_states.pop(req_id)
217
- if not engine_core_output.finished:
218
- # If req not finished in EngineCore, but Detokenizer
219
- # detected stop string, abort needed in EngineCore.
220
- reqs_to_abort.append(req_id)
306
+ # Free completed requests.
307
+ if finish_reason is not None:
308
+ self.request_states.pop(req_id)
309
+ if not engine_core_output.finished:
310
+ # If req not finished in EngineCore, but Detokenizer
311
+ # detected stop string, abort needed in EngineCore.
312
+ reqs_to_abort.append(req_id)
313
222
- # Track per-request stats
223
- self._update_stats_from_finished(req_state, request_output,
224
- finish_reason,
225
- iteration_stats)
314
+ # Track per-request stats
315
+ self._update_stats_from_finished(req_state, finish_reason,
316
+ iteration_stats)
317
318
self.lora_states.update_iteration_stats(iteration_stats)
319
@@ -249,7 +340,6 @@ class OutputProcessor:
340
req_state.stats, lora_stats)
341
342
def _update_stats_from_finished(self, req_state: RequestState,
252
- request_output: RequestOutput,
343
finish_reason: Optional[FinishReason],
344
iteration_stats: Optional[IterationStats]):
345
if iteration_stats is None:
@@ -257,55 +347,8 @@ class OutputProcessor:
347
348
assert finish_reason is not None
349
assert req_state.stats is not None
260
- iteration_stats.update_from_finished_request(finish_reason,
261
- request_output,
262
- req_state.stats)
350
+ iteration_stats.update_from_finished_request(
351
+ finish_reason=finish_reason,
352
+ num_prompt_tokens=len(req_state.prompt_token_ids),
353
+ req_stats=req_state.stats)
354
self.lora_states.finish_request(req_state)
264
-
265
- @staticmethod
266
- def _make_request_output(
267
- request_state: RequestState,
268
- new_token_ids: list[int],
269
- finish_reason: Optional[FinishReason],
270
- stop_reason: Union[int, str, None],
271
- ) -> Optional[RequestOutput]:
272
-
273
- finished = finish_reason is not None
274
- output_kind = request_state.output_kind
275
- # In follow up, we will switch to invariant where EngineCore
276
- # does not stream partial prefills.
277
- if not finished and (request_state.is_prefilling
278
- or output_kind == RequestOutputKind.FINAL_ONLY):
279
- # Only the final output is required in FINAL_ONLY mode.
280
- return None
281
-
282
- detokenizer = request_state.detokenizer
283
- logprobs_processor = request_state.logprobs_processor
284
-
285
- delta = output_kind == RequestOutputKind.DELTA
286
- logprobs = logprobs_processor.logprobs
287
- if delta:
288
- if logprobs:
289
- logprobs = logprobs[-len(new_token_ids):]
290
- # Side effect: logprobs processor forgets prompt logprobs
291
- prompt_logprobs = logprobs_processor.pop_prompt_logprobs()
292
- else:
293
- prompt_logprobs = logprobs_processor.prompt_logprobs
294
-
295
- request_output = RequestOutput.new(
296
- request_id=request_state.request_id,
297
- prompt=request_state.prompt,
298
- prompt_token_ids=request_state.prompt_token_ids,
299
- text=detokenizer.get_next_output_text(finished, delta),
300
- token_ids=new_token_ids if delta else detokenizer.output_token_ids,
301
- logprobs=logprobs,
302
- prompt_logprobs=prompt_logprobs,
303
- cumulative_logprob=logprobs_processor.cumulative_logprob,
304
- finished=finished,
305
- )
306
- if finished:
307
- completion_output = request_output.outputs[0]
308
- completion_output.finish_reason = str(finish_reason)
309
- completion_output.stop_reason = stop_reason
310
-
311
- return request_output
vllm/v1/engine/parallel_sampling.py
+44
-300
@@ -1,69 +1,46 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
-from collections.abc import AsyncGenerator, Mapping
3
from copy import copy
5
-from typing import Optional, Protocol, Union
4
+from typing import Callable, Optional, Union
5
7
-from vllm.inputs import PromptType
8
-from vllm.lora.request import LoRARequest
9
-from vllm.outputs import RequestOutput
6
+from vllm.outputs import CompletionOutput, RequestOutput
7
from vllm.pooling_params import PoolingParams
11
-from vllm.prompt_adapter.request import PromptAdapterRequest
12
-from vllm.sampling_params import RequestOutputKind, SamplingParams
13
-from vllm.utils import merge_async_iterators
8
+from vllm.sampling_params import SamplingParams
9
10
16
-class AsyncGenerateMethodType(Protocol):
17
-
18
- def __call__(self,
19
- prompt: PromptType,
20
- sampling_params: SamplingParams,
21
- request_id: str,
22
- lora_request: Optional[LoRARequest] = None,
23
- trace_headers: Optional[Mapping[str, str]] = None,
24
- prompt_adapter_request: Optional[PromptAdapterRequest] = None,
25
- priority: int = 0) -> AsyncGenerator[RequestOutput, None]:
26
- ...
27
-
28
-
29
-class SyncAddRequestMethodType(Protocol):
30
-
31
- def __call__(self,
32
- request_id: str,
33
- prompt: PromptType,
34
- params: Union[SamplingParams, PoolingParams],
35
- arrival_time: Optional[float] = None,
36
- lora_request: Optional[LoRARequest] = None,
37
- trace_headers: Optional[Mapping[str, str]] = None,
38
- prompt_adapter_request: Optional[PromptAdapterRequest] = None,
39
- priority: int = 0) -> None:
40
- ...
41
-
42
-
43
-class ParallelSamplingRequest:
11
+class ParentRequest:
12
"""Info, state & processing for parallel sampling request.
45
-
13
+
14
Store parent request ID and sampling params.
15
Facilitate generating child request sampling params.
48
- Transform child request outputs into parent request
49
- outputs.
50
- When stream mode is disabled, then `self.request_output`
51
- aggregates child request completions.
16
"""
17
18
request_id: str
19
sampling_params: SamplingParams
20
+
21
+ # To aggregate child completions when not streaming
22
+ output_aggregator: Optional[RequestOutput]
23
+
24
+ # To efficiently obtain child sampling params
25
cached_child_sampling_params: Optional[SamplingParams]
57
- request_output: Optional[RequestOutput]
58
- num_finished_completions: int
26
27
def __init__(self, request_id: str,
28
sampling_params: SamplingParams) -> None:
29
self.request_id = request_id
30
self.sampling_params = sampling_params
31
+
32
+ self.output_aggregator = None
33
self.cached_child_sampling_params = None
65
- self.request_output = None
66
- self.num_finished_completions = 0
34
+
35
+ @classmethod
36
+ def from_params(
37
+ cls,
38
+ request_id: str,
39
+ params: Union[SamplingParams, PoolingParams],
40
+ ) -> Optional['ParentRequest']:
41
+ if not isinstance(params, SamplingParams) or params.n == 1:
42
+ return None
43
+ return cls(request_id, params)
44
45
def _get_child_sampling_params(
46
self,
@@ -96,47 +73,6 @@ class ParallelSamplingRequest:
73
child_sampling_params.seed = seed + index
74
return child_sampling_params
75
99
- def _add_output(
100
- self,
101
- child_req_output: RequestOutput,
102
- index: int,
103
- ) -> None:
104
- """Aggregate a parallel sampling child
105
- request output.
106
-
107
- Non-stream-mode (`output_kind == FINAL_ONLY`)
108
- only. Inject correct parent request ID and
109
- completion index.
110
-
111
- Args:
112
- child_req_output: a single request output
113
- from a parallel sampling
114
- child request.
115
- index: index within `n` child
116
- """
117
- self.num_finished_completions += 1
118
- new_completion = child_req_output.outputs[0]
119
- new_completion.index = index
120
- if self.request_output is None:
121
- # Save the first request output; reinstate
122
- # original request ID; metrics are not
123
- # supported for parallel sampling
124
- child_req_output.request_id = self.request_id
125
- child_req_output.metrics = None
126
- self.request_output = child_req_output
127
- else:
128
- # Aggregate additional completion into request output
129
- # Note: will be sorted by index later
130
- self.request_output.outputs.append(new_completion)
131
-
132
- def _get_final_request_output(self) -> RequestOutput:
133
- """Invariant: parent completion outputs sorted by index"""
134
- assert self.request_output is not None
135
- self.request_output.finished = True
136
- self.request_output.outputs = sorted(self.request_output.outputs,
137
- key=lambda x: x.index)
138
- return self.request_output
139
-
76
def get_child_info(self, index: int) -> tuple[str, SamplingParams]:
77
"""Get child request ID and sampling params.
78
@@ -149,227 +85,35 @@ class ParallelSamplingRequest:
85
return (f"{index}_{self.request_id}",
86
self._get_child_sampling_params(index))
87
152
- def process_output(
153
- self,
154
- child_req_output: RequestOutput,
155
- index: int,
156
- ) -> Optional[RequestOutput]:
157
- """Filter, aggregate and transform parallel sampling
158
- child request outputs.
159
-
160
- If the parent request has `stream=false`
161
- (`output_kind == FINAL_ONLY`), each child will also have
162
- `output_kind == FINAL_ONLY`. All child request outputs
163
- must be aggregated into a single request output, with
164
- multiple completions. This request output is only returned
165
- once `n` completions are aggregated.
166
-
167
- If the parent request has `stream=true`
168
- (`output_kind == DELTA`), each child will also have
169
- `output_kind == DELTA`. All child request outputs
170
- must be streamed directly to the caller.
171
-
172
- Args:
173
- child_req_output: a single child request output
174
- index: index within `n` child requests
175
-
176
- Returns:
177
- `None`, unless a processed request output is ready to
178
- send back to the caller.
179
- """
180
- if self.output_kind != RequestOutputKind.FINAL_ONLY:
181
- # stream=true: return child completions immediately
182
- child_req_output.request_id = self.request_id
183
- child_req_output.outputs[0].index = index
184
- if child_req_output.finished:
185
- # Parent request is complete if all child requests are
186
- # complete.
187
- self.num_finished_completions += 1
188
- child_req_output.finished = (
189
- self.num_finished_completions == self.n)
190
- return child_req_output
191
-
192
- # stream=false: aggregate child completions
193
- self._add_output(child_req_output, index)
194
- if self.num_finished_completions == self.n:
195
- # Return aggregated request output after obtaining
196
- # all completions
197
- return self._get_final_request_output()
198
- return None
199
-
200
- async def wrap_child_async_generator(
201
- self,
202
- child_gen: AsyncGenerator[RequestOutput, None],
203
- index: int,
204
- ) -> AsyncGenerator[RequestOutput, None]:
205
- """Output generator for a single parallel sampling
206
- child request.
207
-
208
- Each parallel sampling request triggers at
209
- least two child requests. This generator
210
- yields zero or more request outputs to
211
- return to the caller, as they become
212
- available.
213
-
214
- Args:
215
- child_gen: generator for child request
216
- outputs.
217
- index: index within the `n` child requests
218
-
219
- Returns:
220
- Yields zero or more request outputs to return
221
- to the caller.
222
- """
223
- async for out in child_gen:
224
- if req_out := self.process_output(out, index):
225
- yield req_out
226
-
88
@property
89
def n(self) -> int:
90
return self.sampling_params.n
91
231
- @property
232
- def output_kind(self) -> RequestOutputKind:
233
- return self.sampling_params.output_kind
234
-
235
-
236
-class SyncParallelSamplingManager:
237
-
238
- def __init__(self):
239
- # Parent req ID -> parent request manager
240
- self.parent_reqs: dict[str, ParallelSamplingRequest] = {}
241
- # Child req ID -> (child req index, parent req ID)
242
- self.child_reqs: dict[str, tuple[int, str]] = {}
243
-
244
- def _register_parent_request(self, req: ParallelSamplingRequest) -> None:
245
- """Register parallel sampling parent request."""
246
- self.parent_reqs[req.request_id] = req
247
-
248
- def _register_child_request(self, req_id: str, child_req_id: str,
249
- index: int) -> None:
250
- """Register parallel sampling child request with parent.
251
-
252
- Args:
253
- req_id: parent request ID
254
- child_req_id: child request ID
255
- index: child request index within `n` child requests
256
- """
257
- self.child_reqs[child_req_id] = (index, req_id)
258
-
259
- def get_num_unfinished_requests(self, num_core_reqs: int) -> int:
260
- """Get the number of unfinished requests, correcting for parallel
261
- sampling.
262
-
263
- Args:
264
- num_core_reqs: The number of unfinished requests in the engine core.
265
-
266
- Returns:
267
- Number of unfinished requests, where each parallel sampling req
268
- counts as 1
269
- """
270
- return num_core_reqs + len(self.parent_reqs) - len(self.child_reqs)
271
-
272
- def add_request_parallel_sampling(
92
+ def make_request_output(
93
self,
274
- add_request: SyncAddRequestMethodType,
275
- request_id: str,
276
- prompt: PromptType,
277
- params: Union[SamplingParams, PoolingParams],
278
- arrival_time: Optional[float] = None,
279
- lora_request: Optional[LoRARequest] = None,
280
- trace_headers: Optional[Mapping[str, str]] = None,
281
- prompt_adapter_request: Optional[PromptAdapterRequest] = None,
282
- priority: int = 0,
283
- ) -> None:
284
- """Add sync parallel sampling request."""
285
- req = ParallelSamplingRequest(request_id, params)
286
- self._register_parent_request(req)
287
- # Add n child requests with unique request IDs & random seeds and n=1
288
- for idx in range(req.n):
289
- child_req_id, child_params = req.get_child_info(idx)
290
- self._register_child_request(request_id, child_req_id, idx)
291
- add_request(request_id=child_req_id,
292
- prompt=prompt,
293
- params=child_params,
294
- arrival_time=arrival_time,
295
- lora_request=lora_request,
296
- trace_headers=trace_headers,
297
- prompt_adapter_request=prompt_adapter_request,
298
- priority=priority) # type: ignore
299
-
300
- def step(
301
- self,
302
- outputs: list[RequestOutput],
303
- ) -> list[RequestOutput]:
304
- """Build parallel sampling request outputs.
305
-
306
- Extract child request outputs, aggregate them
307
- into parent request output, and return parent
308
- output when complete.
309
-
310
- Do not modify `n=1` requests.
311
-
312
- Args:
313
- outputs: step request outputs. Mix of child request
314
- outputs & `n=1` request outputs.
94
+ final_only: bool,
95
+ completion_output: CompletionOutput,
96
+ new_request_output: Callable[[str], RequestOutput],
97
+ ) -> Optional[RequestOutput]:
98
+ # Use an existing RequestOutput if we're aggregating
99
+ request_output = self.output_aggregator
100
316
- Return:
317
- List of parallel sampling parent request outputs &
318
- unmodified `n=1` request outputs passed-thru from input.
319
- """
320
- if not (self.parent_reqs and outputs):
321
- # Return unmodified
322
- return outputs
323
- agg_outputs = []
324
- for output in outputs:
325
- req_id = output.request_id
326
- if child_req_entry := self.child_reqs.get(req_id, None):
327
- # For each parallel sampling child request output:
328
- (index, parent_req_id) = child_req_entry
329
- req = self.parent_reqs[parent_req_id]
330
- # Update parallel sampling request
331
- if out := req.process_output(output, index):
332
- # Return parent request output if complete;
333
- # cleanup parent request bookkeeping.
334
- agg_outputs.append(out)
335
- del self.parent_reqs[parent_req_id]
336
- # Cleanup child request bookkeeping.
337
- del self.child_reqs[req_id]
338
- else:
339
- # Not a parallel sampling request output
340
- agg_outputs.append(output)
341
- return agg_outputs
101
+ # Make new RequestOutput otherwise
102
+ if request_output is None:
103
+ request_output = new_request_output(self.request_id)
104
105
+ # Add a new completion
106
+ request_output.outputs.append(completion_output)
107
344
-async def generate_parallel_sampling_async(
345
- generate: AsyncGenerateMethodType,
346
- prompt: PromptType,
347
- sampling_params: SamplingParams,
348
- request_id: str,
349
- lora_request: Optional[LoRARequest] = None,
350
- trace_headers: Optional[Mapping[str, str]] = None,
351
- prompt_adapter_request: Optional[PromptAdapterRequest] = None,
352
- priority: int = 0,
353
-) -> AsyncGenerator[RequestOutput, None]:
354
- """Generate completions for async parallel sampling requests."""
355
- parent_req = ParallelSamplingRequest(request_id, sampling_params)
108
+ # If not streaming, aggregate until all child requests complete
109
+ if final_only and len(request_output.outputs) != self.n:
110
+ self.output_aggregator = request_output
111
+ return None
112
357
- # Aggregate generators for n child requests
358
- gens: list[AsyncGenerator[RequestOutput, None]] = []
359
- for idx in range(parent_req.n):
360
- child_req_id, child_params = parent_req.get_child_info(idx)
361
- child_gen = generate(
362
- prompt=prompt,
363
- sampling_params=child_params,
364
- request_id=child_req_id,
365
- lora_request=lora_request,
366
- trace_headers=trace_headers,
367
- prompt_adapter_request=prompt_adapter_request,
368
- priority=priority,
369
- ) # type: ignore
370
- gen = parent_req.wrap_child_async_generator(child_gen, idx)
371
- gens.append(gen)
113
+ # We're done aggregating
114
+ self.output_aggregator = None
115
373
- # Merge generators
374
- async for _, out in merge_async_iterators(*gens):
375
- yield out
116
+ # Parent completion output list must be sorted by index
117
+ request_output.outputs = sorted(request_output.outputs,
118
+ key=lambda x: x.index)
119
+ return request_output
vllm/v1/metrics/stats.py
+2
-3
@@ -5,7 +5,6 @@ from dataclasses import dataclass, field
5
from typing import TYPE_CHECKING, Optional
6
7
if TYPE_CHECKING:
8
- from vllm.outputs import RequestOutput
8
from vllm.v1.engine import EngineCoreEvent, EngineCoreOutput, FinishReason
9
from vllm.v1.output_processor import RequestState
10
@@ -150,7 +149,7 @@ class IterationStats:
149
self.num_preempted_reqs += 1
150
151
def update_from_finished_request(self, finish_reason: "FinishReason",
153
- request_output: "RequestOutput",
152
+ num_prompt_tokens: int,
153
req_stats: RequestStateStats):
154
e2e_latency = self._time_since(req_stats.arrival_time)
155
@@ -172,7 +171,7 @@ class IterationStats:
171
finished_req = \
172
FinishedRequestStats(finish_reason=finish_reason,
173
e2e_latency=e2e_latency,
175
- num_prompt_tokens=len(request_output.prompt_token_ids),
174
+ num_prompt_tokens=num_prompt_tokens,
175
num_generation_tokens=req_stats.num_generation_tokens,
176
queued_time=queued_time,
177
prefill_time=prefill_time,