| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # Copyright 2025 The Baidu team. |
| 5 | # Copyright 2023 The vLLM team. |
| 6 | # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. |
| 7 | # |
| 8 | # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX |
| 9 | # and OPT implementations in this library. It has been modified from its |
| 10 | # original forms to accommodate minor architectural differences compared |
| 11 | # to GPT-NeoX and OPT used by the Meta AI team that trained the model. |
| 12 | # |
| 13 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 14 | # you may not use this file except in compliance with the License. |
| 15 | # You may obtain a copy of the License at |
| 16 | # |
| 17 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 18 | # |
| 19 | # Unless required by applicable law or agreed to in writing, software |
| 20 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 21 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 22 | # See the License for the specific language governing permissions and |
| 23 | # limitations under the License. |
| 24 | """Inference-only Ernie VL model compatible with HuggingFace weights.""" |
| 25 | |
| 26 | import math |
| 27 | from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence |
| 28 | from functools import partial |
| 29 | from typing import Annotated, Any, Literal |
| 30 | |
| 31 | import numpy as np |
| 32 | import torch |
| 33 | import torch.nn as nn |
| 34 | import torch.nn.functional as F |
| 35 | from einops import rearrange |
| 36 | from transformers import BaseImageProcessor, BatchFeature |
| 37 | |
| 38 | from vllm.config import VllmConfig |
| 39 | from vllm.config.multimodal import MultiModalDummyOptions, VideoDummyOptions |
| 40 | from vllm.distributed import parallel_state |
| 41 | from vllm.distributed import utils as dist_utils |
| 42 | from vllm.inputs import MultiModalDataDict |
| 43 | from vllm.logger import init_logger |
| 44 | from vllm.model_executor.layers.activation import QuickGELU |
| 45 | from vllm.model_executor.layers.attention import ( |
| 46 | MMEncoderAttention, |
| 47 | ) |
| 48 | from vllm.model_executor.layers.layernorm import RMSNorm |
| 49 | from vllm.model_executor.layers.linear import ( |
| 50 | ColumnParallelLinear, |
| 51 | QKVParallelLinear, |
| 52 | RowParallelLinear, |
| 53 | ) |
| 54 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 55 | from vllm.model_executor.layers.rotary_embedding.common import ( |
| 56 | ApplyRotaryEmb, |
| 57 | ) |
| 58 | from vllm.model_executor.model_loader.weight_utils import default_weight_loader |
| 59 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 60 | from vllm.multimodal.inputs import ( |
| 61 | MultiModalFeatureSpec, |
| 62 | MultiModalFieldConfig, |
| 63 | MultiModalKwargsItems, |
| 64 | ) |
| 65 | from vllm.multimodal.parse import ImageSize, MultiModalDataItems, MultiModalDataParser |
| 66 | from vllm.multimodal.processing import ( |
| 67 | BaseDummyInputsBuilder, |
| 68 | BaseMultiModalProcessor, |
| 69 | BaseProcessingInfo, |
| 70 | PromptReplacement, |
| 71 | PromptUpdate, |
| 72 | cached_encode, |
| 73 | ) |
| 74 | from vllm.multimodal.processing.processor import HFMultiModalInputs |
| 75 | from vllm.sequence import IntermediateTensors |
| 76 | from vllm.utils.gpu_sync_debug import gpu_sync_allowed |
| 77 | from vllm.utils.tensor_schema import TensorSchema, TensorShape |
| 78 | from vllm.utils.torch_utils import async_tensor_h2d |
| 79 | from vllm.v1.attention.backends.registry import AttentionBackendEnum |
| 80 | |
| 81 | from .ernie45_vl_moe import Ernie4_5_VLMoeForCausalLM |
| 82 | from .interfaces import ( |
| 83 | MultiModalEmbeddings, |
| 84 | SupportsEncoderCudaGraph, |
| 85 | SupportsLoRA, |
| 86 | SupportsMRoPE, |
| 87 | SupportsMultiModal, |
| 88 | SupportsPP, |
| 89 | ) |
| 90 | from .utils import ( |
| 91 | AutoWeightsLoader, |
| 92 | WeightsMapper, |
| 93 | maybe_prefix, |
| 94 | scatter_output_slices, |
| 95 | ) |
| 96 | from .vision import get_vit_attn_backend |
| 97 | |
| 98 | logger = init_logger(__name__) |
| 99 | |
| 100 | # === Vision Transformer === # |
| 101 | |
| 102 | |
| 103 | def all_gather_interleave(local_tensor, hidden_size: int, tp_size: int): |
| 104 | """All-gather the input tensor interleavely across model parallel group.""" |
| 105 | import torch.distributed as dist |
| 106 | |
| 107 | gathered_tensors = [torch.zeros_like(local_tensor) for _ in range(tp_size)] |
| 108 | dist.all_gather( |
| 109 | gathered_tensors, local_tensor, group=parallel_state.get_tp_group().device_group |
| 110 | ) |
| 111 | |
| 112 | gathered_tensors_split = [ |
| 113 | torch.split(tensor, hidden_size // tp_size, -1) for tensor in gathered_tensors |
| 114 | ] |
| 115 | ordered_tensors = [ |
| 116 | tensor for pair in zip(*gathered_tensors_split) for tensor in pair |
| 117 | ] |
| 118 | result_tensor = torch.cat(ordered_tensors, dim=-1) |
| 119 | return result_tensor |
| 120 | |
| 121 | |
| 122 | class Ernie4_5_VisionAttention(nn.Module): |
| 123 | """VisionAttention using VLLM framework APIs.""" |
| 124 | |
| 125 | def __init__( |
| 126 | self, |
| 127 | embed_dim: int, |
| 128 | num_heads: int, |
| 129 | projection_size: int, |
| 130 | quant_config: QuantizationConfig | None = None, |
| 131 | prefix: str = "", |
| 132 | ) -> None: |
| 133 | super().__init__() |
| 134 | # Per attention head and per partition values. |
| 135 | self.tp_size = parallel_state.get_tensor_model_parallel_world_size() |
| 136 | self.tp_rank = parallel_state.get_tensor_model_parallel_rank() |
| 137 | self.hidden_size_per_attention_head = dist_utils.divide( |
| 138 | projection_size, num_heads |
| 139 | ) |
| 140 | self.num_attention_heads_per_partition = dist_utils.divide( |
| 141 | num_heads, self.tp_size |
| 142 | ) |
| 143 | |
| 144 | self.qkv = QKVParallelLinear( |
| 145 | hidden_size=embed_dim, |
| 146 | head_size=self.hidden_size_per_attention_head, |
| 147 | total_num_heads=num_heads, |
| 148 | total_num_kv_heads=num_heads, |
| 149 | bias=True, |
| 150 | quant_config=quant_config, |
| 151 | prefix=f"{prefix}.qkv", |
| 152 | ) |
| 153 | self.proj = RowParallelLinear( |
| 154 | input_size=projection_size, |
| 155 | output_size=embed_dim, |
| 156 | quant_config=quant_config, |
| 157 | prefix=f"{prefix}.proj", |
| 158 | ) |
| 159 | |
| 160 | self.attn = MMEncoderAttention( |
| 161 | num_heads=self.num_attention_heads_per_partition, |
| 162 | head_size=self.hidden_size_per_attention_head, |
| 163 | scale=self.hidden_size_per_attention_head**-0.5, |
| 164 | prefix=f"{prefix}.attn", |
| 165 | ) |
| 166 | |
| 167 | self.apply_rotary_emb = ApplyRotaryEmb( |
| 168 | enforce_enable=True, |
| 169 | enable_fp32_compute=True, |
| 170 | ) |
| 171 | |
| 172 | def split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]: |
| 173 | # [s, b, 3 * head * head_dim] |
| 174 | seq_len, bs, _ = qkv.shape |
| 175 | if self.tp_size > 1: |
| 176 | qkv = all_gather_interleave(qkv, self.qkv.hidden_size, self.tp_size) |
| 177 | |
| 178 | # [s, b, 3 * head * head_dim] -> 3 * [s, b, head * head_dim] |
| 179 | q, k, v = qkv.chunk(3, dim=2) |
| 180 | |
| 181 | # 3 * [s, b, head * head_dim] |
| 182 | if self.tp_size > 1: |
| 183 | splitter = partial( |
| 184 | dist_utils.split_tensor_along_last_dim, num_partitions=self.tp_size |
| 185 | ) |
| 186 | q = splitter(q)[self.tp_rank] |
| 187 | k = splitter(k)[self.tp_rank] |
| 188 | v = splitter(v)[self.tp_rank] |
| 189 | |
| 190 | # 3 * [s, b, head * head_dim] -> 3 * [s, b, head, head_dim] |
| 191 | new_shape = ( |
| 192 | seq_len, |
| 193 | bs, |
| 194 | self.num_attention_heads_per_partition, |
| 195 | self.hidden_size_per_attention_head, |
| 196 | ) |
| 197 | q, k, v = (x.view(*new_shape) for x in (q, k, v)) |
| 198 | return q, k, v |
| 199 | |
| 200 | def forward( |
| 201 | self, |
| 202 | x: torch.Tensor, |
| 203 | cu_seqlens: torch.Tensor, |
| 204 | rotary_pos_emb: torch.Tensor, |
| 205 | max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention |
| 206 | ) -> torch.Tensor: |
| 207 | # [s, b, c] --> [s, b, head * 3 * head_dim] |
| 208 | x, _ = self.qkv(x) |
| 209 | |
| 210 | # [s, b, 3 * head * head_dim] -> 3 * [s, b, head, head_dim] |
| 211 | q, k, v = self.split_qkv(x) |
| 212 | |
| 213 | q, k, v = (rearrange(x, "s b ... -> b s ...").contiguous() for x in (q, k, v)) |
| 214 | if rotary_pos_emb is not None: |
| 215 | qk_concat = torch.cat([q, k], dim=0) |
| 216 | qk_rotated = self.apply_rotary_emb( |
| 217 | qk_concat, |
| 218 | rotary_pos_emb.cos(), |
| 219 | rotary_pos_emb.sin(), |
| 220 | ) |
| 221 | q, k = torch.chunk(qk_rotated, 2, dim=0) |
| 222 | |
| 223 | output = self.attn( |
| 224 | query=q, |
| 225 | key=k, |
| 226 | value=v, |
| 227 | cu_seqlens=cu_seqlens, |
| 228 | max_seqlen=max_seqlen, |
| 229 | ) |
| 230 | context_layer = rearrange(output, "b s h d -> s b (h d)").contiguous() |
| 231 | |
| 232 | output, _ = self.proj(context_layer) |
| 233 | return output |
| 234 | |
| 235 | |
| 236 | class Ernie4_5_VisionMLP(nn.Module): |
| 237 | def __init__( |
| 238 | self, |
| 239 | in_features: int, |
| 240 | hidden_features: int, |
| 241 | act_layer: type[nn.Module] = QuickGELU, |
| 242 | quant_config: QuantizationConfig | None = None, |
| 243 | prefix: str = "", |
| 244 | ): |
| 245 | super().__init__() |
| 246 | self.fc1 = ColumnParallelLinear( |
| 247 | in_features, |
| 248 | hidden_features, |
| 249 | quant_config=quant_config, |
| 250 | prefix=f"{prefix}.fc1", |
| 251 | ) |
| 252 | self.act = act_layer() |
| 253 | self.fc2 = RowParallelLinear( |
| 254 | hidden_features, |
| 255 | in_features, |
| 256 | quant_config=quant_config, |
| 257 | prefix=f"{prefix}.fc2", |
| 258 | ) |
| 259 | |
| 260 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 261 | x_parallel, _ = self.fc1(x) |
| 262 | x_parallel = self.act(x_parallel) |
| 263 | x, _ = self.fc2(x_parallel) |
| 264 | return x |
| 265 | |
| 266 | |
| 267 | class Ernie4_5_VisionBlock(nn.Module): |
| 268 | def __init__( |
| 269 | self, |
| 270 | dim: int, |
| 271 | num_heads: int, |
| 272 | mlp_ratio: float, |
| 273 | act_layer: type[nn.Module] = QuickGELU, |
| 274 | norm_layer: Callable[[int], nn.Module] | None = None, |
| 275 | quant_config: QuantizationConfig | None = None, |
| 276 | prefix: str = "", |
| 277 | ) -> None: |
| 278 | super().__init__() |
| 279 | |
| 280 | if norm_layer is None: |
| 281 | norm_layer = partial(nn.LayerNorm, eps=1e-6) |
| 282 | self.norm1 = norm_layer(dim) |
| 283 | self.norm2 = norm_layer(dim) |
| 284 | mlp_hidden_dim = int(dim * mlp_ratio) |
| 285 | |
| 286 | self.attn = Ernie4_5_VisionAttention( |
| 287 | embed_dim=dim, |
| 288 | num_heads=num_heads, |
| 289 | projection_size=dim, |
| 290 | quant_config=quant_config, |
| 291 | prefix=f"{prefix}.attn", |
| 292 | ) |
| 293 | |
| 294 | self.mlp = Ernie4_5_VisionMLP( |
| 295 | dim, |
| 296 | mlp_hidden_dim, |
| 297 | act_layer=act_layer, |
| 298 | quant_config=quant_config, |
| 299 | prefix=f"{prefix}.mlp", |
| 300 | ) |
| 301 | |
| 302 | def forward( |
| 303 | self, |
| 304 | hidden_states: torch.Tensor, |
| 305 | cu_seqlens: torch.Tensor, |
| 306 | rotary_pos_emb: torch.Tensor, |
| 307 | max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention |
| 308 | ) -> torch.Tensor: |
| 309 | hidden_states = hidden_states + self.attn( |
| 310 | self.norm1(hidden_states), |
| 311 | cu_seqlens=cu_seqlens, |
| 312 | rotary_pos_emb=rotary_pos_emb, |
| 313 | max_seqlen=max_seqlen, |
| 314 | ) |
| 315 | hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)) |
| 316 | return hidden_states |
| 317 | |
| 318 | |
| 319 | class Ernie4_5_VisionPatchEmbed(nn.Module): |
| 320 | def __init__( |
| 321 | self, |
| 322 | patch_size: int = 14, |
| 323 | in_channels: int = 3, |
| 324 | embed_dim: int = 1280, |
| 325 | prefix="", |
| 326 | ) -> None: |
| 327 | super().__init__() |
| 328 | self.patch_size = patch_size |
| 329 | self.in_channels = in_channels |
| 330 | self.embed_dim = embed_dim |
| 331 | |
| 332 | self.proj = nn.Linear( |
| 333 | in_channels * patch_size * patch_size, embed_dim, bias=False |
| 334 | ) |
| 335 | |
| 336 | def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 337 | target_dtype = self.proj.weight.dtype |
| 338 | hidden_states = hidden_states.to(target_dtype) |
| 339 | hidden_states = self.proj(hidden_states) |
| 340 | |
| 341 | return hidden_states |
| 342 | |
| 343 | |
| 344 | class Ernie4_5_VisionRotaryEmbedding(nn.Module): |
| 345 | def __init__(self, dim: int, theta: float = 10000.0) -> None: |
| 346 | super().__init__() |
| 347 | self.register_buffer( |
| 348 | "inv_freq", |
| 349 | 1.0 |
| 350 | / theta |
| 351 | ** (torch.arange(start=0, end=dim, step=2, dtype=torch.float32) / dim), |
| 352 | persistent=False, |
| 353 | ) |
| 354 | |
| 355 | def forward(self, seqlen: int) -> torch.Tensor: |
| 356 | seq = torch.arange( |
| 357 | seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype |
| 358 | ) |
| 359 | freqs = torch.outer(input=seq, vec2=self.inv_freq) |
| 360 | return freqs |
| 361 | |
| 362 | |
| 363 | class Ernie4_5_VisionTransformer(nn.Module): |
| 364 | def __init__( |
| 365 | self, |
| 366 | vision_config, |
| 367 | norm_eps: float = 1e-6, |
| 368 | quant_config: QuantizationConfig | None = None, |
| 369 | prefix: str = "", |
| 370 | ) -> None: |
| 371 | super().__init__() |
| 372 | patch_size = vision_config.patch_size |
| 373 | spatial_merge_size = vision_config.spatial_merge_size |
| 374 | in_channels = vision_config.in_channels |
| 375 | hidden_size = vision_config.hidden_size |
| 376 | embed_dim = vision_config.embed_dim |
| 377 | depth = vision_config.depth |
| 378 | num_heads = vision_config.num_heads |
| 379 | mlp_ratio = vision_config.mlp_ratio |
| 380 | |
| 381 | self.spatial_merge_size = spatial_merge_size |
| 382 | self.num_heads = num_heads |
| 383 | self.embed_dim = embed_dim |
| 384 | |
| 385 | self.patch_embed = Ernie4_5_VisionPatchEmbed( |
| 386 | patch_size=patch_size, |
| 387 | in_channels=in_channels, |
| 388 | embed_dim=embed_dim, |
| 389 | prefix=f"{prefix}.patch_embed", |
| 390 | ) |
| 391 | |
| 392 | norm_layer = partial(nn.LayerNorm, eps=norm_eps) |
| 393 | head_dim = embed_dim // num_heads |
| 394 | self.rotary_pos_emb = Ernie4_5_VisionRotaryEmbedding(head_dim // 2) |
| 395 | |
| 396 | self.blocks = nn.ModuleList( |
| 397 | [ |
| 398 | Ernie4_5_VisionBlock( |
| 399 | dim=embed_dim, |
| 400 | num_heads=num_heads, |
| 401 | mlp_ratio=mlp_ratio, |
| 402 | norm_layer=norm_layer, |
| 403 | quant_config=quant_config, |
| 404 | prefix=f"{prefix}.blocks.{layer_idx}", |
| 405 | ) |
| 406 | for layer_idx in range(depth) |
| 407 | ] |
| 408 | ) |
| 409 | |
| 410 | assert hidden_size == embed_dim, ( |
| 411 | "vit's config.hidden must be equal to config.embed_dim" |
| 412 | ) |
| 413 | self.ln = nn.LayerNorm(hidden_size, eps=1e-6) |
| 414 | |
| 415 | self.attn_backend = get_vit_attn_backend( |
| 416 | head_size=head_dim, |
| 417 | dtype=torch.get_default_dtype(), |
| 418 | ) |
| 419 | |
| 420 | @property |
| 421 | def dtype(self) -> torch.dtype: |
| 422 | return self.patch_embed.proj.weight.dtype |
| 423 | |
| 424 | @property |
| 425 | def device(self) -> torch.device: |
| 426 | return self.patch_embed.proj.weight.device |
| 427 | |
| 428 | def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor: |
| 429 | device = self.device |
| 430 | pos_ids = [] |
| 431 | for t, h, w in grid_thw: |
| 432 | hpos_ids = torch.arange(h, device=device).unsqueeze(1).expand(-1, w) |
| 433 | wpos_ids = torch.arange(w, device=device).unsqueeze(0).expand(h, -1) |
| 434 | hpos_ids = ( |
| 435 | hpos_ids.reshape( |
| 436 | h // self.spatial_merge_size, |
| 437 | self.spatial_merge_size, |
| 438 | w // self.spatial_merge_size, |
| 439 | self.spatial_merge_size, |
| 440 | ) |
| 441 | .permute(0, 2, 1, 3) |
| 442 | .flatten() |
| 443 | ) |
| 444 | wpos_ids = ( |
| 445 | wpos_ids.reshape( |
| 446 | h // self.spatial_merge_size, |
| 447 | self.spatial_merge_size, |
| 448 | w // self.spatial_merge_size, |
| 449 | self.spatial_merge_size, |
| 450 | ) |
| 451 | .permute(0, 2, 1, 3) |
| 452 | .flatten() |
| 453 | ) |
| 454 | pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1)) |
| 455 | pos_ids = torch.cat(pos_ids, dim=0) |
| 456 | max_grid_size = grid_thw[:, 1:].max() |
| 457 | rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size) |
| 458 | rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1) |
| 459 | return rotary_pos_emb |
| 460 | |
| 461 | def compute_attn_mask_seqlen(self, cu_seqlens: torch.Tensor) -> torch.Tensor | None: |
| 462 | max_seqlen = None |
| 463 | if self.attn_backend in { |
| 464 | AttentionBackendEnum.FLASH_ATTN, |
| 465 | AttentionBackendEnum.ROCM_AITER_FA, |
| 466 | AttentionBackendEnum.TRITON_ATTN, |
| 467 | }: |
| 468 | max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max() |
| 469 | return max_seqlen |
| 470 | |
| 471 | def prepare_encoder_metadata( |
| 472 | self, |
| 473 | grid_thw: torch.Tensor | list[list[int]], |
| 474 | *, |
| 475 | num_pad: int = 0, |
| 476 | max_batch_size: int | None = None, |
| 477 | max_seqlen_override: int | None = None, |
| 478 | device: torch.device | None = None, |
| 479 | ) -> dict[str, torch.Tensor]: |
| 480 | """Compute encoder metadata outside the CUDA graph. |
| 481 | |
| 482 | Splits the rotary embeddings, ``cu_seqlens`` and ``max_seqlen`` out of |
| 483 | `forward` so they can be precomputed on the host side and fed into a |
| 484 | captured graph through fixed buffers. Shared by the eager path, CUDA |
| 485 | graph capture and CUDA graph replay |
| 486 | |
| 487 | Args: |
| 488 | grid_thw: Per-frame grid sizes as a ``[num_frames, 3]`` tensor or |
| 489 | an equivalent list of ``[t, h, w]``. |
| 490 | num_pad: Legacy ``cu_seqlens`` padding |
| 491 | max_batch_size: If set, pad ``cu_seqlens`` to ``max_batch_size + 1`` |
| 492 | entries for a fixed graph buffer shape. |
| 493 | max_seqlen_override: If set, use this value for ``max_seqlen`` |
| 494 | instead of deriving it from ``cu_seqlens``. CUDA graph capture |
| 495 | bakes ``max_seqlen`` at capture time, so it must cover the |
| 496 | worst-case replay. |
| 497 | device: Device for the returned tensors. Defaults to ``self.device``. |
| 498 | |
| 499 | Returns: |
| 500 | Dict with ``rotary_pos_emb``, ``cu_seqlens`` and ``max_seqlen``. |
| 501 | |
| 502 | """ |
| 503 | if device is None: |
| 504 | device = self.device |
| 505 | if not isinstance(grid_thw, torch.Tensor): |
| 506 | grid_thw = torch.tensor(grid_thw, dtype=torch.int32) |
| 507 | |
| 508 | rotary_pos_emb = self.rot_pos_emb(grid_thw).to(device) |
| 509 | |
| 510 | cu_seqlens = torch.repeat_interleave( |
| 511 | grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0] |
| 512 | ).cumsum(dim=0, dtype=torch.int32) |
| 513 | |
| 514 | zeros = cu_seqlens.new_zeros(1) |
| 515 | if num_pad > 0: |
| 516 | cu_seqlens = torch.cat([zeros, cu_seqlens, zeros]) |
| 517 | cu_seqlens[-1] = cu_seqlens[-2] + num_pad |
| 518 | else: |
| 519 | cu_seqlens = torch.cat([zeros, cu_seqlens]) |
| 520 | |
| 521 | # Pad cu_seqlens to a fixed number of sequences so the CUDA graph |
| 522 | # replay buffer keeps a constant shape across batches. |
| 523 | if max_batch_size is not None: |
| 524 | num_seqs = cu_seqlens.numel() - 1 |
| 525 | if num_seqs < max_batch_size: |
| 526 | cu_seqlens = torch.cat( |
| 527 | [ |
| 528 | cu_seqlens, |
| 529 | cu_seqlens.new_full( |
| 530 | (max_batch_size - num_seqs,), cu_seqlens[-1] |
| 531 | ), |
| 532 | ] |
| 533 | ) |
| 534 | |
| 535 | # max_seqlen is baked into the graph at capture time, so capture passes |
| 536 | # a worst-case override. compute_attn_mask_seqlen keeps it on CPU |
| 537 | # (consumed via .item()), avoiding a captured D2H copy. |
| 538 | if max_seqlen_override is not None: |
| 539 | max_seqlen = torch.tensor(max_seqlen_override, dtype=torch.int32) |
| 540 | else: |
| 541 | max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens) |
| 542 | |
| 543 | cu_seqlens = async_tensor_h2d(cu_seqlens, device) |
| 544 | |
| 545 | return { |
| 546 | "rotary_pos_emb": rotary_pos_emb, |
| 547 | "cu_seqlens": cu_seqlens, |
| 548 | "max_seqlen": max_seqlen, |
| 549 | } |
| 550 | |
| 551 | def forward( |
| 552 | self, |
| 553 | hidden_states: torch.Tensor, |
| 554 | grid_thw: torch.Tensor | None = None, |
| 555 | num_pad: int = 0, |
| 556 | *, |
| 557 | encoder_metadata: dict[str, torch.Tensor] | None = None, |
| 558 | ) -> torch.Tensor: |
| 559 | hidden_states = self.patch_embed(hidden_states) |
| 560 | |
| 561 | if encoder_metadata is None: |
| 562 | # Eager path: compute metadata inline |
| 563 | encoder_metadata = self.prepare_encoder_metadata( |
| 564 | grid_thw, num_pad=num_pad, device=hidden_states.device |
| 565 | ) |
| 566 | rotary_pos_emb = encoder_metadata["rotary_pos_emb"] |
| 567 | cu_seqlens = encoder_metadata["cu_seqlens"] |
| 568 | max_seqlen = encoder_metadata["max_seqlen"] |
| 569 | |
| 570 | # add batch size |
| 571 | if hidden_states.ndim == 2: |
| 572 | hidden_states = hidden_states.unsqueeze(dim=1) |
| 573 | |
| 574 | for blk in self.blocks: |
| 575 | hidden_states = blk( |
| 576 | hidden_states, |
| 577 | cu_seqlens=cu_seqlens, |
| 578 | rotary_pos_emb=rotary_pos_emb, |
| 579 | max_seqlen=max_seqlen, |
| 580 | ) |
| 581 | |
| 582 | final_output = self.ln(hidden_states) |
| 583 | |
| 584 | if final_output.ndim == 3: |
| 585 | final_output = final_output.squeeze(dim=1) |
| 586 | |
| 587 | return final_output |
| 588 | |
| 589 | def load_weights(self, weights) -> set[str]: |
| 590 | params_dict = dict(self.named_parameters(remove_duplicate=False)) |
| 591 | loaded_params: set[str] = set() |
| 592 | |
| 593 | for name, loaded_weight in weights: |
| 594 | param = params_dict[name] |
| 595 | weight_loader = getattr(param, "weight_loader", default_weight_loader) |
| 596 | weight_loader(param, loaded_weight) |
| 597 | loaded_params.add(name) |
| 598 | return loaded_params |
| 599 | |
| 600 | |
| 601 | # === Vision Inputs === # |
| 602 | |
| 603 | |
| 604 | class Ernie4_5_VLImagePixelInputs(TensorSchema): |
| 605 | """Dimensions: |
| 606 | - np: The total number of patches over each image over each prompt in |
| 607 | the batch |
| 608 | - ni: Number of images |
| 609 | - cps: Number of channels * patch_size * patch_size |
| 610 | """ |
| 611 | |
| 612 | type: Literal["pixel_values"] |
| 613 | |
| 614 | pixel_values: Annotated[torch.Tensor, TensorShape("np", "cps")] |
| 615 | image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)] |
| 616 | |
| 617 | |
| 618 | Ernie4_5_VLImageInputs = Ernie4_5_VLImagePixelInputs |
| 619 | |
| 620 | |
| 621 | class Ernie4_5_VLVideoPixelInputs(TensorSchema): |
| 622 | """Dimensions: |
| 623 | - np: The total number of patches over each image over each prompt in |
| 624 | the batch |
| 625 | - ni: Number of images |
| 626 | - cps: Number of channels * temporal_patch_size * patch_size * |
| 627 | patch_size |
| 628 | """ |
| 629 | |
| 630 | type: Literal["pixel_values_videos"] |
| 631 | pixel_values_videos: Annotated[torch.Tensor, TensorShape("np", "cps")] |
| 632 | video_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)] |
| 633 | |
| 634 | |
| 635 | Ernie4_5_VLVideoInputs = Ernie4_5_VLVideoPixelInputs |
| 636 | |
| 637 | # === Vision Processor === # |
| 638 | |
| 639 | |
| 640 | def round_by_factor(number: int | float, factor: int) -> int: |
| 641 | return round(number / factor) * factor |
| 642 | |
| 643 | |
| 644 | def ceil_by_factor(number: int | float, factor: int) -> int: |
| 645 | return math.ceil(number / factor) * factor |
| 646 | |
| 647 | |
| 648 | def floor_by_factor(number: int | float, factor: int) -> int: |
| 649 | return math.floor(number / factor) * factor |
| 650 | |
| 651 | |
| 652 | def smart_resize( |
| 653 | height: int, |
| 654 | width: int, |
| 655 | factor: int = 28, |
| 656 | min_pixels: int = 4 * 28 * 28, |
| 657 | max_pixels: int = 16384 * 28 * 28, |
| 658 | ): |
| 659 | MAX_RATIO = 200 |
| 660 | if max(height, width) / min(height, width) > MAX_RATIO: |
| 661 | if height > width: |
| 662 | new_width = max(factor, round_by_factor(width, factor)) |
| 663 | new_height = floor_by_factor(new_width * MAX_RATIO, factor) |
| 664 | else: |
| 665 | new_height = max(factor, round_by_factor(height, factor)) |
| 666 | new_width = floor_by_factor(new_height * MAX_RATIO, factor) |
| 667 | |
| 668 | height = new_height |
| 669 | width = new_width |
| 670 | |
| 671 | h_bar = max(factor, round_by_factor(height, factor)) |
| 672 | w_bar = max(factor, round_by_factor(width, factor)) |
| 673 | if h_bar * w_bar > max_pixels: |
| 674 | beta = math.sqrt((height * width) / max_pixels) |
| 675 | h_bar = floor_by_factor(height / beta, factor) |
| 676 | w_bar = floor_by_factor(width / beta, factor) |
| 677 | elif h_bar * w_bar < min_pixels: |
| 678 | beta = math.sqrt(min_pixels / (height * width)) |
| 679 | h_bar = ceil_by_factor(height * beta, factor) |
| 680 | w_bar = ceil_by_factor(width * beta, factor) |
| 681 | |
| 682 | if min_pixels > h_bar * w_bar or h_bar * w_bar > max_pixels: |
| 683 | raise ValueError( |
| 684 | f"Invalid h_bar={h_bar}, w_bar={w_bar}: " |
| 685 | f"h_bar * w_bar must be >= min_pixels ({min_pixels}) " |
| 686 | f"and <= max_pixels ({max_pixels})." |
| 687 | ) |
| 688 | |
| 689 | return h_bar, w_bar |
| 690 | |
| 691 | |
| 692 | class VariableResolutionResamplerModel(nn.Module): |
| 693 | def __init__( |
| 694 | self, |
| 695 | in_dim, |
| 696 | out_dim, |
| 697 | spatial_conv_size, |
| 698 | temporal_conv_size, |
| 699 | config, |
| 700 | prefix: str = "", |
| 701 | ) -> None: |
| 702 | super().__init__() |
| 703 | self.in_dim = in_dim |
| 704 | self.out_dim = out_dim |
| 705 | self.config = config |
| 706 | self.spatial_conv_size = spatial_conv_size |
| 707 | self.temporal_conv_size = temporal_conv_size |
| 708 | self.use_temporal_conv = config.use_temporal_conv |
| 709 | |
| 710 | # compress 2d conv(picture) to 1d |
| 711 | self.spatial_dim = self.in_dim * self.spatial_conv_size * self.spatial_conv_size |
| 712 | # compress 3d conv(video) to 1d |
| 713 | self.temporal_dim = ( |
| 714 | self.in_dim |
| 715 | * self.spatial_conv_size |
| 716 | * self.spatial_conv_size |
| 717 | * self.temporal_conv_size |
| 718 | ) |
| 719 | |
| 720 | self.spatial_linear1 = ColumnParallelLinear( |
| 721 | self.spatial_dim, |
| 722 | self.spatial_dim, |
| 723 | bias=True, |
| 724 | gather_output=True, |
| 725 | quant_config=getattr(config, "quant_config", None), |
| 726 | prefix=f"{prefix}.spatial_linear1", |
| 727 | ) |
| 728 | |
| 729 | self.spatial_gelu = nn.GELU() |
| 730 | |
| 731 | self.spatial_linear2 = ColumnParallelLinear( |
| 732 | self.spatial_dim, |
| 733 | self.spatial_dim, |
| 734 | bias=True, |
| 735 | gather_output=True, |
| 736 | quant_config=getattr(config, "quant_config", None), |
| 737 | prefix=f"{prefix}.spatial_linear2", |
| 738 | ) |
| 739 | |
| 740 | self.spatial_norm = nn.LayerNorm(self.spatial_dim, eps=1e-6) |
| 741 | |
| 742 | if self.use_temporal_conv: |
| 743 | self.temporal_linear1 = ColumnParallelLinear( |
| 744 | self.temporal_dim, |
| 745 | self.spatial_dim, |
| 746 | bias=True, |
| 747 | gather_output=True, |
| 748 | quant_config=getattr(config, "quant_config", None), |
| 749 | prefix=f"{prefix}.temporal_linear1", |
| 750 | ) |
| 751 | |
| 752 | self.temporal_gelu = nn.GELU() |
| 753 | |
| 754 | self.temporal_linear2 = ColumnParallelLinear( |
| 755 | self.spatial_dim, |
| 756 | self.spatial_dim, |
| 757 | bias=True, |
| 758 | gather_output=True, |
| 759 | quant_config=getattr(config, "quant_config", None), |
| 760 | prefix=f"{prefix}.temporal_linear2", |
| 761 | ) |
| 762 | |
| 763 | self.temporal_norm = nn.LayerNorm(self.spatial_dim, eps=1e-6) |
| 764 | |
| 765 | self.mlp = ColumnParallelLinear( |
| 766 | self.spatial_dim, |
| 767 | self.out_dim, |
| 768 | bias=True, |
| 769 | gather_output=True, |
| 770 | quant_config=getattr(config, "quant_config", None), |
| 771 | prefix=f"{prefix}.mlp", |
| 772 | ) |
| 773 | |
| 774 | self.after_norm = RMSNorm( |
| 775 | hidden_size=out_dim, eps=getattr(config, "rms_norm_eps", 1e-6) |
| 776 | ) |
| 777 | |
| 778 | def spatial_conv_reshape(self, x, spatial_conv_size): |
| 779 | S, C = x.shape |
| 780 | x = x.reshape([-1, C * (spatial_conv_size**2)]) |
| 781 | return x |
| 782 | |
| 783 | def forward(self, x, grid_thw): |
| 784 | def fwd_spatial(x): |
| 785 | x = self.spatial_conv_reshape(x, self.spatial_conv_size) |
| 786 | |
| 787 | x, _ = self.spatial_linear1(x) |
| 788 | x = self.spatial_gelu(x) |
| 789 | x, _ = self.spatial_linear2(x) |
| 790 | x = self.spatial_norm(x) |
| 791 | |
| 792 | return x |
| 793 | |
| 794 | def fwd_placeholder(x, grid_thw, to_tensor=False): |
| 795 | # The per-image grids drive the Python-level offset arithmetic |
| 796 | # below; the same read is wrapped on the cudagraph path. |
| 797 | with gpu_sync_allowed(): |
| 798 | grid_thw_cpu = grid_thw.cpu().numpy() |
| 799 | grid_t, grid_hw = grid_thw_cpu[:, 0], grid_thw_cpu[:, 1:] |
| 800 | grid_hw_after_conv = grid_hw.prod(-1) // (self.spatial_conv_size**2) |
| 801 | |
| 802 | tokens_per_img_or_vid = grid_thw_cpu.prod(-1) // (self.spatial_conv_size**2) |
| 803 | batch_offset = np.empty( |
| 804 | tokens_per_img_or_vid.size, dtype=tokens_per_img_or_vid.dtype |
| 805 | ) |
| 806 | batch_offset[0] = 0 |
| 807 | batch_offset[1:] = tokens_per_img_or_vid.cumsum()[:-1] |
| 808 | |
| 809 | slice_offsets = [] |
| 810 | for temporoal_size, spatial_size, b_offset in zip( |
| 811 | grid_t, grid_hw_after_conv, batch_offset |
| 812 | ): |
| 813 | for temp_offset in range(0, temporoal_size, 2): |
| 814 | slice_offsets.append( |
| 815 | np.arange( |
| 816 | b_offset + (temp_offset) * spatial_size, |
| 817 | b_offset + (temp_offset + 1) * spatial_size, |
| 818 | ) |
| 819 | ) |
| 820 | slice_offsets = async_tensor_h2d( |
| 821 | np.concatenate(slice_offsets, axis=-1), device=x.device |
| 822 | ) |
| 823 | |
| 824 | slice_offsets2 = [] |
| 825 | for temporoal_size, spatial_size, b_offset in zip( |
| 826 | grid_t, grid_hw_after_conv, batch_offset |
| 827 | ): |
| 828 | for temp_offset in range( |
| 829 | 1 if temporoal_size > 1 else 0, temporoal_size, 2 |
| 830 | ): |
| 831 | slice_offsets2.append( |
| 832 | np.arange( |
| 833 | b_offset + (temp_offset) * spatial_size, |
| 834 | b_offset + (temp_offset + 1) * spatial_size, |
| 835 | ) |
| 836 | ) |
| 837 | slice_offsets2 = async_tensor_h2d( |
| 838 | np.concatenate(slice_offsets2, axis=-1), device=x.device |
| 839 | ) |
| 840 | |
| 841 | x_timestep_1 = torch.index_select(x, dim=0, index=slice_offsets) |
| 842 | x_timestep_2 = torch.index_select(x, dim=0, index=slice_offsets2) |
| 843 | x = torch.concat([x_timestep_1, x_timestep_2], dim=-1) |
| 844 | return x |
| 845 | |
| 846 | def fwd_temporal(x): |
| 847 | x, _ = self.temporal_linear1(x) |
| 848 | x = self.temporal_gelu(x) |
| 849 | x, _ = self.temporal_linear2(x) |
| 850 | x = self.temporal_norm(x) |
| 851 | return x |
| 852 | |
| 853 | def fwd_mlp(x): |
| 854 | x, _ = self.mlp(x) |
| 855 | x = self.after_norm(x) |
| 856 | return x |
| 857 | |
| 858 | x = fwd_spatial(x) |
| 859 | if self.use_temporal_conv: |
| 860 | x = fwd_placeholder(x, grid_thw) |
| 861 | x = fwd_temporal(x) |
| 862 | x = fwd_mlp(x) |
| 863 | return x |
| 864 | |
| 865 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 866 | params_dict = dict(self.named_parameters(remove_duplicate=False)) |
| 867 | loaded_params: set[str] = set() |
| 868 | |
| 869 | for name, loaded_weight in weights: |
| 870 | if name not in params_dict: |
| 871 | continue |
| 872 | param = params_dict[name] |
| 873 | weight_loader = getattr(param, "weight_loader", default_weight_loader) |
| 874 | weight_loader(param, loaded_weight) |
| 875 | loaded_params.add(name) |
| 876 | return loaded_params |
| 877 | |
| 878 | |
| 879 | class Ernie4_5_VLProcessingInfo(BaseProcessingInfo): |
| 880 | def get_hf_config(self): |
| 881 | return self.ctx.model_config.hf_config |
| 882 | |
| 883 | def get_hf_processor(self, **kwargs: object): |
| 884 | return self.ctx.get_hf_processor(**kwargs) |
| 885 | |
| 886 | def get_image_processor(self, **kwargs: object): |
| 887 | return self.get_hf_processor(**kwargs).image_processor |
| 888 | |
| 889 | def get_data_parser(self): |
| 890 | return MultiModalDataParser( |
| 891 | video_needs_metadata=True, |
| 892 | expected_hidden_size=self._get_expected_hidden_size(), |
| 893 | ) |
| 894 | |
| 895 | def get_supported_mm_limits(self) -> Mapping[str, int | None]: |
| 896 | return {"image": None, "video": None} |
| 897 | |
| 898 | def get_mm_max_tokens_per_item( |
| 899 | self, |
| 900 | seq_len: int, |
| 901 | mm_counts: Mapping[str, int], |
| 902 | ) -> Mapping[str, int]: |
| 903 | max_image_tokens = self.get_max_image_tokens() |
| 904 | max_video_tokens = self.get_max_video_tokens(seq_len, mm_counts) |
| 905 | return {"image": max_image_tokens, "video": max_video_tokens} |
| 906 | |
| 907 | def _get_vision_info( |
| 908 | self, |
| 909 | *, |
| 910 | image_width: int, |
| 911 | image_height: int, |
| 912 | num_frames: int = 1, |
| 913 | do_resize: bool = True, |
| 914 | image_processor: BaseImageProcessor, |
| 915 | mm_kwargs: Mapping[str, object], |
| 916 | ) -> tuple[ImageSize, int]: |
| 917 | hf_config = self.get_hf_config() |
| 918 | vision_config = hf_config.vision_config |
| 919 | |
| 920 | patch_size = vision_config.patch_size |
| 921 | spatial_conv_size = hf_config.spatial_conv_size |
| 922 | temporal_conv_size = hf_config.temporal_conv_size |
| 923 | |
| 924 | if self.ctx.model_config.trust_remote_code: |
| 925 | # Defined in HF Hub repo |
| 926 | min_pixels_key = "min_pixels" |
| 927 | max_pixels_key = "max_pixels" |
| 928 | else: |
| 929 | # Defined in Transformers library (requires v5.0 or above) |
| 930 | min_pixels_key = "shortest_edge" |
| 931 | max_pixels_key = "longest_edge" |
| 932 | |
| 933 | mm_kwargs = self.ctx.get_merged_mm_kwargs(mm_kwargs) |
| 934 | size = image_processor.size |
| 935 | if override_size := mm_kwargs.get("size"): |
| 936 | size = size | override_size |
| 937 | if (override_min_pixels := mm_kwargs.get("min_pixels")) is not None: |
| 938 | size = size | {min_pixels_key: override_min_pixels} |
| 939 | if (override_max_pixels := mm_kwargs.get("max_pixels")) is not None: |
| 940 | size = size | {max_pixels_key: override_max_pixels} |
| 941 | |
| 942 | if do_resize: |
| 943 | resized_height, resized_width = smart_resize( |
| 944 | height=image_height, |
| 945 | width=image_width, |
| 946 | factor=patch_size * spatial_conv_size, |
| 947 | min_pixels=size[min_pixels_key], |
| 948 | max_pixels=size[max_pixels_key], |
| 949 | ) |
| 950 | preprocessed_size = ImageSize(width=resized_width, height=resized_height) |
| 951 | else: |
| 952 | preprocessed_size = ImageSize(width=image_width, height=image_height) |
| 953 | |
| 954 | grid_t = max(num_frames // temporal_conv_size, 1) |
| 955 | grid_h = preprocessed_size.height // patch_size |
| 956 | grid_w = preprocessed_size.width // patch_size |
| 957 | |
| 958 | num_patches = grid_t * grid_h * grid_w |
| 959 | num_vision_tokens = num_patches // (spatial_conv_size**2) |
| 960 | |
| 961 | return preprocessed_size, num_vision_tokens |
| 962 | |
| 963 | def get_num_image_tokens( |
| 964 | self, |
| 965 | *, |
| 966 | image_width: int, |
| 967 | image_height: int, |
| 968 | image_processor: BaseImageProcessor, |
| 969 | mm_kwargs: Mapping[str, object], |
| 970 | ) -> int: |
| 971 | _, num_image_tokens = self._get_vision_info( |
| 972 | image_width=image_width, |
| 973 | image_height=image_height, |
| 974 | image_processor=image_processor, |
| 975 | mm_kwargs=mm_kwargs, |
| 976 | ) |
| 977 | return num_image_tokens |
| 978 | |
| 979 | def get_num_video_tokens( |
| 980 | self, |
| 981 | *, |
| 982 | image_width: int, |
| 983 | image_height: int, |
| 984 | num_frames: int, |
| 985 | image_processor: BaseImageProcessor, |
| 986 | mm_kwargs: Mapping[str, object], |
| 987 | ) -> int: |
| 988 | _, num_video_tokens = self._get_vision_info( |
| 989 | image_width=image_width, |
| 990 | image_height=image_height, |
| 991 | num_frames=num_frames, |
| 992 | image_processor=image_processor, |
| 993 | mm_kwargs=mm_kwargs, |
| 994 | ) |
| 995 | return num_video_tokens |
| 996 | |
| 997 | def get_image_size_with_most_features(self) -> ImageSize: |
| 998 | image_processor = self.get_image_processor() |
| 999 | |
| 1000 | max_image_size, _ = self._get_vision_info( |
| 1001 | image_width=9999999, |
| 1002 | image_height=9999999, |
| 1003 | image_processor=image_processor, |
| 1004 | mm_kwargs={}, |
| 1005 | ) |
| 1006 | return max_image_size |
| 1007 | |
| 1008 | def get_max_image_tokens(self) -> int: |
| 1009 | image_processor = self.get_image_processor() |
| 1010 | target_width, target_height = self.get_image_size_with_most_features() |
| 1011 | |
| 1012 | num_image_tokens = self.get_num_image_tokens( |
| 1013 | image_width=target_width, |
| 1014 | image_height=target_height, |
| 1015 | image_processor=image_processor, |
| 1016 | mm_kwargs={}, |
| 1017 | ) |
| 1018 | return num_image_tokens |
| 1019 | |
| 1020 | def _get_max_video_frames(self, max_tokens: int) -> int: |
| 1021 | image_processor = self.get_image_processor() |
| 1022 | target_width, target_height = self.get_image_size_with_most_features() |
| 1023 | |
| 1024 | num_frames = 0 |
| 1025 | |
| 1026 | while True: |
| 1027 | next_num_frames = num_frames + 1 |
| 1028 | next_max_tokens = self.get_num_video_tokens( |
| 1029 | image_width=target_width, |
| 1030 | image_height=target_height, |
| 1031 | num_frames=next_num_frames, |
| 1032 | image_processor=image_processor, |
| 1033 | mm_kwargs={}, |
| 1034 | ) |
| 1035 | |
| 1036 | if next_max_tokens > max_tokens: |
| 1037 | break |
| 1038 | |
| 1039 | num_frames = next_num_frames |
| 1040 | |
| 1041 | # If the number of frames is odd, discard one frame. |
| 1042 | if num_frames % 2 != 0: |
| 1043 | num_frames -= 1 |
| 1044 | |
| 1045 | return num_frames |
| 1046 | |
| 1047 | def get_num_frames_with_most_features( |
| 1048 | self, |
| 1049 | seq_len: int, |
| 1050 | mm_counts: Mapping[str, int], |
| 1051 | ) -> int: |
| 1052 | max_images = mm_counts.get("image", 0) |
| 1053 | max_videos = mm_counts.get("video", 0) |
| 1054 | |
| 1055 | max_image_tokens = self.get_max_image_tokens() * max_images |
| 1056 | max_total_frames = self._get_max_video_frames(seq_len - max_image_tokens) |
| 1057 | max_frames_per_video = max_total_frames // max(max_videos, 1) |
| 1058 | |
| 1059 | return max(max_frames_per_video, 2) |
| 1060 | |
| 1061 | def get_max_video_tokens( |
| 1062 | self, |
| 1063 | seq_len: int, |
| 1064 | mm_counts: Mapping[str, int], |
| 1065 | ) -> int: |
| 1066 | image_processor = self.get_image_processor() |
| 1067 | target_width, target_height = self.get_image_size_with_most_features() |
| 1068 | |
| 1069 | return self.get_num_video_tokens( |
| 1070 | image_width=target_width, |
| 1071 | image_height=target_height, |
| 1072 | num_frames=self.get_num_frames_with_most_features(seq_len, mm_counts), |
| 1073 | image_processor=image_processor, |
| 1074 | mm_kwargs={}, |
| 1075 | ) |
| 1076 | |
| 1077 | |
| 1078 | class Ernie4_5VLMultiModalProcessor(BaseMultiModalProcessor[Ernie4_5_VLProcessingInfo]): |
| 1079 | def _pixel_values_norm( |
| 1080 | self, |
| 1081 | pixel_values: torch.Tensor, |
| 1082 | mm_kwargs: Mapping[str, object], |
| 1083 | ) -> torch.Tensor: |
| 1084 | hf_config = self.info.get_hf_config() |
| 1085 | vision_config = hf_config.vision_config |
| 1086 | image_processor = self.info.get_image_processor(**mm_kwargs) |
| 1087 | image_mean_tensor = torch.tensor( |
| 1088 | image_processor.image_mean, dtype=torch.float32 |
| 1089 | ).reshape([1, 3, 1, 1]) |
| 1090 | image_std_tensor = torch.tensor( |
| 1091 | image_processor.image_std, dtype=torch.float32 |
| 1092 | ).reshape([1, 3, 1, 1]) |
| 1093 | rescale_factor = torch.tensor( |
| 1094 | image_processor.rescale_factor, dtype=torch.float32 |
| 1095 | ) |
| 1096 | patch_size_squared = vision_config.patch_size**2 |
| 1097 | |
| 1098 | image_mean_tensor = image_mean_tensor.squeeze([-2, -1]).repeat_interleave( |
| 1099 | patch_size_squared, -1 |
| 1100 | ) |
| 1101 | image_std_tensor = image_std_tensor.squeeze([-2, -1]).repeat_interleave( |
| 1102 | patch_size_squared, -1 |
| 1103 | ) |
| 1104 | |
| 1105 | if not image_mean_tensor.is_contiguous(): |
| 1106 | image_mean_tensor = image_mean_tensor.contiguous() |
| 1107 | if not image_std_tensor.is_contiguous(): |
| 1108 | image_std_tensor = image_std_tensor.contiguous() |
| 1109 | |
| 1110 | pixel_values = ( |
| 1111 | rescale_factor * pixel_values.to(torch.float32) - image_mean_tensor |
| 1112 | ) / image_std_tensor |
| 1113 | pixel_values = pixel_values.to(hf_config.dtype) |
| 1114 | return pixel_values |
| 1115 | |
| 1116 | def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: |
| 1117 | return self.dummy_inputs.get_dummy_text(mm_counts) |
| 1118 | |
| 1119 | def _get_hf_mm_inputs( |
| 1120 | self, |
| 1121 | mm_items: MultiModalDataItems, |
| 1122 | hf_kwargs: Mapping[str, object], |
| 1123 | ) -> HFMultiModalInputs: |
| 1124 | hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) |
| 1125 | hf_data = hf_inputs.hf_data |
| 1126 | |
| 1127 | if "images" not in hf_data and "videos" not in hf_data: |
| 1128 | return hf_inputs |
| 1129 | |
| 1130 | prompt_text = hf_data["text"] |
| 1131 | hf_data["text"] = [prompt_text] |
| 1132 | hf_data.setdefault("images", []) |
| 1133 | hf_data.setdefault("videos", []) |
| 1134 | |
| 1135 | # Check if HF processor supports video metadata |
| 1136 | hf_processor = self.info.get_hf_processor(**hf_inputs.hf_kwargs) |
| 1137 | supports_video_metadata = getattr( |
| 1138 | hf_processor, "supports_video_metadata", False |
| 1139 | ) |
| 1140 | |
| 1141 | videos = hf_data["videos"] |
| 1142 | assert isinstance(videos, Sequence) |
| 1143 | if videos and not supports_video_metadata: |
| 1144 | # Old HF processor, unwrap tuple to pure frames |
| 1145 | logger.warning_once( |
| 1146 | "HF processor doesn't support video metadata. " |
| 1147 | "Timestamps will NOT be rendered. Please upgrade the model." |
| 1148 | ) |
| 1149 | hf_data["videos"] = [v[0] if isinstance(v, tuple) else v for v in videos] |
| 1150 | |
| 1151 | return hf_inputs |
| 1152 | |
| 1153 | def _postprocess_hf_mm_data( |
| 1154 | self, |
| 1155 | hf_data: Mapping[str, object], |
| 1156 | hf_kwargs: Mapping[str, object], |
| 1157 | processed_data: BatchFeature, |
| 1158 | ) -> BatchFeature: |
| 1159 | if "images" not in hf_data and "videos" not in hf_data: |
| 1160 | return processed_data |
| 1161 | |
| 1162 | # Divide the processor_output into two modalities: image and video. |
| 1163 | pixel_values = processed_data["images"] |
| 1164 | if pixel_values is not None: |
| 1165 | processed_data["images"] = self._pixel_values_norm(pixel_values, hf_kwargs) |
| 1166 | for key in list(processed_data.keys()): |
| 1167 | if processed_data[key] is None: |
| 1168 | del processed_data[key] |
| 1169 | continue |
| 1170 | if key == "grid_thw": |
| 1171 | grid_thw = processed_data["grid_thw"] |
| 1172 | pixel_values_all = processed_data["images"] |
| 1173 | # Identify elements where the first |
| 1174 | # dimension is greater than 1 and |
| 1175 | # treat them as the video modality |
| 1176 | mask = grid_thw[:, 0] > 1 |
| 1177 | processed_data["video_grid_thw"] = grid_thw[mask] |
| 1178 | processed_data["image_grid_thw"] = grid_thw[~mask] |
| 1179 | image_patch_num = processed_data["image_grid_thw"].prod(dim=1).sum() |
| 1180 | processed_data["pixel_values"] = pixel_values_all[:image_patch_num] |
| 1181 | processed_data["pixel_values_videos"] = pixel_values_all[ |
| 1182 | image_patch_num: |
| 1183 | ] |
| 1184 | del processed_data["images"] |
| 1185 | |
| 1186 | return processed_data |
| 1187 | |
| 1188 | def _get_prompt_updates( |
| 1189 | self, |
| 1190 | mm_items: MultiModalDataItems, |
| 1191 | hf_processor_mm_kwargs: Mapping[str, Any], |
| 1192 | out_mm_kwargs: MultiModalKwargsItems, |
| 1193 | ) -> Sequence[PromptUpdate]: |
| 1194 | hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) |
| 1195 | tokenizer = self.info.get_tokenizer() |
| 1196 | |
| 1197 | before_placeholder = { |
| 1198 | "image": "<|image@placeholder|>", |
| 1199 | "video": "<|video@placeholder|>", |
| 1200 | } |
| 1201 | |
| 1202 | after_placeholder = { |
| 1203 | # image and video have same placeholder |
| 1204 | "image": "<|IMAGE_PLACEHOLDER|>", |
| 1205 | "video": "<|IMAGE_PLACEHOLDER|>", |
| 1206 | } |
| 1207 | |
| 1208 | merge_length = hf_processor.spatial_conv_size**2 |
| 1209 | |
| 1210 | def get_replacement_ernie45vl(item_idx: int, modality: str): |
| 1211 | out_item = out_mm_kwargs[modality][item_idx] |
| 1212 | grid_thw = out_item[f"{modality}_grid_thw"].data |
| 1213 | assert isinstance(grid_thw, torch.Tensor) |
| 1214 | if modality == "video": |
| 1215 | num_tokens = ( |
| 1216 | int(grid_thw.prod()) |
| 1217 | // hf_processor.temporal_conv_size |
| 1218 | // merge_length |
| 1219 | ) |
| 1220 | else: |
| 1221 | num_tokens = int(grid_thw.prod()) // merge_length |
| 1222 | placeholder_ids = cached_encode( |
| 1223 | tokenizer, after_placeholder[modality], add_special_tokens=False |
| 1224 | ) |
| 1225 | return placeholder_ids * num_tokens |
| 1226 | |
| 1227 | return [ |
| 1228 | PromptReplacement( |
| 1229 | modality=modality, |
| 1230 | target=cached_encode( |
| 1231 | tokenizer, before_placeholder[modality], add_special_tokens=False |
| 1232 | ), |
| 1233 | replacement=partial(get_replacement_ernie45vl, modality=modality), |
| 1234 | ) |
| 1235 | for modality in ("image", "video") |
| 1236 | ] |
| 1237 | |
| 1238 | def _get_mm_fields_config( |
| 1239 | self, |
| 1240 | hf_inputs: BatchFeature, |
| 1241 | hf_processor_mm_kwargs: Mapping[str, object], |
| 1242 | ) -> Mapping[str, MultiModalFieldConfig]: |
| 1243 | image_grid_thw = hf_inputs.get("image_grid_thw", torch.empty((0, 3))) |
| 1244 | image_grid_sizes = image_grid_thw.prod(-1) |
| 1245 | |
| 1246 | video_grid_thw = hf_inputs.get("video_grid_thw", torch.empty((0, 3))) |
| 1247 | video_grid_sizes = video_grid_thw.prod(-1) |
| 1248 | |
| 1249 | return dict( |
| 1250 | pixel_values=MultiModalFieldConfig.flat_from_sizes( |
| 1251 | "image", image_grid_sizes |
| 1252 | ), |
| 1253 | image_grid_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True), |
| 1254 | pixel_values_videos=MultiModalFieldConfig.flat_from_sizes( |
| 1255 | "video", video_grid_sizes |
| 1256 | ), |
| 1257 | video_grid_thw=MultiModalFieldConfig.batched("video", keep_on_cpu=True), |
| 1258 | ) |
| 1259 | |
| 1260 | |
| 1261 | class Ernie4_5_VLDummyInputsBuilder(BaseDummyInputsBuilder[Ernie4_5_VLProcessingInfo]): |
| 1262 | def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: |
| 1263 | num_images = mm_counts.get("image", 0) |
| 1264 | num_videos = mm_counts.get("video", 0) |
| 1265 | prompt = "" |
| 1266 | for i in range(num_images): |
| 1267 | prompt += ( |
| 1268 | f"Picture {i + 1}:<|IMAGE_START|><|image@placeholder|><|IMAGE_END|>" |
| 1269 | ) |
| 1270 | |
| 1271 | for i in range(num_videos): |
| 1272 | prompt += f"Video {i + 1}:<|VIDEO_START|><|video@placeholder|><|VIDEO_END|>" |
| 1273 | return prompt |
| 1274 | |
| 1275 | def get_dummy_mm_data( |
| 1276 | self, |
| 1277 | seq_len: int, |
| 1278 | mm_counts: Mapping[str, int], |
| 1279 | mm_options: MultiModalDummyOptions, |
| 1280 | ) -> MultiModalDataDict: |
| 1281 | target_width, target_height = self.info.get_image_size_with_most_features() |
| 1282 | target_num_frames = self.info.get_num_frames_with_most_features( |
| 1283 | seq_len, mm_counts |
| 1284 | ) |
| 1285 | |
| 1286 | return { |
| 1287 | "image": self._get_dummy_images( |
| 1288 | width=target_width, |
| 1289 | height=target_height, |
| 1290 | num_images=mm_counts.get("image", 0), |
| 1291 | overrides=mm_options.get("image"), |
| 1292 | ), |
| 1293 | "video": self._get_dummy_videos( |
| 1294 | width=target_width, |
| 1295 | height=target_height, |
| 1296 | num_frames=target_num_frames, |
| 1297 | num_videos=mm_counts.get("video", 0), |
| 1298 | overrides=mm_options.get("video"), |
| 1299 | ), |
| 1300 | } |
| 1301 | |
| 1302 | def _get_dummy_videos( |
| 1303 | self, |
| 1304 | *, |
| 1305 | width: int, |
| 1306 | height: int, |
| 1307 | num_frames: int, |
| 1308 | num_videos: int, |
| 1309 | overrides: VideoDummyOptions | None = None, |
| 1310 | ): |
| 1311 | # ernie4.5-vl requires at least 2 frames |
| 1312 | num_frames = max(num_frames, 2) |
| 1313 | if overrides and overrides.num_frames: |
| 1314 | overrides.num_frames = max(overrides.num_frames, 2) |
| 1315 | |
| 1316 | videos = super()._get_dummy_videos( |
| 1317 | width=width, |
| 1318 | height=height, |
| 1319 | num_frames=num_frames, |
| 1320 | num_videos=num_videos, |
| 1321 | overrides=overrides, |
| 1322 | ) |
| 1323 | videos = [v.copy() for v in videos] |
| 1324 | |
| 1325 | video_items = [] |
| 1326 | for video in videos: |
| 1327 | video_num_frames = video.shape[0] |
| 1328 | video_metadata = { |
| 1329 | "fps": 2.0, |
| 1330 | "duration": video_num_frames / 2.0, |
| 1331 | "total_num_frames": video_num_frames, |
| 1332 | "frames_indices": list(range(video_num_frames)), |
| 1333 | "video_backend": "opencv", |
| 1334 | "do_sample_frames": False, |
| 1335 | } |
| 1336 | video_items.append((video, video_metadata)) |
| 1337 | |
| 1338 | return video_items |
| 1339 | |
| 1340 | |
| 1341 | @MULTIMODAL_REGISTRY.register_processor( |
| 1342 | Ernie4_5VLMultiModalProcessor, |
| 1343 | info=Ernie4_5_VLProcessingInfo, |
| 1344 | dummy_inputs=Ernie4_5_VLDummyInputsBuilder, |
| 1345 | ) |
| 1346 | class Ernie4_5_VLMoeForConditionalGeneration( |
| 1347 | nn.Module, |
| 1348 | SupportsMultiModal, |
| 1349 | SupportsLoRA, |
| 1350 | SupportsPP, |
| 1351 | SupportsMRoPE, |
| 1352 | SupportsEncoderCudaGraph, |
| 1353 | ): |
| 1354 | packed_modules_mapping = { |
| 1355 | "qkv_proj": [ |
| 1356 | "q_proj", |
| 1357 | "k_proj", |
| 1358 | "v_proj", |
| 1359 | ], |
| 1360 | "gate_up_proj": [ |
| 1361 | "gate_proj", |
| 1362 | "up_proj", |
| 1363 | ], |
| 1364 | } |
| 1365 | |
| 1366 | # To ensure correct weight loading and mapping. |
| 1367 | hf_to_vllm_mapper = WeightsMapper( |
| 1368 | orig_to_new_prefix={ |
| 1369 | "lm_head.": "language_model.lm_head.", |
| 1370 | "model.": "language_model.model.", |
| 1371 | # model.resampler_model.-> language_model.model.resampler_model. |
| 1372 | # language_model.model.resampler_model. -> resampler_model. |
| 1373 | "language_model.model.resampler_model.": "resampler_model.", |
| 1374 | }, |
| 1375 | # resampler_weight_mappings |
| 1376 | orig_to_new_substr={ |
| 1377 | "spatial_linear.0.": "spatial_linear1.", |
| 1378 | "spatial_linear.2.": "spatial_linear2.", |
| 1379 | "spatial_linear.3.": "spatial_norm.", |
| 1380 | "temporal_linear.0.": "temporal_linear1.", |
| 1381 | "temporal_linear.2.": "temporal_linear2.", |
| 1382 | "temporal_linear.3.": "temporal_norm.", |
| 1383 | }, |
| 1384 | ) |
| 1385 | |
| 1386 | @classmethod |
| 1387 | def get_placeholder_str(cls, modality: str, i: int) -> str | None: |
| 1388 | if modality.startswith("image"): |
| 1389 | return "<|IMAGE_START|><|image@placeholder|><|IMAGE_END|>" |
| 1390 | if modality.startswith("video"): |
| 1391 | return "<|VIDEO_START|><|video@placeholder|><|VIDEO_END|>" |
| 1392 | |
| 1393 | raise ValueError("Only image or video modality is supported") |
| 1394 | |
| 1395 | def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: |
| 1396 | super().__init__() |
| 1397 | config = vllm_config.model_config.hf_config |
| 1398 | quant_config = vllm_config.quant_config |
| 1399 | multimodal_config = vllm_config.model_config.multimodal_config |
| 1400 | |
| 1401 | self.config = config |
| 1402 | self.multimodal_config = multimodal_config |
| 1403 | |
| 1404 | with self._mark_tower_model(vllm_config, {"image", "video"}): |
| 1405 | self.vision_model = Ernie4_5_VisionTransformer( |
| 1406 | config.vision_config, |
| 1407 | norm_eps=getattr(config, "rms_norm_eps", 1e-6), |
| 1408 | quant_config=quant_config, |
| 1409 | prefix=maybe_prefix(prefix, "vision_model"), |
| 1410 | ) |
| 1411 | self.resampler_model = VariableResolutionResamplerModel( |
| 1412 | self.config.pixel_hidden_size, |
| 1413 | self.config.hidden_size, |
| 1414 | self.config.spatial_conv_size, |
| 1415 | self.config.temporal_conv_size, |
| 1416 | config=self.config, |
| 1417 | prefix=maybe_prefix(prefix, "resampler_model"), |
| 1418 | ) |
| 1419 | |
| 1420 | with self._mark_language_model(vllm_config): |
| 1421 | self.language_model = Ernie4_5_VLMoeForCausalLM( |
| 1422 | vllm_config=vllm_config, |
| 1423 | prefix=maybe_prefix(prefix, "language_model"), |
| 1424 | ) |
| 1425 | |
| 1426 | self.visual_token_mask = None |
| 1427 | self.make_empty_intermediate_tensors = ( |
| 1428 | self.language_model.make_empty_intermediate_tensors |
| 1429 | ) |
| 1430 | if getattr(self.config, "im_patch_id", None): |
| 1431 | visual_token_ids = [ |
| 1432 | token_id |
| 1433 | for token_id in [ |
| 1434 | self.config.im_patch_id, |
| 1435 | getattr(self.config, "image_start_token_id", None), |
| 1436 | getattr(self.config, "image_end_token_id", None), |
| 1437 | getattr(self.config, "video_start_token_id", None), |
| 1438 | getattr(self.config, "video_end_token_id", None), |
| 1439 | ] |
| 1440 | if token_id is not None |
| 1441 | ] |
| 1442 | visual_token_ids_tensor_cache = torch.tensor( |
| 1443 | visual_token_ids, dtype=torch.long |
| 1444 | ) |
| 1445 | else: |
| 1446 | visual_token_ids_tensor_cache = None |
| 1447 | self.register_buffer( |
| 1448 | "_visual_token_ids_tensor_cache", |
| 1449 | visual_token_ids_tensor_cache, |
| 1450 | persistent=False, |
| 1451 | ) |
| 1452 | |
| 1453 | def compute_logits( |
| 1454 | self, |
| 1455 | hidden_states: torch.Tensor, |
| 1456 | ) -> torch.Tensor | None: |
| 1457 | return self.language_model.compute_logits(hidden_states) |
| 1458 | |
| 1459 | def _vision_forward( |
| 1460 | self, |
| 1461 | pixel_values: torch.Tensor, |
| 1462 | grid_thw: torch.Tensor, |
| 1463 | ) -> torch.Tensor: |
| 1464 | if grid_thw is not None: |
| 1465 | grid_thw = grid_thw[grid_thw > 0] |
| 1466 | if grid_thw.numel() % 3 != 0: |
| 1467 | raise ValueError( |
| 1468 | f"grid_thw has {grid_thw.numel()} elements after filtering," |
| 1469 | "which is not divisible by 3." |
| 1470 | ) |
| 1471 | grid_thw = grid_thw.reshape(-1, 3) |
| 1472 | # example: [[1,64,64],[2,80,80]] -> [[1,64,64],[1,80,80],[1,80,80]] |
| 1473 | grid_thw = F.pad( |
| 1474 | torch.repeat_interleave(grid_thw[:, 1:], grid_thw[:, 0], 0), |
| 1475 | [1, 0, 0, 0], |
| 1476 | value=1, |
| 1477 | ) |
| 1478 | image_features = self.vision_model(pixel_values, grid_thw) |
| 1479 | return image_features |
| 1480 | |
| 1481 | def _set_visual_token_mask(self, input_ids: torch.Tensor) -> None: |
| 1482 | """Set mask for visual tokens (image/video patches and delimiters).""" |
| 1483 | if self._visual_token_ids_tensor_cache is None: |
| 1484 | self.visual_token_mask = None |
| 1485 | return |
| 1486 | # Create tensor on the correct device |
| 1487 | visual_token_ids_tensor = self._visual_token_ids_tensor_cache.to( |
| 1488 | device=input_ids.device, |
| 1489 | dtype=input_ids.dtype, |
| 1490 | ) |
| 1491 | |
| 1492 | self.visual_token_mask = torch.isin(input_ids, visual_token_ids_tensor).reshape( |
| 1493 | -1, 1 |
| 1494 | ) |
| 1495 | |
| 1496 | def get_mrope_input_positions( |
| 1497 | self, |
| 1498 | input_tokens: list[int], |
| 1499 | mm_features: list[MultiModalFeatureSpec], |
| 1500 | ) -> tuple[torch.Tensor, int]: |
| 1501 | llm_pos_ids_list: list = [] |
| 1502 | st = 0 |
| 1503 | |
| 1504 | for ( |
| 1505 | offset, |
| 1506 | llm_grid_t, |
| 1507 | llm_grid_h, |
| 1508 | llm_grid_w, |
| 1509 | ) in self.iter_mm_grid_thw(mm_features): |
| 1510 | text_len = offset - st |
| 1511 | st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 |
| 1512 | llm_pos_ids_list.append( |
| 1513 | np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx |
| 1514 | ) |
| 1515 | |
| 1516 | grid_indices = np.indices((llm_grid_t, llm_grid_h, llm_grid_w)).reshape( |
| 1517 | 3, -1 |
| 1518 | ) |
| 1519 | llm_pos_ids_list.append(grid_indices + text_len + st_idx) |
| 1520 | st = offset + llm_grid_t * llm_grid_h * llm_grid_w |
| 1521 | |
| 1522 | if st < len(input_tokens): |
| 1523 | text_len = len(input_tokens) - st |
| 1524 | st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 |
| 1525 | llm_pos_ids_list.append( |
| 1526 | np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx |
| 1527 | ) |
| 1528 | |
| 1529 | llm_positions = np.concatenate(llm_pos_ids_list, axis=1).reshape(3, -1) |
| 1530 | mrope_position_delta = (llm_positions.max() + 1 - len(input_tokens)).item() |
| 1531 | return torch.from_numpy(llm_positions), mrope_position_delta |
| 1532 | |
| 1533 | def iter_mm_grid_thw( |
| 1534 | self, mm_features: list[MultiModalFeatureSpec] |
| 1535 | ) -> Iterator[tuple[int, int, int, int]]: |
| 1536 | spatial_conv_size = self.config.spatial_conv_size |
| 1537 | temporal_conv_size = self.config.temporal_conv_size |
| 1538 | |
| 1539 | for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset): |
| 1540 | if mm_feature.data is None: |
| 1541 | raise ValueError("M-RoPE calculation requires multimodal feature data") |
| 1542 | |
| 1543 | offset = mm_feature.mm_position.offset |
| 1544 | if mm_feature.modality == "image": |
| 1545 | grid_thw = mm_feature.data["image_grid_thw"].data |
| 1546 | assert isinstance(grid_thw, torch.Tensor) |
| 1547 | t, h, w = grid_thw.tolist() |
| 1548 | yield offset, t, h // spatial_conv_size, w // spatial_conv_size |
| 1549 | elif mm_feature.modality == "video": |
| 1550 | grid_thw = mm_feature.data["video_grid_thw"].data |
| 1551 | assert isinstance(grid_thw, torch.Tensor) |
| 1552 | t, h, w = grid_thw.tolist() |
| 1553 | yield ( |
| 1554 | offset, |
| 1555 | t // temporal_conv_size, |
| 1556 | h // spatial_conv_size, |
| 1557 | w // spatial_conv_size, |
| 1558 | ) |
| 1559 | else: |
| 1560 | raise ValueError(f"Unsupported modality: {mm_feature.modality}") |
| 1561 | |
| 1562 | def _parse_and_validate_image_input( |
| 1563 | self, **kwargs: object |
| 1564 | ) -> Ernie4_5_VLImageInputs | None: |
| 1565 | pixel_values = kwargs.pop("pixel_values", None) |
| 1566 | image_grid_thw = kwargs.pop("image_grid_thw", None) |
| 1567 | |
| 1568 | if pixel_values is None: |
| 1569 | return None |
| 1570 | |
| 1571 | return Ernie4_5_VLImagePixelInputs( |
| 1572 | type="pixel_values", |
| 1573 | pixel_values=pixel_values, |
| 1574 | image_grid_thw=image_grid_thw, |
| 1575 | ) |
| 1576 | |
| 1577 | def _parse_and_validate_video_input( |
| 1578 | self, **kwargs: object |
| 1579 | ) -> Ernie4_5_VLVideoInputs | None: |
| 1580 | pixel_values_videos = kwargs.pop("pixel_values_videos", None) |
| 1581 | video_grid_thw = kwargs.pop("video_grid_thw", None) |
| 1582 | |
| 1583 | if pixel_values_videos is None: |
| 1584 | return None |
| 1585 | |
| 1586 | return Ernie4_5_VLVideoPixelInputs( |
| 1587 | type="pixel_values_videos", |
| 1588 | pixel_values_videos=pixel_values_videos, |
| 1589 | video_grid_thw=video_grid_thw, |
| 1590 | ) |
| 1591 | |
| 1592 | # -- SupportsEncoderCudaGraph protocol methods -- |
| 1593 | # |
| 1594 | # Only the vision transformer (self.vision_model) is captured into the |
| 1595 | # CUDA graph. The VariableResolutionResamplerModel (spatial merge + |
| 1596 | # projection) runs eagerly in encoder_eager_forward / postprocess_ |
| 1597 | # encoder_output because its temporal path performs host-side (numpy) |
| 1598 | # indexing that cannot be captured. Image only: video uses a temporal |
| 1599 | # conv that changes the output token count, so it falls back to the |
| 1600 | # eager multimodal path for now. |
| 1601 | |
| 1602 | def get_encoder_cudagraph_config(self): |
| 1603 | from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphConfig |
| 1604 | |
| 1605 | return EncoderCudaGraphConfig( |
| 1606 | modalities=["image"], |
| 1607 | # Ernie consumes a single rotary freqs tensor (not cos/sin) |
| 1608 | buffer_keys=[ |
| 1609 | "pixel_values", |
| 1610 | "rotary_pos_emb", |
| 1611 | "cu_seqlens", |
| 1612 | "max_seqlen", |
| 1613 | ], |
| 1614 | # Post-merge embeddings, produced in postprocess_encoder_output, |
| 1615 | # are at the LM hidden dim, used only for DP gather sizing. |
| 1616 | out_hidden_size=self.config.hidden_size, |
| 1617 | ) |
| 1618 | |
| 1619 | def get_encoder_cudagraph_budget_range(self, vllm_config) -> tuple[int, int]: |
| 1620 | # Min: a 224x224 image -> 16x16 patches (patch_size=14), |
| 1621 | # spatial_merge_size=2 -> 8x8 = 64 output tokens. |
| 1622 | min_budget = 64 |
| 1623 | max_budget = min( |
| 1624 | vllm_config.scheduler_config.max_num_batched_tokens, |
| 1625 | vllm_config.model_config.max_model_len, |
| 1626 | ) |
| 1627 | return (min_budget, max_budget) |
| 1628 | |
| 1629 | def get_encoder_cudagraph_item_specs(self, mm_kwargs: dict[str, Any]): |
| 1630 | from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec |
| 1631 | |
| 1632 | m = self.vision_model.spatial_merge_size |
| 1633 | grid_thw = mm_kwargs["image_grid_thw"].tolist() |
| 1634 | return [ |
| 1635 | EncoderItemSpec( |
| 1636 | input_size=t * h * w, |
| 1637 | output_tokens=(t * h * w) // m // m, |
| 1638 | ) |
| 1639 | for t, h, w in grid_thw |
| 1640 | ] |
| 1641 | |
| 1642 | def select_encoder_cudagraph_items( |
| 1643 | self, mm_kwargs: dict[str, Any], indices: list[int] |
| 1644 | ) -> dict[str, Any]: |
| 1645 | grid_thw = mm_kwargs["image_grid_thw"] |
| 1646 | pixel_values = mm_kwargs["pixel_values"] |
| 1647 | |
| 1648 | if len(indices) == 0: |
| 1649 | return {"pixel_values": pixel_values[:0], "image_grid_thw": grid_thw[:0]} |
| 1650 | |
| 1651 | # Cumulative patch offsets for slicing the concatenated pixel_values. |
| 1652 | patches_per_item = (grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2]).tolist() |
| 1653 | cum_patches = [0] |
| 1654 | for p in patches_per_item: |
| 1655 | cum_patches.append(cum_patches[-1] + p) |
| 1656 | selected_pv = torch.cat( |
| 1657 | [pixel_values[cum_patches[i] : cum_patches[i + 1]] for i in indices] |
| 1658 | ) |
| 1659 | return {"pixel_values": selected_pv, "image_grid_thw": grid_thw[indices]} |
| 1660 | |
| 1661 | def prepare_encoder_cudagraph_capture_inputs( |
| 1662 | self, |
| 1663 | token_budget, |
| 1664 | max_batch_size, |
| 1665 | max_frames_per_batch, |
| 1666 | device, |
| 1667 | dtype, |
| 1668 | path: str = "default", |
| 1669 | axis_keys: tuple[Hashable, ...] | None = None, |
| 1670 | ): |
| 1671 | from vllm.v1.worker.encoder_cudagraph_defs import ( |
| 1672 | EncoderCudaGraphCaptureInputs, |
| 1673 | ) |
| 1674 | |
| 1675 | m = self.vision_model.spatial_merge_size |
| 1676 | # Ceil so the buffer fits the worst case of one item using the full |
| 1677 | # budget. |
| 1678 | per_mm_item_output = (token_budget + max_batch_size - 1) // max_batch_size |
| 1679 | # Image-format grid (T=1): h=m, w=per_mm_item_output*m, so the merged |
| 1680 | # token count per item equals per_mm_item_output. |
| 1681 | grid_config = [[1, m, per_mm_item_output * m] for _ in range(max_batch_size)] |
| 1682 | |
| 1683 | patch_embed = self.vision_model.patch_embed |
| 1684 | in_channels = patch_embed.in_channels |
| 1685 | patch_size = patch_embed.patch_size |
| 1686 | total_patches = sum(t * h * w for t, h, w in grid_config) |
| 1687 | flattened_patch_size = in_channels * patch_size * patch_size |
| 1688 | dummy_pixel_values = torch.randn( |
| 1689 | total_patches, flattened_patch_size, device=device, dtype=dtype |
| 1690 | ) |
| 1691 | |
| 1692 | # max_seqlen is baked at capture: worst case is one item consuming the |
| 1693 | # full budget - seq_len = token_budget * spatial_merge_size**2. |
| 1694 | metadata = self.vision_model.prepare_encoder_metadata( |
| 1695 | grid_config, |
| 1696 | max_batch_size=max_batch_size, |
| 1697 | max_seqlen_override=token_budget * (m**2), |
| 1698 | device=device, |
| 1699 | ) |
| 1700 | values = metadata | {"pixel_values": dummy_pixel_values} |
| 1701 | return EncoderCudaGraphCaptureInputs(values=values) |
| 1702 | |
| 1703 | def prepare_encoder_cudagraph_replay_buffers( |
| 1704 | self, mm_kwargs, max_batch_size, max_frames_per_batch, path: str = "default" |
| 1705 | ): |
| 1706 | from vllm.v1.worker.encoder_cudagraph_defs import ( |
| 1707 | EncoderCudaGraphReplayBuffers, |
| 1708 | ) |
| 1709 | |
| 1710 | # The per-image grids are needed as Python ints to size the buffers. |
| 1711 | with gpu_sync_allowed(): |
| 1712 | grid_thw_list = mm_kwargs["image_grid_thw"].tolist() |
| 1713 | metadata = self.vision_model.prepare_encoder_metadata( |
| 1714 | grid_thw_list, max_batch_size=max_batch_size |
| 1715 | ) |
| 1716 | values = metadata | { |
| 1717 | "pixel_values": mm_kwargs["pixel_values"], |
| 1718 | } |
| 1719 | return EncoderCudaGraphReplayBuffers(values=values) |
| 1720 | |
| 1721 | def encoder_cudagraph_forward( |
| 1722 | self, values: dict[str, torch.Tensor], path: str = "default" |
| 1723 | ) -> torch.Tensor: |
| 1724 | # Graph captures the ViT only, and the resampler runs in |
| 1725 | # postprocess_encoder_output. |
| 1726 | pixel_values = values.pop("pixel_values") |
| 1727 | return self.vision_model(pixel_values, encoder_metadata=values) |
| 1728 | |
| 1729 | def encoder_eager_forward( |
| 1730 | self, mm_kwargs: dict[str, Any], path: str = "default" |
| 1731 | ) -> torch.Tensor: |
| 1732 | # Eager fallback: run the full pipeline (ViT + resampler). The result |
| 1733 | # is scattered directly, so it must be the post-merge embeddings. |
| 1734 | pixel_values = mm_kwargs["pixel_values"].type(self.vision_model.dtype) |
| 1735 | grid_thw = async_tensor_h2d( |
| 1736 | mm_kwargs["image_grid_thw"], self.vision_model.device |
| 1737 | ) |
| 1738 | image_features = self.vision_model(pixel_values, grid_thw) |
| 1739 | return self.resampler_model(image_features, grid_thw) |
| 1740 | |
| 1741 | def postprocess_encoder_output( |
| 1742 | self, |
| 1743 | outputs: dict[str, torch.Tensor], |
| 1744 | indices: list[int], |
| 1745 | per_item_out_tokens: list[int], |
| 1746 | dest, |
| 1747 | clone: bool = False, |
| 1748 | batch_mm_kwargs: dict[str, Any] | None = None, |
| 1749 | ) -> None: |
| 1750 | # The graph output is the raw ViT features (pre-merge), padded to the |
| 1751 | # token budget. Run the resampler eagerly on the valid portion using |
| 1752 | # the actual batch grid_thw, then scatter the post-merge embeddings. |
| 1753 | # Ernie only uses the single "default" encoder path. |
| 1754 | output = outputs["default"] |
| 1755 | assert batch_mm_kwargs is not None |
| 1756 | grid_thw_cpu = batch_mm_kwargs["image_grid_thw"] |
| 1757 | grid_thw = async_tensor_h2d(grid_thw_cpu, output.device) |
| 1758 | # The valid token count slices the graph output for the eager |
| 1759 | # resampler call, so it has to come back to the host. |
| 1760 | num_valid = int( |
| 1761 | (grid_thw_cpu[:, 0] * grid_thw_cpu[:, 1] * grid_thw_cpu[:, 2]).sum() |
| 1762 | ) |
| 1763 | image_embeds = self.resampler_model(output[:num_valid], grid_thw) |
| 1764 | scatter_output_slices(image_embeds, indices, per_item_out_tokens, dest, clone) |
| 1765 | |
| 1766 | def _process_image_input( |
| 1767 | self, image_input: Ernie4_5_VLImageInputs |
| 1768 | ) -> tuple[torch.Tensor, ...]: |
| 1769 | grid_thw = image_input["image_grid_thw"] |
| 1770 | assert grid_thw.ndim == 2 |
| 1771 | |
| 1772 | pixel_values = image_input["pixel_values"].type(self.vision_model.dtype) |
| 1773 | image_features = self._vision_forward( |
| 1774 | pixel_values=pixel_values, grid_thw=grid_thw |
| 1775 | ) |
| 1776 | image_embeds = self.resampler_model(image_features, grid_thw) |
| 1777 | |
| 1778 | merge_size = self.vision_model.spatial_merge_size |
| 1779 | sizes = grid_thw.prod(-1) // merge_size // merge_size |
| 1780 | |
| 1781 | return image_embeds.split(sizes.tolist()) |
| 1782 | |
| 1783 | def _process_video_input( |
| 1784 | self, video_input: Ernie4_5_VLVideoInputs |
| 1785 | ) -> tuple[torch.Tensor, ...]: |
| 1786 | grid_thw = video_input["video_grid_thw"] |
| 1787 | assert grid_thw.ndim == 2 |
| 1788 | |
| 1789 | pixel_values_videos = video_input["pixel_values_videos"].type( |
| 1790 | self.vision_model.dtype |
| 1791 | ) |
| 1792 | video_features = self._vision_forward( |
| 1793 | pixel_values=pixel_values_videos, grid_thw=grid_thw |
| 1794 | ) |
| 1795 | video_embeds = self.resampler_model(video_features, grid_thw) |
| 1796 | |
| 1797 | merge_size = self.vision_model.spatial_merge_size |
| 1798 | sizes = ( |
| 1799 | (grid_thw.prod(-1) // self.config.temporal_conv_size) |
| 1800 | // merge_size |
| 1801 | // merge_size |
| 1802 | ) |
| 1803 | |
| 1804 | return video_embeds.split(sizes.tolist()) |
| 1805 | |
| 1806 | def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict: |
| 1807 | modalities: dict[ |
| 1808 | str, Ernie4_5_VLImageInputs | Ernie4_5_VLVideoInputs | None |
| 1809 | ] = {} |
| 1810 | |
| 1811 | # Preserve the order of modalities if there are multiple of them |
| 1812 | # from the order of kwargs. |
| 1813 | for input_key in kwargs: |
| 1814 | if ( |
| 1815 | input_key in ("pixel_values", "image_embeds") |
| 1816 | and "images" not in modalities |
| 1817 | ): |
| 1818 | modalities["images"] = self._parse_and_validate_image_input(**kwargs) |
| 1819 | if ( |
| 1820 | input_key in ("pixel_values_videos", "video_embeds") |
| 1821 | and "videos" not in modalities |
| 1822 | ): |
| 1823 | modalities["videos"] = self._parse_and_validate_video_input(**kwargs) |
| 1824 | |
| 1825 | return modalities |
| 1826 | |
| 1827 | def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings | None: |
| 1828 | modalities = self._parse_and_validate_multimodal_inputs(**kwargs) |
| 1829 | if not modalities: |
| 1830 | return None |
| 1831 | |
| 1832 | # The result multimodal_embeddings is tuple of tensors, with each |
| 1833 | # tensor corresponding to a multimodal data item (image or video). |
| 1834 | multimodal_embeddings: tuple[torch.Tensor, ...] = () |
| 1835 | |
| 1836 | # NOTE: It is important to iterate over the keys in this dictionary |
| 1837 | # to preserve the order of the modalities. |
| 1838 | for modality in modalities: |
| 1839 | if modality == "images": |
| 1840 | image_input = modalities["images"] |
| 1841 | image_embeddings = self._process_image_input(image_input) |
| 1842 | multimodal_embeddings += tuple(image_embeddings) |
| 1843 | if modality == "videos": |
| 1844 | video_input = modalities["videos"] |
| 1845 | video_embeddings = self._process_video_input(video_input) |
| 1846 | multimodal_embeddings += tuple(video_embeddings) |
| 1847 | |
| 1848 | return multimodal_embeddings |
| 1849 | |
| 1850 | def embed_input_ids( |
| 1851 | self, |
| 1852 | input_ids: torch.Tensor, |
| 1853 | multimodal_embeddings: MultiModalEmbeddings | None = None, |
| 1854 | *, |
| 1855 | is_multimodal: torch.Tensor | None = None, |
| 1856 | ) -> torch.Tensor: |
| 1857 | if multimodal_embeddings is not None and len(multimodal_embeddings) > 0: |
| 1858 | self._set_visual_token_mask(input_ids) |
| 1859 | |
| 1860 | # This is to satisfy the type checker for each overload |
| 1861 | if multimodal_embeddings is None or is_multimodal is None: |
| 1862 | return super().embed_input_ids(input_ids) |
| 1863 | |
| 1864 | return super().embed_input_ids( |
| 1865 | input_ids, |
| 1866 | multimodal_embeddings=multimodal_embeddings, |
| 1867 | is_multimodal=is_multimodal, |
| 1868 | ) |
| 1869 | |
| 1870 | def forward( |
| 1871 | self, |
| 1872 | input_ids: torch.Tensor | None, |
| 1873 | positions: torch.Tensor, |
| 1874 | intermediate_tensors: IntermediateTensors | None = None, |
| 1875 | inputs_embeds: torch.Tensor | None = None, |
| 1876 | **kwargs, |
| 1877 | ): |
| 1878 | forward_kwargs = { |
| 1879 | "input_ids": input_ids, |
| 1880 | "positions": positions, |
| 1881 | "intermediate_tensors": intermediate_tensors, |
| 1882 | "inputs_embeds": inputs_embeds, |
| 1883 | } |
| 1884 | |
| 1885 | if self.visual_token_mask is not None: |
| 1886 | if self.visual_token_mask.shape[0] != inputs_embeds.shape[0]: |
| 1887 | padding_len = inputs_embeds.shape[0] - self.visual_token_mask.shape[0] |
| 1888 | # right pad False |
| 1889 | pad = torch.zeros( |
| 1890 | (padding_len, self.visual_token_mask.shape[1]), |
| 1891 | dtype=self.visual_token_mask.dtype, |
| 1892 | device=self.visual_token_mask.device, |
| 1893 | ) |
| 1894 | self.visual_token_mask = torch.cat([self.visual_token_mask, pad], dim=0) |
| 1895 | |
| 1896 | forward_kwargs.update({"visual_token_mask": self.visual_token_mask}) |
| 1897 | self.visual_token_mask = None |
| 1898 | |
| 1899 | hidden_states = self.language_model.model( |
| 1900 | **forward_kwargs, |
| 1901 | **kwargs, |
| 1902 | ) |
| 1903 | |
| 1904 | return hidden_states |
| 1905 | |
| 1906 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 1907 | loader = AutoWeightsLoader(self) |
| 1908 | return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) |