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 9, 01:15 UTC
[Refactor] Clean up pooling serial utils (#33665)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Feb 3, 2026 at 18:29 UTC
83449a5ff04f70a20c24e8e6fc719881b29e10ac
9 files changed
+407
-322
examples/pooling/embed/embedding_requests_base64_online.py
+2
-6
@@ -12,11 +12,7 @@ import base64
12
import requests
13
import torch
14
15
-from vllm.utils.serial_utils import (
16
- EMBED_DTYPE_TO_TORCH_DTYPE,
17
- ENDIANNESS,
18
- binary2tensor,
19
-)
15
+from vllm.utils.serial_utils import EMBED_DTYPES, ENDIANNESS, binary2tensor
16
17
18
def post_http_request(prompt: dict, api_url: str) -> requests.Response:
@@ -45,7 +41,7 @@ def main(args):
41
] * 2
42
43
# The OpenAI client does not support the embed_dtype and endianness parameters.
48
- for embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE:
44
+ for embed_dtype in EMBED_DTYPES:
45
for endianness in ENDIANNESS:
46
prompt = {
47
"model": model,
examples/pooling/embed/embedding_requests_bytes_online.py
+4
-5
@@ -12,13 +12,12 @@ import json
12
import requests
13
import torch
14
15
-from vllm.utils.serial_utils import (
16
- EMBED_DTYPE_TO_TORCH_DTYPE,
17
- ENDIANNESS,
15
+from vllm.entrypoints.pooling.utils import (
16
MetadataItem,
17
build_metadata_items,
18
decode_pooling_output,
19
)
20
+from vllm.utils.serial_utils import EMBED_DTYPES, ENDIANNESS
21
22
23
def post_http_request(prompt: dict, api_url: str) -> requests.Response:
@@ -51,7 +50,7 @@ def main(args):
50
51
# The OpenAI client does not support the bytes encoding_format.
52
# The OpenAI client does not support the embed_dtype and endianness parameters.
54
- for embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE:
53
+ for embed_dtype in EMBED_DTYPES:
54
for endianness in ENDIANNESS:
55
prompt = {
56
"model": model,
@@ -74,7 +73,7 @@ def main(args):
73
# The vllm server always sorts the returned embeddings in the order of input. So
74
# returning metadata is not necessary. You can set encoding_format to bytes_only
75
# to let the server not return metadata.
77
- for embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE:
76
+ for embed_dtype in EMBED_DTYPES:
77
for endianness in ENDIANNESS:
78
prompt = {
79
"model": model,
tests/entrypoints/pooling/embed/test_online.py
+7
-9
@@ -17,16 +17,14 @@ from tests.models.utils import check_embeddings_close
17
from tests.utils import RemoteOpenAIServer
18
from vllm.entrypoints.pooling.embed.protocol import EmbeddingResponse
19
from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse
20
-from vllm.platforms import current_platform
21
-from vllm.tokenizers import get_tokenizer
22
-from vllm.utils.serial_utils import (
23
- EMBED_DTYPE_TO_TORCH_DTYPE,
24
- ENDIANNESS,
20
+from vllm.entrypoints.pooling.utils import (
21
MetadataItem,
26
- binary2tensor,
22
build_metadata_items,
23
decode_pooling_output,
24
)
25
+from vllm.platforms import current_platform
26
+from vllm.tokenizers import get_tokenizer
27
+from vllm.utils.serial_utils import EMBED_DTYPES, ENDIANNESS, binary2tensor
28
29
MODEL_NAME = "intfloat/multilingual-e5-small"
30
DUMMY_CHAT_TEMPLATE = """{% for message in messages %}{{message['role'] + ': ' + message['content'] + '\\n'}}{% endfor %}""" # noqa: E501
@@ -535,7 +533,7 @@ async def test_base64_embed_dtype_and_endianness(
533
)
534
float_data = [d.embedding for d in responses_float.data]
535
538
- for embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE:
536
+ for embed_dtype in EMBED_DTYPES:
537
for endianness in ENDIANNESS:
538
responses_base64 = requests.post(
539
server.url_for("/v1/embeddings"),
@@ -574,7 +572,7 @@ async def test_bytes_embed_dtype_and_endianness(
572
)
573
float_data = [d.embedding for d in responses_float.data]
574
577
- for embed_dtype in list(EMBED_DTYPE_TO_TORCH_DTYPE.keys()):
575
+ for embed_dtype in EMBED_DTYPES:
576
for endianness in ENDIANNESS:
577
responses_bytes = requests.post(
578
server.url_for("/v1/embeddings"),
@@ -618,7 +616,7 @@ async def test_bytes_only_embed_dtype_and_endianness(
616
float_data = [d.embedding for d in responses_float.data]
617
embedding_size = len(float_data[0])
618
621
- for embed_dtype in list(EMBED_DTYPE_TO_TORCH_DTYPE.keys()):
619
+ for embed_dtype in EMBED_DTYPES:
620
for endianness in ENDIANNESS:
621
responses_bytes = requests.post(
622
server.url_for("/v1/embeddings"),
tests/entrypoints/pooling/pooling/test_online.py
+6
-8
@@ -12,15 +12,13 @@ import torch
12
from tests.models.utils import check_embeddings_close
13
from tests.utils import RemoteOpenAIServer
14
from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse
15
-from vllm.tokenizers import get_tokenizer
16
-from vllm.utils.serial_utils import (
17
- EMBED_DTYPE_TO_TORCH_DTYPE,
18
- ENDIANNESS,
15
+from vllm.entrypoints.pooling.utils import (
16
MetadataItem,
20
- binary2tensor,
17
build_metadata_items,
18
decode_pooling_output,
19
)
20
+from vllm.tokenizers import get_tokenizer
21
+from vllm.utils.serial_utils import EMBED_DTYPES, ENDIANNESS, binary2tensor
22
23
MODEL_NAME = "internlm/internlm2-1_8b-reward"
24
DUMMY_CHAT_TEMPLATE = """{% for message in messages %}{{message['role'] + ': ' + message['content'] + '\\n'}}{% endfor %}""" # noqa: E501
@@ -342,7 +340,7 @@ async def test_base64_embed_dtype_and_endianness(
340
responses_float = PoolingResponse.model_validate(float_response.json())
341
float_data = [np.array(d.data).squeeze(-1).tolist() for d in responses_float.data]
342
345
- for embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE:
343
+ for embed_dtype in EMBED_DTYPES:
344
for endianness in ENDIANNESS:
345
responses_base64 = requests.post(
346
url,
@@ -389,7 +387,7 @@ async def test_bytes_embed_dtype_and_endianness(
387
responses_float = PoolingResponse.model_validate(float_response.json())
388
float_data = [np.array(d.data).squeeze(-1).tolist() for d in responses_float.data]
389
392
- for embed_dtype in list(EMBED_DTYPE_TO_TORCH_DTYPE.keys()):
390
+ for embed_dtype in EMBED_DTYPES:
391
for endianness in ENDIANNESS:
392
responses_bytes = requests.post(
393
url,
@@ -438,7 +436,7 @@ async def test_bytes_only_embed_dtype_and_endianness(
436
float_data = [np.array(d.data).squeeze(-1).tolist() for d in responses_float.data]
437
n_tokens = responses_float.usage.prompt_tokens // len(input_texts)
438
441
- for embed_dtype in list(EMBED_DTYPE_TO_TORCH_DTYPE.keys()):
439
+ for embed_dtype in EMBED_DTYPES:
440
for endianness in ENDIANNESS:
441
responses_bytes = requests.post(
442
url,
tests/utils_/test_serial_utils.py
+5
-3
@@ -5,17 +5,19 @@ import torch
5
6
from tests.models.utils import check_embeddings_close
7
from vllm.utils.serial_utils import (
8
- EMBED_DTYPE_TO_TORCH_DTYPE,
8
+ EMBED_DTYPES,
9
ENDIANNESS,
10
+ EmbedDType,
11
+ Endianness,
12
binary2tensor,
13
tensor2binary,
14
)
15
16
17
@pytest.mark.parametrize("endianness", ENDIANNESS)
16
-@pytest.mark.parametrize("embed_dtype", EMBED_DTYPE_TO_TORCH_DTYPE.keys())
18
+@pytest.mark.parametrize("embed_dtype", EMBED_DTYPES.keys())
19
@torch.inference_mode()
18
-def test_encode_and_decode(embed_dtype: str, endianness: str):
20
+def test_encode_and_decode(embed_dtype: EmbedDType, endianness: Endianness):
21
for i in range(10):
22
tensor = torch.rand(2, 3, 5, 7, 11, 13, device="cpu", dtype=torch.float32)
23
shape = tensor.shape
vllm/entrypoints/pooling/embed/serving.py
+111
-67
@@ -1,8 +1,9 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import json
4
-from collections.abc import AsyncGenerator, Mapping
5
-from typing import Any, Final, TypeAlias
4
+from collections.abc import AsyncGenerator, Callable, Mapping
5
+from functools import partial
6
+from typing import Any, Final, Literal, TypeAlias, cast
7
8
import torch
9
from fastapi import Request
@@ -22,16 +23,18 @@ from vllm.entrypoints.pooling.embed.protocol import (
23
EmbeddingResponse,
24
EmbeddingResponseData,
25
)
26
+from vllm.entrypoints.pooling.utils import (
27
+ encode_pooling_bytes,
28
+ encode_pooling_output_base64,
29
+ encode_pooling_output_float,
30
+)
31
from vllm.inputs.data import EmbedsPrompt, TokensPrompt
32
from vllm.logger import init_logger
33
from vllm.outputs import PoolingOutput, PoolingRequestOutput
34
from vllm.pooling_params import PoolingParams
35
from vllm.utils.async_utils import merge_async_iterators
36
from vllm.utils.collection_utils import chunk_list
31
-from vllm.utils.serial_utils import (
32
- encode_pooling_bytes,
33
- encode_pooling_output,
34
-)
37
+from vllm.utils.serial_utils import EmbedDType, Endianness
38
39
logger = init_logger(__name__)
40
@@ -113,79 +116,120 @@ class OpenAIServingEmbedding(OpenAIServing):
116
logger.exception("Error in preprocessing prompt inputs")
117
return self.create_error_response(str(e))
118
119
+ def request_output_to_embed_json_response(
120
+ self,
121
+ final_res_batch: list[PoolingRequestOutput],
122
+ request_id: str,
123
+ created_time: int,
124
+ model_name: str,
125
+ encoding_format: Literal["float", "base64"],
126
+ embed_dtype: EmbedDType,
127
+ endianness: Endianness,
128
+ ) -> EmbeddingResponse:
129
+ encode_fn = cast(
130
+ Callable[[PoolingRequestOutput], list[float] | str],
131
+ (
132
+ encode_pooling_output_float
133
+ if encoding_format == "float"
134
+ else partial(
135
+ encode_pooling_output_base64,
136
+ embed_dtype=embed_dtype,
137
+ endianness=endianness,
138
+ )
139
+ ),
140
+ )
141
+
142
+ items: list[EmbeddingResponseData] = []
143
+ num_prompt_tokens = 0
144
+
145
+ for idx, final_res in enumerate(final_res_batch):
146
+ item = EmbeddingResponseData(
147
+ index=idx,
148
+ embedding=encode_fn(final_res),
149
+ )
150
+ prompt_token_ids = final_res.prompt_token_ids
151
+
152
+ items.append(item)
153
+ num_prompt_tokens += len(prompt_token_ids)
154
+
155
+ usage = UsageInfo(
156
+ prompt_tokens=num_prompt_tokens,
157
+ total_tokens=num_prompt_tokens,
158
+ )
159
+
160
+ return EmbeddingResponse(
161
+ id=request_id,
162
+ created=created_time,
163
+ model=model_name,
164
+ data=items,
165
+ usage=usage,
166
+ )
167
+
168
+ def request_output_to_embed_bytes_response(
169
+ self,
170
+ final_res_batch: list[PoolingRequestOutput],
171
+ request_id: str,
172
+ created_time: int,
173
+ model_name: str,
174
+ encoding_format: Literal["bytes", "bytes_only"],
175
+ embed_dtype: EmbedDType,
176
+ endianness: Endianness,
177
+ ) -> EmbeddingBytesResponse:
178
+ content, items, usage = encode_pooling_bytes(
179
+ pooling_outputs=final_res_batch,
180
+ embed_dtype=embed_dtype,
181
+ endianness=endianness,
182
+ )
183
+
184
+ headers = (
185
+ None
186
+ if encoding_format == "bytes_only"
187
+ else {
188
+ "metadata": json.dumps(
189
+ {
190
+ "id": request_id,
191
+ "created": created_time,
192
+ "model": model_name,
193
+ "data": items,
194
+ "usage": usage,
195
+ }
196
+ )
197
+ }
198
+ )
199
+
200
+ return EmbeddingBytesResponse(content=content, headers=headers)
201
+
202
def _build_response(
203
self,
204
ctx: EmbeddingServeContext,
205
) -> EmbeddingResponse | EmbeddingBytesResponse | ErrorResponse:
120
- final_res_batch_checked = ctx.final_res_batch
121
-
206
encoding_format = ctx.request.encoding_format
207
embed_dtype = ctx.request.embed_dtype
208
endianness = ctx.request.endianness
209
126
- def encode_float_base64():
127
- items: list[EmbeddingResponseData] = []
128
- num_prompt_tokens = 0
129
-
130
- for idx, final_res in enumerate(final_res_batch_checked):
131
- item = EmbeddingResponseData(
132
- index=idx,
133
- embedding=encode_pooling_output(
134
- final_res,
135
- encoding_format=encoding_format,
136
- embed_dtype=embed_dtype,
137
- endianness=endianness,
138
- ),
139
- )
140
- prompt_token_ids = final_res.prompt_token_ids
141
-
142
- items.append(item)
143
- num_prompt_tokens += len(prompt_token_ids)
144
-
145
- usage = UsageInfo(
146
- prompt_tokens=num_prompt_tokens,
147
- total_tokens=num_prompt_tokens,
148
- )
149
-
150
- return EmbeddingResponse(
151
- id=ctx.request_id,
152
- created=ctx.created_time,
153
- model=ctx.model_name,
154
- data=items,
155
- usage=usage,
210
+ if encoding_format == "float" or encoding_format == "base64":
211
+ return self.request_output_to_embed_json_response(
212
+ ctx.final_res_batch,
213
+ ctx.request_id,
214
+ ctx.created_time,
215
+ ctx.model_name,
216
+ encoding_format,
217
+ embed_dtype,
218
+ endianness,
219
)
220
158
- def encode_bytes(bytes_only: bool) -> EmbeddingBytesResponse:
159
- content, items, usage = encode_pooling_bytes(
160
- pooling_outputs=final_res_batch_checked,
161
- embed_dtype=embed_dtype,
162
- endianness=endianness,
221
+ if encoding_format == "bytes" or encoding_format == "bytes_only":
222
+ return self.request_output_to_embed_bytes_response(
223
+ ctx.final_res_batch,
224
+ ctx.request_id,
225
+ ctx.created_time,
226
+ ctx.model_name,
227
+ encoding_format,
228
+ embed_dtype,
229
+ endianness,
230
)
231
165
- headers = (
166
- None
167
- if bytes_only
168
- else {
169
- "metadata": json.dumps(
170
- {
171
- "id": ctx.request_id,
172
- "created": ctx.created_time,
173
- "model": ctx.model_name,
174
- "data": items,
175
- "usage": usage,
176
- }
177
- )
178
- }
179
- )
180
-
181
- return EmbeddingBytesResponse(content=content, headers=headers)
182
-
183
- if encoding_format == "float" or encoding_format == "base64":
184
- return encode_float_base64()
185
- elif encoding_format == "bytes" or encoding_format == "bytes_only":
186
- return encode_bytes(bytes_only=encoding_format == "bytes_only")
187
- else:
188
- assert_never(encoding_format)
232
+ assert_never(encoding_format)
233
234
def _get_max_position_embeddings(self) -> int:
235
"""Get the model's effective maximum sequence length for chunking."""
vllm/entrypoints/pooling/pooling/serving.py
+107
-67
@@ -4,8 +4,9 @@
4
import asyncio
5
import json
6
import time
7
-from collections.abc import AsyncGenerator, Sequence
8
-from typing import Any, Final, cast
7
+from collections.abc import AsyncGenerator, Callable, Sequence
8
+from functools import partial
9
+from typing import Any, Final, Literal, cast
10
11
import jinja2
12
from fastapi import Request
@@ -27,17 +28,16 @@ from vllm.entrypoints.pooling.pooling.protocol import (
28
PoolingResponse,
29
PoolingResponseData,
30
)
31
+from vllm.entrypoints.pooling.utils import (
32
+ encode_pooling_bytes,
33
+ encode_pooling_output_base64,
34
+ encode_pooling_output_float,
35
+)
36
from vllm.logger import init_logger
37
from vllm.outputs import PoolingRequestOutput
38
from vllm.tasks import PoolingTask, SupportedTask
39
from vllm.utils.async_utils import merge_async_iterators
34
-from vllm.utils.serial_utils import (
35
- EmbedDType,
36
- EncodingFormat,
37
- Endianness,
38
- encode_pooling_bytes,
39
- encode_pooling_output,
40
-)
40
+from vllm.utils.serial_utils import EmbedDType, EncodingFormat, Endianness
41
42
logger = init_logger(__name__)
43
@@ -256,79 +256,119 @@ class OpenAIServingPooling(OpenAIServing):
256
257
return response
258
259
- def request_output_to_pooling_response(
259
+ def request_output_to_pooling_json_response(
260
self,
261
final_res_batch: list[PoolingRequestOutput],
262
request_id: str,
263
created_time: int,
264
model_name: str,
265
- encoding_format: EncodingFormat,
265
+ encoding_format: Literal["float", "base64"],
266
embed_dtype: EmbedDType,
267
endianness: Endianness,
268
- ) -> PoolingResponse | PoolingBytesResponse:
269
- def encode_float_base64():
270
- items: list[PoolingResponseData] = []
271
- num_prompt_tokens = 0
272
-
273
- for idx, final_res in enumerate(final_res_batch):
274
- item = PoolingResponseData(
275
- index=idx,
276
- data=encode_pooling_output(
277
- final_res,
278
- encoding_format=encoding_format,
279
- embed_dtype=embed_dtype,
280
- endianness=endianness,
281
- ),
268
+ ) -> PoolingResponse:
269
+ encode_fn = cast(
270
+ Callable[[PoolingRequestOutput], list[float] | str],
271
+ (
272
+ encode_pooling_output_float
273
+ if encoding_format == "float"
274
+ else partial(
275
+ encode_pooling_output_base64,
276
+ embed_dtype=embed_dtype,
277
+ endianness=endianness,
278
)
283
- prompt_token_ids = final_res.prompt_token_ids
279
+ ),
280
+ )
281
285
- items.append(item)
286
- num_prompt_tokens += len(prompt_token_ids)
282
+ items: list[PoolingResponseData] = []
283
+ num_prompt_tokens = 0
284
288
- usage = UsageInfo(
289
- prompt_tokens=num_prompt_tokens,
290
- total_tokens=num_prompt_tokens,
285
+ for idx, final_res in enumerate(final_res_batch):
286
+ item = PoolingResponseData(
287
+ index=idx,
288
+ data=encode_fn(final_res),
289
)
290
+ prompt_token_ids = final_res.prompt_token_ids
291
293
- return PoolingResponse(
294
- id=request_id,
295
- created=created_time,
296
- model=model_name,
297
- data=items,
298
- usage=usage,
299
- )
292
+ items.append(item)
293
+ num_prompt_tokens += len(prompt_token_ids)
294
301
- def encode_bytes(bytes_only: bool) -> PoolingBytesResponse:
302
- content, items, usage = encode_pooling_bytes(
303
- pooling_outputs=final_res_batch,
304
- embed_dtype=embed_dtype,
305
- endianness=endianness,
306
- )
295
+ usage = UsageInfo(
296
+ prompt_tokens=num_prompt_tokens,
297
+ total_tokens=num_prompt_tokens,
298
+ )
299
308
- headers = (
309
- None
310
- if bytes_only
311
- else {
312
- "metadata": json.dumps(
313
- {
314
- "id": request_id,
315
- "created": created_time,
316
- "model": model_name,
317
- "data": items,
318
- "usage": usage,
319
- }
320
- )
321
- }
300
+ return PoolingResponse(
301
+ id=request_id,
302
+ created=created_time,
303
+ model=model_name,
304
+ data=items,
305
+ usage=usage,
306
+ )
307
+
308
+ def request_output_to_pooling_bytes_response(
309
+ self,
310
+ final_res_batch: list[PoolingRequestOutput],
311
+ request_id: str,
312
+ created_time: int,
313
+ model_name: str,
314
+ encoding_format: Literal["bytes", "bytes_only"],
315
+ embed_dtype: EmbedDType,
316
+ endianness: Endianness,
317
+ ) -> 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
+ headers = (
325
+ None
326
+ if encoding_format == "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(content=content, headers=headers)
341
+
342
+ def request_output_to_pooling_response(
343
+ self,
344
+ final_res_batch: list[PoolingRequestOutput],
345
+ request_id: str,
346
+ created_time: int,
347
+ model_name: str,
348
+ encoding_format: EncodingFormat,
349
+ embed_dtype: EmbedDType,
350
+ endianness: Endianness,
351
+ ) -> PoolingResponse | PoolingBytesResponse:
352
+ if encoding_format == "float" or encoding_format == "base64":
353
+ return self.request_output_to_pooling_json_response(
354
+ final_res_batch,
355
+ request_id,
356
+ created_time,
357
+ model_name,
358
+ encoding_format,
359
+ embed_dtype,
360
+ endianness,
361
)
362
324
- return PoolingBytesResponse(
325
- content=content,
326
- headers=headers,
363
+ if encoding_format == "bytes" or encoding_format == "bytes_only":
364
+ return self.request_output_to_pooling_bytes_response(
365
+ final_res_batch,
366
+ request_id,
367
+ created_time,
368
+ model_name,
369
+ encoding_format,
370
+ embed_dtype,
371
+ endianness,
372
)
373
329
- if encoding_format == "float" or encoding_format == "base64":
330
- return encode_float_base64()
331
- elif encoding_format == "bytes" or encoding_format == "bytes_only":
332
- return encode_bytes(bytes_only=encoding_format == "bytes_only")
333
- else:
334
- assert_never(encoding_format)
374
+ assert_never(encoding_format)
vllm/entrypoints/pooling/utils.py
new
+124
@@ -0,0 +1,124 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import math
4
+from dataclasses import dataclass
5
+from typing import Any
6
+
7
+import pybase64
8
+import torch
9
+
10
+from vllm.outputs import PoolingRequestOutput
11
+from vllm.utils.serial_utils import (
12
+ EMBED_DTYPES,
13
+ EmbedDType,
14
+ Endianness,
15
+ binary2tensor,
16
+ tensor2binary,
17
+)
18
+
19
+
20
+@dataclass
21
+class MetadataItem:
22
+ index: int
23
+ embed_dtype: EmbedDType
24
+ endianness: Endianness
25
+ start: int
26
+ end: int
27
+ shape: tuple[int, ...]
28
+
29
+
30
+def build_metadata_items(
31
+ embed_dtype: EmbedDType,
32
+ endianness: Endianness,
33
+ shape: tuple[int, ...],
34
+ n_request: int,
35
+) -> list[MetadataItem]:
36
+ n_bytes = EMBED_DTYPES[embed_dtype].nbytes
37
+ size = math.prod(shape)
38
+
39
+ return [
40
+ MetadataItem(
41
+ index=i,
42
+ embed_dtype=embed_dtype,
43
+ endianness=endianness,
44
+ start=i * size * n_bytes,
45
+ end=(i + 1) * size * n_bytes,
46
+ shape=shape,
47
+ )
48
+ for i in range(n_request)
49
+ ]
50
+
51
+
52
+def encode_pooling_output_float(output: PoolingRequestOutput) -> list[float]:
53
+ return output.outputs.data.tolist()
54
+
55
+
56
+def encode_pooling_output_binary(
57
+ output: PoolingRequestOutput,
58
+ embed_dtype: EmbedDType,
59
+ endianness: Endianness,
60
+) -> bytes:
61
+ return tensor2binary(output.outputs.data, embed_dtype, endianness)
62
+
63
+
64
+def encode_pooling_output_base64(
65
+ output: PoolingRequestOutput,
66
+ embed_dtype: EmbedDType,
67
+ endianness: Endianness,
68
+) -> str:
69
+ embedding_bytes = tensor2binary(output.outputs.data, embed_dtype, endianness)
70
+ return pybase64.b64encode(embedding_bytes).decode("utf-8")
71
+
72
+
73
+def encode_pooling_bytes(
74
+ pooling_outputs: list[PoolingRequestOutput],
75
+ embed_dtype: EmbedDType,
76
+ endianness: Endianness,
77
+) -> tuple[list[bytes], list[dict[str, Any]], dict[str, Any]]:
78
+ num_prompt_tokens = 0
79
+ items: list[dict[str, Any]] = []
80
+ body: list[bytes] = []
81
+ offset = 0
82
+ for idx, output in enumerate(pooling_outputs):
83
+ binary = tensor2binary(
84
+ tensor=output.outputs.data,
85
+ embed_dtype=embed_dtype,
86
+ endianness=endianness,
87
+ )
88
+ size = len(binary)
89
+
90
+ # Dictionary form of MetadataItem
91
+ item = dict(
92
+ index=idx,
93
+ embed_dtype=embed_dtype,
94
+ endianness=endianness,
95
+ start=offset,
96
+ end=offset + size,
97
+ shape=output.outputs.data.shape,
98
+ )
99
+
100
+ body.append(binary)
101
+ items.append(item)
102
+ prompt_token_ids = output.prompt_token_ids
103
+ num_prompt_tokens += len(prompt_token_ids)
104
+ offset += size
105
+
106
+ # Dictionary form of UsageInfo
107
+ usage = dict(
108
+ prompt_tokens=num_prompt_tokens,
109
+ total_tokens=num_prompt_tokens,
110
+ )
111
+
112
+ return body, items, usage
113
+
114
+
115
+def decode_pooling_output(items: list[MetadataItem], body: bytes) -> list[torch.Tensor]:
116
+ return [
117
+ binary2tensor(
118
+ body[item.start : item.end],
119
+ item.shape,
120
+ item.embed_dtype,
121
+ item.endianness,
122
+ )
123
+ for item in sorted(items, key=lambda x: x.index)
124
+ ]
vllm/utils/serial_utils.py
+41
-157
@@ -1,69 +1,49 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import base64
3
import io
5
-import math
4
import sys
5
+from collections.abc import Mapping
6
from dataclasses import dataclass
8
-from typing import TYPE_CHECKING, Any, Literal
7
+from typing import Literal, get_args
8
9
import numpy as np
10
+import numpy.typing as npt
11
+import pybase64
12
import torch
12
-from typing_extensions import assert_never
13
-
14
-if TYPE_CHECKING:
15
- from vllm import PoolingRequestOutput
16
-else:
17
- PoolingRequestOutput = Any
13
14
sys_byteorder = sys.byteorder
15
16
22
-EMBED_DTYPE_TO_TORCH_DTYPE = {
23
- "float32": torch.float32,
24
- "float16": torch.float16,
25
- "bfloat16": torch.bfloat16,
26
- # I'm not sure if other platforms' CPUs support the fp8 data format.
27
- # EMBED_DTYPE only uses the fp8 data representation,
28
- # does not use fp8 computation, and only occurs on the CPU.
29
- # Apologize for any possible break.
30
- "fp8_e4m3": torch.float8_e4m3fn,
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
-
17
+@dataclass(frozen=True)
18
+class DTypeInfo:
19
+ torch_dtype: torch.dtype
20
43
-EMBED_DTYPE_TO_TORCH_DTYPE_VIEW = {
44
- "float32": torch.float32,
45
- "float16": torch.float16,
46
- # numpy does not support bfloat16 and fp8
47
- "bfloat16": torch.float16,
48
- "fp8_e4m3": torch.uint8,
49
- "fp8_e5m2": torch.uint8,
50
-}
21
+ torch_view_dtype: torch.dtype
22
+ numpy_view_dtype: npt.DTypeLike
23
52
-EMBED_DTYPE_TO_NUMPY_DTYPE_VIEW = {
53
- "float32": np.float32,
54
- "float16": np.float16,
55
- # numpy does not support bfloat16 and fp8
56
- "bfloat16": np.float16,
57
- "fp8_e4m3": np.uint8,
58
- "fp8_e5m2": np.uint8,
59
-}
24
+ @property
25
+ def nbytes(self) -> int:
26
+ return self.torch_dtype.itemsize
27
61
-ENDIANNESS = ["native", "big", "little"]
28
29
EmbedDType = Literal["float32", "float16", "bfloat16", "fp8_e4m3", "fp8_e5m2"]
30
Endianness = Literal["native", "big", "little"]
31
EncodingFormat = Literal["float", "base64", "bytes", "bytes_only"]
32
33
+# I'm not sure if other platforms' CPUs support the fp8 data format.
34
+# EMBED_DTYPE only uses the fp8 data representation,
35
+# does not use fp8 computation, and only occurs on the CPU.
36
+# Apologize for any possible break.
37
+# NOTE: numpy does not support bfloat16 and fp8
38
+EMBED_DTYPES: Mapping[EmbedDType, DTypeInfo] = {
39
+ "float32": DTypeInfo(torch.float32, torch.float32, np.float32),
40
+ "float16": DTypeInfo(torch.float16, torch.float16, np.float16),
41
+ "bfloat16": DTypeInfo(torch.bfloat16, torch.float16, np.float16),
42
+ "fp8_e4m3": DTypeInfo(torch.float8_e4m3fn, torch.uint8, np.uint8),
43
+ "fp8_e5m2": DTypeInfo(torch.float8_e5m2, torch.uint8, np.uint8),
44
+}
45
+ENDIANNESS: tuple[Endianness, ...] = get_args(Endianness)
46
+
47
48
def tensor2base64(x: torch.Tensor) -> str:
49
with io.BytesIO() as buf:
@@ -71,21 +51,26 @@ def tensor2base64(x: torch.Tensor) -> str:
51
buf.seek(0)
52
binary_data = buf.read()
53
74
- return base64.b64encode(binary_data).decode("utf-8")
54
+ return pybase64.b64encode(binary_data).decode("utf-8")
55
56
57
def tensor2binary(
78
- tensor: torch.Tensor, embed_dtype: EmbedDType, endianness: Endianness
58
+ tensor: torch.Tensor,
59
+ embed_dtype: EmbedDType,
60
+ endianness: Endianness,
61
) -> bytes:
62
assert isinstance(tensor, torch.Tensor)
81
- assert embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE
63
+ assert embed_dtype in EMBED_DTYPES
64
assert endianness in ENDIANNESS
65
84
- torch_dtype = EMBED_DTYPE_TO_TORCH_DTYPE[embed_dtype]
85
- torch_view_dtype = EMBED_DTYPE_TO_TORCH_DTYPE_VIEW[embed_dtype]
66
+ dtype_info = EMBED_DTYPES[embed_dtype]
67
68
np_array = (
88
- tensor.to(torch_dtype).flatten().contiguous().view(torch_view_dtype).numpy()
69
+ tensor.to(dtype_info.torch_dtype)
70
+ .flatten()
71
+ .contiguous()
72
+ .view(dtype_info.torch_view_dtype)
73
+ .numpy()
74
)
75
76
if endianness != "native" and endianness != sys_byteorder:
@@ -100,115 +85,14 @@ def binary2tensor(
85
embed_dtype: EmbedDType,
86
endianness: Endianness,
87
) -> torch.Tensor:
103
- assert embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE
104
- assert embed_dtype in EMBED_DTYPE_TO_NUMPY_DTYPE_VIEW
88
+ assert embed_dtype in EMBED_DTYPES
89
assert endianness in ENDIANNESS
90
107
- torch_dtype = EMBED_DTYPE_TO_TORCH_DTYPE[embed_dtype]
108
- np_dtype = EMBED_DTYPE_TO_NUMPY_DTYPE_VIEW[embed_dtype]
91
+ dtype_info = EMBED_DTYPES[embed_dtype]
92
110
- np_array = np.frombuffer(binary, dtype=np_dtype).reshape(shape)
93
+ np_array = np.frombuffer(binary, dtype=dtype_info.numpy_view_dtype).reshape(shape)
94
95
if endianness != "native" and endianness != sys_byteorder:
96
np_array = np_array.byteswap()
97
115
- return torch.from_numpy(np_array).view(torch_dtype)
116
-
117
-
118
-def encode_pooling_output(
119
- output: PoolingRequestOutput,
120
- encoding_format: EncodingFormat,
121
- embed_dtype: EmbedDType,
122
- endianness: Endianness,
123
-) -> list[float] | str | bytes:
124
- if encoding_format == "float":
125
- return output.outputs.data.tolist()
126
- elif encoding_format == "base64":
127
- embedding_bytes = tensor2binary(output.outputs.data, embed_dtype, endianness)
128
- return base64.b64encode(embedding_bytes).decode("utf-8")
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
-
133
-
134
-@dataclass
135
-class MetadataItem:
136
- index: int
137
- embed_dtype: EmbedDType
138
- endianness: Endianness
139
- start: int
140
- end: int
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,
170
- endianness: Endianness,
171
-):
172
- num_prompt_tokens = 0
173
- items: list[dict[str, MetadataItem]] = []
174
- body = []
175
- offset = 0
176
- for idx, output in enumerate(pooling_outputs):
177
- binary = tensor2binary(
178
- tensor=output.outputs.data,
179
- embed_dtype=embed_dtype,
180
- endianness=endianness,
181
- )
182
- size = len(binary)
183
-
184
- item = {
185
- "index": idx,
186
- "embed_dtype": embed_dtype,
187
- "endianness": endianness,
188
- "start": offset,
189
- "end": offset + size,
190
- "shape": output.outputs.data.shape,
191
- }
192
-
193
- body.append(binary)
194
- items.append(item)
195
- prompt_token_ids = output.prompt_token_ids
196
- num_prompt_tokens += len(prompt_token_ids)
197
- offset += size
198
-
199
- usage = {
200
- "prompt_tokens": num_prompt_tokens,
201
- "total_tokens": num_prompt_tokens,
202
- }
203
- return body, items, usage
204
-
205
-
206
-def decode_pooling_output(items: list[MetadataItem], body: bytes) -> list[torch.Tensor]:
207
- items.sort(key=lambda x: x.index)
208
-
209
- tensor_list: list[torch.Tensor] = []
210
- for item in items:
211
- binary = body[item.start : item.end]
212
- tensor = binary2tensor(binary, item.shape, item.embed_dtype, item.endianness)
213
- tensor_list.append(tensor)
214
- return tensor_list
98
+ return torch.from_numpy(np_array).view(dtype_info.torch_dtype)