interfaces.py 38.4 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
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
class LMMissingLayer(nn.Module):
    packed_modules_mapping: dict[str, list[str]] = {}

    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):
    packed_modules_mapping: dict[str, list[str]] = {}

    def __init__(self, modalities: set[str]) -> None:
        super().__init__()

        self.modalities = modalities

    def __call__(self, *args, **kwargs):
        raise RuntimeError(f"The following modalities are disabled: {self.modalities}")


@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


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

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

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

126
127
128
129
130
131
    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.
    """

132
133
134
135
136
137
    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
138
139
140
141
142
143
    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.
    """

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

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

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

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

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

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

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

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

        Returns:
            torch.nn.Module: The core language model component.
        """
188
189
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
        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
246

247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
    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.
        """
        ...

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

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

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

            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)

296
        return embed_input_ids(input_ids)
297

298
    def embed_input_ids(
299
        self,
300
        input_ids: Tensor,
301
        multimodal_embeddings: MultiModalEmbeddings | None = None,
302
        *,
303
        is_multimodal: Tensor | None = None,
304
        handle_oov_mm_token: bool = False,
305
    ) -> Tensor:
306
        """
307
308
309
310
311
312
313
        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`
314
        to avoid calling the language model's `embed_input_ids` method
315
316
        on those tokens. Note however that doing so increases memory usage
        as an additional buffer is needed to hold the input embeddings.
317
        """
318
319
        from .utils import _merge_multimodal_embeddings

320
        inputs_embeds = self._embed_text_input_ids(
321
            input_ids,
322
            self.get_language_model().embed_input_ids,
323
324
325
326
327
328
329
330
331
332
            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,
333
            is_multimodal=_require_is_multimodal(is_multimodal),
334
        )
335

336

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

344
345
346
    supports_multimodal_pruning: ClassVar[Literal[True]] = True

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


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


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


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


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

393

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


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


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


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


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


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

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


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


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


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


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

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


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

488
489
    packed_modules_mapping: dict[str, list[str]]
    embedding_modules: dict[str, str]
490
491
492


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


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


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

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

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

    return result


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

    return isinstance(model, SupportsLoRA)
536
537


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

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

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


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

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


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


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


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

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

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

    return supports_attributes and supports_inspect


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

    return isinstance(model, SupportsPP)


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

651
    return supports_kw(model_forward, "intermediate_tensors")
652
653


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


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


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


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


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


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


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


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

719
720
721
    @classmethod
    def get_mamba_state_shape_from_config(
        cls,
722
        vllm_config: VllmConfig,
723
724
725
726
727
728
729
730
731
732
733
734
735
    ) -> 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
        """
        ...

736
737

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


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


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


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

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

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

801
802
803
804
805
806
807
808
809
810
811
812
813
        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.
        """
814
815
816
817
818
819
820
821
822
        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,
            )
823

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

830
831

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


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


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


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


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


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
@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)


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


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


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


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


913
914
915
class SupportsQuant:
    """The interface required for all models that support quantization."""

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

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

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

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

937
938
939
        return instance

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


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

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

    supports_transcription: ClassVar[Literal[True]] = True

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

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

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

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

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

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

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

1053
1054
1055

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


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


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


1070
@runtime_checkable
1071
1072
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
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):
1132
    """The interface required for models that support
1133
    EAGLE-3 speculative decoding."""
1134
1135
1136

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

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

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

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


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


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


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


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

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

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

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

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


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


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


def supports_mrope(
1226
1227
    model: type[object] | object,
) -> TypeIs[type[SupportsMRoPE]] | TypeIs[SupportsMRoPE]:
1228
    return isinstance(model, SupportsMRoPE)
1229
1230
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


@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)