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

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, 01:00 UTC

Add tokenization_kwargs to encode for embedding model truncation (#21033)

Wang Yijun committed Jul 22, 2025 at 23:24 UTC 44554a00685b7e60eeed5883986749405e538806

3 files changed +20 -3

vllm/engine/async_llm_engine.py

@@ -438,6 +438,7 @@ class _AsyncLLMEngine(LLMEngine):

438 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

439 priority: int = 0,

440 data_parallel_rank: Optional[int] = None,

441 + tokenization_kwargs: Optional[dict[str, Any]] = None,

442 ) -> None:

443 """

444 Async version of

@@ -468,6 +469,7 @@ class _AsyncLLMEngine(LLMEngine):

469 prompt,

470 lora_request=lora_request,

471 prompt_adapter_request=prompt_adapter_request,

472 + tokenization_kwargs=tokenization_kwargs,

473 )

474

475 if isinstance(params, SamplingParams) and \

@@ -862,6 +864,7 @@ class AsyncLLMEngine(EngineClient):

864 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

865 priority: int = 0,

866 data_parallel_rank: Optional[int] = None,

867 + tokenization_kwargs: Optional[dict[str, Any]] = None,

868 ) -> AsyncGenerator[Union[RequestOutput, PoolingRequestOutput], None]:

869 if not self.is_running:

870 if self.start_engine_loop:

@@ -889,6 +892,7 @@ class AsyncLLMEngine(EngineClient):

892 prompt_adapter_request=prompt_adapter_request,

893 priority=priority,

894 data_parallel_rank=data_parallel_rank,

895 + tokenization_kwargs=tokenization_kwargs,

896 )

897

898 return stream.generator()

@@ -996,6 +1000,7 @@ class AsyncLLMEngine(EngineClient):

1000 lora_request: Optional[LoRARequest] = None,

1001 trace_headers: Optional[Mapping[str, str]] = None,

1002 priority: int = 0,

1003 + tokenization_kwargs: Optional[dict[str, Any]] = None,

1004 ) -> AsyncGenerator[PoolingRequestOutput, None]:

1005 """Generate outputs for a request from a pooling model.

1006

@@ -1070,6 +1075,7 @@ class AsyncLLMEngine(EngineClient):

1075 lora_request=lora_request,

1076 trace_headers=trace_headers,

1077 priority=priority,

1078 + tokenization_kwargs=tokenization_kwargs,

1079 ):

1080 yield LLMEngine.validate_output(output, PoolingRequestOutput)

1081 except asyncio.CancelledError:

vllm/entrypoints/llm.py

+12 -3

@@ -965,6 +965,7 @@ class LLM:

965 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

966 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

967 pooling_task: PoolingTask = "encode",

968 + tokenization_kwargs: Optional[dict[str, Any]] = None,

969 ) -> list[PoolingRequestOutput]:

970 ...

971

@@ -981,6 +982,7 @@ class LLM:

982 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

983 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

984 pooling_task: PoolingTask = "encode",

985 + tokenization_kwargs: Optional[dict[str, Any]] = None,

986 ) -> list[PoolingRequestOutput]:

987 ...

988

@@ -997,6 +999,7 @@ class LLM:

999 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

1000 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

1001 pooling_task: PoolingTask = "encode",

1002 + tokenization_kwargs: Optional[dict[str, Any]] = None,

1003 ) -> list[PoolingRequestOutput]:

1004 ...

1005

@@ -1014,6 +1017,7 @@ class LLM:

1017 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

1018 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

1019 pooling_task: PoolingTask = "encode",

1020 + tokenization_kwargs: Optional[dict[str, Any]] = None,

1021 ) -> list[PoolingRequestOutput]:

1022 ...

1023

@@ -1031,6 +1035,7 @@ class LLM:

1035 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

1036 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

1037 pooling_task: PoolingTask = "encode",

1038 + tokenization_kwargs: Optional[dict[str, Any]] = None,

1039 ) -> list[PoolingRequestOutput]:

1040 ...

1041

@@ -1046,6 +1051,7 @@ class LLM:

1051 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

1052 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

1053 pooling_task: PoolingTask = "encode",

1054 + tokenization_kwargs: Optional[dict[str, Any]] = None,

1055 ) -> list[PoolingRequestOutput]:

1056 ...

1057

@@ -1066,6 +1072,7 @@ class LLM:

1072 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

1073 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

1074 pooling_task: PoolingTask = "encode",

1075 + tokenization_kwargs: Optional[dict[str, Any]] = None,

1076 ) -> list[PoolingRequestOutput]:

1077 """Apply pooling to the hidden states corresponding to the input

1078 prompts.

@@ -1131,9 +1138,11 @@ class LLM:

1138 for pooling_param in pooling_params:

1139 pooling_param.verify(pooling_task, model_config)

1140

1134 - tokenization_kwargs = dict[str, Any]()

1135 - _validate_truncation_size(model_config.max_model_len,

1136 - truncate_prompt_tokens, tokenization_kwargs)

1141 + if tokenization_kwargs is None:

1142 + tokenization_kwargs = dict[str, Any]()

1143 + _validate_truncation_size(model_config.max_model_len,

1144 + truncate_prompt_tokens,

1145 + tokenization_kwargs)

1146

1147 self._validate_and_add_requests(

1148 prompts=parsed_prompts,

vllm/v1/engine/async_llm.py

@@ -437,6 +437,7 @@ class AsyncLLM(EngineClient):

437 lora_request: Optional[LoRARequest] = None,

438 trace_headers: Optional[Mapping[str, str]] = None,

439 priority: int = 0,

440 + tokenization_kwargs: Optional[dict[str, Any]] = None,

441 ) -> AsyncGenerator[PoolingRequestOutput, None]:

442 """

443 Main function called by the API server to kick off a request

@@ -465,6 +466,7 @@ class AsyncLLM(EngineClient):

466 lora_request=lora_request,

467 trace_headers=trace_headers,

468 priority=priority,

469 + tokenization_kwargs=tokenization_kwargs,

470 )

471

472 # The output_handler task pushes items into the queue.