From 0f17394564fa2fccd332cf63321314884c15ee37 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Fri, 31 Jul 2026 00:15:36 -0400 Subject: [PATCH] [Model Runner V2] Enable encoder token classification (#50293) Signed-off-by: Taneem Ibrahim --- .../pooling/test_splade_sparse_pooler.py | 41 +++++++++++++++++ .../pooling/test_token_classification.py | 22 ++++++--- vllm/v1/worker/gpu/model_runner.py | 4 +- vllm/v1/worker/gpu/pool/pooling_runner.py | 45 +++++++++++++------ 4 files changed, 92 insertions(+), 20 deletions(-) diff --git a/tests/models/language/pooling/test_splade_sparse_pooler.py b/tests/models/language/pooling/test_splade_sparse_pooler.py index b77b494435..4a43bc1991 100644 --- a/tests/models/language/pooling/test_splade_sparse_pooler.py +++ b/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"} diff --git a/tests/models/language/pooling/test_token_classification.py b/tests/models/language/pooling/test_token_classification.py index 2e95fb8f70..5b18bdf1a5 100644 --- a/tests/models/language/pooling/test_token_classification.py +++ b/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"]) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 9b0380093d..2183b0df3b 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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: diff --git a/vllm/v1/worker/gpu/pool/pooling_runner.py b/vllm/v1/worker/gpu/pool/pooling_runner.py index f2fd1170ba..11075ed5d9 100644 --- a/vllm/v1/worker/gpu/pool/pooling_runner.py +++ b/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,