interfaces.py 35 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, Set
5
6
7
8
9
from typing import (
    TYPE_CHECKING,
    ClassVar,
    Literal,
    Protocol,
10
    TypeAlias,
11
12
13
    overload,
    runtime_checkable,
)
14

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

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

29
from .interfaces_base import VllmModel, is_pooling_model
30

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

44
45
logger = init_logger(__name__)

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

55
56

@runtime_checkable
57
class SupportsMultiModal(Protocol):
58
    """The interface required for all multi-modal models."""
59

60
    supports_multimodal: ClassVar[Literal[True]] = True
61
    """
62
    A flag that indicates this model supports multi-modal inputs.
63
64
65
66
67

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

69
70
71
72
73
74
    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.
    """

75
76
77
78
79
80
    supports_encoder_tp_data: ClassVar[bool] = False
    """
    A flag that indicates whether this model supports
    `multimodal_config.mm_encoder_tp_mode="data"`.
    """

81
    merge_by_field_config: ClassVar[bool | None] = None
82
    """
83
    [DEPRECATED] A flag that indicates which implementation of
84
85
86
    `vllm.multimodal.utils.group_mm_kwargs_by_modality` to use.
    """

87
    multimodal_cpu_fields: ClassVar[Set[str] | None] = None
88
    """
89
    [DEPRECATED] A set indicating CPU-only multimodal fields.
90
91
    """

92
93
94
95
96
    _processor_factory: ClassVar[_ProcessorFactories]
    """
    Set internally by `MultiModalRegistry.register_processor`.
    """

97
    @classmethod
98
    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
99
100
101
102
103
        """
        Get the placeholder text for the `i`th `modality` item in the prompt.
        """
        ...

104
    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
105
        """
106
        Returns multimodal embeddings generated from multimodal kwargs
107
        to be merged with text embeddings.
108

109
        Note:
110
111
            The returned multimodal embeddings must be in the same order as
            the appearances of their corresponding multimodal data item in the
112
            input prompt.
113
        """
114
        ...
115

116
    def get_language_model(self) -> VllmModel:
117
118
119
        """
        Returns the underlying language model used for text generation.

120
        This is typically the `torch.nn.Module` instance responsible for
121
122
123
124
125
126
127
        processing the merged multimodal embeddings and producing hidden states

        Returns:
            torch.nn.Module: The core language model component.
        """
        ...

128
    @overload
129
    def embed_input_ids(self, input_ids: Tensor) -> Tensor: ...
130
131

    @overload
132
    def embed_input_ids(
133
134
135
136
137
138
        self,
        input_ids: Tensor,
        multimodal_embeddings: MultiModalEmbeddings,
        *,
        is_multimodal: torch.Tensor,
        handle_oov_mm_token: bool = False,
139
    ) -> Tensor: ...
140

141
    def _embed_text_input_ids(
142
143
        self,
        input_ids: Tensor,
144
        embed_input_ids: Callable[[Tensor], Tensor],
145
        *,
146
        is_multimodal: Tensor | None,
147
148
149
150
        handle_oov_mm_token: bool,
    ) -> Tensor:
        if handle_oov_mm_token and is_multimodal is not None:
            is_text = ~is_multimodal
151
            text_embeds = embed_input_ids(input_ids[is_text])
152
153
154
155
156
157
158

            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)

159
        return embed_input_ids(input_ids)
160

161
    def embed_input_ids(
162
        self,
163
        input_ids: Tensor,
164
        multimodal_embeddings: MultiModalEmbeddings | None = None,
165
        *,
166
        is_multimodal: Tensor | None = None,
167
        handle_oov_mm_token: bool = False,
168
    ) -> Tensor:
169
        """
170
171
172
173
174
175
176
        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`
177
        to avoid calling the language model's `embed_input_ids` method
178
179
        on those tokens. Note however that doing so increases memory usage
        as an additional buffer is needed to hold the input embeddings.
180
        """
181
182
        from .utils import _merge_multimodal_embeddings

183
        inputs_embeds = self._embed_text_input_ids(
184
            input_ids,
185
            self.get_language_model().embed_input_ids,
186
187
188
189
190
191
192
            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

193
        assert is_multimodal is not None
194
195
196
197
198
199

        return _merge_multimodal_embeddings(
            inputs_embeds=inputs_embeds,
            multimodal_embeddings=multimodal_embeddings,
            is_multimodal=is_multimodal,
        )
200

201

202
203
204
205
206
207
@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.
    """
208

209
210
211
    supports_multimodal_pruning: ClassVar[Literal[True]] = True

    def recompute_mrope_positions(
212
213
214
215
216
        self,
        input_ids: list[int],
        multimodal_embeddings: MultiModalEmbeddings,
        mrope_positions: torch.LongTensor,
        num_computed_tokens: int,
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
    ) -> 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).
        """
        ...


241
@overload
242
def supports_multimodal(model: type[object]) -> TypeIs[type[SupportsMultiModal]]: ...
243
244
245


@overload
246
def supports_multimodal(model: object) -> TypeIs[SupportsMultiModal]: ...
247
248


249
def supports_multimodal(
250
251
    model: type[object] | object,
) -> TypeIs[type[SupportsMultiModal]] | TypeIs[SupportsMultiModal]:
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
    res = getattr(model, "supports_multimodal", False)

    if res:
        # We can remove this starting from v0.14
        merge_by_field_config = getattr(model, "merge_by_field_config", None)
        if merge_by_field_config is False:
            raise ValueError(
                "`merge_by_field_config=False` is no longer effective, "
                "please update your model to consider the new batching logic "
                "in `group_mm_kwargs_by_modality` (refer to "
                "https://github.com/vllm-project/vllm/issues/26149), "
                "and then remove the override from your model."
            )
        if merge_by_field_config is True:
            logger.warning_once(
                "`merge_by_field_config=True` is redundant, "
                "please remove the override from your model."
            )

271
272
273
274
275
276
277
278
279
        multimodal_cpu_fields = getattr(model, "multimodal_cpu_fields", None)
        if multimodal_cpu_fields is not None:
            raise ValueError(
                "`multimodal_cpu_fields` is no longer effective, "
                "please set `keep_on_cpu=True` in `MultiModalFieldConfig` "
                "(refer to https://github.com/vllm-project/vllm/pull/30181), "
                "and then remove the override from your model."
            )

280
    return res
281
282


283
def supports_multimodal_raw_input_only(model: type[object] | object) -> bool:
284
    return getattr(model, "supports_multimodal_raw_input_only", False)
285

286

287
def supports_multimodal_encoder_tp_data(model: type[object] | object) -> bool:
288
    return getattr(model, "supports_encoder_tp_data", False)
289
290


291
292
@overload
def supports_multimodal_pruning(
293
294
    model: type[object],
) -> TypeIs[type[SupportsMultiModalPruning]]: ...
295
296
297


@overload
298
def supports_multimodal_pruning(model: object) -> TypeIs[SupportsMultiModalPruning]: ...
299
300
301


def supports_multimodal_pruning(
302
303
    model: type[object] | object,
) -> TypeIs[type[SupportsMultiModalPruning]] | TypeIs[SupportsMultiModalPruning]:
304
305
306
    return getattr(model, "supports_multimodal_pruning", False)


307
308
309
310
311
312
313
314
315
316
317
318
319
320
@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
321
    def get_score_template(cls, query: str, document: str) -> str | None:
322
323
        """
        Generate a full prompt by populating the score template with query and document content.
324
        """  # noqa: E501
325
326
327
328
329
330
331
332
333
334
335
336
        ...

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


@overload
def supports_score_template(
337
338
    model: type[object],
) -> TypeIs[type[SupportsScoreTemplate]]: ...
339
340
341


@overload
342
def supports_score_template(model: object) -> TypeIs[SupportsScoreTemplate]: ...
343
344
345


def supports_score_template(
346
347
    model: type[object] | object,
) -> TypeIs[type[SupportsScoreTemplate]] | TypeIs[SupportsScoreTemplate]:
348
    return getattr(model, "supports_score_template", False)
349
350


351
352
353
354
@runtime_checkable
class SupportsLoRA(Protocol):
    """The interface required for all models that support LoRA."""

355
356
357
358
359
360
361
362
    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.
    """
363
    is_3d_moe_weight: ClassVar[bool] = False
364
365
    # The `embedding_module` and `embedding_padding_modules`
    # are empty by default.
366
    embedding_modules: ClassVar[dict[str, str]] = {}
367
    packed_modules_mapping: dict[str, list[str]] = {}
368
369
370
371
372
373
374
375


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

376
377
    packed_modules_mapping: dict[str, list[str]]
    embedding_modules: dict[str, str]
378
379
380


@overload
381
def supports_lora(model: type[object]) -> TypeIs[type[SupportsLoRA]]: ...
382
383
384


@overload
385
def supports_lora(model: object) -> TypeIs[SupportsLoRA]: ...
386
387
388


def supports_lora(
389
390
    model: type[object] | object,
) -> TypeIs[type[SupportsLoRA]] | TypeIs[SupportsLoRA]:
391
392
393
394
395
396
397
    result = _supports_lora(model)

    if not result:
        lora_attrs = (
            "packed_modules_mapping",
            "embedding_modules",
        )
398
        missing_attrs = tuple(attr for attr in lora_attrs if not hasattr(model, attr))
399
400
401
402
403
404
405
406
407
408
409
410
411

        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, "
412
413
414
                    "but does not set `supports_lora=True`.",
                    model,
                )
415
416
417
418

    return result


419
def _supports_lora(model: type[object] | object) -> bool:
420
421
422
423
    if isinstance(model, type):
        return isinstance(model, _SupportsLoRAType)

    return isinstance(model, SupportsLoRA)
424
425


426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
@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,
444
    ) -> IntermediateTensors:
445
446
447
448
449
        """Called when PP rank > 0 for profiling purposes."""
        ...

    def forward(
        self,
450
        *,
451
452
        intermediate_tensors: IntermediateTensors | None,
    ) -> IntermediateTensors | None:
453
        """
454
455
        Accept [`IntermediateTensors`][vllm.sequence.IntermediateTensors] when
        PP rank > 0.
456

457
458
        Return [`IntermediateTensors`][vllm.sequence.IntermediateTensors] only
        for the last PP rank.
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
        """
        ...


# 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,
474
    ) -> IntermediateTensors: ...
475
476
477

    def forward(
        self,
478
        *,
479
480
        intermediate_tensors: IntermediateTensors | None,
    ) -> Tensor | IntermediateTensors: ...
481
482
483


@overload
484
def supports_pp(model: type[object]) -> TypeIs[type[SupportsPP]]: ...
485
486
487


@overload
488
def supports_pp(model: object) -> TypeIs[SupportsPP]: ...
489
490
491


def supports_pp(
492
493
    model: type[object] | object,
) -> bool | TypeIs[type[SupportsPP]] | TypeIs[SupportsPP]:
494
495
496
497
498
499
    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 "
500
501
502
            "`intermediate_tensors` in its `forward` method",
            model,
        )
503
504

    if not supports_attributes:
505
506
        pp_attrs = ("make_empty_intermediate_tensors",)
        missing_attrs = tuple(attr for attr in pp_attrs if not hasattr(model, attr))
507
508
509
510
511
512
513
514
515
516
517
518
519

        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, "
520
521
522
                    "but does not set `supports_pp=True`.",
                    model,
                )
523
524
525
526

    return supports_attributes and supports_inspect


527
def _supports_pp_attributes(model: type[object] | object) -> bool:
528
529
530
531
532
533
    if isinstance(model, type):
        return isinstance(model, _SupportsPPType)

    return isinstance(model, SupportsPP)


534
def _supports_pp_inspect(model: type[object] | object) -> bool:
535
536
537
538
    model_forward = getattr(model, "forward", None)
    if not callable(model_forward):
        return False

539
    return supports_kw(model_forward, "intermediate_tensors")
540
541


542
543
544
545
546
547
548
549
@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
550
        for max_num_seqs, etc. True for e.g. both Mamba and Jamba.
551
552
553
554
    """


@overload
555
def has_inner_state(model: object) -> TypeIs[HasInnerState]: ...
556
557
558


@overload
559
def has_inner_state(model: type[object]) -> TypeIs[type[HasInnerState]]: ...
560
561
562


def has_inner_state(
563
564
    model: type[object] | object,
) -> TypeIs[type[HasInnerState]] | TypeIs[HasInnerState]:
565
    return getattr(model, "has_inner_state", False)
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581


@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
582
def is_attention_free(model: object) -> TypeIs[IsAttentionFree]: ...
583
584
585


@overload
586
def is_attention_free(model: type[object]) -> TypeIs[type[IsAttentionFree]]: ...
587
588
589


def is_attention_free(
590
591
    model: type[object] | object,
) -> TypeIs[type[IsAttentionFree]] | TypeIs[IsAttentionFree]:
592
    return getattr(model, "is_attention_free", False)
593
594


595
596
597
@runtime_checkable
class IsHybrid(Protocol):
    """The interface required for all models like Jamba that have both
598
    attention and mamba blocks, indicates that
599
600
601
602
603
604
605
606
    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' """

607
608
609
    @classmethod
    def get_mamba_state_shape_from_config(
        cls,
610
        vllm_config: VllmConfig,
611
612
613
614
615
616
617
618
619
620
621
622
623
    ) -> 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
        """
        ...

624
625

@overload
626
def is_hybrid(model: object) -> TypeIs[IsHybrid]: ...
627
628
629


@overload
630
def is_hybrid(model: type[object]) -> TypeIs[type[IsHybrid]]: ...
631
632
633


def is_hybrid(
634
635
    model: type[object] | object,
) -> TypeIs[type[IsHybrid]] | TypeIs[IsHybrid]:
636
    return getattr(model, "is_hybrid", False)
637
638


639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
@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."""

677
678
679
    moe_layers: Iterable[nn.Module]
    """List of MoE layers in this model."""

680
681
682
683
684
685
686
687
    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.
688

689
690
691
692
693
694
695
696
697
698
699
700
701
        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.
        """
702
703
704
705
706
707
708
709
710
        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,
            )
711

712
713
714
715
    def update_physical_experts_metadata(
        self,
        num_physical_experts: int,
        num_local_physical_experts: int,
716
    ) -> None: ...
717

718
719

def is_mixture_of_experts(model: object) -> TypeIs[MixtureOfExperts]:
720
721
722
    return (
        isinstance(model, MixtureOfExperts) and getattr(model, "num_moe_layers", 0) > 0
    )
723
724


725
726
727
728
729
730
@runtime_checkable
class HasNoOps(Protocol):
    has_noops: ClassVar[Literal[True]] = True


@overload
731
def has_noops(model: object) -> TypeIs[HasNoOps]: ...
732
733
734


@overload
735
def has_noops(model: type[object]) -> TypeIs[type[HasNoOps]]: ...
736
737
738


def has_noops(
739
740
    model: type[object] | object,
) -> TypeIs[type[HasNoOps]] | TypeIs[HasNoOps]:
741
    return getattr(model, "has_noops", False)
742
743


744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
@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)


772
773
774
775
776
777
778
779
780
@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(
781
782
    model: type[object],
) -> TypeIs[type[SupportsCrossEncoding]]: ...
783
784
785


@overload
786
def supports_cross_encoding(model: object) -> TypeIs[SupportsCrossEncoding]: ...
787
788
789


def _supports_cross_encoding(
790
791
    model: type[object] | object,
) -> TypeIs[type[SupportsCrossEncoding]] | TypeIs[SupportsCrossEncoding]:
792
    return getattr(model, "supports_cross_encoding", False)
793
794
795


def supports_cross_encoding(
796
797
    model: type[object] | object,
) -> TypeIs[type[SupportsCrossEncoding]] | TypeIs[SupportsCrossEncoding]:
798
    return is_pooling_model(model) and _supports_cross_encoding(model)
799
800


801
802
803
class SupportsQuant:
    """The interface required for all models that support quantization."""

804
805
806
    hf_to_vllm_mapper: ClassVar[WeightsMapper | None] = None
    packed_modules_mapping: ClassVar[dict[str, list[str]] | None] = None
    quant_config: QuantizationConfig | None = None
807

808
    def __new__(cls, *args, **kwargs) -> Self:
809
        instance = super().__new__(cls)
810
811

        # find config passed in arguments
812
813
        quant_config = cls._find_quant_config(*args, **kwargs)
        if quant_config is not None:
814
            # attach config to model for general use
815
            instance.quant_config = quant_config
816
817

            # apply model mappings to config for proper config-model matching
818
819
820
            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:
821
                instance.quant_config.packed_modules_mapping.update(
822
823
                    instance.packed_modules_mapping
                )
824

825
826
827
        return instance

    @staticmethod
828
    def _find_quant_config(*args, **kwargs) -> QuantizationConfig | None:
829
        """Find quant config passed through model constructor args"""
830
831
832
833
834
835
836
837
838
839
840
841
842
        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


843
844
845
@runtime_checkable
class SupportsTranscription(Protocol):
    """The interface required for all models that support transcription."""
846

847
848
    # Mapping from ISO639_1 language codes: language names
    supported_languages: ClassVar[Mapping[str, str]]
849
850
851

    supports_transcription: ClassVar[Literal[True]] = True

852
853
854
855
856
    supports_transcription_only: ClassVar[bool] = False
    """
    Transcription models can opt out of text generation by setting this to
    `True`.
    """
857
858
859
860
    supports_segment_timestamp: ClassVar[bool] = False
    """
    Enables the segment timestamp option for supported models by setting this to `True`.
    """
861

862
863
864
865
866
867
868
869
870
    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. "
871
872
                f"Valid choices are: {sorted(LANGUAGES.keys())}"
            )
873

874
    @classmethod
875
876
877
878
    def get_generation_prompt(
        cls,
        audio: np.ndarray,
        stt_config: SpeechToTextConfig,
879
        model_config: ModelConfig,
880
        language: str | None,
881
882
        task_type: Literal["transcribe", "translate"],
        request_prompt: str,
883
        to_language: str | None,
884
    ) -> PromptType:
885
886
887
        """Get the prompt for the ASR model.
        The model has control over the construction, as long as it
        returns a valid PromptType."""
888
889
890
        ...

    @classmethod
891
892
    def get_other_languages(cls) -> Mapping[str, str]:
        # other possible language codes from the whisper map
893
        return {k: v for k, v in LANGUAGES.items() if k not in cls.supported_languages}
894
895

    @classmethod
896
    def validate_language(cls, language: str | None) -> str | None:
897
        """
898
899
900
        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
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
        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 "
917
918
                f"{list(cls.supported_languages.keys())}."
            )
919

920
921
    @classmethod
    def get_speech_to_text_config(
922
        cls, model_config: ModelConfig, task_type: Literal["transcribe", "translate"]
923
    ) -> SpeechToTextConfig:
924
925
926
927
        """Get the speech to text config for the ASR model."""
        ...

    @classmethod
928
929
930
931
    def get_num_audio_tokens(
        cls,
        audio_duration_s: float,
        stt_config: SpeechToTextConfig,
932
        model_config: ModelConfig,
933
    ) -> int | None:
934
        """
935
        Map from audio duration to number of audio tokens produced by the ASR
936
937
938
939
940
        model, without running a forward pass.
        This is used for estimating the amount of processing for this audio.
        """
        return None

941
942
943

@overload
def supports_transcription(
944
945
    model: type[object],
) -> TypeIs[type[SupportsTranscription]]: ...
946
947
948


@overload
949
def supports_transcription(model: object) -> TypeIs[SupportsTranscription]: ...
950
951
952


def supports_transcription(
953
954
    model: type[object] | object,
) -> TypeIs[type[SupportsTranscription]] | TypeIs[SupportsTranscription]:
955
    return getattr(model, "supports_transcription", False)
956
957


958
@runtime_checkable
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
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):
1020
    """The interface required for models that support
1021
    EAGLE-3 speculative decoding."""
1022
1023
1024

    supports_eagle3: ClassVar[Literal[True]] = True
    """
1025
    A flag that indicates this model supports EAGLE-3 
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
    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
1036
        hidden states for EAGLE-3.
1037

1038
1039
        Args:
            layers: Tuple of layer indices that should output auxiliary
1040
                hidden states.
1041
1042
1043
1044
1045
1046
        """
        ...

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

1049
1050
1051
1052
1053
1054
1055
        Returns:
            Tuple of layer indices for auxiliary hidden state outputs.
        """
        ...


@overload
1056
def supports_eagle3(model: type[object]) -> TypeIs[type[SupportsEagle3]]: ...
1057
1058
1059


@overload
1060
def supports_eagle3(model: object) -> TypeIs[SupportsEagle3]: ...
1061
1062
1063


def supports_eagle3(
1064
1065
    model: type[object] | object,
) -> TypeIs[type[SupportsEagle3]] | TypeIs[SupportsEagle3]:
1066
    return isinstance(model, SupportsEagle3)
1067
1068
1069
1070
1071
1072
1073
1074
1075


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

1077
1078
1079
1080
1081
1082
1083
1084
    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],
1085
        mm_features: list["MultiModalFeatureSpec"],
1086
1087
1088
    ) -> tuple[torch.Tensor, int]:
        """
        Get M-RoPE input positions and delta value for this specific model.
1089

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

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

1097
        Returns:
1098
1099
            Tuple of `(llm_positions, mrope_position_delta)`
            - llm_positions: Tensor of shape `[3, num_tokens]` with T/H/W positions
1100
1101
1102
1103
1104
1105
            - mrope_position_delta: Delta for position calculations
        """
        ...


@overload
1106
def supports_mrope(model: type[object]) -> TypeIs[type[SupportsMRoPE]]: ...
1107
1108
1109


@overload
1110
def supports_mrope(model: object) -> TypeIs[SupportsMRoPE]: ...
1111
1112
1113


def supports_mrope(
1114
1115
    model: type[object] | object,
) -> TypeIs[type[SupportsMRoPE]] | TypeIs[SupportsMRoPE]:
1116
    return isinstance(model, SupportsMRoPE)
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165


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