vllm - 💡(How to fix) Fix [RFC]: Add `max_tokens_per_doc` support for rerank and scoring endpoints [2 comments, 2 participants]

Motivation

vLLM's reranking and scoring endpoints currently have no mechanism to truncate documents before processing. When users send long documents that exceed the model's context window, requests fail or produce degraded results.

Both the Cohere Rerank API and Jina Reranker API support a max_tokens_per_doc (or equivalent) parameter that truncates each document to a specified token limit before scoring. This is a standard feature in production reranking APIs that vLLM should support.

PR #33315 by @hustxiayang introduced an initial implementation of this feature. This RFC formalizes the design — particularly around offline support, PoolingParams integration, and score template compatibility — to align on the approach and move toward merging.

Proposed Change

Add a max_tokens_per_doc parameter to the rerank/score request schema that truncates each document's token representation before model inference. The implementation should support both online (HTTP API) and offline (LLM.score() / LLM.embed()) usage paths.

API Surface

Online (HTTP): Add max_tokens_per_doc: Optional[int] to the rerank and score request schemas.

{
  "model": "BAAI/bge-reranker-v2-m3",
  "query": "What is deep learning?",
  "documents": ["A very long document...", "Another long document..."],
  "max_tokens_per_doc": 512
}

Offline: Expose via PoolingParams so it is available in both online and offline scenarios:

from vllm import LLM

llm = LLM(model="BAAI/bge-reranker-v2-m3", task="score")
results = llm.score(
    "What is deep learning?",
    ["A very long document...", "Another long document..."],
    pooling_params=PoolingParams(max_tokens_per_doc=512),
)

Truncation Behavior

• Tokenize each document using tokenizer.encode() (per DarkLight1337's suggestion in #33315)
• Truncate the token list to max_tokens_per_doc tokens
• Apply truncation before prompt/template assembly to ensure it respects the per-document limit
• If max_tokens_per_doc is None or not set, no truncation occurs (backward compatible)

Files to Modify

Compatibility

This must work correctly with score templates introduced in #30550 and #31335. Specifically:

• Truncation must happen per-document, before template expansion — the template tokens (query prefix, separators, etc.) should not count toward the max_tokens_per_doc limit
• Tests should cover both template-based and non-template scoring models

Feedback Period.

No response

CC List.

@noooop @DarkLight1337 @hustxiayang

Any Other Things.

• Prior art: PR #33315 contains a working implementation that handles bi-encoder and cross-encoder architectures. The core truncation logic can be reused.
• The Cohere /v2/rerank documentation defaults max_tokens_per_doc to 4096. We should not set a default and leave it as None to preserve backward compatibility.
• Future extension: this pattern could also apply to embedding endpoints (/v1/embeddings) for long-document truncation, but that is out of scope for this RFC.

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.