[ModelRunner v2] Support Transformers pooling model (#52425)

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

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")

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)