@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d2a7938582
Fork 0

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 14, 05:30 UTC

[Frontend][1/N] Improve all pooling task | Support FP16 Embedding Base64 (Still uses fp32 by default). (#26414)

Signed-off-by: wang.yuqi <noooop@126.com> Co-authored-by: Maximilien de Bayser <maxdebayser@gmail.com> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>

wang.yuqi committed Oct 14, 2025 at 03:06 UTC d2a7938582011cb0dd117fec468221dc5e3d22a2
8 files changed +312 -30
examples/online_serving/pooling/README.md
+6
@@ -6,6 +6,12 @@
6 python examples/online_serving/pooling/cohere_rerank_client.py
7 ```
8
9 +## Embedding embed_dtype usage
10 +
11 +```bash
12 +python examples/online_serving/pooling/embedding_embed_dtype_client.py
13 +```
14 +
15 ## Jinaai rerank usage
16
17 ```bash
examples/online_serving/pooling/embedding_embed_dtype_client.py new
+59
@@ -0,0 +1,59 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Example Python client for embedding API using vLLM API server
4 +NOTE:
5 + start a supported embeddings model server with `vllm serve`, e.g.
6 + vllm serve intfloat/e5-small
7 +"""
8 +
9 +import argparse
10 +import base64
11 +
12 +import requests
13 +import torch
14 +
15 +from vllm.entrypoints.openai.protocol import EMBED_DTYPE_TO_TORCH_DTYPE
16 +
17 +
18 +def post_http_request(prompt: dict, api_url: str) -> requests.Response:
19 + headers = {"User-Agent": "Test Client"}
20 + response = requests.post(api_url, headers=headers, json=prompt)
21 + return response
22 +
23 +
24 +def parse_args():
25 + parser = argparse.ArgumentParser()
26 + parser.add_argument("--host", type=str, default="localhost")
27 + parser.add_argument("--port", type=int, default=8000)
28 + parser.add_argument("--model", type=str, default="intfloat/e5-small")
29 +
30 + return parser.parse_args()
31 +
32 +
33 +def main(args):
34 + api_url = f"http://{args.host}:{args.port}/v1/embeddings"
35 + model_name = args.model
36 +
37 + for embed_dtype, torch_dtype in EMBED_DTYPE_TO_TORCH_DTYPE.items():
38 + prompt = {
39 + "model": model_name,
40 + "input": "vLLM is great!",
41 + "encoding_format": "base64",
42 + "embed_dtype": embed_dtype,
43 + }
44 + response = post_http_request(prompt=prompt, api_url=api_url)
45 +
46 + embedding = []
47 + for data in response.json()["data"]:
48 + embedding.append(
49 + torch.frombuffer(
50 + base64.b64decode(data["embedding"]), dtype=torch_dtype
51 + ).to(torch.float32)
52 + )
53 + embedding = torch.cat(embedding)
54 + print(embed_dtype, embedding.shape)
55 +
56 +
57 +if __name__ == "__main__":
58 + args = parse_args()
59 + main(args)
tests/entrypoints/pooling/openai/test_embedding.py
+73 -1
@@ -14,7 +14,10 @@ import torch.nn.functional as F
14 from tests.models.language.pooling.embed_utils import run_embedding_correctness_test
15 from tests.models.utils import check_embeddings_close
16 from tests.utils import RemoteOpenAIServer
17 -from vllm.entrypoints.openai.protocol import EmbeddingResponse
17 +from vllm.entrypoints.openai.protocol import (
18 + EMBED_DTYPE_TO_TORCH_DTYPE,
19 + EmbeddingResponse,
20 +)
21 from vllm.transformers_utils.tokenizer import get_tokenizer
22
23 MODEL_NAME = "intfloat/multilingual-e5-small"
@@ -244,6 +247,75 @@ async def test_batch_base64_embedding(
247 run_embedding_correctness_test(hf_model, input_texts, default_data)
248
249
250 +@pytest.mark.asyncio
251 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
252 +async def test_base64_embed_dtype(
253 + hf_model, server: RemoteOpenAIServer, client: openai.AsyncOpenAI, model_name: str
254 +):
255 + input_texts = [
256 + "The best thing about vLLM is that it supports many different models",
257 + ]
258 +
259 + responses_float = await client.embeddings.create(
260 + input=input_texts, model=model_name, encoding_format="float"
261 + )
262 + float_data = [d.embedding for d in responses_float.data]
263 +
264 + for embed_dtype, torch_dtype in EMBED_DTYPE_TO_TORCH_DTYPE.items():
265 + responses_base64 = requests.post(
266 + server.url_for("/v1/embeddings"),
267 + json={
268 + "model": model_name,
269 + "input": input_texts,
270 + "encoding_format": "base64",
271 + "embed_dtype": embed_dtype,
272 + },
273 + )
274 +
275 + base64_data = []
276 + for data in responses_base64.json()["data"]:
277 + base64_data.append(
278 + torch.frombuffer(base64.b64decode(data["embedding"]), dtype=torch_dtype)
279 + .to(torch.float32)
280 + .tolist()
281 + )
282 +
283 + check_embeddings_close(
284 + embeddings_0_lst=float_data,
285 + embeddings_1_lst=base64_data,
286 + name_0="float_data",
287 + name_1="base64_data",
288 + tol=1e-2,
289 + )
290 +
291 +
292 +@pytest.mark.asyncio
293 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
294 +async def test_base64_embed_dtype_not_supported(
295 + hf_model, server: RemoteOpenAIServer, model_name: str
296 +):
297 + input_texts = [
298 + "The best thing about vLLM is that it supports many different models",
299 + ]
300 +
301 + bad_embed_dtype = "bad_embed_dtype"
302 +
303 + responses_base64 = requests.post(
304 + server.url_for("/v1/embeddings"),
305 + json={
306 + "model": model_name,
307 + "input": input_texts,
308 + "encoding_format": "base64",
309 + "embed_dtype": bad_embed_dtype,
310 + },
311 + )
312 +
313 + assert responses_base64.status_code == 400
314 + assert responses_base64.json()["error"]["message"].startswith(
315 + f"embed_dtype={bad_embed_dtype!r} is not supported."
316 + )
317 +
318 +
319 @pytest.mark.asyncio
320 @pytest.mark.parametrize("model_name", [MODEL_NAME])
321 async def test_single_embedding_truncation(client: openai.AsyncOpenAI, model_name: str):
tests/entrypoints/pooling/openai/test_pooling.py
+76 -1
@@ -6,10 +6,11 @@ import base64
6 import numpy as np
7 import pytest
8 import requests
9 +import torch
10
11 from tests.models.utils import check_embeddings_close
12 from tests.utils import RemoteOpenAIServer
12 -from vllm.entrypoints.openai.protocol import PoolingResponse
13 +from vllm.entrypoints.openai.protocol import EMBED_DTYPE_TO_TORCH_DTYPE, PoolingResponse
14 from vllm.transformers_utils.tokenizer import get_tokenizer
15
16 MODEL_NAME = "internlm/internlm2-1_8b-reward"
@@ -248,6 +249,80 @@ async def test_batch_base64_pooling(server: RemoteOpenAIServer, model_name: str)
249 )
250
251
252 +@pytest.mark.asyncio
253 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
254 +async def test_base64_embed_dtype(server: RemoteOpenAIServer, model_name: str):
255 + input_texts = [
256 + "The best thing about vLLM is that it supports many different models",
257 + ]
258 +
259 + url = server.url_for("pooling")
260 + float_response = requests.post(
261 + url,
262 + json={
263 + "model": model_name,
264 + "input": input_texts,
265 + "encoding_format": "float",
266 + },
267 + )
268 + responses_float = PoolingResponse.model_validate(float_response.json())
269 + float_data = [np.array(d.data).squeeze(-1).tolist() for d in responses_float.data]
270 +
271 + for embed_dtype, torch_dtype in EMBED_DTYPE_TO_TORCH_DTYPE.items():
272 + responses_base64 = requests.post(
273 + url,
274 + json={
275 + "model": model_name,
276 + "input": input_texts,
277 + "encoding_format": "base64",
278 + "embed_dtype": embed_dtype,
279 + },
280 + )
281 +
282 + base64_data = []
283 + for data in responses_base64.json()["data"]:
284 + base64_data.append(
285 + torch.frombuffer(base64.b64decode(data["data"]), dtype=torch_dtype)
286 + .to(torch.float32)
287 + .tolist()
288 + )
289 +
290 + check_embeddings_close(
291 + embeddings_0_lst=float_data,
292 + embeddings_1_lst=base64_data,
293 + name_0="float_data",
294 + name_1="base64_data",
295 + tol=1e-2,
296 + )
297 +
298 +
299 +@pytest.mark.asyncio
300 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
301 +async def test_base64_embed_dtype_not_supported(
302 + server: RemoteOpenAIServer, model_name: str
303 +):
304 + input_texts = [
305 + "The best thing about vLLM is that it supports many different models",
306 + ]
307 +
308 + bad_embed_dtype = "bad_embed_dtype"
309 +
310 + responses_base64 = requests.post(
311 + server.url_for("pooling"),
312 + json={
313 + "model": model_name,
314 + "input": input_texts,
315 + "encoding_format": "base64",
316 + "embed_dtype": bad_embed_dtype,
317 + },
318 + )
319 +
320 + assert responses_base64.status_code == 400
321 + assert responses_base64.json()["error"]["message"].startswith(
322 + f"embed_dtype={bad_embed_dtype!r} is not supported."
323 + )
324 +
325 +
326 @pytest.mark.asyncio
327 async def test_invocations(server: RemoteOpenAIServer):
328 input_texts = [
vllm/entrypoints/openai/protocol.py
+42 -3
@@ -83,6 +83,18 @@ from vllm.sampling_params import (
83 )
84 from vllm.utils import random_uuid, resolve_obj_by_qualname
85
86 +EMBED_DTYPE_TO_TORCH_DTYPE = {
87 + "float32": torch.float32,
88 + "float16": torch.float16,
89 + "bfloat16": torch.bfloat16,
90 + # I'm not sure if other platforms' CPUs support the fp8 data format.
91 + # EMBED_DTYPE only uses the fp8 data representation,
92 + # does not use fp8 computation, and only occurs on the CPU.
93 + # Apologize for any possible break.
94 + "fp8_e4m3": torch.float8_e4m3fn,
95 + "fp8_e5m2": torch.float8_e5m2,
96 +}
97 +
98 logger = init_logger(__name__)
99
100 _LONG_INFO = torch.iinfo(torch.long)
@@ -1517,8 +1529,17 @@ class EmbeddingCompletionRequest(OpenAIBaseModel):
1529 "through out the inference process and return in response."
1530 ),
1531 )
1520 - normalize: bool | None = None
1521 -
1532 + normalize: bool | None = Field(
1533 + default=None,
1534 + description="Whether to normalize the embeddings outputs. Default is True.",
1535 + )
1536 + embed_dtype: str = Field(
1537 + default="float32",
1538 + description=(
1539 + "What dtype to use for base64 encoding. Default to using "
1540 + "float32 for base64 encoding to match the OpenAI python client behavior."
1541 + ),
1542 + )
1543 # --8<-- [end:embedding-extra-params]
1544
1545 def to_pooling_params(self):
@@ -1594,7 +1615,17 @@ class EmbeddingChatRequest(OpenAIBaseModel):
1615 "through out the inference process and return in response."
1616 ),
1617 )
1597 - normalize: bool | None = None
1618 + normalize: bool | None = Field(
1619 + default=None,
1620 + description="Whether to normalize the embeddings outputs. Default is True.",
1621 + )
1622 + embed_dtype: str = Field(
1623 + default="float32",
1624 + description=(
1625 + "Which dtype to use for base64 encoding. Defaults to float32 "
1626 + "to match OpenAI API."
1627 + ),
1628 + )
1629 # --8<-- [end:chat-embedding-extra-params]
1630
1631 @model_validator(mode="before")
@@ -1639,6 +1670,14 @@ class IOProcessorRequest(OpenAIBaseModel, Generic[T]):
1670 """
1671 softmax: bool = True
1672
1673 + embed_dtype: str = Field(
1674 + default="float32",
1675 + description=(
1676 + "What dtype to use for base64 encoding. Default to using "
1677 + "float32 for base64 encoding to match the OpenAI python client behavior."
1678 + ),
1679 + )
1680 +
1681 def to_pooling_params(self):
1682 return PoolingParams(task="encode", softmax=self.softmax)
1683
vllm/entrypoints/openai/serving_embedding.py
+12 -24
@@ -1,19 +1,18 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -import base64
4 from collections.abc import AsyncGenerator, Mapping
6 -from typing import Any, Final, Literal, cast
5 +from typing import Any, Final, cast
6
8 -import numpy as np
7 import torch
8 from fastapi import Request
11 -from typing_extensions import assert_never, override
9 +from typing_extensions import override
10
11 from vllm.engine.protocol import EngineClient
12 from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption
13 from vllm.entrypoints.logger import RequestLogger
14 from vllm.entrypoints.openai.protocol import (
15 + EMBED_DTYPE_TO_TORCH_DTYPE,
16 EmbeddingChatRequest,
17 EmbeddingCompletionRequest,
18 EmbeddingRequest,
@@ -29,11 +28,11 @@ from vllm.entrypoints.openai.serving_engine import (
28 TextTokensPrompt,
29 )
30 from vllm.entrypoints.openai.serving_models import OpenAIServingModels
31 +from vllm.entrypoints.openai.utils import encoding_pooling_output
32 from vllm.entrypoints.renderer import RenderConfig
33 from vllm.inputs.data import TokensPrompt as EngineTokensPrompt
34 from vllm.logger import init_logger
35 from vllm.outputs import (
36 - EmbeddingOutput,
36 EmbeddingRequestOutput,
37 PoolingOutput,
38 PoolingRequestOutput,
@@ -45,21 +44,6 @@ from vllm.utils import chunk_list
44 logger = init_logger(__name__)
45
46
48 -def _get_embedding(
49 - output: EmbeddingOutput,
50 - encoding_format: Literal["float", "base64"],
51 -) -> list[float] | str:
52 - if encoding_format == "float":
53 - return output.embedding
54 - elif encoding_format == "base64":
55 - # Force to use float32 for base64 encoding
56 - # to match the OpenAI python client behavior
57 - embedding_bytes = np.array(output.embedding, dtype="float32").tobytes()
58 - return base64.b64encode(embedding_bytes).decode("utf-8")
59 -
60 - assert_never(encoding_format)
61 -
62 -
47 class EmbeddingMixin(OpenAIServing):
48 def __init__(self, *args, **kwargs):
49 super().__init__(*args, **kwargs)
@@ -83,6 +67,12 @@ class EmbeddingMixin(OpenAIServing):
67 ) -> ErrorResponse | None:
68 ctx = cast(EmbeddingServeContext, ctx)
69 try:
70 + if ctx.request.embed_dtype not in EMBED_DTYPE_TO_TORCH_DTYPE:
71 + return self.create_error_response(
72 + f"embed_dtype={ctx.request.embed_dtype!r} is not supported. "
73 + f"Supported types: {EMBED_DTYPE_TO_TORCH_DTYPE.keys()}"
74 + )
75 +
76 ctx.lora_request = self._maybe_get_adapters(ctx.request)
77
78 tokenizer = await self.engine_client.get_tokenizer()
@@ -137,12 +127,10 @@ class EmbeddingMixin(OpenAIServing):
127 final_res_batch_checked = cast(list[PoolingRequestOutput], ctx.final_res_batch)
128
129 for idx, final_res in enumerate(final_res_batch_checked):
140 - embedding_res = EmbeddingRequestOutput.from_base(final_res)
141 -
130 item = EmbeddingResponseData(
131 index=idx,
144 - embedding=_get_embedding(
145 - embedding_res.outputs, ctx.request.encoding_format
132 + embedding=encoding_pooling_output(
133 + final_res, ctx.request.encoding_format, ctx.request.embed_dtype
134 ),
135 )
136 prompt_token_ids = final_res.prompt_token_ids
vllm/entrypoints/openai/serving_pooling.py
+11 -1
@@ -17,6 +17,7 @@ from vllm.engine.protocol import EngineClient
17 from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption
18 from vllm.entrypoints.logger import RequestLogger
19 from vllm.entrypoints.openai.protocol import (
20 + EMBED_DTYPE_TO_TORCH_DTYPE,
21 ErrorResponse,
22 IOProcessorRequest,
23 IOProcessorResponse,
@@ -29,6 +30,7 @@ from vllm.entrypoints.openai.protocol import (
30 )
31 from vllm.entrypoints.openai.serving_engine import OpenAIServing
32 from vllm.entrypoints.openai.serving_models import OpenAIServingModels
33 +from vllm.entrypoints.openai.utils import encoding_pooling_output
34 from vllm.entrypoints.renderer import RenderConfig
35 from vllm.entrypoints.utils import _validate_truncation_size
36 from vllm.logger import init_logger
@@ -90,6 +92,12 @@ class OpenAIServingPooling(OpenAIServing):
92 if error_check_ret is not None:
93 return error_check_ret
94
95 + if request.embed_dtype not in EMBED_DTYPE_TO_TORCH_DTYPE:
96 + return self.create_error_response(
97 + f"embed_dtype={request.embed_dtype!r} is not supported. "
98 + f"Supported types: {EMBED_DTYPE_TO_TORCH_DTYPE.keys()}"
99 + )
100 +
101 model_name = self.models.model_name()
102
103 request_id = f"pool-{self._base_request_id(raw_request)}"
@@ -235,6 +243,7 @@ class OpenAIServingPooling(OpenAIServing):
243 created_time,
244 model_name,
245 request.encoding_format,
246 + request.embed_dtype,
247 )
248 except asyncio.CancelledError:
249 return self.create_error_response("Client disconnected")
@@ -251,6 +260,7 @@ class OpenAIServingPooling(OpenAIServing):
260 created_time: int,
261 model_name: str,
262 encoding_format: Literal["float", "base64"],
263 + embed_dtype: str,
264 ) -> PoolingResponse:
265 items: list[PoolingResponseData] = []
266 num_prompt_tokens = 0
@@ -258,7 +268,7 @@ class OpenAIServingPooling(OpenAIServing):
268 for idx, final_res in enumerate(final_res_batch):
269 item = PoolingResponseData(
270 index=idx,
261 - data=_get_data(final_res.outputs, encoding_format),
271 + data=encoding_pooling_output(final_res, encoding_format, embed_dtype),
272 )
273 prompt_token_ids = final_res.prompt_token_ids
274
vllm/entrypoints/openai/utils.py new
+33
@@ -0,0 +1,33 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import base64
4 +from typing import Literal
5 +
6 +import torch
7 +from typing_extensions import assert_never
8 +
9 +from vllm import PoolingRequestOutput
10 +from vllm.entrypoints.openai.protocol import EMBED_DTYPE_TO_TORCH_DTYPE
11 +
12 +
13 +def encoding_pooling_output(
14 + output: PoolingRequestOutput,
15 + encoding_format: Literal["float", "base64"],
16 + embed_dtype: str,
17 +) -> list[float] | str:
18 + if encoding_format == "float":
19 + return output.outputs.data.tolist()
20 + elif encoding_format == "base64":
21 + assert embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE
22 + torch_dtype = EMBED_DTYPE_TO_TORCH_DTYPE[embed_dtype]
23 + embedding_bytes = (
24 + output.outputs.data.to(torch_dtype)
25 + .flatten()
26 + .contiguous()
27 + .view(torch.uint8)
28 + .numpy()
29 + .tobytes()
30 + )
31 + return base64.b64encode(embedding_bytes).decode("utf-8")
32 +
33 + assert_never(encoding_format)