From e7f365655598c9a6b2383d72b9483ebc53645c64 Mon Sep 17 00:00:00 2001 From: Tzu-Ling Date: Tue, 2 Jun 2026 14:55:58 -0500 Subject: [PATCH 1/2] perf(embedding): build embedding base64 directly from the tensor (numpy.tobytes) Build the worker's base64 embedding payload straight from the pooling tensor via numpy.tobytes, instead of tensor.tolist() + struct.pack("<{N}f", *floats). For a batch-15 x 3072-dim response that removes materializing 46,080 Python float objects and unpacking them as 46,080 positional args into struct.pack -- an O(N) interpreter-level pass replaced by a single C-level copy. Emitted base64 bytes are identical on little-endian hosts. Shared detach/cpu/flatten/float32 tensor-prep is factored into _flatten_pooling_tensor, reused by _pooling_output_to_list so the tensor handling isn't duplicated; the fast path still avoids .tolist(). Measured on GB200 (Qwen3-Embedding-0.6B, dim 3072): -30% at batch 15 (124 -> 87 ms) and -35% at batch 64 (470 -> 307 ms) per-request latency. Works cross-node (no shared memory). Complements #10139 (base64 wire format): that made shipping the payload cheap, this makes producing it cheap. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Tzu-Ling --- components/src/dynamo/vllm/handlers.py | 77 +++++++++++++++++++------- 1 file changed, 56 insertions(+), 21 deletions(-) diff --git a/components/src/dynamo/vllm/handlers.py b/components/src/dynamo/vllm/handlers.py index a7b534cf4f9a..74fa97319983 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, } ) @@ -3150,16 +3144,27 @@ def _classify_embedding_input(input_field: Any) -> list[Any]: ) +def _flatten_pooling_tensor(data: "torch.Tensor") -> "torch.Tensor": + """Flatten a vLLM ``PoolingOutput.data`` tensor to a 1-D float32 CPU tensor. + + Shared by :func:`_pooling_output_to_list` and + :func:`_pooling_output_to_base64` so the detach/cpu/flatten/cast step isn't + duplicated. ``float32`` matches the OpenAI base64 f32 wire format. + + vLLM's pooling pipeline can return a tensor with a singleton batch dim + (shape ``(1, hidden_dim)``) instead of a 1D vector; we flatten unconditionally. + """ + return data.detach().cpu().flatten().to(torch.float32) + + def _pooling_output_to_list(data: Any) -> list[float]: """Convert a vLLM PoolingOutput.data tensor (or list) to a flat list[float]. - vLLM's pooling pipeline can return a tensor with a singleton batch dim - (shape ``(1, hidden_dim)``) instead of a 1D vector (shape ``(hidden_dim,)``). The OpenAI ``/v1/embeddings`` response expects ``data[].embedding`` to be a - flat array of floats, so we flatten unconditionally. + flat array of floats. """ if isinstance(data, torch.Tensor): - return data.detach().cpu().flatten().tolist() + return _flatten_pooling_tensor(data).tolist() if isinstance(data, (list, tuple)): # Already a list — flatten one level if it's a list-of-lists. if data and isinstance(data[0], (list, tuple)): @@ -3182,3 +3187,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 = _flatten_pooling_tensor(data) + 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) From d73aa0f16b7c759dc1ea2f5849b8f8a89cf941f6 Mon Sep 17 00:00:00 2001 From: Tzu-Ling Date: Tue, 2 Jun 2026 15:00:51 -0500 Subject: [PATCH 2/2] refactor(embedding): inline the pooling-tensor flatten Drop the _flatten_pooling_tensor helper; inline data.detach().cpu().flatten().to(torch.float32) at its single use in _pooling_output_to_base64. _pooling_output_to_list reverts to its original form, so the net diff vs main is just the response-loop switch + the new _pooling_output_to_base64 helper. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Tzu-Ling --- components/src/dynamo/vllm/handlers.py | 21 +++++---------------- 1 file changed, 5 insertions(+), 16 deletions(-) diff --git a/components/src/dynamo/vllm/handlers.py b/components/src/dynamo/vllm/handlers.py index 74fa97319983..d1d32c6240a9 100644 --- a/components/src/dynamo/vllm/handlers.py +++ b/components/src/dynamo/vllm/handlers.py @@ -3144,27 +3144,16 @@ def _classify_embedding_input(input_field: Any) -> list[Any]: ) -def _flatten_pooling_tensor(data: "torch.Tensor") -> "torch.Tensor": - """Flatten a vLLM ``PoolingOutput.data`` tensor to a 1-D float32 CPU tensor. - - Shared by :func:`_pooling_output_to_list` and - :func:`_pooling_output_to_base64` so the detach/cpu/flatten/cast step isn't - duplicated. ``float32`` matches the OpenAI base64 f32 wire format. - - vLLM's pooling pipeline can return a tensor with a singleton batch dim - (shape ``(1, hidden_dim)``) instead of a 1D vector; we flatten unconditionally. - """ - return data.detach().cpu().flatten().to(torch.float32) - - def _pooling_output_to_list(data: Any) -> list[float]: """Convert a vLLM PoolingOutput.data tensor (or list) to a flat list[float]. + vLLM's pooling pipeline can return a tensor with a singleton batch dim + (shape ``(1, hidden_dim)``) instead of a 1D vector (shape ``(hidden_dim,)``). The OpenAI ``/v1/embeddings`` response expects ``data[].embedding`` to be a - flat array of floats. + flat array of floats, so we flatten unconditionally. """ if isinstance(data, torch.Tensor): - return _flatten_pooling_tensor(data).tolist() + return data.detach().cpu().flatten().tolist() if isinstance(data, (list, tuple)): # Already a list — flatten one level if it's a list-of-lists. if data and isinstance(data[0], (list, tuple)): @@ -3198,7 +3187,7 @@ def _pooling_output_to_base64(data: Any, dimensions: int | None = None) -> str: are identical to the ``struct``-based path on little-endian hosts. """ if isinstance(data, torch.Tensor): - vec = _flatten_pooling_tensor(data) + vec = data.detach().cpu().flatten().to(torch.float32) if dimensions is not None: if dimensions > vec.numel(): raise ValueError(