| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from typing import cast |
| 4 | |
| 5 | import numpy as np |
| 6 | import torch |
| 7 | import torch.nn as nn |
| 8 | |
| 9 | from vllm.config import VllmConfig |
| 10 | from vllm.model_executor.models import VllmModelForPooling, is_pooling_model |
| 11 | from vllm.pooling_params import PoolingParams |
| 12 | from vllm.tasks import PoolingTask |
| 13 | from vllm.utils.torch_utils import PIN_MEMORY |
| 14 | from vllm.v1.outputs import PoolerOutput |
| 15 | from vllm.v1.pool.late_interaction_runner import LateInteractionRunner |
| 16 | from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates |
| 17 | from vllm.v1.worker.gpu.input_batch import InputBatch |
| 18 | from vllm.v1.worker.gpu.states import RequestState |
| 19 | |
| 20 | _SUPPORTED_TASKS: frozenset[PoolingTask] = frozenset( |
| 21 | {"embed", "classify", "token_embed", "token_classify", "embed&token_classify"} |
| 22 | ) |
| 23 | |
| 24 | |
| 25 | class PoolingRunner: |
| 26 | def __init__(self, model: nn.Module, vllm_config: VllmConfig): |
| 27 | self.model = cast(VllmModelForPooling, model) |
| 28 | self.model_config = vllm_config.model_config |
| 29 | self.max_num_reqs = vllm_config.scheduler_config.max_num_seqs |
| 30 | model_tasks = tuple(sorted(self.model.pooler.get_supported_tasks())) |
| 31 | selected_task = self.model_config.get_pooling_task(model_tasks) |
| 32 | if selected_task not in _SUPPORTED_TASKS: |
| 33 | hint = ( |
| 34 | "Set an explicitly supported task or VLLM_USE_V2_MODEL_RUNNER=0." |
| 35 | if _SUPPORTED_TASKS.intersection(model_tasks) |
| 36 | else "Set VLLM_USE_V2_MODEL_RUNNER=0 to use this model." |
| 37 | ) |
| 38 | raise ValueError( |
| 39 | "Model Runner V2 supports pooling tasks " |
| 40 | f"{sorted(_SUPPORTED_TASKS)}, but this model selects " |
| 41 | f"{selected_task!r} from {list(model_tasks)}. {hint}" |
| 42 | ) |
| 43 | self.supported_tasks = frozenset(self.get_supported_tasks(model)) |
| 44 | if not self.supported_tasks: |
| 45 | raise ValueError( |
| 46 | "Model Runner V2 supports pooling tasks " |
| 47 | f"{sorted(_SUPPORTED_TASKS)}, but this model supports " |
| 48 | f"{list(model_tasks)}. " |
| 49 | "Set VLLM_USE_V2_MODEL_RUNNER=0 to use this model." |
| 50 | ) |
| 51 | self.pooling_params: dict[int, PoolingParams] = {} |
| 52 | self.pooling_states: dict[int, PoolingStates] = {} |
| 53 | self.prompt_token_ids: dict[int, torch.Tensor] = {} |
| 54 | pooler_config = self.model_config.pooler_config |
| 55 | self.late_interaction_runner = LateInteractionRunner( |
| 56 | enable_flash=pooler_config is None |
| 57 | or getattr(pooler_config, "enable_flash_late_interaction", True) |
| 58 | ) |
| 59 | |
| 60 | @staticmethod |
| 61 | def get_supported_tasks(model: nn.Module) -> list[PoolingTask]: |
| 62 | if not is_pooling_model(model): |
| 63 | return [] |
| 64 | return sorted(model.pooler.get_supported_tasks() & _SUPPORTED_TASKS) |
| 65 | |
| 66 | def add_request( |
| 67 | self, |
| 68 | req_id: str, |
| 69 | req_index: int, |
| 70 | pooling_params: PoolingParams, |
| 71 | prompt_token_ids: list[int] | None, |
| 72 | ) -> None: |
| 73 | task = pooling_params.task |
| 74 | if task not in self.supported_tasks: |
| 75 | raise ValueError( |
| 76 | f"Unsupported task: {task!r}. " |
| 77 | f"Supported tasks: {sorted(self.supported_tasks)}" |
| 78 | ) |
| 79 | self.model.pooler.get_pooling_updates(task).apply(pooling_params) |
| 80 | self.late_interaction_runner.register_request(req_id, pooling_params) |
| 81 | self.pooling_params[req_index] = pooling_params |
| 82 | self.pooling_states[req_index] = PoolingStates() |
| 83 | if pooling_params.requires_token_ids: |
| 84 | assert prompt_token_ids is not None |
| 85 | self.prompt_token_ids[req_index] = torch.tensor( |
| 86 | prompt_token_ids, dtype=torch.int64 |
| 87 | ) |
| 88 | |
| 89 | def remove_request(self, req_index: int) -> None: |
| 90 | self.pooling_params.pop(req_index, None) |
| 91 | if state := self.pooling_states.pop(req_index, None): |
| 92 | state.clean() |
| 93 | self.prompt_token_ids.pop(req_index, None) |
| 94 | |
| 95 | def on_requests_finished(self, req_ids: set[str]) -> None: |
| 96 | self.late_interaction_runner.on_requests_finished(req_ids) |
| 97 | |
| 98 | def clear(self) -> None: |
| 99 | self.late_interaction_runner.clear() |
| 100 | |
| 101 | def release_late_interaction_query_cache(self, query_keys: list[str]) -> None: |
| 102 | self.late_interaction_runner.release_queries(query_keys) |
| 103 | |
| 104 | def _get_pooling_metadata( |
| 105 | self, |
| 106 | input_batch: InputBatch, |
| 107 | req_states: RequestState, |
| 108 | device: torch.device, |
| 109 | ) -> PoolingMetadata: |
| 110 | req_indices = input_batch.idx_mapping_np.tolist() |
| 111 | pooling_params = [self.pooling_params[i] for i in req_indices] |
| 112 | pooling_states = [self.pooling_states[i] for i in req_indices] |
| 113 | prompt_lens = torch.from_numpy( |
| 114 | req_states.prompt_len.np[input_batch.idx_mapping_np] |
| 115 | ) |
| 116 | |
| 117 | prompt_token_ids_cpu = None |
| 118 | if any(params.requires_token_ids for params in pooling_params): |
| 119 | max_prompt_len = int(prompt_lens.max()) |
| 120 | prompt_token_ids_cpu = torch.zeros( |
| 121 | (input_batch.num_reqs, max_prompt_len), |
| 122 | dtype=torch.int64, |
| 123 | pin_memory=PIN_MEMORY and device.type != "cpu", |
| 124 | ) |
| 125 | for i, (req_index, params) in enumerate(zip(req_indices, pooling_params)): |
| 126 | if not params.requires_token_ids: |
| 127 | continue |
| 128 | token_ids = self.prompt_token_ids[req_index] |
| 129 | prompt_token_ids_cpu[i, : token_ids.numel()] = token_ids |
| 130 | |
| 131 | return PoolingMetadata( |
| 132 | prompt_lens=prompt_lens, |
| 133 | prompt_token_ids=None, |
| 134 | prompt_token_ids_cpu=prompt_token_ids_cpu, |
| 135 | pooling_params=pooling_params, |
| 136 | pooling_states=pooling_states, |
| 137 | ) |
| 138 | |
| 139 | def pool( |
| 140 | self, |
| 141 | hidden_states: torch.Tensor, |
| 142 | input_batch: InputBatch, |
| 143 | req_states: RequestState, |
| 144 | ) -> tuple[PoolerOutput, list[bool]]: |
| 145 | hidden_states = hidden_states[: input_batch.num_tokens] |
| 146 | pooling_metadata = self._get_pooling_metadata( |
| 147 | input_batch, req_states, hidden_states.device |
| 148 | ) |
| 149 | num_reqs = input_batch.num_reqs |
| 150 | # Pooling has no speculative tokens, so this CPU upper bound is exact. |
| 151 | seq_lens_cpu = input_batch.seq_lens_cpu_upper_bound[:num_reqs] |
| 152 | pooling_metadata.build_pooling_cursor( |
| 153 | input_batch.num_scheduled_tokens, |
| 154 | seq_lens_cpu, |
| 155 | device=hidden_states.device, |
| 156 | query_start_loc_gpu=input_batch.query_start_loc[: num_reqs + 1], |
| 157 | ) |
| 158 | pooler_output = self.model.pooler(hidden_states, pooling_metadata) |
| 159 | finished_mask = pooling_metadata.get_pooling_cursor().get_finished_mask() |
| 160 | pooler_output = self.late_interaction_runner.postprocess_pooler_output( |
| 161 | raw_pooler_output=pooler_output, |
| 162 | pooling_params=pooling_metadata.pooling_params, |
| 163 | req_ids=input_batch.req_ids, |
| 164 | finished_mask=finished_mask, |
| 165 | ) |
| 166 | return pooler_output, finished_mask |
| 167 | |
| 168 | def _dummy_pooler_run_task( |
| 169 | self, hidden_states: torch.Tensor, task: PoolingTask |
| 170 | ) -> PoolerOutput: |
| 171 | num_tokens = hidden_states.shape[0] |
| 172 | num_reqs = min(num_tokens, self.max_num_reqs) |
| 173 | base_tokens = num_tokens // num_reqs |
| 174 | num_extra = num_tokens % num_reqs |
| 175 | num_scheduled_tokens = np.full(num_reqs, base_tokens, dtype=np.int32) |
| 176 | if num_extra > 0: |
| 177 | num_scheduled_tokens[-num_extra:] += 1 |
| 178 | prompt_lens = torch.from_numpy(num_scheduled_tokens) |
| 179 | |
| 180 | pooling_params = PoolingParams(task=task) |
| 181 | pooling_params.verify(self.model_config) |
| 182 | self.model.pooler.get_pooling_updates(task).apply(pooling_params) |
| 183 | prompt_token_ids_cpu = None |
| 184 | if pooling_params.requires_token_ids: |
| 185 | prompt_token_ids_cpu = torch.zeros( |
| 186 | (num_reqs, int(prompt_lens.max())), |
| 187 | dtype=torch.int64, |
| 188 | ) |
| 189 | pooling_metadata = PoolingMetadata( |
| 190 | prompt_lens=prompt_lens, |
| 191 | prompt_token_ids=None, |
| 192 | prompt_token_ids_cpu=prompt_token_ids_cpu, |
| 193 | pooling_params=[pooling_params] * num_reqs, |
| 194 | pooling_states=[PoolingStates() for _ in range(num_reqs)], |
| 195 | ) |
| 196 | pooling_metadata.build_pooling_cursor( |
| 197 | num_scheduled_tokens, |
| 198 | seq_lens_cpu=prompt_lens, |
| 199 | device=hidden_states.device, |
| 200 | ) |
| 201 | try: |
| 202 | return self.model.pooler(hidden_states, pooling_metadata) |
| 203 | except RuntimeError as e: |
| 204 | if "out of memory" not in str(e): |
| 205 | raise |
| 206 | raise RuntimeError( |
| 207 | "CUDA out of memory occurred when warming up pooler " |
| 208 | f"({task=}) with {num_reqs} dummy requests. Please try " |
| 209 | "lowering `max_num_seqs` or `gpu_memory_utilization` when " |
| 210 | "initializing the engine." |
| 211 | ) from e |
| 212 | |
| 213 | def dummy_pooler_run(self, hidden_states: torch.Tensor) -> None: |
| 214 | for task in sorted(self.supported_tasks): |
| 215 | self._dummy_pooler_run_task(hidden_states, task) |