@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / edbc1abd1c
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 5, 06:00 UTC

feat: add max_tokens_per_doc in rerank request. (#38827)

Signed-off-by: Jesus Federico <jefp@amazon.com> Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: wang.yuqi <yuqi.wang@daocloud.io>

Jesus Federico committed Apr 13, 2026 at 04:24 UTC edbc1abd1c206db43dd4d297d773532edbb226a2
9 files changed +544 -27
tests/entrypoints/pooling/scoring/test_cross_encoder_offline.py
+29
@@ -112,6 +112,35 @@ def test_classify(llm):
112 assert len(outputs[0].outputs.data) == 1
113
114
115 +@pytest.mark.skip_global_cleanup
116 +def test_max_tokens_per_doc(llm: LLM):
117 + """Test max_tokens_per_doc via PoolingParams.extra_kwargs (offline)."""
118 + long_doc = "The capital of France is Paris. " * 20
119 +
120 + # Without truncation
121 + outputs_no_limit = llm.score(
122 + TEXTS_1[0],
123 + long_doc,
124 + use_tqdm=False,
125 + )
126 +
127 + # With truncation via extra_kwargs
128 + outputs_with_limit = llm.score(
129 + TEXTS_1[0],
130 + long_doc,
131 + pooling_params=PoolingParams(extra_kwargs={"max_tokens_per_doc": 10}),
132 + use_tqdm=False,
133 + )
134 +
135 + assert len(outputs_no_limit) == 1
136 + assert len(outputs_with_limit) == 1
137 +
138 + # Truncated version should have fewer prompt tokens
139 + no_limit_tokens = len(outputs_no_limit[0].prompt_token_ids)
140 + with_limit_tokens = len(outputs_with_limit[0].prompt_token_ids)
141 + assert with_limit_tokens < no_limit_tokens
142 +
143 +
144 def test_pooling_params(llm: LLM):
145 def get_outputs(use_activation):
146 outputs = llm.score(
tests/entrypoints/pooling/scoring/test_cross_encoder_online.py
+72
@@ -471,6 +471,78 @@ async def test_pooling_token_classify(server: RemoteOpenAIServer):
471 assert len(poolings.data[0].data[0]) == 1
472
473
474 +@pytest.mark.asyncio
475 +async def test_rerank_max_tokens_per_doc(
476 + server: RemoteOpenAIServer,
477 +):
478 + """Test that max_tokens_per_doc actually reduces the token count."""
479 + query = "What is the capital of France?"
480 + # Use a doc that fits within max_model_len=100 (query ~8 tokens + 4 special)
481 + long_doc = "The capital of France is Paris. " * 10 # ~70 tokens
482 +
483 + # Without max_tokens_per_doc
484 + response_no_limit = requests.post(
485 + server.url_for("rerank"),
486 + json={
487 + "model": MODEL_NAME,
488 + "query": query,
489 + "documents": [long_doc],
490 + "truncate_prompt_tokens": 99,
491 + },
492 + )
493 + response_no_limit.raise_for_status()
494 + rerank_no_limit = RerankResponse.model_validate(response_no_limit.json())
495 +
496 + # With max_tokens_per_doc
497 + response_with_limit = requests.post(
498 + server.url_for("rerank"),
499 + json={
500 + "model": MODEL_NAME,
501 + "query": query,
502 + "documents": [long_doc],
503 + "max_tokens_per_doc": 10,
504 + },
505 + )
506 + response_with_limit.raise_for_status()
507 + rerank_with_limit = RerankResponse.model_validate(response_with_limit.json())
508 +
509 + assert rerank_with_limit.usage.prompt_tokens < rerank_no_limit.usage.prompt_tokens
510 +
511 +
512 +@pytest.mark.asyncio
513 +async def test_rerank_max_tokens_per_doc_validation(
514 + server: RemoteOpenAIServer,
515 +):
516 + """Test that max_tokens_per_doc validation works correctly."""
517 + query = "What is the capital of France?"
518 + documents = ["The capital of France is Paris."]
519 +
520 + # Test with max_tokens_per_doc=0 (should succeed — means no truncation)
521 + response = requests.post(
522 + server.url_for("rerank"),
523 + json={
524 + "model": MODEL_NAME,
525 + "query": query,
526 + "documents": documents,
527 + "max_tokens_per_doc": 0,
528 + },
529 + )
530 + response.raise_for_status()
531 +
532 + # Test with invalid max_tokens_per_doc (negative)
533 + response = requests.post(
534 + server.url_for("rerank"),
535 + json={
536 + "model": MODEL_NAME,
537 + "query": query,
538 + "documents": documents,
539 + "max_tokens_per_doc": -5,
540 + },
541 + )
542 + assert response.status_code == 400
543 + assert "max_tokens_per_doc must be a non-negative integer" in response.text
544 +
545 +
546 @pytest.mark.asyncio
547 @pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
548 async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str):
tests/models/language/pooling/test_max_tokens_per_doc.py new
+209
@@ -0,0 +1,209 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Tests for max_tokens_per_doc and max_tokens_per_query.
5 +"""
6 +
7 +import json
8 +import os
9 +from dataclasses import dataclass
10 +
11 +import pytest
12 +import requests
13 +
14 +from tests.utils import VLLM_PATH, RemoteOpenAIServer
15 +from vllm.entrypoints.pooling.scoring.protocol import RerankResponse
16 +
17 +os.environ["VLLM_LOGGING_LEVEL"] = "WARNING"
18 +
19 +TEMPLATE_DIR = str(VLLM_PATH / "examples/pooling/score/template")
20 +
21 +long_query = "What is the capital of France?" * 20
22 +long_doc = "The capital of France is Paris. " * 20
23 +
24 +
25 +@dataclass
26 +class TestConfig:
27 + model: str
28 + args: list[str]
29 + without_truncated_prompt_tokens: int
30 + with_max_tokens_per_query_prompt_tokens: int
31 + with_max_tokens_per_doc_prompt_tokens: int
32 + with_max_tokens_per_query_and_doc_prompt_tokens: int
33 +
34 +
35 +RERANK_CONFIGS = [
36 + # 1. cross-encoder
37 + TestConfig(
38 + model="jinaai/jina-reranker-v2-base-multilingual",
39 + args=[
40 + "--enforce-eager",
41 + "--max-model-len",
42 + "1024",
43 + "--trust-remote-code",
44 + ],
45 + without_truncated_prompt_tokens=284,
46 + with_max_tokens_per_query_prompt_tokens=154,
47 + with_max_tokens_per_doc_prompt_tokens=154,
48 + with_max_tokens_per_query_and_doc_prompt_tokens=24,
49 + ),
50 + # 2. cross-encoder + score template
51 + TestConfig(
52 + model="Qwen/Qwen3-Reranker-0.6B",
53 + args=[
54 + "--enforce-eager",
55 + "--max-model-len",
56 + "1024",
57 + "--hf-overrides",
58 + json.dumps(
59 + {
60 + "architectures": ["Qwen3ForSequenceClassification"],
61 + "classifier_from_token": ["no", "yes"],
62 + "is_original_qwen3_reranker": True,
63 + }
64 + ),
65 + "--chat-template",
66 + os.path.join(TEMPLATE_DIR, "qwen3_reranker.jinja"),
67 + ],
68 + without_truncated_prompt_tokens=352,
69 + with_max_tokens_per_query_prompt_tokens=223,
70 + with_max_tokens_per_doc_prompt_tokens=221,
71 + with_max_tokens_per_query_and_doc_prompt_tokens=92,
72 + ),
73 + # 3. bi-encoder
74 + TestConfig(
75 + model="intfloat/multilingual-e5-small",
76 + args=[
77 + "--enforce-eager",
78 + "--max-model-len",
79 + "512",
80 + "--trust-remote-code",
81 + ],
82 + without_truncated_prompt_tokens=286,
83 + with_max_tokens_per_query_prompt_tokens=156,
84 + with_max_tokens_per_doc_prompt_tokens=155,
85 + with_max_tokens_per_query_and_doc_prompt_tokens=25,
86 + ),
87 + # 4. late-interaction
88 + TestConfig(
89 + model="answerdotai/answerai-colbert-small-v1",
90 + args=[
91 + "--enforce-eager",
92 + "--max-model-len",
93 + "512",
94 + "--trust-remote-code",
95 + ],
96 + without_truncated_prompt_tokens=285,
97 + with_max_tokens_per_query_prompt_tokens=155,
98 + with_max_tokens_per_doc_prompt_tokens=155,
99 + with_max_tokens_per_query_and_doc_prompt_tokens=25,
100 + ),
101 + # 5. jinaai/jina-reranker-v3
102 + TestConfig(
103 + model="jinaai/jina-reranker-v3",
104 + args=[
105 + "--enforce-eager",
106 + "--max-model-len",
107 + "1024",
108 + "--trust-remote-code",
109 + ],
110 + without_truncated_prompt_tokens=567,
111 + with_max_tokens_per_query_prompt_tokens=308,
112 + with_max_tokens_per_doc_prompt_tokens=436,
113 + with_max_tokens_per_query_and_doc_prompt_tokens=177,
114 + ),
115 +]
116 +
117 +
118 +@pytest.fixture(scope="module", params=RERANK_CONFIGS, ids=lambda c: c.model)
119 +def server(request):
120 + config: TestConfig = request.param
121 + with RemoteOpenAIServer(config.model, config.args) as remote_server:
122 + yield config, remote_server
123 +
124 +
125 +def test_without_truncated(server):
126 + """Test that max_tokens_per_doc truncates documents correctly."""
127 + config, remote_server = server
128 +
129 + response = requests.post(
130 + remote_server.url_for("rerank"),
131 + json={"model": config.model, "query": long_query, "documents": [long_doc]},
132 + )
133 + response.raise_for_status()
134 + rerank = RerankResponse.model_validate(response.json())
135 +
136 + assert rerank.id is not None
137 + assert rerank.results is not None
138 + assert len(rerank.results) == 1
139 + assert rerank.usage.prompt_tokens == config.without_truncated_prompt_tokens
140 +
141 +
142 +def test_max_tokens_per_query(server):
143 + """Test that max_tokens_per_doc truncates documents correctly."""
144 + config, remote_server = server
145 +
146 + response = requests.post(
147 + remote_server.url_for("rerank"),
148 + json={
149 + "model": config.model,
150 + "query": long_query,
151 + "documents": [long_doc],
152 + "max_tokens_per_query": 10,
153 + },
154 + )
155 + response.raise_for_status()
156 + rerank = RerankResponse.model_validate(response.json())
157 +
158 + assert rerank.id is not None
159 + assert rerank.results is not None
160 + assert len(rerank.results) == 1
161 + assert rerank.usage.prompt_tokens == config.with_max_tokens_per_query_prompt_tokens
162 +
163 +
164 +def test_max_tokens_per_doc(server):
165 + """Test that max_tokens_per_doc truncates documents correctly."""
166 + config, remote_server = server
167 +
168 + response = requests.post(
169 + remote_server.url_for("rerank"),
170 + json={
171 + "model": config.model,
172 + "query": long_query,
173 + "documents": [long_doc],
174 + "max_tokens_per_doc": 10,
175 + },
176 + )
177 + response.raise_for_status()
178 + rerank = RerankResponse.model_validate(response.json())
179 +
180 + assert rerank.id is not None
181 + assert rerank.results is not None
182 + assert len(rerank.results) == 1
183 + assert rerank.usage.prompt_tokens == config.with_max_tokens_per_doc_prompt_tokens
184 +
185 +
186 +def test_max_tokens_per_query_and_doc(server):
187 + """Test that max_tokens_per_doc truncates documents correctly."""
188 + config, remote_server = server
189 +
190 + response = requests.post(
191 + remote_server.url_for("rerank"),
192 + json={
193 + "model": config.model,
194 + "query": long_query,
195 + "documents": [long_doc],
196 + "max_tokens_per_query": 10,
197 + "max_tokens_per_doc": 10,
198 + },
199 + )
200 + response.raise_for_status()
201 + rerank = RerankResponse.model_validate(response.json())
202 +
203 + assert rerank.id is not None
204 + assert rerank.results is not None
205 + assert len(rerank.results) == 1
206 + assert (
207 + rerank.usage.prompt_tokens
208 + == config.with_max_tokens_per_query_and_doc_prompt_tokens
209 + )
vllm/entrypoints/pooling/base/serving.py
+1
@@ -232,6 +232,7 @@ class PoolingServingBase(ABC):
232 "greater than max_model_len."
233 " Please request a smaller truncation size."
234 )
235 +
236 return None
237
238 async def _get_trace_headers(
vllm/entrypoints/pooling/scoring/io_processor.py
+179 -8
@@ -25,8 +25,10 @@ from .typing import ScoreData, ScoreInput, ScoringData
25 from .utils import (
26 compress_token_type_ids,
27 compute_maxsim_score,
28 + get_num_special_tokens_for_pair,
29 parse_score_data,
30 score_data_to_prompts,
31 + truncate_text_to_tokens,
32 validate_score_input,
33 )
34
@@ -48,6 +50,64 @@ class ScoringIOProcessor(PoolingIOProcessor):
50 def create_pooling_params(self, request):
51 return request.to_pooling_params(self.pooling_task)
52
53 + def _validate_token_limit(self, value: int, name: str) -> None:
54 + if value < 0:
55 + raise ValueError(f"{name} must be a non-negative integer")
56 + if value >= self.model_config.max_model_len:
57 + raise ValueError(
58 + f"{name} ({value}) must be less "
59 + f"than max_model_len ({self.model_config.max_model_len})."
60 + )
61 +
62 + def _get_token_limits(
63 + self,
64 + request: ScoringRequest | None = None,
65 + pooling_params: PoolingParams | None = None,
66 + ) -> tuple[int, int]:
67 + """Extract and validate token limits from request or pooling_params."""
68 + if request is not None:
69 + max_tokens_per_query = getattr(request, "max_tokens_per_query", 0)
70 + max_tokens_per_doc = getattr(request, "max_tokens_per_doc", 0)
71 + else:
72 + extra = (
73 + (pooling_params.extra_kwargs or {})
74 + if pooling_params is not None
75 + else {}
76 + )
77 + max_tokens_per_query = extra.get("max_tokens_per_query", 0)
78 + max_tokens_per_doc = extra.get("max_tokens_per_doc", 0)
79 +
80 + if max_tokens_per_query != 0:
81 + self._validate_token_limit(max_tokens_per_query, "max_tokens_per_query")
82 + if max_tokens_per_doc != 0:
83 + self._validate_token_limit(max_tokens_per_doc, "max_tokens_per_doc")
84 + return max_tokens_per_query, max_tokens_per_doc
85 +
86 + def _truncate_scoring_data(
87 + self,
88 + scoring_data: ScoringData,
89 + max_tokens_per_query: int = 0,
90 + max_tokens_per_doc: int = 0,
91 + ) -> ScoringData:
92 + """Truncate query/document texts to token limits."""
93 + data_1 = scoring_data.data_1
94 + data_2 = scoring_data.data_2
95 + if max_tokens_per_query > 0:
96 + data_1 = [
97 + truncate_text_to_tokens(d, self.tokenizer, max_tokens_per_query)
98 + if isinstance(d, str)
99 + else d
100 + for d in data_1
101 + ]
102 + if max_tokens_per_doc > 0:
103 + data_2 = [
104 + truncate_text_to_tokens(d, self.tokenizer, max_tokens_per_doc)
105 + if isinstance(d, str)
106 + else d
107 + for d in data_2
108 + ]
109 + return ScoringData(data_1=data_1, data_2=data_2)
110 +
111 def valid_inputs(
112 self,
113 data_1: ScoreInput | list[ScoreInput],
@@ -82,6 +142,15 @@ class BiEncoderIOProcessor(ScoringIOProcessor):
142 raise ValueError(f"Invalid {self.name} request type")
143
144 scoring_data = self.valid_inputs(data_1, data_2)
145 +
146 + max_tokens_per_query, max_tokens_per_doc = self._get_token_limits(
147 + request=request
148 + )
149 + if max_tokens_per_query > 0 or max_tokens_per_doc > 0:
150 + scoring_data = self._truncate_scoring_data(
151 + scoring_data, max_tokens_per_query, max_tokens_per_doc
152 + )
153 +
154 tok_params = request.build_tok_params(self.model_config)
155 engine_inputs = self._pre_process(
156 scoring_data,
@@ -112,10 +181,23 @@ class BiEncoderIOProcessor(ScoringIOProcessor):
181
182 def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]:
183 assert isinstance(ctx.prompts, ScoringData)
184 + assert not isinstance(ctx.pooling_params, Sequence)
185 +
186 tok_params = self.renderer.default_cmpl_tok_params.with_kwargs(
187 **(ctx.tokenization_kwargs or {})
188 )
118 - return self._pre_process(ctx.prompts, tok_params)
189 +
190 + max_tokens_per_query, max_tokens_per_doc = self._get_token_limits(
191 + pooling_params=ctx.pooling_params
192 + )
193 +
194 + scoring_data = ctx.prompts
195 + if max_tokens_per_query > 0 or max_tokens_per_doc > 0:
196 + scoring_data = self._truncate_scoring_data(
197 + scoring_data, max_tokens_per_query, max_tokens_per_doc
198 + )
199 +
200 + return self._pre_process(scoring_data, tok_params)
201
202 def post_process_offline(
203 self,
@@ -217,8 +299,38 @@ class LateInteractionIOProcessor(BiEncoderIOProcessor):
299 class FlashLateInteractionIOProcessor(LateInteractionIOProcessor):
300 name = "flash-late-interaction"
301
220 - def _post_process(self, outputs: list[PoolingRequestOutput], n_queries: int):
221 - return outputs
302 + def post_process_online(
303 + self,
304 + ctx: ScoringServeContext,
305 + ):
306 + assert ctx.query_final_res_batch is not None
307 + assert ctx.final_res_batch is not None
308 + assert isinstance(ctx.n_queries, int)
309 +
310 + # Expand queries if 1:N scoring
311 + if len(ctx.query_final_res_batch) == 1:
312 + ctx.query_final_res_batch = ctx.query_final_res_batch * len(
313 + ctx.final_res_batch
314 + )
315 +
316 + final_res_batch: list[PoolingRequestOutput] = []
317 + for d1, d2 in zip(ctx.query_final_res_batch, ctx.final_res_batch):
318 + padding: list[int] = []
319 + if (pad_token_id := self.pad_token_id) is not None:
320 + padding = [pad_token_id]
321 +
322 + tokens = d1.prompt_token_ids + padding + d2.prompt_token_ids
323 +
324 + final_res_batch.append(
325 + PoolingRequestOutput(
326 + request_id=f"{d1.request_id}_{d2.request_id}",
327 + outputs=d2.outputs,
328 + prompt_token_ids=tokens,
329 + num_cached_tokens=d1.num_cached_tokens + d2.num_cached_tokens,
330 + finished=True,
331 + )
332 + )
333 + ctx.final_res_batch = final_res_batch
334
335
336 class CrossEncoderIOProcessor(ScoringIOProcessor):
@@ -255,6 +367,11 @@ class CrossEncoderIOProcessor(ScoringIOProcessor):
367 raise ValueError(f"Invalid {self.name} request type")
368
369 scoring_data = self.valid_inputs(data_1, data_2)
370 +
371 + max_tokens_per_query, max_tokens_per_doc = self._get_token_limits(
372 + request=request
373 + )
374 +
375 tok_params = request.build_tok_params(self.model_config)
376 pooling_params = self.create_pooling_params(request)
377
@@ -263,6 +380,8 @@ class CrossEncoderIOProcessor(ScoringIOProcessor):
380 tok_params,
381 pooling_params,
382 chat_template=self.chat_template,
383 + max_tokens_per_query=max_tokens_per_query,
384 + max_tokens_per_doc=max_tokens_per_doc,
385 prompt_extras={
386 k: v
387 for k in ("mm_processor_kwargs", "cache_salt")
@@ -283,8 +402,18 @@ class CrossEncoderIOProcessor(ScoringIOProcessor):
402 tok_params = self.renderer.default_cmpl_tok_params.with_kwargs(
403 **(ctx.tokenization_kwargs or {})
404 )
405 +
406 + max_tokens_per_query, max_tokens_per_doc = self._get_token_limits(
407 + pooling_params=ctx.pooling_params
408 + )
409 +
410 engine_inputs, pooling_params_list = self._pre_process(
287 - ctx.prompts, tok_params, ctx.pooling_params, ctx.chat_template
411 + ctx.prompts,
412 + tok_params,
413 + ctx.pooling_params,
414 + ctx.chat_template,
415 + max_tokens_per_query=max_tokens_per_query,
416 + max_tokens_per_doc=max_tokens_per_doc,
417 )
418 ctx.pooling_params = pooling_params_list
419 return engine_inputs
@@ -298,6 +427,8 @@ class CrossEncoderIOProcessor(ScoringIOProcessor):
427 tok_params: TokenizeParams,
428 pooling_params: PoolingParams | None,
429 chat_template: str | None = None,
430 + max_tokens_per_query: int = 0,
431 + max_tokens_per_doc: int = 0,
432 prompt_extras: dict[str, Any] | None = None,
433 ) -> tuple[Sequence[EngineInput], list[PoolingParams]]:
434 # todo: support prompt_extras
@@ -320,6 +451,8 @@ class CrossEncoderIOProcessor(ScoringIOProcessor):
451 data_2=d,
452 encode_kwargs=tok_params.get_encode_kwargs(),
453 chat_template=chat_template,
454 + max_tokens_per_query=max_tokens_per_query,
455 + max_tokens_per_doc=max_tokens_per_doc,
456 )
457
458 if token_type_ids := engine_prompt.pop("token_type_ids", None):
@@ -342,6 +475,8 @@ class CrossEncoderIOProcessor(ScoringIOProcessor):
475 data_2: ScoreData,
476 encode_kwargs: dict[str, Any],
477 chat_template: str | None = None,
478 + max_tokens_per_query: int = 0,
479 + max_tokens_per_doc: int = 0,
480 ):
481 model_config = self.model_config
482 tokenizer = self.tokenizer
@@ -352,25 +487,61 @@ class CrossEncoderIOProcessor(ScoringIOProcessor):
487 model_config,
488 )
489
490 + # Apply truncation before defining closures
491 + if max_tokens_per_query > 0 and isinstance(prompt_1, str):
492 + prompt_1 = truncate_text_to_tokens(
493 + prompt_1, tokenizer, max_tokens_per_query
494 + )
495 + if max_tokens_per_doc > 0 and isinstance(prompt_2, str):
496 + prompt_2 = truncate_text_to_tokens(prompt_2, tokenizer, max_tokens_per_doc)
497 +
498 def default_tokenizer_encode():
499 + local_kwargs = encode_kwargs.copy()
500 +
501 if self.supports_score_template:
502 assert self.model is not None
503 full_prompt = self.model.get_score_template(prompt_1, prompt_2)
504 if full_prompt is None:
505 raise ValueError("Get empty score template from model")
506
362 - prompt_inputs = tokenizer(full_prompt, **encode_kwargs)
507 + prompt_inputs = tokenizer(full_prompt, **local_kwargs)
508 else:
509 if self.use_sep_token:
510 # cross_encoder models defaults to using separating token.
511 + if max_tokens_per_doc > 0 and isinstance(prompt_2, str):
512 + query_tokens = tokenizer.encode(
513 + prompt_1, add_special_tokens=False
514 + )
515 + num_special = get_num_special_tokens_for_pair(tokenizer)
516 + doc_limit_max_length = (
517 + len(query_tokens) + max_tokens_per_doc + num_special
518 + )
519 + existing_max_length = local_kwargs.get("max_length")
520 + if existing_max_length is not None:
521 + effective_max_length = min(
522 + doc_limit_max_length, existing_max_length
523 + )
524 + else:
525 + effective_max_length = doc_limit_max_length
526 + local_kwargs["truncation"] = "only_second"
527 + local_kwargs["max_length"] = effective_max_length
528 +
529 prompt_inputs = tokenizer(
367 - text=prompt_1, text_pair=prompt_2, **encode_kwargs
530 + text=prompt_1, text_pair=prompt_2, **local_kwargs
531 )
532 full_prompt = tokenizer.decode(prompt_inputs["input_ids"])
533 else:
534 # `llm as reranker` defaults to not using separating token.
372 - full_prompt = prompt_1 + prompt_2
373 - prompt_inputs = tokenizer(text=full_prompt, **encode_kwargs)
535 + if max_tokens_per_doc > 0 and isinstance(prompt_2, str):
536 + query_ids = tokenizer.encode(prompt_1, add_special_tokens=False)
537 + doc_ids = tokenizer.encode(prompt_2, add_special_tokens=False)
538 + doc_ids = doc_ids[:max_tokens_per_doc]
539 + input_ids = query_ids + doc_ids
540 + full_prompt = tokenizer.decode(input_ids)
541 + prompt_inputs = {"input_ids": input_ids}
542 + else:
543 + full_prompt = prompt_1 + prompt_2
544 + prompt_inputs = tokenizer(text=full_prompt, **local_kwargs)
545 return full_prompt, prompt_inputs
546
547 # FIXME: For now, we only apply a template when one is explicitly provided.
vllm/entrypoints/pooling/scoring/protocol.py
+19 -19
@@ -20,6 +20,24 @@ from .typing import ScoreContentPartParam, ScoreInput
20
21
22 class ScoreRequestMixin(PoolingBasicRequestMixin, ClassifyRequestMixin):
23 + max_tokens_per_query: int = Field(
24 + default=0,
25 + description=(
26 + "Maximum number of tokens per query. Queries longer than "
27 + "this will be truncated to this length. 0 means no "
28 + "query-level truncation is applied."
29 + ),
30 + )
31 + max_tokens_per_doc: int = Field(
32 + default=0,
33 + description=(
34 + "Maximum number of tokens per document. Documents longer than "
35 + "this will be truncated to this length. 0 means no "
36 + "document-level truncation is applied (only truncate_prompt_tokens "
37 + "applies to the combined query+document)."
38 + ),
39 + )
40 +
41 def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
42 encoder_config = model_config.encoder_config or {}
43
@@ -91,29 +109,11 @@ ScoreRequest: TypeAlias = (
109 )
110
111
94 -class RerankRequest(PoolingBasicRequestMixin, ClassifyRequestMixin):
112 +class RerankRequest(ScoreRequestMixin):
113 query: ScoreInput
114 documents: ScoreInput | list[ScoreInput]
115 top_n: int = Field(default_factory=lambda: 0)
116
99 - def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
100 - encoder_config = model_config.encoder_config or {}
101 -
102 - return TokenizeParams(
103 - max_total_tokens=model_config.max_model_len,
104 - max_output_tokens=0,
105 - truncate_prompt_tokens=self.truncate_prompt_tokens,
106 - truncation_side=self.truncation_side,
107 - do_lower_case=encoder_config.get("do_lower_case", False),
108 - max_total_tokens_param="max_model_len",
109 - )
110 -
111 - def to_pooling_params(self, task: PoolingTask = "classify"):
112 - return PoolingParams(
113 - task=task,
114 - use_activation=self.use_activation,
115 - )
116 -
117
118 ScoringRequest: TypeAlias = ScoreRequest | RerankRequest
119
vllm/entrypoints/pooling/scoring/serving.py
+1
@@ -237,6 +237,7 @@ class ServingScores(PoolingServing):
237
238 await self._prepare_generators(query_ctx)
239 await self._collect_batch(query_ctx)
240 + ctx.query_final_res_batch = query_ctx.final_res_batch
241
242 async def _flash_late_interaction_encode_docs(self, ctx: ScoringServeContext):
243 assert ctx.n_queries is not None
vllm/entrypoints/pooling/scoring/utils.py
+31
@@ -25,6 +25,37 @@ from .typing import (
25 )
26
27
28 +def get_num_special_tokens_for_pair(tokenizer) -> int:
29 + """Get number of special tokens added for a text pair encoding."""
30 + method = getattr(tokenizer, "num_special_tokens_to_add", None)
31 + if method is not None:
32 + try:
33 + return method(pair=True)
34 + except TypeError:
35 + pass
36 + # Fallback: compute by tokenizing empty strings
37 + empty_encoding = tokenizer("", text_pair="", add_special_tokens=True)
38 + return len(empty_encoding["input_ids"])
39 +
40 +
41 +def truncate_text_to_tokens(
42 + text: str,
43 + tokenizer,
44 + max_tokens: int,
45 +) -> str:
46 + """Truncate text to a maximum number of content tokens.
47 +
48 + Uses offset_mapping to slice the original text at the exact character
49 + boundary, avoiding lossy encode→decode round-trips that can shift
50 + the token count by 1-3 tokens due to BPE merge boundary changes.
51 + """
52 + encoding = tokenizer(text, add_special_tokens=False, return_offsets_mapping=True)
53 + if len(encoding["input_ids"]) <= max_tokens:
54 + return text
55 + char_end = encoding["offset_mapping"][max_tokens - 1][1]
56 + return text[:char_end]
57 +
58 +
59 def compute_maxsim_score(q_emb: torch.Tensor, d_emb: torch.Tensor) -> torch.Tensor:
60 """
61 Compute ColBERT MaxSim score.
vllm/entrypoints/pooling/typing.py
+3
@@ -89,6 +89,9 @@ class PoolingServeContext(Generic[PoolingRequestT]):
89 ## for IOProcessorResponse
90 response: Any | None = None
91
92 + ## for flash-late-interaction
93 + query_final_res_batch: list[PoolingRequestOutput] | None = None
94 +
95
96 @dataclass
97 class OfflineInputsContext: