From 8bf064f8d3408ca89cabc2f071adc696314c867e Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Sat, 27 Jun 2026 16:57:47 -0400 Subject: [PATCH] Fixed chunked embedding aggregation with request-id metadata (#46782) Signed-off-by: Taneem Ibrahim --- .../pooling/embed/test_io_processor.py | 92 ++++++++++++++ .../entrypoints/pooling/embed/io_processor.py | 120 +++++++++--------- vllm/entrypoints/pooling/typing.py | 7 + 3 files changed, 157 insertions(+), 62 deletions(-) diff --git a/tests/entrypoints/pooling/embed/test_io_processor.py b/tests/entrypoints/pooling/embed/test_io_processor.py index fbee91fc48..5a7a8aab2a 100644 --- a/tests/entrypoints/pooling/embed/test_io_processor.py +++ b/tests/entrypoints/pooling/embed/test_io_processor.py @@ -3,6 +3,7 @@ """Unit tests for EmbedIOProcessor.""" import pytest +import torch from pydantic import TypeAdapter, ValidationError from vllm import PoolingParams @@ -19,6 +20,7 @@ from vllm.entrypoints.pooling.embed.protocol import ( EmbeddingRequest, ) from vllm.entrypoints.pooling.typing import PoolingServeContext +from vllm.outputs import PoolingOutput, PoolingRequestOutput class TestEmbeddingRequestParsing: @@ -398,6 +400,96 @@ class TestValidateInputType: handler._validate_input_type("z") +class TestChunkedEmbeddingProcessing: + """Unit tests for chunked embedding aggregation.""" + + class _FakeModelConfig: + max_model_len = 3 + + @classmethod + def _make_handler(cls): + handler = object.__new__(EmbedIOProcessor) + handler.model_config = cls._FakeModelConfig() + return handler + + @staticmethod + def _make_context() -> PoolingServeContext[EmbeddingCompletionRequest]: + request = TypeAdapter(EmbeddingRequest).validate_python( + { + "model": "test", + "input": [[0, 1, 2, 3, 4], [10, 11]], + } + ) + assert isinstance(request, EmbeddingCompletionRequest) + return PoolingServeContext( + request=request, + pooling_params=PoolingParams(), + model_name="test", + request_id="embd-client-prompt-999-chunk-888", + engine_inputs=[ + {"prompt_token_ids": [0, 1, 2, 3, 4]}, + {"prompt_token_ids": [10, 11]}, + ], + ) + + @staticmethod + def _make_output( + request_id: str, + prompt_token_ids: list[int], + embedding: list[float], + ) -> PoolingRequestOutput: + return PoolingRequestOutput( + request_id=request_id, + outputs=PoolingOutput(data=torch.tensor(embedding)), + prompt_token_ids=prompt_token_ids, + num_cached_tokens=0, + finished=True, + ) + + def test_aggregation_uses_metadata_not_request_id_parsing(self): + handler = self._make_handler() + ctx = self._make_context() + + handler._pre_process_chunked(ctx) + + assert ctx.prompt_request_ids == [ + "embd-client-prompt-999-chunk-888-prompt-0-chunk-0", + "embd-client-prompt-999-chunk-888-prompt-0-chunk-1", + "embd-client-prompt-999-chunk-888-prompt-1-chunk-0", + ] + assert ctx.chunked_embedding_metadata is not None + assert [ + (item.prompt_index, item.chunk_index) + for item in ctx.chunked_embedding_metadata + ] == [(0, 0), (0, 1), (1, 0)] + + ctx.final_res_batch = [ + self._make_output(ctx.prompt_request_ids[0], [0, 1, 2], [1.0, 1.0]), + self._make_output(ctx.prompt_request_ids[1], [3, 4], [4.0, 7.0]), + self._make_output(ctx.prompt_request_ids[2], [10, 11], [9.0, 9.0]), + ] + + handler._post_process_chunked(ctx) + + assert len(ctx.final_res_batch) == 2 + assert ctx.final_res_batch[0].request_id == ( + "embd-client-prompt-999-chunk-888-prompt-0" + ) + assert ctx.final_res_batch[0].prompt_token_ids == [0, 1, 2, 3, 4] + assert torch.allclose( + ctx.final_res_batch[0].outputs.data, + torch.tensor([2.2, 3.4]), + ) + assert ctx.final_res_batch[1].request_id == ( + "embd-client-prompt-999-chunk-888-prompt-1" + ) + assert ctx.final_res_batch[1].prompt_token_ids == [10, 11] + assert torch.allclose( + ctx.final_res_batch[1].outputs.data, + torch.tensor([9.0, 9.0]), + ) + + class TestPreProcessCohereOnline: """Unit tests for EmbedIOProcessor._pre_process_cohere_online.""" diff --git a/vllm/entrypoints/pooling/embed/io_processor.py b/vllm/entrypoints/pooling/embed/io_processor.py index d2e6f23c14..ec52e2efd6 100644 --- a/vllm/entrypoints/pooling/embed/io_processor.py +++ b/vllm/entrypoints/pooling/embed/io_processor.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Sequence +from dataclasses import dataclass from typing import Any, Literal, cast import torch @@ -27,6 +28,7 @@ from vllm.utils.mistral import is_mistral_tokenizer from ..base.io_processor import PoolingIOProcessor from ..scoring.io_processor import JinaRankingIOProcessorMixin from ..typing import ( + ChunkedEmbeddingMetadata, OfflineInputsContext, PoolingChatLikeRequest, PoolingCompletionLikeRequest, @@ -46,6 +48,12 @@ from .protocol import ( logger = init_logger(__name__) +@dataclass +class _ChunkedPromptAggregator: + weighted_sum: torch.Tensor | None = None + total_weight: int = 0 + + class EmbedIOProcessor(PoolingIOProcessor): name = "embed" @@ -113,6 +121,7 @@ class EmbedIOProcessor(PoolingIOProcessor): max_model_len = self.model_config.max_model_len chunked_engine_inputs: list[EngineInput] = [] prompt_request_ids: list[str] = [] + chunked_embedding_metadata: list[ChunkedEmbeddingMetadata] = [] for prompt_idx, engine_input in enumerate(ctx.engine_inputs): token_ids = engine_input.get("prompt_token_ids", None) if token_ids is None: @@ -132,9 +141,16 @@ class EmbedIOProcessor(PoolingIOProcessor): prompt_request_ids.append( f"{request_id}-prompt-{prompt_idx}-chunk-{chunk_idx}" ) + chunked_embedding_metadata.append( + ChunkedEmbeddingMetadata( + prompt_index=prompt_idx, + chunk_index=chunk_idx, + ) + ) ctx.engine_inputs = chunked_engine_inputs ctx.prompt_request_ids = prompt_request_ids + ctx.chunked_embedding_metadata = chunked_embedding_metadata return None @@ -142,66 +158,48 @@ class EmbedIOProcessor(PoolingIOProcessor): # Online aggregation for chunked requests to # minimize memory usage # Track aggregation state for each prompt - prompt_aggregators: dict[int, dict[str, Any]] = {} - short_prompts_results: dict[int, PoolingRequestOutput] = {} - for result_idx, result in enumerate(ctx.final_res_batch): - if "-chunk-" not in result.request_id: - # Non-chunked result - extract prompt_idx from request_id - parts = result.request_id.split("-") - try: - # Last part should be prompt index - prompt_idx = int(parts[-1]) - except (ValueError, IndexError): - prompt_idx = result_idx # Fallback to result_idx + if ctx.chunked_embedding_metadata is None: + raise ValueError("Chunked embedding metadata not available") + if len(ctx.chunked_embedding_metadata) != len(ctx.final_res_batch): + raise ValueError( + "Chunked embedding metadata count does not match result count" + ) - short_prompts_results[prompt_idx] = result + prompt_aggregators: dict[int, _ChunkedPromptAggregator] = {} + for result, chunk_metadata in zip( + ctx.final_res_batch, ctx.chunked_embedding_metadata + ): + prompt_idx = chunk_metadata.prompt_index + aggregator = prompt_aggregators.setdefault( + prompt_idx, _ChunkedPromptAggregator() + ) + + # MEAN pooling with online weighted averaging + # Ensure result is PoolingRequestOutput + # for embedding processing + if not isinstance(result, PoolingRequestOutput): + raise ValueError( + f"Expected PoolingRequestOutput for " + f"chunked embedding, got " + f"{type(result).__name__}" + ) + if result.prompt_token_ids is None: + raise ValueError( + "prompt_token_ids cannot be None for chunked processing" + ) + + weight = len(result.prompt_token_ids) + embedding_data = result.outputs.data + weighted_embedding = embedding_data.to(dtype=torch.float32) * weight + + if aggregator.weighted_sum is None: + # First chunk + aggregator.weighted_sum = weighted_embedding else: - # Extract prompt_idx from chunked request_id - parts = result.request_id.split("-") - try: - prompt_idx = int(parts[parts.index("prompt") + 1]) - except (ValueError, IndexError): - # Fallback: extract from result_idx if parsing fails - prompt_idx = result_idx + # Accumulate + aggregator.weighted_sum += weighted_embedding - # Initialize aggregator for this prompt if needed - if prompt_idx not in prompt_aggregators: - prompt_aggregators[prompt_idx] = { - "weighted_sum": None, - "total_weight": 0, - "chunk_count": 0, - "request_id": result.request_id.split("-chunk-")[0], - } - - aggregator = prompt_aggregators[prompt_idx] - - # MEAN pooling with online weighted averaging - # Ensure result is PoolingRequestOutput - # for embedding processing - if not isinstance(result, PoolingRequestOutput): - raise ValueError( - f"Expected PoolingRequestOutput for " - f"chunked embedding, got " - f"{type(result).__name__}" - ) - if result.prompt_token_ids is None: - raise ValueError( - "prompt_token_ids cannot be None for chunked processing" - ) - - weight = len(result.prompt_token_ids) - embedding_data = result.outputs.data - weighted_embedding = embedding_data.to(dtype=torch.float32) * weight - - if aggregator["weighted_sum"] is None: - # First chunk - aggregator["weighted_sum"] = weighted_embedding - else: - # Accumulate - aggregator["weighted_sum"] += weighted_embedding - - aggregator["total_weight"] += weight - aggregator["chunk_count"] += 1 + aggregator.total_weight += weight if ctx.original_engine_inputs is None: raise ValueError("Original engine inputs not available") @@ -216,8 +214,8 @@ class EmbedIOProcessor(PoolingIOProcessor): # Finalize MEAN aggregation for this chunked prompt aggregator = prompt_aggregators[prompt_idx] - weighted_sum = aggregator["weighted_sum"] - total_weight = aggregator["total_weight"] + weighted_sum = aggregator.weighted_sum + total_weight = aggregator.total_weight if ( weighted_sum is not None @@ -243,7 +241,7 @@ class EmbedIOProcessor(PoolingIOProcessor): original_token_ids = cast(list[int], token_ids) pooling_request_output = PoolingRequestOutput( - request_id=aggregator["request_id"], + request_id=f"{ctx.request_id}-prompt-{prompt_idx}", prompt_token_ids=original_token_ids, outputs=pooling_output_data, num_cached_tokens=0, @@ -255,8 +253,6 @@ class EmbedIOProcessor(PoolingIOProcessor): raise ValueError( f"Failed to aggregate chunks for prompt {prompt_idx}" ) - elif prompt_idx in short_prompts_results: - final_res_batch.append(short_prompts_results[prompt_idx]) else: raise ValueError(f"Result not found for prompt {prompt_idx}") diff --git a/vllm/entrypoints/pooling/typing.py b/vllm/entrypoints/pooling/typing.py index 2cf3849005..54d02c5b61 100644 --- a/vllm/entrypoints/pooling/typing.py +++ b/vllm/entrypoints/pooling/typing.py @@ -70,6 +70,12 @@ AnyPoolingResponse: TypeAlias = ( PoolingRequestT = TypeVar("PoolingRequestT", bound=AnyPoolingRequest) +@dataclass(kw_only=True) +class ChunkedEmbeddingMetadata: + prompt_index: int + chunk_index: int + + @dataclass(kw_only=True) class PoolingServeContext(Generic[PoolingRequestT]): model_config = ConfigDict(arbitrary_types_allowed=True) @@ -91,6 +97,7 @@ class PoolingServeContext(Generic[PoolingRequestT]): ## for Long Text Embedding with Chunked Processing original_engine_inputs: Sequence[EngineInput] | None = None + chunked_embedding_metadata: list[ChunkedEmbeddingMetadata] | None = None ## for bi-encoder & late-interaction n_queries: int | None = None