PYTHAI/Kimi-K2.7-Code-fork
022
1# coding=utf-82# Copyright 2025-2026 The Moonshot AI Team, DeepSeek-AI, and HuggingFace Inc. team. All rights reserved.3#4# The code is based on llava (llava/modeling_llava.py) and DeepSeek-V3 (DeepSeek-V3/modeling_deepseek.py), but modified for Kimi-K2.5.5#6# Licensing Information:7# - 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.8# - Other parts of the code are licensed under the MIT License.9#10# Apache License, Version 2.0:11# Licensed under the Apache License, Version 2.0 (the "License");12# you may not use this file except in compliance with the License.13# You may obtain a copy of the License at14#15# http://www.apache.org/licenses/LICENSE-2.016#17# Unless required by applicable law or agreed to in writing, software18# distributed under the License is distributed on an "AS IS" BASIS,19# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.20# See the License for the specific language governing permissions and21# limitations under the License.22#23# MIT License:24# Permission is hereby granted, free of charge, to any person obtaining a copy25# of this software and associated documentation files (the "Software"), to deal26# in the Software without restriction, including without limitation the rights27# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell28# copies of the Software, and to permit persons to whom the Software is29# furnished to do so, subject to the following conditions:30#31# The above copyright notice and this permission notice shall be included in all32# copies or substantial portions of the Software.33#34# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR35# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,36# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE37# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER38# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,39# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE40# SOFTWARE.41import math42from collections.abc import Sequence43from copy import deepcopy44from typing import Optional45 46import numpy as np47import torch48import torch.nn as nn49import torch.nn.functional as F50from transformers import activations51 52try:53 from transformers.activations import PytorchGELUTanh54except ImportError:55 from transformers.activations import GELUTanh56 activations.PytorchGELUTanh = GELUTanh57 PytorchGELUTanh = GELUTanh58from transformers.activations import PytorchGELUTanh59from transformers.cache_utils import Cache60from transformers.configuration_utils import PretrainedConfig61from transformers.modeling_utils import PreTrainedModel62from transformers.models.llava.modeling_llava import \63 LlavaCausalLMOutputWithPast64from transformers.utils import is_flash_attn_2_available65 66from .configuration_kimi_k25 import KimiK25Config67from .modeling_deepseek import DeepseekV3ForCausalLM68 69# Flash attention imports70if is_flash_attn_2_available():71 from flash_attn import flash_attn_varlen_func72else:73 flash_attn_varlen_func = None74 75 76def multihead_attention(77 q: torch.Tensor,78 k: torch.Tensor,79 v: torch.Tensor,80 q_cu_seqlens: torch.Tensor | None = None,81 k_cu_seqlens: torch.Tensor | None = None,82 max_seqlen_q: int | None = None,83 max_seqlen_k: int | None = None,84 deterministic: bool = False,85):86 """Multi-head attention using flash attention 2.87 88 Args:89 q, k, v: tensor of shape (batch_size, seqlen, num_heads, head_dim),90 or (tot_seqlens, num_heads, head_dim) if packing.91 q_cu_seqlens (torch.Tensor): cumulative sequence lengths of q.92 The first element should be 0 and the last element should be q.shape[0].93 k_cu_seqlens (torch.Tensor): cumulative sequence lengths of k.94 The first element should be 0 and the last element should be k.shape[0].95 96 Returns:97 output: shape (batch_size, seqlen, dim) or (tot_seqlens, dim) if packing,98 where dim = num_heads * head_dim99 """100 attn_out = flash_attn_varlen_func(101 q,102 k,103 v,104 q_cu_seqlens,105 k_cu_seqlens,106 max_seqlen_q,107 max_seqlen_k,108 causal=False,109 deterministic=deterministic,110 )111 if isinstance(attn_out, tuple):112 attn_out = attn_out[0]113 114 attn_out = attn_out.flatten(start_dim=-2)115 116 return attn_out117 118 119def eager_attention(120 q: torch.Tensor,121 k: torch.Tensor,122 v: torch.Tensor,123 q_cu_seqlens: Optional[torch.Tensor] = None,124 k_cu_seqlens: Optional[torch.Tensor] = None,125 **kwargs,126) -> torch.Tensor:127 seq_length = q.shape[0]128 attention_mask = torch.zeros([1, seq_length, seq_length],129 device=q.device,130 dtype=torch.bool)131 for i in range(1, len(q_cu_seqlens)):132 attention_mask[133 ...,134 q_cu_seqlens[i - 1]:q_cu_seqlens[i],135 q_cu_seqlens[i - 1]:q_cu_seqlens[i],136 ] = True137 q = q.transpose(0, 1)138 k = k.transpose(0, 1)139 v = v.transpose(0, 1)140 141 attn_weight = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1])142 attn_weight += attention_mask143 attn_weight = torch.softmax(attn_weight, dim=-1,144 dtype=torch.float32).to(q.dtype)145 146 attn_output = attn_weight @ v147 attn_output = attn_output.transpose(0, 1)148 attn_output = attn_output.reshape(seq_length, -1)149 return attn_output150 151 152VL_VISION_ATTENTION_FUNCTIONS = {153 "flash_attention_2": multihead_attention,154 "eager": eager_attention,155}156 157 158def _apply_rope_input_validation(x, freqs_cis):159 assert x.ndim == freqs_cis.ndim + 1, (x.shape, freqs_cis.shape)160 assert x.shape[:-2] == freqs_cis.shape[:-1], (x.shape, freqs_cis.shape)161 assert x.shape[-1] == 2 * freqs_cis.shape[-1], (x.shape, freqs_cis.shape)162 assert freqs_cis.dtype == torch.complex64, freqs_cis.dtype163 164 165def get_rope_shape_decorate(func):166 _get_rope_shape_first_call_flag = set()167 168 def wrapper(org, interpolation_mode, shape):169 key = (org.requires_grad, torch.is_grad_enabled(), interpolation_mode)170 if key not in _get_rope_shape_first_call_flag:171 _get_rope_shape_first_call_flag.add(key)172 _ = func(org, interpolation_mode, shape=(64, 64))173 return func(org, interpolation_mode, shape)174 175 return wrapper176 177 178@get_rope_shape_decorate179@torch.compile(dynamic=True)180def get_rope_shape(org, interpolation_mode, shape):181 return (F.interpolate(182 org.permute((2, 0, 1)).unsqueeze(0),183 size=shape,184 mode=interpolation_mode,185 ).squeeze(0).permute((1, 2, 0)).flatten(end_dim=1))186 187 188def apply_rope(xq: torch.Tensor, xk: torch.Tensor,189 freqs_cis: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:190 """191 Args: (The leading dimensions of all inputs should be the same)192 xq: query, tensor of shape (..., num_heads, head_dim)193 xk: key, tensor of shape (..., num_heads, head_dim)194 freqs_cis: tensor of shape (..., head_dim/2), dtype=torch.complex64. It contains the precomputed cis(freqs) for each position in the 2D grid.195 Returns:196 xq_out, xk_out: tensors of shape (..., num_heads, head_dim)197 """198 _apply_rope_input_validation(xq, freqs_cis)199 _apply_rope_input_validation(xk, freqs_cis)200 201 freqs_cis = freqs_cis.unsqueeze(-2) # ..., 1, head_dim/2202 # ..., num_heads, head_dim/2203 xq_ = torch.view_as_complex(xq.float().view(*xq.shape[:-1], -1, 2))204 xk_ = torch.view_as_complex(xk.float().view(*xq.shape[:-1], -1, 2))205 xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(206 -2) # ..., num_heads, head_dim207 xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(208 -2) # ..., num_heads, head_dim209 return xq_out.type_as(xq), xk_out.type_as(xk)210 211 212def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):213 """214 From:215 https://github.com/OpenGVLab/InternVideo/blob/421f6d2361fc8f61a3394244571f2601a4e99e29/InternVideo2/multi_modality/models/backbones/internvideo2/pos_embed.py#L86216 embed_dim: output dimension for each position217 pos: a list of positions to be encoded: size (M,)218 out: (M, D)219 """220 assert embed_dim % 2 == 0221 omega = np.arange(embed_dim // 2, dtype=np.float32)222 omega /= embed_dim / 2.0223 omega = 1.0 / 10000**omega # (D/2,)224 225 pos = pos.reshape(-1) # (M,)226 out = np.einsum('m,d->md', pos, omega) # (M, D/2), outer product227 228 emb_sin = np.sin(out) # (M, D/2)229 emb_cos = np.cos(out) # (M, D/2)230 231 emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D)232 return emb233 234 235def get_1d_sincos_pos_embed(embed_dim, t_size, cls_token=False):236 """237 t_size: int of the temporal size238 return:239 pos_embed: [t_size, embed_dim] or [1+t_size, embed_dim] (w/ or w/o cls_token)240 """241 grid_t = np.arange(t_size, dtype=np.float32)242 pos_embed = get_1d_sincos_pos_embed_from_grid(embed_dim, grid_t)243 if cls_token:244 pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed],245 axis=0)246 return pos_embed247 248 249def _first_layer_key_first_token_vector(past_key_values):250 """``past_key_values[0][0][..., 0]`` for LLaVA-style cache masking (shape ``[batch, heads, seq]``).251 Legacy caches are ``list`` of ``(key, value)`` per layer. Transformers v4.36+ / v5 use ``Cache`` (e.g.252 ``DynamicCache``) with per-layer ``.keys`` tensors instead of subscripting ``[0][0]``.253 """254 if isinstance(past_key_values, Cache):255 layers = getattr(past_key_values, "layers", None) or []256 if not layers:257 return None258 layer0 = layers[0]259 keys = getattr(layer0, "keys", None)260 if keys is None or keys.numel() == 0 or keys.ndim < 4:261 return None262 return keys[:, :, :, 0]263 return past_key_values[0][0][:, :, :, 0]264 265 266def _first_layer_past_seq_length(past_key_values):267 """Layer-0 KV cache sequence length (BHSD keys: ``shape[2] == seq_len``).268 """269 if isinstance(past_key_values, Cache):270 try:271 return int(past_key_values.get_seq_length(0))272 except Exception:273 return None274 try:275 k0 = past_key_values[0][0]276 if k0 is None or k0.ndim < 3:277 return None278 return int(k0.shape[2])279 except Exception:280 return None281 282 283class Learnable2DInterpPosEmbDivided_fixed(nn.Module):284 285 def __init__(self,286 height: int,287 width: int,288 num_frames: int,289 dim: int,290 interpolation_mode: str = 'bicubic') -> None:291 super().__init__()292 self.height = height293 self.width = width294 self.num_frames = num_frames295 self.dim = dim296 self.interpolation_mode = interpolation_mode297 self.weight = nn.Parameter(torch.empty(height, width, dim))298 self.register_buffer('time_weight',299 torch.from_numpy(300 get_1d_sincos_pos_embed(301 self.dim,302 self.num_frames)).float().unsqueeze(1),303 persistent=False)304 305 self.reset_parameters()306 307 def reset_parameters(self):308 nn.init.normal_(self.weight)309 310 def forward(self, x: torch.Tensor,311 grid_thws: torch.Tensor) -> torch.Tensor:312 pos_embs = []313 for t, h, w in grid_thws.tolist():314 assert t <= self.num_frames, f't:{t} > self.num_frames:{self.num_frames}'315 if (h, w) == self.weight.shape[:-1]:316 pos_emb_2d = self.weight.flatten(end_dim=1)317 else:318 pos_emb_2d = get_rope_shape(319 self.weight,320 interpolation_mode=self.interpolation_mode,321 shape=(h, w),322 )323 324 if t == 1:325 pos_emb_3d = pos_emb_2d326 else:327 pos_emb_3d = pos_emb_2d.unsqueeze(0).repeat(328 t, 1, 1) + self.time_weight[0:t]329 330 pos_embs.append(pos_emb_3d.reshape(-1, pos_emb_3d.shape[-1]))331 332 out = x + torch.cat(pos_embs)333 return out334 335 336class MoonVision3dPatchEmbed(nn.Module):337 338 def __init__(self,339 out_dim: int,340 in_dim: int = 3,341 patch_size: int | tuple[int, int] = (14, 14),342 pos_emb_height: int = 14,343 pos_emb_width: int = 14,344 pos_emb_time: int = 4,345 pos_emb_type: str = 'divided_fixed'):346 super().__init__()347 assert isinstance(348 patch_size,349 int | Sequence), f'Invalid patch_size type: {type(patch_size)}'350 if isinstance(patch_size, int):351 patch_size = (patch_size, patch_size)352 assert (len(patch_size) == 2353 ), f'Expected patch_size to be a tuple of 2, got {patch_size}'354 self.patch_size = patch_size355 356 self.proj = nn.Conv2d(in_dim,357 out_dim,358 kernel_size=patch_size,359 stride=patch_size)360 361 if pos_emb_type == 'divided_fixed':362 self.pos_emb = Learnable2DInterpPosEmbDivided_fixed(363 height=pos_emb_height,364 width=pos_emb_width,365 num_frames=pos_emb_time,366 dim=out_dim)367 else:368 raise NotImplementedError(369 f'Not support pos_emb_type: {pos_emb_type}')370 371 def forward(self, x: torch.Tensor,372 grid_thws: torch.Tensor) -> torch.Tensor:373 """374 Args:375 x (L, Channels): input tensor376 grid_hws (N, 3): temporal, height and width377 378 Returns:379 (L, Cout) tensor380 """381 x = self.proj(x).view(x.size(0), -1)382 # apply positional embedding383 x = self.pos_emb(x, grid_thws)384 return x385 386 387class Rope2DPosEmbRepeated(nn.Module):388 """2D rotary position embedding with multi-resolution support.389 390 This class is intended to be used in the following way:391 1. Before training, create an instance of Rope2DPosEmb. This instance will hold the precomputed cis.392 2. Before each forward pass, call `get_freqs_cis_by_*` to get the `freqs_cis` tensor for this iteration.393 3. During the forward pass, pass the `freqs_cis` tensor to each attention layer, and call `apply` just before each attention operation.394 The rope is shared across all attention layers and all heads.395 396 Refs:397 - RoFormer: https://arxiv.org/abs/2104.09864398 - VisionLLaMA: https://arxiv.org/abs/2403.00522399 - https://github.com/Meituan-AutoML/VisionLLaMA/blob/main/dit/models.py400 401 Args:402 dim (int): usually the multi-head attention dimension, should be divisible by 4 (TODO: relax this constraint if needed)403 max_height (int): the maximum height of the 2D grid404 max_width (int): the maximum width of the 2D grid405 theta_base (float): the base of the theta406 device (str): the device to store the precomputed cis407 """408 409 def __init__(self,410 dim: int,411 max_height: int,412 max_width: int,413 theta_base=10000):414 super().__init__()415 self.dim = dim416 assert self.dim % 4 == 0, 'dim must be divisible by 4'417 self.max_height = max_height418 self.max_width = max_width419 self.theta_base = theta_base420 421 def extra_repr(self):422 return f'dim={self.dim}, max_height={self.max_height}, max_width={self.max_width}, theta_base={self.theta_base}'423 424 def _precompute_freqs_cis(self, device: torch.device) -> torch.Tensor:425 """Calculate the cis(freqs) for each position in the 2D grid.426 427 Return: complex tensor of shape (max_height, max_width, dim//2) and value:428 height axis: ret[h, w, 2*i] = cis(h * theta_base**(-4*i/dim))429 weight axis: ret[h, w, 2*i+1] = cis(w * theta_base**(-4*i/dim)) with (i in [0, dim//4))430 note: `cis` is a mathematical notation defined by cis x = cos x + i sin x,431 """432 N = self.max_height * self.max_width433 flat_pos = torch.arange(0, N).float().to(device)434 x_pos = flat_pos % self.max_width435 y_pos = flat_pos // self.max_width436 dim_range = (torch.arange(0, self.dim,437 4)[:(self.dim // 4)].float().to(device)438 ) # C/4439 freqs = 1.0 / (self.theta_base**(dim_range / self.dim))440 x_freqs = torch.outer(x_pos, freqs).float() # N, C/4441 y_freqs = torch.outer(y_pos, freqs).float() # N, C/4442 x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs) # N, C/4443 y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs) # N, C/4444 # N, C/4, 2445 freqs_cis = torch.cat(446 [x_cis.unsqueeze(dim=-1),447 y_cis.unsqueeze(dim=-1)], dim=-1)448 # max_height, max_width, C/2449 freqs_cis = freqs_cis.reshape(self.max_height, self.max_width, -1)450 return freqs_cis451 452 def get_freqs_cis(self, grid_thws: torch.Tensor,453 device: torch.device) -> torch.Tensor:454 """455 Args:456 grid_thws (torch.Tensor): grid time, height and width457 458 Returns:459 freqs_cis: tensor of shape (sum(t * height * width), dim//2)460 """461 if not hasattr(self, 'freqs_cis'):462 self.register_buffer('freqs_cis',463 self._precompute_freqs_cis(device),464 persistent=False)465 466 shapes = grid_thws.tolist()467 assert all(1 <= h <= self.max_height and 1 <= w <= self.max_width468 for t, h, w in shapes), (469 shapes,470 self.max_height,471 self.max_width,472 )473 freqs_cis = torch.cat(474 [475 self.freqs_cis[:h, :w].reshape(-1, self.dim // 2).repeat(t, 1)476 for t, h, w in shapes477 ],478 dim=0,479 )480 return freqs_cis481 482 483class MLP2(nn.Module):484 """485 Args:486 dims: [in_dim, hidden_dim, out_dim]487 bias: whether to use bias in linear layer.488 """489 490 def __init__(self, dims: list[int], activation, bias=True):491 super().__init__()492 assert len(dims) == 3493 self.fc0 = nn.Linear(dims[0], dims[1], bias=bias)494 self.fc1 = nn.Linear(dims[1], dims[2], bias=bias)495 self.activation = activation496 for m in [self.fc0, self.fc1]:497 nn.init.trunc_normal_(m.weight, std=math.sqrt(2 / m.in_features))498 if m.bias is not None:499 nn.init.zeros_(m.bias)500 501 def forward(self, x: torch.Tensor) -> torch.Tensor:502 x = self.fc0(x)503 x = self.activation(x)504 return self.fc1(x)505 506 507class MoonViTEncoderLayer(nn.Module):508 509 def __init__(510 self,511 num_heads: int,512 hidden_dim: int,513 mlp_dim: int,514 *,515 attn_implementation: str = 'flash_attention_2',516 activation=F.gelu,517 attn_bias: bool = False,518 use_deterministic_attn: bool = False,519 ):520 super().__init__()521 self.num_heads = num_heads522 self.hidden_dim = hidden_dim523 self.hidden_size_per_attention_head = self.hidden_dim // self.num_heads524 self.attn_implementation = attn_implementation525 self.use_deterministic_attn = use_deterministic_attn526 527 self.norm0 = nn.LayerNorm(hidden_dim)528 self.norm1 = nn.LayerNorm(hidden_dim)529 self.mlp = MLP2([hidden_dim, mlp_dim, hidden_dim], activation)530 self.wqkv = nn.Linear(hidden_dim, hidden_dim * 3, bias=attn_bias)531 self.wo = nn.Linear(hidden_dim, hidden_dim, bias=attn_bias)532 533 def attention_qkvpacked(534 self,535 x: torch.Tensor,536 cu_seqlens: torch.Tensor,537 max_seqlen: torch.Tensor,538 rope_freqs_cis: torch.Tensor | None = None,539 ):540 """541 Args:542 x (torch.Tensor): (batch_size, seqlen, hidden_dim)543 cu_seqlens (torch.Tensor):544 """545 xqkv = self.wqkv(x)546 547 qkv_shape = xqkv.size()[:-1] + (548 3,549 self.num_heads,550 self.hidden_size_per_attention_head,551 )552 # xqkv: (batch_size, seqlen, 3, nheads, headdim)553 xqkv = xqkv.view(*qkv_shape)554 xq, xk, xv = torch.unbind(xqkv, dim=-3)555 556 xq, xk = apply_rope(xq, xk, rope_freqs_cis)557 558 attn_func = VL_VISION_ATTENTION_FUNCTIONS[self.attn_implementation]559 attn_out = attn_func(xq,560 xk,561 xv,562 q_cu_seqlens=cu_seqlens,563 k_cu_seqlens=cu_seqlens,564 max_seqlen_k=max_seqlen,565 max_seqlen_q=max_seqlen,566 deterministic=self.use_deterministic_attn)567 568 attn_out = self.wo(attn_out)569 return attn_out570 571 def forward(572 self,573 hidden_states: torch.Tensor,574 cu_seqlens: torch.Tensor,575 max_seqlen: int,576 rope_freqs_cis: torch.Tensor | None = None,577 ):578 residual = hidden_states579 hidden_states = self.norm0(hidden_states)580 581 hidden_states = self.attention_qkvpacked(hidden_states, cu_seqlens,582 max_seqlen, rope_freqs_cis)583 hidden_states = residual + hidden_states584 585 residual = hidden_states586 hidden_states = self.norm1(hidden_states)587 hidden_states = self.mlp(hidden_states)588 hidden_states = residual + hidden_states589 590 return hidden_states591 592 593class MoonViT3dEncoder(nn.Module):594 595 def __init__(self,596 hidden_dim: int,597 num_layers: int,598 block_cfg: dict,599 video_attn_type: str = 'spatial_temporal',600 use_deterministic_attn: bool = False) -> None:601 super().__init__()602 603 assert video_attn_type == 'spatial_temporal', f'video_attn_type must be "spatial_temporal", got {video_attn_type}'604 self.video_attn_type = video_attn_type605 self.rope_2d = Rope2DPosEmbRepeated(606 block_cfg['hidden_dim'] // block_cfg['num_heads'], 512, 512)607 self.blocks = nn.ModuleList([608 MoonViTEncoderLayer(**block_cfg,609 use_deterministic_attn=use_deterministic_attn)610 for _ in range(num_layers)611 ])612 self.final_layernorm = nn.LayerNorm(hidden_dim)613 614 def forward(615 self,616 hidden_states: torch.Tensor,617 grid_thws: torch.Tensor,618 ) -> torch.Tensor:619 rope_freqs_cis = self.rope_2d.get_freqs_cis(620 grid_thws=grid_thws, device=hidden_states.device)621 622 lengths = torch.cat((623 torch.zeros(1, dtype=grid_thws.dtype, device=grid_thws.device),624 grid_thws[:, 0] * grid_thws[:, 1] * grid_thws[:, 2],625 ))626 627 max_seqlen = lengths.max()628 cu_seqlens = lengths.to(hidden_states.device).cumsum(dim=0,629 dtype=torch.int32)630 for block in self.blocks:631 hidden_states = block(hidden_states,632 cu_seqlens,633 max_seqlen,634 rope_freqs_cis=rope_freqs_cis)635 636 hidden_states = self.final_layernorm(hidden_states)637 return hidden_states638 639 640def tpool_patch_merger(641 x: torch.Tensor,642 grid_thws: torch.Tensor,643 merge_kernel_size: tuple[int, int] = (2, 2),644) -> list[torch.Tensor]:645 d_model = x.size(-1)646 647 outputs = []648 pre_sum = 0649 for t, h, w in grid_thws.tolist():650 # Get the current sequence651 seq = x[pre_sum:pre_sum + t * h * w]652 # Reshape along self.merge_kernel_size and concat to the last dimension653 kernel_height, kernel_width = merge_kernel_size654 new_height, new_width = h // kernel_height, w // kernel_width655 reshaped_seq = seq.view(t, new_height, kernel_height, new_width,656 kernel_width, d_model)657 reshaped_seq = reshaped_seq.permute(0, 1,658 3, 2, 4, 5).contiguous().mean(659 dim=0) # temporal pooling660 padded_seq = reshaped_seq.view(new_height * new_width,661 kernel_height * kernel_width, -1)662 outputs.append(padded_seq)663 pre_sum += t * h * w664 665 return outputs666 667 668class MoonViT3dPretrainedModel(PreTrainedModel):669 config_class = None670 model_type = 'moonvit3d'671 _no_split_modules = ['PackingTransformer']672 _supports_flash_attn_2 = True673 _supports_flash_attn = True674 _supports_sdpa = True675 676 def __init__(self, config, *inputs, **kwargs):677 super().__init__(config, *inputs, **kwargs)678 config = deepcopy(config)679 self.merge_kernel_size = config.merge_kernel_size680 self.patch_size = config.patch_size681 self.merge_type = config.merge_type682 683 self.patch_embed = MoonVision3dPatchEmbed(684 out_dim=config.hidden_size,685 patch_size=config.patch_size,686 pos_emb_height=config.init_pos_emb_height,687 pos_emb_width=config.init_pos_emb_width,688 pos_emb_time=config.init_pos_emb_time,689 pos_emb_type=config.pos_emb_type,690 )691 692 self.encoder = MoonViT3dEncoder(hidden_dim=config.hidden_size,693 num_layers=config.num_hidden_layers,694 block_cfg={695 'num_heads':696 config.num_attention_heads,697 'hidden_dim':698 config.hidden_size,699 'mlp_dim':700 config.intermediate_size,701 'activation':702 PytorchGELUTanh(),703 'attn_bias':704 True,705 'attn_implementation':706 config._attn_implementation,707 },708 video_attn_type=config.video_attn_type)709 710 def forward(self, pixel_values: torch.Tensor,711 grid_thws: torch.Tensor) -> torch.Tensor:712 """713 Args:714 pixel_values (torch.Tensor): The input pixel values.715 grid_thws (torch.Tensor): Temporal, height and width.716 717 Returns:718 torch.Tensor: The output tokens.719 """720 # grid_thws = grid_thws.to('cpu')721 assert grid_thws.ndim == 2, f'grid_thws should be 2D, got {grid_thws.ndim}'722 assert grid_thws.size(1) == 3, f'No support for thw: {grid_thws}'723 hidden_states = self.patch_embed(pixel_values, grid_thws)724 hidden_states = self.encoder(hidden_states, grid_thws)725 if self.merge_type == 'sd2_tpool': # spatial downsampling 2x with temporal pooling all726 hidden_states = tpool_patch_merger(727 hidden_states,728 grid_thws,729 merge_kernel_size=self.merge_kernel_size)730 else:731 raise NotImplementedError(f'Not support {self.merge_type}')732 733 return hidden_states734 735 736# ============================================================================737# MM Projector Helper Classes (from mm_projector/modeling_mm_projectors.py)738# ============================================================================739 740 741class IdentityMap(nn.Module):742 743 def __init__(self):744 super().__init__()745 746 def forward(self, x, *args, **kwargs):747 return x748 749 750class MLP(nn.Module):751 752 def __init__(self, config):753 super().__init__()754 # TODO, use faster LayerNorm755 self.pre_norm = nn.LayerNorm(config.mm_hidden_size)756 self.proj = nn.Sequential(757 nn.Linear(config.mm_hidden_size, config.hidden_size), nn.GELU(),758 nn.Linear(config.hidden_size, config.hidden_size))759 760 def forward(self, x, *args, **kwargs):761 assert isinstance(x,762 list | tuple), f'x is not a list or tuple: {type(x)}'763 lengths = [item.shape[0] for item in x]764 x = torch.cat(x, dim=0)765 x = self.pre_norm(x)766 x = self.proj(x)767 x = torch.split(x, lengths, dim=0)768 769 return x770 771 772class PatchMergerMLP(nn.Module):773 774 def __init__(self, config):775 super().__init__()776 eps = config.projector_ln_eps777 self.hidden_size = config.mm_hidden_size * (778 config.merge_kernel_size[0] * config.merge_kernel_size[1])779 self.pre_norm = nn.LayerNorm(config.mm_hidden_size, eps=eps)780 self.proj = nn.Sequential(781 nn.Linear(self.hidden_size, self.hidden_size),782 nn.GELU(),783 nn.Linear(self.hidden_size, config.hidden_size),784 )785 786 def forward(self, x, *args, **kwargs):787 if isinstance(x, list) or isinstance(x, tuple):788 x = [789 self.proj(self.pre_norm(item).view(item.shape[0], -1))790 for item in x791 ]792 else:793 # B, N, N_k, C = x.shape794 B = x.shape[0]795 x = self.proj(self.pre_norm(x).view(B, -1, self.hidden_size))796 return x797 798 799class KimiK25PreTrainedModel(PreTrainedModel):800 config_class = KimiK25Config801 base_model_prefix = "model"802 _no_split_modules = [803 "MoonViT3dPretrainedModel",804 "MoonViTEncoderLayer",805 "DeepseekDecoderLayer",806 "PatchMergerMLP",807 ]808 _skip_keys_device_placement = "past_key_values"809 _supports_flash_attn_2 = True810 _supports_flash_attn = True811 _supports_sdpa = False812 813 def _init_weights(self, module):814 # important: this ported version of Llava isn't meant for training from scratch - only815 # inference and fine-tuning - so the proper init weights code has been removed - the original codebase816 # https://github.com/haotian-liu/LLaVA/tree/main/llava should serve for that purpose817 std = (self.config.initializer_range if hasattr(818 self.config, "initializer_range") else819 self.config.text_config.initializer_range)820 821 if hasattr(module, "class_embedding"):822 module.class_embedding.data.normal_(mean=0.0, std=std)823 824 if isinstance(module, (nn.Linear, nn.Conv2d)):825 module.weight.data.normal_(mean=0.0, std=std)826 if module.bias is not None:827 module.bias.data.zero_()828 elif isinstance(module, nn.Embedding):829 module.weight.data.normal_(mean=0.0, std=std)830 if module.padding_idx is not None:831 module.weight.data[module.padding_idx].zero_()832 833 834class VisionTowerConfig(PretrainedConfig):835 model_type = 'moonvit3d'836 837 def __init__(self, config: KimiK25Config, **kwargs):838 super().__init__(**kwargs)839 self.patch_size = config.patch_size840 self.init_pos_emb_height = config.init_pos_emb_height841 self.init_pos_emb_width = config.init_pos_emb_width842 self.init_pos_emb_time = config.init_pos_emb_time843 self.pos_emb_type = config.pos_emb_type844 self.num_attention_heads = config.vt_num_attention_heads845 self.num_hidden_layers = config.vt_num_hidden_layers846 self.hidden_size = config.vt_hidden_size847 self.intermediate_size = config.vt_intermediate_size848 self.merge_kernel_size = config.merge_kernel_size849 self.video_attn_type = config.video_attn_type850 self.merge_type = config.merge_type851 self._attn_implementation = config._attn_implementation852 853 854class ProjectorConfig:855 856 def __init__(self, config: KimiK25Config):857 self.mm_projector_type = config.mm_projector_type858 self.mm_hidden_size = config.mm_hidden_size859 self.hidden_size = config.text_hidden_size860 self.merge_kernel_size = config.merge_kernel_size861 self.projector_hidden_act = config.projector_hidden_act862 self.projector_ln_eps = config.projector_ln_eps863 864 865# ref https://github.com/huggingface/transformers/blob/78b2929c0554b79e0489b451ce4ece14d265ead2/src/transformers/models/llava/modeling_llava.py#L240866class KimiK25ForConditionalGeneration(KimiK25PreTrainedModel):867 868 def __init__(self, config: KimiK25Config):869 super().__init__(config)870 871 vt_config = VisionTowerConfig(config.vision_config)872 self.vision_tower = MoonViT3dPretrainedModel(vt_config)873 874 proj_config = ProjectorConfig(config.vision_config)875 if proj_config.mm_projector_type == 'identity':876 self.mm_projector = IdentityMap()877 elif proj_config.mm_projector_type == 'mlp':878 self.mm_projector = MLP(proj_config)879 elif proj_config.mm_projector_type == 'patchmerger':880 self.mm_projector = PatchMergerMLP(proj_config)881 else:882 raise ValueError(883 f"Unsupported mm_projector_type: {proj_config.mm_projector_type}"884 )885 886 self.language_model = DeepseekV3ForCausalLM(config.text_config)887 self.post_init()888 889 if hasattr(self.language_model, 'dtype'):890 target_dtype = self.language_model.dtype891 self.vision_tower = self.vision_tower.to(dtype=target_dtype)892 self.mm_projector = self.mm_projector.to(dtype=target_dtype)893 894 def get_input_embeddings(self):895 return self.language_model.get_input_embeddings()896 897 def set_input_embeddings(self, value):898 self.language_model.set_input_embeddings(value)899 900 def get_output_embeddings(self):901 return self.language_model.get_output_embeddings()902 903 def set_output_embeddings(self, new_embeddings):904 self.language_model.set_output_embeddings(new_embeddings)905 906 def set_decoder(self, decoder):907 self.language_model.set_decoder(decoder)908 909 def get_decoder(self):910 return self.language_model.get_decoder()911 912 def tie_weights(self, *args, **kwargs):913 # Transformers >=5 passes ``missing_keys`` / ``recompute_mapping``; forward for the text backbone only.914 return self.language_model.tie_weights(*args, **kwargs)915 916 def resize_token_embeddings(self,917 new_num_tokens: int | None = None,918 pad_to_multiple_of=None) -> nn.Embedding:919 model_embeds = self.language_model.resize_token_embeddings(920 new_num_tokens, pad_to_multiple_of)921 # update vocab size922 self.config.text_config.vocab_size = model_embeds.num_embeddings923 self.vocab_size = model_embeds.num_embeddings924 return model_embeds925 926 def _merge_input_ids_with_image_features(927 self,928 image_features: list[torch.Tensor],929 inputs_embeds: torch.Tensor,930 input_ids: torch.Tensor,931 attention_mask: torch.Tensor,932 labels: torch.Tensor | None = None,933 ):934 """935 Args:936 image_features (:obj:`torch.Tensor` of shape :obj:`(num_image_tokens, embed_dim)`):937 The image features to merge with the input embeddings.938 inputs_embeds (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length, embed_dim)`):939 The input embeddings.940 input_ids (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`):941 The input ids.942 attention_mask (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`):943 The attention mask.944 labels (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`, *optional*):945 The labels.946 """947 _, embed_dim = image_features[0].shape948 feature_lengths = [x.shape[0] for x in image_features]949 image_features = torch.cat(image_features, dim=0)950 951 image_token_index: int = self.config.media_placeholder_token_id952 pad_token_id: int = self.config.pad_token_id953 ignore_index: int = self.config.ignore_index954 955 batch_size, sequence_length = input_ids.shape956 left_padding = not torch.sum(957 input_ids[:, -1] == torch.tensor(pad_token_id))958 959 # 1. Create a mask to know where special image tokens are960 _token_occupation_table = torch.ones_like(input_ids.flatten())961 _token_occupation_table[input_ids.flatten() ==962 image_token_index] = torch.tensor(963 feature_lengths,964 dtype=torch.long,965 device=input_ids.device)966 _token_occupation_table = _token_occupation_table.reshape(967 input_ids.shape)968 969 max_embed_dim = _token_occupation_table.sum(-1).max().item()970 assert (971 max_embed_dim >= sequence_length972 ), f"The maximum embedding dimension ({max_embed_dim}) is less than the sequence length ({sequence_length})"973 batch_indices, non_image_indices = torch.where(974 input_ids != image_token_index)975 976 # 2. Compute the positions where text should be written977 # Calculate new positions for text tokens in merged image-text sequence.978 new_token_positions = torch.cumsum(_token_occupation_table, -1) - 1979 nb_image_pad = max_embed_dim - 1 - new_token_positions[:, -1]980 if left_padding:981 new_token_positions += nb_image_pad[:,982 None] # offset for left padding983 text_to_overwrite = new_token_positions[batch_indices,984 non_image_indices]985 986 # 3. Create the full embedding, already padded to the maximum position987 final_embedding = torch.zeros(988 batch_size,989 max_embed_dim,990 embed_dim,991 dtype=inputs_embeds.dtype,992 device=inputs_embeds.device,993 )994 final_attention_mask = torch.zeros(batch_size,995 max_embed_dim,996 dtype=attention_mask.dtype,997 device=inputs_embeds.device)998 if labels is not None:999 final_labels = torch.full(1000 (batch_size, max_embed_dim),1001 ignore_index,1002 dtype=input_ids.dtype,1003 device=input_ids.device,1004 )1005 # In case the Vision model or the Language model has been offloaded to CPU, we need to manually1006 # set the corresponding tensors into their correct target device.1007 target_device = inputs_embeds.device1008 batch_indices, non_image_indices, text_to_overwrite = (1009 batch_indices.to(target_device),1010 non_image_indices.to(target_device),1011 text_to_overwrite.to(target_device),1012 )1013 attention_mask = attention_mask.to(target_device)1014 1015 # 4. Fill the embeddings based on the mask.1016 final_embedding[batch_indices,1017 text_to_overwrite] = inputs_embeds[batch_indices,1018 non_image_indices]1019 final_attention_mask[batch_indices,1020 text_to_overwrite] = attention_mask[1021 batch_indices, non_image_indices]1022 if labels is not None:1023 final_labels[batch_indices,1024 text_to_overwrite] = labels[batch_indices,1025 non_image_indices]1026 1027 # 5. Fill the embeddings corresponding to the images. Anything that is not `text_positions` needs filling (#29835)1028 image_to_overwrite = torch.full((batch_size, max_embed_dim),1029 True,1030 dtype=torch.bool,1031 device=inputs_embeds.device)1032 image_to_overwrite[batch_indices, text_to_overwrite] = False1033 image_to_overwrite &= image_to_overwrite.cumsum(1034 -1) - 1 >= nb_image_pad[:, None].to(target_device)1035 1036 if image_to_overwrite.sum() != image_features.shape[:-1].numel():1037 raise ValueError(1038 f"The input provided to the model are wrong. The number of image tokens is {image_to_overwrite.sum()} while"1039 f" the number of image features given to the model is {image_features.shape[:-1].numel()}. "1040 "This prevents correct indexing and breaks batch generation.")1041 1042 final_embedding[image_to_overwrite] = (1043 image_features.contiguous().reshape(-1,1044 embed_dim).to(target_device))1045 final_attention_mask |= image_to_overwrite1046 position_ids = (final_attention_mask.cumsum(-1) - 1).masked_fill_(1047 (final_attention_mask == 0), 1)1048 1049 # 6. Mask out the embedding at padding positions, as we later use the past_key_value value to determine the non-attended tokens.1050 batch_indices, pad_indices = torch.where(input_ids == pad_token_id)1051 indices_to_mask = new_token_positions[batch_indices, pad_indices]1052 1053 final_embedding[batch_indices, indices_to_mask] = 01054 1055 if labels is None:1056 final_labels = None1057 1058 return final_embedding, final_attention_mask, final_labels, position_ids1059 1060 def _extract_image_features(self, pixel_values: torch.Tensor,1061 grid_thws: torch.Tensor) -> list[torch.Tensor]:1062 """1063 Args:1064 pixel_values (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, num_channels, height, width)`):1065 The pixel values of the images processed by image processor.1066 grid_thws (:obj:`torch.Tensor` of shape :obj:`(batch_size, 3)`):1067 The grid, height, width of the images.1068 1069 Returns:1070 selected_image_feature (:obj:`torch.FloatTensor` of shape :obj:`(num_image_tokens, embed_dim)`):1071 The selected image features to use as input to the projector head.1072 1073 """1074 1075 target_dtype = self.vision_tower.patch_embed.proj.weight.dtype1076 pixel_values = pixel_values.to(target_dtype)1077 1078 image_features = self.vision_tower(pixel_values, grid_thws)1079 return image_features1080 1081 def forward(1082 self,1083 input_ids: torch.LongTensor | None = None,1084 pixel_values: torch.FloatTensor | list[torch.FloatTensor]1085 | None = None,1086 grid_thws: torch.Tensor | None = None,1087 attention_mask: torch.Tensor | None = None,1088 position_ids: torch.LongTensor | None = None,1089 past_key_values: list[torch.FloatTensor] | None = None,1090 inputs_embeds: torch.FloatTensor | None = None,1091 labels: torch.LongTensor | None = None,1092 use_cache: bool | None = None,1093 output_attentions: bool | None = None,1094 output_hidden_states: bool | None = None,1095 return_dict: bool | None = None,1096 ) -> tuple | LlavaCausalLMOutputWithPast:1097 r"""1098 Args:1099 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):1100 Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,1101 config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored1102 (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.1103 1104 ```"""1105 assert self.vision_tower is not None, "vision_tower is not loaded"1106 output_attentions = (output_attentions if output_attentions is not None1107 else self.config.output_attentions)1108 output_hidden_states = (output_hidden_states1109 if output_hidden_states is not None else1110 self.config.output_hidden_states)1111 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1112 1113 if inputs_embeds is None:1114 # 1. Extra the input embeddings1115 inputs_embeds = self.get_input_embeddings()(input_ids)1116 1117 # 2. Merge text and images1118 if pixel_values is not None and len(1119 pixel_values) > 0 and input_ids.shape[1] != 1:1120 image_features = self._extract_image_features(1121 pixel_values, grid_thws)1122 if self.mm_projector:1123 image_features = self.mm_projector(image_features)1124 1125 inputs_embeds = inputs_embeds.to(1126 image_features[0].dtype) # num_tokens, embed_dim1127 inputs_embeds, attention_mask, labels, position_ids = (1128 self._merge_input_ids_with_image_features(1129 image_features,1130 inputs_embeds,1131 input_ids,1132 attention_mask,1133 labels,1134 ))1135 1136 # In case input_ids.shape[1] == 1 & pixel_values==None & past_key_values != None, we are in the case of1137 # generation with cache1138 elif (past_key_values is not None and pixel_values is not None1139 and input_ids.shape[1] == 1):1140 first_layer_past_key_value = _first_layer_key_first_token_vector(1141 past_key_values)1142 if first_layer_past_key_value is not None:1143 # Sum all dimensions of head_dim (-2) to avoid random errors such as: https://github.com/huggingface/transformers/pull/28032#issuecomment-18636919411144 batch_index, non_attended_tokens = torch.where(1145 first_layer_past_key_value.float().sum(-2) == 0)1146 1147 # Get the target length1148 target_length = input_ids.shape[1]1149 past_length = _first_layer_past_seq_length(past_key_values)1150 if past_length is None:1151 past_length = int(first_layer_past_key_value.shape[-1])1152 1153 extended_attention_mask = torch.ones(1154 (attention_mask.shape[0], past_length),1155 dtype=attention_mask.dtype,1156 device=attention_mask.device,1157 )1158 1159 # Filter out only the tokens that can be un-attended, this can happen1160 # if one uses Llava + Fused modules where the cache on the1161 # first iteration is already big enough, or if one passes custom cache1162 valid_indices = non_attended_tokens < extended_attention_mask.size(1163 -1)1164 new_batch_index = batch_index[valid_indices]1165 new_non_attended_tokens = non_attended_tokens[1166 valid_indices]1167 1168 # Zero-out the places where we don't need to attend1169 extended_attention_mask[new_batch_index,1170 new_non_attended_tokens] = 01171 1172 attention_mask = torch.cat(1173 (extended_attention_mask,1174 attention_mask[:, -target_length:]),1175 dim=1)1176 position_ids = torch.sum(attention_mask,1177 dim=1).unsqueeze(-1) - 11178 1179 outputs = self.language_model(1180 attention_mask=attention_mask,1181 position_ids=position_ids,1182 past_key_values=past_key_values,1183 inputs_embeds=inputs_embeds,1184 use_cache=use_cache,1185 output_attentions=output_attentions,1186 output_hidden_states=output_hidden_states,1187 return_dict=return_dict,1188 )1189 1190 logits = outputs[0]1191 1192 loss = None1193 if labels is not None:1194 # Shift so that tokens < n predict n1195 if attention_mask is not None:1196 shift_attention_mask = attention_mask[..., 1:]1197 shift_logits = logits[..., :-1, :][shift_attention_mask.to(1198 logits.device) != 0].contiguous()1199 shift_labels = labels[..., 1:][shift_attention_mask.to(1200 labels.device) != 0].contiguous()