[Model Runner V2] Enable encoder token classification (#50293)

Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>

tests/models/language/pooling/test_splade_sparse_pooler.py

@@ -159,6 +159,7 @@ def test_pooling_runner_rejects_unsupported_selected_task() -> None:
     }
     vllm_config = MagicMock()
     vllm_config.scheduler_config.max_num_seqs = 2
+    vllm_config.model_config.attn_type = "encoder_only"
     vllm_config.model_config.get_pooling_task.return_value = "embed&token_classify"
 
     with (
@@ -166,3 +167,43 @@ def test_pooling_runner_rejects_unsupported_selected_task() -> None:
         pytest.raises(ValueError, match="selects 'embed&token_classify'"),
     ):
         PoolingRunner(model, vllm_config)
+
+
+def test_pooling_runner_supports_encoder_token_classification() -> None:
+    model = MagicMock()
+    model.pooler.get_supported_tasks.return_value = {"token_classify"}
+    vllm_config = MagicMock()
+    vllm_config.scheduler_config.max_num_seqs = 2
+    vllm_config.model_config.attn_type = "encoder_only"
+    vllm_config.model_config.get_pooling_task.return_value = "token_classify"
+
+    runner = PoolingRunner(model, vllm_config)
+
+    assert runner.supported_tasks == {"token_classify"}
+
+
+def test_pooling_runner_rejects_decoder_token_classification() -> None:
+    model = MagicMock()
+    model.pooler.get_supported_tasks.return_value = {"token_classify"}
+    vllm_config = MagicMock()
+    vllm_config.scheduler_config.max_num_seqs = 2
+    vllm_config.model_config.attn_type = "decoder"
+    vllm_config.model_config.get_pooling_task.return_value = "token_classify"
+
+    with pytest.raises(ValueError, match="selects 'token_classify'") as exc_info:
+        PoolingRunner(model, vllm_config)
+
+    assert "Set an explicitly supported task" not in str(exc_info.value)
+
+
+def test_pooling_runner_filters_decoder_token_classification() -> None:
+    model = MagicMock()
+    model.pooler.get_supported_tasks.return_value = {"embed", "token_classify"}
+    vllm_config = MagicMock()
+    vllm_config.scheduler_config.max_num_seqs = 2
+    vllm_config.model_config.attn_type = "decoder"
+    vllm_config.model_config.get_pooling_task.return_value = "embed"
+
+    runner = PoolingRunner(model, vllm_config)
+
+    assert runner.supported_tasks == {"embed"}

tests/models/language/pooling/test_token_classification.py

@@ -29,16 +29,24 @@ def seed_everything():
 )
 # The float32 is required for this tiny model to pass the test.
 @pytest.mark.parametrize("dtype", ["float"])
+@pytest.mark.core_model
 @torch.inference_mode
-def test_bert_models(
+def test_bert_model_runner_v2(
     hf_runner,
     vllm_runner,
     example_prompts,
+    monkeypatch,
     model: str,
     dtype: str,
 ) -> None:
+    prompt_batches = [[example_prompts[0]], example_prompts]
+
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
     with vllm_runner(model, max_model_len=None, dtype=dtype) as vllm_model:
-        vllm_outputs = vllm_model.token_classify(example_prompts)
+        assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner
+        vllm_output_batches = [
+            vllm_model.token_classify(prompts) for prompts in prompt_batches
+        ]
 
     # Use eager attention on ROCm to avoid HF Transformers flash attention
     # accuracy issues: https://github.com/vllm-project/vllm/issues/30167
@@ -59,12 +67,14 @@ def test_bert_models(
             inputs = hf_model.wrap_device(inputs)
             output = hf_model.model(**inputs)
             hf_outputs.append(softmax(output.logits[0]))
+        hf_output_batches = [[hf_outputs[0]], hf_outputs]
 
     # check logits difference
-    for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
-        hf_output = hf_output.detach().clone().cpu().float()
-        vllm_output = vllm_output.detach().clone().cpu().float()
-        torch.testing.assert_close(hf_output, vllm_output, atol=3.2e-2, rtol=1e-3)
+    for hf_outputs, vllm_outputs in zip(hf_output_batches, vllm_output_batches):
+        for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
+            hf_output = hf_output.detach().clone().cpu().float()
+            vllm_output = vllm_output.detach().clone().cpu().float()
+            torch.testing.assert_close(hf_output, vllm_output, atol=3.2e-2, rtol=1e-3)
 
 
 @pytest.mark.parametrize("model", ["disham993/electrical-ner-ModernBERT-base"])

vllm/v1/worker/gpu/model_runner.py

@@ -291,7 +291,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             # Do not rely on pooling_runner here, since this information is needed
             # on the first PP rank, while pooling_runner is only initialized
             # on the last PP rank.
-            tasks.extend(PoolingRunner.get_supported_tasks(self.model))
+            tasks.extend(
+                PoolingRunner.get_supported_tasks(self.model, self.model_config)
+            )
         return tuple(tasks)
 
     def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None:

vllm/v1/worker/gpu/pool/pooling_runner.py

@@ -6,7 +6,7 @@ import numpy as np
 import torch
 import torch.nn as nn
 
-from vllm.config import VllmConfig
+from vllm.config import ModelConfig, VllmConfig
 from vllm.model_executor.models import VllmModelForPooling, is_pooling_model
 from vllm.pooling_params import PoolingParams
 from vllm.tasks import PoolingTask
@@ -16,7 +16,9 @@ from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates
 from vllm.v1.worker.gpu.input_batch import InputBatch
 from vllm.v1.worker.gpu.states import RequestState
 
-_SUPPORTED_TASKS: frozenset[PoolingTask] = frozenset({"embed", "classify"})
+_SUPPORTED_TASKS: frozenset[PoolingTask] = frozenset(
+    {"embed", "classify", "token_classify"}
+)
 
 
 class PoolingRunner:
@@ -25,19 +27,26 @@ class PoolingRunner:
         self.model_config = vllm_config.model_config
         self.max_num_reqs = vllm_config.scheduler_config.max_num_seqs
         model_tasks = tuple(sorted(self.model.pooler.get_supported_tasks()))
+        enabled_tasks = self._get_enabled_tasks(self.model_config)
         selected_task = self.model_config.get_pooling_task(model_tasks)
-        if selected_task not in _SUPPORTED_TASKS:
-            raise ValueError(
-                "Model Runner V2 supports only sequence-level pooling tasks "
-                f"{sorted(_SUPPORTED_TASKS)}, but this model selects "
-                f"{selected_task!r} from {list(model_tasks)}. Set an explicitly "
-                "supported task or VLLM_USE_V2_MODEL_RUNNER=0."
+        if selected_task not in enabled_tasks:
+            hint = (
+                "Set an explicitly supported task or VLLM_USE_V2_MODEL_RUNNER=0."
+                if enabled_tasks.intersection(model_tasks)
+                else "Set VLLM_USE_V2_MODEL_RUNNER=0 to use this model."
             )
-        self.supported_tasks = frozenset(self.get_supported_tasks(model))
+            raise ValueError(
+                "Model Runner V2 supports pooling tasks "
+                f"{sorted(enabled_tasks)}, but this model selects "
+                f"{selected_task!r} from {list(model_tasks)}. {hint}"
+            )
+        self.supported_tasks = frozenset(
+            self.get_supported_tasks(model, self.model_config)
+        )
         if not self.supported_tasks:
             raise ValueError(
-                "Model Runner V2 supports only sequence-level pooling tasks "
-                f"{sorted(_SUPPORTED_TASKS)}, but this model supports "
+                "Model Runner V2 supports pooling tasks "
+                f"{sorted(enabled_tasks)}, but this model supports "
                 f"{list(model_tasks)}. "
                 "Set VLLM_USE_V2_MODEL_RUNNER=0 to use this model."
             )
@@ -46,10 +55,20 @@ class PoolingRunner:
         self.prompt_token_ids: dict[int, torch.Tensor] = {}
 
     @staticmethod
-    def get_supported_tasks(model: nn.Module) -> list[PoolingTask]:
+    def _get_enabled_tasks(model_config: ModelConfig) -> frozenset[PoolingTask]:
+        # Token classification is validated only for unchunked encoder-only prefill.
+        if model_config.attn_type == "encoder_only":
+            return _SUPPORTED_TASKS
+        return _SUPPORTED_TASKS - {"token_classify"}
+
+    @staticmethod
+    def get_supported_tasks(
+        model: nn.Module, model_config: ModelConfig
+    ) -> list[PoolingTask]:
         if not is_pooling_model(model):
             return []
-        return sorted(model.pooler.get_supported_tasks() & _SUPPORTED_TASKS)
+        enabled_tasks = PoolingRunner._get_enabled_tasks(model_config)
+        return sorted(model.pooler.get_supported_tasks() & enabled_tasks)
 
     def add_request(
         self,