From 76e90cc9ff3eb3bf54425d463c9af341d968d5d2 Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Fri, 19 Jun 2026 16:13:07 +0530 Subject: [PATCH 01/33] Enable Kimi K25 Vision model with ConditionalGeneration class For DualQPC and upgrade TF Signed-off-by: Mamta Singh --- QEfficient/generation/vlm_generation.py | 1 - .../models/gemma3/modeling_gemma3.py | 2 +- .../transformers/models/kimi_k25/__init__.py | 7 + .../models/kimi_k25/configuration_kimi_k25.py | 122 ++ .../models/kimi_k25/modeling_kimi_k25.py | 1542 +++++++++++++++++ .../transformers/models/pytorch_transforms.py | 20 + examples/kimi_k2/export_kimi_k25_vision.py | 423 +++++ pyproject.toml | 2 +- 8 files changed, 2116 insertions(+), 3 deletions(-) create mode 100644 QEfficient/transformers/models/kimi_k25/__init__.py create mode 100644 QEfficient/transformers/models/kimi_k25/configuration_kimi_k25.py create mode 100644 QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py create mode 100644 examples/kimi_k2/export_kimi_k25_vision.py diff --git a/QEfficient/generation/vlm_generation.py b/QEfficient/generation/vlm_generation.py index 127ef432fd..1e87408234 100755 --- a/QEfficient/generation/vlm_generation.py +++ b/QEfficient/generation/vlm_generation.py @@ -736,7 +736,6 @@ def _generate_multi_frame_specialization( generation_len: int = None, stream: List[str] = None, ): - exec_batch_size = self.batch_size max_gen_length = self._ctx_len if not generation_len else max(self._ctx_len, generation_len) self.initialize_decode_inputs( diff --git a/QEfficient/transformers/models/gemma3/modeling_gemma3.py b/QEfficient/transformers/models/gemma3/modeling_gemma3.py index 96bf6cd1a0..cbb3902ec9 100644 --- a/QEfficient/transformers/models/gemma3/modeling_gemma3.py +++ b/QEfficient/transformers/models/gemma3/modeling_gemma3.py @@ -21,7 +21,6 @@ Gemma3ForConditionalGeneration, Gemma3TextConfig, Gemma3TextModel, - logger, repeat_kv, rotate_half, ) @@ -32,6 +31,7 @@ from QEfficient.utils import constants from QEfficient.utils._utils import IOInfo from QEfficient.utils.constants import MIN_MASKED_ATTENTION_VALUE +from QEfficient.utils.logging_utils import logger class GemmaRMSNormFunc(torch.autograd.Function): diff --git a/QEfficient/transformers/models/kimi_k25/__init__.py b/QEfficient/transformers/models/kimi_k25/__init__.py new file mode 100644 index 0000000000..da26921c50 --- /dev/null +++ b/QEfficient/transformers/models/kimi_k25/__init__.py @@ -0,0 +1,7 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) 2024 Qualcomm Innovation Center, Inc. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + diff --git a/QEfficient/transformers/models/kimi_k25/configuration_kimi_k25.py b/QEfficient/transformers/models/kimi_k25/configuration_kimi_k25.py new file mode 100644 index 0000000000..74c9a30a8b --- /dev/null +++ b/QEfficient/transformers/models/kimi_k25/configuration_kimi_k25.py @@ -0,0 +1,122 @@ +from transformers.configuration_utils import PretrainedConfig + +try: + from QEfficient.transformers.models.deepseek_v3.configuration_deepseek import DeepseekV3Config +except ImportError: + from QEfficient.transformers.models.deepseek_v3.configuration_deepseek import DeepseekV3Config + + +class KimiK25VisionConfig(PretrainedConfig): + def __init__( + self, + patch_size: int = 14, + init_pos_emb_height: int = 64, + init_pos_emb_width: int = 64, + init_pos_emb_time: int = 4, + pos_emb_type: str = "divided_fixed", + vt_num_attention_heads: int = 16, + vt_num_hidden_layers: int = 27, + vt_hidden_size: int = 1152, + vt_intermediate_size: int = 4304, + merge_kernel_size: tuple = (2, 2), + video_attn_type: str = "spatial_temporal", + merge_type: str = "sd2_tpool", + _attn_implementation: str = "flash_attention_2", + # MM Projector parameters + mm_projector_type: str = "patchmerger", + mm_hidden_size: int | None = None, + projector_hidden_act: str = "gelu", + projector_ln_eps: float = 1e-5, + # Other parameters + ignore_index: int = -100, + media_placeholder_token_id: int = 163605, + pad_token_id: int = 0, + use_unified_vision_chunk: bool = True, + video_placeholder="<|kimi_k25_video_placeholder|>", + text_hidden_size=7168, + **vision_config_kwargs, + ): + self.patch_size = patch_size + self.init_pos_emb_height = init_pos_emb_height + self.init_pos_emb_width = init_pos_emb_width + self.init_pos_emb_time = init_pos_emb_time + self.pos_emb_type = pos_emb_type + self.vt_num_attention_heads = vt_num_attention_heads + self.vt_num_hidden_layers = vt_num_hidden_layers + self.vt_hidden_size = vt_hidden_size + self.vt_intermediate_size = vt_intermediate_size + self.merge_kernel_size = merge_kernel_size + self.video_attn_type = video_attn_type + self.merge_type = merge_type + self._attn_implementation = _attn_implementation + + # MM Projector config + self.mm_projector_type = mm_projector_type + self.mm_hidden_size = mm_hidden_size if mm_hidden_size is not None else vt_hidden_size + self.projector_hidden_act = projector_hidden_act + self.projector_ln_eps = projector_ln_eps + self.text_hidden_size = text_hidden_size + + +class KimiK25Config(PretrainedConfig): + """Kimi-K2.5 model configuration. + + Args: + text_config (dict | DeepseekV3Config): Configuration for the text model. + + Vision Tower Parameters (from MoonViT3dConfig): + patch_size (int): Patch size for vision tower. + init_pos_emb_height (int): Initial position embedding height. + init_pos_emb_width (int): Initial position embedding width. + init_pos_emb_time (int): Initial position embedding time dimension. + pos_emb_type (str): Type of position embedding. + vt_num_attention_heads (int): Number of attention heads in vision tower. + vt_num_hidden_layers (int): Number of hidden layers in vision tower. + vt_hidden_size (int): Hidden size of vision tower. + vt_intermediate_size (int): Intermediate size in vision tower FFN. + merge_kernel_size (tuple): Kernel size for patch merging. + video_attn_type (str): Type of video attention. + merge_type (str): Type of merge operation. + _attn_implementation (str): Attention implementation type. + + MM Projector Parameters (from MultiModalProjectorConfig): + mm_projector_type (str): Type of multimodal projector. + mm_hidden_size (int): Hidden size from vision tower (should match vt_hidden_size). + projector_hidden_act (str): Activation function for projector. + projector_ln_eps (float): Layer norm epsilon for projector. + + Other Parameters: + ignore_index (int): The ignore index for the loss function. + media_placeholder_token_id (int): The token ID to use for media placeholders. + pad_token_id (int): The token ID to use for padding. + """ + + model_type = "kimi_k25" + + def __init__( + self, + text_config: dict | DeepseekV3Config = None, + vision_config: dict | KimiK25VisionConfig = None, + # Other parameters + ignore_index: int = -100, + media_placeholder_token_id: int = 163605, + pad_token_id: int = 0, + use_unified_vision_chunk: bool = True, + video_placeholder="<|kimi_k25_video_placeholder|>", + **kwargs, + ): + if isinstance(text_config, dict): + text_config = DeepseekV3Config(**text_config) + if isinstance(vision_config, dict): + vision_config = KimiK25VisionConfig(**vision_config) + self.text_config = text_config + self.vision_config = vision_config + # Other config + self.ignore_index = ignore_index + self.media_placeholder_token_id = media_placeholder_token_id + self.use_unified_vision_chunk = use_unified_vision_chunk + self.video_placeholder = video_placeholder + if getattr(self.text_config, "quantization_config", None) is not None: + self.quantization_config = self.text_config.quantization_config + + super().__init__(pad_token_id=pad_token_id, **kwargs) diff --git a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py new file mode 100644 index 0000000000..63dfdc69f6 --- /dev/null +++ b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py @@ -0,0 +1,1542 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ---------------------------------------------------------------------------- + +import math +import sys as _sys +from collections.abc import Sequence +from copy import deepcopy +from io import BytesIO +from pathlib import Path + +# from QEfficient import QEFFAutoModelForImageTextToText +from typing import List, Optional, Tuple, Type, Union + +import numpy as np +import requests +import torch +import torch.nn as nn +import torch.nn.functional as F +from PIL import Image +from transformers import AutoProcessor, activations + +from QEfficient.transformers.models.deepseek_v3.modeling_deepseek import QEffDeepseekV3ForCausalLM + +try: + from transformers.activations import PytorchGELUTanh +except ImportError: + from transformers.activations import GELUTanh + + activations.PytorchGELUTanh = GELUTanh + PytorchGELUTanh = GELUTanh +from transformers.activations import PytorchGELUTanh +from transformers.cache_utils import Cache +from transformers.configuration_utils import PretrainedConfig +from transformers.modeling_outputs import CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.models.llava.modeling_llava import LlavaCausalLMOutputWithPast + +from QEfficient.utils import constants + +from .configuration_kimi_k25 import KimiK25Config + + +def eager_attention_forward(q, k, v, **kwargs): + q = q.transpose(0, 1) # (num_heads, seq_len, head_dim) + k = k.transpose(0, 1) + v = v.transpose(0, 1) + attn_weight = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1]) + attn_weight = torch.softmax(attn_weight, dim=-1, dtype=torch.float32).to(q.dtype) + attn_out = attn_weight @ v + attn_out = attn_out.transpose(0, 1) # (seq_len, num_heads, head_dim) + return attn_out.reshape(attn_out.shape[0], -1) + + +VL_VISION_ATTENTION_FUNCTIONS = {"eager": eager_attention_forward} + + +def get_rope_shape_decorate(func): + _get_rope_shape_first_call_flag = set() + + def wrapper(org, interpolation_mode, shape): + key = (org.requires_grad, torch.is_grad_enabled(), interpolation_mode) + if key not in _get_rope_shape_first_call_flag: + _get_rope_shape_first_call_flag.add(key) + _ = func(org, interpolation_mode, shape=(64, 64)) + return func(org, interpolation_mode, shape) + + return wrapper + + +@get_rope_shape_decorate +@torch.compile(dynamic=True) # Remove this +def get_rope_shape(org, interpolation_mode, shape): + return ( + F.interpolate( + org.permute((2, 0, 1)).unsqueeze(0), + size=shape, + mode=interpolation_mode, + ) + .squeeze(0) + .permute((1, 2, 0)) + .flatten(end_dim=1) + ) + + +def apply_rope(xq, xk, freqs_cis): + # freqs_cis shape: (2, seq_len, dim//2) + # xq shape: (seq_len, num_heads, head_dim) + # Need freqs to broadcast: (seq_len, 1, dim//2) + freqs_cos = freqs_cis[0].unsqueeze(-2) # (seq_len, 1, dim//2) + freqs_sin = freqs_cis[1].unsqueeze(-2) # (seq_len, 1, dim//2) + xq_r = xq.float().view(*xq.shape[:-1], -1, 2) + xq_r0, xq_r1 = xq_r[..., 0], xq_r[..., 1] + xq_out_r = xq_r0 * freqs_cos - xq_r1 * freqs_sin + xq_out_i = xq_r0 * freqs_sin + xq_r1 * freqs_cos + xq_out = torch.stack([xq_out_r, xq_out_i], dim=-1).flatten(-2) + xk_r = xk.float().view(*xk.shape[:-1], -1, 2) + xk_r0, xk_r1 = xk_r[..., 0], xk_r[..., 1] + xk_out_r = xk_r0 * freqs_cos - xk_r1 * freqs_sin + xk_out_i = xk_r0 * freqs_sin + xk_r1 * freqs_cos + xk_out = torch.stack([xk_out_r, xk_out_i], dim=-1).flatten(-2) + return xq_out.type_as(xq), xk_out.type_as(xk) + + +def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): + """ + From: + https://github.com/OpenGVLab/InternVideo/blob/421f6d2361fc8f61a3394244571f2601a4e99e29/InternVideo2/multi_modality/models/backbones/internvideo2/pos_embed.py#L86 + embed_dim: output dimension for each position + pos: a list of positions to be encoded: size (M,) + out: (M, D) + """ + assert embed_dim % 2 == 0 + omega = np.arange(embed_dim // 2, dtype=np.float32) + omega /= embed_dim / 2.0 + omega = 1.0 / 10000**omega # (D/2,) + + pos = pos.reshape(-1) # (M,) + out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product + + emb_sin = np.sin(out) # (M, D/2) + emb_cos = np.cos(out) # (M, D/2) + + emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D) + return emb + + +def get_1d_sincos_pos_embed(embed_dim, t_size, cls_token=False): + """ + t_size: int of the temporal size + return: + pos_embed: [t_size, embed_dim] or [1+t_size, embed_dim] (w/ or w/o cls_token) + """ + grid_t = np.arange(t_size, dtype=np.float32) + pos_embed = get_1d_sincos_pos_embed_from_grid(embed_dim, grid_t) + if cls_token: + pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0) + return pos_embed + + +class QEffLearnable2DInterpPosEmbDivided_fixed(nn.Module): + def __qeff_init__(self): + self.interpolation_mode = "bilinear" + + """def __qeff_init__(self, + height: int, + width: int, + num_frames: int, + dim: int, + interpolation_mode: str = 'bilinear') -> None: + super().__init__() + self.height = height + self.width = width + self.num_frames = num_frames + self.dim = dim + self.interpolation_mode = interpolation_mode + self.weight = nn.Parameter(torch.empty(height, width, dim)) + self.register_buffer('time_weight', + torch.from_numpy( + get_1d_sincos_pos_embed( + self.dim, + self.num_frames)).float().unsqueeze(1), + persistent=False) + + self.reset_parameters() + """ + + def reset_parameters(self): + nn.init.normal_(self.weight) + + def forward(self, x: torch.Tensor, grid_thws: torch.Tensor) -> torch.Tensor: + pos_embs = [] + for t, h, w in grid_thws.tolist(): + assert t <= self.num_frames, f"t:{t} > self.num_frames:{self.num_frames}" + if (h, w) == self.weight.shape[:-1]: + pos_emb_2d = self.weight.flatten(end_dim=1) + else: + pos_emb_2d = get_rope_shape( + self.weight, + interpolation_mode=self.interpolation_mode, + shape=(h, w), + ) + + if t == 1: + pos_emb_3d = pos_emb_2d + else: + pos_emb_3d = pos_emb_2d.unsqueeze(0).repeat(t, 1, 1) + self.time_weight[0:t] + + pos_embs.append(pos_emb_3d.reshape(-1, pos_emb_3d.shape[-1])) + + out = x + torch.cat(pos_embs) + return out + + +class MoonVision3dPatchEmbed(nn.Module): + def __init__( + self, + out_dim: int, + in_dim: int = 3, + patch_size: int | tuple[int, int] = (14, 14), + pos_emb_height: int = 14, + pos_emb_width: int = 14, + pos_emb_time: int = 4, + pos_emb_type: str = "divided_fixed", + ): + super().__init__() + assert isinstance(patch_size, int | Sequence), f"Invalid patch_size type: {type(patch_size)}" + if isinstance(patch_size, int): + patch_size = (patch_size, patch_size) + assert len(patch_size) == 2, f"Expected patch_size to be a tuple of 2, got {patch_size}" + self.patch_size = patch_size + + self.proj = nn.Conv2d(in_dim, out_dim, kernel_size=patch_size, stride=patch_size) + + if pos_emb_type == "divided_fixed": + self.pos_emb = Learnable2DInterpPosEmbDivided_fixed( + height=pos_emb_height, width=pos_emb_width, num_frames=pos_emb_time, dim=out_dim + ) + else: + raise NotImplementedError(f"Not support pos_emb_type: {pos_emb_type}") + + def forward(self, x: torch.Tensor, grid_thws: torch.Tensor) -> torch.Tensor: + """ + Args: + x (L, Channels): input tensor + grid_hws (N, 3): temporal, height and width + Returns: + (L, Cout) tensor + """ + x = self.proj(x).view(x.size(0), -1) + # apply positional embedding + x = self.pos_emb(x, grid_thws) + return x + + +class Rope2DPosEmbRepeated(nn.Module): + """2D rotary position embedding with multi-resolution support. + This class is intended to be used in the following way: + 1. Before training, create an instance of Rope2DPosEmb. This instance will hold the precomputed cis. + 2. Before each forward pass, call `get_freqs_cis_by_*` to get the `freqs_cis` tensor for this iteration. + 3. During the forward pass, pass the `freqs_cis` tensor to each attention layer, and call `apply` just before each attention operation. + The rope is shared across all attention layers and all heads. + Refs: + - RoFormer: https://arxiv.org/abs/2104.09864 + - VisionLLaMA: https://arxiv.org/abs/2403.00522 + - https://github.com/Meituan-AutoML/VisionLLaMA/blob/main/dit/models.py + Args: + dim (int): usually the multi-head attention dimension, should be divisible by 4 (TODO: relax this constraint if needed) + max_height (int): the maximum height of the 2D grid + max_width (int): the maximum width of the 2D grid + theta_base (float): the base of the theta + device (str): the device to store the precomputed cis + """ + + def __init__(self, dim: int, max_height: int, max_width: int, theta_base=10000): + super().__init__() + self.dim = dim + assert self.dim % 4 == 0, "dim must be divisible by 4" + self.max_height = max_height + self.max_width = max_width + self.theta_base = theta_base + + def extra_repr(self): + return f"dim={self.dim}, max_height={self.max_height}, max_width={self.max_width}, theta_base={self.theta_base}" + + def _precompute_freqs_cis(self, device: torch.device) -> torch.Tensor: + """Calculate the cis(freqs) for each position in the 2D grid. + Return: complex tensor of shape (max_height, max_width, dim//2) and value: + height axis: ret[h, w, 2*i] = cis(h * theta_base**(-4*i/dim)) + weight axis: ret[h, w, 2*i+1] = cis(w * theta_base**(-4*i/dim)) with (i in [0, dim//4)) + note: `cis` is a mathematical notation defined by cis x = cos x + i sin x, + """ + N = self.max_height * self.max_width + flat_pos = torch.arange(0, N).float().to(device) + x_pos = flat_pos % self.max_width + y_pos = flat_pos // self.max_width + dim_range = torch.arange(0, self.dim, 4)[: (self.dim // 4)].float().to(device) # C/4 + freqs = 1.0 / (self.theta_base ** (dim_range / self.dim)) + x_freqs = torch.outer(x_pos, freqs).float() # N, C/4 + y_freqs = torch.outer(y_pos, freqs).float() # N, C/4 + x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs) # N, C/4 + y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs) # N, C/4 + # N, C/4, 2 + freqs_cis = torch.cat([x_cis.unsqueeze(dim=-1), y_cis.unsqueeze(dim=-1)], dim=-1) + # max_height, max_width, C/2 + freqs_cis = freqs_cis.reshape(self.max_height, self.max_width, -1) + return freqs_cis + + def get_freqs_cis(self, grid_thws: torch.Tensor, device: torch.device) -> torch.Tensor: + """ + Args: + grid_thws (torch.Tensor): grid time, height and width + Returns: + freqs_cis: tensor of shape (sum(t * height * width), dim//2) + """ + if not hasattr(self, "freqs_cis"): + self.register_buffer("freqs_cis", self._precompute_freqs_cis(device), persistent=False) + + shapes = grid_thws.tolist() + assert all(1 <= h <= self.max_height and 1 <= w <= self.max_width for t, h, w in shapes), ( + shapes, + self.max_height, + self.max_width, + ) + freqs_cis = torch.cat( + [self.freqs_cis[:h, :w].reshape(-1, self.dim // 2).repeat(t, 1) for t, h, w in shapes], + dim=0, + ) + return freqs_cis + + +class MLP2(nn.Module): + """ + Args: + dims: [in_dim, hidden_dim, out_dim] + bias: whether to use bias in linear layer. + """ + + def __init__(self, dims: list[int], activation, bias=True): + super().__init__() + assert len(dims) == 3 + self.fc0 = nn.Linear(dims[0], dims[1], bias=bias) + self.fc1 = nn.Linear(dims[1], dims[2], bias=bias) + self.activation = activation + for m in [self.fc0, self.fc1]: + nn.init.trunc_normal_(m.weight, std=math.sqrt(2 / m.in_features)) + if m.bias is not None: + nn.init.zeros_(m.bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.fc0(x) + x = self.activation(x) + return self.fc1(x) + + +class MoonViTEncoderLayer(nn.Module): + def __init__( + self, + num_heads: int, + hidden_dim: int, + mlp_dim: int, + *, + attn_implementation: str = "flash_attention_2", + activation=F.gelu, + attn_bias: bool = False, + use_deterministic_attn: bool = False, + ): + super().__init__() + self.num_heads = num_heads + self.hidden_dim = hidden_dim + self.hidden_size_per_attention_head = self.hidden_dim // self.num_heads + self.attn_implementation = attn_implementation + self.use_deterministic_attn = use_deterministic_attn + + self.norm0 = nn.LayerNorm(hidden_dim) + self.norm1 = nn.LayerNorm(hidden_dim) + self.mlp = MLP2([hidden_dim, mlp_dim, hidden_dim], activation) + self.wqkv = nn.Linear(hidden_dim, hidden_dim * 3, bias=attn_bias) + self.wo = nn.Linear(hidden_dim, hidden_dim, bias=attn_bias) + + def attention_qkvpacked( + self, + x: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: torch.Tensor, + rope_freqs_cis: torch.Tensor | None = None, + ): + """ + Args: + x (torch.Tensor): (batch_size, seqlen, hidden_dim) + cu_seqlens (torch.Tensor): + """ + xqkv = self.wqkv(x) + + qkv_shape = xqkv.size()[:-1] + ( + 3, + self.num_heads, + self.hidden_size_per_attention_head, + ) + # xqkv: (batch_size, seqlen, 3, nheads, headdim) + xqkv = xqkv.view(*qkv_shape) + xq, xk, xv = torch.unbind(xqkv, dim=-3) + + xq, xk = apply_rope(xq, xk, rope_freqs_cis) + + attn_func = VL_VISION_ATTENTION_FUNCTIONS[self.attn_implementation] + attn_out = attn_func( + xq, + xk, + xv, + q_cu_seqlens=cu_seqlens, + k_cu_seqlens=cu_seqlens, + max_seqlen_k=max_seqlen, + max_seqlen_q=max_seqlen, + deterministic=self.use_deterministic_attn, + ) + + attn_out = self.wo(attn_out) + return attn_out + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: int, + rope_freqs_cis: torch.Tensor | None = None, + ): + residual = hidden_states + hidden_states = self.norm0(hidden_states) + + hidden_states = self.attention_qkvpacked(hidden_states, cu_seqlens, max_seqlen, rope_freqs_cis) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.norm1(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +class QEffMoonViT3dEncoder(nn.Module): + def __qeff_init__(self): + self.block_cfg = { + "num_heads": 64, # config.num_attention_heads, + "hidden_dim": 7168, # config.hidden_size, + "mlp_dim": 18432, # config.intermediate_size, + "activation": PytorchGELUTanh(), + "attn_bias": True, + "attn_implementation": "eager", # config._attn_implementation, + } + self.rope_2d = Rope2DPosEmbRepeated(self.block_cfg["hidden_dim"] // self.block_cfg["num_heads"], 512, 512) + # if not hasattr(self.rope_2d, 'freqs_cos'): + # self.rope_2d.register_buffer('freqs_cos', self.rope_2d.freqs_cis.real.contiguous(), persistent=False) + # self.rope_2d.register_buffer('freqs_sin', self.rope_2d.freqs_cis.imag.contiguous(), persistent=False) + self.blocks = nn.ModuleList( + [ + MoonViTEncoderLayer(**self.block_cfg, use_deterministic_attn=False) + for _ in range(2) # config.num_hidden_layers) + ] + ) + + def forward( + self, + hidden_states: torch.Tensor, + grid_thws: torch.Tensor, + ) -> torch.Tensor: + # if not hasattr(self.rope_2d, 'freqs_cos'): + # self.rope_2d.register_buffer('freqs_cos', self.rope_2d.freqs_cis.real.contiguous(), persistent=False) + # self.rope_2d.register_buffer('freqs_sin', self.rope_2d.freqs_cis.imag.contiguous(), persistent=False) + rope_freqs_cis = self.rope_2d.get_freqs_cis(grid_thws=grid_thws, device=hidden_states.device) + + lengths = torch.cat( + ( + torch.zeros(1, dtype=grid_thws.dtype, device=grid_thws.device), + grid_thws[:, 0] * grid_thws[:, 1] * grid_thws[:, 2], + ) + ) + + max_seqlen = lengths.max() + cu_seqlens = lengths.to(hidden_states.device).cumsum(dim=0, dtype=torch.int32) + for block in self.blocks: + hidden_states = block(hidden_states, cu_seqlens, max_seqlen, rope_freqs_cis=rope_freqs_cis) + + hidden_states = self.final_layernorm(hidden_states) + return hidden_states + + +def tpool_patch_merger( + x: torch.Tensor, + grid_thws: torch.Tensor, + merge_kernel_size: tuple[int, int] = (2, 2), +) -> list[torch.Tensor]: + d_model = x.size(-1) + + outputs = [] + pre_sum = 0 + for t, h, w in grid_thws.tolist(): + # Get the current sequence + seq = x[pre_sum : pre_sum + t * h * w] + # Reshape along self.merge_kernel_size and concat to the last dimension + kernel_height, kernel_width = merge_kernel_size + new_height, new_width = h // kernel_height, w // kernel_width + reshaped_seq = seq.view(t, new_height, kernel_height, new_width, kernel_width, d_model) + reshaped_seq = reshaped_seq.permute(0, 1, 3, 2, 4, 5).contiguous().mean(dim=0) # temporal pooling + padded_seq = reshaped_seq.view(new_height * new_width, kernel_height * kernel_width, -1) + outputs.append(padded_seq) + pre_sum += t * h * w + + return outputs + + +class MoonViT3dPretrainedModel(PreTrainedModel): + config_class = None + model_type = "moonvit3d" + _no_split_modules = ["PackingTransformer"] + _supports_flash_attn_2 = True + _supports_sdpa = True + + def __init__(self, config, *inputs, **kwargs): + super().__init__(config, *inputs, **kwargs) + config = deepcopy(config) + self.merge_kernel_size = config.merge_kernel_size + self.patch_size = config.patch_size + self.merge_type = config.merge_type + + self.patch_embed = MoonVision3dPatchEmbed( + out_dim=config.hidden_size, + patch_size=config.patch_size, + pos_emb_height=config.init_pos_emb_height, + pos_emb_width=config.init_pos_emb_width, + pos_emb_time=config.init_pos_emb_time, + pos_emb_type=config.pos_emb_type, + ) + + self.encoder = MoonViT3dEncoder( + hidden_dim=config.hidden_size, + num_layers=config.num_hidden_layers, + block_cfg={ + "num_heads": config.num_attention_heads, + "hidden_dim": config.hidden_size, + "mlp_dim": config.intermediate_size, + "activation": PytorchGELUTanh(), + "attn_bias": True, + "attn_implementation": config._attn_implementation, + }, + video_attn_type=config.video_attn_type, + ) + + def forward(self, pixel_values: torch.Tensor, grid_thws: torch.Tensor) -> torch.Tensor: + """ + Args: + pixel_values (torch.Tensor): The input pixel values. + grid_thws (torch.Tensor): Temporal, height and width. + Returns: + torch.Tensor: The output tokens. + """ + # grid_thws = grid_thws.to('cpu') + assert grid_thws.ndim == 2, f"grid_thws should be 2D, got {grid_thws.ndim}" + assert grid_thws.size(1) == 3, f"No support for thw: {grid_thws}" + hidden_states = self.patch_embed(pixel_values, grid_thws) + hidden_states = self.encoder(hidden_states, grid_thws) + if self.merge_type == "sd2_tpool": # spatial downsampling 2x with temporal pooling all + hidden_states = tpool_patch_merger(hidden_states, grid_thws, merge_kernel_size=self.merge_kernel_size) + else: + raise NotImplementedError(f"Not support {self.merge_type}") + + return hidden_states + + +# ============================================================================ +# MM Projector Helper Classes (from mm_projector/modeling_mm_projectors.py) +# ============================================================================ + + +class IdentityMap(nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x, *args, **kwargs): + return x + + +class MLP(nn.Module): + def __init__(self, config): + super().__init__() + # TODO, use faster LayerNorm + self.pre_norm = nn.LayerNorm(config.mm_hidden_size) + self.proj = nn.Sequential( + nn.Linear(config.mm_hidden_size, config.hidden_size), + nn.GELU(), + nn.Linear(config.hidden_size, config.hidden_size), + ) + + def forward(self, x, *args, **kwargs): + assert isinstance(x, list | tuple), f"x is not a list or tuple: {type(x)}" + lengths = [item.shape[0] for item in x] + x = torch.cat(x, dim=0) + x = self.pre_norm(x) + x = self.proj(x) + x = torch.split(x, lengths, dim=0) + + return x + + +class PatchMergerMLP(nn.Module): + def __init__(self, config): + super().__init__() + eps = config.projector_ln_eps + self.hidden_size = config.mm_hidden_size * (config.merge_kernel_size[0] * config.merge_kernel_size[1]) + self.pre_norm = nn.LayerNorm(config.mm_hidden_size, eps=eps) + self.proj = nn.Sequential( + nn.Linear(self.hidden_size, self.hidden_size), + nn.GELU(), + nn.Linear(self.hidden_size, config.hidden_size), + ) + + def forward(self, x, *args, **kwargs): + if isinstance(x, list) or isinstance(x, tuple): + x = [self.proj(self.pre_norm(item).view(item.shape[0], -1)) for item in x] + else: + # B, N, N_k, C = x.shape + B = x.shape[0] + x = self.proj(self.pre_norm(x).view(B, -1, self.hidden_size)) + return x + + +class KimiK25PreTrainedModel(PreTrainedModel): + config_class = KimiK25Config + base_model_prefix = "model" + _no_split_modules = [ + "MoonViT3dPretrainedModel", + "MoonViTEncoderLayer", + "DeepseekDecoderLayer", + "PatchMergerMLP", + ] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn_2 = True + _supports_sdpa = False + + def _init_weights(self, module): + # important: this ported version of Llava isn't meant for training from scratch - only + # inference and fine-tuning - so the proper init weights code has been removed - the original codebase + # https://github.com/haotian-liu/LLaVA/tree/main/llava should serve for that purpose + std = ( + self.config.initializer_range + if hasattr(self.config, "initializer_range") + else self.config.text_config.initializer_range + ) + + if hasattr(module, "class_embedding"): + module.class_embedding.data.normal_(mean=0.0, std=std) + + if isinstance(module, (nn.Linear, nn.Conv2d)): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +class VisionTowerConfig(PretrainedConfig): + model_type = "moonvit3d" + + def __init__(self, config: KimiK25Config, **kwargs): + super().__init__(**kwargs) + self.patch_size = config.patch_size + self.init_pos_emb_height = config.init_pos_emb_height + self.init_pos_emb_width = config.init_pos_emb_width + self.init_pos_emb_time = config.init_pos_emb_time + self.pos_emb_type = config.pos_emb_type + self.num_attention_heads = config.vt_num_attention_heads + self.num_hidden_layers = config.vt_num_hidden_layers + self.hidden_size = config.vt_hidden_size + self.intermediate_size = config.vt_intermediate_size + self.merge_kernel_size = config.merge_kernel_size + self.video_attn_type = config.video_attn_type + self.merge_type = config.merge_type + self._attn_implementation = config._attn_implementation + + +class ProjectorConfig: + def __init__(self, config: KimiK25Config): + self.mm_projector_type = config.mm_projector_type + self.mm_hidden_size = config.mm_hidden_size + self.hidden_size = config.text_hidden_size + self.merge_kernel_size = config.merge_kernel_size + self.projector_hidden_act = config.projector_hidden_act + self.projector_ln_eps = config.projector_ln_eps + + +class QEffKimiK25EncoderWrapper(nn.Module): + def __init__(self, model): + super().__init__() + self.model = model + self.config = self.model.config + # _orig_apply_rope = _kimi_module.apply_rope + # _kimi_module.apply_rope = _apply_rope_real + # _kimi_module = _sys.modules[type(model).__module__] + + # Restore original apply_rope and attention functions + # _kimi_module.apply_rope = _orig_apply_rope + # _kimi_module.VL_VISION_ATTENTION_FUNCTIONS.update(_orig_attn_functions) + + # _orig_attn_functions = _kimi_module.VL_VISION_ATTENTION_FUNCTIONS.copy() + # _kimi_module.VL_VISION_ATTENTION_FUNCTIONS["eager"] = _full_attention_forward + + def get_submodules_for_export(self) -> Type[nn.Module]: + """ + Return the set of class used as the repeated layer across the model for subfunction extraction. + Notes: + This method should return the *class object* (not an instance). + Downstream code can use this to find/build subfunctions for repeated blocks. + """ + return {self.model.vision_model.model.layers[0].__class__} + # return {self.model.layers[0].__class__} + + def forward_only_image(self, pixel_values: torch.Tensor, grid_thws: torch.Tensor) -> list[torch.Tensor]: + """ + Run only the vision tower and mm_projector to extract image embeddings. + + Args: + pixel_values: Preprocessed image pixel values. + grid_thws: Grid temporal/height/width info for the images. + + Returns: + image_embeds: List of projected image embedding tensors, one per image. + """ + image_features = self._extract_image_features(pixel_values, grid_thws) + if self.mm_projector: + image_features = self.mm_projector(image_features) + return image_features + + def forward_only_image_for_export( + self, pixel_values: torch.Tensor, h_shape: torch.Tensor, w_shape: torch.Tensor + ) -> torch.Tensor: + """ + ONNX-exportable forward that runs only the vision tower and mm_projector. + Uses h_shape and w_shape (int64 ones tensors of length h and w respectively) + to derive spatial dimensions via .shape[0], enabling dynamic axis export. + + Args: + pixel_values: Preprocessed image pixel values (num_patches, channels*patch_h*patch_w). + h_shape: int64 ones tensor of shape (h,) encoding number of patch rows. + w_shape: int64 ones tensor of shape (w,) encoding number of patch columns. + + Returns: + image_embeds: Projected image embeddings as a single tensor. + """ + _kimi_module = _sys.modules[type(self).__module__] + _get_rope_shape = _kimi_module.get_rope_shape + + h_shape = h_shape.to(pixel_values.device) + w_shape = w_shape.to(pixel_values.device) + + # Keep them in ONNX graph + dummy = (h_shape.float().sum() + w_shape.float().sum()) * 0.0 + pixel_values = pixel_values + dummy + + h = h_shape.shape[0] + w = w_shape.shape[0] + + target_dtype = self.vision_tower.patch_embed.proj.weight.dtype + pixel_values = pixel_values.to(target_dtype) + + # --- Patch embedding --- + x = self.vision_tower.patch_embed.proj(pixel_values).view(pixel_values.size(0), -1) + + # Positional embedding (single image, t=1) + pos_emb_module = self.vision_tower.patch_embed.pos_emb + pos_emb_2d = _get_rope_shape( + pos_emb_module.weight, + interpolation_mode=pos_emb_module.interpolation_mode, + shape=(h, w), + ) + x = x + pos_emb_2d + + # --- Encoder --- + # For single image with t=1: cu_seqlens = [0, h*w] + num_tokens = h * w + cu_seqlens = torch.zeros(2, dtype=torch.int32, device=x.device) + cu_seqlens[1] = num_tokens + max_seqlen = num_tokens + + # RoPE frequencies for single image (stacked real/imag to avoid complex tensors in ONNX) + # Shape: (2, num_tokens, dim//2) where [0]=cos, [1]=sin + rope_2d = self.vision_tower.encoder.rope_2d + freqs_cos = rope_2d.freqs_cos[:h, :w].reshape(-1, rope_2d.dim // 2) + freqs_sin = rope_2d.freqs_sin[:h, :w].reshape(-1, rope_2d.dim // 2) + freqs_cis = torch.stack([freqs_cos, freqs_sin], dim=0) + + for block in self.vision_tower.encoder.blocks: + x = block(x, cu_seqlens, max_seqlen, rope_freqs_cis=freqs_cis) + + x = self.vision_tower.encoder.final_layernorm(x) + + # --- tpool_patch_merger (single image, t=1) --- + merge_kernel_size = self.vision_tower.merge_kernel_size + kernel_height, kernel_width = merge_kernel_size + d_model = x.size(-1) + new_height = h // kernel_height + new_width = w // kernel_width + reshaped = x.view(1, new_height, kernel_height, new_width, kernel_width, d_model) + reshaped = reshaped.permute(0, 1, 3, 2, 4, 5).contiguous().mean(dim=0) + merged = reshaped.view(new_height * new_width, kernel_height * kernel_width, -1) + + # --- mm_projector (PatchMergerMLP on single tensor) --- + image_embeds = self.mm_projector.proj(self.mm_projector.pre_norm(merged).view(merged.shape[0], -1)) + return image_embeds + + +class QEffKimiK25DecoderWrapper(nn.Module): + def __init__(self, model): + super().__init__() + self.model = model + self.language_model = self.model.language_model + # self.language_model = QEffDeepseekV3ForCausalLM#(config.text_config) + self.config = self.model.config + + def get_submodules_for_export(self) -> Type[nn.Module]: + """ + Return the set of class used as the repeated layer across the model for subfunction extraction. + Notes: + This method should return the *class object* (not an instance). + Downstream code can use this to find/build subfunctions for repeated blocks. + """ + return {self.model.layers[0].__class__} + + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + compressed_kvs: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + batch_index: Optional[torch.LongTensor] = None, + image_embeds: Optional[List[torch.FloatTensor]] = None, + # inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + num_logits_to_keep: int = 0, + **kwargs, + ) -> Union[Tuple, CausalLMOutputWithPast]: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model.language_model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + compressed_kvs=compressed_kvs, + past_key_values=past_key_values, + batch_index=batch_index, + inputs_embeds=image_embeds, # inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + **kwargs, + ) + + hidden_states = outputs[0] + logit_index = position_ids.to(torch.int32).argmax(1, keepdim=True) + hidden_states = hidden_states[torch.arange(position_ids.shape[0]).view(-1, 1), logit_index] + logits = self.lm_head(hidden_states).float() + + loss = None + if labels is not None: + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + loss_fct = nn.CrossEntropyLoss() + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1).to(shift_logits.device) + loss = loss_fct(shift_logits, shift_labels) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +# ref https://github.com/huggingface/transformers/blob/78b2929c0554b79e0489b451ce4ece14d265ead2/src/transformers/models/llava/modeling_llava.py#L240 +class QEffKimiK25ForConditionalGeneration(KimiK25PreTrainedModel): + def __init__(self, config: KimiK25Config): + super().__init__(config) + + vt_config = VisionTowerConfig(config.vision_config) + self.vision_tower = MoonViT3dPretrainedModel(vt_config) + + proj_config = ProjectorConfig(config.vision_config) + if proj_config.mm_projector_type == "identity": + self.mm_projector = IdentityMap() + elif proj_config.mm_projector_type == "mlp": + self.mm_projector = MLP(proj_config) + elif proj_config.mm_projector_type == "patchmerger": + self.mm_projector = PatchMergerMLP(proj_config) + else: + raise ValueError(f"Unsupported mm_projector_type: {proj_config.mm_projector_type}") + + self.language_model = QEffDeepseekV3ForCausalLM(config.text_config) + self.post_init() + + if hasattr(self.language_model, "dtype"): + target_dtype = self.language_model.dtype + self.vision_tower = self.vision_tower.to(dtype=target_dtype) + self.mm_projector = self.mm_projector.to(dtype=target_dtype) + + def get_qeff_vision_encoder(self): + return QEffKimiK25EncoderWrapper(self) + + def get_qeff_language_decoder(self): + return QEffKimiK25DecoderWrapper(self) + + def get_input_embeddings(self): + return self.language_model.get_input_embeddings() + + def set_input_embeddings(self, value): + self.language_model.set_input_embeddings(value) + + def get_output_embeddings(self): + return self.language_model.get_output_embeddings() + + def set_output_embeddings(self, new_embeddings): + self.language_model.set_output_embeddings(new_embeddings) + + def set_decoder(self, decoder): + self.language_model.set_decoder(decoder) + + def get_decoder(self): + return self.language_model.get_decoder() + + def tie_weights(self): + return self.language_model.tie_weights() + + def resize_token_embeddings(self, new_num_tokens: int | None = None, pad_to_multiple_of=None) -> nn.Embedding: + model_embeds = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of) + # update vocab size + self.config.text_config.vocab_size = model_embeds.num_embeddings + self.vocab_size = model_embeds.num_embeddings + return model_embeds + + def _merge_input_ids_with_image_features( + self, + image_features: list[torch.Tensor], + inputs_embeds: torch.Tensor, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + labels: torch.Tensor | None = None, + ): + """ + Args: + image_features (:obj:`torch.Tensor` of shape :obj:`(num_image_tokens, embed_dim)`): + The image features to merge with the input embeddings. + inputs_embeds (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length, embed_dim)`): + The input embeddings. + input_ids (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`): + The input ids. + attention_mask (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`): + The attention mask. + labels (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`, *optional*): + The labels. + """ + _, embed_dim = image_features[0].shape + feature_lengths = [x.shape[0] for x in image_features] + image_features = torch.cat(image_features, dim=0) + + image_token_index: int = self.config.media_placeholder_token_id + pad_token_id: int = self.config.pad_token_id + ignore_index: int = self.config.ignore_index + + batch_size, sequence_length = input_ids.shape + left_padding = not torch.sum(input_ids[:, -1] == torch.tensor(pad_token_id)) + + # 1. Create a mask to know where special image tokens are + _token_occupation_table = torch.ones_like(input_ids.flatten()) + _token_occupation_table[input_ids.flatten() == image_token_index] = torch.tensor( + feature_lengths, dtype=torch.long, device=input_ids.device + ) + _token_occupation_table = _token_occupation_table.reshape(input_ids.shape) + + max_embed_dim = _token_occupation_table.sum(-1).max().item() + assert max_embed_dim >= sequence_length, ( + f"The maximum embedding dimension ({max_embed_dim}) is less than the sequence length ({sequence_length})" + ) + batch_indices, non_image_indices = torch.where(input_ids != image_token_index) + + # 2. Compute the positions where text should be written + # Calculate new positions for text tokens in merged image-text sequence. + new_token_positions = torch.cumsum(_token_occupation_table, -1) - 1 + nb_image_pad = max_embed_dim - 1 - new_token_positions[:, -1] + if left_padding: + new_token_positions += nb_image_pad[:, None] # offset for left padding + text_to_overwrite = new_token_positions[batch_indices, non_image_indices] + + # 3. Create the full embedding, already padded to the maximum position + final_embedding = torch.zeros( + batch_size, + max_embed_dim, + embed_dim, + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) + final_attention_mask = torch.zeros( + batch_size, max_embed_dim, dtype=attention_mask.dtype, device=inputs_embeds.device + ) + if labels is not None: + final_labels = torch.full( + (batch_size, max_embed_dim), + ignore_index, + dtype=input_ids.dtype, + device=input_ids.device, + ) + # In case the Vision model or the Language model has been offloaded to CPU, we need to manually + # set the corresponding tensors into their correct target device. + target_device = inputs_embeds.device + batch_indices, non_image_indices, text_to_overwrite = ( + batch_indices.to(target_device), + non_image_indices.to(target_device), + text_to_overwrite.to(target_device), + ) + attention_mask = attention_mask.to(target_device) + + # 4. Fill the embeddings based on the mask. + final_embedding[batch_indices, text_to_overwrite] = inputs_embeds[batch_indices, non_image_indices] + final_attention_mask[batch_indices, text_to_overwrite] = attention_mask[batch_indices, non_image_indices] + if labels is not None: + final_labels[batch_indices, text_to_overwrite] = labels[batch_indices, non_image_indices] + + # 5. Fill the embeddings corresponding to the images. Anything that is not `text_positions` needs filling (#29835) + image_to_overwrite = torch.full( + (batch_size, max_embed_dim), True, dtype=torch.bool, device=inputs_embeds.device + ) + image_to_overwrite[batch_indices, text_to_overwrite] = False + image_to_overwrite &= image_to_overwrite.cumsum(-1) - 1 >= nb_image_pad[:, None].to(target_device) + + if image_to_overwrite.sum() != image_features.shape[:-1].numel(): + raise ValueError( + f"The input provided to the model are wrong. The number of image tokens is {image_to_overwrite.sum()} while" + f" the number of image features given to the model is {image_features.shape[:-1].numel()}. " + "This prevents correct indexing and breaks batch generation." + ) + + final_embedding[image_to_overwrite] = image_features.contiguous().reshape(-1, embed_dim).to(target_device) + final_attention_mask |= image_to_overwrite + position_ids = (final_attention_mask.cumsum(-1) - 1).masked_fill_((final_attention_mask == 0), 1) + + # 6. Mask out the embedding at padding positions, as we later use the past_key_value value to determine the non-attended tokens. + batch_indices, pad_indices = torch.where(input_ids == pad_token_id) + indices_to_mask = new_token_positions[batch_indices, pad_indices] + + final_embedding[batch_indices, indices_to_mask] = 0 + + if labels is None: + final_labels = None + + return final_embedding, final_attention_mask, final_labels, position_ids + + def _extract_image_features(self, pixel_values: torch.Tensor, grid_thws: torch.Tensor) -> list[torch.Tensor]: + """ + Args: + pixel_values (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, num_channels, height, width)`): + The pixel values of the images processed by image processor. + grid_thws (:obj:`torch.Tensor` of shape :obj:`(batch_size, 3)`): + The grid, height, width of the images. + Returns: + selected_image_feature (:obj:`torch.FloatTensor` of shape :obj:`(num_image_tokens, embed_dim)`): + The selected image features to use as input to the projector head. + """ + + target_dtype = self.vision_tower.patch_embed.proj.weight.dtype + pixel_values = pixel_values.to(target_dtype) + + image_features = self.vision_tower(pixel_values, grid_thws) + return image_features + + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | list[torch.FloatTensor] | None = None, + grid_thws: torch.Tensor | None = None, + # h_shape: torch.Tensor | None = None, + # w_shape: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + compressed_kvs: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + ) -> tuple | LlavaCausalLMOutputWithPast: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + ```""" + assert self.vision_tower is not None, "vision_tower is not loaded" + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if inputs_embeds is None: + # 1. Extra the input embeddings + inputs_embeds = self.get_input_embeddings()(input_ids) + + # 2. Merge text and images + if pixel_values is not None and len(pixel_values) > 0 and input_ids.shape[1] != 1: + image_features = self._extract_image_features(pixel_values, grid_thws) + if self.mm_projector: + image_features = self.mm_projector(image_features) + + inputs_embeds = inputs_embeds.to(image_features[0].dtype) # num_tokens, embed_dim + inputs_embeds, attention_mask, labels, position_ids = self._merge_input_ids_with_image_features( + image_features, + inputs_embeds, + input_ids, + attention_mask, + labels, + ) + + # In case input_ids.shape[1] == 1 & pixel_values==None & past_key_values != None, we are in the case of + # generation with cache + elif past_key_values is not None and pixel_values is not None and input_ids.shape[1] == 1: + # Retrieve the first layer to inspect the logits and mask out the hidden states + # that are set to 0 + first_layer_past_key_value = past_key_values[0][0][:, :, :, 0] + + # Sum all dimensions of head_dim (-2) to avoid random errors such as: https://github.com/huggingface/transformers/pull/28032#issuecomment-1863691941 + batch_index, non_attended_tokens = torch.where(first_layer_past_key_value.float().sum(-2) == 0) + + # Get the target length + target_length = input_ids.shape[1] + past_length = first_layer_past_key_value.shape[-1] + + extended_attention_mask = torch.ones( + (attention_mask.shape[0], past_length), + dtype=attention_mask.dtype, + device=attention_mask.device, + ) + + # Filter out only the tokens that can be un-attended, this can happen + # if one uses Llava + Fused modules where the cache on the + # first iteration is already big enough, or if one passes custom cache + valid_indices = non_attended_tokens < extended_attention_mask.size(-1) + new_batch_index = batch_index[valid_indices] + new_non_attended_tokens = non_attended_tokens[valid_indices] + + # Zero-out the places where we don't need to attend + extended_attention_mask[new_batch_index, new_non_attended_tokens] = 0 + + attention_mask = torch.cat((extended_attention_mask, attention_mask[:, -target_length:]), dim=1) + position_ids = torch.sum(attention_mask, dim=1).unsqueeze(-1) - 1 + + outputs = self.language_model( + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + compressed_kvs=compressed_kvs, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + logits = outputs[0] + + loss = None + if labels is not None: + # Shift so that tokens < n predict n + if attention_mask is not None: + shift_attention_mask = attention_mask[..., 1:] + shift_logits = logits[..., :-1, :][shift_attention_mask.to(logits.device) != 0].contiguous() + shift_labels = labels[..., 1:][shift_attention_mask.to(labels.device) != 0].contiguous() + else: + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = nn.CrossEntropyLoss() + loss = loss_fct( + shift_logits.view(-1, shift_logits.size(-1)), + shift_labels.view(-1).to(shift_logits.device), + ) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return LlavaCausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + inputs_embeds=None, + pixel_values=None, + grid_thws=None, + attention_mask=None, + **kwargs, + ): + if past_key_values is not None: + if isinstance(past_key_values, Cache): + cache_length = past_key_values.get_seq_length() + past_length = getattr(past_key_values, "seen_tokens", cache_length) + else: + cache_length = past_length = past_key_values[0][0].shape[2] + + # Keep only the unprocessed tokens: + # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where + # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as + # input) + if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]: + input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :] + # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard + # input_ids based on the past_length. + elif past_length < input_ids.shape[1]: + input_ids = input_ids[:, past_length:] + # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens. + elif self.config.media_placeholder_token_id in input_ids: + input_ids = input_ids[:, input_ids.shape[1] - 1 :] + # If the cache has seen more tokens than it can hold, then the cache has a size limit. Let's discard the + # older attention values, as their corresponding values are not part of the input. + if cache_length < past_length and attention_mask is not None: + attention_mask = attention_mask[:, -(cache_length + input_ids.shape[1]) :] + + position_ids = kwargs.get("position_ids", None) + if attention_mask is not None and position_ids is None: + # create position_ids on the fly for batch generation + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + if past_key_values: + position_ids = position_ids[:, -input_ids.shape[1] :] + + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step + if inputs_embeds is not None and past_key_values is None: + model_inputs = {"inputs_embeds": inputs_embeds} + else: + model_inputs = {"input_ids": input_ids} + + model_inputs.update( + { + "position_ids": position_ids, + "past_key_values": past_key_values, + "use_cache": kwargs.get("use_cache"), + "attention_mask": attention_mask, + "pixel_values": pixel_values, + "grid_thws": grid_thws, + } + ) + return model_inputs + + def _reorder_cache(self, *args, **kwargs): + return self.language_model._reorder_cache(*args, **kwargs) + + def get_output_names(self, kv_offload: bool = False): + vision_output_names = ["image_embeds"] + lang_output_names = ["logits"] + + mla_absorption = getattr(self.language_model, "mla_absorption", None) + if mla_absorption is not None: + cache_compressed = mla_absorption.get("cache_compressed", False) + else: + cache_compressed = False + + if cache_compressed: + for i in range(self.language_model.config.num_hidden_layers): + lang_output_names.append(f"compressed_kv.{i}_RetainedState") + lang_output_names.append(f"k_pe.{i}_RetainedState") + else: + for i in range(self.language_model.config.num_hidden_layers): + for kv in ["key", "value"]: + lang_output_names.append(f"past_{kv}.{i}_RetainedState") + + output_names = {} + if kv_offload: + output_names["vision"] = vision_output_names + output_names["lang"] = lang_output_names + else: + return lang_output_names + return output_names + + def get_dummy_inputs( + self, + comp_ctx_lengths: Optional[List[int]] = None, + kv_offload: bool = False, + continuous_batching: bool = False, + **kwargs, + ): + prefill_seq_len = kwargs.get("prefill_seq_len") + if prefill_seq_len is None: + prefill_seq_len = constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN + prefill_seq_len = int(prefill_seq_len) + + bs: int = constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE + fbs: int = constants.ONNX_EXPORT_EXAMPLE_FBS + + model_path = Path( + "/home/huggingface_hub/models--moonshotai--Kimi-K2.5/snapshots/4d01dfe0332d63057c186e0b262165819efb6611" + ) + processor = AutoProcessor.from_pretrained(str(model_path), trust_remote_code=True) + image_url = "https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png" + image = Image.open(BytesIO(requests.get(image_url).content)).convert("RGB") + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": image}, + {"type": "text", "text": "Tell me about yourself."}, + ], + } + ] + inputs = processor( + messages=messages, + add_generation_prompt=True, + tokenize=False, + return_tensors="pt", + ) + inputs = {k: v.to(self.language_model.device) if hasattr(v, "to") else v for k, v in inputs.items()} + + # Build h_shape and w_shape from grid_thws (single image, t=1) + grid_thws_val = inputs["grid_thws"] + h_val = int(grid_thws_val[0, 1].item()) + w_val = int(grid_thws_val[0, 2].item()) + h_shape_tensor = torch.ones(h_val, dtype=torch.int64, device=self.language_model.device) + w_shape_tensor = torch.ones(w_val, dtype=torch.int64, device=self.language_model.device) + + vision_inputs = { + "pixel_values": inputs["pixel_values"], + "h_shape": h_shape_tensor, + "w_shape": w_shape_tensor, + } + + lang_inputs = { + "input_ids": torch.zeros((bs, prefill_seq_len), dtype=torch.int64), + "position_ids": torch.arange(prefill_seq_len, dtype=torch.int64).view(1, prefill_seq_len).repeat(bs, 1), + } + + mla_absorption = getattr(self.language_model, "mla_absorption", None) + if mla_absorption is not None: + cache_compressed = mla_absorption.get("cache_compressed", False) + else: + cache_compressed = False + + pkv_cache = self.language_model.get_dummy_pkv_cache( + config=self.language_model.config, + batch_size=fbs if continuous_batching else bs, + seq_len=prefill_seq_len, + ) + + if cache_compressed: + lang_inputs["compressed_kvs"] = [[] for _ in range(self.language_model.config.num_hidden_layers)] + for i in range(self.language_model.config.num_hidden_layers): + lang_inputs["compressed_kvs"][i].append( + torch.zeros(pkv_cache[0][0].shape, dtype=self.language_model.config.torch_dtype) + ) + lang_inputs["compressed_kvs"][i].append( + torch.zeros(pkv_cache[0][1].shape, dtype=self.language_model.config.torch_dtype) + ) + else: + lang_inputs["past_key_values"] = [[] for _ in range(self.language_model.config.num_hidden_layers)] + for i in range(self.language_model.config.num_hidden_layers): + lang_inputs["past_key_values"][i].append( + torch.zeros(pkv_cache[0][0].shape, dtype=self.language_model.config.torch_dtype) + ) + lang_inputs["past_key_values"][i].append( + torch.zeros(pkv_cache[0][1].shape, dtype=self.language_model.config.torch_dtype) + ) + + if continuous_batching: + lang_inputs["batch_index"] = torch.arange(bs).view(bs, 1) + + if comp_ctx_lengths is not None: + lang_inputs["comp_ctx_lengths"] = torch.randint(0, 100, (40,), dtype=torch.int64) + + inputs = {} + if kv_offload: + inputs["vision"] = vision_inputs + inputs["lang"] = lang_inputs + else: + lang_inputs.pop("image_embeds") + inputs = {**vision_inputs, **lang_inputs} + + return inputs + + def get_onnx_dynamic_axes( + self, comp_ctx_lengths: Optional[List[int]] = None, kv_offload: bool = False, continuous_batching: bool = False + ): + vision_dynamic_axes = {} + lang_dynamic_axes = {} + lang_dynamic_axes["input_ids"] = {0: "batch_size", 1: "seq_len"} + lang_dynamic_axes["position_ids"] = {0: "batch_size", 1: "seq_len"} + lang_dynamic_axes["image_embeds"] = {0: "num_image_tokens"} + if continuous_batching: + lang_dynamic_axes["batch_index"] = {0: "batch_size"} + vision_dynamic_axes = { + "pixel_values": {0: "num_patches"}, + "h_shape": {0: "h"}, + "w_shape": {0: "w"}, + "image_embeds": {0: "num_image_tokens"}, + } + + mla_absorption = getattr(self.language_model, "mla_absorption", None) + if mla_absorption is not None: + cache_compressed = mla_absorption.get("cache_compressed", False) + else: + cache_compressed = False + + if cache_compressed: + for i in range(self.language_model.config.num_hidden_layers): + lang_dynamic_axes[f"compressed_kv.{i}"] = {0: "batch_size", 2: "ctx_len"} + lang_dynamic_axes[f"k_pe.{i}"] = {0: "batch_size", 2: "ctx_len"} + else: + for i in range(self.language_model.config.num_hidden_layers): + for kv in ["key", "value"]: + lang_dynamic_axes[f"past_{kv}.{i}"] = {0: "batch_size", 2: "ctx_len"} + + if comp_ctx_lengths is not None: + lang_dynamic_axes["comp_ctx_lengths"] = {0: "comp_ctx_lengths"} + + dynamic_axes = {} + if kv_offload: + dynamic_axes["vision"] = vision_dynamic_axes + dynamic_axes["lang"] = lang_dynamic_axes + else: + lang_dynamic_axes.pop("image_embeds") + dynamic_axes = {**vision_dynamic_axes, **lang_dynamic_axes} + return dynamic_axes + + def get_specializations( + self, + batch_size: int, + prefill_seq_len: int, + ctx_len: int, + kv_offload: bool = False, + continuous_batching: bool = False, + kv_cache_batch_size: Optional[int] = None, + full_batch_size: Optional[int] = None, + **compiler_options, + ): + comp_ctx_lengths_prefill = compiler_options.pop("comp_ctx_lengths_prefill", None) + comp_ctx_lengths_decode = compiler_options.pop("comp_ctx_lengths_decode", None) + num_patches = compiler_options.pop("num_patches", None) + h = compiler_options.pop("h", None) + w = compiler_options.pop("w", None) + num_image_tokens = compiler_options.pop("num_image_tokens", None) + + prefill_seq_len = prefill_seq_len if prefill_seq_len else 32 # constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN + ctx_len = ctx_len if ctx_len else 32 # constants.ONNX_EXPORT_EXAMPLE_CTX_LEN + + vision = [ + { + "num_patches": 2400, # num_patches + "h": 30, # h + "w": 80, # w + "num_image_tokens": 600, # num_image_tokens + } + ] + + if comp_ctx_lengths_prefill is not None: + lang = [] + + for i in range(0, len(comp_ctx_lengths_prefill)): + lang_prefill = { + "batch_size": 1 if continuous_batching else batch_size, + "seq_len": prefill_seq_len, + "ctx_len": ctx_len, + "num_image_tokens": 600, # num_image_tokens + } + if continuous_batching: + lang_prefill["full_batch_size"] = kv_cache_batch_size + else: + lang_prefill["batch_size"] = kv_cache_batch_size + if full_batch_size: + lang_prefill["full_batch_exec_size"] = full_batch_size + + lang.append(lang_prefill) + + for i in range(0, len(comp_ctx_lengths_decode)): + lang_decode = { + "batch_size": full_batch_size if continuous_batching else batch_size, + "seq_len": "1", + "ctx_len": ctx_len, + "num_image_tokens": 600, # num_image_tokens + } + + if continuous_batching: + lang_decode["full_batch_size"] = kv_cache_batch_size + else: + lang_decode["batch_size"] = kv_cache_batch_size + + lang.append(lang_decode) + + else: + lang_prefill = { + "batch_size": 1 if continuous_batching else batch_size, + "seq_len": prefill_seq_len, + "ctx_len": ctx_len, + "num_image_tokens": 600, # num_image_tokens + } + if continuous_batching: + lang_prefill["full_batch_size"] = kv_cache_batch_size + else: + lang_prefill["batch_size"] = kv_cache_batch_size + if full_batch_size: + lang_prefill["full_batch_exec_size"] = full_batch_size + + lang_decode = { + "batch_size": full_batch_size if continuous_batching else batch_size, + "seq_len": 1, + "ctx_len": ctx_len, + "num_image_tokens": 600, # num_image_tokens + } + + if continuous_batching: + lang_decode["full_batch_size"] = kv_cache_batch_size + else: + lang_decode["batch_size"] = kv_cache_batch_size + + lang = [lang_prefill, lang_decode] + + specializations = {} + + if kv_offload: + specializations["vision"] = vision + specializations["lang"] = lang + return specializations, compiler_options + else: + lang[0].pop("vision_size") + lang[1].pop("vision_size") + return lang, compiler_options diff --git a/QEfficient/transformers/models/pytorch_transforms.py b/QEfficient/transformers/models/pytorch_transforms.py index 9117ff0e20..9d56f5b8e6 100755 --- a/QEfficient/transformers/models/pytorch_transforms.py +++ b/QEfficient/transformers/models/pytorch_transforms.py @@ -435,6 +435,11 @@ QEffInternVisionEmbeddings, QEffInternVLModel, ) +from QEfficient.transformers.models.kimi_k25.modeling_kimi_k25 import ( + QEffKimiK25ForConditionalGeneration, + QEffLearnable2DInterpPosEmbDivided_fixed, + QEffMoonViT3dEncoder, +) from QEfficient.transformers.models.llama.modeling_llama import ( QEffLlamaAttention, QEffLlamaDecoderLayer, @@ -1226,6 +1231,21 @@ class KVCacheExternalModuleMapperTransform(ExternalModuleMapperTransform): "RMSNorm": { "forward": QEFFGrok1CustomRMSNormAIC.forward, }, + "KimiK25ForConditionalGeneration": { + # "forward": QEffKimiK25ForConditionalGeneration.forward_only_image_for_export, + "get_qeff_vision_encoder": QEffKimiK25ForConditionalGeneration.get_qeff_vision_encoder, + "get_qeff_language_decoder": QEffKimiK25ForConditionalGeneration.get_qeff_language_decoder, + "get_specializations": QEffKimiK25ForConditionalGeneration.get_specializations, + "get_onnx_dynamic_axes": QEffKimiK25ForConditionalGeneration.get_onnx_dynamic_axes, + "get_output_names": QEffKimiK25ForConditionalGeneration.get_output_names, + "get_dummy_inputs": QEffKimiK25ForConditionalGeneration.get_dummy_inputs, + }, + "MoonViT3dEncoder": { + "__qeff_init__": QEffMoonViT3dEncoder.__qeff_init__, + }, + "Learnable2DInterpPosEmbDivided_fixed": { + "__qeff_init__": QEffLearnable2DInterpPosEmbDivided_fixed.__qeff_init__, + }, "DeepseekV3ForCausalLM": { "forward": QEffDeepseekV3ForCausalLM.forward, "get_submodules_for_export": QEffDeepseekV3ForCausalLM.get_submodules_for_export, diff --git a/examples/kimi_k2/export_kimi_k25_vision.py b/examples/kimi_k2/export_kimi_k25_vision.py new file mode 100644 index 0000000000..37aac14c0e --- /dev/null +++ b/examples/kimi_k2/export_kimi_k25_vision.py @@ -0,0 +1,423 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import argparse +import copy +import json +import re +import sys +import tempfile +from collections import defaultdict +from io import BytesIO +from pathlib import Path + +import requests +from PIL import Image +from safetensors import safe_open +from safetensors.torch import save_file +from transformers import AutoConfig, AutoProcessor, AutoTokenizer, TextStreamer +from transformers.dynamic_module_utils import get_class_from_dynamic_module + +from QEfficient import QEFFAutoModelForImageTextToText + +MODEL_PATH = Path( + "/home/huggingface_hub/models--moonshotai--Kimi-K2.5/snapshots/4d01dfe0332d63057c186e0b262165819efb6611" +) +NUM_VISION_LAYERS = 2 +NUM_TEXT_LAYERS = 2 +LOADED_EXPERT_IDS = (0, 1, 2, 3) +NUM_EXPERTS_PER_TOKEN = 2 + +EXPERT_KEY_PATTERN = re.compile(r"^(language_model\.model\.layers\.\d+\.mlp\.experts\.)(\d+)(\..+)$") + + +def _parse_expert_ids(value: str): + expert_ids = tuple(int(expert_id) for expert_id in value.split(",") if expert_id.strip()) + if not expert_ids: + raise argparse.ArgumentTypeError("At least one expert id must be provided.") + return expert_ids + + +def _validate_kimi_k25_config(kimi_config): + if kimi_config.model_type != "kimi_k25": + raise ValueError(f"This script only supports Kimi K2.5 config, got {kimi_config.model_type!r}.") + architectures = getattr(kimi_config, "architectures", None) or [] + if "KimiK25ForConditionalGeneration" not in architectures: + raise ValueError(f"Expected KimiK25ForConditionalGeneration, got architectures={architectures!r}.") + + +def _validate_layer_count(name, requested_count, available_count): + if requested_count < 1: + raise ValueError(f"{name} must be >= 1, got {requested_count}.") + if requested_count > available_count: + raise ValueError(f"{name}={requested_count} exceeds available layers={available_count}.") + + +def _validate_expert_subset(loaded_expert_ids, num_experts_per_tok, total_experts): + expert_ids = tuple(loaded_expert_ids) + if len(expert_ids) != 4: + raise ValueError(f"Expected exactly 4 routed experts, got {expert_ids!r}.") + if len(set(expert_ids)) != len(expert_ids): + raise ValueError(f"Expert ids must be unique, got {expert_ids!r}.") + invalid_ids = [expert_id for expert_id in expert_ids if expert_id < 0 or expert_id >= total_experts] + if invalid_ids: + raise ValueError(f"Expert ids {invalid_ids!r} are outside the valid range [0, {total_experts - 1}].") + if num_experts_per_tok > len(expert_ids): + raise ValueError(f"num_experts_per_tok={num_experts_per_tok} cannot exceed {len(expert_ids)} loaded experts.") + return expert_ids + + +def _extend_quantization_ignore(config, patterns): + quantization_config = getattr(config, "quantization_config", None) + if not quantization_config: + return + + ignored_modules = quantization_config.setdefault("ignore", []) + for pattern in patterns: + if pattern not in ignored_modules: + ignored_modules.append(pattern) + + +def _remap_checkpoint_key(checkpoint_key, expert_index_map): + match = EXPERT_KEY_PATTERN.match(checkpoint_key) + if not match: + return checkpoint_key + + original_expert_idx = int(match.group(2)) + remapped_expert_idx = expert_index_map.get(original_expert_idx) + if remapped_expert_idx is None: + return None + return f"{match.group(1)}{remapped_expert_idx}{match.group(3)}" + + +def _is_routed_gate_weight(checkpoint_key): + return checkpoint_key.endswith(".mlp.gate.weight") + + +def _is_routed_gate_bias(checkpoint_key): + return checkpoint_key.endswith(".mlp.gate.e_score_correction_bias") + + +def _materialize_subset_checkpoint( + model_path: Path, + temp_model_path: Path, + weight_map, + allowed_prefixes, + loaded_expert_ids, +): + expert_index_map = {expert_id: remapped_idx for remapped_idx, expert_id in enumerate(loaded_expert_ids)} + selected_by_shard = defaultdict(list) + + for checkpoint_key, shard_name in weight_map.items(): + if not any(checkpoint_key.startswith(prefix) for prefix in allowed_prefixes): + continue + + remapped_key = _remap_checkpoint_key(checkpoint_key, expert_index_map) + if remapped_key is None: + continue + selected_by_shard[shard_name].append((checkpoint_key, remapped_key)) + + if not selected_by_shard: + raise RuntimeError("No multimodal weights were selected from the Kimi K2.5 checkpoint.") + + filtered_weight_map = {} + subset_shards = [] + for shard_idx, (source_shard_name, shard_entries) in enumerate(sorted(selected_by_shard.items())): + tensors = {} + with safe_open(model_path / source_shard_name, framework="pt", device="cpu") as shard_reader: + for checkpoint_key, remapped_key in shard_entries: + tensor = shard_reader.get_tensor(checkpoint_key) + if _is_routed_gate_weight(checkpoint_key): + tensor = tensor[list(loaded_expert_ids), :].contiguous() + elif _is_routed_gate_bias(checkpoint_key): + tensor = tensor[list(loaded_expert_ids)].contiguous() + tensors[remapped_key] = tensor + + subset_shard_name = f"model-subset-{shard_idx:05d}.safetensors" + save_file(tensors, str(temp_model_path / subset_shard_name)) + subset_shards.append(subset_shard_name) + filtered_weight_map.update({remapped_key: subset_shard_name for _, remapped_key in shard_entries}) + + return filtered_weight_map, subset_shards + + +def _allowed_prefixes(num_vision_layers: int, num_text_layers: int): + prefixes = [ + "vision_tower.patch_embed.", + "vision_tower.encoder.final_layernorm.", + "mm_projector.", + "language_model.model.embed_tokens.", + "language_model.model.norm.", + "language_model.lm_head.", + ] + prefixes.extend(f"vision_tower.encoder.blocks.{layer_idx}." for layer_idx in range(num_vision_layers)) + prefixes.extend(f"language_model.model.layers.{layer_idx}." for layer_idx in range(num_text_layers)) + return prefixes + + +def load_kimi_k25_vision_moe_strip( + model_path: Path, + num_vision_layers: int, + num_text_layers: int, + loaded_expert_ids=LOADED_EXPERT_IDS, + num_experts_per_tok: int = NUM_EXPERTS_PER_TOKEN, +): + kimi_config = AutoConfig.from_pretrained("moonshotai/Kimi-K2.5", trust_remote_code=True) + _validate_kimi_k25_config(kimi_config) + + stripped_config = copy.deepcopy(kimi_config) + text_config = stripped_config.text_config + vision_config = stripped_config.vision_config + stripped_config._attn_implementation = "eager" + text_config._attn_implementation = "eager" + vision_config._attn_implementation = "eager" + quantization_ignores = [ + "re:vision_tower.*", + "re:mm_projector.*", + "re:language_model.lm_head.*", + ] + _extend_quantization_ignore(stripped_config, quantization_ignores) + _extend_quantization_ignore(text_config, quantization_ignores) + + _validate_layer_count("num_text_layers", num_text_layers, text_config.num_hidden_layers) + _validate_layer_count("num_vision_layers", num_vision_layers, vision_config.vt_num_hidden_layers) + + text_config.num_hidden_layers = num_text_layers + vision_config.vt_num_hidden_layers = num_vision_layers + loaded_expert_ids = _validate_expert_subset( + loaded_expert_ids, + num_experts_per_tok, + text_config.n_routed_experts, + ) + text_config.n_routed_experts = len(loaded_expert_ids) + text_config.num_experts_per_tok = num_experts_per_tok + text_config.n_group = 1 + text_config.topk_group = 1 + + kimi_cls = get_class_from_dynamic_module("modeling_kimi_k25.KimiK25ForConditionalGeneration", str(model_path)) + kimi_module = sys.modules[kimi_cls.__module__] + kimi_module.MoonViT3dEncoder.use_deterministic_attn = False + + checkpoint_index = json.loads((model_path / "model.safetensors.index.json").read_text()) + weight_map = checkpoint_index["weight_map"] + + with tempfile.TemporaryDirectory() as tmpdir: + temp_model_path = Path(tmpdir) + filtered_weight_map, subset_shards = _materialize_subset_checkpoint( + model_path=model_path, + temp_model_path=temp_model_path, + weight_map=weight_map, + allowed_prefixes=_allowed_prefixes(num_vision_layers, num_text_layers), + loaded_expert_ids=loaded_expert_ids, + ) + (temp_model_path / "config.json").write_text(stripped_config.to_json_string(use_diff=False)) + (temp_model_path / "model.safetensors.index.json").write_text( + json.dumps( + { + "metadata": { + "total_size": sum((temp_model_path / shard_name).stat().st_size for shard_name in subset_shards) + }, + "weight_map": filtered_weight_map, + } + ) + ) + + original_base_model_prefix = kimi_cls.base_model_prefix + kimi_cls.base_model_prefix = "" + try: + model, loading_info = kimi_cls.from_pretrained( + str(temp_model_path), + config=stripped_config, + local_files_only=True, + attn_implementation="eager", + output_loading_info=True, + ) + finally: + kimi_cls.base_model_prefix = original_base_model_prefix + + unexpected_keys = loading_info["unexpected_keys"] + missing_keys = loading_info["missing_keys"] + mismatched_keys = loading_info["mismatched_keys"] + if unexpected_keys or missing_keys or mismatched_keys: + raise RuntimeError( + "Failed to load the stripped Kimi K2.5 multimodal checkpoint slice cleanly. " + f"missing={missing_keys}, unexpected={unexpected_keys}, mismatched={mismatched_keys}" + ) + + model.eval() + tokenizer = AutoTokenizer.from_pretrained(str(model_path), trust_remote_code=True) + processor = AutoProcessor.from_pretrained(str(model_path), trust_remote_code=True) + return model, tokenizer, processor + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Load a stripped Kimi K2.5 multimodal checkpoint with configurable vision/text layers and 4 routed experts." + ) + parser.add_argument("--model-path", type=Path, default=MODEL_PATH) + parser.add_argument("--num-vision-layers", type=int, default=NUM_VISION_LAYERS) + parser.add_argument("--num-text-layers", type=int, default=NUM_TEXT_LAYERS) + parser.add_argument("--expert-ids", type=_parse_expert_ids, default=LOADED_EXPERT_IDS) + parser.add_argument("--num-experts-per-token", type=int, default=NUM_EXPERTS_PER_TOKEN) + parser.add_argument( + "--image-url", + type=str, + default="https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png", + ) + parser.add_argument("--prompt", type=str, default="Describe this image.") + parser.add_argument("--test", action="store_true", help="Validate ONNX output matches PyTorch image-only forward.") + return parser.parse_args() + + +def main(): + args = parse_args() + model, tokenizer, processor = load_kimi_k25_vision_moe_strip( + args.model_path, + args.num_vision_layers, + args.num_text_layers, + loaded_expert_ids=args.expert_ids, + num_experts_per_tok=args.num_experts_per_token, + ) + print( + f"Loaded {type(model).__name__} with " + f"{model.config.vision_config.vt_num_hidden_layers} vision layers, " + f"{model.config.text_config.num_hidden_layers} text layers, " + f"{model.config.text_config.n_routed_experts} routed experts." + ) + print(f"Tokenizer vocab size: {tokenizer.vocab_size}") + print(f"Processor type: {type(processor).__name__}") + + mla_absorption = {"cache_compressed": True, "absorption": False, "online": False} + qaic_config = { + "mla_absorption": mla_absorption + } # , "enable_blocking": True, "blocking_mode": "par", "par_num_split": 4, "num_kv_blocks": 8} + breakpoint() + qeff_model = QEFFAutoModelForImageTextToText(model) # , qaic_config=qaic_config) + breakpoint() + + skip_vision = True + + if skip_vision: + ## TEXT-ONLY MODE ## + + ## STEP 3: Compile Model for Text-Only Execution + # Set skip_vision=True to bypass image processing + qeff_model.compile( + qaic_config=qaic_config, + prefill_seq_len=32, + ctx_len=1024, + num_cores=16, + num_devices=2, + mxfp6_matmul=False, + mxint8_kv_cache=False, + aic_enable_depth_first=False, + skip_vision=True, # Skip vision encoder for text-only inference + mos=1, + num_patches=2400, # num_patches + h=30, # h + w=80, # w + num_image_tokens=600, # num_image_tokens + ) + breakpoint() + ## STEP 4: Prepare Text-Only Input + # Create a text-only message without any image + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Tell me about yourself."}, + ], + }, + ] + + ## STEP 5: Process Input with Chat Template + inputs = processor.apply_chat_template( + messages, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors="pt", + ) + + ## STEP 6: Run Text-Only Inference + streamer = TextStreamer(tokenizer) + output = qeff_model.generate(inputs=inputs, device_ids=[0, 1], generation_len=10) + + ## STEP 7: Display Results + print(output.generated_ids) + print(tokenizer.batch_decode(output.generated_ids)) + print(output) + + else: + ## VISION + TEXT MODE ## + + ## STEP 3: Compile Model for Vision+Text Execution + # Do not set skip_vision (defaults to False) to enable image processing + qeff_model.compile( + qaic_config=qaic_config, + prefill_seq_len=32, + ctx_len=1024, + num_cores=16, + num_devices=2, + mxfp6_matmul=False, + mxint8_kv_cache=False, + aic_enable_depth_first=False, + skip_vision=True, # Skip vision encoder for text-only inference + mos=1, + num_patches=2400, # num_patches + h=30, # h + w=80, # w + num_image_tokens=600, # num_image_tokens + ) + + breakpoint() + ## STEP 4: Prepare Image and Text Input + # Define the image URL to process + image_url = "https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png" + image = Image.open(BytesIO(requests.get(args.image_url).content)).convert("RGB") + + # Create a message with both image and text + messages = [ + { + "role": "user", + "content": [ + {"type": "image", "url": image}, + {"type": "text", "text": "Can you describe the image in detail."}, + ], + }, + ] + + ## STEP 5: Process Input with Chat Template + """inputs = processor.apply_chat_template( + messages, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors="pt", + ) + """ + inputs = processor( + messages=messages, + add_generation_prompt=True, + tokenize=False, + return_tensors="pt", + ) + # Convert pixel values to float32 for processing + inputs["pixel_values"] = inputs["pixel_values"].to(qeff_model.model.config.torch_dtype) + + ## STEP 6: Run Vision+Text Inference + streamer = TextStreamer(tokenizer) + output = qeff_model.generate(inputs=inputs, device_ids=[0, 1], generation_len=100) + + ## STEP 7: Display Results + print(output.generated_ids) + print(tokenizer.batch_decode(output.generated_ids)) + print(output) + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index ee4ca672b9..5255ffc3d0 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,7 +19,7 @@ classifiers = [ ] requires-python = ">=3.8,<3.13" dependencies = [ - "transformers==5.5.4", + "transformers==5.8", "diffusers==0.38.0", "huggingface-hub==1.7.1", "hf_transfer==0.1.9", From 71d0a7d4960efc1f41444f1a23251a024effa40d Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Sun, 21 Jun 2026 18:55:35 +0530 Subject: [PATCH 02/33] update load script Signed-off-by: Mamta Singh --- examples/kimi_k2/export_kimi_k25_vision.py | 834 +++++++++--------- .../kimi_k2/export_kimi_k25_vision_4.57.3.py | 423 +++++++++ pyproject.toml | 2 +- 3 files changed, 835 insertions(+), 424 deletions(-) create mode 100644 examples/kimi_k2/export_kimi_k25_vision_4.57.3.py diff --git a/examples/kimi_k2/export_kimi_k25_vision.py b/examples/kimi_k2/export_kimi_k25_vision.py index 37aac14c0e..ba34938ed4 100644 --- a/examples/kimi_k2/export_kimi_k25_vision.py +++ b/examples/kimi_k2/export_kimi_k25_vision.py @@ -1,423 +1,411 @@ -# ----------------------------------------------------------------------------- -# -# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. -# SPDX-License-Identifier: BSD-3-Clause -# -# ----------------------------------------------------------------------------- - -import argparse -import copy -import json -import re -import sys -import tempfile -from collections import defaultdict -from io import BytesIO -from pathlib import Path - -import requests -from PIL import Image -from safetensors import safe_open -from safetensors.torch import save_file -from transformers import AutoConfig, AutoProcessor, AutoTokenizer, TextStreamer -from transformers.dynamic_module_utils import get_class_from_dynamic_module - -from QEfficient import QEFFAutoModelForImageTextToText - -MODEL_PATH = Path( - "/home/huggingface_hub/models--moonshotai--Kimi-K2.5/snapshots/4d01dfe0332d63057c186e0b262165819efb6611" -) -NUM_VISION_LAYERS = 2 -NUM_TEXT_LAYERS = 2 -LOADED_EXPERT_IDS = (0, 1, 2, 3) -NUM_EXPERTS_PER_TOKEN = 2 - -EXPERT_KEY_PATTERN = re.compile(r"^(language_model\.model\.layers\.\d+\.mlp\.experts\.)(\d+)(\..+)$") - - -def _parse_expert_ids(value: str): - expert_ids = tuple(int(expert_id) for expert_id in value.split(",") if expert_id.strip()) - if not expert_ids: - raise argparse.ArgumentTypeError("At least one expert id must be provided.") - return expert_ids - - -def _validate_kimi_k25_config(kimi_config): - if kimi_config.model_type != "kimi_k25": - raise ValueError(f"This script only supports Kimi K2.5 config, got {kimi_config.model_type!r}.") - architectures = getattr(kimi_config, "architectures", None) or [] - if "KimiK25ForConditionalGeneration" not in architectures: - raise ValueError(f"Expected KimiK25ForConditionalGeneration, got architectures={architectures!r}.") - - -def _validate_layer_count(name, requested_count, available_count): - if requested_count < 1: - raise ValueError(f"{name} must be >= 1, got {requested_count}.") - if requested_count > available_count: - raise ValueError(f"{name}={requested_count} exceeds available layers={available_count}.") - - -def _validate_expert_subset(loaded_expert_ids, num_experts_per_tok, total_experts): - expert_ids = tuple(loaded_expert_ids) - if len(expert_ids) != 4: - raise ValueError(f"Expected exactly 4 routed experts, got {expert_ids!r}.") - if len(set(expert_ids)) != len(expert_ids): - raise ValueError(f"Expert ids must be unique, got {expert_ids!r}.") - invalid_ids = [expert_id for expert_id in expert_ids if expert_id < 0 or expert_id >= total_experts] - if invalid_ids: - raise ValueError(f"Expert ids {invalid_ids!r} are outside the valid range [0, {total_experts - 1}].") - if num_experts_per_tok > len(expert_ids): - raise ValueError(f"num_experts_per_tok={num_experts_per_tok} cannot exceed {len(expert_ids)} loaded experts.") - return expert_ids - - -def _extend_quantization_ignore(config, patterns): - quantization_config = getattr(config, "quantization_config", None) - if not quantization_config: - return - - ignored_modules = quantization_config.setdefault("ignore", []) - for pattern in patterns: - if pattern not in ignored_modules: - ignored_modules.append(pattern) - - -def _remap_checkpoint_key(checkpoint_key, expert_index_map): - match = EXPERT_KEY_PATTERN.match(checkpoint_key) - if not match: - return checkpoint_key - - original_expert_idx = int(match.group(2)) - remapped_expert_idx = expert_index_map.get(original_expert_idx) - if remapped_expert_idx is None: - return None - return f"{match.group(1)}{remapped_expert_idx}{match.group(3)}" - - -def _is_routed_gate_weight(checkpoint_key): - return checkpoint_key.endswith(".mlp.gate.weight") - - -def _is_routed_gate_bias(checkpoint_key): - return checkpoint_key.endswith(".mlp.gate.e_score_correction_bias") - - -def _materialize_subset_checkpoint( - model_path: Path, - temp_model_path: Path, - weight_map, - allowed_prefixes, - loaded_expert_ids, -): - expert_index_map = {expert_id: remapped_idx for remapped_idx, expert_id in enumerate(loaded_expert_ids)} - selected_by_shard = defaultdict(list) - - for checkpoint_key, shard_name in weight_map.items(): - if not any(checkpoint_key.startswith(prefix) for prefix in allowed_prefixes): - continue - - remapped_key = _remap_checkpoint_key(checkpoint_key, expert_index_map) - if remapped_key is None: - continue - selected_by_shard[shard_name].append((checkpoint_key, remapped_key)) - - if not selected_by_shard: - raise RuntimeError("No multimodal weights were selected from the Kimi K2.5 checkpoint.") - - filtered_weight_map = {} - subset_shards = [] - for shard_idx, (source_shard_name, shard_entries) in enumerate(sorted(selected_by_shard.items())): - tensors = {} - with safe_open(model_path / source_shard_name, framework="pt", device="cpu") as shard_reader: - for checkpoint_key, remapped_key in shard_entries: - tensor = shard_reader.get_tensor(checkpoint_key) - if _is_routed_gate_weight(checkpoint_key): - tensor = tensor[list(loaded_expert_ids), :].contiguous() - elif _is_routed_gate_bias(checkpoint_key): - tensor = tensor[list(loaded_expert_ids)].contiguous() - tensors[remapped_key] = tensor - - subset_shard_name = f"model-subset-{shard_idx:05d}.safetensors" - save_file(tensors, str(temp_model_path / subset_shard_name)) - subset_shards.append(subset_shard_name) - filtered_weight_map.update({remapped_key: subset_shard_name for _, remapped_key in shard_entries}) - - return filtered_weight_map, subset_shards - - -def _allowed_prefixes(num_vision_layers: int, num_text_layers: int): - prefixes = [ - "vision_tower.patch_embed.", - "vision_tower.encoder.final_layernorm.", - "mm_projector.", - "language_model.model.embed_tokens.", - "language_model.model.norm.", - "language_model.lm_head.", - ] - prefixes.extend(f"vision_tower.encoder.blocks.{layer_idx}." for layer_idx in range(num_vision_layers)) - prefixes.extend(f"language_model.model.layers.{layer_idx}." for layer_idx in range(num_text_layers)) - return prefixes - - -def load_kimi_k25_vision_moe_strip( - model_path: Path, - num_vision_layers: int, - num_text_layers: int, - loaded_expert_ids=LOADED_EXPERT_IDS, - num_experts_per_tok: int = NUM_EXPERTS_PER_TOKEN, -): - kimi_config = AutoConfig.from_pretrained("moonshotai/Kimi-K2.5", trust_remote_code=True) - _validate_kimi_k25_config(kimi_config) - - stripped_config = copy.deepcopy(kimi_config) - text_config = stripped_config.text_config - vision_config = stripped_config.vision_config - stripped_config._attn_implementation = "eager" - text_config._attn_implementation = "eager" - vision_config._attn_implementation = "eager" - quantization_ignores = [ - "re:vision_tower.*", - "re:mm_projector.*", - "re:language_model.lm_head.*", - ] - _extend_quantization_ignore(stripped_config, quantization_ignores) - _extend_quantization_ignore(text_config, quantization_ignores) - - _validate_layer_count("num_text_layers", num_text_layers, text_config.num_hidden_layers) - _validate_layer_count("num_vision_layers", num_vision_layers, vision_config.vt_num_hidden_layers) - - text_config.num_hidden_layers = num_text_layers - vision_config.vt_num_hidden_layers = num_vision_layers - loaded_expert_ids = _validate_expert_subset( - loaded_expert_ids, - num_experts_per_tok, - text_config.n_routed_experts, - ) - text_config.n_routed_experts = len(loaded_expert_ids) - text_config.num_experts_per_tok = num_experts_per_tok - text_config.n_group = 1 - text_config.topk_group = 1 - - kimi_cls = get_class_from_dynamic_module("modeling_kimi_k25.KimiK25ForConditionalGeneration", str(model_path)) - kimi_module = sys.modules[kimi_cls.__module__] - kimi_module.MoonViT3dEncoder.use_deterministic_attn = False - - checkpoint_index = json.loads((model_path / "model.safetensors.index.json").read_text()) - weight_map = checkpoint_index["weight_map"] - - with tempfile.TemporaryDirectory() as tmpdir: - temp_model_path = Path(tmpdir) - filtered_weight_map, subset_shards = _materialize_subset_checkpoint( - model_path=model_path, - temp_model_path=temp_model_path, - weight_map=weight_map, - allowed_prefixes=_allowed_prefixes(num_vision_layers, num_text_layers), - loaded_expert_ids=loaded_expert_ids, - ) - (temp_model_path / "config.json").write_text(stripped_config.to_json_string(use_diff=False)) - (temp_model_path / "model.safetensors.index.json").write_text( - json.dumps( - { - "metadata": { - "total_size": sum((temp_model_path / shard_name).stat().st_size for shard_name in subset_shards) - }, - "weight_map": filtered_weight_map, - } - ) - ) - - original_base_model_prefix = kimi_cls.base_model_prefix - kimi_cls.base_model_prefix = "" - try: - model, loading_info = kimi_cls.from_pretrained( - str(temp_model_path), - config=stripped_config, - local_files_only=True, - attn_implementation="eager", - output_loading_info=True, - ) - finally: - kimi_cls.base_model_prefix = original_base_model_prefix - - unexpected_keys = loading_info["unexpected_keys"] - missing_keys = loading_info["missing_keys"] - mismatched_keys = loading_info["mismatched_keys"] - if unexpected_keys or missing_keys or mismatched_keys: - raise RuntimeError( - "Failed to load the stripped Kimi K2.5 multimodal checkpoint slice cleanly. " - f"missing={missing_keys}, unexpected={unexpected_keys}, mismatched={mismatched_keys}" - ) - - model.eval() - tokenizer = AutoTokenizer.from_pretrained(str(model_path), trust_remote_code=True) - processor = AutoProcessor.from_pretrained(str(model_path), trust_remote_code=True) - return model, tokenizer, processor - - -def parse_args(): - parser = argparse.ArgumentParser( - description="Load a stripped Kimi K2.5 multimodal checkpoint with configurable vision/text layers and 4 routed experts." - ) - parser.add_argument("--model-path", type=Path, default=MODEL_PATH) - parser.add_argument("--num-vision-layers", type=int, default=NUM_VISION_LAYERS) - parser.add_argument("--num-text-layers", type=int, default=NUM_TEXT_LAYERS) - parser.add_argument("--expert-ids", type=_parse_expert_ids, default=LOADED_EXPERT_IDS) - parser.add_argument("--num-experts-per-token", type=int, default=NUM_EXPERTS_PER_TOKEN) - parser.add_argument( - "--image-url", - type=str, - default="https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png", - ) - parser.add_argument("--prompt", type=str, default="Describe this image.") - parser.add_argument("--test", action="store_true", help="Validate ONNX output matches PyTorch image-only forward.") - return parser.parse_args() - - -def main(): - args = parse_args() - model, tokenizer, processor = load_kimi_k25_vision_moe_strip( - args.model_path, - args.num_vision_layers, - args.num_text_layers, - loaded_expert_ids=args.expert_ids, - num_experts_per_tok=args.num_experts_per_token, - ) - print( - f"Loaded {type(model).__name__} with " - f"{model.config.vision_config.vt_num_hidden_layers} vision layers, " - f"{model.config.text_config.num_hidden_layers} text layers, " - f"{model.config.text_config.n_routed_experts} routed experts." - ) - print(f"Tokenizer vocab size: {tokenizer.vocab_size}") - print(f"Processor type: {type(processor).__name__}") - - mla_absorption = {"cache_compressed": True, "absorption": False, "online": False} - qaic_config = { - "mla_absorption": mla_absorption - } # , "enable_blocking": True, "blocking_mode": "par", "par_num_split": 4, "num_kv_blocks": 8} - breakpoint() - qeff_model = QEFFAutoModelForImageTextToText(model) # , qaic_config=qaic_config) - breakpoint() - - skip_vision = True - - if skip_vision: - ## TEXT-ONLY MODE ## - - ## STEP 3: Compile Model for Text-Only Execution - # Set skip_vision=True to bypass image processing - qeff_model.compile( - qaic_config=qaic_config, - prefill_seq_len=32, - ctx_len=1024, - num_cores=16, - num_devices=2, - mxfp6_matmul=False, - mxint8_kv_cache=False, - aic_enable_depth_first=False, - skip_vision=True, # Skip vision encoder for text-only inference - mos=1, - num_patches=2400, # num_patches - h=30, # h - w=80, # w - num_image_tokens=600, # num_image_tokens - ) - breakpoint() - ## STEP 4: Prepare Text-Only Input - # Create a text-only message without any image - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "Tell me about yourself."}, - ], - }, - ] - - ## STEP 5: Process Input with Chat Template - inputs = processor.apply_chat_template( - messages, - add_generation_prompt=True, - tokenize=True, - return_dict=True, - return_tensors="pt", - ) - - ## STEP 6: Run Text-Only Inference - streamer = TextStreamer(tokenizer) - output = qeff_model.generate(inputs=inputs, device_ids=[0, 1], generation_len=10) - - ## STEP 7: Display Results - print(output.generated_ids) - print(tokenizer.batch_decode(output.generated_ids)) - print(output) - - else: - ## VISION + TEXT MODE ## - - ## STEP 3: Compile Model for Vision+Text Execution - # Do not set skip_vision (defaults to False) to enable image processing - qeff_model.compile( - qaic_config=qaic_config, - prefill_seq_len=32, - ctx_len=1024, - num_cores=16, - num_devices=2, - mxfp6_matmul=False, - mxint8_kv_cache=False, - aic_enable_depth_first=False, - skip_vision=True, # Skip vision encoder for text-only inference - mos=1, - num_patches=2400, # num_patches - h=30, # h - w=80, # w - num_image_tokens=600, # num_image_tokens - ) - - breakpoint() - ## STEP 4: Prepare Image and Text Input - # Define the image URL to process - image_url = "https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png" - image = Image.open(BytesIO(requests.get(args.image_url).content)).convert("RGB") - - # Create a message with both image and text - messages = [ - { - "role": "user", - "content": [ - {"type": "image", "url": image}, - {"type": "text", "text": "Can you describe the image in detail."}, - ], - }, - ] - - ## STEP 5: Process Input with Chat Template - """inputs = processor.apply_chat_template( - messages, - add_generation_prompt=True, - tokenize=True, - return_dict=True, - return_tensors="pt", - ) - """ - inputs = processor( - messages=messages, - add_generation_prompt=True, - tokenize=False, - return_tensors="pt", - ) - # Convert pixel values to float32 for processing - inputs["pixel_values"] = inputs["pixel_values"].to(qeff_model.model.config.torch_dtype) - - ## STEP 6: Run Vision+Text Inference - streamer = TextStreamer(tokenizer) - output = qeff_model.generate(inputs=inputs, device_ids=[0, 1], generation_len=100) - - ## STEP 7: Display Results - print(output.generated_ids) - print(tokenizer.batch_decode(output.generated_ids)) - print(output) - - -if __name__ == "__main__": - main() +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import argparse +import copy +import inspect +import json +import re +import sys +import tempfile +from pathlib import Path + +import torch +from safetensors import safe_open +from safetensors.torch import save_file +from transformers import AutoConfig, AutoProcessor, AutoTokenizer +from transformers.dynamic_module_utils import get_class_from_dynamic_module +from transformers.utils import import_utils as hf_import_utils + + +MODEL_PATH = Path( + "/home/huggingface_hub/models--moonshotai--Kimi-K2.5/snapshots/4d01dfe0332d63057c186e0b262165819efb6611" +) +NUM_VISION_LAYERS = 2 +NUM_TEXT_LAYERS = 2 +LOADED_EXPERT_IDS = (0, 1, 2, 3) +NUM_EXPERTS_PER_TOKEN = 2 + +EXPERT_KEY_PATTERN = re.compile(r"^(language_model\.model\.layers\.\d+\.mlp\.experts\.)(\d+)(\..+)$") + + +def _ensure_torch_fx_import_compatibility(): + """Backfill `is_torch_fx_available` for remote model code expecting older Transformers APIs.""" + if hasattr(hf_import_utils, "is_torch_fx_available"): + return + + def _is_torch_fx_available() -> bool: + if not hf_import_utils.is_torch_available(): + return False + try: + import torch.fx # noqa: F401 + + return True + except Exception: + return False + + hf_import_utils.is_torch_fx_available = _is_torch_fx_available + + +def _patch_kimi_tie_weights_compat(kimi_cls): + tie_signature = inspect.signature(kimi_cls.tie_weights) + if tuple(tie_signature.parameters) != ("self",): + return + + def _tie_weights_compat(self, missing_keys=None, recompute_mapping=True): + lm_tie_weights = getattr(self.language_model, "tie_weights") + try: + return lm_tie_weights(missing_keys=missing_keys, recompute_mapping=recompute_mapping) + except TypeError: + return lm_tie_weights() + + kimi_cls.tie_weights = _tie_weights_compat + + +def _patch_deepseek_init_weights_compat(kimi_cls): + module_prefix, _ = kimi_cls.__module__.rsplit(".", maxsplit=1) + deepseek_module_name = f"{module_prefix}.modeling_deepseek" + deepseek_module = sys.modules.get(deepseek_module_name) + if deepseek_module is None or not hasattr(deepseek_module, "DeepseekV3PreTrainedModel"): + return + + deepseek_cls = deepseek_module.DeepseekV3PreTrainedModel + if getattr(deepseek_cls, "_qeff_t55_init_weights_patched", False): + return + + def _init_weights_compat(self, module): + std = self.config.initializer_range + if isinstance(module, torch.nn.Linear): + if hasattr(module, "weight") and module.weight is not None: + module.weight.data.normal_(mean=0.0, std=std) + if hasattr(module, "bias") and module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, torch.nn.Embedding): + if hasattr(module, "weight") and module.weight is not None: + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + deepseek_cls._init_weights = _init_weights_compat + deepseek_cls._qeff_t55_init_weights_patched = True + + +def _prepare_config(model_path: Path): + config = AutoConfig.from_pretrained(str(model_path), trust_remote_code=True) + + # Avoid FA2 dispatch checks in Transformers 5.5.x for this remote code. + config._attn_implementation = "eager" + if hasattr(config, "text_config"): + config.text_config._attn_implementation = "eager" + if hasattr(config, "vision_config"): + config.vision_config._attn_implementation = "eager" + return config + + +def _parse_expert_ids(value: str): + expert_ids = tuple(int(expert_id) for expert_id in value.split(",") if expert_id.strip()) + if not expert_ids: + raise argparse.ArgumentTypeError("At least one expert id must be provided.") + return expert_ids + + +def _validate_layer_count(name, requested_count, available_count): + if requested_count < 1: + raise ValueError(f"{name} must be >= 1, got {requested_count}.") + if requested_count > available_count: + raise ValueError(f"{name}={requested_count} exceeds available layers={available_count}.") + + +def _validate_expert_subset(loaded_expert_ids, num_experts_per_tok, total_experts): + expert_ids = tuple(loaded_expert_ids) + if len(expert_ids) != 4: + raise ValueError(f"Expected exactly 4 routed experts, got {expert_ids!r}.") + if len(set(expert_ids)) != len(expert_ids): + raise ValueError(f"Expert ids must be unique, got {expert_ids!r}.") + invalid_ids = [expert_id for expert_id in expert_ids if expert_id < 0 or expert_id >= total_experts] + if invalid_ids: + raise ValueError(f"Expert ids {invalid_ids!r} are outside the valid range [0, {total_experts - 1}].") + if num_experts_per_tok > len(expert_ids): + raise ValueError(f"num_experts_per_tok={num_experts_per_tok} cannot exceed {len(expert_ids)} loaded experts.") + return expert_ids + + +def _remap_checkpoint_key(checkpoint_key, expert_index_map): + match = EXPERT_KEY_PATTERN.match(checkpoint_key) + if not match: + return checkpoint_key + + original_expert_idx = int(match.group(2)) + remapped_expert_idx = expert_index_map.get(original_expert_idx) + if remapped_expert_idx is None: + return None + return f"{match.group(1)}{remapped_expert_idx}{match.group(3)}" + + +def _is_routed_gate_weight(checkpoint_key): + return checkpoint_key.endswith(".mlp.gate.weight") + + +def _is_routed_gate_bias(checkpoint_key): + return checkpoint_key.endswith(".mlp.gate.e_score_correction_bias") + + +def _allowed_prefixes(num_vision_layers: int, num_text_layers: int): + prefixes = [ + "vision_tower.patch_embed.", + "vision_tower.encoder.final_layernorm.", + "mm_projector.", + "language_model.model.embed_tokens.", + "language_model.model.norm.", + "language_model.lm_head.", + ] + prefixes.extend(f"vision_tower.encoder.blocks.{layer_idx}." for layer_idx in range(num_vision_layers)) + prefixes.extend(f"language_model.model.layers.{layer_idx}." for layer_idx in range(num_text_layers)) + return prefixes + + +def _build_layer_subset_config(config, num_vision_layers, num_text_layers, loaded_expert_ids, num_experts_per_tok): + stripped_config = copy.deepcopy(config) + text_config = stripped_config.text_config + vision_config = stripped_config.vision_config + + _validate_layer_count("num_text_layers", num_text_layers, text_config.num_hidden_layers) + _validate_layer_count("num_vision_layers", num_vision_layers, vision_config.vt_num_hidden_layers) + + text_config.num_hidden_layers = num_text_layers + vision_config.vt_num_hidden_layers = num_vision_layers + + loaded_expert_ids = _validate_expert_subset( + loaded_expert_ids, + num_experts_per_tok, + text_config.n_routed_experts, + ) + text_config.n_routed_experts = len(loaded_expert_ids) + text_config.num_experts_per_tok = num_experts_per_tok + text_config.n_group = 1 + text_config.topk_group = 1 + return stripped_config, loaded_expert_ids + + +def _materialize_subset_checkpoint(model_path: Path, temp_model_path: Path, weight_map, allowed_prefixes, loaded_expert_ids): + expert_index_map = {expert_id: remapped_idx for remapped_idx, expert_id in enumerate(loaded_expert_ids)} + shard_to_entries = {} + for checkpoint_key, source_shard_name in weight_map.items(): + if not checkpoint_key.startswith(tuple(allowed_prefixes)): + continue + + remapped_key = _remap_checkpoint_key(checkpoint_key, expert_index_map) + if remapped_key is None: + continue + shard_to_entries.setdefault(source_shard_name, []).append((checkpoint_key, remapped_key)) + + filtered_weight_map = {} + subset_shards = [] + for shard_idx, (source_shard_name, shard_entries) in enumerate(sorted(shard_to_entries.items())): + tensors = {} + with safe_open(model_path / source_shard_name, framework="pt", device="cpu") as shard_reader: + for checkpoint_key, remapped_key in shard_entries: + tensor = shard_reader.get_tensor(checkpoint_key) + if _is_routed_gate_weight(checkpoint_key): + tensor = tensor[list(loaded_expert_ids), :].contiguous() + elif _is_routed_gate_bias(checkpoint_key): + tensor = tensor[list(loaded_expert_ids)].contiguous() + tensors[remapped_key] = tensor + + subset_shard_name = f"model-subset-{shard_idx:05d}.safetensors" + save_file(tensors, str(temp_model_path / subset_shard_name)) + subset_shards.append(subset_shard_name) + filtered_weight_map.update({remapped_key: subset_shard_name for _, remapped_key in shard_entries}) + + return filtered_weight_map, subset_shards + + +def _load_layer_subset_model( + *, + model_path: Path, + kimi_cls, + config, + num_vision_layers: int, + num_text_layers: int, + loaded_expert_ids, + num_experts_per_tok: int, + local_files_only: bool, + dtype, +): + checkpoint_index = json.loads((model_path / "model.safetensors.index.json").read_text()) + weight_map = checkpoint_index["weight_map"] + stripped_config, loaded_expert_ids = _build_layer_subset_config( + config, + num_vision_layers=num_vision_layers, + num_text_layers=num_text_layers, + loaded_expert_ids=loaded_expert_ids, + num_experts_per_tok=num_experts_per_tok, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + temp_model_path = Path(tmpdir) + filtered_weight_map, subset_shards = _materialize_subset_checkpoint( + model_path=model_path, + temp_model_path=temp_model_path, + weight_map=weight_map, + allowed_prefixes=_allowed_prefixes(num_vision_layers, num_text_layers), + loaded_expert_ids=loaded_expert_ids, + ) + (temp_model_path / "config.json").write_text(stripped_config.to_json_string(use_diff=False)) + (temp_model_path / "model.safetensors.index.json").write_text( + json.dumps( + { + "metadata": {"total_size": sum((temp_model_path / shard_name).stat().st_size for shard_name in subset_shards)}, + "weight_map": filtered_weight_map, + } + ) + ) + + model_kwargs = { + "config": stripped_config, + "trust_remote_code": True, + "attn_implementation": "eager", + "local_files_only": local_files_only, + "output_loading_info": True, + } + if dtype is not None: + model_kwargs["torch_dtype"] = dtype + + original_base_model_prefix = kimi_cls.base_model_prefix + kimi_cls.base_model_prefix = "" + try: + model, loading_info = kimi_cls.from_pretrained(str(temp_model_path), **model_kwargs) + finally: + kimi_cls.base_model_prefix = original_base_model_prefix + + unexpected_keys = loading_info["unexpected_keys"] + missing_keys = loading_info["missing_keys"] + mismatched_keys = loading_info["mismatched_keys"] + if unexpected_keys or missing_keys or mismatched_keys: + raise RuntimeError( + "Failed to load the stripped Kimi K2.5 checkpoint slice cleanly. " + f"missing={missing_keys}, unexpected={unexpected_keys}, mismatched={mismatched_keys}" + ) + + return model + + +def parse_args(): + parser = argparse.ArgumentParser(description="Load Kimi K2.5 with runtime compatibility for transformers==5.5.4.") + parser.add_argument("--model-path", type=Path, default=MODEL_PATH) + parser.add_argument("--local-files-only", action="store_true", help="Only read from local cache/files.") + parser.add_argument( + "--full-model", + action="store_true", + help="Load the full model. By default, the script loads a small layer subset for faster startup.", + ) + parser.add_argument( + "--num-vision-layers", + type=int, + default=NUM_VISION_LAYERS, + help=( + "Load only the first N vision layers. " + f"Defaults to {NUM_VISION_LAYERS} for faster startup. Ignored when --full-model is set." + ), + ) + parser.add_argument( + "--num-text-layers", + type=int, + default=NUM_TEXT_LAYERS, + help=( + "Load only the first N text layers. " + f"Defaults to {NUM_TEXT_LAYERS} for faster startup. Ignored when --full-model is set." + ), + ) + parser.add_argument("--expert-ids", type=_parse_expert_ids, default=LOADED_EXPERT_IDS) + parser.add_argument("--num-experts-per-token", type=int, default=NUM_EXPERTS_PER_TOKEN) + parser.add_argument( + "--dtype", + choices=["auto", "float32", "float16", "bfloat16"], + default="auto", + help="torch_dtype used by from_pretrained.", + ) + parser.add_argument("--skip-processor", action="store_true", help="Skip loading processor.") + parser.add_argument("--skip-tokenizer", action="store_true", help="Skip loading tokenizer.") + return parser.parse_args() + + +def _dtype_from_arg(dtype_arg: str): + if dtype_arg == "auto": + return None + if dtype_arg == "float16": + return torch.float16 + if dtype_arg == "bfloat16": + return torch.bfloat16 + return torch.float32 + + +def main(): + args = parse_args() + _ensure_torch_fx_import_compatibility() + + config = _prepare_config(args.model_path) + kimi_cls = get_class_from_dynamic_module( + "modeling_kimi_k25.KimiK25ForConditionalGeneration", + str(args.model_path), + ) + _patch_kimi_tie_weights_compat(kimi_cls) + _patch_deepseek_init_weights_compat(kimi_cls) + + kimi_module = sys.modules[kimi_cls.__module__] + if hasattr(kimi_module, "MoonViT3dEncoder"): + kimi_module.MoonViT3dEncoder.use_deterministic_attn = False + + model_kwargs = { + "config": config, + "trust_remote_code": True, + "attn_implementation": "eager", + "local_files_only": args.local_files_only, + } + dtype = _dtype_from_arg(args.dtype) + if dtype is not None: + model_kwargs["torch_dtype"] = dtype + + if args.full_model: + model = kimi_cls.from_pretrained(str(args.model_path), **model_kwargs) + elif args.num_vision_layers is not None and args.num_text_layers is not None: + model = _load_layer_subset_model( + model_path=args.model_path, + kimi_cls=kimi_cls, + config=config, + num_vision_layers=args.num_vision_layers, + num_text_layers=args.num_text_layers, + loaded_expert_ids=args.expert_ids, + num_experts_per_tok=args.num_experts_per_token, + local_files_only=args.local_files_only, + dtype=dtype, + ) + print( + "Loaded layer subset: " + f"vision={model.config.vision_config.vt_num_hidden_layers}, " + f"text={model.config.text_config.num_hidden_layers}, " + f"experts={model.config.text_config.n_routed_experts}" + ) + else: + raise ValueError("Pass both --num-vision-layers and --num-text-layers to load a layer subset.") + print(f"Loaded model: {type(model).__name__}") + + if not args.skip_tokenizer: + tokenizer = AutoTokenizer.from_pretrained( + str(args.model_path), trust_remote_code=True, local_files_only=args.local_files_only + ) + print(f"Tokenizer vocab size: {tokenizer.vocab_size}") + + if not args.skip_processor: + processor = AutoProcessor.from_pretrained( + str(args.model_path), trust_remote_code=True, local_files_only=args.local_files_only + ) + print(f"Processor type: {type(processor).__name__}") + + +if __name__ == "__main__": + main() diff --git a/examples/kimi_k2/export_kimi_k25_vision_4.57.3.py b/examples/kimi_k2/export_kimi_k25_vision_4.57.3.py new file mode 100644 index 0000000000..bc04b4b294 --- /dev/null +++ b/examples/kimi_k2/export_kimi_k25_vision_4.57.3.py @@ -0,0 +1,423 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import argparse +import copy +import json +import re +import sys +import tempfile +from collections import defaultdict +from io import BytesIO +from pathlib import Path + +import requests +from PIL import Image +from safetensors import safe_open +from safetensors.torch import save_file +from transformers import AutoConfig, AutoProcessor, AutoTokenizer, TextStreamer +from transformers.dynamic_module_utils import get_class_from_dynamic_module + +from QEfficient import QEFFAutoModelForImageTextToText + +MODEL_PATH = Path( + "/home/huggingface_hub/models--moonshotai--Kimi-K2.5/snapshots/4d01dfe0332d63057c186e0b262165819efb6611" +) +NUM_VISION_LAYERS = 2 +NUM_TEXT_LAYERS = 2 +LOADED_EXPERT_IDS = (0, 1, 2, 3) +NUM_EXPERTS_PER_TOKEN = 2 + +EXPERT_KEY_PATTERN = re.compile(r"^(language_model\.model\.layers\.\d+\.mlp\.experts\.)(\d+)(\..+)$") + + +def _parse_expert_ids(value: str): + expert_ids = tuple(int(expert_id) for expert_id in value.split(",") if expert_id.strip()) + if not expert_ids: + raise argparse.ArgumentTypeError("At least one expert id must be provided.") + return expert_ids + + +def _validate_kimi_k25_config(kimi_config): + if kimi_config.model_type != "kimi_k25": + raise ValueError(f"This script only supports Kimi K2.5 config, got {kimi_config.model_type!r}.") + architectures = getattr(kimi_config, "architectures", None) or [] + if "KimiK25ForConditionalGeneration" not in architectures: + raise ValueError(f"Expected KimiK25ForConditionalGeneration, got architectures={architectures!r}.") + + +def _validate_layer_count(name, requested_count, available_count): + if requested_count < 1: + raise ValueError(f"{name} must be >= 1, got {requested_count}.") + if requested_count > available_count: + raise ValueError(f"{name}={requested_count} exceeds available layers={available_count}.") + + +def _validate_expert_subset(loaded_expert_ids, num_experts_per_tok, total_experts): + expert_ids = tuple(loaded_expert_ids) + if len(expert_ids) != 4: + raise ValueError(f"Expected exactly 4 routed experts, got {expert_ids!r}.") + if len(set(expert_ids)) != len(expert_ids): + raise ValueError(f"Expert ids must be unique, got {expert_ids!r}.") + invalid_ids = [expert_id for expert_id in expert_ids if expert_id < 0 or expert_id >= total_experts] + if invalid_ids: + raise ValueError(f"Expert ids {invalid_ids!r} are outside the valid range [0, {total_experts - 1}].") + if num_experts_per_tok > len(expert_ids): + raise ValueError(f"num_experts_per_tok={num_experts_per_tok} cannot exceed {len(expert_ids)} loaded experts.") + return expert_ids + + +def _extend_quantization_ignore(config, patterns): + quantization_config = getattr(config, "quantization_config", None) + if not quantization_config: + return + + ignored_modules = quantization_config.setdefault("ignore", []) + for pattern in patterns: + if pattern not in ignored_modules: + ignored_modules.append(pattern) + + +def _remap_checkpoint_key(checkpoint_key, expert_index_map): + match = EXPERT_KEY_PATTERN.match(checkpoint_key) + if not match: + return checkpoint_key + + original_expert_idx = int(match.group(2)) + remapped_expert_idx = expert_index_map.get(original_expert_idx) + if remapped_expert_idx is None: + return None + return f"{match.group(1)}{remapped_expert_idx}{match.group(3)}" + + +def _is_routed_gate_weight(checkpoint_key): + return checkpoint_key.endswith(".mlp.gate.weight") + + +def _is_routed_gate_bias(checkpoint_key): + return checkpoint_key.endswith(".mlp.gate.e_score_correction_bias") + + +def _materialize_subset_checkpoint( + model_path: Path, + temp_model_path: Path, + weight_map, + allowed_prefixes, + loaded_expert_ids, +): + expert_index_map = {expert_id: remapped_idx for remapped_idx, expert_id in enumerate(loaded_expert_ids)} + selected_by_shard = defaultdict(list) + + for checkpoint_key, shard_name in weight_map.items(): + if not any(checkpoint_key.startswith(prefix) for prefix in allowed_prefixes): + continue + + remapped_key = _remap_checkpoint_key(checkpoint_key, expert_index_map) + if remapped_key is None: + continue + selected_by_shard[shard_name].append((checkpoint_key, remapped_key)) + + if not selected_by_shard: + raise RuntimeError("No multimodal weights were selected from the Kimi K2.5 checkpoint.") + + filtered_weight_map = {} + subset_shards = [] + for shard_idx, (source_shard_name, shard_entries) in enumerate(sorted(selected_by_shard.items())): + tensors = {} + with safe_open(model_path / source_shard_name, framework="pt", device="cpu") as shard_reader: + for checkpoint_key, remapped_key in shard_entries: + tensor = shard_reader.get_tensor(checkpoint_key) + if _is_routed_gate_weight(checkpoint_key): + tensor = tensor[list(loaded_expert_ids), :].contiguous() + elif _is_routed_gate_bias(checkpoint_key): + tensor = tensor[list(loaded_expert_ids)].contiguous() + tensors[remapped_key] = tensor + + subset_shard_name = f"model-subset-{shard_idx:05d}.safetensors" + save_file(tensors, str(temp_model_path / subset_shard_name)) + subset_shards.append(subset_shard_name) + filtered_weight_map.update({remapped_key: subset_shard_name for _, remapped_key in shard_entries}) + + return filtered_weight_map, subset_shards + + +def _allowed_prefixes(num_vision_layers: int, num_text_layers: int): + prefixes = [ + "vision_tower.patch_embed.", + "vision_tower.encoder.final_layernorm.", + "mm_projector.", + "language_model.model.embed_tokens.", + "language_model.model.norm.", + "language_model.lm_head.", + ] + prefixes.extend(f"vision_tower.encoder.blocks.{layer_idx}." for layer_idx in range(num_vision_layers)) + prefixes.extend(f"language_model.model.layers.{layer_idx}." for layer_idx in range(num_text_layers)) + return prefixes + + +def load_kimi_k25_vision_moe_strip( + model_path: Path, + num_vision_layers: int, + num_text_layers: int, + loaded_expert_ids=LOADED_EXPERT_IDS, + num_experts_per_tok: int = NUM_EXPERTS_PER_TOKEN, +): + kimi_config = AutoConfig.from_pretrained(str(model_path), trust_remote_code=True) + _validate_kimi_k25_config(kimi_config) + + stripped_config = copy.deepcopy(kimi_config) + text_config = stripped_config.text_config + vision_config = stripped_config.vision_config + stripped_config._attn_implementation = "eager" + text_config._attn_implementation = "eager" + vision_config._attn_implementation = "eager" + quantization_ignores = [ + "re:vision_tower.*", + "re:mm_projector.*", + "re:language_model.lm_head.*", + ] + _extend_quantization_ignore(stripped_config, quantization_ignores) + _extend_quantization_ignore(text_config, quantization_ignores) + + _validate_layer_count("num_text_layers", num_text_layers, text_config.num_hidden_layers) + _validate_layer_count("num_vision_layers", num_vision_layers, vision_config.vt_num_hidden_layers) + + text_config.num_hidden_layers = num_text_layers + vision_config.vt_num_hidden_layers = num_vision_layers + loaded_expert_ids = _validate_expert_subset( + loaded_expert_ids, + num_experts_per_tok, + text_config.n_routed_experts, + ) + text_config.n_routed_experts = len(loaded_expert_ids) + text_config.num_experts_per_tok = num_experts_per_tok + text_config.n_group = 1 + text_config.topk_group = 1 + + kimi_cls = get_class_from_dynamic_module("modeling_kimi_k25.KimiK25ForConditionalGeneration", str(model_path)) + kimi_module = sys.modules[kimi_cls.__module__] + kimi_module.MoonViT3dEncoder.use_deterministic_attn = False + + checkpoint_index = json.loads((model_path / "model.safetensors.index.json").read_text()) + weight_map = checkpoint_index["weight_map"] + + with tempfile.TemporaryDirectory() as tmpdir: + temp_model_path = Path(tmpdir) + filtered_weight_map, subset_shards = _materialize_subset_checkpoint( + model_path=model_path, + temp_model_path=temp_model_path, + weight_map=weight_map, + allowed_prefixes=_allowed_prefixes(num_vision_layers, num_text_layers), + loaded_expert_ids=loaded_expert_ids, + ) + (temp_model_path / "config.json").write_text(stripped_config.to_json_string(use_diff=False)) + (temp_model_path / "model.safetensors.index.json").write_text( + json.dumps( + { + "metadata": { + "total_size": sum((temp_model_path / shard_name).stat().st_size for shard_name in subset_shards) + }, + "weight_map": filtered_weight_map, + } + ) + ) + + original_base_model_prefix = kimi_cls.base_model_prefix + kimi_cls.base_model_prefix = "" + try: + model, loading_info = kimi_cls.from_pretrained( + str(temp_model_path), + config=stripped_config, + local_files_only=True, + attn_implementation="eager", + output_loading_info=True, + ) + finally: + kimi_cls.base_model_prefix = original_base_model_prefix + + unexpected_keys = loading_info["unexpected_keys"] + missing_keys = loading_info["missing_keys"] + mismatched_keys = loading_info["mismatched_keys"] + if unexpected_keys or missing_keys or mismatched_keys: + raise RuntimeError( + "Failed to load the stripped Kimi K2.5 multimodal checkpoint slice cleanly. " + f"missing={missing_keys}, unexpected={unexpected_keys}, mismatched={mismatched_keys}" + ) + + model.eval() + tokenizer = AutoTokenizer.from_pretrained(str(model_path), trust_remote_code=True) + processor = AutoProcessor.from_pretrained(str(model_path), trust_remote_code=True) + return model, tokenizer, processor + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Load a stripped Kimi K2.5 multimodal checkpoint with configurable vision/text layers and 4 routed experts." + ) + parser.add_argument("--model-path", type=Path, default=MODEL_PATH) + parser.add_argument("--num-vision-layers", type=int, default=NUM_VISION_LAYERS) + parser.add_argument("--num-text-layers", type=int, default=NUM_TEXT_LAYERS) + parser.add_argument("--expert-ids", type=_parse_expert_ids, default=LOADED_EXPERT_IDS) + parser.add_argument("--num-experts-per-token", type=int, default=NUM_EXPERTS_PER_TOKEN) + parser.add_argument( + "--image-url", + type=str, + default="https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png", + ) + parser.add_argument("--prompt", type=str, default="Describe this image.") + parser.add_argument("--test", action="store_true", help="Validate ONNX output matches PyTorch image-only forward.") + return parser.parse_args() + + +def main(): + args = parse_args() + model, tokenizer, processor = load_kimi_k25_vision_moe_strip( + args.model_path, + args.num_vision_layers, + args.num_text_layers, + loaded_expert_ids=args.expert_ids, + num_experts_per_tok=args.num_experts_per_token, + ) + print( + f"Loaded {type(model).__name__} with " + f"{model.config.vision_config.vt_num_hidden_layers} vision layers, " + f"{model.config.text_config.num_hidden_layers} text layers, " + f"{model.config.text_config.n_routed_experts} routed experts." + ) + print(f"Tokenizer vocab size: {tokenizer.vocab_size}") + print(f"Processor type: {type(processor).__name__}") + + mla_absorption = {"cache_compressed": True, "absorption": False, "online": False} + qaic_config = { + "mla_absorption": mla_absorption + } # , "enable_blocking": True, "blocking_mode": "par", "par_num_split": 4, "num_kv_blocks": 8} + breakpoint() + qeff_model = QEFFAutoModelForImageTextToText(model) # , qaic_config=qaic_config) + breakpoint() + + skip_vision = True + + if skip_vision: + ## TEXT-ONLY MODE ## + + ## STEP 3: Compile Model for Text-Only Execution + # Set skip_vision=True to bypass image processing + qeff_model.compile( + qaic_config=qaic_config, + prefill_seq_len=32, + ctx_len=1024, + num_cores=16, + num_devices=2, + mxfp6_matmul=False, + mxint8_kv_cache=False, + aic_enable_depth_first=False, + skip_vision=True, # Skip vision encoder for text-only inference + mos=1, + num_patches=2400, # num_patches + h=30, # h + w=80, # w + num_image_tokens=600, # num_image_tokens + ) + breakpoint() + ## STEP 4: Prepare Text-Only Input + # Create a text-only message without any image + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Tell me about yourself."}, + ], + }, + ] + + ## STEP 5: Process Input with Chat Template + inputs = processor.apply_chat_template( + messages, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors="pt", + ) + + ## STEP 6: Run Text-Only Inference + streamer = TextStreamer(tokenizer) + output = qeff_model.generate(inputs=inputs, device_ids=[0, 1], generation_len=10) + + ## STEP 7: Display Results + print(output.generated_ids) + print(tokenizer.batch_decode(output.generated_ids)) + print(output) + + else: + ## VISION + TEXT MODE ## + + ## STEP 3: Compile Model for Vision+Text Execution + # Do not set skip_vision (defaults to False) to enable image processing + qeff_model.compile( + qaic_config=qaic_config, + prefill_seq_len=32, + ctx_len=1024, + num_cores=16, + num_devices=2, + mxfp6_matmul=False, + mxint8_kv_cache=False, + aic_enable_depth_first=False, + skip_vision=True, # Skip vision encoder for text-only inference + mos=1, + num_patches=2400, # num_patches + h=30, # h + w=80, # w + num_image_tokens=600, # num_image_tokens + ) + + breakpoint() + ## STEP 4: Prepare Image and Text Input + # Define the image URL to process + image_url = "https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png" + image = Image.open(BytesIO(requests.get(args.image_url).content)).convert("RGB") + + # Create a message with both image and text + messages = [ + { + "role": "user", + "content": [ + {"type": "image", "url": image}, + {"type": "text", "text": "Can you describe the image in detail."}, + ], + }, + ] + + ## STEP 5: Process Input with Chat Template + """inputs = processor.apply_chat_template( + messages, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors="pt", + ) + """ + inputs = processor( + messages=messages, + add_generation_prompt=True, + tokenize=False, + return_tensors="pt", + ) + # Convert pixel values to float32 for processing + inputs["pixel_values"] = inputs["pixel_values"].to(qeff_model.model.config.torch_dtype) + + ## STEP 6: Run Vision+Text Inference + streamer = TextStreamer(tokenizer) + output = qeff_model.generate(inputs=inputs, device_ids=[0, 1], generation_len=100) + + ## STEP 7: Display Results + print(output.generated_ids) + print(tokenizer.batch_decode(output.generated_ids)) + print(output) + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index 5255ffc3d0..ee4ca672b9 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,7 +19,7 @@ classifiers = [ ] requires-python = ">=3.8,<3.13" dependencies = [ - "transformers==5.8", + "transformers==5.5.4", "diffusers==0.38.0", "huggingface-hub==1.7.1", "hf_transfer==0.1.9", From d48d8b1acfd08fa0906614518a259f8e54883045 Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Mon, 22 Jun 2026 16:05:28 +0530 Subject: [PATCH 03/33] local_changes Signed-off-by: Mamta Singh --- QEfficient/base/modeling_qeff.py | 2 + .../models/deepseek_v3/modeling_deepseek.py | 31 ++++- .../models/kimi_k25/modeling_kimi_k25.py | 108 +++++++-------- examples/kimi_k2/export_kimi_k25_vision.py | 129 ++++++++++++++++++ .../kimi_k2/export_kimi_k25_vision_4.57.3.py | 1 + 5 files changed, 204 insertions(+), 67 deletions(-) diff --git a/QEfficient/base/modeling_qeff.py b/QEfficient/base/modeling_qeff.py index 28e2780aff..f998d70fc9 100755 --- a/QEfficient/base/modeling_qeff.py +++ b/QEfficient/base/modeling_qeff.py @@ -957,6 +957,8 @@ def _compile( command.append(f"-mdp-load-partition-config={mdp_ts_json_path}") for key, value in compiler_options.items(): + if value is None: + continue option = "-" + key.replace("_", "-") if isinstance(value, bool): if value: diff --git a/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py b/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py index 3dad27103a..df07871d45 100644 --- a/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py +++ b/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py @@ -26,6 +26,23 @@ from QEfficient.utils.constants import MAX_POSITION_EMBEDDINGS, MIN_MASKED_ATTENTION_VALUE +def _get_linear_weight(linear_module: nn.Module) -> torch.Tensor: + """Return a linear layer weight, decompressing compressed-tensors modules if needed.""" + weight = getattr(linear_module, "weight", None) + if weight is not None: + return weight + + if hasattr(linear_module, "weight_packed") and hasattr(linear_module, "weight_scale"): + from compressed_tensors.compressors.base import decompress_module as ct_decompress_module + + ct_decompress_module(linear_module) + weight = getattr(linear_module, "weight", None) + if weight is not None: + return weight + + raise AttributeError(f"{linear_module.__class__.__name__!s} object has no attribute 'weight'") + + def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] @@ -768,23 +785,25 @@ class QEffDeepseekV3MoE(nn.Module): def __qeff_init__( self, ): + #breakpoint() self.all_gate_proj = torch.nn.Parameter( torch.cat( - [exp.gate_proj.compressor.decompress_module(exp.gate_proj).T.unsqueeze(0) for exp in self.experts], + [_get_linear_weight(exp.gate_proj).T.unsqueeze(0) for exp in self.experts], dim=0, ) ) self.all_up_proj = torch.nn.Parameter( torch.cat( - [exp.up_proj.compressor.decompress_module(exp.up_proj).T.unsqueeze(0) for exp in self.experts], dim=0 + [_get_linear_weight(exp.up_proj).T.unsqueeze(0) for exp in self.experts], ) ) self.all_down_proj = torch.nn.Parameter( torch.cat( - [exp.down_proj.compressor.decompress_module(exp.down_proj).T.unsqueeze(0) for exp in self.experts], + [_get_linear_weight(exp.down_proj).T.unsqueeze(0) for exp in self.experts], dim=0, ) ) + self.act_fn = self.experts[0].act_fn def moe( @@ -833,9 +852,9 @@ def __qeff_init__( up_proj = torch.nn.Linear(self.config.hidden_size, self.config.moe_intermediate_size, bias=False) down_proj = torch.nn.Linear(self.config.moe_intermediate_size, self.config.hidden_size, bias=False) - gate_proj.weight = torch.nn.Parameter(exp.gate_proj.compressor.decompress_module(exp.gate_proj)) - up_proj.weight = torch.nn.Parameter(exp.up_proj.compressor.decompress_module(exp.up_proj)) - down_proj.weight = torch.nn.Parameter(exp.down_proj.compressor.decompress_module(exp.down_proj)) + gate_proj.weight = torch.nn.Parameter(_get_linear_weight(exp.gate_proj).detach().clone()) + up_proj.weight = torch.nn.Parameter(_get_linear_weight(exp.up_proj).detach().clone()) + down_proj.weight = torch.nn.Parameter(_get_linear_weight(exp.down_proj).detach().clone()) setattr(exp, "gate_proj", gate_proj) setattr(exp, "up_proj", up_proj) diff --git a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py index 63dfdc69f6..bf828ea7b7 100644 --- a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py +++ b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py @@ -74,16 +74,11 @@ def wrapper(org, interpolation_mode, shape): @get_rope_shape_decorate @torch.compile(dynamic=True) # Remove this def get_rope_shape(org, interpolation_mode, shape): - return ( - F.interpolate( - org.permute((2, 0, 1)).unsqueeze(0), - size=shape, - mode=interpolation_mode, - ) - .squeeze(0) - .permute((1, 2, 0)) - .flatten(end_dim=1) - ) + return (F.interpolate( + org.permute((2, 0, 1)).unsqueeze(0), + size=shape, + mode=interpolation_mode, + ).squeeze(0).permute((1, 2, 0)).flatten(end_dim=1)) def apply_rope(xq, xk, freqs_cis): @@ -183,7 +178,6 @@ def forward(self, x: torch.Tensor, grid_thws: torch.Tensor) -> torch.Tensor: interpolation_mode=self.interpolation_mode, shape=(h, w), ) - if t == 1: pos_emb_3d = pos_emb_2d else: @@ -716,7 +710,7 @@ def forward_only_image(self, pixel_values: torch.Tensor, grid_thws: torch.Tensor image_features = self.mm_projector(image_features) return image_features - def forward_only_image_for_export( + def forward( self, pixel_values: torch.Tensor, h_shape: torch.Tensor, w_shape: torch.Tensor ) -> torch.Tensor: """ @@ -745,14 +739,14 @@ def forward_only_image_for_export( h = h_shape.shape[0] w = w_shape.shape[0] - target_dtype = self.vision_tower.patch_embed.proj.weight.dtype + target_dtype = self.model.vision_tower.patch_embed.proj.weight.dtype pixel_values = pixel_values.to(target_dtype) # --- Patch embedding --- - x = self.vision_tower.patch_embed.proj(pixel_values).view(pixel_values.size(0), -1) + x = self.model.vision_tower.patch_embed.proj(pixel_values).view(pixel_values.size(0), -1) # Positional embedding (single image, t=1) - pos_emb_module = self.vision_tower.patch_embed.pos_emb + pos_emb_module = self.model.vision_tower.patch_embed.pos_emb pos_emb_2d = _get_rope_shape( pos_emb_module.weight, interpolation_mode=pos_emb_module.interpolation_mode, @@ -769,18 +763,18 @@ def forward_only_image_for_export( # RoPE frequencies for single image (stacked real/imag to avoid complex tensors in ONNX) # Shape: (2, num_tokens, dim//2) where [0]=cos, [1]=sin - rope_2d = self.vision_tower.encoder.rope_2d + rope_2d = self.model.vision_tower.encoder.rope_2d freqs_cos = rope_2d.freqs_cos[:h, :w].reshape(-1, rope_2d.dim // 2) freqs_sin = rope_2d.freqs_sin[:h, :w].reshape(-1, rope_2d.dim // 2) freqs_cis = torch.stack([freqs_cos, freqs_sin], dim=0) - for block in self.vision_tower.encoder.blocks: + for block in self.model.vision_tower.encoder.blocks: x = block(x, cu_seqlens, max_seqlen, rope_freqs_cis=freqs_cis) - x = self.vision_tower.encoder.final_layernorm(x) + x = self.model.vision_tower.encoder.final_layernorm(x) # --- tpool_patch_merger (single image, t=1) --- - merge_kernel_size = self.vision_tower.merge_kernel_size + merge_kernel_size = self.model.vision_tower.merge_kernel_size kernel_height, kernel_width = merge_kernel_size d_model = x.size(-1) new_height = h // kernel_height @@ -790,7 +784,7 @@ def forward_only_image_for_export( merged = reshaped.view(new_height * new_width, kernel_height * kernel_width, -1) # --- mm_projector (PatchMergerMLP on single tensor) --- - image_embeds = self.mm_projector.proj(self.mm_projector.pre_norm(merged).view(merged.shape[0], -1)) + image_embeds = self.model.mm_projector.proj(self.model.mm_projector.pre_norm(merged).view(merged.shape[0], -1)) return image_embeds @@ -799,7 +793,7 @@ def __init__(self, model): super().__init__() self.model = model self.language_model = self.model.language_model - # self.language_model = QEffDeepseekV3ForCausalLM#(config.text_config) + self.lm_head = self.language_model.lm_head self.config = self.model.config def get_submodules_for_export(self) -> Type[nn.Module]: @@ -814,41 +808,50 @@ def get_submodules_for_export(self) -> Type[nn.Module]: def forward( self, input_ids: torch.LongTensor = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + image_embeds: Optional[torch.FloatTensor] = None, + image_idx: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, compressed_kvs: Optional[List[torch.FloatTensor]] = None, past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, batch_index: Optional[torch.LongTensor] = None, - image_embeds: Optional[List[torch.FloatTensor]] = None, - # inputs_embeds: Optional[torch.FloatTensor] = None, - labels: Optional[torch.LongTensor] = None, + comp_ctx_lengths: Optional[List[int]] = None, use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - cache_position: Optional[torch.LongTensor] = None, - num_logits_to_keep: int = 0, **kwargs, - ) -> Union[Tuple, CausalLMOutputWithPast]: - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict + ) -> Tuple: + del labels, output_attentions, output_hidden_states, return_dict, cache_position, num_logits_to_keep + if inputs_embeds is None: + inputs_embeds = self.model.get_input_embeddings()(input_ids) + + if image_idx is None: + image_idx = torch.zeros((1, 1), dtype=torch.int64, device=inputs_embeds.device) + + if image_embeds is not None and input_ids is not None and input_ids.shape[1] != 1: + if image_embeds.dim() == 2: + image_embeds = image_embeds.unsqueeze(0) + + image_embeds = image_embeds.to(device=inputs_embeds.device, dtype=inputs_embeds.dtype) + _, _, hidden_size = inputs_embeds.shape + selected = input_ids == self.config.media_placeholder_token_id + indices1 = selected.to(torch.int64).cumsum(1) - 1 + indices1 = torch.where(indices1 != -1, indices1 + image_idx.to(indices1.device), indices1) + indices0 = torch.arange(selected.shape[0], device=selected.device).view(-1, 1) + safe_indices1 = torch.where(indices1 < 0, torch.zeros_like(indices1), indices1) + image_features_expanded = image_embeds.reshape(-1, hidden_size).unsqueeze(0)[indices0, safe_indices1] + image_input_embeds = torch.where(selected.unsqueeze(-1), image_features_expanded, inputs_embeds) + inputs_embeds = torch.where(input_ids.shape[1] == torch.tensor(1), inputs_embeds, image_input_embeds) + image_idx = (indices1.max() + 1).reshape(1, 1) outputs = self.model.language_model( - input_ids=input_ids, + inputs_embeds=inputs_embeds, attention_mask=attention_mask, position_ids=position_ids, compressed_kvs=compressed_kvs, past_key_values=past_key_values, batch_index=batch_index, - inputs_embeds=image_embeds, # inputs_embeds, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - cache_position=cache_position, + comp_ctx_lengths=comp_ctx_lengths, + use_cache=True if use_cache is None else use_cache, **kwargs, ) @@ -858,25 +861,8 @@ def forward( logits = self.lm_head(hidden_states).float() loss = None - if labels is not None: - shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - loss_fct = nn.CrossEntropyLoss() - shift_logits = shift_logits.view(-1, self.config.vocab_size) - shift_labels = shift_labels.view(-1).to(shift_logits.device) - loss = loss_fct(shift_logits, shift_labels) - - if not return_dict: - output = (logits,) + outputs[1:] - return (loss,) + output if loss is not None else output - - return CausalLMOutputWithPast( - loss=loss, - logits=logits, - past_key_values=outputs.past_key_values, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) + output = (logits,) + outputs[1:] + return logits, image_embeds, image_idx, outputs.compressed_kvs, outputs.past_key_values # ref https://github.com/huggingface/transformers/blob/78b2929c0554b79e0489b451ce4ece14d265ead2/src/transformers/models/llava/modeling_llava.py#L240 diff --git a/examples/kimi_k2/export_kimi_k25_vision.py b/examples/kimi_k2/export_kimi_k25_vision.py index ba34938ed4..4d2c77fbf5 100644 --- a/examples/kimi_k2/export_kimi_k25_vision.py +++ b/examples/kimi_k2/export_kimi_k25_vision.py @@ -21,6 +21,8 @@ from transformers.dynamic_module_utils import get_class_from_dynamic_module from transformers.utils import import_utils as hf_import_utils +from QEfficient import QEFFAutoModelForImageTextToText + MODEL_PATH = Path( "/home/huggingface_hub/models--moonshotai--Kimi-K2.5/snapshots/4d01dfe0332d63057c186e0b262165819efb6611" @@ -405,6 +407,133 @@ def main(): str(args.model_path), trust_remote_code=True, local_files_only=args.local_files_only ) print(f"Processor type: {type(processor).__name__}") + + mla_absorption = {"cache_compressed": True, "absorption": False, "online": False} + qaic_config = { + "mla_absorption": mla_absorption + } # , "enable_blocking": True, "blocking_mode": "par", "par_num_split": 4, "num_kv_blocks": 8} + + qeff_model = QEFFAutoModelForImageTextToText(model) # , qaic_config=qaic_config) + + skip_vision = False + + if skip_vision: + ## TEXT-ONLY MODE ## + + ## STEP 3: Compile Model for Text-Only Execution + # Set skip_vision=True to bypass image processing + qeff_model.compile( + qaic_config=qaic_config, + prefill_seq_len=32, + ctx_len=1024, + num_cores=16, + num_devices=2, + mxfp6_matmul=False, + mxint8_kv_cache=False, + aic_enable_depth_first=False, + skip_vision=True, # Skip vision encoder for text-only inference + mos=1, + num_patches=2400, # num_patches + h=30, # h + w=80, # w + num_image_tokens=600, # num_image_tokens + ) + breakpoint() + ## STEP 4: Prepare Text-Only Input + # Create a text-only message without any image + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Tell me about yourself."}, + ], + }, + ] + + ## STEP 5: Process Input with Chat Template + inputs = processor.apply_chat_template( + messages, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors="pt", + ) + + ## STEP 6: Run Text-Only Inference + streamer = TextStreamer(tokenizer) + output = qeff_model.generate(inputs=inputs, device_ids=[0, 1], generation_len=10) + + ## STEP 7: Display Results + print(output.generated_ids) + print(tokenizer.batch_decode(output.generated_ids)) + print(output) + + else: + ## VISION + TEXT MODE ## + + ## STEP 3: Compile Model for Vision+Text Execution + # Do not set skip_vision (defaults to False) to enable image processing + qeff_model.compile( + qaic_config=qaic_config, + prefill_seq_len=1, + ctx_len=1024, + num_cores=16, + num_devices=2, + mxfp6_matmul=False, + mxint8_kv_cache=False, + aic_enable_depth_first=False, + #skip_vision=True, # Skip vision encoder for text-only inference + mos=1, + num_patches=2400, # num_patches + h=30, # h + w=80, # w + num_image_tokens=600, # num_image_tokens + ) + + breakpoint() + ## STEP 4: Prepare Image and Text Input + # Define the image URL to process + image_url = "https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png" + image = Image.open(BytesIO(requests.get(args.image_url).content)).convert("RGB") + + # Create a message with both image and text + messages = [ + { + "role": "user", + "content": [ + {"type": "image", "url": image}, + {"type": "text", "text": "Can you describe the image in detail."}, + ], + }, + ] + + ## STEP 5: Process Input with Chat Template + """inputs = processor.apply_chat_template( + messages, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors="pt", + ) + """ + inputs = processor( + messages=messages, + add_generation_prompt=True, + tokenize=False, + return_tensors="pt", + ) + # Convert pixel values to float32 for processing + inputs["pixel_values"] = inputs["pixel_values"].to(qeff_model.model.config.torch_dtype) + + ## STEP 6: Run Vision+Text Inference + streamer = TextStreamer(tokenizer) + output = qeff_model.generate(inputs=inputs, device_ids=[0, 1], generation_len=100) + + ## STEP 7: Display Results + print(output.generated_ids) + print(tokenizer.batch_decode(output.generated_ids)) + print(output) + if __name__ == "__main__": diff --git a/examples/kimi_k2/export_kimi_k25_vision_4.57.3.py b/examples/kimi_k2/export_kimi_k25_vision_4.57.3.py index bc04b4b294..b55c77f4d3 100644 --- a/examples/kimi_k2/export_kimi_k25_vision_4.57.3.py +++ b/examples/kimi_k2/export_kimi_k25_vision_4.57.3.py @@ -24,6 +24,7 @@ from QEfficient import QEFFAutoModelForImageTextToText + MODEL_PATH = Path( "/home/huggingface_hub/models--moonshotai--Kimi-K2.5/snapshots/4d01dfe0332d63057c186e0b262165819efb6611" ) From 5b1d648cd079bce3057e6b1fc5e72c045bdb748a Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Mon, 22 Jun 2026 17:43:41 +0530 Subject: [PATCH 04/33] vision export Signed-off-by: Mamta Singh --- .../models/kimi_k25/modeling_kimi_k25.py | 3103 +++++++++-------- .../transformers/models/modeling_auto.py | 2 + examples/kimi_k2/export_kimi_k25_vision.py | 6 + 3 files changed, 1583 insertions(+), 1528 deletions(-) diff --git a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py index bf828ea7b7..d7ab06acdf 100644 --- a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py +++ b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py @@ -1,1528 +1,1575 @@ -# ----------------------------------------------------------------------------- -# -# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. -# SPDX-License-Identifier: BSD-3-Clause -# -# ---------------------------------------------------------------------------- - -import math -import sys as _sys -from collections.abc import Sequence -from copy import deepcopy -from io import BytesIO -from pathlib import Path - -# from QEfficient import QEFFAutoModelForImageTextToText -from typing import List, Optional, Tuple, Type, Union - -import numpy as np -import requests -import torch -import torch.nn as nn -import torch.nn.functional as F -from PIL import Image -from transformers import AutoProcessor, activations - -from QEfficient.transformers.models.deepseek_v3.modeling_deepseek import QEffDeepseekV3ForCausalLM - -try: - from transformers.activations import PytorchGELUTanh -except ImportError: - from transformers.activations import GELUTanh - - activations.PytorchGELUTanh = GELUTanh - PytorchGELUTanh = GELUTanh -from transformers.activations import PytorchGELUTanh -from transformers.cache_utils import Cache -from transformers.configuration_utils import PretrainedConfig -from transformers.modeling_outputs import CausalLMOutputWithPast -from transformers.modeling_utils import PreTrainedModel -from transformers.models.llava.modeling_llava import LlavaCausalLMOutputWithPast - -from QEfficient.utils import constants - -from .configuration_kimi_k25 import KimiK25Config - - -def eager_attention_forward(q, k, v, **kwargs): - q = q.transpose(0, 1) # (num_heads, seq_len, head_dim) - k = k.transpose(0, 1) - v = v.transpose(0, 1) - attn_weight = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1]) - attn_weight = torch.softmax(attn_weight, dim=-1, dtype=torch.float32).to(q.dtype) - attn_out = attn_weight @ v - attn_out = attn_out.transpose(0, 1) # (seq_len, num_heads, head_dim) - return attn_out.reshape(attn_out.shape[0], -1) - - -VL_VISION_ATTENTION_FUNCTIONS = {"eager": eager_attention_forward} - - -def get_rope_shape_decorate(func): - _get_rope_shape_first_call_flag = set() - - def wrapper(org, interpolation_mode, shape): - key = (org.requires_grad, torch.is_grad_enabled(), interpolation_mode) - if key not in _get_rope_shape_first_call_flag: - _get_rope_shape_first_call_flag.add(key) - _ = func(org, interpolation_mode, shape=(64, 64)) - return func(org, interpolation_mode, shape) - - return wrapper - - -@get_rope_shape_decorate -@torch.compile(dynamic=True) # Remove this -def get_rope_shape(org, interpolation_mode, shape): - return (F.interpolate( - org.permute((2, 0, 1)).unsqueeze(0), - size=shape, - mode=interpolation_mode, - ).squeeze(0).permute((1, 2, 0)).flatten(end_dim=1)) - - -def apply_rope(xq, xk, freqs_cis): - # freqs_cis shape: (2, seq_len, dim//2) - # xq shape: (seq_len, num_heads, head_dim) - # Need freqs to broadcast: (seq_len, 1, dim//2) - freqs_cos = freqs_cis[0].unsqueeze(-2) # (seq_len, 1, dim//2) - freqs_sin = freqs_cis[1].unsqueeze(-2) # (seq_len, 1, dim//2) - xq_r = xq.float().view(*xq.shape[:-1], -1, 2) - xq_r0, xq_r1 = xq_r[..., 0], xq_r[..., 1] - xq_out_r = xq_r0 * freqs_cos - xq_r1 * freqs_sin - xq_out_i = xq_r0 * freqs_sin + xq_r1 * freqs_cos - xq_out = torch.stack([xq_out_r, xq_out_i], dim=-1).flatten(-2) - xk_r = xk.float().view(*xk.shape[:-1], -1, 2) - xk_r0, xk_r1 = xk_r[..., 0], xk_r[..., 1] - xk_out_r = xk_r0 * freqs_cos - xk_r1 * freqs_sin - xk_out_i = xk_r0 * freqs_sin + xk_r1 * freqs_cos - xk_out = torch.stack([xk_out_r, xk_out_i], dim=-1).flatten(-2) - return xq_out.type_as(xq), xk_out.type_as(xk) - - -def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): - """ - From: - https://github.com/OpenGVLab/InternVideo/blob/421f6d2361fc8f61a3394244571f2601a4e99e29/InternVideo2/multi_modality/models/backbones/internvideo2/pos_embed.py#L86 - embed_dim: output dimension for each position - pos: a list of positions to be encoded: size (M,) - out: (M, D) - """ - assert embed_dim % 2 == 0 - omega = np.arange(embed_dim // 2, dtype=np.float32) - omega /= embed_dim / 2.0 - omega = 1.0 / 10000**omega # (D/2,) - - pos = pos.reshape(-1) # (M,) - out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product - - emb_sin = np.sin(out) # (M, D/2) - emb_cos = np.cos(out) # (M, D/2) - - emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D) - return emb - - -def get_1d_sincos_pos_embed(embed_dim, t_size, cls_token=False): - """ - t_size: int of the temporal size - return: - pos_embed: [t_size, embed_dim] or [1+t_size, embed_dim] (w/ or w/o cls_token) - """ - grid_t = np.arange(t_size, dtype=np.float32) - pos_embed = get_1d_sincos_pos_embed_from_grid(embed_dim, grid_t) - if cls_token: - pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0) - return pos_embed - - -class QEffLearnable2DInterpPosEmbDivided_fixed(nn.Module): - def __qeff_init__(self): - self.interpolation_mode = "bilinear" - - """def __qeff_init__(self, - height: int, - width: int, - num_frames: int, - dim: int, - interpolation_mode: str = 'bilinear') -> None: - super().__init__() - self.height = height - self.width = width - self.num_frames = num_frames - self.dim = dim - self.interpolation_mode = interpolation_mode - self.weight = nn.Parameter(torch.empty(height, width, dim)) - self.register_buffer('time_weight', - torch.from_numpy( - get_1d_sincos_pos_embed( - self.dim, - self.num_frames)).float().unsqueeze(1), - persistent=False) - - self.reset_parameters() - """ - - def reset_parameters(self): - nn.init.normal_(self.weight) - - def forward(self, x: torch.Tensor, grid_thws: torch.Tensor) -> torch.Tensor: - pos_embs = [] - for t, h, w in grid_thws.tolist(): - assert t <= self.num_frames, f"t:{t} > self.num_frames:{self.num_frames}" - if (h, w) == self.weight.shape[:-1]: - pos_emb_2d = self.weight.flatten(end_dim=1) - else: - pos_emb_2d = get_rope_shape( - self.weight, - interpolation_mode=self.interpolation_mode, - shape=(h, w), - ) - if t == 1: - pos_emb_3d = pos_emb_2d - else: - pos_emb_3d = pos_emb_2d.unsqueeze(0).repeat(t, 1, 1) + self.time_weight[0:t] - - pos_embs.append(pos_emb_3d.reshape(-1, pos_emb_3d.shape[-1])) - - out = x + torch.cat(pos_embs) - return out - - -class MoonVision3dPatchEmbed(nn.Module): - def __init__( - self, - out_dim: int, - in_dim: int = 3, - patch_size: int | tuple[int, int] = (14, 14), - pos_emb_height: int = 14, - pos_emb_width: int = 14, - pos_emb_time: int = 4, - pos_emb_type: str = "divided_fixed", - ): - super().__init__() - assert isinstance(patch_size, int | Sequence), f"Invalid patch_size type: {type(patch_size)}" - if isinstance(patch_size, int): - patch_size = (patch_size, patch_size) - assert len(patch_size) == 2, f"Expected patch_size to be a tuple of 2, got {patch_size}" - self.patch_size = patch_size - - self.proj = nn.Conv2d(in_dim, out_dim, kernel_size=patch_size, stride=patch_size) - - if pos_emb_type == "divided_fixed": - self.pos_emb = Learnable2DInterpPosEmbDivided_fixed( - height=pos_emb_height, width=pos_emb_width, num_frames=pos_emb_time, dim=out_dim - ) - else: - raise NotImplementedError(f"Not support pos_emb_type: {pos_emb_type}") - - def forward(self, x: torch.Tensor, grid_thws: torch.Tensor) -> torch.Tensor: - """ - Args: - x (L, Channels): input tensor - grid_hws (N, 3): temporal, height and width - Returns: - (L, Cout) tensor - """ - x = self.proj(x).view(x.size(0), -1) - # apply positional embedding - x = self.pos_emb(x, grid_thws) - return x - - -class Rope2DPosEmbRepeated(nn.Module): - """2D rotary position embedding with multi-resolution support. - This class is intended to be used in the following way: - 1. Before training, create an instance of Rope2DPosEmb. This instance will hold the precomputed cis. - 2. Before each forward pass, call `get_freqs_cis_by_*` to get the `freqs_cis` tensor for this iteration. - 3. During the forward pass, pass the `freqs_cis` tensor to each attention layer, and call `apply` just before each attention operation. - The rope is shared across all attention layers and all heads. - Refs: - - RoFormer: https://arxiv.org/abs/2104.09864 - - VisionLLaMA: https://arxiv.org/abs/2403.00522 - - https://github.com/Meituan-AutoML/VisionLLaMA/blob/main/dit/models.py - Args: - dim (int): usually the multi-head attention dimension, should be divisible by 4 (TODO: relax this constraint if needed) - max_height (int): the maximum height of the 2D grid - max_width (int): the maximum width of the 2D grid - theta_base (float): the base of the theta - device (str): the device to store the precomputed cis - """ - - def __init__(self, dim: int, max_height: int, max_width: int, theta_base=10000): - super().__init__() - self.dim = dim - assert self.dim % 4 == 0, "dim must be divisible by 4" - self.max_height = max_height - self.max_width = max_width - self.theta_base = theta_base - - def extra_repr(self): - return f"dim={self.dim}, max_height={self.max_height}, max_width={self.max_width}, theta_base={self.theta_base}" - - def _precompute_freqs_cis(self, device: torch.device) -> torch.Tensor: - """Calculate the cis(freqs) for each position in the 2D grid. - Return: complex tensor of shape (max_height, max_width, dim//2) and value: - height axis: ret[h, w, 2*i] = cis(h * theta_base**(-4*i/dim)) - weight axis: ret[h, w, 2*i+1] = cis(w * theta_base**(-4*i/dim)) with (i in [0, dim//4)) - note: `cis` is a mathematical notation defined by cis x = cos x + i sin x, - """ - N = self.max_height * self.max_width - flat_pos = torch.arange(0, N).float().to(device) - x_pos = flat_pos % self.max_width - y_pos = flat_pos // self.max_width - dim_range = torch.arange(0, self.dim, 4)[: (self.dim // 4)].float().to(device) # C/4 - freqs = 1.0 / (self.theta_base ** (dim_range / self.dim)) - x_freqs = torch.outer(x_pos, freqs).float() # N, C/4 - y_freqs = torch.outer(y_pos, freqs).float() # N, C/4 - x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs) # N, C/4 - y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs) # N, C/4 - # N, C/4, 2 - freqs_cis = torch.cat([x_cis.unsqueeze(dim=-1), y_cis.unsqueeze(dim=-1)], dim=-1) - # max_height, max_width, C/2 - freqs_cis = freqs_cis.reshape(self.max_height, self.max_width, -1) - return freqs_cis - - def get_freqs_cis(self, grid_thws: torch.Tensor, device: torch.device) -> torch.Tensor: - """ - Args: - grid_thws (torch.Tensor): grid time, height and width - Returns: - freqs_cis: tensor of shape (sum(t * height * width), dim//2) - """ - if not hasattr(self, "freqs_cis"): - self.register_buffer("freqs_cis", self._precompute_freqs_cis(device), persistent=False) - - shapes = grid_thws.tolist() - assert all(1 <= h <= self.max_height and 1 <= w <= self.max_width for t, h, w in shapes), ( - shapes, - self.max_height, - self.max_width, - ) - freqs_cis = torch.cat( - [self.freqs_cis[:h, :w].reshape(-1, self.dim // 2).repeat(t, 1) for t, h, w in shapes], - dim=0, - ) - return freqs_cis - - -class MLP2(nn.Module): - """ - Args: - dims: [in_dim, hidden_dim, out_dim] - bias: whether to use bias in linear layer. - """ - - def __init__(self, dims: list[int], activation, bias=True): - super().__init__() - assert len(dims) == 3 - self.fc0 = nn.Linear(dims[0], dims[1], bias=bias) - self.fc1 = nn.Linear(dims[1], dims[2], bias=bias) - self.activation = activation - for m in [self.fc0, self.fc1]: - nn.init.trunc_normal_(m.weight, std=math.sqrt(2 / m.in_features)) - if m.bias is not None: - nn.init.zeros_(m.bias) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = self.fc0(x) - x = self.activation(x) - return self.fc1(x) - - -class MoonViTEncoderLayer(nn.Module): - def __init__( - self, - num_heads: int, - hidden_dim: int, - mlp_dim: int, - *, - attn_implementation: str = "flash_attention_2", - activation=F.gelu, - attn_bias: bool = False, - use_deterministic_attn: bool = False, - ): - super().__init__() - self.num_heads = num_heads - self.hidden_dim = hidden_dim - self.hidden_size_per_attention_head = self.hidden_dim // self.num_heads - self.attn_implementation = attn_implementation - self.use_deterministic_attn = use_deterministic_attn - - self.norm0 = nn.LayerNorm(hidden_dim) - self.norm1 = nn.LayerNorm(hidden_dim) - self.mlp = MLP2([hidden_dim, mlp_dim, hidden_dim], activation) - self.wqkv = nn.Linear(hidden_dim, hidden_dim * 3, bias=attn_bias) - self.wo = nn.Linear(hidden_dim, hidden_dim, bias=attn_bias) - - def attention_qkvpacked( - self, - x: torch.Tensor, - cu_seqlens: torch.Tensor, - max_seqlen: torch.Tensor, - rope_freqs_cis: torch.Tensor | None = None, - ): - """ - Args: - x (torch.Tensor): (batch_size, seqlen, hidden_dim) - cu_seqlens (torch.Tensor): - """ - xqkv = self.wqkv(x) - - qkv_shape = xqkv.size()[:-1] + ( - 3, - self.num_heads, - self.hidden_size_per_attention_head, - ) - # xqkv: (batch_size, seqlen, 3, nheads, headdim) - xqkv = xqkv.view(*qkv_shape) - xq, xk, xv = torch.unbind(xqkv, dim=-3) - - xq, xk = apply_rope(xq, xk, rope_freqs_cis) - - attn_func = VL_VISION_ATTENTION_FUNCTIONS[self.attn_implementation] - attn_out = attn_func( - xq, - xk, - xv, - q_cu_seqlens=cu_seqlens, - k_cu_seqlens=cu_seqlens, - max_seqlen_k=max_seqlen, - max_seqlen_q=max_seqlen, - deterministic=self.use_deterministic_attn, - ) - - attn_out = self.wo(attn_out) - return attn_out - - def forward( - self, - hidden_states: torch.Tensor, - cu_seqlens: torch.Tensor, - max_seqlen: int, - rope_freqs_cis: torch.Tensor | None = None, - ): - residual = hidden_states - hidden_states = self.norm0(hidden_states) - - hidden_states = self.attention_qkvpacked(hidden_states, cu_seqlens, max_seqlen, rope_freqs_cis) - hidden_states = residual + hidden_states - - residual = hidden_states - hidden_states = self.norm1(hidden_states) - hidden_states = self.mlp(hidden_states) - hidden_states = residual + hidden_states - - return hidden_states - - -class QEffMoonViT3dEncoder(nn.Module): - def __qeff_init__(self): - self.block_cfg = { - "num_heads": 64, # config.num_attention_heads, - "hidden_dim": 7168, # config.hidden_size, - "mlp_dim": 18432, # config.intermediate_size, - "activation": PytorchGELUTanh(), - "attn_bias": True, - "attn_implementation": "eager", # config._attn_implementation, - } - self.rope_2d = Rope2DPosEmbRepeated(self.block_cfg["hidden_dim"] // self.block_cfg["num_heads"], 512, 512) - # if not hasattr(self.rope_2d, 'freqs_cos'): - # self.rope_2d.register_buffer('freqs_cos', self.rope_2d.freqs_cis.real.contiguous(), persistent=False) - # self.rope_2d.register_buffer('freqs_sin', self.rope_2d.freqs_cis.imag.contiguous(), persistent=False) - self.blocks = nn.ModuleList( - [ - MoonViTEncoderLayer(**self.block_cfg, use_deterministic_attn=False) - for _ in range(2) # config.num_hidden_layers) - ] - ) - - def forward( - self, - hidden_states: torch.Tensor, - grid_thws: torch.Tensor, - ) -> torch.Tensor: - # if not hasattr(self.rope_2d, 'freqs_cos'): - # self.rope_2d.register_buffer('freqs_cos', self.rope_2d.freqs_cis.real.contiguous(), persistent=False) - # self.rope_2d.register_buffer('freqs_sin', self.rope_2d.freqs_cis.imag.contiguous(), persistent=False) - rope_freqs_cis = self.rope_2d.get_freqs_cis(grid_thws=grid_thws, device=hidden_states.device) - - lengths = torch.cat( - ( - torch.zeros(1, dtype=grid_thws.dtype, device=grid_thws.device), - grid_thws[:, 0] * grid_thws[:, 1] * grid_thws[:, 2], - ) - ) - - max_seqlen = lengths.max() - cu_seqlens = lengths.to(hidden_states.device).cumsum(dim=0, dtype=torch.int32) - for block in self.blocks: - hidden_states = block(hidden_states, cu_seqlens, max_seqlen, rope_freqs_cis=rope_freqs_cis) - - hidden_states = self.final_layernorm(hidden_states) - return hidden_states - - -def tpool_patch_merger( - x: torch.Tensor, - grid_thws: torch.Tensor, - merge_kernel_size: tuple[int, int] = (2, 2), -) -> list[torch.Tensor]: - d_model = x.size(-1) - - outputs = [] - pre_sum = 0 - for t, h, w in grid_thws.tolist(): - # Get the current sequence - seq = x[pre_sum : pre_sum + t * h * w] - # Reshape along self.merge_kernel_size and concat to the last dimension - kernel_height, kernel_width = merge_kernel_size - new_height, new_width = h // kernel_height, w // kernel_width - reshaped_seq = seq.view(t, new_height, kernel_height, new_width, kernel_width, d_model) - reshaped_seq = reshaped_seq.permute(0, 1, 3, 2, 4, 5).contiguous().mean(dim=0) # temporal pooling - padded_seq = reshaped_seq.view(new_height * new_width, kernel_height * kernel_width, -1) - outputs.append(padded_seq) - pre_sum += t * h * w - - return outputs - - -class MoonViT3dPretrainedModel(PreTrainedModel): - config_class = None - model_type = "moonvit3d" - _no_split_modules = ["PackingTransformer"] - _supports_flash_attn_2 = True - _supports_sdpa = True - - def __init__(self, config, *inputs, **kwargs): - super().__init__(config, *inputs, **kwargs) - config = deepcopy(config) - self.merge_kernel_size = config.merge_kernel_size - self.patch_size = config.patch_size - self.merge_type = config.merge_type - - self.patch_embed = MoonVision3dPatchEmbed( - out_dim=config.hidden_size, - patch_size=config.patch_size, - pos_emb_height=config.init_pos_emb_height, - pos_emb_width=config.init_pos_emb_width, - pos_emb_time=config.init_pos_emb_time, - pos_emb_type=config.pos_emb_type, - ) - - self.encoder = MoonViT3dEncoder( - hidden_dim=config.hidden_size, - num_layers=config.num_hidden_layers, - block_cfg={ - "num_heads": config.num_attention_heads, - "hidden_dim": config.hidden_size, - "mlp_dim": config.intermediate_size, - "activation": PytorchGELUTanh(), - "attn_bias": True, - "attn_implementation": config._attn_implementation, - }, - video_attn_type=config.video_attn_type, - ) - - def forward(self, pixel_values: torch.Tensor, grid_thws: torch.Tensor) -> torch.Tensor: - """ - Args: - pixel_values (torch.Tensor): The input pixel values. - grid_thws (torch.Tensor): Temporal, height and width. - Returns: - torch.Tensor: The output tokens. - """ - # grid_thws = grid_thws.to('cpu') - assert grid_thws.ndim == 2, f"grid_thws should be 2D, got {grid_thws.ndim}" - assert grid_thws.size(1) == 3, f"No support for thw: {grid_thws}" - hidden_states = self.patch_embed(pixel_values, grid_thws) - hidden_states = self.encoder(hidden_states, grid_thws) - if self.merge_type == "sd2_tpool": # spatial downsampling 2x with temporal pooling all - hidden_states = tpool_patch_merger(hidden_states, grid_thws, merge_kernel_size=self.merge_kernel_size) - else: - raise NotImplementedError(f"Not support {self.merge_type}") - - return hidden_states - - -# ============================================================================ -# MM Projector Helper Classes (from mm_projector/modeling_mm_projectors.py) -# ============================================================================ - - -class IdentityMap(nn.Module): - def __init__(self): - super().__init__() - - def forward(self, x, *args, **kwargs): - return x - - -class MLP(nn.Module): - def __init__(self, config): - super().__init__() - # TODO, use faster LayerNorm - self.pre_norm = nn.LayerNorm(config.mm_hidden_size) - self.proj = nn.Sequential( - nn.Linear(config.mm_hidden_size, config.hidden_size), - nn.GELU(), - nn.Linear(config.hidden_size, config.hidden_size), - ) - - def forward(self, x, *args, **kwargs): - assert isinstance(x, list | tuple), f"x is not a list or tuple: {type(x)}" - lengths = [item.shape[0] for item in x] - x = torch.cat(x, dim=0) - x = self.pre_norm(x) - x = self.proj(x) - x = torch.split(x, lengths, dim=0) - - return x - - -class PatchMergerMLP(nn.Module): - def __init__(self, config): - super().__init__() - eps = config.projector_ln_eps - self.hidden_size = config.mm_hidden_size * (config.merge_kernel_size[0] * config.merge_kernel_size[1]) - self.pre_norm = nn.LayerNorm(config.mm_hidden_size, eps=eps) - self.proj = nn.Sequential( - nn.Linear(self.hidden_size, self.hidden_size), - nn.GELU(), - nn.Linear(self.hidden_size, config.hidden_size), - ) - - def forward(self, x, *args, **kwargs): - if isinstance(x, list) or isinstance(x, tuple): - x = [self.proj(self.pre_norm(item).view(item.shape[0], -1)) for item in x] - else: - # B, N, N_k, C = x.shape - B = x.shape[0] - x = self.proj(self.pre_norm(x).view(B, -1, self.hidden_size)) - return x - - -class KimiK25PreTrainedModel(PreTrainedModel): - config_class = KimiK25Config - base_model_prefix = "model" - _no_split_modules = [ - "MoonViT3dPretrainedModel", - "MoonViTEncoderLayer", - "DeepseekDecoderLayer", - "PatchMergerMLP", - ] - _skip_keys_device_placement = "past_key_values" - _supports_flash_attn_2 = True - _supports_sdpa = False - - def _init_weights(self, module): - # important: this ported version of Llava isn't meant for training from scratch - only - # inference and fine-tuning - so the proper init weights code has been removed - the original codebase - # https://github.com/haotian-liu/LLaVA/tree/main/llava should serve for that purpose - std = ( - self.config.initializer_range - if hasattr(self.config, "initializer_range") - else self.config.text_config.initializer_range - ) - - if hasattr(module, "class_embedding"): - module.class_embedding.data.normal_(mean=0.0, std=std) - - if isinstance(module, (nn.Linear, nn.Conv2d)): - module.weight.data.normal_(mean=0.0, std=std) - if module.bias is not None: - module.bias.data.zero_() - elif isinstance(module, nn.Embedding): - module.weight.data.normal_(mean=0.0, std=std) - if module.padding_idx is not None: - module.weight.data[module.padding_idx].zero_() - - -class VisionTowerConfig(PretrainedConfig): - model_type = "moonvit3d" - - def __init__(self, config: KimiK25Config, **kwargs): - super().__init__(**kwargs) - self.patch_size = config.patch_size - self.init_pos_emb_height = config.init_pos_emb_height - self.init_pos_emb_width = config.init_pos_emb_width - self.init_pos_emb_time = config.init_pos_emb_time - self.pos_emb_type = config.pos_emb_type - self.num_attention_heads = config.vt_num_attention_heads - self.num_hidden_layers = config.vt_num_hidden_layers - self.hidden_size = config.vt_hidden_size - self.intermediate_size = config.vt_intermediate_size - self.merge_kernel_size = config.merge_kernel_size - self.video_attn_type = config.video_attn_type - self.merge_type = config.merge_type - self._attn_implementation = config._attn_implementation - - -class ProjectorConfig: - def __init__(self, config: KimiK25Config): - self.mm_projector_type = config.mm_projector_type - self.mm_hidden_size = config.mm_hidden_size - self.hidden_size = config.text_hidden_size - self.merge_kernel_size = config.merge_kernel_size - self.projector_hidden_act = config.projector_hidden_act - self.projector_ln_eps = config.projector_ln_eps - - -class QEffKimiK25EncoderWrapper(nn.Module): - def __init__(self, model): - super().__init__() - self.model = model - self.config = self.model.config - # _orig_apply_rope = _kimi_module.apply_rope - # _kimi_module.apply_rope = _apply_rope_real - # _kimi_module = _sys.modules[type(model).__module__] - - # Restore original apply_rope and attention functions - # _kimi_module.apply_rope = _orig_apply_rope - # _kimi_module.VL_VISION_ATTENTION_FUNCTIONS.update(_orig_attn_functions) - - # _orig_attn_functions = _kimi_module.VL_VISION_ATTENTION_FUNCTIONS.copy() - # _kimi_module.VL_VISION_ATTENTION_FUNCTIONS["eager"] = _full_attention_forward - - def get_submodules_for_export(self) -> Type[nn.Module]: - """ - Return the set of class used as the repeated layer across the model for subfunction extraction. - Notes: - This method should return the *class object* (not an instance). - Downstream code can use this to find/build subfunctions for repeated blocks. - """ - return {self.model.vision_model.model.layers[0].__class__} - # return {self.model.layers[0].__class__} - - def forward_only_image(self, pixel_values: torch.Tensor, grid_thws: torch.Tensor) -> list[torch.Tensor]: - """ - Run only the vision tower and mm_projector to extract image embeddings. - - Args: - pixel_values: Preprocessed image pixel values. - grid_thws: Grid temporal/height/width info for the images. - - Returns: - image_embeds: List of projected image embedding tensors, one per image. - """ - image_features = self._extract_image_features(pixel_values, grid_thws) - if self.mm_projector: - image_features = self.mm_projector(image_features) - return image_features - - def forward( - self, pixel_values: torch.Tensor, h_shape: torch.Tensor, w_shape: torch.Tensor - ) -> torch.Tensor: - """ - ONNX-exportable forward that runs only the vision tower and mm_projector. - Uses h_shape and w_shape (int64 ones tensors of length h and w respectively) - to derive spatial dimensions via .shape[0], enabling dynamic axis export. - - Args: - pixel_values: Preprocessed image pixel values (num_patches, channels*patch_h*patch_w). - h_shape: int64 ones tensor of shape (h,) encoding number of patch rows. - w_shape: int64 ones tensor of shape (w,) encoding number of patch columns. - - Returns: - image_embeds: Projected image embeddings as a single tensor. - """ - _kimi_module = _sys.modules[type(self).__module__] - _get_rope_shape = _kimi_module.get_rope_shape - - h_shape = h_shape.to(pixel_values.device) - w_shape = w_shape.to(pixel_values.device) - - # Keep them in ONNX graph - dummy = (h_shape.float().sum() + w_shape.float().sum()) * 0.0 - pixel_values = pixel_values + dummy - - h = h_shape.shape[0] - w = w_shape.shape[0] - - target_dtype = self.model.vision_tower.patch_embed.proj.weight.dtype - pixel_values = pixel_values.to(target_dtype) - - # --- Patch embedding --- - x = self.model.vision_tower.patch_embed.proj(pixel_values).view(pixel_values.size(0), -1) - - # Positional embedding (single image, t=1) - pos_emb_module = self.model.vision_tower.patch_embed.pos_emb - pos_emb_2d = _get_rope_shape( - pos_emb_module.weight, - interpolation_mode=pos_emb_module.interpolation_mode, - shape=(h, w), - ) - x = x + pos_emb_2d - - # --- Encoder --- - # For single image with t=1: cu_seqlens = [0, h*w] - num_tokens = h * w - cu_seqlens = torch.zeros(2, dtype=torch.int32, device=x.device) - cu_seqlens[1] = num_tokens - max_seqlen = num_tokens - - # RoPE frequencies for single image (stacked real/imag to avoid complex tensors in ONNX) - # Shape: (2, num_tokens, dim//2) where [0]=cos, [1]=sin - rope_2d = self.model.vision_tower.encoder.rope_2d - freqs_cos = rope_2d.freqs_cos[:h, :w].reshape(-1, rope_2d.dim // 2) - freqs_sin = rope_2d.freqs_sin[:h, :w].reshape(-1, rope_2d.dim // 2) - freqs_cis = torch.stack([freqs_cos, freqs_sin], dim=0) - - for block in self.model.vision_tower.encoder.blocks: - x = block(x, cu_seqlens, max_seqlen, rope_freqs_cis=freqs_cis) - - x = self.model.vision_tower.encoder.final_layernorm(x) - - # --- tpool_patch_merger (single image, t=1) --- - merge_kernel_size = self.model.vision_tower.merge_kernel_size - kernel_height, kernel_width = merge_kernel_size - d_model = x.size(-1) - new_height = h // kernel_height - new_width = w // kernel_width - reshaped = x.view(1, new_height, kernel_height, new_width, kernel_width, d_model) - reshaped = reshaped.permute(0, 1, 3, 2, 4, 5).contiguous().mean(dim=0) - merged = reshaped.view(new_height * new_width, kernel_height * kernel_width, -1) - - # --- mm_projector (PatchMergerMLP on single tensor) --- - image_embeds = self.model.mm_projector.proj(self.model.mm_projector.pre_norm(merged).view(merged.shape[0], -1)) - return image_embeds - - -class QEffKimiK25DecoderWrapper(nn.Module): - def __init__(self, model): - super().__init__() - self.model = model - self.language_model = self.model.language_model - self.lm_head = self.language_model.lm_head - self.config = self.model.config - - def get_submodules_for_export(self) -> Type[nn.Module]: - """ - Return the set of class used as the repeated layer across the model for subfunction extraction. - Notes: - This method should return the *class object* (not an instance). - Downstream code can use this to find/build subfunctions for repeated blocks. - """ - return {self.model.layers[0].__class__} - - def forward( - self, - input_ids: torch.LongTensor = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - image_embeds: Optional[torch.FloatTensor] = None, - image_idx: Optional[torch.LongTensor] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - compressed_kvs: Optional[List[torch.FloatTensor]] = None, - past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, - batch_index: Optional[torch.LongTensor] = None, - comp_ctx_lengths: Optional[List[int]] = None, - use_cache: Optional[bool] = None, - **kwargs, - ) -> Tuple: - del labels, output_attentions, output_hidden_states, return_dict, cache_position, num_logits_to_keep - if inputs_embeds is None: - inputs_embeds = self.model.get_input_embeddings()(input_ids) - - if image_idx is None: - image_idx = torch.zeros((1, 1), dtype=torch.int64, device=inputs_embeds.device) - - if image_embeds is not None and input_ids is not None and input_ids.shape[1] != 1: - if image_embeds.dim() == 2: - image_embeds = image_embeds.unsqueeze(0) - - image_embeds = image_embeds.to(device=inputs_embeds.device, dtype=inputs_embeds.dtype) - _, _, hidden_size = inputs_embeds.shape - selected = input_ids == self.config.media_placeholder_token_id - indices1 = selected.to(torch.int64).cumsum(1) - 1 - indices1 = torch.where(indices1 != -1, indices1 + image_idx.to(indices1.device), indices1) - indices0 = torch.arange(selected.shape[0], device=selected.device).view(-1, 1) - safe_indices1 = torch.where(indices1 < 0, torch.zeros_like(indices1), indices1) - image_features_expanded = image_embeds.reshape(-1, hidden_size).unsqueeze(0)[indices0, safe_indices1] - image_input_embeds = torch.where(selected.unsqueeze(-1), image_features_expanded, inputs_embeds) - inputs_embeds = torch.where(input_ids.shape[1] == torch.tensor(1), inputs_embeds, image_input_embeds) - image_idx = (indices1.max() + 1).reshape(1, 1) - - outputs = self.model.language_model( - inputs_embeds=inputs_embeds, - attention_mask=attention_mask, - position_ids=position_ids, - compressed_kvs=compressed_kvs, - past_key_values=past_key_values, - batch_index=batch_index, - comp_ctx_lengths=comp_ctx_lengths, - use_cache=True if use_cache is None else use_cache, - **kwargs, - ) - - hidden_states = outputs[0] - logit_index = position_ids.to(torch.int32).argmax(1, keepdim=True) - hidden_states = hidden_states[torch.arange(position_ids.shape[0]).view(-1, 1), logit_index] - logits = self.lm_head(hidden_states).float() - - loss = None - output = (logits,) + outputs[1:] - return logits, image_embeds, image_idx, outputs.compressed_kvs, outputs.past_key_values - - -# ref https://github.com/huggingface/transformers/blob/78b2929c0554b79e0489b451ce4ece14d265ead2/src/transformers/models/llava/modeling_llava.py#L240 -class QEffKimiK25ForConditionalGeneration(KimiK25PreTrainedModel): - def __init__(self, config: KimiK25Config): - super().__init__(config) - - vt_config = VisionTowerConfig(config.vision_config) - self.vision_tower = MoonViT3dPretrainedModel(vt_config) - - proj_config = ProjectorConfig(config.vision_config) - if proj_config.mm_projector_type == "identity": - self.mm_projector = IdentityMap() - elif proj_config.mm_projector_type == "mlp": - self.mm_projector = MLP(proj_config) - elif proj_config.mm_projector_type == "patchmerger": - self.mm_projector = PatchMergerMLP(proj_config) - else: - raise ValueError(f"Unsupported mm_projector_type: {proj_config.mm_projector_type}") - - self.language_model = QEffDeepseekV3ForCausalLM(config.text_config) - self.post_init() - - if hasattr(self.language_model, "dtype"): - target_dtype = self.language_model.dtype - self.vision_tower = self.vision_tower.to(dtype=target_dtype) - self.mm_projector = self.mm_projector.to(dtype=target_dtype) - - def get_qeff_vision_encoder(self): - return QEffKimiK25EncoderWrapper(self) - - def get_qeff_language_decoder(self): - return QEffKimiK25DecoderWrapper(self) - - def get_input_embeddings(self): - return self.language_model.get_input_embeddings() - - def set_input_embeddings(self, value): - self.language_model.set_input_embeddings(value) - - def get_output_embeddings(self): - return self.language_model.get_output_embeddings() - - def set_output_embeddings(self, new_embeddings): - self.language_model.set_output_embeddings(new_embeddings) - - def set_decoder(self, decoder): - self.language_model.set_decoder(decoder) - - def get_decoder(self): - return self.language_model.get_decoder() - - def tie_weights(self): - return self.language_model.tie_weights() - - def resize_token_embeddings(self, new_num_tokens: int | None = None, pad_to_multiple_of=None) -> nn.Embedding: - model_embeds = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of) - # update vocab size - self.config.text_config.vocab_size = model_embeds.num_embeddings - self.vocab_size = model_embeds.num_embeddings - return model_embeds - - def _merge_input_ids_with_image_features( - self, - image_features: list[torch.Tensor], - inputs_embeds: torch.Tensor, - input_ids: torch.Tensor, - attention_mask: torch.Tensor, - labels: torch.Tensor | None = None, - ): - """ - Args: - image_features (:obj:`torch.Tensor` of shape :obj:`(num_image_tokens, embed_dim)`): - The image features to merge with the input embeddings. - inputs_embeds (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length, embed_dim)`): - The input embeddings. - input_ids (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`): - The input ids. - attention_mask (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`): - The attention mask. - labels (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`, *optional*): - The labels. - """ - _, embed_dim = image_features[0].shape - feature_lengths = [x.shape[0] for x in image_features] - image_features = torch.cat(image_features, dim=0) - - image_token_index: int = self.config.media_placeholder_token_id - pad_token_id: int = self.config.pad_token_id - ignore_index: int = self.config.ignore_index - - batch_size, sequence_length = input_ids.shape - left_padding = not torch.sum(input_ids[:, -1] == torch.tensor(pad_token_id)) - - # 1. Create a mask to know where special image tokens are - _token_occupation_table = torch.ones_like(input_ids.flatten()) - _token_occupation_table[input_ids.flatten() == image_token_index] = torch.tensor( - feature_lengths, dtype=torch.long, device=input_ids.device - ) - _token_occupation_table = _token_occupation_table.reshape(input_ids.shape) - - max_embed_dim = _token_occupation_table.sum(-1).max().item() - assert max_embed_dim >= sequence_length, ( - f"The maximum embedding dimension ({max_embed_dim}) is less than the sequence length ({sequence_length})" - ) - batch_indices, non_image_indices = torch.where(input_ids != image_token_index) - - # 2. Compute the positions where text should be written - # Calculate new positions for text tokens in merged image-text sequence. - new_token_positions = torch.cumsum(_token_occupation_table, -1) - 1 - nb_image_pad = max_embed_dim - 1 - new_token_positions[:, -1] - if left_padding: - new_token_positions += nb_image_pad[:, None] # offset for left padding - text_to_overwrite = new_token_positions[batch_indices, non_image_indices] - - # 3. Create the full embedding, already padded to the maximum position - final_embedding = torch.zeros( - batch_size, - max_embed_dim, - embed_dim, - dtype=inputs_embeds.dtype, - device=inputs_embeds.device, - ) - final_attention_mask = torch.zeros( - batch_size, max_embed_dim, dtype=attention_mask.dtype, device=inputs_embeds.device - ) - if labels is not None: - final_labels = torch.full( - (batch_size, max_embed_dim), - ignore_index, - dtype=input_ids.dtype, - device=input_ids.device, - ) - # In case the Vision model or the Language model has been offloaded to CPU, we need to manually - # set the corresponding tensors into their correct target device. - target_device = inputs_embeds.device - batch_indices, non_image_indices, text_to_overwrite = ( - batch_indices.to(target_device), - non_image_indices.to(target_device), - text_to_overwrite.to(target_device), - ) - attention_mask = attention_mask.to(target_device) - - # 4. Fill the embeddings based on the mask. - final_embedding[batch_indices, text_to_overwrite] = inputs_embeds[batch_indices, non_image_indices] - final_attention_mask[batch_indices, text_to_overwrite] = attention_mask[batch_indices, non_image_indices] - if labels is not None: - final_labels[batch_indices, text_to_overwrite] = labels[batch_indices, non_image_indices] - - # 5. Fill the embeddings corresponding to the images. Anything that is not `text_positions` needs filling (#29835) - image_to_overwrite = torch.full( - (batch_size, max_embed_dim), True, dtype=torch.bool, device=inputs_embeds.device - ) - image_to_overwrite[batch_indices, text_to_overwrite] = False - image_to_overwrite &= image_to_overwrite.cumsum(-1) - 1 >= nb_image_pad[:, None].to(target_device) - - if image_to_overwrite.sum() != image_features.shape[:-1].numel(): - raise ValueError( - f"The input provided to the model are wrong. The number of image tokens is {image_to_overwrite.sum()} while" - f" the number of image features given to the model is {image_features.shape[:-1].numel()}. " - "This prevents correct indexing and breaks batch generation." - ) - - final_embedding[image_to_overwrite] = image_features.contiguous().reshape(-1, embed_dim).to(target_device) - final_attention_mask |= image_to_overwrite - position_ids = (final_attention_mask.cumsum(-1) - 1).masked_fill_((final_attention_mask == 0), 1) - - # 6. Mask out the embedding at padding positions, as we later use the past_key_value value to determine the non-attended tokens. - batch_indices, pad_indices = torch.where(input_ids == pad_token_id) - indices_to_mask = new_token_positions[batch_indices, pad_indices] - - final_embedding[batch_indices, indices_to_mask] = 0 - - if labels is None: - final_labels = None - - return final_embedding, final_attention_mask, final_labels, position_ids - - def _extract_image_features(self, pixel_values: torch.Tensor, grid_thws: torch.Tensor) -> list[torch.Tensor]: - """ - Args: - pixel_values (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, num_channels, height, width)`): - The pixel values of the images processed by image processor. - grid_thws (:obj:`torch.Tensor` of shape :obj:`(batch_size, 3)`): - The grid, height, width of the images. - Returns: - selected_image_feature (:obj:`torch.FloatTensor` of shape :obj:`(num_image_tokens, embed_dim)`): - The selected image features to use as input to the projector head. - """ - - target_dtype = self.vision_tower.patch_embed.proj.weight.dtype - pixel_values = pixel_values.to(target_dtype) - - image_features = self.vision_tower(pixel_values, grid_thws) - return image_features - - def forward( - self, - input_ids: torch.LongTensor | None = None, - pixel_values: torch.FloatTensor | list[torch.FloatTensor] | None = None, - grid_thws: torch.Tensor | None = None, - # h_shape: torch.Tensor | None = None, - # w_shape: torch.Tensor | None = None, - attention_mask: torch.Tensor | None = None, - position_ids: torch.LongTensor | None = None, - past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, - compressed_kvs: Optional[List[torch.FloatTensor]] = None, - inputs_embeds: torch.FloatTensor | None = None, - labels: torch.LongTensor | None = None, - use_cache: bool | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - return_dict: bool | None = None, - ) -> tuple | LlavaCausalLMOutputWithPast: - r""" - Args: - labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): - Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., - config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored - (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. - ```""" - assert self.vision_tower is not None, "vision_tower is not loaded" - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if inputs_embeds is None: - # 1. Extra the input embeddings - inputs_embeds = self.get_input_embeddings()(input_ids) - - # 2. Merge text and images - if pixel_values is not None and len(pixel_values) > 0 and input_ids.shape[1] != 1: - image_features = self._extract_image_features(pixel_values, grid_thws) - if self.mm_projector: - image_features = self.mm_projector(image_features) - - inputs_embeds = inputs_embeds.to(image_features[0].dtype) # num_tokens, embed_dim - inputs_embeds, attention_mask, labels, position_ids = self._merge_input_ids_with_image_features( - image_features, - inputs_embeds, - input_ids, - attention_mask, - labels, - ) - - # In case input_ids.shape[1] == 1 & pixel_values==None & past_key_values != None, we are in the case of - # generation with cache - elif past_key_values is not None and pixel_values is not None and input_ids.shape[1] == 1: - # Retrieve the first layer to inspect the logits and mask out the hidden states - # that are set to 0 - first_layer_past_key_value = past_key_values[0][0][:, :, :, 0] - - # Sum all dimensions of head_dim (-2) to avoid random errors such as: https://github.com/huggingface/transformers/pull/28032#issuecomment-1863691941 - batch_index, non_attended_tokens = torch.where(first_layer_past_key_value.float().sum(-2) == 0) - - # Get the target length - target_length = input_ids.shape[1] - past_length = first_layer_past_key_value.shape[-1] - - extended_attention_mask = torch.ones( - (attention_mask.shape[0], past_length), - dtype=attention_mask.dtype, - device=attention_mask.device, - ) - - # Filter out only the tokens that can be un-attended, this can happen - # if one uses Llava + Fused modules where the cache on the - # first iteration is already big enough, or if one passes custom cache - valid_indices = non_attended_tokens < extended_attention_mask.size(-1) - new_batch_index = batch_index[valid_indices] - new_non_attended_tokens = non_attended_tokens[valid_indices] - - # Zero-out the places where we don't need to attend - extended_attention_mask[new_batch_index, new_non_attended_tokens] = 0 - - attention_mask = torch.cat((extended_attention_mask, attention_mask[:, -target_length:]), dim=1) - position_ids = torch.sum(attention_mask, dim=1).unsqueeze(-1) - 1 - - outputs = self.language_model( - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - compressed_kvs=compressed_kvs, - inputs_embeds=inputs_embeds, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - - logits = outputs[0] - - loss = None - if labels is not None: - # Shift so that tokens < n predict n - if attention_mask is not None: - shift_attention_mask = attention_mask[..., 1:] - shift_logits = logits[..., :-1, :][shift_attention_mask.to(logits.device) != 0].contiguous() - shift_labels = labels[..., 1:][shift_attention_mask.to(labels.device) != 0].contiguous() - else: - shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - # Flatten the tokens - loss_fct = nn.CrossEntropyLoss() - loss = loss_fct( - shift_logits.view(-1, shift_logits.size(-1)), - shift_labels.view(-1).to(shift_logits.device), - ) - - if not return_dict: - output = (logits,) + outputs[1:] - return (loss,) + output if loss is not None else output - - return LlavaCausalLMOutputWithPast( - loss=loss, - logits=logits, - past_key_values=outputs.past_key_values, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) - - def prepare_inputs_for_generation( - self, - input_ids, - past_key_values=None, - inputs_embeds=None, - pixel_values=None, - grid_thws=None, - attention_mask=None, - **kwargs, - ): - if past_key_values is not None: - if isinstance(past_key_values, Cache): - cache_length = past_key_values.get_seq_length() - past_length = getattr(past_key_values, "seen_tokens", cache_length) - else: - cache_length = past_length = past_key_values[0][0].shape[2] - - # Keep only the unprocessed tokens: - # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where - # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as - # input) - if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]: - input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :] - # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard - # input_ids based on the past_length. - elif past_length < input_ids.shape[1]: - input_ids = input_ids[:, past_length:] - # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens. - elif self.config.media_placeholder_token_id in input_ids: - input_ids = input_ids[:, input_ids.shape[1] - 1 :] - # If the cache has seen more tokens than it can hold, then the cache has a size limit. Let's discard the - # older attention values, as their corresponding values are not part of the input. - if cache_length < past_length and attention_mask is not None: - attention_mask = attention_mask[:, -(cache_length + input_ids.shape[1]) :] - - position_ids = kwargs.get("position_ids", None) - if attention_mask is not None and position_ids is None: - # create position_ids on the fly for batch generation - position_ids = attention_mask.long().cumsum(-1) - 1 - position_ids.masked_fill_(attention_mask == 0, 1) - if past_key_values: - position_ids = position_ids[:, -input_ids.shape[1] :] - - # if `inputs_embeds` are passed, we only want to use them in the 1st generation step - if inputs_embeds is not None and past_key_values is None: - model_inputs = {"inputs_embeds": inputs_embeds} - else: - model_inputs = {"input_ids": input_ids} - - model_inputs.update( - { - "position_ids": position_ids, - "past_key_values": past_key_values, - "use_cache": kwargs.get("use_cache"), - "attention_mask": attention_mask, - "pixel_values": pixel_values, - "grid_thws": grid_thws, - } - ) - return model_inputs - - def _reorder_cache(self, *args, **kwargs): - return self.language_model._reorder_cache(*args, **kwargs) - - def get_output_names(self, kv_offload: bool = False): - vision_output_names = ["image_embeds"] - lang_output_names = ["logits"] - - mla_absorption = getattr(self.language_model, "mla_absorption", None) - if mla_absorption is not None: - cache_compressed = mla_absorption.get("cache_compressed", False) - else: - cache_compressed = False - - if cache_compressed: - for i in range(self.language_model.config.num_hidden_layers): - lang_output_names.append(f"compressed_kv.{i}_RetainedState") - lang_output_names.append(f"k_pe.{i}_RetainedState") - else: - for i in range(self.language_model.config.num_hidden_layers): - for kv in ["key", "value"]: - lang_output_names.append(f"past_{kv}.{i}_RetainedState") - - output_names = {} - if kv_offload: - output_names["vision"] = vision_output_names - output_names["lang"] = lang_output_names - else: - return lang_output_names - return output_names - - def get_dummy_inputs( - self, - comp_ctx_lengths: Optional[List[int]] = None, - kv_offload: bool = False, - continuous_batching: bool = False, - **kwargs, - ): - prefill_seq_len = kwargs.get("prefill_seq_len") - if prefill_seq_len is None: - prefill_seq_len = constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN - prefill_seq_len = int(prefill_seq_len) - - bs: int = constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE - fbs: int = constants.ONNX_EXPORT_EXAMPLE_FBS - - model_path = Path( - "/home/huggingface_hub/models--moonshotai--Kimi-K2.5/snapshots/4d01dfe0332d63057c186e0b262165819efb6611" - ) - processor = AutoProcessor.from_pretrained(str(model_path), trust_remote_code=True) - image_url = "https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png" - image = Image.open(BytesIO(requests.get(image_url).content)).convert("RGB") - messages = [ - { - "role": "user", - "content": [ - {"type": "image_url", "image_url": image}, - {"type": "text", "text": "Tell me about yourself."}, - ], - } - ] - inputs = processor( - messages=messages, - add_generation_prompt=True, - tokenize=False, - return_tensors="pt", - ) - inputs = {k: v.to(self.language_model.device) if hasattr(v, "to") else v for k, v in inputs.items()} - - # Build h_shape and w_shape from grid_thws (single image, t=1) - grid_thws_val = inputs["grid_thws"] - h_val = int(grid_thws_val[0, 1].item()) - w_val = int(grid_thws_val[0, 2].item()) - h_shape_tensor = torch.ones(h_val, dtype=torch.int64, device=self.language_model.device) - w_shape_tensor = torch.ones(w_val, dtype=torch.int64, device=self.language_model.device) - - vision_inputs = { - "pixel_values": inputs["pixel_values"], - "h_shape": h_shape_tensor, - "w_shape": w_shape_tensor, - } - - lang_inputs = { - "input_ids": torch.zeros((bs, prefill_seq_len), dtype=torch.int64), - "position_ids": torch.arange(prefill_seq_len, dtype=torch.int64).view(1, prefill_seq_len).repeat(bs, 1), - } - - mla_absorption = getattr(self.language_model, "mla_absorption", None) - if mla_absorption is not None: - cache_compressed = mla_absorption.get("cache_compressed", False) - else: - cache_compressed = False - - pkv_cache = self.language_model.get_dummy_pkv_cache( - config=self.language_model.config, - batch_size=fbs if continuous_batching else bs, - seq_len=prefill_seq_len, - ) - - if cache_compressed: - lang_inputs["compressed_kvs"] = [[] for _ in range(self.language_model.config.num_hidden_layers)] - for i in range(self.language_model.config.num_hidden_layers): - lang_inputs["compressed_kvs"][i].append( - torch.zeros(pkv_cache[0][0].shape, dtype=self.language_model.config.torch_dtype) - ) - lang_inputs["compressed_kvs"][i].append( - torch.zeros(pkv_cache[0][1].shape, dtype=self.language_model.config.torch_dtype) - ) - else: - lang_inputs["past_key_values"] = [[] for _ in range(self.language_model.config.num_hidden_layers)] - for i in range(self.language_model.config.num_hidden_layers): - lang_inputs["past_key_values"][i].append( - torch.zeros(pkv_cache[0][0].shape, dtype=self.language_model.config.torch_dtype) - ) - lang_inputs["past_key_values"][i].append( - torch.zeros(pkv_cache[0][1].shape, dtype=self.language_model.config.torch_dtype) - ) - - if continuous_batching: - lang_inputs["batch_index"] = torch.arange(bs).view(bs, 1) - - if comp_ctx_lengths is not None: - lang_inputs["comp_ctx_lengths"] = torch.randint(0, 100, (40,), dtype=torch.int64) - - inputs = {} - if kv_offload: - inputs["vision"] = vision_inputs - inputs["lang"] = lang_inputs - else: - lang_inputs.pop("image_embeds") - inputs = {**vision_inputs, **lang_inputs} - - return inputs - - def get_onnx_dynamic_axes( - self, comp_ctx_lengths: Optional[List[int]] = None, kv_offload: bool = False, continuous_batching: bool = False - ): - vision_dynamic_axes = {} - lang_dynamic_axes = {} - lang_dynamic_axes["input_ids"] = {0: "batch_size", 1: "seq_len"} - lang_dynamic_axes["position_ids"] = {0: "batch_size", 1: "seq_len"} - lang_dynamic_axes["image_embeds"] = {0: "num_image_tokens"} - if continuous_batching: - lang_dynamic_axes["batch_index"] = {0: "batch_size"} - vision_dynamic_axes = { - "pixel_values": {0: "num_patches"}, - "h_shape": {0: "h"}, - "w_shape": {0: "w"}, - "image_embeds": {0: "num_image_tokens"}, - } - - mla_absorption = getattr(self.language_model, "mla_absorption", None) - if mla_absorption is not None: - cache_compressed = mla_absorption.get("cache_compressed", False) - else: - cache_compressed = False - - if cache_compressed: - for i in range(self.language_model.config.num_hidden_layers): - lang_dynamic_axes[f"compressed_kv.{i}"] = {0: "batch_size", 2: "ctx_len"} - lang_dynamic_axes[f"k_pe.{i}"] = {0: "batch_size", 2: "ctx_len"} - else: - for i in range(self.language_model.config.num_hidden_layers): - for kv in ["key", "value"]: - lang_dynamic_axes[f"past_{kv}.{i}"] = {0: "batch_size", 2: "ctx_len"} - - if comp_ctx_lengths is not None: - lang_dynamic_axes["comp_ctx_lengths"] = {0: "comp_ctx_lengths"} - - dynamic_axes = {} - if kv_offload: - dynamic_axes["vision"] = vision_dynamic_axes - dynamic_axes["lang"] = lang_dynamic_axes - else: - lang_dynamic_axes.pop("image_embeds") - dynamic_axes = {**vision_dynamic_axes, **lang_dynamic_axes} - return dynamic_axes - - def get_specializations( - self, - batch_size: int, - prefill_seq_len: int, - ctx_len: int, - kv_offload: bool = False, - continuous_batching: bool = False, - kv_cache_batch_size: Optional[int] = None, - full_batch_size: Optional[int] = None, - **compiler_options, - ): - comp_ctx_lengths_prefill = compiler_options.pop("comp_ctx_lengths_prefill", None) - comp_ctx_lengths_decode = compiler_options.pop("comp_ctx_lengths_decode", None) - num_patches = compiler_options.pop("num_patches", None) - h = compiler_options.pop("h", None) - w = compiler_options.pop("w", None) - num_image_tokens = compiler_options.pop("num_image_tokens", None) - - prefill_seq_len = prefill_seq_len if prefill_seq_len else 32 # constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN - ctx_len = ctx_len if ctx_len else 32 # constants.ONNX_EXPORT_EXAMPLE_CTX_LEN - - vision = [ - { - "num_patches": 2400, # num_patches - "h": 30, # h - "w": 80, # w - "num_image_tokens": 600, # num_image_tokens - } - ] - - if comp_ctx_lengths_prefill is not None: - lang = [] - - for i in range(0, len(comp_ctx_lengths_prefill)): - lang_prefill = { - "batch_size": 1 if continuous_batching else batch_size, - "seq_len": prefill_seq_len, - "ctx_len": ctx_len, - "num_image_tokens": 600, # num_image_tokens - } - if continuous_batching: - lang_prefill["full_batch_size"] = kv_cache_batch_size - else: - lang_prefill["batch_size"] = kv_cache_batch_size - if full_batch_size: - lang_prefill["full_batch_exec_size"] = full_batch_size - - lang.append(lang_prefill) - - for i in range(0, len(comp_ctx_lengths_decode)): - lang_decode = { - "batch_size": full_batch_size if continuous_batching else batch_size, - "seq_len": "1", - "ctx_len": ctx_len, - "num_image_tokens": 600, # num_image_tokens - } - - if continuous_batching: - lang_decode["full_batch_size"] = kv_cache_batch_size - else: - lang_decode["batch_size"] = kv_cache_batch_size - - lang.append(lang_decode) - - else: - lang_prefill = { - "batch_size": 1 if continuous_batching else batch_size, - "seq_len": prefill_seq_len, - "ctx_len": ctx_len, - "num_image_tokens": 600, # num_image_tokens - } - if continuous_batching: - lang_prefill["full_batch_size"] = kv_cache_batch_size - else: - lang_prefill["batch_size"] = kv_cache_batch_size - if full_batch_size: - lang_prefill["full_batch_exec_size"] = full_batch_size - - lang_decode = { - "batch_size": full_batch_size if continuous_batching else batch_size, - "seq_len": 1, - "ctx_len": ctx_len, - "num_image_tokens": 600, # num_image_tokens - } - - if continuous_batching: - lang_decode["full_batch_size"] = kv_cache_batch_size - else: - lang_decode["batch_size"] = kv_cache_batch_size - - lang = [lang_prefill, lang_decode] - - specializations = {} - - if kv_offload: - specializations["vision"] = vision - specializations["lang"] = lang - return specializations, compiler_options - else: - lang[0].pop("vision_size") - lang[1].pop("vision_size") - return lang, compiler_options +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ---------------------------------------------------------------------------- + +import math +import sys as _sys +from collections.abc import Sequence +from copy import deepcopy +from io import BytesIO +from pathlib import Path + +# from QEfficient import QEFFAutoModelForImageTextToText +from typing import List, Optional, Tuple, Type, Union + +import numpy as np +import requests +import torch +import torch.nn as nn +import torch.nn.functional as F +from PIL import Image +from transformers import AutoProcessor, activations + +from QEfficient.transformers.models.deepseek_v3.modeling_deepseek import QEffDeepseekV3ForCausalLM + +try: + from transformers.activations import PytorchGELUTanh +except ImportError: + from transformers.activations import GELUTanh + + activations.PytorchGELUTanh = GELUTanh + PytorchGELUTanh = GELUTanh +from transformers.activations import PytorchGELUTanh +from transformers.cache_utils import Cache +from transformers.configuration_utils import PretrainedConfig +from transformers.modeling_outputs import CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.models.llava.modeling_llava import LlavaCausalLMOutputWithPast + +from QEfficient.utils import constants + +from .configuration_kimi_k25 import KimiK25Config + + +def eager_attention_forward(q, k, v, **kwargs): + q = q.transpose(0, 1) # (num_heads, seq_len, head_dim) + k = k.transpose(0, 1) + v = v.transpose(0, 1) + attn_weight = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1]) + attn_weight = torch.softmax(attn_weight, dim=-1, dtype=torch.float32).to(q.dtype) + attn_out = attn_weight @ v + attn_out = attn_out.transpose(0, 1) # (seq_len, num_heads, head_dim) + return attn_out.reshape(attn_out.shape[0], -1) + + +VL_VISION_ATTENTION_FUNCTIONS = {"eager": eager_attention_forward} + + +def get_rope_shape_decorate(func): + _get_rope_shape_first_call_flag = set() + + def wrapper(org, interpolation_mode, shape): + key = (org.requires_grad, torch.is_grad_enabled(), interpolation_mode) + if key not in _get_rope_shape_first_call_flag: + _get_rope_shape_first_call_flag.add(key) + _ = func(org, interpolation_mode, shape=(64, 64)) + return func(org, interpolation_mode, shape) + + return wrapper + + +@get_rope_shape_decorate + +def get_rope_shape(org, interpolation_mode, shape): + return (F.interpolate( + org.permute((2, 0, 1)).unsqueeze(0), + size=shape, + mode=interpolation_mode, + ).squeeze(0).permute((1, 2, 0)).flatten(end_dim=1)) + + +def apply_rope(xq, xk, freqs_cis): + # Support both complex freqs (..., dim//2) and stacked real/imag (2, ..., dim//2) + if torch.is_complex(freqs_cis): + freqs_cis = freqs_cis.unsqueeze(-2) + xq_ = torch.view_as_complex(xq.float().view(*xq.shape[:-1], -1, 2)) + xk_ = torch.view_as_complex(xk.float().view(*xk.shape[:-1], -1, 2)) + xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(-2) + xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(-2) + return xq_out.type_as(xq), xk_out.type_as(xk) + + # freqs_cis shape: (2, seq_len, dim//2) + # xq shape: (seq_len, num_heads, head_dim) + freqs_cos = freqs_cis[0].unsqueeze(-2) # (seq_len, 1, dim//2) + freqs_sin = freqs_cis[1].unsqueeze(-2) # (seq_len, 1, dim//2) + xq_r = xq.float().view(*xq.shape[:-1], -1, 2) + xq_r0, xq_r1 = xq_r[..., 0], xq_r[..., 1] + xq_out_r = xq_r0 * freqs_cos - xq_r1 * freqs_sin + xq_out_i = xq_r0 * freqs_sin + xq_r1 * freqs_cos + xq_out = torch.stack([xq_out_r, xq_out_i], dim=-1).flatten(-2) + xk_r = xk.float().view(*xk.shape[:-1], -1, 2) + xk_r0, xk_r1 = xk_r[..., 0], xk_r[..., 1] + xk_out_r = xk_r0 * freqs_cos - xk_r1 * freqs_sin + xk_out_i = xk_r0 * freqs_sin + xk_r1 * freqs_cos + xk_out = torch.stack([xk_out_r, xk_out_i], dim=-1).flatten(-2) + return xq_out.type_as(xq), xk_out.type_as(xk) + + +def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): + """ + From: + https://github.com/OpenGVLab/InternVideo/blob/421f6d2361fc8f61a3394244571f2601a4e99e29/InternVideo2/multi_modality/models/backbones/internvideo2/pos_embed.py#L86 + embed_dim: output dimension for each position + pos: a list of positions to be encoded: size (M,) + out: (M, D) + """ + assert embed_dim % 2 == 0 + omega = np.arange(embed_dim // 2, dtype=np.float32) + omega /= embed_dim / 2.0 + omega = 1.0 / 10000**omega # (D/2,) + + pos = pos.reshape(-1) # (M,) + out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product + + emb_sin = np.sin(out) # (M, D/2) + emb_cos = np.cos(out) # (M, D/2) + + emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D) + return emb + + +def get_1d_sincos_pos_embed(embed_dim, t_size, cls_token=False): + """ + t_size: int of the temporal size + return: + pos_embed: [t_size, embed_dim] or [1+t_size, embed_dim] (w/ or w/o cls_token) + """ + grid_t = np.arange(t_size, dtype=np.float32) + pos_embed = get_1d_sincos_pos_embed_from_grid(embed_dim, grid_t) + if cls_token: + pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0) + return pos_embed + + +class QEffLearnable2DInterpPosEmbDivided_fixed(nn.Module): + def __qeff_init__(self): + self.interpolation_mode = "bilinear" + + """def __qeff_init__(self, + height: int, + width: int, + num_frames: int, + dim: int, + interpolation_mode: str = 'bilinear') -> None: + super().__init__() + self.height = height + self.width = width + self.num_frames = num_frames + self.dim = dim + self.interpolation_mode = interpolation_mode + self.weight = nn.Parameter(torch.empty(height, width, dim)) + self.register_buffer('time_weight', + torch.from_numpy( + get_1d_sincos_pos_embed( + self.dim, + self.num_frames)).float().unsqueeze(1), + persistent=False) + + self.reset_parameters() + """ + + def reset_parameters(self): + nn.init.normal_(self.weight) + + def forward(self, x: torch.Tensor, grid_thws: torch.Tensor) -> torch.Tensor: + pos_embs = [] + for t, h, w in grid_thws.tolist(): + assert t <= self.num_frames, f"t:{t} > self.num_frames:{self.num_frames}" + if (h, w) == self.weight.shape[:-1]: + pos_emb_2d = self.weight.flatten(end_dim=1) + else: + pos_emb_2d = get_rope_shape( + self.weight, + interpolation_mode=self.interpolation_mode, + shape=(h, w), + ) + if t == 1: + pos_emb_3d = pos_emb_2d + else: + pos_emb_3d = pos_emb_2d.unsqueeze(0).repeat(t, 1, 1) + self.time_weight[0:t] + + pos_embs.append(pos_emb_3d.reshape(-1, pos_emb_3d.shape[-1])) + + out = x + torch.cat(pos_embs) + return out + + +class MoonVision3dPatchEmbed(nn.Module): + def __init__( + self, + out_dim: int, + in_dim: int = 3, + patch_size: int | tuple[int, int] = (14, 14), + pos_emb_height: int = 14, + pos_emb_width: int = 14, + pos_emb_time: int = 4, + pos_emb_type: str = "divided_fixed", + ): + super().__init__() + assert isinstance(patch_size, int | Sequence), f"Invalid patch_size type: {type(patch_size)}" + if isinstance(patch_size, int): + patch_size = (patch_size, patch_size) + assert len(patch_size) == 2, f"Expected patch_size to be a tuple of 2, got {patch_size}" + self.patch_size = patch_size + + self.proj = nn.Conv2d(in_dim, out_dim, kernel_size=patch_size, stride=patch_size) + + if pos_emb_type == "divided_fixed": + self.pos_emb = Learnable2DInterpPosEmbDivided_fixed( + height=pos_emb_height, width=pos_emb_width, num_frames=pos_emb_time, dim=out_dim + ) + else: + raise NotImplementedError(f"Not support pos_emb_type: {pos_emb_type}") + + def forward(self, x: torch.Tensor, grid_thws: torch.Tensor) -> torch.Tensor: + """ + Args: + x (L, Channels): input tensor + grid_hws (N, 3): temporal, height and width + Returns: + (L, Cout) tensor + """ + x = self.proj(x).view(x.size(0), -1) + # apply positional embedding + x = self.pos_emb(x, grid_thws) + return x + + +class Rope2DPosEmbRepeated(nn.Module): + """2D rotary position embedding with multi-resolution support. + This class is intended to be used in the following way: + 1. Before training, create an instance of Rope2DPosEmb. This instance will hold the precomputed cis. + 2. Before each forward pass, call `get_freqs_cis_by_*` to get the `freqs_cis` tensor for this iteration. + 3. During the forward pass, pass the `freqs_cis` tensor to each attention layer, and call `apply` just before each attention operation. + The rope is shared across all attention layers and all heads. + Refs: + - RoFormer: https://arxiv.org/abs/2104.09864 + - VisionLLaMA: https://arxiv.org/abs/2403.00522 + - https://github.com/Meituan-AutoML/VisionLLaMA/blob/main/dit/models.py + Args: + dim (int): usually the multi-head attention dimension, should be divisible by 4 (TODO: relax this constraint if needed) + max_height (int): the maximum height of the 2D grid + max_width (int): the maximum width of the 2D grid + theta_base (float): the base of the theta + device (str): the device to store the precomputed cis + """ + + def __init__(self, dim: int, max_height: int, max_width: int, theta_base=10000): + super().__init__() + self.dim = dim + assert self.dim % 4 == 0, "dim must be divisible by 4" + self.max_height = max_height + self.max_width = max_width + self.theta_base = theta_base + + def extra_repr(self): + return f"dim={self.dim}, max_height={self.max_height}, max_width={self.max_width}, theta_base={self.theta_base}" + + def _ensure_precomputed_freqs(self, device: torch.device) -> None: + if not hasattr(self, "freqs_cis"): + self.register_buffer("freqs_cis", self._precompute_freqs_cis(device), persistent=False) + elif self.freqs_cis.device != device: + self.freqs_cis = self._precompute_freqs_cis(device) + + if not hasattr(self, "freqs_cos"): + self.register_buffer("freqs_cos", self.freqs_cis.real.contiguous(), persistent=False) + elif self.freqs_cos.device != device: + self.freqs_cos = self.freqs_cis.real.contiguous() + + if not hasattr(self, "freqs_sin"): + self.register_buffer("freqs_sin", self.freqs_cis.imag.contiguous(), persistent=False) + elif self.freqs_sin.device != device: + self.freqs_sin = self.freqs_cis.imag.contiguous() + + def _precompute_freqs_cis(self, device: torch.device) -> torch.Tensor: + """Calculate the cis(freqs) for each position in the 2D grid. + Return: complex tensor of shape (max_height, max_width, dim//2) and value: + height axis: ret[h, w, 2*i] = cis(h * theta_base**(-4*i/dim)) + weight axis: ret[h, w, 2*i+1] = cis(w * theta_base**(-4*i/dim)) with (i in [0, dim//4)) + note: `cis` is a mathematical notation defined by cis x = cos x + i sin x, + """ + N = self.max_height * self.max_width + flat_pos = torch.arange(0, N).float().to(device) + x_pos = flat_pos % self.max_width + y_pos = flat_pos // self.max_width + dim_range = torch.arange(0, self.dim, 4)[: (self.dim // 4)].float().to(device) # C/4 + freqs = 1.0 / (self.theta_base ** (dim_range / self.dim)) + x_freqs = torch.outer(x_pos, freqs).float() # N, C/4 + y_freqs = torch.outer(y_pos, freqs).float() # N, C/4 + x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs) # N, C/4 + y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs) # N, C/4 + # N, C/4, 2 + freqs_cis = torch.cat([x_cis.unsqueeze(dim=-1), y_cis.unsqueeze(dim=-1)], dim=-1) + # max_height, max_width, C/2 + freqs_cis = freqs_cis.reshape(self.max_height, self.max_width, -1) + return freqs_cis + + def get_freqs_cis(self, grid_thws: torch.Tensor, device: torch.device) -> torch.Tensor: + """ + Args: + grid_thws (torch.Tensor): grid time, height and width + Returns: + freqs_cis: tensor of shape (sum(t * height * width), dim//2) + """ + self._ensure_precomputed_freqs(device) + + shapes = grid_thws.tolist() + assert all(1 <= h <= self.max_height and 1 <= w <= self.max_width for t, h, w in shapes), ( + shapes, + self.max_height, + self.max_width, + ) + freqs_cis = torch.cat( + [self.freqs_cis[:h, :w].reshape(-1, self.dim // 2).repeat(t, 1) for t, h, w in shapes], + dim=0, + ) + return freqs_cis + + +class MLP2(nn.Module): + """ + Args: + dims: [in_dim, hidden_dim, out_dim] + bias: whether to use bias in linear layer. + """ + + def __init__(self, dims: list[int], activation, bias=True): + super().__init__() + assert len(dims) == 3 + self.fc0 = nn.Linear(dims[0], dims[1], bias=bias) + self.fc1 = nn.Linear(dims[1], dims[2], bias=bias) + self.activation = activation + for m in [self.fc0, self.fc1]: + nn.init.trunc_normal_(m.weight, std=math.sqrt(2 / m.in_features)) + if m.bias is not None: + nn.init.zeros_(m.bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.fc0(x) + x = self.activation(x) + return self.fc1(x) + + +class MoonViTEncoderLayer(nn.Module): + def __init__( + self, + num_heads: int, + hidden_dim: int, + mlp_dim: int, + *, + attn_implementation: str = "flash_attention_2", + activation=F.gelu, + attn_bias: bool = False, + use_deterministic_attn: bool = False, + ): + super().__init__() + self.num_heads = num_heads + self.hidden_dim = hidden_dim + self.hidden_size_per_attention_head = self.hidden_dim // self.num_heads + self.attn_implementation = attn_implementation + self.use_deterministic_attn = use_deterministic_attn + + self.norm0 = nn.LayerNorm(hidden_dim) + self.norm1 = nn.LayerNorm(hidden_dim) + self.mlp = MLP2([hidden_dim, mlp_dim, hidden_dim], activation) + self.wqkv = nn.Linear(hidden_dim, hidden_dim * 3, bias=attn_bias) + self.wo = nn.Linear(hidden_dim, hidden_dim, bias=attn_bias) + + def attention_qkvpacked( + self, + x: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: torch.Tensor, + rope_freqs_cis: torch.Tensor | None = None, + ): + """ + Args: + x (torch.Tensor): (batch_size, seqlen, hidden_dim) + cu_seqlens (torch.Tensor): + """ + xqkv = self.wqkv(x) + + qkv_shape = xqkv.size()[:-1] + ( + 3, + self.num_heads, + self.hidden_size_per_attention_head, + ) + # xqkv: (batch_size, seqlen, 3, nheads, headdim) + xqkv = xqkv.view(*qkv_shape) + xq, xk, xv = torch.unbind(xqkv, dim=-3) + + xq, xk = apply_rope(xq, xk, rope_freqs_cis) + + attn_func = VL_VISION_ATTENTION_FUNCTIONS[self.attn_implementation] + attn_out = attn_func( + xq, + xk, + xv, + q_cu_seqlens=cu_seqlens, + k_cu_seqlens=cu_seqlens, + max_seqlen_k=max_seqlen, + max_seqlen_q=max_seqlen, + deterministic=self.use_deterministic_attn, + ) + + attn_out = self.wo(attn_out) + return attn_out + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: int, + rope_freqs_cis: torch.Tensor | None = None, + ): + residual = hidden_states + hidden_states = self.norm0(hidden_states) + + hidden_states = self.attention_qkvpacked(hidden_states, cu_seqlens, max_seqlen, rope_freqs_cis) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.norm1(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +class QEffMoonViT3dEncoder(nn.Module): + def __qeff_init__(self): + if not hasattr(self, "blocks") or len(self.blocks) == 0: + return + + first_block = self.blocks[0] + self.block_cfg = { + "num_heads": first_block.num_heads, + "hidden_dim": first_block.hidden_dim, + "mlp_dim": first_block.mlp.fc0.out_features, + "activation": PytorchGELUTanh(), + "attn_bias": first_block.wqkv.bias is not None, + "attn_implementation": "eager", + } + + head_dim = first_block.hidden_size_per_attention_head + max_height = getattr(self.rope_2d, "max_height", 512) + max_width = getattr(self.rope_2d, "max_width", 512) + theta_base = getattr(self.rope_2d, "theta_base", 10000) + self.rope_2d = Rope2DPosEmbRepeated(head_dim, max_height, max_width, theta_base=theta_base) + + self.blocks = nn.ModuleList( + [MoonViTEncoderLayer(**self.block_cfg, use_deterministic_attn=False) for _ in range(len(self.blocks))] + ) + + def forward( + self, + hidden_states: torch.Tensor, + grid_thws: torch.Tensor, + ) -> torch.Tensor: + # if not hasattr(self.rope_2d, 'freqs_cos'): + # self.rope_2d.register_buffer('freqs_cos', self.rope_2d.freqs_cis.real.contiguous(), persistent=False) + # self.rope_2d.register_buffer('freqs_sin', self.rope_2d.freqs_cis.imag.contiguous(), persistent=False) + rope_freqs_cis = self.rope_2d.get_freqs_cis(grid_thws=grid_thws, device=hidden_states.device) + + lengths = torch.cat( + ( + torch.zeros(1, dtype=grid_thws.dtype, device=grid_thws.device), + grid_thws[:, 0] * grid_thws[:, 1] * grid_thws[:, 2], + ) + ) + + max_seqlen = lengths.max() + cu_seqlens = lengths.to(hidden_states.device).cumsum(dim=0, dtype=torch.int32) + for block in self.blocks: + hidden_states = block(hidden_states, cu_seqlens, max_seqlen, rope_freqs_cis=rope_freqs_cis) + + hidden_states = self.final_layernorm(hidden_states) + return hidden_states + + +def tpool_patch_merger( + x: torch.Tensor, + grid_thws: torch.Tensor, + merge_kernel_size: tuple[int, int] = (2, 2), +) -> list[torch.Tensor]: + d_model = x.size(-1) + + outputs = [] + pre_sum = 0 + for t, h, w in grid_thws.tolist(): + # Get the current sequence + seq = x[pre_sum : pre_sum + t * h * w] + # Reshape along self.merge_kernel_size and concat to the last dimension + kernel_height, kernel_width = merge_kernel_size + new_height, new_width = h // kernel_height, w // kernel_width + reshaped_seq = seq.view(t, new_height, kernel_height, new_width, kernel_width, d_model) + reshaped_seq = reshaped_seq.permute(0, 1, 3, 2, 4, 5).contiguous().mean(dim=0) # temporal pooling + padded_seq = reshaped_seq.view(new_height * new_width, kernel_height * kernel_width, -1) + outputs.append(padded_seq) + pre_sum += t * h * w + + return outputs + + +class MoonViT3dPretrainedModel(PreTrainedModel): + config_class = None + model_type = "moonvit3d" + _no_split_modules = ["PackingTransformer"] + _supports_flash_attn_2 = True + _supports_sdpa = True + + def __init__(self, config, *inputs, **kwargs): + super().__init__(config, *inputs, **kwargs) + config = deepcopy(config) + self.merge_kernel_size = config.merge_kernel_size + self.patch_size = config.patch_size + self.merge_type = config.merge_type + + self.patch_embed = MoonVision3dPatchEmbed( + out_dim=config.hidden_size, + patch_size=config.patch_size, + pos_emb_height=config.init_pos_emb_height, + pos_emb_width=config.init_pos_emb_width, + pos_emb_time=config.init_pos_emb_time, + pos_emb_type=config.pos_emb_type, + ) + + self.encoder = MoonViT3dEncoder( + hidden_dim=config.hidden_size, + num_layers=config.num_hidden_layers, + block_cfg={ + "num_heads": config.num_attention_heads, + "hidden_dim": config.hidden_size, + "mlp_dim": config.intermediate_size, + "activation": PytorchGELUTanh(), + "attn_bias": True, + "attn_implementation": config._attn_implementation, + }, + video_attn_type=config.video_attn_type, + ) + + def forward(self, pixel_values: torch.Tensor, grid_thws: torch.Tensor) -> torch.Tensor: + """ + Args: + pixel_values (torch.Tensor): The input pixel values. + grid_thws (torch.Tensor): Temporal, height and width. + Returns: + torch.Tensor: The output tokens. + """ + # grid_thws = grid_thws.to('cpu') + assert grid_thws.ndim == 2, f"grid_thws should be 2D, got {grid_thws.ndim}" + assert grid_thws.size(1) == 3, f"No support for thw: {grid_thws}" + hidden_states = self.patch_embed(pixel_values, grid_thws) + hidden_states = self.encoder(hidden_states, grid_thws) + if self.merge_type == "sd2_tpool": # spatial downsampling 2x with temporal pooling all + hidden_states = tpool_patch_merger(hidden_states, grid_thws, merge_kernel_size=self.merge_kernel_size) + else: + raise NotImplementedError(f"Not support {self.merge_type}") + + return hidden_states + + +# ============================================================================ +# MM Projector Helper Classes (from mm_projector/modeling_mm_projectors.py) +# ============================================================================ + + +class IdentityMap(nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x, *args, **kwargs): + return x + + +class MLP(nn.Module): + def __init__(self, config): + super().__init__() + # TODO, use faster LayerNorm + self.pre_norm = nn.LayerNorm(config.mm_hidden_size) + self.proj = nn.Sequential( + nn.Linear(config.mm_hidden_size, config.hidden_size), + nn.GELU(), + nn.Linear(config.hidden_size, config.hidden_size), + ) + + def forward(self, x, *args, **kwargs): + assert isinstance(x, list | tuple), f"x is not a list or tuple: {type(x)}" + lengths = [item.shape[0] for item in x] + x = torch.cat(x, dim=0) + x = self.pre_norm(x) + x = self.proj(x) + x = torch.split(x, lengths, dim=0) + + return x + + +class PatchMergerMLP(nn.Module): + def __init__(self, config): + super().__init__() + eps = config.projector_ln_eps + self.hidden_size = config.mm_hidden_size * (config.merge_kernel_size[0] * config.merge_kernel_size[1]) + self.pre_norm = nn.LayerNorm(config.mm_hidden_size, eps=eps) + self.proj = nn.Sequential( + nn.Linear(self.hidden_size, self.hidden_size), + nn.GELU(), + nn.Linear(self.hidden_size, config.hidden_size), + ) + + def forward(self, x, *args, **kwargs): + if isinstance(x, list) or isinstance(x, tuple): + x = [self.proj(self.pre_norm(item).view(item.shape[0], -1)) for item in x] + else: + # B, N, N_k, C = x.shape + B = x.shape[0] + x = self.proj(self.pre_norm(x).view(B, -1, self.hidden_size)) + return x + + +class KimiK25PreTrainedModel(PreTrainedModel): + config_class = KimiK25Config + base_model_prefix = "model" + _no_split_modules = [ + "MoonViT3dPretrainedModel", + "MoonViTEncoderLayer", + "DeepseekDecoderLayer", + "PatchMergerMLP", + ] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn_2 = True + _supports_sdpa = False + + def _init_weights(self, module): + # important: this ported version of Llava isn't meant for training from scratch - only + # inference and fine-tuning - so the proper init weights code has been removed - the original codebase + # https://github.com/haotian-liu/LLaVA/tree/main/llava should serve for that purpose + std = ( + self.config.initializer_range + if hasattr(self.config, "initializer_range") + else self.config.text_config.initializer_range + ) + + if hasattr(module, "class_embedding"): + module.class_embedding.data.normal_(mean=0.0, std=std) + + if isinstance(module, (nn.Linear, nn.Conv2d)): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +class VisionTowerConfig(PretrainedConfig): + model_type = "moonvit3d" + + def __init__(self, config: KimiK25Config, **kwargs): + super().__init__(**kwargs) + self.patch_size = config.patch_size + self.init_pos_emb_height = config.init_pos_emb_height + self.init_pos_emb_width = config.init_pos_emb_width + self.init_pos_emb_time = config.init_pos_emb_time + self.pos_emb_type = config.pos_emb_type + self.num_attention_heads = config.vt_num_attention_heads + self.num_hidden_layers = config.vt_num_hidden_layers + self.hidden_size = config.vt_hidden_size + self.intermediate_size = config.vt_intermediate_size + self.merge_kernel_size = config.merge_kernel_size + self.video_attn_type = config.video_attn_type + self.merge_type = config.merge_type + self._attn_implementation = config._attn_implementation + + +class ProjectorConfig: + def __init__(self, config: KimiK25Config): + self.mm_projector_type = config.mm_projector_type + self.mm_hidden_size = config.mm_hidden_size + self.hidden_size = config.text_hidden_size + self.merge_kernel_size = config.merge_kernel_size + self.projector_hidden_act = config.projector_hidden_act + self.projector_ln_eps = config.projector_ln_eps + + +class QEffKimiK25EncoderWrapper(nn.Module): + def __init__(self, model): + super().__init__() + self.model = model + self.config = self.model.config + # _orig_apply_rope = _kimi_module.apply_rope + # _kimi_module.apply_rope = _apply_rope_real + # _kimi_module = _sys.modules[type(model).__module__] + + # Restore original apply_rope and attention functions + # _kimi_module.apply_rope = _orig_apply_rope + # _kimi_module.VL_VISION_ATTENTION_FUNCTIONS.update(_orig_attn_functions) + + # _orig_attn_functions = _kimi_module.VL_VISION_ATTENTION_FUNCTIONS.copy() + # _kimi_module.VL_VISION_ATTENTION_FUNCTIONS["eager"] = _full_attention_forward + + def get_submodules_for_export(self) -> Type[nn.Module]: + """ + Return the set of class used as the repeated layer across the model for subfunction extraction. + Notes: + This method should return the *class object* (not an instance). + Downstream code can use this to find/build subfunctions for repeated blocks. + """ + return {self.model.vision_model.model.layers[0].__class__} + # return {self.model.layers[0].__class__} + + def forward_only_image(self, pixel_values: torch.Tensor, grid_thws: torch.Tensor) -> list[torch.Tensor]: + """ + Run only the vision tower and mm_projector to extract image embeddings. + + Args: + pixel_values: Preprocessed image pixel values. + grid_thws: Grid temporal/height/width info for the images. + + Returns: + image_embeds: List of projected image embedding tensors, one per image. + """ + image_features = self._extract_image_features(pixel_values, grid_thws) + if self.mm_projector: + image_features = self.mm_projector(image_features) + return image_features + + def forward( + self, pixel_values: torch.Tensor, h_shape: torch.Tensor, w_shape: torch.Tensor + ) -> torch.Tensor: + """ + ONNX-exportable forward that runs only the vision tower and mm_projector. + Uses h_shape and w_shape (int64 ones tensors of length h and w respectively) + to derive spatial dimensions via .shape[0], enabling dynamic axis export. + + Args: + pixel_values: Preprocessed image pixel values (num_patches, channels*patch_h*patch_w). + h_shape: int64 ones tensor of shape (h,) encoding number of patch rows. + w_shape: int64 ones tensor of shape (w,) encoding number of patch columns. + + Returns: + image_embeds: Projected image embeddings as a single tensor. + """ + _kimi_module = _sys.modules[type(self).__module__] + _get_rope_shape = _kimi_module.get_rope_shape + + h_shape = h_shape.to(pixel_values.device) + w_shape = w_shape.to(pixel_values.device) + + # Keep them in ONNX graph + dummy = (h_shape.float().sum() + w_shape.float().sum()) * 0.0 + pixel_values = pixel_values + dummy + + h = h_shape.shape[0] + w = w_shape.shape[0] + + target_dtype = self.model.vision_tower.patch_embed.proj.weight.dtype + pixel_values = pixel_values.to(target_dtype) + + # --- Patch embedding --- + x = self.model.vision_tower.patch_embed.proj(pixel_values).view(pixel_values.size(0), -1) + + # Positional embedding (single image, t=1) + pos_emb_module = self.model.vision_tower.patch_embed.pos_emb + pos_emb_2d = _get_rope_shape( + pos_emb_module.weight, + interpolation_mode=pos_emb_module.interpolation_mode, + shape=(h, w), + ) + x = x + pos_emb_2d + + # --- Encoder --- + # For single image with t=1: cu_seqlens = [0, h*w] + num_tokens = h * w + cu_seqlens = torch.zeros(2, dtype=torch.int32, device=x.device) + cu_seqlens[1] = num_tokens + max_seqlen = num_tokens + + # RoPE frequencies for single image (stacked real/imag to avoid complex tensors in ONNX) + # Shape: (2, num_tokens, dim//2) where [0]=cos, [1]=sin + rope_2d = self.model.vision_tower.encoder.rope_2d + rope_2d._ensure_precomputed_freqs(x.device) + freqs_cos = rope_2d.freqs_cos[:h, :w].reshape(-1, rope_2d.dim // 2) + freqs_sin = rope_2d.freqs_sin[:h, :w].reshape(-1, rope_2d.dim // 2) + freqs_cis = torch.stack([freqs_cos, freqs_sin], dim=0) + + encoder_dtype = self.model.vision_tower.encoder.blocks[0].wqkv.weight.dtype + x = x.to(encoder_dtype) + freqs_cis = freqs_cis.to(encoder_dtype) + + for block in self.model.vision_tower.encoder.blocks: + x = block(x, cu_seqlens, max_seqlen, rope_freqs_cis=freqs_cis) + + final_ln_dtype = self.model.vision_tower.encoder.final_layernorm.weight.dtype + x = self.model.vision_tower.encoder.final_layernorm(x.to(final_ln_dtype)) + + # --- tpool_patch_merger (single image, t=1) --- + merge_kernel_size = self.model.vision_tower.merge_kernel_size + kernel_height, kernel_width = merge_kernel_size + d_model = x.size(-1) + new_height = h // kernel_height + new_width = w // kernel_width + reshaped = x.view(1, new_height, kernel_height, new_width, kernel_width, d_model) + reshaped = reshaped.permute(0, 1, 3, 2, 4, 5).contiguous().mean(dim=0) + merged = reshaped.view(new_height * new_width, kernel_height * kernel_width, -1) + + # --- mm_projector (PatchMergerMLP on single tensor) --- + pre_norm_dtype = self.model.mm_projector.pre_norm.weight.dtype + merged = merged.to(pre_norm_dtype) + image_embeds = self.model.mm_projector.proj(self.model.mm_projector.pre_norm(merged).view(merged.shape[0], -1)) + return image_embeds + + +class QEffKimiK25DecoderWrapper(nn.Module): + def __init__(self, model): + super().__init__() + self.model = model + self.language_model = self.model.language_model + self.lm_head = self.language_model.lm_head + self.config = self.model.config + + def get_submodules_for_export(self) -> Type[nn.Module]: + """ + Return the set of class used as the repeated layer across the model for subfunction extraction. + Notes: + This method should return the *class object* (not an instance). + Downstream code can use this to find/build subfunctions for repeated blocks. + """ + return {self.model.layers[0].__class__} + + def forward( + self, + input_ids: torch.LongTensor = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + image_embeds: Optional[torch.FloatTensor] = None, + image_idx: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + compressed_kvs: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + batch_index: Optional[torch.LongTensor] = None, + comp_ctx_lengths: Optional[List[int]] = None, + use_cache: Optional[bool] = None, + labels: Optional[torch.LongTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + num_logits_to_keep: Optional[int] = None, + **kwargs, + ) -> Tuple: + del labels, output_attentions, output_hidden_states, return_dict, cache_position, num_logits_to_keep + if inputs_embeds is None: + inputs_embeds = self.model.get_input_embeddings()(input_ids) + + if image_idx is None: + image_idx = torch.zeros((1, 1), dtype=torch.int64, device=inputs_embeds.device) + + if image_embeds is not None and input_ids is not None and input_ids.shape[1] != 1: + if image_embeds.dim() == 2: + image_embeds = image_embeds.unsqueeze(0) + + image_embeds = image_embeds.to(device=inputs_embeds.device, dtype=inputs_embeds.dtype) + _, _, hidden_size = inputs_embeds.shape + selected = input_ids == self.config.media_placeholder_token_id + indices1 = selected.to(torch.int64).cumsum(1) - 1 + indices1 = torch.where(indices1 != -1, indices1 + image_idx.to(indices1.device), indices1) + indices0 = torch.arange(selected.shape[0], device=selected.device).view(-1, 1) + safe_indices1 = torch.where(indices1 < 0, torch.zeros_like(indices1), indices1) + image_features_expanded = image_embeds.reshape(-1, hidden_size).unsqueeze(0)[indices0, safe_indices1] + image_input_embeds = torch.where(selected.unsqueeze(-1), image_features_expanded, inputs_embeds) + inputs_embeds = torch.where(input_ids.shape[1] == torch.tensor(1), inputs_embeds, image_input_embeds) + image_idx = (indices1.max() + 1).reshape(1, 1) + + outputs = self.model.language_model( + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + compressed_kvs=compressed_kvs, + past_key_values=past_key_values, + batch_index=batch_index, + comp_ctx_lengths=comp_ctx_lengths, + use_cache=True if use_cache is None else use_cache, + **kwargs, + ) + + hidden_states = outputs[0] + logit_index = position_ids.to(torch.int32).argmax(1, keepdim=True) + max_index = hidden_states.shape[1] - 1 + safe_logit_index = torch.clamp(logit_index, min=0, max=max_index) + hidden_states = hidden_states[torch.arange(position_ids.shape[0]).view(-1, 1), safe_logit_index] + + if hidden_states.shape[-1] == self.lm_head.in_features: + logits = self.lm_head(hidden_states).float() + else: + logits = hidden_states.float() + + output_compressed_kvs = getattr(outputs, "compressed_kvs", compressed_kvs) + output_past_key_values = getattr(outputs, "past_key_values", None) + return logits, image_embeds, image_idx, output_compressed_kvs, output_past_key_values + + +# ref https://github.com/huggingface/transformers/blob/78b2929c0554b79e0489b451ce4ece14d265ead2/src/transformers/models/llava/modeling_llava.py#L240 +class QEffKimiK25ForConditionalGeneration(KimiK25PreTrainedModel): + def __init__(self, config: KimiK25Config): + super().__init__(config) + + vt_config = VisionTowerConfig(config.vision_config) + self.vision_tower = MoonViT3dPretrainedModel(vt_config) + + proj_config = ProjectorConfig(config.vision_config) + if proj_config.mm_projector_type == "identity": + self.mm_projector = IdentityMap() + elif proj_config.mm_projector_type == "mlp": + self.mm_projector = MLP(proj_config) + elif proj_config.mm_projector_type == "patchmerger": + self.mm_projector = PatchMergerMLP(proj_config) + else: + raise ValueError(f"Unsupported mm_projector_type: {proj_config.mm_projector_type}") + + self.language_model = QEffDeepseekV3ForCausalLM(config.text_config) + self.post_init() + + if hasattr(self.language_model, "dtype"): + target_dtype = self.language_model.dtype + self.vision_tower = self.vision_tower.to(dtype=target_dtype) + self.mm_projector = self.mm_projector.to(dtype=target_dtype) + + def get_qeff_vision_encoder(self): + return QEffKimiK25EncoderWrapper(self) + + def get_qeff_language_decoder(self): + return QEffKimiK25DecoderWrapper(self) + + def get_input_embeddings(self): + return self.language_model.get_input_embeddings() + + def set_input_embeddings(self, value): + self.language_model.set_input_embeddings(value) + + def get_output_embeddings(self): + return self.language_model.get_output_embeddings() + + def set_output_embeddings(self, new_embeddings): + self.language_model.set_output_embeddings(new_embeddings) + + def set_decoder(self, decoder): + self.language_model.set_decoder(decoder) + + def get_decoder(self): + return self.language_model.get_decoder() + + def tie_weights(self): + return self.language_model.tie_weights() + + def resize_token_embeddings(self, new_num_tokens: int | None = None, pad_to_multiple_of=None) -> nn.Embedding: + model_embeds = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of) + # update vocab size + self.config.text_config.vocab_size = model_embeds.num_embeddings + self.vocab_size = model_embeds.num_embeddings + return model_embeds + + def _merge_input_ids_with_image_features( + self, + image_features: list[torch.Tensor], + inputs_embeds: torch.Tensor, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + labels: torch.Tensor | None = None, + ): + """ + Args: + image_features (:obj:`torch.Tensor` of shape :obj:`(num_image_tokens, embed_dim)`): + The image features to merge with the input embeddings. + inputs_embeds (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length, embed_dim)`): + The input embeddings. + input_ids (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`): + The input ids. + attention_mask (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`): + The attention mask. + labels (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`, *optional*): + The labels. + """ + _, embed_dim = image_features[0].shape + feature_lengths = [x.shape[0] for x in image_features] + image_features = torch.cat(image_features, dim=0) + + image_token_index: int = self.config.media_placeholder_token_id + pad_token_id: int = self.config.pad_token_id + ignore_index: int = self.config.ignore_index + + batch_size, sequence_length = input_ids.shape + left_padding = not torch.sum(input_ids[:, -1] == torch.tensor(pad_token_id)) + + # 1. Create a mask to know where special image tokens are + _token_occupation_table = torch.ones_like(input_ids.flatten()) + _token_occupation_table[input_ids.flatten() == image_token_index] = torch.tensor( + feature_lengths, dtype=torch.long, device=input_ids.device + ) + _token_occupation_table = _token_occupation_table.reshape(input_ids.shape) + + max_embed_dim = _token_occupation_table.sum(-1).max().item() + assert max_embed_dim >= sequence_length, ( + f"The maximum embedding dimension ({max_embed_dim}) is less than the sequence length ({sequence_length})" + ) + batch_indices, non_image_indices = torch.where(input_ids != image_token_index) + + # 2. Compute the positions where text should be written + # Calculate new positions for text tokens in merged image-text sequence. + new_token_positions = torch.cumsum(_token_occupation_table, -1) - 1 + nb_image_pad = max_embed_dim - 1 - new_token_positions[:, -1] + if left_padding: + new_token_positions += nb_image_pad[:, None] # offset for left padding + text_to_overwrite = new_token_positions[batch_indices, non_image_indices] + + # 3. Create the full embedding, already padded to the maximum position + final_embedding = torch.zeros( + batch_size, + max_embed_dim, + embed_dim, + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) + final_attention_mask = torch.zeros( + batch_size, max_embed_dim, dtype=attention_mask.dtype, device=inputs_embeds.device + ) + if labels is not None: + final_labels = torch.full( + (batch_size, max_embed_dim), + ignore_index, + dtype=input_ids.dtype, + device=input_ids.device, + ) + # In case the Vision model or the Language model has been offloaded to CPU, we need to manually + # set the corresponding tensors into their correct target device. + target_device = inputs_embeds.device + batch_indices, non_image_indices, text_to_overwrite = ( + batch_indices.to(target_device), + non_image_indices.to(target_device), + text_to_overwrite.to(target_device), + ) + attention_mask = attention_mask.to(target_device) + + # 4. Fill the embeddings based on the mask. + final_embedding[batch_indices, text_to_overwrite] = inputs_embeds[batch_indices, non_image_indices] + final_attention_mask[batch_indices, text_to_overwrite] = attention_mask[batch_indices, non_image_indices] + if labels is not None: + final_labels[batch_indices, text_to_overwrite] = labels[batch_indices, non_image_indices] + + # 5. Fill the embeddings corresponding to the images. Anything that is not `text_positions` needs filling (#29835) + image_to_overwrite = torch.full( + (batch_size, max_embed_dim), True, dtype=torch.bool, device=inputs_embeds.device + ) + image_to_overwrite[batch_indices, text_to_overwrite] = False + image_to_overwrite &= image_to_overwrite.cumsum(-1) - 1 >= nb_image_pad[:, None].to(target_device) + + if image_to_overwrite.sum() != image_features.shape[:-1].numel(): + raise ValueError( + f"The input provided to the model are wrong. The number of image tokens is {image_to_overwrite.sum()} while" + f" the number of image features given to the model is {image_features.shape[:-1].numel()}. " + "This prevents correct indexing and breaks batch generation." + ) + + final_embedding[image_to_overwrite] = image_features.contiguous().reshape(-1, embed_dim).to(target_device) + final_attention_mask |= image_to_overwrite + position_ids = (final_attention_mask.cumsum(-1) - 1).masked_fill_((final_attention_mask == 0), 1) + + # 6. Mask out the embedding at padding positions, as we later use the past_key_value value to determine the non-attended tokens. + batch_indices, pad_indices = torch.where(input_ids == pad_token_id) + indices_to_mask = new_token_positions[batch_indices, pad_indices] + + final_embedding[batch_indices, indices_to_mask] = 0 + + if labels is None: + final_labels = None + + return final_embedding, final_attention_mask, final_labels, position_ids + + def _extract_image_features(self, pixel_values: torch.Tensor, grid_thws: torch.Tensor) -> list[torch.Tensor]: + """ + Args: + pixel_values (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, num_channels, height, width)`): + The pixel values of the images processed by image processor. + grid_thws (:obj:`torch.Tensor` of shape :obj:`(batch_size, 3)`): + The grid, height, width of the images. + Returns: + selected_image_feature (:obj:`torch.FloatTensor` of shape :obj:`(num_image_tokens, embed_dim)`): + The selected image features to use as input to the projector head. + """ + + target_dtype = self.vision_tower.patch_embed.proj.weight.dtype + pixel_values = pixel_values.to(target_dtype) + + image_features = self.vision_tower(pixel_values, grid_thws) + return image_features + + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | list[torch.FloatTensor] | None = None, + grid_thws: torch.Tensor | None = None, + # h_shape: torch.Tensor | None = None, + # w_shape: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + compressed_kvs: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + ) -> tuple | LlavaCausalLMOutputWithPast: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + ```""" + assert self.vision_tower is not None, "vision_tower is not loaded" + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if inputs_embeds is None: + # 1. Extra the input embeddings + inputs_embeds = self.get_input_embeddings()(input_ids) + + # 2. Merge text and images + if pixel_values is not None and len(pixel_values) > 0 and input_ids.shape[1] != 1: + image_features = self._extract_image_features(pixel_values, grid_thws) + if self.mm_projector: + image_features = self.mm_projector(image_features) + + inputs_embeds = inputs_embeds.to(image_features[0].dtype) # num_tokens, embed_dim + inputs_embeds, attention_mask, labels, position_ids = self._merge_input_ids_with_image_features( + image_features, + inputs_embeds, + input_ids, + attention_mask, + labels, + ) + + # In case input_ids.shape[1] == 1 & pixel_values==None & past_key_values != None, we are in the case of + # generation with cache + elif past_key_values is not None and pixel_values is not None and input_ids.shape[1] == 1: + # Retrieve the first layer to inspect the logits and mask out the hidden states + # that are set to 0 + first_layer_past_key_value = past_key_values[0][0][:, :, :, 0] + + # Sum all dimensions of head_dim (-2) to avoid random errors such as: https://github.com/huggingface/transformers/pull/28032#issuecomment-1863691941 + batch_index, non_attended_tokens = torch.where(first_layer_past_key_value.float().sum(-2) == 0) + + # Get the target length + target_length = input_ids.shape[1] + past_length = first_layer_past_key_value.shape[-1] + + extended_attention_mask = torch.ones( + (attention_mask.shape[0], past_length), + dtype=attention_mask.dtype, + device=attention_mask.device, + ) + + # Filter out only the tokens that can be un-attended, this can happen + # if one uses Llava + Fused modules where the cache on the + # first iteration is already big enough, or if one passes custom cache + valid_indices = non_attended_tokens < extended_attention_mask.size(-1) + new_batch_index = batch_index[valid_indices] + new_non_attended_tokens = non_attended_tokens[valid_indices] + + # Zero-out the places where we don't need to attend + extended_attention_mask[new_batch_index, new_non_attended_tokens] = 0 + + attention_mask = torch.cat((extended_attention_mask, attention_mask[:, -target_length:]), dim=1) + position_ids = torch.sum(attention_mask, dim=1).unsqueeze(-1) - 1 + + outputs = self.language_model( + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + compressed_kvs=compressed_kvs, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + logits = outputs[0] + + loss = None + if labels is not None: + # Shift so that tokens < n predict n + if attention_mask is not None: + shift_attention_mask = attention_mask[..., 1:] + shift_logits = logits[..., :-1, :][shift_attention_mask.to(logits.device) != 0].contiguous() + shift_labels = labels[..., 1:][shift_attention_mask.to(labels.device) != 0].contiguous() + else: + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = nn.CrossEntropyLoss() + loss = loss_fct( + shift_logits.view(-1, shift_logits.size(-1)), + shift_labels.view(-1).to(shift_logits.device), + ) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return LlavaCausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + inputs_embeds=None, + pixel_values=None, + grid_thws=None, + attention_mask=None, + **kwargs, + ): + if past_key_values is not None: + if isinstance(past_key_values, Cache): + cache_length = past_key_values.get_seq_length() + past_length = getattr(past_key_values, "seen_tokens", cache_length) + else: + cache_length = past_length = past_key_values[0][0].shape[2] + + # Keep only the unprocessed tokens: + # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where + # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as + # input) + if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]: + input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :] + # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard + # input_ids based on the past_length. + elif past_length < input_ids.shape[1]: + input_ids = input_ids[:, past_length:] + # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens. + elif self.config.media_placeholder_token_id in input_ids: + input_ids = input_ids[:, input_ids.shape[1] - 1 :] + # If the cache has seen more tokens than it can hold, then the cache has a size limit. Let's discard the + # older attention values, as their corresponding values are not part of the input. + if cache_length < past_length and attention_mask is not None: + attention_mask = attention_mask[:, -(cache_length + input_ids.shape[1]) :] + + position_ids = kwargs.get("position_ids", None) + if attention_mask is not None and position_ids is None: + # create position_ids on the fly for batch generation + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + if past_key_values: + position_ids = position_ids[:, -input_ids.shape[1] :] + + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step + if inputs_embeds is not None and past_key_values is None: + model_inputs = {"inputs_embeds": inputs_embeds} + else: + model_inputs = {"input_ids": input_ids} + + model_inputs.update( + { + "position_ids": position_ids, + "past_key_values": past_key_values, + "use_cache": kwargs.get("use_cache"), + "attention_mask": attention_mask, + "pixel_values": pixel_values, + "grid_thws": grid_thws, + } + ) + return model_inputs + + def _reorder_cache(self, *args, **kwargs): + return self.language_model._reorder_cache(*args, **kwargs) + + def get_output_names(self, kv_offload: bool = False): + vision_output_names = ["image_embeds"] + lang_output_names = ["logits"] + + mla_absorption = getattr(self.language_model, "mla_absorption", None) + if mla_absorption is not None: + cache_compressed = mla_absorption.get("cache_compressed", False) + else: + cache_compressed = False + + if cache_compressed: + for i in range(self.language_model.config.num_hidden_layers): + lang_output_names.append(f"compressed_kv.{i}_RetainedState") + lang_output_names.append(f"k_pe.{i}_RetainedState") + else: + for i in range(self.language_model.config.num_hidden_layers): + for kv in ["key", "value"]: + lang_output_names.append(f"past_{kv}.{i}_RetainedState") + + output_names = {} + if kv_offload: + output_names["vision"] = vision_output_names + output_names["lang"] = lang_output_names + else: + return lang_output_names + return output_names + + def get_dummy_inputs( + self, + comp_ctx_lengths: Optional[List[int]] = None, + kv_offload: bool = False, + continuous_batching: bool = False, + **kwargs, + ): + prefill_seq_len = kwargs.get("prefill_seq_len") + if prefill_seq_len is None: + prefill_seq_len = constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN + prefill_seq_len = int(prefill_seq_len) + + bs: int = constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE + fbs: int = constants.ONNX_EXPORT_EXAMPLE_FBS + + model_path = Path( + "/home/huggingface_hub/models--moonshotai--Kimi-K2.5/snapshots/4d01dfe0332d63057c186e0b262165819efb6611" + ) + processor = AutoProcessor.from_pretrained(str(model_path), trust_remote_code=True) + image_url = "https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png" + image = Image.open(BytesIO(requests.get(image_url).content)).convert("RGB") + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": image}, + {"type": "text", "text": "Tell me about yourself."}, + ], + } + ] + inputs = processor( + messages=messages, + add_generation_prompt=True, + tokenize=False, + return_tensors="pt", + ) + inputs = {k: v.to(self.language_model.device) if hasattr(v, "to") else v for k, v in inputs.items()} + + # Build h_shape and w_shape from grid_thws (single image, t=1) + grid_thws_val = inputs["grid_thws"] + h_val = int(grid_thws_val[0, 1].item()) + w_val = int(grid_thws_val[0, 2].item()) + h_shape_tensor = torch.ones(h_val, dtype=torch.int64, device=self.language_model.device) + w_shape_tensor = torch.ones(w_val, dtype=torch.int64, device=self.language_model.device) + + vision_inputs = { + "pixel_values": inputs["pixel_values"], + "h_shape": h_shape_tensor, + "w_shape": w_shape_tensor, + } + + lang_inputs = { + "input_ids": torch.zeros((bs, prefill_seq_len), dtype=torch.int64), + "position_ids": torch.arange(prefill_seq_len, dtype=torch.int64).view(1, prefill_seq_len).repeat(bs, 1), + } + + mla_absorption = getattr(self.language_model, "mla_absorption", None) + if mla_absorption is not None: + cache_compressed = mla_absorption.get("cache_compressed", False) + else: + cache_compressed = False + + pkv_cache = self.language_model.get_dummy_pkv_cache( + config=self.language_model.config, + batch_size=fbs if continuous_batching else bs, + seq_len=prefill_seq_len, + ) + + if cache_compressed: + lang_inputs["compressed_kvs"] = [[] for _ in range(self.language_model.config.num_hidden_layers)] + for i in range(self.language_model.config.num_hidden_layers): + lang_inputs["compressed_kvs"][i].append( + torch.zeros(pkv_cache[0][0].shape, dtype=self.language_model.config.torch_dtype) + ) + lang_inputs["compressed_kvs"][i].append( + torch.zeros(pkv_cache[0][1].shape, dtype=self.language_model.config.torch_dtype) + ) + else: + lang_inputs["past_key_values"] = [[] for _ in range(self.language_model.config.num_hidden_layers)] + for i in range(self.language_model.config.num_hidden_layers): + lang_inputs["past_key_values"][i].append( + torch.zeros(pkv_cache[0][0].shape, dtype=self.language_model.config.torch_dtype) + ) + lang_inputs["past_key_values"][i].append( + torch.zeros(pkv_cache[0][1].shape, dtype=self.language_model.config.torch_dtype) + ) + + if continuous_batching: + lang_inputs["batch_index"] = torch.arange(bs).view(bs, 1) + + if comp_ctx_lengths is not None: + lang_inputs["comp_ctx_lengths"] = torch.randint(0, 100, (40,), dtype=torch.int64) + + inputs = {} + if kv_offload: + inputs["vision"] = vision_inputs + inputs["lang"] = lang_inputs + else: + lang_inputs.pop("image_embeds") + inputs = {**vision_inputs, **lang_inputs} + + return inputs + + def get_onnx_dynamic_axes( + self, comp_ctx_lengths: Optional[List[int]] = None, kv_offload: bool = False, continuous_batching: bool = False + ): + vision_dynamic_axes = {} + lang_dynamic_axes = {} + lang_dynamic_axes["input_ids"] = {0: "batch_size", 1: "seq_len"} + lang_dynamic_axes["position_ids"] = {0: "batch_size", 1: "seq_len"} + lang_dynamic_axes["image_embeds"] = {0: "num_image_tokens"} + if continuous_batching: + lang_dynamic_axes["batch_index"] = {0: "batch_size"} + vision_dynamic_axes = { + "pixel_values": {0: "num_patches"}, + "h_shape": {0: "h"}, + "w_shape": {0: "w"}, + "image_embeds": {0: "num_image_tokens"}, + } + + mla_absorption = getattr(self.language_model, "mla_absorption", None) + if mla_absorption is not None: + cache_compressed = mla_absorption.get("cache_compressed", False) + else: + cache_compressed = False + + if cache_compressed: + for i in range(self.language_model.config.num_hidden_layers): + lang_dynamic_axes[f"compressed_kv.{i}"] = {0: "batch_size", 2: "ctx_len"} + lang_dynamic_axes[f"k_pe.{i}"] = {0: "batch_size", 2: "ctx_len"} + else: + for i in range(self.language_model.config.num_hidden_layers): + for kv in ["key", "value"]: + lang_dynamic_axes[f"past_{kv}.{i}"] = {0: "batch_size", 2: "ctx_len"} + + if comp_ctx_lengths is not None: + lang_dynamic_axes["comp_ctx_lengths"] = {0: "comp_ctx_lengths"} + + dynamic_axes = {} + if kv_offload: + dynamic_axes["vision"] = vision_dynamic_axes + dynamic_axes["lang"] = lang_dynamic_axes + else: + lang_dynamic_axes.pop("image_embeds") + dynamic_axes = {**vision_dynamic_axes, **lang_dynamic_axes} + return dynamic_axes + + def get_specializations( + self, + batch_size: int, + prefill_seq_len: int, + ctx_len: int, + kv_offload: bool = False, + continuous_batching: bool = False, + kv_cache_batch_size: Optional[int] = None, + full_batch_size: Optional[int] = None, + **compiler_options, + ): + comp_ctx_lengths_prefill = compiler_options.pop("comp_ctx_lengths_prefill", None) + comp_ctx_lengths_decode = compiler_options.pop("comp_ctx_lengths_decode", None) + num_patches = compiler_options.pop("num_patches", None) + h = compiler_options.pop("h", None) + w = compiler_options.pop("w", None) + num_image_tokens = compiler_options.pop("num_image_tokens", None) + + prefill_seq_len = prefill_seq_len if prefill_seq_len else 32 # constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN + ctx_len = ctx_len if ctx_len else 32 # constants.ONNX_EXPORT_EXAMPLE_CTX_LEN + + vision = [ + { + "num_patches": 2400, # num_patches + "h": 30, # h + "w": 80, # w + "num_image_tokens": 600, # num_image_tokens + } + ] + + if comp_ctx_lengths_prefill is not None: + lang = [] + + for i in range(0, len(comp_ctx_lengths_prefill)): + lang_prefill = { + "batch_size": 1 if continuous_batching else batch_size, + "seq_len": prefill_seq_len, + "ctx_len": ctx_len, + "num_image_tokens": 600, # num_image_tokens + } + if continuous_batching: + lang_prefill["full_batch_size"] = kv_cache_batch_size + else: + lang_prefill["batch_size"] = kv_cache_batch_size + if full_batch_size: + lang_prefill["full_batch_exec_size"] = full_batch_size + + lang.append(lang_prefill) + + for i in range(0, len(comp_ctx_lengths_decode)): + lang_decode = { + "batch_size": full_batch_size if continuous_batching else batch_size, + "seq_len": "1", + "ctx_len": ctx_len, + "num_image_tokens": 600, # num_image_tokens + } + + if continuous_batching: + lang_decode["full_batch_size"] = kv_cache_batch_size + else: + lang_decode["batch_size"] = kv_cache_batch_size + + lang.append(lang_decode) + + else: + lang_prefill = { + "batch_size": 1 if continuous_batching else batch_size, + "seq_len": prefill_seq_len, + "ctx_len": ctx_len, + "num_image_tokens": 600, # num_image_tokens + } + if continuous_batching: + lang_prefill["full_batch_size"] = kv_cache_batch_size + else: + lang_prefill["batch_size"] = kv_cache_batch_size + if full_batch_size: + lang_prefill["full_batch_exec_size"] = full_batch_size + + lang_decode = { + "batch_size": full_batch_size if continuous_batching else batch_size, + "seq_len": 1, + "ctx_len": ctx_len, + "num_image_tokens": 600, # num_image_tokens + } + + if continuous_batching: + lang_decode["full_batch_size"] = kv_cache_batch_size + else: + lang_decode["batch_size"] = kv_cache_batch_size + + lang = [lang_prefill, lang_decode] + + specializations = {} + + if kv_offload: + specializations["vision"] = vision + specializations["lang"] = lang + return specializations, compiler_options + else: + lang[0].pop("vision_size") + lang[1].pop("vision_size") + return lang, compiler_options diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index 2cabf04526..7afe944e62 100755 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -1983,6 +1983,8 @@ def compile( custom_io_vision = {} target_dtype = getattr(self.model.config, "torch_dtype", torch.float32) + if target_dtype == torch.bfloat16 and constants.DEFAULT_AIC_HW_VERSION != "ai200": + target_dtype = torch.float16 kv_cache_dtype = "mxint8" if mxint8_kv_cache else CUSTOM_IO_DTYPE_MAP[target_dtype] molmo = hasattr(self.model.config, "model_type") and self.model.config.model_type == "molmo" if molmo: diff --git a/examples/kimi_k2/export_kimi_k25_vision.py b/examples/kimi_k2/export_kimi_k25_vision.py index 4d2c77fbf5..361c0c54af 100644 --- a/examples/kimi_k2/export_kimi_k25_vision.py +++ b/examples/kimi_k2/export_kimi_k25_vision.py @@ -333,6 +333,11 @@ def parse_args(): ) parser.add_argument("--skip-processor", action="store_true", help="Skip loading processor.") parser.add_argument("--skip-tokenizer", action="store_true", help="Skip loading tokenizer.") + parser.add_argument( + "--skip-lang-compile", + action="store_true", + help="Compile/export only the vision encoder and skip language decoder compile.", + ) return parser.parse_args() @@ -482,6 +487,7 @@ def main(): mxfp6_matmul=False, mxint8_kv_cache=False, aic_enable_depth_first=False, + skip_lang=args.skip_lang_compile, #skip_vision=True, # Skip vision encoder for text-only inference mos=1, num_patches=2400, # num_patches From 3b278e0f706526077c66332c95a5e7f7380e93dc Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Mon, 22 Jun 2026 18:41:41 +0530 Subject: [PATCH 05/33] update example script Signed-off-by: Mamta Singh --- .../models/deepseek_v3/modeling_deepseek.py | 3 +- .../models/kimi_k25/modeling_kimi_k25.py | 21 +-- examples/kimi_k2/export_kimi_k25_vision.py | 132 +++++------------- .../kimi_k2/export_kimi_k25_vision_4.57.3.py | 1 - 4 files changed, 49 insertions(+), 108 deletions(-) diff --git a/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py b/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py index df07871d45..daddfe048f 100644 --- a/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py +++ b/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py @@ -785,7 +785,6 @@ class QEffDeepseekV3MoE(nn.Module): def __qeff_init__( self, ): - #breakpoint() self.all_gate_proj = torch.nn.Parameter( torch.cat( [_get_linear_weight(exp.gate_proj).T.unsqueeze(0) for exp in self.experts], @@ -803,7 +802,7 @@ def __qeff_init__( dim=0, ) ) - + self.act_fn = self.experts[0].act_fn def moe( diff --git a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py index d7ab06acdf..b72e152860 100644 --- a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py +++ b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py @@ -35,7 +35,6 @@ from transformers.activations import PytorchGELUTanh from transformers.cache_utils import Cache from transformers.configuration_utils import PretrainedConfig -from transformers.modeling_outputs import CausalLMOutputWithPast from transformers.modeling_utils import PreTrainedModel from transformers.models.llava.modeling_llava import LlavaCausalLMOutputWithPast @@ -72,13 +71,17 @@ def wrapper(org, interpolation_mode, shape): @get_rope_shape_decorate - def get_rope_shape(org, interpolation_mode, shape): - return (F.interpolate( - org.permute((2, 0, 1)).unsqueeze(0), - size=shape, - mode=interpolation_mode, - ).squeeze(0).permute((1, 2, 0)).flatten(end_dim=1)) + return ( + F.interpolate( + org.permute((2, 0, 1)).unsqueeze(0), + size=shape, + mode=interpolation_mode, + ) + .squeeze(0) + .permute((1, 2, 0)) + .flatten(end_dim=1) + ) def apply_rope(xq, xk, freqs_cis): @@ -737,9 +740,7 @@ def forward_only_image(self, pixel_values: torch.Tensor, grid_thws: torch.Tensor image_features = self.mm_projector(image_features) return image_features - def forward( - self, pixel_values: torch.Tensor, h_shape: torch.Tensor, w_shape: torch.Tensor - ) -> torch.Tensor: + def forward(self, pixel_values: torch.Tensor, h_shape: torch.Tensor, w_shape: torch.Tensor) -> torch.Tensor: """ ONNX-exportable forward that runs only the vision tower and mm_projector. Uses h_shape and w_shape (int64 ones tensors of length h and w respectively) diff --git a/examples/kimi_k2/export_kimi_k25_vision.py b/examples/kimi_k2/export_kimi_k25_vision.py index 361c0c54af..40e968fdd0 100644 --- a/examples/kimi_k2/export_kimi_k25_vision.py +++ b/examples/kimi_k2/export_kimi_k25_vision.py @@ -12,9 +12,12 @@ import re import sys import tempfile +from io import BytesIO from pathlib import Path +import requests import torch +from PIL import Image from safetensors import safe_open from safetensors.torch import save_file from transformers import AutoConfig, AutoProcessor, AutoTokenizer @@ -23,7 +26,6 @@ from QEfficient import QEFFAutoModelForImageTextToText - MODEL_PATH = Path( "/home/huggingface_hub/models--moonshotai--Kimi-K2.5/snapshots/4d01dfe0332d63057c186e0b262165819efb6611" ) @@ -193,7 +195,9 @@ def _build_layer_subset_config(config, num_vision_layers, num_text_layers, loade return stripped_config, loaded_expert_ids -def _materialize_subset_checkpoint(model_path: Path, temp_model_path: Path, weight_map, allowed_prefixes, loaded_expert_ids): +def _materialize_subset_checkpoint( + model_path: Path, temp_model_path: Path, weight_map, allowed_prefixes, loaded_expert_ids +): expert_index_map = {expert_id: remapped_idx for remapped_idx, expert_id in enumerate(loaded_expert_ids)} shard_to_entries = {} for checkpoint_key, source_shard_name in weight_map.items(): @@ -235,7 +239,6 @@ def _load_layer_subset_model( num_text_layers: int, loaded_expert_ids, num_experts_per_tok: int, - local_files_only: bool, dtype, ): checkpoint_index = json.loads((model_path / "model.safetensors.index.json").read_text()) @@ -261,7 +264,9 @@ def _load_layer_subset_model( (temp_model_path / "model.safetensors.index.json").write_text( json.dumps( { - "metadata": {"total_size": sum((temp_model_path / shard_name).stat().st_size for shard_name in subset_shards)}, + "metadata": { + "total_size": sum((temp_model_path / shard_name).stat().st_size for shard_name in subset_shards) + }, "weight_map": filtered_weight_map, } ) @@ -271,7 +276,6 @@ def _load_layer_subset_model( "config": stripped_config, "trust_remote_code": True, "attn_implementation": "eager", - "local_files_only": local_files_only, "output_loading_info": True, } if dtype is not None: @@ -292,65 +296,39 @@ def _load_layer_subset_model( "Failed to load the stripped Kimi K2.5 checkpoint slice cleanly. " f"missing={missing_keys}, unexpected={unexpected_keys}, mismatched={mismatched_keys}" ) + model.eval() + tokenizer = AutoTokenizer.from_pretrained(str(model_path), trust_remote_code=True) + processor = AutoProcessor.from_pretrained(str(model_path), trust_remote_code=True) + + print(f"Loaded model: {type(model).__name__}") + print(f"Tokenizer vocab size: {tokenizer.vocab_size}") + print(f"Processor type: {type(processor).__name__}") - return model + return model, tokenizer, processor def parse_args(): parser = argparse.ArgumentParser(description="Load Kimi K2.5 with runtime compatibility for transformers==5.5.4.") parser.add_argument("--model-path", type=Path, default=MODEL_PATH) - parser.add_argument("--local-files-only", action="store_true", help="Only read from local cache/files.") parser.add_argument( "--full-model", action="store_true", help="Load the full model. By default, the script loads a small layer subset for faster startup.", ) - parser.add_argument( - "--num-vision-layers", - type=int, - default=NUM_VISION_LAYERS, - help=( - "Load only the first N vision layers. " - f"Defaults to {NUM_VISION_LAYERS} for faster startup. Ignored when --full-model is set." - ), - ) - parser.add_argument( - "--num-text-layers", - type=int, - default=NUM_TEXT_LAYERS, - help=( - "Load only the first N text layers. " - f"Defaults to {NUM_TEXT_LAYERS} for faster startup. Ignored when --full-model is set." - ), - ) + parser.add_argument("--num-vision-layers", type=int, default=NUM_VISION_LAYERS) + parser.add_argument("--num-text-layers", type=int, default=NUM_TEXT_LAYERS) parser.add_argument("--expert-ids", type=_parse_expert_ids, default=LOADED_EXPERT_IDS) parser.add_argument("--num-experts-per-token", type=int, default=NUM_EXPERTS_PER_TOKEN) parser.add_argument( - "--dtype", - choices=["auto", "float32", "float16", "bfloat16"], - default="auto", - help="torch_dtype used by from_pretrained.", - ) - parser.add_argument("--skip-processor", action="store_true", help="Skip loading processor.") - parser.add_argument("--skip-tokenizer", action="store_true", help="Skip loading tokenizer.") - parser.add_argument( - "--skip-lang-compile", - action="store_true", - help="Compile/export only the vision encoder and skip language decoder compile.", + "--image-url", + type=str, + default="https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png", ) + parser.add_argument("--prompt", type=str, default="Describe this image.") + parser.add_argument("--test", action="store_true", help="Validate ONNX output matches PyTorch image-only forward.") return parser.parse_args() -def _dtype_from_arg(dtype_arg: str): - if dtype_arg == "auto": - return None - if dtype_arg == "float16": - return torch.float16 - if dtype_arg == "bfloat16": - return torch.bfloat16 - return torch.float32 - - def main(): args = parse_args() _ensure_torch_fx_import_compatibility() @@ -371,16 +349,13 @@ def main(): "config": config, "trust_remote_code": True, "attn_implementation": "eager", - "local_files_only": args.local_files_only, + "torch_dtype": torch.float32, } - dtype = _dtype_from_arg(args.dtype) - if dtype is not None: - model_kwargs["torch_dtype"] = dtype if args.full_model: - model = kimi_cls.from_pretrained(str(args.model_path), **model_kwargs) + model, tokenizer, processor = kimi_cls.from_pretrained(str(args.model_path), **model_kwargs) elif args.num_vision_layers is not None and args.num_text_layers is not None: - model = _load_layer_subset_model( + model, tokenizer, processor = _load_layer_subset_model( model_path=args.model_path, kimi_cls=kimi_cls, config=config, @@ -388,8 +363,7 @@ def main(): num_text_layers=args.num_text_layers, loaded_expert_ids=args.expert_ids, num_experts_per_tok=args.num_experts_per_token, - local_files_only=args.local_files_only, - dtype=dtype, + dtype=torch.float32, ) print( "Loaded layer subset: " @@ -399,26 +373,10 @@ def main(): ) else: raise ValueError("Pass both --num-vision-layers and --num-text-layers to load a layer subset.") - print(f"Loaded model: {type(model).__name__}") - - if not args.skip_tokenizer: - tokenizer = AutoTokenizer.from_pretrained( - str(args.model_path), trust_remote_code=True, local_files_only=args.local_files_only - ) - print(f"Tokenizer vocab size: {tokenizer.vocab_size}") - if not args.skip_processor: - processor = AutoProcessor.from_pretrained( - str(args.model_path), trust_remote_code=True, local_files_only=args.local_files_only - ) - print(f"Processor type: {type(processor).__name__}") - - mla_absorption = {"cache_compressed": True, "absorption": False, "online": False} - qaic_config = { - "mla_absorption": mla_absorption - } # , "enable_blocking": True, "blocking_mode": "par", "par_num_split": 4, "num_kv_blocks": 8} + qaic_config = {"mla_absorption": {"cache_compressed": True, "absorption": False, "online": False}} - qeff_model = QEFFAutoModelForImageTextToText(model) # , qaic_config=qaic_config) + qeff_model = QEFFAutoModelForImageTextToText(model) skip_vision = False @@ -429,10 +387,10 @@ def main(): # Set skip_vision=True to bypass image processing qeff_model.compile( qaic_config=qaic_config, - prefill_seq_len=32, + prefill_seq_len=1, ctx_len=1024, num_cores=16, - num_devices=2, + num_devices=1, mxfp6_matmul=False, mxint8_kv_cache=False, aic_enable_depth_first=False, @@ -443,7 +401,6 @@ def main(): w=80, # w num_image_tokens=600, # num_image_tokens ) - breakpoint() ## STEP 4: Prepare Text-Only Input # Create a text-only message without any image messages = [ @@ -465,7 +422,6 @@ def main(): ) ## STEP 6: Run Text-Only Inference - streamer = TextStreamer(tokenizer) output = qeff_model.generate(inputs=inputs, device_ids=[0, 1], generation_len=10) ## STEP 7: Display Results @@ -483,12 +439,11 @@ def main(): prefill_seq_len=1, ctx_len=1024, num_cores=16, - num_devices=2, + num_devices=1, mxfp6_matmul=False, mxint8_kv_cache=False, aic_enable_depth_first=False, - skip_lang=args.skip_lang_compile, - #skip_vision=True, # Skip vision encoder for text-only inference + # skip_lang=True, mos=1, num_patches=2400, # num_patches h=30, # h @@ -496,10 +451,7 @@ def main(): num_image_tokens=600, # num_image_tokens ) - breakpoint() ## STEP 4: Prepare Image and Text Input - # Define the image URL to process - image_url = "https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png" image = Image.open(BytesIO(requests.get(args.image_url).content)).convert("RGB") # Create a message with both image and text @@ -507,21 +459,13 @@ def main(): { "role": "user", "content": [ - {"type": "image", "url": image}, - {"type": "text", "text": "Can you describe the image in detail."}, + {"type": "image_url", "image_url": image}, + {"type": "text", "text": args.prompt}, ], }, ] - ## STEP 5: Process Input with Chat Template - """inputs = processor.apply_chat_template( - messages, - add_generation_prompt=True, - tokenize=True, - return_dict=True, - return_tensors="pt", - ) - """ + ## STEP 5: Process Input inputs = processor( messages=messages, add_generation_prompt=True, @@ -532,7 +476,6 @@ def main(): inputs["pixel_values"] = inputs["pixel_values"].to(qeff_model.model.config.torch_dtype) ## STEP 6: Run Vision+Text Inference - streamer = TextStreamer(tokenizer) output = qeff_model.generate(inputs=inputs, device_ids=[0, 1], generation_len=100) ## STEP 7: Display Results @@ -541,6 +484,5 @@ def main(): print(output) - if __name__ == "__main__": main() diff --git a/examples/kimi_k2/export_kimi_k25_vision_4.57.3.py b/examples/kimi_k2/export_kimi_k25_vision_4.57.3.py index b55c77f4d3..bc04b4b294 100644 --- a/examples/kimi_k2/export_kimi_k25_vision_4.57.3.py +++ b/examples/kimi_k2/export_kimi_k25_vision_4.57.3.py @@ -24,7 +24,6 @@ from QEfficient import QEFFAutoModelForImageTextToText - MODEL_PATH = Path( "/home/huggingface_hub/models--moonshotai--Kimi-K2.5/snapshots/4d01dfe0332d63057c186e0b262165819efb6611" ) From 2e6bc245efe0c69374d9467a234f45acd0bacb83 Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Mon, 22 Jun 2026 21:17:03 +0530 Subject: [PATCH 06/33] fix MAD Signed-off-by: Mamta Singh --- .../models/kimi_k25/modeling_kimi_k25.py | 14 +- examples/kimi_k2/export_kimi_k25_vision.py | 3 +- .../kimi_k2/verify_kimi_k25_vision_parity.py | 269 ++++++++++++++++++ 3 files changed, 280 insertions(+), 6 deletions(-) create mode 100644 examples/kimi_k2/verify_kimi_k25_vision_parity.py diff --git a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py index b72e152860..7074707b9c 100644 --- a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py +++ b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py @@ -447,7 +447,8 @@ def __qeff_init__(self): if not hasattr(self, "blocks") or len(self.blocks) == 0: return - first_block = self.blocks[0] + old_blocks = list(self.blocks) + first_block = old_blocks[0] self.block_cfg = { "num_heads": first_block.num_heads, "hidden_dim": first_block.hidden_dim, @@ -463,9 +464,12 @@ def __qeff_init__(self): theta_base = getattr(self.rope_2d, "theta_base", 10000) self.rope_2d = Rope2DPosEmbRepeated(head_dim, max_height, max_width, theta_base=theta_base) - self.blocks = nn.ModuleList( - [MoonViTEncoderLayer(**self.block_cfg, use_deterministic_attn=False) for _ in range(len(self.blocks))] - ) + new_blocks = [] + for old_block in old_blocks: + new_block = MoonViTEncoderLayer(**self.block_cfg, use_deterministic_attn=False) + new_block.load_state_dict(old_block.state_dict()) + new_blocks.append(new_block.to(device=old_block.wqkv.weight.device, dtype=old_block.wqkv.weight.dtype)) + self.blocks = nn.ModuleList(new_blocks) def forward( self, @@ -908,7 +912,7 @@ def forward( else: logits = hidden_states.float() - output_compressed_kvs = getattr(outputs, "compressed_kvs", compressed_kvs) + output_compressed_kvs = getattr(outputs, "compressed_kvs", None) output_past_key_values = getattr(outputs, "past_key_values", None) return logits, image_embeds, image_idx, output_compressed_kvs, output_past_key_values diff --git a/examples/kimi_k2/export_kimi_k25_vision.py b/examples/kimi_k2/export_kimi_k25_vision.py index 40e968fdd0..ecf8c9e9b6 100644 --- a/examples/kimi_k2/export_kimi_k25_vision.py +++ b/examples/kimi_k2/export_kimi_k25_vision.py @@ -443,7 +443,7 @@ def main(): mxfp6_matmul=False, mxint8_kv_cache=False, aic_enable_depth_first=False, - # skip_lang=True, + #skip_lang=True, mos=1, num_patches=2400, # num_patches h=30, # h @@ -451,6 +451,7 @@ def main(): num_image_tokens=600, # num_image_tokens ) + #exit() ## STEP 4: Prepare Image and Text Input image = Image.open(BytesIO(requests.get(args.image_url).content)).convert("RGB") diff --git a/examples/kimi_k2/verify_kimi_k25_vision_parity.py b/examples/kimi_k2/verify_kimi_k25_vision_parity.py new file mode 100644 index 0000000000..649c93e65f --- /dev/null +++ b/examples/kimi_k2/verify_kimi_k25_vision_parity.py @@ -0,0 +1,269 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import argparse +import copy +import shutil +import subprocess +import tempfile +from io import BytesIO +from pathlib import Path + +import numpy as np +import onnxruntime as ort +import requests +import torch +from PIL import Image +from transformers.dynamic_module_utils import get_class_from_dynamic_module + +from QEfficient.generation.cloud_infer import QAICInferenceSession +from QEfficient.transformers.models.modeling_auto import QEffVisionEncoderForTextImageToTextModel +from export_kimi_k25_vision import ( + LOADED_EXPERT_IDS, + MODEL_PATH, + NUM_EXPERTS_PER_TOKEN, + NUM_TEXT_LAYERS, + NUM_VISION_LAYERS, + _ensure_torch_fx_import_compatibility, + _load_layer_subset_model, + _patch_deepseek_init_weights_compat, + _patch_kimi_tie_weights_compat, + _prepare_config, +) + + +def parse_args(): + parser = argparse.ArgumentParser(description="Verify Kimi K2.5 vision parity: HF PyTorch vs QEff PyTorch vs ONNXRuntime.") + parser.add_argument("--model-path", type=Path, default=MODEL_PATH) + parser.add_argument("--num-vision-layers", type=int, default=NUM_VISION_LAYERS) + parser.add_argument("--num-text-layers", type=int, default=NUM_TEXT_LAYERS) + parser.add_argument("--expert-ids", type=str, default=",".join(str(x) for x in LOADED_EXPERT_IDS)) + parser.add_argument("--num-experts-per-token", type=int, default=NUM_EXPERTS_PER_TOKEN) + parser.add_argument( + "--image-url", + type=str, + default="https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png", + ) + parser.add_argument("--prompt", type=str, default="Tell me about yourself.") + parser.add_argument("--atol", type=float, default=1e-3) + parser.add_argument("--rtol", type=float, default=1e-3) + parser.add_argument( + "--keep-onnx", + action="store_true", + help="Keep exported ONNX under ./aisand/onnx_export_verify instead of temp directory.", + ) + parser.add_argument("--run-qpc", action="store_true", help="Compile and execute QPC, then compare QPC output vs HF.") + parser.add_argument( + "--qaic-compile-bin", + type=str, + default="/opt/qti-aic/exec/qaic-compile", + help="Path to qaic-compile binary.", + ) + parser.add_argument( + "--keep-qpc", + action="store_true", + help="Keep compiled QPC under --qpc-dir (used only with --run-qpc).", + ) + parser.add_argument( + "--qpc-dir", + type=Path, + default=Path("aisand/qpc_verify"), + help="Output directory for QPC when --keep-qpc is set.", + ) + parser.add_argument("--aic-num-cores", type=int, default=16, help="Number of QAIC cores for qpc compile.") + parser.add_argument("--mos", type=int, default=1, help="-mos value for qpc compile.") + parser.add_argument( + "--no-convert-to-fp16", + action="store_true", + help="Disable -convert-to-fp16 during qpc compile (defaults to enabled).", + ) + return parser.parse_args() + + +def _parse_expert_ids(value: str): + return tuple(int(x.strip()) for x in value.split(",") if x.strip()) + + +def _stats(a: torch.Tensor, b: torch.Tensor, atol: float, rtol: float): + diff = (a - b).abs() + return { + "max_abs": float(diff.max().item()), + "mean_abs": float(diff.mean().item()), + "rmse": float(torch.sqrt(torch.mean((a - b) ** 2)).item()), + "allclose": bool(torch.allclose(a, b, atol=atol, rtol=rtol)), + } + + +def _hf_image_embeds(model, pixel_values, grid_thws): + target_dtype = model.vision_tower.patch_embed.proj.weight.dtype + pixel_values = pixel_values.to(target_dtype) + image_features = model.vision_tower(pixel_values, grid_thws) + if model.mm_projector: + image_features = model.mm_projector(image_features) + return image_features[0] if isinstance(image_features, list) else image_features + + +def _run_onnx(qeff_vision_wrapper, qeff_base_model, pixel_values, h_shape, w_shape, keep_onnx: bool): + output_names = qeff_base_model.get_output_names(kv_offload=True)["vision"] + dynamic_axes = qeff_base_model.get_onnx_dynamic_axes(kv_offload=True)["vision"] + + if keep_onnx: + export_dir = Path("aisand/onnx_export_verify") + export_dir.mkdir(parents=True, exist_ok=True) + else: + export_dir = Path(tempfile.mkdtemp(prefix="kimi_k25_verify_onnx_")) + + onnx_path = qeff_vision_wrapper.export( + {"pixel_values": pixel_values, "h_shape": h_shape, "w_shape": w_shape}, + output_names, + dynamic_axes, + export_dir=export_dir, + offload_pt_weights=False, + ) + + session = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"]) + ort_inputs = { + "pixel_values": pixel_values.detach().cpu().numpy().astype(np.float32), + "h_shape": h_shape.detach().cpu().numpy().astype(np.int64), + "w_shape": w_shape.detach().cpu().numpy().astype(np.int64), + } + ort_out = session.run(None, ort_inputs)[0] + return torch.from_numpy(ort_out).float(), Path(onnx_path) + + +def _run_qpc(onnx_path: Path, pixel_values, h_shape, w_shape, args): + if args.keep_qpc: + qpc_dir = args.qpc_dir + if qpc_dir.exists(): + shutil.rmtree(qpc_dir) + qpc_dir.parent.mkdir(parents=True, exist_ok=True) + else: + qpc_dir = Path(tempfile.mkdtemp(prefix="kimi_k25_verify_qpc_")) + + compile_cmd = [ + args.qaic_compile_bin, + f"-m={onnx_path}", + f"-onnx-define-symbol=num_patches,{pixel_values.shape[0]}", + f"-onnx-define-symbol=h,{h_shape.shape[0]}", + f"-onnx-define-symbol=w,{w_shape.shape[0]}", + f"-aic-num-cores={args.aic_num_cores}", + f"-mos={args.mos}", + f"-aic-binary-dir={qpc_dir}", + ] + if not args.no_convert_to_fp16: + compile_cmd.insert(1, "-convert-to-fp16") + compile_result = subprocess.run(compile_cmd, capture_output=True, text=True) + if compile_result.returncode != 0: + raise RuntimeError( + "QPC compile failed.\n" + f"Command: {' '.join(compile_cmd)}\n" + f"stdout:\n{compile_result.stdout}\n" + f"stderr:\n{compile_result.stderr}" + ) + + session = QAICInferenceSession(str(qpc_dir)) + qpc_inputs = {} + for name in session.input_names: + if name == "pixel_values": + qpc_inputs[name] = pixel_values.detach().cpu().numpy().astype(np.float32) + elif name == "h_shape": + qpc_inputs[name] = h_shape.detach().cpu().numpy().astype(np.int64) + elif name == "w_shape": + qpc_inputs[name] = w_shape.detach().cpu().numpy().astype(np.int64) + else: + raise RuntimeError(f"Unexpected QPC input name: {name}") + + qpc_outputs = session.run(qpc_inputs) + output_name = session.output_names[0] + qpc_out = qpc_outputs[output_name] + return torch.from_numpy(qpc_out).float(), qpc_dir + + +if __name__ == "__main__": + args = parse_args() + expert_ids = _parse_expert_ids(args.expert_ids) + + torch.manual_seed(123) + torch.set_grad_enabled(False) + + _ensure_torch_fx_import_compatibility() + config = _prepare_config(args.model_path) + kimi_cls = get_class_from_dynamic_module("modeling_kimi_k25.KimiK25ForConditionalGeneration", str(args.model_path)) + _patch_kimi_tie_weights_compat(kimi_cls) + _patch_deepseek_init_weights_compat(kimi_cls) + + hf_model, _, processor = _load_layer_subset_model( + model_path=args.model_path, + kimi_cls=kimi_cls, + config=config, + num_vision_layers=args.num_vision_layers, + num_text_layers=args.num_text_layers, + loaded_expert_ids=expert_ids, + num_experts_per_tok=args.num_experts_per_token, + dtype=torch.float32, + ) + + # Required for HW compatibility and parity in this flow. + hf_model.vision_tower.patch_embed.pos_emb.interpolation_mode = "bilinear" + hf_model.eval() + + qeff_base_model = copy.deepcopy(hf_model).eval() + qeff_vision_wrapper = QEffVisionEncoderForTextImageToTextModel(qeff_base_model) + qeff_encoder = qeff_vision_wrapper.model.eval() + + image = Image.open(BytesIO(requests.get(args.image_url, timeout=30).content)).convert("RGB") + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": image}, + {"type": "text", "text": args.prompt}, + ], + } + ] + inputs = processor(messages=messages, add_generation_prompt=True, tokenize=False, return_tensors="pt") + + pixel_values = inputs["pixel_values"] + grid_thws = inputs["grid_thws"] + h_shape = torch.ones(int(grid_thws[0, 1].item()), dtype=torch.int64) + w_shape = torch.ones(int(grid_thws[0, 2].item()), dtype=torch.int64) + + hf_out = _hf_image_embeds(hf_model, pixel_values, grid_thws).detach().cpu().float() + qeff_out = qeff_encoder(pixel_values, h_shape, w_shape).detach().cpu().float() + ort_out, onnx_path = _run_onnx(qeff_vision_wrapper, qeff_base_model, pixel_values, h_shape, w_shape, args.keep_onnx) + + hf_vs_qeff = _stats(hf_out, qeff_out, args.atol, args.rtol) + hf_vs_ort = _stats(hf_out, ort_out, args.atol, args.rtol) + qeff_vs_ort = _stats(qeff_out, ort_out, args.atol, args.rtol) + + print(f"ONNX path: {onnx_path}") + print("hf_vs_qeff", hf_vs_qeff) + print("hf_vs_onnxrt", hf_vs_ort) + print("qeff_vs_onnxrt", qeff_vs_ort) + + all_ok = hf_vs_qeff["allclose"] and hf_vs_ort["allclose"] and qeff_vs_ort["allclose"] + + if args.run_qpc: + atol = 1e-2 + rtol = 1e-2 + qpc_out, qpc_dir = _run_qpc(onnx_path, pixel_values, h_shape, w_shape, args) + hf_vs_qpc = _stats(hf_out, qpc_out, atol, rtol) + qeff_vs_qpc = _stats(qeff_out, qpc_out, atol, rtol) + ort_vs_qpc = _stats(ort_out, qpc_out, atol, rtol) + print(f"QPC dir: {qpc_dir}") + print("hf_vs_qpc", hf_vs_qpc) + print("qeff_vs_qpc", qeff_vs_qpc) + print("onnxrt_vs_qpc", ort_vs_qpc) + all_ok = all_ok and hf_vs_qpc["allclose"] and qeff_vs_qpc["allclose"] and ort_vs_qpc["allclose"] + + if not all_ok: + raise SystemExit( + f"Parity check failed for atol={atol}, rtol={rtol}. " + f"See metrics above for details." + ) + + #print(f"PASS: requested parity checks match within atol={args.atol}, rtol={args.rtol}") \ No newline at end of file From 95e415eab4fa422aae716d4710143c131c6615ff Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Mon, 22 Jun 2026 23:36:20 +0530 Subject: [PATCH 07/33] language model export and compile Signed-off-by: Mamta Singh --- QEfficient/generation/vlm_generation.py | 1 + .../models/gemma3/modeling_gemma3.py | 2 +- .../models/kimi_k25/modeling_kimi_k25.py | 19 +- .../transformers/models/pytorch_transforms.py | 1 - examples/kimi_k2/export_kimi_k25_vision.py | 10 +- .../kimi_k2/export_kimi_k25_vision_4.57.3.py | 423 ------------------ ...verify_kimi_k25_language_decoder_parity.py | 327 ++++++++++++++ .../kimi_k2/verify_kimi_k25_vision_parity.py | 25 +- 8 files changed, 349 insertions(+), 459 deletions(-) delete mode 100644 examples/kimi_k2/export_kimi_k25_vision_4.57.3.py create mode 100644 examples/kimi_k2/verify_kimi_k25_language_decoder_parity.py diff --git a/QEfficient/generation/vlm_generation.py b/QEfficient/generation/vlm_generation.py index 1e87408234..127ef432fd 100755 --- a/QEfficient/generation/vlm_generation.py +++ b/QEfficient/generation/vlm_generation.py @@ -736,6 +736,7 @@ def _generate_multi_frame_specialization( generation_len: int = None, stream: List[str] = None, ): + exec_batch_size = self.batch_size max_gen_length = self._ctx_len if not generation_len else max(self._ctx_len, generation_len) self.initialize_decode_inputs( diff --git a/QEfficient/transformers/models/gemma3/modeling_gemma3.py b/QEfficient/transformers/models/gemma3/modeling_gemma3.py index cbb3902ec9..96bf6cd1a0 100644 --- a/QEfficient/transformers/models/gemma3/modeling_gemma3.py +++ b/QEfficient/transformers/models/gemma3/modeling_gemma3.py @@ -21,6 +21,7 @@ Gemma3ForConditionalGeneration, Gemma3TextConfig, Gemma3TextModel, + logger, repeat_kv, rotate_half, ) @@ -31,7 +32,6 @@ from QEfficient.utils import constants from QEfficient.utils._utils import IOInfo from QEfficient.utils.constants import MIN_MASKED_ATTENTION_VALUE -from QEfficient.utils.logging_utils import logger class GemmaRMSNormFunc(torch.autograd.Function): diff --git a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py index 7074707b9c..b9905602ff 100644 --- a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py +++ b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py @@ -858,15 +858,8 @@ def forward( batch_index: Optional[torch.LongTensor] = None, comp_ctx_lengths: Optional[List[int]] = None, use_cache: Optional[bool] = None, - labels: Optional[torch.LongTensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - cache_position: Optional[torch.LongTensor] = None, - num_logits_to_keep: Optional[int] = None, **kwargs, ) -> Tuple: - del labels, output_attentions, output_hidden_states, return_dict, cache_position, num_logits_to_keep if inputs_embeds is None: inputs_embeds = self.model.get_input_embeddings()(input_ids) @@ -901,16 +894,8 @@ def forward( **kwargs, ) - hidden_states = outputs[0] - logit_index = position_ids.to(torch.int32).argmax(1, keepdim=True) - max_index = hidden_states.shape[1] - 1 - safe_logit_index = torch.clamp(logit_index, min=0, max=max_index) - hidden_states = hidden_states[torch.arange(position_ids.shape[0]).view(-1, 1), safe_logit_index] - - if hidden_states.shape[-1] == self.lm_head.in_features: - logits = self.lm_head(hidden_states).float() - else: - logits = hidden_states.float() + # QEff Deepseek language_model.forward already returns final logits. + logits = outputs[0].float() output_compressed_kvs = getattr(outputs, "compressed_kvs", None) output_past_key_values = getattr(outputs, "past_key_values", None) diff --git a/QEfficient/transformers/models/pytorch_transforms.py b/QEfficient/transformers/models/pytorch_transforms.py index 9d56f5b8e6..247340a94e 100755 --- a/QEfficient/transformers/models/pytorch_transforms.py +++ b/QEfficient/transformers/models/pytorch_transforms.py @@ -1232,7 +1232,6 @@ class KVCacheExternalModuleMapperTransform(ExternalModuleMapperTransform): "forward": QEFFGrok1CustomRMSNormAIC.forward, }, "KimiK25ForConditionalGeneration": { - # "forward": QEffKimiK25ForConditionalGeneration.forward_only_image_for_export, "get_qeff_vision_encoder": QEffKimiK25ForConditionalGeneration.get_qeff_vision_encoder, "get_qeff_language_decoder": QEffKimiK25ForConditionalGeneration.get_qeff_language_decoder, "get_specializations": QEffKimiK25ForConditionalGeneration.get_specializations, diff --git a/examples/kimi_k2/export_kimi_k25_vision.py b/examples/kimi_k2/export_kimi_k25_vision.py index ecf8c9e9b6..70d58cb1bc 100644 --- a/examples/kimi_k2/export_kimi_k25_vision.py +++ b/examples/kimi_k2/export_kimi_k25_vision.py @@ -374,7 +374,7 @@ def main(): else: raise ValueError("Pass both --num-vision-layers and --num-text-layers to load a layer subset.") - qaic_config = {"mla_absorption": {"cache_compressed": True, "absorption": False, "online": False}} + # qaic_config = {"mla_absorption": {"cache_compressed": True, "absorption": False, "online": False}} qeff_model = QEFFAutoModelForImageTextToText(model) @@ -386,7 +386,7 @@ def main(): ## STEP 3: Compile Model for Text-Only Execution # Set skip_vision=True to bypass image processing qeff_model.compile( - qaic_config=qaic_config, + # qaic_config=qaic_config, prefill_seq_len=1, ctx_len=1024, num_cores=16, @@ -435,7 +435,7 @@ def main(): ## STEP 3: Compile Model for Vision+Text Execution # Do not set skip_vision (defaults to False) to enable image processing qeff_model.compile( - qaic_config=qaic_config, + # qaic_config=qaic_config, prefill_seq_len=1, ctx_len=1024, num_cores=16, @@ -443,7 +443,7 @@ def main(): mxfp6_matmul=False, mxint8_kv_cache=False, aic_enable_depth_first=False, - #skip_lang=True, + # skip_lang=True, mos=1, num_patches=2400, # num_patches h=30, # h @@ -451,7 +451,7 @@ def main(): num_image_tokens=600, # num_image_tokens ) - #exit() + # exit() ## STEP 4: Prepare Image and Text Input image = Image.open(BytesIO(requests.get(args.image_url).content)).convert("RGB") diff --git a/examples/kimi_k2/export_kimi_k25_vision_4.57.3.py b/examples/kimi_k2/export_kimi_k25_vision_4.57.3.py deleted file mode 100644 index bc04b4b294..0000000000 --- a/examples/kimi_k2/export_kimi_k25_vision_4.57.3.py +++ /dev/null @@ -1,423 +0,0 @@ -# ----------------------------------------------------------------------------- -# -# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. -# SPDX-License-Identifier: BSD-3-Clause -# -# ----------------------------------------------------------------------------- - -import argparse -import copy -import json -import re -import sys -import tempfile -from collections import defaultdict -from io import BytesIO -from pathlib import Path - -import requests -from PIL import Image -from safetensors import safe_open -from safetensors.torch import save_file -from transformers import AutoConfig, AutoProcessor, AutoTokenizer, TextStreamer -from transformers.dynamic_module_utils import get_class_from_dynamic_module - -from QEfficient import QEFFAutoModelForImageTextToText - -MODEL_PATH = Path( - "/home/huggingface_hub/models--moonshotai--Kimi-K2.5/snapshots/4d01dfe0332d63057c186e0b262165819efb6611" -) -NUM_VISION_LAYERS = 2 -NUM_TEXT_LAYERS = 2 -LOADED_EXPERT_IDS = (0, 1, 2, 3) -NUM_EXPERTS_PER_TOKEN = 2 - -EXPERT_KEY_PATTERN = re.compile(r"^(language_model\.model\.layers\.\d+\.mlp\.experts\.)(\d+)(\..+)$") - - -def _parse_expert_ids(value: str): - expert_ids = tuple(int(expert_id) for expert_id in value.split(",") if expert_id.strip()) - if not expert_ids: - raise argparse.ArgumentTypeError("At least one expert id must be provided.") - return expert_ids - - -def _validate_kimi_k25_config(kimi_config): - if kimi_config.model_type != "kimi_k25": - raise ValueError(f"This script only supports Kimi K2.5 config, got {kimi_config.model_type!r}.") - architectures = getattr(kimi_config, "architectures", None) or [] - if "KimiK25ForConditionalGeneration" not in architectures: - raise ValueError(f"Expected KimiK25ForConditionalGeneration, got architectures={architectures!r}.") - - -def _validate_layer_count(name, requested_count, available_count): - if requested_count < 1: - raise ValueError(f"{name} must be >= 1, got {requested_count}.") - if requested_count > available_count: - raise ValueError(f"{name}={requested_count} exceeds available layers={available_count}.") - - -def _validate_expert_subset(loaded_expert_ids, num_experts_per_tok, total_experts): - expert_ids = tuple(loaded_expert_ids) - if len(expert_ids) != 4: - raise ValueError(f"Expected exactly 4 routed experts, got {expert_ids!r}.") - if len(set(expert_ids)) != len(expert_ids): - raise ValueError(f"Expert ids must be unique, got {expert_ids!r}.") - invalid_ids = [expert_id for expert_id in expert_ids if expert_id < 0 or expert_id >= total_experts] - if invalid_ids: - raise ValueError(f"Expert ids {invalid_ids!r} are outside the valid range [0, {total_experts - 1}].") - if num_experts_per_tok > len(expert_ids): - raise ValueError(f"num_experts_per_tok={num_experts_per_tok} cannot exceed {len(expert_ids)} loaded experts.") - return expert_ids - - -def _extend_quantization_ignore(config, patterns): - quantization_config = getattr(config, "quantization_config", None) - if not quantization_config: - return - - ignored_modules = quantization_config.setdefault("ignore", []) - for pattern in patterns: - if pattern not in ignored_modules: - ignored_modules.append(pattern) - - -def _remap_checkpoint_key(checkpoint_key, expert_index_map): - match = EXPERT_KEY_PATTERN.match(checkpoint_key) - if not match: - return checkpoint_key - - original_expert_idx = int(match.group(2)) - remapped_expert_idx = expert_index_map.get(original_expert_idx) - if remapped_expert_idx is None: - return None - return f"{match.group(1)}{remapped_expert_idx}{match.group(3)}" - - -def _is_routed_gate_weight(checkpoint_key): - return checkpoint_key.endswith(".mlp.gate.weight") - - -def _is_routed_gate_bias(checkpoint_key): - return checkpoint_key.endswith(".mlp.gate.e_score_correction_bias") - - -def _materialize_subset_checkpoint( - model_path: Path, - temp_model_path: Path, - weight_map, - allowed_prefixes, - loaded_expert_ids, -): - expert_index_map = {expert_id: remapped_idx for remapped_idx, expert_id in enumerate(loaded_expert_ids)} - selected_by_shard = defaultdict(list) - - for checkpoint_key, shard_name in weight_map.items(): - if not any(checkpoint_key.startswith(prefix) for prefix in allowed_prefixes): - continue - - remapped_key = _remap_checkpoint_key(checkpoint_key, expert_index_map) - if remapped_key is None: - continue - selected_by_shard[shard_name].append((checkpoint_key, remapped_key)) - - if not selected_by_shard: - raise RuntimeError("No multimodal weights were selected from the Kimi K2.5 checkpoint.") - - filtered_weight_map = {} - subset_shards = [] - for shard_idx, (source_shard_name, shard_entries) in enumerate(sorted(selected_by_shard.items())): - tensors = {} - with safe_open(model_path / source_shard_name, framework="pt", device="cpu") as shard_reader: - for checkpoint_key, remapped_key in shard_entries: - tensor = shard_reader.get_tensor(checkpoint_key) - if _is_routed_gate_weight(checkpoint_key): - tensor = tensor[list(loaded_expert_ids), :].contiguous() - elif _is_routed_gate_bias(checkpoint_key): - tensor = tensor[list(loaded_expert_ids)].contiguous() - tensors[remapped_key] = tensor - - subset_shard_name = f"model-subset-{shard_idx:05d}.safetensors" - save_file(tensors, str(temp_model_path / subset_shard_name)) - subset_shards.append(subset_shard_name) - filtered_weight_map.update({remapped_key: subset_shard_name for _, remapped_key in shard_entries}) - - return filtered_weight_map, subset_shards - - -def _allowed_prefixes(num_vision_layers: int, num_text_layers: int): - prefixes = [ - "vision_tower.patch_embed.", - "vision_tower.encoder.final_layernorm.", - "mm_projector.", - "language_model.model.embed_tokens.", - "language_model.model.norm.", - "language_model.lm_head.", - ] - prefixes.extend(f"vision_tower.encoder.blocks.{layer_idx}." for layer_idx in range(num_vision_layers)) - prefixes.extend(f"language_model.model.layers.{layer_idx}." for layer_idx in range(num_text_layers)) - return prefixes - - -def load_kimi_k25_vision_moe_strip( - model_path: Path, - num_vision_layers: int, - num_text_layers: int, - loaded_expert_ids=LOADED_EXPERT_IDS, - num_experts_per_tok: int = NUM_EXPERTS_PER_TOKEN, -): - kimi_config = AutoConfig.from_pretrained(str(model_path), trust_remote_code=True) - _validate_kimi_k25_config(kimi_config) - - stripped_config = copy.deepcopy(kimi_config) - text_config = stripped_config.text_config - vision_config = stripped_config.vision_config - stripped_config._attn_implementation = "eager" - text_config._attn_implementation = "eager" - vision_config._attn_implementation = "eager" - quantization_ignores = [ - "re:vision_tower.*", - "re:mm_projector.*", - "re:language_model.lm_head.*", - ] - _extend_quantization_ignore(stripped_config, quantization_ignores) - _extend_quantization_ignore(text_config, quantization_ignores) - - _validate_layer_count("num_text_layers", num_text_layers, text_config.num_hidden_layers) - _validate_layer_count("num_vision_layers", num_vision_layers, vision_config.vt_num_hidden_layers) - - text_config.num_hidden_layers = num_text_layers - vision_config.vt_num_hidden_layers = num_vision_layers - loaded_expert_ids = _validate_expert_subset( - loaded_expert_ids, - num_experts_per_tok, - text_config.n_routed_experts, - ) - text_config.n_routed_experts = len(loaded_expert_ids) - text_config.num_experts_per_tok = num_experts_per_tok - text_config.n_group = 1 - text_config.topk_group = 1 - - kimi_cls = get_class_from_dynamic_module("modeling_kimi_k25.KimiK25ForConditionalGeneration", str(model_path)) - kimi_module = sys.modules[kimi_cls.__module__] - kimi_module.MoonViT3dEncoder.use_deterministic_attn = False - - checkpoint_index = json.loads((model_path / "model.safetensors.index.json").read_text()) - weight_map = checkpoint_index["weight_map"] - - with tempfile.TemporaryDirectory() as tmpdir: - temp_model_path = Path(tmpdir) - filtered_weight_map, subset_shards = _materialize_subset_checkpoint( - model_path=model_path, - temp_model_path=temp_model_path, - weight_map=weight_map, - allowed_prefixes=_allowed_prefixes(num_vision_layers, num_text_layers), - loaded_expert_ids=loaded_expert_ids, - ) - (temp_model_path / "config.json").write_text(stripped_config.to_json_string(use_diff=False)) - (temp_model_path / "model.safetensors.index.json").write_text( - json.dumps( - { - "metadata": { - "total_size": sum((temp_model_path / shard_name).stat().st_size for shard_name in subset_shards) - }, - "weight_map": filtered_weight_map, - } - ) - ) - - original_base_model_prefix = kimi_cls.base_model_prefix - kimi_cls.base_model_prefix = "" - try: - model, loading_info = kimi_cls.from_pretrained( - str(temp_model_path), - config=stripped_config, - local_files_only=True, - attn_implementation="eager", - output_loading_info=True, - ) - finally: - kimi_cls.base_model_prefix = original_base_model_prefix - - unexpected_keys = loading_info["unexpected_keys"] - missing_keys = loading_info["missing_keys"] - mismatched_keys = loading_info["mismatched_keys"] - if unexpected_keys or missing_keys or mismatched_keys: - raise RuntimeError( - "Failed to load the stripped Kimi K2.5 multimodal checkpoint slice cleanly. " - f"missing={missing_keys}, unexpected={unexpected_keys}, mismatched={mismatched_keys}" - ) - - model.eval() - tokenizer = AutoTokenizer.from_pretrained(str(model_path), trust_remote_code=True) - processor = AutoProcessor.from_pretrained(str(model_path), trust_remote_code=True) - return model, tokenizer, processor - - -def parse_args(): - parser = argparse.ArgumentParser( - description="Load a stripped Kimi K2.5 multimodal checkpoint with configurable vision/text layers and 4 routed experts." - ) - parser.add_argument("--model-path", type=Path, default=MODEL_PATH) - parser.add_argument("--num-vision-layers", type=int, default=NUM_VISION_LAYERS) - parser.add_argument("--num-text-layers", type=int, default=NUM_TEXT_LAYERS) - parser.add_argument("--expert-ids", type=_parse_expert_ids, default=LOADED_EXPERT_IDS) - parser.add_argument("--num-experts-per-token", type=int, default=NUM_EXPERTS_PER_TOKEN) - parser.add_argument( - "--image-url", - type=str, - default="https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png", - ) - parser.add_argument("--prompt", type=str, default="Describe this image.") - parser.add_argument("--test", action="store_true", help="Validate ONNX output matches PyTorch image-only forward.") - return parser.parse_args() - - -def main(): - args = parse_args() - model, tokenizer, processor = load_kimi_k25_vision_moe_strip( - args.model_path, - args.num_vision_layers, - args.num_text_layers, - loaded_expert_ids=args.expert_ids, - num_experts_per_tok=args.num_experts_per_token, - ) - print( - f"Loaded {type(model).__name__} with " - f"{model.config.vision_config.vt_num_hidden_layers} vision layers, " - f"{model.config.text_config.num_hidden_layers} text layers, " - f"{model.config.text_config.n_routed_experts} routed experts." - ) - print(f"Tokenizer vocab size: {tokenizer.vocab_size}") - print(f"Processor type: {type(processor).__name__}") - - mla_absorption = {"cache_compressed": True, "absorption": False, "online": False} - qaic_config = { - "mla_absorption": mla_absorption - } # , "enable_blocking": True, "blocking_mode": "par", "par_num_split": 4, "num_kv_blocks": 8} - breakpoint() - qeff_model = QEFFAutoModelForImageTextToText(model) # , qaic_config=qaic_config) - breakpoint() - - skip_vision = True - - if skip_vision: - ## TEXT-ONLY MODE ## - - ## STEP 3: Compile Model for Text-Only Execution - # Set skip_vision=True to bypass image processing - qeff_model.compile( - qaic_config=qaic_config, - prefill_seq_len=32, - ctx_len=1024, - num_cores=16, - num_devices=2, - mxfp6_matmul=False, - mxint8_kv_cache=False, - aic_enable_depth_first=False, - skip_vision=True, # Skip vision encoder for text-only inference - mos=1, - num_patches=2400, # num_patches - h=30, # h - w=80, # w - num_image_tokens=600, # num_image_tokens - ) - breakpoint() - ## STEP 4: Prepare Text-Only Input - # Create a text-only message without any image - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "Tell me about yourself."}, - ], - }, - ] - - ## STEP 5: Process Input with Chat Template - inputs = processor.apply_chat_template( - messages, - add_generation_prompt=True, - tokenize=True, - return_dict=True, - return_tensors="pt", - ) - - ## STEP 6: Run Text-Only Inference - streamer = TextStreamer(tokenizer) - output = qeff_model.generate(inputs=inputs, device_ids=[0, 1], generation_len=10) - - ## STEP 7: Display Results - print(output.generated_ids) - print(tokenizer.batch_decode(output.generated_ids)) - print(output) - - else: - ## VISION + TEXT MODE ## - - ## STEP 3: Compile Model for Vision+Text Execution - # Do not set skip_vision (defaults to False) to enable image processing - qeff_model.compile( - qaic_config=qaic_config, - prefill_seq_len=32, - ctx_len=1024, - num_cores=16, - num_devices=2, - mxfp6_matmul=False, - mxint8_kv_cache=False, - aic_enable_depth_first=False, - skip_vision=True, # Skip vision encoder for text-only inference - mos=1, - num_patches=2400, # num_patches - h=30, # h - w=80, # w - num_image_tokens=600, # num_image_tokens - ) - - breakpoint() - ## STEP 4: Prepare Image and Text Input - # Define the image URL to process - image_url = "https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png" - image = Image.open(BytesIO(requests.get(args.image_url).content)).convert("RGB") - - # Create a message with both image and text - messages = [ - { - "role": "user", - "content": [ - {"type": "image", "url": image}, - {"type": "text", "text": "Can you describe the image in detail."}, - ], - }, - ] - - ## STEP 5: Process Input with Chat Template - """inputs = processor.apply_chat_template( - messages, - add_generation_prompt=True, - tokenize=True, - return_dict=True, - return_tensors="pt", - ) - """ - inputs = processor( - messages=messages, - add_generation_prompt=True, - tokenize=False, - return_tensors="pt", - ) - # Convert pixel values to float32 for processing - inputs["pixel_values"] = inputs["pixel_values"].to(qeff_model.model.config.torch_dtype) - - ## STEP 6: Run Vision+Text Inference - streamer = TextStreamer(tokenizer) - output = qeff_model.generate(inputs=inputs, device_ids=[0, 1], generation_len=100) - - ## STEP 7: Display Results - print(output.generated_ids) - print(tokenizer.batch_decode(output.generated_ids)) - print(output) - - -if __name__ == "__main__": - main() diff --git a/examples/kimi_k2/verify_kimi_k25_language_decoder_parity.py b/examples/kimi_k2/verify_kimi_k25_language_decoder_parity.py new file mode 100644 index 0000000000..70cc0c35e1 --- /dev/null +++ b/examples/kimi_k2/verify_kimi_k25_language_decoder_parity.py @@ -0,0 +1,327 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import argparse +import copy +import shutil +import subprocess +import tempfile +from pathlib import Path + +import numpy as np +import onnxruntime as ort +import torch +from export_kimi_k25_vision import ( + LOADED_EXPERT_IDS, + MODEL_PATH, + NUM_EXPERTS_PER_TOKEN, + _ensure_torch_fx_import_compatibility, + _load_layer_subset_model, + _patch_deepseek_init_weights_compat, + _patch_kimi_tie_weights_compat, + _prepare_config, +) +from transformers.dynamic_module_utils import get_class_from_dynamic_module + +from QEfficient import QEFFAutoModelForImageTextToText +from QEfficient.generation.cloud_infer import QAICInferenceSession +from QEfficient.transformers.models.modeling_auto import QEffCausalLMForTextImageToTextModel + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Verify Kimi-K2.5 language decoder parity: HF LM logits vs transformed decoder-wrapper logits." + ) + parser.add_argument("--model-path", type=Path, default=MODEL_PATH) + parser.add_argument("--num-text-layers", type=int, default=2) + parser.add_argument("--num-vision-layers", type=int, default=1) + parser.add_argument("--expert-ids", type=str, default=",".join(str(x) for x in LOADED_EXPERT_IDS)) + parser.add_argument("--num-experts-per-token", type=int, default=NUM_EXPERTS_PER_TOKEN) + parser.add_argument("--prompt", type=str, default="Describe this image.") + parser.add_argument("--batch-size", type=int, default=1) + parser.add_argument("--atol", type=float, default=1e-4) + parser.add_argument("--rtol", type=float, default=1e-4) + parser.add_argument("--skip-onnx", action="store_true", help="Skip ONNX export/runtime parity check.") + parser.add_argument( + "--keep-onnx", + action="store_true", + help="Keep exported ONNX under ./aisand/onnx_export_verify_lang instead of a temporary directory.", + ) + parser.add_argument("--run-qpc", action="store_true", help="Compile and execute QPC, then compare vs references.") + parser.add_argument( + "--qaic-compile-bin", + type=str, + default="/opt/qti-aic/exec/qaic-compile", + help="Path to qaic-compile binary.", + ) + parser.add_argument( + "--keep-qpc", + action="store_true", + help="Keep compiled QPC under --qpc-dir (used only with --run-qpc).", + ) + parser.add_argument( + "--qpc-dir", + type=Path, + default=Path("aisand/qpc_verify_lang"), + help="Output directory for QPC when --keep-qpc is set.", + ) + parser.add_argument("--aic-num-cores", type=int, default=16, help="Number of QAIC cores for qpc compile.") + parser.add_argument("--mos", type=int, default=1, help="-mos value for qpc compile.") + parser.add_argument( + "--no-convert-to-fp16", + action="store_true", + help="Disable -convert-to-fp16 during qpc compile (defaults to enabled).", + ) + parser.add_argument("--qpc-atol", type=float, default=1e-2) + parser.add_argument("--qpc-rtol", type=float, default=1e-2) + return parser.parse_args() + + +def _parse_expert_ids(value: str): + return tuple(int(x.strip()) for x in value.split(",") if x.strip()) + + +def _stats(a: torch.Tensor, b: torch.Tensor, atol: float, rtol: float): + diff = (a - b).abs() + return { + "shape": tuple(a.shape), + "max_abs": float(diff.max().item()), + "mean_abs": float(diff.mean().item()), + "rmse": float(torch.sqrt(torch.mean((a - b) ** 2)).item()), + "allclose": bool(torch.allclose(a, b, atol=atol, rtol=rtol)), + } + + +def _language_model_reference_logits(model, input_ids, attention_mask, position_ids, past_key_values): + inputs_embeds = model.get_input_embeddings()(input_ids) + outputs = model.language_model( + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=True, + ) + return outputs[0].float() + + +def _create_prompt_inputs(tokenizer, prompt: str, batch_size: int): + prompts = [prompt] * batch_size + tokenized = tokenizer(prompts, return_tensors="pt", padding=True) + + input_ids = tokenized["input_ids"].to(torch.long) + attention_mask = tokenized["attention_mask"].to(torch.long) + position_ids = torch.where(attention_mask.bool(), attention_mask.cumsum(-1) - 1, -1) + return input_ids, attention_mask, position_ids + + +def _export_and_run_onnx( + transformed_model, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + past_key_values, + keep_onnx: bool, +): + lang_exporter = QEffCausalLMForTextImageToTextModel(copy.deepcopy(transformed_model)) + lang_inputs = transformed_model.get_dummy_inputs(kv_offload=True)["lang"] + output_names = transformed_model.get_output_names(kv_offload=True)["lang"] + dynamic_axes = transformed_model.get_onnx_dynamic_axes(kv_offload=True)["lang"] + + if keep_onnx: + export_dir = Path("aisand/onnx_export_verify_lang") + export_dir.mkdir(parents=True, exist_ok=True) + else: + export_dir = Path(tempfile.mkdtemp(prefix="kimi_k25_verify_lang_onnx_")) + + onnx_path = lang_exporter.export( + inputs=lang_inputs, + output_names=output_names, + dynamic_axes=dynamic_axes, + export_dir=export_dir, + offload_pt_weights=False, + ) + + session = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"]) + ort_inputs = { + "input_ids": input_ids.detach().cpu().numpy().astype(np.int64), + "position_ids": position_ids.detach().cpu().numpy().astype(np.int64), + } + for layer_idx, (past_key, past_value) in enumerate(past_key_values): + ort_inputs[f"past_key.{layer_idx}"] = past_key.detach().cpu().numpy().astype(np.float32) + ort_inputs[f"past_value.{layer_idx}"] = past_value.detach().cpu().numpy().astype(np.float32) + + ort_logits = torch.from_numpy(session.run(["logits"], ort_inputs)[0]).float() + return ort_logits, Path(onnx_path) + + +def _run_qpc(onnx_path: Path, input_ids: torch.Tensor, position_ids: torch.Tensor, past_key_values, args): + if args.keep_qpc: + qpc_dir = args.qpc_dir + if qpc_dir.exists(): + shutil.rmtree(qpc_dir) + qpc_dir.parent.mkdir(parents=True, exist_ok=True) + else: + qpc_dir = Path(tempfile.mkdtemp(prefix="kimi_k25_verify_lang_qpc_")) + shutil.rmtree(qpc_dir) + + compile_cmd = [ + args.qaic_compile_bin, + f"-m={onnx_path}", + f"-onnx-define-symbol=batch_size,{input_ids.shape[0]}", + f"-onnx-define-symbol=seq_len,{input_ids.shape[1]}", + f"-onnx-define-symbol=ctx_len,{past_key_values[0][0].shape[2]}", + f"-aic-num-cores={args.aic_num_cores}", + f"-mos={args.mos}", + f"-aic-binary-dir={qpc_dir}", + ] + if not args.no_convert_to_fp16: + compile_cmd.insert(1, "-convert-to-fp16") + + compile_result = subprocess.run(compile_cmd, capture_output=True, text=True) + if compile_result.returncode != 0: + raise RuntimeError( + "QPC compile failed.\n" + f"Command: {' '.join(compile_cmd)}\n" + f"stdout:\n{compile_result.stdout}\n" + f"stderr:\n{compile_result.stderr}" + ) + + session = QAICInferenceSession(str(qpc_dir)) + qpc_inputs = { + "input_ids": input_ids.detach().cpu().numpy().astype(np.int64), + "position_ids": position_ids.detach().cpu().numpy().astype(np.int64), + } + for layer_idx, (past_key, past_value) in enumerate(past_key_values): + qpc_inputs[f"past_key.{layer_idx}"] = past_key.detach().cpu().numpy().astype(np.float32) + qpc_inputs[f"past_value.{layer_idx}"] = past_value.detach().cpu().numpy().astype(np.float32) + + qpc_feed = {name: qpc_inputs[name] for name in session.input_names} + qpc_outputs = session.run(qpc_feed) + output_name = "logits" if "logits" in qpc_outputs else session.output_names[0] + qpc_logits = torch.from_numpy(qpc_outputs[output_name]).float() + return qpc_logits, qpc_dir + + +if __name__ == "__main__": + args = parse_args() + expert_ids = _parse_expert_ids(args.expert_ids) + + torch.set_grad_enabled(False) + _ensure_torch_fx_import_compatibility() + + config = _prepare_config(args.model_path) + kimi_cls = get_class_from_dynamic_module("modeling_kimi_k25.KimiK25ForConditionalGeneration", str(args.model_path)) + _patch_kimi_tie_weights_compat(kimi_cls) + _patch_deepseek_init_weights_compat(kimi_cls) + + hf_model, tokenizer, _ = _load_layer_subset_model( + model_path=args.model_path, + kimi_cls=kimi_cls, + config=config, + num_vision_layers=args.num_vision_layers, + num_text_layers=args.num_text_layers, + loaded_expert_ids=expert_ids, + num_experts_per_tok=args.num_experts_per_token, + dtype=torch.float32, + ) + hf_model = hf_model.eval() + + transformed_wrapper = QEFFAutoModelForImageTextToText(copy.deepcopy(hf_model)) + transformed_model = transformed_wrapper.model.eval() + decoder_wrapper = transformed_model.get_qeff_language_decoder().eval() + + input_ids, attention_mask, position_ids = _create_prompt_inputs( + tokenizer, + prompt=args.prompt, + batch_size=args.batch_size, + ) + + past_key_values = transformed_model.language_model.get_dummy_pkv_cache( + transformed_model.config.text_config, + args.batch_size, + input_ids.shape[1], + ) + + reference_logits = ( + _language_model_reference_logits( + transformed_model, + input_ids, + attention_mask, + position_ids, + past_key_values=copy.deepcopy(past_key_values), + ) + .detach() + .cpu() + ) + transformed_decoder_logits = ( + decoder_wrapper( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + image_embeds=None, + image_idx=None, + past_key_values=copy.deepcopy(past_key_values), + use_cache=True, + )[0] + .detach() + .cpu() + .float() + ) + + lm_vs_wrapper = _stats(reference_logits, transformed_decoder_logits, args.atol, args.rtol) + print( + "Loaded subset:", + f"vision_layers={transformed_model.config.vision_config.vt_num_hidden_layers}", + f"text_layers={transformed_model.config.text_config.num_hidden_layers}", + ) + print("Prompt:", repr(args.prompt)) + print("Input shape:", tuple(input_ids.shape)) + print("Language model vs decoder-wrapper:", lm_vs_wrapper) + + if not lm_vs_wrapper["allclose"]: + raise SystemExit( + "Parity check failed: transformed decoder-wrapper logits do not match transformed language-model logits." + ) + + ort_logits = None + onnx_path = None + if not args.skip_onnx or args.run_qpc: + ort_logits, onnx_path = _export_and_run_onnx( + transformed_model, + input_ids=input_ids, + position_ids=position_ids, + past_key_values=copy.deepcopy(past_key_values), + keep_onnx=args.keep_onnx, + ) + + if ort_logits is not None: + lm_vs_onnx = _stats(reference_logits, ort_logits, args.atol, args.rtol) + wrapper_vs_onnx = _stats(transformed_decoder_logits, ort_logits, args.atol, args.rtol) + print("Language ONNX path:", onnx_path) + print("Language model vs ONNXRuntime:", lm_vs_onnx) + print("Decoder-wrapper vs ONNXRuntime:", wrapper_vs_onnx) + + if not lm_vs_onnx["allclose"] or not wrapper_vs_onnx["allclose"]: + raise SystemExit("Parity check failed: ONNXRuntime logits do not match transformed language references.") + + if args.run_qpc: + qpc_logits, qpc_dir = _run_qpc( + onnx_path, + input_ids=input_ids, + position_ids=position_ids, + past_key_values=copy.deepcopy(past_key_values), + args=args, + ) + lm_vs_qpc = _stats(reference_logits, qpc_logits, args.qpc_atol, args.qpc_rtol) + wrapper_vs_qpc = _stats(transformed_decoder_logits, qpc_logits, args.qpc_atol, args.qpc_rtol) + print("Language QPC dir:", qpc_dir) + print(f"Language model vs QPC (atol={args.qpc_atol}, rtol={args.qpc_rtol}):", lm_vs_qpc) + print(f"Decoder-wrapper vs QPC (atol={args.qpc_atol}, rtol={args.qpc_rtol}):", wrapper_vs_qpc) + + if not lm_vs_qpc["allclose"] or not wrapper_vs_qpc["allclose"]: + raise SystemExit("Parity check failed: QPC logits do not match transformed language references.") + + print("Parity check passed.") diff --git a/examples/kimi_k2/verify_kimi_k25_vision_parity.py b/examples/kimi_k2/verify_kimi_k25_vision_parity.py index 649c93e65f..729b22edec 100644 --- a/examples/kimi_k2/verify_kimi_k25_vision_parity.py +++ b/examples/kimi_k2/verify_kimi_k25_vision_parity.py @@ -17,11 +17,6 @@ import onnxruntime as ort import requests import torch -from PIL import Image -from transformers.dynamic_module_utils import get_class_from_dynamic_module - -from QEfficient.generation.cloud_infer import QAICInferenceSession -from QEfficient.transformers.models.modeling_auto import QEffVisionEncoderForTextImageToTextModel from export_kimi_k25_vision import ( LOADED_EXPERT_IDS, MODEL_PATH, @@ -34,10 +29,17 @@ _patch_kimi_tie_weights_compat, _prepare_config, ) +from PIL import Image +from transformers.dynamic_module_utils import get_class_from_dynamic_module + +from QEfficient.generation.cloud_infer import QAICInferenceSession +from QEfficient.transformers.models.modeling_auto import QEffVisionEncoderForTextImageToTextModel def parse_args(): - parser = argparse.ArgumentParser(description="Verify Kimi K2.5 vision parity: HF PyTorch vs QEff PyTorch vs ONNXRuntime.") + parser = argparse.ArgumentParser( + description="Verify Kimi K2.5 vision parity: HF PyTorch vs QEff PyTorch vs ONNXRuntime." + ) parser.add_argument("--model-path", type=Path, default=MODEL_PATH) parser.add_argument("--num-vision-layers", type=int, default=NUM_VISION_LAYERS) parser.add_argument("--num-text-layers", type=int, default=NUM_TEXT_LAYERS) @@ -56,7 +58,9 @@ def parse_args(): action="store_true", help="Keep exported ONNX under ./aisand/onnx_export_verify instead of temp directory.", ) - parser.add_argument("--run-qpc", action="store_true", help="Compile and execute QPC, then compare QPC output vs HF.") + parser.add_argument( + "--run-qpc", action="store_true", help="Compile and execute QPC, then compare QPC output vs HF." + ) parser.add_argument( "--qaic-compile-bin", type=str, @@ -261,9 +265,6 @@ def _run_qpc(onnx_path: Path, pixel_values, h_shape, w_shape, args): all_ok = all_ok and hf_vs_qpc["allclose"] and qeff_vs_qpc["allclose"] and ort_vs_qpc["allclose"] if not all_ok: - raise SystemExit( - f"Parity check failed for atol={atol}, rtol={rtol}. " - f"See metrics above for details." - ) + raise SystemExit(f"Parity check failed for atol={atol}, rtol={rtol}. See metrics above for details.") - #print(f"PASS: requested parity checks match within atol={args.atol}, rtol={args.rtol}") \ No newline at end of file + print("Parity check passed.") From ad3ddfda20e65c83bba7985763228ff254b23ce0 Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Tue, 23 Jun 2026 13:00:48 +0530 Subject: [PATCH 08/33] vision + lang execute Signed-off-by: Mamta Singh --- .../models/kimi_k25/modeling_kimi_k25.py | 111 ++-- .../transformers/models/modeling_auto.py | 6 + examples/kimi_k2/export_kimi_k25_vision.py | 29 +- .../kimi_k2/verify_kimi_k25_vision_parity.py | 1 + .../verify_kimi_k25_vl_decoder_parity.py | 499 ++++++++++++++++++ 5 files changed, 594 insertions(+), 52 deletions(-) create mode 100644 examples/kimi_k2/verify_kimi_k25_vl_decoder_parity.py diff --git a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py index b9905602ff..7583a94a11 100644 --- a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py +++ b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py @@ -221,7 +221,7 @@ def __init__( self.proj = nn.Conv2d(in_dim, out_dim, kernel_size=patch_size, stride=patch_size) if pos_emb_type == "divided_fixed": - self.pos_emb = Learnable2DInterpPosEmbDivided_fixed( + self.pos_emb = QEffLearnable2DInterpPosEmbDivided_fixed( height=pos_emb_height, width=pos_emb_width, num_frames=pos_emb_time, dim=out_dim ) else: @@ -544,7 +544,7 @@ def __init__(self, config, *inputs, **kwargs): pos_emb_type=config.pos_emb_type, ) - self.encoder = MoonViT3dEncoder( + self.encoder = QEffMoonViT3dEncoder( hidden_dim=config.hidden_size, num_layers=config.num_hidden_layers, block_cfg={ @@ -707,16 +707,6 @@ def __init__(self, model): super().__init__() self.model = model self.config = self.model.config - # _orig_apply_rope = _kimi_module.apply_rope - # _kimi_module.apply_rope = _apply_rope_real - # _kimi_module = _sys.modules[type(model).__module__] - - # Restore original apply_rope and attention functions - # _kimi_module.apply_rope = _orig_apply_rope - # _kimi_module.VL_VISION_ATTENTION_FUNCTIONS.update(_orig_attn_functions) - - # _orig_attn_functions = _kimi_module.VL_VISION_ATTENTION_FUNCTIONS.copy() - # _kimi_module.VL_VISION_ATTENTION_FUNCTIONS["eager"] = _full_attention_forward def get_submodules_for_export(self) -> Type[nn.Module]: """ @@ -864,23 +854,39 @@ def forward( inputs_embeds = self.model.get_input_embeddings()(input_ids) if image_idx is None: - image_idx = torch.zeros((1, 1), dtype=torch.int64, device=inputs_embeds.device) - - if image_embeds is not None and input_ids is not None and input_ids.shape[1] != 1: - if image_embeds.dim() == 2: - image_embeds = image_embeds.unsqueeze(0) + image_idx = torch.zeros((inputs_embeds.shape[0], 1), dtype=torch.int64, device=inputs_embeds.device) + + image_embeds_for_state = None + if image_embeds is not None: + if image_embeds.dim() == 3: + if image_embeds.shape[0] != 1: + raise ValueError(f"Expected image_embeds batch dim to be 1, got shape {tuple(image_embeds.shape)}") + image_embeds_for_state = image_embeds[0].to(device=inputs_embeds.device, dtype=inputs_embeds.dtype) + elif image_embeds.dim() == 2: + image_embeds_for_state = image_embeds.to(device=inputs_embeds.device, dtype=inputs_embeds.dtype) + else: + raise ValueError(f"Expected image_embeds rank 2 or 3, got {image_embeds.dim()}.") - image_embeds = image_embeds.to(device=inputs_embeds.device, dtype=inputs_embeds.dtype) - _, _, hidden_size = inputs_embeds.shape + if image_embeds_for_state is not None and input_ids is not None and input_ids.shape[1] != 1: selected = input_ids == self.config.media_placeholder_token_id - indices1 = selected.to(torch.int64).cumsum(1) - 1 - indices1 = torch.where(indices1 != -1, indices1 + image_idx.to(indices1.device), indices1) - indices0 = torch.arange(selected.shape[0], device=selected.device).view(-1, 1) - safe_indices1 = torch.where(indices1 < 0, torch.zeros_like(indices1), indices1) - image_features_expanded = image_embeds.reshape(-1, hidden_size).unsqueeze(0)[indices0, safe_indices1] - image_input_embeds = torch.where(selected.unsqueeze(-1), image_features_expanded, inputs_embeds) - inputs_embeds = torch.where(input_ids.shape[1] == torch.tensor(1), inputs_embeds, image_input_embeds) - image_idx = (indices1.max() + 1).reshape(1, 1) + if selected.any(): + if attention_mask is None: + attention_mask = torch.ones_like(input_ids, dtype=torch.long, device=input_ids.device) + + image_features = [image_embeds_for_state] + inputs_embeds = inputs_embeds.to(image_embeds_for_state.dtype) + inputs_embeds, attention_mask, _, position_ids = self.model._merge_input_ids_with_image_features( + image_features=image_features, + inputs_embeds=inputs_embeds, + input_ids=input_ids, + attention_mask=attention_mask, + labels=None, + ) + image_idx = image_idx + torch.tensor( + [[sum(feature.shape[0] for feature in image_features)]], + dtype=torch.int64, + device=image_idx.device, + ) outputs = self.model.language_model( inputs_embeds=inputs_embeds, @@ -897,9 +903,17 @@ def forward( # QEff Deepseek language_model.forward already returns final logits. logits = outputs[0].float() - output_compressed_kvs = getattr(outputs, "compressed_kvs", None) - output_past_key_values = getattr(outputs, "past_key_values", None) - return logits, image_embeds, image_idx, output_compressed_kvs, output_past_key_values + mla_absorption = getattr(self.model, "mla_absorption", None) + if mla_absorption is not None: + cache_compressed = mla_absorption.get("cache_compressed", False) + else: + cache_compressed = False + + if cache_compressed: + output_kvs = getattr(outputs, "compressed_kvs", None) + else: + output_kvs = getattr(outputs, "past_key_values", None) + return logits, image_embeds_for_state, image_idx, output_kvs # ref https://github.com/huggingface/transformers/blob/78b2929c0554b79e0489b451ce4ece14d265ead2/src/transformers/models/llava/modeling_llava.py#L240 @@ -1308,9 +1322,13 @@ def get_output_names(self, kv_offload: bool = False): output_names = {} if kv_offload: + lang_output_names.insert(1, "image_embeds_RetainedState") + lang_output_names.insert(2, "image_idx_output") output_names["vision"] = vision_output_names output_names["lang"] = lang_output_names else: + lang_output_names.insert(1, "pixel_values_RetainedState") + lang_output_names.insert(2, "image_idx_output") return lang_output_names return output_names @@ -1401,6 +1419,18 @@ def get_dummy_inputs( torch.zeros(pkv_cache[0][1].shape, dtype=self.language_model.config.torch_dtype) ) + inputs_shapes = {} + # Decoder expects image embeddings indexed by token position. Keep this as + # [num_image_tokens, hidden_size] so `num_image_tokens` maps to axis 0. + inputs_shapes["image_embeds"] = (600, 7168) + inputs_shapes["pixel_values"] = (2400, 3, 14, 14) + inputs_shapes["image_idx"] = (1, 1) + lang_inputs["image_embeds"] = torch.zeros( + inputs_shapes["image_embeds"], + dtype=self.language_model.config.torch_dtype, + ) + lang_inputs["image_idx"] = torch.zeros(inputs_shapes["image_idx"], dtype=torch.int64) + if continuous_batching: lang_inputs["batch_index"] = torch.arange(bs).view(bs, 1) @@ -1482,12 +1512,17 @@ def get_specializations( prefill_seq_len = prefill_seq_len if prefill_seq_len else 32 # constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN ctx_len = ctx_len if ctx_len else 32 # constants.ONNX_EXPORT_EXAMPLE_CTX_LEN + num_patches = num_patches if num_patches is not None else 2400 + h = h if h is not None else 30 + w = w if w is not None else 80 + num_image_tokens = num_image_tokens if num_image_tokens is not None else 600 + vision = [ { - "num_patches": 2400, # num_patches - "h": 30, # h - "w": 80, # w - "num_image_tokens": 600, # num_image_tokens + "num_patches": num_patches, + "h": h, + "w": w, + "num_image_tokens": num_image_tokens, } ] @@ -1499,7 +1534,7 @@ def get_specializations( "batch_size": 1 if continuous_batching else batch_size, "seq_len": prefill_seq_len, "ctx_len": ctx_len, - "num_image_tokens": 600, # num_image_tokens + "num_image_tokens": num_image_tokens, } if continuous_batching: lang_prefill["full_batch_size"] = kv_cache_batch_size @@ -1515,7 +1550,7 @@ def get_specializations( "batch_size": full_batch_size if continuous_batching else batch_size, "seq_len": "1", "ctx_len": ctx_len, - "num_image_tokens": 600, # num_image_tokens + "num_image_tokens": num_image_tokens, } if continuous_batching: @@ -1530,7 +1565,7 @@ def get_specializations( "batch_size": 1 if continuous_batching else batch_size, "seq_len": prefill_seq_len, "ctx_len": ctx_len, - "num_image_tokens": 600, # num_image_tokens + "num_image_tokens": num_image_tokens, } if continuous_batching: lang_prefill["full_batch_size"] = kv_cache_batch_size @@ -1543,7 +1578,7 @@ def get_specializations( "batch_size": full_batch_size if continuous_batching else batch_size, "seq_len": 1, "ctx_len": ctx_len, - "num_image_tokens": 600, # num_image_tokens + "num_image_tokens": num_image_tokens, } if continuous_batching: diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index 7afe944e62..15e98a6b88 100755 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -1233,6 +1233,11 @@ def __init__(self, model, qaic_config: Optional[dict] = None, **kwargs): self.model.qaic_config = qaic_config self.hash_params["qeff_auto_class"] = self.__class__.__name__ self.continuous_batching = False + if qaic_config: + self.ccl_enabled = qaic_config.get("ccl_enabled", False) + if mla_absorption := qaic_config.get("mla_absorption", None): + self.hash_params["mla_absorption"] = mla_absorption + setattr(self.model.language_model, "mla_absorption", mla_absorption) def __update_prefill_transform( self, @@ -1435,6 +1440,7 @@ def __init__( self.ccl_enabled = False if qaic_config: self.ccl_enabled = qaic_config.get("ccl_enabled", False) + self.comp_ctx_lengths_prefill, self.comp_ctx_lengths_decode = None, None self.input_shapes, self.output_names = None, None # ---Sampling--- diff --git a/examples/kimi_k2/export_kimi_k25_vision.py b/examples/kimi_k2/export_kimi_k25_vision.py index 70d58cb1bc..b877235f69 100644 --- a/examples/kimi_k2/export_kimi_k25_vision.py +++ b/examples/kimi_k2/export_kimi_k25_vision.py @@ -374,9 +374,9 @@ def main(): else: raise ValueError("Pass both --num-vision-layers and --num-text-layers to load a layer subset.") - # qaic_config = {"mla_absorption": {"cache_compressed": True, "absorption": False, "online": False}} + qaic_config = {"mla_absorption": {"cache_compressed": True, "absorption": False, "online": False}} - qeff_model = QEFFAutoModelForImageTextToText(model) + qeff_model = QEFFAutoModelForImageTextToText(model, qaic_config=qaic_config) skip_vision = False @@ -386,9 +386,9 @@ def main(): ## STEP 3: Compile Model for Text-Only Execution # Set skip_vision=True to bypass image processing qeff_model.compile( - # qaic_config=qaic_config, + qaic_config=qaic_config, prefill_seq_len=1, - ctx_len=1024, + ctx_len=4096, num_cores=16, num_devices=1, mxfp6_matmul=False, @@ -407,7 +407,7 @@ def main(): { "role": "user", "content": [ - {"type": "text", "text": "Tell me about yourself."}, + {"type": "text", "text": args.prompt}, ], }, ] @@ -435,23 +435,23 @@ def main(): ## STEP 3: Compile Model for Vision+Text Execution # Do not set skip_vision (defaults to False) to enable image processing qeff_model.compile( - # qaic_config=qaic_config, + qaic_config=qaic_config, prefill_seq_len=1, - ctx_len=1024, + ctx_len=4096, num_cores=16, num_devices=1, mxfp6_matmul=False, mxint8_kv_cache=False, aic_enable_depth_first=False, - # skip_lang=True, mos=1, - num_patches=2400, # num_patches - h=30, # h - w=80, # w - num_image_tokens=600, # num_image_tokens + num_patches=2400, + h=30, + w=80, + # Keep language-side image embedding specialization tightly bounded to + # actual single-image token count to avoid oversized dynamic VA mapping. + num_image_tokens=600, ) - # exit() ## STEP 4: Prepare Image and Text Input image = Image.open(BytesIO(requests.get(args.image_url).content)).convert("RGB") @@ -476,8 +476,9 @@ def main(): # Convert pixel values to float32 for processing inputs["pixel_values"] = inputs["pixel_values"].to(qeff_model.model.config.torch_dtype) + breakpoint() ## STEP 6: Run Vision+Text Inference - output = qeff_model.generate(inputs=inputs, device_ids=[0, 1], generation_len=100) + output = qeff_model.generate(inputs=inputs, device_ids=[0], generation_len=10) ## STEP 7: Display Results print(output.generated_ids) diff --git a/examples/kimi_k2/verify_kimi_k25_vision_parity.py b/examples/kimi_k2/verify_kimi_k25_vision_parity.py index 729b22edec..034ea2ca72 100644 --- a/examples/kimi_k2/verify_kimi_k25_vision_parity.py +++ b/examples/kimi_k2/verify_kimi_k25_vision_parity.py @@ -147,6 +147,7 @@ def _run_qpc(onnx_path: Path, pixel_values, h_shape, w_shape, args): qpc_dir.parent.mkdir(parents=True, exist_ok=True) else: qpc_dir = Path(tempfile.mkdtemp(prefix="kimi_k25_verify_qpc_")) + shutil.rmtree(qpc_dir) compile_cmd = [ args.qaic_compile_bin, diff --git a/examples/kimi_k2/verify_kimi_k25_vl_decoder_parity.py b/examples/kimi_k2/verify_kimi_k25_vl_decoder_parity.py new file mode 100644 index 0000000000..83c4ea9d72 --- /dev/null +++ b/examples/kimi_k2/verify_kimi_k25_vl_decoder_parity.py @@ -0,0 +1,499 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import argparse +import copy +import os +import shutil +import subprocess +import sys +import tempfile +from io import BytesIO +from pathlib import Path + +import numpy as np +import onnxruntime as ort +import requests +import torch +from export_kimi_k25_vision import ( + LOADED_EXPERT_IDS, + MODEL_PATH, + NUM_EXPERTS_PER_TOKEN, + _ensure_torch_fx_import_compatibility, + _load_layer_subset_model, + _patch_deepseek_init_weights_compat, + _patch_kimi_tie_weights_compat, + _prepare_config, +) +from huggingface_hub import snapshot_download +from PIL import Image +from transformers.dynamic_module_utils import get_class_from_dynamic_module + +from QEfficient import QEFFAutoModelForImageTextToText +from QEfficient.generation.cloud_infer import QAICInferenceSession +from QEfficient.transformers.models.modeling_auto import QEffCausalLMForTextImageToTextModel + + +def parse_args(): + parser = argparse.ArgumentParser( + description=("Verify Kimi-K2.5 V+L decoder parity on CPU with a 2-layer vision + 2-layer text subset.") + ) + parser.add_argument( + "--model-path", + type=str, + default="moonshotai/Kimi-K2.5", + help=( + "Local model path or HF repo id (defaults to moonshotai/Kimi-K2.5; " + "downloads/resolves using HF_HUB_CACHE when needed)." + ), + ) + parser.add_argument("--num-text-layers", type=int, default=2) + parser.add_argument("--num-vision-layers", type=int, default=2) + parser.add_argument("--expert-ids", type=str, default=",".join(str(x) for x in LOADED_EXPERT_IDS)) + parser.add_argument("--num-experts-per-token", type=int, default=NUM_EXPERTS_PER_TOKEN) + parser.add_argument( + "--image-url", + type=str, + default="https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png", + ) + parser.add_argument("--prompt", type=str, default="Describe this image in one sentence.") + parser.add_argument("--atol", type=float, default=1e-4) + parser.add_argument("--rtol", type=float, default=1e-4) + parser.add_argument("--skip-onnx", action="store_true", help="Skip ONNX export/runtime parity check.") + parser.add_argument( + "--keep-onnx", + action="store_true", + help="Keep exported ONNX under ./aisand/onnx_export_verify_vl instead of a temporary directory.", + ) + parser.add_argument("--run-qpc", action="store_true", help="Compile and execute QPC, then compare vs references.") + parser.add_argument( + "--qaic-compile-bin", + type=str, + default="/opt/qti-aic/exec/qaic-compile", + help="Path to qaic-compile binary.", + ) + parser.add_argument( + "--keep-qpc", + action="store_true", + help="Keep compiled QPC under --qpc-dir (used only with --run-qpc).", + ) + parser.add_argument( + "--qpc-dir", + type=Path, + default=Path("aisand/qpc_verify_vl"), + help="Output directory for QPC when --keep-qpc is set.", + ) + parser.add_argument("--aic-num-cores", type=int, default=16, help="Number of QAIC cores for qpc compile.") + parser.add_argument("--mos", type=int, default=1, help="-mos value for qpc compile.") + parser.add_argument( + "--no-convert-to-fp16", + action="store_true", + help="Disable -convert-to-fp16 during qpc compile (defaults to enabled).", + ) + parser.add_argument("--qpc-atol", type=float, default=1e-2) + parser.add_argument("--qpc-rtol", type=float, default=1e-2) + return parser.parse_args() + + +def _parse_expert_ids(value: str): + return tuple(int(x.strip()) for x in value.split(",") if x.strip()) + + +def _resolve_model_path(model_path_or_id: str) -> Path: + candidate = Path(model_path_or_id) + if candidate.exists(): + return candidate + + cache_dir = os.environ.get("HF_HUB_CACHE") + return Path(snapshot_download(repo_id=model_path_or_id, cache_dir=cache_dir)) + + +def _stats(a: torch.Tensor, b: torch.Tensor, atol: float, rtol: float): + diff = (a - b).abs() + return { + "shape": tuple(a.shape), + "max_abs": float(diff.max().item()), + "mean_abs": float(diff.mean().item()), + "rmse": float(torch.sqrt(torch.mean((a - b) ** 2)).item()), + "allclose": bool(torch.allclose(a, b, atol=atol, rtol=rtol)), + } + + +def _hf_image_embeds(model, pixel_values: torch.Tensor, grid_thws: torch.Tensor) -> torch.Tensor: + target_dtype = model.vision_tower.patch_embed.proj.weight.dtype + pixel_values = pixel_values.to(target_dtype) + image_features = model.vision_tower(pixel_values, grid_thws) + if model.mm_projector: + image_features = model.mm_projector(image_features) + return image_features[0] if isinstance(image_features, list) else image_features + + +def _get_onnx_export_dir(keep_onnx: bool): + if keep_onnx: + base_dir = Path("aisand") + base_dir.mkdir(parents=True, exist_ok=True) + for old in base_dir.glob("onnx_export_verify_vl-*"): + shutil.rmtree(old, ignore_errors=True) + export_dir = base_dir / "onnx_export_verify_vl" + if export_dir.exists(): + shutil.rmtree(export_dir, ignore_errors=True) + export_dir.mkdir(parents=True, exist_ok=True) + return export_dir + return Path(tempfile.mkdtemp(prefix="kimi_k25_verify_vl_onnx_")) + + +def _export_and_run_lang_onnx( + transformed_model, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + past_key_values, + keep_onnx: bool, +): + lang_exporter = QEffCausalLMForTextImageToTextModel(copy.deepcopy(transformed_model)) + lang_inputs = transformed_model.get_dummy_inputs(kv_offload=True)["lang"] + output_names = transformed_model.get_output_names(kv_offload=True)["lang"] + dynamic_axes = transformed_model.get_onnx_dynamic_axes(kv_offload=True)["lang"] + + export_dir = _get_onnx_export_dir(keep_onnx) + onnx_path = lang_exporter.export( + inputs=lang_inputs, + output_names=output_names, + dynamic_axes=dynamic_axes, + export_dir=export_dir, + offload_pt_weights=False, + ) + + session = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"]) + ort_inputs = { + "input_ids": input_ids.detach().cpu().numpy().astype(np.int64), + "position_ids": position_ids.detach().cpu().numpy().astype(np.int64), + } + for layer_idx, (past_key, past_value) in enumerate(past_key_values): + ort_inputs[f"past_key.{layer_idx}"] = past_key.detach().cpu().numpy().astype(np.float32) + ort_inputs[f"past_value.{layer_idx}"] = past_value.detach().cpu().numpy().astype(np.float32) + + ort_logits = torch.from_numpy(session.run(["logits"], ort_inputs)[0]).float() + return ort_logits, Path(onnx_path) + + +def _prepare_qpc_dir(args, suffix: str): + if args.keep_qpc: + qpc_dir = args.qpc_dir / suffix + if qpc_dir.exists(): + shutil.rmtree(qpc_dir) + qpc_dir.parent.mkdir(parents=True, exist_ok=True) + else: + qpc_dir = Path(tempfile.mkdtemp(prefix=f"kimi_k25_verify_vl_{suffix}_qpc_")) + shutil.rmtree(qpc_dir) + return qpc_dir + + +def _run_vision_qpc(onnx_path: Path, pixel_values: torch.Tensor, h_shape: torch.Tensor, w_shape: torch.Tensor, args): + qpc_dir = _prepare_qpc_dir(args, "vision") + compile_cmd = [ + args.qaic_compile_bin, + f"-m={onnx_path}", + f"-onnx-define-symbol=num_patches,{pixel_values.shape[0]}", + f"-onnx-define-symbol=h,{h_shape.shape[0]}", + f"-onnx-define-symbol=w,{w_shape.shape[0]}", + f"-aic-num-cores={args.aic_num_cores}", + f"-mos={args.mos}", + f"-aic-binary-dir={qpc_dir}", + ] + if not args.no_convert_to_fp16: + compile_cmd.insert(1, "-convert-to-fp16") + + compile_result = subprocess.run(compile_cmd, capture_output=True, text=True) + if compile_result.returncode != 0: + raise RuntimeError( + "Vision QPC compile failed.\n" + f"Command: {' '.join(compile_cmd)}\n" + f"stdout:\n{compile_result.stdout}\n" + f"stderr:\n{compile_result.stderr}" + ) + + session = QAICInferenceSession(str(qpc_dir)) + qpc_inputs = {} + for name in session.input_names: + if name == "pixel_values": + qpc_inputs[name] = pixel_values.detach().cpu().numpy().astype(np.float32) + elif name == "h_shape": + qpc_inputs[name] = h_shape.detach().cpu().numpy().astype(np.int64) + elif name == "w_shape": + qpc_inputs[name] = w_shape.detach().cpu().numpy().astype(np.int64) + else: + raise RuntimeError(f"Unexpected vision QPC input name: {name}") + + qpc_outputs = session.run(qpc_inputs) + output_name = session.output_names[0] + qpc_out = qpc_outputs[output_name] + return torch.from_numpy(qpc_out).float(), qpc_dir + + +def _run_lang_qpc(onnx_path: Path, input_ids: torch.Tensor, position_ids: torch.Tensor, past_key_values, args): + qpc_dir = _prepare_qpc_dir(args, "lang") + compile_cmd = [ + args.qaic_compile_bin, + f"-m={onnx_path}", + f"-onnx-define-symbol=batch_size,{input_ids.shape[0]}", + f"-onnx-define-symbol=seq_len,{input_ids.shape[1]}", + f"-onnx-define-symbol=ctx_len,{past_key_values[0][0].shape[2]}", + f"-aic-num-cores={args.aic_num_cores}", + f"-mos={args.mos}", + f"-aic-binary-dir={qpc_dir}", + ] + if not args.no_convert_to_fp16: + compile_cmd.insert(1, "-convert-to-fp16") + + compile_result = subprocess.run(compile_cmd, capture_output=True, text=True) + if compile_result.returncode != 0: + raise RuntimeError( + "QPC compile failed.\n" + f"Command: {' '.join(compile_cmd)}\n" + f"stdout:\n{compile_result.stdout}\n" + f"stderr:\n{compile_result.stderr}" + ) + + session = QAICInferenceSession(str(qpc_dir)) + qpc_inputs = { + "input_ids": input_ids.detach().cpu().numpy().astype(np.int64), + "position_ids": position_ids.detach().cpu().numpy().astype(np.int64), + } + for layer_idx, (past_key, past_value) in enumerate(past_key_values): + qpc_inputs[f"past_key.{layer_idx}"] = past_key.detach().cpu().numpy().astype(np.float32) + qpc_inputs[f"past_value.{layer_idx}"] = past_value.detach().cpu().numpy().astype(np.float32) + + qpc_feed = {name: qpc_inputs[name] for name in session.input_names} + qpc_outputs = session.run(qpc_feed) + output_name = "logits" if "logits" in qpc_outputs else session.output_names[0] + qpc_logits = torch.from_numpy(qpc_outputs[output_name]).float() + return qpc_logits, qpc_dir + + +if __name__ == "__main__": + args = parse_args() + if args.model_path is None: + args.model_path = MODEL_PATH + args.model_path = _resolve_model_path(str(args.model_path)) + expert_ids = _parse_expert_ids(args.expert_ids) + + torch.set_grad_enabled(False) + _ensure_torch_fx_import_compatibility() + + config = _prepare_config(args.model_path) + kimi_cls = get_class_from_dynamic_module("modeling_kimi_k25.KimiK25ForConditionalGeneration", str(args.model_path)) + _patch_kimi_tie_weights_compat(kimi_cls) + _patch_deepseek_init_weights_compat(kimi_cls) + + hf_model, tokenizer, processor = _load_layer_subset_model( + model_path=args.model_path, + kimi_cls=kimi_cls, + config=config, + num_vision_layers=args.num_vision_layers, + num_text_layers=args.num_text_layers, + loaded_expert_ids=expert_ids, + num_experts_per_tok=args.num_experts_per_token, + dtype=torch.float32, + ) + # Required for HW compatibility and stable vision parity in this flow. + hf_model.vision_tower.patch_embed.pos_emb.interpolation_mode = "bilinear" + hf_model = hf_model.eval() + + transformed_wrapper = QEFFAutoModelForImageTextToText(copy.deepcopy(hf_model)) + transformed_model = transformed_wrapper.model.eval() + decoder_wrapper = transformed_model.get_qeff_language_decoder().eval() + + image = Image.open(BytesIO(requests.get(args.image_url, timeout=30).content)).convert("RGB") + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": image}, + {"type": "text", "text": args.prompt}, + ], + } + ] + inputs = processor(messages=messages, add_generation_prompt=True, tokenize=False, return_tensors="pt") + + input_ids = inputs["input_ids"].to(torch.long) + attention_mask = inputs["attention_mask"].to(torch.long) + pixel_values = inputs["pixel_values"].to(torch.float32) + grid_thws = inputs["grid_thws"].to(torch.long) + h_shape = torch.ones(int(grid_thws[0, 1].item()), dtype=torch.int64) + w_shape = torch.ones(int(grid_thws[0, 2].item()), dtype=torch.int64) + position_ids = torch.where(attention_mask.bool(), attention_mask.cumsum(-1) - 1, -1) + + image_features = transformed_model._extract_image_features(pixel_values, grid_thws) + if transformed_model.mm_projector: + image_features = transformed_model.mm_projector(image_features) + image_embeds = image_features[0] if isinstance(image_features, list) else image_features + vision_reference = _hf_image_embeds(hf_model, pixel_values, grid_thws).detach().cpu().float() + + vision_onnx_path = None + if not args.skip_onnx or args.run_qpc: + vision_cmd = [ + str(Path(sys.executable)), + str(Path(__file__).with_name("verify_kimi_k25_vision_parity.py")), + "--model-path", + str(args.model_path), + "--num-vision-layers", + str(args.num_vision_layers), + "--num-text-layers", + str(args.num_text_layers), + "--expert-ids", + args.expert_ids, + "--num-experts-per-token", + str(args.num_experts_per_token), + "--image-url", + args.image_url, + "--prompt", + args.prompt, + "--atol", + str(args.atol), + "--rtol", + str(args.rtol), + ] + keep_vision_onnx = args.keep_onnx or args.run_qpc + if keep_vision_onnx: + vision_cmd.append("--keep-onnx") + subprocess.run(vision_cmd, check=True) + + if keep_vision_onnx: + vision_matches = sorted(Path("aisand").glob("onnx_export_verify-*/KimiK25EncoderWrapper.onnx")) + if not vision_matches: + raise RuntimeError("Vision ONNX not found under aisand/onnx_export_verify-*/KimiK25EncoderWrapper.onnx") + vision_onnx_path = vision_matches[-1] + + merged_inputs_embeds, merged_attention_mask, _, merged_position_ids = ( + transformed_model._merge_input_ids_with_image_features( + image_features=image_features, + inputs_embeds=transformed_model.get_input_embeddings()(input_ids), + input_ids=input_ids, + attention_mask=attention_mask, + labels=None, + ) + ) + + past_key_values = transformed_model.language_model.get_dummy_pkv_cache( + transformed_model.config.text_config, + input_ids.shape[0], + merged_inputs_embeds.shape[1], + ) + + model_logits = ( + transformed_model.language_model( + inputs_embeds=merged_inputs_embeds, + attention_mask=merged_attention_mask, + position_ids=merged_position_ids, + past_key_values=copy.deepcopy(past_key_values), + use_cache=True, + )[0] + .detach() + .cpu() + .float() + ) + + decoder_logits = ( + decoder_wrapper( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + image_embeds=image_embeds, + image_idx=torch.zeros((input_ids.shape[0], 1), dtype=torch.int64), + past_key_values=copy.deepcopy(past_key_values), + use_cache=True, + )[0] + .detach() + .cpu() + .float() + ) + + decoder_vs_model = _stats(model_logits, decoder_logits, args.atol, args.rtol) + + print( + "Loaded subset:", + f"vision_layers={transformed_model.config.vision_config.vt_num_hidden_layers}", + f"text_layers={transformed_model.config.text_config.num_hidden_layers}", + ) + print("Tokenizer vocab size:", tokenizer.vocab_size) + print("Prompt:", repr(args.prompt)) + print("Input shape:", tuple(input_ids.shape)) + print("Model logits vs decoder-wrapper logits:", decoder_vs_model) + + if not decoder_vs_model["allclose"]: + raise SystemExit( + "Parity check failed: decoder-wrapper logits do not match transformed model logits for image+text input." + ) + + # ONNX/QPC parity runs the language decoder in decode mode (seq_len=1) + # with multimodal context represented in past_key_values. + onnx_input_ids = input_ids[:, -1:] + onnx_position_ids = merged_position_ids[:, -1:] + onnx_reference_logits = ( + transformed_model.language_model( + inputs_embeds=transformed_model.get_input_embeddings()(onnx_input_ids), + position_ids=onnx_position_ids, + past_key_values=copy.deepcopy(past_key_values), + use_cache=True, + )[0] + .detach() + .cpu() + .float() + ) + + lang_ort = None + lang_onnx_path = None + if not args.skip_onnx or args.run_qpc: + lang_ort, lang_onnx_path = _export_and_run_lang_onnx( + transformed_model, + input_ids=onnx_input_ids, + position_ids=onnx_position_ids, + past_key_values=copy.deepcopy(past_key_values), + keep_onnx=args.keep_onnx, + ) + + if vision_onnx_path is not None: + print("Vision ONNX path:", vision_onnx_path) + + if lang_ort is not None: + lang_ref_vs_onnx = _stats(onnx_reference_logits, lang_ort, args.atol, args.rtol) + print("Language ONNX path:", lang_onnx_path) + print("Decode-step reference vs ONNXRuntime:", lang_ref_vs_onnx) + + if not lang_ref_vs_onnx["allclose"]: + raise SystemExit("Parity check failed: language ONNXRuntime logits do not match decode-step reference.") + + if args.run_qpc: + vision_qpc, vision_qpc_dir = _run_vision_qpc( + vision_onnx_path, + pixel_values=pixel_values, + h_shape=h_shape, + w_shape=w_shape, + args=args, + ) + vision_ref_vs_qpc = _stats(vision_reference, vision_qpc, args.qpc_atol, args.qpc_rtol) + print("Vision QPC dir:", vision_qpc_dir) + print(f"Vision reference vs QPC (atol={args.qpc_atol}, rtol={args.qpc_rtol}):", vision_ref_vs_qpc) + + if not vision_ref_vs_qpc["allclose"]: + raise SystemExit("Parity check failed: vision QPC output does not match vision reference.") + + lang_qpc, lang_qpc_dir = _run_lang_qpc( + lang_onnx_path, + input_ids=onnx_input_ids, + position_ids=onnx_position_ids, + past_key_values=copy.deepcopy(past_key_values), + args=args, + ) + lang_ref_vs_qpc = _stats(onnx_reference_logits, lang_qpc, args.qpc_atol, args.qpc_rtol) + print("Language QPC dir:", lang_qpc_dir) + print(f"Decode-step reference vs QPC (atol={args.qpc_atol}, rtol={args.qpc_rtol}):", lang_ref_vs_qpc) + + if not lang_ref_vs_qpc["allclose"]: + raise SystemExit("Parity check failed: language QPC logits do not match decode-step reference.") + + print("Parity check passed.") From 68c58116075385a59359eff687168bb48f130b37 Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Wed, 24 Jun 2026 13:28:12 +0530 Subject: [PATCH 09/33] add tests Signed-off-by: Mamta Singh --- .../models/deepseek_v3/modeling_deepseek.py | 1 + .../models/kimi_k25/configuration_kimi_k25.py | 7 + .../transformers/models/modeling_auto.py | 2 +- examples/kimi_k2/export_kimi_k25_vision.py | 14 +- examples/kimi_k2/test_kimi_k25.py | 185 +++++++ examples/kimi_k2/test_kimi_k25_text.py | 430 +++++++++++++++ ...verify_kimi_k25_language_decoder_parity.py | 327 ------------ .../verify_kimi_k25_vl_decoder_parity.py | 499 ------------------ 8 files changed, 636 insertions(+), 829 deletions(-) create mode 100644 examples/kimi_k2/test_kimi_k25.py create mode 100644 examples/kimi_k2/test_kimi_k25_text.py delete mode 100644 examples/kimi_k2/verify_kimi_k25_language_decoder_parity.py delete mode 100644 examples/kimi_k2/verify_kimi_k25_vl_decoder_parity.py diff --git a/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py b/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py index daddfe048f..8a8b3efcf9 100644 --- a/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py +++ b/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py @@ -794,6 +794,7 @@ def __qeff_init__( self.all_up_proj = torch.nn.Parameter( torch.cat( [_get_linear_weight(exp.up_proj).T.unsqueeze(0) for exp in self.experts], + dim=0, ) ) self.all_down_proj = torch.nn.Parameter( diff --git a/QEfficient/transformers/models/kimi_k25/configuration_kimi_k25.py b/QEfficient/transformers/models/kimi_k25/configuration_kimi_k25.py index 74c9a30a8b..dde386b53a 100644 --- a/QEfficient/transformers/models/kimi_k25/configuration_kimi_k25.py +++ b/QEfficient/transformers/models/kimi_k25/configuration_kimi_k25.py @@ -1,3 +1,10 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + from transformers.configuration_utils import PretrainedConfig try: diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index 15e98a6b88..74c2c54348 100755 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -1440,7 +1440,7 @@ def __init__( self.ccl_enabled = False if qaic_config: self.ccl_enabled = qaic_config.get("ccl_enabled", False) - + self.comp_ctx_lengths_prefill, self.comp_ctx_lengths_decode = None, None self.input_shapes, self.output_names = None, None # ---Sampling--- diff --git a/examples/kimi_k2/export_kimi_k25_vision.py b/examples/kimi_k2/export_kimi_k25_vision.py index b877235f69..c5fc366c1e 100644 --- a/examples/kimi_k2/export_kimi_k25_vision.py +++ b/examples/kimi_k2/export_kimi_k25_vision.py @@ -36,6 +36,16 @@ EXPERT_KEY_PATTERN = re.compile(r"^(language_model\.model\.layers\.\d+\.mlp\.experts\.)(\d+)(\..+)$") +def _set_deterministic(seed: int): + import random + import numpy as np + + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + torch.use_deterministic_algorithms(True) def _ensure_torch_fx_import_compatibility(): """Backfill `is_torch_fx_available` for remote model code expecting older Transformers APIs.""" @@ -331,8 +341,9 @@ def parse_args(): def main(): args = parse_args() - _ensure_torch_fx_import_compatibility() + _set_deterministic(1234) + _ensure_torch_fx_import_compatibility() config = _prepare_config(args.model_path) kimi_cls = get_class_from_dynamic_module( "modeling_kimi_k25.KimiK25ForConditionalGeneration", @@ -476,7 +487,6 @@ def main(): # Convert pixel values to float32 for processing inputs["pixel_values"] = inputs["pixel_values"].to(qeff_model.model.config.torch_dtype) - breakpoint() ## STEP 6: Run Vision+Text Inference output = qeff_model.generate(inputs=inputs, device_ids=[0], generation_len=10) diff --git a/examples/kimi_k2/test_kimi_k25.py b/examples/kimi_k2/test_kimi_k25.py new file mode 100644 index 0000000000..821b6434c7 --- /dev/null +++ b/examples/kimi_k2/test_kimi_k25.py @@ -0,0 +1,185 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import os +from io import BytesIO +from pathlib import Path + +import requests +import torch +from export_kimi_k25_vision import ( + LOADED_EXPERT_IDS, + NUM_EXPERTS_PER_TOKEN, + _ensure_torch_fx_import_compatibility, + _load_layer_subset_model, + _patch_deepseek_init_weights_compat, + _patch_kimi_tie_weights_compat, + _prepare_config, +) +from huggingface_hub import snapshot_download +from PIL import Image +from transformers.dynamic_module_utils import get_class_from_dynamic_module + +from QEfficient import QEFFAutoModelForImageTextToText + +MODEL_NAME = "moonshotai/Kimi-K2.5" +IMAGE_URL = "https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png" +TEXT_PROMPT = "Describe this image." +NUM_VISION_LAYERS = 2 +NUM_TEXT_LAYERS = 2 +NEW_GENERATION_TOKENS = 10 +CTX_LEN = 4096 + +def _set_deterministic(seed: int): + import random + import numpy as np + + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + torch.use_deterministic_algorithms(True) + +def _resolve_model_path() -> Path: + return Path(snapshot_download(repo_id=MODEL_NAME, cache_dir=os.environ.get("HF_HUB_CACHE"))) + +def _prepare_inputs(processor): + image = Image.open(BytesIO(requests.get(IMAGE_URL, timeout=30).content)).convert("RGB") + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": image}, + {"type": "text", "text": TEXT_PROMPT}, + ], + } + ] + return processor( + messages=messages, + add_generation_prompt=True, + tokenize=False, + return_tensors="pt", + ) + +def _greedy_generate_cpu(model, inputs, max_new_tokens: int): + generated_ids = inputs["input_ids"] + attention_mask = inputs["attention_mask"] + pixel_values = inputs["pixel_values"] + grid_thws = inputs["grid_thws"] + + eos_token_id = getattr(model.config, "eos_token_id", None) + if eos_token_id is None and hasattr(model.config, "text_config"): + eos_token_id = getattr(model.config.text_config, "eos_token_id", None) + + for _ in range(max_new_tokens): + outputs = model( + input_ids=generated_ids, + attention_mask=attention_mask, + pixel_values=pixel_values, + grid_thws=grid_thws, + use_cache=False, + return_dict=True, + ) + logits = outputs[0] if isinstance(outputs, tuple) else outputs.logits + next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True) + + generated_ids = torch.cat([generated_ids, next_token], dim=1) + attention_mask = torch.cat( + [ + attention_mask, + torch.ones((attention_mask.shape[0], 1), dtype=attention_mask.dtype, device=attention_mask.device), + ], + dim=1, + ) + + if eos_token_id is not None and torch.all(next_token == eos_token_id): + break + + return generated_ids[:, -max_new_tokens:] + + +def check_kimi_k25_pytorch_vs_ai100(): + _set_deterministic(1234) + _ensure_torch_fx_import_compatibility() + model_path = _resolve_model_path() + config = _prepare_config(model_path) + kimi_cls = get_class_from_dynamic_module("modeling_kimi_k25.KimiK25ForConditionalGeneration", str(model_path)) + _patch_kimi_tie_weights_compat(kimi_cls) + _patch_deepseek_init_weights_compat(kimi_cls) + + model, tokenizer, processor = _load_layer_subset_model( + model_path=model_path, + kimi_cls=kimi_cls, + config=config, + num_vision_layers=NUM_VISION_LAYERS, + num_text_layers=NUM_TEXT_LAYERS, + loaded_expert_ids=LOADED_EXPERT_IDS, + num_experts_per_tok=NUM_EXPERTS_PER_TOKEN, + dtype=torch.float32, + ) + + model.vision_tower.patch_embed.pos_emb.interpolation_mode = "bilinear" + model = model.eval().to("cpu") + + inputs = _prepare_inputs(processor) + inputs = {k: (v.to("cpu") if torch.is_tensor(v) else v) for k, v in inputs.items()} + + with torch.inference_mode(): + generated_ids = _greedy_generate_cpu(model, inputs, max_new_tokens=NEW_GENERATION_TOKENS) + + decoded = tokenizer.batch_decode(generated_ids, skip_special_tokens=True) + + print( + "Loaded subset:", + f"vision_layers={model.config.vision_config.vt_num_hidden_layers}", + f"text_layers={model.config.text_config.num_hidden_layers}", + ) + print("Prompt:", repr(TEXT_PROMPT)) + print("Generated IDs shape:", tuple(generated_ids.shape)) + print("Output:") + print(decoded[0] if decoded else "") + + qeff_model = QEFFAutoModelForImageTextToText( + model, + kv_offload=True, + torch_dtype=torch.float32, + ) + + _ = qeff_model.export() + qeff_model.compile( + num_devices=1, + prefill_seq_len=1, # prompt_len, + ctx_len=CTX_LEN, + mxfp6=False, + num_patches=int(inputs["pixel_values"].shape[0]), + h=int(inputs["grid_thws"][0, 1].item()), + w=int(inputs["grid_thws"][0, 2].item()), + num_image_tokens=600, + num_cores=16, + ) + + exec_info = qeff_model.generate(inputs=inputs, generation_len=NEW_GENERATION_TOKENS) + cloud_ai_100_tokens = exec_info.generated_ids[:, :-1] + decoded = tokenizer.batch_decode(cloud_ai_100_tokens, skip_special_tokens=True) + + print( + "Loaded subset:", + f"vision_layers={model.config.vision_config.vt_num_hidden_layers}", + f"text_layers={model.config.text_config.num_hidden_layers}", + ) + print("Prompt:", repr(TEXT_PROMPT)) + print("Generated IDs shape:", tuple(cloud_ai_100_tokens.shape)) + print("Output:") + print(decoded[0] if decoded else "") + + assert (generated_ids == cloud_ai_100_tokens).all(), "Tokens don't match for pytorch HF output and QPC output" + + +def test_kimi_k25_image_text_to_text_pytorch_vs_ai100(): + torch.manual_seed(42) + check_kimi_k25_pytorch_vs_ai100() diff --git a/examples/kimi_k2/test_kimi_k25_text.py b/examples/kimi_k2/test_kimi_k25_text.py new file mode 100644 index 0000000000..6d4166a39f --- /dev/null +++ b/examples/kimi_k2/test_kimi_k25_text.py @@ -0,0 +1,430 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ---------------------------------------------------------------------------- + +import copy +import json +import re +import tempfile +from collections import defaultdict +from pathlib import Path + +import numpy as np +import pytest +import torch +from export_kimi_k25_vision import ( + _ensure_torch_fx_import_compatibility, + _patch_deepseek_init_weights_compat, + _patch_kimi_tie_weights_compat, +) +from safetensors import safe_open +from safetensors.torch import save_file +from transformers import AutoConfig, AutoTokenizer +from transformers.cache_utils import Cache, DynamicCache +from transformers.dynamic_module_utils import get_class_from_dynamic_module + +from QEfficient import QEFFAutoModelForCausalLM + +MODEL_PATH = Path( + "/home/huggingface_hub/models--moonshotai--Kimi-K2.5/snapshots/4d01dfe0332d63057c186e0b262165819efb6611" +) +NUM_HIDDEN_LAYERS = 2 +LOADED_EXPERT_IDS = (0, 1, 2, 3) +NUM_EXPERTS_PER_TOKEN = 2 + +EXPERT_KEY_PATTERN = re.compile(r"^(language_model\.model\.layers\.\d+\.mlp\.experts\.)(\d+)(\..+)$") + + +def _patch_dynamic_cache_compat(): + if not hasattr(DynamicCache, "from_legacy_cache"): + + @classmethod + def _from_legacy_cache(cls, legacy_cache): + if legacy_cache is None: + return cls() + if isinstance(legacy_cache, cls): + return legacy_cache + if isinstance(legacy_cache, Cache): + return cls(ddp_cache_data=[tuple(layer[:2]) for layer in legacy_cache]) + + ddp_cache_data = [] + for layer in legacy_cache: + if layer is None: + continue + if len(layer) < 2: + raise ValueError("Each legacy cache layer must provide key/value tensors.") + ddp_cache_data.append((layer[0], layer[1])) + return cls(ddp_cache_data=ddp_cache_data) + + DynamicCache.from_legacy_cache = _from_legacy_cache + + if not hasattr(DynamicCache, "to_legacy_cache"): + + def _to_legacy_cache(self): + return tuple((layer[0], layer[1]) for layer in self) + + DynamicCache.to_legacy_cache = _to_legacy_cache + + if not hasattr(DynamicCache, "get_max_length"): + + def _get_max_length(self): + return None + + DynamicCache.get_max_length = _get_max_length + + +def _validate_expert_subset(loaded_expert_ids, num_experts_per_tok, total_experts): + expert_ids = tuple(loaded_expert_ids) + if len(expert_ids) != 4: + raise ValueError(f"Expected exactly 4 routed experts, got {expert_ids!r}.") + if len(set(expert_ids)) != len(expert_ids): + raise ValueError(f"Expert ids must be unique, got {expert_ids!r}.") + invalid_ids = [expert_id for expert_id in expert_ids if expert_id < 0 or expert_id >= total_experts] + if invalid_ids: + raise ValueError(f"Expert ids {invalid_ids!r} are outside the valid range [0, {total_experts - 1}].") + if num_experts_per_tok > len(expert_ids): + raise ValueError(f"num_experts_per_tok={num_experts_per_tok} cannot exceed {len(expert_ids)} loaded experts.") + return expert_ids + + +def _remap_checkpoint_key(checkpoint_key, expert_index_map): + match = EXPERT_KEY_PATTERN.match(checkpoint_key) + if not match: + return checkpoint_key + + original_expert_idx = int(match.group(2)) + remapped_expert_idx = expert_index_map.get(original_expert_idx) + if remapped_expert_idx is None: + return None + return f"{match.group(1)}{remapped_expert_idx}{match.group(3)}" + + +def _is_routed_gate_weight(checkpoint_key): + return checkpoint_key.endswith(".mlp.gate.weight") + + +def _is_routed_gate_bias(checkpoint_key): + return checkpoint_key.endswith(".mlp.gate.e_score_correction_bias") + + +def _materialize_subset_checkpoint( + model_path: Path, + temp_model_path: Path, + weight_map, + allowed_prefixes, + loaded_expert_ids, +): + expert_index_map = {expert_id: remapped_idx for remapped_idx, expert_id in enumerate(loaded_expert_ids)} + selected_by_shard = defaultdict(list) + + for checkpoint_key, shard_name in weight_map.items(): + if not any(checkpoint_key.startswith(prefix) for prefix in allowed_prefixes): + continue + + remapped_key = _remap_checkpoint_key(checkpoint_key, expert_index_map) + if remapped_key is None: + continue + selected_by_shard[shard_name].append((checkpoint_key, remapped_key)) + + if not selected_by_shard: + raise RuntimeError("No text-only weights were selected from the Kimi K2.5 checkpoint.") + + filtered_weight_map = {} + subset_shards = [] + for shard_idx, (source_shard_name, shard_entries) in enumerate(sorted(selected_by_shard.items())): + tensors = {} + with safe_open(model_path / source_shard_name, framework="pt", device="cpu") as shard_reader: + for checkpoint_key, remapped_key in shard_entries: + tensor = shard_reader.get_tensor(checkpoint_key) + if _is_routed_gate_weight(checkpoint_key): + tensor = tensor[list(loaded_expert_ids), :].contiguous() + elif _is_routed_gate_bias(checkpoint_key): + tensor = tensor[list(loaded_expert_ids)].contiguous() + tensors[remapped_key] = tensor + + subset_shard_name = f"model-subset-{shard_idx:05d}.safetensors" + save_file(tensors, str(temp_model_path / subset_shard_name)) + subset_shards.append(subset_shard_name) + filtered_weight_map.update({remapped_key: subset_shard_name for _, remapped_key in shard_entries}) + + return filtered_weight_map, subset_shards + + +def load_text_only_kimi( + model_path: Path, + num_hidden_layers: int, + loaded_expert_ids=LOADED_EXPERT_IDS, + num_experts_per_tok: int = NUM_EXPERTS_PER_TOKEN, +): + kimi_config = AutoConfig.from_pretrained(str(model_path), trust_remote_code=True) + + # Kimi K2.5 is multimodal, so the text depth must be overridden on text_config. + text_config = copy.deepcopy(kimi_config.text_config) + text_config.num_hidden_layers = num_hidden_layers + loaded_expert_ids = _validate_expert_subset( + loaded_expert_ids, + num_experts_per_tok, + text_config.n_routed_experts, + ) + text_config.n_routed_experts = len(loaded_expert_ids) + text_config.num_experts_per_tok = num_experts_per_tok + text_config.n_group = 1 + text_config.topk_group = 1 + + deepseek_cls = get_class_from_dynamic_module("modeling_deepseek.DeepseekV3ForCausalLM", str(model_path)) + + checkpoint_index = json.loads((model_path / "model.safetensors.index.json").read_text()) + weight_map = checkpoint_index["weight_map"] + + allowed_prefixes = [ + "language_model.model.embed_tokens.", + "language_model.model.norm.", + "language_model.lm_head.", + ] + allowed_prefixes.extend(f"language_model.model.layers.{layer_idx}." for layer_idx in range(num_hidden_layers)) + + with tempfile.TemporaryDirectory() as tmpdir: + temp_model_path = Path(tmpdir) + filtered_weight_map, subset_shards = _materialize_subset_checkpoint( + model_path=model_path, + temp_model_path=temp_model_path, + weight_map=weight_map, + allowed_prefixes=allowed_prefixes, + loaded_expert_ids=loaded_expert_ids, + ) + (temp_model_path / "config.json").write_text(text_config.to_json_string(use_diff=False)) + (temp_model_path / "model.safetensors.index.json").write_text( + json.dumps( + { + "metadata": { + "total_size": sum((temp_model_path / shard_name).stat().st_size for shard_name in subset_shards) + }, + "weight_map": filtered_weight_map, + } + ) + ) + + # We are loading a task checkpoint into the base text model, so disable the + # base/task prefix heuristic and let `key_mapping` strip `language_model.`. + original_base_model_prefix = deepseek_cls.base_model_prefix + deepseek_cls.base_model_prefix = "" + try: + model, loading_info = deepseek_cls.from_pretrained( + str(temp_model_path), + torch_dtype=torch.float32, + config=text_config, + local_files_only=True, + key_mapping={r"^language_model\.": ""}, + output_loading_info=True, + ) + finally: + deepseek_cls.base_model_prefix = original_base_model_prefix + + unexpected_keys = loading_info["unexpected_keys"] + missing_keys = loading_info["missing_keys"] + mismatched_keys = loading_info["mismatched_keys"] + if unexpected_keys or missing_keys or mismatched_keys: + raise RuntimeError( + "Failed to load the text-only Kimi K2.5 checkpoint slice cleanly. " + f"missing={missing_keys}, unexpected={unexpected_keys}, mismatched={mismatched_keys}" + ) + + model.eval() + tokenizer = AutoTokenizer.from_pretrained(str(model_path), trust_remote_code=True) + return model, tokenizer + + +""" +mla_absorption = {"cache_compressed": False, "absorption": False, "online": False} +qaic_configs=[ + None, # Full PKV Cache + {"mla_absorption": mla_absorption}, + {"enable_blocking": True, "blocking_mode": "h"}, # Full PKV Cache with Head Blocking + {"mla_absorption": mla_absorption, "enable_blocking": True, "blocking_mode": "h"}, # Full PKV Cache with Head Blocking +] +""" +TS = 1 +mla_absorption = {"cache_compressed": True, "absorption": False, "online": False} +# mla_absorption = {"cache_compressed": True, "absorption": True, "online": False} +# mla_absorption = {"cache_compressed": True, "absorption": True, "online": True} +qaic_configs = [ + {"mla_absorption": mla_absorption}, # for No Blocking + # {"mla_absorption": mla_absorption, "num_kv_heads_repeat": TS}, # No blocking with kv head replication + # {"mla_absorption": mla_absorption, "enable_blocking": True, "blocking_mode": "kv"}, # for KV blocking + # {"mla_absorption": mla_absorption, "enable_blocking": True, "blocking_mode": "kv", "num_kv_heads_repeat":TS}, # for KV blocking with kv head replication + # {"mla_absorption": mla_absorption, "enable_blocking": True, "blocking_mode": "h","num_kv_heads_repeat": TS}, # for h blocking with kv head replication +] + + +def _set_deterministic(seed: int): + import random + + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + torch.use_deterministic_algorithms(True) + + +@pytest.mark.parametrize("qaic_config", qaic_configs, ids=lambda x: f"qaic_{x}") +def test_qeff_generation(qaic_config): + prompt = "Once upon a time," + PREFILL_SEQ_LEN = 512 + CTX_LEN = 1024 + + _set_deterministic(1234) + _ensure_torch_fx_import_compatibility() + _patch_dynamic_cache_compat() + #config = _prepare_config(MODEL_PATH) + kimi_cls = get_class_from_dynamic_module("modeling_deepseek.DeepseekV3ForCausalLM", str(MODEL_PATH)) + _patch_kimi_tie_weights_compat(kimi_cls) + _patch_deepseek_init_weights_compat(kimi_cls) + model, tokenizer = load_text_only_kimi( + MODEL_PATH, + NUM_HIDDEN_LAYERS, + loaded_expert_ids=LOADED_EXPERT_IDS, + num_experts_per_tok=NUM_EXPERTS_PER_TOKEN, + ) + + inputs = tokenizer(prompt, return_tensors="pt", padding=True) + padded_len = inputs["input_ids"].shape[1] + num_chunks = -(padded_len // -PREFILL_SEQ_LEN) # ceil divide without float + padded_len = num_chunks * PREFILL_SEQ_LEN # Convert to a multiple of prompt_len + + model.to(torch.float32) + # Recompute reference logits after dtype conversion so we compare + # transformed fp32 path against fp32 reference. + with torch.no_grad(): + out = model(**inputs) + + qeff_model = QEFFAutoModelForCausalLM(model, qaic_config=qaic_config) + qeff_model.transform(ctx_len=CTX_LEN, seq_len=PREFILL_SEQ_LEN, bs=1, num_devices=1, qaic_config=qaic_config) + qeff_model.prefill(enable=True) + + inputs = tokenizer(prompt, return_tensors="np", padding="max_length", max_length=padded_len) + inputs["position_ids"] = np.where(inputs.pop("attention_mask"), np.arange(padded_len), -1) + inputs.pop("token_type_ids", None) + inputs = {k: torch.from_numpy(v) for k, v in inputs.items()} + + pad_shape_k = ( + 1, + model.config.num_attention_heads, + CTX_LEN, + model.config.qk_nope_head_dim + model.config.qk_rope_head_dim, + ) + pad_shape_v = (1, model.config.num_attention_heads, CTX_LEN, model.config.v_head_dim) + + num_heads = model.model.layers[0].self_attn.kv_a_proj_with_mqa.weight.shape[0] // ( + model.config.kv_lora_rank + model.config.qk_rope_head_dim + ) + pad_shape_ckv = (1, num_heads, CTX_LEN, model.config.kv_lora_rank) + pad_shape_k_pe = (1, num_heads, CTX_LEN, model.config.qk_rope_head_dim) + + past_key_values = [] + compressed_kvs = [] + + for i in range(model.config.num_hidden_layers): + past_key = torch.zeros((pad_shape_k), dtype=torch.float32) + past_value = torch.zeros((pad_shape_v), dtype=torch.float32) + pkv = (past_key, past_value) + past_key_values.append(pkv) + + ckv = torch.zeros((pad_shape_ckv), dtype=torch.float32) + k_pe = torch.zeros((pad_shape_k_pe), dtype=torch.float32) + x = (ckv, k_pe) + compressed_kvs.append(x) + + cache_compressed = mla_absorption.get("cache_compressed", False) + if cache_compressed: + inputs["compressed_kvs"] = compressed_kvs + else: + inputs["past_key_values"] = past_key_values + setattr(qeff_model.model.model, "num_q_blocks_ffn", PREFILL_SEQ_LEN // 256) + prefill_qeff_out = qeff_model.model(**inputs) + + x = prefill_qeff_out.logits + y = out.logits[:, -1, :] + assert (x - y).abs().max() < 2e-4 + +''' + # Tokenize once; reuse for both reference and qeff model + raw_inputs = tokenizer(prompt, return_tensors="np", padding=True) + padded_len = raw_inputs["input_ids"].shape[1] + num_chunks = -(padded_len // -PREFILL_SEQ_LEN) # ceil divide without float + padded_len = num_chunks * PREFILL_SEQ_LEN # Convert to a multiple of prompt_len + + raw_inputs = tokenizer(prompt, return_tensors="np", padding="max_length", max_length=padded_len) + raw_inputs["position_ids"] = np.where(raw_inputs.pop("attention_mask"), np.arange(padded_len), -1) + raw_inputs.pop("token_type_ids", None) + + inputs = {k: torch.from_numpy(v).to(model.device) for k, v in raw_inputs.items()} + + config = qeff_model.model.config + + inputs = {k: torch.from_numpy(v) for k, v in raw_inputs.items()} + past_key_values = [] + + pad_shape_k = ( + 1, + model.config.num_attention_heads, + CTX_LEN, + model.config.qk_nope_head_dim + model.config.qk_rope_head_dim, + ) + pad_shape_v = (1, model.config.num_attention_heads, CTX_LEN, model.config.v_head_dim) + + num_heads = model.model.layers[0].self_attn.kv_a_proj_with_mqa.weight.shape[0] // ( + model.config.kv_lora_rank + model.config.qk_rope_head_dim + ) + pad_shape_ckv = (1, num_heads, CTX_LEN, model.config.kv_lora_rank) + pad_shape_k_pe = (1, num_heads, CTX_LEN, model.config.qk_rope_head_dim) + + past_key_values = [] + compressed_kvs = [] + + for i in range(model.config.num_hidden_layers): + past_key = torch.zeros((pad_shape_k), dtype=torch.float32) + past_value = torch.zeros((pad_shape_v), dtype=torch.float32) + pkv = (past_key, past_value) + past_key_values.append(pkv) + + ckv = torch.zeros((pad_shape_ckv), dtype=torch.float32) + k_pe = torch.zeros((pad_shape_k_pe), dtype=torch.float32) + x = (ckv, k_pe) + compressed_kvs.append(x) + cache_compressed = mla_absorption.get("cache_compressed", False) + if cache_compressed: + inputs["compressed_kvs"] = compressed_kvs + else: + inputs["past_key_values"] = past_key_values + + prefill_qpc_path = qeff_model.compile( + prefill_seq_len=PREFILL_SEQ_LEN, + ctx_len=CTX_LEN, + num_cores=16, + mxfp6_matmul=False, + mxint8_kv_cache=False, + num_devices=1, + mos=1, + aic_enable_depth_first=True, + num_speculative_tokens=None, + prefill_only=True, + qaic_config=qaic_config, + ) + + from QEfficient.generation.cloud_infer import QAICInferenceSession + + prefill_session = QAICInferenceSession(prefill_qpc_path) + logits_out_placeholder = np.zeros((1, 1, config.vocab_size), dtype=np.float32) + prefill_session.set_buffers({"logits": logits_out_placeholder}) + if cache_compressed: + inputs.pop("compressed_kvs") + else: + inputs.pop("past_key_values") + inputs = {k: v.detach().numpy() for k, v in inputs.items()} + qpc_out = prefill_session.run(inputs) + assert (torch.from_numpy(qpc_out["logits"]) - prefill_qeff_out.logits).abs().max() < 5e-2 +''' \ No newline at end of file diff --git a/examples/kimi_k2/verify_kimi_k25_language_decoder_parity.py b/examples/kimi_k2/verify_kimi_k25_language_decoder_parity.py deleted file mode 100644 index 70cc0c35e1..0000000000 --- a/examples/kimi_k2/verify_kimi_k25_language_decoder_parity.py +++ /dev/null @@ -1,327 +0,0 @@ -# ----------------------------------------------------------------------------- -# -# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. -# SPDX-License-Identifier: BSD-3-Clause -# -# ----------------------------------------------------------------------------- - -import argparse -import copy -import shutil -import subprocess -import tempfile -from pathlib import Path - -import numpy as np -import onnxruntime as ort -import torch -from export_kimi_k25_vision import ( - LOADED_EXPERT_IDS, - MODEL_PATH, - NUM_EXPERTS_PER_TOKEN, - _ensure_torch_fx_import_compatibility, - _load_layer_subset_model, - _patch_deepseek_init_weights_compat, - _patch_kimi_tie_weights_compat, - _prepare_config, -) -from transformers.dynamic_module_utils import get_class_from_dynamic_module - -from QEfficient import QEFFAutoModelForImageTextToText -from QEfficient.generation.cloud_infer import QAICInferenceSession -from QEfficient.transformers.models.modeling_auto import QEffCausalLMForTextImageToTextModel - - -def parse_args(): - parser = argparse.ArgumentParser( - description="Verify Kimi-K2.5 language decoder parity: HF LM logits vs transformed decoder-wrapper logits." - ) - parser.add_argument("--model-path", type=Path, default=MODEL_PATH) - parser.add_argument("--num-text-layers", type=int, default=2) - parser.add_argument("--num-vision-layers", type=int, default=1) - parser.add_argument("--expert-ids", type=str, default=",".join(str(x) for x in LOADED_EXPERT_IDS)) - parser.add_argument("--num-experts-per-token", type=int, default=NUM_EXPERTS_PER_TOKEN) - parser.add_argument("--prompt", type=str, default="Describe this image.") - parser.add_argument("--batch-size", type=int, default=1) - parser.add_argument("--atol", type=float, default=1e-4) - parser.add_argument("--rtol", type=float, default=1e-4) - parser.add_argument("--skip-onnx", action="store_true", help="Skip ONNX export/runtime parity check.") - parser.add_argument( - "--keep-onnx", - action="store_true", - help="Keep exported ONNX under ./aisand/onnx_export_verify_lang instead of a temporary directory.", - ) - parser.add_argument("--run-qpc", action="store_true", help="Compile and execute QPC, then compare vs references.") - parser.add_argument( - "--qaic-compile-bin", - type=str, - default="/opt/qti-aic/exec/qaic-compile", - help="Path to qaic-compile binary.", - ) - parser.add_argument( - "--keep-qpc", - action="store_true", - help="Keep compiled QPC under --qpc-dir (used only with --run-qpc).", - ) - parser.add_argument( - "--qpc-dir", - type=Path, - default=Path("aisand/qpc_verify_lang"), - help="Output directory for QPC when --keep-qpc is set.", - ) - parser.add_argument("--aic-num-cores", type=int, default=16, help="Number of QAIC cores for qpc compile.") - parser.add_argument("--mos", type=int, default=1, help="-mos value for qpc compile.") - parser.add_argument( - "--no-convert-to-fp16", - action="store_true", - help="Disable -convert-to-fp16 during qpc compile (defaults to enabled).", - ) - parser.add_argument("--qpc-atol", type=float, default=1e-2) - parser.add_argument("--qpc-rtol", type=float, default=1e-2) - return parser.parse_args() - - -def _parse_expert_ids(value: str): - return tuple(int(x.strip()) for x in value.split(",") if x.strip()) - - -def _stats(a: torch.Tensor, b: torch.Tensor, atol: float, rtol: float): - diff = (a - b).abs() - return { - "shape": tuple(a.shape), - "max_abs": float(diff.max().item()), - "mean_abs": float(diff.mean().item()), - "rmse": float(torch.sqrt(torch.mean((a - b) ** 2)).item()), - "allclose": bool(torch.allclose(a, b, atol=atol, rtol=rtol)), - } - - -def _language_model_reference_logits(model, input_ids, attention_mask, position_ids, past_key_values): - inputs_embeds = model.get_input_embeddings()(input_ids) - outputs = model.language_model( - inputs_embeds=inputs_embeds, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - use_cache=True, - ) - return outputs[0].float() - - -def _create_prompt_inputs(tokenizer, prompt: str, batch_size: int): - prompts = [prompt] * batch_size - tokenized = tokenizer(prompts, return_tensors="pt", padding=True) - - input_ids = tokenized["input_ids"].to(torch.long) - attention_mask = tokenized["attention_mask"].to(torch.long) - position_ids = torch.where(attention_mask.bool(), attention_mask.cumsum(-1) - 1, -1) - return input_ids, attention_mask, position_ids - - -def _export_and_run_onnx( - transformed_model, - input_ids: torch.Tensor, - position_ids: torch.Tensor, - past_key_values, - keep_onnx: bool, -): - lang_exporter = QEffCausalLMForTextImageToTextModel(copy.deepcopy(transformed_model)) - lang_inputs = transformed_model.get_dummy_inputs(kv_offload=True)["lang"] - output_names = transformed_model.get_output_names(kv_offload=True)["lang"] - dynamic_axes = transformed_model.get_onnx_dynamic_axes(kv_offload=True)["lang"] - - if keep_onnx: - export_dir = Path("aisand/onnx_export_verify_lang") - export_dir.mkdir(parents=True, exist_ok=True) - else: - export_dir = Path(tempfile.mkdtemp(prefix="kimi_k25_verify_lang_onnx_")) - - onnx_path = lang_exporter.export( - inputs=lang_inputs, - output_names=output_names, - dynamic_axes=dynamic_axes, - export_dir=export_dir, - offload_pt_weights=False, - ) - - session = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"]) - ort_inputs = { - "input_ids": input_ids.detach().cpu().numpy().astype(np.int64), - "position_ids": position_ids.detach().cpu().numpy().astype(np.int64), - } - for layer_idx, (past_key, past_value) in enumerate(past_key_values): - ort_inputs[f"past_key.{layer_idx}"] = past_key.detach().cpu().numpy().astype(np.float32) - ort_inputs[f"past_value.{layer_idx}"] = past_value.detach().cpu().numpy().astype(np.float32) - - ort_logits = torch.from_numpy(session.run(["logits"], ort_inputs)[0]).float() - return ort_logits, Path(onnx_path) - - -def _run_qpc(onnx_path: Path, input_ids: torch.Tensor, position_ids: torch.Tensor, past_key_values, args): - if args.keep_qpc: - qpc_dir = args.qpc_dir - if qpc_dir.exists(): - shutil.rmtree(qpc_dir) - qpc_dir.parent.mkdir(parents=True, exist_ok=True) - else: - qpc_dir = Path(tempfile.mkdtemp(prefix="kimi_k25_verify_lang_qpc_")) - shutil.rmtree(qpc_dir) - - compile_cmd = [ - args.qaic_compile_bin, - f"-m={onnx_path}", - f"-onnx-define-symbol=batch_size,{input_ids.shape[0]}", - f"-onnx-define-symbol=seq_len,{input_ids.shape[1]}", - f"-onnx-define-symbol=ctx_len,{past_key_values[0][0].shape[2]}", - f"-aic-num-cores={args.aic_num_cores}", - f"-mos={args.mos}", - f"-aic-binary-dir={qpc_dir}", - ] - if not args.no_convert_to_fp16: - compile_cmd.insert(1, "-convert-to-fp16") - - compile_result = subprocess.run(compile_cmd, capture_output=True, text=True) - if compile_result.returncode != 0: - raise RuntimeError( - "QPC compile failed.\n" - f"Command: {' '.join(compile_cmd)}\n" - f"stdout:\n{compile_result.stdout}\n" - f"stderr:\n{compile_result.stderr}" - ) - - session = QAICInferenceSession(str(qpc_dir)) - qpc_inputs = { - "input_ids": input_ids.detach().cpu().numpy().astype(np.int64), - "position_ids": position_ids.detach().cpu().numpy().astype(np.int64), - } - for layer_idx, (past_key, past_value) in enumerate(past_key_values): - qpc_inputs[f"past_key.{layer_idx}"] = past_key.detach().cpu().numpy().astype(np.float32) - qpc_inputs[f"past_value.{layer_idx}"] = past_value.detach().cpu().numpy().astype(np.float32) - - qpc_feed = {name: qpc_inputs[name] for name in session.input_names} - qpc_outputs = session.run(qpc_feed) - output_name = "logits" if "logits" in qpc_outputs else session.output_names[0] - qpc_logits = torch.from_numpy(qpc_outputs[output_name]).float() - return qpc_logits, qpc_dir - - -if __name__ == "__main__": - args = parse_args() - expert_ids = _parse_expert_ids(args.expert_ids) - - torch.set_grad_enabled(False) - _ensure_torch_fx_import_compatibility() - - config = _prepare_config(args.model_path) - kimi_cls = get_class_from_dynamic_module("modeling_kimi_k25.KimiK25ForConditionalGeneration", str(args.model_path)) - _patch_kimi_tie_weights_compat(kimi_cls) - _patch_deepseek_init_weights_compat(kimi_cls) - - hf_model, tokenizer, _ = _load_layer_subset_model( - model_path=args.model_path, - kimi_cls=kimi_cls, - config=config, - num_vision_layers=args.num_vision_layers, - num_text_layers=args.num_text_layers, - loaded_expert_ids=expert_ids, - num_experts_per_tok=args.num_experts_per_token, - dtype=torch.float32, - ) - hf_model = hf_model.eval() - - transformed_wrapper = QEFFAutoModelForImageTextToText(copy.deepcopy(hf_model)) - transformed_model = transformed_wrapper.model.eval() - decoder_wrapper = transformed_model.get_qeff_language_decoder().eval() - - input_ids, attention_mask, position_ids = _create_prompt_inputs( - tokenizer, - prompt=args.prompt, - batch_size=args.batch_size, - ) - - past_key_values = transformed_model.language_model.get_dummy_pkv_cache( - transformed_model.config.text_config, - args.batch_size, - input_ids.shape[1], - ) - - reference_logits = ( - _language_model_reference_logits( - transformed_model, - input_ids, - attention_mask, - position_ids, - past_key_values=copy.deepcopy(past_key_values), - ) - .detach() - .cpu() - ) - transformed_decoder_logits = ( - decoder_wrapper( - input_ids=input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - image_embeds=None, - image_idx=None, - past_key_values=copy.deepcopy(past_key_values), - use_cache=True, - )[0] - .detach() - .cpu() - .float() - ) - - lm_vs_wrapper = _stats(reference_logits, transformed_decoder_logits, args.atol, args.rtol) - print( - "Loaded subset:", - f"vision_layers={transformed_model.config.vision_config.vt_num_hidden_layers}", - f"text_layers={transformed_model.config.text_config.num_hidden_layers}", - ) - print("Prompt:", repr(args.prompt)) - print("Input shape:", tuple(input_ids.shape)) - print("Language model vs decoder-wrapper:", lm_vs_wrapper) - - if not lm_vs_wrapper["allclose"]: - raise SystemExit( - "Parity check failed: transformed decoder-wrapper logits do not match transformed language-model logits." - ) - - ort_logits = None - onnx_path = None - if not args.skip_onnx or args.run_qpc: - ort_logits, onnx_path = _export_and_run_onnx( - transformed_model, - input_ids=input_ids, - position_ids=position_ids, - past_key_values=copy.deepcopy(past_key_values), - keep_onnx=args.keep_onnx, - ) - - if ort_logits is not None: - lm_vs_onnx = _stats(reference_logits, ort_logits, args.atol, args.rtol) - wrapper_vs_onnx = _stats(transformed_decoder_logits, ort_logits, args.atol, args.rtol) - print("Language ONNX path:", onnx_path) - print("Language model vs ONNXRuntime:", lm_vs_onnx) - print("Decoder-wrapper vs ONNXRuntime:", wrapper_vs_onnx) - - if not lm_vs_onnx["allclose"] or not wrapper_vs_onnx["allclose"]: - raise SystemExit("Parity check failed: ONNXRuntime logits do not match transformed language references.") - - if args.run_qpc: - qpc_logits, qpc_dir = _run_qpc( - onnx_path, - input_ids=input_ids, - position_ids=position_ids, - past_key_values=copy.deepcopy(past_key_values), - args=args, - ) - lm_vs_qpc = _stats(reference_logits, qpc_logits, args.qpc_atol, args.qpc_rtol) - wrapper_vs_qpc = _stats(transformed_decoder_logits, qpc_logits, args.qpc_atol, args.qpc_rtol) - print("Language QPC dir:", qpc_dir) - print(f"Language model vs QPC (atol={args.qpc_atol}, rtol={args.qpc_rtol}):", lm_vs_qpc) - print(f"Decoder-wrapper vs QPC (atol={args.qpc_atol}, rtol={args.qpc_rtol}):", wrapper_vs_qpc) - - if not lm_vs_qpc["allclose"] or not wrapper_vs_qpc["allclose"]: - raise SystemExit("Parity check failed: QPC logits do not match transformed language references.") - - print("Parity check passed.") diff --git a/examples/kimi_k2/verify_kimi_k25_vl_decoder_parity.py b/examples/kimi_k2/verify_kimi_k25_vl_decoder_parity.py deleted file mode 100644 index 83c4ea9d72..0000000000 --- a/examples/kimi_k2/verify_kimi_k25_vl_decoder_parity.py +++ /dev/null @@ -1,499 +0,0 @@ -# ----------------------------------------------------------------------------- -# -# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. -# SPDX-License-Identifier: BSD-3-Clause -# -# ----------------------------------------------------------------------------- - -import argparse -import copy -import os -import shutil -import subprocess -import sys -import tempfile -from io import BytesIO -from pathlib import Path - -import numpy as np -import onnxruntime as ort -import requests -import torch -from export_kimi_k25_vision import ( - LOADED_EXPERT_IDS, - MODEL_PATH, - NUM_EXPERTS_PER_TOKEN, - _ensure_torch_fx_import_compatibility, - _load_layer_subset_model, - _patch_deepseek_init_weights_compat, - _patch_kimi_tie_weights_compat, - _prepare_config, -) -from huggingface_hub import snapshot_download -from PIL import Image -from transformers.dynamic_module_utils import get_class_from_dynamic_module - -from QEfficient import QEFFAutoModelForImageTextToText -from QEfficient.generation.cloud_infer import QAICInferenceSession -from QEfficient.transformers.models.modeling_auto import QEffCausalLMForTextImageToTextModel - - -def parse_args(): - parser = argparse.ArgumentParser( - description=("Verify Kimi-K2.5 V+L decoder parity on CPU with a 2-layer vision + 2-layer text subset.") - ) - parser.add_argument( - "--model-path", - type=str, - default="moonshotai/Kimi-K2.5", - help=( - "Local model path or HF repo id (defaults to moonshotai/Kimi-K2.5; " - "downloads/resolves using HF_HUB_CACHE when needed)." - ), - ) - parser.add_argument("--num-text-layers", type=int, default=2) - parser.add_argument("--num-vision-layers", type=int, default=2) - parser.add_argument("--expert-ids", type=str, default=",".join(str(x) for x in LOADED_EXPERT_IDS)) - parser.add_argument("--num-experts-per-token", type=int, default=NUM_EXPERTS_PER_TOKEN) - parser.add_argument( - "--image-url", - type=str, - default="https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png", - ) - parser.add_argument("--prompt", type=str, default="Describe this image in one sentence.") - parser.add_argument("--atol", type=float, default=1e-4) - parser.add_argument("--rtol", type=float, default=1e-4) - parser.add_argument("--skip-onnx", action="store_true", help="Skip ONNX export/runtime parity check.") - parser.add_argument( - "--keep-onnx", - action="store_true", - help="Keep exported ONNX under ./aisand/onnx_export_verify_vl instead of a temporary directory.", - ) - parser.add_argument("--run-qpc", action="store_true", help="Compile and execute QPC, then compare vs references.") - parser.add_argument( - "--qaic-compile-bin", - type=str, - default="/opt/qti-aic/exec/qaic-compile", - help="Path to qaic-compile binary.", - ) - parser.add_argument( - "--keep-qpc", - action="store_true", - help="Keep compiled QPC under --qpc-dir (used only with --run-qpc).", - ) - parser.add_argument( - "--qpc-dir", - type=Path, - default=Path("aisand/qpc_verify_vl"), - help="Output directory for QPC when --keep-qpc is set.", - ) - parser.add_argument("--aic-num-cores", type=int, default=16, help="Number of QAIC cores for qpc compile.") - parser.add_argument("--mos", type=int, default=1, help="-mos value for qpc compile.") - parser.add_argument( - "--no-convert-to-fp16", - action="store_true", - help="Disable -convert-to-fp16 during qpc compile (defaults to enabled).", - ) - parser.add_argument("--qpc-atol", type=float, default=1e-2) - parser.add_argument("--qpc-rtol", type=float, default=1e-2) - return parser.parse_args() - - -def _parse_expert_ids(value: str): - return tuple(int(x.strip()) for x in value.split(",") if x.strip()) - - -def _resolve_model_path(model_path_or_id: str) -> Path: - candidate = Path(model_path_or_id) - if candidate.exists(): - return candidate - - cache_dir = os.environ.get("HF_HUB_CACHE") - return Path(snapshot_download(repo_id=model_path_or_id, cache_dir=cache_dir)) - - -def _stats(a: torch.Tensor, b: torch.Tensor, atol: float, rtol: float): - diff = (a - b).abs() - return { - "shape": tuple(a.shape), - "max_abs": float(diff.max().item()), - "mean_abs": float(diff.mean().item()), - "rmse": float(torch.sqrt(torch.mean((a - b) ** 2)).item()), - "allclose": bool(torch.allclose(a, b, atol=atol, rtol=rtol)), - } - - -def _hf_image_embeds(model, pixel_values: torch.Tensor, grid_thws: torch.Tensor) -> torch.Tensor: - target_dtype = model.vision_tower.patch_embed.proj.weight.dtype - pixel_values = pixel_values.to(target_dtype) - image_features = model.vision_tower(pixel_values, grid_thws) - if model.mm_projector: - image_features = model.mm_projector(image_features) - return image_features[0] if isinstance(image_features, list) else image_features - - -def _get_onnx_export_dir(keep_onnx: bool): - if keep_onnx: - base_dir = Path("aisand") - base_dir.mkdir(parents=True, exist_ok=True) - for old in base_dir.glob("onnx_export_verify_vl-*"): - shutil.rmtree(old, ignore_errors=True) - export_dir = base_dir / "onnx_export_verify_vl" - if export_dir.exists(): - shutil.rmtree(export_dir, ignore_errors=True) - export_dir.mkdir(parents=True, exist_ok=True) - return export_dir - return Path(tempfile.mkdtemp(prefix="kimi_k25_verify_vl_onnx_")) - - -def _export_and_run_lang_onnx( - transformed_model, - input_ids: torch.Tensor, - position_ids: torch.Tensor, - past_key_values, - keep_onnx: bool, -): - lang_exporter = QEffCausalLMForTextImageToTextModel(copy.deepcopy(transformed_model)) - lang_inputs = transformed_model.get_dummy_inputs(kv_offload=True)["lang"] - output_names = transformed_model.get_output_names(kv_offload=True)["lang"] - dynamic_axes = transformed_model.get_onnx_dynamic_axes(kv_offload=True)["lang"] - - export_dir = _get_onnx_export_dir(keep_onnx) - onnx_path = lang_exporter.export( - inputs=lang_inputs, - output_names=output_names, - dynamic_axes=dynamic_axes, - export_dir=export_dir, - offload_pt_weights=False, - ) - - session = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"]) - ort_inputs = { - "input_ids": input_ids.detach().cpu().numpy().astype(np.int64), - "position_ids": position_ids.detach().cpu().numpy().astype(np.int64), - } - for layer_idx, (past_key, past_value) in enumerate(past_key_values): - ort_inputs[f"past_key.{layer_idx}"] = past_key.detach().cpu().numpy().astype(np.float32) - ort_inputs[f"past_value.{layer_idx}"] = past_value.detach().cpu().numpy().astype(np.float32) - - ort_logits = torch.from_numpy(session.run(["logits"], ort_inputs)[0]).float() - return ort_logits, Path(onnx_path) - - -def _prepare_qpc_dir(args, suffix: str): - if args.keep_qpc: - qpc_dir = args.qpc_dir / suffix - if qpc_dir.exists(): - shutil.rmtree(qpc_dir) - qpc_dir.parent.mkdir(parents=True, exist_ok=True) - else: - qpc_dir = Path(tempfile.mkdtemp(prefix=f"kimi_k25_verify_vl_{suffix}_qpc_")) - shutil.rmtree(qpc_dir) - return qpc_dir - - -def _run_vision_qpc(onnx_path: Path, pixel_values: torch.Tensor, h_shape: torch.Tensor, w_shape: torch.Tensor, args): - qpc_dir = _prepare_qpc_dir(args, "vision") - compile_cmd = [ - args.qaic_compile_bin, - f"-m={onnx_path}", - f"-onnx-define-symbol=num_patches,{pixel_values.shape[0]}", - f"-onnx-define-symbol=h,{h_shape.shape[0]}", - f"-onnx-define-symbol=w,{w_shape.shape[0]}", - f"-aic-num-cores={args.aic_num_cores}", - f"-mos={args.mos}", - f"-aic-binary-dir={qpc_dir}", - ] - if not args.no_convert_to_fp16: - compile_cmd.insert(1, "-convert-to-fp16") - - compile_result = subprocess.run(compile_cmd, capture_output=True, text=True) - if compile_result.returncode != 0: - raise RuntimeError( - "Vision QPC compile failed.\n" - f"Command: {' '.join(compile_cmd)}\n" - f"stdout:\n{compile_result.stdout}\n" - f"stderr:\n{compile_result.stderr}" - ) - - session = QAICInferenceSession(str(qpc_dir)) - qpc_inputs = {} - for name in session.input_names: - if name == "pixel_values": - qpc_inputs[name] = pixel_values.detach().cpu().numpy().astype(np.float32) - elif name == "h_shape": - qpc_inputs[name] = h_shape.detach().cpu().numpy().astype(np.int64) - elif name == "w_shape": - qpc_inputs[name] = w_shape.detach().cpu().numpy().astype(np.int64) - else: - raise RuntimeError(f"Unexpected vision QPC input name: {name}") - - qpc_outputs = session.run(qpc_inputs) - output_name = session.output_names[0] - qpc_out = qpc_outputs[output_name] - return torch.from_numpy(qpc_out).float(), qpc_dir - - -def _run_lang_qpc(onnx_path: Path, input_ids: torch.Tensor, position_ids: torch.Tensor, past_key_values, args): - qpc_dir = _prepare_qpc_dir(args, "lang") - compile_cmd = [ - args.qaic_compile_bin, - f"-m={onnx_path}", - f"-onnx-define-symbol=batch_size,{input_ids.shape[0]}", - f"-onnx-define-symbol=seq_len,{input_ids.shape[1]}", - f"-onnx-define-symbol=ctx_len,{past_key_values[0][0].shape[2]}", - f"-aic-num-cores={args.aic_num_cores}", - f"-mos={args.mos}", - f"-aic-binary-dir={qpc_dir}", - ] - if not args.no_convert_to_fp16: - compile_cmd.insert(1, "-convert-to-fp16") - - compile_result = subprocess.run(compile_cmd, capture_output=True, text=True) - if compile_result.returncode != 0: - raise RuntimeError( - "QPC compile failed.\n" - f"Command: {' '.join(compile_cmd)}\n" - f"stdout:\n{compile_result.stdout}\n" - f"stderr:\n{compile_result.stderr}" - ) - - session = QAICInferenceSession(str(qpc_dir)) - qpc_inputs = { - "input_ids": input_ids.detach().cpu().numpy().astype(np.int64), - "position_ids": position_ids.detach().cpu().numpy().astype(np.int64), - } - for layer_idx, (past_key, past_value) in enumerate(past_key_values): - qpc_inputs[f"past_key.{layer_idx}"] = past_key.detach().cpu().numpy().astype(np.float32) - qpc_inputs[f"past_value.{layer_idx}"] = past_value.detach().cpu().numpy().astype(np.float32) - - qpc_feed = {name: qpc_inputs[name] for name in session.input_names} - qpc_outputs = session.run(qpc_feed) - output_name = "logits" if "logits" in qpc_outputs else session.output_names[0] - qpc_logits = torch.from_numpy(qpc_outputs[output_name]).float() - return qpc_logits, qpc_dir - - -if __name__ == "__main__": - args = parse_args() - if args.model_path is None: - args.model_path = MODEL_PATH - args.model_path = _resolve_model_path(str(args.model_path)) - expert_ids = _parse_expert_ids(args.expert_ids) - - torch.set_grad_enabled(False) - _ensure_torch_fx_import_compatibility() - - config = _prepare_config(args.model_path) - kimi_cls = get_class_from_dynamic_module("modeling_kimi_k25.KimiK25ForConditionalGeneration", str(args.model_path)) - _patch_kimi_tie_weights_compat(kimi_cls) - _patch_deepseek_init_weights_compat(kimi_cls) - - hf_model, tokenizer, processor = _load_layer_subset_model( - model_path=args.model_path, - kimi_cls=kimi_cls, - config=config, - num_vision_layers=args.num_vision_layers, - num_text_layers=args.num_text_layers, - loaded_expert_ids=expert_ids, - num_experts_per_tok=args.num_experts_per_token, - dtype=torch.float32, - ) - # Required for HW compatibility and stable vision parity in this flow. - hf_model.vision_tower.patch_embed.pos_emb.interpolation_mode = "bilinear" - hf_model = hf_model.eval() - - transformed_wrapper = QEFFAutoModelForImageTextToText(copy.deepcopy(hf_model)) - transformed_model = transformed_wrapper.model.eval() - decoder_wrapper = transformed_model.get_qeff_language_decoder().eval() - - image = Image.open(BytesIO(requests.get(args.image_url, timeout=30).content)).convert("RGB") - messages = [ - { - "role": "user", - "content": [ - {"type": "image_url", "image_url": image}, - {"type": "text", "text": args.prompt}, - ], - } - ] - inputs = processor(messages=messages, add_generation_prompt=True, tokenize=False, return_tensors="pt") - - input_ids = inputs["input_ids"].to(torch.long) - attention_mask = inputs["attention_mask"].to(torch.long) - pixel_values = inputs["pixel_values"].to(torch.float32) - grid_thws = inputs["grid_thws"].to(torch.long) - h_shape = torch.ones(int(grid_thws[0, 1].item()), dtype=torch.int64) - w_shape = torch.ones(int(grid_thws[0, 2].item()), dtype=torch.int64) - position_ids = torch.where(attention_mask.bool(), attention_mask.cumsum(-1) - 1, -1) - - image_features = transformed_model._extract_image_features(pixel_values, grid_thws) - if transformed_model.mm_projector: - image_features = transformed_model.mm_projector(image_features) - image_embeds = image_features[0] if isinstance(image_features, list) else image_features - vision_reference = _hf_image_embeds(hf_model, pixel_values, grid_thws).detach().cpu().float() - - vision_onnx_path = None - if not args.skip_onnx or args.run_qpc: - vision_cmd = [ - str(Path(sys.executable)), - str(Path(__file__).with_name("verify_kimi_k25_vision_parity.py")), - "--model-path", - str(args.model_path), - "--num-vision-layers", - str(args.num_vision_layers), - "--num-text-layers", - str(args.num_text_layers), - "--expert-ids", - args.expert_ids, - "--num-experts-per-token", - str(args.num_experts_per_token), - "--image-url", - args.image_url, - "--prompt", - args.prompt, - "--atol", - str(args.atol), - "--rtol", - str(args.rtol), - ] - keep_vision_onnx = args.keep_onnx or args.run_qpc - if keep_vision_onnx: - vision_cmd.append("--keep-onnx") - subprocess.run(vision_cmd, check=True) - - if keep_vision_onnx: - vision_matches = sorted(Path("aisand").glob("onnx_export_verify-*/KimiK25EncoderWrapper.onnx")) - if not vision_matches: - raise RuntimeError("Vision ONNX not found under aisand/onnx_export_verify-*/KimiK25EncoderWrapper.onnx") - vision_onnx_path = vision_matches[-1] - - merged_inputs_embeds, merged_attention_mask, _, merged_position_ids = ( - transformed_model._merge_input_ids_with_image_features( - image_features=image_features, - inputs_embeds=transformed_model.get_input_embeddings()(input_ids), - input_ids=input_ids, - attention_mask=attention_mask, - labels=None, - ) - ) - - past_key_values = transformed_model.language_model.get_dummy_pkv_cache( - transformed_model.config.text_config, - input_ids.shape[0], - merged_inputs_embeds.shape[1], - ) - - model_logits = ( - transformed_model.language_model( - inputs_embeds=merged_inputs_embeds, - attention_mask=merged_attention_mask, - position_ids=merged_position_ids, - past_key_values=copy.deepcopy(past_key_values), - use_cache=True, - )[0] - .detach() - .cpu() - .float() - ) - - decoder_logits = ( - decoder_wrapper( - input_ids=input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - image_embeds=image_embeds, - image_idx=torch.zeros((input_ids.shape[0], 1), dtype=torch.int64), - past_key_values=copy.deepcopy(past_key_values), - use_cache=True, - )[0] - .detach() - .cpu() - .float() - ) - - decoder_vs_model = _stats(model_logits, decoder_logits, args.atol, args.rtol) - - print( - "Loaded subset:", - f"vision_layers={transformed_model.config.vision_config.vt_num_hidden_layers}", - f"text_layers={transformed_model.config.text_config.num_hidden_layers}", - ) - print("Tokenizer vocab size:", tokenizer.vocab_size) - print("Prompt:", repr(args.prompt)) - print("Input shape:", tuple(input_ids.shape)) - print("Model logits vs decoder-wrapper logits:", decoder_vs_model) - - if not decoder_vs_model["allclose"]: - raise SystemExit( - "Parity check failed: decoder-wrapper logits do not match transformed model logits for image+text input." - ) - - # ONNX/QPC parity runs the language decoder in decode mode (seq_len=1) - # with multimodal context represented in past_key_values. - onnx_input_ids = input_ids[:, -1:] - onnx_position_ids = merged_position_ids[:, -1:] - onnx_reference_logits = ( - transformed_model.language_model( - inputs_embeds=transformed_model.get_input_embeddings()(onnx_input_ids), - position_ids=onnx_position_ids, - past_key_values=copy.deepcopy(past_key_values), - use_cache=True, - )[0] - .detach() - .cpu() - .float() - ) - - lang_ort = None - lang_onnx_path = None - if not args.skip_onnx or args.run_qpc: - lang_ort, lang_onnx_path = _export_and_run_lang_onnx( - transformed_model, - input_ids=onnx_input_ids, - position_ids=onnx_position_ids, - past_key_values=copy.deepcopy(past_key_values), - keep_onnx=args.keep_onnx, - ) - - if vision_onnx_path is not None: - print("Vision ONNX path:", vision_onnx_path) - - if lang_ort is not None: - lang_ref_vs_onnx = _stats(onnx_reference_logits, lang_ort, args.atol, args.rtol) - print("Language ONNX path:", lang_onnx_path) - print("Decode-step reference vs ONNXRuntime:", lang_ref_vs_onnx) - - if not lang_ref_vs_onnx["allclose"]: - raise SystemExit("Parity check failed: language ONNXRuntime logits do not match decode-step reference.") - - if args.run_qpc: - vision_qpc, vision_qpc_dir = _run_vision_qpc( - vision_onnx_path, - pixel_values=pixel_values, - h_shape=h_shape, - w_shape=w_shape, - args=args, - ) - vision_ref_vs_qpc = _stats(vision_reference, vision_qpc, args.qpc_atol, args.qpc_rtol) - print("Vision QPC dir:", vision_qpc_dir) - print(f"Vision reference vs QPC (atol={args.qpc_atol}, rtol={args.qpc_rtol}):", vision_ref_vs_qpc) - - if not vision_ref_vs_qpc["allclose"]: - raise SystemExit("Parity check failed: vision QPC output does not match vision reference.") - - lang_qpc, lang_qpc_dir = _run_lang_qpc( - lang_onnx_path, - input_ids=onnx_input_ids, - position_ids=onnx_position_ids, - past_key_values=copy.deepcopy(past_key_values), - args=args, - ) - lang_ref_vs_qpc = _stats(onnx_reference_logits, lang_qpc, args.qpc_atol, args.qpc_rtol) - print("Language QPC dir:", lang_qpc_dir) - print(f"Decode-step reference vs QPC (atol={args.qpc_atol}, rtol={args.qpc_rtol}):", lang_ref_vs_qpc) - - if not lang_ref_vs_qpc["allclose"]: - raise SystemExit("Parity check failed: language QPC logits do not match decode-step reference.") - - print("Parity check passed.") From 585d4d96d4aa0b929afa49e5b716b10eb2428316 Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Tue, 30 Jun 2026 09:17:31 +0530 Subject: [PATCH 10/33] fix merge input_embeds and position ids Signed-off-by: Mamta Singh --- .../models/kimi_k25/configuration_kimi_k25.py | 5 +- .../models/kimi_k25/modeling_kimi_k25.py | 746 ++++-------------- .../transformers/models/modeling_auto.py | 17 +- .../transformers/models/pytorch_transforms.py | 1 + examples/kimi_k2/export_kimi_k25_vision.py | 15 +- examples/kimi_k2/test_kimi_k25.py | 631 ++++++++++----- examples/kimi_k2/test_kimi_k25_text.py | 430 ---------- .../kimi_k2/verify_kimi_k25_vision_parity.py | 271 ------- 8 files changed, 604 insertions(+), 1512 deletions(-) delete mode 100644 examples/kimi_k2/test_kimi_k25_text.py delete mode 100644 examples/kimi_k2/verify_kimi_k25_vision_parity.py diff --git a/QEfficient/transformers/models/kimi_k25/configuration_kimi_k25.py b/QEfficient/transformers/models/kimi_k25/configuration_kimi_k25.py index dde386b53a..2f6d3b2467 100644 --- a/QEfficient/transformers/models/kimi_k25/configuration_kimi_k25.py +++ b/QEfficient/transformers/models/kimi_k25/configuration_kimi_k25.py @@ -7,10 +7,7 @@ from transformers.configuration_utils import PretrainedConfig -try: - from QEfficient.transformers.models.deepseek_v3.configuration_deepseek import DeepseekV3Config -except ImportError: - from QEfficient.transformers.models.deepseek_v3.configuration_deepseek import DeepseekV3Config +from QEfficient.transformers.models.deepseek_v3.configuration_deepseek import DeepseekV3Config class KimiK25VisionConfig(PretrainedConfig): diff --git a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py index 7583a94a11..f2df4e272c 100644 --- a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py +++ b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py @@ -7,15 +7,10 @@ import math import sys as _sys -from collections.abc import Sequence -from copy import deepcopy from io import BytesIO from pathlib import Path - -# from QEfficient import QEFFAutoModelForImageTextToText from typing import List, Optional, Tuple, Type, Union -import numpy as np import requests import torch import torch.nn as nn @@ -23,8 +18,6 @@ from PIL import Image from transformers import AutoProcessor, activations -from QEfficient.transformers.models.deepseek_v3.modeling_deepseek import QEffDeepseekV3ForCausalLM - try: from transformers.activations import PytorchGELUTanh except ImportError: @@ -34,14 +27,10 @@ PytorchGELUTanh = GELUTanh from transformers.activations import PytorchGELUTanh from transformers.cache_utils import Cache -from transformers.configuration_utils import PretrainedConfig -from transformers.modeling_utils import PreTrainedModel from transformers.models.llava.modeling_llava import LlavaCausalLMOutputWithPast from QEfficient.utils import constants -from .configuration_kimi_k25 import KimiK25Config - def eager_attention_forward(q, k, v, **kwargs): q = q.transpose(0, 1) # (num_heads, seq_len, head_dim) @@ -111,135 +100,10 @@ def apply_rope(xq, xk, freqs_cis): return xq_out.type_as(xq), xk_out.type_as(xk) -def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): - """ - From: - https://github.com/OpenGVLab/InternVideo/blob/421f6d2361fc8f61a3394244571f2601a4e99e29/InternVideo2/multi_modality/models/backbones/internvideo2/pos_embed.py#L86 - embed_dim: output dimension for each position - pos: a list of positions to be encoded: size (M,) - out: (M, D) - """ - assert embed_dim % 2 == 0 - omega = np.arange(embed_dim // 2, dtype=np.float32) - omega /= embed_dim / 2.0 - omega = 1.0 / 10000**omega # (D/2,) - - pos = pos.reshape(-1) # (M,) - out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product - - emb_sin = np.sin(out) # (M, D/2) - emb_cos = np.cos(out) # (M, D/2) - - emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D) - return emb - - -def get_1d_sincos_pos_embed(embed_dim, t_size, cls_token=False): - """ - t_size: int of the temporal size - return: - pos_embed: [t_size, embed_dim] or [1+t_size, embed_dim] (w/ or w/o cls_token) - """ - grid_t = np.arange(t_size, dtype=np.float32) - pos_embed = get_1d_sincos_pos_embed_from_grid(embed_dim, grid_t) - if cls_token: - pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0) - return pos_embed - - class QEffLearnable2DInterpPosEmbDivided_fixed(nn.Module): def __qeff_init__(self): self.interpolation_mode = "bilinear" - """def __qeff_init__(self, - height: int, - width: int, - num_frames: int, - dim: int, - interpolation_mode: str = 'bilinear') -> None: - super().__init__() - self.height = height - self.width = width - self.num_frames = num_frames - self.dim = dim - self.interpolation_mode = interpolation_mode - self.weight = nn.Parameter(torch.empty(height, width, dim)) - self.register_buffer('time_weight', - torch.from_numpy( - get_1d_sincos_pos_embed( - self.dim, - self.num_frames)).float().unsqueeze(1), - persistent=False) - - self.reset_parameters() - """ - - def reset_parameters(self): - nn.init.normal_(self.weight) - - def forward(self, x: torch.Tensor, grid_thws: torch.Tensor) -> torch.Tensor: - pos_embs = [] - for t, h, w in grid_thws.tolist(): - assert t <= self.num_frames, f"t:{t} > self.num_frames:{self.num_frames}" - if (h, w) == self.weight.shape[:-1]: - pos_emb_2d = self.weight.flatten(end_dim=1) - else: - pos_emb_2d = get_rope_shape( - self.weight, - interpolation_mode=self.interpolation_mode, - shape=(h, w), - ) - if t == 1: - pos_emb_3d = pos_emb_2d - else: - pos_emb_3d = pos_emb_2d.unsqueeze(0).repeat(t, 1, 1) + self.time_weight[0:t] - - pos_embs.append(pos_emb_3d.reshape(-1, pos_emb_3d.shape[-1])) - - out = x + torch.cat(pos_embs) - return out - - -class MoonVision3dPatchEmbed(nn.Module): - def __init__( - self, - out_dim: int, - in_dim: int = 3, - patch_size: int | tuple[int, int] = (14, 14), - pos_emb_height: int = 14, - pos_emb_width: int = 14, - pos_emb_time: int = 4, - pos_emb_type: str = "divided_fixed", - ): - super().__init__() - assert isinstance(patch_size, int | Sequence), f"Invalid patch_size type: {type(patch_size)}" - if isinstance(patch_size, int): - patch_size = (patch_size, patch_size) - assert len(patch_size) == 2, f"Expected patch_size to be a tuple of 2, got {patch_size}" - self.patch_size = patch_size - - self.proj = nn.Conv2d(in_dim, out_dim, kernel_size=patch_size, stride=patch_size) - - if pos_emb_type == "divided_fixed": - self.pos_emb = QEffLearnable2DInterpPosEmbDivided_fixed( - height=pos_emb_height, width=pos_emb_width, num_frames=pos_emb_time, dim=out_dim - ) - else: - raise NotImplementedError(f"Not support pos_emb_type: {pos_emb_type}") - - def forward(self, x: torch.Tensor, grid_thws: torch.Tensor) -> torch.Tensor: - """ - Args: - x (L, Channels): input tensor - grid_hws (N, 3): temporal, height and width - Returns: - (L, Cout) tensor - """ - x = self.proj(x).view(x.size(0), -1) - # apply positional embedding - x = self.pos_emb(x, grid_thws) - return x - class Rope2DPosEmbRepeated(nn.Module): """2D rotary position embedding with multi-resolution support. @@ -471,236 +335,6 @@ def __qeff_init__(self): new_blocks.append(new_block.to(device=old_block.wqkv.weight.device, dtype=old_block.wqkv.weight.dtype)) self.blocks = nn.ModuleList(new_blocks) - def forward( - self, - hidden_states: torch.Tensor, - grid_thws: torch.Tensor, - ) -> torch.Tensor: - # if not hasattr(self.rope_2d, 'freqs_cos'): - # self.rope_2d.register_buffer('freqs_cos', self.rope_2d.freqs_cis.real.contiguous(), persistent=False) - # self.rope_2d.register_buffer('freqs_sin', self.rope_2d.freqs_cis.imag.contiguous(), persistent=False) - rope_freqs_cis = self.rope_2d.get_freqs_cis(grid_thws=grid_thws, device=hidden_states.device) - - lengths = torch.cat( - ( - torch.zeros(1, dtype=grid_thws.dtype, device=grid_thws.device), - grid_thws[:, 0] * grid_thws[:, 1] * grid_thws[:, 2], - ) - ) - - max_seqlen = lengths.max() - cu_seqlens = lengths.to(hidden_states.device).cumsum(dim=0, dtype=torch.int32) - for block in self.blocks: - hidden_states = block(hidden_states, cu_seqlens, max_seqlen, rope_freqs_cis=rope_freqs_cis) - - hidden_states = self.final_layernorm(hidden_states) - return hidden_states - - -def tpool_patch_merger( - x: torch.Tensor, - grid_thws: torch.Tensor, - merge_kernel_size: tuple[int, int] = (2, 2), -) -> list[torch.Tensor]: - d_model = x.size(-1) - - outputs = [] - pre_sum = 0 - for t, h, w in grid_thws.tolist(): - # Get the current sequence - seq = x[pre_sum : pre_sum + t * h * w] - # Reshape along self.merge_kernel_size and concat to the last dimension - kernel_height, kernel_width = merge_kernel_size - new_height, new_width = h // kernel_height, w // kernel_width - reshaped_seq = seq.view(t, new_height, kernel_height, new_width, kernel_width, d_model) - reshaped_seq = reshaped_seq.permute(0, 1, 3, 2, 4, 5).contiguous().mean(dim=0) # temporal pooling - padded_seq = reshaped_seq.view(new_height * new_width, kernel_height * kernel_width, -1) - outputs.append(padded_seq) - pre_sum += t * h * w - - return outputs - - -class MoonViT3dPretrainedModel(PreTrainedModel): - config_class = None - model_type = "moonvit3d" - _no_split_modules = ["PackingTransformer"] - _supports_flash_attn_2 = True - _supports_sdpa = True - - def __init__(self, config, *inputs, **kwargs): - super().__init__(config, *inputs, **kwargs) - config = deepcopy(config) - self.merge_kernel_size = config.merge_kernel_size - self.patch_size = config.patch_size - self.merge_type = config.merge_type - - self.patch_embed = MoonVision3dPatchEmbed( - out_dim=config.hidden_size, - patch_size=config.patch_size, - pos_emb_height=config.init_pos_emb_height, - pos_emb_width=config.init_pos_emb_width, - pos_emb_time=config.init_pos_emb_time, - pos_emb_type=config.pos_emb_type, - ) - - self.encoder = QEffMoonViT3dEncoder( - hidden_dim=config.hidden_size, - num_layers=config.num_hidden_layers, - block_cfg={ - "num_heads": config.num_attention_heads, - "hidden_dim": config.hidden_size, - "mlp_dim": config.intermediate_size, - "activation": PytorchGELUTanh(), - "attn_bias": True, - "attn_implementation": config._attn_implementation, - }, - video_attn_type=config.video_attn_type, - ) - - def forward(self, pixel_values: torch.Tensor, grid_thws: torch.Tensor) -> torch.Tensor: - """ - Args: - pixel_values (torch.Tensor): The input pixel values. - grid_thws (torch.Tensor): Temporal, height and width. - Returns: - torch.Tensor: The output tokens. - """ - # grid_thws = grid_thws.to('cpu') - assert grid_thws.ndim == 2, f"grid_thws should be 2D, got {grid_thws.ndim}" - assert grid_thws.size(1) == 3, f"No support for thw: {grid_thws}" - hidden_states = self.patch_embed(pixel_values, grid_thws) - hidden_states = self.encoder(hidden_states, grid_thws) - if self.merge_type == "sd2_tpool": # spatial downsampling 2x with temporal pooling all - hidden_states = tpool_patch_merger(hidden_states, grid_thws, merge_kernel_size=self.merge_kernel_size) - else: - raise NotImplementedError(f"Not support {self.merge_type}") - - return hidden_states - - -# ============================================================================ -# MM Projector Helper Classes (from mm_projector/modeling_mm_projectors.py) -# ============================================================================ - - -class IdentityMap(nn.Module): - def __init__(self): - super().__init__() - - def forward(self, x, *args, **kwargs): - return x - - -class MLP(nn.Module): - def __init__(self, config): - super().__init__() - # TODO, use faster LayerNorm - self.pre_norm = nn.LayerNorm(config.mm_hidden_size) - self.proj = nn.Sequential( - nn.Linear(config.mm_hidden_size, config.hidden_size), - nn.GELU(), - nn.Linear(config.hidden_size, config.hidden_size), - ) - - def forward(self, x, *args, **kwargs): - assert isinstance(x, list | tuple), f"x is not a list or tuple: {type(x)}" - lengths = [item.shape[0] for item in x] - x = torch.cat(x, dim=0) - x = self.pre_norm(x) - x = self.proj(x) - x = torch.split(x, lengths, dim=0) - - return x - - -class PatchMergerMLP(nn.Module): - def __init__(self, config): - super().__init__() - eps = config.projector_ln_eps - self.hidden_size = config.mm_hidden_size * (config.merge_kernel_size[0] * config.merge_kernel_size[1]) - self.pre_norm = nn.LayerNorm(config.mm_hidden_size, eps=eps) - self.proj = nn.Sequential( - nn.Linear(self.hidden_size, self.hidden_size), - nn.GELU(), - nn.Linear(self.hidden_size, config.hidden_size), - ) - - def forward(self, x, *args, **kwargs): - if isinstance(x, list) or isinstance(x, tuple): - x = [self.proj(self.pre_norm(item).view(item.shape[0], -1)) for item in x] - else: - # B, N, N_k, C = x.shape - B = x.shape[0] - x = self.proj(self.pre_norm(x).view(B, -1, self.hidden_size)) - return x - - -class KimiK25PreTrainedModel(PreTrainedModel): - config_class = KimiK25Config - base_model_prefix = "model" - _no_split_modules = [ - "MoonViT3dPretrainedModel", - "MoonViTEncoderLayer", - "DeepseekDecoderLayer", - "PatchMergerMLP", - ] - _skip_keys_device_placement = "past_key_values" - _supports_flash_attn_2 = True - _supports_sdpa = False - - def _init_weights(self, module): - # important: this ported version of Llava isn't meant for training from scratch - only - # inference and fine-tuning - so the proper init weights code has been removed - the original codebase - # https://github.com/haotian-liu/LLaVA/tree/main/llava should serve for that purpose - std = ( - self.config.initializer_range - if hasattr(self.config, "initializer_range") - else self.config.text_config.initializer_range - ) - - if hasattr(module, "class_embedding"): - module.class_embedding.data.normal_(mean=0.0, std=std) - - if isinstance(module, (nn.Linear, nn.Conv2d)): - module.weight.data.normal_(mean=0.0, std=std) - if module.bias is not None: - module.bias.data.zero_() - elif isinstance(module, nn.Embedding): - module.weight.data.normal_(mean=0.0, std=std) - if module.padding_idx is not None: - module.weight.data[module.padding_idx].zero_() - - -class VisionTowerConfig(PretrainedConfig): - model_type = "moonvit3d" - - def __init__(self, config: KimiK25Config, **kwargs): - super().__init__(**kwargs) - self.patch_size = config.patch_size - self.init_pos_emb_height = config.init_pos_emb_height - self.init_pos_emb_width = config.init_pos_emb_width - self.init_pos_emb_time = config.init_pos_emb_time - self.pos_emb_type = config.pos_emb_type - self.num_attention_heads = config.vt_num_attention_heads - self.num_hidden_layers = config.vt_num_hidden_layers - self.hidden_size = config.vt_hidden_size - self.intermediate_size = config.vt_intermediate_size - self.merge_kernel_size = config.merge_kernel_size - self.video_attn_type = config.video_attn_type - self.merge_type = config.merge_type - self._attn_implementation = config._attn_implementation - - -class ProjectorConfig: - def __init__(self, config: KimiK25Config): - self.mm_projector_type = config.mm_projector_type - self.mm_hidden_size = config.mm_hidden_size - self.hidden_size = config.text_hidden_size - self.merge_kernel_size = config.merge_kernel_size - self.projector_hidden_act = config.projector_hidden_act - self.projector_ln_eps = config.projector_ln_eps - class QEffKimiK25EncoderWrapper(nn.Module): def __init__(self, model): @@ -718,22 +352,6 @@ def get_submodules_for_export(self) -> Type[nn.Module]: return {self.model.vision_model.model.layers[0].__class__} # return {self.model.layers[0].__class__} - def forward_only_image(self, pixel_values: torch.Tensor, grid_thws: torch.Tensor) -> list[torch.Tensor]: - """ - Run only the vision tower and mm_projector to extract image embeddings. - - Args: - pixel_values: Preprocessed image pixel values. - grid_thws: Grid temporal/height/width info for the images. - - Returns: - image_embeds: List of projected image embedding tensors, one per image. - """ - image_features = self._extract_image_features(pixel_values, grid_thws) - if self.mm_projector: - image_features = self.mm_projector(image_features) - return image_features - def forward(self, pixel_values: torch.Tensor, h_shape: torch.Tensor, w_shape: torch.Tensor) -> torch.Tensor: """ ONNX-exportable forward that runs only the vision tower and mm_projector. @@ -867,27 +485,69 @@ def forward( else: raise ValueError(f"Expected image_embeds rank 2 or 3, got {image_embeds.dim()}.") - if image_embeds_for_state is not None and input_ids is not None and input_ids.shape[1] != 1: + if image_embeds_for_state is not None and input_ids is not None: + if attention_mask is None: + # ONNX prefill/decode callers may omit attention_mask but still + # provide position_ids with `-1` on padded slots. + if position_ids is not None: + attention_mask = (position_ids >= 0).to(dtype=torch.long, device=input_ids.device) + else: + pad_token_id = getattr(self.config, "pad_token_id", None) + if pad_token_id is None: + attention_mask = torch.ones_like(input_ids, dtype=torch.long, device=input_ids.device) + else: + attention_mask = (input_ids != pad_token_id).to(dtype=torch.long, device=input_ids.device) + selected = input_ids == self.config.media_placeholder_token_id - if selected.any(): - if attention_mask is None: - attention_mask = torch.ones_like(input_ids, dtype=torch.long, device=input_ids.device) + selected_any = selected.any(dim=1, keepdim=True) + selected_image_tokens = selected.to(torch.int64).sum(dim=1, keepdim=True) - image_features = [image_embeds_for_state] - inputs_embeds = inputs_embeds.to(image_embeds_for_state.dtype) - inputs_embeds, attention_mask, _, position_ids = self.model._merge_input_ids_with_image_features( + image_features = [image_embeds_for_state] + inputs_embeds = inputs_embeds.to(image_embeds_for_state.dtype) + + merged_inputs_embeds, merged_attention_mask, _, merged_position_ids = ( + self.model._merge_input_ids_with_image_features( image_features=image_features, inputs_embeds=inputs_embeds, input_ids=input_ids, attention_mask=attention_mask, labels=None, ) - image_idx = image_idx + torch.tensor( - [[sum(feature.shape[0] for feature in image_features)]], - dtype=torch.int64, - device=image_idx.device, + ) + + inputs_embeds = merged_inputs_embeds + attention_mask = merged_attention_mask + if position_ids is None: + position_ids = merged_position_ids + else: + # Preserve caller-provided absolute position offset (needed for + # chunked prefill/decode parity) while using merged sequence + # positions for image-expanded tokens. + position_offset = position_ids[:, :1] + image_idx.to( + device=position_ids.device, dtype=position_ids.dtype + ) + position_ids = torch.where( + merged_attention_mask > 0, + merged_position_ids + position_offset, + torch.full_like(merged_position_ids, -1), ) + merged_image_tokens = torch.tensor( + [[sum(feature.shape[0] for feature in image_features)]], + dtype=torch.int64, + device=image_idx.device, + ) + image_position_delta = torch.clamp(merged_image_tokens - selected_image_tokens, min=0) + image_idx = image_idx + selected_any.to(torch.int64) * image_position_delta + + if position_ids is None and attention_mask is not None: + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids = torch.where(attention_mask > 0, position_ids, torch.full_like(position_ids, -1)) + elif position_ids is not None and attention_mask is not None: + # Preserve `-1` on padded slots so cache updates skip them without + # creating dynamic Slice nodes unsupported by compile. + position_ids = torch.where(attention_mask > 0, position_ids, torch.full_like(position_ids, -1)) + outputs = self.model.language_model( inputs_embeds=inputs_embeds, attention_mask=attention_mask, @@ -917,66 +577,14 @@ def forward( # ref https://github.com/huggingface/transformers/blob/78b2929c0554b79e0489b451ce4ece14d265ead2/src/transformers/models/llava/modeling_llava.py#L240 -class QEffKimiK25ForConditionalGeneration(KimiK25PreTrainedModel): - def __init__(self, config: KimiK25Config): - super().__init__(config) - - vt_config = VisionTowerConfig(config.vision_config) - self.vision_tower = MoonViT3dPretrainedModel(vt_config) - - proj_config = ProjectorConfig(config.vision_config) - if proj_config.mm_projector_type == "identity": - self.mm_projector = IdentityMap() - elif proj_config.mm_projector_type == "mlp": - self.mm_projector = MLP(proj_config) - elif proj_config.mm_projector_type == "patchmerger": - self.mm_projector = PatchMergerMLP(proj_config) - else: - raise ValueError(f"Unsupported mm_projector_type: {proj_config.mm_projector_type}") - - self.language_model = QEffDeepseekV3ForCausalLM(config.text_config) - self.post_init() - - if hasattr(self.language_model, "dtype"): - target_dtype = self.language_model.dtype - self.vision_tower = self.vision_tower.to(dtype=target_dtype) - self.mm_projector = self.mm_projector.to(dtype=target_dtype) - +class QEffKimiK25ForConditionalGeneration(nn.Module): def get_qeff_vision_encoder(self): return QEffKimiK25EncoderWrapper(self) def get_qeff_language_decoder(self): return QEffKimiK25DecoderWrapper(self) - def get_input_embeddings(self): - return self.language_model.get_input_embeddings() - - def set_input_embeddings(self, value): - self.language_model.set_input_embeddings(value) - - def get_output_embeddings(self): - return self.language_model.get_output_embeddings() - - def set_output_embeddings(self, new_embeddings): - self.language_model.set_output_embeddings(new_embeddings) - - def set_decoder(self, decoder): - self.language_model.set_decoder(decoder) - - def get_decoder(self): - return self.language_model.get_decoder() - - def tie_weights(self): - return self.language_model.tie_weights() - - def resize_token_embeddings(self, new_num_tokens: int | None = None, pad_to_multiple_of=None) -> nn.Embedding: - model_embeds = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of) - # update vocab size - self.config.text_config.vocab_size = model_embeds.num_embeddings - self.vocab_size = model_embeds.num_embeddings - return model_embeds - - def _merge_input_ids_with_image_features( + def _qeff_merge_input_ids_with_image_features( self, image_features: list[torch.Tensor], inputs_embeds: torch.Tensor, @@ -997,118 +605,90 @@ def _merge_input_ids_with_image_features( labels (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`, *optional*): The labels. """ + if len(image_features) == 0: + raise ValueError("At least one image_features tensor is required.") + _, embed_dim = image_features[0].shape feature_lengths = [x.shape[0] for x in image_features] - image_features = torch.cat(image_features, dim=0) + image_features_cat = torch.cat(image_features, dim=0) image_token_index: int = self.config.media_placeholder_token_id pad_token_id: int = self.config.pad_token_id - ignore_index: int = self.config.ignore_index - batch_size, sequence_length = input_ids.shape - left_padding = not torch.sum(input_ids[:, -1] == torch.tensor(pad_token_id)) + batch_size, _ = input_ids.shape + target_device = inputs_embeds.device - # 1. Create a mask to know where special image tokens are - _token_occupation_table = torch.ones_like(input_ids.flatten()) - _token_occupation_table[input_ids.flatten() == image_token_index] = torch.tensor( - feature_lengths, dtype=torch.long, device=input_ids.device - ) - _token_occupation_table = _token_occupation_table.reshape(input_ids.shape) + left_padding = (~(input_ids[:, -1] == pad_token_id)).sum() == 0 - max_embed_dim = _token_occupation_table.sum(-1).max().item() - assert max_embed_dim >= sequence_length, ( - f"The maximum embedding dimension ({max_embed_dim}) is less than the sequence length ({sequence_length})" - ) - batch_indices, non_image_indices = torch.where(input_ids != image_token_index) + image_token_mask = input_ids == image_token_index + non_image_mask = ~image_token_mask - # 2. Compute the positions where text should be written - # Calculate new positions for text tokens in merged image-text sequence. - new_token_positions = torch.cumsum(_token_occupation_table, -1) - 1 - nb_image_pad = max_embed_dim - 1 - new_token_positions[:, -1] - if left_padding: - new_token_positions += nb_image_pad[:, None] # offset for left padding - text_to_overwrite = new_token_positions[batch_indices, non_image_indices] - - # 3. Create the full embedding, already padded to the maximum position - final_embedding = torch.zeros( - batch_size, - max_embed_dim, - embed_dim, - dtype=inputs_embeds.dtype, - device=inputs_embeds.device, - ) - final_attention_mask = torch.zeros( - batch_size, max_embed_dim, dtype=attention_mask.dtype, device=inputs_embeds.device - ) - if labels is not None: - final_labels = torch.full( - (batch_size, max_embed_dim), - ignore_index, - dtype=input_ids.dtype, - device=input_ids.device, + flat_image_token_mask = image_token_mask.reshape(-1).to(torch.long) + flat_image_order = torch.cumsum(flat_image_token_mask, dim=0) + feature_lengths_tensor = torch.tensor(feature_lengths, dtype=input_ids.dtype, device=input_ids.device) + flat_image_lengths = torch.zeros_like(flat_image_order, dtype=input_ids.dtype) + for idx in range(len(feature_lengths)): + flat_image_lengths = ( + flat_image_lengths + (flat_image_order == (idx + 1)).to(input_ids.dtype) * (feature_lengths_tensor[idx]) ) - # In case the Vision model or the Language model has been offloaded to CPU, we need to manually - # set the corresponding tensors into their correct target device. - target_device = inputs_embeds.device - batch_indices, non_image_indices, text_to_overwrite = ( - batch_indices.to(target_device), - non_image_indices.to(target_device), - text_to_overwrite.to(target_device), - ) - attention_mask = attention_mask.to(target_device) - - # 4. Fill the embeddings based on the mask. - final_embedding[batch_indices, text_to_overwrite] = inputs_embeds[batch_indices, non_image_indices] - final_attention_mask[batch_indices, text_to_overwrite] = attention_mask[batch_indices, non_image_indices] - if labels is not None: - final_labels[batch_indices, text_to_overwrite] = labels[batch_indices, non_image_indices] - # 5. Fill the embeddings corresponding to the images. Anything that is not `text_positions` needs filling (#29835) - image_to_overwrite = torch.full( - (batch_size, max_embed_dim), True, dtype=torch.bool, device=inputs_embeds.device - ) - image_to_overwrite[batch_indices, text_to_overwrite] = False - image_to_overwrite &= image_to_overwrite.cumsum(-1) - 1 >= nb_image_pad[:, None].to(target_device) - - if image_to_overwrite.sum() != image_features.shape[:-1].numel(): - raise ValueError( - f"The input provided to the model are wrong. The number of image tokens is {image_to_overwrite.sum()} while" - f" the number of image features given to the model is {image_features.shape[:-1].numel()}. " - "This prevents correct indexing and breaks batch generation." - ) + token_occupation_table = torch.ones_like(input_ids) + token_occupation_table = token_occupation_table.reshape(-1) + token_occupation_table = torch.where(flat_image_token_mask.bool(), flat_image_lengths, token_occupation_table) + token_occupation_table = token_occupation_table.reshape(input_ids.shape) - final_embedding[image_to_overwrite] = image_features.contiguous().reshape(-1, embed_dim).to(target_device) - final_attention_mask |= image_to_overwrite - position_ids = (final_attention_mask.cumsum(-1) - 1).masked_fill_((final_attention_mask == 0), 1) + max_embed_dim = int(token_occupation_table.sum(-1).max().item()) - # 6. Mask out the embedding at padding positions, as we later use the past_key_value value to determine the non-attended tokens. - batch_indices, pad_indices = torch.where(input_ids == pad_token_id) - indices_to_mask = new_token_positions[batch_indices, pad_indices] + new_token_positions = torch.cumsum(token_occupation_table, 1) - 1 + nb_image_pad = max_embed_dim - 1 - new_token_positions[:, -1] + if bool(left_padding): + new_token_positions = new_token_positions + nb_image_pad[:, None] + + merged_positions = torch.arange(max_embed_dim, device=input_ids.device).view(1, 1, -1) + text_position_one_hot = merged_positions == new_token_positions.unsqueeze(-1) + text_position_one_hot = torch.logical_and(text_position_one_hot, non_image_mask.unsqueeze(-1)) + + final_embedding = ( + text_position_one_hot.to(inputs_embeds.dtype).unsqueeze(-1) * inputs_embeds.unsqueeze(2) + ).sum(dim=1) + final_attention_mask = (text_position_one_hot.to(attention_mask.dtype) * attention_mask.unsqueeze(-1)).sum( + dim=1 + ) - final_embedding[batch_indices, indices_to_mask] = 0 + image_to_overwrite = torch.logical_not(text_position_one_hot.any(dim=1)) + image_to_overwrite_cumsum = torch.cumsum(image_to_overwrite.to(torch.int64), dim=1) + image_to_overwrite = torch.logical_and( + image_to_overwrite, + image_to_overwrite_cumsum - 1 >= nb_image_pad[:, None].to(target_device), + ) - if labels is None: - final_labels = None + image_slot_ids = torch.cumsum(image_to_overwrite.to(torch.int64), dim=1) - 1 + if image_features_cat.shape[0] > 0: + valid_image_slots = torch.logical_and(image_to_overwrite, image_slot_ids < image_features_cat.shape[0]) + image_slot_ids = torch.clamp(image_slot_ids, min=0, max=image_features_cat.shape[0] - 1) + gathered_image_embeddings = image_features_cat.to(target_device)[image_slot_ids] + final_embedding = torch.where(valid_image_slots.unsqueeze(-1), gathered_image_embeddings, final_embedding) + image_to_overwrite = valid_image_slots + else: + image_to_overwrite = torch.zeros_like(image_to_overwrite) - return final_embedding, final_attention_mask, final_labels, position_ids + final_attention_mask = torch.logical_or(final_attention_mask.bool(), image_to_overwrite).to( + final_attention_mask.dtype + ) - def _extract_image_features(self, pixel_values: torch.Tensor, grid_thws: torch.Tensor) -> list[torch.Tensor]: - """ - Args: - pixel_values (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, num_channels, height, width)`): - The pixel values of the images processed by image processor. - grid_thws (:obj:`torch.Tensor` of shape :obj:`(batch_size, 3)`): - The grid, height, width of the images. - Returns: - selected_image_feature (:obj:`torch.FloatTensor` of shape :obj:`(num_image_tokens, embed_dim)`): - The selected image features to use as input to the projector head. - """ + position_ids = torch.cumsum(final_attention_mask, dim=1) - 1 + position_ids = torch.where(final_attention_mask == 0, torch.ones_like(position_ids), position_ids) - target_dtype = self.vision_tower.patch_embed.proj.weight.dtype - pixel_values = pixel_values.to(target_dtype) + pad_token_mask = input_ids == pad_token_id + pad_position_one_hot = torch.logical_and( + merged_positions == new_token_positions.unsqueeze(-1), pad_token_mask.unsqueeze(-1) + ) + pad_positions_mask = pad_position_one_hot.any(dim=1) + final_embedding = torch.where( + pad_positions_mask.unsqueeze(-1), torch.zeros_like(final_embedding), final_embedding + ) - image_features = self.vision_tower(pixel_values, grid_thws) - return image_features + return final_embedding, final_attention_mask, labels, position_ids def forward( self, @@ -1237,70 +817,6 @@ def forward( attentions=outputs.attentions, ) - def prepare_inputs_for_generation( - self, - input_ids, - past_key_values=None, - inputs_embeds=None, - pixel_values=None, - grid_thws=None, - attention_mask=None, - **kwargs, - ): - if past_key_values is not None: - if isinstance(past_key_values, Cache): - cache_length = past_key_values.get_seq_length() - past_length = getattr(past_key_values, "seen_tokens", cache_length) - else: - cache_length = past_length = past_key_values[0][0].shape[2] - - # Keep only the unprocessed tokens: - # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where - # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as - # input) - if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]: - input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :] - # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard - # input_ids based on the past_length. - elif past_length < input_ids.shape[1]: - input_ids = input_ids[:, past_length:] - # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens. - elif self.config.media_placeholder_token_id in input_ids: - input_ids = input_ids[:, input_ids.shape[1] - 1 :] - # If the cache has seen more tokens than it can hold, then the cache has a size limit. Let's discard the - # older attention values, as their corresponding values are not part of the input. - if cache_length < past_length and attention_mask is not None: - attention_mask = attention_mask[:, -(cache_length + input_ids.shape[1]) :] - - position_ids = kwargs.get("position_ids", None) - if attention_mask is not None and position_ids is None: - # create position_ids on the fly for batch generation - position_ids = attention_mask.long().cumsum(-1) - 1 - position_ids.masked_fill_(attention_mask == 0, 1) - if past_key_values: - position_ids = position_ids[:, -input_ids.shape[1] :] - - # if `inputs_embeds` are passed, we only want to use them in the 1st generation step - if inputs_embeds is not None and past_key_values is None: - model_inputs = {"inputs_embeds": inputs_embeds} - else: - model_inputs = {"input_ids": input_ids} - - model_inputs.update( - { - "position_ids": position_ids, - "past_key_values": past_key_values, - "use_cache": kwargs.get("use_cache"), - "attention_mask": attention_mask, - "pixel_values": pixel_values, - "grid_thws": grid_thws, - } - ) - return model_inputs - - def _reorder_cache(self, *args, **kwargs): - return self.language_model._reorder_cache(*args, **kwargs) - def get_output_names(self, kv_offload: bool = False): vision_output_names = ["image_embeds"] lang_output_names = ["logits"] @@ -1358,7 +874,7 @@ def get_dummy_inputs( "role": "user", "content": [ {"type": "image_url", "image_url": image}, - {"type": "text", "text": "Tell me about yourself."}, + {"type": "text", "text": "Describe this image."}, ], } ] @@ -1368,12 +884,17 @@ def get_dummy_inputs( tokenize=False, return_tensors="pt", ) - inputs = {k: v.to(self.language_model.device) if hasattr(v, "to") else v for k, v in inputs.items()} - # Build h_shape and w_shape from grid_thws (single image, t=1) + # Read grid dimensions before moving tensors to model device. During + # export flows that materialize modules on `meta`, calling `.item()` on + # meta tensors raises `RuntimeError`. grid_thws_val = inputs["grid_thws"] h_val = int(grid_thws_val[0, 1].item()) w_val = int(grid_thws_val[0, 2].item()) + + inputs = {k: v.to(self.language_model.device) if hasattr(v, "to") else v for k, v in inputs.items()} + + # Build h_shape and w_shape from grid_thws (single image, t=1) h_shape_tensor = torch.ones(h_val, dtype=torch.int64, device=self.language_model.device) w_shape_tensor = torch.ones(w_val, dtype=torch.int64, device=self.language_model.device) @@ -1383,8 +904,13 @@ def get_dummy_inputs( "w_shape": w_shape_tensor, } + # Ensure ONNX export traces the multimodal merge path in decoder wrapper. + # A dummy sequence without the media placeholder token can dead-code this + # branch and produce a decoder graph that ignores `image_embeds`. + input_ids = torch.ones((bs, prefill_seq_len), dtype=torch.int64) + input_ids[:, 0] = self.config.media_placeholder_token_id lang_inputs = { - "input_ids": torch.zeros((bs, prefill_seq_len), dtype=torch.int64), + "input_ids": input_ids, "position_ids": torch.arange(prefill_seq_len, dtype=torch.int64).view(1, prefill_seq_len).repeat(bs, 1), } @@ -1397,7 +923,7 @@ def get_dummy_inputs( pkv_cache = self.language_model.get_dummy_pkv_cache( config=self.language_model.config, batch_size=fbs if continuous_batching else bs, - seq_len=prefill_seq_len, + seq_len=1024, ) if cache_compressed: @@ -1510,7 +1036,7 @@ def get_specializations( num_image_tokens = compiler_options.pop("num_image_tokens", None) prefill_seq_len = prefill_seq_len if prefill_seq_len else 32 # constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN - ctx_len = ctx_len if ctx_len else 32 # constants.ONNX_EXPORT_EXAMPLE_CTX_LEN + ctx_len = ctx_len if ctx_len else 1024 # constants.ONNX_EXPORT_EXAMPLE_CTX_LEN num_patches = num_patches if num_patches is not None else 2400 h = h if h is not None else 30 diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index 74c2c54348..7cce2ed492 100755 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -1234,7 +1234,6 @@ def __init__(self, model, qaic_config: Optional[dict] = None, **kwargs): self.hash_params["qeff_auto_class"] = self.__class__.__name__ self.continuous_batching = False if qaic_config: - self.ccl_enabled = qaic_config.get("ccl_enabled", False) if mla_absorption := qaic_config.get("mla_absorption", None): self.hash_params["mla_absorption"] = mla_absorption setattr(self.model.language_model, "mla_absorption", mla_absorption) @@ -1989,8 +1988,6 @@ def compile( custom_io_vision = {} target_dtype = getattr(self.model.config, "torch_dtype", torch.float32) - if target_dtype == torch.bfloat16 and constants.DEFAULT_AIC_HW_VERSION != "ai200": - target_dtype = torch.float16 kv_cache_dtype = "mxint8" if mxint8_kv_cache else CUSTOM_IO_DTYPE_MAP[target_dtype] molmo = hasattr(self.model.config, "model_type") and self.model.config.model_type == "molmo" if molmo: @@ -2328,6 +2325,13 @@ def kv_offload_generate( vision_inputs_fp16 = {"pixel_values", "image_masks"} vision_inputs.update({k: vision_inputs[k].astype("float16") for k in vision_inputs_fp16 if k in vision_inputs}) + grid_thws_val = inputs.pop("grid_thws", None) + if grid_thws_val is not None: + h_val = int(grid_thws_val[0, 1].item()) + w_val = int(grid_thws_val[0, 2].item()) + vision_inputs["h_shape"] = np.ones((h_val), dtype=np.int64) + vision_inputs["w_shape"] = np.ones((w_val), dtype=np.int64) + vision_start = perf_counter() vision_outputs = {} @@ -2438,6 +2442,9 @@ def kv_offload_generate( if self._write_io_dir is not None: write_io_files(lang_inputs, outputs, self._write_io_dir, "prefill", "aic_batch_io", True, False) + if "image_idx_output" in outputs: + lang_inputs["image_idx"] = chunk_inputs["image_idx"] + prefill_time = perf_counter() - lang_start + vision_end - vision_start # Skip inputs/outputs again lang_session.skip_buffers( @@ -3064,7 +3071,9 @@ def cloud_ai_100_generate( prefill_time = perf_counter() - prefill_start # Get first token inputs["input_ids"] = outputs["logits"].argmax(2) - inputs["position_ids"] = input_len.numpy() + inputs["position_ids"] = np.max(inputs["position_ids"], axis=-1, keepdims=True) + 1 + if "image_idx_output" in outputs: + inputs["image_idx"] = chunk_inputs["image_idx"] if "cross_attention_mask" in inputs: bs, _, num_images, img_tiles = inputs["cross_attention_mask"].shape diff --git a/QEfficient/transformers/models/pytorch_transforms.py b/QEfficient/transformers/models/pytorch_transforms.py index 247340a94e..142c80f992 100755 --- a/QEfficient/transformers/models/pytorch_transforms.py +++ b/QEfficient/transformers/models/pytorch_transforms.py @@ -1232,6 +1232,7 @@ class KVCacheExternalModuleMapperTransform(ExternalModuleMapperTransform): "forward": QEFFGrok1CustomRMSNormAIC.forward, }, "KimiK25ForConditionalGeneration": { + "_merge_input_ids_with_image_features": QEffKimiK25ForConditionalGeneration._qeff_merge_input_ids_with_image_features, "get_qeff_vision_encoder": QEffKimiK25ForConditionalGeneration.get_qeff_vision_encoder, "get_qeff_language_decoder": QEffKimiK25ForConditionalGeneration.get_qeff_language_decoder, "get_specializations": QEffKimiK25ForConditionalGeneration.get_specializations, diff --git a/examples/kimi_k2/export_kimi_k25_vision.py b/examples/kimi_k2/export_kimi_k25_vision.py index c5fc366c1e..5412d6a725 100644 --- a/examples/kimi_k2/export_kimi_k25_vision.py +++ b/examples/kimi_k2/export_kimi_k25_vision.py @@ -36,8 +36,10 @@ EXPERT_KEY_PATTERN = re.compile(r"^(language_model\.model\.layers\.\d+\.mlp\.experts\.)(\d+)(\..+)$") + def _set_deterministic(seed: int): import random + import numpy as np random.seed(seed) @@ -47,6 +49,7 @@ def _set_deterministic(seed: int): torch.cuda.manual_seed_all(seed) torch.use_deterministic_algorithms(True) + def _ensure_torch_fx_import_compatibility(): """Backfill `is_torch_fx_available` for remote model code expecting older Transformers APIs.""" if hasattr(hf_import_utils, "is_torch_fx_available"): @@ -399,7 +402,7 @@ def main(): qeff_model.compile( qaic_config=qaic_config, prefill_seq_len=1, - ctx_len=4096, + ctx_len=1024, num_cores=16, num_devices=1, mxfp6_matmul=False, @@ -407,10 +410,6 @@ def main(): aic_enable_depth_first=False, skip_vision=True, # Skip vision encoder for text-only inference mos=1, - num_patches=2400, # num_patches - h=30, # h - w=80, # w - num_image_tokens=600, # num_image_tokens ) ## STEP 4: Prepare Text-Only Input # Create a text-only message without any image @@ -433,7 +432,7 @@ def main(): ) ## STEP 6: Run Text-Only Inference - output = qeff_model.generate(inputs=inputs, device_ids=[0, 1], generation_len=10) + output = qeff_model.generate(inputs=inputs, device_ids=[0], generation_len=10) ## STEP 7: Display Results print(output.generated_ids) @@ -448,7 +447,7 @@ def main(): qeff_model.compile( qaic_config=qaic_config, prefill_seq_len=1, - ctx_len=4096, + ctx_len=1024, num_cores=16, num_devices=1, mxfp6_matmul=False, @@ -488,7 +487,7 @@ def main(): inputs["pixel_values"] = inputs["pixel_values"].to(qeff_model.model.config.torch_dtype) ## STEP 6: Run Vision+Text Inference - output = qeff_model.generate(inputs=inputs, device_ids=[0], generation_len=10) + output = qeff_model.generate(inputs=inputs, generation_len=10) ## STEP 7: Display Results print(output.generated_ids) diff --git a/examples/kimi_k2/test_kimi_k25.py b/examples/kimi_k2/test_kimi_k25.py index 821b6434c7..7e19ac4083 100644 --- a/examples/kimi_k2/test_kimi_k25.py +++ b/examples/kimi_k2/test_kimi_k25.py @@ -1,185 +1,446 @@ -# ----------------------------------------------------------------------------- -# -# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. -# SPDX-License-Identifier: BSD-3-Clause -# -# ----------------------------------------------------------------------------- - -import os -from io import BytesIO -from pathlib import Path - -import requests -import torch -from export_kimi_k25_vision import ( - LOADED_EXPERT_IDS, - NUM_EXPERTS_PER_TOKEN, - _ensure_torch_fx_import_compatibility, - _load_layer_subset_model, - _patch_deepseek_init_weights_compat, - _patch_kimi_tie_weights_compat, - _prepare_config, -) -from huggingface_hub import snapshot_download -from PIL import Image -from transformers.dynamic_module_utils import get_class_from_dynamic_module - -from QEfficient import QEFFAutoModelForImageTextToText - -MODEL_NAME = "moonshotai/Kimi-K2.5" -IMAGE_URL = "https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png" -TEXT_PROMPT = "Describe this image." -NUM_VISION_LAYERS = 2 -NUM_TEXT_LAYERS = 2 -NEW_GENERATION_TOKENS = 10 -CTX_LEN = 4096 - -def _set_deterministic(seed: int): - import random - import numpy as np - - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(seed) - torch.use_deterministic_algorithms(True) - -def _resolve_model_path() -> Path: - return Path(snapshot_download(repo_id=MODEL_NAME, cache_dir=os.environ.get("HF_HUB_CACHE"))) - -def _prepare_inputs(processor): - image = Image.open(BytesIO(requests.get(IMAGE_URL, timeout=30).content)).convert("RGB") - messages = [ - { - "role": "user", - "content": [ - {"type": "image_url", "image_url": image}, - {"type": "text", "text": TEXT_PROMPT}, - ], - } - ] - return processor( - messages=messages, - add_generation_prompt=True, - tokenize=False, - return_tensors="pt", - ) - -def _greedy_generate_cpu(model, inputs, max_new_tokens: int): - generated_ids = inputs["input_ids"] - attention_mask = inputs["attention_mask"] - pixel_values = inputs["pixel_values"] - grid_thws = inputs["grid_thws"] - - eos_token_id = getattr(model.config, "eos_token_id", None) - if eos_token_id is None and hasattr(model.config, "text_config"): - eos_token_id = getattr(model.config.text_config, "eos_token_id", None) - - for _ in range(max_new_tokens): - outputs = model( - input_ids=generated_ids, - attention_mask=attention_mask, - pixel_values=pixel_values, - grid_thws=grid_thws, - use_cache=False, - return_dict=True, - ) - logits = outputs[0] if isinstance(outputs, tuple) else outputs.logits - next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True) - - generated_ids = torch.cat([generated_ids, next_token], dim=1) - attention_mask = torch.cat( - [ - attention_mask, - torch.ones((attention_mask.shape[0], 1), dtype=attention_mask.dtype, device=attention_mask.device), - ], - dim=1, - ) - - if eos_token_id is not None and torch.all(next_token == eos_token_id): - break - - return generated_ids[:, -max_new_tokens:] - - -def check_kimi_k25_pytorch_vs_ai100(): - _set_deterministic(1234) - _ensure_torch_fx_import_compatibility() - model_path = _resolve_model_path() - config = _prepare_config(model_path) - kimi_cls = get_class_from_dynamic_module("modeling_kimi_k25.KimiK25ForConditionalGeneration", str(model_path)) - _patch_kimi_tie_weights_compat(kimi_cls) - _patch_deepseek_init_weights_compat(kimi_cls) - - model, tokenizer, processor = _load_layer_subset_model( - model_path=model_path, - kimi_cls=kimi_cls, - config=config, - num_vision_layers=NUM_VISION_LAYERS, - num_text_layers=NUM_TEXT_LAYERS, - loaded_expert_ids=LOADED_EXPERT_IDS, - num_experts_per_tok=NUM_EXPERTS_PER_TOKEN, - dtype=torch.float32, - ) - - model.vision_tower.patch_embed.pos_emb.interpolation_mode = "bilinear" - model = model.eval().to("cpu") - - inputs = _prepare_inputs(processor) - inputs = {k: (v.to("cpu") if torch.is_tensor(v) else v) for k, v in inputs.items()} - - with torch.inference_mode(): - generated_ids = _greedy_generate_cpu(model, inputs, max_new_tokens=NEW_GENERATION_TOKENS) - - decoded = tokenizer.batch_decode(generated_ids, skip_special_tokens=True) - - print( - "Loaded subset:", - f"vision_layers={model.config.vision_config.vt_num_hidden_layers}", - f"text_layers={model.config.text_config.num_hidden_layers}", - ) - print("Prompt:", repr(TEXT_PROMPT)) - print("Generated IDs shape:", tuple(generated_ids.shape)) - print("Output:") - print(decoded[0] if decoded else "") - - qeff_model = QEFFAutoModelForImageTextToText( - model, - kv_offload=True, - torch_dtype=torch.float32, - ) - - _ = qeff_model.export() - qeff_model.compile( - num_devices=1, - prefill_seq_len=1, # prompt_len, - ctx_len=CTX_LEN, - mxfp6=False, - num_patches=int(inputs["pixel_values"].shape[0]), - h=int(inputs["grid_thws"][0, 1].item()), - w=int(inputs["grid_thws"][0, 2].item()), - num_image_tokens=600, - num_cores=16, - ) - - exec_info = qeff_model.generate(inputs=inputs, generation_len=NEW_GENERATION_TOKENS) - cloud_ai_100_tokens = exec_info.generated_ids[:, :-1] - decoded = tokenizer.batch_decode(cloud_ai_100_tokens, skip_special_tokens=True) - - print( - "Loaded subset:", - f"vision_layers={model.config.vision_config.vt_num_hidden_layers}", - f"text_layers={model.config.text_config.num_hidden_layers}", - ) - print("Prompt:", repr(TEXT_PROMPT)) - print("Generated IDs shape:", tuple(cloud_ai_100_tokens.shape)) - print("Output:") - print(decoded[0] if decoded else "") - - assert (generated_ids == cloud_ai_100_tokens).all(), "Tokens don't match for pytorch HF output and QPC output" - - -def test_kimi_k25_image_text_to_text_pytorch_vs_ai100(): - torch.manual_seed(42) - check_kimi_k25_pytorch_vs_ai100() +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import copy +import os +import re +from io import BytesIO +from pathlib import Path + +import numpy as np +import onnx +import onnxruntime as ort +import requests +import torch +from export_kimi_k25_vision import ( + LOADED_EXPERT_IDS, + NUM_EXPERTS_PER_TOKEN, + _ensure_torch_fx_import_compatibility, + _load_layer_subset_model, + _patch_deepseek_init_weights_compat, + _patch_kimi_tie_weights_compat, + _prepare_config, +) +from huggingface_hub import snapshot_download +from PIL import Image +from transformers.dynamic_module_utils import get_class_from_dynamic_module + +from QEfficient import QEFFAutoModelForImageTextToText +from QEfficient.generation.cloud_infer import QAICInferenceSession + +MODEL_NAME = "moonshotai/Kimi-K2.5" +IMAGE_URL = "https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png" +TEXT_PROMPT = "Describe this image." +NUM_VISION_LAYERS = 2 +NUM_TEXT_LAYERS = 2 +NEW_GENERATION_TOKENS = 10 +CTX_LEN = 1024 + + +def _materialize_missing_linear_weights(module: torch.nn.Module): + """Decompress compressed linear modules that expose packed weights but no `.weight`.""" + try: + from compressed_tensors.compressors.base import decompress_module as ct_decompress_module + except Exception: + return + + for child in module.modules(): + if not isinstance(child, torch.nn.Linear): + continue + if hasattr(child, "weight"): + continue + if not (hasattr(child, "weight_packed") and hasattr(child, "weight_scale")): + continue + ct_decompress_module(child) + + +def _disable_compressed_tensor_forward_pre_hooks(module: torch.nn.Module): + """Avoid double-decompress hooks that can raise KeyError('weight_packed').""" + if hasattr(module, "_forward_pre_hooks"): + module._forward_pre_hooks.clear() + + +def _has_qaic_runtime_access() -> bool: + try: + _ = QAICInferenceSession + except Exception: + return False + try: + import qaicrt + + _ctx = qaicrt.Context() + return True + except Exception: + return False + + +def _set_deterministic(seed: int): + import random + + import numpy as np + + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + torch.use_deterministic_algorithms(True) + + +def _resolve_model_path() -> Path: + return Path(snapshot_download(repo_id=MODEL_NAME, cache_dir=os.environ.get("HF_HUB_CACHE"))) + + +def _prepare_inputs(processor): + image = Image.open(BytesIO(requests.get(IMAGE_URL, timeout=30).content)).convert("RGB") + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": image}, + {"type": "text", "text": TEXT_PROMPT}, + ], + } + ] + return processor( + messages=messages, + add_generation_prompt=True, + tokenize=False, + return_tensors="pt", + ) + + +def _decode_tokens(tokenizer, token_ids: torch.Tensor) -> str: + decoded = tokenizer.batch_decode(token_ids, skip_special_tokens=True) + return decoded[0] if decoded else "" + + +def _clone_inputs(inputs): + return {k: (v.clone() if torch.is_tensor(v) else copy.deepcopy(v)) for k, v in inputs.items()} + + +def _greedy_generate_hf(model, inputs, max_new_tokens: int): + generated_ids = inputs["input_ids"] + attention_mask = inputs["attention_mask"] + pixel_values = inputs["pixel_values"] + grid_thws = inputs["grid_thws"] + + eos_token_id = getattr(model.config, "eos_token_id", None) + if eos_token_id is None and hasattr(model.config, "text_config"): + eos_token_id = getattr(model.config.text_config, "eos_token_id", None) + + for _ in range(max_new_tokens): + outputs = model( + input_ids=generated_ids, + attention_mask=attention_mask, + pixel_values=pixel_values, + grid_thws=grid_thws, + use_cache=False, + return_dict=True, + ) + logits = outputs[0] if isinstance(outputs, tuple) else outputs.logits + next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True) + + generated_ids = torch.cat([generated_ids, next_token], dim=1) + attention_mask = torch.cat( + [ + attention_mask, + torch.ones((attention_mask.shape[0], 1), dtype=attention_mask.dtype, device=attention_mask.device), + ], + dim=1, + ) + + if eos_token_id is not None and torch.all(next_token == eos_token_id): + break + + return generated_ids[:, -max_new_tokens:] + + +def _match_indexed_name(name: str, base: str): + pattern = rf"^{re.escape(base)}\.(\d+)$" + match = re.match(pattern, name) + return int(match.group(1)) if match else None + + +def _build_cache_inputs_from_dummy(transformed_model, input_names, batch_size: int): + pkv = transformed_model.language_model.get_dummy_pkv_cache( + transformed_model.config.text_config, + batch_size, + CTX_LEN, + ) + cache_inputs = {} + for name in input_names: + key_idx = _match_indexed_name(name, "past_key") + value_idx = _match_indexed_name(name, "past_value") + compressed_idx = _match_indexed_name(name, "compressed_kv") + k_pe_idx = _match_indexed_name(name, "k_pe") + + if key_idx is not None: + cache_inputs[name] = pkv[key_idx][0].detach().cpu().numpy().astype(np.float32) + elif value_idx is not None: + cache_inputs[name] = pkv[value_idx][1].detach().cpu().numpy().astype(np.float32) + elif compressed_idx is not None: + cache_inputs[name] = pkv[compressed_idx][0].detach().cpu().numpy().astype(np.float32) + elif k_pe_idx is not None: + cache_inputs[name] = pkv[k_pe_idx][1].detach().cpu().numpy().astype(np.float32) + return cache_inputs + + +def _greedy_generate_qeff_wrapper(transformed_model, inputs, max_new_tokens: int): + qeff_encoder = transformed_model.get_qeff_vision_encoder().eval() + decoder_wrapper = transformed_model.get_qeff_language_decoder().eval() + + grid_thws = inputs["grid_thws"].to(torch.long) + h_shape = torch.ones(int(grid_thws[0, 1].item()), dtype=torch.int64) + w_shape = torch.ones(int(grid_thws[0, 2].item()), dtype=torch.int64) + image_embeds = qeff_encoder(inputs["pixel_values"].to(torch.float32), h_shape, w_shape).detach() + + generated_ids = inputs["input_ids"].to(torch.long) + attention_mask = inputs["attention_mask"].to(torch.long) + + for _ in range(max_new_tokens): + position_ids = torch.where(attention_mask.bool(), attention_mask.cumsum(-1) - 1, -1) + past_key_values = transformed_model.language_model.get_dummy_pkv_cache( + transformed_model.config.text_config, + generated_ids.shape[0], + CTX_LEN, + ) + logits = decoder_wrapper( + input_ids=generated_ids, + attention_mask=attention_mask, + position_ids=position_ids, + image_embeds=image_embeds, + image_idx=torch.zeros((generated_ids.shape[0], 1), dtype=torch.int64), + past_key_values=past_key_values, + use_cache=True, + )[0] + next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True) + + generated_ids = torch.cat([generated_ids, next_token], dim=1) + attention_mask = torch.cat( + [attention_mask, torch.ones((attention_mask.shape[0], 1), dtype=attention_mask.dtype)], + dim=1, + ) + + return generated_ids[:, -max_new_tokens:] + + +def _greedy_generate_onnx(transformed_model, onnx_paths, inputs, max_new_tokens, session_options): + vision_onnx_path, lang_onnx_path = onnx_paths + vision_session = ort.InferenceSession(str(vision_onnx_path)) # , providers=["CPUExecutionProvider"]) + lang_session = ort.InferenceSession(str(lang_onnx_path), session_options) # providers=["CPUExecutionProvider"]) + + pixel_values = inputs["pixel_values"].detach().cpu().numpy().astype(np.float32) + grid_thws = inputs["grid_thws"].detach().cpu().numpy().astype(np.int64) + h = int(grid_thws[0, 1]) + w = int(grid_thws[0, 2]) + + vision_outputs = vision_session.run( + None, + { + "pixel_values": pixel_values, + "h_shape": np.ones((h,), dtype=np.int64), + "w_shape": np.ones((w,), dtype=np.int64), + }, + ) + vision_output_names = [out.name for out in vision_session.get_outputs()] + image_embeds = {name: value for name, value in zip(vision_output_names, vision_outputs)}.get( + "image_embeds", vision_outputs[0] + ) + + lang_input_names = [meta.name for meta in lang_session.get_inputs()] + lang_output_names = [meta.name for meta in lang_session.get_outputs()] + + prompt_input_ids = inputs["input_ids"].detach().cpu().numpy().astype(np.int64) + prompt_attention_mask = inputs["attention_mask"].detach().cpu().numpy().astype(np.int64) + prompt_len = prompt_input_ids.shape[1] + + prefill_seq_len = 32 + num_chunks = -(prompt_len // -prefill_seq_len) + padded_len = num_chunks * prefill_seq_len + + pad_token_id = 1 + padded_input_ids = np.pad(prompt_input_ids, ((0, 0), (0, padded_len - prompt_len)), constant_values=pad_token_id) + padded_attention = np.pad(prompt_attention_mask, ((0, 0), (0, padded_len - prompt_len)), constant_values=0) + padded_position_ids = np.where(padded_attention > 0, np.arange(padded_len), -1).astype(np.int64) + + cache_inputs = _build_cache_inputs_from_dummy( + transformed_model=transformed_model, + input_names=lang_input_names, + batch_size=prompt_input_ids.shape[0], + ) + image_idx = np.zeros((prompt_input_ids.shape[0], 1), dtype=np.int64) + + output_map = None + for chunk_idx in range(num_chunks): + start = chunk_idx * prefill_seq_len + end = (chunk_idx + 1) * prefill_seq_len + ort_inputs = { + "input_ids": padded_input_ids[:, start:end], + "position_ids": padded_position_ids[:, start:end], + "image_embeds": image_embeds.astype(np.float32), + "image_idx": image_idx, + **cache_inputs, + } + ort_outputs = lang_session.run(None, ort_inputs) + output_map = {name: value for name, value in zip(lang_output_names, ort_outputs)} + if "image_idx_output" in output_map: + image_idx = output_map["image_idx_output"].astype(np.int64) + for key in list(cache_inputs.keys()): + retained_name = f"{key}_RetainedState" + if retained_name in output_map: + cache_inputs[key] = output_map[retained_name] + + if output_map is None: + raise RuntimeError("ONNX prefill did not execute.") + + next_token = np.argmax(output_map["logits"][:, -1, :], axis=-1, keepdims=True).astype(np.int64) + generated_new_tokens = [next_token] + + decode_input_ids = next_token + decode_position_ids = np.max(padded_position_ids, axis=1, keepdims=True).astype(np.int64) + 1 + + for _ in range(1, max_new_tokens): + ort_inputs = { + "input_ids": decode_input_ids, + "position_ids": decode_position_ids, + "image_embeds": image_embeds.astype(np.float32), + "image_idx": image_idx, + **cache_inputs, + } + ort_outputs = lang_session.run(None, ort_inputs) + output_map = {name: value for name, value in zip(lang_output_names, ort_outputs)} + + decode_input_ids = np.argmax(output_map["logits"][:, -1, :], axis=-1, keepdims=True).astype(np.int64) + generated_new_tokens.append(decode_input_ids) + decode_position_ids = decode_position_ids + 1 + + if "image_idx_output" in output_map: + image_idx = output_map["image_idx_output"].astype(np.int64) + for key in list(cache_inputs.keys()): + retained_name = f"{key}_RetainedState" + if retained_name in output_map: + cache_inputs[key] = output_map[retained_name] + + return torch.from_numpy(np.concatenate(generated_new_tokens, axis=1)) + + +def check_kimi_k25_pytorch_vs_ai100(): + _set_deterministic(1234) + _ensure_torch_fx_import_compatibility() + model_path = _resolve_model_path() + config = _prepare_config(model_path) + kimi_cls = get_class_from_dynamic_module("modeling_kimi_k25.KimiK25ForConditionalGeneration", str(model_path)) + _patch_kimi_tie_weights_compat(kimi_cls) + _patch_deepseek_init_weights_compat(kimi_cls) + + model, tokenizer, processor = _load_layer_subset_model( + model_path=model_path, + kimi_cls=kimi_cls, + config=config, + num_vision_layers=NUM_VISION_LAYERS, + num_text_layers=NUM_TEXT_LAYERS, + loaded_expert_ids=LOADED_EXPERT_IDS, + num_experts_per_tok=NUM_EXPERTS_PER_TOKEN, + dtype=torch.float32, + ) + + _materialize_missing_linear_weights(model.language_model) + _disable_compressed_tensor_forward_pre_hooks(model) + model.vision_tower.patch_embed.pos_emb.interpolation_mode = "bilinear" + model = model.eval().to("cpu") + + inputs = _prepare_inputs(processor) + inputs = {k: (v.to("cpu") if torch.is_tensor(v) else v) for k, v in inputs.items()} + + hf_tokens = _greedy_generate_hf(model, _clone_inputs(inputs), max_new_tokens=NEW_GENERATION_TOKENS) + print("HF:", _decode_tokens(tokenizer, hf_tokens), "\n", hf_tokens) + + qeff_model = QEFFAutoModelForImageTextToText( + copy.deepcopy(model), + kv_offload=True, + torch_dtype=torch.float32, + ) + _disable_compressed_tensor_forward_pre_hooks(qeff_model.model) + + qeff_tokens = _greedy_generate_qeff_wrapper( + transformed_model=qeff_model.model, + inputs=_clone_inputs(inputs), + max_new_tokens=NEW_GENERATION_TOKENS, + ) + print("QEFF:", _decode_tokens(tokenizer, qeff_tokens), "\n", qeff_tokens) + + onnx_paths = qeff_model.export() + + # Replace invalid index value for INT32 max to 0 using add_initializer + m = onnx.load(onnx_paths[1], load_external_data=False) + # NOTE: OrtValue objects should be kept around until the session is run, hence this dict is required + added_initializers = {} + for node in m.graph.node: + if node.op_type == "Constant": + np_tensor = onnx.numpy_helper.to_array(node.attribute[0].t, os.path.dirname(onnx_paths[1])) + if len(np_tensor.shape) == 0 and np_tensor.item() == 2147483647: + added_initializers[node.output[0]] = ort.OrtValue.ortvalue_from_numpy(np.array(0, np_tensor.dtype)) + + session_options = ort.SessionOptions() + for name, value in added_initializers.items(): + session_options.add_initializer(name, value) + + try: + onnx_tokens = _greedy_generate_onnx( + transformed_model=qeff_model.model, + onnx_paths=onnx_paths, + inputs=_clone_inputs(inputs), + max_new_tokens=NEW_GENERATION_TOKENS, + session_options=session_options, + ) + print("ONNX:", _decode_tokens(tokenizer, onnx_tokens) if onnx_tokens is not None else "") + except Exception as exc: + print(f"ONNX generation failed: {exc}") + onnx_tokens = None + + qeff_model.compile( + num_devices=1, + prefill_seq_len=1, + ctx_len=CTX_LEN, + mxfp6=False, + num_patches=int(inputs["pixel_values"].shape[0]), + h=int(inputs["grid_thws"][0, 1].item()), + w=int(inputs["grid_thws"][0, 2].item()), + num_image_tokens=600, + num_cores=16, + ) + + qaic_tokens = None + if _has_qaic_runtime_access(): + qaic_tokens = qeff_model.generate( + inputs=_clone_inputs(inputs), generation_len=NEW_GENERATION_TOKENS + ).generated_ids[:, :-1] + print("QAIC:", _decode_tokens(tokenizer, qaic_tokens) if qaic_tokens is not None else "") + else: + print("QAIC generation skipped: no QAIC runtime access (user/group permissions).") + + print( + "Loaded subset:", + f"vision_layers={model.config.vision_config.vt_num_hidden_layers}", + f"text_layers={model.config.text_config.num_hidden_layers}", + ) + print("Prompt:", repr(TEXT_PROMPT)) + print("HF:", _decode_tokens(tokenizer, hf_tokens)) + print("QEFF:", _decode_tokens(tokenizer, qeff_tokens)) + print("ONNX:", _decode_tokens(tokenizer, onnx_tokens) if onnx_tokens is not None else "") + print("QAIC:", _decode_tokens(tokenizer, qaic_tokens) if qaic_tokens is not None else "") + + assert torch.equal(hf_tokens, qeff_tokens), "HF and QEFF(Pytorch runtime) tokens do not match" + if onnx_tokens is not None: + assert torch.equal(hf_tokens, onnx_tokens), "HF and ONNXRuntime tokens do not match" + if qaic_tokens is not None: + assert torch.equal(hf_tokens, torch.as_tensor(qaic_tokens)), "HF and QAIC tokens do not match" + + +def test_kimi_k25_image_text_to_text_pytorch_vs_ai100(): + check_kimi_k25_pytorch_vs_ai100() diff --git a/examples/kimi_k2/test_kimi_k25_text.py b/examples/kimi_k2/test_kimi_k25_text.py deleted file mode 100644 index 6d4166a39f..0000000000 --- a/examples/kimi_k2/test_kimi_k25_text.py +++ /dev/null @@ -1,430 +0,0 @@ -# ----------------------------------------------------------------------------- -# -# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. -# SPDX-License-Identifier: BSD-3-Clause -# -# ---------------------------------------------------------------------------- - -import copy -import json -import re -import tempfile -from collections import defaultdict -from pathlib import Path - -import numpy as np -import pytest -import torch -from export_kimi_k25_vision import ( - _ensure_torch_fx_import_compatibility, - _patch_deepseek_init_weights_compat, - _patch_kimi_tie_weights_compat, -) -from safetensors import safe_open -from safetensors.torch import save_file -from transformers import AutoConfig, AutoTokenizer -from transformers.cache_utils import Cache, DynamicCache -from transformers.dynamic_module_utils import get_class_from_dynamic_module - -from QEfficient import QEFFAutoModelForCausalLM - -MODEL_PATH = Path( - "/home/huggingface_hub/models--moonshotai--Kimi-K2.5/snapshots/4d01dfe0332d63057c186e0b262165819efb6611" -) -NUM_HIDDEN_LAYERS = 2 -LOADED_EXPERT_IDS = (0, 1, 2, 3) -NUM_EXPERTS_PER_TOKEN = 2 - -EXPERT_KEY_PATTERN = re.compile(r"^(language_model\.model\.layers\.\d+\.mlp\.experts\.)(\d+)(\..+)$") - - -def _patch_dynamic_cache_compat(): - if not hasattr(DynamicCache, "from_legacy_cache"): - - @classmethod - def _from_legacy_cache(cls, legacy_cache): - if legacy_cache is None: - return cls() - if isinstance(legacy_cache, cls): - return legacy_cache - if isinstance(legacy_cache, Cache): - return cls(ddp_cache_data=[tuple(layer[:2]) for layer in legacy_cache]) - - ddp_cache_data = [] - for layer in legacy_cache: - if layer is None: - continue - if len(layer) < 2: - raise ValueError("Each legacy cache layer must provide key/value tensors.") - ddp_cache_data.append((layer[0], layer[1])) - return cls(ddp_cache_data=ddp_cache_data) - - DynamicCache.from_legacy_cache = _from_legacy_cache - - if not hasattr(DynamicCache, "to_legacy_cache"): - - def _to_legacy_cache(self): - return tuple((layer[0], layer[1]) for layer in self) - - DynamicCache.to_legacy_cache = _to_legacy_cache - - if not hasattr(DynamicCache, "get_max_length"): - - def _get_max_length(self): - return None - - DynamicCache.get_max_length = _get_max_length - - -def _validate_expert_subset(loaded_expert_ids, num_experts_per_tok, total_experts): - expert_ids = tuple(loaded_expert_ids) - if len(expert_ids) != 4: - raise ValueError(f"Expected exactly 4 routed experts, got {expert_ids!r}.") - if len(set(expert_ids)) != len(expert_ids): - raise ValueError(f"Expert ids must be unique, got {expert_ids!r}.") - invalid_ids = [expert_id for expert_id in expert_ids if expert_id < 0 or expert_id >= total_experts] - if invalid_ids: - raise ValueError(f"Expert ids {invalid_ids!r} are outside the valid range [0, {total_experts - 1}].") - if num_experts_per_tok > len(expert_ids): - raise ValueError(f"num_experts_per_tok={num_experts_per_tok} cannot exceed {len(expert_ids)} loaded experts.") - return expert_ids - - -def _remap_checkpoint_key(checkpoint_key, expert_index_map): - match = EXPERT_KEY_PATTERN.match(checkpoint_key) - if not match: - return checkpoint_key - - original_expert_idx = int(match.group(2)) - remapped_expert_idx = expert_index_map.get(original_expert_idx) - if remapped_expert_idx is None: - return None - return f"{match.group(1)}{remapped_expert_idx}{match.group(3)}" - - -def _is_routed_gate_weight(checkpoint_key): - return checkpoint_key.endswith(".mlp.gate.weight") - - -def _is_routed_gate_bias(checkpoint_key): - return checkpoint_key.endswith(".mlp.gate.e_score_correction_bias") - - -def _materialize_subset_checkpoint( - model_path: Path, - temp_model_path: Path, - weight_map, - allowed_prefixes, - loaded_expert_ids, -): - expert_index_map = {expert_id: remapped_idx for remapped_idx, expert_id in enumerate(loaded_expert_ids)} - selected_by_shard = defaultdict(list) - - for checkpoint_key, shard_name in weight_map.items(): - if not any(checkpoint_key.startswith(prefix) for prefix in allowed_prefixes): - continue - - remapped_key = _remap_checkpoint_key(checkpoint_key, expert_index_map) - if remapped_key is None: - continue - selected_by_shard[shard_name].append((checkpoint_key, remapped_key)) - - if not selected_by_shard: - raise RuntimeError("No text-only weights were selected from the Kimi K2.5 checkpoint.") - - filtered_weight_map = {} - subset_shards = [] - for shard_idx, (source_shard_name, shard_entries) in enumerate(sorted(selected_by_shard.items())): - tensors = {} - with safe_open(model_path / source_shard_name, framework="pt", device="cpu") as shard_reader: - for checkpoint_key, remapped_key in shard_entries: - tensor = shard_reader.get_tensor(checkpoint_key) - if _is_routed_gate_weight(checkpoint_key): - tensor = tensor[list(loaded_expert_ids), :].contiguous() - elif _is_routed_gate_bias(checkpoint_key): - tensor = tensor[list(loaded_expert_ids)].contiguous() - tensors[remapped_key] = tensor - - subset_shard_name = f"model-subset-{shard_idx:05d}.safetensors" - save_file(tensors, str(temp_model_path / subset_shard_name)) - subset_shards.append(subset_shard_name) - filtered_weight_map.update({remapped_key: subset_shard_name for _, remapped_key in shard_entries}) - - return filtered_weight_map, subset_shards - - -def load_text_only_kimi( - model_path: Path, - num_hidden_layers: int, - loaded_expert_ids=LOADED_EXPERT_IDS, - num_experts_per_tok: int = NUM_EXPERTS_PER_TOKEN, -): - kimi_config = AutoConfig.from_pretrained(str(model_path), trust_remote_code=True) - - # Kimi K2.5 is multimodal, so the text depth must be overridden on text_config. - text_config = copy.deepcopy(kimi_config.text_config) - text_config.num_hidden_layers = num_hidden_layers - loaded_expert_ids = _validate_expert_subset( - loaded_expert_ids, - num_experts_per_tok, - text_config.n_routed_experts, - ) - text_config.n_routed_experts = len(loaded_expert_ids) - text_config.num_experts_per_tok = num_experts_per_tok - text_config.n_group = 1 - text_config.topk_group = 1 - - deepseek_cls = get_class_from_dynamic_module("modeling_deepseek.DeepseekV3ForCausalLM", str(model_path)) - - checkpoint_index = json.loads((model_path / "model.safetensors.index.json").read_text()) - weight_map = checkpoint_index["weight_map"] - - allowed_prefixes = [ - "language_model.model.embed_tokens.", - "language_model.model.norm.", - "language_model.lm_head.", - ] - allowed_prefixes.extend(f"language_model.model.layers.{layer_idx}." for layer_idx in range(num_hidden_layers)) - - with tempfile.TemporaryDirectory() as tmpdir: - temp_model_path = Path(tmpdir) - filtered_weight_map, subset_shards = _materialize_subset_checkpoint( - model_path=model_path, - temp_model_path=temp_model_path, - weight_map=weight_map, - allowed_prefixes=allowed_prefixes, - loaded_expert_ids=loaded_expert_ids, - ) - (temp_model_path / "config.json").write_text(text_config.to_json_string(use_diff=False)) - (temp_model_path / "model.safetensors.index.json").write_text( - json.dumps( - { - "metadata": { - "total_size": sum((temp_model_path / shard_name).stat().st_size for shard_name in subset_shards) - }, - "weight_map": filtered_weight_map, - } - ) - ) - - # We are loading a task checkpoint into the base text model, so disable the - # base/task prefix heuristic and let `key_mapping` strip `language_model.`. - original_base_model_prefix = deepseek_cls.base_model_prefix - deepseek_cls.base_model_prefix = "" - try: - model, loading_info = deepseek_cls.from_pretrained( - str(temp_model_path), - torch_dtype=torch.float32, - config=text_config, - local_files_only=True, - key_mapping={r"^language_model\.": ""}, - output_loading_info=True, - ) - finally: - deepseek_cls.base_model_prefix = original_base_model_prefix - - unexpected_keys = loading_info["unexpected_keys"] - missing_keys = loading_info["missing_keys"] - mismatched_keys = loading_info["mismatched_keys"] - if unexpected_keys or missing_keys or mismatched_keys: - raise RuntimeError( - "Failed to load the text-only Kimi K2.5 checkpoint slice cleanly. " - f"missing={missing_keys}, unexpected={unexpected_keys}, mismatched={mismatched_keys}" - ) - - model.eval() - tokenizer = AutoTokenizer.from_pretrained(str(model_path), trust_remote_code=True) - return model, tokenizer - - -""" -mla_absorption = {"cache_compressed": False, "absorption": False, "online": False} -qaic_configs=[ - None, # Full PKV Cache - {"mla_absorption": mla_absorption}, - {"enable_blocking": True, "blocking_mode": "h"}, # Full PKV Cache with Head Blocking - {"mla_absorption": mla_absorption, "enable_blocking": True, "blocking_mode": "h"}, # Full PKV Cache with Head Blocking -] -""" -TS = 1 -mla_absorption = {"cache_compressed": True, "absorption": False, "online": False} -# mla_absorption = {"cache_compressed": True, "absorption": True, "online": False} -# mla_absorption = {"cache_compressed": True, "absorption": True, "online": True} -qaic_configs = [ - {"mla_absorption": mla_absorption}, # for No Blocking - # {"mla_absorption": mla_absorption, "num_kv_heads_repeat": TS}, # No blocking with kv head replication - # {"mla_absorption": mla_absorption, "enable_blocking": True, "blocking_mode": "kv"}, # for KV blocking - # {"mla_absorption": mla_absorption, "enable_blocking": True, "blocking_mode": "kv", "num_kv_heads_repeat":TS}, # for KV blocking with kv head replication - # {"mla_absorption": mla_absorption, "enable_blocking": True, "blocking_mode": "h","num_kv_heads_repeat": TS}, # for h blocking with kv head replication -] - - -def _set_deterministic(seed: int): - import random - - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(seed) - torch.use_deterministic_algorithms(True) - - -@pytest.mark.parametrize("qaic_config", qaic_configs, ids=lambda x: f"qaic_{x}") -def test_qeff_generation(qaic_config): - prompt = "Once upon a time," - PREFILL_SEQ_LEN = 512 - CTX_LEN = 1024 - - _set_deterministic(1234) - _ensure_torch_fx_import_compatibility() - _patch_dynamic_cache_compat() - #config = _prepare_config(MODEL_PATH) - kimi_cls = get_class_from_dynamic_module("modeling_deepseek.DeepseekV3ForCausalLM", str(MODEL_PATH)) - _patch_kimi_tie_weights_compat(kimi_cls) - _patch_deepseek_init_weights_compat(kimi_cls) - model, tokenizer = load_text_only_kimi( - MODEL_PATH, - NUM_HIDDEN_LAYERS, - loaded_expert_ids=LOADED_EXPERT_IDS, - num_experts_per_tok=NUM_EXPERTS_PER_TOKEN, - ) - - inputs = tokenizer(prompt, return_tensors="pt", padding=True) - padded_len = inputs["input_ids"].shape[1] - num_chunks = -(padded_len // -PREFILL_SEQ_LEN) # ceil divide without float - padded_len = num_chunks * PREFILL_SEQ_LEN # Convert to a multiple of prompt_len - - model.to(torch.float32) - # Recompute reference logits after dtype conversion so we compare - # transformed fp32 path against fp32 reference. - with torch.no_grad(): - out = model(**inputs) - - qeff_model = QEFFAutoModelForCausalLM(model, qaic_config=qaic_config) - qeff_model.transform(ctx_len=CTX_LEN, seq_len=PREFILL_SEQ_LEN, bs=1, num_devices=1, qaic_config=qaic_config) - qeff_model.prefill(enable=True) - - inputs = tokenizer(prompt, return_tensors="np", padding="max_length", max_length=padded_len) - inputs["position_ids"] = np.where(inputs.pop("attention_mask"), np.arange(padded_len), -1) - inputs.pop("token_type_ids", None) - inputs = {k: torch.from_numpy(v) for k, v in inputs.items()} - - pad_shape_k = ( - 1, - model.config.num_attention_heads, - CTX_LEN, - model.config.qk_nope_head_dim + model.config.qk_rope_head_dim, - ) - pad_shape_v = (1, model.config.num_attention_heads, CTX_LEN, model.config.v_head_dim) - - num_heads = model.model.layers[0].self_attn.kv_a_proj_with_mqa.weight.shape[0] // ( - model.config.kv_lora_rank + model.config.qk_rope_head_dim - ) - pad_shape_ckv = (1, num_heads, CTX_LEN, model.config.kv_lora_rank) - pad_shape_k_pe = (1, num_heads, CTX_LEN, model.config.qk_rope_head_dim) - - past_key_values = [] - compressed_kvs = [] - - for i in range(model.config.num_hidden_layers): - past_key = torch.zeros((pad_shape_k), dtype=torch.float32) - past_value = torch.zeros((pad_shape_v), dtype=torch.float32) - pkv = (past_key, past_value) - past_key_values.append(pkv) - - ckv = torch.zeros((pad_shape_ckv), dtype=torch.float32) - k_pe = torch.zeros((pad_shape_k_pe), dtype=torch.float32) - x = (ckv, k_pe) - compressed_kvs.append(x) - - cache_compressed = mla_absorption.get("cache_compressed", False) - if cache_compressed: - inputs["compressed_kvs"] = compressed_kvs - else: - inputs["past_key_values"] = past_key_values - setattr(qeff_model.model.model, "num_q_blocks_ffn", PREFILL_SEQ_LEN // 256) - prefill_qeff_out = qeff_model.model(**inputs) - - x = prefill_qeff_out.logits - y = out.logits[:, -1, :] - assert (x - y).abs().max() < 2e-4 - -''' - # Tokenize once; reuse for both reference and qeff model - raw_inputs = tokenizer(prompt, return_tensors="np", padding=True) - padded_len = raw_inputs["input_ids"].shape[1] - num_chunks = -(padded_len // -PREFILL_SEQ_LEN) # ceil divide without float - padded_len = num_chunks * PREFILL_SEQ_LEN # Convert to a multiple of prompt_len - - raw_inputs = tokenizer(prompt, return_tensors="np", padding="max_length", max_length=padded_len) - raw_inputs["position_ids"] = np.where(raw_inputs.pop("attention_mask"), np.arange(padded_len), -1) - raw_inputs.pop("token_type_ids", None) - - inputs = {k: torch.from_numpy(v).to(model.device) for k, v in raw_inputs.items()} - - config = qeff_model.model.config - - inputs = {k: torch.from_numpy(v) for k, v in raw_inputs.items()} - past_key_values = [] - - pad_shape_k = ( - 1, - model.config.num_attention_heads, - CTX_LEN, - model.config.qk_nope_head_dim + model.config.qk_rope_head_dim, - ) - pad_shape_v = (1, model.config.num_attention_heads, CTX_LEN, model.config.v_head_dim) - - num_heads = model.model.layers[0].self_attn.kv_a_proj_with_mqa.weight.shape[0] // ( - model.config.kv_lora_rank + model.config.qk_rope_head_dim - ) - pad_shape_ckv = (1, num_heads, CTX_LEN, model.config.kv_lora_rank) - pad_shape_k_pe = (1, num_heads, CTX_LEN, model.config.qk_rope_head_dim) - - past_key_values = [] - compressed_kvs = [] - - for i in range(model.config.num_hidden_layers): - past_key = torch.zeros((pad_shape_k), dtype=torch.float32) - past_value = torch.zeros((pad_shape_v), dtype=torch.float32) - pkv = (past_key, past_value) - past_key_values.append(pkv) - - ckv = torch.zeros((pad_shape_ckv), dtype=torch.float32) - k_pe = torch.zeros((pad_shape_k_pe), dtype=torch.float32) - x = (ckv, k_pe) - compressed_kvs.append(x) - cache_compressed = mla_absorption.get("cache_compressed", False) - if cache_compressed: - inputs["compressed_kvs"] = compressed_kvs - else: - inputs["past_key_values"] = past_key_values - - prefill_qpc_path = qeff_model.compile( - prefill_seq_len=PREFILL_SEQ_LEN, - ctx_len=CTX_LEN, - num_cores=16, - mxfp6_matmul=False, - mxint8_kv_cache=False, - num_devices=1, - mos=1, - aic_enable_depth_first=True, - num_speculative_tokens=None, - prefill_only=True, - qaic_config=qaic_config, - ) - - from QEfficient.generation.cloud_infer import QAICInferenceSession - - prefill_session = QAICInferenceSession(prefill_qpc_path) - logits_out_placeholder = np.zeros((1, 1, config.vocab_size), dtype=np.float32) - prefill_session.set_buffers({"logits": logits_out_placeholder}) - if cache_compressed: - inputs.pop("compressed_kvs") - else: - inputs.pop("past_key_values") - inputs = {k: v.detach().numpy() for k, v in inputs.items()} - qpc_out = prefill_session.run(inputs) - assert (torch.from_numpy(qpc_out["logits"]) - prefill_qeff_out.logits).abs().max() < 5e-2 -''' \ No newline at end of file diff --git a/examples/kimi_k2/verify_kimi_k25_vision_parity.py b/examples/kimi_k2/verify_kimi_k25_vision_parity.py deleted file mode 100644 index 034ea2ca72..0000000000 --- a/examples/kimi_k2/verify_kimi_k25_vision_parity.py +++ /dev/null @@ -1,271 +0,0 @@ -# ----------------------------------------------------------------------------- -# -# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. -# SPDX-License-Identifier: BSD-3-Clause -# -# ----------------------------------------------------------------------------- - -import argparse -import copy -import shutil -import subprocess -import tempfile -from io import BytesIO -from pathlib import Path - -import numpy as np -import onnxruntime as ort -import requests -import torch -from export_kimi_k25_vision import ( - LOADED_EXPERT_IDS, - MODEL_PATH, - NUM_EXPERTS_PER_TOKEN, - NUM_TEXT_LAYERS, - NUM_VISION_LAYERS, - _ensure_torch_fx_import_compatibility, - _load_layer_subset_model, - _patch_deepseek_init_weights_compat, - _patch_kimi_tie_weights_compat, - _prepare_config, -) -from PIL import Image -from transformers.dynamic_module_utils import get_class_from_dynamic_module - -from QEfficient.generation.cloud_infer import QAICInferenceSession -from QEfficient.transformers.models.modeling_auto import QEffVisionEncoderForTextImageToTextModel - - -def parse_args(): - parser = argparse.ArgumentParser( - description="Verify Kimi K2.5 vision parity: HF PyTorch vs QEff PyTorch vs ONNXRuntime." - ) - parser.add_argument("--model-path", type=Path, default=MODEL_PATH) - parser.add_argument("--num-vision-layers", type=int, default=NUM_VISION_LAYERS) - parser.add_argument("--num-text-layers", type=int, default=NUM_TEXT_LAYERS) - parser.add_argument("--expert-ids", type=str, default=",".join(str(x) for x in LOADED_EXPERT_IDS)) - parser.add_argument("--num-experts-per-token", type=int, default=NUM_EXPERTS_PER_TOKEN) - parser.add_argument( - "--image-url", - type=str, - default="https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png", - ) - parser.add_argument("--prompt", type=str, default="Tell me about yourself.") - parser.add_argument("--atol", type=float, default=1e-3) - parser.add_argument("--rtol", type=float, default=1e-3) - parser.add_argument( - "--keep-onnx", - action="store_true", - help="Keep exported ONNX under ./aisand/onnx_export_verify instead of temp directory.", - ) - parser.add_argument( - "--run-qpc", action="store_true", help="Compile and execute QPC, then compare QPC output vs HF." - ) - parser.add_argument( - "--qaic-compile-bin", - type=str, - default="/opt/qti-aic/exec/qaic-compile", - help="Path to qaic-compile binary.", - ) - parser.add_argument( - "--keep-qpc", - action="store_true", - help="Keep compiled QPC under --qpc-dir (used only with --run-qpc).", - ) - parser.add_argument( - "--qpc-dir", - type=Path, - default=Path("aisand/qpc_verify"), - help="Output directory for QPC when --keep-qpc is set.", - ) - parser.add_argument("--aic-num-cores", type=int, default=16, help="Number of QAIC cores for qpc compile.") - parser.add_argument("--mos", type=int, default=1, help="-mos value for qpc compile.") - parser.add_argument( - "--no-convert-to-fp16", - action="store_true", - help="Disable -convert-to-fp16 during qpc compile (defaults to enabled).", - ) - return parser.parse_args() - - -def _parse_expert_ids(value: str): - return tuple(int(x.strip()) for x in value.split(",") if x.strip()) - - -def _stats(a: torch.Tensor, b: torch.Tensor, atol: float, rtol: float): - diff = (a - b).abs() - return { - "max_abs": float(diff.max().item()), - "mean_abs": float(diff.mean().item()), - "rmse": float(torch.sqrt(torch.mean((a - b) ** 2)).item()), - "allclose": bool(torch.allclose(a, b, atol=atol, rtol=rtol)), - } - - -def _hf_image_embeds(model, pixel_values, grid_thws): - target_dtype = model.vision_tower.patch_embed.proj.weight.dtype - pixel_values = pixel_values.to(target_dtype) - image_features = model.vision_tower(pixel_values, grid_thws) - if model.mm_projector: - image_features = model.mm_projector(image_features) - return image_features[0] if isinstance(image_features, list) else image_features - - -def _run_onnx(qeff_vision_wrapper, qeff_base_model, pixel_values, h_shape, w_shape, keep_onnx: bool): - output_names = qeff_base_model.get_output_names(kv_offload=True)["vision"] - dynamic_axes = qeff_base_model.get_onnx_dynamic_axes(kv_offload=True)["vision"] - - if keep_onnx: - export_dir = Path("aisand/onnx_export_verify") - export_dir.mkdir(parents=True, exist_ok=True) - else: - export_dir = Path(tempfile.mkdtemp(prefix="kimi_k25_verify_onnx_")) - - onnx_path = qeff_vision_wrapper.export( - {"pixel_values": pixel_values, "h_shape": h_shape, "w_shape": w_shape}, - output_names, - dynamic_axes, - export_dir=export_dir, - offload_pt_weights=False, - ) - - session = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"]) - ort_inputs = { - "pixel_values": pixel_values.detach().cpu().numpy().astype(np.float32), - "h_shape": h_shape.detach().cpu().numpy().astype(np.int64), - "w_shape": w_shape.detach().cpu().numpy().astype(np.int64), - } - ort_out = session.run(None, ort_inputs)[0] - return torch.from_numpy(ort_out).float(), Path(onnx_path) - - -def _run_qpc(onnx_path: Path, pixel_values, h_shape, w_shape, args): - if args.keep_qpc: - qpc_dir = args.qpc_dir - if qpc_dir.exists(): - shutil.rmtree(qpc_dir) - qpc_dir.parent.mkdir(parents=True, exist_ok=True) - else: - qpc_dir = Path(tempfile.mkdtemp(prefix="kimi_k25_verify_qpc_")) - shutil.rmtree(qpc_dir) - - compile_cmd = [ - args.qaic_compile_bin, - f"-m={onnx_path}", - f"-onnx-define-symbol=num_patches,{pixel_values.shape[0]}", - f"-onnx-define-symbol=h,{h_shape.shape[0]}", - f"-onnx-define-symbol=w,{w_shape.shape[0]}", - f"-aic-num-cores={args.aic_num_cores}", - f"-mos={args.mos}", - f"-aic-binary-dir={qpc_dir}", - ] - if not args.no_convert_to_fp16: - compile_cmd.insert(1, "-convert-to-fp16") - compile_result = subprocess.run(compile_cmd, capture_output=True, text=True) - if compile_result.returncode != 0: - raise RuntimeError( - "QPC compile failed.\n" - f"Command: {' '.join(compile_cmd)}\n" - f"stdout:\n{compile_result.stdout}\n" - f"stderr:\n{compile_result.stderr}" - ) - - session = QAICInferenceSession(str(qpc_dir)) - qpc_inputs = {} - for name in session.input_names: - if name == "pixel_values": - qpc_inputs[name] = pixel_values.detach().cpu().numpy().astype(np.float32) - elif name == "h_shape": - qpc_inputs[name] = h_shape.detach().cpu().numpy().astype(np.int64) - elif name == "w_shape": - qpc_inputs[name] = w_shape.detach().cpu().numpy().astype(np.int64) - else: - raise RuntimeError(f"Unexpected QPC input name: {name}") - - qpc_outputs = session.run(qpc_inputs) - output_name = session.output_names[0] - qpc_out = qpc_outputs[output_name] - return torch.from_numpy(qpc_out).float(), qpc_dir - - -if __name__ == "__main__": - args = parse_args() - expert_ids = _parse_expert_ids(args.expert_ids) - - torch.manual_seed(123) - torch.set_grad_enabled(False) - - _ensure_torch_fx_import_compatibility() - config = _prepare_config(args.model_path) - kimi_cls = get_class_from_dynamic_module("modeling_kimi_k25.KimiK25ForConditionalGeneration", str(args.model_path)) - _patch_kimi_tie_weights_compat(kimi_cls) - _patch_deepseek_init_weights_compat(kimi_cls) - - hf_model, _, processor = _load_layer_subset_model( - model_path=args.model_path, - kimi_cls=kimi_cls, - config=config, - num_vision_layers=args.num_vision_layers, - num_text_layers=args.num_text_layers, - loaded_expert_ids=expert_ids, - num_experts_per_tok=args.num_experts_per_token, - dtype=torch.float32, - ) - - # Required for HW compatibility and parity in this flow. - hf_model.vision_tower.patch_embed.pos_emb.interpolation_mode = "bilinear" - hf_model.eval() - - qeff_base_model = copy.deepcopy(hf_model).eval() - qeff_vision_wrapper = QEffVisionEncoderForTextImageToTextModel(qeff_base_model) - qeff_encoder = qeff_vision_wrapper.model.eval() - - image = Image.open(BytesIO(requests.get(args.image_url, timeout=30).content)).convert("RGB") - messages = [ - { - "role": "user", - "content": [ - {"type": "image_url", "image_url": image}, - {"type": "text", "text": args.prompt}, - ], - } - ] - inputs = processor(messages=messages, add_generation_prompt=True, tokenize=False, return_tensors="pt") - - pixel_values = inputs["pixel_values"] - grid_thws = inputs["grid_thws"] - h_shape = torch.ones(int(grid_thws[0, 1].item()), dtype=torch.int64) - w_shape = torch.ones(int(grid_thws[0, 2].item()), dtype=torch.int64) - - hf_out = _hf_image_embeds(hf_model, pixel_values, grid_thws).detach().cpu().float() - qeff_out = qeff_encoder(pixel_values, h_shape, w_shape).detach().cpu().float() - ort_out, onnx_path = _run_onnx(qeff_vision_wrapper, qeff_base_model, pixel_values, h_shape, w_shape, args.keep_onnx) - - hf_vs_qeff = _stats(hf_out, qeff_out, args.atol, args.rtol) - hf_vs_ort = _stats(hf_out, ort_out, args.atol, args.rtol) - qeff_vs_ort = _stats(qeff_out, ort_out, args.atol, args.rtol) - - print(f"ONNX path: {onnx_path}") - print("hf_vs_qeff", hf_vs_qeff) - print("hf_vs_onnxrt", hf_vs_ort) - print("qeff_vs_onnxrt", qeff_vs_ort) - - all_ok = hf_vs_qeff["allclose"] and hf_vs_ort["allclose"] and qeff_vs_ort["allclose"] - - if args.run_qpc: - atol = 1e-2 - rtol = 1e-2 - qpc_out, qpc_dir = _run_qpc(onnx_path, pixel_values, h_shape, w_shape, args) - hf_vs_qpc = _stats(hf_out, qpc_out, atol, rtol) - qeff_vs_qpc = _stats(qeff_out, qpc_out, atol, rtol) - ort_vs_qpc = _stats(ort_out, qpc_out, atol, rtol) - print(f"QPC dir: {qpc_dir}") - print("hf_vs_qpc", hf_vs_qpc) - print("qeff_vs_qpc", qeff_vs_qpc) - print("onnxrt_vs_qpc", ort_vs_qpc) - all_ok = all_ok and hf_vs_qpc["allclose"] and qeff_vs_qpc["allclose"] and ort_vs_qpc["allclose"] - - if not all_ok: - raise SystemExit(f"Parity check failed for atol={atol}, rtol={rtol}. See metrics above for details.") - - print("Parity check passed.") From a6ef204067bea37ffffbcc848f74dc1acb01d387 Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Tue, 30 Jun 2026 23:19:50 +0530 Subject: [PATCH 11/33] cleanup Signed-off-by: Mamta Singh --- .../models/kimi_k25/modeling_kimi_k25.py | 111 +++++------------- QEfficient/utils/constants.py | 7 +- examples/kimi_k2/export_kimi_k25_vision.py | 4 - 3 files changed, 33 insertions(+), 89 deletions(-) diff --git a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py index f2df4e272c..5d4516b083 100644 --- a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py +++ b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py @@ -6,17 +6,12 @@ # ---------------------------------------------------------------------------- import math -import sys as _sys -from io import BytesIO -from pathlib import Path from typing import List, Optional, Tuple, Type, Union -import requests import torch import torch.nn as nn import torch.nn.functional as F -from PIL import Image -from transformers import AutoProcessor, activations +from transformers import activations try: from transformers.activations import PytorchGELUTanh @@ -366,8 +361,6 @@ def forward(self, pixel_values: torch.Tensor, h_shape: torch.Tensor, w_shape: to Returns: image_embeds: Projected image embeddings as a single tensor. """ - _kimi_module = _sys.modules[type(self).__module__] - _get_rope_shape = _kimi_module.get_rope_shape h_shape = h_shape.to(pixel_values.device) w_shape = w_shape.to(pixel_values.device) @@ -387,7 +380,7 @@ def forward(self, pixel_values: torch.Tensor, h_shape: torch.Tensor, w_shape: to # Positional embedding (single image, t=1) pos_emb_module = self.model.vision_tower.patch_embed.pos_emb - pos_emb_2d = _get_rope_shape( + pos_emb_2d = get_rope_shape( pos_emb_module.weight, interpolation_mode=pos_emb_module.interpolation_mode, shape=(h, w), @@ -468,12 +461,7 @@ def forward( use_cache: Optional[bool] = None, **kwargs, ) -> Tuple: - if inputs_embeds is None: - inputs_embeds = self.model.get_input_embeddings()(input_ids) - - if image_idx is None: - image_idx = torch.zeros((inputs_embeds.shape[0], 1), dtype=torch.int64, device=inputs_embeds.device) - + inputs_embeds = self.model.get_input_embeddings()(input_ids) image_embeds_for_state = None if image_embeds is not None: if image_embeds.dim() == 3: @@ -487,17 +475,7 @@ def forward( if image_embeds_for_state is not None and input_ids is not None: if attention_mask is None: - # ONNX prefill/decode callers may omit attention_mask but still - # provide position_ids with `-1` on padded slots. - if position_ids is not None: - attention_mask = (position_ids >= 0).to(dtype=torch.long, device=input_ids.device) - else: - pad_token_id = getattr(self.config, "pad_token_id", None) - if pad_token_id is None: - attention_mask = torch.ones_like(input_ids, dtype=torch.long, device=input_ids.device) - else: - attention_mask = (input_ids != pad_token_id).to(dtype=torch.long, device=input_ids.device) - + attention_mask = (position_ids >= 0).to(dtype=torch.long, device=input_ids.device) selected = input_ids == self.config.media_placeholder_token_id selected_any = selected.any(dim=1, keepdim=True) selected_image_tokens = selected.to(torch.int64).sum(dim=1, keepdim=True) @@ -544,13 +522,10 @@ def forward( position_ids = attention_mask.long().cumsum(-1) - 1 position_ids = torch.where(attention_mask > 0, position_ids, torch.full_like(position_ids, -1)) elif position_ids is not None and attention_mask is not None: - # Preserve `-1` on padded slots so cache updates skip them without - # creating dynamic Slice nodes unsupported by compile. position_ids = torch.where(attention_mask > 0, position_ids, torch.full_like(position_ids, -1)) outputs = self.model.language_model( inputs_embeds=inputs_embeds, - attention_mask=attention_mask, position_ids=position_ids, compressed_kvs=compressed_kvs, past_key_values=past_key_values, @@ -559,8 +534,6 @@ def forward( use_cache=True if use_cache is None else use_cache, **kwargs, ) - - # QEff Deepseek language_model.forward already returns final logits. logits = outputs[0].float() mla_absorption = getattr(self.model, "mla_absorption", None) @@ -695,8 +668,6 @@ def forward( input_ids: torch.LongTensor | None = None, pixel_values: torch.FloatTensor | list[torch.FloatTensor] | None = None, grid_thws: torch.Tensor | None = None, - # h_shape: torch.Tensor | None = None, - # w_shape: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, @@ -863,45 +834,25 @@ def get_dummy_inputs( bs: int = constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE fbs: int = constants.ONNX_EXPORT_EXAMPLE_FBS - model_path = Path( - "/home/huggingface_hub/models--moonshotai--Kimi-K2.5/snapshots/4d01dfe0332d63057c186e0b262165819efb6611" + inputs_shapes = {} + inputs_shapes["pixel_values"] = ( + constants.KIMI_NUM_PATCHES, + constants.ONNX_EXPORT_IMAGE_DEPTH, + constants.KIMI_PATCH_SIZE, + constants.KIMI_PATCH_SIZE, ) - processor = AutoProcessor.from_pretrained(str(model_path), trust_remote_code=True) - image_url = "https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png" - image = Image.open(BytesIO(requests.get(image_url).content)).convert("RGB") - messages = [ - { - "role": "user", - "content": [ - {"type": "image_url", "image_url": image}, - {"type": "text", "text": "Describe this image."}, - ], - } - ] - inputs = processor( - messages=messages, - add_generation_prompt=True, - tokenize=False, - return_tensors="pt", + inputs_shapes["image_embeds"] = ( + constants.KIMI_NUM_IMAGE_TOKENS, + self.language_model.config.hidden_size, ) - - # Read grid dimensions before moving tensors to model device. During - # export flows that materialize modules on `meta`, calling `.item()` on - # meta tensors raises `RuntimeError`. - grid_thws_val = inputs["grid_thws"] - h_val = int(grid_thws_val[0, 1].item()) - w_val = int(grid_thws_val[0, 2].item()) - - inputs = {k: v.to(self.language_model.device) if hasattr(v, "to") else v for k, v in inputs.items()} - - # Build h_shape and w_shape from grid_thws (single image, t=1) - h_shape_tensor = torch.ones(h_val, dtype=torch.int64, device=self.language_model.device) - w_shape_tensor = torch.ones(w_val, dtype=torch.int64, device=self.language_model.device) + inputs_shapes["image_idx"] = (1, 1) + inputs_shapes["h"] = constants.KIMI_IMAGE_HEIGHT + inputs_shapes["w"] = constants.KIMI_IMAGE_WIDTH vision_inputs = { - "pixel_values": inputs["pixel_values"], - "h_shape": h_shape_tensor, - "w_shape": w_shape_tensor, + "pixel_values": torch.zeros((inputs_shapes["pixel_values"]), dtype=self.config.torch_dtype), + "h_shape": torch.zeros((inputs_shapes["h"]), dtype=torch.int64), + "w_shape": torch.zeros((inputs_shapes["w"]), dtype=torch.int64), } # Ensure ONNX export traces the multimodal merge path in decoder wrapper. @@ -909,6 +860,7 @@ def get_dummy_inputs( # branch and produce a decoder graph that ignores `image_embeds`. input_ids = torch.ones((bs, prefill_seq_len), dtype=torch.int64) input_ids[:, 0] = self.config.media_placeholder_token_id + lang_inputs = { "input_ids": input_ids, "position_ids": torch.arange(prefill_seq_len, dtype=torch.int64).view(1, prefill_seq_len).repeat(bs, 1), @@ -923,7 +875,7 @@ def get_dummy_inputs( pkv_cache = self.language_model.get_dummy_pkv_cache( config=self.language_model.config, batch_size=fbs if continuous_batching else bs, - seq_len=1024, + seq_len=constants.ONNX_EXPORT_CTX_LEN, ) if cache_compressed: @@ -945,12 +897,6 @@ def get_dummy_inputs( torch.zeros(pkv_cache[0][1].shape, dtype=self.language_model.config.torch_dtype) ) - inputs_shapes = {} - # Decoder expects image embeddings indexed by token position. Keep this as - # [num_image_tokens, hidden_size] so `num_image_tokens` maps to axis 0. - inputs_shapes["image_embeds"] = (600, 7168) - inputs_shapes["pixel_values"] = (2400, 3, 14, 14) - inputs_shapes["image_idx"] = (1, 1) lang_inputs["image_embeds"] = torch.zeros( inputs_shapes["image_embeds"], dtype=self.language_model.config.torch_dtype, @@ -960,9 +906,6 @@ def get_dummy_inputs( if continuous_batching: lang_inputs["batch_index"] = torch.arange(bs).view(bs, 1) - if comp_ctx_lengths is not None: - lang_inputs["comp_ctx_lengths"] = torch.randint(0, 100, (40,), dtype=torch.int64) - inputs = {} if kv_offload: inputs["vision"] = vision_inputs @@ -1035,13 +978,13 @@ def get_specializations( w = compiler_options.pop("w", None) num_image_tokens = compiler_options.pop("num_image_tokens", None) - prefill_seq_len = prefill_seq_len if prefill_seq_len else 32 # constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN - ctx_len = ctx_len if ctx_len else 1024 # constants.ONNX_EXPORT_EXAMPLE_CTX_LEN + prefill_seq_len = prefill_seq_len if prefill_seq_len else constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN + ctx_len = ctx_len if ctx_len else constants.ONNX_EXPORT_CTX_LEN - num_patches = num_patches if num_patches is not None else 2400 - h = h if h is not None else 30 - w = w if w is not None else 80 - num_image_tokens = num_image_tokens if num_image_tokens is not None else 600 + num_patches = num_patches if num_patches is not None else constants.KIMI_NUM_PATCHES + h = h if h is not None else constants.KIMI_IMAGE_HEIGHT + w = w if w is not None else constants.KIMI_IMAGE_WIDTH + num_image_tokens = num_image_tokens if num_image_tokens is not None else constants.KIMI_NUM_IMAGE_TOKENS vision = [ { diff --git a/QEfficient/utils/constants.py b/QEfficient/utils/constants.py index 4e1032a40f..bf68f58df0 100644 --- a/QEfficient/utils/constants.py +++ b/QEfficient/utils/constants.py @@ -172,7 +172,12 @@ def get_default_aic_hw_version() -> str: LLAMA4_ATTENTION_CHUNK_SIZE = 8192 LLAMA4_MAX_POSITION_EMBEDDINGS = 65536 -# DeepSeek Kimi-k2 Constant +# DeepSeek Kimi-k2.5 Constants +KIMI_NUM_PATCHES = 2400 +KIMI_PATCH_SIZE = 14 +KIMI_NUM_IMAGE_TOKENS = 600 +KIMI_IMAGE_HEIGHT = 30 +KIMI_IMAGE_WIDTH = 80 MAX_POSITION_EMBEDDINGS = 32768 FP16_BYTES = 2 DEFAULT_NUM_HEADS = 64 diff --git a/examples/kimi_k2/export_kimi_k25_vision.py b/examples/kimi_k2/export_kimi_k25_vision.py index 5412d6a725..54d24e663d 100644 --- a/examples/kimi_k2/export_kimi_k25_vision.py +++ b/examples/kimi_k2/export_kimi_k25_vision.py @@ -355,10 +355,6 @@ def main(): _patch_kimi_tie_weights_compat(kimi_cls) _patch_deepseek_init_weights_compat(kimi_cls) - kimi_module = sys.modules[kimi_cls.__module__] - if hasattr(kimi_module, "MoonViT3dEncoder"): - kimi_module.MoonViT3dEncoder.use_deterministic_attn = False - model_kwargs = { "config": config, "trust_remote_code": True, From 1b25282725aba8710a4c3869466d07dcaab75cc8 Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Mon, 6 Jul 2026 13:19:09 +0530 Subject: [PATCH 12/33] tests and multi image support Signed-off-by: Mamta Singh --- .../models/kimi_k25/modeling_kimi_k25.py | 94 +++++--- QEfficient/utils/test_utils.py | 3 + tests/configs/image_text_model_configs.json | 51 +++++ .../test_image_text_to_text_models.py | 205 +++++++++++++++++- 4 files changed, 316 insertions(+), 37 deletions(-) diff --git a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py index 5d4516b083..7076850235 100644 --- a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py +++ b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py @@ -28,10 +28,25 @@ def eager_attention_forward(q, k, v, **kwargs): + q_cu_seqlens = kwargs.get("q_cu_seqlens") + seq_length = q.shape[0] + if q_cu_seqlens is not None: + attention_mask = torch.zeros([1, seq_length, seq_length], device=q.device, dtype=torch.bool) + for idx in range(1, len(q_cu_seqlens)): + attention_mask[ + ..., + q_cu_seqlens[idx - 1] : q_cu_seqlens[idx], + q_cu_seqlens[idx - 1] : q_cu_seqlens[idx], + ] = True + else: + attention_mask = None + q = q.transpose(0, 1) # (num_heads, seq_len, head_dim) k = k.transpose(0, 1) v = v.transpose(0, 1) attn_weight = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1]) + if attention_mask is not None: + attn_weight = attn_weight + attention_mask attn_weight = torch.softmax(attn_weight, dim=-1, dtype=torch.float32).to(q.dtype) attn_out = attn_weight @ v attn_out = attn_out.transpose(0, 1) # (seq_len, num_heads, head_dim) @@ -344,8 +359,7 @@ def get_submodules_for_export(self) -> Type[nn.Module]: This method should return the *class object* (not an instance). Downstream code can use this to find/build subfunctions for repeated blocks. """ - return {self.model.vision_model.model.layers[0].__class__} - # return {self.model.layers[0].__class__} + return {self.model.vision_tower.encoder.blocks[0].__class__} def forward(self, pixel_values: torch.Tensor, h_shape: torch.Tensor, w_shape: torch.Tensor) -> torch.Tensor: """ @@ -359,7 +373,8 @@ def forward(self, pixel_values: torch.Tensor, h_shape: torch.Tensor, w_shape: to w_shape: int64 ones tensor of shape (w,) encoding number of patch columns. Returns: - image_embeds: Projected image embeddings as a single tensor. + image_embeds: Projected image embeddings as a single tensor. For multiple + same-sized images, embeddings are concatenated in input order. """ h_shape = h_shape.to(pixel_values.device) @@ -371,58 +386,60 @@ def forward(self, pixel_values: torch.Tensor, h_shape: torch.Tensor, w_shape: to h = h_shape.shape[0] w = w_shape.shape[0] + num_tokens_per_image = h * w + if pixel_values.shape[0] % num_tokens_per_image != 0: + raise ValueError( + f"pixel_values first dimension ({pixel_values.shape[0]}) must be divisible by h*w ({num_tokens_per_image})." + ) + num_images = pixel_values.shape[0] // num_tokens_per_image target_dtype = self.model.vision_tower.patch_embed.proj.weight.dtype pixel_values = pixel_values.to(target_dtype) - # --- Patch embedding --- - x = self.model.vision_tower.patch_embed.proj(pixel_values).view(pixel_values.size(0), -1) + hidden_states = self.model.vision_tower.patch_embed.proj(pixel_values).view(pixel_values.size(0), -1) - # Positional embedding (single image, t=1) pos_emb_module = self.model.vision_tower.patch_embed.pos_emb pos_emb_2d = get_rope_shape( pos_emb_module.weight, interpolation_mode=pos_emb_module.interpolation_mode, shape=(h, w), ) - x = x + pos_emb_2d - - # --- Encoder --- - # For single image with t=1: cu_seqlens = [0, h*w] - num_tokens = h * w - cu_seqlens = torch.zeros(2, dtype=torch.int32, device=x.device) - cu_seqlens[1] = num_tokens - max_seqlen = num_tokens + hidden_states = hidden_states + pos_emb_2d.repeat(num_images, 1) - # RoPE frequencies for single image (stacked real/imag to avoid complex tensors in ONNX) - # Shape: (2, num_tokens, dim//2) where [0]=cos, [1]=sin rope_2d = self.model.vision_tower.encoder.rope_2d - rope_2d._ensure_precomputed_freqs(x.device) - freqs_cos = rope_2d.freqs_cos[:h, :w].reshape(-1, rope_2d.dim // 2) - freqs_sin = rope_2d.freqs_sin[:h, :w].reshape(-1, rope_2d.dim // 2) + rope_2d._ensure_precomputed_freqs(pixel_values.device) + freqs_cos = rope_2d.freqs_cos[:h, :w].reshape(-1, rope_2d.dim // 2).repeat(num_images, 1) + freqs_sin = rope_2d.freqs_sin[:h, :w].reshape(-1, rope_2d.dim // 2).repeat(num_images, 1) freqs_cis = torch.stack([freqs_cos, freqs_sin], dim=0) + if num_images == 1: + cu_seqlens = None + else: + lengths = torch.full( + (num_images + 1,), num_tokens_per_image, dtype=torch.int64, device=hidden_states.device + ) + lengths[0] = 0 + cu_seqlens = lengths.cumsum(dim=0, dtype=torch.int64) + max_seqlen = num_tokens_per_image + encoder_dtype = self.model.vision_tower.encoder.blocks[0].wqkv.weight.dtype - x = x.to(encoder_dtype) + hidden_states = hidden_states.to(encoder_dtype) freqs_cis = freqs_cis.to(encoder_dtype) - for block in self.model.vision_tower.encoder.blocks: - x = block(x, cu_seqlens, max_seqlen, rope_freqs_cis=freqs_cis) + hidden_states = block(hidden_states, cu_seqlens, max_seqlen, rope_freqs_cis=freqs_cis) final_ln_dtype = self.model.vision_tower.encoder.final_layernorm.weight.dtype - x = self.model.vision_tower.encoder.final_layernorm(x.to(final_ln_dtype)) + hidden_states = self.model.vision_tower.encoder.final_layernorm(hidden_states.to(final_ln_dtype)) - # --- tpool_patch_merger (single image, t=1) --- merge_kernel_size = self.model.vision_tower.merge_kernel_size kernel_height, kernel_width = merge_kernel_size - d_model = x.size(-1) new_height = h // kernel_height new_width = w // kernel_width - reshaped = x.view(1, new_height, kernel_height, new_width, kernel_width, d_model) - reshaped = reshaped.permute(0, 1, 3, 2, 4, 5).contiguous().mean(dim=0) - merged = reshaped.view(new_height * new_width, kernel_height * kernel_width, -1) + d_model = hidden_states.size(-1) + reshaped = hidden_states.view(num_images, new_height, kernel_height, new_width, kernel_width, d_model) + reshaped = reshaped.permute(0, 1, 3, 2, 4, 5).contiguous() + merged = reshaped.view(num_images * new_height * new_width, kernel_height * kernel_width, -1) - # --- mm_projector (PatchMergerMLP on single tensor) --- pre_norm_dtype = self.model.mm_projector.pre_norm.weight.dtype merged = merged.to(pre_norm_dtype) image_embeds = self.model.mm_projector.proj(self.model.mm_projector.pre_norm(merged).view(merged.shape[0], -1)) @@ -475,12 +492,21 @@ def forward( if image_embeds_for_state is not None and input_ids is not None: if attention_mask is None: - attention_mask = (position_ids >= 0).to(dtype=torch.long, device=input_ids.device) + if position_ids is None: + attention_mask = torch.ones_like(input_ids, dtype=torch.long, device=input_ids.device) + else: + attention_mask = (position_ids >= 0).to(dtype=torch.long, device=input_ids.device) + if image_idx is None: + image_idx = torch.zeros((input_ids.shape[0], 1), dtype=torch.int64, device=input_ids.device) selected = input_ids == self.config.media_placeholder_token_id selected_any = selected.any(dim=1, keepdim=True) selected_image_tokens = selected.to(torch.int64).sum(dim=1, keepdim=True) - image_features = [image_embeds_for_state] + num_selected_images = int(selected_image_tokens.max().item()) + if num_selected_images > 0 and image_embeds_for_state.shape[0] % num_selected_images == 0: + image_features = list(torch.chunk(image_embeds_for_state, num_selected_images, dim=0)) + else: + image_features = [image_embeds_for_state] inputs_embeds = inputs_embeds.to(image_embeds_for_state.dtype) merged_inputs_embeds, merged_attention_mask, _, merged_position_ids = ( @@ -821,7 +847,6 @@ def get_output_names(self, kv_offload: bool = False): def get_dummy_inputs( self, - comp_ctx_lengths: Optional[List[int]] = None, kv_offload: bool = False, continuous_batching: bool = False, **kwargs, @@ -855,10 +880,7 @@ def get_dummy_inputs( "w_shape": torch.zeros((inputs_shapes["w"]), dtype=torch.int64), } - # Ensure ONNX export traces the multimodal merge path in decoder wrapper. - # A dummy sequence without the media placeholder token can dead-code this - # branch and produce a decoder graph that ignores `image_embeds`. - input_ids = torch.ones((bs, prefill_seq_len), dtype=torch.int64) + input_ids = torch.zeros((bs, prefill_seq_len), dtype=torch.int64) input_ids[:, 0] = self.config.media_placeholder_token_id lang_inputs = { diff --git a/QEfficient/utils/test_utils.py b/QEfficient/utils/test_utils.py index 131ff59e26..e00e2d187b 100644 --- a/QEfficient/utils/test_utils.py +++ b/QEfficient/utils/test_utils.py @@ -98,6 +98,8 @@ def set_num_layers_vlm(config: AutoConfig, n_layer: int = -1): elif hasattr(config, "text_config"): config.text_config.num_hidden_layers = n_layer config.vision_config.num_hidden_layers = n_layer + if hasattr(config.vision_config, "vt_num_hidden_layers"): + config.vision_config.vt_num_hidden_layers = n_layer if hasattr(config.vision_config, "depth"): config.vision_config.depth = n_layer if hasattr(config.vision_config, "deepstack_visual_indexes"): @@ -490,6 +492,7 @@ class ModelConfig: "Qwen/Qwen3-VL-Reranker-8B", "Qwen/Qwen3.5-0.8B", "Qwen/Qwen3.6-35B-A3B", + "moonshotai/Kimi-K2.5", } EXTERNAL_MODELS = { diff --git a/tests/configs/image_text_model_configs.json b/tests/configs/image_text_model_configs.json index 85df559970..060615e0d5 100644 --- a/tests/configs/image_text_model_configs.json +++ b/tests/configs/image_text_model_configs.json @@ -608,6 +608,57 @@ "patch_size": 14 } } + }, + { + "model_name": "moonshotai/Kimi-K2.5", + "model_type": "kimi_k25", + "batch_size": 1, + "prompt_len": 32, + "ctx_len": 1024, + "img_size": null, + "img_url": "https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png", + "text_prompt": "Describe this image.", + "num_layers": 2, + "img_url_list": [ + "https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png", + "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/datasets/cat_style_layout.png" + ], + "text_prompt_list": [ + "Describe this image.", + "Can you describe the image in detail?" + ], + "full_batch_size": 2, + "additional_params": { + "text_config": { + "hidden_size": 64, + "intermediate_size": 128, + "kv_lora_rank": 16, + "max_position_embeddings": 1024, + "model_type": "kimi_k2", + "moe_intermediate_size": 32, + "n_group": 1, + "n_routed_experts": 4, + "n_shared_experts": 1, + "num_attention_heads": 4, + "num_experts_per_tok": 2, + "num_hidden_layers": 2, + "num_key_value_heads": 4, + "q_lora_rank": 32, + "qk_nope_head_dim": 8, + "qk_rope_head_dim": 8, + "topk_group": 1, + "v_head_dim": 16, + "vocab_size": 163840 + }, + "vision_config": { + "mm_hidden_size": 64, + "text_hidden_size": 64, + "vt_hidden_size": 64, + "vt_intermediate_size": 128, + "vt_num_attention_heads": 4, + "vt_num_hidden_layers": 2 + } + } } ], "image_text_subfunction_models": [ diff --git a/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py b/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py index 9b9e662e52..17f38da365 100644 --- a/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py +++ b/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py @@ -6,8 +6,10 @@ # ---------------------------------------------------------------------------- import copy +import inspect import json import os +import sys from io import BytesIO from typing import List, Optional @@ -22,6 +24,8 @@ GenerationConfig, TextStreamer, ) +from transformers.dynamic_module_utils import get_class_from_dynamic_module +from transformers.utils import import_utils as hf_import_utils from QEfficient import QEFFAutoModelForCausalLM, QEFFAutoModelForImageTextToText from QEfficient.utils._utils import create_json @@ -45,6 +49,150 @@ model_config_dict = {model["model_name"]: model for model in multimodal_models} NEW_GENERATION_TOKENS = 10 +KIMI_K25_MODEL_NAME = "moonshotai/Kimi-K2.5" +KIMI_K25_QAIC_CONFIG = {"mla_absorption": {"cache_compressed": True, "absorption": False, "online": False}} + + +def _is_kimi_k25(model_name: str) -> bool: + return model_name == KIMI_K25_MODEL_NAME + + +def _ensure_torch_fx_import_compatibility(): + if hasattr(hf_import_utils, "is_torch_fx_available"): + return + + def _is_torch_fx_available() -> bool: + if not hf_import_utils.is_torch_available(): + return False + try: + import torch.fx # noqa: F401 + + return True + except Exception: + return False + + hf_import_utils.is_torch_fx_available = _is_torch_fx_available + + +def _patch_kimi_k25_tie_weights_compat(kimi_cls): + tie_signature = inspect.signature(kimi_cls.tie_weights) + if tuple(tie_signature.parameters) != ("self",): + return + + def _tie_weights_compat(self, missing_keys=None, recompute_mapping=True): + lm_tie_weights = getattr(self.language_model, "tie_weights") + try: + return lm_tie_weights(missing_keys=missing_keys, recompute_mapping=recompute_mapping) + except TypeError: + return lm_tie_weights() + + kimi_cls.tie_weights = _tie_weights_compat + + +def _patch_kimi_k25_deepseek_init_weights_compat(kimi_cls): + module_prefix, _ = kimi_cls.__module__.rsplit(".", maxsplit=1) + deepseek_module = sys.modules.get(f"{module_prefix}.modeling_deepseek") + if deepseek_module is None or not hasattr(deepseek_module, "DeepseekV3PreTrainedModel"): + return + + deepseek_cls = deepseek_module.DeepseekV3PreTrainedModel + if getattr(deepseek_cls, "_qeff_test_init_weights_patched", False): + return + + def _init_weights_compat(self, module): + std = self.config.initializer_range + if isinstance(module, torch.nn.Linear): + if hasattr(module, "weight") and module.weight is not None: + module.weight.data.normal_(mean=0.0, std=std) + if hasattr(module, "bias") and module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, torch.nn.Embedding): + if hasattr(module, "weight") and module.weight is not None: + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + deepseek_cls._init_weights = _init_weights_compat + deepseek_cls._qeff_test_init_weights_patched = True + + +def _patch_kimi_k25_remote_code_compat(config): + _ensure_torch_fx_import_compatibility() + kimi_cls = get_class_from_dynamic_module("modeling_kimi_k25.KimiK25ForConditionalGeneration", config._name_or_path) + _patch_kimi_k25_tie_weights_compat(kimi_cls) + _patch_kimi_k25_deepseek_init_weights_compat(kimi_cls) + + +def _get_kimi_k25_test_config(model_name: str): + config = AutoConfig.from_pretrained(model_name, trust_remote_code=True) + config._attn_implementation = "eager" + config.torch_dtype = torch.float32 + config.dtype = torch.float32 + additional_params = model_config_dict[model_name]["additional_params"] + + for attr, value in additional_params["text_config"].items(): + setattr(config.text_config, attr, value) + config.text_config._attn_implementation = "eager" + config.text_config.torch_dtype = torch.float32 + config.text_config.dtype = torch.float32 + + for attr, value in additional_params["vision_config"].items(): + setattr(config.vision_config, attr, value) + config.vision_config._attn_implementation = "eager" + config.vision_config.torch_dtype = torch.float32 + config.vision_config.dtype = torch.float32 + + _patch_kimi_k25_remote_code_compat(config) + return config + + +def _get_kimi_k25_num_image_tokens(config, grid_thws): + merge_height, merge_width = config.vision_config.merge_kernel_size + return int(grid_thws[0, 1].item() // merge_height) * int(grid_thws[0, 2].item() // merge_width) + + +@torch.no_grad() +def _run_kimi_k25_hf_model_on_pytorch(model, processor, inputs, max_gen_len): + generated_ids = inputs["input_ids"] + attention_mask = inputs["attention_mask"] + pixel_values = inputs["pixel_values"] + grid_thws = inputs["grid_thws"] + new_tokens = [] + + eos_token_id = getattr(model.config, "eos_token_id", None) + if eos_token_id is None and hasattr(model.config, "text_config"): + eos_token_id = getattr(model.config.text_config, "eos_token_id", None) + + for _ in range(max_gen_len): + outputs = model( + input_ids=generated_ids, + attention_mask=attention_mask, + pixel_values=pixel_values, + grid_thws=grid_thws, + use_cache=False, + return_dict=True, + ) + logits = outputs[0] if isinstance(outputs, tuple) else outputs.logits + next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True) + new_tokens.append(next_token) + + generated_ids = torch.cat([generated_ids, next_token], dim=1) + attention_mask = torch.cat( + [ + attention_mask, + torch.ones((attention_mask.shape[0], 1), dtype=attention_mask.dtype, device=attention_mask.device), + ], + dim=1, + ) + + if eos_token_id is not None and torch.all(next_token == eos_token_id): + break + + output_tokens = torch.cat(new_tokens, dim=1).squeeze(0) + py_output = processor.tokenizer.decode(output_tokens.tolist()).strip() + print("Original HF Model Outputs (Torch CPU):") + print("Completion:", repr(py_output)) + return output_tokens def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100( @@ -70,6 +218,8 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100( pytorch_kv_tokens = None ort_tokens = None n_layer = num_hidden_layers + if _is_kimi_k25(model_name) and config is None: + config = _get_kimi_k25_test_config(model_name) if config is None: config = AutoConfig.from_pretrained( model_name, trust_remote_code=True, padding=model_name not in ModelConfig.MOLMO_MODELS @@ -104,9 +254,11 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100( ) else: model_hf = load_vlm_model_from_config(config) + qaic_config = KIMI_K25_QAIC_CONFIG if _is_kimi_k25(model_name) else None qeff_model = QEFFAutoModelForImageTextToText( copy.deepcopy(model_hf), kv_offload=kv_offload, + qaic_config=qaic_config, config=model_hf.config, torch_dtype=torch_dtype, ) @@ -188,6 +340,53 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100( inputs["pixel_values"] = inputs.pop("images") compile_kwargs["img_size"] = img_size + elif _is_kimi_k25(model_name): + processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True, padding=True) + image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB") + conversation = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": image}, + {"type": "text", "text": query}, + ], + }, + ] + prompt = processor.apply_chat_template(conversation, add_generation_prompt=True) + api_runner = ApiRunnerVlm( + batch_size, + processor, + config, + image, + conversation, + prompt, + prompt_len, + ctx_len, + max_gen_len, + num_hidden_layers, + ) + inputs = processor( + messages=conversation, + add_generation_prompt=True, + tokenize=False, + return_tensors="pt", + ) + pytorch_hf_tokens = _run_kimi_k25_hf_model_on_pytorch(model_hf, processor, inputs, max_gen_len) + inputs = processor( + messages=conversation, + add_generation_prompt=True, + tokenize=False, + return_tensors="pt", + ) + compile_kwargs.update( + { + "num_patches": int(inputs["pixel_values"].shape[0]), + "h": int(inputs["grid_thws"][0, 1].item()), + "w": int(inputs["grid_thws"][0, 2].item()), + "num_image_tokens": _get_kimi_k25_num_image_tokens(config, inputs["grid_thws"]), + } + ) + else: processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True, padding=True) image = Image.open(requests.get(img_url, stream=True).raw) @@ -351,7 +550,11 @@ def test_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100_qnn(model_name, kv_off ``Mandatory`` Args: :model_name (str): Hugging Face Model Card name, Example: ``gpt2`` """ - if model_name == "meta-llama/Llama-4-Scout-17B-16E-Instruct" or model_name == "google/gemma-3-4b-it": + if model_name in { + "meta-llama/Llama-4-Scout-17B-16E-Instruct", + "google/gemma-3-4b-it", + KIMI_K25_MODEL_NAME, + }: pytest.skip("QNN is not supported for these models yet.") qnn_config_json_path = os.path.join(os.getcwd(), "qnn_config.json") From 8f275332b058313839d7482c2c70ad23cb7522be Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Tue, 7 Jul 2026 15:41:34 +0530 Subject: [PATCH 13/33] update test_image_text_to_text_models.py Signed-off-by: Mamta Singh --- .../test_image_text_to_text_models.py | 71 +++++++++++++++++-- 1 file changed, 65 insertions(+), 6 deletions(-) diff --git a/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py b/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py index 45612a531f..91087c593e 100644 --- a/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py +++ b/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py @@ -51,7 +51,6 @@ NEW_GENERATION_TOKENS = 10 KIMI_K25_MODEL_NAME = "moonshotai/Kimi-K2.5" -KIMI_K25_QAIC_CONFIG = {"mla_absorption": {"cache_compressed": True, "absorption": False, "online": False}} def _is_kimi_k25(model_name: str) -> bool: @@ -152,6 +151,56 @@ def _get_kimi_k25_num_image_tokens(config, grid_thws): return int(grid_thws[0, 1].item() // merge_height) * int(grid_thws[0, 2].item() // merge_width) +def _disable_compressed_tensor_forward_pre_hooks(module: torch.nn.Module): + if hasattr(module, "_forward_pre_hooks"): + module._forward_pre_hooks.clear() + + +def _load_kimi_k25_layer_subset_model(): + examples_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../examples/kimi_k2")) + if examples_dir not in sys.path: + sys.path.insert(0, examples_dir) + + from test_kimi_k25 import ( # noqa: PLC0415 + LOADED_EXPERT_IDS, + NUM_EXPERTS_PER_TOKEN, + NUM_TEXT_LAYERS, + NUM_VISION_LAYERS, + _disable_compressed_tensor_forward_pre_hooks as _example_disable_compressed_tensor_forward_pre_hooks, + _load_layer_subset_model, + _materialize_missing_linear_weights, + _patch_deepseek_init_weights_compat, + _patch_kimi_tie_weights_compat, + _prepare_config, + _resolve_model_path, + _set_deterministic, + ) + + _set_deterministic(1234) + _ensure_torch_fx_import_compatibility() + model_path = _resolve_model_path() + config = _prepare_config(model_path) + kimi_cls = get_class_from_dynamic_module("modeling_kimi_k25.KimiK25ForConditionalGeneration", str(model_path)) + _patch_kimi_tie_weights_compat(kimi_cls) + _patch_deepseek_init_weights_compat(kimi_cls) + + model, tokenizer, processor = _load_layer_subset_model( + model_path=model_path, + kimi_cls=kimi_cls, + config=config, + num_vision_layers=NUM_VISION_LAYERS, + num_text_layers=NUM_TEXT_LAYERS, + loaded_expert_ids=LOADED_EXPERT_IDS, + num_experts_per_tok=NUM_EXPERTS_PER_TOKEN, + dtype=torch.float32, + ) + _materialize_missing_linear_weights(model.language_model) + _example_disable_compressed_tensor_forward_pre_hooks(model) + model.vision_tower.patch_embed.pos_emb.interpolation_mode = "bilinear" + model = model.eval().to("cpu") + return model, tokenizer, processor + + @torch.no_grad() def _run_kimi_k25_hf_model_on_pytorch(model, processor, inputs, max_gen_len): generated_ids = inputs["input_ids"] @@ -223,9 +272,19 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100( pytorch_kv_tokens = None ort_tokens = None n_layer = num_hidden_layers + kimi_tokenizer = None + kimi_processor = None if _is_kimi_k25(model_name) and config is None: - config = _get_kimi_k25_test_config(model_name) - if config is None: + model_hf, kimi_tokenizer, kimi_processor = _load_kimi_k25_layer_subset_model() + config = model_hf.config + qeff_model = QEFFAutoModelForImageTextToText( + copy.deepcopy(model_hf), + kv_offload=kv_offload, + config=model_hf.config, + torch_dtype=torch_dtype, + ) + _disable_compressed_tensor_forward_pre_hooks(qeff_model.model) + elif config is None: config = AutoConfig.from_pretrained( model_name, trust_remote_code=True, padding=model_name not in ModelConfig.MOLMO_MODELS ) @@ -259,11 +318,9 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100( ) else: model_hf = load_vlm_model_from_config(config) - qaic_config = KIMI_K25_QAIC_CONFIG if _is_kimi_k25(model_name) else None qeff_model = QEFFAutoModelForImageTextToText( copy.deepcopy(model_hf), kv_offload=kv_offload, - qaic_config=qaic_config, config=model_hf.config, torch_dtype=torch_dtype, ) @@ -353,7 +410,8 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100( compile_kwargs["img_size"] = img_size elif _is_kimi_k25(model_name): - processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True, padding=True) + processor = kimi_processor + tokenizer = kimi_tokenizer image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB") conversation = [ { @@ -392,6 +450,7 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100( ) compile_kwargs.update( { + "prefill_seq_len": 1, "num_patches": int(inputs["pixel_values"].shape[0]), "h": int(inputs["grid_thws"][0, 1].item()), "w": int(inputs["grid_thws"][0, 2].item()), From 445269c4a35a19e4b6b7aa582999fcd805e7e433 Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Thu, 9 Jul 2026 21:50:50 +0530 Subject: [PATCH 14/33] subfunction and 3 qpc flow verified Signed-off-by: Mamta Singh --- .../models/kimi_k25/modeling_kimi_k25.py | 2 +- .../test_kimi_k25_disagg.py | 310 ++++++++++++++++++ 2 files changed, 311 insertions(+), 1 deletion(-) create mode 100644 tests/transformers/models/image_text_to_text/test_kimi_k25_disagg.py diff --git a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py index 7076850235..ef3e66eb73 100644 --- a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py +++ b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py @@ -461,7 +461,7 @@ def get_submodules_for_export(self) -> Type[nn.Module]: This method should return the *class object* (not an instance). Downstream code can use this to find/build subfunctions for repeated blocks. """ - return {self.model.layers[0].__class__} + return {self.language_model.model.layers[0].__class__} def forward( self, diff --git a/tests/transformers/models/image_text_to_text/test_kimi_k25_disagg.py b/tests/transformers/models/image_text_to_text/test_kimi_k25_disagg.py new file mode 100644 index 0000000000..70c0855948 --- /dev/null +++ b/tests/transformers/models/image_text_to_text/test_kimi_k25_disagg.py @@ -0,0 +1,310 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ---------------------------------------------------------------------------- + +from pathlib import Path + +import numpy as np +import pytest +import torch +from export_kimi_k25_vision import ( + LOADED_EXPERT_IDS, + NUM_EXPERTS_PER_TOKEN, + _ensure_torch_fx_import_compatibility, + _load_layer_subset_model, + _patch_deepseek_init_weights_compat, + _patch_kimi_tie_weights_compat, + _prepare_config, +) +from test_kimi_k25 import ( + _clone_inputs, + _decode_tokens, + _disable_compressed_tensor_forward_pre_hooks, + _greedy_generate_hf, + _has_qaic_runtime_access, + _materialize_missing_linear_weights, + _prepare_inputs, + _resolve_model_path, + _set_deterministic, +) +from transformers.dynamic_module_utils import get_class_from_dynamic_module + +from QEfficient import QEFFAutoModelForImageTextToText +from QEfficient.generation.cloud_infer import QAICInferenceSession + +PREFILL_SEQ_LEN = 32 +CTX_LEN = 1024 +BATCH_SIZE = 1 +GENERATION_LEN = 4 +NUM_VISION_LAYERS = 2 +NUM_TEXT_LAYERS = 2 + + +def _assert_onnx_path(onnx_path, label: str) -> Path: + assert onnx_path is not None, f"{label} compile did not set an ONNX path" + onnx_path = Path(onnx_path) + assert onnx_path.is_file(), f"{label} ONNX path does not exist: {onnx_path}" + assert onnx_path.suffix == ".onnx", f"{label} path is not an ONNX file: {onnx_path}" + return onnx_path.resolve() + + +def _assert_distinct_onnx_paths(onnx_paths: dict[str, Path]): + unique_paths = {str(path) for path in onnx_paths.values()} + assert len(unique_paths) == len(onnx_paths), f"Expected distinct ONNX paths per compile, got: {onnx_paths}" + + +def _numpy(value): + if torch.is_tensor(value): + return value.detach().cpu().numpy() + return np.asarray(value) + + +def _session_input_names(session: QAICInferenceSession) -> set[str]: + input_names = set(session.input_names) + input_names.update(name.rsplit("/", 1)[-1] for name in session.input_names) + return input_names + + +def _cast_for_session(session: QAICInferenceSession, name: str, value: np.ndarray) -> np.ndarray: + binding_index = session.binding_index_map.get(name) + if binding_index is None: + return value + dtype = session.aic_to_np_dtype_mapping[session.bindings[binding_index].type] + return value.astype(dtype, copy=False) + + +def _filter_session_inputs(session: QAICInferenceSession, inputs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: + input_names = _session_input_names(session) + return {name: _cast_for_session(session, name, value) for name, value in inputs.items() if name in input_names} + + +def _get_next_token_ids(logits: np.ndarray) -> np.ndarray: + logits = np.asarray(logits) + return logits[:, -1, :].argmax(axis=-1).astype(np.int64).reshape(BATCH_SIZE, 1) + + +def _update_retained_states(target_inputs: dict[str, np.ndarray], source_outputs: dict[str, np.ndarray]): + for output_name, value in source_outputs.items(): + output_basename = output_name.rsplit("/", 1)[-1] + if not output_basename.endswith("_RetainedState"): + continue + target_inputs[output_basename.removesuffix("_RetainedState")] = value + + +def _load_kimi_subset_model(): + _set_deterministic(1234) + _ensure_torch_fx_import_compatibility() + model_path = _resolve_model_path() + config = _prepare_config(model_path) + kimi_cls = get_class_from_dynamic_module("modeling_kimi_k25.KimiK25ForConditionalGeneration", str(model_path)) + _patch_kimi_tie_weights_compat(kimi_cls) + _patch_deepseek_init_weights_compat(kimi_cls) + + model, tokenizer, processor = _load_layer_subset_model( + model_path=model_path, + kimi_cls=kimi_cls, + config=config, + num_vision_layers=NUM_VISION_LAYERS, + num_text_layers=NUM_TEXT_LAYERS, + loaded_expert_ids=LOADED_EXPERT_IDS, + num_experts_per_tok=NUM_EXPERTS_PER_TOKEN, + dtype=torch.float32, + ) + _materialize_missing_linear_weights(model.language_model) + _disable_compressed_tensor_forward_pre_hooks(model) + model.vision_tower.patch_embed.pos_emb.interpolation_mode = "bilinear" + return model.eval().to("cpu"), tokenizer, processor + + +def _get_image_compile_dims(model, inputs: dict[str, torch.Tensor]) -> dict[str, int]: + grid_thws = inputs["grid_thws"].to(torch.long) + h = int(grid_thws[0, 1].item()) + w = int(grid_thws[0, 2].item()) + num_patches = int(inputs["pixel_values"].shape[0]) + merge_height, merge_width = model.vision_tower.merge_kernel_size + num_images = max(num_patches // (h * w), 1) + num_image_tokens = num_images * (h // merge_height) * (w // merge_width) + return {"num_patches": num_patches, "h": h, "w": w, "num_image_tokens": num_image_tokens} + + +def _compile_disagg_qpcs(qeff_model: QEFFAutoModelForImageTextToText, compile_dims: dict[str, int]): + common_compile_kwargs = { + "batch_size": BATCH_SIZE, + "ctx_len": CTX_LEN, + "num_devices": 1, + "num_cores": 16, + "mxfp6_matmul": False, + "split_model_io": True, + "mos": 1, + "aic_enable_depth_first": True, + "use_onnx_subfunctions": True, + "layerwise": False, + **compile_dims, + } + + compiled_onnx_paths = {} + vision_qpc_path = qeff_model.compile( + prefill_seq_len=PREFILL_SEQ_LEN, + skip_vision=False, + skip_lang=True, + **common_compile_kwargs, + ) + compiled_onnx_paths["vision"] = _assert_onnx_path(qeff_model.vision_model.onnx_path, "vision") + + prefill_qpc_path = qeff_model.compile( + prefill_seq_len=PREFILL_SEQ_LEN, + retain_full_kv=True, + prefill_only=True, + enable_chunking=True, + skip_vision=True, + skip_lang=False, + **common_compile_kwargs, + ) + compiled_onnx_paths["prefill"] = _assert_onnx_path(qeff_model.lang_model.onnx_path, "prefill") + + decode_qpc_path = qeff_model.compile( + prefill_seq_len=1, + prefill_only=False, + skip_vision=True, + skip_lang=False, + **common_compile_kwargs, + ) + compiled_onnx_paths["decode"] = _assert_onnx_path(qeff_model.lang_model.onnx_path, "decode") + _assert_distinct_onnx_paths(compiled_onnx_paths) + + return vision_qpc_path, prefill_qpc_path, decode_qpc_path, compiled_onnx_paths + + +def _run_disagg_qaic_generation( + common_inputs: dict[str, torch.Tensor], + vision_session: QAICInferenceSession, + prefill_session: QAICInferenceSession, + decode_session: QAICInferenceSession, +) -> np.ndarray: + inputs = {name: _numpy(value) for name, value in _clone_inputs(common_inputs).items()} + input_ids_length = inputs["input_ids"].shape[1] + num_chunks = -(input_ids_length // -PREFILL_SEQ_LEN) + padded_len = num_chunks * PREFILL_SEQ_LEN + + inputs["input_ids"] = np.pad( + inputs["input_ids"], + ((0, 0), (0, padded_len - input_ids_length)), + constant_values=1, + ) + inputs["attention_mask"] = np.pad( + inputs["attention_mask"], + ((0, 0), (0, padded_len - input_ids_length)), + constant_values=0, + ) + + grid_thws = inputs["grid_thws"].astype(np.int64) + h = int(grid_thws[0, 1]) + w = int(grid_thws[0, 2]) + vision_inputs = { + "pixel_values": inputs["pixel_values"], + "h_shape": np.ones((h,), dtype=np.int64), + "w_shape": np.ones((w,), dtype=np.int64), + } + vision_outputs = vision_session.run(_filter_session_inputs(vision_session, vision_inputs)) + vision_session.deactivate() + + image_embeds = vision_outputs.get("image_embeds") + assert image_embeds is not None, f"Vision QPC did not return image_embeds. Outputs: {vision_outputs.keys()}" + + lang_inputs = { + "input_ids": inputs["input_ids"].astype(np.int64), + "position_ids": np.where(inputs["attention_mask"] > 0, np.arange(padded_len), -1).astype(np.int64), + "image_embeds": image_embeds, + "image_idx": np.zeros((BATCH_SIZE, 1), dtype=np.int64), + } + + prefill_session.set_buffers(vision_outputs) + chunk_inputs = lang_inputs.copy() + prefill_outputs = None + for chunk_idx in range(num_chunks): + start = chunk_idx * PREFILL_SEQ_LEN + end = (chunk_idx + 1) * PREFILL_SEQ_LEN + chunk_inputs["input_ids"] = lang_inputs["input_ids"][:, start:end] + chunk_inputs["position_ids"] = lang_inputs["position_ids"][:, start:end] + prefill_outputs = prefill_session.run(_filter_session_inputs(prefill_session, chunk_inputs)) + _update_retained_states(chunk_inputs, prefill_outputs) + if "image_idx_output" in prefill_outputs: + chunk_inputs["image_idx"] = prefill_outputs["image_idx_output"].astype(np.int64) + + prefill_session.deactivate() + assert prefill_outputs is not None, "QAIC prefill did not execute." + + generated_ids = [_get_next_token_ids(prefill_outputs["logits"])] + decode_inputs = { + "input_ids": generated_ids[-1], + "position_ids": np.max(lang_inputs["position_ids"], axis=-1, keepdims=True).astype(np.int64) + 1, + "image_embeds": chunk_inputs.get("image_embeds", image_embeds), + "image_idx": chunk_inputs.get("image_idx", np.zeros((BATCH_SIZE, 1), dtype=np.int64)), + } + _update_retained_states(decode_inputs, prefill_outputs) + + for _ in range(1, GENERATION_LEN): + decode_outputs = decode_session.run(_filter_session_inputs(decode_session, decode_inputs)) + generated_ids.append(_get_next_token_ids(decode_outputs["logits"])) + decode_inputs["input_ids"] = generated_ids[-1] + decode_inputs["position_ids"] = decode_inputs["position_ids"] + 1 + if "image_idx_output" in decode_outputs: + decode_inputs["image_idx"] = decode_outputs["image_idx_output"].astype(np.int64) + _update_retained_states(decode_inputs, decode_outputs) + + return np.concatenate(generated_ids, axis=1) + + +@pytest.mark.on_qaic +@pytest.mark.multimodal +def test_kimi_k25_disagg_qaic_vs_hf_fp32(): + if not _has_qaic_runtime_access(): + pytest.skip("QAIC runtime is not available.") + + model, tokenizer, processor = _load_kimi_subset_model() + inputs = _prepare_inputs(processor) + inputs = {name: (value.to("cpu") if torch.is_tensor(value) else value) for name, value in inputs.items()} + hf_tokens = _greedy_generate_hf(model, _clone_inputs(inputs), max_new_tokens=GENERATION_LEN) + + qeff_model = QEFFAutoModelForImageTextToText( + model, + kv_offload=True, + config=model.config, + torch_dtype=torch.float32, + layerwise=False, + ) + _disable_compressed_tensor_forward_pre_hooks(qeff_model.model) + + compile_dims = _get_image_compile_dims(qeff_model.model, inputs) + vision_qpc_path, prefill_qpc_path, decode_qpc_path, compiled_onnx_paths = _compile_disagg_qpcs( + qeff_model, + compile_dims, + ) + print(f"Kimi-K2.5 disagg ONNX paths: {compiled_onnx_paths}") + + sessions = [] + try: + vision_session = QAICInferenceSession(vision_qpc_path.get("vision_qpc_path")) + prefill_session = QAICInferenceSession(prefill_qpc_path.get("lang_prefill_qpc_path")) + decode_session = QAICInferenceSession(decode_qpc_path.get("lang_decode_qpc_path")) + sessions.extend([vision_session, prefill_session, decode_session]) + qaic_tokens = _run_disagg_qaic_generation( + common_inputs=inputs, + vision_session=vision_session, + prefill_session=prefill_session, + decode_session=decode_session, + ) + finally: + for session in sessions: + session.deactivate() + + print("HF:", _decode_tokens(tokenizer, hf_tokens), "\n", hf_tokens) + print("Disagg QAIC:", _decode_tokens(tokenizer, torch.as_tensor(qaic_tokens)), "\n", qaic_tokens) + + assert qaic_tokens.shape == (BATCH_SIZE, GENERATION_LEN) + assert hf_tokens.shape == (BATCH_SIZE, GENERATION_LEN) + assert np.issubdtype(qaic_tokens.dtype, np.integer) + assert torch.equal(hf_tokens, torch.as_tensor(qaic_tokens)), "HF and disagg QAIC tokens do not match" From 344ff995103bf9ca2e149b34408820da674c0412 Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Fri, 10 Jul 2026 00:46:24 +0530 Subject: [PATCH 15/33] int4 changes for kimi Signed-off-by: Mamta Singh --- QEfficient/base/onnx_transforms.py | 5 +- QEfficient/base/pytorch_transforms.py | 12 +- QEfficient/customop/matmulnbits.py | 88 +- QEfficient/customop/quantization_ops.py | 149 +++ QEfficient/transformers/cache_utils.py | 58 +- .../models/deepseek_v3/modeling_deepseek.py | 845 +++++++++++++++--- .../transformers/models/modeling_auto.py | 36 +- .../transformers/models/pytorch_transforms.py | 13 +- .../quantizers/quant_transforms.py | 80 ++ .../quantizer_compressed_tensors.py | 158 +++- QEfficient/utils/constants.py | 2 + 11 files changed, 1269 insertions(+), 177 deletions(-) create mode 100644 QEfficient/customop/quantization_ops.py diff --git a/QEfficient/base/onnx_transforms.py b/QEfficient/base/onnx_transforms.py index 271d1670e1..ee044642b9 100644 --- a/QEfficient/base/onnx_transforms.py +++ b/QEfficient/base/onnx_transforms.py @@ -43,8 +43,7 @@ CtxScatterFuncCB, CtxScatterFuncCB3D, ) - -# from QEfficient.customop.quantization_ops import CastToUInt4, CastToUInt4Func +from QEfficient.customop.quantization_ops import CastToUInt4, CastToUInt4Func from QEfficient.customop.rms_norm import CustomRMSNorm, CustomRMSNormFunc from QEfficient.utils.constants import FILE_CHUNK_SIZE_DEFAULT, ONNX_EXPORT_OPSET, SIZE_THRESHOLD_DEFAULT @@ -109,7 +108,7 @@ class CustomOpTransform(BaseOnnxTransform): "CtxGatherFuncBlockedKVCB": (CtxGatherFuncBlockedKVCB, CtxGatherBlockedKVCB), "CtxScatterFuncCB": (CtxScatterFuncCB, CtxScatterCB), "CtxGatherFuncCB": (CtxGatherFuncCB, CtxGatherCB), - # "CastToUInt4": (CastToUInt4Func, CastToUInt4), + "CastToUInt4": (CastToUInt4Func, CastToUInt4), } @classmethod diff --git a/QEfficient/base/pytorch_transforms.py b/QEfficient/base/pytorch_transforms.py index 812177eac2..a716805d36 100644 --- a/QEfficient/base/pytorch_transforms.py +++ b/QEfficient/base/pytorch_transforms.py @@ -5,7 +5,7 @@ # # ---------------------------------------------------------------------------- from types import MethodType -from typing import Callable, Dict, Tuple, Type +from typing import Callable, Dict, Optional, Tuple, Type from torch import nn @@ -97,6 +97,7 @@ class ModuleMutatorTransform(PytorchTransform): """ _match_class: nn.Module + _match_string: Optional[str] = None @classmethod def apply(cls, model: nn.Module) -> Tuple[nn.Module, bool]: @@ -135,7 +136,14 @@ def apply(cls, model: nn.Module) -> Tuple[nn.Module, bool]: repl_method_map := cls._match_string_replace_method.get(module.__class__.__name__) ): for orig_method_name, mapped_method in repl_method_map.items(): - setattr(module, orig_method_name, MethodType(mapped_method, module)) + parts = orig_method_name.split(".") + if len(parts) > 1: + target = module + for part in parts[:-1]: + target = getattr(target, part) + setattr(target, parts[-1], MethodType(mapped_method, target)) + else: + setattr(module, orig_method_name, MethodType(mapped_method, module)) if hasattr(module, "__qeff_init__"): module.__qeff_init__() diff --git a/QEfficient/customop/matmulnbits.py b/QEfficient/customop/matmulnbits.py index d8cc0e8f1b..72cf898ad5 100644 --- a/QEfficient/customop/matmulnbits.py +++ b/QEfficient/customop/matmulnbits.py @@ -14,7 +14,10 @@ class QuantLinearTorchFunction(torch.autograd.Function): @staticmethod def symbolic(g, x, qself_qweight, qself_scales, qself_qzeros, g_idx, bits, group_size, in_features, out_features): - input_tuple = (x, qself_qweight, qself_scales, qself_qzeros) + if qself_qzeros is None: + input_tuple = (x, qself_qweight, qself_scales) + else: + input_tuple = (x, qself_qweight, qself_scales, qself_qzeros) input_tuple += (g_idx,) if g_idx is not None else () return g.op( "com.microsoft::MatMulNBits", @@ -28,7 +31,10 @@ def symbolic(g, x, qself_qweight, qself_scales, qself_qzeros, g_idx, bits, group @staticmethod def forward(ctx, x, qself_qweight, qself_scales, qself_qzeros, g_idx, bits, group_size, in_features, out_features): + if qself_qzeros is None: + qself_qzeros = 2 ^ (bits - 1) if torch.onnx.is_in_onnx_export(): + # For faster export return torch.zeros(x.shape[:-1] + (out_features,), dtype=x.dtype).float() fp_weight = dequantize_blockwise_bits( qself_qweight, qself_scales, qself_qzeros, bits, group_size, g_idx, in_features, out_features @@ -40,8 +46,7 @@ def forward(ctx, x, qself_qweight, qself_scales, qself_qzeros, g_idx, bits, grou def dequantize_blockwise_bits(quant_values, scale, zero_point, bits, group_size, g_idx, rows, cols): if bits != 4: raise ValueError("Only bits=4 is supported for executing quantized model") - if group_size != 128: - raise ValueError("Only group_size=128 is supported for executing quantized model") + expand_quant_value = (quant_values.unsqueeze(-1) >> torch.tensor([[[[0, 4]]]], dtype=torch.int32)) & 0x0F expand_quant_value = expand_quant_value.reshape(*quant_values.shape[:-1], -1) aligned_scale = scale.reshape(*quant_values.shape[:-1], 1) @@ -88,20 +93,20 @@ def __init__(self, bits, group_size, in_features, out_features, bias): q_rows = in_features // self.group_size self.register_buffer( "qweight", - torch.zeros((out_features, q_rows, self.group_size // (8 // bits)), dtype=torch.uint8), + torch.empty((out_features, q_rows, self.group_size // (8 // bits)), dtype=torch.uint8), ) self.register_buffer( "qzeros", - torch.zeros((q_rows + (q_rows & 1)) * (out_features // 8 * self.bits), dtype=torch.uint8), + torch.empty((q_rows + (q_rows & 1)) * (out_features // 8 * self.bits), dtype=torch.uint8), ) self.register_buffer( - "scales", torch.zeros((math.ceil(in_features / self.group_size) * out_features), dtype=torch.float16) + "scales", torch.empty((math.ceil(in_features / self.group_size) * out_features), dtype=torch.float16) ) self.register_buffer( "g_idx", torch.tensor([i // self.group_size for i in range(in_features)], dtype=torch.int32) ) if bias: - self.register_buffer("bias", torch.zeros((out_features), dtype=torch.float16)) + self.register_buffer("bias", torch.empty((out_features), dtype=torch.float16)) else: self.bias = None @@ -183,3 +188,72 @@ def forward(self, inputs): ) out = out + self.bias if self.bias is not None else out return out + + +class QMOE(torch.autograd.Function): + @staticmethod + def symbolic( + g, + x, + router_weights, + fc1_experts_weights, + fc1_scales, + fc2_experts_weights, + fc2_scales, + fc3_experts_weights, + fc3_scales, + router_probs, + activation_type, + block_size, + expert_weight_bits, + k, + ): + qmoe_out = g.op( + "com.microsoft::QMoE", + x, + router_weights, + router_probs, + fc1_experts_weights, + fc1_scales, + fc2_experts_weights, + fc2_scales, + fc3_experts_weights, + fc3_scales, + outputs=1, + activation_type_s=activation_type, # <-- _s suffix for string + block_size_i=block_size, + expert_weight_bits_i=expert_weight_bits, + k_i=k, + ) + + # # Create axes=-1 as an explicit int64 constant tensor + # axes = g.op("Constant", value_t=torch.tensor([-1], dtype=torch.int64)) + + # # Compute mean of router_probs along the last axis, keepdims for broadcasting + # router_probs_mean = g.op("ReduceMean", router_probs, axes, keepdims_i=1) + + # Multiply qmoe_out with the averaged router_probs + return qmoe_out + # return g.op("Mul", qmoe_out, router_probs_mean)) + + @staticmethod + def forward( + ctx, + x, + router_weights, + fc1_experts_weights, + fc1_scales, + fc2_experts_weights, + fc2_scales, + fc3_experts_weights, + fc3_scales, + router_probs, + activation_type, + block_size, + expert_weight_bits, + k, + ): + # Dummy forward: simulate qmoe_out as zeros_like(x), then apply ReduceMean * Mul + qmoe_out = torch.zeros_like(x) + router_probs_mean = router_probs.mean(dim=-1, keepdim=True) + return qmoe_out * router_probs_mean diff --git a/QEfficient/customop/quantization_ops.py b/QEfficient/customop/quantization_ops.py new file mode 100644 index 0000000000..3878804c7d --- /dev/null +++ b/QEfficient/customop/quantization_ops.py @@ -0,0 +1,149 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import onnxscript +import torch +from onnx import TensorProto + +from QEfficient.utils import constants + +ops = getattr(onnxscript, "opset" + str(constants.ONNX_EXPORT_OPSET)) + + +@onnxscript.script(onnxscript.values.Opset("com.qti.aisw.onnx", 1)) +def CastToUInt4(weight_packed: onnxscript.UINT8) -> onnxscript.UINT8: + """ + Unpack packed uint8 weights into uint4 values and cast output to UINT4. + Supports N-D input: all leading dimensions are preserved; only the last + dimension (in_features // 2) is doubled to (in_features). + + Input: (..., in_features // 2) UINT8 + Each byte holds two nibbles: byte = (w_y << 4) | (w_x & 0x0F) + Output: (..., in_features) UINT4, values in [0, 15] + + Operations: + w_x = weight_packed % 16 (lower nibble) + w_y = (weight_packed >> 4) % 16 (upper nibble) + stacked = concat([w_x, w_y], axis=-1) after unsqueeze + → (..., in//2, 2) + leading_dims = shape[:-1] + new_shape = [...leading_dims, last_dim * 2] + reshaped = reshape(stacked, new_shape) + output = Cast(reshaped, to=UINT4) + """ + sixteen = ops.CastLike(ops.Constant(value_ints=[16]), weight_packed) + + # Lower nibble: weight_packed & 0x0F = weight_packed % 16 + w_x = ops.Mod(weight_packed, sixteen) + + # Upper nibble: (weight_packed >> 4) & 0x0F + shift = ops.CastLike(ops.Constant(value_ints=[4]), weight_packed) + w_shifted = ops.BitShift(weight_packed, shift, direction="RIGHT") + w_y = ops.Mod(w_shifted, sixteen) + + # Stack along a new last dim → (..., in_features//2, 2) + w_x_unsq = ops.Unsqueeze(w_x, [-1]) + w_y_unsq = ops.Unsqueeze(w_y, [-1]) + stacked = ops.Concat(w_x_unsq, w_y_unsq, axis=-1) + + # N-D aware reshape: preserve all leading dims, double the last dim. + # packed_shape = [d0, d1, ..., last_dim] + packed_shape = ops.Shape(weight_packed) + # All dims except the last: [d0, d1, ...] + leading_dims = ops.Slice(packed_shape, starts=[0], ends=[-1], axes=[0]) + # Last dim only: [last_dim] + last_dim = ops.Slice(packed_shape, starts=[-1], ends=[2147483647], axes=[0]) + # Double the last dim: [last_dim * 2] + last_dim_doubled = ops.Mul(last_dim, ops.Constant(value_ints=[2])) + # New shape: [d0, d1, ..., last_dim * 2] + new_shape = ops.Concat(leading_dims, last_dim_doubled, axis=0) + reshaped = ops.Reshape(stacked, new_shape) + + # Cast to UINT4 — data_type value is version-dependent (21 in ONNX 1.18, 23 in newer) + return ops.Cast(reshaped, to=int(TensorProto.UINT4)) + + +class CastToUInt4Func(torch.autograd.Function): + """ + Custom op: unpacks packed uint8 → uint8 (values 0-15) in PyTorch. + In ONNX the custom op subgraph includes a Cast → UINT4 as its last step. + Supports N-D input: all leading dimensions are preserved. + + PyTorch forward : packed uint8 (..., in//2) → uint8 (..., in), values [0, 15] + ONNX symbolic : emits CastToUInt4 node (com.qti.aisw.onnx) + The subgraph ends with Cast → UINT4. + """ + + @staticmethod + def forward(weight_packed: torch.Tensor) -> torch.Tensor: + w_x = weight_packed & 0x0F # lower nibble, (..., in//2), range [0, 15] + w_y = (weight_packed >> 4) & 0x0F # upper nibble, (..., in//2), range [0, 15] + # New shape: all leading dims unchanged, last dim doubled + new_shape = list(weight_packed.shape[:-1]) + [weight_packed.shape[-1] * 2] + return torch.stack( + [w_x, w_y], dim=-1 + ).reshape( + new_shape + ) # Can't add a cast operation to uint4 here, as its not supported in pytorch; The ONNX export will handle the cast to IINT4 in the symbolic method. + + @staticmethod + def setup_context(ctx, inputs, outputs): + pass + + @staticmethod + def symbolic(g: torch.Graph, weight_packed: torch.Value) -> torch.Value: + output = g.onnxscript_op(CastToUInt4, weight_packed) + return output + + +class DequantizeLinearFunc(torch.autograd.Function): + """ + Emits a standard ONNX DequantizeLinear node (ai.onnx domain, not custom). + + Symmetric blockwise quantization — no zero_point: + output = x * scale (per block along the last axis) + + Supports N-D input: + weight_unpacked : (..., in_features) — quantized values + scale : (..., num_blocks) — per-block scales + block_size : int — elements per block + + PyTorch forward : expand blockwise scale along last dim, multiply + ONNX symbolic : DequantizeLinear(weight_unpacked, scale, + axis=2, block_size=block_size) + axis=2 for 3D input (2, out_features, in_features). + No zero_point input (symmetric). + """ + + @staticmethod + def forward( + weight_unpacked: torch.Tensor, scale: torch.Tensor, zeros: torch.Tensor, block_size: int + ) -> torch.Tensor: + # Expand per-block scale → per-element scale along last dim + scale_expanded = scale.repeat_interleave(block_size, dim=-1) + zeros_expanded = zeros.repeat_interleave(block_size, dim=-1) + return (weight_unpacked.to(torch.int8) - zeros_expanded.to(torch.int8)) * scale_expanded + + @staticmethod + def setup_context(ctx, inputs, outputs): + pass + + @staticmethod + def symbolic( + g: torch.Graph, weight_unpacked: torch.Value, scale: torch.Value, zeros: torch.Value, block_size: int + ) -> torch.Value: + # Standard DequantizeLinear: symmetric (no zero_point), blockwise. + # Input is 3D: (2, out_features, in_features) → axis=2 (last dim). + # DequantizeLinear natively supports batch dimensions. + return g.op( + "DequantizeLinear", + weight_unpacked, + scale, + zeros, + axis_i=2, + block_size_i=block_size, + ) diff --git a/QEfficient/transformers/cache_utils.py b/QEfficient/transformers/cache_utils.py index f6c2b61282..fd3ac24310 100755 --- a/QEfficient/transformers/cache_utils.py +++ b/QEfficient/transformers/cache_utils.py @@ -420,8 +420,14 @@ def __init__(self, ckv, k_pe): def update_ckv(self, compressed_kv, cache_kwargs): position_ids = cache_kwargs.get("position_ids") + batch_index = cache_kwargs.get("batch_index", None) # Check and fetch batch index value form the kwargs - self.ckv = CtxScatterFunc.apply(self.ckv, position_ids, compressed_kv) + if batch_index is not None: + invalid_scatter_index = torch.iinfo(torch.int32).max + scatter_position_ids = torch.where(position_ids < 0, invalid_scatter_index, position_ids) + self.ckv = CtxScatterFuncCB.apply(self.ckv, batch_index, scatter_position_ids, compressed_kv) + else: + self.ckv = CtxScatterFunc.apply(self.ckv, position_ids, compressed_kv) ckv_out = self.ckv ctx_len = ckv_out.shape[-2] @@ -434,14 +440,23 @@ def update_ckv(self, compressed_kv, cache_kwargs): invalid_idx_value = 0 ctx_indices = torch.where(invalid_mask, invalid_idx_value, ctx_indices) - ckv_out = CtxGatherFunc.apply(ckv_out, ctx_indices, ctx_len) + if batch_index is not None: + ckv_out = CtxGatherFuncCB.apply(ckv_out, batch_index, ctx_indices, ctx_len) + else: + ckv_out = CtxGatherFunc.apply(ckv_out, ctx_indices, ctx_len) ckv_out = torch.where(invalid_mask.unsqueeze(-1), torch.tensor(0.0, dtype=torch.float32), ckv_out) return ckv_out def update_k_pe(self, k_pe_cache, cache_kwargs): position_ids = cache_kwargs.get("position_ids") + batch_index = cache_kwargs.get("batch_index", None) # Check and fetch batch index value form the kwargs - self.k_pe = CtxScatterFunc.apply(self.k_pe, position_ids, k_pe_cache) + if batch_index is not None: + invalid_scatter_index = torch.iinfo(torch.int32).max + scatter_position_ids = torch.where(position_ids < 0, invalid_scatter_index, position_ids) + self.k_pe = CtxScatterFuncCB.apply(self.k_pe, batch_index, scatter_position_ids, k_pe_cache) + else: + self.k_pe = CtxScatterFunc.apply(self.k_pe, position_ids, k_pe_cache) k_pe_out = self.k_pe ctx_len = k_pe_out.shape[-2] @@ -454,7 +469,10 @@ def update_k_pe(self, k_pe_cache, cache_kwargs): invalid_idx_value = 0 ctx_indices = torch.where(invalid_mask, invalid_idx_value, ctx_indices) - k_pe_out = CtxGatherFunc.apply(k_pe_out, ctx_indices, ctx_len) + if batch_index is not None: + k_pe_out = CtxGatherFuncCB.apply(k_pe_out, batch_index, ctx_indices, ctx_len) + else: + k_pe_out = CtxGatherFunc.apply(k_pe_out, ctx_indices, ctx_len) k_pe_out = torch.where(invalid_mask.unsqueeze(-1), torch.tensor(0.0, dtype=torch.float32), k_pe_out) return k_pe_out @@ -462,6 +480,7 @@ def read_only_blocked_ckv(self, start_index, end_index, cache_kwargs): # Gather ckv_out = self.ckv position_ids = cache_kwargs.get("position_ids") + batch_index = cache_kwargs.get("batch_index", None) batch, num_kv_heads, _, _ = ckv_out.shape ctx_indices = torch.arange(start=start_index, end=end_index)[None, None, ...] gather_limit = position_ids.max(1, keepdim=True).values.unsqueeze(1) @@ -474,8 +493,11 @@ def read_only_blocked_ckv(self, start_index, end_index, cache_kwargs): ctx_indices = torch.where(invalid_mask, invalid_idx_value, ctx_indices) - ctx_indices = ctx_indices.expand(batch, num_kv_heads, ctx_indices.shape[-1]) - ckv_out = CtxGatherFuncBlockedKV.apply(ckv_out, ctx_indices) + if batch_index is not None: + ckv_out = CtxGatherFuncBlockedKVCB.apply(ckv_out, batch_index, ctx_indices) + else: + ctx_indices = ctx_indices.expand(batch, num_kv_heads, ctx_indices.shape[-1]) + ckv_out = CtxGatherFuncBlockedKV.apply(ckv_out, ctx_indices) ckv_out = torch.where(invalid_mask.unsqueeze(-1), torch.tensor(0.0, dtype=torch.float32), ckv_out) return ckv_out @@ -484,6 +506,7 @@ def read_only_blocked_k_pe(self, start_index, end_index, cache_kwargs): # Gather k_pe_out = self.k_pe position_ids = cache_kwargs.get("position_ids") + batch_index = cache_kwargs.get("batch_index", None) batch, num_kv_heads, _, _ = k_pe_out.shape ctx_indices = torch.arange(start=start_index, end=end_index)[None, None, ...] gather_limit = position_ids.max(1, keepdim=True).values.unsqueeze(1) @@ -496,22 +519,37 @@ def read_only_blocked_k_pe(self, start_index, end_index, cache_kwargs): ctx_indices = torch.where(invalid_mask, invalid_idx_value, ctx_indices) - ctx_indices = ctx_indices.expand(batch, num_kv_heads, ctx_indices.shape[-1]) - k_pe_out = CtxGatherFuncBlockedKV.apply(k_pe_out, ctx_indices) + if batch_index is not None: + k_pe_out = CtxGatherFuncBlockedKVCB.apply(k_pe_out, batch_index, ctx_indices) + else: + ctx_indices = ctx_indices.expand(batch, num_kv_heads, ctx_indices.shape[-1]) + k_pe_out = CtxGatherFuncBlockedKV.apply(k_pe_out, ctx_indices) k_pe_out = torch.where(invalid_mask.unsqueeze(-1), torch.tensor(0.0, dtype=torch.float32), k_pe_out) return k_pe_out def write_only_k_pe(self, k_pe_cache, cache_kwargs): position_ids = cache_kwargs.get("position_ids") + batch_index = cache_kwargs.get("batch_index", None) # Check and fetch batch index value form the kwargs - self.k_pe = CtxScatterFunc.apply(self.k_pe, position_ids, k_pe_cache) + if batch_index is not None: + invalid_scatter_index = torch.iinfo(torch.int32).max + scatter_position_ids = torch.where(position_ids < 0, invalid_scatter_index, position_ids) + self.k_pe = CtxScatterFuncCB.apply(self.k_pe, batch_index, scatter_position_ids, k_pe_cache) + else: + self.k_pe = CtxScatterFunc.apply(self.k_pe, position_ids, k_pe_cache) return self.k_pe def write_only_ckv(self, compressed_kv, cache_kwargs): position_ids = cache_kwargs.get("position_ids") + batch_index = cache_kwargs.get("batch_index", None) # Check and fetch batch index value form the kwargs - self.ckv = CtxScatterFunc.apply(self.ckv, position_ids, compressed_kv) + if batch_index is not None: + invalid_scatter_index = torch.iinfo(torch.int32).max + scatter_position_ids = torch.where(position_ids < 0, invalid_scatter_index, position_ids) + self.ckv = CtxScatterFuncCB.apply(self.ckv, batch_index, scatter_position_ids, compressed_kv) + else: + self.ckv = CtxScatterFunc.apply(self.ckv, position_ids, compressed_kv) return self.ckv diff --git a/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py b/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py index 8a8b3efcf9..2eab9ff5bf 100644 --- a/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py +++ b/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py @@ -20,6 +20,13 @@ generic_blocked_attention_interface, generic_blocked_mla_attention_interface, ) +from QEfficient.customop.ctx_scatter_gather import ( + CtxGatherFunc3DGeneralized, + CtxScatterFunc3DGeneralized, + CtxScatterFunc3DInt, +) +from QEfficient.customop.matmulnbits import QMOE, QuantLinearTorchFunction +from QEfficient.customop.quantization_ops import CastToUInt4Func, DequantizeLinearFunc from QEfficient.customop.rms_norm import CustomRMSNormFunc from QEfficient.transformers.cache_utils import QEffDynamicCache, QEffDynamicCompressedKVRopeCache from QEfficient.transformers.modeling_attn_mask_utils import _create_causal_mask @@ -132,15 +139,16 @@ def _set_cos_sin_cache(self, seq_len, device, dtype): self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False) self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False) - def forward(self, x, seq_len=None): - # x: [bs, num_attention_heads, seq_len, head_size] - if self.max_seq_len_cached is None or seq_len > self.max_seq_len_cached: - self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype) - return ( - self.cos_cached[:seq_len].to(dtype=x.dtype), - self.sin_cached[:seq_len].to(dtype=x.dtype), - ) +# def forward(self, x, seq_len=None): +# # x: [bs, num_attention_heads, seq_len, head_size] +# if self.max_seq_len_cached is None or seq_len > self.max_seq_len_cached: +# self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype) +# +# return ( +# self.cos_cached[:seq_len].to(dtype=x.dtype), +# self.sin_cached[:seq_len].to(dtype=x.dtype), +# ) class DeepseekV3YarnRotaryEmbedding(DeepseekV3RotaryEmbedding): @@ -201,7 +209,7 @@ def _set_cos_sin_cache(self, seq_len, device, dtype): # Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb -def orig_apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1): +def orig_apply_rotary_pos_emb(q, k, cos, sin): # , position_ids, unsqueeze_dim=1): """Applies Rotary Position Embedding to the query and key tensors. Args: @@ -222,8 +230,8 @@ def orig_apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1): Returns: `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. """ - cos = cos[position_ids].unsqueeze(unsqueeze_dim) - sin = sin[position_ids].unsqueeze(unsqueeze_dim) + # cos = cos[position_ids].unsqueeze(unsqueeze_dim) + # sin = sin[position_ids].unsqueeze(unsqueeze_dim) b, h, s, d = q.shape q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) @@ -289,6 +297,8 @@ def fused_forward_h_blocking( use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, mla_absorption: Optional[Dict[str, bool]] = None, + cos_cached: Optional[torch.Tensor] = None, + sin_cached: Optional[torch.Tensor] = None, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: bsz, q_len, _ = hidden_states.size() @@ -310,8 +320,7 @@ def fused_forward_h_blocking( if compressed_kvs is not None: kva = compressed_kvs.update_ckv(kva, self.layer_idx, cache_kwargs) - cos, sin = self.rotary_emb(kva, seq_len=32 * 1024) - q_pe, k_pe = orig_apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids) + q_pe, k_pe = orig_apply_rotary_pos_emb(q_pe, k_pe, cos_cached, sin_cached) if compressed_kvs is not None: k_pe = compressed_kvs.update_k_pe(k_pe, self.layer_idx, cache_kwargs) @@ -326,6 +335,7 @@ def fused_forward_h_blocking( q_pe=q_pe, kva=kva, k_pe=k_pe, + batch_index=batch_index, per_head_q_up=self.per_head_q_up, per_head_k_up=self.per_head_k_up, per_head_v_up=self.per_head_v_up, @@ -356,10 +366,13 @@ def fused_forward_kv_blocking( use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, mla_absorption: Optional[Dict[str, bool]] = None, + cos_cached: Optional[torch.Tensor] = None, + sin_cached: Optional[torch.Tensor] = None, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: bsz, q_len, _ = hidden_states.size() + # -- KV compression (write to cache) ---------------------------------- compressed_kv = self.kv_a_proj_with_mqa(hidden_states) compressed_kv = compressed_kv.view(bsz, q_len, -1, self.kv_lora_rank + self.qk_rope_head_dim).transpose(1, 2) @@ -377,8 +390,7 @@ def fused_forward_kv_blocking( if compressed_kvs is not None: compressed_kvs.write_only_ckv(kva, self.layer_idx, cache_kwargs) - cos, sin = self.rotary_emb(hidden_states, seq_len=32 * 1024) - q_pe, k_pe = orig_apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids) + q_pe, k_pe = orig_apply_rotary_pos_emb(q_pe, k_pe, cos_cached, sin_cached) if compressed_kvs is not None: compressed_kvs.write_only_k_pe(k_pe, self.layer_idx, cache_kwargs) @@ -395,13 +407,11 @@ def fused_forward_kv_blocking( dq_qup_kupT = torch.matmul(q_a_proj_out, qup_kupT) else: dq_qup_kupT = torch.matmul(q_a_proj_out, self.fusedqk) - qkupTrope_nope = torch.cat((dq_qup_kupT, q_pe), dim=-1) - query = qkupTrope_nope + query = torch.cat((dq_qup_kupT, q_pe), dim=-1) # [B, num_heads, q_len, d_abs] else: q_nope = torch.bmm(q_a_proj_out, self.q_up) q_nope = q_nope.view(bsz, q_len, self.num_heads, self.qk_nope_head_dim).transpose(1, 2) - qnope_rope = torch.cat((q_nope, q_pe), dim=-1) - query = qnope_rope + query = torch.cat((q_nope, q_pe), dim=-1) blocking_config = getattr(self, "attn_blocking_config", AttentionBlockingConfig()) @@ -411,6 +421,7 @@ def fused_forward_kv_blocking( per_head_k_up_normal=self.per_head_k_up_normal, per_head_v_up=self.per_head_v_up, attention_mask=attention_mask, + batch_index=batch_index, scaling=self.softmax_scale, layer_idx=self.layer_idx, compressed_kvs=compressed_kvs, @@ -438,6 +449,8 @@ def fused_forward_orig( use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, mla_absorption: Optional[Dict[str, bool]] = None, + cos_cached: Optional[torch.Tensor] = None, + sin_cached: Optional[torch.Tensor] = None, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: bsz, q_len, _ = hidden_states.size() @@ -469,8 +482,7 @@ def fused_forward_orig( absorption = False # ---- Rotary ---- - cos, sin = self.rotary_emb(q_pe, seq_len=32 * 1024) # Doesn't need q_pe as head_dim is initialized - q_pe, k_pe = orig_apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids) + q_pe, k_pe = orig_apply_rotary_pos_emb(q_pe, k_pe, cos_cached, sin_cached) if compressed_kvs is not None: k_pe = compressed_kvs.update_k_pe(k_pe, self.layer_idx, cache_kwargs) @@ -520,7 +532,7 @@ def fused_forward_orig( self.k_up.squeeze(0).view(self.kv_lora_rank, self.num_heads, self.qk_nope_head_dim).permute(1, 0, 2) ) k_nope = torch.matmul(kva_expanded, k_up_per_head) - if k_heads <= 1: + if k_heads == 1: k_pe_expanded = ( k_pe_expanded.unsqueeze(1) .expand(-1, self.num_heads, -1, -1, -1) @@ -557,6 +569,8 @@ def fused_forward( use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, mla_absorption: Optional[Dict[str, bool]] = None, + cos_cached: Optional[torch.Tensor] = None, + sin_cached: Optional[torch.Tensor] = None, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: blocking_config = getattr(self, "attn_blocking_config", AttentionBlockingConfig()) @@ -573,6 +587,8 @@ def fused_forward( use_cache, cache_position, mla_absorption, + cos_cached, + sin_cached, **kwargs, ) elif getattr(blocking_config, "mode", None) == "kv": @@ -588,6 +604,8 @@ def fused_forward( use_cache, cache_position, mla_absorption, + cos_cached, + sin_cached, **kwargs, ) else: @@ -603,6 +621,8 @@ def fused_forward( use_cache, cache_position, mla_absorption, + cos_cached, + sin_cached, **kwargs, ) @@ -617,6 +637,8 @@ def forward_full_kv( output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, + cos_cached: Optional[torch.Tensor] = None, + sin_cached: Optional[torch.Tensor] = None, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: bsz, q_len, _ = hidden_states.size() @@ -644,8 +666,7 @@ def forward_full_kv( k_nope = kv[:, :, :, : self.qk_nope_head_dim] value_states = kv[:, :, :, self.qk_nope_head_dim :] - cos, sin = self.rotary_emb(value_states, seq_len=32 * 1024) - q_pe, k_pe = orig_apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids) + q_pe, k_pe = orig_apply_rotary_pos_emb(q_pe, k_pe, cos_cached, sin_cached) query_states = torch.cat((q_nope, q_pe), -1) k_pe_new = k_pe.expand(-1, self.num_heads, -1, -1) @@ -682,6 +703,8 @@ def forward_full_kv_h_blocking( output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, + cos_cached: Optional[torch.Tensor] = None, + sin_cached: Optional[torch.Tensor] = None, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: bsz, q_len, _ = hidden_states.size() @@ -712,8 +735,7 @@ def forward_full_kv_h_blocking( k_nope = kv[:, :, :, : self.qk_nope_head_dim] value_states = kv[:, :, :, self.qk_nope_head_dim :] - cos, sin = self.rotary_emb(value_states, seq_len=32 * 1024) - q_pe, k_pe = orig_apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids) + q_pe, k_pe = orig_apply_rotary_pos_emb(q_pe, k_pe, cos_cached, sin_cached) query_states = torch.cat((q_nope, q_pe), -1) k_pe_new = k_pe.expand(-1, self.num_heads, -1, -1) @@ -750,6 +772,8 @@ def forward( output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, + cos_cached: Optional[torch.Tensor] = None, + sin_cached: Optional[torch.Tensor] = None, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: blocking_config = getattr(self, "attn_blocking_config", AttentionBlockingConfig()) @@ -764,6 +788,8 @@ def forward( output_attentions, use_cache, cache_position, + cos_cached, + sin_cached, **kwargs, ) else: @@ -777,36 +803,172 @@ def forward( output_attentions, use_cache, cache_position, + cos_cached, + sin_cached, **kwargs, ) +EXPERT_BLOCKING_NUM_NSP = int(os.environ.get("EXPERT_BLOCKING_NUM_NSP", "4")) +EXPERT_BLOCKING_PACKED_CHUNK_SIZE = int(os.environ.get("EXPERT_BLOCKING_PACKED_CHUNK_SIZE", "256")) + + +def _build_matched_idx_from_cumsum(T2Ei: torch.Tensor) -> torch.Tensor: + """Build packed->original token index""" + batch_size, seq_len = T2Ei.shape + int32_max = torch.iinfo(torch.int32).max + int32_max_scalar = torch.tensor(int32_max, dtype=torch.int32, device=T2Ei.device) + token_idx = torch.arange(seq_len, dtype=torch.int32, device=T2Ei.device).unsqueeze(0).expand(batch_size, -1) + valid_prefix = torch.cumsum(T2Ei.to(torch.int32), dim=1) + valid_dest = valid_prefix - 1 + scatter_pos = torch.where(T2Ei, valid_dest, int32_max_scalar) + matched_idx = torch.full_like(token_idx, int32_max) + matched_idx = CtxScatterFunc3DInt.apply( + matched_idx.unsqueeze(-1), + scatter_pos, + token_idx.unsqueeze(-1), + ).squeeze(-1) + return matched_idx + + +class QEffDeepseekMoEGate(nn.Module): + def forward(self, hidden_states): + bsz, seq_len, h = hidden_states.shape + ### compute gating score + hidden_states = hidden_states.view(-1, h) + logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32), None) + if self.scoring_func == "sigmoid": + scores = logits.sigmoid() + else: + raise NotImplementedError(f"insupportable scoring function for MoE gating: {self.scoring_func}") + + ### select top-k experts + if self.topk_method == "noaux_tc": + assert not self.training + scores_for_choice = scores.view(bsz * seq_len, -1) + self.e_score_correction_bias.unsqueeze(0) + group_scores = torch.einsum( + "abc->ab", scores_for_choice.view(bsz * seq_len, self.n_group, -1).topk(2, dim=-1)[0] + ) # [n, n_group] + group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1] # [n, top_k_group] + group_mask = torch.zeros_like(group_scores) # [n, n_group] + group_mask.scatter_(1, group_idx, 1) # [n, n_group] + score_mask = ( + group_mask.unsqueeze(-1) + .expand(bsz * seq_len, self.n_group, self.n_routed_experts // self.n_group) + .reshape(bsz * seq_len, -1) + ) # [n, e] + tmp_scores = scores_for_choice.masked_fill(~score_mask.bool(), 0.0) # [n, e] + _, topk_idx = torch.topk(tmp_scores, k=self.top_k, dim=-1, sorted=False) + topk_weight = scores.gather(1, topk_idx) + else: + raise NotImplementedError(f"insupportable TopK function for MoE gating: {self.topk_method}") + + ### norm gate to sum 1 + if self.top_k > 1 and self.norm_topk_prob: + denominator = torch.einsum("bi->b", topk_weight).unsqueeze(-1) + 1e-20 + + topk_weight = topk_weight / denominator + topk_weight = topk_weight * self.routed_scaling_factor # must multiply the scaling factor + + router_probs = tmp_scores + router_weights = scores + return topk_idx, topk_weight, router_probs, router_weights + + class QEffDeepseekV3MoE(nn.Module): def __qeff_init__( self, ): - self.all_gate_proj = torch.nn.Parameter( - torch.cat( - [_get_linear_weight(exp.gate_proj).T.unsqueeze(0) for exp in self.experts], - dim=0, - ) + # Get common parameters from first expert + first_expert = self.experts[0] + self.bits = first_expert.gate_proj.bits + self.group_size = first_expert.gate_proj.group_size + self.act_fn = first_expert.act_fn + assert first_expert.gate_proj.act_order == first_expert.up_proj.act_order == first_expert.down_proj.act_order, ( + "act_order mismatch" ) - self.all_up_proj = torch.nn.Parameter( - torch.cat( - [_get_linear_weight(exp.up_proj).T.unsqueeze(0) for exp in self.experts], - dim=0, - ) + self.act_order = first_expert.gate_proj.act_order + + # Store dimensions for dequantization + self.in_features_gate, self.out_features_gate = ( + first_expert.gate_proj.in_features, + first_expert.gate_proj.out_features, ) - self.all_down_proj = torch.nn.Parameter( - torch.cat( - [_get_linear_weight(exp.down_proj).T.unsqueeze(0) for exp in self.experts], - dim=0, - ) + self.in_features_up, self.out_features_up = first_expert.up_proj.in_features, first_expert.up_proj.out_features + self.in_features_down, self.out_features_down = ( + first_expert.down_proj.in_features, + first_expert.down_proj.out_features, ) - self.act_fn = self.experts[0].act_fn + # Stack all parameters along a new dimension (expert dimension) + self.all_gate_qweight = torch.nn.Parameter( + torch.stack([exp.gate_proj.qweight for exp in self.experts], dim=0).reshape( + -1, self.out_features_gate, self.in_features_gate // 2 + ), + requires_grad=False, + ) + self.all_gate_scales = torch.nn.Parameter( + torch.stack([exp.gate_proj.scales for exp in self.experts], dim=0).reshape( + -1, self.out_features_gate, self.in_features_gate // self.group_size + ), + requires_grad=False, + ) + # TODO: Since we know qzeros is always 8 -> Just embed this once into the operator as parameter -> explore this later + self.all_gate_qzeros = torch.nn.Parameter( + torch.stack([exp.gate_proj.qzeros for exp in self.experts], dim=0).reshape( + -1, self.out_features_gate, self.in_features_gate // (self.group_size * 2) + ), + requires_grad=False, + ) + self.all_gate_gidx = torch.nn.Parameter( + torch.stack([exp.gate_proj.g_idx for exp in self.experts], dim=0), requires_grad=False + ) + + self.all_up_qweight = torch.nn.Parameter( + torch.stack([exp.up_proj.qweight for exp in self.experts], dim=0).reshape( + -1, self.out_features_up, self.in_features_up // 2 + ), + requires_grad=False, + ) + self.all_up_scales = torch.nn.Parameter( + torch.stack([exp.up_proj.scales for exp in self.experts], dim=0).reshape( + -1, self.out_features_up, self.in_features_up // self.group_size + ), + requires_grad=False, + ) + self.all_up_qzeros = torch.nn.Parameter( + torch.stack([exp.up_proj.qzeros for exp in self.experts], dim=0).reshape( + -1, self.out_features_up, self.in_features_up // (self.group_size * 2) + ), + requires_grad=False, + ) + self.all_up_gidx = torch.nn.Parameter( + torch.stack([exp.up_proj.g_idx for exp in self.experts], dim=0), requires_grad=False + ) + + self.all_down_qweight = torch.nn.Parameter( + torch.stack([exp.down_proj.qweight for exp in self.experts], dim=0).reshape( + -1, self.out_features_down, self.in_features_down // 2 + ), + requires_grad=False, + ) + self.all_down_scales = torch.nn.Parameter( + torch.stack([exp.down_proj.scales for exp in self.experts], dim=0).reshape( + -1, self.out_features_down, self.in_features_down // self.group_size + ), + requires_grad=False, + ) + self.all_down_qzeros = torch.nn.Parameter( + torch.stack([exp.down_proj.qzeros for exp in self.experts], dim=0).reshape( + -1, self.out_features_down, self.in_features_down // (self.group_size * 2) + ), + requires_grad=False, + ) + self.all_down_gidx = torch.nn.Parameter( + torch.stack([exp.down_proj.g_idx for exp in self.experts], dim=0), requires_grad=False + ) - def moe( + def moe_old( self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, @@ -816,29 +978,195 @@ def moe( hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) final_hidden_states = torch.zeros_like(hidden_states, dtype=topk_weights.dtype) - gate_proj = self.all_gate_proj[topk_indices.flatten()] - up_proj = self.all_up_proj[topk_indices.flatten()] - down_proj = self.all_down_proj[topk_indices.flatten()] + for i in range(self.gate.top_k): + expert_idx = topk_indices[:, i] + curr_weight = topk_weights[:, i] + gate_qweight = self.all_gate_qweight[expert_idx].reshape( + seq_len * self.out_features_gate, + self.in_features_gate // self.group_size, + (self.group_size * self.bits) // 8, + ) + gate_scales = self.all_gate_scales[expert_idx].reshape( + seq_len * self.out_features_gate * (self.in_features_gate // self.group_size) + ) + gate_qzeros = self.all_gate_qzeros[expert_idx].reshape( + seq_len * self.out_features_gate, self.in_features_gate // self.group_size + ) + gate_gidx = self.all_gate_gidx[expert_idx].reshape(seq_len * self.in_features_gate) + + up_qweight = self.all_up_qweight[expert_idx].reshape( + seq_len * self.out_features_up, + self.in_features_up // self.group_size, + (self.group_size * self.bits) // 8, + ) + up_scales = self.all_up_scales[expert_idx].reshape( + seq_len * self.out_features_up * (self.in_features_up // self.group_size) + ) + up_qzeros = self.all_up_qzeros[expert_idx].reshape( + seq_len * self.out_features_up, self.in_features_up // self.group_size + ) + up_gidx = self.all_up_gidx[expert_idx].reshape(seq_len * self.in_features_up) + + down_qweight = self.all_down_qweight[expert_idx].reshape( + seq_len * self.out_features_down, + self.in_features_down // self.group_size, + (self.group_size * self.bits) // 8, + ) + down_scales = self.all_down_scales[expert_idx].reshape( + seq_len * self.out_features_down * (self.in_features_down // self.group_size) + ) + down_qzeros = self.all_down_qzeros[expert_idx].reshape( + seq_len * self.out_features_down, self.in_features_down // self.group_size + ) + down_gidx = self.all_down_gidx[expert_idx].reshape(seq_len * self.in_features_down) + + gate_out = QuantLinearTorchFunction.apply( + hidden_states, + gate_qweight, + gate_scales, + gate_qzeros, + gate_gidx if self.act_order else None, + self.bits, + self.group_size, + self.in_features_gate, + self.out_features_gate * seq_len, + ) + + up_out = QuantLinearTorchFunction.apply( + hidden_states, + up_qweight, + up_scales, + up_qzeros, + up_gidx if self.act_order else None, + self.bits, + self.group_size, + self.in_features_up, + self.out_features_up * seq_len, + ) + + hidden = self.act_fn(gate_out) * up_out + down_out = QuantLinearTorchFunction.apply( + hidden, + down_qweight, + down_scales, + down_qzeros, + down_gidx if self.act_order else None, + self.bits, + self.group_size, + self.in_features_down, + self.out_features_down, + ) + down_out = down_out.reshape(seq_len, self.out_features_down) + final_hidden_states += down_out * curr_weight.unsqueeze(1) + + return final_hidden_states + + def moe_weights_as_activations(self, hidden_states, router_probs, router_weights): + return QMOE.apply( + hidden_states, + router_weights, + self.fc1_experts_weights, + self.fc1_scales, + self.fc2_experts_weights, + self.fc2_scales, + self.fc3_experts_weights, + self.fc3_scales, + router_probs, + self.config.hidden_act, + self.group_size, + self.bits, + self.num_experts_per_tok, + ) + + @torch.no_grad() + def original_moe(self, x, topk_ids, topk_weight): + cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts))) + cnts.scatter_(1, topk_ids, 1) + tokens_per_expert = cnts.sum(dim=0) + idxs = topk_ids.view(-1).argsort() + sorted_tokens = x[idxs // topk_ids.shape[1]] + # sorted_tokens_shape = sorted_tokens.shape + tokens_per_expert = tokens_per_expert.cpu().numpy() + + outputs = [] + start_idx = 0 + for i, num_tokens in enumerate(tokens_per_expert): + end_idx = start_idx + num_tokens + if num_tokens == 0: + continue + expert = self.experts[i + self.ep_rank * self.experts_per_rank] + tokens_for_this_expert = sorted_tokens[start_idx:end_idx] + expert_out = expert(tokens_for_this_expert) + outputs.append(expert_out) + start_idx = end_idx + + outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) + + new_x = torch.empty_like(outs) + new_x[idxs] = outs + final_out = ( + new_x.view(*topk_ids.shape, -1) + .type(topk_weight.dtype) + .mul_(topk_weight.unsqueeze(dim=-1)) + .sum(dim=1) + .type(new_x.dtype) + ) + return final_out + + def moe_waa_unpack(self, hidden_states, topk_indices, topk_weights): + # GATHER - collect weights for selected experts + gate_proj_qweight = self.all_gate_qweight[topk_indices.flatten()] + gate_proj_scales = self.all_gate_scales[topk_indices.flatten()] + gate_proj_qzeros = self.all_gate_qzeros[topk_indices.flatten()] + + up_proj_qweight = self.all_up_qweight[topk_indices.flatten()] + up_proj_scales = self.all_up_scales[topk_indices.flatten()] + up_proj_qzeros = self.all_up_qzeros[topk_indices.flatten()] + + down_proj_qweight = self.all_down_qweight[topk_indices.flatten()] + down_proj_scales = self.all_down_scales[topk_indices.flatten()] + down_proj_qzeros = self.all_down_qzeros[topk_indices.flatten()] + + gate_proj_unpacked = CastToUInt4Func.apply(gate_proj_qweight) + gate_zeros_unpacked = CastToUInt4Func.apply(gate_proj_qzeros) + gate_proj_dq = DequantizeLinearFunc.apply( + gate_proj_unpacked, gate_proj_scales, gate_zeros_unpacked, self.group_size + ) + + up_proj_unpacked = CastToUInt4Func.apply(up_proj_qweight) + up_zeros_unpacked = CastToUInt4Func.apply(up_proj_qzeros) + up_proj_dq = DequantizeLinearFunc.apply(up_proj_unpacked, up_proj_scales, up_zeros_unpacked, self.group_size) + + down_proj_unpacked = CastToUInt4Func.apply(down_proj_qweight) + down_zeros_unpacked = CastToUInt4Func.apply(down_proj_qzeros) + down_proj_dq = DequantizeLinearFunc.apply( + down_proj_unpacked, down_proj_scales, down_zeros_unpacked, self.group_size + ) + + # Reshape for bmm: (bs*seq_len*top_k, 1, hidden_size) expert_in = ( - hidden_states.unsqueeze(1).expand(-1, self.gate.top_k, -1).contiguous().view(-1, 1, self.config.hidden_size) + hidden_states.unsqueeze(1).expand(-1, self.gate.top_k, -1).contiguous().view(-1, 1, self.in_features_gate) ) - gate_out = torch.bmm(expert_in, gate_proj) - up_out = torch.bmm(expert_in, up_proj) + + gate_out = torch.bmm(expert_in, gate_proj_dq.transpose(1, 2).to(expert_in.dtype)) + up_out = torch.bmm(expert_in, up_proj_dq.transpose(1, 2).to(expert_in.dtype)) hidden = self.act_fn(gate_out) * up_out - expert_output = torch.bmm(hidden, down_proj) - experts_out = expert_output.view(seq_len, self.gate.top_k, self.config.hidden_size) - experts_out = experts_out * topk_weights.unsqueeze(-1) + down_out = torch.bmm(hidden, down_proj_dq.transpose(1, 2).to(expert_in.dtype)) + + down_out = down_out.view(-1, self.gate.top_k, self.out_features_down) - final_hidden_states = torch.einsum("abc->ac", experts_out) + down_out = down_out * topk_weights.unsqueeze(-1) - return final_hidden_states.type(hidden_states.dtype) + return torch.einsum("abc-> ac", down_out) def forward(self, hidden_states): + print("Using new MoE forward with weights as activations") residuals = hidden_states orig_shape = hidden_states.shape - topk_indices, topk_weights = self.gate(hidden_states) + topk_indices, topk_weights, router_probs, router_weights = self.gate(hidden_states) hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) - hidden_states = self.moe(hidden_states, topk_indices, topk_weights).view(*orig_shape) + # hidden_states = self.moe_weights_as_activations(hidden_states, router_probs, router_weights).view(*orig_shape) + hidden_states = self.moe_waa_unpack(hidden_states, topk_indices, topk_weights).view(*orig_shape) hidden_states = hidden_states + self.shared_experts(residuals) return hidden_states @@ -847,83 +1175,333 @@ class QEffPrefillOnlyDeepseekV3MoE(nn.Module): def __qeff_init__( self, ): - for exp in self.experts: - gate_proj = torch.nn.Linear(self.config.hidden_size, self.config.moe_intermediate_size, bias=False) - up_proj = torch.nn.Linear(self.config.hidden_size, self.config.moe_intermediate_size, bias=False) - down_proj = torch.nn.Linear(self.config.moe_intermediate_size, self.config.hidden_size, bias=False) + # Get common parameters from first expert + first_expert = self.experts[0] + self.bits = first_expert.gate_proj.bits + self.group_size = first_expert.gate_proj.group_size + self.act_fn = first_expert.act_fn + assert first_expert.gate_proj.act_order == first_expert.up_proj.act_order == first_expert.down_proj.act_order, ( + "act_order mismatch" + ) + self.act_order = first_expert.gate_proj.act_order - gate_proj.weight = torch.nn.Parameter(_get_linear_weight(exp.gate_proj).detach().clone()) - up_proj.weight = torch.nn.Parameter(_get_linear_weight(exp.up_proj).detach().clone()) - down_proj.weight = torch.nn.Parameter(_get_linear_weight(exp.down_proj).detach().clone()) + # Store dimensions for dequantization + self.in_features_gate, self.out_features_gate = ( + first_expert.gate_proj.in_features, + first_expert.gate_proj.out_features, + ) + self.in_features_up, self.out_features_up = first_expert.up_proj.in_features, first_expert.up_proj.out_features + self.in_features_down, self.out_features_down = ( + first_expert.down_proj.in_features, + first_expert.down_proj.out_features, + ) - setattr(exp, "gate_proj", gate_proj) - setattr(exp, "up_proj", up_proj) - setattr(exp, "down_proj", down_proj) + # Stack all parameters along a new dimension (expert dimension) + self.all_gate_qweight = torch.nn.Parameter( + torch.stack([exp.gate_proj.qweight for exp in self.experts], dim=0).reshape( + -1, self.out_features_gate, self.in_features_gate // 2 + ), + requires_grad=False, + ) + self.all_gate_scales = torch.nn.Parameter( + torch.stack([exp.gate_proj.scales for exp in self.experts], dim=0) + .reshape(-1, self.out_features_gate, self.in_features_gate // self.group_size) + .to(torch.float16), + requires_grad=False, + ) + # TODO: Since we know qzeros is always 8 -> Just embed this once into the operator as parameter -> explore this later + self.all_gate_qzeros = torch.nn.Parameter( + torch.stack([exp.gate_proj.qzeros for exp in self.experts], dim=0).reshape( + -1, self.out_features_gate, self.in_features_gate // (self.group_size * 2) + ), + requires_grad=False, + ) + self.all_gate_gidx = torch.nn.Parameter( + torch.stack([exp.gate_proj.g_idx for exp in self.experts], dim=0), requires_grad=False + ) - def moe(self, hidden_states: torch.Tensor, topk_weights: torch.Tensor, expert_mask: torch.Tensor, num_experts: int): - final_hidden_states = torch.zeros_like(hidden_states, dtype=topk_weights.dtype) - for expert_idx in range(num_experts): - expert = self.experts[expert_idx] - gate_out = expert.gate_proj(hidden_states) - up_out = expert.up_proj(hidden_states) - hidden = expert.act_fn(gate_out) * up_out - expert_output = expert.down_proj(hidden) - current_hidden_states = expert_output * expert_mask[:, expert_idx].unsqueeze(-1) - final_hidden_states += current_hidden_states + self.all_up_qweight = torch.nn.Parameter( + torch.stack([exp.up_proj.qweight for exp in self.experts], dim=0).reshape( + -1, self.out_features_up, self.in_features_up // 2 + ), + requires_grad=False, + ) + self.all_up_scales = torch.nn.Parameter( + torch.stack([exp.up_proj.scales for exp in self.experts], dim=0) + .reshape(-1, self.out_features_up, self.in_features_up // self.group_size) + .to(torch.float16), + requires_grad=False, + ) + self.all_up_qzeros = torch.nn.Parameter( + torch.stack([exp.up_proj.qzeros for exp in self.experts], dim=0).reshape( + -1, self.out_features_up, self.in_features_up // (self.group_size * 2) + ), + requires_grad=False, + ) + self.all_up_gidx = torch.nn.Parameter( + torch.stack([exp.up_proj.g_idx for exp in self.experts], dim=0), requires_grad=False + ) - return final_hidden_states.type(hidden_states.dtype) + self.all_down_qweight = torch.nn.Parameter( + torch.stack([exp.down_proj.qweight for exp in self.experts], dim=0).reshape( + -1, self.out_features_down, self.in_features_down // 2 + ), + requires_grad=False, + ) + self.all_down_scales = torch.nn.Parameter( + torch.stack([exp.down_proj.scales for exp in self.experts], dim=0) + .reshape(-1, self.out_features_down, self.in_features_down // self.group_size) + .to(torch.float16), + requires_grad=False, + ) + self.all_down_qzeros = torch.nn.Parameter( + torch.stack([exp.down_proj.qzeros for exp in self.experts], dim=0).reshape( + -1, self.out_features_down, self.in_features_down // (self.group_size * 2) + ), + requires_grad=False, + ) + self.all_down_gidx = torch.nn.Parameter( + torch.stack([exp.down_proj.g_idx for exp in self.experts], dim=0), requires_grad=False + ) - def orig_moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): - r""" - CALL FOR CONTRIBUTION! I don't have time to optimise this right now, but expert weights need to be fused - to not have to do a loop here (deepseek has 256 experts soooo yeah). + def _cumsum_scatter_gather_update_expert_blocked( + self, + x: torch.Tensor, + T2Ei: torch.Tensor, + slot_gate_qweight: torch.Tensor, + slot_gate_scales: torch.Tensor, + slot_gate_qzeros: torch.Tensor, + slot_up_qweight: torch.Tensor, + slot_up_scales: torch.Tensor, + slot_up_qzeros: torch.Tensor, + slot_down_qweight: torch.Tensor, + slot_down_scales: torch.Tensor, + slot_down_qzeros: torch.Tensor, + routing_weight: torch.Tensor, + expert_out: torch.Tensor, + packed_chunk_size: int, + num_q_ffn_blocks: Optional[int] = None, + ) -> torch.Tensor: + """Cumsum-scatter-gather-update expert helper for NSP-blocked dispatch. + + Accumulates one local expert's contribution in-place onto ``expert_out``. + Uses a packed/cumsum layout so the MLP runs only over active rows, then + scatters the weighted output back to original token positions. + + Shapes: + x : [T, H] + T2Ei : [num_nsp, T] (bool) + W_g, W_u : [num_nsp, H, I] + W_d : [num_nsp, I, H] + routing_weight : [num_nsp, T] + expert_out : [num_nsp, T, H] (accumulator, in-out) """ - hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) - final_hidden_states = torch.zeros_like(hidden_states, dtype=topk_weights.dtype) - expert_mask = torch.nn.functional.one_hot(topk_indices, num_classes=len(self.experts)) - expert_mask = expert_mask.permute(2, 0, 1) - for expert_idx in range(len(self.experts)): - expert = self.experts[expert_idx] - mask = expert_mask[expert_idx] - token_indices, weight_indices = torch.where(mask) + batch_size, seq_len = T2Ei.shape + packed_chunk_size = max(1, min(packed_chunk_size, seq_len)) - if token_indices.numel() > 0: - expert_weights = topk_weights[token_indices, weight_indices] - expert_input = hidden_states[token_indices] - expert_output = expert(expert_input) - weighted_output = expert_output * expert_weights.unsqueeze(-1) - final_hidden_states.index_add_(0, token_indices, weighted_output) + if num_q_ffn_blocks is not None: + assert seq_len % num_q_ffn_blocks == 0, "Something went wrong" + packed_chunk_size = seq_len // num_q_ffn_blocks + else: + num_q_ffn_blocks = seq_len // packed_chunk_size + import ipdb + + ipdb.set_trace() + matched_idx = _build_matched_idx_from_cumsum(T2Ei) + valid_rows = torch.einsum("bi->b", T2Ei.to(torch.int32)).unsqueeze(1) + row_range = torch.arange(packed_chunk_size, dtype=torch.int32, device=x.device).unsqueeze(0) + x_expanded = x.unsqueeze(0).expand(batch_size, -1, -1) + + for chunk_idx in range(num_q_ffn_blocks): + print("executing chunk", chunk_idx) + packed_start = chunk_idx * packed_chunk_size + if chunk_idx == num_q_ffn_blocks - 1: + packed_stop = seq_len + else: + packed_stop = packed_start + packed_chunk_size - # in original deepseek, the output of the experts are gathered once we leave this module - # thus the moe module is itelsf an IsolatedParallel module - # and all expert are "local" meaning we shard but we don't gather - return final_hidden_states.type(hidden_states.dtype) + chunk_matched_idx = matched_idx[:, packed_start:packed_stop] - def forward(self, hidden_states): - """ - Forward pass of MoE block. - """ - residuals = hidden_states - orig_shape = hidden_states.shape - topk_indices, topk_weights = self.gate(hidden_states) - # orig_out = self.orig_moe(hidden_states, topk_indices, topk_weights).view(*orig_shape) + x_chunk = CtxGatherFunc3DGeneralized.apply(x_expanded, chunk_matched_idx) - hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) - mask = torch.zeros(hidden_states.shape[0], self.config.n_routed_experts) - mask.scatter_(1, topk_indices, topk_weights) - if os.environ.get("NUM_FFN_BLOCKS", None) is not None and os.environ.get("FFN_W_BLOCK_SIZE", None) is not None: - hidden_states = self.moe_blocked_weights_forward( - hidden_states, topk_weights, mask, self.config.n_routed_experts - ).view(*orig_shape) - elif os.environ.get("NUM_FFN_BLOCKS", None) is not None: - hidden_states = self.moe_blocked_forward( - hidden_states, topk_weights, mask, self.config.n_routed_experts - ).view(*orig_shape) - else: - hidden_states = self.moe(hidden_states, topk_weights, mask, self.config.n_routed_experts).view(*orig_shape) + gate_proj_unpacked = CastToUInt4Func.apply(slot_gate_qweight) + gate_zeros_unpacked = CastToUInt4Func.apply(slot_gate_qzeros) + gate_proj_dq = DequantizeLinearFunc.apply( + gate_proj_unpacked, slot_gate_scales, gate_zeros_unpacked, self.group_size + ) - hidden_states = hidden_states + self.shared_experts(residuals) - return hidden_states + up_proj_unpacked = CastToUInt4Func.apply(slot_up_qweight) + up_zeros_unpacked = CastToUInt4Func.apply(slot_up_qzeros) + up_proj_dq = DequantizeLinearFunc.apply( + up_proj_unpacked, slot_up_scales, up_zeros_unpacked, self.group_size + ) + + down_proj_unpacked = CastToUInt4Func.apply(slot_down_qweight) + down_zeros_unpacked = CastToUInt4Func.apply(slot_down_qzeros) + + down_proj_dq = DequantizeLinearFunc.apply( + down_proj_unpacked, slot_down_scales, down_zeros_unpacked, self.group_size + ) + + gate_out = torch.bmm(x_chunk, gate_proj_dq.transpose(1, 2).to(x_chunk.dtype)) + up_out = torch.bmm(x_chunk, up_proj_dq.transpose(1, 2).to(x_chunk.dtype)) + hidden = self.act_fn(gate_out) * up_out + down_out = torch.bmm(hidden, down_proj_dq.transpose(1, 2).to(x_chunk.dtype)) + + rw_chunk = CtxGatherFunc3DGeneralized.apply(routing_weight, chunk_matched_idx) + old_expert_out = CtxGatherFunc3DGeneralized.apply(expert_out, chunk_matched_idx) + chunk_valid_rows = torch.clamp(valid_rows - packed_start, min=0, max=packed_chunk_size) + current_expert_out = ( + torch.where( + (row_range < chunk_valid_rows).unsqueeze(-1), + down_out, + torch.zeros_like(down_out), + ) + * rw_chunk + ) + updated_chunk = old_expert_out + current_expert_out + expert_out = CtxScatterFunc3DGeneralized.apply(expert_out, chunk_matched_idx, updated_chunk) + + return expert_out + + def _forward_expert_blocked( + self, + x: torch.Tensor, + local_T2E: torch.Tensor, + routing_weights: torch.Tensor, + num_q_ffn_blocks: Optional[int] = None, + ) -> torch.Tensor: + T, H = x.shape + num_nsp = EXPERT_BLOCKING_NUM_NSP + if len(self.experts) % num_nsp != 0: + raise ValueError( + f"num_experts ({len(self.experts)}) must be divisible by EXPERT_BLOCKING_NUM_NSP ({num_nsp})" + ) + local_experts = len(self.experts) // num_nsp + routing_weights_unsqueezed = routing_weights.unsqueeze(-1) + expert_out = x.new_zeros((num_nsp, T, H)) + + local_gate_qweight = ( + self.all_gate_qweight.view(local_experts, num_nsp, self.out_features_gate, self.in_features_gate // 2) + .transpose(0, 1) + .contiguous() + ) + local_gate_scales = ( + self.all_gate_scales.view( + local_experts, num_nsp, self.out_features_gate, self.in_features_gate // self.group_size + ) + .transpose(0, 1) + .contiguous() + ) + local_gate_qzeros = ( + self.all_gate_qzeros.view( + local_experts, num_nsp, self.out_features_gate, self.in_features_gate // (self.group_size * 2) + ) + .transpose(0, 1) + .contiguous() + ) + + local_up_qweight = ( + self.all_up_qweight.view(local_experts, num_nsp, self.out_features_up, self.in_features_up // 2) + .transpose(0, 1) + .contiguous() + ) + local_up_scales = ( + self.all_up_scales.view( + local_experts, num_nsp, self.out_features_up, self.in_features_up // self.group_size + ) + .transpose(0, 1) + .contiguous() + ) + local_up_qzeros = ( + self.all_up_qzeros.view( + local_experts, num_nsp, self.out_features_up, self.in_features_up // (self.group_size * 2) + ) + .transpose(0, 1) + .contiguous() + ) + + local_down_qweight = ( + self.all_down_qweight.view(local_experts, num_nsp, self.out_features_down, self.in_features_down // 2) + .transpose(0, 1) + .contiguous() + ) + local_down_scales = ( + self.all_down_scales.view( + local_experts, num_nsp, self.out_features_down, self.in_features_down // self.group_size + ) + .transpose(0, 1) + .contiguous() + ) + local_down_qzeros = ( + self.all_down_qzeros.view( + local_experts, num_nsp, self.out_features_down, self.in_features_down // (self.group_size * 2) + ) + .transpose(0, 1) + .contiguous() + ) + for slot in range(local_experts): + print(f"executing slot {slot}") + T2Ei = local_T2E[:, slot, :] + expert_out = self._cumsum_scatter_gather_update_expert_blocked( + x=x, + T2Ei=T2Ei, + slot_gate_qweight=local_gate_qweight[:, slot], + slot_gate_scales=local_gate_scales[:, slot], + slot_gate_qzeros=local_gate_qzeros[:, slot], + slot_up_qweight=local_up_qweight[:, slot], + slot_up_scales=local_up_scales[:, slot], + slot_up_qzeros=local_up_qzeros[:, slot], + slot_down_qweight=local_down_qweight[:, slot], + slot_down_scales=local_down_scales[:, slot], + slot_down_qzeros=local_down_qzeros[:, slot], + routing_weight=routing_weights_unsqueezed[:, slot], + expert_out=expert_out, + packed_chunk_size=EXPERT_BLOCKING_PACKED_CHUNK_SIZE, + num_q_ffn_blocks=num_q_ffn_blocks, + ) + return torch.einsum("bij->ij", expert_out) + + def forward( + self, hidden_states: torch.Tensor, num_q_ffn_blocks: Optional[int] = None + ) -> tuple[torch.Tensor, torch.Tensor]: + topk_idx, topk_weight, _, _ = self.gate(hidden_states) + B, S, H = hidden_states.shape + T = B * S + x = hidden_states.view(T, H) + + if len(self.experts) % EXPERT_BLOCKING_NUM_NSP == 0: + expert_ids = torch.arange( + len(self.experts) // EXPERT_BLOCKING_NUM_NSP, + device=x.device, + dtype=topk_idx.dtype, + ).unsqueeze(0) * EXPERT_BLOCKING_NUM_NSP + torch.arange( + EXPERT_BLOCKING_NUM_NSP, device=x.device, dtype=topk_idx.dtype + ).unsqueeze(1) # [N, L] + eq = topk_idx.unsqueeze(0).unsqueeze(0) == expert_ids.unsqueeze(-1).unsqueeze(-1) + local_T2E = eq.to(topk_idx.dtype).sum(dim=-1) > 0 # [N, L, T] + routing_weights = (eq.to(topk_weight.dtype) * topk_weight.unsqueeze(0).unsqueeze(0)).sum(dim=-1) + + expert_out = self._forward_expert_blocked( + x=x, local_T2E=local_T2E, routing_weights=routing_weights, num_q_ffn_blocks=num_q_ffn_blocks + ) + self.shared_experts(hidden_states) + return expert_out.view(B, S, H) + + routing_weights = torch.zeros(T, self.config.n_routed_experts, device=x.device, dtype=topk_weight.dtype) + routing_weights.scatter_(1, topk_idx, topk_weight) + + final_hidden_states = x.new_zeros((T, H)) + for expert_idx in range(self.config.n_routed_experts): + expert = self.experts[expert_idx] + gate_out = expert.gate_proj(hidden_states) + up_out = expert.up_proj(hidden_states) + hidden = expert.act_fn(gate_out) * up_out + expert_output = expert.down_proj(hidden) + current_hidden_states = expert_output * routing_weights[:, expert_idx].unsqueeze(-1) + final_hidden_states = final_hidden_states.view(B, S, H) + current_hidden_states + + final_hidden_states = final_hidden_states + self.shared_experts(hidden_states) + return final_hidden_states.view(B, S, H) class QEffDeepseekV3DecoderLayer(nn.Module): @@ -942,6 +1520,9 @@ def forward( cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, mla_absorption: Optional[Dict[str, bool]] = None, + num_q_ffn_blocks: Optional[int] = None, + sin_cached=None, + cos_cached=None, **kwargs, ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: residual = hidden_states @@ -963,6 +1544,8 @@ def forward( use_cache=use_cache, cache_position=cache_position, mla_absorption=mla_absorption, + sin_cached=sin_cached, + cos_cached=cos_cached, **kwargs, ) else: @@ -976,13 +1559,18 @@ def forward( output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, + sin_cached=sin_cached, + cos_cached=cos_cached, **kwargs, ) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) + if num_q_ffn_blocks is not None and self.mlp.__class__.__name__ == "DeepseekV3MoE": + hidden_states = self.mlp(hidden_states, num_q_ffn_blocks) + else: + hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states outputs = (hidden_states,) @@ -1021,6 +1609,8 @@ def __qeff_init__(self): base=self.config.rope_theta, **kwargs, ) + self.sin_cached = torch.nn.Parameter(self.rotary_emb.sin_cached.to(torch.float16)) + self.cos_cached = torch.nn.Parameter(self.rotary_emb.cos_cached.to(torch.float16)) def forward( self, @@ -1075,7 +1665,8 @@ def forward( if position_ids is None: position_ids = cache_position.unsqueeze(0) - causal_mask = _create_causal_mask(position_ids=position_ids, target_length=target_len) + # ctx_len = compressed_kvs.layers[0].ckv.shape[-2] + causal_mask = _create_causal_mask(position_ids=position_ids, target_length=target_len) # ctx_len) hidden_states = inputs_embeds position_embeddings = None @@ -1083,6 +1674,9 @@ def forward( all_self_attns = () if output_attentions else None next_decoder_cache = None + sin = self.sin_cached[position_ids].unsqueeze(1) + cos = self.cos_cached[position_ids].unsqueeze(1) + num_q_ffn_blocks = getattr(self, "num_q_blocks_ffn", None) for decoder_layer in self.layers: if output_hidden_states: all_hidden_states += (hidden_states,) @@ -1099,6 +1693,9 @@ def forward( cache_position=cache_position, position_embeddings=position_embeddings, mla_absorption=mla_absorption, + num_q_ffn_blocks=num_q_ffn_blocks, + sin_cached=sin, + cos_cached=cos, **kwargs, ) diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index 93c7d2eba9..f068e7b495 100755 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -71,6 +71,7 @@ FP8DeQuantLinearToLinearTransform, GPTQToMatmulNbitsTransform, Mxfp4GptOssExpertDequantizeTransform, + PackQuantizedInt4ToMatMulNBitsTransform, ) from QEfficient.utils import ( apply_kv_cache_prefix, @@ -1070,6 +1071,7 @@ class QEffVisionEncoderForTextImageToTextModel(QEFFBaseModel): _pytorch_transforms = [ AwqToMatmulNbitsTransform, GPTQToMatmulNbitsTransform, + PackQuantizedInt4ToMatMulNBitsTransform, CustomOpsTransform, KVCacheTransform, KVCacheExternalModuleMapperTransform, @@ -1204,6 +1206,7 @@ class QEffCausalLMForTextImageToTextModel(QEFFBaseModel): _pytorch_transforms = [ AwqToMatmulNbitsTransform, GPTQToMatmulNbitsTransform, + PackQuantizedInt4ToMatMulNBitsTransform, FP8BlockWiseDequantQwen3VLMoeTextExpertsToQwen3VLMoeTextExpertsTransform, FP8BlockWiseDequantLinearToLinearTransform, CustomOpsTransform, @@ -3381,6 +3384,7 @@ class QEFFAutoModelForCausalLM(QEFFBaseModel): AwqToMatmulNbitsTransform, GPTQToMatmulNbitsTransform, FP8DeQuantLinearToLinearTransform, + PackQuantizedInt4ToMatMulNBitsTransform, Mxfp4GptOssExpertDequantizeTransform, CustomOpsTransform, KVCacheTransform, @@ -3664,7 +3668,9 @@ def get_seq_len_and_handle_specialized_prefill_model( self.hash_params["moe_prefill_num_packed_chunks"] = num_packed_chunks if self.model.config.model_type in {"qwen3_moe", "gpt_oss", "glm4_moe"}: return max(prefill_seq_len or 0, constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN) - return constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN + seq_len = max(prefill_seq_len or 0, constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN) + self.hash_params["chunking_seq_len"] = seq_len + return seq_len num_q_blocks = ( self.hash_params["blocking_config"].num_q_blocks if self.hash_params.get("blocking_kwargs", None) else None @@ -3818,13 +3824,13 @@ def export( # TODO: move this to a DA Serving utility class if self.model.config.model_type in SPECIALIZED_DISAGG_SERVING_MODEL_ARCH: if prefill_only: - if not enable_chunking and self.continuous_batching: - raise NotImplementedError( - "Looks like you are trying to run prefix-caching without chunking, this feature is not available yet!" - ) self.__update_prefill_transform(enable=True, enable_chunking=enable_chunking) self.hash_params.pop("retain_full_kv", None) if "DeepseekV3ForCausalLM" not in (getattr(self.model.config, "architectures", None) or []): + if not enable_chunking and self.continuous_batching: + raise NotImplementedError( + "Looks like you are trying to run prefix-caching without chunking, this feature is not available yet!" + ) seq_len = self.get_seq_len_and_handle_specialized_prefill_model( prefill_seq_len=prefill_seq_len, enable_chunking=enable_chunking, @@ -3835,6 +3841,10 @@ def export( kv_cache_shape[2] = ( seq_len + (sliding_window if sliding_window is not None else 0) if enable_chunking else seq_len ) + else: + seq_len = max(prefill_seq_len or 0, constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN) + self.hash_params["chunking_seq_len"] = seq_len + else: self.__update_prefill_transform(False, retain_full_kv=kwargs.get("retain_full_kv", False)) self.hash_params.pop("prefill_only", None) @@ -3945,8 +3955,14 @@ def export( example_inputs["compressed_kvs"][i].append( torch.zeros(pkv_cache[0][1].shape, dtype=self.model.config.torch_dtype) ) - dynamic_axes[f"compressed_kv.{i}"] = {0: "batch_size", 2: "ctx_len"} - dynamic_axes[f"k_pe.{i}"] = {0: "batch_size", 2: "ctx_len"} + dynamic_axes[f"compressed_kv.{i}"] = { + 0: "full_batch_size" if self.continuous_batching else "batch_size", + 2: "ctx_len", + } + dynamic_axes[f"k_pe.{i}"] = { + 0: "full_batch_size" if self.continuous_batching else "batch_size", + 2: "ctx_len", + } output_names.append(f"compressed_kv.{i}_RetainedState") output_names.append(f"k_pe.{i}_RetainedState") else: @@ -4034,6 +4050,12 @@ def _legacyify_cache(obj): output_names = apply_kv_cache_prefix(output_names, kv_cache_prefix) self.hash_params["kv_cache_prefix"] = kv_cache_prefix + if prefill_only: + assert prefill_seq_len is not None, "prefill_seq_len must be provided when prefill_only is True" + num_q_blocks_ffn = prefill_seq_len // constants.EXPERT_BLOCKING_PACKED_CHUNK_SIZE + num_q_blocks_ffn = num_q_blocks_ffn if num_q_blocks_ffn > 0 else 1 + setattr(self.model.model, "num_q_blocks_ffn", num_q_blocks_ffn) + if QEFFBaseModel._layerwise_active: return self._export_layerwise( example_inputs, diff --git a/QEfficient/transformers/models/pytorch_transforms.py b/QEfficient/transformers/models/pytorch_transforms.py index 12fec63085..e4ca5e49f2 100755 --- a/QEfficient/transformers/models/pytorch_transforms.py +++ b/QEfficient/transformers/models/pytorch_transforms.py @@ -321,6 +321,7 @@ QEffDisentangledSelfAttention, ) from QEfficient.transformers.models.deepseek_v3.modeling_deepseek import ( + QEffDeepseekMoEGate, QEffDeepseekV3Attention, QEffDeepseekV3CustomRMSNormAIC, QEffDeepseekV3DecoderLayer, @@ -1051,6 +1052,8 @@ def apply(cls, model: nn.Module, num_kv_heads_repeat: int = 1) -> nn.Module: transformed = True for block in text_model.layers: attn = getattr(block, "cross_attn", getattr(block, "self_attn", None)) + if not attn: + continue attn.num_key_value_heads = new_kv_heads head_dim = attn.kv_lora_rank + attn.qk_rope_head_dim @@ -1263,8 +1266,11 @@ class KVCacheExternalModuleMapperTransform(ExternalModuleMapperTransform): }, "DeepseekV3MoE": { "forward": QEffDeepseekV3MoE.forward, - "moe": QEffDeepseekV3MoE.moe, + "moe_weights_as_activations": QEffDeepseekV3MoE.moe_weights_as_activations, + "moe_waa_unpack": QEffDeepseekV3MoE.moe_waa_unpack, + "original_moe": QEffDeepseekV3MoE.original_moe, "__qeff_init__": QEffDeepseekV3MoE.__qeff_init__, + "gate.forward": QEffDeepseekMoEGate.forward, }, "DeepseekV3Attention": { "forward": QEffDeepseekV3Attention.forward, @@ -1287,8 +1293,9 @@ class PrefillOnlyExternalModuleMapperTransform(ExternalModuleMapperTransform): _match_string_replace_method = { "DeepseekV3MoE": { "forward": QEffPrefillOnlyDeepseekV3MoE.forward, - "moe": QEffPrefillOnlyDeepseekV3MoE.moe, "__qeff_init__": QEffPrefillOnlyDeepseekV3MoE.__qeff_init__, + "_forward_expert_blocked": QEffPrefillOnlyDeepseekV3MoE._forward_expert_blocked, + "_cumsum_scatter_gather_update_expert_blocked": QEffPrefillOnlyDeepseekV3MoE._cumsum_scatter_gather_update_expert_blocked, }, } @@ -1298,7 +1305,7 @@ class RevertPrefillOnlyExternalModuleMapperTransform(ExternalModuleMapperTransfo _match_string_replace_method = { "DeepseekV3MoE": { "forward": QEffDeepseekV3MoE.forward, - "moe": QEffDeepseekV3MoE.moe, + "moe": QEffDeepseekV3MoE.moe_waa_unpack, "__qeff_init__": QEffDeepseekV3MoE.__qeff_init__, }, } diff --git a/QEfficient/transformers/quantizers/quant_transforms.py b/QEfficient/transformers/quantizers/quant_transforms.py index c3f204e629..8c2b997868 100644 --- a/QEfficient/transformers/quantizers/quant_transforms.py +++ b/QEfficient/transformers/quantizers/quant_transforms.py @@ -6,6 +6,8 @@ # ----------------------------------------------------------------------------- import torch +from compressed_tensors.compressors import PackedQuantizationCompressor +from compressed_tensors.linear.compressed_linear import CompressedLinear from torch import nn from transformers import AutoConfig from transformers.models.gpt_oss.modeling_gpt_oss import GptOssExperts @@ -69,6 +71,84 @@ def mutate(cls, original_module: nn.Module, parent_module: nn.Module): return new_module +class PackQuantizedInt4ToMatMulNBitsTransform(ModuleMutatorTransform): + """ + This transform is used to pack the quantized int4 weights into a format that can be used by the MatMulNBits kernel. + It is used for the ONNX export of the quantized model. + """ + + _match_class = CompressedLinear + + @classmethod + def _is_pack_quantized_int4_linear(cls, module): + quantization_scheme = getattr(module, "quantization_scheme", None) + quantization_args = getattr(quantization_scheme, "weights", None) + return ( + isinstance(module, nn.Linear) + and hasattr(module, "weight_packed") + and hasattr(module, "weight_scale") + and quantization_args is not None + and getattr(quantization_args, "num_bits", None) == 4 + and getattr(quantization_args, "type", None) == "int" + and getattr(quantization_args, "strategy", None) == "group" + ) + + @classmethod + def _matches(cls, module): + return isinstance(module, cls._match_class) or cls._is_pack_quantized_int4_linear(module) + + @classmethod + def apply(cls, model: nn.Module): + transformed = False + for name, module in model.named_children(): + if cls._matches(module): + setattr(model, name, cls.mutate(module, model)) + transformed = True + else: + _, child_transformed = cls.apply(module) + transformed = transformed or child_transformed + + if cls._matches(model): + model = cls.mutate(model, None) + transformed = True + + return model, transformed + + @classmethod + def mutate(cls, original_module, parent_module): + # add compressor.decompress to get the decompressed weight + # and then package into matmulnbit + if isinstance(original_module, CompressedLinear): + assert isinstance(original_module.compressor, PackedQuantizationCompressor), ( + f"Only {PackedQuantizationCompressor} supported for now" + ) + fp_weight = original_module.compressor.decompress_module(original_module) + else: + from compressed_tensors.compressors.base import decompress_module as ct_decompress_module + + ct_decompress_module(original_module) + fp_weight = original_module.weight + scales = original_module.weight_scale + # assuming symmetric quantization + quantization_args = original_module.quantization_scheme.weights + zeros = (torch.zeros_like(scales) + pow(2, (quantization_args.num_bits - 1))).to(torch.uint8) + g_idx = torch.arange(original_module.in_features // quantization_args.group_size).repeat_interleave( + quantization_args.group_size + ) + original_module.weight = torch.nn.Parameter(fp_weight) + assert quantization_args.type == "int", "uint is not tested yet" + new_module = QuantLinearORT( + quantization_args.num_bits, + quantization_args.group_size, + original_module.in_features, + original_module.out_features, + original_module.bias is not None, + ) + new_module.bias = original_module.bias if original_module.bias is not None else None + new_module.pack(original_module, scales, zeros, g_idx) + return new_module + + class GPTQToMatmulNbitsTransform(ModuleMutatorTransform): """ A transformation class that mutates a ``QuantLinearGPTQ`` module to a ``QuantLinearORT`` diff --git a/QEfficient/transformers/quantizers/quantizer_compressed_tensors.py b/QEfficient/transformers/quantizers/quantizer_compressed_tensors.py index 81c6b81cdf..8fff7e5f70 100644 --- a/QEfficient/transformers/quantizers/quantizer_compressed_tensors.py +++ b/QEfficient/transformers/quantizers/quantizer_compressed_tensors.py @@ -12,6 +12,7 @@ import torch from transformers.models.qwen3_vl_moe.modeling_qwen3_vl_moe import Qwen3VLMoeTextExperts from transformers.quantizers.quantizer_compressed_tensors import CompressedTensorsHfQuantizer +from transformers.utils import is_compressed_tensors_available from transformers.utils.quantization_config import CompressedTensorsConfig, QuantizationConfigMixin, QuantizationMethod from QEfficient.transformers.quantizers.quantizer_utils import blockwise_dequantize, get_keys_to_not_convert @@ -390,7 +391,55 @@ def update_unexpected_keys(self, model, unexpected_keys: List[str], prefix: str class QEffCompressedTensorsConfig(CompressedTensorsConfig): - def __init__( + def handle_pack_quantized_init( + self, + config_groups=None, + format="dense", + quantization_status="initialized", + kv_cache_scheme=None, + global_compression_ratio=None, + ignore=None, + sparsity_config=None, + quant_method="compressed-tensors", + run_compressed: bool = True, + **kwargs, + ): + if is_compressed_tensors_available(): + from compressed_tensors.config import SparsityCompressionConfig + from compressed_tensors.quantization import QuantizationConfig + else: + raise ImportError( + "compressed_tensors is not installed and is required for compressed-tensors quantization. Please install it with `pip install compressed-tensors`." + ) + self.quantization_config = None + self.sparsity_config = None + + self.run_compressed = run_compressed + assert self.run_compressed, "pack-quantized needs to have run_compressed set to True" + + # parse from dict to load nested QuantizationScheme objects + if config_groups or kv_cache_scheme: + self.quantization_config = QuantizationConfig.model_validate( + { + "config_groups": config_groups, + "quant_method": quant_method, + "format": format, + "quantization_status": quantization_status, + "kv_cache_scheme": kv_cache_scheme, + "global_compression_ratio": global_compression_ratio, + "ignore": ignore, + **kwargs, + } + ) + + if sparsity_config: + self.sparsity_config = SparsityCompressionConfig.load_from_registry( + sparsity_config.get("format"), **sparsity_config + ) + + self.quant_method = QuantizationMethod.COMPRESSED_TENSORS + + def handle_fp8_init( self, config_groups=None, format="dense", @@ -480,7 +529,50 @@ def __init__( self.quant_method = QuantizationMethod.COMPRESSED_TENSORS + def __init__( + self, + config_groups=None, + format="dense", + quantization_status="initialized", + kv_cache_scheme=None, + global_compression_ratio=None, + ignore=None, + sparsity_config=None, + quant_method="compressed-tensors", + run_compressed: bool = None, + **kwargs, + ): + if format == "pack-quantized": + self.handle_pack_quantized_init( + config_groups=config_groups, + format=format, + quantization_status=quantization_status, + kv_cache_scheme=kv_cache_scheme, + global_compression_ratio=global_compression_ratio, + ignore=ignore, + sparsity_config=sparsity_config, + quant_method=quant_method, + run_compressed=True if run_compressed is None else run_compressed, + **kwargs, + ) + else: + self.handle_fp8_init( + config_groups=config_groups, + format=format, + quantization_status=quantization_status, + kv_cache_scheme=kv_cache_scheme, + global_compression_ratio=global_compression_ratio, + ignore=ignore, + sparsity_config=sparsity_config, + quant_method=quant_method, + run_compressed=False if run_compressed is None else run_compressed, + **kwargs, + ) + def to_dict(self): + if self.quantization_config.format == "pack-quantized": + return super().to_dict() + return { "quantization_config": { "config_groups": self.config_groups, @@ -501,39 +593,58 @@ def to_dict(self): class QEffCompressedTensorsFP8Quantizer(CompressedTensorsHfQuantizer): requires_calibration = False - def __init__(self, quantization_config, **kwargs): - # TODO: check if more checks are required - if not isinstance(quantization_config, QEffCompressedTensorsConfig): - raise TypeError( - f"Only {QEffCompressedTensorsConfig} is supported for initialization got {type(quantization_config)}" - ) - self.run_compressed = quantization_config.run_compressed - self.quantization_config = quantization_config - - # -- Handle extra kwargs below -- - self.modules_to_not_convert = kwargs.pop("modules_to_not_convert", []) - self.modules_to_not_convert = list( - set(self.modules_to_not_convert if self.modules_to_not_convert else []) - | set(self.quantization_config.ignore if self.quantization_config.ignore else []) + @staticmethod + def is_pack_quantized(quant_config): + return ( + hasattr(quant_config, "quantization_config") + and hasattr(quant_config.quantization_config, "format") + and quant_config.quantization_config.format == "pack-quantized" ) - self.pre_quantized = kwargs.pop("pre_quantized", True) - if not self.pre_quantized and self.requires_calibration: - raise ValueError( - f"The quantization method {quantization_config.quant_method} does require the model to be pre-quantized." - f" You explicitly passed `pre_quantized=False` meaning your model weights are not quantized. Make sure to " - f"pass `pre_quantized=True` while knowing what you are doing." + def __init__(self, quantization_config, **kwargs): + if self.is_pack_quantized(quantization_config): + super().__init__(quantization_config, **kwargs) + else: + if not isinstance(quantization_config, QEffCompressedTensorsConfig): + raise TypeError( + f"Only {QEffCompressedTensorsConfig} is supported for initialization got {type(quantization_config)}" + ) + self.run_compressed = quantization_config.run_compressed + self.quantization_config = quantization_config + + # -- Handle extra kwargs below -- + self.modules_to_not_convert = kwargs.pop("modules_to_not_convert", []) + self.modules_to_not_convert = list( + set(self.modules_to_not_convert if self.modules_to_not_convert else []) + | set(self.quantization_config.ignore if self.quantization_config.ignore else []) ) + self.pre_quantized = kwargs.pop("pre_quantized", True) + + if not self.pre_quantized and self.requires_calibration: + raise ValueError( + f"The quantization method {quantization_config.quant_method} does require the model to be pre-quantized." + f" You explicitly passed `pre_quantized=False` meaning your model weights are not quantized. Make sure to " + f"pass `pre_quantized=True` while knowing what you are doing." + ) def validate_environment(self, *args, **kwargs): + if self.is_pack_quantized(self.quantization_config): + return super().validate_environment(*args, **kwargs) return True def update_torch_dtype(self, torch_dtype): + if self.is_pack_quantized(self.quantization_config): + return super().update_torch_dtype(torch_dtype) + if torch_dtype not in [None, torch.float32]: logger.warning(f"Requested dtype {torch_dtype} is not supported, overriding to float32") return torch.float32 def _process_model_before_weight_loading(self, model, **kwargs): + if self.is_pack_quantized(self.quantization_config): + super()._process_model_before_weight_loading(model, **kwargs) + return + if self.quantization_config.targets != ["Linear"]: raise NotImplementedError( f"Only Linear layer with FP8 quantization are supported got targets = {self.quantization_config.targets}" @@ -561,9 +672,14 @@ def replace_linear_with_fp8_dequant_layer(module): replace_linear_with_fp8_dequant_layer(model) def _process_model_after_weight_loading(self, model, **kwargs): + if self.is_pack_quantized(self.quantization_config): + super()._process_model_after_weight_loading(model, **kwargs) + return pass def update_missing_keys_after_loading(self, model, missing_keys: List[str], prefix: str) -> List[str]: + if self.is_pack_quantized(self.quantization_config): + return super().update_missing_keys_after_loading(model, missing_keys=missing_keys, prefix=prefix) return missing_keys def update_unexpected_keys(self, model, unexpected_keys: List[str], prefix: str = None) -> List[str]: diff --git a/QEfficient/utils/constants.py b/QEfficient/utils/constants.py index bf68f58df0..bf3e8f9676 100644 --- a/QEfficient/utils/constants.py +++ b/QEfficient/utils/constants.py @@ -140,6 +140,8 @@ def get_default_aic_hw_version() -> str: DEFAULT_AIC_HW_VERSION = get_default_aic_hw_version() ONNX_TRANSFORM_MEMORY_CLEANUP_INTERVAL = 100 +EXPERT_BLOCKING_PACKED_CHUNK_SIZE = int(os.environ.get("EXPERT_BLOCKING_PACKED_CHUNK_SIZE", "256")) + # InternVL constants # Fixing the feature size with reference to OpenGVLab/InternVL2_5-1B, OpenGVLab/InternVL2_5-38B and OpenGVLab/InternVL2_5-78B INTERN_FEATURE_SIZE = 256 From 9463dc37cbff0b53c56e856139930d7f0b876762 Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Fri, 10 Jul 2026 00:48:55 +0530 Subject: [PATCH 16/33] update tests Signed-off-by: Mamta Singh --- .../test_image_text_to_text_models.py | 4 +- .../test_quantization_transforms.py | 59 +++++++++++++++++++ 2 files changed, 62 insertions(+), 1 deletion(-) diff --git a/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py b/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py index 91087c593e..c16c009285 100644 --- a/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py +++ b/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py @@ -166,7 +166,6 @@ def _load_kimi_k25_layer_subset_model(): NUM_EXPERTS_PER_TOKEN, NUM_TEXT_LAYERS, NUM_VISION_LAYERS, - _disable_compressed_tensor_forward_pre_hooks as _example_disable_compressed_tensor_forward_pre_hooks, _load_layer_subset_model, _materialize_missing_linear_weights, _patch_deepseek_init_weights_compat, @@ -175,6 +174,9 @@ def _load_kimi_k25_layer_subset_model(): _resolve_model_path, _set_deterministic, ) + from test_kimi_k25 import ( + _disable_compressed_tensor_forward_pre_hooks as _example_disable_compressed_tensor_forward_pre_hooks, + ) _set_deterministic(1234) _ensure_torch_fx_import_compatibility() diff --git a/tests/unit_test/transforms/test_quantization_transforms.py b/tests/unit_test/transforms/test_quantization_transforms.py index b7fa03c1d2..e6c317c344 100644 --- a/tests/unit_test/transforms/test_quantization_transforms.py +++ b/tests/unit_test/transforms/test_quantization_transforms.py @@ -93,6 +93,7 @@ def test_all_transforms_have_mutate_classmethod(self): FP8DeQuantLinearToLinearTransform, GPTQToMatmulNbitsTransform, Mxfp4GptOssExpertDequantizeTransform, + PackQuantizedInt4ToMatMulNBitsTransform, ) for cls in [ @@ -100,6 +101,7 @@ def test_all_transforms_have_mutate_classmethod(self): GPTQToMatmulNbitsTransform, FP8DeQuantLinearToLinearTransform, Mxfp4GptOssExpertDequantizeTransform, + PackQuantizedInt4ToMatMulNBitsTransform, ]: assert hasattr(cls, "mutate"), f"{cls.__name__} missing mutate method" assert callable(cls.mutate), f"{cls.__name__}.mutate is not callable" @@ -111,6 +113,7 @@ def test_all_transforms_are_subclasses_of_module_mutator(self): FP8DeQuantLinearToLinearTransform, GPTQToMatmulNbitsTransform, Mxfp4GptOssExpertDequantizeTransform, + PackQuantizedInt4ToMatMulNBitsTransform, ) for cls in [ @@ -118,6 +121,7 @@ def test_all_transforms_are_subclasses_of_module_mutator(self): GPTQToMatmulNbitsTransform, FP8DeQuantLinearToLinearTransform, Mxfp4GptOssExpertDequantizeTransform, + PackQuantizedInt4ToMatMulNBitsTransform, ]: assert issubclass(cls, ModuleMutatorTransform), ( f"{cls.__name__} must be a subclass of ModuleMutatorTransform" @@ -328,6 +332,55 @@ def test_quantization_transforms_come_before_kv_cache_transform(self): f"AwqToMatmulNbitsTransform (idx={awq_idx}) must come before KVCacheTransform (idx={kv_idx})" ) + def test_image_text_wrappers_include_pack_quantized_int4_transform(self): + from QEfficient.transformers.models.modeling_auto import ( + QEffCausalLMForTextImageToTextModel, + QEffVisionEncoderForTextImageToTextModel, + ) + from QEfficient.transformers.models.pytorch_transforms import CustomOpsTransform + from QEfficient.transformers.quantizers.quant_transforms import PackQuantizedInt4ToMatMulNBitsTransform + + for wrapper_cls in [QEffVisionEncoderForTextImageToTextModel, QEffCausalLMForTextImageToTextModel]: + transforms = wrapper_cls._pytorch_transforms + pack_idx = transforms.index(PackQuantizedInt4ToMatMulNBitsTransform) + custom_ops_idx = transforms.index(CustomOpsTransform) + assert pack_idx < custom_ops_idx + + def test_pack_quantized_int4_transform_matches_compressed_linear_metadata(self, monkeypatch): + from types import SimpleNamespace + import importlib + + import torch + + from QEfficient.customop.matmulnbits import QuantLinearORT + from QEfficient.transformers.quantizers.quant_transforms import PackQuantizedInt4ToMatMulNBitsTransform + + linear = torch.nn.Linear(8, 8, bias=False) + del linear._parameters["weight"] + linear.register_parameter( + "weight_packed", torch.nn.Parameter(torch.zeros((8, 4), dtype=torch.int32), requires_grad=False) + ) + linear.register_parameter( + "weight_scale", torch.nn.Parameter(torch.ones((8, 2), dtype=torch.float32), requires_grad=False) + ) + linear.quantization_scheme = SimpleNamespace( + weights=SimpleNamespace(num_bits=4, type="int", strategy="group", group_size=4) + ) + + def fake_decompress_module(module): + module.weight = torch.nn.Parameter(torch.zeros(module.out_features, module.in_features), requires_grad=False) + + compressor_base = importlib.import_module("compressed_tensors.compressors.base") + monkeypatch.setattr(compressor_base, "decompress_module", fake_decompress_module) + + model = torch.nn.Sequential(linear) + transformed_model, transformed = PackQuantizedInt4ToMatMulNBitsTransform.apply(model) + + assert transformed + assert isinstance(transformed_model[0], QuantLinearORT) + assert transformed_model[0].bits == 4 + assert transformed_model[0].group_size == 4 + def test_non_quantized_model_not_affected_by_quant_transforms(self): """Applying quantization transforms to a non-quantized model must not change it.""" import torch @@ -336,6 +389,7 @@ def test_non_quantized_model_not_affected_by_quant_transforms(self): from QEfficient.transformers.quantizers.quant_transforms import ( AwqToMatmulNbitsTransform, GPTQToMatmulNbitsTransform, + PackQuantizedInt4ToMatMulNBitsTransform, ) cfg = GPT2Config(n_layer=1, n_head=2, n_embd=64, vocab_size=500, n_positions=32, n_ctx=32) @@ -349,9 +403,14 @@ def test_non_quantized_model_not_affected_by_quant_transforms(self): model_gptq, applied_gptq = GPTQToMatmulNbitsTransform.apply(model) assert not applied_gptq, "GPTQToMatmulNbitsTransform must not apply to non-quantized model" + model_pack, applied_pack = PackQuantizedInt4ToMatMulNBitsTransform.apply(model) + assert not applied_pack, "PackQuantizedInt4ToMatMulNBitsTransform must not apply to non-quantized model" + # Model output must be unchanged input_ids = torch.randint(0, 500, (1, 8)) with torch.no_grad(): original_logits = model(input_ids=input_ids).logits awq_logits = model_awq(input_ids=input_ids).logits + pack_logits = model_pack(input_ids=input_ids).logits assert torch.allclose(original_logits, awq_logits), "AWQ transform must not change non-quantized model output" + assert torch.allclose(original_logits, pack_logits), "Packed-int4 transform must not change non-quantized model output" From bd24583cf4ea4ad4014a3820ded12b1301885134 Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Mon, 13 Jul 2026 19:07:31 +0530 Subject: [PATCH 17/33] fix mismatch Signed-off-by: Mamta Singh --- .../models/deepseek_v3/modeling_deepseek.py | 7 +++- examples/kimi_k2/test_kimi_k25.py | 41 ++++--------------- .../test_image_text_to_text_models.py | 12 ------ .../test_quantization_transforms.py | 10 +++-- 4 files changed, 21 insertions(+), 49 deletions(-) diff --git a/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py b/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py index 2eab9ff5bf..c3b8343a5f 100644 --- a/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py +++ b/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py @@ -673,7 +673,12 @@ def forward_full_kv( key_states = torch.cat((k_nope, k_pe_new), -1) if past_key_value is not None: - cache_kwargs = {"sin": sin, "cos": cos, "batch_index": batch_index, "position_ids": position_ids} + cache_kwargs = { + "sin": sin_cached, + "cos": cos_cached, + "batch_index": batch_index, + "position_ids": position_ids, + } key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.softmax_scale diff --git a/examples/kimi_k2/test_kimi_k25.py b/examples/kimi_k2/test_kimi_k25.py index 7e19ac4083..82e54331b4 100644 --- a/examples/kimi_k2/test_kimi_k25.py +++ b/examples/kimi_k2/test_kimi_k25.py @@ -41,29 +41,6 @@ CTX_LEN = 1024 -def _materialize_missing_linear_weights(module: torch.nn.Module): - """Decompress compressed linear modules that expose packed weights but no `.weight`.""" - try: - from compressed_tensors.compressors.base import decompress_module as ct_decompress_module - except Exception: - return - - for child in module.modules(): - if not isinstance(child, torch.nn.Linear): - continue - if hasattr(child, "weight"): - continue - if not (hasattr(child, "weight_packed") and hasattr(child, "weight_scale")): - continue - ct_decompress_module(child) - - -def _disable_compressed_tensor_forward_pre_hooks(module: torch.nn.Module): - """Avoid double-decompress hooks that can raise KeyError('weight_packed').""" - if hasattr(module, "_forward_pre_hooks"): - module._forward_pre_hooks.clear() - - def _has_qaic_runtime_access() -> bool: try: _ = QAICInferenceSession @@ -231,8 +208,8 @@ def _greedy_generate_qeff_wrapper(transformed_model, inputs, max_new_tokens: int def _greedy_generate_onnx(transformed_model, onnx_paths, inputs, max_new_tokens, session_options): vision_onnx_path, lang_onnx_path = onnx_paths - vision_session = ort.InferenceSession(str(vision_onnx_path)) # , providers=["CPUExecutionProvider"]) - lang_session = ort.InferenceSession(str(lang_onnx_path), session_options) # providers=["CPUExecutionProvider"]) + vision_session = ort.InferenceSession(str(vision_onnx_path)) + lang_session = ort.InferenceSession(str(lang_onnx_path), session_options) pixel_values = inputs["pixel_values"].detach().cpu().numpy().astype(np.float32) grid_thws = inputs["grid_thws"].detach().cpu().numpy().astype(np.int64) @@ -349,8 +326,6 @@ def check_kimi_k25_pytorch_vs_ai100(): dtype=torch.float32, ) - _materialize_missing_linear_weights(model.language_model) - _disable_compressed_tensor_forward_pre_hooks(model) model.vision_tower.patch_embed.pos_emb.interpolation_mode = "bilinear" model = model.eval().to("cpu") @@ -360,12 +335,9 @@ def check_kimi_k25_pytorch_vs_ai100(): hf_tokens = _greedy_generate_hf(model, _clone_inputs(inputs), max_new_tokens=NEW_GENERATION_TOKENS) print("HF:", _decode_tokens(tokenizer, hf_tokens), "\n", hf_tokens) - qeff_model = QEFFAutoModelForImageTextToText( - copy.deepcopy(model), - kv_offload=True, - torch_dtype=torch.float32, - ) - _disable_compressed_tensor_forward_pre_hooks(qeff_model.model) + qaic_config = {"mla_absorption": {"cache_compressed": True, "absorption": False, "online": False}} + + qeff_model = QEFFAutoModelForImageTextToText(model, qaic_config=qaic_config) qeff_tokens = _greedy_generate_qeff_wrapper( transformed_model=qeff_model.model, @@ -404,6 +376,7 @@ def check_kimi_k25_pytorch_vs_ai100(): onnx_tokens = None qeff_model.compile( + qaic_config=qaic_config, num_devices=1, prefill_seq_len=1, ctx_len=CTX_LEN, @@ -417,6 +390,7 @@ def check_kimi_k25_pytorch_vs_ai100(): qaic_tokens = None if _has_qaic_runtime_access(): + inputs["pixel_values"] = inputs["pixel_values"].to(qeff_model.model.config.torch_dtype) qaic_tokens = qeff_model.generate( inputs=_clone_inputs(inputs), generation_len=NEW_GENERATION_TOKENS ).generated_ids[:, :-1] @@ -429,6 +403,7 @@ def check_kimi_k25_pytorch_vs_ai100(): f"vision_layers={model.config.vision_config.vt_num_hidden_layers}", f"text_layers={model.config.text_config.num_hidden_layers}", ) + print("Prompt:", repr(TEXT_PROMPT)) print("HF:", _decode_tokens(tokenizer, hf_tokens)) print("QEFF:", _decode_tokens(tokenizer, qeff_tokens)) diff --git a/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py b/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py index c16c009285..02e08f6527 100644 --- a/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py +++ b/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py @@ -151,11 +151,6 @@ def _get_kimi_k25_num_image_tokens(config, grid_thws): return int(grid_thws[0, 1].item() // merge_height) * int(grid_thws[0, 2].item() // merge_width) -def _disable_compressed_tensor_forward_pre_hooks(module: torch.nn.Module): - if hasattr(module, "_forward_pre_hooks"): - module._forward_pre_hooks.clear() - - def _load_kimi_k25_layer_subset_model(): examples_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../examples/kimi_k2")) if examples_dir not in sys.path: @@ -167,16 +162,12 @@ def _load_kimi_k25_layer_subset_model(): NUM_TEXT_LAYERS, NUM_VISION_LAYERS, _load_layer_subset_model, - _materialize_missing_linear_weights, _patch_deepseek_init_weights_compat, _patch_kimi_tie_weights_compat, _prepare_config, _resolve_model_path, _set_deterministic, ) - from test_kimi_k25 import ( - _disable_compressed_tensor_forward_pre_hooks as _example_disable_compressed_tensor_forward_pre_hooks, - ) _set_deterministic(1234) _ensure_torch_fx_import_compatibility() @@ -196,8 +187,6 @@ def _load_kimi_k25_layer_subset_model(): num_experts_per_tok=NUM_EXPERTS_PER_TOKEN, dtype=torch.float32, ) - _materialize_missing_linear_weights(model.language_model) - _example_disable_compressed_tensor_forward_pre_hooks(model) model.vision_tower.patch_embed.pos_emb.interpolation_mode = "bilinear" model = model.eval().to("cpu") return model, tokenizer, processor @@ -285,7 +274,6 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100( config=model_hf.config, torch_dtype=torch_dtype, ) - _disable_compressed_tensor_forward_pre_hooks(qeff_model.model) elif config is None: config = AutoConfig.from_pretrained( model_name, trust_remote_code=True, padding=model_name not in ModelConfig.MOLMO_MODELS diff --git a/tests/unit_test/transforms/test_quantization_transforms.py b/tests/unit_test/transforms/test_quantization_transforms.py index e6c317c344..e89432873c 100644 --- a/tests/unit_test/transforms/test_quantization_transforms.py +++ b/tests/unit_test/transforms/test_quantization_transforms.py @@ -347,8 +347,8 @@ def test_image_text_wrappers_include_pack_quantized_int4_transform(self): assert pack_idx < custom_ops_idx def test_pack_quantized_int4_transform_matches_compressed_linear_metadata(self, monkeypatch): - from types import SimpleNamespace import importlib + from types import SimpleNamespace import torch @@ -368,7 +368,9 @@ def test_pack_quantized_int4_transform_matches_compressed_linear_metadata(self, ) def fake_decompress_module(module): - module.weight = torch.nn.Parameter(torch.zeros(module.out_features, module.in_features), requires_grad=False) + module.weight = torch.nn.Parameter( + torch.zeros(module.out_features, module.in_features), requires_grad=False + ) compressor_base = importlib.import_module("compressed_tensors.compressors.base") monkeypatch.setattr(compressor_base, "decompress_module", fake_decompress_module) @@ -413,4 +415,6 @@ def test_non_quantized_model_not_affected_by_quant_transforms(self): awq_logits = model_awq(input_ids=input_ids).logits pack_logits = model_pack(input_ids=input_ids).logits assert torch.allclose(original_logits, awq_logits), "AWQ transform must not change non-quantized model output" - assert torch.allclose(original_logits, pack_logits), "Packed-int4 transform must not change non-quantized model output" + assert torch.allclose(original_logits, pack_logits), ( + "Packed-int4 transform must not change non-quantized model output" + ) From bb57f5461c9bf1bdc340e64bf52c242b74c9ee10 Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Tue, 14 Jul 2026 01:37:09 +0530 Subject: [PATCH 18/33] update tests Signed-off-by: Mamta Singh --- examples/kimi_k2/test_kimi_k25.py | 8 ++++---- .../kimi_k2}/test_kimi_k25_disagg.py | 8 ++------ tests/configs/image_text_model_configs.json | 3 +++ .../image_text_to_text/test_image_text_to_text_models.py | 2 +- 4 files changed, 10 insertions(+), 11 deletions(-) rename {tests/transformers/models/image_text_to_text => examples/kimi_k2}/test_kimi_k25_disagg.py (96%) diff --git a/examples/kimi_k2/test_kimi_k25.py b/examples/kimi_k2/test_kimi_k25.py index 82e54331b4..d69c366e4f 100644 --- a/examples/kimi_k2/test_kimi_k25.py +++ b/examples/kimi_k2/test_kimi_k25.py @@ -332,12 +332,12 @@ def check_kimi_k25_pytorch_vs_ai100(): inputs = _prepare_inputs(processor) inputs = {k: (v.to("cpu") if torch.is_tensor(v) else v) for k, v in inputs.items()} - hf_tokens = _greedy_generate_hf(model, _clone_inputs(inputs), max_new_tokens=NEW_GENERATION_TOKENS) + hf_tokens = _greedy_generate_hf(copy.deepcopy(model), _clone_inputs(inputs), max_new_tokens=NEW_GENERATION_TOKENS) print("HF:", _decode_tokens(tokenizer, hf_tokens), "\n", hf_tokens) - qaic_config = {"mla_absorption": {"cache_compressed": True, "absorption": False, "online": False}} + # qaic_config = {"mla_absorption": {"cache_compressed": True, "absorption": False, "online": False}} - qeff_model = QEFFAutoModelForImageTextToText(model, qaic_config=qaic_config) + qeff_model = QEFFAutoModelForImageTextToText(model) # , qaic_config=qaic_config) qeff_tokens = _greedy_generate_qeff_wrapper( transformed_model=qeff_model.model, @@ -376,7 +376,7 @@ def check_kimi_k25_pytorch_vs_ai100(): onnx_tokens = None qeff_model.compile( - qaic_config=qaic_config, + # qaic_config=qaic_config, num_devices=1, prefill_seq_len=1, ctx_len=CTX_LEN, diff --git a/tests/transformers/models/image_text_to_text/test_kimi_k25_disagg.py b/examples/kimi_k2/test_kimi_k25_disagg.py similarity index 96% rename from tests/transformers/models/image_text_to_text/test_kimi_k25_disagg.py rename to examples/kimi_k2/test_kimi_k25_disagg.py index 70c0855948..1965bb50b4 100644 --- a/tests/transformers/models/image_text_to_text/test_kimi_k25_disagg.py +++ b/examples/kimi_k2/test_kimi_k25_disagg.py @@ -5,6 +5,7 @@ # # ---------------------------------------------------------------------------- +import copy from pathlib import Path import numpy as np @@ -22,10 +23,8 @@ from test_kimi_k25 import ( _clone_inputs, _decode_tokens, - _disable_compressed_tensor_forward_pre_hooks, _greedy_generate_hf, _has_qaic_runtime_access, - _materialize_missing_linear_weights, _prepare_inputs, _resolve_model_path, _set_deterministic, @@ -113,8 +112,6 @@ def _load_kimi_subset_model(): num_experts_per_tok=NUM_EXPERTS_PER_TOKEN, dtype=torch.float32, ) - _materialize_missing_linear_weights(model.language_model) - _disable_compressed_tensor_forward_pre_hooks(model) model.vision_tower.patch_embed.pos_emb.interpolation_mode = "bilinear" return model.eval().to("cpu"), tokenizer, processor @@ -267,7 +264,7 @@ def test_kimi_k25_disagg_qaic_vs_hf_fp32(): model, tokenizer, processor = _load_kimi_subset_model() inputs = _prepare_inputs(processor) inputs = {name: (value.to("cpu") if torch.is_tensor(value) else value) for name, value in inputs.items()} - hf_tokens = _greedy_generate_hf(model, _clone_inputs(inputs), max_new_tokens=GENERATION_LEN) + hf_tokens = _greedy_generate_hf(copy.deepcopy(model), _clone_inputs(inputs), max_new_tokens=GENERATION_LEN) qeff_model = QEFFAutoModelForImageTextToText( model, @@ -276,7 +273,6 @@ def test_kimi_k25_disagg_qaic_vs_hf_fp32(): torch_dtype=torch.float32, layerwise=False, ) - _disable_compressed_tensor_forward_pre_hooks(qeff_model.model) compile_dims = _get_image_compile_dims(qeff_model.model, inputs) vision_qpc_path, prefill_qpc_path, decode_qpc_path, compiled_onnx_paths = _compile_disagg_qpcs( diff --git a/tests/configs/image_text_model_configs.json b/tests/configs/image_text_model_configs.json index cc4cc343d9..5a301b0dcd 100644 --- a/tests/configs/image_text_model_configs.json +++ b/tests/configs/image_text_model_configs.json @@ -658,6 +658,9 @@ "vt_num_attention_heads": 4, "vt_num_hidden_layers": 2 } + } + }, + { "model_name": "tiny-random/gemma-4-dense", "model_type": "gemma4", "batch_size": 1, diff --git a/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py b/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py index 49cdeedb82..df923b22be 100644 --- a/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py +++ b/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py @@ -437,7 +437,7 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100( tokenize=False, return_tensors="pt", ) - pytorch_hf_tokens = _run_kimi_k25_hf_model_on_pytorch(model_hf, processor, inputs, max_gen_len) + pytorch_hf_tokens = _run_kimi_k25_hf_model_on_pytorch(copy.deepcopy(model_hf), processor, inputs, max_gen_len) inputs = processor( messages=conversation, add_generation_prompt=True, From 9d4b48d4d51c0a09eb16165d9d81133d1c49a3a3 Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Wed, 15 Jul 2026 00:51:05 +0530 Subject: [PATCH 19/33] cleanup Signed-off-by: Mamta Singh --- .../models/deepseek_v3/modeling_deepseek.py | 40 +--- .../models/kimi_k25/modeling_kimi_k25.py | 176 ++++++++++++++++++ examples/kimi_k2/test_kimi_k25.py | 162 ---------------- examples/kimi_k2/test_kimi_k25_disagg.py | 10 +- 4 files changed, 184 insertions(+), 204 deletions(-) diff --git a/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py b/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py index 1c6241a0fb..5940cc1a51 100644 --- a/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py +++ b/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py @@ -33,23 +33,6 @@ from QEfficient.utils.constants import MAX_POSITION_EMBEDDINGS, MIN_MASKED_ATTENTION_VALUE -def _get_linear_weight(linear_module: nn.Module) -> torch.Tensor: - """Return a linear layer weight, decompressing compressed-tensors modules if needed.""" - weight = getattr(linear_module, "weight", None) - if weight is not None: - return weight - - if hasattr(linear_module, "weight_packed") and hasattr(linear_module, "weight_scale"): - from compressed_tensors.compressors.base import decompress_module as ct_decompress_module - - ct_decompress_module(linear_module) - weight = getattr(linear_module, "weight", None) - if weight is not None: - return weight - - raise AttributeError(f"{linear_module.__class__.__name__!s} object has no attribute 'weight'") - - def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] @@ -140,17 +123,6 @@ def _set_cos_sin_cache(self, seq_len, device, dtype): self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False) -# def forward(self, x, seq_len=None): -# # x: [bs, num_attention_heads, seq_len, head_size] -# if self.max_seq_len_cached is None or seq_len > self.max_seq_len_cached: -# self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype) -# -# return ( -# self.cos_cached[:seq_len].to(dtype=x.dtype), -# self.sin_cached[:seq_len].to(dtype=x.dtype), -# ) - - class DeepseekV3YarnRotaryEmbedding(DeepseekV3RotaryEmbedding): def __init__( self, @@ -230,8 +202,6 @@ def orig_apply_rotary_pos_emb(q, k, cos, sin): # , position_ids, unsqueeze_dim= Returns: `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. """ - # cos = cos[position_ids].unsqueeze(unsqueeze_dim) - # sin = sin[position_ids].unsqueeze(unsqueeze_dim) b, h, s, d = q.shape q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) @@ -731,9 +701,7 @@ def forward_full_kv_h_blocking( kv = ( self.kv_b_proj(self.kv_a_layernorm(kva)) - .view( - bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim - ) # TODO : split this matmul #with k_up and v_up + .view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim) .transpose(1, 2) ) @@ -814,7 +782,7 @@ def forward( ) -EXPERT_BLOCKING_NUM_NSP = int(os.environ.get("EXPERT_BLOCKING_NUM_NSP", "4")) +EXPERT_BLOCKING_NUM_NSP = int(os.environ.get("EXPERT_BLOCKING_NUM_NSP", "16")) EXPERT_BLOCKING_PACKED_CHUNK_SIZE = int(os.environ.get("EXPERT_BLOCKING_PACKED_CHUNK_SIZE", "256")) @@ -1090,7 +1058,6 @@ def original_moe(self, x, topk_ids, topk_weight): tokens_per_expert = cnts.sum(dim=0) idxs = topk_ids.view(-1).argsort() sorted_tokens = x[idxs // topk_ids.shape[1]] - # sorted_tokens_shape = sorted_tokens.shape tokens_per_expert = tokens_per_expert.cpu().numpy() outputs = [] @@ -1668,8 +1635,7 @@ def forward( if position_ids is None: position_ids = cache_position.unsqueeze(0) - # ctx_len = compressed_kvs.layers[0].ckv.shape[-2] - causal_mask = _create_causal_mask(position_ids=position_ids, target_length=target_len) # ctx_len) + causal_mask = _create_causal_mask(position_ids=position_ids, target_length=target_len) hidden_states = inputs_embeds position_embeddings = None diff --git a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py index ef3e66eb73..2352132e21 100644 --- a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py +++ b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py @@ -1089,3 +1089,179 @@ def get_specializations( lang[0].pop("vision_size") lang[1].pop("vision_size") return lang, compiler_options + + def get_specializations_multi_res( + self, + batch_size: int, + prefill_seq_len: int, + ctx_len: int, + img_size: None, + h: int | List[int] = None, + w: int | List[int] = None, + num_frames: int | List[int] = 1, + kv_offload: bool = False, + continuous_batching: bool = False, + kv_cache_batch_size: Optional[int] = None, + full_batch_size: Optional[int] = None, + **compiler_options, + ): + comp_ctx_lengths_prefill = compiler_options.pop("comp_ctx_lengths_prefill", None) + comp_ctx_lengths_decode = compiler_options.pop("comp_ctx_lengths_decode", None) + num_patches = compiler_options.pop("num_patches", None) + h = compiler_options.pop("h", None) + w = compiler_options.pop("w", None) + num_image_tokens = compiler_options.pop("num_image_tokens", None) + + height = [h] if isinstance(h, int) else h + width = [w] if isinstance(w, int) else w + num_frames = [num_frames] * len(h) if isinstance(num_frames, int) else num_frames + + prefill_seq_len = prefill_seq_len if prefill_seq_len else constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN + ctx_len = ctx_len if ctx_len else constants.ONNX_EXPORT_CTX_LEN + + channel = 3 + patch_size = self.config.vision_config.patch_size + #temporal_patch_size = self.config.vision_config.temporal_patch_size + + IMAGE_FACTOR = constants.IMAGE_FACTOR_QWEN_2_5 + IMAGE_MIN_TOKEN_NUM = constants.IMAGE_MIN_TOKEN_NUM + IMAGE_MAX_TOKEN_NUM = constants.IMAGE_MAX_TOKEN_NUM + min_pixels = IMAGE_MIN_TOKEN_NUM * IMAGE_FACTOR**2 + max_pixels = IMAGE_MAX_TOKEN_NUM * IMAGE_FACTOR**2 + breakpoint() + mm_processor_kwargs = compiler_options.pop("mm_processor_kwargs", None) + if mm_processor_kwargs: + min_pixels = mm_processor_kwargs.get("min_pixels", min_pixels) + max_pixels = mm_processor_kwargs.get("max_pixels", max_pixels) + + vision = [] + max_vision_size = 0 + user_vision_size = compiler_options.pop("vision_size", None) + if user_vision_size: + assert user_vision_size < ctx_len, "vision_size must be less than ctx_len" + max_vision_size = user_vision_size + + for h, w, f in zip(height, width, num_frames): + resized_height, resized_width = smart_resize( + height=h, width=w, factor=IMAGE_FACTOR, min_pixels=min_pixels, max_pixels=max_pixels + ) + grid_h, grid_w = resized_height // patch_size, resized_width // patch_size + grid_height = grid_h * grid_w + grid_width = patch_size * patch_size * temporal_patch_size * channel + vision_size = grid_height // 4 + grid_height = grid_height * batch_size + if not user_vision_size: + max_vision_size = max(max_vision_size, vision_size * f) + assert max_vision_size < ctx_len, ( + f"Computed vision_size of {vision_size * f} tokens " + f"(vision_size={vision_size}, num_frames={f}) for image resolution " + f"(width={w}, height={h}) must be less than ctx_len. Please adjust the image " + "resolution." + ) + else: + if vision_size * f > user_vision_size: + logger.warning_once( + f"Computed vision_size of {vision_size * f} tokens " + f"(vision_size={vision_size}, num_frames={f}) for image resolution " + f"(width={w}, height={h}) exceeds the provided " + f"vision_size={user_vision_size}. " + f"Vision embedding needs to be chunked during prefill." + ) + + vision.append( + { + "batch_size": batch_size, + "vision_size": vision_size, + "grid_height": grid_height, + "grid_width": grid_width, + "grid_h": grid_h, + "grid_w": grid_w, + } + ) + + num_patches = num_patches if num_patches is not None else constants.KIMI_NUM_PATCHES + h = h if h is not None else constants.KIMI_IMAGE_HEIGHT + w = w if w is not None else constants.KIMI_IMAGE_WIDTH + num_image_tokens = num_image_tokens if num_image_tokens is not None else constants.KIMI_NUM_IMAGE_TOKENS + + vision = [ + { + "num_patches": num_patches, + "h": h, + "w": w, + "num_image_tokens": num_image_tokens, + } + ] + + if comp_ctx_lengths_prefill is not None: + lang = [] + + for i in range(0, len(comp_ctx_lengths_prefill)): + lang_prefill = { + "batch_size": 1 if continuous_batching else batch_size, + "seq_len": prefill_seq_len, + "ctx_len": ctx_len, + "num_image_tokens": num_image_tokens, + } + if continuous_batching: + lang_prefill["full_batch_size"] = kv_cache_batch_size + else: + lang_prefill["batch_size"] = kv_cache_batch_size + if full_batch_size: + lang_prefill["full_batch_exec_size"] = full_batch_size + + lang.append(lang_prefill) + + for i in range(0, len(comp_ctx_lengths_decode)): + lang_decode = { + "batch_size": full_batch_size if continuous_batching else batch_size, + "seq_len": "1", + "ctx_len": ctx_len, + "num_image_tokens": num_image_tokens, + } + + if continuous_batching: + lang_decode["full_batch_size"] = kv_cache_batch_size + else: + lang_decode["batch_size"] = kv_cache_batch_size + + lang.append(lang_decode) + + else: + lang_prefill = { + "batch_size": 1 if continuous_batching else batch_size, + "seq_len": prefill_seq_len, + "ctx_len": ctx_len, + "num_image_tokens": num_image_tokens, + } + if continuous_batching: + lang_prefill["full_batch_size"] = kv_cache_batch_size + else: + lang_prefill["batch_size"] = kv_cache_batch_size + if full_batch_size: + lang_prefill["full_batch_exec_size"] = full_batch_size + + lang_decode = { + "batch_size": full_batch_size if continuous_batching else batch_size, + "seq_len": 1, + "ctx_len": ctx_len, + "num_image_tokens": num_image_tokens, + } + + if continuous_batching: + lang_decode["full_batch_size"] = kv_cache_batch_size + else: + lang_decode["batch_size"] = kv_cache_batch_size + + lang = [lang_prefill, lang_decode] + + specializations = {} + + if kv_offload: + specializations["vision"] = vision + specializations["lang"] = lang + return specializations, compiler_options + else: + lang[0].pop("vision_size") + lang[1].pop("vision_size") + return lang, compiler_options diff --git a/examples/kimi_k2/test_kimi_k25.py b/examples/kimi_k2/test_kimi_k25.py index d69c366e4f..40b26b31eb 100644 --- a/examples/kimi_k2/test_kimi_k25.py +++ b/examples/kimi_k2/test_kimi_k25.py @@ -137,36 +137,6 @@ def _greedy_generate_hf(model, inputs, max_new_tokens: int): return generated_ids[:, -max_new_tokens:] -def _match_indexed_name(name: str, base: str): - pattern = rf"^{re.escape(base)}\.(\d+)$" - match = re.match(pattern, name) - return int(match.group(1)) if match else None - - -def _build_cache_inputs_from_dummy(transformed_model, input_names, batch_size: int): - pkv = transformed_model.language_model.get_dummy_pkv_cache( - transformed_model.config.text_config, - batch_size, - CTX_LEN, - ) - cache_inputs = {} - for name in input_names: - key_idx = _match_indexed_name(name, "past_key") - value_idx = _match_indexed_name(name, "past_value") - compressed_idx = _match_indexed_name(name, "compressed_kv") - k_pe_idx = _match_indexed_name(name, "k_pe") - - if key_idx is not None: - cache_inputs[name] = pkv[key_idx][0].detach().cpu().numpy().astype(np.float32) - elif value_idx is not None: - cache_inputs[name] = pkv[value_idx][1].detach().cpu().numpy().astype(np.float32) - elif compressed_idx is not None: - cache_inputs[name] = pkv[compressed_idx][0].detach().cpu().numpy().astype(np.float32) - elif k_pe_idx is not None: - cache_inputs[name] = pkv[k_pe_idx][1].detach().cpu().numpy().astype(np.float32) - return cache_inputs - - def _greedy_generate_qeff_wrapper(transformed_model, inputs, max_new_tokens: int): qeff_encoder = transformed_model.get_qeff_vision_encoder().eval() decoder_wrapper = transformed_model.get_qeff_language_decoder().eval() @@ -206,106 +176,6 @@ def _greedy_generate_qeff_wrapper(transformed_model, inputs, max_new_tokens: int return generated_ids[:, -max_new_tokens:] -def _greedy_generate_onnx(transformed_model, onnx_paths, inputs, max_new_tokens, session_options): - vision_onnx_path, lang_onnx_path = onnx_paths - vision_session = ort.InferenceSession(str(vision_onnx_path)) - lang_session = ort.InferenceSession(str(lang_onnx_path), session_options) - - pixel_values = inputs["pixel_values"].detach().cpu().numpy().astype(np.float32) - grid_thws = inputs["grid_thws"].detach().cpu().numpy().astype(np.int64) - h = int(grid_thws[0, 1]) - w = int(grid_thws[0, 2]) - - vision_outputs = vision_session.run( - None, - { - "pixel_values": pixel_values, - "h_shape": np.ones((h,), dtype=np.int64), - "w_shape": np.ones((w,), dtype=np.int64), - }, - ) - vision_output_names = [out.name for out in vision_session.get_outputs()] - image_embeds = {name: value for name, value in zip(vision_output_names, vision_outputs)}.get( - "image_embeds", vision_outputs[0] - ) - - lang_input_names = [meta.name for meta in lang_session.get_inputs()] - lang_output_names = [meta.name for meta in lang_session.get_outputs()] - - prompt_input_ids = inputs["input_ids"].detach().cpu().numpy().astype(np.int64) - prompt_attention_mask = inputs["attention_mask"].detach().cpu().numpy().astype(np.int64) - prompt_len = prompt_input_ids.shape[1] - - prefill_seq_len = 32 - num_chunks = -(prompt_len // -prefill_seq_len) - padded_len = num_chunks * prefill_seq_len - - pad_token_id = 1 - padded_input_ids = np.pad(prompt_input_ids, ((0, 0), (0, padded_len - prompt_len)), constant_values=pad_token_id) - padded_attention = np.pad(prompt_attention_mask, ((0, 0), (0, padded_len - prompt_len)), constant_values=0) - padded_position_ids = np.where(padded_attention > 0, np.arange(padded_len), -1).astype(np.int64) - - cache_inputs = _build_cache_inputs_from_dummy( - transformed_model=transformed_model, - input_names=lang_input_names, - batch_size=prompt_input_ids.shape[0], - ) - image_idx = np.zeros((prompt_input_ids.shape[0], 1), dtype=np.int64) - - output_map = None - for chunk_idx in range(num_chunks): - start = chunk_idx * prefill_seq_len - end = (chunk_idx + 1) * prefill_seq_len - ort_inputs = { - "input_ids": padded_input_ids[:, start:end], - "position_ids": padded_position_ids[:, start:end], - "image_embeds": image_embeds.astype(np.float32), - "image_idx": image_idx, - **cache_inputs, - } - ort_outputs = lang_session.run(None, ort_inputs) - output_map = {name: value for name, value in zip(lang_output_names, ort_outputs)} - if "image_idx_output" in output_map: - image_idx = output_map["image_idx_output"].astype(np.int64) - for key in list(cache_inputs.keys()): - retained_name = f"{key}_RetainedState" - if retained_name in output_map: - cache_inputs[key] = output_map[retained_name] - - if output_map is None: - raise RuntimeError("ONNX prefill did not execute.") - - next_token = np.argmax(output_map["logits"][:, -1, :], axis=-1, keepdims=True).astype(np.int64) - generated_new_tokens = [next_token] - - decode_input_ids = next_token - decode_position_ids = np.max(padded_position_ids, axis=1, keepdims=True).astype(np.int64) + 1 - - for _ in range(1, max_new_tokens): - ort_inputs = { - "input_ids": decode_input_ids, - "position_ids": decode_position_ids, - "image_embeds": image_embeds.astype(np.float32), - "image_idx": image_idx, - **cache_inputs, - } - ort_outputs = lang_session.run(None, ort_inputs) - output_map = {name: value for name, value in zip(lang_output_names, ort_outputs)} - - decode_input_ids = np.argmax(output_map["logits"][:, -1, :], axis=-1, keepdims=True).astype(np.int64) - generated_new_tokens.append(decode_input_ids) - decode_position_ids = decode_position_ids + 1 - - if "image_idx_output" in output_map: - image_idx = output_map["image_idx_output"].astype(np.int64) - for key in list(cache_inputs.keys()): - retained_name = f"{key}_RetainedState" - if retained_name in output_map: - cache_inputs[key] = output_map[retained_name] - - return torch.from_numpy(np.concatenate(generated_new_tokens, axis=1)) - - def check_kimi_k25_pytorch_vs_ai100(): _set_deterministic(1234) _ensure_torch_fx_import_compatibility() @@ -346,35 +216,6 @@ def check_kimi_k25_pytorch_vs_ai100(): ) print("QEFF:", _decode_tokens(tokenizer, qeff_tokens), "\n", qeff_tokens) - onnx_paths = qeff_model.export() - - # Replace invalid index value for INT32 max to 0 using add_initializer - m = onnx.load(onnx_paths[1], load_external_data=False) - # NOTE: OrtValue objects should be kept around until the session is run, hence this dict is required - added_initializers = {} - for node in m.graph.node: - if node.op_type == "Constant": - np_tensor = onnx.numpy_helper.to_array(node.attribute[0].t, os.path.dirname(onnx_paths[1])) - if len(np_tensor.shape) == 0 and np_tensor.item() == 2147483647: - added_initializers[node.output[0]] = ort.OrtValue.ortvalue_from_numpy(np.array(0, np_tensor.dtype)) - - session_options = ort.SessionOptions() - for name, value in added_initializers.items(): - session_options.add_initializer(name, value) - - try: - onnx_tokens = _greedy_generate_onnx( - transformed_model=qeff_model.model, - onnx_paths=onnx_paths, - inputs=_clone_inputs(inputs), - max_new_tokens=NEW_GENERATION_TOKENS, - session_options=session_options, - ) - print("ONNX:", _decode_tokens(tokenizer, onnx_tokens) if onnx_tokens is not None else "") - except Exception as exc: - print(f"ONNX generation failed: {exc}") - onnx_tokens = None - qeff_model.compile( # qaic_config=qaic_config, num_devices=1, @@ -407,12 +248,9 @@ def check_kimi_k25_pytorch_vs_ai100(): print("Prompt:", repr(TEXT_PROMPT)) print("HF:", _decode_tokens(tokenizer, hf_tokens)) print("QEFF:", _decode_tokens(tokenizer, qeff_tokens)) - print("ONNX:", _decode_tokens(tokenizer, onnx_tokens) if onnx_tokens is not None else "") print("QAIC:", _decode_tokens(tokenizer, qaic_tokens) if qaic_tokens is not None else "") assert torch.equal(hf_tokens, qeff_tokens), "HF and QEFF(Pytorch runtime) tokens do not match" - if onnx_tokens is not None: - assert torch.equal(hf_tokens, onnx_tokens), "HF and ONNXRuntime tokens do not match" if qaic_tokens is not None: assert torch.equal(hf_tokens, torch.as_tensor(qaic_tokens)), "HF and QAIC tokens do not match" diff --git a/examples/kimi_k2/test_kimi_k25_disagg.py b/examples/kimi_k2/test_kimi_k25_disagg.py index 1965bb50b4..14a414d84e 100644 --- a/examples/kimi_k2/test_kimi_k25_disagg.py +++ b/examples/kimi_k2/test_kimi_k25_disagg.py @@ -34,10 +34,10 @@ from QEfficient import QEFFAutoModelForImageTextToText from QEfficient.generation.cloud_infer import QAICInferenceSession -PREFILL_SEQ_LEN = 32 -CTX_LEN = 1024 +PREFILL_SEQ_LEN = 512 +CTX_LEN = 2048 BATCH_SIZE = 1 -GENERATION_LEN = 4 +GENERATION_LEN = 10 NUM_VISION_LAYERS = 2 NUM_TEXT_LAYERS = 2 @@ -153,9 +153,9 @@ def _compile_disagg_qpcs(qeff_model: QEFFAutoModelForImageTextToText, compile_di prefill_qpc_path = qeff_model.compile( prefill_seq_len=PREFILL_SEQ_LEN, - retain_full_kv=True, + #retain_full_kv=True, prefill_only=True, - enable_chunking=True, + #enable_chunking=True, skip_vision=True, skip_lang=False, **common_compile_kwargs, From 82a1a09c2ec63ca62ddd95c31ac140708fd7360d Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Fri, 17 Jul 2026 02:13:59 +0530 Subject: [PATCH 20/33] Enable CB Signed-off-by: Mamta Singh --- QEfficient/generation/embedding_handler.py | 35 ++- .../generation/text_generation_inference.py | 32 +- QEfficient/generation/vlm_generation.py | 7 +- .../models/kimi_k25/modeling_kimi_k25.py | 184 +---------- examples/kimi_k2/test_kimi_k25.py | 5 - examples/kimi_k2/test_kimi_k25_disagg.py | 4 +- .../test_continuous_batching.py | 295 +++++++++++++++++- 7 files changed, 365 insertions(+), 197 deletions(-) diff --git a/QEfficient/generation/embedding_handler.py b/QEfficient/generation/embedding_handler.py index 2a5a61f6b8..15c553ab91 100755 --- a/QEfficient/generation/embedding_handler.py +++ b/QEfficient/generation/embedding_handler.py @@ -235,18 +235,23 @@ def prepare_vlm_inputs(self, image_url: str, query: str, prefill_seq_len: int) - image = image.resize( (constants.GRANITEVISION_IMG_SIZE_HEIGHT, constants.GRANITEVISION_IMG_SIZE_WIDTH) ) + model_type = getattr(getattr(self._qeff_model, "model", None).config, "model_type", "") + # Gemma4 expects the processor-rendered prompt with the image placeholder ahead of user text. is_gemma4 = ( hasattr(self._qeff_model.model.config, "model_type") and self._qeff_model.model.config.model_type == "gemma4" ) + is_kimi_k25 = model_type == "kimi_k25" # Prepare conversation format conversation = [ { "role": "user", "content": ( - [{"type": "image"}, {"type": "text", "text": query}] + [{"type": "image_url", "image_url": image}, {"type": "text", "text": query}] + if is_kimi_k25 + else [{"type": "image"}, {"type": "text", "text": query}] if is_gemma4 else [{"type": "text", "text": query}, {"type": "image"}] ), @@ -254,22 +259,29 @@ def prepare_vlm_inputs(self, image_url: str, query: str, prefill_seq_len: int) - ] # Apply chat template - if is_gemma4: + if is_kimi_k25: + inputs = self._processor( + messages=conversation, + add_generation_prompt=True, + tokenize=False, + return_tensors="pt", + ) + elif is_gemma4: prompt = self._processor.apply_chat_template( conversation, tokenize=False, add_generation_prompt=True, enable_thinking=False, ) + inputs = self._processor(images=image, text=prompt, return_tensors="pt") else: prompt = self._processor.apply_chat_template( conversation, tokenize=False, add_generation_prompt=True, ) - # Process image and text - inputs = self._processor(images=image, text=prompt, return_tensors="pt") - model_type = getattr(getattr(self._qeff_model, "model", None).config, "model_type", "") + inputs = self._processor(images=image, text=prompt, return_tensors="pt") + if model_type in { "qwen2_5_vl", "qwen3_vl_moe", @@ -299,6 +311,13 @@ def prepare_vlm_inputs(self, image_url: str, query: str, prefill_seq_len: int) - }: vision_inputs[k] = np.array(v) + if is_kimi_k25: + grid_thws = inputs.get("grid_thws") + if grid_thws is None: + raise ValueError("Kimi-K2.5 processor output must include grid_thws for vision export.") + vision_inputs["h_shape"] = np.ones(int(grid_thws[0, 1].item()), dtype=np.int64) + vision_inputs["w_shape"] = np.ones(int(grid_thws[0, 2].item()), dtype=np.int64) + # Convert specific inputs to float16 vision_inputs_fp16 = {"pixel_values", "image_masks"} for k in vision_inputs_fp16: @@ -432,7 +451,11 @@ def setup_vision_buffers(self): buffers = {} for output_name, shape in shapes.items(): # Create placeholder with appropriate dtype - if "vision_embeds" in output_name or "deepstack_features" in output_name: + if ( + "vision_embeds" in output_name + or "image_embeds" in output_name + or "deepstack_features" in output_name + ): buffers[output_name] = np.zeros(shape, dtype=np.float16) else: buffers[output_name] = np.zeros(shape, dtype=np.float32) diff --git a/QEfficient/generation/text_generation_inference.py b/QEfficient/generation/text_generation_inference.py index 17c992064c..f80fd4d77e 100755 --- a/QEfficient/generation/text_generation_inference.py +++ b/QEfficient/generation/text_generation_inference.py @@ -462,6 +462,7 @@ def __init__( self.include_guided_decoding = include_guided_decoding self.sampling_params = sampling_params self._qpc_path = qpc_path # Store qpc_path for later use + self.tokenizer = tokenizer # Load QPC self._session = QAICInferenceSession( @@ -496,7 +497,6 @@ def __init__( self.decode_pos_ids = None self.generation_len = None - self.tokenizer = tokenizer self._set_tokenizer_params() # set tokenizer params # Skip inputs/outputs self._session.skip_buffers( @@ -569,7 +569,18 @@ def _fetch_decode_seq_len( decode_seq_len = min( [x[self._session.binding_index_map["input_ids"]][1][1] for x in self._session.allowed_shapes] ) - return decode_seq_len + elif "input_ids" in self._session.binding_index_map: + dims = self._session.bindings[self._session.binding_index_map["input_ids"]].dims + if len(dims) > 1 and dims[1] not in (None, -1): + decode_seq_len = dims[1] + + if decode_seq_len is None: + decode_seq_len = getattr(self, "_prefill_seq_len", None) + + if decode_seq_len is None: + raise ValueError("Unable to determine decode sequence length from QPC input_ids binding.") + + return int(decode_seq_len) def _fetch_vocab_size( self, @@ -588,6 +599,16 @@ def _fetch_vocab_size( ) if self._session.allowed_shapes: return [x[self._session.binding_index_map[key]] for x in self._session.allowed_shapes][0][1][2] + if key in self._session.binding_index_map: + dims = self._session.bindings[self._session.binding_index_map[key]].dims + if len(dims) >= 3 and dims[2] not in (None, -1): + return dims[2] + if self.tokenizer is not None: + vocab_size = getattr(self.tokenizer, "vocab_size", None) + if vocab_size is not None: + return vocab_size + return len(self.tokenizer) + return None return self._session.bindings[self._session.binding_index_map[key]].dims[2] @@ -685,6 +706,7 @@ def initialize_decode_inputs(self, num_prompts, execution_batch_size, max_gen_le self.generated_ids = np.full((num_prompts, max_gen_length), self.tokenizer.pad_token_id) self.decode_input_ids = np.zeros((execution_batch_size, 1), np.int64) self.decode_pos_ids = np.zeros((execution_batch_size, 1), np.int64) + self.decode_image_idx = np.zeros((execution_batch_size, 1), np.int64) self.generation_len = np.zeros((execution_batch_size, 1), np.int64) def initialize_lora_id_mapping(self, prompt_to_lora_id_mapping): @@ -722,6 +744,8 @@ def update_decode_input(self, outputs, position_ids, generation_len, decode_batc decode_batch = decode_batch_id if decode_batch_id is not None else slice(None) self.decode_input_ids[decode_batch] = next_token_id self.decode_pos_ids[decode_batch] = position_ids + if "image_idx_output" in outputs: + self.decode_image_idx[decode_batch] = outputs["image_idx_output"] self.generated_ids[decode_batch, 0] = next_token_id.squeeze() self.generation_len[decode_batch] = generation_len return next_token_id @@ -764,6 +788,10 @@ def _set_output_buffers(self, batch_size: int = 1, sequence_length: int = 1): logits_out_placeholder = np.zeros((batch_size, sequence_length, self._vocab_size), dtype=np.float32) self._session.set_buffers({"logits": logits_out_placeholder}) + if "image_idx_output" in getattr(self._session, "binding_index_map", {}): + image_idx_out_placeholder = np.zeros((batch_size, 1), dtype=np.int64) + self._session.set_buffers({"image_idx_output": image_idx_out_placeholder}) + def run_prefill(self, prompt, generation_len, prefill_logit_bs=1, decode_batch_id=None): """ Runs prefill for a given prompt and generation length. diff --git a/QEfficient/generation/vlm_generation.py b/QEfficient/generation/vlm_generation.py index df731f951a..013452fbc1 100755 --- a/QEfficient/generation/vlm_generation.py +++ b/QEfficient/generation/vlm_generation.py @@ -989,7 +989,12 @@ def prepare_decode_inputs(self): if "image_idx" in getattr(self._session, "binding_index_map", {}): idx = self._session.binding_index_map["image_idx"] dims = tuple(self._session.bindings[idx].dims) - decode_inputs["image_idx"] = np.zeros(dims, dtype=np.int64) + if dims == tuple(self.decode_image_idx.shape): + decode_inputs["image_idx"] = self.decode_image_idx + elif dims[0] == 1: + decode_inputs["image_idx"] = self.decode_image_idx[:1] + else: + decode_inputs["image_idx"] = np.broadcast_to(self.decode_image_idx[:1], dims).copy() else: decode_inputs["image_idx"] = np.array([[0]], dtype=np.int64) except Exception: diff --git a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py index 2352132e21..ea068df862 100644 --- a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py +++ b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py @@ -961,14 +961,16 @@ def get_onnx_dynamic_axes( else: cache_compressed = False + cache_batch_axis = "full_batch_size" if continuous_batching else "batch_size" + if cache_compressed: for i in range(self.language_model.config.num_hidden_layers): - lang_dynamic_axes[f"compressed_kv.{i}"] = {0: "batch_size", 2: "ctx_len"} - lang_dynamic_axes[f"k_pe.{i}"] = {0: "batch_size", 2: "ctx_len"} + lang_dynamic_axes[f"compressed_kv.{i}"] = {0: cache_batch_axis, 2: "ctx_len"} + lang_dynamic_axes[f"k_pe.{i}"] = {0: cache_batch_axis, 2: "ctx_len"} else: for i in range(self.language_model.config.num_hidden_layers): for kv in ["key", "value"]: - lang_dynamic_axes[f"past_{kv}.{i}"] = {0: "batch_size", 2: "ctx_len"} + lang_dynamic_axes[f"past_{kv}.{i}"] = {0: cache_batch_axis, 2: "ctx_len"} if comp_ctx_lengths is not None: lang_dynamic_axes["comp_ctx_lengths"] = {0: "comp_ctx_lengths"} @@ -1052,19 +1054,6 @@ def get_specializations( lang.append(lang_decode) else: - lang_prefill = { - "batch_size": 1 if continuous_batching else batch_size, - "seq_len": prefill_seq_len, - "ctx_len": ctx_len, - "num_image_tokens": num_image_tokens, - } - if continuous_batching: - lang_prefill["full_batch_size"] = kv_cache_batch_size - else: - lang_prefill["batch_size"] = kv_cache_batch_size - if full_batch_size: - lang_prefill["full_batch_exec_size"] = full_batch_size - lang_decode = { "batch_size": full_batch_size if continuous_batching else batch_size, "seq_len": 1, @@ -1077,157 +1066,6 @@ def get_specializations( else: lang_decode["batch_size"] = kv_cache_batch_size - lang = [lang_prefill, lang_decode] - - specializations = {} - - if kv_offload: - specializations["vision"] = vision - specializations["lang"] = lang - return specializations, compiler_options - else: - lang[0].pop("vision_size") - lang[1].pop("vision_size") - return lang, compiler_options - - def get_specializations_multi_res( - self, - batch_size: int, - prefill_seq_len: int, - ctx_len: int, - img_size: None, - h: int | List[int] = None, - w: int | List[int] = None, - num_frames: int | List[int] = 1, - kv_offload: bool = False, - continuous_batching: bool = False, - kv_cache_batch_size: Optional[int] = None, - full_batch_size: Optional[int] = None, - **compiler_options, - ): - comp_ctx_lengths_prefill = compiler_options.pop("comp_ctx_lengths_prefill", None) - comp_ctx_lengths_decode = compiler_options.pop("comp_ctx_lengths_decode", None) - num_patches = compiler_options.pop("num_patches", None) - h = compiler_options.pop("h", None) - w = compiler_options.pop("w", None) - num_image_tokens = compiler_options.pop("num_image_tokens", None) - - height = [h] if isinstance(h, int) else h - width = [w] if isinstance(w, int) else w - num_frames = [num_frames] * len(h) if isinstance(num_frames, int) else num_frames - - prefill_seq_len = prefill_seq_len if prefill_seq_len else constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN - ctx_len = ctx_len if ctx_len else constants.ONNX_EXPORT_CTX_LEN - - channel = 3 - patch_size = self.config.vision_config.patch_size - #temporal_patch_size = self.config.vision_config.temporal_patch_size - - IMAGE_FACTOR = constants.IMAGE_FACTOR_QWEN_2_5 - IMAGE_MIN_TOKEN_NUM = constants.IMAGE_MIN_TOKEN_NUM - IMAGE_MAX_TOKEN_NUM = constants.IMAGE_MAX_TOKEN_NUM - min_pixels = IMAGE_MIN_TOKEN_NUM * IMAGE_FACTOR**2 - max_pixels = IMAGE_MAX_TOKEN_NUM * IMAGE_FACTOR**2 - breakpoint() - mm_processor_kwargs = compiler_options.pop("mm_processor_kwargs", None) - if mm_processor_kwargs: - min_pixels = mm_processor_kwargs.get("min_pixels", min_pixels) - max_pixels = mm_processor_kwargs.get("max_pixels", max_pixels) - - vision = [] - max_vision_size = 0 - user_vision_size = compiler_options.pop("vision_size", None) - if user_vision_size: - assert user_vision_size < ctx_len, "vision_size must be less than ctx_len" - max_vision_size = user_vision_size - - for h, w, f in zip(height, width, num_frames): - resized_height, resized_width = smart_resize( - height=h, width=w, factor=IMAGE_FACTOR, min_pixels=min_pixels, max_pixels=max_pixels - ) - grid_h, grid_w = resized_height // patch_size, resized_width // patch_size - grid_height = grid_h * grid_w - grid_width = patch_size * patch_size * temporal_patch_size * channel - vision_size = grid_height // 4 - grid_height = grid_height * batch_size - if not user_vision_size: - max_vision_size = max(max_vision_size, vision_size * f) - assert max_vision_size < ctx_len, ( - f"Computed vision_size of {vision_size * f} tokens " - f"(vision_size={vision_size}, num_frames={f}) for image resolution " - f"(width={w}, height={h}) must be less than ctx_len. Please adjust the image " - "resolution." - ) - else: - if vision_size * f > user_vision_size: - logger.warning_once( - f"Computed vision_size of {vision_size * f} tokens " - f"(vision_size={vision_size}, num_frames={f}) for image resolution " - f"(width={w}, height={h}) exceeds the provided " - f"vision_size={user_vision_size}. " - f"Vision embedding needs to be chunked during prefill." - ) - - vision.append( - { - "batch_size": batch_size, - "vision_size": vision_size, - "grid_height": grid_height, - "grid_width": grid_width, - "grid_h": grid_h, - "grid_w": grid_w, - } - ) - - num_patches = num_patches if num_patches is not None else constants.KIMI_NUM_PATCHES - h = h if h is not None else constants.KIMI_IMAGE_HEIGHT - w = w if w is not None else constants.KIMI_IMAGE_WIDTH - num_image_tokens = num_image_tokens if num_image_tokens is not None else constants.KIMI_NUM_IMAGE_TOKENS - - vision = [ - { - "num_patches": num_patches, - "h": h, - "w": w, - "num_image_tokens": num_image_tokens, - } - ] - - if comp_ctx_lengths_prefill is not None: - lang = [] - - for i in range(0, len(comp_ctx_lengths_prefill)): - lang_prefill = { - "batch_size": 1 if continuous_batching else batch_size, - "seq_len": prefill_seq_len, - "ctx_len": ctx_len, - "num_image_tokens": num_image_tokens, - } - if continuous_batching: - lang_prefill["full_batch_size"] = kv_cache_batch_size - else: - lang_prefill["batch_size"] = kv_cache_batch_size - if full_batch_size: - lang_prefill["full_batch_exec_size"] = full_batch_size - - lang.append(lang_prefill) - - for i in range(0, len(comp_ctx_lengths_decode)): - lang_decode = { - "batch_size": full_batch_size if continuous_batching else batch_size, - "seq_len": "1", - "ctx_len": ctx_len, - "num_image_tokens": num_image_tokens, - } - - if continuous_batching: - lang_decode["full_batch_size"] = kv_cache_batch_size - else: - lang_decode["batch_size"] = kv_cache_batch_size - - lang.append(lang_decode) - - else: lang_prefill = { "batch_size": 1 if continuous_batching else batch_size, "seq_len": prefill_seq_len, @@ -1241,18 +1079,6 @@ def get_specializations_multi_res( if full_batch_size: lang_prefill["full_batch_exec_size"] = full_batch_size - lang_decode = { - "batch_size": full_batch_size if continuous_batching else batch_size, - "seq_len": 1, - "ctx_len": ctx_len, - "num_image_tokens": num_image_tokens, - } - - if continuous_batching: - lang_decode["full_batch_size"] = kv_cache_batch_size - else: - lang_decode["batch_size"] = kv_cache_batch_size - lang = [lang_prefill, lang_decode] specializations = {} diff --git a/examples/kimi_k2/test_kimi_k25.py b/examples/kimi_k2/test_kimi_k25.py index 40b26b31eb..ec3fc4a2c7 100644 --- a/examples/kimi_k2/test_kimi_k25.py +++ b/examples/kimi_k2/test_kimi_k25.py @@ -7,13 +7,10 @@ import copy import os -import re from io import BytesIO from pathlib import Path import numpy as np -import onnx -import onnxruntime as ort import requests import torch from export_kimi_k25_vision import ( @@ -58,8 +55,6 @@ def _has_qaic_runtime_access() -> bool: def _set_deterministic(seed: int): import random - import numpy as np - random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) diff --git a/examples/kimi_k2/test_kimi_k25_disagg.py b/examples/kimi_k2/test_kimi_k25_disagg.py index 14a414d84e..4efd432a05 100644 --- a/examples/kimi_k2/test_kimi_k25_disagg.py +++ b/examples/kimi_k2/test_kimi_k25_disagg.py @@ -153,9 +153,9 @@ def _compile_disagg_qpcs(qeff_model: QEFFAutoModelForImageTextToText, compile_di prefill_qpc_path = qeff_model.compile( prefill_seq_len=PREFILL_SEQ_LEN, - #retain_full_kv=True, + # retain_full_kv=True, prefill_only=True, - #enable_chunking=True, + # enable_chunking=True, skip_vision=True, skip_lang=False, **common_compile_kwargs, diff --git a/tests/transformers/models/image_text_to_text/test_continuous_batching.py b/tests/transformers/models/image_text_to_text/test_continuous_batching.py index 6dbeb5fd04..f5d3927172 100644 --- a/tests/transformers/models/image_text_to_text/test_continuous_batching.py +++ b/tests/transformers/models/image_text_to_text/test_continuous_batching.py @@ -6,8 +6,10 @@ # ---------------------------------------------------------------------------- import copy +import inspect import json import os +import sys from io import BytesIO from typing import Optional @@ -21,6 +23,8 @@ AutoTokenizer, GenerationConfig, ) +from transformers.dynamic_module_utils import get_class_from_dynamic_module +from transformers.utils import import_utils as hf_import_utils from QEfficient import QEFFAutoModelForCausalLM, QEFFAutoModelForImageTextToText from QEfficient.utils.run_utils import ApiRunnerInternVL, ApiRunnerMolmo, ApiRunnerVlm @@ -40,6 +44,216 @@ model_config_dict = {model["model_name"]: model for model in multimodal_models} NEW_GENERATION_TOKENS = 10 +KIMI_K25_MODEL_NAME = "moonshotai/Kimi-K2.5" + + +def _is_kimi_k25(model_name: str) -> bool: + return model_name == KIMI_K25_MODEL_NAME + + +def _ensure_torch_fx_import_compatibility(): + if hasattr(hf_import_utils, "is_torch_fx_available"): + return + + def _is_torch_fx_available() -> bool: + if not hf_import_utils.is_torch_available(): + return False + try: + import torch.fx # noqa: F401 + + return True + except Exception: + return False + + hf_import_utils.is_torch_fx_available = _is_torch_fx_available + + +def _patch_kimi_k25_tie_weights_compat(kimi_cls): + tie_signature = inspect.signature(kimi_cls.tie_weights) + if tuple(tie_signature.parameters) != ("self",): + return + + def _tie_weights_compat(self, missing_keys=None, recompute_mapping=True): + lm_tie_weights = getattr(self.language_model, "tie_weights") + try: + return lm_tie_weights(missing_keys=missing_keys, recompute_mapping=recompute_mapping) + except TypeError: + return lm_tie_weights() + + kimi_cls.tie_weights = _tie_weights_compat + + +def _patch_kimi_k25_deepseek_init_weights_compat(kimi_cls): + module_prefix, _ = kimi_cls.__module__.rsplit(".", maxsplit=1) + deepseek_module = sys.modules.get(f"{module_prefix}.modeling_deepseek") + if deepseek_module is None or not hasattr(deepseek_module, "DeepseekV3PreTrainedModel"): + return + + deepseek_cls = deepseek_module.DeepseekV3PreTrainedModel + if getattr(deepseek_cls, "_qeff_test_init_weights_patched", False): + return + + def _init_weights_compat(self, module): + std = self.config.initializer_range + if isinstance(module, torch.nn.Linear): + if hasattr(module, "weight") and module.weight is not None: + module.weight.data.normal_(mean=0.0, std=std) + if hasattr(module, "bias") and module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, torch.nn.Embedding): + if hasattr(module, "weight") and module.weight is not None: + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + deepseek_cls._init_weights = _init_weights_compat + deepseek_cls._qeff_test_init_weights_patched = True + + +def _patch_kimi_k25_remote_code_compat(config): + _ensure_torch_fx_import_compatibility() + kimi_cls = get_class_from_dynamic_module("modeling_kimi_k25.KimiK25ForConditionalGeneration", config._name_or_path) + _patch_kimi_k25_tie_weights_compat(kimi_cls) + _patch_kimi_k25_deepseek_init_weights_compat(kimi_cls) + return kimi_cls + + +def _get_kimi_k25_test_config(model_name: str): + config = AutoConfig.from_pretrained(model_name, trust_remote_code=True) + config._attn_implementation = "eager" + config.torch_dtype = torch.float32 + config.dtype = torch.float32 + additional_params = model_config_dict[model_name]["additional_params"] + + for attr, value in additional_params["text_config"].items(): + setattr(config.text_config, attr, value) + config.text_config._attn_implementation = "eager" + config.text_config.torch_dtype = torch.float32 + config.text_config.dtype = torch.float32 + + for attr, value in additional_params["vision_config"].items(): + setattr(config.vision_config, attr, value) + config.vision_config._attn_implementation = "eager" + config.vision_config.torch_dtype = torch.float32 + config.vision_config.dtype = torch.float32 + + _patch_kimi_k25_remote_code_compat(config) + return config + + +def _load_kimi_k25_model_from_config(config): + kimi_cls = _patch_kimi_k25_remote_code_compat(config) + model = kimi_cls._from_config(config) + torch_dtype = getattr(model.config, "torch_dtype", None) + if torch_dtype == torch.bfloat16 or torch_dtype == torch.float16: + model = model.to(torch.float32) + return model.eval() + + +def _get_kimi_k25_num_image_tokens(config, grid_thws): + merge_height, merge_width = config.vision_config.merge_kernel_size + return int(grid_thws[0, 1].item() // merge_height) * int(grid_thws[0, 2].item() // merge_width) + + +@torch.no_grad() +def _run_kimi_k25_hf_model_on_pytorch_CB(model, processor, images, queries, max_gen_len): + generated_tokens = [] + + eos_token_id = getattr(model.config, "eos_token_id", None) + if eos_token_id is None and hasattr(model.config, "text_config"): + eos_token_id = getattr(model.config.text_config, "eos_token_id", None) + + for idx, (image, query) in enumerate(zip(images, queries)): + conversation = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": image}, + {"type": "text", "text": query}, + ], + }, + ] + inputs = processor(messages=conversation, add_generation_prompt=True, tokenize=False, return_tensors="pt") + generated_ids = inputs["input_ids"] + attention_mask = inputs["attention_mask"] + pixel_values = inputs["pixel_values"] + grid_thws = inputs["grid_thws"] + new_tokens = [] + + for _ in range(max_gen_len): + outputs = model( + input_ids=generated_ids, + attention_mask=attention_mask, + pixel_values=pixel_values, + grid_thws=grid_thws, + use_cache=False, + return_dict=True, + ) + logits = outputs[0] if isinstance(outputs, tuple) else outputs.logits + next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True) + new_tokens.append(next_token) + + generated_ids = torch.cat([generated_ids, next_token], dim=1) + attention_mask = torch.cat( + [ + attention_mask, + torch.ones((attention_mask.shape[0], 1), dtype=attention_mask.dtype, device=attention_mask.device), + ], + dim=1, + ) + + if eos_token_id is not None and torch.all(next_token == eos_token_id): + break + + output_tokens = torch.cat(new_tokens, dim=1).squeeze(0) + py_output = processor.tokenizer.decode(output_tokens.tolist()).strip() + print(f"Original HF Model Outputs (Torch CPU) for prompt {idx}:") + print("Query:", repr(query)) + print("Completion:", repr(py_output)) + generated_tokens.append(output_tokens.numpy()) + + return generated_tokens + + +def _load_kimi_k25_layer_subset_model(): + examples_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../examples/kimi_k2")) + if examples_dir not in sys.path: + sys.path.insert(0, examples_dir) + + from test_kimi_k25 import ( # noqa: PLC0415 + LOADED_EXPERT_IDS, + NUM_EXPERTS_PER_TOKEN, + NUM_TEXT_LAYERS, + NUM_VISION_LAYERS, + _load_layer_subset_model, + _patch_deepseek_init_weights_compat, + _patch_kimi_tie_weights_compat, + _prepare_config, + _resolve_model_path, + _set_deterministic, + ) + + _set_deterministic(1234) + _ensure_torch_fx_import_compatibility() + model_path = _resolve_model_path() + config = _prepare_config(model_path) + kimi_cls = get_class_from_dynamic_module("modeling_kimi_k25.KimiK25ForConditionalGeneration", str(model_path)) + _patch_kimi_tie_weights_compat(kimi_cls) + _patch_deepseek_init_weights_compat(kimi_cls) + + model, tokenizer, processor = _load_layer_subset_model( + model_path=model_path, + kimi_cls=kimi_cls, + config=config, + num_vision_layers=min(NUM_VISION_LAYERS, 1), + num_text_layers=min(NUM_TEXT_LAYERS, 1), + loaded_expert_ids=LOADED_EXPERT_IDS, + num_experts_per_tok=NUM_EXPERTS_PER_TOKEN, + dtype=torch.float32, + ) + model.vision_tower.patch_embed.pos_emb.interpolation_mode = "bilinear" + model = model.eval().to("cpu") + return model, tokenizer, processor def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100_CB( @@ -62,8 +276,33 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100_CB( batch_size = model_config_dict[model_name]["batch_size"] full_batch_size = model_config_dict[model_name]["full_batch_size"] max_gen_len = NEW_GENERATION_TOKENS + kimi_tokenizer = None + kimi_processor = None + if _is_kimi_k25(model_name): + full_batch_size = 1 - if config is None: + if _is_kimi_k25(model_name) and config is None: + model_hf, kimi_tokenizer, kimi_processor = _load_kimi_k25_layer_subset_model() + config = model_hf.config + qeff_model = QEFFAutoModelForImageTextToText( + copy.deepcopy(model_hf), + kv_offload=kv_offload, + config=model_hf.config, + torch_dtype=torch.float32, + continuous_batching=True, + ) + elif _is_kimi_k25(model_name): + if config is None: + config = _get_kimi_k25_test_config(model_name) + model_hf = _load_kimi_k25_model_from_config(config) + model_hf.vision_tower.patch_embed.pos_emb.interpolation_mode = "bilinear" + qeff_model = QEFFAutoModelForImageTextToText( + copy.deepcopy(model_hf), + kv_offload=kv_offload, + config=model_hf.config, + continuous_batching=True, + ) + elif config is None: config = AutoConfig.from_pretrained( model_name, trust_remote_code=True, padding=model_name not in ModelConfig.MOLMO_MODELS ) @@ -184,6 +423,50 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100_CB( model_hf, image_list, prompt_list, generation_config ) compile_kwargs["img_size"] = img_size + elif _is_kimi_k25(model_name): + processor = kimi_processor or AutoProcessor.from_pretrained(model_name, trust_remote_code=True) + tokenizer = kimi_tokenizer or AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) + image_height = None + image_width = None + image_urls = [image_urls[0]] * len(queries) + num_patches = [] + image_heights = [] + image_widths = [] + num_image_tokens = [] + for img_url in image_urls: + image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB") + images.append(image) + + for image, query in zip(images, queries): + conversation = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": image}, + {"type": "text", "text": query}, + ], + }, + ] + inputs = processor(messages=conversation, add_generation_prompt=True, tokenize=False, return_tensors="pt") + num_patches.append(int(inputs["pixel_values"].shape[0])) + image_heights.append(int(inputs["grid_thws"][0, 1].item())) + image_widths.append(int(inputs["grid_thws"][0, 2].item())) + num_image_tokens.append(_get_kimi_k25_num_image_tokens(config, inputs["grid_thws"])) + + image_list = [images[0]] * full_batch_size + prompt_list = [queries[0]] * full_batch_size + pytorch_hf_tokens = _run_kimi_k25_hf_model_on_pytorch_CB( + copy.deepcopy(model_hf), processor, image_list, prompt_list, max_gen_len + ) + compile_kwargs.update( + { + "prefill_seq_len": 1, + "num_patches": num_patches[0], + "h": image_heights[0], + "w": image_widths[0], + "num_image_tokens": num_image_tokens[0], + } + ) else: processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True, padding=True) tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) @@ -226,6 +509,9 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100_CB( qeff_model.export() qeff_model.compile(**compile_kwargs) + # if _is_kimi_k25(model_name): + # manual_cleanup(qeff_model.onnx_path) + # return print("QPC Outputs (QAIC):") exec_info = qeff_model.generate( tokenizer=tokenizer, @@ -248,7 +534,12 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100_CB( model_hf, images, queries, generation_config=generation_config ) else: - pytorch_hf_tokens = api_runner.run_vlm_hf_model_on_pytorch_CB(model_hf, images, queries) + if _is_kimi_k25(model_name): + pytorch_hf_tokens = _run_kimi_k25_hf_model_on_pytorch_CB( + copy.deepcopy(model_hf), processor, images, queries, max_gen_len + ) + else: + pytorch_hf_tokens = api_runner.run_vlm_hf_model_on_pytorch_CB(model_hf, images, queries) print("QPC Outputs (QAIC):") exec_info = qeff_model.generate( From 089265eea57ffc00b81d4012fa12c07cf0a5b442 Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Fri, 17 Jul 2026 13:55:08 +0530 Subject: [PATCH 21/33] add support for multi resolution Signed-off-by: Mamta Singh --- .../models/kimi_k25/modeling_kimi_k25.py | 148 ++++++++++++++---- .../unit_test/models/test_model_quickcheck.py | 40 +++++ 2 files changed, 157 insertions(+), 31 deletions(-) diff --git a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py index ea068df862..a0e7cc8af5 100644 --- a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py +++ b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py @@ -997,27 +997,115 @@ def get_specializations( ): comp_ctx_lengths_prefill = compiler_options.pop("comp_ctx_lengths_prefill", None) comp_ctx_lengths_decode = compiler_options.pop("comp_ctx_lengths_decode", None) + compiler_options.pop("img_size", None) num_patches = compiler_options.pop("num_patches", None) + height = compiler_options.pop("height", None) + width = compiler_options.pop("width", None) h = compiler_options.pop("h", None) w = compiler_options.pop("w", None) + num_frames = compiler_options.pop("num_frames", 1) num_image_tokens = compiler_options.pop("num_image_tokens", None) + mm_processor_kwargs = compiler_options.pop("mm_processor_kwargs", None) or {} prefill_seq_len = prefill_seq_len if prefill_seq_len else constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN ctx_len = ctx_len if ctx_len else constants.ONNX_EXPORT_CTX_LEN - num_patches = num_patches if num_patches is not None else constants.KIMI_NUM_PATCHES - h = h if h is not None else constants.KIMI_IMAGE_HEIGHT - w = w if w is not None else constants.KIMI_IMAGE_WIDTH - num_image_tokens = num_image_tokens if num_image_tokens is not None else constants.KIMI_NUM_IMAGE_TOKENS - - vision = [ - { - "num_patches": num_patches, - "h": h, - "w": w, - "num_image_tokens": num_image_tokens, - } - ] + def normalize_list(value, default): + if value is None: + return [default] + if isinstance(value, int): + return [value] + return list(value) + + def normalize_sized_list(value, count, name): + if value is None: + return None + if isinstance(value, int): + return [value] * count + value = list(value) + if len(value) != count: + raise ValueError(f"Expected {name} to contain {count} entries, got {len(value)}.") + return value + + def validate_dimension_lists(heights, widths, height_name, width_name): + if len(heights) != len(widths): + raise ValueError( + f"Expected {height_name} and {width_name} to contain the same number of entries, " + f"got {len(heights)} and {len(widths)}." + ) + + patch_size = getattr(self.config.vision_config, "patch_size", constants.KIMI_PATCH_SIZE) + merge_kernel_size = getattr(self.config.vision_config, "merge_kernel_size", (2, 2)) + if isinstance(merge_kernel_size, int): + kernel_height = kernel_width = merge_kernel_size + merge_kernel_size = (merge_kernel_size, merge_kernel_size) + else: + kernel_height, kernel_width = merge_kernel_size + + if h is not None or w is not None: + heights = normalize_list(h, constants.KIMI_IMAGE_HEIGHT) + widths = normalize_list(w, constants.KIMI_IMAGE_WIDTH) + validate_dimension_lists(heights, widths, "h", "w") + elif height is not None or width is not None: + pixel_heights = normalize_list(height, constants.KIMI_IMAGE_HEIGHT * patch_size) + pixel_widths = normalize_list(width, constants.KIMI_IMAGE_WIDTH * patch_size) + validate_dimension_lists(pixel_heights, pixel_widths, "height", "width") + + in_patch_limit = mm_processor_kwargs.get("in_patch_limit", 16384) + patch_limit_on_one_side = mm_processor_kwargs.get("patch_limit_on_one_side", 512) + factor_height = kernel_height * patch_size + factor_width = kernel_width * patch_size + heights = [] + widths = [] + for pixel_height, pixel_width in zip(pixel_heights, pixel_widths): + scale = min( + 1.0, + math.sqrt( + in_patch_limit / (max(1.0, pixel_width // patch_size) * max(1.0, pixel_height // patch_size)) + ), + patch_limit_on_one_side * patch_size / pixel_width, + patch_limit_on_one_side * patch_size / pixel_height, + ) + resized_height = min(max(1, int(pixel_height * scale)), patch_limit_on_one_side * patch_size) + resized_width = min(max(1, int(pixel_width * scale)), patch_limit_on_one_side * patch_size) + pad_height = (factor_height - resized_height % factor_height) % factor_height + pad_width = (factor_width - resized_width % factor_width) % factor_width + heights.append((resized_height + pad_height) // patch_size) + widths.append((resized_width + pad_width) // patch_size) + else: + heights = [constants.KIMI_IMAGE_HEIGHT] + widths = [constants.KIMI_IMAGE_WIDTH] + + num_frames = normalize_sized_list(1 if num_frames is None else num_frames, len(heights), "num_frames") + explicit_num_patches = normalize_sized_list(num_patches, len(heights), "num_patches") + explicit_num_image_tokens = normalize_sized_list(num_image_tokens, len(heights), "num_image_tokens") + + vision = [] + max_num_image_tokens = 0 + for index, (height, width, frames) in enumerate(zip(heights, widths, num_frames)): + if height % kernel_height != 0 or width % kernel_width != 0: + raise ValueError( + f"Kimi image grid h={height}, w={width} must be divisible by merge_kernel_size={merge_kernel_size}." + ) + + computed_num_patches = height * width * frames + computed_num_image_tokens = (height // kernel_height) * (width // kernel_width) * frames + resolved_num_patches = ( + explicit_num_patches[index] if explicit_num_patches is not None else computed_num_patches + ) + resolved_num_image_tokens = ( + explicit_num_image_tokens[index] if explicit_num_image_tokens is not None else computed_num_image_tokens + ) + max_num_image_tokens = max(max_num_image_tokens, resolved_num_image_tokens) + + vision.append( + { + "num_patches": resolved_num_patches, + "h": height, + "w": width, + "num_image_tokens": resolved_num_image_tokens, + } + ) if comp_ctx_lengths_prefill is not None: lang = [] @@ -1027,7 +1115,7 @@ def get_specializations( "batch_size": 1 if continuous_batching else batch_size, "seq_len": prefill_seq_len, "ctx_len": ctx_len, - "num_image_tokens": num_image_tokens, + "num_image_tokens": max_num_image_tokens, } if continuous_batching: lang_prefill["full_batch_size"] = kv_cache_batch_size @@ -1043,7 +1131,7 @@ def get_specializations( "batch_size": full_batch_size if continuous_batching else batch_size, "seq_len": "1", "ctx_len": ctx_len, - "num_image_tokens": num_image_tokens, + "num_image_tokens": max_num_image_tokens, } if continuous_batching: @@ -1054,23 +1142,11 @@ def get_specializations( lang.append(lang_decode) else: - lang_decode = { - "batch_size": full_batch_size if continuous_batching else batch_size, - "seq_len": 1, - "ctx_len": ctx_len, - "num_image_tokens": num_image_tokens, - } - - if continuous_batching: - lang_decode["full_batch_size"] = kv_cache_batch_size - else: - lang_decode["batch_size"] = kv_cache_batch_size - lang_prefill = { "batch_size": 1 if continuous_batching else batch_size, "seq_len": prefill_seq_len, "ctx_len": ctx_len, - "num_image_tokens": num_image_tokens, + "num_image_tokens": max_num_image_tokens, } if continuous_batching: lang_prefill["full_batch_size"] = kv_cache_batch_size @@ -1079,6 +1155,18 @@ def get_specializations( if full_batch_size: lang_prefill["full_batch_exec_size"] = full_batch_size + lang_decode = { + "batch_size": full_batch_size if continuous_batching else batch_size, + "seq_len": 1, + "ctx_len": ctx_len, + "num_image_tokens": max_num_image_tokens, + } + + if continuous_batching: + lang_decode["full_batch_size"] = kv_cache_batch_size + else: + lang_decode["batch_size"] = kv_cache_batch_size + lang = [lang_prefill, lang_decode] specializations = {} @@ -1088,6 +1176,4 @@ def get_specializations( specializations["lang"] = lang return specializations, compiler_options else: - lang[0].pop("vision_size") - lang[1].pop("vision_size") - return lang, compiler_options + return [{**vision_spec, **lang_spec} for vision_spec in vision for lang_spec in lang], compiler_options diff --git a/tests/unit_test/models/test_model_quickcheck.py b/tests/unit_test/models/test_model_quickcheck.py index f77a6a49b8..5afbd25d0c 100644 --- a/tests/unit_test/models/test_model_quickcheck.py +++ b/tests/unit_test/models/test_model_quickcheck.py @@ -3489,3 +3489,43 @@ def test_layerwise_export_default_names_unchanged(tmp_path): assert f"past_key.{window}" in captured["input_names"] assert all("_vllmKvCache" not in n and "_VLLM" not in n for n in captured["output_names"]) assert all("_vllmKvCache" not in n and "_VLLM" not in n for n in captured["input_names"]) + + +def test_kimi_k25_get_specializations_supports_multi_resolution_grid_sizes(): + """Kimi K2.5 accepts list-valued raw-pixel and patch-grid sizes for multi-resolution specs.""" + from types import SimpleNamespace + + from QEfficient.transformers.models.kimi_k25.modeling_kimi_k25 import QEffKimiK25ForConditionalGeneration + + model = QEffKimiK25ForConditionalGeneration.__new__(QEffKimiK25ForConditionalGeneration) + model.config = SimpleNamespace(vision_config=SimpleNamespace(patch_size=14, merge_kernel_size=(2, 2))) + + specs, _ = model.get_specializations( + batch_size=1, + prefill_seq_len=64, + ctx_len=4096, + height=[512, 448], + width=[910, 448], + num_frames=[1, 1], + kv_offload=True, + ) + assert specs["vision"] == [ + {"num_patches": 2508, "h": 38, "w": 66, "num_image_tokens": 627}, + {"num_patches": 1024, "h": 32, "w": 32, "num_image_tokens": 256}, + ] + assert all(spec["num_image_tokens"] == 627 for spec in specs["lang"]) + + specs, _ = model.get_specializations( + batch_size=1, + prefill_seq_len=64, + ctx_len=4096, + h=[30, 32], + w=[80, 64], + num_frames=[1, 2], + kv_offload=True, + ) + assert specs["vision"] == [ + {"num_patches": 2400, "h": 30, "w": 80, "num_image_tokens": 600}, + {"num_patches": 4096, "h": 32, "w": 64, "num_image_tokens": 1024}, + ] + assert all(spec["num_image_tokens"] == 1024 for spec in specs["lang"]) From b8972004a892c3f0d7994faed2a741431456f438 Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Fri, 17 Jul 2026 14:57:55 +0530 Subject: [PATCH 22/33] refactor code and update documentation Signed-off-by: Mamta Singh --- QEfficient/utils/load_kimi_utils.py | 515 ++++++++++++++++++ README.md | 1 + docs/source/introduction.md | 1 + docs/source/release_docs.md | 9 +- docs/source/validate.md | 1 + examples/image_text_to_text/README.md | 16 + examples/kimi_k2/export_kimi_k25_vision.py | 334 +----------- examples/kimi_k2/test_kimi_k25.py | 56 +- examples/kimi_k2/test_kimi_k25_disagg.py | 28 +- .../test_continuous_batching.py | 226 +------- .../test_image_text_to_text_models.py | 201 +------ 11 files changed, 639 insertions(+), 749 deletions(-) create mode 100644 QEfficient/utils/load_kimi_utils.py diff --git a/QEfficient/utils/load_kimi_utils.py b/QEfficient/utils/load_kimi_utils.py new file mode 100644 index 0000000000..4aec58ceea --- /dev/null +++ b/QEfficient/utils/load_kimi_utils.py @@ -0,0 +1,515 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import argparse +import copy +import inspect +import json +import os +import random +import re +import sys +import tempfile +from pathlib import Path + +import torch +from huggingface_hub import snapshot_download +from safetensors import safe_open +from safetensors.torch import save_file +from transformers import AutoConfig, AutoProcessor, AutoTokenizer +from transformers.dynamic_module_utils import get_class_from_dynamic_module +from transformers.utils import import_utils as hf_import_utils + +KIMI_K25_MODEL_NAME = "moonshotai/Kimi-K2.5" +DEFAULT_MODEL_PATH = Path( + "/home/huggingface_hub/models--moonshotai--Kimi-K2.5/snapshots/4d01dfe0332d63057c186e0b262165819efb6611" +) +NUM_VISION_LAYERS = 2 +NUM_TEXT_LAYERS = 2 +LOADED_EXPERT_IDS = (0, 1, 2, 3) +NUM_EXPERTS_PER_TOKEN = 2 + +EXPERT_KEY_PATTERN = re.compile(r"^(language_model\.model\.layers\.\d+\.mlp\.experts\.)(\d+)(\..+)$") + + +def is_kimi_k25(model_name: str) -> bool: + return model_name == KIMI_K25_MODEL_NAME + + +def set_deterministic(seed: int): + import numpy as np + + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + torch.use_deterministic_algorithms(True) + + +def resolve_model_path(model_name: str = KIMI_K25_MODEL_NAME) -> Path: + os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1") + return Path(snapshot_download(repo_id=model_name, cache_dir=os.environ.get("HF_HUB_CACHE"))) + + +def ensure_torch_fx_import_compatibility(): + if hasattr(hf_import_utils, "is_torch_fx_available"): + return + + def _is_torch_fx_available() -> bool: + if not hf_import_utils.is_torch_available(): + return False + try: + import torch.fx # noqa: F401 + + return True + except Exception: + return False + + hf_import_utils.is_torch_fx_available = _is_torch_fx_available + + +def patch_kimi_tie_weights_compat(kimi_cls): + tie_signature = inspect.signature(kimi_cls.tie_weights) + if tuple(tie_signature.parameters) != ("self",): + return + + def _tie_weights_compat(self, missing_keys=None, recompute_mapping=True): + lm_tie_weights = getattr(self.language_model, "tie_weights") + try: + return lm_tie_weights(missing_keys=missing_keys, recompute_mapping=recompute_mapping) + except TypeError: + return lm_tie_weights() + + kimi_cls.tie_weights = _tie_weights_compat + + +def patch_deepseek_init_weights_compat(kimi_cls): + module_prefix, _ = kimi_cls.__module__.rsplit(".", maxsplit=1) + deepseek_module = sys.modules.get(f"{module_prefix}.modeling_deepseek") + if deepseek_module is None or not hasattr(deepseek_module, "DeepseekV3PreTrainedModel"): + return + + deepseek_cls = deepseek_module.DeepseekV3PreTrainedModel + if ( + getattr(deepseek_cls, "_qeff_kimi_k25_init_weights_patched", False) + or getattr(deepseek_cls, "_qeff_test_init_weights_patched", False) + or getattr(deepseek_cls, "_qeff_t55_init_weights_patched", False) + ): + return + + def _init_weights_compat(self, module): + std = self.config.initializer_range + if isinstance(module, torch.nn.Linear): + if hasattr(module, "weight") and module.weight is not None: + module.weight.data.normal_(mean=0.0, std=std) + if hasattr(module, "bias") and module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, torch.nn.Embedding): + if hasattr(module, "weight") and module.weight is not None: + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + deepseek_cls._init_weights = _init_weights_compat + deepseek_cls._qeff_kimi_k25_init_weights_patched = True + deepseek_cls._qeff_test_init_weights_patched = True + deepseek_cls._qeff_t55_init_weights_patched = True + + +def load_kimi_k25_class(model_path_or_name): + ensure_torch_fx_import_compatibility() + kimi_cls = get_class_from_dynamic_module( + "modeling_kimi_k25.KimiK25ForConditionalGeneration", + str(model_path_or_name), + ) + patch_kimi_tie_weights_compat(kimi_cls) + patch_deepseek_init_weights_compat(kimi_cls) + return kimi_cls + + +def patch_kimi_k25_remote_code_compat(config): + return load_kimi_k25_class(config._name_or_path) + + +def prepare_config(model_path: Path): + config = AutoConfig.from_pretrained(str(model_path), trust_remote_code=True) + + config._attn_implementation = "eager" + if hasattr(config, "text_config"): + config.text_config._attn_implementation = "eager" + if hasattr(config, "vision_config"): + config.vision_config._attn_implementation = "eager" + return config + + +def get_kimi_k25_test_config(model_name: str, model_config_dict): + config = AutoConfig.from_pretrained(model_name, trust_remote_code=True) + config._attn_implementation = "eager" + config.torch_dtype = torch.float32 + config.dtype = torch.float32 + additional_params = model_config_dict[model_name]["additional_params"] + + for attr, value in additional_params["text_config"].items(): + setattr(config.text_config, attr, value) + config.text_config._attn_implementation = "eager" + config.text_config.torch_dtype = torch.float32 + config.text_config.dtype = torch.float32 + + for attr, value in additional_params["vision_config"].items(): + setattr(config.vision_config, attr, value) + config.vision_config._attn_implementation = "eager" + config.vision_config.torch_dtype = torch.float32 + config.vision_config.dtype = torch.float32 + + patch_kimi_k25_remote_code_compat(config) + return config + + +def load_kimi_k25_model_from_config(config): + kimi_cls = patch_kimi_k25_remote_code_compat(config) + model = kimi_cls._from_config(config) + torch_dtype = getattr(model.config, "torch_dtype", None) + if torch_dtype == torch.bfloat16 or torch_dtype == torch.float16: + model = model.to(torch.float32) + return model.eval() + + +def get_kimi_k25_num_image_tokens(config, grid_thws): + merge_height, merge_width = config.vision_config.merge_kernel_size + return int(grid_thws[0, 1].item() // merge_height) * int(grid_thws[0, 2].item() // merge_width) + + +def parse_expert_ids(value: str): + expert_ids = tuple(int(expert_id) for expert_id in value.split(",") if expert_id.strip()) + if not expert_ids: + raise argparse.ArgumentTypeError("At least one expert id must be provided.") + return expert_ids + + +def _validate_layer_count(name, requested_count, available_count): + if requested_count < 1: + raise ValueError(f"{name} must be >= 1, got {requested_count}.") + if requested_count > available_count: + raise ValueError(f"{name}={requested_count} exceeds available layers={available_count}.") + + +def _validate_expert_subset(loaded_expert_ids, num_experts_per_tok, total_experts): + expert_ids = tuple(loaded_expert_ids) + if len(expert_ids) != 4: + raise ValueError(f"Expected exactly 4 routed experts, got {expert_ids!r}.") + if len(set(expert_ids)) != len(expert_ids): + raise ValueError(f"Expert ids must be unique, got {expert_ids!r}.") + invalid_ids = [expert_id for expert_id in expert_ids if expert_id < 0 or expert_id >= total_experts] + if invalid_ids: + raise ValueError(f"Expert ids {invalid_ids!r} are outside the valid range [0, {total_experts - 1}].") + if num_experts_per_tok > len(expert_ids): + raise ValueError(f"num_experts_per_tok={num_experts_per_tok} cannot exceed {len(expert_ids)} loaded experts.") + return expert_ids + + +def _remap_checkpoint_key(checkpoint_key, expert_index_map): + match = EXPERT_KEY_PATTERN.match(checkpoint_key) + if not match: + return checkpoint_key + + original_expert_idx = int(match.group(2)) + remapped_expert_idx = expert_index_map.get(original_expert_idx) + if remapped_expert_idx is None: + return None + return f"{match.group(1)}{remapped_expert_idx}{match.group(3)}" + + +def _is_routed_gate_weight(checkpoint_key): + return checkpoint_key.endswith(".mlp.gate.weight") + + +def _is_routed_gate_bias(checkpoint_key): + return checkpoint_key.endswith(".mlp.gate.e_score_correction_bias") + + +def allowed_prefixes(num_vision_layers: int, num_text_layers: int): + prefixes = [ + "vision_tower.patch_embed.", + "vision_tower.encoder.final_layernorm.", + "mm_projector.", + "language_model.model.embed_tokens.", + "language_model.model.norm.", + "language_model.lm_head.", + ] + prefixes.extend(f"vision_tower.encoder.blocks.{layer_idx}." for layer_idx in range(num_vision_layers)) + prefixes.extend(f"language_model.model.layers.{layer_idx}." for layer_idx in range(num_text_layers)) + return prefixes + + +def build_layer_subset_config(config, num_vision_layers, num_text_layers, loaded_expert_ids, num_experts_per_tok): + stripped_config = copy.deepcopy(config) + text_config = stripped_config.text_config + vision_config = stripped_config.vision_config + + _validate_layer_count("num_text_layers", num_text_layers, text_config.num_hidden_layers) + _validate_layer_count("num_vision_layers", num_vision_layers, vision_config.vt_num_hidden_layers) + + text_config.num_hidden_layers = num_text_layers + vision_config.vt_num_hidden_layers = num_vision_layers + + loaded_expert_ids = _validate_expert_subset( + loaded_expert_ids, + num_experts_per_tok, + text_config.n_routed_experts, + ) + text_config.n_routed_experts = len(loaded_expert_ids) + text_config.num_experts_per_tok = num_experts_per_tok + text_config.n_group = 1 + text_config.topk_group = 1 + return stripped_config, loaded_expert_ids + + +def materialize_subset_checkpoint( + model_path: Path, temp_model_path: Path, weight_map, allowed_weight_prefixes, loaded_expert_ids +): + expert_index_map = {expert_id: remapped_idx for remapped_idx, expert_id in enumerate(loaded_expert_ids)} + shard_to_entries = {} + for checkpoint_key, source_shard_name in weight_map.items(): + if not checkpoint_key.startswith(tuple(allowed_weight_prefixes)): + continue + + remapped_key = _remap_checkpoint_key(checkpoint_key, expert_index_map) + if remapped_key is None: + continue + shard_to_entries.setdefault(source_shard_name, []).append((checkpoint_key, remapped_key)) + + filtered_weight_map = {} + subset_shards = [] + for shard_idx, (source_shard_name, shard_entries) in enumerate(sorted(shard_to_entries.items())): + tensors = {} + with safe_open(model_path / source_shard_name, framework="pt", device="cpu") as shard_reader: + for checkpoint_key, remapped_key in shard_entries: + tensor = shard_reader.get_tensor(checkpoint_key) + if _is_routed_gate_weight(checkpoint_key): + tensor = tensor[list(loaded_expert_ids), :].contiguous() + elif _is_routed_gate_bias(checkpoint_key): + tensor = tensor[list(loaded_expert_ids)].contiguous() + tensors[remapped_key] = tensor + + subset_shard_name = f"model-subset-{shard_idx:05d}.safetensors" + save_file(tensors, str(temp_model_path / subset_shard_name)) + subset_shards.append(subset_shard_name) + filtered_weight_map.update({remapped_key: subset_shard_name for _, remapped_key in shard_entries}) + + return filtered_weight_map, subset_shards + + +def load_layer_subset_model( + *, + model_path: Path, + kimi_cls, + config, + num_vision_layers: int, + num_text_layers: int, + loaded_expert_ids, + num_experts_per_tok: int, + dtype, +): + checkpoint_index = json.loads((model_path / "model.safetensors.index.json").read_text()) + weight_map = checkpoint_index["weight_map"] + stripped_config, loaded_expert_ids = build_layer_subset_config( + config, + num_vision_layers=num_vision_layers, + num_text_layers=num_text_layers, + loaded_expert_ids=loaded_expert_ids, + num_experts_per_tok=num_experts_per_tok, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + temp_model_path = Path(tmpdir) + filtered_weight_map, subset_shards = materialize_subset_checkpoint( + model_path=model_path, + temp_model_path=temp_model_path, + weight_map=weight_map, + allowed_weight_prefixes=allowed_prefixes(num_vision_layers, num_text_layers), + loaded_expert_ids=loaded_expert_ids, + ) + (temp_model_path / "config.json").write_text(stripped_config.to_json_string(use_diff=False)) + (temp_model_path / "model.safetensors.index.json").write_text( + json.dumps( + { + "metadata": { + "total_size": sum((temp_model_path / shard_name).stat().st_size for shard_name in subset_shards) + }, + "weight_map": filtered_weight_map, + } + ) + ) + + model_kwargs = { + "config": stripped_config, + "trust_remote_code": True, + "attn_implementation": "eager", + "output_loading_info": True, + } + if dtype is not None: + model_kwargs["torch_dtype"] = dtype + + original_base_model_prefix = kimi_cls.base_model_prefix + kimi_cls.base_model_prefix = "" + try: + model, loading_info = kimi_cls.from_pretrained(str(temp_model_path), **model_kwargs) + finally: + kimi_cls.base_model_prefix = original_base_model_prefix + + unexpected_keys = loading_info["unexpected_keys"] + missing_keys = loading_info["missing_keys"] + mismatched_keys = loading_info["mismatched_keys"] + if unexpected_keys or missing_keys or mismatched_keys: + raise RuntimeError( + "Failed to load the stripped Kimi K2.5 checkpoint slice cleanly. " + f"missing={missing_keys}, unexpected={unexpected_keys}, mismatched={mismatched_keys}" + ) + model.eval() + tokenizer = AutoTokenizer.from_pretrained(str(model_path), trust_remote_code=True) + processor = AutoProcessor.from_pretrained(str(model_path), trust_remote_code=True) + + print(f"Loaded model: {type(model).__name__}") + print(f"Tokenizer vocab size: {tokenizer.vocab_size}") + print(f"Processor type: {type(processor).__name__}") + + return model, tokenizer, processor + + +def load_kimi_k25_layer_subset_model( + *, + model_path: Path | None = None, + num_vision_layers: int = NUM_VISION_LAYERS, + num_text_layers: int = NUM_TEXT_LAYERS, + loaded_expert_ids=LOADED_EXPERT_IDS, + num_experts_per_tok: int = NUM_EXPERTS_PER_TOKEN, + dtype=torch.float32, + seed: int = 1234, +): + set_deterministic(seed) + ensure_torch_fx_import_compatibility() + resolved_model_path = Path(model_path) if model_path is not None else resolve_model_path() + config = prepare_config(resolved_model_path) + kimi_cls = load_kimi_k25_class(resolved_model_path) + + model, tokenizer, processor = load_layer_subset_model( + model_path=resolved_model_path, + kimi_cls=kimi_cls, + config=config, + num_vision_layers=num_vision_layers, + num_text_layers=num_text_layers, + loaded_expert_ids=loaded_expert_ids, + num_experts_per_tok=num_experts_per_tok, + dtype=dtype, + ) + model.vision_tower.patch_embed.pos_emb.interpolation_mode = "bilinear" + return model.eval().to("cpu"), tokenizer, processor + + +@torch.no_grad() +def run_kimi_k25_hf_model_on_pytorch(model, processor, inputs, max_gen_len): + generated_ids = inputs["input_ids"] + attention_mask = inputs["attention_mask"] + pixel_values = inputs["pixel_values"] + grid_thws = inputs["grid_thws"] + new_tokens = [] + + eos_token_id = getattr(model.config, "eos_token_id", None) + if eos_token_id is None and hasattr(model.config, "text_config"): + eos_token_id = getattr(model.config.text_config, "eos_token_id", None) + + for _ in range(max_gen_len): + outputs = model( + input_ids=generated_ids, + attention_mask=attention_mask, + pixel_values=pixel_values, + grid_thws=grid_thws, + use_cache=False, + return_dict=True, + ) + logits = outputs[0] if isinstance(outputs, tuple) else outputs.logits + next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True) + new_tokens.append(next_token) + + generated_ids = torch.cat([generated_ids, next_token], dim=1) + attention_mask = torch.cat( + [ + attention_mask, + torch.ones((attention_mask.shape[0], 1), dtype=attention_mask.dtype, device=attention_mask.device), + ], + dim=1, + ) + + if eos_token_id is not None and torch.all(next_token == eos_token_id): + break + + output_tokens = torch.cat(new_tokens, dim=1).squeeze(0) + py_output = processor.tokenizer.decode(output_tokens.tolist()).strip() + print("Original HF Model Outputs (Torch CPU):") + print("Completion:", repr(py_output)) + return output_tokens + + +@torch.no_grad() +def run_kimi_k25_hf_model_on_pytorch_cb(model, processor, images, queries, max_gen_len): + generated_tokens = [] + + eos_token_id = getattr(model.config, "eos_token_id", None) + if eos_token_id is None and hasattr(model.config, "text_config"): + eos_token_id = getattr(model.config.text_config, "eos_token_id", None) + + for idx, (image, query) in enumerate(zip(images, queries)): + conversation = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": image}, + {"type": "text", "text": query}, + ], + }, + ] + inputs = processor(messages=conversation, add_generation_prompt=True, tokenize=False, return_tensors="pt") + generated_ids = inputs["input_ids"] + attention_mask = inputs["attention_mask"] + pixel_values = inputs["pixel_values"] + grid_thws = inputs["grid_thws"] + new_tokens = [] + + for _ in range(max_gen_len): + outputs = model( + input_ids=generated_ids, + attention_mask=attention_mask, + pixel_values=pixel_values, + grid_thws=grid_thws, + use_cache=False, + return_dict=True, + ) + logits = outputs[0] if isinstance(outputs, tuple) else outputs.logits + next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True) + new_tokens.append(next_token) + + generated_ids = torch.cat([generated_ids, next_token], dim=1) + attention_mask = torch.cat( + [ + attention_mask, + torch.ones((attention_mask.shape[0], 1), dtype=attention_mask.dtype, device=attention_mask.device), + ], + dim=1, + ) + + if eos_token_id is not None and torch.all(next_token == eos_token_id): + break + + output_tokens = torch.cat(new_tokens, dim=1).squeeze(0) + py_output = processor.tokenizer.decode(output_tokens.tolist()).strip() + print(f"Original HF Model Outputs (Torch CPU) for prompt {idx}:") + print("Query:", repr(query)) + print("Completion:", repr(py_output)) + generated_tokens.append(output_tokens.numpy()) + + return generated_tokens diff --git a/README.md b/README.md index be0857f40c..484d21c0f0 100644 --- a/README.md +++ b/README.md @@ -6,6 +6,7 @@ --- *Latest news* :fire:
+- [07/2026] Added support for Kimi-K2.5 vision-language model [moonshotai/Kimi-K2.5](https://huggingface.co/moonshotai/Kimi-K2.5) via `QEFFAutoModelForImageTextToText`. - [04/2026] Added WAN non-unified execution support in `QEffWanPipeline` with separate `transformer_high` and `transformer_low` modules - [04/2026] Added first-block-cache support for WAN non-unified mode and FLUX (`QEffWanPipeline`, `QEffFluxPipeline`) - [12/2025] Enabled [disaggregated serving](examples/disagg_serving) for GPT-OSS model diff --git a/docs/source/introduction.md b/docs/source/introduction.md index 971bbc3c37..f1bdc7fc07 100644 --- a/docs/source/introduction.md +++ b/docs/source/introduction.md @@ -23,6 +23,7 @@ For other models, there is comprehensive documentation to inspire upon the chang ***Latest news*** :
- [coming soon] Support for more popular [models](models_coming_soon)
+- [07/2026] Added support for Kimi-K2.5 vision-language model [moonshotai/Kimi-K2.5](https://huggingface.co/moonshotai/Kimi-K2.5) via `QEFFAutoModelForImageTextToText`. - [04/2026] Added WAN non-unified execution support in `QEffWanPipeline` with separate `transformer_high` and `transformer_low` modules - [04/2026] Added first-block-cache support for WAN non-unified mode and FLUX (`QEffWanPipeline`, `QEffFluxPipeline`) - [12/2025] Enabled [disaggregated serving](https://github.com/quic/efficient-transformers/tree/main/examples/disagg_serving) for GPT-OSS model diff --git a/docs/source/release_docs.md b/docs/source/release_docs.md index 880c3a4e4c..71678f1173 100644 --- a/docs/source/release_docs.md +++ b/docs/source/release_docs.md @@ -26,6 +26,13 @@ Welcome to the official release of **Efficient Transformer Library v1.21.0**! Th - Executable via [`QEFFAutoModelForImageTextToText`](#QEFFAutoModelForImageTextToText) - [Mistral-3.1 Example Script](https://github.com/quic/efficient-transformers/blob/main/examples/image_text_to_text/models/mistral_vision/mistral3_example.py) +- **Kimi-K2.5 (Vision Language)** + - Executable via [`QEFFAutoModelForImageTextToText`](#QEFFAutoModelForImageTextToText) + - Supports `moonshotai/Kimi-K2.5` with `KimiK25ForConditionalGeneration` + - Supports single-QPC and dual-QPC (`kv_offload=True`) vision-language compilation + - Uses Kimi-specific image-grid specializations: `num_patches`, `h`, `w`, and `num_image_tokens` + - [Kimi-K2.5 Vision Example Script](https://github.com/quic/efficient-transformers/blob/main/examples/kimi_k2/export_kimi_k25_vision.py) + - **Disaggregated serving ready via vLLM GPT-OSS** > **Note**: If running GPT-OSS models natively via vLLM, PR-685 of the qefficient library is required for Python 3.12 compatibility. @@ -78,7 +85,7 @@ Welcome to the official release of **Efficient Transformer Library v1.21.0**! Th - **Compute-Context-Length (CCL) support**: To optimize the throughput when handling very large context lengths - **Prefill/Decode Separation**: Support for GPT OSS using disaggregate serving models - **Continuous Batching (VLMs)**: Extended to Vision Language Models with multi-image handling - - Supported models: Llava, Llava_Next, Gemma3, Mistral3, InternVL2_5, InternVL3_5, Molmo + - Supported models: Llava, Llava_Next, Gemma3, Mistral3, InternVL2_5, InternVL3_5, Molmo, Kimi-K2.5 - **ONNX Sub-Functions**: Feature enabling more efficient model compilation and execution on hardware. Users can enable the feature by passing `use_onnx_subfunctions=True` during export - **Memory Profiling**: Built-in utilities for optimization analysis - **Extend on-device Sampling**: Extend on-device sampling to dual QPC VLMs and Guided decoding for on-device sampling diff --git a/docs/source/validate.md b/docs/source/validate.md index 6fd58484a6..1c3e682b44 100644 --- a/docs/source/validate.md +++ b/docs/source/validate.md @@ -90,6 +90,7 @@ | **Qwen3_5MoeForConditionalGeneration** | Qwen3.5 | [Qwen/Qwen3.5-122B-A10B](https://huggingface.co/Qwen/Qwen3.5-122B-A10B)
[Qwen/Qwen3.5-35B-A3B](https://huggingface.co/Qwen/Qwen3.5-35B-A3B) | ✕ | ✔️ | ✕ | ✔️ | | **Qwen3_5ForConditionalGeneration** | Qwen3.6 | [Qwen/Qwen3.6-27B](https://huggingface.co/Qwen/Qwen3.6-27B) | ✕ | ✔️ | ✕ | ✔️ | | **Qwen3_5MoeForConditionalGeneration** | Qwen3.6 | [Qwen/Qwen3.6-35B-A3B](https://huggingface.co/Qwen/Qwen3.6-35B-A3B) | ✕ | ✔️ | ✕ | ✔️ | +| **KimiK25ForConditionalGeneration** | Kimi-K2.5 | [moonshotai/Kimi-K2.5](https://huggingface.co/moonshotai/Kimi-K2.5) | ✔️ | ✔️ | ✕ | ✕ | | **Mistral3ForConditionalGeneration** | Mistral3| [mistralai/Mistral-Small-3.1-24B-Instruct-2503](https://huggingface.co/mistralai/Mistral-Small-3.1-24B-Instruct-2503)| ✕ | ✔️ | ✕ | ✕ | ### Vision-Language Reranker Models (Text + Image Scoring) diff --git a/examples/image_text_to_text/README.md b/examples/image_text_to_text/README.md index 612f9d400a..6aac140698 100644 --- a/examples/image_text_to_text/README.md +++ b/examples/image_text_to_text/README.md @@ -91,6 +91,7 @@ Popular model families include: - InternVL - Molmo - LLaVA +- Kimi-K2.5 ### Model-Specific Examples @@ -107,6 +108,21 @@ Some models have specialized examples demonstrating advanced features: | **Granite** | [models/granite_vision/](models/granite_vision/) | | **InternVL** | [models/internvl/](models/internvl/) | | **Molmo** | [models/molmo/](models/molmo/) | +| **Kimi-K2.5** | [../kimi_k2/export_kimi_k25_vision.py](../kimi_k2/export_kimi_k25_vision.py) | + +### Kimi-K2.5 Vision + +Kimi-K2.5 uses the `moonshotai/Kimi-K2.5` remote-code model with the `KimiK25ForConditionalGeneration` architecture. The example script loads either the full checkpoint or a smaller layer subset, wraps it with `QEFFAutoModelForImageTextToText`, and compiles the vision-language path with Kimi-specific image-grid specializations. + +```bash + +python ../kimi_k2/export_kimi_k25_vision.py \ + --model-path "$HF_HUB_CACHE/models--moonshotai--Kimi-K2.5/snapshots/" \ + --image-url "https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png" \ + --prompt "Describe this image." +``` + +For custom image sizes, pass matching compile specializations for the processed image grid: `num_patches`, `h`, `w`, and `num_image_tokens`. The defaults in `../kimi_k2/export_kimi_k25_vision.py` demonstrate the expected values for the sample image and keep the language-side image embedding shape bounded to the actual image token count. For reranker examples, see [../reranker/](../reranker/). diff --git a/examples/kimi_k2/export_kimi_k25_vision.py b/examples/kimi_k2/export_kimi_k25_vision.py index 54d24e663d..96f9d086ed 100644 --- a/examples/kimi_k2/export_kimi_k25_vision.py +++ b/examples/kimi_k2/export_kimi_k25_vision.py @@ -6,318 +6,39 @@ # ----------------------------------------------------------------------------- import argparse -import copy -import inspect -import json -import re -import sys -import tempfile from io import BytesIO from pathlib import Path import requests import torch from PIL import Image -from safetensors import safe_open -from safetensors.torch import save_file -from transformers import AutoConfig, AutoProcessor, AutoTokenizer -from transformers.dynamic_module_utils import get_class_from_dynamic_module -from transformers.utils import import_utils as hf_import_utils from QEfficient import QEFFAutoModelForImageTextToText - -MODEL_PATH = Path( - "/home/huggingface_hub/models--moonshotai--Kimi-K2.5/snapshots/4d01dfe0332d63057c186e0b262165819efb6611" +from QEfficient.utils.load_kimi_utils import ( + DEFAULT_MODEL_PATH as MODEL_PATH, +) +from QEfficient.utils.load_kimi_utils import ( + LOADED_EXPERT_IDS, + NUM_EXPERTS_PER_TOKEN, + NUM_TEXT_LAYERS, + NUM_VISION_LAYERS, + load_kimi_k25_class, +) +from QEfficient.utils.load_kimi_utils import ( + ensure_torch_fx_import_compatibility as _ensure_torch_fx_import_compatibility, +) +from QEfficient.utils.load_kimi_utils import ( + load_layer_subset_model as _load_layer_subset_model, +) +from QEfficient.utils.load_kimi_utils import ( + parse_expert_ids as _parse_expert_ids, +) +from QEfficient.utils.load_kimi_utils import ( + prepare_config as _prepare_config, +) +from QEfficient.utils.load_kimi_utils import ( + set_deterministic as _set_deterministic, ) -NUM_VISION_LAYERS = 2 -NUM_TEXT_LAYERS = 2 -LOADED_EXPERT_IDS = (0, 1, 2, 3) -NUM_EXPERTS_PER_TOKEN = 2 - -EXPERT_KEY_PATTERN = re.compile(r"^(language_model\.model\.layers\.\d+\.mlp\.experts\.)(\d+)(\..+)$") - - -def _set_deterministic(seed: int): - import random - - import numpy as np - - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(seed) - torch.use_deterministic_algorithms(True) - - -def _ensure_torch_fx_import_compatibility(): - """Backfill `is_torch_fx_available` for remote model code expecting older Transformers APIs.""" - if hasattr(hf_import_utils, "is_torch_fx_available"): - return - - def _is_torch_fx_available() -> bool: - if not hf_import_utils.is_torch_available(): - return False - try: - import torch.fx # noqa: F401 - - return True - except Exception: - return False - - hf_import_utils.is_torch_fx_available = _is_torch_fx_available - - -def _patch_kimi_tie_weights_compat(kimi_cls): - tie_signature = inspect.signature(kimi_cls.tie_weights) - if tuple(tie_signature.parameters) != ("self",): - return - - def _tie_weights_compat(self, missing_keys=None, recompute_mapping=True): - lm_tie_weights = getattr(self.language_model, "tie_weights") - try: - return lm_tie_weights(missing_keys=missing_keys, recompute_mapping=recompute_mapping) - except TypeError: - return lm_tie_weights() - - kimi_cls.tie_weights = _tie_weights_compat - - -def _patch_deepseek_init_weights_compat(kimi_cls): - module_prefix, _ = kimi_cls.__module__.rsplit(".", maxsplit=1) - deepseek_module_name = f"{module_prefix}.modeling_deepseek" - deepseek_module = sys.modules.get(deepseek_module_name) - if deepseek_module is None or not hasattr(deepseek_module, "DeepseekV3PreTrainedModel"): - return - - deepseek_cls = deepseek_module.DeepseekV3PreTrainedModel - if getattr(deepseek_cls, "_qeff_t55_init_weights_patched", False): - return - - def _init_weights_compat(self, module): - std = self.config.initializer_range - if isinstance(module, torch.nn.Linear): - if hasattr(module, "weight") and module.weight is not None: - module.weight.data.normal_(mean=0.0, std=std) - if hasattr(module, "bias") and module.bias is not None: - module.bias.data.zero_() - elif isinstance(module, torch.nn.Embedding): - if hasattr(module, "weight") and module.weight is not None: - module.weight.data.normal_(mean=0.0, std=std) - if module.padding_idx is not None: - module.weight.data[module.padding_idx].zero_() - - deepseek_cls._init_weights = _init_weights_compat - deepseek_cls._qeff_t55_init_weights_patched = True - - -def _prepare_config(model_path: Path): - config = AutoConfig.from_pretrained(str(model_path), trust_remote_code=True) - - # Avoid FA2 dispatch checks in Transformers 5.5.x for this remote code. - config._attn_implementation = "eager" - if hasattr(config, "text_config"): - config.text_config._attn_implementation = "eager" - if hasattr(config, "vision_config"): - config.vision_config._attn_implementation = "eager" - return config - - -def _parse_expert_ids(value: str): - expert_ids = tuple(int(expert_id) for expert_id in value.split(",") if expert_id.strip()) - if not expert_ids: - raise argparse.ArgumentTypeError("At least one expert id must be provided.") - return expert_ids - - -def _validate_layer_count(name, requested_count, available_count): - if requested_count < 1: - raise ValueError(f"{name} must be >= 1, got {requested_count}.") - if requested_count > available_count: - raise ValueError(f"{name}={requested_count} exceeds available layers={available_count}.") - - -def _validate_expert_subset(loaded_expert_ids, num_experts_per_tok, total_experts): - expert_ids = tuple(loaded_expert_ids) - if len(expert_ids) != 4: - raise ValueError(f"Expected exactly 4 routed experts, got {expert_ids!r}.") - if len(set(expert_ids)) != len(expert_ids): - raise ValueError(f"Expert ids must be unique, got {expert_ids!r}.") - invalid_ids = [expert_id for expert_id in expert_ids if expert_id < 0 or expert_id >= total_experts] - if invalid_ids: - raise ValueError(f"Expert ids {invalid_ids!r} are outside the valid range [0, {total_experts - 1}].") - if num_experts_per_tok > len(expert_ids): - raise ValueError(f"num_experts_per_tok={num_experts_per_tok} cannot exceed {len(expert_ids)} loaded experts.") - return expert_ids - - -def _remap_checkpoint_key(checkpoint_key, expert_index_map): - match = EXPERT_KEY_PATTERN.match(checkpoint_key) - if not match: - return checkpoint_key - - original_expert_idx = int(match.group(2)) - remapped_expert_idx = expert_index_map.get(original_expert_idx) - if remapped_expert_idx is None: - return None - return f"{match.group(1)}{remapped_expert_idx}{match.group(3)}" - - -def _is_routed_gate_weight(checkpoint_key): - return checkpoint_key.endswith(".mlp.gate.weight") - - -def _is_routed_gate_bias(checkpoint_key): - return checkpoint_key.endswith(".mlp.gate.e_score_correction_bias") - - -def _allowed_prefixes(num_vision_layers: int, num_text_layers: int): - prefixes = [ - "vision_tower.patch_embed.", - "vision_tower.encoder.final_layernorm.", - "mm_projector.", - "language_model.model.embed_tokens.", - "language_model.model.norm.", - "language_model.lm_head.", - ] - prefixes.extend(f"vision_tower.encoder.blocks.{layer_idx}." for layer_idx in range(num_vision_layers)) - prefixes.extend(f"language_model.model.layers.{layer_idx}." for layer_idx in range(num_text_layers)) - return prefixes - - -def _build_layer_subset_config(config, num_vision_layers, num_text_layers, loaded_expert_ids, num_experts_per_tok): - stripped_config = copy.deepcopy(config) - text_config = stripped_config.text_config - vision_config = stripped_config.vision_config - - _validate_layer_count("num_text_layers", num_text_layers, text_config.num_hidden_layers) - _validate_layer_count("num_vision_layers", num_vision_layers, vision_config.vt_num_hidden_layers) - - text_config.num_hidden_layers = num_text_layers - vision_config.vt_num_hidden_layers = num_vision_layers - - loaded_expert_ids = _validate_expert_subset( - loaded_expert_ids, - num_experts_per_tok, - text_config.n_routed_experts, - ) - text_config.n_routed_experts = len(loaded_expert_ids) - text_config.num_experts_per_tok = num_experts_per_tok - text_config.n_group = 1 - text_config.topk_group = 1 - return stripped_config, loaded_expert_ids - - -def _materialize_subset_checkpoint( - model_path: Path, temp_model_path: Path, weight_map, allowed_prefixes, loaded_expert_ids -): - expert_index_map = {expert_id: remapped_idx for remapped_idx, expert_id in enumerate(loaded_expert_ids)} - shard_to_entries = {} - for checkpoint_key, source_shard_name in weight_map.items(): - if not checkpoint_key.startswith(tuple(allowed_prefixes)): - continue - - remapped_key = _remap_checkpoint_key(checkpoint_key, expert_index_map) - if remapped_key is None: - continue - shard_to_entries.setdefault(source_shard_name, []).append((checkpoint_key, remapped_key)) - - filtered_weight_map = {} - subset_shards = [] - for shard_idx, (source_shard_name, shard_entries) in enumerate(sorted(shard_to_entries.items())): - tensors = {} - with safe_open(model_path / source_shard_name, framework="pt", device="cpu") as shard_reader: - for checkpoint_key, remapped_key in shard_entries: - tensor = shard_reader.get_tensor(checkpoint_key) - if _is_routed_gate_weight(checkpoint_key): - tensor = tensor[list(loaded_expert_ids), :].contiguous() - elif _is_routed_gate_bias(checkpoint_key): - tensor = tensor[list(loaded_expert_ids)].contiguous() - tensors[remapped_key] = tensor - - subset_shard_name = f"model-subset-{shard_idx:05d}.safetensors" - save_file(tensors, str(temp_model_path / subset_shard_name)) - subset_shards.append(subset_shard_name) - filtered_weight_map.update({remapped_key: subset_shard_name for _, remapped_key in shard_entries}) - - return filtered_weight_map, subset_shards - - -def _load_layer_subset_model( - *, - model_path: Path, - kimi_cls, - config, - num_vision_layers: int, - num_text_layers: int, - loaded_expert_ids, - num_experts_per_tok: int, - dtype, -): - checkpoint_index = json.loads((model_path / "model.safetensors.index.json").read_text()) - weight_map = checkpoint_index["weight_map"] - stripped_config, loaded_expert_ids = _build_layer_subset_config( - config, - num_vision_layers=num_vision_layers, - num_text_layers=num_text_layers, - loaded_expert_ids=loaded_expert_ids, - num_experts_per_tok=num_experts_per_tok, - ) - - with tempfile.TemporaryDirectory() as tmpdir: - temp_model_path = Path(tmpdir) - filtered_weight_map, subset_shards = _materialize_subset_checkpoint( - model_path=model_path, - temp_model_path=temp_model_path, - weight_map=weight_map, - allowed_prefixes=_allowed_prefixes(num_vision_layers, num_text_layers), - loaded_expert_ids=loaded_expert_ids, - ) - (temp_model_path / "config.json").write_text(stripped_config.to_json_string(use_diff=False)) - (temp_model_path / "model.safetensors.index.json").write_text( - json.dumps( - { - "metadata": { - "total_size": sum((temp_model_path / shard_name).stat().st_size for shard_name in subset_shards) - }, - "weight_map": filtered_weight_map, - } - ) - ) - - model_kwargs = { - "config": stripped_config, - "trust_remote_code": True, - "attn_implementation": "eager", - "output_loading_info": True, - } - if dtype is not None: - model_kwargs["torch_dtype"] = dtype - - original_base_model_prefix = kimi_cls.base_model_prefix - kimi_cls.base_model_prefix = "" - try: - model, loading_info = kimi_cls.from_pretrained(str(temp_model_path), **model_kwargs) - finally: - kimi_cls.base_model_prefix = original_base_model_prefix - - unexpected_keys = loading_info["unexpected_keys"] - missing_keys = loading_info["missing_keys"] - mismatched_keys = loading_info["mismatched_keys"] - if unexpected_keys or missing_keys or mismatched_keys: - raise RuntimeError( - "Failed to load the stripped Kimi K2.5 checkpoint slice cleanly. " - f"missing={missing_keys}, unexpected={unexpected_keys}, mismatched={mismatched_keys}" - ) - model.eval() - tokenizer = AutoTokenizer.from_pretrained(str(model_path), trust_remote_code=True) - processor = AutoProcessor.from_pretrained(str(model_path), trust_remote_code=True) - - print(f"Loaded model: {type(model).__name__}") - print(f"Tokenizer vocab size: {tokenizer.vocab_size}") - print(f"Processor type: {type(processor).__name__}") - - return model, tokenizer, processor def parse_args(): @@ -348,12 +69,7 @@ def main(): _set_deterministic(1234) _ensure_torch_fx_import_compatibility() config = _prepare_config(args.model_path) - kimi_cls = get_class_from_dynamic_module( - "modeling_kimi_k25.KimiK25ForConditionalGeneration", - str(args.model_path), - ) - _patch_kimi_tie_weights_compat(kimi_cls) - _patch_deepseek_init_weights_compat(kimi_cls) + kimi_cls = load_kimi_k25_class(args.model_path) model_kwargs = { "config": config, diff --git a/examples/kimi_k2/test_kimi_k25.py b/examples/kimi_k2/test_kimi_k25.py index ec3fc4a2c7..0a13072779 100644 --- a/examples/kimi_k2/test_kimi_k25.py +++ b/examples/kimi_k2/test_kimi_k25.py @@ -6,30 +6,37 @@ # ----------------------------------------------------------------------------- import copy -import os from io import BytesIO -from pathlib import Path -import numpy as np import requests import torch -from export_kimi_k25_vision import ( - LOADED_EXPERT_IDS, - NUM_EXPERTS_PER_TOKEN, - _ensure_torch_fx_import_compatibility, - _load_layer_subset_model, - _patch_deepseek_init_weights_compat, - _patch_kimi_tie_weights_compat, - _prepare_config, -) -from huggingface_hub import snapshot_download from PIL import Image -from transformers.dynamic_module_utils import get_class_from_dynamic_module from QEfficient import QEFFAutoModelForImageTextToText from QEfficient.generation.cloud_infer import QAICInferenceSession +from QEfficient.utils.load_kimi_utils import ( + KIMI_K25_MODEL_NAME, + LOADED_EXPERT_IDS, + NUM_EXPERTS_PER_TOKEN, + load_kimi_k25_class, +) +from QEfficient.utils.load_kimi_utils import ( + ensure_torch_fx_import_compatibility as _ensure_torch_fx_import_compatibility, +) +from QEfficient.utils.load_kimi_utils import ( + load_layer_subset_model as _load_layer_subset_model, +) +from QEfficient.utils.load_kimi_utils import ( + prepare_config as _prepare_config, +) +from QEfficient.utils.load_kimi_utils import ( + resolve_model_path as _resolve_model_path, +) +from QEfficient.utils.load_kimi_utils import ( + set_deterministic as _set_deterministic, +) -MODEL_NAME = "moonshotai/Kimi-K2.5" +MODEL_NAME = KIMI_K25_MODEL_NAME IMAGE_URL = "https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png" TEXT_PROMPT = "Describe this image." NUM_VISION_LAYERS = 2 @@ -52,21 +59,6 @@ def _has_qaic_runtime_access() -> bool: return False -def _set_deterministic(seed: int): - import random - - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(seed) - torch.use_deterministic_algorithms(True) - - -def _resolve_model_path() -> Path: - return Path(snapshot_download(repo_id=MODEL_NAME, cache_dir=os.environ.get("HF_HUB_CACHE"))) - - def _prepare_inputs(processor): image = Image.open(BytesIO(requests.get(IMAGE_URL, timeout=30).content)).convert("RGB") messages = [ @@ -176,9 +168,7 @@ def check_kimi_k25_pytorch_vs_ai100(): _ensure_torch_fx_import_compatibility() model_path = _resolve_model_path() config = _prepare_config(model_path) - kimi_cls = get_class_from_dynamic_module("modeling_kimi_k25.KimiK25ForConditionalGeneration", str(model_path)) - _patch_kimi_tie_weights_compat(kimi_cls) - _patch_deepseek_init_weights_compat(kimi_cls) + kimi_cls = load_kimi_k25_class(model_path) model, tokenizer, processor = _load_layer_subset_model( model_path=model_path, diff --git a/examples/kimi_k2/test_kimi_k25_disagg.py b/examples/kimi_k2/test_kimi_k25_disagg.py index 4efd432a05..c83ee01b63 100644 --- a/examples/kimi_k2/test_kimi_k25_disagg.py +++ b/examples/kimi_k2/test_kimi_k25_disagg.py @@ -11,15 +11,6 @@ import numpy as np import pytest import torch -from export_kimi_k25_vision import ( - LOADED_EXPERT_IDS, - NUM_EXPERTS_PER_TOKEN, - _ensure_torch_fx_import_compatibility, - _load_layer_subset_model, - _patch_deepseek_init_weights_compat, - _patch_kimi_tie_weights_compat, - _prepare_config, -) from test_kimi_k25 import ( _clone_inputs, _decode_tokens, @@ -29,10 +20,23 @@ _resolve_model_path, _set_deterministic, ) -from transformers.dynamic_module_utils import get_class_from_dynamic_module from QEfficient import QEFFAutoModelForImageTextToText from QEfficient.generation.cloud_infer import QAICInferenceSession +from QEfficient.utils.load_kimi_utils import ( + LOADED_EXPERT_IDS, + NUM_EXPERTS_PER_TOKEN, + load_kimi_k25_class, +) +from QEfficient.utils.load_kimi_utils import ( + ensure_torch_fx_import_compatibility as _ensure_torch_fx_import_compatibility, +) +from QEfficient.utils.load_kimi_utils import ( + load_layer_subset_model as _load_layer_subset_model, +) +from QEfficient.utils.load_kimi_utils import ( + prepare_config as _prepare_config, +) PREFILL_SEQ_LEN = 512 CTX_LEN = 2048 @@ -98,9 +102,7 @@ def _load_kimi_subset_model(): _ensure_torch_fx_import_compatibility() model_path = _resolve_model_path() config = _prepare_config(model_path) - kimi_cls = get_class_from_dynamic_module("modeling_kimi_k25.KimiK25ForConditionalGeneration", str(model_path)) - _patch_kimi_tie_weights_compat(kimi_cls) - _patch_deepseek_init_weights_compat(kimi_cls) + kimi_cls = load_kimi_k25_class(model_path) model, tokenizer, processor = _load_layer_subset_model( model_path=model_path, diff --git a/tests/transformers/models/image_text_to_text/test_continuous_batching.py b/tests/transformers/models/image_text_to_text/test_continuous_batching.py index f5d3927172..8258fb3a36 100644 --- a/tests/transformers/models/image_text_to_text/test_continuous_batching.py +++ b/tests/transformers/models/image_text_to_text/test_continuous_batching.py @@ -6,10 +6,8 @@ # ---------------------------------------------------------------------------- import copy -import inspect import json import os -import sys from io import BytesIO from typing import Optional @@ -23,10 +21,24 @@ AutoTokenizer, GenerationConfig, ) -from transformers.dynamic_module_utils import get_class_from_dynamic_module -from transformers.utils import import_utils as hf_import_utils from QEfficient import QEFFAutoModelForCausalLM, QEFFAutoModelForImageTextToText +from QEfficient.utils.load_kimi_utils import ( + get_kimi_k25_num_image_tokens as _get_kimi_k25_num_image_tokens, +) +from QEfficient.utils.load_kimi_utils import ( + get_kimi_k25_test_config, + load_kimi_k25_layer_subset_model, +) +from QEfficient.utils.load_kimi_utils import ( + is_kimi_k25 as _is_kimi_k25, +) +from QEfficient.utils.load_kimi_utils import ( + load_kimi_k25_model_from_config as _load_kimi_k25_model_from_config, +) +from QEfficient.utils.load_kimi_utils import ( + run_kimi_k25_hf_model_on_pytorch_cb as _run_kimi_k25_hf_model_on_pytorch_CB, +) from QEfficient.utils.run_utils import ApiRunnerInternVL, ApiRunnerMolmo, ApiRunnerVlm from QEfficient.utils.test_utils import ( InternProcessor, @@ -44,216 +56,14 @@ model_config_dict = {model["model_name"]: model for model in multimodal_models} NEW_GENERATION_TOKENS = 10 -KIMI_K25_MODEL_NAME = "moonshotai/Kimi-K2.5" - - -def _is_kimi_k25(model_name: str) -> bool: - return model_name == KIMI_K25_MODEL_NAME - - -def _ensure_torch_fx_import_compatibility(): - if hasattr(hf_import_utils, "is_torch_fx_available"): - return - - def _is_torch_fx_available() -> bool: - if not hf_import_utils.is_torch_available(): - return False - try: - import torch.fx # noqa: F401 - - return True - except Exception: - return False - - hf_import_utils.is_torch_fx_available = _is_torch_fx_available - - -def _patch_kimi_k25_tie_weights_compat(kimi_cls): - tie_signature = inspect.signature(kimi_cls.tie_weights) - if tuple(tie_signature.parameters) != ("self",): - return - - def _tie_weights_compat(self, missing_keys=None, recompute_mapping=True): - lm_tie_weights = getattr(self.language_model, "tie_weights") - try: - return lm_tie_weights(missing_keys=missing_keys, recompute_mapping=recompute_mapping) - except TypeError: - return lm_tie_weights() - - kimi_cls.tie_weights = _tie_weights_compat - - -def _patch_kimi_k25_deepseek_init_weights_compat(kimi_cls): - module_prefix, _ = kimi_cls.__module__.rsplit(".", maxsplit=1) - deepseek_module = sys.modules.get(f"{module_prefix}.modeling_deepseek") - if deepseek_module is None or not hasattr(deepseek_module, "DeepseekV3PreTrainedModel"): - return - - deepseek_cls = deepseek_module.DeepseekV3PreTrainedModel - if getattr(deepseek_cls, "_qeff_test_init_weights_patched", False): - return - - def _init_weights_compat(self, module): - std = self.config.initializer_range - if isinstance(module, torch.nn.Linear): - if hasattr(module, "weight") and module.weight is not None: - module.weight.data.normal_(mean=0.0, std=std) - if hasattr(module, "bias") and module.bias is not None: - module.bias.data.zero_() - elif isinstance(module, torch.nn.Embedding): - if hasattr(module, "weight") and module.weight is not None: - module.weight.data.normal_(mean=0.0, std=std) - if module.padding_idx is not None: - module.weight.data[module.padding_idx].zero_() - - deepseek_cls._init_weights = _init_weights_compat - deepseek_cls._qeff_test_init_weights_patched = True - - -def _patch_kimi_k25_remote_code_compat(config): - _ensure_torch_fx_import_compatibility() - kimi_cls = get_class_from_dynamic_module("modeling_kimi_k25.KimiK25ForConditionalGeneration", config._name_or_path) - _patch_kimi_k25_tie_weights_compat(kimi_cls) - _patch_kimi_k25_deepseek_init_weights_compat(kimi_cls) - return kimi_cls def _get_kimi_k25_test_config(model_name: str): - config = AutoConfig.from_pretrained(model_name, trust_remote_code=True) - config._attn_implementation = "eager" - config.torch_dtype = torch.float32 - config.dtype = torch.float32 - additional_params = model_config_dict[model_name]["additional_params"] - - for attr, value in additional_params["text_config"].items(): - setattr(config.text_config, attr, value) - config.text_config._attn_implementation = "eager" - config.text_config.torch_dtype = torch.float32 - config.text_config.dtype = torch.float32 - - for attr, value in additional_params["vision_config"].items(): - setattr(config.vision_config, attr, value) - config.vision_config._attn_implementation = "eager" - config.vision_config.torch_dtype = torch.float32 - config.vision_config.dtype = torch.float32 - - _patch_kimi_k25_remote_code_compat(config) - return config - - -def _load_kimi_k25_model_from_config(config): - kimi_cls = _patch_kimi_k25_remote_code_compat(config) - model = kimi_cls._from_config(config) - torch_dtype = getattr(model.config, "torch_dtype", None) - if torch_dtype == torch.bfloat16 or torch_dtype == torch.float16: - model = model.to(torch.float32) - return model.eval() - - -def _get_kimi_k25_num_image_tokens(config, grid_thws): - merge_height, merge_width = config.vision_config.merge_kernel_size - return int(grid_thws[0, 1].item() // merge_height) * int(grid_thws[0, 2].item() // merge_width) - - -@torch.no_grad() -def _run_kimi_k25_hf_model_on_pytorch_CB(model, processor, images, queries, max_gen_len): - generated_tokens = [] - - eos_token_id = getattr(model.config, "eos_token_id", None) - if eos_token_id is None and hasattr(model.config, "text_config"): - eos_token_id = getattr(model.config.text_config, "eos_token_id", None) - - for idx, (image, query) in enumerate(zip(images, queries)): - conversation = [ - { - "role": "user", - "content": [ - {"type": "image_url", "image_url": image}, - {"type": "text", "text": query}, - ], - }, - ] - inputs = processor(messages=conversation, add_generation_prompt=True, tokenize=False, return_tensors="pt") - generated_ids = inputs["input_ids"] - attention_mask = inputs["attention_mask"] - pixel_values = inputs["pixel_values"] - grid_thws = inputs["grid_thws"] - new_tokens = [] - - for _ in range(max_gen_len): - outputs = model( - input_ids=generated_ids, - attention_mask=attention_mask, - pixel_values=pixel_values, - grid_thws=grid_thws, - use_cache=False, - return_dict=True, - ) - logits = outputs[0] if isinstance(outputs, tuple) else outputs.logits - next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True) - new_tokens.append(next_token) - - generated_ids = torch.cat([generated_ids, next_token], dim=1) - attention_mask = torch.cat( - [ - attention_mask, - torch.ones((attention_mask.shape[0], 1), dtype=attention_mask.dtype, device=attention_mask.device), - ], - dim=1, - ) - - if eos_token_id is not None and torch.all(next_token == eos_token_id): - break - - output_tokens = torch.cat(new_tokens, dim=1).squeeze(0) - py_output = processor.tokenizer.decode(output_tokens.tolist()).strip() - print(f"Original HF Model Outputs (Torch CPU) for prompt {idx}:") - print("Query:", repr(query)) - print("Completion:", repr(py_output)) - generated_tokens.append(output_tokens.numpy()) - - return generated_tokens + return get_kimi_k25_test_config(model_name, model_config_dict) def _load_kimi_k25_layer_subset_model(): - examples_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../examples/kimi_k2")) - if examples_dir not in sys.path: - sys.path.insert(0, examples_dir) - - from test_kimi_k25 import ( # noqa: PLC0415 - LOADED_EXPERT_IDS, - NUM_EXPERTS_PER_TOKEN, - NUM_TEXT_LAYERS, - NUM_VISION_LAYERS, - _load_layer_subset_model, - _patch_deepseek_init_weights_compat, - _patch_kimi_tie_weights_compat, - _prepare_config, - _resolve_model_path, - _set_deterministic, - ) - - _set_deterministic(1234) - _ensure_torch_fx_import_compatibility() - model_path = _resolve_model_path() - config = _prepare_config(model_path) - kimi_cls = get_class_from_dynamic_module("modeling_kimi_k25.KimiK25ForConditionalGeneration", str(model_path)) - _patch_kimi_tie_weights_compat(kimi_cls) - _patch_deepseek_init_weights_compat(kimi_cls) - - model, tokenizer, processor = _load_layer_subset_model( - model_path=model_path, - kimi_cls=kimi_cls, - config=config, - num_vision_layers=min(NUM_VISION_LAYERS, 1), - num_text_layers=min(NUM_TEXT_LAYERS, 1), - loaded_expert_ids=LOADED_EXPERT_IDS, - num_experts_per_tok=NUM_EXPERTS_PER_TOKEN, - dtype=torch.float32, - ) - model.vision_tower.patch_embed.pos_emb.interpolation_mode = "bilinear" - model = model.eval().to("cpu") - return model, tokenizer, processor + return load_kimi_k25_layer_subset_model(num_vision_layers=1, num_text_layers=1) def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100_CB( diff --git a/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py b/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py index df923b22be..d26a1c2a14 100644 --- a/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py +++ b/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py @@ -6,10 +6,8 @@ # ---------------------------------------------------------------------------- import copy -import inspect import json import os -import sys from io import BytesIO from typing import List, Optional @@ -24,12 +22,25 @@ GenerationConfig, TextStreamer, ) -from transformers.dynamic_module_utils import get_class_from_dynamic_module -from transformers.utils import import_utils as hf_import_utils from QEfficient import QEFFAutoModelForCausalLM, QEFFAutoModelForImageTextToText from QEfficient.utils._utils import create_json from QEfficient.utils.constants import QnnConstants +from QEfficient.utils.load_kimi_utils import ( + get_kimi_k25_num_image_tokens as _get_kimi_k25_num_image_tokens, +) +from QEfficient.utils.load_kimi_utils import ( + get_kimi_k25_test_config, +) +from QEfficient.utils.load_kimi_utils import ( + is_kimi_k25 as _is_kimi_k25, +) +from QEfficient.utils.load_kimi_utils import ( + load_kimi_k25_layer_subset_model as _load_kimi_k25_layer_subset_model, +) +from QEfficient.utils.load_kimi_utils import ( + run_kimi_k25_hf_model_on_pytorch as _run_kimi_k25_hf_model_on_pytorch, +) from QEfficient.utils.run_utils import ApiRunnerInternVL, ApiRunnerMolmo, ApiRunnerVlm from QEfficient.utils.test_utils import ( InternProcessor, @@ -50,190 +61,10 @@ test_mm_moe_models = [model["model_name"] for model in multimodal_models if "moe" in model.get("model_type", "")] NEW_GENERATION_TOKENS = 10 -KIMI_K25_MODEL_NAME = "moonshotai/Kimi-K2.5" - - -def _is_kimi_k25(model_name: str) -> bool: - return model_name == KIMI_K25_MODEL_NAME - - -def _ensure_torch_fx_import_compatibility(): - if hasattr(hf_import_utils, "is_torch_fx_available"): - return - - def _is_torch_fx_available() -> bool: - if not hf_import_utils.is_torch_available(): - return False - try: - import torch.fx # noqa: F401 - - return True - except Exception: - return False - - hf_import_utils.is_torch_fx_available = _is_torch_fx_available - - -def _patch_kimi_k25_tie_weights_compat(kimi_cls): - tie_signature = inspect.signature(kimi_cls.tie_weights) - if tuple(tie_signature.parameters) != ("self",): - return - - def _tie_weights_compat(self, missing_keys=None, recompute_mapping=True): - lm_tie_weights = getattr(self.language_model, "tie_weights") - try: - return lm_tie_weights(missing_keys=missing_keys, recompute_mapping=recompute_mapping) - except TypeError: - return lm_tie_weights() - - kimi_cls.tie_weights = _tie_weights_compat - - -def _patch_kimi_k25_deepseek_init_weights_compat(kimi_cls): - module_prefix, _ = kimi_cls.__module__.rsplit(".", maxsplit=1) - deepseek_module = sys.modules.get(f"{module_prefix}.modeling_deepseek") - if deepseek_module is None or not hasattr(deepseek_module, "DeepseekV3PreTrainedModel"): - return - - deepseek_cls = deepseek_module.DeepseekV3PreTrainedModel - if getattr(deepseek_cls, "_qeff_test_init_weights_patched", False): - return - - def _init_weights_compat(self, module): - std = self.config.initializer_range - if isinstance(module, torch.nn.Linear): - if hasattr(module, "weight") and module.weight is not None: - module.weight.data.normal_(mean=0.0, std=std) - if hasattr(module, "bias") and module.bias is not None: - module.bias.data.zero_() - elif isinstance(module, torch.nn.Embedding): - if hasattr(module, "weight") and module.weight is not None: - module.weight.data.normal_(mean=0.0, std=std) - if module.padding_idx is not None: - module.weight.data[module.padding_idx].zero_() - - deepseek_cls._init_weights = _init_weights_compat - deepseek_cls._qeff_test_init_weights_patched = True - - -def _patch_kimi_k25_remote_code_compat(config): - _ensure_torch_fx_import_compatibility() - kimi_cls = get_class_from_dynamic_module("modeling_kimi_k25.KimiK25ForConditionalGeneration", config._name_or_path) - _patch_kimi_k25_tie_weights_compat(kimi_cls) - _patch_kimi_k25_deepseek_init_weights_compat(kimi_cls) def _get_kimi_k25_test_config(model_name: str): - config = AutoConfig.from_pretrained(model_name, trust_remote_code=True) - config._attn_implementation = "eager" - config.torch_dtype = torch.float32 - config.dtype = torch.float32 - additional_params = model_config_dict[model_name]["additional_params"] - - for attr, value in additional_params["text_config"].items(): - setattr(config.text_config, attr, value) - config.text_config._attn_implementation = "eager" - config.text_config.torch_dtype = torch.float32 - config.text_config.dtype = torch.float32 - - for attr, value in additional_params["vision_config"].items(): - setattr(config.vision_config, attr, value) - config.vision_config._attn_implementation = "eager" - config.vision_config.torch_dtype = torch.float32 - config.vision_config.dtype = torch.float32 - - _patch_kimi_k25_remote_code_compat(config) - return config - - -def _get_kimi_k25_num_image_tokens(config, grid_thws): - merge_height, merge_width = config.vision_config.merge_kernel_size - return int(grid_thws[0, 1].item() // merge_height) * int(grid_thws[0, 2].item() // merge_width) - - -def _load_kimi_k25_layer_subset_model(): - examples_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../examples/kimi_k2")) - if examples_dir not in sys.path: - sys.path.insert(0, examples_dir) - - from test_kimi_k25 import ( # noqa: PLC0415 - LOADED_EXPERT_IDS, - NUM_EXPERTS_PER_TOKEN, - NUM_TEXT_LAYERS, - NUM_VISION_LAYERS, - _load_layer_subset_model, - _patch_deepseek_init_weights_compat, - _patch_kimi_tie_weights_compat, - _prepare_config, - _resolve_model_path, - _set_deterministic, - ) - - _set_deterministic(1234) - _ensure_torch_fx_import_compatibility() - model_path = _resolve_model_path() - config = _prepare_config(model_path) - kimi_cls = get_class_from_dynamic_module("modeling_kimi_k25.KimiK25ForConditionalGeneration", str(model_path)) - _patch_kimi_tie_weights_compat(kimi_cls) - _patch_deepseek_init_weights_compat(kimi_cls) - - model, tokenizer, processor = _load_layer_subset_model( - model_path=model_path, - kimi_cls=kimi_cls, - config=config, - num_vision_layers=NUM_VISION_LAYERS, - num_text_layers=NUM_TEXT_LAYERS, - loaded_expert_ids=LOADED_EXPERT_IDS, - num_experts_per_tok=NUM_EXPERTS_PER_TOKEN, - dtype=torch.float32, - ) - model.vision_tower.patch_embed.pos_emb.interpolation_mode = "bilinear" - model = model.eval().to("cpu") - return model, tokenizer, processor - - -@torch.no_grad() -def _run_kimi_k25_hf_model_on_pytorch(model, processor, inputs, max_gen_len): - generated_ids = inputs["input_ids"] - attention_mask = inputs["attention_mask"] - pixel_values = inputs["pixel_values"] - grid_thws = inputs["grid_thws"] - new_tokens = [] - - eos_token_id = getattr(model.config, "eos_token_id", None) - if eos_token_id is None and hasattr(model.config, "text_config"): - eos_token_id = getattr(model.config.text_config, "eos_token_id", None) - - for _ in range(max_gen_len): - outputs = model( - input_ids=generated_ids, - attention_mask=attention_mask, - pixel_values=pixel_values, - grid_thws=grid_thws, - use_cache=False, - return_dict=True, - ) - logits = outputs[0] if isinstance(outputs, tuple) else outputs.logits - next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True) - new_tokens.append(next_token) - - generated_ids = torch.cat([generated_ids, next_token], dim=1) - attention_mask = torch.cat( - [ - attention_mask, - torch.ones((attention_mask.shape[0], 1), dtype=attention_mask.dtype, device=attention_mask.device), - ], - dim=1, - ) - - if eos_token_id is not None and torch.all(next_token == eos_token_id): - break - - output_tokens = torch.cat(new_tokens, dim=1).squeeze(0) - py_output = processor.tokenizer.decode(output_tokens.tolist()).strip() - print("Original HF Model Outputs (Torch CPU):") - print("Completion:", repr(py_output)) - return output_tokens + return get_kimi_k25_test_config(model_name, model_config_dict) def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100( From 3976e31d8ebb7a4c89b6e8e36d5b3cd98f0446bb Mon Sep 17 00:00:00 2001 From: Mamta Singh <168400541+quic-mamta@users.noreply.github.com> Date: Fri, 17 Jul 2026 16:03:28 +0530 Subject: [PATCH 23/33] Update release_docs.md Signed-off-by: Mamta Singh <168400541+quic-mamta@users.noreply.github.com> --- docs/source/release_docs.md | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/docs/source/release_docs.md b/docs/source/release_docs.md index 71678f1173..df0bfe7fa8 100644 --- a/docs/source/release_docs.md +++ b/docs/source/release_docs.md @@ -29,8 +29,7 @@ Welcome to the official release of **Efficient Transformer Library v1.21.0**! Th - **Kimi-K2.5 (Vision Language)** - Executable via [`QEFFAutoModelForImageTextToText`](#QEFFAutoModelForImageTextToText) - Supports `moonshotai/Kimi-K2.5` with `KimiK25ForConditionalGeneration` - - Supports single-QPC and dual-QPC (`kv_offload=True`) vision-language compilation - - Uses Kimi-specific image-grid specializations: `num_patches`, `h`, `w`, and `num_image_tokens` + - Supports dual-QPC (`kv_offload=True`) vision-language compilation - [Kimi-K2.5 Vision Example Script](https://github.com/quic/efficient-transformers/blob/main/examples/kimi_k2/export_kimi_k25_vision.py) From 3fe235cd0cfaec16a8ec0ef31780163af9ffdc98 Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Sun, 19 Jul 2026 20:58:05 +0530 Subject: [PATCH 24/33] rename image_embeds to vision_embeds Signed-off-by: Mamta Singh --- QEfficient/generation/embedding_handler.py | 6 +- QEfficient/generation/vlm_generation.py | 1 - .../models/kimi_k25/modeling_kimi_k25.py | 141 +++++++++++++----- .../transformers/models/pytorch_transforms.py | 1 + examples/kimi_k2/export_kimi_k25_vision.py | 81 +++++----- examples/kimi_k2/test_kimi_k25.py | 24 ++- .../test_kimi_k25_disagg.py | 107 +++++++++++-- 7 files changed, 257 insertions(+), 104 deletions(-) rename {examples/kimi_k2 => tests/transformers/models/image_text_to_text}/test_kimi_k25_disagg.py (78%) diff --git a/QEfficient/generation/embedding_handler.py b/QEfficient/generation/embedding_handler.py index 15c553ab91..c8f3528094 100755 --- a/QEfficient/generation/embedding_handler.py +++ b/QEfficient/generation/embedding_handler.py @@ -451,11 +451,7 @@ def setup_vision_buffers(self): buffers = {} for output_name, shape in shapes.items(): # Create placeholder with appropriate dtype - if ( - "vision_embeds" in output_name - or "image_embeds" in output_name - or "deepstack_features" in output_name - ): + if "vision_embeds" in output_name or "deepstack_features" in output_name: buffers[output_name] = np.zeros(shape, dtype=np.float16) else: buffers[output_name] = np.zeros(shape, dtype=np.float32) diff --git a/QEfficient/generation/vlm_generation.py b/QEfficient/generation/vlm_generation.py index 013452fbc1..3d16b7f16c 100755 --- a/QEfficient/generation/vlm_generation.py +++ b/QEfficient/generation/vlm_generation.py @@ -737,7 +737,6 @@ def _generate_multi_frame_specialization( generation_len: int = None, stream: List[str] = None, ): - exec_batch_size = self.batch_size max_gen_length = self._ctx_len if not generation_len else max(self._ctx_len, generation_len) self.initialize_decode_inputs( diff --git a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py index a0e7cc8af5..1635b50912 100644 --- a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py +++ b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py @@ -373,7 +373,7 @@ def forward(self, pixel_values: torch.Tensor, h_shape: torch.Tensor, w_shape: to w_shape: int64 ones tensor of shape (w,) encoding number of patch columns. Returns: - image_embeds: Projected image embeddings as a single tensor. For multiple + vision_embeds: Projected image embeddings as a single tensor. For multiple same-sized images, embeddings are concatenated in input order. """ @@ -442,8 +442,8 @@ def forward(self, pixel_values: torch.Tensor, h_shape: torch.Tensor, w_shape: to pre_norm_dtype = self.model.mm_projector.pre_norm.weight.dtype merged = merged.to(pre_norm_dtype) - image_embeds = self.model.mm_projector.proj(self.model.mm_projector.pre_norm(merged).view(merged.shape[0], -1)) - return image_embeds + vision_embeds = self.model.mm_projector.proj(self.model.mm_projector.pre_norm(merged).view(merged.shape[0], -1)) + return vision_embeds class QEffKimiK25DecoderWrapper(nn.Module): @@ -467,7 +467,7 @@ def forward( self, input_ids: torch.LongTensor = None, inputs_embeds: Optional[torch.FloatTensor] = None, - image_embeds: Optional[torch.FloatTensor] = None, + vision_embeds: Optional[torch.FloatTensor] = None, image_idx: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, @@ -479,18 +479,20 @@ def forward( **kwargs, ) -> Tuple: inputs_embeds = self.model.get_input_embeddings()(input_ids) - image_embeds_for_state = None - if image_embeds is not None: - if image_embeds.dim() == 3: - if image_embeds.shape[0] != 1: - raise ValueError(f"Expected image_embeds batch dim to be 1, got shape {tuple(image_embeds.shape)}") - image_embeds_for_state = image_embeds[0].to(device=inputs_embeds.device, dtype=inputs_embeds.dtype) - elif image_embeds.dim() == 2: - image_embeds_for_state = image_embeds.to(device=inputs_embeds.device, dtype=inputs_embeds.dtype) + vision_embeds_for_state = None + if vision_embeds is not None: + if vision_embeds.dim() == 3: + if vision_embeds.shape[0] != 1: + raise ValueError( + f"Expected vision_embeds batch dim to be 1, got shape {tuple(vision_embeds.shape)}" + ) + vision_embeds_for_state = vision_embeds[0].to(device=inputs_embeds.device, dtype=inputs_embeds.dtype) + elif vision_embeds.dim() == 2: + vision_embeds_for_state = vision_embeds.to(device=inputs_embeds.device, dtype=inputs_embeds.dtype) else: - raise ValueError(f"Expected image_embeds rank 2 or 3, got {image_embeds.dim()}.") + raise ValueError(f"Expected vision_embeds rank 2 or 3, got {vision_embeds.dim()}.") - if image_embeds_for_state is not None and input_ids is not None: + if vision_embeds_for_state is not None and input_ids is not None: if attention_mask is None: if position_ids is None: attention_mask = torch.ones_like(input_ids, dtype=torch.long, device=input_ids.device) @@ -501,17 +503,11 @@ def forward( selected = input_ids == self.config.media_placeholder_token_id selected_any = selected.any(dim=1, keepdim=True) selected_image_tokens = selected.to(torch.int64).sum(dim=1, keepdim=True) - - num_selected_images = int(selected_image_tokens.max().item()) - if num_selected_images > 0 and image_embeds_for_state.shape[0] % num_selected_images == 0: - image_features = list(torch.chunk(image_embeds_for_state, num_selected_images, dim=0)) - else: - image_features = [image_embeds_for_state] - inputs_embeds = inputs_embeds.to(image_embeds_for_state.dtype) + inputs_embeds = inputs_embeds.to(vision_embeds_for_state.dtype) merged_inputs_embeds, merged_attention_mask, _, merged_position_ids = ( - self.model._merge_input_ids_with_image_features( - image_features=image_features, + self.model._qeff_merge_single_image_symbolic( + image_features=vision_embeds_for_state, inputs_embeds=inputs_embeds, input_ids=input_ids, attention_mask=attention_mask, @@ -536,10 +532,10 @@ def forward( torch.full_like(merged_position_ids, -1), ) - merged_image_tokens = torch.tensor( - [[sum(feature.shape[0] for feature in image_features)]], - dtype=torch.int64, - device=image_idx.device, + merged_image_tokens = ( + torch._shape_as_tensor(vision_embeds_for_state)[:1] + .view(1, 1) + .to(device=image_idx.device, dtype=torch.int64) ) image_position_delta = torch.clamp(merged_image_tokens - selected_image_tokens, min=0) image_idx = image_idx + selected_any.to(torch.int64) * image_position_delta @@ -572,7 +568,7 @@ def forward( output_kvs = getattr(outputs, "compressed_kvs", None) else: output_kvs = getattr(outputs, "past_key_values", None) - return logits, image_embeds_for_state, image_idx, output_kvs + return logits, vision_embeds_for_state, image_idx, output_kvs # ref https://github.com/huggingface/transformers/blob/78b2929c0554b79e0489b451ce4ece14d265ead2/src/transformers/models/llava/modeling_llava.py#L240 @@ -689,6 +685,77 @@ def _qeff_merge_input_ids_with_image_features( return final_embedding, final_attention_mask, labels, position_ids + def _qeff_merge_single_image_symbolic( + self, + image_features: torch.Tensor, + inputs_embeds: torch.Tensor, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + labels: torch.Tensor | None = None, + ): + image_token_index: int = self.config.media_placeholder_token_id + pad_token_id: int = self.config.pad_token_id + + target_device = inputs_embeds.device + image_features = image_features.to(target_device) + image_shape = torch._shape_as_tensor(image_features).to(device=input_ids.device, dtype=input_ids.dtype) + num_image_tokens = image_shape[0] + + image_token_mask = input_ids == image_token_index + non_image_mask = ~image_token_mask + image_token_count = image_token_mask.to(input_ids.dtype).sum(dim=1, keepdim=True) + has_image = image_token_count > 0 + + token_occupation = torch.where( + image_token_mask, + num_image_tokens.view(1, 1).expand_as(input_ids), + torch.ones_like(input_ids), + ) + new_token_positions = torch.cumsum(token_occupation, dim=1) - 1 + + final_seq_len = input_ids.shape[1] + image_features.shape[0] - 1 + merged_positions = torch.arange(final_seq_len, device=input_ids.device, dtype=input_ids.dtype).view(1, 1, -1) + text_position_one_hot = merged_positions == new_token_positions.unsqueeze(-1) + text_position_one_hot = torch.logical_and(text_position_one_hot, non_image_mask.unsqueeze(-1)) + + final_embedding = ( + text_position_one_hot.to(inputs_embeds.dtype).unsqueeze(-1) * inputs_embeds.unsqueeze(2) + ).sum(dim=1) + final_attention_mask = (text_position_one_hot.to(attention_mask.dtype) * attention_mask.unsqueeze(-1)).sum( + dim=1 + ) + + image_start_positions = torch.where( + image_token_mask, + new_token_positions - num_image_tokens.view(1, 1) + 1, + torch.zeros_like(new_token_positions), + ) + image_start = image_start_positions.max(dim=1, keepdim=True).values + image_positions = merged_positions.squeeze(1) - image_start + max_image_index = num_image_tokens.view(1, 1) - 1 + safe_image_positions = torch.minimum(torch.clamp(image_positions, min=0), max_image_index) + image_slots = torch.logical_and(image_positions >= 0, image_positions < num_image_tokens.view(1, 1)) + image_slots = torch.logical_and(image_slots, has_image) + image_slots = torch.logical_and(image_slots, torch.logical_not(text_position_one_hot.any(dim=1))) + + gathered_image_embeddings = image_features[safe_image_positions.to(torch.long)] + final_embedding = torch.where(image_slots.unsqueeze(-1), gathered_image_embeddings, final_embedding) + final_attention_mask = torch.logical_or(final_attention_mask.bool(), image_slots).to(final_attention_mask.dtype) + + position_ids = torch.cumsum(final_attention_mask, dim=1) - 1 + position_ids = torch.where(final_attention_mask == 0, torch.full_like(position_ids, -1), position_ids) + + pad_token_mask = input_ids == pad_token_id + pad_position_one_hot = torch.logical_and( + merged_positions == new_token_positions.unsqueeze(-1), pad_token_mask.unsqueeze(-1) + ) + pad_positions_mask = pad_position_one_hot.any(dim=1) + final_embedding = torch.where( + pad_positions_mask.unsqueeze(-1), torch.zeros_like(final_embedding), final_embedding + ) + + return final_embedding, final_attention_mask, labels, position_ids + def forward( self, input_ids: torch.LongTensor | None = None, @@ -815,7 +882,7 @@ def forward( ) def get_output_names(self, kv_offload: bool = False): - vision_output_names = ["image_embeds"] + vision_output_names = ["vision_embeds"] lang_output_names = ["logits"] mla_absorption = getattr(self.language_model, "mla_absorption", None) @@ -835,7 +902,7 @@ def get_output_names(self, kv_offload: bool = False): output_names = {} if kv_offload: - lang_output_names.insert(1, "image_embeds_RetainedState") + lang_output_names.insert(1, "vision_embeds_RetainedState") lang_output_names.insert(2, "image_idx_output") output_names["vision"] = vision_output_names output_names["lang"] = lang_output_names @@ -866,7 +933,7 @@ def get_dummy_inputs( constants.KIMI_PATCH_SIZE, constants.KIMI_PATCH_SIZE, ) - inputs_shapes["image_embeds"] = ( + inputs_shapes["vision_embeds"] = ( constants.KIMI_NUM_IMAGE_TOKENS, self.language_model.config.hidden_size, ) @@ -919,8 +986,8 @@ def get_dummy_inputs( torch.zeros(pkv_cache[0][1].shape, dtype=self.language_model.config.torch_dtype) ) - lang_inputs["image_embeds"] = torch.zeros( - inputs_shapes["image_embeds"], + lang_inputs["vision_embeds"] = torch.zeros( + inputs_shapes["vision_embeds"], dtype=self.language_model.config.torch_dtype, ) lang_inputs["image_idx"] = torch.zeros(inputs_shapes["image_idx"], dtype=torch.int64) @@ -933,7 +1000,7 @@ def get_dummy_inputs( inputs["vision"] = vision_inputs inputs["lang"] = lang_inputs else: - lang_inputs.pop("image_embeds") + lang_inputs.pop("vision_embeds") inputs = {**vision_inputs, **lang_inputs} return inputs @@ -945,14 +1012,14 @@ def get_onnx_dynamic_axes( lang_dynamic_axes = {} lang_dynamic_axes["input_ids"] = {0: "batch_size", 1: "seq_len"} lang_dynamic_axes["position_ids"] = {0: "batch_size", 1: "seq_len"} - lang_dynamic_axes["image_embeds"] = {0: "num_image_tokens"} + lang_dynamic_axes["vision_embeds"] = {0: "num_image_tokens"} if continuous_batching: lang_dynamic_axes["batch_index"] = {0: "batch_size"} vision_dynamic_axes = { "pixel_values": {0: "num_patches"}, "h_shape": {0: "h"}, "w_shape": {0: "w"}, - "image_embeds": {0: "num_image_tokens"}, + "vision_embeds": {0: "num_image_tokens"}, } mla_absorption = getattr(self.language_model, "mla_absorption", None) @@ -980,7 +1047,7 @@ def get_onnx_dynamic_axes( dynamic_axes["vision"] = vision_dynamic_axes dynamic_axes["lang"] = lang_dynamic_axes else: - lang_dynamic_axes.pop("image_embeds") + lang_dynamic_axes.pop("vision_embeds") dynamic_axes = {**vision_dynamic_axes, **lang_dynamic_axes} return dynamic_axes diff --git a/QEfficient/transformers/models/pytorch_transforms.py b/QEfficient/transformers/models/pytorch_transforms.py index e4ca5e49f2..41dfe41bd8 100755 --- a/QEfficient/transformers/models/pytorch_transforms.py +++ b/QEfficient/transformers/models/pytorch_transforms.py @@ -1242,6 +1242,7 @@ class KVCacheExternalModuleMapperTransform(ExternalModuleMapperTransform): }, "KimiK25ForConditionalGeneration": { "_merge_input_ids_with_image_features": QEffKimiK25ForConditionalGeneration._qeff_merge_input_ids_with_image_features, + "_qeff_merge_single_image_symbolic": QEffKimiK25ForConditionalGeneration._qeff_merge_single_image_symbolic, "get_qeff_vision_encoder": QEffKimiK25ForConditionalGeneration.get_qeff_vision_encoder, "get_qeff_language_decoder": QEffKimiK25ForConditionalGeneration.get_qeff_language_decoder, "get_specializations": QEffKimiK25ForConditionalGeneration.get_specializations, diff --git a/examples/kimi_k2/export_kimi_k25_vision.py b/examples/kimi_k2/export_kimi_k25_vision.py index 96f9d086ed..59a3204417 100644 --- a/examples/kimi_k2/export_kimi_k25_vision.py +++ b/examples/kimi_k2/export_kimi_k25_vision.py @@ -109,22 +109,6 @@ def main(): if skip_vision: ## TEXT-ONLY MODE ## - ## STEP 3: Compile Model for Text-Only Execution - # Set skip_vision=True to bypass image processing - qeff_model.compile( - qaic_config=qaic_config, - prefill_seq_len=1, - ctx_len=1024, - num_cores=16, - num_devices=1, - mxfp6_matmul=False, - mxint8_kv_cache=False, - aic_enable_depth_first=False, - skip_vision=True, # Skip vision encoder for text-only inference - mos=1, - ) - ## STEP 4: Prepare Text-Only Input - # Create a text-only message without any image messages = [ { "role": "user", @@ -134,7 +118,6 @@ def main(): }, ] - ## STEP 5: Process Input with Chat Template inputs = processor.apply_chat_template( messages, add_generation_prompt=True, @@ -143,19 +126,6 @@ def main(): return_tensors="pt", ) - ## STEP 6: Run Text-Only Inference - output = qeff_model.generate(inputs=inputs, device_ids=[0], generation_len=10) - - ## STEP 7: Display Results - print(output.generated_ids) - print(tokenizer.batch_decode(output.generated_ids)) - print(output) - - else: - ## VISION + TEXT MODE ## - - ## STEP 3: Compile Model for Vision+Text Execution - # Do not set skip_vision (defaults to False) to enable image processing qeff_model.compile( qaic_config=qaic_config, prefill_seq_len=1, @@ -165,19 +135,21 @@ def main(): mxfp6_matmul=False, mxint8_kv_cache=False, aic_enable_depth_first=False, + skip_vision=True, # Skip vision encoder for text-only inference mos=1, - num_patches=2400, - h=30, - w=80, - # Keep language-side image embedding specialization tightly bounded to - # actual single-image token count to avoid oversized dynamic VA mapping. - num_image_tokens=600, ) - ## STEP 4: Prepare Image and Text Input + output = qeff_model.generate(inputs=inputs, device_ids=[0], generation_len=10) + + print(output.generated_ids) + print(tokenizer.batch_decode(output.generated_ids)) + print(output) + + else: + ## VISION + TEXT MODE ## + image = Image.open(BytesIO(requests.get(args.image_url).content)).convert("RGB") - # Create a message with both image and text messages = [ { "role": "user", @@ -188,20 +160,45 @@ def main(): }, ] - ## STEP 5: Process Input inputs = processor( messages=messages, add_generation_prompt=True, tokenize=False, return_tensors="pt", ) - # Convert pixel values to float32 for processing + inputs["pixel_values"] = inputs["pixel_values"].to(qeff_model.model.config.torch_dtype) - ## STEP 6: Run Vision+Text Inference + merge_kernel_size = getattr(model.config.vision_config, "merge_kernel_size", (2, 2)) + if isinstance(merge_kernel_size, int): + kernel_height = kernel_width = merge_kernel_size + merge_kernel_size = (merge_kernel_size, merge_kernel_size) + else: + kernel_height, kernel_width = merge_kernel_size + + num_patches = int(inputs["pixel_values"].shape[0]) + h = int(inputs["grid_thws"][0, 1].item()) + w = int(inputs["grid_thws"][0, 2].item()) + num_image_tokens = int(inputs["pixel_values"].shape[0] // (kernel_height * kernel_width)) + + qeff_model.compile( + qaic_config=qaic_config, + prefill_seq_len=1, + ctx_len=1024, + num_cores=16, + num_devices=1, + mxfp6_matmul=False, + mxint8_kv_cache=False, + aic_enable_depth_first=False, + mos=1, + num_patches=num_patches, + h=h, + w=w, + num_image_tokens=num_image_tokens, + ) + output = qeff_model.generate(inputs=inputs, generation_len=10) - ## STEP 7: Display Results print(output.generated_ids) print(tokenizer.batch_decode(output.generated_ids)) print(output) diff --git a/examples/kimi_k2/test_kimi_k25.py b/examples/kimi_k2/test_kimi_k25.py index 0a13072779..9235c32794 100644 --- a/examples/kimi_k2/test_kimi_k25.py +++ b/examples/kimi_k2/test_kimi_k25.py @@ -131,7 +131,7 @@ def _greedy_generate_qeff_wrapper(transformed_model, inputs, max_new_tokens: int grid_thws = inputs["grid_thws"].to(torch.long) h_shape = torch.ones(int(grid_thws[0, 1].item()), dtype=torch.int64) w_shape = torch.ones(int(grid_thws[0, 2].item()), dtype=torch.int64) - image_embeds = qeff_encoder(inputs["pixel_values"].to(torch.float32), h_shape, w_shape).detach() + vision_embeds = qeff_encoder(inputs["pixel_values"].to(torch.float32), h_shape, w_shape).detach() generated_ids = inputs["input_ids"].to(torch.long) attention_mask = inputs["attention_mask"].to(torch.long) @@ -147,7 +147,7 @@ def _greedy_generate_qeff_wrapper(transformed_model, inputs, max_new_tokens: int input_ids=generated_ids, attention_mask=attention_mask, position_ids=position_ids, - image_embeds=image_embeds, + vision_embeds=vision_embeds, image_idx=torch.zeros((generated_ids.shape[0], 1), dtype=torch.int64), past_key_values=past_key_values, use_cache=True, @@ -201,16 +201,28 @@ def check_kimi_k25_pytorch_vs_ai100(): ) print("QEFF:", _decode_tokens(tokenizer, qeff_tokens), "\n", qeff_tokens) + merge_kernel_size = getattr(model.config.vision_config, "merge_kernel_size", (2, 2)) + if isinstance(merge_kernel_size, int): + kernel_height = kernel_width = merge_kernel_size + merge_kernel_size = (merge_kernel_size, merge_kernel_size) + else: + kernel_height, kernel_width = merge_kernel_size + + num_patches = int(inputs["pixel_values"].shape[0]) + h = int(inputs["grid_thws"][0, 1].item()) + w = int(inputs["grid_thws"][0, 2].item()) + num_image_tokens = int(inputs["pixel_values"].shape[0] // (kernel_height * kernel_width)) + qeff_model.compile( # qaic_config=qaic_config, num_devices=1, prefill_seq_len=1, ctx_len=CTX_LEN, mxfp6=False, - num_patches=int(inputs["pixel_values"].shape[0]), - h=int(inputs["grid_thws"][0, 1].item()), - w=int(inputs["grid_thws"][0, 2].item()), - num_image_tokens=600, + num_patches=num_patches, + h=h, + w=w, + num_image_tokens=num_image_tokens, num_cores=16, ) diff --git a/examples/kimi_k2/test_kimi_k25_disagg.py b/tests/transformers/models/image_text_to_text/test_kimi_k25_disagg.py similarity index 78% rename from examples/kimi_k2/test_kimi_k25_disagg.py rename to tests/transformers/models/image_text_to_text/test_kimi_k25_disagg.py index c83ee01b63..b576104219 100644 --- a/examples/kimi_k2/test_kimi_k25_disagg.py +++ b/tests/transformers/models/image_text_to_text/test_kimi_k25_disagg.py @@ -6,20 +6,14 @@ # ---------------------------------------------------------------------------- import copy +from io import BytesIO from pathlib import Path import numpy as np import pytest +import requests import torch -from test_kimi_k25 import ( - _clone_inputs, - _decode_tokens, - _greedy_generate_hf, - _has_qaic_runtime_access, - _prepare_inputs, - _resolve_model_path, - _set_deterministic, -) +from PIL import Image from QEfficient import QEFFAutoModelForImageTextToText from QEfficient.generation.cloud_infer import QAICInferenceSession @@ -37,6 +31,12 @@ from QEfficient.utils.load_kimi_utils import ( prepare_config as _prepare_config, ) +from QEfficient.utils.load_kimi_utils import ( + resolve_model_path as _resolve_model_path, +) +from QEfficient.utils.load_kimi_utils import ( + set_deterministic as _set_deterministic, +) PREFILL_SEQ_LEN = 512 CTX_LEN = 2048 @@ -44,6 +44,87 @@ GENERATION_LEN = 10 NUM_VISION_LAYERS = 2 NUM_TEXT_LAYERS = 2 +IMAGE_URL = "https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png" +TEXT_PROMPT = "Describe this image." + + +def _has_qaic_runtime_access() -> bool: + try: + _ = QAICInferenceSession + except Exception: + return False + try: + import qaicrt + + _ctx = qaicrt.Context() + return True + except Exception: + return False + + +def _prepare_inputs(processor): + image = Image.open(BytesIO(requests.get(IMAGE_URL, timeout=30).content)).convert("RGB") + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": image}, + {"type": "text", "text": TEXT_PROMPT}, + ], + } + ] + return processor( + messages=messages, + add_generation_prompt=True, + tokenize=False, + return_tensors="pt", + ) + + +def _decode_tokens(tokenizer, token_ids: torch.Tensor) -> str: + decoded = tokenizer.batch_decode(token_ids, skip_special_tokens=True) + return decoded[0] if decoded else "" + + +def _clone_inputs(inputs): + return {k: (v.clone() if torch.is_tensor(v) else copy.deepcopy(v)) for k, v in inputs.items()} + + +def _greedy_generate_hf(model, inputs, max_new_tokens: int): + generated_ids = inputs["input_ids"] + attention_mask = inputs["attention_mask"] + pixel_values = inputs["pixel_values"] + grid_thws = inputs["grid_thws"] + + eos_token_id = getattr(model.config, "eos_token_id", None) + if eos_token_id is None and hasattr(model.config, "text_config"): + eos_token_id = getattr(model.config.text_config, "eos_token_id", None) + + for _ in range(max_new_tokens): + outputs = model( + input_ids=generated_ids, + attention_mask=attention_mask, + pixel_values=pixel_values, + grid_thws=grid_thws, + use_cache=False, + return_dict=True, + ) + logits = outputs[0] if isinstance(outputs, tuple) else outputs.logits + next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True) + + generated_ids = torch.cat([generated_ids, next_token], dim=1) + attention_mask = torch.cat( + [ + attention_mask, + torch.ones((attention_mask.shape[0], 1), dtype=attention_mask.dtype, device=attention_mask.device), + ], + dim=1, + ) + + if eos_token_id is not None and torch.all(next_token == eos_token_id): + break + + return generated_ids[:, -max_new_tokens:] def _assert_onnx_path(onnx_path, label: str) -> Path: @@ -210,13 +291,13 @@ def _run_disagg_qaic_generation( vision_outputs = vision_session.run(_filter_session_inputs(vision_session, vision_inputs)) vision_session.deactivate() - image_embeds = vision_outputs.get("image_embeds") - assert image_embeds is not None, f"Vision QPC did not return image_embeds. Outputs: {vision_outputs.keys()}" + vision_embeds = vision_outputs.get("vision_embeds") + assert vision_embeds is not None, f"Vision QPC did not return vision_embeds. Outputs: {vision_outputs.keys()}" lang_inputs = { "input_ids": inputs["input_ids"].astype(np.int64), "position_ids": np.where(inputs["attention_mask"] > 0, np.arange(padded_len), -1).astype(np.int64), - "image_embeds": image_embeds, + "vision_embeds": vision_embeds, "image_idx": np.zeros((BATCH_SIZE, 1), dtype=np.int64), } @@ -240,7 +321,7 @@ def _run_disagg_qaic_generation( decode_inputs = { "input_ids": generated_ids[-1], "position_ids": np.max(lang_inputs["position_ids"], axis=-1, keepdims=True).astype(np.int64) + 1, - "image_embeds": chunk_inputs.get("image_embeds", image_embeds), + "vision_embeds": chunk_inputs.get("vision_embeds", vision_embeds), "image_idx": chunk_inputs.get("image_idx", np.zeros((BATCH_SIZE, 1), dtype=np.int64)), } _update_retained_states(decode_inputs, prefill_outputs) From 27af3627915d62261f1ccf16a04bb5930b002de3 Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Mon, 20 Jul 2026 13:51:29 +0530 Subject: [PATCH 25/33] subfunction fix Signed-off-by: Mamta Singh --- QEfficient/base/modeling_qeff.py | 20 +- .../models/kimi_k25/modeling_kimi_k25.py | 21 ++ QEfficient/utils/constants.py | 1 + QEfficient/utils/export_utils.py | 11 +- examples/kimi_k2/test_kimi_k25.py | 256 ------------------ 5 files changed, 48 insertions(+), 261 deletions(-) delete mode 100644 examples/kimi_k2/test_kimi_k25.py diff --git a/QEfficient/base/modeling_qeff.py b/QEfficient/base/modeling_qeff.py index 2430042a2a..d994148ac7 100755 --- a/QEfficient/base/modeling_qeff.py +++ b/QEfficient/base/modeling_qeff.py @@ -97,6 +97,19 @@ def _restore_output_names_exact(model: onnx.ModelProto, output_names: List[str]) _rename_graph_value(model.graph, current_name, expected_name) +def _has_exported_model_subfunctions(onnx_path: Union[str, Path]) -> bool: + """Return True when the ONNX has non-custom local functions for compiler subfunction mode.""" + try: + model = onnx.load(onnx_path, load_external_data=False) + except Exception: + return True + + custom_function_names = { + onnxscript_func.to_function_proto().name for _, onnxscript_func in CustomOpTransform._custom_ops.values() + } + return any(function.name not in custom_function_names for function in model.functions) + + class QEFFBaseModel(ABC): """ Base class for all the model classes (i.e. LLMs, SD, quantized etc.). @@ -941,6 +954,11 @@ def _compile( return self.qpc_path + use_compiler_subfunctions = use_onnx_subfunctions and _has_exported_model_subfunctions(onnx_path) + + if use_compiler_subfunctions and prefill_only is True: + compiler_options.setdefault("elf_va_limit", constants.DEFAULT_ONNX_SUBFUNCTION_PREFILL_ELF_VA_LIMIT_MB) + command = ( constants.COMPILER + [ @@ -1026,7 +1044,7 @@ def _compile( except Exception: pass - if use_onnx_subfunctions: + if use_compiler_subfunctions: logger.info("Using ONNX subfunctions for compilation.") command.append("-sub-functions") diff --git a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py index 1635b50912..f2d1a7fb8c 100644 --- a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py +++ b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py @@ -461,8 +461,29 @@ def get_submodules_for_export(self) -> Type[nn.Module]: This method should return the *class object* (not an instance). Downstream code can use this to find/build subfunctions for repeated blocks. """ + if self._has_uint4_weight_as_activation_moe(): + return set() return {self.language_model.model.layers[0].__class__} + def _has_uint4_weight_as_activation_moe(self) -> bool: + for layer in self.language_model.model.layers: + mlp = getattr(layer, "mlp", None) + if mlp is None: + continue + if all( + hasattr(mlp, attr) + for attr in ( + "all_gate_qweight", + "all_gate_qzeros", + "all_up_qweight", + "all_up_qzeros", + "all_down_qweight", + "all_down_qzeros", + ) + ): + return True + return False + def forward( self, input_ids: torch.LongTensor = None, diff --git a/QEfficient/utils/constants.py b/QEfficient/utils/constants.py index bf3e8f9676..68ecddd13f 100644 --- a/QEfficient/utils/constants.py +++ b/QEfficient/utils/constants.py @@ -42,6 +42,7 @@ # Compiler defaults DEFAULT_AIC_NUM_CORES = 16 DEFAULT_AIC_MXPF6_MATMUL = False +DEFAULT_ONNX_SUBFUNCTION_PREFILL_ELF_VA_LIMIT_MB = 2048 # Hashing defaults HASH_HEXDIGEST_STR_LEN = 16 KWARGS_INCLUSION_LIST = [ diff --git a/QEfficient/utils/export_utils.py b/QEfficient/utils/export_utils.py index 9c3e1abdda..f219f62de3 100755 --- a/QEfficient/utils/export_utils.py +++ b/QEfficient/utils/export_utils.py @@ -177,8 +177,12 @@ def _setup_onnx_subfunctions(qeff_model, args, kwargs): apply_torch_patches() InvalidIndexProvider.SUBFUNC_ENABLED = True - # Transform output names for subfunction compatibility - if "output_names" in kwargs: + submodule_classes = qeff_model.model.get_submodules_for_export() + + # Transform output names for decoder-layer subfunction compatibility. If a model opts out of + # exporting decoder layers as ONNX functions, keep the public retained-state names because there + # are no function outputs for RenameFunctionOutputsTransform to map back. + if "output_names" in kwargs and submodule_classes: kwargs["output_names"] = [ re.sub("_RetainedState", "_InternalRetainedState", name) if name.endswith("_RetainedState") @@ -193,7 +197,7 @@ def _setup_onnx_subfunctions(qeff_model, args, kwargs): else name for name in kwargs["output_names"] ] - else: + elif "output_names" not in kwargs: warnings.warn( "ONNX subfunctions are enabled, but no retained-state output names were found to rewrite. " "Ensure `output_names` includes key/value retained states if subfunction compatibility is required." @@ -208,7 +212,6 @@ def _setup_onnx_subfunctions(qeff_model, args, kwargs): if CustomOpTransform not in qeff_model._onnx_transforms: qeff_model._onnx_transforms.append(CustomOpTransform) - submodule_classes = qeff_model.model.get_submodules_for_export() if submodule_classes: kwargs["export_modules_as_functions"] = submodule_classes return args, kwargs, {"onnx_transforms": original_transforms} diff --git a/examples/kimi_k2/test_kimi_k25.py b/examples/kimi_k2/test_kimi_k25.py deleted file mode 100644 index 9235c32794..0000000000 --- a/examples/kimi_k2/test_kimi_k25.py +++ /dev/null @@ -1,256 +0,0 @@ -# ----------------------------------------------------------------------------- -# -# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. -# SPDX-License-Identifier: BSD-3-Clause -# -# ----------------------------------------------------------------------------- - -import copy -from io import BytesIO - -import requests -import torch -from PIL import Image - -from QEfficient import QEFFAutoModelForImageTextToText -from QEfficient.generation.cloud_infer import QAICInferenceSession -from QEfficient.utils.load_kimi_utils import ( - KIMI_K25_MODEL_NAME, - LOADED_EXPERT_IDS, - NUM_EXPERTS_PER_TOKEN, - load_kimi_k25_class, -) -from QEfficient.utils.load_kimi_utils import ( - ensure_torch_fx_import_compatibility as _ensure_torch_fx_import_compatibility, -) -from QEfficient.utils.load_kimi_utils import ( - load_layer_subset_model as _load_layer_subset_model, -) -from QEfficient.utils.load_kimi_utils import ( - prepare_config as _prepare_config, -) -from QEfficient.utils.load_kimi_utils import ( - resolve_model_path as _resolve_model_path, -) -from QEfficient.utils.load_kimi_utils import ( - set_deterministic as _set_deterministic, -) - -MODEL_NAME = KIMI_K25_MODEL_NAME -IMAGE_URL = "https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png" -TEXT_PROMPT = "Describe this image." -NUM_VISION_LAYERS = 2 -NUM_TEXT_LAYERS = 2 -NEW_GENERATION_TOKENS = 10 -CTX_LEN = 1024 - - -def _has_qaic_runtime_access() -> bool: - try: - _ = QAICInferenceSession - except Exception: - return False - try: - import qaicrt - - _ctx = qaicrt.Context() - return True - except Exception: - return False - - -def _prepare_inputs(processor): - image = Image.open(BytesIO(requests.get(IMAGE_URL, timeout=30).content)).convert("RGB") - messages = [ - { - "role": "user", - "content": [ - {"type": "image_url", "image_url": image}, - {"type": "text", "text": TEXT_PROMPT}, - ], - } - ] - return processor( - messages=messages, - add_generation_prompt=True, - tokenize=False, - return_tensors="pt", - ) - - -def _decode_tokens(tokenizer, token_ids: torch.Tensor) -> str: - decoded = tokenizer.batch_decode(token_ids, skip_special_tokens=True) - return decoded[0] if decoded else "" - - -def _clone_inputs(inputs): - return {k: (v.clone() if torch.is_tensor(v) else copy.deepcopy(v)) for k, v in inputs.items()} - - -def _greedy_generate_hf(model, inputs, max_new_tokens: int): - generated_ids = inputs["input_ids"] - attention_mask = inputs["attention_mask"] - pixel_values = inputs["pixel_values"] - grid_thws = inputs["grid_thws"] - - eos_token_id = getattr(model.config, "eos_token_id", None) - if eos_token_id is None and hasattr(model.config, "text_config"): - eos_token_id = getattr(model.config.text_config, "eos_token_id", None) - - for _ in range(max_new_tokens): - outputs = model( - input_ids=generated_ids, - attention_mask=attention_mask, - pixel_values=pixel_values, - grid_thws=grid_thws, - use_cache=False, - return_dict=True, - ) - logits = outputs[0] if isinstance(outputs, tuple) else outputs.logits - next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True) - - generated_ids = torch.cat([generated_ids, next_token], dim=1) - attention_mask = torch.cat( - [ - attention_mask, - torch.ones((attention_mask.shape[0], 1), dtype=attention_mask.dtype, device=attention_mask.device), - ], - dim=1, - ) - - if eos_token_id is not None and torch.all(next_token == eos_token_id): - break - - return generated_ids[:, -max_new_tokens:] - - -def _greedy_generate_qeff_wrapper(transformed_model, inputs, max_new_tokens: int): - qeff_encoder = transformed_model.get_qeff_vision_encoder().eval() - decoder_wrapper = transformed_model.get_qeff_language_decoder().eval() - - grid_thws = inputs["grid_thws"].to(torch.long) - h_shape = torch.ones(int(grid_thws[0, 1].item()), dtype=torch.int64) - w_shape = torch.ones(int(grid_thws[0, 2].item()), dtype=torch.int64) - vision_embeds = qeff_encoder(inputs["pixel_values"].to(torch.float32), h_shape, w_shape).detach() - - generated_ids = inputs["input_ids"].to(torch.long) - attention_mask = inputs["attention_mask"].to(torch.long) - - for _ in range(max_new_tokens): - position_ids = torch.where(attention_mask.bool(), attention_mask.cumsum(-1) - 1, -1) - past_key_values = transformed_model.language_model.get_dummy_pkv_cache( - transformed_model.config.text_config, - generated_ids.shape[0], - CTX_LEN, - ) - logits = decoder_wrapper( - input_ids=generated_ids, - attention_mask=attention_mask, - position_ids=position_ids, - vision_embeds=vision_embeds, - image_idx=torch.zeros((generated_ids.shape[0], 1), dtype=torch.int64), - past_key_values=past_key_values, - use_cache=True, - )[0] - next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True) - - generated_ids = torch.cat([generated_ids, next_token], dim=1) - attention_mask = torch.cat( - [attention_mask, torch.ones((attention_mask.shape[0], 1), dtype=attention_mask.dtype)], - dim=1, - ) - - return generated_ids[:, -max_new_tokens:] - - -def check_kimi_k25_pytorch_vs_ai100(): - _set_deterministic(1234) - _ensure_torch_fx_import_compatibility() - model_path = _resolve_model_path() - config = _prepare_config(model_path) - kimi_cls = load_kimi_k25_class(model_path) - - model, tokenizer, processor = _load_layer_subset_model( - model_path=model_path, - kimi_cls=kimi_cls, - config=config, - num_vision_layers=NUM_VISION_LAYERS, - num_text_layers=NUM_TEXT_LAYERS, - loaded_expert_ids=LOADED_EXPERT_IDS, - num_experts_per_tok=NUM_EXPERTS_PER_TOKEN, - dtype=torch.float32, - ) - - model.vision_tower.patch_embed.pos_emb.interpolation_mode = "bilinear" - model = model.eval().to("cpu") - - inputs = _prepare_inputs(processor) - inputs = {k: (v.to("cpu") if torch.is_tensor(v) else v) for k, v in inputs.items()} - - hf_tokens = _greedy_generate_hf(copy.deepcopy(model), _clone_inputs(inputs), max_new_tokens=NEW_GENERATION_TOKENS) - print("HF:", _decode_tokens(tokenizer, hf_tokens), "\n", hf_tokens) - - # qaic_config = {"mla_absorption": {"cache_compressed": True, "absorption": False, "online": False}} - - qeff_model = QEFFAutoModelForImageTextToText(model) # , qaic_config=qaic_config) - - qeff_tokens = _greedy_generate_qeff_wrapper( - transformed_model=qeff_model.model, - inputs=_clone_inputs(inputs), - max_new_tokens=NEW_GENERATION_TOKENS, - ) - print("QEFF:", _decode_tokens(tokenizer, qeff_tokens), "\n", qeff_tokens) - - merge_kernel_size = getattr(model.config.vision_config, "merge_kernel_size", (2, 2)) - if isinstance(merge_kernel_size, int): - kernel_height = kernel_width = merge_kernel_size - merge_kernel_size = (merge_kernel_size, merge_kernel_size) - else: - kernel_height, kernel_width = merge_kernel_size - - num_patches = int(inputs["pixel_values"].shape[0]) - h = int(inputs["grid_thws"][0, 1].item()) - w = int(inputs["grid_thws"][0, 2].item()) - num_image_tokens = int(inputs["pixel_values"].shape[0] // (kernel_height * kernel_width)) - - qeff_model.compile( - # qaic_config=qaic_config, - num_devices=1, - prefill_seq_len=1, - ctx_len=CTX_LEN, - mxfp6=False, - num_patches=num_patches, - h=h, - w=w, - num_image_tokens=num_image_tokens, - num_cores=16, - ) - - qaic_tokens = None - if _has_qaic_runtime_access(): - inputs["pixel_values"] = inputs["pixel_values"].to(qeff_model.model.config.torch_dtype) - qaic_tokens = qeff_model.generate( - inputs=_clone_inputs(inputs), generation_len=NEW_GENERATION_TOKENS - ).generated_ids[:, :-1] - print("QAIC:", _decode_tokens(tokenizer, qaic_tokens) if qaic_tokens is not None else "") - else: - print("QAIC generation skipped: no QAIC runtime access (user/group permissions).") - - print( - "Loaded subset:", - f"vision_layers={model.config.vision_config.vt_num_hidden_layers}", - f"text_layers={model.config.text_config.num_hidden_layers}", - ) - - print("Prompt:", repr(TEXT_PROMPT)) - print("HF:", _decode_tokens(tokenizer, hf_tokens)) - print("QEFF:", _decode_tokens(tokenizer, qeff_tokens)) - print("QAIC:", _decode_tokens(tokenizer, qaic_tokens) if qaic_tokens is not None else "") - - assert torch.equal(hf_tokens, qeff_tokens), "HF and QEFF(Pytorch runtime) tokens do not match" - if qaic_tokens is not None: - assert torch.equal(hf_tokens, torch.as_tensor(qaic_tokens)), "HF and QAIC tokens do not match" - - -def test_kimi_k25_image_text_to_text_pytorch_vs_ai100(): - check_kimi_k25_pytorch_vs_ai100() From f0bcc06aaff9757c7efd0ea73c6a6964dbe08094 Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Mon, 20 Jul 2026 15:47:23 +0530 Subject: [PATCH 26/33] cleanup tests Signed-off-by: Mamta Singh --- .../models/kimi_k25/modeling_kimi_k25.py | 128 ++---------------- .../transformers/models/modeling_auto.py | 4 +- .../transformers/models/pytorch_transforms.py | 3 +- QEfficient/utils/constants.py | 7 +- QEfficient/utils/load_kimi_utils.py | 8 +- examples/kimi_k2/export_kimi_k25_vision.py | 40 ++---- .../test_continuous_batching.py | 56 ++------ .../test_image_text_to_text_models.py | 53 ++------ 8 files changed, 60 insertions(+), 239 deletions(-) diff --git a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py index f2d1a7fb8c..6e0af3c2c1 100644 --- a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py +++ b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py @@ -527,7 +527,7 @@ def forward( inputs_embeds = inputs_embeds.to(vision_embeds_for_state.dtype) merged_inputs_embeds, merged_attention_mask, _, merged_position_ids = ( - self.model._qeff_merge_single_image_symbolic( + self.model._qeff_merge_input_ids_with_image_features( image_features=vision_embeds_for_state, inputs_embeds=inputs_embeds, input_ids=input_ids, @@ -601,112 +601,6 @@ def get_qeff_language_decoder(self): return QEffKimiK25DecoderWrapper(self) def _qeff_merge_input_ids_with_image_features( - self, - image_features: list[torch.Tensor], - inputs_embeds: torch.Tensor, - input_ids: torch.Tensor, - attention_mask: torch.Tensor, - labels: torch.Tensor | None = None, - ): - """ - Args: - image_features (:obj:`torch.Tensor` of shape :obj:`(num_image_tokens, embed_dim)`): - The image features to merge with the input embeddings. - inputs_embeds (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length, embed_dim)`): - The input embeddings. - input_ids (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`): - The input ids. - attention_mask (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`): - The attention mask. - labels (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`, *optional*): - The labels. - """ - if len(image_features) == 0: - raise ValueError("At least one image_features tensor is required.") - - _, embed_dim = image_features[0].shape - feature_lengths = [x.shape[0] for x in image_features] - image_features_cat = torch.cat(image_features, dim=0) - - image_token_index: int = self.config.media_placeholder_token_id - pad_token_id: int = self.config.pad_token_id - - batch_size, _ = input_ids.shape - target_device = inputs_embeds.device - - left_padding = (~(input_ids[:, -1] == pad_token_id)).sum() == 0 - - image_token_mask = input_ids == image_token_index - non_image_mask = ~image_token_mask - - flat_image_token_mask = image_token_mask.reshape(-1).to(torch.long) - flat_image_order = torch.cumsum(flat_image_token_mask, dim=0) - feature_lengths_tensor = torch.tensor(feature_lengths, dtype=input_ids.dtype, device=input_ids.device) - flat_image_lengths = torch.zeros_like(flat_image_order, dtype=input_ids.dtype) - for idx in range(len(feature_lengths)): - flat_image_lengths = ( - flat_image_lengths + (flat_image_order == (idx + 1)).to(input_ids.dtype) * (feature_lengths_tensor[idx]) - ) - - token_occupation_table = torch.ones_like(input_ids) - token_occupation_table = token_occupation_table.reshape(-1) - token_occupation_table = torch.where(flat_image_token_mask.bool(), flat_image_lengths, token_occupation_table) - token_occupation_table = token_occupation_table.reshape(input_ids.shape) - - max_embed_dim = int(token_occupation_table.sum(-1).max().item()) - - new_token_positions = torch.cumsum(token_occupation_table, 1) - 1 - nb_image_pad = max_embed_dim - 1 - new_token_positions[:, -1] - if bool(left_padding): - new_token_positions = new_token_positions + nb_image_pad[:, None] - - merged_positions = torch.arange(max_embed_dim, device=input_ids.device).view(1, 1, -1) - text_position_one_hot = merged_positions == new_token_positions.unsqueeze(-1) - text_position_one_hot = torch.logical_and(text_position_one_hot, non_image_mask.unsqueeze(-1)) - - final_embedding = ( - text_position_one_hot.to(inputs_embeds.dtype).unsqueeze(-1) * inputs_embeds.unsqueeze(2) - ).sum(dim=1) - final_attention_mask = (text_position_one_hot.to(attention_mask.dtype) * attention_mask.unsqueeze(-1)).sum( - dim=1 - ) - - image_to_overwrite = torch.logical_not(text_position_one_hot.any(dim=1)) - image_to_overwrite_cumsum = torch.cumsum(image_to_overwrite.to(torch.int64), dim=1) - image_to_overwrite = torch.logical_and( - image_to_overwrite, - image_to_overwrite_cumsum - 1 >= nb_image_pad[:, None].to(target_device), - ) - - image_slot_ids = torch.cumsum(image_to_overwrite.to(torch.int64), dim=1) - 1 - if image_features_cat.shape[0] > 0: - valid_image_slots = torch.logical_and(image_to_overwrite, image_slot_ids < image_features_cat.shape[0]) - image_slot_ids = torch.clamp(image_slot_ids, min=0, max=image_features_cat.shape[0] - 1) - gathered_image_embeddings = image_features_cat.to(target_device)[image_slot_ids] - final_embedding = torch.where(valid_image_slots.unsqueeze(-1), gathered_image_embeddings, final_embedding) - image_to_overwrite = valid_image_slots - else: - image_to_overwrite = torch.zeros_like(image_to_overwrite) - - final_attention_mask = torch.logical_or(final_attention_mask.bool(), image_to_overwrite).to( - final_attention_mask.dtype - ) - - position_ids = torch.cumsum(final_attention_mask, dim=1) - 1 - position_ids = torch.where(final_attention_mask == 0, torch.ones_like(position_ids), position_ids) - - pad_token_mask = input_ids == pad_token_id - pad_position_one_hot = torch.logical_and( - merged_positions == new_token_positions.unsqueeze(-1), pad_token_mask.unsqueeze(-1) - ) - pad_positions_mask = pad_position_one_hot.any(dim=1) - final_embedding = torch.where( - pad_positions_mask.unsqueeze(-1), torch.zeros_like(final_embedding), final_embedding - ) - - return final_embedding, final_attention_mask, labels, position_ids - - def _qeff_merge_single_image_symbolic( self, image_features: torch.Tensor, inputs_embeds: torch.Tensor, @@ -949,18 +843,18 @@ def get_dummy_inputs( inputs_shapes = {} inputs_shapes["pixel_values"] = ( - constants.KIMI_NUM_PATCHES, + constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_HEIGHT * constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_WIDTH, constants.ONNX_EXPORT_IMAGE_DEPTH, constants.KIMI_PATCH_SIZE, constants.KIMI_PATCH_SIZE, ) inputs_shapes["vision_embeds"] = ( - constants.KIMI_NUM_IMAGE_TOKENS, + constants.KIMI_EXAMPLE_IMAGE_NUM_IMAGE_TOKENS, self.language_model.config.hidden_size, ) inputs_shapes["image_idx"] = (1, 1) - inputs_shapes["h"] = constants.KIMI_IMAGE_HEIGHT - inputs_shapes["w"] = constants.KIMI_IMAGE_WIDTH + inputs_shapes["h"] = constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_HEIGHT + inputs_shapes["w"] = constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_WIDTH vision_inputs = { "pixel_values": torch.zeros((inputs_shapes["pixel_values"]), dtype=self.config.torch_dtype), @@ -1131,12 +1025,12 @@ def validate_dimension_lists(heights, widths, height_name, width_name): kernel_height, kernel_width = merge_kernel_size if h is not None or w is not None: - heights = normalize_list(h, constants.KIMI_IMAGE_HEIGHT) - widths = normalize_list(w, constants.KIMI_IMAGE_WIDTH) + heights = normalize_list(h, constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_HEIGHT) + widths = normalize_list(w, constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_WIDTH) validate_dimension_lists(heights, widths, "h", "w") elif height is not None or width is not None: - pixel_heights = normalize_list(height, constants.KIMI_IMAGE_HEIGHT * patch_size) - pixel_widths = normalize_list(width, constants.KIMI_IMAGE_WIDTH * patch_size) + pixel_heights = normalize_list(height, constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_HEIGHT * patch_size) + pixel_widths = normalize_list(width, constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_WIDTH * patch_size) validate_dimension_lists(pixel_heights, pixel_widths, "height", "width") in_patch_limit = mm_processor_kwargs.get("in_patch_limit", 16384) @@ -1161,8 +1055,8 @@ def validate_dimension_lists(heights, widths, height_name, width_name): heights.append((resized_height + pad_height) // patch_size) widths.append((resized_width + pad_width) // patch_size) else: - heights = [constants.KIMI_IMAGE_HEIGHT] - widths = [constants.KIMI_IMAGE_WIDTH] + heights = [constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_HEIGHT] + widths = [constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_WIDTH] num_frames = normalize_sized_list(1 if num_frames is None else num_frames, len(heights), "num_frames") explicit_num_patches = normalize_sized_list(num_patches, len(heights), "num_patches") diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index fa4fb11ade..881a3762c7 100755 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -2097,7 +2097,9 @@ def compile( if prefill_only: specializations = specializations["lang"][:1] qpc_key = "lang_prefill_qpc_path" - elif prefill_seq_len == 1: + elif prefill_seq_len == 1 and not ( + self.continuous_batching and full_batch_size is not None and full_batch_size != batch_size + ): specializations = specializations["lang"][-1:] qpc_key = "lang_decode_qpc_path" else: diff --git a/QEfficient/transformers/models/pytorch_transforms.py b/QEfficient/transformers/models/pytorch_transforms.py index a8e7ecbfef..36c5cf8377 100755 --- a/QEfficient/transformers/models/pytorch_transforms.py +++ b/QEfficient/transformers/models/pytorch_transforms.py @@ -1322,8 +1322,7 @@ class KVCacheExternalModuleMapperTransform(ExternalModuleMapperTransform): "forward": QEFFGrok1CustomRMSNormAIC.forward, }, "KimiK25ForConditionalGeneration": { - "_merge_input_ids_with_image_features": QEffKimiK25ForConditionalGeneration._qeff_merge_input_ids_with_image_features, - "_qeff_merge_single_image_symbolic": QEffKimiK25ForConditionalGeneration._qeff_merge_single_image_symbolic, + "_qeff_merge_input_ids_with_image_features": QEffKimiK25ForConditionalGeneration._qeff_merge_input_ids_with_image_features, "get_qeff_vision_encoder": QEffKimiK25ForConditionalGeneration.get_qeff_vision_encoder, "get_qeff_language_decoder": QEffKimiK25ForConditionalGeneration.get_qeff_language_decoder, "get_specializations": QEffKimiK25ForConditionalGeneration.get_specializations, diff --git a/QEfficient/utils/constants.py b/QEfficient/utils/constants.py index 3479dc3ffb..1a1720abe4 100644 --- a/QEfficient/utils/constants.py +++ b/QEfficient/utils/constants.py @@ -180,11 +180,10 @@ def get_default_aic_hw_version() -> str: LLAMA4_MAX_POSITION_EMBEDDINGS = 65536 # DeepSeek Kimi-k2.5 Constants -KIMI_NUM_PATCHES = 2400 KIMI_PATCH_SIZE = 14 -KIMI_NUM_IMAGE_TOKENS = 600 -KIMI_IMAGE_HEIGHT = 30 -KIMI_IMAGE_WIDTH = 80 +KIMI_EXAMPLE_IMAGE_NUM_IMAGE_TOKENS = 600 +KIMI_EXAMPLE_IMAGE_NUM_PATCHES_HEIGHT = 30 +KIMI_EXAMPLE_IMAGE_NUM_PATCHES_WIDTH = 80 MAX_POSITION_EMBEDDINGS = 32768 FP16_BYTES = 2 DEFAULT_NUM_HEADS = 64 diff --git a/QEfficient/utils/load_kimi_utils.py b/QEfficient/utils/load_kimi_utils.py index 4aec58ceea..7cc387a9be 100644 --- a/QEfficient/utils/load_kimi_utils.py +++ b/QEfficient/utils/load_kimi_utils.py @@ -176,7 +176,11 @@ def load_kimi_k25_model_from_config(config): torch_dtype = getattr(model.config, "torch_dtype", None) if torch_dtype == torch.bfloat16 or torch_dtype == torch.float16: model = model.to(torch.float32) - return model.eval() + model.vision_tower.patch_embed.pos_emb.interpolation_mode = "bilinear" + model.eval() + tokenizer = AutoTokenizer.from_pretrained(KIMI_K25_MODEL_NAME, trust_remote_code=True) + processor = AutoProcessor.from_pretrained(KIMI_K25_MODEL_NAME, trust_remote_code=True) + return model, tokenizer, processor def get_kimi_k25_num_image_tokens(config, grid_thws): @@ -456,7 +460,7 @@ def run_kimi_k25_hf_model_on_pytorch(model, processor, inputs, max_gen_len): @torch.no_grad() -def run_kimi_k25_hf_model_on_pytorch_cb(model, processor, images, queries, max_gen_len): +def run_kimi_k25_hf_model_on_pytorch_CB(model, processor, images, queries, max_gen_len): generated_tokens = [] eos_token_id = getattr(model.config, "eos_token_id", None) diff --git a/examples/kimi_k2/export_kimi_k25_vision.py b/examples/kimi_k2/export_kimi_k25_vision.py index 59a3204417..774d55d6c6 100644 --- a/examples/kimi_k2/export_kimi_k25_vision.py +++ b/examples/kimi_k2/export_kimi_k25_vision.py @@ -22,22 +22,13 @@ NUM_EXPERTS_PER_TOKEN, NUM_TEXT_LAYERS, NUM_VISION_LAYERS, + ensure_torch_fx_import_compatibility, + get_kimi_k25_num_image_tokens, load_kimi_k25_class, -) -from QEfficient.utils.load_kimi_utils import ( - ensure_torch_fx_import_compatibility as _ensure_torch_fx_import_compatibility, -) -from QEfficient.utils.load_kimi_utils import ( - load_layer_subset_model as _load_layer_subset_model, -) -from QEfficient.utils.load_kimi_utils import ( - parse_expert_ids as _parse_expert_ids, -) -from QEfficient.utils.load_kimi_utils import ( - prepare_config as _prepare_config, -) -from QEfficient.utils.load_kimi_utils import ( - set_deterministic as _set_deterministic, + load_layer_subset_model, + parse_expert_ids, + prepare_config, + set_deterministic, ) @@ -51,7 +42,7 @@ def parse_args(): ) parser.add_argument("--num-vision-layers", type=int, default=NUM_VISION_LAYERS) parser.add_argument("--num-text-layers", type=int, default=NUM_TEXT_LAYERS) - parser.add_argument("--expert-ids", type=_parse_expert_ids, default=LOADED_EXPERT_IDS) + parser.add_argument("--expert-ids", type=parse_expert_ids, default=LOADED_EXPERT_IDS) parser.add_argument("--num-experts-per-token", type=int, default=NUM_EXPERTS_PER_TOKEN) parser.add_argument( "--image-url", @@ -66,9 +57,9 @@ def parse_args(): def main(): args = parse_args() - _set_deterministic(1234) - _ensure_torch_fx_import_compatibility() - config = _prepare_config(args.model_path) + set_deterministic(1234) + ensure_torch_fx_import_compatibility() + config = prepare_config(args.model_path) kimi_cls = load_kimi_k25_class(args.model_path) model_kwargs = { @@ -81,7 +72,7 @@ def main(): if args.full_model: model, tokenizer, processor = kimi_cls.from_pretrained(str(args.model_path), **model_kwargs) elif args.num_vision_layers is not None and args.num_text_layers is not None: - model, tokenizer, processor = _load_layer_subset_model( + model, tokenizer, processor = load_layer_subset_model( model_path=args.model_path, kimi_cls=kimi_cls, config=config, @@ -169,17 +160,10 @@ def main(): inputs["pixel_values"] = inputs["pixel_values"].to(qeff_model.model.config.torch_dtype) - merge_kernel_size = getattr(model.config.vision_config, "merge_kernel_size", (2, 2)) - if isinstance(merge_kernel_size, int): - kernel_height = kernel_width = merge_kernel_size - merge_kernel_size = (merge_kernel_size, merge_kernel_size) - else: - kernel_height, kernel_width = merge_kernel_size - num_patches = int(inputs["pixel_values"].shape[0]) h = int(inputs["grid_thws"][0, 1].item()) w = int(inputs["grid_thws"][0, 2].item()) - num_image_tokens = int(inputs["pixel_values"].shape[0] // (kernel_height * kernel_width)) + num_image_tokens = (get_kimi_k25_num_image_tokens(config, inputs["grid_thws"]),) qeff_model.compile( qaic_config=qaic_config, diff --git a/tests/transformers/models/image_text_to_text/test_continuous_batching.py b/tests/transformers/models/image_text_to_text/test_continuous_batching.py index 4645f0a39a..d1a2f3fc65 100644 --- a/tests/transformers/models/image_text_to_text/test_continuous_batching.py +++ b/tests/transformers/models/image_text_to_text/test_continuous_batching.py @@ -24,20 +24,12 @@ from QEfficient import QEFFAutoModelForCausalLM, QEFFAutoModelForImageTextToText from QEfficient.utils.load_kimi_utils import ( - get_kimi_k25_num_image_tokens as _get_kimi_k25_num_image_tokens, -) -from QEfficient.utils.load_kimi_utils import ( + get_kimi_k25_num_image_tokens, get_kimi_k25_test_config, + is_kimi_k25, load_kimi_k25_layer_subset_model, -) -from QEfficient.utils.load_kimi_utils import ( - is_kimi_k25 as _is_kimi_k25, -) -from QEfficient.utils.load_kimi_utils import ( - load_kimi_k25_model_from_config as _load_kimi_k25_model_from_config, -) -from QEfficient.utils.load_kimi_utils import ( - run_kimi_k25_hf_model_on_pytorch_cb as _run_kimi_k25_hf_model_on_pytorch_CB, + load_kimi_k25_model_from_config, + run_kimi_k25_hf_model_on_pytorch_CB, ) from QEfficient.utils.run_utils import ApiRunnerInternVL, ApiRunnerMolmo, ApiRunnerVlm from QEfficient.utils.test_utils import ( @@ -58,14 +50,6 @@ NEW_GENERATION_TOKENS = 10 -def _get_kimi_k25_test_config(model_name: str): - return get_kimi_k25_test_config(model_name, model_config_dict) - - -def _load_kimi_k25_layer_subset_model(): - return load_kimi_k25_layer_subset_model(num_vision_layers=1, num_text_layers=1) - - def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100_CB( model_name: str, manual_cleanup: callable, @@ -78,7 +62,6 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100_CB( ): prompt_len = model_config_dict[model_name]["prompt_len"] ctx_len = model_config_dict[model_name]["ctx_len"] - max_gen_len = (NEW_GENERATION_TOKENS,) img_size = model_config_dict[model_name].get("img_size") image_urls = model_config_dict[model_name]["img_url_list"] queries = model_config_dict[model_name]["text_prompt_list"] @@ -86,13 +69,9 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100_CB( batch_size = model_config_dict[model_name]["batch_size"] full_batch_size = model_config_dict[model_name]["full_batch_size"] max_gen_len = NEW_GENERATION_TOKENS - kimi_tokenizer = None - kimi_processor = None - if _is_kimi_k25(model_name): - full_batch_size = 1 - if _is_kimi_k25(model_name) and config is None: - model_hf, kimi_tokenizer, kimi_processor = _load_kimi_k25_layer_subset_model() + if is_kimi_k25(model_name) and config is None: + model_hf, tokenizer, processor = load_kimi_k25_layer_subset_model() config = model_hf.config qeff_model = QEFFAutoModelForImageTextToText( copy.deepcopy(model_hf), @@ -101,11 +80,10 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100_CB( torch_dtype=torch.float32, continuous_batching=True, ) - elif _is_kimi_k25(model_name): + elif is_kimi_k25(model_name): if config is None: - config = _get_kimi_k25_test_config(model_name) - model_hf = _load_kimi_k25_model_from_config(config) - model_hf.vision_tower.patch_embed.pos_emb.interpolation_mode = "bilinear" + config = get_kimi_k25_test_config(model_name, model_config_dict) + model_hf, tokenizer, processor = load_kimi_k25_model_from_config(config) qeff_model = QEFFAutoModelForImageTextToText( copy.deepcopy(model_hf), kv_offload=kv_offload, @@ -233,9 +211,7 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100_CB( model_hf, image_list, prompt_list, generation_config ) compile_kwargs["img_size"] = img_size - elif _is_kimi_k25(model_name): - processor = kimi_processor or AutoProcessor.from_pretrained(model_name, trust_remote_code=True) - tokenizer = kimi_tokenizer or AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) + elif is_kimi_k25(model_name): image_height = None image_width = None image_urls = [image_urls[0]] * len(queries) @@ -261,11 +237,11 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100_CB( num_patches.append(int(inputs["pixel_values"].shape[0])) image_heights.append(int(inputs["grid_thws"][0, 1].item())) image_widths.append(int(inputs["grid_thws"][0, 2].item())) - num_image_tokens.append(_get_kimi_k25_num_image_tokens(config, inputs["grid_thws"])) + num_image_tokens.append(get_kimi_k25_num_image_tokens(config, inputs["grid_thws"])) image_list = [images[0]] * full_batch_size prompt_list = [queries[0]] * full_batch_size - pytorch_hf_tokens = _run_kimi_k25_hf_model_on_pytorch_CB( + pytorch_hf_tokens = run_kimi_k25_hf_model_on_pytorch_CB( copy.deepcopy(model_hf), processor, image_list, prompt_list, max_gen_len ) compile_kwargs.update( @@ -318,9 +294,7 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100_CB( compile_kwargs["img_size"] = img_size qeff_model.compile(**compile_kwargs) - # if _is_kimi_k25(model_name): - # manual_cleanup(qeff_model.onnx_path) - # return + print("QPC Outputs (QAIC):") exec_info = qeff_model.generate( tokenizer=tokenizer, @@ -343,8 +317,8 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100_CB( model_hf, images, queries, generation_config=generation_config ) else: - if _is_kimi_k25(model_name): - pytorch_hf_tokens = _run_kimi_k25_hf_model_on_pytorch_CB( + if is_kimi_k25(model_name): + pytorch_hf_tokens = run_kimi_k25_hf_model_on_pytorch_CB( copy.deepcopy(model_hf), processor, images, queries, max_gen_len ) else: diff --git a/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py b/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py index a43b6f90e1..19cd2d6ed0 100644 --- a/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py +++ b/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py @@ -27,19 +27,10 @@ from QEfficient.utils._utils import create_json from QEfficient.utils.constants import QnnConstants from QEfficient.utils.load_kimi_utils import ( - get_kimi_k25_num_image_tokens as _get_kimi_k25_num_image_tokens, -) -from QEfficient.utils.load_kimi_utils import ( - get_kimi_k25_test_config, -) -from QEfficient.utils.load_kimi_utils import ( - is_kimi_k25 as _is_kimi_k25, -) -from QEfficient.utils.load_kimi_utils import ( - load_kimi_k25_layer_subset_model as _load_kimi_k25_layer_subset_model, -) -from QEfficient.utils.load_kimi_utils import ( - run_kimi_k25_hf_model_on_pytorch as _run_kimi_k25_hf_model_on_pytorch, + get_kimi_k25_num_image_tokens, + is_kimi_k25, + load_kimi_k25_layer_subset_model, + run_kimi_k25_hf_model_on_pytorch, ) from QEfficient.utils.run_utils import ApiRunnerInternVL, ApiRunnerMolmo, ApiRunnerVlm from QEfficient.utils.test_utils import ( @@ -64,10 +55,6 @@ NEW_GENERATION_TOKENS = 10 -def _get_kimi_k25_test_config(model_name: str): - return get_kimi_k25_test_config(model_name, model_config_dict) - - def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100( model_name: str, manual_cleanup: callable, @@ -99,10 +86,8 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100( n_layer = num_hidden_layers qaic_config = copy.deepcopy(qaic_config) if qaic_config is not None else None - kimi_tokenizer = None - kimi_processor = None - if _is_kimi_k25(model_name) and config is None: - model_hf, kimi_tokenizer, kimi_processor = _load_kimi_k25_layer_subset_model() + if is_kimi_k25(model_name) and config is None: + model_hf, tokenizer, processor = load_kimi_k25_layer_subset_model() config = model_hf.config qeff_model = QEFFAutoModelForImageTextToText( copy.deepcopy(model_hf), @@ -259,9 +244,7 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100( inputs["pixel_values"] = inputs.pop("images") compile_kwargs["img_size"] = img_size - elif _is_kimi_k25(model_name): - processor = kimi_processor - tokenizer = kimi_tokenizer + elif is_kimi_k25(model_name): image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB") conversation = [ { @@ -273,38 +256,20 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100( }, ] prompt = processor.apply_chat_template(conversation, add_generation_prompt=True) - api_runner = ApiRunnerVlm( - batch_size, - processor, - config, - image, - conversation, - prompt, - prompt_len, - ctx_len, - max_gen_len, - num_hidden_layers, - ) - inputs = processor( - messages=conversation, - add_generation_prompt=True, - tokenize=False, - return_tensors="pt", - ) - pytorch_hf_tokens = _run_kimi_k25_hf_model_on_pytorch(copy.deepcopy(model_hf), processor, inputs, max_gen_len) inputs = processor( messages=conversation, add_generation_prompt=True, tokenize=False, return_tensors="pt", ) + pytorch_hf_tokens = run_kimi_k25_hf_model_on_pytorch(copy.deepcopy(model_hf), processor, inputs, max_gen_len) compile_kwargs.update( { "prefill_seq_len": 1, "num_patches": int(inputs["pixel_values"].shape[0]), "h": int(inputs["grid_thws"][0, 1].item()), "w": int(inputs["grid_thws"][0, 2].item()), - "num_image_tokens": _get_kimi_k25_num_image_tokens(config, inputs["grid_thws"]), + "num_image_tokens": get_kimi_k25_num_image_tokens(config, inputs["grid_thws"]), } ) From d7b308dffae7ac0486555c3dae687d728fb265b2 Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Fri, 24 Jul 2026 11:03:57 +0530 Subject: [PATCH 27/33] CB modeling change Signed-off-by: Mamta Singh --- .../models/deepseek_v3/modeling_deepseek.py | 52 +++++++------------ 1 file changed, 20 insertions(+), 32 deletions(-) diff --git a/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py b/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py index 5940cc1a51..3356c6b95d 100644 --- a/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py +++ b/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py @@ -1086,50 +1086,38 @@ def original_moe(self, x, topk_ids, topk_weight): return final_out def moe_waa_unpack(self, hidden_states, topk_indices, topk_weights): - # GATHER - collect weights for selected experts - gate_proj_qweight = self.all_gate_qweight[topk_indices.flatten()] - gate_proj_scales = self.all_gate_scales[topk_indices.flatten()] - gate_proj_qzeros = self.all_gate_qzeros[topk_indices.flatten()] - - up_proj_qweight = self.all_up_qweight[topk_indices.flatten()] - up_proj_scales = self.all_up_scales[topk_indices.flatten()] - up_proj_qzeros = self.all_up_qzeros[topk_indices.flatten()] - - down_proj_qweight = self.all_down_qweight[topk_indices.flatten()] - down_proj_scales = self.all_down_scales[topk_indices.flatten()] - down_proj_qzeros = self.all_down_qzeros[topk_indices.flatten()] - - gate_proj_unpacked = CastToUInt4Func.apply(gate_proj_qweight) - gate_zeros_unpacked = CastToUInt4Func.apply(gate_proj_qzeros) + gate_proj_unpacked = CastToUInt4Func.apply(self.all_gate_qweight) + gate_zeros_unpacked = CastToUInt4Func.apply(self.all_gate_qzeros) gate_proj_dq = DequantizeLinearFunc.apply( - gate_proj_unpacked, gate_proj_scales, gate_zeros_unpacked, self.group_size + gate_proj_unpacked, self.all_gate_scales, gate_zeros_unpacked, self.group_size ) - up_proj_unpacked = CastToUInt4Func.apply(up_proj_qweight) - up_zeros_unpacked = CastToUInt4Func.apply(up_proj_qzeros) - up_proj_dq = DequantizeLinearFunc.apply(up_proj_unpacked, up_proj_scales, up_zeros_unpacked, self.group_size) - - down_proj_unpacked = CastToUInt4Func.apply(down_proj_qweight) - down_zeros_unpacked = CastToUInt4Func.apply(down_proj_qzeros) - down_proj_dq = DequantizeLinearFunc.apply( - down_proj_unpacked, down_proj_scales, down_zeros_unpacked, self.group_size + up_proj_unpacked = CastToUInt4Func.apply(self.all_up_qweight) + up_zeros_unpacked = CastToUInt4Func.apply(self.all_up_qzeros) + up_proj_dq = DequantizeLinearFunc.apply( + up_proj_unpacked, self.all_up_scales, up_zeros_unpacked, self.group_size ) - # Reshape for bmm: (bs*seq_len*top_k, 1, hidden_size) - expert_in = ( - hidden_states.unsqueeze(1).expand(-1, self.gate.top_k, -1).contiguous().view(-1, 1, self.in_features_gate) + down_proj_unpacked = CastToUInt4Func.apply(self.all_down_qweight) + down_zeros_unpacked = CastToUInt4Func.apply(self.all_down_qzeros) + down_proj_dq = DequantizeLinearFunc.apply( + down_proj_unpacked, self.all_down_scales, down_zeros_unpacked, self.group_size ) + num_experts = self.all_gate_qweight.shape[0] + expert_in = hidden_states.unsqueeze(0).expand(num_experts, -1, -1) gate_out = torch.bmm(expert_in, gate_proj_dq.transpose(1, 2).to(expert_in.dtype)) up_out = torch.bmm(expert_in, up_proj_dq.transpose(1, 2).to(expert_in.dtype)) hidden = self.act_fn(gate_out) * up_out down_out = torch.bmm(hidden, down_proj_dq.transpose(1, 2).to(expert_in.dtype)) - down_out = down_out.view(-1, self.gate.top_k, self.out_features_down) - - down_out = down_out * topk_weights.unsqueeze(-1) - - return torch.einsum("abc-> ac", down_out) + routed_out = down_out.transpose(0, 1) + selected_out = torch.gather( + routed_out, + 1, + topk_indices.unsqueeze(-1).expand(-1, self.gate.top_k, self.out_features_down), + ) + return torch.einsum("abc,ab->ac", selected_out, topk_weights) def forward(self, hidden_states): print("Using new MoE forward with weights as activations") From f87faaf6624a22fd450c1b9013a99110113488bd Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Fri, 24 Jul 2026 21:19:07 +0530 Subject: [PATCH 28/33] 3 qpc mdp export Signed-off-by: Mamta Singh --- .../test_kimi_k25_disagg.py | 99 ++++--------------- 1 file changed, 19 insertions(+), 80 deletions(-) diff --git a/tests/transformers/models/image_text_to_text/test_kimi_k25_disagg.py b/tests/transformers/models/image_text_to_text/test_kimi_k25_disagg.py index b576104219..ac2801dd9f 100644 --- a/tests/transformers/models/image_text_to_text/test_kimi_k25_disagg.py +++ b/tests/transformers/models/image_text_to_text/test_kimi_k25_disagg.py @@ -20,48 +20,25 @@ from QEfficient.utils.load_kimi_utils import ( LOADED_EXPERT_IDS, NUM_EXPERTS_PER_TOKEN, + ensure_torch_fx_import_compatibility, load_kimi_k25_class, -) -from QEfficient.utils.load_kimi_utils import ( - ensure_torch_fx_import_compatibility as _ensure_torch_fx_import_compatibility, -) -from QEfficient.utils.load_kimi_utils import ( - load_layer_subset_model as _load_layer_subset_model, -) -from QEfficient.utils.load_kimi_utils import ( - prepare_config as _prepare_config, -) -from QEfficient.utils.load_kimi_utils import ( - resolve_model_path as _resolve_model_path, -) -from QEfficient.utils.load_kimi_utils import ( - set_deterministic as _set_deterministic, + load_layer_subset_model, + prepare_config, + resolve_model_path, + run_kimi_k25_hf_model_on_pytorch, + set_deterministic, ) PREFILL_SEQ_LEN = 512 CTX_LEN = 2048 BATCH_SIZE = 1 GENERATION_LEN = 10 -NUM_VISION_LAYERS = 2 -NUM_TEXT_LAYERS = 2 +NUM_VISION_LAYERS = 4 +NUM_TEXT_LAYERS = 4 IMAGE_URL = "https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png" TEXT_PROMPT = "Describe this image." -def _has_qaic_runtime_access() -> bool: - try: - _ = QAICInferenceSession - except Exception: - return False - try: - import qaicrt - - _ctx = qaicrt.Context() - return True - except Exception: - return False - - def _prepare_inputs(processor): image = Image.open(BytesIO(requests.get(IMAGE_URL, timeout=30).content)).convert("RGB") messages = [ @@ -90,43 +67,6 @@ def _clone_inputs(inputs): return {k: (v.clone() if torch.is_tensor(v) else copy.deepcopy(v)) for k, v in inputs.items()} -def _greedy_generate_hf(model, inputs, max_new_tokens: int): - generated_ids = inputs["input_ids"] - attention_mask = inputs["attention_mask"] - pixel_values = inputs["pixel_values"] - grid_thws = inputs["grid_thws"] - - eos_token_id = getattr(model.config, "eos_token_id", None) - if eos_token_id is None and hasattr(model.config, "text_config"): - eos_token_id = getattr(model.config.text_config, "eos_token_id", None) - - for _ in range(max_new_tokens): - outputs = model( - input_ids=generated_ids, - attention_mask=attention_mask, - pixel_values=pixel_values, - grid_thws=grid_thws, - use_cache=False, - return_dict=True, - ) - logits = outputs[0] if isinstance(outputs, tuple) else outputs.logits - next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True) - - generated_ids = torch.cat([generated_ids, next_token], dim=1) - attention_mask = torch.cat( - [ - attention_mask, - torch.ones((attention_mask.shape[0], 1), dtype=attention_mask.dtype, device=attention_mask.device), - ], - dim=1, - ) - - if eos_token_id is not None and torch.all(next_token == eos_token_id): - break - - return generated_ids[:, -max_new_tokens:] - - def _assert_onnx_path(onnx_path, label: str) -> Path: assert onnx_path is not None, f"{label} compile did not set an ONNX path" onnx_path = Path(onnx_path) @@ -179,13 +119,13 @@ def _update_retained_states(target_inputs: dict[str, np.ndarray], source_outputs def _load_kimi_subset_model(): - _set_deterministic(1234) - _ensure_torch_fx_import_compatibility() - model_path = _resolve_model_path() - config = _prepare_config(model_path) + set_deterministic(1234) + ensure_torch_fx_import_compatibility() + model_path = resolve_model_path() + config = prepare_config(model_path) kimi_cls = load_kimi_k25_class(model_path) - model, tokenizer, processor = _load_layer_subset_model( + model, tokenizer, processor = load_layer_subset_model( model_path=model_path, kimi_cls=kimi_cls, config=config, @@ -214,13 +154,11 @@ def _compile_disagg_qpcs(qeff_model: QEFFAutoModelForImageTextToText, compile_di common_compile_kwargs = { "batch_size": BATCH_SIZE, "ctx_len": CTX_LEN, - "num_devices": 1, "num_cores": 16, "mxfp6_matmul": False, "split_model_io": True, "mos": 1, "aic_enable_depth_first": True, - "use_onnx_subfunctions": True, "layerwise": False, **compile_dims, } @@ -230,17 +168,17 @@ def _compile_disagg_qpcs(qeff_model: QEFFAutoModelForImageTextToText, compile_di prefill_seq_len=PREFILL_SEQ_LEN, skip_vision=False, skip_lang=True, + num_devices=1, **common_compile_kwargs, ) compiled_onnx_paths["vision"] = _assert_onnx_path(qeff_model.vision_model.onnx_path, "vision") prefill_qpc_path = qeff_model.compile( prefill_seq_len=PREFILL_SEQ_LEN, - # retain_full_kv=True, prefill_only=True, - # enable_chunking=True, skip_vision=True, skip_lang=False, + num_devices=4, **common_compile_kwargs, ) compiled_onnx_paths["prefill"] = _assert_onnx_path(qeff_model.lang_model.onnx_path, "prefill") @@ -250,6 +188,7 @@ def _compile_disagg_qpcs(qeff_model: QEFFAutoModelForImageTextToText, compile_di prefill_only=False, skip_vision=True, skip_lang=False, + num_devices=4, **common_compile_kwargs, ) compiled_onnx_paths["decode"] = _assert_onnx_path(qeff_model.lang_model.onnx_path, "decode") @@ -341,13 +280,13 @@ def _run_disagg_qaic_generation( @pytest.mark.on_qaic @pytest.mark.multimodal def test_kimi_k25_disagg_qaic_vs_hf_fp32(): - if not _has_qaic_runtime_access(): - pytest.skip("QAIC runtime is not available.") model, tokenizer, processor = _load_kimi_subset_model() inputs = _prepare_inputs(processor) inputs = {name: (value.to("cpu") if torch.is_tensor(value) else value) for name, value in inputs.items()} - hf_tokens = _greedy_generate_hf(copy.deepcopy(model), _clone_inputs(inputs), max_new_tokens=GENERATION_LEN) + hf_tokens = run_kimi_k25_hf_model_on_pytorch( + copy.deepcopy(model), processor, _clone_inputs(inputs), max_gen_len=GENERATION_LEN + ) qeff_model = QEFFAutoModelForImageTextToText( model, From 1a568514e3b0e22c01b6deac97850c64a8f7b4e9 Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Fri, 24 Jul 2026 21:35:18 +0530 Subject: [PATCH 29/33] subfunction changes Signed-off-by: Mamta Singh --- QEfficient/base/modeling_qeff.py | 20 +----------- QEfficient/base/onnx_transforms.py | 2 +- .../generation/text_generation_inference.py | 25 ++------------- QEfficient/generation/vlm_generation.py | 1 + QEfficient/transformers/cache_utils.py | 32 +++++++++---------- .../models/kimi_k25/modeling_kimi_k25.py | 22 +------------ QEfficient/utils/constants.py | 1 - docs/source/introduction.md | 2 +- docs/source/validate.md | 2 +- examples/kimi_k2/export_kimi_k25_vision.py | 6 ++-- .../test_kimi_k25_disagg.py | 2 +- 11 files changed, 27 insertions(+), 88 deletions(-) diff --git a/QEfficient/base/modeling_qeff.py b/QEfficient/base/modeling_qeff.py index 5a3c3d935f..1a5b77ec38 100755 --- a/QEfficient/base/modeling_qeff.py +++ b/QEfficient/base/modeling_qeff.py @@ -101,19 +101,6 @@ def _restore_output_names_exact(model: onnx.ModelProto, output_names: List[str]) _rename_graph_value(model.graph, current_name, expected_name) -def _has_exported_model_subfunctions(onnx_path: Union[str, Path]) -> bool: - """Return True when the ONNX has non-custom local functions for compiler subfunction mode.""" - try: - model = onnx.load(onnx_path, load_external_data=False) - except Exception: - return True - - custom_function_names = { - onnxscript_func.to_function_proto().name for _, onnxscript_func in CustomOpTransform._custom_ops.values() - } - return any(function.name not in custom_function_names for function in model.functions) - - class QEFFBaseModel(ABC): """ Base class for all the model classes (i.e. LLMs, SD, quantized etc.). @@ -1084,11 +1071,6 @@ def _compile( return self.qpc_path - use_compiler_subfunctions = use_onnx_subfunctions and _has_exported_model_subfunctions(onnx_path) - - if use_compiler_subfunctions and prefill_only is True: - compiler_options.setdefault("elf_va_limit", constants.DEFAULT_ONNX_SUBFUNCTION_PREFILL_ELF_VA_LIMIT_MB) - command = ( constants.COMPILER + [ @@ -1172,7 +1154,7 @@ def _compile( except Exception: pass - if use_compiler_subfunctions: + if use_onnx_subfunctions: logger.info("Using ONNX subfunctions for compilation.") command.append("-sub-functions") diff --git a/QEfficient/base/onnx_transforms.py b/QEfficient/base/onnx_transforms.py index e56baaab8b..b24b4ebcb5 100644 --- a/QEfficient/base/onnx_transforms.py +++ b/QEfficient/base/onnx_transforms.py @@ -44,8 +44,8 @@ CtxScatterFuncCB, CtxScatterFuncCB3D, ) -from QEfficient.customop.quantization_ops import CastToUInt4, CastToUInt4Func from QEfficient.customop.onnxscript_utils import get_onnxscript_func +from QEfficient.customop.quantization_ops import CastToUInt4, CastToUInt4Func from QEfficient.customop.rms_norm import CustomRMSNorm, CustomRMSNormFunc from QEfficient.utils import constants from QEfficient.utils.constants import FILE_CHUNK_SIZE_DEFAULT, SIZE_THRESHOLD_DEFAULT diff --git a/QEfficient/generation/text_generation_inference.py b/QEfficient/generation/text_generation_inference.py index f80fd4d77e..4e4b71366e 100755 --- a/QEfficient/generation/text_generation_inference.py +++ b/QEfficient/generation/text_generation_inference.py @@ -462,7 +462,6 @@ def __init__( self.include_guided_decoding = include_guided_decoding self.sampling_params = sampling_params self._qpc_path = qpc_path # Store qpc_path for later use - self.tokenizer = tokenizer # Load QPC self._session = QAICInferenceSession( @@ -497,6 +496,7 @@ def __init__( self.decode_pos_ids = None self.generation_len = None + self.tokenizer = tokenizer self._set_tokenizer_params() # set tokenizer params # Skip inputs/outputs self._session.skip_buffers( @@ -569,18 +569,7 @@ def _fetch_decode_seq_len( decode_seq_len = min( [x[self._session.binding_index_map["input_ids"]][1][1] for x in self._session.allowed_shapes] ) - elif "input_ids" in self._session.binding_index_map: - dims = self._session.bindings[self._session.binding_index_map["input_ids"]].dims - if len(dims) > 1 and dims[1] not in (None, -1): - decode_seq_len = dims[1] - - if decode_seq_len is None: - decode_seq_len = getattr(self, "_prefill_seq_len", None) - - if decode_seq_len is None: - raise ValueError("Unable to determine decode sequence length from QPC input_ids binding.") - - return int(decode_seq_len) + return decode_seq_len def _fetch_vocab_size( self, @@ -599,16 +588,6 @@ def _fetch_vocab_size( ) if self._session.allowed_shapes: return [x[self._session.binding_index_map[key]] for x in self._session.allowed_shapes][0][1][2] - if key in self._session.binding_index_map: - dims = self._session.bindings[self._session.binding_index_map[key]].dims - if len(dims) >= 3 and dims[2] not in (None, -1): - return dims[2] - if self.tokenizer is not None: - vocab_size = getattr(self.tokenizer, "vocab_size", None) - if vocab_size is not None: - return vocab_size - return len(self.tokenizer) - return None return self._session.bindings[self._session.binding_index_map[key]].dims[2] diff --git a/QEfficient/generation/vlm_generation.py b/QEfficient/generation/vlm_generation.py index 3d16b7f16c..013452fbc1 100755 --- a/QEfficient/generation/vlm_generation.py +++ b/QEfficient/generation/vlm_generation.py @@ -737,6 +737,7 @@ def _generate_multi_frame_specialization( generation_len: int = None, stream: List[str] = None, ): + exec_batch_size = self.batch_size max_gen_length = self._ctx_len if not generation_len else max(self._ctx_len, generation_len) self.initialize_decode_inputs( diff --git a/QEfficient/transformers/cache_utils.py b/QEfficient/transformers/cache_utils.py index c5cf88f0a1..7031ac11bb 100755 --- a/QEfficient/transformers/cache_utils.py +++ b/QEfficient/transformers/cache_utils.py @@ -428,9 +428,9 @@ def update_ckv(self, compressed_kv, cache_kwargs): if batch_index is not None: invalid_scatter_index = torch.iinfo(torch.int32).max scatter_position_ids = torch.where(position_ids < 0, invalid_scatter_index, position_ids) - self.ckv = CtxScatterFuncCB.apply(self.ckv, batch_index, scatter_position_ids, compressed_kv) + self.ckv = ctx_scatter_cb.apply(self.ckv, batch_index, scatter_position_ids, compressed_kv) else: - self.ckv = CtxScatterFunc.apply(self.ckv, position_ids, compressed_kv) + self.ckv = ctx_scatter.apply(self.ckv, position_ids, compressed_kv) ckv_out = self.ckv ctx_len = ckv_out.shape[-2] @@ -441,9 +441,9 @@ def update_ckv(self, compressed_kv, cache_kwargs): ctx_indices = torch.where(invalid_mask, invalid_idx_value, ctx_indices) if batch_index is not None: - ckv_out = CtxGatherFuncCB.apply(ckv_out, batch_index, ctx_indices, ctx_len) + ckv_out = ctx_gather_cb.apply(ckv_out, batch_index, ctx_indices, ctx_len) else: - ckv_out = CtxGatherFunc.apply(ckv_out, ctx_indices, ctx_len) + ckv_out = ctx_gather.apply(ckv_out, ctx_indices, ctx_len) ckv_out = torch.where(invalid_mask.unsqueeze(-1), torch.tensor(0.0, dtype=torch.float32), ckv_out) return ckv_out @@ -454,9 +454,9 @@ def update_k_pe(self, k_pe_cache, cache_kwargs): if batch_index is not None: invalid_scatter_index = torch.iinfo(torch.int32).max scatter_position_ids = torch.where(position_ids < 0, invalid_scatter_index, position_ids) - self.k_pe = CtxScatterFuncCB.apply(self.k_pe, batch_index, scatter_position_ids, k_pe_cache) + self.k_pe = ctx_scatter_cb.apply(self.k_pe, batch_index, scatter_position_ids, k_pe_cache) else: - self.k_pe = CtxScatterFunc.apply(self.k_pe, position_ids, k_pe_cache) + self.k_pe = ctx_scatter.apply(self.k_pe, position_ids, k_pe_cache) k_pe_out = self.k_pe ctx_len = k_pe_out.shape[-2] @@ -467,9 +467,9 @@ def update_k_pe(self, k_pe_cache, cache_kwargs): ctx_indices = torch.where(invalid_mask, invalid_idx_value, ctx_indices) if batch_index is not None: - k_pe_out = CtxGatherFuncCB.apply(k_pe_out, batch_index, ctx_indices, ctx_len) + k_pe_out = ctx_gather_cb.apply(k_pe_out, batch_index, ctx_indices, ctx_len) else: - k_pe_out = CtxGatherFunc.apply(k_pe_out, ctx_indices, ctx_len) + k_pe_out = ctx_gather.apply(k_pe_out, ctx_indices, ctx_len) k_pe_out = torch.where(invalid_mask.unsqueeze(-1), torch.tensor(0.0, dtype=torch.float32), k_pe_out) return k_pe_out @@ -488,10 +488,10 @@ def read_only_blocked_ckv(self, start_index, end_index, cache_kwargs): ctx_indices = torch.where(invalid_mask, invalid_idx_value, ctx_indices) if batch_index is not None: - ckv_out = CtxGatherFuncBlockedKVCB.apply(ckv_out, batch_index, ctx_indices) + ckv_out = ctx_gather_blocked_kv_cb.apply(ckv_out, batch_index, ctx_indices) else: ctx_indices = ctx_indices.expand(batch, num_kv_heads, ctx_indices.shape[-1]) - ckv_out = CtxGatherFuncBlockedKV.apply(ckv_out, ctx_indices) + ckv_out = ctx_gather_blocked_kv.apply(ckv_out, ctx_indices) ckv_out = torch.where(invalid_mask.unsqueeze(-1), torch.zeros_like(ckv_out, dtype=ckv_out.dtype), ckv_out) return ckv_out @@ -511,10 +511,10 @@ def read_only_blocked_k_pe(self, start_index, end_index, cache_kwargs): ctx_indices = torch.where(invalid_mask, invalid_idx_value, ctx_indices) if batch_index is not None: - k_pe_out = CtxGatherFuncBlockedKVCB.apply(k_pe_out, batch_index, ctx_indices) + k_pe_out = ctx_gather_blocked_kv_cb.apply(k_pe_out, batch_index, ctx_indices) else: ctx_indices = ctx_indices.expand(batch, num_kv_heads, ctx_indices.shape[-1]) - k_pe_out = CtxGatherFuncBlockedKV.apply(k_pe_out, ctx_indices) + k_pe_out = ctx_gather_blocked_kv.apply(k_pe_out, ctx_indices) k_pe_out = torch.where(invalid_mask.unsqueeze(-1), torch.zeros_like(k_pe_out, dtype=k_pe_out.dtype), k_pe_out) return k_pe_out @@ -526,9 +526,9 @@ def write_only_k_pe(self, k_pe_cache, cache_kwargs): if batch_index is not None: invalid_scatter_index = torch.iinfo(torch.int32).max scatter_position_ids = torch.where(position_ids < 0, invalid_scatter_index, position_ids) - self.k_pe = CtxScatterFuncCB.apply(self.k_pe, batch_index, scatter_position_ids, k_pe_cache) + self.k_pe = ctx_scatter_cb.apply(self.k_pe, batch_index, scatter_position_ids, k_pe_cache) else: - self.k_pe = CtxScatterFunc.apply(self.k_pe, position_ids, k_pe_cache) + self.k_pe = ctx_scatter.apply(self.k_pe, position_ids, k_pe_cache) return self.k_pe def write_only_ckv(self, compressed_kv, cache_kwargs): @@ -538,9 +538,9 @@ def write_only_ckv(self, compressed_kv, cache_kwargs): if batch_index is not None: invalid_scatter_index = torch.iinfo(torch.int32).max scatter_position_ids = torch.where(position_ids < 0, invalid_scatter_index, position_ids) - self.ckv = CtxScatterFuncCB.apply(self.ckv, batch_index, scatter_position_ids, compressed_kv) + self.ckv = ctx_scatter_cb.apply(self.ckv, batch_index, scatter_position_ids, compressed_kv) else: - self.ckv = CtxScatterFunc.apply(self.ckv, position_ids, compressed_kv) + self.ckv = ctx_scatter.apply(self.ckv, position_ids, compressed_kv) return self.ckv diff --git a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py index 6e0af3c2c1..dca3de48f2 100644 --- a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py +++ b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py @@ -461,29 +461,8 @@ def get_submodules_for_export(self) -> Type[nn.Module]: This method should return the *class object* (not an instance). Downstream code can use this to find/build subfunctions for repeated blocks. """ - if self._has_uint4_weight_as_activation_moe(): - return set() return {self.language_model.model.layers[0].__class__} - def _has_uint4_weight_as_activation_moe(self) -> bool: - for layer in self.language_model.model.layers: - mlp = getattr(layer, "mlp", None) - if mlp is None: - continue - if all( - hasattr(mlp, attr) - for attr in ( - "all_gate_qweight", - "all_gate_qzeros", - "all_up_qweight", - "all_up_qzeros", - "all_down_qweight", - "all_down_qzeros", - ) - ): - return True - return False - def forward( self, input_ids: torch.LongTensor = None, @@ -560,6 +539,7 @@ def forward( ) image_position_delta = torch.clamp(merged_image_tokens - selected_image_tokens, min=0) image_idx = image_idx + selected_any.to(torch.int64) * image_position_delta + image_idx = image_idx + torch.zeros_like(input_ids[:, :1], dtype=image_idx.dtype, device=image_idx.device) if position_ids is None and attention_mask is not None: position_ids = attention_mask.long().cumsum(-1) - 1 diff --git a/QEfficient/utils/constants.py b/QEfficient/utils/constants.py index c15d68acaf..f26b09087b 100644 --- a/QEfficient/utils/constants.py +++ b/QEfficient/utils/constants.py @@ -44,7 +44,6 @@ # Compiler defaults DEFAULT_AIC_NUM_CORES = 16 DEFAULT_AIC_MXPF6_MATMUL = False -DEFAULT_ONNX_SUBFUNCTION_PREFILL_ELF_VA_LIMIT_MB = 2048 # Hashing defaults HASH_HEXDIGEST_STR_LEN = 16 KWARGS_INCLUSION_LIST = [ diff --git a/docs/source/introduction.md b/docs/source/introduction.md index 930b1c66fc..616d5bbd84 100644 --- a/docs/source/introduction.md +++ b/docs/source/introduction.md @@ -21,8 +21,8 @@ For other models, there is comprehensive documentation to inspire upon the chang 8. Unit test templates. -- [07/2026] Added support for Kimi-K2.5 vision-language model [moonshotai/Kimi-K2.5](https://huggingface.co/moonshotai/Kimi-K2.5) via `QEFFAutoModelForImageTextToText`. *Latest news* :
+- [07/2026] Added support for Kimi-K2.5 vision-language model [moonshotai/Kimi-K2.5](https://huggingface.co/moonshotai/Kimi-K2.5) via `QEFFAutoModelForImageTextToText`. - [07/2026] Added `dynamo` flag to `QEFFAutoModelForCausalLM.export()` to support `torch.onnx.export` dynamo-based ONNX export for CausalLM models - [06/2026] Added support for Gemma4 models, [google/gemma-4-E2B-it](https://huggingface.co/google/gemma-4-E2B), [google/gemma-4-26B-A4B-it](https://huggingface.co/google/gemma-4-26B-A4B-it) - [06/2026] Added support for Qwen3.6 model [Qwen/Qwen3.6-35B-A3B](https://huggingface.co/Qwen/Qwen3.6-35B-A3B) diff --git a/docs/source/validate.md b/docs/source/validate.md index 05bf81f895..be5009aac7 100644 --- a/docs/source/validate.md +++ b/docs/source/validate.md @@ -90,7 +90,7 @@ | **Qwen3_5MoeForConditionalGeneration** | Qwen3.5 | [Qwen/Qwen3.5-122B-A10B](https://huggingface.co/Qwen/Qwen3.5-122B-A10B)
[Qwen/Qwen3.5-35B-A3B](https://huggingface.co/Qwen/Qwen3.5-35B-A3B) | ✕ | ✔️ | ✕ | ✔️ | | **Qwen3_5ForConditionalGeneration** | Qwen3.6 | [Qwen/Qwen3.6-27B](https://huggingface.co/Qwen/Qwen3.6-27B) | ✕ | ✔️ | ✕ | ✔️ | | **Qwen3_5MoeForConditionalGeneration** | Qwen3.6 | [Qwen/Qwen3.6-35B-A3B](https://huggingface.co/Qwen/Qwen3.6-35B-A3B) | ✕ | ✔️ | ✕ | ✔️ | -| **KimiK25ForConditionalGeneration** | Kimi-K2.5 | [moonshotai/Kimi-K2.5](https://huggingface.co/moonshotai/Kimi-K2.5) | ✔️ | ✔️ | ✕ | ✕ | +| **KimiK25ForConditionalGeneration** | Kimi-K2.5 | [moonshotai/Kimi-K2.5](https://huggingface.co/moonshotai/Kimi-K2.5) | ✕ | ✔️ | ✕ | ✔️ | | **Mistral3ForConditionalGeneration** | Mistral3| [mistralai/Mistral-Small-3.1-24B-Instruct-2503](https://huggingface.co/mistralai/Mistral-Small-3.1-24B-Instruct-2503)| ✕ | ✔️ | ✕ | ✕ | ### Vision-Language Reranker Models (Text + Image Scoring) diff --git a/examples/kimi_k2/export_kimi_k25_vision.py b/examples/kimi_k2/export_kimi_k25_vision.py index 774d55d6c6..553d0afd66 100644 --- a/examples/kimi_k2/export_kimi_k25_vision.py +++ b/examples/kimi_k2/export_kimi_k25_vision.py @@ -15,9 +15,7 @@ from QEfficient import QEFFAutoModelForImageTextToText from QEfficient.utils.load_kimi_utils import ( - DEFAULT_MODEL_PATH as MODEL_PATH, -) -from QEfficient.utils.load_kimi_utils import ( + DEFAULT_MODEL_PATH, LOADED_EXPERT_IDS, NUM_EXPERTS_PER_TOKEN, NUM_TEXT_LAYERS, @@ -34,7 +32,7 @@ def parse_args(): parser = argparse.ArgumentParser(description="Load Kimi K2.5 with runtime compatibility for transformers==5.5.4.") - parser.add_argument("--model-path", type=Path, default=MODEL_PATH) + parser.add_argument("--model-path", type=Path, default=DEFAULT_MODEL_PATH) parser.add_argument( "--full-model", action="store_true", diff --git a/tests/transformers/models/image_text_to_text/test_kimi_k25_disagg.py b/tests/transformers/models/image_text_to_text/test_kimi_k25_disagg.py index ac2801dd9f..3bf0c42872 100644 --- a/tests/transformers/models/image_text_to_text/test_kimi_k25_disagg.py +++ b/tests/transformers/models/image_text_to_text/test_kimi_k25_disagg.py @@ -159,6 +159,7 @@ def _compile_disagg_qpcs(qeff_model: QEFFAutoModelForImageTextToText, compile_di "split_model_io": True, "mos": 1, "aic_enable_depth_first": True, + "use_onnx_subfunctions": True, "layerwise": False, **compile_dims, } @@ -280,7 +281,6 @@ def _run_disagg_qaic_generation( @pytest.mark.on_qaic @pytest.mark.multimodal def test_kimi_k25_disagg_qaic_vs_hf_fp32(): - model, tokenizer, processor = _load_kimi_subset_model() inputs = _prepare_inputs(processor) inputs = {name: (value.to("cpu") if torch.is_tensor(value) else value) for name, value in inputs.items()} From cbd5dec27df73b90e94ef0cc2c4f4b84c786b612 Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Sun, 26 Jul 2026 17:28:37 +0530 Subject: [PATCH 30/33] remove grid h and w from compile params Signed-off-by: Mamta Singh --- QEfficient/base/modeling_qeff.py | 2 - QEfficient/transformers/cache_utils.py | 32 +++---- .../models/kimi_k25/modeling_kimi_k25.py | 90 +++++++++---------- examples/image_text_to_text/README.md | 6 +- examples/kimi_k2/export_kimi_k25_vision.py | 32 ++++--- .../test_continuous_batching.py | 29 +----- .../test_image_text_to_text_models.py | 7 +- .../unit_test/models/test_model_quickcheck.py | 45 ++++++---- 8 files changed, 116 insertions(+), 127 deletions(-) diff --git a/QEfficient/base/modeling_qeff.py b/QEfficient/base/modeling_qeff.py index 1a5b77ec38..1c4c0ed6b7 100755 --- a/QEfficient/base/modeling_qeff.py +++ b/QEfficient/base/modeling_qeff.py @@ -1120,8 +1120,6 @@ def _compile( ) for key, value in compiler_options.items(): - if value is None: - continue option = "-" + key.replace("_", "-") if isinstance(value, bool): if value: diff --git a/QEfficient/transformers/cache_utils.py b/QEfficient/transformers/cache_utils.py index 7031ac11bb..4f016b2d49 100755 --- a/QEfficient/transformers/cache_utils.py +++ b/QEfficient/transformers/cache_utils.py @@ -428,9 +428,9 @@ def update_ckv(self, compressed_kv, cache_kwargs): if batch_index is not None: invalid_scatter_index = torch.iinfo(torch.int32).max scatter_position_ids = torch.where(position_ids < 0, invalid_scatter_index, position_ids) - self.ckv = ctx_scatter_cb.apply(self.ckv, batch_index, scatter_position_ids, compressed_kv) + self.ckv = ctx_scatter_cb(self.ckv, batch_index, scatter_position_ids, compressed_kv) else: - self.ckv = ctx_scatter.apply(self.ckv, position_ids, compressed_kv) + self.ckv = ctx_scatter(self.ckv, position_ids, compressed_kv) ckv_out = self.ckv ctx_len = ckv_out.shape[-2] @@ -441,9 +441,9 @@ def update_ckv(self, compressed_kv, cache_kwargs): ctx_indices = torch.where(invalid_mask, invalid_idx_value, ctx_indices) if batch_index is not None: - ckv_out = ctx_gather_cb.apply(ckv_out, batch_index, ctx_indices, ctx_len) + ckv_out = ctx_gather_cb(ckv_out, batch_index, ctx_indices, ctx_len) else: - ckv_out = ctx_gather.apply(ckv_out, ctx_indices, ctx_len) + ckv_out = ctx_gather(ckv_out, ctx_indices, ctx_len) ckv_out = torch.where(invalid_mask.unsqueeze(-1), torch.tensor(0.0, dtype=torch.float32), ckv_out) return ckv_out @@ -454,9 +454,9 @@ def update_k_pe(self, k_pe_cache, cache_kwargs): if batch_index is not None: invalid_scatter_index = torch.iinfo(torch.int32).max scatter_position_ids = torch.where(position_ids < 0, invalid_scatter_index, position_ids) - self.k_pe = ctx_scatter_cb.apply(self.k_pe, batch_index, scatter_position_ids, k_pe_cache) + self.k_pe = ctx_scatter_cb(self.k_pe, batch_index, scatter_position_ids, k_pe_cache) else: - self.k_pe = ctx_scatter.apply(self.k_pe, position_ids, k_pe_cache) + self.k_pe = ctx_scatter(self.k_pe, position_ids, k_pe_cache) k_pe_out = self.k_pe ctx_len = k_pe_out.shape[-2] @@ -467,9 +467,9 @@ def update_k_pe(self, k_pe_cache, cache_kwargs): ctx_indices = torch.where(invalid_mask, invalid_idx_value, ctx_indices) if batch_index is not None: - k_pe_out = ctx_gather_cb.apply(k_pe_out, batch_index, ctx_indices, ctx_len) + k_pe_out = ctx_gather_cb(k_pe_out, batch_index, ctx_indices, ctx_len) else: - k_pe_out = ctx_gather.apply(k_pe_out, ctx_indices, ctx_len) + k_pe_out = ctx_gather(k_pe_out, ctx_indices, ctx_len) k_pe_out = torch.where(invalid_mask.unsqueeze(-1), torch.tensor(0.0, dtype=torch.float32), k_pe_out) return k_pe_out @@ -488,10 +488,10 @@ def read_only_blocked_ckv(self, start_index, end_index, cache_kwargs): ctx_indices = torch.where(invalid_mask, invalid_idx_value, ctx_indices) if batch_index is not None: - ckv_out = ctx_gather_blocked_kv_cb.apply(ckv_out, batch_index, ctx_indices) + ckv_out = ctx_gather_blocked_kv_cb(ckv_out, batch_index, ctx_indices) else: ctx_indices = ctx_indices.expand(batch, num_kv_heads, ctx_indices.shape[-1]) - ckv_out = ctx_gather_blocked_kv.apply(ckv_out, ctx_indices) + ckv_out = ctx_gather_blocked_kv(ckv_out, ctx_indices) ckv_out = torch.where(invalid_mask.unsqueeze(-1), torch.zeros_like(ckv_out, dtype=ckv_out.dtype), ckv_out) return ckv_out @@ -511,10 +511,10 @@ def read_only_blocked_k_pe(self, start_index, end_index, cache_kwargs): ctx_indices = torch.where(invalid_mask, invalid_idx_value, ctx_indices) if batch_index is not None: - k_pe_out = ctx_gather_blocked_kv_cb.apply(k_pe_out, batch_index, ctx_indices) + k_pe_out = ctx_gather_blocked_kv_cb(k_pe_out, batch_index, ctx_indices) else: ctx_indices = ctx_indices.expand(batch, num_kv_heads, ctx_indices.shape[-1]) - k_pe_out = ctx_gather_blocked_kv.apply(k_pe_out, ctx_indices) + k_pe_out = ctx_gather_blocked_kv(k_pe_out, ctx_indices) k_pe_out = torch.where(invalid_mask.unsqueeze(-1), torch.zeros_like(k_pe_out, dtype=k_pe_out.dtype), k_pe_out) return k_pe_out @@ -526,9 +526,9 @@ def write_only_k_pe(self, k_pe_cache, cache_kwargs): if batch_index is not None: invalid_scatter_index = torch.iinfo(torch.int32).max scatter_position_ids = torch.where(position_ids < 0, invalid_scatter_index, position_ids) - self.k_pe = ctx_scatter_cb.apply(self.k_pe, batch_index, scatter_position_ids, k_pe_cache) + self.k_pe = ctx_scatter_cb(self.k_pe, batch_index, scatter_position_ids, k_pe_cache) else: - self.k_pe = ctx_scatter.apply(self.k_pe, position_ids, k_pe_cache) + self.k_pe = ctx_scatter(self.k_pe, position_ids, k_pe_cache) return self.k_pe def write_only_ckv(self, compressed_kv, cache_kwargs): @@ -538,9 +538,9 @@ def write_only_ckv(self, compressed_kv, cache_kwargs): if batch_index is not None: invalid_scatter_index = torch.iinfo(torch.int32).max scatter_position_ids = torch.where(position_ids < 0, invalid_scatter_index, position_ids) - self.ckv = ctx_scatter_cb.apply(self.ckv, batch_index, scatter_position_ids, compressed_kv) + self.ckv = ctx_scatter_cb(self.ckv, batch_index, scatter_position_ids, compressed_kv) else: - self.ckv = ctx_scatter.apply(self.ckv, position_ids, compressed_kv) + self.ckv = ctx_scatter(self.ckv, position_ids, compressed_kv) return self.ckv diff --git a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py index dca3de48f2..cbfc86d228 100644 --- a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py +++ b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py @@ -539,7 +539,6 @@ def forward( ) image_position_delta = torch.clamp(merged_image_tokens - selected_image_tokens, min=0) image_idx = image_idx + selected_any.to(torch.int64) * image_position_delta - image_idx = image_idx + torch.zeros_like(input_ids[:, :1], dtype=image_idx.dtype, device=image_idx.device) if position_ids is None and attention_mask is not None: position_ids = attention_mask.long().cumsum(-1) - 1 @@ -833,13 +832,13 @@ def get_dummy_inputs( self.language_model.config.hidden_size, ) inputs_shapes["image_idx"] = (1, 1) - inputs_shapes["h"] = constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_HEIGHT - inputs_shapes["w"] = constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_WIDTH + inputs_shapes["grid_h"] = constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_HEIGHT + inputs_shapes["grid_w"] = constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_WIDTH vision_inputs = { "pixel_values": torch.zeros((inputs_shapes["pixel_values"]), dtype=self.config.torch_dtype), - "h_shape": torch.zeros((inputs_shapes["h"]), dtype=torch.int64), - "w_shape": torch.zeros((inputs_shapes["w"]), dtype=torch.int64), + "h_shape": torch.zeros((inputs_shapes["grid_h"]), dtype=torch.int64), + "w_shape": torch.zeros((inputs_shapes["grid_w"]), dtype=torch.int64), } input_ids = torch.zeros((bs, prefill_seq_len), dtype=torch.int64) @@ -912,8 +911,8 @@ def get_onnx_dynamic_axes( lang_dynamic_axes["batch_index"] = {0: "batch_size"} vision_dynamic_axes = { "pixel_values": {0: "num_patches"}, - "h_shape": {0: "h"}, - "w_shape": {0: "w"}, + "h_shape": {0: "grid_h"}, + "w_shape": {0: "grid_w"}, "vision_embeds": {0: "num_image_tokens"}, } @@ -960,13 +959,18 @@ def get_specializations( comp_ctx_lengths_prefill = compiler_options.pop("comp_ctx_lengths_prefill", None) comp_ctx_lengths_decode = compiler_options.pop("comp_ctx_lengths_decode", None) compiler_options.pop("img_size", None) - num_patches = compiler_options.pop("num_patches", None) - height = compiler_options.pop("height", None) - width = compiler_options.pop("width", None) - h = compiler_options.pop("h", None) - w = compiler_options.pop("w", None) + image_height = compiler_options.pop("image_height", None) + image_width = compiler_options.pop("image_width", None) + unsupported_shape_args = { + name: compiler_options.pop(name, None) + for name in ("height", "width", "h", "w", "num_patches", "num_image_tokens") + } + if any(value is not None for value in unsupported_shape_args.values()): + raise ValueError( + "Kimi-K2.5 compile expects image_height and image_width. " + "Do not pass height/width, h/w grid dimensions, num_patches, or num_image_tokens." + ) num_frames = compiler_options.pop("num_frames", 1) - num_image_tokens = compiler_options.pop("num_image_tokens", None) mm_processor_kwargs = compiler_options.pop("mm_processor_kwargs", None) or {} prefill_seq_len = prefill_seq_len if prefill_seq_len else constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN @@ -995,6 +999,8 @@ def validate_dimension_lists(heights, widths, height_name, width_name): f"Expected {height_name} and {width_name} to contain the same number of entries, " f"got {len(heights)} and {len(widths)}." ) + if any(height <= 0 for height in heights) or any(width <= 0 for width in widths): + raise ValueError(f"Expected {height_name} and {width_name} to contain positive values.") patch_size = getattr(self.config.vision_config, "patch_size", constants.KIMI_PATCH_SIZE) merge_kernel_size = getattr(self.config.vision_config, "merge_kernel_size", (2, 2)) @@ -1004,21 +1010,20 @@ def validate_dimension_lists(heights, widths, height_name, width_name): else: kernel_height, kernel_width = merge_kernel_size - if h is not None or w is not None: - heights = normalize_list(h, constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_HEIGHT) - widths = normalize_list(w, constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_WIDTH) - validate_dimension_lists(heights, widths, "h", "w") - elif height is not None or width is not None: - pixel_heights = normalize_list(height, constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_HEIGHT * patch_size) - pixel_widths = normalize_list(width, constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_WIDTH * patch_size) - validate_dimension_lists(pixel_heights, pixel_widths, "height", "width") + if (image_height is None) != (image_width is None): + raise ValueError("Kimi-K2.5 compile expects image_height and image_width to be provided together.") + + if image_height is not None: + pixel_heights = normalize_list(image_height, constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_HEIGHT * patch_size) + pixel_widths = normalize_list(image_width, constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_WIDTH * patch_size) + validate_dimension_lists(pixel_heights, pixel_widths, "image_height", "image_width") in_patch_limit = mm_processor_kwargs.get("in_patch_limit", 16384) patch_limit_on_one_side = mm_processor_kwargs.get("patch_limit_on_one_side", 512) factor_height = kernel_height * patch_size factor_width = kernel_width * patch_size - heights = [] - widths = [] + grid_heights = [] + grid_widths = [] for pixel_height, pixel_width in zip(pixel_heights, pixel_widths): scale = min( 1.0, @@ -1032,40 +1037,33 @@ def validate_dimension_lists(heights, widths, height_name, width_name): resized_width = min(max(1, int(pixel_width * scale)), patch_limit_on_one_side * patch_size) pad_height = (factor_height - resized_height % factor_height) % factor_height pad_width = (factor_width - resized_width % factor_width) % factor_width - heights.append((resized_height + pad_height) // patch_size) - widths.append((resized_width + pad_width) // patch_size) + grid_heights.append((resized_height + pad_height) // patch_size) + grid_widths.append((resized_width + pad_width) // patch_size) else: - heights = [constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_HEIGHT] - widths = [constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_WIDTH] + grid_heights = [constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_HEIGHT] + grid_widths = [constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_WIDTH] - num_frames = normalize_sized_list(1 if num_frames is None else num_frames, len(heights), "num_frames") - explicit_num_patches = normalize_sized_list(num_patches, len(heights), "num_patches") - explicit_num_image_tokens = normalize_sized_list(num_image_tokens, len(heights), "num_image_tokens") + num_frames = normalize_sized_list(1 if num_frames is None else num_frames, len(grid_heights), "num_frames") vision = [] max_num_image_tokens = 0 - for index, (height, width, frames) in enumerate(zip(heights, widths, num_frames)): - if height % kernel_height != 0 or width % kernel_width != 0: + for index, (grid_height, grid_width, frames) in enumerate(zip(grid_heights, grid_widths, num_frames)): + if grid_height % kernel_height != 0 or grid_width % kernel_width != 0: raise ValueError( - f"Kimi image grid h={height}, w={width} must be divisible by merge_kernel_size={merge_kernel_size}." + f"Kimi image grid h={grid_height}, w={grid_width} must be divisible by " + f"merge_kernel_size={merge_kernel_size}." ) - computed_num_patches = height * width * frames - computed_num_image_tokens = (height // kernel_height) * (width // kernel_width) * frames - resolved_num_patches = ( - explicit_num_patches[index] if explicit_num_patches is not None else computed_num_patches - ) - resolved_num_image_tokens = ( - explicit_num_image_tokens[index] if explicit_num_image_tokens is not None else computed_num_image_tokens - ) - max_num_image_tokens = max(max_num_image_tokens, resolved_num_image_tokens) + computed_num_patches = grid_height * grid_width * frames + computed_num_image_tokens = (grid_height // kernel_height) * (grid_width // kernel_width) * frames + max_num_image_tokens = max(max_num_image_tokens, computed_num_image_tokens) vision.append( { - "num_patches": resolved_num_patches, - "h": height, - "w": width, - "num_image_tokens": resolved_num_image_tokens, + "num_patches": computed_num_patches, + "grid_h": grid_height, + "grid_w": grid_width, + "num_image_tokens": computed_num_image_tokens, } ) diff --git a/examples/image_text_to_text/README.md b/examples/image_text_to_text/README.md index 6aac140698..e9504b6b0c 100644 --- a/examples/image_text_to_text/README.md +++ b/examples/image_text_to_text/README.md @@ -112,17 +112,19 @@ Some models have specialized examples demonstrating advanced features: ### Kimi-K2.5 Vision -Kimi-K2.5 uses the `moonshotai/Kimi-K2.5` remote-code model with the `KimiK25ForConditionalGeneration` architecture. The example script loads either the full checkpoint or a smaller layer subset, wraps it with `QEFFAutoModelForImageTextToText`, and compiles the vision-language path with Kimi-specific image-grid specializations. +Kimi-K2.5 uses the `moonshotai/Kimi-K2.5` remote-code model with the `KimiK25ForConditionalGeneration` architecture. The example script loads either the full checkpoint or a smaller layer subset, wraps it with `QEFFAutoModelForImageTextToText`, and compiles the vision-language path with Kimi-specific image-size specializations. ```bash python ../kimi_k2/export_kimi_k25_vision.py \ --model-path "$HF_HUB_CACHE/models--moonshotai--Kimi-K2.5/snapshots/" \ --image-url "https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png" \ + --image-height 512 \ + --image-width 910 \ --prompt "Describe this image." ``` -For custom image sizes, pass matching compile specializations for the processed image grid: `num_patches`, `h`, `w`, and `num_image_tokens`. The defaults in `../kimi_k2/export_kimi_k25_vision.py` demonstrate the expected values for the sample image and keep the language-side image embedding shape bounded to the actual image token count. +For custom image sizes, pass only the input image pixel dimensions with `--image-height` and `--image-width`. Kimi-K2.5 compile derives the internal grid height, grid width, patch count, and image-token count from those pixel dimensions. For reranker examples, see [../reranker/](../reranker/). diff --git a/examples/kimi_k2/export_kimi_k25_vision.py b/examples/kimi_k2/export_kimi_k25_vision.py index 553d0afd66..09bcd6b948 100644 --- a/examples/kimi_k2/export_kimi_k25_vision.py +++ b/examples/kimi_k2/export_kimi_k25_vision.py @@ -21,7 +21,6 @@ NUM_TEXT_LAYERS, NUM_VISION_LAYERS, ensure_torch_fx_import_compatibility, - get_kimi_k25_num_image_tokens, load_kimi_k25_class, load_layer_subset_model, parse_expert_ids, @@ -47,9 +46,23 @@ def parse_args(): type=str, default="https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png", ) + parser.add_argument( + "--image-height", + type=int, + default=None, + help="Image height in pixels for Kimi-K2.5 vision compile. Defaults to the loaded image height.", + ) + parser.add_argument( + "--image-width", + type=int, + default=None, + help="Image width in pixels for Kimi-K2.5 vision compile. Defaults to the loaded image width.", + ) parser.add_argument("--prompt", type=str, default="Describe this image.") - parser.add_argument("--test", action="store_true", help="Validate ONNX output matches PyTorch image-only forward.") - return parser.parse_args() + args = parser.parse_args() + if (args.image_height is None) != (args.image_width is None): + parser.error("--image-height and --image-width must be provided together.") + return args def main(): @@ -138,6 +151,8 @@ def main(): ## VISION + TEXT MODE ## image = Image.open(BytesIO(requests.get(args.image_url).content)).convert("RGB") + if args.image_height is not None: + image = image.resize((args.image_width, args.image_height)) messages = [ { @@ -158,11 +173,6 @@ def main(): inputs["pixel_values"] = inputs["pixel_values"].to(qeff_model.model.config.torch_dtype) - num_patches = int(inputs["pixel_values"].shape[0]) - h = int(inputs["grid_thws"][0, 1].item()) - w = int(inputs["grid_thws"][0, 2].item()) - num_image_tokens = (get_kimi_k25_num_image_tokens(config, inputs["grid_thws"]),) - qeff_model.compile( qaic_config=qaic_config, prefill_seq_len=1, @@ -173,10 +183,8 @@ def main(): mxint8_kv_cache=False, aic_enable_depth_first=False, mos=1, - num_patches=num_patches, - h=h, - w=w, - num_image_tokens=num_image_tokens, + image_height=image.height, + image_width=image.width, ) output = qeff_model.generate(inputs=inputs, generation_len=10) diff --git a/tests/transformers/models/image_text_to_text/test_continuous_batching.py b/tests/transformers/models/image_text_to_text/test_continuous_batching.py index 0e52f87dda..3bd0198aec 100644 --- a/tests/transformers/models/image_text_to_text/test_continuous_batching.py +++ b/tests/transformers/models/image_text_to_text/test_continuous_batching.py @@ -26,7 +26,6 @@ from QEfficient import QEFFAutoModelForCausalLM, QEFFAutoModelForImageTextToText from QEfficient.utils.load_kimi_utils import ( - get_kimi_k25_num_image_tokens, get_kimi_k25_test_config, is_kimi_k25, load_kimi_k25_layer_subset_model, @@ -220,42 +219,22 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100_CB( image_height = None image_width = None image_urls = [image_urls[0]] * len(queries) - num_patches = [] - image_heights = [] - image_widths = [] - num_image_tokens = [] for img_url in image_urls: image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB") images.append(image) - for image, query in zip(images, queries): - conversation = [ - { - "role": "user", - "content": [ - {"type": "image_url", "image_url": image}, - {"type": "text", "text": query}, - ], - }, - ] - inputs = processor(messages=conversation, add_generation_prompt=True, tokenize=False, return_tensors="pt") - num_patches.append(int(inputs["pixel_values"].shape[0])) - image_heights.append(int(inputs["grid_thws"][0, 1].item())) - image_widths.append(int(inputs["grid_thws"][0, 2].item())) - num_image_tokens.append(get_kimi_k25_num_image_tokens(config, inputs["grid_thws"])) - image_list = [images[0]] * full_batch_size prompt_list = [queries[0]] * full_batch_size pytorch_hf_tokens = run_kimi_k25_hf_model_on_pytorch_CB( copy.deepcopy(model_hf), processor, image_list, prompt_list, max_gen_len ) + image_height = images[0].height + image_width = images[0].width compile_kwargs.update( { "prefill_seq_len": 1, - "num_patches": num_patches[0], - "h": image_heights[0], - "w": image_widths[0], - "num_image_tokens": num_image_tokens[0], + "image_height": image_height, + "image_width": image_width, } ) else: diff --git a/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py b/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py index 869fe9b11b..12bb3fe06e 100644 --- a/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py +++ b/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py @@ -29,7 +29,6 @@ from QEfficient.utils._utils import create_json from QEfficient.utils.constants import QnnConstants from QEfficient.utils.load_kimi_utils import ( - get_kimi_k25_num_image_tokens, is_kimi_k25, load_kimi_k25_layer_subset_model, run_kimi_k25_hf_model_on_pytorch, @@ -271,10 +270,8 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100( compile_kwargs.update( { "prefill_seq_len": 1, - "num_patches": int(inputs["pixel_values"].shape[0]), - "h": int(inputs["grid_thws"][0, 1].item()), - "w": int(inputs["grid_thws"][0, 2].item()), - "num_image_tokens": get_kimi_k25_num_image_tokens(config, inputs["grid_thws"]), + "image_height": image.height, + "image_width": image.width, } ) diff --git a/tests/unit_test/models/test_model_quickcheck.py b/tests/unit_test/models/test_model_quickcheck.py index 38c78cd009..cd8813c343 100644 --- a/tests/unit_test/models/test_model_quickcheck.py +++ b/tests/unit_test/models/test_model_quickcheck.py @@ -3554,7 +3554,7 @@ def test_layerwise_export_default_names_unchanged(tmp_path): def test_kimi_k25_get_specializations_supports_multi_resolution_grid_sizes(): - """Kimi K2.5 accepts list-valued raw-pixel and patch-grid sizes for multi-resolution specs.""" + """Kimi K2.5 accepts list-valued image sizes for multi-resolution specs.""" from types import SimpleNamespace from QEfficient.transformers.models.kimi_k25.modeling_kimi_k25 import QEffKimiK25ForConditionalGeneration @@ -3566,28 +3566,35 @@ def test_kimi_k25_get_specializations_supports_multi_resolution_grid_sizes(): batch_size=1, prefill_seq_len=64, ctx_len=4096, - height=[512, 448], - width=[910, 448], + image_height=[512, 448], + image_width=[910, 448], num_frames=[1, 1], kv_offload=True, ) assert specs["vision"] == [ - {"num_patches": 2508, "h": 38, "w": 66, "num_image_tokens": 627}, - {"num_patches": 1024, "h": 32, "w": 32, "num_image_tokens": 256}, + {"num_patches": 2508, "grid_h": 38, "grid_w": 66, "num_image_tokens": 627}, + {"num_patches": 1024, "grid_h": 32, "grid_w": 32, "num_image_tokens": 256}, ] assert all(spec["num_image_tokens"] == 627 for spec in specs["lang"]) - specs, _ = model.get_specializations( - batch_size=1, - prefill_seq_len=64, - ctx_len=4096, - h=[30, 32], - w=[80, 64], - num_frames=[1, 2], - kv_offload=True, - ) - assert specs["vision"] == [ - {"num_patches": 2400, "h": 30, "w": 80, "num_image_tokens": 600}, - {"num_patches": 4096, "h": 32, "w": 64, "num_image_tokens": 1024}, - ] - assert all(spec["num_image_tokens"] == 1024 for spec in specs["lang"]) + with pytest.raises(ValueError, match="image_height and image_width"): + model.get_specializations( + batch_size=1, + prefill_seq_len=64, + ctx_len=4096, + h=[30, 32], + w=[80, 64], + num_frames=[1, 2], + kv_offload=True, + ) + + with pytest.raises(ValueError, match="num_patches"): + model.get_specializations( + batch_size=1, + prefill_seq_len=64, + ctx_len=4096, + image_height=512, + image_width=910, + num_patches=2508, + kv_offload=True, + ) From c502602e0339cf4015090191eb56d3ce01d21ab0 Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Sun, 26 Jul 2026 17:53:36 +0530 Subject: [PATCH 31/33] remove redundant runtime changes Signed-off-by: Mamta Singh --- .../generation/text_generation_inference.py | 7 ------ QEfficient/generation/vlm_generation.py | 7 +----- .../models/kimi_k25/modeling_kimi_k25.py | 22 ++++++++++++------- README.md | 1 + .../test_continuous_batching.py | 2 -- 5 files changed, 16 insertions(+), 23 deletions(-) diff --git a/QEfficient/generation/text_generation_inference.py b/QEfficient/generation/text_generation_inference.py index 4e4b71366e..17c992064c 100755 --- a/QEfficient/generation/text_generation_inference.py +++ b/QEfficient/generation/text_generation_inference.py @@ -685,7 +685,6 @@ def initialize_decode_inputs(self, num_prompts, execution_batch_size, max_gen_le self.generated_ids = np.full((num_prompts, max_gen_length), self.tokenizer.pad_token_id) self.decode_input_ids = np.zeros((execution_batch_size, 1), np.int64) self.decode_pos_ids = np.zeros((execution_batch_size, 1), np.int64) - self.decode_image_idx = np.zeros((execution_batch_size, 1), np.int64) self.generation_len = np.zeros((execution_batch_size, 1), np.int64) def initialize_lora_id_mapping(self, prompt_to_lora_id_mapping): @@ -723,8 +722,6 @@ def update_decode_input(self, outputs, position_ids, generation_len, decode_batc decode_batch = decode_batch_id if decode_batch_id is not None else slice(None) self.decode_input_ids[decode_batch] = next_token_id self.decode_pos_ids[decode_batch] = position_ids - if "image_idx_output" in outputs: - self.decode_image_idx[decode_batch] = outputs["image_idx_output"] self.generated_ids[decode_batch, 0] = next_token_id.squeeze() self.generation_len[decode_batch] = generation_len return next_token_id @@ -767,10 +764,6 @@ def _set_output_buffers(self, batch_size: int = 1, sequence_length: int = 1): logits_out_placeholder = np.zeros((batch_size, sequence_length, self._vocab_size), dtype=np.float32) self._session.set_buffers({"logits": logits_out_placeholder}) - if "image_idx_output" in getattr(self._session, "binding_index_map", {}): - image_idx_out_placeholder = np.zeros((batch_size, 1), dtype=np.int64) - self._session.set_buffers({"image_idx_output": image_idx_out_placeholder}) - def run_prefill(self, prompt, generation_len, prefill_logit_bs=1, decode_batch_id=None): """ Runs prefill for a given prompt and generation length. diff --git a/QEfficient/generation/vlm_generation.py b/QEfficient/generation/vlm_generation.py index 013452fbc1..df731f951a 100755 --- a/QEfficient/generation/vlm_generation.py +++ b/QEfficient/generation/vlm_generation.py @@ -989,12 +989,7 @@ def prepare_decode_inputs(self): if "image_idx" in getattr(self._session, "binding_index_map", {}): idx = self._session.binding_index_map["image_idx"] dims = tuple(self._session.bindings[idx].dims) - if dims == tuple(self.decode_image_idx.shape): - decode_inputs["image_idx"] = self.decode_image_idx - elif dims[0] == 1: - decode_inputs["image_idx"] = self.decode_image_idx[:1] - else: - decode_inputs["image_idx"] = np.broadcast_to(self.decode_image_idx[:1], dims).copy() + decode_inputs["image_idx"] = np.zeros(dims, dtype=np.int64) else: decode_inputs["image_idx"] = np.array([[0]], dtype=np.int64) except Exception: diff --git a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py index cbfc86d228..abdb1a5701 100644 --- a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py +++ b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py @@ -517,13 +517,23 @@ def forward( inputs_embeds = merged_inputs_embeds attention_mask = merged_attention_mask + merged_image_tokens = ( + torch._shape_as_tensor(vision_embeds_for_state)[:1] + .view(1, 1) + .to(device=image_idx.device, dtype=torch.int64) + ) + default_image_idx = torch.clamp(merged_image_tokens - 1, min=0) + input_batch = torch._shape_as_tensor(input_ids)[:1].view(1, 1).to(device=image_idx.device) + image_idx_batch = torch._shape_as_tensor(image_idx)[:1].view(1, 1).to(device=image_idx.device) + use_default_image_idx = torch.logical_and(torch.logical_not(selected_any), input_batch != image_idx_batch) + effective_image_idx = torch.where(use_default_image_idx, default_image_idx, image_idx) if position_ids is None: position_ids = merged_position_ids else: # Preserve caller-provided absolute position offset (needed for # chunked prefill/decode parity) while using merged sequence # positions for image-expanded tokens. - position_offset = position_ids[:, :1] + image_idx.to( + position_offset = position_ids[:, :1] + effective_image_idx.to( device=position_ids.device, dtype=position_ids.dtype ) position_ids = torch.where( @@ -532,13 +542,8 @@ def forward( torch.full_like(merged_position_ids, -1), ) - merged_image_tokens = ( - torch._shape_as_tensor(vision_embeds_for_state)[:1] - .view(1, 1) - .to(device=image_idx.device, dtype=torch.int64) - ) image_position_delta = torch.clamp(merged_image_tokens - selected_image_tokens, min=0) - image_idx = image_idx + selected_any.to(torch.int64) * image_position_delta + image_idx = effective_image_idx + selected_any.to(torch.int64) * image_position_delta if position_ids is None and attention_mask is not None: position_ids = attention_mask.long().cumsum(-1) - 1 @@ -568,7 +573,8 @@ def forward( output_kvs = getattr(outputs, "compressed_kvs", None) else: output_kvs = getattr(outputs, "past_key_values", None) - return logits, vision_embeds_for_state, image_idx, output_kvs + image_idx_output = image_idx[:1, :1] if image_idx is not None else image_idx + return logits, vision_embeds_for_state, image_idx_output, output_kvs # ref https://github.com/huggingface/transformers/blob/78b2929c0554b79e0489b451ce4ece14d265ead2/src/transformers/models/llava/modeling_llava.py#L240 diff --git a/README.md b/README.md index 764feb54fb..81b6bee2b5 100644 --- a/README.md +++ b/README.md @@ -12,6 +12,7 @@ --- *Latest news* :fire:
+ - [07/2026] Added support for Kimi-K2.5 vision-language model [moonshotai/Kimi-K2.5](https://huggingface.co/moonshotai/Kimi-K2.5) via `QEFFAutoModelForImageTextToText`. - [07/2026] Added `dynamo` flag to `QEFFAutoModelForCausalLM.export()` to support `torch.onnx.export` dynamo-based ONNX export for CausalLM models - [06/2026] Added support for Gemma4 models, [google/gemma-4-E2B-it](https://huggingface.co/google/gemma-4-E2B), [google/gemma-4-26B-A4B-it](https://huggingface.co/google/gemma-4-26B-A4B-it) diff --git a/tests/transformers/models/image_text_to_text/test_continuous_batching.py b/tests/transformers/models/image_text_to_text/test_continuous_batching.py index 3bd0198aec..133bfa475d 100644 --- a/tests/transformers/models/image_text_to_text/test_continuous_batching.py +++ b/tests/transformers/models/image_text_to_text/test_continuous_batching.py @@ -216,8 +216,6 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100_CB( ) compile_kwargs["img_size"] = img_size elif is_kimi_k25(model_name): - image_height = None - image_width = None image_urls = [image_urls[0]] * len(queries) for img_url in image_urls: image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB") From c92fdc688b1f90ca1d7c508695df8c93d5c1c2f9 Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Mon, 27 Jul 2026 17:33:38 +0530 Subject: [PATCH 32/33] address review comments Signed-off-by: Mamta Singh --- QEfficient/__init__.py | 21 + .../models/kimi_k25/configuration_kimi_k25.py | 126 ------ .../transformers/models/modeling_auto.py | 1 + examples/kimi_k2/export_kimi_k25_vision.py | 391 +++++++++++++----- .../test_continuous_batching.py | 14 +- .../test_image_text_to_text_models.py | 10 +- .../test_kimi_k25_disagg.py | 4 +- .../unit_test/models/test_model_quickcheck.py | 70 ++++ .../utils/load_kimi_utils.py | 23 -- 9 files changed, 401 insertions(+), 259 deletions(-) delete mode 100644 QEfficient/transformers/models/kimi_k25/configuration_kimi_k25.py rename {QEfficient => tests}/utils/load_kimi_utils.py (96%) diff --git a/QEfficient/__init__.py b/QEfficient/__init__.py index 393ee7d4ad..fb85183144 100755 --- a/QEfficient/__init__.py +++ b/QEfficient/__init__.py @@ -19,6 +19,7 @@ import warnings # noqa: I001 import transformers import transformers.utils as transformers_utils +from transformers.utils import import_utils as hf_import_utils try: from transformers import HybridCache as _TransformersHybridCache # noqa: F401 @@ -117,3 +118,23 @@ def check_qaic_sdk(): if not check_qaic_sdk(): logger.warning("QAIC SDK is not installed, eager mode features won't be available!") + + +def ensure_torch_fx_import_compatibility(): + if hasattr(hf_import_utils, "is_torch_fx_available"): + return + + def _is_torch_fx_available() -> bool: + if not hf_import_utils.is_torch_available(): + return False + try: + import torch.fx # noqa: F401 + + return True + except Exception: + return False + + hf_import_utils.is_torch_fx_available = _is_torch_fx_available + + +ensure_torch_fx_import_compatibility() diff --git a/QEfficient/transformers/models/kimi_k25/configuration_kimi_k25.py b/QEfficient/transformers/models/kimi_k25/configuration_kimi_k25.py deleted file mode 100644 index 2f6d3b2467..0000000000 --- a/QEfficient/transformers/models/kimi_k25/configuration_kimi_k25.py +++ /dev/null @@ -1,126 +0,0 @@ -# ----------------------------------------------------------------------------- -# -# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. -# SPDX-License-Identifier: BSD-3-Clause -# -# ----------------------------------------------------------------------------- - -from transformers.configuration_utils import PretrainedConfig - -from QEfficient.transformers.models.deepseek_v3.configuration_deepseek import DeepseekV3Config - - -class KimiK25VisionConfig(PretrainedConfig): - def __init__( - self, - patch_size: int = 14, - init_pos_emb_height: int = 64, - init_pos_emb_width: int = 64, - init_pos_emb_time: int = 4, - pos_emb_type: str = "divided_fixed", - vt_num_attention_heads: int = 16, - vt_num_hidden_layers: int = 27, - vt_hidden_size: int = 1152, - vt_intermediate_size: int = 4304, - merge_kernel_size: tuple = (2, 2), - video_attn_type: str = "spatial_temporal", - merge_type: str = "sd2_tpool", - _attn_implementation: str = "flash_attention_2", - # MM Projector parameters - mm_projector_type: str = "patchmerger", - mm_hidden_size: int | None = None, - projector_hidden_act: str = "gelu", - projector_ln_eps: float = 1e-5, - # Other parameters - ignore_index: int = -100, - media_placeholder_token_id: int = 163605, - pad_token_id: int = 0, - use_unified_vision_chunk: bool = True, - video_placeholder="<|kimi_k25_video_placeholder|>", - text_hidden_size=7168, - **vision_config_kwargs, - ): - self.patch_size = patch_size - self.init_pos_emb_height = init_pos_emb_height - self.init_pos_emb_width = init_pos_emb_width - self.init_pos_emb_time = init_pos_emb_time - self.pos_emb_type = pos_emb_type - self.vt_num_attention_heads = vt_num_attention_heads - self.vt_num_hidden_layers = vt_num_hidden_layers - self.vt_hidden_size = vt_hidden_size - self.vt_intermediate_size = vt_intermediate_size - self.merge_kernel_size = merge_kernel_size - self.video_attn_type = video_attn_type - self.merge_type = merge_type - self._attn_implementation = _attn_implementation - - # MM Projector config - self.mm_projector_type = mm_projector_type - self.mm_hidden_size = mm_hidden_size if mm_hidden_size is not None else vt_hidden_size - self.projector_hidden_act = projector_hidden_act - self.projector_ln_eps = projector_ln_eps - self.text_hidden_size = text_hidden_size - - -class KimiK25Config(PretrainedConfig): - """Kimi-K2.5 model configuration. - - Args: - text_config (dict | DeepseekV3Config): Configuration for the text model. - - Vision Tower Parameters (from MoonViT3dConfig): - patch_size (int): Patch size for vision tower. - init_pos_emb_height (int): Initial position embedding height. - init_pos_emb_width (int): Initial position embedding width. - init_pos_emb_time (int): Initial position embedding time dimension. - pos_emb_type (str): Type of position embedding. - vt_num_attention_heads (int): Number of attention heads in vision tower. - vt_num_hidden_layers (int): Number of hidden layers in vision tower. - vt_hidden_size (int): Hidden size of vision tower. - vt_intermediate_size (int): Intermediate size in vision tower FFN. - merge_kernel_size (tuple): Kernel size for patch merging. - video_attn_type (str): Type of video attention. - merge_type (str): Type of merge operation. - _attn_implementation (str): Attention implementation type. - - MM Projector Parameters (from MultiModalProjectorConfig): - mm_projector_type (str): Type of multimodal projector. - mm_hidden_size (int): Hidden size from vision tower (should match vt_hidden_size). - projector_hidden_act (str): Activation function for projector. - projector_ln_eps (float): Layer norm epsilon for projector. - - Other Parameters: - ignore_index (int): The ignore index for the loss function. - media_placeholder_token_id (int): The token ID to use for media placeholders. - pad_token_id (int): The token ID to use for padding. - """ - - model_type = "kimi_k25" - - def __init__( - self, - text_config: dict | DeepseekV3Config = None, - vision_config: dict | KimiK25VisionConfig = None, - # Other parameters - ignore_index: int = -100, - media_placeholder_token_id: int = 163605, - pad_token_id: int = 0, - use_unified_vision_chunk: bool = True, - video_placeholder="<|kimi_k25_video_placeholder|>", - **kwargs, - ): - if isinstance(text_config, dict): - text_config = DeepseekV3Config(**text_config) - if isinstance(vision_config, dict): - vision_config = KimiK25VisionConfig(**vision_config) - self.text_config = text_config - self.vision_config = vision_config - # Other config - self.ignore_index = ignore_index - self.media_placeholder_token_id = media_placeholder_token_id - self.use_unified_vision_chunk = use_unified_vision_chunk - self.video_placeholder = video_placeholder - if getattr(self.text_config, "quantization_config", None) is not None: - self.quantization_config = self.text_config.quantization_config - - super().__init__(pad_token_id=pad_token_id, **kwargs) diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index 245082789a..078ae9f4fd 100755 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -2340,6 +2340,7 @@ def kv_offload_generate( vision_inputs_fp16 = {"pixel_values", "image_masks"} vision_inputs.update({k: vision_inputs[k].astype("float16") for k in vision_inputs_fp16 if k in vision_inputs}) + # Required for KIMI-K25 grid_thws_val = inputs.pop("grid_thws", None) if grid_thws_val is not None: h_val = int(grid_thws_val[0, 1].item()) diff --git a/examples/kimi_k2/export_kimi_k25_vision.py b/examples/kimi_k2/export_kimi_k25_vision.py index 09bcd6b948..65a71168af 100644 --- a/examples/kimi_k2/export_kimi_k25_vision.py +++ b/examples/kimi_k2/export_kimi_k25_vision.py @@ -6,31 +6,48 @@ # ----------------------------------------------------------------------------- import argparse +import copy +import importlib.util from io import BytesIO from pathlib import Path +from time import perf_counter +import numpy as np import requests import torch from PIL import Image from QEfficient import QEFFAutoModelForImageTextToText -from QEfficient.utils.load_kimi_utils import ( - DEFAULT_MODEL_PATH, - LOADED_EXPERT_IDS, - NUM_EXPERTS_PER_TOKEN, - NUM_TEXT_LAYERS, - NUM_VISION_LAYERS, - ensure_torch_fx_import_compatibility, - load_kimi_k25_class, - load_layer_subset_model, - parse_expert_ids, - prepare_config, - set_deterministic, +from QEfficient.generation.cloud_infer import QAICInferenceSession + +LOAD_KIMI_UTILS_PATH = Path(__file__).resolve().parents[2] / "tests" / "utils" / "load_kimi_utils.py" +_load_kimi_spec = importlib.util.spec_from_file_location("load_kimi_utils", LOAD_KIMI_UTILS_PATH) +if _load_kimi_spec is None or _load_kimi_spec.loader is None: + raise ImportError(f"Unable to load Kimi helpers from {LOAD_KIMI_UTILS_PATH}") +load_kimi_utils = importlib.util.module_from_spec(_load_kimi_spec) +_load_kimi_spec.loader.exec_module(load_kimi_utils) + +LOADED_EXPERT_IDS = load_kimi_utils.LOADED_EXPERT_IDS +NUM_EXPERTS_PER_TOKEN = load_kimi_utils.NUM_EXPERTS_PER_TOKEN +NUM_TEXT_LAYERS = load_kimi_utils.NUM_TEXT_LAYERS +NUM_VISION_LAYERS = load_kimi_utils.NUM_VISION_LAYERS +load_kimi_k25_class = load_kimi_utils.load_kimi_k25_class +load_layer_subset_model = load_kimi_utils.load_layer_subset_model +parse_expert_ids = load_kimi_utils.parse_expert_ids +prepare_config = load_kimi_utils.prepare_config +set_deterministic = load_kimi_utils.set_deterministic + +DEFAULT_MODEL_PATH = Path( + "/home/huggingface_hub/models--moonshotai--Kimi-K2.5/snapshots/4d01dfe0332d63057c186e0b262165819efb6611" ) +PREFILL_SEQ_LEN = 512 +CTX_LEN = 2048 +BS = 1 +GENERATION_LEN = 10 def parse_args(): - parser = argparse.ArgumentParser(description="Load Kimi K2.5 with runtime compatibility for transformers==5.5.4.") + parser = argparse.ArgumentParser(description="Run Kimi K2.5 vision disaggregated vision -> prefill -> decode flow.") parser.add_argument("--model-path", type=Path, default=DEFAULT_MODEL_PATH) parser.add_argument( "--full-model", @@ -59,17 +76,221 @@ def parse_args(): help="Image width in pixels for Kimi-K2.5 vision compile. Defaults to the loaded image width.", ) parser.add_argument("--prompt", type=str, default="Describe this image.") + parser.add_argument("--prefill-seq-len", type=int, default=PREFILL_SEQ_LEN) + parser.add_argument("--ctx-len", type=int, default=CTX_LEN) + parser.add_argument("--generation-len", type=int, default=GENERATION_LEN) + parser.add_argument("--num-cores", type=int, default=16) + parser.add_argument("--vision-num-devices", type=int, default=1) + parser.add_argument("--lang-num-devices", type=int, default=4) + parser.add_argument("--mxfp6-matmul", action="store_true") + parser.add_argument("--mxint8-kv-cache", action="store_true") args = parser.parse_args() if (args.image_height is None) != (args.image_width is None): parser.error("--image-height and --image-width must be provided together.") return args -def main(): - args = parse_args() +def _clone_inputs(inputs): + return {key: (value.clone() if torch.is_tensor(value) else copy.deepcopy(value)) for key, value in inputs.items()} + + +def _numpy(value): + if torch.is_tensor(value): + return value.detach().cpu().numpy() + return np.asarray(value) + + +def _session_input_names(session: QAICInferenceSession) -> set[str]: + input_names = set(session.input_names) + input_names.update(name.rsplit("/", 1)[-1] for name in session.input_names) + return input_names + + +def _cast_for_session(session: QAICInferenceSession, name: str, value: np.ndarray) -> np.ndarray: + binding_index = session.binding_index_map.get(name) + if binding_index is None: + return value + dtype = session.aic_to_np_dtype_mapping[session.bindings[binding_index].type] + return value.astype(dtype, copy=False) + + +def _filter_session_inputs(session: QAICInferenceSession, inputs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: + input_names = _session_input_names(session) + return {name: _cast_for_session(session, name, value) for name, value in inputs.items() if name in input_names} + + +def _resolve_qpc_path(qpc_paths, key: str): + if isinstance(qpc_paths, dict): + qpc_path = qpc_paths.get(key) + if qpc_path is None: + raise KeyError(f"Missing {key!r} in compile output keys: {list(qpc_paths.keys())}") + return qpc_path + return qpc_paths + + +def _update_retained_states(target_inputs: dict[str, np.ndarray], source_outputs: dict[str, np.ndarray]): + for output_name, value in source_outputs.items(): + output_basename = output_name.rsplit("/", 1)[-1] + if output_basename.endswith("_RetainedState"): + target_inputs[output_basename.removesuffix("_RetainedState")] = value + + +def _get_next_token_ids(logits: np.ndarray) -> np.ndarray: + logits = np.asarray(logits) + return logits[:, -1, :].argmax(axis=-1).astype(np.int64).reshape(BS, 1) + + +def _compile_disagg_qpcs(qeff_model: QEFFAutoModelForImageTextToText, args, image: Image.Image): + qaic_config = {"mla_absorption": {"cache_compressed": True, "absorption": False, "online": False}} + common_compile_kwargs = { + "qaic_config": qaic_config, + "batch_size": BS, + "ctx_len": args.ctx_len, + "image_height": image.height, + "image_width": image.width, + "num_cores": args.num_cores, + "mxfp6_matmul": args.mxfp6_matmul, + "mxint8_kv_cache": args.mxint8_kv_cache, + "split_model_io": True, + "mos": 1, + "aic_enable_depth_first": True, + "use_onnx_subfunctions": True, + "layerwise": False, + } + + print("Compiling vision QPC...") + vision_qpc_path = qeff_model.compile( + prefill_seq_len=args.prefill_seq_len, + skip_vision=False, + skip_lang=True, + num_devices=args.vision_num_devices, + **common_compile_kwargs, + ) + + print("Compiling prefill QPC...") + prefill_qpc_path = qeff_model.compile( + prefill_seq_len=args.prefill_seq_len, + prefill_only=True, + enable_chunking=True, + skip_vision=True, + skip_lang=False, + num_devices=args.lang_num_devices, + **common_compile_kwargs, + ) + + print("Compiling decode QPC...") + decode_qpc_path = qeff_model.compile( + prefill_seq_len=1, + prefill_only=False, + skip_vision=True, + skip_lang=False, + num_devices=args.lang_num_devices, + **common_compile_kwargs, + ) + + return vision_qpc_path, prefill_qpc_path, decode_qpc_path + +def _run_disagg_generation( + inputs: dict[str, torch.Tensor], + vision_session: QAICInferenceSession, + prefill_session: QAICInferenceSession, + decode_session: QAICInferenceSession, + *, + prefill_seq_len: int, + generation_len: int, +) -> np.ndarray: + inputs = {name: _numpy(value) for name, value in _clone_inputs(inputs).items()} + input_ids_length = inputs["input_ids"].shape[1] + num_chunks = -(input_ids_length // -prefill_seq_len) + padded_len = num_chunks * prefill_seq_len + + inputs["input_ids"] = np.pad( + inputs["input_ids"], + ((0, 0), (0, padded_len - input_ids_length)), + constant_values=1, + ) + inputs["attention_mask"] = np.pad( + inputs["attention_mask"], + ((0, 0), (0, padded_len - input_ids_length)), + constant_values=0, + ) + + grid_thws = inputs.pop("grid_thws").astype(np.int64) + h = int(grid_thws[0, 1]) + w = int(grid_thws[0, 2]) + vision_inputs = { + "pixel_values": inputs["pixel_values"], + "h_shape": np.ones((h,), dtype=np.int64), + "w_shape": np.ones((w,), dtype=np.int64), + } + + vision_start = perf_counter() + print("Running vision QPC...") + vision_outputs = vision_session.run(_filter_session_inputs(vision_session, vision_inputs)) + vision_session.deactivate() + vision_time = perf_counter() - vision_start + + vision_embeds = vision_outputs.get("vision_embeds") + if vision_embeds is None: + raise RuntimeError(f"Vision QPC did not return vision_embeds. Outputs: {vision_outputs.keys()}") + + lang_inputs = { + "input_ids": inputs["input_ids"].astype(np.int64), + "position_ids": np.where(inputs["attention_mask"] > 0, np.arange(padded_len), -1).astype(np.int64), + "vision_embeds": vision_embeds, + "image_idx": np.zeros((BS, 1), dtype=np.int64), + } + + prefill_start = perf_counter() + print("Running prefill QPC...") + prefill_session.set_buffers(vision_outputs) + chunk_inputs = lang_inputs.copy() + prefill_outputs = None + for chunk_idx in range(num_chunks): + start = chunk_idx * prefill_seq_len + end = (chunk_idx + 1) * prefill_seq_len + chunk_inputs["input_ids"] = lang_inputs["input_ids"][:, start:end] + chunk_inputs["position_ids"] = lang_inputs["position_ids"][:, start:end] + prefill_outputs = prefill_session.run(_filter_session_inputs(prefill_session, chunk_inputs)) + _update_retained_states(chunk_inputs, prefill_outputs) + if "image_idx_output" in prefill_outputs: + chunk_inputs["image_idx"] = prefill_outputs["image_idx_output"].astype(np.int64) + + prefill_session.deactivate() + if prefill_outputs is None: + raise RuntimeError("QAIC prefill did not execute.") + prefill_time = perf_counter() - prefill_start + vision_time + print(f"Prefill time, including vision: {prefill_time:.2f} secs") + + generated_ids = [_get_next_token_ids(prefill_outputs["logits"])] + decode_inputs = { + "input_ids": generated_ids[-1], + "position_ids": np.max(lang_inputs["position_ids"], axis=-1, keepdims=True).astype(np.int64) + 1, + "vision_embeds": chunk_inputs.get("vision_embeds", vision_embeds), + "image_idx": chunk_inputs.get("image_idx", np.zeros((BS, 1), dtype=np.int64)), + } + _update_retained_states(decode_inputs, prefill_outputs) + + print("Running decode QPC...") + decode_start = perf_counter() + for _ in range(1, generation_len): + decode_outputs = decode_session.run(_filter_session_inputs(decode_session, decode_inputs)) + generated_ids.append(_get_next_token_ids(decode_outputs["logits"])) + decode_inputs["input_ids"] = generated_ids[-1] + decode_inputs["position_ids"] = decode_inputs["position_ids"] + 1 + if "image_idx_output" in decode_outputs: + decode_inputs["image_idx"] = decode_outputs["image_idx_output"].astype(np.int64) + _update_retained_states(decode_inputs, decode_outputs) + + decode_time = perf_counter() - decode_start + if generation_len > 1: + print(f"Decode tok/sec: {(generation_len - 1) / decode_time:.2f}") + return np.concatenate(generated_ids, axis=1) + + +def _load_model(args): set_deterministic(1234) - ensure_torch_fx_import_compatibility() config = prepare_config(args.model_path) kimi_cls = load_kimi_k25_class(args.model_path) @@ -102,96 +323,76 @@ def main(): else: raise ValueError("Pass both --num-vision-layers and --num-text-layers to load a layer subset.") - qaic_config = {"mla_absorption": {"cache_compressed": True, "absorption": False, "online": False}} - - qeff_model = QEFFAutoModelForImageTextToText(model, qaic_config=qaic_config) + model.vision_tower.patch_embed.pos_emb.interpolation_mode = "bilinear" + return model.eval().to("cpu"), tokenizer, processor - skip_vision = False - if skip_vision: - ## TEXT-ONLY MODE ## +def _prepare_inputs(processor, args): + image = Image.open(BytesIO(requests.get(args.image_url, timeout=30).content)).convert("RGB") + if args.image_height is not None: + image = image.resize((args.image_width, args.image_height)) - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": args.prompt}, - ], - }, - ] + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": image}, + {"type": "text", "text": args.prompt}, + ], + }, + ] + inputs = processor( + messages=messages, + add_generation_prompt=True, + tokenize=False, + return_tensors="pt", + ) + inputs = {name: (value.to("cpu") if torch.is_tensor(value) else value) for name, value in inputs.items()} + return image, inputs - inputs = processor.apply_chat_template( - messages, - add_generation_prompt=True, - tokenize=True, - return_dict=True, - return_tensors="pt", - ) - qeff_model.compile( - qaic_config=qaic_config, - prefill_seq_len=1, - ctx_len=1024, - num_cores=16, - num_devices=1, - mxfp6_matmul=False, - mxint8_kv_cache=False, - aic_enable_depth_first=False, - skip_vision=True, # Skip vision encoder for text-only inference - mos=1, - ) +def main(): + args = parse_args() + model, tokenizer, processor = _load_model(args) + qaic_config = {"mla_absorption": {"cache_compressed": True, "absorption": False, "online": False}} + qeff_model = QEFFAutoModelForImageTextToText( + model, + kv_offload=True, + config=model.config, + torch_dtype=torch.float32, + qaic_config=qaic_config, + layerwise=False, + ) - output = qeff_model.generate(inputs=inputs, device_ids=[0], generation_len=10) + image, inputs = _prepare_inputs(processor, args) + inputs["pixel_values"] = inputs["pixel_values"].to(qeff_model.model.config.torch_dtype) - print(output.generated_ids) - print(tokenizer.batch_decode(output.generated_ids)) - print(output) + vision_qpc_path, prefill_qpc_path, decode_qpc_path = _compile_disagg_qpcs(qeff_model, args, image) + print(f"Vision QPC path: {vision_qpc_path}") + print(f"Prefill QPC path: {prefill_qpc_path}") + print(f"Decode QPC path: {decode_qpc_path}") - else: - ## VISION + TEXT MODE ## - - image = Image.open(BytesIO(requests.get(args.image_url).content)).convert("RGB") - if args.image_height is not None: - image = image.resize((args.image_width, args.image_height)) - - messages = [ - { - "role": "user", - "content": [ - {"type": "image_url", "image_url": image}, - {"type": "text", "text": args.prompt}, - ], - }, - ] - - inputs = processor( - messages=messages, - add_generation_prompt=True, - tokenize=False, - return_tensors="pt", - ) + sessions = [] + try: + vision_session = QAICInferenceSession(_resolve_qpc_path(vision_qpc_path, "vision_qpc_path")) + prefill_session = QAICInferenceSession(_resolve_qpc_path(prefill_qpc_path, "lang_prefill_qpc_path")) + decode_session = QAICInferenceSession(_resolve_qpc_path(decode_qpc_path, "lang_decode_qpc_path")) + sessions.extend([vision_session, prefill_session, decode_session]) - inputs["pixel_values"] = inputs["pixel_values"].to(qeff_model.model.config.torch_dtype) - - qeff_model.compile( - qaic_config=qaic_config, - prefill_seq_len=1, - ctx_len=1024, - num_cores=16, - num_devices=1, - mxfp6_matmul=False, - mxint8_kv_cache=False, - aic_enable_depth_first=False, - mos=1, - image_height=image.height, - image_width=image.width, + generated_ids = _run_disagg_generation( + inputs, + vision_session, + prefill_session, + decode_session, + prefill_seq_len=args.prefill_seq_len, + generation_len=args.generation_len, ) + finally: + for session in sessions: + session.deactivate() - output = qeff_model.generate(inputs=inputs, generation_len=10) - - print(output.generated_ids) - print(tokenizer.batch_decode(output.generated_ids)) - print(output) + print(generated_ids) + print(tokenizer.batch_decode(torch.as_tensor(generated_ids), skip_special_tokens=True)) if __name__ == "__main__": diff --git a/tests/transformers/models/image_text_to_text/test_continuous_batching.py b/tests/transformers/models/image_text_to_text/test_continuous_batching.py index 133bfa475d..5f9ab5bff6 100644 --- a/tests/transformers/models/image_text_to_text/test_continuous_batching.py +++ b/tests/transformers/models/image_text_to_text/test_continuous_batching.py @@ -25,13 +25,6 @@ from urllib3.util.retry import Retry from QEfficient import QEFFAutoModelForCausalLM, QEFFAutoModelForImageTextToText -from QEfficient.utils.load_kimi_utils import ( - get_kimi_k25_test_config, - is_kimi_k25, - load_kimi_k25_layer_subset_model, - load_kimi_k25_model_from_config, - run_kimi_k25_hf_model_on_pytorch_CB, -) from QEfficient.utils.run_utils import ApiRunnerInternVL, ApiRunnerMolmo, ApiRunnerVlm from QEfficient.utils.test_utils import ( InternProcessor, @@ -40,6 +33,13 @@ load_vlm_model_from_config, set_num_layers_vlm, ) +from tests.utils.load_kimi_utils import ( + get_kimi_k25_test_config, + is_kimi_k25, + load_kimi_k25_layer_subset_model, + load_kimi_k25_model_from_config, + run_kimi_k25_hf_model_on_pytorch_CB, +) _session = requests.Session() _session.mount("https://", HTTPAdapter(max_retries=Retry(total=3, backoff_factor=1))) diff --git a/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py b/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py index 12bb3fe06e..8e4203c291 100644 --- a/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py +++ b/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py @@ -28,11 +28,6 @@ from QEfficient import QEFFAutoModelForCausalLM, QEFFAutoModelForImageTextToText from QEfficient.utils._utils import create_json from QEfficient.utils.constants import QnnConstants -from QEfficient.utils.load_kimi_utils import ( - is_kimi_k25, - load_kimi_k25_layer_subset_model, - run_kimi_k25_hf_model_on_pytorch, -) from QEfficient.utils.run_utils import ApiRunnerInternVL, ApiRunnerMolmo, ApiRunnerVlm from QEfficient.utils.test_utils import ( InternProcessor, @@ -42,6 +37,11 @@ load_vlm_model_from_config, set_num_layers_vlm, ) +from tests.utils.load_kimi_utils import ( + is_kimi_k25, + load_kimi_k25_layer_subset_model, + run_kimi_k25_hf_model_on_pytorch, +) from ..check_model_results import dump_and_compare_results diff --git a/tests/transformers/models/image_text_to_text/test_kimi_k25_disagg.py b/tests/transformers/models/image_text_to_text/test_kimi_k25_disagg.py index 3bf0c42872..33ce915ae8 100644 --- a/tests/transformers/models/image_text_to_text/test_kimi_k25_disagg.py +++ b/tests/transformers/models/image_text_to_text/test_kimi_k25_disagg.py @@ -17,10 +17,9 @@ from QEfficient import QEFFAutoModelForImageTextToText from QEfficient.generation.cloud_infer import QAICInferenceSession -from QEfficient.utils.load_kimi_utils import ( +from tests.utils.load_kimi_utils import ( LOADED_EXPERT_IDS, NUM_EXPERTS_PER_TOKEN, - ensure_torch_fx_import_compatibility, load_kimi_k25_class, load_layer_subset_model, prepare_config, @@ -120,7 +119,6 @@ def _update_retained_states(target_inputs: dict[str, np.ndarray], source_outputs def _load_kimi_subset_model(): set_deterministic(1234) - ensure_torch_fx_import_compatibility() model_path = resolve_model_path() config = prepare_config(model_path) kimi_cls = load_kimi_k25_class(model_path) diff --git a/tests/unit_test/models/test_model_quickcheck.py b/tests/unit_test/models/test_model_quickcheck.py index cd8813c343..9b599f0d24 100644 --- a/tests/unit_test/models/test_model_quickcheck.py +++ b/tests/unit_test/models/test_model_quickcheck.py @@ -927,6 +927,76 @@ def _assert_qwen_hf_qeff_ort_parity(model_type: str, tmp_path, *, prefill_only: assert np.allclose(qeff_logits, ort_logits, atol=atol, rtol=1e-4) +def _kimi_k25_prompt_inputs(processor): + from PIL import Image + + image = Image.new("RGB", (64, 64), color=(128, 64, 32)) + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": image}, + {"type": "text", "text": "Describe."}, + ], + }, + ] + inputs = processor( + messages=messages, + add_generation_prompt=True, + tokenize=False, + return_tensors="pt", + ) + return {name: (value.to("cpu") if torch.is_tensor(value) else value) for name, value in inputs.items()} + + +def _kimi_k25_qeff_lang_inputs(qeff_model, inputs, vision_embeds): + seq_len = inputs["input_ids"].shape[1] + lang_inputs = deepcopy(qeff_model.model.get_dummy_inputs(kv_offload=True, prefill_seq_len=seq_len)["lang"]) + lang_inputs["input_ids"] = inputs["input_ids"].clone() + lang_inputs["position_ids"] = torch.where( + inputs["attention_mask"] > 0, + torch.arange(seq_len, dtype=torch.int64).view(1, seq_len), + -1, + ) + lang_inputs["vision_embeds"] = vision_embeds + lang_inputs["image_idx"] = torch.zeros((1, 1), dtype=torch.int64) + return lang_inputs + + +@pytest.mark.llm_model +def test_kimi_k25_quickcheck_hf_qeff_vision_logits_parity(): + from tests.utils.load_kimi_utils import load_kimi_k25_layer_subset_model + + model_id = "moonshotai/Kimi-K2.5" + try: + model_hf, _, processor = load_kimi_k25_layer_subset_model(num_vision_layers=1, num_text_layers=1) + except Exception as exc: + _skip_on_model_fetch_error(exc, model_id) + + inputs = _kimi_k25_prompt_inputs(processor) + with torch.no_grad(): + hf_outputs = model_hf(**inputs, use_cache=False, return_dict=True) + hf_logits = hf_outputs.logits[:, -1:, :].detach().float().numpy() + + qeff_model = QEFFAutoModelForImageTextToText( + deepcopy(model_hf), + kv_offload=True, + config=model_hf.config, + torch_dtype=torch.float32, + ) + grid_thws = inputs["grid_thws"].to(torch.int64) + h_shape = torch.ones((int(grid_thws[0, 1].item()),), dtype=torch.int64) + w_shape = torch.ones((int(grid_thws[0, 2].item()),), dtype=torch.int64) + + with torch.no_grad(): + vision_embeds = qeff_model.vision_model.model(inputs["pixel_values"], h_shape, w_shape) + lang_inputs = _kimi_k25_qeff_lang_inputs(qeff_model, inputs, vision_embeds) + qeff_logits = qeff_model.lang_model.model(**lang_inputs)[0].detach().float().numpy() + + assert qeff_logits.shape == hf_logits.shape + assert np.allclose(hf_logits, qeff_logits, atol=1e-4, rtol=1e-4) + + @pytest.mark.llm_model @pytest.mark.parametrize( ("model_type", "model_id"), diff --git a/QEfficient/utils/load_kimi_utils.py b/tests/utils/load_kimi_utils.py similarity index 96% rename from QEfficient/utils/load_kimi_utils.py rename to tests/utils/load_kimi_utils.py index 7cc387a9be..af1f8c259b 100644 --- a/QEfficient/utils/load_kimi_utils.py +++ b/tests/utils/load_kimi_utils.py @@ -22,12 +22,8 @@ from safetensors.torch import save_file from transformers import AutoConfig, AutoProcessor, AutoTokenizer from transformers.dynamic_module_utils import get_class_from_dynamic_module -from transformers.utils import import_utils as hf_import_utils KIMI_K25_MODEL_NAME = "moonshotai/Kimi-K2.5" -DEFAULT_MODEL_PATH = Path( - "/home/huggingface_hub/models--moonshotai--Kimi-K2.5/snapshots/4d01dfe0332d63057c186e0b262165819efb6611" -) NUM_VISION_LAYERS = 2 NUM_TEXT_LAYERS = 2 LOADED_EXPERT_IDS = (0, 1, 2, 3) @@ -56,23 +52,6 @@ def resolve_model_path(model_name: str = KIMI_K25_MODEL_NAME) -> Path: return Path(snapshot_download(repo_id=model_name, cache_dir=os.environ.get("HF_HUB_CACHE"))) -def ensure_torch_fx_import_compatibility(): - if hasattr(hf_import_utils, "is_torch_fx_available"): - return - - def _is_torch_fx_available() -> bool: - if not hf_import_utils.is_torch_available(): - return False - try: - import torch.fx # noqa: F401 - - return True - except Exception: - return False - - hf_import_utils.is_torch_fx_available = _is_torch_fx_available - - def patch_kimi_tie_weights_compat(kimi_cls): tie_signature = inspect.signature(kimi_cls.tie_weights) if tuple(tie_signature.parameters) != ("self",): @@ -122,7 +101,6 @@ def _init_weights_compat(self, module): def load_kimi_k25_class(model_path_or_name): - ensure_torch_fx_import_compatibility() kimi_cls = get_class_from_dynamic_module( "modeling_kimi_k25.KimiK25ForConditionalGeneration", str(model_path_or_name), @@ -396,7 +374,6 @@ def load_kimi_k25_layer_subset_model( seed: int = 1234, ): set_deterministic(seed) - ensure_torch_fx_import_compatibility() resolved_model_path = Path(model_path) if model_path is not None else resolve_model_path() config = prepare_config(resolved_model_path) kimi_cls = load_kimi_k25_class(resolved_model_path) From aaed1b7fc4a3d830b803b0ae42a6fc008f3c096e Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Tue, 28 Jul 2026 13:16:01 +0530 Subject: [PATCH 33/33] Remove old pointers to weights Signed-off-by: Mamta Singh --- .../deepseek_v3/configuration_deepseek.py | 219 ------------------ .../models/deepseek_v3/modeling_deepseek.py | 18 +- examples/kimi_k2/export_kimi_k25_vision.py | 5 +- 3 files changed, 17 insertions(+), 225 deletions(-) delete mode 100644 QEfficient/transformers/models/deepseek_v3/configuration_deepseek.py diff --git a/QEfficient/transformers/models/deepseek_v3/configuration_deepseek.py b/QEfficient/transformers/models/deepseek_v3/configuration_deepseek.py deleted file mode 100644 index 7f68c3d86e..0000000000 --- a/QEfficient/transformers/models/deepseek_v3/configuration_deepseek.py +++ /dev/null @@ -1,219 +0,0 @@ -# ----------------------------------------------------------------------------- -# -# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. -# SPDX-License-Identifier: BSD-3-Clause -# -# ---------------------------------------------------------------------------- - -from transformers.configuration_utils import PretrainedConfig -from transformers.utils import logging - -logger = logging.get_logger(__name__) - -DEEPSEEK_PRETRAINED_CONFIG_ARCHIVE_MAP = {} - - -class DeepseekV3Config(PretrainedConfig): - r""" - This is the configuration class to store the configuration of a [`DeepseekV3Model`]. It is used to instantiate an DeepSeek - model according to the specified arguments, defining the model architecture. Instantiating a configuration with the - defaults will yield a similar configuration to that of the DeepSeek-V3. - - Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the - documentation from [`PretrainedConfig`] for more information. - - - Args: - vocab_size (`int`, *optional*, defaults to 129280): - Vocabulary size of the Deep model. Defines the number of different tokens that can be represented by the - `inputs_ids` passed when calling [`DeepseekV3Model`] - hidden_size (`int`, *optional*, defaults to 4096): - Dimension of the hidden representations. - intermediate_size (`int`, *optional*, defaults to 11008): - Dimension of the MLP representations. - moe_intermediate_size (`int`, *optional*, defaults to 1407): - Dimension of the MoE representations. - num_hidden_layers (`int`, *optional*, defaults to 32): - Number of hidden layers in the Transformer decoder. - num_nextn_predict_layers (`int`, *optional*, defaults to 1): - Number of nextn predict layers in the DeepSeekV3 Model. - num_attention_heads (`int`, *optional*, defaults to 32): - Number of attention heads for each attention layer in the Transformer decoder. - n_shared_experts (`int`, *optional*, defaults to None): - Number of shared experts, None means dense model. - n_routed_experts (`int`, *optional*, defaults to None): - Number of routed experts, None means dense model. - routed_scaling_factor (`float`, *optional*, defaults to 1.0): - Scaling factor or routed experts. - topk_method (`str`, *optional*, defaults to `gready`): - Topk method used in routed gate. - n_group (`int`, *optional*, defaults to None): - Number of groups for routed experts. - topk_group (`int`, *optional*, defaults to None): - Number of selected groups for each token(for each token, ensuring the selected experts is only within `topk_group` groups). - num_experts_per_tok (`int`, *optional*, defaults to None): - Number of selected experts, None means dense model. - moe_layer_freq (`int`, *optional*, defaults to 1): - The frequency of the MoE layer: one expert layer for every `moe_layer_freq - 1` dense layers. - first_k_dense_replace (`int`, *optional*, defaults to 0): - Number of dense layers in shallow layers(embed->dense->dense->...->dense->moe->moe...->lm_head). - \--k dense layers--/ - norm_topk_prob (`bool`, *optional*, defaults to False): - Whether to normalize the weights of the routed experts. - scoring_func (`str`, *optional*, defaults to 'softmax'): - Method of computing expert weights. - aux_loss_alpha (`float`, *optional*, defaults to 0.001): - Auxiliary loss weight coefficient. - seq_aux = (`bool`, *optional*, defaults to True): - Whether to compute the auxiliary loss for each individual sample. - num_key_value_heads (`int`, *optional*): - This is the number of key_value heads that should be used to implement Grouped Query Attention. If - `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if - `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When - converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed - by meanpooling all the original heads within that group. For more details checkout [this - paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to - `num_attention_heads`. - hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): - The non-linear activation function (function or string) in the decoder. - max_position_embeddings (`int`, *optional*, defaults to 2048): - The maximum sequence length that this model might ever be used with. - initializer_range (`float`, *optional*, defaults to 0.02): - The standard deviation of the truncated_normal_initializer for initializing all weight matrices. - rms_norm_eps (`float`, *optional*, defaults to 1e-06): - The epsilon used by the rms normalization layers. - use_cache (`bool`, *optional*, defaults to `True`): - Whether or not the model should return the last key/values attentions (not used by all models). Only - relevant if `config.is_decoder=True`. - pad_token_id (`int`, *optional*): - Padding token id. - bos_token_id (`int`, *optional*, defaults to 1): - Beginning of stream token id. - eos_token_id (`int`, *optional*, defaults to 2): - End of stream token id. - pretraining_tp (`int`, *optional*, defaults to 1): - Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this - document](https://huggingface.co/docs/transformers/parallelism) to understand more about it. This value is - necessary to ensure exact reproducibility of the pretraining results. Please refer to [this - issue](https://github.com/pytorch/pytorch/issues/76232). - tie_word_embeddings (`bool`, *optional*, defaults to `False`): - Whether to tie weight embeddings - rope_theta (`float`, *optional*, defaults to 10000.0): - The base period of the RoPE embeddings. - rope_scaling (`Dict`, *optional*): - Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling - strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is - `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update - `max_position_embeddings` to the expected new maximum. - attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`): - Whether to use a bias in the query, key, value and output projection layers during self-attention. - attention_dropout (`float`, *optional*, defaults to 0.0): - The dropout ratio for the attention probabilities. - - ```python - >>> from transformers import DeepseekV3Model, DeepseekV3Config - - >>> # Initializing a Deepseek-V3 style configuration - >>> configuration = DeepseekV3Config() - - >>> # Accessing the model configuration - >>> configuration = model.config - ```""" - - model_type = "deepseek_v3" - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - vocab_size=129280, - hidden_size=7168, - intermediate_size=18432, - moe_intermediate_size=2048, - num_hidden_layers=61, - num_nextn_predict_layers=1, - num_attention_heads=128, - num_key_value_heads=128, - n_shared_experts=1, - n_routed_experts=256, - ep_size=1, - routed_scaling_factor=2.5, - kv_lora_rank=512, - q_lora_rank=1536, - qk_rope_head_dim=64, - v_head_dim=128, - qk_nope_head_dim=128, - topk_method="noaux_tc", - n_group=8, - topk_group=4, - num_experts_per_tok=8, - moe_layer_freq=1, - first_k_dense_replace=3, - norm_topk_prob=True, - scoring_func="sigmoid", - aux_loss_alpha=0.001, - seq_aux=True, - hidden_act="silu", - max_position_embeddings=4096, - initializer_range=0.02, - rms_norm_eps=1e-6, - use_cache=True, - pad_token_id=None, - bos_token_id=0, - eos_token_id=1, - pretraining_tp=1, - tie_word_embeddings=False, - rope_theta=10000.0, - rope_scaling=None, - attention_bias=False, - attention_dropout=0.0, - **kwargs, - ): - self.vocab_size = vocab_size - self.max_position_embeddings = max_position_embeddings - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.moe_intermediate_size = moe_intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_nextn_predict_layers = num_nextn_predict_layers - self.num_attention_heads = num_attention_heads - self.n_shared_experts = n_shared_experts - self.n_routed_experts = n_routed_experts - self.ep_size = ep_size - self.routed_scaling_factor = routed_scaling_factor - self.kv_lora_rank = kv_lora_rank - self.q_lora_rank = q_lora_rank - self.qk_rope_head_dim = qk_rope_head_dim - self.v_head_dim = v_head_dim - self.qk_nope_head_dim = qk_nope_head_dim - self.topk_method = topk_method - self.n_group = n_group - self.topk_group = topk_group - self.num_experts_per_tok = num_experts_per_tok - self.moe_layer_freq = moe_layer_freq - self.first_k_dense_replace = first_k_dense_replace - self.norm_topk_prob = norm_topk_prob - self.scoring_func = scoring_func - self.aux_loss_alpha = aux_loss_alpha - self.seq_aux = seq_aux - # for backward compatibility - if num_key_value_heads is None: - num_key_value_heads = num_attention_heads - - self.num_key_value_heads = num_key_value_heads - self.hidden_act = hidden_act - self.initializer_range = initializer_range - self.rms_norm_eps = rms_norm_eps - self.pretraining_tp = pretraining_tp - self.use_cache = use_cache - self.rope_theta = rope_theta - self.rope_scaling = rope_scaling - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - - super().__init__( - pad_token_id=pad_token_id, - bos_token_id=bos_token_id, - eos_token_id=eos_token_id, - tie_word_embeddings=tie_word_embeddings, - **kwargs, - ) diff --git a/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py b/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py index 6c77562c8d..fdc8edb4d3 100644 --- a/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py +++ b/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py @@ -853,8 +853,15 @@ class QEffDeepseekV3MoE(nn.Module): def __qeff_init__( self, ): + if hasattr(self, "all_gate_qweight") and not hasattr(self, "experts"): + return + # Get common parameters from first expert first_expert = self.experts[0] + self._qeff_quantized_experts = hasattr(first_expert.gate_proj, "bits") + if not self._qeff_quantized_experts: + return + self.bits = first_expert.gate_proj.bits self.group_size = first_expert.gate_proj.group_size self.act_fn = first_expert.act_fn @@ -941,6 +948,7 @@ def __qeff_init__( self.all_down_gidx = torch.nn.Parameter( torch.stack([exp.down_proj.g_idx for exp in self.experts], dim=0), requires_grad=False ) + del self.experts def moe_old( self, @@ -1121,12 +1129,10 @@ def moe_waa_unpack(self, hidden_states, topk_indices, topk_weights): return torch.einsum("abc,ab->ac", selected_out, topk_weights) def forward(self, hidden_states): - print("Using new MoE forward with weights as activations") residuals = hidden_states orig_shape = hidden_states.shape topk_indices, topk_weights, router_probs, router_weights = self.gate(hidden_states) hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) - # hidden_states = self.moe_weights_as_activations(hidden_states, router_probs, router_weights).view(*orig_shape) hidden_states = self.moe_waa_unpack(hidden_states, topk_indices, topk_weights).view(*orig_shape) hidden_states = hidden_states + self.shared_experts(residuals) return hidden_states @@ -1136,8 +1142,15 @@ class QEffPrefillOnlyDeepseekV3MoE(nn.Module): def __qeff_init__( self, ): + if hasattr(self, "all_gate_qweight") and not hasattr(self, "experts"): + return + # Get common parameters from first expert first_expert = self.experts[0] + self._qeff_quantized_experts = hasattr(first_expert.gate_proj, "bits") + if not self._qeff_quantized_experts: + return + self.bits = first_expert.gate_proj.bits self.group_size = first_expert.gate_proj.group_size self.act_fn = first_expert.act_fn @@ -1224,6 +1237,7 @@ def __qeff_init__( self.all_down_gidx = torch.nn.Parameter( torch.stack([exp.down_proj.g_idx for exp in self.experts], dim=0), requires_grad=False ) + del self.experts def _cumsum_scatter_gather_update_expert_blocked( self, diff --git a/examples/kimi_k2/export_kimi_k25_vision.py b/examples/kimi_k2/export_kimi_k25_vision.py index 65a71168af..93f46cd6b3 100644 --- a/examples/kimi_k2/export_kimi_k25_vision.py +++ b/examples/kimi_k2/export_kimi_k25_vision.py @@ -37,9 +37,6 @@ prepare_config = load_kimi_utils.prepare_config set_deterministic = load_kimi_utils.set_deterministic -DEFAULT_MODEL_PATH = Path( - "/home/huggingface_hub/models--moonshotai--Kimi-K2.5/snapshots/4d01dfe0332d63057c186e0b262165819efb6611" -) PREFILL_SEQ_LEN = 512 CTX_LEN = 2048 BS = 1 @@ -48,7 +45,7 @@ def parse_args(): parser = argparse.ArgumentParser(description="Run Kimi K2.5 vision disaggregated vision -> prefill -> decode flow.") - parser.add_argument("--model-path", type=Path, default=DEFAULT_MODEL_PATH) + parser.add_argument("--model-path", type=Path, required=True) parser.add_argument( "--full-model", action="store_true",