Fix the rope mutation in a more natural way

Co-authored-by: N8 <n8@n8programs.com>
Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>

.github/workflows/pull_request.yml

@@ -40,5 +40,5 @@ jobs:
         run: |
           curl -o test_data.zip -L https://github.com/ml-explore/mlx-lm/releases/download/test_data/test_data.zip
           unzip test_data.zip
-          HF_HOME="." python -m xmlrunner discover -v tests -o test-results/
+          METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 HF_HOME="." python -m xmlrunner discover -v tests -o test-results/
           mlx.launch -n 2 tests/model_parallel_tests.py

mlx_lm/SERVER.md

@@ -72,12 +72,24 @@ curl localhost:8080/v1/chat/completions \
 - `min_p`: (Optional) A float specifying the min-p sampling parameter.
   Defaults to `0.0` (disabled).
 
-- `repetition_penalty`: (Optional) Applies a penalty to repeated tokens.
-  Defaults to `1.0`.
+- `repetition_penalty`: (Optional) Applies a multiplicative penalty to repeated
+  tokens. Defaults to `0.0` (disabled).
 
 - `repetition_context_size`: (Optional) The size of the context window for
   applying repetition penalty. Defaults to `20`.
 
+- `presence_penalty`: (Optional) Applies an additive penalty to tokens
+  that appeared before. Defaults to `0.0` (disabled).
+
+- `presence_context_size`: (Optional) The size of the context window for
+  applying presence penalty. Defaults to `20`.
+
+- `frequency_penalty`: (Optional) Applies an additive penalty proportional to
+  how many times a token appeared previously. Defaults to `0.0` (disabled).
+
+- `frequency_context_size`: (Optional) The size of the context window for
+  applying frequency penalty. Defaults to `20`.
+
 - `logit_bias`: (Optional) A dictionary mapping token IDs to their bias
   values. Defaults to `None`.
 

mlx_lm/_version.py

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

mlx_lm/benchmark.py

@@ -1,6 +1,7 @@
 # Copyright © 2025 Apple Inc.
 
 import argparse
+import time
 
 import mlx.core as mx
 
@@ -60,6 +61,18 @@ def setup_arg_parser():
         action="store_true",
         help="Quantize activations using the same quantization config as the corresponding layer.",
     )
+    parser.add_argument(
+        "--prefill-step-size",
+        type=int,
+        default=2048,
+        help="Step size for prefill processing (default: 2048)",
+    )
+    parser.add_argument(
+        "--delay",
+        type=int,
+        default=0,
+        help="Delay between each test in seconds (default: 0)",
+    )
     return parser
 
 
@@ -103,14 +116,22 @@ def main():
 
     def single_bench():
         for response in stream_generate(
-            model, tokenizer, prompt, max_tokens=generation_tokens
+            model,
+            tokenizer,
+            prompt,
+            max_tokens=generation_tokens,
+            prefill_step_size=args.prefill_step_size,
         ):
             pass
         return response
 
     def batch_bench():
         return batch_generate(
-            model, tokenizer, prompts, max_tokens=generation_tokens
+            model,
+            tokenizer,
+            prompts,
+            max_tokens=generation_tokens,
+            prefill_step_size=args.prefill_step_size,
         ).stats
 
     if batch_size == 1:
@@ -125,6 +146,8 @@ def main():
     rprint(f"Timing with {prompt_tokens=}, {generation_tokens=}, {batch_size=}.")
     responses = []
     for i in range(args.num_trials):
+        if args.delay > 0:
+            time.sleep(args.delay)
         response = _bench()
         responses.append(response)
         results = [(k, getattr(response, k)) for k in report_keys]

mlx_lm/cli.py

@@ -22,6 +22,7 @@ def main():
         "gptq",
         "server",
         "upload",
+        "share",
     )
     subpackages = {
         "awq": "quant",

mlx_lm/convert.py

@@ -72,7 +72,7 @@ def mixed_quant_predicate_builder(
         if "lm_head" in path:
             return {"group_size": group_size, "bits": high_bits, "mode": mode}
 
-        return {"group_size": group_size, "bits": low_bits}
+        return {"group_size": group_size, "bits": low_bits, "mode": mode}
 
     return mixed_quant_predicate
 
@@ -86,8 +86,8 @@ def convert(
     hf_path: str,
     mlx_path: str = "mlx_model",
     quantize: bool = False,
-    q_group_size: int = 64,
-    q_bits: int = 4,
+    q_group_size: Optional[int] = None,
+    q_bits: Optional[int] = None,
     q_mode: str = "affine",
     dtype: Optional[str] = None,
     upload_repo: str = None,
@@ -128,6 +128,8 @@ def convert(
 
     if dtype is None:
         dtype = config.get("torch_dtype", None)
+    if dtype is None and (text_config := config.get("text_config", None)):
+        dtype = text_config.get("dtype", None)
     if dtype in MODEL_CONVERSION_DTYPES:
         print("[INFO] Using dtype:", dtype)
         dtype = getattr(mx, dtype)

mlx_lm/generate.py

@@ -927,6 +927,11 @@ def _merge_caches(caches):
     return batch_cache
 
 
+def _lazy_extract_cache(cache, i):
+    # Generators like lambdas are late bound so we can't just use it in the loop
+    return (c.extract(i) for c in cache)
+
+
 class BatchGenerator:
     @dataclass
     class Response:
@@ -948,6 +953,9 @@ class BatchGenerator:
         completion_batch_size: int = 32,
         prefill_batch_size: int = 8,
         prefill_step_size: int = 2048,
+        prompt_checkpoint_callback: Optional[
+            Callable[[List[Tuple[int, int, List[Any]]]], None]
+        ] = None,
         prompt_progress_callback: Optional[
             Callable[[List[Tuple[int, int, int]]], None]
         ] = None,
@@ -963,8 +971,10 @@ class BatchGenerator:
         self.prefill_step_size = prefill_step_size
         self.prefill_batch_size = prefill_batch_size
         self.completion_batch_size = max(completion_batch_size, prefill_batch_size)
+        self.prompt_checkpoint_callback = prompt_checkpoint_callback
         self.prompt_progress_callback = prompt_progress_callback or (lambda *_: None)
         self._stats = BatchStats()
+        self._next_count = 0
         self.max_kv_size = max_kv_size
 
         self.active_batch = None
@@ -992,12 +1002,16 @@ class BatchGenerator:
         caches=None,
         samplers: list | None = None,
         logits_processors: list | None = None,
+        prompt_checkpoints: list | int | None = None,
     ):
         uids = []
 
         if max_tokens is None or isinstance(max_tokens, int):
             max_tokens = [max_tokens or self.max_tokens] * len(prompts)
 
+        if prompt_checkpoints is None or isinstance(prompt_checkpoints, int):
+            prompt_checkpoints = [prompt_checkpoints or -1] * len(prompts)
+
         if caches is None:
             caches = [None] * len(prompts)
         for i in range(len(prompts)):
@@ -1007,10 +1021,10 @@ class BatchGenerator:
         samplers = samplers or [None] * len(prompts)
         logits_processors = logits_processors or [self.logits_processors] * len(prompts)
 
-        for p, m, c, s, lp in zip(
-            prompts, max_tokens, caches, samplers, logits_processors
+        for p, m, c, s, lp, pc in zip(
+            prompts, max_tokens, caches, samplers, logits_processors, prompt_checkpoints
         ):
-            self.unprocessed_prompts.append((self.uid_count, p, m, c, s, lp))
+            self.unprocessed_prompts.append((self.uid_count, p, m, c, s, lp, pc))
             uids.append(self.uid_count)
             self.uid_count += 1
         # Sort in ascending order of length
@@ -1043,13 +1057,36 @@ class BatchGenerator:
         if return_prompt_caches:
             return caches
 
+    @property
+    def prompt_cache_nbytes(self):
+        total = sum(c.nbytes for p in self.unprocessed_prompts for c in p[3])
+        if self.active_batch is not None:
+            total += sum(c.nbytes for c in self.active_batch.cache)
+        return total
+
     def _process_prompts(self, prompts):
-        uids, inputs, max_tokens, caches, samplers, logits_processors = zip(*prompts)
+        (
+            uids,
+            inputs,
+            max_tokens,
+            caches,
+            samplers,
+            logits_processors,
+            prompt_checkpoints,
+        ) = zip(*prompts)
 
         lengths = [len(p) for p in inputs]
         max_length = max(lengths)
         padding = [max_length - l for l in lengths]
 
+        # Get the checkpoint token as an offset from the end of each prompt.
+        # Then select the largest one so that we perform the checkpoint at
+        # least `pc` before the end.
+        prompt_checkpoints = [
+            (l - pc if pc > 0 else -pc) for l, pc in zip(lengths, prompt_checkpoints)
+        ]
+        prompt_checkpoint = max(1, max(prompt_checkpoints))
+
         self._stats.prompt_tokens += sum(lengths)
 
         tokens = [mx.array(inp) for inp in inputs]
@@ -1062,8 +1099,10 @@ class BatchGenerator:
             inputs = _left_pad_prompts(inputs, max_length=max_length)
             prompt_cache = _make_cache(self.model, padding, self.max_kv_size)
 
-            while inputs.shape[1] > 1:
-                n_to_process = min(self.prefill_step_size, inputs.shape[1] - 1)
+            while inputs.shape[1] > prompt_checkpoint:
+                n_to_process = min(
+                    self.prefill_step_size, inputs.shape[1] - prompt_checkpoint
+                )
                 self.model(inputs[:, :n_to_process], cache=prompt_cache)
                 mx.eval([c.state for c in prompt_cache])
                 inputs = inputs[:, n_to_process:]
@@ -1074,6 +1113,7 @@ class BatchGenerator:
                         for uid, length in zip(uids, lengths)
                     ]
                 )
+                mx.clear_cache()
 
         # Further prompt processing so we need to
         #   1. Merge the KV caches and prepare for right padded prompts
@@ -1081,16 +1121,22 @@ class BatchGenerator:
         #   2. Process
         #   3. Finalize the KV caches so they are left padded again
         else:
-            last_inputs = mx.array([p[-1:] for p in inputs])
+            last_inputs = mx.array([p[-prompt_checkpoint:] for p in inputs])
             inputs = _right_pad_prompts(inputs, max_length=max_length)
             prompt_cache = _merge_caches(caches)
 
             for c in prompt_cache:
-                # subtract one from lengths since we don't process the last token during prefill
-                c.prepare(lengths=[l - 1 for l in lengths], right_padding=padding)
+                # subtract from lengths since we don't process the last
+                # `prompt_checkpoint` tokens during prefill
+                c.prepare(
+                    lengths=[l - prompt_checkpoint for l in lengths],
+                    right_padding=padding,
+                )
 
-            while inputs.shape[1] > 1:
-                n_to_process = min(self.prefill_step_size, inputs.shape[1] - 1)
+            while inputs.shape[1] > prompt_checkpoint:
+                n_to_process = min(
+                    self.prefill_step_size, inputs.shape[1] - prompt_checkpoint
+                )
                 self.model(inputs[:, :n_to_process], cache=prompt_cache)
                 mx.eval([c.state for c in prompt_cache])
                 inputs = inputs[:, n_to_process:]
@@ -1108,6 +1154,20 @@ class BatchGenerator:
 
         for c in prompt_cache:
             c.finalize()
+
+        # We processed L - prompt_checkpoint tokens so call the checkpoint
+        # callback.
+        if self.prompt_checkpoint_callback is not None:
+            self.prompt_checkpoint_callback(
+                [
+                    (uid, prompt_checkpoint, _lazy_extract_cache(prompt_cache, i))
+                    for i, uid in enumerate(uids)
+                ]
+            )
+        # Process the remaining prompt_checkpoint-1 tokens
+        if prompt_checkpoint > 1:
+            self.model(inputs[:, : prompt_checkpoint - 1], cache=prompt_cache)
+            mx.eval([c.state for c in prompt_cache])
         mx.clear_cache()
 
         y, logprobs = self._step(
@@ -1220,7 +1280,7 @@ class BatchGenerator:
             batch.tokens,
         )
 
-        mx.async_eval(batch.y, batch.logprobs)
+        mx.async_eval(batch.y, batch.logprobs, batch.tokens)
 
         y = y.tolist()
         toc = time.perf_counter()
@@ -1258,6 +1318,9 @@ class BatchGenerator:
             else:
                 self.active_batch = None
 
+        self._next_count += 1
+        if self._next_count % 512 == 0:
+            mx.clear_cache()
         self._stats.generation_tokens += len(responses)
         return responses
 

mlx_lm/lora.py

@@ -21,7 +21,7 @@ from .tuner.utils import (
     load_adapters,
     print_trainable_parameters,
 )
-from .utils import load, save_config
+from .utils import _parse_size, load, save_config
 
 yaml_loader = yaml.SafeLoader
 yaml_loader.add_implicit_resolver(
@@ -69,6 +69,7 @@ CONFIG_DEFAULTS = {
     "config": None,
     "grad_checkpoint": False,
     "grad_accumulation_steps": 1,
+    "clear_cache_threshold": 0,
     "lr_schedule": None,
     "lora_parameters": {"rank": 8, "dropout": 0.0, "scale": 20.0},
     "mask_prompt": False,
@@ -190,6 +191,12 @@ def build_parser():
         help="Use gradient checkpointing to reduce memory use.",
         default=None,
     )
+    parser.add_argument(
+        "--clear-cache-threshold",
+        type=_parse_size,
+        default=0,
+        help="Clear the allocator cache between steps if it grows too large.",
+    )
     parser.add_argument(
         "--report-to",
         type=str,

mlx_lm/models/cache.py

@@ -153,6 +153,11 @@ class _BaseCache:
         """
         return 0
 
+    @property
+    def nbytes(self):
+        """Return the size of this cache in bytes"""
+        raise NotImplementedError("Cache sub-class must implement nbytes")
+
     def empty(self):
         """
         Return if the cache is empty or not.
@@ -215,6 +220,12 @@ class ConcatenateKVCache(_BaseCache):
     def empty(self):
         return self.keys is None
 
+    @property
+    def nbytes(self):
+        if self.keys is None:
+            return 0
+        return self.keys.nbytes + self.values.nbytes
+
 
 class QuantizedKVCache(_BaseCache):
     step = 256
@@ -304,6 +315,10 @@ class QuantizedKVCache(_BaseCache):
     def empty(self):
         return self.keys is None
 
+    @property
+    def nbytes(self):
+        return tree_reduce(lambda a, x: a + x.nbytes, (self.keys, self.values), 0)
+
 
 class KVCache(_BaseCache):
     step = 256
@@ -383,6 +398,12 @@ class KVCache(_BaseCache):
     def empty(self):
         return self.keys is None
 
+    @property
+    def nbytes(self):
+        if self.keys is None:
+            return 0
+        return self.keys.nbytes + self.values.nbytes
+
 
 class RotatingKVCache(_BaseCache):
     step = 256
@@ -561,6 +582,12 @@ class RotatingKVCache(_BaseCache):
     def empty(self):
         return self.keys is None
 
+    @property
+    def nbytes(self):
+        if self.keys is None:
+            return 0
+        return self.keys.nbytes + self.values.nbytes
+
 
 class ArraysCache(_BaseCache):
     def __new__(cls, *args, **kwargs):
@@ -647,6 +674,10 @@ class ArraysCache(_BaseCache):
     def empty(self):
         return self.cache[0] is None
 
+    @property
+    def nbytes(self):
+        return sum(c.nbytes for c in self.cache if c is not None)
+
 
 class ChunkedKVCache(_BaseCache):
     step = 256
@@ -724,6 +755,12 @@ class ChunkedKVCache(_BaseCache):
     def empty(self):
         return self.keys is None
 
+    @property
+    def nbytes(self):
+        if self.keys is None:
+            return 0
+        return self.keys.nbytes + self.values.nbytes
+
 
 class CacheList(_BaseCache):
     def __init__(self, *caches):
@@ -742,16 +779,24 @@ class CacheList(_BaseCache):
 
     @property
     def state(self):
-        return [s for c in self.caches for s in c.state]
+        return [c.state for c in self.caches]
 
     @state.setter
     def state(self, v):
-        state_lens = [len(c.state) for c in self.caches]
-        start = 0
-        for c in self.caches:
-            l = len(c.state)
-            c.state = v[start : start + l]
-            start += l
+        for c, s in zip(self.caches, v):
+            c.state = s
+
+    @property
+    def meta_state(self):
+        return (
+            [type(c).__name__ for c in self.caches],
+            [c.meta_state for c in self.caches],
+        )
+
+    @meta_state.setter
+    def meta_state(self, v):
+        for c, m in zip(self.caches, v[1]):
+            c.meta_state = m
 
     def filter(self, batch_indices):
         """
@@ -793,6 +838,18 @@ class CacheList(_BaseCache):
     def empty(self):
         return self.caches[0].empty()
 
+    @property
+    def nbytes(self):
+        return sum(c.nbytes for c in self.caches)
+
+    @classmethod
+    def from_state(cls, state, meta_state):
+        obj = cls.__new__(cls)
+        obj.caches = [
+            globals()[c].from_state(s, m) for s, c, m in zip(state, *meta_state)
+        ]
+        return obj
+
 
 def dynamic_roll(x, shifts, axis):
     n = x.shape[axis]
@@ -991,6 +1048,12 @@ class BatchKVCache(_BaseCache):
     def empty(self):
         return self.keys is None
 
+    @property
+    def nbytes(self):
+        if self.keys is None:
+            return 0
+        return self.keys.nbytes + self.values.nbytes
+
 
 class BatchRotatingKVCache(_BaseCache):
     step = 256
@@ -1061,6 +1124,10 @@ class BatchRotatingKVCache(_BaseCache):
         self.offset += keys.shape[2]
         self._offset += keys.shape[2]
         self._idx = self.keys.shape[2]
+
+        # Make sure left_padding and offset are evaluated
+        self.keys = mx.depends(self.keys, (self.left_padding, self.offset))
+
         return self.keys, self.values
 
     def _update_in_place(self, keys, values):
@@ -1111,6 +1178,9 @@ class BatchRotatingKVCache(_BaseCache):
         self.offset += S
         self._idx += S
 
+        # Make sure left_padding and offset are evaluated
+        self.keys = mx.depends(self.keys, (self.left_padding, self.offset))
+
         # If the buffer is not full, slice off the end
         if self._offset < self.max_size:
             return (
@@ -1305,3 +1375,9 @@ class BatchRotatingKVCache(_BaseCache):
 
     def empty(self):
         return self.keys is None
+
+    @property
+    def nbytes(self):
+        if self.keys is None:
+            return 0
+        return self.keys.nbytes + self.values.nbytes

mlx_lm/models/gated_delta.py

@@ -7,9 +7,7 @@ import mlx.nn as nn
 
 @partial(mx.compile, shapeless=True)
 def compute_g(A_log, a, dt_bias):
-    return mx.exp(-mx.exp(A_log.astype(mx.float32)) * nn.softplus(a + dt_bias)).astype(
-        A_log.dtype
-    )
+    return mx.exp(-mx.exp(A_log.astype(mx.float32)) * nn.softplus(a + dt_bias))
 
 
 def _make_gated_delta_kernel(has_mask=False, vectorized=False):
@@ -94,7 +92,7 @@ def _make_gated_delta_kernel(has_mask=False, vectorized=False):
         }}
         for (int i = 0; i < n_per_t; ++i) {{
           auto s_idx = n_per_t * dk_idx + i;
-          o_state[s_idx] = static_cast<InT>(state[i]);
+          o_state[s_idx] = static_cast<StT>(state[i]);
         }}
     """
     inputs = ["q", "k", "v", "g", "beta", "state_in", "T"]
@@ -165,7 +163,7 @@ def _gated_delta_step_ops(
     if mask is not None:
         mask = mx.expand_dims(mask, axis=(1, 2, 3))
         state = mx.where(mask, state, old_state)
-    return y, state
+    return y.astype(q.dtype), state
 
 
 def gated_delta_kernel(
@@ -180,6 +178,7 @@ def gated_delta_kernel(
     B, T, Hk, Dk = k.shape
     Hv, Dv = v.shape[2:]
     input_type = q.dtype
+    state_type = state.dtype
     if g.ndim == 4:
         kernel = _gated_delta_kernel_vec
         inputs = [q, k, v, g, beta, state, T]
@@ -197,6 +196,7 @@ def gated_delta_kernel(
         inputs=inputs,
         template=[
             ("InT", input_type),
+            ("StT", state_type),
             ("Dk", Dk),
             ("Dv", Dv),
             ("Hk", Hk),
@@ -205,7 +205,7 @@ def gated_delta_kernel(
         grid=(32, Dv, B * Hv),
         threadgroup=(32, 4, 1),
         output_shapes=[(B, T, Hv, Dv), state.shape],
-        output_dtypes=[input_type, input_type],
+        output_dtypes=[input_type, state_type],
     )
 
 
@@ -235,7 +235,7 @@ def gated_delta_ops(
     B, T, Hk, Dk = q.shape
     Hv, Dv = v.shape[-2:]
     if state is None:
-        state = mx.zeros((B, Hv, Dv, Dk), dtype=q.dtype)
+        state = mx.zeros((B, Hv, Dv, Dk), dtype=mx.float32)
 
     if (repeat_factor := Hv // Hk) > 1:
         q = mx.repeat(q, repeat_factor, -2)
@@ -269,13 +269,12 @@ def gated_delta_update(
     mask: Optional[mx.array] = None,
     use_kernel: bool = True,
 ) -> Tuple[mx.array, mx.array]:
-
     beta = mx.sigmoid(b)
     g = compute_g(A_log, a, dt_bias)
     if state is None:
         B, _, Hk, Dk = q.shape
         Hv, Dv = v.shape[-2:]
-        state = mx.zeros((B, Hv, Dv, Dk), dtype=q.dtype)
+        state = mx.zeros((B, Hv, Dv, Dk), dtype=mx.float32)
 
     if not use_kernel or mx.default_device() != mx.gpu or not mx.metal.is_available():
         return gated_delta_ops(q, k, v, g, beta, state, mask)

mlx_lm/models/lfm2.py

@@ -32,11 +32,14 @@ class ModelArgs(BaseModelArgs):
     block_multiple_of: int
     block_ffn_dim_multiplier: float
     block_auto_adjust_ff_dim: bool
-    rope_theta: float
+    rope_theta: float = 1000000.0
+    rope_parameters: Optional[dict] = None
     full_attn_idxs: Optional[List[int]] = None
     layer_types: Optional[List[str]] = None
 
     def __post_init__(self):
+        if self.rope_parameters is not None and "rope_theta" in self.rope_parameters:
+            self.rope_theta = self.rope_parameters["rope_theta"]
         if self.num_key_value_heads is None:
             self.num_key_value_heads = self.num_attention_heads
         if self.full_attn_idxs is None:

mlx_lm/models/lfm2_moe.py

@@ -35,11 +35,14 @@ class ModelArgs(BaseModelArgs):
     norm_eps: float
     conv_bias: bool
     conv_L_cache: int
-    rope_theta: float
+    rope_theta: float = 1000000.0
+    rope_parameters: Optional[dict] = None
     full_attn_idxs: Optional[List[int]] = None
     layer_types: Optional[List[str]] = None
 
     def __post_init__(self):
+        if self.rope_parameters is not None and "rope_theta" in self.rope_parameters:
+            self.rope_theta = self.rope_parameters["rope_theta"]
         if self.full_attn_idxs is None:
             self.full_attn_idxs = [
                 i

mlx_lm/models/minimax.py

@@ -1,6 +1,7 @@
 # Copyright © 2025 Apple Inc.
 
 from dataclasses import dataclass
+from functools import lru_cache
 from typing import Any, List, Optional
 
 import mlx.core as mx
@@ -33,6 +34,55 @@ class ModelArgs(BaseModelArgs):
     use_qk_norm: bool = True
 
 
+@lru_cache
+def sharded_rms_norm(group):
+    @mx.compile
+    def _cast_square_sum(x):
+        return x.astype(mx.float32).square().sum(-1, keepdims=True)
+
+    @mx.compile
+    def _normalize(x, norm2, w, eps):
+        norm2 = mx.distributed.all_sum(norm2, group=group)
+        norm = mx.rsqrt(norm2 / (x.shape[-1] * group.size()) + eps)
+        return (x.astype(mx.float32) * norm * w).astype(x.dtype)
+
+    # Split the compile so that x upcasting doesn't break the compile and we
+    # have 2 kernels generated 1 for f(x) = square(upcast(x)) and another
+    # g(x) = downcast(upcast(x) * norm * w)
+    def _inner_sharded_rms_norm(x, w, eps):
+        return _normalize(x, _cast_square_sum(x), w, eps)
+
+    return _inner_sharded_rms_norm
+
+
+class ShardedRMSNorm(nn.Module):
+    def __init__(
+        self, dims: int, eps: float = 1e-5, group: Optional[mx.distributed.Group] = None
+    ):
+        super().__init__()
+        group = group or mx.distributed.init()
+        self.weight = mx.ones((dims // group.size(),))
+        self.group = group
+        self.eps = eps
+
+    def _extra_repr(self):
+        return f"{self.weight.shape[0] * self.group.size()}, eps={self.eps}"
+
+    def __call__(self, x):
+        return sharded_rms_norm(self.group)(x, self["weight"], self.eps)
+
+    @classmethod
+    def from_rms_norm(
+        cls, norm_module, *, group: Optional[mx.distributed.Group] = None
+    ):
+        sn = cls(norm_module.weight.shape[0], norm_module.eps, group=group)
+        sn.weight = mx.contiguous(
+            mx.split(norm_module.weight, group.size(), axis=-1)[group.rank()]
+        )
+
+        return sn
+
+
 class MiniMaxAttention(nn.Module):
     def __init__(self, args: ModelArgs):
         super().__init__()
@@ -295,12 +345,12 @@ class Model(nn.Module):
                 layer.self_attn.o_proj, "sharded-to-all", group=group
             )
             if layer.self_attn.use_qk_norm:
-                layer.self_attn.q_norm.weight = layer.self_attn.q_norm.weight.split(
-                    N, axis=-1
-                )[rank]
-                layer.self_attn.k_norm.weight = layer.self_attn.k_norm.weight.split(
-                    N, axis=-1
-                )[rank]
+                layer.self_attn.q_norm = ShardedRMSNorm.from_rms_norm(
+                    layer.self_attn.q_norm, group=group
+                )
+                layer.self_attn.k_norm = ShardedRMSNorm.from_rms_norm(
+                    layer.self_attn.k_norm, group=group
+                )
 
             layer.self_attn.num_attention_heads //= N
             layer.self_attn.num_key_value_heads //= N

mlx_lm/models/mixtral.py

@@ -1,4 +1,4 @@
-# Copyright © 2023-2024 Apple Inc.
+# Copyright © 2026 Apple Inc.
 
 from dataclasses import dataclass
 from typing import Any, Dict, Optional, Union
@@ -25,6 +25,7 @@ class ModelArgs(BaseModelArgs):
     rope_theta: float = 1e6
     rope_traditional: bool = False
     rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    tie_word_embeddings: bool = False
 
     def __post_init__(self):
         if self.num_key_value_heads is None:
@@ -162,8 +163,12 @@ class MixtralModel(nn.Module):
         self,
         inputs: mx.array,
         cache=None,
-    ):
-        h = self.embed_tokens(inputs)
+        input_embeddings: Optional[mx.array] = None,
+    ) -> mx.array:
+        if input_embeddings is not None:
+            h = input_embeddings
+        else:
+            h = self.embed_tokens(inputs)
 
         if cache is None:
             cache = [None] * len(self.layers)
@@ -179,20 +184,27 @@ class MixtralModel(nn.Module):
 class Model(nn.Module):
     def __init__(self, args: ModelArgs):
         super().__init__()
+        self.args = args
         self.model_type = args.model_type
         self.model = MixtralModel(args)
-        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
-        self.args = args
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
 
     def __call__(
         self,
         inputs: mx.array,
         cache=None,
-    ):
-        out = self.model(inputs, cache)
-        return self.lm_head(out)
+        input_embeddings: Optional[mx.array] = None,
+    ) -> mx.array:
+        out = self.model(inputs, cache, input_embeddings)
+        if self.args.tie_word_embeddings:
+            return self.model.embed_tokens.as_linear(out)
+        else:
+            return self.lm_head(out)
 
     def sanitize(self, weights):
+        if self.args.tie_word_embeddings:
+            weights.pop("lm_head.weight", None)
         if "model.layers.0.block_sparse_moe.experts.0.w1.weight" not in weights:
             return weights
         for l in range(self.args.num_hidden_layers):

mlx_lm/models/mla.py

@@ -26,7 +26,7 @@ class MultiLinear(nn.Module):
         self,
         group_size: int,
         bits: int,
-        mode: str,
+        mode: str = "affine",
     ):
         num_heads, output_dims, input_dims = self.weight.shape
         ql = QuantizedMultiLinear(

mlx_lm/models/nemotron_h.py

@@ -40,10 +40,12 @@ class ModelArgs(BaseModelArgs):
     layer_norm_epsilon: float
     use_bias: bool
     use_conv_bias: bool
-    hybrid_override_pattern: List[str]
+    hybrid_override_pattern: Optional[List[str]] = None
+    layers_block_type: Optional[List[str]] = None
     head_dim: Optional[int] = None
     moe_intermediate_size: Optional[int] = None
     moe_shared_expert_intermediate_size: Optional[int] = None
+    moe_latent_size: Optional[int] = None
     n_group: Optional[int] = None
     n_routed_experts: Optional[int] = None
     n_shared_experts: Optional[int] = None
@@ -55,13 +57,20 @@ class ModelArgs(BaseModelArgs):
     time_step_min: Optional[float] = None
     time_step_max: Optional[float] = None
 
+    # Map from layers_block_type names to single-char pattern codes
+    _block_type_to_char = {"mamba": "M", "attention": "*", "moe": "E", "mlp": "-"}
+
     def __post_init__(self):
-        if (
-            self.time_step_limit is None
-            and self.time_step_min is not None
-            and self.time_step_max is not None
-        ):
-            self.time_step_limit = (self.time_step_min, self.time_step_max)
+        if self.time_step_limit is None and self.time_step_min is not None:
+            self.time_step_limit = (self.time_step_min, float("inf"))
+
+        # Normalize to hybrid_override_pattern (single-char list)
+        if self.hybrid_override_pattern is None and self.layers_block_type is not None:
+            self.hybrid_override_pattern = [
+                self._block_type_to_char[t] for t in self.layers_block_type
+            ]
+        if self.hybrid_override_pattern is not None:
+            self.num_hidden_layers = len(self.hybrid_override_pattern)
 
 
 class MambaRMSNormGated(nn.Module):
@@ -365,8 +374,16 @@ class NemotronHMoE(nn.Module):
         super().__init__()
         self.config = config
         self.num_experts_per_tok = config.num_experts_per_tok
+        self.moe_latent_size = config.moe_latent_size
+
+        # When latent projection is used, experts operate on the latent dim
+        expert_input_dim = (
+            config.moe_latent_size
+            if config.moe_latent_size is not None
+            else config.hidden_size
+        )
         self.switch_mlp = SwitchMLP(
-            config.hidden_size,
+            expert_input_dim,
             config.moe_intermediate_size,
             config.n_routed_experts,
             activation=nn.ReLU2(),
@@ -379,12 +396,30 @@ class NemotronHMoE(nn.Module):
                 config, intermediate_size=intermediate_size
             )
 
+        # Latent projection layers for dimensionality reduction before/after experts
+        if config.moe_latent_size is not None:
+            self.fc1_latent_proj = nn.Linear(
+                config.hidden_size, config.moe_latent_size, bias=config.mlp_bias
+            )
+            self.fc2_latent_proj = nn.Linear(
+                config.moe_latent_size, config.hidden_size, bias=config.mlp_bias
+            )
+
     def __call__(self, x):
+        residuals = x
         inds, scores = self.gate(x)
+
+        if self.moe_latent_size is not None:
+            x = self.fc1_latent_proj(x)
+
         y = self.switch_mlp(x, inds)
         y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype)
+
+        if self.moe_latent_size is not None:
+            y = self.fc2_latent_proj(y)
+
         if self.config.n_shared_experts is not None:
-            y = y + self.shared_experts(x)
+            y = y + self.shared_experts(residuals)
 
         return y
 
@@ -501,6 +536,7 @@ class Model(nn.Module):
         return caches
 
     def sanitize(self, weights):
+        weights = {k: v for (k, v) in weights.items() if not k.startswith("mtp.")}
         for k, v in weights.items():
             if "conv1d.weight" in k and v.shape[-1] != 1:
                 weights[k] = v.moveaxis(2, 1)

mlx_lm/models/qwen3_5.py

@@ -5,7 +5,8 @@ from typing import Any, Dict, List, Optional, Union
 
 import mlx.core as mx
 import mlx.nn as nn
-from mlx.utils import tree_flatten, tree_unflatten
+from mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients
+from mlx.utils import tree_map
 
 from .base import (
     BaseModelArgs,
@@ -126,6 +127,8 @@ class GatedDeltaNet(nn.Module):
 
         self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)
 
+        self.sharding_group = None
+
     def __call__(
         self,
         inputs: mx.array,
@@ -134,6 +137,9 @@ class GatedDeltaNet(nn.Module):
     ) -> mx.array:
         B, S, _ = inputs.shape
 
+        if self.sharding_group is not None:
+            inputs = sum_gradients(self.sharding_group)(inputs)
+
         qkv = self.in_proj_qkv(inputs)
         z = self.in_proj_z(inputs).reshape(B, S, self.num_v_heads, self.head_v_dim)
         b = self.in_proj_b(inputs)
@@ -185,7 +191,12 @@ class GatedDeltaNet(nn.Module):
             cache[1] = state
 
         out = self.norm(out, z)
-        return self.out_proj(out.reshape(B, S, -1))
+        out = self.out_proj(out.reshape(B, S, -1))
+
+        if self.sharding_group is not None:
+            out = mx.distributed.all_sum(out, group=self.sharding_group)
+
+        return out
 
 
 class DecoderLayer(nn.Module):
@@ -324,6 +335,15 @@ class TextModel(nn.Module):
 
         return predicate
 
+    @property
+    def cast_predicate(self):
+        def predicate(path: str):
+            if path.endswith("A_log"):
+                return False
+            return True
+
+        return predicate
+
 
 @dataclass
 class ModelArgs(BaseModelArgs):
@@ -355,11 +375,10 @@ class Model(nn.Module):
         )
 
     def sanitize(self, weights):
-        weights = tree_unflatten(list(weights.items()))
-        weights = dict(tree_flatten(weights))
-
         sanitized = {}
         for key, value in weights.items():
+            if key.startswith("vision_tower") or key.startswith("model.visual"):
+                continue
             if key.startswith("model.visual"):
                 continue
             if key.startswith("model.language_model"):
@@ -371,6 +390,124 @@ class Model(nn.Module):
             sanitized[key] = value
         return self.language_model.sanitize(sanitized)
 
+    def shard(self, group=None):
+        group = group or mx.distributed.init()
+        N = group.size()
+        rank = group.rank()
+
+        # A sharding factory for the convolution in gated delta net
+        def conv_sharding(key_dim):
+            return lambda p, w: (0, [key_dim, 2 * key_dim])
+
+        def repeat_kv_layer_inplace(layer, h):
+            # No repeat needed cause we have more heads than nodes
+            if N <= h:
+                return
+
+            # Repeat function to apply to the layer weights
+            def _repeat(p):
+                s = p.shape
+                p = p.reshape(h, s[0] // h, *s[1:])
+                p = mx.repeat(p, N // h, axis=0)
+                p = p.reshape(-1, *s[1:])
+                return p
+
+            layer.update(tree_map(_repeat, layer.parameters()))
+
+        for layer in self.layers:
+            # Linear attention
+            if layer.is_linear:
+                kd = layer.linear_attn.key_dim
+                layer.linear_attn.sharding_group = group
+                shard_inplace(layer.linear_attn.conv1d, conv_sharding(kd), group=group)
+                layer.linear_attn.conv1d.groups //= N
+                shard_inplace(
+                    layer.linear_attn.in_proj_qkv,
+                    "all-to-sharded",
+                    segments=[kd, 2 * kd],
+                    group=group,
+                )
+                shard_inplace(
+                    layer.linear_attn.in_proj_z, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.linear_attn.in_proj_b, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.linear_attn.in_proj_a, "all-to-sharded", group=group
+                )
+                layer.linear_attn.dt_bias = mx.contiguous(
+                    mx.split(layer.linear_attn.dt_bias, N)[rank]
+                )
+                layer.linear_attn.A_log = mx.contiguous(
+                    mx.split(layer.linear_attn.A_log, N)[rank]
+                )
+                shard_inplace(layer.linear_attn.out_proj, "sharded-to-all", group=group)
+                layer.linear_attn.num_k_heads //= N
+                layer.linear_attn.num_v_heads //= N
+                layer.linear_attn.key_dim //= N
+                layer.linear_attn.value_dim //= N
+                layer.linear_attn.conv_dim //= N
+
+            # Softmax attention
+            else:
+                layer.self_attn.o_proj = shard_linear(
+                    layer.self_attn.o_proj, "sharded-to-all", group=group
+                )
+                layer.self_attn.q_proj = shard_linear(
+                    layer.self_attn.q_proj, "all-to-sharded", group=group
+                )
+                repeat_kv_layer_inplace(
+                    layer.self_attn.k_proj, layer.self_attn.num_key_value_heads
+                )
+                repeat_kv_layer_inplace(
+                    layer.self_attn.v_proj, layer.self_attn.num_key_value_heads
+                )
+                layer.self_attn.k_proj = shard_linear(
+                    layer.self_attn.k_proj, "all-to-sharded", group=group
+                )
+                layer.self_attn.v_proj = shard_linear(
+                    layer.self_attn.v_proj, "all-to-sharded", group=group
+                )
+                layer.self_attn.num_attention_heads //= N
+                layer.self_attn.num_key_value_heads = max(
+                    1, layer.self_attn.num_key_value_heads // N
+                )
+
+            # MLP
+            if isinstance(layer.mlp, MLP):
+                layer.mlp.gate_proj = shard_linear(
+                    layer.mlp.gate_proj, "all-to-sharded", group=group
+                )
+                layer.mlp.down_proj = shard_linear(
+                    layer.mlp.down_proj, "sharded-to-all", group=group
+                )
+                layer.mlp.up_proj = shard_linear(
+                    layer.mlp.up_proj, "all-to-sharded", group=group
+                )
+
+            # MoE
+            else:
+                layer.mlp.sharding_group = group
+                shard_inplace(
+                    layer.mlp.shared_expert.gate_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.shared_expert.down_proj, "sharded-to-all", group=group
+                )
+                shard_inplace(
+                    layer.mlp.shared_expert.up_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.gate_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.down_proj, "sharded-to-all", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.up_proj, "all-to-sharded", group=group
+                )
+
     @property
     def layers(self):
         return self.language_model.model.layers
@@ -381,3 +518,7 @@ class Model(nn.Module):
     @property
     def quant_predicate(self):
         return self.language_model.quant_predicate
+
+    @property
+    def cast_predicate(self):
+        return self.language_model.cast_predicate

mlx_lm/models/qwen3_5_moe.py

@@ -2,8 +2,6 @@
 
 from dataclasses import dataclass
 
-from mlx.utils import tree_flatten, tree_unflatten
-
 from .base import BaseModelArgs
 from .qwen3_5 import Model as Qwen3_5Model
 
@@ -23,12 +21,9 @@ class ModelArgs(BaseModelArgs):
 class Model(Qwen3_5Model):
 
     def sanitize(self, weights):
-        weights = tree_unflatten(list(weights.items()))
-        weights = dict(tree_flatten(weights))
-
         new_weights = {}
         for key, value in weights.items():
-            if key.startswith("model.visual"):
+            if key.startswith("vision_tower") or key.startswith("model.visual"):
                 continue
             if key.startswith("model.language_model"):
                 key = key.replace("model.language_model", "language_model.model")

mlx_lm/models/qwen3_moe.py

@@ -1,4 +1,4 @@
-# Copyright © 2025 Apple Inc.
+# Copyright © 2026 Apple Inc.
 
 from dataclasses import dataclass
 from typing import Any, Dict, List, Optional, Union
@@ -123,7 +123,7 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
     def __call__(
         self,
         x: mx.array,
-    ):
+    ) -> mx.array:
         gates = self.gate(x)
         gates = mx.softmax(gates, axis=-1, precise=True)
 
@@ -190,7 +190,7 @@ class Qwen3MoeModel(nn.Module):
         inputs: mx.array,
         cache=None,
         input_embeddings: Optional[mx.array] = None,
-    ):
+    ) -> mx.array:
         if input_embeddings is not None:
             h = input_embeddings
         else:
@@ -213,15 +213,25 @@ class Model(nn.Module):
         self.args = args
         self.model_type = args.model_type
         self.model = Qwen3MoeModel(args)
-        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
 
     def __call__(
-        self, inputs: mx.array, cache=None, input_embeddings: Optional[mx.array] = None
-    ):
+        self,
+        inputs: mx.array,
+        cache=None,
+        input_embeddings: Optional[mx.array] = None,
+    ) -> mx.array:
         out = self.model(inputs, cache, input_embeddings)
-        return self.lm_head(out)
+        if self.args.tie_word_embeddings:
+            out = self.model.embed_tokens.as_linear(out)
+        else:
+            out = self.lm_head(out)
+        return out
 
     def sanitize(self, weights):
+        if self.args.tie_word_embeddings:
+            weights.pop("lm_head.weight", None)
         if "model.layers.0.mlp.experts.0.up_proj.weight" not in weights:
             return weights
         for l in range(self.args.num_hidden_layers):

mlx_lm/models/qwen3_next.py

@@ -3,10 +3,12 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from functools import partial
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import mlx.core as mx
 import mlx.nn as nn
+from mlx.nn.layers.distributed import sum_gradients
 
 from .activations import swiglu
 from .base import (
@@ -53,6 +55,13 @@ class ModelArgs(BaseModelArgs):
     full_attention_interval: int = 4
 
 
+@partial(mx.compile, shapeless=True)
+def _precise_swiglu(h, gate, x):
+    gate = nn.silu(gate.astype(mx.float32))
+    x = x.astype(mx.float32)
+    return (gate * x).astype(h.dtype)
+
+
 class Qwen3NextRMSNormGated(nn.Module):
     def __init__(self, hidden_size: int, eps: float = 1e-6):
         super().__init__()
@@ -64,8 +73,9 @@ class Qwen3NextRMSNormGated(nn.Module):
     ) -> mx.array:
         x = mx.fast.rms_norm(hidden_states, self.weight, self.eps)
         if gate is not None:
-            x = swiglu(gate, x)
-        return x
+            return _precise_swiglu(hidden_states, gate, x)
+        else:
+            return x.astype(hidden_states.dtype)
 
 
 class Qwen3NextAttention(nn.Module):
@@ -312,10 +322,15 @@ class Qwen3NextSparseMoeBlock(nn.Module):
         self.shared_expert = Qwen3NextMLP(dim, shared_expert_intermediate_size)
         self.shared_expert_gate = nn.Linear(dim, 1, bias=False)
 
+        self.sharding_group = None
+
     def __call__(
         self,
         x: mx.array,
     ) -> mx.array:
+        if self.sharding_group is not None:
+            x = sum_gradients(self.sharding_group)(x)
+
         gates = self.gate(x)
         gates = mx.softmax(gates, axis=-1, precise=True)
 
@@ -331,7 +346,12 @@ class Qwen3NextSparseMoeBlock(nn.Module):
         shared_y = self.shared_expert(x)
         shared_y = mx.sigmoid(self.shared_expert_gate(x)) * shared_y
 
-        return y + shared_y
+        y = y + shared_y
+
+        if self.sharding_group is not None:
+            y = mx.distributed.all_sum(y, group=self.sharding_group)
+
+        return y
 
 
 class Qwen3NextDecoderLayer(nn.Module):

mlx_lm/models/rope_utils.py

@@ -58,9 +58,8 @@ class SuScaledRoPE(nn.Module):
         self._scale = long_mscale or (1.0 if factor <= 1.0 else default_scale(factor))
 
     def __call__(self, x, offset: Union[int, mx.array] = 0):
-        x[..., : self.dim] = self._scale * x[..., : self.dim]
         return mx.fast.rope(
-            x,
+            x.at[..., : self.dim].multiply(self._scale),
             self.dim,
             traditional=False,
             base=None,
@@ -71,7 +70,6 @@ class SuScaledRoPE(nn.Module):
 
 
 class Llama3RoPE(nn.Module):
-
     def __init__(
         self,
         dims: int,
@@ -183,7 +181,7 @@ class YarnRoPE(nn.Module):
 
     def __call__(self, x, offset=0):
         if self.mscale != 1.0:
-            x[..., : self.dims] = self.mscale * x[..., : self.dims]
+            x = x.at[..., : self.dims].multiply(self.mscale)
         return mx.fast.rope(
             x,
             self.dims,

mlx_lm/models/ssm.py

@@ -6,6 +6,7 @@ import mlx.nn as nn
 
 @mx.compile
 def compute_dt(dt, dt_bias, time_step_limit):
+    dt = dt.astype(mx.float32)
     dt = nn.softplus(dt + dt_bias)
     return mx.clip(dt, time_step_limit[0], time_step_limit[1])
 
@@ -44,7 +45,7 @@ def make_ssm_kernel():
             auto idx = d_idx * Ds + s_idx;
             auto dB_by_x = x_ * dt_ * static_cast<float>(B_[s_idx]);
             auto state = dA * i_state[idx] + dB_by_x;
-            o_state[idx] = static_cast<T>(state);
+            o_state[idx] = static_cast<U>(state);
             acc += state * C_[s_idx];
         }
         acc = simd_sum(acc);
@@ -76,15 +77,23 @@ def ssm_update_kernel(
 ):
     n, _, h, d = hidden_states.shape
     input_type = hidden_states.dtype
+    state_type = state.dtype
     hb, ds = B.shape[-2:]
     dt = compute_dt(dt, dt_bias, time_step_limit)
     return _ssm_kernel(
         inputs=[hidden_states, A_log, B, C, D, dt, state],
-        template=[("T", input_type), ("Dh", d), ("Ds", ds), ("H", h), ("G", h // hb)],
+        template=[
+            ("T", input_type),
+            ("U", state_type),
+            ("Dh", d),
+            ("Ds", ds),
+            ("H", h),
+            ("G", h // hb),
+        ],
         grid=(32, d, h * n),
         threadgroup=(32, 8, 1),
         output_shapes=[(n, 1, h, d), state.shape],
-        output_dtypes=[input_type, input_type],
+        output_dtypes=[input_type, state_type],
     )
 
 
@@ -186,7 +195,7 @@ def ssm_attn(
                 mx.expand_dims(lengths < 0, (1, 2, 3)), state, next_state
             )
 
-        return y, next_state
+        return y.astype(x.dtype), next_state
 
     ys = []
     for i in range(0, l, step):

mlx_lm/models/step3p5.py

@@ -10,7 +10,7 @@ from mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients
 
 from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
-from .cache import KVCache
+from .cache import KVCache, RotatingKVCache
 from .rope_utils import initialize_rope
 from .switch_layers import SwiGLU, SwitchGLU
 
@@ -394,7 +394,14 @@ class Model(nn.Module):
         return self.model.layers
 
     def make_cache(self):
-        return [KVCache() for _ in self.layers]
+        return [
+            (
+                RotatingKVCache(max_size=self.args.sliding_window)
+                if layer.is_sliding
+                else KVCache()
+            )
+            for layer in self.layers
+        ]
 
     def sanitize(self, weights):
         remappings = [

mlx_lm/perplexity.py

@@ -106,6 +106,11 @@ def main():
         required=True,
         help="Path to model or Hugging Face model ID",
     )
+    parser.add_argument(
+        "--trust-remote-code",
+        action="store_true",
+        help="Enable trusting remote code for tokenizer/model loading from Hugging Face.",
+    )
     parser.add_argument(
         "--batch-size", type=int, default=8, help="Batch size for evaluation"
     )
@@ -139,7 +144,8 @@ def main():
 
     # Load model
     print(f"Loading model from {args.model}...")
-    model, tokenizer = load(args.model)
+    tokenizer_config = {"trust_remote_code": True if args.trust_remote_code else None}
+    model, tokenizer = load(args.model, tokenizer_config=tokenizer_config)
 
     # Count parameters
     total_params = get_total_parameters(model)

mlx_lm/quant/dwq.py

@@ -383,7 +383,7 @@ def main():
         del model
 
     if mx.metal.is_available():
-        max_rec_size = mx.metal.device_info()["max_recommended_working_set_size"]
+        max_rec_size = mx.device_info()["max_recommended_working_set_size"]
         mx.set_wired_limit(max_rec_size)
 
     opt = optimizers.Adam(learning_rate=args.learning_rate, bias_correction=True)

mlx_lm/sample_utils.py

@@ -73,15 +73,28 @@ def make_logits_processors(
     logit_bias: Optional[Dict[int, float]] = None,
     repetition_penalty: Optional[float] = None,
     repetition_context_size: Optional[int] = 20,
+    presence_penalty: Optional[float] = None,
+    presence_context_size: Optional[int] = 20,
+    frequency_penalty: Optional[float] = None,
+    frequency_context_size: Optional[int] = 20,
 ):
     """
     Make logits processors for use with ``generate_step``.
 
     Args:
-        repetition_penalty (float, optional): The penalty factor for repeating
-          tokens.
+        repetition_penalty (float, optional): A (sign-aware) multiplicative
+          penalty for repeating tokens.
         repetition_context_size (int, optional): The number of tokens to
           consider for repetition penalty. Default: ``20``.
+        presence_penalty (float, optional): An additive penalty to reduce
+          repeating tokens.
+        presence_context_size (int, optional): The number of tokens to consider
+          for the presence penalty. Default: ``20``.
+        frequency_penalty (float, optional): An additive penalty to reduce
+          repeating tokens. The tokens are penalized proportionally to their
+          frequency.
+        frequency_context_size (int, optional): The number of tokens to consider
+          for the frequency penalty. Default: ``20``.
         logit_bias (dictionary, optional): Additive logit bias.
 
     Returns:
@@ -96,15 +109,20 @@ def make_logits_processors(
         values = mx.array(list(logit_bias.values()))
 
         def logit_bias_processor(_, logits):
-            logits[:, indices] += values
-            return logits
+            return logits.at[:, indices].add(values)
 
         logits_processors.append(logit_bias_processor)
 
-    if repetition_penalty and repetition_penalty != 0.0:
-        logits_processors.append(
-            make_repetition_penalty(repetition_penalty, repetition_context_size)
-        )
+    repetition_penalties = [
+        (make_repetition_penalty, repetition_penalty, repetition_context_size),
+        (make_presence_penalty, presence_penalty, presence_context_size),
+        (make_frequency_penalty, frequency_penalty, frequency_context_size),
+    ]
+
+    for make_penalty, penalty, context_size in repetition_penalties:
+        if penalty is not None and penalty != 0:
+            logits_processors.append(make_penalty(penalty, context_size))
+
     return logits_processors
 
 
@@ -307,3 +325,58 @@ def make_repetition_penalty(penalty: float, context_size: int = 20):
         return logits
 
     return repetition_penalty_processor
+
+
+def make_presence_penalty(penalty: float, context_size: int = 20):
+    """
+    Make a presence penalty processor.
+
+    Corresponds to the OpenAI option with the same name. Namely, subtracts
+    ``penalty`` from a logit if the token has occured at least once in the
+    ``context_size`` previous tokens.
+
+    Args:
+        penalty (float): The presence penalty to be applied.
+        context_size (int): The number of previous tokens to use.
+            Default: ``20``.
+
+    Returns:
+        Callable[[mx.array, List[int]], mx.array]
+    """
+
+    def presence_penalty_processor(tokens, logits):
+        if len(tokens) > 0:
+            tokens = tokens[-context_size:]
+            logits[:, tokens] -= penalty
+        return logits
+
+    return presence_penalty_processor
+
+
+def make_frequency_penalty(penalty: float, context_size: int = 20):
+    """
+    Make a frequency penalty processor.
+
+    Corresponds to the OpenAI option with the same name. Namely, subtracts
+    ``penalty`` from a logit for every time that the token has occured in the
+    ``context_size`` previous tokens.
+
+    The difference with the presence penalty is that the more often a token
+    occurs the more it will be penalized.
+
+    Args:
+        penalty (float): The frequency penalty to be applied.
+        context_size (int): The number of previous tokens to use.
+            Default: ``20``.
+
+    Returns:
+        Callable[[mx.array, List[int]], mx.array]
+    """
+
+    def frequency_penalty_processor(tokens, logits):
+        if len(tokens) > 0:
+            tokens = tokens[-context_size:]
+            logits = logits.at[:, tokens].subtract(penalty)
+        return logits
+
+    return frequency_penalty_processor

mlx_lm/server.py

@@ -2,6 +2,7 @@
 
 import argparse
 import copy
+import heapq
 import json
 import logging
 import pickle
@@ -41,7 +42,7 @@ from .models.cache import (
     trim_prompt_cache,
 )
 from .sample_utils import make_logits_processors, make_sampler
-from .utils import load, sharded_load
+from .utils import _parse_size, load, sharded_load
 
 
 def get_system_fingerprint():
@@ -171,11 +172,34 @@ def process_message_content(messages):
 
 
 class LRUPromptCache:
-
     @dataclass
     class CacheEntry:
         prompt_cache: List[Any]
-        count: int
+        nbytes: int
+
+    class CacheOrder:
+        def __init__(self):
+            self._lru_checkpoints = deque()
+            self._lru = deque()
+
+        def __len__(self):
+            return len(self._lru) + len(self._lru_checkpoints)
+
+        def push(self, model, tokens, checkpoint: bool = False):
+            c = self._lru_checkpoints if checkpoint else self._lru
+            c.append((model, tokens))
+
+        def remove(self, model, tokens):
+            try:
+                self._lru.remove((model, tokens))
+            except ValueError:
+                self._lru_checkpoints.remove((model, tokens))
+
+        def pop(self):
+            if len(self._lru) >= len(self._lru_checkpoints):
+                return self._lru.popleft()
+            else:
+                return self._lru_checkpoints.popleft()
 
     @dataclass
     class SearchResult:
@@ -185,10 +209,19 @@ class LRUPromptCache:
         longer: List[int]
         common_prefix: int
 
-    def __init__(self, max_size: int = 10):
+    def __init__(self, max_size: int = 10, max_bytes: int = 1 << 63):
         self.max_size = max_size
+        self.max_bytes = max_bytes
         self._cache = {}
-        self._lru = deque()
+        self._lru = self.CacheOrder()
+        self._n_bytes = 0
+
+    def __len__(self):
+        return len(self._lru)
+
+    @property
+    def nbytes(self):
+        return self._n_bytes
 
     def _search(self, model, tokens):
         """Search the cache for a prompt cache. Return exact or close match."""
@@ -217,7 +250,7 @@ class LRUPromptCache:
         # Check for caches that are longer
         longer = None
         common_prefix = index
-        if index > 0 and last_cache_index <= 0:
+        if index > 0:
             best = None
             stack = [(current, [])]
             while stack:
@@ -241,6 +274,8 @@ class LRUPromptCache:
         path = [self._cache[model]]
         for tok in tokens:
             path.append(path[-1][tok])
+        cache_bytes = path[-1]["cache"].nbytes
+        self._n_bytes -= cache_bytes
         del path[-1]["cache"]
         for i in reversed(range(len(tokens))):
             d_prev, d, t = path[i], path[i + 1], tokens[i]
@@ -248,63 +283,81 @@ class LRUPromptCache:
                 break
             del d_prev[t]
 
-    def _extract(self, model, tokens):
-        cache_entry = self._get(model, tokens)
-        if cache_entry.count == 1:
-            self._delete(model, tokens)
-            self._lru.remove((model, tokens))
-            return cache_entry
-
-        cache_entry.count -= 1
-        return self.CacheEntry(
-            copy.deepcopy(cache_entry.prompt_cache),
-            1,
-        )
-
     def fetch_nearest_cache(self, model, tokens):
         result = self._search(model, tokens)
         if result.exact is not None:
-            cache_entry = self._extract(result.model, result.exact)
-            return cache_entry.prompt_cache, []
+            cache_entry = self._get(result.model, result.exact)
+            return copy.deepcopy(cache_entry.prompt_cache), []
 
-        if result.shorter is not None:
-            cache_entry = self._extract(result.model, result.shorter)
-            prefix_len = len(result.shorter)
-            return cache_entry.prompt_cache, tokens[prefix_len:]
-
-        if result.longer is not None:
+        short_length = len(result.shorter) if result.shorter is not None else 0
+        if result.longer is not None and result.common_prefix > short_length:
             cache_entry = self._get(result.model, result.longer)
             if can_trim_prompt_cache(cache_entry.prompt_cache):
-                cache_entry = self.CacheEntry(
-                    copy.deepcopy(cache_entry.prompt_cache),
-                    1,
-                )
+                cache = copy.deepcopy(cache_entry.prompt_cache)
                 prefix = min(len(tokens) - 1, result.common_prefix)
                 num_to_trim = len(result.longer) - prefix
-                trim_prompt_cache(cache_entry.prompt_cache, num_to_trim)
-                return cache_entry.prompt_cache, tokens[prefix:]
+                trim_prompt_cache(cache, num_to_trim)
+                return cache, tokens[prefix:]
+
+        if short_length > 0:
+            cache_entry = self._get(result.model, result.shorter)
+            return copy.deepcopy(cache_entry.prompt_cache), tokens[short_length:]
 
         return None, tokens
 
-    def insert_cache(self, model, tokens, prompt_cache):
+    def insert_cache(self, model, tokens, prompt_cache, checkpoint: bool = False):
+        is_trimmable = can_trim_prompt_cache(prompt_cache)
+
         if model not in self._cache:
             self._cache[model] = {}
         current = self._cache[model]
-        for tok in tokens:
+        for i, tok in enumerate(tokens):
             if tok not in current:
                 current[tok] = {}
+            if is_trimmable and "cache" in current:
+                self._n_bytes -= current["cache"].nbytes
+                del current["cache"]
+                self._lru.remove(model, tokens[:i])
             current = current[tok]
 
         if "cache" in current:
-            current["cache"].count += 1
-            self._lru.remove((model, tokens))
+            self._lru.remove(model, tokens)
         else:
-            current["cache"] = self.CacheEntry(prompt_cache, 1)
+            cache_bytes = sum(c.nbytes for c in prompt_cache)
+            current["cache"] = self.CacheEntry(prompt_cache, cache_bytes)
+            self._n_bytes += cache_bytes
 
-        self._lru.append((model, tokens))
+        self._lru.push(model, tokens, checkpoint=checkpoint)
         if len(self._lru) > self.max_size:
-            model, tokens = self._lru.popleft()
+            model, tokens = self._lru.pop()
             self._delete(model, tokens)
+        while self._n_bytes > self.max_bytes and len(self._lru) > 1:
+            model, tokens = self._lru.pop()
+            self._delete(model, tokens)
+
+    def trim_to(
+        self, *, n_sequences: Optional[int] = None, n_bytes: Optional[int] = None
+    ):
+        n_sequences = max(0, n_sequences) if n_sequences is not None else 1 << 63
+        n_bytes = max(0, n_bytes) if n_bytes is not None else 1 << 63
+
+        while len(self._lru) > n_sequences:
+            model, tokens = self._lru.pop()
+            self._delete(model, tokens)
+        while self._n_bytes > n_bytes:
+            model, tokens = self._lru.pop()
+            self._delete(model, tokens)
+
+    def log_cache_stats(self):
+        ncaches, nbytes = len(self), self.nbytes
+        ntok = (
+            len(self._lru._lru_checkpoints[-1][1])
+            if len(self._lru._lru_checkpoints) > 0
+            else 0
+        )
+        logging.info(
+            f"KV Caches: {ncaches} seq, {nbytes / 1e9:.2f} GB, latest user cache {ntok} tokens"
+        )
 
 
 @dataclass
@@ -329,6 +382,10 @@ class LogitsProcessorArguments:
     logit_bias: Optional[Dict[int, float]]
     repetition_penalty: float
     repetition_context_size: int
+    presence_penalty: float
+    presence_context_size: int
+    frequency_penalty: float
+    frequency_context_size: int
 
 
 @dataclass
@@ -371,6 +428,7 @@ class GenerationContext:
     eos_token_ids: set
     stop_token_sequences: List[List[int]]
     prompt: List[int]
+    prompt_cache_count: int = -1
 
     _should_stop: bool = False
 
@@ -556,6 +614,10 @@ def _make_logits_processors(args):
         args.logits.logit_bias,
         args.logits.repetition_penalty,
         args.logits.repetition_context_size,
+        args.logits.presence_penalty,
+        args.logits.presence_context_size,
+        args.logits.frequency_penalty,
+        args.logits.frequency_context_size,
     )
 
 
@@ -672,6 +734,24 @@ class ResponseGenerator:
         else:
             return tokenizer.encode(request.prompt)
 
+    def _compute_prompt_checkpoint(self, tokenizer, request, prompt):
+        if request.request_type != "chat":
+            return False, -1
+        if request.messages[-1]["role"] != "user":
+            return False, -1
+
+        # Save the KV cache at the end of the prompt just before
+        # the think start token which will likely be removed in the
+        # next turn.
+        prompt_checkpoint = -1
+        if tokenizer.has_thinking:
+            for i in range(1, min(11, len(prompt)) - 1, 1):
+                if prompt[-i] == tokenizer.think_start_id:
+                    prompt_checkpoint = -i - 1
+                    break
+
+        return True, prompt_checkpoint
+
     def _is_batchable(self, args):
         if not self.model_provider.is_batchable:
             return False
@@ -702,6 +782,18 @@ class ResponseGenerator:
                 if uid in batch_results:
                     batch_results[uid]["rqueue"].put((min(processed, total), total))
 
+        def checkpoint_callback(prompts):
+            for uid, prompt_end, cache in prompts:
+                rs = batch_results[uid]
+                if not rs["checkpoint"]:
+                    continue
+                self.prompt_cache.insert_cache(
+                    current_model_key,
+                    rs["cache_key"][:-prompt_end],
+                    list(cache),
+                    checkpoint=True,
+                )
+
         if self._is_distributed:
             seed = mx.distributed.all_sum(mx.random.state[0]).view(mx.uint64).item()
             mx.random.seed(seed)
@@ -750,25 +842,40 @@ class ResponseGenerator:
                     )
                     rqueue.put(ctx)
 
+                    self.prompt_cache.log_cache_stats()
                     cache, rest = self.prompt_cache.fetch_nearest_cache(
                         current_model_key, prompt
                     )
+                    ctx.prompt_cache_count = len(prompt) - len(rest)
                     if cache is None:
                         cache = make_prompt_cache(self.model_provider.model)
 
+                    do_checkpoint, checkpoint_position = (
+                        self._compute_prompt_checkpoint(tokenizer, request, prompt)
+                    )
+
                     (uid,) = batch_generator.insert(
                         [rest],
                         args.max_tokens,
                         caches=[cache],
                         samplers=[_make_sampler(args, tokenizer)],
                         logits_processors=[_make_logits_processors(args)],
+                        prompt_checkpoints=[checkpoint_position],
                     )
                     batch_results[uid] = {
                         "ctx": ctx,
                         "cache_key": prompt[:],
                         "rqueue": rqueue,
                         "detokenizer": tokenizer.detokenizer,
+                        "checkpoint": do_checkpoint,
                     }
+                    # just making sure we don't leave a reference around
+                    del cache
+
+                    if self.model_provider.cli_args.prompt_cache_bytes is not None:
+                        total = self.model_provider.cli_args.prompt_cache_bytes
+                        active = batch_generator.prompt_cache_nbytes
+                        self.prompt_cache.trim_to(n_bytes=total - active)
                     continue
 
                 # No batch generator. Load the model and if it's not
@@ -796,7 +903,9 @@ class ResponseGenerator:
                         stop_tokens=tokenizer.eos_token_ids,
                         completion_batch_size=self.cli_args.decode_concurrency,
                         prefill_batch_size=self.cli_args.prompt_concurrency,
+                        prefill_step_size=self.cli_args.prefill_step_size,
                         prompt_progress_callback=progress_callback,
+                        prompt_checkpoint_callback=checkpoint_callback,
                     )
                     unprocessed_requests.append((rqueue, request, args))
                     continue
@@ -914,9 +1023,11 @@ class ResponseGenerator:
             logits_processors = _make_logits_processors(args)
 
             # Load the KV cache
+            self.prompt_cache.log_cache_stats()
             cache, rest = self.prompt_cache.fetch_nearest_cache(
                 self.model_provider.model_key, prompt
             )
+            ctx.prompt_cache_count = len(prompt) - len(rest)
             cache_key = prompt[:]
             if cache is None:
                 cache = make_prompt_cache(self.model_provider.model)
@@ -935,6 +1046,7 @@ class ResponseGenerator:
                 draft_model=draft_model,
                 num_draft_tokens=args.num_draft_tokens,
                 prompt_progress_callback=progress,
+                prefill_step_size=self.cli_args.prefill_step_size,
             ):
                 rqueue.put(
                     Response(
@@ -1014,7 +1126,13 @@ class APIHandler(BaseHTTPRequestHandler):
         super().__init__(*args, **kwargs)
 
     def _set_cors_headers(self):
-        self.send_header("Access-Control-Allow-Origin", "*")
+        allowed_origins = self.response_generator.cli_args.allowed_origins
+        origin = self.headers.get("Origin")
+        if "*" in allowed_origins:
+            self.send_header("Access-Control-Allow-Origin", "*")
+        elif origin in allowed_origins:
+            self.send_header("Access-Control-Allow-Origin", origin)
+            self.send_header("Vary", "Origin")
         self.send_header("Access-Control-Allow-Methods", "*")
         self.send_header("Access-Control-Allow-Headers", "*")
 
@@ -1050,7 +1168,23 @@ class APIHandler(BaseHTTPRequestHandler):
             return
 
         # Fetch and parse request body
-        content_length = int(self.headers["Content-Length"])
+        content_length = self.headers.get("Content-Length")
+        if content_length is None:
+            self._set_completion_headers(411)
+            self.end_headers()
+            self.wfile.write(
+                json.dumps({"error": "Content-Length header is required"}).encode()
+            )
+            return
+        try:
+            content_length = int(content_length)
+        except ValueError:
+            self._set_completion_headers(400)
+            self.end_headers()
+            self.wfile.write(
+                json.dumps({"error": "Invalid Content-Length header"}).encode()
+            )
+            return
         raw_body = self.rfile.read(content_length)
         try:
             self.body = json.loads(raw_body.decode())
@@ -1091,6 +1225,10 @@ class APIHandler(BaseHTTPRequestHandler):
         self.min_p = self.body.get("min_p", self.response_generator.cli_args.min_p)
         self.repetition_penalty = self.body.get("repetition_penalty", 0.0)
         self.repetition_context_size = self.body.get("repetition_context_size", 20)
+        self.presence_penalty = self.body.get("presence_penalty", 0.0)
+        self.presence_context_size = self.body.get("presence_context_size", 20)
+        self.frequency_penalty = self.body.get("frequency_penalty", 0.0)
+        self.frequency_context_size = self.body.get("frequency_context_size", 20)
         self.xtc_probability = self.body.get("xtc_probability", 0.0)
         self.xtc_threshold = self.body.get("xtc_threshold", 0.0)
         self.logit_bias = self.body.get("logit_bias", None)
@@ -1139,6 +1277,25 @@ class APIHandler(BaseHTTPRequestHandler):
             or self.repetition_penalty < 0
         ):
             raise ValueError("repetition_penalty must be a non-negative float")
+        if (
+            not isinstance(self.repetition_context_size, int)
+            or self.repetition_context_size < 0
+        ):
+            raise ValueError("repetition_context_size must be a non-negative integer")
+        if not isinstance(self.presence_penalty, (float, int)):
+            raise ValueError("Presence penalty must be must be a float")
+        if (
+            not isinstance(self.presence_context_size, int)
+            or self.presence_context_size < 0
+        ):
+            raise ValueError("presence_context_size must be a non-negative integer")
+        if not isinstance(self.frequency_penalty, (float, int)):
+            raise ValueError("Presence penalty must be must be a float")
+        if (
+            not isinstance(self.frequency_context_size, int)
+            or self.frequency_context_size < 0
+        ):
+            raise ValueError("frequency_context_size must be a non-negative integer")
 
         if not isinstance(self.logprobs, bool):
             raise ValueError("logprobs must be a boolean")
@@ -1148,12 +1305,6 @@ class APIHandler(BaseHTTPRequestHandler):
                 f"top_logprobs must be between 1 and 10 but got {self.top_logprobs:,}"
             )
 
-        if (
-            not isinstance(self.repetition_context_size, int)
-            or self.repetition_context_size < 0
-        ):
-            raise ValueError("repetition_context_size must be a non-negative integer")
-
         if self.logit_bias is not None:
             if not isinstance(self.logit_bias, dict):
                 raise ValueError("logit_bias must be a dict of int to float")
@@ -1184,6 +1335,7 @@ class APIHandler(BaseHTTPRequestHandler):
         finish_reason: Union[Literal["length", "stop"], None],
         prompt_token_count: Optional[int] = None,
         completion_token_count: Optional[int] = None,
+        prompt_cache_count: Optional[int] = None,
         token_logprobs: Optional[List[float]] = None,
         top_tokens: Optional[List[Tuple[Dict[str, Any]]]] = None,
         tokens: Optional[List[int]] = None,
@@ -1202,6 +1354,8 @@ class APIHandler(BaseHTTPRequestHandler):
               used to populate the "usage" field (not used when stream).
             completion_token_count (Optional[int]): The number of tokens in the
               response, used to populate the "usage" field (not used when stream).
+            prompt_cache_count (Optional[int]): The portion of prompt_token_count
+              that was found in the cache when servicing the request.
             token_logprobs (Optional[List[float]]): The log probabilities per token,
               in token order.
             top_tokens (Optional[List[Tuple[Dict[str, Any]]]]): List of outputs from
@@ -1260,6 +1414,10 @@ class APIHandler(BaseHTTPRequestHandler):
                 "completion_tokens": completion_token_count,
                 "total_tokens": prompt_token_count + completion_token_count,
             }
+            if prompt_cache_count is not None and prompt_cache_count >= 0:
+                response["usage"]["prompt_tokens_details"] = {
+                    "cached_tokens": prompt_cache_count,
+                }
 
         choice = response["choices"][0]
 
@@ -1306,6 +1464,10 @@ class APIHandler(BaseHTTPRequestHandler):
                 logit_bias=self.logit_bias,
                 repetition_penalty=self.repetition_penalty,
                 repetition_context_size=self.repetition_context_size,
+                presence_penalty=self.presence_penalty,
+                presence_context_size=self.presence_context_size,
+                frequency_penalty=self.frequency_penalty,
+                frequency_context_size=self.frequency_context_size,
             ),
             stop_words=stop_words,
             max_tokens=self.max_tokens,
@@ -1501,7 +1663,11 @@ class APIHandler(BaseHTTPRequestHandler):
             self.wfile.write(f"data: {json.dumps(response)}\n\n".encode())
             self.wfile.flush()
             if self.stream_options is not None and self.stream_options["include_usage"]:
-                response = self.completion_usage_response(len(ctx.prompt), len(tokens))
+                response = self.completion_usage_response(
+                    len(ctx.prompt),
+                    len(tokens),
+                    ctx.prompt_cache_count,
+                )
                 self.wfile.write(f"data: {json.dumps(response)}\n\n".encode())
                 self.wfile.flush()
             self.wfile.write("data: [DONE]\n\n".encode())
@@ -1512,6 +1678,7 @@ class APIHandler(BaseHTTPRequestHandler):
                 finish_reason,
                 len(ctx.prompt),
                 len(tokens),
+                ctx.prompt_cache_count,
                 token_logprobs=token_logprobs,
                 top_tokens=top_tokens,
                 tokens=tokens,
@@ -1532,6 +1699,7 @@ class APIHandler(BaseHTTPRequestHandler):
         self,
         prompt_token_count: Optional[int] = None,
         completion_token_count: Optional[int] = None,
+        prompt_cache_count: Optional[int] = None,
     ):
         response = {
             "id": self.request_id,
@@ -1546,6 +1714,10 @@ class APIHandler(BaseHTTPRequestHandler):
                 "total_tokens": prompt_token_count + completion_token_count,
             },
         }
+        if prompt_cache_count is not None and prompt_cache_count >= 0:
+            response["usage"]["prompt_tokens_details"] = {
+                "cached_tokens": prompt_cache_count,
+            }
         return response
 
     def handle_chat_completions(self) -> CompletionRequest:
@@ -1712,7 +1884,8 @@ def run(
     handler_class=APIHandler,
 ):
     group = mx.distributed.init()
-    response_generator = ResponseGenerator(model_provider, LRUPromptCache())
+    prompt_cache = LRUPromptCache(model_provider.cli_args.prompt_cache_size)
+    response_generator = ResponseGenerator(model_provider, prompt_cache)
     if group.rank() == 0:
         _run_http_server(host, port, response_generator)
     else:
@@ -1743,6 +1916,12 @@ def main():
         default=8080,
         help="Port for the HTTP server (default: 8080)",
     )
+    parser.add_argument(
+        "--allowed-origins",
+        type=lambda x: x.split(","),
+        default="*",
+        help="Allowed origins (default: *)",
+    )
     parser.add_argument(
         "--draft-model",
         type=str,
@@ -1827,6 +2006,23 @@ def main():
         default=8,
         help="When a request is batchable then process that many prompts in parallel",
     )
+    parser.add_argument(
+        "--prefill-step-size",
+        type=int,
+        default=2048,
+        help="Step size for prefill processing (default: 2048)",
+    )
+    parser.add_argument(
+        "--prompt-cache-size",
+        type=int,
+        default=10,
+        help="Maximum number of distinct KV caches to hold in the prompt cache",
+    )
+    parser.add_argument(
+        "--prompt-cache-bytes",
+        type=_parse_size,
+        help="Maximum size in bytes of the KV caches",
+    )
     parser.add_argument(
         "--pipeline",
         action="store_true",

mlx_lm/share.py

@@ -0,0 +1,290 @@
+# Copyright © 2026 Apple Inc.
+
+import argparse
+import os
+import pickle
+import sys
+import time
+from dataclasses import dataclass
+from functools import partial, total_ordering
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Literal, Optional
+
+import mlx.core as mx
+from huggingface_hub.errors import LocalEntryNotFoundError
+from mlx._distributed_utils.common import Hostfile
+from mlx._distributed_utils.launch import launch_jaccl, launch_ring
+from tqdm import tqdm
+
+from .utils import hf_repo_to_path
+
+CHUNK_SIZE = 100 * 1024 * 1024
+
+
+@total_ordering
+@dataclass
+class DirectoryEntry:
+    entry_type: Literal["directory", "symlink", "file"]
+    path: str
+    dst: Optional[str]
+
+    def __lt__(self, other):
+        order_type = dict(directory=0, symlink=1, file=2)
+        o1 = order_type[self.entry_type]
+        o2 = order_type[other.entry_type]
+        return o1 < o2 or (o1 == o2 and self.path < other.path)
+
+    def __eq__(self, other):
+        return (
+            self.entry_type == other.entry_type
+            and self.path == other.path
+            and self.dst == other.dst
+        )
+
+    @classmethod
+    def from_path(cls, root, path):
+        entry_type = {
+            (True, False): "directory",
+            (False, True): "symlink",
+            (False, False): "file",
+        }[path.is_dir(), path.is_symlink()]
+        dst = path.readlink() if path.is_symlink() else None
+
+        return cls(entry_type, str(path.relative_to(root)), str(dst))
+
+
+def error(*args, **kwargs):
+    kwargs["file"] = sys.stderr
+    print("\033[31m[ERROR]", *args, "\033[0m", **kwargs)
+
+
+def launch(args):
+    if args.hostfile is None:
+        raise ValueError("No hostfile provided")
+
+    hostfile = Hostfile.from_file(args.hostfile)
+    if hostfile.backend == "":
+        raise ValueError("Backend needs to be defined in the hostfile.")
+    if len(hostfile.hosts) == 1:
+        raise ValueError("More than one node needs to be in the hostfile")
+
+    launch_args = argparse.Namespace(
+        backend=hostfile.backend,
+        cwd=str(Path.cwd()),
+        env=hostfile.envs,
+        verbose=False,
+        python=None,
+        starting_port=32323,
+        connections_per_ip=1,
+    )
+    cmd = [
+        sys.executable,
+        "-m",
+        "mlx_lm",
+        "share",
+    ]
+    if args.path is not None:
+        cmd += ["--path", args.path]
+    if args.model is not None:
+        cmd += ["--model", args.model]
+    if args.tmpdir is not None:
+        cmd += ["--tmpdir", args.tmpdir]
+    if args.dst is not None:
+        cmd += ["--dst", args.dst]
+
+    if hostfile.backend == "ring":
+        launch_ring(None, hostfile.hosts, launch_args, cmd)
+    elif hostfile.backend == "jaccl" or hostfile.backend == "jaccl-ring":
+        launch_jaccl(None, hostfile.hosts, launch_args, cmd)
+    else:
+        raise ValueError("Only ring, jaccl and jaccl-ring backends are supported.")
+
+
+def get_files(path):
+    if not path.is_dir():
+        return path.parent, [DirectoryEntry.from_path(path.parent, path)]
+
+    files = [DirectoryEntry.from_path(path, f) for f in path.rglob("*")]
+    return path, sorted(files)
+
+
+def format_bw(x):
+    if x >= 1e9:
+        return f"{x / 1e9:.2} GB/s"
+    if x >= 1e6:
+        return f"{x / 1e6:.2} MB/s"
+    if x >= 1e3:
+        return f"{x / 1e3:.2} KB/s"
+    return f"{x:.2} B/s"
+
+
+def share_file(path, file, src, group=None):
+    group = group or mx.distributed.init()
+    all_sum = partial(mx.distributed.all_sum, group=group)
+    total_size = 0
+    start_time = time.time()
+
+    if group.rank() == src:
+        with open(path / file, "rb") as f:
+            f.seek(0, 2)
+            total_size = f.tell()
+            f.seek(0)
+
+            pbar = tqdm(
+                total=total_size,
+                unit="B",
+                unit_scale=True,
+                desc=file,
+                position=1,
+                leave=False,
+            )
+            while True:
+                data = f.read(CHUNK_SIZE)
+                if not data:
+                    mx.eval(all_sum(0))
+                    break
+
+                mx.eval(all_sum(len(data)))
+                mx.async_eval(all_sum(data))
+                pbar.update(len(data))
+            pbar.close()
+
+    else:
+        with open(path / file, "wb") as f:
+            data = None
+            chunk_size = all_sum(0).item()
+            if chunk_size > 0:
+                data = all_sum(mx.zeros(chunk_size, dtype=mx.uint8))
+                mx.eval(data)
+
+            while chunk_size > 0:
+                next_data = None
+                chunk_size = all_sum(0).item()
+                if chunk_size > 0:
+                    next_data = all_sum(mx.zeros(chunk_size, dtype=mx.uint8))
+                    mx.async_eval(next_data)
+
+                f.write(bytes(data))
+                data = next_data
+
+    return total_size, time.time() - start_time
+
+
+def share_files(path, files, src, group=None):
+    group = group or mx.distributed.init()
+    all_sum = partial(mx.distributed.all_sum, group=group)
+
+    if group.rank() == src:
+        # Share the list first
+        file_list = pickle.dumps(files)
+        mx.eval(all_sum(len(file_list)))
+        mx.eval(all_sum(file_list))
+
+    else:
+        # Get the list first
+        file_list_size = all_sum(0).item()
+        data = all_sum(mx.zeros(file_list_size, dtype=mx.uint8))
+        files = pickle.loads(bytes(data))
+
+        # Make the directories and symlinks
+        for file in files:
+            if file.entry_type == "directory":
+                (path / file.path).mkdir()
+            elif file.entry_type == "symlink":
+                (path / file.path).symlink_to(file.dst)
+
+    # Everybody shares the files
+    total_size = 0
+    total_time = 1e-6
+    pbar = tqdm(total=len(files), desc="Files", position=0, disable=group.rank() != src)
+    for file in files:
+        if file.entry_type == "file":
+            s, t = share_file(path, file.path, src, group)
+            total_size += s
+            total_time += t
+        pbar.update(1)
+        pbar.set_postfix(speed=format_bw(total_size / total_time))
+    pbar.close()
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Distribute a model to other nodes using MLX distributed."
+    )
+    parser.add_argument("--path", type=str, help="Path to a file or folder to share.")
+    parser.add_argument(
+        "--model", type=str, help="The path to a local model or Hugging Face repo"
+    )
+    parser.add_argument(
+        "--hostfile",
+        type=str,
+        help="The file containing the hosts and connection information",
+    )
+    parser.add_argument(
+        "--dst",
+        type=str,
+        help="The destination path in other nodes (defaults to --path or --model)",
+    )
+    parser.add_argument(
+        "--tmpdir",
+        type=str,
+        help="Intermediate temporary directory to ensure successfull transfer",
+    )
+
+    args = parser.parse_args()
+
+    if args.path is args.model is None:
+        parser.error("One of --path or --model must be provided")
+
+    mx.set_default_device(mx.cpu)
+    world = mx.distributed.init()
+
+    if world.size() == 1:
+        launch(args)
+        return
+
+    # Check if any node has the data
+    path = None
+    files = []
+    if args.path is not None and (path := Path(args.path)).exists():
+        path, files = get_files(path)
+    elif args.model is not None:
+        try:
+            path = hf_repo_to_path(args.model)
+            if path.parent.name != "snapshots":
+                raise ValueError(
+                    f"The model repository appears to be corrupted, it resolved to {str(path)}"
+                )
+            path, files = get_files(path.parent.parent)
+        except Exception as e:
+            pass
+    has_file = mx.distributed.all_gather(len(files) > 0)
+    src = has_file.argmax().item()
+    has_file = has_file.any().item()
+
+    if not has_file:
+        error("The --path needs to exist in at least one node.")
+        error("If it is a remote repository download it first with `hf download`")
+        sys.exit(1)
+
+    # Share the path that is resolved
+    if args.dst is None:
+        if world.rank() == src:
+            data = str(path).encode("utf-8")
+            mx.eval(mx.distributed.all_sum(len(data)))
+            mx.eval(mx.distributed.all_sum(data))
+        else:
+            data_size = mx.distributed.all_sum(0).item()
+            data = mx.distributed.all_sum(mx.zeros(data_size, dtype=mx.uint8))
+            path = Path(bytes(data).decode("utf-8"))
+    elif world.rank() != src:
+        path = Path(args.dst)
+
+    with TemporaryDirectory(dir=args.tmpdir) as tmp:
+        if world.rank() == src:
+            share_files(path, files, src, world)
+        else:
+            share_files(Path(tmp), files, src, world)
+            path.mkdir(parents=True, exist_ok=True)
+            os.rename(tmp, path)

mlx_lm/tool_parsers/qwen3_coder.py

@@ -4,6 +4,7 @@
 Modified from:
 https://huggingface.co/Qwen/Qwen3-Coder-30B-A3B-Instruct/blob/main/qwen3coder_tool_parser.py
 """
+
 import ast
 import json
 from typing import Any, Optional
@@ -70,7 +71,10 @@ def _convert_param_value(param_value: str, param_name: str, param_config: dict)
             or param_type.startswith("dict")
             or param_type.startswith("list")
         ):
-            return json.loads(param_value)
+            try:
+                return json.loads(param_value)
+            except json.JSONDecodeError:
+                return ast.literal_eval(param_value)
 
         return ast.literal_eval(param_value)
 

mlx_lm/tuner/datasets.py

@@ -116,7 +116,7 @@ class CompletionsDataset:
         if self.mask_prompt:
             offset = len(
                 self.tokenizer.apply_chat_template(
-                    messages[0],
+                    messages[:-1],
                     tools=tools,
                     add_generation_prompt=True,
                     return_dict=False,

mlx_lm/tuner/trainer.py

@@ -17,6 +17,11 @@ from .callbacks import TrainingCallback
 from .datasets import CacheDataset
 
 
+def _clear_cache(threshold: int):
+    if mx.get_cache_memory() > threshold:
+        mx.clear_cache()
+
+
 def grad_checkpoint(layer):
     """
     Update all instances of type(layer) to use gradient checkpointing.
@@ -70,6 +75,12 @@ class TrainingArgs:
             "help": "Number of steps to accumulate gradients before applying an optimizer update."
         },
     )
+    clear_cache_threshold: int = field(
+        default=0,
+        metadata={
+            "help": "Clear the allocator cache between steps if it grows too large."
+        },
+    )
 
 
 def default_loss(model, batch, lengths):
@@ -170,6 +181,7 @@ def evaluate(
     max_seq_length=2048,
     loss: callable = default_loss,
     iterate_batches: callable = iterate_batches,
+    clear_cache_threshold: int = 0,
 ):
     model.eval()
     all_losses = mx.array(0.0)
@@ -194,11 +206,13 @@ def evaluate(
         all_losses += losses * toks
         ntokens += toks
         mx.eval(all_losses, ntokens)
+        _clear_cache(clear_cache_threshold)
 
     all_losses = mx.distributed.all_sum(all_losses, stream=mx.cpu)
     ntokens = mx.distributed.all_sum(ntokens, stream=mx.cpu)
+    avg_loss = (all_losses / ntokens).item()
 
-    return (all_losses / ntokens).item()
+    return avg_loss
 
 
 def train(
@@ -212,7 +226,7 @@ def train(
     training_callback: TrainingCallback = None,
 ):
     if mx.metal.is_available():
-        mx.set_wired_limit(mx.metal.device_info()["max_recommended_working_set_size"])
+        mx.set_wired_limit(mx.device_info()["max_recommended_working_set_size"])
     print(f"Starting training..., iters: {args.iters}")
     world = mx.distributed.init()
     world_size = world.size()
@@ -312,6 +326,7 @@ def train(
         n_tokens += toks
         steps += 1
         mx.eval(state, losses, n_tokens, grad_accum)
+        _clear_cache(args.clear_cache_threshold)
         train_time += time.perf_counter() - tic
 
         # Report training loss if needed

mlx_lm/utils.py

@@ -47,6 +47,7 @@ MODEL_REMAPPING = {
     "llava": "mistral3",
     "phi-msft": "phixtral",
     "falcon_mamba": "mamba",
+    "joyai_llm_flash": "deepseek_v3",
     "kimi_k2": "deepseek_v3",
     "qwen2_5_vl": "qwen2_vl",
     "minimax_m2": "minimax",
@@ -56,6 +57,18 @@ MODEL_REMAPPING = {
 MAX_FILE_SIZE_GB = 5
 
 
+def _parse_size(x):
+    sizes = {"M": 1e6, "G": 1e9, "MB": 1e6, "GB": 1e9, "": 1}
+    split = 0
+    for xi in x:
+        if not (xi.isdigit() or xi == "."):
+            break
+        split += 1
+    digits = float(x[:split])
+    size = (x[split:]).strip().upper()
+    return int(digits * sizes[size])
+
+
 def _unpack_awq_weights(qweight: mx.array) -> mx.array:
     bits = 4
     pack_factor = 32 // bits
@@ -514,7 +527,7 @@ def sharded_load(
     # weights we need to download.
     model, config = load_model(model_path, lazy=True, strict=False)
 
-    has_pipelining = hasattr(model.model, "pipeline")
+    has_pipelining = hasattr(model, "model") and hasattr(model.model, "pipeline")
     has_tensor_parallel = hasattr(model, "shard")
 
     if pipeline_group is not None and not has_pipelining:

setup.py

@@ -66,6 +66,7 @@ setup(
             "mlx_lm.lora = mlx_lm.lora:main",
             "mlx_lm.perplexity = mlx_lm.perplexity:main",
             "mlx_lm.server = mlx_lm.server:main",
+            "mlx_lm.share = mlx_lm.share:main",
             "mlx_lm.manage = mlx_lm.manage:main",
             "mlx_lm.upload = mlx_lm.upload:main",
         ]

tests/test_datsets.py

@@ -61,6 +61,37 @@ class TestDatasets(unittest.TestCase):
         self.assertTrue(len(valid[0]) > 0)
         self.assertTrue(isinstance(train, datasets.CompletionsDataset))
 
+    def test_completions_mask_prompt(self):
+        data = {"prompt": "What is the capital of France?", "completion": "Paris."}
+        self.save_data(4 * [data])
+        args = types.SimpleNamespace(
+            train=True, test=False, data=self.test_dir, mask_prompt=True
+        )
+        tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_PATH, local_files_only=True)
+        train, valid, test = datasets.load_dataset(args, tokenizer)
+        self.assertEqual(len(train), 4)
+        self.assertEqual(len(valid), 4)
+        self.assertEqual(len(test), 0)
+        expected_prompt_tokens = tokenizer.apply_chat_template(
+            [{"role": "user", "content": data["prompt"]}],
+            add_generation_prompt=True,
+            return_dict=False,
+        )
+        expected_offset = len(expected_prompt_tokens)
+
+        train_tokens, train_offset = train.process(train[0])
+        valid_tokens, valid_offset = valid.process(valid[0])
+
+        self.assertTrue(len(train_tokens) > 0)
+        self.assertTrue(len(valid_tokens) > 0)
+        self.assertEqual(train_offset, expected_offset)
+        self.assertEqual(valid_offset, expected_offset)
+        self.assertLess(train_offset, len(train_tokens))
+        self.assertLess(valid_offset, len(valid_tokens))
+        self.assertEqual(train_tokens[:train_offset], expected_prompt_tokens)
+        self.assertEqual(valid_tokens[:valid_offset], expected_prompt_tokens)
+        self.assertTrue(isinstance(train, datasets.CompletionsDataset))
+
     def test_chat(self):
         data = {
             "messages": [

tests/test_losses.py

@@ -13,21 +13,21 @@ class TestLosses(unittest.TestCase):
         self.assertTrue(can_run_metal())
         mx.random.seed(0)
 
-        logits_q = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
-        logits_p = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
+        logits_q = mx.random.normal((2, 4, 4000))
+        logits_p = mx.random.normal((2, 4, 4000))
 
         with mx.stream(mx.cpu):
             expected = kl_div_loss(logits_q, logits_p)
         kl = kl_div_loss(logits_q, logits_p)
 
-        self.assertTrue(mx.allclose(kl, expected, rtol=1e-4))
+        self.assertTrue(mx.allclose(kl, expected))
 
     def test_js_div_loss(self):
         self.assertTrue(can_run_metal())
         mx.random.seed(0)
 
-        logits_q = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
-        logits_p = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
+        logits_q = mx.random.normal((2, 4, 4000))
+        logits_p = mx.random.normal((2, 4, 4000))
 
         with mx.stream(mx.cpu):
             expected = js_div_loss(logits_q, logits_p)
@@ -39,9 +39,9 @@ class TestLosses(unittest.TestCase):
         self.assertTrue(can_run_metal())
         mx.random.seed(0)
 
-        logits_q = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
-        logits_p = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
-        cotan = mx.random.uniform(shape=(4, 8), dtype=mx.float32)
+        logits_q = mx.random.normal((2, 4, 4000))
+        logits_p = mx.random.normal((2, 4, 4000))
+        cotan = mx.random.normal((2, 4))
 
         with mx.stream(mx.cpu):
             expected = mx.vjp(kl_div_loss, [logits_q, logits_p], [cotan])[1][0]
@@ -53,9 +53,9 @@ class TestLosses(unittest.TestCase):
         self.assertTrue(can_run_metal())
         mx.random.seed(0)
 
-        logits_q = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
-        logits_p = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
-        cotan = mx.random.uniform(shape=(4, 8), dtype=mx.float32)
+        logits_q = mx.random.normal((2, 4, 4000))
+        logits_p = mx.random.normal((2, 4, 4000))
+        cotan = mx.random.normal((2, 4))
 
         with mx.stream(mx.cpu):
             expected = mx.vjp(js_div_loss, [logits_q, logits_p], [cotan])[1][0]

tests/test_models.py

@@ -10,7 +10,11 @@ from mlx.utils import tree_map
 from mlx_lm.models import rope_utils
 from mlx_lm.models.base import create_causal_mask, scaled_dot_product_attention
 from mlx_lm.models.cache import KVCache, RotatingKVCache, make_prompt_cache
-from mlx_lm.models.gated_delta import gated_delta_kernel, gated_delta_ops
+from mlx_lm.models.gated_delta import (
+    gated_delta_kernel,
+    gated_delta_ops,
+    gated_delta_update,
+)
 from mlx_lm.models.ssm import ssm_attn, ssm_update
 
 
@@ -238,6 +242,30 @@ class TestModels(unittest.TestCase):
         )
         self.assertTrue(isinstance(rope, rope_utils.Llama3RoPE))
 
+    def test_su_scaled_rope_no_mutation(self):
+        rope = rope_utils.SuScaledRoPE(
+            dims=8,
+            max_position_embeddings=131072,
+            original_max_position_embeddings=4096,
+            long_factor=[1.0] * 4,
+        )
+        x = mx.ones((1, 2, 4, 8))
+        rope(x)
+        mx.eval(x)
+        self.assertTrue((x == 1).all())
+
+    def test_yarn_rope_no_mutation(self):
+        rope = rope_utils.YarnRoPE(
+            dims=8,
+            scaling_factor=2.0,
+            mscale=1.0,
+            mscale_all_dim=0,
+        )
+        x = mx.ones((1, 2, 4, 8))
+        rope(x)
+        mx.eval(x)
+        self.assertTrue((x == 1).all())
+
     def test_quantized_sdpa(self):
         cache = KVCache()
 
@@ -762,6 +790,104 @@ class TestModels(unittest.TestCase):
             model, args.model_type, args.vocab_size, args.num_hidden_layers
         )
 
+    def test_step3p5_make_cache_uses_rotating_for_sliding_layers(self):
+        from mlx_lm.models import step3p5
+
+        args = step3p5.ModelArgs(
+            model_type="step3p5",
+            hidden_size=256,
+            num_hidden_layers=4,
+            vocab_size=1024,
+            num_attention_heads=4,
+            num_attention_groups=2,
+            head_dim=64,
+            intermediate_size=512,
+            rms_norm_eps=1e-5,
+            rope_theta=[10000.0, 10000.0, 10000.0, 10000.0],
+            sliding_window=4,
+            layer_types=[
+                "full_attention",
+                "sliding_attention",
+                "sliding_attention",
+                "full_attention",
+            ],
+            partial_rotary_factors=[0.5, 1.0, 1.0, 0.5],
+            attention_other_setting={
+                "num_attention_heads": 8,
+                "num_attention_groups": 2,
+            },
+            use_head_wise_attn_gate=True,
+            moe_num_experts=4,
+            moe_top_k=2,
+            moe_intermediate_size=256,
+            share_expert_dim=256,
+            moe_layers_enum="1,2,3",
+        )
+        model = step3p5.Model(args)
+
+        caches = model.make_cache()
+        self.assertIsInstance(caches[0], KVCache)
+        self.assertIsInstance(caches[1], RotatingKVCache)
+        self.assertIsInstance(caches[2], RotatingKVCache)
+        self.assertIsInstance(caches[3], KVCache)
+
+        tokens = mx.array([[1, 2, 3, 4, 5, 6, 7]], dtype=mx.int32)
+        step = model(tokens[:, :3], cache=caches)
+        mx.eval(step)
+        for i in range(3, 7):
+            step = model(tokens[:, i : i + 1], cache=caches)
+            mx.eval(step)
+
+        self.assertEqual(caches[0].size(), 7)
+        self.assertEqual(caches[1].size(), args.sliding_window)
+        self.assertEqual(caches[2].size(), args.sliding_window)
+        self.assertEqual(caches[3].size(), 7)
+
+    def test_step3p5_make_cache_uses_fallback_sliding_pattern(self):
+        from mlx_lm.models import step3p5
+
+        args = step3p5.ModelArgs(
+            model_type="step3p5",
+            hidden_size=256,
+            num_hidden_layers=5,
+            vocab_size=1024,
+            num_attention_heads=4,
+            num_attention_groups=2,
+            head_dim=64,
+            intermediate_size=512,
+            rms_norm_eps=1e-5,
+            rope_theta=10000.0,
+            sliding_window=4,
+            partial_rotary_factors=[1.0] * 5,
+            use_head_wise_attn_gate=True,
+            moe_num_experts=4,
+            moe_top_k=2,
+            moe_intermediate_size=256,
+            share_expert_dim=256,
+            moe_layers_enum="1,2,3,4",
+        )
+        model = step3p5.Model(args)
+
+        caches = model.make_cache()
+        self.assertIsInstance(caches[0], RotatingKVCache)
+        self.assertIsInstance(caches[1], KVCache)
+        self.assertIsInstance(caches[2], RotatingKVCache)
+        self.assertIsInstance(caches[3], KVCache)
+        self.assertIsInstance(caches[4], RotatingKVCache)
+
+        tokens = mx.array([[1, 2, 3, 4, 5, 6]], dtype=mx.int32)
+        step = model(tokens[:, :2], cache=caches)
+        mx.eval(step)
+        for i in range(2, 6):
+            step = model(tokens[:, i : i + 1], cache=caches)
+            mx.eval(step)
+
+        self.assertEqual(caches[0].size(), args.sliding_window)
+        self.assertEqual(caches[1].size(), 6)
+        self.assertEqual(caches[2].size(), args.sliding_window)
+        self.assertEqual(caches[3].size(), 6)
+        self.assertEqual(caches[4].size(), args.sliding_window)
+
     def test_cohere(self):
         from mlx_lm.models import cohere
 
@@ -1585,7 +1711,7 @@ class TestModels(unittest.TestCase):
                 "rms_norm_eps": 1e-5,
                 "vocab_size": 1000,
                 "num_key_value_heads": 2,
-                "partial_rotary_factor": 0,
+                "partial_rotary_factor": 0.5,
                 "rope_theta": 1000,
             },
             {
@@ -1613,7 +1739,7 @@ class TestModels(unittest.TestCase):
                 "use_qk_norm": True,
                 "tie_word_embeddings": False,
                 "attention_bias": False,
-                "partial_rotary_factor": 0.0,
+                "partial_rotary_factor": 0.5,
             },
             {
                 "model_type": "glm4_moe_lite",
@@ -2532,6 +2658,60 @@ class TestModels(unittest.TestCase):
                 self.assertTrue(mx.allclose(y_op, y_c, rtol=1e-4, atol=1e-4))
                 self.assertTrue(mx.allclose(st_op, st_c, rtol=1e-4, atol=1e-4))
 
+    def test_gated_delta_precision(self):
+        mx.random.seed(42)
+
+        N_STEPS = 512
+        B = 1
+        Hk = 4
+        Hv = 4
+        Dk = 64
+        Dv = 64
+
+        A_log = mx.zeros((Hv,))
+        dt_bias = mx.ones((Hv,))
+
+        all_q = mx.random.normal(shape=(N_STEPS, B, 1, Hk, Dk)) * 0.1
+        all_k = mx.random.normal(shape=(N_STEPS, B, 1, Hk, Dk)) * 0.1
+        all_v = mx.random.normal(shape=(N_STEPS, B, 1, Hv, Dv)) * 0.1
+        all_a = -7.0 + mx.random.normal(shape=(N_STEPS, B, 1, Hv)) * 0.3
+        all_b = mx.random.normal(shape=(N_STEPS, B, 1, Hv))
+        mx.eval(all_q, all_k, all_v, all_a, all_b, A_log, dt_bias)
+
+        state_ref = mx.zeros((B, Hv, Dv, Dk), dtype=mx.float32)
+        for t in range(N_STEPS):
+            y_ref, state_ref = gated_delta_update(
+                all_q[t],
+                all_k[t],
+                all_v[t],
+                all_a[t],
+                all_b[t],
+                A_log,
+                dt_bias,
+                state_ref,
+                use_kernel=False,
+            )
+            mx.eval(y_ref, state_ref)
+
+        for use_kernel in (False, True):
+            state_lo = mx.zeros((B, Hv, Dv, Dk), dtype=mx.bfloat16)
+            for t in range(N_STEPS):
+                y_lo, state_lo = gated_delta_update(
+                    all_q[t].astype(mx.bfloat16),
+                    all_k[t].astype(mx.bfloat16),
+                    all_v[t].astype(mx.bfloat16),
+                    all_a[t].astype(mx.bfloat16),
+                    all_b[t].astype(mx.bfloat16),
+                    A_log,
+                    dt_bias,
+                    state_lo,
+                    use_kernel=use_kernel,
+                )
+                mx.eval(y_lo, state_lo)
+
+            self.assertTrue(mx.allclose(state_lo, state_ref, rtol=0.05, atol=0.01))
+            self.assertTrue(mx.allclose(y_lo, y_ref, rtol=0.05, atol=0.01))
+
     def test_gated_delta_masked(self):
         B = 1
         T = 3

tests/test_prompt_cache.py

@@ -132,6 +132,41 @@ class TestPromptCache(unittest.TestCase):
                 self.assertTrue(mx.array_equal(k, lk))
                 self.assertTrue(mx.array_equal(v, lv))
 
+    def test_save_load_cache_list(self):
+        cache_file = os.path.join(self.test_dir, "prompt_cache.safetensors")
+
+        cache = [
+            ArraysCache(size=2),
+            KVCache(),
+            RotatingKVCache(8),
+            ArraysCache(size=2),
+            ChunkedKVCache(256),
+        ]
+        for c in cache:
+            if isinstance(c, ArraysCache):
+                c[0] = mx.random.uniform(shape=(4, 4, 4))
+                c[1] = mx.random.uniform(shape=(4, 4, 4))
+            else:
+                x = mx.random.uniform(shape=(4, 4, 7, 4))
+                y = mx.random.uniform(shape=(4, 4, 7, 4))
+                c.update_and_fetch(x, y)
+        cache = [CacheList(*cache)]
+
+        save_prompt_cache(cache_file, cache)
+        loaded_cache = load_prompt_cache(cache_file)
+        for c, lc in zip(cache[0].caches, loaded_cache[0].caches):
+            if isinstance(c, ArraysCache):
+                self.assertTrue(mx.array_equal(c[0], lc[0]))
+                self.assertTrue(mx.array_equal(c[1], lc[1]))
+            else:
+                x = mx.random.uniform(shape=(4, 4, 1, 4))
+                y = mx.random.uniform(shape=(4, 4, 1, 4))
+                k, v = c.update_and_fetch(x, y)
+                lk, lv = lc.update_and_fetch(x, y)
+                self.assertEqual(c.offset, lc.offset)
+                self.assertTrue(mx.array_equal(k, lk))
+                self.assertTrue(mx.array_equal(v, lv))
+
     def test_save_load_arrays_cache(self):
         cache_file = os.path.join(self.test_dir, "prompt_cache.safetensors")
 

tests/test_sample_utils.py

@@ -116,6 +116,64 @@ class TestSampleUtils(unittest.TestCase):
         new_probs = mx.softmax(apply_xtc(mx.log(probs), 0, 0.1, [0]), -1)
         self.assertTrue(mx.allclose(new_probs, probs))
 
+    def test_presence_penalty(self):
+        from mlx_lm.sample_utils import make_presence_penalty
+
+        # Token appears multiple times - penalty applied once
+        tokens = mx.array([0, 0, 0, 1, 1])
+        logits = mx.zeros((1, 4))
+        processor = make_presence_penalty(0.5, context_size=5)
+        result = processor(tokens, logits)
+        # Token 0 appears 3 times, token 1 appears 2 times - both penalized once
+        self.assertAlmostEqual(result[0, 0].item(), -0.5)
+        self.assertAlmostEqual(result[0, 1].item(), -0.5)
+        # Tokens not in context not penalized
+        self.assertAlmostEqual(result[0, 2].item(), 0.0)
+        self.assertAlmostEqual(result[0, 3].item(), 0.0)
+
+    def test_frequency_penalty(self):
+        from mlx_lm.sample_utils import make_frequency_penalty
+
+        # Token appears multiple times - penalty applied proportionally
+        tokens = mx.array([0, 0, 0, 1, 1])
+        logits = mx.zeros((1, 4))
+        processor = make_frequency_penalty(0.5, context_size=5)
+        result = processor(tokens, logits)
+        # Token 0 appears 3 times -> 3 * 0.5 = 1.5 penalty
+        self.assertAlmostEqual(result[0, 0].item(), -1.5)
+        # Token 1 appears 2 times -> 2 * 0.5 = 1.0 penalty
+        self.assertAlmostEqual(result[0, 1].item(), -1.0)
+        # Tokens not in context not penalized
+        self.assertAlmostEqual(result[0, 2].item(), 0.0)
+        self.assertAlmostEqual(result[0, 3].item(), 0.0)
+
+    def test_make_logits_processors(self):
+        from mlx_lm.sample_utils import make_logits_processors
+
+        # Create processors with all three penalty types
+        tokens = mx.array([0, 0, 0, 1, 1])
+        # Use non-zero logits so repetition penalty has effect
+        logits = mx.array([[1.0, 0.5, 0.0, -0.5]])
+        processors = make_logits_processors(
+            repetition_penalty=1.5,
+            repetition_context_size=5,
+            presence_penalty=0.5,
+            presence_context_size=5,
+            frequency_penalty=0.25,
+            frequency_context_size=5,
+        )
+        # Apply all processors
+        for processor in processors:
+            logits = processor(tokens, logits)
+        # Token 0 (appears 3x): 1.0/1.5 - 0.5 - 0.75 = -0.5833
+        # Token 1 (appears 2x): 0.5/1.5 - 0.5 - 0.5 = -0.6667
+        # Token 2 (not in context): 0.0 (no penalty)
+        # Token 3 (not in context): -0.5 (no penalty)
+        self.assertAlmostEqual(logits[0, 0].item(), -0.5833, places=4)
+        self.assertAlmostEqual(logits[0, 1].item(), -0.6667, places=4)
+        self.assertAlmostEqual(logits[0, 2].item(), 0.0, places=4)
+        self.assertAlmostEqual(logits[0, 3].item(), -0.5, places=4)
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_server.py

@@ -43,6 +43,11 @@ class DummyModelProvider:
                 "model": None,
                 "decode_concurrency": 32,
                 "prompt_concurrency": 8,
+                "prefill_step_size": 2048,
+                "prompt_cache_size": 10,
+                "prompt_cache_bytes": 1 << 63,
+                "prompt_cache_total_bytes": None,
+                "allowed_origins": ["*"],
             },
         )
 
@@ -57,6 +62,26 @@ class DummyModelProvider:
         return self.model, self.tokenizer
 
 
+class MockCache:
+    def __init__(self, value, is_trimmable: bool = True):
+        self.value = value
+        self._is_trimmable = is_trimmable
+
+    @property
+    def nbytes(self):
+        return len(self.value)
+
+    def __eq__(self, other):
+        return other.value == self.value
+
+    def is_trimmable(self):
+        return self._is_trimmable
+
+    def trim(self, n):
+        assert self._is_trimmable
+        return n
+
+
 class TestServer(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
@@ -354,7 +379,6 @@ class TestServerWithDraftModel(unittest.TestCase):
 
 
 class TestKeepalive(unittest.TestCase):
-
     def test_keepalive_callback(self):
         """Test keepalive callback sends SSE comments and handles errors"""
         from unittest.mock import Mock
@@ -404,7 +428,6 @@ class TestKeepalive(unittest.TestCase):
 
 
 class TestLRUPromptCache(unittest.TestCase):
-
     def test_caching(self):
         cache = LRUPromptCache(max_size=10)
 
@@ -423,18 +446,23 @@ class TestLRUPromptCache(unittest.TestCase):
         c[0].update_and_fetch(*get_kv(24))
         cache.insert_cache(model, t, c)
 
+        # Fetching a cache that is strictly a prefix doesn't remove it from the
+        # lru cache
         tokens = tokens + [20] * 5
         c, t = cache.fetch_nearest_cache(model, tokens)
         k, v = c[0].state
         self.assertTrue((k == v).all().item())
         self.assertTrue((k.flatten() == mx.arange(24)).all().item())
         self.assertEqual(t, [20] * 5)
-        self.assertEqual(len(cache._lru), 0)
+        self.assertEqual(len(cache), 1)
 
+        # Inserting a trimmable cache with shared prefix removes the prefixes
         tokens = tokens + [30] * 3
         c[0].update_and_fetch(*get_kv(8))
         cache.insert_cache(model, tokens, c)
+        self.assertEqual(len(cache), 1)
 
+        # Fetching a cache with a shared prefix doesn't remove it either
         tokens = tokens[:26] + [40] * 8
         c, t = cache.fetch_nearest_cache(model, tokens)
         k, v = c[0].state
@@ -443,38 +471,95 @@ class TestLRUPromptCache(unittest.TestCase):
             (k.flatten() == mx.concatenate([mx.arange(24), mx.arange(2)])).all().item()
         )
         self.assertEqual(t, [40] * 8)
-        self.assertEqual(len(cache._lru), 1)
+        self.assertEqual(len(cache), 1)
+
+        # Inserting a diverged cache actually creates another entry
+        c[0].update_and_fetch(*get_kv(8))
+        cache.insert_cache(model, tokens, c)
+        self.assertEqual(len(cache), 2)
 
     def test_lru(self):
         cache = LRUPromptCache(max_size=2)
         model = ("test", None, None)
-        cache.insert_cache(model, [1, 2], ["test1"])
-        cache.insert_cache(model, [1, 2], ["test1"])
+        cache.insert_cache(model, [1, 2], [MockCache("test1")])
+        cache.insert_cache(model, [2, 3], [MockCache("test2")])
 
         c, t = cache.fetch_nearest_cache(model, [1, 2])
-        self.assertEqual(c, ["test1"])
+        self.assertEqual(c, [MockCache("test1")])
         self.assertEqual(t, [])
-        c, t = cache.fetch_nearest_cache(model, [1, 2])
-        self.assertEqual(c, ["test1"])
-        self.assertEqual(t, [])
-        c, t = cache.fetch_nearest_cache(model, [1, 2])
-        self.assertEqual(c, None)
-        self.assertEqual(t, [1, 2])
+        c, t = cache.fetch_nearest_cache(model, [1])
+        self.assertEqual(c, [MockCache("test1")])
+        self.assertEqual(t, [1])
+        c, t = cache.fetch_nearest_cache(model, [1, 3, 4])
+        self.assertEqual(c, [MockCache("test1")])
+        self.assertEqual(t, [3, 4])
+        c, t = cache.fetch_nearest_cache(model, [2, 3, 4])
+        self.assertEqual(c, [MockCache("test2")])
+        self.assertEqual(t, [4])
+        c, t = cache.fetch_nearest_cache(model, [2, 4, 5])
+        self.assertEqual(c, [MockCache("test2")])
+        self.assertEqual(t, [4, 5])
 
-        cache.insert_cache(model, [1, 2], ["test1"])
-        cache.insert_cache(model, [2, 3], ["test2"])
-        cache.insert_cache(model, [3, 4], ["test3"])
+        cache.insert_cache(model, [1, 2], [MockCache("test1")])
+        cache.insert_cache(model, [2, 3], [MockCache("test2")])
+        cache.insert_cache(model, [3, 4], [MockCache("test3")])
 
         c, t = cache.fetch_nearest_cache(model, [1, 2])
         self.assertEqual(c, None)
         self.assertEqual(t, [1, 2])
         c, t = cache.fetch_nearest_cache(model, [2, 3])
-        self.assertEqual(c, ["test2"])
+        self.assertEqual(c, [MockCache("test2")])
         self.assertEqual(t, [])
         c, t = cache.fetch_nearest_cache(model, [3, 4])
-        self.assertEqual(c, ["test3"])
+        self.assertEqual(c, [MockCache("test3")])
         self.assertEqual(t, [])
 
+        cache.insert_cache(model, [4, 5], [MockCache("test4")], checkpoint=True)
+        c, t = cache.fetch_nearest_cache(model, [2, 3])
+        self.assertEqual(c, None)
+        self.assertEqual(t, [2, 3])
+        c, t = cache.fetch_nearest_cache(model, [3, 4])
+        self.assertEqual(c, [MockCache("test3")])
+        self.assertEqual(t, [])
+        c, t = cache.fetch_nearest_cache(model, [4, 5])
+        self.assertEqual(c, [MockCache("test4")])
+        self.assertEqual(t, [])
+
+        cache.insert_cache(model, [5, 6], [MockCache("test5")])
+        cache.insert_cache(model, [6, 7], [MockCache("test6")])
+        c, t = cache.fetch_nearest_cache(model, [5, 6])
+        self.assertEqual(c, None)
+        self.assertEqual(t, [5, 6])
+        c, t = cache.fetch_nearest_cache(model, [6, 7])
+        self.assertEqual(c, [MockCache("test6")])
+        self.assertEqual(t, [])
+        c, t = cache.fetch_nearest_cache(model, [4, 5])
+        self.assertEqual(c, [MockCache("test4")])
+        self.assertEqual(t, [])
+
+    def test_lru_bytes(self):
+        cache = LRUPromptCache(max_size=100, max_bytes=10)
+        model = ("test", None, None)
+
+        cache.insert_cache(model, [1, 2], [MockCache("aaa")])
+        cache.insert_cache(model, [3, 4], [MockCache("bbb")])
+        cache.insert_cache(model, [4, 5], [MockCache("ccc")])
+        cache.insert_cache(model, [6, 7], [MockCache("ddd")])
+
+        self.assertEqual(len(cache), 3)
+        self.assertEqual(cache.nbytes, 9)
+
+        cache.trim_to(n_bytes=7)
+        self.assertEqual(len(cache), 2)
+        self.assertEqual(cache.nbytes, 6)
+
+        c, t = cache.fetch_nearest_cache(model, [1, 2])
+        self.assertEqual(c, None)
+        self.assertEqual(t, [1, 2])
+        c, t = cache.fetch_nearest_cache(model, [3, 4])
+        self.assertEqual(c, None)
+        self.assertEqual(t, [3, 4])
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_tool_parsing.py

@@ -15,7 +15,6 @@ from mlx_lm.tool_parsers import (
 
 
 class TestToolParsing(unittest.TestCase):
-
     def test_parsers(self):
         test_cases = [
             ("call:multiply{a:12234585,b:48838483920}", function_gemma),
@@ -149,6 +148,49 @@ class TestToolParsing(unittest.TestCase):
                 }
                 self.assertEqual(tool_call, expected)
 
+    def test_qwen3_coder_single_quoted_params(self):
+        tools = [
+            {
+                "type": "function",
+                "function": {
+                    "name": "search",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "filters": {"type": "object"},
+                            "tags": {"type": "array"},
+                        },
+                    },
+                },
+            }
+        ]
+
+        # single-quoted dict (python-style, not valid JSON)
+        test_case = (
+            "<function=search>"
+            "<parameter=filters>{'category': 'books', 'in_stock': True}</parameter>"
+            "<parameter=tags>['fiction', 'new']</parameter>"
+            "</function>"
+        )
+        tool_call = qwen3_coder.parse_tool_call(test_case, tools)
+        self.assertEqual(tool_call["name"], "search")
+        self.assertEqual(
+            tool_call["arguments"]["filters"],
+            {"category": "books", "in_stock": True},
+        )
+        self.assertEqual(tool_call["arguments"]["tags"], ["fiction", "new"])
+
+        # valid JSON (double-quoted) should still work
+        test_case = (
+            "<function=search>"
+            '<parameter=filters>{"category": "books"}</parameter>'
+            '<parameter=tags>["fiction", "new"]</parameter>'
+            "</function>"
+        )
+        tool_call = qwen3_coder.parse_tool_call(test_case, tools)
+        self.assertEqual(tool_call["arguments"]["filters"], {"category": "books"})
+        self.assertEqual(tool_call["arguments"]["tags"], ["fiction", "new"])
+
     def test_kimi_k2(self):
         # Single tool call
         test_case = (