diff --git a/docs/models/pooling_models/token_embed.md b/docs/models/pooling_models/token_embed.md index 0c2a322e80..2bc51f4bec 100644 --- a/docs/models/pooling_models/token_embed.md +++ b/docs/models/pooling_models/token_embed.md @@ -90,6 +90,7 @@ The following [pooling parameters][vllm.PoolingParams] are supported. ```python --8<-- "vllm/pooling_params.py:common-pooling-params" --8<-- "vllm/pooling_params.py:embed-pooling-params" +--8<-- "vllm/pooling_params.py:token-embed-pooling-params" ``` ### `LLM.encode` @@ -108,6 +109,39 @@ data = output.outputs.data print(f"Data: {data!r}") ``` +### Fixed-length late chunking + +Late chunking encodes the full input before averaging contextual token representations within each chunk. The chunk means pass through the model's embedding head, including its projection and optional normalization. + +This mode is currently available through the offline `LLM.encode` API for plain-text input. With an instance configured for `token_embed`, enable it per request: + +```python +from vllm import PoolingParams +from vllm.pooling_params import LateChunkingParams + +outputs = llm.encode( + [document], + pooling_task="token_embed", + pooling_params=PoolingParams( + late_chunking_params=LateChunkingParams(chunk_size=256) + ), +) +vectors = outputs[0].outputs.data # [num_chunks, embedding_dimension] +chunks = outputs[0].late_chunking.chunks +``` + +For `N` model-input tokens and chunk size `C`, the output contains `ceil(N / C)` vectors. The last chunk uses its actual length. Special tokens and instruction prefixes count toward both chunk boundaries and means. Omitting `late_chunking_params` preserves ordinary token embeddings. + +Each vector row corresponds to one entry in `chunks`. Its `token_range` and `char_range` are zero-based, half-open ranges. Character ranges index the submitted Python string; use `document[start:end]` to retrieve the source text. A chunk containing only special tokens has `char_range=None`. + +Use nonempty text within the model's context window and a fast tokenizer that provides aligned source offsets. Input truncation, explicit padding and renderer text normalization are not supported. The custom document/query outputs of `JinaForRanking` are also unsupported. Late chunking does not extend the model's context window or reduce its attention computation. + +The [offline example](../../../examples/pooling/token_embed/late_chunking_offline.py) loads an E5 model and prints the vector shape, token ranges and source excerpts for two documents: + +```bash +python examples/pooling/token_embed/late_chunking_offline.py --chunk-size 16 +``` + ### `LLM.score` The [score][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.score] method outputs similarity scores between sentence pairs. diff --git a/examples/pooling/token_embed/late_chunking_offline.py b/examples/pooling/token_embed/late_chunking_offline.py new file mode 100644 index 0000000000..4f82b17cce --- /dev/null +++ b/examples/pooling/token_embed/late_chunking_offline.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from argparse import Namespace + +from vllm import LLM, EngineArgs, PoolingParams +from vllm.config import PoolerConfig +from vllm.pooling_params import LateChunkingParams +from vllm.utils.argparse_utils import FlexibleArgumentParser + + +def parse_args(): + parser = FlexibleArgumentParser() + parser = EngineArgs.add_cli_args(parser) + parser.add_argument( + "--chunk-size", + type=int, + default=16, + help="Number of model-input tokens per chunk, including special tokens.", + ) + parser.set_defaults( + model="intfloat/multilingual-e5-small", + runner="pooling", + pooler_config=PoolerConfig(task="token_embed"), + max_model_len=512, + enforce_eager=True, + ) + return parser.parse_args() + + +def main(args: Namespace): + pooling_params = PoolingParams( + late_chunking_params=LateChunkingParams(chunk_size=args.chunk_size) + ) + engine_args = vars(args).copy() + del engine_args["chunk_size"] + llm = LLM(**engine_args) + + # E5 document inputs use the "passage: " prefix; it also counts toward chunks. + documents = [ + "passage: Berlin is the capital of Germany. The city has many museums. " + "They display art and historical collections from around the world.", + "passage: Kyoto is a city in Japan. Its historic temples attract visitors " + "throughout the year.", + ] + outputs = llm.encode( + documents, + pooling_task="token_embed", + pooling_params=pooling_params, + ) + + for doc_index, (document, output) in enumerate(zip(documents, outputs)): + vectors = output.outputs.data + metadata = output.late_chunking + assert metadata is not None + print(f"\nDocument {doc_index}: {metadata.input_tokens} input tokens") + print(f"Chunk vectors: {tuple(vectors.shape)}") + for index, chunk in enumerate(metadata.chunks): + # vectors[index] corresponds to this chunk and its source-text range. + if chunk.char_range is None: + excerpt = "" + else: + start, end = chunk.char_range + excerpt = document[start:end] + print( + f" Chunk {index}: tokens={chunk.token_range}, chars={chunk.char_range}" + ) + print(f" {excerpt!r}") + + +if __name__ == "__main__": + main(parse_args()) diff --git a/tests/entrypoints/pooling/test_late_chunking.py b/tests/entrypoints/pooling/test_late_chunking.py new file mode 100644 index 0000000000..eaec62ac19 --- /dev/null +++ b/tests/entrypoints/pooling/test_late_chunking.py @@ -0,0 +1,314 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest import mock +from unittest.mock import Mock + +import pytest +import torch + +from vllm import PoolingParams +from vllm.config import PoolerConfig +from vllm.entrypoints.pooling.embed.io_processor import ( + JinaRankingTokenEmbedIOProcessor, + TokenEmbedIOProcessor, +) +from vllm.entrypoints.pooling.late_chunking import build_late_chunking_metadata +from vllm.entrypoints.pooling.offline import PoolingOfflineMixin +from vllm.entrypoints.pooling.typing import ( + OfflineEncodeInputsContext, + OfflineOutputsContext, +) +from vllm.exceptions import VLLMValidationError +from vllm.outputs import ( + LateChunk, + PoolingOutput, + PoolingRequestOutput, + RequestError, +) +from vllm.pooling_params import LateChunkingParams +from vllm.renderers import TokenizeParams + + +@pytest.fixture +def late_chunk_processor(): + processor = object.__new__(TokenEmbedIOProcessor) + processor.model_config = SimpleNamespace( + is_encoder_decoder=False, + architecture="NomicBertModel", + model_impl="auto", + is_matryoshka=False, + hf_config=SimpleNamespace(), + pooler_config=PoolerConfig(seq_pooling_type="MEAN", tok_pooling_type="ALL"), + ) + processor.vllm_config = SimpleNamespace( + model_config=processor.model_config, + cache_config=SimpleNamespace(enable_prefix_caching=False), + scheduler_config=SimpleNamespace(enable_chunked_prefill=False), + lora_config=None, + ) + processor.renderer = Mock() + processor.renderer.default_cmpl_tok_params = TokenizeParams(max_total_tokens=32) + processor.renderer.render_cmpl.side_effect = lambda **_: [ + { + "type": "token", + "prompt_token_ids": [101, 1, 2, 102], + "prompt_token_offsets": [(0, 0), (0, 1), (2, 3), (0, 0)], + } + ] + return processor + + +def _late_chunk_context(prompts="a b", **kwargs): + return OfflineEncodeInputsContext( + pooling_task="token_embed", + tokenization_kwargs=kwargs or None, + lora_request=None, + priorities=None, + prompts=prompts, + pooling_params=PoolingParams( + late_chunking_params=LateChunkingParams(chunk_size=2) + ), + ) + + +def test_late_chunk_render_requests_offsets_once(late_chunk_processor): + processor = late_chunk_processor + ctx = _late_chunk_context() + factory, count = processor.get_request_factory_offline(ctx) + render_params = next(factory()) + result = processor.render(render_params) + assert count == 1 + processor.renderer.render_cmpl.assert_called_once() + assert processor.renderer.render_cmpl.call_args.kwargs[ + "tok_params" + ].return_token_offsets + assert not render_params["tok_params"].return_token_offsets + assert "prompt_token_offsets" not in result["prompts"] + assert "late_chunking" not in result + assert [c.char_range for c in ctx.late_chunking[0].metadata.chunks] == [ + (0, 1), + (2, 3), + ] + assert ctx.late_chunking[0].metadata.input_tokens == 4 + + +@pytest.mark.parametrize("as_sequence", [False, True]) +@pytest.mark.parametrize("chunk_size", [None, 2]) +def test_jina_ranking_rejects_chunking_before_tokenization( + late_chunk_processor, as_sequence, chunk_size +): + processor = object.__new__(JinaRankingTokenEmbedIOProcessor) + processor.__dict__.update(late_chunk_processor.__dict__) + ctx = _late_chunk_context(prompts=["document", "query"]) + params = PoolingParams( + late_chunking_params=LateChunkingParams(chunk_size) if chunk_size else None + ) + ctx.pooling_params = [params] if as_sequence else params + factory, count = processor.get_request_factory_offline(ctx) + assert count == 1 + request = next(factory()) + if chunk_size is not None: + with pytest.raises(VLLMValidationError, match="JinaForRanking.*late chunking"): + processor.render(request) + late_chunk_processor.renderer.render_cmpl.assert_not_called() + else: + result = processor.render(request) + assert result["prompts"]["type"] == "token" + assert result["prompts"]["prompt_token_ids"] == [101, 1, 2, 102] + late_chunk_processor.renderer.render_cmpl.assert_called_once() + assert not request["tok_params"].return_token_offsets + + +@pytest.mark.parametrize( + "prompts, kwargs, error", + [ + pytest.param("", {}, "plain-text", id="empty"), + pytest.param([1, 2], {}, "plain-text", id="token-ids"), + pytest.param( + {"prompt": "abc", "multi_modal_data": {}}, + {}, + "plain-text", + id="multimodal", + ), + ] + + [ + pytest.param("abc", {key: value}, "does not support", id=key) + for key, value in ( + ("truncate_prompt_tokens", 2), + ("truncation", True), + ("pad_prompt_tokens", 8), + ("padding", True), + ("do_lower_case", True), + ) + ], +) +def test_late_chunking_rejects_unsupported_input_before_render( + late_chunk_processor, prompts, kwargs, error +): + processor = late_chunk_processor + with pytest.raises(VLLMValidationError, match=error): + factory, _ = processor.get_request_factory_offline( + _late_chunk_context(prompts, **kwargs) + ) + processor.render(next(factory())) + processor.renderer.render_cmpl.assert_not_called() + + +def test_late_chunk_ranges_keep_unicode_overlaps_and_special_only_chunks(): + text = "中 😀 e\u0301" + offsets = [(0, 0), (0, 1), (2, 3), (2, 3), (4, 6), (0, 0)] + metadata = build_late_chunking_metadata(text, len(offsets), offsets, 1) + assert [c.char_range for c in metadata.chunks] == [ + None, + (0, 1), + (2, 3), + (2, 3), + (4, 6), + None, + ] + assert [ + text[slice(*c.char_range)] if c.char_range else None for c in metadata.chunks + ] == [None, "中", "😀", "😀", "e\u0301", None] + chunks = build_late_chunking_metadata(text, len(offsets), offsets, 4).chunks + assert [(c.token_range, c.char_range) for c in chunks] == [ + ((0, 4), (0, 3)), + ((4, 6), (4, 6)), + ] + + +@pytest.mark.parametrize( + "offsets", [None, [(0, 1)], [(-1, 1), (0, 0)], [(0, 4), (0, 0)], [(2, 3), (0, 1)]] +) +def test_late_chunk_ranges_reject_missing_or_invalid_offsets(offsets): + with pytest.raises(VLLMValidationError, match="offsets"): + build_late_chunking_metadata("abc", 2, offsets, 2) + + +def _mock_chunk_llm(processor, outputs): + llm = mock.Mock(spec=PoolingOfflineMixin) + llm.encode = PoolingOfflineMixin.encode.__get__(llm) + llm._run_tiling_engine = PoolingOfflineMixin._run_tiling_engine.__get__(llm) + llm.pooling_io_processors = {"token_embed": processor} + llm._executor = SimpleNamespace(map=map) + llm.llm_engine = mock.Mock() + llm.llm_engine.vllm_config.scheduler_config.max_num_seqs = 2 + llm.llm_engine.has_unfinished_requests.return_value = False + llm.llm_engine.step.side_effect = outputs + llm._render_and_add_requests = mock.Mock( + side_effect=lambda **kw: [ + str(i) + "-internal" for i in range(len(kw["params"])) + ] + ) + return llm + + +def _chunk_output(request_id, rows=2, **kwargs): + return PoolingRequestOutput( + str(request_id), + PoolingOutput(torch.ones(rows, 4)), + [101, 1, 2, 102], + 0, + True, + **kwargs, + ) + + +def test_late_chunk_mapping_preserves_order_mixed_requests_and_errors( + late_chunk_processor, +): + error = RequestError("test_error", "original error") + failed = _chunk_output(3, rows=0, error=error) + llm = _mock_chunk_llm(late_chunk_processor, RuntimeError("step failed")) + params = [ + PoolingParams(late_chunking_params=LateChunkingParams(2)), + PoolingParams(), + PoolingParams(late_chunking_params=LateChunkingParams(1)), + PoolingParams(late_chunking_params=LateChunkingParams(2)), + ] + # A failed encode call must not leave metadata for the following batch. + with pytest.raises(RuntimeError, match="step failed"): + llm.encode( + ["a b"] * 2, + pooling_task="token_embed", + pooling_params=params[:2], + use_tqdm=False, + ) + llm.llm_engine.step.side_effect = [ + [ + _chunk_output(2, rows=4), + failed, + _chunk_output(1, rows=4), + _chunk_output(0), + ] + ] + outputs = llm.encode( + ["a b"] * 4, pooling_task="token_embed", pooling_params=params, use_tqdm=False + ) + assert outputs[0].late_chunking.chunks == [ + LateChunk((0, 2), (0, 1)), + LateChunk((2, 4), (2, 3)), + ] + assert outputs[1].late_chunking is None + assert len(outputs[2].late_chunking.chunks) == 4 + assert outputs[3].late_chunking is None + assert outputs[3].error is error + + +def test_late_chunk_mapping_rejects_successful_output_with_wrong_row_count( + late_chunk_processor, +): + llm = _mock_chunk_llm(late_chunk_processor, [[_chunk_output(0, rows=0)]]) + with pytest.raises(ValueError, match="does not match"): + llm.encode( + "a b", + pooling_task="token_embed", + pooling_params=PoolingParams(late_chunking_params=LateChunkingParams(2)), + use_tqdm=False, + ) + # All engine requests have finished before post-processing validates the output. + llm.llm_engine.abort_request.assert_not_called() + + +def test_late_chunk_contexts_remain_isolated_when_rendering_interleaves( + late_chunk_processor, +): + processor = late_chunk_processor + first, second = _late_chunk_context("xy z"), _late_chunk_context() + shared = PoolingParams(late_chunking_params=LateChunkingParams(2)) + first.pooling_params = second.pooling_params = shared + first_factory, _ = processor.get_request_factory_offline(first) + second_factory, _ = processor.get_request_factory_offline(second) + first_request, second_request = next(first_factory()), next(second_factory()) + processor.render(second_request) + assert first.late_chunking[0].metadata is None + processor.renderer.render_cmpl.side_effect = lambda **_: [ + { + "type": "token", + "prompt_token_ids": [101, 1, 2, 102], + "prompt_token_offsets": [(0, 0), (0, 2), (3, 4), (0, 0)], + } + ] + processor.render(first_request) + first_output = processor.post_process_offline( + OfflineOutputsContext([_chunk_output(0)], late_chunking=first.late_chunking) + )[0] + second_output = processor.post_process_offline( + OfflineOutputsContext([_chunk_output(0)], late_chunking=second.late_chunking) + )[0] + assert first_output.late_chunking.chunks[0].char_range == (0, 2) + assert second_output.late_chunking.chunks[0].char_range == (0, 1) + + +def test_plain_token_embed_skips_chunk_metadata(late_chunk_processor): + ctx = _late_chunk_context() + ctx.pooling_params = PoolingParams() + factory, _ = late_chunk_processor.get_request_factory_offline(ctx) + request = next(factory()) + result = late_chunk_processor.render(request) + assert not ctx.late_chunking + assert result["params"].late_chunking_params is None + assert not late_chunk_processor.renderer.render_cmpl.call_args.kwargs[ + "tok_params" + ].return_token_offsets diff --git a/tests/entrypoints/pooling/token_embed/test_offline.py b/tests/entrypoints/pooling/token_embed/test_offline.py index e0bda65422..b65ef6ff8b 100644 --- a/tests/entrypoints/pooling/token_embed/test_offline.py +++ b/tests/entrypoints/pooling/token_embed/test_offline.py @@ -3,9 +3,13 @@ import weakref import pytest +import torch -from vllm import LLM, PoolingRequestOutput +from tests.models.utils import check_embeddings_close +from vllm import LLM, PoolingParams, PoolingRequestOutput from vllm.config import PoolerConfig +from vllm.exceptions import VLLMValidationError +from vllm.pooling_params import LateChunkingParams from vllm.tasks import PoolingTask MODEL_NAME = "intfloat/multilingual-e5-small" @@ -61,3 +65,218 @@ def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm): with pytest.raises(ValueError, match=err_msg): llm.encode(prompt, pooling_task=task, use_tqdm=False) + + +def test_bge_m3_late_chunking_keeps_special_tokens(vllm_runner): + with vllm_runner( + "BAAI/bge-m3", + hf_overrides={"architectures": ["BgeM3EmbeddingModel"]}, + pooler_config=PoolerConfig(task="token_embed"), + dtype="float32", + max_model_len=128, + # Keep encoder batch size fixed for the numerical reference. + max_num_seqs=1, + gpu_memory_utilization=0.5, + enforce_eager=True, + ) as runner: + llm = runner.llm + text = "A document about Berlin and its museums." + sizes = (1, 3, 128, None) + tokens, *outputs = llm.encode( + [text] * len(sizes), + pooling_task="token_embed", + pooling_params=[ + PoolingParams( + use_activation=size not in (1, None), + late_chunking_params=LateChunkingParams(size) if size else None, + ) + for size in sizes + ], + use_tqdm=False, + ) + assert tokens.outputs.data.shape[0] == len(tokens.prompt_token_ids) + assert tokens.late_chunking.chunks[0].char_range is None + for size, output in zip(sizes[1:], outputs): + if size is None: + # Ordinary BGE-M3 token embeddings still drop BOS only. + expected = tokens.outputs.data[1:] + assert output.late_chunking is None + else: + # BGE-M3's projector is linear, so its unnormalized token + # outputs also provide a reference for projected chunk means. + expected = torch.nn.functional.normalize( + torch.stack([p.mean(0) for p in tokens.outputs.data.split(size)]), + dim=-1, + ) + assert len(output.late_chunking.chunks) == len(expected) + torch.testing.assert_close( + output.outputs.data, expected, atol=2e-5, rtol=2e-4 + ) + + +def test_nomic_late_chunking_without_special_tokens(vllm_runner): + """Keep source ranges without special tokens and reject empty tokenization.""" + with vllm_runner( + "nomic-ai/nomic-embed-text-v1", + revision="720244025c1a7e15661a174c63cce63c8218e52b", + trust_remote_code=True, + runner="pooling", + pooler_config=PoolerConfig(task="token_embed"), + dtype="float32", + max_model_len=128, + max_num_seqs=2, + gpu_memory_utilization=0.3, + enforce_eager=True, + enable_prefix_caching=False, + enable_chunked_prefill=False, + ) as runner: + llm = runner.llm + text = "search_document: 中文 😀 café e\u0301" + params = PoolingParams( + late_chunking_params=LateChunkingParams(3), use_activation=False + ) + kwargs = {"add_special_tokens": False} + raw = llm.encode( + text, + pooling_task="token_embed", + use_tqdm=False, + pooling_params=PoolingParams(use_activation=False), + tokenization_kwargs=kwargs, + )[0] + actual = llm.encode( + text, + pooling_task="token_embed", + use_tqdm=False, + pooling_params=params, + tokenization_kwargs=kwargs, + )[0] + expected = torch.stack([part.mean(0) for part in raw.outputs.data.split(3)]) + torch.testing.assert_close(actual.outputs.data, expected, atol=1e-5, rtol=1e-4) + encoded = llm.get_tokenizer()(text, return_offsets_mapping=True, **kwargs) + assert actual.prompt_token_ids == encoded["input_ids"] + for chunk in actual.late_chunking.chunks: + start, end = chunk.token_range + offsets = [(a, b) for a, b in encoded["offset_mapping"][start:end] if a < b] + assert chunk.char_range == (offsets[0][0], offsets[-1][1]) + for rejected in ("long " * 256, " "): + with pytest.raises(VLLMValidationError): + llm.encode( + rejected, + pooling_task="token_embed", + use_tqdm=False, + pooling_params=params, + tokenization_kwargs=kwargs, + ) + assert ( + llm.encode(text, pooling_task="token_embed", use_tqdm=False)[ + 0 + ].late_chunking + is None + ) + + +@pytest.mark.skip_global_cleanup +@pytest.mark.parametrize("normalize", [False, True]) +def test_e5_late_chunking_mixed_requests(llm: LLM, normalize: bool): + texts = [prompt, "Short text.", "中文 😀 café"] + raw = llm.encode( + texts, + pooling_task="token_embed", + use_tqdm=False, + pooling_params=PoolingParams(use_activation=False), + ) + sizes = [3, None, 100] + params = [ + PoolingParams( + late_chunking_params=LateChunkingParams(size) if size is not None else None, + use_activation=normalize, + ) + for size in sizes + ] + actual = llm.encode( + texts, pooling_task="token_embed", pooling_params=params, use_tqdm=False + ) + for original, result, size in zip(raw, actual, sizes): + expected = original.outputs.data + if size is not None: + expected = torch.stack( + [part.float().mean(0) for part in expected.split(size)] + ) + assert result.late_chunking is not None + assert len(result.late_chunking.chunks) == len(expected) + assert result.late_chunking.input_tokens == len(result.prompt_token_ids) + else: + assert result.late_chunking is None + if normalize: + expected = torch.nn.functional.normalize(expected, dim=-1) + check_embeddings_close( + embeddings_0_lst=result.outputs.data.tolist(), + embeddings_1_lst=expected.tolist(), + name_0="late_chunking", + name_1="token_reference", + ) + torch.testing.assert_close( + result.outputs.data.norm(dim=-1), + expected.norm(dim=-1), + atol=1e-3, + rtol=1e-3, + ) + + +@pytest.mark.parametrize("use_v2_runner", [False, True]) +def test_late_chunking_recomputes_cached_prefix_with_chunked_prefill( + vllm_runner, monkeypatch, use_v2_runner +): + # Dummy weights suffice for testing scheduling and reduction of the same states. + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_v2_runner else "0") + with vllm_runner( + "openai-community/gpt2", + runner="pooling", + convert="embed", + load_format="dummy", + pooler_config=PoolerConfig( + task="token_embed", seq_pooling_type="LAST", tok_pooling_type="ALL" + ), + dtype="float32", + max_model_len=128, + max_num_batched_tokens=16, + max_num_seqs=2, + enable_prefix_caching=True, + enable_chunked_prefill=True, + enforce_eager=True, + gpu_memory_utilization=0.3, + seed=0, + ) as runner: + text = "The document has several important words. " * 6 + llm = runner.llm + raw = llm.encode( + text, + pooling_task="token_embed", + use_tqdm=False, + pooling_params=PoolingParams(use_activation=False), + )[0] + assert len(raw.prompt_token_ids) > 16 + cached = llm.encode( + text, + pooling_task="token_embed", + use_tqdm=False, + pooling_params=PoolingParams( + use_activation=False, skip_reading_prefix_cache=False + ), + )[0] + assert cached.num_cached_tokens > 0 + params = PoolingParams( + late_chunking_params=LateChunkingParams(7), + use_activation=False, + skip_reading_prefix_cache=False, + ) + actual = llm.encode( + text, pooling_task="token_embed", pooling_params=params, use_tqdm=False + )[0] + assert actual.num_cached_tokens == 0 + assert actual.prompt_token_ids == raw.prompt_token_ids + expected = torch.stack([part.mean(0) for part in raw.outputs.data.split(7)]) + torch.testing.assert_close(actual.outputs.data, expected, atol=1e-5, rtol=1e-4) + assert actual.late_chunking.input_tokens == len(raw.prompt_token_ids) + assert params.skip_reading_prefix_cache is False + assert params.late_chunking_params is not None diff --git a/tests/model_executor/layers/test_pooler_heads.py b/tests/model_executor/layers/test_pooler_heads.py index a80c5b8344..00091ec0c2 100644 --- a/tests/model_executor/layers/test_pooler_heads.py +++ b/tests/model_executor/layers/test_pooler_heads.py @@ -308,6 +308,7 @@ class TestTokenEmbeddingPoolerHead: x = torch.randn(5, _HIDDEN) param = PoolingParams(task="token_embed") out = head.forward_chunk(x, param) + assert out is x assert torch.equal(out, x) def test_none_chunked_prefill(self): diff --git a/tests/model_executor/layers/test_pooler_methods.py b/tests/model_executor/layers/test_pooler_methods.py index 75bbed4f37..d6b03dee19 100644 --- a/tests/model_executor/layers/test_pooler_methods.py +++ b/tests/model_executor/layers/test_pooler_methods.py @@ -13,19 +13,22 @@ from transformers import PreTrainedConfig from vllm.config import PoolerConfig, set_current_vllm_config from vllm.model_executor.layers.pooler import PoolingParamsUpdate +from vllm.model_executor.layers.pooler.activations import PoolerNormalize from vllm.model_executor.layers.pooler.seqwise.methods import ( CLSPool, LastPool, MeanPool, get_seq_pooling_method, ) -from vllm.model_executor.layers.pooler.special import DispatchPooler +from vllm.model_executor.layers.pooler.special import BOSEOSFilter, DispatchPooler +from vllm.model_executor.layers.pooler.tokwise.heads import TokenEmbeddingPoolerHead from vllm.model_executor.layers.pooler.tokwise.methods import ( AllPool, StepPool, get_tok_pooling_method, ) -from vllm.pooling_params import PoolingParams +from vllm.model_executor.layers.pooler.tokwise.poolers import TokenPooler +from vllm.pooling_params import LateChunkingParams, PoolingParams from vllm.tasks import PoolingTask from vllm.v1.pool.metadata import PoolingCursor, PoolingMetadata, PoolingStates @@ -606,3 +609,139 @@ class TestGetTokPoolingMethod: def test_unknown_raises(self): with pytest.raises(NotImplementedError, match="UNKNOWN"): get_tok_pooling_method("UNKNOWN") + + +@pytest.mark.parametrize( + ("chunk_size", "dtype", "head_dtype"), + [ + (1, torch.float32, None), + (2, torch.float16, None), + (2, torch.bfloat16, torch.float32), + (7, torch.float32, None), + (10, torch.float32, None), + (2**63 - 1, torch.float32, None), + ], +) +def test_late_chunk_pool_matches_fp32_means_before_projection_and_normalization( + chunk_size, dtype, head_dtype +): + # Unequal norms expose the incorrect normalize-then-average implementation. + hidden = (torch.arange(28).reshape(7, 4) - 10).to(dtype) + projector = torch.nn.Linear(4, 3, bias=True).to(head_dtype or dtype) + head = TokenEmbeddingPoolerHead(head_dtype, projector, PoolerNormalize()) + pooler = TokenPooler(TestAllPool._make_all_pool(), head) + params = PoolingParams( + task="token_embed", + late_chunking_params=LateChunkingParams(chunk_size=chunk_size), + use_activation=True, + ) + metadata = _make_metadata([7], pooling_params=[params]) + actual = pooler(hidden, metadata)[0] + means = torch.stack( + [hidden[a : a + chunk_size].float().mean(0) for a in range(0, 7, chunk_size)] + ).to(head_dtype or dtype) + expected = torch.nn.functional.normalize(projector(means), p=2, dim=-1) + torch.testing.assert_close(actual, expected) + assert actual.dtype == (head_dtype or dtype) + + +def test_late_chunk_pool_reduces_before_nonlinear_projector(): + hidden = torch.tensor([[2.0, 0.0], [0.0, 1.0]]) + projector = torch.nn.Tanh() + head = TokenEmbeddingPoolerHead(projector=projector, activation=PoolerNormalize()) + params = PoolingParams( + task="token_embed", + late_chunking_params=LateChunkingParams(chunk_size=2), + use_activation=True, + ) + actual = head.forward_chunk(hidden, params) + expected = torch.nn.functional.normalize(projector(hidden.mean(0, keepdim=True))) + wrong_order = torch.nn.functional.normalize(projector(hidden).mean(0, keepdim=True)) + torch.testing.assert_close(actual, expected) + assert not torch.allclose(actual, wrong_order) + + +def test_late_chunk_pool_isolates_mixed_requests_and_owns_finished_storage(): + hidden = torch.arange(36, dtype=torch.float32).reshape(9, 4) + params = [ + PoolingParams( + task="token_embed", late_chunking_params=LateChunkingParams(chunk_size=2) + ), + PoolingParams(task="token_embed"), + PoolingParams( + task="token_embed", late_chunking_params=LateChunkingParams(chunk_size=1) + ), + ] + pooler = TokenPooler( + TestAllPool._make_all_pool(async_scheduling=True), TokenEmbeddingPoolerHead() + ) + outputs = pooler(hidden, _make_metadata([3, 4, 2], pooling_params=params)) + expected = [ + torch.stack([hidden[:2].mean(0), hidden[2]]), + hidden[3:7].clone(), + hidden[7:].clone(), + ] + hidden.fill_(-999) + for actual, reference in zip(outputs, expected): + torch.testing.assert_close(actual, reference) + + +@pytest.mark.parametrize("chunk_size", [1, 2, 8]) +@pytest.mark.parametrize("embedding_size", [1, 4]) +def test_bos_eos_filter_preserves_chunk_vectors_in_mixed_batch( + chunk_size, embedding_size +): + hidden = torch.arange(10 * embedding_size, dtype=torch.float32).reshape( + 10, embedding_size + ) + pooler = BOSEOSFilter( + TokenPooler(TestAllPool._make_all_pool(), TokenEmbeddingPoolerHead()), + bos_token_id=101, + eos_token_id=102, + ) + metadata = _make_metadata( + [5, 5], + token_ids=[[101, 1, 2, 3, 102]] * 2, + pooling_params=[ + PoolingParams( + task="token_embed", late_chunking_params=LateChunkingParams(chunk_size) + ), + PoolingParams(task="token_embed"), + ], + ) + chunked, ordinary = pooler(hidden, metadata) + expected = torch.stack([part.mean(0) for part in hidden[:5].split(chunk_size)]) + torch.testing.assert_close(chunked, expected) + torch.testing.assert_close(ordinary, hidden[6:9].squeeze(-1)) + + +def test_bos_eos_filter_waits_for_complete_chunk_vectors(): + pooler = BOSEOSFilter( + TokenPooler( + TestAllPool._make_all_pool(chunked=True), TokenEmbeddingPoolerHead() + ), + bos_token_id=101, + ) + params = [ + PoolingParams(task="token_embed", late_chunking_params=LateChunkingParams(3)) + ] + first = _make_metadata( + [4], + token_ids=[[101, 1, 2, 102]], + pooling_params=params, + num_scheduled_tokens=[2], + seq_lens=[2], + ) + hidden = torch.arange(8, dtype=torch.float32).reshape(4, 2) + assert pooler(hidden[:2], first) == [None] + last = _make_metadata( + [4], + token_ids=[[101, 1, 2, 102]], + pooling_params=params, + num_scheduled_tokens=[2], + seq_lens=[4], + ) + last.pooling_states = first.pooling_states + actual = pooler(hidden[2:], last)[0] + torch.testing.assert_close(actual, torch.stack([hidden[:3].mean(0), hidden[3]])) + assert not last.pooling_states[0].hidden_states_cache diff --git a/tests/test_pooling_params.py b/tests/test_pooling_params.py index f34270d0da..81af02b502 100644 --- a/tests/test_pooling_params.py +++ b/tests/test_pooling_params.py @@ -1,8 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from dataclasses import dataclass +from types import SimpleNamespace from typing import Any +import msgspec import pytest from pydantic import TypeAdapter, ValidationError @@ -13,6 +15,7 @@ from vllm.entrypoints.pooling.classify.protocol import ClassificationRequest from vllm.entrypoints.pooling.embed.protocol import EmbeddingRequest from vllm.entrypoints.pooling.pooling.protocol import PoolingRequest from vllm.exceptions import VLLMValidationError +from vllm.pooling_params import LateChunkingParams EMBEDDING_MODELS = [ EmbedModelInfo("intfloat/multilingual-e5-small", is_matryoshka=False), @@ -211,3 +214,92 @@ def test_token_classify(pooling_type: str): with pytest.raises(VLLMValidationError): pooling_params = PoolingParams(task=task, **{p: True}) pooling_params.verify(model_config) + + +@pytest.mark.parametrize("value", [0, -1, True, False, 1.5, "2"]) +def test_late_chunk_size_rejects_non_positive_integers(value): + with pytest.raises(VLLMValidationError, match="positive integer"): + PoolingParams(late_chunking_params=LateChunkingParams(chunk_size=value)) + + +def _late_chunking_model_config(): + return SimpleNamespace( + architecture="NomicBertModel", + model_impl="auto", + is_matryoshka=False, + hf_config=SimpleNamespace(), + pooler_config=PoolerConfig(seq_pooling_type="MEAN", tok_pooling_type="ALL"), + ) + + +def test_late_chunking_params_preserve_defaults_clone_and_wire_format(): + model_config = _late_chunking_model_config() + params = PoolingParams( + task="token_embed", late_chunking_params=LateChunkingParams(chunk_size=3) + ) + params.verify(model_config) + assert params.skip_reading_prefix_cache is True + assert params.use_activation is True + # Chunk means must include cached prefixes even when the user opts into reads. + explicit_cache_flag = PoolingParams( + task="token_embed", + late_chunking_params=LateChunkingParams(chunk_size=3), + skip_reading_prefix_cache=False, + ) + explicit_cache_flag.verify(model_config) + assert explicit_cache_flag.skip_reading_prefix_cache is True + clone = params.clone() + assert clone.late_chunking_params is not None + clone.late_chunking_params.chunk_size = 7 + assert params.late_chunking_params is not None + assert params.late_chunking_params.chunk_size == 3 + # Workers receive only the fixed chunk size, never source ranges. + assert msgspec.msgpack.decode( + msgspec.msgpack.encode(params.late_chunking_params) + ) == [3] + assert ( + msgspec.msgpack.decode(msgspec.msgpack.encode(params), type=PoolingParams) + == params + ) + # A message written before the appended field still uses the old default. + wire = msgspec.msgpack.decode(msgspec.msgpack.encode(PoolingParams())) + assert msgspec.convert(wire[:-1], type=PoolingParams).late_chunking_params is None + + +@pytest.mark.parametrize( + "task", ["embed", "classify", "token_classify", "plugin", None] +) +def test_late_chunking_rejects_other_tasks(task): + with pytest.raises(VLLMValidationError, match="requires token_embed"): + PoolingParams( + task=task, late_chunking_params=LateChunkingParams(chunk_size=2) + ).verify(_late_chunking_model_config()) + + +def test_late_chunking_uses_existing_dimension_validation(): + config = MockMatryoshkaModelConfig( + pooler_config=PoolerConfig(tok_pooling_type="ALL", use_activation=False), + ) + params = PoolingParams( + task="token_embed", late_chunking_params=LateChunkingParams(2), dimensions=4 + ) + params.verify(config) + assert params.dimensions == 4 + assert params.use_activation is False + params.dimensions = config.embedding_size + 1 + with pytest.raises(VLLMValidationError, match="dimensions in range"): + params.verify(config) + + +@pytest.mark.parametrize("from_config", [False, True]) +def test_late_chunking_rejects_token_filtering_after_resolving_defaults(from_config): + config = MockModelConfig(pooler_config=PoolerConfig(tok_pooling_type="STEP")) + params = PoolingParams( + task="token_embed", late_chunking_params=LateChunkingParams(2) + ) + if from_config: + config.pooler_config.step_tag_id = 42 + else: + params.step_tag_id = 42 + with pytest.raises(VLLMValidationError, match="all token rows"): + params.verify(config) diff --git a/vllm/entrypoints/pooling/embed/io_processor.py b/vllm/entrypoints/pooling/embed/io_processor.py index 94f84f09b6..b37a009c92 100644 --- a/vllm/entrypoints/pooling/embed/io_processor.py +++ b/vllm/entrypoints/pooling/embed/io_processor.py @@ -12,6 +12,7 @@ from openai.types.chat import ( from openai.types.chat.chat_completion_content_part_image_param import ImageURL from vllm import PoolingParams +from vllm.exceptions import VLLMValidationError from vllm.inputs import tokens_input from vllm.logger import init_logger from vllm.outputs import PoolingOutput, PoolingRequestOutput @@ -26,18 +27,26 @@ from vllm.utils.collection_utils import chunk_list from vllm.utils.mistral import is_mistral_tokenizer from ..base.io_processor import PoolingIOProcessor +from ..late_chunking import ( + attach_late_chunking_metadata, + build_late_chunking_metadata, + prepare_late_chunking_input, +) from ..scoring.io_processor import JinaRankingIOProcessorMixin from ..typing import ( AnyOfflineInputsContext, AnyRenderParam, ChunkedEmbeddingMetadata, EncodeChatRenderParams, + LateChunkingSlot, OfflineEncodeInputsContext, + OfflineOutputsContext, PoolingChatLikeRequest, PoolingCompletionLikeRequest, PoolingEngineInput, PoolingServeContext, RequestFactory, + RequestGenerator, ) from .protocol import ( CohereEmbedContent, @@ -657,10 +666,75 @@ class EmbedIOProcessor(PoolingIOProcessor): class TokenEmbedIOProcessor(PoolingIOProcessor): name = "token_embed" + def get_request_factory_offline( + self, ctx: AnyOfflineInputsContext + ) -> tuple[RequestFactory, int]: + factory, num_requests = super().get_request_factory_offline(ctx) + assert isinstance(ctx, OfflineEncodeInputsContext) + if ctx.pooling_params is None or not any( + p.late_chunking_params is not None + for p in self._params_to_seq(ctx.pooling_params, num_requests) + ): + return factory, num_requests + + if (ctx.tokenization_kwargs or {}).get("padding") not in ( + None, + False, + "do_not_pad", + ): + raise VLLMValidationError("Late chunking does not support input padding") + + def request_factory() -> RequestGenerator: + for request in factory(): + slot = None + if request["params"].late_chunking_params is not None: + slot = LateChunkingSlot() + request["late_chunking"] = slot + ctx.late_chunking.append(slot) + yield request + + return request_factory, num_requests + + def render(self, render_params: AnyRenderParam) -> PoolingEngineInput: + late_chunking = render_params["params"].late_chunking_params + if late_chunking is None: + return super().render(render_params) + + text, tok_params = prepare_late_chunking_input(self.vllm_config, render_params) + render_params = render_params.copy() + render_params["tok_params"] = tok_params + result = super().render(render_params) + prompt = result["prompts"] + if prompt["type"] != "token": + raise VLLMValidationError("Late chunking requires tokenized text") + render_params["late_chunking"].metadata = build_late_chunking_metadata( + text, + len(prompt["prompt_token_ids"]), + prompt.pop("prompt_token_offsets", None), + late_chunking.chunk_size, + ) + return result + + def post_process_offline(self, ctx: OfflineOutputsContext): + if ctx.late_chunking: + for output, slot in zip(ctx.outputs, ctx.late_chunking, strict=True): + if slot is not None: + assert slot.metadata is not None + attach_late_chunking_metadata(output, slot.metadata) + return ctx.outputs + class JinaRankingTokenEmbedIOProcessor( TokenEmbedIOProcessor, JinaRankingIOProcessorMixin ): + def render(self, render_params: AnyRenderParam) -> PoolingEngineInput: + if render_params["params"].late_chunking_params is not None: + raise VLLMValidationError( + "JinaForRanking returns document/query vectors and does not support " + "late chunking" + ) + return super().render(render_params) + def get_request_factory_online( self, ctx: PoolingServeContext ) -> Sequence[AnyRenderParam]: diff --git a/vllm/entrypoints/pooling/late_chunking.py b/vllm/entrypoints/pooling/late_chunking.py new file mode 100644 index 0000000000..d848c3257f --- /dev/null +++ b/vllm/entrypoints/pooling/late_chunking.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Sequence +from dataclasses import replace +from typing import cast + +from vllm.config import VllmConfig +from vllm.exceptions import VLLMValidationError +from vllm.outputs import LateChunk, LateChunkingMetadata, PoolingRequestOutput +from vllm.renderers import TokenizeParams + +from .typing import AnyRenderParam, EncodeCMPLRenderParams + + +def prepare_late_chunking_input( + config: VllmConfig, render_params: AnyRenderParam +) -> tuple[str, TokenizeParams]: + """Validate full-document text input before requesting tokenizer offsets.""" + params = render_params["params"] + params.verify(config.model_config) + if "prompts" not in render_params: + raise VLLMValidationError("Late chunking requires plain-text input") + prompt = cast(EncodeCMPLRenderParams, render_params)["prompts"] + if ( + not isinstance(prompt, dict) + or set(prompt) - {"prompt", "cache_salt"} + or not isinstance(text := prompt.get("prompt"), str) + or not text + ): + raise VLLMValidationError("Late chunking requires nonempty plain-text input") + tok_params = render_params["tok_params"] + if ( + tok_params.pad_prompt_tokens is not None + or tok_params.truncate_prompt_tokens is not None + or tok_params.do_lower_case + ): + raise VLLMValidationError( + "Late chunking does not support input padding, truncation or " + "renderer text normalization" + ) + return text, replace(tok_params, return_token_offsets=True) + + +def build_late_chunking_metadata( + text: str, + num_tokens: int, + offsets: Sequence[tuple[int, int]] | None, + chunk_size: int, +) -> LateChunkingMetadata: + """Map chunks to original text, keeping special-token-only chunks.""" + if num_tokens <= 0 or offsets is None or len(offsets) != num_tokens: + raise VLLMValidationError( + "Late chunking requires a fast tokenizer with offsets aligned to " + "all input tokens" + ) + previous_start = previous_end = 0 + for start, end in offsets: + if not 0 <= start <= end <= len(text): + raise VLLMValidationError("Token offsets are not aligned to the input text") + if start < end: + if start < previous_start or end < previous_end: + raise VLLMValidationError("Token offsets are not in source-text order") + previous_start, previous_end = start, end + + chunks = [] + for start in range(0, num_tokens, chunk_size): + end = min(start + chunk_size, num_tokens) + source = [(a, b) for a, b in offsets[start:end] if a < b] + char_range = ( + (min(a for a, _ in source), max(b for _, b in source)) if source else None + ) + chunks.append(LateChunk(token_range=(start, end), char_range=char_range)) + return LateChunkingMetadata( + chunk_size=chunk_size, input_tokens=num_tokens, chunks=chunks + ) + + +def attach_late_chunking_metadata( + output: PoolingRequestOutput, metadata: LateChunkingMetadata +) -> None: + """Attach ranges only to successful, complete chunk-vector results.""" + if not output.finished or output.error is not None: + return + data = output.outputs.data + if ( + data.ndim != 2 + or data.shape[0] != len(metadata.chunks) + or len(output.prompt_token_ids) != metadata.input_tokens + ): + raise ValueError("Late-chunking output does not match its token/source ranges") + output.late_chunking = metadata diff --git a/vllm/entrypoints/pooling/offline.py b/vllm/entrypoints/pooling/offline.py index 19be1eddde..edf92791a8 100644 --- a/vllm/entrypoints/pooling/offline.py +++ b/vllm/entrypoints/pooling/offline.py @@ -129,7 +129,7 @@ class PoolingOfflineMixin(OfflineInferenceMixin): io_processor, request_factory, num_requests, use_tqdm=use_tqdm ) outputs = io_processor.post_process_offline( - ctx=OfflineOutputsContext(outputs=outputs) + ctx=OfflineOutputsContext(outputs=outputs, late_chunking=ctx.late_chunking) ) return outputs diff --git a/vllm/entrypoints/pooling/typing.py b/vllm/entrypoints/pooling/typing.py index a299fc657d..33aa659989 100644 --- a/vllm/entrypoints/pooling/typing.py +++ b/vllm/entrypoints/pooling/typing.py @@ -3,7 +3,7 @@ import time from collections.abc import AsyncGenerator, Callable, Generator, Sequence from dataclasses import dataclass, field -from typing import Any, Generic, TypeAlias, TypedDict, TypeVar +from typing import Any, Generic, NotRequired, TypeAlias, TypedDict, TypeVar from fastapi import Request from pydantic import ConfigDict @@ -11,6 +11,7 @@ from pydantic import ConfigDict from vllm import PoolingParams, PoolingRequestOutput, PromptType from vllm.inputs import DataPrompt, EngineInput from vllm.lora.request import LoRARequest +from vllm.outputs import LateChunkingMetadata from vllm.renderers import ChatParams, TokenizeParams from vllm.renderers.chat_utils import ChatCompletionMessageParam from vllm.renderers.inputs import DictPrompt @@ -80,6 +81,13 @@ class ChunkedEmbeddingMetadata: chunk_index: int +@dataclass +class LateChunkingSlot: + """Per-document source ranges shared by frontend rendering and output contexts.""" + + metadata: LateChunkingMetadata | None = None + + @dataclass(kw_only=True) class PoolingServeContext(Generic[PoolingRequestT]): model_config = ConfigDict(arbitrary_types_allowed=True) @@ -124,6 +132,9 @@ class OfflineInputsContext: tokenization_kwargs: dict[str, Any] | None lora_request: Sequence[LoRARequest | None] | None priorities: int | Sequence[int] | None + late_chunking: list[LateChunkingSlot | None] = field( + default_factory=list, kw_only=True + ) @dataclass @@ -158,6 +169,7 @@ class OfflineOutputsContext: ## for bi-encoder & late-interaction n_queries: int | None = None + late_chunking: list[LateChunkingSlot | None] = field(default_factory=list) class RenderParams(TypedDict): @@ -168,6 +180,7 @@ class RenderParams(TypedDict): params: PoolingParams lora_requests: LoRARequest | None priorities: int + late_chunking: NotRequired[LateChunkingSlot] class EncodeCMPLRenderParams(RenderParams): diff --git a/vllm/model_executor/layers/pooler/special.py b/vllm/model_executor/layers/pooler/special.py index 23354d66f3..f464f28300 100644 --- a/vllm/model_executor/layers/pooler/special.py +++ b/vllm/model_executor/layers/pooler/special.py @@ -198,6 +198,11 @@ class BOSEOSFilter(Pooler): for i, (prompt_len, token_ids) in enumerate( zip(pooling_metadata.prompt_lens, prompt_token_ids) ): + # Late chunking returns one row per chunk, with special tokens + # already included in the means. Skip BOS/EOS row removal for these + # requests; ordinary token outputs still use the filtering below. + if pooling_metadata.pooling_params[i].late_chunking_params is not None: + continue pooled_data = pooled_outputs[i] assert ( isinstance(pooled_data, torch.Tensor) diff --git a/vllm/model_executor/layers/pooler/tokwise/heads.py b/vllm/model_executor/layers/pooler/tokwise/heads.py index 78ec82a538..95242ee470 100644 --- a/vllm/model_executor/layers/pooler/tokwise/heads.py +++ b/vllm/model_executor/layers/pooler/tokwise/heads.py @@ -17,6 +17,23 @@ from .methods import TokenPoolingMethodOutputItem TokenPoolerHeadOutputItem: TypeAlias = torch.Tensor | None +def _mean_pool_chunks(data: torch.Tensor, chunk_size: int) -> torch.Tensor: + """Mean-pool complete contextual states, accumulating in FP32.""" + chunk_size = min(chunk_size, len(data)) + num_full, remainder = divmod(len(data), chunk_size) + full_end = num_full * chunk_size + full_chunks = data[:full_end].reshape(num_full, chunk_size, data.shape[-1]) + if not remainder: + return full_chunks.mean(dim=1, dtype=torch.float32) + + return torch.cat( + [ + full_chunks.mean(dim=1, dtype=torch.float32), + data[full_end:].mean(dim=0, keepdim=True, dtype=torch.float32), + ] + ) + + class TokenPoolerHead(nn.Module, ABC): @abstractmethod def get_supported_tasks(self) -> Set[PoolingTask]: @@ -80,6 +97,12 @@ class TokenEmbeddingPoolerHead(TokenPoolerHead): if pooled_data is None: return None + if (late_chunking := pooling_param.late_chunking_params) is not None: + means = _mean_pool_chunks(pooled_data, late_chunking.chunk_size) + pooled_data = ( + means.to(pooled_data.dtype) if self.head_dtype is None else means + ) + if self.head_dtype is not None: pooled_data = pooled_data.to(self.head_dtype) # pooled_data shape: [n_tokens, hidden_size] diff --git a/vllm/outputs.py b/vllm/outputs.py index 145b7d510f..887ebc5933 100644 --- a/vllm/outputs.py +++ b/vllm/outputs.py @@ -277,6 +277,20 @@ STREAM_FINISHED = RequestOutput( _O = TypeVar("_O", default=PoolingOutput) +@dataclass(frozen=True) +class LateChunk: + token_range: tuple[int, int] + char_range: tuple[int, int] | None + + +@dataclass(frozen=True) +class LateChunkingMetadata: + chunk_size: int + input_tokens: int + chunks: list[LateChunk] + strategy: str = "fixed_tokens" + + class PoolingRequestOutput(Generic[_O]): """The output data of a pooling request to the LLM. @@ -287,6 +301,7 @@ class PoolingRequestOutput(Generic[_O]): num_cached_tokens: The number of tokens with prefix cache hit. finished (bool): A flag indicating whether the pooling is completed. error: Structured request-level error information, if the request failed. + late_chunking: Token and source ranges for fixed-length late chunking. """ @@ -299,8 +314,10 @@ class PoolingRequestOutput(Generic[_O]): finished: bool, *, error: RequestError | None = None, + late_chunking: LateChunkingMetadata | None = None, ): self.request_id = request_id + self.late_chunking = late_chunking self.prompt_token_ids = prompt_token_ids self.num_cached_tokens = num_cached_tokens self.finished = finished diff --git a/vllm/pooling_params.py b/vllm/pooling_params.py index 319b77d598..45c34afbe3 100644 --- a/vllm/pooling_params.py +++ b/vllm/pooling_params.py @@ -36,6 +36,28 @@ class LateInteractionParams( query_uses: int | None = None +class LateChunkingParams( + msgspec.Struct, + omit_defaults=True, # type: ignore[call-arg] + array_like=True, +): # type: ignore[call-arg] + """Fixed-length mean pooling before the token embedding head. + + Requests recompute all prompt tokens, bypassing prefix cache reads. + Special tokens count toward chunk boundaries and are included in chunk means, + including the BOS token normally filtered from BGE-M3 token embeddings. + """ + + chunk_size: int + + def __post_init__(self) -> None: + self.verify() + + def verify(self) -> None: + if type(self.chunk_size) is not int or self.chunk_size <= 0: + raise VLLMValidationError("chunk_size must be a positive integer") + + class PoolingParams( msgspec.Struct, omit_defaults=True, # type: ignore[call-arg] @@ -48,6 +70,8 @@ class PoolingParams( `None` uses the pooler's default, which is `True` in most cases. dimensions: Reduce the dimensions of embeddings if model support matryoshka representation. + late_chunking_params: Mean-pool contextual token states into fixed-size + chunks for the token embedding task. """ @@ -72,16 +96,22 @@ class PoolingParams( extra_kwargs: dict[str, Any] | None = None output_kind: RequestOutputKind = RequestOutputKind.FINAL_ONLY + ## for token embedding models + # --8<-- [start:token-embed-pooling-params] + late_chunking_params: LateChunkingParams | None = None + """Mean-pool contextual tokens into fixed-size chunks for `token_embed`.""" + # --8<-- [end:token-embed-pooling-params] + @property def all_parameters(self) -> list[str]: - return ["dimensions", "use_activation"] + return ["dimensions", "use_activation", "late_chunking_params"] @property def valid_parameters(self): return { "embed": ["dimensions", "use_activation"], "classify": ["use_activation"], - "token_embed": ["dimensions", "use_activation"], + "token_embed": ["dimensions", "use_activation", "late_chunking_params"], "token_classify": ["use_activation"], } @@ -90,6 +120,15 @@ class PoolingParams( return deepcopy(self) def verify(self, model_config: ModelConfig) -> None: + if self.late_chunking_params is not None: + self.late_chunking_params.verify() + if self.task != "token_embed": + raise VLLMValidationError("Late chunking requires token_embed") + if self.late_interaction_params is not None: + raise VLLMValidationError( + "Late chunking cannot be combined with late-interaction scoring" + ) + # plugin task uses io_processor.parse_data to verify inputs, # skipping PoolingParams verify if self.task == "plugin": @@ -107,6 +146,11 @@ class PoolingParams( self._merge_default_parameters(model_config) self._set_default_parameters(model_config) self._verify_valid_parameters() + if self.late_chunking_params is not None: + if self.step_tag_id is not None: + raise VLLMValidationError("Late chunking requires all token rows") + # KV cache hits omit hidden states needed for complete chunk means. + self.skip_reading_prefix_cache = True def _merge_default_parameters(self, model_config: ModelConfig) -> None: pooler_config = model_config.pooler_config @@ -228,6 +272,7 @@ class PoolingParams( f"requires_token_ids={self.requires_token_ids}, " f"skip_reading_prefix_cache={self.skip_reading_prefix_cache}, " f"late_interaction_params={self.late_interaction_params}, " + f"late_chunking_params={self.late_chunking_params}, " f"extra_kwargs={self.extra_kwargs})" )