vllm.model_executor.models.interfaces ¶

@dataclass(frozen=True, slots=True)
class DiarizedTranscriptionSegment:
    """A timestamped, speaker-attributed segment produced by an ASR model."""

    start: float
    end: float
    speaker: str
    text: str

@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
        for max_num_seqs, etc. True for e.g. both Mamba and Jamba.
    """

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

@runtime_checkable
class IsHybrid(Protocol):
    """The interface required for all models like Jamba that have both
    attention and mamba blocks, indicates that
    hf_config has 'layers_block_type'"""

    is_hybrid: ClassVar[Literal[True]] = True
    """
        A flag that indicates this model has both mamba and attention blocks
        , also indicates that the model's hf_config has 
        'layers_block_type' """

    @classmethod
    def get_mamba_state_shape_from_config(
        cls,
        vllm_config: "VllmConfig",
    ) -> MambaStateShapes:
        """Calculate shapes for Mamba's convolutional and state caches.

        Args:
            vllm_config: vLLM config

        Returns:
            Shapes for each state cache used by the model.
        """
        ...

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

    @classmethod
    def get_mamba_state_copy_funcs(
        cls,
        mamba_types: set["MambaAttentionBackendEnum"],
    ) -> "MambaStateCopyFuncsByType":
        copy_funcs = cls.get_mamba_state_copy_func()
        return {mamba_type: copy_funcs for mamba_type in mamba_types}

class LocalArgmaxMixin:
    """Mixin for draft model heads in speculative decoding.

    Provides a D2T-aware ``get_top_tokens`` that preserves the
    local-argmax communication reduction even when the draft vocabulary
    is smaller than the target vocabulary.

    When ``draft_id_to_target_id`` is present (shape ``(draft_vocab_size,)``,
    containing per-token offset to target vocab id), the draft argmax index
    ``k`` is mapped to the target vocab id via::

        target_id = k + draft_id_to_target_id[k]

    This is mathematically equivalent to computing the full-vocab scatter
    logits and taking the global argmax, but requires only
    O(batch * 2 * tp_size) communication instead of O(batch * vocab_size).

    Requires the subclass to expose:
        ``self.logits_processor``: LogitsProcessor
        ``self.lm_head``: ParallelLMHead
        ``self.draft_id_to_target_id`` (optional): nn.Parameter
    """

    logits_processor: "LogitsProcessor"
    lm_head: "ParallelLMHead"

    def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """Vocab-parallel argmax with optional D2T remapping."""
        top = self.logits_processor.get_top_tokens(
            self.lm_head,
            hidden_states,
        )
        d2t = getattr(self, "draft_id_to_target_id", None)
        if d2t is not None:
            top = top + d2t[top]
        return top

@runtime_checkable
class MixtureOfExperts(Protocol):
    """
    Check if the model is a mixture of experts (MoE) model.
    """

    expert_weights: MutableSequence[Sequence[Tensor]]
    """
    Expert weights saved in this rank.

    The first dimension is the layer, and the second dimension is different
    parameters in the layer, e.g. up/down projection weights.
    """

    num_moe_layers: int
    """Number of MoE layers in this model."""

    num_expert_groups: int
    """Number of expert groups in this model."""

    num_logical_experts: int
    """Number of logical experts in this model."""

    num_physical_experts: int
    """Number of physical experts in this model."""

    num_local_physical_experts: int
    """Number of local physical experts in this model."""

    num_routed_experts: int
    """Number of routed experts in this model."""

    num_shared_experts: int
    """Number of shared experts in this model."""

    num_redundant_experts: int
    """Number of redundant experts in this model."""

    moe_layers: Sequence["MoERunner"]
    """List of MoE layers in this model."""

    def set_eplb_state(
        self,
        expert_load_view: Tensor,
        logical_to_physical_map: Tensor,
        logical_replica_count: Tensor,
    ) -> None:
        """
        Register the EPLB state in the MoE model.

        Since these are views of the actual EPLB state, any changes made by
        the EPLB algorithm are automatically reflected in the model's behavior
        without requiring additional method calls to set new states.

        You should also collect model's `expert_weights` here instead of in
        the weight loader, since after initial weight loading, further
        processing like quantization may be applied to the weights.

        Args:
            expert_load_view: A view of the expert load metrics tensor.
            logical_to_physical_map: Mapping from logical to physical experts.
            logical_replica_count: Count of replicas for each logical expert.
        """
        self.expert_weights = []
        for layer_idx, layer in enumerate(self.moe_layers):
            # Register the expert weights.
            self.expert_weights.append(list(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,
            )

    def update_physical_experts_metadata(
        self,
        num_physical_experts: int,
        num_local_physical_experts: int,
    ) -> None: ...

class StreamingTranscriptionPostProcessor:
    """Stateful streaming post-processor for transcription deltas."""

    def process_delta(self, text_delta: str, finished: bool) -> str:
        return text_delta

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

    score_type: ClassVar["ScoreType"] = "cross-encoder"

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

@runtime_checkable
class SupportsEagle3(SupportsEagleBase, Protocol):
    """The interface required for models that support
    EAGLE-3 speculative decoding."""

    supports_eagle3: ClassVar[Literal[True]] = True
    """
    A flag that indicates this model supports EAGLE-3 
    speculative decoding.

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

    def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
        """
        Set which layers should output auxiliary hidden states for EAGLE-3.

        Args:
            layers: Tuple of layer indices that should output auxiliary
                hidden states.
        """
        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
        # A multimodal model that builds its decoder inside
        # `_mark_language_model` has get_language_model() return the inner
        # decoder itself, which IS the EagleModelMixin and has no further
        # `.model`. Unwrap only when there is something to unwrap.
        holder = getattr(parent_ref, "model", parent_ref)
        assert isinstance(holder, EagleModelMixin), (
            "Model instance must inherit from EagleModelMixin to set auxiliary layers"
        )
        holder._set_aux_hidden_state_layers(layers)

    def get_eagle3_default_aux_hidden_state_layers(self) -> tuple[int, ...]:
        """
        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.

        Returns:
            Tuple of layer indices for auxiliary hidden state outputs.
        """
        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
        # Same unwrap-only-if-needed rule as set_aux_hidden_state_layers.
        holder = getattr(parent_ref, "model", parent_ref)
        assert hasattr(holder, "layers"), (
            "Model instance must have 'layers' attribute to get number of layers"
        )
        num_layers = len(holder.layers)
        return (2, num_layers // 2, num_layers - 3)

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

@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_input_modality(
        self,
        mm_kwargs: dict[str, Any],
    ) -> str:
        """Return the modality of the inputs (default: image-only)."""
        return "image"

    def get_max_frames_per_video(
        self,
    ) -> int:
        """Return model-specific max frames per video."""
        ...

    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_vision_items_per_batch`` are not explicitly
        specified by the user.
        """
        ...

    def get_encoder_cudagraph_item_specs(
        self,
        mm_kwargs: dict[str, Any],
    ) -> list["EncoderItemSpec"]:
        """Return specs describing each item in the batch.

        Replaces the former separate methods for num_items,
        per_item_output_tokens, and per_item_input_sizes.
        The manager derives all three from this single return value.
        """
        ...

    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.

        Models that configure ``EncoderCudaGraphConfig.capture_axes`` must
        additionally store the resolved per-axis keys (one key per axis, in
        order) under ``ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG`` in the returned
        dict; the manager pops it before the kwargs are used elsewhere.
        """
        ...

    def postprocess_encoder_output(
        self,
        outputs: dict[str, torch.Tensor],
        indices: list[int],
        per_item_out_tokens: list[int],
        dest: dict[int, torch.Tensor] | list[torch.Tensor | None],
        clone: bool = False,
        batch_mm_kwargs: dict[str, Any] | None = None,
    ) -> None:
        """
        Post-process encoder output, directly call scatter_output_slices by default.

        By default, delegates directly to scatter_output_slices.
        Override this for models that require additional processing on the raw
        encoder output prior to scattering, e.g. Step3-VL, which merges features
        according to dynamic patch counts before scattering.
        """
        from vllm.model_executor.models.utils import scatter_output_slices

        scatter_output_slices(
            outputs["default"], indices, per_item_out_tokens, dest, clone
        )

    def prepare_encoder_cudagraph_capture_inputs(
        self,
        token_budget: int,
        max_batch_size: int,
        max_frames_per_batch: int,
        device: torch.device,
        dtype: torch.dtype,
        path: str = "default",
        axis_keys: tuple[Hashable, ...] | None = None,
    ) -> "EncoderCudaGraphCaptureInputs":
        """Create dummy inputs and buffers for CUDA graph capture.

        Args:
            axis_keys: The resolved capture-axis keys (one per axis of
                ``EncoderCudaGraphConfig.capture_axes``) this capture is for.
                None or empty when no capture axes are configured; models
                without capture axes ignore it.
        """
        ...

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

    def encoder_cudagraph_forward(
        self,
        inputs: dict[str, torch.Tensor],
        path: str = "default",
    ) -> 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],
        path: str = "default",
    ) -> torch.Tensor:
        """Run the encoder forward pass without precomputed buffers.

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

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

    score_type: ClassVar["ScoreType"] = "late-interaction"

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

    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.
    """
    is_3d_moe_weight: ClassVar[bool] = False
    is_non_gated_moe: ClassVar[bool] = False
    # The `embedding_module` and `embedding_padding_modules`
    # are empty by default.
    embedding_modules: ClassVar[dict[str, str]] = {}
    packed_modules_mapping: dict[str, list[str]] = {}
    # Module prefixes to skip during LoRA loading (e.g., ["mtp."] for MTP layers)
    lora_skip_prefixes: ClassVar[list[str]] = []
    lora_manager: "LoRAModelManager | None"

@runtime_checkable
class SupportsMRoPE(Protocol):
    """The interface required for all models that support M-RoPE."""

    supports_mrope: ClassVar[Literal[True]] = True
    """
    A flag that indicates this model supports M-RoPE.

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

    def get_mrope_input_positions(
        self,
        input_tokens: list[int],
        mm_features: list["MultiModalFeatureSpec"],
    ) -> tuple[torch.Tensor, int]:
        """
        Get M-RoPE input positions and delta value for this specific model.

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

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

        Returns:
            llm_positions: Tensor of shape `[num_dims, num_tokens]`, one row
                per M-RoPE position channel (e.g. T/H/W).
            mrope_position_delta: Delta for position calculations.
        """
        ...

@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

    @classmethod
    def get_mamba_state_copy_func(cls) -> tuple["MambaStateCopyFunc", ...]:
        """Return copy functions for the model's Mamba states."""
        ...

    @classmethod
    def get_mamba_state_copy_funcs(
        cls,
        mamba_types: set["MambaAttentionBackendEnum"],
    ) -> "MambaStateCopyFuncsByType":
        """Map legacy copy functions to each requested Mamba backend."""
        copy_funcs = cls.get_mamba_state_copy_func()
        return {mamba_type: copy_funcs for mamba_type in mamba_types}