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

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 9, 13:33 UTC

[Frontend] Cleanup serving engine (#33103)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Cyrus Leung committed Jan 27, 2026 at 12:47 UTC e0b005d9cf9aff4d2f45d9bc9de401c807544e03

5 files changed +57 -40

vllm/entrypoints/openai/chat_completion/serving.py

+4 -5

@@ -68,6 +68,7 @@ from vllm.entrypoints.openai.parser.harmony_utils import (

68 from vllm.entrypoints.openai.utils import maybe_filter_parallel_tool_calls

69 from vllm.entrypoints.utils import get_max_tokens, should_include_usage

70 from vllm.inputs.data import TokensPrompt

71 +from vllm.inputs.parse import get_prompt_components

72 from vllm.logger import init_logger

73 from vllm.logprobs import Logprob

74 from vllm.outputs import CompletionOutput, RequestOutput

@@ -374,20 +375,18 @@ class OpenAIServingChat(OpenAIServing):

375 generators: list[AsyncGenerator[RequestOutput, None]] = []

376 try:

377 for i, engine_prompt in enumerate(engine_prompts):

377 - prompt_text, _, _ = self._get_prompt_components(engine_prompt)

378 + prompt_text, _, _ = get_prompt_components(engine_prompt)

379 +

380 # If we are creating sub requests for multiple prompts, ensure that they

381 # have unique request ids.

382 sub_request_id = (

383 request_id if len(engine_prompts) == 1 else f"{request_id}_{i}"

384 )

385

384 - if self.default_sampling_params is None:

385 - self.default_sampling_params = {}

386 -

386 max_tokens = get_max_tokens(

387 max_model_len=self.max_model_len,

388 request=request,

390 - input_length=len(engine_prompt["prompt_token_ids"]),

389 + prompt=engine_prompt,

390 default_sampling_params=self.default_sampling_params,

391 )

392

vllm/entrypoints/openai/completion/serving.py

+3 -15

@@ -36,6 +36,7 @@ from vllm.entrypoints.renderer import RenderConfig

36 from vllm.entrypoints.utils import get_max_tokens, should_include_usage

37 from vllm.exceptions import VLLMValidationError

38 from vllm.inputs.data import EmbedsPrompt, TokensPrompt, is_embeds_prompt

39 +from vllm.inputs.parse import get_prompt_components

40 from vllm.logger import init_logger

41 from vllm.logprobs import Logprob

42 from vllm.outputs import RequestOutput

@@ -162,25 +163,12 @@ class OpenAIServingCompletion(OpenAIServing):

163 generators: list[AsyncGenerator[RequestOutput, None]] = []

164 try:

165 for i, engine_prompt in enumerate(engine_prompts):

165 - prompt_text, prompt_token_ids, prompt_embeds = (

166 - self._get_prompt_components(engine_prompt)

167 - )

168 -

169 - input_length = None

170 - if prompt_token_ids is not None:

171 - input_length = len(prompt_token_ids)

172 - elif prompt_embeds is not None:

173 - input_length = len(prompt_embeds)

174 - else:

175 - raise NotImplementedError

176 -

177 - if self.default_sampling_params is None:

178 - self.default_sampling_params = {}

166 + prompt_text, _, _ = get_prompt_components(engine_prompt)

167

168 max_tokens = get_max_tokens(

169 max_model_len=self.max_model_len,

170 request=request,

183 - input_length=input_length,

171 + prompt=engine_prompt,

172 default_sampling_params=self.default_sampling_params,

173 )

174

vllm/entrypoints/openai/engine/serving.py

+20 -15

@@ -94,11 +94,14 @@ from vllm.entrypoints.serve.tokenize.protocol import (

94 TokenizeCompletionRequest,

95 TokenizeResponse,

96 )

97 -from vllm.entrypoints.utils import _validate_truncation_size, sanitize_message

97 +from vllm.entrypoints.utils import (

98 + _validate_truncation_size,

99 + get_max_tokens,

100 + sanitize_message,

101 +)

102 from vllm.exceptions import VLLMValidationError

103 from vllm.inputs.data import PromptType, TokensPrompt

104 from vllm.inputs.parse import (

101 - PromptComponents,

105 get_prompt_components,

106 is_explicit_encoder_decoder_prompt,

107 )

@@ -1287,7 +1290,7 @@ class OpenAIServing:

1290 priority: int = 0,

1291 **kwargs,

1292 ):

1290 - prompt_text, _, _ = self._get_prompt_components(engine_prompt)

1293 + prompt_text, _, _ = get_prompt_components(engine_prompt)

1294

1295 orig_priority = priority

1296 sub_request = 0

@@ -1338,10 +1341,12 @@ class OpenAIServing:

1341 # yield context

1342

1343 # Create inputs for the next turn.

1341 - # Render the next prompt token ids.

1344 + # Render the next prompt token ids and update sampling_params.

1345 if isinstance(context, (HarmonyContext, StreamingHarmonyContext)):

1343 - prompt_token_ids = context.render_for_completion()

1344 - engine_prompt = TokensPrompt(prompt_token_ids=prompt_token_ids)

1346 + token_ids = context.render_for_completion()

1347 + engine_prompt = TokensPrompt(prompt_token_ids=token_ids)

1348 +

1349 + sampling_params.max_tokens = self.max_model_len - len(token_ids)

1350 elif isinstance(context, ParsableContext):

1351 engine_prompts = await self._render_next_turn(

1352 context.request,

@@ -1353,19 +1358,19 @@ class OpenAIServing:

1358 context.chat_template_content_format,

1359 )

1360 engine_prompt = engine_prompts[0]

1356 - prompt_text, _, _ = self._get_prompt_components(engine_prompt)

1361 + prompt_text, _, _ = get_prompt_components(engine_prompt)

1362 +

1363 + sampling_params.max_tokens = get_max_tokens(

1364 + self.max_model_len,

1365 + context.request,

1366 + engine_prompt,

1367 + self.default_sampling_params, # type: ignore

1368 + )

1369

1358 - # Update the sampling params.

1359 - sampling_params.max_tokens = self.max_model_len - len(

1360 - engine_prompt["prompt_token_ids"]

1361 - )

1370 # OPTIMIZATION

1371 priority = orig_priority - 1

1372 sub_request += 1

1373

1366 - def _get_prompt_components(self, prompt: PromptType) -> PromptComponents:

1367 - return get_prompt_components(prompt)

1368 -

1374 def _log_inputs(

1375 self,

1376 request_id: str,

@@ -1376,7 +1381,7 @@ class OpenAIServing:

1381 if self.request_logger is None:

1382 return

1383

1379 - prompt, prompt_token_ids, prompt_embeds = self._get_prompt_components(inputs)

1384 + prompt, prompt_token_ids, prompt_embeds = get_prompt_components(inputs)

1385

1386 self.request_logger.log_inputs(

1387 request_id,

vllm/entrypoints/openai/responses/serving.py

+6 -2

@@ -116,6 +116,7 @@ from vllm.entrypoints.openai.responses.utils import (

116 extract_tool_types,

117 should_continue_final_message,

118 )

119 +from vllm.entrypoints.utils import get_max_tokens

120 from vllm.exceptions import VLLMValidationError

121 from vllm.inputs.data import TokensPrompt

122 from vllm.logger import init_logger

@@ -435,8 +436,11 @@ class OpenAIServingResponses(OpenAIServing):

436 if maybe_error is not None:

437 return maybe_error

438

438 - default_max_tokens = self.max_model_len - len(

439 - engine_prompt["prompt_token_ids"]

439 + default_max_tokens = get_max_tokens(

440 + self.max_model_len,

441 + request,

442 + engine_prompt,

443 + self.default_sampling_params,

444 )

445

446 sampling_params = request.to_sampling_params(

vllm/entrypoints/utils.py

+24 -3

@@ -17,8 +17,10 @@ from starlette.background import BackgroundTask, BackgroundTasks

18 from vllm import envs

19 from vllm.engine.arg_utils import EngineArgs

20 +from vllm.inputs import EmbedsPrompt, TokensPrompt

21 from vllm.logger import current_formatter_type, init_logger

22 from vllm.platforms import current_platform

23 +from vllm.utils import length_from_prompt_token_ids_or_embeds

24 from vllm.utils.argparse_utils import FlexibleArgumentParser

26 if TYPE_CHECKING:

@@ -32,11 +34,15 @@ if TYPE_CHECKING:

34 StreamOptions,

35 )

36 from vllm.entrypoints.openai.models.protocol import LoRAModulePath

37 + from vllm.entrypoints.openai.responses.protocol import (

38 + ResponsesRequest,

39 + )

40 else:

41 ChatCompletionRequest = object

42 CompletionRequest = object

43 StreamOptions = object

44 LoRAModulePath = object

45 + ResponsesRequest = object

48 logger = init_logger(__name__)

@@ -211,11 +217,26 @@ def _validate_truncation_size(

217

218 def get_max_tokens(

219 max_model_len: int,

214 - request: "ChatCompletionRequest | CompletionRequest",

215 - input_length: int,

220 + request: "CompletionRequest | ChatCompletionRequest | ResponsesRequest",

221 + prompt: TokensPrompt | EmbedsPrompt,

222 default_sampling_params: dict,

223 ) -> int:

218 - max_tokens = getattr(request, "max_completion_tokens", None) or request.max_tokens

224 + # NOTE: Avoid isinstance() for better efficiency

225 + max_tokens: int | None = None

226 + if max_tokens is None:

227 + # ChatCompletionRequest

228 + max_tokens = getattr(request, "max_completion_tokens", None)

229 + if max_tokens is None:

230 + # ResponsesRequest

231 + max_tokens = getattr(request, "max_output_tokens", None)

232 + if max_tokens is None:

233 + # CompletionRequest (also a fallback for ChatCompletionRequest)

234 + max_tokens = getattr(request, "max_tokens", None)

235 +

236 + input_length = length_from_prompt_token_ids_or_embeds(

237 + prompt.get("prompt_token_ids"), # type: ignore[arg-type]

238 + prompt.get("prompt_embeds"), # type: ignore[arg-type]

239 + )

240 default_max_tokens = max_model_len - input_length

241 max_output_tokens = current_platform.get_max_output_tokens(input_length)

242