interfaces.py 45.3 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
import asyncio
5
6
7
8
9
10
11
12
from collections.abc import (
    AsyncGenerator,
    Callable,
    Iterable,
    Mapping,
    MutableSequence,
    Sequence,
)
13
from contextlib import ExitStack, contextmanager, nullcontext
14
15
16
17
18
from typing import (
    TYPE_CHECKING,
    ClassVar,
    Literal,
    Protocol,
19
    TypeAlias,
20
21
22
    overload,
    runtime_checkable,
)
23

24
import numpy as np
25
import torch
26
import torch.nn as nn
27
from torch import Tensor
28
from transformers.models.whisper.tokenization_whisper import LANGUAGES
29
from typing_extensions import Self, TypeIs
30

31
from vllm.config import ModelConfig, SpeechToTextConfig
32
from vllm.inputs import TokensPrompt
33
from vllm.inputs.data import PromptType
34
from vllm.logger import init_logger
35
from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc
36
from vllm.model_executor.layers.quantization import QuantizationConfig
37
from vllm.utils.collection_utils import common_prefix
38
from vllm.utils.func_utils import supports_kw
39

40
from .interfaces_base import VllmModel, is_pooling_model
41

42
if TYPE_CHECKING:
43
    from vllm.config import VllmConfig
44
    from vllm.model_executor.models.utils import WeightsMapper
45
    from vllm.multimodal.inputs import MultiModalFeatureSpec
46
    from vllm.multimodal.registry import _ProcessorFactories
47
    from vllm.sequence import IntermediateTensors
48
49
50
else:
    VllmConfig = object
    WeightsMapper = object
51
    MultiModalFeatureSpec = object
52
    _ProcessorFactories = object
53
    IntermediateTensors = object
54

55
56
logger = init_logger(__name__)

57
MultiModalEmbeddings: TypeAlias = list[Tensor] | Tensor | tuple[Tensor, ...]
58
59
60
61
62
63
64
"""
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.
"""
65

66

67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
def _require_is_multimodal(is_multimodal: Tensor | None) -> Tensor:
    """
    A helper function to be used in the context of
    [vllm.model_executor.models.interfaces.SupportsMultiModal.embed_input_ids][]
    to provide a better error message.
    """
    if is_multimodal is None:
        raise ValueError(
            "`embed_input_ids` now requires `is_multimodal` arg, "
            "please update your model runner according to "
            "https://github.com/vllm-project/vllm/pull/16229."
        )

    return is_multimodal


83
84
# Cache results of `SupportsMultiModal.get_language_model`
_language_model_by_module = dict[nn.Module, VllmModel]()
85
86


87
@runtime_checkable
88
class SupportsMultiModal(Protocol):
89
    """The interface required for all multi-modal models."""
90

91
    supports_multimodal: ClassVar[Literal[True]] = True
92
    """
93
    A flag that indicates this model supports multi-modal inputs.
94
95
96
97
98

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

100
101
102
103
104
105
    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.
    """

106
107
108
109
110
111
    supports_encoder_tp_data: ClassVar[bool] = False
    """
    A flag that indicates whether this model supports
    `multimodal_config.mm_encoder_tp_mode="data"`.
    """

Patrick von Platen's avatar
Patrick von Platen committed
112
113
114
115
116
117
    requires_raw_input_tokens: ClassVar[bool] = False
    """
    A flag that indicates this model processes input id tokens
    in their raw form and not input embeddings.
    """

118
119
120
121
122
    _processor_factory: ClassVar[_ProcessorFactories]
    """
    Set internally by `MultiModalRegistry.register_processor`.
    """

123
124
125
126
127
128
129
130
131
132
    _language_model_names: list[str] = []
    """
    Set internally by `_mark_language_model`.
    """

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

133
134
135
136
137
138
139
    _has_oov_mm_tokens: bool = False
    """
    In general, this should be set at init time by invoking
    `configure_mm_token_handling` models & passing all potentially
    OOV multimodal tokens.
    """

140
    @classmethod
141
    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
142
143
144
145
146
        """
        Get the placeholder text for the `i`th `modality` item in the prompt.
        """
        ...

147
    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
148
        """
149
        Returns multimodal embeddings generated from multimodal kwargs
150
        to be merged with text embeddings.
151

152
        Note:
153
154
            The returned multimodal embeddings must be in the same order as
            the appearances of their corresponding multimodal data item in the
155
            input prompt.
156
        """
157
        ...
158

159
160
161
162
163
164
165
166
167
168
169
    def configure_mm_token_handling(self, vocab_size: int, mm_token_ids: list[int]):
        """Check if any multimodal tokens are out of vocabulary. If so, we will
        explicitly mask all multimodal tokens out when computing text embeddings,
        since the multimodal embeddings will be scattered over the results.
        """
        self._has_oov_mm_tokens = any(tok_id >= vocab_size for tok_id in mm_token_ids)
        logger.info(
            "Contains out of vocabulary multimodal tokens? %s",
            self._has_oov_mm_tokens,
        )

170
    def get_language_model(self) -> VllmModel:
171
172
173
        """
        Returns the underlying language model used for text generation.

174
        This is typically the `torch.nn.Module` instance responsible for
175
176
177
178
179
        processing the merged multimodal embeddings and producing hidden states

        Returns:
            torch.nn.Module: The core language model component.
        """
180
181
182
183
        # Cached
        if self in _language_model_by_module:
            return _language_model_by_module[self]

184
        if self._language_model_names:
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
            mod = self
            for attr in common_prefix(
                [name.split(".") for name in self._language_model_names]
            ):
                if attr:
                    mod = getattr(mod, attr)

            if mod is not self and hasattr(mod, "embed_input_ids"):
                _language_model_by_module[self] = mod
                return mod

        # Fallback
        for mod in self.children():
            if hasattr(mod, "embed_input_ids"):
                _language_model_by_module[self] = mod
                return mod
201
202
203

        raise NotImplementedError(
            f"No language model found in {type(self).__name__}! "
204
            "You should initialize it via `_mark_language_model`."
205
206
207
        )

    @contextmanager
208
209
210
211
212
213
    def _mark_language_model(
        self,
        vllm_config: VllmConfig,
        *,
        targets: type[nn.Module] | tuple[type[nn.Module], ...] | None = None,
    ):
214
        """
215
216
217
218
219
        Mark each child module that was assigned to this model during this context
        as a language model component.

        Language model components are automatically skipped in `--mm-encoder-only`
        mode.
220

221
222
223
224
        If `targets` is set, instead include descendants that are an instance
        of `targets`, even if they aren't direct children.
        """
        from .utils import StageMissingLayer, collect_children, no_init_weights
225

226
        mm_config = vllm_config.model_config.multimodal_config
227

228
        with collect_children(self, targets=targets) as children_names:  # noqa: SIM117
229
            with (
230
231
232
233
234
                no_init_weights(
                    self,
                    lambda mod: StageMissingLayer("language_model", mod),
                    targets=targets,
                )
235
236
237
238
239
240
241
242
                if mm_config.mm_encoder_only
                else nullcontext()
            ):
                yield

        self._language_model_names = children_names

    @contextmanager
243
244
245
246
247
248
249
    def _mark_tower_model(
        self,
        vllm_config: VllmConfig,
        modalities: set[str] | str,
        *,
        targets: type[nn.Module] | tuple[type[nn.Module], ...] | None = None,
    ):
250
        """
251
252
253
254
255
256
257
258
        Mark each child module that was assigned to this model during this context
        as a tower model component.

        Tower model components are automatically skipped when `--limit-mm-per-prompt`
        is set to zero for all of their modalities.

        If `targets` is set, instead include descendants that are an instance
        of `targets`, even if they aren't direct children.
259
        """
260
261
        from .utils import StageMissingLayer, collect_children, no_init_weights

262
263
264
        if isinstance(modalities, str):
            modalities = {modalities}

265
266
267
268
        if modalities == {"image", "video"}:
            stage_name = "vision_tower"
        else:
            stage_name = "_".join([*modalities, "tower"])
269

270
        mm_config = vllm_config.model_config.multimodal_config
271

272
        with collect_children(self, targets=targets) as children_names:  # noqa: SIM117
273
            with (
274
275
276
277
278
                no_init_weights(
                    self,
                    lambda mod: StageMissingLayer(stage_name, mod),
                    targets=targets,
                )
279
280
281
282
283
284
                if all(mm_config.get_limit_per_prompt(m) == 0 for m in modalities)
                else nullcontext()
            ):
                yield

        self._tower_model_names = children_names
285

286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
    @contextmanager
    def _mark_composite_model(
        self,
        vllm_config: VllmConfig,
        *,
        language_targets: type[nn.Module] | tuple[type[nn.Module], ...],
        tower_targets: dict[str, type[nn.Module] | tuple[type[nn.Module], ...]],
    ):
        """
        Composite wrapper over `_mark_language_model` and
        `_mark_tower_model` by modality.
        """
        with ExitStack() as stack:
            stack.enter_context(
                self._mark_language_model(
                    vllm_config,
                    targets=language_targets,
                )
            )

            for modality, modality_targets in tower_targets.items():
                stack.enter_context(
                    self._mark_tower_model(
                        vllm_config,
                        modality,
                        targets=modality_targets,
                    )
                )

            yield

317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
    def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
        """
        Implement this function to enable LoRA support
        for the tower module of the multi-modal model.
        Given the number of image tokens, output the number of
        multi-modal encoder tokens.
        """
        ...

    def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
        """
        Implement this function to enable LoRA support
        for the connector module of the multi-modal model.
        Given the number of vision tokens, output the number of
        multi-modal connector tokens.
        """
        ...

335
    @overload
336
    def embed_input_ids(self, input_ids: Tensor) -> Tensor: ...
337
338

    @overload
339
    def embed_input_ids(
340
341
342
343
344
        self,
        input_ids: Tensor,
        multimodal_embeddings: MultiModalEmbeddings,
        *,
        is_multimodal: torch.Tensor,
345
    ) -> Tensor: ...
346

347
    def _embed_text_input_ids(
348
349
        self,
        input_ids: Tensor,
350
        embed_input_ids: Callable[[Tensor], Tensor],
351
        *,
352
        is_multimodal: Tensor | None,
353
    ) -> Tensor:
354
355
356
357
358
359
360
        if is_multimodal is not None and self._has_oov_mm_tokens:
            # Force all input IDs to be in vocab; we do this instead of squeezing
            # to ensure that any external configuration requiring offset tracking,
            # e.g., LoRA, are applied correctly regardless of whether or not
            # we have multimodal tokens.
            in_vocab_ids = input_ids.masked_fill(is_multimodal, 0)
            return embed_input_ids(in_vocab_ids)
361

362
        return embed_input_ids(input_ids)
363

364
    def embed_input_ids(
365
        self,
366
        input_ids: Tensor,
367
        multimodal_embeddings: MultiModalEmbeddings | None = None,
368
        *,
369
        is_multimodal: Tensor | None = None,
370
    ) -> Tensor:
371
        """
372
373
374
375
376
        Apply token embeddings to `input_ids`.

        If `multimodal_embeddings` is passed, scatter them into
        `input_ids` according to the mask `is_multimodal`.

377
378
379
        NOTE: If this model has multimodal tokens that are of vocabulary
        (i.e., self._has_oov_mm_tokens=True), the input_ids will be copied
        and masked to 0 during the forward pass for the text embeddings.
380
        """
381
382
        from .utils import _merge_multimodal_embeddings

383
384
385
        # Get text embeddings first; multimodal embeddings will clobber
        # any invalid contents in the indices of multimodal embeddings
        # for the in vocabulary and out of vocabulary case.
386
        inputs_embeds = self._embed_text_input_ids(
387
            input_ids,
388
            self.get_language_model().embed_input_ids,
389
390
391
392
393
394
395
396
397
            is_multimodal=is_multimodal,
        )

        if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
            return inputs_embeds

        return _merge_multimodal_embeddings(
            inputs_embeds=inputs_embeds,
            multimodal_embeddings=multimodal_embeddings,
398
            is_multimodal=_require_is_multimodal(is_multimodal),
399
        )
400

401

402
403
404
405
406
407
@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.
    """
408

409
410
411
    supports_multimodal_pruning: ClassVar[Literal[True]] = True

    def recompute_mrope_positions(
412
413
414
415
416
        self,
        input_ids: list[int],
        multimodal_embeddings: MultiModalEmbeddings,
        mrope_positions: torch.LongTensor,
        num_computed_tokens: int,
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
    ) -> 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).
        """
        ...


441
@overload
442
def supports_multimodal(model: type[object]) -> TypeIs[type[SupportsMultiModal]]: ...
443
444
445


@overload
446
def supports_multimodal(model: object) -> TypeIs[SupportsMultiModal]: ...
447
448


449
def supports_multimodal(
450
451
    model: type[object] | object,
) -> TypeIs[type[SupportsMultiModal]] | TypeIs[SupportsMultiModal]:
452
    return getattr(model, "supports_multimodal", False)
453
454


455
def supports_multimodal_raw_input_only(model: type[object] | object) -> bool:
456
    return getattr(model, "supports_multimodal_raw_input_only", False)
457

458

Patrick von Platen's avatar
Patrick von Platen committed
459
460
461
462
def requires_raw_input_tokens(model: type[object] | object) -> bool:
    return getattr(model, "requires_raw_input_tokens", False)


463
def supports_multimodal_encoder_tp_data(model: type[object] | object) -> bool:
464
    return getattr(model, "supports_encoder_tp_data", False)
465
466


467
468
@overload
def supports_multimodal_pruning(
469
470
    model: type[object],
) -> TypeIs[type[SupportsMultiModalPruning]]: ...
471
472
473


@overload
474
def supports_multimodal_pruning(model: object) -> TypeIs[SupportsMultiModalPruning]: ...
475
476
477


def supports_multimodal_pruning(
478
479
    model: type[object] | object,
) -> TypeIs[type[SupportsMultiModalPruning]] | TypeIs[SupportsMultiModalPruning]:
480
481
482
    return getattr(model, "supports_multimodal_pruning", False)


483
484
485
486
487
488
489
490
491
492
493
494
495
496
@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
497
    def get_score_template(cls, query: str, document: str) -> str | None:
498
499
        """
        Generate a full prompt by populating the score template with query and document content.
500
        """  # noqa: E501
501
502
503
504
505
506
507
508
509
510
511
512
        ...

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


@overload
def supports_score_template(
513
514
    model: type[object],
) -> TypeIs[type[SupportsScoreTemplate]]: ...
515
516
517


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


def supports_score_template(
522
523
    model: type[object] | object,
) -> TypeIs[type[SupportsScoreTemplate]] | TypeIs[SupportsScoreTemplate]:
524
    return getattr(model, "supports_score_template", False)
525
526


527
528
529
530
@runtime_checkable
class SupportsLoRA(Protocol):
    """The interface required for all models that support LoRA."""

531
532
533
534
535
536
537
538
    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.
    """
539
    is_3d_moe_weight: ClassVar[bool] = False
540
    is_non_gated_moe: ClassVar[bool] = False
541
542
    # The `embedding_module` and `embedding_padding_modules`
    # are empty by default.
543
    embedding_modules: ClassVar[dict[str, str]] = {}
544
    packed_modules_mapping: dict[str, list[str]] = {}
545
546
    # Module prefixes to skip during LoRA loading (e.g., ["mtp."] for MTP layers)
    lora_skip_prefixes: ClassVar[list[str]] = []
547
548
549
550
551
552
553
554


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

555
556
    packed_modules_mapping: dict[str, list[str]]
    embedding_modules: dict[str, str]
557
558
559


@overload
560
def supports_lora(model: type[object]) -> TypeIs[type[SupportsLoRA]]: ...
561
562
563


@overload
564
def supports_lora(model: object) -> TypeIs[SupportsLoRA]: ...
565
566
567


def supports_lora(
568
569
    model: type[object] | object,
) -> TypeIs[type[SupportsLoRA]] | TypeIs[SupportsLoRA]:
570
571
572
573
574
575
576
    result = _supports_lora(model)

    if not result:
        lora_attrs = (
            "packed_modules_mapping",
            "embedding_modules",
        )
577
        missing_attrs = tuple(attr for attr in lora_attrs if not hasattr(model, attr))
578
579
580
581
582
583
584
585
586
587
588
589
590

        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, "
591
592
593
                    "but does not set `supports_lora=True`.",
                    model,
                )
594
595
596
597

    return result


598
def _supports_lora(model: type[object] | object) -> bool:
599
600
601
602
    if isinstance(model, type):
        return isinstance(model, _SupportsLoRAType)

    return isinstance(model, SupportsLoRA)
603
604


605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
@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,
623
    ) -> IntermediateTensors:
624
625
626
627
628
        """Called when PP rank > 0 for profiling purposes."""
        ...

    def forward(
        self,
629
630
        input_ids: Tensor | None,
        positions: Tensor,
631
        *,
632
633
        intermediate_tensors: IntermediateTensors | None,
    ) -> IntermediateTensors | None:
634
        """
635
636
        Accept [`IntermediateTensors`][vllm.sequence.IntermediateTensors] when
        PP rank > 0.
637

638
639
        Return [`IntermediateTensors`][vllm.sequence.IntermediateTensors] only
        for the last PP rank.
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
        """
        ...


# 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,
655
    ) -> IntermediateTensors: ...
656
657
658

    def forward(
        self,
659
660
        input_ids: Tensor | None,
        positions: Tensor,
661
        *,
662
663
        intermediate_tensors: IntermediateTensors | None,
    ) -> Tensor | IntermediateTensors: ...
664
665
666


@overload
667
def supports_pp(model: type[object]) -> TypeIs[type[SupportsPP]]: ...
668
669
670


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


def supports_pp(
675
676
    model: type[object] | object,
) -> bool | TypeIs[type[SupportsPP]] | TypeIs[SupportsPP]:
677
678
679
680
681
682
    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 "
683
684
685
            "`intermediate_tensors` in its `forward` method",
            model,
        )
686
687

    if not supports_attributes:
688
689
        pp_attrs = ("make_empty_intermediate_tensors",)
        missing_attrs = tuple(attr for attr in pp_attrs if not hasattr(model, attr))
690
691
692
693
694
695
696
697
698
699
700
701
702

        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, "
703
704
705
                    "but does not set `supports_pp=True`.",
                    model,
                )
706
707
708
709

    return supports_attributes and supports_inspect


710
def _supports_pp_attributes(model: type[object] | object) -> bool:
711
712
713
714
715
716
    if isinstance(model, type):
        return isinstance(model, _SupportsPPType)

    return isinstance(model, SupportsPP)


717
def _supports_pp_inspect(model: type[object] | object) -> bool:
718
719
720
721
    model_forward = getattr(model, "forward", None)
    if not callable(model_forward):
        return False

722
    return supports_kw(model_forward, "intermediate_tensors")
723
724


725
726
727
728
729
730
731
732
@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
733
        for max_num_seqs, etc. True for e.g. both Mamba and Jamba.
734
735
736
737
    """


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


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


def has_inner_state(
746
747
    model: type[object] | object,
) -> TypeIs[type[HasInnerState]] | TypeIs[HasInnerState]:
748
    return getattr(model, "has_inner_state", False)
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764


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


@overload
769
def is_attention_free(model: type[object]) -> TypeIs[type[IsAttentionFree]]: ...
770
771
772


def is_attention_free(
773
774
    model: type[object] | object,
) -> TypeIs[type[IsAttentionFree]] | TypeIs[IsAttentionFree]:
775
    return getattr(model, "is_attention_free", False)
776
777


778
779
780
@runtime_checkable
class IsHybrid(Protocol):
    """The interface required for all models like Jamba that have both
781
    attention and mamba blocks, indicates that
782
783
784
785
786
787
788
789
    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' """

790
791
792
    @classmethod
    def get_mamba_state_shape_from_config(
        cls,
793
        vllm_config: VllmConfig,
794
795
796
797
798
799
800
801
802
803
804
805
806
    ) -> 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
        """
        ...

807
808
809
810
811
812
813
814
815
816
817
818
819
    @classmethod
    def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, ...]:
        """Calculate copy-function callables for each Mamba state.

        Returns:
            A tuple of MambaStateCopyFunc callables that correspond, in order,
            to the Mamba states produced by the model. Each callable accepts
            (state, block_ids, cur_block_idx, num_accepted_tokens) and returns
            a MambaCopySpec describing the memory-copy parameters for prefix
            caching in align mode.
        """
        ...

820
821

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


@overload
826
def is_hybrid(model: type[object]) -> TypeIs[type[IsHybrid]]: ...
827
828
829


def is_hybrid(
830
831
    model: type[object] | object,
) -> TypeIs[type[IsHybrid]] | TypeIs[IsHybrid]:
832
    return getattr(model, "is_hybrid", False)
833
834


835
836
837
838
839
840
@runtime_checkable
class MixtureOfExperts(Protocol):
    """
    Check if the model is a mixture of experts (MoE) model.
    """

841
    expert_weights: MutableSequence[Sequence[Tensor]]
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
    """
    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."""

873
874
875
    moe_layers: Iterable[nn.Module]
    """List of MoE layers in this model."""

876
877
878
879
880
881
882
883
    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.
884

885
886
887
888
889
890
891
892
893
894
895
896
897
        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.
        """
898
899
900
901
902
903
904
905
906
        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,
            )
907

908
909
910
911
    def update_physical_experts_metadata(
        self,
        num_physical_experts: int,
        num_local_physical_experts: int,
912
    ) -> None: ...
913

914
915

def is_mixture_of_experts(model: object) -> TypeIs[MixtureOfExperts]:
916
917
918
    return (
        isinstance(model, MixtureOfExperts) and getattr(model, "num_moe_layers", 0) > 0
    )
919
920


921
922
923
924
925
926
@runtime_checkable
class HasNoOps(Protocol):
    has_noops: ClassVar[Literal[True]] = True


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


@overload
931
def has_noops(model: type[object]) -> TypeIs[type[HasNoOps]]: ...
932
933
934


def has_noops(
935
936
    model: type[object] | object,
) -> TypeIs[type[HasNoOps]] | TypeIs[HasNoOps]:
937
    return getattr(model, "has_noops", False)
938
939


940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
@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)


968
969
970
971
972
973
974
975
976
@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(
977
978
    model: type[object],
) -> TypeIs[type[SupportsCrossEncoding]]: ...
979
980
981


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


def _supports_cross_encoding(
986
987
    model: type[object] | object,
) -> TypeIs[type[SupportsCrossEncoding]] | TypeIs[SupportsCrossEncoding]:
988
    return getattr(model, "supports_cross_encoding", False)
989
990
991


def supports_cross_encoding(
992
993
    model: type[object] | object,
) -> TypeIs[type[SupportsCrossEncoding]] | TypeIs[SupportsCrossEncoding]:
994
    return is_pooling_model(model) and _supports_cross_encoding(model)
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
@runtime_checkable
class SupportsLateInteraction(Protocol):
    """The interface required for all models that support late interaction.

    Late interaction models (like ColBERT) encode queries and documents
    separately into per-token embeddings, then compute similarity via
    MaxSim (max over document tokens, sum over query tokens).
    """

    supports_late_interaction: ClassVar[Literal[True]] = True


@overload
def supports_late_interaction(
    model: type[object],
) -> TypeIs[type[SupportsLateInteraction]]: ...


@overload
def supports_late_interaction(model: object) -> TypeIs[SupportsLateInteraction]: ...


def _supports_late_interaction(
    model: type[object] | object,
) -> TypeIs[type[SupportsLateInteraction]] | TypeIs[SupportsLateInteraction]:
    return getattr(model, "supports_late_interaction", False)


def supports_late_interaction(
    model: type[object] | object,
) -> TypeIs[type[SupportsLateInteraction]] | TypeIs[SupportsLateInteraction]:
    return is_pooling_model(model) and _supports_late_interaction(model)


1031
1032
1033
class SupportsQuant:
    """The interface required for all models that support quantization."""

1034
1035
1036
    hf_to_vllm_mapper: ClassVar[WeightsMapper | None] = None
    packed_modules_mapping: ClassVar[dict[str, list[str]] | None] = None
    quant_config: QuantizationConfig | None = None
1037

1038
    def __new__(cls, *args, **kwargs) -> Self:
1039
        instance = super().__new__(cls)
1040
1041

        # find config passed in arguments
1042
1043
        quant_config = cls._find_quant_config(*args, **kwargs)
        if quant_config is not None:
1044
            # attach config to model for general use
1045
            instance.quant_config = quant_config
1046
1047

            # apply model mappings to config for proper config-model matching
1048
1049
1050
            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:
1051
                instance.quant_config.packed_modules_mapping.update(
1052
1053
                    instance.packed_modules_mapping
                )
1054

1055
1056
1057
        return instance

    @staticmethod
1058
    def _find_quant_config(*args, **kwargs) -> QuantizationConfig | None:
1059
        """Find quant config passed through model constructor args"""
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
        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


1073
1074
1075
1076
1077
1078
@runtime_checkable
class SupportsRealtime(Protocol):
    """The interface required for all models that support transcription."""

    supports_realtime: ClassVar[Literal[True]] = True

1079
1080
1081
1082
    realtime_max_tokens: ClassVar[int] = 1
    """Maximum tokens to generate per streaming audio segment.
    Override in subclasses based on the model's expected output length."""

1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
    @classmethod
    async def buffer_realtime_audio(
        cls,
        audio_stream: AsyncGenerator[np.ndarray, None],
        input_stream: asyncio.Queue[list[int]],
        model_config: ModelConfig,
    ) -> AsyncGenerator[PromptType, None]: ...


@overload
def supports_realtime(
    model: type[object],
) -> TypeIs[type[SupportsRealtime]]: ...


@overload
def supports_realtime(model: object) -> TypeIs[SupportsRealtime]: ...


def supports_realtime(
    model: type[object] | object,
) -> TypeIs[type[SupportsRealtime]] | TypeIs[SupportsRealtime]:
    return getattr(model, "supports_realtime", False)


1108
1109
1110
@runtime_checkable
class SupportsTranscription(Protocol):
    """The interface required for all models that support transcription."""
1111

1112
1113
    # Mapping from ISO639_1 language codes: language names
    supported_languages: ClassVar[Mapping[str, str]]
1114
1115
1116

    supports_transcription: ClassVar[Literal[True]] = True

1117
1118
1119
1120
1121
    supports_transcription_only: ClassVar[bool] = False
    """
    Transcription models can opt out of text generation by setting this to
    `True`.
    """
1122
1123
1124
1125
    supports_segment_timestamp: ClassVar[bool] = False
    """
    Enables the segment timestamp option for supported models by setting this to `True`.
    """
1126

1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
    supports_explicit_language_detection: ClassVar[bool] = False
    """
    Transcription models that require an explicit language detection step
    (e.g. Whisper needs a separate forward pass to predict the language
    token) should set this to ``True`` and implement
    :meth:`get_language_detection_prompt` and
    :meth:`parse_language_detection_output` and
    :meth:`get_language_token_ids`.
    """

1137
1138
1139
1140
1141
1142
1143
1144
1145
    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. "
1146
1147
                f"Valid choices are: {sorted(LANGUAGES.keys())}"
            )
1148

1149
    @classmethod
1150
1151
1152
1153
    def get_generation_prompt(
        cls,
        audio: np.ndarray,
        stt_config: SpeechToTextConfig,
1154
        model_config: ModelConfig,
1155
        language: str | None,
1156
1157
        task_type: Literal["transcribe", "translate"],
        request_prompt: str,
1158
        to_language: str | None,
1159
    ) -> PromptType:
1160
1161
1162
        """Get the prompt for the ASR model.
        The model has control over the construction, as long as it
        returns a valid PromptType."""
1163
1164
1165
        ...

    @classmethod
1166
1167
    def get_other_languages(cls) -> Mapping[str, str]:
        # other possible language codes from the whisper map
1168
        return {k: v for k, v in LANGUAGES.items() if k not in cls.supported_languages}
1169
1170

    @classmethod
1171
    def validate_language(cls, language: str | None) -> str | None:
1172
        """
1173
1174
1175
        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
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
        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 "
1192
1193
                f"{list(cls.supported_languages.keys())}."
            )
1194

1195
1196
    @classmethod
    def get_speech_to_text_config(
1197
        cls, model_config: ModelConfig, task_type: Literal["transcribe", "translate"]
1198
    ) -> SpeechToTextConfig:
1199
1200
1201
1202
        """Get the speech to text config for the ASR model."""
        ...

    @classmethod
1203
1204
1205
1206
    def get_num_audio_tokens(
        cls,
        audio_duration_s: float,
        stt_config: SpeechToTextConfig,
1207
        model_config: ModelConfig,
1208
    ) -> int | None:
1209
        """
1210
        Map from audio duration to number of audio tokens produced by the ASR
1211
1212
1213
1214
1215
        model, without running a forward pass.
        This is used for estimating the amount of processing for this audio.
        """
        return None

1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
    @classmethod
    def post_process_output(cls, text: str) -> str:
        """
        Post-process the raw model output text.

        Some ASR models output structured formats (e.g., language tags,
        special tokens) that need to be stripped before returning to the user.

        Args:
            text: Raw decoded text from the model.

        Returns:
            Cleaned transcription text.
        """
        return text

1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
    @classmethod
    def get_language_detection_prompt(
        cls,
        audio: np.ndarray,
        stt_config: SpeechToTextConfig,
    ) -> PromptType:
        """Return a prompt that triggers language detection.

        Only needs to be implemented when
        ``supports_explicit_language_detection`` is ``True``.
        """
        raise NotImplementedError

    @classmethod
    def parse_language_detection_output(
        cls,
        token_ids: list[int],
        tokenizer: object,
    ) -> str:
        """Parse the detected language from model output token IDs.

        Only needs to be implemented when
        ``supports_explicit_language_detection`` is ``True``.
        """
        raise NotImplementedError

    @classmethod
    def get_language_token_ids(
        cls,
        tokenizer: object,
    ) -> list[int] | None:
        """Return token IDs that represent valid language tokens.

        Used to constrain language detection to only produce valid language tokens.

        Only needs to be implemented when
        ``supports_explicit_language_detection`` is ``True``.
        """
        raise NotImplementedError

1272
1273
1274

@overload
def supports_transcription(
1275
1276
    model: type[object],
) -> TypeIs[type[SupportsTranscription]]: ...
1277
1278
1279


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


def supports_transcription(
1284
1285
    model: type[object] | object,
) -> TypeIs[type[SupportsTranscription]] | TypeIs[SupportsTranscription]:
1286
    return getattr(model, "supports_transcription", False)
1287
1288


1289
@runtime_checkable
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
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):
1351
    """The interface required for models that support
1352
    EAGLE-3 speculative decoding."""
1353
1354
1355

    supports_eagle3: ClassVar[Literal[True]] = True
    """
1356
    A flag that indicates this model supports EAGLE-3 
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
    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
1367
        hidden states for EAGLE-3.
1368

1369
1370
        Args:
            layers: Tuple of layer indices that should output auxiliary
1371
                hidden states.
1372
1373
1374
1375
1376
1377
        """
        ...

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

1380
1381
1382
1383
1384
1385
1386
        Returns:
            Tuple of layer indices for auxiliary hidden state outputs.
        """
        ...


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


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


def supports_eagle3(
1395
1396
    model: type[object] | object,
) -> TypeIs[type[SupportsEagle3]] | TypeIs[SupportsEagle3]:
1397
    return isinstance(model, SupportsEagle3)
1398
1399
1400
1401
1402
1403
1404
1405
1406


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

1408
1409
1410
1411
1412
1413
1414
1415
    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],
1416
        mm_features: list["MultiModalFeatureSpec"],
1417
1418
1419
    ) -> tuple[torch.Tensor, int]:
        """
        Get M-RoPE input positions and delta value for this specific model.
1420

1421
1422
        This method should be implemented by each model that supports M-RoPE
        to provide model-specific logic for computing input positions.
1423

1424
1425
        Args:
            input_tokens: List of input token IDs
1426
            mm_features: Information about each multi-modal data item
1427

1428
        Returns:
1429
1430
            Tuple of `(llm_positions, mrope_position_delta)`
            - llm_positions: Tensor of shape `[3, num_tokens]` with T/H/W positions
1431
1432
1433
1434
1435
1436
            - mrope_position_delta: Delta for position calculations
        """
        ...


@overload
1437
def supports_mrope(model: type[object]) -> TypeIs[type[SupportsMRoPE]]: ...
1438
1439
1440


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


def supports_mrope(
1445
1446
    model: type[object] | object,
) -> TypeIs[type[SupportsMRoPE]] | TypeIs[SupportsMRoPE]:
1447
    return isinstance(model, SupportsMRoPE)
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496


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