Fix regression #71

Author: turboderp

examples/multimodal.py

@@ -18,7 +18,7 @@
         model_dir = "/mnt/str/models/gemma3-4b-it/exl3/5.0bpw/"
     case "mistral3":
         prompt_format = "mistral"
-        model_dir = "/mnt/str/models/mistral-small-3.1-24b-instruct/exl3/8.0bpw_H8"
+        model_dir = "/mnt/str/models/mistral-small-3.1-24b-instruct-2503/exl3/4.0bpw/"
 
 images = [
     # Cat

exllamav3/architecture/gemma3.py

@@ -582,4 +582,9 @@ def get_image_embeddings(
 
             mmes.append(mme)
 
-        return mmes if return_batch else mmes[0]
+        return mmes if return_batch else mmes[0]
+
+
+    @override
+    def prepare_inputs(self, input_ids: torch.Tensor, params: dict) -> torch.Tensor:
+        return input_ids

exllamav3/architecture/mistral3.py

@@ -517,3 +517,8 @@ def get_image_embeddings(
         })
 
         return mme
+
+
+    @override
+    def prepare_inputs(self, input_ids: torch.Tensor, params: dict) -> torch.Tensor:
+        return input_ids

exllamav3/modules/mlp.py

@@ -20,6 +20,7 @@ def __init__(
         key: str,
         hidden_size: int,
         intermediate_size: int,
+        out_size: int | None = None,
         key_up: str | None = None,
         key_down: str | None = None,
         qmap: str | None = None,
@@ -40,6 +41,7 @@ def __init__(
         self.activation_fn = activation_fn
         self.intermediate_size = intermediate_size
         self.intermediate_split_size = intermediate_split_size
+        self.out_size = out_size or hidden_size
 
         fkey, frange_up = None, None
 
@@ -99,9 +101,9 @@ def __init__(
                     config = config,
                     key = s_key_d,
                     in_features = b - a,
-                    out_features = hidden_size,
+                    out_features = self.out_size,
                     full_in_features = intermediate_size,
-                    full_out_features = hidden_size,
+                    full_out_features = self.out_size,
                     first_in_feature = a,
                     first_out_feature = 0,
                     qmap = qmap + ".down",