diff --git a/components/src/dynamo/vllm/handlers.py b/components/src/dynamo/vllm/handlers.py index a7b534cf4f9a..d1d32c6240a9 100644 --- a/components/src/dynamo/vllm/handlers.py +++ b/components/src/dynamo/vllm/handlers.py @@ -3035,29 +3035,23 @@ async def _encode_one(idx: int, prompt: Any): if pending: await asyncio.gather(*pending, return_exceptions=True) + # Always emit base64 over the worker->frontend wire format. The Rust + # frontend decodes back to float when the client's ``encoding_format`` + # is float (or unset); base64 bytes are ~3x smaller and far faster to + # (de)serialize than a JSON float array, and the client-visible wire + # format is preserved because Rust converts at the HTTP boundary. The + # base64 is built straight from the pooling tensor (torch -> + # numpy.tobytes), skipping the per-embedding Python float list + + # struct.pack varargs. Bytes are identical on little-endian hosts. embedding_objects: list[Dict[str, Any]] = [] prompt_tokens = 0 for idx, final_output in enumerate(outputs): - embedding = _pooling_output_to_list(final_output.outputs.data) - if dimensions is not None: - if dimensions > len(embedding): - raise ValueError( - f"dimensions={dimensions} exceeds model embedding " - f"dimension {len(embedding)}" - ) - embedding = embedding[:dimensions] - - # Always emit base64 over the worker->frontend wire format. The - # Rust frontend decodes back to float when the client's - # ``encoding_format`` is float (or unset). 15x1024-float responses - # serialized as JSON arrays cost ~110 ms in Python json.dumps + - # Rust serde parse; base64 bytes are ~3x smaller and ~10x faster - # to (de)serialize. Client-visible wire format is preserved - # because Rust converts at the HTTP boundary. embedding_objects.append( { "object": "embedding", - "embedding": _encode_floats_to_base64(embedding), + "embedding": _pooling_output_to_base64( + final_output.outputs.data, dimensions + ), "index": idx, } ) @@ -3182,3 +3176,33 @@ def _encode_floats_to_base64(floats: list[float]) -> str: """ packed = struct.pack(f"<{len(floats)}f", *floats) return base64.b64encode(packed).decode("ascii") + + +def _pooling_output_to_base64(data: Any, dimensions: int | None = None) -> str: + """Serialize a vLLM ``PoolingOutput.data`` tensor straight to a base64 + float32 string, skipping the intermediate Python ``list[float]`` and the + ``struct.pack("<{N}f", *floats)`` varargs expansion. + + ``torch -> numpy.tobytes -> base64`` keeps the heavy work in C; output bytes + are identical to the ``struct``-based path on little-endian hosts. + """ + if isinstance(data, torch.Tensor): + vec = data.detach().cpu().flatten().to(torch.float32) + if dimensions is not None: + if dimensions > vec.numel(): + raise ValueError( + f"dimensions={dimensions} exceeds model embedding " + f"dimension {vec.numel()}" + ) + vec = vec[:dimensions] + return base64.b64encode(vec.contiguous().numpy().tobytes()).decode("ascii") + # Fallback for non-tensor pooling outputs (rare): reuse the list path. + floats = _pooling_output_to_list(data) + if dimensions is not None: + if dimensions > len(floats): + raise ValueError( + f"dimensions={dimensions} exceeds model embedding " + f"dimension {len(floats)}" + ) + floats = floats[:dimensions] + return _encode_floats_to_base64(floats)