[ModelRunner V2] Enable sequence pooling for embedding and classification models (#48791)

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

tests/models/language/pooling/test_all_pooling_plus_chunked_prefill.py

@@ -15,7 +15,9 @@ from vllm.config import PoolerConfig
     ["Qwen/Qwen3-Embedding-0.6B"],
 )
 @torch.inference_mode
-def test_embed_models(hf_runner, vllm_runner, model: str):
+def test_embed_models(hf_runner, vllm_runner, monkeypatch, model: str):
+    # Keep token_embed on MRV1 when sequence pooling becomes MRV2 by default.
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
     chunk_size = 10
     n_prompt_tokens = [55, 56, 57]
     token_prompts = [[1024 + i for i in range(n)] for n in n_prompt_tokens]
@@ -104,3 +106,44 @@ def test_last_pool_score_chunked_prefill_matches_unchunked(vllm_runner, model: s
     assert chunked == pytest.approx(unchunked, abs=5e-2), (
         f"chunked score {chunked} diverged from unchunked {unchunked}"
     )
+
+
+@torch.inference_mode
+def test_sequence_embed_model_runner_v2(hf_runner, vllm_runner, monkeypatch) -> None:
+    model = "Qwen/Qwen3-Embedding-0.6B"
+    chunk_size = 10
+    token_prompts = [[1024 + i for i in range(n)] for n in (25, 27)]
+    prompts = [TokensPrompt(prompt_token_ids=t) for t in token_prompts]
+
+    with hf_runner(model, auto_cls=AutoModel) as hf_model:
+        hf_outputs = []
+        for token_prompt in token_prompts:
+            inputs = hf_model.wrap_device({"input_ids": torch.tensor([token_prompt])})
+            output = hf_model.model(inputs["input_ids"])
+            embedding = torch.nn.functional.normalize(
+                output.last_hidden_state.float()[0, -1], dim=0
+            )
+            hf_outputs.append(embedding.cpu().tolist())
+
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
+    with vllm_runner(
+        model,
+        runner="pooling",
+        pooler_config=PoolerConfig(task="embed"),
+        max_model_len=64,
+        max_num_batched_tokens=chunk_size,
+        max_num_seqs=2,
+        gpu_memory_utilization=0.25,
+        enforce_eager=True,
+        enable_chunked_prefill=True,
+    ) as vllm_model:
+        assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner
+        vllm_outputs = vllm_model.embed(prompts)
+
+    check_embeddings_close(
+        embeddings_0_lst=hf_outputs,
+        embeddings_1_lst=vllm_outputs,
+        name_0="hf",
+        name_1="vllm",
+        tol=1e-2,
+    )

tests/models/language/pooling/test_classification.py

@@ -49,3 +49,60 @@ def test_models(
             vllm_output,
             rtol=2e-3 if dtype == "float" else 1e-2,
         )
+
+
+@pytest.mark.core_model
+def test_bert_model_runner_v2(hf_runner, vllm_runner, monkeypatch) -> None:
+    model = "cross-encoder/ms-marco-TinyBERT-L-2-v2"
+    score_inputs = (
+        "What is the capital of France?",
+        [
+            "Paris.",
+            "Paris is the capital and largest city of France.",
+            "William Shakespeare wrote Hamlet in the early seventeenth century.",
+        ],
+    )
+    prompt_batches = [
+        ["short input"],
+        [
+            "short input",
+            "a longer input that exercises mixed sequence lengths",
+        ],
+    ]
+
+    with hf_runner(
+        model, dtype="half", auto_cls=AutoModelForSequenceClassification
+    ) as hf_model:
+        # HfRunner uses problem_type to preserve the model's
+        # sbert_ce_default_activation_function=Identity raw logits.
+        hf_model.config.problem_type = "regression"
+        hf_outputs = [hf_model.classify(prompts) for prompts in prompt_batches]
+
+    text_1, text_2 = score_inputs
+    text_pairs = [[text_1, document] for document in text_2]
+    with hf_runner(model, dtype="half", is_cross_encoder=True) as hf_model:
+        hf_scores = hf_model.predict(text_pairs).tolist()
+
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
+    with vllm_runner(
+        model,
+        runner="pooling",
+        dtype="half",
+        max_model_len=64,
+    ) as vllm_model:
+        assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner
+        vllm_outputs = [vllm_model.classify(prompts) for prompts in prompt_batches]
+        vllm_scores = vllm_model.score(*score_inputs)
+
+    for hf_batch, vllm_batch in zip(hf_outputs, vllm_outputs):
+        hf_tensor = torch.tensor(hf_batch)
+        vllm_tensor = torch.tensor(vllm_batch)
+        assert vllm_tensor.shape == hf_tensor.shape
+        assert torch.allclose(vllm_tensor, hf_tensor, rtol=1e-2, atol=1e-4)
+
+    assert torch.allclose(
+        torch.tensor(vllm_scores),
+        torch.tensor(hf_scores),
+        rtol=1e-2,
+        atol=1e-4,
+    )

tests/models/language/pooling/test_embedding.py

@@ -5,6 +5,7 @@ import pytest
 import torch
 from transformers import AutoModel
 
+from vllm import PoolingParams
 from vllm.config import PoolerConfig
 
 from ...utils import check_embeddings_close
@@ -132,6 +133,7 @@ def test_encoder_only_model_runner_v2_attention(
             task="embed", seq_pooling_type="LAST", use_activation=True
         ),
     ) as vllm_model:
+        assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner
         vllm_outputs = vllm_model.embed(prompts)
 
     check_embeddings_close(
@@ -141,3 +143,77 @@ def test_encoder_only_model_runner_v2_attention(
         name_1="vllm",
         tol=1e-2,
     )
+
+
+@pytest.mark.core_model
+def test_encoder_model_runner_v2(hf_runner, vllm_runner, monkeypatch) -> None:
+    model = "sentence-transformers/all-MiniLM-L6-v2"
+    prompt_batches = [
+        ["short input"],
+        [
+            "short input",
+            "a longer input that exercises mixed sequence lengths",
+        ],
+    ]
+
+    with hf_runner(model, is_sentence_transformer=True) as hf_model:
+        hf_outputs = [hf_model.encode(prompts) for prompts in prompt_batches]
+
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
+    with vllm_runner(
+        model,
+        runner="pooling",
+        max_model_len=64,
+    ) as vllm_model:
+        assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner
+        vllm_outputs = [vllm_model.embed(prompts) for prompts in prompt_batches]
+
+    for hf_batch, vllm_batch in zip(hf_outputs, vllm_outputs):
+        check_embeddings_close(
+            embeddings_0_lst=hf_batch,
+            embeddings_1_lst=vllm_batch,
+            name_0="hf",
+            name_1="vllm",
+            tol=1e-2,
+        )
+
+
+@pytest.mark.core_model
+def test_matryoshka_dimensions_model_runner_v2(
+    hf_runner, vllm_runner, monkeypatch
+) -> None:
+    model = "Snowflake/snowflake-arctic-embed-m-v1.5"
+    prompts = ["short input", "a longer input for a different output width"]
+    dimensions = [None, 256]
+
+    with hf_runner(model, is_sentence_transformer=True) as hf_model:
+        hf_outputs = hf_model.encode(prompts)
+
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
+    with vllm_runner(
+        model,
+        runner="pooling",
+        max_model_len=64,
+        gpu_memory_utilization=0.25,
+    ) as vllm_model:
+        assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner
+        vllm_outputs = vllm_model.embed(
+            prompts,
+            pooling_params=[PoolingParams(dimensions=d) for d in dimensions],
+        )
+
+    expected_outputs = []
+    for output, dimension in zip(hf_outputs, dimensions):
+        output = torch.as_tensor(output)
+        if dimension is not None:
+            output = torch.nn.functional.normalize(output[:dimension], dim=0)
+        expected_outputs.append(output.tolist())
+
+    assert [len(output) for output in vllm_outputs] == [768, 256]
+    check_embeddings_close(
+        embeddings_0_lst=expected_outputs,
+        embeddings_1_lst=vllm_outputs,
+        name_0="hf",
+        name_1="vllm",
+        tol=1e-2,
+    )

tests/models/language/pooling/test_splade_sparse_pooler.py

@@ -1,6 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from unittest.mock import MagicMock, patch
+
+import numpy as np
 import pytest
 import torch
 import torch.nn as nn
@@ -9,8 +12,13 @@ from vllm.model_executor.models.bert import (
     BertMLMHead,
     SPLADESparsePooler,
 )
+from vllm.platforms import current_platform
 from vllm.pooling_params import PoolingParams
+from vllm.utils.torch_utils import PIN_MEMORY
 from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates
+from vllm.v1.worker.gpu.input_batch import InputBatch
+from vllm.v1.worker.gpu.pool.pooling_runner import PoolingRunner
+from vllm.v1.worker.gpu.states import RequestState
 
 # ---------------------------------------------------------------------
 # Functional test: SPLADE formula correctness (no HF download needed)
@@ -91,3 +99,70 @@ def test_splade_pooler_matches_reference_formula(B, T, H, V):
         rtol=1e-4,
         atol=1e-4,
     )
+
+
+def test_pooling_runner_gathers_required_token_ids() -> None:
+    runner = PoolingRunner.__new__(PoolingRunner)
+    pooling_params = PoolingParams(task="embed", requires_token_ids=True)
+    runner.pooling_params = {1: pooling_params, 3: pooling_params}
+    runner.pooling_states = {1: PoolingStates(), 3: PoolingStates()}
+    runner.prompt_token_ids = {
+        1: torch.tensor([101, 102]),
+        3: torch.tensor([101, 11, 102]),
+    }
+
+    input_batch = MagicMock(spec=InputBatch)
+    input_batch.idx_mapping_np = np.array([3, 1], dtype=np.int32)
+    input_batch.num_reqs = 2
+    req_states = MagicMock(spec=RequestState)
+    req_states.prompt_len = MagicMock(np=np.array([0, 2, 0, 3], dtype=np.int32))
+    metadata = runner._get_pooling_metadata(
+        input_batch, req_states, torch.device(current_platform.device_type)
+    )
+
+    expected = torch.tensor([[101, 11, 102], [101, 102, 0]])
+    assert metadata.prompt_token_ids_cpu is not None
+    assert metadata.prompt_token_ids is not None
+    assert metadata.prompt_token_ids_cpu.is_pinned() == PIN_MEMORY
+    torch.testing.assert_close(
+        metadata.prompt_lens, torch.tensor([3, 2], dtype=torch.int32)
+    )
+    torch.testing.assert_close(metadata.prompt_token_ids_cpu, expected)
+    torch.testing.assert_close(metadata.prompt_token_ids.cpu(), expected)
+
+
+def test_pooling_runner_stores_only_required_token_ids() -> None:
+    runner = PoolingRunner.__new__(PoolingRunner)
+    runner.model = MagicMock()
+    runner.supported_tasks = frozenset({"embed"})
+    runner.pooling_params = {}
+    runner.pooling_states = {}
+    runner.prompt_token_ids = {}
+
+    runner.add_request(1, PoolingParams(task="embed"), [101, 102])
+    runner.add_request(
+        2,
+        PoolingParams(task="embed", requires_token_ids=True),
+        [101, 11, 102],
+    )
+
+    assert 1 not in runner.prompt_token_ids
+    torch.testing.assert_close(runner.prompt_token_ids[2], torch.tensor([101, 11, 102]))
+
+
+def test_pooling_runner_rejects_unsupported_selected_task() -> None:
+    model = MagicMock()
+    model.pooler.get_supported_tasks.return_value = {
+        "embed",
+        "embed&token_classify",
+        "token_classify",
+    }
+    vllm_config = MagicMock()
+    vllm_config.scheduler_config.max_num_seqs = 2
+    vllm_config.model_config.get_pooling_task.return_value = "embed&token_classify"
+
+    with (
+        patch.object(PoolingRunner, "get_supported_tasks", return_value=["embed"]),
+        pytest.raises(ValueError, match="selects 'embed&token_classify'"),
+    ):
+        PoolingRunner(model, vllm_config)

tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py

@@ -41,6 +41,7 @@ def mock_model_runner_with_req_states():
     runner.sampler = None
     runner.prompt_logprobs_worker = None
     runner.is_last_pp_rank = False
+    runner.pooling_runner = None
 
     # Mock staged writes — they use Triton kernels that require GPU
     runner.req_states.apply_staged_writes = Mock()

vllm/v1/worker/gpu/async_utils.py

@@ -6,7 +6,12 @@ import numpy as np
 import torch
 
 from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager
-from vllm.v1.outputs import AsyncModelRunnerOutput, LogprobsTensors, ModelRunnerOutput
+from vllm.v1.outputs import (
+    AsyncModelRunnerOutput,
+    LogprobsTensors,
+    ModelRunnerOutput,
+    PoolerOutput,
+)
 from vllm.v1.worker.gpu.sample.output import SamplerOutput
 
 
@@ -89,34 +94,42 @@ class AsyncPoolingOutput(AsyncModelRunnerOutput):
     def __init__(
         self,
         model_runner_output: ModelRunnerOutput,
-        pooler_output: torch.Tensor,
-        is_valid: torch.Tensor | None,
+        pooler_output: PoolerOutput,
+        finished_mask: list[bool],
         main_stream: torch.cuda.Stream,
         copy_stream: torch.cuda.Stream,
     ):
         self.model_runner_output = model_runner_output
         self.pooler_output = pooler_output
-        self.is_valid = is_valid
         # Blocking (sleep) event to avoid busy-polling the CUDA driver lock.
         self.copy_event = torch.cuda.Event(blocking=True)
 
         with stream(copy_stream, main_stream):
             copy_stream.wait_stream(main_stream)
-            self.pooler_output_cpu = self.pooler_output.to("cpu", non_blocking=True)
-            if self.is_valid is not None:
-                self.is_valid_cpu = self.is_valid.to("cpu", non_blocking=True)
+            if isinstance(self.pooler_output, torch.Tensor) and all(finished_mask):
+                self.pooler_output_cpu: PoolerOutput = self.pooler_output.to(
+                    "cpu", non_blocking=True
+                )
             else:
-                self.is_valid_cpu = None
+                outputs = (
+                    self.pooler_output.unbind()
+                    if isinstance(self.pooler_output, torch.Tensor)
+                    else self.pooler_output
+                )
+                self.pooler_output_cpu = [
+                    None
+                    if output is None or not is_finished
+                    else output.to("cpu", non_blocking=True)
+                    for output, is_finished in zip(outputs, finished_mask, strict=True)
+                ]
             self.copy_event.record(copy_stream)
 
     def get_output(self) -> ModelRunnerOutput:
-        pooler_output = list(self.pooler_output_cpu.unbind(dim=0))
+        if isinstance(self.pooler_output_cpu, torch.Tensor):
+            pooler_output = list(self.pooler_output_cpu.unbind(dim=0))
+        else:
+            pooler_output = self.pooler_output_cpu
         self.copy_event.synchronize()
-        if self.is_valid_cpu is not None:
-            is_valid_cpu = self.is_valid_cpu.tolist()
-            for i, is_valid in enumerate(is_valid_cpu):
-                if not is_valid:
-                    pooler_output[i] = None
         self.model_runner_output.pooler_output = pooler_output
         return self.model_runner_output
 

vllm/v1/worker/gpu/model_runner.py

@@ -366,7 +366,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             )
 
         if self.is_pooling_model and self.is_last_pp_rank:
-            self.pooling_runner = PoolingRunner(self.model)
+            self.pooling_runner = PoolingRunner(self.model, self.vllm_config)
         eplb_models_added |= self.eplb.maybe_register_model(
             self.model,
             self.model_config,
@@ -787,6 +787,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         req_idx = self.req_states.remove_request(req_id)
         if req_idx is None:
             return False
+        if self.pooling_runner is not None:
+            self.pooling_runner.remove_request(req_idx)
         if self.pp_handler is not None:
             self.pp_handler.on_req_idx_freed(req_idx)
         if self.encoder_cache is not None:
@@ -839,6 +841,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             )
             req_index = self.req_states.req_id_to_index[req_id]
 
+            if self.pooling_runner is not None:
+                assert new_req_data.pooling_params is not None
+                self.pooling_runner.add_request(
+                    req_index,
+                    new_req_data.pooling_params,
+                    new_req_data.prompt_token_ids,
+                )
+
             if self.encoder_cache is not None:
                 self.encoder_cache.add_request(req_id, new_req_data.mm_features)
 
@@ -1591,7 +1601,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)
 
         assert self.pooling_runner is not None
-        pooler_output, is_valid = self.pooling_runner.pool(
+        pooler_output, finished_mask = self.pooling_runner.pool(
             hidden_states, input_batch, self.req_states
         )
 
@@ -1604,7 +1614,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         async_output = AsyncPoolingOutput(
             model_runner_output=model_runner_output,
             pooler_output=pooler_output,
-            is_valid=is_valid,
+            finished_mask=finished_mask,
             main_stream=self.main_stream,
             copy_stream=self.output_copy_stream,
         )

vllm/v1/worker/gpu/model_states/encoder_only.py

@@ -8,12 +8,13 @@ import torch.nn as nn
 from vllm.config import VllmConfig, get_layers_from_vllm_config
 from vllm.config.compilation import CUDAGraphMode
 from vllm.model_executor.layers.attention import Attention
-from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
+from vllm.utils.torch_utils import PIN_MEMORY, STR_DTYPE_TO_TORCH_DTYPE
 from vllm.v1.attention.backend import (
     AttentionCGSupport,
     AttentionType,
     CommonAttentionMetadata,
 )
+from vllm.v1.core.sched.output import NewRequestData
 from vllm.v1.kv_cache_interface import (
     AttentionSpec,
     EncoderOnlyAttentionSpec,
@@ -22,6 +23,7 @@ from vllm.v1.kv_cache_interface import (
 from vllm.v1.worker.gpu.input_batch import InputBatch
 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
 from vllm.v1.worker.gpu.model_states.default import DefaultModelState
+from vllm.v1.worker.gpu.states import RequestState
 from vllm.v1.worker.utils import AttentionGroup
 
 
@@ -95,6 +97,52 @@ class EncoderOnlyModelState(DefaultModelState):
         self._dummy_slot_mapping = torch.zeros(
             self.max_num_tokens, dtype=torch.int64, device=device
         )
+        self.token_type_ids: dict[str, torch.Tensor] = {}
+
+    def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
+        super().add_request(req_index, new_req_data)
+        pooling_params = new_req_data.pooling_params
+        if pooling_params is None or pooling_params.extra_kwargs is None:
+            return
+
+        token_type_start = pooling_params.extra_kwargs.get("compressed_token_type_ids")
+        if token_type_start is not None:
+            assert new_req_data.prompt_token_ids is not None
+            self.token_type_ids[new_req_data.req_id] = (
+                torch.arange(len(new_req_data.prompt_token_ids), dtype=torch.int32)
+                >= token_type_start
+            ).to(torch.int32)
+
+    def remove_request(self, req_id: str) -> None:
+        super().remove_request(req_id)
+        self.token_type_ids.pop(req_id, None)
+
+    def prepare_inputs(
+        self, input_batch: InputBatch, req_states: RequestState
+    ) -> dict[str, torch.Tensor | None]:
+        model_inputs = super().prepare_inputs(input_batch, req_states)
+        if not self.token_type_ids:
+            return model_inputs
+
+        token_type_ids_cpu = torch.zeros(
+            input_batch.num_tokens_after_padding,
+            dtype=torch.int32,
+            pin_memory=PIN_MEMORY,
+        )
+        offset = 0
+        for i, req_id in enumerate(input_batch.req_ids):
+            num_tokens = int(input_batch.num_scheduled_tokens[i])
+            request_token_type_ids = self.token_type_ids.get(req_id)
+            if request_token_type_ids is not None:
+                start = int(input_batch.num_computed_tokens_np[i])
+                token_type_ids_cpu[offset : offset + num_tokens].copy_(
+                    request_token_type_ids[start : start + num_tokens]
+                )
+            offset += num_tokens
+        model_inputs["token_type_ids"] = token_type_ids_cpu.to(
+            self.device, non_blocking=True
+        )
+        return model_inputs
 
     def get_additional_cg_support(self) -> tuple[AttentionCGSupport, str | None]:
         # Encoder groups are built here rather than in init_attn_backend, so

vllm/v1/worker/gpu/pool/pooling_runner.py

@@ -2,45 +2,190 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from typing import cast
 
+import numpy as np
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
+from vllm.config import VllmConfig
 from vllm.model_executor.models import VllmModelForPooling, is_pooling_model
+from vllm.pooling_params import PoolingParams
 from vllm.tasks import PoolingTask
+from vllm.utils.torch_utils import PIN_MEMORY
+from vllm.v1.outputs import PoolerOutput
+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"})
+
 
-# NOTE(woosuk): Currently, this class only supports the "LAST" pooling task
-# on decoder-only models. How to support other pooling tasks and models
-# is to be determined.
 class PoolingRunner:
-    def __init__(self, model: nn.Module):
+    def __init__(self, model: nn.Module, vllm_config: VllmConfig):
         self.model = cast(VllmModelForPooling, model)
+        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()))
+        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."
+            )
+        self.supported_tasks = frozenset(self.get_supported_tasks(model))
+        if not self.supported_tasks:
+            raise ValueError(
+                "Model Runner V2 supports only sequence-level pooling tasks "
+                f"{sorted(_SUPPORTED_TASKS)}, but this model supports "
+                f"{list(model_tasks)}. "
+                "Set VLLM_USE_V2_MODEL_RUNNER=0 to use this model."
+            )
+        self.pooling_params: dict[int, PoolingParams] = {}
+        self.pooling_states: dict[int, PoolingStates] = {}
+        self.prompt_token_ids: dict[int, torch.Tensor] = {}
 
     @staticmethod
     def get_supported_tasks(model: nn.Module) -> list[PoolingTask]:
         if not is_pooling_model(model):
             return []
-        assert "embed" in model.pooler.get_supported_tasks()
-        return ["embed"]
+        return sorted(model.pooler.get_supported_tasks() & _SUPPORTED_TASKS)
+
+    def add_request(
+        self,
+        req_index: int,
+        pooling_params: PoolingParams,
+        prompt_token_ids: list[int],
+    ) -> None:
+        task = pooling_params.task
+        if task not in self.supported_tasks:
+            raise ValueError(
+                f"Unsupported task: {task!r}. "
+                f"Supported tasks: {sorted(self.supported_tasks)}"
+            )
+        self.model.pooler.get_pooling_updates(task).apply(pooling_params)
+        self.pooling_params[req_index] = pooling_params
+        self.pooling_states[req_index] = PoolingStates()
+        if pooling_params.requires_token_ids:
+            self.prompt_token_ids[req_index] = torch.tensor(
+                prompt_token_ids, dtype=torch.int64
+            )
+
+    def remove_request(self, req_index: int) -> None:
+        self.pooling_params.pop(req_index, None)
+        if state := self.pooling_states.pop(req_index, None):
+            state.clean()
+        self.prompt_token_ids.pop(req_index, None)
+
+    def _get_pooling_metadata(
+        self,
+        input_batch: InputBatch,
+        req_states: RequestState,
+        device: torch.device,
+    ) -> PoolingMetadata:
+        req_indices = input_batch.idx_mapping_np.tolist()
+        pooling_params = [self.pooling_params[i] for i in req_indices]
+        pooling_states = [self.pooling_states[i] for i in req_indices]
+        prompt_lens = torch.from_numpy(
+            req_states.prompt_len.np[input_batch.idx_mapping_np].copy()
+        )
+
+        prompt_token_ids_cpu = None
+        prompt_token_ids = None
+        if any(params.requires_token_ids for params in pooling_params):
+            max_prompt_len = int(prompt_lens.max())
+            prompt_token_ids_cpu = torch.zeros(
+                (input_batch.num_reqs, max_prompt_len),
+                dtype=torch.int64,
+                pin_memory=PIN_MEMORY,
+            )
+            for i, (req_index, params) in enumerate(zip(req_indices, pooling_params)):
+                if not params.requires_token_ids:
+                    continue
+                token_ids = self.prompt_token_ids[req_index]
+                prompt_token_ids_cpu[i, : token_ids.numel()] = token_ids
+            prompt_token_ids = prompt_token_ids_cpu.to(device, non_blocking=True)
+
+        return PoolingMetadata(
+            prompt_lens=prompt_lens,
+            prompt_token_ids=prompt_token_ids,
+            prompt_token_ids_cpu=prompt_token_ids_cpu,
+            pooling_params=pooling_params,
+            pooling_states=pooling_states,
+        )
 
     def pool(
         self,
         hidden_states: torch.Tensor,
         input_batch: InputBatch,
         req_states: RequestState,
-    ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        # TODO(woosuk): Support different types of pooling tasks.
-        last_hidden_states = hidden_states[input_batch.logits_indices]
-        # TODO(woosuk): Make normalization optional.
-        last_hidden_states = F.normalize(last_hidden_states, p=2, dim=-1)
+    ) -> tuple[PoolerOutput, list[bool]]:
+        hidden_states = hidden_states[: input_batch.num_tokens]
+        pooling_metadata = self._get_pooling_metadata(
+            input_batch, req_states, hidden_states.device
+        )
+        num_reqs = input_batch.num_reqs
+        # Pooling has no speculative tokens, so this CPU upper bound is exact.
+        seq_lens_cpu = input_batch.seq_lens_cpu_upper_bound[:num_reqs]
+        pooling_metadata.build_pooling_cursor(
+            input_batch.num_scheduled_tokens,
+            seq_lens_cpu,
+            device=hidden_states.device,
+            query_start_loc_gpu=input_batch.query_start_loc[: num_reqs + 1],
+        )
+        pooler_output = self.model.pooler(hidden_states, pooling_metadata)
 
-        prompt_len = req_states.prompt_len.gpu[input_batch.idx_mapping]
-        is_valid = input_batch.seq_lens == prompt_len
-        return last_hidden_states, is_valid
+        finished_mask = pooling_metadata.get_pooling_cursor().is_finished().tolist()
+        return pooler_output, finished_mask
+
+    def _dummy_pooler_run_task(
+        self, hidden_states: torch.Tensor, task: PoolingTask
+    ) -> PoolerOutput:
+        num_tokens = hidden_states.shape[0]
+        num_reqs = min(num_tokens, self.max_num_reqs)
+        base_tokens = num_tokens // num_reqs
+        num_extra = num_tokens % num_reqs
+        num_scheduled_tokens = np.full(num_reqs, base_tokens, dtype=np.int32)
+        if num_extra > 0:
+            num_scheduled_tokens[-num_extra:] += 1
+        prompt_lens = torch.from_numpy(num_scheduled_tokens)
+
+        pooling_params = PoolingParams(task=task)
+        pooling_params.verify(self.model_config)
+        self.model.pooler.get_pooling_updates(task).apply(pooling_params)
+        prompt_token_ids = None
+        if pooling_params.requires_token_ids:
+            prompt_token_ids = torch.zeros(
+                (num_reqs, int(prompt_lens.max())),
+                dtype=torch.int64,
+                device=hidden_states.device,
+            )
+        pooling_metadata = PoolingMetadata(
+            prompt_lens=prompt_lens,
+            prompt_token_ids=prompt_token_ids,
+            prompt_token_ids_cpu=None
+            if prompt_token_ids is None
+            else prompt_token_ids.cpu(),
+            pooling_params=[pooling_params] * num_reqs,
+            pooling_states=[PoolingStates() for _ in range(num_reqs)],
+        )
+        pooling_metadata.build_pooling_cursor(
+            num_scheduled_tokens,
+            seq_lens_cpu=prompt_lens,
+            device=hidden_states.device,
+        )
+        try:
+            return self.model.pooler(hidden_states, pooling_metadata)
+        except RuntimeError as e:
+            if "out of memory" not in str(e):
+                raise
+            raise RuntimeError(
+                "CUDA out of memory occurred when warming up pooler "
+                f"({task=}) with {num_reqs} dummy requests. Please try "
+                "lowering `max_num_seqs` or `gpu_memory_utilization` when "
+                "initializing the engine."
+            ) from e
 
     def dummy_pooler_run(self, hidden_states: torch.Tensor) -> None:
-        F.normalize(hidden_states, p=2, dim=-1)
-        return
+        for task in sorted(self.supported_tasks):
+            self._dummy_pooler_run_task(hidden_states, task)

vllm/v1/worker/gpu/warmup.py

@@ -230,7 +230,11 @@ def warmup_kernels(
     # SamplingParams exercising all sampling features.
     if model_runner.is_pooling_model:
         sampling_params = None
-        pooling_params = PoolingParams()
+        pooling_task = model_runner.model_config.get_pooling_task(
+            model_runner.get_supported_tasks()
+        )
+        pooling_params = PoolingParams(task=pooling_task)
+        pooling_params.verify(model_runner.model_config)
     else:
         sampling_params = SamplingParams.for_sampler_warmup()
         pooling_params = None