From fdab2b10bcac00a16c406f8b17a75a1c3f729e59 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Sun, 16 Aug 2026 14:25:49 -0500 Subject: [PATCH] [ModelRunner v2] Support Transformers pooling model (#52425) Signed-off-by: Taneem Ibrahim --- tests/models/transformers/test_backend.py | 11 ++++++++++- vllm/v1/worker/gpu/model_states/__init__.py | 12 ++++++++---- 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/tests/models/transformers/test_backend.py b/tests/models/transformers/test_backend.py index 4b6b2796af..5fd39a5b9d 100644 --- a/tests/models/transformers/test_backend.py +++ b/tests/models/transformers/test_backend.py @@ -263,7 +263,16 @@ def test_embed_loading(vllm_runner, model): @pytest.mark.parametrize( "arch", ["TransformersEmbeddingModel", "TransformersForSequenceClassification"] ) -def test_pooling(hf_runner, vllm_runner, example_prompts, arch): +@pytest.mark.parametrize("use_v2_model_runner", [False, True], ids=["v1", "v2"]) +def test_pooling( + hf_runner, + vllm_runner, + example_prompts, + arch, + monkeypatch, + use_v2_model_runner, +): + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", str(int(use_v2_model_runner))) model = get_model(arch) vllm_kwargs = dict(max_model_len=None, model_impl="transformers") diff --git a/vllm/v1/worker/gpu/model_states/__init__.py b/vllm/v1/worker/gpu/model_states/__init__.py index 3717b804ee..ed95d82856 100644 --- a/vllm/v1/worker/gpu/model_states/__init__.py +++ b/vllm/v1/worker/gpu/model_states/__init__.py @@ -3,8 +3,9 @@ import torch import torch.nn as nn -from vllm.config import VllmConfig -from vllm.model_executor.layers.attention import CrossAttention, EncoderOnlyAttention +from vllm.config import VllmConfig, get_layers_from_vllm_config +from vllm.model_executor.layers.attention import Attention, CrossAttention +from vllm.v1.attention.backend import AttentionType from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache @@ -27,8 +28,11 @@ def init_model_state( return EncoderDecoderModelState(vllm_config, model, encoder_cache, device) - # Encoder-only models (BERT/RoBERTa): non-causal self-attention, no KV cache. - if any(isinstance(m, EncoderOnlyAttention) for m in model.modules()): + # Encoder-only attention is non-causal and needs no KV cache. + if any( + layer.attn_type == AttentionType.ENCODER_ONLY + for layer in get_layers_from_vllm_config(vllm_config, Attention).values() + ): from vllm.v1.worker.gpu.model_states.encoder_only import EncoderOnlyModelState return EncoderOnlyModelState(vllm_config, model, encoder_cache, device)