fix for rnj-1 (#657)

Co-authored-by: Awni Hannun <awni@apple.com>

mlx_lm/_version.py

@@ -1,3 +1,3 @@
 # Copyright © 2023-2025 Apple Inc.
 
-__version__ = "0.28.4"
+__version__ = "0.29.0"

mlx_lm/generate.py

@@ -948,6 +948,22 @@ class BatchGenerator:
 
         self.active_batch = None
 
+        if mx.metal.is_available():
+            self._old_wired_limit = mx.set_wired_limit(
+                mx.metal.device_info()["max_recommended_working_set_size"]
+            )
+        else:
+            self._old_wired_limit = None
+
+    def close(self):
+        if self._old_wired_limit is not None:
+            mx.synchronize(generation_stream)
+            mx.set_wired_limit(self._old_wired_limit)
+            self._old_wired_limit = None
+
+    def __del__(self):
+        self.close()
+
     def insert(
         self, prompts, max_tokens: Union[List[int], int, None] = None, caches=None
     ):
@@ -1196,23 +1212,23 @@ def batch_generate(
     if verbose:
         print(f"[batch_generate] Finished processing 0/{num_samples} ...", end="\r")
 
-    with wired_limit(model, [generation_stream]):
-        uids = gen.insert(prompts, max_tokens, caches=prompt_caches)
-        results = {uid: [] for uid in uids}
-        prompt_caches = {}
-        while responses := gen.next():
-            for r in responses:
-                if r.finish_reason is not None:
-                    if return_prompt_caches:
-                        prompt_caches[r.uid] = r.prompt_cache
-                    if verbose:
-                        fin += 1
-                        print(
-                            f"[batch_generate] Finished processing {fin}/{num_samples} ...",
-                            end="\r",
-                        )
-                if r.finish_reason != "stop":
-                    results[r.uid].append(r.token)
+    uids = gen.insert(prompts, max_tokens, caches=prompt_caches)
+    results = {uid: [] for uid in uids}
+    prompt_caches = {}
+    while responses := gen.next():
+        for r in responses:
+            if r.finish_reason is not None:
+                if return_prompt_caches:
+                    prompt_caches[r.uid] = r.prompt_cache
+                if verbose:
+                    fin += 1
+                    print(
+                        f"[batch_generate] Finished processing {fin}/{num_samples} ...",
+                        end="\r",
+                    )
+            if r.finish_reason != "stop":
+                results[r.uid].append(r.token)
+    gen.close()
     if verbose:
         print(f"[batch_generate] Finished processing {fin}/{num_samples}")
 

mlx_lm/models/gemma3_text.py

@@ -2,13 +2,14 @@
 
 from dataclasses import dataclass
 from functools import partial
-from typing import Any, Optional
+from typing import Any, Dict, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
 
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .cache import KVCache, RotatingKVCache
+from .rope_utils import initialize_rope
 
 
 @dataclass
@@ -22,12 +23,13 @@ class ModelArgs(BaseModelArgs):
     rms_norm_eps: float = 1.0e-6
     vocab_size: int = 262144
     num_key_value_heads: int = 1
-    rope_global_base_freq: float = 1_000_000.0
+    rope_theta: float = 1_000_000.0
     rope_local_base_freq: float = 10_000.0
-    rope_traditional: bool = False
     query_pre_attn_scalar: float = 256
     sliding_window: int = 512
     sliding_window_pattern: int = 6
+    max_position_embeddings: int = 32768
+    rope_scaling: Dict = None
 
 
 class Attention(nn.Module):
@@ -52,14 +54,12 @@ class Attention(nn.Module):
         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 = nn.RoPE(
-            head_dim,
-            traditional=args.rope_traditional,
-            base=(
-                args.rope_local_base_freq
-                if self.is_sliding
-                else args.rope_global_base_freq
-            ),
+        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,
         )
 
     def __call__(
@@ -187,11 +187,14 @@ class Gemma3Model(nn.Module):
 
         global_mask = create_attention_mask(h, cache[self.sliding_window_pattern - 1])
 
-        sliding_window_mask = create_attention_mask(
-            h,
-            cache[0],
-            window_size=self.window_size,
-        )
+        if self.sliding_window_pattern > 1:
+            sliding_window_mask = create_attention_mask(
+                h,
+                cache[0],
+                window_size=self.window_size,
+            )
+        else:
+            sliding_window_mask = None
         for i, (layer, c) in enumerate(zip(self.layers, cache)):
             is_global = (
                 i % self.sliding_window_pattern == self.sliding_window_pattern - 1

mlx_lm/server.py

@@ -631,6 +631,7 @@ class ResponseGenerator:
                         current_sampling = None
                         current_tokenizer = None
                         current_model_key = None
+                        batch_generator.close()
                         batch_generator = None
                         drain_batch = False
                     continue