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 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
@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 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 
39
40
41
42
43
44
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 
46
47
48
49
50
51
def forward(
    self,
    input_ids: torch.Tensor,
    positions: torch.Tensor,
) -> T_co:
    ...

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 
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
@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 
 96
 97
 98
 99
100
101
102
103
104
105
106
@runtime_checkable
class VllmModelForTextGeneration(VllmModel[T], Protocol[T]):
    """The interface required for all generative models in vLLM."""

    def compute_logits(
        self,
        hidden_states: T,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[T]:
        """Return `None` if TP rank > 0."""
        ...

compute_logits 

compute_logits(
    hidden_states: T, sampling_metadata: SamplingMetadata
) -> Optional[T]

Return None if TP rank > 0.

Source code in vllm/model_executor/models/interfaces_base.py 
100
101
102
103
104
105
106
def compute_logits(
    self,
    hidden_states: T,
    sampling_metadata: SamplingMetadata,
) -> Optional[T]:
    """Return `None` if TP rank > 0."""
    ...

_check_vllm_model_forward 

_check_vllm_model_forward(
    model: Union[type[object], object],
) -> bool
Source code in vllm/model_executor/models/interfaces_base.py 
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
def _check_vllm_model_forward(model: Union[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_init 

_check_vllm_model_init(
    model: Union[type[object], object],
) -> bool
Source code in vllm/model_executor/models/interfaces_base.py 
54
55
56
def _check_vllm_model_init(model: Union[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 
184
185
186
187
188
189
190
191
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: Union[type[object], object],
) -> str
Source code in vllm/model_executor/models/interfaces_base.py 
194
195
def get_default_pooling_type(model: Union[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: Union[type[object], object],
) -> Union[
    TypeIs[type[VllmModelForPooling]],
    TypeIs[VllmModelForPooling],
]
Source code in vllm/model_executor/models/interfaces_base.py 
172
173
174
175
176
177
178
def is_pooling_model(
    model: Union[type[object], object],
) -> Union[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: Union[type[object], object],
) -> Union[
    TypeIs[type[VllmModelForTextGeneration]],
    TypeIs[VllmModelForTextGeneration],
]
Source code in vllm/model_executor/models/interfaces_base.py 
121
122
123
124
125
126
127
128
129
130
131
def is_text_generation_model(
    model: Union[type[object], object],
) -> Union[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: Union[type[object], object],
) -> Union[TypeIs[type[VllmModel]], TypeIs[VllmModel]]
Source code in vllm/model_executor/models/interfaces_base.py 
90
91
92
93
def is_vllm_model(
    model: Union[type[object], object],
) -> Union[TypeIs[type[VllmModel]], TypeIs[VllmModel]]:
    return _check_vllm_model_init(model) and _check_vllm_model_forward(model)