@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f5d3acd474

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, 19:45 UTC

[BugFix][V1] Fix parallel sampling finishing/aborts (#14512)

Signed-off-by: Nick Hill <nhill@redhat.com>

Nick Hill committed Mar 12, 2025 at 13:29 UTC f5d3acd47466f094beb36f7a5d05520466713f93

7 files changed +137 -113

tests/v1/engine/test_async_llm.py

+50 -6

@@ -46,6 +46,7 @@ async def generate(engine: AsyncLLM,

46 prompt: PromptType,

47 output_kind: RequestOutputKind,

48 max_tokens: int,

49 + n: int = 1,

50 prompt_logprobs: Optional[int] = None) -> tuple[int, str]:

51 # Ensure generate doesn't complete too fast for cancellation test.

52 await asyncio.sleep(0.2)

@@ -54,13 +55,15 @@ async def generate(engine: AsyncLLM,

55 sampling_params = SamplingParams(max_tokens=max_tokens,

56 ignore_eos=True,

57 output_kind=output_kind,

57 - temperature=0,

58 + temperature=0.5,

59 + seed=33,

60 + n=n,

61 prompt_logprobs=prompt_logprobs)

62 async for out in engine.generate(request_id=request_id,

63 prompt=prompt,

64 sampling_params=sampling_params):

63 - num_tokens = len(out.outputs[0].token_ids)

66 + num_tokens = sum(len(output.token_ids) for output in out.outputs)

67 if output_kind == RequestOutputKind.DELTA:

68 count += num_tokens

69 else:

@@ -136,17 +139,22 @@ async def test_abort(monkeypatch, output_kind: RequestOutputKind,

139

140 NUM_REQUESTS = 100

141 NUM_EXPECTED_TOKENS = 100

142 + NUM_EXPECTED_TOKENS_LONG = 50000

143 REQUEST_IDS_TO_ABORT = range(1, 100, 10)

144 + PARALLEL_SAMPLE_REQ_IDS = range(1, 100, 15)

145

146 request_ids = [f"request-{i}" for i in range(NUM_REQUESTS)]

147

148 # Create concurrent requests.

149 tasks: list[asyncio.Task] = []

145 - for request_id in request_ids:

150 + for idx, request_id in enumerate(request_ids):

151 + max_tokens = NUM_EXPECTED_TOKENS_LONG if (

152 + idx in REQUEST_IDS_TO_ABORT) else NUM_EXPECTED_TOKENS

153 + n = 3 if idx in PARALLEL_SAMPLE_REQ_IDS else 1

154 tasks.append(

155 asyncio.create_task(

156 generate(engine, request_id, prompt, output_kind,

149 - NUM_EXPECTED_TOKENS)))

157 + max_tokens, n)))

158

159 # API server cancels requests when they disconnect.

160 for idx in REQUEST_IDS_TO_ABORT:

@@ -162,10 +170,13 @@ async def test_abort(monkeypatch, output_kind: RequestOutputKind,

170 else:

171 # Otherwise, make sure the request was not impacted.

172 num_generated_tokens, request_id = await task

165 - assert num_generated_tokens == NUM_EXPECTED_TOKENS, (

173 + n = 3 if idx in PARALLEL_SAMPLE_REQ_IDS else 1

174 + expected_tokens = NUM_EXPECTED_TOKENS * n

175 + assert num_generated_tokens == expected_tokens, (

176 f"{request_id} generated {num_generated_tokens} but "

167 - f"expected {NUM_EXPECTED_TOKENS}")

177 + f"expected {expected_tokens}")

178

179 + # Make sure all aborted requests were really aborted.

180 assert not engine.output_processor.has_unfinished_requests()

181

182 # Confirm we can do another generation.

@@ -176,3 +187,36 @@ async def test_abort(monkeypatch, output_kind: RequestOutputKind,

187 num_generated_tokens, request_id = await task

188 assert num_generated_tokens == NUM_EXPECTED_TOKENS

189 assert not engine.output_processor.has_unfinished_requests()

190 +

191 +

192 +@pytest.mark.parametrize("n", [1, 3])

193 +@pytest.mark.parametrize("engine_args_and_prompt",

194 + [(TEXT_ENGINE_ARGS, TEXT_PROMPT),

195 + (VISION_ENGINE_ARGS, VISION_PROMPT)])

196 +@pytest.mark.asyncio

197 +async def test_finished_flag(monkeypatch, n: int,

198 + engine_args_and_prompt: tuple[AsyncEngineArgs,

199 + PromptType]):

200 +

201 + with monkeypatch.context() as m, ExitStack() as after:

202 + m.setenv("VLLM_USE_V1", "1")

203 + engine_args, prompt = engine_args_and_prompt

204 +

205 + engine = AsyncLLM.from_engine_args(engine_args)

206 + after.callback(engine.shutdown)

207 +

208 + sampling_params = SamplingParams(max_tokens=100,

209 + output_kind=RequestOutputKind.DELTA,

210 + temperature=1.0,

211 + seed=33,

212 + n=n)

213 + outputs = [

214 + out

215 + async for out in engine.generate(request_id="request-33",

216 + prompt=prompt,

217 + sampling_params=sampling_params)

218 + ]

219 +

220 + # Assert only the last output has the finished flag set

221 + assert all(not out.finished for out in outputs[:-1])

222 + assert outputs[-1].finished

tests/v1/entrypoints/openai/test_completion.py

+15 -6

@@ -263,15 +263,16 @@ async def test_parallel_no_streaming(client: openai.AsyncOpenAI,

263

264 prompt = "What is an LLM?"

265 n = 3

266 - max_tokens = 5

266 + max_tokens = 50 # we want some to finish earlier than others

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,

273 + temperature=1.0,

274 stream=False,

275 + logprobs=0,

276 seed=42)

277

278 # Assert `n` completions

@@ -279,6 +280,7 @@ async def test_parallel_no_streaming(client: openai.AsyncOpenAI,

280 assert num_completions == n, (

281 f"Num completions {num_completions} but expected {n}.")

282 completion_repeats: dict[str, int] = {}

283 + output_token_lengths = set()

284 for idx, choice in enumerate(completion.choices):

285 # Assert correct completion index & some finish reason.

286 assert choice.index == idx, (

@@ -287,6 +289,9 @@ async def test_parallel_no_streaming(client: openai.AsyncOpenAI,

289 "None finish_reason is invalid.")

290 text = choice.text

291 completion_repeats[text] = completion_repeats.get(text, 0) + 1

292 + output_token_lengths.add(len(choice.logprobs.tokens))

293 + # Assert subrequests finished at different times

294 + assert len(output_token_lengths) > 1

295 # Assert `n` unique completions

296 num_unique = len(completion_repeats)

297 if num_unique != n:

@@ -312,16 +317,16 @@ async def test_parallel_streaming(client: openai.AsyncOpenAI, model_name: str):

317

318 prompt = "What is an LLM?"

319 n = 3

315 - max_tokens = 5

320 + max_tokens = 50 # we want some to finish earlier than others

321

322 stream = await client.completions.create(model=model_name,

323 prompt=prompt,

324 max_tokens=max_tokens,

325 n=n,

321 - temperature=0.95,

326 + temperature=1.0,

327 stream=True,

328 seed=42)

324 - chunks: list[list[str]] = [[] for i in range(n)]

329 + chunks: list[list[str]] = [[] for _ in range(n)]

330 finish_reason_count = 0

331 async for chunk in stream:

332 index = chunk.choices[0].index

@@ -333,14 +338,18 @@ async def test_parallel_streaming(client: openai.AsyncOpenAI, model_name: str):

338 assert finish_reason_count == n, (

339 f"Expected {n} completions with valid indices and finish_reason.")

340 completion_repeats: dict[str, int] = {}

341 + chunk_lengths = set()

342 for chunk in chunks:

343 chunk_len = len(chunk)

344 # Assert correct number of completion tokens

339 - assert chunk_len == max_tokens, (

345 + chunk_lengths.add(chunk_len)

346 + assert chunk_len <= max_tokens, (

347 f"max_tokens={max_tokens} but chunk len is {chunk_len}.")

348 text = "".join(chunk)

349 completion_repeats[text] = completion_repeats.get(text, 0) + 1

350 print(text)

351 + # Assert subrequests finished at different times

352 + assert len(chunk_lengths) > 1

353 # Assert `n` unique completions

354 num_unique = len(completion_repeats)

355 if num_unique != n:

vllm/outputs.py

+18 -46

@@ -134,57 +134,29 @@ class RequestOutput:

134 self.encoder_prompt_token_ids = encoder_prompt_token_ids

135 self.num_cached_tokens = num_cached_tokens

136

137 - @classmethod

138 - def new(

139 - cls,

140 - request_id: str,

141 - prompt: Optional[str],

142 - prompt_token_ids: Optional[list[int]],

143 - text: str,

144 - token_ids: list[int],

145 - logprobs: Optional[SampleLogprobs],

146 - prompt_logprobs: Optional[PromptLogprobs],

147 - cumulative_logprob: Optional[float],

148 - finished: bool = False,

149 - ) -> "RequestOutput":

150 - """Initialize a new RequestOutput object."""

151 -

152 - # TODO: Support `n` > 1.

153 - completion_output = CompletionOutput(

154 - index=0,

155 - text=text,

156 - token_ids=token_ids,

157 - cumulative_logprob=cumulative_logprob,

158 - logprobs=logprobs)

159 -

160 - return RequestOutput(

161 - request_id=request_id,

162 - prompt=prompt,

163 - prompt_token_ids=prompt_token_ids,

164 - prompt_logprobs=prompt_logprobs,

165 - outputs=[completion_output],

166 - finished=finished,

167 - )

168 -

137 def add(self, next_output: "RequestOutput") -> None:

138 """Merge subsequent RequestOutput into this one"""

139

172 - self.prompt = next_output.prompt

173 - self.prompt_token_ids = next_output.prompt_token_ids

174 - self.prompt_logprobs = next_output.prompt_logprobs

140 self.finished |= next_output.finished

141

177 - #TODO assuming n == 1 for now

178 - completion = self.outputs[0]

179 - next_completion = next_output.outputs[0]

180 - completion.text += next_completion.text

181 - if not isinstance(completion.token_ids, MutableSequence):

182 - completion.token_ids = list(completion.token_ids)

183 - completion.token_ids.extend(next_completion.token_ids)

184 - if next_completion.logprobs:

185 - assert completion.logprobs is not None

186 - completion.logprobs.extend(next_completion.logprobs)

187 - completion.cumulative_logprob = next_completion.cumulative_logprob

142 + for next_completion in next_output.outputs:

143 + for completion in self.outputs:

144 + if completion.index == next_completion.index:

145 + # Merge outputs with same index

146 + completion.text += next_completion.text

147 + if not isinstance(completion.token_ids, MutableSequence):

148 + completion.token_ids = list(completion.token_ids)

149 + completion.token_ids.extend(next_completion.token_ids)

150 + if next_completion.logprobs:

151 + assert completion.logprobs is not None

152 + completion.logprobs.extend(next_completion.logprobs)

153 + completion.cumulative_logprob = (

154 + next_completion.cumulative_logprob)

155 + completion.finish_reason = next_completion.finish_reason

156 + completion.stop_reason = next_completion.stop_reason

157 + break

158 + else:

159 + self.outputs.append(next_completion)

160

161 @classmethod

162 def from_seq_group(

vllm/v1/engine/async_llm.py

+1 -2

@@ -298,9 +298,8 @@ class AsyncLLM(EngineClient):

298 async def abort(self, request_id: str) -> None:

299 """Abort RequestId in OutputProcessor and EngineCore."""

300

301 - request_ids = [request_id]

301 + request_ids = self.output_processor.abort_requests((request_id, ))

302 await self.engine_core.abort_requests_async(request_ids)

303 - self.output_processor.abort_requests(request_ids)

303

304 if self.log_requests:

305 logger.info("Aborted request %s.", request_id)

vllm/v1/engine/llm_engine.py

+1 -1

@@ -137,8 +137,8 @@ class LLMEngine:

137 def abort_request(self, request_ids: list[str]) -> None:

138 """Remove request_ids from EngineCore and Detokenizer."""

139

140 + request_ids = self.output_processor.abort_requests(request_ids)

141 self.engine_core.abort_requests(request_ids)

141 - self.output_processor.abort_requests(request_ids)

142

143 def add_request(

144 self,

vllm/v1/engine/output_processor.py

+30 -19

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

3 import asyncio

4 +from collections.abc import Iterable

5 from dataclasses import dataclass

6 from typing import Optional, Union

@@ -102,8 +103,7 @@ class RequestState:

103 ) -> Optional[RequestOutput]:

104

105 finished = finish_reason is not None

105 - output_kind = self.output_kind

106 - final_only = output_kind == RequestOutputKind.FINAL_ONLY

106 + final_only = self.output_kind == RequestOutputKind.FINAL_ONLY

107

108 # In follow up, we will switch to invariant where EngineCore

109 # does not stream partial prefills.

@@ -111,24 +111,24 @@ class RequestState:

111 # Only the final output is required in FINAL_ONLY mode.

112 return None

113

114 - def new_request_output(request_id: str) -> RequestOutput:

115 - return self._new_request_output(request_id, finished)

116 -

114 completion_output = self._new_completion_output(

115 new_token_ids, finish_reason, stop_reason)

116

120 - if self.parent_req is not None:

121 - return self.parent_req.make_request_output(final_only,

122 - completion_output,

123 - new_request_output)

117 + request_id = self.request_id

118 + if self.parent_req is None:

119 + outputs = [completion_output]

120 + else:

121 + request_id, outputs, finished = self.parent_req.get_outputs(

122 + request_id, completion_output)

123 + if not outputs:

124 + return None

125

125 - request_output = new_request_output(self.request_id)

126 - request_output.outputs.append(completion_output)

127 - return request_output

126 + return self._new_request_output(request_id, outputs, finished)

127

128 def _new_request_output(

129 self,

130 request_id: str,

131 + outputs: list[CompletionOutput],

132 finished: bool,

133 ) -> RequestOutput:

134

@@ -143,7 +143,7 @@ class RequestState:

143 prompt=self.prompt,

144 prompt_token_ids=self.prompt_token_ids,

145 prompt_logprobs=prompt_logprobs,

146 - outputs=[],

146 + outputs=outputs,

147 finished=finished,

148 )

149

@@ -188,6 +188,7 @@ class OutputProcessor:

188 self.log_stats = log_stats

189 self.tokenizer = tokenizer

190 self.request_states: dict[str, RequestState] = {}

191 + self.parent_requests: dict[str, ParentRequest] = {}

192 self.lora_states = LoRARequestStates()

193

194 def get_num_unfinished_requests(self):

@@ -198,14 +199,20 @@ class OutputProcessor:

199

200 def abort_requests(

201 self,

201 - request_ids: list[str],

202 - ) -> None:

202 + request_ids: Iterable[str],

203 + ) -> list[str]:

204 + request_ids_to_abort = []

205 for request_id in request_ids:

206 req_state = self.request_states.pop(request_id, None)

207 if req_state is not None:

208 self.lora_states.abort_request(req_state)

207 - if req_state.parent_req is not None:

208 - req_state.parent_req.finish_child_request(request_id)

209 + request_ids_to_abort.append(request_id)

210 + else:

211 + parent = self.parent_requests.pop(request_id, None)

212 + if parent and parent.child_requests:

213 + self.abort_requests(parent.child_requests)

214 + request_ids_to_abort.extend(parent.child_requests)

215 + return request_ids_to_abort

216

217 def add_request(

218 self,

@@ -227,6 +234,8 @@ class OutputProcessor:

234 log_stats=self.log_stats)

235 self.request_states[request_id] = req_state

236 self.lora_states.add_request(req_state)

237 + if parent_req:

238 + self.parent_requests[parent_req.request_id] = parent_req

239

240 def process_outputs(

241 self,

@@ -314,12 +323,14 @@ class OutputProcessor:

323 # Free completed requests.

324 if finish_reason is not None:

325 self.request_states.pop(req_id)

326 + # Remove parent request if applicable.

327 + parent_req = req_state.parent_req

328 + if parent_req and not parent_req.child_requests:

329 + self.parent_requests.pop(parent_req.request_id, None)

330 if not engine_core_output.finished:

331 # If req not finished in EngineCore, but Detokenizer

332 # detected stop string, abort needed in EngineCore.

333 reqs_to_abort.append(req_id)

321 - if req_state.parent_req is not None:

322 - req_state.parent_req.finish_child_request(req_id)

334

335 # Track per-request stats

336 self._update_stats_from_finished(req_state, finish_reason,

vllm/v1/engine/parallel_sampling.py

+22 -33

@@ -1,11 +1,11 @@

1 # SPDX-License-Identifier: Apache-2.0

3 from copy import copy

4 -from typing import Callable, Optional, Union

4 +from typing import Optional, Union

6 -from vllm.outputs import CompletionOutput, RequestOutput

6 +from vllm.outputs import CompletionOutput

7 from vllm.pooling_params import PoolingParams

8 -from vllm.sampling_params import SamplingParams

8 +from vllm.sampling_params import RequestOutputKind, SamplingParams

9 from vllm.v1.metrics.stats import IterationStats

@@ -23,7 +23,7 @@ class ParentRequest:

23 child_requests: set[str]

25 # To aggregate child completions when not streaming

26 - output_aggregator: Optional[RequestOutput]

26 + output_aggregator: list[CompletionOutput]

28 # To find the max number of generated tokens across all children

29 max_num_generation_tokens: int

@@ -37,7 +37,9 @@ class ParentRequest:

37 self.sampling_params = sampling_params

39 self.child_requests = set()

40 - self.output_aggregator = None

40 + self.output_aggregator = [None] * sampling_params.n if (

41 + sampling_params.output_kind

42 + == RequestOutputKind.FINAL_ONLY) else []

43 self.max_num_generation_tokens = 0

44 self.cached_child_sampling_params = None

@@ -93,43 +95,30 @@ class ParentRequest:

95 """

96 child_req_id = f"{index}_{self.request_id}"

97 self.child_requests.add(child_req_id)

96 - return (child_req_id, self._get_child_sampling_params(index))

97 -

98 - def finish_child_request(self, req_id: str):

99 - self.child_requests.remove(req_id)

98 + return child_req_id, self._get_child_sampling_params(index)

100 @property

101 def n(self) -> int:

102 return self.sampling_params.n

103

105 - def make_request_output(

104 + def get_outputs(

105 self,

107 - final_only: bool,

106 + child_request_id: str,

107 completion_output: CompletionOutput,

109 - new_request_output: Callable[[str], RequestOutput],

110 - ) -> Optional[RequestOutput]:

111 - # Use an existing RequestOutput if we're aggregating

112 - request_output = self.output_aggregator

113 -

114 - # Make new RequestOutput otherwise

115 - if request_output is None:

116 - request_output = new_request_output(self.request_id)

117 -

118 - # Add a new completion

119 - request_output.outputs.append(completion_output)

108 + ) -> tuple[str, list[CompletionOutput], bool]:

109 + if completion_output.finished():

110 + self.child_requests.remove(child_request_id)

111

121 - # If not streaming, aggregate until all child requests complete

122 - if final_only and len(request_output.outputs) != self.n:

123 - self.output_aggregator = request_output

124 - return None

125 -

126 - # We're done aggregating

127 - self.output_aggregator = None

112 + if self.sampling_params.output_kind != RequestOutputKind.FINAL_ONLY:

113 + # If streaming, just return the current output.

114 + outputs = [completion_output]

115 + else:

116 + # If not streaming, aggregate the n final outputs.

117 + self.output_aggregator[completion_output.index] = completion_output

118 + outputs = [] if self.child_requests else self.output_aggregator

119

129 - # Parent completion output list must be sorted by index

130 - request_output.outputs = sorted(request_output.outputs,

131 - key=lambda x: x.index)

132 - return request_output

120 + finished = not self.child_requests

121 + return self.request_id, outputs, finished

122

123 def observe_num_generation_tokens(self, num_generation_tokens: int):

124 self.max_num_generation_tokens = max(num_generation_tokens,