From 5af3e6fb66eb14073acb8cb6b95353cc5214f6a9 Mon Sep 17 00:00:00 2001 From: Tzu-Ling Date: Mon, 1 Jun 2026 22:14:33 -0500 Subject: [PATCH 1/2] =?UTF-8?q?perf(embedding):=20opt-2=20on=20SHM=20?= =?UTF-8?q?=E2=80=94=20direct=20tensor->base64=20for=20the=20non-SHM=20pat?= =?UTF-8?q?h?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Stacks on the SHM branch. The SHM fast path already passes raw f32 bytes (no base64), so opt-2's remaining target is the base64 fallback used when DYN_EMBEDDING_SHM is off: build the base64 straight from the pooling tensor via torch -> numpy.tobytes, dropping the per-embedding Python float-list materialization and the struct.pack("<{N}f", *floats) varargs (DIS-2154 #3). Output bytes are unchanged on little-endian hosts. Net effect: - DYN_EMBEDDING_SHM=1: identical to the SHM branch (raw f32 over /dev/shm). - DYN_EMBEDDING_SHM=0: faster base64 serialization than the SHM branch. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Tzu-Ling --- components/src/dynamo/vllm/handlers.py | 63 +++++++++++++++++++------- 1 file changed, 46 insertions(+), 17 deletions(-) diff --git a/components/src/dynamo/vllm/handlers.py b/components/src/dynamo/vllm/handlers.py index e548d9cadf34..31c60dd8e026 100644 --- a/components/src/dynamo/vllm/handlers.py +++ b/components/src/dynamo/vllm/handlers.py @@ -3015,23 +3015,20 @@ async def _encode_one(idx: int, prompt: Any): # 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: list[Dict[str, Any]] = [] - 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] - embedding_objects.append( - { - "object": "embedding", - "embedding": _encode_floats_to_base64(embedding), - "index": idx, - } - ) + # opt-2: build the base64 straight from the pooling tensor + # (torch -> numpy.tobytes), skipping the per-embedding Python float list + # + struct.pack("<{N}f", *floats) varargs (DIS-2154 #3). Bytes are + # identical to the struct path on little-endian hosts. + embedding_objects: list[Dict[str, Any]] = [ + { + "object": "embedding", + "embedding": _pooling_output_to_base64( + final_output.outputs.data, dimensions + ), + "index": idx, + } + for idx, final_output in enumerate(outputs) + ] yield { "object": "list", @@ -3150,6 +3147,38 @@ def _encode_floats_to_base64(floats: list[float]) -> str: 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 (DIS-2154 #3). + + ``torch -> numpy.tobytes -> base64`` keeps the heavy work in C; output bytes + are identical to the ``struct``-based path on little-endian hosts. + ``.to(torch.float32)`` makes bf16/fp16 pooling outputs match the + ``struct.pack(" 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) + + def _write_embeddings_to_shm(outputs: list[Any], dimensions: int | None) -> dict[str, Any]: """Stage all embedding vectors as one contiguous ``[count, dim]`` little-endian f32 buffer in a POSIX shared-memory segment; return the From 14a42718c89e8d0c85b825afb272afb4ba962871 Mon Sep 17 00:00:00 2001 From: Tzu-Ling Date: Tue, 2 Jun 2026 14:07:12 -0500 Subject: [PATCH 2/2] refactor(embedding): share pooling-tensor flatten between list + base64 helpers Factor the detach/cpu/flatten/float32 step into `_flatten_pooling_tensor`, reused by both `_pooling_output_to_list` and `_pooling_output_to_base64` so the tensor-prep isn't duplicated. The fast base64 path still avoids `.tolist()` (it goes tensor -> numpy.tobytes), so no per-element Python work is reintroduced. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Tzu-Ling --- components/src/dynamo/vllm/handlers.py | 40 +++++++++++++++----------- 1 file changed, 23 insertions(+), 17 deletions(-) diff --git a/components/src/dynamo/vllm/handlers.py b/components/src/dynamo/vllm/handlers.py index 31c60dd8e026..6ef293976881 100644 --- a/components/src/dynamo/vllm/handlers.py +++ b/components/src/dynamo/vllm/handlers.py @@ -3010,15 +3010,12 @@ async def _encode_one(idx: int, prompt: Any): return # Default wire format: always base64 (the Rust frontend decodes back to - # float when the client's ``encoding_format`` is float or unset). - # 15x1024-float 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. - # opt-2: build the base64 straight from the pooling tensor - # (torch -> numpy.tobytes), skipping the per-embedding Python float list - # + struct.pack("<{N}f", *floats) varargs (DIS-2154 #3). Bytes are - # identical to the struct path on little-endian hosts. + # 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; 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]] = [ { "object": "embedding", @@ -3113,16 +3110,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)): @@ -3150,15 +3158,13 @@ def _encode_floats_to_base64(floats: list[float]) -> str: 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 (DIS-2154 #3). + ``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. - ``.to(torch.float32)`` makes bf16/fp16 pooling outputs match the - ``struct.pack(" vec.numel(): raise ValueError(