[Misc] Improve Matryoshka pooling dimensions validation (#48057)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Taneem Ibrahim committed
Jul 13, 2026 at 12:57 UTC
62286308c9e30adfef3780c6fbefeca5cf8f36ae
1 file changed
+17
-16
vllm/pooling_params.py
+17
-16
@@ -164,28 +164,29 @@ class PoolingParams(
164
self.use_activation = True
165
166
if self.dimensions is not None:
167
+ dimensions = self.dimensions
168
+ model_name = model_config.served_model_name
169
+ embedding_size = model_config.embedding_size
170
+ valid_range = f"[1, {embedding_size}]"
171
+ dimensions_in_range = 1 <= dimensions <= embedding_size
172
if not model_config.is_matryoshka:
173
raise ValueError(
169
- f'Model "{model_config.served_model_name}" does not '
170
- f"support matryoshka representation, "
171
- f"changing output dimensions will lead to poor results."
174
+ f"Model {model_name!r} does not support Matryoshka "
175
+ f"embeddings; dimensions must be unset "
176
+ f"(received dimensions={dimensions})."
177
+ )
178
+
179
+ if not dimensions_in_range:
180
+ raise ValueError(
181
+ f"Model {model_name!r} only supports dimensions in "
182
+ f"range {valid_range}, got {dimensions}."
183
)
184
185
mds = model_config.matryoshka_dimensions
175
- if mds is not None:
176
- if self.dimensions not in mds:
177
- raise ValueError(
178
- f"Model {model_config.served_model_name!r} "
179
- f"only supports {str(mds)} matryoshka dimensions, "
180
- f"use other output dimensions will "
181
- f"lead to poor results."
182
- )
183
- elif self.dimensions < 1:
184
- raise ValueError("Dimensions must be greater than 0")
185
- elif self.dimensions > model_config.embedding_size:
186
+ if mds is not None and dimensions not in mds:
187
raise ValueError(
187
- "Dimensions must be less than or equal to the model's "
188
- f"embedding size ({model_config.embedding_size})"
188
+ f"Model {model_name!r} only supports Matryoshka "
189
+ f"dimensions {str(mds)}, got {dimensions}."
190
)
191
192
elif self.task in ["classify", "token_classify"]: