Fix the rope mutation in a more natural way

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

.github/workflows/pull_request.yml

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

ACKNOWLEDGMENTS.md

@@ -10,11 +10,11 @@ MLX LM was developed with contributions from the following individuals:
 - Shunta Saito: Added support for PLaMo models.
 - Gökdeniz Gülmez: Added support for the following architectures: 
 OpenBMB's `MiniCPM` and `MiniCPM3`, Kyutai's `Helium`, State-Space's `Mamba v1` and 
-`Mamba v2`, Z.ai & THUKEG's `GLM`, `GLM4`, Rednote `dots.llm1`, Baidu's `Ernie4.5 MoE`, 
+`Mamba v2`, Z.ai & THUKEG's `GLM`, `GLM4`, `GLM5 (GLM MoE DSA)`, Rednote `dots.llm1`, Baidu's `Ernie4.5 MoE`, 
 inclusionAI's `Bailing MoE e.g. Ling-family`, `Bailing MoE Linear e.g. Ling-Linear-family`, 
 Klear team - Kuaishou Technology's `Klear`, AI21 Lab's `Jamba` IBM's `Granite MoE`, 
 Meituan's `LongCat`, Nvidia's `Nemotron H`, Swiss-AI's `Apertus`, Nikity's `Lille130m`, 
-Alibaba Qwen's `Qwen3Next`, and Allenai's `OLMoE` and `Olmo 3`; 
+Alibaba Qwen's `Qwen3Next`, Tele-AI's `TeleChat3`, and Allenai's `OLMoE` and `Olmo 3`; 
 Helped add support for the following model architectures: 
 Alibaba Qwen's `Qwen3 & Qwen3MoE)`; Added support for the following training algorithms: 
 `Full Weight Fine-Tuning`, and the `Muon` optimizer; 
@@ -26,4 +26,8 @@ Added support for the following other features:
   MoonshotAI's `Kimi-Linear`, LiquidAI's `LFM2` and `LFM2 MoE`, 
   Google DeepMind's `Gemma 3`, TII's `Falcon H1` and InterLM's `InternLM 2.5`.
 - Ivan Fioravanti: Added support for the following architectures:
- ServiceNow-AI's `Apriel 1.5`, Tencent's `Hunyuan Dense V1` and `Hunyuan MoE V1`.
+ ServiceNow-AI's `Apriel 1.5`, Tencent's `Hunyuan Dense V1` and `Hunyuan MoE V1`.
+- Tarjei Mandt: Added support for the following architectures: `Step 3.5 Flash`, 
+MoonshotAI's `Kimi K2.5`, Upstage's `Solar Open`, LG AI Research's `K-Exaone MoE`, 
+Meituan's `LongCat Flash Lite` Helped add support for the following model architectures: 
+Z.ai & THUKEG's `GLM5 (GLM MoE DSA)`

mlx_lm/LORA.md

@@ -66,9 +66,10 @@ mlx_lm.lora \
 To fine-tune the full model weights, add the `--fine-tune-type full` flag.
 Currently supported fine-tuning types are `lora` (default), `dora`, and `full`.
 
-The `--data` argument must specify a path to a `train.jsonl`, `valid.jsonl`
-when using `--train` and a path to a `test.jsonl` when using `--test`. For more
-details on the data format see the section on [Data](#Data).
+The `--data` argument must specify a path to a `train.jsonl` when using
+`--train` and a path to a `test.jsonl` when using `--test`. A `valid.jsonl` is
+optional; if provided, validation loss will be reported during training. For
+more details on the data format see the section on [Data](#Data).
 
 For example, to fine-tune a Mistral 7B you can use `--model
 mistralai/Mistral-7B-v0.1`.
@@ -184,9 +185,10 @@ Face.
 
 ### Local Datasets
 
-For fine-tuning (`--train`), the data loader expects a `train.jsonl` and a
-`valid.jsonl` to be in the data directory. For evaluation (`--test`), the data
-loader expects a `test.jsonl` in the data directory. 
+For fine-tuning (`--train`), the data loader expects a `train.jsonl` to be in
+the data directory. A `valid.jsonl` is optional; if present, validation loss
+will be reported periodically during training. For evaluation (`--test`), the
+data loader expects a `test.jsonl` in the data directory. 
 
 Currently, `*.jsonl` files support `chat`, `tools`, `completions`, and `text`
 data formats. Here are examples of these formats:

mlx_lm/SERVER.md

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

mlx_lm/__main__.py

@@ -1,36 +1,6 @@
 # Copyright © 2025 Apple Inc.
 
-import importlib
-import sys
-
 if __name__ == "__main__":
-    subcommands = {
-        "quant.awq",
-        "quant.dwq",
-        "quant.dynamic_quant",
-        "quant.gptq",
-        "benchmark",
-        "cache_prompt",
-        "chat",
-        "convert",
-        "evaluate",
-        "fuse",
-        "generate",
-        "lora",
-        "perplexity",
-        "server",
-        "manage",
-        "upload",
-    }
-    if len(sys.argv) < 2:
-        raise ValueError(f"CLI requires a subcommand in {subcommands}")
-    subcommand = sys.argv.pop(1)
-    if subcommand in subcommands:
-        submodule = importlib.import_module(f"mlx_lm.{subcommand}")
-        submodule.main()
-    elif subcommand == "--version":
-        from mlx_lm import __version__
+    from . import cli
 
-        print(__version__)
-    else:
-        raise ValueError(f"CLI requires a subcommand in {subcommands}")
+    cli.main()

mlx_lm/_version.py

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

mlx_lm/benchmark.py

@@ -1,6 +1,7 @@
 # Copyright © 2025 Apple Inc.
 
 import argparse
+import time
 
 import mlx.core as mx
 
@@ -54,6 +55,24 @@ def setup_arg_parser():
         action="store_true",
         help="Use pipelining instead of tensor parallelism",
     )
+    parser.add_argument(
+        "--quantize-activations",
+        "-qa",
+        action="store_true",
+        help="Quantize activations using the same quantization config as the corresponding layer.",
+    )
+    parser.add_argument(
+        "--prefill-step-size",
+        type=int,
+        default=2048,
+        help="Step size for prefill processing (default: 2048)",
+    )
+    parser.add_argument(
+        "--delay",
+        type=int,
+        default=0,
+        help="Delay between each test in seconds (default: 0)",
+    )
     return parser
 
 
@@ -79,7 +98,10 @@ def main():
         )
     else:
         model, tokenizer, config = load(
-            model_path, return_config=True, tokenizer_config={"trust_remote_code": True}
+            model_path,
+            return_config=True,
+            tokenizer_config={"trust_remote_code": True},
+            model_config={"quantize_activations": args.quantize_activations},
         )
 
     # Empty to avoid early stopping
@@ -94,14 +116,22 @@ def main():
 
     def single_bench():
         for response in stream_generate(
-            model, tokenizer, prompt, max_tokens=generation_tokens
+            model,
+            tokenizer,
+            prompt,
+            max_tokens=generation_tokens,
+            prefill_step_size=args.prefill_step_size,
         ):
             pass
         return response
 
     def batch_bench():
         return batch_generate(
-            model, tokenizer, prompts, max_tokens=generation_tokens
+            model,
+            tokenizer,
+            prompts,
+            max_tokens=generation_tokens,
+            prefill_step_size=args.prefill_step_size,
         ).stats
 
     if batch_size == 1:
@@ -116,6 +146,8 @@ def main():
     rprint(f"Timing with {prompt_tokens=}, {generation_tokens=}, {batch_size=}.")
     responses = []
     for i in range(args.num_trials):
+        if args.delay > 0:
+            time.sleep(args.delay)
         response = _bench()
         responses.append(response)
         results = [(k, getattr(response, k)) for k in report_keys]

mlx_lm/chat.py

@@ -13,7 +13,7 @@ DEFAULT_TEMP = 0.0
 DEFAULT_TOP_P = 1.0
 DEFAULT_XTC_PROBABILITY = 0.0
 DEFAULT_XTC_THRESHOLD = 0.0
-DEFAULT_SEED = None
+DEFAULT_SEED = 0
 DEFAULT_MAX_TOKENS = 256
 DEFAULT_MODEL = "mlx-community/Llama-3.2-3B-Instruct-4bit"
 
@@ -100,8 +100,7 @@ def main():
         if rank == 0:
             print(*args, **kwargs)
 
-    if args.seed is not None:
-        mx.random.seed(args.seed)
+    mx.random.seed(args.seed)
 
     if group.size() > 1:
         if args.adapter_path:

mlx_lm/chat_templates/deepseek_v32.py

@@ -3,8 +3,11 @@
 import copy
 import json
 import re
+from inspect import isfunction
 from typing import Any, Dict, List, Optional, Tuple, Union
 
+from transformers.utils.chat_template_utils import get_json_schema
+
 TOOLS_SYSTEM_TEMPLATE = """## Tools
 
 You have access to a set of tools you can use to answer the user's question.
@@ -70,7 +73,12 @@ def to_json(value: Any) -> str:
 
 
 def tools_from_openai_format(tools):
-    return [tool["function"] for tool in tools]
+    def normalize_tool(tool):
+        if isfunction(tool):
+            return get_json_schema(tool)
+        return tool["function"]
+
+    return [normalize_tool(tool) for tool in tools]
 
 
 def tool_calls_from_openai_format(tool_calls):
@@ -141,7 +149,10 @@ def find_last_user_index(messages: List[Dict[str, Any]]) -> int:
 
 
 def render_message(
-    index: int, messages: List[Dict[str, Any]], thinking_mode: str
+    index: int,
+    messages: List[Dict[str, Any]],
+    thinking_mode: str,
+    tools: Any = None,
 ) -> str:
     assert 0 <= index < len(messages)
     assert thinking_mode in [
@@ -155,20 +166,18 @@ def render_message(
 
     role = msg.get("role")
     content = msg.get("content")
-    tools = msg.get("tools")
+    tools = tools or msg.get("tools")
     response_format = msg.get("response_format")
     tool_calls = msg.get("tool_calls")
     reasoning_content = msg.get("reasoning_content")
 
-    if tools:
-        tools = tools_from_openai_format(tools)
     if tool_calls:
         tool_calls = tool_calls_from_openai_format(tool_calls)
 
     if role == "system":
         prompt += system_msg_template.format(content=content or "")
         if tools:
-            prompt += "\n\n" + render_tools(tools)
+            prompt += "\n\n" + render_tools(tools_from_openai_format(tools))
 
         if response_format:
             prompt += "\n\n" + response_format_template.format(
@@ -179,7 +188,7 @@ def render_message(
         assert content, f"Invalid message for role `{role}`: {msg}"
         content_developer = ""
         if tools:
-            content_developer += "\n\n" + render_tools(tools)
+            content_developer += "\n\n" + render_tools(tools_from_openai_format(tools))
 
         if response_format:
             content_developer += "\n\n" + response_format_template.format(
@@ -301,6 +310,7 @@ def encode_messages(
     context: Optional[List[Dict[str, Any]]] = None,
     drop_thinking: bool = True,
     add_default_bos_token: bool = True,
+    tools: Any = None,
 ) -> str:
     context = context if context else []
     full_messages = context + messages
@@ -311,7 +321,10 @@ def encode_messages(
 
     for idx in range(len(messages)):
         prompt += render_message(
-            idx + len(context), full_messages, thinking_mode=thinking_mode
+            idx + len(context),
+            full_messages,
+            thinking_mode=thinking_mode,
+            tools=tools,
         )
 
     return prompt

mlx_lm/cli.py

@@ -0,0 +1,53 @@
+# Copyright © 2025 Apple Inc.
+
+import importlib
+import sys
+
+
+def main():
+    subcommands = (
+        "benchmark",
+        "cache_prompt",
+        "chat",
+        "convert",
+        "evaluate",
+        "fuse",
+        "generate",
+        "lora",
+        "manage",
+        "perplexity",
+        "awq",
+        "dwq",
+        "dynamic_quant",
+        "gptq",
+        "server",
+        "upload",
+        "share",
+    )
+    subpackages = {
+        "awq": "quant",
+        "dwq": "quant",
+        "dynamic_quant": "quant",
+        "gptq": "quant",
+    }
+    if len(sys.argv) < 2:
+        raise ValueError(f"CLI requires a subcommand in {subcommands}")
+    subcommand = sys.argv.pop(1)
+    if subcommand in subcommands:
+        if subpackage := subpackages.get(subcommand):
+            subcommand = f"{subpackage}.{subcommand}"
+        submodule = importlib.import_module(f"mlx_lm.{subcommand}")
+        submodule.main()
+    elif subcommand == "--version":
+        from mlx_lm import __version__
+
+        print(__version__)
+    elif subcommand in ("-h", "--help"):
+        print(f"The supported subcommands are {subcommands}")
+        print()
+        print(
+            "For help on an individual subcommand, pass --help "
+            "to the subcommand. For example: mlx_lm.generate --help"
+        )
+    else:
+        raise ValueError(f"CLI requires a subcommand in {subcommands}")

mlx_lm/convert.py

@@ -20,6 +20,7 @@ from .utils import (
 def mixed_quant_predicate_builder(
     recipe: str, model: nn.Module, group_size: int = 64
 ) -> Callable[[str, nn.Module, dict], Union[bool, dict]]:
+    mode = "affine"
     high_bits = 6
 
     if recipe == "mixed_2_6":
@@ -65,13 +66,13 @@ def mixed_quant_predicate_builder(
         if (
             "v_proj" in path or "v_a_proj" in path or "v_b_proj" in path
         ) and use_more_bits:
-            return {"group_size": group_size, "bits": high_bits}
+            return {"group_size": group_size, "bits": high_bits, "mode": mode}
         if "down_proj" in path and use_more_bits:
-            return {"group_size": group_size, "bits": high_bits}
+            return {"group_size": group_size, "bits": high_bits, "mode": mode}
         if "lm_head" in path:
-            return {"group_size": group_size, "bits": high_bits}
+            return {"group_size": group_size, "bits": high_bits, "mode": mode}
 
-        return {"group_size": group_size, "bits": low_bits}
+        return {"group_size": group_size, "bits": low_bits, "mode": mode}
 
     return mixed_quant_predicate
 
@@ -85,8 +86,8 @@ def convert(
     hf_path: str,
     mlx_path: str = "mlx_model",
     quantize: bool = False,
-    q_group_size: int = 64,
-    q_bits: int = 4,
+    q_group_size: Optional[int] = None,
+    q_bits: Optional[int] = None,
     q_mode: str = "affine",
     dtype: Optional[str] = None,
     upload_repo: str = None,
@@ -117,12 +118,18 @@ def convert(
     )
 
     if isinstance(quant_predicate, str):
+        if q_mode != "affine":
+            raise ValueError(f"Quant predicates only support 'affine' quantization.")
         quant_predicate = mixed_quant_predicate_builder(
-            quant_predicate, model, q_group_size
+            quant_predicate,
+            model,
+            q_group_size,
         )
 
     if dtype is None:
         dtype = config.get("torch_dtype", None)
+    if dtype is None and (text_config := config.get("text_config", None)):
+        dtype = text_config.get("dtype", None)
     if dtype in MODEL_CONVERSION_DTYPES:
         print("[INFO] Using dtype:", dtype)
         dtype = getattr(mx, dtype)

mlx_lm/evaluate.py

@@ -20,7 +20,6 @@ import mlx.nn as nn
 import numpy as np
 from lm_eval.api.model import LM
 from lm_eval.api.registry import register_model
-from lm_eval.models import huggingface
 from tqdm import tqdm
 
 from .generate import batch_generate
@@ -72,7 +71,6 @@ def chat_template_fn(**extra_kwargs):
 @register_model("mlxlm")
 class MLXLM(LM):
 
-    tokenizer_name = huggingface.HFLM.tokenizer_name
     apply_chat_template = chat_template_fn()
 
     def __init__(
@@ -147,6 +145,10 @@ class MLXLM(LM):
             for t in texts
         ]
 
+    @property
+    def tokenizer_name(self) -> str:
+        return self.tokenizer.name_or_path.replace("/", "__")
+
     def loglikelihood(self, requests) -> list[tuple[float, bool]]:
         """Compute log-likelihood of generating a continuation from a context.
         Downstream tasks should attempt to use loglikelihood instead of other

mlx_lm/generate.py

@@ -178,6 +178,12 @@ def setup_arg_parser():
         default=None,
         help="A file containing saved KV caches to avoid recomputing them",
     )
+    parser.add_argument(
+        "--quantize-activations",
+        "-qa",
+        action="store_true",
+        help="Quantize activations using the same quantization config as the corresponding layer.",
+    )
     parser.add_argument(
         "--kv-bits",
         type=int,
@@ -234,7 +240,7 @@ def wired_limit(model: nn.Module, streams: Optional[List[mx.Stream]] = None):
         model_bytes = tree_reduce(
             lambda acc, x: acc + x.nbytes if isinstance(x, mx.array) else acc, model, 0
         )
-        max_rec_size = mx.metal.device_info()["max_recommended_working_set_size"]
+        max_rec_size = mx.device_info()["max_recommended_working_set_size"]
         if model_bytes > 0.9 * max_rec_size:
             model_mb = model_bytes // 2**20
             max_rec_mb = max_rec_size // 2**20
@@ -877,7 +883,7 @@ class Batch:
         return [c.extract(idx) for c in self.cache]
 
 
-def _make_cache(model, left_padding):
+def _make_cache(model, left_padding, max_kv_size):
     """
     Convert a list of regular caches into their corresponding
     batch-aware caches.
@@ -902,6 +908,10 @@ def _make_cache(model, left_padding):
         cache = model.make_cache()
         return [to_batch_cache(c) for c in cache]
     else:
+        if max_kv_size is not None:
+            return [
+                BatchRotatingKVCache(max_kv_size, left_padding) for _ in model.layers
+            ]
         return [BatchKVCache(left_padding) for _ in model.layers]
 
 
@@ -917,6 +927,11 @@ def _merge_caches(caches):
     return batch_cache
 
 
+def _lazy_extract_cache(cache, i):
+    # Generators like lambdas are late bound so we can't just use it in the loop
+    return (c.extract(i) for c in cache)
+
+
 class BatchGenerator:
     @dataclass
     class Response:
@@ -938,9 +953,13 @@ class BatchGenerator:
         completion_batch_size: int = 32,
         prefill_batch_size: int = 8,
         prefill_step_size: int = 2048,
+        prompt_checkpoint_callback: Optional[
+            Callable[[List[Tuple[int, int, List[Any]]]], None]
+        ] = None,
         prompt_progress_callback: Optional[
             Callable[[List[Tuple[int, int, int]]], None]
         ] = None,
+        max_kv_size: Optional[int] = None,
     ):
         self.model = model
         self.unprocessed_prompts = []
@@ -952,14 +971,17 @@ class BatchGenerator:
         self.prefill_step_size = prefill_step_size
         self.prefill_batch_size = prefill_batch_size
         self.completion_batch_size = max(completion_batch_size, prefill_batch_size)
+        self.prompt_checkpoint_callback = prompt_checkpoint_callback
         self.prompt_progress_callback = prompt_progress_callback or (lambda *_: None)
         self._stats = BatchStats()
+        self._next_count = 0
+        self.max_kv_size = max_kv_size
 
         self.active_batch = None
 
         if mx.metal.is_available():
             self._old_wired_limit = mx.set_wired_limit(
-                mx.metal.device_info()["max_recommended_working_set_size"]
+                mx.device_info()["max_recommended_working_set_size"]
             )
         else:
             self._old_wired_limit = None
@@ -980,12 +1002,16 @@ class BatchGenerator:
         caches=None,
         samplers: list | None = None,
         logits_processors: list | None = None,
+        prompt_checkpoints: list | int | None = None,
     ):
         uids = []
 
         if max_tokens is None or isinstance(max_tokens, int):
             max_tokens = [max_tokens or self.max_tokens] * len(prompts)
 
+        if prompt_checkpoints is None or isinstance(prompt_checkpoints, int):
+            prompt_checkpoints = [prompt_checkpoints or -1] * len(prompts)
+
         if caches is None:
             caches = [None] * len(prompts)
         for i in range(len(prompts)):
@@ -995,10 +1021,10 @@ class BatchGenerator:
         samplers = samplers or [None] * len(prompts)
         logits_processors = logits_processors or [self.logits_processors] * len(prompts)
 
-        for p, m, c, s, lp in zip(
-            prompts, max_tokens, caches, samplers, logits_processors
+        for p, m, c, s, lp, pc in zip(
+            prompts, max_tokens, caches, samplers, logits_processors, prompt_checkpoints
         ):
-            self.unprocessed_prompts.append((self.uid_count, p, m, c, s, lp))
+            self.unprocessed_prompts.append((self.uid_count, p, m, c, s, lp, pc))
             uids.append(self.uid_count)
             self.uid_count += 1
         # Sort in ascending order of length
@@ -1008,10 +1034,16 @@ class BatchGenerator:
         )
         return uids
 
-    def remove(self, uids: List[int]):
+    def remove(self, uids: List[int], return_prompt_caches: bool = False):
+        caches = {}
         uids = set(uids)
         if self.active_batch is not None:
             batch = self.active_batch
+            if return_prompt_caches:
+                for e, uid in enumerate(batch.uids):
+                    if uid not in uids:
+                        continue
+                    caches[uid] = batch.extract_cache(e)
             keep_idx = [e for e, uid in enumerate(batch.uids) if uid not in uids]
             if len(keep_idx) > 0:
                 batch.filter(keep_idx)
@@ -1022,13 +1054,39 @@ class BatchGenerator:
             if self.unprocessed_prompts[i][0] in uids:
                 self.unprocessed_prompts.pop(i)
 
+        if return_prompt_caches:
+            return caches
+
+    @property
+    def prompt_cache_nbytes(self):
+        total = sum(c.nbytes for p in self.unprocessed_prompts for c in p[3])
+        if self.active_batch is not None:
+            total += sum(c.nbytes for c in self.active_batch.cache)
+        return total
+
     def _process_prompts(self, prompts):
-        uids, inputs, max_tokens, caches, samplers, logits_processors = zip(*prompts)
+        (
+            uids,
+            inputs,
+            max_tokens,
+            caches,
+            samplers,
+            logits_processors,
+            prompt_checkpoints,
+        ) = zip(*prompts)
 
         lengths = [len(p) for p in inputs]
         max_length = max(lengths)
         padding = [max_length - l for l in lengths]
 
+        # Get the checkpoint token as an offset from the end of each prompt.
+        # Then select the largest one so that we perform the checkpoint at
+        # least `pc` before the end.
+        prompt_checkpoints = [
+            (l - pc if pc > 0 else -pc) for l, pc in zip(lengths, prompt_checkpoints)
+        ]
+        prompt_checkpoint = max(1, max(prompt_checkpoints))
+
         self._stats.prompt_tokens += sum(lengths)
 
         tokens = [mx.array(inp) for inp in inputs]
@@ -1039,10 +1097,12 @@ class BatchGenerator:
         #   2. Process
         if all(c[0].empty() for c in caches):
             inputs = _left_pad_prompts(inputs, max_length=max_length)
-            prompt_cache = _make_cache(self.model, padding)
+            prompt_cache = _make_cache(self.model, padding, self.max_kv_size)
 
-            while inputs.shape[1] > 1:
-                n_to_process = min(self.prefill_step_size, inputs.shape[1] - 1)
+            while inputs.shape[1] > prompt_checkpoint:
+                n_to_process = min(
+                    self.prefill_step_size, inputs.shape[1] - prompt_checkpoint
+                )
                 self.model(inputs[:, :n_to_process], cache=prompt_cache)
                 mx.eval([c.state for c in prompt_cache])
                 inputs = inputs[:, n_to_process:]
@@ -1053,6 +1113,7 @@ class BatchGenerator:
                         for uid, length in zip(uids, lengths)
                     ]
                 )
+                mx.clear_cache()
 
         # Further prompt processing so we need to
         #   1. Merge the KV caches and prepare for right padded prompts
@@ -1060,16 +1121,22 @@ class BatchGenerator:
         #   2. Process
         #   3. Finalize the KV caches so they are left padded again
         else:
-            last_inputs = mx.array([p[-1:] for p in inputs])
+            last_inputs = mx.array([p[-prompt_checkpoint:] for p in inputs])
             inputs = _right_pad_prompts(inputs, max_length=max_length)
             prompt_cache = _merge_caches(caches)
 
             for c in prompt_cache:
-                # subtract one from lengths since we don't process the last token during prefill
-                c.prepare(lengths=[l - 1 for l in lengths], right_padding=padding)
+                # subtract from lengths since we don't process the last
+                # `prompt_checkpoint` tokens during prefill
+                c.prepare(
+                    lengths=[l - prompt_checkpoint for l in lengths],
+                    right_padding=padding,
+                )
 
-            while inputs.shape[1] > 1:
-                n_to_process = min(self.prefill_step_size, inputs.shape[1] - 1)
+            while inputs.shape[1] > prompt_checkpoint:
+                n_to_process = min(
+                    self.prefill_step_size, inputs.shape[1] - prompt_checkpoint
+                )
                 self.model(inputs[:, :n_to_process], cache=prompt_cache)
                 mx.eval([c.state for c in prompt_cache])
                 inputs = inputs[:, n_to_process:]
@@ -1087,6 +1154,20 @@ class BatchGenerator:
 
         for c in prompt_cache:
             c.finalize()
+
+        # We processed L - prompt_checkpoint tokens so call the checkpoint
+        # callback.
+        if self.prompt_checkpoint_callback is not None:
+            self.prompt_checkpoint_callback(
+                [
+                    (uid, prompt_checkpoint, _lazy_extract_cache(prompt_cache, i))
+                    for i, uid in enumerate(uids)
+                ]
+            )
+        # Process the remaining prompt_checkpoint-1 tokens
+        if prompt_checkpoint > 1:
+            self.model(inputs[:, : prompt_checkpoint - 1], cache=prompt_cache)
+            mx.eval([c.state for c in prompt_cache])
         mx.clear_cache()
 
         y, logprobs = self._step(
@@ -1199,7 +1280,7 @@ class BatchGenerator:
             batch.tokens,
         )
 
-        mx.async_eval(batch.y, batch.logprobs)
+        mx.async_eval(batch.y, batch.logprobs, batch.tokens)
 
         y = y.tolist()
         toc = time.perf_counter()
@@ -1237,6 +1318,9 @@ class BatchGenerator:
             else:
                 self.active_batch = None
 
+        self._next_count += 1
+        if self._next_count % 512 == 0:
+            mx.clear_cache()
         self._stats.generation_tokens += len(responses)
         return responses
 
@@ -1370,6 +1454,7 @@ def main():
         model_path,
         adapter_path=args.adapter_path,
         tokenizer_config=tokenizer_config,
+        model_config={"quantize_activations": args.quantize_activations},
     )
     for eos_token in args.extra_eos_token:
         tokenizer.add_eos_token(eos_token)

mlx_lm/lora.py

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

mlx_lm/models/activations.py

@@ -9,3 +9,35 @@ import mlx.nn as nn
 @partial(mx.compile, shapeless=True)
 def swiglu(gate, x):
     return nn.silu(gate) * x
+
+
+@partial(mx.compile, shapeless=True)
+def xielu(x, alpha_p, alpha_n, beta, eps):
+    alpha_p = nn.softplus(alpha_p)
+    alpha_n = beta + nn.softplus(alpha_n)
+    return mx.where(
+        x > 0,
+        alpha_p * mx.square(x) + beta * x,
+        (mx.expm1(mx.minimum(x, eps)) - x) * alpha_n + beta * x,
+    )
+
+
+class XieLU(nn.Module):
+    def __init__(
+        self,
+        alpha_p_init=0.8,
+        alpha_n_init=0.8,
+        beta=0.5,
+        eps=-1e-6,
+    ):
+        super().__init__()
+        alpha_p_tensor = mx.array(alpha_p_init)
+        alpha_n_tensor = mx.array(alpha_n_init - beta)
+        self.alpha_p = mx.log(mx.exp(alpha_p_tensor) - 1)
+        self.alpha_n = mx.log(mx.exp(alpha_n_tensor) - 1)
+
+        self.beta = mx.array(beta)
+        self.eps = mx.array(eps)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return xielu(x, self.alpha_p, self.alpha_n, self.beta, self.eps)

mlx_lm/models/apertus.py

@@ -7,6 +7,7 @@ from typing import Any, Dict, Optional, Union
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import XieLU
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .rope_utils import initialize_rope
 
@@ -32,38 +33,6 @@ class ModelArgs(BaseModelArgs):
     rope_scaling: Optional[Dict[str, Union[float, str]]] = None
 
 
-@partial(mx.compile, shapeless=True)
-def xielu(x, alpha_p, alpha_n, beta, eps):
-    alpha_p = nn.softplus(alpha_p)
-    alpha_n = beta + nn.softplus(alpha_n)
-    return mx.where(
-        x > 0,
-        alpha_p * mx.square(x) + beta * x,
-        (mx.expm1(mx.minimum(x, eps)) - x) * alpha_n + beta * x,
-    )
-
-
-class XieLU(nn.Module):
-    def __init__(
-        self,
-        alpha_p_init=0.8,
-        alpha_n_init=0.8,
-        beta=0.5,
-        eps=-1e-6,
-    ):
-        super().__init__()
-        alpha_p_tensor = mx.array(alpha_p_init)
-        alpha_n_tensor = mx.array(alpha_n_init - beta)
-        self.alpha_p = mx.log(mx.exp(alpha_p_tensor) - 1)
-        self.alpha_n = mx.log(mx.exp(alpha_n_tensor) - 1)
-
-        self.beta = mx.array(beta)
-        self.eps = mx.array(eps)
-
-    def __call__(self, x: mx.array) -> mx.array:
-        return xielu(x, self.alpha_p, self.alpha_n, self.beta, self.eps)
-
-
 class ApertusMLP(nn.Module):
     def __init__(self, args: ModelArgs):
         super().__init__()

mlx_lm/models/baichuan_m1.py

@@ -8,7 +8,7 @@ import mlx.nn as nn
 
 from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
-from .cache import CacheList, KVCache, MambaCache, RotatingKVCache
+from .cache import ArraysCache, CacheList, KVCache, RotatingKVCache
 
 
 @dataclass
@@ -223,7 +223,7 @@ class Model(nn.Module):
         caches = []
         for i, layer in enumerate(self.model.layers):
             is_swa = i in self.config.sliding_window_layers
-            conv_cache = MambaCache()
+            conv_cache = ArraysCache(size=2)
             if is_swa:
                 kv_cache = RotatingKVCache(max_size=self.config.sliding_window)
             else:

mlx_lm/models/cache.py

@@ -112,9 +112,11 @@ def trim_prompt_cache(cache: List[Any], num_tokens: int) -> List[Any]:
 def create_attention_mask(
     N: int, offset: int, return_array: bool, window_size: Optional[int]
 ):
-    if N == 1:
+    if window_size is not None:
+        return create_causal_mask(N, offset, window_size=window_size)
+    elif N == 1:
         return None
-    if return_array:
+    elif return_array:
         return create_causal_mask(N, offset, window_size=window_size)
     else:
         return "causal"
@@ -151,6 +153,11 @@ class _BaseCache:
         """
         return 0
 
+    @property
+    def nbytes(self):
+        """Return the size of this cache in bytes"""
+        raise NotImplementedError("Cache sub-class must implement nbytes")
+
     def empty(self):
         """
         Return if the cache is empty or not.
@@ -213,6 +220,12 @@ class ConcatenateKVCache(_BaseCache):
     def empty(self):
         return self.keys is None
 
+    @property
+    def nbytes(self):
+        if self.keys is None:
+            return 0
+        return self.keys.nbytes + self.values.nbytes
+
 
 class QuantizedKVCache(_BaseCache):
     step = 256
@@ -302,6 +315,10 @@ class QuantizedKVCache(_BaseCache):
     def empty(self):
         return self.keys is None
 
+    @property
+    def nbytes(self):
+        return tree_reduce(lambda a, x: a + x.nbytes, (self.keys, self.values), 0)
+
 
 class KVCache(_BaseCache):
     step = 256
@@ -381,6 +398,12 @@ class KVCache(_BaseCache):
     def empty(self):
         return self.keys is None
 
+    @property
+    def nbytes(self):
+        if self.keys is None:
+            return 0
+        return self.keys.nbytes + self.values.nbytes
+
 
 class RotatingKVCache(_BaseCache):
     step = 256
@@ -559,12 +582,24 @@ class RotatingKVCache(_BaseCache):
     def empty(self):
         return self.keys is None
 
+    @property
+    def nbytes(self):
+        if self.keys is None:
+            return 0
+        return self.keys.nbytes + self.values.nbytes
+
 
 class ArraysCache(_BaseCache):
+    def __new__(cls, *args, **kwargs):
+        instance = super().__new__(cls)
+        instance.left_padding = None
+        instance.lengths = None
+        return instance
+
     def __init__(self, size, left_padding: Optional[List[int]] = None):
         self.cache = [None] * size
-        self.left_padding = mx.array(left_padding) if left_padding else None
-        self.lengths = None
+        if left_padding:
+            self.left_padding = mx.array(left_padding)
 
     def __setitem__(self, idx, value):
         self.cache[idx] = value
@@ -639,10 +674,9 @@ class ArraysCache(_BaseCache):
     def empty(self):
         return self.cache[0] is None
 
-
-class MambaCache(ArraysCache):
-    def __init__(self, left_padding: Optional[List[int]] = None):
-        super().__init__(size=2, left_padding=left_padding)
+    @property
+    def nbytes(self):
+        return sum(c.nbytes for c in self.cache if c is not None)
 
 
 class ChunkedKVCache(_BaseCache):
@@ -721,6 +755,12 @@ class ChunkedKVCache(_BaseCache):
     def empty(self):
         return self.keys is None
 
+    @property
+    def nbytes(self):
+        if self.keys is None:
+            return 0
+        return self.keys.nbytes + self.values.nbytes
+
 
 class CacheList(_BaseCache):
     def __init__(self, *caches):
@@ -739,16 +779,24 @@ class CacheList(_BaseCache):
 
     @property
     def state(self):
-        return [s for c in self.caches for s in c.state]
+        return [c.state for c in self.caches]
 
     @state.setter
     def state(self, v):
-        state_lens = [len(c.state) for c in self.caches]
-        start = 0
-        for c in self.caches:
-            l = len(c.state)
-            c.state = v[start : start + l]
-            start += l
+        for c, s in zip(self.caches, v):
+            c.state = s
+
+    @property
+    def meta_state(self):
+        return (
+            [type(c).__name__ for c in self.caches],
+            [c.meta_state for c in self.caches],
+        )
+
+    @meta_state.setter
+    def meta_state(self, v):
+        for c, m in zip(self.caches, v[1]):
+            c.meta_state = m
 
     def filter(self, batch_indices):
         """
@@ -790,6 +838,18 @@ class CacheList(_BaseCache):
     def empty(self):
         return self.caches[0].empty()
 
+    @property
+    def nbytes(self):
+        return sum(c.nbytes for c in self.caches)
+
+    @classmethod
+    def from_state(cls, state, meta_state):
+        obj = cls.__new__(cls)
+        obj.caches = [
+            globals()[c].from_state(s, m) for s, c, m in zip(state, *meta_state)
+        ]
+        return obj
+
 
 def dynamic_roll(x, shifts, axis):
     n = x.shape[axis]
@@ -988,6 +1048,12 @@ class BatchKVCache(_BaseCache):
     def empty(self):
         return self.keys is None
 
+    @property
+    def nbytes(self):
+        if self.keys is None:
+            return 0
+        return self.keys.nbytes + self.values.nbytes
+
 
 class BatchRotatingKVCache(_BaseCache):
     step = 256
@@ -1058,6 +1124,10 @@ class BatchRotatingKVCache(_BaseCache):
         self.offset += keys.shape[2]
         self._offset += keys.shape[2]
         self._idx = self.keys.shape[2]
+
+        # Make sure left_padding and offset are evaluated
+        self.keys = mx.depends(self.keys, (self.left_padding, self.offset))
+
         return self.keys, self.values
 
     def _update_in_place(self, keys, values):
@@ -1108,6 +1178,9 @@ class BatchRotatingKVCache(_BaseCache):
         self.offset += S
         self._idx += S
 
+        # Make sure left_padding and offset are evaluated
+        self.keys = mx.depends(self.keys, (self.left_padding, self.offset))
+
         # If the buffer is not full, slice off the end
         if self._offset < self.max_size:
             return (
@@ -1302,3 +1375,9 @@ class BatchRotatingKVCache(_BaseCache):
 
     def empty(self):
         return self.keys is None
+
+    @property
+    def nbytes(self):
+        if self.keys is None:
+            return 0
+        return self.keys.nbytes + self.values.nbytes

mlx_lm/models/deepseek_v3.py

@@ -11,6 +11,7 @@ from mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients
 
 from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .mla import MultiLinear
 from .pipeline import PipelineMixin
 from .rope_utils import initialize_rope
 from .switch_layers import SwitchGLU
@@ -85,11 +86,11 @@ class DeepseekV3Attention(nn.Module):
             bias=config.attention_bias,
         )
         self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank, eps=1e-6)
-        self.kv_b_proj = nn.Linear(
-            self.kv_lora_rank,
-            self.num_heads
-            * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),
-            bias=False,
+        self.embed_q = MultiLinear(
+            self.qk_nope_head_dim, self.kv_lora_rank, self.num_heads
+        )
+        self.unembed_out = MultiLinear(
+            self.kv_lora_rank, self.v_head_dim, self.num_heads
         )
 
         self.o_proj = nn.Linear(
@@ -132,29 +133,38 @@ class DeepseekV3Attention(nn.Module):
         compressed_kv = self.kv_a_proj_with_mqa(x)
         compressed_kv, k_pe = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1)
         k_pe = k_pe.reshape(B, L, 1, self.qk_rope_head_dim).transpose(0, 2, 1, 3)
-        kv = self.kv_b_proj(self.kv_a_layernorm(compressed_kv))
-        kv = kv.reshape(B, L, self.num_heads, -1).transpose(0, 2, 1, 3)
+        kv_latent = self.kv_a_layernorm(compressed_kv)
 
-        k_nope, values = mx.split(kv, [self.qk_nope_head_dim], axis=-1)
+        offset = cache.offset if cache is not None else 0
+        q_pe = self.rope(q_pe, offset)
+        k_pe = self.rope(k_pe, offset)
+
+        kv_latent = mx.expand_dims(kv_latent, axis=1)
 
         if cache is not None:
-            q_pe = self.rope(q_pe, cache.offset)
-            k_pe = self.rope(k_pe, cache.offset)
-            k_pe = mx.repeat(k_pe, self.num_heads, axis=1)
-            keys, values = cache.update_and_fetch(
-                mx.concatenate([k_nope, k_pe], axis=-1), values
-            )
-        else:
-            q_pe = self.rope(q_pe)
-            k_pe = self.rope(k_pe)
-            k_pe = mx.repeat(k_pe, self.num_heads, axis=1)
-            keys = mx.concatenate([k_nope, k_pe], axis=-1)
+            kv_latent, k_pe = cache.update_and_fetch(kv_latent, k_pe)
 
-        queries = mx.concatenate([q_nope, q_pe], axis=-1)
+        pe_scores = (q_pe * self.scale) @ k_pe.swapaxes(-1, -2)
+        if mask is not None:
+            pe_scores = mx.where(
+                mask,
+                pe_scores,
+                mx.array(mx.finfo(pe_scores.dtype).min, pe_scores.dtype),
+            )
+
+        if L == 1:
+            q_nope = self.embed_q(q_nope)
+            k = v = kv_latent
+        else:
+            k = self.embed_q(kv_latent, transpose=False)
+            v = self.unembed_out(kv_latent)
 
         output = scaled_dot_product_attention(
-            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+            q_nope, k, v, cache=cache, scale=self.scale, mask=pe_scores
         )
+        if L == 1:
+            output = self.unembed_out(output)
+
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
         return self.o_proj(output)
 
@@ -329,7 +339,7 @@ class DeepseekV3Model(PipelineMixin, nn.Module):
 
         if cache is None:
             cache = [None] * len(self.pipeline_layers)
-        mask = create_attention_mask(h, cache[0])
+        mask = create_attention_mask(h, cache[0], return_array=True)
 
         # Receive from the previous process in the pipeline
         if pipeline_rank < pipeline_size - 1:
@@ -423,6 +433,42 @@ class Model(nn.Module):
                             for e in range(self.args.n_routed_experts)
                         ]
                         weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
+            prefix = f"model.layers.{l}.self_attn"
+            if f"{prefix}.kv_b_proj.weight" in weights:
+                layer = self.model.layers[l].self_attn.embed_q
+                quantized = f"{prefix}.kv_b_proj.scales" in weights
+                v = weights.pop(f"{prefix}.kv_b_proj.weight")
+                head_dim = self.args.qk_nope_head_dim + self.args.v_head_dim
+
+                if quantized:
+                    dims = self.args.kv_lora_rank
+                    scales = weights.pop(f"{prefix}.kv_b_proj.scales")
+                    biases = weights.pop(f"{prefix}.kv_b_proj.biases")
+                    # Try to infer bits and group size
+                    bits = (v.shape[-1] * 32) // dims
+                    group_size = dims // scales.shape[-1]
+                    v = mx.dequantize(
+                        v, scales, biases, bits=bits, group_size=group_size
+                    )
+                num_heads = self.args.num_attention_heads
+                v = v.reshape(num_heads, head_dim, -1)
+                wk = mx.contiguous(
+                    v[:, : self.args.qk_nope_head_dim, :].swapaxes(-1, -2)
+                )
+                wv = mx.contiguous(v[:, self.args.qk_nope_head_dim :, :])
+                if quantized:
+                    wk, wk_scales, wk_biases = mx.quantize(
+                        wk, bits=bits, group_size=group_size
+                    )
+                    wv, wv_scales, wv_biases = mx.quantize(
+                        wv, bits=bits, group_size=group_size
+                    )
+                    weights[f"{prefix}.embed_q.scales"] = wk_scales
+                    weights[f"{prefix}.unembed_out.scales"] = wv_scales
+                    weights[f"{prefix}.embed_q.biases"] = wk_biases
+                    weights[f"{prefix}.unembed_out.biases"] = wv_biases
+                weights[f"{prefix}.embed_q.weight"] = wk
+                weights[f"{prefix}.unembed_out.weight"] = wv
 
         # Remove multi-token prediction layer and any unused precomputed rotary freqs
         return {
@@ -434,6 +480,7 @@ class Model(nn.Module):
     def shard(self, group: Optional[mx.distributed.Group] = None):
         group = group or mx.distributed.init()
         N = group.size()
+        rank = group.rank()
         for layer in self.model.layers:
             # Shard the self attention
             if layer.self_attn.q_lora_rank is None:
@@ -444,13 +491,20 @@ class Model(nn.Module):
                 layer.self_attn.q_b_proj = shard_linear(
                     layer.self_attn.q_b_proj, "all-to-sharded", group=group
                 )
-            layer.self_attn.kv_b_proj = shard_linear(
-                layer.self_attn.kv_b_proj, "all-to-sharded", group=group
-            )
+            layer.self_attn.num_heads //= N
+            num_heads = layer.self_attn.num_heads
+            sh = rank * num_heads
+            eh = sh + num_heads
+
+            def shard_heads(w):
+                return w[sh:eh]
+
+            layer.self_attn.embed_q.apply(shard_heads)
+            layer.self_attn.unembed_out.apply(shard_heads)
+
             layer.self_attn.o_proj = shard_linear(
                 layer.self_attn.o_proj, "sharded-to-all", group=group
             )
-            layer.self_attn.num_heads //= N
 
             # Shard the MLP
             if isinstance(layer.mlp, DeepseekV3MLP):

mlx_lm/models/deepseek_v32.py

@@ -11,6 +11,7 @@ from mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients
 from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .cache import CacheList, KVCache
+from .mla import MultiLinear
 from .rope_utils import initialize_rope
 from .switch_layers import SwitchGLU
 
@@ -70,7 +71,7 @@ class Indexer(nn.Module):
         self.rope = initialize_rope(
             dims=args.qk_rope_head_dim,
             base=args.rope_theta,
-            traditional=False,
+            traditional=True,
             max_position_embeddings=args.max_position_embeddings,
             scaling_config=args.rope_scaling,
         )
@@ -108,7 +109,7 @@ class Indexer(nn.Module):
         weights = self.weights_proj(x) * (self.n_heads**-0.5 * self.softmax_scale)
         weights = weights.swapaxes(-1, -2)[..., None]
         scores = scores * weights
-        scores = scores.sum(axis=1)
+        scores = scores.sum(axis=1, keepdims=True)
         if mask is not None:
             scores = mx.where(mask, scores, -float("inf"))
         return mx.argpartition(scores, kth=-self.index_topk, axis=-1)[
@@ -147,11 +148,11 @@ class DeepseekV32Attention(nn.Module):
             bias=config.attention_bias,
         )
         self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank, eps=1e-6)
-        self.kv_b_proj = nn.Linear(
-            self.kv_lora_rank,
-            self.num_heads
-            * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),
-            bias=False,
+        self.embed_q = MultiLinear(
+            self.qk_nope_head_dim, self.kv_lora_rank, self.num_heads
+        )
+        self.unembed_out = MultiLinear(
+            self.kv_lora_rank, self.v_head_dim, self.num_heads
         )
 
         self.o_proj = nn.Linear(
@@ -193,45 +194,70 @@ class DeepseekV32Attention(nn.Module):
         compressed_kv = self.kv_a_proj_with_mqa(x)
         compressed_kv, k_pe = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1)
         k_pe = k_pe.reshape(B, L, 1, self.qk_rope_head_dim).transpose(0, 2, 1, 3)
-        kv = self.kv_b_proj(self.kv_a_layernorm(compressed_kv))
-        kv = kv.reshape(B, L, self.num_heads, -1).transpose(0, 2, 1, 3)
+        kv_latent = self.kv_a_layernorm(compressed_kv)
 
-        k_nope, values = mx.split(kv, [self.qk_nope_head_dim], axis=-1)
+        offset = cache[0].offset if cache is not None else 0
+        q_pe = self.rope(q_pe, offset)
+        k_pe = self.rope(k_pe, offset)
+
+        kv_latent = mx.expand_dims(kv_latent, axis=1)
 
         if cache is not None:
-            q_pe = self.rope(q_pe, cache[0].offset)
-            k_pe = self.rope(k_pe, cache[0].offset)
-            k_pe = mx.repeat(k_pe, self.num_heads, axis=1)
-            keys, values = cache[0].update_and_fetch(
-                mx.concatenate([k_nope, k_pe], axis=-1), values
-            )
+            kv_latent, k_pe = cache[0].update_and_fetch(kv_latent, k_pe)
         else:
             cache = [None] * 2
-            q_pe = self.rope(q_pe)
-            k_pe = self.rope(k_pe)
-            k_pe = mx.repeat(k_pe, self.num_heads, axis=1)
-            keys = mx.concatenate([k_nope, k_pe], axis=-1)
 
-        queries = mx.concatenate([q_nope, q_pe], axis=-1)
         topk_indices = self.indexer(x, qr, mask, cache=cache[1])
         if topk_indices is not None:
-            k_seq = keys.shape[2]
-            sparse_mask = mx.zeros((B, L, k_seq), dtype=mx.bool_)
-            sparse_mask = mx.put_along_axis(
-                sparse_mask, topk_indices, mx.array(True), axis=-1
-            )
-            sparse_mask = sparse_mask[:, None, :, :]
-            if mask is not None:
-                sparse_mask = sparse_mask & mask
-            mask = sparse_mask
+            if L == 1:
+                idx = topk_indices[:, :, 0, :, None]
+                kv_latent = mx.take_along_axis(
+                    kv_latent,
+                    mx.broadcast_to(idx, idx.shape[:-1] + (kv_latent.shape[-1],)),
+                    axis=2,
+                )
+                k_pe = mx.take_along_axis(
+                    k_pe,
+                    mx.broadcast_to(idx, idx.shape[:-1] + (k_pe.shape[-1],)),
+                    axis=2,
+                )
+                mask = None
+            else:
+                shape = list(topk_indices.shape)
+                shape[-1] = kv_latent.shape[2]
+                sparse_mask = mx.zeros(shape, dtype=mx.bool_)
+                sparse_mask = mx.put_along_axis(
+                    sparse_mask, topk_indices, mx.array(True), axis=-1
+                )
+                if mask is not None:
+                    sparse_mask = sparse_mask & mask
+                mask = sparse_mask
         # Ensure the indexer cache is evaluated even if the topk_indices are unused
         # to keep the graph from getting too large
         if cache is not None and cache[0] is not None:
             cache[0].keys = mx.depends(cache[0].keys, (cache[1].keys, cache[1].values))
 
+        pe_scores = (q_pe * self.scale) @ k_pe.swapaxes(-1, -2)
+        if mask is not None:
+            pe_scores = mx.where(
+                mask,
+                pe_scores,
+                mx.array(mx.finfo(pe_scores.dtype).min, pe_scores.dtype),
+            )
+
+        if L == 1:
+            q_nope = self.embed_q(q_nope)
+            k = v = kv_latent
+        else:
+            k = self.embed_q(kv_latent, transpose=False)
+            v = self.unembed_out(kv_latent)
+
         output = scaled_dot_product_attention(
-            queries, keys, values, cache=cache[0], scale=self.scale, mask=mask
+            q_nope, k, v, cache=cache, scale=self.scale, mask=pe_scores
         )
+        if L == 1:
+            output = self.unembed_out(output)
+
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
         return self.o_proj(output)
 
@@ -469,6 +495,16 @@ class Model(nn.Module):
         return self.lm_head(out)
 
     def sanitize(self, weights):
+        # Remove multi-token prediction layers
+        mpt_layer = self.args.num_hidden_layers
+        new_weights = {}
+        for k, v in weights.items():
+            parts = k.split(".")
+            if len(parts) >= 3 and parts[1] == "layers" and int(parts[2]) >= mpt_layer:
+                continue
+            new_weights[k] = v
+        weights = new_weights
+
         def dequant(weight, scale_inv):
             dtype = mx.bfloat16
             weight = mx.from_fp8(weight, dtype=mx.bfloat16)
@@ -509,28 +545,67 @@ class Model(nn.Module):
                             for e in range(self.args.n_routed_experts)
                         ]
                         weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
+            prefix = f"model.layers.{l}.self_attn"
+            if f"{prefix}.kv_b_proj.weight" in weights:
+                layer = self.model.layers[l].self_attn.embed_q
+                quantized = f"{prefix}.kv_b_proj.scales" in weights
+                v = weights.pop(f"{prefix}.kv_b_proj.weight")
+                head_dim = self.args.qk_nope_head_dim + self.args.v_head_dim
 
-        # Remove multi-token prediction layer and any unused precomputed rotary freqs
-        return {
-            k: v
-            for k, v in weights.items()
-            if not k.startswith("model.layers.61") and "rotary_emb.inv_freq" not in k
-        }
+                if quantized:
+                    dims = self.args.kv_lora_rank
+                    scales = weights.pop(f"{prefix}.kv_b_proj.scales")
+                    biases = weights.pop(f"{prefix}.kv_b_proj.biases")
+                    # Try to infer bits and group size
+                    bits = (v.shape[-1] * 32) // dims
+                    group_size = dims // scales.shape[-1]
+                    v = mx.dequantize(
+                        v, scales, biases, bits=bits, group_size=group_size
+                    )
+                num_heads = self.args.num_attention_heads
+                v = v.reshape(num_heads, head_dim, -1)
+                wk = mx.contiguous(
+                    v[:, : self.args.qk_nope_head_dim, :].swapaxes(-1, -2)
+                )
+                wv = mx.contiguous(v[:, self.args.qk_nope_head_dim :, :])
+                if quantized:
+                    wk, wk_scales, wk_biases = mx.quantize(
+                        wk, bits=bits, group_size=group_size
+                    )
+                    wv, wv_scales, wv_biases = mx.quantize(
+                        wv, bits=bits, group_size=group_size
+                    )
+                    weights[f"{prefix}.embed_q.scales"] = wk_scales
+                    weights[f"{prefix}.unembed_out.scales"] = wv_scales
+                    weights[f"{prefix}.embed_q.biases"] = wk_biases
+                    weights[f"{prefix}.unembed_out.biases"] = wv_biases
+                weights[f"{prefix}.embed_q.weight"] = wk
+                weights[f"{prefix}.unembed_out.weight"] = wv
+
+        return weights
 
     def shard(self, group: Optional[mx.distributed.Group] = None):
         group = group or mx.distributed.init()
         N = group.size()
+        rank = group.rank()
         for layer in self.model.layers:
             layer.self_attn.q_b_proj = shard_linear(
                 layer.self_attn.q_b_proj, "all-to-sharded", group=group
             )
-            layer.self_attn.kv_b_proj = shard_linear(
-                layer.self_attn.kv_b_proj, "all-to-sharded", group=group
-            )
+
             layer.self_attn.o_proj = shard_linear(
                 layer.self_attn.o_proj, "sharded-to-all", group=group
             )
             layer.self_attn.num_heads //= N
+            num_heads = layer.self_attn.num_heads
+            sh = rank * num_heads
+            eh = sh + num_heads
+
+            def shard_heads(w):
+                return w[sh:eh]
+
+            layer.self_attn.embed_q.apply(shard_heads)
+            layer.self_attn.unembed_out.apply(shard_heads)
 
             # Shard the MLP
             if isinstance(layer.mlp, DeepseekV32MLP):

mlx_lm/models/falcon_h1.py

@@ -13,7 +13,7 @@ from .base import (
     create_ssm_mask,
     scaled_dot_product_attention,
 )
-from .cache import CacheList, KVCache, MambaCache
+from .cache import ArraysCache, CacheList, KVCache
 from .rope_utils import initialize_rope
 from .ssm import ssm_update
 
@@ -236,7 +236,7 @@ class FalconH1Mixer(nn.Module):
     def _conv(
         self,
         conv_input: mx.array,
-        cache: Optional[MambaCache],
+        cache: Optional[ArraysCache],
         mask: Optional[mx.array],
     ) -> mx.array:
         if mask is not None:
@@ -273,7 +273,7 @@ class FalconH1Mixer(nn.Module):
         B: mx.array,
         C: mx.array,
         dt: mx.array,
-        cache: Optional[MambaCache],
+        cache: Optional[ArraysCache],
         mask: Optional[mx.array],
     ) -> mx.array:
         batch_size, seq_len, _ = hidden_states.shape
@@ -495,7 +495,7 @@ class Model(nn.Module):
 
     def make_cache(self):
         return [
-            CacheList(MambaCache(), KVCache())
+            CacheList(ArraysCache(size=2), KVCache())
             for _ in range(self.args.num_hidden_layers)
         ]
 

mlx_lm/models/gated_delta.py

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

mlx_lm/models/glm4_moe_lite.py

@@ -10,6 +10,7 @@ from mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients
 
 from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .mla import MultiLinear
 from .pipeline import PipelineMixin
 from .rope_utils import initialize_rope
 from .switch_layers import SwitchGLU
@@ -45,12 +46,12 @@ class ModelArgs(BaseModelArgs):
     rms_norm_eps: float = 1e-5
     rope_theta: float = 1_000_000.0
     rope_scaling: Optional[Dict] = None
-    rope_traditional: bool = True
     attention_bias: bool = False
     attention_dropout: float = 0.0
     partial_rotary_factor: float = 1.0
     tie_word_embeddings: bool = False
     num_nextn_predict_layers: int = 1
+    quantization: Optional[Dict[str, Any]] = None
 
 
 class Glm4MoeLiteAttention(nn.Module):
@@ -60,6 +61,7 @@ class Glm4MoeLiteAttention(nn.Module):
         self.hidden_size = config.hidden_size
         self.num_heads = config.num_attention_heads
         self.max_position_embeddings = config.max_position_embeddings
+        rope_params = config.rope_scaling
         self.rope_theta = config.rope_theta
         self.q_lora_rank = config.q_lora_rank
         self.qk_rope_head_dim = config.qk_rope_head_dim
@@ -89,11 +91,12 @@ class Glm4MoeLiteAttention(nn.Module):
             bias=config.attention_bias,
         )
         self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)
-        self.kv_b_proj = nn.Linear(
-            self.kv_lora_rank,
-            self.num_heads
-            * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),
-            bias=False,
+        head_dim = self.qk_nope_head_dim + self.v_head_dim
+        self.embed_q = MultiLinear(
+            self.qk_nope_head_dim, self.kv_lora_rank, self.num_heads
+        )
+        self.unembed_out = MultiLinear(
+            self.kv_lora_rank, self.v_head_dim, self.num_heads
         )
 
         self.o_proj = nn.Linear(
@@ -102,10 +105,10 @@ class Glm4MoeLiteAttention(nn.Module):
             bias=config.attention_bias,
         )
 
-        if self.config.rope_scaling is not None:
-            mscale_all_dim = self.config.rope_scaling.get("mscale_all_dim", 0)
+        if rope_params is not None:
+            mscale_all_dim = rope_params.get("mscale_all_dim", 0)
             if mscale_all_dim:
-                scaling_factor = self.config.rope_scaling["factor"]
+                scaling_factor = rope_params["factor"]
                 if scaling_factor > 1:
                     s = 0.1 * mscale_all_dim * math.log(scaling_factor) + 1.0
                     self.scale = self.scale * s * s
@@ -113,9 +116,9 @@ class Glm4MoeLiteAttention(nn.Module):
         self.rope = initialize_rope(
             dims=self.qk_rope_head_dim,
             base=self.rope_theta,
-            traditional=self.config.rope_traditional,
+            traditional=True,
             max_position_embeddings=self.max_position_embeddings,
-            scaling_config=self.config.rope_scaling,
+            scaling_config=rope_params,
         )
 
     def __call__(
@@ -136,29 +139,37 @@ class Glm4MoeLiteAttention(nn.Module):
         compressed_kv = self.kv_a_proj_with_mqa(x)
         compressed_kv, k_pe = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1)
         k_pe = k_pe.reshape(B, L, 1, self.qk_rope_head_dim).transpose(0, 2, 1, 3)
-        kv = self.kv_b_proj(self.kv_a_layernorm(compressed_kv))
-        kv = kv.reshape(B, L, self.num_heads, -1).transpose(0, 2, 1, 3)
+        kv_latent = self.kv_a_layernorm(compressed_kv)
 
-        k_nope, values = mx.split(kv, [self.qk_nope_head_dim], axis=-1)
+        offset = cache.offset if cache is not None else 0
+        q_pe = self.rope(q_pe, offset)
+        k_pe = self.rope(k_pe, offset)
+
+        kv_latent = mx.expand_dims(kv_latent, axis=1)
 
         if cache is not None:
-            q_pe = self.rope(q_pe, cache.offset)
-            k_pe = self.rope(k_pe, cache.offset)
-            k_pe = mx.repeat(k_pe, self.num_heads, axis=1)
-            keys, values = cache.update_and_fetch(
-                mx.concatenate([k_nope, k_pe], axis=-1), values
+            kv_latent, k_pe = cache.update_and_fetch(kv_latent, k_pe)
+
+        pe_scores = (q_pe * self.scale) @ k_pe.swapaxes(-1, -2)
+        if mask is not None:
+            pe_scores = mx.where(
+                mask,
+                pe_scores,
+                mx.array(mx.finfo(pe_scores.dtype).min, pe_scores.dtype),
             )
+
+        if L == 1:
+            q_nope = self.embed_q(q_nope)
+            k = v = kv_latent
         else:
-            q_pe = self.rope(q_pe)
-            k_pe = self.rope(k_pe)
-            k_pe = mx.repeat(k_pe, self.num_heads, axis=1)
-            keys = mx.concatenate([k_nope, k_pe], axis=-1)
-
-        queries = mx.concatenate([q_nope, q_pe], axis=-1)
-
+            k = self.embed_q(kv_latent, transpose=False)
+            v = self.unembed_out(kv_latent)
         output = scaled_dot_product_attention(
-            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+            q_nope, k, v, cache=cache, scale=self.scale, mask=pe_scores
         )
+        if L == 1:
+            output = self.unembed_out(output)
+
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
         return self.o_proj(output)
 
@@ -211,7 +222,7 @@ def group_expert_select(
     scores = mx.take_along_axis(orig_scores, inds, axis=-1)
     if top_k > 1 and norm_topk_prob:
         denominator = scores.sum(axis=-1, keepdims=True)
-        scores = scores / denominator
+        scores = scores / (denominator + 1e-20)
     scores = scores * routed_scaling_factor
 
     return inds, scores
@@ -283,15 +294,12 @@ class Glm4MoeLiteDecoderLayer(nn.Module):
     def __init__(self, config: ModelArgs, layer_idx: int):
         super().__init__()
         self.self_attn = Glm4MoeLiteAttention(config)
-        self.mlp = (
-            Glm4MoeLiteMoE(config)
-            if (
-                config.n_routed_experts is not None
-                and layer_idx >= config.first_k_dense_replace
-                and layer_idx % config.moe_layer_freq == 0
-            )
-            else Glm4MoeLiteMLP(config)
+        use_moe = (
+            config.n_routed_experts is not None
+            and layer_idx >= config.first_k_dense_replace
+            and layer_idx % config.moe_layer_freq == 0
         )
+        self.mlp = Glm4MoeLiteMoE(config) if use_moe else Glm4MoeLiteMLP(config)
         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
@@ -332,7 +340,7 @@ class Glm4MoeLiteModel(PipelineMixin, nn.Module):
 
         if cache is None:
             cache = [None] * len(self.pipeline_layers)
-        mask = create_attention_mask(h, cache[0])
+        mask = create_attention_mask(h, cache[0], return_array=True)
 
         # Receive from the previous process in the pipeline
         if pipeline_rank < pipeline_size - 1:
@@ -371,6 +379,25 @@ class Model(nn.Module):
         return self.lm_head(out)
 
     def sanitize(self, weights):
+        def is_mpt_layer(key):
+            subkeys = key.split(".")
+            if len(subkeys) < 3:
+                return False
+            if (
+                subkeys[1] == "layers"
+                and int(subkeys[2]) >= self.args.num_hidden_layers
+            ):
+                return True
+            return False
+
+        new_weights = {}
+        for k, v in weights.items():
+            if is_mpt_layer(k):
+                continue
+            else:
+                new_weights[k] = v
+        weights = new_weights
+
         # Stack experts
         for l in range(self.args.num_hidden_layers):
             prefix = f"model.layers.{l}"
@@ -382,24 +409,48 @@ class Model(nn.Module):
                             for e in range(self.args.n_routed_experts)
                         ]
                         weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
+            prefix = f"model.layers.{l}.self_attn"
+            if f"{prefix}.kv_b_proj.weight" in weights:
+                layer = self.layers[l].self_attn.embed_q
+                quantized = f"{prefix}.kv_b_proj.scales" in weights
+                v = weights.pop(f"{prefix}.kv_b_proj.weight")
+                head_dim = self.args.qk_nope_head_dim + self.args.v_head_dim
 
-        num_mpt_layers = getattr(self.args, "num_nextn_predict_layers", 0) or 0
-        if num_mpt_layers:
-
-            def _is_mpt_layer(key: str) -> bool:
-                for idx in range(num_mpt_layers):
-                    if key.startswith(
-                        f"model.layers.{self.args.num_hidden_layers + idx}"
-                    ):
-                        return True
-                return False
-
-            weights = {k: v for k, v in weights.items() if not _is_mpt_layer(k)}
+                if quantized:
+                    dims = self.args.kv_lora_rank
+                    scales = weights.pop(f"{prefix}.kv_b_proj.scales")
+                    biases = weights.pop(f"{prefix}.kv_b_proj.biases")
+                    # Try to infer bits and group size
+                    bits = (v.shape[-1] * 32) // dims
+                    group_size = dims // scales.shape[-1]
+                    v = mx.dequantize(
+                        v, scales, biases, bits=bits, group_size=group_size
+                    )
+                num_heads = self.args.num_attention_heads
+                v = v.reshape(num_heads, head_dim, -1)
+                wk = mx.contiguous(
+                    v[:, : self.args.qk_nope_head_dim, :].swapaxes(-1, -2)
+                )
+                wv = mx.contiguous(v[:, self.args.qk_nope_head_dim :, :])
+                if quantized:
+                    wk, wk_scales, wk_biases = mx.quantize(
+                        wk, bits=bits, group_size=group_size
+                    )
+                    wv, wv_scales, wv_biases = mx.quantize(
+                        wv, bits=bits, group_size=group_size
+                    )
+                    weights[f"{prefix}.embed_q.scales"] = wk_scales
+                    weights[f"{prefix}.unembed_out.scales"] = wv_scales
+                    weights[f"{prefix}.embed_q.biases"] = wk_biases
+                    weights[f"{prefix}.unembed_out.biases"] = wv_biases
+                weights[f"{prefix}.embed_q.weight"] = wk
+                weights[f"{prefix}.unembed_out.weight"] = wv
 
         return weights
 
     def shard(self, group: Optional[mx.distributed.Group] = None):
         group = group or mx.distributed.init()
+        rank = group.rank()
         N = group.size()
         for layer in self.model.layers:
             # Shard the self attention
@@ -411,13 +462,20 @@ class Model(nn.Module):
                 layer.self_attn.q_b_proj = shard_linear(
                     layer.self_attn.q_b_proj, "all-to-sharded", group=group
                 )
-            layer.self_attn.kv_b_proj = shard_linear(
-                layer.self_attn.kv_b_proj, "all-to-sharded", group=group
-            )
+            layer.self_attn.num_heads //= N
+            num_heads = layer.self_attn.num_heads
+            sh = rank * num_heads
+            eh = sh + num_heads
+
+            def shard_heads(w):
+                return w[sh:eh]
+
+            layer.self_attn.embed_q.apply(shard_heads)
+            layer.self_attn.unembed_out.apply(shard_heads)
+
             layer.self_attn.o_proj = shard_linear(
                 layer.self_attn.o_proj, "sharded-to-all", group=group
             )
-            layer.self_attn.num_heads //= N
 
             # Shard the MLP
             if isinstance(layer.mlp, Glm4MoeLiteMLP):

mlx_lm/models/glm_moe_dsa.py

@@ -0,0 +1,53 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+from .base import BaseModelArgs
+from .deepseek_v32 import Model as DSV32Model
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    vocab_size: int
+    hidden_size: int
+    index_head_dim: int
+    index_n_heads: int
+    index_topk: int
+    intermediate_size: int
+    moe_intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    n_shared_experts: Optional[int]
+    n_routed_experts: Optional[int]
+    routed_scaling_factor: float
+    kv_lora_rank: int
+    q_lora_rank: int
+    qk_rope_head_dim: int
+    v_head_dim: int
+    qk_nope_head_dim: int
+    topk_method: str
+    scoring_func: str
+    norm_topk_prob: bool
+    n_group: int
+    topk_group: int
+    num_experts_per_tok: int
+    moe_layer_freq: int
+    first_k_dense_replace: int
+    max_position_embeddings: int
+    rms_norm_eps: float
+    rope_parameters: Dict
+    attention_bias: bool
+    rope_scaling: Dict = None
+    rope_theta: Optional[float] = None
+
+    def __post_init__(self):
+        self.rope_scaling = self.rope_parameters
+        self.rope_theta = self.rope_parameters["rope_theta"]
+
+
+class Model(DSV32Model):
+    def __init__(self, config: ModelArgs):
+        super().__init__(config)

mlx_lm/models/granitemoehybrid.py

@@ -13,7 +13,7 @@ from .base import (
     create_ssm_mask,
     scaled_dot_product_attention,
 )
-from .cache import KVCache, MambaCache
+from .cache import ArraysCache, KVCache
 from .rope_utils import initialize_rope
 from .ssm import ssm_update
 from .switch_layers import SwitchGLU
@@ -123,7 +123,7 @@ class GraniteMoeHybridMamba2Mixer(nn.Module):
     def _conv(
         self,
         conv_input: mx.array,
-        cache: Optional[MambaCache],
+        cache: Optional[ArraysCache],
         mask: Optional[mx.array],
     ) -> mx.array:
         if mask is not None:
@@ -160,7 +160,7 @@ class GraniteMoeHybridMamba2Mixer(nn.Module):
         B: mx.array,
         C: mx.array,
         dt: mx.array,
-        cache: Optional[MambaCache],
+        cache: Optional[ArraysCache],
         mask: Optional[mx.array],
     ) -> mx.array:
         batch_size, seq_len, _ = hidden_states.shape
@@ -197,7 +197,7 @@ class GraniteMoeHybridMamba2Mixer(nn.Module):
         self,
         hidden_states: mx.array,
         mask: Optional[mx.array],
-        cache: Optional[MambaCache] = None,
+        cache: Optional[ArraysCache] = None,
     ) -> mx.array:
 
         projected = self.in_proj(hidden_states)
@@ -496,7 +496,7 @@ class Model(nn.Module):
         caches = []
         for layer in self.layers:
             if layer.layer_type == "mamba":
-                caches.append(MambaCache())
+                caches.append(ArraysCache(size=2))
             elif layer.layer_type == "attention":
                 caches.append(KVCache())
         return caches

mlx_lm/models/jamba.py

@@ -14,7 +14,7 @@ from .base import (
     create_ssm_mask,
     scaled_dot_product_attention,
 )
-from .cache import KVCache, MambaCache
+from .cache import ArraysCache, KVCache
 from .switch_layers import SwitchGLU
 
 
@@ -341,7 +341,7 @@ class Model(nn.Module):
             if layer.is_attn:
                 caches.append(KVCache())
             else:
-                caches.append(MambaCache())
+                caches.append(ArraysCache(size=2))
         return caches
 
     def sanitize(self, weights):

mlx_lm/models/kimi_k25.py

@@ -0,0 +1,83 @@
+# Copyright © 2026 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_flatten, tree_unflatten
+
+from .base import BaseModelArgs
+from .deepseek_v3 import DeepseekV3Model
+from .deepseek_v3 import Model as DeepseekV3LM
+from .deepseek_v3 import ModelArgs as TextConfig
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    text_config: Union[TextConfig, dict]
+    model_type: str = "kimi_k25"
+
+    def __post_init__(self):
+        if isinstance(self.text_config, dict):
+            self.text_config = TextConfig.from_dict(self.text_config)
+
+
+class LanguageModel(nn.Module):
+    def __init__(self, config: TextConfig):
+        super().__init__()
+        self.args = config
+        self.model = DeepseekV3Model(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ):
+        out = self.model(inputs, cache)
+        return self.lm_head(out)
+
+
+class Model(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.args = config
+        self.model_type = config.model_type
+        self.language_model = LanguageModel(config.text_config)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ):
+        return self.language_model(inputs, cache)
+
+    def sanitize(self, weights):
+        weights = tree_unflatten(list(weights.items()))
+        weights.pop("vision_tower", None)
+        weights.pop("vision_model", None)
+        weights.pop("multi_modal_projector", None)
+        weights.pop("mm_projector", None)
+        lm_weights = dict(tree_flatten(weights["language_model"]))
+        lm_weights = DeepseekV3LM.sanitize(self.language_model, lm_weights)
+        weights["language_model"] = tree_unflatten(list(lm_weights.items()))
+        return dict(tree_flatten(weights))
+
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        DeepseekV3LM.shard(self.language_model, group)
+
+    @property
+    def model(self):
+        return self.language_model.model
+
+    @property
+    def layers(self):
+        return self.language_model.model.pipeline_layers
+
+    @property
+    def cast_predicate(self):
+        def predicate(k):
+            return "e_score_correction_bias" not in k
+
+        return predicate

mlx_lm/models/kimi_linear.py

@@ -13,9 +13,9 @@ from .base import (
     create_ssm_mask,
     scaled_dot_product_attention,
 )
-from .cache import KVCache, MambaCache
+from .cache import ArraysCache, KVCache
 from .gated_delta import gated_delta_update
-from .rope_utils import initialize_rope
+from .mla import MultiLinear
 from .switch_layers import SwitchGLU
 
 
@@ -165,6 +165,7 @@ class KimiMLAAttention(nn.Module):
         self.qk_rope_head_dim = args.qk_rope_head_dim or 0
         self.q_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
         self.v_head_dim = args.v_head_dim or args.head_dim
+        self.kv_lora_rank = args.kv_lora_rank
         self.scale = self.q_head_dim**-0.5
 
         hidden = args.hidden_size
@@ -175,23 +176,14 @@ class KimiMLAAttention(nn.Module):
             bias=False,
         )
         self.kv_a_layernorm = nn.RMSNorm(args.kv_lora_rank, eps=args.rms_norm_eps)
-        self.kv_b_proj = nn.Linear(
-            args.kv_lora_rank,
-            self.num_heads
-            * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),
-            bias=False,
+        self.embed_q = MultiLinear(
+            self.qk_nope_head_dim, args.kv_lora_rank, self.num_heads
+        )
+        self.unembed_out = MultiLinear(
+            args.kv_lora_rank, self.v_head_dim, self.num_heads
         )
         self.o_proj = nn.Linear(self.num_heads * self.v_head_dim, hidden, bias=False)
 
-        rope_dim = self.qk_rope_head_dim or self.q_head_dim
-        self.rope = initialize_rope(
-            rope_dim,
-            base=args.rope_theta,
-            traditional=False,
-            scaling_config=args.rope_scaling,
-            max_position_embeddings=args.model_max_length,
-        )
-
     def __call__(
         self,
         x: mx.array,
@@ -199,51 +191,45 @@ class KimiMLAAttention(nn.Module):
         cache: Optional[KVCache] = None,
     ) -> mx.array:
         B, L, _ = x.shape
-        q_states = self.q_proj(x).reshape(B, L, self.num_heads, self.q_head_dim)
-        q_pass, q_rot = mx.split(q_states, [self.qk_nope_head_dim], axis=-1)
 
-        compressed = self.kv_a_proj_with_mqa(x)
-        k_pass, k_rot = mx.split(
-            compressed, [compressed.shape[-1] - self.qk_rope_head_dim], axis=-1
-        )
-        k_pass = self.kv_a_layernorm(k_pass)
-        kv = self.kv_b_proj(k_pass)
-        kv = kv.reshape(
-            B,
-            L,
-            self.num_heads,
-            self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim,
-        )
-        k_pass, v_states = mx.split(kv, [self.qk_nope_head_dim], axis=-1)
+        q = self.q_proj(x).reshape(B, L, self.num_heads, self.q_head_dim)
+        q = q.transpose(0, 2, 1, 3)
+        q_nope, q_pe = mx.split(q, [self.qk_nope_head_dim], axis=-1)
 
-        if self.qk_rope_head_dim:
-            k_rot = mx.reshape(k_rot, (B, L, 1, self.qk_rope_head_dim))
-            k_rot = mx.broadcast_to(k_rot, (*k_pass.shape[:-1], self.qk_rope_head_dim))
-        else:
-            k_rot = mx.zeros((*k_pass.shape[:-1], 0), dtype=k_pass.dtype)
+        compressed_kv = self.kv_a_proj_with_mqa(x)
+        compressed_kv, k_pe = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1)
+        k_pe = k_pe.reshape(B, L, 1, self.qk_rope_head_dim).transpose(0, 2, 1, 3)
+        kv_latent = self.kv_a_layernorm(compressed_kv)
 
-        queries = mx.concatenate([q_pass, q_rot], axis=-1).transpose(0, 2, 1, 3)
-        keys = mx.concatenate([k_pass, k_rot], axis=-1).transpose(0, 2, 1, 3)
-        values = v_states.transpose(0, 2, 1, 3)
+        kv_latent = mx.expand_dims(kv_latent, axis=1)
 
         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)
+            kv_latent, k_pe = cache.update_and_fetch(kv_latent, k_pe)
 
-        out = scaled_dot_product_attention(
-            queries,
-            keys,
-            values,
-            cache,
-            scale=self.scale,
-            mask=mask,
+        pe_scores = (q_pe * self.scale) @ k_pe.swapaxes(-1, -2)
+        if mask is not None:
+            pe_scores = mx.where(
+                mask,
+                pe_scores,
+                mx.array(mx.finfo(pe_scores.dtype).min, pe_scores.dtype),
+            )
+
+        if L == 1:
+            q_nope = self.embed_q(q_nope)
+            k = v = kv_latent
+        else:
+            k = self.embed_q(kv_latent, transpose=False)
+            v = self.unembed_out(kv_latent)
+
+        output = scaled_dot_product_attention(
+            q_nope, k, v, cache=cache, scale=self.scale, mask=pe_scores
         )
-        out = out.transpose(0, 2, 1, 3).reshape(B, L, -1)
-        return self.o_proj(out)
+
+        if L == 1:
+            output = self.unembed_out(output)
+
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
 
 
 class ShortConv1d(nn.Module):
@@ -277,7 +263,7 @@ class ShortConv1d(nn.Module):
         out = nn.silu(self.conv(conv_input))
         n_keep = self.kernel_size - 1
         if lengths is not None:
-            ends = mx.clip(cache.lengths, 0, x.shape[1])
+            ends = mx.clip(lengths, 0, x.shape[1])
             positions = (ends[:, None] + mx.arange(n_keep))[..., None]
             new_state = mx.take_along_axis(conv_input, positions, axis=1)
         else:
@@ -335,39 +321,37 @@ class KimiDeltaAttention(nn.Module):
         dtype = x.dtype
 
         if cache is not None:
-            conv_state, ssm_state = cache
+            q_state, k_state, v_state, ssm_state = cache
             lengths = cache.lengths
         else:
-            conv_state = None
+            q_state = None
+            k_state = None
+            v_state = None
             ssm_state = None
             lengths = None
 
-        if conv_state is None:
+        if q_state is None:
             s = mx.zeros((B, self.conv_kernel - 1, self.projection_dim), dtype=dtype)
             q_state = s
             k_state = s
             v_state = s
-        else:
-            q_state, k_state, v_state = conv_state
 
         q_conv, q_state = self.q_conv(self.q_proj(x), q_state, mask, lengths)
         k_conv, k_state = self.k_conv(self.k_proj(x), k_state, mask, lengths)
         v_conv, v_state = self.v_conv(self.v_proj(x), v_state, mask, lengths)
 
         if cache is not None:
-            cache[0] = (q_state, k_state, v_state)
+            cache[0] = q_state
+            cache[1] = k_state
+            cache[2] = v_state
 
         q = q_conv.reshape(B, T, self.num_heads, self.head_dim)
         k = k_conv.reshape(B, T, self.num_heads, self.head_dim)
         v = v_conv.reshape(B, T, self.num_heads, self.head_dim)
 
-        def _l2norm(x, eps=1e-6):
-            norm = mx.linalg.norm(x, axis=-1, keepdims=True)
-            return x / (norm + eps)
-
-        q = _l2norm(q)
-        k = _l2norm(k)
-        q = q * self.scale
+        inv_scale = self.scale
+        q = (inv_scale**2) * mx.fast.rms_norm(q, None, 1e-6)
+        k = inv_scale * mx.fast.rms_norm(k, None, 1e-6)
 
         a_logits = self.f_b_proj(self.f_a_proj(x)).reshape(
             B, T, self.num_heads, self.head_dim
@@ -388,7 +372,7 @@ class KimiDeltaAttention(nn.Module):
         )
 
         if cache is not None:
-            cache[1] = ssm_state
+            cache[3] = ssm_state
             cache.advance(T)
 
         gate = self.g_b_proj(self.g_a_proj(x)).reshape(
@@ -462,7 +446,7 @@ class KimiLinearModel(nn.Module):
             cache = [None] * len(self.layers)
 
         ssm_mask = create_ssm_mask(h, cache[self.ssm_idx])
-        attn_mask = create_attention_mask(h, cache[self.attn_idx])
+        attn_mask = create_attention_mask(h, cache[self.attn_idx], return_array=True)
 
         for layer, layer_cache in zip(self.layers, cache):
             mask = ssm_mask if layer.is_linear else attn_mask
@@ -500,7 +484,7 @@ class Model(nn.Module):
         caches: List[Any] = []
         for layer in self.layers:
             if layer.is_linear:
-                caches.append(MambaCache())
+                caches.append(ArraysCache(size=4))
             else:
                 caches.append(KVCache())
         return caches
@@ -568,6 +552,42 @@ class Model(nn.Module):
                     if weights[dt_key].ndim > 1:
                         weights[dt_key] = mx.reshape(weights[dt_key], (-1,))
 
+            attn_prefix = f"{prefix}.self_attn"
+            kv_b_key = f"{attn_prefix}.kv_b_proj.weight"
+            if kv_b_key in weights:
+                qk_nope = self.args.qk_nope_head_dim or self.args.head_dim
+                v_head = self.args.v_head_dim or self.args.head_dim
+                head_dim = qk_nope + v_head
+                num_heads = self.args.num_attention_heads
+
+                quantized = f"{attn_prefix}.kv_b_proj.scales" in weights
+                v = weights.pop(kv_b_key)
+
+                if quantized:
+                    dims = self.args.kv_lora_rank
+                    scales = weights.pop(f"{attn_prefix}.kv_b_proj.scales")
+                    biases = weights.pop(f"{attn_prefix}.kv_b_proj.biases")
+                    bits = (v.shape[-1] * 32) // dims
+                    group_size = dims // scales.shape[-1]
+                    v = mx.dequantize(
+                        v, scales, biases, bits=bits, group_size=group_size
+                    )
+
+                v = v.reshape(num_heads, head_dim, -1)
+                wk = mx.contiguous(v[:, :qk_nope, :].swapaxes(-1, -2))
+                wv = mx.contiguous(v[:, qk_nope:, :])
+
+                if quantized:
+                    wk, wk_s, wk_b = mx.quantize(wk, bits=bits, group_size=group_size)
+                    wv, wv_s, wv_b = mx.quantize(wv, bits=bits, group_size=group_size)
+                    weights[f"{attn_prefix}.embed_q.scales"] = wk_s
+                    weights[f"{attn_prefix}.embed_q.biases"] = wk_b
+                    weights[f"{attn_prefix}.unembed_out.scales"] = wv_s
+                    weights[f"{attn_prefix}.unembed_out.biases"] = wv_b
+
+                weights[f"{attn_prefix}.embed_q.weight"] = wk
+                weights[f"{attn_prefix}.unembed_out.weight"] = wv
+
         return weights
 
     @property

mlx_lm/models/lfm2.py

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

mlx_lm/models/lfm2_moe.py

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

mlx_lm/models/longcat_flash.py

@@ -9,6 +9,7 @@ from mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients
 from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .cache import CacheList, KVCache
+from .mla import MultiLinear
 from .rope_utils import initialize_rope
 from .switch_layers import SwitchGLU
 
@@ -80,10 +81,11 @@ class LongcatFlashMLA(nn.Module):
             bias=args.attention_bias,
         )
         self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank)
-        self.kv_b_proj = nn.Linear(
-            self.kv_lora_rank,
-            self.num_attention_heads * (self.qk_nope_head_dim + args.v_head_dim),
-            bias=False,
+        self.embed_q = MultiLinear(
+            self.qk_nope_head_dim, self.kv_lora_rank, self.num_attention_heads
+        )
+        self.unembed_out = MultiLinear(
+            self.kv_lora_rank, self.v_head_dim, self.num_attention_heads
         )
 
         self.o_proj = nn.Linear(
@@ -122,56 +124,59 @@ class LongcatFlashMLA(nn.Module):
         B, L, _ = x.shape
 
         if self.q_lora_rank is None:
-            q_states = self.q_proj(x)
+            q = self.q_proj(x)
         else:
-            q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(x)))
+            q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(x)))
 
-        q_states = q_states.reshape(B, L, -1, self.qk_head_dim).transpose(0, 2, 1, 3)
-
-        if self.mla_scale_q_lora is not None:
-            q_states = q_states * self.mla_scale_q_lora
-
-        q_pass, q_rot = mx.split(q_states, [self.qk_nope_head_dim], axis=-1)
-
-        compressed_kv = self.kv_a_proj_with_mqa(x)
-        k_pass, k_rot = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1)
-        k_pass = self.kv_a_layernorm(k_pass)
-
-        if self.mla_scale_kv_lora is not None:
-            k_pass = k_pass * self.mla_scale_kv_lora
-
-        key_shape = (B, L, -1, self.qk_nope_head_dim + self.v_head_dim)
-        k_pass = self.kv_b_proj(k_pass).reshape(*key_shape).transpose(0, 2, 1, 3)
-        k_pass, value_states = mx.split(k_pass, [self.qk_nope_head_dim], axis=-1)
-
-        k_rot = k_rot.reshape(B, 1, L, self.qk_rope_head_dim)
-
-        if cache is not None:
-            q_rot = self.rope(q_rot, cache.offset)
-            k_rot = self.rope(k_rot, cache.offset)
-        else:
-            q_rot = self.rope(q_rot)
-            k_rot = self.rope(k_rot)
-
-        k_rot = mx.broadcast_to(k_rot, (*k_pass.shape[:-1], k_rot.shape[-1]))
-
-        query_states = mx.concatenate([q_pass, q_rot], axis=-1)
-        key_states = mx.concatenate([k_pass, k_rot], axis=-1)
-
-        if cache is not None:
-            key_states, value_states = cache.update_and_fetch(key_states, value_states)
-
-        attn_output = scaled_dot_product_attention(
-            query_states,
-            key_states,
-            value_states,
-            cache=cache,
-            scale=self.scale,
-            mask=mask,
+        q = q.reshape(B, L, self.num_attention_heads, self.qk_head_dim).transpose(
+            0, 2, 1, 3
         )
 
-        attn_output = attn_output.transpose(0, 2, 1, 3).reshape(B, L, -1)
-        return self.o_proj(attn_output)
+        if self.mla_scale_q_lora is not None:
+            q = q * self.mla_scale_q_lora
+
+        q_nope, q_pe = mx.split(q, [self.qk_nope_head_dim], axis=-1)
+
+        compressed_kv = self.kv_a_proj_with_mqa(x)
+        compressed_kv, k_pe = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1)
+        k_pe = k_pe.reshape(B, L, 1, self.qk_rope_head_dim).transpose(0, 2, 1, 3)
+        kv_latent = self.kv_a_layernorm(compressed_kv)
+
+        if self.mla_scale_kv_lora is not None:
+            kv_latent = kv_latent * self.mla_scale_kv_lora
+
+        offset = cache.offset if cache is not None else 0
+        q_pe = self.rope(q_pe, offset)
+        k_pe = self.rope(k_pe, offset)
+
+        kv_latent = mx.expand_dims(kv_latent, axis=1)
+
+        if cache is not None:
+            kv_latent, k_pe = cache.update_and_fetch(kv_latent, k_pe)
+
+        pe_scores = (q_pe * self.scale) @ k_pe.swapaxes(-1, -2)
+        if mask is not None:
+            pe_scores = mx.where(
+                mask,
+                pe_scores,
+                mx.array(mx.finfo(pe_scores.dtype).min, pe_scores.dtype),
+            )
+
+        if L == 1:
+            q_nope = self.embed_q(q_nope)
+            k = v = kv_latent
+        else:
+            k = self.embed_q(kv_latent, transpose=False)
+            v = self.unembed_out(kv_latent)
+
+        output = scaled_dot_product_attention(
+            q_nope, k, v, cache=cache, scale=self.scale, mask=pe_scores
+        )
+        if L == 1:
+            output = self.unembed_out(output)
+
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
 
 
 class LongcatFlashMLP(nn.Module):
@@ -339,7 +344,7 @@ class LongcatFlashModel(nn.Module):
         if cache is None:
             cache = [(None, None)] * self.num_layers
 
-        mask = create_attention_mask(h, cache[0][0])
+        mask = create_attention_mask(h, cache[0][0], return_array=True)
 
         for layer, c in zip(self.layers, cache):
             h = layer(h, mask, cache=c)
@@ -395,6 +400,47 @@ class Model(nn.Module):
                         ]
                         weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
 
+        for l in range(self.args.num_layers):
+            for i in range(2):
+                prefix = f"model.layers.{l}.self_attn.{i}"
+                kv_b_key = f"{prefix}.kv_b_proj.weight"
+                if kv_b_key in weights:
+                    num_heads = self.args.num_attention_heads
+                    head_dim = self.args.qk_nope_head_dim + self.args.v_head_dim
+                    quantized = f"{prefix}.kv_b_proj.scales" in weights
+                    v = weights.pop(kv_b_key)
+
+                    if quantized:
+                        dims = self.args.kv_lora_rank
+                        scales = weights.pop(f"{prefix}.kv_b_proj.scales")
+                        biases = weights.pop(f"{prefix}.kv_b_proj.biases")
+                        bits = (v.shape[-1] * 32) // dims
+                        group_size = dims // scales.shape[-1]
+                        v = mx.dequantize(
+                            v, scales, biases, bits=bits, group_size=group_size
+                        )
+
+                    v = v.reshape(num_heads, head_dim, -1)
+                    wk = mx.contiguous(
+                        v[:, : self.args.qk_nope_head_dim, :].swapaxes(-1, -2)
+                    )
+                    wv = mx.contiguous(v[:, self.args.qk_nope_head_dim :, :])
+
+                    if quantized:
+                        wk, wk_s, wk_b = mx.quantize(
+                            wk, bits=bits, group_size=group_size
+                        )
+                        wv, wv_s, wv_b = mx.quantize(
+                            wv, bits=bits, group_size=group_size
+                        )
+                        weights[f"{prefix}.embed_q.scales"] = wk_s
+                        weights[f"{prefix}.embed_q.biases"] = wk_b
+                        weights[f"{prefix}.unembed_out.scales"] = wv_s
+                        weights[f"{prefix}.unembed_out.biases"] = wv_b
+
+                    weights[f"{prefix}.embed_q.weight"] = wk
+                    weights[f"{prefix}.unembed_out.weight"] = wv
+
         new_weights = {}
         for k, v in weights.items():
             if k.startswith("model.mtp"):
@@ -408,6 +454,7 @@ class Model(nn.Module):
     def shard(self, group: Optional[mx.distributed.Group] = None):
         group = group or mx.distributed.init()
         N = group.size()
+        rank = group.rank()
 
         for layer in self.model.layers:
             for attn in layer.self_attn:
@@ -419,11 +466,17 @@ class Model(nn.Module):
                     attn.q_b_proj = shard_linear(
                         attn.q_b_proj, "all-to-sharded", group=group
                     )
-                attn.kv_b_proj = shard_linear(
-                    attn.kv_b_proj, "all-to-sharded", group=group
-                )
                 attn.o_proj = shard_linear(attn.o_proj, "sharded-to-all", group=group)
                 attn.num_attention_heads //= N
+                num_heads = attn.num_attention_heads
+                sh = rank * num_heads
+                eh = sh + num_heads
+
+                def shard_heads(w):
+                    return w[sh:eh]
+
+                attn.embed_q.apply(shard_heads)
+                attn.unembed_out.apply(shard_heads)
 
             for mlp in layer.mlps:
                 mlp.gate_proj = shard_linear(

mlx_lm/models/longcat_flash_ngram.py

@@ -0,0 +1,214 @@
+# Copyright © 2026 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask
+from .cache import ArraysCache, CacheList, KVCache
+from .longcat_flash import LongcatFlashDecoderLayer
+from .longcat_flash import Model as LongcatFlashLM
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    ffn_hidden_size: int
+    moe_topk: int
+    expert_ffn_hidden_size: int
+    n_routed_experts: int
+    zero_expert_num: int
+    num_layers: int
+    vocab_size: int
+    max_position_embeddings: int
+    num_attention_heads: int
+    kv_lora_rank: int
+    q_lora_rank: int
+    qk_rope_head_dim: int
+    qk_nope_head_dim: int
+    v_head_dim: int
+    routed_scaling_factor: float
+    rms_norm_eps: float
+    rope_theta: float
+    mla_scale_q_lora: bool
+    mla_scale_kv_lora: bool
+    attention_bias: bool = False
+    zero_expert_type: str = "identity"
+    ngram_vocab_size_ratio: int = 78
+    emb_neighbor_num: int = 4
+    emb_split_num: int = 4
+    norm_topk_prob: bool = False
+    router_bias: bool = False
+    rope_scaling: Optional[Dict] = None
+
+
+class NgramEmbedding(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.vocab_size = args.vocab_size
+        self.hidden_size = args.hidden_size
+        self.m = args.ngram_vocab_size_ratio * args.vocab_size
+        self.k = args.emb_split_num
+        self.n = args.emb_neighbor_num
+
+        self.word_embeddings = nn.Embedding(args.vocab_size, args.hidden_size)
+
+        num_embedders = self.k * (self.n - 1)
+        emb_dim = args.hidden_size // num_embedders
+
+        self.embedders = []
+        self.post_projs = []
+        for i in range(num_embedders):
+            emb_vocab_size = int(self.m + i * 2 + 1)
+            self.embedders.append(nn.Embedding(emb_vocab_size, emb_dim))
+            self.post_projs.append(nn.Linear(emb_dim, args.hidden_size, bias=False))
+        self._compute_vocab_mods()
+
+    def _compute_vocab_mods(self):
+        vocab_mods = {}
+        for i in range(2, self.n + 1):
+            for j in range(self.k):
+                index = (i - 2) * self.k + j
+                emb_vocab_dim = int(self.m + index * 2 + 1)
+                mods = []
+                power_mod = 1
+                for _ in range(i - 1):
+                    power_mod = (power_mod * self.vocab_size) % emb_vocab_dim
+                    mods.append(power_mod)
+                vocab_mods[(i, j)] = mods
+        self._vocab_mods = vocab_mods
+
+    def _shift_right(self, x: mx.array, n: int) -> mx.array:
+        if n <= 0:
+            return x
+        batch_size, seq_len = x.shape
+        if seq_len <= n:
+            return mx.zeros_like(x)
+        return mx.concatenate(
+            [mx.zeros((batch_size, n), dtype=x.dtype), x[..., :-n]], axis=-1
+        )
+
+    def _get_ngram_ids(
+        self,
+        input_ids: mx.array,
+        shifted_ids: Dict[int, mx.array],
+        vocab_mods: List[int],
+        ngram: int,
+    ) -> mx.array:
+        ngram_ids = input_ids
+        for k in range(2, ngram + 1):
+            ngram_ids = ngram_ids + shifted_ids[k] * vocab_mods[k - 2]
+        return ngram_ids
+
+    def __call__(
+        self,
+        input_ids: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        seq_len = input_ids.shape[-1]
+
+        input_ids = input_ids.astype(mx.int64)
+        if cache is not None:
+            context = cache[0]
+            if context is None:
+                context = input_ids
+            else:
+                context = mx.concatenate([context, input_ids], axis=-1)
+            cache[0] = context[..., max(0, context.shape[-1] - self.n + 1) :]
+        else:
+            context = input_ids
+
+        x = self.word_embeddings(input_ids)
+        vocab_mods = self._vocab_mods
+
+        shifted_ids = {}
+        for i in range(2, self.n + 1):
+            shifted_ids[i] = self._shift_right(context, i - 1)
+
+        for i in range(2, self.n + 1):
+            for j in range(self.k):
+                index = (i - 2) * self.k + j
+                emb_vocab_dim = int(self.m + index * 2 + 1)
+                ngram_ids = self._get_ngram_ids(
+                    context, shifted_ids, vocab_mods[(i, j)], ngram=i
+                )
+                new_ids = (ngram_ids % emb_vocab_dim)[..., -seq_len:]
+                x_ngram = self.embedders[index](new_ids)
+                x_proj = self.post_projs[index](x_ngram)
+                x = x + x_proj
+
+        return x / (1 + self.k * (self.n - 1))
+
+
+class LongcatFlashNgramModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_layers = args.num_layers
+        self.ngram_embeddings = NgramEmbedding(args)
+        self.layers = [LongcatFlashDecoderLayer(args) for _ in range(args.num_layers)]
+        self.norm = nn.RMSNorm(args.hidden_size, args.rms_norm_eps)
+
+    def __call__(
+        self,
+        input_ids: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        if cache is None:
+            cache = [None] + [(None, None)] * self.num_layers
+
+        h = self.ngram_embeddings(input_ids, cache=cache[0])
+
+        mask = create_attention_mask(h, cache[1][0], return_array=True)
+
+        for layer, c in zip(self.layers, cache[1:]):
+            h = layer(h, mask, cache=c)
+
+        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 = LongcatFlashNgramModel(args)
+        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ):
+        out = self.model(inputs, cache)
+        return self.lm_head(out)
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    @property
+    def quant_predicate(self):
+        return LongcatFlashLM.quant_predicate.fget(self)
+
+    @property
+    def cast_predicate(self):
+        return LongcatFlashLM.cast_predicate.fget(self)
+
+    def sanitize(self, weights):
+        weights = LongcatFlashLM.sanitize(self, weights)
+        if "model.embed_tokens.weight" in weights:
+            weights["model.ngram_embeddings.word_embeddings.weight"] = weights.pop(
+                "model.embed_tokens.weight"
+            )
+        return weights
+
+    def make_cache(self):
+        return [ArraysCache(size=1)] + [
+            CacheList(KVCache(), KVCache()) for _ in self.model.layers
+        ]
+
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        LongcatFlashLM.shard(self, group)

mlx_lm/models/mamba.py

@@ -8,7 +8,7 @@ import mlx.nn as nn
 
 from .activations import swiglu
 from .base import BaseModelArgs
-from .cache import MambaCache
+from .cache import ArraysCache
 
 
 @dataclass
@@ -153,7 +153,7 @@ class MambaBlock(nn.Module):
             x, conv_cache, state_cache
         )
 
-        if isinstance(cache, MambaCache):
+        if isinstance(cache, ArraysCache):
             cache[0] = new_conv_cache
             cache[1] = new_state_cache
 
@@ -208,7 +208,7 @@ class Model(nn.Module):
         return logits
 
     def make_cache(self):
-        return [MambaCache() for _ in range(len(self.layers))]
+        return [ArraysCache(size=2) for _ in range(len(self.layers))]
 
     @property
     def layers(self):

mlx_lm/models/mamba2.py

@@ -9,7 +9,7 @@ import mlx.nn as nn
 
 from .activations import swiglu
 from .base import BaseModelArgs, create_ssm_mask
-from .cache import MambaCache
+from .cache import ArraysCache
 from .ssm import ssm_update
 
 
@@ -97,7 +97,7 @@ class Mamba2Block(nn.Module):
     def _conv(
         self,
         conv_input: mx.array,
-        cache: Optional[MambaCache],
+        cache: Optional[ArraysCache],
         mask: Optional[mx.array],
     ) -> mx.array:
         if mask is not None:
@@ -134,7 +134,7 @@ class Mamba2Block(nn.Module):
         B: mx.array,
         C: mx.array,
         dt: mx.array,
-        cache: Optional[MambaCache],
+        cache: Optional[ArraysCache],
         mask: Optional[mx.array],
     ) -> mx.array:
         batch_size, seq_len, _ = hidden_states.shape
@@ -169,7 +169,7 @@ class Mamba2Block(nn.Module):
         self,
         hidden_states: mx.array,
         mask: Optional[mx.array],
-        cache: Optional[MambaCache] = None,
+        cache: Optional[ArraysCache] = None,
     ) -> mx.array:
         projected = self.in_proj(hidden_states)
         gate, conv_input, dt = mx.split(
@@ -200,7 +200,7 @@ class ResidualBlock(nn.Module):
         self.norm = nn.RMSNorm(args.hidden_size)
 
     def __call__(
-        self, x: mx.array, mask: Optional[mx.array], cache: Optional[MambaCache] = None
+        self, x: mx.array, mask: Optional[mx.array], cache: Optional[ArraysCache] = None
     ) -> mx.array:
         output = self.mixer(self.norm(x), mask, cache)
         return output + x
@@ -215,7 +215,7 @@ class Mamba2(nn.Module):
         self.norm_f = nn.RMSNorm(args.hidden_size, eps=args.layer_norm_epsilon)
 
     def __call__(
-        self, x: mx.array, cache: Optional[list[MambaCache]] = None
+        self, x: mx.array, cache: Optional[list[ArraysCache]] = None
     ) -> mx.array:
         hidden = self.embeddings(x)
 
@@ -240,7 +240,7 @@ class Model(nn.Module):
             self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
 
     def __call__(
-        self, inputs: mx.array, cache: Optional[list[MambaCache]] = None
+        self, inputs: mx.array, cache: Optional[list[ArraysCache]] = None
     ) -> mx.array:
         hidden = self.backbone(inputs, cache)
 
@@ -250,8 +250,8 @@ class Model(nn.Module):
             logits = self.lm_head(hidden)
         return logits
 
-    def make_cache(self, batch_size: int = 1) -> list[MambaCache]:
-        return [MambaCache() for _ in range(self.args.num_hidden_layers)]
+    def make_cache(self, batch_size: int = 1) -> list[ArraysCache]:
+        return [ArraysCache(size=2) for _ in range(self.args.num_hidden_layers)]
 
     @property
     def layers(self):

mlx_lm/models/minimax.py

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

mlx_lm/models/mixtral.py

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

mlx_lm/models/mla.py

@@ -0,0 +1,85 @@
+# Copyright © 2026 Apple Inc.
+
+import math
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class MultiLinear(nn.Module):
+    def __init__(self, input_dims: int, output_dims: int, num_heads: int) -> None:
+        super().__init__()
+        scale = math.sqrt(1.0 / input_dims)
+        self.weight = mx.random.uniform(
+            low=-scale,
+            high=scale,
+            shape=(num_heads, output_dims, input_dims),
+        )
+
+    def __call__(self, x, transpose=True):
+        if transpose:
+            return x @ self.weight.swapaxes(-1, -2)
+        else:
+            return x @ self.weight
+
+    def to_quantized(
+        self,
+        group_size: int,
+        bits: int,
+        mode: str = "affine",
+    ):
+        num_heads, output_dims, input_dims = self.weight.shape
+        ql = QuantizedMultiLinear(
+            input_dims, output_dims, num_heads, group_size, bits, mode
+        )
+        ql.weight, ql.scales, *biases = mx.quantize(
+            self.weight,
+            group_size,
+            bits,
+            mode=mode,
+        )
+        ql.biases = biases[0] if biases else None
+        return ql
+
+
+class QuantizedMultiLinear(nn.Module):
+    def __init__(
+        self,
+        input_dims: int,
+        output_dims: int,
+        num_heads: int,
+        group_size: int,
+        bits: int,
+        mode: str,
+    ):
+        super().__init__()
+
+        self.group_size = group_size
+        self.bits = bits
+        self.mode = mode
+
+        # Initialize the quantized weight
+        scale = math.sqrt(1 / input_dims)
+        weight = mx.random.uniform(
+            low=-scale,
+            high=scale,
+            shape=(num_heads, output_dims, input_dims),
+        )
+        self.weight, self.scales, *biases = mx.quantize(
+            weight, group_size, bits, mode=mode
+        )
+        self.biases = biases[0] if biases else None
+
+        self.freeze()
+
+    def __call__(self, x, transpose=True):
+        return mx.quantized_matmul(
+            x,
+            self["weight"],
+            scales=self["scales"],
+            biases=self.get("biases"),
+            transpose=transpose,
+            group_size=self.group_size,
+            bits=self.bits,
+            mode=self.mode,
+        )

mlx_lm/models/nemotron_h.py

@@ -14,7 +14,7 @@ from .base import (
     create_ssm_mask,
     scaled_dot_product_attention,
 )
-from .cache import KVCache, MambaCache
+from .cache import ArraysCache, KVCache
 from .ssm import ssm_update
 from .switch_layers import SwitchMLP
 
@@ -36,15 +36,16 @@ class ModelArgs(BaseModelArgs):
     ssm_state_size: int
     conv_kernel: int
     n_groups: int
-    time_step_limit: Tuple[float, float]
     mlp_bias: bool
     layer_norm_epsilon: float
     use_bias: bool
     use_conv_bias: bool
-    hybrid_override_pattern: List[str]
+    hybrid_override_pattern: Optional[List[str]] = None
+    layers_block_type: Optional[List[str]] = None
     head_dim: Optional[int] = None
     moe_intermediate_size: Optional[int] = None
     moe_shared_expert_intermediate_size: Optional[int] = None
+    moe_latent_size: Optional[int] = None
     n_group: Optional[int] = None
     n_routed_experts: Optional[int] = None
     n_shared_experts: Optional[int] = None
@@ -52,6 +53,24 @@ class ModelArgs(BaseModelArgs):
     num_experts_per_tok: Optional[int] = None
     norm_topk_prob: Optional[bool] = None
     routed_scaling_factor: Optional[float] = None
+    time_step_limit: Optional[Tuple[float, float]] = None
+    time_step_min: Optional[float] = None
+    time_step_max: Optional[float] = None
+
+    # Map from layers_block_type names to single-char pattern codes
+    _block_type_to_char = {"mamba": "M", "attention": "*", "moe": "E", "mlp": "-"}
+
+    def __post_init__(self):
+        if self.time_step_limit is None and self.time_step_min is not None:
+            self.time_step_limit = (self.time_step_min, float("inf"))
+
+        # Normalize to hybrid_override_pattern (single-char list)
+        if self.hybrid_override_pattern is None and self.layers_block_type is not None:
+            self.hybrid_override_pattern = [
+                self._block_type_to_char[t] for t in self.layers_block_type
+            ]
+        if self.hybrid_override_pattern is not None:
+            self.num_hidden_layers = len(self.hybrid_override_pattern)
 
 
 class MambaRMSNormGated(nn.Module):
@@ -115,7 +134,7 @@ class NemotronHMamba2Mixer(nn.Module):
     def _conv(
         self,
         conv_input: mx.array,
-        cache: Optional[MambaCache],
+        cache: Optional[ArraysCache],
         mask: Optional[mx.array],
     ) -> mx.array:
         if mask is not None:
@@ -152,7 +171,7 @@ class NemotronHMamba2Mixer(nn.Module):
         B: mx.array,
         C: mx.array,
         dt: mx.array,
-        cache: Optional[MambaCache],
+        cache: Optional[ArraysCache],
         mask: Optional[mx.array],
     ) -> mx.array:
         batch_size, seq_len, _ = hidden_states.shape
@@ -189,7 +208,7 @@ class NemotronHMamba2Mixer(nn.Module):
         self,
         hidden_states: mx.array,
         mask: Optional[mx.array],
-        cache: Optional[MambaCache] = None,
+        cache: Optional[ArraysCache] = None,
     ) -> mx.array:
 
         projected = self.in_proj(hidden_states)
@@ -355,8 +374,16 @@ class NemotronHMoE(nn.Module):
         super().__init__()
         self.config = config
         self.num_experts_per_tok = config.num_experts_per_tok
+        self.moe_latent_size = config.moe_latent_size
+
+        # When latent projection is used, experts operate on the latent dim
+        expert_input_dim = (
+            config.moe_latent_size
+            if config.moe_latent_size is not None
+            else config.hidden_size
+        )
         self.switch_mlp = SwitchMLP(
-            config.hidden_size,
+            expert_input_dim,
             config.moe_intermediate_size,
             config.n_routed_experts,
             activation=nn.ReLU2(),
@@ -369,12 +396,30 @@ class NemotronHMoE(nn.Module):
                 config, intermediate_size=intermediate_size
             )
 
+        # Latent projection layers for dimensionality reduction before/after experts
+        if config.moe_latent_size is not None:
+            self.fc1_latent_proj = nn.Linear(
+                config.hidden_size, config.moe_latent_size, bias=config.mlp_bias
+            )
+            self.fc2_latent_proj = nn.Linear(
+                config.moe_latent_size, config.hidden_size, bias=config.mlp_bias
+            )
+
     def __call__(self, x):
+        residuals = x
         inds, scores = self.gate(x)
+
+        if self.moe_latent_size is not None:
+            x = self.fc1_latent_proj(x)
+
         y = self.switch_mlp(x, inds)
         y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype)
+
+        if self.moe_latent_size is not None:
+            y = self.fc2_latent_proj(y)
+
         if self.config.n_shared_experts is not None:
-            y = y + self.shared_experts(x)
+            y = y + self.shared_experts(residuals)
 
         return y
 
@@ -485,12 +530,13 @@ class Model(nn.Module):
         caches = []
         for l in self.layers:
             if l.block_type == "M":
-                caches.append(MambaCache())
+                caches.append(ArraysCache(size=2))
             elif l.block_type == "*":
                 caches.append(KVCache())
         return caches
 
     def sanitize(self, weights):
+        weights = {k: v for (k, v) in weights.items() if not k.startswith("mtp.")}
         for k, v in weights.items():
             if "conv1d.weight" in k and v.shape[-1] != 1:
                 weights[k] = v.moveaxis(2, 1)

mlx_lm/models/plamo2.py

@@ -10,7 +10,7 @@ import mlx.nn as nn
 from mlx_lm.models.base import BaseModelArgs, create_attention_mask, create_ssm_mask
 
 from .activations import swiglu
-from .cache import KVCache, MambaCache
+from .cache import ArraysCache, KVCache
 from .ssm import ssm_update
 
 
@@ -101,7 +101,7 @@ class Mamba(nn.Module):
     def _conv(
         self,
         conv_input: mx.array,
-        cache: Optional[MambaCache],
+        cache: Optional[ArraysCache],
         mask: Optional[mx.array],
     ) -> mx.array:
         if mask is not None:
@@ -459,7 +459,7 @@ class Model(nn.Module):
     def make_cache(self):
         # TODO use RotatingKVCache is not full_attn
         # full_attn = self.layer_idx in self.config.full_attention_idx
-        return [MambaCache() if l.is_mamba else KVCache() for l in self.layers]
+        return [ArraysCache(size=2) if l.is_mamba else KVCache() for l in self.layers]
 
     def __call__(self, inputs: mx.array, cache=None) -> mx.array:
         outputs = self.model(

mlx_lm/models/qwen3_5.py

@@ -0,0 +1,524 @@
+# Copyright © 2026 Apple Inc.
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients
+from mlx.utils import tree_map
+
+from .base import (
+    BaseModelArgs,
+    create_attention_mask,
+    create_ssm_mask,
+)
+from .cache import ArraysCache, KVCache
+from .gated_delta import gated_delta_update
+from .qwen3_next import Qwen3NextAttention as Attention
+from .qwen3_next import Qwen3NextMLP as MLP
+from .qwen3_next import Qwen3NextRMSNormGated as RMSNormGated
+from .qwen3_next import Qwen3NextSparseMoeBlock as SparseMoeBlock
+
+
+@dataclass
+class TextModelArgs(BaseModelArgs):
+    model_type: str = ""
+    hidden_size: int = 4096
+    intermediate_size: int = 14336
+    num_hidden_layers: int = 32
+    num_attention_heads: int = 32
+    rms_norm_eps: float = 1e-6
+    vocab_size: int = 151936
+    num_key_value_heads: int = 8
+    max_position_embeddings: int = 131072
+    linear_num_value_heads: int = 64
+    linear_num_key_heads: int = 16
+    linear_key_head_dim: int = 192
+    linear_value_head_dim: int = 128
+    linear_conv_kernel_dim: int = 4
+    tie_word_embeddings: bool = False
+    attention_bias: bool = False
+    head_dim: Optional[int] = None
+    full_attention_interval: int = 4
+
+    # MoE fields (optional, for Qwen3_5MoeForConditionalGeneration)
+    num_experts: int = 0
+    num_experts_per_tok: int = 0
+    decoder_sparse_step: int = 1
+    shared_expert_intermediate_size: int = 0
+    moe_intermediate_size: int = 0
+    norm_topk_prob: bool = True
+
+    # Rope parameters
+    rope_parameters: Optional[Dict[str, Union[float, str, bool, List[int]]]] = field(
+        default_factory=lambda: {
+            "type": "default",
+            "mrope_section": [11, 11, 10],
+            "rope_theta": 100000,
+            "partial_rotary_factor": 0.25,
+        }
+    )
+
+    # Derived from rope_parameters (set in __post_init__)
+    partial_rotary_factor: float = 0.25
+    rope_theta: float = 100000.0
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+
+    def __post_init__(self):
+        if self.head_dim is None:
+            self.head_dim = self.hidden_size // self.num_attention_heads
+
+        if self.rope_parameters:
+            if (
+                "type" not in self.rope_parameters
+                and "rope_type" in self.rope_parameters
+            ):
+                self.rope_parameters["type"] = self.rope_parameters.pop("rope_type")
+
+            self.partial_rotary_factor = self.rope_parameters.get(
+                "partial_rotary_factor", 0.25
+            )
+            self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0)
+            self.rope_scaling = self.rope_parameters
+
+
+class GatedDeltaNet(nn.Module):
+    def __init__(self, config: TextModelArgs):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.num_v_heads = config.linear_num_value_heads
+        self.num_k_heads = config.linear_num_key_heads
+        self.head_k_dim = config.linear_key_head_dim
+        self.head_v_dim = config.linear_value_head_dim
+        self.key_dim = self.head_k_dim * self.num_k_heads
+        self.value_dim = self.head_v_dim * self.num_v_heads
+        if self.num_v_heads % self.num_k_heads != 0:
+            raise ValueError(
+                f"num_v_heads ({self.num_v_heads}) must be divisible by num_k_heads ({self.num_k_heads})"
+            )
+
+        self.conv_kernel_size = config.linear_conv_kernel_dim
+        self.layer_norm_epsilon = config.rms_norm_eps
+
+        self.conv_dim = self.key_dim * 2 + self.value_dim
+        self.conv1d = nn.Conv1d(
+            in_channels=self.conv_dim,
+            out_channels=self.conv_dim,
+            bias=False,
+            kernel_size=self.conv_kernel_size,
+            groups=self.conv_dim,
+            padding=0,
+        )
+
+        self.in_proj_qkv = nn.Linear(
+            self.hidden_size, self.key_dim * 2 + self.value_dim, bias=False
+        )
+        self.in_proj_z = nn.Linear(self.hidden_size, self.value_dim, bias=False)
+        self.in_proj_b = nn.Linear(self.hidden_size, self.num_v_heads, bias=False)
+        self.in_proj_a = nn.Linear(self.hidden_size, self.num_v_heads, bias=False)
+
+        self.dt_bias = mx.ones(self.num_v_heads)
+
+        A = mx.random.uniform(low=0, high=16, shape=(self.num_v_heads,))
+        self.A_log = mx.log(A)
+
+        self.norm = RMSNormGated(self.head_v_dim, eps=self.layer_norm_epsilon)
+
+        self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)
+
+        self.sharding_group = None
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, S, _ = inputs.shape
+
+        if self.sharding_group is not None:
+            inputs = sum_gradients(self.sharding_group)(inputs)
+
+        qkv = self.in_proj_qkv(inputs)
+        z = self.in_proj_z(inputs).reshape(B, S, self.num_v_heads, self.head_v_dim)
+        b = self.in_proj_b(inputs)
+        a = self.in_proj_a(inputs)
+
+        if cache is not None and cache[0] is not None:
+            conv_state = cache[0]
+        else:
+            conv_state = mx.zeros(
+                (B, self.conv_kernel_size - 1, self.conv_dim),
+                dtype=inputs.dtype,
+            )
+
+        if mask is not None:
+            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) :]
+        conv_out = nn.silu(self.conv1d(conv_input))
+
+        q, k, v = [
+            t.reshape(B, S, h, d)
+            for t, h, d in zip(
+                mx.split(conv_out, [self.key_dim, 2 * self.key_dim], -1),
+                [self.num_k_heads, self.num_k_heads, self.num_v_heads],
+                [self.head_k_dim, self.head_k_dim, self.head_v_dim],
+            )
+        ]
+
+        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,
+            mask,
+            use_kernel=not self.training,
+        )
+
+        if cache is not None:
+            cache[1] = state
+
+        out = self.norm(out, z)
+        out = self.out_proj(out.reshape(B, S, -1))
+
+        if self.sharding_group is not None:
+            out = mx.distributed.all_sum(out, group=self.sharding_group)
+
+        return out
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, args: TextModelArgs, layer_idx: int):
+        super().__init__()
+        self.is_linear = (layer_idx + 1) % args.full_attention_interval != 0
+        if self.is_linear:
+            self.linear_attn = GatedDeltaNet(args)
+        else:
+            self.self_attn = Attention(args)
+
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+
+        if args.num_experts > 0:
+            self.mlp = SparseMoeBlock(args)
+        else:
+            self.mlp = MLP(args.hidden_size, args.intermediate_size)
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        if self.is_linear:
+            r = self.linear_attn(self.input_layernorm(x), mask, cache)
+        else:
+            r = self.self_attn(self.input_layernorm(x), mask, cache)
+        h = x + r
+        out = h + self.mlp(self.post_attention_layernorm(h))
+        return out
+
+
+class Qwen3_5TextModel(nn.Module):
+    def __init__(self, args: TextModelArgs):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            DecoderLayer(args=args, layer_idx=i) 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__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+        input_embeddings: Optional[mx.array] = None,
+    ) -> mx.array:
+        if input_embeddings is not None:
+            hidden_states = input_embeddings
+        else:
+            hidden_states = self.embed_tokens(inputs)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        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)
+
+
+class TextModel(nn.Module):
+    def __init__(self, args: TextModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = Qwen3_5TextModel(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+        input_embeddings: Optional[mx.array] = None,
+    ) -> mx.array:
+        out = self.model(inputs, cache, input_embeddings=input_embeddings)
+        if self.args.tie_word_embeddings:
+            out = self.model.embed_tokens.as_linear(out)
+        else:
+            out = self.lm_head(out)
+        return out
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    def make_cache(self):
+        return [ArraysCache(size=2) if l.is_linear else KVCache() for l in self.layers]
+
+    def sanitize(self, weights):
+        has_mtp_weights = any("mtp." in k for k in weights)
+        has_unsanitized_conv1d = any(
+            "conv1d.weight" in k and v.shape[-1] != 1 for k, v in weights.items()
+        )
+        should_shift_norm_weights = has_mtp_weights or has_unsanitized_conv1d
+        weights = {k: v for k, v in weights.items() if "mtp." not in k}
+
+        if self.args.tie_word_embeddings:
+            weights.pop("lm_head.weight", None)
+
+        norm_keys = (
+            ".input_layernorm.weight",
+            ".post_attention_layernorm.weight",
+            "model.norm.weight",
+            ".q_norm.weight",
+            ".k_norm.weight",
+        )
+        for k, v in weights.items():
+            if "conv1d.weight" in k and v.shape[-1] != 1:
+                weights[k] = v.moveaxis(2, 1)
+            if should_shift_norm_weights and any(k.endswith(sfx) for sfx in norm_keys):
+                if v.ndim == 1:
+                    weights[k] = v + 1.0
+        return weights
+
+    @property
+    def quant_predicate(self):
+        if self.args.num_experts <= 0:
+            return None
+
+        def predicate(path, _):
+            if path.endswith("mlp.gate") or path.endswith("shared_expert_gate"):
+                return {"group_size": 64, "bits": 8}
+            return True
+
+        return predicate
+
+    @property
+    def cast_predicate(self):
+        def predicate(path: str):
+            if path.endswith("A_log"):
+                return False
+            return True
+
+        return predicate
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    text_config: dict
+
+    @classmethod
+    def from_dict(cls, params):
+        if "text_config" not in params:
+            return cls(model_type=params["model_type"], text_config=params)
+        return super().from_dict(params)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.language_model = TextModel(TextModelArgs.from_dict(args.text_config))
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+        input_embeddings: Optional[mx.array] = None,
+    ):
+        return self.language_model(
+            inputs, cache=cache, input_embeddings=input_embeddings
+        )
+
+    def sanitize(self, weights):
+        sanitized = {}
+        for key, value in weights.items():
+            if key.startswith("vision_tower") or key.startswith("model.visual"):
+                continue
+            if key.startswith("model.visual"):
+                continue
+            if key.startswith("model.language_model"):
+                key = key.replace("model.language_model", "language_model.model")
+            elif key.startswith("language_model."):
+                pass
+            else:
+                key = "language_model." + key
+            sanitized[key] = value
+        return self.language_model.sanitize(sanitized)
+
+    def shard(self, group=None):
+        group = group or mx.distributed.init()
+        N = group.size()
+        rank = group.rank()
+
+        # A sharding factory for the convolution in gated delta net
+        def conv_sharding(key_dim):
+            return lambda p, w: (0, [key_dim, 2 * key_dim])
+
+        def repeat_kv_layer_inplace(layer, h):
+            # No repeat needed cause we have more heads than nodes
+            if N <= h:
+                return
+
+            # Repeat function to apply to the layer weights
+            def _repeat(p):
+                s = p.shape
+                p = p.reshape(h, s[0] // h, *s[1:])
+                p = mx.repeat(p, N // h, axis=0)
+                p = p.reshape(-1, *s[1:])
+                return p
+
+            layer.update(tree_map(_repeat, layer.parameters()))
+
+        for layer in self.layers:
+            # Linear attention
+            if layer.is_linear:
+                kd = layer.linear_attn.key_dim
+                layer.linear_attn.sharding_group = group
+                shard_inplace(layer.linear_attn.conv1d, conv_sharding(kd), group=group)
+                layer.linear_attn.conv1d.groups //= N
+                shard_inplace(
+                    layer.linear_attn.in_proj_qkv,
+                    "all-to-sharded",
+                    segments=[kd, 2 * kd],
+                    group=group,
+                )
+                shard_inplace(
+                    layer.linear_attn.in_proj_z, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.linear_attn.in_proj_b, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.linear_attn.in_proj_a, "all-to-sharded", group=group
+                )
+                layer.linear_attn.dt_bias = mx.contiguous(
+                    mx.split(layer.linear_attn.dt_bias, N)[rank]
+                )
+                layer.linear_attn.A_log = mx.contiguous(
+                    mx.split(layer.linear_attn.A_log, N)[rank]
+                )
+                shard_inplace(layer.linear_attn.out_proj, "sharded-to-all", group=group)
+                layer.linear_attn.num_k_heads //= N
+                layer.linear_attn.num_v_heads //= N
+                layer.linear_attn.key_dim //= N
+                layer.linear_attn.value_dim //= N
+                layer.linear_attn.conv_dim //= N
+
+            # Softmax attention
+            else:
+                layer.self_attn.o_proj = shard_linear(
+                    layer.self_attn.o_proj, "sharded-to-all", group=group
+                )
+                layer.self_attn.q_proj = shard_linear(
+                    layer.self_attn.q_proj, "all-to-sharded", group=group
+                )
+                repeat_kv_layer_inplace(
+                    layer.self_attn.k_proj, layer.self_attn.num_key_value_heads
+                )
+                repeat_kv_layer_inplace(
+                    layer.self_attn.v_proj, layer.self_attn.num_key_value_heads
+                )
+                layer.self_attn.k_proj = shard_linear(
+                    layer.self_attn.k_proj, "all-to-sharded", group=group
+                )
+                layer.self_attn.v_proj = shard_linear(
+                    layer.self_attn.v_proj, "all-to-sharded", group=group
+                )
+                layer.self_attn.num_attention_heads //= N
+                layer.self_attn.num_key_value_heads = max(
+                    1, layer.self_attn.num_key_value_heads // N
+                )
+
+            # MLP
+            if isinstance(layer.mlp, MLP):
+                layer.mlp.gate_proj = shard_linear(
+                    layer.mlp.gate_proj, "all-to-sharded", group=group
+                )
+                layer.mlp.down_proj = shard_linear(
+                    layer.mlp.down_proj, "sharded-to-all", group=group
+                )
+                layer.mlp.up_proj = shard_linear(
+                    layer.mlp.up_proj, "all-to-sharded", group=group
+                )
+
+            # MoE
+            else:
+                layer.mlp.sharding_group = group
+                shard_inplace(
+                    layer.mlp.shared_expert.gate_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.shared_expert.down_proj, "sharded-to-all", group=group
+                )
+                shard_inplace(
+                    layer.mlp.shared_expert.up_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.gate_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.down_proj, "sharded-to-all", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.up_proj, "all-to-sharded", group=group
+                )
+
+    @property
+    def layers(self):
+        return self.language_model.model.layers
+
+    def make_cache(self):
+        return self.language_model.make_cache()
+
+    @property
+    def quant_predicate(self):
+        return self.language_model.quant_predicate
+
+    @property
+    def cast_predicate(self):
+        return self.language_model.cast_predicate

mlx_lm/models/qwen3_5_moe.py

@@ -0,0 +1,52 @@
+# Copyright © 2026 Apple Inc.
+
+from dataclasses import dataclass
+
+from .base import BaseModelArgs
+from .qwen3_5 import Model as Qwen3_5Model
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    text_config: dict
+
+    @classmethod
+    def from_dict(cls, params):
+        if "text_config" not in params:
+            return cls(model_type=params["model_type"], text_config=params)
+        return super().from_dict(params)
+
+
+class Model(Qwen3_5Model):
+
+    def sanitize(self, weights):
+        new_weights = {}
+        for key, value in weights.items():
+            if key.startswith("vision_tower") or key.startswith("model.visual"):
+                continue
+            if key.startswith("model.language_model"):
+                key = key.replace("model.language_model", "language_model.model")
+            elif key.startswith("language_model."):
+                pass
+            else:
+                key = "language_model." + key
+            new_weights[key] = value
+
+        for l in range(self.language_model.args.num_hidden_layers):
+            prefix = f"language_model.model.layers.{l}.mlp"
+            gate_up_key = f"{prefix}.experts.gate_up_proj"
+            if gate_up_key in new_weights:
+                gate_up = new_weights.pop(gate_up_key)
+                mid = gate_up.shape[-2] // 2
+                new_weights[f"{prefix}.switch_mlp.gate_proj.weight"] = gate_up[
+                    ..., :mid, :
+                ]
+                new_weights[f"{prefix}.switch_mlp.up_proj.weight"] = gate_up[
+                    ..., mid:, :
+                ]
+                new_weights[f"{prefix}.switch_mlp.down_proj.weight"] = new_weights.pop(
+                    f"{prefix}.experts.down_proj"
+                )
+
+        return self.language_model.sanitize(new_weights)

mlx_lm/models/qwen3_moe.py

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

mlx_lm/models/qwen3_next.py

@@ -3,10 +3,12 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from functools import partial
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import mlx.core as mx
 import mlx.nn as nn
+from mlx.nn.layers.distributed import sum_gradients
 
 from .activations import swiglu
 from .base import (
@@ -15,7 +17,7 @@ from .base import (
     create_ssm_mask,
     scaled_dot_product_attention,
 )
-from .cache import KVCache, MambaCache
+from .cache import ArraysCache, KVCache
 from .gated_delta import gated_delta_update
 from .rope_utils import initialize_rope
 from .switch_layers import SwitchGLU
@@ -53,6 +55,13 @@ class ModelArgs(BaseModelArgs):
     full_attention_interval: int = 4
 
 
+@partial(mx.compile, shapeless=True)
+def _precise_swiglu(h, gate, x):
+    gate = nn.silu(gate.astype(mx.float32))
+    x = x.astype(mx.float32)
+    return (gate * x).astype(h.dtype)
+
+
 class Qwen3NextRMSNormGated(nn.Module):
     def __init__(self, hidden_size: int, eps: float = 1e-6):
         super().__init__()
@@ -64,8 +73,9 @@ class Qwen3NextRMSNormGated(nn.Module):
     ) -> mx.array:
         x = mx.fast.rms_norm(hidden_states, self.weight, self.eps)
         if gate is not None:
-            x = swiglu(gate, x)
-        return x
+            return _precise_swiglu(hidden_states, gate, x)
+        else:
+            return x.astype(hidden_states.dtype)
 
 
 class Qwen3NextAttention(nn.Module):
@@ -312,10 +322,15 @@ class Qwen3NextSparseMoeBlock(nn.Module):
         self.shared_expert = Qwen3NextMLP(dim, shared_expert_intermediate_size)
         self.shared_expert_gate = nn.Linear(dim, 1, bias=False)
 
+        self.sharding_group = None
+
     def __call__(
         self,
         x: mx.array,
     ) -> mx.array:
+        if self.sharding_group is not None:
+            x = sum_gradients(self.sharding_group)(x)
+
         gates = self.gate(x)
         gates = mx.softmax(gates, axis=-1, precise=True)
 
@@ -331,7 +346,12 @@ class Qwen3NextSparseMoeBlock(nn.Module):
         shared_y = self.shared_expert(x)
         shared_y = mx.sigmoid(self.shared_expert_gate(x)) * shared_y
 
-        return y + shared_y
+        y = y + shared_y
+
+        if self.sharding_group is not None:
+            y = mx.distributed.all_sum(y, group=self.sharding_group)
+
+        return y
 
 
 class Qwen3NextDecoderLayer(nn.Module):
@@ -427,7 +447,7 @@ class Model(nn.Module):
         return self.model.layers
 
     def make_cache(self):
-        return [MambaCache() if l.is_linear else KVCache() for l in self.layers]
+        return [ArraysCache(size=2) if l.is_linear else KVCache() for l in self.layers]
 
     def sanitize(self, weights):
         if "model.layers.0.mlp.experts.0.up_proj.weight" not in weights:

mlx_lm/models/recurrent_gemma.py

@@ -8,7 +8,7 @@ import mlx.core as mx
 import mlx.nn as nn
 
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
-from .cache import MambaCache, RotatingKVCache
+from .cache import ArraysCache, RotatingKVCache
 
 
 @dataclass
@@ -446,7 +446,7 @@ class Model(nn.Module):
         cache = []
         for layer in self.layers:
             if layer.temporal_block_type == "recurrent":
-                cache.append(MambaCache())
+                cache.append(ArraysCache(size=2))
             else:
                 cache.append(RotatingKVCache(max_size=self.args.attention_window_size))
         return cache

mlx_lm/models/rope_utils.py

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

mlx_lm/models/ssm.py

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

mlx_lm/models/step3p5.py

@@ -0,0 +1,512 @@
+# Copyright © 2026 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 mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients
+
+from .activations import swiglu
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .cache import KVCache, RotatingKVCache
+from .rope_utils import initialize_rope
+from .switch_layers import SwiGLU, SwitchGLU
+
+
+@partial(mx.compile, shapeless=True)
+def clamped_swiglu(x, gate, limit):
+    gate = mx.clip(nn.silu(gate), a_min=None, a_max=limit)
+    x = mx.clip(x, a_min=-limit, a_max=limit)
+    return gate * x
+
+
+class ClampedSwiGLU(nn.Module):
+    def __init__(self, limit: float):
+        super().__init__()
+        self.limit = limit
+
+    def __call__(self, x: mx.array, gate: mx.array) -> mx.array:
+        return clamped_swiglu(x, gate, self.limit)
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    num_hidden_layers: int
+    vocab_size: int
+    num_attention_heads: int
+    num_attention_groups: int
+    head_dim: int
+    intermediate_size: int
+    rms_norm_eps: float = 1e-5
+    rope_theta: float = 10000.0
+    rope_scaling: Optional[Dict] = None
+    max_position_embeddings: int = 262144
+    sliding_window: int = 512
+    layer_types: Optional[List[str]] = None
+    yarn_only_types: Optional[List[str]] = None
+    partial_rotary_factors: Optional[List[float]] = None
+    attention_other_setting: Optional[Dict] = None
+    use_head_wise_attn_gate: bool = True
+    moe_num_experts: int = 288
+    moe_top_k: int = 8
+    moe_intermediate_size: int = 1280
+    share_expert_dim: int = 1280
+    moe_layers_enum: Optional[str] = None
+    moe_router_scaling_factor: float = 3.0
+    norm_expert_weight: bool = True
+    swiglu_limits: Optional[List[float]] = None
+    swiglu_limits_shared: Optional[List[float]] = None
+    tie_word_embeddings: bool = False
+
+
+class ZeroCenteredRMSNorm(nn.Module):
+    def __init__(self, dims: int, eps: float = 1e-5):
+        super().__init__()
+        self.weight = mx.ones((dims,))
+        self.eps = eps
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return mx.fast.rms_norm(x, self.weight, self.eps)
+
+
+class Step3p5MLP(nn.Module):
+    def __init__(
+        self, args: ModelArgs, intermediate_size: int, swiglu_limit: float = 0
+    ):
+        super().__init__()
+        self.hidden_size = args.hidden_size
+        self.intermediate_size = intermediate_size
+
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+
+        self.limit = swiglu_limit if swiglu_limit and swiglu_limit > 0 else None
+
+    def __call__(self, x: mx.array) -> mx.array:
+        if self.limit is not None:
+            return self.down_proj(
+                clamped_swiglu(self.up_proj(x), self.gate_proj(x), self.limit)
+            )
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
+
+
+@mx.compile
+def moe_gate_select(gates, router_bias, top_k, routed_scaling_factor, norm_topk_prob):
+    scores = mx.sigmoid(gates.astype(mx.float32))
+    corrected_scores = scores + router_bias
+
+    topk_indices = mx.argpartition(-corrected_scores, kth=top_k - 1, axis=-1)[
+        ..., :top_k
+    ]
+    topk_weights = mx.take_along_axis(scores, topk_indices, axis=-1)
+
+    if norm_topk_prob:
+        topk_weights = topk_weights / (
+            mx.sum(topk_weights, axis=-1, keepdims=True) + 1e-20
+        )
+
+    return topk_indices, topk_weights * routed_scaling_factor
+
+
+class Step3p5MoEGate(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.top_k = args.moe_top_k
+        self.n_routed_experts = args.moe_num_experts
+        self.routed_scaling_factor = args.moe_router_scaling_factor
+        self.norm_topk_prob = args.norm_expert_weight
+
+        self.gate = nn.Linear(args.hidden_size, self.n_routed_experts, bias=False)
+        self.router_bias = mx.zeros((self.n_routed_experts,))
+
+    def __call__(self, x: mx.array):
+        return moe_gate_select(
+            self.gate(x),
+            self.router_bias,
+            self.top_k,
+            self.routed_scaling_factor,
+            self.norm_topk_prob,
+        )
+
+
+class Step3p5MoE(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+
+        swiglu_limit = 0
+        if args.swiglu_limits and layer_idx < len(args.swiglu_limits):
+            swiglu_limit = args.swiglu_limits[layer_idx] or 0
+
+        swiglu_limit_shared = 0
+        if args.swiglu_limits_shared and layer_idx < len(args.swiglu_limits_shared):
+            swiglu_limit_shared = args.swiglu_limits_shared[layer_idx] or 0
+
+        self.gate = Step3p5MoEGate(args)
+
+        activation = ClampedSwiGLU(swiglu_limit) if swiglu_limit > 0 else SwiGLU()
+        self.switch_mlp = SwitchGLU(
+            args.hidden_size,
+            args.moe_intermediate_size,
+            args.moe_num_experts,
+            activation=activation,
+        )
+
+        self.share_expert = Step3p5MLP(
+            args,
+            intermediate_size=args.share_expert_dim,
+            swiglu_limit=swiglu_limit_shared,
+        )
+
+        self.sharding_group = None
+
+    def __call__(self, x: mx.array) -> mx.array:
+        if self.sharding_group is not None:
+            x = sum_gradients(self.sharding_group)(x)
+
+        topk_indices, topk_weights = self.gate(x)
+        routed_output = self.switch_mlp(x, topk_indices)
+        routed_output = (
+            (routed_output * topk_weights[..., None])
+            .sum(axis=-2)
+            .astype(routed_output.dtype)
+        )
+        y = routed_output + self.share_expert(x)
+
+        if self.sharding_group is not None:
+            y = mx.distributed.all_sum(y, group=self.sharding_group)
+
+        return y
+
+
+class Step3p5Attention(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        dim = args.hidden_size
+
+        layer_types = args.layer_types or []
+        if layer_types:
+            self.is_sliding = layer_types[layer_idx] == "sliding_attention"
+        else:
+            self.is_sliding = layer_idx % 2 == 0
+
+        if self.is_sliding and args.attention_other_setting:
+            self.num_heads = args.attention_other_setting["num_attention_heads"]
+            self.num_kv_heads = args.attention_other_setting["num_attention_groups"]
+        else:
+            self.num_heads = args.num_attention_heads
+            self.num_kv_heads = args.num_attention_groups
+
+        self.head_dim = args.head_dim
+        self.scale = self.head_dim**-0.5
+
+        self.q_proj = nn.Linear(dim, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(dim, self.num_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(dim, self.num_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, dim, bias=False)
+
+        self.q_norm = ZeroCenteredRMSNorm(self.head_dim, eps=args.rms_norm_eps)
+        self.k_norm = ZeroCenteredRMSNorm(self.head_dim, eps=args.rms_norm_eps)
+
+        self.use_head_wise_attn_gate = args.use_head_wise_attn_gate
+        if self.use_head_wise_attn_gate:
+            self.g_proj = nn.Linear(dim, self.num_heads, bias=False)
+
+        rope_theta = args.rope_theta
+        if isinstance(rope_theta, list):
+            rope_theta = rope_theta[layer_idx]
+
+        partial_rotary_factor = 1.0
+        if args.partial_rotary_factors and layer_idx < len(args.partial_rotary_factors):
+            partial_rotary_factor = args.partial_rotary_factors[layer_idx]
+
+        rope_dims = int(self.head_dim * partial_rotary_factor)
+
+        yarn_only_types = args.yarn_only_types or []
+        layer_type = layer_types[layer_idx] if layer_types else "full_attention"
+        if yarn_only_types and layer_type not in yarn_only_types:
+            rope_scaling = None
+        else:
+            rope_scaling = args.rope_scaling
+
+        self.rope = initialize_rope(
+            dims=rope_dims,
+            base=rope_theta,
+            traditional=False,
+            scaling_config=rope_scaling,
+            max_position_embeddings=args.max_position_embeddings,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, _ = x.shape
+
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        queries = self.q_norm(queries.reshape(B, L, self.num_heads, -1)).transpose(
+            0, 2, 1, 3
+        )
+        keys = self.k_norm(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, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3)
+
+        if self.use_head_wise_attn_gate:
+            output = output * mx.sigmoid(self.g_proj(x))[..., None]
+
+        return self.o_proj(output.reshape(B, L, -1))
+
+
+class Step3p5DecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+
+        self.self_attn = Step3p5Attention(args, layer_idx)
+        self.is_sliding = self.self_attn.is_sliding
+
+        moe_layers_idx = set()
+        if args.moe_layers_enum:
+            moe_layers_idx = {int(i) for i in args.moe_layers_enum.strip().split(",")}
+        else:
+            moe_layers_idx = set(range(1, args.num_hidden_layers))
+
+        self.is_moe_layer = layer_idx in moe_layers_idx
+
+        if self.is_moe_layer:
+            self.mlp = Step3p5MoE(args, layer_idx)
+        else:
+            swiglu_limit = 0
+            if args.swiglu_limits_shared and layer_idx < len(args.swiglu_limits_shared):
+                swiglu_limit = args.swiglu_limits_shared[layer_idx] or 0
+            self.mlp = Step3p5MLP(
+                args,
+                intermediate_size=args.intermediate_size,
+                swiglu_limit=swiglu_limit,
+            )
+
+        self.input_layernorm = ZeroCenteredRMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+        self.post_attention_layernorm = ZeroCenteredRMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.self_attn(self.input_layernorm(x), mask=mask, cache=cache)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        return h + r
+
+
+class Step3p5Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.vocab_size = args.vocab_size
+        self.num_layers = args.num_hidden_layers
+
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            Step3p5DecoderLayer(args, layer_idx)
+            for layer_idx in range(args.num_hidden_layers)
+        ]
+        self.norm = ZeroCenteredRMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+        self._swa_idx = next(
+            (i for i, l in enumerate(self.layers) if l.is_sliding), None
+        )
+        self._full_idx = next(
+            (i for i, l in enumerate(self.layers) if not l.is_sliding), None
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        cache: Optional[List[Any]] = None,
+    ) -> mx.array:
+        h = self.embed_tokens(x)
+
+        if cache is None:
+            cache = [None] * self.num_layers
+
+        full_mask = None
+        swa_mask = None
+
+        if self._full_idx is not None:
+            full_mask = create_attention_mask(h, cache[self._full_idx])
+
+        if self._swa_idx is not None:
+            swa_mask = create_attention_mask(
+                h, cache[self._swa_idx], window_size=self.args.sliding_window
+            )
+
+        for layer, c in zip(self.layers, cache):
+            mask = swa_mask if layer.is_sliding else full_mask
+            h = layer(h, mask=mask, cache=c)
+
+        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 = Step3p5Model(args)
+        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[List[Any]] = None,
+    ):
+        out = self.model(inputs, cache)
+        return self.lm_head(out)
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    def make_cache(self):
+        return [
+            (
+                RotatingKVCache(max_size=self.args.sliding_window)
+                if layer.is_sliding
+                else KVCache()
+            )
+            for layer in self.layers
+        ]
+
+    def sanitize(self, weights):
+        remappings = [
+            (".moe.gate_proj.", ".mlp.switch_mlp.gate_proj."),
+            (".moe.up_proj.", ".mlp.switch_mlp.up_proj."),
+            (".moe.down_proj.", ".mlp.switch_mlp.down_proj."),
+            (".moe.gate.", ".mlp.gate.gate."),
+            (".moe.router_bias", ".mlp.gate.router_bias"),
+            (".share_expert.", ".mlp.share_expert."),
+        ]
+
+        is_vanilla = any(
+            src in k and dst not in k for k in weights for src, dst in remappings
+        )
+
+        new_weights = {}
+        for k, v in weights.items():
+            if ".mtp" in k:
+                continue
+            if "model.layers." in k:
+                parts = k.split(".")
+                if len(parts) > 2 and parts[2].isdigit():
+                    if int(parts[2]) >= self.args.num_hidden_layers:
+                        continue
+
+            for src, dst in remappings:
+                if src in k and dst not in k:
+                    k = k.replace(src, dst)
+                    break
+
+            if is_vanilla and k.endswith(".weight") and "norm" in k:
+                v = v + 1
+
+            new_weights[k] = v
+
+        return new_weights
+
+    @property
+    def cast_predicate(self):
+        def predicate(k):
+            return "router_bias" not in k
+
+        return predicate
+
+    @property
+    def quant_predicate(self):
+        def predicate(path, _):
+            if "mlp.gate.gate" in path:
+                return {"group_size": 64, "bits": 8}
+            return True
+
+        return predicate
+
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        group = group or mx.distributed.init()
+        N = group.size()
+
+        for layer in self.model.layers:
+            layer.self_attn.q_proj = shard_linear(
+                layer.self_attn.q_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.k_proj = shard_linear(
+                layer.self_attn.k_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.v_proj = shard_linear(
+                layer.self_attn.v_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.o_proj = shard_linear(
+                layer.self_attn.o_proj, "sharded-to-all", group=group
+            )
+            layer.self_attn.num_heads //= N
+            layer.self_attn.num_kv_heads //= N
+
+            if layer.self_attn.use_head_wise_attn_gate:
+                layer.self_attn.g_proj = shard_linear(
+                    layer.self_attn.g_proj, "all-to-sharded", group=group
+                )
+
+            if isinstance(layer.mlp, Step3p5MLP):
+                layer.mlp.gate_proj = shard_linear(
+                    layer.mlp.gate_proj, "all-to-sharded", group=group
+                )
+                layer.mlp.up_proj = shard_linear(
+                    layer.mlp.up_proj, "all-to-sharded", group=group
+                )
+                layer.mlp.down_proj = shard_linear(
+                    layer.mlp.down_proj, "sharded-to-all", group=group
+                )
+            else:
+                layer.mlp.sharding_group = group
+                shard_inplace(
+                    layer.mlp.share_expert.gate_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.share_expert.up_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.share_expert.down_proj, "sharded-to-all", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.gate_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.up_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.down_proj, "sharded-to-all", group=group
+                )

mlx_lm/models/telechat3.py

@@ -0,0 +1,202 @@
+# Copyright © 2026 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .activations import swiglu
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .rope_utils import initialize_rope
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    intermediate_size: int
+    max_position_embeddings: int
+    num_attention_heads: int
+    num_hidden_layers: int
+    num_key_value_heads: int
+    rms_norm_eps: float
+    vocab_size: int
+    rope_theta: float
+    mlp_bias: bool = False
+    attention_bias: bool = False
+    head_dim: Optional[int] = None
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    tie_word_embeddings: bool = False
+
+
+class Telechat3Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_attention_heads = args.num_attention_heads
+        self.num_key_value_heads = args.num_key_value_heads
+
+        self.head_dim = args.head_dim or args.hidden_size // args.num_attention_heads
+        self.scale = self.head_dim**-0.5
+
+        self.q_proj = nn.Linear(
+            args.hidden_size,
+            args.num_attention_heads * self.head_dim,
+            bias=args.attention_bias,
+        )
+        self.k_proj = nn.Linear(
+            args.hidden_size,
+            args.num_key_value_heads * self.head_dim,
+            bias=args.attention_bias,
+        )
+        self.v_proj = nn.Linear(
+            args.hidden_size,
+            args.num_key_value_heads * self.head_dim,
+            bias=args.attention_bias,
+        )
+        self.o_proj = nn.Linear(
+            args.num_attention_heads * self.head_dim,
+            args.hidden_size,
+            bias=args.attention_bias,
+        )
+
+        self.rope = initialize_rope(
+            self.head_dim,
+            traditional=False,
+            base=args.rope_theta,
+            scaling_config=args.rope_scaling,
+            max_position_embeddings=args.max_position_embeddings,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, _ = x.shape
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        queries = queries.reshape(B, L, self.num_attention_heads, -1).transpose(
+            0, 2, 1, 3
+        )
+        keys = keys.reshape(B, L, self.num_key_value_heads, -1).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.num_key_value_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, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class Telechat3MLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.gate_proj = nn.Linear(
+            args.hidden_size, args.intermediate_size, bias=args.mlp_bias
+        )
+        self.down_proj = nn.Linear(
+            args.intermediate_size, args.hidden_size, bias=args.mlp_bias
+        )
+        self.up_proj = nn.Linear(
+            args.hidden_size, args.intermediate_size, bias=args.mlp_bias
+        )
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
+
+
+class Telechat3DecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_attention_heads = args.num_attention_heads
+        self.hidden_size = args.hidden_size
+        self.self_attn = Telechat3Attention(args)
+        self.mlp = Telechat3MLP(args)
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+        self.args = args
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+        h = x + r
+        out = h + self.mlp(self.post_attention_layernorm(h))
+        return out
+
+
+class Telechat3Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.vocab_size = args.vocab_size
+        self.num_hidden_layers = args.num_hidden_layers
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            Telechat3DecoderLayer(args=args) for _ in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+        input_embeddings: Optional[mx.array] = None,
+    ) -> mx.array:
+        if input_embeddings is not None:
+            h = input_embeddings
+        else:
+            h = self.embed_tokens(inputs)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+        mask = create_attention_mask(h, cache[0])
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, c)
+
+        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 = Telechat3Model(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+        input_embeddings: Optional[mx.array] = None,
+    ) -> mx.array:
+        out = self.model(inputs, cache, input_embeddings)
+        if self.args.tie_word_embeddings:
+            out = self.model.embed_tokens.as_linear(out)
+        else:
+            out = self.lm_head(out)
+        return out
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/perplexity.py

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

mlx_lm/quant/dwq.py

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

mlx_lm/sample_utils.py

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

mlx_lm/share.py

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

mlx_lm/tokenizer_utils.py

@@ -91,7 +91,11 @@ class NaiveStreamingDetokenizer(StreamingDetokenizer):
     def text(self):
         if self._current_tokens:
             self._current_text = self._tokenizer.decode(self._current_tokens)
-            if self._current_text.endswith("\ufffd"):
+            if self._current_text.endswith("\ufffd") or (
+                self._tokenizer.clean_up_tokenization_spaces
+                and len(self._current_text) > 0
+                and self._current_text[-1] == " "
+            ):
                 self._current_text = self._current_text[:-1]
         if self._current_text and self._current_text[-1] == "\n":
             self._text += self._current_text
@@ -159,6 +163,8 @@ class BPEStreamingDetokenizer(StreamingDetokenizer):
     _space_matches = (".", "?", "!", ",", "n't", "'m", "'s", "'ve", "'re")
 
     def __init__(self, tokenizer):
+        self.clean_spaces = tokenizer.clean_up_tokenization_spaces
+
         # Extract the tokens in a list from id to text
         self.tokenmap = [None] * len(tokenizer.vocab)
         for value, tokenid in tokenizer.vocab.items():
@@ -193,6 +199,8 @@ class BPEStreamingDetokenizer(StreamingDetokenizer):
             return current_text
         elif not self.text:
             return current_text[1:]
+        elif self.clean_spaces and current_text[1:].startswith(self._space_matches):
+            return current_text[1:]
         return current_text
 
     def add_token(self, token):
@@ -202,7 +210,10 @@ class BPEStreamingDetokenizer(StreamingDetokenizer):
         text = self._decode_bytes(self._unflushed)
 
         # For multi-byte utf-8 wait until they are complete
-        if not text.endswith("\ufffd"):
+        # For single spaces wait until the next token to clean it if needed
+        if not text.endswith("\ufffd") and not (
+            len(v) == 1 and self._byte_decoder.get(v[0]) == 32
+        ):
             self.text += self._maybe_trim_space(text)
             self._unflushed = ""
 
@@ -280,7 +291,10 @@ class TokenizerWrapper:
         self._tool_call_end = tool_call_end
 
         vocab = tokenizer.get_vocab()
-        THINK_TOKENS = [("<think>", "</think>")]
+        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
@@ -461,10 +475,21 @@ def _infer_tool_parser(chat_template):
         return "minimax_m2"
     elif "<start_function_call>" in chat_template:
         return "function_gemma"
+    elif "<longcat_tool_call>" in chat_template:
+        return "longcat"
     elif "<arg_key>" in chat_template:
         return "glm47"
-    elif "<tool_call>\n<function=" in chat_template:
+    elif "<|tool_list_start|>" in chat_template:
+        return "pythonic"
+    elif (
+        "<tool_call>\\n<function=" in chat_template
+        or "<tool_call>\n<function=" in chat_template
+    ):
         return "qwen3_coder"
+    elif "<|tool_calls_section_begin|>" in chat_template:
+        return "kimi_k2"
+    elif "[TOOL_CALLS]" in chat_template:
+        return "mistral"
     elif "<tool_call>" in chat_template and "tool_call.name" in chat_template:
         return "json_tools"
     return None

mlx_lm/tool_parsers/glm47.py

@@ -7,6 +7,7 @@ https://github.com/vllm-project/vllm/blob/main/vllm/tool_parsers/glm4_moe_tool_p
 
 import ast
 import json
+import shlex
 from typing import Any
 
 import regex as re
@@ -21,22 +22,21 @@ tool_call_start = "<tool_call>"
 tool_call_end = "</tool_call>"
 
 
-def _is_string_type(
-    tool_name: str,
-    arg_name: str,
-    tools: list[Any] | None,
-) -> bool:
+def _get_string_arg_names(tool_name: str, tools: list[Any] | None) -> set[str]:
     if tools is None:
-        return False
+        return set()
     for tool in tools:
-        func = tool["function"]
-        if func["name"] == tool_name:
-            params = func["parameters"]
-            if params is None:
-                return False
-            arg_type = params.get("properties", {}).get(arg_name, {}).get("type", None)
-            return arg_type == "string"
-    return False
+        func = tool.get("function")
+        if not func or func.get("name") != tool_name:
+            continue
+        params = func.get("parameters") or {}
+        properties = params.get("properties") or {}
+        return {
+            name
+            for name, schema in properties.items()
+            if schema.get("type") == "string"
+        }
+    return set()
 
 
 def _deserialize(value: str) -> Any:
@@ -52,14 +52,179 @@ def _deserialize(value: str) -> Any:
     return value
 
 
+# Normalize argument values based on tool schema types.
+def _normalize_arguments(
+    func_name: str,
+    arguments: dict[str, Any],
+    tools: list[Any] | None,
+    string_args: set[str] | None = None,
+) -> dict[str, Any]:
+    if string_args is None:
+        string_args = _get_string_arg_names(func_name, tools)
+    normalized = {}
+    for key, value in arguments.items():
+        # Preserve declared string types; coerce others when values are strings.
+        if key in string_args:
+            normalized[key] = value if isinstance(value, str) else str(value)
+            continue
+        if isinstance(value, str):
+            normalized[key] = _deserialize(value)
+        else:
+            normalized[key] = value
+    return normalized
+
+
+# Parse JSON tool call payloads used by some GLM outputs.
+def _parse_json_tool_call(text: str, tools: list[Any] | None):
+    try:
+        parsed = json.loads(text.strip())
+    except Exception:
+        return None
+
+    if isinstance(parsed, list) and parsed:
+        if isinstance(parsed[0], dict):
+            parsed = parsed[0]
+    if not isinstance(parsed, dict):
+        return None
+
+    # Pull out name/arguments from known JSON shapes.
+    name = None
+    arguments = None
+    if "name" in parsed and "arguments" in parsed:
+        name = parsed.get("name")
+        arguments = parsed.get("arguments")
+    elif "function" in parsed and "arguments" in parsed:
+        name = parsed.get("function")
+        arguments = parsed.get("arguments")
+    elif "tool" in parsed and isinstance(parsed.get("tool"), dict):
+        tool = parsed["tool"]
+        name = tool.get("name")
+        arguments = tool.get("arguments")
+
+    if isinstance(name, dict):
+        arguments = arguments or name.get("arguments")
+        name = name.get("name")
+
+    if isinstance(arguments, str):
+        arguments = _deserialize(arguments)
+
+    string_args = _get_string_arg_names(name, tools) if isinstance(name, str) else None
+
+    if isinstance(name, str) and arguments is None:
+        return dict(name=name, arguments={})
+    if isinstance(name, str) and isinstance(arguments, dict):
+        return dict(
+            name=name,
+            arguments=_normalize_arguments(
+                name, arguments, tools, string_args=string_args
+            ),
+        )
+
+    return None
+
+
+# Parse key=value tokens into an arguments dict.
+def _parse_key_value_pairs(
+    text: str,
+    func_name: str,
+    tools: list[Any] | None,
+    string_args: set[str] | None = None,
+) -> dict[str, Any] | None:
+    if "=" not in text:
+        return None
+    try:
+        tokens = shlex.split(text)
+    except ValueError:
+        return None
+    if not tokens:
+        return None
+
+    if string_args is None:
+        string_args = _get_string_arg_names(func_name, tools)
+
+    arguments = {}
+    for token in tokens:
+        # Require key=value tokens to avoid mis-parsing freeform text.
+        if "=" not in token:
+            return None
+        key, value = token.split("=", 1)
+        key = key.strip()
+        if not key:
+            return None
+        if key in string_args:
+            arguments[key] = value
+        else:
+            arguments[key] = _deserialize(value)
+    return arguments
+
+
+# Parse plain text tool calls like "name a=1 b=2" or "name {json}".
+def _parse_plain_text_tool_call(text: str, tools: list[Any] | None):
+    stripped = text.strip()
+    if not stripped:
+        return None
+
+    # Handle "name\\n{...}" style payloads.
+    if "\n" in stripped:
+        first_line, rest = stripped.split("\n", 1)
+        name = first_line.strip()
+        rest = rest.strip()
+        if name and rest:
+            string_args = _get_string_arg_names(name, tools)
+            arguments = _deserialize(rest)
+            if isinstance(arguments, dict):
+                return dict(
+                    name=name,
+                    arguments=_normalize_arguments(
+                        name, arguments, tools, string_args=string_args
+                    ),
+                )
+
+    # Split on whitespace to get name + arguments segment.
+    name, _, rest = stripped.partition(" ")
+    if not name:
+        return None
+    rest = rest.strip()
+    if not rest:
+        return dict(name=name, arguments={})
+
+    string_args = _get_string_arg_names(name, tools)
+    arguments = _deserialize(rest)
+    if isinstance(arguments, dict):
+        return dict(
+            name=name,
+            arguments=_normalize_arguments(
+                name, arguments, tools, string_args=string_args
+            ),
+        )
+
+    kv_arguments = _parse_key_value_pairs(rest, name, tools, string_args=string_args)
+    if kv_arguments is not None:
+        return dict(name=name, arguments=kv_arguments)
+
+    return dict(name=name, arguments={"raw": rest})
+
+
 def parse_tool_call(text: str, tools: list[Any] | None = None):
-    func_name = _func_name_regex.search(text).group(1)
-    pairs = _func_arg_regex.findall(text)
+    """Parse a GLM 4.7 tool call string into a name and arguments dict."""
+    match = _func_name_regex.search(text)
+    if not match:
+        # Fallbacks for alternate formats seen in GLM tool calls.
+        fallback = _parse_json_tool_call(text, tools)
+        if fallback is not None:
+            return fallback
+        fallback = _parse_plain_text_tool_call(text, tools)
+        if fallback is not None:
+            return fallback
+        return dict(name="unknown", arguments={"raw": text.strip()})
+
+    func_name = match.group(1)
+    string_args = _get_string_arg_names(func_name, tools)
     arg_dct = {}
-    for key, value in pairs:
-        arg_key = key.strip()
-        arg_val = value.strip()
-        if not _is_string_type(func_name, arg_key, tools):
+    for match in _func_arg_regex.finditer(text):
+        arg_key = match.group(1).strip()
+        arg_val = match.group(2).strip()
+        if arg_key not in string_args:
             arg_val = _deserialize(arg_val)
         arg_dct[arg_key] = arg_val
     return dict(name=func_name, arguments=arg_dct)

mlx_lm/tool_parsers/kimi_k2.py

@@ -0,0 +1,61 @@
+# Copyright © 2026 Apple Inc.
+"""
+Modified from:
+https://github.com/vllm-project/vllm/blob/main/vllm/tool_parsers/kimi_k2_tool_parser.py
+"""
+
+import ast
+import json
+from typing import Any
+
+import regex as re
+
+# kimi has a fixed function naming scheme, with a json formatted arg
+#   functions.multiply:0<|tool_call_argument_begin|>{"a": 2, "b": 3}
+_func_name_regex = re.compile(
+    r"^\s*((?:functions\.)?(.+?):\d+)\s*<\|tool_call_argument_begin\|>", re.DOTALL
+)
+_func_arg_regex = re.compile(r"<\|tool_call_argument_begin\|>\s*(.*)\s*", re.DOTALL)
+_tool_call_split_regex = re.compile(
+    r"<\|tool_call_begin\|>(.*?)<\|tool_call_end\|>", re.DOTALL
+)
+
+tool_call_start = "<|tool_calls_section_begin|>"
+tool_call_end = "<|tool_calls_section_end|>"
+
+
+def _deserialize(value: str) -> Any:
+    try:
+        return json.loads(value)
+    except Exception:
+        pass
+
+    try:
+        return ast.literal_eval(value)
+    except Exception:
+        pass
+    return value
+
+
+def _parse_single_tool(text: str) -> dict:
+    func_name_match = _func_name_regex.search(text)
+    if func_name_match is None:
+        raise ValueError("No tool call found.")
+    tool_call_id = func_name_match.group(1)  # e.g. "functions.get_weather:0"
+    func_name = func_name_match.group(2)  # e.g. "get_weather"
+
+    func_args_match = _func_arg_regex.search(text)
+    if func_args_match is None:
+        raise ValueError("No tool call arguments found.")
+    func_args = func_args_match.group(1)
+    arg_dct = _deserialize(func_args)
+
+    return dict(id=tool_call_id, name=func_name, arguments=arg_dct)
+
+
+def parse_tool_call(text: str, tools: Any | None = None):
+    tool_matches = _tool_call_split_regex.findall(text)
+    if tool_matches:
+        return [_parse_single_tool(match) for match in tool_matches]
+    else:
+        return [_parse_single_tool(text)]

mlx_lm/tool_parsers/longcat.py

@@ -0,0 +1,68 @@
+# Copyright © 2026 Apple Inc.
+
+import ast
+import json
+from typing import Any
+
+import regex as re
+
+_func_name_regex = re.compile(r"^(.*?)<longcat_arg_key>", re.DOTALL)
+_func_arg_regex = re.compile(
+    r"<longcat_arg_key>(.*?)</longcat_arg_key>(?:\\n|\s)*<longcat_arg_value>(.*?)</longcat_arg_value>",
+    re.DOTALL,
+)
+
+tool_call_start = "<longcat_tool_call>"
+tool_call_end = "</longcat_tool_call>"
+
+
+def _is_string_type(
+    tool_name: str,
+    arg_name: str,
+    tools: list[Any] | None,
+) -> bool:
+    if tools is None:
+        return False
+    for tool in tools:
+        func = tool["function"]
+        if func["name"] == tool_name:
+            params = func["parameters"]
+            if params is None:
+                return False
+            arg_type = params.get("properties", {}).get(arg_name, {}).get("type", None)
+            return arg_type == "string"
+    return False
+
+
+def _deserialize(value: str) -> Any:
+    try:
+        return json.loads(value)
+    except Exception:
+        pass
+
+    try:
+        return ast.literal_eval(value)
+    except Exception:
+        pass
+    return value
+
+
+def parse_tool_call(text: str, tools: list[Any] | None = None):
+    text = text.strip()
+
+    if text.startswith("{"):
+        try:
+            return json.loads(text)
+        except json.JSONDecodeError:
+            pass
+
+    func_name = _func_name_regex.search(text).group(1).strip()
+    pairs = _func_arg_regex.findall(text)
+    arg_dct = {}
+    for key, value in pairs:
+        arg_key = key.strip()
+        arg_val = value.strip()
+        if not _is_string_type(func_name, arg_key, tools):
+            arg_val = _deserialize(arg_val)
+        arg_dct[arg_key] = arg_val
+    return dict(name=func_name, arguments=arg_dct)

mlx_lm/tool_parsers/mistral.py

@@ -0,0 +1,20 @@
+# Copyright © 2026 Apple Inc.
+
+import json
+from typing import Any
+
+import regex as re
+
+_tool_call_regex = re.compile(r"\s*(\w+)\[ARGS\]\s*(\{.*\})", re.DOTALL)
+
+tool_call_start = "[TOOL_CALLS]"
+tool_call_end = ""
+
+
+def parse_tool_call(text: str, tools: Any | None = None):
+    match = _tool_call_regex.search(text)
+    if match is None:
+        raise ValueError(f"Could not parse tool call from: {text}")
+    func_name = match.group(1)
+    func_args = json.loads(match.group(2))
+    return dict(name=func_name, arguments=func_args)

mlx_lm/tool_parsers/pythonic.py

@@ -0,0 +1,49 @@
+# Copyright © 2026 Apple Inc.
+
+import ast
+from typing import Any, Dict, List
+
+import regex as re
+
+"""
+Tool parser for Pythonic function call formats.
+
+Parses assistant responses containing tool calls in formats like:
+<|tool_call_start|>[function_name(arg1="value1", arg2=2)]<|tool_call_end|>
+"""
+
+
+_tool_call_regex = re.compile(r"\[(\w+)\((.*?)\)\]", re.DOTALL)
+_tool_args_regex = re.compile(r'(\w+)=(?:"([^"]*)"|([^,]+))(?:,\s*|$)', re.DOTALL)
+
+
+def parse_tool_call(text: str, tools: Any | None = None):
+    match = _tool_call_regex.search(text)
+    if not match:
+        raise ValueError("No function provided.")
+
+    func_name = match.group(1)
+    args_str = match.group(2)
+
+    arguments = {}
+    if args_str:
+        matches = _tool_args_regex.findall(args_str)
+        for pair in matches:
+            key = pair[0].strip()
+            # pair[1] is quoted value, pair[2] is unquoted value
+            value = pair[1] if pair[1] else pair[2].strip()
+
+            # Try to parse the value using ast.literal_eval
+            try:
+                value = ast.literal_eval(value)
+            except (ValueError, SyntaxError):
+                # If parsing fails, keep as string
+                pass
+
+            arguments[key] = value
+
+    return dict(name=func_name, arguments=arguments)
+
+
+tool_call_start = "<|tool_call_start|>"
+tool_call_end = "<|tool_call_end|>"

mlx_lm/tool_parsers/qwen3_coder.py

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

mlx_lm/tuner/datasets.py

@@ -116,7 +116,7 @@ class CompletionsDataset:
         if self.mask_prompt:
             offset = len(
                 self.tokenizer.apply_chat_template(
-                    messages[0],
+                    messages[:-1],
                     tools=tools,
                     add_generation_prompt=True,
                     return_dict=False,
@@ -322,8 +322,8 @@ def load_dataset(args, tokenizer: PreTrainedTokenizer):
             "Training set not found or empty. Must provide training set for fine-tuning."
         )
     if args.train and len(valid) == 0:
-        raise ValueError(
-            "Validation set not found or empty. Must provide validation set for fine-tuning."
+        print(
+            "Warning: Validation set not found or empty. Training will proceed without validation."
         )
     if args.test and len(test) == 0:
         raise ValueError(

mlx_lm/tuner/trainer.py

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

mlx_lm/utils.py

@@ -47,6 +47,7 @@ MODEL_REMAPPING = {
     "llava": "mistral3",
     "phi-msft": "phixtral",
     "falcon_mamba": "mamba",
+    "joyai_llm_flash": "deepseek_v3",
     "kimi_k2": "deepseek_v3",
     "qwen2_5_vl": "qwen2_vl",
     "minimax_m2": "minimax",
@@ -56,6 +57,18 @@ MODEL_REMAPPING = {
 MAX_FILE_SIZE_GB = 5
 
 
+def _parse_size(x):
+    sizes = {"M": 1e6, "G": 1e9, "MB": 1e6, "GB": 1e9, "": 1}
+    split = 0
+    for xi in x:
+        if not (xi.isdigit() or xi == "."):
+            break
+        split += 1
+    digits = float(x[:split])
+    size = (x[split:]).strip().upper()
+    return int(digits * sizes[size])
+
+
 def _unpack_awq_weights(qweight: mx.array) -> mx.array:
     bits = 4
     pack_factor = 32 // bits
@@ -303,16 +316,30 @@ def load_model(
     weight_files = glob.glob(str(model_path / "model*.safetensors"))
 
     if not weight_files and strict:
-        logging.error(f"No safetensors found in {model_path}")
         raise FileNotFoundError(f"No safetensors found in {model_path}")
 
     weights = {}
     for wf in weight_files:
         weights.update(mx.load(wf))
 
-    model_class, model_args_class = get_model_classes(config=config)
+    if (model_file := config.get("model_file")) is not None:
+        spec = importlib.util.spec_from_file_location(
+            "custom_model",
+            model_path / model_file,
+        )
+        arch = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(arch)
+        model_class, model_args_class = arch.Model, arch.ModelArgs
+    else:
+        model_class, model_args_class = get_model_classes(config=config)
+
+    if "quantization_config" not in config:
+        text_config = config.get("text_config", {})
+        if "quantization_config" in text_config:
+            config["quantization_config"] = text_config["quantization_config"]
 
     model_args = model_args_class.from_dict(config)
+
     model = model_class(model_args)
 
     if hasattr(model, "sanitize"):
@@ -362,12 +389,34 @@ def load_model(
             config["quantization_config"] = quantization
             _quantize(quantization)
 
+    if config.get("quantize_activations", False):
+
+        def _maybe_qq(m):
+            if isinstance(m, nn.QuantizedLinear):
+                if m.mode not in ("nvfp4", "mxfp8"):
+                    raise ValueError(
+                        "Mode ({m.mode}) does not support activation quantization"
+                    )
+                if m.get("bias", False):
+                    raise ValueError(
+                        "Linear layer with bias does not support activation quantization"
+                    )
+                out_dims, in_dims = m.weight.shape
+                in_dims *= 32 // m.bits
+                return nn.QQLinear(in_dims, out_dims, m.group_size, m.bits, m.mode)
+            else:
+                return m
+
+        leaves = tree_map(_maybe_qq, model.leaf_modules(), is_leaf=nn.Module.is_module)
+
+        model.update_modules(leaves)
+
+    model.eval()
     model.load_weights(list(weights.items()), strict=strict)
 
     if not lazy:
         mx.eval(model.parameters())
 
-    model.eval()
     return model, config
 
 
@@ -478,7 +527,7 @@ def sharded_load(
     # weights we need to download.
     model, config = load_model(model_path, lazy=True, strict=False)
 
-    has_pipelining = hasattr(model.model, "pipeline")
+    has_pipelining = hasattr(model, "model") and hasattr(model.model, "pipeline")
     has_tensor_parallel = hasattr(model, "shard")
 
     if pipeline_group is not None and not has_pipelining:

setup.py

@@ -10,7 +10,7 @@ sys.path.append(str(package_dir))
 
 from _version import __version__
 
-MIN_MLX_VERSION = "0.30.3"
+MIN_MLX_VERSION = "0.30.4"
 
 setup(
     name="mlx-lm",
@@ -26,7 +26,7 @@ setup(
     install_requires=[
         f"mlx>={MIN_MLX_VERSION}; platform_system == 'Darwin'",
         "numpy",
-        "transformers==5.0.0rc1",
+        "transformers>=5.0.0",
         "sentencepiece",
         "protobuf",
         "pyyaml",
@@ -51,6 +51,7 @@ setup(
     },
     entry_points={
         "console_scripts": [
+            "mlx_lm = mlx_lm.cli:main",
             "mlx_lm.awq = mlx_lm.quant.awq:main",
             "mlx_lm.dwq = mlx_lm.quant.dwq:main",
             "mlx_lm.dynamic_quant = mlx_lm.quant.dynamic_quant:main",
@@ -65,6 +66,7 @@ setup(
             "mlx_lm.lora = mlx_lm.lora:main",
             "mlx_lm.perplexity = mlx_lm.perplexity:main",
             "mlx_lm.server = mlx_lm.server:main",
+            "mlx_lm.share = mlx_lm.share:main",
             "mlx_lm.manage = mlx_lm.manage:main",
             "mlx_lm.upload = mlx_lm.upload:main",
         ]

tests/model_parallel_tests.py

@@ -0,0 +1,114 @@
+# Copyright © 2026 Apple Inc.
+
+import importlib
+import unittest
+
+import mlx.core as mx
+
+import mlx_lm
+
+
+class TestModelParallel(unittest.TestCase):
+
+    def test_shard(self):
+        test_configs = [
+            {
+                "model_type": "deepseek_v3",
+                "vocab_size": 1024,
+                "hidden_size": 128,
+                "intermediate_size": 256,
+                "moe_intermediate_size": 256,
+                "num_hidden_layers": 4,
+                "num_attention_heads": 4,
+                "num_key_value_heads": 2,
+                "n_routed_experts": 4,
+                "n_group": 2,
+                "topk_group": 1,
+                "num_experts_per_tok": 2,
+                "n_shared_experts": 1,
+                "kv_lora_rank": 4,
+                "q_lora_rank": 4,
+                "qk_rope_head_dim": 32,
+                "v_head_dim": 16,
+                "qk_nope_head_dim": 32,
+                "rope_scaling": {
+                    "beta_fast": 32,
+                    "beta_slow": 1,
+                    "factor": 40,
+                    "mscale": 1.0,
+                    "mscale_all_dim": 1.0,
+                    "original_max_position_embeddings": 4096,
+                    "type": "yarn",
+                },
+            },
+            {
+                "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_type": "glm4_moe_lite",
+                "vocab_size": 1000,
+                "hidden_size": 64,
+                "intermediate_size": 128,
+                "moe_intermediate_size": 32,
+                "num_hidden_layers": 4,
+                "num_attention_heads": 4,
+                "num_key_value_heads": 4,
+                "n_shared_experts": 1,
+                "n_routed_experts": 4,
+                "routed_scaling_factor": 1.0,
+                "kv_lora_rank": 8,
+                "q_lora_rank": 8,
+                "qk_rope_head_dim": 8,
+                "qk_nope_head_dim": 16,
+                "v_head_dim": 8,
+                "topk_method": "noaux_tc",
+                "scoring_func": "sigmoid",
+                "norm_topk_prob": True,
+                "n_group": 1,
+                "topk_group": 1,
+                "num_experts_per_tok": 2,
+                "moe_layer_freq": 1,
+                "first_k_dense_replace": 1,
+                "max_position_embeddings": 256,
+                "rms_norm_eps": 1e-5,
+                "rope_theta": 1000,
+                "rope_scaling": None,
+                "attention_bias": False,
+                "partial_rotary_factor": 1.0,
+                "tie_word_embeddings": False,
+                "num_nextn_predict_layers": 1,
+            },
+        ]
+        mx.random.seed(0)
+        for config in test_configs:
+            model_type = config["model_type"]
+            with self.subTest(f"Testing {model_type}", model_type=model_type):
+                arch = importlib.import_module(f"mlx_lm.models.{model_type}")
+                args = arch.ModelArgs.from_dict(config)
+                model = arch.Model(args)
+                vocab_size = args.vocab_size
+                x = mx.random.randint(0, vocab_size, shape=(32, 4))
+                expected = model(x)
+                model.shard()
+                out = model(x)
+                self.assertTrue(mx.allclose(expected, out, rtol=1e-3, atol=1e-3))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_datsets.py

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

tests/test_losses.py

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

tests/test_models.py

@@ -10,7 +10,11 @@ from mlx.utils import tree_map
 from mlx_lm.models import rope_utils
 from mlx_lm.models.base import create_causal_mask, scaled_dot_product_attention
 from mlx_lm.models.cache import KVCache, RotatingKVCache, make_prompt_cache
-from mlx_lm.models.gated_delta import gated_delta_kernel, gated_delta_ops
+from mlx_lm.models.gated_delta import (
+    gated_delta_kernel,
+    gated_delta_ops,
+    gated_delta_update,
+)
 from mlx_lm.models.ssm import ssm_attn, ssm_update
 
 
@@ -238,6 +242,30 @@ class TestModels(unittest.TestCase):
         )
         self.assertTrue(isinstance(rope, rope_utils.Llama3RoPE))
 
+    def test_su_scaled_rope_no_mutation(self):
+        rope = rope_utils.SuScaledRoPE(
+            dims=8,
+            max_position_embeddings=131072,
+            original_max_position_embeddings=4096,
+            long_factor=[1.0] * 4,
+        )
+        x = mx.ones((1, 2, 4, 8))
+        rope(x)
+        mx.eval(x)
+        self.assertTrue((x == 1).all())
+
+    def test_yarn_rope_no_mutation(self):
+        rope = rope_utils.YarnRoPE(
+            dims=8,
+            scaling_factor=2.0,
+            mscale=1.0,
+            mscale_all_dim=0,
+        )
+        x = mx.ones((1, 2, 4, 8))
+        rope(x)
+        mx.eval(x)
+        self.assertTrue((x == 1).all())
+
     def test_quantized_sdpa(self):
         cache = KVCache()
 
@@ -531,6 +559,59 @@ class TestModels(unittest.TestCase):
             model, args.model_type, args.vocab_size, args.num_hidden_layers
         )
 
+    def test_qwen3_5_family_convert_then_load_norm_not_shift_twice(self):
+        text_config = {
+            "hidden_size": 8,
+            "intermediate_size": 16,
+            "num_hidden_layers": 1,
+            "num_attention_heads": 1,
+            "num_key_value_heads": 1,
+            "rms_norm_eps": 1e-5,
+            "vocab_size": 32,
+            "linear_num_value_heads": 1,
+            "linear_num_key_heads": 1,
+            "linear_key_head_dim": 4,
+            "linear_value_head_dim": 4,
+            "linear_conv_kernel_dim": 1,
+            "full_attention_interval": 1,
+            "tie_word_embeddings": False,
+            "max_position_embeddings": 64,
+        }
+        hf_norm_key = "model.language_model.layers.0.input_layernorm.weight"
+        mlx_norm_key = "language_model.model.layers.0.input_layernorm.weight"
+
+        for model_type, hf_mtp_key in (
+            ("qwen3_5", "mtp.fc.weights"),
+            ("qwen3_5_moe", "mtp.fc.weight"),
+        ):
+            module = importlib.import_module(f"mlx_lm.models.{model_type}")
+            args = module.ModelArgs.from_dict(
+                {
+                    "model_type": model_type,
+                    "text_config": {"model_type": model_type, **text_config},
+                }
+            )
+            model = module.Model(args)
+
+            base = mx.arange(8, dtype=mx.float32)
+
+            # Simulate convert sanitize on HF-style keys.
+            converted = model.sanitize(
+                {
+                    hf_norm_key: base,
+                    hf_mtp_key: mx.zeros((1,), dtype=mx.float32),
+                }
+            )
+            self.assertIn(mlx_norm_key, converted)
+            self.assertTrue(mx.array_equal(converted[mlx_norm_key], base + 1.0))
+            self.assertFalse(any("mtp." in k for k in converted))
+
+            # Simulate load sanitize on already-converted keys.
+            loaded = model.sanitize(converted)
+            self.assertTrue(
+                mx.array_equal(loaded[mlx_norm_key], converted[mlx_norm_key])
+            )
+
     def test_qwen2_moe(self):
         from mlx_lm.models import qwen2_moe
 
@@ -671,6 +752,142 @@ class TestModels(unittest.TestCase):
             model, args.model_type, args.vocab_size, args.num_hidden_layers
         )
 
+    def test_step3p5(self):
+        from mlx_lm.models import step3p5
+
+        args = step3p5.ModelArgs(
+            model_type="step3p5",
+            hidden_size=256,
+            num_hidden_layers=4,
+            vocab_size=1024,
+            num_attention_heads=4,
+            num_attention_groups=2,
+            head_dim=64,
+            intermediate_size=512,
+            rms_norm_eps=1e-5,
+            rope_theta=[10000.0, 10000.0, 10000.0, 10000.0],
+            sliding_window=64,
+            layer_types=[
+                "full_attention",
+                "sliding_attention",
+                "sliding_attention",
+                "full_attention",
+            ],
+            partial_rotary_factors=[0.5, 1.0, 1.0, 0.5],
+            attention_other_setting={
+                "num_attention_heads": 8,
+                "num_attention_groups": 2,
+            },
+            use_head_wise_attn_gate=True,
+            moe_num_experts=4,
+            moe_top_k=2,
+            moe_intermediate_size=256,
+            share_expert_dim=256,
+            moe_layers_enum="1,2,3",
+        )
+        model = step3p5.Model(args)
+        self.model_test_runner(
+            model, args.model_type, args.vocab_size, args.num_hidden_layers
+        )
+
+    def test_step3p5_make_cache_uses_rotating_for_sliding_layers(self):
+        from mlx_lm.models import step3p5
+
+        args = step3p5.ModelArgs(
+            model_type="step3p5",
+            hidden_size=256,
+            num_hidden_layers=4,
+            vocab_size=1024,
+            num_attention_heads=4,
+            num_attention_groups=2,
+            head_dim=64,
+            intermediate_size=512,
+            rms_norm_eps=1e-5,
+            rope_theta=[10000.0, 10000.0, 10000.0, 10000.0],
+            sliding_window=4,
+            layer_types=[
+                "full_attention",
+                "sliding_attention",
+                "sliding_attention",
+                "full_attention",
+            ],
+            partial_rotary_factors=[0.5, 1.0, 1.0, 0.5],
+            attention_other_setting={
+                "num_attention_heads": 8,
+                "num_attention_groups": 2,
+            },
+            use_head_wise_attn_gate=True,
+            moe_num_experts=4,
+            moe_top_k=2,
+            moe_intermediate_size=256,
+            share_expert_dim=256,
+            moe_layers_enum="1,2,3",
+        )
+        model = step3p5.Model(args)
+
+        caches = model.make_cache()
+        self.assertIsInstance(caches[0], KVCache)
+        self.assertIsInstance(caches[1], RotatingKVCache)
+        self.assertIsInstance(caches[2], RotatingKVCache)
+        self.assertIsInstance(caches[3], KVCache)
+
+        tokens = mx.array([[1, 2, 3, 4, 5, 6, 7]], dtype=mx.int32)
+        step = model(tokens[:, :3], cache=caches)
+        mx.eval(step)
+        for i in range(3, 7):
+            step = model(tokens[:, i : i + 1], cache=caches)
+            mx.eval(step)
+
+        self.assertEqual(caches[0].size(), 7)
+        self.assertEqual(caches[1].size(), args.sliding_window)
+        self.assertEqual(caches[2].size(), args.sliding_window)
+        self.assertEqual(caches[3].size(), 7)
+
+    def test_step3p5_make_cache_uses_fallback_sliding_pattern(self):
+        from mlx_lm.models import step3p5
+
+        args = step3p5.ModelArgs(
+            model_type="step3p5",
+            hidden_size=256,
+            num_hidden_layers=5,
+            vocab_size=1024,
+            num_attention_heads=4,
+            num_attention_groups=2,
+            head_dim=64,
+            intermediate_size=512,
+            rms_norm_eps=1e-5,
+            rope_theta=10000.0,
+            sliding_window=4,
+            partial_rotary_factors=[1.0] * 5,
+            use_head_wise_attn_gate=True,
+            moe_num_experts=4,
+            moe_top_k=2,
+            moe_intermediate_size=256,
+            share_expert_dim=256,
+            moe_layers_enum="1,2,3,4",
+        )
+        model = step3p5.Model(args)
+
+        caches = model.make_cache()
+        self.assertIsInstance(caches[0], RotatingKVCache)
+        self.assertIsInstance(caches[1], KVCache)
+        self.assertIsInstance(caches[2], RotatingKVCache)
+        self.assertIsInstance(caches[3], KVCache)
+        self.assertIsInstance(caches[4], RotatingKVCache)
+
+        tokens = mx.array([[1, 2, 3, 4, 5, 6]], dtype=mx.int32)
+        step = model(tokens[:, :2], cache=caches)
+        mx.eval(step)
+        for i in range(2, 6):
+            step = model(tokens[:, i : i + 1], cache=caches)
+            mx.eval(step)
+
+        self.assertEqual(caches[0].size(), args.sliding_window)
+        self.assertEqual(caches[1].size(), 6)
+        self.assertEqual(caches[2].size(), args.sliding_window)
+        self.assertEqual(caches[3].size(), 6)
+        self.assertEqual(caches[4].size(), args.sliding_window)
+
     def test_cohere(self):
         from mlx_lm.models import cohere
 
@@ -1494,7 +1711,7 @@ class TestModels(unittest.TestCase):
                 "rms_norm_eps": 1e-5,
                 "vocab_size": 1000,
                 "num_key_value_heads": 2,
-                "partial_rotary_factor": 0,
+                "partial_rotary_factor": 0.5,
                 "rope_theta": 1000,
             },
             {
@@ -1522,7 +1739,7 @@ class TestModels(unittest.TestCase):
                 "use_qk_norm": True,
                 "tie_word_embeddings": False,
                 "attention_bias": False,
-                "partial_rotary_factor": 0.0,
+                "partial_rotary_factor": 0.5,
             },
             {
                 "model_type": "glm4_moe_lite",
@@ -1682,6 +1899,33 @@ class TestModels(unittest.TestCase):
                 "num_hidden_layers": 4,
                 "vocab_size": 1000,
             },
+            {
+                "model_type": "longcat_flash_ngram",
+                "attention_method": "MLA",
+                "zero_expert_type": "identity",
+                "hidden_size": 128,
+                "ffn_hidden_size": 128,
+                "moe_topk": 2,
+                "expert_ffn_hidden_size": 128,
+                "n_routed_experts": 2,
+                "zero_expert_num": 2,
+                "num_layers": 4,
+                "num_hidden_layers": 4,
+                "vocab_size": 1000,
+                "max_position_embeddings": 1000,
+                "num_attention_heads": 4,
+                "kv_lora_rank": 16,
+                "q_lora_rank": 16,
+                "qk_rope_head_dim": 8,
+                "qk_nope_head_dim": 8,
+                "v_head_dim": 8,
+                "routed_scaling_factor": 1.0,
+                "rms_norm_eps": 1e-5,
+                "rope_theta": 1000,
+                "mla_scale_q_lora": True,
+                "mla_scale_kv_lora": True,
+                "attention_bias": False,
+            },
             {
                 "model_type": "longcat_flash",
                 "attention_method": "MLA",
@@ -2056,6 +2300,47 @@ class TestModels(unittest.TestCase):
                 "partial_rotary_factor": 0.5,
                 "max_position_embeddings": 1000,
             },
+            {
+                "model_type": "qwen3_5",
+                "hidden_size": 128,
+                "num_hidden_layers": 4,
+                "intermediate_size": 128,
+                "num_attention_heads": 8,
+                "num_key_value_heads": 4,
+                "vocab_size": 1000,
+                "linear_num_value_heads": 4,
+                "linear_num_key_heads": 4,
+                "linear_key_head_dim": 32,
+                "linear_value_head_dim": 32,
+                "linear_conv_kernel_dim": 3,
+                "rms_norm_eps": 1e-5,
+                "head_dim": 64,
+                "rope_theta": 1000.0,
+                "partial_rotary_factor": 0.5,
+                "max_position_embeddings": 1000,
+            },
+            {
+                "model_type": "qwen3_5_moe",
+                "hidden_size": 128,
+                "num_hidden_layers": 4,
+                "num_attention_heads": 8,
+                "num_key_value_heads": 4,
+                "vocab_size": 1000,
+                "linear_num_value_heads": 4,
+                "linear_num_key_heads": 4,
+                "linear_key_head_dim": 32,
+                "linear_value_head_dim": 32,
+                "linear_conv_kernel_dim": 3,
+                "num_experts": 4,
+                "num_experts_per_tok": 2,
+                "shared_expert_intermediate_size": 128,
+                "moe_intermediate_size": 128,
+                "rms_norm_eps": 1e-5,
+                "head_dim": 64,
+                "rope_theta": 1000.0,
+                "partial_rotary_factor": 0.5,
+                "max_position_embeddings": 1000,
+            },
             {
                 "model_type": "kimi_linear",
                 "vocab_size": 1000,
@@ -2076,6 +2361,9 @@ class TestModels(unittest.TestCase):
                 "num_experts": 2,
                 "moe_intermediate_size": 128,
                 "kv_lora_rank": 8,
+                "qk_nope_head_dim": 16,
+                "qk_rope_head_dim": 16,
+                "v_head_dim": 16,
             },
             {
                 "model_type": "afmoe",
@@ -2213,6 +2501,18 @@ class TestModels(unittest.TestCase):
                 "kv_lora_rank": 128,
                 "q_lora_rank": 256,
             },
+            {
+                "model_type": "telechat3",
+                "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,
+                "rope_theta": 10000.0,
+                "max_position_embeddings": 1000,
+            },
         ]
         for config in test_configs:
             model_type = config["model_type"]
@@ -2358,6 +2658,60 @@ class TestModels(unittest.TestCase):
                 self.assertTrue(mx.allclose(y_op, y_c, rtol=1e-4, atol=1e-4))
                 self.assertTrue(mx.allclose(st_op, st_c, rtol=1e-4, atol=1e-4))
 
+    def test_gated_delta_precision(self):
+        mx.random.seed(42)
+
+        N_STEPS = 512
+        B = 1
+        Hk = 4
+        Hv = 4
+        Dk = 64
+        Dv = 64
+
+        A_log = mx.zeros((Hv,))
+        dt_bias = mx.ones((Hv,))
+
+        all_q = mx.random.normal(shape=(N_STEPS, B, 1, Hk, Dk)) * 0.1
+        all_k = mx.random.normal(shape=(N_STEPS, B, 1, Hk, Dk)) * 0.1
+        all_v = mx.random.normal(shape=(N_STEPS, B, 1, Hv, Dv)) * 0.1
+        all_a = -7.0 + mx.random.normal(shape=(N_STEPS, B, 1, Hv)) * 0.3
+        all_b = mx.random.normal(shape=(N_STEPS, B, 1, Hv))
+        mx.eval(all_q, all_k, all_v, all_a, all_b, A_log, dt_bias)
+
+        state_ref = mx.zeros((B, Hv, Dv, Dk), dtype=mx.float32)
+        for t in range(N_STEPS):
+            y_ref, state_ref = gated_delta_update(
+                all_q[t],
+                all_k[t],
+                all_v[t],
+                all_a[t],
+                all_b[t],
+                A_log,
+                dt_bias,
+                state_ref,
+                use_kernel=False,
+            )
+            mx.eval(y_ref, state_ref)
+
+        for use_kernel in (False, True):
+            state_lo = mx.zeros((B, Hv, Dv, Dk), dtype=mx.bfloat16)
+            for t in range(N_STEPS):
+                y_lo, state_lo = gated_delta_update(
+                    all_q[t].astype(mx.bfloat16),
+                    all_k[t].astype(mx.bfloat16),
+                    all_v[t].astype(mx.bfloat16),
+                    all_a[t].astype(mx.bfloat16),
+                    all_b[t].astype(mx.bfloat16),
+                    A_log,
+                    dt_bias,
+                    state_lo,
+                    use_kernel=use_kernel,
+                )
+                mx.eval(y_lo, state_lo)
+
+            self.assertTrue(mx.allclose(state_lo, state_ref, rtol=0.05, atol=0.01))
+            self.assertTrue(mx.allclose(y_lo, y_ref, rtol=0.05, atol=0.01))
+
     def test_gated_delta_masked(self):
         B = 1
         T = 3

tests/test_prompt_cache.py

@@ -16,7 +16,6 @@ from mlx_lm.models.cache import (
     CacheList,
     ChunkedKVCache,
     KVCache,
-    MambaCache,
     QuantizedKVCache,
     RotatingKVCache,
     load_prompt_cache,
@@ -103,14 +102,14 @@ class TestPromptCache(unittest.TestCase):
         cache_file = os.path.join(self.test_dir, "prompt_cache.safetensors")
 
         cache = [
-            MambaCache(),
+            ArraysCache(size=2),
             KVCache(),
             RotatingKVCache(8),
-            MambaCache(),
+            ArraysCache(size=2),
             ChunkedKVCache(256),
         ]
         for c in cache:
-            if isinstance(c, MambaCache):
+            if isinstance(c, ArraysCache):
                 c[0] = mx.random.uniform(shape=(4, 4, 4))
                 c[1] = mx.random.uniform(shape=(4, 4, 4))
             else:
@@ -121,7 +120,7 @@ class TestPromptCache(unittest.TestCase):
         save_prompt_cache(cache_file, cache)
         loaded_cache = load_prompt_cache(cache_file)
         for c, lc in zip(cache, loaded_cache):
-            if isinstance(c, MambaCache):
+            if isinstance(c, ArraysCache):
                 self.assertTrue(mx.array_equal(c[0], lc[0]))
                 self.assertTrue(mx.array_equal(c[1], lc[1]))
             else:
@@ -133,6 +132,54 @@ class TestPromptCache(unittest.TestCase):
                 self.assertTrue(mx.array_equal(k, lk))
                 self.assertTrue(mx.array_equal(v, lv))
 
+    def test_save_load_cache_list(self):
+        cache_file = os.path.join(self.test_dir, "prompt_cache.safetensors")
+
+        cache = [
+            ArraysCache(size=2),
+            KVCache(),
+            RotatingKVCache(8),
+            ArraysCache(size=2),
+            ChunkedKVCache(256),
+        ]
+        for c in cache:
+            if isinstance(c, ArraysCache):
+                c[0] = mx.random.uniform(shape=(4, 4, 4))
+                c[1] = mx.random.uniform(shape=(4, 4, 4))
+            else:
+                x = mx.random.uniform(shape=(4, 4, 7, 4))
+                y = mx.random.uniform(shape=(4, 4, 7, 4))
+                c.update_and_fetch(x, y)
+        cache = [CacheList(*cache)]
+
+        save_prompt_cache(cache_file, cache)
+        loaded_cache = load_prompt_cache(cache_file)
+        for c, lc in zip(cache[0].caches, loaded_cache[0].caches):
+            if isinstance(c, ArraysCache):
+                self.assertTrue(mx.array_equal(c[0], lc[0]))
+                self.assertTrue(mx.array_equal(c[1], lc[1]))
+            else:
+                x = mx.random.uniform(shape=(4, 4, 1, 4))
+                y = mx.random.uniform(shape=(4, 4, 1, 4))
+                k, v = c.update_and_fetch(x, y)
+                lk, lv = lc.update_and_fetch(x, y)
+                self.assertEqual(c.offset, lc.offset)
+                self.assertTrue(mx.array_equal(k, lk))
+                self.assertTrue(mx.array_equal(v, lv))
+
+    def test_save_load_arrays_cache(self):
+        cache_file = os.path.join(self.test_dir, "prompt_cache.safetensors")
+
+        cache = [ArraysCache(size=2)]
+        cache[0][0] = mx.zeros((1, 4, 4))
+        cache[0][1] = mx.zeros((1, 4, 4))
+
+        save_prompt_cache(cache_file, cache)
+        loaded = load_prompt_cache(cache_file)
+
+        # Try to make a mask
+        mask = loaded[0].make_mask(4)
+
     def test_cache_with_generate(self):
         model, tokenizer = self.model, self.tokenizer
         prompt = tokenizer.encode("this is a prompt", return_tensors="mlx")[0]
@@ -169,16 +216,18 @@ class TestPromptCache(unittest.TestCase):
         num_trimmed = trim_prompt_cache(cache, 4)
         self.assertEqual(num_trimmed, 3)
 
-        # Can't trim mamba cache
-        cache = [MambaCache() for _ in range(2)]
+        # Can't trim arrays cache
+        cache = [ArraysCache(size=2) for _ in range(2)]
         for c in cache:
-            c.state = mx.zeros((5, 5))
+            c[0] = mx.zeros((5, 5))
+            c[1] = mx.zeros((5, 5))
         num_trimmed = trim_prompt_cache(cache, 7)
         self.assertEqual(num_trimmed, 0)
 
         # All cache's have to be trimmable
-        cache = [MambaCache(), KVCache()]
-        cache[0].state = mx.zeros((5, 5))
+        cache = [ArraysCache(size=2), KVCache()]
+        cache[0][0] = mx.zeros((5, 5))
+        cache[0][1] = mx.zeros((5, 5))
         x = mx.random.uniform(shape=(1, 8, 10, 4))
         cache[1].update_and_fetch(x, x)
         num_trimmed = trim_prompt_cache(cache, 1)
@@ -325,7 +374,7 @@ class TestPromptCache(unittest.TestCase):
         m = c.trim(5)
         self.assertEqual(m, 5)
 
-        c = CacheList(MambaCache(), KVCache())
+        c = CacheList(ArraysCache(size=2), KVCache())
         self.assertFalse(c.is_trimmable())
 
         c1 = CacheList(ArraysCache(size=1), KVCache())
@@ -557,12 +606,12 @@ class TestPromptCache(unittest.TestCase):
         cache_file = os.path.join(self.test_dir, "prompt_cache.safetensors")
 
         cache = [
-            MambaCache(left_padding=[1, 2]),
+            ArraysCache(size=2, 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):
+            if isinstance(c, ArraysCache):
                 c[0] = mx.random.uniform(shape=(4, 4, 4))
                 c[1] = mx.random.uniform(shape=(4, 4, 4))
             else:
@@ -613,6 +662,16 @@ class TestPromptCache(unittest.TestCase):
         c_out = KVCache.merge((c1, c2))
         self.assertEqual(c_out.keys.shape, (2, 4, 4, 4))
 
+    def test_window_mask_with_full_kv_cache(self):
+        c = KVCache()
+        kv = mx.zeros((1, 1, 32, 128))
+        c.update_and_fetch(kv, kv)
+
+        h = mx.zeros((1, 1, 1, 128))
+        mask = create_attention_mask(h, c, window_size=4)
+        expected = create_causal_mask(1, offset=32, window_size=4)
+        self.assertTrue(mx.array_equal(mask, expected))
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_sample_utils.py

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

tests/test_server.py

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

tests/test_tool_parsing.py

@@ -5,23 +5,58 @@ from mlx_lm.tool_parsers import (
     function_gemma,
     glm47,
     json_tools,
+    kimi_k2,
+    longcat,
     minimax_m2,
+    mistral,
+    pythonic,
     qwen3_coder,
 )
 
 
 class TestToolParsing(unittest.TestCase):
-
     def test_parsers(self):
-        parsers = [function_gemma, glm47, json_tools, minimax_m2, qwen3_coder]
-
         test_cases = [
-            "call:multiply{a:12234585,b:48838483920}",
-            "multiply<arg_key>a</arg_key><arg_value>12234585</arg_value><arg_key>b</arg_key><arg_value>48838483920</arg_value>",
-            '{"name": "multiply", "arguments": {"a": 12234585, "b": 48838483920}}',
-            '<invoke name="multiply">\n<parameter name="a">12234585</parameter>\n<parameter name="b">48838483920</parameter>\n</invoke>',
-            "<function=multiply>\n<parameter=a>\n12234585\n</parameter>\n<parameter=b>\n48838483920\n</parameter>\n</function>",
+            ("call:multiply{a:12234585,b:48838483920}", function_gemma),
+            (
+                '{"name": "multiply", "arguments": {"a": 12234585, "b": 48838483920}}',
+                glm47,
+            ),
+            ("multiply a=12234585 b=48838483920", glm47),
+            (
+                "multiply<arg_key>a</arg_key><arg_value>12234585</arg_value><arg_key>b</arg_key><arg_value>48838483920</arg_value>",
+                glm47,
+            ),
+            (
+                '{"name": "multiply", "arguments": {"a": 12234585, "b": 48838483920}}',
+                json_tools,
+            ),
+            (
+                '<invoke name="multiply">\n<parameter name="a">12234585</parameter>\n<parameter name="b">48838483920</parameter>\n</invoke>',
+                minimax_m2,
+            ),
+            (
+                "<function=multiply>\n<parameter=a>\n12234585\n</parameter>\n<parameter=b>\n48838483920\n</parameter>\n</function>",
+                qwen3_coder,
+            ),
+            (
+                "multiply<longcat_arg_key>a</longcat_arg_key>\n<longcat_arg_value>12234585</longcat_arg_value>\n<longcat_arg_key>b</longcat_arg_key>\n<longcat_arg_value>48838483920</longcat_arg_value>",
+                longcat,
+            ),
+            (
+                '{"name": "multiply", "arguments": {"a": 12234585, "b": 48838483920}}',
+                longcat,
+            ),
+            (
+                "[multiply(a=12234585, b=48838483920)]",
+                pythonic,
+            ),
+            (
+                'multiply[ARGS]{"a": 12234585, "b": 48838483920}',
+                mistral,
+            ),
         ]
+
         tools = [
             {
                 "type": "function",
@@ -40,7 +75,7 @@ class TestToolParsing(unittest.TestCase):
             }
         ]
 
-        for parser, test_case in zip(parsers, test_cases):
+        for test_case, parser in test_cases:
             with self.subTest(parser=parser):
                 tool_call = parser.parse_tool_call(test_case, tools)
                 expected = {
@@ -50,11 +85,42 @@ class TestToolParsing(unittest.TestCase):
                 self.assertEqual(tool_call, expected)
 
         test_cases = [
-            "call:get_current_temperature{location:<escape>London<escape>}",
-            'get_current_temperature<arg_key>location</arg_key><arg_value>"London"</arg_value>',
-            '{"name": "get_current_temperature", "arguments": {"location": "London"}}',
-            '<invoke name="get_current_temperature">\n<parameter name="location">London</parameter>\n</invoke>',
-            "<function=get_current_temperature>\n<parameter=location>\nLondon\n</parameter>\n</function>",
+            (
+                "call:get_current_temperature{location:<escape>London<escape>}",
+                function_gemma,
+            ),
+            (
+                'get_current_temperature<arg_key>location</arg_key><arg_value>"London"</arg_value>',
+                glm47,
+            ),
+            (
+                '{"name": "get_current_temperature", "arguments": {"location": "London"}}',
+                json_tools,
+            ),
+            (
+                '<invoke name="get_current_temperature">\n<parameter name="location">London</parameter>\n</invoke>',
+                minimax_m2,
+            ),
+            (
+                "<function=get_current_temperature>\n<parameter=location>\nLondon\n</parameter>\n</function>",
+                qwen3_coder,
+            ),
+            (
+                "get_current_temperature<longcat_arg_key>location</longcat_arg_key>\n<longcat_arg_value>London</longcat_arg_value>",
+                longcat,
+            ),
+            (
+                '{"name": "get_current_temperature", "arguments": {"location": "London"}}',
+                longcat,
+            ),
+            (
+                '[get_current_temperature(location="London")]',
+                pythonic,
+            ),
+            (
+                'get_current_temperature[ARGS]{"location": "London"}',
+                mistral,
+            ),
         ]
         tools = [
             {
@@ -73,7 +139,7 @@ class TestToolParsing(unittest.TestCase):
             }
         ]
 
-        for parser, test_case in zip(parsers, test_cases):
+        for test_case, parser in test_cases:
             with self.subTest(parser=parser):
                 tool_call = parser.parse_tool_call(test_case, tools)
                 expected = {
@@ -82,6 +148,87 @@ class TestToolParsing(unittest.TestCase):
                 }
                 self.assertEqual(tool_call, expected)
 
+    def test_qwen3_coder_single_quoted_params(self):
+        tools = [
+            {
+                "type": "function",
+                "function": {
+                    "name": "search",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "filters": {"type": "object"},
+                            "tags": {"type": "array"},
+                        },
+                    },
+                },
+            }
+        ]
+
+        # single-quoted dict (python-style, not valid JSON)
+        test_case = (
+            "<function=search>"
+            "<parameter=filters>{'category': 'books', 'in_stock': True}</parameter>"
+            "<parameter=tags>['fiction', 'new']</parameter>"
+            "</function>"
+        )
+        tool_call = qwen3_coder.parse_tool_call(test_case, tools)
+        self.assertEqual(tool_call["name"], "search")
+        self.assertEqual(
+            tool_call["arguments"]["filters"],
+            {"category": "books", "in_stock": True},
+        )
+        self.assertEqual(tool_call["arguments"]["tags"], ["fiction", "new"])
+
+        # valid JSON (double-quoted) should still work
+        test_case = (
+            "<function=search>"
+            '<parameter=filters>{"category": "books"}</parameter>'
+            '<parameter=tags>["fiction", "new"]</parameter>'
+            "</function>"
+        )
+        tool_call = qwen3_coder.parse_tool_call(test_case, tools)
+        self.assertEqual(tool_call["arguments"]["filters"], {"category": "books"})
+        self.assertEqual(tool_call["arguments"]["tags"], ["fiction", "new"])
+
+    def test_kimi_k2(self):
+        # Single tool call
+        test_case = (
+            "<|tool_call_begin|>functions.multiply:0<|tool_call_argument_begin|>"
+            '{"a": 12234585, "b": 48838483920}<|tool_call_end|>'
+        )
+        tool_calls = kimi_k2.parse_tool_call(test_case, None)
+        expected = [
+            {
+                "id": "functions.multiply:0",
+                "name": "multiply",
+                "arguments": {"a": 12234585, "b": 48838483920},
+            }
+        ]
+        self.assertEqual(tool_calls, expected)
+
+        # Multiple tool calls
+        test_case = (
+            "<|tool_call_begin|>functions.search:0<|tool_call_argument_begin|>"
+            '{"query": "weather"}<|tool_call_end|>'
+            "<|tool_call_begin|>functions.read_file:1<|tool_call_argument_begin|>"
+            '{"path": "/tmp/test.txt"}<|tool_call_end|>'
+        )
+        tool_calls = kimi_k2.parse_tool_call(test_case, None)
+        expected = [
+            {
+                "id": "functions.search:0",
+                "name": "search",
+                "arguments": {"query": "weather"},
+            },
+            {
+                "id": "functions.read_file:1",
+                "name": "read_file",
+                "arguments": {"path": "/tmp/test.txt"},
+            },
+        ]
+        self.assertEqual(tool_calls, expected)
+
 
 if __name__ == "__main__":
     unittest.main()