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 3, 19:00 UTC
[chore] Refactor pooling metadata token ID accessors (#42368)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Taneem Ibrahim committed
May 13, 2026 at 02:08 UTC
503697c9ceb5200cd2731dbcfbdc539fa142aafa
4 files changed
+142
-142
vllm/entrypoints/pooling/embed/serving.py
+25
-65
@@ -1,14 +1,10 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import json
4
-from collections.abc import Callable
5
-from functools import partial
6
-from typing import Literal, TypeAlias, cast
3
+from typing import TypeAlias, cast
4
5
from fastapi.responses import JSONResponse, Response, StreamingResponse
6
from typing_extensions import assert_never
7
11
-from vllm.entrypoints.openai.engine.protocol import UsageInfo
8
from vllm.logger import init_logger
9
from vllm.outputs import PoolingRequestOutput
10
from vllm.utils.serial_utils import EmbedDType, Endianness
@@ -16,11 +12,14 @@ from vllm.utils.serial_utils import EmbedDType, Endianness
12
from ..base.serving import PoolingServing
13
from ..typing import PoolingServeContext
14
from ..utils import (
19
- encode_pooling_bytes,
20
- encode_pooling_output_base64,
15
+ BytesEncodingFormat,
16
+ JsonEncodingFormat,
17
+ build_pooling_bytes_streaming_response,
18
encode_pooling_output_float,
19
encode_pooling_output_float_or_ndarray,
20
get_json_response_cls,
21
+ get_pooling_output_encoder,
22
+ get_pooling_usage,
23
)
24
from .io_processor import EmbedIOProcessor
25
from .protocol import (
@@ -28,7 +27,6 @@ from .protocol import (
27
CohereEmbedRequest,
28
CohereEmbedResponse,
29
CohereMeta,
31
- EmbeddingBytesResponse,
30
EmbeddingRequest,
31
EmbeddingResponse,
32
EmbeddingResponseData,
@@ -101,7 +99,7 @@ class ServingEmbedding(PoolingServing):
99
request_id: str,
100
created_time: int,
101
model_name: str,
104
- encoding_format: Literal["float", "base64"],
102
+ encoding_format: JsonEncodingFormat,
103
embed_dtype: EmbedDType,
104
endianness: Endianness,
105
) -> JSONResponse:
@@ -111,7 +109,6 @@ class ServingEmbedding(PoolingServing):
109
)
110
if use_ndarray_response:
111
ndarray_items: list[dict[str, object]] = []
114
- ndarray_num_tokens = 0
112
113
for idx, final_res in enumerate(final_res_batch):
114
item_dict = EmbeddingResponseData(
@@ -122,60 +119,39 @@ class ServingEmbedding(PoolingServing):
119
final_res
120
)
121
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
- )
122
ndarray_response = EmbeddingResponse(
123
id=request_id,
124
created=created_time,
125
model=model_name,
126
data=[], # type: ignore[arg-type]
136
- usage=ndarray_usage,
127
+ usage=get_pooling_usage(final_res_batch),
128
).model_dump()
129
ndarray_response["data"] = ndarray_items
130
131
return self.json_response_cls(content=ndarray_response)
132
142
- encode_fn = cast(
143
- Callable[[PoolingRequestOutput], list[float] | str],
144
- (
145
- encode_pooling_output_float
146
- if encoding_format == "float"
147
- else partial(
148
- encode_pooling_output_base64,
149
- embed_dtype=embed_dtype,
150
- endianness=endianness,
151
- )
152
- ),
133
+ encode_fn = get_pooling_output_encoder(
134
+ encoding_format=encoding_format,
135
+ embed_dtype=embed_dtype,
136
+ endianness=endianness,
137
)
138
139
items: list[EmbeddingResponseData] = []
156
- num_prompt_tokens = 0
140
141
for idx, final_res in enumerate(final_res_batch):
142
item = EmbeddingResponseData(
143
index=idx,
161
- embedding=encode_fn(final_res),
144
+ embedding=cast(list[float] | str, encode_fn(final_res)),
145
)
163
- prompt_token_ids = final_res.prompt_token_ids
146
147
items.append(item)
166
- num_prompt_tokens += len(prompt_token_ids)
167
-
168
- usage = UsageInfo(
169
- prompt_tokens=num_prompt_tokens,
170
- total_tokens=num_prompt_tokens,
171
- )
148
149
response = EmbeddingResponse(
150
id=request_id,
151
created=created_time,
152
model=model_name,
153
data=items,
178
- usage=usage,
154
+ usage=get_pooling_usage(final_res_batch),
155
)
156
return self.json_response_cls(content=response.model_dump())
157
@@ -185,39 +161,20 @@ class ServingEmbedding(PoolingServing):
161
request_id: str,
162
created_time: int,
163
model_name: str,
188
- encoding_format: Literal["bytes", "bytes_only"],
164
+ encoding_format: BytesEncodingFormat,
165
embed_dtype: EmbedDType,
166
endianness: Endianness,
167
) -> StreamingResponse:
192
- content, items, usage = encode_pooling_bytes(
168
+ return build_pooling_bytes_streaming_response(
169
pooling_outputs=final_res_batch,
170
+ request_id=request_id,
171
+ created_time=created_time,
172
+ model_name=model_name,
173
+ encoding_format=encoding_format,
174
embed_dtype=embed_dtype,
175
endianness=endianness,
176
)
177
198
- headers = (
199
- None
200
- if encoding_format == "bytes_only"
201
- else {
202
- "metadata": json.dumps(
203
- {
204
- "id": request_id,
205
- "created": created_time,
206
- "model": model_name,
207
- "data": items,
208
- "usage": usage,
209
- }
210
- )
211
- }
212
- )
213
-
214
- response = EmbeddingBytesResponse(content=content, headers=headers)
215
- return StreamingResponse(
216
- content=response.content,
217
- headers=response.headers,
218
- media_type=response.media_type,
219
- )
220
-
178
def _build_cohere_response_from_ctx(
179
self,
180
ctx: PoolingServeContext,
@@ -225,8 +182,11 @@ class ServingEmbedding(PoolingServing):
182
request = ctx.request
183
assert isinstance(request, CohereEmbedRequest)
184
228
- all_floats = [encode_pooling_output_float(out) for out in ctx.final_res_batch]
229
- total_tokens = sum(len(out.prompt_token_ids) for out in ctx.final_res_batch)
185
+ all_floats = [
186
+ cast(list[float], encode_pooling_output_float(out))
187
+ for out in ctx.final_res_batch
188
+ ]
189
+ total_tokens = get_pooling_usage(ctx.final_res_batch).prompt_tokens
190
191
image_tokens = total_tokens if request.images is not None else 0
192
texts_echo = request.texts
vllm/entrypoints/pooling/pooling/serving.py
+17
-59
@@ -1,16 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-
5
-import json
6
-from collections.abc import Callable
7
-from functools import partial
8
-from typing import Literal, cast
9
-
3
from fastapi.responses import JSONResponse, Response, StreamingResponse
4
from typing_extensions import assert_never
5
13
-from vllm.entrypoints.openai.engine.protocol import UsageInfo
6
from vllm.logger import init_logger
7
from vllm.outputs import PoolingRequestOutput
8
from vllm.tasks import SupportedTask
@@ -21,14 +13,15 @@ from ..base.serving import PoolingServingBase
13
from ..factories import init_pooling_io_processors
14
from ..typing import AnyPoolingRequest, PoolingServeContext
15
from ..utils import (
24
- encode_pooling_bytes,
25
- encode_pooling_output_base64,
26
- encode_pooling_output_float,
16
+ BytesEncodingFormat,
17
+ JsonEncodingFormat,
18
+ build_pooling_bytes_streaming_response,
19
get_json_response_cls,
20
+ get_pooling_output_encoder,
21
+ get_pooling_usage,
22
)
23
from .protocol import (
24
IOProcessorRequest,
31
- PoolingBytesResponse,
25
PoolingRequest,
26
PoolingResponse,
27
PoolingResponseData,
@@ -141,47 +134,32 @@ class ServingPooling(PoolingServingBase):
134
request_id: str,
135
created_time: int,
136
model_name: str,
144
- encoding_format: Literal["float", "base64"],
137
+ encoding_format: JsonEncodingFormat,
138
embed_dtype: EmbedDType,
139
endianness: Endianness,
140
) -> JSONResponse:
148
- encode_fn = cast(
149
- Callable[[PoolingRequestOutput], list[float] | str],
150
- (
151
- encode_pooling_output_float
152
- if encoding_format == "float"
153
- else partial(
154
- encode_pooling_output_base64,
155
- embed_dtype=embed_dtype,
156
- endianness=endianness,
157
- )
158
- ),
141
+ encode_fn = get_pooling_output_encoder(
142
+ encoding_format=encoding_format,
143
+ embed_dtype=embed_dtype,
144
+ endianness=endianness,
145
)
146
147
items: list[PoolingResponseData] = []
162
- num_prompt_tokens = 0
148
149
for idx, final_res in enumerate(final_res_batch):
150
item = PoolingResponseData(
151
index=idx,
152
data=encode_fn(final_res),
153
)
169
- prompt_token_ids = final_res.prompt_token_ids
154
155
items.append(item)
172
- num_prompt_tokens += len(prompt_token_ids)
173
-
174
- usage = UsageInfo(
175
- prompt_tokens=num_prompt_tokens,
176
- total_tokens=num_prompt_tokens,
177
- )
156
157
response = PoolingResponse(
158
id=request_id,
159
created=created_time,
160
model=model_name,
161
data=items,
184
- usage=usage,
162
+ usage=get_pooling_usage(final_res_batch),
163
)
164
return self.json_response_cls(content=response.model_dump())
165
@@ -191,36 +169,16 @@ class ServingPooling(PoolingServingBase):
169
request_id: str,
170
created_time: int,
171
model_name: str,
194
- encoding_format: Literal["bytes", "bytes_only"],
172
+ encoding_format: BytesEncodingFormat,
173
embed_dtype: EmbedDType,
174
endianness: Endianness,
175
) -> StreamingResponse:
198
- content, items, usage = encode_pooling_bytes(
176
+ return build_pooling_bytes_streaming_response(
177
pooling_outputs=final_res_batch,
178
+ request_id=request_id,
179
+ created_time=created_time,
180
+ model_name=model_name,
181
+ encoding_format=encoding_format,
182
embed_dtype=embed_dtype,
183
endianness=endianness,
184
)
203
-
204
- headers = (
205
- None
206
- if encoding_format == "bytes_only"
207
- else {
208
- "metadata": json.dumps(
209
- {
210
- "id": request_id,
211
- "created": created_time,
212
- "model": model_name,
213
- "data": items,
214
- "usage": usage,
215
- }
216
- )
217
- }
218
- )
219
-
220
- response = PoolingBytesResponse(content=content, headers=headers)
221
-
222
- return StreamingResponse(
223
- content=response.content,
224
- headers=response.headers,
225
- media_type=response.media_type,
226
- )
vllm/entrypoints/pooling/utils.py
+92
-11
@@ -2,16 +2,19 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import importlib.util
5
+import json
6
import math
7
+from collections.abc import Callable, Sequence
8
from dataclasses import dataclass
7
-from functools import lru_cache
8
-from typing import Any
9
+from functools import lru_cache, partial
10
+from typing import Any, Literal, cast
11
12
import pybase64
13
import torch
12
-from fastapi.responses import JSONResponse
14
+from fastapi.responses import JSONResponse, StreamingResponse
15
16
from vllm.config import ModelConfig
17
+from vllm.entrypoints.openai.engine.protocol import UsageInfo
18
from vllm.logger import init_logger
19
from vllm.outputs import PoolingRequestOutput
20
from vllm.tasks import SupportedTask
@@ -25,6 +28,11 @@ from vllm.utils.serial_utils import (
28
29
logger = init_logger(__name__)
30
31
+JsonEncodingFormat = Literal["float", "base64"]
32
+BytesEncodingFormat = Literal["bytes", "bytes_only"]
33
+FloatEncodedPoolingOutput = list[float] | list[list[float]]
34
+JsonEncodedPoolingOutput = FloatEncodedPoolingOutput | str
35
+
36
37
@dataclass
38
class MetadataItem:
@@ -58,7 +66,9 @@ def build_metadata_items(
66
]
67
68
61
-def encode_pooling_output_float(output: PoolingRequestOutput) -> list[float]:
69
+def encode_pooling_output_float(
70
+ output: PoolingRequestOutput,
71
+) -> FloatEncodedPoolingOutput:
72
return output.outputs.data.tolist()
73
74
@@ -86,8 +96,7 @@ def encode_pooling_bytes(
96
pooling_outputs: list[PoolingRequestOutput],
97
embed_dtype: EmbedDType,
98
endianness: Endianness,
89
-) -> tuple[list[bytes], list[dict[str, Any]], dict[str, Any]]:
90
- num_prompt_tokens = 0
99
+) -> tuple[list[bytes], list[dict[str, Any]]]:
100
items: list[dict[str, Any]] = []
101
body: list[bytes] = []
102
offset = 0
@@ -111,17 +120,89 @@ def encode_pooling_bytes(
120
121
body.append(binary)
122
items.append(item)
114
- prompt_token_ids = output.prompt_token_ids
115
- num_prompt_tokens += len(prompt_token_ids)
123
offset += size
124
118
- # Dictionary form of UsageInfo
119
- usage = dict(
125
+ return body, items
126
+
127
+
128
+def get_pooling_output_encoder(
129
+ encoding_format: JsonEncodingFormat,
130
+ embed_dtype: EmbedDType,
131
+ endianness: Endianness,
132
+) -> Callable[[PoolingRequestOutput], JsonEncodedPoolingOutput]:
133
+ return cast(
134
+ Callable[[PoolingRequestOutput], JsonEncodedPoolingOutput],
135
+ (
136
+ encode_pooling_output_float
137
+ if encoding_format == "float"
138
+ else partial(
139
+ encode_pooling_output_base64,
140
+ embed_dtype=embed_dtype,
141
+ endianness=endianness,
142
+ )
143
+ ),
144
+ )
145
+
146
+
147
+def get_pooling_usage(
148
+ pooling_outputs: Sequence[PoolingRequestOutput],
149
+) -> UsageInfo:
150
+ num_prompt_tokens = sum(
151
+ len(output.prompt_token_ids) if output.prompt_token_ids is not None else 0
152
+ for output in pooling_outputs
153
+ )
154
+ return UsageInfo(
155
prompt_tokens=num_prompt_tokens,
156
total_tokens=num_prompt_tokens,
157
)
158
124
- return body, items, usage
159
+
160
+def get_pooling_usage_payload(
161
+ pooling_outputs: Sequence[PoolingRequestOutput],
162
+) -> dict[str, int]:
163
+ usage = get_pooling_usage(pooling_outputs)
164
+ return {
165
+ "prompt_tokens": usage.prompt_tokens,
166
+ "total_tokens": usage.total_tokens,
167
+ }
168
+
169
+
170
+def build_pooling_bytes_streaming_response(
171
+ pooling_outputs: list[PoolingRequestOutput],
172
+ request_id: str,
173
+ created_time: int,
174
+ model_name: str,
175
+ encoding_format: BytesEncodingFormat,
176
+ embed_dtype: EmbedDType,
177
+ endianness: Endianness,
178
+) -> StreamingResponse:
179
+ content, items = encode_pooling_bytes(
180
+ pooling_outputs=pooling_outputs,
181
+ embed_dtype=embed_dtype,
182
+ endianness=endianness,
183
+ )
184
+
185
+ headers = (
186
+ None
187
+ if encoding_format == "bytes_only"
188
+ else {
189
+ "metadata": json.dumps(
190
+ {
191
+ "id": request_id,
192
+ "created": created_time,
193
+ "model": model_name,
194
+ "data": items,
195
+ "usage": get_pooling_usage_payload(pooling_outputs),
196
+ }
197
+ )
198
+ }
199
+ )
200
+
201
+ return StreamingResponse(
202
+ content=content,
203
+ headers=headers,
204
+ media_type="application/octet-stream",
205
+ )
206
207
208
def decode_pooling_output(items: list[MetadataItem], body: bytes) -> list[torch.Tensor]:
vllm/v1/pool/metadata.py
+8
-7
@@ -84,19 +84,20 @@ class PoolingMetadata:
84
else self.pooling_cursor[indices],
85
)
86
87
- def get_prompt_token_ids(self) -> list[torch.Tensor]:
88
- prompt_token_ids = self.prompt_token_ids
87
+ def _get_prompt_token_ids(
88
+ self,
89
+ prompt_token_ids: torch.Tensor | None,
90
+ ) -> list[torch.Tensor]:
91
assert prompt_token_ids is not None, (
92
"Please set `requires_token_ids=True` in `get_pooling_updates`"
93
)
94
return [prompt_token_ids[i, :num] for i, num in enumerate(self.prompt_lens)]
95
96
+ def get_prompt_token_ids(self) -> list[torch.Tensor]:
97
+ return self._get_prompt_token_ids(self.prompt_token_ids)
98
+
99
def get_prompt_token_ids_cpu(self) -> list[torch.Tensor]:
95
- prompt_token_ids = self.prompt_token_ids_cpu
96
- assert prompt_token_ids is not None, (
97
- "Please set `requires_token_ids=True` in `get_pooling_updates`"
98
- )
99
- return [prompt_token_ids[i, :num] for i, num in enumerate(self.prompt_lens)]
100
+ return self._get_prompt_token_ids(self.prompt_token_ids_cpu)
101
102
def get_pooling_cursor(self) -> PoolingCursor:
103
pooling_cursor = self.pooling_cursor