Fix gemma3 and cohere2 sliding window (#39)

* fix gemma3 and cohere2 sliding window

* fix

mlx_lm/models/cohere2.py

@@ -170,10 +170,22 @@ class CohereModel(nn.Module):
 
         if mask is None:
             j = self.args.sliding_window_pattern
-            mask = create_attention_mask(h, cache[j - 1 : j])
+            full_mask = create_attention_mask(h, cache[j - 1 : j])
+            sliding_window_mask = create_attention_mask(h, cache)
 
-        for layer, c in zip(self.layers, cache):
-            h = layer(h, mask, c)
+        for i, (layer, c) in enumerate(zip(self.layers, cache)):
+            is_global = (
+                i % self.args.sliding_window_pattern
+                == self.args.sliding_window_pattern - 1
+            )
+
+            local_mask = mask
+            if mask is None and is_global:
+                local_mask = full_mask
+            elif mask is None:
+                local_mask = sliding_window_mask
+
+            h = layer(h, local_mask, c)
 
         return self.norm(h)
 

mlx_lm/models/gemma3_text.py

@@ -179,17 +179,18 @@ class Gemma3Model(nn.Module):
             sliding_window_mask = create_attention_mask(h, cache)
 
         for i, (layer, c) in enumerate(zip(self.layers, cache)):
-            is_sliding = (
+            is_global = (
                 i % self.args.sliding_window_pattern
                 == self.args.sliding_window_pattern - 1
             )
 
-            if mask is None and is_sliding:
-                mask = sliding_window_mask
+            local_mask = mask
+            if mask is None and is_global:
+                local_mask = full_mask
             elif mask is None:
-                mask = full_mask
+                local_mask = sliding_window_mask
 
-            h = layer(h, mask, c)
+            h = layer(h, local_mask, c)
 
         return self.norm(h)