perf(embedding): same-node shared-memory fast path for the embedding response

[tzulingk]

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, the SHM transport provides no measurable benefit at a real 1024-dim embedding; the earlier wins 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

Status: parked — not recommended for text embeddings. Benchmarks (GB200, Qwen3-Embedding-0.6B, batch 15 → 1024) show this SHM path matches the much simpler numpy.tobytes serialization fix in #10221 — same latency — while adding a same-node constraint (frontend + worker must share /dev/shm), per-request shm_open/unlink syscalls, and a leak-on-error risk. There is also no way for it to compound with #10221: when SHM is on, the base64 path #10221 optimizes is never executed.

Recommendation: ship #10221 for text embeddings; keep this PR for the large-payload / multimodal path (e.g. frontend→worker image mm_kwargs), where the wire transfer itself dominates and shared memory pulls meaningfully ahead. Holding as draft until there's such a consumer.

Overview:

For small embedding models the GPU forward pass is sub-millisecond, so per-request framework overhead dominates end-to-end latency. A large slice of that overhead is serializing the embedding vectors on the worker → frontend hop: the Python worker turns the pooling tensor into a base64 string and ships it as JSON, and the Rust frontend parses + decodes it.

This PR adds an opt-in, same-node shared-memory fast path for the embedding response. Instead of serializing the vectors onto the wire, the worker writes them as raw f32 bytes into a POSIX shared-memory segment (/dev/shm) and sends only a tiny {name, count, dim} handle in the normal response; the frontend reads the bytes directly and unlinks the segment. This is the same multiprocessing.shared_memory handle-passing pattern as the multimodal mm_kwargs_transfer path, adapted to the worker→frontend direction with a Rust reader.

Enabled with DYN_EMBEDDING_SHM=1. Measured on GB200 (Qwen3-Embedding-0.6B, dim 3072): for small batches it matches the cheaper serialization fix; for larger response payloads (e.g. batch 64, ~786 KB) it begins to pull ahead, and the gap grows with payload size.

Architecture:

The frontend (Rust HTTP) and worker (Python/vLLM) stay as two separate processes — no merged CLI, no single-process mode. The only change is how the response payload crosses between them when they're co-located (sharing /dev/shm):

sequenceDiagram
    autonumber
    participant C as Client
    participant F as Frontend<br/>(Rust, HTTP)
    participant W as Worker<br/>(Python, vLLM)
    participant S as /dev/shm<br/>(POSIX shared memory)

    C->>F: POST /v1/embeddings
    F->>W: request (token ids) over request plane
    W->>W: engine.encode → pooling tensor [count, dim]

    alt DYN_EMBEDDING_SHM=1  (same-node fast path)
        W->>S: write [count, dim] f32 (numpy.tobytes)
        W-->>F: response { data: [], data_shm: {name, count, dim} }
        Note over W,F: only a tiny handle crosses the wire
        F->>S: std::fs::read("/dev/shm/<name>")
        F->>S: remove_file  (unlink — frontend owns cleanup)
        F->>F: build vectors in client's encoding (float / base64)
    else default  (base64 over the wire)
        W->>W: tensor → base64
        W-->>F: response { data: [ {embedding: "<base64>"}, ... ] }
        F->>F: base64 → float (if client asked for float)
    end

    F-->>C: embeddings JSON

Loading

Key properties:

• Opt-in & same-node only. When DYN_EMBEDDING_SHM is unset, the existing base64-over-the-wire path is used unchanged. SHM requires the frontend and worker to share /dev/shm (one pod / shared IPC); cross-node deployments keep the normal path.
• No protocol break. The handle rides on an internal data_shm field that is never serialized to the client; the public /v1/embeddings response is identical.
• No new deps. POSIX shm is tmpfs at /dev/shm, so the frontend just uses std::fs::read / remove_file — no libc/mmap.

Details:

Worker (Python — components/src/dynamo/vllm/handlers.py)

• _write_embeddings_to_shm() stacks the batch into one contiguous [count, dim] little-endian f32 buffer (torch.stack → numpy.tobytes), writes it to a fresh multiprocessing.shared_memory segment, and returns the {name, count, dim} handle.
• The embedding handler yields { "data": [], "data_shm": {...} } when DYN_EMBEDDING_SHM is set, else the unchanged base64 list.
• The segment is unregistered from resource_tracker so the worker doesn't race the frontend to unlink it — the frontend owns cleanup.

Frontend (Rust)

• lib/llm/src/protocols/openai/embeddings.rs: NvCreateEmbeddingResponse gains an internal data_shm: Option<EmbeddingShmHandle> (#[serde(skip_serializing_if)], so never sent to the client).
• .../embeddings/aggregator.rs: merge carries the handle through the stream fold.
• lib/llm/src/http/service/openai.rs: read_shm_embeddings() does std::fs::read("/dev/shm/<name>") → count × dim f32 → vectors in the client's requested encoding → remove_file (unlink). The embeddings handler invokes it before the existing base64-decode step.

Benchmark (GB200, Qwen3-Embedding-0.6B, dim 3072, ISL 80):

Per-request latency (ms, avg). σ = standard deviation of per-request latency across the run (run noise, not an error bar on the mean); % figures compare the avg. Small batches: rate-limited (no queueing), 200/100 reqs. Large batches: closed-loop --concurrency 1, 40/25 reqs.

batch	response (raw f32)	baseline (tolist+struct.pack)	numpy.tobytes (#10221)	SHM (this PR)	numpy.tobytes + SHM
15	~184 KB	124.25	86.94	85.49	85.31
64	~786 KB	469.54	306.57	290.69	294.72
128	~1.5 MB	934.16	584.31	573.98	566.64
256	~3 MB	1,860.85	1,182.62	1,128.59	1,126.00
512	~6 MB	3,684.53	2,364.84	2,301.73	2,200.67
1024	~12 MB	7,491.99	4,694.72	4,615.99	4,447.78

(avg ms.) Both serialization-killers beat the baseline by ~30–37%, growing with batch. SHM's edge over the much simpler numpy.tobytes (#10221) is only ~2–5% and within run-to-run noise at every size (the two ran on different nodes) — i.e. statistically indistinguishable for text embeddings. SHM's real advantage is reserved for payloads large enough that the wire transfer dominates (multimodal), which is why this PR is parked rather than recommended for the embedding path.

Where should the reviewer start?

• lib/llm/src/http/service/openai.rs — read_shm_embeddings + the data_shm branch in the embeddings handler (the read + unlink).
• lib/llm/src/protocols/openai/embeddings.rs and .../embeddings/aggregator.rs — the internal data_shm field and the fold carry-through.
• components/src/dynamo/vllm/handlers.py — _write_embeddings_to_shm and the gated yield.

Caveats / follow-ups:

• Leak-on-error: if the frontend dies after the worker writes but before it unlinks, the segment lingers in /dev/shm until the container restarts. A worker-side reaper / TTL is wanted before this is used outside benchmarking.
• Same-node constraint is enforced only by operator config today (the env var); auto-detecting co-location is a follow-up.
• For small text-embedding payloads, a simpler change — building the base64 directly from the tensor with numpy.tobytes — captures most of the same win without shared memory and works cross-node; this SHM path is aimed at larger payloads where the wire transfer itself dominates.

🤖 Generated with Claude Code

[tzulingk]

Corrected benchmark — the earlier embedding-serialization gains were a pooling bug, not a real win

While 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)

EmbeddingWorkerHandler built PoolingParams() with no task. On vLLM 0.21, encode() with task=None returns per-token output (the full n_tokens × hidden hidden-state matrix, un-normalized) instead of one pooled, L2-normalized vector. So /v1/embeddings returned n_tokens × 1024 floats — e.g. dim 2048 for "hi", 3072 for "hello world", ~81,920 for an 80-token prompt — un-normalized.

This is why the prior benchmarks looked the way they did: at ISL≈80, dynamo was serializing ~80× more floats than bare vllm serve (which pools to 1024). The "30–37% opt-2 win" and "SHM ≈ opt-2" results were entirely the cost of (de)serializing that oversized payload. Fix = PoolingParams(task="embed") → fixed 1024-dim, ‖v‖=1.0, matching bare vLLM and DIS-2154's "15 × 1024 floats".

Corrected same-node matrix

GB200 (single node), vLLM 0.21, Qwen/Qwen3-Embedding-0.6B, dim 1024, ISL 80, warmup 15. Small batches (15/64) rate-limited (rate 10); large batches (128–1024) closed-loop --concurrency 1. All configs run back-to-back on the same node with identical params.

p50 latency (ms):

batch	bare vLLM	dyn baseline	dyn opt-2	dyn SHM	dyn opt-2+SHM
15	28.3	26.9	32.0	27.6	25.7
64	46.3	57.4	65.0	54.6	50.2
128	81.5	87.6	93.4	78.1	75.8
256	144.0	161.4	154.0	151.2	142.4
512	261.6	313.9	293.8	285.0	269.4
1024	542.4	631.7	604.0	589.2	573.0

p90 latency (ms):

batch	bare vLLM	dyn baseline	dyn opt-2	dyn SHM	dyn opt-2+SHM
15	30.6	28.2	34.4	29.4	26.9
64	51.7	82.3	81.5	68.4	65.1
128	88.1	100.5	95.7	80.3	89.6
256	176.3	164.8	193.1	156.4	175.1
512	265.8	618.5	670.4	616.2	597.2
1024	674.2	977.1	1029.6	962.5	955.8

p99 latency (ms):

batch	bare vLLM	dyn baseline	dyn opt-2	dyn SHM	dyn opt-2+SHM
15	32.6	36.3	36.2	31.2	29.4
64	56.2	284.0	392.4	76.1	203.5
128	91.2	375.1	164.8	154.1	150.4
256	187.8	463.2	532.1	455.0	449.2
512	271.0	661.3	694.6	632.4	608.7
1024	737.0	1001.5	1056.7	958.3	959.0

Findings

• opt-2 ≈ baseline ≈ SHM ≈ opt-2+SHM at p50 — all four dynamo serialization paths are within run-to-run noise of each other. opt-2 is not faster than baseline (slightly slower at small batches); SHM / opt-2+SHM are marginally lower but within noise. At a real 1024-dim embedding, the tolist()+struct.pack cost is negligible, so optimizing it changes nothing.
• bare vllm serve has the tightest tail (p90/p99 ≈ p50). Every dynamo variant has a fat tail — the frontend↔worker IPC hop + Python GIL add occasional slow requests (see p99 at batch 64/512/1024). At p50 dynamo is within ~10–20% of vLLM; on the tail vLLM wins clearly.
• The dim/normalization bug, not serialization, was the whole story. DIS-2154's 1.95× (dynamo 123.7 ms / vLLM 63.4 ms) was on vLLM 0.19; the larger gaps in later runs were the per-token payload on vLLM 0.21, not a regression.

Recommendation

Ship the correctness fix #10248 (PoolingParams(task="embed")). The serialization optimizations — #10231 (numpy.tobytes) and #10220 (SHM) — do not justify on perf grounds for realistic text-embedding sizes and are recommended for closure.

[tzulingk]

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 (PoolingParams(task="embed")). Closing in favor of that.