fix(vllm): pool embedding worker output via PoolingParams(task="embed")

[tzulingk]

Overview:

The dynamo text-embedding worker returns the wrong output on vLLM 0.21: it emits the full per-token hidden-state matrix (n_tokens × hidden, un-normalized) instead of one pooled, L2-normalized embedding per input. This one-line fix sets the pooling task to "embed", restoring correct OpenAI /v1/embeddings output.

Details:

EmbeddingWorkerHandler.generate builds pooling_params = PoolingParams() with no task. On vLLM 0.21, engine.encode() with task=None resolves to per-token output, so the worker serializes the entire hidden-state sequence. The /v1/embeddings response dimension then scales with input length instead of being the model's fixed embedding dimension.

• Symptom: for Qwen/Qwen3-Embedding-0.6B (hidden_size=1024), input=["hi"] returns dim 2048, "hello world" → 3072, an 80-token prompt → ~81,920 — i.e. n_tokens × 1024 — and the vectors are un-normalized.
• Expected: one 1024-dim, L2-normalized vector per input, matching bare vllm serve and the model's sentence-transformers config (pooling_mode_lasttoken, 2_Normalize).
• Cause: vLLM's default pooling-task behavior changed between 0.19 (pooled by default) and 0.21 (per-token by default); the handler relied on the old default.
• Fix: PoolingParams(task="embed") selects the pooled + normalized embedding, matching vLLM's own embedding server.
• Follow-up (this PR): also forward the OpenAI dimensions parameter into PoolingParams(task="embed", dimensions=...) rather than truncating the pooled vector in the handler. vLLM's pooler then performs the Matryoshka reduction correctly (truncate then L2-renormalize) and validates model support, matching bare vllm serve; the previous post-hoc slice did not re-normalize (returned ‖v‖ < 1). Models whose HF config doesn't declare Matryoshka (e.g. Qwen3-Embedding) are launched with --hf-overrides '{"is_matryoshka": true}' in agg_embed.sh.

Verification (GB200, Qwen/Qwen3-Embedding-0.6B):

input	~tokens	dim before	dim after
"hi"	2	2048	1024
"hello world"	3	3072	1024
40 words	47	48128	1024
100 words	102	104448	1024

After the fix the response is a fixed 1024-dim vector for any input length, with ‖v‖ = 1.0.

Where should the reviewer start?

• components/src/dynamo/vllm/handlers.py — EmbeddingWorkerHandler.generate, the PoolingParams(task="embed") line.

Related:

• Surfaced while benchmarking embedding serialization (DIS-2177 / DIS-2178). The corrected, same-node benchmark there shows the serialization optimizations in perf(embedding): build embedding base64 directly from the tensor (numpy.tobytes) #10231 (numpy.tobytes) and perf(embedding): same-node shared-memory fast path for the embedding response #10220 (SHM) provide no measurable benefit once embeddings are the correct size — the apparent 30–37% gains were artifacts of this oversized per-token payload.

vLLM version (updated):

main now pins vLLM 0.23.0 (#10723) and this branch is rebased onto it. The pooling/MRL internals the fix relies on — PoolingParams(task="embed", dimensions=...), the is_matryoshka gate in PoolingParams.verify(), and the pooler's truncate-then-normalize order in EmbeddingPoolerHead — are unchanged between 0.22.0 and 0.23.0, and the fix is functionally verified on 0.23.0 in CI (the embedding_agg serve test, including the dimensions:128 case). The 0.21 / 0.22 tables further down are retained as historical; the 0.23.0 latency refresh is in the section immediately below.

vLLM 0.23.0 latency (refreshed):

Rebuilt a dynamo image on the 0.23.0 base (arm64) and re-ran the sweep on a single GB200. Same node/GPU, back-to-back, identical engine args (--dtype bfloat16 --max-model-len 3096 --runner pooling --pooler-config '{"pooling_type":"MEAN","use_activation":false}' --no-enable-prefix-caching) and identical aiperf params (Qwen/Qwen3-Embedding-0.6B, dim 1024, ISL 80, --request-count 200, warmup 15). Small batches (15/64) rate-limited (--request-rate 10); large (128–1024) closed-loop --concurrency 1. dyn = dynamo embedding worker (pooling-fixed, task="embed"), baseline serialization.

p50 latency (ms):

batch	vLLM 0.23	dyn 0.23
15	42.3	35.0
64	83.3	72.9
128	152.5	111.7
256	286.3	179.4
512	541.9	420.9
1024	1075.0	780.8

p90 latency (ms):

batch	vLLM 0.23	dyn 0.23
15	44.9	38.0
64	87.4	76.8
128	170.1	115.8
256	298.8	227.6
512	567.1	548.4
1024	1098.5	1723.0

p99 latency (ms):

batch	vLLM 0.23	dyn 0.23
15	51.2	40.1
64	92.7	279.6
128	180.0	409.9
256	309.7	1097.7
512	579.8	1459.4
1024	1128.6	2092.6

Takeaways (0.23.0): the pooling-fixed dynamo worker is at/below bare vllm serve at p50 and p90 for every batch (e.g. batch-1024 p50 781 vs 1075 ms) — no median-latency regression on the shipped version, output dim = 1024. bare vLLM keeps a much tighter tail: dynamo's p99 grows at batch ≥64 (the frontend↔worker IPC hop + queueing), consistent with the 0.21/0.22 finding. The dyn↔vLLM comparison is apples-to-apples within this same-node 0.23.0 run; absolute values are not directly comparable to the historical 0.21/0.22 tables below (different harness/node).

Verified on vLLM 0.22.0 + 0.21 vs 0.22 sweep (historical):

Built a dynamo image on the 0.22.0 base and re-ran the verification + the full uniform-count sweep:

• Bug reproduces on 0.22.0 (the version pyproject.toml pins): unpatched worker returns per-token output — "hi" → dim 2048, "hello world" → 3072, 100-word → ~104,448 (= n_tokens × 1024), un-pooled. With task="embed": a fixed 1024-dim for any input length, ‖v‖=1.0. So this is a live correctness fix on the shipped version, not just hardening.
• 0.21 vs 0.22 latency is within noise (0.22 marginally tighter at large batch), so the change carries no perf regression.

Same node (single GB200), Qwen/Qwen3-Embedding-0.6B, dim 1024, ISL 80, --request-count 200 (uniform). Small batches (15/64) rate-limited; large (128–1024) closed-loop --concurrency 1. dyn = dynamo embedding worker, pooling-fixed (task="embed"), baseline serialization.

p50 latency (ms):

batch	vLLM 0.21	vLLM 0.22	dyn 0.21	dyn 0.22
15	28.7	28.6	27.9	26.5
64	47.0	45.3	69.9	54.7
128	79.1	79.7	99.6	89.3
256	139.1	140.5	177.5	160.3
512	271.9	266.6	335.3	324.4
1024	522.9	502.8	681.9	689.8

p90 latency (ms):

batch	vLLM 0.21	vLLM 0.22	dyn 0.21	dyn 0.22
15	31.7	30.5	32.9	28.3
64	55.9	49.1	73.1	57.4
128	81.4	80.8	103.4	101.4
256	142.7	157.5	210.0	168.0
512	280.8	272.7	421.5	505.8
1024	573.8	512.3	1288.1	1314.5

p99 latency (ms):

batch	vLLM 0.21	vLLM 0.22	dyn 0.21	dyn 0.22
15	46.2	32.5	37.4	31.9
64	71.7	54.2	79.5	63.5
128	84.1	84.6	474.1	410.3
256	146.9	161.7	618.7	531.5
512	310.7	275.7	884.7	827.5
1024	639.6	518.9	1458.0	1417.5

Takeaways: 0.21 ≈ 0.22 for both bare vLLM and the (pooling-fixed) dynamo worker. At p50 dynamo is within ~10–25% of bare vllm serve; bare vLLM keeps a much tighter tail (p99@1024 ≈ 0.52 s vs dynamo ≈ 1.4 s — the frontend↔worker IPC hop). The pooling fix is the only change that matters here; serialization optimizations (the now-closed #10231 / #10220) stay within noise at correct embedding size.

🤖 Generated with Claude Code

[coderabbitai]

Walkthrough

The embeddings request handler in the vLLM integration now explicitly configures vLLM's pooling behavior by passing PoolingParams(task="embed") instead of default pooling parameters. This change ensures the /v1/embeddings endpoint returns one pooled, normalized embedding vector per input rather than per-token embeddings, aligning with OpenAI API response expectations.

Changes

Embeddings Pooling Configuration

Layer / File(s)	Summary
vLLM pooling parameter configuration components/src/dynamo/vllm/handlers.py	The embeddings handler sets PoolingParams(task="embed") to select vLLM's sentence-level pooling behavior, ensuring the response contains one pooled, normalized embedding vector per input rather than per-token embeddings.

Estimated code review effort

🎯 1 (Trivial) | ⏱️ ~3 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The PR description is detailed, but it lacks the required Related Issues section with either a closing issue link or a no-issue confirmation.	Add the required Related Issues section, then keep only one path: either 'Closes #XXXX' or the checked 'Confirmed — no related issue' item.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: adding PoolingParams(task='embed') to the vLLM embedding worker output, which directly corresponds to the core fix described in the changeset.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[tzulingk]

Confirmed the bug is live on both vLLM 0.21.0 and 0.22.0 (main pins 0.22.0). Built a dynamo image on the 0.22.0 base and ran the unpatched worker — Qwen/Qwen3-Embedding-0.6B:

input	~tokens	dim (unpatched, 0.22)	dim (this fix)
"hi"	2	2048	1024
"hello world"	3	3072	1024
100 words	102	104448	1024

So it returns n_tokens × 1024 (per-token, un-pooled) on the shipped version too; with task="embed" it's a fixed 1024-dim, L2-normalized vector for any input length. This is a live correctness fix, not just hardening. Latency on 0.21 vs 0.22 is within noise, so no perf regression from the change.

[treydock]

Can confirm this resolves embedding model issues on 1.3.0.dev.1 images. Using this patch on 1.3.0 dev.1 image:

"""Applies fix from https://github.com/ai-dynamo/dynamo/pull/10248.
"""
import os
import sys

if len(sys.argv) > 1:
    if '1.2.0' in sys.argv[1]:
        sys.exit(0)

SITE = "/usr/local/lib/python3.12/dist-packages/dynamo/vllm/"


def patch(path, replacements):
    with open(path) as f:
        src = f.read()
    for old, new in replacements:
        assert old in src, f"Pattern not found in {path}:\n{old!r}"
        src = src.replace(old, new, 1)
    with open(path, "w") as f:
        f.write(src)
    print(f"patched {path}")


patch(
    f"{SITE}/handlers.py",
    [
        (
            'from vllm import PoolingParams',
            '',
        ),
        (
            'import torch',
            'import torch\n'
            'from vllm import PoolingParams',
        ),
        (
            'pooling_params = PoolingParams()',
            'pooling_params = PoolingParams(task="embed")',
        ),
    ],
)

Usage in Dockerfile:

COPY patch-10248.py .
ARG VERSION=1.3.0-dev.1-cuda13
RUN python3 patch-10248.py $VERSION