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 3, 20:15 UTC
[Refactor] Clean up pooling models `build_tok_params` logic (#42341)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
May 12, 2026 at 11:05 UTC
a1b2d874988ad2c43afddb17be0e2f7d26a1a8f0
5 files changed
+115
-136
vllm/entrypoints/pooling/base/protocol.py
+84
-1
@@ -6,13 +6,14 @@ from typing import Annotated, Any, Literal
6
7
from pydantic import Field, model_validator
8
9
+from vllm.config import ModelConfig
10
from vllm.entrypoints.chat_utils import (
11
ChatCompletionMessageParam,
12
ChatTemplateContentFormatOption,
13
)
14
from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel
15
from vllm.exceptions import VLLMValidationError
15
-from vllm.renderers import ChatParams, merge_kwargs
16
+from vllm.renderers import ChatParams, TokenizeParams, merge_kwargs
17
from vllm.utils import random_uuid
18
from vllm.utils.serial_utils import EmbedDType, EncodingFormat, Endianness
19
@@ -68,6 +69,88 @@ class PoolingBasicRequestMixin(OpenAIBaseModel):
69
)
70
# --8<-- [end:pooling-common-extra-params]
71
72
+ def _build_pooling_tok_params(
73
+ self,
74
+ model_config: ModelConfig,
75
+ *,
76
+ add_special_tokens: bool,
77
+ max_total_tokens: int | None,
78
+ max_output_tokens: int,
79
+ max_total_tokens_param: str = "max_model_len",
80
+ max_output_tokens_param: str | None = None,
81
+ ) -> TokenizeParams:
82
+ encoder_config = model_config.encoder_config or {}
83
+ if max_output_tokens_param is None:
84
+ return TokenizeParams(
85
+ max_total_tokens=max_total_tokens,
86
+ max_output_tokens=max_output_tokens,
87
+ truncate_prompt_tokens=self.truncate_prompt_tokens,
88
+ truncation_side=self.truncation_side,
89
+ do_lower_case=encoder_config.get("do_lower_case", False),
90
+ add_special_tokens=add_special_tokens,
91
+ max_total_tokens_param=max_total_tokens_param,
92
+ )
93
+
94
+ return TokenizeParams(
95
+ max_total_tokens=max_total_tokens,
96
+ max_output_tokens=max_output_tokens,
97
+ truncate_prompt_tokens=self.truncate_prompt_tokens,
98
+ truncation_side=self.truncation_side,
99
+ do_lower_case=encoder_config.get("do_lower_case", False),
100
+ add_special_tokens=add_special_tokens,
101
+ max_total_tokens_param=max_total_tokens_param,
102
+ max_output_tokens_param=max_output_tokens_param,
103
+ )
104
+
105
+
106
+class PoolingTokenizeParamsMixin:
107
+ add_special_tokens: bool
108
+
109
+ def _build_pooling_tok_params(
110
+ self,
111
+ model_config: ModelConfig,
112
+ *,
113
+ add_special_tokens: bool,
114
+ max_total_tokens: int | None,
115
+ max_output_tokens: int,
116
+ max_total_tokens_param: str = "max_model_len",
117
+ max_output_tokens_param: str | None = None,
118
+ ) -> TokenizeParams:
119
+ raise NotImplementedError
120
+
121
+
122
+class FixedMaxLenTokenizeParamsMixin(PoolingTokenizeParamsMixin):
123
+ def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
124
+ return self._build_pooling_tok_params(
125
+ model_config,
126
+ add_special_tokens=self.add_special_tokens,
127
+ max_total_tokens=model_config.max_model_len,
128
+ max_output_tokens=0,
129
+ )
130
+
131
+
132
+class EmbeddingTokenizeParamsMixin(PoolingTokenizeParamsMixin):
133
+ def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
134
+ default_max_total_tokens = model_config.max_model_len
135
+ max_total_tokens: int | None = default_max_total_tokens
136
+ max_output_tokens = 0
137
+
138
+ pooler_config = model_config.pooler_config
139
+ if pooler_config is not None:
140
+ if pooler_config.enable_chunked_processing:
141
+ max_total_tokens = None
142
+ else:
143
+ max_embed_len = pooler_config.max_embed_len or default_max_total_tokens
144
+ max_output_tokens = default_max_total_tokens - max_embed_len
145
+
146
+ return self._build_pooling_tok_params(
147
+ model_config,
148
+ add_special_tokens=self.add_special_tokens,
149
+ max_total_tokens=max_total_tokens,
150
+ max_output_tokens=max_output_tokens,
151
+ max_output_tokens_param="max_model_len - max_embed_len",
152
+ )
153
+
154
155
class CompletionRequestMixin(OpenAIBaseModel):
156
# --8<-- [start:completion-params]
vllm/entrypoints/pooling/classify/protocol.py
+9
-30
@@ -7,16 +7,15 @@ from typing import TypeAlias
7
from pydantic import Field
8
9
from vllm import PoolingParams
10
-from vllm.config import ModelConfig
10
from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel, UsageInfo
11
from vllm.logger import init_logger
13
-from vllm.renderers import TokenizeParams
12
from vllm.utils import random_uuid
13
14
from ..base.protocol import (
15
ChatRequestMixin,
16
ClassifyRequestMixin,
17
CompletionRequestMixin,
18
+ FixedMaxLenTokenizeParamsMixin,
19
PoolingBasicRequestMixin,
20
)
21
@@ -24,21 +23,11 @@ logger = init_logger(__name__)
23
24
25
class ClassificationCompletionRequest(
27
- PoolingBasicRequestMixin, CompletionRequestMixin, ClassifyRequestMixin
26
+ PoolingBasicRequestMixin,
27
+ CompletionRequestMixin,
28
+ ClassifyRequestMixin,
29
+ FixedMaxLenTokenizeParamsMixin,
30
):
29
- def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
30
- encoder_config = model_config.encoder_config or {}
31
-
32
- return TokenizeParams(
33
- max_total_tokens=model_config.max_model_len,
34
- max_output_tokens=0,
35
- truncate_prompt_tokens=self.truncate_prompt_tokens,
36
- truncation_side=self.truncation_side,
37
- do_lower_case=encoder_config.get("do_lower_case", False),
38
- add_special_tokens=self.add_special_tokens,
39
- max_total_tokens_param="max_model_len",
40
- )
41
-
31
def to_pooling_params(self):
32
return PoolingParams(
33
task="classify",
@@ -47,21 +36,11 @@ class ClassificationCompletionRequest(
36
37
38
class ClassificationChatRequest(
50
- PoolingBasicRequestMixin, ChatRequestMixin, ClassifyRequestMixin
39
+ PoolingBasicRequestMixin,
40
+ ChatRequestMixin,
41
+ ClassifyRequestMixin,
42
+ FixedMaxLenTokenizeParamsMixin,
43
):
52
- def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
53
- encoder_config = model_config.encoder_config or {}
54
-
55
- return TokenizeParams(
56
- max_total_tokens=model_config.max_model_len,
57
- max_output_tokens=0,
58
- truncate_prompt_tokens=self.truncate_prompt_tokens,
59
- truncation_side=self.truncation_side,
60
- do_lower_case=encoder_config.get("do_lower_case", False),
61
- add_special_tokens=self.add_special_tokens,
62
- max_total_tokens_param="max_model_len",
63
- )
64
-
44
def to_pooling_params(self):
45
return PoolingParams(
46
task="classify",
vllm/entrypoints/pooling/embed/protocol.py
+9
-63
@@ -16,62 +16,24 @@ import pybase64 as base64
16
from pydantic import BaseModel, Field
17
18
from vllm import PoolingParams
19
-from vllm.config import ModelConfig
19
from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel, UsageInfo
21
-from vllm.renderers import TokenizeParams
20
from vllm.utils import random_uuid
21
22
from ..base.protocol import (
23
ChatRequestMixin,
24
CompletionRequestMixin,
25
+ EmbeddingTokenizeParamsMixin,
26
EmbedRequestMixin,
27
PoolingBasicRequestMixin,
28
)
29
31
-# ---------------------------------------------------------------------------
32
-# OpenAI /v1/embeddings — request models
33
-# ---------------------------------------------------------------------------
34
-
35
-
36
-def _get_max_total_output_tokens(
37
- model_config: ModelConfig,
38
-) -> tuple[int | None, int]:
39
- max_total_tokens = model_config.max_model_len
40
- pooler_config = model_config.pooler_config
41
-
42
- if pooler_config is None:
43
- return max_total_tokens, 0
44
-
45
- if pooler_config.enable_chunked_processing:
46
- return None, 0
47
-
48
- max_embed_len = pooler_config.max_embed_len or max_total_tokens
49
- max_output_tokens = max_total_tokens - max_embed_len
50
- return max_total_tokens, max_output_tokens
51
-
30
31
class EmbeddingCompletionRequest(
54
- PoolingBasicRequestMixin, CompletionRequestMixin, EmbedRequestMixin
32
+ PoolingBasicRequestMixin,
33
+ CompletionRequestMixin,
34
+ EmbedRequestMixin,
35
+ EmbeddingTokenizeParamsMixin,
36
):
56
- def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
57
- encoder_config = model_config.encoder_config or {}
58
-
59
- (
60
- max_total_tokens,
61
- max_output_tokens,
62
- ) = _get_max_total_output_tokens(model_config)
63
-
64
- return TokenizeParams(
65
- max_total_tokens=max_total_tokens,
66
- max_output_tokens=max_output_tokens,
67
- truncate_prompt_tokens=self.truncate_prompt_tokens,
68
- truncation_side=self.truncation_side,
69
- do_lower_case=encoder_config.get("do_lower_case", False),
70
- add_special_tokens=self.add_special_tokens,
71
- max_total_tokens_param="max_model_len",
72
- max_output_tokens_param="max_model_len - max_embed_len",
73
- )
74
-
37
def to_pooling_params(self):
38
return PoolingParams(
39
task="embed",
@@ -81,27 +43,11 @@ class EmbeddingCompletionRequest(
43
44
45
class EmbeddingChatRequest(
84
- PoolingBasicRequestMixin, ChatRequestMixin, EmbedRequestMixin
46
+ PoolingBasicRequestMixin,
47
+ ChatRequestMixin,
48
+ EmbedRequestMixin,
49
+ EmbeddingTokenizeParamsMixin,
50
):
86
- def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
87
- encoder_config = model_config.encoder_config or {}
88
-
89
- (
90
- max_total_tokens,
91
- max_output_tokens,
92
- ) = _get_max_total_output_tokens(model_config)
93
-
94
- return TokenizeParams(
95
- max_total_tokens=max_total_tokens,
96
- max_output_tokens=max_output_tokens,
97
- truncate_prompt_tokens=self.truncate_prompt_tokens,
98
- truncation_side=self.truncation_side,
99
- do_lower_case=encoder_config.get("do_lower_case", False),
100
- add_special_tokens=self.add_special_tokens,
101
- max_total_tokens_param="max_model_len",
102
- max_output_tokens_param="max_model_len - max_embed_len",
103
- )
104
-
51
def to_pooling_params(self):
52
return PoolingParams(
53
task="embed",
vllm/entrypoints/pooling/pooling/protocol.py
+10
-35
@@ -18,6 +18,7 @@ from ..base.protocol import (
18
CompletionRequestMixin,
19
EmbedRequestMixin,
20
EncodingRequestMixin,
21
+ FixedMaxLenTokenizeParamsMixin,
22
PoolingBasicRequestMixin,
23
)
24
@@ -27,22 +28,10 @@ class PoolingCompletionRequest(
28
CompletionRequestMixin,
29
EmbedRequestMixin,
30
ClassifyRequestMixin,
31
+ FixedMaxLenTokenizeParamsMixin,
32
):
33
task: PoolingTask | None = None
34
33
- def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
34
- encoder_config = model_config.encoder_config or {}
35
-
36
- return TokenizeParams(
37
- max_total_tokens=model_config.max_model_len,
38
- max_output_tokens=0,
39
- truncate_prompt_tokens=self.truncate_prompt_tokens,
40
- truncation_side=self.truncation_side,
41
- do_lower_case=encoder_config.get("do_lower_case", False),
42
- add_special_tokens=self.add_special_tokens,
43
- max_total_tokens_param="max_model_len",
44
- )
45
-
35
def to_pooling_params(self):
36
return PoolingParams(
37
task=self.task,
@@ -52,23 +41,14 @@ class PoolingCompletionRequest(
41
42
43
class PoolingChatRequest(
55
- PoolingBasicRequestMixin, ChatRequestMixin, EmbedRequestMixin, ClassifyRequestMixin
44
+ PoolingBasicRequestMixin,
45
+ ChatRequestMixin,
46
+ EmbedRequestMixin,
47
+ ClassifyRequestMixin,
48
+ FixedMaxLenTokenizeParamsMixin,
49
):
50
task: PoolingTask | None = None
51
59
- def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
60
- encoder_config = model_config.encoder_config or {}
61
-
62
- return TokenizeParams(
63
- max_total_tokens=model_config.max_model_len,
64
- max_output_tokens=0,
65
- truncate_prompt_tokens=self.truncate_prompt_tokens,
66
- truncation_side=self.truncation_side,
67
- do_lower_case=encoder_config.get("do_lower_case", False),
68
- add_special_tokens=self.add_special_tokens,
69
- max_total_tokens_param="max_model_len",
70
- )
71
-
52
def to_pooling_params(self):
53
return PoolingParams(
54
task=self.task,
@@ -85,16 +65,11 @@ class IOProcessorRequest(PoolingBasicRequestMixin, EncodingRequestMixin, Generic
65
task: PoolingTask = "plugin"
66
67
def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
88
- encoder_config = model_config.encoder_config or {}
89
-
90
- return TokenizeParams(
68
+ return self._build_pooling_tok_params(
69
+ model_config,
70
+ add_special_tokens=not model_config.is_encoder_decoder,
71
max_total_tokens=model_config.max_model_len,
72
max_output_tokens=0,
93
- truncate_prompt_tokens=self.truncate_prompt_tokens,
94
- truncation_side=self.truncation_side,
95
- do_lower_case=encoder_config.get("do_lower_case", False),
96
- add_special_tokens=not model_config.is_encoder_decoder,
97
- max_total_tokens_param="max_model_len",
73
)
74
75
def to_pooling_params(self):
vllm/entrypoints/pooling/scoring/protocol.py
+3
-7
@@ -38,15 +38,11 @@ class ScoringRequestMixin(PoolingBasicRequestMixin, ClassifyRequestMixin):
38
# --8<-- [end:scoring-common-params]
39
40
def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
41
- encoder_config = model_config.encoder_config or {}
42
-
43
- return TokenizeParams(
41
+ return self._build_pooling_tok_params(
42
+ model_config,
43
+ add_special_tokens=True,
44
max_total_tokens=model_config.max_model_len,
45
max_output_tokens=0,
46
- truncate_prompt_tokens=self.truncate_prompt_tokens,
47
- truncation_side=self.truncation_side,
48
- do_lower_case=encoder_config.get("do_lower_case", False),
49
- max_total_tokens_param="max_model_len",
46
)
47
48
def to_pooling_params(self, task: PoolingTask = "classify"):