| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from abc import ABC, abstractmethod |
| 4 | from typing import Any, ClassVar, cast |
| 5 | |
| 6 | import torch |
| 7 | import torch.nn as nn |
| 8 | |
| 9 | from vllm.config import VllmConfig |
| 10 | from vllm.config.compilation import CUDAGraphMode |
| 11 | from vllm.model_executor.models.interfaces import ( |
| 12 | SupportsEncoderCudaGraph, |
| 13 | supports_encoder_cudagraph, |
| 14 | ) |
| 15 | from vllm.tasks import GenerationTask |
| 16 | from vllm.v1.attention.backend import AttentionCGSupport |
| 17 | from vllm.v1.core.sched.output import NewRequestData |
| 18 | from vllm.v1.kv_cache_interface import KVCacheConfig |
| 19 | from vllm.v1.worker.encoder_cudagraph import EncoderCudaGraphManager |
| 20 | from vllm.v1.worker.gpu.input_batch import InputBatch |
| 21 | from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache |
| 22 | from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner |
| 23 | from vllm.v1.worker.gpu.states import RequestState |
| 24 | from vllm.v1.worker.utils import AttentionGroup |
| 25 | |
| 26 | |
| 27 | class ModelSpecificAttnMetadata: |
| 28 | """Base class for model-specific attention metadata.""" |
| 29 | |
| 30 | def get_extra_common_attn_kwargs( |
| 31 | self, kv_cache_group_id: int, num_reqs: int |
| 32 | ) -> dict[str, Any]: |
| 33 | return {} |
| 34 | |
| 35 | def get_extra_attn_kwargs( |
| 36 | self, attn_metadata_builder: Any, num_reqs: int |
| 37 | ) -> dict[str, Any]: |
| 38 | return {} |
| 39 | |
| 40 | |
| 41 | class ModelState(ABC): |
| 42 | supports_prompt_embeds: ClassVar[bool] = False |
| 43 | """Whether this state implements user-provided prompt embeddings.""" |
| 44 | |
| 45 | def __init__( |
| 46 | self, |
| 47 | vllm_config: VllmConfig, |
| 48 | model: nn.Module, |
| 49 | encoder_cache: EncoderCache | None, |
| 50 | device: torch.device, |
| 51 | ) -> None: |
| 52 | self.vllm_config = vllm_config |
| 53 | self.model_config = vllm_config.model_config |
| 54 | self.scheduler_config = vllm_config.scheduler_config |
| 55 | self.model = model |
| 56 | self.device = device |
| 57 | |
| 58 | self.max_num_reqs = self.scheduler_config.max_num_seqs |
| 59 | self.max_num_tokens = self.scheduler_config.max_num_batched_tokens |
| 60 | self.inputs_embeds_size = self.model_config.get_inputs_embeds_size() |
| 61 | self.dtype = self.model_config.dtype |
| 62 | |
| 63 | self.supports_mm_inputs = encoder_cache is not None |
| 64 | if encoder_cache is not None: |
| 65 | enable_encoder_cuda_graph = ( |
| 66 | not self.model_config.enforce_eager |
| 67 | and vllm_config.compilation_config.cudagraph_mm_encoder |
| 68 | and supports_encoder_cudagraph(model) |
| 69 | ) |
| 70 | cudagraph_manager = ( |
| 71 | EncoderCudaGraphManager( |
| 72 | vllm_config=vllm_config, |
| 73 | device=device, |
| 74 | dtype=self.dtype, |
| 75 | model=cast(SupportsEncoderCudaGraph, model), |
| 76 | ) |
| 77 | if enable_encoder_cuda_graph |
| 78 | else None |
| 79 | ) |
| 80 | |
| 81 | self.encoder_cache = encoder_cache |
| 82 | observability_config = vllm_config.observability_config |
| 83 | self.encoder_runner = EncoderRunner( |
| 84 | model=self.model, |
| 85 | max_num_tokens=self.max_num_tokens, |
| 86 | hidden_size=self.inputs_embeds_size, |
| 87 | encoder_cache=encoder_cache, |
| 88 | dtype=self.dtype, |
| 89 | device=self.device, |
| 90 | cudagraph_manager=cudagraph_manager, |
| 91 | enable_timing=bool( |
| 92 | observability_config |
| 93 | and observability_config.enable_mm_processor_stats |
| 94 | ), |
| 95 | ) |
| 96 | |
| 97 | @property |
| 98 | def max_model_len(self) -> int: |
| 99 | # Auto-fit can reduce the limit after model state initialization. |
| 100 | return self.model_config.max_model_len |
| 101 | |
| 102 | def get_supported_generation_tasks(self) -> tuple[GenerationTask, ...]: |
| 103 | from vllm.model_executor.models.interfaces import ( |
| 104 | supports_realtime, |
| 105 | supports_transcription, |
| 106 | ) |
| 107 | from vllm.model_executor.models.interfaces_base import is_text_generation_model |
| 108 | |
| 109 | supported_tasks = list[GenerationTask]() |
| 110 | if is_text_generation_model(self.model): |
| 111 | supported_tasks.append("generate") |
| 112 | if supports_transcription(self.model): |
| 113 | if self.model.supports_transcription_only: |
| 114 | return ("transcription",) |
| 115 | supported_tasks.append("transcription") |
| 116 | if supports_realtime(self.model): |
| 117 | supported_tasks.append("realtime") |
| 118 | return tuple(supported_tasks) |
| 119 | |
| 120 | def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: |
| 121 | return None |
| 122 | |
| 123 | def remove_request(self, req_id: str) -> None: |
| 124 | return None |
| 125 | |
| 126 | def apply_staged_writes(self) -> None: |
| 127 | return None |
| 128 | |
| 129 | def get_additional_cg_support(self) -> tuple[AttentionCGSupport, str | None]: |
| 130 | """Cudagraph support of attention groups this ModelState builds outside |
| 131 | ``init_attn_backend`` (e.g. encoder-only layers). |
| 132 | |
| 133 | Returns the minimum support level and its backend name. The default of |
| 134 | ``ALWAYS`` imposes no extra constraint on the runner's cudagraph mode. |
| 135 | """ |
| 136 | return AttentionCGSupport.ALWAYS, None |
| 137 | |
| 138 | def preprocess_state( |
| 139 | self, |
| 140 | input_batch: InputBatch, |
| 141 | block_tables: tuple[torch.Tensor, ...], |
| 142 | kv_cache_config: KVCacheConfig, |
| 143 | num_computed_tokens: torch.Tensor, |
| 144 | ) -> None: |
| 145 | """Hook run on real batches before the forward pass (after block tables |
| 146 | are gathered). Used by mamba "align" prefix caching to pre-copy state |
| 147 | across block boundaries. No-op by default.""" |
| 148 | return None |
| 149 | |
| 150 | def postprocess_state( |
| 151 | self, |
| 152 | idx_mapping: torch.Tensor, |
| 153 | num_sampled: torch.Tensor, |
| 154 | num_computed_tokens: torch.Tensor | None = None, |
| 155 | ) -> None: |
| 156 | return None |
| 157 | |
| 158 | @abstractmethod |
| 159 | def prepare_inputs_embeds( |
| 160 | self, |
| 161 | scheduled_encoder_inputs: dict[str, list[int]], |
| 162 | input_batch: InputBatch, |
| 163 | req_states: RequestState, |
| 164 | ) -> torch.Tensor | None: |
| 165 | """Prepare the ``inputs_embeds`` tensor for the current forward pass.""" |
| 166 | raise NotImplementedError |
| 167 | |
| 168 | def dummy_inputs_embeds(self, num_tokens: int) -> torch.Tensor | None: |
| 169 | """Pre-allocated inputs_embeds buffer for dummy runs (contents unused).""" |
| 170 | return None |
| 171 | |
| 172 | def execute_mm_encoder( |
| 173 | self, scheduled_encoder_inputs: dict[str, list[int]] |
| 174 | ) -> None: |
| 175 | """Run the multi-modal encoder and cache its outputs by `mm_hash`. |
| 176 | |
| 177 | The encode half of `get_mm_embeddings`, without the gather, for callers |
| 178 | that run no language model. |
| 179 | """ |
| 180 | mm_hashes, mm_kwargs = self.encoder_runner.prepare_mm_inputs( |
| 181 | scheduled_encoder_inputs |
| 182 | ) |
| 183 | if mm_kwargs: |
| 184 | with self.encoder_runner.timed_encoder_operation( |
| 185 | scheduled_encoder_inputs.keys() |
| 186 | ): |
| 187 | encoder_outputs = self.encoder_runner.execute_mm_encoder(mm_kwargs) |
| 188 | self.encoder_cache.encoder_outputs.update(zip(mm_hashes, encoder_outputs)) |
| 189 | |
| 190 | def gather_mm_embeddings( |
| 191 | self, input_batch: InputBatch, draft_lookahead: int = 0 |
| 192 | ) -> tuple[list[torch.Tensor], torch.Tensor]: |
| 193 | """Gather cached multimodal embeddings.""" |
| 194 | return self.encoder_runner.gather_mm_embeddings( |
| 195 | input_batch.req_ids, |
| 196 | input_batch.num_tokens, |
| 197 | input_batch.num_scheduled_tokens, |
| 198 | input_batch.query_start_loc_np, |
| 199 | input_batch.prefill_len_np, |
| 200 | input_batch.num_computed_tokens_np, |
| 201 | draft_lookahead=draft_lookahead, |
| 202 | ) |
| 203 | |
| 204 | @abstractmethod |
| 205 | def prepare_inputs( |
| 206 | self, input_batch: InputBatch, req_states: RequestState |
| 207 | ) -> dict[str, Any]: |
| 208 | raise NotImplementedError |
| 209 | |
| 210 | @abstractmethod |
| 211 | def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]: |
| 212 | raise NotImplementedError |
| 213 | |
| 214 | @abstractmethod |
| 215 | def prepare_attn( |
| 216 | self, |
| 217 | input_batch: InputBatch, |
| 218 | cudagraph_mode: CUDAGraphMode, |
| 219 | block_tables: tuple[torch.Tensor, ...], |
| 220 | slot_mappings: torch.Tensor, |
| 221 | attn_groups: list[list[AttentionGroup]], |
| 222 | kv_cache_config: KVCacheConfig, |
| 223 | for_capture: bool = False, |
| 224 | ubatch_idx: int = 0, |
| 225 | ) -> dict[str, Any]: |
| 226 | raise NotImplementedError |
| 227 | |
| 228 | def custom_sampler(self, sampler: Any) -> tuple[Any, Any] | None: |
| 229 | """Wrap or replace the default sampler. |
| 230 | |
| 231 | Called after model loading with the already-constructed base |
| 232 | ``Sampler``. Return ``None`` to keep the defaults, or |
| 233 | ``(sampler, rejection_sampler | None)`` to override. |
| 234 | """ |
| 235 | return None |
| 236 | |
| 237 | num_new_sampled_tokens_per_step: int = 1 |
| 238 | """New tokens sampled on each decode step |
| 239 | (excluding accepted draft tokens, a.k.a num bonus tokens).""" |