[Feature][Pooling] Serve late chunking through /v1/embeddings (#61772)

Signed-off-by: Linze-Shi <linzeshi0@gmail.com>
Co-authored-by: vllm-agent <claw@inferact.ai>
Co-authored-by: wang.yuqi <yuqi.wang@daocloud.io>

docs/models/pooling_models/README.md

@@ -130,7 +130,7 @@ enabling the corresponding APIs.
 | Classification Usages       | `LLM.classify(...)` | `classify`                        | `cross-encoder` (see note) | linear classifier        |
 | Embedding Usages            | `LLM.embed(...)`    | `embed`                           | `bi-encoder`               | cosine similarity        |
 | Token Classification Usages | N/A                 | `token_classify`                  | N/A                        | N/A                      |
-| Token Embedding Usages      | N/A                 | `token_embed`                     | `late-interaction`         | late interaction(MaxSim) |
+| Token Embedding Usages      | `LLM.embed(...)`    | `token_embed`                     | `late-interaction`         | late interaction(MaxSim) |
 | Reward Usages               | N/A                 | `classify` & `token_classify`     | N/A                        | N/A                      |
 | Scoring Usages              | `LLM.score(...)`    | N/A                               | N/A                        | N/A                      |
 | Plugins Usages              | N/A                 | `plugin`                          | N/A                        | N/A                      |
@@ -146,9 +146,9 @@ For more information about `LLM.classify`, see [this page](classify.md#offline-i
 
 ### `LLM.embed`
 
-The [embed][vllm.LLM.embed] method outputs an embedding vector for each prompt.
+The [embed][vllm.LLM.embed] method outputs one embedding vector per prompt with the `embed` task, or a matrix of token vectors with the `token_embed` task. Eligible token embedding models also support [late chunking](token_embed.md#fixed-length-late-chunking) through this method.
 It is primarily designed for [embedding models](embed.md).
-For more information about `LLM.embed`, see [this page](embed.md#offline-inference).
+For more information about `LLM.embed`, see the [sequence embedding](embed.md#offline-inference) and [token embedding](token_embed.md#offline-inference) guides.
 
 ### `LLM.score`
 
@@ -178,8 +178,9 @@ print(f"Data: {data!r}")
 Our online Server provides endpoints that correspond to the offline APIs:
 
 - Corresponding to `LLM.embed`:
-    - [Cohere Embed API](embed.md#cohere-embed-api) (`/v2/embed`)
+    - [Cohere Embed API](embed.md#cohere-embed-api) (`/v2/embed`) for the `embed` task
     - [OpenAI-compatible Embeddings API](embed.md#openai-compatible-embeddings-api) (`/v1/embeddings`)
+    - [Token Embeddings API](token_embed.md#online-serving) (`/v1/embeddings`) for the `token_embed` task, returning token or chunk vectors as a vLLM extension
 - Corresponding to `LLM.classify`:
     - [Classification API](classify.md#online-serving)(`/classify`)
 - Corresponding to `LLM.score`:

docs/models/pooling_models/token_embed.md

@@ -5,9 +5,11 @@
 - Model Usage: Token classification models
 - Pooling Tasks: `token_embed`
 - Offline APIs:
+    - `LLM.embed(...)` with the `token_embed` pooling task
     - `LLM.encode(..., pooling_task="token_embed")`
 - Online APIs:
     - Pooling API (`/pooling`)
+    - Embedding API (`/v1/embeddings`), with optional [late chunking](#online-late-chunking)
 
 The difference between the (sequence) embedding task and the token embedding task is that (sequence) embedding outputs one embedding for each sequence, while token embedding outputs an embedding for each token.
 
@@ -98,7 +100,7 @@ The following [pooling parameters][vllm.PoolingParams] are supported.
 
 The [encode][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.encode] method is available to all pooling models in vLLM.
 
-Set `pooling_task="token_embed"` when using `LLM.encode` for token embedding Models:
+Set `pooling_task="token_embed"` when using `LLM.encode` for token embedding models:
 
 ```python
 from vllm import LLM
@@ -110,19 +112,32 @@ data = output.outputs.data
 print(f"Data: {data!r}")
 ```
 
+### `LLM.embed`
+
+The [embed][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.embed] method uses the model's configured pooling task. For `token_embed`, it returns a `PoolingRequestOutput` per input, with a token-vector matrix in `outputs.data`. Ordinary `embed` models keep returning `EmbeddingRequestOutput` with one vector in `outputs.embedding`.
+
+```python
+from vllm import LLM
+
+llm = LLM(model="answerdotai/answerai-colbert-small-v1", runner="pooling")
+(output,) = llm.embed("Hello, my name is")
+
+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:
+With an instance configured for `token_embed`, enable late chunking per request through `LLM.embed` for plain-text input:
 
 ```python
 from vllm import PoolingParams
 from vllm.pooling_params import LateChunkingParams
 
-outputs = llm.encode(
+outputs = llm.embed(
     [document],
-    pooling_task="token_embed",
     pooling_params=PoolingParams(
         late_chunking_params=LateChunkingParams(chunk_size=256)
     ),
@@ -131,6 +146,8 @@ vectors = outputs[0].outputs.data  # [num_chunks, embedding_dimension]
 chunks = outputs[0].late_chunking.chunks
 ```
 
+To use `LLM.encode` instead, pass the same pooling parameters and set `pooling_task="token_embed"`.
+
 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`.
@@ -164,7 +181,72 @@ print(f"Score: {score}")
 
 ## Online Serving
 
-Please refer to the [Pooling API](README.md#pooling-api) and use `"task":"token_embed"`.
+### Embedding API
+
+On a server configured for `token_embed`, `/v1/embeddings` returns token vectors by default and chunk vectors when `late_chunking_params` is provided.
+
+#### Online late chunking
+
+Start a token embedding server, for example:
+
+```bash
+vllm serve intfloat/multilingual-e5-small \
+    --runner pooling --pooler-config.task token_embed --max-model-len 512
+```
+
+Set `late_chunking_params.chunk_size` to a positive integer less than `2**63`:
+
+```bash
+curl http://localhost:8000/v1/embeddings \
+    -H 'Content-Type: application/json' \
+    -d '{
+      "model": "intfloat/multilingual-e5-small",
+      "input": "passage: Berlin has many museums. They display art from around the world.",
+      "late_chunking_params": {"chunk_size": 16},
+      "encoding_format": "float"
+    }'
+```
+
+Each response item contains `embedding`, a `[num_chunks, embedding_dimension]` matrix, and `late_chunking`, with these fields:
+
+| Field | Meaning |
+| --- | --- |
+| `strategy` | `"fixed_tokens"` |
+| `chunk_size` | Requested number of model-input tokens per chunk |
+| `input_tokens` | Number of tokens in this document, including special tokens |
+| `chunks` | One entry per vector row, containing `token_range` and `char_range` |
+
+Ranges have the same semantics as the offline API. JSON represents them as two-element arrays; `char_range` is `null` for a special-token-only chunk. Character indices count Unicode code points, so clients using UTF-16 string indices must convert them before slicing.
+
+Use a string or a nonempty list of strings as `input`. A batch shares one chunk size, returns one item per document in input order, and counts input tokens once in `usage`. Values of `chunk_size` larger than the input token count produce one chunk.
+
+Omit `late_chunking_params` or set it to `null` to return ordinary token vectors through the same endpoint. Each item then contains a token-vector matrix in `embedding` and omits `late_chunking`. This mode also accepts token IDs. With late chunking enabled, input must be plain text and must not request truncation or padding, because source ranges refer to the complete submitted text. Only `encoding_format="float"` is supported; chat messages and `dimensions` are unsupported by this endpoint.
+
+The matrix-valued `embedding` and `late_chunking` fields are a vLLM extension, not the standard OpenAI embedding response. Use an HTTP client that handles this format. Ordinary `embed` servers keep their existing response format; they do not enable late chunking and may ignore an unrecognized `late_chunking_params` field. Cohere `/v2/embed`, the batch API, and `/invocations` are not extended with late chunking.
+
+#### Retrieve source text from chunk vectors
+
+The [retrieval example](../../../examples/pooling/token_embed/late_chunking_retrieval_online.py) uses the server above to:
+
+1. Encode documents with E5's `passage:` prefix and request contextual chunks.
+2. Build an in-memory index containing vectors, document IDs, chunk indices and source ranges.
+3. Encode a question with the `query:` prefix, using one chunk covering the whole query.
+4. Rank stored vectors by cosine similarity and print the matching source excerpts.
+
+```bash
+python examples/pooling/token_embed/late_chunking_retrieval_online.py \
+    --chunk-size 16 --query "Which city has museums with art collections?"
+```
+
+Add `--index-path late_chunking_index.json` to save the index. The example creates no index file by default.
+
+For E5, one chunk covering the complete query produces its mean-pooled sequence embedding. Other models may require a different query head. Keep the model revision, prefixes, projection and normalization consistent between indexing and querying.
+
+Source ranges refer to the exact submitted text, including the prefix. The example excludes special-token-only chunks from retrieval. Fixed token boundaries may split sentences; late chunking does not guarantee better retrieval on every dataset.
+
+### Pooling API
+
+Use the [Pooling API](README.md#pooling-api) with `"task": "token_embed"` to return token vectors. The Pooling API does not support late chunking.
 
 ## More examples
 

examples/pooling/token_embed/late_chunking_offline.py

@@ -43,9 +43,8 @@ def main(args: Namespace):
         "passage: Kyoto is a city in Japan. Its historic temples attract visitors "
         "throughout the year.",
     ]
-    outputs = llm.encode(
+    outputs = llm.embed(
         documents,
-        pooling_task="token_embed",
         pooling_params=pooling_params,
     )
 

examples/pooling/token_embed/late_chunking_retrieval_online.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Index contextual chunks and retrieve source excerpts with multilingual E5.
+
+Start the server with:
+    vllm serve intfloat/multilingual-e5-small --runner pooling \
+        --pooler-config.task token_embed --max-model-len 512
+
+Then run this file to search an in-memory index. Use --index-path to save it.
+"""
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import requests
+
+
+def encode(session, args, texts, chunk_size):
+    response = session.post(
+        f"{args.url.rstrip('/')}/v1/embeddings",
+        json={
+            "model": args.model,
+            "input": texts,
+            "encoding_format": "float",
+            "use_activation": True,
+            "late_chunking_params": {"chunk_size": chunk_size},
+        },
+        timeout=120,
+    )
+    response.raise_for_status()
+    return response.json()["data"]
+
+
+def main(args):
+    # Ranges refer to these exact submitted strings, including the E5 prefix.
+    documents = {
+        "berlin": "passage: Berlin is the capital of Germany. The city has many "
+        "museums. They display art and historical collections from around the world.",
+        "kyoto": "passage: Kyoto is a city in Japan. Its historic temples attract "
+        "visitors throughout the year.",
+        "storage": "passage: Database backups protect data from accidental loss. "
+        "They should be tested regularly to ensure that recovery works.",
+    }
+    chunks = []
+    with requests.Session() as session:
+        outputs = encode(session, args, list(documents.values()), args.chunk_size)
+        for doc_id, output in zip(documents, outputs):
+            for index, (vector, chunk) in enumerate(
+                zip(output["embedding"], output["late_chunking"]["chunks"])
+            ):
+                if chunk["char_range"] is None:
+                    continue  # No source text to retrieve for special-only chunks.
+                chunks.append(
+                    {
+                        "document_id": doc_id,
+                        "chunk_index": index,
+                        "vector": vector,
+                        **chunk,
+                    }
+                )
+
+        # E5 uses mean pooling. One chunk covering the query gives its sequence
+        # embedding, with the same model, head and normalization as the documents.
+        # The server rejects queries longer than its 512-token context window.
+        query = encode(session, args, [f"query: {args.query}"], chunk_size=512)[0]
+        (query_embedding,) = query["embedding"]
+        query_vector = np.asarray(query_embedding, dtype=np.float32)
+
+    if args.index_path is not None:
+        args.index_path.write_text(
+            json.dumps(
+                {
+                    "model": args.model,
+                    "chunk_size": args.chunk_size,
+                    "documents": documents,
+                    "chunks": chunks,
+                },
+                ensure_ascii=False,
+            ),
+            encoding="utf-8",
+        )
+        print(f"Saved {len(chunks)} chunks to {args.index_path}")
+    vectors = np.asarray([chunk["vector"] for chunk in chunks], dtype=np.float32)
+    scores = vectors @ query_vector  # Cosine similarity for normalized vectors.
+    for index in np.argsort(-scores)[: args.top_k]:
+        chunk = chunks[index]
+        start, end = chunk["char_range"]
+        excerpt = documents[chunk["document_id"]][start:end]
+        print(
+            f"{scores[index]:.4f} {chunk['document_id']} "
+            f"chunk {chunk['chunk_index']}: {excerpt!r}"
+        )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--url", default="http://localhost:8000")
+    parser.add_argument("--model", default="intfloat/multilingual-e5-small")
+    parser.add_argument("--chunk-size", type=int, default=16)
+    parser.add_argument(
+        "--query", default="Which city has museums with art collections?"
+    )
+    parser.add_argument("--top-k", type=int, default=3)
+    parser.add_argument(
+        "--index-path", type=Path, help="Optionally save the index as JSON."
+    )
+    main(parser.parse_args())

tests/entrypoints/pooling/embed/test_offline.py

@@ -58,10 +58,10 @@ def test_list_prompts(llm: LLM):
     prompts: list[PromptType] = [prompt, prompt_token_ids]
     outputs = llm.embed(prompts, use_tqdm=False)
     assert len(outputs) == 2
-    for i in range(len(outputs)):
-        assert isinstance(outputs[i], EmbeddingRequestOutput)
-        assert outputs[i].prompt_token_ids == prompt_token_ids
-        assert len(outputs[i].outputs.embedding) == embedding_size
+    for output in outputs:
+        assert isinstance(output, EmbeddingRequestOutput)
+        assert output.prompt_token_ids == prompt_token_ids
+        assert len(output.outputs.embedding) == embedding_size
 
 
 @pytest.mark.skip_global_cleanup
@@ -72,7 +72,11 @@ def test_pooling_params(llm: LLM):
             pooling_params=PoolingParams(use_activation=normalize),
             use_tqdm=False,
         )
-        return torch.tensor([x.outputs.embedding for x in outputs])
+        embeddings = []
+        for output in outputs:
+            assert isinstance(output, EmbeddingRequestOutput)
+            embeddings.append(output.outputs.embedding)
+        return torch.tensor(embeddings)
 
     default = get_outputs(normalize=None)
     w_normal = get_outputs(normalize=True)

tests/entrypoints/pooling/test_factories.py

@@ -1,16 +1,25 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from types import SimpleNamespace
 from unittest.mock import MagicMock
 
 import pytest
+from fastapi import FastAPI
+from fastapi.responses import JSONResponse
+from fastapi.testclient import TestClient
 
+from vllm import PoolingParams
 from vllm.entrypoints.pooling import factories
 from vllm.entrypoints.pooling.factories import init_pooling_io_processors
 from vllm.entrypoints.pooling.pooling.io_processor import (
     PluginWithIOProcessorPlugins,
     UnsupportedCombinedTaskIOProcessor,
 )
+from vllm.entrypoints.pooling.token_embed.serving import ServingTokenEmbedding
+from vllm.entrypoints.serve.exception_handling.register import init_exception_handler
+from vllm.exceptions import VLLMValidationError
+from vllm.pooling_params import LateChunkingParams
 
 
 def _bge_m3_config(io_processor_plugin=None):
@@ -92,3 +101,93 @@ def test_combined_task_plain_pooling_request_has_actionable_error(monkeypatch):
     io_processor = serving.get_io_processor(request)
     with pytest.raises(ValueError, match="plugin request with a 'data' field"):
         io_processor.create_pooling_params(request)
+
+
+@pytest.mark.parametrize("task", ["embed", "token_embed"])
+def test_embedding_route_validates_the_configured_task(task):
+    """Invalid chunk requests must never fall back to the ordinary schema."""
+    app = FastAPI()
+    app.state.args = SimpleNamespace(log_error_stack=False)
+    init_exception_handler(app)
+    model_config = MagicMock()
+    model_config.get_pooling_task.return_value = task
+    model_config.is_cross_encoder = False
+    factories.register_pooling_api_routers(app, ("embed", "token_embed"), model_config)
+    received = []
+
+    async def handle(request, raw_request):
+        received.append(request)
+        return JSONResponse({"task": task})
+
+    if task == "token_embed":
+        app.state.serving_token_embedding = handle
+    else:
+        app.state.serving_embedding = handle
+
+    with TestClient(app) as client:
+        if task == "embed":
+            assert client.post("/v1/embeddings", json={"input": "text"}).json() == {
+                "task": "embed"
+            }
+            assert client.post("/v2/embed", json={"texts": ["text"]}).status_code == 200
+            return
+
+        assert client.post("/v2/embed", json={"texts": ["text"]}).status_code == 404
+        bad_requests: list[dict[str, object]] = [
+            {"late_chunking_params": {}},
+            {"late_chunking_params": {"chunk_size": 2, "metadata": {}}},
+            {"encoding_format": "base64"},
+            {"dimensions": 128},
+        ]
+        bad_requests.extend(
+            {"late_chunking_params": {"chunk_size": value}}
+            for value in (0, -1, True, 1.5, "2", 2**63, 2**64)
+        )
+        for overrides in bad_requests:
+            body: dict[str, object] = {"input": "text", **overrides}
+            if "late_chunking_params" not in overrides and overrides:
+                body["late_chunking_params"] = {"chunk_size": 2}
+            response = client.post("/v1/embeddings", json=body)
+            assert response.status_code == 400, response.text
+        assert not received
+
+        for body in (
+            {"input": "text"},
+            {"input": ["text", "another text"], "late_chunking_params": None},
+            {"input": [1, 2]},
+        ):
+            response = client.post("/v1/embeddings", json=body)
+            assert response.json() == {"task": "token_embed"}
+            params = received[-1].to_pooling_params()
+            assert params.task == "token_embed"
+            assert params.late_chunking_params is None
+
+        for chunk_size in (1, 1024, 2**63 - 1):
+            response = client.post(
+                "/v1/embeddings",
+                json={
+                    "input": "text",
+                    "late_chunking_params": {"chunk_size": chunk_size},
+                },
+            )
+            assert response.json() == {"task": "token_embed"}
+            params = received[-1].to_pooling_params()
+            assert params.task == "token_embed"
+            assert params.late_chunking_params.chunk_size == chunk_size
+
+
+def test_token_embedding_keeps_model_specific_chunking_rejection():
+    vllm_config, renderer, chat_template_config = _bge_m3_config()
+    vllm_config.model_config.get_pooling_task.return_value = "token_embed"
+    vllm_config.model_config.architecture = "JinaForRanking"
+    serving = object.__new__(ServingTokenEmbedding)
+    processor = serving.init_io_processor(
+        vllm_config=vllm_config,
+        renderer=renderer,
+        chat_template_config=chat_template_config,
+    )
+    with pytest.raises(VLLMValidationError, match="JinaForRanking.*late chunking"):
+        processor.render(
+            {"params": PoolingParams(late_chunking_params=LateChunkingParams(2))}
+        )
+    renderer.render_cmpl.assert_not_called()

tests/entrypoints/pooling/token_embed/test_late_chunking.py

@@ -1,12 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import asyncio
+import json
 from types import SimpleNamespace
 from unittest import mock
 from unittest.mock import Mock
 
 import pytest
 import torch
+from fastapi.responses import JSONResponse, ORJSONResponse
 
 from vllm import PoolingParams
 from vllm.config import PoolerConfig
@@ -14,14 +17,23 @@ 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.token_embed.late_chunking import (
+    build_late_chunking_metadata,
+)
+from vllm.entrypoints.pooling.token_embed.protocol import (
+    TokenEmbeddingRequest,
+    TokenEmbeddingResponse,
+)
+from vllm.entrypoints.pooling.token_embed.serving import ServingTokenEmbedding
 from vllm.entrypoints.pooling.typing import (
     OfflineEncodeInputsContext,
     OfflineOutputsContext,
+    PoolingServeContext,
 )
-from vllm.exceptions import VLLMValidationError
+from vllm.exceptions import GenerationError, VLLMValidationError
 from vllm.outputs import (
+    EmbeddingRequestOutput,
     LateChunk,
     PoolingOutput,
     PoolingRequestOutput,
@@ -36,6 +48,8 @@ def late_chunk_processor():
     processor = object.__new__(TokenEmbedIOProcessor)
     processor.model_config = SimpleNamespace(
         is_encoder_decoder=False,
+        encoder_config={},
+        max_model_len=32,
         architecture="NomicBertModel",
         model_impl="auto",
         is_matryoshka=False,
@@ -188,6 +202,8 @@ def test_late_chunk_ranges_reject_missing_or_invalid_offsets(offsets):
 
 def _mock_chunk_llm(processor, outputs):
     llm = mock.Mock(spec=PoolingOfflineMixin)
+    llm.pooling_task = "token_embed"
+    llm.embed = PoolingOfflineMixin.embed.__get__(llm)
     llm.encode = PoolingOfflineMixin.encode.__get__(llm)
     llm._run_tiling_engine = PoolingOfflineMixin._run_tiling_engine.__get__(llm)
     llm.pooling_io_processors = {"token_embed": processor}
@@ -243,9 +259,7 @@ def test_late_chunk_mapping_preserves_order_mixed_requests_and_errors(
             _chunk_output(0),
         ]
     ]
-    outputs = llm.encode(
-        ["a b"] * 4, pooling_task="token_embed", pooling_params=params, use_tqdm=False
-    )
+    outputs = llm.embed(["a b"] * 4, pooling_params=params, use_tqdm=False)
     assert outputs[0].late_chunking.chunks == [
         LateChunk((0, 2), (0, 1)),
         LateChunk((2, 4), (2, 3)),
@@ -312,3 +326,215 @@ def test_plain_token_embed_skips_chunk_metadata(late_chunk_processor):
     assert not late_chunk_processor.renderer.render_cmpl.call_args.kwargs[
         "tok_params"
     ].return_token_offsets
+
+
+def _online_chunk_context(texts, chunk_size=2):
+    request = TokenEmbeddingRequest(
+        input=texts,
+        late_chunking_params=(
+            {"chunk_size": chunk_size} if chunk_size is not None else None
+        ),
+    )
+    return PoolingServeContext(
+        request=request,
+        model_name="model",
+        request_id="same-client-id",
+        pooling_params=request.to_pooling_params(),
+        lora_request=None,
+        priorities=0,
+        prompt_extras=None,
+    )
+
+
+@pytest.mark.parametrize("task", ["embed", "token_embed"])
+def test_embed_preserves_output_types(task):
+    llm = mock.Mock(spec=PoolingOfflineMixin)
+    llm.pooling_task = task
+    output = _chunk_output(0)
+    if task == "embed":
+        output.outputs.data = output.outputs.data[0]
+    llm.encode.return_value = [output]
+
+    result = PoolingOfflineMixin.embed(llm, "text")[0]
+    assert llm.encode.call_args.kwargs["pooling_task"] == task
+    if task == "token_embed":
+        assert result is output
+    else:
+        assert isinstance(result, EmbeddingRequestOutput)
+        assert result.outputs.embedding == output.outputs.data.tolist()
+
+
+@pytest.mark.parametrize(
+    "runner, task, supported_tasks, error",
+    [
+        ("pooling", "classify", ("classify",), "--convert embed"),
+        (
+            "pooling",
+            "classify",
+            ("classify", "embed"),
+            r'PoolerConfig\(task="embed"\)',
+        ),
+        ("generate", None, ("generate",), "--runner pooling"),
+    ],
+)
+def test_embed_preserves_task_validation(runner, task, supported_tasks, error):
+    llm = mock.Mock(spec=PoolingOfflineMixin)
+    llm.runner_type = runner
+    llm.pooling_task = task
+    llm.supported_tasks = supported_tasks
+    llm.encode = PoolingOfflineMixin.encode.__get__(llm)
+    llm._verify_pooling_task = PoolingOfflineMixin._verify_pooling_task.__get__(llm)
+
+    with pytest.raises(ValueError, match=error):
+        PoolingOfflineMixin.embed(llm, "text", use_tqdm=False)
+    llm._run_tiling_engine.assert_not_called()
+
+
+@pytest.mark.parametrize("chunk_size", [None, 1])
+@pytest.mark.parametrize(
+    "response_cls, dtype",
+    [
+        (JSONResponse, torch.float32),
+        (ORJSONResponse, torch.float32),
+        (ORJSONResponse, torch.bfloat16),
+    ],
+)
+def test_token_embedding_response_preserves_ndarray_and_optional_ranges(
+    late_chunk_processor, chunk_size, response_cls, dtype
+):
+    if response_cls is ORJSONResponse:
+        pytest.importorskip("orjson")
+    ctx = _online_chunk_context("a b", chunk_size=chunk_size)
+    processor = late_chunk_processor
+    ctx.engine_inputs = [
+        processor.render(r) for r in processor.get_request_factory_online(ctx)
+    ]
+    output = _chunk_output(0, rows=4)
+    # Exercise the existing encoder's contiguous-copy and dtype fallback paths.
+    data = (torch.arange(16).reshape(4, 4).to(dtype) / 10).T
+    output.outputs.data = data
+    ctx.final_res_batch = [output]
+    processor.post_process_online(ctx)
+    serving = object.__new__(ServingTokenEmbedding)
+    serving.json_response_cls = response_cls
+
+    with mock.patch.object(
+        response_cls, "render", autospec=True, side_effect=response_cls.render
+    ) as render:
+        response = serving._build_response(ctx)
+    encoded = render.call_args.args[1]["data"][0]["embedding"]
+    if response_cls is ORJSONResponse and dtype == torch.float32:
+        # An ndarray must reach the renderer, without Pydantic expanding it.
+        assert torch.equal(torch.from_numpy(encoded), data)
+    body = json.loads(response.body)
+    TokenEmbeddingResponse.model_validate(body)
+    item = body["data"][0]
+    torch.testing.assert_close(torch.tensor(item["embedding"], dtype=dtype), data)
+    if chunk_size is None:
+        assert "late_chunking" not in item
+        assert not ctx.late_chunking
+        assert not processor.renderer.render_cmpl.call_args.kwargs[
+            "tok_params"
+        ].return_token_offsets
+    else:
+        assert item["late_chunking"]["chunks"][0]["char_range"] is None
+
+
+@pytest.mark.parametrize("response_cls", [JSONResponse, ORJSONResponse])
+def test_online_chunk_ranges_are_per_document(late_chunk_processor, response_cls):
+    if response_cls is ORJSONResponse:
+        pytest.importorskip("orjson")
+    processor = late_chunk_processor
+    ctx = _online_chunk_context(["a b", "xy z"])
+
+    def render_cmpl(**kwargs):
+        text = kwargs["prompts"][0]["prompt"]
+        split = text.index(" ")
+        return [
+            {
+                "type": "token",
+                "prompt_token_ids": [101, 1, 2, 102],
+                "prompt_token_offsets": [
+                    (0, 0),
+                    (0, split),
+                    (split + 1, len(text)),
+                    (0, 0),
+                ],
+            }
+        ]
+
+    processor.renderer.render_cmpl.side_effect = render_cmpl
+    requests = processor.get_request_factory_online(ctx)
+    # Render out of order, as the online thread pool can do.
+    second = processor.render(requests[1])
+    first = processor.render(requests[0])
+    ctx.engine_inputs = [first, second]
+    ctx.final_res_batch = [_chunk_output(0), _chunk_output(1)]
+    processor.post_process_online(ctx)
+
+    serving = object.__new__(ServingTokenEmbedding)
+    serving.json_response_cls = response_cls
+    body = json.loads(serving._build_response(ctx).body)
+    assert body["usage"]["prompt_tokens"] == 8
+    assert [item["index"] for item in body["data"]] == [0, 1]
+    assert [item["embedding"] for item in body["data"]] == [
+        output.outputs.data.tolist() for output in ctx.final_res_batch
+    ]
+    assert [
+        item["late_chunking"]["chunks"][0]["char_range"] for item in body["data"]
+    ] == [[0, 1], [0, 2]]
+    assert first["params"] is second["params"] is ctx.pooling_params
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("outcome", ["success", "error", "cancel"])
+async def test_online_chunk_mapping_survives_reordered_completion_and_failure(
+    late_chunk_processor, outcome
+):
+    processor = late_chunk_processor
+    serving = object.__new__(ServingTokenEmbedding)
+    ctx = _online_chunk_context(["a b", "a b"])
+    ctx.engine_inputs = [
+        processor.render(r) for r in processor.get_request_factory_online(ctx)
+    ]
+    entered = asyncio.Event()
+    closed = asyncio.Event()
+    error = RequestError("test_error", "original error")
+
+    async def results():
+        try:
+            yield 1, _chunk_output(1)
+            if outcome == "cancel":
+                entered.set()
+                await asyncio.Event().wait()
+            yield 0, _chunk_output(0, error=error if outcome == "error" else None)
+        finally:
+            closed.set()
+
+    ctx.result_generator = results()
+    task = asyncio.create_task(serving._collect_batch(ctx))
+    if outcome == "cancel":
+        await entered.wait()
+        task.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await task
+    elif outcome == "error":
+        with pytest.raises(GenerationError, match="original error"):
+            await task
+        await ctx.result_generator.aclose()
+    else:
+        await task
+        processor.post_process_online(ctx)
+        assert [result.request_id for result in ctx.final_res_batch] == ["0", "1"]
+        assert all(result.late_chunking is not None for result in ctx.final_res_batch)
+    assert closed.is_set()
+
+    # A later request with the same caller ID must get its own mapping.
+    following = _online_chunk_context("a b", chunk_size=1)
+    following.engine_inputs = [
+        processor.render(r) for r in processor.get_request_factory_online(following)
+    ]
+    following.final_res_batch = [_chunk_output(0, rows=4)]
+    processor.post_process_online(following)
+    assert following.final_res_batch[0].late_chunking.chunk_size == 1
+    assert len(following.final_res_batch[0].late_chunking.chunks) == 4

tests/entrypoints/pooling/token_embed/test_offline.py

@@ -9,6 +9,7 @@ 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.outputs import PoolingOutput
 from vllm.pooling_params import LateChunkingParams
 from vllm.tasks import PoolingTask
 
@@ -39,16 +40,20 @@ def llm(vllm_runner):
 
 
 @pytest.mark.skip_global_cleanup
-def test_str_prompts(llm: LLM):
-    outputs = llm.encode(prompt, pooling_task="token_embed", use_tqdm=False)
+@pytest.mark.parametrize("method", ["encode", "embed"])
+def test_str_prompts(llm: LLM, method: str):
+    kwargs = {"pooling_task": "token_embed"} if method == "encode" else {}
+    outputs = getattr(llm, method)(prompt, use_tqdm=False, **kwargs)
     assert len(outputs) == 1
     assert isinstance(outputs[0], PoolingRequestOutput)
     assert outputs[0].outputs.data.shape == (11, 384)
 
 
 @pytest.mark.skip_global_cleanup
-def test_token_ids_prompts(llm: LLM):
-    outputs = llm.encode([prompt_token_ids], pooling_task="token_embed", use_tqdm=False)
+@pytest.mark.parametrize("method", ["encode", "embed"])
+def test_token_ids_prompts(llm: LLM, method: str):
+    kwargs = {"pooling_task": "token_embed"} if method == "encode" else {}
+    outputs = getattr(llm, method)([prompt_token_ids], use_tqdm=False, **kwargs)
     assert len(outputs) == 1
     assert isinstance(outputs[0], PoolingRequestOutput)
     assert outputs[0].outputs.data.shape == (11, 384)
@@ -193,10 +198,10 @@ def test_e5_late_chunking_mixed_requests(llm: LLM, normalize: bool):
         )
         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):
+    actual = llm.embed(texts, pooling_params=params, use_tqdm=False)
+    for index, (original, size) in enumerate(zip(raw, sizes)):
+        result = actual[index]
+        assert isinstance(result.outputs, PoolingOutput)
         expected = original.outputs.data
         if size is not None:
             expected = torch.stack(

tests/entrypoints/pooling/token_embed/test_online.py

@@ -1,12 +1,18 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import asyncio
 
+import httpx
 import pytest
 import requests
+import torch
+from transformers import AutoTokenizer
 
+from tests.models.utils import check_embeddings_close
 from tests.utils import RemoteOpenAIServer
 from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse
+from vllm.entrypoints.pooling.token_embed.protocol import TokenEmbeddingResponse
 
 MODEL_NAME = "intfloat/multilingual-e5-small"
 DTYPE = "bfloat16"
@@ -31,6 +37,13 @@ input_tokens = [
     2,
 ]
 
+LATE_CHUNK_CASES = [(1, False), (5, True), (2**63 - 1, False)]
+LATE_CHUNK_TEXTS = [
+    "passage: Berlin has museums. They display art.",
+    "passage: 中文 😀 e\u0301",
+    "   ",
+]
+
 
 @pytest.fixture(scope="module")
 def server():
@@ -99,3 +112,152 @@ async def test_pooling_not_supported(
     else:
         err_msg = f"Unsupported task: {task!r}"
     assert response.json()["error"]["message"].startswith(err_msg)
+
+
+@pytest.mark.parametrize("inputs", [input_text, input_tokens, [input_text, "Short."]])
+def test_embeddings_without_chunking_match_pooling(server, inputs):
+    body = {"model": MODEL_NAME, "input": inputs, "use_activation": False}
+    expected = requests.post(
+        server.url_for("pooling"), json={**body, "task": "token_embed"}
+    )
+    expected.raise_for_status()
+    response = requests.post(server.url_for("v1/embeddings"), json=body)
+    response.raise_for_status()
+    result = response.json()
+    TokenEmbeddingResponse.model_validate(result)
+    reference = expected.json()
+    assert result["usage"] == reference["usage"]
+    assert len(result["data"]) == len(reference["data"])
+    for index, (actual, raw) in enumerate(zip(result["data"], reference["data"])):
+        assert actual["index"] == index
+        assert "late_chunking" not in actual
+        torch.testing.assert_close(
+            torch.tensor(actual["embedding"]), torch.tensor(raw["data"])
+        )
+
+
+@pytest.mark.parametrize("chunk_size, normalize", LATE_CHUNK_CASES)
+def test_late_chunking_vectors_and_source_ranges(server, chunk_size, normalize):
+    texts = LATE_CHUNK_TEXTS
+    common = {"model": MODEL_NAME, "input": texts}
+    raw_response = requests.post(
+        server.url_for("pooling"),
+        json={**common, "task": "token_embed", "use_activation": False},
+    )
+    raw_response.raise_for_status()
+    raw = raw_response.json()
+    response = requests.post(
+        server.url_for("v1/embeddings"),
+        json={
+            **common,
+            "use_activation": normalize,
+            "late_chunking_params": {"chunk_size": chunk_size},
+        },
+    )
+    response.raise_for_status()
+    result = response.json()
+    TokenEmbeddingResponse.model_validate(result)
+    assert len(result["data"]) == len(raw["data"]) == len(texts)
+    assert result["usage"] == raw["usage"]
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    for index, (text, tokens, output) in enumerate(
+        zip(texts, raw["data"], result["data"])
+    ):
+        assert "late_chunking" not in tokens
+        expected = torch.stack(
+            [part.mean(0) for part in torch.tensor(tokens["data"]).split(chunk_size)]
+        )
+        if normalize:
+            expected = torch.nn.functional.normalize(expected, dim=-1)
+        check_embeddings_close(
+            embeddings_0_lst=output["embedding"],
+            embeddings_1_lst=expected.tolist(),
+            name_0="online_chunks",
+            name_1="contextual_token_reference",
+        )
+        torch.testing.assert_close(
+            torch.tensor(output["embedding"]).norm(dim=-1),
+            expected.norm(dim=-1),
+            rtol=1e-2,
+            atol=1e-3,
+        )
+        assert output["index"] == index
+        metadata = output["late_chunking"]
+        assert metadata["strategy"] == "fixed_tokens"
+        encoded = tokenizer(text, return_offsets_mapping=True)
+        assert metadata["input_tokens"] == len(encoded["input_ids"])
+        assert metadata["chunk_size"] == chunk_size
+        assert len(metadata["chunks"]) == len(expected)
+        for i, chunk in enumerate(metadata["chunks"]):
+            start, end = (
+                i * chunk_size,
+                min((i + 1) * chunk_size, len(encoded["input_ids"])),
+            )
+            assert chunk["token_range"] == [start, end]
+            offsets = [(a, b) for a, b in encoded["offset_mapping"][start:end] if a < b]
+            assert chunk["char_range"] == (
+                [min(a for a, _ in offsets), max(b for _, b in offsets)]
+                if offsets
+                else None
+            )
+
+
+@pytest.mark.parametrize(
+    "overrides",
+    [
+        {"input": [0, 581, 2]},
+        {"input": ""},
+        {"truncate_prompt_tokens": 4},
+        {"padding": "max_length"},
+        {"input": "word " * 600},
+        {"late_chunking_params": {"chunk_size": 2**64}},
+    ],
+)
+def test_late_chunking_http_rejects_unsupported_requests(server, overrides):
+    response = requests.post(
+        server.url_for("v1/embeddings"),
+        json={
+            "model": MODEL_NAME,
+            "input": input_text,
+            "late_chunking_params": {"chunk_size": 5},
+            **overrides,
+        },
+    )
+    assert response.status_code == 400
+    assert response.json()["error"]["message"]
+
+
+@pytest.mark.asyncio
+async def test_late_chunking_concurrent_requests_do_not_share_ranges(server):
+    texts = [
+        "passage: A short text.",
+        "passage: 一个较长的中文文档，内容不同。",
+        input_text,
+    ]
+    sizes = [1, 5, None]
+    async with httpx.AsyncClient(base_url=server.url_root, timeout=60) as client:
+
+        async def send(text, size):
+            body = {
+                "model": MODEL_NAME,
+                "input": text,
+                "request_id": "reused-client-id",
+            }
+            if size:
+                body["late_chunking_params"] = {"chunk_size": size}
+            response = await client.post("/v1/embeddings", json=body)
+            response.raise_for_status()
+            return response.json()["data"][0]
+
+        serial = [await send(text, size) for text, size in zip(texts, sizes)]
+        concurrent = await asyncio.gather(
+            *(send(text, size) for text, size in zip(texts, sizes))
+        )
+    for expected, actual in zip(serial, concurrent):
+        assert actual.get("late_chunking") == expected.get("late_chunking")
+        check_embeddings_close(
+            embeddings_0_lst=actual.get("embedding", actual.get("data")),
+            embeddings_1_lst=expected.get("embedding", expected.get("data")),
+            name_0="concurrent",
+            name_1="serial",
+        )

tests/models/language/pooling/test_colbert.py

@@ -315,10 +315,10 @@ class TestColbertSharedEngine:
         assert scores[0] > scores[1], "ML doc should score higher than Python doc"
         assert scores[2] > scores[1], "DL doc should score higher than Python doc"
 
-    def test_colbert_embed_not_supported(self, colbert_model):
+    def test_colbert_embed_task_not_supported(self, colbert_model):
         """Test that ColBERT model does not support the embed task."""
         with pytest.raises(ValueError, match="Embedding API is not supported"):
-            colbert_model.embed([TEXTS_1[0]])
+            colbert_model.llm.encode([TEXTS_1[0]], pooling_task="embed")
 
 
 @pytest.mark.parametrize(

tests/plugins_tests/test_oot_registration_offline.py

@@ -4,7 +4,7 @@
 import pytest
 
 from tests.utils import create_new_process_for_each_test
-from vllm import LLM, SamplingParams
+from vllm import LLM, EmbeddingRequestOutput, SamplingParams
 from vllm.assets.image import ImageAsset
 from vllm.multimodal.image import convert_image_mode
 
@@ -55,6 +55,7 @@ def test_oot_registration_embedding(
         outputs = llm.embed(prompts)
 
         for output in outputs:
+            assert isinstance(output, EmbeddingRequestOutput)
             assert all(v == 0 for v in output.outputs.embedding)
 
 

vllm/entrypoints/pooling/embed/io_processor.py

@@ -27,12 +27,12 @@ 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 (
+from ..scoring.io_processor import JinaRankingIOProcessorMixin
+from ..token_embed.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,
@@ -675,6 +675,25 @@ class EmbedIOProcessor(PoolingIOProcessor):
 class TokenEmbedIOProcessor(PoolingIOProcessor):
     name = "token_embed"
 
+    def get_request_factory_online(
+        self, ctx: PoolingServeContext
+    ) -> Sequence[AnyRenderParam]:
+        requests = super().get_request_factory_online(ctx)
+        if ctx.pooling_params.late_chunking_params is not None:
+            for request in requests:
+                slot = LateChunkingSlot()
+                ctx.late_chunking.append(slot)
+                request["late_chunking"] = slot
+        return requests
+
+    def post_process_online(self, ctx: PoolingServeContext):
+        if ctx.late_chunking:
+            for output, slot in zip(
+                ctx.final_res_batch, ctx.late_chunking, strict=True
+            ):
+                assert slot.metadata is not None
+                attach_late_chunking_metadata(output, slot.metadata)
+
     def get_request_factory_offline(
         self, ctx: AnyOfflineInputsContext
     ) -> tuple[RequestFactory, int]:

vllm/entrypoints/pooling/factories.py

@@ -129,7 +129,11 @@ def register_pooling_api_routers(
 
         app.include_router(classify_router)
 
-    if "embed" in supported_tasks:
+    if pooling_task == "token_embed":
+        from .token_embed.api_router import router as token_embed_router
+
+        app.include_router(token_embed_router)
+    elif "embed" in supported_tasks:
         from .embed.api_router import router as embed_router
 
         app.include_router(embed_router)
@@ -158,6 +162,7 @@ def init_pooling_state(
     from .embed.serving import ServingEmbedding
     from .pooling.serving import ServingPooling
     from .scoring.serving import ServingScores
+    from .token_embed.serving import ServingTokenEmbedding
 
     resolved_chat_template = load_chat_template(args.chat_template)
     pooling_task = model_config.get_pooling_task(supported_tasks)
@@ -192,6 +197,16 @@ def init_pooling_state(
         if pooling_task == "embed"
         else None
     )
+    state.serving_token_embedding = (
+        ServingTokenEmbedding(
+            engine_client,
+            state.openai_serving_models,
+            request_logger=request_logger,
+            chat_template_config=chat_template_config,
+        )
+        if pooling_task == "token_embed"
+        else None
+    )
     state.serving_classification = (
         ServingClassification(
             engine_client,

vllm/entrypoints/pooling/offline.py

@@ -202,8 +202,8 @@ class PoolingOfflineMixin(OfflineInferenceMixin):
         pooling_params: PoolingParams | Sequence[PoolingParams] | None = None,
         lora_request: list[LoRARequest] | LoRARequest | None = None,
         tokenization_kwargs: dict[str, Any] | None = None,
-    ) -> list[EmbeddingRequestOutput]:
-        """Generate an embedding vector for each prompt.
+    ) -> list[EmbeddingRequestOutput] | list[PoolingRequestOutput]:
+        """Generate embeddings using the model's configured pooling task.
 
         This class automatically batches the given prompts, considering
         the memory constraint. For the best performance, put all of your prompts
@@ -223,8 +223,11 @@ class PoolingOfflineMixin(OfflineInferenceMixin):
             tokenization_kwargs: Overrides for `tokenizer.encode`.
 
         Returns:
-            A list of `EmbeddingRequestOutput` objects containing the
-            embedding vectors in the same order as the input prompts.
+            For `embed`, `EmbeddingRequestOutput` objects with one vector per
+            prompt in `outputs.embedding`. For `token_embed`,
+            `PoolingRequestOutput` objects with a matrix in `outputs.data`.
+            Rows are token vectors, or chunk vectors when late chunking is
+            enabled. Results follow input order and preserve chunk metadata.
 
         """
         items = self.encode(
@@ -232,10 +235,14 @@ class PoolingOfflineMixin(OfflineInferenceMixin):
             use_tqdm=use_tqdm,
             pooling_params=pooling_params,
             lora_request=lora_request,
-            pooling_task="embed",
+            pooling_task="token_embed"
+            if self.pooling_task == "token_embed"
+            else "embed",
             tokenization_kwargs=tokenization_kwargs,
         )
 
+        if self.pooling_task == "token_embed":
+            return items
         return [EmbeddingRequestOutput.from_base(item) for item in items]
 
     def classify(

vllm/entrypoints/pooling/token_embed/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/entrypoints/pooling/token_embed/api_router.py

@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from http import HTTPStatus
+
+from fastapi import APIRouter, Depends, Request
+
+from vllm.entrypoints.serve.engine.protocol import ErrorResponse
+from vllm.entrypoints.serve.utils.api_utils import (
+    load_aware_call,
+    validate_json_request,
+    with_cancellation,
+)
+
+from .protocol import TokenEmbeddingRequest
+from .serving import ServingTokenEmbedding
+
+router = APIRouter()
+
+
+def token_embedding(request: Request) -> ServingTokenEmbedding:
+    return request.app.state.serving_token_embedding
+
+
+@router.post(
+    "/v1/embeddings",
+    dependencies=[Depends(validate_json_request)],
+    responses={
+        HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
+        HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
+    },
+)
+@with_cancellation
+@load_aware_call
+async def create_token_embedding(
+    request: TokenEmbeddingRequest,
+    raw_request: Request,
+):
+    handler = token_embedding(raw_request)
+    return await handler(request, raw_request)

vllm/entrypoints/pooling/token_embed/late_chunking.py

@@ -10,7 +10,7 @@ from vllm.exceptions import VLLMValidationError
 from vllm.outputs import LateChunk, LateChunkingMetadata, PoolingRequestOutput
 from vllm.renderers import TokenizeParams
 
-from .typing import AnyRenderParam, EncodeCMPLRenderParams
+from ..typing import AnyRenderParam, EncodeCMPLRenderParams
 
 
 def prepare_late_chunking_input(

vllm/entrypoints/pooling/token_embed/protocol.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import time
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel, UsageInfo
+from vllm.outputs import LateChunkingMetadata
+from vllm.pooling_params import LateChunkingParams, PoolingParams
+from vllm.utils import random_uuid
+
+from ..base.protocol import (
+    CompletionRequestMixin,
+    EmbedRequestMixin,
+    FixedMaxLenTokenizeParamsMixin,
+    PoolingBasicRequestMixin,
+    reject_removed_pooling_parameters,
+)
+
+
+class LateChunkingConfig(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    chunk_size: int = Field(gt=0, le=2**63 - 1, strict=True)
+
+
+class TokenEmbeddingRequest(
+    PoolingBasicRequestMixin,
+    CompletionRequestMixin,
+    EmbedRequestMixin,
+    FixedMaxLenTokenizeParamsMixin,
+):
+    """Token embeddings with optional full-document late chunking."""
+
+    late_chunking_params: LateChunkingConfig | None = None
+    encoding_format: Literal["float"] = "float"
+    dimensions: None = None
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_removed_parameters(cls, data):
+        return reject_removed_pooling_parameters(data)
+
+    def to_pooling_params(self):
+        return PoolingParams(
+            task="token_embed",
+            use_activation=self.use_activation,
+            late_chunking_params=(
+                LateChunkingParams(chunk_size=self.late_chunking_params.chunk_size)
+                if self.late_chunking_params is not None
+                else None
+            ),
+        )
+
+
+class TokenEmbeddingResponseData(OpenAIBaseModel):
+    index: int
+    object: str = "embedding"
+    embedding: list[list[float]]
+    late_chunking: LateChunkingMetadata | None = None
+
+
+class TokenEmbeddingResponse(OpenAIBaseModel):
+    """vLLM extension returning one matrix of token or chunk vectors per input."""
+
+    id: str = Field(default_factory=lambda: f"embd-{random_uuid()}")
+    object: str = "list"
+    created: int = Field(default_factory=lambda: int(time.time()))
+    model: str | None = None
+    data: list[TokenEmbeddingResponseData]
+    usage: UsageInfo

vllm/entrypoints/pooling/token_embed/serving.py

@@ -0,0 +1,58 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from fastapi.responses import JSONResponse
+
+from ..base.serving import PoolingServing
+from ..factories import init_pooling_io_processors
+from ..typing import PoolingServeContext
+from ..utils import (
+    encode_pooling_output_float,
+    encode_pooling_output_float_or_ndarray,
+    get_json_response_cls,
+    get_pooling_usage,
+)
+from .protocol import TokenEmbeddingResponse, TokenEmbeddingResponseData
+
+
+class ServingTokenEmbedding(PoolingServing):
+    """Serve token embeddings and late chunking through /v1/embeddings."""
+
+    request_id_prefix = "embd"
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.json_response_cls = get_json_response_cls()
+
+    def init_io_processor(self, *args, **kwargs):
+        return init_pooling_io_processors(("token_embed",), *args, **kwargs)[
+            "token_embed"
+        ]
+
+    def _build_response(self, ctx: PoolingServeContext) -> JSONResponse:
+        encode = (
+            encode_pooling_output_float_or_ndarray
+            if self.json_response_cls.__name__ == "ORJSONResponse"
+            else encode_pooling_output_float
+        )
+        items = []
+        for index, output in enumerate(ctx.final_res_batch):
+            item = TokenEmbeddingResponseData(
+                index=index,
+                embedding=[],
+                late_chunking=output.late_chunking,
+            ).model_dump(
+                exclude={"late_chunking"} if output.late_chunking is None else None
+            )
+            # Keep the vector matrix out of Pydantic's list conversion.
+            item["embedding"] = encode(output)
+            items.append(item)
+        response = TokenEmbeddingResponse(
+            id=ctx.request_id,
+            created=ctx.created_time,
+            model=ctx.model_name,
+            data=[],
+            usage=get_pooling_usage(ctx.final_res_batch),
+        ).model_dump()
+        response["data"] = items
+        return self.json_response_cls(content=response)

vllm/entrypoints/pooling/typing.py

@@ -40,11 +40,13 @@ from .pooling.protocol import (
 )
 from .scoring.protocol import ScoringRequest, ScoringResponse
 from .scoring.typing import ScoreData, ScoringData
+from .token_embed.protocol import TokenEmbeddingRequest, TokenEmbeddingResponse
 
 PoolingCompletionLikeRequest: TypeAlias = (
     EmbeddingCompletionRequest
     | ClassificationCompletionRequest
     | PoolingCompletionRequest
+    | TokenEmbeddingRequest
 )
 
 PoolingChatLikeRequest: TypeAlias = (
@@ -70,6 +72,7 @@ AnyPoolingResponse: TypeAlias = (
     | PoolingResponse
     | PoolingBytesResponse
     | ScoringResponse
+    | TokenEmbeddingResponse
 )
 
 PoolingRequestT = TypeVar("PoolingRequestT", bound=AnyPoolingRequest)
@@ -109,6 +112,7 @@ class PoolingServeContext(Generic[PoolingRequestT]):
         None
     )
     final_res_batch: list[PoolingRequestOutput] = field(default_factory=list)
+    late_chunking: list[LateChunkingSlot] = field(default_factory=list)
 
     ## for Long Text Embedding with Chunked Processing
     original_engine_inputs: Sequence["PoolingEngineInput"] | None = None