| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from typing import Any, Literal, get_args |
| 5 | |
| 6 | from pydantic import model_validator |
| 7 | from pydantic_core import ArgsKwargs |
| 8 | |
| 9 | from vllm.config.utils import config |
| 10 | from vllm.logger import init_logger |
| 11 | from vllm.tasks import PoolingTask, check_removed_pooling_task |
| 12 | from vllm.utils.hashing import safe_hash |
| 13 | |
| 14 | logger = init_logger(__name__) |
| 15 | |
| 16 | SequencePoolingType = Literal["CLS", "LAST", "MEAN"] |
| 17 | SEQ_POOLING_TYPES: tuple[SequencePoolingType, ...] = get_args(SequencePoolingType) |
| 18 | |
| 19 | TokenPoolingType = Literal["ALL", "STEP"] |
| 20 | TOK_POOLING_TYPES: tuple[TokenPoolingType, ...] = get_args(TokenPoolingType) |
| 21 | |
| 22 | POOLER_CONFIG_LOG_FIELDS = ( |
| 23 | "seq_pooling_type", |
| 24 | "tok_pooling_type", |
| 25 | "use_activation", |
| 26 | ) |
| 27 | |
| 28 | |
| 29 | @config |
| 30 | class PoolerConfig: |
| 31 | """Controls the behavior of output pooling in pooling models.""" |
| 32 | |
| 33 | task: PoolingTask | None = None |
| 34 | """ |
| 35 | The task used for pooling. |
| 36 | """ |
| 37 | |
| 38 | pooling_type: SequencePoolingType | TokenPoolingType | None = None |
| 39 | """ |
| 40 | The pooling method used for pooling. |
| 41 | |
| 42 | If set, `seq_pooling_type` or `tok_pooling_type` are automatically populated |
| 43 | with this field. Alternatively, users can set `seq_pooling_type` and |
| 44 | `tok_pooling_type` explicitly. |
| 45 | |
| 46 | This field is mainly for user convenience. Internal code should always use |
| 47 | `seq_pooling_type` or `tok_pooling_type` instead of `pooling_type`. |
| 48 | """ |
| 49 | |
| 50 | seq_pooling_type: SequencePoolingType | None = None |
| 51 | """ |
| 52 | The pooling method used for sequence pooling. |
| 53 | """ |
| 54 | |
| 55 | tok_pooling_type: TokenPoolingType | None = None |
| 56 | """ |
| 57 | The pooling method used for tokenwise pooling. |
| 58 | """ |
| 59 | |
| 60 | use_activation: bool | None = None |
| 61 | """ |
| 62 | Whether to apply activation function to the pooler outputs. |
| 63 | `None` uses the pooler's default, which is `True` in most cases. |
| 64 | """ |
| 65 | |
| 66 | enable_flash_late_interaction: bool = True |
| 67 | """ |
| 68 | Whether the engine-side late-interaction scorer may use the fused |
| 69 | flash-maxsim Triton kernel. Disabled automatically when the API server |
| 70 | is started with `--no-enable-flash-late-interaction`; the reference |
| 71 | scorer is used instead. |
| 72 | """ |
| 73 | |
| 74 | ## for embedding models |
| 75 | dimensions: int | None = None |
| 76 | """ |
| 77 | Reduce the dimensions of embeddings if model |
| 78 | support matryoshka representation. Defaults to None. |
| 79 | """ |
| 80 | enable_chunked_processing: bool = False |
| 81 | """ |
| 82 | Whether to enable chunked processing for long inputs that exceed the model's |
| 83 | maximum position embeddings. When enabled, long inputs will be split into |
| 84 | chunks, processed separately, and then aggregated using weighted averaging. |
| 85 | This allows embedding models to handle arbitrarily long text without CUDA |
| 86 | errors. Defaults to False. |
| 87 | """ |
| 88 | max_embed_len: int | None = None |
| 89 | """ |
| 90 | Maximum input length allowed for embedding generation. When set, allows |
| 91 | inputs longer than max_embed_len to be accepted for embedding models. |
| 92 | When an input exceeds max_embed_len, it will be handled according to |
| 93 | the original max_model_len validation logic. |
| 94 | Defaults to None (i.e. set to max_model_len). |
| 95 | """ |
| 96 | |
| 97 | ## for classification models — affine score calibration |
| 98 | logit_mean: float | None = None |
| 99 | """ |
| 100 | If provided, subtract this value from classification logits before |
| 101 | activation. Used for affine score calibration (Platt scaling): |
| 102 | activation((logit - logit_mean) / logit_sigma). Defaults to None. |
| 103 | """ |
| 104 | |
| 105 | logit_sigma: float | None = None |
| 106 | """ |
| 107 | If provided, divide the classification logits by this value after |
| 108 | mean subtraction. Used for affine score calibration (Platt scaling): |
| 109 | activation((logit - logit_mean) / logit_sigma). Defaults to None. |
| 110 | """ |
| 111 | |
| 112 | ## for reward models |
| 113 | step_tag_id: int | None = None |
| 114 | """ |
| 115 | If set, only the score corresponding to the `step_tag_id` in the |
| 116 | generated sentence should be returned. Otherwise, the scores for all tokens |
| 117 | are returned. |
| 118 | """ |
| 119 | returned_token_ids: list[int] | None = None |
| 120 | """ |
| 121 | A list of indices for the vocabulary dimensions to be extracted, |
| 122 | such as the token IDs of `good_token` and `bad_token` in the |
| 123 | `math-shepherd-mistral-7b-prm` model. |
| 124 | """ |
| 125 | |
| 126 | @model_validator(mode="before") |
| 127 | @classmethod |
| 128 | def reject_removed_parameters(cls, data): |
| 129 | values = data.kwargs if isinstance(data, ArgsKwargs) else data |
| 130 | if not isinstance(values, dict): |
| 131 | return data |
| 132 | if "normalize" in values: |
| 133 | raise ValueError( |
| 134 | "Parameter `normalize` was removed; use `use_activation` instead." |
| 135 | ) |
| 136 | check_removed_pooling_task(values.get("task")) |
| 137 | return data |
| 138 | |
| 139 | def __post_init__(self) -> None: |
| 140 | if self.logit_sigma is not None and self.logit_sigma == 0: |
| 141 | raise ValueError("logit_sigma cannot be 0 (division by zero)") |
| 142 | |
| 143 | if pooling_type := self.pooling_type: |
| 144 | if self.seq_pooling_type is not None: |
| 145 | raise ValueError( |
| 146 | "Cannot set both `pooling_type` and `seq_pooling_type`" |
| 147 | ) |
| 148 | if self.tok_pooling_type is not None: |
| 149 | raise ValueError( |
| 150 | "Cannot set both `pooling_type` and `tok_pooling_type`" |
| 151 | ) |
| 152 | |
| 153 | if pooling_type in SEQ_POOLING_TYPES: |
| 154 | logger.debug( |
| 155 | "Resolved `pooling_type=%r` to `seq_pooling_type=%r`.", |
| 156 | pooling_type, |
| 157 | pooling_type, |
| 158 | ) |
| 159 | self.seq_pooling_type = pooling_type # type: ignore[assignment] |
| 160 | elif pooling_type in TOK_POOLING_TYPES: |
| 161 | logger.debug( |
| 162 | "Resolved `pooling_type=%r` to `tok_pooling_type=%r`.", |
| 163 | pooling_type, |
| 164 | pooling_type, |
| 165 | ) |
| 166 | self.tok_pooling_type = pooling_type # type: ignore[assignment] |
| 167 | else: |
| 168 | raise NotImplementedError(pooling_type) |
| 169 | |
| 170 | def get_seq_pooling_type(self) -> SequencePoolingType: |
| 171 | if self.seq_pooling_type is None: |
| 172 | raise ValueError( |
| 173 | "seq_pooling_type is not set; it should be resolved by" |
| 174 | " ModelConfig before calling get_seq_pooling_type()" |
| 175 | ) |
| 176 | return self.seq_pooling_type |
| 177 | |
| 178 | def get_tok_pooling_type(self) -> TokenPoolingType: |
| 179 | if self.tok_pooling_type is None: |
| 180 | raise ValueError( |
| 181 | "tok_pooling_type is not set; it should be resolved by" |
| 182 | " ModelConfig before calling get_tok_pooling_type()" |
| 183 | ) |
| 184 | return self.tok_pooling_type |
| 185 | |
| 186 | def compute_hash(self) -> str: |
| 187 | """WARNING: Whenever a new field is added to this config, |
| 188 | ensure that it is included in the factors list if |
| 189 | it affects the computation graph. |
| 190 | |
| 191 | Provide a hash that uniquely identifies all the configs |
| 192 | that affect the structure of the computation |
| 193 | graph from input ids/embeddings to the final hidden states, |
| 194 | excluding anything before input ids/embeddings and after |
| 195 | the final hidden states. |
| 196 | """ |
| 197 | # no factors to consider. |
| 198 | # this config will not affect the computation graph. |
| 199 | factors: list[Any] = [] |
| 200 | hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest() |
| 201 | return hash_str |