[Bugfix] Reject matryoshka embedding dimensions above hidden size (#46313)
Signed-off-by: EazyReal <8047065+EazyReal@users.noreply.github.com> Co-authored-by: EazyReal <8047065+EazyReal@users.noreply.github.com>
Maxwill Lin committed
Jun 22, 2026 at 03:16 UTC
78739e3bda0466ecbd63de1bda07d5ac09d88dca
2 files changed
+28
tests/test_pooling_params.py
+23
@@ -74,6 +74,29 @@ def test_embed_dimensions(model_info: EmbedModelInfo):
74
pooling_params.verify(model_config)
75
76
77
+@dataclass()
78
+class MockMatryoshkaModelConfig:
79
+ pooler_config: PoolerConfig
80
+ is_matryoshka: bool = True
81
+ matryoshka_dimensions: list[int] | None = None
82
+ served_model_name: str = "mock-matryoshka-model"
83
+ embedding_size: int = 32
84
+
85
+
86
+def test_embed_dimensions_matryoshka_without_list_upper_bound():
87
+ task = "embed"
88
+ model_config = MockMatryoshkaModelConfig(
89
+ pooler_config=PoolerConfig(seq_pooling_type="CLS"),
90
+ matryoshka_dimensions=None,
91
+ embedding_size=32,
92
+ )
93
+
94
+ PoolingParams(task=task, dimensions=16).verify(model_config)
95
+
96
+ with pytest.raises(ValueError):
97
+ PoolingParams(task=task, dimensions=64).verify(model_config)
98
+
99
+
100
@pytest.mark.parametrize("task", ["classify"])
101
def test_classify(task):
102
model_config = MockModelConfig(pooler_config=PoolerConfig(seq_pooling_type="CLS"))
vllm/pooling_params.py
+5
@@ -182,6 +182,11 @@ class PoolingParams(
182
)
183
elif self.dimensions < 1:
184
raise ValueError("Dimensions must be greater than 0")
185
+ elif self.dimensions > model_config.embedding_size:
186
+ raise ValueError(
187
+ "Dimensions must be less than or equal to the model's "
188
+ f"embedding size ({model_config.embedding_size})"
189
+ )
190
191
elif self.task in ["classify", "token_classify"]:
192
if self.use_activation is None: