Fix the rope mutation in a more natural way

2026-03-18 13:34:13 -07:00
37 changed files with 1440 additions and 4570 deletions
@@ -1,3 +1,3 @@
 # Copyright © 2023-2025 Apple Inc.

-__version__ = "0.31.3"
+__version__ = "0.31.2"
@@ -148,13 +148,10 @@ def main():
    for i in range(args.num_trials):
        if args.delay > 0:
            time.sleep(args.delay)
-        tic = time.perf_counter()
        response = _bench()
-        toc = time.perf_counter()
        responses.append(response)
        results = [(k, getattr(response, k)) for k in report_keys]
        results = [f"{k}={v:.3f}" for k, v in results]
-        results.append(f"total_time={toc - tic:.3f}")
        rprint(f"Trial {i+1}:  " + ", ".join(results))

    def avg(k):
@@ -1,16 +1,9 @@
 # Copyright © 2023-2024 Apple Inc.

 import argparse
-import sys

 import mlx.core as mx

-from .cli_ui import (
-    make_console,
-    make_corridor_prompt,
-    print_chat_help,
-    print_header_panel,
-)
 from .generate import stream_generate
 from .models.cache import make_prompt_cache
 from .sample_utils import make_sampler
@@ -25,19 +18,6 @@ DEFAULT_MAX_TOKENS = 256
 DEFAULT_MODEL = "mlx-community/Llama-3.2-3B-Instruct-4bit"


-def _print_chat_header(args, console):
-    rows = [("model", str(args.model))]
-    if args.adapter_path:
-        rows.append(("adapter", str(args.adapter_path)))
-    rows.append(("max tokens", f"{args.max_tokens:,}"))
-    if args.system_prompt:
-        sp = args.system_prompt
-        if len(sp) > 60:
-            sp = sp[:57] + "..."
-        rows.append(("system", sp))
-    print_header_panel(console, "mlx_lm.chat", rows)
-
-
 def setup_arg_parser():
    """Set up and return the argument parser."""
    parser = argparse.ArgumentParser(description="Chat with an LLM")
@@ -116,8 +96,6 @@ def main():
    pipeline_group = group if args.pipeline else None
    tensor_group = group if not args.pipeline else None

-    console = make_console()
-
    def rprint(*args, **kwargs):
        if rank == 0:
            print(*args, **kwargs)
@@ -137,38 +115,24 @@ def main():
            },
        )

-    if rank == 0:
-        _print_chat_header(args, console)
-        print_chat_help(console)
-
-    if rank == 0:
-        prompt_console = make_console(force_terminal=True, color_system="truecolor")
-        _draw_corridor_prompt = make_corridor_prompt(prompt_console)
-    else:
-        _draw_corridor_prompt = lambda: ""
+    def print_help():
+        rprint("The command list:")
+        rprint("- 'q' to exit")
+        rprint("- 'r' to reset the chat")
+        rprint("- 'h' to display these commands")

+    rprint(f"[INFO] Starting chat session with {args.model}.")
+    print_help()
    prompt_cache = make_prompt_cache(model, args.max_kv_size)
    while True:
-        prompt = _draw_corridor_prompt()
-        query = input(prompt)
-        if rank == 0:
-            # Cursor is now on the bottom-rule row; advance past it.
-            sys.stdout.write("\n")
-            sys.stdout.flush()
+        query = input(">> " if rank == 0 else "")
        if query == "q":
-            if rank == 0:
-                console.print("[ui.muted]bye[/ui.muted]")
            break
        if query == "r":
            prompt_cache = make_prompt_cache(model, args.max_kv_size)
-            if rank == 0:
-                console.print(
-                    "  [ui.good]reset[/ui.good] [ui.muted]context cleared[/ui.muted]"
-                )
            continue
        if query == "h":
-            if rank == 0:
-                print_chat_help(console)
+            print_help()
            continue
        messages = []
        if args.system_prompt is not None:
@@ -178,7 +142,6 @@ def main():
            messages,
            add_generation_prompt=True,
        )
-        last_response = None
        for response in stream_generate(
            model,
            tokenizer,
@@ -196,15 +159,7 @@ def main():
            prompt_cache=prompt_cache,
        ):
            rprint(response.text, flush=True, end="")
-            last_response = response
        rprint()
-        if rank == 0 and last_response is not None:
-            console.print(
-                f"  [ui.muted]{last_response.generation_tokens} tokens · "
-                f"{last_response.generation_tps:.1f} tok/s · "
-                f"prompt {last_response.prompt_tps:.1f} tok/s · "
-                f"peak {last_response.peak_memory:.2f} GB[/ui.muted]"
-            )


 if __name__ == "__main__":
@@ -1,172 +0,0 @@
-# Copyright © 2024 Apple Inc.
-
-"""Shared UI helpers for the mlx_lm command-line tools.
-
-Centralizes the rich-based panel/progress/prompt rendering used by the chat
-and training entry points. The theme is hardcoded for a light terminal
-background.
-"""
-
-import os
-import re
-import shutil
-import sys
-
-from rich.box import ROUNDED
-from rich.console import Console
-from rich.panel import Panel
-from rich.progress import (
-    Progress,
-    ProgressColumn,
-    TextColumn,
-)
-from rich.text import Text
-from rich.theme import Theme
-
-
-def _terminal_width(default: int = 120) -> int:
-    """Best-effort terminal width.
-
-    Under launchers like ``mlx.launch`` the worker's stdout is a pipe, so
-    Rich's auto-detection falls back to 80 columns. Honor an explicit
-    ``MLX_LM_WIDTH`` override, then ``COLUMNS``, then a real TTY query, and
-    finally a generous default that's nicer than 80 on modern terminals.
-    """
-    for var in ("MLX_LM_WIDTH", "COLUMNS"):
-        value = os.environ.get(var)
-        if value and value.isdigit():
-            return int(value)
-    width = shutil.get_terminal_size(fallback=(0, 0)).columns
-    return width if width > 0 else default
-
-
-def _make_theme() -> Theme:
-    return Theme(
-        {
-            "ui.strong": "bold #000000",
-            "ui.label": "#2a2a2a",
-            "ui.muted": "grey42",
-            "ui.heading": "bold #1a1a1a",
-            "ui.dim": "grey62",
-            "ui.accent": "bold purple",
-            "ui.border": "blue",
-            "ui.good": "bold green",
-            "ui.warn": "yellow",
-            "progress.percentage": "bold blue",
-        }
-    )
-
-
-def make_console(**kwargs) -> Console:
-    """Return a rich Console pre-loaded with the mlx_lm theme."""
-    kwargs.setdefault("highlight", False)
-    kwargs.setdefault("color_system", "truecolor")
-    kwargs.setdefault("width", _terminal_width())
-    return Console(theme=_make_theme(), **kwargs)
-
-
-def print_header_panel(
-    console: Console, title: str, rows: list[tuple[str, str]]
-) -> None:
-    """Render the boxed header used by the chat and training entry points."""
-    label_w = max(len(k) for k, _ in rows)
-    body = "\n".join(
-        f"  [ui.label]{k.ljust(label_w)}[/ui.label]   [ui.strong]{v}[/ui.strong]"
-        for k, v in rows
-    )
-    console.print(
-        Panel(
-            body,
-            title=f"[ui.accent]{title}[/ui.accent]",
-            title_align="left",
-            border_style="ui.border",
-            box=ROUNDED,
-            padding=(0, 2),
-        )
-    )
-
-
-def print_chat_help(console: Console) -> None:
-    console.print(
-        "  [ui.label]commands[/ui.label]    "
-        "[ui.strong]q[/ui.strong] [ui.muted]exit[/ui.muted]   "
-        "[ui.strong]r[/ui.strong] [ui.muted]reset[/ui.muted]   "
-        "[ui.strong]h[/ui.strong] [ui.muted]help[/ui.muted]"
-    )
-
-
-def make_corridor_prompt(console: Console):
-    """Return a callable that draws the chat input corridor.
-
-    The returned callable draws the top/bottom rules around the input line,
-    repositions the cursor onto the middle line, and returns the styled
-    "›" prompt string. Pass that string to ``input()`` so readline treats
-    the marker as part of the prompt — otherwise backspace will erase it.
-    """
-
-    _ANSI_RE = re.compile(r"\x1b\[[0-9;]*[A-Za-z]")
-
-    def _readline_safe(text: str) -> str:
-        # Wrap escape sequences in \x01..\x02 so readline doesn't count
-        # them when computing the prompt's visible width.
-        return _ANSI_RE.sub(lambda m: f"\x01{m.group(0)}\x02", text)
-
-    def _draw() -> str:
-        width = console.width
-        dashes = "─" * max(width - 1, 10)
-        with console.capture() as cap:
-            console.print(f"[ui.muted]{dashes}[/ui.muted]")
-            console.print()
-            console.print(f"[ui.muted]{dashes}[/ui.muted]")
-        sys.stdout.write(cap.get())
-        # Move the cursor up two rows back onto the blank middle line.
-        sys.stdout.write("\033[2A\r")
-        sys.stdout.flush()
-        with console.capture() as cap2:
-            console.print("[ui.accent]›[/ui.accent] ", end="")
-        return _readline_safe(cap2.get())
-
-    return _draw
-
-
-class SquareBar(ProgressColumn):
-    """Progress bar rendered with █/░ blocks plus eighth-block sub-precision."""
-
-    _EIGHTHS = "▏▎▍▌▋▊▉"  # 1/8 .. 7/8
-
-    def __init__(self, bar_width: int = 40, complete_style: str = "blue"):
-        super().__init__()
-        self.bar_width = bar_width
-        self.complete_style = complete_style
-
-    def render(self, task):
-        if not task.total:
-            return Text("░" * self.bar_width, style="ui.dim")
-        pct = min(max(task.completed / task.total, 0.0), 1.0)
-        total_eighths = int(pct * self.bar_width * 8)
-        full = total_eighths // 8
-        rem = total_eighths % 8
-        text = Text()
-        text.append("█" * full, style=self.complete_style)
-        used = full
-        if rem > 0 and full < self.bar_width:
-            text.append(self._EIGHTHS[rem - 1], style=self.complete_style)
-            used += 1
-        text.append("░" * (self.bar_width - used), style="ui.dim")
-        return text
-
-
-def make_train_progress(console: Console, *, disable: bool = False) -> Progress:
-    return Progress(
-        TextColumn("[bold blue]train[/bold blue]"),
-        SquareBar(bar_width=30, complete_style="blue"),
-        TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
-        TextColumn("[ui.muted]·[/ui.muted]"),
-        TextColumn(
-            "[bold blue]{task.completed:>5,}[/bold blue]"
-            "[ui.muted]/{task.total:,}[/ui.muted]"
-        ),
-        console=console,
-        transient=False,
-        disable=disable,
-    )
@@ -27,7 +27,7 @@ prompts = [

 # Set `verbose=True` to see generation statistics
 result = batch_generate(
-    model, tokenizer, prompts, verbose=False, return_prompt_caches=True, max_tokens=2048
+    model, tokenizer, prompts, verbose=False, return_prompt_caches=True
 )
 print(result.texts[-1])

@@ -12,7 +12,6 @@ import mlx.optimizers as optim
 import numpy as np
 import yaml

-from .cli_ui import make_console, print_header_panel
 from .tuner.callbacks import get_reporting_callbacks
 from .tuner.datasets import CacheDataset, load_dataset
 from .tuner.trainer import TrainingArgs, TrainingCallback, evaluate, train
@@ -24,12 +23,6 @@ from .tuner.utils import (
 )
 from .utils import _parse_size, load, save_config

-
-def printf(*args, **kwargs):
-    if mx.distributed.init().rank() == 0:
-        print(*args, **kwargs)
-
-
 yaml_loader = yaml.SafeLoader
 yaml_loader.add_implicit_resolver(
    "tag:yaml.org,2002:float",
@@ -253,7 +246,7 @@ def train_model(

    # Resume from weights if provided
    if args.resume_adapter_file is not None:
-        printf(f"Loading fine-tuned weights from {args.resume_adapter_file}")
+        print(f"Loading fine-tuned weights from {args.resume_adapter_file}")
        model.load_weights(args.resume_adapter_file, strict=False)

    print_trainable_parameters(model)
@@ -298,8 +291,6 @@ def train_model(

    opt = opt_class(learning_rate=lr, **optimizer_config)

-    _print_run_header(args)
-
    # Train model
    train(
        model=model,
@@ -311,50 +302,18 @@ def train_model(
    )


-def _print_run_header(args):
-    rank = mx.distributed.init().rank()
-    if rank != 0:
-        return
-
-    type_label = args.fine_tune_type
-    if args.fine_tune_type in ("lora", "dora"):
-        rank = args.lora_parameters.get("rank", "?")
-        type_label = f"{args.fine_tune_type} · {args.num_layers} layers · rank {rank}"
-    elif args.fine_tune_type == "full":
-        type_label = f"full · {args.num_layers} layers"
-
-    lr = (
-        args.learning_rate
-        if isinstance(args.learning_rate, (int, float))
-        else "schedule"
-    )
-    lr_str = f"{lr:.1e}" if isinstance(lr, (int, float)) else lr
-
-    rows = [
-        ("model", str(args.model)),
-        ("type", type_label),
-        ("dataset", str(args.data)),
-        ("optimizer", f"{args.optimizer} · lr {lr_str}"),
-        ("batch · iters", f"{args.batch_size} · {args.iters:,}"),
-        ("max seq", f"{args.max_seq_length:,}"),
-    ]
-    print_header_panel(make_console(), "mlx_lm.lora", rows)
-
-
 def evaluate_model(args, model: nn.Module, test_set):
-    rank = mx.distributed.init().rank()
    test_loss = evaluate(
        model=model,
        dataset=CacheDataset(test_set),
        batch_size=args.batch_size,
        num_batches=args.test_batches,
        max_seq_length=args.max_seq_length,
-        progress=(rank == 0),
    )

    test_ppl = math.exp(test_loss)

-    printf(f"Test loss {test_loss:.3f}, Test ppl {test_ppl:.3f}.")
+    print(f"Test loss {test_loss:.3f}, Test ppl {test_ppl:.3f}.")


 def run(args, training_callback: TrainingCallback = None):
@@ -366,10 +325,10 @@ def run(args, training_callback: TrainingCallback = None):
        config=vars(args),
    )

-    printf("Loading pretrained model")
+    print("Loading pretrained model")
    model, tokenizer = load(args.model, tokenizer_config={"trust_remote_code": True})

-    printf("Loading datasets")
+    print("Loading datasets")
    train_set, valid_set, test_set = load_dataset(args, tokenizer)

    if args.test and not args.train:
@@ -378,13 +337,13 @@ def run(args, training_callback: TrainingCallback = None):
            load_adapters(model, args.adapter_path)

    elif args.train:
-        printf("Training")
+        print("Training")
        train_model(args, model, train_set, valid_set, training_callback)
    else:
        raise ValueError("Must provide at least one of --train or --test")

    if args.test:
-        printf("Testing")
+        print("Testing")
        evaluate_model(args, model, test_set)


@@ -392,11 +351,10 @@ def main():
    os.environ["TOKENIZERS_PARALLELISM"] = "true"
    parser = build_parser()
    args = parser.parse_args()
-
    config = args.config
    args = vars(args)
    if config:
-        printf("Loading configuration file", config)
+        print("Loading configuration file", config)
        with open(config, "r") as file:
            config = yaml.load(file, yaml_loader)
        # Prefer parameters from command-line arguments
@@ -167,8 +167,7 @@ class Model(nn.Module):
        self.args = args
        self.model_type = args.model_type
        self.model = ApertusModel(args)
-        if not args.tie_word_embeddings:
-            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)

    def __call__(
        self,
@@ -176,18 +175,12 @@ class Model(nn.Module):
        cache: Optional[Any] = None,
    ) -> mx.array:
        out = self.model(inputs, cache)
-        if self.args.tie_word_embeddings:
-            out = self.model.embed_tokens.as_linear(out)
-        else:
-            out = self.lm_head(out)
-        return out
+        return self.lm_head(out)

    def sanitize(self, weights):
        for k, v in weights.items():
            if k.endswith("alpha_p") or k.endswith("alpha_n"):
                weights[k] = v.squeeze()
-        if self.args.tie_word_embeddings:
-            weights.pop("lm_head.weight", None)
        return weights

    @property
@@ -1,13 +1,11 @@
 # Copyright © 2023-2024 Apple Inc.

 import copy
-from collections import deque
-from dataclasses import dataclass
 from typing import Any, Dict, List, Optional

 import mlx.core as mx
 import mlx.nn as nn
-from mlx.utils import tree_flatten, tree_map, tree_reduce, tree_unflatten
+from mlx.utils import tree_flatten, tree_map, tree_unflatten

 from .base import create_causal_mask

@@ -603,18 +601,6 @@ class ArraysCache(_BaseCache):
        if left_padding:
            self.left_padding = mx.array(left_padding)

-    @property
-    def batch_size(self):
-        for c in self.cache:
-            if c is not None:
-                return c.shape[0]
-        if self.left_padding is not None:
-            return self.left_padding.size
-        elif self.lengths is not None:
-            return self.lengths.size
-        else:
-            return 1
-
    def __setitem__(self, idx, value):
        self.cache[idx] = value

@@ -633,42 +619,13 @@ class ArraysCache(_BaseCache):
        """
        In-place filter to keep just the given indices in the cache.
        """
-        self.cache = [c[batch_indices] if c is not None else None for c in self.cache]
-        if self.left_padding is not None:
-            self.left_padding = self.left_padding[batch_indices]
-        if self.lengths is not None:
-            self.lengths = self.lengths[batch_indices]
+        self.cache = [c[batch_indices] for c in self.cache]

    def extend(self, other):
        """
        In-place extend this cache with the other cache.
        """
-
-        a_batch = self.batch_size
-        b_batch = other.batch_size
-
-        def cat(a, b):
-            shape = dtype = None
-            if a is not None:
-                shape = a.shape
-                dtype = a.dtype
-            if b is not None:
-                shape = b.shape
-                dtype = b.dtype
-
-            if shape is None:
-                return None
-
-            if a is None:
-                a = mx.zeros((a_batch,) + shape[1:], dtype=dtype)
-            if b is None:
-                b = mx.zeros((b_batch,) + shape[1:], dtype=dtype)
-
-            return mx.concatenate([a, b])
-
-        self.cache = [cat(c, o) for c, o in zip(self.cache, other.cache)]
-        self.left_padding = cat(self.left_padding, other.left_padding)
-        self.lengths = cat(self.lengths, other.lengths)
+        self.cache = [mx.concatenate([c, o]) for c, o in zip(self.cache, other.cache)]

    def extract(self, idx):
        cache = ArraysCache(len(self.cache))
@@ -703,12 +660,6 @@ class ArraysCache(_BaseCache):
        n_state = len(caches[0].cache)
        B = len(caches)
        cache = cls(n_state)
-
-        # All caches are empty so return early
-        if all(c.empty() for c in caches):
-            cache.left_padding = mx.array([0] * B)
-            return cache
-
        for e in range(n_state):
            c_init = next(iter(c[e] for c in caches if c[e] is not None))
            shape = list(c_init.shape)
@@ -1017,18 +968,16 @@ class BatchKVCache(_BaseCache):
        """
        In-place filter to keep just the given indices in the cache.
        """
-        if self.keys is not None:
-            self.keys = self.keys[batch_indices]
-            self.values = self.values[batch_indices]
+        self.keys = self.keys[batch_indices]
+        self.values = self.values[batch_indices]
        self.offset = self.offset[batch_indices]
        self.left_padding = self.left_padding[batch_indices]

        # Shift left to reduce padding
        min_left_pad = self.left_padding.min().item()
        if min_left_pad > 0:
-            if self.keys is not None:
-                self.keys = self.keys[..., min_left_pad:, :]
-                self.values = self.values[..., min_left_pad:, :]
+            self.keys = self.keys[..., min_left_pad:, :]
+            self.values = self.values[..., min_left_pad:, :]
            self._idx -= min_left_pad
            self.left_padding -= min_left_pad

@@ -1036,31 +985,15 @@ class BatchKVCache(_BaseCache):
        """
        In-place extend this cache with the other cache.
        """
-        if self.keys is None and other.keys is None:
-            self.left_padding = mx.concatenate([self.left_padding, other.left_padding])
-            self.offset = mx.concatenate([self.offset, other.offset])
-            return
-
        max_idx = max(self._idx, other._idx)
-        L1 = L2 = 0
-        if self.keys is not None:
-            B, H, L1, D = self.keys.shape
-            M = self.values.shape[3]
-        if other.keys is not None:
-            B, H, L2, D = other.keys.shape
-            M = other.values.shape[3]
-        max_size = max(L1, L2)
+        max_size = max(self.keys.shape[2], other.keys.shape[2])

        # Pad the keys and values so they are right-justified
        # with the index and the same size
        def pad(c):
-            k, v = c.keys, c.values
-            if k is None:
-                Bc = c.offset.shape[0]
-                k = mx.array([]).reshape(Bc, H, 0, D)
-                v = mx.array([]).reshape(Bc, H, 0, M)
            left = max_idx - c._idx
-            right = max_size - k.shape[2] - left
+            right = max_size - c.keys.shape[2] - left
+            k, v = c.keys, c.values
            if right < 0:
                k = k[..., :right, :]
                v = v[..., :right, :]
@@ -1089,11 +1022,6 @@ class BatchKVCache(_BaseCache):
    def merge(cls, caches):
        lengths = [c.size() for c in caches]
        max_length = max(lengths)
-
-        # No cache has content so make an empty one
-        if max_length == 0:
-            return BatchKVCache([0] * len(caches))
-
        padding = [max_length - l for l in lengths]
        B = len(caches)
        H = max(c.keys.shape[1] for c in caches if c.keys is not None)
@@ -1117,9 +1045,6 @@ class BatchKVCache(_BaseCache):

        return cache

-    def size(self):
-        return self._idx
-
    def empty(self):
        return self.keys is None

@@ -1360,9 +1285,8 @@ class BatchRotatingKVCache(_BaseCache):
        """
        In-place filter to keep just the given indices in the cache.
        """
-        if self.keys is not None:
-            self.keys = self.keys[batch_indices]
-            self.values = self.values[batch_indices]
+        self.keys = self.keys[batch_indices]
+        self.values = self.values[batch_indices]
        self.offset = self.offset[batch_indices]
        self.left_padding = self.left_padding[batch_indices]

@@ -1370,33 +1294,17 @@ class BatchRotatingKVCache(_BaseCache):
        """
        In-place extend this cache with the other cache.
        """
-        if self.keys is None and other.keys is None:
-            self.left_padding = mx.concatenate([self.left_padding, other.left_padding])
-            self.offset = mx.concatenate([self.offset, other.offset])
-            return
-
        if (self.rotated != other.rotated) or self._idx != other._idx:
            self._temporal_order()
            other._temporal_order()

        max_idx = max(self._idx, other._idx)
-        L1 = L2 = 0
-        if self.keys is not None:
-            B, H, L1, D = self.keys.shape
-            M = self.values.shape[3]
-        if other.keys is not None:
-            B, H, L2, D = other.keys.shape
-            M = other.values.shape[3]
-        max_size = max(L1, L2)
+        max_size = max(self.keys.shape[2], other.keys.shape[2])

        def pad(c):
            left = max_idx - c._idx
+            right = max_size - c.keys.shape[2] - left
            k, v = c.keys, c.values
-            if k is None:
-                Bc = c.offset.shape[0]
-                k = mx.array([]).reshape(Bc, H, 0, D)
-                v = mx.array([]).reshape(Bc, H, 0, M)
-            right = max_size - k.shape[2] - left
            if right < 0:
                k = k[..., :right, :]
                v = v[..., :right, :]
@@ -1415,10 +1323,9 @@ class BatchRotatingKVCache(_BaseCache):
        self._offset = max(self._offset, other._offset)

    def extract(self, idx):
-        mx.eval(self.left_padding, self.offset)
        cache = RotatingKVCache(self.max_size)
-        padding = max(0, self.left_padding.tolist()[idx])
-        offset = self.offset.tolist()[idx]
+        padding = self.left_padding[idx].item()
+        offset = self.offset[idx].item()
        cache.keys = self.keys[idx : idx + 1]
        cache.values = self.values[idx : idx + 1]
        cache._idx = self._idx
@@ -1442,11 +1349,6 @@ class BatchRotatingKVCache(_BaseCache):
        offsets = [c.offset for c in caches]
        lengths = [c.size() for c in caches]
        max_length = max(lengths)
-
-        # No cache has content so make an empty one
-        if max_length == 0:
-            return cls(caches[0].max_size, [0] * len(caches))
-
        padding = [max_length - l for l in lengths]
        B = len(caches)
        H = max(c.keys.shape[1] for c in caches if c.keys is not None)
@@ -1456,11 +1358,11 @@ class BatchRotatingKVCache(_BaseCache):

        keys = mx.zeros((B, H, max_length, Dk), dtype=dt)
        values = mx.zeros((B, H, max_length, Dv), dtype=dt)
-        for i, (p, l, c) in enumerate(zip(padding, lengths, caches)):
+        for i, (p, c) in enumerate(zip(padding, caches)):
            if c.keys is None:
                continue
-            keys[i : i + 1, :, p : p + l] = c._temporal_order(c.keys)[..., -l:, :]
-            values[i : i + 1, :, p : p + l] = c._temporal_order(c.values)[..., -l:, :]
+            keys[i : i + 1, :, p : p + c._idx] = c._temporal_order(c.keys)
+            values[i : i + 1, :, p : p + c._idx] = c._temporal_order(c.values)

        cache = cls(caches[0].max_size, padding)
        cache.keys = keys
@@ -1471,9 +1373,6 @@ class BatchRotatingKVCache(_BaseCache):

        return cache

-    def size(self):
-        return min(self._offset, self.max_size)
-
    def empty(self):
        return self.keys is None

@@ -1482,282 +1381,3 @@ class BatchRotatingKVCache(_BaseCache):
        if self.keys is None:
            return 0
        return self.keys.nbytes + self.values.nbytes
-
-
-class TokenBuffer:
-    """A simple token buffer that can be efficiently appended to in a similar
-    fashion to the KVCache.
-
-    Perhaps these could share some logic in the future.
-    """
-
-    step = 256
-
-    def __init__(self, tokens=[]):
-        self._buffer = mx.array(tokens, dtype=mx.int32)
-        self._size = len(tokens)
-
-    def update_and_fetch(self, tokens):
-        start = self._size
-        end = start + len(tokens)
-
-        new_size = ((end + self.step - 1) // self.step) * self.step
-        if new_size > self._buffer.size:
-            self._buffer = mx.concatenate(
-                [self._buffer, mx.zeros(new_size - self._buffer.size, dtype=mx.int32)]
-            )
-        self._buffer[start:end] = tokens
-        self._size = end
-
-        return self._buffer[:end]
-
-    @property
-    def state(self):
-        return self._buffer
-
-    @property
-    def tokens(self):
-        return self._buffer[: self._size]
-
-
-@dataclass
-class PromptTrieResult:
-    model: Any
-    exact: Optional[List[int]]  # Exact match found
-    shorter: Optional[List[int]]  # Longest prefix with a value
-    longer: Optional[List[int]]  # Shortest value that extends beyond tokens
-    common_prefix: int  # Length of common prefix with any path
-
-
-class PromptTrie:
-    def __init__(self):
-        self._trie = {}
-
-    def add(self, model: Any, tokens: List[int], value: Any):
-        if model not in self._trie:
-            self._trie[model] = {}
-
-        current = self._trie[model]
-        for tok in tokens:
-            if tok not in current:
-                current[tok] = {}
-            current = current[tok]
-        prev = current.get("__value__", None)
-        current["__value__"] = value
-        return prev
-
-    def get(self, model: Any, tokens: List[int]):
-        current = self._trie[model]
-        for tok in tokens:
-            current = current[tok]
-        return current["__value__"]
-
-    def pop(self, model: Any, tokens: List[int]):
-        path = [self._trie[model]]
-        for tok in tokens:
-            path.append(path[-1][tok])
-        value = path[-1].pop("__value__")
-        for i in range(len(tokens), 0, -1):
-            node = path[i]
-            parent = path[i - 1]
-            tok = tokens[i - 1]
-            if len(node) > 0:
-                break
-            del parent[tok]
-        return value
-
-    def pop_prefixes(self, model: Any, tokens: List[int]):
-        values = []
-        current = self._trie[model]
-        for i, tok in enumerate(tokens):
-            if "__value__" in current:
-                values.append((i, current.pop("__value__")))
-            current = current[tok]
-        return values
-
-    def search(self, model: Any, tokens: List[int]) -> PromptTrieResult:
-        if model not in self._trie:
-            return PromptTrieResult(model, None, None, None, 0)
-
-        current = self._trie[model]
-
-        if not tokens and "__value__" in current:
-            return PromptTrieResult(model, [], None, None, 0)
-
-        # Walk the tokens as far as we can
-        last_index = -1
-        index = 0
-        while index < len(tokens) and tokens[index] in current:
-            current = current[tokens[index]]
-            if "__value__" in current:
-                last_index = index
-            index += 1
-
-        # Got an exact match
-        if last_index == len(tokens) - 1 >= 0:
-            return PromptTrieResult(model, tokens, None, None, 0)
-
-        # Check if we found a prefix at any point
-        shorter = None
-        if last_index > 0:
-            shorter = tokens[: last_index + 1]
-
-        # Check for sequences that are longer
-        longer = None
-        common_prefix = index
-        if index > 0:
-            best = None
-            stack = [(current, [])]
-            while stack:
-                current, extra = stack.pop()
-                if "__value__" in current:
-                    if best is None or len(extra) < len(best):
-                        best = extra
-                elif best is None or len(extra) < len(best):
-                    for tok in current:
-                        stack.append((current[tok], extra + [tok]))
-            longer = tokens[:index] + best
-        return PromptTrieResult(model, None, shorter, longer, common_prefix)
-
-
-class LRUPromptCache:
-    @dataclass
-    class CacheEntry:
-        prompt_cache: List[Any]
-        nbytes: int
-        cache_type: str
-
-    class CacheOrder:
-        def __init__(self, ordering: List[str] = ["assistant", "user", "system"]):
-            self._ordering = ordering
-            self._lrus = {k: deque() for k in ordering}
-
-        def __len__(self):
-            return sum(len(lru) for lru in self._lrus.values())
-
-        def push(self, model: Any, tokens: List[Any], cache_type: str = "assistant"):
-            self._lrus[cache_type].append((model, tokens))
-
-        def remove(self, model: Any, tokens: List[Any]):
-            for cache_type in self._ordering:
-                try:
-                    self._lrus[cache_type].remove((model, tokens))
-                    break
-                except ValueError:
-                    pass
-
-        def pop(self):
-            i = 0
-            while i + 1 < len(self._ordering):
-                lru_a = self._lrus[self._ordering[i]]
-                lru_b = self._lrus[self._ordering[i + 1]]
-                if lru_a and len(lru_a) >= len(lru_b):
-                    return lru_a.popleft()
-                i += 1
-            return lru_b.popleft()
-
-    def __init__(self, max_size: int = 10, max_bytes: int = 1 << 63):
-        self.max_size = max_size
-        self.max_bytes = max_bytes
-        self._trie = PromptTrie()
-        self._lru = LRUPromptCache.CacheOrder()
-        self._n_bytes = 0
-        self._n_bytes_by_type = {k: 0 for k in self._lru._ordering}
-
-    def __len__(self):
-        return len(self._lru)
-
-    @property
-    def nbytes(self):
-        return self._n_bytes
-
-    def fetch_nearest_cache(self, model: Any, tokens: List[int]):
-        result = self._trie.search(model, tokens)
-        if result.exact is not None:
-            cache_entry = self._trie.get(result.model, result.exact)
-            return copy.deepcopy(cache_entry.prompt_cache), []
-
-        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._trie.get(result.model, result.longer)
-            if can_trim_prompt_cache(cache_entry.prompt_cache):
-                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, num_to_trim)
-                return cache, tokens[prefix:]
-
-        if short_length > 0:
-            cache_entry = self._trie.get(result.model, result.shorter)
-            return copy.deepcopy(cache_entry.prompt_cache), tokens[short_length:]
-
-        return None, tokens
-
-    def insert_cache(
-        self,
-        model: Any,
-        tokens: List[int],
-        prompt_cache: List[Any],
-        *,
-        cache_type: str = "assistant",
-    ):
-        # Make the cache entry
-        entry = LRUPromptCache.CacheEntry(
-            prompt_cache, sum(c.nbytes for c in prompt_cache), cache_type
-        )
-
-        # Insert into the trie and update the byte counter and lru position
-        self._n_bytes += entry.nbytes
-        self._n_bytes_by_type[cache_type] += entry.nbytes
-        prev = self._trie.add(model, tokens, entry)
-        if prev is not None:
-            self._n_bytes -= prev.nbytes
-            self._n_bytes_by_type[prev.cache_type] -= prev.nbytes
-            self._lru.remove(model, tokens)
-        self._lru.push(model, tokens, cache_type)
-
-        # If it is a trimmable cache remove all prefixes cause they just take
-        # space
-        if can_trim_prompt_cache(prompt_cache):
-            for prefix_len, entry in self._trie.pop_prefixes(model, tokens):
-                self._n_bytes -= entry.nbytes
-                self._n_bytes_by_type[entry.cache_type] -= entry.nbytes
-                self._lru.remove(model, tokens[:prefix_len])
-
-        # Ensure we match the constraints
-        if len(self._lru) > self.max_size:
-            model, tokens = self._lru.pop()
-            entry = self._trie.pop(model, tokens)
-            self._n_bytes -= entry.nbytes
-            self._n_bytes_by_type[entry.cache_type] -= entry.nbytes
-        while self._n_bytes > self.max_bytes:
-            model, tokens = self._lru.pop()
-            entry = self._trie.pop(model, tokens)
-            self._n_bytes -= entry.nbytes
-            self._n_bytes_by_type[entry.cache_type] -= entry.nbytes
-
-    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()
-            entry = self._trie.pop(model, tokens)
-            self._n_bytes -= entry.nbytes
-            self._n_bytes_by_type[entry.cache_type] -= entry.nbytes
-        while self._n_bytes > n_bytes:
-            model, tokens = self._lru.pop()
-            entry = self._trie.pop(model, tokens)
-            self._n_bytes -= entry.nbytes
-            self._n_bytes_by_type[entry.cache_type] -= entry.nbytes
-
-    def stats_by_type(self):
-        result = {}
-        for cache_type in self._lru._ordering:
-            result[cache_type] = {
-                "n_sequences": len(self._lru._lrus[cache_type]),
-                "n_bytes": self._n_bytes_by_type[cache_type],
-            }
-        return result
@@ -87,15 +87,19 @@ class Indexer(nn.Module):
        b, s, _ = x.shape
        q = self.wq_b(qr)
        q = q.reshape(b, s, self.n_heads, self.head_dim).swapaxes(1, 2)
-        k = self.wk(x)
-        k = self.k_norm(k)
-        k = mx.reshape(k, (b, 1, s, self.head_dim))
+        q_pe, q_nope = mx.split(q, [self.rope_head_dim], axis=-1)

        offset = cache.offset if cache is not None else 0

-        q = self.rope(q, offset=offset)
-        k = self.rope(k, offset=offset)
+        q_pe = self.rope(q_pe, offset=offset)
+        q = mx.concatenate([q_pe, q_nope], axis=-1)

+        k = self.wk(x)
+        k = self.k_norm(k)
+        k = mx.reshape(k, (b, 1, s, self.head_dim))
+        k_pe, k_nope = mx.split(k, [self.rope_head_dim], axis=-1)
+        k_pe = self.rope(k_pe, offset=offset)
+        k = mx.concatenate([k_pe, k_nope], axis=-1)
        if cache is not None:
            k, _ = cache.update_and_fetch(k, mx.zeros([b, 1, s, 0]))
        if k.shape[2] <= self.index_topk:
@@ -217,8 +221,7 @@ class DeepseekV32Attention(nn.Module):
                    mx.broadcast_to(idx, idx.shape[:-1] + (k_pe.shape[-1],)),
                    axis=2,
                )
-                if mask is not None:
-                    mask = mx.take_along_axis(mask, topk_indices, axis=-1)
+                mask = None
            else:
                shape = list(topk_indices.shape)
                shape[-1] = kv_latent.shape[2]
@@ -81,8 +81,6 @@ def _make_gated_delta_kernel(has_mask=False, vectorized=False):
            if (thread_index_in_simdgroup == 0) {{
              y[dv_idx] = static_cast<InT>(out);
            }}
-          }} else {{
-            y[dv_idx] = static_cast<InT>(0);
          }}
          // Increment data pointers to next time step
          q_ += Hk * Dk;
@@ -1,92 +0,0 @@
-# Copyright © 2025 Apple Inc.
-
-from dataclasses import dataclass
-from typing import Optional
-
-import mlx.core as mx
-import mlx.nn as nn
-from mlx.utils import tree_flatten, tree_unflatten
-
-from . import gemma4_text
-from .base import BaseModelArgs
-
-
-@dataclass
-class ModelArgs(BaseModelArgs):
-    model_type: str = "gemma4"
-    text_config: dict = None
-    vocab_size: int = 262144
-
-    def __post_init__(self):
-        if self.text_config is None:
-            self.text_config = {}
-        self.text_config["vocab_size"] = self.vocab_size
-        self.text_config["num_attention_heads"] = self.text_config.get(
-            "num_attention_heads", 8
-        )
-        self.text_config["num_key_value_heads"] = self.text_config.get(
-            "num_key_value_heads", 1
-        )
-
-
-class Model(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.args = args
-        self.model_type = args.model_type
-        self.language_model = gemma4_text.Model(
-            gemma4_text.ModelArgs.from_dict(args.text_config)
-        )
-
-    def __call__(
-        self,
-        inputs: mx.array,
-        cache=None,
-        input_embeddings: Optional[mx.array] = None,
-        per_layer_inputs: Optional[mx.array] = None,
-    ):
-        return self.language_model(
-            inputs,
-            cache=cache,
-            input_embeddings=input_embeddings,
-            per_layer_inputs=per_layer_inputs,
-        )
-
-    def sanitize(self, weights):
-        new_weights = {}
-        for k, v in weights.items():
-            starts_w_model = k.startswith("model.")
-
-            k = k.removeprefix("model.")
-            if k.startswith(
-                (
-                    "vision_tower",
-                    "multi_modal_projector",
-                    "audio_tower",
-                    "embed_audio",
-                    "embed_vision",
-                )
-            ):
-                continue
-
-            if not starts_w_model:
-                new_weights[k] = v
-                continue
-
-            if k.startswith("language_model"):
-                k = k.replace("language_model.", "language_model.model.")
-
-            new_weights[k] = v
-
-        return self.language_model.sanitize(new_weights)
-
-    @property
-    def layers(self):
-        return self.language_model.layers
-
-    @property
-    def quant_predicate(self):
-        return self.language_model.quant_predicate
-
-    def make_cache(self):
-        return self.language_model.make_cache()
@@ -1,688 +0,0 @@
-# Copyright © 2025 Apple Inc.
-
-from dataclasses import dataclass
-from functools import partial
-from typing import Any, Dict, List, Optional
-
-import mlx.core as mx
-import mlx.nn as nn
-
-from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
-from .cache import KVCache, RotatingKVCache, _BaseCache
-from .rope_utils import initialize_rope
-from .switch_layers import SwitchGLU
-
-
-@dataclass
-class ModelArgs(BaseModelArgs):
-    model_type: str = "gemma4_text"
-    hidden_size: int = 1536
-    num_hidden_layers: int = 35
-    intermediate_size: int = 6144
-    num_attention_heads: int = 8
-    head_dim: int = 256
-    global_head_dim: int = 512
-    global_partial_rotary_factor: float = 0.25
-    rms_norm_eps: float = 1e-6
-    vocab_size: int = 262144
-    vocab_size_per_layer_input: int = 262144
-    num_key_value_heads: int = 1
-    num_global_key_value_heads: Optional[int] = None
-    num_kv_shared_layers: int = 20
-    pad_token_id: int = 0
-    hidden_size_per_layer_input: int = 256
-    rope_traditional: bool = False
-    partial_rotary_factor: float = 1.0
-    rope_parameters: Optional[Dict] = None
-    sliding_window: int = 512
-    sliding_window_pattern: int = 5
-    max_position_embeddings: int = 131072
-    attention_k_eq_v: bool = False
-    final_logit_softcapping: float = 30.0
-    use_double_wide_mlp: bool = True
-    enable_moe_block: bool = False
-    num_experts: Optional[int] = None
-    top_k_experts: Optional[int] = None
-    moe_intermediate_size: Optional[int] = None
-    layer_types: Optional[List[str]] = None
-    tie_word_embeddings: bool = True
-
-    def __post_init__(self):
-        if self.rope_parameters is None:
-            self.rope_parameters = {
-                "full_attention": {
-                    "partial_rotary_factor": 0.25,
-                    "rope_theta": 1000000.0,
-                    "rope_type": "proportional",
-                },
-                "sliding_attention": {
-                    "partial_rotary_factor": 1.0,
-                    "rope_theta": 10000.0,
-                    "rope_type": "default",
-                },
-            }
-        if self.layer_types is None:
-            pattern = ["sliding_attention"] * (self.sliding_window_pattern - 1) + [
-                "full_attention"
-            ]
-            self.layer_types = (pattern * (self.num_hidden_layers // len(pattern) + 1))[
-                : self.num_hidden_layers
-            ]
-
-
-class RMSNormNoScale(nn.Module):
-    """RMSNorm without learnable scale."""
-
-    def __init__(self, dim: int, eps: float = 1e-6):
-        super().__init__()
-        self.eps = eps
-
-    def __call__(self, x: mx.array) -> mx.array:
-        return mx.fast.rms_norm(x, None, self.eps)
-
-
-@partial(mx.compile, shapeless=True)
-def logit_softcap(softcap, x):
-    return mx.tanh(x / softcap) * softcap
-
-
-@partial(mx.compile, shapeless=True)
-def _complete_square(x2, y2, xy):
-    return x2 + mx.expand_dims(y2, -1) - 2 * xy
-
-
-@partial(mx.compile, shapeless=True)
-def geglu(gate, x):
-    return nn.gelu_approx(gate) * x
-
-
-class MLP(nn.Module):
-    def __init__(self, config: ModelArgs, layer_idx: int = 0):
-        super().__init__()
-        first_kv_shared_layer_idx = (
-            config.num_hidden_layers - config.num_kv_shared_layers
-        )
-        is_kv_shared_layer = layer_idx >= first_kv_shared_layer_idx > 0
-        use_double_wide = config.use_double_wide_mlp and is_kv_shared_layer
-        intermediate_size = config.intermediate_size * (2 if use_double_wide else 1)
-
-        self.gate_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
-        self.down_proj = nn.Linear(intermediate_size, config.hidden_size, bias=False)
-        self.up_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
-
-    def __call__(self, x: mx.array) -> mx.array:
-        return self.down_proj(geglu(self.gate_proj(x), self.up_proj(x)))
-
-
-class Router(nn.Module):
-    """Expert router: norm -> scale -> project -> top-k -> renormalize."""
-
-    def __init__(self, config: ModelArgs):
-        super().__init__()
-        self.config = config
-        self.eps = config.rms_norm_eps
-        self.proj = nn.Linear(config.hidden_size, config.num_experts, bias=False)
-        self.scale = mx.ones((config.hidden_size,))
-        self.per_expert_scale = mx.ones((config.num_experts,))
-        self._root_size = config.hidden_size**-0.5
-
-    def __call__(self, x: mx.array):
-        x = mx.fast.rms_norm(x, self.scale * self._root_size, self.eps)
-
-        expert_scores = self.proj(x)
-
-        top_k_indices = mx.argpartition(
-            expert_scores, kth=-self.config.top_k_experts, axis=-1
-        )
-        top_k_indices = top_k_indices[..., -self.config.top_k_experts :]
-
-        top_k_weights = mx.take_along_axis(expert_scores, top_k_indices, axis=-1)
-        top_k_weights = mx.softmax(top_k_weights, axis=-1)
-        top_k_weights = top_k_weights * self.per_expert_scale[top_k_indices]
-
-        return top_k_indices, top_k_weights
-
-
-class GeGLU(nn.Module):
-    """GELU-gated linear unit activation for SwitchGLU."""
-
-    def __call__(self, x, gate):
-        return geglu(gate, x)
-
-
-class Experts(nn.Module):
-    """Sparse MoE using SwitchGLU with gather_mm."""
-
-    def __init__(self, config: ModelArgs):
-        super().__init__()
-
-        self.switch_glu = SwitchGLU(
-            input_dims=config.hidden_size,
-            hidden_dims=config.moe_intermediate_size,
-            num_experts=config.num_experts,
-            activation=GeGLU(),
-            bias=False,
-        )
-
-    def __call__(
-        self, x: mx.array, top_k_indices: mx.array, top_k_weights: mx.array
-    ) -> mx.array:
-        w = mx.expand_dims(top_k_weights, -1)
-        y = self.switch_glu(x, top_k_indices)
-
-        return (w * y).sum(-2)
-
-
-class Attention(nn.Module):
-    def __init__(self, config: ModelArgs, layer_idx: int):
-        super().__init__()
-        self.config = config
-        self.layer_idx = layer_idx
-        self.layer_type = config.layer_types[layer_idx]
-        self.is_sliding = self.layer_type == "sliding_attention"
-        self.has_kv = layer_idx < config.num_hidden_layers - config.num_kv_shared_layers
-
-        self.head_dim = (
-            config.global_head_dim
-            if self.layer_type == "full_attention"
-            and hasattr(config, "global_head_dim")
-            and config.global_head_dim
-            else config.head_dim
-        )
-
-        dim = config.hidden_size
-        self.n_heads = config.num_attention_heads
-
-        # K-eq-V for full attention layers (26B/31B models)
-        self.use_k_eq_v = config.attention_k_eq_v and not self.is_sliding
-        if self.use_k_eq_v and config.num_global_key_value_heads is not None:
-            self.n_kv_heads = config.num_global_key_value_heads
-        else:
-            self.n_kv_heads = config.num_key_value_heads
-
-        self.scale = 1.0
-
-        self.q_proj = nn.Linear(dim, self.n_heads * self.head_dim, bias=False)
-        if self.has_kv:
-            self.k_proj = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=False)
-            if not self.use_k_eq_v:
-                self.v_proj = nn.Linear(
-                    dim, self.n_kv_heads * self.head_dim, bias=False
-                )
-        self.o_proj = nn.Linear(self.n_heads * self.head_dim, dim, bias=False)
-
-        self.q_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps)
-        if self.has_kv:
-            self.k_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps)
-            self.v_norm = RMSNormNoScale(self.head_dim, eps=config.rms_norm_eps)
-
-        # RoPE (with partial rotation support)
-        layer_key = "sliding_attention" if self.is_sliding else "full_attention"
-        rope_params = config.rope_parameters.get(layer_key, {})
-        rope_theta = rope_params.get("rope_theta", 10000.0)
-        self.rope = initialize_rope(
-            dims=self.head_dim,
-            traditional=config.rope_traditional,
-            base=rope_theta,
-            scaling_config=rope_params,
-            max_position_embeddings=config.max_position_embeddings,
-        )
-
-    def __call__(
-        self,
-        x: mx.array,
-        mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
-        shared_kv: Optional[tuple] = None,
-        offset: Optional[Any] = None,
-    ) -> mx.array:
-        B, L, _ = x.shape
-
-        queries = self.q_proj(x).reshape(B, L, self.n_heads, self.head_dim)
-        queries = self.q_norm(queries)
-
-        if shared_kv is not None:
-            keys, values = shared_kv
-        elif not self.has_kv:
-            raise ValueError(
-                f"Layer {self.layer_idx} is a KV-shared layer but received no shared_kv"
-            )
-        else:
-            keys = self.k_proj(x).reshape(B, L, self.n_kv_heads, self.head_dim)
-            values = keys
-            if not self.use_k_eq_v:
-                values = self.v_proj(x).reshape(B, L, self.n_kv_heads, self.head_dim)
-
-            offset = mx.array(cache.offset) if cache is not None else 0
-
-            keys = self.k_norm(keys)
-            keys = keys.transpose(0, 2, 1, 3)
-            keys = self.rope(keys, offset=offset)
-
-            values = self.v_norm(values)
-            values = values.transpose(0, 2, 1, 3)
-
-        queries = queries.transpose(0, 2, 1, 3)
-        queries = self.rope(queries, offset=offset)
-
-        if cache is not None:
-            keys, values = cache.update_and_fetch(keys, values)
-
-        output = scaled_dot_product_attention(
-            queries, keys, values, cache=cache, scale=self.scale, mask=mask
-        )
-        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
-
-        return self.o_proj(output), (keys, values), offset
-
-
-class DecoderLayer(nn.Module):
-    def __init__(self, config: ModelArgs, layer_idx: int):
-        super().__init__()
-        self.config = config
-        self.layer_idx = layer_idx
-        self.layer_type = config.layer_types[layer_idx]
-        self.self_attn = Attention(config, layer_idx)
-        self.mlp = MLP(config, layer_idx)
-        self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_attention_layernorm = nn.RMSNorm(
-            config.hidden_size, eps=config.rms_norm_eps
-        )
-        self.pre_feedforward_layernorm = nn.RMSNorm(
-            config.hidden_size, eps=config.rms_norm_eps
-        )
-        self.post_feedforward_layernorm = nn.RMSNorm(
-            config.hidden_size, eps=config.rms_norm_eps
-        )
-
-        # MoE (26B model)
-        self.enable_moe = config.enable_moe_block
-        if self.enable_moe:
-            self.router = Router(config)
-            self.experts = Experts(config)
-            self.post_feedforward_layernorm_1 = nn.RMSNorm(
-                config.hidden_size, eps=config.rms_norm_eps
-            )
-            self.post_feedforward_layernorm_2 = nn.RMSNorm(
-                config.hidden_size, eps=config.rms_norm_eps
-            )
-            self.pre_feedforward_layernorm_2 = nn.RMSNorm(
-                config.hidden_size, eps=config.rms_norm_eps
-            )
-
-        # Per-layer input gating (2B/4B models)
-        self.hidden_size_per_layer_input = config.hidden_size_per_layer_input
-        if self.hidden_size_per_layer_input:
-            self.per_layer_input_gate = nn.Linear(
-                config.hidden_size, self.hidden_size_per_layer_input, bias=False
-            )
-            self.per_layer_projection = nn.Linear(
-                self.hidden_size_per_layer_input, config.hidden_size, bias=False
-            )
-            self.post_per_layer_input_norm = nn.RMSNorm(
-                config.hidden_size, eps=config.rms_norm_eps
-            )
-        else:
-            self.per_layer_input_gate = None
-            self.per_layer_projection = None
-            self.post_per_layer_input_norm = None
-
-        # Layer scalar
-        self.layer_scalar = mx.ones((1,))
-
-    def __call__(
-        self,
-        x: mx.array,
-        mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
-        per_layer_input: Optional[mx.array] = None,
-        shared_kv: Optional[tuple] = None,
-        offset: Optional[Any] = None,
-    ) -> mx.array:
-        residual = x
-
-        h = self.input_layernorm(x)
-        h, shared_kv, offset = self.self_attn(
-            h, mask, cache, shared_kv=shared_kv, offset=offset
-        )
-        h = self.post_attention_layernorm(h)
-        h = residual + h
-
-        residual = h
-
-        if self.enable_moe:
-            h1 = self.pre_feedforward_layernorm(h)
-            h1 = self.mlp(h1)
-            h1 = self.post_feedforward_layernorm_1(h1)
-
-            top_k_indices, top_k_weights = self.router(h)
-            h2 = self.pre_feedforward_layernorm_2(h)
-            h2 = self.experts(h2, top_k_indices, top_k_weights)
-            h2 = self.post_feedforward_layernorm_2(h2)
-
-            h = h1 + h2
-        else:
-            h = self.pre_feedforward_layernorm(h)
-            h = self.mlp(h)
-
-        h = self.post_feedforward_layernorm(h)
-        h = residual + h
-
-        # Per-layer input gating
-        if (
-            self.per_layer_input_gate is not None
-            and self.per_layer_projection is not None
-            and self.post_per_layer_input_norm is not None
-            and per_layer_input is not None
-        ):
-            residual = h
-            gate = self.per_layer_input_gate(h)
-            gate = nn.gelu_approx(gate)
-            gate = mx.multiply(gate, per_layer_input)
-            gate = self.per_layer_projection(gate)
-            gate = self.post_per_layer_input_norm(gate)
-            h = residual + gate
-
-        if self.layer_scalar is not None:
-            h = h * self.layer_scalar
-
-        return h, shared_kv, offset
-
-
-class Gemma4TextModel(nn.Module):
-    def __init__(self, config: ModelArgs):
-        super().__init__()
-        self.config = config
-        self.vocab_size = config.vocab_size
-        self.window_size = config.sliding_window
-        self.sliding_window_pattern = config.sliding_window_pattern
-        self.num_hidden_layers = config.num_hidden_layers
-
-        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
-        self.embed_scale = config.hidden_size**0.5
-        self.layers = [
-            DecoderLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)
-        ]
-        self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-
-        # Per-layer input embeddings (2B/4B models)
-        self.hidden_size_per_layer_input = config.hidden_size_per_layer_input
-        if self.hidden_size_per_layer_input:
-            self.embed_tokens_per_layer = nn.Embedding(
-                config.vocab_size_per_layer_input,
-                config.num_hidden_layers * config.hidden_size_per_layer_input,
-            )
-            self.embed_tokens_per_layer_scale = config.hidden_size_per_layer_input**0.5
-            self.per_layer_input_scale = 2.0**-0.5
-            self.per_layer_projection_scale = config.hidden_size**-0.5
-            self.per_layer_model_projection = nn.Linear(
-                config.hidden_size,
-                config.num_hidden_layers * config.hidden_size_per_layer_input,
-                bias=False,
-            )
-            self.per_layer_projection_norm = nn.RMSNorm(
-                config.hidden_size_per_layer_input, eps=config.rms_norm_eps
-            )
-        else:
-            self.embed_tokens_per_layer = None
-            self.per_layer_input_scale = None
-            self.per_layer_projection_scale = None
-            self.per_layer_model_projection = None
-            self.per_layer_projection_norm = None
-
-        # Arrange for shared KVs
-        self.previous_kvs = list(range(len(self.layers)))
-        if config.num_kv_shared_layers > 0:
-            N = len(self.layers)
-            M = N - config.num_kv_shared_layers
-            kvs_by_type = {}
-            for i in range(M):
-                kvs_by_type[self.layers[i].layer_type] = i
-            for j in range(M, N):
-                self.previous_kvs[j] = kvs_by_type[self.layers[j].layer_type]
-
-    def _get_per_layer_inputs(
-        self,
-        input_ids: Optional[mx.array],
-        input_embeddings: Optional[mx.array] = None,
-    ) -> mx.array:
-        if input_ids is None:
-            if input_embeddings is None:
-                raise RuntimeError(
-                    "input_embeddings must be provided when input_ids are omitted."
-                )
-
-            # Split the sequence dimension if this still holds too much
-            # memory. 260k vocab means the distance tensor would be ~1GB
-            # per 2k tokens in bf16.
-            #
-            # If the embedding is quantized we have to dequantize it anyway to
-            # perform the match test.
-            norms_embedding = self.embed_tokens.weight.square().sum(-1)
-            norms_input = input_embeddings.square().sum(-1)
-            distance = _complete_square(
-                norms_embedding,
-                norms_input,
-                self.embed_tokens.as_linear(input_embeddings),
-            )
-
-            # Checks can be added if needed but they necessarily break the GPU
-            # pipelining and force an eval.
-            #
-            #   match_counts = (distance < eps).sum(-1)
-            #
-            input_ids = mx.argmin(distance, -1)
-
-        result = self.embed_tokens_per_layer(input_ids)
-        result = result * self.embed_tokens_per_layer_scale
-        return mx.unflatten(
-            result,
-            -1,
-            (self.config.num_hidden_layers, self.hidden_size_per_layer_input),
-        )
-
-    def _project_per_layer_inputs(
-        self,
-        input_embeddings: mx.array,
-        per_layer_inputs: Optional[mx.array] = None,
-    ) -> mx.array:
-        per_layer_projection = self.per_layer_model_projection(input_embeddings)
-        per_layer_projection = per_layer_projection * self.per_layer_projection_scale
-        per_layer_projection = mx.unflatten(
-            per_layer_projection,
-            -1,
-            (self.config.num_hidden_layers, self.hidden_size_per_layer_input),
-        )
-        per_layer_projection = self.per_layer_projection_norm(per_layer_projection)
-
-        if per_layer_inputs is None:
-            return per_layer_projection
-
-        return (per_layer_projection + per_layer_inputs) * self.per_layer_input_scale
-
-    def _make_masks(self, h, cache):
-        mask = {}
-        masks = []
-        for l, c in zip(self.layers, cache):
-            if l.layer_type not in mask:
-                if l.layer_type == "full_attention":
-                    mask["full_attention"] = create_attention_mask(h, c)
-                elif l.layer_type == "sliding_attention":
-                    mask["sliding_attention"] = create_attention_mask(
-                        h, c, window_size=self.window_size
-                    )
-            masks.append(mask[l.layer_type])
-        return masks
-
-    def __call__(
-        self,
-        inputs: mx.array = None,
-        cache=None,
-        input_embeddings: Optional[mx.array] = None,
-        per_layer_inputs: Optional[mx.array] = None,
-    ):
-        # Make the initial hidden state
-        if input_embeddings is None:
-            input_embeddings = self.embed_tokens(inputs)
-        h = input_embeddings
-        h = h * self.embed_scale
-
-        # Get the extra inputs per layer if we have per layer embeddings
-        if self.hidden_size_per_layer_input:
-            if per_layer_inputs is None:
-                per_layer_inputs = self._get_per_layer_inputs(inputs, input_embeddings)
-            per_layer_inputs = self._project_per_layer_inputs(h, per_layer_inputs)
-        if per_layer_inputs is not None:
-            per_layer_inputs = [
-                per_layer_inputs[:, :, i, :] for i, _ in enumerate(self.layers)
-            ]
-        else:
-            per_layer_inputs = [None] * len(self.layers)
-
-        # Make the kv cache list, be sure to append None for all the shared kv
-        # layers
-        if cache is None:
-            cache = [None] * len(self.layers)
-        else:
-            cache = cache + [None] * (len(self.layers) - len(cache))
-
-        # Apply each layer. We save all intermediate kvs and offset and grab
-        # the previous one for the shared kv layers.
-        masks = self._make_masks(h, cache)
-        intermediates = [(None, None)] * len(self.layers)
-        for idx, (layer, c, mask, prev_idx, per_layer_input) in enumerate(
-            zip(
-                self.layers,
-                cache,
-                masks,
-                self.previous_kvs,
-                per_layer_inputs,
-            )
-        ):
-            kvs, offset = intermediates[prev_idx]
-
-            h, kvs, offset = layer(
-                h,
-                mask,
-                c,
-                per_layer_input=per_layer_input,
-                shared_kv=kvs,
-                offset=offset,
-            )
-
-            intermediates[idx] = (kvs, offset)
-
-        return self.norm(h)
-
-
-class Model(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.args = args
-        self.model_type = args.model_type
-        self.model = Gemma4TextModel(args)
-        self.final_logit_softcapping = args.final_logit_softcapping
-        self.tie_word_embeddings = args.tie_word_embeddings
-        if not self.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,
-        per_layer_inputs: Optional[mx.array] = None,
-    ):
-        out = self.model(
-            inputs,
-            cache=cache,
-            input_embeddings=input_embeddings,
-            per_layer_inputs=per_layer_inputs,
-        )
-        if self.tie_word_embeddings:
-            out = self.model.embed_tokens.as_linear(out)
-        else:
-            out = self.lm_head(out)
-        if self.final_logit_softcapping is not None:
-            out = logit_softcap(self.final_logit_softcapping, out)
-        return out
-
-    def sanitize(self, weights):
-        sanitized = {}
-        first_kv_shared = self.args.num_hidden_layers - self.args.num_kv_shared_layers
-        for k, v in weights.items():
-            if any(
-                s in k
-                for s in (
-                    "self_attn.rotary_emb",
-                    "input_max",
-                    "input_min",
-                    "output_max",
-                    "output_min",
-                )
-            ):
-                continue
-
-            # KV-shared layers reuse K/V from earlier layers — drop their projections
-            if any(
-                s in k
-                for s in (".self_attn.k_proj", ".self_attn.v_proj", ".self_attn.k_norm")
-            ):
-                try:
-                    layer_idx = int(k.split("layers.")[1].split(".")[0])
-                    if layer_idx >= first_kv_shared:
-                        continue
-                except (IndexError, ValueError):
-                    pass
-
-            if k.endswith(".experts.gate_up_proj"):
-                base = k.removesuffix(".gate_up_proj")
-                gate, up = map(mx.contiguous, mx.split(v, 2, axis=-2))
-                sanitized[f"{base}.switch_glu.gate_proj.weight"] = gate
-                sanitized[f"{base}.switch_glu.up_proj.weight"] = up
-                continue
-
-            if k.endswith(".experts.down_proj"):
-                base = k.removesuffix(".down_proj")
-                sanitized[f"{base}.switch_glu.down_proj.weight"] = v
-                continue
-
-            sanitized[k] = v
-
-        return sanitized
-
-    @property
-    def quant_predicate(self):
-        def predicate(path, _):
-            if path.endswith("router.proj"):
-                return {"group_size": 64, "bits": 8}
-            return True
-
-        return predicate
-
-    @property
-    def layers(self):
-        return self.model.layers
-
-    @property
-    def head_dim(self):
-        return self.args.head_dim
-
-    @property
-    def n_kv_heads(self):
-        return self.args.num_key_value_heads
-
-    def make_cache(self):
-        first_kv_shared = self.args.num_hidden_layers - self.args.num_kv_shared_layers
-        caches = []
-        for i in range(first_kv_shared):
-            if self.args.layer_types[i] == "full_attention":
-                caches.append(KVCache())
-            else:
-                caches.append(
-                    RotatingKVCache(
-                        max_size=self.args.sliding_window,
-                        keep=0,
-                    )
-                )
-        return caches
@@ -267,7 +267,7 @@ class ShortConv1d(nn.Module):
            positions = (ends[:, None] + mx.arange(n_keep))[..., None]
            new_state = mx.take_along_axis(conv_input, positions, axis=1)
        else:
-            new_state = mx.contiguous(conv_input[:, -n_keep:, :])
+            new_state = conv_input[:, -n_keep:, :]

        return out, new_state

@@ -61,8 +61,8 @@ class ModelArgs(BaseModelArgs):
    _block_type_to_char = {"mamba": "M", "attention": "*", "moe": "E", "mlp": "-"}

    def __post_init__(self):
-        if self.time_step_limit is None:
-            self.time_step_limit = (0.0, float("inf"))
+        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:
@@ -157,13 +157,7 @@ class GatedDeltaNet(nn.Module):
            qkv = mx.where(mask[..., None], qkv, 0)
        conv_input = mx.concatenate([conv_state, qkv], axis=1)
        if cache is not None:
-            n_keep = self.conv_kernel_size - 1
-            if cache.lengths is not None:
-                ends = mx.clip(cache.lengths, 0, S)
-                positions = (ends[:, None] + mx.arange(n_keep))[..., None]
-                cache[0] = mx.take_along_axis(conv_input, positions, axis=1)
-            else:
-                cache[0] = mx.contiguous(conv_input[:, -n_keep:, :])
+            cache[0] = conv_input[:, -(self.conv_kernel_size - 1) :]
        conv_out = nn.silu(self.conv1d(conv_input))

        q, k, v = [
@@ -195,7 +189,6 @@ class GatedDeltaNet(nn.Module):

        if cache is not None:
            cache[1] = state
-            cache.advance(S)

        out = self.norm(out, z)
        out = self.out_proj(out.reshape(B, S, -1))
@@ -266,7 +266,7 @@ class Qwen3NextGatedDeltaNet(nn.Module):
                positions = (ends[:, None] + mx.arange(n_keep))[..., None]
                cache[0] = mx.take_along_axis(conv_input, positions, axis=1)
            else:
-                cache[0] = mx.contiguous(conv_input[:, -n_keep:, :])
+                cache[0] = conv_input[:, -n_keep:, :]

        conv_out = nn.silu(self.conv1d(conv_input))

@@ -58,10 +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 = x[...]
-        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,
@@ -183,44 +181,7 @@ class YarnRoPE(nn.Module):

    def __call__(self, x, offset=0):
        if self.mscale != 1.0:
-            x = x[...]
-            x[..., : self.dims] = self.mscale * x[..., : self.dims]
-        return mx.fast.rope(
-            x,
-            self.dims,
-            traditional=self.traditional,
-            base=None,
-            scale=1.0,
-            offset=offset,
-            freqs=self._freqs,
-        )
-
-
-class ProportionalRoPE(nn.Module):
-    def __init__(
-        self,
-        dims: int,
-        rotated_dims: int,
-        traditional: bool = False,
-        base: float = 10000.0,
-        factor: float = 1.0,
-    ):
-        super().__init__()
-        self.dims = dims
-        self.traditional = traditional
-
-        if rotated_dims > dims:
-            raise ValueError("rotated_dims should be smaller than dims")
-
-        exponents = mx.arange(0, rotated_dims, 2, dtype=mx.float32) / dims
-        self._freqs = mx.concatenate(
-            [
-                factor * (base**exponents),
-                mx.full(((dims - rotated_dims) // 2,), mx.inf),
-            ]
-        )
-
-    def __call__(self, x, offset=0):
+            x = x.at[..., : self.dims].multiply(self.mscale)
        return mx.fast.rope(
            x,
            self.dims,
@@ -290,14 +251,6 @@ def initialize_rope(
            short_factor=scaling_config["short_factor"],
            long_factor=scaling_config["long_factor"],
        )
-    elif rope_type == "proportional":
-        return ProportionalRoPE(
-            dims=dims,
-            rotated_dims=int(dims * scaling_config.get("partial_rotary_factor", 1.0)),
-            traditional=traditional,
-            base=base,
-            factor=scaling_config.get("factor", 1.0),
-        )
    elif rope_type == "mrope":
        mrope_section = scaling_config.get("mrope_section", [])
        assert (
@@ -314,11 +314,7 @@ def main():

    if args.target_dir is not None:
        target_dir = Path(args.target_dir)
-        has_targets = (
-            target_dir.is_dir()
-            and any((target_dir / "train").glob("*.safetensors"))
-            and any((target_dir / "valid").glob("*.safetensors"))
-        )
+        has_targets = target_dir.exists()
    else:
        has_targets = False
        target_dir = None
@@ -181,24 +181,39 @@ def apply_min_p(
        raise ValueError(
            f"`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}"
        )
+    # reference implementation: https://github.com/huggingface/transformers/blob/main/src/transformers/generation/logits_process.py#L531-L605

-    # Mask tokens that have a probability less than the max(p) * min_p
-    top_logprobs = mx.max(logprobs, axis=-1, keepdims=True)
+    # Indices sorted in decreasing order
+    sorted_indices = mx.argsort(-logprobs, axis=-1)
+    sorted_logprobs = mx.take_along_axis(logprobs, sorted_indices, axis=-1)
+
+    # Top probability
+    top_logprobs = sorted_logprobs[:, 0:1]
+
+    # Calculate the min_p threshold
    scaled_min_p = top_logprobs + math.log(min_p)
-    tokens_to_remove = logprobs < scaled_min_p

-    # Ensure at least min_tokens_to_keep survive the filter
-    if min_tokens_to_keep > 1:
-        top_indices = mx.argpartition(logprobs, kth=-min_tokens_to_keep, axis=-1)
-        top_indices = top_indices[..., -min_tokens_to_keep:]
-        tokens_to_remove = mx.put_along_axis(
-            tokens_to_remove,
-            top_indices,
-            False,
-            axis=-1,
-        )
+    # Mask tokens that have a probability less than the scaled min_p
+    tokens_to_remove = sorted_logprobs < scaled_min_p
+    tokens_to_remove[..., :min_tokens_to_keep] = False

-    return mx.where(tokens_to_remove, -float("inf"), logprobs)
+    # Create pool of tokens with probability less than scaled min_p
+    selected_logprobs = mx.where(tokens_to_remove, -float("inf"), sorted_logprobs)
+
+    # Create a mapping to rearrange back to original indices
+    inverse_indices = mx.put_along_axis(
+        mx.zeros_like(sorted_indices),
+        sorted_indices,
+        mx.arange(sorted_indices.shape[-1], dtype=sorted_indices.dtype),
+        axis=-1,
+    )
+
+    # Rearrange selected_logprobs back to original order
+    original_order_logprobs = mx.take_along_axis(
+        selected_logprobs, inverse_indices, axis=-1
+    )
+
+    return original_order_logprobs


@partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)
@@ -253,37 +253,6 @@ class BPEStreamingDetokenizer(StreamingDetokenizer):
        cls._byte_decoder = char_to_bytes


-def _infer_thinking(tokenizer):
-    vocab = tokenizer.get_vocab()
-    THINK_TOKENS = [
-        ("<think>", "</think>"),
-        ("<longcat_think>", "</longcat_think>"),
-    ]
-
-    # Single token thinking modes
-    for think_start, think_end in THINK_TOKENS:
-        if think_start in vocab and think_end in vocab:
-            return (
-                think_start,
-                think_end,
-                (vocab[think_start],),
-                (vocab[think_end],),
-            )
-
-    # Multi token thinking modes
-    if "<|channel>" in vocab and "<channel|>" in vocab:
-        think_start = "<|channel>thought"
-        think_end = "<channel|>"
-        return (
-            think_start,
-            think_end,
-            tuple(tokenizer.encode(think_start, add_special_tokens=False)),
-            tuple(tokenizer.encode(think_end, add_special_tokens=False)),
-        )
-
-    return (None, None, None, None)
-
-
 class TokenizerWrapper:
    """A wrapper that combines an HF tokenizer and a detokenizer.

@@ -308,12 +277,10 @@ class TokenizerWrapper:
            if eos_token_ids is not None
            else {tokenizer.eos_token_id}
        )
-        (
-            self._think_start,
-            self._think_end,
-            self._think_start_tokens,
-            self._think_end_tokens,
-        ) = _infer_thinking(tokenizer)
+        self._think_start = None
+        self._think_end = None
+        self._think_start_id = None
+        self._think_end_id = None

        self._chat_template = chat_template
        self.has_chat_template = (
@@ -322,20 +289,29 @@ class TokenizerWrapper:
        self._tool_parser = tool_parser
        self._tool_call_start = tool_call_start
        self._tool_call_end = tool_call_end
-        self._tool_call_start_tokens = None
-        self._tool_call_end_tokens = None
-        if tool_call_start is not None:
-            self._tool_call_start_tokens = tuple(
-                tokenizer.encode(tool_call_start, add_special_tokens=False)
-            )
-            self._tool_call_end_tokens = tuple(
-                tokenizer.encode(tool_call_end, add_special_tokens=False)
-            )
+
+        vocab = tokenizer.get_vocab()
+        THINK_TOKENS = [
+            ("<think>", "</think>"),
+            ("<longcat_think>", "</longcat_think>"),
+        ]
+        for think_start, think_end in THINK_TOKENS:
+            if think_start in vocab and think_end in vocab:
+                self._think_start = think_start
+                self._think_end = think_end
+                self._think_start_id = vocab[think_start]
+                self._think_end_id = vocab[think_end]
+                break
+
+        # Disable tool calling if tool call tokens aren't in vocab
+        if (tool_call_start and tool_call_start not in vocab) or (
+            tool_call_end and tool_call_end not in vocab
+        ):
+            self._tool_call_start = None
+            self._tool_call_end = None
+            self._tool_parser = None

    def apply_chat_template(self, *args, tokenize=True, **kwargs):
-        if "enable_thinking" not in kwargs:
-            kwargs["enable_thinking"] = self.has_thinking
-
        if self._chat_template is not None:
            out = self._chat_template(*args, **kwargs)
            if tokenize:
@@ -357,36 +333,6 @@ class TokenizerWrapper:

        self._eos_token_ids.add(token_id)

-    def _find(self, tokens, sequence, start=None, end=None, reverse=False):
-        start = start or 0
-        end = end or len(tokens)
-        outer_loop = (
-            range(end - len(sequence), start - 1, -1)
-            if reverse
-            else range(start, end - len(sequence) + 1)
-        )
-        for i in outer_loop:
-            if tokens[i] == sequence[0]:
-                if all(tokens[i + j] == sequence[j] for j in range(1, len(sequence))):
-                    return i
-        return -1
-
-    def find_think_start(self, tokens, start=None, end=None):
-        return self._find(tokens, self._think_start_tokens, start=start, end=end)
-
-    def rfind_think_start(self, tokens, start=None, end=None):
-        return self._find(
-            tokens, self._think_start_tokens, start=start, end=end, reverse=True
-        )
-
-    def find_think_end(self, tokens, start=None, end=None):
-        return self._find(tokens, self._think_end_tokens, start=start, end=end)
-
-    def rfind_think_end(self, tokens, start=None, end=None):
-        return self._find(
-            tokens, self._think_end_tokens, start=start, end=end, reverse=True
-        )
-
    @property
    def has_thinking(self):
        return self._think_start is not None
@@ -397,15 +343,7 @@ class TokenizerWrapper:

    @property
    def think_start_id(self):
-        if self._think_start_tokens is None:
-            return None
-        if len(self._think_start_tokens) > 1:
-            raise ValueError("The start thinking sequence is more than 1 token")
-        return self._think_start_tokens[0]
-
-    @property
-    def think_start_tokens(self):
-        return self._think_start_tokens
+        return self._think_start_id

    @property
    def think_end(self):
@@ -413,15 +351,7 @@ class TokenizerWrapper:

    @property
    def think_end_id(self):
-        if self._think_end_tokens is None:
-            return None
-        if len(self._think_end_tokens) > 1:
-            raise ValueError("The end thinking sequence is more than 1 token")
-        return self._think_end_tokens[0]
-
-    @property
-    def think_end_tokens(self):
-        return self._think_end_tokens
+        return self._think_end_id

    @property
    def has_tool_calling(self):
@@ -431,18 +361,10 @@ class TokenizerWrapper:
    def tool_call_start(self):
        return self._tool_call_start

-    @property
-    def tool_call_start_tokens(self):
-        return self._tool_call_start_tokens
-
    @property
    def tool_call_end(self):
        return self._tool_call_end

-    @property
-    def tool_call_end_tokens(self):
-        return self._tool_call_end_tokens
-
    @property
    def tool_parser(self):
        return self._tool_parser
@@ -551,8 +473,6 @@ def _infer_tool_parser(chat_template):
        return None
    elif "<minimax:tool_call>" in chat_template:
        return "minimax_m2"
-    elif "<|tool_call>" in chat_template and "<tool_call|>" in chat_template:
-        return "gemma4"
    elif "<start_function_call>" in chat_template:
        return "function_gemma"
    elif "<longcat_tool_call>" in chat_template:
@@ -1,65 +0,0 @@
-# Copyright © 2025 Apple Inc.
-
-import json
-from typing import Any, Optional
-
-import regex as re
-
-# Matches <|"|>...<|"|> string literals (Gemma 4's string delimiter).
-_GEMMA4_STR = r'<\|"\|>(?:(?!<\|"\|>)[\s\S])*?<\|"\|>'
-
-# Matches call:name{...} with balanced braces via the regex module's
-# recursive (?R)-style support.  The inner alternatives handle:
-#   [^{}<]          – any char that is not a brace or start of <|"|>
-#   <(?!\|"\|>)     – a lone '<' that is NOT the start of <|"|>
-#   <|"|>...<|"|>   – a complete string literal (braces inside are ignored)
-#   (?2)            – recursively balanced nested brace group
-_tool_call_regex = re.compile(
-    r"call:([\w-]+)(\{(?:[^{}<]|<(?!\|\"\|>)|" + _GEMMA4_STR + r"|(?2))*\})",
-    re.DOTALL,
-)
-
-
-def _gemma4_args_to_json(text: str) -> str:
-    """Convert Gemma 4 tool call args to valid JSON.
-
-    Gemma 4 uses unquoted keys and <|"|> as string delimiters
-    instead of standard double quotes.
-    """
-    strings = []
-
-    def _capture(m):
-        strings.append(m.group(1))
-        return f"\x00{len(strings) - 1}\x00"
-
-    # Extract <|"|>-delimited strings and replace with placeholders
-    text = re.sub(r'<\|"\|>(.*?)<\|"\|>', _capture, text, flags=re.DOTALL)
-    # Quote bare keys
-    text = re.sub(r"(?<=[{,])(\w+):", r'"\1":', text)
-    # Restore captured strings as properly escaped JSON strings
-    for i, s in enumerate(strings):
-        text = text.replace(f"\x00{i}\x00", json.dumps(s))
-
-    return text
-
-
-def _parse_single(match: re.Match) -> dict:
-    """Parse a single call:name{args} regex match into a tool call dict."""
-    func_name = match.group(1)
-    args_str = match.group(2)
-    json_str = _gemma4_args_to_json(args_str)
-    arguments = json.loads(json_str)
-    return dict(name=func_name, arguments=arguments)
-
-
-def parse_tool_call(text: str, _: Optional[Any] = None):
-    matches = list(_tool_call_regex.finditer(text))
-    if not matches:
-        raise ValueError("No function provided.")
-    if len(matches) == 1:
-        return _parse_single(matches[0])
-    return [_parse_single(m) for m in matches]
-
-
-tool_call_start = "<|tool_call>"
-tool_call_end = "<tool_call|>"
@@ -157,44 +157,43 @@ def _get_param_types_from_config(param_name: str, param_config: dict) -> list[st


 def parse_tool_call(text: str, tools: list | None = None):
-    invoke_matches = _invoke_complete_regex.findall(text)
-    if not invoke_matches:
+    invoke_match = _invoke_complete_regex.findall(text)
+    if not invoke_match:
        raise ValueError("No tool call found")
+    invoke_text = invoke_match[0]

-    param_config_for = {}
+    name_match = re.search(r"^([^>]+)", invoke_text)
+    if not name_match:
+        return None
+
+    function_name = _extract_name(name_match.group(1))
+
+    # Get parameter configuration
+    param_config = {}
    if tools:
        for tool in tools:
            if func := tool.get("function", False):
+                if func["name"] != function_name:
+                    continue
                if params := func.get("parameters", False):
-                    param_config_for[func["name"]] = params.get("properties", {})
+                    param_config = params.get("properties", {})

-    calls = []
-    for invoke_text in invoke_matches:
-        name_match = re.search(r"^([^>]+)", invoke_text)
-        if not name_match:
-            continue
-        function_name = _extract_name(name_match.group(1))
-        param_config = param_config_for.get(function_name, {})
+    # Extract parameters
+    param_dict = {}
+    for match in _parameter_complete_regex.findall(invoke_text):
+        param_match = re.search(r"^([^>]+)>(.*)", match, re.DOTALL)
+        if param_match:
+            param_name = _extract_name(param_match.group(1))
+            param_value = param_match.group(2).strip()
+            if param_value.startswith("\n"):
+                param_value = param_value[1:]
+            if param_value.endswith("\n"):
+                param_value = param_value[:-1]

-        param_dict = {}
-        for match in _parameter_complete_regex.findall(invoke_text):
-            param_match = re.search(r"^([^>]+)>(.*)", match, re.DOTALL)
-            if param_match:
-                param_name = _extract_name(param_match.group(1))
-                param_value = param_match.group(2).strip()
-                if param_value.startswith("\n"):
-                    param_value = param_value[1:]
-                if param_value.endswith("\n"):
-                    param_value = param_value[:-1]
+            param_type = _get_param_types_from_config(param_name, param_config)

-                param_type = _get_param_types_from_config(param_name, param_config)
+            param_dict[param_name] = _convert_param_value_with_types(
+                param_value, param_type
+            )

-                param_dict[param_name] = _convert_param_value_with_types(
-                    param_value, param_type
-                )
-
-        calls.append(dict(name=function_name, arguments=param_dict))
-
-    if len(calls) == 1:
-        return calls[0]
-    return calls
+    return dict(name=function_name, arguments=param_dict)
@@ -5,15 +5,9 @@ import types
 from pathlib import Path
 from typing import Any, Dict, List

-import mlx.core as mx
 from transformers import PreTrainedTokenizer


-def printf(*args, **kwargs):
-    if mx.distributed.init().rank() == 0:
-        print(*args, **kwargs)
-
-
 class TextDataset:
    """
    Light-weight wrapper to hold a dataset.
@@ -270,7 +264,7 @@ def load_custom_hf_dataset(args, tokenizer: PreTrainedTokenizer):
    collection = []
    for ds in dataset_collection:
        ds_path = ds["path"]
-        printf(f"Loading Hugging Face dataset {ds_path}.")
+        print(f"Loading Hugging Face dataset {ds_path}.")
        ds["mask_prompt"] = getattr(args, "mask_prompt", False)
        config = types.SimpleNamespace(**ds)
        hf_config = ds.get("config", {})
@@ -320,7 +314,7 @@ def load_dataset(args, tokenizer: PreTrainedTokenizer):
        if data_path.exists():
            train, valid, test = load_local_dataset(data_path, tokenizer, args)
        else:
-            printf(f"Loading Hugging Face dataset {args.data}.")
+            print(f"Loading Hugging Face dataset {args.data}.")
            train, valid, test = load_hf_dataset(args.data, tokenizer, args)

    if args.train and len(train) == 0:
@@ -328,7 +322,7 @@ def load_dataset(args, tokenizer: PreTrainedTokenizer):
            "Training set not found or empty. Must provide training set for fine-tuning."
        )
    if args.train and len(valid) == 0:
-        printf(
+        print(
            "Warning: Validation set not found or empty. Training will proceed without validation."
        )
    if args.test and len(test) == 0:
@@ -13,16 +13,10 @@ from mlx.nn.utils import average_gradients
 from mlx.utils import tree_flatten, tree_map
 from tqdm import tqdm

-from ..cli_ui import make_console, make_train_progress
 from .callbacks import TrainingCallback
 from .datasets import CacheDataset


-def printf(*args, **kwargs):
-    if mx.distributed.init().rank() == 0:
-        print(*args, **kwargs)
-
-
 def _clear_cache(threshold: int):
    if mx.get_cache_memory() > threshold:
        mx.clear_cache()
@@ -153,7 +147,7 @@ def iterate_batches(
                offsets = [0] * len(batch)
            lengths = [len(x) for x in batch]
            if max(lengths) > max_seq_length:
-                printf(
+                print(
                    f"[WARNING] Some sequences are longer than {max_seq_length} tokens. "
                    f"The longest sentence {max(lengths)} will be truncated to {max_seq_length}. "
                    "Consider pre-splitting your data to save memory."
@@ -188,8 +182,6 @@ def evaluate(
    loss: callable = default_loss,
    iterate_batches: callable = iterate_batches,
    clear_cache_threshold: int = 0,
-    progress: bool = True,
-    batch_callback: callable = None,
 ):
    model.eval()
    all_losses = mx.array(0.0)
@@ -197,30 +189,24 @@ def evaluate(

    index_iterator = iter(range(num_batches)) if num_batches != -1 else iter(int, 1)

-    batch_iter = zip(
-        index_iterator,
-        iterate_batches(
-            dataset=dataset,
-            batch_size=batch_size,
-            max_seq_length=max_seq_length,
-            comm_group=mx.distributed.init(),
+    for _, batch in tqdm(
+        zip(
+            index_iterator,
+            iterate_batches(
+                dataset=dataset,
+                batch_size=batch_size,
+                max_seq_length=max_seq_length,
+                comm_group=mx.distributed.init(),
+            ),
        ),
-    )
-    if progress:
-        batch_iter = tqdm(
-            batch_iter,
-            desc="Calculating loss...",
-            total=min(len(dataset) // batch_size, num_batches),
-        )
-
-    for _, batch in batch_iter:
+        desc="Calculating loss...",
+        total=min(len(dataset) // batch_size, num_batches),
+    ):
        losses, toks = loss(model, *batch)
        all_losses += losses * toks
        ntokens += toks
        mx.eval(all_losses, ntokens)
        _clear_cache(clear_cache_threshold)
-        if batch_callback is not None:
-            batch_callback()

    all_losses = mx.distributed.all_sum(all_losses, stream=mx.cpu)
    ntokens = mx.distributed.all_sum(ntokens, stream=mx.cpu)
@@ -241,13 +227,12 @@ def train(
 ):
    if mx.metal.is_available():
        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()
    rank = world.rank()
-
-    console = make_console()
-    if rank == 0 and world_size > 1:
-        console.print(f"[ui.muted]node {rank} of {world_size}[/ui.muted]")
+    if world_size > 1:
+        print(f"Node {rank} of {world_size}")

    if args.grad_checkpoint:
        grad_checkpoint(model.layers[0])
@@ -284,153 +269,119 @@ def train(
    train_time = 0
    grad_accum = None

-    progress = make_train_progress(console, disable=(rank != 0))
-    task = progress.add_task("train", total=args.iters)
-
-    if rank == 0:
-        console.print(
-            "  [ui.heading]iter   train_loss     tok/s     tokens[/ui.heading]"
-        )
-    progress.start()
-
-    prev_train_loss = None
-    try:
-        # Main training loop
-        for it, batch in zip(
-            range(1, args.iters + 1),
-            iterate_batches(
-                dataset=train_dataset,
-                batch_size=args.batch_size,
-                max_seq_length=args.max_seq_length,
-                loop=True,
-                comm_group=world,
-            ),
+    # Main training loop
+    for it, batch in zip(
+        range(1, args.iters + 1),
+        iterate_batches(
+            dataset=train_dataset,
+            batch_size=args.batch_size,
+            max_seq_length=args.max_seq_length,
+            loop=True,
+            comm_group=world,
+        ),
+    ):
+        tic = time.perf_counter()
+        # Report validation loss if needed, the first validation loss
+        # is always measured before any training.
+        if val_dataset and (
+            it == 1 or it % args.steps_per_eval == 0 or it == args.iters
        ):
            tic = time.perf_counter()
-            # Run validation periodically (skip iter 1).
-            if val_dataset and (it % args.steps_per_eval == 0 or it == args.iters):
-                if args.val_batches == -1:
-                    val_total = len(val_dataset) // args.batch_size
-                else:
-                    val_total = min(
-                        len(val_dataset) // args.batch_size, args.val_batches
-                    )
-                val_task = (
-                    progress.add_task("[bold blue]val  [/bold blue]", total=val_total)
-                    if rank == 0
-                    else None
-                )
-
-                def _advance_val():
-                    if val_task is not None:
-                        progress.advance(val_task)
-
-                tic = time.perf_counter()
-                val_loss = evaluate(
-                    model=model,
-                    dataset=val_dataset,
-                    loss=loss,
-                    batch_size=args.batch_size,
-                    num_batches=args.val_batches,
-                    max_seq_length=args.max_seq_length,
-                    iterate_batches=iterate_batches,
-                    progress=False,
-                    batch_callback=_advance_val,
-                )
-                model.train()
-                val_time = time.perf_counter() - tic
-                if val_task is not None:
-                    progress.remove_task(val_task)
-                if rank == 0:
-                    progress.console.print(
-                        f"  [ui.muted]{it:>4}[/ui.muted]    "
-                        f"[ui.accent]val[/ui.accent] "
-                        f"[ui.strong]{val_loss:>5.3f}[/ui.strong]    "
-                        f"[ui.muted]({val_time:.2f}s)[/ui.muted]"
-                    )
-
-                if training_callback is not None:
-                    val_info = {
-                        "iteration": it - 1,
-                        "val_loss": val_loss,
-                        "val_time": val_time,
-                    }
-                    training_callback.on_val_loss_report(val_info)
-
-                tic = time.perf_counter()
-
-            lvalue, toks, grad_accum = step(
-                batch,
-                grad_accum,
-                it % grad_accum_steps == 0,
+            val_loss = evaluate(
+                model=model,
+                dataset=val_dataset,
+                loss=loss,
+                batch_size=args.batch_size,
+                num_batches=args.val_batches,
+                max_seq_length=args.max_seq_length,
+                iterate_batches=iterate_batches,
            )
-
-            losses += lvalue
-            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
-
-            progress.advance(task)
-
-            # Report training loss if needed
-            if it % args.steps_per_report == 0 or it == args.iters:
-                train_loss = mx.distributed.all_sum(losses, stream=mx.cpu).item()
-                train_loss /= steps * world_size
-                n_tokens = mx.distributed.all_sum(n_tokens, stream=mx.cpu).item()
-                learning_rate = optimizer.learning_rate.item()
-                it_sec = args.steps_per_report / train_time
-                tokens_sec = float(n_tokens) / train_time
-                trained_tokens += n_tokens
-                peak_mem = mx.get_peak_memory() / 1e9
-                if rank == 0:
-                    if prev_train_loss is None or train_loss <= prev_train_loss:
-                        arrow, arrow_style = "▼", "green"
-                    else:
-                        arrow, arrow_style = "▲", "yellow"
-                    prev_train_loss = train_loss
-                    progress.console.print(
-                        f"  [ui.muted]{it:>4}[/ui.muted]    "
-                        f"[bold {arrow_style}]{train_loss:>5.3f} {arrow}"
-                        f"[/bold {arrow_style}]    "
-                        f"[ui.strong]{tokens_sec:>5,.0f}[/ui.strong]    "
-                        f"[ui.muted]{trained_tokens / 1000:>5.1f}k[/ui.muted]"
-                    )
-
-                if training_callback is not None:
-                    train_info = {
-                        "iteration": it,
-                        "train_loss": train_loss,
-                        "learning_rate": learning_rate,
-                        "iterations_per_second": it_sec,
-                        "tokens_per_second": tokens_sec,
-                        "trained_tokens": trained_tokens,
-                        "peak_memory": peak_mem,
-                    }
-                    training_callback.on_train_loss_report(train_info)
-
-                losses = 0
-                n_tokens = 0
-                steps = 0
-                train_time = 0
-
-            # Save adapter weights
-            if it % args.steps_per_save == 0 and rank == 0:
-                adapter_weights = dict(tree_flatten(model.trainable_parameters()))
-                mx.save_safetensors(str(args.adapter_file), adapter_weights)
-                checkpoint = (
-                    Path(args.adapter_file).parent / f"{it:07d}_adapters.safetensors"
+            model.train()
+            val_time = time.perf_counter() - tic
+            if rank == 0:
+                print(
+                    f"Iter {it}: "
+                    f"Val loss {val_loss:.3f}, "
+                    f"Val took {val_time:.3f}s",
+                    flush=True,
                )
-                mx.save_safetensors(str(checkpoint), adapter_weights)
-                progress.console.print(
-                    f"  [ui.good]save[/ui.good]  "
-                    f"[ui.muted]{checkpoint.name}[/ui.muted]"
+
+            if training_callback is not None:
+                val_info = {
+                    "iteration": it - 1,
+                    "val_loss": val_loss,
+                    "val_time": val_time,
+                }
+                training_callback.on_val_loss_report(val_info)
+
+            tic = time.perf_counter()
+
+        lvalue, toks, grad_accum = step(
+            batch,
+            grad_accum,
+            it % grad_accum_steps == 0,
+        )
+
+        losses += lvalue
+        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
+        if it % args.steps_per_report == 0 or it == args.iters:
+            train_loss = mx.distributed.all_sum(losses, stream=mx.cpu).item()
+            train_loss /= steps * world_size
+            n_tokens = mx.distributed.all_sum(n_tokens, stream=mx.cpu).item()
+            learning_rate = optimizer.learning_rate.item()
+            it_sec = args.steps_per_report / train_time
+            tokens_sec = float(n_tokens) / train_time
+            trained_tokens += n_tokens
+            peak_mem = mx.get_peak_memory() / 1e9
+            if rank == 0:
+                print(
+                    f"Iter {it}: Train loss {train_loss:.3f}, "
+                    f"Learning Rate {learning_rate:.3e}, "
+                    f"It/sec {it_sec:.3f}, "
+                    f"Tokens/sec {tokens_sec:.3f}, "
+                    f"Trained Tokens {trained_tokens}, "
+                    f"Peak mem {peak_mem:.3f} GB",
+                    flush=True,
                )
-    finally:
-        progress.stop()
+
+            if training_callback is not None:
+                train_info = {
+                    "iteration": it,
+                    "train_loss": train_loss,
+                    "learning_rate": learning_rate,
+                    "iterations_per_second": it_sec,
+                    "tokens_per_second": tokens_sec,
+                    "trained_tokens": trained_tokens,
+                    "peak_memory": peak_mem,
+                }
+                training_callback.on_train_loss_report(train_info)
+
+            losses = 0
+            n_tokens = 0
+            steps = 0
+            train_time = 0
+
+        # Save adapter weights
+        if it % args.steps_per_save == 0 and rank == 0:
+            adapter_weights = dict(tree_flatten(model.trainable_parameters()))
+            mx.save_safetensors(str(args.adapter_file), adapter_weights)
+            checkpoint = (
+                Path(args.adapter_file).parent / f"{it:07d}_adapters.safetensors"
+            )
+            mx.save_safetensors(str(checkpoint), adapter_weights)
+            print(
+                f"Iter {it}: Saved adapter weights to "
+                f"{args.adapter_file} and {checkpoint}."
+            )

    # Save final weights
    if rank == 0:
        adapter_weights = dict(tree_flatten(model.trainable_parameters()))
        mx.save_safetensors(str(args.adapter_file), adapter_weights)
+        print(f"Saved final weights to {args.adapter_file}.")
@@ -15,11 +15,6 @@ from .dora import DoRAEmbedding, DoRALinear
 from .lora import LoRAEmbedding, LoRALinear, LoRASwitchLinear


-def printf(*args, **kwargs):
-    if mx.distributed.init().rank() == 0:
-        print(*args, **kwargs)
-
-
 def build_schedule(schedule_config: Dict):
    """
    Build a learning rate schedule from the given config.
@@ -167,7 +162,7 @@ def print_trainable_parameters(model):
    trainable_p = (
        sum(v.size for _, v in tree_flatten(model.trainable_parameters())) / 1e6
    )
-    printf(
+    print(
        f"Trainable parameters: {(trainable_p * 100 / total_p):.3f}% "
        f"({trainable_p:.3f}M/{total_p:.3f}M)"
    )
@@ -342,7 +342,7 @@ def load_model(

    model = model_class(model_args)

-    if weights and hasattr(model, "sanitize"):
+    if hasattr(model, "sanitize"):
        weights = model.sanitize(weights)

    def _quantize(quantization):
@@ -507,8 +507,6 @@ def sharded_load(
    pipeline_group: Optional[mx.distributed.Group] = None,
    tensor_group: Optional[mx.distributed.Group] = None,
    return_config: bool = False,
-    *,
-    tokenizer_config: Optional[Dict[str, Any]] = None,
 ):
    # Get model path with everything but weight safetensors
    model_path = _download(
@@ -573,7 +571,7 @@ def sharded_load(
    # Load and shard the model, and load the weights
    tokenizer = load_tokenizer(
        model_path,
-        tokenizer_config or {"trust_remote_code": True},
+        {"trust_remote_code": True},
        eos_token_ids=config.get("eos_token_id", None),
    )
    model, _ = load_model(model_path, lazy=True, strict=False)
@@ -10,7 +10,7 @@ sys.path.append(str(package_dir))

 from _version import __version__

-MIN_MLX_VERSION = "0.31.2"
+MIN_MLX_VERSION = "0.30.4"

 setup(
    name="mlx-lm",
@@ -9,13 +9,12 @@ import mlx.core as mx
 from mlx_lm.generate import (
    BatchGenerator,
    GenerationResponse,
-    SequenceStateMachine,
    batch_generate,
    generate,
    generate_step,
    stream_generate,
 )
-from mlx_lm.models.cache import KVCache, RotatingKVCache
+from mlx_lm.models.cache import RotatingKVCache
 from mlx_lm.sample_utils import make_logits_processors, make_sampler
 from mlx_lm.utils import load

@@ -200,7 +199,7 @@ class TestGenerate(unittest.TestCase):
            self.model, stop_tokens=self.tokenizer.eos_token_ids, max_tokens=1
        )
        uids = gen.insert(prompts)
-        batch_responses = {r.uid: r for r in gen.next_generated()}
+        batch_responses = {r.uid: r for r in gen.next()}

        # Do a test for each prompt the logits are close
        for e, prompt in enumerate(prompts):
@@ -242,7 +241,7 @@ class TestGenerate(unittest.TestCase):
        batch_responses = {}
        not_in = True
        iters = 0
-        while responses := gen.next_generated():
+        while responses := gen.next():
            for r in responses:
                not_in &= r.uid not in batch_responses
                batch_responses[r.uid] = r
@@ -290,7 +289,7 @@ class TestGenerate(unittest.TestCase):
        num_toks = [2, 3, 4, 5]
        uids = gen.insert(prompts, max_tokens=num_toks)
        batch_responses = {uid: [] for uid in uids}
-        while responses := gen.next_generated():
+        while responses := gen.next():
            for r in responses:
                batch_responses[r.uid].append(r.token)

@@ -338,7 +337,7 @@ class TestGenerate(unittest.TestCase):
        )
        uids = batch_gen.insert(prompts)
        batch_responses = {uid: [] for uid in uids}
-        while responses := batch_gen.next_generated():
+        while responses := batch_gen.next():
            for r in responses:
                batch_responses[r.uid].append(r.logprobs)

@@ -371,7 +370,7 @@ class TestGenerate(unittest.TestCase):
        )
        prompt = self.tokenizer.encode("hello")
        uids = batch_gen.insert([prompt])
-        response = batch_gen.next_generated()[0]
+        response = batch_gen.next()[0]
        logprobs = response.logprobs
        self.assertEqual(logprobs[0].item(), 0.0)
        self.assertEqual(logprobs.argmin().item(), 1)
@@ -396,48 +395,12 @@ class TestGenerate(unittest.TestCase):
        processors = make_logits_processors(logit_bias)
        (uid2,) = batch_gen.insert([prompt], logits_processors=[processors])

-        responses = batch_gen.next_generated()
+        responses = batch_gen.next()
        responses = {response.uid: response for response in responses}
        self.assertEqual(responses[uid0].logprobs[0].item(), 0.0)
        self.assertEqual(responses[uid1].logprobs[1].item(), 0.0)
        self.assertEqual(responses[uid2].logprobs[2].item(), 0.0)

-    def test_batch_generate_processor_tokens_match_prompt_on_first_step(self):
-        prompt = self.tokenizer.encode("hello")
-        seen = []
-
-        def processor(tokens, logits):
-            seen.append(tokens)
-            return logits
-
-        batch_gen = BatchGenerator(
-            self.model,
-            max_tokens=1,
-            logits_processors=[processor],
-        )
-        batch_gen.insert([prompt])
-        batch_gen.next_generated()
-
-        self.assertTrue(hasattr(seen[0], "shape"))
-        self.assertEqual(seen[0].tolist(), prompt)
-
-    def test_batch_generate_function_with_logits_processors(self):
-        """Test that batch_generate function with logits_processors produces correct results."""
-        logit_bias = {0: 2000.0, 1: -2000.0}
-        processors = make_logits_processors(logit_bias)
-
-        prompts = [self.tokenizer.encode("hello")]
-        response = batch_generate(
-            self.model,
-            self.tokenizer,
-            prompts,
-            max_tokens=1,
-            logits_processors=processors,
-        )
-        self.assertEqual(len(response.texts), 1)
-        generated_token = self.tokenizer.encode(response.texts[0])[0]
-        self.assertEqual(generated_token, 0)
-
    def test_batch_generate_with_samplers(self):
        """Test that batch_generate with logits_processors produces correct results."""
        batch_gen = BatchGenerator(
@@ -447,7 +410,7 @@ class TestGenerate(unittest.TestCase):
        )
        prompt = self.tokenizer.encode("hello")
        uids = batch_gen.insert([prompt])
-        response = batch_gen.next_generated()[0]
+        response = batch_gen.next()[0]
        self.assertEqual(response.token, 1)

        del batch_gen
@@ -464,47 +427,12 @@ class TestGenerate(unittest.TestCase):
            samplers=[lambda _: mx.array([2]), lambda _: mx.array([3])],
        )

-        responses = batch_gen.next_generated()
+        responses = batch_gen.next()
        responses = {response.uid: response for response in responses}
        self.assertEqual(responses[uid0].token, 1)
        self.assertEqual(responses[uid1].token, 2)
        self.assertEqual(responses[uid2].token, 3)

-    def test_batch_generate_with_state_machines(self):
-        """Test that batch_generate with per-sequence state_machines stops on different tokens."""
-        batch_gen = BatchGenerator(
-            self.model,
-            max_tokens=10,
-        )
-        prompt = self.tokenizer.encode("hello")
-
-        sm_0 = SequenceStateMachine({"normal": [([0], None)]}, initial="normal")
-        sm_1 = SequenceStateMachine({"normal": [([1], None)]}, initial="normal")
-        sm_2 = SequenceStateMachine({"normal": [([2], None)]}, initial="normal")
-
-        processor_0 = make_logits_processors({0: 2000.0})
-        processor_1 = make_logits_processors({1: 2000.0})
-        processor_2 = make_logits_processors({2: 2000.0})
-
-        uid0, uid1, uid2 = batch_gen.insert(
-            [prompt, prompt, prompt],
-            logits_processors=[processor_0, processor_1, processor_2],
-            state_machines=[sm_0, sm_1, sm_2],
-        )
-
-        responses = batch_gen.next_generated()
-        responses = {response.uid: response for response in responses}
-
-        self.assertEqual(responses[uid0].token, 0)
-        self.assertEqual(responses[uid1].token, 1)
-        self.assertEqual(responses[uid2].token, 2)
-        self.assertEqual(responses[uid0].finish_reason, "stop")
-        self.assertEqual(responses[uid1].finish_reason, "stop")
-        self.assertEqual(responses[uid2].finish_reason, "stop")
-        self.assertEqual(responses[uid0].match_sequence, (0,))
-        self.assertEqual(responses[uid1].match_sequence, (1,))
-        self.assertEqual(responses[uid2].match_sequence, (2,))
-
    def test_batch_continued_generation(self):
        for rotating in [False, True]:
            if rotating:
@@ -553,7 +481,7 @@ class TestGenerate(unittest.TestCase):
            )
            uids = batch_gen.insert(prompts_a)
            caches = {uid: None for uid in uids}
-            while responses := batch_gen.next_generated():
+            while responses := batch_gen.next():
                for r in responses:
                    if r.finish_reason is not None:
                        caches[r.uid] = r.prompt_cache
@@ -562,7 +490,7 @@ class TestGenerate(unittest.TestCase):
            # Generate the 2nd time
            uids = batch_gen.insert(prompts_b, caches=caches)
            batch_responses = {uid: [] for uid in uids}
-            while responses := batch_gen.next_generated():
+            while responses := batch_gen.next():
                for r in responses:
                    batch_responses[r.uid].append(r.logprobs)

@@ -615,7 +543,7 @@ class TestGenerate(unittest.TestCase):

        uids = batch_gen.insert(prompts_a)
        caches = {uid: None for uid in uids}
-        while responses := batch_gen.next_generated():
+        while responses := batch_gen.next():
            for r in responses:
                if r.finish_reason is not None:
                    caches[r.uid] = r.prompt_cache
@@ -625,7 +553,7 @@ class TestGenerate(unittest.TestCase):
        # Generate the 2nd time
        uids = batch_gen.insert(prompts_b, caches=caches)
        batch_responses = {uid: [] for uid in uids}
-        while responses := batch_gen.next_generated():
+        while responses := batch_gen.next():
            for r in responses:
                batch_responses[r.uid].append(r.logprobs)

@@ -704,108 +632,6 @@ class TestGenerate(unittest.TestCase):
        model = qwen3_next.Model(args)
        self._continued_generation_test_helper(model)

-    def test_extend_cache_with_empty(self):
-        from mlx_lm.generate import _extend_cache
-        from mlx_lm.models.cache import make_prompt_cache
-
-        cache_a = make_prompt_cache(self.model)
-
-        prompt = mx.array([[1, 2, 3]])
-        self.model(prompt, cache=cache_a)
-        mx.eval([c.state for c in cache_a])
-
-        result = _extend_cache(cache_a, [])
-        self.assertEqual(len(result), len(cache_a))
-        for c in result:
-            self.assertGreater(c.offset, 0)
-
-        result = _extend_cache([], cache_a)
-        self.assertEqual(len(result), len(cache_a))
-        for c in result:
-            self.assertGreater(c.offset, 0)
-
-    def test_remove_prompt_batch_updates_currently_processing(self):
-        prompt_a = self.tokenizer.encode("Write a long story about a cat")
-        prompt_b = self.tokenizer.encode("Write a long story about a dog")
-
-        gen = BatchGenerator(
-            self.model,
-            max_tokens=5,
-            prefill_batch_size=2,
-            prefill_step_size=4,
-            completion_batch_size=4,
-        )
-        uid_a, uid_b = gen.insert([prompt_a, prompt_b])
-
-        gen.next()
-
-        found = gen._find_uids([uid_a, uid_b])
-        for uid in [uid_a, uid_b]:
-            self.assertIn(uid, found)
-            self.assertEqual(found[uid][0], 1)
-
-        gen.remove([uid_a])
-
-        self.assertEqual(len(gen._currently_processing), len(gen._prompt_batch))
-
-        found = gen._find_uids([uid_b])
-        self.assertIn(uid_b, found)
-
-        while responses := gen.next_generated():
-            if all(r.finish_reason is not None for r in responses):
-                break
-
-    def test_batch_max_kv_size_creates_rotating_cache(self):
-        max_kv_size = 256
-        gen = BatchGenerator(
-            self.model,
-            max_tokens=1,
-            max_kv_size=max_kv_size,
-        )
-
-        prompt = self.tokenizer.encode("Write a long story about a cat")
-        gen.insert([prompt])
-
-        for r in gen.next_generated():
-            if r.finish_reason is not None:
-                for cache in r.prompt_cache:
-                    self.assertIsInstance(cache, RotatingKVCache)
-                    self.assertEqual(cache.max_size, max_kv_size)
-
-    def test_batch_max_kv_size_limits_cache_growth(self):
-        max_kv_size = 5
-        gen = BatchGenerator(
-            self.model,
-            max_tokens=10,
-            max_kv_size=max_kv_size,
-            prefill_batch_size=1,
-            prefill_step_size=128,
-            completion_batch_size=1,
-        )
-
-        prompt = self.tokenizer.encode("Write a long story about a cat")
-        gen.insert([prompt])
-
-        for r in gen.next_generated():
-            if r.finish_reason is not None:
-                for cache in r.prompt_cache:
-                    self.assertLessEqual(cache.keys.shape[2], max_kv_size)
-
-    def test_batch_max_kv_size_none_creates_regular_cache(self):
-        gen = BatchGenerator(
-            self.model,
-            max_tokens=1,
-            max_kv_size=None,
-        )
-
-        prompt = self.tokenizer.encode("Write a long story about a cat")
-        gen.insert([prompt])
-
-        for r in gen.next_generated():
-            if r.finish_reason is not None:
-                for cache in r.prompt_cache:
-                    self.assertIsInstance(cache, KVCache)
-

 if __name__ == "__main__":
    unittest.main()
@@ -35,9 +35,6 @@ class TestConvertToGGUFWithoutMocks(unittest.TestCase):
        mock_tokenizer.get_vocab.return_value = {"<pad>": 0, "hello": 1, "world": 2}
        mock_tokenizer.all_special_tokens = ["<pad>"]
        mock_tokenizer.all_special_ids = [0]
-        mock_tokenizer.bos_token_id = None
-        mock_tokenizer.eos_token_id = None
-        mock_tokenizer.unk_token_id = None
        mock_from_pretrained.return_value = mock_tokenizer

        model_path = Path(self.test_dir)
@@ -5,7 +5,7 @@ import unittest

 import mlx.core as mx
 import mlx.nn as nn
-from mlx.utils import tree_flatten, tree_map
+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
@@ -242,43 +242,6 @@ class TestModels(unittest.TestCase):
        )
        self.assertTrue(isinstance(rope, rope_utils.Llama3RoPE))

-        rope = rope_utils.initialize_rope(
-            16,
-            base=100.0,
-            traditional=False,
-            scaling_config={
-                "rope_type": "proportional",
-                "partial_rotary_factor": 0.5,
-            },
-        )
-        self.assertTrue(isinstance(rope, rope_utils.ProportionalRoPE))
-        expected_freqs = 100.0 ** (mx.arange(0, 8, 2, dtype=mx.float32) / 16)
-        self.assertTrue(mx.allclose(rope._freqs[:4], expected_freqs))
-        self.assertTrue(mx.all(mx.isinf(rope._freqs[4:])))
-
-        x = mx.arange(16, dtype=mx.float32).reshape(1, 1, 1, 16)
-        y = rope(x, offset=1)
-        expected_rotated = mx.fast.rope(
-            mx.concatenate([x[..., :4], x[..., 8:12]], axis=-1),
-            8,
-            traditional=False,
-            base=None,
-            scale=1.0,
-            offset=1,
-            freqs=expected_freqs,
-        )
-        expected = mx.concatenate(
-            [
-                expected_rotated[..., :4],
-                x[..., 4:8],
-                expected_rotated[..., 4:],
-                x[..., 12:],
-            ],
-            axis=-1,
-        )
-        mx.eval(y, expected)
-        self.assertTrue(mx.allclose(y, expected))
-
    def test_su_scaled_rope_no_mutation(self):
        rope = rope_utils.SuScaledRoPE(
            dims=8,
@@ -649,199 +612,6 @@ class TestModels(unittest.TestCase):
                mx.array_equal(loaded[mlx_norm_key], converted[mlx_norm_key])
            )

-    def test_gemma4_convert_then_load_keeps_language_model_prefix(self):
-        from mlx_lm.models import gemma4
-
-        args = gemma4.ModelArgs.from_dict(
-            {
-                "model_type": "gemma4",
-                "vocab_size": 32,
-                "text_config": {
-                    "model_type": "gemma4_text",
-                    "hidden_size": 8,
-                    "num_hidden_layers": 1,
-                    "intermediate_size": 16,
-                    "num_attention_heads": 1,
-                    "num_key_value_heads": 1,
-                    "num_global_key_value_heads": 1,
-                    "head_dim": 8,
-                    "global_head_dim": 8,
-                    "sliding_window": 8,
-                    "sliding_window_pattern": 1,
-                    "layer_types": ["full_attention"],
-                    "hidden_size_per_layer_input": 0,
-                    "num_kv_shared_layers": 0,
-                    "tie_word_embeddings": True,
-                },
-            }
-        )
-        model = gemma4.Model(args)
-
-        base = mx.arange(8, dtype=mx.float32)
-        hf_norm_key = "model.language_model.layers.0.input_layernorm.weight"
-        mlx_norm_key = "language_model.model.layers.0.input_layernorm.weight"
-
-        converted = model.sanitize(
-            {
-                hf_norm_key: base,
-                "model.vision_tower.stub": mx.zeros((1,), dtype=mx.float32),
-            }
-        )
-        self.assertIn(mlx_norm_key, converted)
-        self.assertNotIn(
-            "language_model.model.model.layers.0.input_layernorm.weight", converted
-        )
-        self.assertTrue(mx.array_equal(converted[mlx_norm_key], base))
-        self.assertFalse(any("vision_tower" in k for k in converted))
-
-        loaded = model.sanitize({mlx_norm_key: base})
-        self.assertIn(mlx_norm_key, loaded)
-        self.assertNotIn(
-            "language_model.model.model.layers.0.input_layernorm.weight", loaded
-        )
-        self.assertTrue(mx.array_equal(loaded[mlx_norm_key], base))
-
-    def test_gemma4_raw_hf_language_model_prefixes_model(self):
-        from mlx_lm.models import gemma4
-
-        args = gemma4.ModelArgs.from_dict(
-            {
-                "model_type": "gemma4",
-                "vocab_size": 32,
-                "text_config": {
-                    "model_type": "gemma4_text",
-                    "hidden_size": 8,
-                    "num_hidden_layers": 1,
-                    "intermediate_size": 16,
-                    "num_attention_heads": 1,
-                    "num_key_value_heads": 1,
-                    "num_global_key_value_heads": 1,
-                    "head_dim": 8,
-                    "global_head_dim": 8,
-                    "sliding_window": 8,
-                    "sliding_window_pattern": 1,
-                    "layer_types": ["full_attention"],
-                    "hidden_size_per_layer_input": 0,
-                    "num_kv_shared_layers": 0,
-                    "tie_word_embeddings": True,
-                },
-            }
-        )
-        model = gemma4.Model(args)
-
-        base = mx.arange(8, dtype=mx.float32)
-        hf_norm_key = "model.language_model.layers.0.input_layernorm.weight"
-        mlx_norm_key = "language_model.model.layers.0.input_layernorm.weight"
-
-        converted = model.sanitize({hf_norm_key: base})
-        self.assertIn(mlx_norm_key, converted)
-        self.assertTrue(mx.array_equal(converted[mlx_norm_key], base))
-
-    def test_gemma4_raw_hf_moe_expert_weights_split_for_switch_glu(self):
-        from mlx_lm.models import gemma4
-
-        args = gemma4.ModelArgs.from_dict(
-            {
-                "model_type": "gemma4",
-                "vocab_size": 32,
-                "text_config": {
-                    "model_type": "gemma4_text",
-                    "hidden_size": 8,
-                    "num_hidden_layers": 1,
-                    "intermediate_size": 16,
-                    "num_attention_heads": 1,
-                    "num_key_value_heads": 1,
-                    "num_global_key_value_heads": 1,
-                    "head_dim": 8,
-                    "global_head_dim": 8,
-                    "sliding_window": 8,
-                    "sliding_window_pattern": 1,
-                    "layer_types": ["full_attention"],
-                    "hidden_size_per_layer_input": 0,
-                    "num_kv_shared_layers": 0,
-                    "tie_word_embeddings": True,
-                    "enable_moe_block": True,
-                    "num_experts": 2,
-                    "top_k_experts": 1,
-                    "moe_intermediate_size": 3,
-                },
-            }
-        )
-        model = gemma4.Model(args)
-
-        gate_up = mx.arange(2 * 6 * 8, dtype=mx.float32).reshape(2, 6, 8)
-        down = mx.arange(2 * 8 * 3, dtype=mx.float32).reshape(2, 8, 3)
-
-        converted = model.sanitize(
-            {
-                "model.language_model.layers.0.experts.gate_up_proj": gate_up,
-                "model.language_model.layers.0.experts.down_proj": down,
-            }
-        )
-
-        gate_key = "language_model.model.layers.0.experts.switch_glu.gate_proj.weight"
-        up_key = "language_model.model.layers.0.experts.switch_glu.up_proj.weight"
-        down_key = "language_model.model.layers.0.experts.switch_glu.down_proj.weight"
-
-        self.assertIn(gate_key, converted)
-        self.assertIn(up_key, converted)
-        self.assertIn(down_key, converted)
-        self.assertTrue(mx.array_equal(converted[gate_key], gate_up[:, :3, :]))
-        self.assertTrue(mx.array_equal(converted[up_key], gate_up[:, 3:, :]))
-        self.assertTrue(mx.array_equal(converted[down_key], down))
-        self.assertFalse(any("gate_up_proj" in k for k in converted))
-
-    def test_gemma4_moe_router_quantizes_to_8bit(self):
-        from mlx_lm.models import gemma4
-        from mlx_lm.models.switch_layers import QuantizedSwitchLinear
-        from mlx_lm.utils import quantize_model
-
-        args = gemma4.ModelArgs.from_dict(
-            {
-                "model_type": "gemma4",
-                "vocab_size": 64,
-                "text_config": {
-                    "model_type": "gemma4_text",
-                    "hidden_size": 64,
-                    "num_hidden_layers": 1,
-                    "intermediate_size": 128,
-                    "moe_intermediate_size": 128,
-                    "num_attention_heads": 1,
-                    "num_key_value_heads": 1,
-                    "num_global_key_value_heads": 1,
-                    "head_dim": 64,
-                    "global_head_dim": 64,
-                    "sliding_window": 8,
-                    "sliding_window_pattern": 1,
-                    "layer_types": ["full_attention"],
-                    "hidden_size_per_layer_input": 0,
-                    "num_kv_shared_layers": 0,
-                    "tie_word_embeddings": True,
-                    "enable_moe_block": True,
-                    "num_experts": 8,
-                    "top_k_experts": 2,
-                },
-            }
-        )
-        model = gemma4.Model(args)
-        model, config = quantize_model(
-            model,
-            {"model_type": "gemma4", "text_config": copy.deepcopy(args.text_config)},
-            group_size=64,
-            bits=4,
-        )
-
-        layer = model.language_model.model.layers[0]
-        self.assertIsInstance(layer.router.proj, nn.QuantizedLinear)
-        self.assertEqual(layer.router.proj.bits, 8)
-        self.assertIsInstance(layer.experts.switch_glu.gate_proj, QuantizedSwitchLinear)
-        self.assertEqual(layer.experts.switch_glu.gate_proj.bits, 4)
-        self.assertEqual(
-            config["quantization"]["language_model.model.layers.0.router.proj"]["bits"],
-            8,
-        )
-        self.assertEqual(config["quantization"]["bits"], 4)
-
    def test_qwen2_moe(self):
        from mlx_lm.models import qwen2_moe

@@ -1461,206 +1231,6 @@ class TestModels(unittest.TestCase):
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

-    def test_gemma4_text(self):
-        from mlx_lm.models import gemma4_text
-
-        args = gemma4_text.ModelArgs(
-            model_type="gemma4_text",
-            hidden_size=128,
-            num_hidden_layers=10,
-            intermediate_size=256,
-            num_attention_heads=4,
-            head_dim=32,
-            global_head_dim=64,
-            rms_norm_eps=1e-6,
-            vocab_size=1000,
-            vocab_size_per_layer_input=1000,
-            num_key_value_heads=1,
-            num_kv_shared_layers=4,
-            hidden_size_per_layer_input=32,
-            sliding_window=8,
-            sliding_window_pattern=5,
-            final_logit_softcapping=30.0,
-            layer_types=[
-                "sliding_attention",
-                "sliding_attention",
-                "sliding_attention",
-                "sliding_attention",
-                "full_attention",
-                "sliding_attention",
-                "sliding_attention",
-                "sliding_attention",
-                "sliding_attention",
-                "full_attention",
-            ],
-            rope_parameters={
-                "full_attention": {
-                    "partial_rotary_factor": 0.25,
-                    "rope_theta": 1000000.0,
-                },
-                "sliding_attention": {
-                    "rope_theta": 10000.0,
-                },
-            },
-        )
-        model = gemma4_text.Model(args)
-        self.model_test_runner(
-            model, args.model_type, args.vocab_size, args.num_hidden_layers
-        )
-
-    def test_gemma4_quantized_embedding_preserves_lookup_scale(self):
-        from mlx_lm.models import gemma4_text
-
-        args = gemma4_text.ModelArgs(
-            model_type="gemma4_text",
-            hidden_size=32,
-            num_hidden_layers=1,
-            intermediate_size=64,
-            num_attention_heads=2,
-            num_key_value_heads=1,
-            num_global_key_value_heads=1,
-            head_dim=16,
-            global_head_dim=16,
-            sliding_window=8,
-            sliding_window_pattern=1,
-            layer_types=["full_attention"],
-            hidden_size_per_layer_input=0,
-            vocab_size=4,
-            num_kv_shared_layers=0,
-        )
-        model = gemma4_text.Gemma4TextModel(args)
-        model.embed_tokens.weight = mx.ones((4, 32), dtype=mx.float32)
-        model.embed_tokens = model.embed_tokens.to_quantized(group_size=32, bits=8)
-
-        token_ids = mx.array([[0, 1]], dtype=mx.int32)
-        lookup = model.embed_tokens(token_ids) * model.embed_scale
-        logits = model.embed_tokens.as_linear(mx.ones((1, 1, 32), dtype=mx.float32))
-        mx.eval(lookup, logits)
-
-        self.assertTrue(
-            mx.allclose(
-                lookup,
-                mx.ones((1, 2, 32), dtype=mx.float32) * (32.0**0.5),
-            )
-        )
-        self.assertTrue(
-            mx.allclose(logits, mx.ones((1, 1, 4), dtype=mx.float32) * 32.0)
-        )
-
-    def test_gemma4_kv_shared_layers_omit_kv_projections(self):
-        """KV-shared layers must not create k_proj/v_proj/k_norm/v_norm so that
-        models saved without redundant weights (e.g. via transformers
-        save_pretrained) can be loaded with strict=True."""
-        from mlx_lm.models import gemma4_text
-
-        args = gemma4_text.ModelArgs(
-            model_type="gemma4_text",
-            hidden_size=128,
-            num_hidden_layers=10,
-            intermediate_size=256,
-            num_attention_heads=4,
-            head_dim=32,
-            global_head_dim=64,
-            rms_norm_eps=1e-6,
-            vocab_size=1000,
-            vocab_size_per_layer_input=1000,
-            num_key_value_heads=1,
-            num_kv_shared_layers=4,
-            hidden_size_per_layer_input=32,
-            sliding_window=8,
-            sliding_window_pattern=5,
-            final_logit_softcapping=30.0,
-            layer_types=[
-                "sliding_attention",
-                "sliding_attention",
-                "sliding_attention",
-                "sliding_attention",
-                "full_attention",
-                "sliding_attention",
-                "sliding_attention",
-                "sliding_attention",
-                "sliding_attention",
-                "full_attention",
-            ],
-            rope_parameters={
-                "full_attention": {
-                    "partial_rotary_factor": 0.25,
-                    "rope_theta": 1000000.0,
-                },
-                "sliding_attention": {
-                    "rope_theta": 10000.0,
-                },
-            },
-        )
-        model = gemma4_text.Model(args)
-
-        # Non-shared layers (0-5) should have KV projections
-        for i in range(6):
-            attn = model.model.layers[i].self_attn
-            self.assertTrue(attn.has_kv)
-            self.assertTrue(hasattr(attn, "k_proj"))
-            self.assertTrue(hasattr(attn, "k_norm"))
-
-        # Shared layers (6-9) should NOT have KV projections
-        for i in range(6, 10):
-            attn = model.model.layers[i].self_attn
-            self.assertFalse(attn.has_kv)
-            self.assertFalse(hasattr(attn, "k_proj"))
-            self.assertFalse(hasattr(attn, "k_norm"))
-            self.assertFalse(hasattr(attn, "v_proj"))
-
-        # Verify the model can load weights that omit shared-layer KV params
-        weights = dict(tree_flatten(model.parameters()))
-        kv_keys = [
-            k for k in weights if "k_proj" in k or "v_proj" in k or "k_norm" in k
-        ]
-        for k in kv_keys:
-            # All KV keys should belong to non-shared layers (0-5)
-            layer_idx = int(k.split("layers.")[1].split(".")[0])
-            self.assertLess(layer_idx, 6)
-
-    def test_gemma4_input_embeddings_reconstruct_per_layer_inputs(self):
-        from mlx_lm.models import gemma4_text
-
-        args = gemma4_text.ModelArgs(
-            model_type="gemma4_text",
-            hidden_size=32,
-            num_hidden_layers=2,
-            intermediate_size=64,
-            num_attention_heads=2,
-            num_key_value_heads=1,
-            num_global_key_value_heads=1,
-            head_dim=16,
-            global_head_dim=16,
-            sliding_window=8,
-            sliding_window_pattern=1,
-            layer_types=["full_attention", "full_attention"],
-            hidden_size_per_layer_input=8,
-            vocab_size=32,
-            vocab_size_per_layer_input=32,
-            num_kv_shared_layers=0,
-        )
-        model = gemma4_text.Model(args)
-        tokens = mx.array([[1, 2, 3]], dtype=mx.int32)
-        embeddings = model.model.embed_tokens(tokens)
-        per_layer_inputs = model.model._get_per_layer_inputs(tokens)
-
-        direct = model(tokens)
-        from_embeddings = model(None, input_embeddings=embeddings)
-        explicit = model(
-            None,
-            input_embeddings=embeddings,
-            per_layer_inputs=per_layer_inputs,
-        )
-        mx.eval(direct, from_embeddings, explicit)
-
-        self.assertTrue(
-            mx.allclose(direct.astype(mx.float32), from_embeddings.astype(mx.float32))
-        )
-        self.assertTrue(
-            mx.allclose(direct.astype(mx.float32), explicit.astype(mx.float32))
-        )
-
    def test_gpt_bigcode(self):
        from mlx_lm.models import gpt_bigcode

@@ -2094,50 +1664,6 @@ class TestModels(unittest.TestCase):
                "sliding_window": 8,
                "sliding_window_pattern": "LLGL",
            },
-            {
-                "model_type": "gemma4",
-                "num_hidden_layers": 10,
-                "vocab_size": 1000,
-                "text_config": {
-                    "model_type": "gemma4_text",
-                    "hidden_size": 128,
-                    "num_hidden_layers": 10,
-                    "intermediate_size": 128,
-                    "num_attention_heads": 4,
-                    "head_dim": 32,
-                    "global_head_dim": 64,
-                    "rms_norm_eps": 1e-6,
-                    "vocab_size": 1000,
-                    "vocab_size_per_layer_input": 1000,
-                    "num_key_value_heads": 1,
-                    "num_kv_shared_layers": 4,
-                    "hidden_size_per_layer_input": 32,
-                    "sliding_window": 8,
-                    "sliding_window_pattern": 5,
-                    "final_logit_softcapping": 30.0,
-                    "layer_types": [
-                        "sliding_attention",
-                        "sliding_attention",
-                        "sliding_attention",
-                        "sliding_attention",
-                        "full_attention",
-                        "sliding_attention",
-                        "sliding_attention",
-                        "sliding_attention",
-                        "sliding_attention",
-                        "full_attention",
-                    ],
-                    "rope_parameters": {
-                        "full_attention": {
-                            "partial_rotary_factor": 0.25,
-                            "rope_theta": 1000000.0,
-                        },
-                        "sliding_attention": {
-                            "rope_theta": 10000.0,
-                        },
-                    },
-                },
-            },
            {
                "model_type": "gemma3n",
                "num_hidden_layers": 4,
@@ -662,102 +662,6 @@ class TestPromptCache(unittest.TestCase):
        c_out = KVCache.merge((c1, c2))
        self.assertEqual(c_out.keys.shape, (2, 4, 4, 4))

-    def test_extend_with_empty_and_nonempty_batch_caches(self):
-        """Extending a batch cache when one side has keys=None should use the
-        correct batch size for the placeholder, not the batch size from the
-        non-None side. Regression test for broadcast error in dynamic_roll."""
-        H, D = 8, 64
-        max_size = 512
-
-        # -- BatchRotatingKVCache --
-        # Create 2 caches with content and 3 empty caches
-        c1 = RotatingKVCache(max_size=max_size)
-        c2 = RotatingKVCache(max_size=max_size)
-        c1.update_and_fetch(mx.ones((1, H, 5, D)), mx.ones((1, H, 5, D)))
-        c2.update_and_fetch(mx.ones((1, H, 3, D)), mx.ones((1, H, 3, D)))
-        batch_full = BatchRotatingKVCache.merge([c1, c2])
-
-        empty_caches = [RotatingKVCache(max_size=max_size) for _ in range(3)]
-        batch_empty = BatchRotatingKVCache.merge(empty_caches)
-
-        # Extend non-empty with empty (different batch sizes)
-        batch_full.extend(batch_empty)
-        self.assertEqual(batch_full.keys.shape[0], 5)
-        self.assertEqual(batch_full.offset.shape[0], 5)
-
-        # Prompt processing with right padding should not crash
-        batch_full.prepare(lengths=[10, 8, 12, 7, 11], right_padding=[2, 4, 0, 5, 1])
-        new_kv = mx.ones((5, H, 12, D))
-        batch_full.update_and_fetch(new_kv, new_kv)
-
-        # Also test empty extending non-empty
-        batch_full2 = BatchRotatingKVCache.merge(
-            [RotatingKVCache(max_size=max_size) for _ in range(3)]
-        )
-        c3 = RotatingKVCache(max_size=max_size)
-        c4 = RotatingKVCache(max_size=max_size)
-        c3.update_and_fetch(mx.ones((1, H, 4, D)), mx.ones((1, H, 4, D)))
-        c4.update_and_fetch(mx.ones((1, H, 6, D)), mx.ones((1, H, 6, D)))
-        batch_content = BatchRotatingKVCache.merge([c3, c4])
-        batch_full2.extend(batch_content)
-        self.assertEqual(batch_full2.keys.shape[0], 5)
-        self.assertEqual(batch_full2.offset.shape[0], 5)
-
-        # -- BatchKVCache --
-        c1 = KVCache()
-        c2 = KVCache()
-        c1.update_and_fetch(mx.ones((1, H, 5, D)), mx.ones((1, H, 5, D)))
-        c2.update_and_fetch(mx.ones((1, H, 3, D)), mx.ones((1, H, 3, D)))
-        batch_full = BatchKVCache.merge([c1, c2])
-
-        empty_caches = [KVCache() for _ in range(3)]
-        batch_empty = BatchKVCache.merge(empty_caches)
-
-        batch_full.extend(batch_empty)
-        self.assertEqual(batch_full.keys.shape[0], 5)
-        self.assertEqual(batch_full.offset.shape[0], 5)
-
-    def test_arrays_cache_extend_with_empty(self):
-        # test simple merge
-        c1 = ArraysCache(2)
-        c2 = ArraysCache(2)
-        c1[0] = mx.zeros((1, 4, 8))
-        c1[1] = mx.zeros((1, 4))
-        c2[0] = mx.zeros((1, 4, 8))
-        c2[1] = mx.zeros((1, 4))
-        full = ArraysCache.merge((c1, c2))
-        self.assertEqual(full[0].shape, (2, 4, 8))
-
-        # extend with empty
-        empty = ArraysCache.merge((ArraysCache(2),))
-        full.extend(empty)
-        self.assertEqual(full[0].shape, (3, 4, 8))
-        self.assertEqual(full[1].shape, (3, 4))
-        self.assertTrue(mx.all(full[0][2:] == 0))
-
-        # making an empty cache with 2 sequences and merging it with
-        # another one with 2 sequences
-        empty2 = ArraysCache.merge((ArraysCache(2), ArraysCache(2)))
-        content = ArraysCache.merge((c1, c2))
-        empty2.extend(content)
-        self.assertEqual(empty2[0].shape, (4, 4, 8))
-        self.assertEqual(empty2[1].shape, (4, 4))
-
-        # Extend content with empty
-        content = ArraysCache.merge((c1, c2))
-        empty2 = ArraysCache.merge((ArraysCache(2), ArraysCache(2)))
-        content.extend(empty2)
-        self.assertEqual(content[0].shape, (4, 4, 8))
-        self.assertEqual(content[1].shape, (4, 4))
-        self.assertEqual(content.make_mask(10).shape, (4, 10))
-
-        # multiple empty extensions accumulate correctly
-        stepwise = ArraysCache.merge((c1,))
-        stepwise.extend(ArraysCache(2))
-        stepwise.extend(ArraysCache.merge((ArraysCache(2), ArraysCache(2))))
-        self.assertEqual(stepwise[0].shape, (4, 4, 8))
-        self.assertEqual(stepwise[1].shape, (4, 4))
-
    def test_window_mask_with_full_kv_cache(self):
        c = KVCache()
        kv = mx.zeros((1, 1, 32, 128))
@@ -4,20 +4,13 @@ import http
 import io
 import json
 import threading
-import types
 import unittest

 import mlx.core as mx
 import requests

 from mlx_lm.models.cache import KVCache
-from mlx_lm.server import (
-    APIHandler,
-    LRUPromptCache,
-    Response,
-    ResponseGenerator,
-    _process_control_tokens,
-)
+from mlx_lm.server import APIHandler, LRUPromptCache, ResponseGenerator
 from mlx_lm.utils import load


@@ -68,9 +61,6 @@ class DummyModelProvider:
        assert model in ["default_model", "chat_model"]
        return self.model, self.tokenizer

-    def load_default(self):
-        return self.load("default_model", None, "default_model")
-

 class MockCache:
    def __init__(self, value, is_trimmable: bool = True):
@@ -92,71 +82,6 @@ class MockCache:
        return n


-class TestProcessControlTokens(unittest.TestCase):
-    @staticmethod
-    def _r(text, state, match=None):
-        return Response(text, 0, state, match, 0.0, None, ())
-
-    def test_single_tool_call_passes_body_with_open_and_close_crossings(self):
-        r = self._r
-        stream = [
-            r("hi ", "normal"),
-            r("<tool_call>", "tool", match=(0,)),
-            r("body", "tool"),
-            r("</tool_call>", "normal", match=(1,)),
-            r(" bye", "normal"),
-        ]
-        ctx = types.SimpleNamespace(
-            sequences={(0,): "<tool_call>", (1,): "</tool_call>"}
-        )
-        out = list(_process_control_tokens(ctx, iter(stream)))
-
-        self.assertEqual("".join(t.text for t in out), "hi body bye")
-        states = [t.state for t in out]
-        self.assertEqual(sum(1 for a, b in zip(states, states[1:]) if a != b), 2)
-
-    def test_back_to_back_tool_calls_emit_state_crossings(self):
-        r = self._r
-        stream = [
-            r("<tool_call>", "tool", match=(0,)),
-            r("call1_body", "tool"),
-            r("</tool_call>", "normal", match=(1,)),
-            r("<tool_call>", "tool", match=(0,)),
-            r("call2_body", "tool"),
-            r("</tool_call>", "normal", match=(1,)),
-        ]
-        ctx = types.SimpleNamespace(
-            sequences={(0,): "<tool_call>", (1,): "</tool_call>"}
-        )
-        out = list(_process_control_tokens(ctx, iter(stream)))
-
-        self.assertEqual("".join(t.text for t in out), "call1_bodycall2_body")
-        states = [t.state for t in out]
-        crossings = sum(
-            1 for a, b in zip(states, states[1:]) if a == "tool" and b == "normal"
-        )
-        self.assertEqual(crossings, 2)
-
-    def test_multi_token_match_preserves_order(self):
-        r = self._r
-        match = (10, 11, 12)
-        stream = [
-            r("body", "tool"),
-            r("</", "tool"),
-            r("tool", "tool"),
-            r("_call>", "normal", match=match),
-            r(" ok", "normal"),
-        ]
-        ctx = types.SimpleNamespace(sequences={match: "</tool_call>"})
-        out = list(_process_control_tokens(ctx, iter(stream)))
-
-        self.assertEqual([t.text for t in out], ["body", "", "", "", " ok"])
-        self.assertEqual(
-            [t.state for t in out],
-            ["tool", "tool", "tool", "normal", "normal"],
-        )
-
-
 class TestServer(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
@@ -280,33 +205,6 @@ class TestServer(unittest.TestCase):
        self.assertIn("id", response_body)
        self.assertIn("choices", response_body)

-    def test_make_state_machine_empty_tool_call_end(self):
-        class FakeTokenizer:
-            has_thinking = False
-            has_tool_calling = True
-            tool_call_start = "[TOOL_CALLS]"
-            tool_call_end = ""
-            tool_call_start_tokens = (100,)
-            tool_call_end_tokens = ()
-            eos_token_ids = [2]
-
-            def convert_ids_to_tokens(self, t):
-                return f"<eos{t}>"
-
-        sm, _ = self.response_generator._make_state_machine(
-            ("fake-empty-end", None, None),
-            FakeTokenizer(),
-            stop_words=[],
-        )
-        state = sm.make_state()
-        state, _, s = sm.match(state, 100)
-        self.assertEqual(s, "tool")
-        for tok in [42, 43, 44]:
-            state, _, s = sm.match(state, tok)
-            self.assertEqual(s, "tool")
-        state, _, s = sm.match(state, 2)
-        self.assertIsNone(s)
-
    def test_handle_models(self):
        url = f"http://localhost:{self.port}/v1/models"
        response = requests.get(url)
@@ -320,6 +218,18 @@ class TestServer(unittest.TestCase):
        self.assertEqual(model["object"], "model")
        self.assertIn("created", model)

+    def test_sequence_overlap(self):
+        from mlx_lm.server import sequence_overlap
+
+        self.assertTrue(sequence_overlap([1], [1]))
+        self.assertTrue(sequence_overlap([1, 2], [1, 2]))
+        self.assertTrue(sequence_overlap([1, 3], [3, 4]))
+        self.assertTrue(sequence_overlap([1, 2, 3], [2, 3]))
+
+        self.assertFalse(sequence_overlap([1], [2]))
+        self.assertFalse(sequence_overlap([1, 2], [3, 4]))
+        self.assertFalse(sequence_overlap([1, 2, 3], [4, 1, 2, 3]))
+

 class TestServerWithDraftModel(unittest.TestCase):
    @classmethod
@@ -604,7 +514,7 @@ class TestLRUPromptCache(unittest.TestCase):
        self.assertEqual(c, [MockCache("test3")])
        self.assertEqual(t, [])

-        cache.insert_cache(model, [4, 5], [MockCache("test4")], cache_type="user")
+        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])
@@ -627,41 +537,6 @@ class TestLRUPromptCache(unittest.TestCase):
        self.assertEqual(c, [MockCache("test4")])
        self.assertEqual(t, [])

-    def test_insert_trimmable_cache_removes_immediate_prefix(self):
-        cache = LRUPromptCache(max_size=10)
-        model = ("test", None, None)
-
-        cache.insert_cache(model, [1, 2], [MockCache("ab")])
-        self.assertEqual(len(cache), 1)
-        self.assertEqual(cache.nbytes, 2)
-
-        cache.insert_cache(model, [1, 2, 3], [MockCache("abc")])
-        self.assertEqual(len(cache), 1)
-        self.assertEqual(cache.nbytes, 3)
-
-    def test_insert_empty_tokens_does_not_self_destruct(self):
-        cache = LRUPromptCache(max_size=10)
-        model = ("test", None, None)
-
-        cache.insert_cache(model, [], [MockCache("root")])
-        self.assertEqual(len(cache), 1)
-        self.assertEqual(cache.nbytes, 4)
-
-        c, t = cache.fetch_nearest_cache(model, [])
-        self.assertIsNotNone(c)
-        self.assertEqual(t, [])
-
-    def test_fetch_empty_tokens_after_root_eviction(self):
-        cache = LRUPromptCache(max_size=10)
-        model = ("test", None, None)
-
-        cache.insert_cache(model, [], [MockCache("root")])
-        cache.insert_cache(model, [1], [MockCache("a")])
-
-        c, t = cache.fetch_nearest_cache(model, [])
-        self.assertIsNone(c)
-        self.assertEqual(t, [])
-
    def test_lru_bytes(self):
        cache = LRUPromptCache(max_size=100, max_bytes=10)
        model = ("test", None, None)
@@ -101,14 +101,6 @@ class TestTokenizers(unittest.TestCase):
        self.assertEqual(tokenizer.think_start, "<think>")
        self.assertEqual(tokenizer.think_end, "</think>")

-        tokenizer_repo = "mlx-community/Llama-3.2-1B-Instruct-4bit"
-        tokenizer = load_tokenizer(tokenizer_repo)
-        self.assertFalse(tokenizer.has_thinking)
-        self.assertIsNone(tokenizer.think_start)
-        self.assertIsNone(tokenizer.think_end)
-        self.assertIsNone(tokenizer.think_start_id)
-        self.assertIsNone(tokenizer.think_end_id)
-

 if __name__ == "__main__":
    unittest.main()
@@ -3,7 +3,6 @@ from pathlib import Path

 from mlx_lm.tool_parsers import (
    function_gemma,
-    gemma4,
    glm47,
    json_tools,
    kimi_k2,
@@ -19,7 +18,6 @@ class TestToolParsing(unittest.TestCase):
    def test_parsers(self):
        test_cases = [
            ("call:multiply{a:12234585,b:48838483920}", function_gemma),
-            ("call:multiply{a:12234585,b:48838483920}", gemma4),
            (
                '{"name": "multiply", "arguments": {"a": 12234585, "b": 48838483920}}',
                glm47,
@@ -91,10 +89,6 @@ class TestToolParsing(unittest.TestCase):
                "call:get_current_temperature{location:<escape>London<escape>}",
                function_gemma,
            ),
-            (
-                'call:get_current_temperature{location:<|"|>London<|"|>}',
-                gemma4,
-            ),
            (
                'get_current_temperature<arg_key>location</arg_key><arg_value>"London"</arg_value>',
                glm47,
@@ -197,84 +191,6 @@ class TestToolParsing(unittest.TestCase):
        self.assertEqual(tool_call["arguments"]["filters"], {"category": "books"})
        self.assertEqual(tool_call["arguments"]["tags"], ["fiction", "new"])

-    def test_gemma4(self):
-        # Nested object
-        test_case = 'call:configure{settings:{enabled:true,name:<|"|>test<|"|>}}'
-        tool_call = gemma4.parse_tool_call(test_case, None)
-        self.assertEqual(tool_call["name"], "configure")
-        self.assertEqual(
-            tool_call["arguments"],
-            {"settings": {"enabled": True, "name": "test"}},
-        )
-
-        # Array of strings
-        test_case = 'call:tag{items:[<|"|>foo<|"|>,<|"|>bar<|"|>]}'
-        tool_call = gemma4.parse_tool_call(test_case, None)
-        self.assertEqual(tool_call["name"], "tag")
-        self.assertEqual(tool_call["arguments"], {"items": ["foo", "bar"]})
-
-        # Mixed types
-        test_case = 'call:search{query:<|"|>hello world<|"|>,limit:10,verbose:false}'
-        tool_call = gemma4.parse_tool_call(test_case, None)
-        self.assertEqual(tool_call["name"], "search")
-        self.assertEqual(
-            tool_call["arguments"],
-            {"query": "hello world", "limit": 10, "verbose": False},
-        )
-
-        # Multiple tool calls in a single block (no delimiter between them)
-        test_case = (
-            'call:glob{pattern:<|"|>README*.md<|"|>}'
-            'call:glob{pattern:<|"|>CONTRIBUTING.md<|"|>}'
-        )
-        tool_calls = gemma4.parse_tool_call(test_case, None)
-        self.assertIsInstance(tool_calls, list)
-        self.assertEqual(len(tool_calls), 2)
-        self.assertEqual(tool_calls[0]["name"], "glob")
-        self.assertEqual(tool_calls[0]["arguments"], {"pattern": "README*.md"})
-        self.assertEqual(tool_calls[1]["name"], "glob")
-        self.assertEqual(tool_calls[1]["arguments"], {"pattern": "CONTRIBUTING.md"})
-
-        # Multiple tool calls with nested args
-        test_case = (
-            'call:search{query:<|"|>weather<|"|>,limit:5}'
-            'call:configure{settings:{enabled:true,name:<|"|>test<|"|>}}'
-        )
-        tool_calls = gemma4.parse_tool_call(test_case, None)
-        self.assertIsInstance(tool_calls, list)
-        self.assertEqual(len(tool_calls), 2)
-        self.assertEqual(tool_calls[0]["name"], "search")
-        self.assertEqual(
-            tool_calls[0]["arguments"],
-            {"query": "weather", "limit": 5},
-        )
-        self.assertEqual(tool_calls[1]["name"], "configure")
-        self.assertEqual(
-            tool_calls[1]["arguments"],
-            {"settings": {"enabled": True, "name": "test"}},
-        )
-
-        # Hyphenated function name (e.g. manim-video)
-        test_case = (
-            'call:manim-video{mode:<|"|>plan<|"|>,prompt:<|"|>explain KV caching<|"|>}'
-        )
-        tool_call = gemma4.parse_tool_call(test_case, None)
-        self.assertEqual(tool_call["name"], "manim-video")
-        self.assertEqual(
-            tool_call["arguments"],
-            {"mode": "plan", "prompt": "explain KV caching"},
-        )
-
-        # Braces inside a string argument (e.g. code snippets or markdown in content)
-        test_case = (
-            'call:skill_manage{action:<|"|>create<|"|>,'
-            'content:<|"|>use a dict like {key: value} in your code<|"|>}'
-        )
-        tool_call = gemma4.parse_tool_call(test_case, None)
-        self.assertEqual(tool_call["name"], "skill_manage")
-        self.assertEqual(tool_call["arguments"]["action"], "create")
-        self.assertIn("{", tool_call["arguments"]["content"])
-
    def test_kimi_k2(self):
        # Single tool call
        test_case = (
@@ -313,22 +229,6 @@ class TestToolParsing(unittest.TestCase):
        ]
        self.assertEqual(tool_calls, expected)

-    def test_minimax_m2(self):
-        test_case = (
-            '<invoke name="search">\n'
-            '<parameter name="query">weather</parameter>\n'
-            "</invoke>\n"
-            '<invoke name="read_file">\n'
-            '<parameter name="path">/tmp/test.txt</parameter>\n'
-            "</invoke>"
-        )
-        expected = [
-            {"name": "search", "arguments": {"query": "weather"}},
-            {"name": "read_file", "arguments": {"path": "/tmp/test.txt"}},
-        ]
-        tool_calls = minimax_m2.parse_tool_call(test_case, None)
-        self.assertEqual(expected, tool_calls)
-

 if __name__ == "__main__":
    unittest.main()
@@ -3,7 +3,6 @@
 import os
 import tempfile
 import unittest
-from pathlib import Path

 import mlx.core as mx
 import mlx.nn as nn
@@ -124,65 +123,6 @@ class TestUtils(unittest.TestCase):
        self.assertEqual(model.custom_attribute, "This is a custom model")
        self.assertTrue(hasattr(model, "qwenWeights"))

-    def test_load_model_gemma4_with_per_layer_projection_quantization(self):
-        from mlx_lm.models import gemma4
-
-        args = gemma4.ModelArgs.from_dict(
-            {
-                "model_type": "gemma4",
-                "vocab_size": 32,
-                "text_config": {
-                    "model_type": "gemma4_text",
-                    "hidden_size": 32,
-                    "num_hidden_layers": 2,
-                    "intermediate_size": 64,
-                    "num_attention_heads": 2,
-                    "num_key_value_heads": 1,
-                    "num_global_key_value_heads": 1,
-                    "head_dim": 16,
-                    "global_head_dim": 16,
-                    "sliding_window": 8,
-                    "sliding_window_pattern": 1,
-                    "layer_types": ["full_attention", "full_attention"],
-                    "hidden_size_per_layer_input": 32,
-                    "vocab_size_per_layer_input": 32,
-                    "num_kv_shared_layers": 0,
-                    "tie_word_embeddings": True,
-                },
-            }
-        )
-        model = gemma4.Model(args)
-        model, config = utils.quantize_model(
-            model,
-            {
-                "model_type": "gemma4",
-                "vocab_size": args.vocab_size,
-                "text_config": args.text_config,
-            },
-            group_size=32,
-            bits=4,
-        )
-
-        config["quantization"]["language_model.model.per_layer_model_projection"] = {
-            "group_size": 32,
-            "bits": 4,
-        }
-
-        with tempfile.TemporaryDirectory(dir=self.test_dir) as mlx_path:
-            utils.save_model(mlx_path, model)
-            utils.save_config(config, os.path.join(mlx_path, "config.json"))
-
-            loaded, loaded_config = utils.load_model(Path(mlx_path))
-
-            self.assertIn(
-                "language_model.model.per_layer_model_projection",
-                loaded_config["quantization"],
-            )
-
-            logits = loaded(mx.array([[1, 2, 3]], dtype=mx.int32))
-            mx.eval(logits)
-            self.assertEqual(logits.shape, (1, 3, args.vocab_size))
-

 if __name__ == "__main__":
    unittest.main()