[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>

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

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,
-        )

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]:

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