@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1c2bc7ead0
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, 12:15 UTC

Truncation control for embedding models (#14776)

Signed-off-by: Gabriel Marinho <gmarinho@ibm.com> Signed-off-by: Max de Bayser <mbayser@br.ibm.com> Co-authored-by: Max de Bayser <mbayser@br.ibm.com>

Gabriel Marinho committed Apr 29, 2025 at 22:24 UTC 1c2bc7ead019cdf5b04b2f1d07b00982352f85ef
21 files changed +332 -70
tests/entrypoints/openai/test_truncation.py new
+103
@@ -0,0 +1,103 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from typing import Any
3 +
4 +import openai
5 +import pytest
6 +import pytest_asyncio
7 +
8 +from tests.utils import RemoteOpenAIServer
9 +
10 +MODEL_NAME = "sentence-transformers/all-MiniLM-L12-v2"
11 +max_model_len = 128
12 +
13 +input = """Immerse yourself in the enchanting chronicle of calculus, a
14 + mathematical domain that has radically transformed our comprehension of
15 + change and motion. Despite its roots in ancient civilizations, the
16 + formal birth of calculus predominantly occurred in the 17th century,
17 + primarily under the influential guidance of Sir Isaac Newton and Gottfried
18 + Wilhelm Leibniz. The earliest traces of calculus concepts are found in
19 + ancient Greek mathematics,most notably in the works of Eudoxus and
20 + Archimedes, around 300 BCE. They utilized the 'method of exhaustion'—a
21 + technique for computing areas and volumes through the use of finite sums.
22 + This methodology laid crucial foundational work for integral calculus.
23 + In the 17th century, both Newton and Leibniz independently pioneered
24 + calculus, each contributing unique perspectives that would shape this new
25 + field."""
26 +
27 +
28 +@pytest.fixture(scope="module")
29 +def server():
30 + args = [
31 + "--task",
32 + "embed",
33 + "--dtype",
34 + "bfloat16",
35 + "--enforce-eager",
36 + "--max-model-len",
37 + str(max_model_len),
38 + ]
39 +
40 + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
41 + yield remote_server
42 +
43 +
44 +@pytest_asyncio.fixture
45 +async def client(server):
46 + async with server.get_async_client() as async_client:
47 + yield async_client
48 +
49 +
50 +@pytest.mark.asyncio
51 +async def test_smaller_truncation_size(client: openai.AsyncOpenAI):
52 + truncation_size = 10
53 + kwargs: dict[str, Any] = {
54 + "model": MODEL_NAME,
55 + "input": input,
56 + "truncate_prompt_tokens": truncation_size
57 + }
58 +
59 + response = await client.post(path="embeddings",
60 + cast_to=object,
61 + body={**kwargs})
62 +
63 + assert response["usage"]["prompt_tokens"] == truncation_size
64 +
65 +
66 +@pytest.mark.asyncio
67 +async def test_bigger_truncation_size(client: openai.AsyncOpenAI):
68 + truncation_size = max_model_len + 1
69 + kwargs: dict[str, Any] = {
70 + "model": MODEL_NAME,
71 + "input": input,
72 + "truncate_prompt_tokens": truncation_size
73 + }
74 +
75 + with pytest.raises(openai.BadRequestError) as err:
76 + err = await client.post(path="embeddings",
77 + cast_to=object,
78 + body={**kwargs})
79 +
80 + assert str(err) == f"""openai.BadRequestError:
81 + Error code: 400 - {{'object': 'error',
82 + 'message': 'truncate_prompt_tokens value
83 + ({truncation_size})
84 + is greater than max_model_len ({max_model_len}).
85 + Please, select a smaller truncation size.',
86 + 'type': 'BadRequestError',
87 + 'param': None, 'code': 400}}"""
88 +
89 +
90 +@pytest.mark.asyncio
91 +async def test_max_truncation_size(client: openai.AsyncOpenAI):
92 + truncation_size = -1
93 + kwargs: dict[str, Any] = {
94 + "model": MODEL_NAME,
95 + "input": input,
96 + "truncate_prompt_tokens": truncation_size
97 + }
98 +
99 + response = await client.post(path="embeddings",
100 + cast_to=object,
101 + body={**kwargs})
102 +
103 + assert response["usage"]["prompt_tokens"] == max_model_len
tests/models/embedding/language/test_truncation_control.py new
+69
@@ -0,0 +1,69 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import pytest
3 +
4 +MODEL_NAME = "sentence-transformers/all-MiniLM-L12-v2"
5 +max_model_len = 128
6 +
7 +input_str = """Immerse yourself in the enchanting chronicle of calculus, a
8 + mathematical domain that has radically transformed our comprehension of
9 + change and motion. Despite its roots in ancient civilizations, the
10 + formal birth of calculus predominantly occurred in the 17th century,
11 + primarily under the influential guidance of Sir Isaac Newton and Gottfried
12 + Wilhelm Leibniz. The earliest traces of calculus concepts are found in
13 + ancient Greek mathematics,most notably in the works of Eudoxus and
14 + Archimedes, around 300 BCE. They utilized the 'method of exhaustion'—a
15 + technique for computing areas and volumes through the use of finite sums.
16 + This methodology laid crucial foundational work for integral calculus.
17 + In the 17th century, both Newton and Leibniz independently pioneered
18 + calculus, each contributing unique perspectives that would shape this new
19 + field."""
20 +
21 +
22 +def test_smaller_truncation_size(vllm_runner,
23 + model_name=MODEL_NAME,
24 + input_str=input_str):
25 +
26 + truncate_prompt_tokens = 10
27 +
28 + with vllm_runner(model_name, task="embed",
29 + max_model_len=max_model_len) as vllm_model:
30 + vllm_output = vllm_model.model.encode(
31 + input_str, truncate_prompt_tokens=truncate_prompt_tokens)
32 +
33 + prompt_tokens = vllm_output[0].prompt_token_ids
34 +
35 + assert len(prompt_tokens) == truncate_prompt_tokens
36 +
37 +
38 +def test_max_truncation_size(vllm_runner,
39 + model_name=MODEL_NAME,
40 + input_str=input_str):
41 + truncate_prompt_tokens = -1
42 +
43 + with vllm_runner(model_name, task="embed",
44 + max_model_len=max_model_len) as vllm_model:
45 + vllm_output = vllm_model.model.encode(
46 + input_str, truncate_prompt_tokens=truncate_prompt_tokens)
47 +
48 + prompt_tokens = vllm_output[0].prompt_token_ids
49 +
50 + assert len(prompt_tokens) == max_model_len
51 +
52 +
53 +def test_bigger_truncation_size(vllm_runner,
54 + model_name=MODEL_NAME,
55 + input_str=input_str):
56 +
57 + truncate_prompt_tokens = max_model_len + 1
58 +
59 + with pytest.raises(ValueError), vllm_runner(
60 + model_name, task="embed",
61 + max_model_len=max_model_len) as vllm_model:
62 +
63 + llm_output = vllm_model.model.encode(
64 + input_str, truncate_prompt_tokens=truncate_prompt_tokens)
65 +
66 + assert llm_output == f"""truncate_prompt_tokens value
67 + ({truncate_prompt_tokens}) is greater than
68 + max_model_len ({max_model_len}). Please, select
69 + a smaller truncation size."""
vllm/engine/llm_engine.py
+3
@@ -645,6 +645,7 @@ class LLMEngine:
645 params: Union[SamplingParams, PoolingParams],
646 arrival_time: Optional[float] = None,
647 lora_request: Optional[LoRARequest] = None,
648 + tokenization_kwargs: Optional[dict[str, Any]] = None,
649 trace_headers: Optional[Mapping[str, str]] = None,
650 prompt_adapter_request: Optional[PromptAdapterRequest] = None,
651 priority: int = 0,
@@ -678,6 +679,7 @@ class LLMEngine:
679 params: Optional[Union[SamplingParams, PoolingParams]] = None,
680 arrival_time: Optional[float] = None,
681 lora_request: Optional[LoRARequest] = None,
682 + tokenization_kwargs: Optional[dict[str, Any]] = None,
683 trace_headers: Optional[Mapping[str, str]] = None,
684 prompt_adapter_request: Optional[PromptAdapterRequest] = None,
685 priority: int = 0,
@@ -758,6 +760,7 @@ class LLMEngine:
760
761 processed_inputs = self.input_preprocessor.preprocess(
762 prompt,
763 + tokenization_kwargs=tokenization_kwargs,
764 lora_request=lora_request,
765 prompt_adapter_request=prompt_adapter_request,
766 )
vllm/engine/protocol.py
+2 -2
@@ -2,7 +2,7 @@
2
3 import asyncio
4 from abc import ABC, abstractmethod
5 -from typing import AsyncGenerator, List, Mapping, Optional
5 +from typing import AsyncGenerator, Mapping, Optional
6
7 from vllm.beam_search import BeamSearchSequence, create_sort_beams_key_function
8 from vllm.config import DecodingConfig, ModelConfig, VllmConfig
@@ -256,7 +256,7 @@ class EngineClient(ABC):
256 async def do_log_stats(
257 self,
258 scheduler_outputs: Optional[SchedulerOutputs] = None,
259 - model_output: Optional[List[SamplerOutput]] = None,
259 + model_output: Optional[list[SamplerOutput]] = None,
260 ) -> None:
261 ...
262
vllm/entrypoints/llm.py
+22 -3
@@ -25,6 +25,7 @@ from vllm.entrypoints.chat_utils import (ChatCompletionMessageParam,
25 resolve_chat_template_content_format)
26 from vllm.entrypoints.score_utils import (_cosine_similarity,
27 _validate_score_input_lens)
28 +from vllm.entrypoints.utils import _validate_truncation_size
29 from vllm.inputs import PromptType, SingletonPrompt, TextPrompt, TokensPrompt
30 from vllm.inputs.parse import is_token_prompt, parse_and_batch_prompt
31 from vllm.logger import init_logger
@@ -793,6 +794,7 @@ class LLM:
794 pooling_params: Optional[Union[PoolingParams,
795 Sequence[PoolingParams]]] = None,
796 *,
797 + truncate_prompt_tokens: Optional[int] = None,
798 use_tqdm: bool = True,
799 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
800 prompt_adapter_request: Optional[PromptAdapterRequest] = None,
@@ -807,6 +809,7 @@ class LLM:
809 pooling_params: Optional[Union[PoolingParams,
810 Sequence[PoolingParams]]] = None,
811 prompt_token_ids: Optional[list[int]] = None,
812 + truncate_prompt_tokens: Optional[int] = None,
813 use_tqdm: bool = True,
814 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
815 prompt_adapter_request: Optional[PromptAdapterRequest] = None,
@@ -821,6 +824,7 @@ class LLM:
824 pooling_params: Optional[Union[PoolingParams,
825 Sequence[PoolingParams]]] = None,
826 prompt_token_ids: Optional[list[list[int]]] = None,
827 + truncate_prompt_tokens: Optional[int] = None,
828 use_tqdm: bool = True,
829 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
830 prompt_adapter_request: Optional[PromptAdapterRequest] = None,
@@ -836,6 +840,7 @@ class LLM:
840 Sequence[PoolingParams]]] = None,
841 *,
842 prompt_token_ids: list[int],
843 + truncate_prompt_tokens: Optional[int] = None,
844 use_tqdm: bool = True,
845 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
846 prompt_adapter_request: Optional[PromptAdapterRequest] = None,
@@ -851,6 +856,7 @@ class LLM:
856 Sequence[PoolingParams]]] = None,
857 *,
858 prompt_token_ids: list[list[int]],
859 + truncate_prompt_tokens: Optional[int] = None,
860 use_tqdm: bool = True,
861 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
862 prompt_adapter_request: Optional[PromptAdapterRequest] = None,
@@ -864,6 +870,7 @@ class LLM:
870 prompts: None,
871 pooling_params: None,
872 prompt_token_ids: Union[list[int], list[list[int]]],
873 + truncate_prompt_tokens: Optional[int] = None,
874 use_tqdm: bool = True,
875 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
876 prompt_adapter_request: Optional[PromptAdapterRequest] = None,
@@ -882,6 +889,7 @@ class LLM:
889 pooling_params: Optional[Union[PoolingParams,
890 Sequence[PoolingParams]]] = None,
891 prompt_token_ids: Optional[Union[list[int], list[list[int]]]] = None,
892 + truncate_prompt_tokens: Optional[int] = None,
893 use_tqdm: bool = True,
894 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
895 prompt_adapter_request: Optional[PromptAdapterRequest] = None,
@@ -946,10 +954,15 @@ class LLM:
954 for pooling_param in pooling_params:
955 pooling_param.verify(self.llm_engine.model_config)
956
957 + tokenization_kwargs: dict[str, Any] = {}
958 + _validate_truncation_size(self.llm_engine.model_config.max_model_len,
959 + truncate_prompt_tokens, tokenization_kwargs)
960 +
961 self._validate_and_add_requests(
962 prompts=parsed_prompts,
963 params=pooling_params,
964 lora_request=lora_request,
965 + tokenization_kwargs=tokenization_kwargs,
966 prompt_adapter_request=prompt_adapter_request,
967 )
968
@@ -962,6 +975,7 @@ class LLM:
975 prompts: Union[PromptType, Sequence[PromptType]],
976 /,
977 *,
978 + truncate_prompt_tokens: Optional[int] = None,
979 use_tqdm: bool = True,
980 pooling_params: Optional[Union[PoolingParams,
981 Sequence[PoolingParams]]] = None,
@@ -995,6 +1009,7 @@ class LLM:
1009 "Embedding API is only enabled for `--task embed`")
1010
1011 items = self.encode(prompts,
1012 + truncate_prompt_tokens=truncate_prompt_tokens,
1013 use_tqdm=use_tqdm,
1014 pooling_params=pooling_params,
1015 lora_request=lora_request,
@@ -1055,6 +1070,7 @@ class LLM:
1070
1071 encoded_output: list[PoolingRequestOutput] = self.encode(
1072 text_1 + text_2,
1073 + truncate_prompt_tokens=truncate_prompt_tokens,
1074 use_tqdm=use_tqdm,
1075 lora_request=lora_request,
1076 prompt_adapter_request=prompt_adapter_request)
@@ -1098,9 +1114,8 @@ class LLM:
1114 pooling_params = PoolingParams()
1115
1116 tokenization_kwargs: dict[str, Any] = {}
1101 - if truncate_prompt_tokens is not None:
1102 - tokenization_kwargs["truncation"] = True
1103 - tokenization_kwargs["max_length"] = truncate_prompt_tokens
1117 + _validate_truncation_size(self.llm_engine.model_config.max_model_len,
1118 + truncate_prompt_tokens, tokenization_kwargs)
1119
1120 parsed_prompts = []
1121
@@ -1323,6 +1338,7 @@ class LLM:
1338 Sequence[PoolingParams]],
1339 lora_request: Optional[Union[Sequence[LoRARequest], LoRARequest]],
1340 prompt_adapter_request: Optional[PromptAdapterRequest],
1341 + tokenization_kwargs: Optional[dict[str, Any]] = None,
1342 guided_options: Optional[GuidedDecodingRequest] = None,
1343 priority: Optional[list[int]] = None,
1344 ) -> None:
@@ -1359,6 +1375,7 @@ class LLM:
1375 self._add_request(
1376 prompt,
1377 params[i] if isinstance(params, Sequence) else params,
1378 + tokenization_kwargs=tokenization_kwargs,
1379 lora_request=lora_request[i] if isinstance(
1380 lora_request, Sequence) else lora_request,
1381 prompt_adapter_request=prompt_adapter_request,
@@ -1369,6 +1386,7 @@ class LLM:
1386 self,
1387 prompt: PromptType,
1388 params: Union[SamplingParams, PoolingParams],
1389 + tokenization_kwargs: Optional[dict[str, Any]] = None,
1390 lora_request: Optional[LoRARequest] = None,
1391 prompt_adapter_request: Optional[PromptAdapterRequest] = None,
1392 priority: int = 0,
@@ -1379,6 +1397,7 @@ class LLM:
1397 prompt,
1398 params,
1399 lora_request=lora_request,
1400 + tokenization_kwargs=tokenization_kwargs,
1401 prompt_adapter_request=prompt_adapter_request,
1402 priority=priority,
1403 )
vllm/entrypoints/openai/protocol.py
+4 -4
@@ -1014,7 +1014,7 @@ class EmbeddingCompletionRequest(OpenAIBaseModel):
1014 encoding_format: Literal["float", "base64"] = "float"
1015 dimensions: Optional[int] = None
1016 user: Optional[str] = None
1017 - truncate_prompt_tokens: Optional[Annotated[int, Field(ge=1)]] = None
1017 + truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = None
1018
1019 # doc: begin-embedding-pooling-params
1020 additional_data: Optional[Any] = None
@@ -1049,7 +1049,7 @@ class EmbeddingChatRequest(OpenAIBaseModel):
1049 encoding_format: Literal["float", "base64"] = "float"
1050 dimensions: Optional[int] = None
1051 user: Optional[str] = None
1052 - truncate_prompt_tokens: Optional[Annotated[int, Field(ge=1)]] = None
1052 + truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = None
1053
1054 # doc: begin-chat-embedding-pooling-params
1055 additional_data: Optional[Any] = None
@@ -1116,7 +1116,7 @@ class ScoreRequest(OpenAIBaseModel):
1116 model: Optional[str] = None
1117 text_1: Union[list[str], str]
1118 text_2: Union[list[str], str]
1119 - truncate_prompt_tokens: Optional[Annotated[int, Field(ge=1)]] = None
1119 + truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = None
1120
1121 # doc: begin-score-pooling-params
1122 additional_data: Optional[Any] = None
@@ -1142,7 +1142,7 @@ class RerankRequest(OpenAIBaseModel):
1142 query: str
1143 documents: list[str]
1144 top_n: int = Field(default_factory=lambda: 0)
1145 - truncate_prompt_tokens: Optional[Annotated[int, Field(ge=1)]] = None
1145 + truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = None
1146
1147 # doc: begin-rerank-pooling-params
1148 additional_data: Optional[Any] = None
vllm/entrypoints/openai/serving_embedding.py
+4 -10
@@ -21,6 +21,7 @@ from vllm.entrypoints.openai.protocol import (EmbeddingChatRequest,
21 ErrorResponse, UsageInfo)
22 from vllm.entrypoints.openai.serving_engine import OpenAIServing
23 from vllm.entrypoints.openai.serving_models import OpenAIServingModels
24 +from vllm.entrypoints.utils import _validate_truncation_size
25 from vllm.logger import init_logger
26 from vllm.outputs import (EmbeddingOutput, EmbeddingRequestOutput,
27 PoolingRequestOutput)
@@ -85,16 +86,7 @@ class OpenAIServingEmbedding(OpenAIServing):
86 request_id = f"embd-{self._base_request_id(raw_request)}"
87 created_time = int(time.time())
88
88 - truncate_prompt_tokens = None
89 -
90 - if request.truncate_prompt_tokens is not None:
91 - if request.truncate_prompt_tokens <= self.max_model_len:
92 - truncate_prompt_tokens = request.truncate_prompt_tokens
93 - else:
94 - return self.create_error_response(
95 - "truncate_prompt_tokens value is "
96 - "greater than max_model_len."
97 - " Please, select a smaller truncation size.")
89 + truncate_prompt_tokens = request.truncate_prompt_tokens
90
91 pooling_params = request.to_pooling_params()
92
@@ -104,6 +96,8 @@ class OpenAIServingEmbedding(OpenAIServing):
96 return self.create_error_response(str(e))
97
98 try:
99 + truncate_prompt_tokens = _validate_truncation_size(
100 + self.max_model_len, truncate_prompt_tokens)
101 (
102 lora_request,
103 prompt_adapter_request,
vllm/entrypoints/openai/serving_engine.py
+5 -5
@@ -173,7 +173,7 @@ class OpenAIServing:
173 request: AnyRequest,
174 tokenizer: AnyTokenizer,
175 prompt: str,
176 - truncate_prompt_tokens: Optional[Annotated[int, Field(ge=1)]],
176 + truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]],
177 add_special_tokens: bool,
178 ) -> TextTokensPrompt:
179 if (self.model_config.encoder_config is not None
@@ -271,7 +271,7 @@ class OpenAIServing:
271 request: AnyRequest,
272 tokenizer: AnyTokenizer,
273 prompt_input: Union[str, list[int]],
274 - truncate_prompt_tokens: Optional[Annotated[int, Field(ge=1)]] = None,
274 + truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = None,
275 add_special_tokens: bool = True,
276 ) -> TextTokensPrompt:
277 """
@@ -292,7 +292,7 @@ class OpenAIServing:
292 request: AnyRequest,
293 tokenizer: AnyTokenizer,
294 prompt_inputs: Iterable[Union[str, list[int]]],
295 - truncate_prompt_tokens: Optional[Annotated[int, Field(ge=1)]] = None,
295 + truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = None,
296 add_special_tokens: bool = True,
297 ) -> Iterator[TextTokensPrompt]:
298 """
@@ -321,7 +321,7 @@ class OpenAIServing:
321 request: AnyRequest,
322 tokenizer: AnyTokenizer,
323 input_or_inputs: Union[str, list[str], list[int], list[list[int]]],
324 - truncate_prompt_tokens: Optional[Annotated[int, Field(ge=1)]] = None,
324 + truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = None,
325 add_special_tokens: bool = True,
326 ) -> list[TextTokensPrompt]:
327 """
@@ -356,7 +356,7 @@ class OpenAIServing:
356 request: CompletionLikeRequest,
357 tokenizer: AnyTokenizer,
358 input_or_inputs: Union[str, list[str], list[int], list[list[int]]],
359 - truncate_prompt_tokens: Optional[Annotated[int, Field(ge=1)]] = None,
359 + truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = None,
360 add_special_tokens: bool = True,
361 ) -> tuple[list[TextTokensPrompt], list[TokensPrompt]]:
362 request_prompts = await self._tokenize_prompt_input_or_inputs_async(
vllm/entrypoints/openai/serving_pooling.py
+4 -10
@@ -21,6 +21,7 @@ from vllm.entrypoints.openai.protocol import (ErrorResponse,
21 PoolingResponseData, UsageInfo)
22 from vllm.entrypoints.openai.serving_engine import OpenAIServing
23 from vllm.entrypoints.openai.serving_models import OpenAIServingModels
24 +from vllm.entrypoints.utils import _validate_truncation_size
25 from vllm.logger import init_logger
26 from vllm.outputs import PoolingOutput, PoolingRequestOutput
27 from vllm.utils import merge_async_iterators
@@ -85,18 +86,11 @@ class OpenAIServingPooling(OpenAIServing):
86 request_id = f"pool-{self._base_request_id(raw_request)}"
87 created_time = int(time.time())
88
88 - truncate_prompt_tokens = None
89 -
90 - if request.truncate_prompt_tokens is not None:
91 - if request.truncate_prompt_tokens <= self.max_model_len:
92 - truncate_prompt_tokens = request.truncate_prompt_tokens
93 - else:
94 - return self.create_error_response(
95 - "truncate_prompt_tokens value is "
96 - "greater than max_model_len."
97 - " Please, select a smaller truncation size.")
89 + truncate_prompt_tokens = request.truncate_prompt_tokens
90
91 try:
92 + truncate_prompt_tokens = _validate_truncation_size(
93 + self.max_model_len, truncate_prompt_tokens)
94 (
95 lora_request,
96 prompt_adapter_request,
vllm/entrypoints/openai/serving_score.py
+4 -11
@@ -18,6 +18,7 @@ from vllm.entrypoints.openai.serving_engine import OpenAIServing
18 from vllm.entrypoints.openai.serving_models import OpenAIServingModels
19 from vllm.entrypoints.score_utils import (_cosine_similarity,
20 _validate_score_input_lens)
21 +from vllm.entrypoints.utils import _validate_truncation_size
22 from vllm.inputs.data import TokensPrompt
23 from vllm.logger import init_logger
24 from vllm.lora.request import LoRARequest
@@ -231,11 +232,6 @@ class ServingScores(OpenAIServing):
232 truncate_prompt_tokens: Optional[int] = None,
233 ) -> list[PoolingRequestOutput]:
234
234 - tokenization_kwargs: dict[str, Any] = {}
235 - if truncate_prompt_tokens is not None:
236 - tokenization_kwargs["truncation"] = True
237 - tokenization_kwargs["max_length"] = truncate_prompt_tokens
238 -
235 (
236 lora_request,
237 prompt_adapter_request,
@@ -247,12 +243,9 @@ class ServingScores(OpenAIServing):
243
244 tokenizer = await self.engine_client.get_tokenizer(lora_request)
245
250 - if truncate_prompt_tokens is not None and \
251 - truncate_prompt_tokens > self.max_model_len:
252 - raise ValueError(
253 - f"truncate_prompt_tokens value ({truncate_prompt_tokens}) "
254 - f"is greater than max_model_len ({self.max_model_len})."
255 - f" Please, select a smaller truncation size.")
246 + tokenization_kwargs: dict[str, Any] = {}
247 + _validate_truncation_size(self.max_model_len, truncate_prompt_tokens,
248 + tokenization_kwargs)
249
250 trace_headers = (None if raw_request is None else await
251 self._get_trace_headers(raw_request.headers))
vllm/entrypoints/score_utils.py
+1 -1
@@ -46,4 +46,4 @@ def _validate_score_input_lens(
46 if len(texts_1) == 0:
47 raise ValueError("At least one text element must be given")
48 if len(texts_2) == 0:
49 - raise ValueError("At least one text_pair element must be given")
49 + raise ValueError("At least one text_pair element must be given")
\ No newline at end of file
vllm/entrypoints/utils.py
+24
@@ -3,6 +3,7 @@
3 import asyncio
4 import functools
5 import os
6 +from typing import Any, Optional
7
8 from fastapi import Request
9 from fastapi.responses import JSONResponse, StreamingResponse
@@ -134,3 +135,26 @@ def cli_env_setup():
135 if "VLLM_WORKER_MULTIPROC_METHOD" not in os.environ:
136 logger.debug("Setting VLLM_WORKER_MULTIPROC_METHOD to 'spawn'")
137 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
138 +
139 +
140 +def _validate_truncation_size(
141 + max_model_len: int,
142 + truncate_prompt_tokens: Optional[int],
143 + tokenization_kwargs: Optional[dict[str, Any]] = None,
144 +) -> Optional[int]:
145 +
146 + if truncate_prompt_tokens is not None:
147 + if truncate_prompt_tokens <= -1:
148 + truncate_prompt_tokens = max_model_len
149 +
150 + if truncate_prompt_tokens > max_model_len:
151 + raise ValueError(
152 + f"truncate_prompt_tokens value ({truncate_prompt_tokens}) "
153 + f"is greater than max_model_len ({max_model_len})."
154 + f" Please, select a smaller truncation size.")
155 +
156 + if tokenization_kwargs is not None:
157 + tokenization_kwargs["truncation"] = True
158 + tokenization_kwargs["max_length"] = truncate_prompt_tokens
159 +
160 + return truncate_prompt_tokens
vllm/inputs/preprocess.py
+48 -15
@@ -2,7 +2,7 @@
2
3 import asyncio
4 from collections.abc import Mapping
5 -from typing import Optional, Union, cast
5 +from typing import Any, Optional, Union, cast
6
7 from typing_extensions import assert_never
8
@@ -183,18 +183,21 @@ class InputPreprocessor:
183 self,
184 prompt: str,
185 lora_request: Optional[LoRARequest],
186 + tokenization_kwargs: Optional[dict[str, Any]] = None,
187 ) -> list[int]:
188 """
189 Apply the model's tokenizer to a text prompt, returning the
190 corresponding token IDs.
191 """
192 tokenizer = self.get_tokenizer_group()
192 - add_special_tokens = None
193 + if tokenization_kwargs is None:
194 + tokenization_kwargs = {}
195 +
196 if self.model_config.hf_config.model_type == "whisper":
197 # For Whisper, special tokens should be provided by the user based
198 # on the task and language of their request. Also needed to avoid
199 # appending an EOS token to the prompt which disrupts generation.
197 - add_special_tokens = False
200 + tokenization_kwargs["add_special_tokens"] = False
201
202 if (self.model_config.encoder_config is not None
203 and self.model_config.encoder_config.get(
@@ -203,25 +206,27 @@ class InputPreprocessor:
206
207 return tokenizer.encode(prompt=prompt,
208 lora_request=lora_request,
206 - add_special_tokens=add_special_tokens)
209 + **tokenization_kwargs)
210
211 async def _tokenize_prompt_async(
212 self,
213 prompt: str,
214 lora_request: Optional[LoRARequest],
215 + tokenization_kwargs: Optional[dict[str, Any]] = None,
216 ) -> list[int]:
217 """Async version of :meth:`_tokenize_prompt`."""
218 tokenizer = self.get_tokenizer_group()
215 - add_special_tokens = None
219 + if tokenization_kwargs is None:
220 + tokenization_kwargs = {}
221 +
222 if self.model_config.hf_config.model_type == "whisper":
223 # For Whisper, special tokens should be provided by the user based
224 # on the task and language of their request. Also needed to avoid
225 # appending an EOS token to the prompt which disrupts generation.
220 - add_special_tokens = False
221 - return await tokenizer.encode_async(
222 - prompt=prompt,
223 - lora_request=lora_request,
224 - add_special_tokens=add_special_tokens)
226 + tokenization_kwargs["add_special_tokens"] = False
227 + return await tokenizer.encode_async(prompt=prompt,
228 + lora_request=lora_request,
229 + **tokenization_kwargs)
230
231 def _process_multimodal(
232 self,
@@ -281,6 +286,7 @@ class InputPreprocessor:
286 def _prompt_to_llm_inputs(
287 self,
288 prompt: SingletonPrompt,
289 + tokenization_kwargs: Optional[dict[str, Any]] = None,
290 lora_request: Optional[LoRARequest] = None,
291 return_mm_hashes: bool = False,
292 ) -> SingletonInputs:
@@ -304,6 +310,7 @@ class InputPreprocessor:
310 prompt_token_ids = self._tokenize_prompt(
311 prompt_text,
312 lora_request=lora_request,
313 + tokenization_kwargs=tokenization_kwargs,
314 )
315
316 return token_inputs(
@@ -352,6 +359,7 @@ class InputPreprocessor:
359 prompt_token_ids = self._tokenize_prompt(
360 prompt_text,
361 lora_request=lora_request,
362 + tokenization_kwargs=tokenization_kwargs,
363 )
364
365 return token_inputs(
@@ -364,6 +372,7 @@ class InputPreprocessor:
372 async def _prompt_to_llm_inputs_async(
373 self,
374 prompt: SingletonPrompt,
375 + tokenization_kwargs: Optional[dict[str, Any]] = None,
376 lora_request: Optional[LoRARequest] = None,
377 return_mm_hashes: bool = False,
378 ) -> SingletonInputs:
@@ -375,6 +384,7 @@ class InputPreprocessor:
384 prompt_token_ids = await self._tokenize_prompt_async(
385 prompt_text,
386 lora_request=lora_request,
387 + tokenization_kwargs=tokenization_kwargs,
388 )
389
390 return token_inputs(
@@ -517,6 +527,7 @@ class InputPreprocessor:
527 def _process_encoder_decoder_prompt(
528 self,
529 prompt: PromptType,
530 + tokenization_kwargs: Optional[dict[str, Any]] = None,
531 ) -> EncoderDecoderInputs:
532 """
533 For encoder/decoder models only:
@@ -553,7 +564,9 @@ class InputPreprocessor:
564
565 if is_explicit_encoder_decoder_prompt(prompt):
566 encoder_inputs = self._prompt_to_llm_inputs(
556 - prompt["encoder_prompt"])
567 + prompt["encoder_prompt"],
568 + tokenization_kwargs=tokenization_kwargs,
569 + )
570 if (decoder_input := prompt["decoder_prompt"]) is None:
571 decoder_inputs = None
572 else:
@@ -565,7 +578,10 @@ class InputPreprocessor:
578 self._separate_enc_dec_inputs_from_mm_processor_outputs(
579 encoder_inputs, decoder_inputs))
580 else:
568 - inputs = self._prompt_to_llm_inputs(prompt)
581 + inputs = self._prompt_to_llm_inputs(
582 + prompt,
583 + tokenization_kwargs=tokenization_kwargs,
584 + )
585 if self.model_config.is_multimodal_model:
586 # Encoder-Decoder Multimodal model
587 encoder_inputs, decoder_inputs = (
@@ -581,6 +597,7 @@ class InputPreprocessor:
597 async def _process_encoder_decoder_prompt_async(
598 self,
599 prompt: PromptType,
600 + tokenization_kwargs: Optional[dict[str, Any]] = None,
601 ) -> EncoderDecoderInputs:
602 """Async version of :meth:`_process_encoder_decoder_prompt`."""
603 encoder_inputs: SingletonInputs
@@ -588,13 +605,18 @@ class InputPreprocessor:
605
606 if is_explicit_encoder_decoder_prompt(prompt):
607 encoder_task = self._prompt_to_llm_inputs_async(
591 - prompt["encoder_prompt"])
608 + prompt["encoder_prompt"],
609 + tokenization_kwargs=tokenization_kwargs,
610 + )
611
612 if (decoder_input := prompt["decoder_prompt"]) is None:
613 encoder_inputs = await encoder_task
614 decoder_inputs = None
615 else:
597 - decoder_task = self._prompt_to_llm_inputs_async(decoder_input)
616 + decoder_task = self._prompt_to_llm_inputs_async(
617 + decoder_input,
618 + tokenization_kwargs=tokenization_kwargs,
619 + )
620
621 encoder_inputs, decoder_inputs = await asyncio.gather(
622 encoder_task, decoder_task)
@@ -606,7 +628,10 @@ class InputPreprocessor:
628 self._separate_enc_dec_inputs_from_mm_processor_outputs(
629 encoder_inputs, decoder_inputs))
630 else:
609 - inputs = await self._prompt_to_llm_inputs_async(prompt)
631 + inputs = await self._prompt_to_llm_inputs_async(
632 + prompt,
633 + tokenization_kwargs=tokenization_kwargs,
634 + )
635 if self.model_config.is_multimodal_model:
636 # Encoder-Decoder Multimodal model
637 encoder_inputs, decoder_inputs = (
@@ -638,6 +663,7 @@ class InputPreprocessor:
663 def _process_decoder_only_prompt(
664 self,
665 prompt: SingletonPrompt,
666 + tokenization_kwargs: Optional[dict[str, Any]] = None,
667 lora_request: Optional[LoRARequest] = None,
668 prompt_adapter_request: Optional[PromptAdapterRequest] = None,
669 return_mm_hashes: bool = False,
@@ -660,6 +686,7 @@ class InputPreprocessor:
686
687 prompt_comps = self._prompt_to_llm_inputs(
688 prompt,
689 + tokenization_kwargs=tokenization_kwargs,
690 lora_request=lora_request,
691 return_mm_hashes=return_mm_hashes,
692 )
@@ -672,6 +699,7 @@ class InputPreprocessor:
699 async def _process_decoder_only_prompt_async(
700 self,
701 prompt: SingletonPrompt,
702 + tokenization_kwargs: Optional[dict[str, Any]] = None,
703 lora_request: Optional[LoRARequest] = None,
704 prompt_adapter_request: Optional[PromptAdapterRequest] = None,
705 return_mm_hashes: bool = False,
@@ -679,6 +707,7 @@ class InputPreprocessor:
707 """Async version of :meth:`_process_decoder_only_prompt`."""
708 prompt_comps = await self._prompt_to_llm_inputs_async(
709 prompt,
710 + tokenization_kwargs=tokenization_kwargs,
711 lora_request=lora_request,
712 return_mm_hashes=return_mm_hashes,
713 )
@@ -691,6 +720,7 @@ class InputPreprocessor:
720 def preprocess(
721 self,
722 prompt: PromptType,
723 + tokenization_kwargs: Optional[dict[str, Any]] = None,
724 lora_request: Optional[LoRARequest] = None,
725 prompt_adapter_request: Optional[PromptAdapterRequest] = None,
726 return_mm_hashes: bool = False,
@@ -711,6 +741,7 @@ class InputPreprocessor:
741 # Decoder-only operation
742 return self._process_decoder_only_prompt(
743 prompt,
744 + tokenization_kwargs=tokenization_kwargs,
745 lora_request=lora_request,
746 prompt_adapter_request=prompt_adapter_request,
747 return_mm_hashes=return_mm_hashes,
@@ -719,6 +750,7 @@ class InputPreprocessor:
750 async def preprocess_async(
751 self,
752 prompt: PromptType,
753 + tokenization_kwargs: Optional[dict[str, Any]] = None,
754 lora_request: Optional[LoRARequest] = None,
755 prompt_adapter_request: Optional[PromptAdapterRequest] = None,
756 return_mm_hashes: bool = False,
@@ -739,6 +771,7 @@ class InputPreprocessor:
771 # Decoder-only operation
772 return await self._process_decoder_only_prompt_async(
773 prompt,
774 + tokenization_kwargs=tokenization_kwargs,
775 lora_request=lora_request,
776 prompt_adapter_request=prompt_adapter_request,
777 return_mm_hashes=return_mm_hashes,
vllm/sampling_params.py
+4 -3
@@ -186,9 +186,10 @@ class SamplingParams(
186 logits_processors: list of functions that modify logits based on
187 previously generated tokens, and optionally prompt tokens as
188 a first argument.
189 - truncate_prompt_tokens: If set to an integer k, will use only the last k
190 - tokens from the prompt (i.e., left truncation). Defaults to None
191 - (i.e., no truncation).
189 + truncate_prompt_tokens: If set to -1, will use the truncation size
190 + supported by the model. If set to an integer k, will use only
191 + the last k tokens from the prompt (i.e., left truncation).
192 + Defaults to None (i.e., no truncation).
193 guided_decoding: If provided, the engine will construct a guided
194 decoding logits processor from these parameters. Defaults to None.
195 logit_bias: If provided, the engine will construct a logits processor
vllm/transformers_utils/tokenizer.py
+12 -2
@@ -55,6 +55,8 @@ def encode_tokens(
55 tokenizer: AnyTokenizer,
56 text: str,
57 *,
58 + truncation: Optional[bool] = None,
59 + max_length: Optional[int] = None,
60 add_special_tokens: Optional[bool] = None,
61 ) -> list[int]:
62 """
@@ -64,10 +66,18 @@ def encode_tokens(
66 :code:`add_special_tokens=None` means to use the backend's default
67 settings.
68 """
69 +
70 + kw_args: dict[str, Any] = {}
71 + if max_length is not None:
72 + kw_args["max_length"] = max_length
73 +
74 + if truncation is not None:
75 + kw_args["truncation"] = truncation
76 +
77 if add_special_tokens is not None:
68 - return tokenizer.encode(text, add_special_tokens=add_special_tokens)
78 + kw_args["add_special_tokens"] = add_special_tokens
79
70 - return tokenizer.encode(text)
80 + return tokenizer.encode(text, **kw_args)
81
82
83 def get_cached_tokenizer(tokenizer: AnyTokenizer) -> AnyTokenizer:
vllm/transformers_utils/tokenizer_base.py
+2
@@ -94,6 +94,8 @@ class TokenizerBase(ABC):
94 @abstractmethod
95 def encode(self,
96 text: str,
97 + truncation: Optional[bool] = None,
98 + max_length: Optional[int] = None,
99 add_special_tokens: Optional[bool] = None) -> list[int]:
100 raise NotImplementedError()
101
vllm/transformers_utils/tokenizer_group.py
+9
@@ -45,11 +45,16 @@ class TokenizerGroup:
45
46 def encode(self,
47 prompt: str,
48 + max_length: Optional[int] = None,
49 + truncation: Optional[bool] = None,
50 lora_request: Optional[LoRARequest] = None,
51 add_special_tokens: Optional[bool] = None) -> List[int]:
52 +
53 tokenizer = self.get_lora_tokenizer(lora_request)
54 ret = encode_tokens(tokenizer,
55 prompt,
56 + max_length=max_length,
57 + truncation=truncation,
58 add_special_tokens=add_special_tokens)
59 self._raise_if_input_too_long(ret, lora_request)
60 return ret
@@ -57,11 +62,15 @@ class TokenizerGroup:
62 async def encode_async(
63 self,
64 prompt: str,
65 + max_length: Optional[int] = None,
66 + truncation: Optional[bool] = None,
67 lora_request: Optional[LoRARequest] = None,
68 add_special_tokens: Optional[bool] = None) -> List[int]:
69 tokenizer = await self.get_lora_tokenizer_async(lora_request)
70 ret = encode_tokens(tokenizer,
71 prompt,
72 + max_length=max_length,
73 + truncation=truncation,
74 add_special_tokens=add_special_tokens)
75 self._raise_if_input_too_long(ret, lora_request)
76 return ret
vllm/transformers_utils/tokenizers/mistral.py
+2
@@ -359,6 +359,8 @@ class MistralTokenizer(TokenizerBase):
359
360 def encode(self,
361 text: str,
362 + truncation: Optional[bool] = None,
363 + max_length: Optional[int] = None,
364 add_special_tokens: Optional[bool] = None) -> List[int]:
365 # `encode` should only be used for prompt completion
366 # it should never be used for chat_completion.
vllm/v1/engine/async_llm.py
+4 -2
@@ -2,7 +2,7 @@
2 import asyncio
3 from collections.abc import AsyncGenerator, Mapping
4 from copy import copy
5 -from typing import Optional, Union
5 +from typing import Any, Optional, Union
6
7 import numpy as np
8
@@ -201,6 +201,7 @@ class AsyncLLM(EngineClient):
201 params: Union[SamplingParams, PoolingParams],
202 arrival_time: Optional[float] = None,
203 lora_request: Optional[LoRARequest] = None,
204 + tokenization_kwargs: Optional[dict[str, Any]] = None,
205 trace_headers: Optional[Mapping[str, str]] = None,
206 prompt_adapter_request: Optional[PromptAdapterRequest] = None,
207 priority: int = 0,
@@ -219,7 +220,8 @@ class AsyncLLM(EngineClient):
220 # Convert Input --> Request.
221 prompt_str, request = self.processor.process_inputs(
222 request_id, prompt, params, arrival_time, lora_request,
222 - trace_headers, prompt_adapter_request, priority)
223 + tokenization_kwargs, trace_headers, prompt_adapter_request,
224 + priority)
225
226 if params.n == 1:
227 await self._add_request(request, prompt_str, None, 0, queue)
vllm/v1/engine/llm_engine.py
+3 -1
@@ -175,6 +175,7 @@ class LLMEngine:
175 params: Union[SamplingParams, PoolingParams],
176 arrival_time: Optional[float] = None,
177 lora_request: Optional[LoRARequest] = None,
178 + tokenization_kwargs: Optional[dict[str, Any]] = None,
179 trace_headers: Optional[Mapping[str, str]] = None,
180 prompt_adapter_request: Optional[PromptAdapterRequest] = None,
181 priority: int = 0,
@@ -182,7 +183,8 @@ class LLMEngine:
183 # Process raw inputs into the request.
184 prompt_str, request = self.processor.process_inputs(
185 request_id, prompt, params, arrival_time, lora_request,
185 - trace_headers, prompt_adapter_request, priority)
186 + tokenization_kwargs, trace_headers, prompt_adapter_request,
187 + priority)
188
189 n = params.n if isinstance(params, SamplingParams) else 1
190
vllm/v1/engine/processor.py
+3 -1
@@ -2,7 +2,7 @@
2
3 import time
4 from collections.abc import Mapping, Sequence
5 -from typing import Literal, Optional, Union
5 +from typing import Any, Literal, Optional, Union
6
7 from vllm.config import VllmConfig
8 from vllm.inputs import ProcessorInputs, PromptType, SingletonInputs
@@ -198,6 +198,7 @@ class Processor:
198 params: Union[SamplingParams, PoolingParams],
199 arrival_time: Optional[float] = None,
200 lora_request: Optional[LoRARequest] = None,
201 + tokenization_kwargs: Optional[dict[str, Any]] = None,
202 trace_headers: Optional[Mapping[str, str]] = None,
203 prompt_adapter_request: Optional[PromptAdapterRequest] = None,
204 priority: int = 0,
@@ -224,6 +225,7 @@ class Processor:
225 # 3. Apply prompt adapter to prompt token ids if one exists.
226 processed_inputs: ProcessorInputs = self.input_preprocessor.preprocess(
227 prompt,
228 + tokenization_kwargs=tokenization_kwargs,
229 lora_request=lora_request,
230 prompt_adapter_request=prompt_adapter_request,
231 return_mm_hashes=self.use_hash,