| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import asyncio |
| 5 | import weakref |
| 6 | from bisect import bisect_right |
| 7 | from collections.abc import ( |
| 8 | AsyncGenerator, |
| 9 | Callable, |
| 10 | Hashable, |
| 11 | Mapping, |
| 12 | MutableSequence, |
| 13 | Sequence, |
| 14 | ) |
| 15 | from contextlib import ExitStack, contextmanager, nullcontext |
| 16 | from dataclasses import dataclass |
| 17 | from operator import attrgetter |
| 18 | from typing import ( |
| 19 | TYPE_CHECKING, |
| 20 | Any, |
| 21 | ClassVar, |
| 22 | Literal, |
| 23 | Protocol, |
| 24 | TypeAlias, |
| 25 | overload, |
| 26 | runtime_checkable, |
| 27 | ) |
| 28 | |
| 29 | import numpy as np |
| 30 | import torch |
| 31 | import torch.nn as nn |
| 32 | from torch import Tensor |
| 33 | from transformers.models.whisper.tokenization_whisper import LANGUAGES |
| 34 | from typing_extensions import Self, TypeIs, deprecated |
| 35 | |
| 36 | from vllm.logger import init_logger |
| 37 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 38 | from vllm.model_executor.models.interfaces_base import VllmModel |
| 39 | from vllm.utils.collection_utils import common_prefix |
| 40 | from vllm.utils.func_utils import supports_kw |
| 41 | |
| 42 | if TYPE_CHECKING: |
| 43 | from vllm.config import ( |
| 44 | ModelConfig, |
| 45 | SpeechToTextConfig, |
| 46 | SpeechToTextParams, |
| 47 | VllmConfig, |
| 48 | ) |
| 49 | from vllm.config.multimodal import VideoPruningMethod |
| 50 | from vllm.inputs import PromptType, TokensPrompt |
| 51 | from vllm.lora.model_manager import LoRAModelManager |
| 52 | from vllm.model_executor.layers.fused_moe import MoERunner |
| 53 | from vllm.model_executor.layers.logits_processor import LogitsProcessor |
| 54 | from vllm.model_executor.layers.mamba.mamba_utils import ( |
| 55 | MambaStateCopyFunc, |
| 56 | MambaStateCopyFuncsByType, |
| 57 | ) |
| 58 | from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead |
| 59 | from vllm.model_executor.models.utils import WeightsMapper |
| 60 | from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalKwargsItem |
| 61 | from vllm.multimodal.registry import _ProcessorFactories |
| 62 | from vllm.sequence import IntermediateTensors |
| 63 | from vllm.tasks import ScoreType |
| 64 | from vllm.tokenizers import TokenizerLike |
| 65 | from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum |
| 66 | from vllm.v1.worker.encoder_cudagraph_defs import ( |
| 67 | EncoderCudaGraphCaptureInputs, |
| 68 | EncoderCudaGraphConfig, |
| 69 | EncoderCudaGraphReplayBuffers, |
| 70 | EncoderItemSpec, |
| 71 | ) |
| 72 | |
| 73 | logger = init_logger(__name__) |
| 74 | |
| 75 | MultiModalEmbeddings: TypeAlias = list[Tensor] | Tensor | tuple[Tensor, ...] |
| 76 | """ |
| 77 | The output embeddings must be one of the following formats: |
| 78 | |
| 79 | - A list or tuple of 2D tensors, where each tensor corresponds to |
| 80 | each input multimodal data item (e.g, image). |
| 81 | - A single 3D tensor, with the batch dimension grouping the 2D tensors. |
| 82 | """ |
| 83 | |
| 84 | MambaStateShapes: TypeAlias = ( |
| 85 | tuple[tuple[int, int]] |
| 86 | | tuple[tuple[int, int, int]] |
| 87 | | tuple[tuple[int, int], tuple[int, int]] |
| 88 | | tuple[tuple[int, int], tuple[int, int, int]] |
| 89 | | tuple[ |
| 90 | tuple[int, int], |
| 91 | tuple[int, int, int], |
| 92 | tuple[int, int, int], |
| 93 | tuple[int, int, int], |
| 94 | ] |
| 95 | | tuple[ |
| 96 | tuple[int, int], |
| 97 | tuple[int, int, int], |
| 98 | tuple[int, int, int], |
| 99 | tuple[int, int], |
| 100 | tuple[int, int, int], |
| 101 | ] |
| 102 | ) |
| 103 | |
| 104 | |
| 105 | @dataclass(frozen=True, slots=True) |
| 106 | class DiarizedTranscriptionSegment: |
| 107 | """A timestamped, speaker-attributed segment produced by an ASR model.""" |
| 108 | |
| 109 | start: float |
| 110 | end: float |
| 111 | speaker: str |
| 112 | text: str |
| 113 | |
| 114 | |
| 115 | class StreamingTranscriptionPostProcessor: |
| 116 | """Stateful streaming post-processor for transcription deltas.""" |
| 117 | |
| 118 | def process_delta(self, text_delta: str, finished: bool) -> str: |
| 119 | return text_delta |
| 120 | |
| 121 | |
| 122 | def _require_is_multimodal(is_multimodal: Tensor | None) -> Tensor: |
| 123 | """A helper function to be used in the context of |
| 124 | [vllm.model_executor.models.interfaces.SupportsMultiModal.embed_input_ids][] |
| 125 | to provide a better error message. |
| 126 | """ |
| 127 | if is_multimodal is None: |
| 128 | raise ValueError( |
| 129 | "`embed_input_ids` now requires `is_multimodal` arg, " |
| 130 | "please update your model runner according to " |
| 131 | "https://github.com/vllm-project/vllm/pull/16229." |
| 132 | ) |
| 133 | |
| 134 | return is_multimodal |
| 135 | |
| 136 | |
| 137 | # Cache results of `SupportsMultiModal.get_language_model` |
| 138 | _language_model_by_module: "weakref.WeakKeyDictionary[nn.Module, VllmModel]" = ( |
| 139 | weakref.WeakKeyDictionary() |
| 140 | ) |
| 141 | |
| 142 | |
| 143 | @runtime_checkable |
| 144 | class SupportsMultiModalEmbeddings(Protocol): |
| 145 | """The interface for models that can merge external multimodal embeddings.""" |
| 146 | |
| 147 | supports_multimodal_embeddings: ClassVar[Literal[True]] = True |
| 148 | |
| 149 | @overload |
| 150 | def embed_input_ids( |
| 151 | self, |
| 152 | input_ids: Tensor, |
| 153 | multimodal_embeddings: None = None, |
| 154 | *, |
| 155 | is_multimodal: Tensor | None = None, |
| 156 | ) -> Tensor: ... |
| 157 | |
| 158 | @overload |
| 159 | def embed_input_ids( |
| 160 | self, |
| 161 | input_ids: Tensor, |
| 162 | multimodal_embeddings: MultiModalEmbeddings, |
| 163 | *, |
| 164 | is_multimodal: Tensor, |
| 165 | ) -> Tensor: ... |
| 166 | |
| 167 | |
| 168 | @runtime_checkable |
| 169 | class SupportsMultiModal(SupportsMultiModalEmbeddings, Protocol): |
| 170 | """The interface required for all multi-modal models.""" |
| 171 | |
| 172 | supports_multimodal: ClassVar[Literal[True]] = True |
| 173 | """ |
| 174 | A flag that indicates this model supports multi-modal inputs. |
| 175 | |
| 176 | Note: |
| 177 | There is no need to redefine this flag if this class is in the |
| 178 | MRO of your model class. |
| 179 | """ |
| 180 | |
| 181 | supports_multimodal_raw_input_only: ClassVar[bool] = False |
| 182 | """ |
| 183 | A flag that indicates this model supports multi-modal inputs and processes |
| 184 | them in their raw form and not embeddings. |
| 185 | """ |
| 186 | |
| 187 | supports_encoder_tp_data: ClassVar[bool] = False |
| 188 | """ |
| 189 | A flag that indicates whether this model supports |
| 190 | `multimodal_config.mm_encoder_tp_mode="data"`. |
| 191 | """ |
| 192 | |
| 193 | supports_mm_device_do_normalize: ClassVar[bool] = False |
| 194 | """ |
| 195 | A flag that indicates whether this model supports |
| 196 | `multimodal_config.mm_device_do_normalize`. |
| 197 | """ |
| 198 | |
| 199 | supports_tower_connector_lora: ClassVar[bool] = False |
| 200 | """ |
| 201 | A flag that indicates whether this model supports |
| 202 | `lora_config.enable_tower_connector_lora`. |
| 203 | """ |
| 204 | |
| 205 | requires_raw_input_tokens: ClassVar[bool] = False |
| 206 | """ |
| 207 | A flag that indicates this model processes input id tokens |
| 208 | in their raw form and not input embeddings. |
| 209 | """ |
| 210 | |
| 211 | _processor_factory: ClassVar["_ProcessorFactories"] |
| 212 | """ |
| 213 | Set internally by `MultiModalRegistry.register_processor`. |
| 214 | """ |
| 215 | |
| 216 | _language_model_names: list[str] = [] |
| 217 | """ |
| 218 | Set internally by `_mark_language_model`. |
| 219 | """ |
| 220 | |
| 221 | _tower_model_names: list[str] = [] |
| 222 | """ |
| 223 | Set internally by `_mark_tower_model`. |
| 224 | """ |
| 225 | |
| 226 | _has_oov_mm_tokens: bool = False |
| 227 | """ |
| 228 | In general, this should be set at init time by invoking |
| 229 | `configure_mm_token_handling` models & passing all potentially |
| 230 | OOV multimodal tokens. |
| 231 | """ |
| 232 | |
| 233 | @classmethod |
| 234 | def get_placeholder_str(cls, modality: str, i: int) -> str | None: |
| 235 | """Get the placeholder text for the `i`th `modality` item in the prompt.""" |
| 236 | ... |
| 237 | |
| 238 | def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: |
| 239 | """Returns multimodal embeddings generated from multimodal kwargs |
| 240 | to be merged with text embeddings. |
| 241 | |
| 242 | Note: |
| 243 | The returned multimodal embeddings must be in the same order as |
| 244 | the appearances of their corresponding multimodal data item in the |
| 245 | input prompt. |
| 246 | |
| 247 | """ |
| 248 | ... |
| 249 | |
| 250 | def configure_mm_token_handling(self, vocab_size: int, mm_token_ids: list[int]): |
| 251 | """Check if any multimodal tokens are out of vocabulary. If so, we will |
| 252 | explicitly mask all multimodal tokens out when computing text embeddings, |
| 253 | since the multimodal embeddings will be scattered over the results. |
| 254 | """ |
| 255 | self._has_oov_mm_tokens = any(tok_id >= vocab_size for tok_id in mm_token_ids) |
| 256 | logger.info( |
| 257 | "Contains out of vocabulary multimodal tokens? %s", |
| 258 | self._has_oov_mm_tokens, |
| 259 | ) |
| 260 | |
| 261 | def get_language_model(self) -> "VllmModel": |
| 262 | """Returns the underlying language model used for text generation. |
| 263 | |
| 264 | This is typically the `torch.nn.Module` instance responsible for |
| 265 | processing the merged multimodal embeddings and producing hidden states |
| 266 | |
| 267 | Returns: |
| 268 | torch.nn.Module: The core language model component. |
| 269 | |
| 270 | """ |
| 271 | # Cached |
| 272 | assert isinstance(self, nn.Module) |
| 273 | self_module = self |
| 274 | if self_module in _language_model_by_module: |
| 275 | return _language_model_by_module[self_module] |
| 276 | |
| 277 | if self._language_model_names: |
| 278 | mod = self_module |
| 279 | for attr in common_prefix( |
| 280 | [name.split(".") for name in self._language_model_names] |
| 281 | ): |
| 282 | if attr: |
| 283 | mod = getattr(mod, attr) |
| 284 | |
| 285 | candidate: object = mod |
| 286 | if candidate is not self_module and isinstance(candidate, VllmModel): |
| 287 | _language_model_by_module[self_module] = candidate |
| 288 | return candidate |
| 289 | |
| 290 | # Fallback |
| 291 | for mod in self_module.children(): |
| 292 | candidate = mod |
| 293 | if isinstance(candidate, VllmModel): |
| 294 | _language_model_by_module[self_module] = candidate |
| 295 | return candidate |
| 296 | |
| 297 | raise NotImplementedError( |
| 298 | f"No language model found in {type(self).__name__}! " |
| 299 | "You should initialize it via `_mark_language_model`, " |
| 300 | "and make sure `embed_input_ids` is implemented." |
| 301 | ) |
| 302 | |
| 303 | @contextmanager |
| 304 | def _mark_language_model( |
| 305 | self, |
| 306 | vllm_config: "VllmConfig", |
| 307 | *, |
| 308 | targets: type[nn.Module] | tuple[type[nn.Module], ...] | None = None, |
| 309 | ): |
| 310 | """Mark each child module that was assigned to this model during this context |
| 311 | as a language model component. |
| 312 | |
| 313 | Language model components are automatically skipped in `--mm-encoder-only` |
| 314 | mode. |
| 315 | |
| 316 | If `targets` is set, instead include descendants that are an instance |
| 317 | of `targets`, even if they aren't direct children. |
| 318 | """ |
| 319 | from .utils import StageMissingLayer, collect_children, no_init_weights |
| 320 | |
| 321 | mm_config = vllm_config.model_config.get_multimodal_config() |
| 322 | |
| 323 | with collect_children(self, targets=targets) as children_names: # noqa: SIM117 |
| 324 | with ( |
| 325 | no_init_weights( |
| 326 | self, |
| 327 | lambda mod: StageMissingLayer("language_model", mod), |
| 328 | targets=targets, |
| 329 | ) |
| 330 | if mm_config.mm_encoder_only |
| 331 | else nullcontext() |
| 332 | ): |
| 333 | yield |
| 334 | |
| 335 | self._language_model_names = children_names |
| 336 | |
| 337 | @contextmanager |
| 338 | def _mark_tower_model( |
| 339 | self, |
| 340 | vllm_config: "VllmConfig", |
| 341 | modalities: set[str] | str, |
| 342 | *, |
| 343 | targets: type[nn.Module] | tuple[type[nn.Module], ...] | None = None, |
| 344 | ): |
| 345 | """Mark each child module that was assigned to this model during this context |
| 346 | as a tower model component. |
| 347 | |
| 348 | Tower model components are automatically skipped when `--limit-mm-per-prompt` |
| 349 | is set to zero for all of their modalities. |
| 350 | |
| 351 | If `targets` is set, instead include descendants that are an instance |
| 352 | of `targets`, even if they aren't direct children. |
| 353 | |
| 354 | Marked components are also routed through the active offloader (when |
| 355 | it supports tower offloading), since `make_layers` only ever sees the |
| 356 | decoder layer stack. |
| 357 | """ |
| 358 | from vllm.model_executor.offloader import get_offloader |
| 359 | |
| 360 | from .utils import StageMissingLayer, collect_children, no_init_weights |
| 361 | |
| 362 | if isinstance(modalities, str): |
| 363 | modalities = {modalities} |
| 364 | |
| 365 | if modalities == {"image", "video"}: |
| 366 | stage_name = "vision_tower" |
| 367 | else: |
| 368 | stage_name = "_".join([*modalities, "tower"]) |
| 369 | |
| 370 | mm_config = vllm_config.model_config.get_multimodal_config() |
| 371 | |
| 372 | with collect_children(self, targets=targets) as children_names: # noqa: SIM117 |
| 373 | with ( |
| 374 | no_init_weights( |
| 375 | self, |
| 376 | lambda mod: StageMissingLayer(stage_name, mod), |
| 377 | targets=targets, |
| 378 | ) |
| 379 | if all(mm_config.get_limit_per_prompt(m) == 0 for m in modalities) |
| 380 | else nullcontext() |
| 381 | ): |
| 382 | yield |
| 383 | |
| 384 | self._tower_model_names = children_names |
| 385 | |
| 386 | # Towers are constructed directly, so `make_layers` never routes them |
| 387 | # through the offloader. Do it here, at the same construction stage, |
| 388 | # so offloaded tower weights are never allocated on the device. |
| 389 | offloader = get_offloader() |
| 390 | if offloader.supports_tower_offload: |
| 391 | for name in children_names: |
| 392 | modules = (module for module in [attrgetter(name)(self)]) |
| 393 | offloader.wrap_modules(modules, prefix=name) |
| 394 | |
| 395 | @contextmanager |
| 396 | def _mark_composite_model( |
| 397 | self, |
| 398 | vllm_config: "VllmConfig", |
| 399 | *, |
| 400 | language_targets: type[nn.Module] | tuple[type[nn.Module], ...], |
| 401 | tower_targets: dict[str, type[nn.Module] | tuple[type[nn.Module], ...]], |
| 402 | ): |
| 403 | """Composite wrapper over `_mark_language_model` and |
| 404 | `_mark_tower_model` by modality. |
| 405 | """ |
| 406 | with ExitStack() as stack: |
| 407 | stack.enter_context( |
| 408 | self._mark_language_model( |
| 409 | vllm_config, |
| 410 | targets=language_targets, |
| 411 | ) |
| 412 | ) |
| 413 | |
| 414 | for modality, modality_targets in tower_targets.items(): |
| 415 | stack.enter_context( |
| 416 | self._mark_tower_model( |
| 417 | vllm_config, |
| 418 | modality, |
| 419 | targets=modality_targets, |
| 420 | ) |
| 421 | ) |
| 422 | |
| 423 | yield |
| 424 | |
| 425 | @deprecated( |
| 426 | "get_num_mm_encoder_tokens is deprecated; use get_mm_lora_token_counts instead." |
| 427 | ) |
| 428 | def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: ... |
| 429 | |
| 430 | @deprecated( |
| 431 | "get_num_mm_connector_tokens is deprecated; use " |
| 432 | "get_mm_lora_token_counts instead." |
| 433 | ) |
| 434 | def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: ... |
| 435 | |
| 436 | def get_mm_lora_token_counts( |
| 437 | self, |
| 438 | *, |
| 439 | modality: str, |
| 440 | mm_kwargs: "MultiModalKwargsItem | None", |
| 441 | num_mm_embeds: int, |
| 442 | ) -> tuple[int, int | None]: |
| 443 | """Return ``(tower_tokens, connector_tokens)`` for multimodal LoRA mappings. |
| 444 | |
| 445 | MM LoRA uses these counts to build adapter mappings for the tower and |
| 446 | connector forwards. Models with multiple modalities can override this |
| 447 | when each modality has different encoder padding or pooling behavior. |
| 448 | """ |
| 449 | del modality, mm_kwargs |
| 450 | num_encoder_tokens = self.get_num_mm_encoder_tokens(num_mm_embeds) |
| 451 | num_connector_tokens = self.get_num_mm_connector_tokens(num_encoder_tokens) |
| 452 | return ( |
| 453 | num_encoder_tokens, |
| 454 | num_connector_tokens if isinstance(num_connector_tokens, int) else None, |
| 455 | ) |
| 456 | |
| 457 | def _embed_text_input_ids( |
| 458 | self, |
| 459 | input_ids: Tensor, |
| 460 | embed_input_ids: Callable[[Tensor], Tensor], |
| 461 | *, |
| 462 | is_multimodal: Tensor | None, |
| 463 | ) -> Tensor: |
| 464 | if is_multimodal is not None and self._has_oov_mm_tokens: |
| 465 | # Force all input IDs to be in vocab; we do this instead of squeezing |
| 466 | # to ensure that any external configuration requiring offset tracking, |
| 467 | # e.g., LoRA, are applied correctly regardless of whether or not |
| 468 | # we have multimodal tokens. |
| 469 | in_vocab_ids = input_ids.masked_fill( |
| 470 | is_multimodal.to(device=input_ids.device, non_blocking=True), 0 |
| 471 | ) |
| 472 | return embed_input_ids(in_vocab_ids) |
| 473 | |
| 474 | return embed_input_ids(input_ids) |
| 475 | |
| 476 | @overload |
| 477 | def embed_input_ids( |
| 478 | self, |
| 479 | input_ids: Tensor, |
| 480 | multimodal_embeddings: None = None, |
| 481 | *, |
| 482 | is_multimodal: Tensor | None = None, |
| 483 | ) -> Tensor: ... |
| 484 | |
| 485 | @overload |
| 486 | def embed_input_ids( |
| 487 | self, |
| 488 | input_ids: Tensor, |
| 489 | multimodal_embeddings: MultiModalEmbeddings, |
| 490 | *, |
| 491 | is_multimodal: Tensor, |
| 492 | ) -> Tensor: ... |
| 493 | |
| 494 | def embed_input_ids( |
| 495 | self, |
| 496 | input_ids: Tensor, |
| 497 | multimodal_embeddings: MultiModalEmbeddings | None = None, |
| 498 | *, |
| 499 | is_multimodal: Tensor | None = None, |
| 500 | ) -> Tensor: |
| 501 | """Apply token embeddings to `input_ids`. |
| 502 | |
| 503 | If `multimodal_embeddings` is passed, scatter them into |
| 504 | `input_ids` according to the mask `is_multimodal`. |
| 505 | |
| 506 | NOTE: If this model has multimodal tokens that are of vocabulary |
| 507 | (i.e., self._has_oov_mm_tokens=True), the input_ids will be copied |
| 508 | and masked to 0 during the forward pass for the text embeddings. |
| 509 | """ |
| 510 | from .utils import _merge_multimodal_embeddings |
| 511 | |
| 512 | # Get text embeddings first; multimodal embeddings will clobber |
| 513 | # any invalid contents in the indices of multimodal embeddings |
| 514 | # for the in vocabulary and out of vocabulary case. |
| 515 | inputs_embeds = self._embed_text_input_ids( |
| 516 | input_ids, |
| 517 | self.get_language_model().embed_input_ids, |
| 518 | is_multimodal=is_multimodal, |
| 519 | ) |
| 520 | |
| 521 | if multimodal_embeddings is None or len(multimodal_embeddings) == 0: |
| 522 | return inputs_embeds |
| 523 | |
| 524 | return _merge_multimodal_embeddings( |
| 525 | inputs_embeds=inputs_embeds, |
| 526 | multimodal_embeddings=multimodal_embeddings, |
| 527 | is_multimodal=_require_is_multimodal(is_multimodal), |
| 528 | ) |
| 529 | |
| 530 | |
| 531 | @runtime_checkable |
| 532 | class SupportsMultiModalPruning(Protocol): |
| 533 | """The interface required for models that support returning both input |
| 534 | embeddings and positions. Model may require custom positions for dynamic |
| 535 | pruning of multimodal embeddings. |
| 536 | """ |
| 537 | |
| 538 | supports_multimodal_pruning: ClassVar[Literal[True]] = True |
| 539 | |
| 540 | supported_video_pruning_methods: ClassVar[tuple["VideoPruningMethod", ...]] = ( |
| 541 | "evs", |
| 542 | ) |
| 543 | """Video pruning methods (as reported by |
| 544 | `MultiModalConfig.get_video_pruning_spec`) implemented by this model. |
| 545 | Models supporting methods beyond EVS should override this.""" |
| 546 | |
| 547 | def recompute_mrope_positions( |
| 548 | self, |
| 549 | input_ids: list[int] | torch.Tensor, |
| 550 | multimodal_embeddings: Sequence[torch.Tensor], |
| 551 | mrope_positions: torch.LongTensor, |
| 552 | num_computed_tokens: int, |
| 553 | ) -> tuple[Sequence[torch.Tensor], Tensor, int]: |
| 554 | """Update part of input mrope positions (starting with |
| 555 | num_computed_tokens index). Original mrope_positions are computed |
| 556 | for unpruned sequence and becomes incorrect once pruning occurs, |
| 557 | so once we prune media tokens we should reflect this in the |
| 558 | mrope_positions before we feed it to LLM. |
| 559 | |
| 560 | Args: |
| 561 | input_ids: (N,) All input tokens of the prompt containing |
| 562 | entire sequence. Either a host-side list or an already |
| 563 | device-resident tensor. |
| 564 | multimodal_embeddings: Sequence of multimodal embeddings that |
| 565 | fits into the prefill chunk that is being processed. |
| 566 | mrope_positions: Existing mrope positions (3, N) for entire |
| 567 | sequence |
| 568 | num_computed_tokens: A number of computed tokens so far. |
| 569 | |
| 570 | Returns: |
| 571 | Tuple of (multimodal_embeddings, mrope_positions, |
| 572 | mrope_position_delta). |
| 573 | |
| 574 | """ |
| 575 | ... |
| 576 | |
| 577 | |
| 578 | @overload |
| 579 | def supports_multimodal(model: type[object]) -> TypeIs[type[SupportsMultiModal]]: ... |
| 580 | |
| 581 | |
| 582 | @overload |
| 583 | def supports_multimodal(model: object) -> TypeIs[SupportsMultiModal]: ... |
| 584 | |
| 585 | |
| 586 | def supports_multimodal( |
| 587 | model: type[object] | object, |
| 588 | ) -> TypeIs[type[SupportsMultiModal]] | TypeIs[SupportsMultiModal]: |
| 589 | return getattr(model, "supports_multimodal", False) |
| 590 | |
| 591 | |
| 592 | @overload |
| 593 | def supports_multimodal_embeddings( |
| 594 | model: type[object], |
| 595 | ) -> TypeIs[type[SupportsMultiModalEmbeddings]]: ... |
| 596 | |
| 597 | |
| 598 | @overload |
| 599 | def supports_multimodal_embeddings( |
| 600 | model: object, |
| 601 | ) -> TypeIs[SupportsMultiModalEmbeddings]: ... |
| 602 | |
| 603 | |
| 604 | def supports_multimodal_embeddings( |
| 605 | model: type[object] | object, |
| 606 | ) -> TypeIs[type[SupportsMultiModalEmbeddings]] | TypeIs[SupportsMultiModalEmbeddings]: |
| 607 | return getattr(model, "supports_multimodal_embeddings", False) |
| 608 | |
| 609 | |
| 610 | def supports_multimodal_raw_input_only(model: type[object] | object) -> bool: |
| 611 | return getattr(model, "supports_multimodal_raw_input_only", False) |
| 612 | |
| 613 | |
| 614 | def requires_raw_input_tokens(model: type[object] | object) -> bool: |
| 615 | return getattr(model, "requires_raw_input_tokens", False) |
| 616 | |
| 617 | |
| 618 | def supports_multimodal_encoder_tp_data(model: type[object] | object) -> bool: |
| 619 | return getattr(model, "supports_encoder_tp_data", False) |
| 620 | |
| 621 | |
| 622 | @overload |
| 623 | def supports_multimodal_pruning( |
| 624 | model: type[object], |
| 625 | ) -> TypeIs[type[SupportsMultiModalPruning]]: ... |
| 626 | |
| 627 | |
| 628 | @overload |
| 629 | def supports_multimodal_pruning(model: object) -> TypeIs[SupportsMultiModalPruning]: ... |
| 630 | |
| 631 | |
| 632 | def supports_multimodal_pruning( |
| 633 | model: type[object] | object, |
| 634 | ) -> TypeIs[type[SupportsMultiModalPruning]] | TypeIs[SupportsMultiModalPruning]: |
| 635 | return getattr(model, "supports_multimodal_pruning", False) |
| 636 | |
| 637 | |
| 638 | @runtime_checkable |
| 639 | class SupportsScoreTemplate(Protocol): |
| 640 | """The interface required for all models that support score template.""" |
| 641 | |
| 642 | supports_score_template: ClassVar[Literal[True]] = True |
| 643 | """ |
| 644 | A flag that indicates this model supports score template. |
| 645 | |
| 646 | Note: |
| 647 | There is no need to redefine this flag if this class is in the |
| 648 | MRO of your model class. |
| 649 | """ |
| 650 | |
| 651 | @classmethod |
| 652 | def get_score_template(cls, query: str, document: str) -> str | None: |
| 653 | """Generate a full prompt by populating the score template with query and document content.""" # noqa: E501 |
| 654 | ... |
| 655 | |
| 656 | @classmethod |
| 657 | def post_process_tokens(cls, prompt: "TokensPrompt") -> None: |
| 658 | """Perform architecture-specific manipulations on the input tokens.""" |
| 659 | ... |
| 660 | |
| 661 | |
| 662 | @overload |
| 663 | def supports_score_template( |
| 664 | model: type[object], |
| 665 | ) -> TypeIs[type[SupportsScoreTemplate]]: ... |
| 666 | |
| 667 | |
| 668 | @overload |
| 669 | def supports_score_template(model: object) -> TypeIs[SupportsScoreTemplate]: ... |
| 670 | |
| 671 | |
| 672 | def supports_score_template( |
| 673 | model: type[object] | object, |
| 674 | ) -> TypeIs[type[SupportsScoreTemplate]] | TypeIs[SupportsScoreTemplate]: |
| 675 | return getattr(model, "supports_score_template", False) |
| 676 | |
| 677 | |
| 678 | @runtime_checkable |
| 679 | class SupportsLoRA(Protocol): |
| 680 | """The interface required for all models that support LoRA.""" |
| 681 | |
| 682 | supports_lora: ClassVar[Literal[True]] = True |
| 683 | """ |
| 684 | A flag that indicates this model supports LoRA. |
| 685 | |
| 686 | Note: |
| 687 | There is no need to redefine this flag if this class is in the |
| 688 | MRO of your model class. |
| 689 | """ |
| 690 | is_3d_moe_weight: ClassVar[bool] = False |
| 691 | is_non_gated_moe: ClassVar[bool] = False |
| 692 | # The `embedding_module` and `embedding_padding_modules` |
| 693 | # are empty by default. |
| 694 | embedding_modules: ClassVar[dict[str, str]] = {} |
| 695 | packed_modules_mapping: dict[str, list[str]] = {} |
| 696 | # Module prefixes to skip during LoRA loading (e.g., ["mtp."] for MTP layers) |
| 697 | lora_skip_prefixes: ClassVar[list[str]] = [] |
| 698 | lora_manager: "LoRAModelManager | None" |
| 699 | |
| 700 | |
| 701 | # We can't use runtime_checkable with ClassVar for issubclass checks |
| 702 | # so we need to treat the class as an instance and use isinstance instead |
| 703 | @runtime_checkable |
| 704 | class _SupportsLoRAType(Protocol): |
| 705 | supports_lora: Literal[True] |
| 706 | |
| 707 | packed_modules_mapping: dict[str, list[str]] |
| 708 | embedding_modules: dict[str, str] |
| 709 | |
| 710 | |
| 711 | @overload |
| 712 | def supports_lora(model: type[object]) -> TypeIs[type[SupportsLoRA]]: ... |
| 713 | |
| 714 | |
| 715 | @overload |
| 716 | def supports_lora(model: object) -> TypeIs[SupportsLoRA]: ... |
| 717 | |
| 718 | |
| 719 | def supports_lora( |
| 720 | model: type[object] | object, |
| 721 | ) -> TypeIs[type[SupportsLoRA]] | TypeIs[SupportsLoRA]: |
| 722 | result = _supports_lora(model) |
| 723 | |
| 724 | if not result: |
| 725 | lora_attrs = ( |
| 726 | "packed_modules_mapping", |
| 727 | "embedding_modules", |
| 728 | ) |
| 729 | missing_attrs = tuple(attr for attr in lora_attrs if not hasattr(model, attr)) |
| 730 | |
| 731 | if getattr(model, "supports_lora", False): |
| 732 | if missing_attrs: |
| 733 | logger.warning( |
| 734 | "The model (%s) sets `supports_lora=True`, " |
| 735 | "but is missing LoRA-specific attributes: %s", |
| 736 | model, |
| 737 | missing_attrs, |
| 738 | ) |
| 739 | else: |
| 740 | if not missing_attrs: |
| 741 | logger.warning( |
| 742 | "The model (%s) contains all LoRA-specific attributes, " |
| 743 | "but does not set `supports_lora=True`.", |
| 744 | model, |
| 745 | ) |
| 746 | |
| 747 | return result |
| 748 | |
| 749 | |
| 750 | def _supports_lora(model: type[object] | object) -> bool: |
| 751 | if isinstance(model, type): |
| 752 | return isinstance(model, _SupportsLoRAType) |
| 753 | |
| 754 | return isinstance(model, SupportsLoRA) |
| 755 | |
| 756 | |
| 757 | class _MakeEmptyIntermediateTensors(Protocol): |
| 758 | def __call__( |
| 759 | self, |
| 760 | batch_size: int, |
| 761 | dtype: torch.dtype, |
| 762 | device: torch.device, |
| 763 | ) -> "IntermediateTensors": ... |
| 764 | |
| 765 | |
| 766 | @runtime_checkable |
| 767 | class SupportsPP(Protocol): |
| 768 | """The interface required for all models that support pipeline parallel.""" |
| 769 | |
| 770 | supports_pp: ClassVar[Literal[True]] = True |
| 771 | """ |
| 772 | A flag that indicates this model supports pipeline parallel. |
| 773 | |
| 774 | Note: |
| 775 | There is no need to redefine this flag if this class is in the |
| 776 | MRO of your model class. |
| 777 | """ |
| 778 | |
| 779 | make_empty_intermediate_tensors: _MakeEmptyIntermediateTensors |
| 780 | """Called when PP rank > 0 for profiling purposes.""" |
| 781 | |
| 782 | def forward( |
| 783 | self, |
| 784 | input_ids: Tensor | None, |
| 785 | positions: Tensor, |
| 786 | *, |
| 787 | intermediate_tensors: "IntermediateTensors | None", |
| 788 | ) -> "Tensor | IntermediateTensors | tuple[Tensor, list[Tensor]]": |
| 789 | """Accept [`IntermediateTensors`][vllm.sequence.IntermediateTensors] when |
| 790 | PP rank > 0. |
| 791 | |
| 792 | Return [`IntermediateTensors`][vllm.sequence.IntermediateTensors] only |
| 793 | for the last PP rank. |
| 794 | """ |
| 795 | ... |
| 796 | |
| 797 | |
| 798 | # We can't use runtime_checkable with ClassVar for issubclass checks |
| 799 | # so we need to treat the class as an instance and use isinstance instead |
| 800 | @runtime_checkable |
| 801 | class _SupportsPPType(Protocol): |
| 802 | supports_pp: Literal[True] |
| 803 | |
| 804 | make_empty_intermediate_tensors: _MakeEmptyIntermediateTensors |
| 805 | |
| 806 | def forward( |
| 807 | self, |
| 808 | input_ids: Tensor | None, |
| 809 | positions: Tensor, |
| 810 | *, |
| 811 | intermediate_tensors: "IntermediateTensors | None", |
| 812 | ) -> "Tensor | IntermediateTensors | tuple[Tensor, list[Tensor]]": ... |
| 813 | |
| 814 | |
| 815 | @overload |
| 816 | def supports_pp(model: type[object]) -> TypeIs[type[SupportsPP]]: ... |
| 817 | |
| 818 | |
| 819 | @overload |
| 820 | def supports_pp(model: object) -> TypeIs[SupportsPP]: ... |
| 821 | |
| 822 | |
| 823 | def supports_pp( |
| 824 | model: type[object] | object, |
| 825 | ) -> bool | TypeIs[type[SupportsPP]] | TypeIs[SupportsPP]: |
| 826 | supports_attributes = _supports_pp_attributes(model) |
| 827 | supports_inspect = _supports_pp_inspect(model) |
| 828 | |
| 829 | if supports_attributes and not supports_inspect: |
| 830 | logger.warning( |
| 831 | "The model (%s) sets `supports_pp=True`, but does not accept " |
| 832 | "`intermediate_tensors` in its `forward` method", |
| 833 | model, |
| 834 | ) |
| 835 | |
| 836 | if not supports_attributes: |
| 837 | pp_attrs = ("make_empty_intermediate_tensors",) |
| 838 | missing_attrs = tuple(attr for attr in pp_attrs if not hasattr(model, attr)) |
| 839 | |
| 840 | if getattr(model, "supports_pp", False): |
| 841 | if missing_attrs: |
| 842 | logger.warning( |
| 843 | "The model (%s) sets `supports_pp=True`, " |
| 844 | "but is missing PP-specific attributes: %s", |
| 845 | model, |
| 846 | missing_attrs, |
| 847 | ) |
| 848 | else: |
| 849 | if not missing_attrs: |
| 850 | logger.warning( |
| 851 | "The model (%s) contains all PP-specific attributes, " |
| 852 | "but does not set `supports_pp=True`.", |
| 853 | model, |
| 854 | ) |
| 855 | |
| 856 | return supports_attributes and supports_inspect |
| 857 | |
| 858 | |
| 859 | def _supports_pp_attributes(model: type[object] | object) -> bool: |
| 860 | if isinstance(model, type): |
| 861 | return SupportsPP in model.__mro__ or isinstance(model, _SupportsPPType) |
| 862 | |
| 863 | return isinstance(model, SupportsPP) |
| 864 | |
| 865 | |
| 866 | def _supports_pp_inspect(model: type[object] | object) -> bool: |
| 867 | model_forward = getattr(model, "forward", None) |
| 868 | if not callable(model_forward): |
| 869 | return False |
| 870 | |
| 871 | return supports_kw(model_forward, "intermediate_tensors") |
| 872 | |
| 873 | |
| 874 | @runtime_checkable |
| 875 | class HasInnerState(Protocol): |
| 876 | """The interface required for all models that has inner state.""" |
| 877 | |
| 878 | has_inner_state: ClassVar[Literal[True]] = True |
| 879 | """ |
| 880 | A flag that indicates this model has inner state. |
| 881 | Models that has inner state usually need access to the scheduler_config |
| 882 | for max_num_seqs, etc. True for e.g. both Mamba and Jamba. |
| 883 | """ |
| 884 | |
| 885 | |
| 886 | @overload |
| 887 | def has_inner_state(model: type[object]) -> TypeIs[type[HasInnerState]]: ... |
| 888 | |
| 889 | |
| 890 | @overload |
| 891 | def has_inner_state(model: object) -> TypeIs[HasInnerState]: ... |
| 892 | |
| 893 | |
| 894 | def has_inner_state( |
| 895 | model: type[object] | object, |
| 896 | ) -> TypeIs[type[HasInnerState]] | TypeIs[HasInnerState]: |
| 897 | return getattr(model, "has_inner_state", False) |
| 898 | |
| 899 | |
| 900 | @runtime_checkable |
| 901 | class IsAttentionFree(Protocol): |
| 902 | """The interface required for all models like Mamba that lack attention, |
| 903 | but do have state whose size is constant wrt the number of tokens.""" |
| 904 | |
| 905 | is_attention_free: ClassVar[Literal[True]] = True |
| 906 | """ |
| 907 | A flag that indicates this model has no attention. |
| 908 | Used for block manager and attention backend selection. |
| 909 | True for Mamba but not Jamba. |
| 910 | """ |
| 911 | |
| 912 | |
| 913 | @overload |
| 914 | def is_attention_free(model: type[object]) -> TypeIs[type[IsAttentionFree]]: ... |
| 915 | |
| 916 | |
| 917 | @overload |
| 918 | def is_attention_free(model: object) -> TypeIs[IsAttentionFree]: ... |
| 919 | |
| 920 | |
| 921 | def is_attention_free( |
| 922 | model: type[object] | object, |
| 923 | ) -> TypeIs[type[IsAttentionFree]] | TypeIs[IsAttentionFree]: |
| 924 | return getattr(model, "is_attention_free", False) |
| 925 | |
| 926 | |
| 927 | @runtime_checkable |
| 928 | class IsHybrid(Protocol): |
| 929 | """The interface required for all models like Jamba that have both |
| 930 | attention and mamba blocks, indicates that |
| 931 | hf_config has 'layers_block_type'""" |
| 932 | |
| 933 | is_hybrid: ClassVar[Literal[True]] = True |
| 934 | """ |
| 935 | A flag that indicates this model has both mamba and attention blocks |
| 936 | , also indicates that the model's hf_config has |
| 937 | 'layers_block_type' """ |
| 938 | |
| 939 | @classmethod |
| 940 | def get_mamba_state_shape_from_config( |
| 941 | cls, |
| 942 | vllm_config: "VllmConfig", |
| 943 | ) -> MambaStateShapes: |
| 944 | """Calculate shapes for Mamba's convolutional and state caches. |
| 945 | |
| 946 | Args: |
| 947 | vllm_config: vLLM config |
| 948 | |
| 949 | Returns: |
| 950 | Shapes for each state cache used by the model. |
| 951 | |
| 952 | """ |
| 953 | ... |
| 954 | |
| 955 | @classmethod |
| 956 | def get_mamba_state_copy_func(cls) -> tuple["MambaStateCopyFunc", ...]: |
| 957 | """Calculate copy-function callables for each Mamba state. |
| 958 | |
| 959 | Returns: |
| 960 | A tuple of MambaStateCopyFunc callables that correspond, in order, |
| 961 | to the Mamba states produced by the model. Each callable accepts |
| 962 | (state, block_ids, cur_block_idx, num_accepted_tokens) and returns |
| 963 | a MambaCopySpec describing the memory-copy parameters for prefix |
| 964 | caching in align mode. |
| 965 | |
| 966 | """ |
| 967 | ... |
| 968 | |
| 969 | @classmethod |
| 970 | def get_mamba_state_copy_funcs( |
| 971 | cls, |
| 972 | mamba_types: set["MambaAttentionBackendEnum"], |
| 973 | ) -> "MambaStateCopyFuncsByType": |
| 974 | copy_funcs = cls.get_mamba_state_copy_func() |
| 975 | return {mamba_type: copy_funcs for mamba_type in mamba_types} |
| 976 | |
| 977 | |
| 978 | @overload |
| 979 | def is_hybrid(model: type[object]) -> TypeIs[type[IsHybrid]]: ... |
| 980 | |
| 981 | |
| 982 | @overload |
| 983 | def is_hybrid(model: object) -> TypeIs[IsHybrid]: ... |
| 984 | |
| 985 | |
| 986 | def is_hybrid( |
| 987 | model: type[object] | object, |
| 988 | ) -> TypeIs[type[IsHybrid]] | TypeIs[IsHybrid]: |
| 989 | return getattr(model, "is_hybrid", False) |
| 990 | |
| 991 | |
| 992 | @runtime_checkable |
| 993 | class MixtureOfExperts(Protocol): |
| 994 | """Check if the model is a mixture of experts (MoE) model.""" |
| 995 | |
| 996 | expert_weights: MutableSequence[Sequence[Tensor]] |
| 997 | """ |
| 998 | Expert weights saved in this rank. |
| 999 | |
| 1000 | The first dimension is the layer, and the second dimension is different |
| 1001 | parameters in the layer, e.g. up/down projection weights. |
| 1002 | """ |
| 1003 | |
| 1004 | num_moe_layers: int |
| 1005 | """Number of MoE layers in this model.""" |
| 1006 | |
| 1007 | num_expert_groups: int |
| 1008 | """Number of expert groups in this model.""" |
| 1009 | |
| 1010 | num_logical_experts: int |
| 1011 | """Number of logical experts in this model.""" |
| 1012 | |
| 1013 | num_physical_experts: int |
| 1014 | """Number of physical experts in this model.""" |
| 1015 | |
| 1016 | num_local_physical_experts: int |
| 1017 | """Number of local physical experts in this model.""" |
| 1018 | |
| 1019 | num_routed_experts: int |
| 1020 | """Number of routed experts in this model.""" |
| 1021 | |
| 1022 | num_shared_experts: int |
| 1023 | """Number of shared experts in this model.""" |
| 1024 | |
| 1025 | num_redundant_experts: int |
| 1026 | """Number of redundant experts in this model.""" |
| 1027 | |
| 1028 | moe_layers: Sequence["MoERunner"] |
| 1029 | """List of MoE layers in this model.""" |
| 1030 | |
| 1031 | def set_eplb_state( |
| 1032 | self, |
| 1033 | expert_load_view: Tensor, |
| 1034 | logical_to_physical_map: Tensor, |
| 1035 | logical_replica_count: Tensor, |
| 1036 | ) -> None: |
| 1037 | """Register the EPLB state in the MoE model. |
| 1038 | |
| 1039 | Since these are views of the actual EPLB state, any changes made by |
| 1040 | the EPLB algorithm are automatically reflected in the model's behavior |
| 1041 | without requiring additional method calls to set new states. |
| 1042 | |
| 1043 | You should also collect model's `expert_weights` here instead of in |
| 1044 | the weight loader, since after initial weight loading, further |
| 1045 | processing like quantization may be applied to the weights. |
| 1046 | |
| 1047 | Args: |
| 1048 | expert_load_view: A view of the expert load metrics tensor. |
| 1049 | logical_to_physical_map: Mapping from logical to physical experts. |
| 1050 | logical_replica_count: Count of replicas for each logical expert. |
| 1051 | |
| 1052 | """ |
| 1053 | self.expert_weights = [] |
| 1054 | for layer_idx, layer in enumerate(self.moe_layers): |
| 1055 | # Register the expert weights. |
| 1056 | self.expert_weights.append(list(layer.get_expert_weights())) |
| 1057 | layer.set_eplb_state( |
| 1058 | moe_layer_idx=layer_idx, |
| 1059 | expert_load_view=expert_load_view, |
| 1060 | logical_to_physical_map=logical_to_physical_map, |
| 1061 | logical_replica_count=logical_replica_count, |
| 1062 | ) |
| 1063 | |
| 1064 | def update_physical_experts_metadata( |
| 1065 | self, |
| 1066 | num_physical_experts: int, |
| 1067 | num_local_physical_experts: int, |
| 1068 | ) -> None: ... |
| 1069 | |
| 1070 | |
| 1071 | def get_mixture_of_experts_model(model: object) -> MixtureOfExperts | None: |
| 1072 | """Return the MixtureOfExperts contained within an arbitrary model. |
| 1073 | |
| 1074 | - If the model itself is a MixtureOfExperts, return the model directly. |
| 1075 | - If the model is a multi-modal model, and its `language_model` is a |
| 1076 | MixtureOfExperts, return the `language_model`. |
| 1077 | - If neither, return None. |
| 1078 | |
| 1079 | Args: |
| 1080 | model: Model being served. |
| 1081 | |
| 1082 | Returns: |
| 1083 | The MixtureOfExperts instance contained within the model, or None. |
| 1084 | |
| 1085 | """ |
| 1086 | if is_mixture_of_experts(model): |
| 1087 | return model |
| 1088 | |
| 1089 | if isinstance(model, SupportsMultiModal): |
| 1090 | try: |
| 1091 | mm_language_model = model.get_language_model() |
| 1092 | return ( |
| 1093 | mm_language_model if is_mixture_of_experts(mm_language_model) else None |
| 1094 | ) |
| 1095 | except NotImplementedError: |
| 1096 | logger.info_once("Cannot fetch language_model from MultiModal model") |
| 1097 | return None |
| 1098 | |
| 1099 | return None |
| 1100 | |
| 1101 | |
| 1102 | def is_mixture_of_experts(model: object) -> TypeIs[MixtureOfExperts]: |
| 1103 | return ( |
| 1104 | isinstance(model, MixtureOfExperts) and getattr(model, "num_moe_layers", 0) > 0 |
| 1105 | ) |
| 1106 | |
| 1107 | |
| 1108 | @runtime_checkable |
| 1109 | class HasNoOps(Protocol): |
| 1110 | has_noops: ClassVar[Literal[True]] = True |
| 1111 | |
| 1112 | |
| 1113 | @overload |
| 1114 | def has_noops(model: type[object]) -> TypeIs[type[HasNoOps]]: ... |
| 1115 | |
| 1116 | |
| 1117 | @overload |
| 1118 | def has_noops(model: object) -> TypeIs[HasNoOps]: ... |
| 1119 | |
| 1120 | |
| 1121 | def has_noops( |
| 1122 | model: type[object] | object, |
| 1123 | ) -> TypeIs[type[HasNoOps]] | TypeIs[HasNoOps]: |
| 1124 | return getattr(model, "has_noops", False) |
| 1125 | |
| 1126 | |
| 1127 | @runtime_checkable |
| 1128 | class SupportsMambaPrefixCaching(Protocol): |
| 1129 | """The interface for models whose mamba layers support prefix caching. |
| 1130 | |
| 1131 | This is currently experimental. |
| 1132 | """ |
| 1133 | |
| 1134 | supports_mamba_prefix_caching: ClassVar[Literal[True]] = True |
| 1135 | |
| 1136 | @classmethod |
| 1137 | def get_mamba_state_copy_func(cls) -> tuple["MambaStateCopyFunc", ...]: |
| 1138 | """Return copy functions for the model's Mamba states.""" |
| 1139 | ... |
| 1140 | |
| 1141 | @classmethod |
| 1142 | def get_mamba_state_copy_funcs( |
| 1143 | cls, |
| 1144 | mamba_types: set["MambaAttentionBackendEnum"], |
| 1145 | ) -> "MambaStateCopyFuncsByType": |
| 1146 | """Map legacy copy functions to each requested Mamba backend.""" |
| 1147 | copy_funcs = cls.get_mamba_state_copy_func() |
| 1148 | return {mamba_type: copy_funcs for mamba_type in mamba_types} |
| 1149 | |
| 1150 | |
| 1151 | @overload |
| 1152 | def supports_mamba_prefix_caching( |
| 1153 | model: type[object], |
| 1154 | ) -> TypeIs[type[SupportsMambaPrefixCaching]]: ... |
| 1155 | |
| 1156 | |
| 1157 | @overload |
| 1158 | def supports_mamba_prefix_caching( |
| 1159 | model: object, |
| 1160 | ) -> TypeIs[SupportsMambaPrefixCaching]: ... |
| 1161 | |
| 1162 | |
| 1163 | def supports_mamba_prefix_caching( |
| 1164 | model: type[object] | object, |
| 1165 | ) -> TypeIs[type[SupportsMambaPrefixCaching]] | TypeIs[SupportsMambaPrefixCaching]: |
| 1166 | return getattr(model, "supports_mamba_prefix_caching", False) |
| 1167 | |
| 1168 | |
| 1169 | @runtime_checkable |
| 1170 | class SupportsReplaySSM(Protocol): |
| 1171 | """The interface for models whose recurrent layers support ReplaySSM |
| 1172 | cached decode. |
| 1173 | |
| 1174 | This is currently experimental. |
| 1175 | """ |
| 1176 | |
| 1177 | supports_replayssm: ClassVar[Literal[True]] = True |
| 1178 | |
| 1179 | |
| 1180 | @overload |
| 1181 | def supports_replayssm(model: type[object]) -> TypeIs[type[SupportsReplaySSM]]: ... |
| 1182 | |
| 1183 | |
| 1184 | @overload |
| 1185 | def supports_replayssm(model: object) -> TypeIs[SupportsReplaySSM]: ... |
| 1186 | |
| 1187 | |
| 1188 | def supports_replayssm( |
| 1189 | model: type[object] | object, |
| 1190 | ) -> TypeIs[type[SupportsReplaySSM]] | TypeIs[SupportsReplaySSM]: |
| 1191 | return getattr(model, "supports_replayssm", False) |
| 1192 | |
| 1193 | |
| 1194 | @runtime_checkable |
| 1195 | class SupportsCrossEncoding(Protocol): |
| 1196 | """The interface required for all models that support cross encoding.""" |
| 1197 | |
| 1198 | score_type: ClassVar["ScoreType"] = "cross-encoder" |
| 1199 | |
| 1200 | |
| 1201 | @runtime_checkable |
| 1202 | class SupportsLateInteraction(Protocol): |
| 1203 | """The interface required for all models that support late interaction. |
| 1204 | |
| 1205 | Late interaction models (like ColBERT) encode queries and documents |
| 1206 | separately into per-token embeddings, then compute similarity via |
| 1207 | MaxSim (max over document tokens, sum over query tokens). |
| 1208 | """ |
| 1209 | |
| 1210 | score_type: ClassVar["ScoreType"] = "late-interaction" |
| 1211 | |
| 1212 | |
| 1213 | class SupportsQuant: |
| 1214 | """The interface required for all models that support quantization.""" |
| 1215 | |
| 1216 | hf_to_vllm_mapper: "WeightsMapper | None" = None |
| 1217 | packed_modules_mapping: dict[str, list[str]] |
| 1218 | quant_config: QuantizationConfig | None = None |
| 1219 | |
| 1220 | def __new__(cls, *args, **kwargs) -> Self: |
| 1221 | instance = super().__new__(cls) |
| 1222 | |
| 1223 | # find config passed in arguments and attach it to model for general use |
| 1224 | instance.quant_config = cls._find_quant_config(*args, **kwargs) |
| 1225 | |
| 1226 | cls._maybe_apply_model_mapping(instance) |
| 1227 | |
| 1228 | return instance |
| 1229 | |
| 1230 | @staticmethod |
| 1231 | def _find_quant_config(*args, **kwargs) -> QuantizationConfig | None: |
| 1232 | """Find quant config passed through model constructor args""" |
| 1233 | from vllm.config import VllmConfig # avoid circular import |
| 1234 | |
| 1235 | args_values = list(args) + list(kwargs.values()) |
| 1236 | for arg in args_values: |
| 1237 | if isinstance(arg, VllmConfig): |
| 1238 | return arg.quant_config |
| 1239 | |
| 1240 | if isinstance(arg, QuantizationConfig): |
| 1241 | return arg |
| 1242 | |
| 1243 | return None |
| 1244 | |
| 1245 | def _maybe_apply_model_mapping(self): |
| 1246 | """Apply model mappings to config for proper config-model matching""" |
| 1247 | if self.quant_config is None: |
| 1248 | return |
| 1249 | if (hf_to_vllm_mapper := self.hf_to_vllm_mapper) is not None: |
| 1250 | self.quant_config.apply_vllm_mapper(hf_to_vllm_mapper.get_rename_mapper()) |
| 1251 | if packed_modules_mapping := getattr(self, "packed_modules_mapping", None): |
| 1252 | self.quant_config.packed_modules_mapping.update(packed_modules_mapping) |
| 1253 | |
| 1254 | |
| 1255 | @runtime_checkable |
| 1256 | class SupportsRealtime(Protocol): |
| 1257 | """The interface required for all models that support transcription.""" |
| 1258 | |
| 1259 | supports_realtime: ClassVar[Literal[True]] = True |
| 1260 | |
| 1261 | realtime_max_tokens: ClassVar[int] = 1 |
| 1262 | """Maximum tokens to generate per streaming audio segment. |
| 1263 | Override in subclasses based on the model's expected output length.""" |
| 1264 | |
| 1265 | @classmethod |
| 1266 | def buffer_realtime_audio( |
| 1267 | cls, |
| 1268 | audio_stream: AsyncGenerator[np.ndarray, None], |
| 1269 | input_stream: asyncio.Queue[list[int]], |
| 1270 | model_config: "ModelConfig", |
| 1271 | ) -> AsyncGenerator["PromptType", None]: ... |
| 1272 | |
| 1273 | |
| 1274 | @overload |
| 1275 | def supports_realtime( |
| 1276 | model: type[object], |
| 1277 | ) -> TypeIs[type[SupportsRealtime]]: ... |
| 1278 | |
| 1279 | |
| 1280 | @overload |
| 1281 | def supports_realtime(model: object) -> TypeIs[SupportsRealtime]: ... |
| 1282 | |
| 1283 | |
| 1284 | def supports_realtime( |
| 1285 | model: type[object] | object, |
| 1286 | ) -> TypeIs[type[SupportsRealtime]] | TypeIs[SupportsRealtime]: |
| 1287 | return getattr(model, "supports_realtime", False) |
| 1288 | |
| 1289 | |
| 1290 | @runtime_checkable |
| 1291 | class SupportsTranscription(Protocol): |
| 1292 | """The interface required for all models that support transcription.""" |
| 1293 | |
| 1294 | # Mapping from ISO639_1 language codes: language names |
| 1295 | supported_languages: ClassVar[Mapping[str, str]] |
| 1296 | |
| 1297 | supports_transcription: ClassVar[Literal[True]] = True |
| 1298 | |
| 1299 | supports_transcription_only: ClassVar[bool] = False |
| 1300 | """ |
| 1301 | Transcription models can opt out of text generation by setting this to |
| 1302 | `True`. |
| 1303 | """ |
| 1304 | supports_segment_timestamp: ClassVar[bool] = False |
| 1305 | """ |
| 1306 | Enables the segment timestamp option for supported models by setting this to `True`. |
| 1307 | """ |
| 1308 | |
| 1309 | supports_diarized_transcription: ClassVar[bool] = False |
| 1310 | """Enables the ``diarized_json`` response format for the model.""" |
| 1311 | |
| 1312 | supports_explicit_language_detection: ClassVar[bool] = False |
| 1313 | """ |
| 1314 | Transcription models that require an explicit language detection step |
| 1315 | (e.g. Whisper needs a separate forward pass to predict the language |
| 1316 | token) should set this to ``True`` and implement |
| 1317 | :meth:`get_language_detection_prompt` and |
| 1318 | :meth:`parse_language_detection_output` and |
| 1319 | :meth:`get_language_token_ids`. |
| 1320 | """ |
| 1321 | |
| 1322 | no_space_languages: ClassVar[set[str]] = {"ja", "zh"} |
| 1323 | """ |
| 1324 | Languages that don't need a space between words. |
| 1325 | For example, Japanese (ja) and Chinese (zh) don't need a space between words. |
| 1326 | """ |
| 1327 | |
| 1328 | def __init_subclass__(cls, **kwargs): |
| 1329 | super().__init_subclass__(**kwargs) |
| 1330 | # language codes in supported_languages |
| 1331 | # that don't exist in the full language map |
| 1332 | invalid = set(cls.supported_languages) - set(LANGUAGES.keys()) |
| 1333 | if invalid: |
| 1334 | raise ValueError( |
| 1335 | f"{cls.__name__}.supported_languages contains invalid " |
| 1336 | f"language codes: {sorted(invalid)}\n. " |
| 1337 | f"Valid choices are: {sorted(LANGUAGES.keys())}" |
| 1338 | ) |
| 1339 | |
| 1340 | @classmethod |
| 1341 | def get_generation_prompt( |
| 1342 | cls, |
| 1343 | stt_params: "SpeechToTextParams", |
| 1344 | ) -> "PromptType": |
| 1345 | """Get the prompt for the ASR model. |
| 1346 | The model has control over the construction, as long as it |
| 1347 | returns a valid PromptType.""" |
| 1348 | ... |
| 1349 | |
| 1350 | @classmethod |
| 1351 | def get_other_languages(cls) -> Mapping[str, str]: |
| 1352 | # other possible language codes from the whisper map |
| 1353 | return {k: v for k, v in LANGUAGES.items() if k not in cls.supported_languages} |
| 1354 | |
| 1355 | @classmethod |
| 1356 | def validate_language(cls, language: str | None) -> str | None: |
| 1357 | """Ensure the language specified in the transcription request |
| 1358 | is a valid ISO 639-1 language code. If the request language is |
| 1359 | valid, but not natively supported by the model, trigger a |
| 1360 | warning (but not an exception). |
| 1361 | """ |
| 1362 | if language is None or language in cls.supported_languages: |
| 1363 | return language |
| 1364 | elif language in cls.get_other_languages(): |
| 1365 | logger.warning( |
| 1366 | "Language %r is not natively supported by %s; " |
| 1367 | "results may be less accurate. Supported languages: %r", |
| 1368 | language, |
| 1369 | cls.__name__, |
| 1370 | list(cls.supported_languages.keys()), |
| 1371 | ) |
| 1372 | return language |
| 1373 | else: |
| 1374 | raise ValueError( |
| 1375 | f"Unsupported language: {language!r}. Must be one of " |
| 1376 | f"{list(cls.supported_languages.keys())}." |
| 1377 | ) |
| 1378 | |
| 1379 | @classmethod |
| 1380 | def get_speech_to_text_config( |
| 1381 | cls, model_config: "ModelConfig", task_type: Literal["transcribe", "translate"] |
| 1382 | ) -> "SpeechToTextConfig": |
| 1383 | """Get the speech to text config for the ASR model.""" |
| 1384 | ... |
| 1385 | |
| 1386 | @classmethod |
| 1387 | def get_num_audio_tokens( |
| 1388 | cls, |
| 1389 | audio_duration_s: float, |
| 1390 | stt_config: "SpeechToTextConfig", |
| 1391 | model_config: "ModelConfig", |
| 1392 | ) -> int | None: |
| 1393 | """Map from audio duration to number of audio tokens produced by the ASR |
| 1394 | model, without running a forward pass. |
| 1395 | This is used for estimating the amount of processing for this audio. |
| 1396 | """ |
| 1397 | return None |
| 1398 | |
| 1399 | @classmethod |
| 1400 | def post_process_output(cls, text: str) -> str: |
| 1401 | """Post-process the raw model output text. |
| 1402 | |
| 1403 | Some ASR models output structured formats (e.g., language tags, |
| 1404 | special tokens) that need to be stripped before returning to the user. |
| 1405 | |
| 1406 | Args: |
| 1407 | text: Raw decoded text from the model. |
| 1408 | |
| 1409 | Returns: |
| 1410 | Cleaned transcription text. |
| 1411 | |
| 1412 | """ |
| 1413 | return text |
| 1414 | |
| 1415 | @classmethod |
| 1416 | def parse_diarized_transcript(cls, text: str) -> list[DiarizedTranscriptionSegment]: |
| 1417 | """Parse the model-specific diarized transcript format. |
| 1418 | |
| 1419 | Only models that set ``supports_diarized_transcription`` must override |
| 1420 | this method. |
| 1421 | """ |
| 1422 | raise NotImplementedError |
| 1423 | |
| 1424 | @classmethod |
| 1425 | def get_streaming_post_processor_cls( |
| 1426 | cls, |
| 1427 | ) -> type[StreamingTranscriptionPostProcessor]: |
| 1428 | """Return a stateful post-processor class for streaming output deltas. |
| 1429 | |
| 1430 | Each instance receives the next decoded text delta and whether the |
| 1431 | request output is final. It returns the cleaned delta that should be |
| 1432 | sent to the client. |
| 1433 | """ |
| 1434 | return StreamingTranscriptionPostProcessor |
| 1435 | |
| 1436 | @classmethod |
| 1437 | def get_language_detection_prompt( |
| 1438 | cls, |
| 1439 | audio: np.ndarray, |
| 1440 | stt_config: "SpeechToTextConfig", |
| 1441 | ) -> "PromptType": |
| 1442 | """Return a prompt that triggers language detection. |
| 1443 | |
| 1444 | Only needs to be implemented when |
| 1445 | ``supports_explicit_language_detection`` is ``True``. |
| 1446 | """ |
| 1447 | raise NotImplementedError |
| 1448 | |
| 1449 | @classmethod |
| 1450 | def parse_language_detection_output( |
| 1451 | cls, |
| 1452 | token_ids: list[int], |
| 1453 | tokenizer: "TokenizerLike", |
| 1454 | ) -> str: |
| 1455 | """Parse the detected language from model output token IDs. |
| 1456 | |
| 1457 | Only needs to be implemented when |
| 1458 | ``supports_explicit_language_detection`` is ``True``. |
| 1459 | """ |
| 1460 | raise NotImplementedError |
| 1461 | |
| 1462 | @classmethod |
| 1463 | def get_language_token_ids( |
| 1464 | cls, |
| 1465 | tokenizer: "TokenizerLike", |
| 1466 | ) -> list[int] | None: |
| 1467 | """Return token IDs that represent valid language tokens. |
| 1468 | |
| 1469 | Used to constrain language detection to only produce valid language tokens. |
| 1470 | |
| 1471 | Only needs to be implemented when |
| 1472 | ``supports_explicit_language_detection`` is ``True``. |
| 1473 | """ |
| 1474 | raise NotImplementedError |
| 1475 | |
| 1476 | |
| 1477 | @overload |
| 1478 | def supports_transcription( |
| 1479 | model: type[object], |
| 1480 | ) -> TypeIs[type[SupportsTranscription]]: ... |
| 1481 | |
| 1482 | |
| 1483 | @overload |
| 1484 | def supports_transcription(model: object) -> TypeIs[SupportsTranscription]: ... |
| 1485 | |
| 1486 | |
| 1487 | def supports_transcription( |
| 1488 | model: type[object] | object, |
| 1489 | ) -> TypeIs[type[SupportsTranscription]] | TypeIs[SupportsTranscription]: |
| 1490 | return getattr(model, "supports_transcription", False) |
| 1491 | |
| 1492 | |
| 1493 | @runtime_checkable |
| 1494 | class SupportsEagleBase(Protocol): |
| 1495 | """Base interface for models that support EAGLE-based speculative decoding.""" |
| 1496 | |
| 1497 | has_own_lm_head: bool = False |
| 1498 | """ |
| 1499 | A flag that indicates this model has trained its own lm_head. |
| 1500 | """ |
| 1501 | |
| 1502 | has_own_embed_tokens: bool = False |
| 1503 | """ |
| 1504 | A flag that indicates this model has trained its own input embeddings. |
| 1505 | """ |
| 1506 | |
| 1507 | |
| 1508 | @overload |
| 1509 | def supports_any_eagle(model: type[object]) -> TypeIs[type[SupportsEagleBase]]: ... |
| 1510 | |
| 1511 | |
| 1512 | @overload |
| 1513 | def supports_any_eagle(model: object) -> TypeIs[SupportsEagleBase]: ... |
| 1514 | |
| 1515 | |
| 1516 | def supports_any_eagle( |
| 1517 | model: type[object] | object, |
| 1518 | ) -> TypeIs[type[SupportsEagleBase]] | TypeIs[SupportsEagleBase]: |
| 1519 | """Check if model supports any EAGLE variant (1, 2, or 3).""" |
| 1520 | return supports_eagle(model) or supports_eagle3(model) |
| 1521 | |
| 1522 | |
| 1523 | class LocalArgmaxMixin: |
| 1524 | """Mixin for draft model heads in speculative decoding. |
| 1525 | |
| 1526 | Provides a D2T-aware ``get_top_tokens`` that preserves the |
| 1527 | local-argmax communication reduction even when the draft vocabulary |
| 1528 | is smaller than the target vocabulary. |
| 1529 | |
| 1530 | When ``draft_id_to_target_id`` is present (shape ``(draft_vocab_size,)``, |
| 1531 | containing per-token offset to target vocab id), the draft argmax index |
| 1532 | ``k`` is mapped to the target vocab id via:: |
| 1533 | |
| 1534 | target_id = k + draft_id_to_target_id[k] |
| 1535 | |
| 1536 | This is mathematically equivalent to computing the full-vocab scatter |
| 1537 | logits and taking the global argmax, but requires only |
| 1538 | O(batch * 2 * tp_size) communication instead of O(batch * vocab_size). |
| 1539 | |
| 1540 | Requires the subclass to expose: |
| 1541 | ``self.logits_processor``: LogitsProcessor |
| 1542 | ``self.lm_head``: ParallelLMHead |
| 1543 | ``self.draft_id_to_target_id`` (optional): nn.Parameter |
| 1544 | """ |
| 1545 | |
| 1546 | logits_processor: "LogitsProcessor" |
| 1547 | lm_head: "ParallelLMHead" |
| 1548 | |
| 1549 | def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 1550 | """Vocab-parallel argmax with optional D2T remapping.""" |
| 1551 | top = self.logits_processor.get_top_tokens( |
| 1552 | self.lm_head, |
| 1553 | hidden_states, |
| 1554 | ) |
| 1555 | d2t = getattr(self, "draft_id_to_target_id", None) |
| 1556 | if d2t is not None: |
| 1557 | top = top + d2t[top] |
| 1558 | return top |
| 1559 | |
| 1560 | |
| 1561 | class EagleModelMixin: |
| 1562 | start_layer: int |
| 1563 | aux_hidden_state_layers: tuple[int, ...] = () |
| 1564 | supports_aux_hidden_states_over_pp: ClassVar[bool] = False |
| 1565 | AUX_HIDDEN_STATE_KEY: ClassVar[str] = "aux_hidden_states_" |
| 1566 | _aux_slot_base_cached: int = 0 |
| 1567 | _aux_upstream_total_cached: int = 0 |
| 1568 | |
| 1569 | def _set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: |
| 1570 | self.aux_hidden_state_layers = tuple(sorted(set(layers))) |
| 1571 | self._aux_slot_base_cached = 0 |
| 1572 | self._aux_upstream_total_cached = 0 |
| 1573 | self._cache_aux_pp_layout() |
| 1574 | |
| 1575 | def _cache_aux_pp_layout(self) -> None: |
| 1576 | from vllm.distributed.parallel_state import ( |
| 1577 | get_pp_group, |
| 1578 | model_parallel_is_initialized, |
| 1579 | ) |
| 1580 | |
| 1581 | if not model_parallel_is_initialized(): |
| 1582 | return |
| 1583 | pp = get_pp_group() |
| 1584 | if pp.world_size < 2: |
| 1585 | return |
| 1586 | if not pp.is_first_rank: |
| 1587 | self._aux_slot_base_cached = bisect_right( |
| 1588 | self.aux_hidden_state_layers, self.start_layer |
| 1589 | ) |
| 1590 | if pp.is_last_rank: |
| 1591 | self._aux_upstream_total_cached = self._aux_slot_base_cached |
| 1592 | |
| 1593 | def _maybe_add_hidden_state( |
| 1594 | self, |
| 1595 | aux_hidden_states: list[torch.Tensor], |
| 1596 | layer_idx: int, |
| 1597 | hidden_states: torch.Tensor, |
| 1598 | residual: torch.Tensor | None, |
| 1599 | ) -> list[torch.Tensor]: |
| 1600 | if layer_idx in self.aux_hidden_state_layers: |
| 1601 | value = hidden_states + residual if residual is not None else hidden_states |
| 1602 | aux_hidden_states.append(value) |
| 1603 | return aux_hidden_states |
| 1604 | |
| 1605 | def pack_local_aux_hidden_states( |
| 1606 | self, aux_hidden_states: list[torch.Tensor] |
| 1607 | ) -> dict[str, torch.Tensor]: |
| 1608 | if not aux_hidden_states: |
| 1609 | return {} |
| 1610 | base = self._aux_slot_base_cached |
| 1611 | return { |
| 1612 | f"{self.AUX_HIDDEN_STATE_KEY}{base + i}": t |
| 1613 | for i, t in enumerate(aux_hidden_states) |
| 1614 | } |
| 1615 | |
| 1616 | def collect_remote_aux_hidden_states( |
| 1617 | self, intermediate_tensors: "IntermediateTensors | None" |
| 1618 | ) -> list[torch.Tensor]: |
| 1619 | total = self._aux_upstream_total_cached |
| 1620 | if total == 0: |
| 1621 | return [] |
| 1622 | |
| 1623 | assert intermediate_tensors is not None |
| 1624 | out: list[torch.Tensor] = [] |
| 1625 | for i in range(total): |
| 1626 | key = f"{self.AUX_HIDDEN_STATE_KEY}{i}" |
| 1627 | if key not in intermediate_tensors.tensors: |
| 1628 | raise RuntimeError( |
| 1629 | f"Missing {key} from PP intermediate tensors: " |
| 1630 | f"{sorted(intermediate_tensors.tensors)}" |
| 1631 | ) |
| 1632 | out.append(intermediate_tensors[key]) |
| 1633 | return out |
| 1634 | |
| 1635 | |
| 1636 | @runtime_checkable |
| 1637 | class SupportsEagle(SupportsEagleBase, Protocol): |
| 1638 | """The interface required for models that support |
| 1639 | EAGLE-1 and EAGLE-2 speculative decoding.""" |
| 1640 | |
| 1641 | supports_eagle: ClassVar[Literal[True]] = True |
| 1642 | """ |
| 1643 | A flag that indicates this model supports EAGLE-1 and EAGLE-2 |
| 1644 | speculative decoding. |
| 1645 | |
| 1646 | Note: |
| 1647 | There is no need to redefine this flag if this class is in the |
| 1648 | MRO of your model class. |
| 1649 | """ |
| 1650 | |
| 1651 | |
| 1652 | @overload |
| 1653 | def supports_eagle(model: type[object]) -> TypeIs[type[SupportsEagle]]: ... |
| 1654 | |
| 1655 | |
| 1656 | @overload |
| 1657 | def supports_eagle(model: object) -> TypeIs[SupportsEagle]: ... |
| 1658 | |
| 1659 | |
| 1660 | def supports_eagle( |
| 1661 | model: type[object] | object, |
| 1662 | ) -> TypeIs[type[SupportsEagle]] | TypeIs[SupportsEagle]: |
| 1663 | return isinstance(model, SupportsEagle) |
| 1664 | |
| 1665 | |
| 1666 | @runtime_checkable |
| 1667 | class SupportsEagle3(SupportsEagleBase, Protocol): |
| 1668 | """The interface required for models that support |
| 1669 | EAGLE-3 speculative decoding.""" |
| 1670 | |
| 1671 | supports_eagle3: ClassVar[Literal[True]] = True |
| 1672 | """ |
| 1673 | A flag that indicates this model supports EAGLE-3 |
| 1674 | speculative decoding. |
| 1675 | |
| 1676 | Note: |
| 1677 | There is no need to redefine this flag if this class is in the |
| 1678 | MRO of your model class. |
| 1679 | """ |
| 1680 | |
| 1681 | def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: |
| 1682 | """Set which layers should output auxiliary hidden states for EAGLE-3. |
| 1683 | |
| 1684 | Args: |
| 1685 | layers: Tuple of layer indices that should output auxiliary |
| 1686 | hidden states. |
| 1687 | |
| 1688 | """ |
| 1689 | parent_ref = self |
| 1690 | if hasattr(self, "get_language_model"): |
| 1691 | parent_ref = self.get_language_model() |
| 1692 | elif hasattr(self, "language_model"): |
| 1693 | parent_ref = self.language_model |
| 1694 | # A multimodal model that builds its decoder inside |
| 1695 | # `_mark_language_model` has get_language_model() return the inner |
| 1696 | # decoder itself, which IS the EagleModelMixin and has no further |
| 1697 | # `.model`. Unwrap only when there is something to unwrap. |
| 1698 | holder = getattr(parent_ref, "model", parent_ref) |
| 1699 | assert isinstance(holder, EagleModelMixin), ( |
| 1700 | "Model instance must inherit from EagleModelMixin to set auxiliary layers" |
| 1701 | ) |
| 1702 | holder._set_aux_hidden_state_layers(layers) |
| 1703 | |
| 1704 | def get_eagle3_default_aux_hidden_state_layers(self) -> tuple[int, ...]: |
| 1705 | """Get the default layer indices that should output auxiliary hidden states |
| 1706 | for EAGLE-3 for this model. Models can override this method to provide |
| 1707 | different default layers based on their architecture, but it is encouraged |
| 1708 | to instead include the layer specification in the model's config if possible. |
| 1709 | |
| 1710 | Returns: |
| 1711 | Tuple of layer indices for auxiliary hidden state outputs. |
| 1712 | |
| 1713 | """ |
| 1714 | parent_ref = self |
| 1715 | if hasattr(self, "get_language_model"): |
| 1716 | parent_ref = self.get_language_model() |
| 1717 | elif hasattr(self, "language_model"): |
| 1718 | parent_ref = self.language_model |
| 1719 | # Same unwrap-only-if-needed rule as set_aux_hidden_state_layers. |
| 1720 | holder = getattr(parent_ref, "model", parent_ref) |
| 1721 | assert hasattr(holder, "layers"), ( |
| 1722 | "Model instance must have 'layers' attribute to get number of layers" |
| 1723 | ) |
| 1724 | num_layers = len(holder.layers) |
| 1725 | return (2, num_layers // 2, num_layers - 3) |
| 1726 | |
| 1727 | |
| 1728 | @overload |
| 1729 | def supports_eagle3(model: type[object]) -> TypeIs[type[SupportsEagle3]]: ... |
| 1730 | |
| 1731 | |
| 1732 | @overload |
| 1733 | def supports_eagle3(model: object) -> TypeIs[SupportsEagle3]: ... |
| 1734 | |
| 1735 | |
| 1736 | def supports_eagle3( |
| 1737 | model: type[object] | object, |
| 1738 | ) -> TypeIs[type[SupportsEagle3]] | TypeIs[SupportsEagle3]: |
| 1739 | return isinstance(model, SupportsEagle3) |
| 1740 | |
| 1741 | |
| 1742 | @runtime_checkable |
| 1743 | class SupportsMRoPE(Protocol): |
| 1744 | """The interface required for all models that support M-RoPE.""" |
| 1745 | |
| 1746 | supports_mrope: ClassVar[Literal[True]] = True |
| 1747 | """ |
| 1748 | A flag that indicates this model supports M-RoPE. |
| 1749 | |
| 1750 | Note: |
| 1751 | There is no need to redefine this flag if this class is in the |
| 1752 | MRO of your model class. |
| 1753 | """ |
| 1754 | |
| 1755 | def get_mrope_input_positions( |
| 1756 | self, |
| 1757 | input_tokens: list[int], |
| 1758 | mm_features: list["MultiModalFeatureSpec"], |
| 1759 | ) -> tuple[torch.Tensor, int]: |
| 1760 | """Get M-RoPE input positions and delta value for this specific model. |
| 1761 | |
| 1762 | This method should be implemented by each model that supports M-RoPE |
| 1763 | to provide model-specific logic for computing input positions. |
| 1764 | |
| 1765 | Args: |
| 1766 | input_tokens: List of input token IDs |
| 1767 | mm_features: Information about each multi-modal data item |
| 1768 | |
| 1769 | Returns: |
| 1770 | llm_positions: Tensor of shape `[num_dims, num_tokens]`, one row |
| 1771 | per M-RoPE position channel (e.g. T/H/W). |
| 1772 | mrope_position_delta: Delta for position calculations. |
| 1773 | |
| 1774 | """ |
| 1775 | ... |
| 1776 | |
| 1777 | |
| 1778 | @overload |
| 1779 | def supports_mrope(model: type[object]) -> TypeIs[type[SupportsMRoPE]]: ... |
| 1780 | |
| 1781 | |
| 1782 | @overload |
| 1783 | def supports_mrope(model: object) -> TypeIs[SupportsMRoPE]: ... |
| 1784 | |
| 1785 | |
| 1786 | def supports_mrope( |
| 1787 | model: type[object] | object, |
| 1788 | ) -> TypeIs[type[SupportsMRoPE]] | TypeIs[SupportsMRoPE]: |
| 1789 | return isinstance(model, SupportsMRoPE) |
| 1790 | |
| 1791 | |
| 1792 | @runtime_checkable |
| 1793 | class SupportsEncoderCudaGraph(Protocol): |
| 1794 | """Interface for models whose vision encoder supports CUDA graph |
| 1795 | capture/replay. |
| 1796 | |
| 1797 | Models implement these methods to provide the |
| 1798 | :class:`EncoderCudaGraphManager` with all model-specific logic |
| 1799 | (input handling, metadata computation, forward pass) without the |
| 1800 | manager needing to know model internals. |
| 1801 | """ |
| 1802 | |
| 1803 | supports_encoder_cudagraph: ClassVar[Literal[True]] = True |
| 1804 | |
| 1805 | def get_encoder_cudagraph_config(self) -> "EncoderCudaGraphConfig": ... |
| 1806 | |
| 1807 | def get_input_modality( |
| 1808 | self, |
| 1809 | mm_kwargs: dict[str, Any], |
| 1810 | ) -> str: |
| 1811 | """Return the modality of the inputs (default: image-only).""" |
| 1812 | return "image" |
| 1813 | |
| 1814 | def get_max_frames_per_video( |
| 1815 | self, |
| 1816 | ) -> int: |
| 1817 | """Return model-specific max frames per video.""" |
| 1818 | ... |
| 1819 | |
| 1820 | def get_encoder_cudagraph_budget_range( |
| 1821 | self, |
| 1822 | vllm_config: "VllmConfig", |
| 1823 | ) -> tuple[int, int]: |
| 1824 | """Return (min_token_budget, max_token_budget) for auto-inference. |
| 1825 | |
| 1826 | - min_token_budget: estimated smallest possible encoder input |
| 1827 | (e.g. 64 for a 224x224 image) |
| 1828 | - max_token_budget: estimated largest budget worth capturing |
| 1829 | (e.g. max_num_batched_tokens) |
| 1830 | |
| 1831 | Used when ``encoder_cudagraph_token_budgets`` and/or |
| 1832 | ``encoder_cudagraph_max_vision_items_per_batch`` are not explicitly |
| 1833 | specified by the user. |
| 1834 | """ |
| 1835 | ... |
| 1836 | |
| 1837 | def get_encoder_cudagraph_item_specs( |
| 1838 | self, |
| 1839 | mm_kwargs: dict[str, Any], |
| 1840 | ) -> list["EncoderItemSpec"]: |
| 1841 | """Return specs describing each item in the batch. |
| 1842 | |
| 1843 | Replaces the former separate methods for num_items, |
| 1844 | per_item_output_tokens, and per_item_input_sizes. |
| 1845 | The manager derives all three from this single return value. |
| 1846 | """ |
| 1847 | ... |
| 1848 | |
| 1849 | def select_encoder_cudagraph_items( |
| 1850 | self, |
| 1851 | mm_kwargs: dict[str, Any], |
| 1852 | indices: list[int], |
| 1853 | ) -> dict[str, Any]: |
| 1854 | """Select a subset of items and return mm_kwargs for the sub-batch. |
| 1855 | |
| 1856 | Called by the manager during greedy packing and DP sharding to |
| 1857 | extract inputs for a specific set of items (e.g. images at |
| 1858 | indices [0, 3, 5]). The implementation is model-specific |
| 1859 | because input formats differ: |
| 1860 | |
| 1861 | - Qwen-family: slice concatenated pixel_values by cumulative |
| 1862 | patch offsets, subset grid_thw by indices. |
| 1863 | - Batched models (CLIP): index pixel_values along dim 0. |
| 1864 | |
| 1865 | Models that configure ``EncoderCudaGraphConfig.capture_axes`` must |
| 1866 | additionally store the resolved per-axis keys (one key per axis, in |
| 1867 | order) under ``ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG`` in the returned |
| 1868 | dict; the manager pops it before the kwargs are used elsewhere. |
| 1869 | """ |
| 1870 | ... |
| 1871 | |
| 1872 | def postprocess_encoder_output( |
| 1873 | self, |
| 1874 | outputs: dict[str, torch.Tensor], |
| 1875 | indices: list[int], |
| 1876 | per_item_out_tokens: list[int], |
| 1877 | dest: dict[int, torch.Tensor] | list[torch.Tensor | None], |
| 1878 | clone: bool = False, |
| 1879 | batch_mm_kwargs: dict[str, Any] | None = None, |
| 1880 | ) -> None: |
| 1881 | """Post-process encoder output, directly call scatter_output_slices by default. |
| 1882 | |
| 1883 | By default, delegates directly to scatter_output_slices. |
| 1884 | Override this for models that require additional processing on the raw |
| 1885 | encoder output prior to scattering, e.g. Step3-VL, which merges features |
| 1886 | according to dynamic patch counts before scattering. |
| 1887 | """ |
| 1888 | from vllm.model_executor.models.utils import scatter_output_slices |
| 1889 | |
| 1890 | scatter_output_slices( |
| 1891 | outputs["default"], indices, per_item_out_tokens, dest, clone |
| 1892 | ) |
| 1893 | |
| 1894 | def prepare_encoder_cudagraph_capture_inputs( |
| 1895 | self, |
| 1896 | token_budget: int, |
| 1897 | max_batch_size: int, |
| 1898 | max_frames_per_batch: int, |
| 1899 | device: torch.device, |
| 1900 | dtype: torch.dtype, |
| 1901 | path: str = "default", |
| 1902 | axis_keys: tuple[Hashable, ...] | None = None, |
| 1903 | ) -> "EncoderCudaGraphCaptureInputs": |
| 1904 | """Create dummy inputs and buffers for CUDA graph capture. |
| 1905 | |
| 1906 | Args: |
| 1907 | token_budget: Token budget the capture is sized for. |
| 1908 | max_batch_size: Maximum number of items in a captured batch. |
| 1909 | max_frames_per_batch: Maximum number of frames in a captured batch. |
| 1910 | device: Device the dummy inputs and buffers are created on. |
| 1911 | dtype: Dtype of the dummy inputs and buffers. |
| 1912 | path: Configured encoder path. |
| 1913 | axis_keys: The resolved capture-axis keys (one per axis of |
| 1914 | ``EncoderCudaGraphConfig.capture_axes``) this capture is for. |
| 1915 | None or empty when no capture axes are configured; models |
| 1916 | without capture axes ignore it. |
| 1917 | |
| 1918 | """ |
| 1919 | ... |
| 1920 | |
| 1921 | def prepare_encoder_cudagraph_replay_buffers( |
| 1922 | self, |
| 1923 | mm_kwargs: dict[str, Any], |
| 1924 | max_batch_size: int, |
| 1925 | max_frames_per_batch: int, |
| 1926 | path: str = "default", |
| 1927 | ) -> "EncoderCudaGraphReplayBuffers": |
| 1928 | """Compute buffer values from actual batch inputs for replay.""" |
| 1929 | ... |
| 1930 | |
| 1931 | def encoder_cudagraph_forward( |
| 1932 | self, |
| 1933 | inputs: dict[str, torch.Tensor], |
| 1934 | path: str = "default", |
| 1935 | ) -> torch.Tensor: |
| 1936 | """Run the encoder forward pass with precomputed buffers. |
| 1937 | |
| 1938 | Used during both CUDA graph capture and replay. |
| 1939 | """ |
| 1940 | ... |
| 1941 | |
| 1942 | def encoder_eager_forward( |
| 1943 | self, |
| 1944 | mm_kwargs: dict[str, Any], |
| 1945 | path: str = "default", |
| 1946 | ) -> torch.Tensor: |
| 1947 | """Run the encoder forward pass without precomputed buffers. |
| 1948 | |
| 1949 | Used as eager fallback when inputs exceed all budgets. |
| 1950 | """ |
| 1951 | ... |
| 1952 | |
| 1953 | |
| 1954 | @overload |
| 1955 | def supports_encoder_cudagraph( |
| 1956 | model: type[object], |
| 1957 | ) -> TypeIs[type[SupportsEncoderCudaGraph]]: ... |
| 1958 | |
| 1959 | |
| 1960 | @overload |
| 1961 | def supports_encoder_cudagraph( |
| 1962 | model: object, |
| 1963 | ) -> TypeIs[SupportsEncoderCudaGraph]: ... |
| 1964 | |
| 1965 | |
| 1966 | def supports_encoder_cudagraph( |
| 1967 | model: type[object] | object, |
| 1968 | ) -> TypeIs[type[SupportsEncoderCudaGraph]] | TypeIs[SupportsEncoderCudaGraph]: |
| 1969 | return isinstance(model, SupportsEncoderCudaGraph) |