Add support for encoder embedding models using MultiModal args

tests/models/language/pooling/test_embedding.py

@@ -38,19 +38,13 @@ def v1(run_with_both_engines):
                      marks=[pytest.mark.skip_v0, pytest.mark.cpu_model]),
         # [Encoder-only]
         pytest.param("BAAI/bge-base-en-v1.5",
-                     marks=[
-                         pytest.mark.core_model, pytest.mark.cpu_model,
-                         pytest.mark.skip_v1
-                     ]),
-        pytest.param("sentence-transformers/all-MiniLM-L12-v2",
-                     marks=[pytest.mark.skip_v1]),
-        pytest.param("intfloat/multilingual-e5-small",
-                     marks=[pytest.mark.skip_v1]),
+                     marks=[pytest.mark.core_model, pytest.mark.cpu_model]),
+        pytest.param("sentence-transformers/all-MiniLM-L12-v2"),
+        pytest.param("intfloat/multilingual-e5-small"),
         pytest.param("Alibaba-NLP/gte-Qwen2-1.5B-instruct",
-                     marks=[pytest.mark.skip_v1]),
+                     marks=[pytest.mark.skip_v0]),
         # [Cross-Encoder]
-        pytest.param("sentence-transformers/stsb-roberta-base-v2",
-                     marks=[pytest.mark.skip_v1]),
+        pytest.param("sentence-transformers/stsb-roberta-base-v2"),
     ],
 )
 def test_models(

tests/models/language/pooling/test_jina.py

@@ -26,6 +26,14 @@
 ]
 
 
+@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
+
+
 @pytest.mark.parametrize("model_info", EMBEDDING_MODELS)
 def test_embed_models_mteb(hf_runner, vllm_runner,
                            model_info: EmbedModelInfo) -> None:

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"
 
 

tests/v1/core/test_kv_cache_utils.py

@@ -916,4 +916,4 @@ def test_get_kv_cache_config():
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())
-        ])
+        ])

vllm/config.py

@@ -716,6 +716,11 @@ def _init_pooler_config(self) -> Optional["PoolerConfig"]:
                 self.override_pooler_config = PoolerConfig(
                     **self.override_pooler_config)
 
+            # WIP: currently cuda graphs are not working for encoder models.
+            logger.warning("CUDA graph is not supported for pooling yet, "
+                           "fallback to the eager mode.")
+            self.enforce_eager = True
+
             pooler_config = self.override_pooler_config or PoolerConfig()
 
             base_config = get_pooling_config(self.model, self.revision)

vllm/engine/arg_utils.py

@@ -1664,7 +1664,8 @@ def _set_default_args_v1(self, usage_context: UsageContext,
 
         if (self.max_num_seqs is None
                 and usage_context in default_max_num_seqs):
-            self.max_num_seqs = default_max_num_seqs[usage_context]
+            self.max_num_seqs = min(default_max_num_seqs[usage_context],
+                                    self.max_num_batched_tokens)
 
             logger.debug("Setting max_num_seqs to %d for %s usage context.",
                          self.max_num_seqs, use_context_value)

vllm/entrypoints/openai/serving_score.py

@@ -7,6 +7,7 @@
 
 from fastapi import Request
 
+import vllm.envs as envs
 from vllm.config import ModelConfig
 from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.logger import RequestLogger
@@ -180,9 +181,17 @@
             input_ids = prompt_inputs["input_ids"]
             text_token_prompt = \
                 self._validate_input(request, input_ids, request_prompt)
+
+            token_type_ids = prompt_inputs.get("token_type_ids")
+            mm_data = None
+            if envs.VLLM_USE_V1 and token_type_ids is not None:
+                mm_data = {"token_type_ids": token_type_ids}
+                token_type_ids = None
+
             engine_prompt = TokensPrompt(
                 prompt_token_ids=text_token_prompt["prompt_token_ids"],
-                token_type_ids=prompt_inputs.get("token_type_ids"))
+                token_type_ids=token_type_ids,
+                multi_modal_data=mm_data)
 
             request_prompts.append(request_prompt)
             engine_prompts.append(engine_prompt)