286 lines
11 KiB
Python
286 lines
11 KiB
Python
# 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}"
|
|
)
|