fix dsv32 and gemma3 (#664)

Author: awni

mlx_lm/models/deepseek_v32.py

@@ -98,7 +98,7 @@ def __call__(
         k_pe = self.rope(k_pe, offset=offset)
         k = mx.concatenate([k_pe, k_nope], axis=-1)
         if cache is not None:
-            k, _ = cache.update_and_fetch(k, mx.zeros([b, 1, s, 0]))
+            k, _ = cache.update_and_fetch(mx.zeros_like(k), mx.zeros([b, 1, s, 0]))
         if k.shape[2] <= self.index_topk:
             return None
         scores = q @ k.swapaxes(-1, -2)
@@ -222,6 +222,11 @@ def __call__(
             if mask is not None:
                 sparse_mask = sparse_mask & mask
             mask = sparse_mask
+        # Ensure the indexer cache is evaluated even if the topk_indices are unused
+        # to keep the graph from getting too large
+        if cache is not None:
+            cache[0].keys = mx.depends(cache[0].keys, (cache[1].keys, cache[1].values))
+
         output = scaled_dot_product_attention(
             queries, keys, values, cache=cache[0], scale=self.scale, mask=mask
         )

mlx_lm/models/gemma3_text.py

@@ -54,13 +54,20 @@ def __init__(self, args: ModelArgs, layer_idx: int):
         self.k_norm = RMSNorm(dims=head_dim, eps=args.rms_norm_eps)
         self.is_sliding = (layer_idx + 1) % args.sliding_window_pattern != 0
 
-        self.rope = initialize_rope(
-            dims=head_dim,
-            base=(args.rope_local_base_freq if self.is_sliding else args.rope_theta),
-            traditional=False,
-            max_position_embeddings=args.max_position_embeddings,
-            scaling_config=args.rope_scaling,
-        )
+        if self.is_sliding:
+            self.rope = initialize_rope(
+                dims=head_dim,
+                base=args.rope_local_base_freq,
+                traditional=False,
+            )
+        else:
+            self.rope = initialize_rope(
+                dims=head_dim,
+                base=args.rope_theta,
+                traditional=False,
+                max_position_embeddings=args.max_position_embeddings,
+                scaling_config=args.rope_scaling,
+            )
 
     def __call__(
         self,