| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import threading |
| 4 | import time |
| 5 | from collections.abc import Collection |
| 6 | from contextlib import contextmanager |
| 7 | from typing import TYPE_CHECKING |
| 8 | |
| 9 | import numpy as np |
| 10 | import torch |
| 11 | |
| 12 | from vllm.logger import init_logger |
| 13 | from vllm.model_executor.models.interfaces import SupportsMultiModal, supports_realtime |
| 14 | from vllm.multimodal.encoder_budget import MultiModalBudget |
| 15 | from vllm.multimodal.inputs import MultiModalKwargsItem |
| 16 | from vllm.multimodal.utils import ( |
| 17 | get_mm_features_in_window, |
| 18 | group_and_batch_mm_kwargs, |
| 19 | set_mm_embedding_modality, |
| 20 | ) |
| 21 | from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d |
| 22 | from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache |
| 23 | from vllm.v1.worker.utils import ( |
| 24 | EncoderTimingStats, |
| 25 | sanity_check_mm_encoder_outputs, |
| 26 | ) |
| 27 | |
| 28 | if TYPE_CHECKING: |
| 29 | from vllm.v1.worker.encoder_cudagraph import EncoderCudaGraphManager |
| 30 | |
| 31 | logger = init_logger(__name__) |
| 32 | |
| 33 | |
| 34 | class EncoderRunner: |
| 35 | def __init__( |
| 36 | self, |
| 37 | model: SupportsMultiModal, |
| 38 | max_num_tokens: int, |
| 39 | hidden_size: int, |
| 40 | encoder_cache: EncoderCache, |
| 41 | dtype: torch.dtype, |
| 42 | device: torch.device, |
| 43 | cudagraph_manager: "EncoderCudaGraphManager | None" = None, |
| 44 | enable_timing: bool = False, |
| 45 | ): |
| 46 | self.model = model |
| 47 | self.max_num_tokens = max_num_tokens |
| 48 | self.hidden_size = hidden_size |
| 49 | self.encoder_cache = encoder_cache |
| 50 | self.dtype = dtype |
| 51 | self.device = device |
| 52 | self.is_realtime = supports_realtime(model) |
| 53 | self.cudagraph_manager = cudagraph_manager |
| 54 | self.enable_timing = enable_timing |
| 55 | self.encoder_timing_registry: dict[str, EncoderTimingStats] = {} |
| 56 | self._timing_lock = threading.Lock() |
| 57 | |
| 58 | self.inputs_embeds = torch.zeros( |
| 59 | max_num_tokens, hidden_size, dtype=dtype, device=device |
| 60 | ) |
| 61 | |
| 62 | def has_cudagraph(self) -> bool: |
| 63 | return self.cudagraph_manager is not None |
| 64 | |
| 65 | @torch.inference_mode() |
| 66 | def capture(self) -> None: |
| 67 | manager = self.cudagraph_manager |
| 68 | assert manager is not None |
| 69 | |
| 70 | from vllm.distributed.parallel_state import graph_capture |
| 71 | from vllm.platforms import current_platform |
| 72 | |
| 73 | with graph_capture(device=self.device): |
| 74 | manager.capture(graph_pool=current_platform.graph_pool_handle()) |
| 75 | torch.accelerator.synchronize() |
| 76 | |
| 77 | def clear(self) -> None: |
| 78 | if self.cudagraph_manager is not None: |
| 79 | self.cudagraph_manager.clear() |
| 80 | |
| 81 | def prepare_mm_inputs( |
| 82 | self, scheduled_encoder_inputs: dict[str, list[int]] |
| 83 | ) -> tuple[list[str], list[tuple[str, MultiModalKwargsItem]]]: |
| 84 | mm_hashes: list[str] = [] |
| 85 | mm_kwargs: list[tuple[str, MultiModalKwargsItem]] = [] |
| 86 | for req_id, encoder_input_ids in scheduled_encoder_inputs.items(): |
| 87 | mm_features = self.encoder_cache.mm_features[req_id] |
| 88 | for mm_input_id in encoder_input_ids: |
| 89 | mm_feature = mm_features[mm_input_id] |
| 90 | if mm_feature.data is None: |
| 91 | continue |
| 92 | if mm_feature.identifier in self.encoder_cache.encoder_outputs: |
| 93 | continue |
| 94 | if mm_feature.modality == "prompt_embeds": |
| 95 | # Passthrough modality: the tensor is already in the |
| 96 | # model's embedding space, so no encoder runs. Cache it |
| 97 | # directly so gather_mm_embeddings splices it via the |
| 98 | # standard is_mm_embed path. |
| 99 | embeds = mm_feature.data["embedding"].data |
| 100 | assert isinstance(embeds, torch.Tensor) |
| 101 | self.encoder_cache.encoder_outputs[mm_feature.identifier] = ( |
| 102 | async_tensor_h2d(embeds, device=self.device) |
| 103 | ) |
| 104 | continue |
| 105 | mm_hashes.append(mm_feature.identifier) |
| 106 | mm_kwargs.append((mm_feature.modality, mm_feature.data)) |
| 107 | |
| 108 | return mm_hashes, mm_kwargs |
| 109 | |
| 110 | @torch.inference_mode() |
| 111 | def profile_encoder_cache( |
| 112 | self, |
| 113 | dummy_mm_inputs: list[tuple[str, MultiModalKwargsItem]], |
| 114 | budget: MultiModalBudget, |
| 115 | ) -> None: |
| 116 | """Profile multimodal encoder and temporary encoder cache memory.""" |
| 117 | if (encoder_budget := budget.get_encoder_budget()) <= 0: |
| 118 | return |
| 119 | |
| 120 | if not budget.mm_max_toks_per_item: |
| 121 | logger.info( |
| 122 | "Skipping encoder profiling for embedding-only mode " |
| 123 | "(all modality limits=0 with enable_mm_embeds=True).", |
| 124 | ) |
| 125 | return |
| 126 | |
| 127 | assert dummy_mm_inputs, "Dummy inputs should be generated for encoder profiling" |
| 128 | dummy_modality = dummy_mm_inputs[0][0] |
| 129 | max_mm_items_per_batch = len(dummy_mm_inputs) |
| 130 | |
| 131 | logger.info_once( |
| 132 | "Encoder cache will be initialized with a budget of %s tokens, " |
| 133 | "and profiled with %s %s items of the maximum feature size.", |
| 134 | encoder_budget, |
| 135 | max_mm_items_per_batch, |
| 136 | dummy_modality, |
| 137 | ) |
| 138 | |
| 139 | dummy_encoder_outputs = self.execute_mm_encoder(dummy_mm_inputs) |
| 140 | |
| 141 | sanity_check_mm_encoder_outputs( |
| 142 | dummy_encoder_outputs, |
| 143 | expected_num_items=max_mm_items_per_batch, |
| 144 | ) |
| 145 | self.encoder_cache.encoder_outputs.update( |
| 146 | (f"tmp_{i}", output) for i, output in enumerate(dummy_encoder_outputs) |
| 147 | ) |
| 148 | |
| 149 | @torch.inference_mode() |
| 150 | def execute_mm_encoder( |
| 151 | self, mm_kwargs: list[tuple[str, MultiModalKwargsItem]] |
| 152 | ) -> list[torch.Tensor]: |
| 153 | encoder_outputs: list[torch.Tensor] = [] |
| 154 | for modality, num_items, mm_kwargs_batch in group_and_batch_mm_kwargs( |
| 155 | mm_kwargs, device=self.device, pin_memory=PIN_MEMORY |
| 156 | ): |
| 157 | cg_manager = self.cudagraph_manager |
| 158 | cudagraph_output = ( |
| 159 | cg_manager.execute(mm_kwargs_batch) |
| 160 | if cg_manager is not None |
| 161 | and cg_manager.is_captured() |
| 162 | and cg_manager.supports_modality(modality) |
| 163 | else None |
| 164 | ) |
| 165 | batch_outputs = ( |
| 166 | cudagraph_output |
| 167 | if cudagraph_output is not None |
| 168 | else self.model.embed_multimodal(**mm_kwargs_batch) |
| 169 | ) |
| 170 | sanity_check_mm_encoder_outputs(batch_outputs, expected_num_items=num_items) |
| 171 | encoder_outputs.extend(batch_outputs) |
| 172 | return encoder_outputs |
| 173 | |
| 174 | @contextmanager |
| 175 | def timed_encoder_operation(self, request_ids: Collection[str]): |
| 176 | if not (self.enable_timing and request_ids): |
| 177 | yield |
| 178 | return |
| 179 | |
| 180 | torch.accelerator.synchronize() |
| 181 | start_time = time.perf_counter() |
| 182 | try: |
| 183 | yield |
| 184 | finally: |
| 185 | torch.accelerator.synchronize() |
| 186 | per_request_time = (time.perf_counter() - start_time) / len(request_ids) |
| 187 | with self._timing_lock: |
| 188 | for req_id in request_ids: |
| 189 | stats = self.encoder_timing_registry.setdefault( |
| 190 | req_id, EncoderTimingStats() |
| 191 | ) |
| 192 | stats.encoder_forward_secs += per_request_time |
| 193 | stats.num_encoder_calls += 1 |
| 194 | |
| 195 | def get_encoder_timing_stats(self) -> dict[str, dict[str, float | int]]: |
| 196 | with self._timing_lock: |
| 197 | stats = { |
| 198 | req_id: stats_obj.to_dict() |
| 199 | for req_id, stats_obj in self.encoder_timing_registry.items() |
| 200 | } |
| 201 | self.encoder_timing_registry.clear() |
| 202 | return stats |
| 203 | |
| 204 | def gather_mm_embeddings( |
| 205 | self, |
| 206 | req_ids: list[str], |
| 207 | total_num_scheduled_tokens: int, |
| 208 | num_scheduled_tokens: np.ndarray, |
| 209 | query_start_loc: np.ndarray, |
| 210 | prefill_lens: np.ndarray, |
| 211 | num_computed_tokens: np.ndarray, |
| 212 | draft_lookahead: int = 0, |
| 213 | ) -> tuple[list[torch.Tensor], torch.Tensor]: |
| 214 | if draft_lookahead: |
| 215 | num_computed_tokens = num_computed_tokens + draft_lookahead |
| 216 | |
| 217 | is_mm_embed = torch.zeros( |
| 218 | total_num_scheduled_tokens, |
| 219 | dtype=torch.bool, |
| 220 | device="cpu", |
| 221 | pin_memory=PIN_MEMORY, |
| 222 | ) |
| 223 | |
| 224 | # Whether to gather media embeddings this step. |
| 225 | exclude_embeddings: list[bool] | None = None |
| 226 | if not self.is_realtime: |
| 227 | # Non-realtime models only have media embeddings within the prompt. |
| 228 | is_decode = num_computed_tokens >= prefill_lens |
| 229 | if is_decode.all(): |
| 230 | # All decode requests, so no need to gather any embeddings. |
| 231 | return [], is_mm_embed |
| 232 | exclude_embeddings = is_decode.tolist() |
| 233 | |
| 234 | query_start = num_computed_tokens.tolist() |
| 235 | query_end = (num_computed_tokens + num_scheduled_tokens).tolist() |
| 236 | |
| 237 | mm_embeds: list[torch.Tensor] = [] |
| 238 | for i, req_id in enumerate(req_ids): |
| 239 | if exclude_embeddings is not None and exclude_embeddings[i]: |
| 240 | continue |
| 241 | |
| 242 | cur_query_start = query_start[i] |
| 243 | cur_query_end = query_end[i] |
| 244 | |
| 245 | mm_features = self.encoder_cache.mm_features[req_id] |
| 246 | lo, hi = get_mm_features_in_window( |
| 247 | mm_features, start=cur_query_start, end=cur_query_end |
| 248 | ) |
| 249 | for idx in range(lo, hi): |
| 250 | mm_feature = mm_features[idx] |
| 251 | pos_info = mm_feature.mm_position |
| 252 | start_pos = pos_info.offset |
| 253 | num_encoder_tokens = pos_info.length |
| 254 | |
| 255 | start_idx = max(cur_query_start - start_pos, 0) |
| 256 | end_idx = min(cur_query_end - start_pos, num_encoder_tokens) |
| 257 | assert start_idx < end_idx |
| 258 | curr_embeds_start, curr_embeds_end = ( |
| 259 | pos_info.get_embeds_indices_in_range(start_idx, end_idx) |
| 260 | ) |
| 261 | # If there are no embeddings in the current range, we skip |
| 262 | # gathering the embeddings. |
| 263 | if curr_embeds_start == curr_embeds_end: |
| 264 | continue |
| 265 | |
| 266 | mm_hash = mm_feature.identifier |
| 267 | encoder_output = self.encoder_cache.encoder_outputs.get(mm_hash, None) |
| 268 | if encoder_output is None: |
| 269 | # A feature starting at/after the processed boundary is only |
| 270 | # reached via the drafter's +1 look-ahead and might not be |
| 271 | # encoded yet; fall back to the token embedding for drafting. |
| 272 | if start_pos + draft_lookahead >= cur_query_end: |
| 273 | continue |
| 274 | raise RuntimeError(f"Encoder cache miss for {mm_hash}.") |
| 275 | |
| 276 | if (is_embed := pos_info.is_embed) is not None: |
| 277 | is_embed = is_embed[start_idx:end_idx] |
| 278 | mm_embeds_item = encoder_output[curr_embeds_start:curr_embeds_end] |
| 279 | else: |
| 280 | mm_embeds_item = encoder_output[start_idx:end_idx] |
| 281 | |
| 282 | # Attach modality for Omni interleaved merge (collected on demand). |
| 283 | set_mm_embedding_modality(mm_embeds_item, mm_feature.modality) |
| 284 | |
| 285 | req_start_pos = query_start_loc[i] + start_pos - cur_query_start |
| 286 | is_mm_embed[req_start_pos + start_idx : req_start_pos + end_idx] |= ( |
| 287 | True if is_embed is None else is_embed |
| 288 | ) |
| 289 | mm_embeds.append(mm_embeds_item) |
| 290 | |
| 291 | return mm_embeds, is_mm_embed |
| 292 | |
| 293 | @torch.inference_mode() |
| 294 | def get_inputs_embeds( |
| 295 | self, |
| 296 | input_ids: torch.Tensor, |
| 297 | mm_embeds: list[torch.Tensor], |
| 298 | is_mm_embed: torch.Tensor, |
| 299 | ) -> torch.Tensor: |
| 300 | x = self.model.embed_input_ids( |
| 301 | input_ids, multimodal_embeddings=mm_embeds, is_multimodal=is_mm_embed |
| 302 | ) |
| 303 | # Copy to the pre-allocated buffer for CUDA graphs. |
| 304 | self.inputs_embeds[: x.shape[0]] = x |
| 305 | return self.inputs_embeds |