| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import importlib |
| 4 | import inspect |
| 5 | import itertools |
| 6 | from abc import abstractmethod |
| 7 | from collections.abc import Sequence |
| 8 | from functools import lru_cache, partial |
| 9 | from typing import TYPE_CHECKING |
| 10 | |
| 11 | import torch |
| 12 | |
| 13 | from vllm.exceptions import VLLMValidationError |
| 14 | from vllm.logger import init_logger |
| 15 | from vllm.logits_process import LogitsProcessor as RequestLogitsProcessor |
| 16 | from vllm.sampling_params import SamplingParams |
| 17 | from vllm.utils.torch_utils import guard_cuda_initialization |
| 18 | from vllm.v1.sample.logits_processor.builtin import ( |
| 19 | LogitBiasLogitsProcessor, |
| 20 | MinPLogitsProcessor, |
| 21 | MinTokensLogitsProcessor, |
| 22 | process_dict_updates, |
| 23 | ) |
| 24 | from vllm.v1.sample.logits_processor.interface import ( |
| 25 | BatchUpdate, |
| 26 | LogitsProcessor, |
| 27 | MoveDirectionality, |
| 28 | ) |
| 29 | from vllm.v1.sample.logits_processor.state import BatchUpdateBuilder, LogitsProcessors |
| 30 | |
| 31 | if TYPE_CHECKING: |
| 32 | from vllm.config import VllmConfig |
| 33 | |
| 34 | logger = init_logger(__name__) |
| 35 | |
| 36 | # Error message when the user tries to initialize vLLM with a pooling model |
| 37 | # and custom logitsproces |
| 38 | STR_POOLING_REJECTS_LOGITSPROCS = ( |
| 39 | "Pooling models do not support custom logits processors." |
| 40 | ) |
| 41 | |
| 42 | # Error message when the user tries to initialize vLLM with a speculative |
| 43 | # decoding enabled and custom logitsproces |
| 44 | STR_SPEC_DEC_REJECTS_LOGITSPROCS = ( |
| 45 | "Custom logits processors are not supported when speculative decoding is enabled." |
| 46 | ) |
| 47 | |
| 48 | LOGITSPROCS_GROUP = "vllm.logits_processors" |
| 49 | |
| 50 | BUILTIN_LOGITS_PROCESSORS: list[type[LogitsProcessor]] = [ |
| 51 | MinTokensLogitsProcessor, |
| 52 | LogitBiasLogitsProcessor, |
| 53 | MinPLogitsProcessor, |
| 54 | ] |
| 55 | |
| 56 | |
| 57 | def _load_logitsprocs_plugins() -> list[type[LogitsProcessor]]: |
| 58 | """Load all installed logit processor plugins.""" |
| 59 | from importlib.metadata import entry_points |
| 60 | |
| 61 | installed_logitsprocs_plugins = entry_points(group=LOGITSPROCS_GROUP) |
| 62 | if len(installed_logitsprocs_plugins) == 0: |
| 63 | logger.debug("No logitsprocs plugins installed (group %s).", LOGITSPROCS_GROUP) |
| 64 | return [] |
| 65 | |
| 66 | # Load logitsprocs plugins |
| 67 | logger.debug("Loading installed logitsprocs plugins (group %s):", LOGITSPROCS_GROUP) |
| 68 | classes: list[type[LogitsProcessor]] = [] |
| 69 | for entrypoint in installed_logitsprocs_plugins: |
| 70 | try: |
| 71 | logger.debug( |
| 72 | "- Loading logitproc plugin entrypoint=%s target=%s", |
| 73 | entrypoint.name, |
| 74 | entrypoint.value, |
| 75 | ) |
| 76 | with guard_cuda_initialization(): |
| 77 | classes.append(entrypoint.load()) |
| 78 | except Exception as e: |
| 79 | logger.error("Failed to load LogitsProcessor plugin %s: %s", entrypoint, e) |
| 80 | raise RuntimeError( |
| 81 | f"Failed to load LogitsProcessor plugin {entrypoint}" |
| 82 | ) from e |
| 83 | return classes |
| 84 | |
| 85 | |
| 86 | def _load_logitsprocs_by_fqcns( |
| 87 | logits_processors: Sequence[str | type[LogitsProcessor]] | None, |
| 88 | ) -> list[type[LogitsProcessor]]: |
| 89 | """Load logit processor types, identifying them by fully-qualified class |
| 90 | names (FQCNs). |
| 91 | |
| 92 | Effectively, a mixed list of logitproc types and FQCN strings is converted |
| 93 | into a list of entirely logitproc types, by loading from the FQCNs. |
| 94 | |
| 95 | FQCN syntax is <module>:<type> i.e. x.y.z:CustomLogitProc |
| 96 | |
| 97 | Already-loaded logitproc types must be subclasses of LogitsProcessor |
| 98 | |
| 99 | Args: |
| 100 | logits_processors: Potentially mixed list of logitsprocs types and FQCN |
| 101 | strings for logitproc types |
| 102 | |
| 103 | Returns: |
| 104 | List of logitproc types |
| 105 | |
| 106 | """ |
| 107 | if not logits_processors: |
| 108 | return [] |
| 109 | |
| 110 | logger.debug( |
| 111 | "%s additional custom logits processors specified, checking whether " |
| 112 | "they need to be loaded.", |
| 113 | len(logits_processors), |
| 114 | ) |
| 115 | |
| 116 | classes: list[type[LogitsProcessor]] = [] |
| 117 | for ldx, logitproc in enumerate(logits_processors): |
| 118 | if isinstance(logitproc, type): |
| 119 | logger.debug(" - Already-loaded logit processor: %s", logitproc.__name__) |
| 120 | if not issubclass(logitproc, LogitsProcessor): |
| 121 | raise ValueError( |
| 122 | f"{logitproc.__name__} is not a subclass of LogitsProcessor" |
| 123 | ) |
| 124 | classes.append(logitproc) |
| 125 | continue |
| 126 | |
| 127 | logger.debug("- Loading logits processor %s", logitproc) |
| 128 | parts = logitproc.split(":") |
| 129 | if len(parts) != 2: |
| 130 | raise ValueError( |
| 131 | f"Invalid logits processor FQCN {logitproc!r}. " |
| 132 | "Expected format: '<module>:<type>'" |
| 133 | ) |
| 134 | module_path, qualname = parts |
| 135 | |
| 136 | try: |
| 137 | # Load module |
| 138 | with guard_cuda_initialization(): |
| 139 | module = importlib.import_module(module_path) |
| 140 | except Exception as e: |
| 141 | logger.error( |
| 142 | "Failed to load %sth LogitsProcessor plugin %s: %s", |
| 143 | ldx, |
| 144 | logitproc, |
| 145 | e, |
| 146 | ) |
| 147 | raise RuntimeError( |
| 148 | f"Failed to load {ldx}th LogitsProcessor plugin {logitproc}" |
| 149 | ) from e |
| 150 | |
| 151 | # Walk down dotted name to get logitproc class |
| 152 | obj = module |
| 153 | for attr in qualname.split("."): |
| 154 | obj = getattr(obj, attr) |
| 155 | if not isinstance(obj, type): |
| 156 | raise ValueError("Loaded logit processor must be a type.") |
| 157 | if not issubclass(obj, LogitsProcessor): |
| 158 | raise ValueError(f"{obj.__name__} must be a subclass of LogitsProcessor") |
| 159 | classes.append(obj) |
| 160 | |
| 161 | return classes |
| 162 | |
| 163 | |
| 164 | def _load_custom_logitsprocs( |
| 165 | logits_processors: Sequence[str | type[LogitsProcessor]] | None, |
| 166 | ) -> list[type[LogitsProcessor]]: |
| 167 | """Load all custom logits processors. |
| 168 | |
| 169 | * First load all installed logitproc plugins |
| 170 | * Second load custom logitsprocs pass by the user at initialization time |
| 171 | |
| 172 | Args: |
| 173 | logits_processors: potentially mixed list of logitproc types and |
| 174 | logitproc type fully-qualified names (FQCNs) |
| 175 | which need to be loaded |
| 176 | |
| 177 | Returns: |
| 178 | A list of all loaded logitproc types |
| 179 | |
| 180 | """ |
| 181 | from vllm.platforms import current_platform |
| 182 | |
| 183 | if current_platform.is_tpu(): |
| 184 | # No logitsprocs specified by caller |
| 185 | # TODO(andy) - vLLM V1 on TPU does not support custom logitsprocs |
| 186 | return [] |
| 187 | |
| 188 | return _load_logitsprocs_plugins() + _load_logitsprocs_by_fqcns(logits_processors) |
| 189 | |
| 190 | |
| 191 | def build_logitsprocs( |
| 192 | vllm_config: "VllmConfig", |
| 193 | device: torch.device, |
| 194 | is_pin_memory: bool, |
| 195 | is_pooling_model: bool, |
| 196 | custom_logitsprocs: Sequence[str | type[LogitsProcessor]] = (), |
| 197 | ) -> LogitsProcessors: |
| 198 | if is_pooling_model: |
| 199 | if custom_logitsprocs: |
| 200 | raise ValueError(STR_POOLING_REJECTS_LOGITSPROCS) |
| 201 | logger.debug( |
| 202 | "Skipping logits processor loading because pooling models" |
| 203 | " do not support logits processors." |
| 204 | ) |
| 205 | return LogitsProcessors() |
| 206 | |
| 207 | # Check if speculative decoding is enabled. |
| 208 | if vllm_config.speculative_config: |
| 209 | if custom_logitsprocs: |
| 210 | raise ValueError(STR_SPEC_DEC_REJECTS_LOGITSPROCS) |
| 211 | logger.warning( |
| 212 | "min_p and logit_bias parameters won't work with speculative decoding." |
| 213 | ) |
| 214 | return LogitsProcessors( |
| 215 | [MinTokensLogitsProcessor(vllm_config, device, is_pin_memory)] |
| 216 | ) |
| 217 | |
| 218 | custom_logitsprocs_classes = _load_custom_logitsprocs(custom_logitsprocs) |
| 219 | return LogitsProcessors( |
| 220 | ctor(vllm_config, device, is_pin_memory) |
| 221 | for ctor in itertools.chain( |
| 222 | BUILTIN_LOGITS_PROCESSORS, custom_logitsprocs_classes |
| 223 | ) |
| 224 | ) |
| 225 | |
| 226 | |
| 227 | cached_load_custom_logitsprocs = lru_cache(_load_custom_logitsprocs) |
| 228 | |
| 229 | |
| 230 | def validate_logits_processors_parameters( |
| 231 | logits_processors: Sequence[str | type[LogitsProcessor]] | None, |
| 232 | sampling_params: SamplingParams, |
| 233 | ): |
| 234 | logits_processors = ( |
| 235 | tuple(logits_processors) if logits_processors is not None else None |
| 236 | ) |
| 237 | for logits_procs in cached_load_custom_logitsprocs(logits_processors): |
| 238 | try: |
| 239 | logits_procs.validate_params(sampling_params) |
| 240 | except ValueError as e: |
| 241 | # Legacy custom logitsprocs may still raise ValueError from |
| 242 | # validate_params; convert for backward compatibility. |
| 243 | raise VLLMValidationError(str(e)) from e |
| 244 | |
| 245 | |
| 246 | class AdapterLogitsProcessor(LogitsProcessor): |
| 247 | """Wrapper for per-request logits processors. |
| 248 | |
| 249 | To wrap a specific per-request logits processor, |
| 250 | * Subclass `AdapterLogitsProcessor` |
| 251 | * Implement `self.is_argmax_invariant()` base-class method |
| 252 | * Implement `self.new_req_logits_processor(params)` |
| 253 | |
| 254 | `self.__init__(vllm_config, device, is_pin_memory)` does not need to be |
| 255 | overridden in general. However, to implement custom constructor behavior - |
| 256 | especially any logic which operates on or stores `vllm_config`, `device`, |
| 257 | or `is_pin_memory` - `self.__init__(vllm_config, device, is_pin_memory)` |
| 258 | must be overridden and the override must call |
| 259 | `super().__init__(vllm_config, device, is_pin_memory)` |
| 260 | """ |
| 261 | |
| 262 | def __init__( |
| 263 | self, vllm_config: "VllmConfig", device: torch.device, is_pin_memory: bool |
| 264 | ): |
| 265 | """Subclass must invoke |
| 266 | `super().__init__(vllm_config, device, is_pin_memory)`. |
| 267 | |
| 268 | Subclass constructor may find it useful to utilize the `vllm_config`, |
| 269 | `device` and `is_pin_memory` argument. However regardless of whether |
| 270 | these arguments are used, the vLLM logits processor interface requires |
| 271 | all three arguments to be present. |
| 272 | """ |
| 273 | # Map req index -> logits processor state |
| 274 | # |
| 275 | # State representation is a partial[Tensor] comprising a request-level |
| 276 | # logits processor with the output token ids argument and (if required) |
| 277 | # the prompt token ids argument pre-populated |
| 278 | # |
| 279 | # Note that the partial carries a *reference* to output token ids, and |
| 280 | # will thus always operate on the list as it is currently, not as it |
| 281 | # was when the partial was created. |
| 282 | self.req_info: dict[int, partial[torch.Tensor]] = {} |
| 283 | |
| 284 | @abstractmethod |
| 285 | def new_req_logits_processor( |
| 286 | self, |
| 287 | params: SamplingParams, |
| 288 | ) -> RequestLogitsProcessor | None: |
| 289 | """Consume request info; return a per-request logits processor. |
| 290 | |
| 291 | Return None if logits processor does not need to be applied to request |
| 292 | |
| 293 | Args: |
| 294 | params: request sampling params |
| 295 | |
| 296 | Returns: |
| 297 | None if logits processor should not be applied to request; otherwise |
| 298 | returns a `RequestLogitsProcessor` instance |
| 299 | |
| 300 | """ |
| 301 | raise NotImplementedError |
| 302 | |
| 303 | def _new_state( |
| 304 | self, |
| 305 | params: SamplingParams, |
| 306 | prompt_ids: list[int] | None, |
| 307 | output_ids: list[int], |
| 308 | ) -> partial[torch.Tensor] | None: |
| 309 | """Return state representation for new request. |
| 310 | |
| 311 | Returns None if logits processor is not applicable to request |
| 312 | |
| 313 | Args: |
| 314 | params: request sampling params |
| 315 | prompt_ids: request prompt token ids |
| 316 | output_ids: decoded tokens so far for this request |
| 317 | |
| 318 | Returns: |
| 319 | logits processor partial[Tensor] or None |
| 320 | |
| 321 | """ |
| 322 | if req_lp := self.new_req_logits_processor(params): |
| 323 | if len(inspect.signature(req_lp).parameters) == 3: |
| 324 | if prompt_ids is None: |
| 325 | raise ValueError( |
| 326 | "Prompt token ids are required for this " |
| 327 | "logits processor but were not provided." |
| 328 | ) |
| 329 | args = [prompt_ids, output_ids] |
| 330 | else: |
| 331 | args = [output_ids] |
| 332 | return partial(req_lp, *args) |
| 333 | return None |
| 334 | |
| 335 | def update_state(self, batch_update: BatchUpdate | None): |
| 336 | process_dict_updates( |
| 337 | self.req_info, |
| 338 | batch_update, |
| 339 | self._new_state, |
| 340 | ) |
| 341 | |
| 342 | def apply(self, logits: torch.Tensor) -> torch.Tensor: |
| 343 | if self.req_info: |
| 344 | # Apply per-request logits processors to corresponding rows of |
| 345 | # logits tensor |
| 346 | for req_idx, req_lp in self.req_info.items(): |
| 347 | req_logits = logits[req_idx] |
| 348 | new_logits = req_lp(req_logits) |
| 349 | if new_logits is not req_logits: |
| 350 | # Modify logits tensor row in-place if necessary |
| 351 | logits[req_idx] = new_logits |
| 352 | return logits |
| 353 | |
| 354 | |
| 355 | __all__ = [ |
| 356 | "LogitsProcessor", |
| 357 | "LogitBiasLogitsProcessor", |
| 358 | "MinPLogitsProcessor", |
| 359 | "MinTokensLogitsProcessor", |
| 360 | "BatchUpdate", |
| 361 | "BatchUpdateBuilder", |
| 362 | "MoveDirectionality", |
| 363 | "LogitsProcessors", |
| 364 | "build_logitsprocs", |
| 365 | "STR_POOLING_REJECTS_LOGITSPROCS", |
| 366 | "LOGITSPROCS_GROUP", |
| 367 | "AdapterLogitsProcessor", |
| 368 | ] |