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

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

[Feat][CLI] enforce-include-usage (#19695)

Signed-off-by: Max Wittig <max.wittig@siemens.com>

Max Wittig committed Jun 25, 2025 at 07:43 UTC f59fc60fb317e7e04456de50b7abce99a9017225

5 files changed +34 -9

vllm/entrypoints/openai/api_server.py

@@ -1190,6 +1190,7 @@ async def init_app_state(

1190 tool_parser=args.tool_call_parser,

1191 reasoning_parser=args.reasoning_parser,

1192 enable_prompt_tokens_details=args.enable_prompt_tokens_details,

1193 + enable_force_include_usage=args.enable_force_include_usage,

1194 ) if model_config.runner_type == "generate" else None

1195 state.openai_serving_completion = OpenAIServingCompletion(

1196 engine_client,

@@ -1197,6 +1198,7 @@ async def init_app_state(

1198 state.openai_serving_models,

1199 request_logger=request_logger,

1200 return_tokens_as_token_ids=args.return_tokens_as_token_ids,

1201 + enable_force_include_usage=args.enable_force_include_usage,

1202 ) if model_config.runner_type == "generate" else None

1203 state.openai_serving_pooling = OpenAIServingPooling(

1204 engine_client,

vllm/entrypoints/openai/cli_args.py

@@ -272,6 +272,11 @@ def make_arg_parser(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:

272 action='store_true',

273 default=False,

274 help="If set to True, enable prompt_tokens_details in usage.")

275 + parser.add_argument(

276 + "--enable-force-include-usage",

277 + action='store_true',

278 + default=False,

279 + help="If set to True, including usage on every request.")

280 parser.add_argument(

281 "--enable-server-load-tracking",

282 action='store_true',

vllm/entrypoints/openai/serving_chat.py

+15 -4

@@ -64,12 +64,14 @@ class OpenAIServingChat(OpenAIServing):

64 enable_auto_tools: bool = False,

65 tool_parser: Optional[str] = None,

66 enable_prompt_tokens_details: bool = False,

67 + enable_force_include_usage: bool = False,

68 ) -> None:

69 super().__init__(engine_client=engine_client,

70 model_config=model_config,

71 models=models,

72 request_logger=request_logger,

72 - return_tokens_as_token_ids=return_tokens_as_token_ids)

73 + return_tokens_as_token_ids=return_tokens_as_token_ids,

74 + enable_force_include_usage=enable_force_include_usage)

76 self.response_role = response_role

77 self.chat_template = chat_template

@@ -110,6 +112,7 @@ class OpenAIServingChat(OpenAIServing):

112 "been registered") from e

113

114 self.enable_prompt_tokens_details = enable_prompt_tokens_details

115 + self.enable_force_include_usage = enable_force_include_usage

116 self.default_sampling_params = (

117 self.model_config.get_diff_sampling_param())

118 if self.default_sampling_params:

@@ -261,8 +264,14 @@ class OpenAIServingChat(OpenAIServing):

264 # Streaming response

265 if request.stream:

266 return self.chat_completion_stream_generator(

264 - request, result_generator, request_id, model_name,

265 - conversation, tokenizer, request_metadata)

267 + request,

268 + result_generator,

269 + request_id,

270 + model_name,

271 + conversation,

272 + tokenizer,

273 + request_metadata,

274 + enable_force_include_usage=self.enable_force_include_usage)

275

276 try:

277 return await self.chat_completion_full_generator(

@@ -405,6 +414,7 @@ class OpenAIServingChat(OpenAIServing):

414 conversation: list[ConversationMessage],

415 tokenizer: AnyTokenizer,

416 request_metadata: RequestResponseMetadata,

417 + enable_force_include_usage: bool,

418 ) -> AsyncGenerator[str, None]:

419 created_time = int(time.time())

420 chunk_object_type: Final = "chat.completion.chunk"

@@ -471,7 +481,8 @@ class OpenAIServingChat(OpenAIServing):

481

482 stream_options = request.stream_options

483 if stream_options:

474 - include_usage = stream_options.include_usage

484 + include_usage = stream_options.include_usage \

485 + or enable_force_include_usage

486 include_continuous_usage = include_usage and \

487 stream_options.continuous_usage_stats

488 else:

vllm/entrypoints/openai/serving_completion.py

+8 -3

@@ -52,12 +52,14 @@ class OpenAIServingCompletion(OpenAIServing):

52 *,

53 request_logger: Optional[RequestLogger],

54 return_tokens_as_token_ids: bool = False,

55 + enable_force_include_usage: bool = False,

56 ):

57 super().__init__(engine_client=engine_client,

58 model_config=model_config,

59 models=models,

60 request_logger=request_logger,

60 - return_tokens_as_token_ids=return_tokens_as_token_ids)

61 + return_tokens_as_token_ids=return_tokens_as_token_ids,

62 + enable_force_include_usage=enable_force_include_usage)

63 self.default_sampling_params = (

64 self.model_config.get_diff_sampling_param())

65 if self.default_sampling_params:

@@ -227,7 +229,8 @@ class OpenAIServingCompletion(OpenAIServing):

229 model_name,

230 num_prompts=num_prompts,

231 tokenizer=tokenizer,

230 - request_metadata=request_metadata)

232 + request_metadata=request_metadata,

233 + enable_force_include_usage=self.enable_force_include_usage)

234

235 # Non-streaming response

236 final_res_batch: list[Optional[RequestOutput]] = [None] * num_prompts

@@ -289,6 +292,7 @@ class OpenAIServingCompletion(OpenAIServing):

292 num_prompts: int,

293 tokenizer: AnyTokenizer,

294 request_metadata: RequestResponseMetadata,

295 + enable_force_include_usage: bool,

296 ) -> AsyncGenerator[str, None]:

297 num_choices = 1 if request.n is None else request.n

298 previous_text_lens = [0] * num_choices * num_prompts

@@ -298,7 +302,8 @@ class OpenAIServingCompletion(OpenAIServing):

302

303 stream_options = request.stream_options

304 if stream_options:

301 - include_usage = stream_options.include_usage

305 + include_usage = stream_options.include_usage or \

306 + enable_force_include_usage

307 include_continuous_usage = include_usage and \

308 stream_options.continuous_usage_stats

309 else:

vllm/entrypoints/openai/serving_engine.py

+4 -2

@@ -132,7 +132,7 @@ RequestT = TypeVar("RequestT", bound=AnyRequest)

132

133 class RequestProcessingMixin(BaseModel):

134 """

135 - Mixin for request processing,

135 + Mixin for request processing,

136 handling prompt preparation and engine input.

137 """

138 request_prompts: Optional[Sequence[RequestPrompt]] = []

@@ -144,7 +144,7 @@ class RequestProcessingMixin(BaseModel):

144

145 class ResponseGenerationMixin(BaseModel):

146 """

147 - Mixin for response generation,

147 + Mixin for response generation,

148 managing result generators and final batch results.

149 """

150 result_generator: Optional[AsyncGenerator[tuple[int, Union[

@@ -208,6 +208,7 @@ class OpenAIServing:

208 *,

209 request_logger: Optional[RequestLogger],

210 return_tokens_as_token_ids: bool = False,

211 + enable_force_include_usage: bool = False,

212 ):

213 super().__init__()

214

@@ -219,6 +220,7 @@ class OpenAIServing:

220

221 self.request_logger = request_logger

222 self.return_tokens_as_token_ids = return_tokens_as_token_ids

223 + self.enable_force_include_usage = enable_force_include_usage

224

225 self._tokenizer_executor = ThreadPoolExecutor(max_workers=1)

226