@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 232214b2ae
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 10, 06:30 UTC

[Bugfix] Replace `PoolingParams.normalize` with `use_activation` (#32243)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Cyrus Leung committed Jan 13, 2026 at 18:45 UTC 232214b2aece00201a2355ed7368452f761b76db
21 files changed +68 -70
examples/pooling/embed/openai_embedding_long_text/README.md
+1 -1
@@ -47,7 +47,7 @@ The key parameters for chunked processing are in the `--pooler-config`:
47 ```json
48 {
49 "pooling_type": "auto",
50 - "normalize": true,
50 + "use_activation": true,
51 "enable_chunked_processing": true,
52 "max_embed_len": 3072000
53 }
examples/pooling/embed/openai_embedding_long_text/client.py
+2 -2
@@ -14,7 +14,7 @@ Prerequisites:
14 # MEAN pooling (processes all chunks, recommended for complete coverage)
15 vllm serve intfloat/multilingual-e5-large \
16 --pooler-config \
17 - '{"pooling_type": "MEAN", "normalize": true, ' \
17 + '{"pooling_type": "MEAN", "use_activation": true, ' \
18 '"enable_chunked_processing": true, "max_embed_len": 3072000}' \
19 --served-model-name multilingual-e5-large \
20 --trust-remote-code \
@@ -24,7 +24,7 @@ Prerequisites:
24 # OR CLS pooling (native CLS within chunks, MEAN aggregation across chunks)
25 vllm serve BAAI/bge-large-en-v1.5 \
26 --pooler-config \
27 - '{"pooling_type": "CLS", "normalize": true, ' \
27 + '{"pooling_type": "CLS", "use_activation": true, ' \
28 '"enable_chunked_processing": true, "max_embed_len": 1048576}' \
29 --served-model-name bge-large-en-v1.5 \
30 --trust-remote-code \
examples/pooling/embed/openai_embedding_long_text/service.sh
+1 -1
@@ -96,7 +96,7 @@ echo ""
96 echo "🔧 Starting server with enhanced chunked processing configuration..."
97
98 # Build pooler config JSON
99 -POOLER_CONFIG="{\"pooling_type\": \"$POOLING_TYPE\", \"normalize\": true, \"enable_chunked_processing\": ${VLLM_ENABLE_CHUNKED_PROCESSING}, \"max_embed_len\": ${MAX_EMBED_LEN}}"
99 +POOLER_CONFIG="{\"pooling_type\": \"$POOLING_TYPE\", \"use_activation\": true, \"enable_chunked_processing\": ${VLLM_ENABLE_CHUNKED_PROCESSING}, \"max_embed_len\": ${MAX_EMBED_LEN}}"
100
101 # Start vLLM server with enhanced chunked processing
102 vllm serve "$MODEL_NAME" \
tests/entrypoints/pooling/embed/test_offline.py
+3 -1
@@ -53,7 +53,9 @@ def test_token_embed(llm: LLM):
53 def test_pooling_params(llm: LLM):
54 def get_outputs(normalize):
55 outputs = llm.embed(
56 - prompts, pooling_params=PoolingParams(normalize=normalize), use_tqdm=False
56 + prompts,
57 + pooling_params=PoolingParams(use_activation=normalize),
58 + use_tqdm=False,
59 )
60 return torch.tensor([x.outputs.embedding for x in outputs])
61
tests/entrypoints/pooling/embed/test_online_long_text.py
+1 -1
@@ -216,7 +216,7 @@ def server_with_chunked_processing():
216 "512", # Set smaller max_model_len to trigger chunking mechanism
217 "--pooler-config",
218 (
219 - '{"pooling_type": "MEAN", "normalize": true, '
219 + '{"pooling_type": "MEAN", "use_activation": true, '
220 '"enable_chunked_processing": true, "max_embed_len": 10000}'
221 ),
222 "--gpu-memory-utilization",
tests/entrypoints/pooling/score/test_online_score.py
+4 -13
@@ -236,17 +236,14 @@ class TestModel:
236 "use_activation": use_activation,
237 },
238 )
239 - if response.status_code != 200:
240 - return response
241 -
239 outputs = response.json()
240 return torch.tensor([x["score"] for x in outputs["data"]])
241
245 - if model["is_cross_encoder"]:
246 - default = get_outputs(use_activation=None)
247 - w_activation = get_outputs(use_activation=True)
248 - wo_activation = get_outputs(use_activation=False)
242 + default = get_outputs(use_activation=None)
243 + w_activation = get_outputs(use_activation=True)
244 + wo_activation = get_outputs(use_activation=False)
245
246 + if model["is_cross_encoder"]:
247 assert torch.allclose(default, w_activation, atol=1e-2), (
248 "Default should use activation."
249 )
@@ -256,9 +253,3 @@ class TestModel:
253 assert torch.allclose(F.sigmoid(wo_activation), w_activation, atol=1e-2), (
254 "w_activation should be close to activation(wo_activation)."
255 )
259 - else:
260 - get_outputs(use_activation=None)
261 -
262 - # The activation parameter only works for the is_cross_encoder model
263 - response = get_outputs(use_activation=True)
264 - assert response.status_code == 400
tests/model_executor/test_model_load_with_params.py
+2 -2
@@ -48,7 +48,7 @@ def test_model_loading_with_params(vllm_runner, monkeypatch):
48 # asserts on the pooling config files
49 assert model_config.pooler_config.seq_pooling_type == "CLS"
50 assert model_config.pooler_config.tok_pooling_type == "ALL"
51 - assert model_config.pooler_config.normalize
51 + assert model_config.pooler_config.use_activation
52
53 # asserts on the tokenizer loaded
54 assert model_config.tokenizer == "BAAI/bge-base-en-v1.5"
@@ -93,7 +93,7 @@ def test_roberta_model_loading_with_params(vllm_runner, monkeypatch):
93 # asserts on the pooling config files
94 assert model_config.pooler_config.seq_pooling_type == "MEAN"
95 assert model_config.pooler_config.tok_pooling_type == "ALL"
96 - assert model_config.pooler_config.normalize
96 + assert model_config.pooler_config.use_activation
97
98 # asserts on the tokenizer loaded
99 assert model_config.tokenizer == "intfloat/multilingual-e5-base"
tests/models/language/pooling/test_pooler_config_init_behaviour.py
+4 -4
@@ -66,7 +66,7 @@ def test_embed_models_using_normalize(
66 model,
67 max_model_len=512,
68 dtype=dtype,
69 - pooler_config=PoolerConfig(normalize=False),
69 + pooler_config=PoolerConfig(use_activation=False),
70 ) as vllm_model:
71 wo_normalize = torch.tensor(vllm_model.embed(example_prompts))
72
@@ -74,7 +74,7 @@ def test_embed_models_using_normalize(
74 model,
75 max_model_len=512,
76 dtype=dtype,
77 - pooler_config=PoolerConfig(normalize=True),
77 + pooler_config=PoolerConfig(use_activation=True),
78 ) as vllm_model:
79 w_normalize = torch.tensor(vllm_model.embed(example_prompts))
80
@@ -146,7 +146,7 @@ def test_multi_vector_retrieval_models_using_normalize(
146 model,
147 max_model_len=512,
148 dtype=dtype,
149 - pooler_config=PoolerConfig(normalize=False),
149 + pooler_config=PoolerConfig(use_activation=False),
150 ) as vllm_model:
151 wo_normalize = vllm_model.token_embed(example_prompts)
152
@@ -154,7 +154,7 @@ def test_multi_vector_retrieval_models_using_normalize(
154 model,
155 max_model_len=512,
156 dtype=dtype,
157 - pooler_config=PoolerConfig(normalize=True),
157 + pooler_config=PoolerConfig(use_activation=True),
158 ) as vllm_model:
159 w_normalize = vllm_model.token_embed(example_prompts)
160
tests/test_config.py
+1 -1
@@ -162,7 +162,7 @@ def test_get_pooling_config():
162 model_config = ModelConfig(model_id)
163
164 assert model_config.pooler_config is not None
165 - assert model_config.pooler_config.normalize
165 + assert model_config.pooler_config.use_activation
166 assert model_config.pooler_config.seq_pooling_type == "MEAN"
167 assert model_config.pooler_config.tok_pooling_type == "ALL"
168
tests/test_pooling_params.py
+11 -11
@@ -18,7 +18,7 @@ EMBEDDING_MODELS = [
18 ]
19
20 classify_parameters = ["use_activation"]
21 -embed_parameters = ["dimensions", "normalize"]
21 +embed_parameters = ["dimensions", "use_activation"]
22 step_pooling_parameters = ["step_tag_id", "returned_token_ids"]
23
24
@@ -42,17 +42,17 @@ def test_embed():
42 task = "embed"
43 model_config = MockModelConfig(pooler_config=PoolerConfig(seq_pooling_type="CLS"))
44
45 - pooling_params = PoolingParams(normalize=None)
45 + pooling_params = PoolingParams(use_activation=None)
46 pooling_params.verify(task=task, model_config=model_config)
47
48 - pooling_params = PoolingParams(normalize=True)
48 + pooling_params = PoolingParams(use_activation=True)
49 pooling_params.verify(task=task, model_config=model_config)
50
51 - pooling_params = PoolingParams(normalize=False)
51 + pooling_params = PoolingParams(use_activation=False)
52 pooling_params.verify(task=task, model_config=model_config)
53
54 invalid_parameters = classify_parameters + step_pooling_parameters
55 - for p in invalid_parameters:
55 + for p in set(invalid_parameters) - set(embed_parameters):
56 with pytest.raises(ValueError):
57 pooling_params = PoolingParams(**{p: True})
58 pooling_params.verify(task=task, model_config=model_config)
@@ -98,7 +98,7 @@ def test_classify(task):
98 pooling_params.verify(task=task, model_config=model_config)
99
100 invalid_parameters = embed_parameters + step_pooling_parameters
101 - for p in invalid_parameters:
101 + for p in set(invalid_parameters) - set(classify_parameters):
102 with pytest.raises(ValueError):
103 pooling_params = PoolingParams(**{p: True})
104 pooling_params.verify(task=task, model_config=model_config)
@@ -111,20 +111,20 @@ def test_token_embed(pooling_type: str):
111 pooler_config=PoolerConfig(tok_pooling_type=pooling_type)
112 )
113
114 - pooling_params = PoolingParams(normalize=None)
114 + pooling_params = PoolingParams(use_activation=None)
115 pooling_params.verify(task=task, model_config=model_config)
116
117 - pooling_params = PoolingParams(normalize=True)
117 + pooling_params = PoolingParams(use_activation=True)
118 pooling_params.verify(task=task, model_config=model_config)
119
120 - pooling_params = PoolingParams(normalize=False)
120 + pooling_params = PoolingParams(use_activation=False)
121 pooling_params.verify(task=task, model_config=model_config)
122
123 invalid_parameters = classify_parameters
124 if pooling_type != "STEP":
125 invalid_parameters = classify_parameters + step_pooling_parameters
126
127 - for p in invalid_parameters:
127 + for p in set(invalid_parameters) - set(embed_parameters):
128 with pytest.raises(ValueError):
129 pooling_params = PoolingParams(**{p: True})
130 pooling_params.verify(task=task, model_config=model_config)
@@ -150,7 +150,7 @@ def test_token_classify(pooling_type: str):
150 if pooling_type != "STEP":
151 invalid_parameters = embed_parameters + step_pooling_parameters
152
153 - for p in invalid_parameters:
153 + for p in set(invalid_parameters) - set(classify_parameters):
154 with pytest.raises(ValueError):
155 pooling_params = PoolingParams(**{p: True})
156 pooling_params.verify(task=task, model_config=model_config)
vllm/config/pooler.py
+16 -9
@@ -48,7 +48,7 @@ class PoolerConfig:
48 ## for embeddings models
49 normalize: bool | None = None
50 """
51 - Whether to normalize the embeddings outputs. Defaults to True.
51 + DEPRECATED: please use `use_activation` instead.
52 """
53 dimensions: int | None = None
54 """
@@ -75,11 +75,11 @@ class PoolerConfig:
75 ## for classification models
76 softmax: float | None = None
77 """
78 - softmax will be deprecated, please use use_activation instead.
78 + DEPRECATED: please use `use_activation` instead.
79 """
80 activation: float | None = None
81 """
82 - activation will be deprecated, please use use_activation instead.
82 + DEPRECATED: please use `use_activation` instead.
83 """
84 use_activation: bool | None = None
85 """
@@ -164,17 +164,24 @@ class PoolerConfig:
164
165
166 def get_use_activation(o: object):
167 - if softmax := getattr(o, "softmax", None) is not None:
167 + if (normalize := getattr(o, "normalize", None)) is not None:
168 logger.warning_once(
169 - "softmax will be deprecated and will be removed in v0.15. "
170 - "Please use use_activation instead."
169 + "`normalize` is deprecated and will be removed in v0.15. "
170 + "Please use `use_activation` instead."
171 + )
172 + return normalize
173 +
174 + if (softmax := getattr(o, "softmax", None)) is not None:
175 + logger.warning_once(
176 + "`softmax` is deprecated and will be removed in v0.15. "
177 + "Please use `use_activation` instead."
178 )
179 return softmax
180
174 - if activation := getattr(o, "activation", None) is not None:
181 + if (activation := getattr(o, "activation", None)) is not None:
182 logger.warning_once(
176 - "activation will be deprecated and will be removed in v0.15. "
177 - "Please use use_activation instead."
183 + "`activation` is deprecated and will be removed in v0.15. "
184 + "Please use `use_activation` instead."
185 )
186 return activation
187
vllm/entrypoints/pooling/embed/protocol.py
+2 -2
@@ -75,7 +75,7 @@ class EmbeddingCompletionRequest(OpenAIBaseModel):
75 return PoolingParams(
76 truncate_prompt_tokens=self.truncate_prompt_tokens,
77 dimensions=self.dimensions,
78 - normalize=self.normalize,
78 + use_activation=self.normalize,
79 )
80
81
@@ -189,7 +189,7 @@ class EmbeddingChatRequest(OpenAIBaseModel):
189 return PoolingParams(
190 truncate_prompt_tokens=self.truncate_prompt_tokens,
191 dimensions=self.dimensions,
192 - normalize=self.normalize,
192 + use_activation=self.normalize,
193 )
194
195
vllm/entrypoints/pooling/pooling/protocol.py
-2
@@ -40,7 +40,6 @@ class PoolingCompletionRequest(EmbeddingCompletionRequest):
40 return PoolingParams(
41 truncate_prompt_tokens=self.truncate_prompt_tokens,
42 dimensions=self.dimensions,
43 - normalize=self.normalize,
43 use_activation=get_use_activation(self),
44 )
45
@@ -66,7 +65,6 @@ class PoolingChatRequest(EmbeddingChatRequest):
65 return PoolingParams(
66 truncate_prompt_tokens=self.truncate_prompt_tokens,
67 dimensions=self.dimensions,
69 - normalize=self.normalize,
68 use_activation=get_use_activation(self),
69 )
70
vllm/model_executor/layers/pooler/seqwise/heads.py
+1 -1
@@ -83,7 +83,7 @@ class EmbeddingPoolerHead(SequencePoolerHead):
83
84 # for normalize
85 if self.activation is not None:
86 - flags = [p.normalize for p in pooling_params]
86 + flags = [p.use_activation for p in pooling_params]
87 if len(set(flags)) == 1:
88 if flags[0]:
89 pooled_data = self.activation(pooled_data)
vllm/model_executor/layers/pooler/seqwise/poolers.py
+2 -2
@@ -95,8 +95,8 @@ def pooler_for_embed(pooler_config: PoolerConfig):
95 vllm_config = get_current_vllm_config()
96 model_config = vllm_config.model_config
97 head = EmbeddingPoolerHead(
98 - projector=_load_st_projector(model_config),
98 head_dtype=model_config.head_dtype,
99 + projector=_load_st_projector(model_config),
100 activation=PoolerNormalize(),
101 )
102
@@ -116,9 +116,9 @@ def pooler_for_classify(
116 vllm_config = get_current_vllm_config()
117 model_config = vllm_config.model_config
118 head = ClassifierPoolerHead(
119 + head_dtype=model_config.head_dtype,
120 classifier=classifier,
121 logit_bias=model_config.pooler_config.logit_bias,
121 - head_dtype=model_config.head_dtype,
122 activation=resolve_classifier_act_fn(
123 model_config, static_num_labels=True, act_fn=act_fn
124 ),
vllm/model_executor/layers/pooler/tokwise/heads.py
+3 -3
@@ -44,14 +44,14 @@ class TokenPoolerHead(nn.Module, ABC):
44 class TokenEmbeddingPoolerHead(TokenPoolerHead):
45 def __init__(
46 self,
47 - projector: ProjectorFn | None = None,
47 head_dtype: torch.dtype | str | None = None,
48 + projector: ProjectorFn | None = None,
49 activation: ActivationFn | None = None,
50 ) -> None:
51 super().__init__()
52
53 - self.projector = projector
53 self.head_dtype = head_dtype
54 + self.projector = projector
55 self.activation = activation
56
57 def get_supported_tasks(self) -> Set[PoolingTask]:
@@ -79,7 +79,7 @@ class TokenEmbeddingPoolerHead(TokenPoolerHead):
79 pooled_data = pooled_data[..., : pooling_param.dimensions]
80
81 # for normalize
82 - if self.activation is not None and pooling_param.normalize:
82 + if self.activation is not None and pooling_param.use_activation:
83 pooled_data = self.activation(pooled_data)
84
85 # pooled_data shape: [n_tokens, embedding_dimension]
vllm/model_executor/layers/pooler/tokwise/poolers.py
+2 -2
@@ -95,8 +95,8 @@ def pooler_for_token_embed(pooler_config: PoolerConfig):
95 vllm_config = get_current_vllm_config()
96 model_config = vllm_config.model_config
97 head = TokenEmbeddingPoolerHead(
98 - projector=_load_st_projector(model_config),
98 head_dtype=model_config.head_dtype,
99 + projector=_load_st_projector(model_config),
100 activation=PoolerNormalize(),
101 )
102
@@ -116,9 +116,9 @@ def pooler_for_token_classify(
116 vllm_config = get_current_vllm_config()
117 model_config = vllm_config.model_config
118 head = TokenClassifierPoolerHead(
119 + head_dtype=model_config.head_dtype,
120 classifier=classifier,
121 logit_bias=model_config.pooler_config.logit_bias,
121 - head_dtype=model_config.head_dtype,
122 activation=resolve_classifier_act_fn(
123 model_config, static_num_labels=False, act_fn=act_fn
124 ),
vllm/model_executor/models/bert.py
+1 -1
@@ -116,8 +116,8 @@ class BertPooler(SequencePooler):
116
117 # Use lambdas so that weights are not registered under `self.head`
118 self.head = EmbeddingPoolerHead(
119 - projector=lambda x: self.dense(x),
119 head_dtype=head_dtype,
120 + projector=lambda x: self.dense(x),
121 activation=LambdaPoolerActivation(self.act_fn),
122 )
123
vllm/model_executor/models/modernbert.py
+2 -1
@@ -309,12 +309,13 @@ class ModernBertPooler(SequencePooler):
309 config.hidden_size,
310 eps=config.norm_eps,
311 bias=config.norm_bias,
312 + dtype=head_dtype,
313 )
314
315 # Use lambdas so that weights are not registered under `self.head`
316 self.head = EmbeddingPoolerHead(
316 - projector=lambda x: self.dense(x),
317 head_dtype=head_dtype,
318 + projector=lambda x: self.dense(x),
319 activation=LambdaPoolerActivation(lambda x: self.norm(self.act(x))),
320 )
321
vllm/pooling_params.py
+8 -9
@@ -26,9 +26,9 @@ class PoolingParams(
26 Set to None to disable truncation.
27 dimensions: Reduce the dimensions of embeddings
28 if model support matryoshka representation.
29 - normalize: Whether to normalize the embeddings outputs.
30 - softmax: softmax will be deprecated, please use use_activation instead.
31 - activation: activation will be deprecated, please use use_activation instead.
29 + normalize: Deprecated, please use use_activation instead.
30 + softmax: Deprecated, please use use_activation instead.
31 + activation: Deprecated, please use use_activation instead.
32 use_activation: Whether to apply activation function to
33 the classification outputs.
34 """
@@ -63,15 +63,15 @@ class PoolingParams(
63
64 @property
65 def all_parameters(self) -> list[str]:
66 - return ["dimensions", "normalize", "use_activation"]
66 + return ["dimensions", "use_activation"]
67
68 @property
69 def valid_parameters(self):
70 return {
71 - "embed": ["dimensions", "normalize"],
71 + "embed": ["dimensions", "use_activation"],
72 "classify": ["use_activation"],
73 "score": ["use_activation"],
74 - "token_embed": ["dimensions", "normalize"],
74 + "token_embed": ["dimensions", "use_activation"],
75 "token_classify": ["use_activation"],
76 }
77
@@ -162,8 +162,8 @@ class PoolingParams(
162
163 def _set_default_parameters(self, model_config: Optional["ModelConfig"]):
164 if self.task in ["embed", "token_embed"]:
165 - if self.normalize is None:
166 - self.normalize = True
165 + if self.use_activation is None:
166 + self.use_activation = True
167
168 if self.dimensions is not None and model_config is not None:
169 if not model_config.is_matryoshka:
@@ -213,7 +213,6 @@ class PoolingParams(
213 return (
214 f"PoolingParams("
215 f"task={self.task}, "
216 - f"normalize={self.normalize}, "
216 f"dimensions={self.dimensions}, "
217 f"use_activation={self.use_activation}, "
218 f"step_tag_id={self.step_tag_id}, "
vllm/transformers_utils/config.py
+1 -1
@@ -801,7 +801,7 @@ def get_pooling_config(
801
802 logger.info("Found pooling configuration.")
803
804 - config: dict[str, Any] = {"normalize": normalize}
804 + config: dict[str, Any] = {"use_activation": normalize}
805 for key, val in pooling_dict.items():
806 if val is True:
807 pooling_type = parse_pooling_type(key)