| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from copy import deepcopy |
| 5 | from typing import Any |
| 6 | |
| 7 | import msgspec |
| 8 | |
| 9 | from vllm.config import ModelConfig, PoolerConfig |
| 10 | from vllm.exceptions import VLLMValidationError |
| 11 | from vllm.logger import init_logger |
| 12 | from vllm.sampling_params import RequestOutputKind |
| 13 | from vllm.tasks import PoolingTask, check_removed_pooling_task |
| 14 | |
| 15 | logger = init_logger(__name__) |
| 16 | |
| 17 | |
| 18 | class LateInteractionParams( |
| 19 | msgspec.Struct, |
| 20 | omit_defaults=True, # type: ignore[call-arg] |
| 21 | array_like=True, |
| 22 | ): # type: ignore[call-arg] |
| 23 | """Metadata for worker-side late-interaction scoring. |
| 24 | |
| 25 | Attributes: |
| 26 | mode: |
| 27 | - "cache_query": cache query token embeddings |
| 28 | - "score_doc": score a document against a cached query. |
| 29 | query_key: stable key used for both DP routing and worker cache lookup. |
| 30 | query_uses: expected number of document requests |
| 31 | |
| 32 | """ |
| 33 | |
| 34 | mode: str |
| 35 | query_key: str |
| 36 | query_uses: int | None = None |
| 37 | |
| 38 | |
| 39 | class PoolingParams( |
| 40 | msgspec.Struct, |
| 41 | omit_defaults=True, # type: ignore[call-arg] |
| 42 | array_like=True, |
| 43 | ): # type: ignore[call-arg] |
| 44 | """API parameters for pooling models. |
| 45 | |
| 46 | Attributes: |
| 47 | use_activation: Whether to apply activation function to the pooler outputs. |
| 48 | `None` uses the pooler's default, which is `True` in most cases. |
| 49 | dimensions: Reduce the dimensions of embeddings |
| 50 | if model support matryoshka representation. |
| 51 | |
| 52 | """ |
| 53 | |
| 54 | # --8<-- [start:common-pooling-params] |
| 55 | use_activation: bool | None = None |
| 56 | # --8<-- [end:common-pooling-params] |
| 57 | |
| 58 | ## for embeddings models |
| 59 | # --8<-- [start:embed-pooling-params] |
| 60 | dimensions: int | None = None |
| 61 | # --8<-- [end:embed-pooling-params] |
| 62 | |
| 63 | ## for step pooling models |
| 64 | step_tag_id: int | None = None |
| 65 | returned_token_ids: list[int] | None = None |
| 66 | |
| 67 | ## Internal use only |
| 68 | task: PoolingTask | None = None |
| 69 | requires_token_ids: bool = False |
| 70 | skip_reading_prefix_cache: bool | None = None |
| 71 | late_interaction_params: LateInteractionParams | None = None |
| 72 | extra_kwargs: dict[str, Any] | None = None |
| 73 | output_kind: RequestOutputKind = RequestOutputKind.FINAL_ONLY |
| 74 | |
| 75 | @property |
| 76 | def all_parameters(self) -> list[str]: |
| 77 | return ["dimensions", "use_activation"] |
| 78 | |
| 79 | @property |
| 80 | def valid_parameters(self): |
| 81 | return { |
| 82 | "embed": ["dimensions", "use_activation"], |
| 83 | "classify": ["use_activation"], |
| 84 | "token_embed": ["dimensions", "use_activation"], |
| 85 | "token_classify": ["use_activation"], |
| 86 | } |
| 87 | |
| 88 | def clone(self) -> "PoolingParams": |
| 89 | """Returns a deep copy of the PoolingParams instance.""" |
| 90 | return deepcopy(self) |
| 91 | |
| 92 | def verify(self, model_config: ModelConfig) -> None: |
| 93 | # plugin task uses io_processor.parse_data to verify inputs, |
| 94 | # skipping PoolingParams verify |
| 95 | if self.task == "plugin": |
| 96 | if self.skip_reading_prefix_cache is None: |
| 97 | self.skip_reading_prefix_cache = True |
| 98 | return |
| 99 | |
| 100 | # skipping verify, let plugins configure and validate pooling params |
| 101 | if self.task not in self.valid_parameters: |
| 102 | return |
| 103 | |
| 104 | # NOTE: Task validation needs to done against the model instance, |
| 105 | # which is not available in model config. So, it's not included |
| 106 | # in this method |
| 107 | self._merge_default_parameters(model_config) |
| 108 | self._set_default_parameters(model_config) |
| 109 | self._verify_valid_parameters() |
| 110 | |
| 111 | def _merge_default_parameters(self, model_config: ModelConfig) -> None: |
| 112 | pooler_config = model_config.pooler_config |
| 113 | if pooler_config is None: |
| 114 | return |
| 115 | |
| 116 | if self.task is None: |
| 117 | raise ValueError("task must be set before merging parameters") |
| 118 | valid_parameters = self.valid_parameters[self.task] |
| 119 | |
| 120 | for k in valid_parameters: |
| 121 | if getattr(pooler_config, k, None) is None: |
| 122 | continue |
| 123 | |
| 124 | if getattr(self, k, None) is None: |
| 125 | setattr(self, k, getattr(pooler_config, k)) |
| 126 | |
| 127 | if self.skip_reading_prefix_cache is None: |
| 128 | # If prefix caching is enabled, |
| 129 | # the output of all pooling may less than n_prompt_tokens, |
| 130 | # we need to skip reading cache at this request. |
| 131 | if self.task in ["token_embed", "token_classify"]: |
| 132 | self.skip_reading_prefix_cache = True |
| 133 | else: |
| 134 | self.skip_reading_prefix_cache = False |
| 135 | |
| 136 | self._verify_step_pooling(pooler_config, valid_parameters) |
| 137 | |
| 138 | def _verify_step_pooling( |
| 139 | self, |
| 140 | pooler_config: PoolerConfig, |
| 141 | valid_parameters: list[str], |
| 142 | ): |
| 143 | step_pooling_parameters = ["step_tag_id", "returned_token_ids"] |
| 144 | if pooler_config.tok_pooling_type != "STEP": |
| 145 | invalid_parameters = [] |
| 146 | for k in step_pooling_parameters: |
| 147 | if getattr(self, k, None) is not None: |
| 148 | invalid_parameters.append(k) |
| 149 | |
| 150 | if invalid_parameters: |
| 151 | raise VLLMValidationError( |
| 152 | f"Task {self.task} only supports {valid_parameters} " |
| 153 | f"parameters, does not support " |
| 154 | f"{invalid_parameters} parameters" |
| 155 | ) |
| 156 | else: |
| 157 | for k in step_pooling_parameters: |
| 158 | if getattr(pooler_config, k, None) is None: |
| 159 | continue |
| 160 | |
| 161 | if getattr(self, k, None) is None: |
| 162 | setattr(self, k, getattr(pooler_config, k)) |
| 163 | |
| 164 | def _set_default_parameters(self, model_config: ModelConfig): |
| 165 | if self.task in ["embed", "token_embed"]: |
| 166 | if self.use_activation is None: |
| 167 | self.use_activation = True |
| 168 | |
| 169 | if self.dimensions is not None: |
| 170 | dimensions = self.dimensions |
| 171 | model_name = model_config.served_model_name |
| 172 | embedding_size = model_config.embedding_size |
| 173 | valid_range = f"[1, {embedding_size}]" |
| 174 | dimensions_in_range = 1 <= dimensions <= embedding_size |
| 175 | if not model_config.is_matryoshka: |
| 176 | raise VLLMValidationError( |
| 177 | f"Model {model_name!r} does not support Matryoshka " |
| 178 | f"embeddings; dimensions must be unset " |
| 179 | f"(received dimensions={dimensions})." |
| 180 | ) |
| 181 | |
| 182 | if not dimensions_in_range: |
| 183 | raise VLLMValidationError( |
| 184 | f"Model {model_name!r} only supports dimensions in " |
| 185 | f"range {valid_range}, got {dimensions}." |
| 186 | ) |
| 187 | |
| 188 | mds = model_config.matryoshka_dimensions |
| 189 | if mds is not None and dimensions not in mds: |
| 190 | raise VLLMValidationError( |
| 191 | f"Model {model_name!r} only supports Matryoshka " |
| 192 | f"dimensions {str(mds)}, got {dimensions}." |
| 193 | ) |
| 194 | |
| 195 | elif self.task in ["classify", "token_classify"]: |
| 196 | if self.use_activation is None: |
| 197 | self.use_activation = True |
| 198 | else: |
| 199 | raise ValueError(f"Unknown pooling task: {self.task!r}") |
| 200 | |
| 201 | def _verify_valid_parameters(self): |
| 202 | if self.task is None: |
| 203 | raise ValueError("task must be set before verifying parameters") |
| 204 | valid_parameters = self.valid_parameters[self.task] |
| 205 | invalid_parameters = [] |
| 206 | for k in self.all_parameters: |
| 207 | if k in valid_parameters: |
| 208 | continue |
| 209 | |
| 210 | if getattr(self, k, None) is not None: |
| 211 | invalid_parameters.append(k) |
| 212 | |
| 213 | if invalid_parameters: |
| 214 | raise VLLMValidationError( |
| 215 | f"Task {self.task!r} only supports {valid_parameters} " |
| 216 | f"parameters, does not support " |
| 217 | f"{invalid_parameters} parameters" |
| 218 | ) |
| 219 | |
| 220 | def __repr__(self) -> str: |
| 221 | return ( |
| 222 | f"PoolingParams(" |
| 223 | f"task={self.task}, " |
| 224 | f"dimensions={self.dimensions}, " |
| 225 | f"use_activation={self.use_activation}, " |
| 226 | f"step_tag_id={self.step_tag_id}, " |
| 227 | f"returned_token_ids={self.returned_token_ids}, " |
| 228 | f"requires_token_ids={self.requires_token_ids}, " |
| 229 | f"skip_reading_prefix_cache={self.skip_reading_prefix_cache}, " |
| 230 | f"late_interaction_params={self.late_interaction_params}, " |
| 231 | f"extra_kwargs={self.extra_kwargs})" |
| 232 | ) |
| 233 | |
| 234 | def __post_init__(self) -> None: |
| 235 | check_removed_pooling_task(self.task) |
| 236 | if self.output_kind != RequestOutputKind.FINAL_ONLY: |
| 237 | raise VLLMValidationError( |
| 238 | "For pooling output_kind has to be FINAL_ONLY, " |
| 239 | f"got {self.output_kind!r}" |
| 240 | ) |