iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
This commit is contained in:
Taneem Ibrahim
2026-06-27 20:57:47 +00:00
committed by GitHub
parent ea2ead1db3
commit 8bf064f8d3
3 changed files with 157 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/entrypoints/pooling/embed/test_io_processor.py
+92
View File
@@ -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."""
vllm/entrypoints/pooling/embed/io_processor.py
+58 -62
View File
@@ -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}")
vllm/entrypoints/pooling/typing.py
+7
View File
@@ -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