interfaces.py 38.5 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
from collections.abc import Callable, Iterable, Mapping, MutableSequence
5
from contextlib import contextmanager, nullcontext
6
7
8
9
10
from typing import (
    TYPE_CHECKING,
    ClassVar,
    Literal,
    Protocol,
11
    TypeAlias,
12
13
14
    overload,
    runtime_checkable,
)
15

16
import numpy as np
17
import torch
18
import torch.nn as nn
19
from torch import Tensor
20
from transformers.models.whisper.tokenization_whisper import LANGUAGES
21
from typing_extensions import Self, TypeIs
22

23
from vllm.config import ModelConfig, SpeechToTextConfig
24
from vllm.inputs import TokensPrompt
25
from vllm.inputs.data import PromptType
26
from vllm.logger import init_logger
27
from vllm.model_executor.layers.quantization import QuantizationConfig
28
from vllm.utils.func_utils import supports_kw
29

30
from .interfaces_base import VllmModel, is_pooling_model
31

32
if TYPE_CHECKING:
33
    from vllm.config import VllmConfig
34
    from vllm.model_executor.models.utils import WeightsMapper
35
    from vllm.multimodal.inputs import MultiModalFeatureSpec
36
    from vllm.multimodal.registry import _ProcessorFactories
37
    from vllm.sequence import IntermediateTensors
38
39
40
else:
    VllmConfig = object
    WeightsMapper = object
41
    MultiModalFeatureSpec = object
42
    _ProcessorFactories = object
43
    IntermediateTensors = object
44

45
46
logger = init_logger(__name__)

47
MultiModalEmbeddings: TypeAlias = list[Tensor] | Tensor | tuple[Tensor, ...]
48
49
50
51
52
53
54
"""
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.
"""
55

56

57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
def _require_is_multimodal(is_multimodal: Tensor | None) -> Tensor:
    """
    A helper function to be used in the context of
    [vllm.model_executor.models.interfaces.SupportsMultiModal.embed_input_ids][]
    to provide a better error message.
    """
    if is_multimodal is None:
        raise ValueError(
            "`embed_input_ids` now requires `is_multimodal` arg, "
            "please update your model runner according to "
            "https://github.com/vllm-project/vllm/pull/16229."
        )

    return is_multimodal


73
74
75
76
77
78
79
80
81
class LMMissingLayer(nn.Module):
    def make_empty_intermediate_tensors(self, *args, **kwargs):
        raise RuntimeError("This module should not be called in MM encoder-only mode")

    def __call__(self, *args, **kwargs):
        raise RuntimeError("This module should not be called in MM encoder-only mode")


class TowerMissingLayer(nn.Module):
82
83
84
85
    def __init__(self, modalities: set[str] | str) -> None:
        if isinstance(modalities, str):
            modalities = {modalities}

86
87
88
89
90
        super().__init__()

        self.modalities = modalities

    def __call__(self, *args, **kwargs):
91
92
93
94
        raise RuntimeError(
            f"This module should not be called when the following "
            f"modalities are disabled: {self.modalities}"
        )
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114


@contextmanager
def _no_init_weights(module: nn.Module, placeholder: Callable[[], nn.Module]):
    """
    Within this context, prevent weight initialization from using device memory and
    replace direct child assignments to `module` with the result of `placeholder()`.
    """

    def callback(module_, name, submodule):
        if module_ is module:
            return placeholder()

        return submodule

    with torch.nn.modules.module.register_module_module_registration_hook(callback):  # noqa: E501,SIM117
        with torch.device("meta"):
            yield


115
@runtime_checkable
116
class SupportsMultiModal(Protocol):
117
    """The interface required for all multi-modal models."""
118

119
    supports_multimodal: ClassVar[Literal[True]] = True
120
    """
121
    A flag that indicates this model supports multi-modal inputs.
122
123
124
125
126

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

128
129
130
131
132
133
    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.
    """

134
135
136
137
138
139
    supports_encoder_tp_data: ClassVar[bool] = False
    """
    A flag that indicates whether this model supports
    `multimodal_config.mm_encoder_tp_mode="data"`.
    """

Patrick von Platen's avatar
Patrick von Platen committed
140
141
142
143
144
145
    requires_raw_input_tokens: ClassVar[bool] = False
    """
    A flag that indicates this model processes input id tokens
    in their raw form and not input embeddings.
    """

146
147
148
149
150
    _processor_factory: ClassVar[_ProcessorFactories]
    """
    Set internally by `MultiModalRegistry.register_processor`.
    """

151
152
153
154
155
156
157
158
159
160
    _language_model_names: list[str] = []
    """
    Set internally by `_mark_language_model`.
    """

    _tower_model_names: list[str] = []
    """
    Set internally by `_mark_tower_model`.
    """

161
    @classmethod
162
    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
163
164
165
166
167
        """
        Get the placeholder text for the `i`th `modality` item in the prompt.
        """
        ...

168
    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
169
        """
170
        Returns multimodal embeddings generated from multimodal kwargs
171
        to be merged with text embeddings.
172

173
        Note:
174
175
            The returned multimodal embeddings must be in the same order as
            the appearances of their corresponding multimodal data item in the
176
            input prompt.
177
        """
178
        ...
179

180
    def get_language_model(self) -> VllmModel:
181
182
183
        """
        Returns the underlying language model used for text generation.

184
        This is typically the `torch.nn.Module` instance responsible for
185
186
187
188
189
        processing the merged multimodal embeddings and producing hidden states

        Returns:
            torch.nn.Module: The core language model component.
        """
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
        if self._language_model_names:
            return getattr(self, self._language_model_names[0])

        raise NotImplementedError(
            f"No language model found in {type(self).__name__}! "
            "You should initialize it inside `_mark_language_model`."
        )

    @contextmanager
    def _mark_language_model(self, vllm_config: VllmConfig):
        """
        Mark each child module that was assigned to this model
        during this context as a language model component.
        """
        mm_config = vllm_config.model_config.multimodal_config

        children_names = list[str]()

        def callback(module_, name, submodule):
            if module_ is self:
                children_names.append(name)

        with torch.nn.modules.module.register_module_module_registration_hook(callback):  # noqa: E501,SIM117
            with (
                _no_init_weights(self, LMMissingLayer)
                if mm_config.mm_encoder_only
                else nullcontext()
            ):
                yield

        self._language_model_names = children_names

    @contextmanager
    def _mark_tower_model(self, vllm_config: VllmConfig, modalities: set[str] | str):
        """
        Mark each child module that was assigned to this model
        during this context as a tower model component.
        """
        if isinstance(modalities, str):
            modalities = {modalities}

        mm_config = vllm_config.model_config.multimodal_config

        children_names = list[str]()

        def callback(module_, name, submodule):
            if module_ is self:
                children_names.append(name)

        with torch.nn.modules.module.register_module_module_registration_hook(callback):  # noqa: E501,SIM117
            with (
                _no_init_weights(self, lambda: TowerMissingLayer(modalities))
                if all(mm_config.get_limit_per_prompt(m) == 0 for m in modalities)
                else nullcontext()
            ):
                yield

        self._tower_model_names = children_names
248

249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
    def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
        """
        Implement this function to enable LoRA support
        for the tower module of the multi-modal model.
        Given the number of image tokens, output the number of
        multi-modal encoder tokens.
        """
        ...

    def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
        """
        Implement this function to enable LoRA support
        for the connector module of the multi-modal model.
        Given the number of vision tokens, output the number of
        multi-modal connector tokens.
        """
        ...

267
    @overload
268
    def embed_input_ids(self, input_ids: Tensor) -> Tensor: ...
269
270

    @overload
271
    def embed_input_ids(
272
273
274
275
276
277
        self,
        input_ids: Tensor,
        multimodal_embeddings: MultiModalEmbeddings,
        *,
        is_multimodal: torch.Tensor,
        handle_oov_mm_token: bool = False,
278
    ) -> Tensor: ...
279

280
    def _embed_text_input_ids(
281
282
        self,
        input_ids: Tensor,
283
        embed_input_ids: Callable[[Tensor], Tensor],
284
        *,
285
        is_multimodal: Tensor | None,
286
287
288
289
        handle_oov_mm_token: bool,
    ) -> Tensor:
        if handle_oov_mm_token and is_multimodal is not None:
            is_text = ~is_multimodal
290
            text_embeds = embed_input_ids(input_ids[is_text])
291
292
293
294
295
296
297

            return torch.empty(
                (input_ids.shape[0], text_embeds.shape[1]),
                dtype=text_embeds.dtype,
                device=text_embeds.device,
            ).masked_scatter_(is_text.unsqueeze_(-1), text_embeds)

298
        return embed_input_ids(input_ids)
299

300
    def embed_input_ids(
301
        self,
302
        input_ids: Tensor,
303
        multimodal_embeddings: MultiModalEmbeddings | None = None,
304
        *,
305
        is_multimodal: Tensor | None = None,
306
        handle_oov_mm_token: bool = False,
307
    ) -> Tensor:
308
        """
309
310
311
312
313
314
315
        Apply token embeddings to `input_ids`.

        If `multimodal_embeddings` is passed, scatter them into
        `input_ids` according to the mask `is_multimodal`.

        In case the multi-modal token IDs exceed the vocabulary size of
        the language model, you can set `handle_oov_mm_token=False`
316
        to avoid calling the language model's `embed_input_ids` method
317
318
        on those tokens. Note however that doing so increases memory usage
        as an additional buffer is needed to hold the input embeddings.
319
        """
320
321
        from .utils import _merge_multimodal_embeddings

322
        inputs_embeds = self._embed_text_input_ids(
323
            input_ids,
324
            self.get_language_model().embed_input_ids,
325
326
327
328
329
330
331
332
333
334
            is_multimodal=is_multimodal,
            handle_oov_mm_token=handle_oov_mm_token,
        )

        if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
            return inputs_embeds

        return _merge_multimodal_embeddings(
            inputs_embeds=inputs_embeds,
            multimodal_embeddings=multimodal_embeddings,
335
            is_multimodal=_require_is_multimodal(is_multimodal),
336
        )
337

338

339
340
341
342
343
344
@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.
    """
345

346
347
348
    supports_multimodal_pruning: ClassVar[Literal[True]] = True

    def recompute_mrope_positions(
349
350
351
352
353
        self,
        input_ids: list[int],
        multimodal_embeddings: MultiModalEmbeddings,
        mrope_positions: torch.LongTensor,
        num_computed_tokens: int,
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
    ) -> 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).
        """
        ...


378
@overload
379
def supports_multimodal(model: type[object]) -> TypeIs[type[SupportsMultiModal]]: ...
380
381
382


@overload
383
def supports_multimodal(model: object) -> TypeIs[SupportsMultiModal]: ...
384
385


386
def supports_multimodal(
387
388
    model: type[object] | object,
) -> TypeIs[type[SupportsMultiModal]] | TypeIs[SupportsMultiModal]:
389
    return getattr(model, "supports_multimodal", False)
390
391


392
def supports_multimodal_raw_input_only(model: type[object] | object) -> bool:
393
    return getattr(model, "supports_multimodal_raw_input_only", False)
394

395

Patrick von Platen's avatar
Patrick von Platen committed
396
397
398
399
def requires_raw_input_tokens(model: type[object] | object) -> bool:
    return getattr(model, "requires_raw_input_tokens", False)


400
def supports_multimodal_encoder_tp_data(model: type[object] | object) -> bool:
401
    return getattr(model, "supports_encoder_tp_data", False)
402
403


404
405
@overload
def supports_multimodal_pruning(
406
407
    model: type[object],
) -> TypeIs[type[SupportsMultiModalPruning]]: ...
408
409
410


@overload
411
def supports_multimodal_pruning(model: object) -> TypeIs[SupportsMultiModalPruning]: ...
412
413
414


def supports_multimodal_pruning(
415
416
    model: type[object] | object,
) -> TypeIs[type[SupportsMultiModalPruning]] | TypeIs[SupportsMultiModalPruning]:
417
418
419
    return getattr(model, "supports_multimodal_pruning", False)


420
421
422
423
424
425
426
427
428
429
430
431
432
433
@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
434
    def get_score_template(cls, query: str, document: str) -> str | None:
435
436
        """
        Generate a full prompt by populating the score template with query and document content.
437
        """  # noqa: E501
438
439
440
441
442
443
444
445
446
447
448
449
        ...

    @classmethod
    def post_process_tokens(cls, prompt: TokensPrompt) -> None:
        """
        Perform architecture-specific manipulations on the input tokens.
        """
        ...


@overload
def supports_score_template(
450
451
    model: type[object],
) -> TypeIs[type[SupportsScoreTemplate]]: ...
452
453
454


@overload
455
def supports_score_template(model: object) -> TypeIs[SupportsScoreTemplate]: ...
456
457
458


def supports_score_template(
459
460
    model: type[object] | object,
) -> TypeIs[type[SupportsScoreTemplate]] | TypeIs[SupportsScoreTemplate]:
461
    return getattr(model, "supports_score_template", False)
462
463


464
465
466
467
@runtime_checkable
class SupportsLoRA(Protocol):
    """The interface required for all models that support LoRA."""

468
469
470
471
472
473
474
475
    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.
    """
476
    is_3d_moe_weight: ClassVar[bool] = False
477
    is_non_gated_moe: ClassVar[bool] = False
478
479
    # The `embedding_module` and `embedding_padding_modules`
    # are empty by default.
480
    embedding_modules: ClassVar[dict[str, str]] = {}
481
    packed_modules_mapping: dict[str, list[str]] = {}
482
483
484
485
486
487
488
489


# 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]

490
491
    packed_modules_mapping: dict[str, list[str]]
    embedding_modules: dict[str, str]
492
493
494


@overload
495
def supports_lora(model: type[object]) -> TypeIs[type[SupportsLoRA]]: ...
496
497
498


@overload
499
def supports_lora(model: object) -> TypeIs[SupportsLoRA]: ...
500
501
502


def supports_lora(
503
504
    model: type[object] | object,
) -> TypeIs[type[SupportsLoRA]] | TypeIs[SupportsLoRA]:
505
506
507
508
509
510
511
    result = _supports_lora(model)

    if not result:
        lora_attrs = (
            "packed_modules_mapping",
            "embedding_modules",
        )
512
        missing_attrs = tuple(attr for attr in lora_attrs if not hasattr(model, attr))
513
514
515
516
517
518
519
520
521
522
523
524
525

        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, "
526
527
528
                    "but does not set `supports_lora=True`.",
                    model,
                )
529
530
531
532

    return result


533
def _supports_lora(model: type[object] | object) -> bool:
534
535
536
537
    if isinstance(model, type):
        return isinstance(model, _SupportsLoRAType)

    return isinstance(model, SupportsLoRA)
538
539


540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
@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,
558
    ) -> IntermediateTensors:
559
560
561
562
563
        """Called when PP rank > 0 for profiling purposes."""
        ...

    def forward(
        self,
564
        *,
565
566
        intermediate_tensors: IntermediateTensors | None,
    ) -> IntermediateTensors | None:
567
        """
568
569
        Accept [`IntermediateTensors`][vllm.sequence.IntermediateTensors] when
        PP rank > 0.
570

571
572
        Return [`IntermediateTensors`][vllm.sequence.IntermediateTensors] only
        for the last PP rank.
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
        """
        ...


# 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,
588
    ) -> IntermediateTensors: ...
589
590
591

    def forward(
        self,
592
        *,
593
594
        intermediate_tensors: IntermediateTensors | None,
    ) -> Tensor | IntermediateTensors: ...
595
596
597


@overload
598
def supports_pp(model: type[object]) -> TypeIs[type[SupportsPP]]: ...
599
600
601


@overload
602
def supports_pp(model: object) -> TypeIs[SupportsPP]: ...
603
604
605


def supports_pp(
606
607
    model: type[object] | object,
) -> bool | TypeIs[type[SupportsPP]] | TypeIs[SupportsPP]:
608
609
610
611
612
613
    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 "
614
615
616
            "`intermediate_tensors` in its `forward` method",
            model,
        )
617
618

    if not supports_attributes:
619
620
        pp_attrs = ("make_empty_intermediate_tensors",)
        missing_attrs = tuple(attr for attr in pp_attrs if not hasattr(model, attr))
621
622
623
624
625
626
627
628
629
630
631
632
633

        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, "
634
635
636
                    "but does not set `supports_pp=True`.",
                    model,
                )
637
638
639
640

    return supports_attributes and supports_inspect


641
def _supports_pp_attributes(model: type[object] | object) -> bool:
642
643
644
645
646
647
    if isinstance(model, type):
        return isinstance(model, _SupportsPPType)

    return isinstance(model, SupportsPP)


648
def _supports_pp_inspect(model: type[object] | object) -> bool:
649
650
651
652
    model_forward = getattr(model, "forward", None)
    if not callable(model_forward):
        return False

653
    return supports_kw(model_forward, "intermediate_tensors")
654
655


656
657
658
659
660
661
662
663
@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
664
        for max_num_seqs, etc. True for e.g. both Mamba and Jamba.
665
666
667
668
    """


@overload
669
def has_inner_state(model: object) -> TypeIs[HasInnerState]: ...
670
671
672


@overload
673
def has_inner_state(model: type[object]) -> TypeIs[type[HasInnerState]]: ...
674
675
676


def has_inner_state(
677
678
    model: type[object] | object,
) -> TypeIs[type[HasInnerState]] | TypeIs[HasInnerState]:
679
    return getattr(model, "has_inner_state", False)
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695


@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
696
def is_attention_free(model: object) -> TypeIs[IsAttentionFree]: ...
697
698
699


@overload
700
def is_attention_free(model: type[object]) -> TypeIs[type[IsAttentionFree]]: ...
701
702
703


def is_attention_free(
704
705
    model: type[object] | object,
) -> TypeIs[type[IsAttentionFree]] | TypeIs[IsAttentionFree]:
706
    return getattr(model, "is_attention_free", False)
707
708


709
710
711
@runtime_checkable
class IsHybrid(Protocol):
    """The interface required for all models like Jamba that have both
712
    attention and mamba blocks, indicates that
713
714
715
716
717
718
719
720
    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' """

721
722
723
    @classmethod
    def get_mamba_state_shape_from_config(
        cls,
724
        vllm_config: VllmConfig,
725
726
727
728
729
730
731
732
733
734
735
736
737
    ) -> tuple[tuple[int, int], tuple[int, int, int]]:
        """Calculate shapes for Mamba's convolutional and state caches.

        Args:
            vllm_config: vLLM config

        Returns:
            Tuple containing:
            - conv_state_shape: Shape for convolutional state cache
            - temporal_state_shape: Shape for state space model cache
        """
        ...

738
739

@overload
740
def is_hybrid(model: object) -> TypeIs[IsHybrid]: ...
741
742
743


@overload
744
def is_hybrid(model: type[object]) -> TypeIs[type[IsHybrid]]: ...
745
746
747


def is_hybrid(
748
749
    model: type[object] | object,
) -> TypeIs[type[IsHybrid]] | TypeIs[IsHybrid]:
750
    return getattr(model, "is_hybrid", False)
751
752


753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
@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."""

791
792
793
    moe_layers: Iterable[nn.Module]
    """List of MoE layers in this model."""

794
795
796
797
798
799
800
801
    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.
802

803
804
805
806
807
808
809
810
811
812
813
814
815
        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.
        """
816
817
818
819
820
821
822
823
824
        for layer_idx, layer in enumerate(self.moe_layers):
            # Register the expert weights.
            self.expert_weights.append(layer.get_expert_weights())
            layer.set_eplb_state(
                moe_layer_idx=layer_idx,
                expert_load_view=expert_load_view,
                logical_to_physical_map=logical_to_physical_map,
                logical_replica_count=logical_replica_count,
            )
825

826
827
828
829
    def update_physical_experts_metadata(
        self,
        num_physical_experts: int,
        num_local_physical_experts: int,
830
    ) -> None: ...
831

832
833

def is_mixture_of_experts(model: object) -> TypeIs[MixtureOfExperts]:
834
835
836
    return (
        isinstance(model, MixtureOfExperts) and getattr(model, "num_moe_layers", 0) > 0
    )
837
838


839
840
841
842
843
844
@runtime_checkable
class HasNoOps(Protocol):
    has_noops: ClassVar[Literal[True]] = True


@overload
845
def has_noops(model: object) -> TypeIs[HasNoOps]: ...
846
847
848


@overload
849
def has_noops(model: type[object]) -> TypeIs[type[HasNoOps]]: ...
850
851
852


def has_noops(
853
854
    model: type[object] | object,
) -> TypeIs[type[HasNoOps]] | TypeIs[HasNoOps]:
855
    return getattr(model, "has_noops", False)
856
857


858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
@runtime_checkable
class SupportsMambaPrefixCaching(Protocol):
    """The interface for models whose mamba layers support prefix caching.

    This is currently experimental.
    """

    supports_mamba_prefix_caching: ClassVar[Literal[True]] = True


@overload
def supports_mamba_prefix_caching(
    model: object,
) -> TypeIs[SupportsMambaPrefixCaching]: ...


@overload
def supports_mamba_prefix_caching(
    model: type[object],
) -> TypeIs[type[SupportsMambaPrefixCaching]]: ...


def supports_mamba_prefix_caching(
    model: type[object] | object,
) -> TypeIs[type[SupportsMambaPrefixCaching]] | TypeIs[SupportsMambaPrefixCaching]:
    return getattr(model, "supports_mamba_prefix_caching", False)


886
887
888
889
890
891
892
893
894
@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(
895
896
    model: type[object],
) -> TypeIs[type[SupportsCrossEncoding]]: ...
897
898
899


@overload
900
def supports_cross_encoding(model: object) -> TypeIs[SupportsCrossEncoding]: ...
901
902
903


def _supports_cross_encoding(
904
905
    model: type[object] | object,
) -> TypeIs[type[SupportsCrossEncoding]] | TypeIs[SupportsCrossEncoding]:
906
    return getattr(model, "supports_cross_encoding", False)
907
908
909


def supports_cross_encoding(
910
911
    model: type[object] | object,
) -> TypeIs[type[SupportsCrossEncoding]] | TypeIs[SupportsCrossEncoding]:
912
    return is_pooling_model(model) and _supports_cross_encoding(model)
913
914


915
916
917
class SupportsQuant:
    """The interface required for all models that support quantization."""

918
919
920
    hf_to_vllm_mapper: ClassVar[WeightsMapper | None] = None
    packed_modules_mapping: ClassVar[dict[str, list[str]] | None] = None
    quant_config: QuantizationConfig | None = None
921

922
    def __new__(cls, *args, **kwargs) -> Self:
923
        instance = super().__new__(cls)
924
925

        # find config passed in arguments
926
927
        quant_config = cls._find_quant_config(*args, **kwargs)
        if quant_config is not None:
928
            # attach config to model for general use
929
            instance.quant_config = quant_config
930
931

            # apply model mappings to config for proper config-model matching
932
933
934
            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:
935
                instance.quant_config.packed_modules_mapping.update(
936
937
                    instance.packed_modules_mapping
                )
938

939
940
941
        return instance

    @staticmethod
942
    def _find_quant_config(*args, **kwargs) -> QuantizationConfig | None:
943
        """Find quant config passed through model constructor args"""
944
945
946
947
948
949
950
951
952
953
954
955
956
        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


957
958
959
@runtime_checkable
class SupportsTranscription(Protocol):
    """The interface required for all models that support transcription."""
960

961
962
    # Mapping from ISO639_1 language codes: language names
    supported_languages: ClassVar[Mapping[str, str]]
963
964
965

    supports_transcription: ClassVar[Literal[True]] = True

966
967
968
969
970
    supports_transcription_only: ClassVar[bool] = False
    """
    Transcription models can opt out of text generation by setting this to
    `True`.
    """
971
972
973
974
    supports_segment_timestamp: ClassVar[bool] = False
    """
    Enables the segment timestamp option for supported models by setting this to `True`.
    """
975

976
977
978
979
980
981
982
983
984
    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. "
985
986
                f"Valid choices are: {sorted(LANGUAGES.keys())}"
            )
987

988
    @classmethod
989
990
991
992
    def get_generation_prompt(
        cls,
        audio: np.ndarray,
        stt_config: SpeechToTextConfig,
993
        model_config: ModelConfig,
994
        language: str | None,
995
996
        task_type: Literal["transcribe", "translate"],
        request_prompt: str,
997
        to_language: str | None,
998
    ) -> PromptType:
999
1000
1001
        """Get the prompt for the ASR model.
        The model has control over the construction, as long as it
        returns a valid PromptType."""
1002
1003
1004
        ...

    @classmethod
1005
1006
    def get_other_languages(cls) -> Mapping[str, str]:
        # other possible language codes from the whisper map
1007
        return {k: v for k, v in LANGUAGES.items() if k not in cls.supported_languages}
1008
1009

    @classmethod
1010
    def validate_language(cls, language: str | None) -> str | None:
1011
        """
1012
1013
1014
        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
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
        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 "
1031
1032
                f"{list(cls.supported_languages.keys())}."
            )
1033

1034
1035
    @classmethod
    def get_speech_to_text_config(
1036
        cls, model_config: ModelConfig, task_type: Literal["transcribe", "translate"]
1037
    ) -> SpeechToTextConfig:
1038
1039
1040
1041
        """Get the speech to text config for the ASR model."""
        ...

    @classmethod
1042
1043
1044
1045
    def get_num_audio_tokens(
        cls,
        audio_duration_s: float,
        stt_config: SpeechToTextConfig,
1046
        model_config: ModelConfig,
1047
    ) -> int | None:
1048
        """
1049
        Map from audio duration to number of audio tokens produced by the ASR
1050
1051
1052
1053
1054
        model, without running a forward pass.
        This is used for estimating the amount of processing for this audio.
        """
        return None

1055
1056
1057

@overload
def supports_transcription(
1058
1059
    model: type[object],
) -> TypeIs[type[SupportsTranscription]]: ...
1060
1061
1062


@overload
1063
def supports_transcription(model: object) -> TypeIs[SupportsTranscription]: ...
1064
1065
1066


def supports_transcription(
1067
1068
    model: type[object] | object,
) -> TypeIs[type[SupportsTranscription]] | TypeIs[SupportsTranscription]:
1069
    return getattr(model, "supports_transcription", False)
1070
1071


1072
@runtime_checkable
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
class SupportsEagleBase(Protocol):
    """Base interface for models that support EAGLE-based speculative decoding."""

    has_own_lm_head: bool = False
    """
    A flag that indicates this model has trained its own lm_head.
    """

    has_own_embed_tokens: bool = False
    """
    A flag that indicates this model has trained its own input embeddings.
    """


@overload
def supports_any_eagle(model: type[object]) -> TypeIs[type[SupportsEagleBase]]: ...


@overload
def supports_any_eagle(model: object) -> TypeIs[SupportsEagleBase]: ...


def supports_any_eagle(
    model: type[object] | object,
) -> TypeIs[type[SupportsEagleBase]] | TypeIs[SupportsEagleBase]:
    """Check if model supports any EAGLE variant (1, 2, or 3)."""
    return supports_eagle(model) or supports_eagle3(model)


@runtime_checkable
class SupportsEagle(SupportsEagleBase, Protocol):
    """The interface required for models that support
    EAGLE-1 and EAGLE-2 speculative decoding."""

    supports_eagle: ClassVar[Literal[True]] = True
    """
    A flag that indicates this model supports EAGLE-1 and EAGLE-2 
    speculative decoding.

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


@overload
def supports_eagle(model: type[object]) -> TypeIs[type[SupportsEagle]]: ...


@overload
def supports_eagle(model: object) -> TypeIs[SupportsEagle]: ...


def supports_eagle(
    model: type[object] | object,
) -> TypeIs[type[SupportsEagle]] | TypeIs[SupportsEagle]:
    return isinstance(model, SupportsEagle)


@runtime_checkable
class SupportsEagle3(SupportsEagleBase, Protocol):
1134
    """The interface required for models that support
1135
    EAGLE-3 speculative decoding."""
1136
1137
1138

    supports_eagle3: ClassVar[Literal[True]] = True
    """
1139
    A flag that indicates this model supports EAGLE-3 
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
    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
1150
        hidden states for EAGLE-3.
1151

1152
1153
        Args:
            layers: Tuple of layer indices that should output auxiliary
1154
                hidden states.
1155
1156
1157
1158
1159
1160
        """
        ...

    def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
        """
        Get the layer indices that should output auxiliary hidden states
1161
        for EAGLE-3.
1162

1163
1164
1165
1166
1167
1168
1169
        Returns:
            Tuple of layer indices for auxiliary hidden state outputs.
        """
        ...


@overload
1170
def supports_eagle3(model: type[object]) -> TypeIs[type[SupportsEagle3]]: ...
1171
1172
1173


@overload
1174
def supports_eagle3(model: object) -> TypeIs[SupportsEagle3]: ...
1175
1176
1177


def supports_eagle3(
1178
1179
    model: type[object] | object,
) -> TypeIs[type[SupportsEagle3]] | TypeIs[SupportsEagle3]:
1180
    return isinstance(model, SupportsEagle3)
1181
1182
1183
1184
1185
1186
1187
1188
1189


@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.
1190

1191
1192
1193
1194
1195
1196
1197
1198
    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],
1199
        mm_features: list["MultiModalFeatureSpec"],
1200
1201
1202
    ) -> tuple[torch.Tensor, int]:
        """
        Get M-RoPE input positions and delta value for this specific model.
1203

1204
1205
        This method should be implemented by each model that supports M-RoPE
        to provide model-specific logic for computing input positions.
1206

1207
1208
        Args:
            input_tokens: List of input token IDs
1209
            mm_features: Information about each multi-modal data item
1210

1211
        Returns:
1212
1213
            Tuple of `(llm_positions, mrope_position_delta)`
            - llm_positions: Tensor of shape `[3, num_tokens]` with T/H/W positions
1214
1215
1216
1217
1218
1219
            - mrope_position_delta: Delta for position calculations
        """
        ...


@overload
1220
def supports_mrope(model: type[object]) -> TypeIs[type[SupportsMRoPE]]: ...
1221
1222
1223


@overload
1224
def supports_mrope(model: object) -> TypeIs[SupportsMRoPE]: ...
1225
1226
1227


def supports_mrope(
1228
1229
    model: type[object] | object,
) -> TypeIs[type[SupportsMRoPE]] | TypeIs[SupportsMRoPE]:
1230
    return isinstance(model, SupportsMRoPE)
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279


@runtime_checkable
class SupportsXDRoPE(Protocol):
    """The interface required for all models that support XD-RoPE."""

    supports_xdrope: ClassVar[Literal[True]] = True
    """
    A flag that indicates this model supports XD-RoPE.

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

    def get_xdrope_input_positions(
        self,
        input_tokens: list[int],
        mm_features: list["MultiModalFeatureSpec"],
    ) -> torch.Tensor:
        """
        Get XD-RoPE input positions and delta value for this specific model.

        This method should be implemented by each model that supports XD-RoPE
        to provide model-specific logic for computing input positions.

        Args:
            input_tokens: List of input token IDs
            mm_features: Information about each multi-modal data item

        Returns:
            llm_positions: Tensor of shape `[xdrope_dim, num_tokens]` with
            4D(P/W/H/T) or 3D(W/H/T) positions.
        """
        ...


@overload
def supports_xdrope(model: type[object]) -> TypeIs[type[SupportsXDRoPE]]: ...


@overload
def supports_xdrope(model: object) -> TypeIs[SupportsXDRoPE]: ...


def supports_xdrope(
    model: type[object] | object,
) -> TypeIs[type[SupportsXDRoPE]] | TypeIs[SupportsXDRoPE]:
    return isinstance(model, SupportsXDRoPE)