fix(vllm): pool embedding worker output via PoolingParams(task="embed") (#10248) [cherry-pick → release/1.3.0]

components/src/dynamo/vllm/handlers.py

@@ -31,6 +31,7 @@
 )
 
 import torch
+from vllm import PoolingParams
 from vllm.config import ModelConfig, VllmConfig
 from vllm.inputs import EmbedsPrompt, TextPrompt, TokensPrompt
 from vllm.lora.request import LoRARequest
@@ -3878,18 +3879,15 @@ async def generate(
 
         The Rust frontend forwards the request dict directly. Expected keys:
         ``model: str``, ``input: str | list[str] | list[int] | list[list[int]]``.
-        Optional ``dimensions`` (Matryoshka truncation; first N floats of each
-        embedding). Optional ``encoding_format`` (``"float"`` -- default --
-        or ``"base64"``); when ``"base64"`` is requested, each per-input
-        vector is serialized as a base64-encoded string of little-endian
-        ``f32`` bytes per the OpenAI spec, applied after any
-        ``dimensions`` truncation so the byte count matches the requested
+        Optional ``dimensions`` (Matryoshka dimensionality reduction):
+        forwarded to vLLM's pooler, which truncates to N dims and
+        re-normalizes; vLLM requires the model to declare Matryoshka support.
+        Optional ``encoding_format`` (``"float"`` -- default -- or
+        ``"base64"``); when ``"base64"`` is requested, each per-input vector is
+        serialized as a base64-encoded string of little-endian ``f32`` bytes
+        per the OpenAI spec, so the byte count matches the (possibly reduced)
         dimensionality.
         """
-        # Lazy import to avoid pulling PoolingParams into handlers.py at module
-        # load time for non-embedding workers.
-        from vllm import PoolingParams
-
         model_name = request.get("model") or self.config.served_model_name or ""
         input_field = request.get("input")
         if input_field is None:
@@ -3920,7 +3918,29 @@ async def generate(
                 "expected 'float' or 'base64'"
             )
 
-        pooling_params = PoolingParams()
+        # Request the pooled sentence embedding. With no task, vLLM's
+        # encode() resolves to per-token output (the full ``n_tokens x
+        # hidden`` hidden-state matrix), so the OpenAI ``/v1/embeddings``
+        # response ends up with the wrong shape (dim scales with input
+        # length) instead of one vector per input. ``task="embed"`` selects
+        # the pooled embedding and runs the model's configured pooler
+        # (normalization included for models like Qwen3-Embedding), matching
+        # vLLM's own embedding server. ``use_activation`` is intentionally
+        # left at the pooler default so per-model behaviour isn't overridden.
+        #
+        # ``dimensions`` (OpenAI Matryoshka truncation) is forwarded to vLLM
+        # rather than applied here: vLLM's pooler truncates to ``dimensions``
+        # and then re-normalizes (the correct MRL behaviour) and validates
+        # that the model actually supports Matryoshka -- raising rather than
+        # silently returning a degraded, un-normalized vector for models that
+        # don't. This matches bare ``vllm serve``. Models whose HF config
+        # doesn't declare Matryoshka support (e.g. Qwen3-Embedding) must be
+        # launched with ``--hf-overrides '{"is_matryoshka": true}'`` for
+        # ``dimensions`` requests to be accepted.
+        pooling_kwargs: dict[str, Any] = {"task": "embed"}
+        if dimensions is not None:
+            pooling_kwargs["dimensions"] = dimensions
+        pooling_params = PoolingParams(**pooling_kwargs)
         # Use the per-request context id (same as the chat/completion paths
         # in this file) so concurrent embeddings never collide inside
         # ``AsyncLLM``. ``context.trace_id`` is a distributed-trace id
@@ -3973,14 +3993,24 @@ async def _encode_one(idx: int, prompt: Any):
         embedding_objects: list[Dict[str, Any]] = []
         prompt_tokens = 0
         for idx, final_output in enumerate(outputs):
+            # vLLM has already applied any ``dimensions`` Matryoshka reduction
+            # (truncate + re-normalize) inside the pooler, so this is the
+            # final per-input vector -- no post-hoc truncation here.
             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]
+
+            # vLLM rejects an unsupported ``dimensions`` for models that
+            # declare a ``matryoshka_dimensions`` list, but a model enabled
+            # via ``--hf-overrides '{"is_matryoshka": true}'`` (no explicit
+            # list) is only validated for ``dimensions >= 1`` -- the pooler
+            # then silently clamps an oversized request to the model's native
+            # size (``embeddings[..., :dimensions]``). Surface the same clear
+            # error the old post-hoc path raised instead of returning a
+            # shorter-than-requested vector.
+            if dimensions is not None and len(embedding) < dimensions:
+                raise ValueError(
+                    f"dimensions={dimensions} exceeds model embedding "
+                    f"dimension {len(embedding)}"
+                )
 
             # Always emit base64 over the worker->frontend wire format. The
             # Rust frontend decodes back to float when the client's

components/src/dynamo/vllm/tests/test_vllm_worker_handler.py

@@ -1419,6 +1419,99 @@ async def fake_encode(prompt, pooling_params, request_id):
         # cancel-and-await pass must not have touched the engine.
         assert aborted == []
 
+    @pytest.mark.asyncio
+    @pytest.mark.timeout(5)
+    async def test_dimensions_forwarded_to_pooling_params(self):
+        """``dimensions`` is forwarded to vLLM via ``PoolingParams`` rather
+        than applied as post-hoc truncation in the handler.
+
+        vLLM's pooler then does the Matryoshka reduction (truncate +
+        re-normalize) and validates that the model supports it. The handler
+        must NOT slice the returned vector itself, so it emits exactly what
+        the engine produced.
+        """
+        handler = self._make_embedding_handler()
+        context = self._make_context()
+        captured: dict = {}
+        # vLLM's pooler has already reduced to the requested ``dimensions``, so
+        # the stub returns a 128-dim vector (not 3) -- otherwise the handler's
+        # oversized-dimensions guard would (correctly) reject it.
+        vec = [i * 0.01 for i in range(128)]
+
+        async def fake_encode(prompt, pooling_params, request_id):
+            captured["pooling_params"] = pooling_params
+            output = MagicMock()
+            output.outputs.data = torch.tensor(vec)
+            output.prompt_token_ids = [1, 2, 3]
+            yield output
+
+        handler.engine_client.encode = fake_encode
+
+        request = {"input": ["hello"], "model": "test-model", "dimensions": 128}
+        responses = [r async for r in handler.generate(request, context)]
+
+        pp = captured["pooling_params"]
+        assert pp.task == "embed"
+        assert pp.dimensions == 128
+        # No post-hoc truncation: the handler returns exactly the vector vLLM
+        # produced (the 128-float stub here), trusting the pooler to have
+        # already applied the dimensionality reduction.
+        expected_b64 = mod._encode_floats_to_base64(vec)
+        assert responses[0]["data"][0]["embedding"] == expected_b64
+
+    @pytest.mark.asyncio
+    @pytest.mark.timeout(5)
+    async def test_no_dimensions_omits_pooling_dimensions(self):
+        """Without ``dimensions`` the handler requests ``task="embed"`` only,
+        leaving ``PoolingParams.dimensions`` unset so vLLM returns the model's
+        native embedding size.
+        """
+        handler = self._make_embedding_handler()
+        context = self._make_context()
+        captured: dict = {}
+
+        async def fake_encode(prompt, pooling_params, request_id):
+            captured["pooling_params"] = pooling_params
+            output = MagicMock()
+            output.outputs.data = torch.tensor([0.1, 0.2, 0.3])
+            output.prompt_token_ids = [1, 2, 3]
+            yield output
+
+        handler.engine_client.encode = fake_encode
+
+        request = {"input": ["hello"], "model": "test-model"}
+        _ = [r async for r in handler.generate(request, context)]
+
+        pp = captured["pooling_params"]
+        assert pp.task == "embed"
+        assert pp.dimensions is None
+
+    @pytest.mark.asyncio
+    @pytest.mark.timeout(5)
+    async def test_oversized_dimensions_raises(self):
+        """When vLLM silently clamps an oversized ``dimensions`` request (a
+        model enabled via ``--hf-overrides '{"is_matryoshka": true}'`` with no
+        ``matryoshka_dimensions`` list), the handler raises a clear error
+        instead of returning a shorter-than-requested vector.
+        """
+        handler = self._make_embedding_handler()
+        context = self._make_context()
+
+        async def fake_encode(prompt, pooling_params, request_id):
+            output = MagicMock()
+            # vLLM clamped to the model's native size (3 dims here) even though
+            # 2048 was requested.
+            output.outputs.data = torch.tensor([0.1, 0.2, 0.3])
+            output.prompt_token_ids = [1, 2, 3]
+            yield output
+
+        handler.engine_client.encode = fake_encode
+
+        request = {"input": ["hello"], "model": "test-model", "dimensions": 2048}
+        with pytest.raises(ValueError, match="exceeds model embedding dimension"):
+            async for _ in handler.generate(request, context):
+                pass
+
 
 class TestPadMmHashesTo64:
     """The frontend forwards canonical 16-char hex mm_hashes; vLLM must pad

examples/backends/vllm/launch/agg_embed.sh

@@ -69,6 +69,15 @@ python3 -m dynamo.frontend &
 # unusually long embedding inputs.
 MAX_MODEL_LEN="${MAX_MODEL_LEN:-2048}"
 
+# Qwen3-Embedding supports Matryoshka (flexible output dims 32-1024) but its
+# HF config does not declare it, so vLLM rejects OpenAI `dimensions` requests
+# unless told the model is Matryoshka. Inject the flag only for the default
+# model; other models must declare their own support (or omit `dimensions`).
+HF_OVERRIDES_ARGS=()
+if [[ "$MODEL" == "Qwen/Qwen3-Embedding-0.6B" ]]; then
+    HF_OVERRIDES_ARGS=(--hf-overrides '{"is_matryoshka": true}')
+fi
+
 # run worker
 # --runner pooling: required for embedding models.
 # --pooler-config: MEAN pool, no activation — the Qwen3-Embedding default.
@@ -84,6 +93,7 @@ DYN_SYSTEM_PORT=${DYN_SYSTEM_PORT:-8081} \
     --max-model-len "$MAX_MODEL_LEN" \
     --no-enable-prefix-caching \
     --trust-remote-code \
+    "${HF_OVERRIDES_ARGS[@]}" \
     $GPU_MEM_ARGS \
     "${EXTRA_ARGS[@]}" &
 

tests/serve/test_vllm.py

@@ -688,8 +688,12 @@ class VLLMConfig(EngineConfig):
                 repeat_count=1,
                 expected_response=["Generated 3 embeddings with dimension"],
             ),
-            # `dimensions` truncation (Matryoshka). Qwen3-Embedding-0.6B has a
-            # hidden dim of 1024, so the truncated vector should be exactly 128.
+            # `dimensions` reduction (Matryoshka). Qwen3-Embedding-0.6B has a
+            # hidden dim of 1024, so the reduced vector should be exactly 128.
+            # The worker forwards `dimensions` to vLLM's pooler (truncate +
+            # re-normalize); `agg_embed.sh` launches this model with
+            # `--hf-overrides '{"is_matryoshka": true}'` so vLLM accepts the
+            # request (Qwen3-Embedding's config doesn't declare Matryoshka).
             # Built inline because the `embedding_payload()` helper doesn't
             # expose an `extra_body` kwarg yet.
             EmbeddingPayload(