| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # ruff: noqa: E501 |
| 4 | # Adapted from https://huggingface.co/moonshotai/Kimi-VL-A3B-Instruct/blob/main/modeling_kimi_vl.py |
| 5 | # This file is meant to be used in kimi_vl.py only |
| 6 | # Copyright 2025 The Moonshot AI Team, DeepSeek-AI, and HuggingFace Inc. team. All rights reserved. |
| 7 | # |
| 8 | # The code is based on llava (llava/modeling_llava.py) and DeepSeek-V3 (DeepSeek-V3/modeling_deepseek.py), but modified for KimiVL. |
| 9 | # |
| 10 | # Licensing Information: |
| 11 | # - Code derived from llava (llava/modeling_llava.py) and DeepSeek-V3 (DeepSeek-V3/modeling_deepseek.py) is licensed under the Apache License, Version 2.0. |
| 12 | # - Other parts of the code are licensed under the MIT License. |
| 13 | # |
| 14 | # Apache License, Version 2.0: |
| 15 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 16 | # you may not use this file except in compliance with the License. |
| 17 | # You may obtain a copy of the License at |
| 18 | # |
| 19 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 20 | # |
| 21 | # Unless required by applicable law or agreed to in writing, software |
| 22 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 23 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 24 | # See the License for the specific language governing permissions and |
| 25 | # limitations under the License. |
| 26 | # |
| 27 | # MIT License: |
| 28 | # Permission is hereby granted, free of charge, to any person obtaining a copy |
| 29 | # of this software and associated documentation files (the "Software"), to deal |
| 30 | # in the Software without restriction, including without limitation the rights |
| 31 | # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell |
| 32 | # copies of the Software, and to permit persons to whom the Software is |
| 33 | # furnished to do so, subject to the following conditions: |
| 34 | # |
| 35 | # The above copyright notice and this permission notice shall be included in all |
| 36 | # copies or substantial portions of the Software. |
| 37 | # |
| 38 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR |
| 39 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, |
| 40 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE |
| 41 | # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER |
| 42 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, |
| 43 | # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE |
| 44 | # SOFTWARE. |
| 45 | from collections.abc import Sequence |
| 46 | from copy import deepcopy |
| 47 | from functools import cached_property |
| 48 | from typing import Any |
| 49 | |
| 50 | import numpy as np |
| 51 | import torch |
| 52 | import torch.nn as nn |
| 53 | import torch.nn.functional as F |
| 54 | from transformers.activations import ACT2FN |
| 55 | from transformers.modeling_utils import PreTrainedModel |
| 56 | |
| 57 | from vllm.distributed import divide, get_tensor_model_parallel_world_size |
| 58 | from vllm.model_executor.layers.attention import MMEncoderAttention |
| 59 | from vllm.model_executor.layers.conv import Conv2dLayer |
| 60 | from vllm.model_executor.layers.linear import ( |
| 61 | ColumnParallelLinear, |
| 62 | QKVParallelLinear, |
| 63 | RowParallelLinear, |
| 64 | ) |
| 65 | from vllm.model_executor.models.utils import maybe_prefix |
| 66 | from vllm.model_executor.models.vision import is_vit_use_data_parallel |
| 67 | from vllm.platforms import current_platform |
| 68 | from vllm.transformers_utils.configs.moonvit import MoonViTConfig |
| 69 | from vllm.utils.gpu_sync_debug import gpu_sync_allowed |
| 70 | from vllm.utils.torch_utils import async_tensor_h2d |
| 71 | |
| 72 | |
| 73 | def _apply_rope_input_validation(x, freqs_cis): |
| 74 | assert x.ndim == freqs_cis.ndim + 1, (x.shape, freqs_cis.shape) |
| 75 | assert x.shape[:-2] == freqs_cis.shape[:-1], (x.shape, freqs_cis.shape) |
| 76 | assert x.shape[-1] == 2 * freqs_cis.shape[-1], (x.shape, freqs_cis.shape) |
| 77 | assert freqs_cis.dtype == torch.complex64, freqs_cis.dtype |
| 78 | |
| 79 | |
| 80 | def apply_rope( |
| 81 | xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor |
| 82 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 83 | """Args: (The leading dimensions of all inputs should be the same) |
| 84 | xq: query, tensor of shape (..., num_heads, head_dim) |
| 85 | xk: key, tensor of shape (..., num_heads, head_dim) |
| 86 | freqs_cis: tensor of shape (..., head_dim/2), dtype=torch.complex64. It contains the precomputed cis(freqs) for each position in the 2D grid. |
| 87 | |
| 88 | Returns: |
| 89 | xq_out, xk_out: tensors of shape (..., num_heads, head_dim) |
| 90 | |
| 91 | """ |
| 92 | _apply_rope_input_validation(xq, freqs_cis) |
| 93 | _apply_rope_input_validation(xk, freqs_cis) |
| 94 | |
| 95 | freqs_cis = freqs_cis.unsqueeze(-2) # ..., 1, head_dim/2 |
| 96 | # ..., num_heads, head_dim/2 |
| 97 | xq_ = torch.view_as_complex(xq.float().view(*xq.shape[:-1], -1, 2)) |
| 98 | xk_ = torch.view_as_complex(xk.float().view(*xq.shape[:-1], -1, 2)) |
| 99 | xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(-2) # ..., num_heads, head_dim |
| 100 | xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(-2) # ..., num_heads, head_dim |
| 101 | return xq_out.type_as(xq), xk_out.type_as(xk) |
| 102 | |
| 103 | |
| 104 | class Learnable2DInterpPosEmb(nn.Module): |
| 105 | def __init__( |
| 106 | self, height: int, width: int, dim: int, interpolation_mode: str = "bicubic" |
| 107 | ) -> None: |
| 108 | super().__init__() |
| 109 | self.height = height |
| 110 | self.width = width |
| 111 | self.interpolation_mode = interpolation_mode |
| 112 | self.weight = nn.Parameter(torch.empty(height, width, dim)) |
| 113 | self.reset_parameters() |
| 114 | |
| 115 | def reset_parameters(self): |
| 116 | nn.init.normal_(self.weight) |
| 117 | |
| 118 | def get_pos_embeds( |
| 119 | self, |
| 120 | grid_hws_list: list[list[int]] | list[tuple[int, int]], |
| 121 | ) -> torch.Tensor: |
| 122 | """Build packed per-token positional embeddings for a list of grids. |
| 123 | |
| 124 | Returns a tensor of shape ``(sum(h * w), dim)`` formed by interpolating |
| 125 | the learned ``(height, width, dim)`` weight to each ``(h, w)`` grid and |
| 126 | concatenating the flattened results in the same order as |
| 127 | ``grid_hws_list``. Lives outside the captured CUDA graph so the |
| 128 | per-grid Python iteration is safe. |
| 129 | """ |
| 130 | weight_shape = list(self.weight.shape[:-1]) |
| 131 | pos_embs: list[torch.Tensor] = [] |
| 132 | for shape in grid_hws_list: |
| 133 | shape_list = [int(shape[0]), int(shape[1])] |
| 134 | if shape_list == weight_shape: |
| 135 | pos_embs.append(self.weight.flatten(end_dim=1)) |
| 136 | else: |
| 137 | pos_embs.append( |
| 138 | F.interpolate( |
| 139 | self.weight.permute((2, 0, 1)).unsqueeze(0), |
| 140 | size=tuple(shape_list), |
| 141 | mode=self.interpolation_mode, |
| 142 | ) |
| 143 | .squeeze(0) |
| 144 | .permute((1, 2, 0)) |
| 145 | .flatten(end_dim=1) |
| 146 | ) |
| 147 | if not pos_embs: |
| 148 | return self.weight.new_zeros((0, self.weight.shape[-1])) |
| 149 | return torch.cat(pos_embs) |
| 150 | |
| 151 | def forward(self, x: torch.Tensor, grid_hws: torch.Tensor) -> torch.Tensor: |
| 152 | pos_embs = self.get_pos_embeds(grid_hws.tolist()) |
| 153 | out = x + pos_embs |
| 154 | return out |
| 155 | |
| 156 | |
| 157 | class MoonVisionPatchEmbed(nn.Module): |
| 158 | def __init__( |
| 159 | self, |
| 160 | out_dim: int, |
| 161 | in_dim: int = 3, |
| 162 | patch_size: int | tuple[int, int] = (14, 14), |
| 163 | pos_emb_height: int = 14, |
| 164 | pos_emb_width: int = 14, |
| 165 | ): |
| 166 | super().__init__() |
| 167 | assert isinstance(patch_size, (int, Sequence)), ( |
| 168 | f"Invalid patch_size type: {type(patch_size)}" |
| 169 | ) |
| 170 | if isinstance(patch_size, int): |
| 171 | patch_size = (patch_size, patch_size) |
| 172 | assert len(patch_size) == 2, ( |
| 173 | f"Expected patch_size to be a tuple of 2, got {patch_size}" |
| 174 | ) |
| 175 | self.patch_size = patch_size |
| 176 | |
| 177 | self.proj = Conv2dLayer( |
| 178 | in_dim, out_dim, kernel_size=patch_size, stride=patch_size |
| 179 | ) |
| 180 | |
| 181 | self.pos_emb = Learnable2DInterpPosEmb( |
| 182 | height=pos_emb_height, width=pos_emb_width, dim=out_dim |
| 183 | ) |
| 184 | |
| 185 | def forward( |
| 186 | self, |
| 187 | x: torch.Tensor, |
| 188 | grid_hw: torch.Tensor | None = None, |
| 189 | *, |
| 190 | pos_embeds: torch.Tensor | None = None, |
| 191 | ) -> torch.Tensor: |
| 192 | """Args: |
| 193 | x (L, Channels): input tensor |
| 194 | grid_hw (N, 2): grid height and width |
| 195 | pos_embeds: precomputed positional embeddings of shape |
| 196 | ``(L, Cout)``. When provided, ``grid_hw`` is unused and the |
| 197 | CUDA-graph-incompatible interpolation in ``self.pos_emb`` is |
| 198 | skipped. |
| 199 | |
| 200 | Returns: |
| 201 | (L, Cout) tensor |
| 202 | |
| 203 | """ |
| 204 | x = self.proj(x).view(x.size(0), -1) |
| 205 | if pos_embeds is not None: |
| 206 | return x + pos_embeds |
| 207 | return self.pos_emb(x, grid_hw) |
| 208 | |
| 209 | |
| 210 | class Rope2DPosEmb(nn.Module): |
| 211 | """2D rotary position embedding with multi-resolution support. |
| 212 | |
| 213 | This class is intended to be used in the following way: |
| 214 | 1. Before training, create an instance of Rope2DPosEmb. This instance will hold the precomputed cis. |
| 215 | 2. Before each forward pass, call `get_freqs_cis_by_*` to get the `freqs_cis` tensor for this iteration. |
| 216 | 3. During the forward pass, pass the `freqs_cis` tensor to each attention layer, and call `apply` just before each attention operation. |
| 217 | The rope is shared across all attention layers and all heads. |
| 218 | |
| 219 | Refs: |
| 220 | - RoFormer: https://arxiv.org/abs/2104.09864 |
| 221 | - VisionLLaMA: https://arxiv.org/abs/2403.00522 |
| 222 | - https://github.com/Meituan-AutoML/VisionLLaMA/blob/main/dit/models.py |
| 223 | |
| 224 | Args: |
| 225 | dim (int): usually the multi-head attention dimension, should be divisible by 4 (TODO: relax this constraint if needed) |
| 226 | max_height (int): the maximum height of the 2D grid |
| 227 | max_width (int): the maximum width of the 2D grid |
| 228 | theta_base (float): the base of the theta |
| 229 | device (str): the device to store the precomputed cis |
| 230 | |
| 231 | """ |
| 232 | |
| 233 | def __init__( |
| 234 | self, |
| 235 | dim: int, |
| 236 | max_height: int, |
| 237 | max_width: int, |
| 238 | theta_base=10000, |
| 239 | device=current_platform.device_type, |
| 240 | ): |
| 241 | super().__init__() |
| 242 | self.dim = dim |
| 243 | assert self.dim % 4 == 0, "dim must be divisible by 4" |
| 244 | self.max_height = max_height |
| 245 | self.max_width = max_width |
| 246 | self.theta_base = theta_base |
| 247 | self.device = device |
| 248 | |
| 249 | def extra_repr(self): |
| 250 | return f"dim={self.dim}, max_height={self.max_height}, max_width={self.max_width}, theta_base={self.theta_base}" |
| 251 | |
| 252 | @cached_property |
| 253 | def precomputed_freqs_cis(self) -> torch.Tensor: |
| 254 | """Calculate the cis(freqs) for each position in the 2D grid. |
| 255 | |
| 256 | Return: complex tensor of shape (max_height, max_width, dim//2) and value: |
| 257 | height axis: ret[h, w, 2*i] = cis(h * theta_base**(-4*i/dim)) |
| 258 | weight axis: ret[h, w, 2*i+1] = cis(w * theta_base**(-4*i/dim)) with (i in [0, dim//4)) |
| 259 | note: `cis` is a mathematical notation defined by cis x = cos x + i sin x, |
| 260 | """ |
| 261 | N = self.max_height * self.max_width |
| 262 | flat_pos = torch.arange(0, N).float().to(self.device) |
| 263 | x_pos = flat_pos % self.max_width |
| 264 | y_pos = flat_pos // self.max_width |
| 265 | dim_range = ( |
| 266 | torch.arange(0, self.dim, 4)[: (self.dim // 4)].float().to(self.device) |
| 267 | ) # C/4 |
| 268 | freqs = 1.0 / (self.theta_base ** (dim_range / self.dim)) |
| 269 | x_freqs = torch.outer(x_pos, freqs).float() # N, C/4 |
| 270 | y_freqs = torch.outer(y_pos, freqs).float() # N, C/4 |
| 271 | x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs) # N, C/4 |
| 272 | y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs) # N, C/4 |
| 273 | # N, C/4, 2 |
| 274 | freqs_cis = torch.cat( |
| 275 | [x_cis.unsqueeze(dim=-1), y_cis.unsqueeze(dim=-1)], dim=-1 |
| 276 | ) |
| 277 | # max_height, max_width, C/2 |
| 278 | freqs_cis = freqs_cis.reshape(self.max_height, self.max_width, -1) |
| 279 | return freqs_cis |
| 280 | |
| 281 | def get_freqs_cis_by_seqlens_list( |
| 282 | self, |
| 283 | grid_hws_list: list[list[int]] | list[tuple[int, int]], |
| 284 | ) -> torch.Tensor: |
| 285 | """List-based variant of :meth:`get_freqs_cis_by_seqlens`. |
| 286 | |
| 287 | Accepts a Python list of ``(h, w)`` pairs so callers that already |
| 288 | operate outside the captured CUDA graph can avoid materializing a |
| 289 | tensor + ``.tolist()`` round-trip. |
| 290 | """ |
| 291 | assert all( |
| 292 | 1 <= h <= self.max_height and 1 <= w <= self.max_width |
| 293 | for h, w in grid_hws_list |
| 294 | ), ( |
| 295 | grid_hws_list, |
| 296 | self.max_height, |
| 297 | self.max_width, |
| 298 | ) |
| 299 | if not grid_hws_list: |
| 300 | return self.precomputed_freqs_cis.new_zeros((0, self.dim // 2)) |
| 301 | freqs_cis = torch.cat( |
| 302 | [ |
| 303 | self.precomputed_freqs_cis[:h, :w].reshape(-1, self.dim // 2) |
| 304 | for h, w in grid_hws_list |
| 305 | ], |
| 306 | dim=0, |
| 307 | ) |
| 308 | return freqs_cis |
| 309 | |
| 310 | def get_freqs_cis_by_seqlens(self, grid_hws: torch.Tensor) -> torch.Tensor: |
| 311 | """Args: |
| 312 | grid_hws (torch.Tensor): containing list of (height, width) or (t, height, width) tuples. |
| 313 | |
| 314 | Returns: |
| 315 | freqs_cis: tensor of shape (sum(t * height * width), dim//2) |
| 316 | |
| 317 | """ |
| 318 | return self.get_freqs_cis_by_seqlens_list(grid_hws.tolist()) |
| 319 | |
| 320 | def get_freqs_cis_by_idx( |
| 321 | self, pos_idx: torch.Tensor, pos_idx_mask: torch.Tensor |
| 322 | ) -> torch.Tensor: |
| 323 | """Args: |
| 324 | pos_idx: tensor of shape (..., 2), It contains the (h, w) position indices of each 2D token. |
| 325 | pos_idx_mask: a mask of shape (...), the leading dimensions should be the same as pos_idx. |
| 326 | Rope will only be applied to the tokens with True mask. `freqs_cis` for the tokens with False mask with be ones. |
| 327 | |
| 328 | Return: |
| 329 | freqs_cis: tensor of shape (..., dim//2) |
| 330 | |
| 331 | """ |
| 332 | assert ( |
| 333 | pos_idx.shape[:-1] == pos_idx_mask.shape |
| 334 | and pos_idx.shape[-1] == 2 |
| 335 | and pos_idx.ndim == pos_idx_mask.ndim + 1 |
| 336 | ), (pos_idx.shape, pos_idx_mask.shape) |
| 337 | assert pos_idx_mask.dtype == torch.bool, pos_idx_mask.dtype |
| 338 | |
| 339 | shp = pos_idx_mask.shape + (self.dim // 2,) # ..., head_dim/2 |
| 340 | freqs_cis = torch.ones( |
| 341 | shp, dtype=torch.complex64, device=self.device |
| 342 | ) # ..., head_dim/2 |
| 343 | freqs_cis[pos_idx_mask] = self.precomputed_freqs_cis[ |
| 344 | pos_idx[..., 0][pos_idx_mask], pos_idx[..., 1][pos_idx_mask] |
| 345 | ] |
| 346 | return freqs_cis |
| 347 | |
| 348 | |
| 349 | class MLP2(nn.Module): |
| 350 | """Args: |
| 351 | dims: [in_dim, hidden_dim, out_dim] |
| 352 | bias: whether to use bias in linear layer. |
| 353 | |
| 354 | """ |
| 355 | |
| 356 | def __init__( |
| 357 | self, |
| 358 | dims: list[int], |
| 359 | activation, |
| 360 | bias: bool = True, |
| 361 | prefix: str = "", |
| 362 | ): |
| 363 | super().__init__() |
| 364 | assert len(dims) == 3 |
| 365 | self.use_data_parallel = is_vit_use_data_parallel() |
| 366 | self.fc0 = ColumnParallelLinear( |
| 367 | dims[0], |
| 368 | dims[1], |
| 369 | bias=bias, |
| 370 | prefix=maybe_prefix(prefix, "fc0"), |
| 371 | disable_tp=self.use_data_parallel, |
| 372 | ) |
| 373 | self.fc1 = RowParallelLinear( |
| 374 | dims[1], |
| 375 | dims[2], |
| 376 | bias=bias, |
| 377 | prefix=maybe_prefix(prefix, "fc1"), |
| 378 | disable_tp=self.use_data_parallel, |
| 379 | ) |
| 380 | self.activation = activation |
| 381 | |
| 382 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 383 | x, _ = self.fc0(x) |
| 384 | x = self.activation(x) |
| 385 | x, _ = self.fc1(x) |
| 386 | return x |
| 387 | |
| 388 | |
| 389 | class MoonVitEncoderLayer(nn.Module): |
| 390 | def __init__( |
| 391 | self, |
| 392 | num_heads: int, |
| 393 | hidden_dim: int, |
| 394 | mlp_dim: int, |
| 395 | prefix: str = "", |
| 396 | *, |
| 397 | activation=F.gelu, |
| 398 | attn_bias: bool = False, |
| 399 | ): |
| 400 | super().__init__() |
| 401 | self.use_data_parallel = is_vit_use_data_parallel() |
| 402 | |
| 403 | self.num_heads = num_heads |
| 404 | self.hidden_dim = hidden_dim |
| 405 | self.hidden_size_per_attention_head = self.hidden_dim // self.num_heads |
| 406 | self.tp_size = ( |
| 407 | 1 if self.use_data_parallel else get_tensor_model_parallel_world_size() |
| 408 | ) |
| 409 | self.num_attention_heads_per_partition = divide(num_heads, self.tp_size) |
| 410 | |
| 411 | self.norm0 = nn.LayerNorm(hidden_dim) |
| 412 | self.norm1 = nn.LayerNorm(hidden_dim) |
| 413 | self.mlp = MLP2( |
| 414 | [hidden_dim, mlp_dim, hidden_dim], |
| 415 | activation, |
| 416 | prefix=f"{prefix}.mlp", |
| 417 | ) |
| 418 | self.wqkv = QKVParallelLinear( |
| 419 | hidden_size=hidden_dim, |
| 420 | head_size=self.hidden_size_per_attention_head, |
| 421 | total_num_heads=num_heads, |
| 422 | total_num_kv_heads=num_heads, |
| 423 | bias=attn_bias, |
| 424 | prefix=f"{prefix}.wqkv", |
| 425 | disable_tp=self.use_data_parallel, |
| 426 | ) |
| 427 | self.wo = RowParallelLinear( |
| 428 | hidden_dim, |
| 429 | hidden_dim, |
| 430 | bias=attn_bias, |
| 431 | prefix=f"{prefix}.wo", |
| 432 | disable_tp=self.use_data_parallel, |
| 433 | ) |
| 434 | self.attn = MMEncoderAttention( |
| 435 | num_heads=self.num_attention_heads_per_partition, |
| 436 | head_size=self.hidden_size_per_attention_head, |
| 437 | scale=self.hidden_size_per_attention_head**-0.5, |
| 438 | prefix=f"{prefix}.attn", |
| 439 | ) |
| 440 | |
| 441 | def attention_qkvpacked( |
| 442 | self, |
| 443 | x: torch.Tensor, |
| 444 | cu_seqlens: torch.Tensor, |
| 445 | rope_freqs_cis: torch.Tensor | None = None, |
| 446 | max_seqlen: torch.Tensor | None = None, |
| 447 | ): |
| 448 | """Args: |
| 449 | x (torch.Tensor): (seqlen, hidden_dim) |
| 450 | cu_seqlens (torch.Tensor): |
| 451 | max_seqlen: Optional precomputed scalar tensor. When omitted it |
| 452 | is derived from ``cu_seqlens``, which produces a GPU scalar |
| 453 | that breaks CUDA graph capture. |
| 454 | |
| 455 | """ |
| 456 | seq_length = x.size(0) |
| 457 | xqkv, _ = self.wqkv(x) |
| 458 | |
| 459 | qkv_shape = xqkv.size()[:-1] + ( |
| 460 | 3, |
| 461 | self.num_attention_heads_per_partition, |
| 462 | self.hidden_size_per_attention_head, |
| 463 | ) |
| 464 | # xqkv: (batch_size, seqlen, 3, nheads, headdim) |
| 465 | xqkv = xqkv.view(*qkv_shape) |
| 466 | xq, xk, xv = torch.unbind(xqkv, dim=-3) |
| 467 | |
| 468 | xq, xk = apply_rope(xq, xk, rope_freqs_cis) |
| 469 | |
| 470 | if max_seqlen is None: |
| 471 | max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max() |
| 472 | attn_out = self.attn( |
| 473 | xq.unsqueeze(0), |
| 474 | xk.unsqueeze(0), |
| 475 | xv.unsqueeze(0), |
| 476 | cu_seqlens=cu_seqlens, |
| 477 | max_seqlen=max_seqlen, |
| 478 | ) |
| 479 | attn_out = attn_out.reshape( |
| 480 | seq_length, |
| 481 | self.num_attention_heads_per_partition |
| 482 | * self.hidden_size_per_attention_head, |
| 483 | ) |
| 484 | attn_out, _ = self.wo(attn_out) |
| 485 | return attn_out |
| 486 | |
| 487 | def forward( |
| 488 | self, |
| 489 | hidden_states: torch.Tensor, |
| 490 | cu_seqlens: torch.Tensor, |
| 491 | rope_freqs_cis: torch.Tensor | None = None, |
| 492 | max_seqlen: torch.Tensor | None = None, |
| 493 | ) -> torch.Tensor: |
| 494 | """Args: |
| 495 | hidden_states: non-packed (B, N, D) or packed (L, D). if non-packed, seqlens should be None, if packed, seqlens should be set |
| 496 | max_seqlen: optional precomputed max-sequence-length scalar. |
| 497 | |
| 498 | Returns: |
| 499 | output: same shape of input, non-packed (B, N, D) for non-packed input, (L, D) for packed input |
| 500 | |
| 501 | """ |
| 502 | residual = hidden_states |
| 503 | hidden_states = self.norm0(hidden_states) |
| 504 | attn_out = self.attention_qkvpacked( |
| 505 | hidden_states, |
| 506 | cu_seqlens, |
| 507 | rope_freqs_cis=rope_freqs_cis, |
| 508 | max_seqlen=max_seqlen, |
| 509 | ) |
| 510 | hidden_states = residual + attn_out |
| 511 | |
| 512 | residual = hidden_states |
| 513 | hidden_states = self.mlp(self.norm1(hidden_states)) |
| 514 | hidden_states = residual + hidden_states |
| 515 | return hidden_states |
| 516 | |
| 517 | |
| 518 | class MoonVitEncoder(nn.Module): |
| 519 | def __init__( |
| 520 | self, |
| 521 | hidden_dim: int, |
| 522 | num_layers: int, |
| 523 | block_cfg: dict, |
| 524 | prefix: str = "", |
| 525 | ) -> None: |
| 526 | super().__init__() |
| 527 | |
| 528 | self.rope_2d = Rope2DPosEmb( |
| 529 | block_cfg["hidden_dim"] // block_cfg["num_heads"], 512, 512 |
| 530 | ) |
| 531 | self.blocks = nn.ModuleList( |
| 532 | [ |
| 533 | MoonVitEncoderLayer( |
| 534 | prefix=f"{prefix}.blocks.{layer_idx}", |
| 535 | **block_cfg, |
| 536 | ) |
| 537 | for layer_idx in range(num_layers) |
| 538 | ] |
| 539 | ) |
| 540 | self.final_layernorm = nn.LayerNorm(hidden_dim) |
| 541 | |
| 542 | def get_rope_freqs_cis( |
| 543 | self, |
| 544 | grid_hws_list: list[list[int]] | list[tuple[int, int]], |
| 545 | ) -> torch.Tensor: |
| 546 | return self.rope_2d.get_freqs_cis_by_seqlens_list(grid_hws_list) |
| 547 | |
| 548 | def forward( |
| 549 | self, |
| 550 | hidden_states: torch.Tensor, |
| 551 | grid_hw: torch.Tensor | None = None, |
| 552 | *, |
| 553 | cu_seqlens: torch.Tensor | None = None, |
| 554 | rope_freqs_cis: torch.Tensor | None = None, |
| 555 | max_seqlen: torch.Tensor | None = None, |
| 556 | ) -> torch.Tensor: |
| 557 | if rope_freqs_cis is None: |
| 558 | assert grid_hw is not None |
| 559 | rope_freqs_cis = self.rope_2d.get_freqs_cis_by_seqlens(grid_hws=grid_hw) |
| 560 | |
| 561 | if cu_seqlens is None: |
| 562 | assert grid_hw is not None |
| 563 | lengths = torch.cat( |
| 564 | ( |
| 565 | torch.zeros(1, device=hidden_states.device, dtype=grid_hw.dtype), |
| 566 | (grid_hw[:, 0] * grid_hw[:, 1]).to(hidden_states.device), |
| 567 | ) |
| 568 | ) |
| 569 | cu_seqlens = lengths.cumsum(dim=0, dtype=torch.int32) |
| 570 | |
| 571 | for _, block in enumerate(self.blocks): |
| 572 | hidden_states = block( |
| 573 | hidden_states, |
| 574 | cu_seqlens, |
| 575 | rope_freqs_cis=rope_freqs_cis, |
| 576 | max_seqlen=max_seqlen, |
| 577 | ) |
| 578 | |
| 579 | hidden_states = self.final_layernorm(hidden_states) |
| 580 | |
| 581 | return hidden_states |
| 582 | |
| 583 | |
| 584 | def patch_merger( |
| 585 | x: torch.Tensor, |
| 586 | grid_hw: torch.Tensor, |
| 587 | merge_kernel_size: tuple[int, int] = (2, 2), |
| 588 | ) -> list[torch.Tensor]: |
| 589 | d_model = x.size(-1) |
| 590 | |
| 591 | outputs = [] |
| 592 | pre_sum = 0 |
| 593 | for x_shape in grid_hw.tolist(): |
| 594 | height, width = x_shape[0], x_shape[1] |
| 595 | # Get the current sequence |
| 596 | seq = x[pre_sum : pre_sum + height * width] |
| 597 | # Reshape along self.merge_kernel_size and concat to the last dimension |
| 598 | kernel_height, kernel_width = merge_kernel_size |
| 599 | new_height, new_width = height // kernel_height, width // kernel_width |
| 600 | reshaped_seq = seq.view( |
| 601 | new_height, kernel_height, new_width, kernel_width, d_model |
| 602 | ) |
| 603 | reshaped_seq = reshaped_seq.permute(0, 2, 1, 3, 4).contiguous() |
| 604 | padded_seq = reshaped_seq.view( |
| 605 | new_height * new_width, kernel_height * kernel_width, -1 |
| 606 | ) |
| 607 | outputs.append(padded_seq) |
| 608 | pre_sum += height * width |
| 609 | |
| 610 | return outputs |
| 611 | |
| 612 | |
| 613 | def patch_merger_packed( |
| 614 | x: torch.Tensor, |
| 615 | gather_idx: torch.Tensor, |
| 616 | merge_kernel_size: tuple[int, int], |
| 617 | ) -> torch.Tensor: |
| 618 | """CUDA-graph-safe equivalent of :func:`patch_merger`. |
| 619 | |
| 620 | Uses a precomputed index tensor to gather the per-token reshape + |
| 621 | permute that ``patch_merger`` does inside a Python loop. The output is |
| 622 | the concatenated 3D tensor ``(sum(new_h * new_w), kh * kw, d_model)``, |
| 623 | matching what ``torch.cat(patch_merger(...))`` would produce. |
| 624 | """ |
| 625 | kh, kw = merge_kernel_size |
| 626 | d_model = x.size(-1) |
| 627 | return x.index_select(0, gather_idx).view(-1, kh * kw, d_model) |
| 628 | |
| 629 | |
| 630 | def _build_merge_gather_idx( |
| 631 | grid_hws_list: list[list[int]] | list[tuple[int, int]], |
| 632 | merge_kernel_size: tuple[int, int], |
| 633 | ) -> np.ndarray: |
| 634 | """Build the per-token gather indices used by :func:`patch_merger_packed`. |
| 635 | |
| 636 | For each item with grid (h, w) and merge kernel (kh, kw), the output |
| 637 | block at position (nh, nw) gathers the kh*kw input tokens at rows |
| 638 | (nh*kh + ih, nw*kw + iw) of that item, in (ih, iw) row-major order. |
| 639 | """ |
| 640 | kh, kw = merge_kernel_size |
| 641 | parts: list[np.ndarray] = [] |
| 642 | pre_sum = 0 |
| 643 | for h, w in grid_hws_list: |
| 644 | new_h, new_w = h // kh, w // kw |
| 645 | nh = np.arange(new_h, dtype=np.int64).reshape(new_h, 1, 1, 1) |
| 646 | nw = np.arange(new_w, dtype=np.int64).reshape(1, new_w, 1, 1) |
| 647 | ih = np.arange(kh, dtype=np.int64).reshape(1, 1, kh, 1) |
| 648 | iw = np.arange(kw, dtype=np.int64).reshape(1, 1, 1, kw) |
| 649 | # Linearized input row = (nh*kh + ih) * w + (nw*kw + iw), offset by |
| 650 | # the per-item base ``pre_sum``. Output is laid out as |
| 651 | # (new_h, new_w, kh, kw) which patch_merger flattens to |
| 652 | # (new_h*new_w, kh*kw). |
| 653 | idx = pre_sum + (nh * kh + ih) * w + (nw * kw + iw) |
| 654 | parts.append(idx.reshape(-1)) |
| 655 | pre_sum += h * w |
| 656 | if not parts: |
| 657 | return np.zeros(0, dtype=np.int64) |
| 658 | return np.concatenate(parts) |
| 659 | |
| 660 | |
| 661 | class MoonVitPretrainedModel(PreTrainedModel): |
| 662 | config_class = MoonViTConfig |
| 663 | model_type = "moonvit" |
| 664 | _no_split_modules = ["PackingTransformer"] |
| 665 | _supports_flash_attn_2 = True |
| 666 | _supports_sdpa = True |
| 667 | |
| 668 | def __init__( |
| 669 | self, |
| 670 | config: MoonViTConfig, |
| 671 | prefix: str = "", |
| 672 | *inputs, |
| 673 | **kwargs, |
| 674 | ): |
| 675 | super().__init__(config, *inputs, **kwargs) |
| 676 | config = deepcopy(config) |
| 677 | self.merge_kernel_size = config.merge_kernel_size |
| 678 | self.hidden_size = config.hidden_size |
| 679 | self.patch_size = config.patch_size |
| 680 | self.vit_processing_type = "rope_2d" |
| 681 | self.patch_embed = MoonVisionPatchEmbed( |
| 682 | out_dim=config.hidden_size, |
| 683 | patch_size=config.patch_size, |
| 684 | pos_emb_height=config.init_pos_emb_height, |
| 685 | pos_emb_width=config.init_pos_emb_width, |
| 686 | ) |
| 687 | |
| 688 | self.encoder = MoonVitEncoder( |
| 689 | hidden_dim=config.hidden_size, |
| 690 | num_layers=config.num_hidden_layers, |
| 691 | block_cfg={ |
| 692 | "num_heads": config.num_attention_heads, |
| 693 | "hidden_dim": config.hidden_size, |
| 694 | "mlp_dim": config.intermediate_size, |
| 695 | "activation": ACT2FN["gelu_pytorch_tanh"], |
| 696 | "attn_bias": True, |
| 697 | }, |
| 698 | prefix=f"{prefix}.encoder", |
| 699 | ) |
| 700 | |
| 701 | def prepare_encoder_metadata( |
| 702 | self, |
| 703 | grid_hws_list: list[list[int]] | list[tuple[int, int]], |
| 704 | *, |
| 705 | max_batch_size: int | None = None, |
| 706 | max_seqlen_override: int | None = None, |
| 707 | device: torch.device | None = None, |
| 708 | ) -> dict[str, Any]: |
| 709 | """Precompute every grid-dependent input the encoder needs. |
| 710 | |
| 711 | Used by the CUDA graph capture and replay paths to precompute |
| 712 | every grid-dependent input outside the captured graph, so per-grid |
| 713 | Python iteration and ``.tolist()`` round-trips are fine; the |
| 714 | values are then copied into fixed-shape buffers for replay. |
| 715 | |
| 716 | Args: |
| 717 | grid_hws_list: List of ``(h, w)`` patch-grid sizes per image. |
| 718 | max_batch_size: When set, ``cu_seqlens`` is right-padded with |
| 719 | its last value so the buffer covers up to this many |
| 720 | sequences. Required at CUDA graph capture/replay so the |
| 721 | buffer shape matches what was recorded; padding entries |
| 722 | are zero-length sequences and are ignored by varlen |
| 723 | attention. |
| 724 | max_seqlen_override: Override the per-replay max sequence |
| 725 | length scalar. At capture this must be a safe upper bound |
| 726 | (worst case: a single image consuming the full token |
| 727 | budget) because the value is baked into the captured |
| 728 | graph. |
| 729 | device: Device for the metadata tensors. Defaults to the |
| 730 | model's parameter device. |
| 731 | |
| 732 | """ |
| 733 | if device is None: |
| 734 | device = next(self.parameters()).device |
| 735 | |
| 736 | # Normalize to a list of plain Python int pairs so the helpers |
| 737 | # below never need ``.tolist()`` on a tensor. |
| 738 | grid_pairs: list[tuple[int, int]] = [(int(h), int(w)) for h, w in grid_hws_list] |
| 739 | |
| 740 | metadata: dict[str, Any] = {} |
| 741 | |
| 742 | pos_embeds = self.patch_embed.pos_emb.get_pos_embeds(grid_pairs) |
| 743 | metadata["pos_embeds"] = pos_embeds.to(device=device) |
| 744 | |
| 745 | rope_freqs_cis = self.encoder.get_rope_freqs_cis(grid_pairs) |
| 746 | metadata["rope_freqs_cis"] = rope_freqs_cis.to(device=device) |
| 747 | |
| 748 | grid_arr = np.array(grid_pairs, dtype=np.int64) |
| 749 | seq_lens = (grid_arr[:, 0] * grid_arr[:, 1]).astype(np.int32) |
| 750 | cu_seqlens_np = np.concatenate( |
| 751 | [ |
| 752 | np.zeros(1, dtype=np.int32), |
| 753 | seq_lens.cumsum(dtype=np.int32), |
| 754 | ] |
| 755 | ) |
| 756 | |
| 757 | if max_batch_size is not None: |
| 758 | num_seqs = len(cu_seqlens_np) - 1 |
| 759 | if num_seqs < max_batch_size: |
| 760 | cu_seqlens_np = np.concatenate( |
| 761 | [ |
| 762 | cu_seqlens_np, |
| 763 | np.full( |
| 764 | max_batch_size - num_seqs, |
| 765 | cu_seqlens_np[-1], |
| 766 | dtype=np.int32, |
| 767 | ), |
| 768 | ] |
| 769 | ) |
| 770 | metadata["cu_seqlens"] = async_tensor_h2d(cu_seqlens_np, device=device) |
| 771 | |
| 772 | if max_seqlen_override is not None: |
| 773 | max_seqlen_val = int(max_seqlen_override) |
| 774 | else: |
| 775 | max_seqlen_val = int(seq_lens.max()) if len(seq_lens) > 0 else 0 |
| 776 | # Keep on CPU: attention wrappers may call .item() on this scalar |
| 777 | # and we want that materialization to happen outside the captured |
| 778 | # graph (the value is constant per capture anyway). |
| 779 | metadata["max_seqlen"] = torch.tensor(max_seqlen_val, dtype=torch.int32) |
| 780 | |
| 781 | gather_idx_np = _build_merge_gather_idx(grid_pairs, self.merge_kernel_size) |
| 782 | metadata["merge_gather_idx"] = async_tensor_h2d(gather_idx_np, device=device) |
| 783 | |
| 784 | return metadata |
| 785 | |
| 786 | def forward( |
| 787 | self, |
| 788 | pixel_values: torch.Tensor, |
| 789 | grid_hw: torch.Tensor, |
| 790 | *, |
| 791 | encoder_metadata: dict[str, Any] | None = None, |
| 792 | ) -> torch.Tensor | list[torch.Tensor]: |
| 793 | """Args: |
| 794 | pixel_values (torch.Tensor): The input pixel values. |
| 795 | grid_hw (torch.Tensor): The grid height and width. |
| 796 | encoder_metadata: Optional precomputed metadata produced by |
| 797 | :meth:`prepare_encoder_metadata`. When provided every |
| 798 | ``.tolist()`` call in the forward path is skipped, the |
| 799 | returned tensor is the packed |
| 800 | ``(sum(new_h*new_w), kh*kw, hidden_size)`` form (suitable |
| 801 | for CUDA graph capture/replay), and ``grid_hw`` is unused. |
| 802 | When ``None`` the legacy path runs and returns a list of |
| 803 | per-image tensors. |
| 804 | |
| 805 | """ |
| 806 | if encoder_metadata is not None: |
| 807 | hidden_states = self.patch_embed( |
| 808 | pixel_values, pos_embeds=encoder_metadata["pos_embeds"] |
| 809 | ) |
| 810 | hidden_states = self.encoder( |
| 811 | hidden_states, |
| 812 | cu_seqlens=encoder_metadata["cu_seqlens"], |
| 813 | rope_freqs_cis=encoder_metadata["rope_freqs_cis"], |
| 814 | max_seqlen=encoder_metadata["max_seqlen"], |
| 815 | ) |
| 816 | return patch_merger_packed( |
| 817 | hidden_states, |
| 818 | encoder_metadata["merge_gather_idx"], |
| 819 | merge_kernel_size=self.merge_kernel_size, |
| 820 | ) |
| 821 | |
| 822 | # Legacy path: patch_embed, encoder and patch_merger each iterate the |
| 823 | # per-image grids in Python, so all three read `grid_hw` back to the |
| 824 | # host. The `encoder_metadata` path above precomputes that outside the |
| 825 | # graph and does not sync. |
| 826 | with gpu_sync_allowed(): |
| 827 | hidden_states = self.patch_embed(pixel_values, grid_hw) |
| 828 | hidden_states = self.encoder(hidden_states, grid_hw) |
| 829 | hidden_states = patch_merger( |
| 830 | hidden_states, grid_hw, merge_kernel_size=self.merge_kernel_size |
| 831 | ) |
| 832 | return hidden_states |