| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import math |
| 4 | from collections.abc import Hashable, Iterable, Mapping, Sequence |
| 5 | from math import sqrt |
| 6 | from typing import Annotated, Any, Literal, TypeAlias |
| 7 | |
| 8 | import torch |
| 9 | import torch.nn as nn |
| 10 | import torch.nn.functional as F |
| 11 | from transformers import BatchFeature |
| 12 | |
| 13 | from vllm.config import VllmConfig |
| 14 | from vllm.config.multimodal import MultiModalDummyOptions |
| 15 | from vllm.distributed import get_tensor_model_parallel_world_size |
| 16 | from vllm.inputs import MultiModalDataDict |
| 17 | from vllm.model_executor.layers.activation import get_act_fn |
| 18 | from vllm.model_executor.layers.attention import MMEncoderAttention |
| 19 | from vllm.model_executor.layers.conv import Conv2dLayer |
| 20 | from vllm.model_executor.layers.linear import ( |
| 21 | ColumnParallelLinear, |
| 22 | QKVParallelLinear, |
| 23 | RowParallelLinear, |
| 24 | ) |
| 25 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 26 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 27 | from vllm.multimodal.inputs import ( |
| 28 | MultiModalFieldConfig, |
| 29 | MultiModalKwargsItems, |
| 30 | ) |
| 31 | from vllm.multimodal.parse import ImageSize, MultiModalDataItems |
| 32 | from vllm.multimodal.processing import ( |
| 33 | BaseDummyInputsBuilder, |
| 34 | BaseMultiModalProcessor, |
| 35 | BaseProcessingInfo, |
| 36 | PromptReplacement, |
| 37 | PromptUpdate, |
| 38 | PromptUpdateDetails, |
| 39 | ) |
| 40 | from vllm.sequence import IntermediateTensors |
| 41 | from vllm.transformers_utils.configs.step3_vl import Step3VisionEncoderConfig |
| 42 | from vllm.transformers_utils.processors.step3_vl import ( |
| 43 | MAX_IMAGE_SIZE, |
| 44 | Step3VLImageProcessor, |
| 45 | Step3VLProcessor, |
| 46 | ) |
| 47 | from vllm.utils.tensor_schema import TensorSchema, TensorShape |
| 48 | |
| 49 | from .interfaces import ( |
| 50 | MultiModalEmbeddings, |
| 51 | SupportsEncoderCudaGraph, |
| 52 | SupportsMultiModal, |
| 53 | SupportsPP, |
| 54 | supports_pp, |
| 55 | ) |
| 56 | from .utils import ( |
| 57 | AutoWeightsLoader, |
| 58 | WeightsMapper, |
| 59 | init_vllm_registered_model, |
| 60 | maybe_prefix, |
| 61 | ) |
| 62 | from .vision import is_vit_use_data_parallel, run_dp_sharded_vision_model |
| 63 | |
| 64 | |
| 65 | class Step3VLImagePixelInputs(TensorSchema): |
| 66 | """Dimensions: |
| 67 | - bn: Batch size * number of images |
| 68 | - c: Number of channels (3) |
| 69 | - h: Height |
| 70 | - w: Width |
| 71 | - bnp: Batch size * number of images * number of patches |
| 72 | - hp: Height of patch |
| 73 | - wp: Width of patch |
| 74 | """ |
| 75 | |
| 76 | type: Literal["pixel_values"] |
| 77 | pixel_values: Annotated[torch.Tensor, TensorShape("bn", 3, "h", "w")] |
| 78 | patch_pixel_values: Annotated[torch.Tensor, TensorShape("bnp", 3, "hp", "wp")] |
| 79 | num_patches: Annotated[torch.Tensor, TensorShape("bn")] |
| 80 | |
| 81 | |
| 82 | class Step3VLImageEmbeddingInputs(TensorSchema): |
| 83 | """Dimensions: |
| 84 | - bn: Batch size * number of images |
| 85 | - f: Image feature size |
| 86 | - h: Hidden size (must match the hidden size of language model backbone) |
| 87 | """ |
| 88 | |
| 89 | type: Literal["image_embeds"] = "image_embeds" |
| 90 | data: Annotated[torch.Tensor, TensorShape("bn", "f", "h")] |
| 91 | |
| 92 | |
| 93 | Step3VLImageInputs: TypeAlias = Step3VLImagePixelInputs | Step3VLImageEmbeddingInputs |
| 94 | |
| 95 | |
| 96 | class Step3VLProcessingInfo(BaseProcessingInfo): |
| 97 | def get_image_processor(self, **kwargs): |
| 98 | config = self.get_hf_config() |
| 99 | |
| 100 | kwargs.setdefault( |
| 101 | "enable_patch", |
| 102 | getattr(config.vision_config, "enable_patch", True), |
| 103 | ) |
| 104 | |
| 105 | return Step3VLImageProcessor(**kwargs) |
| 106 | |
| 107 | def get_hf_processor(self, **kwargs: object) -> Step3VLProcessor: |
| 108 | return Step3VLProcessor( |
| 109 | tokenizer=self.get_tokenizer(), |
| 110 | image_processor=self.get_image_processor(**kwargs), |
| 111 | ) |
| 112 | |
| 113 | def get_supported_mm_limits(self) -> Mapping[str, int | None]: |
| 114 | return {"image": None} |
| 115 | |
| 116 | def get_max_image_tokens(self) -> int: |
| 117 | image_processor = self.get_image_processor() |
| 118 | target_width, target_height = self.get_image_size_with_most_features() |
| 119 | |
| 120 | return image_processor.get_num_image_tokens(target_width, target_height) |
| 121 | |
| 122 | def get_mm_max_tokens_per_item( |
| 123 | self, |
| 124 | seq_len: int, |
| 125 | mm_counts: Mapping[str, int], |
| 126 | ) -> Mapping[str, int]: |
| 127 | return {"image": self.get_max_image_tokens()} |
| 128 | |
| 129 | def get_image_size_with_most_features(self) -> ImageSize: |
| 130 | return ImageSize(MAX_IMAGE_SIZE, MAX_IMAGE_SIZE) |
| 131 | |
| 132 | |
| 133 | class Step3VLDummyInputsBuilder(BaseDummyInputsBuilder[Step3VLProcessingInfo]): |
| 134 | def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: |
| 135 | num_images = mm_counts.get("image", 0) |
| 136 | return "<im_patch>" * num_images |
| 137 | |
| 138 | def get_dummy_mm_data( |
| 139 | self, |
| 140 | seq_len: int, |
| 141 | mm_counts: Mapping[str, int], |
| 142 | mm_options: MultiModalDummyOptions, |
| 143 | ) -> MultiModalDataDict: |
| 144 | target_width, target_height = self.info.get_image_size_with_most_features() |
| 145 | |
| 146 | return { |
| 147 | "image": self._get_dummy_images( |
| 148 | width=target_width, |
| 149 | height=target_height, |
| 150 | num_images=mm_counts.get("image", 0), |
| 151 | overrides=mm_options.get("image"), |
| 152 | ) |
| 153 | } |
| 154 | |
| 155 | |
| 156 | class Step3VLMultiModalProcessor(BaseMultiModalProcessor[Step3VLProcessingInfo]): |
| 157 | def _get_prompt_updates( |
| 158 | self, |
| 159 | mm_items: MultiModalDataItems, |
| 160 | hf_processor_mm_kwargs: Mapping[str, Any], |
| 161 | out_mm_kwargs: MultiModalKwargsItems, |
| 162 | ) -> Sequence[PromptUpdate]: |
| 163 | hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) |
| 164 | image_placeholder_token_id = hf_processor.image_token_id |
| 165 | |
| 166 | def get_replacement_step1o(item_idx: int): |
| 167 | out_item = out_mm_kwargs["image"][item_idx] |
| 168 | num_patches_data = out_item["num_patches"].data |
| 169 | patch_newline_mask = out_item["patch_newline_mask"].data |
| 170 | assert isinstance(num_patches_data, torch.Tensor) |
| 171 | assert isinstance(patch_newline_mask, torch.Tensor) |
| 172 | num_patches = int(num_patches_data.item()) |
| 173 | image_repl_ids = hf_processor.get_image_repl_feature_ids( |
| 174 | 1, num_patches, patch_newline_mask.tolist() |
| 175 | ) |
| 176 | |
| 177 | return PromptUpdateDetails.select_token_id( |
| 178 | seq=image_repl_ids, |
| 179 | embed_token_id=image_placeholder_token_id, |
| 180 | ) |
| 181 | |
| 182 | return [ |
| 183 | PromptReplacement( |
| 184 | modality="image", |
| 185 | target=[image_placeholder_token_id], |
| 186 | replacement=get_replacement_step1o, |
| 187 | ) |
| 188 | ] |
| 189 | |
| 190 | def _get_mm_fields_config( |
| 191 | self, |
| 192 | hf_inputs: BatchFeature, |
| 193 | hf_processor_mm_kwargs: Mapping[str, object], |
| 194 | ) -> Mapping[str, MultiModalFieldConfig]: |
| 195 | num_patches = hf_inputs.get("num_patches", torch.empty(0)) |
| 196 | |
| 197 | return dict( |
| 198 | pixel_values=MultiModalFieldConfig.batched("image"), |
| 199 | patch_pixel_values=MultiModalFieldConfig.flat_from_sizes( |
| 200 | "image", num_patches |
| 201 | ), |
| 202 | num_patches=MultiModalFieldConfig.batched("image", keep_on_cpu=True), |
| 203 | patch_newline_mask=MultiModalFieldConfig.flat_from_sizes( |
| 204 | "image", num_patches, keep_on_cpu=True |
| 205 | ), |
| 206 | ) |
| 207 | |
| 208 | |
| 209 | def get_abs_pos(abs_pos, tgt_size): |
| 210 | dim = abs_pos.size(-1) |
| 211 | abs_pos_new = abs_pos.squeeze(0) |
| 212 | cls_token, old_pos_embed = abs_pos_new[:1], abs_pos_new[1:] |
| 213 | |
| 214 | src_size = int(math.sqrt(abs_pos_new.shape[0] - 1)) |
| 215 | tgt_size = int(math.sqrt(tgt_size)) |
| 216 | dtype = abs_pos.dtype |
| 217 | |
| 218 | if src_size != tgt_size: |
| 219 | old_pos_embed = ( |
| 220 | old_pos_embed.view(1, src_size, src_size, dim) |
| 221 | .permute(0, 3, 1, 2) |
| 222 | .contiguous() |
| 223 | ) |
| 224 | old_pos_embed = old_pos_embed.to(torch.float32) |
| 225 | new_pos_embed = F.interpolate( |
| 226 | old_pos_embed, |
| 227 | size=(tgt_size, tgt_size), |
| 228 | mode="bicubic", |
| 229 | antialias=True, |
| 230 | align_corners=False, |
| 231 | ).to(dtype) |
| 232 | new_pos_embed = new_pos_embed.permute(0, 2, 3, 1) |
| 233 | new_pos_embed = new_pos_embed.view(tgt_size * tgt_size, dim) |
| 234 | vision_pos_embed = torch.cat([cls_token, new_pos_embed], dim=0) |
| 235 | vision_pos_embed = vision_pos_embed.view(1, tgt_size * tgt_size + 1, dim) |
| 236 | return vision_pos_embed |
| 237 | else: |
| 238 | return abs_pos |
| 239 | |
| 240 | |
| 241 | class Step3VisionEmbeddings(nn.Module): |
| 242 | def __init__(self, config: Step3VisionEncoderConfig): |
| 243 | super().__init__() |
| 244 | self.config = config |
| 245 | self.embed_dim = config.hidden_size |
| 246 | self.image_size = config.image_size |
| 247 | self.patch_size = config.patch_size |
| 248 | |
| 249 | self.class_embedding = nn.Parameter(torch.randn(1, self.embed_dim)) |
| 250 | |
| 251 | self.patch_embedding = Conv2dLayer( |
| 252 | in_channels=config.num_channels, |
| 253 | out_channels=self.embed_dim, |
| 254 | kernel_size=self.patch_size, |
| 255 | stride=self.patch_size, |
| 256 | bias=True, |
| 257 | ) |
| 258 | |
| 259 | self.num_patches = (self.image_size // self.patch_size) ** 2 |
| 260 | self.pad_tp_size = 4 # hard code for padding |
| 261 | # To load the pretrained weights, we still use P+1 as the seqlen |
| 262 | self.position_embedding = torch.nn.Embedding( |
| 263 | self.num_patches + 1, self.embed_dim |
| 264 | ) |
| 265 | self.register_buffer( |
| 266 | "position_ids", |
| 267 | torch.arange(self.num_patches + 1).expand((1, -1)), |
| 268 | persistent=False, |
| 269 | ) |
| 270 | |
| 271 | def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: |
| 272 | batch_size = pixel_values.shape[0] |
| 273 | patch_embeds = self.patch_embedding( |
| 274 | pixel_values |
| 275 | ) # shape = [*, width, grid, grid] |
| 276 | patch_embeds = patch_embeds.flatten(2).transpose(1, 2) |
| 277 | |
| 278 | # pad |
| 279 | class_embeds = self.class_embedding.expand(batch_size, 1, -1) |
| 280 | embeddings = torch.cat([class_embeds, patch_embeds], dim=1) |
| 281 | embeddings = embeddings + get_abs_pos( |
| 282 | self.position_embedding(self.position_ids), patch_embeds.size(1) |
| 283 | ) |
| 284 | embeddings = torch.cat( |
| 285 | [ |
| 286 | embeddings[:, 0, :].unsqueeze(1).repeat(1, self.pad_tp_size - 1, 1), |
| 287 | embeddings, |
| 288 | ], |
| 289 | dim=1, |
| 290 | ) |
| 291 | return embeddings |
| 292 | |
| 293 | |
| 294 | class Step3VisionAttention(nn.Module): |
| 295 | """Multi-headed attention from 'Attention Is All You Need' paper.""" |
| 296 | |
| 297 | def __init__( |
| 298 | self, |
| 299 | config, |
| 300 | quant_config: QuantizationConfig | None = None, |
| 301 | prefix: str = "", |
| 302 | ): |
| 303 | super().__init__() |
| 304 | self.config = config |
| 305 | self.embed_dim = config.hidden_size |
| 306 | self.total_num_heads = config.num_attention_heads |
| 307 | self.head_dim = self.embed_dim // self.total_num_heads |
| 308 | |
| 309 | self.scale = self.head_dim**-0.5 |
| 310 | |
| 311 | use_data_parallel = is_vit_use_data_parallel() |
| 312 | tp_size = 1 if use_data_parallel else get_tensor_model_parallel_world_size() |
| 313 | assert self.total_num_heads % tp_size == 0 |
| 314 | self.num_heads = self.total_num_heads // tp_size |
| 315 | |
| 316 | self.q_size = self.num_heads * self.head_dim |
| 317 | |
| 318 | self.qkv_proj = QKVParallelLinear( |
| 319 | self.embed_dim, |
| 320 | self.head_dim, |
| 321 | self.total_num_heads, |
| 322 | bias=True, |
| 323 | quant_config=quant_config, |
| 324 | prefix=f"{prefix}.qkv_proj", |
| 325 | disable_tp=use_data_parallel, |
| 326 | ) |
| 327 | self.out_proj = RowParallelLinear( |
| 328 | self.embed_dim, |
| 329 | self.embed_dim, |
| 330 | bias=True, |
| 331 | quant_config=quant_config, |
| 332 | prefix=f"{prefix}.out_proj", |
| 333 | disable_tp=use_data_parallel, |
| 334 | ) |
| 335 | |
| 336 | # Use unified MMEncoderAttention with automatic backend selection |
| 337 | self.attn = MMEncoderAttention( |
| 338 | self.num_heads, |
| 339 | self.head_dim, |
| 340 | self.scale, |
| 341 | prefix=f"{prefix}.attn", |
| 342 | ) |
| 343 | |
| 344 | def forward( |
| 345 | self, |
| 346 | hidden_states: torch.Tensor, |
| 347 | ): |
| 348 | """Input shape: Batch x Time x Channel.""" |
| 349 | bsz, tgt_len, _ = hidden_states.size() |
| 350 | |
| 351 | # get query proj |
| 352 | qkv, _ = self.qkv_proj(hidden_states) |
| 353 | q, k, v = qkv.chunk(chunks=3, dim=-1) |
| 354 | |
| 355 | # Use unified MMEncoderAttention with automatic backend selection |
| 356 | attn_output = self.attn(q, k, v) |
| 357 | |
| 358 | attn_output, _ = self.out_proj(attn_output) |
| 359 | |
| 360 | return attn_output |
| 361 | |
| 362 | |
| 363 | class Step3VisionMLP(nn.Module): |
| 364 | def __init__( |
| 365 | self, |
| 366 | config, |
| 367 | quant_config: QuantizationConfig | None = None, |
| 368 | prefix: str = "", |
| 369 | ): |
| 370 | super().__init__() |
| 371 | self.config = config |
| 372 | self.activation_fn = get_act_fn(config.hidden_act) |
| 373 | use_data_parallel = is_vit_use_data_parallel() |
| 374 | self.fc1 = ColumnParallelLinear( |
| 375 | config.hidden_size, |
| 376 | config.intermediate_size, |
| 377 | bias=True, |
| 378 | quant_config=quant_config, |
| 379 | prefix=f"{prefix}.fc1", |
| 380 | disable_tp=use_data_parallel, |
| 381 | ) |
| 382 | self.fc2 = RowParallelLinear( |
| 383 | config.intermediate_size, |
| 384 | config.hidden_size, |
| 385 | bias=True, |
| 386 | quant_config=quant_config, |
| 387 | prefix=f"{prefix}.fc2", |
| 388 | disable_tp=use_data_parallel, |
| 389 | ) |
| 390 | |
| 391 | def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 392 | hidden_states, _ = self.fc1(hidden_states) |
| 393 | hidden_states = self.activation_fn(hidden_states) |
| 394 | hidden_states, _ = self.fc2(hidden_states) |
| 395 | return hidden_states |
| 396 | |
| 397 | |
| 398 | class Step3VisionEncoderLayer(nn.Module): |
| 399 | def __init__( |
| 400 | self, |
| 401 | config: Step3VisionEncoderConfig, |
| 402 | quant_config: QuantizationConfig | None = None, |
| 403 | prefix: str = "", |
| 404 | ): |
| 405 | super().__init__() |
| 406 | self.embed_dim = config.hidden_size |
| 407 | self.self_attn = Step3VisionAttention( |
| 408 | config, |
| 409 | quant_config, |
| 410 | prefix=f"{prefix}.self_attn", |
| 411 | ) |
| 412 | self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) |
| 413 | self.mlp = Step3VisionMLP( |
| 414 | config, |
| 415 | quant_config, |
| 416 | prefix=f"{prefix}.mlp", |
| 417 | ) |
| 418 | self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) |
| 419 | |
| 420 | def forward( |
| 421 | self, |
| 422 | hidden_states: torch.Tensor, |
| 423 | ) -> torch.FloatTensor: |
| 424 | hidden_states = hidden_states + self.layer_norm1(self.self_attn(hidden_states)) |
| 425 | hidden_states = hidden_states + self.layer_norm2(self.mlp(hidden_states)) |
| 426 | return hidden_states |
| 427 | |
| 428 | |
| 429 | class Step3VisionEncoder(nn.Module): |
| 430 | def __init__( |
| 431 | self, |
| 432 | config: Step3VisionEncoderConfig, |
| 433 | quant_config: QuantizationConfig | None = None, |
| 434 | prefix: str = "", |
| 435 | ): |
| 436 | super().__init__() |
| 437 | self.config = config |
| 438 | self.layers = nn.ModuleList( |
| 439 | [ |
| 440 | Step3VisionEncoderLayer( |
| 441 | config, |
| 442 | quant_config, |
| 443 | prefix=f"{prefix}.layers.{i}", |
| 444 | ) |
| 445 | for i in range(config.num_hidden_layers) |
| 446 | ] |
| 447 | ) |
| 448 | |
| 449 | def forward( |
| 450 | self, |
| 451 | inputs_embeds, |
| 452 | ): |
| 453 | hidden_states = inputs_embeds |
| 454 | for encoder_layer in self.layers: |
| 455 | hidden_states = encoder_layer(hidden_states) |
| 456 | return hidden_states |
| 457 | |
| 458 | |
| 459 | class Step3VisionTransformer(nn.Module): |
| 460 | def __init__( |
| 461 | self, |
| 462 | config: Step3VisionEncoderConfig, |
| 463 | quant_config: QuantizationConfig | None = None, |
| 464 | prefix: str = "", |
| 465 | ): |
| 466 | super().__init__() |
| 467 | self.config = config |
| 468 | self.use_data_parallel = is_vit_use_data_parallel() |
| 469 | self.image_size = config.image_size |
| 470 | self.embeddings = Step3VisionEmbeddings(config) |
| 471 | self.transformer = Step3VisionEncoder( |
| 472 | config, |
| 473 | quant_config, |
| 474 | prefix=f"{prefix}.transformer", |
| 475 | ) |
| 476 | |
| 477 | def forward( |
| 478 | self, |
| 479 | pixel_values: torch.Tensor, |
| 480 | ): |
| 481 | hidden_states = self.embeddings(pixel_values) |
| 482 | if self.use_data_parallel: |
| 483 | hidden_states = run_dp_sharded_vision_model(hidden_states, self.transformer) |
| 484 | else: |
| 485 | hidden_states = self.transformer(inputs_embeds=hidden_states) |
| 486 | return hidden_states |
| 487 | |
| 488 | |
| 489 | @MULTIMODAL_REGISTRY.register_processor( |
| 490 | Step3VLMultiModalProcessor, |
| 491 | info=Step3VLProcessingInfo, |
| 492 | dummy_inputs=Step3VLDummyInputsBuilder, |
| 493 | ) |
| 494 | class Step3VLForConditionalGeneration( |
| 495 | nn.Module, SupportsMultiModal, SupportsPP, SupportsEncoderCudaGraph |
| 496 | ): |
| 497 | hf_to_vllm_mapper = WeightsMapper( |
| 498 | orig_to_new_prefix={ |
| 499 | "model.": "language_model.model.", |
| 500 | "lm_head.": "language_model.lm_head.", |
| 501 | } |
| 502 | ) |
| 503 | |
| 504 | supports_encoder_tp_data = True |
| 505 | |
| 506 | @classmethod |
| 507 | def get_placeholder_str(cls, modality: str, i: int) -> str | None: |
| 508 | if modality.startswith("image"): |
| 509 | return "<im_patch>" |
| 510 | |
| 511 | raise ValueError("Only image modality is supported") |
| 512 | |
| 513 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: |
| 514 | super().__init__() |
| 515 | config = vllm_config.model_config.hf_config |
| 516 | multimodal_config = vllm_config.model_config.get_multimodal_config() |
| 517 | |
| 518 | self.config = config |
| 519 | self.model_config = vllm_config.model_config |
| 520 | self.multimodal_config = multimodal_config |
| 521 | self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" |
| 522 | |
| 523 | # NOTE: This behavior is consistent with the previous OOV handling, |
| 524 | # but does not currently handle the start/stop toks around the |
| 525 | # image features (<patch_start> <patch_end> <im_start> <im_end>) |
| 526 | # See: https://huggingface.co/stepfun-ai/step3/blob/main/processing_step3v.py#L323 |
| 527 | # |
| 528 | # If this becomes an issue or we refactor to handle this using the |
| 529 | # processor info in the future, it would probably be best to handle |
| 530 | # those too. |
| 531 | self.configure_mm_token_handling( |
| 532 | self.config.text_config.vocab_size, |
| 533 | [self.config.image_token_id], |
| 534 | ) |
| 535 | |
| 536 | with self._mark_tower_model(vllm_config, "image"): |
| 537 | self.vision_model = Step3VisionTransformer( |
| 538 | config.vision_config, |
| 539 | None, |
| 540 | prefix=maybe_prefix(prefix, "vision_model"), |
| 541 | ) |
| 542 | self.vit_downsampler = Conv2dLayer( |
| 543 | config.vision_config.hidden_size, |
| 544 | config.vision_config.output_hidden_size, |
| 545 | kernel_size=2, |
| 546 | stride=config.understand_projector_stride, |
| 547 | ) |
| 548 | self.vit_downsampler2 = Conv2dLayer( |
| 549 | config.vision_config.output_hidden_size, |
| 550 | config.vision_config.output_hidden_size * 2, |
| 551 | kernel_size=3, |
| 552 | stride=2, |
| 553 | padding=1, |
| 554 | ) |
| 555 | self.vit_large_projector = nn.Linear( |
| 556 | config.vision_config.output_hidden_size * 2, |
| 557 | config.hidden_size, |
| 558 | bias=config.projector_bias, |
| 559 | ) |
| 560 | |
| 561 | with self._mark_language_model(vllm_config): |
| 562 | language_model = init_vllm_registered_model( |
| 563 | vllm_config=vllm_config, |
| 564 | hf_config=config.text_config, |
| 565 | prefix=maybe_prefix(prefix, "language_model"), |
| 566 | ) |
| 567 | self.language_model = language_model |
| 568 | assert supports_pp(language_model) |
| 569 | |
| 570 | self.make_empty_intermediate_tensors = ( |
| 571 | language_model.make_empty_intermediate_tensors |
| 572 | ) |
| 573 | |
| 574 | @property |
| 575 | def device(self): |
| 576 | return next(self.parameters()).device |
| 577 | |
| 578 | @property |
| 579 | def dtype(self): |
| 580 | return next(self.parameters()).dtype |
| 581 | |
| 582 | @staticmethod |
| 583 | def _compute_spatial_tokens(size, patch_size, stride): |
| 584 | # Compute the number of spatial tokens after two rounds of |
| 585 | # downsampling with given patch size and stride. |
| 586 | grid = size // patch_size |
| 587 | vit_tokens = grid * grid |
| 588 | spatial = int(math.sqrt(vit_tokens)) |
| 589 | h1 = (spatial - 2) // stride + 1 |
| 590 | h2 = (h1 - 1) // 2 + 1 |
| 591 | return h2 * h2 |
| 592 | |
| 593 | @property |
| 594 | def img_output_tokens(self) -> int: |
| 595 | return self._compute_spatial_tokens( |
| 596 | self.config.vision_config.image_size, |
| 597 | self.config.vision_config.patch_size, |
| 598 | self.config.understand_projector_stride, |
| 599 | ) |
| 600 | |
| 601 | @property |
| 602 | def patch_output_tokens(self) -> int: |
| 603 | return self._compute_spatial_tokens( |
| 604 | 504, |
| 605 | self.config.vision_config.patch_size, |
| 606 | self.config.understand_projector_stride, |
| 607 | ) |
| 608 | |
| 609 | def _batched_encoder_forward( |
| 610 | self, |
| 611 | pixel_values: torch.Tensor, |
| 612 | ) -> torch.Tensor: |
| 613 | image_features = self._process_image_features( |
| 614 | self._get_vision_model_output(pixel_values) |
| 615 | ) |
| 616 | return image_features.reshape(-1, image_features.shape[-1]) |
| 617 | |
| 618 | def _parse_and_validate_image_input( |
| 619 | self, **kwargs: object |
| 620 | ) -> Step3VLImageInputs | None: |
| 621 | pixel_values = kwargs.pop("pixel_values", None) |
| 622 | patch_pixel_values = kwargs.pop("patch_pixel_values", None) |
| 623 | num_patches = kwargs.pop("num_patches", None) |
| 624 | image_embeds = kwargs.pop("image_embeds", None) |
| 625 | |
| 626 | if pixel_values is None and image_embeds is None: |
| 627 | return None |
| 628 | |
| 629 | if pixel_values is not None and patch_pixel_values is not None: |
| 630 | assert isinstance(pixel_values, torch.Tensor) |
| 631 | assert isinstance(patch_pixel_values, torch.Tensor) |
| 632 | assert isinstance(num_patches, torch.Tensor) |
| 633 | return Step3VLImagePixelInputs( |
| 634 | type="pixel_values", |
| 635 | pixel_values=pixel_values.to(self.dtype), |
| 636 | patch_pixel_values=patch_pixel_values.to(self.dtype), |
| 637 | num_patches=num_patches, |
| 638 | ) |
| 639 | |
| 640 | if image_embeds is not None: |
| 641 | assert isinstance(image_embeds, torch.Tensor) |
| 642 | return Step3VLImageEmbeddingInputs( |
| 643 | type="image_embeds", |
| 644 | data=image_embeds.to(self.dtype), |
| 645 | ) |
| 646 | |
| 647 | raise AssertionError("This line should be unreachable.") |
| 648 | |
| 649 | def _process_image_features(self, image_features: torch.Tensor) -> torch.Tensor: |
| 650 | B, P = image_features.shape[:2] |
| 651 | HW = int(sqrt(P)) |
| 652 | image_features = image_features.permute(0, 2, 1).view(B, -1, HW, HW) |
| 653 | image_features = self.vit_downsampler(image_features) |
| 654 | image_features = self.vit_downsampler2(image_features) |
| 655 | n_dim = image_features.size(1) |
| 656 | image_features = image_features.view(B, n_dim, -1).permute(0, 2, 1) |
| 657 | image_features = self.vit_large_projector(image_features) |
| 658 | return image_features |
| 659 | |
| 660 | def _get_vision_model_output(self, input_tensor: torch.Tensor) -> torch.Tensor: |
| 661 | return self.vision_model(input_tensor)[:, 4:] |
| 662 | |
| 663 | def _process_image_input( |
| 664 | self, image_input: Step3VLImageInputs |
| 665 | ) -> list[torch.Tensor]: |
| 666 | if image_input["type"] == "image_embeds": |
| 667 | image_features = image_input["data"] |
| 668 | return [ |
| 669 | image_features[i].view(-1, image_features.shape[-1]) |
| 670 | for i in range(image_features.shape[0]) |
| 671 | ] |
| 672 | |
| 673 | image_features = self._get_vision_model_output(image_input["pixel_values"]) |
| 674 | patch_image_features = ( |
| 675 | self._get_vision_model_output(image_input["patch_pixel_values"]) |
| 676 | if len(image_input["patch_pixel_values"]) > 0 |
| 677 | else None |
| 678 | ) |
| 679 | num_patches = image_input["num_patches"] |
| 680 | |
| 681 | image_features = self._process_image_features(image_features) |
| 682 | patch_image_features = ( |
| 683 | self._process_image_features(patch_image_features) |
| 684 | if patch_image_features is not None |
| 685 | else None |
| 686 | ) |
| 687 | |
| 688 | merged_image_features = [] |
| 689 | cur_patch_idx = 0 |
| 690 | num_patches_list = num_patches.tolist() |
| 691 | for i, num_patch in enumerate(num_patches_list): |
| 692 | cur_feature = [] |
| 693 | if num_patch > 0: |
| 694 | assert patch_image_features is not None |
| 695 | patch_slice = patch_image_features[ |
| 696 | cur_patch_idx : cur_patch_idx + num_patch |
| 697 | ] |
| 698 | cur_feature.append(patch_slice.view(-1, patch_slice.shape[-1])) |
| 699 | cur_feature.append(image_features[i].view(-1, image_features.shape[-1])) |
| 700 | cur_patch_idx += num_patch |
| 701 | merged_image_features.append( |
| 702 | torch.cat(cur_feature) if len(cur_feature) > 1 else cur_feature[0] |
| 703 | ) |
| 704 | return merged_image_features |
| 705 | |
| 706 | def embed_multimodal(self, **kwargs) -> MultiModalEmbeddings: |
| 707 | image_input = self._parse_and_validate_image_input(**kwargs) |
| 708 | if image_input is None: |
| 709 | return [] |
| 710 | vision_embeddings = self._process_image_input(image_input) |
| 711 | return vision_embeddings |
| 712 | |
| 713 | def embed_input_ids( |
| 714 | self, |
| 715 | input_ids: torch.Tensor, |
| 716 | multimodal_embeddings: MultiModalEmbeddings | None = None, |
| 717 | *, |
| 718 | is_multimodal: torch.Tensor | None = None, |
| 719 | ) -> torch.Tensor: |
| 720 | # This is to satisfy the type checker for each overload |
| 721 | if multimodal_embeddings is None or is_multimodal is None: |
| 722 | return super().embed_input_ids(input_ids) |
| 723 | |
| 724 | return super().embed_input_ids( |
| 725 | input_ids, |
| 726 | multimodal_embeddings=multimodal_embeddings, |
| 727 | is_multimodal=is_multimodal, |
| 728 | ) |
| 729 | |
| 730 | # -- SupportsEncoderCudaGraph protocol methods -- |
| 731 | |
| 732 | def get_encoder_cudagraph_config(self): |
| 733 | from vllm.v1.worker.encoder_cudagraph_defs import ( |
| 734 | EncoderCudaGraphConfig, |
| 735 | EncoderCudaGraphPathConfig, |
| 736 | ) |
| 737 | |
| 738 | return EncoderCudaGraphConfig( |
| 739 | modalities=["image"], |
| 740 | buffer_keys=[ |
| 741 | "pixel_values", |
| 742 | "patch_pixel_values", |
| 743 | ], |
| 744 | out_hidden_size=self.config.hidden_size, |
| 745 | paths={ |
| 746 | "global": EncoderCudaGraphPathConfig( |
| 747 | min_token_budget=self.img_output_tokens |
| 748 | ), |
| 749 | "local": EncoderCudaGraphPathConfig( |
| 750 | min_token_budget=self.patch_output_tokens, |
| 751 | allow_zero_tokens=True, |
| 752 | ), |
| 753 | }, |
| 754 | ) |
| 755 | |
| 756 | def get_encoder_cudagraph_budget_range( |
| 757 | self, |
| 758 | vllm_config: "VllmConfig", |
| 759 | ) -> tuple[int, int]: |
| 760 | min_budget = self.img_output_tokens |
| 761 | max_budget = min( |
| 762 | vllm_config.scheduler_config.max_num_batched_tokens, |
| 763 | self.model_config.max_model_len, |
| 764 | ) |
| 765 | return min_budget, max_budget |
| 766 | |
| 767 | def get_encoder_cudagraph_item_specs( |
| 768 | self, |
| 769 | mm_kwargs: dict[str, Any], |
| 770 | ): |
| 771 | from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec |
| 772 | |
| 773 | num_patches = mm_kwargs.get("num_patches") |
| 774 | assert isinstance(num_patches, torch.Tensor) |
| 775 | |
| 776 | img_grid = ( |
| 777 | self.config.vision_config.image_size // self.config.vision_config.patch_size |
| 778 | ) |
| 779 | patch_grid = 504 // self.config.vision_config.patch_size |
| 780 | total_image_pixel = img_grid * img_grid |
| 781 | total_patch_pixel = patch_grid * patch_grid |
| 782 | |
| 783 | return [ |
| 784 | EncoderItemSpec( |
| 785 | input_size=(total_image_pixel + num_patch * total_patch_pixel), |
| 786 | output_tokens=( |
| 787 | self.img_output_tokens + num_patch * self.patch_output_tokens |
| 788 | ), |
| 789 | path_output_tokens={ |
| 790 | "global": self.img_output_tokens, |
| 791 | "local": num_patch * self.patch_output_tokens, |
| 792 | }, |
| 793 | ) |
| 794 | for num_patch in num_patches |
| 795 | ] |
| 796 | |
| 797 | def select_encoder_cudagraph_items( |
| 798 | self, |
| 799 | mm_kwargs: dict[str, Any], |
| 800 | indices: list[int], |
| 801 | ) -> dict[str, Any]: |
| 802 | pixel_values = mm_kwargs["pixel_values"] |
| 803 | patch_pixel_values = mm_kwargs["patch_pixel_values"] |
| 804 | num_patches = mm_kwargs["num_patches"] |
| 805 | |
| 806 | # calcute the accumulated patch counts |
| 807 | cum_patches = [0] |
| 808 | for p in num_patches: |
| 809 | cum_patches.append(cum_patches[-1] + p) |
| 810 | |
| 811 | if len(indices) == 0: |
| 812 | return { |
| 813 | "pixel_values": pixel_values[:0], |
| 814 | "patch_pixel_values": patch_pixel_values[:0], |
| 815 | "num_patches": num_patches[:0], |
| 816 | } |
| 817 | |
| 818 | selected_pv = pixel_values[indices] |
| 819 | selected_np = num_patches[indices] |
| 820 | selected_ppv = torch.cat( |
| 821 | [patch_pixel_values[cum_patches[i] : cum_patches[i + 1]] for i in indices] |
| 822 | ) |
| 823 | |
| 824 | return { |
| 825 | "pixel_values": selected_pv, |
| 826 | "patch_pixel_values": selected_ppv, |
| 827 | "num_patches": selected_np, |
| 828 | } |
| 829 | |
| 830 | def prepare_encoder_cudagraph_capture_inputs( |
| 831 | self, |
| 832 | token_budget: int, |
| 833 | max_batch_size: int, |
| 834 | max_frames_per_batch: int, |
| 835 | device: torch.device, |
| 836 | dtype: torch.dtype, |
| 837 | path: str = "default", |
| 838 | axis_keys: tuple[Hashable, ...] | None = None, |
| 839 | ): |
| 840 | from vllm.v1.worker.encoder_cudagraph_defs import ( |
| 841 | EncoderCudaGraphCaptureInputs, |
| 842 | ) |
| 843 | |
| 844 | assert path in ("global", "local") |
| 845 | if path == "global": |
| 846 | max_num_images = token_budget // self.img_output_tokens |
| 847 | max_batch_size = min(max_batch_size, max_num_images) |
| 848 | dummy_pixel_values = torch.randn( |
| 849 | max_batch_size, |
| 850 | 3, |
| 851 | self.config.vision_config.image_size, |
| 852 | self.config.vision_config.image_size, |
| 853 | device=device, |
| 854 | dtype=dtype, |
| 855 | ) |
| 856 | values = {"pixel_values": dummy_pixel_values} |
| 857 | else: |
| 858 | max_num_patches = token_budget // self.patch_output_tokens |
| 859 | dummy_patch_pixel_values = torch.randn( |
| 860 | max_num_patches, |
| 861 | 3, |
| 862 | 504, |
| 863 | 504, |
| 864 | device=device, |
| 865 | dtype=dtype, |
| 866 | ) |
| 867 | values = {"patch_pixel_values": dummy_patch_pixel_values} |
| 868 | |
| 869 | return EncoderCudaGraphCaptureInputs( |
| 870 | values=values, |
| 871 | ) |
| 872 | |
| 873 | def encoder_cudagraph_forward( |
| 874 | self, |
| 875 | values: dict[str, torch.Tensor], |
| 876 | path: str = "default", |
| 877 | ) -> torch.Tensor: |
| 878 | assert path in ("global", "local") |
| 879 | if path == "global": |
| 880 | return self._batched_encoder_forward(values["pixel_values"]) |
| 881 | else: |
| 882 | return self._batched_encoder_forward(values["patch_pixel_values"]) |
| 883 | |
| 884 | def encoder_eager_forward( |
| 885 | self, |
| 886 | mm_kwargs: dict[str, Any], |
| 887 | path: str = "default", |
| 888 | ) -> torch.Tensor: |
| 889 | assert path in ("global", "local") |
| 890 | if path == "global": |
| 891 | return self._batched_encoder_forward(mm_kwargs["pixel_values"]) |
| 892 | else: |
| 893 | return self._batched_encoder_forward(mm_kwargs["patch_pixel_values"]) |
| 894 | |
| 895 | def postprocess_encoder_output( |
| 896 | self, |
| 897 | outputs: dict[str, torch.Tensor], |
| 898 | indices: list[int], |
| 899 | per_item_out_tokens: list[int], |
| 900 | dest: dict[int, torch.Tensor] | list[torch.Tensor | None], |
| 901 | clone: bool = False, |
| 902 | batch_mm_kwargs: dict[str, Any] | None = None, |
| 903 | ): |
| 904 | """CPU-side per-item merge after dual-path graph replay. |
| 905 | |
| 906 | ``outputs['global']`` contains global-image features and ``outputs['local']`` |
| 907 | contains local-patch features (or ``None`` when there are no patches). |
| 908 | """ |
| 909 | output = outputs["global"] |
| 910 | local_output = outputs.get("local") |
| 911 | assert batch_mm_kwargs is not None |
| 912 | num_patches = batch_mm_kwargs["num_patches"] |
| 913 | hidden = output.shape[-1] |
| 914 | bsz = len(indices) |
| 915 | |
| 916 | actual_np = [int(np) for np in num_patches] |
| 917 | total_patches = sum(actual_np) |
| 918 | img_tokens = bsz * self.img_output_tokens |
| 919 | patch_tokens = total_patches * self.patch_output_tokens |
| 920 | |
| 921 | global_part = output[:img_tokens].reshape(bsz, self.img_output_tokens, hidden) |
| 922 | if total_patches > 0 and local_output is not None: |
| 923 | patch_part = local_output[:patch_tokens].reshape( |
| 924 | -1, self.patch_output_tokens, hidden |
| 925 | ) |
| 926 | else: |
| 927 | patch_part = None |
| 928 | |
| 929 | merged: dict[int, torch.Tensor] = {} |
| 930 | cur_patch = 0 |
| 931 | for i, idx in enumerate(indices): |
| 932 | np = actual_np[i] |
| 933 | parts: list[torch.Tensor] = [] |
| 934 | if patch_part is not None and np > 0: |
| 935 | parts.append(patch_part[cur_patch : cur_patch + np].reshape(-1, hidden)) |
| 936 | cur_patch += np |
| 937 | parts.append(global_part[i].reshape(-1, hidden)) |
| 938 | merged[idx] = torch.cat(parts, dim=0) if len(parts) > 1 else parts[0] |
| 939 | |
| 940 | out = [merged[i] for i in indices] |
| 941 | for i, idx in enumerate(indices): |
| 942 | dest[idx] = out[i] |
| 943 | |
| 944 | def prepare_encoder_cudagraph_replay_buffers( |
| 945 | self, |
| 946 | mm_kwargs: dict[str, Any], |
| 947 | max_batch_size: int, |
| 948 | max_frames_per_batch: int, |
| 949 | path: str = "default", |
| 950 | ): |
| 951 | from vllm.v1.worker.encoder_cudagraph_defs import ( |
| 952 | EncoderCudaGraphReplayBuffers, |
| 953 | ) |
| 954 | |
| 955 | assert path in ("global", "local") |
| 956 | if path == "global": |
| 957 | values = {"pixel_values": mm_kwargs["pixel_values"]} |
| 958 | else: |
| 959 | values = {"patch_pixel_values": mm_kwargs["patch_pixel_values"]} |
| 960 | |
| 961 | return EncoderCudaGraphReplayBuffers(values=values) |
| 962 | |
| 963 | def forward( |
| 964 | self, |
| 965 | input_ids: torch.Tensor | None, |
| 966 | positions: torch.Tensor, |
| 967 | intermediate_tensors: IntermediateTensors | None = None, |
| 968 | inputs_embeds: torch.Tensor | None = None, |
| 969 | **kwargs: object, |
| 970 | ) -> torch.Tensor | IntermediateTensors: |
| 971 | if intermediate_tensors is not None: |
| 972 | inputs_embeds = None |
| 973 | |
| 974 | hidden_states = self.language_model( |
| 975 | input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds |
| 976 | ) |
| 977 | |
| 978 | return hidden_states |
| 979 | |
| 980 | def compute_logits( |
| 981 | self, |
| 982 | hidden_states: torch.Tensor, |
| 983 | ) -> torch.Tensor | None: |
| 984 | return self.language_model.compute_logits(hidden_states) |
| 985 | |
| 986 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): |
| 987 | loader = AutoWeightsLoader(self) |
| 988 | return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) |