From 465b107c2a55bb36ffa28138e5a7d64ccc419c15 Mon Sep 17 00:00:00 2001 From: Awni Hannun Date: Fri, 21 Mar 2025 08:50:35 -0700 Subject: [PATCH] Fix gemma3 and cohere2 sliding window (#39) * fix gemma3 and cohere2 sliding window * fix --- mlx_lm/models/cohere2.py | 18 +++++++++++++++--- mlx_lm/models/gemma3_text.py | 11 ++++++----- 2 files changed, 21 insertions(+), 8 deletions(-) diff --git a/mlx_lm/models/cohere2.py b/mlx_lm/models/cohere2.py index cdcbe45..b5d05ca 100644 --- a/mlx_lm/models/cohere2.py +++ b/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) diff --git a/mlx_lm/models/gemma3_text.py b/mlx_lm/models/gemma3_text.py index 7da1d4c..275cf6f 100644 --- a/mlx_lm/models/gemma3_text.py +++ b/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)