Fix Gemma 4 sanitize() not stripping KV projections for shared layers (#1240)

Co-authored-by: Tristan Stahnke <tristan@melchior.lan>
Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>

mlx_lm/_version.py

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

mlx_lm/benchmark.py

@@ -148,10 +148,13 @@ 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):

mlx_lm/examples/batch_generate_response.py

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

mlx_lm/models/apertus.py

@@ -167,7 +167,8 @@ class Model(nn.Module):
         self.args = args
         self.model_type = args.model_type
         self.model = ApertusModel(args)
-        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
 
     def __call__(
         self,
@@ -175,12 +176,18 @@ class Model(nn.Module):
         cache: Optional[Any] = None,
     ) -> mx.array:
         out = self.model(inputs, cache)
-        return self.lm_head(out)
+        if self.args.tie_word_embeddings:
+            out = self.model.embed_tokens.as_linear(out)
+        else:
+            out = self.lm_head(out)
+        return out
 
     def sanitize(self, weights):
         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

mlx_lm/models/cache.py

@@ -1,11 +1,13 @@
 # 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_unflatten
+from mlx.utils import tree_flatten, tree_map, tree_reduce, tree_unflatten
 
 from .base import create_causal_mask
 
@@ -601,6 +603,18 @@ 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
 
@@ -619,13 +633,42 @@ class ArraysCache(_BaseCache):
         """
         In-place filter to keep just the given indices in the cache.
         """
-        self.cache = [c[batch_indices] for c in self.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]
 
     def extend(self, other):
         """
         In-place extend this cache with the other cache.
         """
-        self.cache = [mx.concatenate([c, o]) for c, o in zip(self.cache, 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)
 
     def extract(self, idx):
         cache = ArraysCache(len(self.cache))
@@ -660,6 +703,12 @@ 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)
@@ -968,16 +1017,18 @@ class BatchKVCache(_BaseCache):
         """
         In-place filter to keep just the given indices in the cache.
         """
-        self.keys = self.keys[batch_indices]
-        self.values = self.values[batch_indices]
+        if self.keys is not None:
+            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:
-            self.keys = self.keys[..., min_left_pad:, :]
-            self.values = self.values[..., min_left_pad:, :]
+            if self.keys is not None:
+                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
 
@@ -985,15 +1036,31 @@ 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)
-        max_size = max(self.keys.shape[2], other.keys.shape[2])
+        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)
 
         # Pad the keys and values so they are right-justified
         # with the index and the same size
         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)
+            left = max_idx - c._idx
+            right = max_size - k.shape[2] - left
             if right < 0:
                 k = k[..., :right, :]
                 v = v[..., :right, :]
@@ -1022,6 +1089,11 @@ 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)
@@ -1045,6 +1117,9 @@ class BatchKVCache(_BaseCache):
 
         return cache
 
+    def size(self):
+        return self._idx
+
     def empty(self):
         return self.keys is None
 
@@ -1285,8 +1360,9 @@ class BatchRotatingKVCache(_BaseCache):
         """
         In-place filter to keep just the given indices in the cache.
         """
-        self.keys = self.keys[batch_indices]
-        self.values = self.values[batch_indices]
+        if self.keys is not None:
+            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]
 
@@ -1294,17 +1370,33 @@ 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)
-        max_size = max(self.keys.shape[2], other.keys.shape[2])
+        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)
 
         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, :]
@@ -1323,9 +1415,10 @@ 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 = self.left_padding[idx].item()
-        offset = self.offset[idx].item()
+        padding = max(0, self.left_padding.tolist()[idx])
+        offset = self.offset.tolist()[idx]
         cache.keys = self.keys[idx : idx + 1]
         cache.values = self.values[idx : idx + 1]
         cache._idx = self._idx
@@ -1349,6 +1442,11 @@ 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)
@@ -1358,11 +1456,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, c) in enumerate(zip(padding, caches)):
+        for i, (p, l, c) in enumerate(zip(padding, lengths, caches)):
             if c.keys is None:
                 continue
-            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)
+            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:, :]
 
         cache = cls(caches[0].max_size, padding)
         cache.keys = keys
@@ -1373,6 +1471,9 @@ class BatchRotatingKVCache(_BaseCache):
 
         return cache
 
+    def size(self):
+        return min(self._offset, self.max_size)
+
     def empty(self):
         return self.keys is None
 
@@ -1381,3 +1482,282 @@ 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

mlx_lm/models/deepseek_v32.py

@@ -87,19 +87,15 @@ 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)
-        q_pe, q_nope = mx.split(q, [self.rope_head_dim], axis=-1)
-
-        offset = cache.offset if cache is not None else 0
-
-        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)
+
+        offset = cache.offset if cache is not None else 0
+
+        q = self.rope(q, offset=offset)
+        k = self.rope(k, offset=offset)
+
         if cache is not None:
             k, _ = cache.update_and_fetch(k, mx.zeros([b, 1, s, 0]))
         if k.shape[2] <= self.index_topk:
@@ -221,7 +217,8 @@ class DeepseekV32Attention(nn.Module):
                     mx.broadcast_to(idx, idx.shape[:-1] + (k_pe.shape[-1],)),
                     axis=2,
                 )
-                mask = None
+                if mask is not None:
+                    mask = mx.take_along_axis(mask, topk_indices, axis=-1)
             else:
                 shape = list(topk_indices.shape)
                 shape[-1] = kv_latent.shape[2]

mlx_lm/models/gated_delta.py

@@ -81,6 +81,8 @@ 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;

mlx_lm/models/gemma4.py

@@ -0,0 +1,92 @@
+# 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()

mlx_lm/models/gemma4_text.py

@@ -0,0 +1,688 @@
+# 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

mlx_lm/models/kimi_linear.py

@@ -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 = conv_input[:, -n_keep:, :]
+            new_state = mx.contiguous(conv_input[:, -n_keep:, :])
 
         return out, new_state
 

mlx_lm/models/nemotron_h.py

@@ -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 and self.time_step_min is not None:
-            self.time_step_limit = (self.time_step_min, float("inf"))
+        if self.time_step_limit is None:
+            self.time_step_limit = (0.0, 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:

mlx_lm/models/qwen3_5.py

@@ -157,7 +157,13 @@ 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:
-            cache[0] = conv_input[:, -(self.conv_kernel_size - 1) :]
+            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:, :])
         conv_out = nn.silu(self.conv1d(conv_input))
 
         q, k, v = [
@@ -189,6 +195,7 @@ 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))

mlx_lm/models/qwen3_next.py

@@ -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] = conv_input[:, -n_keep:, :]
+                cache[0] = mx.contiguous(conv_input[:, -n_keep:, :])
 
         conv_out = nn.silu(self.conv1d(conv_input))
 

mlx_lm/models/rope_utils.py

@@ -196,6 +196,42 @@ class YarnRoPE(nn.Module):
         )
 
 
+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):
+        return mx.fast.rope(
+            x,
+            self.dims,
+            traditional=self.traditional,
+            base=None,
+            scale=1.0,
+            offset=offset,
+            freqs=self._freqs,
+        )
+
+
 def initialize_rope(
     dims,
     base,
@@ -254,6 +290,14 @@ 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 (

mlx_lm/quant/dwq.py

@@ -314,7 +314,11 @@ def main():
 
     if args.target_dir is not None:
         target_dir = Path(args.target_dir)
-        has_targets = target_dir.exists()
+        has_targets = (
+            target_dir.is_dir()
+            and any((target_dir / "train").glob("*.safetensors"))
+            and any((target_dir / "valid").glob("*.safetensors"))
+        )
     else:
         has_targets = False
         target_dir = None

mlx_lm/sample_utils.py

@@ -181,39 +181,24 @@ 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
 
-    # 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
+    # Mask tokens that have a probability less than the max(p) * min_p
+    top_logprobs = mx.max(logprobs, axis=-1, keepdims=True)
     scaled_min_p = top_logprobs + math.log(min_p)
+    tokens_to_remove = logprobs < scaled_min_p
 
-    # 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
+    # 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,
+        )
 
-    # 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
+    return mx.where(tokens_to_remove, -float("inf"), logprobs)
 
 
 @partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)

mlx_lm/tokenizer_utils.py

@@ -253,6 +253,37 @@ 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.
 
@@ -277,10 +308,12 @@ class TokenizerWrapper:
             if eos_token_ids is not None
             else {tokenizer.eos_token_id}
         )
-        self._think_start = None
-        self._think_end = None
-        self._think_start_id = None
-        self._think_end_id = None
+        (
+            self._think_start,
+            self._think_end,
+            self._think_start_tokens,
+            self._think_end_tokens,
+        ) = _infer_thinking(tokenizer)
 
         self._chat_template = chat_template
         self.has_chat_template = (
@@ -289,29 +322,20 @@ class TokenizerWrapper:
         self._tool_parser = tool_parser
         self._tool_call_start = tool_call_start
         self._tool_call_end = tool_call_end
-
-        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
+        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)
+            )
 
     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:
@@ -333,6 +357,36 @@ 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
@@ -343,7 +397,15 @@ class TokenizerWrapper:
 
     @property
     def think_start_id(self):
-        return self._think_start_id
+        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
 
     @property
     def think_end(self):
@@ -351,7 +413,15 @@ class TokenizerWrapper:
 
     @property
     def think_end_id(self):
-        return self._think_end_id
+        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
 
     @property
     def has_tool_calling(self):
@@ -361,10 +431,18 @@ 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
@@ -473,6 +551,8 @@ 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:

mlx_lm/tool_parsers/gemma4.py

@@ -0,0 +1,65 @@
+# 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|>"

mlx_lm/tool_parsers/minimax_m2.py

@@ -157,43 +157,44 @@ 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_match = _invoke_complete_regex.findall(text)
-    if not invoke_match:
+    invoke_matches = _invoke_complete_regex.findall(text)
+    if not invoke_matches:
         raise ValueError("No tool call found")
-    invoke_text = invoke_match[0]
 
-    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 = {}
+    param_config_for = {}
     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 = params.get("properties", {})
+                    param_config_for[func["name"]] = params.get("properties", {})
 
-    # 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]
+    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, {})
 
-            param_type = _get_param_types_from_config(param_name, param_config)
+        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[param_name] = _convert_param_value_with_types(
-                param_value, param_type
-            )
+                param_type = _get_param_types_from_config(param_name, param_config)
 
-    return dict(name=function_name, arguments=param_dict)
+                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

mlx_lm/utils.py

@@ -507,6 +507,8 @@ 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(
@@ -571,7 +573,7 @@ def sharded_load(
     # Load and shard the model, and load the weights
     tokenizer = load_tokenizer(
         model_path,
-        {"trust_remote_code": True},
+        tokenizer_config or {"trust_remote_code": True},
         eos_token_ids=config.get("eos_token_id", None),
     )
     model, _ = load_model(model_path, lazy=True, strict=False)

setup.py

@@ -10,7 +10,7 @@ sys.path.append(str(package_dir))
 
 from _version import __version__
 
-MIN_MLX_VERSION = "0.30.4"
+MIN_MLX_VERSION = "0.31.2"
 
 setup(
     name="mlx-lm",

tests/test_generate.py

@@ -9,12 +9,13 @@ 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 RotatingKVCache
+from mlx_lm.models.cache import KVCache, RotatingKVCache
 from mlx_lm.sample_utils import make_logits_processors, make_sampler
 from mlx_lm.utils import load
 
@@ -199,7 +200,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()}
+        batch_responses = {r.uid: r for r in gen.next_generated()}
 
         # Do a test for each prompt the logits are close
         for e, prompt in enumerate(prompts):
@@ -241,7 +242,7 @@ class TestGenerate(unittest.TestCase):
         batch_responses = {}
         not_in = True
         iters = 0
-        while responses := gen.next():
+        while responses := gen.next_generated():
             for r in responses:
                 not_in &= r.uid not in batch_responses
                 batch_responses[r.uid] = r
@@ -289,7 +290,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():
+        while responses := gen.next_generated():
             for r in responses:
                 batch_responses[r.uid].append(r.token)
 
@@ -337,7 +338,7 @@ class TestGenerate(unittest.TestCase):
         )
         uids = batch_gen.insert(prompts)
         batch_responses = {uid: [] for uid in uids}
-        while responses := batch_gen.next():
+        while responses := batch_gen.next_generated():
             for r in responses:
                 batch_responses[r.uid].append(r.logprobs)
 
@@ -370,7 +371,7 @@ class TestGenerate(unittest.TestCase):
         )
         prompt = self.tokenizer.encode("hello")
         uids = batch_gen.insert([prompt])
-        response = batch_gen.next()[0]
+        response = batch_gen.next_generated()[0]
         logprobs = response.logprobs
         self.assertEqual(logprobs[0].item(), 0.0)
         self.assertEqual(logprobs.argmin().item(), 1)
@@ -395,12 +396,48 @@ class TestGenerate(unittest.TestCase):
         processors = make_logits_processors(logit_bias)
         (uid2,) = batch_gen.insert([prompt], logits_processors=[processors])
 
-        responses = batch_gen.next()
+        responses = batch_gen.next_generated()
         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(
@@ -410,7 +447,7 @@ class TestGenerate(unittest.TestCase):
         )
         prompt = self.tokenizer.encode("hello")
         uids = batch_gen.insert([prompt])
-        response = batch_gen.next()[0]
+        response = batch_gen.next_generated()[0]
         self.assertEqual(response.token, 1)
 
         del batch_gen
@@ -427,12 +464,47 @@ class TestGenerate(unittest.TestCase):
             samplers=[lambda _: mx.array([2]), lambda _: mx.array([3])],
         )
 
-        responses = batch_gen.next()
+        responses = batch_gen.next_generated()
         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:
@@ -481,7 +553,7 @@ class TestGenerate(unittest.TestCase):
             )
             uids = batch_gen.insert(prompts_a)
             caches = {uid: None for uid in uids}
-            while responses := batch_gen.next():
+            while responses := batch_gen.next_generated():
                 for r in responses:
                     if r.finish_reason is not None:
                         caches[r.uid] = r.prompt_cache
@@ -490,7 +562,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():
+            while responses := batch_gen.next_generated():
                 for r in responses:
                     batch_responses[r.uid].append(r.logprobs)
 
@@ -543,7 +615,7 @@ class TestGenerate(unittest.TestCase):
 
         uids = batch_gen.insert(prompts_a)
         caches = {uid: None for uid in uids}
-        while responses := batch_gen.next():
+        while responses := batch_gen.next_generated():
             for r in responses:
                 if r.finish_reason is not None:
                     caches[r.uid] = r.prompt_cache
@@ -553,7 +625,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():
+        while responses := batch_gen.next_generated():
             for r in responses:
                 batch_responses[r.uid].append(r.logprobs)
 
@@ -632,6 +704,108 @@ 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()

tests/test_gguf.py

@@ -35,6 +35,9 @@ 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)

tests/test_models.py

@@ -5,7 +5,7 @@ import unittest
 
 import mlx.core as mx
 import mlx.nn as nn
-from mlx.utils import tree_map
+from mlx.utils import tree_flatten, tree_map
 
 from mlx_lm.models import rope_utils
 from mlx_lm.models.base import create_causal_mask, scaled_dot_product_attention
@@ -242,6 +242,43 @@ 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,
@@ -612,6 +649,199 @@ 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
 
@@ -1231,6 +1461,206 @@ 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
 
@@ -1664,6 +2094,50 @@ 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,

tests/test_prompt_cache.py

@@ -662,6 +662,102 @@ 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))

tests/test_server.py

@@ -4,13 +4,20 @@ 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, ResponseGenerator
+from mlx_lm.server import (
+    APIHandler,
+    LRUPromptCache,
+    Response,
+    ResponseGenerator,
+    _process_control_tokens,
+)
 from mlx_lm.utils import load
 
 
@@ -61,6 +68,9 @@ 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):
@@ -82,6 +92,71 @@ 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):
@@ -205,6 +280,33 @@ 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)
@@ -218,18 +320,6 @@ 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
@@ -514,7 +604,7 @@ class TestLRUPromptCache(unittest.TestCase):
         self.assertEqual(c, [MockCache("test3")])
         self.assertEqual(t, [])
 
-        cache.insert_cache(model, [4, 5], [MockCache("test4")], checkpoint=True)
+        cache.insert_cache(model, [4, 5], [MockCache("test4")], cache_type="user")
         c, t = cache.fetch_nearest_cache(model, [2, 3])
         self.assertEqual(c, None)
         self.assertEqual(t, [2, 3])
@@ -537,6 +627,41 @@ 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)

tests/test_tokenizers.py

@@ -101,6 +101,14 @@ 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()

tests/test_tool_parsing.py

@@ -3,6 +3,7 @@ from pathlib import Path
 
 from mlx_lm.tool_parsers import (
     function_gemma,
+    gemma4,
     glm47,
     json_tools,
     kimi_k2,
@@ -18,6 +19,7 @@ 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,
@@ -89,6 +91,10 @@ 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,
@@ -191,6 +197,84 @@ 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 = (
@@ -229,6 +313,22 @@ 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()

tests/test_utils.py

@@ -3,6 +3,7 @@
 import os
 import tempfile
 import unittest
+from pathlib import Path
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -123,6 +124,65 @@ 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()