interfaces.py 26.9 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 Iterable, Mapping, MutableSequence
5
from typing import (TYPE_CHECKING, ClassVar, Literal, Optional, Protocol,
6
                    Union, overload, runtime_checkable)
7

8
import numpy as np
9
import torch
10
from torch import Tensor
11
from transformers import PretrainedConfig
12
from transformers.models.whisper.tokenization_whisper import LANGUAGES
13
from typing_extensions import Self, TypeIs
14

15
from vllm.config import ModelConfig, SpeechToTextConfig
16
from vllm.inputs import TokensPrompt
17
from vllm.inputs.data import PromptType
18
from vllm.logger import init_logger
19
20
from vllm.model_executor.layers.quantization.base_config import (
    QuantizationConfig)
21
from vllm.utils import supports_kw
22

23
from .interfaces_base import is_pooling_model
24

25
if TYPE_CHECKING:
26
    from vllm.config import VllmConfig
27
    from vllm.model_executor.models.utils import WeightsMapper
28
29
    from vllm.sequence import IntermediateTensors

30
31
logger = init_logger(__name__)

32
33
34
35
36
37
38
39
MultiModalEmbeddings = Union[list[Tensor], Tensor, tuple[Tensor, ...]]
"""
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.
"""
40

41
42

@runtime_checkable
43
class SupportsMultiModal(Protocol):
44
    """The interface required for all multi-modal models."""
45

46
    supports_multimodal: ClassVar[Literal[True]] = True
47
    """
48
    A flag that indicates this model supports multi-modal inputs.
49
50
51
52
53

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

55
56
57
58
59
60
    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.
    """

61
62
63
64
65
66
    supports_encoder_tp_data: ClassVar[bool] = False
    """
    A flag that indicates whether this model supports
    `multimodal_config.mm_encoder_tp_mode="data"`.
    """

67
68
69
70
71
72
73
    @classmethod
    def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:
        """
        Get the placeholder text for the `i`th `modality` item in the prompt.
        """
        ...

74
75
    def get_multimodal_embeddings(self,
                                  **kwargs: object) -> MultiModalEmbeddings:
76
77
78
        """
        Returns multimodal embeddings generated from multimodal kwargs 
        to be merged with text embeddings.
79

80
        Note:
81
82
            The returned multimodal embeddings must be in the same order as
            the appearances of their corresponding multimodal data item in the
83
            input prompt.
84
85
86
        """
        ...

87
88
89
90
91
92
93
94
95
96
97
98
    def get_language_model(self) -> torch.nn.Module:
        """
        Returns the underlying language model used for text generation.

        This is typically the `torch.nn.Module` instance responsible for 
        processing the merged multimodal embeddings and producing hidden states

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

99
100
    def get_input_embeddings(
        self,
101
        input_ids: Tensor,
102
        multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
103
    ) -> Tensor:
104
105
106
107
108
109
110
        """
        Returns the input embeddings merged from the text embeddings from 
        input_ids and the multimodal embeddings generated from multimodal 
        kwargs.
        """
        ...

111
112

@overload
113
def supports_multimodal(
114
        model: type[object]) -> TypeIs[type[SupportsMultiModal]]:
115
116
117
118
    ...


@overload
119
def supports_multimodal(model: object) -> TypeIs[SupportsMultiModal]:
120
121
122
    ...


123
def supports_multimodal(
124
125
    model: Union[type[object], object],
) -> Union[TypeIs[type[SupportsMultiModal]], TypeIs[SupportsMultiModal]]:
126
    return getattr(model, "supports_multimodal", False)
127
128


129
def supports_multimodal_raw_input_only(
130
        model: Union[type[object], object]) -> bool:
131
    return getattr(model, "supports_multimodal_raw_input_only", False)
132

133

134
135
136
def supports_multimodal_encoder_tp_data(
        model: Union[type[object], object]) -> bool:
    return getattr(model, "supports_encoder_tp_data", False)
137
138


139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
@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
    def get_score_template(cls, query: str, document: str) -> Optional[str]:
        """
        Generate a full prompt by populating the score template with query and document content.
        """ # noqa: E501
        ...

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


@overload
def supports_score_template(
        model: type[object]) -> TypeIs[type[SupportsScoreTemplate]]:
    ...


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


def supports_score_template(
    model: Union[type[object], object],
) -> Union[TypeIs[type[SupportsScoreTemplate]], TypeIs[SupportsScoreTemplate]]:
181
    return getattr(model, "supports_score_template", False)
182
183


184
185
186
187
@runtime_checkable
class SupportsLoRA(Protocol):
    """The interface required for all models that support LoRA."""

188
189
190
191
192
193
194
195
    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.
    """
196
197
    # The `embedding_module` and `embedding_padding_modules`
    # are empty by default.
198
199
200
    embedding_modules: ClassVar[dict[str, str]] = {}
    embedding_padding_modules: ClassVar[list[str]] = []
    packed_modules_mapping: ClassVar[dict[str, list[str]]] = {}
201
202
203
204
205
206
207
208


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

209
210
211
    packed_modules_mapping: dict[str, list[str]]
    embedding_modules: dict[str, str]
    embedding_padding_modules: list[str]
212
213
214


@overload
215
def supports_lora(model: type[object]) -> TypeIs[type[SupportsLoRA]]:
216
217
218
219
    ...


@overload
220
def supports_lora(model: object) -> TypeIs[SupportsLoRA]:
221
222
223
224
    ...


def supports_lora(
225
226
    model: Union[type[object], object],
) -> Union[TypeIs[type[SupportsLoRA]], TypeIs[SupportsLoRA]]:
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
    result = _supports_lora(model)

    if not result:
        lora_attrs = (
            "packed_modules_mapping",
            "embedding_modules",
            "embedding_padding_modules",
        )
        missing_attrs = tuple(attr for attr in lora_attrs
                              if not hasattr(model, attr))

        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, "
                    "but does not set `supports_lora=True`.", model)

    return result


255
def _supports_lora(model: Union[type[object], object]) -> bool:
256
257
258
259
    if isinstance(model, type):
        return isinstance(model, _SupportsLoRAType)

    return isinstance(model, SupportsLoRA)
260
261


262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
@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,
    ) -> "IntermediateTensors":
        """Called when PP rank > 0 for profiling purposes."""
        ...

    def forward(
        self,
286
        *,
287
        intermediate_tensors: Optional["IntermediateTensors"],
288
    ) -> Union[Tensor, "IntermediateTensors"]:
289
        """
290
291
        Accept [`IntermediateTensors`][vllm.sequence.IntermediateTensors] when
        PP rank > 0.
292

293
294
        Return [`IntermediateTensors`][vllm.sequence.IntermediateTensors] only
        for the last PP rank.
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
        """
        ...


# 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,
    ) -> "IntermediateTensors":
        ...

    def forward(
        self,
315
        *,
316
        intermediate_tensors: Optional["IntermediateTensors"],
317
    ) -> Union[Tensor, "IntermediateTensors"]:
318
319
320
321
        ...


@overload
322
def supports_pp(model: type[object]) -> TypeIs[type[SupportsPP]]:
323
324
325
326
327
328
329
330
331
    ...


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


def supports_pp(
332
333
    model: Union[type[object], object],
) -> Union[bool, TypeIs[type[SupportsPP]], TypeIs[SupportsPP]]:
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
    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 "
            "`intermediate_tensors` in its `forward` method", model)

    if not supports_attributes:
        pp_attrs = ("make_empty_intermediate_tensors", )
        missing_attrs = tuple(attr for attr in pp_attrs
                              if not hasattr(model, attr))

        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, "
                    "but does not set `supports_pp=True`.", model)

    return supports_attributes and supports_inspect


364
def _supports_pp_attributes(model: Union[type[object], object]) -> bool:
365
366
367
368
369
370
    if isinstance(model, type):
        return isinstance(model, _SupportsPPType)

    return isinstance(model, SupportsPP)


371
def _supports_pp_inspect(model: Union[type[object], object]) -> bool:
372
373
374
375
    model_forward = getattr(model, "forward", None)
    if not callable(model_forward):
        return False

376
    return supports_kw(model_forward, "intermediate_tensors")
377
378


379
380
381
382
383
384
385
386
@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
387
        for max_num_seqs, etc. True for e.g. both Mamba and Jamba.
388
389
390
391
    """


@overload
392
def has_inner_state(model: object) -> TypeIs[HasInnerState]:
393
394
395
396
    ...


@overload
397
def has_inner_state(model: type[object]) -> TypeIs[type[HasInnerState]]:
398
399
400
401
    ...


def has_inner_state(
402
403
    model: Union[type[object], object]
) -> Union[TypeIs[type[HasInnerState]], TypeIs[HasInnerState]]:
404
    return getattr(model, "has_inner_state", False)
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425


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


@overload
426
def is_attention_free(model: type[object]) -> TypeIs[type[IsAttentionFree]]:
427
428
429
430
    ...


def is_attention_free(
431
432
    model: Union[type[object], object]
) -> Union[TypeIs[type[IsAttentionFree]], TypeIs[IsAttentionFree]]:
433
    return getattr(model, "is_attention_free", False)
434
435


436
437
438
439
440
441
442
443
444
445
446
447
@runtime_checkable
class IsHybrid(Protocol):
    """The interface required for all models like Jamba that have both
    attention and mamba blocks, indicates that 
    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' """

448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
    @classmethod
    def get_mamba_state_shape_from_config(
        cls,
        vllm_config: "VllmConfig",
        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
        """
        ...

467
468
469
470
471
472
473

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


@overload
474
def is_hybrid(model: type[object]) -> TypeIs[type[IsHybrid]]:
475
476
477
478
    ...


def is_hybrid(
479
480
    model: Union[type[object], object]
) -> Union[TypeIs[type[IsHybrid]], TypeIs[IsHybrid]]:
481
    return getattr(model, "is_hybrid", False)
482
483


484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
@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."""

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

546
547
548
549
550
551
552
    def update_physical_experts_metadata(
        self,
        num_physical_experts: int,
        num_local_physical_experts: int,
    ) -> None:
        ...

553
554
555
556
557

def is_mixture_of_experts(model: object) -> TypeIs[MixtureOfExperts]:
    return isinstance(model, MixtureOfExperts)


558
559
560
561
562
563
564
565
566
567
568
@runtime_checkable
class HasNoOps(Protocol):
    has_noops: ClassVar[Literal[True]] = True


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


@overload
569
def has_noops(model: type[object]) -> TypeIs[type[HasNoOps]]:
570
571
572
573
    ...


def has_noops(
574
575
    model: Union[type[object], object]
) -> Union[TypeIs[type[HasNoOps]], TypeIs[HasNoOps]]:
576
    return getattr(model, "has_noops", False)
577
578


579
580
581
582
583
584
585
586
587
@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(
588
        model: type[object]) -> TypeIs[type[SupportsCrossEncoding]]:
589
590
591
592
593
594
595
596
597
    ...


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


def _supports_cross_encoding(
598
599
    model: Union[type[object], object],
) -> Union[TypeIs[type[SupportsCrossEncoding]], TypeIs[SupportsCrossEncoding]]:
600
    return getattr(model, "supports_cross_encoding", False)
601
602
603


def supports_cross_encoding(
604
605
    model: Union[type[object], object],
) -> Union[TypeIs[type[SupportsCrossEncoding]], TypeIs[SupportsCrossEncoding]]:
606
    return is_pooling_model(model) and _supports_cross_encoding(model)
607
608


609
610
611
class SupportsQuant:
    """The interface required for all models that support quantization."""

612
613
    hf_to_vllm_mapper: ClassVar[Optional["WeightsMapper"]] = None
    packed_modules_mapping: ClassVar[Optional[dict[str, list[str]]]] = None
614
615
    quant_config: Optional[QuantizationConfig] = None

616
    def __new__(cls, *args, **kwargs) -> Self:
617
        instance = super().__new__(cls)
618
619

        # find config passed in arguments
620
621
        quant_config = cls._find_quant_config(*args, **kwargs)
        if quant_config is not None:
622
623

            # attach config to model for general use
624
            instance.quant_config = quant_config
625
626

            # apply model mappings to config for proper config-model matching
627
628
629
            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:
630
631
632
                instance.quant_config.packed_modules_mapping.update(
                    instance.packed_modules_mapping)

633
634
635
636
        return instance

    @staticmethod
    def _find_quant_config(*args, **kwargs) -> Optional[QuantizationConfig]:
637
        """Find quant config passed through model constructor args"""
638
639
640
641
642
643
644
645
646
647
648
649
650
        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


651
652
653
@runtime_checkable
class SupportsTranscription(Protocol):
    """The interface required for all models that support transcription."""
654
655
    # Mapping from ISO639_1 language codes: language names
    supported_languages: ClassVar[Mapping[str, str]]
656
657
658

    supports_transcription: ClassVar[Literal[True]] = True

659
660
661
662
663
664
    supports_transcription_only: ClassVar[bool] = False
    """
    Transcription models can opt out of text generation by setting this to
    `True`.
    """

665
666
667
668
669
670
671
672
673
674
675
    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. "
                f"Valid choices are: {sorted(LANGUAGES.keys())}")

676
    @classmethod
677
    def get_generation_prompt(cls, audio: np.ndarray,
Patrick von Platen's avatar
Patrick von Platen committed
678
                              stt_config: SpeechToTextConfig,
679
                              model_config: ModelConfig,
680
681
682
683
                              language: Optional[str],
                              task_type: Literal["transcribe", "translate"],
                              request_prompt: str,
                              to_language: Optional[str]) -> PromptType:
684
685
686
        """Get the prompt for the ASR model.
        The model has control over the construction, as long as it
        returns a valid PromptType."""
687
688
689
        ...

    @classmethod
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
    def get_other_languages(cls) -> Mapping[str, str]:
        # other possible language codes from the whisper map
        return {
            k: v
            for k, v in LANGUAGES.items() if k not in cls.supported_languages
        }

    @classmethod
    def validate_language(cls, language: Optional[str]) -> Optional[str]:
        """
        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 
        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 "
                f"{list(cls.supported_languages.keys())}.")
720

721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
    @classmethod
    def get_speech_to_text_config(
            cls, model_config: ModelConfig,
            task_type: Literal["transcribe",
                               "translate"]) -> SpeechToTextConfig:
        """Get the speech to text config for the ASR model."""
        ...

    @classmethod
    def get_num_audio_tokens(cls, audio_duration_s: float,
                             stt_config: SpeechToTextConfig,
                             model_config: ModelConfig) -> Optional[int]:
        """
        Map from audio duration to number of audio tokens produced by the ASR 
        model, without running a forward pass.
        This is used for estimating the amount of processing for this audio.
        """
        return None

740
741
742

@overload
def supports_transcription(
743
        model: type[object]) -> TypeIs[type[SupportsTranscription]]:
744
745
746
747
748
749
750
751
752
    ...


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


def supports_transcription(
753
754
    model: Union[type[object], object],
) -> Union[TypeIs[type[SupportsTranscription]], TypeIs[SupportsTranscription]]:
755
    return getattr(model, "supports_transcription", False)
756
757
758
759
760
761
762
763
764
765


@runtime_checkable
class SupportsV0Only(Protocol):
    """Models with this interface are not compatible with V1 vLLM."""

    supports_v0_only: ClassVar[Literal[True]] = True


@overload
766
def supports_v0_only(model: type[object]) -> TypeIs[type[SupportsV0Only]]:
767
768
769
770
771
772
773
774
775
    ...


@overload
def supports_v0_only(model: object) -> TypeIs[SupportsV0Only]:
    ...


def supports_v0_only(
776
777
    model: Union[type[object], object],
) -> Union[TypeIs[type[SupportsV0Only]], TypeIs[SupportsV0Only]]:
778
    return getattr(model, "supports_v0_only", False)
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802


@runtime_checkable
class SupportsEagle3(Protocol):
    """The interface required for models that support 
    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.
        
        Args:
            layers: Tuple of layer indices that should output auxiliary
803
                hidden states.
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
        """
        ...

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


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


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


def supports_eagle3(
    model: Union[type[object], object],
) -> Union[TypeIs[type[SupportsEagle3]], TypeIs[SupportsEagle3]]:
    return isinstance(model, SupportsEagle3)
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898


@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,
        image_grid_thw: Optional[Union[list[list[int]], torch.Tensor]],
        video_grid_thw: Optional[Union[list[list[int]], torch.Tensor]],
        second_per_grid_ts: Optional[list[float]] = None,
        context_len: int = 0,
        seq_len: Optional[int] = None,
        audio_feature_lengths: Optional[torch.Tensor] = None,
        use_audio_in_video: bool = False,
    ) -> tuple[torch.Tensor, int]:
        """
        Get M-RoPE input positions and delta value for this specific model.
        
        This method should be implemented by each model that supports M-RoPE
        to provide model-specific logic for computing input positions.
        
        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
            
        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
def supports_mrope(model: type[object]) -> TypeIs[type[SupportsMRoPE]]:
    ...


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


def supports_mrope(
    model: Union[type[object], object],
) -> Union[TypeIs[type[SupportsMRoPE]], TypeIs[SupportsMRoPE]]:
    return isinstance(model, SupportsMRoPE)