Fix Gemma 4 sanitize() not stripping KV projections for shared layers (#1240)

Author: 0k1

mlx_lm/models/gemma4_text.py

@@ -609,6 +609,7 @@ def __call__(
 
     def sanitize(self, weights):
         sanitized = {}
+        first_kv_shared = self.args.num_hidden_layers - self.args.num_kv_shared_layers
         for k, v in weights.items():
             if any(
                 s in k
@@ -622,6 +623,18 @@ def sanitize(self, weights):
             ):
                 continue
 
+            # KV-shared layers reuse K/V from earlier layers — drop their projections
+            if any(
+                s in k
+                for s in (".self_attn.k_proj", ".self_attn.v_proj", ".self_attn.k_norm")
+            ):
+                try:
+                    layer_idx = int(k.split("layers.")[1].split(".")[0])
+                    if layer_idx >= first_kv_shared:
+                        continue
+                except (IndexError, ValueError):
+                    pass
+
             if k.endswith(".experts.gate_up_proj"):
                 base = k.removesuffix(".gate_up_proj")
                 gate, up = map(mx.contiguous, mx.split(v, 2, axis=-2))