interfaces.py 50.2 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
from typing import (
    TYPE_CHECKING,
16
    Any,
17
18
19
    ClassVar,
    Literal,
    Protocol,
20
    TypeAlias,
21
22
23
    overload,
    runtime_checkable,
)
24

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

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

42
from .interfaces_base import VllmModel
43

44
if TYPE_CHECKING:
45
    from vllm.config import VllmConfig
46
    from vllm.model_executor.models.utils import WeightsMapper
47
    from vllm.multimodal.inputs import MultiModalFeatureSpec
48
    from vllm.multimodal.registry import _ProcessorFactories
49
    from vllm.sequence import IntermediateTensors
50
51
52
53
54
    from vllm.v1.worker.gpu.mm.encoder_cudagraph_defs import (
        EncoderCudaGraphCaptureInputs,
        EncoderCudaGraphConfig,
        EncoderCudaGraphReplayBuffers,
    )
55
56
57
else:
    VllmConfig = object
    WeightsMapper = object
58
    MultiModalFeatureSpec = object
59
    _ProcessorFactories = object
60
    IntermediateTensors = object
61

62
63
logger = init_logger(__name__)

64
MultiModalEmbeddings: TypeAlias = list[Tensor] | Tensor | tuple[Tensor, ...]
65
66
67
68
69
70
71
"""
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.
"""
72

73

74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
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


90
91
# Cache results of `SupportsMultiModal.get_language_model`
_language_model_by_module = dict[nn.Module, VllmModel]()
92
93


94
@runtime_checkable
95
class SupportsMultiModal(Protocol):
96
    """The interface required for all multi-modal models."""
97

98
    supports_multimodal: ClassVar[Literal[True]] = True
99
    """
100
    A flag that indicates this model supports multi-modal inputs.
101
102
103
104
105

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

107
108
109
110
111
112
    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.
    """

113
114
115
116
117
118
    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
119
120
121
122
123
124
    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.
    """

125
126
127
128
129
    _processor_factory: ClassVar[_ProcessorFactories]
    """
    Set internally by `MultiModalRegistry.register_processor`.
    """

130
131
132
133
134
135
136
137
138
139
    _language_model_names: list[str] = []
    """
    Set internally by `_mark_language_model`.
    """

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

140
141
142
143
144
145
146
    _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.
    """

147
    @classmethod
148
    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
149
150
151
152
153
        """
        Get the placeholder text for the `i`th `modality` item in the prompt.
        """
        ...

154
    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
155
        """
156
        Returns multimodal embeddings generated from multimodal kwargs
157
        to be merged with text embeddings.
158

159
        Note:
160
161
            The returned multimodal embeddings must be in the same order as
            the appearances of their corresponding multimodal data item in the
162
            input prompt.
163
        """
164
        ...
165

166
167
168
169
170
171
172
173
174
175
176
    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,
        )

177
    def get_language_model(self) -> VllmModel:
178
179
180
        """
        Returns the underlying language model used for text generation.

181
        This is typically the `torch.nn.Module` instance responsible for
182
183
184
185
186
        processing the merged multimodal embeddings and producing hidden states

        Returns:
            torch.nn.Module: The core language model component.
        """
187
188
189
190
        # Cached
        if self in _language_model_by_module:
            return _language_model_by_module[self]

191
        if self._language_model_names:
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
            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
208
209
210

        raise NotImplementedError(
            f"No language model found in {type(self).__name__}! "
211
            "You should initialize it via `_mark_language_model`."
212
213
214
        )

    @contextmanager
215
216
217
218
219
220
    def _mark_language_model(
        self,
        vllm_config: VllmConfig,
        *,
        targets: type[nn.Module] | tuple[type[nn.Module], ...] | None = None,
    ):
221
        """
222
223
224
225
226
        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.
227

228
229
230
231
        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
232

233
        mm_config = vllm_config.model_config.multimodal_config
234

235
        with collect_children(self, targets=targets) as children_names:  # noqa: SIM117
236
            with (
237
238
239
240
241
                no_init_weights(
                    self,
                    lambda mod: StageMissingLayer("language_model", mod),
                    targets=targets,
                )
242
243
244
245
246
247
248
249
                if mm_config.mm_encoder_only
                else nullcontext()
            ):
                yield

        self._language_model_names = children_names

    @contextmanager
250
251
252
253
254
255
256
    def _mark_tower_model(
        self,
        vllm_config: VllmConfig,
        modalities: set[str] | str,
        *,
        targets: type[nn.Module] | tuple[type[nn.Module], ...] | None = None,
    ):
257
        """
258
259
260
261
262
263
264
265
        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.
266
        """
267
268
        from .utils import StageMissingLayer, collect_children, no_init_weights

269
270
271
        if isinstance(modalities, str):
            modalities = {modalities}

272
273
274
275
        if modalities == {"image", "video"}:
            stage_name = "vision_tower"
        else:
            stage_name = "_".join([*modalities, "tower"])
276

277
        mm_config = vllm_config.model_config.multimodal_config
278

279
        with collect_children(self, targets=targets) as children_names:  # noqa: SIM117
280
            with (
281
282
283
284
285
                no_init_weights(
                    self,
                    lambda mod: StageMissingLayer(stage_name, mod),
                    targets=targets,
                )
286
287
288
289
290
291
                if all(mm_config.get_limit_per_prompt(m) == 0 for m in modalities)
                else nullcontext()
            ):
                yield

        self._tower_model_names = children_names
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
317
318
319
320
321
322
323
    @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

324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
    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.
        """
        ...

342
    @overload
343
    def embed_input_ids(self, input_ids: Tensor) -> Tensor: ...
344
345

    @overload
346
    def embed_input_ids(
347
348
349
350
351
        self,
        input_ids: Tensor,
        multimodal_embeddings: MultiModalEmbeddings,
        *,
        is_multimodal: torch.Tensor,
352
    ) -> Tensor: ...
353

354
    def _embed_text_input_ids(
355
356
        self,
        input_ids: Tensor,
357
        embed_input_ids: Callable[[Tensor], Tensor],
358
        *,
359
        is_multimodal: Tensor | None,
360
    ) -> Tensor:
361
362
363
364
365
366
367
        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)
368

369
        return embed_input_ids(input_ids)
370

371
    def embed_input_ids(
372
        self,
373
        input_ids: Tensor,
374
        multimodal_embeddings: MultiModalEmbeddings | None = None,
375
        *,
376
        is_multimodal: Tensor | None = None,
377
    ) -> Tensor:
378
        """
379
380
381
382
383
        Apply token embeddings to `input_ids`.

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

384
385
386
        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.
387
        """
388
389
        from .utils import _merge_multimodal_embeddings

390
391
392
        # 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.
393
        inputs_embeds = self._embed_text_input_ids(
394
            input_ids,
395
            self.get_language_model().embed_input_ids,
396
397
398
399
400
401
402
403
404
            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,
405
            is_multimodal=_require_is_multimodal(is_multimodal),
406
        )
407

408

409
410
411
412
413
414
@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.
    """
415

416
417
418
    supports_multimodal_pruning: ClassVar[Literal[True]] = True

    def recompute_mrope_positions(
419
420
421
422
423
        self,
        input_ids: list[int],
        multimodal_embeddings: MultiModalEmbeddings,
        mrope_positions: torch.LongTensor,
        num_computed_tokens: int,
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
    ) -> 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).
        """
        ...


448
@overload
449
def supports_multimodal(model: type[object]) -> TypeIs[type[SupportsMultiModal]]: ...
450
451
452


@overload
453
def supports_multimodal(model: object) -> TypeIs[SupportsMultiModal]: ...
454
455


456
def supports_multimodal(
457
458
    model: type[object] | object,
) -> TypeIs[type[SupportsMultiModal]] | TypeIs[SupportsMultiModal]:
459
    return getattr(model, "supports_multimodal", False)
460
461


462
def supports_multimodal_raw_input_only(model: type[object] | object) -> bool:
463
    return getattr(model, "supports_multimodal_raw_input_only", False)
464

465

Patrick von Platen's avatar
Patrick von Platen committed
466
467
468
469
def requires_raw_input_tokens(model: type[object] | object) -> bool:
    return getattr(model, "requires_raw_input_tokens", False)


470
def supports_multimodal_encoder_tp_data(model: type[object] | object) -> bool:
471
    return getattr(model, "supports_encoder_tp_data", False)
472
473


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


@overload
481
def supports_multimodal_pruning(model: object) -> TypeIs[SupportsMultiModalPruning]: ...
482
483
484


def supports_multimodal_pruning(
485
486
    model: type[object] | object,
) -> TypeIs[type[SupportsMultiModalPruning]] | TypeIs[SupportsMultiModalPruning]:
487
488
489
    return getattr(model, "supports_multimodal_pruning", False)


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

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


@overload
def supports_score_template(
520
521
    model: type[object],
) -> TypeIs[type[SupportsScoreTemplate]]: ...
522
523
524


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


def supports_score_template(
529
530
    model: type[object] | object,
) -> TypeIs[type[SupportsScoreTemplate]] | TypeIs[SupportsScoreTemplate]:
531
    return getattr(model, "supports_score_template", False)
532
533


534
535
536
537
@runtime_checkable
class SupportsLoRA(Protocol):
    """The interface required for all models that support LoRA."""

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


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

562
563
    packed_modules_mapping: dict[str, list[str]]
    embedding_modules: dict[str, str]
564
565
566


@overload
567
def supports_lora(model: type[object]) -> TypeIs[type[SupportsLoRA]]: ...
568
569
570


@overload
571
def supports_lora(model: object) -> TypeIs[SupportsLoRA]: ...
572
573
574


def supports_lora(
575
576
    model: type[object] | object,
) -> TypeIs[type[SupportsLoRA]] | TypeIs[SupportsLoRA]:
577
578
579
580
581
582
583
    result = _supports_lora(model)

    if not result:
        lora_attrs = (
            "packed_modules_mapping",
            "embedding_modules",
        )
584
        missing_attrs = tuple(attr for attr in lora_attrs if not hasattr(model, attr))
585
586
587
588
589
590
591
592
593
594
595
596
597

        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, "
598
599
600
                    "but does not set `supports_lora=True`.",
                    model,
                )
601
602
603
604

    return result


605
def _supports_lora(model: type[object] | object) -> bool:
606
607
608
609
    if isinstance(model, type):
        return isinstance(model, _SupportsLoRAType)

    return isinstance(model, SupportsLoRA)
610
611


612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
@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,
630
    ) -> IntermediateTensors:
631
632
633
634
635
        """Called when PP rank > 0 for profiling purposes."""
        ...

    def forward(
        self,
636
637
        input_ids: Tensor | None,
        positions: Tensor,
638
        *,
639
640
        intermediate_tensors: IntermediateTensors | None,
    ) -> IntermediateTensors | None:
641
        """
642
643
        Accept [`IntermediateTensors`][vllm.sequence.IntermediateTensors] when
        PP rank > 0.
644

645
646
        Return [`IntermediateTensors`][vllm.sequence.IntermediateTensors] only
        for the last PP rank.
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
        """
        ...


# 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,
662
    ) -> IntermediateTensors: ...
663
664
665

    def forward(
        self,
666
667
        input_ids: Tensor | None,
        positions: Tensor,
668
        *,
669
670
        intermediate_tensors: IntermediateTensors | None,
    ) -> Tensor | IntermediateTensors: ...
671
672
673


@overload
674
def supports_pp(model: type[object]) -> TypeIs[type[SupportsPP]]: ...
675
676
677


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


def supports_pp(
682
683
    model: type[object] | object,
) -> bool | TypeIs[type[SupportsPP]] | TypeIs[SupportsPP]:
684
685
686
687
688
689
    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 "
690
691
692
            "`intermediate_tensors` in its `forward` method",
            model,
        )
693
694

    if not supports_attributes:
695
696
        pp_attrs = ("make_empty_intermediate_tensors",)
        missing_attrs = tuple(attr for attr in pp_attrs if not hasattr(model, attr))
697
698
699
700
701
702
703
704
705
706
707
708
709

        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, "
710
711
712
                    "but does not set `supports_pp=True`.",
                    model,
                )
713
714
715
716

    return supports_attributes and supports_inspect


717
def _supports_pp_attributes(model: type[object] | object) -> bool:
718
719
720
721
722
723
    if isinstance(model, type):
        return isinstance(model, _SupportsPPType)

    return isinstance(model, SupportsPP)


724
def _supports_pp_inspect(model: type[object] | object) -> bool:
725
726
727
728
    model_forward = getattr(model, "forward", None)
    if not callable(model_forward):
        return False

729
    return supports_kw(model_forward, "intermediate_tensors")
730
731


732
733
734
735
736
737
738
739
@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
740
        for max_num_seqs, etc. True for e.g. both Mamba and Jamba.
741
742
743
744
    """


@overload
745
def has_inner_state(model: object) -> TypeIs[HasInnerState]: ...
746
747
748


@overload
749
def has_inner_state(model: type[object]) -> TypeIs[type[HasInnerState]]: ...
750
751
752


def has_inner_state(
753
754
    model: type[object] | object,
) -> TypeIs[type[HasInnerState]] | TypeIs[HasInnerState]:
755
    return getattr(model, "has_inner_state", False)
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771


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


@overload
776
def is_attention_free(model: type[object]) -> TypeIs[type[IsAttentionFree]]: ...
777
778
779


def is_attention_free(
780
781
    model: type[object] | object,
) -> TypeIs[type[IsAttentionFree]] | TypeIs[IsAttentionFree]:
782
    return getattr(model, "is_attention_free", False)
783
784


785
786
787
@runtime_checkable
class IsHybrid(Protocol):
    """The interface required for all models like Jamba that have both
788
    attention and mamba blocks, indicates that
789
790
791
792
793
794
795
796
    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' """

797
798
799
    @classmethod
    def get_mamba_state_shape_from_config(
        cls,
800
        vllm_config: VllmConfig,
801
802
803
804
805
806
807
808
809
810
811
812
813
    ) -> 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
        """
        ...

814
815
816
817
818
819
820
821
822
823
824
825
826
    @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.
        """
        ...

827
828

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


@overload
833
def is_hybrid(model: type[object]) -> TypeIs[type[IsHybrid]]: ...
834
835
836


def is_hybrid(
837
838
    model: type[object] | object,
) -> TypeIs[type[IsHybrid]] | TypeIs[IsHybrid]:
839
    return getattr(model, "is_hybrid", False)
840
841


842
843
844
845
846
847
@runtime_checkable
class MixtureOfExperts(Protocol):
    """
    Check if the model is a mixture of experts (MoE) model.
    """

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

880
881
882
    moe_layers: Iterable[nn.Module]
    """List of MoE layers in this model."""

883
884
885
886
887
888
889
890
    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.
891

892
893
894
895
896
897
898
899
900
901
902
903
904
        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.
        """
905
906
907
908
909
910
911
912
913
        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,
            )
914

915
916
917
918
    def update_physical_experts_metadata(
        self,
        num_physical_experts: int,
        num_local_physical_experts: int,
919
    ) -> None: ...
920

921
922

def is_mixture_of_experts(model: object) -> TypeIs[MixtureOfExperts]:
923
924
925
    return (
        isinstance(model, MixtureOfExperts) and getattr(model, "num_moe_layers", 0) > 0
    )
926
927


928
929
930
931
932
933
@runtime_checkable
class HasNoOps(Protocol):
    has_noops: ClassVar[Literal[True]] = True


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


@overload
938
def has_noops(model: type[object]) -> TypeIs[type[HasNoOps]]: ...
939
940
941


def has_noops(
942
943
    model: type[object] | object,
) -> TypeIs[type[HasNoOps]] | TypeIs[HasNoOps]:
944
    return getattr(model, "has_noops", False)
945
946


947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
@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)


975
976
977
978
@runtime_checkable
class SupportsCrossEncoding(Protocol):
    """The interface required for all models that support cross encoding."""

979
    score_type: ClassVar[ScoreType] = "cross-encoder"
980
981


982
983
984
985
986
987
988
989
990
@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).
    """

991
    score_type: ClassVar[ScoreType] = "late-interaction"
992
993


994
995
996
class SupportsQuant:
    """The interface required for all models that support quantization."""

997
998
999
    hf_to_vllm_mapper: ClassVar[WeightsMapper | None] = None
    packed_modules_mapping: ClassVar[dict[str, list[str]] | None] = None
    quant_config: QuantizationConfig | None = None
1000

1001
    def __new__(cls, *args, **kwargs) -> Self:
1002
        instance = super().__new__(cls)
1003

1004
1005
1006
1007
        # find config passed in arguments and attach it to model for general use
        instance.quant_config = cls._find_quant_config(*args, **kwargs)

        cls._maybe_apply_model_mapping(instance)
1008

1009
1010
1011
        return instance

    @staticmethod
1012
    def _find_quant_config(*args, **kwargs) -> QuantizationConfig | None:
1013
        """Find quant config passed through model constructor args"""
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
        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

1026
1027
1028
1029
1030
1031
1032
1033
1034
    def _maybe_apply_model_mapping(self):
        """Apply model mappings to config for proper config-model matching"""
        if self.quant_config is None:
            return
        if (hf_to_vllm_mapper := self.hf_to_vllm_mapper) is not None:
            self.quant_config.apply_vllm_mapper(hf_to_vllm_mapper)
        if self.packed_modules_mapping is not None:
            self.quant_config.packed_modules_mapping.update(self.packed_modules_mapping)

1035

1036
1037
1038
1039
1040
1041
@runtime_checkable
class SupportsRealtime(Protocol):
    """The interface required for all models that support transcription."""

    supports_realtime: ClassVar[Literal[True]] = True

1042
1043
1044
1045
    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."""

1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
    @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)


1071
1072
1073
@runtime_checkable
class SupportsTranscription(Protocol):
    """The interface required for all models that support transcription."""
1074

1075
1076
    # Mapping from ISO639_1 language codes: language names
    supported_languages: ClassVar[Mapping[str, str]]
1077
1078
1079

    supports_transcription: ClassVar[Literal[True]] = True

1080
1081
1082
1083
1084
    supports_transcription_only: ClassVar[bool] = False
    """
    Transcription models can opt out of text generation by setting this to
    `True`.
    """
1085
1086
1087
1088
    supports_segment_timestamp: ClassVar[bool] = False
    """
    Enables the segment timestamp option for supported models by setting this to `True`.
    """
1089

1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
    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`.
    """

1100
1101
1102
1103
1104
1105
1106
1107
1108
    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. "
1109
1110
                f"Valid choices are: {sorted(LANGUAGES.keys())}"
            )
1111

1112
    @classmethod
1113
1114
1115
1116
    def get_generation_prompt(
        cls,
        audio: np.ndarray,
        stt_config: SpeechToTextConfig,
1117
        model_config: ModelConfig,
1118
        language: str | None,
1119
1120
        task_type: Literal["transcribe", "translate"],
        request_prompt: str,
1121
        to_language: str | None,
1122
    ) -> PromptType:
1123
1124
1125
        """Get the prompt for the ASR model.
        The model has control over the construction, as long as it
        returns a valid PromptType."""
1126
1127
1128
        ...

    @classmethod
1129
1130
    def get_other_languages(cls) -> Mapping[str, str]:
        # other possible language codes from the whisper map
1131
        return {k: v for k, v in LANGUAGES.items() if k not in cls.supported_languages}
1132
1133

    @classmethod
1134
    def validate_language(cls, language: str | None) -> str | None:
1135
        """
1136
1137
1138
        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
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
        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 "
1155
1156
                f"{list(cls.supported_languages.keys())}."
            )
1157

1158
1159
    @classmethod
    def get_speech_to_text_config(
1160
        cls, model_config: ModelConfig, task_type: Literal["transcribe", "translate"]
1161
    ) -> SpeechToTextConfig:
1162
1163
1164
1165
        """Get the speech to text config for the ASR model."""
        ...

    @classmethod
1166
1167
1168
1169
    def get_num_audio_tokens(
        cls,
        audio_duration_s: float,
        stt_config: SpeechToTextConfig,
1170
        model_config: ModelConfig,
1171
    ) -> int | None:
1172
        """
1173
        Map from audio duration to number of audio tokens produced by the ASR
1174
1175
1176
1177
1178
        model, without running a forward pass.
        This is used for estimating the amount of processing for this audio.
        """
        return None

1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
    @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

1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
    @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

1235
1236
1237

@overload
def supports_transcription(
1238
1239
    model: type[object],
) -> TypeIs[type[SupportsTranscription]]: ...
1240
1241
1242


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


def supports_transcription(
1247
1248
    model: type[object] | object,
) -> TypeIs[type[SupportsTranscription]] | TypeIs[SupportsTranscription]:
1249
    return getattr(model, "supports_transcription", False)
1250
1251


1252
@runtime_checkable
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
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)


1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
class EagleModelMixin:
    aux_hidden_state_layers: tuple[int, ...] = ()

    def _set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
        self.aux_hidden_state_layers = layers

    def _maybe_add_hidden_state(
        self,
        aux_hidden_states: list[torch.Tensor],
        layer_idx: int,
        hidden_states: torch.Tensor,
        residual: torch.Tensor,
    ) -> list[torch.Tensor]:
        if layer_idx in self.aux_hidden_state_layers:
            value = hidden_states + residual if residual is not None else hidden_states
            aux_hidden_states.append(value)
        return aux_hidden_states


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
@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):
1333
    """The interface required for models that support
1334
    EAGLE-3 speculative decoding."""
1335
1336
1337

    supports_eagle3: ClassVar[Literal[True]] = True
    """
1338
    A flag that indicates this model supports EAGLE-3 
1339
1340
1341
1342
1343
1344
1345
1346
1347
    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:
        """
1348
        Set which layers should output auxiliary hidden states for EAGLE-3.
1349

1350
1351
        Args:
            layers: Tuple of layer indices that should output auxiliary
1352
                hidden states.
1353
        """
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
        parent_ref = self
        if hasattr(self, "get_language_model"):
            parent_ref = self.get_language_model()
        elif hasattr(self, "language_model"):
            parent_ref = self.language_model
        assert hasattr(parent_ref, "model"), (
            "Model instance must have 'model' attribute to set number of layers"
        )
        assert isinstance(parent_ref.model, EagleModelMixin), (
            "Model instance must inherit from EagleModelMixin to set auxiliary layers"
        )
        parent_ref.model._set_aux_hidden_state_layers(layers)
1366

1367
    def get_eagle3_default_aux_hidden_state_layers(self) -> tuple[int, ...]:
1368
        """
1369
1370
1371
1372
        Get the default layer indices that should output auxiliary hidden states
        for EAGLE-3 for this model. Models can override this method to provide
        different default layers based on their architecture, but it is encouraged
        to instead include the layer specification in the model's config if possible.
1373

1374
1375
1376
        Returns:
            Tuple of layer indices for auxiliary hidden state outputs.
        """
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
        parent_ref = self
        if hasattr(self, "get_language_model"):
            parent_ref = self.get_language_model()
        elif hasattr(self, "language_model"):
            parent_ref = self.language_model
        assert hasattr(parent_ref, "model"), (
            "Model instance must have 'model' attribute to get number of layers"
        )
        assert hasattr(parent_ref.model, "layers"), (
            "Model instance must have 'layers' attribute to get number of layers"
        )
        num_layers = len(parent_ref.model.layers)
        return (2, num_layers // 2, num_layers - 3)
1390
1391
1392


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


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


def supports_eagle3(
1401
1402
    model: type[object] | object,
) -> TypeIs[type[SupportsEagle3]] | TypeIs[SupportsEagle3]:
1403
    return isinstance(model, SupportsEagle3)
1404
1405
1406
1407
1408
1409
1410
1411
1412


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

1414
1415
1416
1417
1418
1419
1420
1421
    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],
1422
        mm_features: list["MultiModalFeatureSpec"],
1423
1424
1425
    ) -> tuple[torch.Tensor, int]:
        """
        Get M-RoPE input positions and delta value for this specific model.
1426

1427
1428
        This method should be implemented by each model that supports M-RoPE
        to provide model-specific logic for computing input positions.
1429

1430
1431
        Args:
            input_tokens: List of input token IDs
1432
            mm_features: Information about each multi-modal data item
1433

1434
        Returns:
1435
1436
            Tuple of `(llm_positions, mrope_position_delta)`
            - llm_positions: Tensor of shape `[3, num_tokens]` with T/H/W positions
1437
1438
1439
1440
1441
1442
            - mrope_position_delta: Delta for position calculations
        """
        ...


@overload
1443
def supports_mrope(model: type[object]) -> TypeIs[type[SupportsMRoPE]]: ...
1444
1445
1446


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


def supports_mrope(
1451
1452
    model: type[object] | object,
) -> TypeIs[type[SupportsMRoPE]] | TypeIs[SupportsMRoPE]:
1453
    return isinstance(model, SupportsMRoPE)
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
1497
1498
1499
1500
1501
1502


@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)
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637


@runtime_checkable
class SupportsEncoderCudaGraph(Protocol):
    """Interface for models whose vision encoder supports CUDA graph
    capture/replay.

    Models implement these methods to provide the
    :class:`EncoderCudaGraphManager` with all model-specific logic
    (input handling, metadata computation, forward pass) without the
    manager needing to know model internals.
    """

    supports_encoder_cudagraph: ClassVar[Literal[True]] = True

    def get_encoder_cudagraph_config(self) -> "EncoderCudaGraphConfig": ...

    def get_encoder_cudagraph_budget_range(
        self,
        vllm_config: "VllmConfig",
    ) -> tuple[int, int]:
        """Return (min_token_budget, max_token_budget) for auto-inference.

        - min_token_budget: estimated smallest possible encoder input
          (e.g. 64 for a 224x224 image)
        - max_token_budget: estimated largest budget worth capturing
          (e.g. max_num_batched_tokens)

        Used when ``encoder_cudagraph_token_budgets`` and/or
        ``encoder_cudagraph_max_images_per_batch`` are not explicitly
        specified by the user.
        """
        ...

    def get_encoder_cudagraph_num_items(
        self,
        mm_kwargs: dict[str, Any],
    ) -> int:
        """Return the number of items (e.g. images) in the batch."""
        ...

    def get_encoder_cudagraph_per_item_output_tokens(
        self,
        mm_kwargs: dict[str, Any],
    ) -> list[int]:
        """Return output token count for each item.

        Used for greedy packing and DP load balancing.
        """
        ...

    def get_encoder_cudagraph_per_item_input_sizes(
        self,
        mm_kwargs: dict[str, Any],
    ) -> list[int]:
        """Return input size (e.g. patch count) for each item.

        Used for input tensor slicing offsets.
        """
        ...

    def select_encoder_cudagraph_items(
        self,
        mm_kwargs: dict[str, Any],
        indices: list[int],
    ) -> dict[str, Any]:
        """Select a subset of items and return mm_kwargs for the sub-batch.

        Called by the manager during greedy packing and DP sharding to
        extract inputs for a specific set of items (e.g. images at
        indices [0, 3, 5]).  The implementation is model-specific
        because input formats differ:

        - Qwen-family: slice concatenated pixel_values by cumulative
          patch offsets, subset grid_thw by indices.
        - Batched models (CLIP): index pixel_values along dim 0.
        """
        ...

    def prepare_encoder_cudagraph_capture_inputs(
        self,
        token_budget: int,
        max_batch_size: int,
        device: torch.device,
        dtype: torch.dtype,
    ) -> "EncoderCudaGraphCaptureInputs":
        """Create dummy inputs and buffers for CUDA graph capture."""
        ...

    def prepare_encoder_cudagraph_replay_buffers(
        self,
        mm_kwargs: dict[str, Any],
        max_batch_size: int,
    ) -> "EncoderCudaGraphReplayBuffers":
        """Compute buffer values from actual batch inputs for replay."""
        ...

    def encoder_cudagraph_forward(
        self,
        mm_kwargs: dict[str, Any],
        buffers: dict[str, torch.Tensor],
    ) -> torch.Tensor:
        """Run the encoder forward pass with precomputed buffers.

        Used during both CUDA graph capture and replay.
        """
        ...

    def encoder_eager_forward(
        self,
        mm_kwargs: dict[str, Any],
    ) -> torch.Tensor:
        """Run the encoder forward pass without precomputed buffers.

        Used as eager fallback when inputs exceed all budgets.
        """
        ...


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


@overload
def supports_encoder_cudagraph(
    model: object,
) -> TypeIs[SupportsEncoderCudaGraph]: ...


def supports_encoder_cudagraph(
    model: type[object] | object,
) -> TypeIs[type[SupportsEncoderCudaGraph]] | TypeIs[SupportsEncoderCudaGraph]:
    return isinstance(model, SupportsEncoderCudaGraph)