@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / befc402d34
Fork 0

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, 22:15 UTC

[V1] V1 engine implements parallel sampling (AsyncLLM and LLMEngine) (#10980)

Signed-off-by: Andrew Feldman <afeldman@neuralmagic.com> Co-authored-by: Nick Hill <nhill@redhat.com>

afeldman-nm committed Feb 24, 2025 at 11:29 UTC befc402d34c2563477ff33bfdd9548ae20a42acd
5 files changed +641 -9
tests/v1/engine/test_llm_engine.py
+98 -5
@@ -1,21 +1,114 @@
1 # SPDX-License-Identifier: Apache-2.0
2
3 +import random
4 +from typing import Dict, List, Optional, Tuple
5 +
6 import pytest
7
8 from tests.v1.engine.utils import PLP_APC_UNSUPPORTED_MSG
9 from vllm import LLM, SamplingParams
10
11 +MODEL = "facebook/opt-125m"
12 +DTYPE = "half"
13
9 -def test_llm_engine_refuses_prompt_logprobs_with_apc(monkeypatch):
10 - """Test passes if LLMEngine raises an exception when it is configured
11 - for automatic prefix caching and it receives a request with
12 - prompt_logprobs enabled, which is incompatible."""
14
15 +def _vllm_model(apc: bool, vllm_runner, monkeypatch):
16 + """Set up VllmRunner instance."""
17 monkeypatch.setenv("VLLM_USE_V1", "1")
18 # TODO(nick): Single-proc to work around a ZMQ shutdown hang for now.
19 monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
20 + return vllm_runner(
21 + MODEL,
22 + dtype=DTYPE,
23 + max_model_len=128,
24 + enforce_eager=True,
25 + enable_prefix_caching=apc,
26 + gpu_memory_utilization=0.5,
27 + )
28 +
29 +
30 +@pytest.fixture(
31 + # Function scope decouples tests & allows
32 + # env var adjustment via monkeypatch
33 + scope="function",
34 + # Prefix caching
35 + params=[False, True])
36 +def vllm_model(vllm_runner, request, monkeypatch):
37 + """VllmRunner test fixture parameterized by APC True/False."""
38 + with _vllm_model(request.param, vllm_runner, monkeypatch) as vllm_model:
39 + yield vllm_model
40 +
41 +
42 +@pytest.fixture(scope="function")
43 +def vllm_model_apc(vllm_runner, monkeypatch):
44 + """VllmRunner test fixture with APC."""
45 + with _vllm_model(True, vllm_runner, monkeypatch) as vllm_model:
46 + yield vllm_model
47 +
48 +
49 +def _get_test_sampling_params(
50 + prompt_list: List[str],
51 + seed: Optional[int] = 42,
52 +) -> Tuple[List[SamplingParams], List[int]]:
53 + """Generate random sampling params for a batch."""
54 +
55 + def get_mostly_n_gt1() -> int:
56 + """Mostly n \in [2,20], ~1/3 n=1"""
57 + x = random.randint(0, 28)
58 + if x < 10:
59 + return 1
60 + else:
61 + return x - 8
62 +
63 + n_list = [get_mostly_n_gt1() for _ in range(len(prompt_list))]
64 + # High temperature to maximize the chance of unique completions
65 + return [
66 + SamplingParams(temperature=0.95, top_p=0.95, n=n, seed=seed)
67 + for n in n_list
68 + ], n_list
69 +
70 +
71 +def test_parallel_sampling(vllm_model, example_prompts) -> None:
72 + """Test passes if parallel sampling `n>1` yields `n` unique completions.
73 +
74 + Args:
75 + vllm_model: VllmRunner instance under test.
76 + example_prompt: test fixture providing prompts for testing.
77 + """
78 + sampling_params_list, n_list = _get_test_sampling_params(example_prompts)
79 + model: LLM = vllm_model.model
80 + outputs = model.generate(example_prompts, sampling_params_list)
81 +
82 + # Validate each request response
83 + for out, n in zip(outputs, n_list):
84 + completion_counts: Dict[str, int] = {}
85 + # Assert correct number of completions
86 + assert len(out.outputs) == n, (
87 + f"{len(out.outputs)} completions; {n} expected.")
88 + for idx in range(n):
89 + comp = out.outputs[idx]
90 + # Assert correct completion indices
91 + assert comp.index == idx, (f"Index {comp.index}; expected {idx}.")
92 + text = comp.text
93 + completion_counts[text] = completion_counts.get(text, 0) + 1
94 + # Assert unique completions
95 + if len(completion_counts) != n:
96 + repeats = {
97 + txt: num
98 + for (txt, num) in completion_counts.items() if num > 1
99 + }
100 + raise AssertionError(
101 + f"{len(completion_counts)} unique completions; expected"
102 + f" {n}. Repeats: {repeats}")
103 +
104 +
105 +def test_llm_engine_refuses_prompt_logprobs_with_apc(vllm_model_apc):
106 + """Test passes if LLMEngine raises an exception when it is configured
107 + for automatic prefix caching and it receives a request with
108 + prompt_logprobs enabled, which is incompatible."""
109 + model: LLM = vllm_model_apc.model
110 with pytest.raises(ValueError) as excinfo:
18 - LLM(model="facebook/opt-125m", enable_prefix_caching=True).generate(
111 + model.generate(
112 "Hello, my name is",
113 SamplingParams(temperature=0.8, top_p=0.95, prompt_logprobs=5))
114
tests/v1/entrypoints/openai/test_completion.py
+102
@@ -250,6 +250,108 @@ async def test_completion_streaming(client: openai.AsyncOpenAI,
250 assert "".join(chunks) == single_output
251
252
253 +@pytest.mark.asyncio
254 +@pytest.mark.parametrize(
255 + "model_name",
256 + [MODEL_NAME],
257 +)
258 +async def test_parallel_no_streaming(client: openai.AsyncOpenAI,
259 + model_name: str):
260 + """Parallel sampling without streaming.
261 + A single request output contains a list of completions.
262 + """
263 +
264 + prompt = "What is an LLM?"
265 + n = 3
266 + max_tokens = 5
267 +
268 + # High temperature to maximize chance of unique completions.
269 + completion = await client.completions.create(model=model_name,
270 + prompt=prompt,
271 + max_tokens=max_tokens,
272 + n=n,
273 + temperature=0.95,
274 + stream=False,
275 + seed=42)
276 +
277 + # Assert `n` completions
278 + num_completions = len(completion.choices)
279 + assert num_completions == n, (
280 + f"Num completions {num_completions} but expected {n}.")
281 + completion_repeats: Dict[str, int] = {}
282 + for idx, choice in enumerate(completion.choices):
283 + # Assert correct completion index & some finish reason.
284 + assert choice.index == idx, (
285 + f"Index {choice.index} but expected {idx}.")
286 + assert choice.finish_reason is not None, (
287 + "None finish_reason is invalid.")
288 + text = choice.text
289 + completion_repeats[text] = completion_repeats.get(text, 0) + 1
290 + # Assert `n` unique completions
291 + num_unique = len(completion_repeats)
292 + if num_unique != n:
293 + repeats = {
294 + txt: num
295 + for (txt, num) in completion_repeats.items() if num > 1
296 + }
297 + raise AssertionError(
298 + f"Expected {n} unique completions, got {num_unique};"
299 + f" repeats: {repeats}.")
300 +
301 +
302 +@pytest.mark.asyncio
303 +@pytest.mark.parametrize(
304 + "model_name",
305 + [MODEL_NAME],
306 +)
307 +async def test_parallel_streaming(client: openai.AsyncOpenAI, model_name: str):
308 + """Streaming for parallel sampling.
309 + The tokens from multiple samples, are flattened into a single stream,
310 + with an index to indicate which sample the token belongs to.
311 + """
312 +
313 + prompt = "What is an LLM?"
314 + n = 3
315 + max_tokens = 5
316 +
317 + stream = await client.completions.create(model=model_name,
318 + prompt=prompt,
319 + max_tokens=max_tokens,
320 + n=n,
321 + temperature=0.95,
322 + stream=True,
323 + seed=42)
324 + chunks: List[List[str]] = [[] for i in range(n)]
325 + finish_reason_count = 0
326 + async for chunk in stream:
327 + index = chunk.choices[0].index
328 + text = chunk.choices[0].text
329 + chunks[index].append(text)
330 + if chunk.choices[0].finish_reason is not None:
331 + finish_reason_count += 1
332 + # Assert `n` completions with correct finish reasons
333 + assert finish_reason_count == n, (
334 + f"Expected {n} completions with valid indices and finish_reason.")
335 + completion_repeats: Dict[str, int] = {}
336 + for chunk in chunks:
337 + chunk_len = len(chunk)
338 + # Assert correct number of completion tokens
339 + assert chunk_len == max_tokens, (
340 + f"max_tokens={max_tokens} but chunk len is {chunk_len}.")
341 + text = "".join(chunk)
342 + completion_repeats[text] = completion_repeats.get(text, 0) + 1
343 + print(text)
344 + # Assert `n` unique completions
345 + num_unique = len(completion_repeats)
346 + if num_unique != n:
347 + repeats = {
348 + txt: num
349 + for (txt, num) in completion_repeats.items() if num > 1
350 + }
351 + raise AssertionError(f"{num_unique} unique completions, expected {n};"
352 + f" repeats: {repeats}")
353 +
354 +
355 @pytest.mark.asyncio
356 @pytest.mark.parametrize(
357 "model_name",
vllm/v1/engine/async_llm.py
+26 -1
@@ -24,6 +24,7 @@ from vllm.usage.usage_lib import UsageContext
24 from vllm.utils import cdiv, kill_process_tree
25 from vllm.v1.engine.core_client import EngineCoreClient
26 from vllm.v1.engine.output_processor import OutputProcessor
27 +from vllm.v1.engine.parallel_sampling import generate_parallel_sampling_async
28 from vllm.v1.engine.processor import Processor
29 from vllm.v1.executor.abstract import Executor
30 from vllm.v1.metrics.loggers import (LoggingStatLogger, PrometheusStatLogger,
@@ -170,7 +171,7 @@ class AsyncLLM(EngineClient):
171 # requests we don't need to send multiple messages to core proc,
172 # and so we don't need multiple streams which then get
173 # re-multiplexed in the API server anyhow.
173 - async def generate(
174 + async def _generate(
175 self,
176 prompt: PromptType,
177 sampling_params: SamplingParams,
@@ -241,6 +242,30 @@ class AsyncLLM(EngineClient):
242 await self.abort(request_id)
243 raise
244
245 + def generate(
246 + self,
247 + prompt: PromptType,
248 + sampling_params: SamplingParams,
249 + request_id: str,
250 + lora_request: Optional[LoRARequest] = None,
251 + trace_headers: Optional[Mapping[str, str]] = None,
252 + prompt_adapter_request: Optional[PromptAdapterRequest] = None,
253 + priority: int = 0,
254 + ) -> AsyncGenerator[RequestOutput, None]:
255 + kwargs = dict(prompt=prompt,
256 + sampling_params=sampling_params,
257 + request_id=request_id,
258 + lora_request=lora_request,
259 + trace_headers=trace_headers,
260 + prompt_adapter_request=prompt_adapter_request,
261 + priority=priority)
262 + if sampling_params.n is None or sampling_params.n == 1:
263 + return self._generate(**kwargs)
264 + else:
265 + # Special handling for parallel sampling requests
266 + return generate_parallel_sampling_async(generate=self._generate,
267 + **kwargs)
268 +
269 async def _run_output_handler(self):
270 """Background loop: pulls from EngineCore and pushes to AsyncStreams."""
271
vllm/v1/engine/llm_engine.py
+40 -3
@@ -21,6 +21,7 @@ from vllm.transformers_utils.tokenizer_group import (
21 from vllm.usage.usage_lib import UsageContext
22 from vllm.v1.engine.core_client import EngineCoreClient
23 from vllm.v1.engine.output_processor import OutputProcessor
24 +from vllm.v1.engine.parallel_sampling import SyncParallelSamplingManager
25 from vllm.v1.engine.processor import Processor
26 from vllm.v1.executor.abstract import Executor
27
@@ -48,6 +49,9 @@ class LLMEngine:
49 self.model_config = vllm_config.model_config
50 self.cache_config = vllm_config.cache_config
51
52 + # Bookkeeping for parallel sampling requests
53 + self.parallel_manager = SyncParallelSamplingManager()
54 +
55 # important: init dp group before init the engine_core
56 self.parallel_config = vllm_config.parallel_config
57 self.dp_enabled = self.parallel_config.data_parallel_size > 1 # noqa
@@ -115,7 +119,8 @@ class LLMEngine:
119 multiprocess_mode=enable_multiprocessing)
120
121 def get_num_unfinished_requests(self) -> int:
118 - return self.output_processor.get_num_unfinished_requests()
122 + return self.parallel_manager.get_num_unfinished_requests(
123 + self.output_processor.get_num_unfinished_requests())
124
125 def has_unfinished_requests(self) -> bool:
126 has_unfinished = self.output_processor.has_unfinished_requests()
@@ -151,7 +156,36 @@ class LLMEngine:
156 prompt_adapter_request: Optional[PromptAdapterRequest] = None,
157 priority: int = 0,
158 ) -> None:
154 -
159 + """Add request."""
160 + kwargs = dict(request_id=request_id,
161 + prompt=prompt,
162 + params=params,
163 + arrival_time=arrival_time,
164 + lora_request=lora_request,
165 + trace_headers=trace_headers,
166 + prompt_adapter_request=prompt_adapter_request,
167 + priority=priority)
168 + # Handle parallel sampling requests differently.
169 + if params is None or isinstance(params,
170 + PoolingParams) or params.n == 1:
171 + self._add_request(**kwargs)
172 + else:
173 + # Special handling for parallel sampling requests
174 + self.parallel_manager.add_request_parallel_sampling(
175 + add_request=self._add_request, **kwargs)
176 +
177 + def _add_request(
178 + self,
179 + request_id: str,
180 + prompt: PromptType,
181 + params: Union[SamplingParams, PoolingParams],
182 + arrival_time: Optional[float] = None,
183 + lora_request: Optional[LoRARequest] = None,
184 + trace_headers: Optional[Mapping[str, str]] = None,
185 + prompt_adapter_request: Optional[PromptAdapterRequest] = None,
186 + priority: int = 0,
187 + ) -> None:
188 + """Add request, `n=1`"""
189 # 1) Process raw inputs into the request.
190 request = self.processor.process_inputs(request_id, prompt, params,
191 arrival_time, lora_request,
@@ -182,7 +216,10 @@ class LLMEngine:
216 # 3) Abort any reqs that finished due to stop strings.
217 self.engine_core.abort_requests(processed_outputs.reqs_to_abort)
218
185 - return processed_outputs.request_outputs
219 + request_outputs = processed_outputs.request_outputs
220 +
221 + # 4) Process unfinished parallel sampling requests
222 + return self.parallel_manager.step(request_outputs)
223
224 def get_model_config(self):
225 return self.model_config
vllm/v1/engine/parallel_sampling.py new
+375
@@ -0,0 +1,375 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +from copy import copy
4 +from typing import (AsyncGenerator, Dict, List, Mapping, Optional, Protocol,
5 + Tuple, Union)
6 +
7 +from vllm.inputs import PromptType
8 +from vllm.lora.request import LoRARequest
9 +from vllm.outputs import RequestOutput
10 +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
14 +
15 +
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:
44 + """Info, state & processing for parallel sampling request.
45 +
46 + Store parent request ID and sampling params.
47 + 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.
52 + """
53 +
54 + request_id: str
55 + sampling_params: SamplingParams
56 + cached_child_sampling_params: Optional[SamplingParams]
57 + request_output: Optional[RequestOutput]
58 + num_finished_completions: int
59 +
60 + def __init__(self, request_id: str,
61 + sampling_params: SamplingParams) -> None:
62 + self.request_id = request_id
63 + self.sampling_params = sampling_params
64 + self.cached_child_sampling_params = None
65 + self.request_output = None
66 + self.num_finished_completions = 0
67 +
68 + def _get_child_sampling_params(
69 + self,
70 + index: int,
71 + ) -> SamplingParams:
72 + """Efficiently obtain child `sampling_params`
73 +
74 + If `sampling_params.seed` is not `None` then
75 + each child request requires a unique clone of
76 + parent `sampling_params` with a unique seed.
77 +
78 + Args:
79 + index: index within `n` child requests
80 +
81 + Returns:
82 + Child `sampling_params` instance.
83 + """
84 + seed = self.sampling_params.seed
85 + if self.cached_child_sampling_params:
86 + # Reuse child sampling_params data structure
87 + return self.cached_child_sampling_params
88 + # Build child sampling_params
89 + child_sampling_params = copy(self.sampling_params)
90 + child_sampling_params.n = 1
91 + if seed is None:
92 + # Cache child sampling_params for later reuse
93 + self.cached_child_sampling_params = child_sampling_params
94 + else:
95 + # Each child gets a clone with a unique seed
96 + child_sampling_params.seed = seed + index
97 + return child_sampling_params
98 +
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 +
140 + def get_child_info(self, index: int) -> Tuple[str, SamplingParams]:
141 + """Get child request ID and sampling params.
142 +
143 + Args:
144 + index: index within `n` child requests.
145 +
146 + Returns:
147 + (request ID, sampling_params) tuple
148 + """
149 + return (f"{index}_{self.request_id}",
150 + self._get_child_sampling_params(index))
151 +
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 +
227 + @property
228 + def n(self) -> int:
229 + return self.sampling_params.n
230 +
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(
273 + 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.
315 +
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
342 +
343 +
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)
356 +
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)
372 +
373 + # Merge generators
374 + async for _, out in merge_async_iterators(*gens):
375 + yield out