Add Support for MiniMax 2.5 Model

QEfficient/base/modeling_qeff.py

@@ -250,15 +250,33 @@ def _export(
         tmp_onnx_path = tmp_onnx_dir / f"{self.model_name}.onnx"
         tmp_onnx_dir.mkdir(parents=True, exist_ok=True)
 
+        def _resolve_pkv_layers(pkv_obj):
+            if isinstance(pkv_obj, (list, tuple)):
+                return pkv_obj
+            if hasattr(pkv_obj, "to_legacy_cache"):
+                return pkv_obj.to_legacy_cache()
+            if hasattr(pkv_obj, "layers"):
+                layers = []
+                for layer in pkv_obj.layers:
+                    keys = getattr(layer, "keys", None)
+                    values = getattr(layer, "values", None)
+                    layers.append((keys, values))
+                return tuple(layers)
+            return None
+
         # Create input_names from example_inputs
         input_names = []
         for param in inspect.signature(self.model.forward).parameters:
             if param in example_inputs:
                 if param == "past_key_values":
-                    for i in range(len(example_inputs["past_key_values"])):
-                        if len(example_inputs["past_key_values"][0]) == 2:
+                    pkv_layers = _resolve_pkv_layers(example_inputs["past_key_values"])
+                    if pkv_layers is None:
+                        input_names.append(param)
+                        continue
+                    for i in range(len(pkv_layers)):
+                        if len(pkv_layers[0]) == 2:
                             input_names.extend([f"past_key.{i}", f"past_value.{i}"])
-                        elif len(example_inputs["past_key_values"][0]) == 4:
+                        elif len(pkv_layers[0]) == 4:
                             input_names.extend(
                                 [
                                     f"past_key_self.{i}",
@@ -269,16 +287,17 @@ def _export(
                             )
                         else:
                             raise ValueError(
-                                f"Unknown shape of past_key_values! Expected length of past_key_values for each layer to be either 2 or 4 but got {len(example_inputs['past_key_values'][0])}"
+                                f"Unknown shape of past_key_values! Expected length of past_key_values for each layer to be either 2 or 4 but got {len(pkv_layers[0])}"
                             )
                 else:
                     input_names.append(param)
 
         try:
             torch.onnx.export(
                 self.model,
-                (example_inputs,),
+                (),
                 str(tmp_onnx_path),
+                kwargs=example_inputs,
                 input_names=input_names,
                 output_names=output_names,
                 dynamic_axes=dynamic_axes,

QEfficient/customop/matmulnbits.py

@@ -55,7 +55,7 @@ def dequantize_blockwise_bits(quant_values, scale, zero_point, bits, group_size,
         except RuntimeError:
             expand_zero_point = expand_zero_point.reshape(quant_values.shape[0], -1, 1)
             expand_zero_point = expand_zero_point[:, : quant_values.shape[1]]
-    if g_idx is not None and g_idx[:32].sum().item() != 0:
+    if g_idx is not None and (not getattr(g_idx, "is_meta", False)) and g_idx[:32].sum().item() != 0:
         float_values = (
             (expand_quant_value.reshape(expand_quant_value.shape[0], -1) - expand_zero_point[:, g_idx, 0])
             * aligned_scale[:, g_idx, 0]
@@ -117,7 +117,10 @@ def pack_on_device(self, int_weight, int_zeros):
             raise ValueError("only 4bit is supported by ONNXRUNTIME for now.")
 
         # Order of groups
-        self.act_order = self.g_idx[: self.group_size // self.bits].sum().item() != 0
+        if getattr(self.g_idx, "is_meta", False):
+            self.act_order = False
+        else:
+            self.act_order = self.g_idx[: self.group_size // self.bits].sum().item() != 0
 
         intzeros_pt = int_zeros.T if int_zeros.dtype == self.scales.dtype else int_zeros.T.byte()
         scales_pt = self.scales.T.to(int_weight.device)