| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Kimi-K2.5 Model Implementation for vLLM. |
| 4 | |
| 5 | Kimi-K2.5 extends Kimi-K2 with vision support. |
| 6 | """ |
| 7 | |
| 8 | from collections.abc import Hashable, Iterable, Mapping, Sequence |
| 9 | from dataclasses import dataclass |
| 10 | from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Literal |
| 11 | |
| 12 | import torch |
| 13 | from torch import nn |
| 14 | from transformers import BatchFeature |
| 15 | |
| 16 | from vllm.config import VllmConfig |
| 17 | from vllm.config.multimodal import MultiModalDummyOptions |
| 18 | from vllm.inputs import MultiModalDataDict |
| 19 | from vllm.logger import init_logger |
| 20 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 21 | from vllm.model_executor.layers.quantization.compressed_tensors import ( |
| 22 | compressed_tensors, |
| 23 | ) |
| 24 | from vllm.model_executor.models.interfaces import ( |
| 25 | SupportsEagle, |
| 26 | SupportsEagle3, |
| 27 | SupportsEncoderCudaGraph, |
| 28 | SupportsMultiModal, |
| 29 | SupportsPP, |
| 30 | SupportsQuant, |
| 31 | ) |
| 32 | |
| 33 | if TYPE_CHECKING: |
| 34 | from vllm.v1.worker.encoder_cudagraph_defs import ( |
| 35 | EncoderCudaGraphCaptureInputs, |
| 36 | EncoderCudaGraphConfig, |
| 37 | EncoderCudaGraphReplayBuffers, |
| 38 | EncoderItemSpec, |
| 39 | ) |
| 40 | from vllm.model_executor.models.kimi_k25_vit import ( |
| 41 | KimiK25MultiModalProjector, |
| 42 | MoonViT3dPretrainedModel, |
| 43 | vision_tower_forward, |
| 44 | ) |
| 45 | from vllm.model_executor.models.vision import is_vit_use_data_parallel |
| 46 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 47 | from vllm.multimodal.inputs import ( |
| 48 | MultiModalFieldConfig, |
| 49 | MultiModalKwargsItems, |
| 50 | NestedTensors, |
| 51 | VisionChunkImage, |
| 52 | VisionChunkVideo, |
| 53 | ) |
| 54 | from vllm.multimodal.parse import MultiModalDataItems, VisionChunkProcessorItems |
| 55 | from vllm.multimodal.processing import ( |
| 56 | BaseDummyInputsBuilder, |
| 57 | BaseMultiModalProcessor, |
| 58 | BaseProcessingInfo, |
| 59 | InputProcessingContext, |
| 60 | PromptReplacement, |
| 61 | PromptUpdate, |
| 62 | ) |
| 63 | from vllm.platforms import current_platform |
| 64 | from vllm.sequence import IntermediateTensors |
| 65 | from vllm.transformers_utils.configs.kimi_k25 import KimiK25Config |
| 66 | from vllm.transformers_utils.processor import cached_get_image_processor |
| 67 | from vllm.transformers_utils.processors.kimi_k25 import KimiK25Processor |
| 68 | from vllm.transformers_utils.processors.kimi_k25_vision_fused import ( |
| 69 | KimiK25FusedVisionProcessor, |
| 70 | ) |
| 71 | from vllm.utils.import_utils import is_numba_available |
| 72 | from vllm.utils.tensor_schema import TensorSchema, TensorShape |
| 73 | |
| 74 | from .utils import ( |
| 75 | AutoWeightsLoader, |
| 76 | WeightsMapper, |
| 77 | init_vllm_registered_model, |
| 78 | maybe_prefix, |
| 79 | ) |
| 80 | |
| 81 | logger = init_logger(__name__) |
| 82 | |
| 83 | |
| 84 | # Dummy input dimensions for profiling. |
| 85 | @dataclass |
| 86 | class MaxImageTokenMeta: |
| 87 | width: int = 3000 |
| 88 | height: int = 3000 |
| 89 | |
| 90 | |
| 91 | class KimiK25MediaPixelInputs(TensorSchema): |
| 92 | """Media input schema for K2-VL model. |
| 93 | |
| 94 | Dimensions: |
| 95 | - np: Number of patches (flattened from all media items) |
| 96 | - ps: Patch size |
| 97 | - nm: Number of media items |
| 98 | """ |
| 99 | |
| 100 | type: Literal["pixel_values"] = "pixel_values" |
| 101 | |
| 102 | pixel_values: Annotated[ |
| 103 | torch.Tensor | list[torch.Tensor], |
| 104 | TensorShape("np", 3, "ps", "ps"), |
| 105 | ] |
| 106 | |
| 107 | grid_thws: Annotated[torch.Tensor, TensorShape("nm", 3)] |
| 108 | |
| 109 | |
| 110 | class KimiK25ProcessingInfo(BaseProcessingInfo): |
| 111 | """Processing information for Kimi-K2.5 model. |
| 112 | |
| 113 | Provides configuration and utilities for processing both |
| 114 | images and video-chunks. |
| 115 | """ |
| 116 | |
| 117 | def __init__(self, ctx: InputProcessingContext) -> None: |
| 118 | super().__init__(ctx) |
| 119 | |
| 120 | self.hf_config = hf_config = self.get_hf_config() |
| 121 | |
| 122 | tokenizer = self.get_tokenizer() |
| 123 | processor_cls = KimiK25FusedVisionProcessor if is_numba_available() else None |
| 124 | logger.info_once( |
| 125 | "Using %s image preprocessing for Kimi-K2.5/K2.6 vision chunks.", |
| 126 | "fused CPU" if processor_cls is not None else "remote HF", |
| 127 | ) |
| 128 | image_processor = cached_get_image_processor( |
| 129 | self.ctx.model_config.model, |
| 130 | revision=self.ctx.model_config.revision, |
| 131 | trust_remote_code=self.ctx.model_config.trust_remote_code, |
| 132 | processor_cls_overrides=processor_cls, |
| 133 | ) |
| 134 | |
| 135 | # Resolve token ID from the tokenizer because transformers v5 |
| 136 | # may remap token IDs vs config.json. |
| 137 | config_token_id = hf_config.media_placeholder_token_id |
| 138 | resolved_token_id = tokenizer.convert_tokens_to_ids("<|media_pad|>") |
| 139 | unk_token_id = getattr(tokenizer, "unk_token_id", None) |
| 140 | is_valid_resolved = isinstance(resolved_token_id, int) and ( |
| 141 | unk_token_id is None or resolved_token_id != unk_token_id |
| 142 | ) |
| 143 | if is_valid_resolved and resolved_token_id != config_token_id: |
| 144 | logger.warning_once( |
| 145 | "Kimi-K2.5 config.media_placeholder_token_id (%d) disagrees " |
| 146 | "with tokenizer mapping for <|media_pad|> (%d). " |
| 147 | "Using tokenizer value.", |
| 148 | config_token_id, |
| 149 | resolved_token_id, |
| 150 | ) |
| 151 | media_token_id = resolved_token_id |
| 152 | # Patch config so downstream code also sees the correct ID. |
| 153 | hf_config.media_placeholder_token_id = resolved_token_id |
| 154 | else: |
| 155 | media_token_id = config_token_id |
| 156 | |
| 157 | self.media_token_id = media_token_id |
| 158 | self.media_token = tokenizer.decode(media_token_id) |
| 159 | |
| 160 | self.image_processor = image_processor |
| 161 | self.hf_processor = KimiK25Processor( |
| 162 | tokenizer=tokenizer, |
| 163 | image_processor=image_processor, |
| 164 | media_token_id=media_token_id, |
| 165 | ) |
| 166 | self.media_tokens_calculator = image_processor.media_tokens_calculator |
| 167 | |
| 168 | def get_hf_processor(self, **kwargs: object): |
| 169 | return self.hf_processor |
| 170 | |
| 171 | def get_hf_config(self): |
| 172 | return self.ctx.get_hf_config(KimiK25Config) |
| 173 | |
| 174 | def get_supported_mm_limits(self) -> Mapping[str, int | None]: |
| 175 | # None means unlimited |
| 176 | return {"vision_chunk": None} |
| 177 | |
| 178 | |
| 179 | class KimiK25DummyInputsBuilder(BaseDummyInputsBuilder[KimiK25ProcessingInfo]): |
| 180 | """Builds dummy inputs for Kimi-K2.5 model profiling.""" |
| 181 | |
| 182 | def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: |
| 183 | num_media = mm_counts.get("vision_chunk", 0) |
| 184 | return self.info.media_token * num_media |
| 185 | |
| 186 | def get_dummy_mm_items(self): |
| 187 | dummy_videos = self._get_dummy_images( |
| 188 | height=MaxImageTokenMeta.height, |
| 189 | width=MaxImageTokenMeta.width, |
| 190 | num_images=self.info.image_processor.num_frames_per_chunk, |
| 191 | ) |
| 192 | |
| 193 | video_chunk_dummy_item = VisionChunkVideo( |
| 194 | type="video_chunk", |
| 195 | video_chunk=dummy_videos, |
| 196 | uuid=None, |
| 197 | prompt="", |
| 198 | video_idx=0, |
| 199 | ) |
| 200 | video_chunk_num_tokens = self.info.media_tokens_calculator( |
| 201 | video_chunk_dummy_item |
| 202 | ) |
| 203 | |
| 204 | image_dummy_item = VisionChunkImage( |
| 205 | type="image", |
| 206 | image=self._get_dummy_images( |
| 207 | height=MaxImageTokenMeta.height, |
| 208 | width=MaxImageTokenMeta.width, |
| 209 | num_images=1, |
| 210 | )[0], |
| 211 | uuid=None, |
| 212 | ) |
| 213 | image_num_tokens = self.info.media_tokens_calculator(image_dummy_item) |
| 214 | # return the larger one |
| 215 | if video_chunk_num_tokens >= image_num_tokens: |
| 216 | return [video_chunk_dummy_item] |
| 217 | else: |
| 218 | return [image_dummy_item] |
| 219 | |
| 220 | def get_dummy_mm_data( |
| 221 | self, |
| 222 | seq_len: int, |
| 223 | mm_counts: Mapping[str, int], |
| 224 | mm_options: MultiModalDummyOptions, |
| 225 | ) -> MultiModalDataDict: |
| 226 | # TODO: Support mm_options for vision_chunk to allow user configuration |
| 227 | dummy_items = self.get_dummy_mm_items() |
| 228 | return {"vision_chunk": dummy_items} |
| 229 | |
| 230 | |
| 231 | class KimiK25MultiModalProcessor(BaseMultiModalProcessor[KimiK25ProcessingInfo]): |
| 232 | """Multi-modal processor for Kimi-K2.5. |
| 233 | |
| 234 | Handles both image and video-chunk modalities. |
| 235 | """ |
| 236 | |
| 237 | def _get_mm_fields_config( |
| 238 | self, |
| 239 | hf_inputs: BatchFeature, |
| 240 | hf_processor_mm_kwargs: Mapping[str, object], |
| 241 | ) -> Mapping[str, MultiModalFieldConfig]: |
| 242 | """Indicates how to slice media input into multiple items. |
| 243 | |
| 244 | pixel_values: [N, 3, patch_size, patch_size], |
| 245 | all patches collected from B medias |
| 246 | grid_thws: [B,3], each item: [N_t, N_h ,N_w], |
| 247 | indicates the grid size in time/height/width direction for current item. |
| 248 | |
| 249 | by multiplying [N_t, N_h ,N_w], we get the number of patches |
| 250 | for each media item, thus we can slice pixel_values by |
| 251 | pixel_values[start:start + N_t*N_h*N_w] to get patches of one item. |
| 252 | |
| 253 | """ |
| 254 | grid_thws = hf_inputs.get("grid_thws", torch.empty((0, 3))) |
| 255 | grid_sizes = grid_thws.prod(-1) |
| 256 | |
| 257 | return dict( |
| 258 | pixel_values=MultiModalFieldConfig.flat_from_sizes( |
| 259 | "vision_chunk", grid_sizes |
| 260 | ), |
| 261 | grid_thws=MultiModalFieldConfig.batched("vision_chunk", keep_on_cpu=True), |
| 262 | ) |
| 263 | |
| 264 | def _get_prompt_updates( |
| 265 | self, |
| 266 | mm_items: MultiModalDataItems, |
| 267 | hf_processor_mm_kwargs: Mapping[str, Any], |
| 268 | out_mm_kwargs: MultiModalKwargsItems, |
| 269 | ) -> Sequence[PromptUpdate]: |
| 270 | media_token_id = self.info.media_token_id |
| 271 | |
| 272 | def get_replacement(item_idx: int): |
| 273 | media = mm_items.get_items("vision_chunk", (VisionChunkProcessorItems,)) |
| 274 | num_media_token = self.info.media_tokens_calculator(media[item_idx]) |
| 275 | return [media_token_id] * num_media_token |
| 276 | |
| 277 | return [ |
| 278 | PromptReplacement( |
| 279 | modality="vision_chunk", |
| 280 | target=[media_token_id], |
| 281 | replacement=get_replacement, |
| 282 | ), |
| 283 | ] |
| 284 | |
| 285 | |
| 286 | @MULTIMODAL_REGISTRY.register_processor( |
| 287 | KimiK25MultiModalProcessor, |
| 288 | info=KimiK25ProcessingInfo, |
| 289 | dummy_inputs=KimiK25DummyInputsBuilder, |
| 290 | ) |
| 291 | class KimiK25ForConditionalGeneration( |
| 292 | nn.Module, |
| 293 | SupportsMultiModal, |
| 294 | SupportsPP, |
| 295 | SupportsQuant, |
| 296 | SupportsEagle, |
| 297 | SupportsEagle3, |
| 298 | SupportsEncoderCudaGraph, |
| 299 | ): |
| 300 | """Kimi-K2.5 model for conditional generation. |
| 301 | |
| 302 | Supports both image and video-chunk modalities. |
| 303 | Video-chunks are temporal segments (typically 4 frames) that are |
| 304 | processed with temporal pooling. |
| 305 | """ |
| 306 | |
| 307 | supports_encoder_tp_data = True |
| 308 | supports_encoder_cudagraph: ClassVar[Literal[True]] = True |
| 309 | |
| 310 | hf_to_vllm_mapper = WeightsMapper( |
| 311 | orig_to_new_prefix={ |
| 312 | # For legacy NVFP4 checkpoint compatibility: |
| 313 | # see https://github.com/vllm-project/vllm/pull/33346#issuecomment-3851475033 |
| 314 | "language_model.layers.": "language_model.model.layers.", |
| 315 | # mm projector |
| 316 | "mm_projector.proj.0": "mm_projector.linear_1", |
| 317 | "mm_projector.proj.2": "mm_projector.linear_2", |
| 318 | } |
| 319 | ) |
| 320 | |
| 321 | @classmethod |
| 322 | def get_placeholder_str(cls, modality: str, i: int) -> str | None: |
| 323 | # Kimi-K2.5 uses video_chunk for all media types |
| 324 | if modality == "image": |
| 325 | return "<|media_begin|>image<|media_content|><|media_pad|><|media_end|>" |
| 326 | elif modality == "video": |
| 327 | # return a placeholder, to be replaced in the future. |
| 328 | return "<|kimi_k25_video_placeholder|>" |
| 329 | |
| 330 | raise ValueError(f"Unsupported modality: {modality}") |
| 331 | |
| 332 | def __init__( |
| 333 | self, |
| 334 | vllm_config: VllmConfig, |
| 335 | prefix: str = "", |
| 336 | ) -> None: |
| 337 | super().__init__() |
| 338 | model_config = vllm_config.model_config |
| 339 | config: KimiK25Config = model_config.hf_config |
| 340 | self.config = config |
| 341 | quant_config = vllm_config.quant_config |
| 342 | |
| 343 | self.use_data_parallel = is_vit_use_data_parallel( |
| 344 | config.vision_config.num_attention_heads |
| 345 | ) |
| 346 | self.hidden_size = config.text_config.hidden_size |
| 347 | self.device = current_platform.current_device() |
| 348 | # Build vision tower directly with KimiK25VisionConfig |
| 349 | with self._mark_tower_model(vllm_config, "vision_chunk"): |
| 350 | self.vision_tower = MoonViT3dPretrainedModel( |
| 351 | config.vision_config, |
| 352 | quant_config=self._maybe_ignore_quant_config(quant_config), |
| 353 | prefix=maybe_prefix(prefix, "vision_tower"), |
| 354 | ) |
| 355 | if self._maybe_ignore_quant_config(quant_config) is not None: |
| 356 | self.vision_tower = self.vision_tower.to(device=self.device) |
| 357 | else: |
| 358 | self.vision_tower = self.vision_tower.to( |
| 359 | device=self.device, dtype=model_config.dtype |
| 360 | ) |
| 361 | |
| 362 | self.mm_projector = KimiK25MultiModalProjector( |
| 363 | config=config.vision_config, |
| 364 | use_data_parallel=self.use_data_parallel, |
| 365 | quant_config=self._maybe_ignore_quant_config(quant_config), |
| 366 | prefix=maybe_prefix(prefix, "mm_projector"), |
| 367 | ) |
| 368 | self.mm_projector = self.mm_projector.to( |
| 369 | device=self.device, dtype=model_config.dtype |
| 370 | ) |
| 371 | |
| 372 | self.quant_config = quant_config |
| 373 | with self._mark_language_model(vllm_config): |
| 374 | self.language_model = init_vllm_registered_model( |
| 375 | vllm_config=vllm_config, |
| 376 | hf_config=config.text_config, |
| 377 | prefix=maybe_prefix(prefix, "language_model"), |
| 378 | architectures=["DeepseekV2ForCausalLM"], |
| 379 | ) |
| 380 | self.make_empty_intermediate_tensors = ( |
| 381 | self.language_model.make_empty_intermediate_tensors |
| 382 | ) |
| 383 | self.media_placeholder: int = self.config.media_placeholder_token_id |
| 384 | |
| 385 | def _maybe_ignore_quant_config(self, quant_config: QuantizationConfig | None): |
| 386 | if isinstance(quant_config, compressed_tensors.CompressedTensorsConfig): |
| 387 | return None |
| 388 | return quant_config |
| 389 | |
| 390 | def _parse_and_validate_media_input( |
| 391 | self, **kwargs: object |
| 392 | ) -> KimiK25MediaPixelInputs | None: |
| 393 | pixel_values = kwargs.pop("pixel_values", None) |
| 394 | grid_thws = kwargs.pop("grid_thws", None) |
| 395 | if pixel_values is None: |
| 396 | return None |
| 397 | |
| 398 | if isinstance(pixel_values, list): |
| 399 | pixel_values = torch.cat(pixel_values, dim=0) |
| 400 | assert isinstance(pixel_values, torch.Tensor) |
| 401 | |
| 402 | if len(pixel_values.shape) == 5 or len(pixel_values.shape) == 3: |
| 403 | pixel_values = pixel_values.reshape( |
| 404 | pixel_values.shape[0] * pixel_values.shape[1], *pixel_values.shape[2:] |
| 405 | ) |
| 406 | |
| 407 | # The batch dimension of pixel_values has been flattened into shape[0] |
| 408 | target_dtype = next(self.vision_tower.parameters()).dtype |
| 409 | pixel_values = pixel_values.to(target_dtype) |
| 410 | assert isinstance(grid_thws, torch.Tensor), ( |
| 411 | f"expect grid_thws to be a tensor, got {type(grid_thws)}" |
| 412 | ) |
| 413 | # In some cases (e.g. with merger), grid_thws has an extra middle dimension |
| 414 | grid_thws = grid_thws.reshape(-1, grid_thws.shape[-1]) |
| 415 | assert grid_thws.ndim == 2 and grid_thws.size(1) == 3, ( |
| 416 | f"unexpected shape for grid_thws: {grid_thws.shape}" |
| 417 | ) |
| 418 | |
| 419 | return KimiK25MediaPixelInputs( |
| 420 | type="pixel_values", |
| 421 | pixel_values=pixel_values, |
| 422 | grid_thws=grid_thws, |
| 423 | ) |
| 424 | |
| 425 | def _process_media_input( |
| 426 | self, media_input: KimiK25MediaPixelInputs |
| 427 | ) -> list[torch.Tensor]: |
| 428 | # NOTE(moyan): This forward will automatically batch the forward pass internally |
| 429 | media_features = vision_tower_forward( |
| 430 | self.vision_tower, |
| 431 | media_input["pixel_values"], |
| 432 | media_input["grid_thws"], |
| 433 | mm_projector=self.mm_projector, |
| 434 | use_data_parallel=self.use_data_parallel, |
| 435 | ) |
| 436 | return media_features |
| 437 | |
| 438 | def embed_multimodal(self, **kwargs: object) -> NestedTensors | None: |
| 439 | # Validate the multimodal input keyword arguments |
| 440 | media_input = self._parse_and_validate_media_input(**kwargs) |
| 441 | if media_input is None: |
| 442 | return None |
| 443 | |
| 444 | # Run multimodal inputs through encoder and projector |
| 445 | vision_embeddings = self._process_media_input(media_input) |
| 446 | return vision_embeddings |
| 447 | |
| 448 | def forward( |
| 449 | self, |
| 450 | input_ids: torch.Tensor, |
| 451 | positions: torch.Tensor, |
| 452 | intermediate_tensors: IntermediateTensors | None = None, |
| 453 | inputs_embeds: torch.Tensor | None = None, |
| 454 | **kwargs: object, |
| 455 | ) -> IntermediateTensors: |
| 456 | if intermediate_tensors is not None: |
| 457 | inputs_embeds = None |
| 458 | hidden_states = self.language_model( |
| 459 | input_ids=input_ids, |
| 460 | positions=positions, |
| 461 | intermediate_tensors=intermediate_tensors, |
| 462 | inputs_embeds=inputs_embeds, |
| 463 | ) |
| 464 | |
| 465 | return hidden_states |
| 466 | |
| 467 | def compute_logits(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor: |
| 468 | logits = self.language_model.compute_logits(hidden_states) |
| 469 | return logits |
| 470 | |
| 471 | def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: |
| 472 | self.language_model.set_aux_hidden_state_layers(layers) |
| 473 | |
| 474 | def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]: |
| 475 | return self.language_model.get_eagle3_aux_hidden_state_layers() |
| 476 | |
| 477 | # ------------------------------------------------------------------ # |
| 478 | # SupportsEncoderCudaGraph protocol # |
| 479 | # Image-only (t == 1). Video chunks (t > 1) fall back to eager. # |
| 480 | # ------------------------------------------------------------------ # |
| 481 | |
| 482 | @property |
| 483 | def _encoder_cudagraph_pad_totals(self) -> dict[int, int]: |
| 484 | """Row count of each captured buffer set, keyed by cu_seqlens ptr.""" |
| 485 | totals = self.__dict__.get("_encoder_cg_pad_totals") |
| 486 | if totals is None: |
| 487 | totals = {} |
| 488 | self.__dict__["_encoder_cg_pad_totals"] = totals |
| 489 | return totals |
| 490 | |
| 491 | def get_encoder_cudagraph_config(self) -> "EncoderCudaGraphConfig": |
| 492 | from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphConfig |
| 493 | |
| 494 | pad_totals = self._encoder_cudagraph_pad_totals |
| 495 | |
| 496 | def pad_cu_seqlens(dst: torch.Tensor, src: torch.Tensor) -> None: |
| 497 | # Varlen attention requires cu_seqlens[-1] to equal the number of |
| 498 | # rows actually passed in. The captured buffers are sized for the |
| 499 | # full token budget, so a smaller real batch has to be completed |
| 500 | # with one trailing padding sequence; declaring fewer rows than the |
| 501 | # buffer holds is undefined behaviour and returns NaN on FlashAttn. |
| 502 | total = pad_totals.get(dst.data_ptr()) |
| 503 | n = min(src.shape[0], dst.shape[0]) |
| 504 | dst[:n].copy_(src[:n]) |
| 505 | dst[n:] = total if total is not None else src[-1] |
| 506 | |
| 507 | return EncoderCudaGraphConfig( |
| 508 | modalities=["vision_chunk"], |
| 509 | buffer_keys=[ |
| 510 | "pixel_values", |
| 511 | "pos_embeds", |
| 512 | "rope_freqs_cis", |
| 513 | "cu_seqlens", |
| 514 | "max_seqlen", |
| 515 | "merge_gather_idx", |
| 516 | ], |
| 517 | out_hidden_size=self.config.text_config.hidden_size, |
| 518 | padding_logics={"cu_seqlens": pad_cu_seqlens}, |
| 519 | ) |
| 520 | |
| 521 | def get_encoder_cudagraph_budget_range( |
| 522 | self, |
| 523 | vllm_config: VllmConfig, |
| 524 | ) -> tuple[int, int]: |
| 525 | # Min: 64 output tokens (e.g. ~128×128 image with patch 14 + merge 2×2). |
| 526 | min_budget = 64 |
| 527 | max_budget = min( |
| 528 | vllm_config.scheduler_config.max_num_batched_tokens, |
| 529 | vllm_config.model_config.max_model_len, |
| 530 | ) |
| 531 | return (min_budget, max_budget) |
| 532 | |
| 533 | def _get_grid_thw_list( |
| 534 | self, |
| 535 | mm_kwargs: dict[str, Any], |
| 536 | ) -> list[list[int]]: |
| 537 | grid_thws = mm_kwargs["grid_thws"] |
| 538 | if isinstance(grid_thws, torch.Tensor): |
| 539 | return [[int(x) for x in row] for row in grid_thws.tolist()] |
| 540 | return [[int(x) for x in row] for row in grid_thws] |
| 541 | |
| 542 | def get_encoder_cudagraph_item_specs( |
| 543 | self, |
| 544 | mm_kwargs: dict[str, Any], |
| 545 | ) -> list["EncoderItemSpec"]: |
| 546 | from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec |
| 547 | |
| 548 | kh, kw = self.vision_tower.merge_kernel_size |
| 549 | specs = [] |
| 550 | for t, h, w in self._get_grid_thw_list(mm_kwargs): |
| 551 | if t != 1: |
| 552 | # Video chunks not supported in encoder CUDA graph; |
| 553 | # sentinel forces eager fallback via the manager. |
| 554 | specs.append(EncoderItemSpec(input_size=t * h * w, output_tokens=2**30)) |
| 555 | else: |
| 556 | specs.append( |
| 557 | EncoderItemSpec( |
| 558 | input_size=h * w, |
| 559 | output_tokens=(h // kh) * (w // kw), |
| 560 | ) |
| 561 | ) |
| 562 | return specs |
| 563 | |
| 564 | def select_encoder_cudagraph_items( |
| 565 | self, |
| 566 | mm_kwargs: dict[str, Any], |
| 567 | indices: list[int], |
| 568 | ) -> dict[str, Any]: |
| 569 | grid_thw_list = self._get_grid_thw_list(mm_kwargs) |
| 570 | pixel_values = mm_kwargs["pixel_values"] |
| 571 | |
| 572 | if len(indices) == 0: |
| 573 | return { |
| 574 | "pixel_values": pixel_values[:0], |
| 575 | "grid_thws": pixel_values.new_zeros((0, 3), dtype=torch.long), |
| 576 | } |
| 577 | |
| 578 | patch_counts = [t * h * w for t, h, w in grid_thw_list] |
| 579 | cum = [0] |
| 580 | for pc in patch_counts: |
| 581 | cum.append(cum[-1] + pc) |
| 582 | |
| 583 | selected_pv = torch.cat( |
| 584 | [pixel_values[cum[i] : cum[i + 1]] for i in indices], dim=0 |
| 585 | ) |
| 586 | selected_grid = torch.tensor( |
| 587 | [grid_thw_list[i] for i in indices], |
| 588 | dtype=torch.long, |
| 589 | device=pixel_values.device, |
| 590 | ) |
| 591 | return {"pixel_values": selected_pv, "grid_thws": selected_grid} |
| 592 | |
| 593 | def prepare_encoder_cudagraph_capture_inputs( |
| 594 | self, |
| 595 | token_budget: int, |
| 596 | max_batch_size: int, |
| 597 | max_frames_per_batch: int, |
| 598 | device: torch.device, |
| 599 | dtype: torch.dtype, |
| 600 | path: str = "default", |
| 601 | axis_keys: tuple[Hashable, ...] | None = None, |
| 602 | ) -> "EncoderCudaGraphCaptureInputs": |
| 603 | from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphCaptureInputs |
| 604 | |
| 605 | kh, kw = self.vision_tower.merge_kernel_size |
| 606 | # Output tokens per item in the dummy grid (ceiling so total >= budget). |
| 607 | per_item_out = (token_budget + max_batch_size - 1) // max_batch_size |
| 608 | |
| 609 | # Fit within RoPE precomputed max dimensions. |
| 610 | rope = self.vision_tower.encoder.rope_2d |
| 611 | max_wo = rope.max_width // kw |
| 612 | wo = min(per_item_out, max_wo) |
| 613 | ho = (per_item_out + wo - 1) // wo |
| 614 | assert ho * kh <= rope.max_height, ( |
| 615 | f"per_item_out={per_item_out} exceeds RoPE grid capacity " |
| 616 | f"(max {(rope.max_height // kh) * (rope.max_width // kw)} tokens)" |
| 617 | ) |
| 618 | |
| 619 | grid_thw_list = [[1, ho * kh, wo * kw] for _ in range(max_batch_size)] |
| 620 | |
| 621 | ps: int | tuple[int, int] = self.vision_tower.patch_size |
| 622 | if isinstance(ps, int): |
| 623 | ps = (ps, ps) |
| 624 | total_patches = max_batch_size * ho * kh * wo * kw |
| 625 | dummy_pixel_values = torch.zeros( |
| 626 | total_patches, 3, ps[0], ps[1], device=device, dtype=dtype |
| 627 | ) |
| 628 | |
| 629 | # max_seqlen must cover the worst case: one item consuming the full |
| 630 | # budget, i.e. token_budget * kh * kw patches. |
| 631 | # max_batch_size + 1 leaves a spare cu_seqlens slot so replay can append |
| 632 | # a padding sequence covering rows the real batch does not fill. |
| 633 | metadata = self.vision_tower.prepare_encoder_cudagraph_metadata( |
| 634 | grid_thw_list, |
| 635 | max_batch_size=max_batch_size + 1, |
| 636 | max_seqlen_override=token_budget * kh * kw, |
| 637 | device=device, |
| 638 | ) |
| 639 | |
| 640 | values: dict[str, torch.Tensor] = {"pixel_values": dummy_pixel_values} |
| 641 | values.update({k: v for k, v in metadata.items() if v is not None}) |
| 642 | |
| 643 | cu_seqlens = values.get("cu_seqlens") |
| 644 | if cu_seqlens is not None: |
| 645 | self._encoder_cudagraph_pad_totals[cu_seqlens.data_ptr()] = total_patches |
| 646 | |
| 647 | return EncoderCudaGraphCaptureInputs(values=values) |
| 648 | |
| 649 | def prepare_encoder_cudagraph_replay_buffers( |
| 650 | self, |
| 651 | mm_kwargs: dict[str, Any], |
| 652 | max_batch_size: int, |
| 653 | max_frames_per_batch: int, |
| 654 | path: str = "default", |
| 655 | ) -> "EncoderCudaGraphReplayBuffers": |
| 656 | from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphReplayBuffers |
| 657 | |
| 658 | grid_thw_list = self._get_grid_thw_list(mm_kwargs) |
| 659 | pixel_values = mm_kwargs["pixel_values"] |
| 660 | |
| 661 | # Unpadded: pad_cu_seqlens completes the tail with the padding sequence |
| 662 | # so cu_seqlens[-1] matches the captured buffer's row count. |
| 663 | metadata = self.vision_tower.prepare_encoder_cudagraph_metadata( |
| 664 | grid_thw_list, |
| 665 | max_batch_size=None, |
| 666 | device=pixel_values.device, |
| 667 | ) |
| 668 | |
| 669 | values: dict[str, torch.Tensor | None] = {"pixel_values": pixel_values} |
| 670 | values.update(metadata) |
| 671 | return EncoderCudaGraphReplayBuffers(values=values) |
| 672 | |
| 673 | def encoder_cudagraph_forward( |
| 674 | self, |
| 675 | inputs: dict[str, torch.Tensor], |
| 676 | path: str = "default", |
| 677 | ) -> torch.Tensor: |
| 678 | pixel_values = inputs.pop("pixel_values") |
| 679 | # Remaining keys (pos_embeds, rope_freqs_cis, cu_seqlens, max_seqlen, |
| 680 | # merge_gather_idx, sequence_lengths) are consumed as encoder_metadata. |
| 681 | encoder_metadata = inputs |
| 682 | |
| 683 | # Fast path: uses precomputed pos_embeds + merge_gather_idx. |
| 684 | vt_output = self.vision_tower( |
| 685 | pixel_values, grid_thws=None, encoder_metadata=encoder_metadata |
| 686 | ) |
| 687 | # vt_output: (total_output_tokens, kh*kw, vit_hidden_dim) |
| 688 | |
| 689 | proj_dtype = next(self.mm_projector.parameters()).dtype |
| 690 | if vt_output.dtype != proj_dtype: |
| 691 | vt_output = vt_output.to(proj_dtype) |
| 692 | |
| 693 | projected = self.mm_projector(vt_output) |
| 694 | return projected.view(-1, self.config.text_config.hidden_size) |
| 695 | |
| 696 | def encoder_eager_forward( |
| 697 | self, |
| 698 | mm_kwargs: dict[str, Any], |
| 699 | path: str = "default", |
| 700 | ) -> torch.Tensor: |
| 701 | media_input = self._parse_and_validate_media_input(**mm_kwargs) |
| 702 | if media_input is None: |
| 703 | proj_dtype = next(self.mm_projector.parameters()).dtype |
| 704 | return torch.zeros( |
| 705 | 0, |
| 706 | self.config.text_config.hidden_size, |
| 707 | device=self.device, |
| 708 | dtype=proj_dtype, |
| 709 | ) |
| 710 | embeddings = self._process_media_input(media_input) |
| 711 | return torch.cat(embeddings, dim=0) |
| 712 | |
| 713 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): |
| 714 | loader = AutoWeightsLoader(self) |
| 715 | return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) |