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, 01:00 UTC
[Frontend] Support cache_salt in /v1/completions and /v1/responses (#20981)
Signed-off-by: Marko Rosenmueller <5467316+dr75@users.noreply.github.com>
Marko Rosenmueller committed
Jul 15, 2025 at 23:01 UTC
19c863068b2d70a452bde25318dbcf04f274ad19
4 files changed
+77
-4
vllm/entrypoints/openai/api_server.py
+1
@@ -1540,6 +1540,7 @@ async def init_app_state(
1540
state.openai_serving_models,
1541
request_logger=request_logger,
1542
return_tokens_as_token_ids=args.return_tokens_as_token_ids,
1543
+ enable_prompt_tokens_details=args.enable_prompt_tokens_details,
1544
enable_force_include_usage=args.enable_force_include_usage,
1545
) if "generate" in model_config.supported_tasks else None
1546
state.openai_serving_pooling = OpenAIServingPooling(
vllm/entrypoints/openai/protocol.py
+49
-3
@@ -290,6 +290,15 @@ class ResponsesRequest(OpenAIBaseModel):
290
"default: 0). Any priority other than 0 will raise an error "
291
"if the served model does not use priority scheduling."),
292
)
293
+ cache_salt: Optional[str] = Field(
294
+ default=None,
295
+ description=(
296
+ "If specified, the prefix cache will be salted with the provided "
297
+ "string to prevent an attacker to guess prompts in multi-user "
298
+ "environments. The salt should be random, protected from "
299
+ "access by 3rd parties, and long enough to be "
300
+ "unpredictable (e.g., 43 characters base64-encoded, corresponding "
301
+ "to 256 bit). Not supported by vLLM engine V0."))
302
# --8<-- [end:responses-extra-params]
303
304
_DEFAULT_SAMPLING_PARAMS = {
@@ -351,6 +360,19 @@ class ResponsesRequest(OpenAIBaseModel):
360
raise ValueError("prompt template is not supported")
361
return data
362
363
+ @model_validator(mode="before")
364
+ def check_cache_salt_support(cls, data):
365
+ if data.get("cache_salt") is not None:
366
+ if not envs.VLLM_USE_V1:
367
+ raise ValueError(
368
+ "Parameter 'cache_salt' is not supported with "
369
+ "this instance of vLLM, which uses engine V0.")
370
+ if not isinstance(data["cache_salt"],
371
+ str) or not data["cache_salt"]:
372
+ raise ValueError("Parameter 'cache_salt' must be a "
373
+ "non-empty string if provided.")
374
+ return data
375
+
376
377
class ChatCompletionRequest(OpenAIBaseModel):
378
# Ordered by official OpenAI API documentation
@@ -1004,6 +1026,16 @@ class CompletionRequest(OpenAIBaseModel):
1026
" as strings of the form 'token_id:{token_id}' so that tokens "
1027
"that are not JSON-encodable can be identified."))
1028
1029
+ cache_salt: Optional[str] = Field(
1030
+ default=None,
1031
+ description=(
1032
+ "If specified, the prefix cache will be salted with the provided "
1033
+ "string to prevent an attacker to guess prompts in multi-user "
1034
+ "environments. The salt should be random, protected from "
1035
+ "access by 3rd parties, and long enough to be "
1036
+ "unpredictable (e.g., 43 characters base64-encoded, corresponding "
1037
+ "to 256 bit). Not supported by vLLM engine V0."))
1038
+
1039
kv_transfer_params: Optional[dict[str, Any]] = Field(
1040
default=None,
1041
description="KVTransfer parameters used for disaggregated serving.")
@@ -1180,6 +1212,20 @@ class CompletionRequest(OpenAIBaseModel):
1212
"At least one of `prompt` or `prompt_embeds` must be set.")
1213
return data
1214
1215
+ @model_validator(mode="before")
1216
+ @classmethod
1217
+ def check_cache_salt_support(cls, data):
1218
+ if data.get("cache_salt") is not None:
1219
+ if not envs.VLLM_USE_V1:
1220
+ raise ValueError(
1221
+ "Parameter 'cache_salt' is not supported with "
1222
+ "this instance of vLLM, which uses engine V0.")
1223
+ if not isinstance(data["cache_salt"],
1224
+ str) or not data["cache_salt"]:
1225
+ raise ValueError("Parameter 'cache_salt' must be a "
1226
+ "non-empty string if provided.")
1227
+ return data
1228
+
1229
1230
class EmbeddingCompletionRequest(OpenAIBaseModel):
1231
# Ordered by official OpenAI API documentation
@@ -1971,7 +2017,7 @@ class TranscriptionRequest(OpenAIBaseModel):
2017
"""
2018
2019
stream: Optional[bool] = False
1974
- """When set, it will enable output to be streamed in a similar fashion
2020
+ """When set, it will enable output to be streamed in a similar fashion
2021
as the Chat Completion endpoint.
2022
"""
2023
# --8<-- [start:transcription-extra-params]
@@ -2233,9 +2279,9 @@ class TranslationRequest(OpenAIBaseModel):
2279
"""
2280
2281
stream: Optional[bool] = False
2236
- """Custom field not present in the original OpenAI definition. When set,
2282
+ """Custom field not present in the original OpenAI definition. When set,
2283
it will enable output to be streamed in a similar fashion as the Chat
2238
- Completion endpoint.
2284
+ Completion endpoint.
2285
"""
2286
# Flattened stream option to simplify form data.
2287
stream_include_usage: Optional[bool] = False
vllm/entrypoints/openai/serving_completion.py
+17
@@ -23,6 +23,7 @@ from vllm.entrypoints.openai.protocol import (CompletionLogProbs,
23
CompletionResponseStreamChoice,
24
CompletionStreamResponse,
25
ErrorResponse,
26
+ PromptTokenUsageInfo,
27
RequestResponseMetadata,
28
UsageInfo)
29
from vllm.entrypoints.openai.serving_engine import (
@@ -56,6 +57,7 @@ class OpenAIServingCompletion(OpenAIServing):
57
*,
58
request_logger: Optional[RequestLogger],
59
return_tokens_as_token_ids: bool = False,
60
+ enable_prompt_tokens_details: bool = False,
61
enable_force_include_usage: bool = False,
62
):
63
super().__init__(engine_client=engine_client,
@@ -64,6 +66,7 @@ class OpenAIServingCompletion(OpenAIServing):
66
request_logger=request_logger,
67
return_tokens_as_token_ids=return_tokens_as_token_ids,
68
enable_force_include_usage=enable_force_include_usage)
69
+ self.enable_prompt_tokens_details = enable_prompt_tokens_details
70
self.default_sampling_params = (
71
self.model_config.get_diff_sampling_param())
72
if self.default_sampling_params:
@@ -313,6 +316,8 @@ class OpenAIServingCompletion(OpenAIServing):
316
previous_num_tokens = [0] * num_choices * num_prompts
317
has_echoed = [False] * num_choices * num_prompts
318
num_prompt_tokens = [0] * num_prompts
319
+ num_cached_tokens = None
320
+ first_iteration = True
321
322
stream_options = request.stream_options
323
if stream_options:
@@ -328,6 +333,10 @@ class OpenAIServingCompletion(OpenAIServing):
333
prompt_token_ids = res.prompt_token_ids
334
prompt_logprobs = res.prompt_logprobs
335
336
+ if first_iteration:
337
+ num_cached_tokens = res.num_cached_tokens
338
+ first_iteration = False
339
+
340
if res.prompt is not None:
341
prompt_text = res.prompt
342
else:
@@ -431,6 +440,10 @@ class OpenAIServingCompletion(OpenAIServing):
440
completion_tokens=total_completion_tokens,
441
total_tokens=total_prompt_tokens + total_completion_tokens)
442
443
+ if self.enable_prompt_tokens_details and num_cached_tokens:
444
+ final_usage_info.prompt_tokens_details = PromptTokenUsageInfo(
445
+ cached_tokens=num_cached_tokens)
446
+
447
if include_usage:
448
final_usage_chunk = CompletionStreamResponse(
449
id=request_id,
@@ -535,6 +548,10 @@ class OpenAIServingCompletion(OpenAIServing):
548
total_tokens=num_prompt_tokens + num_generated_tokens,
549
)
550
551
+ if self.enable_prompt_tokens_details and final_res.num_cached_tokens:
552
+ usage.prompt_tokens_details = PromptTokenUsageInfo(
553
+ cached_tokens=final_res.num_cached_tokens)
554
+
555
request_metadata.final_usage_info = usage
556
557
return CompletionResponse(
vllm/entrypoints/openai/serving_engine.py
+10
-1
@@ -226,7 +226,7 @@ class OpenAIServing:
226
227
def _get_async_tokenizer(self, tokenizer) -> AsyncMicrobatchTokenizer:
228
"""
229
- Return (and cache) an `AsyncMicrobatchTokenizer` bound to the
229
+ Return (and cache) an `AsyncMicrobatchTokenizer` bound to the
230
given tokenizer.
231
"""
232
async_tokenizer = self._async_tokenizer_pool.get(tokenizer)
@@ -811,6 +811,12 @@ class OpenAIServing:
811
prompt_token_ids=request_prompt_text["prompt_token_ids"])
812
for request_prompt_text in request_prompts_text
813
]
814
+ cache_salt = request.cache_salt if (
815
+ hasattr(request, "cache_salt")
816
+ and request.cache_salt is not None) else None
817
+ if cache_salt:
818
+ for prompt_text in engine_prompts_text:
819
+ prompt_text["cache_salt"] = cache_salt
820
821
# This check is equivalent to simply checking if
822
# `request_prompts_embeds` is empty, but it's difficult to propagate
@@ -828,6 +834,9 @@ class OpenAIServing:
834
prompt_embeds=request_prompt_embeds["prompt_embeds"])
835
for request_prompt_embeds in request_prompts_embeds
836
]
837
+ if cache_salt:
838
+ for prompt_embed in engine_prompts_embeds:
839
+ prompt_embed["cache_salt"] = cache_salt
840
841
request_prompts = request_prompts_embeds + request_prompts_text
842
engine_prompts = engine_prompts_embeds + engine_prompts_text