perf(embedding): build embedding base64 directly from the tensor (numpy.tobytes) - #10231
perf(embedding): build embedding base64 directly from the tensor (numpy.tobytes)#10231tzulingk wants to merge 2 commits into
Conversation
…py.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) <noreply@anthropic.com>
Signed-off-by: Tzu-Ling <tzulingk@nvidia.com>
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) <noreply@anthropic.com> Signed-off-by: Tzu-Ling <tzulingk@nvidia.com>
WalkthroughThe PR refactors embedding base64 serialization in ChangesEmbedding Base64 Serialization Refactor
🎯 2 (Simple) | ⏱️ ~10 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Corrected benchmark — the earlier embedding-serialization gains were a pooling bug, not a real winWhile re-benchmarking on GB200 I found the dynamo embedding worker was returning the wrong output on vLLM 0.21, which invalidates the prior numbers on this issue. After fixing it and re-measuring same-node, the serialization optimizations (numpy.tobytes / SHM) provide no measurable benefit at a correct-size embedding. The bug (fixed in #10248)
This is why the prior benchmarks looked the way they did: at ISL≈80, dynamo was serializing ~80× more floats than bare Corrected same-node matrixGB200 (single node), vLLM 0.21, p50 latency (ms):
p90 latency (ms):
p99 latency (ms):
Findings
RecommendationShip the correctness fix #10248 ( |
|
Closing: the corrected same-node benchmark (see the matrix comment above) shows this optimization provides no measurable benefit at a real 1024-dim embedding — the earlier gains were artifacts of a pooling bug that made the worker emit per-token output (~80x oversized payload) on vLLM 0.21. The actual fix is #10248 ( |
Pull request was closed
1 participant
Warning
Superseded — recommend closing. Re-benchmarking on GB200 found the dynamo embedding worker returns per-token output (wrong shape, dim scales with input length) on vLLM 0.21 — fixed in #10248. With that corrected, this PR's
numpy.tobytesoptimization provides no measurable benefit at a real 1024-dim embedding; the 30–37% gains reported below were artifacts of the oversized per-token payload. See the corrected same-node p50/p90/p99 matrix in the comment below. Recommend closing in favor of #10248.Important
Recommended embedding fix. Captures the full −30–35% latency win cross-node with a one-line, behavior-preserving change. Benchmarks (below) show it matches the shared-memory path (#10220) without that PR's same-node constraint or extra failure modes, so it — not SHM — is the one to ship for text embeddings. (Complements the already-merged base64 wire-format change in #10139: that made shipping the payload cheap; this makes producing it cheap.)
Overview:
Speeds up embedding-response serialization on the worker by building the base64 payload directly from the pooling tensor via
numpy.tobytes, instead of going through a Python float list +struct.packvarargs expansion.The current path is
tensor.detach().cpu().flatten().tolist()→struct.pack("<{N}f", *floats)→ base64. For a batch-15 × 3072-dim response that's 46,080 Python float objects materialized and then unpacked as 46,080 positional args intostruct.pack— an O(N) interpreter-level pass over every element. This PR replaces it withtensor → numpy.tobytes() → base64, a single C-level copy. The emitted base64 bytes are identical on little-endian hosts.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. The win grows with batch size, since it scales with the number of floats serialized. It also works cross-node (no shared memory required) and complements the base64-wire-format change in #10139 (this PR makes producing that base64 cheaper; #10139 made shipping it cheaper).
Details:
components/src/dynamo/vllm/handlers.py: new_pooling_output_to_base64()builds base64 straight from the tensor. Shared tensor-prep (detach().cpu().flatten().to(float32)) is factored into a small helper reused by_pooling_output_to_list, so there's no duplicated tensor handling. A list/tuple fallback preserves behavior for non-tensor pooling outputs..to(torch.float32)makes bf16/fp16 pooling outputs match thestruct.pack("<f")width, keeping the emitted bytes identical.Benchmark (GB200, Qwen3-Embedding-0.6B, dim 3072, ISL 80):
Per-request latency (ms, avg). Small batches: rate-limited (no queueing), 200/100 reqs. Large batches: closed-loop
--concurrency 1, 40/25 reqs.tolist+struct.pack)The fix holds a consistent −30 to −37% vs baseline, and the absolute saving grows with batch (37 ms at batch 15 → ~2.8 s at batch 1024) — the baseline's
tolist()+struct.pack("<{N}f", *floats)is an O(N) interpreter-level pass that scales badly. SHM (#10220) tracks this PR within ~2–5% across all sizes (within run-to-run noise; the two configs ran on different nodes), so the shared-memory transport buys little over this serialization fix for text embeddings. Baseline large-batch runs used fewer samples (n=20 for 128/256, n=10 for 512/1024) since each request is so slow.Gap to bare
vllm serve(standalone vLLM v0.21.0 — the same vLLM Dynamo is built on, so this is the floor that isolates Dynamo's overhead):This PR closes ~42% of the Dynamo-over-vLLM overhead at every batch (gap closed =
1 − (thisPR − vLLM)/(baseline − vLLM)). A residual gap remains (Dynamo+this PR is 2.4–5.4× bare vLLM) — the two-process architecture's inherent cost: the extra worker→frontend serialization hop (base64 encode + transport + Rust decode) plus the Rust HTTP frontend, tokenizer, and request-plane, none of which single-processvllm servepays. That residual is what SHM/single-process target separately.Where should the reviewer start?
components/src/dynamo/vllm/handlers.py—_pooling_output_to_base64and its use in the embedding response loop.Related:
🤖 Generated with Claude Code