← Linze Shi

Late chunking in vLLM: one encode, many chunk vectors

September – October 2026 · vLLM pooling

Retrieval systems usually split a document into chunks and embed each chunk on its own. That is simple, and it throws away context. If one paragraph says "its maintenance team" and the company it refers to was named two paragraphs earlier, that chunk's embedding has no way of knowing which company it is.

Late chunking changes the order of operations. Run the whole document through the encoder once, so every token's hidden state has attended to the full text, and only then average the token states inside each chunk. Each chunk vector now carries its surroundings with it. The algorithm is simple and already published. Getting it into vLLM correctly was mostly a question of where in the pipeline the averaging happens, and that turned out to have real consequences.

The order of operations is the whole feature

An embedding model doesn't hand you raw hidden states. After pooling it usually applies a head: sometimes a projection, often a dimension cut, almost always an L2 normalization. Late chunking has to put the chunk mean in the place where the model's normal pooling would go, before that head:

chunk_vector = head( mean( h[start:end] ) )        # correct
chunk_vector = mean( head( h[start:end] ) )        # not the same thing

The two are not interchangeable. A linear projection commutes with the mean, but normalization doesn't: the mean of unit vectors is not a unit vector, and normalizing each token first and then averaging gives a different direction from averaging first, because before normalization the tokens with larger norms pull the mean harder, and after it every token counts the same. Any activation in the head breaks the commutation too. This is the main reason I didn't want users doing it on the client side. vLLM can already return per-token embeddings, but by default they come out normalized, and pooling them afterwards quietly produces a different vector than the model's own pooling would.

Sums need more precision than the model

The hidden states are BF16, which keeps eight bits of precision, a relative rounding error of up to about 0.4%. Adding a few hundred values of similar size in BF16 loses low-order bits on every step, and the error grows with chunk length. So the chunk sums accumulate in FP32, and only the finished mean goes on to the head. It's a one-line decision that you only notice when the chunk vectors from a long document stop matching a reference implementation.

Attention direction matters

The idea assumes every token has seen the whole document. That's true for bidirectional encoders like BERT-style embedding models, which are what late chunking was designed for. In a causal model, token i only attends to tokens before it, so a late-chunked vector carries left context only: the second paragraph knows about the first, but not the other way round. That can still beat chunking in isolation, but it's a different object, which is why the testing focused on bidirectional encoders like E5, Nomic and BGE-M3.

Where the reduction lives

one forward pass over the document
  -> chunk means over contextual token states, accumulated in FP32
  -> the model's own embedding head (projection, dimension cut, normalization)
  -> K chunk vectors returned to the caller

Doing the reduction on the GPU, in the worker, also changes how much data moves. An 8,192-token document with 1,024-dimensional FP32 states is 32 MB of token vectors. Sixty-four chunk vectors of the same width are 256 KB. Exporting every token to the frontend just to average most of it away is a lot of copying for nothing.

From the caller's side it's one option on the existing token-embedding task:

from vllm import PoolingParams
from vllm.pooling_params import LateChunkingParams

output = llm.encode(
    "search_document: Berlin is in Germany. It has many museums.",
    pooling_task="token_embed",
    pooling_params=PoolingParams(
        late_chunking_params=LateChunkingParams(chunk_size=8)
    ),
)[0]
vectors = output.outputs.data          # [num_chunks, embedding_dim]
ranges = output.late_chunking.chunks   # token and source-text ranges

The details that decide whether the vectors are right

Token IDs and character offsets come from a single tokenization pass, so the source-text ranges returned to the caller line up exactly with what was pooled. Tokenizing twice, once to find boundaries and once to encode, is an easy way to end up with chunks that are off by a token at the edges.

Special tokens count. They sit inside chunk boundaries and means like any other token, including BGE-M3's BOS token, and the filter that normally strips BOS and EOS from token outputs leaves chunk outputs alone. The last chunk is averaged over the tokens it actually has instead of being padded out to the chunk size, which would drag its mean toward whatever the padding state is.

Anything that would quietly change the input is rejected up front: truncation, padding, text normalization in the renderer, token-row filtering, late-interaction scoring. One ranking model has a custom pooler that returns document and query vectors. It's rejected before tokenization, because otherwise it would return something that looks like an embedding and isn't.

Testing it

The change landed with 203 parameter, IO, pooler, head and tokenizer regression tests, plus 20 offline GPU and tiling tests covering E5, Nomic and BGE-M3 with prefix caching and chunked prefill, on both model runners. An STS12 run with Nomic matched the reference within the existing tolerance. There's no speed claim. The point is correct context-aware chunk vectors without exporting every token.

Testing Nomic at long context surfaced something unrelated: it was running with the default RoPE instead of its configured dynamic scaling. That went in as a separate bug and fix rather than being folded into this change, so it couldn't hide behind it.

Serving it online

The first release was offline, through LLM.encode. I then extended it to the online /v1/embeddings API, and most of that review was about API shape rather than math: whether chunk vectors belong under embed or under encode. That question sounds cosmetic and isn't, because it decides which defaults and which normalization a user gets without asking for them.