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/base/onnx_transforms.py b/QEfficient/base/onnx_transforms.py
index eb27051b9d..b24b4ebcb5 100644
--- a/QEfficient/base/onnx_transforms.py
+++ b/QEfficient/base/onnx_transforms.py
@@ -44,9 +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
@@ -112,7 +111,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/generation/embedding_handler.py b/QEfficient/generation/embedding_handler.py
index 2a5a61f6b8..c8f3528094 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:
diff --git a/QEfficient/transformers/cache_utils.py b/QEfficient/transformers/cache_utils.py
index 7af4bfbefa..4f016b2d49 100755
--- a/QEfficient/transformers/cache_utils.py
+++ b/QEfficient/transformers/cache_utils.py
@@ -423,8 +423,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 = ctx_scatter(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 = ctx_scatter_cb(self.ckv, batch_index, scatter_position_ids, compressed_kv)
+ else:
+ self.ckv = ctx_scatter(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 = InvalidIndexProvider._get_invalid_idx_value()
ctx_indices = torch.where(invalid_mask, invalid_idx_value, ctx_indices)
- ckv_out = ctx_gather(ckv_out, ctx_indices, ctx_len)
- ckv_out = torch.where(invalid_mask.unsqueeze(-1), torch.zeros_like(ckv_out, dtype=ckv_out.dtype), ckv_out)
+ if batch_index is not None:
+ ckv_out = ctx_gather_cb(ckv_out, batch_index, ctx_indices, ctx_len)
+ else:
+ 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
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 = ctx_scatter(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 = ctx_scatter_cb(self.k_pe, batch_index, scatter_position_ids, k_pe_cache)
+ else:
+ 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]
@@ -451,14 +466,18 @@ def update_k_pe(self, k_pe_cache, cache_kwargs):
invalid_idx_value = InvalidIndexProvider._get_invalid_idx_value()
ctx_indices = torch.where(invalid_mask, invalid_idx_value, ctx_indices)
- k_pe_out = ctx_gather(k_pe_out, ctx_indices, ctx_len)
- k_pe_out = torch.where(invalid_mask.unsqueeze(-1), torch.zeros_like(k_pe_out, dtype=k_pe_out.dtype), k_pe_out)
+ if batch_index is not None:
+ k_pe_out = ctx_gather_cb(k_pe_out, batch_index, ctx_indices, ctx_len)
+ else:
+ 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
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, dtype=position_ids.dtype)[None, None, ...]
gather_limit = position_ids.max(1, keepdim=True).values.unsqueeze(1).to(position_ids.dtype)
@@ -468,8 +487,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 = ctx_gather_blocked_kv(ckv_out, ctx_indices)
+ if batch_index is not None:
+ 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(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
@@ -478,6 +500,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, dtype=position_ids.dtype)[None, None, ...]
gather_limit = position_ids.max(1, keepdim=True).values.unsqueeze(1).to(position_ids.dtype)
@@ -487,22 +510,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 = ctx_gather_blocked_kv(k_pe_out, ctx_indices)
+ if batch_index is not None:
+ 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(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
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 = ctx_scatter(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 = ctx_scatter_cb(self.k_pe, batch_index, scatter_position_ids, k_pe_cache)
+ else:
+ 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):
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 = ctx_scatter(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 = ctx_scatter_cb(self.ckv, batch_index, scatter_position_ids, compressed_kv)
+ else:
+ self.ckv = ctx_scatter(self.ckv, position_ids, compressed_kv)
return self.ckv
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 8f08a3169b..fdc8edb4d3 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.customop.utils import select_interface
from QEfficient.transformers.cache_utils import QEffDynamicCache, QEffDynamicCompressedKVRopeCache
@@ -116,16 +123,6 @@ 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),
- )
-
class DeepseekV3YarnRotaryEmbedding(DeepseekV3RotaryEmbedding):
def __init__(
@@ -185,7 +182,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:
@@ -206,8 +203,6 @@ 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)
b, h, s, d = q.shape
q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)
@@ -273,6 +268,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()
@@ -294,8 +291,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)
@@ -310,6 +306,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,
@@ -340,10 +337,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)
@@ -361,8 +361,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)
@@ -379,13 +378,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())
@@ -395,6 +392,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,
@@ -422,6 +420,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()
@@ -453,8 +453,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)
@@ -504,7 +503,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)
@@ -540,6 +539,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())
@@ -556,6 +557,8 @@ def fused_forward(
use_cache,
cache_position,
mla_absorption,
+ cos_cached,
+ sin_cached,
**kwargs,
)
elif getattr(blocking_config, "mode", None) == "kv":
@@ -571,6 +574,8 @@ def fused_forward(
use_cache,
cache_position,
mla_absorption,
+ cos_cached,
+ sin_cached,
**kwargs,
)
else:
@@ -586,6 +591,8 @@ def fused_forward(
use_cache,
cache_position,
mla_absorption,
+ cos_cached,
+ sin_cached,
**kwargs,
)
@@ -600,6 +607,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()
@@ -627,15 +636,19 @@ 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)
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
@@ -666,6 +679,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()
@@ -687,17 +702,14 @@ 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)
)
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)
@@ -734,6 +746,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())
@@ -748,6 +762,8 @@ def forward(
output_attentions,
use_cache,
cache_position,
+ cos_cached,
+ sin_cached,
**kwargs,
)
else:
@@ -761,34 +777,180 @@ 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", "16"))
+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(
- [exp.gate_proj.compressor.decompress_module(exp.gate_proj).T.unsqueeze(0) for exp in self.experts],
- dim=0,
- )
+ 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
+ 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(
- [exp.up_proj.compressor.decompress_module(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(
- [exp.down_proj.compressor.decompress_module(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,
+ )
+
+ # 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.act_fn = self.experts[0].act_fn
+ 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
+ )
+ del self.experts
- def moe(
+ def moe_old(
self,
hidden_states: torch.Tensor,
topk_indices: torch.Tensor,
@@ -798,29 +960,180 @@ 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()]
- expert_in = (
- hidden_states.unsqueeze(1).expand(-1, self.gate.top_k, -1).contiguous().view(-1, 1, self.config.hidden_size)
+ 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]]
+ 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):
+ 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, self.all_gate_scales, gate_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
)
- gate_out = torch.bmm(expert_in, gate_proj)
- up_out = torch.bmm(expert_in, up_proj)
- 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)
- final_hidden_states = torch.einsum("abc->ac", experts_out)
+ 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
+ )
- return final_hidden_states.type(hidden_states.dtype)
+ 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))
+
+ 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):
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_waa_unpack(hidden_states, topk_indices, topk_weights).view(*orig_shape)
hidden_states = hidden_states + self.shared_experts(residuals)
return hidden_states
@@ -829,83 +1142,339 @@ 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)
+ 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
+ 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(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))
+ # 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
+ )
+ del self.experts
- 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
+
+ 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):
@@ -924,6 +1493,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
@@ -945,6 +1517,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:
@@ -958,13 +1532,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,)
@@ -1003,6 +1582,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,
@@ -1065,6 +1646,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,)
@@ -1081,6 +1665,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/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/modeling_kimi_k25.py b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py
new file mode 100644
index 0000000000..abdb1a5701
--- /dev/null
+++ b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py
@@ -0,0 +1,1145 @@
+# -----------------------------------------------------------------------------
+#
+# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# ----------------------------------------------------------------------------
+
+import math
+from typing import List, Optional, Tuple, Type, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import activations
+
+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.models.llava.modeling_llava import LlavaCausalLMOutputWithPast
+
+from QEfficient.utils import constants
+
+
+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)
+ 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)
+
+
+class QEffLearnable2DInterpPosEmbDivided_fixed(nn.Module):
+ def __qeff_init__(self):
+ self.interpolation_mode = "bilinear"
+
+
+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
+
+ 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,
+ "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)
+
+ 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)
+
+
+class QEffKimiK25EncoderWrapper(nn.Module):
+ def __init__(self, model):
+ super().__init__()
+ self.model = model
+ 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.vision_tower.encoder.blocks[0].__class__}
+
+ 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:
+ vision_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)
+ 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]
+ 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)
+
+ hidden_states = self.model.vision_tower.patch_embed.proj(pixel_values).view(pixel_values.size(0), -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),
+ )
+ hidden_states = hidden_states + pos_emb_2d.repeat(num_images, 1)
+
+ rope_2d = self.model.vision_tower.encoder.rope_2d
+ 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
+ hidden_states = hidden_states.to(encoder_dtype)
+ freqs_cis = freqs_cis.to(encoder_dtype)
+ for block in self.model.vision_tower.encoder.blocks:
+ 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
+ hidden_states = self.model.vision_tower.encoder.final_layernorm(hidden_states.to(final_ln_dtype))
+
+ merge_kernel_size = self.model.vision_tower.merge_kernel_size
+ kernel_height, kernel_width = merge_kernel_size
+ new_height = h // kernel_height
+ new_width = w // kernel_width
+ 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)
+
+ pre_norm_dtype = self.model.mm_projector.pre_norm.weight.dtype
+ merged = merged.to(pre_norm_dtype)
+ 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):
+ 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.language_model.model.layers[0].__class__}
+
+ def forward(
+ self,
+ input_ids: torch.LongTensor = None,
+ inputs_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,
+ 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:
+ inputs_embeds = self.model.get_input_embeddings()(input_ids)
+ 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 vision_embeds rank 2 or 3, got {vision_embeds.dim()}.")
+
+ 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)
+ 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)
+ inputs_embeds = inputs_embeds.to(vision_embeds_for_state.dtype)
+
+ merged_inputs_embeds, merged_attention_mask, _, merged_position_ids = (
+ self.model._qeff_merge_input_ids_with_image_features(
+ image_features=vision_embeds_for_state,
+ inputs_embeds=inputs_embeds,
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ labels=None,
+ )
+ )
+
+ 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] + effective_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),
+ )
+
+ image_position_delta = torch.clamp(merged_image_tokens - selected_image_tokens, min=0)
+ 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
+ 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:
+ position_ids = torch.where(attention_mask > 0, position_ids, torch.full_like(position_ids, -1))
+
+ outputs = self.model.language_model(
+ inputs_embeds=inputs_embeds,
+ 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,
+ )
+ logits = outputs[0].float()
+
+ 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)
+ 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
+class QEffKimiK25ForConditionalGeneration(nn.Module):
+ def get_qeff_vision_encoder(self):
+ return QEffKimiK25EncoderWrapper(self)
+
+ def get_qeff_language_decoder(self):
+ return QEffKimiK25DecoderWrapper(self)
+
+ def _qeff_merge_input_ids_with_image_features(
+ 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,
+ pixel_values: torch.FloatTensor | list[torch.FloatTensor] | None = None,
+ grid_thws: 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 get_output_names(self, kv_offload: bool = False):
+ vision_output_names = ["vision_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:
+ 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
+ else:
+ lang_output_names.insert(1, "pixel_values_RetainedState")
+ lang_output_names.insert(2, "image_idx_output")
+ return lang_output_names
+ return output_names
+
+ def get_dummy_inputs(
+ self,
+ 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
+
+ inputs_shapes = {}
+ inputs_shapes["pixel_values"] = (
+ 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_EXAMPLE_IMAGE_NUM_IMAGE_TOKENS,
+ self.language_model.config.hidden_size,
+ )
+ inputs_shapes["image_idx"] = (1, 1)
+ 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["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)
+ 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),
+ }
+
+ 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=constants.ONNX_EXPORT_CTX_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)
+ )
+
+ 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)
+
+ if continuous_batching:
+ lang_inputs["batch_index"] = torch.arange(bs).view(bs, 1)
+
+ inputs = {}
+ if kv_offload:
+ inputs["vision"] = vision_inputs
+ inputs["lang"] = lang_inputs
+ else:
+ lang_inputs.pop("vision_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["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: "grid_h"},
+ "w_shape": {0: "grid_w"},
+ "vision_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
+
+ 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: 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: cache_batch_axis, 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("vision_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)
+ compiler_options.pop("img_size", 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)
+ 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
+
+ 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)}."
+ )
+ 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))
+ 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 (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
+ grid_heights = []
+ grid_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
+ grid_heights.append((resized_height + pad_height) // patch_size)
+ grid_widths.append((resized_width + pad_width) // patch_size)
+ else:
+ 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(grid_heights), "num_frames")
+
+ vision = []
+ max_num_image_tokens = 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={grid_height}, w={grid_width} must be divisible by "
+ f"merge_kernel_size={merge_kernel_size}."
+ )
+
+ 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": computed_num_patches,
+ "grid_h": grid_height,
+ "grid_w": grid_width,
+ "num_image_tokens": computed_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": max_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": 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.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": max_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": 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 = {}
+
+ if kv_offload:
+ specializations["vision"] = vision
+ specializations["lang"] = lang
+ return specializations, compiler_options
+ else:
+ return [{**vision_spec, **lang_spec} for vision_spec in vision for lang_spec in lang], compiler_options
diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py
index a573bb18d1..078ae9f4fd 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,
@@ -1233,6 +1236,10 @@ 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:
+ 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 +1442,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---
@@ -2089,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:
@@ -2330,6 +2340,14 @@ 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())
+ 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 = {}
@@ -2440,6 +2458,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(
@@ -3066,7 +3087,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
@@ -3363,6 +3386,7 @@ class QEFFAutoModelForCausalLM(QEFFBaseModel):
AwqToMatmulNbitsTransform,
GPTQToMatmulNbitsTransform,
FP8DeQuantLinearToLinearTransform,
+ PackQuantizedInt4ToMatMulNBitsTransform,
Mxfp4GptOssExpertDequantizeTransform,
CustomOpsTransform,
KVCacheTransform,
@@ -3646,7 +3670,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
@@ -3812,13 +3838,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,
@@ -3829,6 +3855,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)
@@ -3939,8 +3969,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:
@@ -4028,6 +4064,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 a5eefb542a..36c5cf8377 100755
--- a/QEfficient/transformers/models/pytorch_transforms.py
+++ b/QEfficient/transformers/models/pytorch_transforms.py
@@ -324,6 +324,7 @@
QEffDisentangledSelfAttention,
)
from QEfficient.transformers.models.deepseek_v3.modeling_deepseek import (
+ QEffDeepseekMoEGate,
QEffDeepseekV3Attention,
QEffDeepseekV3CustomRMSNormAIC,
QEffDeepseekV3DecoderLayer,
@@ -442,6 +443,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,
@@ -1315,6 +1321,21 @@ class KVCacheExternalModuleMapperTransform(ExternalModuleMapperTransform):
"RMSNorm": {
"forward": QEFFGrok1CustomRMSNormAIC.forward,
},
+ "KimiK25ForConditionalGeneration": {
+ "_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,
+ "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,
@@ -1326,8 +1347,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,
@@ -1350,8 +1374,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,
},
}
@@ -1361,7 +1386,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 c8a64adefb..f26b09087b 100644
--- a/QEfficient/utils/constants.py
+++ b/QEfficient/utils/constants.py
@@ -148,6 +148,7 @@ 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"))
# Generic config key aliases used across model families.
ATTENTION_HEAD_CONFIG_KEYS = ("num_attention_heads", "n_head", "n_heads", "num_heads")
KV_HEAD_CONFIG_KEYS = ("num_key_value_heads", "n_kv_heads", "num_kv_heads", "effective_n_kv_heads")
@@ -185,7 +186,11 @@ 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_PATCH_SIZE = 14
+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/test_utils.py b/QEfficient/utils/test_utils.py
index c91248c3e8..cab9300f32 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"):
@@ -499,6 +501,7 @@ class ModelConfig:
"Qwen/Qwen3-VL-Reranker-8B",
"Qwen/Qwen3.5-0.8B",
"Qwen/Qwen3.5-35B-A3B",
+ "moonshotai/Kimi-K2.5",
"tiny-random/gemma-4-dense",
"tiny-random/gemma-4-moe",
}
diff --git a/README.md b/README.md
index b7d8fe4b42..81b6bee2b5 100644
--- a/README.md
+++ b/README.md
@@ -13,6 +13,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)
- [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/introduction.md b/docs/source/introduction.md
index a02bbef8c7..616d5bbd84 100644
--- a/docs/source/introduction.md
+++ b/docs/source/introduction.md
@@ -22,6 +22,7 @@ For other models, there is comprehensive documentation to inspire upon the chang
*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/release_docs.md b/docs/source/release_docs.md
index 2724ab6dd6..09d4e32931 100644
--- a/docs/source/release_docs.md
+++ b/docs/source/release_docs.md
@@ -227,6 +227,12 @@ 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 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)
+
- **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.
@@ -279,7 +285,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 48a0e5b0f8..be5009aac7 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..e9504b6b0c 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,23 @@ 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-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 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
new file mode 100644
index 0000000000..93f46cd6b3
--- /dev/null
+++ b/examples/kimi_k2/export_kimi_k25_vision.py
@@ -0,0 +1,396 @@
+# -----------------------------------------------------------------------------
+#
+# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# -----------------------------------------------------------------------------
+
+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.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
+
+PREFILL_SEQ_LEN = 512
+CTX_LEN = 2048
+BS = 1
+GENERATION_LEN = 10
+
+
+def parse_args():
+ parser = argparse.ArgumentParser(description="Run Kimi K2.5 vision disaggregated vision -> prefill -> decode flow.")
+ parser.add_argument("--model-path", type=Path, required=True)
+ 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)
+ 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(
+ "--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("--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 _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)
+ config = prepare_config(args.model_path)
+ kimi_cls = load_kimi_k25_class(args.model_path)
+
+ model_kwargs = {
+ "config": config,
+ "trust_remote_code": True,
+ "attn_implementation": "eager",
+ "torch_dtype": torch.float32,
+ }
+
+ 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_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,
+ dtype=torch.float32,
+ )
+ 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.")
+
+ model.vision_tower.patch_embed.pos_emb.interpolation_mode = "bilinear"
+ return model.eval().to("cpu"), tokenizer, processor
+
+
+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": "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
+
+
+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,
+ )
+
+ image, inputs = _prepare_inputs(processor, args)
+ inputs["pixel_values"] = inputs["pixel_values"].to(qeff_model.model.config.torch_dtype)
+
+ 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}")
+
+ 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])
+
+ 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()
+
+ print(generated_ids)
+ print(tokenizer.batch_decode(torch.as_tensor(generated_ids), skip_special_tokens=True))
+
+
+if __name__ == "__main__":
+ main()
diff --git a/tests/configs/image_text_model_configs.json b/tests/configs/image_text_model_configs.json
index 1817bf8bc4..5a301b0dcd 100644
--- a/tests/configs/image_text_model_configs.json
+++ b/tests/configs/image_text_model_configs.json
@@ -609,6 +609,57 @@
}
}
},
+ {
+ "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
+ }
+ }
+ },
{
"model_name": "tiny-random/gemma-4-dense",
"model_type": "gemma4",
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 d7289dbb75..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
@@ -33,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)))
@@ -59,7 +66,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"]
@@ -68,7 +74,27 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100_CB(
full_batch_size = model_config_dict[model_name]["full_batch_size"]
max_gen_len = NEW_GENERATION_TOKENS
- if config is None:
+ 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),
+ 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_config_dict)
+ model_hf, tokenizer, processor = load_kimi_k25_model_from_config(config)
+ 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
)
@@ -189,6 +215,26 @@ 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):
+ 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")
+ images.append(image)
+
+ 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,
+ "image_height": image_height,
+ "image_width": image_width,
+ }
+ )
else:
processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True, padding=True)
use_fast = model_name != "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
@@ -231,6 +277,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)
+
print("QPC Outputs (QAIC):")
exec_info = qeff_model.generate(
tokenizer=tokenizer,
@@ -253,7 +300,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(
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 c448cacd16..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
@@ -37,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
@@ -84,7 +89,17 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100(
ort_tokens = None
n_layer = num_hidden_layers
qaic_config = copy.deepcopy(qaic_config) if qaic_config is not None else None
- if config is None:
+
+ 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),
+ kv_offload=kv_offload,
+ config=model_hf.config,
+ torch_dtype=torch_dtype,
+ )
+ elif config is None:
config = AutoConfig.from_pretrained(
model_name, trust_remote_code=True, padding=model_name not in ModelConfig.MOLMO_MODELS
)
@@ -233,6 +248,33 @@ 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):
+ 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)
+ 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,
+ "image_height": image.height,
+ "image_width": image.width,
+ }
+ )
+
else:
processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True, padding=True)
image = Image.open(_session.get(img_url, stream=True).raw)
@@ -491,6 +533,7 @@ def test_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100_qnn(model_name, kv_off
"google/gemma-3-4b-it",
"tiny-random/gemma-4-dense",
"tiny-random/gemma-4-moe",
+ "moonshotai/Kimi-K2.5",
]:
pytest.skip("QNN is not supported for these models yet.")
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..33ce915ae8
--- /dev/null
+++ b/tests/transformers/models/image_text_to_text/test_kimi_k25_disagg.py
@@ -0,0 +1,326 @@
+# -----------------------------------------------------------------------------
+#
+# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# ----------------------------------------------------------------------------
+
+import copy
+from io import BytesIO
+from pathlib import Path
+
+import numpy as np
+import pytest
+import requests
+import torch
+from PIL import Image
+
+from QEfficient import QEFFAutoModelForImageTextToText
+from QEfficient.generation.cloud_infer import QAICInferenceSession
+from tests.utils.load_kimi_utils import (
+ LOADED_EXPERT_IDS,
+ NUM_EXPERTS_PER_TOKEN,
+ load_kimi_k25_class,
+ 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 = 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 _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 _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)
+ 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"
+ 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_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,
+ 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,
+ prefill_only=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")
+
+ decode_qpc_path = qeff_model.compile(
+ prefill_seq_len=1,
+ 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")
+ _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()
+
+ 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),
+ "vision_embeds": vision_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,
+ "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)
+
+ 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():
+ 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 = run_kimi_k25_hf_model_on_pytorch(
+ copy.deepcopy(model), processor, _clone_inputs(inputs), max_gen_len=GENERATION_LEN
+ )
+
+ qeff_model = QEFFAutoModelForImageTextToText(
+ model,
+ kv_offload=True,
+ config=model.config,
+ torch_dtype=torch.float32,
+ layerwise=False,
+ )
+
+ 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"
diff --git a/tests/unit_test/models/test_model_quickcheck.py b/tests/unit_test/models/test_model_quickcheck.py
index 8ebaceaace..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"),
@@ -3551,3 +3621,50 @@ 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 image 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,
+ image_height=[512, 448],
+ image_width=[910, 448],
+ num_frames=[1, 1],
+ kv_offload=True,
+ )
+ assert specs["vision"] == [
+ {"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"])
+
+ 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,
+ )
diff --git a/tests/unit_test/transforms/test_quantization_transforms.py b/tests/unit_test/transforms/test_quantization_transforms.py
index b7fa03c1d2..e89432873c 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,57 @@ 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):
+ import importlib
+ from types import SimpleNamespace
+
+ 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 +391,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 +405,16 @@ 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"
+ )
diff --git a/tests/utils/load_kimi_utils.py b/tests/utils/load_kimi_utils.py
new file mode 100644
index 0000000000..af1f8c259b
--- /dev/null
+++ b/tests/utils/load_kimi_utils.py
@@ -0,0 +1,496 @@
+# -----------------------------------------------------------------------------
+#
+# 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
+
+KIMI_K25_MODEL_NAME = "moonshotai/Kimi-K2.5"
+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 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):
+ 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)
+ 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):
+ 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)
+ 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