[Model Runner V2] Enable encoder token classification (#50293)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Taneem Ibrahim committed
Jul 31, 2026 at 00:15 UTC
0f17394564fa2fccd332cf63321314884c15ee37
4 files changed
+91
-19
tests/models/language/pooling/test_splade_sparse_pooler.py
+41
@@ -159,6 +159,7 @@ def test_pooling_runner_rejects_unsupported_selected_task() -> None:
159
}
160
vllm_config = MagicMock()
161
vllm_config.scheduler_config.max_num_seqs = 2
162
+ vllm_config.model_config.attn_type = "encoder_only"
163
vllm_config.model_config.get_pooling_task.return_value = "embed&token_classify"
164
165
with (
@@ -166,3 +167,43 @@ def test_pooling_runner_rejects_unsupported_selected_task() -> None:
167
pytest.raises(ValueError, match="selects 'embed&token_classify'"),
168
):
169
PoolingRunner(model, vllm_config)
170
+
171
+
172
+def test_pooling_runner_supports_encoder_token_classification() -> None:
173
+ model = MagicMock()
174
+ model.pooler.get_supported_tasks.return_value = {"token_classify"}
175
+ vllm_config = MagicMock()
176
+ vllm_config.scheduler_config.max_num_seqs = 2
177
+ vllm_config.model_config.attn_type = "encoder_only"
178
+ vllm_config.model_config.get_pooling_task.return_value = "token_classify"
179
+
180
+ runner = PoolingRunner(model, vllm_config)
181
+
182
+ assert runner.supported_tasks == {"token_classify"}
183
+
184
+
185
+def test_pooling_runner_rejects_decoder_token_classification() -> None:
186
+ model = MagicMock()
187
+ model.pooler.get_supported_tasks.return_value = {"token_classify"}
188
+ vllm_config = MagicMock()
189
+ vllm_config.scheduler_config.max_num_seqs = 2
190
+ vllm_config.model_config.attn_type = "decoder"
191
+ vllm_config.model_config.get_pooling_task.return_value = "token_classify"
192
+
193
+ with pytest.raises(ValueError, match="selects 'token_classify'") as exc_info:
194
+ PoolingRunner(model, vllm_config)
195
+
196
+ assert "Set an explicitly supported task" not in str(exc_info.value)
197
+
198
+
199
+def test_pooling_runner_filters_decoder_token_classification() -> None:
200
+ model = MagicMock()
201
+ model.pooler.get_supported_tasks.return_value = {"embed", "token_classify"}
202
+ vllm_config = MagicMock()
203
+ vllm_config.scheduler_config.max_num_seqs = 2
204
+ vllm_config.model_config.attn_type = "decoder"
205
+ vllm_config.model_config.get_pooling_task.return_value = "embed"
206
+
207
+ runner = PoolingRunner(model, vllm_config)
208
+
209
+ assert runner.supported_tasks == {"embed"}
tests/models/language/pooling/test_token_classification.py
+16
-6
@@ -29,16 +29,24 @@ def seed_everything():
29
)
30
# The float32 is required for this tiny model to pass the test.
31
@pytest.mark.parametrize("dtype", ["float"])
32
+@pytest.mark.core_model
33
@torch.inference_mode
33
-def test_bert_models(
34
+def test_bert_model_runner_v2(
35
hf_runner,
36
vllm_runner,
37
example_prompts,
38
+ monkeypatch,
39
model: str,
40
dtype: str,
41
) -> None:
42
+ prompt_batches = [[example_prompts[0]], example_prompts]
43
+
44
+ monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
45
with vllm_runner(model, max_model_len=None, dtype=dtype) as vllm_model:
41
- vllm_outputs = vllm_model.token_classify(example_prompts)
46
+ assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner
47
+ vllm_output_batches = [
48
+ vllm_model.token_classify(prompts) for prompts in prompt_batches
49
+ ]
50
51
# Use eager attention on ROCm to avoid HF Transformers flash attention
52
# accuracy issues: https://github.com/vllm-project/vllm/issues/30167
@@ -59,12 +67,14 @@ def test_bert_models(
67
inputs = hf_model.wrap_device(inputs)
68
output = hf_model.model(**inputs)
69
hf_outputs.append(softmax(output.logits[0]))
70
+ hf_output_batches = [[hf_outputs[0]], hf_outputs]
71
72
# check logits difference
64
- for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
65
- hf_output = hf_output.detach().clone().cpu().float()
66
- vllm_output = vllm_output.detach().clone().cpu().float()
67
- torch.testing.assert_close(hf_output, vllm_output, atol=3.2e-2, rtol=1e-3)
73
+ for hf_outputs, vllm_outputs in zip(hf_output_batches, vllm_output_batches):
74
+ for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
75
+ hf_output = hf_output.detach().clone().cpu().float()
76
+ vllm_output = vllm_output.detach().clone().cpu().float()
77
+ torch.testing.assert_close(hf_output, vllm_output, atol=3.2e-2, rtol=1e-3)
78
79
80
@pytest.mark.parametrize("model", ["disham993/electrical-ner-ModernBERT-base"])
vllm/v1/worker/gpu/model_runner.py
+3
-1
@@ -291,7 +291,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
291
# Do not rely on pooling_runner here, since this information is needed
292
# on the first PP rank, while pooling_runner is only initialized
293
# on the last PP rank.
294
- tasks.extend(PoolingRunner.get_supported_tasks(self.model))
294
+ tasks.extend(
295
+ PoolingRunner.get_supported_tasks(self.model, self.model_config)
296
+ )
297
return tuple(tasks)
298
299
def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None:
vllm/v1/worker/gpu/pool/pooling_runner.py
+31
-12
@@ -6,7 +6,7 @@ import numpy as np
6
import torch
7
import torch.nn as nn
8
9
-from vllm.config import VllmConfig
9
+from vllm.config import ModelConfig, VllmConfig
10
from vllm.model_executor.models import VllmModelForPooling, is_pooling_model
11
from vllm.pooling_params import PoolingParams
12
from vllm.tasks import PoolingTask
@@ -16,7 +16,9 @@ from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates
16
from vllm.v1.worker.gpu.input_batch import InputBatch
17
from vllm.v1.worker.gpu.states import RequestState
18
19
-_SUPPORTED_TASKS: frozenset[PoolingTask] = frozenset({"embed", "classify"})
19
+_SUPPORTED_TASKS: frozenset[PoolingTask] = frozenset(
20
+ {"embed", "classify", "token_classify"}
21
+)
22
23
24
class PoolingRunner:
@@ -25,19 +27,26 @@ class PoolingRunner:
27
self.model_config = vllm_config.model_config
28
self.max_num_reqs = vllm_config.scheduler_config.max_num_seqs
29
model_tasks = tuple(sorted(self.model.pooler.get_supported_tasks()))
30
+ enabled_tasks = self._get_enabled_tasks(self.model_config)
31
selected_task = self.model_config.get_pooling_task(model_tasks)
29
- if selected_task not in _SUPPORTED_TASKS:
32
+ if selected_task not in enabled_tasks:
33
+ hint = (
34
+ "Set an explicitly supported task or VLLM_USE_V2_MODEL_RUNNER=0."
35
+ if enabled_tasks.intersection(model_tasks)
36
+ else "Set VLLM_USE_V2_MODEL_RUNNER=0 to use this model."
37
+ )
38
raise ValueError(
31
- "Model Runner V2 supports only sequence-level pooling tasks "
32
- f"{sorted(_SUPPORTED_TASKS)}, but this model selects "
33
- f"{selected_task!r} from {list(model_tasks)}. Set an explicitly "
34
- "supported task or VLLM_USE_V2_MODEL_RUNNER=0."
39
+ "Model Runner V2 supports pooling tasks "
40
+ f"{sorted(enabled_tasks)}, but this model selects "
41
+ f"{selected_task!r} from {list(model_tasks)}. {hint}"
42
)
36
- self.supported_tasks = frozenset(self.get_supported_tasks(model))
43
+ self.supported_tasks = frozenset(
44
+ self.get_supported_tasks(model, self.model_config)
45
+ )
46
if not self.supported_tasks:
47
raise ValueError(
39
- "Model Runner V2 supports only sequence-level pooling tasks "
40
- f"{sorted(_SUPPORTED_TASKS)}, but this model supports "
48
+ "Model Runner V2 supports pooling tasks "
49
+ f"{sorted(enabled_tasks)}, but this model supports "
50
f"{list(model_tasks)}. "
51
"Set VLLM_USE_V2_MODEL_RUNNER=0 to use this model."
52
)
@@ -46,10 +55,20 @@ class PoolingRunner:
55
self.prompt_token_ids: dict[int, torch.Tensor] = {}
56
57
@staticmethod
49
- def get_supported_tasks(model: nn.Module) -> list[PoolingTask]:
58
+ def _get_enabled_tasks(model_config: ModelConfig) -> frozenset[PoolingTask]:
59
+ # Token classification is validated only for unchunked encoder-only prefill.
60
+ if model_config.attn_type == "encoder_only":
61
+ return _SUPPORTED_TASKS
62
+ return _SUPPORTED_TASKS - {"token_classify"}
63
+
64
+ @staticmethod
65
+ def get_supported_tasks(
66
+ model: nn.Module, model_config: ModelConfig
67
+ ) -> list[PoolingTask]:
68
if not is_pooling_model(model):
69
return []
52
- return sorted(model.pooler.get_supported_tasks() & _SUPPORTED_TASKS)
70
+ enabled_tasks = PoolingRunner._get_enabled_tasks(model_config)
71
+ return sorted(model.pooler.get_supported_tasks() & enabled_tasks)
72
73
def add_request(
74
self,