diff --git a/tests/models/language/pooling/test_all_pooling_plus_chunked_prefill.py b/tests/models/language/pooling/test_all_pooling_plus_chunked_prefill.py index 251bc8c1d8a..c987e92caca 100644 --- a/tests/models/language/pooling/test_all_pooling_plus_chunked_prefill.py +++ b/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, + ) diff --git a/tests/models/language/pooling/test_classification.py b/tests/models/language/pooling/test_classification.py index e7128197bfc..e15fc7467ee 100644 --- a/tests/models/language/pooling/test_classification.py +++ b/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, + ) diff --git a/tests/models/language/pooling/test_embedding.py b/tests/models/language/pooling/test_embedding.py index 798aa890016..7d48880e9c0 100644 --- a/tests/models/language/pooling/test_embedding.py +++ b/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, + ) diff --git a/tests/models/language/pooling/test_splade_sparse_pooler.py b/tests/models/language/pooling/test_splade_sparse_pooler.py index 38a90d07abe..b77b4944355 100644 --- a/tests/models/language/pooling/test_splade_sparse_pooler.py +++ b/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) diff --git a/tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py b/tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py index f598d0547f0..489284d6677 100644 --- a/tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py +++ b/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() diff --git a/vllm/v1/worker/gpu/async_utils.py b/vllm/v1/worker/gpu/async_utils.py index 4570d726734..a5d5e87ce17 100644 --- a/vllm/v1/worker/gpu/async_utils.py +++ b/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 diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index b16e75bf45e..12db52961af 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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, ) diff --git a/vllm/v1/worker/gpu/model_states/encoder_only.py b/vllm/v1/worker/gpu/model_states/encoder_only.py index 6595b9e434b..91664b211ed 100644 --- a/vllm/v1/worker/gpu/model_states/encoder_only.py +++ b/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 diff --git a/vllm/v1/worker/gpu/pool/pooling_runner.py b/vllm/v1/worker/gpu/pool/pooling_runner.py index e5864a34d12..f2fd1170ba2 100644 --- a/vllm/v1/worker/gpu/pool/pooling_runner.py +++ b/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) diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 260ca7a5f17..cb39581c200 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/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