perf(embedding): build embedding base64 directly from the tensor (numpy.tobytes)

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)