first commit
This commit is contained in:
959
vllm/model_executor/models/interfaces.py
Normal file
959
vllm/model_executor/models/interfaces.py
Normal file
@@ -0,0 +1,959 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections.abc import Iterable, Mapping, MutableSequence
|
||||
from typing import (TYPE_CHECKING, ClassVar, Literal, Optional, Protocol,
|
||||
Union, overload, runtime_checkable)
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from transformers import PretrainedConfig
|
||||
from transformers.models.whisper.tokenization_whisper import LANGUAGES
|
||||
from typing_extensions import Self, TypeIs
|
||||
|
||||
from vllm.config import ModelConfig, SpeechToTextConfig
|
||||
from vllm.inputs import TokensPrompt
|
||||
from vllm.inputs.data import PromptType
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.quantization.base_config import (
|
||||
QuantizationConfig)
|
||||
from vllm.utils import supports_kw
|
||||
|
||||
from .interfaces_base import is_pooling_model
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.model_executor.models.utils import WeightsMapper
|
||||
from vllm.sequence import IntermediateTensors
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
MultiModalEmbeddings = Union[list[Tensor], Tensor, tuple[Tensor, ...]]
|
||||
"""
|
||||
The output embeddings must be one of the following formats:
|
||||
|
||||
- A list or tuple of 2D tensors, where each tensor corresponds to
|
||||
each input multimodal data item (e.g, image).
|
||||
- A single 3D tensor, with the batch dimension grouping the 2D tensors.
|
||||
"""
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class SupportsMultiModal(Protocol):
|
||||
"""The interface required for all multi-modal models."""
|
||||
|
||||
supports_multimodal: ClassVar[Literal[True]] = True
|
||||
"""
|
||||
A flag that indicates this model supports multi-modal inputs.
|
||||
|
||||
Note:
|
||||
There is no need to redefine this flag if this class is in the
|
||||
MRO of your model class.
|
||||
"""
|
||||
|
||||
supports_multimodal_raw_input_only: ClassVar[bool] = False
|
||||
"""
|
||||
A flag that indicates this model supports multi-modal inputs and processes
|
||||
them in their raw form and not embeddings.
|
||||
"""
|
||||
|
||||
supports_encoder_tp_data: ClassVar[bool] = False
|
||||
"""
|
||||
A flag that indicates whether this model supports
|
||||
`multimodal_config.mm_encoder_tp_mode="data"`.
|
||||
"""
|
||||
|
||||
merge_by_field_config: ClassVar[bool] = False
|
||||
"""
|
||||
A flag that indicates which implementation of
|
||||
`vllm.multimodal.utils.group_mm_kwargs_by_modality` to use.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:
|
||||
"""
|
||||
Get the placeholder text for the `i`th `modality` item in the prompt.
|
||||
"""
|
||||
...
|
||||
|
||||
def get_multimodal_embeddings(self,
|
||||
**kwargs: object) -> MultiModalEmbeddings:
|
||||
"""
|
||||
Returns multimodal embeddings generated from multimodal kwargs
|
||||
to be merged with text embeddings.
|
||||
|
||||
Note:
|
||||
The returned multimodal embeddings must be in the same order as
|
||||
the appearances of their corresponding multimodal data item in the
|
||||
input prompt.
|
||||
"""
|
||||
...
|
||||
|
||||
def get_language_model(self) -> torch.nn.Module:
|
||||
"""
|
||||
Returns the underlying language model used for text generation.
|
||||
|
||||
This is typically the `torch.nn.Module` instance responsible for
|
||||
processing the merged multimodal embeddings and producing hidden states
|
||||
|
||||
Returns:
|
||||
torch.nn.Module: The core language model component.
|
||||
"""
|
||||
...
|
||||
|
||||
def get_input_embeddings(
|
||||
self,
|
||||
input_ids: Tensor,
|
||||
multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
|
||||
) -> Tensor:
|
||||
"""
|
||||
Returns the input embeddings merged from the text embeddings from
|
||||
input_ids and the multimodal embeddings generated from multimodal
|
||||
kwargs.
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class SupportsMultiModalPruning(Protocol):
|
||||
"""The interface required for models that support returning both input
|
||||
embeddings and positions. Model may require custom positions for dynamic
|
||||
pruning of multimodal embeddings.
|
||||
"""
|
||||
supports_multimodal_pruning: ClassVar[Literal[True]] = True
|
||||
|
||||
def recompute_mrope_positions(
|
||||
self, input_ids: list[int],
|
||||
multimodal_embeddings: MultiModalEmbeddings,
|
||||
mrope_positions: torch.LongTensor, num_computed_tokens: int
|
||||
) -> tuple[MultiModalEmbeddings, Tensor, int]:
|
||||
"""
|
||||
Update part of input mrope positions (starting with
|
||||
num_computed_tokens index). Original mrope_positions are computed
|
||||
for unpruned sequence and becomes incorrect once pruning occurs,
|
||||
so once we prune media tokens we should reflect this in the
|
||||
mrope_positions before we feed it to LLM.
|
||||
|
||||
Args:
|
||||
input_ids: (N,) All input tokens of the prompt containing
|
||||
entire sequence.
|
||||
multimodal_embeddings: Tuple of multimodal embeddings that
|
||||
fits into the prefill chunk that is being processed.
|
||||
mrope_positions: Existing mrope positions (3, N) for entire
|
||||
sequence
|
||||
num_computed_tokens: A number of computed tokens so far.
|
||||
|
||||
Returns:
|
||||
Tuple of (multimodal_embeddings, mrope_positions,
|
||||
mrope_position_delta).
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
@overload
|
||||
def supports_multimodal(
|
||||
model: type[object]) -> TypeIs[type[SupportsMultiModal]]:
|
||||
...
|
||||
|
||||
|
||||
@overload
|
||||
def supports_multimodal(model: object) -> TypeIs[SupportsMultiModal]:
|
||||
...
|
||||
|
||||
|
||||
def supports_multimodal(
|
||||
model: Union[type[object], object],
|
||||
) -> Union[TypeIs[type[SupportsMultiModal]], TypeIs[SupportsMultiModal]]:
|
||||
return getattr(model, "supports_multimodal", False)
|
||||
|
||||
|
||||
def supports_multimodal_raw_input_only(
|
||||
model: Union[type[object], object]) -> bool:
|
||||
return getattr(model, "supports_multimodal_raw_input_only", False)
|
||||
|
||||
|
||||
def supports_multimodal_encoder_tp_data(
|
||||
model: Union[type[object], object]) -> bool:
|
||||
return getattr(model, "supports_encoder_tp_data", False)
|
||||
|
||||
|
||||
@overload
|
||||
def supports_multimodal_pruning(
|
||||
model: type[object]) -> TypeIs[type[SupportsMultiModalPruning]]:
|
||||
...
|
||||
|
||||
|
||||
@overload
|
||||
def supports_multimodal_pruning(
|
||||
model: object) -> TypeIs[SupportsMultiModalPruning]:
|
||||
...
|
||||
|
||||
|
||||
def supports_multimodal_pruning(
|
||||
model: Union[type[object], object],
|
||||
) -> Union[TypeIs[type[SupportsMultiModalPruning]],
|
||||
TypeIs[SupportsMultiModalPruning]]:
|
||||
return getattr(model, "supports_multimodal_pruning", False)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class SupportsScoreTemplate(Protocol):
|
||||
"""The interface required for all models that support score template."""
|
||||
|
||||
supports_score_template: ClassVar[Literal[True]] = True
|
||||
"""
|
||||
A flag that indicates this model supports score template.
|
||||
|
||||
Note:
|
||||
There is no need to redefine this flag if this class is in the
|
||||
MRO of your model class.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def get_score_template(cls, query: str, document: str) -> Optional[str]:
|
||||
"""
|
||||
Generate a full prompt by populating the score template with query and document content.
|
||||
""" # noqa: E501
|
||||
...
|
||||
|
||||
@classmethod
|
||||
def post_process_tokens(cls, prompt: TokensPrompt) -> None:
|
||||
"""
|
||||
Perform architecture-specific manipulations on the input tokens.
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
@overload
|
||||
def supports_score_template(
|
||||
model: type[object]) -> TypeIs[type[SupportsScoreTemplate]]:
|
||||
...
|
||||
|
||||
|
||||
@overload
|
||||
def supports_score_template(model: object) -> TypeIs[SupportsScoreTemplate]:
|
||||
...
|
||||
|
||||
|
||||
def supports_score_template(
|
||||
model: Union[type[object], object],
|
||||
) -> Union[TypeIs[type[SupportsScoreTemplate]], TypeIs[SupportsScoreTemplate]]:
|
||||
return getattr(model, "supports_score_template", False)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class SupportsLoRA(Protocol):
|
||||
"""The interface required for all models that support LoRA."""
|
||||
|
||||
supports_lora: ClassVar[Literal[True]] = True
|
||||
"""
|
||||
A flag that indicates this model supports LoRA.
|
||||
|
||||
Note:
|
||||
There is no need to redefine this flag if this class is in the
|
||||
MRO of your model class.
|
||||
"""
|
||||
# The `embedding_module` and `embedding_padding_modules`
|
||||
# are empty by default.
|
||||
embedding_modules: ClassVar[dict[str, str]] = {}
|
||||
embedding_padding_modules: ClassVar[list[str]] = []
|
||||
packed_modules_mapping: ClassVar[dict[str, list[str]]] = {}
|
||||
|
||||
|
||||
# We can't use runtime_checkable with ClassVar for issubclass checks
|
||||
# so we need to treat the class as an instance and use isinstance instead
|
||||
@runtime_checkable
|
||||
class _SupportsLoRAType(Protocol):
|
||||
supports_lora: Literal[True]
|
||||
|
||||
packed_modules_mapping: dict[str, list[str]]
|
||||
embedding_modules: dict[str, str]
|
||||
embedding_padding_modules: list[str]
|
||||
|
||||
|
||||
@overload
|
||||
def supports_lora(model: type[object]) -> TypeIs[type[SupportsLoRA]]:
|
||||
...
|
||||
|
||||
|
||||
@overload
|
||||
def supports_lora(model: object) -> TypeIs[SupportsLoRA]:
|
||||
...
|
||||
|
||||
|
||||
def supports_lora(
|
||||
model: Union[type[object], object],
|
||||
) -> Union[TypeIs[type[SupportsLoRA]], TypeIs[SupportsLoRA]]:
|
||||
result = _supports_lora(model)
|
||||
|
||||
if not result:
|
||||
lora_attrs = (
|
||||
"packed_modules_mapping",
|
||||
"embedding_modules",
|
||||
"embedding_padding_modules",
|
||||
)
|
||||
missing_attrs = tuple(attr for attr in lora_attrs
|
||||
if not hasattr(model, attr))
|
||||
|
||||
if getattr(model, "supports_lora", False):
|
||||
if missing_attrs:
|
||||
logger.warning(
|
||||
"The model (%s) sets `supports_lora=True`, "
|
||||
"but is missing LoRA-specific attributes: %s",
|
||||
model,
|
||||
missing_attrs,
|
||||
)
|
||||
else:
|
||||
if not missing_attrs:
|
||||
logger.warning(
|
||||
"The model (%s) contains all LoRA-specific attributes, "
|
||||
"but does not set `supports_lora=True`.", model)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def _supports_lora(model: Union[type[object], object]) -> bool:
|
||||
if isinstance(model, type):
|
||||
return isinstance(model, _SupportsLoRAType)
|
||||
|
||||
return isinstance(model, SupportsLoRA)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class SupportsPP(Protocol):
|
||||
"""The interface required for all models that support pipeline parallel."""
|
||||
|
||||
supports_pp: ClassVar[Literal[True]] = True
|
||||
"""
|
||||
A flag that indicates this model supports pipeline parallel.
|
||||
|
||||
Note:
|
||||
There is no need to redefine this flag if this class is in the
|
||||
MRO of your model class.
|
||||
"""
|
||||
|
||||
def make_empty_intermediate_tensors(
|
||||
self,
|
||||
batch_size: int,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
) -> "IntermediateTensors":
|
||||
"""Called when PP rank > 0 for profiling purposes."""
|
||||
...
|
||||
|
||||
def forward(
|
||||
self,
|
||||
*,
|
||||
intermediate_tensors: Optional["IntermediateTensors"],
|
||||
) -> Union[Tensor, "IntermediateTensors"]:
|
||||
"""
|
||||
Accept [`IntermediateTensors`][vllm.sequence.IntermediateTensors] when
|
||||
PP rank > 0.
|
||||
|
||||
Return [`IntermediateTensors`][vllm.sequence.IntermediateTensors] only
|
||||
for the last PP rank.
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
# We can't use runtime_checkable with ClassVar for issubclass checks
|
||||
# so we need to treat the class as an instance and use isinstance instead
|
||||
@runtime_checkable
|
||||
class _SupportsPPType(Protocol):
|
||||
supports_pp: Literal[True]
|
||||
|
||||
def make_empty_intermediate_tensors(
|
||||
self,
|
||||
batch_size: int,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
) -> "IntermediateTensors":
|
||||
...
|
||||
|
||||
def forward(
|
||||
self,
|
||||
*,
|
||||
intermediate_tensors: Optional["IntermediateTensors"],
|
||||
) -> Union[Tensor, "IntermediateTensors"]:
|
||||
...
|
||||
|
||||
|
||||
@overload
|
||||
def supports_pp(model: type[object]) -> TypeIs[type[SupportsPP]]:
|
||||
...
|
||||
|
||||
|
||||
@overload
|
||||
def supports_pp(model: object) -> TypeIs[SupportsPP]:
|
||||
...
|
||||
|
||||
|
||||
def supports_pp(
|
||||
model: Union[type[object], object],
|
||||
) -> Union[bool, TypeIs[type[SupportsPP]], TypeIs[SupportsPP]]:
|
||||
supports_attributes = _supports_pp_attributes(model)
|
||||
supports_inspect = _supports_pp_inspect(model)
|
||||
|
||||
if supports_attributes and not supports_inspect:
|
||||
logger.warning(
|
||||
"The model (%s) sets `supports_pp=True`, but does not accept "
|
||||
"`intermediate_tensors` in its `forward` method", model)
|
||||
|
||||
if not supports_attributes:
|
||||
pp_attrs = ("make_empty_intermediate_tensors", )
|
||||
missing_attrs = tuple(attr for attr in pp_attrs
|
||||
if not hasattr(model, attr))
|
||||
|
||||
if getattr(model, "supports_pp", False):
|
||||
if missing_attrs:
|
||||
logger.warning(
|
||||
"The model (%s) sets `supports_pp=True`, "
|
||||
"but is missing PP-specific attributes: %s",
|
||||
model,
|
||||
missing_attrs,
|
||||
)
|
||||
else:
|
||||
if not missing_attrs:
|
||||
logger.warning(
|
||||
"The model (%s) contains all PP-specific attributes, "
|
||||
"but does not set `supports_pp=True`.", model)
|
||||
|
||||
return supports_attributes and supports_inspect
|
||||
|
||||
|
||||
def _supports_pp_attributes(model: Union[type[object], object]) -> bool:
|
||||
if isinstance(model, type):
|
||||
return isinstance(model, _SupportsPPType)
|
||||
|
||||
return isinstance(model, SupportsPP)
|
||||
|
||||
|
||||
def _supports_pp_inspect(model: Union[type[object], object]) -> bool:
|
||||
model_forward = getattr(model, "forward", None)
|
||||
if not callable(model_forward):
|
||||
return False
|
||||
|
||||
return supports_kw(model_forward, "intermediate_tensors")
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class HasInnerState(Protocol):
|
||||
"""The interface required for all models that has inner state."""
|
||||
|
||||
has_inner_state: ClassVar[Literal[True]] = True
|
||||
"""
|
||||
A flag that indicates this model has inner state.
|
||||
Models that has inner state usually need access to the scheduler_config
|
||||
for max_num_seqs, etc. True for e.g. both Mamba and Jamba.
|
||||
"""
|
||||
|
||||
|
||||
@overload
|
||||
def has_inner_state(model: object) -> TypeIs[HasInnerState]:
|
||||
...
|
||||
|
||||
|
||||
@overload
|
||||
def has_inner_state(model: type[object]) -> TypeIs[type[HasInnerState]]:
|
||||
...
|
||||
|
||||
|
||||
def has_inner_state(
|
||||
model: Union[type[object], object]
|
||||
) -> Union[TypeIs[type[HasInnerState]], TypeIs[HasInnerState]]:
|
||||
return getattr(model, "has_inner_state", False)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class IsAttentionFree(Protocol):
|
||||
"""The interface required for all models like Mamba that lack attention,
|
||||
but do have state whose size is constant wrt the number of tokens."""
|
||||
|
||||
is_attention_free: ClassVar[Literal[True]] = True
|
||||
"""
|
||||
A flag that indicates this model has no attention.
|
||||
Used for block manager and attention backend selection.
|
||||
True for Mamba but not Jamba.
|
||||
"""
|
||||
|
||||
|
||||
@overload
|
||||
def is_attention_free(model: object) -> TypeIs[IsAttentionFree]:
|
||||
...
|
||||
|
||||
|
||||
@overload
|
||||
def is_attention_free(model: type[object]) -> TypeIs[type[IsAttentionFree]]:
|
||||
...
|
||||
|
||||
|
||||
def is_attention_free(
|
||||
model: Union[type[object], object]
|
||||
) -> Union[TypeIs[type[IsAttentionFree]], TypeIs[IsAttentionFree]]:
|
||||
return getattr(model, "is_attention_free", False)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class IsHybrid(Protocol):
|
||||
"""The interface required for all models like Jamba that have both
|
||||
attention and mamba blocks, indicates that
|
||||
hf_config has 'layers_block_type'"""
|
||||
|
||||
is_hybrid: ClassVar[Literal[True]] = True
|
||||
"""
|
||||
A flag that indicates this model has both mamba and attention blocks
|
||||
, also indicates that the model's hf_config has
|
||||
'layers_block_type' """
|
||||
|
||||
@classmethod
|
||||
def get_mamba_state_shape_from_config(
|
||||
cls,
|
||||
vllm_config: "VllmConfig",
|
||||
use_v1: bool = True,
|
||||
) -> tuple[tuple[int, int], tuple[int, int, int]]:
|
||||
"""Calculate shapes for Mamba's convolutional and state caches.
|
||||
|
||||
Args:
|
||||
vllm_config: vLLM config
|
||||
use_v1: Get shapes for V1 (or V0)
|
||||
|
||||
Returns:
|
||||
Tuple containing:
|
||||
- conv_state_shape: Shape for convolutional state cache
|
||||
- temporal_state_shape: Shape for state space model cache
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
@overload
|
||||
def is_hybrid(model: object) -> TypeIs[IsHybrid]:
|
||||
...
|
||||
|
||||
|
||||
@overload
|
||||
def is_hybrid(model: type[object]) -> TypeIs[type[IsHybrid]]:
|
||||
...
|
||||
|
||||
|
||||
def is_hybrid(
|
||||
model: Union[type[object], object]
|
||||
) -> Union[TypeIs[type[IsHybrid]], TypeIs[IsHybrid]]:
|
||||
return getattr(model, "is_hybrid", False)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class MixtureOfExperts(Protocol):
|
||||
"""
|
||||
Check if the model is a mixture of experts (MoE) model.
|
||||
"""
|
||||
|
||||
expert_weights: MutableSequence[Iterable[Tensor]]
|
||||
"""
|
||||
Expert weights saved in this rank.
|
||||
|
||||
The first dimension is the layer, and the second dimension is different
|
||||
parameters in the layer, e.g. up/down projection weights.
|
||||
"""
|
||||
|
||||
num_moe_layers: int
|
||||
"""Number of MoE layers in this model."""
|
||||
|
||||
num_expert_groups: int
|
||||
"""Number of expert groups in this model."""
|
||||
|
||||
num_logical_experts: int
|
||||
"""Number of logical experts in this model."""
|
||||
|
||||
num_physical_experts: int
|
||||
"""Number of physical experts in this model."""
|
||||
|
||||
num_local_physical_experts: int
|
||||
"""Number of local physical experts in this model."""
|
||||
|
||||
num_routed_experts: int
|
||||
"""Number of routed experts in this model."""
|
||||
|
||||
num_shared_experts: int
|
||||
"""Number of shared experts in this model."""
|
||||
|
||||
num_redundant_experts: int
|
||||
"""Number of redundant experts in this model."""
|
||||
|
||||
def set_eplb_state(
|
||||
self,
|
||||
expert_load_view: Tensor,
|
||||
logical_to_physical_map: Tensor,
|
||||
logical_replica_count: Tensor,
|
||||
) -> None:
|
||||
"""
|
||||
Register the EPLB state in the MoE model.
|
||||
|
||||
Since these are views of the actual EPLB state, any changes made by
|
||||
the EPLB algorithm are automatically reflected in the model's behavior
|
||||
without requiring additional method calls to set new states.
|
||||
|
||||
You should also collect model's `expert_weights` here instead of in
|
||||
the weight loader, since after initial weight loading, further
|
||||
processing like quantization may be applied to the weights.
|
||||
|
||||
Args:
|
||||
expert_load_view: A view of the expert load metrics tensor.
|
||||
logical_to_physical_map: Mapping from logical to physical experts.
|
||||
logical_replica_count: Count of replicas for each logical expert.
|
||||
"""
|
||||
...
|
||||
|
||||
def update_physical_experts_metadata(
|
||||
self,
|
||||
num_physical_experts: int,
|
||||
num_local_physical_experts: int,
|
||||
) -> None:
|
||||
...
|
||||
|
||||
|
||||
def is_mixture_of_experts(model: object) -> TypeIs[MixtureOfExperts]:
|
||||
return isinstance(model, MixtureOfExperts)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class HasNoOps(Protocol):
|
||||
has_noops: ClassVar[Literal[True]] = True
|
||||
|
||||
|
||||
@overload
|
||||
def has_noops(model: object) -> TypeIs[HasNoOps]:
|
||||
...
|
||||
|
||||
|
||||
@overload
|
||||
def has_noops(model: type[object]) -> TypeIs[type[HasNoOps]]:
|
||||
...
|
||||
|
||||
|
||||
def has_noops(
|
||||
model: Union[type[object], object]
|
||||
) -> Union[TypeIs[type[HasNoOps]], TypeIs[HasNoOps]]:
|
||||
return getattr(model, "has_noops", False)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class SupportsCrossEncoding(Protocol):
|
||||
"""The interface required for all models that support cross encoding."""
|
||||
|
||||
supports_cross_encoding: ClassVar[Literal[True]] = True
|
||||
|
||||
|
||||
@overload
|
||||
def supports_cross_encoding(
|
||||
model: type[object]) -> TypeIs[type[SupportsCrossEncoding]]:
|
||||
...
|
||||
|
||||
|
||||
@overload
|
||||
def supports_cross_encoding(model: object) -> TypeIs[SupportsCrossEncoding]:
|
||||
...
|
||||
|
||||
|
||||
def _supports_cross_encoding(
|
||||
model: Union[type[object], object],
|
||||
) -> Union[TypeIs[type[SupportsCrossEncoding]], TypeIs[SupportsCrossEncoding]]:
|
||||
return getattr(model, "supports_cross_encoding", False)
|
||||
|
||||
|
||||
def supports_cross_encoding(
|
||||
model: Union[type[object], object],
|
||||
) -> Union[TypeIs[type[SupportsCrossEncoding]], TypeIs[SupportsCrossEncoding]]:
|
||||
return is_pooling_model(model) and _supports_cross_encoding(model)
|
||||
|
||||
|
||||
class SupportsQuant:
|
||||
"""The interface required for all models that support quantization."""
|
||||
|
||||
hf_to_vllm_mapper: ClassVar[Optional["WeightsMapper"]] = None
|
||||
packed_modules_mapping: ClassVar[Optional[dict[str, list[str]]]] = None
|
||||
quant_config: Optional[QuantizationConfig] = None
|
||||
|
||||
def __new__(cls, *args, **kwargs) -> Self:
|
||||
instance = super().__new__(cls)
|
||||
|
||||
# find config passed in arguments
|
||||
quant_config = cls._find_quant_config(*args, **kwargs)
|
||||
if quant_config is not None:
|
||||
|
||||
# attach config to model for general use
|
||||
instance.quant_config = quant_config
|
||||
|
||||
# apply model mappings to config for proper config-model matching
|
||||
if (hf_to_vllm_mapper := instance.hf_to_vllm_mapper) is not None:
|
||||
instance.quant_config.apply_vllm_mapper(hf_to_vllm_mapper)
|
||||
if instance.packed_modules_mapping is not None:
|
||||
instance.quant_config.packed_modules_mapping.update(
|
||||
instance.packed_modules_mapping)
|
||||
|
||||
return instance
|
||||
|
||||
@staticmethod
|
||||
def _find_quant_config(*args, **kwargs) -> Optional[QuantizationConfig]:
|
||||
"""Find quant config passed through model constructor args"""
|
||||
from vllm.config import VllmConfig # avoid circular import
|
||||
|
||||
args_values = list(args) + list(kwargs.values())
|
||||
for arg in args_values:
|
||||
if isinstance(arg, VllmConfig):
|
||||
return arg.quant_config
|
||||
|
||||
if isinstance(arg, QuantizationConfig):
|
||||
return arg
|
||||
|
||||
return None
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class SupportsTranscription(Protocol):
|
||||
"""The interface required for all models that support transcription."""
|
||||
# Mapping from ISO639_1 language codes: language names
|
||||
supported_languages: ClassVar[Mapping[str, str]]
|
||||
|
||||
supports_transcription: ClassVar[Literal[True]] = True
|
||||
|
||||
supports_transcription_only: ClassVar[bool] = False
|
||||
"""
|
||||
Transcription models can opt out of text generation by setting this to
|
||||
`True`.
|
||||
"""
|
||||
|
||||
def __init_subclass__(cls, **kwargs):
|
||||
super().__init_subclass__(**kwargs)
|
||||
# language codes in supported_languages
|
||||
# that don't exist in the full language map
|
||||
invalid = set(cls.supported_languages) - set(LANGUAGES.keys())
|
||||
if invalid:
|
||||
raise ValueError(
|
||||
f"{cls.__name__}.supported_languages contains invalid "
|
||||
f"language codes: {sorted(invalid)}\n. "
|
||||
f"Valid choices are: {sorted(LANGUAGES.keys())}")
|
||||
|
||||
@classmethod
|
||||
def get_generation_prompt(cls, audio: np.ndarray,
|
||||
stt_config: SpeechToTextConfig,
|
||||
model_config: ModelConfig,
|
||||
language: Optional[str],
|
||||
task_type: Literal["transcribe", "translate"],
|
||||
request_prompt: str,
|
||||
to_language: Optional[str]) -> PromptType:
|
||||
"""Get the prompt for the ASR model.
|
||||
The model has control over the construction, as long as it
|
||||
returns a valid PromptType."""
|
||||
...
|
||||
|
||||
@classmethod
|
||||
def get_other_languages(cls) -> Mapping[str, str]:
|
||||
# other possible language codes from the whisper map
|
||||
return {
|
||||
k: v
|
||||
for k, v in LANGUAGES.items() if k not in cls.supported_languages
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def validate_language(cls, language: Optional[str]) -> Optional[str]:
|
||||
"""
|
||||
Ensure the language specified in the transcription request
|
||||
is a valid ISO 639-1 language code. If the request language is
|
||||
valid, but not natively supported by the model, trigger a
|
||||
warning (but not an exception).
|
||||
"""
|
||||
if language is None or language in cls.supported_languages:
|
||||
return language
|
||||
elif language in cls.get_other_languages():
|
||||
logger.warning(
|
||||
"Language %r is not natively supported by %s; "
|
||||
"results may be less accurate. Supported languages: %r",
|
||||
language,
|
||||
cls.__name__,
|
||||
list(cls.supported_languages.keys()),
|
||||
)
|
||||
return language
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported language: {language!r}. Must be one of "
|
||||
f"{list(cls.supported_languages.keys())}.")
|
||||
|
||||
@classmethod
|
||||
def get_speech_to_text_config(
|
||||
cls, model_config: ModelConfig,
|
||||
task_type: Literal["transcribe",
|
||||
"translate"]) -> SpeechToTextConfig:
|
||||
"""Get the speech to text config for the ASR model."""
|
||||
...
|
||||
|
||||
@classmethod
|
||||
def get_num_audio_tokens(cls, audio_duration_s: float,
|
||||
stt_config: SpeechToTextConfig,
|
||||
model_config: ModelConfig) -> Optional[int]:
|
||||
"""
|
||||
Map from audio duration to number of audio tokens produced by the ASR
|
||||
model, without running a forward pass.
|
||||
This is used for estimating the amount of processing for this audio.
|
||||
"""
|
||||
return None
|
||||
|
||||
|
||||
@overload
|
||||
def supports_transcription(
|
||||
model: type[object]) -> TypeIs[type[SupportsTranscription]]:
|
||||
...
|
||||
|
||||
|
||||
@overload
|
||||
def supports_transcription(model: object) -> TypeIs[SupportsTranscription]:
|
||||
...
|
||||
|
||||
|
||||
def supports_transcription(
|
||||
model: Union[type[object], object],
|
||||
) -> Union[TypeIs[type[SupportsTranscription]], TypeIs[SupportsTranscription]]:
|
||||
return getattr(model, "supports_transcription", False)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class SupportsV0Only(Protocol):
|
||||
"""Models with this interface are not compatible with V1 vLLM."""
|
||||
|
||||
supports_v0_only: ClassVar[Literal[True]] = True
|
||||
|
||||
|
||||
@overload
|
||||
def supports_v0_only(model: type[object]) -> TypeIs[type[SupportsV0Only]]:
|
||||
...
|
||||
|
||||
|
||||
@overload
|
||||
def supports_v0_only(model: object) -> TypeIs[SupportsV0Only]:
|
||||
...
|
||||
|
||||
|
||||
def supports_v0_only(
|
||||
model: Union[type[object], object],
|
||||
) -> Union[TypeIs[type[SupportsV0Only]], TypeIs[SupportsV0Only]]:
|
||||
return getattr(model, "supports_v0_only", False)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class SupportsEagle3(Protocol):
|
||||
"""The interface required for models that support
|
||||
EAGLE3 speculative decoding."""
|
||||
|
||||
supports_eagle3: ClassVar[Literal[True]] = True
|
||||
"""
|
||||
A flag that indicates this model supports EAGLE3
|
||||
speculative decoding.
|
||||
|
||||
Note:
|
||||
There is no need to redefine this flag if this class is in the
|
||||
MRO of your model class.
|
||||
"""
|
||||
|
||||
def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
|
||||
"""
|
||||
Set which layers should output auxiliary
|
||||
hidden states for EAGLE3.
|
||||
|
||||
Args:
|
||||
layers: Tuple of layer indices that should output auxiliary
|
||||
hidden states.
|
||||
"""
|
||||
...
|
||||
|
||||
def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
|
||||
"""
|
||||
Get the layer indices that should output auxiliary hidden states
|
||||
for EAGLE3.
|
||||
|
||||
Returns:
|
||||
Tuple of layer indices for auxiliary hidden state outputs.
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
@overload
|
||||
def supports_eagle3(model: type[object]) -> TypeIs[type[SupportsEagle3]]:
|
||||
...
|
||||
|
||||
|
||||
@overload
|
||||
def supports_eagle3(model: object) -> TypeIs[SupportsEagle3]:
|
||||
...
|
||||
|
||||
|
||||
def supports_eagle3(
|
||||
model: Union[type[object], object],
|
||||
) -> Union[TypeIs[type[SupportsEagle3]], TypeIs[SupportsEagle3]]:
|
||||
return isinstance(model, SupportsEagle3)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class SupportsMRoPE(Protocol):
|
||||
"""The interface required for all models that support M-RoPE."""
|
||||
|
||||
supports_mrope: ClassVar[Literal[True]] = True
|
||||
"""
|
||||
A flag that indicates this model supports M-RoPE.
|
||||
|
||||
Note:
|
||||
There is no need to redefine this flag if this class is in the
|
||||
MRO of your model class.
|
||||
"""
|
||||
|
||||
def get_mrope_input_positions(
|
||||
self,
|
||||
input_tokens: list[int],
|
||||
hf_config: PretrainedConfig,
|
||||
image_grid_thw: Optional[Union[list[list[int]], torch.Tensor]],
|
||||
video_grid_thw: Optional[Union[list[list[int]], torch.Tensor]],
|
||||
second_per_grid_ts: Optional[list[float]] = None,
|
||||
context_len: int = 0,
|
||||
seq_len: Optional[int] = None,
|
||||
audio_feature_lengths: Optional[torch.Tensor] = None,
|
||||
use_audio_in_video: bool = False,
|
||||
) -> tuple[torch.Tensor, int]:
|
||||
"""
|
||||
Get M-RoPE input positions and delta value for this specific model.
|
||||
|
||||
This method should be implemented by each model that supports M-RoPE
|
||||
to provide model-specific logic for computing input positions.
|
||||
|
||||
Args:
|
||||
input_tokens: List of input token IDs
|
||||
hf_config: HuggingFace model configuration
|
||||
image_grid_thw: Image grid dimensions (t, h, w)
|
||||
video_grid_thw: Video grid dimensions (t, h, w)
|
||||
second_per_grid_ts: Seconds per grid timestep for videos
|
||||
context_len: Context length
|
||||
seq_len: Sequence length
|
||||
audio_feature_lengths: Audio feature lengths for multimodal models
|
||||
use_audio_in_video: Whether to use audio in video for interleaving
|
||||
|
||||
Returns:
|
||||
Tuple of (llm_positions, mrope_position_delta)
|
||||
- llm_positions: Tensor of shape [3, num_tokens]
|
||||
with T/H/W positions
|
||||
- mrope_position_delta: Delta for position calculations
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
@overload
|
||||
def supports_mrope(model: type[object]) -> TypeIs[type[SupportsMRoPE]]:
|
||||
...
|
||||
|
||||
|
||||
@overload
|
||||
def supports_mrope(model: object) -> TypeIs[SupportsMRoPE]:
|
||||
...
|
||||
|
||||
|
||||
def supports_mrope(
|
||||
model: Union[type[object], object],
|
||||
) -> Union[TypeIs[type[SupportsMRoPE]], TypeIs[SupportsMRoPE]]:
|
||||
return isinstance(model, SupportsMRoPE)
|
||||
Reference in New Issue
Block a user