From 503697c9ceb5200cd2731dbcfbdc539fa142aafa Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Wed, 13 May 2026 02:08:01 -0400 Subject: [PATCH] [chore] Refactor pooling metadata token ID accessors (#42368) Signed-off-by: Taneem Ibrahim 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> --- vllm/entrypoints/pooling/embed/serving.py | 90 +++++------------ vllm/entrypoints/pooling/pooling/serving.py | 76 ++++----------- vllm/entrypoints/pooling/utils.py | 103 +++++++++++++++++--- vllm/v1/pool/metadata.py | 15 +-- 4 files changed, 142 insertions(+), 142 deletions(-) diff --git a/vllm/entrypoints/pooling/embed/serving.py b/vllm/entrypoints/pooling/embed/serving.py index d85d237238..5d9616f00c 100644 --- a/vllm/entrypoints/pooling/embed/serving.py +++ b/vllm/entrypoints/pooling/embed/serving.py @@ -1,14 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import json -from collections.abc import Callable -from functools import partial -from typing import Literal, TypeAlias, cast +from typing import TypeAlias, cast from fastapi.responses import JSONResponse, Response, StreamingResponse from typing_extensions import assert_never -from vllm.entrypoints.openai.engine.protocol import UsageInfo from vllm.logger import init_logger from vllm.outputs import PoolingRequestOutput from vllm.utils.serial_utils import EmbedDType, Endianness @@ -16,11 +12,14 @@ from vllm.utils.serial_utils import EmbedDType, Endianness from ..base.serving import PoolingServing from ..typing import PoolingServeContext from ..utils import ( - encode_pooling_bytes, - encode_pooling_output_base64, + BytesEncodingFormat, + JsonEncodingFormat, + build_pooling_bytes_streaming_response, encode_pooling_output_float, encode_pooling_output_float_or_ndarray, get_json_response_cls, + get_pooling_output_encoder, + get_pooling_usage, ) from .io_processor import EmbedIOProcessor from .protocol import ( @@ -28,7 +27,6 @@ from .protocol import ( CohereEmbedRequest, CohereEmbedResponse, CohereMeta, - EmbeddingBytesResponse, EmbeddingRequest, EmbeddingResponse, EmbeddingResponseData, @@ -101,7 +99,7 @@ class ServingEmbedding(PoolingServing): request_id: str, created_time: int, model_name: str, - encoding_format: Literal["float", "base64"], + encoding_format: JsonEncodingFormat, embed_dtype: EmbedDType, endianness: Endianness, ) -> JSONResponse: @@ -111,7 +109,6 @@ class ServingEmbedding(PoolingServing): ) if use_ndarray_response: ndarray_items: list[dict[str, object]] = [] - ndarray_num_tokens = 0 for idx, final_res in enumerate(final_res_batch): item_dict = EmbeddingResponseData( @@ -122,60 +119,39 @@ class ServingEmbedding(PoolingServing): final_res ) ndarray_items.append(item_dict) - ndarray_num_tokens += len(final_res.prompt_token_ids) - - ndarray_usage = UsageInfo( - prompt_tokens=ndarray_num_tokens, - total_tokens=ndarray_num_tokens, - ) ndarray_response = EmbeddingResponse( id=request_id, created=created_time, model=model_name, data=[], # type: ignore[arg-type] - usage=ndarray_usage, + usage=get_pooling_usage(final_res_batch), ).model_dump() ndarray_response["data"] = ndarray_items return self.json_response_cls(content=ndarray_response) - encode_fn = cast( - Callable[[PoolingRequestOutput], list[float] | str], - ( - encode_pooling_output_float - if encoding_format == "float" - else partial( - encode_pooling_output_base64, - embed_dtype=embed_dtype, - endianness=endianness, - ) - ), + encode_fn = get_pooling_output_encoder( + encoding_format=encoding_format, + embed_dtype=embed_dtype, + endianness=endianness, ) items: list[EmbeddingResponseData] = [] - num_prompt_tokens = 0 for idx, final_res in enumerate(final_res_batch): item = EmbeddingResponseData( index=idx, - embedding=encode_fn(final_res), + embedding=cast(list[float] | str, encode_fn(final_res)), ) - prompt_token_ids = final_res.prompt_token_ids items.append(item) - num_prompt_tokens += len(prompt_token_ids) - - usage = UsageInfo( - prompt_tokens=num_prompt_tokens, - total_tokens=num_prompt_tokens, - ) response = EmbeddingResponse( id=request_id, created=created_time, model=model_name, data=items, - usage=usage, + usage=get_pooling_usage(final_res_batch), ) return self.json_response_cls(content=response.model_dump()) @@ -185,39 +161,20 @@ class ServingEmbedding(PoolingServing): request_id: str, created_time: int, model_name: str, - encoding_format: Literal["bytes", "bytes_only"], + encoding_format: BytesEncodingFormat, embed_dtype: EmbedDType, endianness: Endianness, ) -> StreamingResponse: - content, items, usage = encode_pooling_bytes( + return build_pooling_bytes_streaming_response( pooling_outputs=final_res_batch, + request_id=request_id, + created_time=created_time, + model_name=model_name, + encoding_format=encoding_format, embed_dtype=embed_dtype, endianness=endianness, ) - headers = ( - None - if encoding_format == "bytes_only" - else { - "metadata": json.dumps( - { - "id": request_id, - "created": created_time, - "model": model_name, - "data": items, - "usage": usage, - } - ) - } - ) - - response = EmbeddingBytesResponse(content=content, headers=headers) - return StreamingResponse( - content=response.content, - headers=response.headers, - media_type=response.media_type, - ) - def _build_cohere_response_from_ctx( self, ctx: PoolingServeContext, @@ -225,8 +182,11 @@ class ServingEmbedding(PoolingServing): request = ctx.request assert isinstance(request, CohereEmbedRequest) - all_floats = [encode_pooling_output_float(out) for out in ctx.final_res_batch] - total_tokens = sum(len(out.prompt_token_ids) for out in ctx.final_res_batch) + all_floats = [ + cast(list[float], encode_pooling_output_float(out)) + for out in ctx.final_res_batch + ] + total_tokens = get_pooling_usage(ctx.final_res_batch).prompt_tokens image_tokens = total_tokens if request.images is not None else 0 texts_echo = request.texts diff --git a/vllm/entrypoints/pooling/pooling/serving.py b/vllm/entrypoints/pooling/pooling/serving.py index 854fefc88a..1783c861fc 100644 --- a/vllm/entrypoints/pooling/pooling/serving.py +++ b/vllm/entrypoints/pooling/pooling/serving.py @@ -1,16 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project - - -import json -from collections.abc import Callable -from functools import partial -from typing import Literal, cast - from fastapi.responses import JSONResponse, Response, StreamingResponse from typing_extensions import assert_never -from vllm.entrypoints.openai.engine.protocol import UsageInfo from vllm.logger import init_logger from vllm.outputs import PoolingRequestOutput from vllm.tasks import SupportedTask @@ -21,14 +13,15 @@ from ..base.serving import PoolingServingBase from ..factories import init_pooling_io_processors from ..typing import AnyPoolingRequest, PoolingServeContext from ..utils import ( - encode_pooling_bytes, - encode_pooling_output_base64, - encode_pooling_output_float, + BytesEncodingFormat, + JsonEncodingFormat, + build_pooling_bytes_streaming_response, get_json_response_cls, + get_pooling_output_encoder, + get_pooling_usage, ) from .protocol import ( IOProcessorRequest, - PoolingBytesResponse, PoolingRequest, PoolingResponse, PoolingResponseData, @@ -141,47 +134,32 @@ class ServingPooling(PoolingServingBase): request_id: str, created_time: int, model_name: str, - encoding_format: Literal["float", "base64"], + encoding_format: JsonEncodingFormat, embed_dtype: EmbedDType, endianness: Endianness, ) -> JSONResponse: - encode_fn = cast( - Callable[[PoolingRequestOutput], list[float] | str], - ( - encode_pooling_output_float - if encoding_format == "float" - else partial( - encode_pooling_output_base64, - embed_dtype=embed_dtype, - endianness=endianness, - ) - ), + encode_fn = get_pooling_output_encoder( + encoding_format=encoding_format, + embed_dtype=embed_dtype, + endianness=endianness, ) items: list[PoolingResponseData] = [] - num_prompt_tokens = 0 for idx, final_res in enumerate(final_res_batch): item = PoolingResponseData( index=idx, data=encode_fn(final_res), ) - prompt_token_ids = final_res.prompt_token_ids items.append(item) - num_prompt_tokens += len(prompt_token_ids) - - usage = UsageInfo( - prompt_tokens=num_prompt_tokens, - total_tokens=num_prompt_tokens, - ) response = PoolingResponse( id=request_id, created=created_time, model=model_name, data=items, - usage=usage, + usage=get_pooling_usage(final_res_batch), ) return self.json_response_cls(content=response.model_dump()) @@ -191,36 +169,16 @@ class ServingPooling(PoolingServingBase): request_id: str, created_time: int, model_name: str, - encoding_format: Literal["bytes", "bytes_only"], + encoding_format: BytesEncodingFormat, embed_dtype: EmbedDType, endianness: Endianness, ) -> StreamingResponse: - content, items, usage = encode_pooling_bytes( + return build_pooling_bytes_streaming_response( pooling_outputs=final_res_batch, + request_id=request_id, + created_time=created_time, + model_name=model_name, + encoding_format=encoding_format, embed_dtype=embed_dtype, endianness=endianness, ) - - headers = ( - None - if encoding_format == "bytes_only" - else { - "metadata": json.dumps( - { - "id": request_id, - "created": created_time, - "model": model_name, - "data": items, - "usage": usage, - } - ) - } - ) - - response = PoolingBytesResponse(content=content, headers=headers) - - return StreamingResponse( - content=response.content, - headers=response.headers, - media_type=response.media_type, - ) diff --git a/vllm/entrypoints/pooling/utils.py b/vllm/entrypoints/pooling/utils.py index 2036ed7aae..b6d91332ef 100644 --- a/vllm/entrypoints/pooling/utils.py +++ b/vllm/entrypoints/pooling/utils.py @@ -2,16 +2,19 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import importlib.util +import json import math +from collections.abc import Callable, Sequence from dataclasses import dataclass -from functools import lru_cache -from typing import Any +from functools import lru_cache, partial +from typing import Any, Literal, cast import pybase64 import torch -from fastapi.responses import JSONResponse +from fastapi.responses import JSONResponse, StreamingResponse from vllm.config import ModelConfig +from vllm.entrypoints.openai.engine.protocol import UsageInfo from vllm.logger import init_logger from vllm.outputs import PoolingRequestOutput from vllm.tasks import SupportedTask @@ -25,6 +28,11 @@ from vllm.utils.serial_utils import ( logger = init_logger(__name__) +JsonEncodingFormat = Literal["float", "base64"] +BytesEncodingFormat = Literal["bytes", "bytes_only"] +FloatEncodedPoolingOutput = list[float] | list[list[float]] +JsonEncodedPoolingOutput = FloatEncodedPoolingOutput | str + @dataclass class MetadataItem: @@ -58,7 +66,9 @@ def build_metadata_items( ] -def encode_pooling_output_float(output: PoolingRequestOutput) -> list[float]: +def encode_pooling_output_float( + output: PoolingRequestOutput, +) -> FloatEncodedPoolingOutput: return output.outputs.data.tolist() @@ -86,8 +96,7 @@ def encode_pooling_bytes( pooling_outputs: list[PoolingRequestOutput], embed_dtype: EmbedDType, endianness: Endianness, -) -> tuple[list[bytes], list[dict[str, Any]], dict[str, Any]]: - num_prompt_tokens = 0 +) -> tuple[list[bytes], list[dict[str, Any]]]: items: list[dict[str, Any]] = [] body: list[bytes] = [] offset = 0 @@ -111,17 +120,89 @@ def encode_pooling_bytes( body.append(binary) items.append(item) - prompt_token_ids = output.prompt_token_ids - num_prompt_tokens += len(prompt_token_ids) offset += size - # Dictionary form of UsageInfo - usage = dict( + return body, items + + +def get_pooling_output_encoder( + encoding_format: JsonEncodingFormat, + embed_dtype: EmbedDType, + endianness: Endianness, +) -> Callable[[PoolingRequestOutput], JsonEncodedPoolingOutput]: + return cast( + Callable[[PoolingRequestOutput], JsonEncodedPoolingOutput], + ( + encode_pooling_output_float + if encoding_format == "float" + else partial( + encode_pooling_output_base64, + embed_dtype=embed_dtype, + endianness=endianness, + ) + ), + ) + + +def get_pooling_usage( + pooling_outputs: Sequence[PoolingRequestOutput], +) -> UsageInfo: + num_prompt_tokens = sum( + len(output.prompt_token_ids) if output.prompt_token_ids is not None else 0 + for output in pooling_outputs + ) + return UsageInfo( prompt_tokens=num_prompt_tokens, total_tokens=num_prompt_tokens, ) - return body, items, usage + +def get_pooling_usage_payload( + pooling_outputs: Sequence[PoolingRequestOutput], +) -> dict[str, int]: + usage = get_pooling_usage(pooling_outputs) + return { + "prompt_tokens": usage.prompt_tokens, + "total_tokens": usage.total_tokens, + } + + +def build_pooling_bytes_streaming_response( + pooling_outputs: list[PoolingRequestOutput], + request_id: str, + created_time: int, + model_name: str, + encoding_format: BytesEncodingFormat, + embed_dtype: EmbedDType, + endianness: Endianness, +) -> StreamingResponse: + content, items = encode_pooling_bytes( + pooling_outputs=pooling_outputs, + embed_dtype=embed_dtype, + endianness=endianness, + ) + + headers = ( + None + if encoding_format == "bytes_only" + else { + "metadata": json.dumps( + { + "id": request_id, + "created": created_time, + "model": model_name, + "data": items, + "usage": get_pooling_usage_payload(pooling_outputs), + } + ) + } + ) + + return StreamingResponse( + content=content, + headers=headers, + media_type="application/octet-stream", + ) def decode_pooling_output(items: list[MetadataItem], body: bytes) -> list[torch.Tensor]: diff --git a/vllm/v1/pool/metadata.py b/vllm/v1/pool/metadata.py index 076c87526f..ac98f6fe0d 100644 --- a/vllm/v1/pool/metadata.py +++ b/vllm/v1/pool/metadata.py @@ -84,19 +84,20 @@ class PoolingMetadata: else self.pooling_cursor[indices], ) - def get_prompt_token_ids(self) -> list[torch.Tensor]: - prompt_token_ids = self.prompt_token_ids + def _get_prompt_token_ids( + self, + prompt_token_ids: torch.Tensor | None, + ) -> list[torch.Tensor]: assert prompt_token_ids is not None, ( "Please set `requires_token_ids=True` in `get_pooling_updates`" ) return [prompt_token_ids[i, :num] for i, num in enumerate(self.prompt_lens)] + def get_prompt_token_ids(self) -> list[torch.Tensor]: + return self._get_prompt_token_ids(self.prompt_token_ids) + def get_prompt_token_ids_cpu(self) -> list[torch.Tensor]: - prompt_token_ids = self.prompt_token_ids_cpu - assert prompt_token_ids is not None, ( - "Please set `requires_token_ids=True` in `get_pooling_updates`" - ) - return [prompt_token_ids[i, :num] for i, num in enumerate(self.prompt_lens)] + return self._get_prompt_token_ids(self.prompt_token_ids_cpu) def get_pooling_cursor(self) -> PoolingCursor: pooling_cursor = self.pooling_cursor