[Bugfix][Frontend] Truncate pooling prompts before padding them

tests/entrypoints/pooling/scoring/test_io_processor_unit.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for ScoringIOProcessor post-tokenization helpers."""
+
+from dataclasses import dataclass
+
+import pytest
+
+from vllm import TokensPrompt
+from vllm.entrypoints.pooling.scoring.io_processor import (
+    _apply_post_tokenization_to_token_type_ids,
+)
+from vllm.entrypoints.pooling.scoring.utils import compress_token_type_ids
+from vllm.renderers import TokenizeParams
+
+pytestmark = pytest.mark.skip_global_cleanup
+
+
+@dataclass
+class _DummyTokenizer:
+    truncation_side: str = "left"
+    # Outside the range of the prompt ids below, so a test can tell a pad
+    # token apart from a real one.
+    pad_token_id: int = 99999
+
+
+def test_token_type_ids_stay_aligned_with_a_truncated_padded_prompt():
+    """The cross-encoder segment boundary must survive truncate + pad.
+
+    `token_type_ids` are parallel to `prompt_token_ids` and are reduced to a
+    single boundary index by `compress_token_type_ids`. If the two arrays are
+    truncated and padded in different orders they no longer describe the same
+    positions, and the model is told the query segment is empty.
+    """
+    tokenizer = _DummyTokenizer()
+    num_query, num_doc = 20, 30
+    prompt = TokensPrompt(prompt_token_ids=list(range(num_query + num_doc)))
+    token_type_ids = [0] * num_query + [1] * num_doc
+
+    tok_params = TokenizeParams(
+        max_total_tokens=100,
+        pad_prompt_tokens=-1,
+        truncate_prompt_tokens=40,
+        truncation_side="left",
+    )
+
+    prompt = tok_params.apply_post_tokenization(tokenizer, prompt)
+    token_type_ids = _apply_post_tokenization_to_token_type_ids(
+        tokenizer, tok_params, token_type_ids
+    )
+
+    prompt_token_ids = prompt["prompt_token_ids"]
+    assert len(token_type_ids) == len(prompt_token_ids)
+
+    # Keeping the last 40 tokens drops the first 10 query tokens, so 10 query
+    # tokens survive and the document starts at index 10.
+    first_doc = compress_token_type_ids(token_type_ids)
+    assert first_doc == 10
+    assert prompt_token_ids[:first_doc] == list(range(10, num_query))
+    assert prompt_token_ids[first_doc:40] == list(range(num_query, 50))

tests/renderers/test_completions.py

@@ -62,6 +62,9 @@ class MockVllmConfig:
 class DummyTokenizer:
     truncation_side: str = "left"
     max_chars_per_token: int = 1
+    # Deliberately outside the range of ids `encode` returns, so a test can
+    # tell a pad token apart from a real one.
+    pad_token_id: int = 99999
 
     def __post_init__(self) -> None:
         self._captured_encode_kwargs: dict = {}
@@ -428,6 +431,58 @@ def test_explicit_side_right_text(self):
         assert len(results[0]["prompt_token_ids"]) == 5
         assert results[0]["prompt_token_ids"] == list(range(5))
 
+    def test_padding_with_left_truncation_keeps_the_prompt(self):
+        """Padding must not be truncated away.
+
+        `padding` and `truncate_prompt_tokens`/`truncation_side` are all
+        settable on one pooling request. Padding to the full input length
+        before truncating from the left leaves a prompt made entirely of pad
+        tokens, and the request still succeeds -- so the model embeds nothing
+        but padding.
+        """
+        renderer = _build_renderer(MockModelConfig())
+        pad_id = renderer.tokenizer.pad_token_id
+
+        prompts = renderer.render_prompts(
+            _preprocess_prompt(renderer.model_config, "x" * 50)
+        )
+        results = renderer.tokenize_prompts(
+            prompts,
+            TokenizeParams(
+                max_total_tokens=100,
+                pad_prompt_tokens=-1,
+                truncate_prompt_tokens=5,
+                truncation_side="left",
+            ),
+        )
+
+        assert len(results) == 1
+        token_ids = results[0]["prompt_token_ids"]
+
+        # The sentinel: on a padding-first pipeline every surviving id is a
+        # pad token, so the prompt reaches the model with no content at all.
+        assert set(token_ids) != {pad_id}
+
+        # Transformers semantics: truncate to 5, then pad out to the full
+        # input length.
+        assert token_ids[:5] == list(range(45, 50))
+        assert token_ids[5:] == [pad_id] * 95
+
+    def test_padding_without_truncation_is_unchanged(self):
+        renderer = _build_renderer(MockModelConfig())
+        pad_id = renderer.tokenizer.pad_token_id
+
+        prompts = renderer.render_prompts(
+            _preprocess_prompt(renderer.model_config, "x" * 50)
+        )
+        results = renderer.tokenize_prompts(
+            prompts,
+            TokenizeParams(max_total_tokens=100, pad_prompt_tokens=-1),
+        )
+
+        assert len(results) == 1
+        assert results[0]["prompt_token_ids"] == list(range(50)) + [pad_id] * 50
+
     def test_explicit_side_text_pretokenization_guard(self):
         renderer = _build_renderer(MockModelConfig(), max_chars_per_token=1)
 

vllm/entrypoints/pooling/scoring/io_processor.py

@@ -55,6 +55,27 @@ def _apply_post_tokenization_to_token_type_ids(
     tok_params: TokenizeParams,
     token_type_ids: list[int],
 ) -> list[int]:
+    # Must stay in the same order as `TokenizeParams._validate_tokens`, which
+    # truncates before padding. These ids are parallel to the prompt tokens
+    # and locate the query/document boundary for the cross-encoder; applying
+    # the two steps in a different order silently misplaces that boundary.
+    max_length = tok_params.truncate_prompt_tokens
+    if max_length is not None and max_length < 0:
+        max_length = tok_params.max_input_tokens
+
+    if max_length is not None and max_length < len(token_type_ids):
+        if max_length == 0:
+            token_type_ids = token_type_ids[:0]
+        else:
+            side = tok_params.truncation_side or (
+                tokenizer.truncation_side if tokenizer is not None else None
+            )
+            token_type_ids = (
+                token_type_ids[-max_length:]
+                if side == "left"
+                else token_type_ids[:max_length]
+            )
+
     pad_length = tok_params.pad_prompt_tokens
     if pad_length is not None and pad_length < 0:
         pad_length = tok_params.max_input_tokens
@@ -65,22 +86,7 @@ def _apply_post_tokenization_to_token_type_ids(
             pad_length - len(token_type_ids)
         )
 
-    max_length = tok_params.truncate_prompt_tokens
-    if max_length is not None and max_length < 0:
-        max_length = tok_params.max_input_tokens
-
-    if max_length is None or max_length >= len(token_type_ids):
-        return token_type_ids
-    if max_length == 0:
-        return token_type_ids[:0]
-
-    side = tok_params.truncation_side or (
-        tokenizer.truncation_side if tokenizer is not None else None
-    )
-    if side == "left":
-        return token_type_ids[-max_length:]
-
-    return token_type_ids[:max_length]
+    return token_type_ids
 
 
 class ScoringIOProcessor(PoolingIOProcessor):

vllm/renderers/params.py

@@ -462,9 +462,13 @@ def _token_len_check(self, tokenizer: TokenizerLike | None, tokens: _S) -> _S:
 
     def _validate_tokens(self, tokenizer: TokenizerLike | None, tokens: _S) -> _S:
         """Apply all validators to a token sequence."""
+        # Truncation runs before padding, matching the Transformers pipeline
+        # these parameters are named after. Padding first would let a
+        # subsequent left-side truncation keep only the pad tokens it just
+        # appended, discarding the prompt entirely.
         for validator in (
-            self._token_padding,
             self._token_truncation,
+            self._token_padding,
             self._token_len_check,
         ):
             tokens = validator(tokenizer, tokens)