[Bugfix][Pooling] Align CrossEncoder token type ids after truncation (#47772)
Signed-off-by: Ting Sun <suntcrick@gmail.com>
Ting SUN committed
Jul 8, 2026 at 10:59 UTC
0303f37a54dbf0c81e375e654d431a9d33a855a9
2 files changed
+85
-3
tests/entrypoints/pooling/scoring/test_cross_encoder_offline.py
+43
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import weakref
5
+from types import SimpleNamespace
6
7
import pytest
8
import torch
@@ -9,7 +10,10 @@ import torch
10
from tests.models.utils import softmax
11
from vllm import LLM, PoolingParams
12
from vllm.distributed import cleanup_dist_env_and_memory
13
+from vllm.entrypoints.pooling.scoring.io_processor import CrossEncoderIOProcessor
14
+from vllm.entrypoints.pooling.scoring.typing import ScoringData
15
from vllm.platforms import current_platform
16
+from vllm.renderers import TokenizeParams
17
18
MODEL_NAME = "tomaarsen/Qwen3-Reranker-0.6B-seq-cls"
19
PROMPT = "The chef prepared a delicious meal."
@@ -141,6 +145,45 @@ def test_max_tokens_per_doc(llm: LLM):
145
assert with_limit_tokens < no_limit_tokens
146
147
148
+def test_token_type_ids_follow_post_tokenization():
149
+ processor = object.__new__(CrossEncoderIOProcessor)
150
+ processor.tokenizer = SimpleNamespace(truncation_side="right", pad_token_id=-1)
151
+ processor.renderer = SimpleNamespace(process_for_engine=lambda prompt, _: prompt)
152
+ processor.model_config = None
153
+ processor.get_score_prompt = lambda **_: (
154
+ "",
155
+ {
156
+ "prompt_token_ids": list(range(32)),
157
+ "token_type_ids": [0] * 16 + [1] * 16,
158
+ },
159
+ )
160
+
161
+ engine_inputs, pooling_params = processor._pre_process(
162
+ ScoringData(data_1=["query"], data_2=["document"]),
163
+ TokenizeParams(
164
+ max_total_tokens=None,
165
+ truncate_prompt_tokens=16,
166
+ truncation_side="left",
167
+ ),
168
+ PoolingParams(task="classify", extra_kwargs={"cache_salt": "salt"}),
169
+ )
170
+
171
+ assert engine_inputs[0]["prompt_token_ids"] == list(range(16, 32))
172
+ assert pooling_params[0].extra_kwargs == {
173
+ "cache_salt": "salt",
174
+ "compressed_token_type_ids": 0,
175
+ }
176
+
177
+ engine_inputs, pooling_params = processor._pre_process(
178
+ ScoringData(data_1=["query"], data_2=["document"]),
179
+ TokenizeParams(max_total_tokens=None, pad_prompt_tokens=40),
180
+ PoolingParams(task="classify"),
181
+ )
182
+
183
+ assert engine_inputs[0]["prompt_token_ids"] == list(range(32)) + [-1] * 8
184
+ assert pooling_params[0].extra_kwargs == {"compressed_token_type_ids": 16}
185
+
186
+
187
def test_pooling_params(llm: LLM):
188
def get_outputs(use_activation):
189
outputs = llm.score(
vllm/entrypoints/pooling/scoring/io_processor.py
+42
-3
@@ -36,6 +36,39 @@ from .utils import (
36
ScoringServeContext: TypeAlias = PoolingServeContext[ScoringRequest]
37
38
39
+def _apply_post_tokenization_to_token_type_ids(
40
+ tokenizer: Any,
41
+ tok_params: TokenizeParams,
42
+ token_type_ids: list[int],
43
+) -> list[int]:
44
+ pad_length = tok_params.pad_prompt_tokens
45
+ if pad_length is not None and pad_length < 0:
46
+ pad_length = tok_params.max_input_tokens
47
+
48
+ if pad_length is not None and pad_length > len(token_type_ids):
49
+ pad_token_type_id = token_type_ids[-1] if token_type_ids else 0
50
+ token_type_ids = token_type_ids + [pad_token_type_id] * (
51
+ pad_length - len(token_type_ids)
52
+ )
53
+
54
+ max_length = tok_params.truncate_prompt_tokens
55
+ if max_length is not None and max_length < 0:
56
+ max_length = tok_params.max_input_tokens
57
+
58
+ if max_length is None or max_length >= len(token_type_ids):
59
+ return token_type_ids
60
+ if max_length == 0:
61
+ return token_type_ids[:0]
62
+
63
+ side = tok_params.truncation_side or (
64
+ tokenizer.truncation_side if tokenizer is not None else None
65
+ )
66
+ if side == "left":
67
+ return token_type_ids[-max_length:]
68
+
69
+ return token_type_ids[:max_length]
70
+
71
+
72
class ScoringIOProcessor(PoolingIOProcessor):
73
name: str
74
pooling_task: PoolingTask
@@ -469,9 +502,16 @@ class CrossEncoderIOProcessor(ScoringIOProcessor):
502
else None,
503
)
504
472
- if token_type_ids := engine_prompt.pop("token_type_ids", None):
505
+ token_type_ids = engine_prompt.pop("token_type_ids", None)
506
+ tok_params.apply_post_tokenization(self.tokenizer, engine_prompt)
507
+
508
+ if token_type_ids is not None:
509
params = pooling_params.clone()
474
- compressed = compress_token_type_ids(token_type_ids)
510
+ compressed = compress_token_type_ids(
511
+ _apply_post_tokenization_to_token_type_ids(
512
+ self.tokenizer, tok_params, token_type_ids
513
+ )
514
+ )
515
params.extra_kwargs = {
516
**(params.extra_kwargs or {}),
517
"compressed_token_type_ids": compressed,
@@ -480,7 +520,6 @@ class CrossEncoderIOProcessor(ScoringIOProcessor):
520
else:
521
pooling_params_list.append(pooling_params)
522
483
- tok_params.apply_post_tokenization(self.tokenizer, engine_prompt)
523
if engine_prompt_extras:
524
target_prompt = extract_target_prompt(self.model_config, engine_prompt)
525
target_prompt.update(engine_prompt_extras)