[Bugfix][Frontend] Truncate pooling prompts before padding them (#54364)
Signed-off-by: hotragn <hotragn.pettugani_2024@woxsen.edu.in>
Hotragn Pettugani committed
Aug 31, 2026 at 01:12 UTC
687db597440bf2b0cc65404e5f8cc238632127e0
4 files changed
+142
-17
tests/entrypoints/pooling/scoring/test_io_processor_unit.py
new
+60
@@ -0,0 +1,60 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for ScoringIOProcessor post-tokenization helpers."""
4
+
5
+from dataclasses import dataclass
6
+
7
+import pytest
8
+
9
+from vllm import TokensPrompt
10
+from vllm.entrypoints.pooling.scoring.io_processor import (
11
+ _apply_post_tokenization_to_token_type_ids,
12
+)
13
+from vllm.entrypoints.pooling.scoring.utils import compress_token_type_ids
14
+from vllm.renderers import TokenizeParams
15
+
16
+pytestmark = pytest.mark.skip_global_cleanup
17
+
18
+
19
+@dataclass
20
+class _DummyTokenizer:
21
+ truncation_side: str = "left"
22
+ # Outside the range of the prompt ids below, so a test can tell a pad
23
+ # token apart from a real one.
24
+ pad_token_id: int = 99999
25
+
26
+
27
+def test_token_type_ids_stay_aligned_with_a_truncated_padded_prompt():
28
+ """The cross-encoder segment boundary must survive truncate + pad.
29
+
30
+ `token_type_ids` are parallel to `prompt_token_ids` and are reduced to a
31
+ single boundary index by `compress_token_type_ids`. If the two arrays are
32
+ truncated and padded in different orders they no longer describe the same
33
+ positions, and the model is told the query segment is empty.
34
+ """
35
+ tokenizer = _DummyTokenizer()
36
+ num_query, num_doc = 20, 30
37
+ prompt = TokensPrompt(prompt_token_ids=list(range(num_query + num_doc)))
38
+ token_type_ids = [0] * num_query + [1] * num_doc
39
+
40
+ tok_params = TokenizeParams(
41
+ max_total_tokens=100,
42
+ pad_prompt_tokens=-1,
43
+ truncate_prompt_tokens=40,
44
+ truncation_side="left",
45
+ )
46
+
47
+ prompt = tok_params.apply_post_tokenization(tokenizer, prompt)
48
+ token_type_ids = _apply_post_tokenization_to_token_type_ids(
49
+ tokenizer, tok_params, token_type_ids
50
+ )
51
+
52
+ prompt_token_ids = prompt["prompt_token_ids"]
53
+ assert len(token_type_ids) == len(prompt_token_ids)
54
+
55
+ # Keeping the last 40 tokens drops the first 10 query tokens, so 10 query
56
+ # tokens survive and the document starts at index 10.
57
+ first_doc = compress_token_type_ids(token_type_ids)
58
+ assert first_doc == 10
59
+ assert prompt_token_ids[:first_doc] == list(range(10, num_query))
60
+ assert prompt_token_ids[first_doc:40] == list(range(num_query, 50))
tests/renderers/test_completions.py
+55
@@ -62,6 +62,9 @@ class MockVllmConfig:
62
class DummyTokenizer:
63
truncation_side: str = "left"
64
max_chars_per_token: int = 1
65
+ # Deliberately outside the range of ids `encode` returns, so a test can
66
+ # tell a pad token apart from a real one.
67
+ pad_token_id: int = 99999
68
69
def __post_init__(self) -> None:
70
self._captured_encode_kwargs: dict = {}
@@ -428,6 +431,58 @@ class TestRenderPrompt:
431
assert len(results[0]["prompt_token_ids"]) == 5
432
assert results[0]["prompt_token_ids"] == list(range(5))
433
434
+ def test_padding_with_left_truncation_keeps_the_prompt(self):
435
+ """Padding must not be truncated away.
436
+
437
+ `padding` and `truncate_prompt_tokens`/`truncation_side` are all
438
+ settable on one pooling request. Padding to the full input length
439
+ before truncating from the left leaves a prompt made entirely of pad
440
+ tokens, and the request still succeeds -- so the model embeds nothing
441
+ but padding.
442
+ """
443
+ renderer = _build_renderer(MockModelConfig())
444
+ pad_id = renderer.tokenizer.pad_token_id
445
+
446
+ prompts = renderer.render_prompts(
447
+ _preprocess_prompt(renderer.model_config, "x" * 50)
448
+ )
449
+ results = renderer.tokenize_prompts(
450
+ prompts,
451
+ TokenizeParams(
452
+ max_total_tokens=100,
453
+ pad_prompt_tokens=-1,
454
+ truncate_prompt_tokens=5,
455
+ truncation_side="left",
456
+ ),
457
+ )
458
+
459
+ assert len(results) == 1
460
+ token_ids = results[0]["prompt_token_ids"]
461
+
462
+ # The sentinel: on a padding-first pipeline every surviving id is a
463
+ # pad token, so the prompt reaches the model with no content at all.
464
+ assert set(token_ids) != {pad_id}
465
+
466
+ # Transformers semantics: truncate to 5, then pad out to the full
467
+ # input length.
468
+ assert token_ids[:5] == list(range(45, 50))
469
+ assert token_ids[5:] == [pad_id] * 95
470
+
471
+ def test_padding_without_truncation_is_unchanged(self):
472
+ renderer = _build_renderer(MockModelConfig())
473
+ pad_id = renderer.tokenizer.pad_token_id
474
+
475
+ prompts = renderer.render_prompts(
476
+ _preprocess_prompt(renderer.model_config, "x" * 50)
477
+ )
478
+ results = renderer.tokenize_prompts(
479
+ prompts,
480
+ TokenizeParams(max_total_tokens=100, pad_prompt_tokens=-1),
481
+ )
482
+
483
+ assert len(results) == 1
484
+ assert results[0]["prompt_token_ids"] == list(range(50)) + [pad_id] * 50
485
+
486
def test_explicit_side_text_pretokenization_guard(self):
487
renderer = _build_renderer(MockModelConfig(), max_chars_per_token=1)
488
vllm/entrypoints/pooling/scoring/io_processor.py
+22
-16
@@ -55,6 +55,27 @@ def _apply_post_tokenization_to_token_type_ids(
55
tok_params: TokenizeParams,
56
token_type_ids: list[int],
57
) -> list[int]:
58
+ # Must stay in the same order as `TokenizeParams._validate_tokens`, which
59
+ # truncates before padding. These ids are parallel to the prompt tokens
60
+ # and locate the query/document boundary for the cross-encoder; applying
61
+ # the two steps in a different order silently misplaces that boundary.
62
+ max_length = tok_params.truncate_prompt_tokens
63
+ if max_length is not None and max_length < 0:
64
+ max_length = tok_params.max_input_tokens
65
+
66
+ if max_length is not None and max_length < len(token_type_ids):
67
+ if max_length == 0:
68
+ token_type_ids = token_type_ids[:0]
69
+ else:
70
+ side = tok_params.truncation_side or (
71
+ tokenizer.truncation_side if tokenizer is not None else None
72
+ )
73
+ token_type_ids = (
74
+ token_type_ids[-max_length:]
75
+ if side == "left"
76
+ else token_type_ids[:max_length]
77
+ )
78
+
79
pad_length = tok_params.pad_prompt_tokens
80
if pad_length is not None and pad_length < 0:
81
pad_length = tok_params.max_input_tokens
@@ -65,22 +86,7 @@ def _apply_post_tokenization_to_token_type_ids(
86
pad_length - len(token_type_ids)
87
)
88
68
- max_length = tok_params.truncate_prompt_tokens
69
- if max_length is not None and max_length < 0:
70
- max_length = tok_params.max_input_tokens
71
-
72
- if max_length is None or max_length >= len(token_type_ids):
73
- return token_type_ids
74
- if max_length == 0:
75
- return token_type_ids[:0]
76
-
77
- side = tok_params.truncation_side or (
78
- tokenizer.truncation_side if tokenizer is not None else None
79
- )
80
- if side == "left":
81
- return token_type_ids[-max_length:]
82
-
83
- return token_type_ids[:max_length]
89
+ return token_type_ids
90
91
92
class ScoringIOProcessor(PoolingIOProcessor):
vllm/renderers/params.py
+5
-1
@@ -462,9 +462,13 @@ class TokenizeParams:
462
463
def _validate_tokens(self, tokenizer: TokenizerLike | None, tokens: _S) -> _S:
464
"""Apply all validators to a token sequence."""
465
+ # Truncation runs before padding, matching the Transformers pipeline
466
+ # these parameters are named after. Padding first would let a
467
+ # subsequent left-side truncation keep only the pad tokens it just
468
+ # appended, discarding the prompt entirely.
469
for validator in (
466
- self._token_padding,
470
self._token_truncation,
471
+ self._token_padding,
472
self._token_len_check,
473
):
474
tokens = validator(tokenizer, tokens)