patch (#731)

* Parse reasoning in server

* redesign and start to fix tool parsing

* add function gemma

* fix

* fix

* glm47 tools

* add minimax m2 tool parser

* tool_call finish reason

* Keep model wired in the server to reduce ttft

* infer tool parser

.github/workflows/pull_request.yml

@@ -38,4 +38,6 @@ jobs:
       - name: Run tests
         shell: bash -l {0}
         run: |
-          python -m xmlrunner discover -v tests -o test-results/
+          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/

README.md

@@ -71,7 +71,7 @@ prompt = "Write a story about Einstein"
 
 messages = [{"role": "user", "content": prompt}]
 prompt = tokenizer.apply_chat_template(
-    messages, add_generation_prompt=True
+    messages, add_generation_prompt=True,
 )
 
 text = generate(model, tokenizer, prompt=prompt, verbose=True)
@@ -130,7 +130,7 @@ prompt = "Write a story about Einstein"
 
 messages = [{"role": "user", "content": prompt}]
 prompt = tokenizer.apply_chat_template(
-    messages, add_generation_prompt=True
+    messages, add_generation_prompt=True,
 )
 
 for response in stream_generate(model, tokenizer, prompt, max_tokens=512):
@@ -170,7 +170,7 @@ mlx_lm.generate --help
 To quantize a model from the command line run:
 
 ```
-mlx_lm.convert --hf-path mistralai/Mistral-7B-Instruct-v0.3 -q
+mlx_lm.convert --model mistralai/Mistral-7B-Instruct-v0.3 -q
 ```
 
 For more options run:
@@ -185,7 +185,7 @@ You can upload new models to Hugging Face by specifying `--upload-repo` to
 
 ```
 mlx_lm.convert \
-    --hf-path mistralai/Mistral-7B-Instruct-v0.3 \
+    --model mistralai/Mistral-7B-Instruct-v0.3 \
     -q \
     --upload-repo mlx-community/my-4bit-mistral
 ```

mlx_lm/_version.py

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

mlx_lm/benchmark.py

@@ -6,7 +6,7 @@ import mlx.core as mx
 
 from mlx_lm import batch_generate, load, stream_generate
 from mlx_lm.generate import DEFAULT_MODEL
-from mlx_lm.utils import pipeline_load
+from mlx_lm.utils import pipeline_load, sharded_load
 
 
 def setup_arg_parser():
@@ -49,6 +49,11 @@ def setup_arg_parser():
         help="Number of timing trials",
         type=int,
     )
+    parser.add_argument(
+        "--pipeline",
+        action="store_true",
+        help="Use pipelining instead of tensor parallelism",
+    )
     return parser
 
 
@@ -59,6 +64,8 @@ def main():
 
     group = mx.distributed.init()
     rank = group.rank()
+    pipeline_group = group if args.pipeline else None
+    tensor_group = group if not args.pipeline else None
 
     def rprint(*args, **kwargs):
         if rank == 0:
@@ -67,7 +74,9 @@ def main():
     model_path = args.model or DEFAULT_MODEL
 
     if group.size() > 1:
-        model, tokenizer, config = pipeline_load(args.model, return_config=True)
+        model, tokenizer, config = sharded_load(
+            args.model, pipeline_group, tensor_group, return_config=True
+        )
     else:
         model, tokenizer, config = load(
             args.model, return_config=True, tokenizer_config={"trust_remote_code": True}

mlx_lm/cache_prompt.py

@@ -41,16 +41,6 @@ def setup_arg_parser():
         default=None,
         help="End of sequence token for tokenizer",
     )
-    parser.add_argument(
-        "--ignore-chat-template",
-        action="store_true",
-        help="Use the raw prompt without the tokenizer's chat template.",
-    )
-    parser.add_argument(
-        "--use-default-chat-template",
-        action="store_true",
-        help="Use the default chat template",
-    )
     parser.add_argument(
         "--max-kv-size",
         type=int,
@@ -107,14 +97,12 @@ def main():
 
     args.prompt = sys.stdin.read() if args.prompt == "-" else args.prompt
 
-    if args.use_default_chat_template:
-        if tokenizer.chat_template is None:
-            tokenizer.chat_template = tokenizer.default_chat_template
-
-    if not args.ignore_chat_template and tokenizer.chat_template is not None:
+    if tokenizer.has_chat_template:
         messages = [{"role": "user", "content": args.prompt}]
         prompt = tokenizer.apply_chat_template(
-            messages, add_generation_prompt=False, continue_final_message=True
+            messages,
+            add_generation_prompt=False,
+            continue_final_message=True,
         )
 
     else:
@@ -153,7 +141,6 @@ def main():
     print("Saving...")
     metadata = {}
     metadata["model"] = args.model
-    metadata["chat_template"] = json.dumps(tokenizer.chat_template)
     metadata["tokenizer_config"] = json.dumps(tokenizer_config)
     save_prompt_cache(args.prompt_cache_file, cache, metadata)
 

mlx_lm/chat.py

@@ -7,7 +7,7 @@ import mlx.core as mx
 from .generate import stream_generate
 from .models.cache import make_prompt_cache
 from .sample_utils import make_sampler
-from .utils import load
+from .utils import load, sharded_load
 
 DEFAULT_TEMP = 0.0
 DEFAULT_TOP_P = 1.0
@@ -79,6 +79,11 @@ def setup_arg_parser():
         default=None,
         help="System prompt to be used for the chat template",
     )
+    parser.add_argument(
+        "--pipeline",
+        action="store_true",
+        help="Use pipelining instead of tensor parallelism",
+    )
     return parser
 
 
@@ -86,28 +91,42 @@ def main():
     parser = setup_arg_parser()
     args = parser.parse_args()
 
+    group = mx.distributed.init()
+    rank = group.rank()
+    pipeline_group = group if args.pipeline else None
+    tensor_group = group if not args.pipeline else None
+
+    def rprint(*args, **kwargs):
+        if rank == 0:
+            print(*args, **kwargs)
+
     if args.seed is not None:
         mx.random.seed(args.seed)
 
-    model, tokenizer = load(
-        args.model,
-        adapter_path=args.adapter_path,
-        tokenizer_config={
-            "trust_remote_code": True if args.trust_remote_code else None
-        },
-    )
+    if group.size() > 1:
+        if args.adapter_path:
+            parser.error("Adapters not supported in distributed mode")
+        model, tokenizer = sharded_load(args.model, pipeline_group, tensor_group)
+    else:
+        model, tokenizer = load(
+            args.model,
+            adapter_path=args.adapter_path,
+            tokenizer_config={
+                "trust_remote_code": True if args.trust_remote_code else None
+            },
+        )
 
     def print_help():
-        print("The command list:")
-        print("- 'q' to exit")
-        print("- 'r' to reset the chat")
-        print("- 'h' to display these commands")
+        rprint("The command list:")
+        rprint("- 'q' to exit")
+        rprint("- 'r' to reset the chat")
+        rprint("- 'h' to display these commands")
 
-    print(f"[INFO] Starting chat session with {args.model}.")
+    rprint(f"[INFO] Starting chat session with {args.model}.")
     print_help()
     prompt_cache = make_prompt_cache(model, args.max_kv_size)
     while True:
-        query = input(">> ")
+        query = input(">> " if rank == 0 else "")
         if query == "q":
             break
         if query == "r":
@@ -120,7 +139,10 @@ def main():
         if args.system_prompt is not None:
             messages.append({"role": "system", "content": args.system_prompt})
         messages.append({"role": "user", "content": query})
-        prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
+        prompt = tokenizer.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+        )
         for response in stream_generate(
             model,
             tokenizer,
@@ -137,8 +159,8 @@ def main():
             ),
             prompt_cache=prompt_cache,
         ):
-            print(response.text, flush=True, end="")
-        print()
+            rprint(response.text, flush=True, end="")
+        rprint()
 
 
 if __name__ == "__main__":

mlx_lm/chat_templates/deepseek_v32.py

@@ -0,0 +1,332 @@
+# Copyright © 2025 Apple Inc.
+
+import copy
+import json
+import re
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+TOOLS_SYSTEM_TEMPLATE = """## Tools
+
+You have access to a set of tools you can use to answer the user's question.
+You can invoke functions by writing a "<{dsml_token}function_calls>" block like the following as part of your reply to the user:
+<{dsml_token}function_calls>
+<{dsml_token}invoke name="$FUNCTION_NAME">
+<{dsml_token}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</{dsml_token}parameter>
+...
+</{dsml_token}invoke>
+<{dsml_token}invoke name="$FUNCTION_NAME2">
+...
+</{dsml_token}invoke>
+</{dsml_token}function_calls>
+
+String and scalar parameters should be specified as is without any escaping or quotes, while lists and objects should use JSON format. The "string" attribute should be set to "true" for string type parameters and "false" for other types (numbers, booleans, arrays, objects).
+
+If the thinking_mode is enabled, then after function results you should strongly consider outputting a thinking block. Here is an example:
+
+<{dsml_token}function_calls>
+...
+</{dsml_token}function_calls>
+
+<function_results>
+...
+</function_results>
+
+{thinking_start_token}...thinking about results{thinking_end_token}
+
+Here are the functions available in JSONSchema format:
+<functions>
+{tool_schemas}
+</functions>
+"""
+
+bos_token: str = "<｜begin▁of▁sentence｜>"
+eos_token: str = "<｜end▁of▁sentence｜>"
+thinking_start_token: str = "<think>"
+thinking_end_token: str = "</think>"
+dsml_token: str = "｜DSML｜"
+system_msg_template: str = "{content}"
+user_msg_template: str = "<｜User｜>{content}<｜Assistant｜>"
+assistant_msg_template: str = "{reasoning}{content}{tool_calls}<｜end▁of▁sentence｜>"
+thinking_template = "{reasoning_content}"
+
+response_format_template: str = (
+    "## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n{schema}"
+)
+tool_call_template: str = (
+    '<{dsml_token}invoke name="{name}">\n{arguments}\n</{dsml_token}invoke>'
+)
+tool_calls_template = (
+    "<{dsml_token}function_calls>\n{tool_calls}\n</{dsml_token}function_calls>"
+)
+
+tool_output_template: str = "\n<result>{content}</result>"
+
+
+def to_json(value: Any) -> str:
+    try:
+        return json.dumps(value, ensure_ascii=False)
+    except:
+        return json.dumps(value, ensure_ascii=True)
+
+
+def tools_from_openai_format(tools):
+    return [tool["function"] for tool in tools]
+
+
+def tool_calls_from_openai_format(tool_calls):
+    return [
+        {
+            "name": tool_call["function"]["name"],
+            "arguments": tool_call["function"]["arguments"],
+        }
+        for tool_call in tool_calls
+    ]
+
+
+def encode_arguments_to_dsml(tool_call: Dict[str, str]) -> str:
+    p_dsml_template = """<{dsml_token}parameter name="{key}" string="{is_str}">{value}</{dsml_token}parameter>"""
+    P_dsml_strs = []
+
+    arguments = json.loads(tool_call["arguments"])
+
+    for k, v in arguments.items():
+        p_dsml_str = p_dsml_template.format(
+            dsml_token=dsml_token,
+            key=k,
+            is_str="true" if isinstance(v, str) else "false",
+            value=v if isinstance(v, str) else to_json(v),
+        )
+
+        P_dsml_strs.append(p_dsml_str)
+
+    return "\n".join(P_dsml_strs)
+
+
+def decode_dsml_to_arguments(
+    tool_name: str, tool_args: Dict[str, Tuple[str, str]]
+) -> Dict[str, str]:
+    def _decode_value(key: str, value: str, string: str):
+        if string == "true":
+            value = to_json(value)
+        return f"{to_json(key)}: {value}"
+
+    tool_args_json = (
+        "{"
+        + ", ".join(
+            [_decode_value(k, v, string=is_str) for k, (v, is_str) in tool_args.items()]
+        )
+        + "}"
+    )
+    return dict(name=tool_name, arguments=tool_args_json)
+
+
+def render_tools(tools: List[Dict[str, Union[str, Dict[str, Any]]]]) -> str:
+    tools_json = [to_json(t) for t in tools]
+
+    return TOOLS_SYSTEM_TEMPLATE.format(
+        tool_schemas="\n".join(tools_json),
+        dsml_token=dsml_token,
+        thinking_start_token=thinking_start_token,
+        thinking_end_token=thinking_end_token,
+    )
+
+
+def find_last_user_index(messages: List[Dict[str, Any]]) -> int:
+    last_user_index = -1
+    for idx in range(len(messages) - 1, -1, -1):
+        if messages[idx].get("role") in ["user", "developer"]:
+            last_user_index = idx
+            break
+    return last_user_index
+
+
+def render_message(
+    index: int, messages: List[Dict[str, Any]], thinking_mode: str
+) -> str:
+    assert 0 <= index < len(messages)
+    assert thinking_mode in [
+        "chat",
+        "thinking",
+    ], f"Invalid thinking_mode `{thinking_mode}`"
+
+    prompt = ""
+    msg = messages[index]
+    last_user_idx = find_last_user_index(messages)
+
+    role = msg.get("role")
+    content = msg.get("content")
+    tools = 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)
+
+        if response_format:
+            prompt += "\n\n" + response_format_template.format(
+                schema=to_json(response_format)
+            )
+
+    elif role == "developer":
+        assert content, f"Invalid message for role `{role}`: {msg}"
+        content_developer = ""
+        if tools:
+            content_developer += "\n\n" + render_tools(tools)
+
+        if response_format:
+            content_developer += "\n\n" + response_format_template.format(
+                schema=to_json(response_format)
+            )
+
+        content_developer += "\n\n# The user's message is: {}".format(content)
+
+        prompt += user_msg_template.format(content=content_developer)
+        if index == last_user_idx and thinking_mode == "thinking":
+            prompt += thinking_start_token
+        else:
+            prompt += thinking_end_token
+
+    elif role == "user":
+        prompt += user_msg_template.format(content=content)
+
+        if index == last_user_idx and thinking_mode == "thinking":
+            prompt += thinking_start_token
+        else:
+            prompt += thinking_end_token
+
+    elif role == "tool":
+        prev_assistant_idx = index - 1
+        assistant_msg = messages[prev_assistant_idx]
+        while prev_assistant_idx >= 0 and assistant_msg.get("role") == "tool":
+            prev_assistant_idx -= 1
+            assistant_msg = messages[prev_assistant_idx]
+
+        assert (
+            index == 0
+            or prev_assistant_idx >= 0
+            and assistant_msg.get("role") == "assistant"
+        ), f"Invalid messages at {index}:\n{assistant_msg}"
+
+        tool_call_order = index - prev_assistant_idx
+        assistant_tool_calls = assistant_msg.get("tool_calls")
+        assert (
+            assistant_tool_calls and len(assistant_tool_calls) >= tool_call_order
+        ), "No tool calls but found tool output"
+
+        if tool_call_order == 1:
+            prompt += "\n\n<function_results>"
+
+        prompt += tool_output_template.format(content=content)
+
+        if tool_call_order == len(assistant_tool_calls):
+            prompt += "\n</function_results>"
+
+            if index >= last_user_idx and thinking_mode == "thinking":
+                prompt += "\n\n" + thinking_start_token
+            else:
+                prompt += "\n\n" + thinking_end_token
+
+    elif role == "assistant":
+        prev_assistant_idx = index
+        thinking_part = ""
+
+        tool_calls_content = ""
+        if tool_calls:
+            tool_calls = [
+                tool_call_template.format(
+                    dsml_token=dsml_token,
+                    name=tool_call.get("name"),
+                    arguments=encode_arguments_to_dsml(tool_call),
+                )
+                for tool_call in tool_calls
+            ]
+            tool_calls_content += "\n\n" + tool_calls_template.format(
+                dsml_token=dsml_token, tool_calls="\n".join(tool_calls)
+            )
+
+        summary_content = content or ""
+
+        if thinking_mode == "thinking" and index > last_user_idx:
+            assert (
+                reasoning_content or tool_calls
+            ), f"ThinkingMode: {thinking_mode}, invalid message without reasoning_content/tool_calls `{msg}` after last user message"
+            thinking_part = (
+                thinking_template.format(reasoning_content=reasoning_content or "")
+                + thinking_end_token
+            )
+
+        prompt += assistant_msg_template.format(
+            reasoning=thinking_part,
+            content=summary_content,
+            tool_calls=tool_calls_content,
+        )
+    else:
+        raise NotImplementedError(f"Unknown role: {role}")
+
+    return prompt
+
+
+def drop_thinking_messages(
+    messages: List[Dict[str, Any]], last_user_idx: Optional[int] = None
+) -> List[Dict[str, Any]]:
+    messages_wo_thinking: List[Dict[str, Any]] = []
+    last_user_idx = (
+        find_last_user_index(messages) if last_user_idx is None else last_user_idx
+    )
+    for idx, msg in enumerate(messages):
+        role = msg.get("role")
+        if role in ["user", "system", "tool"] or idx >= last_user_idx:
+            messages_wo_thinking.append(msg)
+            continue
+
+        elif role == "assistant":
+            msg_wo_thinking = copy.copy(msg)
+            msg_wo_thinking.pop("reasoning_content", None)
+            messages_wo_thinking.append(msg_wo_thinking)
+
+    return messages_wo_thinking
+
+
+def encode_messages(
+    messages: List[Dict[str, Any]],
+    thinking_mode: str = "thinking",
+    context: Optional[List[Dict[str, Any]]] = None,
+    drop_thinking: bool = True,
+    add_default_bos_token: bool = True,
+) -> str:
+    context = context if context else []
+    full_messages = context + messages
+    prompt = bos_token if add_default_bos_token and len(context) == 0 else ""
+
+    if thinking_mode == "thinking" and drop_thinking:
+        full_messages = drop_thinking_messages(full_messages)
+
+    for idx in range(len(messages)):
+        prompt += render_message(
+            idx + len(context), full_messages, thinking_mode=thinking_mode
+        )
+
+    return prompt
+
+
+def apply_chat_template(
+    messages, continue_final_message=False, add_generation_prompt=False, **kwargs
+):
+    out = encode_messages(messages, **kwargs)
+    if continue_final_message and add_generation_prompt:
+        raise ValueError(
+            "Only one of continue_final_message or add_generation_prompt can be True"
+        )
+    if not add_generation_prompt and messages[-1]["role"] == "user":
+        out = out.removesuffix("<｜Assistant｜><think>")
+    if continue_final_message and messages[-1]["role"] == "assistant":
+        out = out.removesuffix(eos_token)
+    return out

mlx_lm/convert.py

@@ -179,7 +179,12 @@ def configure_parser() -> argparse.ArgumentParser:
         description="Convert Hugging Face model to MLX format"
     )
 
-    parser.add_argument("--hf-path", type=str, help="Path to the Hugging Face model.")
+    parser.add_argument(
+        "--hf-path",
+        "--model",
+        type=str,
+        help="Path to the model. This can be a local path or a Hugging Face Hub model identifier.",
+    )
     parser.add_argument(
         "--mlx-path", type=str, default="mlx_model", help="Path to save the MLX model."
     )
@@ -187,17 +192,23 @@ def configure_parser() -> argparse.ArgumentParser:
         "-q", "--quantize", help="Generate a quantized model.", action="store_true"
     )
     parser.add_argument(
-        "--q-group-size", help="Group size for quantization.", type=int, default=64
+        "--q-group-size",
+        help="Group size for quantization.",
+        type=int,
+        default=None,
     )
     parser.add_argument(
-        "--q-bits", help="Bits per weight for quantization.", type=int, default=4
+        "--q-bits",
+        help="Bits per weight for quantization.",
+        type=int,
+        default=None,
     )
     parser.add_argument(
         "--q-mode",
         help="The quantization mode.",
         type=str,
         default="affine",
-        choices=["affine", "mxfp4"],
+        choices=["affine", "mxfp4", "nvfp4", "mxfp8"],
     )
     parser.add_argument(
         "--quant-predicate",

mlx_lm/examples/chat.py

@@ -15,7 +15,10 @@ prompt_cache = make_prompt_cache(model)
 # User turn
 prompt = "Hi my name is <Name>."
 messages = [{"role": "user", "content": prompt}]
-prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
+prompt = tokenizer.apply_chat_template(
+    messages,
+    add_generation_prompt=True,
+)
 
 # Assistant response
 response = generate(
@@ -29,7 +32,10 @@ response = generate(
 # User turn
 prompt = "What's my name?"
 messages = [{"role": "user", "content": prompt}]
-prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
+prompt = tokenizer.apply_chat_template(
+    messages,
+    add_generation_prompt=True,
+)
 
 # Assistant response
 response = generate(

mlx_lm/examples/generate_response.py

@@ -14,7 +14,8 @@ conversation = [{"role": "user", "content": prompt}]
 
 # Transform the prompt into the chat template
 prompt = tokenizer.apply_chat_template(
-    conversation=conversation, add_generation_prompt=True
+    conversation=conversation,
+    add_generation_prompt=True,
 )
 
 # Specify the maximum number of tokens

mlx_lm/examples/openai_reasoning_content.py

@@ -0,0 +1,40 @@
+from openai import OpenAI
+
+client = OpenAI(
+    api_key="not-needed",
+    base_url="http://localhost:8080/v1",
+)
+
+model = "mlx-community/Qwen3-4B-Thinking-2507-4bit"
+
+messages = [{"role": "user", "content": "9.11 and 9.8, which is greater?"}]
+
+# Non-streaming example
+
+response = client.chat.completions.create(
+    model=model, messages=messages, max_tokens=2048
+)
+
+reasoning = response.choices[0].message.reasoning
+content = response.choices[0].message.content
+
+print("=== reasoning ===\n")
+print(f"\033[37m{reasoning}\033[0m")
+print("=== content ===\n")
+print(content)
+
+# Streaming example
+
+stream = client.chat.completions.create(
+    model=model,
+    messages=messages,
+    stream=True,
+    max_tokens=2048,
+)
+
+for chunk in stream:
+    if (reasoning := chunk.choices[0].delta.reasoning) is not None:
+        print(f"\033[37m{reasoning}\033[0m", end="")
+    if (content := chunk.choices[0].delta.content) is not None:
+        print(f"{content}", end="")
+print()

mlx_lm/examples/openai_tool_use.py

@@ -8,11 +8,13 @@ To run, first start the server:
 
 Then run this script.
 """
+import json
+
 from openai import OpenAI
 
 client = OpenAI(base_url="http://localhost:8080/v1", api_key="not-needed")
 
-model = "mlx-community/qwen3-4b-4bit-DWQ"
+model = "mlx-community/Qwen3-4B-Instruct-2507-4bit"
 messages = [{"role": "user", "content": "What's the weather in Boston?"}]
 
 tools = [

mlx_lm/examples/sharded_generate.py

@@ -1,19 +1,20 @@
-# Copyright © 2024 Apple Inc.
+# Copyright © 2025 Apple Inc.
 
 """
 Run with:
 
 ```
 mlx.launch \
- --hostfile /path/to/hosts.json \
- /path/to/pipeline_generate.py \
- --prompt "hello world"
+    --backend jaccl \
+    --env MLX_METAL_FAST_SYNCH=1 \
+    --hostfile /path/to/hosts.json \
+    /path/to/sharded_generate.py \
+    --prompt 'Hello world'
 ```
 
-Make sure you can run MLX over MPI on two hosts. For more information see the
-documentation:
+For more information on running distributed programs with MLX see the documentation:
 
-https://ml-explore.github.io/mlx/build/html/usage/distributed.html).
+https://ml-explore.github.io/mlx/build/html/usage/distributed.html .
 """
 
 import argparse
@@ -21,13 +22,13 @@ import argparse
 import mlx.core as mx
 
 from mlx_lm import stream_generate
-from mlx_lm.utils import pipeline_load
+from mlx_lm.utils import sharded_load
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="LLM pipelined inference example")
+    parser = argparse.ArgumentParser(description="LLM distributed inference example")
     parser.add_argument(
         "--model",
-        default="mlx-community/DeepSeek-R1-3bit",
+        default="mlx-community/Llama-3.3-70B-Instruct-4bit",
         help="HF repo or path to local model.",
     )
     parser.add_argument(
@@ -43,19 +44,29 @@ if __name__ == "__main__":
         default=256,
         help="Maximum number of tokens to generate",
     )
+    parser.add_argument(
+        "--pipeline",
+        action="store_true",
+        help="Use pipelining instead of tensor parallelism",
+    )
     args = parser.parse_args()
 
     group = mx.distributed.init()
     rank = group.rank()
+    pipeline_group = group if args.pipeline else None
+    tensor_group = group if not args.pipeline else None
 
     def rprint(*args, **kwargs):
         if rank == 0:
             print(*args, **kwargs)
 
-    model, tokenizer = pipeline_load(args.model)
+    model, tokenizer = sharded_load(args.model, pipeline_group, tensor_group)
 
     messages = [{"role": "user", "content": args.prompt}]
-    prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
+    prompt = tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+    )
 
     for response in stream_generate(
         model, tokenizer, prompt, max_tokens=args.max_tokens

mlx_lm/examples/tool_use.py

@@ -6,7 +6,7 @@ from mlx_lm import generate, load
 from mlx_lm.models.cache import make_prompt_cache
 
 # Specify the checkpoint
-checkpoint = "mlx-community/Qwen2.5-32B-Instruct-4bit"
+checkpoint = "mlx-community/Qwen3-4B-Instruct-2507-4bit"
 
 # Load the corresponding model and tokenizer
 model, tokenizer = load(path_or_hf_repo=checkpoint)
@@ -31,7 +31,9 @@ prompt = "Multiply 12234585 and 48838483920."
 messages = [{"role": "user", "content": prompt}]
 
 prompt = tokenizer.apply_chat_template(
-    messages, add_generation_prompt=True, tools=list(tools.values())
+    messages,
+    add_generation_prompt=True,
+    tools=list(tools.values()),
 )
 
 prompt_cache = make_prompt_cache(model)
@@ -47,12 +49,11 @@ response = generate(
 )
 
 # Parse the tool call:
-# (Note, the tool call format is model specific)
-tool_open = "<tool_call>"
-tool_close = "</tool_call>"
-start_tool = response.find(tool_open) + len(tool_open)
-end_tool = response.find(tool_close)
-tool_call = json.loads(response[start_tool:end_tool].strip())
+# - The tool call format is model specific.
+# - The tokenizer's tool parser expects tool call text to be already extracted.
+start_tool = response.find(tokenizer.tool_call_start) + len(tokenizer.tool_call_start)
+end_tool = response.find(tokenizer.tool_call_end)
+tool_call = tokenizer.tool_parser(response[start_tool:end_tool].strip())
 tool_result = tools[tool_call["name"]](**tool_call["arguments"])
 
 # Put the tool result in the prompt

mlx_lm/fuse.py

@@ -76,8 +76,9 @@ def main() -> None:
 
     if args.dequantize:
         print("Dequantizing model")
-        model = dequantize(model)
+        model = dequantize_model(model)
         config.pop("quantization", None)
+        config.pop("quantization_config", None)
 
     save_path = Path(args.save_path)
     save(

mlx_lm/generate.py

@@ -181,8 +181,7 @@ def setup_arg_parser():
     parser.add_argument(
         "--kv-bits",
         type=int,
-        help="Number of bits for KV cache quantization. "
-        "Defaults to no quantization.",
+        help="Number of bits for KV cache quantization. Defaults to no quantization.",
         default=None,
     )
     parser.add_argument(
@@ -548,7 +547,9 @@ def speculative_generate_step(
                     y = y[: -(n_predict - 1)]
                 for i in range(n_predict):
                     prev_tokens = (
-                        mx.concat([prev_tokens, y]) if prev_tokens is not None else y
+                        mx.concatenate([prev_tokens, y])
+                        if prev_tokens is not None
+                        else y
                     )
                     y, logprobs = _process_and_sample(prev_tokens, logits[:, i, :])
                     out_y.append(y)
@@ -840,6 +841,9 @@ class Batch:
     max_tokens: List[int]
     num_tokens: List[int]
     cache: List[Any]
+    samplers: List[Any]
+    logits_processors: List[Any]
+    tokens: List[mx.array]
 
     def __len__(self):
         return len(self.uids)
@@ -849,6 +853,9 @@ class Batch:
         self.logprobs = [self.logprobs[k] for k in keep_idx]
         self.max_tokens = [self.max_tokens[k] for k in keep_idx]
         self.num_tokens = [self.num_tokens[k] for k in keep_idx]
+        self.samplers = [self.samplers[k] for k in keep_idx]
+        self.logits_processors = [self.logits_processors[k] for k in keep_idx]
+        self.tokens = [self.tokens[k] for k in keep_idx]
         keep_idx = mx.array(keep_idx, mx.int32)
         self.y = self.y[keep_idx]
         for c in self.cache:
@@ -860,6 +867,9 @@ class Batch:
         self.logprobs.extend(other.logprobs)
         self.num_tokens.extend(other.num_tokens)
         self.max_tokens.extend(other.max_tokens)
+        self.samplers.extend(other.samplers)
+        self.logits_processors.extend(other.logits_processors)
+        self.tokens.extend(other.tokens)
         for c, o in zip(self.cache, other.cache):
             c.extend(o)
 
@@ -874,7 +884,7 @@ def _make_cache(model, left_padding):
     """
 
     def to_batch_cache(c):
-        if isinstance(c, KVCache):
+        if type(c) is KVCache:
             return BatchKVCache(left_padding)
         elif isinstance(c, ArraysCache):
             c.left_padding = mx.array(left_padding)
@@ -912,7 +922,6 @@ def _merge_caches(caches):
 
 
 class BatchGenerator:
-
     @dataclass
     class Response:
         uid: int
@@ -927,6 +936,9 @@ class BatchGenerator:
         max_tokens: int = 128,
         stop_tokens: Optional[set] = None,
         sampler: Optional[Callable[[mx.array], mx.array]] = None,
+        logits_processors: Optional[
+            List[Callable[[mx.array, mx.array], mx.array]]
+        ] = None,
         completion_batch_size: int = 32,
         prefill_batch_size: int = 8,
         prefill_step_size: int = 2048,
@@ -939,6 +951,7 @@ class BatchGenerator:
         self.max_tokens = max_tokens
         self.stop_tokens = stop_tokens or set()
         self.sampler = sampler or (lambda x: mx.argmax(x, axis=-1))
+        self.logits_processors = logits_processors or []
         self.uid_count = 0
         self.prefill_step_size = prefill_step_size
         self.prefill_batch_size = prefill_batch_size
@@ -965,7 +978,12 @@ class BatchGenerator:
         self.close()
 
     def insert(
-        self, prompts, max_tokens: Union[List[int], int, None] = None, caches=None
+        self,
+        prompts,
+        max_tokens: Union[List[int], int, None] = None,
+        caches=None,
+        samplers: list | None = None,
+        logits_processors: list | None = None,
     ):
         uids = []
 
@@ -978,8 +996,13 @@ class BatchGenerator:
             if caches[i] is None:
                 caches[i] = cache.make_prompt_cache(self.model)
 
-        for p, m, c in zip(prompts, max_tokens, caches):
-            self.unprocessed_prompts.append((self.uid_count, p, m, c))
+        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
+        ):
+            self.unprocessed_prompts.append((self.uid_count, p, m, c, s, lp))
             uids.append(self.uid_count)
             self.uid_count += 1
         # Sort in ascending order of length
@@ -1003,7 +1026,7 @@ class BatchGenerator:
                 self.unprocessed_prompts.pop(i)
 
     def _process_prompts(self, prompts):
-        uids, inputs, max_tokens, caches = zip(*prompts)
+        uids, inputs, max_tokens, caches, samplers, logits_processors = zip(*prompts)
 
         cache_lengths = [cache.cache_length(c) for c in caches]
         max_cache_length = max(cache_lengths)
@@ -1013,6 +1036,7 @@ class BatchGenerator:
 
         self._stats.prompt_tokens += sum(lengths)
 
+        tokens = [mx.array(inp) for inp in inputs]
         processed_tokens = 0
 
         # New prompts so
@@ -1069,17 +1093,56 @@ class BatchGenerator:
             mx.clear_cache()
             inputs = last_inputs
 
-        y, logprobs = self._step(inputs, prompt_cache)
+        y, logprobs = self._step(
+            inputs, prompt_cache, samplers, logits_processors, tokens
+        )
         mx.async_eval(y, logprobs)
+
         return Batch(
-            list(uids), y, logprobs, list(max_tokens), [0] * len(uids), prompt_cache
+            list(uids),
+            y,
+            logprobs,
+            list(max_tokens),
+            [0] * len(uids),
+            prompt_cache,
+            list(samplers),
+            list(logits_processors),
+            tokens,
         )
 
-    def _step(self, input_tokens: mx.array, prompt_cache: List[Any]):
+    def _step(
+        self,
+        input_tokens: mx.array,
+        prompt_cache: List[Any],
+        samplers: list | None,
+        logits_processors: list | None,
+        tokens: List[mx.array],
+    ):
+        batch_size = input_tokens.shape[0]
+
         logits = self.model(input_tokens, cache=prompt_cache)
         logits = logits[:, -1, :]
+
+        if any(logits_processors):
+            processed_logits = []
+            for e in range(batch_size):
+                sample_logits = logits[e : e + 1]
+                for processor in logits_processors[e]:
+                    sample_logits = processor(tokens[e], sample_logits)
+                processed_logits.append(sample_logits)
+            logits = mx.concatenate(processed_logits, axis=0)
+
         logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
-        sampled = self.sampler(logprobs)
+        if any(samplers):
+            all_samples = []
+            for e in range(batch_size):
+                sample_sampler = samplers[e] or self.sampler
+                sampled = sample_sampler(logprobs[e : e + 1])
+                all_samples.append(sampled)
+            sampled = mx.concatenate(all_samples, axis=0)
+        else:
+            sampled = self.sampler(logprobs)
+
         return sampled, list(logprobs)
 
     def stats(self):
@@ -1129,7 +1192,16 @@ class BatchGenerator:
 
         batch = self.active_batch
         y, logprobs = batch.y, batch.logprobs
-        batch.y, batch.logprobs = self._step(y[:, None], batch.cache)
+        for i, toks in enumerate(batch.tokens):
+            batch.tokens[i] = mx.concatenate((toks, y[i : i + 1]))
+        batch.y, batch.logprobs = self._step(
+            y[:, None],
+            batch.cache,
+            batch.samplers,
+            batch.logits_processors,
+            batch.tokens,
+        )
+
         mx.async_eval(batch.y, batch.logprobs)
 
         y = y.tolist()
@@ -1184,6 +1256,7 @@ def batch_generate(
     max_tokens: Union[int, List[int]] = 128,
     verbose: bool = False,
     return_prompt_caches: bool = False,
+    logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = None,
     **kwargs,
 ) -> BatchResponse:
     """
@@ -1202,11 +1275,17 @@ def batch_generate(
           can be per prompt if a list is provided.
        return_prompt_caches (bool): Return the prompt caches in the batch
           responses. Default: ``False``.
+       logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional):
+          A list of functions that take tokens and logits and return the processed logits. Default: ``None``.
        kwargs: The remaining options get passed to :obj:`BatchGenerator`.
           See :obj:`BatchGenerator` for more details.
     """
 
-    gen = BatchGenerator(model, stop_tokens=tokenizer.eos_token_ids, **kwargs)
+    gen = BatchGenerator(
+        model,
+        stop_tokens=tokenizer.eos_token_ids,
+        **kwargs,
+    )
     num_samples = len(prompts)
     fin = 0
     if verbose:
@@ -1302,15 +1381,9 @@ def main():
     if args.chat_template_config is not None:
         template_kwargs = json.loads(args.chat_template_config)
 
-    if args.use_default_chat_template:
-        if tokenizer.chat_template is None:
-            tokenizer.chat_template = tokenizer.default_chat_template
-    elif using_cache:
-        tokenizer.chat_template = json.loads(metadata["chat_template"])
-
     prompt = args.prompt.replace("\\n", "\n").replace("\\t", "\t")
     prompt = sys.stdin.read() if prompt == "-" else prompt
-    if not args.ignore_chat_template and tokenizer.chat_template is not None:
+    if not args.ignore_chat_template and tokenizer.has_chat_template:
         if args.system_prompt is not None:
             messages = [{"role": "system", "content": args.system_prompt}]
         else:

mlx_lm/models/deepseek_v2.py

@@ -6,6 +6,7 @@ from typing import Any, Dict, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
+from mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients
 
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .pipeline import PipelineMixin
@@ -315,13 +316,21 @@ class DeepseekV2MoE(nn.Module):
                 config=config, intermediate_size=intermediate_size
             )
 
+        self.sharding_group = None
+
     def __call__(self, x):
+        if self.sharding_group is not None:
+            x = sum_gradients(self.sharding_group)(x)
+
         inds, scores = self.gate(x)
         y = self.switch_mlp(x, inds)
         y = (y * scores[..., None]).sum(axis=-2)
         if self.config.n_shared_experts is not None:
             y = y + self.shared_experts(x)
 
+        if self.sharding_group is not None:
+            y = mx.distributed.all_sum(y, group=self.sharding_group)
+
         return y
 
 
@@ -395,7 +404,8 @@ class DeepseekV2Model(PipelineMixin, nn.Module):
                 cache[-1].keys = mx.depends(cache[-1].keys, h)
 
         # Broadcast h while keeping it in the graph
-        h = mx.distributed.all_gather(h)[: h.shape[0]]
+        if pipeline_size > 1:
+            h = mx.distributed.all_gather(h)[: h.shape[0]]
 
         return self.norm(h)
 
@@ -429,6 +439,62 @@ class Model(nn.Module):
                         weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
         return weights
 
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        group = group or mx.distributed.init()
+        N = group.size()
+        for layer in self.model.layers:
+            # Shard the self attention
+            if layer.self_attn.q_lora_rank is None:
+                layer.self_attn.q_proj = shard_linear(
+                    layer.self_attn.q_proj, "all-to-sharded", group=group
+                )
+            else:
+                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
+
+            # Shard the MLP
+            if isinstance(layer.mlp, DeepseekV2MLP):
+                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
+                )
+
+            # Shard the MoE. Shard in place since the MoE should be responsible
+            # for aggregating the results.
+            else:
+                layer.mlp.sharding_group = group
+                shard_inplace(
+                    layer.mlp.shared_experts.gate_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.shared_experts.down_proj, "sharded-to-all", group=group
+                )
+                shard_inplace(
+                    layer.mlp.shared_experts.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.model.pipeline_layers

mlx_lm/models/deepseek_v3.py

@@ -7,6 +7,7 @@ from typing import Any, Dict, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
+from mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients
 
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .pipeline import PipelineMixin
@@ -256,13 +257,21 @@ class DeepseekV3MoE(nn.Module):
                 config=config, intermediate_size=intermediate_size
             )
 
+        self.sharding_group = None
+
     def __call__(self, x):
+        if self.sharding_group is not None:
+            x = sum_gradients(self.sharding_group)(x)
+
         inds, scores = self.gate(x)
         y = self.switch_mlp(x, inds)
         y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype)
         if self.config.n_shared_experts is not None:
             y = y + self.shared_experts(x)
 
+        if self.sharding_group is not None:
+            y = mx.distributed.all_sum(y, group=self.sharding_group)
+
         return y
 
 
@@ -335,7 +344,8 @@ class DeepseekV3Model(PipelineMixin, nn.Module):
                 cache[-1].keys = mx.depends(cache[-1].keys, h)
 
         # Broadcast h while keeping it in the graph
-        h = mx.distributed.all_gather(h)[: h.shape[0]]
+        if pipeline_size > 1:
+            h = mx.distributed.all_gather(h)[: h.shape[0]]
 
         return self.norm(h)
 
@@ -419,6 +429,62 @@ class Model(nn.Module):
             if not k.startswith("model.layers.61") and "rotary_emb.inv_freq" not in k
         }
 
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        group = group or mx.distributed.init()
+        N = group.size()
+        for layer in self.model.layers:
+            # Shard the self attention
+            if layer.self_attn.q_lora_rank is None:
+                layer.self_attn.q_proj = shard_linear(
+                    layer.self_attn.q_proj, "all-to-sharded", group=group
+                )
+            else:
+                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
+
+            # Shard the MLP
+            if isinstance(layer.mlp, DeepseekV3MLP):
+                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
+                )
+
+            # Shard the MoE. Shard in place since the MoE should be responsible
+            # for aggregating the results.
+            else:
+                layer.mlp.sharding_group = group
+                shard_inplace(
+                    layer.mlp.shared_experts.gate_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.shared_experts.down_proj, "sharded-to-all", group=group
+                )
+                shard_inplace(
+                    layer.mlp.shared_experts.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.model.pipeline_layers

mlx_lm/models/deepseek_v32.py

@@ -6,6 +6,7 @@ from typing import Any, Dict, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
+from mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients
 
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .cache import CacheList, KVCache
@@ -222,6 +223,11 @@ class DeepseekV32Attention(nn.Module):
             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))
+
         output = scaled_dot_product_attention(
             queries, keys, values, cache=cache[0], scale=self.scale, mask=mask
         )
@@ -328,13 +334,21 @@ class DeepseekV32MoE(nn.Module):
                 config=config, intermediate_size=intermediate_size
             )
 
+        self.sharding_group = None
+
     def __call__(self, x):
+        if self.sharding_group is not None:
+            x = sum_gradients(self.sharding_group)(x)
+
         inds, scores = self.gate(x)
         y = self.switch_mlp(x, inds)
         y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype)
         if self.config.n_shared_experts is not None:
             y = y + self.shared_experts(x)
 
+        if self.sharding_group is not None:
+            y = mx.distributed.all_sum(y, group=self.sharding_group)
+
         return y
 
 
@@ -428,10 +442,11 @@ class DeepseekV32Model(nn.Module):
         if pipeline_rank != 0:
             h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
             if cache[-1] is not None:
-                cache[-1].keys = mx.depends(cache[-1].keys, h)
+                cache[-1][0].keys = mx.depends(cache[-1][0].keys, h)
 
         # Broadcast h while keeping it in the graph
-        h = mx.distributed.all_gather(h)[: h.shape[0]]
+        if pipeline_size > 1:
+            h = mx.distributed.all_gather(h)[: h.shape[0]]
 
         return self.norm(h)
 
@@ -500,6 +515,56 @@ class Model(nn.Module):
             if not k.startswith("model.layers.61") and "rotary_emb.inv_freq" not in k
         }
 
+    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_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
+
+            # Shard the MLP
+            if isinstance(layer.mlp, DeepseekV32MLP):
+                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
+                )
+
+            # Shard the MoE. Shard in place since the MoE should be responsible
+            # for aggregating the results.
+            else:
+                layer.mlp.sharding_group = group = group
+                shard_inplace(
+                    layer.mlp.shared_experts.gate_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.shared_experts.down_proj, "sharded-to-all", group=group
+                )
+                shard_inplace(
+                    layer.mlp.shared_experts.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.model.layers[self.model.start_idx : self.model.end_idx]

mlx_lm/models/exaone_moe.py

@@ -0,0 +1,438 @@
+# Copyright © 2026 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, 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 .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 SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    vocab_size: int
+    hidden_size: int
+    intermediate_size: int
+    moe_intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    head_dim: int
+    num_experts: int
+    num_experts_per_tok: int
+    num_shared_experts: int
+    rms_norm_eps: float
+    max_position_embeddings: int
+    sliding_window: int
+    layer_types: List[str]
+    is_moe_layer: List[bool]
+    n_group: int = 1
+    topk_group: int = 1
+    routed_scaling_factor: float = 2.5
+    norm_topk_prob: bool = True
+    scoring_func: str = "sigmoid"
+    topk_method: str = "noaux_tc"
+    rope_theta: float = 1000000.0
+    rope_scaling: Optional[dict] = None
+    rope_parameters: Optional[dict] = None
+    tie_word_embeddings: bool = False
+
+    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"]
+
+
+@mx.compile
+def group_expert_select(
+    gates,
+    e_score_correction_bias,
+    top_k,
+    n_group,
+    topk_group,
+    routed_scaling_factor,
+    norm_topk_prob,
+):
+    scores = mx.sigmoid(gates.astype(mx.float32))
+    orig_scores = scores
+    scores = scores + e_score_correction_bias
+
+    if n_group > 1:
+        scores = mx.unflatten(scores, axis=-1, shape=(n_group, -1))
+        group_scores = mx.topk(scores, 2, axis=-1).sum(axis=-1, keepdims=True)
+        k = n_group - topk_group
+        group_idx = mx.argpartition(group_scores, kth=k - 1, axis=-2)[..., :k, :]
+        scores = mx.put_along_axis(
+            scores, mx.stop_gradient(group_idx), mx.array(0.0), axis=-2
+        )
+        scores = mx.flatten(scores, -2, -1)
+
+    k = top_k
+    inds = mx.argpartition(-scores, kth=k - 1, axis=-1)[..., :k]
+    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 + 1e-20)
+
+    scores = scores * routed_scaling_factor
+    return inds, scores
+
+
+class MoEGate(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.top_k = args.num_experts_per_tok
+        self.norm_topk_prob = args.norm_topk_prob
+        self.n_routed_experts = args.num_experts
+        self.routed_scaling_factor = args.routed_scaling_factor
+        self.n_group = args.n_group
+        self.topk_group = args.topk_group
+        self.weight = mx.zeros((self.n_routed_experts, args.hidden_size))
+        self.e_score_correction_bias = mx.zeros((self.n_routed_experts,))
+        assert args.topk_method == "noaux_tc", "Unsupported topk method."
+
+    def __call__(self, x):
+        return group_expert_select(
+            x @ self.weight.T,
+            self.e_score_correction_bias,
+            self.top_k,
+            self.n_group,
+            self.topk_group,
+            self.routed_scaling_factor,
+            self.norm_topk_prob,
+        )
+
+
+class MLP(nn.Module):
+    def __init__(self, args: ModelArgs, intermediate_size: Optional[int] = None):
+        super().__init__()
+        hidden_size = args.hidden_size
+        intermediate_size = intermediate_size or args.intermediate_size
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+
+    def __call__(self, x):
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class MoE(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.switch_mlp = SwitchGLU(
+            args.hidden_size,
+            args.moe_intermediate_size,
+            args.num_experts,
+        )
+
+        self.gate = MoEGate(args)
+
+        self.shared_experts = (
+            MLP(
+                args,
+                intermediate_size=args.moe_intermediate_size * args.num_shared_experts,
+            )
+            if args.num_shared_experts is not None and args.num_shared_experts > 0
+            else None
+        )
+
+        self.sharding_group = None
+
+    def __call__(self, x):
+        if self.sharding_group is not None:
+            x = sum_gradients(self.sharding_group)(x)
+
+        inds, scores = self.gate(x)
+        y = self.switch_mlp(x, inds)
+        y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype)
+        if self.shared_experts is not None:
+            y = y + self.shared_experts(x)
+
+        if self.sharding_group is not None:
+            y = mx.distributed.all_sum(y, group=self.sharding_group)
+
+        return y
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+
+        self.hidden_size = args.hidden_size
+        self.n_heads = args.num_attention_heads
+        self.n_kv_heads = args.num_key_value_heads
+        self.head_dim = args.head_dim
+        self.scale = self.head_dim**-0.5
+
+        self.q_proj = nn.Linear(
+            self.hidden_size, self.n_heads * self.head_dim, bias=False
+        )
+        self.k_proj = nn.Linear(
+            self.hidden_size, self.n_kv_heads * self.head_dim, bias=False
+        )
+        self.v_proj = nn.Linear(
+            self.hidden_size, self.n_kv_heads * self.head_dim, bias=False
+        )
+        self.o_proj = nn.Linear(
+            self.n_heads * self.head_dim, self.hidden_size, bias=False
+        )
+
+        self.q_norm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps)
+        self.k_norm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps)
+
+        self.is_sliding_window = args.layer_types[layer_idx] == "sliding_attention"
+        self.apply_rope_all_layers = "sliding_attention" not in args.layer_types
+        self.use_rope = self.is_sliding_window or self.apply_rope_all_layers
+
+        if self.use_rope:
+            self.rope = initialize_rope(
+                self.head_dim,
+                base=args.rope_theta,
+                traditional=False,
+                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, D = 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.n_heads, -1)).transpose(
+            0, 2, 1, 3
+        )
+        keys = self.k_norm(keys.reshape(B, L, self.n_kv_heads, -1)).transpose(
+            0, 2, 1, 3
+        )
+        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            if self.use_rope:
+                queries = self.rope(queries, offset=cache.offset)
+                keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        elif self.use_rope:
+            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 DecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+
+        self.self_attn = Attention(args, layer_idx)
+        self.mlp = MoE(args) if args.is_moe_layer[layer_idx] else MLP(args)
+        self.is_sliding_window = self.self_attn.is_sliding_window
+
+        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
+        )
+
+    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
+        r = self.mlp(self.post_attention_layernorm(h))
+        return h + r
+
+
+class ExaoneMoEModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.vocab_size = args.vocab_size
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [DecoderLayer(args, idx) for idx in range(args.num_hidden_layers)]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+        self.swa_idx = None
+        self.ga_idx = None
+        for i, layer in enumerate(self.layers):
+            if layer.is_sliding_window and self.swa_idx is None:
+                self.swa_idx = i
+            if not layer.is_sliding_window and self.ga_idx is None:
+                self.ga_idx = i
+
+        self.window_size = args.sliding_window
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        h = self.embed_tokens(inputs)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        global_mask = create_attention_mask(
+            h, cache[self.ga_idx] if self.ga_idx is not None else cache[0]
+        )
+        swa_mask = create_attention_mask(
+            h,
+            cache[self.swa_idx] if self.swa_idx is not None else cache[0],
+            window_size=self.window_size,
+        )
+
+        for layer, c in zip(self.layers, cache):
+            mask = swa_mask if layer.is_sliding_window else global_mask
+            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 = ExaoneMoEModel(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,
+    ):
+        out = self.model(inputs, cache)
+        if self.args.tie_word_embeddings:
+            out = self.model.embed_tokens.as_linear(out)
+        else:
+            out = self.lm_head(out)
+        return out
+
+    def sanitize(self, weights):
+        new_weights = {k: v for k, v in weights.items() if not k.startswith("mtp.")}
+        weights = new_weights
+
+        for l in range(self.args.num_hidden_layers):
+            if not self.args.is_moe_layer[l]:
+                continue
+
+            prefix = f"model.layers.{l}"
+
+            bias_key = f"{prefix}.mlp.e_score_correction_bias"
+            if bias_key in weights:
+                weights[f"{prefix}.mlp.gate.e_score_correction_bias"] = weights.pop(
+                    bias_key
+                )
+
+            for m in ["gate_proj", "down_proj", "up_proj"]:
+                for k in ["weight", "scales", "biases"]:
+                    first_key = f"{prefix}.mlp.experts.0.{m}.{k}"
+                    last_key = (
+                        f"{prefix}.mlp.experts.{self.args.num_experts - 1}.{m}.{k}"
+                    )
+                    if first_key in weights and last_key in weights:
+                        to_join = [
+                            weights.pop(f"{prefix}.mlp.experts.{e}.{m}.{k}")
+                            for e in range(self.args.num_experts)
+                        ]
+                        weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
+
+        if self.args.tie_word_embeddings:
+            weights.pop("lm_head.weight", None)
+
+        return weights
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    @property
+    def cast_predicate(self):
+        def predicate(k):
+            return "e_score_correction_bias" not in k
+
+        return predicate
+
+    def make_cache(self):
+        caches = []
+        for layer in self.layers:
+            if layer.is_sliding_window:
+                caches.append(
+                    RotatingKVCache(max_size=self.args.sliding_window, keep=0)
+                )
+            else:
+                caches.append(KVCache())
+        return caches
+
+    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.n_heads //= N
+            layer.self_attn.n_kv_heads //= N
+
+            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
+                )
+            else:
+                layer.mlp.sharding_group = group
+                if layer.mlp.shared_experts is not None:
+                    shard_inplace(
+                        layer.mlp.shared_experts.gate_proj,
+                        "all-to-sharded",
+                        group=group,
+                    )
+                    shard_inplace(
+                        layer.mlp.shared_experts.down_proj,
+                        "sharded-to-all",
+                        group=group,
+                    )
+                    shard_inplace(
+                        layer.mlp.shared_experts.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
+                )

mlx_lm/models/gemma3_text.py

@@ -54,13 +54,20 @@ class Attention(nn.Module):
         self.k_norm = RMSNorm(dims=head_dim, eps=args.rms_norm_eps)
         self.is_sliding = (layer_idx + 1) % args.sliding_window_pattern != 0
 
-        self.rope = initialize_rope(
-            dims=head_dim,
-            base=(args.rope_local_base_freq if self.is_sliding else args.rope_theta),
-            traditional=False,
-            max_position_embeddings=args.max_position_embeddings,
-            scaling_config=args.rope_scaling,
-        )
+        if self.is_sliding:
+            self.rope = initialize_rope(
+                dims=head_dim,
+                base=args.rope_local_base_freq,
+                traditional=False,
+            )
+        else:
+            self.rope = initialize_rope(
+                dims=head_dim,
+                base=args.rope_theta,
+                traditional=False,
+                max_position_embeddings=args.max_position_embeddings,
+                scaling_config=args.rope_scaling,
+            )
 
     def __call__(
         self,

mlx_lm/models/glm4_moe.py

@@ -7,6 +7,7 @@ from typing import Any, Dict, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
+from mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients
 
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .pipeline import PipelineMixin
@@ -205,13 +206,21 @@ class MoE(nn.Module):
                 config=config, intermediate_size=intermediate_size
             )
 
+        self.sharding_group = None
+
     def __call__(self, x):
+        if self.sharding_group is not None:
+            x = sum_gradients(self.sharding_group)(x)
+
         inds, scores = self.gate(x)
         y = self.switch_mlp(x, inds)
         y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype)
         if self.config.n_shared_experts is not None:
             y = y + self.shared_experts(x)
 
+        if self.sharding_group is not None:
+            y = mx.distributed.all_sum(y, group=self.sharding_group)
+
         return y
 
 
@@ -252,10 +261,6 @@ class LanguageModel(PipelineMixin, nn.Module):
         self.layers = [
             DecoderLayer(config, idx) for idx in range(config.num_hidden_layers)
         ]
-        self.start_idx = 0
-        self.end_idx = len(self.layers)
-        self.num_layers = self.end_idx
-
         self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
     def __call__(
@@ -286,7 +291,8 @@ class LanguageModel(PipelineMixin, nn.Module):
                 cache[-1].keys = mx.depends(cache[-1].keys, h)
 
         # Broadcast h while keeping it in the graph
-        h = mx.distributed.all_gather(h)[: h.shape[0]]
+        if pipeline_size > 1:
+            h = mx.distributed.all_gather(h)[: h.shape[0]]
 
         return self.norm(h)
 
@@ -329,6 +335,61 @@ class Model(nn.Module):
             if not k.startswith(f"model.layers.{mpt_layer}")
         }
 
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        group = group or mx.distributed.init()
+        N = group.size()
+        for layer in self.model.layers:
+            # Shard the self attention
+            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.n_heads //= N
+            layer.self_attn.n_kv_heads //= N
+
+            # Shard the 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
+                )
+
+            # Shard the MoE. Shard in place since the MoE should be responsible
+            # for aggregating the results.
+            else:
+                layer.mlp.sharding_group = group
+                shard_inplace(
+                    layer.mlp.shared_experts.gate_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.shared_experts.down_proj, "sharded-to-all", group=group
+                )
+                shard_inplace(
+                    layer.mlp.shared_experts.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.model.pipeline_layers

mlx_lm/models/llama.py

@@ -5,6 +5,7 @@ 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_linear
 
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .cache import KVCache, RotatingKVCache
@@ -226,6 +227,37 @@ class Model(nn.Module):
             weights.pop("lm_head.weight", None)
         return weights
 
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        group = group or mx.distributed.init()
+        N = group.size()
+        for layer in self.model.layers:
+            # Shard the self attention
+            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.n_heads //= N
+            layer.self_attn.n_kv_heads //= N
+
+            # Shard the 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
+            )
+
     @property
     def layers(self):
         return self.model.layers

mlx_lm/models/mimo_v2_flash.py

@@ -0,0 +1,382 @@
+# Copyright © 2024 Apple Inc.
+
+import math
+from dataclasses import dataclass
+from functools import partial
+from typing import Any, Dict, List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .cache import KVCache, RotatingKVCache
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    num_experts_per_tok: int
+    hybrid_layer_pattern: List[int]
+    moe_layer_freq: List[int]
+    add_swa_attention_sink_bias: bool
+    add_full_attention_sink_bias: bool
+    sliding_window_size: int
+    vocab_size: int
+    hidden_size: 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: Optional[float]
+    topk_method: str
+    scoring_func: str
+    norm_topk_prob: bool
+    n_group: int
+    topk_group: int
+    max_position_embeddings: int
+    layernorm_epsilon: float
+    rope_theta: float
+    swa_rope_theta: float
+    swa_num_attention_heads: int
+    swa_num_key_value_heads: int
+    head_dim: int
+    v_head_dim: int
+    swa_head_dim: int
+    swa_v_head_dim: int
+    partial_rotary_factor: int
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs, is_sliding_window: bool):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.is_sliding_window = is_sliding_window
+        if self.is_sliding_window:
+            self.n_heads = n_heads = args.swa_num_attention_heads
+            self.n_kv_heads = n_kv_heads = args.swa_num_key_value_heads
+            self.has_sinks = args.add_swa_attention_sink_bias
+            head_dim = args.swa_head_dim
+            v_head_dim = args.swa_v_head_dim
+            rope_theta = args.swa_rope_theta
+        else:
+            self.n_heads = n_heads = args.num_attention_heads
+            self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+            self.has_sinks = args.add_full_attention_sink_bias
+            head_dim = args.head_dim
+            v_head_dim = args.v_head_dim
+            rope_theta = args.rope_theta
+
+        self.scale = head_dim**-0.5
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=False)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
+        self.v_proj = nn.Linear(dim, n_kv_heads * v_head_dim, bias=False)
+        self.o_proj = nn.Linear(n_heads * v_head_dim, dim, bias=False)
+        if self.has_sinks:
+            self.attention_sink_bias = mx.ones((self.n_heads,))
+        else:
+            self.attention_sink_bias = None
+
+        self.rope = nn.RoPE(
+            int(args.partial_rotary_factor * head_dim),
+            traditional=False,
+            base=rope_theta,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.n_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,
+            sinks=self.attention_sink_bias,
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class MLP(nn.Module):
+    def __init__(
+        self, config: ModelArgs, hidden_size: int = None, intermediate_size: int = None
+    ):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size if hidden_size is None else hidden_size
+        self.intermediate_size = (
+            config.intermediate_size if intermediate_size is None else 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)
+
+    def __call__(self, x):
+        down_proj = self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+
+@mx.compile
+def group_expert_select(
+    gates,
+    e_score_correction_bias,
+    top_k,
+    n_group,
+    topk_group,
+    routed_scaling_factor,
+    norm_topk_prob,
+):
+
+    scores = mx.sigmoid(gates.astype(mx.float32))
+    orig_scores = scores
+    scores = scores + e_score_correction_bias
+    if n_group > 1:
+        scores = mx.unflatten(scores, axis=-1, shape=(n_group, -1))
+        group_scores = mx.topk(scores, 2, axis=-1).sum(axis=-1, keepdims=True)
+        k = n_group - topk_group
+        group_idx = mx.argpartition(group_scores, kth=k - 1, axis=-2)[..., :k, :]
+        scores = mx.put_along_axis(
+            scores, mx.stop_gradient(group_idx), mx.array(0.0), axis=-2
+        )
+        scores = mx.flatten(scores, -2, -1)
+
+    k = top_k
+    inds = mx.argpartition(-scores, kth=k - 1, axis=-1)[..., :k]
+    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 + 1e-20)
+    scores = scores * routed_scaling_factor
+
+    return inds, scores
+
+
+class MoEGate(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.top_k = config.num_experts_per_tok
+        self.norm_topk_prob = config.norm_topk_prob
+        self.n_routed_experts = config.n_routed_experts
+        self.routed_scaling_factor = (
+            config.routed_scaling_factor
+            if config.routed_scaling_factor is not None
+            else 1.0
+        )
+        self.n_group = config.n_group
+        self.topk_group = config.topk_group
+        self.weight = mx.zeros((self.n_routed_experts, config.hidden_size))
+        self.e_score_correction_bias = mx.zeros((self.n_routed_experts,))
+        assert config.topk_method == "noaux_tc", "Unsupported topk method."
+
+    def __call__(self, x):
+        return group_expert_select(
+            x @ self.weight.T,
+            self.e_score_correction_bias,
+            self.top_k,
+            self.n_group,
+            self.topk_group,
+            self.routed_scaling_factor,
+            self.norm_topk_prob,
+        )
+
+
+class MoE(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.num_experts_per_tok = config.num_experts_per_tok
+        self.switch_mlp = SwitchGLU(
+            config.hidden_size,
+            config.moe_intermediate_size,
+            config.n_routed_experts,
+        )
+
+        self.gate = MoEGate(config)
+        if config.n_shared_experts is not None:
+            intermediate_size = config.moe_intermediate_size * config.n_shared_experts
+            self.shared_experts = MLP(
+                config=config, intermediate_size=intermediate_size
+            )
+
+    def __call__(self, x):
+        inds, scores = self.gate(x)
+        y = self.switch_mlp(x, inds)
+        y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype)
+        if self.config.n_shared_experts is not None:
+            y = y + self.shared_experts(x)
+
+        return y
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, config: ModelArgs, is_moe, is_sliding_window):
+        super().__init__()
+        self.self_attn = Attention(config, is_sliding_window)
+        self.mlp = MoE(config) if is_moe else MLP(config)
+        self.is_sliding_window = is_sliding_window
+        self.input_layernorm = nn.RMSNorm(
+            config.hidden_size, eps=config.layernorm_epsilon
+        )
+        self.post_attention_layernorm = nn.RMSNorm(
+            config.hidden_size, eps=config.layernorm_epsilon
+        )
+
+    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
+        r = self.mlp(self.post_attention_layernorm(h))
+        return h + r
+
+
+class LanguageModel(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.vocab_size = config.vocab_size
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = [
+            DecoderLayer(
+                config,
+                is_moe=config.moe_layer_freq[idx] == 1,
+                is_sliding_window=config.hybrid_layer_pattern[idx] == 1,
+            )
+            for idx in range(config.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(config.hidden_size, eps=config.layernorm_epsilon)
+        self.swa_idx = config.hybrid_layer_pattern.index(1)
+        self.ga_idx = config.hybrid_layer_pattern.index(0)
+        self.sliding_window_size = config.sliding_window_size
+
+    def __call__(
+        self,
+        x: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        h = self.embed_tokens(x)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        full_mask = create_attention_mask(x, cache[self.ga_idx])
+        swa_mask = create_attention_mask(
+            x, cache[self.swa_idx], window_size=self.sliding_window_size
+        )
+
+        for l, c in zip(self.layers, cache):
+            mask = swa_mask if l.is_sliding_window else full_mask
+            h = l(h, mask, cache=c)
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.args = config
+        self.model_type = config.model_type
+        self.model = LanguageModel(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)
+
+    def sanitize(self, weights):
+        def dequant(weight, scale_inv):
+            dtype = weight.dtype
+            bs = 128  # block size
+            m, n = weight.shape
+            pad_bottom = bs * scale_inv.shape[0] - m
+            pad_side = bs * scale_inv.shape[1] - n
+            weight = mx.pad(weight, ((0, pad_bottom), (0, pad_side)))
+            weight = weight.reshape(
+                ((m + pad_bottom) // bs, bs, (n + pad_side) // bs, bs)
+            )
+            weight = (weight * scale_inv[:, None, :, None]).reshape(
+                m + pad_bottom, n + pad_side
+            )
+            return weight[:m, :n].astype(dtype)
+
+        # Dequantize fp8
+        new_weights = {}
+        for k, v in weights.items():
+            if "weight_scale_inv" in k:
+                scale_inv = v
+                wk = k.replace("_scale_inv", "")
+                weight = weights[wk]
+                weight = dequant(weight, scale_inv)
+                new_weights[wk] = weight
+            elif k not in new_weights:
+                new_weights[k] = v
+        weights = new_weights
+
+        # Stack experts
+        for l in range(self.args.num_hidden_layers):
+            prefix = f"model.layers.{l}"
+            for n, m in [("w1", "gate_proj"), ("w2", "down_proj"), ("w3", "up_proj")]:
+                for k in ["weight", "scales", "biases"]:
+                    if f"{prefix}.mlp.experts.0.{m}.{k}" in weights:
+                        to_join = [
+                            weights.pop(f"{prefix}.mlp.experts.{e}.{m}.{k}")
+                            for e in range(self.args.n_routed_experts)
+                        ]
+                        weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
+
+        # Remove multi-token prediction layer
+        return {k: v for k, v in weights.items() if not k.startswith("model.mtp")}
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    @property
+    def cast_predicate(self):
+        def predicate(k):
+            return "e_score_correction_bias" not in k
+
+        return predicate
+
+    def make_cache(self):
+        caches = []
+        for l in self.layers:
+            if l.is_sliding_window:
+                caches.append(RotatingKVCache(max_size=self.args.sliding_window_size))
+            else:
+                caches.append(KVCache())
+        return caches

mlx_lm/models/ministral3.py

@@ -5,9 +5,11 @@ 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_linear
 
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .cache import KVCache, RotatingKVCache
+from .pipeline import PipelineMixin
 from .rope_utils import initialize_rope
 
 
@@ -36,13 +38,17 @@ class ModelArgs(BaseModelArgs):
             self.layer_types = ["full_attention"] * self.num_hidden_layers
 
 
-def _get_llama_4_attn_scale(
-    start: int, stop: int, beta: float, max_position_embeddings: int
-):
+def _get_llama_4_attn_scale(size, offset, beta: float, max_position_embeddings: int):
+    if isinstance(offset, mx.array) and offset.ndim > 0:
+        offset = offset[:, None]
+
     scaling = 1 + beta * mx.log(
-        1 + mx.floor(mx.arange(start, stop) / max_position_embeddings)
+        1 + mx.floor((mx.arange(size) + offset) / max_position_embeddings)
     )
-    return scaling[:, None]
+    if scaling.ndim == 2:
+        return scaling[:, None, :, None]
+    else:
+        return scaling[:, None]
 
 
 class Attention(nn.Module):
@@ -146,7 +152,7 @@ class TransformerBlock(nn.Module):
         return out
 
 
-class LanguageModel(nn.Module):
+class LanguageModel(PipelineMixin, nn.Module):
     def __init__(self, args: ModelArgs):
         super().__init__()
         self.args = args
@@ -167,6 +173,18 @@ class LanguageModel(nn.Module):
                 self.swa_idx = e
                 break
 
+    def pipeline(self, group):
+        super().pipeline(group)
+        self.fa_idx = None
+        self.swa_idx = None
+        for e, l in enumerate(self.pipeline_layers):
+            if self.swa_idx is None and l.use_sliding:
+                self.swa_idx = e
+            elif self.fa_idx is None and not l.use_sliding:
+                self.fa_idx = e
+            if self.fa_idx is not None and self.swa_idx is not None:
+                break
+
     def __call__(
         self,
         inputs: mx.array,
@@ -178,28 +196,47 @@ class LanguageModel(nn.Module):
         else:
             h = self.embed_tokens(inputs)
 
+        pipeline_rank = self.pipeline_rank
+        pipeline_size = self.pipeline_size
+
         if cache is None:
-            cache = [None] * len(self.layers)
+            cache = [None] * len(self.pipeline_layers)
             offset = 0
         else:
             offset = cache[0].offset
 
-        fa_mask = create_attention_mask(h, cache[self.fa_idx])
+        swa_mask = fa_mask = None
+        if self.fa_idx is not None:
+            fa_mask = create_attention_mask(h, cache[self.fa_idx])
         if self.swa_idx is not None:
             swa_mask = create_attention_mask(
                 h, cache[self.swa_idx], window_size=self.sliding_window
             )
 
         attn_scale = _get_llama_4_attn_scale(
+            inputs.shape[1],
             offset,
-            offset + inputs.shape[1],
             self.args.rope_parameters["llama_4_scaling_beta"],
             self.args.rope_parameters["original_max_position_embeddings"],
         ).astype(h.dtype)
 
-        for layer, cache in zip(self.layers, cache):
-            mask = swa_mask if layer.use_sliding else fa_mask
-            h = layer(h, attn_scale, mask, cache=cache)
+        # Receive from the previous process in the pipeline
+        if pipeline_rank < pipeline_size - 1:
+            h = mx.distributed.recv_like(h, (pipeline_rank + 1))
+
+        for l, c in zip(self.pipeline_layers, cache):
+            mask = swa_mask if l.use_sliding else fa_mask
+            h = l(h, attn_scale, mask, cache=c)
+
+        # Send to the next process in the pipeline
+        if pipeline_rank != 0:
+            h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
+            if cache[-1] is not None:
+                cache[-1].keys = mx.depends(cache[-1].keys, h)
+
+        # Broadcast h while keeping it in the graph
+        if pipeline_size > 1:
+            h = mx.distributed.all_gather(h)[: h.shape[0]]
 
         return self.norm(h)
 
@@ -249,9 +286,40 @@ class Model(nn.Module):
 
         return weights
 
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        group = group or mx.distributed.init()
+        N = group.size()
+        for layer in self.model.layers:
+            # Shard the self attention
+            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.n_heads //= N
+            layer.self_attn.n_kv_heads //= N
+
+            # Shard the 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
+            )
+
     @property
     def layers(self):
-        return self.model.layers
+        return self.model.pipeline_layers
 
     def make_cache(self):
         return [

mlx_lm/models/nemotron_h.py

@@ -15,6 +15,7 @@ from .base import (
 )
 from .cache import KVCache, MambaCache
 from .ssm import ssm_update
+from .switch_layers import SwitchMLP
 
 
 @dataclass()
@@ -37,24 +38,34 @@ class ModelArgs(BaseModelArgs):
     time_step_limit: Tuple[float, float]
     mlp_bias: bool
     layer_norm_epsilon: float
-    rms_norm_eps: float
     use_bias: bool
     use_conv_bias: bool
-    residual_in_fp32: bool
     hybrid_override_pattern: List[str]
     head_dim: Optional[int] = None
+    moe_intermediate_size: Optional[int] = None
+    moe_shared_expert_intermediate_size: Optional[int] = None
+    n_group: Optional[int] = None
+    n_routed_experts: Optional[int] = None
+    n_shared_experts: Optional[int] = None
+    topk_group: Optional[int] = None
+    num_experts_per_tok: Optional[int] = None
+    norm_topk_prob: Optional[bool] = None
+    routed_scaling_factor: Optional[float] = None
 
 
 class MambaRMSNormGated(nn.Module):
-    def __init__(self, hidden_size: int, eps: float = 1e-6):
+    def __init__(self, hidden_size: int, eps: float, group_size: int):
         super().__init__()
         self.eps = eps
         self.weight = mx.ones(hidden_size)
+        self.group_size = group_size
 
-    def __call__(self, hidden_states: mx.array, gate: mx.array = None) -> mx.array:
+    def __call__(self, x: mx.array, gate: mx.array = None) -> mx.array:
         if gate is not None:
-            hidden_states = hidden_states * nn.silu(gate)
-        return mx.fast.rms_norm(hidden_states, self.weight, self.eps)
+            x = x * nn.silu(gate)
+        x = mx.unflatten(x, axis=-1, shape=(-1, self.group_size))
+        x = mx.fast.rms_norm(x, weight=None, eps=self.eps)
+        return self.weight * x.flatten(-2)
 
 
 class NemotronHMamba2Mixer(nn.Module):
@@ -90,8 +101,11 @@ class NemotronHMamba2Mixer(nn.Module):
         self.A_log = mx.log(mx.arange(1, self.num_heads + 1, dtype=mx.float32))
         self.D = mx.ones(self.num_heads)
 
+        group_size = self.intermediate_size // self.n_groups
         self.norm = MambaRMSNormGated(
-            self.intermediate_size, eps=args.layer_norm_epsilon
+            self.intermediate_size,
+            eps=args.layer_norm_epsilon,
+            group_size=group_size,
         )
         self.out_proj = nn.Linear(
             self.intermediate_size, self.hidden_size, bias=args.mamba_proj_bias
@@ -139,7 +153,7 @@ class NemotronHMamba2Mixer(nn.Module):
             self.A_log,
             B,
             C,
-            self.D,
+            self.D.astype(hidden_states.dtype),
             dt,
             self.dt_bias,
             state,
@@ -245,24 +259,113 @@ class NemotronHAttention(nn.Module):
 
 
 class NemotronHMLP(nn.Module):
-    def __init__(self, args: ModelArgs):
+    def __init__(self, args: ModelArgs, intermediate_size=None):
         super().__init__()
+        intermediate_size = intermediate_size or args.intermediate_size
+
         self.up_proj = nn.Linear(
-            args.hidden_size, args.intermediate_size, bias=args.mlp_bias
+            args.hidden_size, intermediate_size, bias=args.mlp_bias
         )
         self.down_proj = nn.Linear(
-            args.intermediate_size, args.hidden_size, bias=args.mlp_bias
+            intermediate_size, args.hidden_size, bias=args.mlp_bias
         )
 
     def __call__(self, x):
         return self.down_proj(nn.relu2(self.up_proj(x)))
 
 
+@mx.compile
+def group_expert_select(
+    gates,
+    e_score_correction_bias,
+    top_k,
+    n_group,
+    topk_group,
+    routed_scaling_factor,
+    norm_topk_prob,
+):
+
+    orig_scores = scores = mx.sigmoid(gates.astype(mx.float32))
+    scores = scores + e_score_correction_bias
+    if n_group > 1:
+        scores = mx.unflatten(scores, axis=-1, shape=(n_group, -1))
+        group_scores = mx.topk(scores, 2, axis=-1).sum(axis=-1, keepdims=True)
+        k = n_group - topk_group
+        group_idx = mx.argpartition(group_scores, kth=k - 1, axis=-2)[..., :k, :]
+        scores = mx.put_along_axis(
+            scores, mx.stop_gradient(group_idx), mx.array(0.0), axis=-2
+        )
+        scores = mx.flatten(scores, -2, -1)
+
+    k = top_k
+    inds = mx.argpartition(-scores, kth=k - 1, axis=-1)[..., :k]
+    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 + 1e-20)
+    scores = scores * routed_scaling_factor
+
+    return inds, scores
+
+
+class MoEGate(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.top_k = config.num_experts_per_tok
+        self.norm_topk_prob = config.norm_topk_prob
+        self.n_routed_experts = config.n_routed_experts
+        self.routed_scaling_factor = config.routed_scaling_factor
+        self.n_group = config.n_group
+        self.topk_group = config.topk_group
+        self.weight = mx.zeros((self.n_routed_experts, config.hidden_size))
+        self.e_score_correction_bias = mx.zeros((self.n_routed_experts,))
+
+    def __call__(self, x):
+        return group_expert_select(
+            x @ self.weight.T,
+            self.e_score_correction_bias,
+            self.top_k,
+            self.n_group,
+            self.topk_group,
+            self.routed_scaling_factor,
+            self.norm_topk_prob,
+        )
+
+
+class NemotronHMoE(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.num_experts_per_tok = config.num_experts_per_tok
+        self.switch_mlp = SwitchMLP(
+            config.hidden_size,
+            config.moe_intermediate_size,
+            config.n_routed_experts,
+            activation=nn.ReLU2(),
+        )
+
+        self.gate = MoEGate(config)
+        if config.n_shared_experts is not None:
+            intermediate_size = config.moe_shared_expert_intermediate_size
+            self.shared_experts = NemotronHMLP(
+                config, intermediate_size=intermediate_size
+            )
+
+    def __call__(self, x):
+        inds, scores = self.gate(x)
+        y = self.switch_mlp(x, inds)
+        y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype)
+        if self.config.n_shared_experts is not None:
+            y = y + self.shared_experts(x)
+
+        return y
+
+
 class NemotronHBlock(nn.Module):
     def __init__(self, args: ModelArgs, block_type: str):
         super().__init__()
-        self.residual_in_fp32 = args.residual_in_fp32
-        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.layer_norm_epsilon)
 
         self.block_type = block_type
 
@@ -272,6 +375,8 @@ class NemotronHBlock(nn.Module):
             self.mixer = NemotronHAttention(args)
         elif self.block_type == "-":
             self.mixer = NemotronHMLP(args)
+        elif self.block_type == "E":
+            self.mixer = NemotronHMoE(args)
 
     def __call__(
         self,
@@ -296,7 +401,7 @@ class NemotronHModel(nn.Module):
             NemotronHBlock(args, block_type)
             for block_type in args.hybrid_override_pattern
         ]
-        self.norm_f = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.norm_f = nn.RMSNorm(args.hidden_size, eps=args.layer_norm_epsilon)
         self.fa_idx = 0
         self.ssm_idx = 0
         for b in args.hybrid_override_pattern:
@@ -372,4 +477,23 @@ class Model(nn.Module):
         for k, v in weights.items():
             if "conv1d.weight" in k and v.shape[-1] != 1:
                 weights[k] = v.moveaxis(2, 1)
+
+        # Stack experts
+        for l in range(self.args.num_hidden_layers):
+            prefix = f"backbone.layers.{l}.mixer"
+            for m, n in [("down_proj", "fc2"), ("up_proj", "fc1")]:
+                if f"{prefix}.experts.0.{m}.weight" in weights:
+                    to_join = [
+                        weights.pop(f"{prefix}.experts.{e}.{m}.weight")
+                        for e in range(self.args.n_routed_experts)
+                    ]
+                    weights[f"{prefix}.switch_mlp.{n}.weight"] = mx.stack(to_join)
+
         return weights
+
+    @property
+    def cast_predicate(self):
+        def predicate(k):
+            return "e_score_correction_bias" not in k and "A_log" not in k
+
+        return predicate

mlx_lm/models/qwen2.py

@@ -5,6 +5,7 @@ from typing import Any, Dict, Optional, Union
 
 import mlx.core as mx
 import mlx.nn as nn
+from mlx.nn.layers.distributed import shard_linear
 
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .rope_utils import initialize_rope
@@ -183,6 +184,37 @@ class Model(nn.Module):
             k: v for k, v in weights.items() if "self_attn.rotary_emb.inv_freq" not in k
         }
 
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        group = group or mx.distributed.init()
+        N = group.size()
+        for layer in self.model.layers:
+            # Shard the self attention
+            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.n_heads //= N
+            layer.self_attn.n_kv_heads //= N
+
+            # Shard the 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
+            )
+
     @property
     def layers(self):
         return self.model.layers

mlx_lm/models/qwen3.py

@@ -5,6 +5,7 @@ from typing import Any, Dict, Optional, Union
 
 import mlx.core as mx
 import mlx.nn as nn
+from mlx.nn.layers.distributed import shard_linear
 
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .rope_utils import initialize_rope
@@ -185,6 +186,37 @@ class Model(nn.Module):
             weights.pop("lm_head.weight", None)
         return weights
 
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        group = group or mx.distributed.init()
+        N = group.size()
+        for layer in self.model.layers:
+            # Shard the self attention
+            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.n_heads //= N
+            layer.self_attn.n_kv_heads //= N
+
+            # Shard the 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
+            )
+
     @property
     def layers(self):
         return self.model.layers

mlx_lm/models/rope_utils.py

@@ -43,18 +43,22 @@ class SuScaledRoPE(nn.Module):
             long_mscale (float, optional): Scale the input prior to embedding.
         """
         super().__init__()
-        freqs = base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims)
-        self._freqs = mx.array(long_factor, dtype=mx.float32) * freqs
         self.original_max_position_embeddings = original_max_position_embeddings
-        self.scale = long_mscale or math.sqrt(
-            1
-            + math.log(max_position_embeddings / original_max_position_embeddings)
-            / math.log(original_max_position_embeddings)
-        )
         self.dim = dims
 
-    def __call__(self, x, offset: int = 0):
-        x[..., : self.dim] = self.scale * x[..., : self.dim]
+        freqs = base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims)
+        self._freqs = mx.array(long_factor, dtype=mx.float32) * freqs
+
+        def default_scale(factor):
+            return math.sqrt(
+                1 + math.log(factor) / math.log(original_max_position_embeddings)
+            )
+
+        factor = max_position_embeddings / original_max_position_embeddings
+        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,
             self.dim,

mlx_lm/models/rwkv7.py

@@ -0,0 +1,453 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from functools import partial
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs
+from .cache import ArraysCache
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    vocab_size: int
+    hidden_size: int
+    intermediate_size: int
+    norm_eps: float
+    head_dim: int
+    num_hidden_layers: int
+    a_low_rank_dim: int
+    v_low_rank_dim: int
+    gate_low_rank_dim: int
+    decay_low_rank_dim: int
+    tie_word_embeddings: bool = False
+
+
+@partial(mx.compile, shapeless=True)
+def addcmul(x, y, z):
+    return x + y * z
+
+
+@partial(mx.compile, shapeless=True)
+def l2_norm(x):
+    return x / mx.maximum(mx.linalg.norm(x, axis=-1, keepdims=True), 1e-7)
+
+
+@mx.compile
+def _wkv7_step_ops(r, w, k, v, a, b, state):
+    sab = (state @ a[..., None]) @ b[..., None, :]
+    state = state * w[:, :, None, :] + v[..., None] @ k[..., None, :] + sab
+    y = state @ r[..., None]
+    return y, state
+
+
+def _make_wkv7_kernel():
+    if not mx.metal.is_available():
+        return None
+    source = f"""
+        auto n = thread_position_in_grid.z;
+        auto b_idx = n / H;
+        auto h_idx = n % H;
+        constexpr int n_per_t = D / 32;
+
+        // [B, T, H, D]
+        auto r_ = r + b_idx * T * H * D + h_idx * D;
+        auto w_ = w + b_idx * T * H * D + h_idx * D;
+        auto k_ = k + b_idx * T * H * D + h_idx * D;
+        auto v_ = v + b_idx * T * H * D + h_idx * D;
+        auto a_ = a + b_idx * T * H * D + h_idx * D;
+        auto b_ = b + b_idx * T * H * D + h_idx * D;
+        y += b_idx * T * H * D + h_idx * D;
+
+        auto dk_idx = thread_position_in_threadgroup.x;
+        auto dv_idx = thread_position_in_grid.y;
+
+        // state_in, state_out: [B, H, D, D]
+        auto i_state = state_in  + (n * D + dv_idx) * D;
+        auto o_state = state_out + (n * D + dv_idx) * D;
+
+        float state[n_per_t];
+        for (int i = 0; i < n_per_t; ++i) {{
+          auto s_idx = n_per_t * dk_idx + i;
+          state[i] = static_cast<float>(i_state[s_idx]);
+        }}
+
+        for (int t = 0; t < T; ++t) {{
+          float sa = 0.0f;
+          for (int i = 0; i < n_per_t; ++i) {{
+            auto s_idx = n_per_t * dk_idx + i;
+            sa += state[i] * a_[s_idx];
+            state[i] = state[i] * w_[s_idx];
+          }}
+          sa = simd_sum(sa);
+
+          float out = 0.0f;
+          for (int i = 0; i < n_per_t; ++i) {{
+            auto s_idx = n_per_t * dk_idx + i;
+            state[i] = state[i] + k_[s_idx] * v_[dv_idx] + sa * b_[s_idx];
+            out += state[i] * r_[s_idx];
+          }}
+          out = simd_sum(out);
+          if (thread_index_in_simdgroup == 0) {{
+            y[dv_idx] = static_cast<InT>(out);
+          }}
+
+          // Increment data pointers to next time step
+          r_ += H * D;
+          w_ += H * D;
+          k_ += H * D;
+          v_ += H * D;
+          a_ += H * D;
+          b_ += H * D;
+          y  += H * D;
+        }}
+        for (int i = 0; i < n_per_t; ++i) {{
+          auto s_idx = n_per_t * dk_idx + i;
+          o_state[s_idx] = static_cast<InT>(state[i]);
+        }}
+    """
+    inputs = ["r", "w", "k", "v", "a", "b", "state_in", "T"]
+    return mx.fast.metal_kernel(
+        name="wkv7_kernel",
+        input_names=inputs,
+        output_names=["y", "state_out"],
+        source=source,
+    )
+
+
+_wkv7_kernel = _make_wkv7_kernel()
+
+
+def wkv7_kernel(
+    r: mx.array,
+    w: mx.array,
+    k: mx.array,
+    v: mx.array,
+    a: mx.array,
+    b: mx.array,
+    state: mx.array,
+):
+    B, T, H, D = r.shape
+    input_dtype = r.dtype
+
+    return _wkv7_kernel(
+        inputs=[r, w, k, v, a, b, state, T],
+        template=[
+            ("InT", input_dtype),
+            ("H", H),
+            ("D", D),
+        ],
+        grid=(32, D, B * H),
+        threadgroup=(32, 4, 1),
+        output_shapes=[(B, T, H, D), state.shape],
+        output_dtypes=[input_dtype, input_dtype],
+    )
+
+
+class LayerNormPerHead(nn.Module):
+    def __init__(self, head_dim, num_heads, eps):
+        super().__init__()
+        self.weight = mx.zeros((num_heads, head_dim))
+        self.bias = mx.zeros((num_heads, head_dim))
+        self.eps = eps
+
+    def __call__(self, x):
+        return self.weight * mx.fast.layer_norm(x, None, None, self.eps) + self.bias
+
+
+class LoRA(nn.Module):
+    def __init__(
+        self,
+        input_dim: int,
+        output_dim: int,
+        low_rank_dim: int,
+        bias: Optional[bool] = True,
+        activation: Optional[str] = "tanh",
+    ):
+        super().__init__()
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.low_rank_dim = low_rank_dim
+        self.bias = bias
+
+        if activation is None:
+            self.activation = nn.Identity()
+        elif activation == "sigmoid":
+            self.activation = nn.Sigmoid()
+        elif activation == "tanh":
+            self.activation = nn.Tanh()
+        elif activation == "relu":
+            self.activation = nn.ReLU()
+        else:
+            raise ValueError(f"Unsupported activation type: {activation}.")
+
+        self.lora = [
+            nn.Linear(self.input_dim, self.low_rank_dim, bias=False),
+            self.activation,
+            nn.Linear(self.low_rank_dim, self.output_dim, bias=self.bias),
+        ]
+
+    def __call__(self, x) -> mx.array:
+        return self.lora[2](self.lora[1](self.lora[0](x)))
+
+
+class TokenShift(nn.Module):
+    def __call__(self, x, state):
+        B, L, D = x.shape
+        if state is None:
+            state = mx.zeros((B, 1, D), x.dtype)
+        if L == 1:
+            return state
+        else:
+            return mx.concatenate([state, x[:, :-1, :]], axis=1)
+
+
+class Rwkv7ChannelMixing(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        hidden_dim = args.hidden_size
+        intermediate_size = args.intermediate_size
+
+        self.key = nn.Linear(hidden_dim, intermediate_size, bias=False)
+        self.value = nn.Linear(intermediate_size, hidden_dim, bias=False)
+
+        self.x_k = mx.zeros((hidden_dim))
+
+        self.token_shift = TokenShift()
+
+    def __call__(self, x, cache) -> mx.array:
+        state = cache[2] if cache is not None else None
+        x_prev = self.token_shift(x, state)
+        xx = addcmul(x, x_prev - x, self.x_k)
+        if cache is not None:
+            cache[2] = x[:, -1:, :]
+        return self.value(nn.relu2(self.key(xx)))
+
+
+class Rwkv7TimeMixing(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.args = args
+        self.hidden_size = args.hidden_size
+        self.head_dim = args.head_dim
+        self.num_heads = self.hidden_size // self.head_dim
+        self.a_low_rank_dim = args.a_low_rank_dim
+        self.v_low_rank_dim = args.v_low_rank_dim
+        self.gate_low_rank_dim = args.gate_low_rank_dim
+        self.decay_low_rank_dim = args.decay_low_rank_dim
+
+        self.token_shift = TokenShift()
+
+        self.x_r = mx.zeros((1, 1, self.hidden_size))
+        self.x_w = mx.zeros((1, 1, self.hidden_size))
+        self.x_k = mx.zeros((1, 1, self.hidden_size))
+        self.x_v = mx.zeros((1, 1, self.hidden_size))
+        self.x_a = mx.zeros((1, 1, self.hidden_size))
+        self.x_g = mx.zeros((1, 1, self.hidden_size))
+
+        self.k_k = mx.zeros((self.num_heads, self.head_dim))
+        self.k_a = mx.zeros((self.num_heads, self.head_dim))
+        self.r_k = mx.zeros((self.num_heads, self.head_dim))
+
+        self.r_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+        self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+
+        self.g_norm = LayerNormPerHead(self.head_dim, self.num_heads, eps=64e-5)
+
+        self.w_lora = LoRA(
+            self.hidden_size,
+            self.hidden_size,
+            low_rank_dim=self.decay_low_rank_dim,
+            activation="tanh",
+        )
+
+        if self.layer_idx > 0:
+            self.v_lora = LoRA(
+                self.hidden_size,
+                self.hidden_size,
+                low_rank_dim=self.v_low_rank_dim,
+                activation=None,
+            )
+
+        self.a_lora = LoRA(
+            self.hidden_size,
+            self.hidden_size,
+            low_rank_dim=self.a_low_rank_dim,
+            activation=None,
+        )
+
+        self.g_lora = LoRA(
+            self.hidden_size,
+            self.hidden_size,
+            low_rank_dim=self.gate_low_rank_dim,
+            activation="sigmoid",
+            bias=False,
+        )
+
+    def _wkv7(self, r, w, k, v, a, b, state):
+        B, L, _, _ = r.shape
+        if state is None:
+            state = mx.zeros(
+                (B, self.num_heads, self.head_dim, self.head_dim), dtype=r.dtype
+            )
+
+        if mx.default_device() == mx.gpu and mx.metal.is_available():
+            return wkv7_kernel(r, w, k, v, a, b, state)
+        else:
+            ys = []
+            for t in range(L):
+                y, state = _wkv7_step_ops(
+                    r[:, t], w[:, t], k[:, t], v[:, t], a[:, t], b[:, t], state
+                )
+                ys.append(y)
+
+            y = mx.stack(ys, axis=1).astype(r.dtype)
+            return y, state
+
+    def __call__(self, x, v_first, cache):
+        if cache is None:
+            token_shift_cache, state_cache = None, None
+        else:
+            token_shift_cache, state_cache = cache[0], cache[1]
+
+        B, L, D = x.shape
+        x_prev = self.token_shift(x, token_shift_cache)
+        xx = x_prev - x
+
+        xr = addcmul(x, xx, self.x_r)
+        xw = addcmul(x, xx, self.x_w)
+        xk = addcmul(x, xx, self.x_k)
+        xv = addcmul(x, xx, self.x_v)
+        xa = addcmul(x, xx, self.x_a)
+        xg = addcmul(x, xx, self.x_g)
+
+        key = self.k_proj(xk).reshape(B, L, self.num_heads, self.head_dim)
+        value = self.v_proj(xv).reshape(B, L, self.num_heads, self.head_dim)
+        receptance = self.r_proj(xr).reshape(B, L, self.num_heads, self.head_dim)
+        iclr = mx.sigmoid(self.a_lora(xa)).reshape(B, L, self.num_heads, self.head_dim)
+        gate = self.g_lora(xg)
+
+        if self.layer_idx == 0:
+            v_first = value
+        else:
+            vv = mx.sigmoid(self.v_lora(xv)).reshape(
+                B, L, self.num_heads, self.head_dim
+            )
+            value = addcmul(value, v_first - value, vv)
+
+        decay = mx.sigmoid(
+            self.w_lora(xw).reshape(B, L, self.num_heads, self.head_dim)
+        ).astype(mx.float32)
+        decay = mx.exp(-0.606531 * decay).astype(receptance.dtype)
+        kk = l2_norm((key * self.k_k))
+        key = key * (1 + (iclr - 1) * self.k_a)
+        a = -kk
+        b = kk * iclr
+
+        out, new_state_cache = self._wkv7(
+            receptance, decay, key, value, a, b, state_cache
+        )
+        out = self.g_norm(out.reshape(B, L, self.num_heads, self.head_dim))
+        out = (
+            out + (receptance * key * self.r_k).sum(axis=-1, keepdims=True) * value
+        ).reshape([B, L, D])
+
+        if cache is not None:
+            cache[0] = x[:, -1:, :]
+            cache[1] = new_state_cache
+
+        out = self.o_proj(out * gate)
+        return out, v_first
+
+
+class Rwkv7Layer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.layer_idx = layer_idx
+        if self.layer_idx == 0:
+            self.pre_norm = nn.LayerNorm(args.hidden_size, eps=args.norm_eps)
+        self.attn = Rwkv7TimeMixing(args, layer_idx=self.layer_idx)
+        self.ffn = Rwkv7ChannelMixing(args)
+        self.attn_norm = nn.LayerNorm(args.hidden_size, eps=args.norm_eps)
+        self.ffn_norm = nn.LayerNorm(args.hidden_size, eps=args.norm_eps)
+
+    def __call__(self, x, v_first, cache):
+        if self.layer_idx == 0:
+            x = self.pre_norm(x)
+
+        h, v_first = self.attn(self.attn_norm(x), v_first, cache)
+        h = x + h
+        out = h + self.ffn(self.ffn_norm(h), cache)
+        return out, v_first
+
+
+class Rwkv7Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.embeddings = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            Rwkv7Layer(args, layer_idx=i) for i in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.LayerNorm(args.hidden_size, eps=args.norm_eps)
+
+    def __call__(self, x: mx.array, cache):
+        x = self.embeddings(x)
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        v_first = None
+        for layer, c in zip(self.layers, cache):
+            x, v_first = layer(x, v_first, c)
+        return self.norm(x)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = Rwkv7Model(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):
+        x = self.model(inputs, cache)
+
+        if self.args.tie_word_embeddings:
+            logits = self.model.embeddings.as_linear(x)
+        else:
+            logits = self.lm_head(x)
+
+        return logits
+
+    def make_cache(self):
+        return [ArraysCache(size=3) for _ in range(len(self.layers))]
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    def sanitize(self, weights):
+        for k, v in weights.items():
+            if "k_k" in k or "k_a" in k or "g_norm" in k:
+                weights[k] = weights[k].reshape(
+                    self.args.hidden_size // self.args.head_dim, self.args.head_dim
+                )
+        return weights
+
+    @property
+    def quant_predicate(self):
+        def predicate(path, _):
+            if "lora.2" in path or "embeddings" in path:
+                return {"bits": 8}
+            return True
+
+        return predicate

mlx_lm/models/solar_open.py

@@ -0,0 +1,38 @@
+# Copyright © 2026 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Dict, Optional
+
+from .base import BaseModelArgs
+from .glm4_moe import Model  # noqa: F401
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    vocab_size: int
+    hidden_size: int
+    intermediate_size: int
+    moe_intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    head_dim: int
+    n_shared_experts: int
+    n_routed_experts: int
+    routed_scaling_factor: float
+    num_experts_per_tok: int
+    first_k_dense_replace: int
+    norm_topk_prob: bool
+    max_position_embeddings: int
+    rms_norm_eps: float
+    rope_theta: float
+    tie_word_embeddings: bool
+    partial_rotary_factor: float
+    rope_scaling: Optional[Dict] = None
+    attention_bias: bool = False
+    use_qk_norm: bool = False
+    n_group: int = 1
+    topk_group: int = 1
+    scoring_func: str = "sigmoid"
+    topk_method: str = "noaux_tc"

mlx_lm/models/ssm.py

@@ -139,7 +139,7 @@ def ssm_attn(
 
     dt = compute_dt(dt, dt_bias, time_step_limit)
     repeats = h // g
-    A = -mx.exp(A_log)
+    A = -mx.exp(A_log).astype(dt.dtype)
     dtA = dt * A.reshape(1, 1, -1)
     dtx = dt.reshape(b, l, h, 1) * x
 

mlx_lm/models/youtu_llm.py

@@ -0,0 +1,238 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .rope_utils import initialize_rope
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str = "youtu_llm"
+    vocab_size: int = 128256
+    hidden_size: int = 2048
+    intermediate_size: int = 6144
+    num_hidden_layers: int = 32
+    num_attention_heads: int = 16
+    num_key_value_heads: int = 16
+    kv_lora_rank: int = 512
+    q_lora_rank: int = 1536
+    qk_rope_head_dim: int = 64
+    v_head_dim: int = 128
+    qk_nope_head_dim: int = 128
+    max_position_embeddings: int = 131072
+    rms_norm_eps: float = 1e-6
+    rope_theta: float = 1600000
+    rope_traditional: bool = True
+    rope_scaling: Optional[Dict] = None
+    attention_bias: bool = False
+    mlp_bias: bool = False
+    tie_word_embeddings: bool = True
+
+
+class YoutuLLMAttention(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.max_position_embeddings = config.max_position_embeddings
+        self.rope_theta = config.rope_theta
+        self.q_lora_rank = config.q_lora_rank
+        self.qk_rope_head_dim = config.qk_rope_head_dim
+        self.kv_lora_rank = config.kv_lora_rank
+        self.v_head_dim = config.v_head_dim
+        self.qk_nope_head_dim = config.qk_nope_head_dim
+        self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim
+
+        self.scale = self.q_head_dim**-0.5
+
+        if self.q_lora_rank is None:
+            self.q_proj = nn.Linear(
+                self.hidden_size, self.num_heads * self.q_head_dim, bias=False
+            )
+        else:
+            self.q_a_proj = nn.Linear(
+                self.hidden_size, self.q_lora_rank, bias=config.attention_bias
+            )
+            self.q_a_layernorm = nn.RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)
+            self.q_b_proj = nn.Linear(
+                self.q_lora_rank, self.num_heads * self.q_head_dim, bias=False
+            )
+
+        self.kv_a_proj_with_mqa = nn.Linear(
+            self.hidden_size,
+            self.kv_lora_rank + self.qk_rope_head_dim,
+            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,
+        )
+
+        self.o_proj = nn.Linear(
+            self.num_heads * self.v_head_dim,
+            self.hidden_size,
+            bias=config.attention_bias,
+        )
+
+        self.rope = initialize_rope(
+            self.qk_rope_head_dim,
+            base=self.rope_theta,
+            traditional=config.rope_traditional,
+            scaling_config=config.rope_scaling,
+            max_position_embeddings=config.max_position_embeddings,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        if self.q_lora_rank is None:
+            q = self.q_proj(x)
+        else:
+            q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(x)))
+
+        q = q.reshape(B, L, self.num_heads, self.q_head_dim).transpose(0, 2, 1, 3)
+        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 = self.kv_b_proj(self.kv_a_layernorm(compressed_kv))
+        kv = kv.reshape(B, L, self.num_heads, -1).transpose(0, 2, 1, 3)
+
+        k_nope, values = mx.split(kv, [self.qk_nope_head_dim], 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)
+
+        queries = mx.concatenate([q_nope, q_pe], axis=-1)
+
+        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 YoutuLLMMLP(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.gate_proj = nn.Linear(
+            config.hidden_size, config.intermediate_size, bias=config.mlp_bias
+        )
+        self.up_proj = nn.Linear(
+            config.hidden_size, config.intermediate_size, bias=config.mlp_bias
+        )
+        self.down_proj = nn.Linear(
+            config.intermediate_size, config.hidden_size, bias=config.mlp_bias
+        )
+
+    def __call__(self, x) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class YoutuLLMDecoderLayer(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.self_attn = YoutuLLMAttention(config)
+        self.mlp = YoutuLLMMLP(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
+        )
+
+    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
+        r = self.mlp(self.post_attention_layernorm(h))
+        out = h + r
+        return out
+
+
+class YoutuLLMModel(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.vocab_size = config.vocab_size
+        self.num_hidden_layers = config.num_hidden_layers
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = [
+            YoutuLLMDecoderLayer(config=config) for _ in range(config.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+    ):
+        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, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.model_type = config.model_type
+        self.model = YoutuLLMModel(config)
+        if not config.tie_word_embeddings:
+            self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+    ):
+        out = self.model(inputs, cache)
+        if self.config.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.config.tie_word_embeddings:
+            weights.pop("lm_head.weight", None)
+        return weights
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/server.py

@@ -34,7 +34,13 @@ from huggingface_hub import scan_cache_dir
 
 from ._version import __version__
 from .generate import BatchGenerator, stream_generate
-from .models.cache import can_trim_prompt_cache, make_prompt_cache, trim_prompt_cache
+from .models.cache import (
+    KVCache,
+    RotatingKVCache,
+    can_trim_prompt_cache,
+    make_prompt_cache,
+    trim_prompt_cache,
+)
 from .sample_utils import make_logits_processors, make_sampler
 from .utils import load
 
@@ -52,9 +58,9 @@ class StopCondition(NamedTuple):
 
 def stopping_criteria(
     tokens: List[int],
+    eos_token_ids: set,
     stop_id_sequences: List[List[int]],
     stop_words: List[str],
-    eos_token_id: Union[int, None],
 ) -> StopCondition:
     """
     Determines whether the token generation should stop based on predefined
@@ -62,14 +68,14 @@ def stopping_criteria(
 
     Args:
         tokens (List[int]): The current sequence of generated tokens.
+        eos_token_ids (set): The token IDs that represents the
+          end-of-sequence. If the last token in ``tokens`` is in the set,
+          the generation should stop.
         stop_id_sequences (List[List[[int]]): A list of integer lists, each
           representing a sequence of token IDs. If the end of the `tokens`
           list matches any of these sequences, the generation should stop.
         stop_words (List[str]): The stop words that correspond to the
             ``stop_id_sequences``.
-        eos_token_id (Union[int, None]): The token ID that represents the
-          end-of-sequence. If the last token in `tokens` matches this, the
-          generation should stop.
 
     Returns:
         StopCondition: A named tuple indicating whether the stop condition has
@@ -77,7 +83,7 @@ def stopping_criteria(
           end if it has (`trim_length`) as well as the text that should be
           trimmed.
     """
-    if tokens and tokens[-1] == eos_token_id:
+    if tokens and tokens[-1] in eos_token_ids:
         return StopCondition(stop_met=True, trim_length=0, trim_text_length=0)
 
     for stop_ids, stop_word in zip(stop_id_sequences, stop_words):
@@ -158,6 +164,11 @@ def process_message_content(messages):
             message["content"] = "".join(text_fragments)
         elif content is None:
             message["content"] = ""
+        if tool_calls := message.get("tool_calls", False):
+            for tool_call in tool_calls:
+                if func := tool_call.get("function", False):
+                    if args := func.get("arguments", False):
+                        func["arguments"] = json.loads(args)
 
 
 class LRUPromptCache:
@@ -351,7 +362,12 @@ class GenerationContext:
     has_tool_calling: bool
     tool_call_start: str
     tool_call_end: str
-    eos_token_id: int
+    tool_parser: Callable[[str, Any], Dict]
+    has_thinking: bool
+    think_start_id: int
+    think_end_id: int
+    think_end: str
+    eos_token_ids: set
     stop_token_sequences: List[List[int]]
     prompt: List[int]
 
@@ -378,6 +394,7 @@ class ModelProvider:
         self.model = None
         self.tokenizer = None
         self.draft_model = None
+        self.cache_types = set()
 
         # Preload the default model if it is provided
         self.default_model_map = {}
@@ -448,9 +465,41 @@ class ModelProvider:
         elif draft_model_path is not None and draft_model_path != "default_model":
             self.draft_model, draft_tokenizer = load(draft_model_path)
             validate_draft_tokenizer(draft_tokenizer)
+
+        # Figure out the cache types and save them in a set for anybody that
+        # wants to make a decision based on those.
+        for c in make_prompt_cache(self.model):
+            self.cache_types.add(type(c))
+        if self.draft_model is not None:
+            for c in make_prompt_cache(self.draft_model):
+                self.cache_types.add(type(c))
+
         return self.model, self.tokenizer
 
 
+def _make_sampler(args, tokenizer):
+    return make_sampler(
+        args.sampling.temperature,
+        top_p=args.sampling.top_p,
+        top_k=args.sampling.top_k,
+        min_p=args.sampling.min_p,
+        xtc_probability=args.sampling.xtc_probability,
+        xtc_threshold=args.sampling.xtc_threshold,
+        xtc_special_tokens=[
+            tokenizer.eos_token_id,
+            tokenizer.encode("\n"),
+        ],
+    )
+
+
+def _make_logits_processors(args):
+    return make_logits_processors(
+        args.logits.logit_bias,
+        args.logits.repetition_penalty,
+        args.logits.repetition_context_size,
+    )
+
+
 class ResponseGenerator:
     def __init__(self, model_provider: ModelProvider, prompt_cache: LRUPromptCache):
         self.model_provider = model_provider
@@ -471,12 +520,13 @@ class ResponseGenerator:
             tools = request.tools
             role_mapping = request.role_mapping
 
-            if tokenizer.chat_template:
+            if tokenizer.has_chat_template:
                 process_message_content(messages)
                 return tokenizer.apply_chat_template(
                     messages,
                     tools,
                     add_generation_prompt=True,
+                    tokenize=True,
                     **self.model_provider.cli_args.chat_template_args,
                 )
             else:
@@ -490,12 +540,9 @@ class ResponseGenerator:
             or self.model_provider.cli_args.draft_model is not None
         ):
             return False
-        if args.logits.logit_bias is not None:
-            return False
-        if args.logits.repetition_penalty != 0:
-            return False
-        if args.logprobs > 0:
-            return False
+        for c in self.model_provider.cache_types:
+            if c not in (KVCache, RotatingKVCache):
+                return False
         if args.seed is not None:
             return False
 
@@ -512,12 +559,15 @@ class ResponseGenerator:
 
         unprocessed_requests = []
 
-        def get_next_request():
+        def get_next_request(timeout=None):
             if unprocessed_requests:
                 return unprocessed_requests.pop()
             else:
                 try:
-                    return self.requests.get_nowait()
+                    if timeout is not None:
+                        return self.requests.get(timeout=timeout)
+                    else:
+                        return self.requests.get_nowait()
                 except QueueEmpty:
                     return None
 
@@ -529,7 +579,12 @@ class ResponseGenerator:
         while not self._stop:
             request = None
             if not drain_batch:
-                request = get_next_request()
+                timeout = (
+                    None
+                    if (batch_generator is not None and len(batch_results) > 0)
+                    else 0.1
+                )
+                request = get_next_request(timeout=timeout)
 
             # We got a request
             if request is not None:
@@ -541,7 +596,6 @@ class ResponseGenerator:
                 if (
                     batch_generator is not None
                     and current_model == args.model
-                    and current_sampling == args.sampling
                     and is_batchable
                 ):
                     prompt = self._tokenize(current_tokenizer, request)
@@ -549,7 +603,12 @@ class ResponseGenerator:
                         has_tool_calling=tokenizer.has_tool_calling,
                         tool_call_start=tokenizer.tool_call_start,
                         tool_call_end=tokenizer.tool_call_end,
-                        eos_token_id=tokenizer.eos_token_id,
+                        tool_parser=tokenizer.tool_parser,
+                        has_thinking=tokenizer.has_thinking,
+                        think_start_id=tokenizer.think_start_id,
+                        think_end=tokenizer.think_end,
+                        think_end_id=tokenizer.think_end_id,
+                        eos_token_ids=tokenizer.eos_token_ids,
                         stop_token_sequences=[
                             tokenizer.encode(stop_word, add_special_tokens=False)
                             for stop_word in args.stop_words
@@ -565,7 +624,11 @@ class ResponseGenerator:
                         cache = make_prompt_cache(self.model_provider.model)
 
                     (uid,) = batch_generator.insert(
-                        [rest], args.max_tokens, caches=[cache]
+                        [rest],
+                        args.max_tokens,
+                        caches=[cache],
+                        samplers=[_make_sampler(args, tokenizer)],
+                        logits_processors=[_make_logits_processors(args)],
                     )
                     batch_results[uid] = {
                         "ctx": ctx,
@@ -592,25 +655,12 @@ class ResponseGenerator:
                         continue
 
                     current_model = args.model
-                    current_sampling = args.sampling
                     current_tokenizer = tokenizer
                     current_model_key = self.model_provider.model_key
                     batch_results = {}
                     batch_generator = BatchGenerator(
                         model,
                         stop_tokens=tokenizer.eos_token_ids,
-                        sampler=make_sampler(
-                            args.sampling.temperature,
-                            top_p=args.sampling.top_p,
-                            top_k=args.sampling.top_k,
-                            min_p=args.sampling.min_p,
-                            xtc_probability=args.sampling.xtc_probability,
-                            xtc_threshold=args.sampling.xtc_threshold,
-                            xtc_special_tokens=[
-                                tokenizer.eos_token_id,
-                                tokenizer.encode("\n"),
-                            ],
-                        ),
                         prompt_progress_callback=progress_callback,
                     )
                     unprocessed_requests.append((rqueue, request, args))
@@ -650,7 +700,8 @@ class ResponseGenerator:
                     for r in responses:
                         result = batch_results[r.uid]
                         result["cache_key"].append(r.token)
-                        result["detokenizer"].add_token(r.token)
+                        if r.finish_reason != "stop":
+                            result["detokenizer"].add_token(r.token)
 
                         top_tokens = None
                         if args.logprobs > 0:
@@ -708,7 +759,12 @@ class ResponseGenerator:
                 has_tool_calling=tokenizer.has_tool_calling,
                 tool_call_start=tokenizer.tool_call_start,
                 tool_call_end=tokenizer.tool_call_end,
-                eos_token_id=tokenizer.eos_token_id,
+                tool_parser=tokenizer.tool_parser,
+                has_thinking=tokenizer.has_thinking,
+                think_start_id=tokenizer.think_start_id,
+                think_end=tokenizer.think_end,
+                think_end_id=tokenizer.think_end_id,
+                eos_token_ids=tokenizer.eos_token_ids,
                 stop_token_sequences=[
                     tokenizer.encode(stop_word, add_special_tokens=False)
                     for stop_word in args.stop_words
@@ -722,23 +778,8 @@ class ResponseGenerator:
                 mx.random.seed(args.seed)
 
             # Make the sampler and logit processor
-            sampler = make_sampler(
-                args.sampling.temperature,
-                top_p=args.sampling.top_p,
-                top_k=args.sampling.top_k,
-                min_p=args.sampling.min_p,
-                xtc_probability=args.sampling.xtc_probability,
-                xtc_threshold=args.sampling.xtc_threshold,
-                xtc_special_tokens=[
-                    tokenizer.eos_token_id,
-                    tokenizer.encode("\n"),
-                ],
-            )
-            logits_processors = make_logits_processors(
-                args.logits.logit_bias,
-                args.logits.repetition_penalty,
-                args.logits.repetition_context_size,
-            )
+            sampler = _make_sampler(args, tokenizer)
+            logits_processors = _make_logits_processors(args)
 
             # Load the KV cache
             cache, rest = self.prompt_cache.fetch_nearest_cache(
@@ -921,7 +962,7 @@ class APIHandler(BaseHTTPRequestHandler):
         self.top_p = self.body.get("top_p", self.response_generator.cli_args.top_p)
         self.top_k = self.body.get("top_k", self.response_generator.cli_args.top_k)
         self.min_p = self.body.get("min_p", self.response_generator.cli_args.min_p)
-        self.repetition_penalty = self.body.get("repetition_penalty", 1.0)
+        self.repetition_penalty = self.body.get("repetition_penalty", 0.0)
         self.repetition_context_size = self.body.get("repetition_context_size", 20)
         self.xtc_probability = self.body.get("xtc_probability", 0.0)
         self.xtc_threshold = self.body.get("xtc_threshold", 0.0)
@@ -1015,6 +1056,7 @@ class APIHandler(BaseHTTPRequestHandler):
         top_tokens: Optional[List[Dict[int, float]]] = None,
         tokens: Optional[List[int]] = None,
         tool_calls: Optional[List[str]] = None,
+        reasoning_text: Optional[str] = None,
     ) -> dict:
         """
         Generate a single response packet based on response type (stream or
@@ -1034,6 +1076,7 @@ class APIHandler(BaseHTTPRequestHandler):
               tokens to logprobs for the top N tokens at each token position.
             tokens (Optional[List[int]]): List of tokens to return with logprobs structure
             tool_calls (Optional[List[str]]): List of tool calls.
+            reasoning_text (Optional[str]): The reasoning text generated by the model.
 
         Returns:
             dict: A dictionary containing the response, in the same format as
@@ -1043,17 +1086,6 @@ class APIHandler(BaseHTTPRequestHandler):
         top_logprobs = top_tokens or []
         tool_calls = tool_calls or []
 
-        def parse_function(tool_text):
-            tool_call = json.loads(tool_text.strip())
-            return {
-                "function": {
-                    "name": tool_call.get("name", None),
-                    "arguments": json.dumps(tool_call.get("arguments", "")),
-                },
-                "type": "function",
-                "id": None,
-            }
-
         # Static response
         response = {
             "id": self.request_id,
@@ -1099,7 +1131,8 @@ class APIHandler(BaseHTTPRequestHandler):
             choice[key_name] = {
                 "role": "assistant",
                 "content": text,
-                "tool_calls": [parse_function(tool_text) for tool_text in tool_calls],
+                "reasoning": reasoning_text,
+                "tool_calls": tool_calls,
             }
         elif self.object_type == "text_completion":
             choice.update(text=text)
@@ -1184,8 +1217,43 @@ class APIHandler(BaseHTTPRequestHandler):
         # Variables to save the tool calls in as they are being generated by
         # the model.
         in_tool_call = False
+        made_tool_call = False
         tool_calls = []
         tool_text = ""
+        tool_idx = 0
+
+        def parse_single_tool(tool_text):
+            nonlocal tool_idx
+            tool_call = ctx.tool_parser(tool_text, request.tools)
+            tool_call["arguments"] = json.dumps(
+                tool_call["arguments"], ensure_ascii=False
+            )
+            out = {
+                "function": tool_call,
+                "type": "function",
+                "id": str(uuid.uuid4()),
+            }
+            if self.stream:
+                out["index"] = tool_idx
+                tool_idx += 1
+            return out
+
+        def parse_tools(tool_calls):
+            if not tool_calls:
+                return []
+            return [parse_single_tool(tool_text) for tool_text in tool_calls]
+
+        # Start out in reasoning if the model is a reasoning model and the
+        # prompt has an open think token but no closing think token
+        in_reasoning = False
+        if ctx.has_thinking:
+            for i in range(len(ctx.prompt) - 1, -1, -1):
+                if ctx.prompt[i] == ctx.think_end_id:
+                    break
+                elif ctx.prompt[i] == ctx.think_start_id:
+                    in_reasoning = True
+                    break
+        reasoning_text = ""
 
         # Variables to save the generated tokens and the corresponding probs
         tokens = []
@@ -1198,13 +1266,18 @@ class APIHandler(BaseHTTPRequestHandler):
 
         # Well finally save the reason for stopping
         finish_reason = "length"
-
         # Process the generated tokens
         for gen in response:
             logging.debug(gen.text)
 
             # Gather the text in tool calling or text variables
-            if ctx.has_tool_calling and gen.text == ctx.tool_call_start:
+            if in_reasoning:
+                if gen.text == ctx.think_end:
+                    in_reasoning = False
+                else:
+                    reasoning_text += gen.text
+            elif ctx.has_tool_calling and gen.text == ctx.tool_call_start:
+                made_tool_call = True
                 in_tool_call = True
             elif in_tool_call:
                 if gen.text == ctx.tool_call_end:
@@ -1227,10 +1300,13 @@ class APIHandler(BaseHTTPRequestHandler):
 
             # Check if we should stop early
             stop_condition = stopping_criteria(
-                tokens, ctx.stop_token_sequences, stop_words, ctx.eos_token_id
+                tokens,
+                ctx.eos_token_ids,
+                ctx.stop_token_sequences,
+                stop_words,
             )
             if stop_condition.stop_met:
-                finish_reason = "stop"
+                finish_reason = "tool_call" if made_tool_call else "stop"
                 ctx.stop()
                 tokens = tokens[: len(tokens) - stop_condition.trim_length]
                 text = text[: len(text) - stop_condition.trim_text_length]
@@ -1247,12 +1323,16 @@ class APIHandler(BaseHTTPRequestHandler):
                     )
                 ):
                     continue
-                elif segment or tool_calls:
+                elif segment or tool_calls or reasoning_text:
                     response = self.generate_response(
-                        segment, None, tool_calls=tool_calls
+                        segment,
+                        None,
+                        tool_calls=parse_tools(tool_calls),
+                        reasoning_text=reasoning_text,
                     )
                     self.wfile.write(f"data: {json.dumps(response)}\n\n".encode())
                     self.wfile.flush()
+                    reasoning_text = ""
                     segment = ""
                     tool_calls = []
 
@@ -1261,7 +1341,10 @@ class APIHandler(BaseHTTPRequestHandler):
 
         if self.stream:
             response = self.generate_response(
-                segment, finish_reason, tool_calls=tool_calls
+                segment,
+                finish_reason,
+                tool_calls=parse_tools(tool_calls),
+                reasoning_text=reasoning_text,
             )
             self.wfile.write(f"data: {json.dumps(response)}\n\n".encode())
             self.wfile.flush()
@@ -1280,7 +1363,8 @@ class APIHandler(BaseHTTPRequestHandler):
                 token_logprobs=token_logprobs,
                 top_tokens=top_tokens,
                 tokens=tokens,
-                tool_calls=tool_calls,
+                reasoning_text=reasoning_text,
+                tool_calls=parse_tools(tool_calls),
             )
             response_json = json.dumps(response).encode()
             indent = "\t"  # Backslashes can't be inside of f-strings
@@ -1416,6 +1500,18 @@ class APIHandler(BaseHTTPRequestHandler):
             for repo in downloaded_models
         ]
 
+        if self.response_generator.cli_args.model:
+            model_path = Path(self.response_generator.cli_args.model)
+            if model_path.exists():
+                model_id = str(model_path.resolve())
+                models.append(
+                    {
+                        "id": model_id,
+                        "object": "model",
+                        "created": self.created,
+                    }
+                )
+
         response = {"object": "list", "data": models}
 
         response_json = json.dumps(response).encode()
@@ -1550,6 +1646,9 @@ def main():
         default="{}",
     )
     args = parser.parse_args()
+    if mx.metal.is_available():
+        wired_limit = mx.metal.device_info()["max_recommended_working_set_size"]
+        mx.set_wired_limit(wired_limit)
 
     logging.basicConfig(
         level=getattr(logging, args.log_level.upper(), None),

mlx_lm/tokenizer_utils.py

@@ -1,3 +1,4 @@
+import importlib
 import json
 from functools import partial
 from json import JSONDecodeError
@@ -5,6 +6,8 @@ from typing import Any, Dict, List, Optional
 
 from transformers import AutoTokenizer, PreTrainedTokenizerFast
 
+from .tool_parsers.json_tools import parse_tool_call as default_tool_parser
+
 
 class StreamingDetokenizer:
     """The streaming detokenizer interface so that we can detokenize one token at a time.
@@ -89,11 +92,7 @@ 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") or (
-                self._tokenizer.clean_up_tokenization_spaces
-                and len(self._current_text) > 0
-                and self._current_text[-1] == " "
-            ):
+            if self._current_text.endswith("\ufffd"):
                 self._current_text = self._current_text[:-1]
         if self._current_text and self._current_text[-1] == "\n":
             self._text += self._current_text
@@ -161,8 +160,6 @@ 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():
@@ -197,8 +194,6 @@ 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):
@@ -208,10 +203,7 @@ class BPEStreamingDetokenizer(StreamingDetokenizer):
         text = self._decode_bytes(self._unflushed)
 
         # For multi-byte utf-8 wait until they are complete
-        # 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
-        ):
+        if not text.endswith("\ufffd"):
             self.text += self._maybe_trim_space(text)
             self._unflushed = ""
 
@@ -259,7 +251,14 @@ class TokenizerWrapper:
     """
 
     def __init__(
-        self, tokenizer, detokenizer_class=NaiveStreamingDetokenizer, eos_token_ids=None
+        self,
+        tokenizer,
+        detokenizer_class=NaiveStreamingDetokenizer,
+        eos_token_ids=None,
+        chat_template=None,
+        tool_call_start=None,
+        tool_call_end=None,
+        tool_parser=None,
     ):
         self._tokenizer = tokenizer
         self._detokenizer_class = detokenizer_class
@@ -270,24 +269,42 @@ class TokenizerWrapper:
         )
         self._think_start = None
         self._think_end = None
-        self._tool_call_start = None
-        self._tool_call_end = None
+        self._think_start_id = None
+        self._think_end_id = None
 
-        THINK_TOKENS = [("<think>", "</think>")]
-        TOOL_CALL_TOKENS = [("<tool_call>", "</tool_call>")]
+        self._chat_template = chat_template
+        self.has_chat_template = (
+            tokenizer.chat_template is not None or chat_template is not None
+        )
+        self._tool_parser = tool_parser or default_tool_parser
+        self._tool_call_start = tool_call_start
+        self._tool_call_end = tool_call_end
 
         vocab = tokenizer.get_vocab()
+        THINK_TOKENS = [("<think>", "</think>")]
         for think_start, think_end in THINK_TOKENS:
             if think_start in vocab and think_end in vocab:
                 self._think_start = think_start
                 self._think_end = think_end
+                self._think_start_id = vocab[think_start]
+                self._think_end_id = vocab[think_end]
                 break
-        if tokenizer.chat_template and '"tool"' in tokenizer.chat_template:
-            for tool_call_start, tool_call_end in TOOL_CALL_TOKENS:
-                if tool_call_start in vocab and tool_call_end in vocab:
-                    self._tool_call_start = tool_call_start
-                    self._tool_call_end = tool_call_end
-                    break
+
+        # Fallback to defaults if no tool call tokens are provided
+        if tool_call_start and tool_call_start not in vocab:
+            raise ValueError("Tool call start token not in vocab")
+        if tool_call_end and tool_call_end not in vocab:
+            raise ValueError("Tool call end token not in vocab")
+
+    def apply_chat_template(self, *args, tokenize=True, **kwargs):
+        if self._chat_template is not None:
+            out = self._chat_template(*args, **kwargs)
+            if tokenize:
+                out = self._tokenizer.encode(out, add_special_tokens=False)
+            return out
+
+        kwargs["return_dict"] = False
+        return self._tokenizer.apply_chat_template(*args, tokenize=tokenize, **kwargs)
 
     def add_eos_token(self, token: str):
         token_id = None
@@ -309,10 +326,18 @@ class TokenizerWrapper:
     def think_start(self):
         return self._think_start
 
+    @property
+    def think_start_id(self):
+        return self._think_start_id
+
     @property
     def think_end(self):
         return self._think_end
 
+    @property
+    def think_end_id(self):
+        return self._think_end_id
+
     @property
     def has_tool_calling(self):
         return self._tool_call_start is not None
@@ -325,6 +350,10 @@ class TokenizerWrapper:
     def tool_call_end(self):
         return self._tool_call_end
 
+    @property
+    def tool_parser(self):
+        return self._tool_parser
+
     @property
     def detokenizer(self):
         """
@@ -423,10 +452,26 @@ def _is_bpe_decoder(decoder):
     return isinstance(decoder, dict) and decoder.get("type", None) == "ByteLevel"
 
 
+def _infer_tool_parser(chat_template):
+    """Attempt to auto-infer a tool parser from the chat template."""
+    if not isinstance(chat_template, str):
+        return None
+    elif "<minimax:tool_call>" in chat_template:
+        return "minimax_m2"
+    elif "<start_function_call>" in chat_template:
+        return "function_gemma"
+    elif "<arg_key>" in chat_template:
+        return "glm47"
+    elif "<tool_call>\n<function=" in chat_template:
+        return "qwen3_coder"
+    elif "<tool_call>" in chat_template and "tool_call.name" in chat_template:
+        return "json_tools"
+    return None
+
+
 def load(
     model_path,
     tokenizer_config_extra: Optional[Dict[str, Any]] = None,
-    return_tokenizer=True,
     eos_token_ids=None,
 ) -> TokenizerWrapper:
     """Load a huggingface tokenizer and try to infer the type of streaming
@@ -457,15 +502,44 @@ def load(
     if isinstance(eos_token_ids, int):
         eos_token_ids = [eos_token_ids]
 
-    if return_tokenizer:
-        kwargs = tokenizer_config_extra or {}
-        return TokenizerWrapper(
-            AutoTokenizer.from_pretrained(model_path, **kwargs),
-            detokenizer_class,
-            eos_token_ids=eos_token_ids,
-        )
+    tokenizer_config_file = model_path / "tokenizer_config.json"
+    chat_template = None
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_path, **(tokenizer_config_extra or {})
+    )
+
+    tokenizer_config = tokenizer.init_kwargs
+
+    if chat_template_type := tokenizer_config.get("chat_template_type", False):
+        chat_template = importlib.import_module(
+            f"mlx_lm.chat_templates.{chat_template_type}"
+        ).apply_chat_template
+
+    tool_parser_type = tokenizer_config.get(
+        "tool_parser_type", _infer_tool_parser(tokenizer.chat_template)
+    )
+
+    if tool_parser_type is not None:
+        tool_module = importlib.import_module(f"mlx_lm.tool_parsers.{tool_parser_type}")
+        tool_parser = tool_module.parse_tool_call
+        tool_call_start = tool_module.tool_call_start
+        tool_call_end = tool_module.tool_call_end
+        tokenizer_config["tool_parser_type"] = tool_parser_type
     else:
-        return detokenizer_class
+        tool_parser = None
+        tool_call_start = None
+        tool_call_end = None
+
+    return TokenizerWrapper(
+        tokenizer,
+        detokenizer_class,
+        eos_token_ids=eos_token_ids,
+        chat_template=chat_template,
+        tool_parser=tool_parser,
+        tool_call_start=tool_call_start,
+        tool_call_end=tool_call_end,
+    )
 
 
 def no_bos_or_eos(sequence: List, bos: int, eos: int) -> List:

mlx_lm/tool_parsers/function_gemma.py

@@ -0,0 +1,47 @@
+# Copyright © 2025 Apple Inc.
+
+import json
+from typing import Any, Optional
+
+import regex as re
+
+_tool_call_regex = re.compile(r"call:(\w+)\{(.*?)\}", re.DOTALL)
+
+
+def parse_tool_call(text: str, _: Optional[Any] = None):
+    match = _tool_call_regex.findall(text)
+    if not match:
+        raise ValueError("No function provided.")
+    func_name = match[0][0]
+    args_str = match[0][1]
+    arguments = {}
+    escape = "<escape>"
+    while args_str:
+        split = args_str.index(":")
+        key = args_str[:split]
+        args_str = args_str[split + 1 :]
+        # Parse a string
+        if args_str.startswith(escape):
+            args_str = args_str[len(escape) :]
+            split = args_str.index(escape)
+            arguments[key] = args_str[:split]
+            args_str = args_str[split + len(escape) + 1 :]
+            continue
+        if "," in args_str:
+            split = args_str.index(",")
+        else:
+            split = len(args_str)
+
+        value = args_str[:split]
+        args_str = args_str[split + 1 :]
+
+        try:
+            arguments[key] = json.loads(value)
+        except json.JSONDecodeError:
+            arguments[key] = value
+
+    return dict(name=func_name, arguments=arguments)
+
+
+tool_call_start = "<start_function_call>"
+tool_call_end = "<end_function_call>"

mlx_lm/tool_parsers/glm47.py

@@ -0,0 +1,65 @@
+# Copyright © 2025 Apple Inc.
+
+"""
+Modified from:
+https://github.com/vllm-project/vllm/blob/main/vllm/tool_parsers/glm4_moe_tool_parser.py
+"""
+
+import ast
+import json
+from typing import Any
+
+import regex as re
+
+_func_name_regex = re.compile(r"^(.*?)<arg_key>", re.DOTALL)
+_func_arg_regex = re.compile(
+    r"<arg_key>(.*?)</arg_key>(?:\\n|\s)*<arg_value>(.*?)</arg_value>",
+    re.DOTALL,
+)
+
+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:
+    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):
+    func_name = _func_name_regex.search(text).group(1)
+    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/json_tools.py

@@ -0,0 +1,11 @@
+# Copyright © 2025 Apple Inc.
+
+import json
+
+tool_call_start = "<tool_call>"
+
+tool_call_end = "</tool_call>"
+
+
+def parse_tool_call(text, tools=None):
+    return json.loads(text.strip())

mlx_lm/tool_parsers/minimax_m2.py

@@ -0,0 +1,199 @@
+import json
+from typing import Any
+
+import regex as re
+
+tool_call_start: str = "<minimax:tool_call>"
+tool_call_end: str = "</minimax:tool_call>"
+
+_invoke_complete_regex = re.compile(r"<invoke name=(.*?)</invoke>", re.DOTALL)
+
+_parameter_complete_regex = re.compile(r"<parameter name=(.*?)</parameter>", re.DOTALL)
+
+
+def _extract_name(name_str: str) -> str:
+    """Extract name from quoted string."""
+    name_str = name_str.strip()
+    if (
+        name_str.startswith('"')
+        and name_str.endswith('"')
+        or name_str.startswith("'")
+        and name_str.endswith("'")
+    ):
+        return name_str[1:-1]
+    return name_str
+
+
+def _extract_types_from_schema(schema: Any) -> list[str]:
+    """
+    Extract all possible types from a JSON schema definition.
+    Handles anyOf, oneOf, allOf, type arrays, and enum fields.
+
+    Args:
+        schema: The JSON schema definition for a parameter
+
+    Returns:
+        List of type strings (e.g., ["string", "integer", "null"])
+    """
+    if schema is None:
+        return ["string"]
+
+    if not isinstance(schema, dict):
+        return ["string"]
+
+    types: set[str] = set()
+
+    # Handle direct "type" field
+    if "type" in schema:
+        type_value = schema["type"]
+        if isinstance(type_value, str):
+            types.add(type_value)
+        elif isinstance(type_value, list):
+            for t in type_value:
+                if isinstance(t, str):
+                    types.add(t)
+
+    # Handle enum - infer types from enum values
+    if "enum" in schema and isinstance(schema["enum"], list) and schema["enum"]:
+        for value in schema["enum"]:
+            if value is None:
+                types.add("null")
+            elif isinstance(value, bool):
+                types.add("boolean")
+            elif isinstance(value, int):
+                types.add("integer")
+            elif isinstance(value, float):
+                types.add("number")
+            elif isinstance(value, str):
+                types.add("string")
+            elif isinstance(value, list):
+                types.add("array")
+            elif isinstance(value, dict):
+                types.add("object")
+
+    # Handle anyOf, oneOf, allOf - recursively extract types
+    for choice_field in ("anyOf", "oneOf", "allOf"):
+        if choice_field in schema and isinstance(schema[choice_field], list):
+            for choice in schema[choice_field]:
+                extracted = _extract_types_from_schema(choice)
+                types.update(extracted)
+
+    # If no types found, default to string
+    if not types:
+        return ["string"]
+
+    return list(types)
+
+
+def _convert_param_value_with_types(value: str, param_types: list[str]) -> Any:
+    if value.lower() == "null":
+        return None
+
+    # Normalize types
+    normalized_types = [t.lower() for t in param_types]
+
+    # Try null first if it's in the list
+    if "null" in normalized_types or value.lower() in ("null", "none", "nil"):
+        return None
+
+    # Try each type in order of preference (most specific first, string as fallback)
+    # Priority: integer > number > boolean > object > array > string
+    type_priority = [
+        "integer",
+        "int",
+        "number",
+        "float",
+        "boolean",
+        "bool",
+        "object",
+        "array",
+        "string",
+        "str",
+        "text",
+    ]
+
+    for param_type in type_priority:
+        if param_type not in normalized_types:
+            continue
+
+        if param_type in ["string", "str", "text"]:
+            return value
+        elif param_type in ["integer", "int"]:
+            try:
+                return int(value)
+            except (ValueError, TypeError):
+                continue
+        elif param_type in ["number", "float"]:
+            try:
+                val = float(value)
+                return val if val != int(val) else int(val)
+            except (ValueError, TypeError):
+                continue
+        elif param_type in ["boolean", "bool"]:
+            lower_val = value.lower().strip()
+            if lower_val in ["true", "1", "yes", "on"]:
+                return True
+            elif lower_val in ["false", "0", "no", "off"]:
+                return False
+            continue
+        elif param_type in ["object", "array"]:
+            try:
+                return json.loads(value)
+            except json.JSONDecodeError:
+                continue
+
+    # Fallback: try JSON parse, then return as string
+    try:
+        return json.loads(value)
+    except json.JSONDecodeError:
+        return value
+
+
+def _get_param_types_from_config(param_name: str, param_config: dict) -> list[str]:
+    if param_name not in param_config:
+        return ["string"]
+    param_schema = param_config[param_name]
+    return _extract_types_from_schema(param_schema)
+
+
+def parse_tool_call(text: str, tools: list | None = None):
+    invoke_match = _invoke_complete_regex.findall(text)
+    if not invoke_match:
+        raise ValueError("No tool call found")
+    invoke_text = invoke_match[0]
+
+    name_match = re.search(r"^([^>]+)", invoke_text)
+    if not name_match:
+        return None
+
+    function_name = _extract_name(name_match.group(1))
+
+    # Get parameter configuration
+    param_config = {}
+    if tools:
+        for tool in tools:
+            if func := tool.get("function", False):
+                if func["name"] != function_name:
+                    continue
+                if params := func.get("parameters", False):
+                    param_config = params.get("properties", {})
+
+    # Extract parameters
+    param_dict = {}
+    for match in _parameter_complete_regex.findall(invoke_text):
+        param_match = re.search(r"^([^>]+)>(.*)", match, re.DOTALL)
+        if param_match:
+            param_name = _extract_name(param_match.group(1))
+            param_value = param_match.group(2).strip()
+            if param_value.startswith("\n"):
+                param_value = param_value[1:]
+            if param_value.endswith("\n"):
+                param_value = param_value[:-1]
+
+            param_type = _get_param_types_from_config(param_name, param_config)
+
+            param_dict[param_name] = _convert_param_value_with_types(
+                param_value, param_type
+            )
+
+    return dict(name=function_name, arguments=param_dict)

mlx_lm/tool_parsers/qwen3_coder.py

@@ -0,0 +1,111 @@
+# Copyright © 2025 Apple Inc.
+
+"""
+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
+
+import regex as re
+
+_function_regex = re.compile(r"<function=(.*?)</function>$", re.DOTALL)
+_parameter_regex = re.compile(r"<parameter=(.*?)</parameter>", re.DOTALL)
+
+_string_types = {"string", "str", "text", "varchar", "char", "enum"}
+_bool_types = {"boolean", "bool", "binary"}
+_obj_types = {"object", "array", "arr"}
+
+
+def _get_arguments_config(func_name: str, tools: Optional[Any]) -> dict:
+    """Extract argument configuration for a function."""
+    if tools is None:
+        return {}
+    for tool in tools:
+        if not (function := tool.get("function", False)):
+            continue
+        if function["name"] == func_name:
+            if not (params := function.get("parameters", False)):
+                return {}
+            return params.get("properties", {})
+    return {}
+
+
+def _convert_param_value(param_value: str, param_name: str, param_config: dict) -> Any:
+    """Convert parameter value based on its type in the schema."""
+    if param_value.lower() == "null":
+        return None
+
+    if not (param := param_config.get(param_name, False)):
+        return param_value
+
+    if "type" in param:
+        param_type = str(param["type"]).strip().lower()
+    else:
+        param_type = "string"
+    if param_type in _string_types:
+        return param_value
+    elif (
+        param_type.startswith("int")
+        or param_type.startswith("uint")
+        or param_type.startswith("long")
+        or param_type.startswith("short")
+        or param_type.startswith("unsigned")
+    ):
+        return int(param_value)
+    elif param_type.startswith("num") or param_type.startswith("float"):
+        float_param_value = float(param_value)
+        int_param_value = int(float_param_value)
+        return (
+            float_param_value
+            if (float_param_value - int_param_value) != 0
+            else int_param_value
+        )
+    elif param_type in _bool_types:
+        return param_value.lower() == "true"
+    else:
+        if (
+            param_type in _obj_types
+            or param_type.startswith("dict")
+            or param_type.startswith("list")
+        ):
+            return json.loads(param_value)
+
+        return ast.literal_eval(param_value)
+
+
+def _parse_xml_function_call(function_call_str: str, tools: Optional[Any]):
+    end_index = function_call_str.index(">")
+    function_name = function_call_str[:end_index]
+    param_config = _get_arguments_config(function_name, tools)
+    parameters = function_call_str[end_index + 1 :]
+    param_dict = {}
+    for match_text in _parameter_regex.findall(parameters):
+        idx = match_text.index(">")
+        param_name = match_text[:idx]
+        param_value = str(match_text[idx + 1 :])
+        if param_value.startswith("\n"):
+            param_value = param_value[1:]
+        if param_value.endswith("\n"):
+            param_value = param_value[:-1]
+
+        param_dict[param_name] = _convert_param_value(
+            param_value, param_name, param_config
+        )
+    return dict(name=function_name, arguments=param_dict)
+
+
+tool_call_start = "<tool_call>"
+
+tool_call_end = "</tool_call>"
+
+
+def parse_tool_call(
+    model_output: str,
+    tools: Optional[Any] = None,
+):
+    match = _function_regex.findall(model_output)
+    if not match:
+        raise ValueError("No function provided.")
+    return _parse_xml_function_call(match[0], tools)

mlx_lm/tuner/datasets.py

@@ -57,7 +57,11 @@ class ChatDataset:
     def process(self, d):
         messages = d[self.chat_key]
         tools = d.get("tools", None)
-        tokens = self.tokenizer.apply_chat_template(messages, tools=tools)
+        tokens = self.tokenizer.apply_chat_template(
+            messages,
+            tools=tools,
+            return_dict=False,
+        )
         if self.mask_prompt:
             add_generation_prompt = messages[-1].get("role") == "assistant"
             offset = len(
@@ -65,6 +69,7 @@ class ChatDataset:
                     messages[:-1],
                     tools=tools,
                     add_generation_prompt=add_generation_prompt,
+                    return_dict=False,
                 )
             )
             return (tokens, offset)
@@ -105,11 +110,16 @@ class CompletionsDataset:
             {"role": "user", "content": d[self.prompt_key]},
             {"role": "assistant", "content": d[self.completion_key]},
         ]
-        tokens = self.tokenizer.apply_chat_template(messages, tools=tools)
+        tokens = self.tokenizer.apply_chat_template(
+            messages, tools=tools, return_dict=False
+        )
         if self.mask_prompt:
             offset = len(
                 self.tokenizer.apply_chat_template(
-                    messages[0], tools=tools, add_generation_prompt=True
+                    messages[0],
+                    tools=tools,
+                    add_generation_prompt=True,
+                    return_dict=False,
                 )
             )
             return (tokens, offset)

mlx_lm/utils.py

@@ -5,7 +5,6 @@ import glob
 import importlib
 import inspect
 import json
-import logging
 import os
 import resource
 import shutil
@@ -50,6 +49,8 @@ MODEL_REMAPPING = {
     "falcon_mamba": "mamba",
     "kimi_k2": "deepseek_v3",
     "qwen2_5_vl": "qwen2_vl",
+    "minimax_m2": "minimax",
+    "iquestcoder": "llama",
 }
 
 MAX_FILE_SIZE_GB = 5
@@ -71,7 +72,6 @@ def _get_classes(config: dict):
         arch = importlib.import_module(f"mlx_lm.models.{model_type}")
     except ImportError:
         msg = f"Model type {model_type} not supported."
-        logging.error(msg)
         raise ValueError(msg)
 
     return arch.Model, arch.ModelArgs
@@ -145,12 +145,21 @@ def hf_repo_to_path(hf_repo):
 
 
 def load_config(model_path: Path) -> dict:
-    try:
-        with open(model_path / "config.json", "r") as f:
-            config = json.load(f)
-    except FileNotFoundError:
-        logging.error(f"Config file not found in {model_path}")
-        raise
+    with open(model_path / "config.json", "r") as f:
+        config = json.load(f)
+
+    generation_config_file = model_path / "generation_config.json"
+    if generation_config_file.exists():
+        generation_config = {}
+        try:
+            with open(generation_config_file, "r") as f:
+                generation_config = json.load(f)
+        except json.JSONDecodeError:
+            pass
+
+        if eos_token_id := generation_config.get("eos_token_id", False):
+            config["eos_token_id"] = eos_token_id
+
     return config
 
 
@@ -277,7 +286,9 @@ def load_tokenizer(model_path, tokenizer_config_extra=None, eos_token_ids=None):
         ],
     )
     return _load_tokenizer(
-        model_path, tokenizer_config_extra, eos_token_ids=eos_token_ids
+        model_path,
+        tokenizer_config_extra,
+        eos_token_ids=eos_token_ids,
     )
 
 
@@ -333,7 +344,12 @@ def load(
         return model, tokenizer
 
 
-def pipeline_load(repo, return_config=False):
+def sharded_load(
+    repo,
+    pipeline_group: Optional[mx.distributed.Group] = None,
+    tensor_group: Optional[mx.distributed.Group] = None,
+    return_config: bool = False,
+):
     # Get model path with everything but weight safetensors
     model_path = _download(
         repo,
@@ -349,27 +365,50 @@ def pipeline_load(repo, return_config=False):
         ],
     )
 
-    # Lazy load and shard model to figure out which weights we need
+    # Lazy load model to figure out what type of sharding we can do and which
+    # weights we need to download.
     model, config = load_model(model_path, lazy=True, strict=False)
 
-    group = mx.distributed.init()
-    rank = group.rank()
-    model.model.pipeline(group)
+    has_pipelining = hasattr(model.model, "pipeline")
+    has_tensor_parallel = hasattr(model, "shard")
 
-    # Figure out which files we need for the local shard
-    with open(model_path / "model.safetensors.index.json", "r") as fid:
-        weight_index = json.load(fid)["weight_map"]
+    if pipeline_group is not None and not has_pipelining:
+        raise ValueError(
+            "The model does not support pipelining but a pipeline_group was provided"
+        )
+    if tensor_group is not None and not has_tensor_parallel:
+        raise ValueError(
+            "The model does not support tensor parallelism but a tensor_group was provided"
+        )
+    if not has_pipelining and not has_tensor_parallel:
+        raise ValueError("The model does not support any sharding")
 
-    local_files = set()
-    for k, _ in tree_flatten(model.parameters()):
-        if file_name := weight_index.get(k, None) is None:
-            raise ValueError(
-                "Pipeline loading is only supported for MLX converted models."
-            )
-        local_files.add(weight_index[k])
+    if pipeline_group is tensor_group is None:
+        if has_tensor_parallel:
+            tensor_group = mx.distributed.init()
+        elif has_pipelining:
+            pipeline_group = mx.distributed.init()
 
-    # Download weights for local shard
-    _download(repo, allow_patterns=local_files)
+    # If pipelining then figure out which files we need for the local shard
+    if pipeline_group is not None:
+        model.model.pipeline(pipeline_group)
+
+        # Figure out which files we need for the local shard
+        with open(model_path / "model.safetensors.index.json", "r") as fid:
+            weight_index = json.load(fid)["weight_map"]
+
+        local_files = set()
+        for k, _ in tree_flatten(model.parameters()):
+            if file_name := weight_index.get(k, None) is None:
+                raise ValueError(
+                    "Pipeline loading is only supported for MLX converted models."
+                )
+            local_files.add(weight_index[k])
+
+        # Download weights for local shard
+        _download(repo, allow_patterns=local_files)
+    else:
+        _download(repo)
 
     # Load and shard the model, and load the weights
     tokenizer = load_tokenizer(
@@ -378,7 +417,10 @@ def pipeline_load(repo, return_config=False):
         eos_token_ids=config.get("eos_token_id", None),
     )
     model, _ = load_model(model_path, lazy=True, strict=False)
-    model.model.pipeline(group)
+    if tensor_group is not None:
+        model.shard(tensor_group)
+    if pipeline_group is not None:
+        model.model.pipeline(pipeline_group)
     mx.eval(model.parameters())
 
     # Synchronize processes to avoid timeout
@@ -389,6 +431,10 @@ def pipeline_load(repo, return_config=False):
         return model, tokenizer
 
 
+def pipeline_load(repo, return_config=False):
+    return sharded_load(repo, mx.distributed.init(), None, return_config)
+
+
 def make_shards(weights: dict, max_file_size_gb: int = MAX_FILE_SIZE_GB) -> list:
     """
     Splits the weights into smaller shards.
@@ -486,7 +532,7 @@ def upload_to_hub(path: str, upload_repo: str):
         if tokenizer.chat_template is not None:
             messages = [{{"role": "user", "content": prompt}}]
             prompt = tokenizer.apply_chat_template(
-                messages, add_generation_prompt=True
+                messages, add_generation_prompt=True, return_dict=False,
             )
 
         response = generate(model, tokenizer, prompt=prompt, verbose=True)
@@ -568,8 +614,8 @@ def save_model(
 def quantize_model(
     model: nn.Module,
     config: dict,
-    group_size: int,
-    bits: int,
+    group_size: Optional[int],
+    bits: Optional[int],
     mode: str = "affine",
     quant_predicate: Optional[Callable[[str, nn.Module], Union[bool, dict]]] = None,
 ) -> Tuple[nn.Module, dict]:
@@ -579,8 +625,8 @@ def quantize_model(
     Args:
         model (nn.Module): The model to be quantized.
         config (dict): Model configuration.
-        group_size (int): Group size for quantization.
-        bits (int): Bits per weight for quantization.
+        group_size (Optional[int]): Group size for quantization.
+        bits (Optional[int]): Bits per weight for quantization.
         mode (str): The quantization mode.
         quant_predicate (Callable): A callable that decides how to quantize
           each layer based on the path. Accepts the layer `path` and the
@@ -590,9 +636,21 @@ def quantize_model(
     Returns:
         Tuple: Tuple containing quantized model and config.
     """
+
+    def defaults_for_mode(mode, group_size, bits):
+        mode_defaults = {
+            "affine": (64, 4),
+            "mxfp4": (32, 4),
+            "nvfp4": (16, 4),
+            "mxfp8": (32, 8),
+        }
+        default_group_size, default_bits = mode_defaults[mode]
+        return group_size or default_group_size, bits or default_bits
+
     quantized_config = copy.deepcopy(config)
 
     quant_predicate = quant_predicate or getattr(model, "quant_predicate", None)
+    group_size, bits = defaults_for_mode(mode, group_size, bits)
     quant_params = {"group_size": group_size, "bits": bits, "mode": mode}
     if "quantization" in quantized_config:
         # If the model is already partially quantized, return params so that

setup.py

@@ -26,7 +26,7 @@ setup(
     install_requires=[
         f"mlx>={MIN_MLX_VERSION}; platform_system == 'Darwin'",
         "numpy",
-        "transformers>=4.39.3",
+        "transformers==5.0.0rc1",
         "sentencepiece",
         "protobuf",
         "pyyaml",

tests/test_datsets.py

@@ -21,6 +21,8 @@ class TestDatasets(unittest.TestCase):
         cls.test_dir = cls.test_dir_fid.name
         if not os.path.isdir(cls.test_dir):
             os.mkdir(cls.test_dir_fid.name)
+        # Only one HF request
+        AutoTokenizer.from_pretrained(HF_MODEL_PATH)
 
     @classmethod
     def tearDownClass(cls):
@@ -37,7 +39,7 @@ class TestDatasets(unittest.TestCase):
         data = {"text": "This is an example for the model."}
         self.save_data(4 * [data])
         args = types.SimpleNamespace(train=True, test=False, data=self.test_dir)
-        tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_PATH)
+        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)
@@ -50,7 +52,7 @@ class TestDatasets(unittest.TestCase):
         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)
-        tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_PATH)
+        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)
@@ -69,7 +71,7 @@ class TestDatasets(unittest.TestCase):
         }
         self.save_data(4 * [data])
         args = types.SimpleNamespace(train=True, test=False, data=self.test_dir)
-        tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_PATH)
+        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)
@@ -91,7 +93,7 @@ class TestDatasets(unittest.TestCase):
             test=False,
             train=True,
         )
-        tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_PATH)
+        tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_PATH, local_files_only=True)
         train, valid, test = datasets.load_dataset(args, tokenizer)
         self.assertTrue(len(train) > 0)
         self.assertTrue(len(train[0]) > 0)

tests/test_generate.py

@@ -8,6 +8,7 @@ import mlx.core as mx
 from mlx_lm.generate import (
     BatchGenerator,
     GenerationResponse,
+    batch_generate,
     generate,
     stream_generate,
 )
@@ -81,7 +82,7 @@ class TestGenerate(unittest.TestCase):
 
     def test_stream_generate_speculative(self):
         # Use same model as draft model, this is not a speed test
-        draft_model, _ = load(self.HF_MODEL_PATH)
+        draft_model = self.model
 
         results: List[GenerationResponse] = []
         drafted: List[bool] = []
@@ -90,7 +91,8 @@ class TestGenerate(unittest.TestCase):
         sampler = make_sampler(temp=0.0)
         messages = [{"role": "user", "content": "hello"}]
         prompt = self.tokenizer.apply_chat_template(
-            messages, add_generation_prompt=True
+            messages,
+            add_generation_prompt=True,
         )
 
         for generation_result in stream_generate(
@@ -117,7 +119,8 @@ class TestGenerate(unittest.TestCase):
         # get prompt embeddings
         messages = [{"role": "user", "content": "Say 'TEST' and nothing else"}]
         prompt = self.tokenizer.apply_chat_template(
-            messages, add_generation_prompt=True
+            messages,
+            add_generation_prompt=True,
         )
         prompt_embeddings = self.model.model.embed_tokens(prompt)
 
@@ -140,7 +143,8 @@ class TestGenerate(unittest.TestCase):
         # get prompt embeddings
         messages = [{"role": "user", "content": "Say 'TEST' and nothing else"}]
         prompt = self.tokenizer.apply_chat_template(
-            messages, add_generation_prompt=True
+            messages,
+            add_generation_prompt=True,
         )
         prompt_embeddings = self.model.model.embed_tokens(prompt)
 
@@ -352,6 +356,81 @@ class TestGenerate(unittest.TestCase):
 
         del self.model.make_cache
 
+    def test_batch_generate_with_logits_processors(self):
+        """Test that batch_generate with logits_processors produces correct results."""
+        logit_bias = {0: 2000.0, 1: -2000.0}
+        processors = make_logits_processors(logit_bias)
+
+        batch_gen = BatchGenerator(
+            self.model,
+            max_tokens=1,
+            logits_processors=processors,
+        )
+        prompt = self.tokenizer.encode("hello")
+        uids = batch_gen.insert([prompt])
+        response = batch_gen.next()[0]
+        logprobs = response.logprobs
+        self.assertEqual(logprobs[0].item(), 0.0)
+        self.assertEqual(logprobs.argmin().item(), 1)
+
+        del batch_gen
+
+        logit_bias = {0: 2000.0}
+        processors = make_logits_processors(logit_bias)
+        batch_gen = BatchGenerator(
+            self.model,
+            max_tokens=1,
+            logits_processors=processors,
+        )
+
+        (uid0,) = batch_gen.insert([prompt])
+
+        logit_bias = {1: 2000.0}
+        processors = make_logits_processors(logit_bias)
+        (uid1,) = batch_gen.insert([prompt], logits_processors=[processors])
+
+        logit_bias = {2: 2000.0}
+        processors = make_logits_processors(logit_bias)
+        (uid2,) = batch_gen.insert([prompt], logits_processors=[processors])
+
+        responses = batch_gen.next()
+        responses = {response.uid: response for response in responses}
+        self.assertEqual(responses[uid0].logprobs[0].item(), 0.0)
+        self.assertEqual(responses[uid1].logprobs[1].item(), 0.0)
+        self.assertEqual(responses[uid2].logprobs[2].item(), 0.0)
+
+    def test_batch_generate_with_samplers(self):
+        """Test that batch_generate with logits_processors produces correct results."""
+        batch_gen = BatchGenerator(
+            self.model,
+            max_tokens=1,
+            sampler=lambda _: mx.array([1]),
+        )
+        prompt = self.tokenizer.encode("hello")
+        uids = batch_gen.insert([prompt])
+        response = batch_gen.next()[0]
+        self.assertEqual(response.token, 1)
+
+        del batch_gen
+
+        batch_gen = BatchGenerator(
+            self.model,
+            max_tokens=1,
+            sampler=lambda _: mx.array([1]),
+        )
+
+        (uid0,) = batch_gen.insert([prompt])
+        uid1, uid2 = batch_gen.insert(
+            [prompt, prompt],
+            samplers=[lambda _: mx.array([2]), lambda _: mx.array([3])],
+        )
+
+        responses = batch_gen.next()
+        responses = {response.uid: response for response in responses}
+        self.assertEqual(responses[uid0].token, 1)
+        self.assertEqual(responses[uid1].token, 2)
+        self.assertEqual(responses[uid2].token, 3)
+
     def test_batch_continued_generation(self):
         for rotating in [False, True]:
             if rotating:

tests/test_models.py

@@ -2045,6 +2045,93 @@ class TestModels(unittest.TestCase):
                     "type": "yarn",
                 },
             },
+            {
+                "model_type": "mimo_v2_flash",
+                "num_experts_per_tok": 2,
+                "hybrid_layer_pattern": [0, 1, 0, 1],
+                "moe_layer_freq": [0, 1, 0, 1],
+                "add_swa_attention_sink_bias": True,
+                "add_full_attention_sink_bias": False,
+                "sliding_window_size": 32,
+                "vocab_size": 1000,
+                "hidden_size": 512,
+                "intermediate_size": 512,
+                "moe_intermediate_size": 128,
+                "num_hidden_layers": 4,
+                "num_attention_heads": 4,
+                "num_key_value_heads": 2,
+                "n_shared_experts": 1,
+                "n_routed_experts": 8,
+                "routed_scaling_factor": None,
+                "topk_method": "noaux_tc",
+                "scoring_func": "sigmoid",
+                "norm_topk_prob": True,
+                "n_group": 2,
+                "topk_group": 1,
+                "max_position_embeddings": 1000,
+                "layernorm_epsilon": 1e-5,
+                "rope_theta": 1000.0,
+                "swa_rope_theta": 1000.0,
+                "swa_num_attention_heads": 4,
+                "swa_num_key_value_heads": 2,
+                "head_dim": 128,
+                "v_head_dim": 64,
+                "swa_head_dim": 128,
+                "swa_v_head_dim": 64,
+                "partial_rotary_factor": 0.5,
+            },
+            {
+                "model_type": "rwkv7",
+                "vocab_size": 1000,
+                "hidden_size": 128,
+                "intermediate_size": 128,
+                "norm_eps": 1e-5,
+                "head_dim": 32,
+                "num_hidden_layers": 4,
+                "a_low_rank_dim": 16,
+                "v_low_rank_dim": 16,
+                "gate_low_rank_dim": 16,
+                "decay_low_rank_dim": 16,
+            },
+            {
+                "model_type": "exaone_moe",
+                "vocab_size": 1000,
+                "hidden_size": 128,
+                "intermediate_size": 256,
+                "moe_intermediate_size": 64,
+                "num_hidden_layers": 4,
+                "num_attention_heads": 4,
+                "num_key_value_heads": 2,
+                "head_dim": 32,
+                "num_experts": 4,
+                "num_experts_per_tok": 2,
+                "num_shared_experts": 1,
+                "n_group": 1,
+                "topk_group": 1,
+                "routed_scaling_factor": 2.5,
+                "norm_topk_prob": True,
+                "sliding_window": 32,
+                "max_position_embeddings": 1000,
+                "rms_norm_eps": 1e-5,
+                "rope_theta": 1000.0,
+                "layer_types": [
+                    "sliding_attention",
+                    "sliding_attention",
+                    "sliding_attention",
+                    "full_attention",
+                ],
+                "is_moe_layer": [False, True, True, True],
+                "tie_word_embeddings": False,
+            },
+            {
+                "model_type": "youtu_llm",
+                "vocab_size": 1000,
+                "hidden_size": 128,
+                "intermediate_size": 128,
+                "num_hidden_layers": 4,
+                "kv_lora_rank": 128,
+                "q_lora_rank": 256,
+            },
         ]
         for config in test_configs:
             model_type = config["model_type"]

tests/test_prompt_cache.py

@@ -34,6 +34,7 @@ class TestPromptCache(unittest.TestCase):
     def setUpClass(cls):
         cls.test_dir_fid = tempfile.TemporaryDirectory()
         cls.test_dir = cls.test_dir_fid.name
+        cls.model, cls.tokenizer = load(HF_MODEL_PATH)
 
     @classmethod
     def tearDownClass(cls):
@@ -132,7 +133,7 @@ class TestPromptCache(unittest.TestCase):
                 self.assertTrue(mx.array_equal(v, lv))
 
     def test_cache_with_generate(self):
-        model, tokenizer = load(HF_MODEL_PATH)
+        model, tokenizer = self.model, self.tokenizer
         prompt = tokenizer.encode("this is a prompt", return_tensors="mlx")[0]
         results = list(generate_step(prompt, model, max_tokens=4))
         toks, all_logits = zip(*results)
@@ -212,7 +213,7 @@ class TestPromptCache(unittest.TestCase):
         self.assertEqual(num_trimmed, 3)
 
     def test_trim_cache_with_generate(self):
-        model, tokenizer = load(HF_MODEL_PATH)
+        model, tokenizer = self.model, self.tokenizer
         prompt = tokenizer.encode("this is a prompt", return_tensors="mlx")[0]
 
         prompt_cache = make_prompt_cache(model)
@@ -289,7 +290,7 @@ class TestPromptCache(unittest.TestCase):
         self.assertEqual(metadata, loaded_metadata)
 
     def test_cache_to_quantized(self):
-        model, tokenizer = load(HF_MODEL_PATH)
+        model, tokenizer = self.model, self.tokenizer
         prompt = tokenizer.encode("this is a prompt", return_tensors="mlx")[0]
         results = zip(range(4), generate_step(prompt, model))
         toks, all_logits = zip(*(r[1] for r in results))

tests/test_server.py

@@ -19,6 +19,7 @@ class DummyModelProvider:
         HF_MODEL_PATH = "mlx-community/Qwen1.5-0.5B-Chat-4bit"
         self.model, self.tokenizer = load(HF_MODEL_PATH)
         self.model_key = (HF_MODEL_PATH, None)
+        self.cache_types = set([KVCache])
 
         # Add draft model support
         self.draft_model = None
@@ -39,6 +40,7 @@ class DummyModelProvider:
                 "min_p": 0.0,
                 "max_tokens": 512,
                 "chat_template_args": {},
+                "model": None,
             },
         )
 

tests/test_tool_parsing.py

@@ -0,0 +1,87 @@
+import unittest
+from pathlib import Path
+
+from mlx_lm.tool_parsers import (
+    function_gemma,
+    glm47,
+    json_tools,
+    minimax_m2,
+    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>",
+        ]
+        tools = [
+            {
+                "type": "function",
+                "function": {
+                    "name": "multiply",
+                    "description": "Multiply two numbers.",
+                    "parameters": {
+                        "type": "object",
+                        "required": ["a", "b"],
+                        "properties": {
+                            "a": {"type": "number", "description": "a is a number"},
+                            "b": {"type": "number", "description": "b is a number"},
+                        },
+                    },
+                },
+            }
+        ]
+
+        for parser, test_case in zip(parsers, test_cases):
+            with self.subTest(parser=parser):
+                tool_call = parser.parse_tool_call(test_case, tools)
+                expected = {
+                    "name": "multiply",
+                    "arguments": {"a": 12234585, "b": 48838483920},
+                }
+                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>",
+        ]
+        tools = [
+            {
+                "type": "function",
+                "function": {
+                    "name": "get_current_temperature",
+                    "description": "Get the current temperature.",
+                    "parameters": {
+                        "type": "object",
+                        "required": ["location"],
+                        "properties": {
+                            "location": {"type": "str", "description": "The location."},
+                        },
+                    },
+                },
+            }
+        ]
+
+        for parser, test_case in zip(parsers, test_cases):
+            with self.subTest(parser=parser):
+                tool_call = parser.parse_tool_call(test_case, tools)
+                expected = {
+                    "name": "get_current_temperature",
+                    "arguments": {"location": "London"},
+                }
+                self.assertEqual(tool_call, expected)
+
+
+if __name__ == "__main__":
+    unittest.main()