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

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 8, 21:30 UTC

[Perf] Optimize the performance of structured output + reasoning (#33557)

Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>

Chauncey committed Feb 5, 2026 at 15:45 UTC 6abb0454adb531de0b081bbf65ccf907e4bd560d

4 files changed +51 -61

vllm/entrypoints/openai/chat_completion/serving.py

+41 -60

@@ -72,6 +72,7 @@ from vllm.logger import init_logger

72 from vllm.logprobs import Logprob

73 from vllm.outputs import CompletionOutput, RequestOutput

74 from vllm.parser import ParserManager

75 +from vllm.reasoning import ReasoningParser

76 from vllm.sampling_params import BeamSearchParams, SamplingParams

77 from vllm.tokenizers import TokenizerLike

78 from vllm.tokenizers.mistral import (

@@ -132,7 +133,7 @@ class OpenAIServingChat(OpenAIServing):

133 self.logits_processors = self.model_config.logits_processors

134

135 # set up reasoning parser

135 - self.reasoning_parser = ParserManager.get_reasoning_parser(

136 + self.reasoning_parser_cls = ParserManager.get_reasoning_parser(

137 reasoning_parser_name=reasoning_parser

138 )

139 # set up tool use

@@ -330,6 +331,24 @@ class OpenAIServingChat(OpenAIServing):

331 for the API specification. This API mimics the OpenAI

332 Chat Completion API.

333 """

334 + # Streaming response

335 + tokenizer = self.renderer.tokenizer

336 + assert tokenizer is not None

337 + reasoning_parser: ReasoningParser | None = None

338 + try:

339 + if self.reasoning_parser_cls:

340 + # Pass the same chat template kwargs as used in tokenization

341 + chat_template_kwargs = self._prepare_extra_chat_template_kwargs(

342 + request.chat_template_kwargs,

343 + self.default_chat_template_kwargs,

344 + )

345 + reasoning_parser = self.reasoning_parser_cls(

346 + tokenizer,

347 + chat_template_kwargs=chat_template_kwargs, # type: ignore[call-arg]

348 + )

349 + except RuntimeError as e:

350 + logger.exception("Error in reasoning parser creation.")

351 + return self.create_error_response(str(e))

352 result = await self.render_chat_request(request)

353 if isinstance(result, ErrorResponse):

354 return result

@@ -427,7 +446,12 @@ class OpenAIServingChat(OpenAIServing):

446 priority=request.priority,

447 data_parallel_rank=data_parallel_rank,

448 )

430 -

449 + reasoning_ended = None

450 + if reasoning_parser:

451 + reasoning_ended = reasoning_parser.is_reasoning_end(

452 + engine_request.prompt_token_ids or [] # type: ignore[attr-defined]

453 + )

454 + engine_request.reasoning_ended = reasoning_ended

455 generator = self.engine_client.generate(

456 engine_request,

457 sampling_params,

@@ -447,10 +471,6 @@ class OpenAIServingChat(OpenAIServing):

471 assert len(generators) == 1

472 (result_generator,) = generators

473

450 - # Streaming response

451 - tokenizer = self.renderer.tokenizer

452 - assert tokenizer is not None

453 -

474 if request.stream:

475 return self.chat_completion_stream_generator(

476 request,

@@ -460,6 +480,7 @@ class OpenAIServingChat(OpenAIServing):

480 conversation,

481 tokenizer,

482 request_metadata,

483 + reasoning_parser,

484 )

485

486 try:

@@ -471,6 +492,7 @@ class OpenAIServingChat(OpenAIServing):

492 conversation,

493 tokenizer,

494 request_metadata,

495 + reasoning_parser,

496 )

497 except GenerationError as e:

498 return self._convert_generation_error_to_response(e)

@@ -630,6 +652,7 @@ class OpenAIServingChat(OpenAIServing):

652 conversation: list[ConversationMessage],

653 tokenizer: TokenizerLike,

654 request_metadata: RequestResponseMetadata,

655 + reasoning_parser: ReasoningParser | None = None,

656 ) -> AsyncGenerator[str, None]:

657 from vllm.tokenizers.mistral import MistralTokenizer

658

@@ -673,7 +696,7 @@ class OpenAIServingChat(OpenAIServing):

696

697 # Only one of these will be used, thus previous_texts and

698 # all_previous_token_ids will not be used twice in the same iteration.

676 - if tool_choice_auto or self.reasoning_parser:

699 + if tool_choice_auto or reasoning_parser:

700 # These are only required in "auto" tool choice case

701 all_previous_token_ids = [[]] * num_choices

702 # For reasoning parser and tool call all enabled

@@ -683,28 +706,6 @@ class OpenAIServingChat(OpenAIServing):

706 else:

707 all_previous_token_ids = None

708

686 - try:

687 - if self.reasoning_parser:

688 - if tokenizer is None:

689 - raise ValueError(

690 - "Tokenizer not available when `skip_tokenizer_init=True`"

691 - )

692 -

693 - # Pass the same chat template kwargs as used in tokenization

694 - chat_template_kwargs = self._prepare_extra_chat_template_kwargs(

695 - request.chat_template_kwargs,

696 - self.default_chat_template_kwargs,

697 - )

698 - reasoning_parser = self.reasoning_parser(

699 - tokenizer,

700 - chat_template_kwargs=chat_template_kwargs or {}, # type: ignore[call-arg]

701 - )

702 - except RuntimeError as e:

703 - logger.exception("Error in reasoning parser creation.")

704 - data = self.create_streaming_error_response(str(e))

705 - yield f"data: {data}\n\n"

706 - yield "data: [DONE]\n\n"

707 - return

709 # Prepare the tool parser if it's needed

710 try:

711 if tool_choice_auto and self.tool_parser:

@@ -826,7 +827,7 @@ class OpenAIServingChat(OpenAIServing):

827 tool_parser = tool_parsers[i]

828

829 if (

829 - self.reasoning_parser

830 + reasoning_parser

831 and res.prompt_token_ids

832 and prompt_is_reasoning_end_arr[i] is None

833 ):

@@ -888,7 +889,7 @@ class OpenAIServingChat(OpenAIServing):

889 delta_message: DeltaMessage | None

890

891 # just update previous_texts and previous_token_ids

891 - if tool_choice_auto or self.reasoning_parser:

892 + if tool_choice_auto or reasoning_parser:

893 assert previous_texts is not None

894 assert all_previous_token_ids is not None

895 previous_text = previous_texts[i]

@@ -915,7 +916,7 @@ class OpenAIServingChat(OpenAIServing):

916 # handle streaming deltas for tools with named tool_choice

917 elif tool_choice_function_name:

918 if (

918 - self.reasoning_parser

919 + reasoning_parser

920 and not reasoning_end_arr[i]

921 and not reasoning_parser.is_reasoning_end(

922 previous_token_ids

@@ -952,7 +953,7 @@ class OpenAIServingChat(OpenAIServing):

953 current_text = ""

954 else:

955 # Just to add remaining `content`

955 - if self.reasoning_parser:

956 + if reasoning_parser:

957 delta_text = previous_text + delta_text

958 current_text = ""

959

@@ -998,13 +999,13 @@ class OpenAIServingChat(OpenAIServing):

999 output_token_ids = as_list(output.token_ids)

1000

1001 if (

1001 - self.reasoning_parser is not None

1002 + reasoning_parser is not None

1003 and not reasoning_end_arr[i]

1004 and prompt_is_reasoning_end_arr[i]

1005 ):

1006 reasoning_end_arr[i] = True

1007

1007 - if self.reasoning_parser and not reasoning_end_arr[i]:

1008 + if reasoning_parser and not reasoning_end_arr[i]:

1009 delta_message = (

1010 reasoning_parser.extract_reasoning_streaming(

1011 previous_text,

@@ -1047,9 +1048,8 @@ class OpenAIServingChat(OpenAIServing):

1048

1049 # handle streaming deltas for tools with "auto" tool choice

1050 # and reasoning parser

1050 - elif tool_choice_auto and self.reasoning_parser:

1051 + elif tool_choice_auto and reasoning_parser:

1052 assert tool_parser is not None

1052 - assert reasoning_parser is not None

1053 assert added_content_delta_arr is not None

1054 assert reasoning_end_arr is not None

1055 output_token_ids = as_list(output.token_ids)

@@ -1130,7 +1130,7 @@ class OpenAIServingChat(OpenAIServing):

1130 tools_streamed[i] = True

1131

1132 # when only reasoning

1133 - elif self.reasoning_parser:

1133 + elif reasoning_parser:

1134 delta_message = reasoning_parser.extract_reasoning_streaming(

1135 previous_text,

1136 current_text,

@@ -1144,9 +1144,7 @@ class OpenAIServingChat(OpenAIServing):

1144 delta_message = DeltaMessage(content=delta_text)

1145

1146 # update the previous values for the next iteration

1147 - if (

1148 - tool_choice_auto or self.reasoning_parser

1149 - ) and not self.use_harmony:

1147 + if (tool_choice_auto or reasoning_parser) and not self.use_harmony:

1148 assert previous_texts is not None

1149 assert all_previous_token_ids is not None

1150 previous_texts[i] = current_text

@@ -1400,6 +1398,7 @@ class OpenAIServingChat(OpenAIServing):

1398 conversation: list[ConversationMessage],

1399 tokenizer: TokenizerLike,

1400 request_metadata: RequestResponseMetadata,

1401 + reasoning_parser: ReasoningParser | None = None,

1402 ) -> ErrorResponse | ChatCompletionResponse:

1403 from vllm.tokenizers.mistral import MistralTokenizer

1404

@@ -1494,25 +1493,7 @@ class OpenAIServingChat(OpenAIServing):

1493 choices.append(choice_data)

1494 continue

1495

1497 - if self.reasoning_parser:

1498 - try:

1499 - if tokenizer is None:

1500 - raise ValueError(

1501 - "Tokenizer not available when `skip_tokenizer_init=True`"

1502 - )

1503 -

1504 - # Pass the same chat template kwargs as used in tokenization

1505 - chat_template_kwargs = self._prepare_extra_chat_template_kwargs(

1506 - request.chat_template_kwargs,

1507 - self.default_chat_template_kwargs,

1508 - )

1509 - reasoning_parser = self.reasoning_parser(

1510 - tokenizer,

1511 - chat_template_kwargs=chat_template_kwargs, # type: ignore[call-arg]

1512 - )

1513 - except RuntimeError as e:

1514 - logger.exception("Error in reasoning parser creation.")

1515 - return self.create_error_response(str(e))

1496 + if reasoning_parser:

1497 # If the reasoning parser is enabled,

1498 # tool calls are extracted exclusively from the content.

1499 reasoning, content = reasoning_parser.extract_reasoning(

vllm/v1/engine/__init__.py

@@ -83,6 +83,8 @@ class EngineCoreRequest(

83 # Used in outputs and to support abort(req_id, internal=False).

84 external_req_id: str | None = None

86 + reasoning_ended: bool | None = None

87 +

88 @property

89 def params(self) -> SamplingParams | PoolingParams:

90 """Return the processed params (sampling or pooling)."""

vllm/v1/request.py

@@ -74,6 +74,7 @@ class Request:

74 trace_headers: Mapping[str, str] | None = None,

75 block_hasher: Callable[["Request"], list["BlockHash"]] | None = None,

76 resumable: bool = False,

77 + reasoning_ended: bool | None = None,

78 ) -> None:

79 self.request_id = request_id

80 self.client_index = client_index

@@ -86,6 +87,8 @@ class Request:

87 self.structured_output_request = StructuredOutputRequest.from_sampling_params(

88 sampling_params

89 )

90 + if self.structured_output_request is not None:

91 + self.structured_output_request.reasoning_ended = reasoning_ended

92 self.arrival_time = arrival_time if arrival_time is not None else time.time()

94 self.status = RequestStatus.WAITING

@@ -195,6 +198,7 @@ class Request:

198 trace_headers=request.trace_headers,

199 block_hasher=block_hasher,

200 resumable=request.resumable,

201 + reasoning_ended=request.reasoning_ended,

202 )

203

204 def append_output_token_ids(

vllm/v1/structured_output/__init__.py

+4 -1

@@ -284,12 +284,15 @@ class StructuredOutputManager:

284 # NOTE (Hanchen) if enable_in_reasoning is True, it means that

285 # the model needs to be constrained in reasoning. So we should always

286 # enable the bitmask filling.

287 -

287 if self.reasoner is not None:

288 if self.enable_in_reasoning:

289 return True

290 assert request.structured_output_request is not None

291 if request.structured_output_request.reasoning_ended is None:

292 + # This should be removed here, but since `openai_gptoss`

293 + # is an independent code path, it is kept for now.

294 + # After unifying the `openai_gptoss` and non-`openai_gptoss` styles,

295 + # it can be removed.

296 request.structured_output_request.reasoning_ended = (

297 self.reasoner.is_reasoning_end(request.prompt_token_ids or [])

298 )