# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from copy import deepcopy from typing import Any import msgspec from vllm.config import ModelConfig, PoolerConfig from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger from vllm.sampling_params import RequestOutputKind from vllm.tasks import PoolingTask, check_removed_pooling_task logger = init_logger(__name__) class LateInteractionParams( msgspec.Struct, omit_defaults=True, # type: ignore[call-arg] array_like=True, ): # type: ignore[call-arg] """Metadata for worker-side late-interaction scoring. Attributes: mode: - "cache_query": cache query token embeddings - "score_doc": score a document against a cached query. query_key: stable key used for both DP routing and worker cache lookup. query_uses: expected number of document requests """ mode: str query_key: str query_uses: int | None = None class LateChunkingParams( msgspec.Struct, omit_defaults=True, # type: ignore[call-arg] array_like=True, ): # type: ignore[call-arg] """Fixed-length mean pooling before the token embedding head. Requests recompute all prompt tokens, bypassing prefix cache reads. Special tokens count toward chunk boundaries and are included in chunk means, including the BOS token normally filtered from BGE-M3 token embeddings. """ chunk_size: int def __post_init__(self) -> None: self.verify() def verify(self) -> None: if type(self.chunk_size) is not int or self.chunk_size <= 0: raise VLLMValidationError("chunk_size must be a positive integer") class PoolingParams( msgspec.Struct, omit_defaults=True, # type: ignore[call-arg] array_like=True, ): # type: ignore[call-arg] """API parameters for pooling models. Attributes: use_activation: Whether to apply activation function to the pooler outputs. `None` uses the pooler's default, which is `True` in most cases. dimensions: Reduce the dimensions of embeddings if model support matryoshka representation. late_chunking_params: Mean-pool contextual token states into fixed-size chunks for the token embedding task. """ # --8<-- [start:common-pooling-params] use_activation: bool | None = None # --8<-- [end:common-pooling-params] ## for embeddings models # --8<-- [start:embed-pooling-params] dimensions: int | None = None # --8<-- [end:embed-pooling-params] ## for step pooling models step_tag_id: int | None = None returned_token_ids: list[int] | None = None ## Internal use only task: PoolingTask | None = None requires_token_ids: bool = False skip_reading_prefix_cache: bool | None = None late_interaction_params: LateInteractionParams | None = None extra_kwargs: dict[str, Any] | None = None output_kind: RequestOutputKind = RequestOutputKind.FINAL_ONLY ## for token embedding models # --8<-- [start:token-embed-pooling-params] late_chunking_params: LateChunkingParams | None = None """Mean-pool contextual tokens into fixed-size chunks for `token_embed`.""" # --8<-- [end:token-embed-pooling-params] @property def all_parameters(self) -> list[str]: return ["dimensions", "use_activation", "late_chunking_params"] @property def valid_parameters(self): return { "embed": ["dimensions", "use_activation"], "classify": ["use_activation"], "token_embed": ["dimensions", "use_activation", "late_chunking_params"], "token_classify": ["use_activation"], } def clone(self) -> "PoolingParams": """Returns a deep copy of the PoolingParams instance.""" return deepcopy(self) def verify(self, model_config: ModelConfig) -> None: if self.late_chunking_params is not None: self.late_chunking_params.verify() if self.task != "token_embed": raise VLLMValidationError("Late chunking requires token_embed") if self.late_interaction_params is not None: raise VLLMValidationError( "Late chunking cannot be combined with late-interaction scoring" ) # plugin task uses io_processor.parse_data to verify inputs, # skipping PoolingParams verify if self.task == "plugin": if self.skip_reading_prefix_cache is None: self.skip_reading_prefix_cache = True return # skipping verify, let plugins configure and validate pooling params if self.task not in self.valid_parameters: return # NOTE: Task validation needs to done against the model instance, # which is not available in model config. So, it's not included # in this method self._merge_default_parameters(model_config) self._set_default_parameters(model_config) self._verify_valid_parameters() if self.late_chunking_params is not None: if self.step_tag_id is not None: raise VLLMValidationError("Late chunking requires all token rows") # KV cache hits omit hidden states needed for complete chunk means. self.skip_reading_prefix_cache = True def _merge_default_parameters(self, model_config: ModelConfig) -> None: pooler_config = model_config.pooler_config if pooler_config is None: return if self.task is None: raise ValueError("task must be set before merging parameters") valid_parameters = self.valid_parameters[self.task] for k in valid_parameters: if getattr(pooler_config, k, None) is None: continue if getattr(self, k, None) is None: setattr(self, k, getattr(pooler_config, k)) if self.skip_reading_prefix_cache is None: # If prefix caching is enabled, # the output of all pooling may less than n_prompt_tokens, # we need to skip reading cache at this request. if self.task in ["token_embed", "token_classify"]: self.skip_reading_prefix_cache = True else: self.skip_reading_prefix_cache = False self._verify_step_pooling(pooler_config, valid_parameters) def _verify_step_pooling( self, pooler_config: PoolerConfig, valid_parameters: list[str], ): step_pooling_parameters = ["step_tag_id", "returned_token_ids"] if pooler_config.tok_pooling_type != "STEP": invalid_parameters = [] for k in step_pooling_parameters: if getattr(self, k, None) is not None: invalid_parameters.append(k) if invalid_parameters: raise VLLMValidationError( f"Task {self.task} only supports {valid_parameters} " f"parameters, does not support " f"{invalid_parameters} parameters" ) else: for k in step_pooling_parameters: if getattr(pooler_config, k, None) is None: continue if getattr(self, k, None) is None: setattr(self, k, getattr(pooler_config, k)) def _set_default_parameters(self, model_config: ModelConfig): if self.task in ["embed", "token_embed"]: if self.use_activation is None: self.use_activation = True if self.dimensions is not None: dimensions = self.dimensions model_name = model_config.served_model_name embedding_size = model_config.embedding_size valid_range = f"[1, {embedding_size}]" dimensions_in_range = 1 <= dimensions <= embedding_size if not model_config.is_matryoshka: raise VLLMValidationError( f"Model {model_name!r} does not support Matryoshka " f"embeddings; dimensions must be unset " f"(received dimensions={dimensions})." ) if not dimensions_in_range: raise VLLMValidationError( f"Model {model_name!r} only supports dimensions in " f"range {valid_range}, got {dimensions}." ) mds = model_config.matryoshka_dimensions if mds is not None and dimensions not in mds: raise VLLMValidationError( f"Model {model_name!r} only supports Matryoshka " f"dimensions {str(mds)}, got {dimensions}." ) elif self.task in ["classify", "token_classify"]: if self.use_activation is None: self.use_activation = True else: raise ValueError(f"Unknown pooling task: {self.task!r}") def _verify_valid_parameters(self): if self.task is None: raise ValueError("task must be set before verifying parameters") valid_parameters = self.valid_parameters[self.task] invalid_parameters = [] for k in self.all_parameters: if k in valid_parameters: continue if getattr(self, k, None) is not None: invalid_parameters.append(k) if invalid_parameters: raise VLLMValidationError( f"Task {self.task!r} only supports {valid_parameters} " f"parameters, does not support " f"{invalid_parameters} parameters" ) def __repr__(self) -> str: return ( f"PoolingParams(" f"task={self.task}, " f"dimensions={self.dimensions}, " f"use_activation={self.use_activation}, " f"step_tag_id={self.step_tag_id}, " f"returned_token_ids={self.returned_token_ids}, " f"requires_token_ids={self.requires_token_ids}, " f"skip_reading_prefix_cache={self.skip_reading_prefix_cache}, " f"late_interaction_params={self.late_interaction_params}, " f"late_chunking_params={self.late_chunking_params}, " f"extra_kwargs={self.extra_kwargs})" ) def __post_init__(self) -> None: check_removed_pooling_task(self.task) if self.output_kind != RequestOutputKind.FINAL_ONLY: raise VLLMValidationError( "For pooling output_kind has to be FINAL_ONLY, " f"got {self.output_kind!r}" )