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, 13:30 UTC
[Frontend] Using matryoshka_dimensions control the allowed output dimensions. (#16970)
wang.yuqi committed
Apr 24, 2025 at 22:06 UTC
67309a1cb55f865f745b36501707f9f73bb023cc
8 files changed
+177
-81
docs/source/models/pooling_models.md
+5
-5
@@ -159,14 +159,14 @@ For example, setting `dimensions` parameter while using the `BAAI/bge-m3` model
159
160
### Manually enable Matryoshka Embeddings
161
162
-There is currently no official interface for specifying support for Matryoshka Embeddings. In vLLM, we simply check the existence of the fields `is_matryoshka` or `matryoshka_dimensions` inside `config.json`.
162
+There is currently no official interface for specifying support for Matryoshka Embeddings. In vLLM, if `is_matryoshka` is `True` in `config.json,` it is allowed to change the output to arbitrary dimensions. Using `matryoshka_dimensions` can control the allowed output dimensions.
163
164
-For models that support Matryoshka Embeddings but not recognized by vLLM, please manually override the config using `hf_overrides={"is_matryoshka": True}` (offline) or `--hf_overrides '{"is_matryoshka": true}'` (online).
164
+For models that support Matryoshka Embeddings but not recognized by vLLM, please manually override the config using `hf_overrides={"is_matryoshka": True}`, `hf_overrides={"matryoshka_dimensions": [<allowed output dimensions>]}` (offline) or `--hf_overrides '{"is_matryoshka": true}'`, `--hf_overrides '{"matryoshka_dimensions": [<allowed output dimensions>]}'`(online).
165
166
Here is an example to serve a model with Matryoshka Embeddings enabled.
167
168
```text
169
-vllm serve Snowflake/snowflake-arctic-embed-m-v1.5 --hf_overrides '{"is_matryoshka":true}'
169
+vllm serve Snowflake/snowflake-arctic-embed-m-v1.5 --hf_overrides '{"matryoshka_dimensions":[256]}'
170
```
171
172
### Offline Inference
@@ -204,14 +204,14 @@ curl http://127.0.0.1:8000/v1/embeddings \
204
"input": "Follow the white rabbit.",
205
"model": "jinaai/jina-embeddings-v3",
206
"encoding_format": "float",
207
- "dimensions": 1
207
+ "dimensions": 32
208
}'
209
```
210
211
Expected output:
212
213
```json
214
-{"id":"embd-0aab28c384d348c3b8f0eb783109dc5f","object":"list","created":1744195454,"model":"jinaai/jina-embeddings-v3","data":[{"index":0,"object":"embedding","embedding":[-1.0]}],"usage":{"prompt_tokens":10,"total_tokens":10,"completion_tokens":0,"prompt_tokens_details":null}}
214
+{"id":"embd-5c21fc9a5c9d4384a1b021daccaf9f64","object":"list","created":1745476417,"model":"jinaai/jina-embeddings-v3","data":[{"index":0,"object":"embedding","embedding":[-0.3828125,-0.1357421875,0.03759765625,0.125,0.21875,0.09521484375,-0.003662109375,0.1591796875,-0.130859375,-0.0869140625,-0.1982421875,0.1689453125,-0.220703125,0.1728515625,-0.2275390625,-0.0712890625,-0.162109375,-0.283203125,-0.055419921875,-0.0693359375,0.031982421875,-0.04052734375,-0.2734375,0.1826171875,-0.091796875,0.220703125,0.37890625,-0.0888671875,-0.12890625,-0.021484375,-0.0091552734375,0.23046875]}],"usage":{"prompt_tokens":8,"total_tokens":8,"completion_tokens":0,"prompt_tokens_details":null}}
215
```
216
217
A openai client example can be found here: <gh-file:examples/online_serving/openai_embedding_matryoshka_fy.py>
examples/online_serving/openai_embedding_matryoshka_fy.py
+2
-2
@@ -25,11 +25,11 @@ def main():
25
responses = client.embeddings.create(
26
input=["Follow the white rabbit."],
27
model=model,
28
- dimensions=1,
28
+ dimensions=32,
29
)
30
31
for data in responses.data:
32
- print(data.embedding) # List of float of len 1
32
+ print(data.embedding) # List of float of len 32
33
34
35
if __name__ == "__main__":
tests/entrypoints/openai/test_embedding.py
+24
-18
@@ -11,11 +11,12 @@ import requests
11
from vllm.entrypoints.openai.protocol import EmbeddingResponse
12
from vllm.transformers_utils.tokenizer import get_tokenizer
13
14
-from ...models.embedding.utils import check_embeddings_close
14
+from ...models.embedding.utils import correctness_test
15
from ...utils import RemoteOpenAIServer
16
17
MODEL_NAME = "intfloat/multilingual-e5-small"
18
DUMMY_CHAT_TEMPLATE = """{% for message in messages %}{{message['role'] + ': ' + message['content'] + '\\n'}}{% endfor %}""" # noqa: E501
19
+DTYPE = "bfloat16"
20
21
22
@pytest.fixture(scope="module")
@@ -25,7 +26,7 @@ def server():
26
"embed",
27
# use half precision for speed and memory savings in CI environment
28
"--dtype",
28
- "bfloat16",
29
+ DTYPE,
30
"--enforce-eager",
31
"--max-model-len",
32
"512",
@@ -43,9 +44,17 @@ async def client(server):
44
yield async_client
45
46
47
+@pytest.fixture(scope="module")
48
+def hf_model(hf_runner):
49
+ with hf_runner(MODEL_NAME, dtype=DTYPE,
50
+ is_sentence_transformer=True) as hf_model:
51
+ yield hf_model
52
+
53
+
54
@pytest.mark.asyncio
55
@pytest.mark.parametrize("model_name", [MODEL_NAME])
48
-async def test_single_embedding(client: openai.AsyncOpenAI, model_name: str):
56
+async def test_single_embedding(hf_model, client: openai.AsyncOpenAI,
57
+ model_name: str):
58
input_texts = [
59
"The chef prepared a delicious meal.",
60
]
@@ -66,6 +75,9 @@ async def test_single_embedding(client: openai.AsyncOpenAI, model_name: str):
75
assert embeddings.usage.prompt_tokens == 11
76
assert embeddings.usage.total_tokens == 11
77
78
+ vllm_outputs = [d.embedding for d in embeddings.data]
79
+ correctness_test(hf_model, input_texts, vllm_outputs)
80
+
81
# test using token IDs
82
input_tokens = [1, 1, 1, 1, 1]
83
embedding_response = await client.embeddings.create(
@@ -86,7 +98,8 @@ async def test_single_embedding(client: openai.AsyncOpenAI, model_name: str):
98
99
@pytest.mark.asyncio
100
@pytest.mark.parametrize("model_name", [MODEL_NAME])
89
-async def test_batch_embedding(client: openai.AsyncOpenAI, model_name: str):
101
+async def test_batch_embedding(hf_model, client: openai.AsyncOpenAI,
102
+ model_name: str):
103
# test list[str]
104
input_texts = [
105
"The cat sat on the mat.", "A feline was resting on a rug.",
@@ -107,6 +120,9 @@ async def test_batch_embedding(client: openai.AsyncOpenAI, model_name: str):
120
assert embeddings.usage.prompt_tokens == 33
121
assert embeddings.usage.total_tokens == 33
122
123
+ vllm_outputs = [d.embedding for d in embeddings.data]
124
+ correctness_test(hf_model, input_texts, vllm_outputs)
125
+
126
# test list[list[int]]
127
input_tokens = [[4, 5, 7, 9, 20], [15, 29, 499], [24, 24, 24, 24, 24],
128
[25, 32, 64, 77]]
@@ -181,7 +197,7 @@ async def test_conversation_embedding(server: RemoteOpenAIServer,
197
198
@pytest.mark.asyncio
199
@pytest.mark.parametrize("model_name", [MODEL_NAME])
184
-async def test_batch_base64_embedding(client: openai.AsyncOpenAI,
200
+async def test_batch_base64_embedding(hf_model, client: openai.AsyncOpenAI,
201
model_name: str):
202
input_texts = [
203
"Hello my name is",
@@ -192,6 +208,7 @@ async def test_batch_base64_embedding(client: openai.AsyncOpenAI,
208
model=model_name,
209
encoding_format="float")
210
float_data = [d.embedding for d in responses_float.data]
211
+ correctness_test(hf_model, input_texts, float_data)
212
213
responses_base64 = await client.embeddings.create(input=input_texts,
214
model=model_name,
@@ -202,24 +219,13 @@ async def test_batch_base64_embedding(client: openai.AsyncOpenAI,
219
np.frombuffer(base64.b64decode(data.embedding),
220
dtype="float32").tolist())
221
205
- check_embeddings_close(
206
- embeddings_0_lst=float_data,
207
- embeddings_1_lst=base64_data,
208
- name_0="float",
209
- name_1="base64",
210
- )
222
+ correctness_test(hf_model, input_texts, base64_data)
223
224
# Default response is float32 decoded from base64 by OpenAI Client
225
responses_default = await client.embeddings.create(input=input_texts,
226
model=model_name)
227
default_data = [d.embedding for d in responses_default.data]
216
-
217
- check_embeddings_close(
218
- embeddings_0_lst=float_data,
219
- embeddings_1_lst=default_data,
220
- name_0="float",
221
- name_1="default",
222
- )
228
+ correctness_test(hf_model, input_texts, default_data)
229
230
231
@pytest.mark.asyncio
tests/entrypoints/openai/test_embedding_dimensions.py
+91
-43
@@ -3,73 +3,121 @@
3
Run `pytest tests/entrypoints/openai/test_embedding_dimensions.py`.
4
"""
5
6
+from typing import Optional
7
+
8
import openai
9
import pytest
10
11
from vllm.entrypoints.openai.protocol import EmbeddingResponse
12
11
-from ...models.embedding.utils import EmbedModelInfo
13
+from ...conftest import HfRunner
14
+from ...models.embedding.utils import EmbedModelInfo, correctness_test
15
from ...utils import RemoteOpenAIServer
16
17
MODELS = [
15
- EmbedModelInfo(name="BAAI/bge-m3", is_matryoshka=False),
16
- EmbedModelInfo(name="jinaai/jina-embeddings-v3", is_matryoshka=True),
18
+ EmbedModelInfo("intfloat/multilingual-e5-small", is_matryoshka=False),
19
+ EmbedModelInfo("Snowflake/snowflake-arctic-embed-m-v1.5",
20
+ is_matryoshka=True,
21
+ matryoshka_dimensions=[256]),
22
]
23
24
input_texts = [
25
"The chef prepared a delicious meal.",
21
-] * 3
26
+]
27
28
24
-@pytest.mark.asyncio
25
-@pytest.mark.parametrize("model", MODELS)
26
-async def test_validating_dimensions(model: EmbedModelInfo):
29
+@pytest.fixture(scope="module", params=MODELS)
30
+def model_info(request):
31
+ return request.param
32
+
33
+
34
+@pytest.fixture(scope="module", params=["bfloat16"])
35
+def dtype(request):
36
+ return request.param
37
+
38
+
39
+@pytest.fixture(scope="module")
40
+def server(model_info, dtype: str):
41
args = [
42
"--task",
43
"embed",
44
# use half precision for speed and memory savings in CI environment
45
"--dtype",
32
- "bfloat16",
46
+ dtype,
47
"--enforce-eager",
48
"--max-model-len",
35
- "512",
36
- "--trust_remote_code"
49
+ "512"
50
]
38
- with RemoteOpenAIServer(model.name, args) as remote_server:
39
- client = remote_server.get_async_client()
40
-
41
- async def make_request(dimensions):
42
- embedding_response = await client.embeddings.create(
43
- model=model.name,
44
- input=input_texts,
45
- dimensions=dimensions,
46
- encoding_format="float",
47
- )
48
- embeddings = EmbeddingResponse.model_validate(
49
- embedding_response.model_dump(mode="json"))
50
-
51
- assert embeddings.id is not None
52
- assert len(embeddings.data) == 3
53
- assert len(embeddings.data[0].embedding) > 0
54
- assert embeddings.usage.completion_tokens == 0
55
- assert embeddings.usage.prompt_tokens > 0
56
- assert embeddings.usage.total_tokens > 0
57
-
58
- if dimensions is not None:
59
- assert len(embeddings.data[0].embedding) == dimensions
60
-
61
- if model.is_matryoshka:
62
- for dimensions in [None, 16]:
63
- await make_request(dimensions)
51
52
+ if model_info.name == "Snowflake/snowflake-arctic-embed-m-v1.5":
53
+ # Manually enable Matryoshka Embeddings
54
+ args.extend([
55
+ "--trust_remote_code", "--hf_overrides",
56
+ '{"matryoshka_dimensions":[256]}'
57
+ ])
58
+
59
+ with RemoteOpenAIServer(model_info.name, args) as remote_server:
60
+ yield remote_server
61
+
62
+
63
+@pytest.fixture(scope="module")
64
+def hf_model(hf_runner, model_info, dtype: str):
65
+ with hf_runner(model_info.name, dtype=dtype,
66
+ is_sentence_transformer=True) as hf_model:
67
+ yield hf_model
68
+
69
+
70
+@pytest.mark.asyncio
71
+async def test_matryoshka(model_info: EmbedModelInfo,
72
+ server: RemoteOpenAIServer, hf_model: HfRunner):
73
+ client = server.get_async_client()
74
+
75
+ async def make_request_and_correctness_test(dimensions):
76
+ prompts = input_texts * 3
77
+
78
+ embedding_response = await client.embeddings.create(
79
+ model=model_info.name,
80
+ input=prompts,
81
+ dimensions=dimensions,
82
+ encoding_format="float",
83
+ )
84
+ embeddings = EmbeddingResponse.model_validate(
85
+ embedding_response.model_dump(mode="json"))
86
+
87
+ assert embeddings.id is not None
88
+ assert len(embeddings.data) == 3
89
+ assert len(embeddings.data[0].embedding) > 0
90
+ assert embeddings.usage.completion_tokens == 0
91
+ assert embeddings.usage.prompt_tokens > 0
92
+ assert embeddings.usage.total_tokens > 0
93
+
94
+ if dimensions is not None:
95
+ assert len(embeddings.data[0].embedding) == dimensions
96
+
97
+ vllm_outputs = [d.embedding for d in embeddings.data]
98
+ correctness_test(hf_model, prompts, vllm_outputs, dimensions)
99
+
100
+ if model_info.is_matryoshka:
101
+ valid_dimensions: list[Optional[int]] = [None]
102
+ if model_info.matryoshka_dimensions is not None:
103
+ valid_dimensions += model_info.matryoshka_dimensions[:2]
104
+
105
+ for dimensions in valid_dimensions:
106
+ await make_request_and_correctness_test(dimensions)
107
+
108
+ invalid_dimensions: list[Optional[int]] = [-1]
109
+ if model_info.matryoshka_dimensions is not None:
110
+ assert 5 not in model_info.matryoshka_dimensions
111
+ invalid_dimensions.append(5)
112
+
113
+ for dimensions in invalid_dimensions:
114
with pytest.raises(openai.BadRequestError):
66
- for dimensions in [-1]:
67
- await make_request(dimensions)
115
+ await make_request_and_correctness_test(dimensions)
116
69
- else:
70
- for dimensions in [None]:
71
- await make_request(dimensions)
117
+ else:
118
+ for dimensions in [None]:
119
+ await make_request_and_correctness_test(dimensions)
120
121
+ for dimensions in [-1, 16]:
122
with pytest.raises(openai.BadRequestError):
74
- for dimensions in [-1, 16]:
75
- await make_request(dimensions)
123
+ await make_request_and_correctness_test(dimensions)
tests/models/embedding/language/test_jina.py
+21
-11
@@ -153,14 +153,24 @@ def test_matryoshka(
153
154
with vllm_runner(model, task="embed", dtype=dtype,
155
max_model_len=None) as vllm_model:
156
- vllm_outputs = vllm_model.encode(
157
- example_prompts,
158
- pooling_params=PoolingParams(dimensions=dimensions))
159
-
160
- check_embeddings_close(
161
- embeddings_0_lst=hf_outputs,
162
- embeddings_1_lst=vllm_outputs,
163
- name_0="hf",
164
- name_1="vllm",
165
- tol=1e-2,
166
- )
156
+ matryoshka_dimensions = (
157
+ vllm_model.model.llm_engine.model_config.matryoshka_dimensions)
158
+ assert matryoshka_dimensions is not None
159
+
160
+ if dimensions not in matryoshka_dimensions:
161
+ with pytest.raises(ValueError):
162
+ vllm_model.encode(
163
+ example_prompts,
164
+ pooling_params=PoolingParams(dimensions=dimensions))
165
+ else:
166
+ vllm_outputs = vllm_model.encode(
167
+ example_prompts,
168
+ pooling_params=PoolingParams(dimensions=dimensions))
169
+
170
+ check_embeddings_close(
171
+ embeddings_0_lst=hf_outputs,
172
+ embeddings_1_lst=vllm_outputs,
173
+ name_0="hf",
174
+ name_1="vllm",
175
+ tol=1e-2,
176
+ )
tests/models/embedding/utils.py
+20
-1
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
from collections.abc import Sequence
4
-from typing import NamedTuple
4
+from typing import NamedTuple, Optional
5
6
import torch
7
import torch.nn.functional as F
@@ -43,5 +43,24 @@ def matryoshka_fy(tensor, dimensions):
43
class EmbedModelInfo(NamedTuple):
44
name: str
45
is_matryoshka: bool
46
+ matryoshka_dimensions: Optional[list[int]] = None
47
architecture: str = ""
48
enable_test: bool = True
49
+
50
+
51
+def correctness_test(hf_model,
52
+ inputs,
53
+ vllm_outputs: Sequence[list[float]],
54
+ dimensions: Optional[int] = None):
55
+
56
+ hf_outputs = hf_model.encode(inputs)
57
+ if dimensions:
58
+ hf_outputs = matryoshka_fy(hf_outputs, dimensions)
59
+
60
+ check_embeddings_close(
61
+ embeddings_0_lst=hf_outputs,
62
+ embeddings_1_lst=vllm_outputs,
63
+ name_0="hf",
64
+ name_1="vllm",
65
+ tol=1e-2,
66
+ )
vllm/config.py
+4
@@ -1248,6 +1248,10 @@ class ModelConfig:
1248
return (hasattr(self.hf_config, "matryoshka_dimensions")
1249
or getattr(self.hf_config, "is_matryoshka", False))
1250
1251
+ @property
1252
+ def matryoshka_dimensions(self):
1253
+ return getattr(self.hf_config, "matryoshka_dimensions", None)
1254
+
1255
1256
BlockSize = Literal[1, 8, 16, 32, 64, 128]
1257
CacheDType = Literal["auto", "fp8", "fp8_e4m3", "fp8_e5m2"]
vllm/pooling_params.py
+10
-1
@@ -35,7 +35,16 @@ class PoolingParams(
35
f'Model "{model_config.served_model_name}" does not '
36
f'support matryoshka representation, '
37
f'changing output dimensions will lead to poor results.')
38
- if self.dimensions < 1:
38
+
39
+ mds = model_config.matryoshka_dimensions
40
+ if mds is not None:
41
+ if self.dimensions not in mds:
42
+ raise ValueError(
43
+ f'Model "{model_config.served_model_name}" '
44
+ f'only supports {str(mds)} matryoshka dimensions, '
45
+ f'use other output dimensions will '
46
+ f'lead to poor results.')
47
+ elif self.dimensions < 1:
48
raise ValueError("Dimensions must be greater than 0")
49
50
def __repr__(self) -> str: