Support token_type_ids in V1 with less code changes (vllm-project#21985)

Signed-off-by: Max de Bayser <[email protected]>

Author: maxdebayser

tests/entrypoints/openai/test_rerank.py

@@ -126,7 +126,9 @@ def test_invocations(server: RemoteOpenAIServer):
                                                  invocation_output["results"]):
         assert rerank_result.keys() == invocations_result.keys()
         assert rerank_result["relevance_score"] == pytest.approx(
-            invocations_result["relevance_score"], rel=0.01)
+            invocations_result["relevance_score"], rel=0.05)
+        # TODO: reset this tolerance to 0.01 once we find
+        # an alternative to flash_attn with bfloat16
 
 
 @pytest.mark.asyncio

tests/entrypoints/openai/test_score.py

@@ -220,7 +220,9 @@ def test_invocations(self, server: RemoteOpenAIServer, model: dict[str,
                                                invocation_output["data"]):
             assert score_data.keys() == invocation_data.keys()
             assert score_data["score"] == pytest.approx(
-                invocation_data["score"], rel=0.01)
+                invocation_data["score"], rel=0.05)
+            # TODO: reset this tolerance to 0.01 once we find
+            # an alternative to flash_attn with bfloat16
 
     def test_activation(self, server: RemoteOpenAIServer, model: dict[str,
                                                                       Any]):

tests/models/language/pooling/test_scoring.py

@@ -23,6 +23,15 @@
     "The capital of Germany is Berlin.",
 ]
 
+
+@pytest.fixture(autouse=True)
+def v1(run_with_both_engines):
+    # Simple autouse wrapper to run both engines for each test
+    # This can be promoted up to conftest.py to run for every
+    # test in a package
+    pass
+
+
 DTYPE = "half"
 
 

vllm/entrypoints/llm.py

@@ -28,11 +28,15 @@
                                          apply_mistral_chat_template,
                                          parse_chat_messages,
                                          resolve_chat_template_content_format)
+# yapf conflicts with isort for this block
+# yapf: disable
 from vllm.entrypoints.score_utils import (ScoreContentPartParam,
                                           ScoreMultiModalParam,
                                           _cosine_similarity,
                                           _validate_score_input_lens,
+                                          compress_token_type_ids,
                                           get_score_prompt)
+# yapf: enable
 from vllm.entrypoints.utils import (_validate_truncation_size,
                                     log_non_default_args)
 from vllm.inputs import PromptType, SingletonPrompt, TextPrompt, TokensPrompt
@@ -1329,6 +1333,7 @@ def _cross_encoding_score(
 
         model_config = self.llm_engine.model_config
         pooling_params.verify("score", model_config)
+        pooling_params_list = list[PoolingParams]()
 
         tokenization_kwargs: dict[str, Any] = {}
 
@@ -1339,38 +1344,31 @@ def _cross_encoding_score(
 
         input_pairs = [(t1, t2) for t1, t2 in zip(data_1, data_2)]
 
-        if model_config.is_multimodal_model:
-            for q, d in input_pairs:
-                _, engine_prompt = get_score_prompt(
-                    model_config=model_config,
-                    data_1=q,
-                    data_2=d,
-                    tokenizer=tokenizer,
-                    tokenization_kwargs=tokenization_kwargs,
-                )
+        model_config = self.llm_engine.model_config
 
-                parsed_prompts.append(engine_prompt)
-        else:
-            for q, t in input_pairs:
-                if model_config.use_pad_token:
-                    # cross_encoder models defaults to using pad_token.
-                    prompt_inputs = tokenizer(
-                        text=q,  # type: ignore[arg-type]
-                        text_pair=t,  # type: ignore[arg-type]
-                        **tokenization_kwargs)
-                else:
-                    # `llm as reranker` models defaults to not using pad_token.
-                    prompt_inputs = tokenizer(
-                        text=q + t,  # type: ignore[operator]
-                        **tokenization_kwargs)
-                engine_prompt = TokensPrompt(
-                    prompt_token_ids=prompt_inputs["input_ids"],
-                    token_type_ids=prompt_inputs.get("token_type_ids"))
-                parsed_prompts.append(engine_prompt)
+        for q, d in input_pairs:
+            _, engine_prompt = get_score_prompt(
+                model_config=model_config,
+                data_1=q,
+                data_2=d,
+                tokenizer=tokenizer,
+                tokenization_kwargs=tokenization_kwargs,
+            )
+
+            if envs.VLLM_USE_V1 and (token_type_ids := engine_prompt.pop(
+                    "token_type_ids", None)):
+                params = pooling_params.clone()
+                compressed = compress_token_type_ids(token_type_ids)
+                params.extra_kwargs = {"compressed_token_type_ids": compressed}
+                pooling_params_list.append(params)
+            else:
+                pooling_params_list.append(pooling_params)
+
+            parsed_prompts.append(engine_prompt)
 
         self._validate_and_add_requests(
             prompts=parsed_prompts,
-            params=pooling_params,
+            params=pooling_params_list,
             use_tqdm=use_tqdm,
             lora_request=lora_request,
         )

vllm/entrypoints/openai/serving_score.py

@@ -7,6 +7,7 @@
 
 from fastapi import Request
 
+from vllm import envs
 from vllm.config import ModelConfig
 from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.logger import RequestLogger
@@ -17,11 +18,15 @@
                                               ScoreResponseData, UsageInfo)
 from vllm.entrypoints.openai.serving_engine import OpenAIServing
 from vllm.entrypoints.openai.serving_models import OpenAIServingModels
+# yapf conflicts with isort for this block
+# yapf: disable
 from vllm.entrypoints.score_utils import (ScoreContentPartParam,
                                           ScoreMultiModalParam,
                                           _cosine_similarity,
                                           _validate_score_input_lens,
+                                          compress_token_type_ids,
                                           get_score_prompt)
+# yapf: enable
 from vllm.entrypoints.utils import _validate_truncation_size
 from vllm.inputs.data import TokensPrompt
 from vllm.logger import init_logger
@@ -158,6 +163,8 @@ def _preprocess_score(
             tokenizer=tokenizer,
             tokenization_kwargs=tokenization_kwargs,
         )
+        self._validate_input(request, engine_prompt["prompt_token_ids"],
+                             full_prompt)
         if request.mm_processor_kwargs is not None:
             engine_prompt["mm_processor_kwargs"] = request.mm_processor_kwargs
 
@@ -188,64 +195,27 @@ async def _cross_encoding_score(
 
         input_pairs = [(t1, t2) for t1, t2 in zip(data_1, data_2)]
 
-        if self.model_config.is_multimodal_model:
+        preprocess_async = make_async(self._preprocess_score,
+                                      executor=self._tokenizer_executor)
 
-            preprocess_async = make_async(self._preprocess_score,
-                                          executor=self._tokenizer_executor)
+        preprocessed_prompts = await asyncio.gather(
+            *(preprocess_async(request=request,
+                               tokenizer=tokenizer,
+                               tokenization_kwargs=tokenization_kwargs,
+                               data_1=t1,
+                               data_2=t2) for t1, t2 in input_pairs))
 
-            preprocessed_prompts = await asyncio.gather(
-                *(preprocess_async(request=request,
-                                   tokenizer=tokenizer,
-                                   tokenization_kwargs=tokenization_kwargs,
-                                   data_1=t1,
-                                   data_2=t2) for t1, t2 in input_pairs))
-
-            for full_prompt, engine_prompt in preprocessed_prompts:
-                request_prompts.append(full_prompt)
-                engine_prompts.append(engine_prompt)
-
-        else:
-            tokenize_async = make_async(tokenizer.__call__,
-                                        executor=self._tokenizer_executor)
-            use_pad_token = self.model_config.use_pad_token
-
-            if use_pad_token:
-                # cross_encoder models defaults to using pad_token.
-                tokenized_prompts = await asyncio.gather(*(
-                    tokenize_async(
-                        text=t1,  # type: ignore[arg-type]
-                        text_pair=t2,  # type: ignore[arg-type]
-                        **tokenization_kwargs) for t1, t2 in input_pairs))
-            else:
-                # `llm as reranker` models defaults to not using pad_token.
-                tokenized_prompts = await asyncio.gather(*(
-                    tokenize_async(
-                        text=t1 +  # type: ignore[operator]
-                        t2,
-                        **tokenization_kwargs) for t1, t2 in input_pairs))
-
-            for prompt_inputs, (t1, t2) in zip(tokenized_prompts, input_pairs):
-                sep_token = tokenizer.sep_token if (tokenizer.sep_token
-                                                    and use_pad_token) else ''
-                request_prompt = f"{t1}{sep_token}{t2}"
-
-                input_ids = prompt_inputs["input_ids"]
-                text_token_prompt = \
-                    self._validate_input(request, input_ids, request_prompt)
-                engine_prompt = TokensPrompt(
-                    prompt_token_ids=text_token_prompt["prompt_token_ids"],
-                    token_type_ids=prompt_inputs.get("token_type_ids"))
-
-                request_prompts.append(request_prompt)
-                engine_prompts.append(engine_prompt)
+        for full_prompt, engine_prompt in preprocessed_prompts:
+            request_prompts.append(full_prompt)
+            engine_prompts.append(engine_prompt)
 
         # Schedule the request and get the result generator.
         generators: list[AsyncGenerator[PoolingRequestOutput, None]] = []
 
-        pooling_params = request.to_pooling_params()
+        default_pooling_params = request.to_pooling_params()
 
         try:
-            pooling_params.verify("score", self.model_config)
+            default_pooling_params.verify("score", self.model_config)
         except ValueError as e:
             return self.create_error_response(str(e))
 
@@ -254,9 +224,19 @@ async def _cross_encoding_score(
 
             self._log_inputs(request_id_item,
                              request_prompts[i],
-                             params=pooling_params,
+                             params=default_pooling_params,
                              lora_request=lora_request)
 
+            if envs.VLLM_USE_V1 and (token_type_ids := engine_prompt.pop(
+                    "token_type_ids", None)):
+                pooling_params = default_pooling_params.clone()
+                compressed = compress_token_type_ids(token_type_ids)
+                pooling_params.extra_kwargs = {
+                    "compressed_token_type_ids": compressed
+                }
+            else:
+                pooling_params = (default_pooling_params)
+
             generator = self.engine_client.encode(
                 engine_prompt,
                 pooling_params,

vllm/entrypoints/score_utils.py

@@ -184,15 +184,49 @@ def get_score_prompt(
         model_config,
         tokenizer,
     )
+    from vllm.model_executor.model_loader import get_model_cls
 
-    full_prompt = apply_score_template(model_config, prompt_1, prompt_2)
-
-    prompt_inputs = tokenizer(full_prompt, **tokenization_kwargs)
+    model = get_model_cls(model_config)
+    if supports_score_template(model):
+        full_prompt = apply_score_template(model_config, prompt_1, prompt_2)
+        prompt_inputs = tokenizer(full_prompt, **tokenization_kwargs)
+    elif model_config.use_pad_token:
+        # cross_encoder models defaults to using pad_token.
+        prompt_inputs = tokenizer(text=prompt_1,
+                                  text_pair=prompt_2,
+                                  **tokenization_kwargs)
+        full_prompt = tokenizer.decode(prompt_inputs["input_ids"])
+    else:
+        # `llm as reranker` models defaults to not using pad_token.
+        full_prompt = prompt_1 + prompt_2
+        prompt_inputs = tokenizer(text=full_prompt, **tokenization_kwargs)
 
     engine_prompt = TokensPrompt(prompt_token_ids=prompt_inputs["input_ids"])
 
+    if (token_type_ids := prompt_inputs.get("token_type_ids")) is not None:
+        engine_prompt["token_type_ids"] = token_type_ids
+
     post_process_tokens(model_config, engine_prompt)
 
     if mm_data is not None:
         engine_prompt["multi_modal_data"] = mm_data
     return full_prompt, engine_prompt
+
+
+def compress_token_type_ids(token_type_ids: list[int]) -> int:
+    """
+    Return position of the first 1 or the length of the list
+    if not found.
+    """
+    first_one = len(token_type_ids)
+    err_msg = "Token type ids are expected to be a sequence"\
+              " of zeros followed by a sequence of ones"
+    for i, type_id in enumerate(token_type_ids):
+        if type_id == 0 and first_one < i:
+            raise ValueError(err_msg)
+        elif type_id == 1 and first_one > i:
+            first_one = i
+        elif type_id > 1:
+            raise ValueError(err_msg)
+
+    return first_one