[ModelRunner v2] Support Transformers pooling model (#52425)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Taneem Ibrahim committed
Aug 16, 2026 at 14:25 UTC
fdab2b10bcac00a16c406f8b17a75a1c3f729e59
2 files changed
+18
-5
tests/models/transformers/test_backend.py
+10
-1
@@ -263,7 +263,16 @@ def test_embed_loading(vllm_runner, model):
263
@pytest.mark.parametrize(
264
"arch", ["TransformersEmbeddingModel", "TransformersForSequenceClassification"]
265
)
266
-def test_pooling(hf_runner, vllm_runner, example_prompts, arch):
266
+@pytest.mark.parametrize("use_v2_model_runner", [False, True], ids=["v1", "v2"])
267
+def test_pooling(
268
+ hf_runner,
269
+ vllm_runner,
270
+ example_prompts,
271
+ arch,
272
+ monkeypatch,
273
+ use_v2_model_runner,
274
+):
275
+ monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", str(int(use_v2_model_runner)))
276
model = get_model(arch)
277
278
vllm_kwargs = dict(max_model_len=None, model_impl="transformers")
vllm/v1/worker/gpu/model_states/__init__.py
+8
-4
@@ -3,8 +3,9 @@
3
import torch
4
import torch.nn as nn
5
6
-from vllm.config import VllmConfig
7
-from vllm.model_executor.layers.attention import CrossAttention, EncoderOnlyAttention
6
+from vllm.config import VllmConfig, get_layers_from_vllm_config
7
+from vllm.model_executor.layers.attention import Attention, CrossAttention
8
+from vllm.v1.attention.backend import AttentionType
9
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
10
11
@@ -27,8 +28,11 @@ def init_model_state(
28
29
return EncoderDecoderModelState(vllm_config, model, encoder_cache, device)
30
30
- # Encoder-only models (BERT/RoBERTa): non-causal self-attention, no KV cache.
31
- if any(isinstance(m, EncoderOnlyAttention) for m in model.modules()):
31
+ # Encoder-only attention is non-causal and needs no KV cache.
32
+ if any(
33
+ layer.attn_type == AttentionType.ENCODER_ONLY
34
+ for layer in get_layers_from_vllm_config(vllm_config, Attention).values()
35
+ ):
36
from vllm.v1.worker.gpu.model_states.encoder_only import EncoderOnlyModelState
37
38
return EncoderOnlyModelState(vllm_config, model, encoder_cache, device)