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 11, 21:15 UTC
[Frontend] Binary embedding response does not return metadata by setting encoding_format to bytes_only. (#30249)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io> Signed-off-by: wang.yuqi <noooop@126.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
wang.yuqi committed
Dec 8, 2025 at 20:01 UTC
2e660c24349944ac78abcf2c35d17e3caf6cdd08
10 files changed
+230
-41
examples/pooling/embed/embedding_requests_bytes_client.py
+35
-2
@@ -16,6 +16,7 @@ from vllm.utils.serial_utils import (
16
EMBED_DTYPE_TO_TORCH_DTYPE,
17
ENDIANNESS,
18
MetadataItem,
19
+ build_metadata_items,
20
decode_pooling_output,
21
)
22
@@ -38,6 +39,11 @@ def parse_args():
39
def main(args):
40
api_url = f"http://{args.host}:{args.port}/v1/embeddings"
41
model_name = args.model
42
+ embedding_size = 0
43
+
44
+ input_texts = [
45
+ "The best thing about vLLM is that it supports many different models",
46
+ ] * 2
47
48
# The OpenAI client does not support the bytes encoding_format.
49
# The OpenAI client does not support the embed_dtype and endianness parameters.
@@ -45,7 +51,7 @@ def main(args):
51
for endianness in ENDIANNESS:
52
prompt = {
53
"model": model_name,
48
- "input": "vLLM is great!",
54
+ "input": input_texts,
55
"encoding_format": "bytes",
56
"embed_dtype": embed_dtype,
57
"endianness": endianness,
@@ -57,7 +63,34 @@ def main(args):
63
64
embedding = decode_pooling_output(items=items, body=body)
65
embedding = [x.to(torch.float32) for x in embedding]
60
- embedding = torch.cat(embedding)
66
+ embedding = torch.stack(embedding)
67
+ embedding_size = embedding.shape[-1]
68
+ print(embed_dtype, endianness, embedding.shape)
69
+
70
+ # The vllm server always sorts the returned embeddings in the order of input. So
71
+ # returning metadata is not necessary. You can set encoding_format to bytes_only
72
+ # to let the server not return metadata.
73
+ for embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE:
74
+ for endianness in ENDIANNESS:
75
+ prompt = {
76
+ "model": model_name,
77
+ "input": input_texts,
78
+ "encoding_format": "bytes_only",
79
+ "embed_dtype": embed_dtype,
80
+ "endianness": endianness,
81
+ }
82
+ response = post_http_request(prompt=prompt, api_url=api_url)
83
+ body = response.content
84
+
85
+ items = build_metadata_items(
86
+ embed_dtype=embed_dtype,
87
+ endianness=endianness,
88
+ shape=(embedding_size,),
89
+ n_request=len(input_texts),
90
+ )
91
+ embedding = decode_pooling_output(items=items, body=body)
92
+ embedding = [x.to(torch.float32) for x in embedding]
93
+ embedding = torch.stack(embedding)
94
print(embed_dtype, endianness, embedding.shape)
95
96
tests/entrypoints/pooling/embed/test_online.py
+50
@@ -24,6 +24,7 @@ from vllm.utils.serial_utils import (
24
ENDIANNESS,
25
MetadataItem,
26
binary2tensor,
27
+ build_metadata_items,
28
decode_pooling_output,
29
)
30
@@ -344,6 +345,55 @@ async def test_bytes_embed_dtype_and_endianness(
345
)
346
347
348
+@pytest.mark.asyncio
349
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
350
+async def test_bytes_only_embed_dtype_and_endianness(
351
+ server: RemoteOpenAIServer, client: openai.AsyncOpenAI, model_name: str
352
+):
353
+ input_texts = [
354
+ "The best thing about vLLM is that it supports many different models",
355
+ ] * 2
356
+
357
+ responses_float = await client.embeddings.create(
358
+ input=input_texts, model=model_name, encoding_format="float"
359
+ )
360
+ float_data = [d.embedding for d in responses_float.data]
361
+ embedding_size = len(float_data[0])
362
+
363
+ for embed_dtype in list(EMBED_DTYPE_TO_TORCH_DTYPE.keys()):
364
+ for endianness in ENDIANNESS:
365
+ responses_bytes = requests.post(
366
+ server.url_for("/v1/embeddings"),
367
+ json={
368
+ "model": model_name,
369
+ "input": input_texts,
370
+ "encoding_format": "bytes_only",
371
+ "embed_dtype": embed_dtype,
372
+ "endianness": endianness,
373
+ },
374
+ )
375
+
376
+ assert "metadata" not in responses_bytes.headers
377
+ body = responses_bytes.content
378
+ items = build_metadata_items(
379
+ embed_dtype=embed_dtype,
380
+ endianness=endianness,
381
+ shape=(embedding_size,),
382
+ n_request=len(input_texts),
383
+ )
384
+
385
+ bytes_data = decode_pooling_output(items=items, body=body)
386
+ bytes_data = [x.to(torch.float32).tolist() for x in bytes_data]
387
+
388
+ check_embeddings_close(
389
+ embeddings_0_lst=float_data,
390
+ embeddings_1_lst=bytes_data,
391
+ name_0="float_data",
392
+ name_1="bytes_data",
393
+ tol=1e-2,
394
+ )
395
+
396
+
397
@pytest.mark.asyncio
398
@pytest.mark.parametrize("model_name", [MODEL_NAME])
399
@pytest.mark.parametrize("param_name", ["encoding_format", "embed_dtype", "endianness"])
tests/entrypoints/pooling/pooling/test_online.py
+56
@@ -18,6 +18,7 @@ from vllm.utils.serial_utils import (
18
ENDIANNESS,
19
MetadataItem,
20
binary2tensor,
21
+ build_metadata_items,
22
decode_pooling_output,
23
)
24
@@ -352,6 +353,61 @@ async def test_bytes_embed_dtype_and_endianness(
353
)
354
355
356
+@pytest.mark.asyncio
357
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
358
+async def test_bytes_only_embed_dtype_and_endianness(
359
+ server: RemoteOpenAIServer, model_name: str
360
+):
361
+ input_texts = [
362
+ "The best thing about vLLM is that it supports many different models",
363
+ ] * 2
364
+
365
+ url = server.url_for("pooling")
366
+ float_response = requests.post(
367
+ url,
368
+ json={
369
+ "model": model_name,
370
+ "input": input_texts,
371
+ "encoding_format": "float",
372
+ },
373
+ )
374
+ responses_float = PoolingResponse.model_validate(float_response.json())
375
+ float_data = [np.array(d.data).squeeze(-1).tolist() for d in responses_float.data]
376
+ n_tokens = responses_float.usage.prompt_tokens // len(input_texts)
377
+
378
+ for embed_dtype in list(EMBED_DTYPE_TO_TORCH_DTYPE.keys()):
379
+ for endianness in ENDIANNESS:
380
+ responses_bytes = requests.post(
381
+ url,
382
+ json={
383
+ "model": model_name,
384
+ "input": input_texts,
385
+ "encoding_format": "bytes_only",
386
+ "embed_dtype": embed_dtype,
387
+ "endianness": endianness,
388
+ },
389
+ )
390
+
391
+ assert "metadata" not in responses_bytes.headers
392
+ body = responses_bytes.content
393
+ items = build_metadata_items(
394
+ embed_dtype=embed_dtype,
395
+ endianness=endianness,
396
+ shape=(n_tokens, 1),
397
+ n_request=len(input_texts),
398
+ )
399
+ bytes_data = decode_pooling_output(items=items, body=body)
400
+ bytes_data = [x.to(torch.float32).view(-1).tolist() for x in bytes_data]
401
+
402
+ check_embeddings_close(
403
+ embeddings_0_lst=float_data,
404
+ embeddings_1_lst=bytes_data,
405
+ name_0="float_data",
406
+ name_1="bytes_data",
407
+ tol=1e-2,
408
+ )
409
+
410
+
411
@pytest.mark.asyncio
412
@pytest.mark.parametrize("model_name", [MODEL_NAME])
413
@pytest.mark.parametrize("param_name", ["encoding_format", "embed_dtype", "endianness"])
vllm/entrypoints/pooling/embed/api_router.py
+2
-2
@@ -59,8 +59,8 @@ async def create_embedding(
59
return JSONResponse(content=generator.model_dump())
60
elif isinstance(generator, EmbeddingBytesResponse):
61
return StreamingResponse(
62
- content=generator.body,
63
- headers={"metadata": generator.metadata},
62
+ content=generator.content,
63
+ headers=generator.headers,
64
media_type=generator.media_type,
65
)
66
vllm/entrypoints/pooling/embed/protocol.py
+2
-2
@@ -203,6 +203,6 @@ class EmbeddingResponse(OpenAIBaseModel):
203
204
205
class EmbeddingBytesResponse(OpenAIBaseModel):
206
- body: list[bytes]
207
- metadata: str
206
+ content: list[bytes]
207
+ headers: dict[str, str] | None = None
208
media_type: str = "application/octet-stream"
vllm/entrypoints/pooling/embed/serving.py
+20
-14
@@ -163,29 +163,35 @@ class EmbeddingMixin(OpenAIServing):
163
usage=usage,
164
)
165
166
- def encode_bytes():
167
- body, items, usage = encode_pooling_bytes(
166
+ def encode_bytes(bytes_only: bool) -> EmbeddingBytesResponse:
167
+ content, items, usage = encode_pooling_bytes(
168
pooling_outputs=final_res_batch_checked,
169
embed_dtype=embed_dtype,
170
endianness=endianness,
171
)
172
173
- metadata = {
174
- "id": ctx.request_id,
175
- "created": ctx.created_time,
176
- "model": ctx.model_name,
177
- "data": items,
178
- "usage": usage,
179
- }
180
- return EmbeddingBytesResponse(
181
- body=body,
182
- metadata=json.dumps(metadata),
173
+ headers = (
174
+ None
175
+ if bytes_only
176
+ else {
177
+ "metadata": json.dumps(
178
+ {
179
+ "id": ctx.request_id,
180
+ "created": ctx.created_time,
181
+ "model": ctx.model_name,
182
+ "data": items,
183
+ "usage": usage,
184
+ }
185
+ )
186
+ }
187
)
188
189
+ return EmbeddingBytesResponse(content=content, headers=headers)
190
+
191
if encoding_format == "float" or encoding_format == "base64":
192
return encode_float_base64()
187
- elif encoding_format == "bytes":
188
- return encode_bytes()
193
+ elif encoding_format == "bytes" or encoding_format == "bytes_only":
194
+ return encode_bytes(bytes_only=encoding_format == "bytes_only")
195
else:
196
assert_never(encoding_format)
197
vllm/entrypoints/pooling/pooling/api_router.py
+2
-2
@@ -55,8 +55,8 @@ async def create_pooling(request: PoolingRequest, raw_request: Request):
55
return JSONResponse(content=generator.model_dump())
56
elif isinstance(generator, PoolingBytesResponse):
57
return StreamingResponse(
58
- content=generator.body,
59
- headers={"metadata": generator.metadata},
58
+ content=generator.content,
59
+ headers=generator.headers,
60
media_type=generator.media_type,
61
)
62
vllm/entrypoints/pooling/pooling/protocol.py
+2
-2
@@ -143,6 +143,6 @@ class PoolingResponse(OpenAIBaseModel):
143
144
145
class PoolingBytesResponse(OpenAIBaseModel):
146
- body: list[bytes]
147
- metadata: str
146
+ content: list[bytes]
147
+ headers: dict[str, str] | None = None
148
media_type: str = "application/octet-stream"
vllm/entrypoints/pooling/pooling/serving.py
+22
-13
@@ -314,29 +314,38 @@ class OpenAIServingPooling(OpenAIServing):
314
usage=usage,
315
)
316
317
- def encode_bytes():
318
- body, items, usage = encode_pooling_bytes(
317
+ def encode_bytes(bytes_only: bool) -> PoolingBytesResponse:
318
+ content, items, usage = encode_pooling_bytes(
319
pooling_outputs=final_res_batch,
320
embed_dtype=embed_dtype,
321
endianness=endianness,
322
)
323
324
- metadata = {
325
- "id": request_id,
326
- "created": created_time,
327
- "model": model_name,
328
- "data": items,
329
- "usage": usage,
330
- }
324
+ headers = (
325
+ None
326
+ if bytes_only
327
+ else {
328
+ "metadata": json.dumps(
329
+ {
330
+ "id": request_id,
331
+ "created": created_time,
332
+ "model": model_name,
333
+ "data": items,
334
+ "usage": usage,
335
+ }
336
+ )
337
+ }
338
+ )
339
+
340
return PoolingBytesResponse(
332
- body=body,
333
- metadata=json.dumps(metadata),
341
+ content=content,
342
+ headers=headers,
343
)
344
345
if encoding_format == "float" or encoding_format == "base64":
346
return encode_float_base64()
338
- elif encoding_format == "bytes":
339
- return encode_bytes()
347
+ elif encoding_format == "bytes" or encoding_format == "bytes_only":
348
+ return encode_bytes(bytes_only=encoding_format == "bytes_only")
349
else:
350
assert_never(encoding_format)
351
vllm/utils/serial_utils.py
+39
-4
@@ -2,15 +2,19 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import base64
4
import io
5
+import math
6
import sys
7
from dataclasses import dataclass
7
-from typing import Literal
8
+from typing import TYPE_CHECKING, Any, Literal
9
10
import numpy as np
11
import torch
12
from typing_extensions import assert_never
13
13
-from vllm import PoolingRequestOutput
14
+if TYPE_CHECKING:
15
+ from vllm import PoolingRequestOutput
16
+else:
17
+ PoolingRequestOutput = Any
18
19
sys_byteorder = sys.byteorder
20
@@ -27,6 +31,14 @@ EMBED_DTYPE_TO_TORCH_DTYPE = {
31
"fp8_e5m2": torch.float8_e5m2,
32
}
33
34
+EMBED_DTYPE_TO_N_BYTES = {
35
+ "float32": 4,
36
+ "float16": 2,
37
+ "bfloat16": 2,
38
+ "fp8_e4m3": 1,
39
+ "fp8_e5m2": 1,
40
+}
41
+
42
43
EMBED_DTYPE_TO_TORCH_DTYPE_VIEW = {
44
"float32": torch.float32,
@@ -50,7 +62,7 @@ ENDIANNESS = ["native", "big", "little"]
62
63
EmbedDType = Literal["float32", "float16", "bfloat16", "fp8_e4m3", "fp8_e5m2"]
64
Endianness = Literal["native", "big", "little"]
53
-EncodingFormat = Literal["float", "base64", "bytes"]
65
+EncodingFormat = Literal["float", "base64", "bytes", "bytes_only"]
66
67
68
def tensor2base64(x: torch.Tensor) -> str:
@@ -114,7 +126,7 @@ def encode_pooling_output(
126
elif encoding_format == "base64":
127
embedding_bytes = tensor2binary(output.outputs.data, embed_dtype, endianness)
128
return base64.b64encode(embedding_bytes).decode("utf-8")
117
- elif encoding_format == "bytes":
129
+ elif encoding_format == "bytes" or encoding_format == "bytes_only":
130
return tensor2binary(output.outputs.data, embed_dtype, endianness)
131
assert_never(encoding_format)
132
@@ -129,6 +141,29 @@ class MetadataItem:
141
shape: tuple[int, ...]
142
143
144
+def build_metadata_items(
145
+ embed_dtype: EmbedDType,
146
+ endianness: Endianness,
147
+ shape: tuple[int, ...],
148
+ n_request: int,
149
+):
150
+ n_bytes = EMBED_DTYPE_TO_N_BYTES[embed_dtype]
151
+ size = math.prod(shape)
152
+ items = [
153
+ MetadataItem(
154
+ index=i,
155
+ embed_dtype=embed_dtype,
156
+ endianness=endianness,
157
+ start=i * size * n_bytes,
158
+ end=(i + 1) * size * n_bytes,
159
+ shape=shape,
160
+ )
161
+ for i in range(n_request)
162
+ ]
163
+
164
+ return items
165
+
166
+
167
def encode_pooling_bytes(
168
pooling_outputs: list[PoolingRequestOutput],
169
embed_dtype: EmbedDType,