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 4, 03:45 UTC
[Perf] Use numpy zero-copy path for embedding float response serialization (#41681)
Signed-off-by: Shrinav Loka <lokashrinav@gmail.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: wang.yuqi <yuqi.wang@daocloud.io> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Shrinav Loka committed
May 7, 2026 at 19:42 UTC
cd58e30872c25319e43cc530d124539adfee0c2b
3 files changed
+105
tests/entrypoints/pooling/test_utils.py
new
+59
@@ -0,0 +1,59 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import importlib
5
+import importlib.util
6
+import json
7
+import warnings
8
+from types import SimpleNamespace
9
+
10
+import numpy as np
11
+import pytest
12
+import torch
13
+
14
+from vllm.entrypoints.pooling.utils import encode_pooling_output_float_or_ndarray
15
+
16
+
17
+def _pooling_output(data):
18
+ return SimpleNamespace(outputs=SimpleNamespace(data=data))
19
+
20
+
21
+def test_encode_pooling_output_float_or_ndarray_returns_numpy_array():
22
+ output = _pooling_output(torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32))
23
+
24
+ encoded = encode_pooling_output_float_or_ndarray(output)
25
+
26
+ assert isinstance(encoded, np.ndarray)
27
+ np.testing.assert_allclose(encoded, [1.0, 2.0, 3.0])
28
+
29
+
30
+@pytest.mark.skipif(
31
+ importlib.util.find_spec("orjson") is None,
32
+ reason="orjson is not installed",
33
+)
34
+def test_orjson_serializes_numpy_array():
35
+ from fastapi.responses import ORJSONResponse
36
+
37
+ output = _pooling_output(torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32))
38
+ encoded = encode_pooling_output_float_or_ndarray(output)
39
+
40
+ with warnings.catch_warnings():
41
+ warnings.simplefilter("ignore", DeprecationWarning)
42
+ response = ORJSONResponse(content={"embedding": encoded})
43
+ assert json.loads(response.body)["embedding"] == pytest.approx([1.0, 2.0, 3.0])
44
+
45
+
46
+def test_encode_pooling_output_float_or_ndarray_falls_back_to_list():
47
+ class DataWithUnsupportedNumpy:
48
+ def is_contiguous(self):
49
+ return True
50
+
51
+ def numpy(self):
52
+ raise TypeError("unsupported dtype")
53
+
54
+ def tolist(self):
55
+ return [1.0, 2.0, 3.0]
56
+
57
+ output = _pooling_output(DataWithUnsupportedNumpy())
58
+
59
+ assert encode_pooling_output_float_or_ndarray(output) == [1.0, 2.0, 3.0]
vllm/entrypoints/pooling/embed/serving.py
+35
@@ -19,6 +19,7 @@ from ..utils import (
19
encode_pooling_bytes,
20
encode_pooling_output_base64,
21
encode_pooling_output_float,
22
+ encode_pooling_output_float_or_ndarray,
23
get_json_response_cls,
24
)
25
from .io_processor import EmbedIOProcessor
@@ -104,6 +105,40 @@ class ServingEmbedding(PoolingServing):
105
embed_dtype: EmbedDType,
106
endianness: Endianness,
107
) -> JSONResponse:
108
+ use_ndarray_response = (
109
+ encoding_format == "float"
110
+ and self.json_response_cls.__name__ == "ORJSONResponse"
111
+ )
112
+ if use_ndarray_response:
113
+ ndarray_items: list[dict[str, object]] = []
114
+ ndarray_num_tokens = 0
115
+
116
+ for idx, final_res in enumerate(final_res_batch):
117
+ item_dict = EmbeddingResponseData(
118
+ index=idx,
119
+ embedding=[],
120
+ ).model_dump()
121
+ item_dict["embedding"] = encode_pooling_output_float_or_ndarray(
122
+ final_res
123
+ )
124
+ ndarray_items.append(item_dict)
125
+ ndarray_num_tokens += len(final_res.prompt_token_ids)
126
+
127
+ ndarray_usage = UsageInfo(
128
+ prompt_tokens=ndarray_num_tokens,
129
+ total_tokens=ndarray_num_tokens,
130
+ )
131
+ ndarray_response = EmbeddingResponse(
132
+ id=request_id,
133
+ created=created_time,
134
+ model=model_name,
135
+ data=[], # type: ignore[arg-type]
136
+ usage=ndarray_usage,
137
+ ).model_dump()
138
+ ndarray_response["data"] = ndarray_items
139
+
140
+ return self.json_response_cls(content=ndarray_response)
141
+
142
encode_fn = cast(
143
Callable[[PoolingRequestOutput], list[float] | str],
144
(
vllm/entrypoints/pooling/utils.py
+11
@@ -62,6 +62,17 @@ def encode_pooling_output_float(output: PoolingRequestOutput) -> list[float]:
62
return output.outputs.data.tolist()
63
64
65
+def encode_pooling_output_float_or_ndarray(output: PoolingRequestOutput) -> Any:
66
+ """Return an ndarray when the response renderer can serialize NumPy."""
67
+ try:
68
+ data = output.outputs.data
69
+ if not data.is_contiguous():
70
+ data = data.contiguous()
71
+ return data.numpy()
72
+ except (RuntimeError, TypeError):
73
+ return output.outputs.data.tolist()
74
+
75
+
76
def encode_pooling_output_base64(
77
output: PoolingRequestOutput,
78
embed_dtype: EmbedDType,