Use depends in pipeline parallel (#483)

* Add sliding-window support to LLaMA

* nits

* version

---------

Co-authored-by: dnakov <3777433+dnakov@users.noreply.github.com>
Co-authored-by: Awni Hannun <awni@apple.com>

mlx_lm/SERVER.md

@@ -30,8 +30,6 @@ To see a full list of options run:
 mlx_lm.server --help
 ```
 
-## Chat completions API
-
 You can make a request to the model by running:
 
 ```shell
@@ -130,23 +128,7 @@ curl localhost:8080/v1/chat/completions \
     - `completion_tokens`: The number of tokens generated.
     - `total_tokens`: The total number of tokens, i.e. the sum of the above two fields.
 
-## Responses API
-
-The responses API follows the [OpenAI responses API
-spec](https://platform.openai.com/docs/quickstart?api-mode=responses)
-
-To make a request, use the `/reponses` endpoint. For exapmle:
-
-```shell
-curl localhost:8080/responses \
-  -H "Content-Type: application/json" \
-  -d '{
-     "input": [{"role": "user", "content": "Say this is a test!"}],
-     "temperature": 0.7
-   }'
-```
-
-## Models API
+### List Models
 
 Use the `v1/models` endpoint to list available models:
 

mlx_lm/_version.py

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

mlx_lm/benchmark.py

@@ -96,7 +96,10 @@ def main():
             model, tokenizer, prompts, max_tokens=generation_tokens
         ).stats
 
-    _bench = batch_bench
+    if batch_size == 1:
+        _bench = single_bench
+    else:
+        _bench = batch_bench
 
     print("Running warmup..")
     _bench()

mlx_lm/examples/lora_config.yaml

@@ -1,5 +1,5 @@
 # The path to the local model directory or Hugging Face repo.
-model: "mlx-community/Llama-3.2-1B-Instruct"
+model: "mlx-community/Llama-3.2-1B-Instruct-bf16"
 
 # Whether or not to train (boolean)
 train: true

mlx_lm/examples/openai_responses.py

@@ -1,59 +0,0 @@
-# Copyright © 2025 Apple Inc.
-"""
-Examples using the OpenAI responses endpoint with mlx_lm.server.
-
-To run, first start the server:
-
->>> mlx_lm.server
-
-Then run this script.
-
-More documentation on the API spec here:
-https://platform.openai.com/docs/quickstart?api-mode=responses
-"""
-from openai import OpenAI
-
-model = "mlx-community/Qwen3-4B-Instruct-2507-4bit"
-
-### Basic response example
-
-client = OpenAI(base_url="http://localhost:8080/v1", api_key="not-needed")
-
-response = client.responses.create(
-    model=model, input="Write a one-sentence bedtime story about a unicorn."
-)
-print(response.output_text)
-
-### Input with roles
-
-response = client.responses.create(
-    model=model,
-    input=[
-        {
-            "role": "user",
-            "content": [
-                {
-                    "type": "input_text",
-                    "text": "Write a one-sentence bedtime story about a unicorn.",
-                },
-            ],
-        }
-    ],
-)
-print(response.output_text)
-
-### Streaming
-
-stream = client.responses.create(
-    model=model,
-    input=[
-        {
-            "role": "user",
-            "content": "Say 'double bubble bath' ten times fast.",
-        },
-    ],
-    stream=True,
-)
-
-for event in stream:
-    print(event)

mlx_lm/generate.py

@@ -26,8 +26,11 @@ from .models import cache
 from .models.cache import (
     ArraysCache,
     BatchKVCache,
+    BatchRotatingKVCache,
+    CacheList,
     KVCache,
     QuantizedKVCache,
+    RotatingKVCache,
     load_prompt_cache,
 )
 from .sample_utils import make_sampler
@@ -284,16 +287,11 @@ class GenerationResponse:
 
 
 def maybe_quantize_kv_cache(prompt_cache, quantized_kv_start, kv_group_size, kv_bits):
-    if (
-        kv_bits is not None
-        and not isinstance(prompt_cache[0], cache.QuantizedKVCache)
-        and prompt_cache[0].offset > quantized_kv_start
-    ):
-        for i in range(len(prompt_cache)):
-            if isinstance(prompt_cache[i], cache.KVCache):
-                prompt_cache[i] = prompt_cache[i].to_quantized(
-                    group_size=kv_group_size, bits=kv_bits
-                )
+    if kv_bits is None:
+        return
+    for e, c in enumerate(prompt_cache):
+        if hasattr(c, "to_quantized") and c.offset >= quantized_kv_start:
+            prompt_cache[e] = c.to_quantized(group_size=kv_group_size, bits=kv_bits)
 
 
 def generate_step(
@@ -862,18 +860,25 @@ def _make_cache(model, left_padding):
     Convert a list of regular caches into their corresponding
     batch-aware caches.
     """
+
+    def to_batch_cache(c):
+        if isinstance(c, KVCache):
+            return BatchKVCache(left_padding)
+        elif isinstance(c, ArraysCache):
+            c.left_padding = mx.array(left_padding)
+            return c
+        elif isinstance(c, RotatingKVCache):
+            if c.keep > 0:
+                raise ValueError("RotatingKVCache with keep tokens is not supported.")
+            return BatchRotatingKVCache(c.max_size, left_padding)
+        elif isinstance(c, CacheList):
+            return CacheList(*(to_batch_cache(sub_c) for sub_c in c.caches))
+        else:
+            raise ValueError(f"{type(c)} does not yet support batching")
+
     if hasattr(model, "make_cache"):
         cache = model.make_cache()
-        batch_cache = []
-        for c in cache:
-            if isinstance(c, KVCache):
-                batch_cache.append(BatchKVCache(left_padding))
-            elif isinstance(c, ArraysCache):
-                c.left_padding = mx.array(left_padding)
-                batch_cache.append(c)
-            else:
-                raise ValueError(f"{type(c)} does not yet support batching")
-        return batch_cache
+        return [to_batch_cache(c) for c in cache]
     else:
         return [BatchKVCache(left_padding) for _ in model.layers]
 

mlx_lm/models/cache.py

@@ -1,5 +1,6 @@
 # Copyright © 2023-2024 Apple Inc.
 
+import copy
 from typing import Any, Dict, List, Optional
 
 import mlx.core as mx
@@ -73,10 +74,10 @@ def load_prompt_cache(file_name, return_metadata=False):
     arrays = tree_unflatten(list(arrays.items()))
     cache_metadata = tree_unflatten(list(cache_metadata.items()))
     info, metadata, classes = cache_metadata
-    cache = [globals()[c]() for c in classes]
-    for c, state, meta_state in zip(cache, arrays, info):
-        c.state = state
-        c.meta_state = meta_state
+    cache = [
+        globals()[c].from_state(state, meta_state)
+        for c, state, meta_state in zip(classes, arrays, info)
+    ]
     if return_metadata:
         return cache, metadata
     return cache
@@ -141,6 +142,14 @@ class _BaseCache:
     def is_trimmable(self):
         return False
 
+    @classmethod
+    def from_state(cls, state, meta_state):
+        # Create an instance of cls without calling __init__
+        obj = cls.__new__(cls)
+        obj.state = state
+        obj.meta_state = meta_state
+        return obj
+
 
 class ConcatenateKVCache(_BaseCache):
     """ConcatenateKVCache the simplest KV cache implementation.
@@ -188,11 +197,12 @@ class ConcatenateKVCache(_BaseCache):
 
 
 class QuantizedKVCache(_BaseCache):
+    step = 256
+
     def __init__(self, group_size: int = 64, bits: int = 8):
         self.keys = None
         self.values = None
         self.offset = 0
-        self.step = 256
         self.group_size = group_size
         self.bits = bits
 
@@ -254,11 +264,11 @@ class QuantizedKVCache(_BaseCache):
 
     @property
     def meta_state(self):
-        return tuple(map(str, (self.step, self.offset, self.group_size, self.bits)))
+        return tuple(map(str, (self.offset, self.group_size, self.bits)))
 
     @meta_state.setter
     def meta_state(self, v):
-        self.step, self.offset, self.group_size, self.bits = map(int, v)
+        self.offset, self.group_size, self.bits = map(int, v)
 
     def is_trimmable(self):
         return True
@@ -273,11 +283,12 @@ class QuantizedKVCache(_BaseCache):
 
 
 class KVCache(_BaseCache):
+    step = 256
+
     def __init__(self):
         self.keys = None
         self.values = None
         self.offset = 0
-        self.step = 256
 
     def update_and_fetch(self, keys, values):
         prev = self.offset
@@ -341,14 +352,14 @@ class KVCache(_BaseCache):
 
 
 class RotatingKVCache(_BaseCache):
+    step = 256
 
-    def __init__(self, max_size=None, keep=0, step=256):
+    def __init__(self, max_size, keep=0):
         self.keep = keep
         self.keys = None
         self.values = None
         self.offset = 0
         self.max_size = max_size
-        self.step = step
         self._idx = 0
 
     def _trim(self, trim_size, v, append=None):
@@ -388,10 +399,11 @@ class RotatingKVCache(_BaseCache):
             # preserve context
             self.keys = self._temporal_order(self.keys)
             self.values = self._temporal_order(self.values)
+            self._idx = self.keys.shape[2]
 
-            # The largest size is self.max_size + S to ensure
+            # The largest size is self.max_size + S - 1 to ensure
             # every token gets at least self.max_size context
-            trim_size = self._idx - self.max_size
+            trim_size = self._idx - self.max_size + 1
             self.keys = self._trim(trim_size, self.keys, keys)
             self.values = self._trim(trim_size, self.values, values)
         self.offset += keys.shape[2]
@@ -459,13 +471,11 @@ class RotatingKVCache(_BaseCache):
 
     @property
     def meta_state(self):
-        return tuple(
-            map(str, (self.keep, self.max_size, self.step, self.offset, self._idx))
-        )
+        return tuple(map(str, (self.keep, self.max_size, self.offset, self._idx)))
 
     @meta_state.setter
     def meta_state(self, v):
-        self.keep, self.max_size, self.step, self.offset, self._idx = map(
+        self.keep, self.max_size, self.offset, self._idx = map(
             int,
             v,
         )
@@ -487,7 +497,7 @@ class RotatingKVCache(_BaseCache):
     ):
         if N > 1:
             window_size = window_size or self.max_size
-            offset = min(self.max_size, self.offset)
+            offset = min(self.max_size - 1, self.offset)
             if offset + N > window_size or return_array:
                 return create_causal_mask(N, offset, window_size=window_size)
             else:
@@ -500,16 +510,19 @@ class RotatingKVCache(_BaseCache):
                 idx = self._idx
                 if idx >= self.max_size:
                     idx = 0
-                mask_size = min(self.max_size, self.offset)
+                if self.offset < self.max_size:
+                    mask_size = self.offset + 1
+                else:
+                    mask_size = self.max_size
                 mask = mx.arange(mask_size) >= (mask_size - window_size)
                 mask = mx.roll(mask, shift=idx + 1)
-                return mask[:, None]
+                return mask
 
 
 class ArraysCache(_BaseCache):
     def __init__(self, size, left_padding: Optional[List[int]] = None):
         self.cache = [None] * size
-        self.left_padding = left_padding
+        self.left_padding = mx.array(left_padding) if left_padding else None
 
     def __setitem__(self, idx, value):
         self.cache[idx] = value
@@ -552,7 +565,7 @@ class MambaCache(ArraysCache):
 
 
 class ChunkedKVCache(KVCache):
-    def __init__(self, chunk_size=None):
+    def __init__(self, chunk_size):
         super().__init__()
         self.chunk_size = chunk_size
         self.start_position = 0
@@ -603,7 +616,7 @@ class ChunkedKVCache(KVCache):
         self.chunk_size, self.start_position = map(int, v)
 
 
-class CacheList(KVCache):
+class CacheList(_BaseCache):
     def __init__(self, *caches):
         self.caches = caches
 
@@ -631,8 +644,24 @@ class CacheList(KVCache):
             c.state = v[start : start + l]
             start += l
 
+    def filter(self, batch_indices):
+        """
+        In-place filter to keep just the given indices in the cache.
+        """
+        for c in self.caches:
+            c.filter(batch_indices)
+
+    def extend(self, other):
+        """
+        In-place extend this cache with the other cache.
+        """
+        for c in self.caches:
+            c.extend(other)
+
 
 class BatchKVCache(_BaseCache):
+    step = 256
+
     def __init__(self, left_padding: List[int]):
         """
         The BatchKV cache expects inputs to be left-padded.
@@ -657,7 +686,6 @@ class BatchKVCache(_BaseCache):
         self.left_padding = mx.array(left_padding)
         self.offset = mx.array([-l for l in left_padding])
         self._idx = 0
-        self.step = 256
 
     def update_and_fetch(self, keys, values):
         prev = self._idx
@@ -756,3 +784,224 @@ class BatchKVCache(_BaseCache):
             mx.concatenate, zip(*(pad(self), pad(other)))
         )
         self._idx = max_idx
+
+
+class BatchRotatingKVCache(_BaseCache):
+    step = 256
+
+    def __init__(self, max_size, left_padding: List[int]):
+        self.keys = None
+        self.values = None
+
+        self.left_padding = mx.array(left_padding)
+        self.offset = mx.array([-l for l in left_padding])
+
+        self.max_size = max_size
+        self._idx = 0
+        self._offset = 0
+        self.rotated = False
+
+    def _trim(self, trim_size, v, append=None):
+        if trim_size > 0:
+            v = v[..., trim_size:, :]
+        if append is not None:
+            return mx.concatenate([v, append], axis=2)
+        return v
+
+    def _temporal_order(self):
+        """
+        Rearrange the cache into temporal order.
+        """
+        if self.rotated:
+            self.keys = mx.roll(self.keys, -self._idx, axis=2)
+            self.values = mx.roll(self.values, -self._idx, axis=2)
+            self._idx = self.keys.shape[2]
+            self.rotated = False
+
+    def _update_concat(self, keys, values):
+        if self.keys is None:
+            self.keys = keys
+            self.values = values
+        else:
+            # Put the keys/values in temporal order to
+            # preserve context
+            self._temporal_order()
+
+            # Slice off the end if needed
+            if self.keys.shape[2] > self._idx:
+                self.keys = self.keys[..., : self._idx, :]
+                self.values = self.values[..., : self._idx, :]
+
+            # The largest size is self.max_size + S - 1 to ensure
+            # every token gets at least self.max_size context
+            trim_size = self._idx - self.max_size + 1
+            if trim_size > 0:
+                self.left_padding -= trim_size
+            self.keys = self._trim(trim_size, self.keys, keys)
+            self.values = self._trim(trim_size, self.values, values)
+        self.offset += keys.shape[2]
+        self._offset += keys.shape[2]
+        self._idx = self.keys.shape[2]
+        return self.keys, self.values
+
+    def _update_in_place(self, keys, values):
+        # May not have hit the max size yet, so potentially
+        # keep growing the cache
+        B, n_kv_heads, S, k_head_dim = keys.shape
+        prev = self._offset
+        if self.keys is None or (
+            prev >= self.keys.shape[2] and self.keys.shape[2] < self.max_size
+        ):
+            v_head_dim = values.shape[3]
+            new_size = min(self.step, self.max_size - prev)
+            k_shape = (B, n_kv_heads, new_size, k_head_dim)
+            v_shape = (B, n_kv_heads, new_size, v_head_dim)
+            new_k = mx.zeros(k_shape, keys.dtype)
+            new_v = mx.zeros(v_shape, values.dtype)
+            if self.keys is not None:
+                self.keys = mx.concatenate([self.keys, new_k], axis=2)
+                self.values = mx.concatenate([self.values, new_v], axis=2)
+            else:
+                self.keys, self.values = new_k, new_v
+            self._idx = prev
+
+        # Trim if needed
+        trim_size = self.keys.shape[2] - self.max_size
+        if trim_size > 0:
+            self.keys = self._trim(trim_size, self.keys)
+            self.values = self._trim(trim_size, self.values)
+            self._idx = self.max_size
+            self.left_padding -= trim_size
+
+        # Rotate
+        if self._idx == self.max_size:
+            self.rotated = True
+            self._idx = 0
+        if self.rotated:
+            self.left_padding -= S
+
+        # Assign
+        self.keys[..., self._idx : self._idx + S, :] = keys
+        self.values[..., self._idx : self._idx + S, :] = values
+        self._offset += S
+        self.offset += S
+        self._idx += S
+
+        # If the buffer is not full, slice off the end
+        if self._offset < self.max_size:
+            return (
+                self.keys[..., : self._offset, :],
+                self.values[..., : self._offset, :],
+            )
+        return self.keys, self.values
+
+    def update_and_fetch(self, keys, values):
+        if keys.shape[2] == 1:
+            return self._update_in_place(keys, values)
+        return self._update_concat(keys, values)
+
+    @property
+    def state(self):
+        k, v = self.keys, self.values
+        if self._offset < k.shape[2]:
+            k, v = k[..., : self._offset, :], v[..., : self._offset, :]
+        return k, v, self.offset, self.left_padding
+
+    @state.setter
+    def state(self, v):
+        self.keys, self.values, self.offset, self.left_padding = v
+
+    @property
+    def meta_state(self):
+        return tuple(map(str, (self.max_size, self._offset, self._idx, self.rotated)))
+
+    @meta_state.setter
+    def meta_state(self, v):
+        self.max_size, self._offset, self._idx = map(
+            int,
+            v[:3],
+        )
+        self.rotated = bool(v[3])
+
+    def is_trimmable(self):
+        return self._offset < self.max_size
+
+    def trim(self, n):
+        n = min(self._offset, n)
+        self._offset -= n
+        self._idx -= n
+        self.offset -= n
+        return n
+
+    def to_quantized(self, group_size: int = 64, bits: int = 4) -> QuantizedKVCache:
+        raise NotImplementedError("BatchRotatingKVCache Quantization NYI")
+
+    def make_mask(
+        self, N: int, window_size: Optional[int] = None, return_array: bool = False
+    ):
+        left_padding = self.left_padding
+        window_size = window_size or self.max_size
+        offset = min(self.max_size - 1, self._offset)
+        rinds = mx.arange(offset + N)
+        linds = mx.arange(offset, offset + N) if offset else rinds
+        linds = linds[:, None]
+        rinds = rinds[None]
+        mask = linds >= rinds
+        mask &= linds < rinds + window_size
+        if (trim_size := self._idx - self.max_size + int(N > 1)) > 0:
+            left_padding = left_padding - trim_size
+
+        rotated = N == 1 and (self.rotated or self._idx >= self.max_size)
+        if rotated:
+            left_padding = left_padding - 1
+
+        mask = mask & (rinds >= mx.expand_dims(left_padding, (1, 2, 3)))
+
+        if rotated:
+            idx = self._idx
+            if idx >= self.max_size:
+                idx = 0
+            mask = mx.roll(mask, shift=idx + 1, axis=-1)
+
+        return mask
+
+    def filter(self, batch_indices):
+        """
+        In-place filter to keep just the given indices in the cache.
+        """
+        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]
+
+    def extend(self, other):
+        """
+        In-place extend this cache with the other cache.
+        """
+        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])
+
+        def pad(c):
+            left = max_idx - c._idx
+            right = max_size - c.keys.shape[2] - left
+            k, v = c.keys, c.values
+            if right < 0:
+                k = k[..., :right, :]
+                v = v[..., :right, :]
+                right = 0
+            if left != 0 or right != 0:
+                pad = [(0, 0), (0, 0), (left, right), (0, 0)]
+                k = mx.pad(k, pad)
+                v = mx.pad(v, pad)
+            left_padding = c.left_padding + left
+            return k, v, c.offset, left_padding
+
+        self.keys, self.values, self.offset, self.left_padding = map(
+            mx.concatenate, zip(*(pad(self), pad(other)))
+        )
+        self._idx = max_idx
+        self._offset = max(self._offset, other._offset)

mlx_lm/models/deepseek_v2.py

@@ -414,6 +414,8 @@ class DeepseekV2Model(nn.Module):
         # Send to the next process in the pipeline
         if pipeline_rank != 0:
             h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
+            if cache[-1] is not None:
+                cache[-1].keys = mx.depends(cache[-1].keys, h)
 
         # Broadcast h while keeping it in the graph
         h = mx.distributed.all_gather(h)[: h.shape[0]]

mlx_lm/models/deepseek_v3.py

@@ -446,6 +446,8 @@ class DeepseekV3Model(nn.Module):
         # Send to the next process in the pipeline
         if pipeline_rank != 0:
             h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
+            if cache[-1] is not None:
+                cache[-1].keys = mx.depends(cache[-1].keys, h)
 
         # Broadcast h while keeping it in the graph
         h = mx.distributed.all_gather(h)[: h.shape[0]]

mlx_lm/models/falcon_h1.py

@@ -0,0 +1,479 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import (
+    BaseModelArgs,
+    create_attention_mask,
+    create_ssm_mask,
+    scaled_dot_product_attention,
+)
+from .cache import CacheList, KVCache, MambaCache
+from .rope_utils import initialize_rope
+from .ssm import ssm_update
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    attention_bias: bool = False
+    attention_in_multiplier: float = 1.0
+    attention_out_multiplier: float = 0.9375
+    embedding_multiplier: float = 5.656854249492381
+    head_dim: int = 64
+    hidden_size: int = 1024
+    initializer_range: float = 0.02
+    intermediate_size: int = 2048
+    key_multiplier: float = 0.390625
+    lm_head_multiplier: float = 0.0390625
+    mamba_chunk_size: int = 128
+    mamba_conv_bias: bool = True
+    mamba_d_conv: int = 4
+    mamba_d_head: int = 64
+    mamba_d_ssm: int = 1536
+    mamba_d_state: int = 128
+    mamba_expand: int = 2
+    mamba_n_groups: int = 1
+    mamba_n_heads: int = 24
+    mamba_norm_before_gate: bool = False
+    mamba_proj_bias: bool = False
+    mamba_rms_norm: bool = False
+    mamba_use_mlp: bool = True
+    max_position_embeddings: int = 131072
+    mlp_bias: bool = False
+    mlp_expansion_factor: int = 8
+    mlp_multipliers: List[float] = field(
+        default_factory=lambda: [0.8838834764831844, 0.5859375]
+    )
+    model_type: str = "falcon_h1"
+    num_attention_heads: int = 8
+    num_hidden_layers: int = 36
+    num_key_value_heads: int = 2
+    projectors_bias: bool = False
+    rms_norm_eps: float = 1e-05
+    rope_traditional: bool = False
+    rope_scaling: Optional[float] = None
+    rope_theta: float = 100000000000.0
+    ssm_in_multiplier: float = 1.25
+    ssm_multipliers: List[float] = field(
+        default_factory=lambda: [
+            0.3535533905932738,
+            0.25,
+            0.3535533905932738,
+            0.5,
+            0.3535533905932738,
+        ]
+    )
+    ssm_out_multiplier: float = 0.23570226039551587
+    vocab_size: int = 32784
+
+
+class FalconH1RMSNormGated(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6, n_groups=1, norm_before_gate=True):
+        super().__init__()
+        self.weight = mx.ones((hidden_size,))
+        self.variance_epsilon = eps
+        self.n_groups = n_groups
+        self.norm_before_gate = norm_before_gate
+
+    def __call__(self, hidden_states, gate=None):
+        if not self.norm_before_gate and gate is not None:
+            hidden_states = hidden_states * nn.silu(gate)
+
+        hidden_states = mx.fast.rms_norm(
+            hidden_states, self.weight, self.variance_epsilon
+        )
+
+        if self.norm_before_gate and gate is not None:
+            hidden_states = hidden_states * nn.silu(gate)
+        return hidden_states
+
+
+def compute_mup_vector(args):
+    intermediate_size = args.mamba_d_ssm
+    groups_time_state_size = args.mamba_n_groups * args.mamba_d_state
+    num_heads = args.mamba_n_heads
+    sizes = [
+        intermediate_size,
+        intermediate_size,
+        groups_time_state_size,
+        groups_time_state_size,
+        num_heads,
+    ]
+    return mx.concatenate(
+        [
+            mx.broadcast_to(mx.array(m), (s,))
+            for s, m in zip(sizes, args.ssm_multipliers)
+        ]
+    )
+
+
+class FalconH1Attention(nn.Module):
+
+    def __init__(self, args):
+        super().__init__()
+
+        self.hidden_size = args.hidden_size
+        self.num_heads = args.num_attention_heads
+        self.num_kv_heads = args.num_key_value_heads
+        self.head_dim = args.head_dim
+        self.scale = self.head_dim**-0.5
+
+        self.q_proj = nn.Linear(
+            self.hidden_size, self.num_heads * self.head_dim, bias=args.attention_bias
+        )
+        self.k_proj = nn.Linear(
+            self.hidden_size,
+            self.num_kv_heads * self.head_dim,
+            bias=args.attention_bias,
+        )
+        self.v_proj = nn.Linear(
+            self.hidden_size,
+            self.num_kv_heads * self.head_dim,
+            bias=args.attention_bias,
+        )
+        self.o_proj = nn.Linear(
+            self.num_heads * self.head_dim, self.hidden_size, bias=args.attention_bias
+        )
+
+        self.rope = initialize_rope(
+            self.head_dim,
+            args.rope_theta,
+            args.rope_traditional,
+            args.rope_scaling,
+            args.max_position_embeddings,
+        )
+
+    def __call__(self, x, mask=None, cache=None):
+        B, L, _ = x.shape
+
+        queries = self.q_proj(x)
+        keys = self.k_proj(x)
+        values = self.v_proj(x)
+
+        queries = queries.reshape(B, L, self.num_heads, -1).transpose(0, 2, 1, 3)
+        keys = keys.reshape(B, L, self.num_kv_heads, -1).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.num_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, mask=mask, scale=self.scale, cache=cache
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+
+        return self.o_proj(output)
+
+
+class FalconH1Mixer(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.num_heads = args.mamba_n_heads
+        self.hidden_size = args.hidden_size
+        self.ssm_state_size = args.mamba_d_state
+        self.conv_kernel_size = args.mamba_d_conv
+        self.intermediate_size = args.mamba_d_ssm
+        self.use_conv_bias = args.mamba_conv_bias
+
+        self.layer_norm_epsilon = args.rms_norm_eps
+        self.groups_time_state_size = args.mamba_n_groups * self.ssm_state_size
+
+        self.n_groups = args.mamba_n_groups
+        self.head_dim = args.mamba_d_head
+        self.chunk_size = args.mamba_chunk_size
+
+        self.time_step_limit = (0.0, float("inf"))
+        self.time_step_min = 0.001
+        self.time_step_max = 0.1
+
+        self.conv_dim = self.intermediate_size + 2 * self.n_groups * self.ssm_state_size
+        self.conv1d = nn.Conv1d(
+            in_channels=self.conv_dim,
+            out_channels=self.conv_dim,
+            bias=self.use_conv_bias,
+            kernel_size=self.conv_kernel_size,
+            groups=self.conv_dim,
+        )
+
+        projection_size = self.intermediate_size + self.conv_dim + self.num_heads
+        self.in_proj = nn.Linear(
+            self.hidden_size,
+            projection_size,
+            bias=args.mamba_proj_bias,
+        )
+
+        self.dt_bias = mx.ones(self.num_heads)
+
+        A = mx.arange(1, self.num_heads + 1)
+        self.A_log = mx.log(A)
+
+        self.mamba_rms_norm = args.mamba_rms_norm
+        if self.mamba_rms_norm:
+            self.norm = FalconH1RMSNormGated(
+                self.intermediate_size,
+                eps=self.layer_norm_epsilon,
+                n_groups=self.n_groups,
+                norm_before_gate=args.mamba_norm_before_gate,
+            )
+
+        self.D = mx.ones(self.num_heads)
+
+        self.out_proj = nn.Linear(
+            self.intermediate_size, self.hidden_size, bias=args.projectors_bias
+        )
+
+    def _apply_conv(
+        self, conv_input: mx.array, cache: Optional[MambaCache] = None
+    ) -> mx.array:
+        if cache is None or cache[0] is None:
+            conv_state = mx.zeros(
+                (conv_input.shape[0], self.conv_kernel_size - 1, self.conv_dim),
+                dtype=conv_input.dtype,
+            )
+        else:
+            conv_state = cache[0]
+
+        padded_input = mx.concatenate([conv_state, conv_input], axis=1)
+
+        if cache is not None:
+            cache[0] = padded_input[:, -(self.conv_kernel_size - 1) :]
+
+        conv_output = self.conv1d(padded_input)
+        return nn.silu(conv_output)
+
+    def _ssm(
+        self,
+        hidden_states: mx.array,
+        B: mx.array,
+        C: mx.array,
+        dt: mx.array,
+        state: Optional[mx.array] = None,
+        mask: Optional[mx.array] = None,
+    ) -> mx.array:
+        batch_size, seq_len, _ = hidden_states.shape
+
+        hidden_states = hidden_states.reshape(
+            batch_size, seq_len, self.num_heads, self.head_dim
+        )
+        B = B.reshape(batch_size, seq_len, self.n_groups, self.ssm_state_size)
+        C = C.reshape(batch_size, seq_len, self.n_groups, self.ssm_state_size)
+
+        y, state = ssm_update(
+            hidden_states,
+            self.A_log,
+            B,
+            C,
+            self.D,
+            dt,
+            self.dt_bias,
+            state,
+            self.time_step_limit,
+            mask,
+        )
+
+        return y.reshape(batch_size, seq_len, self.intermediate_size), state
+
+    def __call__(self, input_states, cache=None, mask: Optional[mx.array] = None):
+        projected_states = self.in_proj(input_states)
+
+        gate, conv_input, dt = mx.split(
+            projected_states,
+            [self.intermediate_size, self.intermediate_size + self.conv_dim],
+            axis=-1,
+        )
+
+        if mask is not None:
+            conv_input = mx.where(mask[..., None], conv_input, 0)
+        conv_output = self._apply_conv(conv_input, cache)
+
+        hidden_states_ssm, B, C = mx.split(
+            conv_output,
+            [
+                self.intermediate_size,
+                self.intermediate_size + self.n_groups * self.ssm_state_size,
+            ],
+            axis=-1,
+        )
+        state = cache[1] if cache else None
+        y, state = self._ssm(hidden_states_ssm, B, C, dt, state, mask)
+        if cache:
+            cache[1] = state
+
+        if self.mamba_rms_norm:
+            y = self.norm(y, gate)
+        else:
+            y = y * nn.silu(gate)
+
+        return self.out_proj(y)
+
+
+class FalconH1MLP(nn.Module):
+
+    def __init__(self, args):
+        super().__init__()
+
+        hidden_size = args.hidden_size
+        intermediate_size = args.intermediate_size
+
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=args.mlp_bias)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=args.mlp_bias)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=args.mlp_bias)
+
+    def __call__(self, x):
+        y = self.up_proj(x) * nn.silu(self.gate_proj(x))
+        y = self.down_proj(y)
+        return y
+
+
+class FalconH1DecoderLayer(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.feed_forward = FalconH1MLP(args)
+
+        head_dim = args.head_dim
+        self.channels_attn = (
+            args.num_attention_heads * head_dim
+            + 2 * args.num_key_value_heads * head_dim
+        )
+
+        self.mamba = FalconH1Mixer(args=args)
+
+        self.self_attn = FalconH1Attention(args)
+
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.pre_ff_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        h: mx.array,
+        cache,
+        attn_mask: Optional[mx.array],
+        mamba_mask: Optional[mx.array],
+    ) -> mx.array:
+
+        residual = h
+        h = self.input_layernorm(h)
+
+        mamba_h = self.mamba(input_states=h, cache=cache[0], mask=mamba_mask)
+
+        attn_h = self.self_attn(
+            h,
+            mask=attn_mask,
+            cache=cache[1],
+        )
+
+        h = residual + mamba_h + attn_h
+
+        residual = h
+        h = self.pre_ff_layernorm(h)
+        h = self.feed_forward(h)
+        return residual + h
+
+
+class FalconH1Model(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+
+        self.args = args
+        self.vocab_size = args.vocab_size
+        self.hidden_size = args.hidden_size
+
+        self.embed_tokens = nn.Embedding(self.vocab_size, self.hidden_size)
+
+        self._mup_vector = compute_mup_vector(args)
+        self.layers = [
+            FalconH1DecoderLayer(args) for _ in range(args.num_hidden_layers)
+        ]
+        self.final_layernorm = nn.RMSNorm(self.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(self, inputs, cache=None):
+
+        h = self.embed_tokens(inputs)
+
+        h = h
+
+        if cache is None:
+            cache = [(None, None) * len(self.layers)]
+
+        mamba_mask = create_ssm_mask(h, cache[0][0])
+        attn_mask = create_attention_mask(h, cache[0][1])
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(
+                h,
+                cache=c,
+                attn_mask=attn_mask,
+                mamba_mask=mamba_mask,
+            )
+
+        return self.final_layernorm(h)
+
+
+class Model(nn.Module):
+
+    def __init__(self, args):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = FalconH1Model(args=args)
+        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(self, inputs, cache=None):
+        hidden_states = self.model(inputs, cache=cache)
+        return self.lm_head(hidden_states)
+
+    def sanitize(self, weights):
+        # Check if needs sanitization
+        c1d = weights["model.layers.0.mamba.conv1d.weight"]
+        if c1d.shape[-1] <= c1d.shape[1]:
+            return weights
+
+        sanitized_weights = {}
+        args = self.args
+
+        for name, param in weights.items():
+            # Fold-in multipliers
+            if name.endswith("embed_tokens.weight"):
+                param *= args.embedding_multiplier
+            elif name.endswith("lm_head.weight"):
+                param *= args.lm_head_multiplier
+            elif name.endswith("q_proj.weight") or name.endswith("k_proj.weight"):
+                param *= args.attention_in_multiplier
+            elif name.endswith("key_proj.weight"):
+                param *= args.attention_in_multiplier * args.key_multiplier
+            elif name.endswith("o_proj.weight"):
+                param *= args.attention_out_multiplier
+            elif name.endswith("out_proj.weight"):
+                param *= args.ssm_out_multiplier
+            elif name.endswith("gate_proj.weight"):
+                param *= args.mlp_multipliers[0]
+            elif name.endswith("down_proj.weight"):
+                param *= args.mlp_multipliers[1]
+            elif name.endswith("in_proj.weight"):
+                param *= (
+                    args.ssm_in_multiplier
+                    * self.model._mup_vector.astype(param.dtype)[:, None]
+                )
+            elif "conv1d.weight" in name:
+                param = param.transpose(0, 2, 1)
+            sanitized_weights[name] = param
+        return sanitized_weights
+
+    def make_cache(self):
+        return [
+            CacheList(MambaCache(), KVCache())
+            for _ in range(self.args.num_hidden_layers)
+        ]
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/gated_delta.py

@@ -12,10 +12,11 @@ def compute_g(A_log, a, dt_bias):
     )
 
 
-def _make_gated_delta_kernel():
+def _make_gated_delta_kernel(has_mask=False):
     if not mx.metal.is_available():
         return None
-    source = """
+    mask_source = "mask[b_idx * T + t]" if has_mask else "true"
+    source = f"""
         auto n = thread_position_in_grid.z;
         auto b_idx = n / Hv;
         auto hv_idx = n % Hv;
@@ -38,36 +39,38 @@ def _make_gated_delta_kernel():
         auto o_state = state_out + (n * Dv + dv_idx) * Dk;
 
         float state[n_per_t];
-        for (int i = 0; i < n_per_t; ++i) {
-           auto s_idx = n_per_t * dk_idx + i;
-           state[i] = static_cast<float>(i_state[s_idx]);
-        }
+        for (int i = 0; i < n_per_t; ++i) {{
+          auto s_idx = n_per_t * dk_idx + i;
+          state[i] = static_cast<float>(i_state[s_idx]);
+        }}
 
         // beta, g: [B, T, Hv]
         auto g_ = g + b_idx * T * Hv;
         auto beta_ = beta + b_idx * T * Hv;
 
-        for (int t = 0; t < T; ++t) {
-          float kv_mem = 0.0f;
-          for (int i = 0; i < n_per_t; ++i) {
-            auto s_idx = n_per_t * dk_idx + i;
-            state[i] = state[i] * g_[hv_idx];
-            kv_mem += state[i] * k_[s_idx];
-          }
-          kv_mem = simd_sum(kv_mem);
+        for (int t = 0; t < T; ++t) {{
+          if ({mask_source}) {{
+            float kv_mem = 0.0f;
+            for (int i = 0; i < n_per_t; ++i) {{
+              auto s_idx = n_per_t * dk_idx + i;
+              state[i] = state[i] * g_[hv_idx];
+              kv_mem += state[i] * k_[s_idx];
+            }}
+            kv_mem = simd_sum(kv_mem);
 
-          auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx];
+            auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx];
 
-          float out = 0.0f;
-          for (int i = 0; i < n_per_t; ++i) {
-            auto s_idx = n_per_t * dk_idx + i;
-            state[i] = state[i] + k_[s_idx] * delta;
-            out += state[i] * q_[s_idx];
-          }
-          out = simd_sum(out);
-          if (thread_index_in_simdgroup == 0) {
-            y[dv_idx] = static_cast<InT>(out);
-          }
+            float out = 0.0f;
+            for (int i = 0; i < n_per_t; ++i) {{
+              auto s_idx = n_per_t * dk_idx + i;
+              state[i] = state[i] + k_[s_idx] * delta;
+              out += state[i] * q_[s_idx];
+            }}
+            out = simd_sum(out);
+            if (thread_index_in_simdgroup == 0) {{
+              y[dv_idx] = static_cast<InT>(out);
+            }}
+          }}
           // Increment data pointers to next time step
           q_ += Hk * Dk;
           k_ += Hk * Dk;
@@ -75,23 +78,28 @@ def _make_gated_delta_kernel():
           y += Hv * Dv;
           g_ += Hv;
           beta_ += Hv;
-        }
-        for (int i = 0; i < n_per_t; ++i) {
+        }}
+        for (int i = 0; i < n_per_t; ++i) {{
           auto s_idx = n_per_t * dk_idx + i;
           o_state[s_idx] = static_cast<InT>(state[i]);
-        }
+        }}
     """
+    inputs = ["q", "k", "v", "g", "beta", "state_in", "T"]
+    if has_mask:
+        inputs.append("mask")
     return mx.fast.metal_kernel(
-        name="gated_delta_step",
-        input_names=["q", "k", "v", "g", "beta", "state_in", "T"],
+        name="gated_delta_step" + "_mask" if has_mask else "",
+        input_names=inputs,
         output_names=["y", "state_out"],
         source=source,
     )
 
 
 _gated_delta_kernel = _make_gated_delta_kernel()
+_gated_delta_kernel_masked = _make_gated_delta_kernel(True)
 
 
+@mx.compile
 def _gated_delta_step_ops(
     q: mx.array,
     k: mx.array,
@@ -99,6 +107,7 @@ def _gated_delta_step_ops(
     g: mx.array,
     beta: mx.array,
     state: mx.array,
+    mask: Optional[mx.array] = None,
 ) -> Tuple[mx.array, mx.array]:
     """
     Ops-based reference implementation for a single recurrent step.
@@ -114,12 +123,15 @@ def _gated_delta_step_ops(
     """
 
     # Decay
+    old_state = state
     state = state * g[..., None, None]
     kv_mem = (state * k[..., None, :]).sum(axis=-1)  # [B, H, Dv]
     delta = (v - kv_mem) * beta[..., None]  # [B, H, Dv]
     state = state + k[..., None, :] * delta[..., None]
     # Output projection along key dim with q
     y = (state * q[..., None, :]).sum(axis=-1)  # [B, H, Dv]
+    if mask is not None:
+        state = mx.where(mask, state, old_state)
     return y, state
 
 
@@ -130,12 +142,18 @@ def gated_delta_kernel(
     g: mx.array,
     beta: mx.array,
     state: mx.array,
+    mask: Optional[mx.array] = None,
 ) -> Tuple[mx.array, mx.array]:
     B, T, Hk, Dk = k.shape
     Hv, Dv = v.shape[2:]
     input_type = q.dtype
-    return _gated_delta_kernel(
-        inputs=[q, k, v, g, beta, state, T],
+    kernel = _gated_delta_kernel
+    inputs = [q, k, v, g, beta, state, T]
+    if mask is not None:
+        kernel = _gated_delta_kernel_masked
+        inputs.append(mask)
+    return kernel(
+        inputs=inputs,
         template=[
             ("InT", input_type),
             ("Dk", Dk),
@@ -157,6 +175,7 @@ def gated_delta_ops(
     g: mx.array,
     beta: mx.array,
     state: Optional[mx.array] = None,
+    mask: Optional[mx.array] = None,
 ) -> Tuple[mx.array, mx.array]:
     """
     Ops-based reference implementation for prompt prefill (sequential loop).
@@ -181,14 +200,25 @@ def gated_delta_ops(
 
     ys = []
     for t in range(T):
-        y, state = _gated_delta_step_ops(
-            q[:, t],
-            k[:, t],
-            v[:, t],
-            g[:, t],
-            beta[:, t],
-            state,
-        )
+        if mask is not None:
+            y, state = _gated_delta_step_ops(
+                q[:, t],
+                k[:, t],
+                v[:, t],
+                g[:, t],
+                beta[:, t],
+                state,
+                mask[:, t],
+            )
+        else:
+            y, state = _gated_delta_step_ops(
+                q[:, t],
+                k[:, t],
+                v[:, t],
+                g[:, t],
+                beta[:, t],
+                state,
+            )
         ys.append(y)
     y = mx.stack(ys, axis=1)
     return y, state
@@ -203,6 +233,8 @@ def gated_delta_update(
     A_log: mx.array,
     dt_bias: mx.array,
     state: Optional[mx.array] = None,
+    mask: Optional[mx.array] = None,
+    use_kernel: bool = True,
 ) -> Tuple[mx.array, mx.array]:
 
     beta = mx.sigmoid(b)
@@ -213,7 +245,7 @@ def gated_delta_update(
         if state is None:
             state = mx.zeros((B, Hv, Dv, Dk), dtype=q.dtype)
 
-    if mx.default_device() != mx.gpu or not mx.metal.is_available():
-        return gated_delta_ops(q, k, v, g, beta, state)
+    if not use_kernel or mx.default_device() != mx.gpu or not mx.metal.is_available():
+        return gated_delta_ops(q, k, v, g, beta, state, mask)
     else:
-        return gated_delta_kernel(q, k, v, g, beta, state)
+        return gated_delta_kernel(q, k, v, g, beta, state, mask)

mlx_lm/models/gemma3_text.py

@@ -87,8 +87,6 @@ class Attention(nn.Module):
             keys = self.rope(keys)
 
         # Sliding window
-        if isinstance(mask, mx.array) and mask.shape[-1] != keys.shape[-2]:
-            mask = mask[..., -keys.shape[-2] :]
         output = scaled_dot_product_attention(
             queries, keys, values, cache=cache, scale=self.scale, mask=mask
         )
@@ -194,7 +192,6 @@ class Gemma3Model(nn.Module):
             cache[0],
             window_size=self.window_size,
         )
-
         for i, (layer, c) in enumerate(zip(self.layers, cache)):
             is_global = (
                 i % self.sliding_window_pattern == self.sliding_window_pattern - 1
@@ -246,7 +243,5 @@ class Model(nn.Module):
             ):
                 caches.append(KVCache())
             else:
-                caches.append(
-                    RotatingKVCache(max_size=self.args.sliding_window, keep=0)
-                )
+                caches.append(RotatingKVCache(max_size=self.args.sliding_window))
         return caches

mlx_lm/models/lfm2-vl.py

@@ -18,11 +18,6 @@ class ModelArgs(BaseModelArgs):
 
     def __post_init__(self):
         self.text_config["tie_word_embeddings"] = False
-        self.text_config["full_attn_idxs"] = [
-            i
-            for i, layer_type in enumerate(self.text_config["layer_types"])
-            if layer_type == "full_attention"
-        ]
 
 
 class Model(nn.Module):

mlx_lm/models/lfm2.py

@@ -31,8 +31,19 @@ class ModelArgs(BaseModelArgs):
     block_multiple_of: int
     block_ffn_dim_multiplier: float
     block_auto_adjust_ff_dim: bool
-    full_attn_idxs: List[int]
     rope_theta: float
+    full_attn_idxs: Optional[List[int]] = None
+    layer_types: Optional[List[str]] = None
+
+    def __post_init__(self):
+        if self.num_key_value_heads is None:
+            self.num_key_value_heads = self.num_attention_heads
+        if self.full_attn_idxs is None:
+            self.full_attn_idxs = [
+                i
+                for i, layer_type in enumerate(self.layer_types)
+                if layer_type == "full_attention"
+            ]
 
 
 class Attention(nn.Module):

mlx_lm/models/llama.py

@@ -1,12 +1,13 @@
 # Copyright © 2023-2024 Apple Inc.
 
 from dataclasses import dataclass
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 import mlx.core as mx
 import mlx.nn as nn
 
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .cache import KVCache, RotatingKVCache
 from .rope_utils import initialize_rope
 
 
@@ -28,11 +29,16 @@ class ModelArgs(BaseModelArgs):
     rope_traditional: bool = False
     rope_scaling: Optional[Dict[str, Union[float, str]]] = None
     tie_word_embeddings: bool = True
+    layer_types: Optional[List[str]] = None
+    sliding_window: Optional[int] = None
 
     def __post_init__(self):
         if self.num_key_value_heads is None:
             self.num_key_value_heads = self.num_attention_heads
 
+        if self.layer_types is None:
+            self.layer_types = ["full_attention"] * self.num_hidden_layers
+
 
 class Attention(nn.Module):
     def __init__(self, args: ModelArgs):
@@ -114,10 +120,11 @@ class MLP(nn.Module):
 
 
 class TransformerBlock(nn.Module):
-    def __init__(self, args: ModelArgs):
+    def __init__(self, args: ModelArgs, use_sliding: bool = False):
         super().__init__()
         self.num_attention_heads = args.num_attention_heads
         self.hidden_size = args.hidden_size
+        self.use_sliding = use_sliding
         self.self_attn = Attention(args)
         self.mlp = MLP(args)
         self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
@@ -145,12 +152,21 @@ class LlamaModel(nn.Module):
         self.args = args
         self.vocab_size = args.vocab_size
         self.num_hidden_layers = args.num_hidden_layers
+        self.layer_types = args.layer_types
+        self.sliding_window = args.sliding_window
         assert self.vocab_size > 0
         self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
         self.layers = [
-            TransformerBlock(args=args) for _ in range(args.num_hidden_layers)
+            TransformerBlock(args=args, use_sliding=layer_type == "sliding_attention")
+            for layer_type in self.layer_types
         ]
         self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.fa_idx = self.layer_types.index("full_attention")
+        self.swa_idx = None
+        for e, l in enumerate(self.layers):
+            if l.use_sliding:
+                self.swa_idx = e
+                break
 
     def __call__(
         self,
@@ -166,10 +182,15 @@ class LlamaModel(nn.Module):
         if cache is None:
             cache = [None] * len(self.layers)
 
-        mask = create_attention_mask(h, cache[0])
+        fa_mask = create_attention_mask(h, cache[self.fa_idx])
+        if self.swa_idx is not None:
+            swa_mask = create_attention_mask(
+                h, cache[self.swa_idx], window_size=self.sliding_window
+            )
 
-        for layer, c in zip(self.layers, cache):
-            h = layer(h, mask, cache=c)
+        for layer, cache in zip(self.layers, cache):
+            mask = swa_mask if layer.use_sliding else fa_mask
+            h = layer(h, mask, cache=cache)
 
         return self.norm(h)
 
@@ -208,3 +229,13 @@ class Model(nn.Module):
     @property
     def layers(self):
         return self.model.layers
+
+    def make_cache(self):
+        return [
+            (
+                RotatingKVCache(max_size=self.model.sliding_window)
+                if layer.use_sliding
+                else KVCache()
+            )
+            for layer in self.layers
+        ]

mlx_lm/models/qwen2_vl.py

@@ -20,7 +20,7 @@ class ModelArgs(BaseModelArgs):
     def from_dict(cls, params):
         if "text_config" not in params:
             return cls(model_type=params["model_type"], text_config=params)
-        return cls(**params)
+        return super().from_dict(params)
 
 
 class Model(nn.Module):

mlx_lm/models/qwen3_next.py

@@ -6,7 +6,12 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .base import (
+    BaseModelArgs,
+    create_attention_mask,
+    create_ssm_mask,
+    scaled_dot_product_attention,
+)
 from .cache import KVCache, MambaCache
 from .gated_delta import gated_delta_update
 from .rope_utils import initialize_rope
@@ -237,6 +242,8 @@ class Qwen3NextGatedDeltaNet(nn.Module):
         mixed_qkv = mx.concatenate(
             [q.reshape(B, S, -1), k.reshape(B, S, -1), v.reshape(B, S, -1)], axis=-1
         )
+        if mask is not None:
+            mixed_qkv = mx.where(mask[..., None], mixed_qkv, 0)
         conv_input = mx.concatenate([conv_state, mixed_qkv], axis=1)
         if cache is not None:
             cache[0] = conv_input[:, -(self.conv_kernel_size - 1) :]
@@ -251,14 +258,23 @@ class Qwen3NextGatedDeltaNet(nn.Module):
             )
         ]
 
-        if cache is not None:
-            state = cache[1]
-
+        state = cache[1] if cache else None
         inv_scale = k.shape[-1] ** -0.5
         q = (inv_scale**2) * mx.fast.rms_norm(q, None, 1e-6)
         k = inv_scale * mx.fast.rms_norm(k, None, 1e-6)
 
-        out, state = gated_delta_update(q, k, v, a, b, self.A_log, self.dt_bias, state)
+        out, state = gated_delta_update(
+            q,
+            k,
+            v,
+            a,
+            b,
+            self.A_log,
+            self.dt_bias,
+            state,
+            mask,
+            use_kernel=not self.training,
+        )
 
         if cache is not None:
             cache[1] = state
@@ -350,6 +366,7 @@ class Qwen3NextModel(nn.Module):
             for i in range(args.num_hidden_layers)
         ]
         self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.ssm_idx = 0
         self.fa_idx = args.full_attention_interval - 1
 
     def __call__(
@@ -362,9 +379,11 @@ class Qwen3NextModel(nn.Module):
         if cache is None:
             cache = [None] * len(self.layers)
 
-        mask = create_attention_mask(hidden_states, cache[self.fa_idx])
+        fa_mask = create_attention_mask(hidden_states, cache[self.fa_idx])
+        ssm_mask = create_ssm_mask(hidden_states, cache[self.ssm_idx])
 
         for layer, c in zip(self.layers, cache):
+            mask = ssm_mask if layer.is_linear else fa_mask
             hidden_states = layer(hidden_states, mask=mask, cache=c)
 
         return self.norm(hidden_states)

mlx_lm/quant/dwq.py

@@ -324,6 +324,10 @@ def main():
             bits=args.bits,
         )
 
+    if mx.metal.is_available():
+        max_rec_size = mx.metal.device_info()["max_recommended_working_set_size"]
+        mx.set_wired_limit(max_rec_size)
+
     opt = optimizers.Adam(learning_rate=args.learning_rate, bias_correction=True)
     dwq_quantize(
         model,

mlx_lm/server.py

@@ -279,8 +279,6 @@ class APIHandler(BaseHTTPRequestHandler):
         endpoints = {
             "/v1/completions": self.handle_text_completions,
             "/v1/chat/completions": self.handle_chat_completions,
-            "/responses": self.handle_responses,
-            "/v1/responses": self.handle_responses,
             "/chat/completions": self.handle_chat_completions,
         }
 
@@ -330,7 +328,6 @@ class APIHandler(BaseHTTPRequestHandler):
             self.max_tokens = self.body.get(
                 "max_tokens", self.model_provider.cli_args.max_tokens
             )
-
         self.temperature = self.body.get(
             "temperature", self.model_provider.cli_args.temp
         )
@@ -495,18 +492,6 @@ class APIHandler(BaseHTTPRequestHandler):
                 "id": None,
             }
 
-        if self.stream and self.object_type == "response" and finish_reason is None:
-            return {
-                "type": "response.output_text.delta",
-                "delta": text,
-                # TODO, these need valid values
-                "sequence_number": None,
-                "item_id": None,
-                "output_index": 1,
-                "content_index": 0,
-                "logprobs": [],
-            }
-
         # Static response
         response = {
             "id": self.request_id,
@@ -514,27 +499,13 @@ class APIHandler(BaseHTTPRequestHandler):
             "object": self.object_type,
             "model": self.requested_model,
             "created": self.created,
-        }
-
-        if self.object_type == "response":
-            response["output"] = [
+            "choices": [
                 {
-                    "type": "message",
-                    "role": "assistant",
-                    "content": [{"text": text, "type": "output_text"}],
-                }
-            ]
-            if self.stream:
-                return {"response": response, "type": "response.completed"}
-
-            return response
-
-        response["choices"] = [
-            {
-                "index": 0,
-                "finish_reason": finish_reason,
-            },
-        ]
+                    "index": 0,
+                    "finish_reason": finish_reason,
+                },
+            ],
+        }
 
         if token_logprobs or top_logprobs or tokens:
             response["choices"][0]["logprobs"] = {
@@ -893,36 +864,6 @@ class APIHandler(BaseHTTPRequestHandler):
 
         return prompt
 
-    def handle_responses(self) -> List[int]:
-        body = self.body
-        system_prompt = body.get("instructions")
-        prompt = body["input"]
-        tools = body.get("tools")
-        messages = []
-        if system_prompt:
-            messages = [{"role": "system", "content": system_prompt}]
-        if isinstance(prompt, list):
-            for message in prompt:
-                content = message["content"]
-                if isinstance(content, list):
-                    if len(content) != 1 or content[0]["type"] != "input_text":
-                        raise ValueError("Unsupported content type.")
-                    message["content"] = content[0]["text"]
-                messages.append(message)
-        else:
-            messages.append({"role": "user", "content": prompt})
-
-        # Determine response type
-        self.request_id = f"resp_{uuid.uuid4()}"
-        self.object_type = "response"
-        prompt = self.tokenizer.apply_chat_template(
-            messages,
-            tools=tools,
-            add_generation_prompt=True,
-            **self.model_provider.cli_args.chat_template_args,
-        )
-        return prompt
-
     def handle_text_completions(self) -> List[int]:
         """
         Handle a text completion request.

mlx_lm/utils.py

@@ -261,7 +261,7 @@ def load(
     """
     model_path, _ = get_model_path(path_or_hf_repo)
 
-    model, config = load_model(model_path, lazy)
+    model, config = load_model(model_path, lazy, model_config=model_config)
     if adapter_path is not None:
         model = load_adapters(model, adapter_path)
         model.eval()

requirements.txt

@@ -1,4 +1,4 @@
-mlx>=0.29.1
+mlx>=0.29.2
 numpy
 transformers>=4.39.3
 protobuf

tests/test_generate.py

@@ -11,6 +11,7 @@ from mlx_lm.generate import (
     generate,
     stream_generate,
 )
+from mlx_lm.models.cache import RotatingKVCache
 from mlx_lm.sample_utils import make_logits_processors, make_sampler
 from mlx_lm.utils import load
 
@@ -301,6 +302,56 @@ class TestGenerate(unittest.TestCase):
             batch_tokens = batch_responses[uids[e]]
             self.assertEqual(tokens, batch_tokens)
 
+    def test_batch_sliding_window(self):
+        prompts = [
+            "Write a story about Einstein",
+            "Hi",
+            "What time is it?",
+            "How tall is Mt Everest?",
+        ]
+        prompts = [
+            self.tokenizer.apply_chat_template(
+                [{"role": "user", "content": p}],
+                tokenize=True,
+                add_generation_prompt=True,
+            )
+            for p in prompts
+        ]
+
+        self.model.make_cache = lambda: [
+            RotatingKVCache(max_size=4) for _ in self.model.layers
+        ]
+        batch_gen = BatchGenerator(
+            self.model,
+            stop_tokens=self.tokenizer.eos_token_ids,
+            max_tokens=10,
+            prefill_batch_size=1,
+            prefill_step_size=8,
+            completion_batch_size=2,
+        )
+        uids = batch_gen.insert(prompts)
+        batch_responses = {uid: [] for uid in uids}
+        while responses := batch_gen.next():
+            for r in responses:
+                batch_responses[r.uid].append(r.logprobs)
+
+        for e, uid in enumerate(uids):
+            for i, response in enumerate(
+                stream_generate(
+                    self.model,
+                    self.tokenizer,
+                    prompts[e],
+                    max_tokens=10,
+                )
+            ):
+                batch_logprobs = batch_responses[uid][i]
+                logprobs = response.logprobs
+                self.assertTrue(
+                    mx.allclose(batch_logprobs, logprobs, rtol=1e-4, atol=1e-4)
+                )
+
+        del self.model.make_cache
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_models.py

@@ -37,7 +37,7 @@ class TestModels(unittest.TestCase):
 
     def test_rotating_kv_cache(self):
         b, h, d = 1, 2, 32
-        cache = RotatingKVCache(max_size=8, step=4)
+        cache = RotatingKVCache(max_size=8)
 
         k = mx.random.uniform(shape=(b, h, 2, d))
         v = mx.random.uniform(shape=(b, h, 2, d))
@@ -70,7 +70,7 @@ class TestModels(unittest.TestCase):
             idx %= 8
 
         # Try with nonzero keep
-        cache = RotatingKVCache(max_size=8, step=4, keep=2)
+        cache = RotatingKVCache(max_size=8, keep=2)
 
         # Check a large update
         k = mx.random.uniform(shape=(b, h, 20, d))
@@ -98,7 +98,7 @@ class TestModels(unittest.TestCase):
         # alternating prompt/prefill with generation
         d = 4
         h = 2
-        cache = RotatingKVCache(max_size=18, step=4)
+        cache = RotatingKVCache(max_size=18)
 
         x = mx.random.uniform(shape=(1, h, 8, d))
         k, v = cache.update_and_fetch(x, x)
@@ -175,6 +175,49 @@ class TestModels(unittest.TestCase):
         sums = mask.sum(axis=1)
         self.assertTrue(mx.array_equal(sums, expected_sums))
 
+    def test_llama_model_sliding_attention(self):
+        from mlx_lm.models import llama
+
+        args = llama.ModelArgs(
+            model_type="llama",
+            hidden_size=64,
+            num_hidden_layers=4,
+            intermediate_size=256,
+            num_attention_heads=8,
+            num_key_value_heads=4,
+            rms_norm_eps=1e-5,
+            vocab_size=128,
+            sliding_window=4,
+            layer_types=[
+                "full_attention",
+                "sliding_attention",
+                "sliding_attention",
+                "full_attention",
+            ],
+            tie_word_embeddings=False,
+            rope_theta=10000.0,
+        )
+        model = llama.Model(args)
+
+        tokens = mx.array([[1, 2, 3, 4, 5]], dtype=mx.int32)
+        out = model(tokens)
+        mx.eval(out)
+        self.assertEqual(out.shape, (1, 5, args.vocab_size))
+
+        caches = model.make_cache()
+        self.assertIsInstance(caches[0], KVCache)
+        self.assertIsInstance(caches[1], RotatingKVCache)
+        self.assertIsInstance(caches[2], RotatingKVCache)
+        self.assertIsInstance(caches[3], KVCache)
+
+        caches = model.make_cache()
+        step = model(tokens[:, :2], cache=caches)
+        mx.eval(step)
+        step = model(tokens[:, 2:3], cache=caches)
+        mx.eval(step)
+        self.assertEqual(caches[0].offset, 3)
+        self.assertEqual(caches[1].offset, 3)
+
     def test_rope(self):
         rope = rope_utils.initialize_rope(32, base=100, traditional=False)
         self.assertTrue(isinstance(rope, nn.RoPE))
@@ -666,6 +709,19 @@ class TestModels(unittest.TestCase):
             model, args.model_type, args.vocab_size, args.num_hidden_layers
         )
 
+    def test_falcon_h1(self):
+        from mlx_lm.models import falcon_h1
+
+        args = falcon_h1.ModelArgs(
+            model_type="falcon_h1",
+            num_hidden_layers=12,
+            vocab_size=10000,
+        )
+        model = falcon_h1.Model(args)
+        self.model_test_runner(
+            model, args.model_type, args.vocab_size, args.num_hidden_layers
+        )
+
     def test_gpt2(self):
         from mlx_lm.models import gpt2
 
@@ -1849,9 +1905,9 @@ class TestModels(unittest.TestCase):
         self.assertTrue(mx.allclose(out_state, out_state_m, atol=1e-4, rtol=1e-4))
 
     def test_gated_delta(self):
+        mx.random.seed(0)
         for B in [1, 2]:
             for T in [1, 2]:
-                B = 1
                 Hk = 16
                 Hv = 32
                 Dk = 128
@@ -1860,14 +1916,45 @@ class TestModels(unittest.TestCase):
                 q = mx.random.normal(shape=(B, T, Hk, Dk))
                 k = mx.random.normal(shape=(B, T, Hk, Dk))
                 v = mx.random.normal(shape=(B, T, Hv, Dv))
-                g = mx.random.normal(shape=(B, T, Hv))
-                beta = mx.random.normal(shape=(B, T, Hv))
+                g = mx.random.uniform(shape=(B, T, Hv))
+                beta = mx.random.uniform(shape=(B, T, Hv))
                 state = mx.random.normal(shape=(B, Hv, Dk, Dv))
 
                 y_op, st_op = gated_delta_ops(q, k, v, g, beta, state)
                 y_c, st_c = gated_delta_kernel(q, k, v, g, beta, state)
                 self.assertTrue(mx.allclose(y_op, y_c, rtol=1e-4, atol=1e-4))
-                self.assertTrue(mx.allclose(st_op, st_c, rtol=1e-4, atol=1e-3))
+                self.assertTrue(mx.allclose(st_op, st_c, rtol=1e-4, atol=1e-4))
+
+    def test_gated_delta_masked(self):
+        B = 1
+        T = 3
+        Hk = 16
+        Hv = 32
+        Dk = 128
+        Dv = 128
+
+        mx.random.seed(0)
+        q = mx.random.normal(shape=(B, T, Hk, Dk))
+        k = mx.random.normal(shape=(B, T, Hk, Dk))
+        v = mx.random.normal(shape=(B, T, Hv, Dv))
+        g = mx.random.normal(shape=(B, T, Hv))
+        mask = mx.array([[False, True, True]])
+        beta = mx.random.normal(shape=(B, T, Hv))
+        state = mx.random.normal(shape=(B, Hv, Dk, Dv))
+
+        y_gt, st_gt = gated_delta_ops(
+            q[:, 1:],
+            k[:, 1:],
+            v[:, 1:],
+            g[:, 1:],
+            beta[:, 1:],
+            state,
+        )
+        for fn in [gated_delta_ops, gated_delta_kernel]:
+            y, st = fn(q, k, v, g, beta, state, mask)
+            y = y[:, 1:]
+            self.assertTrue(mx.allclose(y, y_gt, rtol=1e-4, atol=1e-4))
+            self.assertTrue(mx.allclose(st, st_gt, rtol=1e-4, atol=1e-3))
 
 
 if __name__ == "__main__":

tests/test_prompt_cache.py

@@ -11,6 +11,7 @@ from mlx_lm.generate import generate_step
 from mlx_lm.models.base import create_attention_mask, create_causal_mask
 from mlx_lm.models.cache import (
     BatchKVCache,
+    BatchRotatingKVCache,
     CacheList,
     ChunkedKVCache,
     KVCache,
@@ -391,7 +392,7 @@ class TestPromptCache(unittest.TestCase):
         kv = mx.zeros((1, 1, 10, 32))
         cache.update_and_fetch(kv, kv)
         mask = cache.make_mask(3, window_size=5)
-        self.assertEqual(mask.shape, (3, 11))
+        self.assertEqual(mask.shape, (3, 10))
         self.assertTrue(mx.all(mask.sum(axis=-1) == 5))
         for i in range(3):
             s = 11 - 3 + i
@@ -405,7 +406,7 @@ class TestPromptCache(unittest.TestCase):
         self.assertEqual(mask, None)
 
         mask = cache.make_mask(1, window_size=5)
-        self.assertEqual(mask.squeeze(1).tolist(), [True] + [False] * 3 + [True] * 4)
+        self.assertEqual(mask.tolist(), [True] + [False] * 3 + [True] * 4)
         cmask = create_attention_mask(mx.zeros((1, 1, 32)), cache, window_size=5)
         self.assertTrue(mx.array_equal(cmask, mask))
 
@@ -413,9 +414,7 @@ class TestPromptCache(unittest.TestCase):
         cache.update_and_fetch(kv, kv)
 
         mask = cache.make_mask(1, window_size=5)
-        self.assertEqual(
-            mask.squeeze(1).tolist(), [True] * 2 + [False] * 3 + [True] * 3
-        )
+        self.assertEqual(mask.tolist(), [True] * 2 + [False] * 3 + [True] * 3)
         cmask = create_attention_mask(mx.zeros((1, 1, 32)), cache, window_size=5)
         self.assertTrue(mx.array_equal(cmask, mask))
 
@@ -460,6 +459,115 @@ class TestPromptCache(unittest.TestCase):
         self.assertEqual(cache_a.offset.tolist(), [6, 7, 6, 1, 4])
         self.assertEqual(cache_a.left_padding.tolist(), [2, 1, 2, 7, 4])
 
+    def test_batch_rotating_kv_cache(self):
+        cache = BatchRotatingKVCache(max_size=4, left_padding=[2, 0])
+        mask = cache.make_mask(4)
+        self.assertFalse(mx.any(mask[0, 0, 0, :]))
+        self.assertTrue(
+            mx.array_equal(mask[1, 0, 0, :], mx.array([True, False, False, False]))
+        )
+
+        # Batch update works
+        k, v = mx.zeros((2, 1, 4, 8)), mx.zeros((2, 1, 4, 8))
+        k, v = cache.update_and_fetch(k, v)
+
+        mask = cache.make_mask(4)
+        k, v = mx.zeros((2, 1, 4, 8)), mx.zeros((2, 1, 4, 8))
+        k, v = cache.update_and_fetch(k, v)
+        self.assertEqual(mask.shape[-2:], (4, k.shape[2]))
+        self.assertEqual(
+            mask[0, 0, 0, :].tolist(), [False, True, True, True, False, False, False]
+        )
+
+        # Single query update works
+        cache = BatchRotatingKVCache(max_size=4, left_padding=[2, 0])
+        k, v = mx.zeros((2, 1, 4, 8)), mx.zeros((2, 1, 4, 8))
+        k, v = cache.update_and_fetch(k, v)
+
+        mask = cache.make_mask(1)
+        k, v = mx.zeros((2, 1, 1, 8)), mx.zeros((2, 1, 1, 8))
+
+        k, v = cache.update_and_fetch(k, v)
+        self.assertEqual(mask.shape[-2:], (1, k.shape[2]))
+        self.assertEqual(mask[0, 0, 0].tolist(), [True, False, True, True])
+        self.assertEqual(mask[1, 0, 0].tolist(), [True, True, True, True])
+
+        # Check filtering
+        cache = BatchRotatingKVCache(max_size=4, left_padding=[2, 0, 3])
+        k, v = mx.zeros((3, 1, 3, 8)), mx.zeros((3, 1, 3, 8))
+        cache.update_and_fetch(k, v)
+        cache.filter(mx.array([1]))
+        self.assertEqual(cache.keys.shape, (1, 1, 3, 8))
+
+        # Check extend
+        cache = BatchRotatingKVCache(max_size=4, left_padding=[2, 1])
+        other = BatchRotatingKVCache(max_size=4, left_padding=[2, 2])
+        k, v = mx.zeros((2, 1, 5, 8)), mx.zeros((2, 1, 5, 8))
+        cache.update_and_fetch(k, v)
+        other.update_and_fetch(k, v)
+        k, v = mx.zeros((2, 1, 1, 8)), mx.zeros((2, 1, 1, 8))
+        cache.update_and_fetch(k, v)
+        cache.extend(other)
+
+        # Check mask when going from prompt -> extend -> prompt
+        cache = BatchRotatingKVCache(max_size=8, left_padding=[4])
+        k, v = mx.zeros((1, 1, 8, 8)), mx.zeros((1, 1, 8, 8))
+        cache.update_and_fetch(k, v)
+
+        mask = cache.make_mask(1)
+        self.assertEqual(
+            mask.squeeze().tolist(), [True, False, False, False, True, True, True, True]
+        )
+
+        k, v = mx.zeros((1, 1, 1, 8)), mx.zeros((1, 1, 1, 8))
+        cache.update_and_fetch(k, v)
+
+        mask = cache.make_mask(2)
+        expected = mx.array(
+            [
+                [False, False, False, True, True, True, True, True, False],
+                [False, False, False, True, True, True, True, True, True],
+            ]
+        )
+        self.assertTrue(mx.array_equal(mask.squeeze(), expected))
+
+    def test_save_load_batch_caches(self):
+        cache_file = os.path.join(self.test_dir, "prompt_cache.safetensors")
+
+        cache = [
+            MambaCache(left_padding=[1, 2]),
+            BatchKVCache(left_padding=[1, 2]),
+            BatchRotatingKVCache(max_size=10, left_padding=[1, 2]),
+        ]
+        for c in cache:
+            if isinstance(c, MambaCache):
+                c[0] = mx.random.uniform(shape=(4, 4, 4))
+                c[1] = mx.random.uniform(shape=(4, 4, 4))
+            else:
+                x = mx.random.uniform(shape=(4, 4, 7, 4))
+                y = mx.random.uniform(shape=(4, 4, 7, 4))
+                c.update_and_fetch(x, y)
+
+        save_prompt_cache(cache_file, cache)
+        loaded_cache = load_prompt_cache(cache_file)
+        left_padding = mx.array([1, 2])
+        for c, lc in zip(cache, loaded_cache):
+            self.assertTrue(mx.array_equal(c.left_padding, left_padding))
+
+    def test_rotating_cache_updates(self):
+        cache = RotatingKVCache(max_size=8)
+        k = v = mx.zeros((1, 1, 10, 1))
+        cache.update_and_fetch(k, v)
+
+        for _ in range(3):
+            k = v = mx.zeros((1, 1, 1, 1))
+            cache.update_and_fetch(k, v)
+
+        k = v = mx.zeros((1, 1, 3, 1))
+        k, v = cache.update_and_fetch(k, v)
+        self.assertEqual(k.shape[2], 10)
+        self.assertEqual(v.shape[2], 10)
+
 
 if __name__ == "__main__":
     unittest.main()