Add Support for Kimi K2.5 Vision

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()

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

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__()

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

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,
+        )