interfaces.py 35.6 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 RendererConfig, 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
117
118
119
120
        if hasattr(self, "get_multimodal_embeddings"):
            logger.warning_once(
                "`get_multimodal_embeddings` for vLLM models is deprecated and will be "
                "removed in v0.13.0 or v1.0.0, whichever is earlier. Please rename "
                "this method to `embed_multimodal`."
            )
            return self.get_multimodal_embeddings(**kwargs)
121

122
    def get_language_model(self) -> VllmModel:
123
124
125
        """
        Returns the underlying language model used for text generation.

126
        This is typically the `torch.nn.Module` instance responsible for
127
128
129
130
131
132
133
        processing the merged multimodal embeddings and producing hidden states

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

134
    @overload
135
    def embed_input_ids(self, input_ids: Tensor) -> Tensor: ...
136
137

    @overload
138
    def embed_input_ids(
139
140
141
142
143
144
        self,
        input_ids: Tensor,
        multimodal_embeddings: MultiModalEmbeddings,
        *,
        is_multimodal: torch.Tensor,
        handle_oov_mm_token: bool = False,
145
    ) -> Tensor: ...
146

147
    def _embed_text_input_ids(
148
149
        self,
        input_ids: Tensor,
150
        embed_input_ids: Callable[[Tensor], Tensor],
151
        *,
152
        is_multimodal: Tensor | None,
153
154
155
156
        handle_oov_mm_token: bool,
    ) -> Tensor:
        if handle_oov_mm_token and is_multimodal is not None:
            is_text = ~is_multimodal
157
            text_embeds = embed_input_ids(input_ids[is_text])
158
159
160
161
162
163
164

            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)

165
        return embed_input_ids(input_ids)
166

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

189
        inputs_embeds = self._embed_text_input_ids(
190
            input_ids,
191
            self.get_language_model().embed_input_ids,
192
193
194
195
196
197
198
199
200
            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

        if is_multimodal is None:
            raise ValueError(
201
                "`embed_input_ids` now requires `is_multimodal` arg, "
202
                "please update your model runner according to "
203
204
                "https://github.com/vllm-project/vllm/pull/16229."
            )
205
206
207
208
209
210

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

212

213
214
215
216
217
218
@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.
    """
219

220
221
222
    supports_multimodal_pruning: ClassVar[Literal[True]] = True

    def recompute_mrope_positions(
223
224
225
226
227
        self,
        input_ids: list[int],
        multimodal_embeddings: MultiModalEmbeddings,
        mrope_positions: torch.LongTensor,
        num_computed_tokens: int,
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
    ) -> 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).
        """
        ...


252
@overload
253
def supports_multimodal(model: type[object]) -> TypeIs[type[SupportsMultiModal]]: ...
254
255
256


@overload
257
def supports_multimodal(model: object) -> TypeIs[SupportsMultiModal]: ...
258
259


260
def supports_multimodal(
261
262
    model: type[object] | object,
) -> TypeIs[type[SupportsMultiModal]] | TypeIs[SupportsMultiModal]:
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
    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."
            )

282
283
284
285
286
287
288
289
290
        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."
            )

291
    return res
292
293


294
def supports_multimodal_raw_input_only(model: type[object] | object) -> bool:
295
    return getattr(model, "supports_multimodal_raw_input_only", False)
296

297

298
def supports_multimodal_encoder_tp_data(model: type[object] | object) -> bool:
299
    return getattr(model, "supports_encoder_tp_data", False)
300
301


302
303
@overload
def supports_multimodal_pruning(
304
305
    model: type[object],
) -> TypeIs[type[SupportsMultiModalPruning]]: ...
306
307
308


@overload
309
def supports_multimodal_pruning(model: object) -> TypeIs[SupportsMultiModalPruning]: ...
310
311
312


def supports_multimodal_pruning(
313
314
    model: type[object] | object,
) -> TypeIs[type[SupportsMultiModalPruning]] | TypeIs[SupportsMultiModalPruning]:
315
316
317
    return getattr(model, "supports_multimodal_pruning", False)


318
319
320
321
322
323
324
325
326
327
328
329
330
331
@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
332
    def get_score_template(cls, query: str, document: str) -> str | None:
333
334
        """
        Generate a full prompt by populating the score template with query and document content.
335
        """  # noqa: E501
336
337
338
339
340
341
342
343
344
345
346
347
        ...

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


@overload
def supports_score_template(
348
349
    model: type[object],
) -> TypeIs[type[SupportsScoreTemplate]]: ...
350
351
352


@overload
353
def supports_score_template(model: object) -> TypeIs[SupportsScoreTemplate]: ...
354
355
356


def supports_score_template(
357
358
    model: type[object] | object,
) -> TypeIs[type[SupportsScoreTemplate]] | TypeIs[SupportsScoreTemplate]:
359
    return getattr(model, "supports_score_template", False)
360
361


362
363
364
365
@runtime_checkable
class SupportsLoRA(Protocol):
    """The interface required for all models that support LoRA."""

366
367
368
369
370
371
372
373
    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.
    """
374
    is_3d_moe_weight: ClassVar[bool] = False
375
376
    # The `embedding_module` and `embedding_padding_modules`
    # are empty by default.
377
    embedding_modules: ClassVar[dict[str, str]] = {}
378
    packed_modules_mapping: dict[str, list[str]] = {}
379
380
381
382
383
384
385
386


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

387
388
    packed_modules_mapping: dict[str, list[str]]
    embedding_modules: dict[str, str]
389
390
391


@overload
392
def supports_lora(model: type[object]) -> TypeIs[type[SupportsLoRA]]: ...
393
394
395


@overload
396
def supports_lora(model: object) -> TypeIs[SupportsLoRA]: ...
397
398
399


def supports_lora(
400
401
    model: type[object] | object,
) -> TypeIs[type[SupportsLoRA]] | TypeIs[SupportsLoRA]:
402
403
404
405
406
407
408
    result = _supports_lora(model)

    if not result:
        lora_attrs = (
            "packed_modules_mapping",
            "embedding_modules",
        )
409
        missing_attrs = tuple(attr for attr in lora_attrs if not hasattr(model, attr))
410
411
412
413
414
415
416
417
418
419
420
421
422

        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, "
423
424
425
                    "but does not set `supports_lora=True`.",
                    model,
                )
426
427
428
429

    return result


430
def _supports_lora(model: type[object] | object) -> bool:
431
432
433
434
    if isinstance(model, type):
        return isinstance(model, _SupportsLoRAType)

    return isinstance(model, SupportsLoRA)
435
436


437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
@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,
455
    ) -> IntermediateTensors:
456
457
458
459
460
        """Called when PP rank > 0 for profiling purposes."""
        ...

    def forward(
        self,
461
        *,
462
463
        intermediate_tensors: IntermediateTensors | None,
    ) -> IntermediateTensors | None:
464
        """
465
466
        Accept [`IntermediateTensors`][vllm.sequence.IntermediateTensors] when
        PP rank > 0.
467

468
469
        Return [`IntermediateTensors`][vllm.sequence.IntermediateTensors] only
        for the last PP rank.
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
        """
        ...


# 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,
485
    ) -> IntermediateTensors: ...
486
487
488

    def forward(
        self,
489
        *,
490
491
        intermediate_tensors: IntermediateTensors | None,
    ) -> Tensor | IntermediateTensors: ...
492
493
494


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


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


def supports_pp(
503
504
    model: type[object] | object,
) -> bool | TypeIs[type[SupportsPP]] | TypeIs[SupportsPP]:
505
506
507
508
509
510
    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 "
511
512
513
            "`intermediate_tensors` in its `forward` method",
            model,
        )
514
515

    if not supports_attributes:
516
517
        pp_attrs = ("make_empty_intermediate_tensors",)
        missing_attrs = tuple(attr for attr in pp_attrs if not hasattr(model, attr))
518
519
520
521
522
523
524
525
526
527
528
529
530

        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, "
531
532
533
                    "but does not set `supports_pp=True`.",
                    model,
                )
534
535
536
537

    return supports_attributes and supports_inspect


538
def _supports_pp_attributes(model: type[object] | object) -> bool:
539
540
541
542
543
544
    if isinstance(model, type):
        return isinstance(model, _SupportsPPType)

    return isinstance(model, SupportsPP)


545
def _supports_pp_inspect(model: type[object] | object) -> bool:
546
547
548
549
    model_forward = getattr(model, "forward", None)
    if not callable(model_forward):
        return False

550
    return supports_kw(model_forward, "intermediate_tensors")
551
552


553
554
555
556
557
558
559
560
@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
561
        for max_num_seqs, etc. True for e.g. both Mamba and Jamba.
562
563
564
565
    """


@overload
566
def has_inner_state(model: object) -> TypeIs[HasInnerState]: ...
567
568
569


@overload
570
def has_inner_state(model: type[object]) -> TypeIs[type[HasInnerState]]: ...
571
572
573


def has_inner_state(
574
575
    model: type[object] | object,
) -> TypeIs[type[HasInnerState]] | TypeIs[HasInnerState]:
576
    return getattr(model, "has_inner_state", False)
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592


@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
593
def is_attention_free(model: object) -> TypeIs[IsAttentionFree]: ...
594
595
596


@overload
597
def is_attention_free(model: type[object]) -> TypeIs[type[IsAttentionFree]]: ...
598
599
600


def is_attention_free(
601
602
    model: type[object] | object,
) -> TypeIs[type[IsAttentionFree]] | TypeIs[IsAttentionFree]:
603
    return getattr(model, "is_attention_free", False)
604
605


606
607
608
@runtime_checkable
class IsHybrid(Protocol):
    """The interface required for all models like Jamba that have both
609
    attention and mamba blocks, indicates that
610
611
612
613
614
615
616
617
    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' """

618
619
620
    @classmethod
    def get_mamba_state_shape_from_config(
        cls,
621
        vllm_config: VllmConfig,
622
623
624
625
626
627
628
629
630
631
632
633
634
    ) -> 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
        """
        ...

635
636

@overload
637
def is_hybrid(model: object) -> TypeIs[IsHybrid]: ...
638
639
640


@overload
641
def is_hybrid(model: type[object]) -> TypeIs[type[IsHybrid]]: ...
642
643
644


def is_hybrid(
645
646
    model: type[object] | object,
) -> TypeIs[type[IsHybrid]] | TypeIs[IsHybrid]:
647
    return getattr(model, "is_hybrid", False)
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
677
678
679
680
681
682
683
684
685
686
687
@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."""

688
689
690
    moe_layers: Iterable[nn.Module]
    """List of MoE layers in this model."""

691
692
693
694
695
696
697
698
    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.
699

700
701
702
703
704
705
706
707
708
709
710
711
712
        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.
        """
713
714
715
716
717
718
719
720
721
        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,
            )
722

723
724
725
726
    def update_physical_experts_metadata(
        self,
        num_physical_experts: int,
        num_local_physical_experts: int,
727
    ) -> None: ...
728

729
730

def is_mixture_of_experts(model: object) -> TypeIs[MixtureOfExperts]:
731
732
733
    return (
        isinstance(model, MixtureOfExperts) and getattr(model, "num_moe_layers", 0) > 0
    )
734
735


736
737
738
739
740
741
@runtime_checkable
class HasNoOps(Protocol):
    has_noops: ClassVar[Literal[True]] = True


@overload
742
def has_noops(model: object) -> TypeIs[HasNoOps]: ...
743
744
745


@overload
746
def has_noops(model: type[object]) -> TypeIs[type[HasNoOps]]: ...
747
748
749


def has_noops(
750
751
    model: type[object] | object,
) -> TypeIs[type[HasNoOps]] | TypeIs[HasNoOps]:
752
    return getattr(model, "has_noops", False)
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
@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)


783
784
785
786
787
788
789
790
791
@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(
792
793
    model: type[object],
) -> TypeIs[type[SupportsCrossEncoding]]: ...
794
795
796


@overload
797
def supports_cross_encoding(model: object) -> TypeIs[SupportsCrossEncoding]: ...
798
799
800


def _supports_cross_encoding(
801
802
    model: type[object] | object,
) -> TypeIs[type[SupportsCrossEncoding]] | TypeIs[SupportsCrossEncoding]:
803
    return getattr(model, "supports_cross_encoding", False)
804
805
806


def supports_cross_encoding(
807
808
    model: type[object] | object,
) -> TypeIs[type[SupportsCrossEncoding]] | TypeIs[SupportsCrossEncoding]:
809
    return is_pooling_model(model) and _supports_cross_encoding(model)
810
811


812
813
814
class SupportsQuant:
    """The interface required for all models that support quantization."""

815
816
817
    hf_to_vllm_mapper: ClassVar[WeightsMapper | None] = None
    packed_modules_mapping: ClassVar[dict[str, list[str]] | None] = None
    quant_config: QuantizationConfig | None = None
818

819
    def __new__(cls, *args, **kwargs) -> Self:
820
        instance = super().__new__(cls)
821
822

        # find config passed in arguments
823
824
        quant_config = cls._find_quant_config(*args, **kwargs)
        if quant_config is not None:
825
            # attach config to model for general use
826
            instance.quant_config = quant_config
827
828

            # apply model mappings to config for proper config-model matching
829
830
831
            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:
832
                instance.quant_config.packed_modules_mapping.update(
833
834
                    instance.packed_modules_mapping
                )
835

836
837
838
        return instance

    @staticmethod
839
    def _find_quant_config(*args, **kwargs) -> QuantizationConfig | None:
840
        """Find quant config passed through model constructor args"""
841
842
843
844
845
846
847
848
849
850
851
852
853
        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


854
855
856
@runtime_checkable
class SupportsTranscription(Protocol):
    """The interface required for all models that support transcription."""
857

858
859
    # Mapping from ISO639_1 language codes: language names
    supported_languages: ClassVar[Mapping[str, str]]
860
861
862

    supports_transcription: ClassVar[Literal[True]] = True

863
864
865
866
867
    supports_transcription_only: ClassVar[bool] = False
    """
    Transcription models can opt out of text generation by setting this to
    `True`.
    """
868
869
870
871
    supports_segment_timestamp: ClassVar[bool] = False
    """
    Enables the segment timestamp option for supported models by setting this to `True`.
    """
872

873
874
875
876
877
878
879
880
881
    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. "
882
883
                f"Valid choices are: {sorted(LANGUAGES.keys())}"
            )
884

885
    @classmethod
886
887
888
889
    def get_generation_prompt(
        cls,
        audio: np.ndarray,
        stt_config: SpeechToTextConfig,
890
        renderer_config: RendererConfig,
891
        language: str | None,
892
893
        task_type: Literal["transcribe", "translate"],
        request_prompt: str,
894
        to_language: str | None,
895
    ) -> PromptType:
896
897
898
        """Get the prompt for the ASR model.
        The model has control over the construction, as long as it
        returns a valid PromptType."""
899
900
901
        ...

    @classmethod
902
903
    def get_other_languages(cls) -> Mapping[str, str]:
        # other possible language codes from the whisper map
904
        return {k: v for k, v in LANGUAGES.items() if k not in cls.supported_languages}
905
906

    @classmethod
907
    def validate_language(cls, language: str | None) -> str | None:
908
        """
909
910
911
        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
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
        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 "
928
929
                f"{list(cls.supported_languages.keys())}."
            )
930

931
932
    @classmethod
    def get_speech_to_text_config(
933
934
935
        cls,
        renderer_config: RendererConfig,
        task_type: Literal["transcribe", "translate"],
936
    ) -> SpeechToTextConfig:
937
938
939
940
        """Get the speech to text config for the ASR model."""
        ...

    @classmethod
941
942
943
944
    def get_num_audio_tokens(
        cls,
        audio_duration_s: float,
        stt_config: SpeechToTextConfig,
945
        renderer_config: RendererConfig,
946
    ) -> int | None:
947
        """
948
        Map from audio duration to number of audio tokens produced by the ASR
949
950
951
952
953
        model, without running a forward pass.
        This is used for estimating the amount of processing for this audio.
        """
        return None

954
955
956

@overload
def supports_transcription(
957
958
    model: type[object],
) -> TypeIs[type[SupportsTranscription]]: ...
959
960
961


@overload
962
def supports_transcription(model: object) -> TypeIs[SupportsTranscription]: ...
963
964
965


def supports_transcription(
966
967
    model: type[object] | object,
) -> TypeIs[type[SupportsTranscription]] | TypeIs[SupportsTranscription]:
968
    return getattr(model, "supports_transcription", False)
969
970


971
@runtime_checkable
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
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
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):
1033
    """The interface required for models that support
1034
    EAGLE-3 speculative decoding."""
1035
1036
1037

    supports_eagle3: ClassVar[Literal[True]] = True
    """
1038
    A flag that indicates this model supports EAGLE-3 
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
    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
1049
        hidden states for EAGLE-3.
1050

1051
1052
        Args:
            layers: Tuple of layer indices that should output auxiliary
1053
                hidden states.
1054
1055
1056
1057
1058
1059
        """
        ...

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

1062
1063
1064
1065
1066
1067
1068
        Returns:
            Tuple of layer indices for auxiliary hidden state outputs.
        """
        ...


@overload
1069
def supports_eagle3(model: type[object]) -> TypeIs[type[SupportsEagle3]]: ...
1070
1071
1072


@overload
1073
def supports_eagle3(model: object) -> TypeIs[SupportsEagle3]: ...
1074
1075
1076


def supports_eagle3(
1077
1078
    model: type[object] | object,
) -> TypeIs[type[SupportsEagle3]] | TypeIs[SupportsEagle3]:
1079
    return isinstance(model, SupportsEagle3)
1080
1081
1082
1083
1084
1085
1086
1087
1088


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

1090
1091
1092
1093
1094
1095
1096
1097
    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],
1098
        mm_features: list["MultiModalFeatureSpec"],
1099
1100
1101
    ) -> tuple[torch.Tensor, int]:
        """
        Get M-RoPE input positions and delta value for this specific model.
1102

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

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

1110
        Returns:
1111
1112
            Tuple of `(llm_positions, mrope_position_delta)`
            - llm_positions: Tensor of shape `[3, num_tokens]` with T/H/W positions
1113
1114
1115
1116
1117
1118
            - mrope_position_delta: Delta for position calculations
        """
        ...


@overload
1119
def supports_mrope(model: type[object]) -> TypeIs[type[SupportsMRoPE]]: ...
1120
1121
1122


@overload
1123
def supports_mrope(model: object) -> TypeIs[SupportsMRoPE]: ...
1124
1125
1126


def supports_mrope(
1127
1128
    model: type[object] | object,
) -> TypeIs[type[SupportsMRoPE]] | TypeIs[SupportsMRoPE]:
1129
    return isinstance(model, SupportsMRoPE)
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
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178


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