[Feature][Pooling] Add offline fixed-length late chunking (#60164)

Signed-off-by: Linze-Shi <linzeshi0@gmail.com>

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.

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 = "<special tokens only>"
+            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())

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

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

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):

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

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)

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]:

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

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
 

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):

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)

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]

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

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})"
         )