@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 8bf064f8d3

Fixed chunked embedding aggregation with request-id metadata (#46782)

Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>

Taneem Ibrahim committed Jun 27, 2026 at 16:57 UTC 8bf064f8d3408ca89cabc2f071adc696314c867e
3 files changed +155 -60
tests/entrypoints/pooling/embed/test_io_processor.py
+92
@@ -3,6 +3,7 @@
3 """Unit tests for EmbedIOProcessor."""
4
5 import pytest
6 +import torch
7 from pydantic import TypeAdapter, ValidationError
8
9 from vllm import PoolingParams
@@ -19,6 +20,7 @@ from vllm.entrypoints.pooling.embed.protocol import (
20 EmbeddingRequest,
21 )
22 from vllm.entrypoints.pooling.typing import PoolingServeContext
23 +from vllm.outputs import PoolingOutput, PoolingRequestOutput
24
25
26 class TestEmbeddingRequestParsing:
@@ -398,6 +400,96 @@ class TestValidateInputType:
400 handler._validate_input_type("z")
401
402
403 +class TestChunkedEmbeddingProcessing:
404 + """Unit tests for chunked embedding aggregation."""
405 +
406 + class _FakeModelConfig:
407 + max_model_len = 3
408 +
409 + @classmethod
410 + def _make_handler(cls):
411 + handler = object.__new__(EmbedIOProcessor)
412 + handler.model_config = cls._FakeModelConfig()
413 + return handler
414 +
415 + @staticmethod
416 + def _make_context() -> PoolingServeContext[EmbeddingCompletionRequest]:
417 + request = TypeAdapter(EmbeddingRequest).validate_python(
418 + {
419 + "model": "test",
420 + "input": [[0, 1, 2, 3, 4], [10, 11]],
421 + }
422 + )
423 + assert isinstance(request, EmbeddingCompletionRequest)
424 + return PoolingServeContext(
425 + request=request,
426 + pooling_params=PoolingParams(),
427 + model_name="test",
428 + request_id="embd-client-prompt-999-chunk-888",
429 + engine_inputs=[
430 + {"prompt_token_ids": [0, 1, 2, 3, 4]},
431 + {"prompt_token_ids": [10, 11]},
432 + ],
433 + )
434 +
435 + @staticmethod
436 + def _make_output(
437 + request_id: str,
438 + prompt_token_ids: list[int],
439 + embedding: list[float],
440 + ) -> PoolingRequestOutput:
441 + return PoolingRequestOutput(
442 + request_id=request_id,
443 + outputs=PoolingOutput(data=torch.tensor(embedding)),
444 + prompt_token_ids=prompt_token_ids,
445 + num_cached_tokens=0,
446 + finished=True,
447 + )
448 +
449 + def test_aggregation_uses_metadata_not_request_id_parsing(self):
450 + handler = self._make_handler()
451 + ctx = self._make_context()
452 +
453 + handler._pre_process_chunked(ctx)
454 +
455 + assert ctx.prompt_request_ids == [
456 + "embd-client-prompt-999-chunk-888-prompt-0-chunk-0",
457 + "embd-client-prompt-999-chunk-888-prompt-0-chunk-1",
458 + "embd-client-prompt-999-chunk-888-prompt-1-chunk-0",
459 + ]
460 + assert ctx.chunked_embedding_metadata is not None
461 + assert [
462 + (item.prompt_index, item.chunk_index)
463 + for item in ctx.chunked_embedding_metadata
464 + ] == [(0, 0), (0, 1), (1, 0)]
465 +
466 + ctx.final_res_batch = [
467 + self._make_output(ctx.prompt_request_ids[0], [0, 1, 2], [1.0, 1.0]),
468 + self._make_output(ctx.prompt_request_ids[1], [3, 4], [4.0, 7.0]),
469 + self._make_output(ctx.prompt_request_ids[2], [10, 11], [9.0, 9.0]),
470 + ]
471 +
472 + handler._post_process_chunked(ctx)
473 +
474 + assert len(ctx.final_res_batch) == 2
475 + assert ctx.final_res_batch[0].request_id == (
476 + "embd-client-prompt-999-chunk-888-prompt-0"
477 + )
478 + assert ctx.final_res_batch[0].prompt_token_ids == [0, 1, 2, 3, 4]
479 + assert torch.allclose(
480 + ctx.final_res_batch[0].outputs.data,
481 + torch.tensor([2.2, 3.4]),
482 + )
483 + assert ctx.final_res_batch[1].request_id == (
484 + "embd-client-prompt-999-chunk-888-prompt-1"
485 + )
486 + assert ctx.final_res_batch[1].prompt_token_ids == [10, 11]
487 + assert torch.allclose(
488 + ctx.final_res_batch[1].outputs.data,
489 + torch.tensor([9.0, 9.0]),
490 + )
491 +
492 +
493 class TestPreProcessCohereOnline:
494 """Unit tests for EmbedIOProcessor._pre_process_cohere_online."""
495
vllm/entrypoints/pooling/embed/io_processor.py
+56 -60
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from collections.abc import Sequence
4 +from dataclasses import dataclass
5 from typing import Any, Literal, cast
6
7 import torch
@@ -27,6 +28,7 @@ from vllm.utils.mistral import is_mistral_tokenizer
28 from ..base.io_processor import PoolingIOProcessor
29 from ..scoring.io_processor import JinaRankingIOProcessorMixin
30 from ..typing import (
31 + ChunkedEmbeddingMetadata,
32 OfflineInputsContext,
33 PoolingChatLikeRequest,
34 PoolingCompletionLikeRequest,
@@ -46,6 +48,12 @@ from .protocol import (
48 logger = init_logger(__name__)
49
50
51 +@dataclass
52 +class _ChunkedPromptAggregator:
53 + weighted_sum: torch.Tensor | None = None
54 + total_weight: int = 0
55 +
56 +
57 class EmbedIOProcessor(PoolingIOProcessor):
58 name = "embed"
59
@@ -113,6 +121,7 @@ class EmbedIOProcessor(PoolingIOProcessor):
121 max_model_len = self.model_config.max_model_len
122 chunked_engine_inputs: list[EngineInput] = []
123 prompt_request_ids: list[str] = []
124 + chunked_embedding_metadata: list[ChunkedEmbeddingMetadata] = []
125 for prompt_idx, engine_input in enumerate(ctx.engine_inputs):
126 token_ids = engine_input.get("prompt_token_ids", None)
127 if token_ids is None:
@@ -132,9 +141,16 @@ class EmbedIOProcessor(PoolingIOProcessor):
141 prompt_request_ids.append(
142 f"{request_id}-prompt-{prompt_idx}-chunk-{chunk_idx}"
143 )
144 + chunked_embedding_metadata.append(
145 + ChunkedEmbeddingMetadata(
146 + prompt_index=prompt_idx,
147 + chunk_index=chunk_idx,
148 + )
149 + )
150
151 ctx.engine_inputs = chunked_engine_inputs
152 ctx.prompt_request_ids = prompt_request_ids
153 + ctx.chunked_embedding_metadata = chunked_embedding_metadata
154
155 return None
156
@@ -142,66 +158,48 @@ class EmbedIOProcessor(PoolingIOProcessor):
158 # Online aggregation for chunked requests to
159 # minimize memory usage
160 # Track aggregation state for each prompt
145 - prompt_aggregators: dict[int, dict[str, Any]] = {}
146 - short_prompts_results: dict[int, PoolingRequestOutput] = {}
147 - for result_idx, result in enumerate(ctx.final_res_batch):
148 - if "-chunk-" not in result.request_id:
149 - # Non-chunked result - extract prompt_idx from request_id
150 - parts = result.request_id.split("-")
151 - try:
152 - # Last part should be prompt index
153 - prompt_idx = int(parts[-1])
154 - except (ValueError, IndexError):
155 - prompt_idx = result_idx # Fallback to result_idx
156 -
157 - short_prompts_results[prompt_idx] = result
158 - else:
159 - # Extract prompt_idx from chunked request_id
160 - parts = result.request_id.split("-")
161 - try:
162 - prompt_idx = int(parts[parts.index("prompt") + 1])
163 - except (ValueError, IndexError):
164 - # Fallback: extract from result_idx if parsing fails
165 - prompt_idx = result_idx
166 -
167 - # Initialize aggregator for this prompt if needed
168 - if prompt_idx not in prompt_aggregators:
169 - prompt_aggregators[prompt_idx] = {
170 - "weighted_sum": None,
171 - "total_weight": 0,
172 - "chunk_count": 0,
173 - "request_id": result.request_id.split("-chunk-")[0],
174 - }
161 + if ctx.chunked_embedding_metadata is None:
162 + raise ValueError("Chunked embedding metadata not available")
163 + if len(ctx.chunked_embedding_metadata) != len(ctx.final_res_batch):
164 + raise ValueError(
165 + "Chunked embedding metadata count does not match result count"
166 + )
167
176 - aggregator = prompt_aggregators[prompt_idx]
168 + prompt_aggregators: dict[int, _ChunkedPromptAggregator] = {}
169 + for result, chunk_metadata in zip(
170 + ctx.final_res_batch, ctx.chunked_embedding_metadata
171 + ):
172 + prompt_idx = chunk_metadata.prompt_index
173 + aggregator = prompt_aggregators.setdefault(
174 + prompt_idx, _ChunkedPromptAggregator()
175 + )
176
178 - # MEAN pooling with online weighted averaging
179 - # Ensure result is PoolingRequestOutput
180 - # for embedding processing
181 - if not isinstance(result, PoolingRequestOutput):
182 - raise ValueError(
183 - f"Expected PoolingRequestOutput for "
184 - f"chunked embedding, got "
185 - f"{type(result).__name__}"
186 - )
187 - if result.prompt_token_ids is None:
188 - raise ValueError(
189 - "prompt_token_ids cannot be None for chunked processing"
190 - )
177 + # MEAN pooling with online weighted averaging
178 + # Ensure result is PoolingRequestOutput
179 + # for embedding processing
180 + if not isinstance(result, PoolingRequestOutput):
181 + raise ValueError(
182 + f"Expected PoolingRequestOutput for "
183 + f"chunked embedding, got "
184 + f"{type(result).__name__}"
185 + )
186 + if result.prompt_token_ids is None:
187 + raise ValueError(
188 + "prompt_token_ids cannot be None for chunked processing"
189 + )
190
192 - weight = len(result.prompt_token_ids)
193 - embedding_data = result.outputs.data
194 - weighted_embedding = embedding_data.to(dtype=torch.float32) * weight
191 + weight = len(result.prompt_token_ids)
192 + embedding_data = result.outputs.data
193 + weighted_embedding = embedding_data.to(dtype=torch.float32) * weight
194
196 - if aggregator["weighted_sum"] is None:
197 - # First chunk
198 - aggregator["weighted_sum"] = weighted_embedding
199 - else:
200 - # Accumulate
201 - aggregator["weighted_sum"] += weighted_embedding
195 + if aggregator.weighted_sum is None:
196 + # First chunk
197 + aggregator.weighted_sum = weighted_embedding
198 + else:
199 + # Accumulate
200 + aggregator.weighted_sum += weighted_embedding
201
203 - aggregator["total_weight"] += weight
204 - aggregator["chunk_count"] += 1
202 + aggregator.total_weight += weight
203
204 if ctx.original_engine_inputs is None:
205 raise ValueError("Original engine inputs not available")
@@ -216,8 +214,8 @@ class EmbedIOProcessor(PoolingIOProcessor):
214 # Finalize MEAN aggregation for this chunked prompt
215 aggregator = prompt_aggregators[prompt_idx]
216
219 - weighted_sum = aggregator["weighted_sum"]
220 - total_weight = aggregator["total_weight"]
217 + weighted_sum = aggregator.weighted_sum
218 + total_weight = aggregator.total_weight
219
220 if (
221 weighted_sum is not None
@@ -243,7 +241,7 @@ class EmbedIOProcessor(PoolingIOProcessor):
241
242 original_token_ids = cast(list[int], token_ids)
243 pooling_request_output = PoolingRequestOutput(
246 - request_id=aggregator["request_id"],
244 + request_id=f"{ctx.request_id}-prompt-{prompt_idx}",
245 prompt_token_ids=original_token_ids,
246 outputs=pooling_output_data,
247 num_cached_tokens=0,
@@ -255,8 +253,6 @@ class EmbedIOProcessor(PoolingIOProcessor):
253 raise ValueError(
254 f"Failed to aggregate chunks for prompt {prompt_idx}"
255 )
258 - elif prompt_idx in short_prompts_results:
259 - final_res_batch.append(short_prompts_results[prompt_idx])
256 else:
257 raise ValueError(f"Result not found for prompt {prompt_idx}")
258
vllm/entrypoints/pooling/typing.py
+7
@@ -70,6 +70,12 @@ AnyPoolingResponse: TypeAlias = (
70 PoolingRequestT = TypeVar("PoolingRequestT", bound=AnyPoolingRequest)
71
72
73 +@dataclass(kw_only=True)
74 +class ChunkedEmbeddingMetadata:
75 + prompt_index: int
76 + chunk_index: int
77 +
78 +
79 @dataclass(kw_only=True)
80 class PoolingServeContext(Generic[PoolingRequestT]):
81 model_config = ConfigDict(arbitrary_types_allowed=True)
@@ -91,6 +97,7 @@ class PoolingServeContext(Generic[PoolingRequestT]):
97
98 ## for Long Text Embedding with Chunked Processing
99 original_engine_inputs: Sequence[EngineInput] | None = None
100 + chunked_embedding_metadata: list[ChunkedEmbeddingMetadata] | None = None
101
102 ## for bi-encoder & late-interaction
103 n_queries: int | None = None