interfaces.py 31.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 import PretrainedConfig
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.sequence import IntermediateTensors
36
37
38
39
else:
    VllmConfig = object
    WeightsMapper = object
    IntermediateTensors = object
40

41
42
logger = init_logger(__name__)

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

52
53

@runtime_checkable
54
class SupportsMultiModal(Protocol):
55
    """The interface required for all multi-modal models."""
56

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

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

66
67
68
69
70
71
    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.
    """

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

78
79
80
81
82
83
    merge_by_field_config: ClassVar[bool] = False
    """
    A flag that indicates which implementation of
    `vllm.multimodal.utils.group_mm_kwargs_by_modality` to use.
    """

84
85
86
87
88
    multimodal_cpu_fields: ClassVar[Set[str]] = frozenset()
    """
    A set indicating CPU-only multimodal fields.
    """

89
    @classmethod
90
    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
91
92
93
94
95
        """
        Get the placeholder text for the `i`th `modality` item in the prompt.
        """
        ...

96
    def get_multimodal_embeddings(self, **kwargs: object) -> MultiModalEmbeddings:
97
        """
98
        Returns multimodal embeddings generated from multimodal kwargs
99
        to be merged with text embeddings.
100

101
        Note:
102
103
            The returned multimodal embeddings must be in the same order as
            the appearances of their corresponding multimodal data item in the
104
            input prompt.
105
106
107
        """
        ...

108
    def get_language_model(self) -> VllmModel:
109
110
111
        """
        Returns the underlying language model used for text generation.

112
        This is typically the `torch.nn.Module` instance responsible for
113
114
115
116
117
118
119
        processing the merged multimodal embeddings and producing hidden states

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

120
    @overload
121
    def get_input_embeddings(self, input_ids: Tensor) -> Tensor: ...
122
123
124
125
126
127
128
129
130

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

    def _get_text_embeddings(
        self,
        input_ids: Tensor,
        get_input_embeddings: Callable[[Tensor], Tensor],
        *,
138
        is_multimodal: Tensor | None,
139
140
141
142
143
144
145
146
147
148
149
150
151
152
        handle_oov_mm_token: bool,
    ) -> Tensor:
        if handle_oov_mm_token and is_multimodal is not None:
            is_text = ~is_multimodal
            text_embeds = get_input_embeddings(input_ids[is_text])

            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)

        return get_input_embeddings(input_ids)

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

        inputs_embeds = self._get_text_embeddings(
            input_ids,
            self.get_language_model().get_input_embeddings,
            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(
                "`get_input_embeddings` now requires `is_multimodal` arg, "
                "please update your model runner according to "
189
190
                "https://github.com/vllm-project/vllm/pull/16229."
            )
191
192
193
194
195
196

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

198

199
200
201
202
203
204
@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.
    """
205

206
207
208
    supports_multimodal_pruning: ClassVar[Literal[True]] = True

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


238
@overload
239
def supports_multimodal(model: type[object]) -> TypeIs[type[SupportsMultiModal]]: ...
240
241
242


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


246
def supports_multimodal(
247
248
    model: type[object] | object,
) -> TypeIs[type[SupportsMultiModal]] | TypeIs[SupportsMultiModal]:
249
    return getattr(model, "supports_multimodal", False)
250
251


252
def supports_multimodal_raw_input_only(model: type[object] | object) -> bool:
253
    return getattr(model, "supports_multimodal_raw_input_only", False)
254

255

256
def supports_multimodal_encoder_tp_data(model: type[object] | object) -> bool:
257
    return getattr(model, "supports_encoder_tp_data", False)
258
259


260
261
@overload
def supports_multimodal_pruning(
262
263
    model: type[object],
) -> TypeIs[type[SupportsMultiModalPruning]]: ...
264
265
266


@overload
267
def supports_multimodal_pruning(model: object) -> TypeIs[SupportsMultiModalPruning]: ...
268
269
270


def supports_multimodal_pruning(
271
272
    model: type[object] | object,
) -> TypeIs[type[SupportsMultiModalPruning]] | TypeIs[SupportsMultiModalPruning]:
273
274
275
    return getattr(model, "supports_multimodal_pruning", False)


276
277
278
279
280
281
282
283
284
285
286
287
288
289
@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
290
    def get_score_template(cls, query: str, document: str) -> str | None:
291
292
        """
        Generate a full prompt by populating the score template with query and document content.
293
        """  # noqa: E501
294
295
296
297
298
299
300
301
302
303
304
305
        ...

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


@overload
def supports_score_template(
306
307
    model: type[object],
) -> TypeIs[type[SupportsScoreTemplate]]: ...
308
309
310


@overload
311
def supports_score_template(model: object) -> TypeIs[SupportsScoreTemplate]: ...
312
313
314


def supports_score_template(
315
316
    model: type[object] | object,
) -> TypeIs[type[SupportsScoreTemplate]] | TypeIs[SupportsScoreTemplate]:
317
    return getattr(model, "supports_score_template", False)
318
319


320
321
322
323
@runtime_checkable
class SupportsLoRA(Protocol):
    """The interface required for all models that support LoRA."""

324
325
326
327
328
329
330
331
    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.
    """
332
333
    # The `embedding_module` and `embedding_padding_modules`
    # are empty by default.
334
335
    embedding_modules: ClassVar[dict[str, str]] = {}
    embedding_padding_modules: ClassVar[list[str]] = []
336
    packed_modules_mapping: dict[str, list[str]] = {}
337
338
339
340
341
342
343
344


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

345
346
347
    packed_modules_mapping: dict[str, list[str]]
    embedding_modules: dict[str, str]
    embedding_padding_modules: list[str]
348
349
350


@overload
351
def supports_lora(model: type[object]) -> TypeIs[type[SupportsLoRA]]: ...
352
353
354


@overload
355
def supports_lora(model: object) -> TypeIs[SupportsLoRA]: ...
356
357
358


def supports_lora(
359
360
    model: type[object] | object,
) -> TypeIs[type[SupportsLoRA]] | TypeIs[SupportsLoRA]:
361
362
363
364
365
366
367
368
    result = _supports_lora(model)

    if not result:
        lora_attrs = (
            "packed_modules_mapping",
            "embedding_modules",
            "embedding_padding_modules",
        )
369
        missing_attrs = tuple(attr for attr in lora_attrs if not hasattr(model, attr))
370
371
372
373
374
375
376
377
378
379
380
381
382

        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, "
383
384
385
                    "but does not set `supports_lora=True`.",
                    model,
                )
386
387
388
389

    return result


390
def _supports_lora(model: type[object] | object) -> bool:
391
392
393
394
    if isinstance(model, type):
        return isinstance(model, _SupportsLoRAType)

    return isinstance(model, SupportsLoRA)
395
396


397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
@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,
415
    ) -> IntermediateTensors:
416
417
418
419
420
        """Called when PP rank > 0 for profiling purposes."""
        ...

    def forward(
        self,
421
        *,
422
423
        intermediate_tensors: IntermediateTensors | None,
    ) -> IntermediateTensors | None:
424
        """
425
426
        Accept [`IntermediateTensors`][vllm.sequence.IntermediateTensors] when
        PP rank > 0.
427

428
429
        Return [`IntermediateTensors`][vllm.sequence.IntermediateTensors] only
        for the last PP rank.
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
        """
        ...


# 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,
445
    ) -> IntermediateTensors: ...
446
447
448

    def forward(
        self,
449
        *,
450
451
        intermediate_tensors: IntermediateTensors | None,
    ) -> Tensor | IntermediateTensors: ...
452
453
454


@overload
455
def supports_pp(model: type[object]) -> TypeIs[type[SupportsPP]]: ...
456
457
458


@overload
459
def supports_pp(model: object) -> TypeIs[SupportsPP]: ...
460
461
462


def supports_pp(
463
464
    model: type[object] | object,
) -> bool | TypeIs[type[SupportsPP]] | TypeIs[SupportsPP]:
465
466
467
468
469
470
    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 "
471
472
473
            "`intermediate_tensors` in its `forward` method",
            model,
        )
474
475

    if not supports_attributes:
476
477
        pp_attrs = ("make_empty_intermediate_tensors",)
        missing_attrs = tuple(attr for attr in pp_attrs if not hasattr(model, attr))
478
479
480
481
482
483
484
485
486
487
488
489
490

        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, "
491
492
493
                    "but does not set `supports_pp=True`.",
                    model,
                )
494
495
496
497

    return supports_attributes and supports_inspect


498
def _supports_pp_attributes(model: type[object] | object) -> bool:
499
500
501
502
503
504
    if isinstance(model, type):
        return isinstance(model, _SupportsPPType)

    return isinstance(model, SupportsPP)


505
def _supports_pp_inspect(model: type[object] | object) -> bool:
506
507
508
509
    model_forward = getattr(model, "forward", None)
    if not callable(model_forward):
        return False

510
    return supports_kw(model_forward, "intermediate_tensors")
511
512


513
514
515
516
517
518
519
520
@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
521
        for max_num_seqs, etc. True for e.g. both Mamba and Jamba.
522
523
524
525
    """


@overload
526
def has_inner_state(model: object) -> TypeIs[HasInnerState]: ...
527
528
529


@overload
530
def has_inner_state(model: type[object]) -> TypeIs[type[HasInnerState]]: ...
531
532
533


def has_inner_state(
534
535
    model: type[object] | object,
) -> TypeIs[type[HasInnerState]] | TypeIs[HasInnerState]:
536
    return getattr(model, "has_inner_state", False)
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552


@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
553
def is_attention_free(model: object) -> TypeIs[IsAttentionFree]: ...
554
555
556


@overload
557
def is_attention_free(model: type[object]) -> TypeIs[type[IsAttentionFree]]: ...
558
559
560


def is_attention_free(
561
562
    model: type[object] | object,
) -> TypeIs[type[IsAttentionFree]] | TypeIs[IsAttentionFree]:
563
    return getattr(model, "is_attention_free", False)
564
565


566
567
568
@runtime_checkable
class IsHybrid(Protocol):
    """The interface required for all models like Jamba that have both
569
    attention and mamba blocks, indicates that
570
571
572
573
574
575
576
577
    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' """

578
579
580
    @classmethod
    def get_mamba_state_shape_from_config(
        cls,
581
        vllm_config: VllmConfig,
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
        use_v1: bool = True,
    ) -> tuple[tuple[int, int], tuple[int, int, int]]:
        """Calculate shapes for Mamba's convolutional and state caches.

        Args:
            vllm_config: vLLM config
            use_v1: Get shapes for V1 (or V0)

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

597
598

@overload
599
def is_hybrid(model: object) -> TypeIs[IsHybrid]: ...
600
601
602


@overload
603
def is_hybrid(model: type[object]) -> TypeIs[type[IsHybrid]]: ...
604
605
606


def is_hybrid(
607
608
    model: type[object] | object,
) -> TypeIs[type[IsHybrid]] | TypeIs[IsHybrid]:
609
    return getattr(model, "is_hybrid", False)
610
611


612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
@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."""

650
651
652
    moe_layers: Iterable[nn.Module]
    """List of MoE layers in this model."""

653
654
655
656
657
658
659
660
    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.
661

662
663
664
665
666
667
668
669
670
671
672
673
674
        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.
        """
675
676
677
678
679
680
681
682
683
        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,
            )
684

685
686
687
688
    def update_physical_experts_metadata(
        self,
        num_physical_experts: int,
        num_local_physical_experts: int,
689
    ) -> None: ...
690

691
692

def is_mixture_of_experts(model: object) -> TypeIs[MixtureOfExperts]:
693
694
695
    return (
        isinstance(model, MixtureOfExperts) and getattr(model, "num_moe_layers", 0) > 0
    )
696
697


698
699
700
701
702
703
@runtime_checkable
class HasNoOps(Protocol):
    has_noops: ClassVar[Literal[True]] = True


@overload
704
def has_noops(model: object) -> TypeIs[HasNoOps]: ...
705
706
707


@overload
708
def has_noops(model: type[object]) -> TypeIs[type[HasNoOps]]: ...
709
710
711


def has_noops(
712
713
    model: type[object] | object,
) -> TypeIs[type[HasNoOps]] | TypeIs[HasNoOps]:
714
    return getattr(model, "has_noops", False)
715
716


717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
@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)


745
746
747
748
749
750
751
752
753
@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(
754
755
    model: type[object],
) -> TypeIs[type[SupportsCrossEncoding]]: ...
756
757
758


@overload
759
def supports_cross_encoding(model: object) -> TypeIs[SupportsCrossEncoding]: ...
760
761
762


def _supports_cross_encoding(
763
764
    model: type[object] | object,
) -> TypeIs[type[SupportsCrossEncoding]] | TypeIs[SupportsCrossEncoding]:
765
    return getattr(model, "supports_cross_encoding", False)
766
767
768


def supports_cross_encoding(
769
770
    model: type[object] | object,
) -> TypeIs[type[SupportsCrossEncoding]] | TypeIs[SupportsCrossEncoding]:
771
    return is_pooling_model(model) and _supports_cross_encoding(model)
772
773


774
775
776
class SupportsQuant:
    """The interface required for all models that support quantization."""

777
778
779
    hf_to_vllm_mapper: ClassVar[WeightsMapper | None] = None
    packed_modules_mapping: ClassVar[dict[str, list[str]] | None] = None
    quant_config: QuantizationConfig | None = None
780

781
    def __new__(cls, *args, **kwargs) -> Self:
782
        instance = super().__new__(cls)
783
784

        # find config passed in arguments
785
786
        quant_config = cls._find_quant_config(*args, **kwargs)
        if quant_config is not None:
787
            # attach config to model for general use
788
            instance.quant_config = quant_config
789
790

            # apply model mappings to config for proper config-model matching
791
792
793
            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:
794
                instance.quant_config.packed_modules_mapping.update(
795
796
                    instance.packed_modules_mapping
                )
797

798
799
800
        return instance

    @staticmethod
801
    def _find_quant_config(*args, **kwargs) -> QuantizationConfig | None:
802
        """Find quant config passed through model constructor args"""
803
804
805
806
807
808
809
810
811
812
813
814
815
        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


816
817
818
@runtime_checkable
class SupportsTranscription(Protocol):
    """The interface required for all models that support transcription."""
819

820
821
    # Mapping from ISO639_1 language codes: language names
    supported_languages: ClassVar[Mapping[str, str]]
822
823
824

    supports_transcription: ClassVar[Literal[True]] = True

825
826
827
828
829
830
    supports_transcription_only: ClassVar[bool] = False
    """
    Transcription models can opt out of text generation by setting this to
    `True`.
    """

831
832
833
834
835
836
837
838
839
    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. "
840
841
                f"Valid choices are: {sorted(LANGUAGES.keys())}"
            )
842

843
    @classmethod
844
845
846
847
848
    def get_generation_prompt(
        cls,
        audio: np.ndarray,
        stt_config: SpeechToTextConfig,
        model_config: ModelConfig,
849
        language: str | None,
850
851
        task_type: Literal["transcribe", "translate"],
        request_prompt: str,
852
        to_language: str | None,
853
    ) -> PromptType:
854
855
856
        """Get the prompt for the ASR model.
        The model has control over the construction, as long as it
        returns a valid PromptType."""
857
858
859
        ...

    @classmethod
860
861
    def get_other_languages(cls) -> Mapping[str, str]:
        # other possible language codes from the whisper map
862
        return {k: v for k, v in LANGUAGES.items() if k not in cls.supported_languages}
863
864

    @classmethod
865
    def validate_language(cls, language: str | None) -> str | None:
866
        """
867
868
869
        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
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
        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 "
886
887
                f"{list(cls.supported_languages.keys())}."
            )
888

889
890
    @classmethod
    def get_speech_to_text_config(
891
892
        cls, model_config: ModelConfig, task_type: Literal["transcribe", "translate"]
    ) -> SpeechToTextConfig:
893
894
895
896
        """Get the speech to text config for the ASR model."""
        ...

    @classmethod
897
898
899
900
901
    def get_num_audio_tokens(
        cls,
        audio_duration_s: float,
        stt_config: SpeechToTextConfig,
        model_config: ModelConfig,
902
    ) -> int | None:
903
        """
904
        Map from audio duration to number of audio tokens produced by the ASR
905
906
907
908
909
        model, without running a forward pass.
        This is used for estimating the amount of processing for this audio.
        """
        return None

910
911
912

@overload
def supports_transcription(
913
914
    model: type[object],
) -> TypeIs[type[SupportsTranscription]]: ...
915
916
917


@overload
918
def supports_transcription(model: object) -> TypeIs[SupportsTranscription]: ...
919
920
921


def supports_transcription(
922
923
    model: type[object] | object,
) -> TypeIs[type[SupportsTranscription]] | TypeIs[SupportsTranscription]:
924
    return getattr(model, "supports_transcription", False)
925
926


927
928
@runtime_checkable
class SupportsEagle3(Protocol):
929
    """The interface required for models that support
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
    EAGLE3 speculative decoding."""

    supports_eagle3: ClassVar[Literal[True]] = True
    """
    A flag that indicates this model supports EAGLE3 
    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
        hidden states for EAGLE3.
946

947
948
        Args:
            layers: Tuple of layer indices that should output auxiliary
949
                hidden states.
950
951
952
953
954
955
956
        """
        ...

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

958
959
960
961
962
963
964
        Returns:
            Tuple of layer indices for auxiliary hidden state outputs.
        """
        ...


@overload
965
def supports_eagle3(model: type[object]) -> TypeIs[type[SupportsEagle3]]: ...
966
967
968


@overload
969
def supports_eagle3(model: object) -> TypeIs[SupportsEagle3]: ...
970
971
972


def supports_eagle3(
973
974
    model: type[object] | object,
) -> TypeIs[type[SupportsEagle3]] | TypeIs[SupportsEagle3]:
975
    return isinstance(model, SupportsEagle3)
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994


@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.
    
    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],
        hf_config: PretrainedConfig,
995
996
997
        image_grid_thw: list[list[int]] | torch.Tensor | None,
        video_grid_thw: list[list[int]] | torch.Tensor | None,
        second_per_grid_ts: list[float] | None = None,
998
        context_len: int = 0,
999
1000
        seq_len: int | None = None,
        audio_feature_lengths: torch.Tensor | None = None,
1001
1002
1003
1004
        use_audio_in_video: bool = False,
    ) -> tuple[torch.Tensor, int]:
        """
        Get M-RoPE input positions and delta value for this specific model.
1005

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

1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
        Args:
            input_tokens: List of input token IDs
            hf_config: HuggingFace model configuration
            image_grid_thw: Image grid dimensions (t, h, w)
            video_grid_thw: Video grid dimensions (t, h, w)
            second_per_grid_ts: Seconds per grid timestep for videos
            context_len: Context length
            seq_len: Sequence length
            audio_feature_lengths: Audio feature lengths for multimodal models
            use_audio_in_video: Whether to use audio in video for interleaving
1019

1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
        Returns:
            Tuple of (llm_positions, mrope_position_delta)
            - llm_positions: Tensor of shape [3, num_tokens]
                with T/H/W positions
            - mrope_position_delta: Delta for position calculations
        """
        ...


@overload
1030
def supports_mrope(model: type[object]) -> TypeIs[type[SupportsMRoPE]]: ...
1031
1032
1033


@overload
1034
def supports_mrope(model: object) -> TypeIs[SupportsMRoPE]: ...
1035
1036
1037


def supports_mrope(
1038
1039
    model: type[object] | object,
) -> TypeIs[type[SupportsMRoPE]] | TypeIs[SupportsMRoPE]:
1040
    return isinstance(model, SupportsMRoPE)