Skip to content

vllm.model_executor.models.interfaces_base

T module-attribute 

T = TypeVar('T', default=Tensor)

T_co module-attribute 

T_co = TypeVar('T_co', default=Tensor, covariant=True)

_T module-attribute 

_T = TypeVar('_T', bound=type[Module])

logger module-attribute 

logger = init_logger(__name__)

VllmModel

Bases: Protocol[T_co]

The interface required for all models in vLLM.

Source code in vllm/model_executor/models/interfaces_base.py 
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
@runtime_checkable
class VllmModel(Protocol[T_co]):
    """The interface required for all models in vLLM."""

    def __init__(
        self,
        vllm_config: VllmConfig,
        prefix: str = "",
    ) -> None: ...

    def get_input_embeddings(
        self,
        input_ids: torch.Tensor,
    ) -> torch.Tensor:
        """Apply token embeddings to `input_ids`."""
        ...

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
    ) -> T_co: ...

__init__ 

__init__(vllm_config: VllmConfig, prefix: str = '') -> None
Source code in vllm/model_executor/models/interfaces_base.py 
44
45
46
47
48
def __init__(
    self,
    vllm_config: VllmConfig,
    prefix: str = "",
) -> None: ...

forward 

forward(input_ids: Tensor, positions: Tensor) -> T_co
Source code in vllm/model_executor/models/interfaces_base.py 
57
58
59
60
61
def forward(
    self,
    input_ids: torch.Tensor,
    positions: torch.Tensor,
) -> T_co: ...

get_input_embeddings 

get_input_embeddings(input_ids: Tensor) -> Tensor

Apply token embeddings to input_ids.

Source code in vllm/model_executor/models/interfaces_base.py 
50
51
52
53
54
55
def get_input_embeddings(
    self,
    input_ids: torch.Tensor,
) -> torch.Tensor:
    """Apply token embeddings to `input_ids`."""
    ...

VllmModelForPooling

Bases: VllmModel[T_co], Protocol[T_co]

The interface required for all pooling models in vLLM.

Source code in vllm/model_executor/models/interfaces_base.py 
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
@runtime_checkable
class VllmModelForPooling(VllmModel[T_co], Protocol[T_co]):
    """The interface required for all pooling models in vLLM."""

    is_pooling_model: ClassVar[Literal[True]] = True
    """
    A flag that indicates this model supports pooling.

    Note:
        There is no need to redefine this flag if this class is in the
        MRO of your model class.
    """

    default_pooling_type: ClassVar[str] = "LAST"
    """
    Indicates the
    [vllm.model_executor.layers.pooler.PoolerConfig.pooling_type][]
    to use by default.

    You can use the
    [vllm.model_executor.models.interfaces_base.default_pooling_type][]
    decorator to conveniently set this field.
    """

    pooler: Pooler
    """The pooler is only called on TP rank 0."""

default_pooling_type class-attribute 

default_pooling_type: str = 'LAST'

Indicates the vllm.model_executor.layers.pooler.PoolerConfig.pooling_type to use by default.

You can use the vllm.model_executor.models.interfaces_base.default_pooling_type decorator to conveniently set this field.

is_pooling_model class-attribute 

is_pooling_model: Literal[True] = True

A flag that indicates this model supports pooling.

Note

There is no need to redefine this flag if this class is in the MRO of your model class.

pooler instance-attribute 

pooler: Pooler

The pooler is only called on TP rank 0.

VllmModelForTextGeneration

Bases: VllmModel[T], Protocol[T]

The interface required for all generative models in vLLM.

Source code in vllm/model_executor/models/interfaces_base.py 
118
119
120
121
122
123
124
125
126
127
@runtime_checkable
class VllmModelForTextGeneration(VllmModel[T], Protocol[T]):
    """The interface required for all generative models in vLLM."""

    def compute_logits(
        self,
        hidden_states: T,
    ) -> T | None:
        """Return `None` if TP rank > 0."""
        ...

compute_logits 

compute_logits(hidden_states: T) -> T | None

Return None if TP rank > 0.

Source code in vllm/model_executor/models/interfaces_base.py 
122
123
124
125
126
127
def compute_logits(
    self,
    hidden_states: T,
) -> T | None:
    """Return `None` if TP rank > 0."""
    ...

_check_vllm_model_forward 

_check_vllm_model_forward(
    model: type[object] | object,
) -> bool
Source code in vllm/model_executor/models/interfaces_base.py 
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
def _check_vllm_model_forward(model: type[object] | object) -> bool:
    model_forward = getattr(model, "forward", None)
    if not callable(model_forward):
        return False

    vllm_kws = ("input_ids", "positions")
    missing_kws = tuple(kw for kw in vllm_kws if not supports_kw(model_forward, kw))

    if missing_kws and (isinstance(model, type) and issubclass(model, nn.Module)):
        logger.warning(
            "The model (%s) is missing "
            "vLLM-specific keywords from its `forward` method: %s",
            model,
            missing_kws,
        )

    return len(missing_kws) == 0

_check_vllm_model_get_input_embeddings 

_check_vllm_model_get_input_embeddings(
    model: type[object] | object,
) -> bool
Source code in vllm/model_executor/models/interfaces_base.py 
69
70
71
72
73
74
75
76
77
78
def _check_vllm_model_get_input_embeddings(model: type[object] | object) -> bool:
    model_get_input_embeddings = getattr(model, "get_input_embeddings", None)
    if not callable(model_get_input_embeddings):
        logger.warning(
            "The model (%s) is missing the `get_input_embeddings` method.",
            model,
        )
        return False

    return True

_check_vllm_model_init 

_check_vllm_model_init(
    model: type[object] | object,
) -> bool
Source code in vllm/model_executor/models/interfaces_base.py 
64
65
66
def _check_vllm_model_init(model: type[object] | object) -> bool:
    model_init = model.__init__
    return supports_kw(model_init, "vllm_config")

default_pooling_type 

default_pooling_type(pooling_type: str)

Decorator to set VllmModelForPooling.default_pooling_type.

Source code in vllm/model_executor/models/interfaces_base.py 
200
201
202
203
204
205
206
207
def default_pooling_type(pooling_type: str):
    """Decorator to set `VllmModelForPooling.default_pooling_type`."""

    def func(model: _T) -> _T:
        model.default_pooling_type = pooling_type  # type: ignore
        return model

    return func

get_default_pooling_type 

get_default_pooling_type(
    model: type[object] | object,
) -> str
Source code in vllm/model_executor/models/interfaces_base.py 
210
211
def get_default_pooling_type(model: type[object] | object) -> str:
    return getattr(model, "default_pooling_type", "LAST")

is_pooling_model 

is_pooling_model(
    model: type[object],
) -> TypeIs[type[VllmModelForPooling]]

is_pooling_model(
    model: object,
) -> TypeIs[VllmModelForPooling]

is_pooling_model(
    model: type[object] | object,
) -> (
    TypeIs[type[VllmModelForPooling]]
    | TypeIs[VllmModelForPooling]
)
Source code in vllm/model_executor/models/interfaces_base.py 
188
189
190
191
192
193
194
def is_pooling_model(
    model: type[object] | object,
) -> TypeIs[type[VllmModelForPooling]] | TypeIs[VllmModelForPooling]:
    if not is_vllm_model(model):
        return False

    return getattr(model, "is_pooling_model", False)

is_text_generation_model 

is_text_generation_model(
    model: type[object],
) -> TypeIs[type[VllmModelForTextGeneration]]

is_text_generation_model(
    model: object,
) -> TypeIs[VllmModelForTextGeneration]

is_text_generation_model(
    model: type[object] | object,
) -> (
    TypeIs[type[VllmModelForTextGeneration]]
    | TypeIs[VllmModelForTextGeneration]
)
Source code in vllm/model_executor/models/interfaces_base.py 
140
141
142
143
144
145
146
147
148
149
def is_text_generation_model(
    model: type[object] | object,
) -> TypeIs[type[VllmModelForTextGeneration]] | TypeIs[VllmModelForTextGeneration]:
    if not is_vllm_model(model):
        return False

    if isinstance(model, type):
        return isinstance(model, VllmModelForTextGeneration)

    return isinstance(model, VllmModelForTextGeneration)

is_vllm_model 

is_vllm_model(
    model: type[object],
) -> TypeIs[type[VllmModel]]

is_vllm_model(model: object) -> TypeIs[VllmModel]

is_vllm_model(
    model: type[object] | object,
) -> TypeIs[type[VllmModel]] | TypeIs[VllmModel]
Source code in vllm/model_executor/models/interfaces_base.py 
108
109
110
111
112
113
114
115
def is_vllm_model(
    model: type[object] | object,
) -> TypeIs[type[VllmModel]] | TypeIs[VllmModel]:
    return (
        _check_vllm_model_init(model)
        and _check_vllm_model_get_input_embeddings(model)
        and _check_vllm_model_forward(model)
    )