Add optional quantization types

* Fix decoding manually added tokens

* fix + test

* nit

* nit

* no lag bpe

---------

Co-authored-by: Awni Hannun <awni@apple.com>

flux/README.md

@@ -188,7 +188,7 @@ The adapters are saved in `mlx_output` and can be used directly by the
 
 ```shell
 python txt2image.py --model dev --save-raw --image-size 512x512 --n-images 1 \
-    --adapter mlx_output/0001200_adapters.safetensors \
+    --adapter mlx_output/final_adapters.safetensors \
     --fuse-adapter \
     --no-t5-padding \
     'A photo of an sks dog lying on the sand at a beach in Greece'

flux/dreambooth.py

@@ -13,7 +13,7 @@ from mlx.nn.utils import average_gradients
 from mlx.utils import tree_flatten, tree_map, tree_reduce
 from PIL import Image
 
-from flux import FluxPipeline, Trainer, load_dataset
+from flux import FluxPipeline, Trainer, load_dataset, save_config
 
 
 def generate_progress_images(iteration, flux, args):
@@ -43,10 +43,10 @@ def generate_progress_images(iteration, flux, args):
     im.save(out_file)
 
 
-def save_adapters(iteration, flux, args):
+def save_adapters(adapter_name, flux, args):
     out_dir = Path(args.output_dir)
     out_dir.mkdir(parents=True, exist_ok=True)
-    out_file = out_dir / f"{iteration:07d}_adapters.safetensors"
+    out_file = out_dir / adapter_name
     print(f"Saving {str(out_file)}")
 
     mx.save_safetensors(
@@ -157,6 +157,10 @@ if __name__ == "__main__":
     parser = setup_arg_parser()
     args = parser.parse_args()
 
+    output_path = Path(args.output_dir)
+    output_path.mkdir(parents=True, exist_ok=True)
+    save_config(vars(args), output_path / "adapter_config.json")
+
     # Load the model and set it up for LoRA training. We use the same random
     # state when creating the LoRA layers so all workers will have the same
     # initial weights.
@@ -278,8 +282,11 @@ if __name__ == "__main__":
             generate_progress_images(i + 1, flux, args)
 
         if (i + 1) % args.checkpoint_every == 0:
-            save_adapters(i + 1, flux, args)
+            save_adapters(f"{i + 1:07d}_adapters.safetensors", flux, args)
 
         if (i + 1) % 10 == 0:
             losses = []
             tic = time.time()
+
+    save_adapters("final_adapters.safetensors", flux, args)
+    print("Training successful.")

flux/flux/__init__.py

@@ -12,4 +12,5 @@ from .utils import (
     load_flow_model,
     load_t5,
     load_t5_tokenizer,
+    save_config,
 )

flux/flux/model.py

@@ -85,6 +85,8 @@ class Flux(nn.Module):
     def sanitize(self, weights):
         new_weights = {}
         for k, w in weights.items():
+            if k.startswith("model.diffusion_model."):
+                k = k[22:]
             if k.endswith(".scale"):
                 k = k[:-6] + ".weight"
             for seq in ["img_mlp", "txt_mlp", "adaLN_modulation"]:

flux/flux/sampler.py

@@ -7,7 +7,7 @@ import mlx.core as mx
 
 
 class FluxSampler:
-    def __init__(self, name: str, base_shift: float = 0.5, max_shift: float = 1.5):
+    def __init__(self, name: str, base_shift: float = 0.5, max_shift: float = 1.15):
         self._base_shift = base_shift
         self._max_shift = max_shift
         self._schnell = "schnell" in name
@@ -25,7 +25,7 @@ class FluxSampler:
     ):
         t = mx.linspace(start, stop, num_steps + 1)
 
-        if self._schnell:
+        if not self._schnell:
             t = self._time_shift(image_sequence_length, t)
 
         return t.tolist()
@@ -50,6 +50,7 @@ class FluxSampler:
             if noise is not None
             else mx.random.normal(x.shape, dtype=x.dtype, key=key)
         )
+        t = t.reshape([-1] + [1] * (x.ndim - 1))
         return x * (1 - t) + t * noise
 
     def step(self, pred, x_t, t, t_prev):

flux/flux/utils.py

@@ -3,7 +3,8 @@
 import json
 import os
 from dataclasses import dataclass
-from typing import Optional
+from pathlib import Path
+from typing import Optional, Union
 
 import mlx.core as mx
 from huggingface_hub import hf_hub_download
@@ -207,3 +208,23 @@ def load_clip_tokenizer(name: str):
 def load_t5_tokenizer(name: str, pad: bool = True):
     model_file = hf_hub_download(configs[name].repo_id, "tokenizer_2/spiece.model")
     return T5Tokenizer(model_file, 256 if "schnell" in name else 512)
+
+
+def save_config(
+    config: dict,
+    config_path: Union[str, Path],
+) -> None:
+    """Save the model configuration to the ``config_path``.
+
+    The final configuration will be sorted before saving for better readability.
+
+    Args:
+        config (dict): The model configuration.
+        config_path (Union[str, Path]): Model configuration file path.
+    """
+    # Sort the config for better readability
+    config = dict(sorted(config.items()))
+
+    # Write the config to the provided file
+    with open(config_path, "w") as fid:
+        json.dump(config, fid, indent=4)

llava/generate.py

@@ -79,10 +79,10 @@ def load_image(image_source):
 def prepare_inputs(processor, image, prompt):
     if isinstance(image, str):
         image = load_image(image)
-    inputs = processor(prompt, image, return_tensors="np")
+    inputs = processor(image, prompt, return_tensors="np")
     pixel_values = mx.array(inputs["pixel_values"])
     input_ids = mx.array(inputs["input_ids"])
-    return input_ids, pixel_values
+    return pixel_values, input_ids
 
 
 def load_model(model_path, tokenizer_config={}):
@@ -126,8 +126,7 @@ def main():
     processor, model = load_model(args.model, tokenizer_config)
 
     prompt = codecs.decode(args.prompt, "unicode_escape")
-
-    input_ids, pixel_values = prepare_inputs(processor, args.image, prompt)
+    pixel_values, input_ids = prepare_inputs(processor, args.image, prompt)
 
     print(prompt)
     generated_text = generate_text(

llava/llava.py

@@ -104,31 +104,21 @@ class LlavaModel(nn.Module):
         self, image_features, inputs_embeds, input_ids
     ):
         image_token_index = self.config.image_token_index
-        num_images, num_image_patches, embed_dim = image_features.shape
+        batch_size, num_image_patches, embed_dim = image_features.shape
 
         # Positions of <image> tokens in input_ids, assuming batch size is 1
-        image_positions = np.where(input_ids[0] == image_token_index)[0].tolist()
+        image_positions = mx.array(
+            np.where(input_ids[0] == image_token_index)[0], mx.uint32
+        )
 
-        if len(image_positions) != num_images:
+        if len(image_positions) != num_image_patches:
             raise ValueError(
                 f"The number of image tokens ({len(image_positions)}) does not "
-                f" match the number of image inputs ({num_images})."
+                f" match the number of image patches ({num_image_patches})."
             )
 
-        text_segments = []
-        start_idx = 0
-
-        for position in image_positions:
-            text_segments.append(inputs_embeds[:, start_idx:position])
-            start_idx = position + 1
-
-        image_embeddings = mx.split(image_features, image_features.shape[0])
-        final_embeddings = [v for p in zip(text_segments, image_embeddings) for v in p]
-        final_embeddings += [inputs_embeds[:, start_idx:]]
-
-        # Create a final embedding of shape
-        # (1, num_image_patches*num_images + sequence_len, embed_dim)
-        return mx.concatenate(final_embeddings, axis=1)
+        inputs_embeds[0, image_positions] = image_features
+        return inputs_embeds
 
     def __call__(self, input_ids: mx.array, pixel_values: mx.array, cache=None):
         input_embddings = self.get_input_embeddings(input_ids, pixel_values)

llms/README.md

@@ -61,7 +61,7 @@ prompt = tokenizer.apply_chat_template(
     messages, tokenize=False, add_generation_prompt=True
 )
 
-response = generate(model, tokenizer, prompt=prompt, verbose=True)
+text = generate(model, tokenizer, prompt=prompt, verbose=True)
 ```
 
 To see a description of all the arguments you can do:
@@ -77,7 +77,7 @@ to see how to use the API in more detail.
 The `mlx-lm` package also comes with functionality to quantize and optionally
 upload models to the Hugging Face Hub.
 
-You can convert models in the Python API with:
+You can convert models using the Python API:
 
 ```python
 from mlx_lm import convert
@@ -100,8 +100,10 @@ To see a description of all the arguments you can do:
 
 #### Streaming
 
-For streaming generation, use the `stream_generate` function. This returns a
-generator object which streams the output text. For example,
+For streaming generation, use the `stream_generate` function. This yields
+a generation response object.
+
+For example,
 
 ```python
 from mlx_lm import load, stream_generate
@@ -116,8 +118,8 @@ prompt = tokenizer.apply_chat_template(
     messages, tokenize=False, add_generation_prompt=True
 )
 
-for t in stream_generate(model, tokenizer, prompt, max_tokens=512):
-    print(t, end="", flush=True)
+for response in stream_generate(model, tokenizer, prompt, max_tokens=512):
+    print(response.text, end="", flush=True)
 print()
 ```
 
@@ -161,6 +163,10 @@ mlx_lm.convert \
     --upload-repo mlx-community/my-4bit-mistral
 ```
 
+Models can also be converted and quantized directly in the
+[mlx-my-repo]https://huggingface.co/spaces/mlx-community/mlx-my-repo) Hugging
+Face Space.
+
 ### Long Prompts and Generations 
 
 `mlx-lm` has some tools to scale efficiently to long prompts and generations:
@@ -221,6 +227,7 @@ Here are a few examples of Hugging Face models that work with this example:
 - [pfnet/plamo-13b-instruct](https://huggingface.co/pfnet/plamo-13b-instruct)
 - [stabilityai/stablelm-2-zephyr-1_6b](https://huggingface.co/stabilityai/stablelm-2-zephyr-1_6b)
 - [internlm/internlm2-7b](https://huggingface.co/internlm/internlm2-7b)
+- [tiiuae/falcon-mamba-7b-instruct](https://huggingface.co/tiiuae/falcon-mamba-7b-instruct)
 
 Most
 [Mistral](https://huggingface.co/models?library=transformers,safetensors&other=mistral&sort=trending),
@@ -248,3 +255,28 @@ model, tokenizer = load(
     tokenizer_config={"eos_token": "<|endoftext|>", "trust_remote_code": True},
 )
 ```
+
+### Large Models
+
+> [!NOTE]
+    This requires macOS 15.0 or higher to work.
+
+Models which are large relative to the total RAM available on the machine can
+be slow. `mlx-lm` will attempt to make them faster by wiring the memory
+occupied by the model and cache. This requires macOS 15 or higher to
+work.
+
+If you see the following warning message:
+
+> [WARNING] Generating with a model that requires ...
+
+then the model will likely be slow on the given machine. If the model fits in
+RAM then it can often be sped up by increasing the system wired memory limit.
+To increase the limit, set the following `sysctl`:
+
+```bash
+sudo sysctl iogpu.wired_limit_mb=N
+```
+
+The value `N` should be larger than the size of the model in megabytes but
+smaller than the memory size of the machine.

llms/mlx_lm/SERVER.md

@@ -92,7 +92,7 @@ curl localhost:8080/v1/chat/completions \
 
 - `system_fingerprint`: A unique identifier for the system.
 
-- `object`: Any of "chat.completions", "chat.completions.chunk" (for
+- `object`: Any of "chat.completion", "chat.completion.chunk" (for
   streaming), or "text.completion".
 
 - `model`: The model repo or path (e.g. `"mlx-community/Llama-3.2-3B-Instruct-4bit"`).

llms/mlx_lm/__init__.py

@@ -1,4 +1,9 @@
 # Copyright © 2023-2024 Apple Inc.
 
+import os
+
 from ._version import __version__
+
+os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "1"
+
 from .utils import convert, generate, load, stream_generate

llms/mlx_lm/_version.py

@@ -1,3 +1,3 @@
 # Copyright © 2023-2024 Apple Inc.
 
-__version__ = "0.19.1"
+__version__ = "0.20.4"

llms/mlx_lm/cache_prompt.py

@@ -8,7 +8,9 @@ import time
 import mlx.core as mx
 
 from .models.cache import make_prompt_cache, save_prompt_cache
-from .utils import load
+from .utils import generate_step, load
+
+DEFAULT_QUANTIZED_KV_START = 5000
 
 
 def setup_arg_parser():
@@ -48,12 +50,6 @@ def setup_arg_parser():
         action="store_true",
         help="Use the default chat template",
     )
-    parser.add_argument(
-        "--cache-limit-gb",
-        type=int,
-        default=None,
-        help="Set the MLX cache limit in GB",
-    )
     parser.add_argument(
         "--max-kv-size",
         type=int,
@@ -70,6 +66,26 @@ def setup_arg_parser():
         required=True,
         help="Message to be processed by the model ('-' reads from stdin)",
     )
+    parser.add_argument(
+        "--kv-bits",
+        type=int,
+        help="Number of bits for KV cache quantization. "
+        "Defaults to no quantization.",
+        default=None,
+    )
+    parser.add_argument(
+        "--kv-group-size",
+        type=int,
+        help="Group size for KV cache quantization.",
+        default=64,
+    )
+    parser.add_argument(
+        "--quantized-kv-start",
+        help="When --kv-bits is set, start quantizing the KV cache "
+        "from this step onwards.",
+        type=int,
+        default=DEFAULT_QUANTIZED_KV_START,
+    )
     return parser
 
 
@@ -77,9 +93,6 @@ def main():
     parser = setup_arg_parser()
     args = parser.parse_args()
 
-    if args.cache_limit_gb is not None:
-        mx.metal.set_cache_limit(args.cache_limit_gb * 1024 * 1024 * 1024)
-
     # Building tokenizer_config
     tokenizer_config = {"trust_remote_code": True if args.trust_remote_code else None}
     if args.eos_token is not None:
@@ -122,29 +135,37 @@ def main():
     y = mx.array(tokenizer.encode(prompt))
 
     # Process the prompt
-    processed = 0
-    step_size = 512
     start = time.time()
     max_msg_len = 0
-    while y.size > 0:
-        model(y[:step_size][None], cache=cache)
-        mx.eval([c.state for c in cache])
-        processed += min(y.size, step_size)
-        y = y[step_size:]
+
+    def callback(processed, total_tokens):
         current = time.time()
         speed = processed / (current - start)
         msg = f"\rProcessed {processed:6d} tokens ({speed:6.2f} tok/s)"
+        nonlocal max_msg_len
         max_msg_len = max(max_msg_len, len(msg))
         print(msg + " " * (max_msg_len - len(msg)), end="", flush=True)
+
+    for _ in generate_step(
+        y,
+        model,
+        max_tokens=0,
+        prompt_cache=cache,
+        kv_bits=args.kv_bits,
+        kv_group_size=args.kv_group_size,
+        quantized_kv_start=args.quantized_kv_start,
+        prompt_progress_callback=callback,
+    ):
+        pass
+
     print()
-    print(f"Peak memory: {mx.metal.get_peak_memory() / 2**30:.3f} GB")
+    print(f"Peak memory: {mx.metal.get_peak_memory() / 1e9:.3f} GB")
 
     print("Saving...")
     metadata = {}
     metadata["model"] = args.model
     metadata["chat_template"] = tokenizer.chat_template
     metadata["tokenizer_config"] = json.dumps(tokenizer_config)
-    print(f"Peak memory: {mx.metal.get_peak_memory() / 2**30:.3f} GB")
     save_prompt_cache(args.prompt_cache_file, cache, metadata)
 
 

llms/mlx_lm/chat.py

@@ -5,12 +5,14 @@ import json
 
 import mlx.core as mx
 
-from .models.cache import load_prompt_cache, make_prompt_cache, save_prompt_cache
+from .models.cache import make_prompt_cache
+from .sample_utils import make_sampler
 from .utils import load, stream_generate
 
 DEFAULT_TEMP = 0.0
 DEFAULT_TOP_P = 1.0
 DEFAULT_SEED = 0
+DEFAULT_MAX_TOKENS = 256
 DEFAULT_MODEL = "mlx-community/Llama-3.2-3B-Instruct-4bit"
 
 
@@ -41,6 +43,13 @@ def setup_arg_parser():
         help="Set the maximum key-value cache size",
         default=None,
     )
+    parser.add_argument(
+        "--max-tokens",
+        "-m",
+        type=int,
+        default=DEFAULT_MAX_TOKENS,
+        help="Maximum number of tokens to generate",
+    )
     return parser
 
 
@@ -56,7 +65,7 @@ def main():
         tokenizer_config={"trust_remote_code": True},
     )
 
-    print(f"[INFO] Starting chat sessiong with {args.model}. To exit, enter 'q'.")
+    print(f"[INFO] Starting chat session with {args.model}. To exit, enter 'q'.")
     prompt_cache = make_prompt_cache(model, args.max_kv_size)
     while True:
         query = input(">> ")
@@ -70,11 +79,11 @@ def main():
             model,
             tokenizer,
             prompt,
-            temp=args.temp,
-            top_p=args.top_p,
+            max_tokens=args.max_tokens,
+            sampler=make_sampler(args.temp, args.top_p),
             prompt_cache=prompt_cache,
         ):
-            print(response, flush=True, end="")
+            print(response.text, flush=True, end="")
         print()
 
 

llms/mlx_lm/convert.py

@@ -29,6 +29,12 @@ def configure_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--q-bits", help="Bits per weight for quantization.", type=int, default=4
     )
+    parser.add_argument(
+        "--q-type",
+        choices=["affine", "affine-packed"],
+        default="affine",
+        help="The type of quantization to apply",
+    )
     parser.add_argument(
         "--dtype",
         help="Type to save the non-quantized parameters.",

llms/mlx_lm/evaluate.py

@@ -0,0 +1,355 @@
+# Adapted from a PyTorch implementation by David Grangier
+
+import argparse
+import json
+import logging
+import os
+from importlib.metadata import version
+from pathlib import Path
+from typing import Optional
+
+import lm_eval
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+from lm_eval.api.model import LM
+from lm_eval.api.registry import register_model
+from tqdm import tqdm
+
+from .models.cache import make_prompt_cache
+from .utils import load, stream_generate
+
+PAD = 0
+
+
+def _len_longest_common_prefix(a, b):
+    l = 0
+    for item_a, item_b in zip(a, b):
+        if item_a != item_b:
+            break
+        l += 1
+    return l
+
+
+def _rstrip_until(s, untils):
+    """Limit a string <s> to the first occurrence of any substring in untils."""
+    l = len(s)
+    f = [s.find(u) for u in untils]
+    f = [l if x < 0 else x for x in f]
+    return s[: min(f)]
+
+
+def _pad_inputs(
+    inputs,
+    maxlen,
+    genlen=0,
+    pad_left=False,
+    pad_multiple=32,
+    truncate=False,
+):
+    # pad the prompts to the left with at least genlen tokens.
+    actual_maxlen = max(len(p) for p in inputs) + genlen
+    if actual_maxlen > maxlen:
+        if not truncate:
+            raise ValueError("Inputs are too long.")
+        else:  # drop begining
+            actual_maxlen = maxlen
+            inputs = [p[max(0, len(p) - maxlen) :] for p in inputs]
+    if pad_multiple > 0:
+        maxlen = (actual_maxlen + pad_multiple - 1) // pad_multiple
+        maxlen *= pad_multiple
+    assert PAD == 0
+    lr = np.array((1, 0) if pad_left else (0, 1))
+    return np.stack(
+        [np.pad(np.array(x, np.int32), lr * (maxlen - len(x))) for x in inputs],
+        axis=0,
+    )
+
+
+@register_model("mlxlm")
+class MLXLM(LM):
+    def __init__(
+        self,
+        path_or_hf_repo: str,
+        batch_size: int = 16,
+        max_tokens: Optional[int] = None,
+    ) -> None:
+        super().__init__()
+        self._batch_size = batch_size
+        self._model, self._tokenizer = load(path_or_hf_repo)
+        self._max_tokens = max_tokens or self._tokenizer.model_max_length
+
+    def _score_fn(self, inputs, tokenize=True, step_size=32):
+        if tokenize:
+            inputs = self._tokenizer.encode(inputs)
+        inputs = _pad_inputs(inputs, self._max_tokens, truncate=False)
+        inputs = mx.array(inputs)
+        inputs, targets = inputs[..., :-1], inputs[..., 1:]
+
+        cache = make_prompt_cache(self._model)
+
+        mask = targets != PAD
+
+        scores, is_greedy = [], []
+        for i in range(0, inputs.shape[1], step_size):
+            logits = self._model(inputs[:, i : i + step_size], cache=cache)
+
+            log_probs = nn.log_softmax(logits.astype(mx.float32))
+            score = mx.take_along_axis(
+                log_probs, targets[:, i : i + step_size, mx.newaxis], axis=-1
+            )[..., 0]
+            ig = mask[:, i : i + step_size] * (
+                targets[:, i : i + step_size] == mx.argmax(logits, axis=-1)
+            )
+
+            mx.eval(score, ig)
+            mx.metal.clear_cache()
+
+            is_greedy.append(ig)
+            scores.append(score)
+
+        scores = mx.concatenate(scores, axis=1)
+        is_greedy = mx.concatenate(is_greedy, axis=1)
+
+        return scores, mask.sum(axis=-1), is_greedy
+
+    def _loglikelihood(self, texts, score_spans=None, tokenize=True):
+        # sort by length to get batches with little padding.
+        sorted_indices = sorted(range(len(texts)), key=lambda i: -len(texts[i]))
+        sorted_inputs = [texts[sorted_indices[i]] for i in range(len(texts))]
+        sorted_spans = None
+        if score_spans is not None:
+            sorted_spans = [score_spans[sorted_indices[i]] for i in range(len(texts))]
+
+        results = []
+        for i in tqdm(range(0, len(sorted_inputs), self._batch_size)):
+            batch = sorted_inputs[i : i + self._batch_size]
+            scores, length, is_greedy = self._score_fn(batch, tokenize=tokenize)
+            for j in range(len(batch)):
+                if sorted_spans is None:  # full sequence score
+                    mask = mx.arange(scores[j].shape[-1]) < length
+                    score = (scores[j].astype(mx.float32) * mask).sum(axis=-1)
+                    ig = (is_greedy[j].astype(mx.int32) * mask).sum(axis=-1)
+                else:  # subsequence score
+                    start, end = sorted_spans[i + j]
+                    score = scores[j][start:end].astype(mx.float32).sum()
+                    ig = is_greedy[j][start:end].astype(mx.int32).sum()
+                    length = end - start
+
+                results.append((score.item(), ig.item(), length))
+
+        # reorder the outputs
+        inv_sort = np.argsort(sorted_indices)
+        results = [results[inv_sort[i]] for i in range(len(results))]
+
+        return results
+
+    def _tokenize(self, texts):
+        return [tuple(self._tokenizer.encode(t)) for t in texts]
+
+    def loglikelihood(self, requests) -> list[tuple[float, bool]]:
+        """Compute log-likelihood of generating a continuation from a context.
+        Downstream tasks should attempt to use loglikelihood instead of other
+        LM calls whenever possible.
+        :param requests: list[Instance]
+            A list of Instance objects, with property `args` which returns a tuple (context, continuation).
+            `context: str`
+                Context string. Implementations of LM must be able to handle an
+                empty context string.
+            `continuation: str`
+                The continuation over which log likelihood will be calculated. If
+                there is a word boundary, the space should be in the continuation.
+                For example, context="hello" continuation=" world" is correct.
+        :return: list[tuple[float, bool]]
+            A list of pairs (logprob, isgreedy)
+            `logprob: float`
+                The log probability of `continuation`.
+            `isgreedy`:
+                Whether `continuation` would be generated by greedy sampling from `context`.
+        """
+        logging.info("Estimating loglikelihood for %d pairs." % len(requests))
+
+        # tokenize prefix and prefix + completion for all requests.
+        tokenized = self._tokenize(
+            [t for r in requests for t in [r.args[0], r.args[0] + r.args[1]]]
+        )
+
+        # max length (prefix + completion) and longest common prefix per question.
+        length_stats = {}
+        for prefix, completed in zip(tokenized[0::2], tokenized[1::2]):
+            max_completed_l, min_prefix_l = length_stats.get(prefix, (0, 1e8))
+            length_stats[prefix] = (
+                max(max_completed_l, len(completed)),
+                min(min_prefix_l, _len_longest_common_prefix(prefix, completed)),
+            )
+
+        # truncate requests for completed sequences longer than model context.
+        shortened = []
+        completion_spans = []
+        long_completions = 0
+        for prefix, completed in zip(tokenized[0::2], tokenized[1::2]):
+            max_completed_l, prefix_l = length_stats[prefix]
+            # compute truncation length
+            truncation = max(0, max_completed_l - self._max_tokens - 1)
+            prefix_l = prefix_l - truncation
+            if prefix_l <= 0:
+                # completion too long, prefix is eliminated for some requests.
+                long_completions += 1
+                truncation = max(0, len(completed) - self._max_tokens - 1)
+                prefix_l = 1
+            # truncate the completed sequence
+            completed = completed[truncation:]
+            shortened.append(completed)
+            # scores do not include initial bos, substract 1 to span bounds
+            completion_spans.append((prefix_l - 1, len(completed) - 1))
+
+        if long_completions > 0:
+            logging.info(
+                f"Prefix eliminated for {long_completions} requests with "
+                + "completion longer than context."
+            )
+
+        # model scoring, returns num_requests x (logp, is_greedy, length).
+        results = self._loglikelihood(
+            shortened,
+            score_spans=completion_spans,
+            tokenize=False,
+        )
+        return [(r[0], r[1] == r[2]) for r in results]
+
+    def loglikelihood_rolling(self, requests) -> list[float]:
+        """Compute full log-likelihood of a string, with no truncation, for perplexity computation
+        - We will use the full max context length of the model.
+        - For inputs that exceed the max context length, we divide the tokenized string into chunks of up to
+        the max context length.
+        - IMPORTANT: Each document's loglikelihood/perplexity is computed *separately*, unlike other implementations
+          which may simply concatenate multiple documents together.
+        - IMPORTANT: We maximize the amount of context for each prediction. Specifically, for inputs that we break into
+          multiple chunks, the last input will still a full-sized context.
+          Example:
+            Input tokens: [ 0 1 2 3 4 5 6 7 8 9 ]
+            Prefix: EOT
+            Max context length: 4
+            Resulting input/prediction pairs:
+                INPUT:  EOT   0   1   2
+                PRED:     0   1   2   3
+                INPUT:    3   4   5   6
+                PRED:     4   5   6   7
+                INPUT:    5   6   7   8
+                PRED:             8   9
+          Observe that:
+            1. Each token is predicted exactly once
+            2. For the last pair, we provide the full context, but only score the last two tokens
+        :param requests: list[Instance]
+            A list of Instance objects with property `args` which returns a tuple (context,).
+            string: str
+                String for which we are computing overall loglikelihood
+        :return: list[tuple[float]]
+            A list of tuples (logprob,)
+            logprob: float
+                The log probability of `context` conditioned on the EOT token.
+        """
+        logging.info(
+            "Estimating loglikelihood rolling for %d sequences." % len(requests)
+        )
+        inputs = [req.args[0] for req in requests]
+        return [t[0] for t in self._loglikelihood(inputs)]
+
+    def generate_until(self, requests) -> list[str]:
+        """Generate greedily until a stopping sequence
+        :param requests: list[Instance]
+            A list of Instance objects with property `args` which returns a tuple (context, until).
+            context: str
+                Context string
+            until: [str]
+                The string sequences to generate until. These string sequences
+                may each span across multiple tokens, or may be part of one token.
+        :return: list[str]
+            A list of strings continuation
+            continuation: str
+                The generated continuation.
+        """
+        logging.info("Generating continuation for %d sequences." % len(requests))
+        contexts, options = zip(*[req.args for req in requests])
+        # contrary to the doc the second element of the tuple contains
+        # {'do_sample': False, 'until': ['\n\n'], 'temperature': 0}
+        keys = list(options[0].keys())
+        assert "until" in keys
+        untils = [x["until"] for x in options]
+        completions = []
+        for context, until in tqdm(zip(contexts, untils), total=len(contexts)):
+            if (
+                hasattr(self._tokenizer, "apply_chat_template")
+                and self._tokenizer.chat_template is not None
+            ):
+                messages = [{"role": "user", "content": context}]
+                context = self._tokenizer.apply_chat_template(
+                    messages, tokenize=False, add_generation_prompt=True
+                )
+
+            max_tokens = min(
+                self._max_tokens,
+                self._tokenizer.model_max_length - len(self._tokenizer.encode(context)),
+            )
+            text = ""
+            for response in stream_generate(
+                self._model, self._tokenizer, prompt=context, max_tokens=max_tokens
+            ):
+                text += response.text
+                if any(u in text for u in until):
+                    text = _rstrip_until(text, until)
+                    completions.append(text)
+                    break
+            else:
+                completions.append(text)
+        return completions
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        "Evaluate an MLX model using lm-evaluation-harness."
+    )
+    parser.add_argument("--model", help="Model to evaluate", required=True)
+    parser.add_argument("--tasks", nargs="+", required=True)
+    parser.add_argument(
+        "--output-dir", default=".", help="Output directory for result files."
+    )
+    parser.add_argument("--batch-size", type=int, default=16, help="Batch size")
+    parser.add_argument("--num-shots", type=int, default=0, help="Number of shots")
+    parser.add_argument(
+        "--max-tokens",
+        type=int,
+        help="Maximum nunber of tokens to generate. Defaults to the model's max context length.",
+    )
+    parser.add_argument("--seed", type=int, default=123, help="Random seed.")
+    args = parser.parse_args()
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Silence tokenizer warnings
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+    mx.random.seed(args.seed)
+
+    lm = MLXLM(args.model, batch_size=args.batch_size, max_tokens=args.max_tokens)
+
+    results = lm_eval.simple_evaluate(
+        model=lm,
+        tasks=args.tasks,
+        num_fewshot=args.num_shots,
+        random_seed=args.seed,
+        numpy_random_seed=args.seed,
+        torch_random_seed=args.seed,
+        fewshot_random_seed=args.seed,
+    )
+
+    model_name = args.model.replace("/", "_")
+    task_names = "_".join(args.tasks)
+    ver = version("lm_eval")
+    filename = f"eval_{model_name}_{task_names}_{args.num_shots:02d}_v_{ver}.json"
+    output_path = output_dir / filename
+    output_path.write_text(json.dumps(results["results"], indent=4))
+    print("Results:")
+    for result in results["results"].values():
+        print(json.dumps(result, indent=4))

llms/mlx_lm/examples/chat.py

@@ -42,7 +42,6 @@ response = generate(
     tokenizer,
     prompt=prompt,
     verbose=True,
-    temp=0.0,
     prompt_cache=prompt_cache,
 )
 

llms/mlx_lm/examples/generate_response.py

@@ -23,14 +23,6 @@ max_tokens = 1_000
 # Specify if tokens and timing information will be printed
 verbose = True
 
-# Some optional arguments for causal language model generation
-generation_args = {
-    "temp": 0.7,
-    "repetition_penalty": 1.2,
-    "repetition_context_size": 20,
-    "top_p": 0.95,
-}
-
 # Generate a response with the specified settings
 response = generate(
     model=model,
@@ -38,5 +30,4 @@ response = generate(
     prompt=prompt,
     max_tokens=max_tokens,
     verbose=verbose,
-    **generation_args,
 )

llms/mlx_lm/examples/lora_config.yaml

@@ -14,7 +14,7 @@ data: "/path/to/training/data"
 seed: 0
 
 # Number of layers to fine-tune
-lora_layers: 16
+num_layers: 16
 
 # Minibatch size.
 batch_size: 4

llms/mlx_lm/generate.py

@@ -1,20 +1,25 @@
 # Copyright © 2023-2024 Apple Inc.
 
 import argparse
+import codecs
 import json
 import sys
 
 import mlx.core as mx
 
-from .models.cache import load_prompt_cache
+from .models.cache import QuantizedKVCache, load_prompt_cache
+from .sample_utils import make_sampler
 from .utils import generate, load
 
 DEFAULT_PROMPT = "hello"
 DEFAULT_MAX_TOKENS = 100
 DEFAULT_TEMP = 0.0
 DEFAULT_TOP_P = 1.0
+DEFAULT_MIN_P = 0.0
+DEFAULT_MIN_TOKENS_TO_KEEP = 1
 DEFAULT_SEED = 0
 DEFAULT_MODEL = "mlx-community/Llama-3.2-3B-Instruct-4bit"
+DEFAULT_QUANTIZED_KV_START = 5000
 
 
 def str2bool(string):
@@ -38,19 +43,20 @@ def setup_arg_parser():
         type=str,
         help="Optional path for the trained adapter weights and config.",
     )
-    parser.add_argument(
-        "--trust-remote-code",
-        action="store_true",
-        help="Enable trusting remote code for tokenizer",
-    )
     parser.add_argument(
         "--eos-token",
         type=str,
         default=None,
         help="End of sequence token for tokenizer",
     )
+    parser.add_argument(
+        "--system-prompt",
+        default=None,
+        help="System prompt to be used for the chat template",
+    )
     parser.add_argument(
         "--prompt",
+        "-p",
         default=DEFAULT_PROMPT,
         help="Message to be processed by the model ('-' reads from stdin)",
     )
@@ -67,6 +73,15 @@ def setup_arg_parser():
     parser.add_argument(
         "--top-p", type=float, default=DEFAULT_TOP_P, help="Sampling top-p"
     )
+    parser.add_argument(
+        "--min-p", type=float, default=DEFAULT_MIN_P, help="Sampling min-p"
+    )
+    parser.add_argument(
+        "--min-tokens-to-keep",
+        type=int,
+        default=DEFAULT_MIN_TOKENS_TO_KEEP,
+        help="Minimum tokens to keep for min-p sampling.",
+    )
     parser.add_argument("--seed", type=int, default=DEFAULT_SEED, help="PRNG seed")
     parser.add_argument(
         "--ignore-chat-template",
@@ -84,17 +99,6 @@ def setup_arg_parser():
         default=True,
         help="Log verbose output when 'True' or 'T' or only print the response when 'False' or 'F'",
     )
-    parser.add_argument(
-        "--colorize",
-        action="store_true",
-        help="Colorize output based on T[0] probability",
-    )
-    parser.add_argument(
-        "--cache-limit-gb",
-        type=int,
-        default=None,
-        help="Set the MLX cache limit in GB",
-    )
     parser.add_argument(
         "--max-kv-size",
         type=int,
@@ -107,58 +111,57 @@ def setup_arg_parser():
         default=None,
         help="A file containing saved KV caches to avoid recomputing them",
     )
+    parser.add_argument(
+        "--kv-bits",
+        type=int,
+        help="Number of bits for KV cache quantization. "
+        "Defaults to no quantization.",
+        default=None,
+    )
+    parser.add_argument(
+        "--kv-group-size",
+        type=int,
+        help="Group size for KV cache quantization.",
+        default=64,
+    )
+    parser.add_argument(
+        "--quantized-kv-start",
+        help="When --kv-bits is set, start quantizing the KV cache "
+        "from this step onwards.",
+        type=int,
+        default=DEFAULT_QUANTIZED_KV_START,
+    )
     return parser
 
 
-def colorprint(color, s):
-    color_codes = {
-        "black": 30,
-        "red": 31,
-        "green": 32,
-        "yellow": 33,
-        "blue": 34,
-        "magenta": 35,
-        "cyan": 36,
-        "white": 39,
-    }
-    ccode = color_codes.get(color, 30)
-    print(f"\033[1m\033[{ccode}m{s}\033[0m", end="", flush=True)
-
-
-def colorprint_by_t0(s, t0):
-    if t0 > 0.95:
-        color = "white"
-    elif t0 > 0.70:
-        color = "green"
-    elif t0 > 0.30:
-        color = "yellow"
-    else:
-        color = "red"
-    colorprint(color, s)
-
-
 def main():
     parser = setup_arg_parser()
     args = parser.parse_args()
 
     mx.random.seed(args.seed)
 
-    if args.cache_limit_gb is not None:
-        mx.metal.set_cache_limit(args.cache_limit_gb * 1024 * 1024 * 1024)
-
     # Load the prompt cache and metadata if a cache file is provided
     using_cache = args.prompt_cache_file is not None
     if using_cache:
         prompt_cache, metadata = load_prompt_cache(
-            args.prompt_cache_file, return_metadata=True
+            args.prompt_cache_file,
+            return_metadata=True,
         )
+        if isinstance(prompt_cache[0], QuantizedKVCache):
+            if args.kv_bits is not None and args.kv_bits != prompt_cache[0].bits:
+                raise ValueError(
+                    "--kv-bits does not match the kv cache loaded from --prompt-cache-file."
+                )
+            if args.kv_group_size != prompt_cache[0].group_size:
+                raise ValueError(
+                    "--kv-group-size does not match the kv cache loaded from --prompt-cache-file."
+                )
 
     # Building tokenizer_config
     tokenizer_config = (
         {} if not using_cache else json.loads(metadata["tokenizer_config"])
     )
-    if args.trust_remote_code:
-        tokenizer_config["trust_remote_code"] = True
+    tokenizer_config["trust_remote_code"] = True
     if args.eos_token is not None:
         tokenizer_config["eos_token"] = args.eos_token
 
@@ -186,16 +189,22 @@ def main():
     elif using_cache:
         tokenizer.chat_template = metadata["chat_template"]
 
+    prompt = codecs.decode(args.prompt, "unicode_escape")
+
     if not args.ignore_chat_template and (
         hasattr(tokenizer, "apply_chat_template")
         and tokenizer.chat_template is not None
     ):
-        messages = [
+        if args.system_prompt is not None:
+            messages = [{"role": "system", "content": args.system_prompt}]
+        else:
+            messages = []
+        messages.append(
             {
                 "role": "user",
-                "content": sys.stdin.read() if args.prompt == "-" else args.prompt,
+                "content": sys.stdin.read() if prompt == "-" else prompt,
             }
-        ]
+        )
         prompt = tokenizer.apply_chat_template(
             messages, tokenize=False, add_generation_prompt=True
         )
@@ -203,30 +212,27 @@ def main():
         # Treat the prompt as a suffix assuming that the prefix is in the
         # stored kv cache.
         if using_cache:
+            messages[-1]["content"] = "<query>"
             test_prompt = tokenizer.apply_chat_template(
-                [{"role": "user", "content": "<query>"}],
+                messages,
                 tokenize=False,
                 add_generation_prompt=True,
             )
             prompt = prompt[test_prompt.index("<query>") :]
-    else:
-        prompt = args.prompt
-
-    if args.colorize and not args.verbose:
-        raise ValueError("Cannot use --colorize with --verbose=False")
-    formatter = colorprint_by_t0 if args.colorize else None
 
+    sampler = make_sampler(args.temp, args.top_p, args.min_p, args.min_tokens_to_keep)
     response = generate(
         model,
         tokenizer,
         prompt,
-        args.max_tokens,
+        max_tokens=args.max_tokens,
         verbose=args.verbose,
-        formatter=formatter,
-        temp=args.temp,
-        top_p=args.top_p,
+        sampler=sampler,
         max_kv_size=args.max_kv_size,
         prompt_cache=prompt_cache if using_cache else None,
+        kv_bits=args.kv_bits,
+        kv_group_size=args.kv_group_size,
+        quantized_kv_start=args.quantized_kv_start,
     )
     if not args.verbose:
         print(response)

llms/mlx_lm/models/base.py

@@ -5,6 +5,9 @@ from dataclasses import dataclass
 from typing import Any, Optional
 
 import mlx.core as mx
+from mlx.utils import tree_map
+
+from .cache import QuantizedKVCache
 
 
 @dataclass
@@ -39,7 +42,7 @@ def create_attention_mask(h: mx.array, cache: Optional[Any] = None):
         if cache is not None and cache[0] is not None:
             c = cache[0]
             if hasattr(c, "max_size"):
-                offset = min(c.max_size - 1, c.offset)
+                offset = min(c.max_size, c.offset)
                 window_size = c.max_size
             else:
                 offset = c.offset
@@ -48,3 +51,63 @@ def create_attention_mask(h: mx.array, cache: Optional[Any] = None):
     else:
         mask = None
     return mask
+
+
+def quantized_scaled_dot_product_attention(
+    queries: mx.array,
+    q_keys: tuple[mx.array, mx.array, mx.array],
+    q_values: tuple[mx.array, mx.array, mx.array],
+    scale: float,
+    mask: Optional[mx.array],
+    group_size: int = 64,
+    bits: int = 8,
+) -> mx.array:
+    B, n_q_heads, L, D = queries.shape
+    n_kv_heads = q_keys[0].shape[-3]
+    n_repeats = n_q_heads // n_kv_heads
+
+    queries *= scale
+
+    if n_repeats > 1:
+        queries = mx.reshape(queries, (B, n_kv_heads, n_repeats, L, D))
+        q_keys = tree_map(lambda x: mx.expand_dims(x, axis=-3), q_keys)
+        q_values = tree_map(lambda x: mx.expand_dims(x, axis=-3), q_values)
+
+    scores = mx.quantized_matmul(
+        queries, *q_keys, transpose=True, group_size=group_size, bits=bits
+    )
+    if mask is not None:
+        scores += mask
+    scores = mx.softmax(scores, axis=-1, precise=True)
+    out = mx.quantized_matmul(
+        scores, *q_values, transpose=False, group_size=group_size, bits=bits
+    )
+
+    if n_repeats > 1:
+        out = mx.reshape(out, (B, n_q_heads, L, D))
+
+    return out
+
+
+def scaled_dot_product_attention(
+    queries,
+    keys,
+    values,
+    cache,
+    scale: float,
+    mask: Optional[mx.array],
+) -> mx.array:
+    if isinstance(cache, QuantizedKVCache):
+        return quantized_scaled_dot_product_attention(
+            queries,
+            keys,
+            values,
+            scale=scale,
+            mask=mask,
+            group_size=cache.group_size,
+            bits=cache.bits,
+        )
+    else:
+        return mx.fast.scaled_dot_product_attention(
+            queries, keys, values, scale=scale, mask=mask
+        )

llms/mlx_lm/models/cache.py

@@ -4,10 +4,13 @@ from typing import Any, Dict, List, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
-from mlx.utils import tree_flatten, tree_unflatten
+from mlx.utils import tree_flatten, tree_map, tree_unflatten
 
 
-def make_prompt_cache(model: nn.Module, max_kv_size: Optional[int] = None) -> List[Any]:
+def make_prompt_cache(
+    model: nn.Module,
+    max_kv_size: Optional[int] = None,
+) -> List[Any]:
     """
     Construct the model's cache for use when cgeneration.
 
@@ -126,6 +129,88 @@ class _BaseCache:
         return False
 
 
+class QuantizedKVCache(_BaseCache):
+    def __init__(self, group_size: int = 64, bits: int = 8):
+        self.keys = None
+        self.values = None
+        self.offset = 0
+        self.step = 256
+        self.group_size = group_size
+        self.bits = bits
+
+    def update_and_fetch(self, keys, values):
+        B, n_kv_heads, num_steps, k_head_dim = keys.shape
+        v_head_dim = values.shape[-1]
+        prev = self.offset
+
+        if self.keys is None or (prev + num_steps) > self.keys[0].shape[-2]:
+            el_per_int = 8 * mx.uint32.size // self.bits
+            new_steps = (self.step + num_steps - 1) // self.step * self.step
+            shape = (B, n_kv_heads, new_steps)
+
+            def init_quant(dim):
+                return (
+                    mx.zeros((*shape, dim // el_per_int), dtype=mx.uint32),
+                    mx.zeros((*shape, dim // self.group_size), dtype=keys.dtype),
+                    mx.zeros((*shape, dim // self.group_size), dtype=keys.dtype),
+                )
+
+            def expand_quant(x):
+                new_x = mx.zeros((*shape, x.shape[-1]), dtype=x.dtype)
+                return mx.concatenate([x, new_x], axis=-2)
+
+            if self.keys is not None:
+                if prev % self.step != 0:
+                    self.keys, self.values = tree_map(
+                        lambda x: x[..., :prev, :], (self.keys, self.values)
+                    )
+
+                self.keys, self.values = tree_map(
+                    expand_quant, (self.keys, self.values)
+                )
+            else:
+                self.keys, self.values = init_quant(k_head_dim), init_quant(v_head_dim)
+
+        self.offset += num_steps
+
+        keys = mx.quantize(keys, group_size=self.group_size, bits=self.bits)
+        values = mx.quantize(values, group_size=self.group_size, bits=self.bits)
+        for i in range(len(self.keys)):
+            self.keys[i][..., prev : self.offset, :] = keys[i]
+            self.values[i][..., prev : self.offset, :] = values[i]
+
+        return tree_map(lambda x: x[..., : self.offset, :], (self.keys, self.values))
+
+    @property
+    def state(self):
+        if self.offset == self.keys[0].shape[2]:
+            return self.keys, self.values
+        else:
+            return tree_map(
+                lambda x: x[..., : self.offset, :], (self.keys, self.values)
+            )
+
+    @state.setter
+    def state(self, v):
+        self.keys, self.values = v
+
+    @property
+    def meta_state(self):
+        return tuple(map(str, (self.step, self.offset, self.group_size, self.bits)))
+
+    @meta_state.setter
+    def meta_state(self, v):
+        self.step, self.offset, self.group_size, self.bits = map(int, v)
+
+    def is_trimmable(self):
+        return True
+
+    def trim(self, n):
+        n = min(self.offset, n)
+        self.offset -= n
+        return n
+
+
 class KVCache(_BaseCache):
     def __init__(self):
         self.keys = None
@@ -180,6 +265,16 @@ class KVCache(_BaseCache):
         self.offset -= n
         return n
 
+    def to_quantized(self, group_size: int = 64, bits: int = 4) -> QuantizedKVCache:
+        quant_cache = QuantizedKVCache(group_size=group_size, bits=bits)
+        quant_cache.offset = self.offset
+        if self.keys is not None:
+            quant_cache.keys = mx.quantize(self.keys, group_size=group_size, bits=bits)
+            quant_cache.values = mx.quantize(
+                self.values, group_size=group_size, bits=bits
+            )
+        return quant_cache
+
 
 class RotatingKVCache(_BaseCache):
 
@@ -230,9 +325,9 @@ class RotatingKVCache(_BaseCache):
             self.keys = self._temporal_order(self.keys)
             self.values = self._temporal_order(self.values)
 
-            # The largest size is self.max_size + S - 1 to ensure
+            # The largest size is self.max_size + S to ensure
             # every token gets at least self.max_size context
-            trim_size = self._idx - self.max_size + 1
+            trim_size = self._idx - self.max_size
             self.keys = self._trim(trim_size, self.keys, keys)
             self.values = self._trim(trim_size, self.values, values)
         self.offset += keys.shape[2]
@@ -320,6 +415,9 @@ class RotatingKVCache(_BaseCache):
         self._idx -= n
         return n
 
+    def to_quantized(self, group_size: int = 64, bits: int = 4) -> QuantizedKVCache:
+        raise NotImplementedError("RotatingKVCache Quantization NYI")
+
 
 class MambaCache(_BaseCache):
     def __init__(self):

llms/mlx_lm/models/cohere.py

@@ -6,7 +6,7 @@ from typing import Any, Optional, Tuple
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs, create_attention_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -93,8 +93,8 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        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)

llms/mlx_lm/models/cohere2.py

@@ -0,0 +1,207 @@
+# Copyright © 2023-2024 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_causal_mask, scaled_dot_product_attention
+from .cache import KVCache, RotatingKVCache
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int = 4096
+    head_dim: int = 128
+    num_hidden_layers: int = 32
+    intermediate_size: int = 14336
+    num_attention_heads: int = 32
+    num_key_value_heads: int = 8
+    rope_theta: float = 50000.0
+    vocab_size: int = 256000
+    layer_norm_eps: float = 1e-05
+    logit_scale: float = 0.0625
+    attention_bias: bool = False
+    layer_norm_bias: bool = False
+    sliding_window: int = 4096
+    sliding_window_pattern: int = 4
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.args = args
+        self.layer_idx = layer_idx
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+        self.head_dim = head_dim = args.head_dim
+        if (head_dim * n_heads) != dim:
+            raise ValueError(
+                f"hidden_size must be divisible by num_heads (got `hidden_size`: {dim}"
+                f" and `num_heads`: {n_heads})."
+            )
+        self.scale = head_dim**-0.5
+
+        attetion_bias = args.attention_bias
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=attetion_bias)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attetion_bias)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attetion_bias)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=attetion_bias)
+
+        self.rope = nn.RoPE(head_dim, traditional=True, base=args.rope_theta)
+
+        self.use_sliding_window = (layer_idx + 1) % args.sliding_window_pattern != 0
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Tuple[mx.array, mx.array]] = 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)
+
+        # Apply RoPE only if sliding window is enabled
+        if self.use_sliding_window:
+            if cache is None:
+                queries = self.rope(queries)
+                keys = self.rope(keys)
+            else:
+                queries = self.rope(queries, offset=cache.offset)
+                keys = self.rope(keys, offset=cache.offset)
+
+        if cache is not None:
+            keys, values = cache.update_and_fetch(keys, values)
+
+        if self.use_sliding_window and mask is not None:
+            key_len = keys.shape[-2]
+            if mask.shape[-1] != key_len:
+                mask = mask[..., -key_len:]
+
+        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 MLP(nn.Module):
+    def __init__(self, dim, hidden_dim):
+        super().__init__()
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
+
+    def __call__(self, x):
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class TransformerBlock(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.self_attn = Attention(args, layer_idx)
+        self.mlp = MLP(args.hidden_size, args.intermediate_size)
+        self.input_layernorm = nn.LayerNorm(
+            args.hidden_size, eps=args.layer_norm_eps, bias=args.layer_norm_bias
+        )
+        self.args = args
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Tuple[mx.array, mx.array]] = None,
+    ) -> mx.array:
+        h = self.input_layernorm(x)
+        attn_h = self.self_attn(h, mask, cache)
+        ff_h = self.mlp(h)
+        return attn_h + ff_h + x
+
+
+class CohereModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.vocab_size = args.vocab_size
+        self.num_hidden_layers = args.num_hidden_layers
+        assert self.vocab_size > 0
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            TransformerBlock(args=args, layer_idx=i)
+            for i in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.LayerNorm(
+            args.hidden_size, eps=args.layer_norm_eps, bias=args.layer_norm_bias
+        )
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+    ):
+        h = self.embed_tokens(inputs)
+
+        T = h.shape[1]
+        if T > 1:
+            offset = cache[0].offset if cache else 0
+            mask = create_causal_mask(T, offset).astype(h.dtype)
+        else:
+            mask = None
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, c)
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.model_type = args.model_type
+        self.model = CohereModel(args)
+        self.args = args
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+    ):
+        out = self.model(inputs, cache)
+        out = self.model.embed_tokens.as_linear(out)
+        out = out * self.model.args.logit_scale
+        return out
+
+    def make_cache(self):
+        caches = []
+        for i in range(self.args.num_hidden_layers):
+            if (
+                i % self.args.sliding_window_pattern
+                == self.args.sliding_window_pattern - 1
+            ):
+                caches.append(KVCache())
+            else:
+                caches.append(
+                    RotatingKVCache(max_size=self.args.sliding_window, keep=0)
+                )
+        return caches
+
+    @property
+    def layers(self):
+        return self.model.layers

llms/mlx_lm/models/dbrx.py

@@ -7,7 +7,7 @@ import mlx.core as mx
 import mlx.nn as nn
 import numpy as np
 
-from .base import BaseModelArgs, create_attention_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -74,8 +74,8 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        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.out_proj(output)

llms/mlx_lm/models/deepseek.py

@@ -4,7 +4,7 @@ from typing import Any, Dict, Optional
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs, create_attention_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .switch_layers import SwitchGLU
 
 
@@ -97,8 +97,8 @@ class DeepseekAttention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        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)

llms/mlx_lm/models/deepseek_v2.py

@@ -7,7 +7,7 @@ from typing import Any, Dict, Optional, Tuple
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs, create_attention_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .switch_layers import SwitchGLU
 
 
@@ -235,8 +235,8 @@ class DeepseekV2Attention(nn.Module):
 
         queries = mx.concatenate([q_nope, q_pe], axis=-1)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        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)

llms/mlx_lm/models/exaone.py

@@ -0,0 +1,163 @@
+# Copyright © 2024 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Union
+
+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
+    hidden_size: int
+    num_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    vocab_size: int
+    rope_theta: float
+    layer_norm_epsilon: float
+    num_key_value_heads: int
+    head_dim: Optional[int] = None
+    max_position_embeddings: Optional[int] = None
+    rope_traditional: bool = False
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    tie_word_embeddings: bool = True
+    attention_bias: bool = False
+    mlp_bias: bool = False
+
+
+class AttentionModule(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+        self.head_dim = head_dim = args.head_dim or (dim // n_heads)
+        self.scale = head_dim**-0.5
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=args.attention_bias)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.attention_bias)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.attention_bias)
+        self.out_proj = nn.Linear(n_heads * head_dim, dim, bias=args.attention_bias)
+
+        self.rope = initialize_rope(
+            self.head_dim,
+            args.rope_theta,
+            args.rope_traditional,
+            args.rope_scaling,
+            args.max_position_embeddings,
+        )
+
+    def __call__(
+        self, x: mx.array, mask: Optional[mx.array] = None, cache: Optional[Any] = None
+    ) -> mx.array:
+        B, L, D = x.shape
+        q = self.q_proj(x).reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        k = self.k_proj(x).reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        v = self.v_proj(x).reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            q = self.rope(q, offset=cache.offset)
+            k = self.rope(k, offset=cache.offset)
+            k, v = cache.update_and_fetch(k, v)
+        else:
+            q = self.rope(q)
+            k = self.rope(k)
+
+        out = scaled_dot_product_attention(
+            q, k, v, cache=cache, scale=self.scale, mask=mask
+        )
+        out = out.transpose(0, 2, 1, 3).reshape(B, L, D)
+        return self.out_proj(out)
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.attention = AttentionModule(args)
+
+
+class MLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        dim = args.hidden_size
+        hidden_dim = args.intermediate_size
+        self.c_fc_0 = nn.Linear(dim, hidden_dim, bias=args.mlp_bias)
+        self.c_fc_1 = nn.Linear(dim, hidden_dim, bias=args.mlp_bias)
+        self.c_proj = nn.Linear(hidden_dim, dim, bias=args.mlp_bias)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.c_proj(nn.silu(self.c_fc_0(x)) * self.c_fc_1(x))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.ln_1 = nn.RMSNorm(args.hidden_size, eps=args.layer_norm_epsilon)
+        self.attn = Attention(args)
+        self.ln_2 = nn.RMSNorm(args.hidden_size, eps=args.layer_norm_epsilon)
+        self.mlp = MLP(args)
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        h = x + self.attn.attention(self.ln_1(x), mask, cache)
+        out = h + self.mlp(self.ln_2(h))
+        return out
+
+
+class ExaoneModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.wte = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.h = [TransformerBlock(args) for _ in range(args.num_layers)]
+        self.ln_f = nn.RMSNorm(args.hidden_size, eps=args.layer_norm_epsilon)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+    ):
+        h = self.wte(inputs)
+        mask = create_attention_mask(h, cache)
+
+        if cache is None:
+            cache = [None] * len(self.h)
+
+        for layer, c in zip(self.h, cache):
+            h = layer(h, mask, cache=c)
+
+        return self.ln_f(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.transformer = ExaoneModel(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+    ):
+        out = self.transformer(inputs, cache)
+        if self.args.tie_word_embeddings:
+            out = self.transformer.wte.as_linear(out)
+        else:
+            out = self.lm_head(out)
+        return out
+
+    @property
+    def layers(self):
+        return self.transformer.h

llms/mlx_lm/models/gemma.py

@@ -6,7 +6,7 @@ from typing import Any, Optional, Tuple
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs, create_attention_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -79,8 +79,8 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        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)

llms/mlx_lm/models/gpt2.py

@@ -7,7 +7,7 @@ import mlx.core as mx
 import mlx.nn as nn
 import numpy as np
 
-from .base import BaseModelArgs, create_attention_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -61,8 +61,8 @@ class Attention(nn.Module):
         if cache is not None:
             keys, values = cache.update_and_fetch(keys, values)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        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)

llms/mlx_lm/models/gpt_bigcode.py

@@ -7,7 +7,7 @@ import mlx.core as mx
 import mlx.nn as nn
 import numpy as np
 
-from .base import BaseModelArgs, create_attention_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -74,8 +74,8 @@ class Attention(nn.Module):
         if cache is not None:
             keys, values = cache.update_and_fetch(keys, values)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        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.c_proj(output)

llms/mlx_lm/models/gpt_neox.py

@@ -7,7 +7,7 @@ import mlx.core as mx
 import mlx.nn as nn
 import numpy as np
 
-from .base import BaseModelArgs, create_attention_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 # Based on the transformers implementation at:
 # https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_neox/modeling_gpt_neox.py
@@ -79,8 +79,8 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        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)

llms/mlx_lm/models/hunyuan.py

@@ -0,0 +1,291 @@
+# Copyright © 2023-2024 Apple Inc.
+
+import math
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Tuple, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    vocab_size: int
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    attention_bias: bool
+    moe_topk: int
+    num_experts: int
+    num_shared_expert: int
+    use_mixed_mlp_moe: bool
+    use_qk_norm: bool
+    rms_norm_eps: float
+    rope_theta: float
+    use_cla: bool
+    cla_share_factor: 2
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    tie_word_embeddings: bool = False
+
+    def __post_init__(self):
+
+        if self.rope_scaling:
+            required_keys = {"factor", "type"}
+            if not all(key in self.rope_scaling for key in required_keys):
+                raise ValueError(f"rope_scaling must contain keys {required_keys}")
+
+
+class DynamicNTKAlphaRoPE(nn.Module):
+    def __init__(
+        self,
+        dims: int,
+        base: float = 10000,
+        scaling_alpha: float = 1.0,
+    ):
+        super().__init__()
+        self.dims = dims
+        base = base * scaling_alpha ** (dims / (dims - 2))
+        self._freqs = base ** (mx.arange(0, self.dims, 2) / self.dims)
+
+    def __call__(self, x, offset: int = 0):
+        return mx.fast.rope(
+            x,
+            self.dims,
+            traditional=False,
+            base=None,
+            scale=1.0,
+            offset=offset,
+            freqs=self._freqs,
+        )
+
+
+class Attention(nn.Module):
+    def __init__(self, kv_proj: bool, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        assert args.num_key_value_heads is not None
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+
+        head_dim = args.hidden_size // n_heads
+        self.scale = head_dim**-0.5
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=args.attention_bias)
+        if kv_proj:
+            self.k_proj = nn.Linear(
+                dim, n_kv_heads * head_dim, bias=args.attention_bias
+            )
+            self.v_proj = nn.Linear(
+                dim, n_kv_heads * head_dim, bias=args.attention_bias
+            )
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=args.attention_bias)
+        self.use_qk_norm = args.use_qk_norm
+        if self.use_qk_norm:
+            self.query_layernorm = nn.RMSNorm(head_dim, args.rms_norm_eps)
+            self.key_layernorm = nn.RMSNorm(head_dim, args.rms_norm_eps)
+
+        self.rope = DynamicNTKAlphaRoPE(
+            head_dim,
+            base=args.rope_theta,
+            scaling_alpha=args.rope_scaling["alpha"],
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+        kv_states=None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries = self.q_proj(x)
+
+        if kv_states is None:
+            keys, values = self.k_proj(x), self.v_proj(x)
+            kv_states = keys, values
+        else:
+            keys, values = kv_states
+
+        # Prepare the queries, keys and values for the attention computation
+        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)
+
+        offset = cache.offset if cache else 0
+        queries = self.rope(queries, offset=offset)
+        keys = self.rope(keys, offset=offset)
+        if self.use_qk_norm:
+            queries = self.query_layernorm(queries)
+            keys = self.key_layernorm(keys)
+
+        if cache is not None:
+            keys, values = cache.update_and_fetch(keys, values)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output), kv_states
+
+
+class MLP(nn.Module):
+    def __init__(self, dim, hidden_dim):
+        super().__init__()
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
+
+    def __call__(self, x) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class Gate(nn.Module):
+    def __init__(self, dim, num_experts):
+        super().__init__()
+        self.wg = nn.Linear(dim, num_experts, bias=False)
+
+    def __call__(self, x) -> mx.array:
+        return self.wg(x)
+
+
+class MoeBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        dim = args.hidden_size
+        intermediate_size = args.intermediate_size
+        self.use_shared_mlp = args.use_mixed_mlp_moe
+
+        if args.use_mixed_mlp_moe:
+            self.shared_mlp = MLP(dim, intermediate_size * args.num_shared_expert)
+
+        self.num_experts = num_experts = args.num_experts
+        self.top_k = args.moe_topk
+
+        self.gate = Gate(dim, num_experts)
+        self.switch_mlp = SwitchGLU(dim, intermediate_size, num_experts)
+
+    def __call__(
+        self,
+        x: mx.array,
+    ):
+        gates = self.gate(x)
+        gates = mx.softmax(gates, axis=-1, precise=True)
+
+        k = self.top_k
+        inds = mx.stop_gradient(mx.argpartition(-gates, kth=k - 1, axis=-1)[..., :k])
+        scores = mx.take_along_axis(gates, inds, axis=-1)
+
+        y = self.switch_mlp(x, inds)
+        y = (y * scores[..., None]).sum(axis=-2)
+
+        if self.use_shared_mlp:
+            shared_expert_output = self.shared_mlp(x)
+            y = y + shared_expert_output
+
+        return y
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, kv_proj: bool):
+        super().__init__()
+        self.hidden_size = args.hidden_size
+        self.self_attn = Attention(kv_proj, args)
+        self.mlp = MoeBlock(args)
+
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+        self.args = args
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+        shared_kv_states: Optional[Tuple[mx.array, mx.array]] = None,
+    ):
+        r, shared_kv_states = self.self_attn(
+            self.input_layernorm(x), mask, cache, shared_kv_states
+        )
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        out = h + r
+        return out, shared_kv_states
+
+
+class HunYuanModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.vocab_size = args.vocab_size
+        self.num_hidden_layers = args.num_hidden_layers
+        assert self.vocab_size > 0
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            DecoderLayer(args=args, kv_proj=(i % args.cla_share_factor) == 0)
+            for i in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+    ):
+        h = self.embed_tokens(inputs)
+
+        mask = create_attention_mask(h, cache)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for i, (layer, c) in enumerate(zip(self.layers, cache)):
+            if i % self.args.cla_share_factor == 0:
+                shared_kv_states = None
+            h, shared_kv_states = layer(h, mask, c, shared_kv_states)
+
+        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 = HunYuanModel(args)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+    ):
+        out = self.model(inputs, cache)
+        return self.model.embed_tokens.as_linear(out)
+
+    def sanitize(self, weights):
+        if "model.layers.0.mlp.experts.0.up_proj.weight" not in weights:
+            return weights
+        for l in range(self.args.num_hidden_layers):
+            prefix = f"model.layers.{l}"
+            for n in ["up_proj", "down_proj", "gate_proj"]:
+                for k in ["weight", "scales", "biases"]:
+                    if f"{prefix}.mlp.experts.0.{n}.{k}" in weights:
+                        to_join = [
+                            weights.pop(f"{prefix}.mlp.experts.{e}.{n}.{k}")
+                            for e in range(self.args.num_experts)
+                        ]
+                        weights[f"{prefix}.mlp.switch_mlp.{n}.{k}"] = mx.stack(to_join)
+        return weights
+
+    @property
+    def layers(self):
+        return self.model.layers

llms/mlx_lm/models/internlm2.py

@@ -6,7 +6,7 @@ from typing import Any, Dict, Optional, Tuple, Union
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs, create_attention_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -141,8 +141,8 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        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.wo(output)

llms/mlx_lm/models/llama.py

@@ -1,12 +1,13 @@
 # Copyright © 2023-2024 Apple Inc.
 
 from dataclasses import dataclass
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any, Dict, Optional, Union
 
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs, create_attention_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .rope_utils import initialize_rope
 
 
 @dataclass
@@ -32,117 +33,6 @@ class ModelArgs(BaseModelArgs):
         if self.num_key_value_heads is None:
             self.num_key_value_heads = self.num_attention_heads
 
-        if self.rope_scaling:
-            if not "factor" in self.rope_scaling:
-                raise ValueError(f"rope_scaling must contain 'factor'")
-            rope_type = self.rope_scaling.get("type") or self.rope_scaling.get(
-                "rope_type"
-            )
-            if rope_type is None:
-                raise ValueError(
-                    f"rope_scaling must contain either 'type' or 'rope_type'"
-                )
-            if rope_type not in ["linear", "dynamic", "llama3"]:
-                raise ValueError(
-                    "rope_scaling 'type' currently only supports 'linear', 'dynamic' or 'llama3'"
-                )
-
-
-class DynamicNTKScalingRoPE(nn.Module):
-    """Implements the rotary positional encoding with Dynamic NTK scaling and Llama 3 RoPE."""
-
-    def __init__(
-        self,
-        dims: int,
-        max_position_embeddings: int = 2048,
-        traditional: bool = False,
-        base: float = 10000,
-        scale: float = 1.0,
-        rope_type: str = "default",
-        rope_scaling: dict = None,
-    ):
-        super().__init__()
-        self.dims = dims
-        self.max_position_embeddings = max_position_embeddings
-        self.traditional = traditional
-        self.scale = scale
-        self.rope_type = rope_type
-        self.rope_scaling = rope_scaling
-        self.base = base
-        self.compute_freqs()
-
-    def compute_freqs(self):
-        if self.rope_type != "llama3":
-            self._freqs = None
-            return
-        factor = self.rope_scaling["factor"]
-        low_freq_factor = self.rope_scaling.get("low_freq_factor", 1.0)
-        high_freq_factor = self.rope_scaling.get("high_freq_factor", 4.0)
-        old_context_len = self.rope_scaling.get(
-            "original_max_position_embeddings",
-            8192,
-        )
-
-        low_freq_wavelen = old_context_len / low_freq_factor
-        high_freq_wavelen = old_context_len / high_freq_factor
-
-        freqs = self.base ** (mx.arange(0, self.dims, 2) / self.dims)
-        wavelens = 2 * mx.pi * freqs
-
-        freqs = mx.where(wavelens > low_freq_wavelen, freqs * factor, freqs)
-        is_medium_freq = (wavelens > high_freq_wavelen) & (wavelens < low_freq_wavelen)
-        smooth_factors = (old_context_len / wavelens - low_freq_factor) / (
-            high_freq_factor - low_freq_factor
-        )
-        smooth_freqs = freqs / ((1 - smooth_factors) / factor + smooth_factors)
-        self._freqs = mx.where(is_medium_freq, smooth_freqs, freqs)
-        self.base = None
-
-    def extra_repr(self):
-        return (
-            f"{self.dims}, traditional={self.traditional}, "
-            f"max_position_embeddings={self.max_position_embeddings}, "
-            f"scaling_factor={self.scale}, rope_type={self.rope_type}"
-        )
-
-    def __call__(self, x, offset: int = 0):
-        return mx.fast.rope(
-            x,
-            self.dims,
-            traditional=self.traditional,
-            base=self.base,
-            scale=self.scale,
-            offset=offset,
-            freqs=self._freqs,
-        )
-
-
-def initialize_rope(args: ModelArgs):
-    head_dim = args.head_dim or args.hidden_size // args.num_attention_heads
-
-    rope_scaling = args.rope_scaling
-    rope_type = "default"
-    rope_scale = 1.0
-
-    if rope_scaling is not None:
-        rope_type = (
-            rope_scaling.get("type") or rope_scaling.get("rope_type") or "default"
-        )
-        if rope_type == "linear":
-            rope_scale = 1 / rope_scaling["factor"]
-        elif rope_type == "llama3":
-            rope_scale = 1.0  # The scaling is handled internally for llama3
-
-    return DynamicNTKScalingRoPE(
-        dims=head_dim,
-        max_position_embeddings=args.max_position_embeddings,
-        traditional=args.rope_traditional,
-        base=args.rope_theta,
-        scale=rope_scale,
-        rope_type=rope_type,
-        rope_scaling=rope_scaling,
-    )
-
 
 class Attention(nn.Module):
     def __init__(self, args: ModelArgs):
@@ -165,7 +55,13 @@ class Attention(nn.Module):
         self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias)
         self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=attention_bias)
 
-        self.rope = initialize_rope(args)
+        self.rope = initialize_rope(
+            self.head_dim,
+            args.rope_theta,
+            args.rope_traditional,
+            args.rope_scaling,
+            args.max_position_embeddings,
+        )
 
     def __call__(
         self,
@@ -190,9 +86,10 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        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)
 

llms/mlx_lm/models/mamba.py

@@ -23,6 +23,8 @@ class ModelArgs(BaseModelArgs):
     use_conv_bias: bool
     time_step_rank: int
     tie_word_embeddings: bool = True
+    use_bcdt_rms: bool = False
+    mixer_rms_eps: float = 1e-6
 
     def __post_init__(self):
         if not hasattr(self, "hidden_size") and hasattr(self, "d_model"):
@@ -44,6 +46,8 @@ class ModelArgs(BaseModelArgs):
 
         if self.time_step_rank == "auto":
             self.time_step_rank = math.ceil(self.hidden_size / 16)
+        if self.model_type == "falcon_mamba":
+            self.use_bcdt_rms = True
 
 
 class DepthWiseConv1d(nn.Module):
@@ -83,6 +87,11 @@ class MambaBlock(nn.Module):
         self.intermediate_size = args.intermediate_size
         self.time_step_rank = int(args.time_step_rank)
         self.use_conv_bias = args.use_conv_bias
+        self.use_bcdt_rms = args.use_bcdt_rms
+        if self.use_bcdt_rms:
+            self.mixer_norm = lambda x: mx.fast.rms_norm(
+                x, mx.ones(x.shape[-1], x.dtype), eps=args.mixer_rms_eps
+            )
 
         self.in_proj = nn.Linear(
             self.hidden_size, self.intermediate_size * 2, bias=args.use_bias
@@ -126,6 +135,8 @@ class MambaBlock(nn.Module):
             ],
             axis=-1,
         )
+        if self.use_bcdt_rms:
+            delta, B, C = map(self.mixer_norm, (delta, B, C))
         delta = nn.softplus(self.dt_proj(delta))
         new_state = mx.expand_dims(delta * x, -1) * mx.expand_dims(B, 1)
         if state is not None:

llms/mlx_lm/models/minicpm.py

@@ -7,7 +7,7 @@ import mlx.core as mx
 import mlx.nn as nn
 import numpy as np
 
-from .base import BaseModelArgs, create_attention_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -105,8 +105,8 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        attn_output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        attn_output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
         )
 
         attn_output = attn_output.transpose(0, 2, 1, 3).reshape(B, L, -1)

llms/mlx_lm/models/mixtral.py

@@ -7,7 +7,7 @@ from typing import Any, Dict, Optional, Tuple, Union
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs, create_attention_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .switch_layers import SwitchGLU
 
 
@@ -87,8 +87,8 @@ class MixtralAttention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        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)

llms/mlx_lm/models/nemotron.py

@@ -7,7 +7,7 @@ from typing import Any, Dict, Optional, Union
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs, create_attention_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -113,8 +113,8 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        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)

llms/mlx_lm/models/olmo2.py

@@ -0,0 +1,209 @@
+# Copyright © 2023-2024 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Union
+
+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
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    rms_norm_eps: float
+    vocab_size: int
+    head_dim: Optional[int] = None
+    max_position_embeddings: Optional[int] = None
+    num_key_value_heads: Optional[int] = None
+    attention_bias: bool = False
+    mlp_bias: bool = False
+    rope_theta: float = 10000
+    rope_traditional: bool = False
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    tie_word_embeddings: bool = True
+
+    def __post_init__(self):
+        if self.num_key_value_heads is None:
+            self.num_key_value_heads = self.num_attention_heads
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+
+        self.head_dim = head_dim = args.head_dim or args.hidden_size // n_heads
+
+        self.scale = head_dim**-0.5
+        if hasattr(args, "attention_bias"):
+            attention_bias = args.attention_bias
+        else:
+            attention_bias = False
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=attention_bias)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=attention_bias)
+
+        self.rope = initialize_rope(
+            self.head_dim,
+            args.rope_theta,
+            args.rope_traditional,
+            args.rope_scaling,
+            args.max_position_embeddings,
+        )
+
+        self.q_norm = nn.RMSNorm(n_heads * head_dim, args.rms_norm_eps)
+        self.k_norm = nn.RMSNorm(n_kv_heads * head_dim, args.rms_norm_eps)
+
+    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)
+        keys = self.k_norm(keys)
+
+        # Prepare the queries, keys and values for the attention computation
+        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
+        )
+
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class MLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        hidden_dim = args.intermediate_size
+        if hasattr(args, "mlp_bias"):
+            mlp_bias = args.mlp_bias
+        else:
+            mlp_bias = False
+
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=mlp_bias)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=mlp_bias)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=mlp_bias)
+
+    def __call__(self, x) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_attention_heads = args.num_attention_heads
+        self.hidden_size = args.hidden_size
+        self.self_attn = Attention(args)
+        self.mlp = MLP(args)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+        self.post_feedforward_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+        self.args = args
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.post_attention_layernorm(self.self_attn(x, mask, cache))
+        h = x + r
+        r = self.post_feedforward_layernorm(self.mlp(h))
+        out = h + r
+        return out
+
+
+class LlamaModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.vocab_size = args.vocab_size
+        self.num_hidden_layers = args.num_hidden_layers
+        assert self.vocab_size > 0
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            TransformerBlock(args=args) for _ in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+    ):
+        h = self.embed_tokens(inputs)
+
+        mask = create_attention_mask(h, cache)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, cache=c)
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = LlamaModel(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+    ):
+        out = self.model(inputs, cache)
+        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):
+        # Remove unused precomputed rotary freqs
+        return {
+            k: v for k, v in weights.items() if "self_attn.rotary_emb.inv_freq" not in k
+        }
+
+    @property
+    def layers(self):
+        return self.model.layers

llms/mlx_lm/models/openelm.py

@@ -6,7 +6,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs, create_attention_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -107,8 +107,8 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        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)

llms/mlx_lm/models/phi.py

@@ -7,7 +7,7 @@ from typing import Tuple
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs, create_attention_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -93,8 +93,13 @@ class PhiAttention(nn.Module):
             keys = self.rope(keys)
 
         scale = math.sqrt(1 / queries.shape[-1])
-        output = mx.fast.scaled_dot_product_attention(
-            queries.astype(mx.float32), keys, values, scale=scale, mask=mask
+        output = scaled_dot_product_attention(
+            queries.astype(mx.float32),
+            keys,
+            values,
+            cache=cache,
+            scale=scale,
+            mask=mask,
         ).astype(values.dtype)
 
         output = output.moveaxis(2, 1).reshape(B, L, -1)

llms/mlx_lm/models/phi3.py

@@ -6,7 +6,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs, create_attention_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .su_rope import SuScaledRotaryEmbedding
 
 
@@ -107,8 +107,8 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        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)

llms/mlx_lm/models/phi3small.py

@@ -8,7 +8,7 @@ from typing import Any, Optional
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs, create_attention_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -188,8 +188,8 @@ class Attention(nn.Module):
                 queries, keys, values, scale=self.scale, mask=mask
             )
         else:
-            output = mx.fast.scaled_dot_product_attention(
-                queries, keys, values, scale=self.scale, mask=mask
+            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.dense(output)

llms/mlx_lm/models/phimoe.py

@@ -6,7 +6,7 @@ from typing import Dict, List, Optional, Union
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs, create_attention_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .su_rope import SuScaledRotaryEmbedding
 from .switch_layers import SwitchGLU
 
@@ -79,8 +79,8 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        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)

llms/mlx_lm/models/phixtral.py

@@ -8,7 +8,7 @@ from typing import Tuple
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import create_attention_mask
+from .base import create_attention_mask, scaled_dot_product_attention
 from .switch_layers import SwitchMLP
 
 
@@ -71,8 +71,13 @@ class RoPEAttention(nn.Module):
         # Finally perform the attention computation
         scale = math.sqrt(1 / queries.shape[-1])
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries.astype(mx.float32), keys, values, scale=scale, mask=mask
+        output = scaled_dot_product_attention(
+            queries.astype(mx.float32),
+            keys,
+            values,
+            cache=cache,
+            scale=scale,
+            mask=mask,
         ).astype(values.dtype)
         output = output.moveaxis(2, 1).reshape(B, L, -1)
 

llms/mlx_lm/models/plamo.py

@@ -7,7 +7,7 @@ import mlx.core as mx
 import mlx.nn as nn
 import numpy as np
 
-from .base import BaseModelArgs, create_attention_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -92,10 +92,11 @@ class Attention(nn.Module):
         keys = mx.tile(keys, [1, self.config.n_shared_head, 1, 1])
         values = mx.tile(values, [1, self.config.n_shared_head, 1, 1])
 
-        output = mx.fast.scaled_dot_product_attention(
+        output = scaled_dot_product_attention(
             queries,
             keys,
             values,
+            cache=cache,
             scale=self.scale,
             mask=attention_mask,
         )

llms/mlx_lm/models/qwen.py

@@ -5,7 +5,7 @@ from dataclasses import dataclass
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs, create_attention_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -64,8 +64,8 @@ class Attention(nn.Module):
             queries = self.rotary_emb(queries)
             keys = self.rotary_emb(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        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)
 

llms/mlx_lm/models/qwen2.py

@@ -6,7 +6,7 @@ from typing import Any, Dict, Optional, Union
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs, create_attention_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -89,8 +89,8 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        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)

llms/mlx_lm/models/qwen2_moe.py

@@ -7,7 +7,7 @@ from typing import Any, Dict, Optional, Union
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs, create_attention_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .switch_layers import SwitchGLU
 
 
@@ -89,8 +89,8 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        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)

llms/mlx_lm/models/recurrent_gemma.py

@@ -7,7 +7,7 @@ from typing import List, Literal, Optional
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs, create_attention_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .cache import MambaCache, RotatingKVCache
 
 
@@ -263,8 +263,8 @@ class LocalAttentionBlock(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        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)

llms/mlx_lm/models/rope_utils.py

@@ -0,0 +1,91 @@
+# Copyright © 2023-2024 Apple Inc.
+
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class Llama3RoPE(nn.Module):
+
+    def __init__(
+        self,
+        dims: int,
+        max_position_embeddings: int = 2048,
+        traditional: bool = False,
+        base: float = 10000,
+        scaling_config: dict = None,
+    ):
+        super().__init__()
+        self.dims = dims
+        self.max_position_embeddings = max_position_embeddings
+        self.traditional = traditional
+
+        factor = scaling_config["factor"]
+        low_freq_factor = scaling_config.get("low_freq_factor", 1.0)
+        high_freq_factor = scaling_config.get("high_freq_factor", 4.0)
+        old_context_len = scaling_config.get(
+            "original_max_position_embeddings",
+            8192,
+        )
+
+        low_freq_wavelen = old_context_len / low_freq_factor
+        high_freq_wavelen = old_context_len / high_freq_factor
+
+        freqs = base ** (mx.arange(0, dims, 2) / dims)
+        wavelens = 2 * mx.pi * freqs
+
+        freqs = mx.where(wavelens > low_freq_wavelen, freqs * factor, freqs)
+        is_medium_freq = (wavelens > high_freq_wavelen) & (wavelens < low_freq_wavelen)
+        smooth_factors = (old_context_len / wavelens - low_freq_factor) / (
+            high_freq_factor - low_freq_factor
+        )
+        smooth_freqs = freqs / ((1 - smooth_factors) / factor + smooth_factors)
+        self._freqs = mx.where(is_medium_freq, smooth_freqs, freqs)
+
+    def extra_repr(self):
+        return (
+            f"{self.dims}, traditional={self.traditional}, "
+            f"max_position_embeddings={self.max_position_embeddings}"
+        )
+
+    def __call__(self, x, offset: int = 0):
+        return mx.fast.rope(
+            x,
+            self.dims,
+            traditional=self.traditional,
+            base=None,
+            scale=1.0,
+            offset=offset,
+            freqs=self._freqs,
+        )
+
+
+def initialize_rope(
+    dims,
+    base,
+    traditional,
+    scaling_config: Optional[dict] = None,
+    max_position_embeddings: Optional[int] = None,
+):
+    if scaling_config is not None:
+        rope_type = scaling_config.get("type") or scaling_config.get(
+            "rope_type", "default"
+        )
+    else:
+        rope_type = "default"
+
+    if rope_type in ["default", "linear"]:
+        scale = 1 / scaling_config["factor"] if rope_type == "linear" else 1.0
+        return nn.RoPE(dims, traditional=traditional, base=base, scale=scale)
+
+    elif rope_type == "llama3":
+        return Llama3RoPE(
+            dims=dims,
+            max_position_embeddings=max_position_embeddings,
+            traditional=traditional,
+            base=base,
+            scaling_config=scaling_config,
+        )
+    else:
+        raise ValueError(f"Unsupported RoPE type {rope_type}")

llms/mlx_lm/models/stablelm.py

@@ -6,7 +6,7 @@ from dataclasses import dataclass
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs, create_attention_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -120,8 +120,8 @@ class Attention(nn.Module):
 
         # Finally perform the attention computation
         scale = math.sqrt(1 / queries.shape[-1])
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=scale, mask=mask
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=scale, mask=mask
         ).astype(values.dtype)
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
         return self.o_proj(output)

llms/mlx_lm/models/starcoder2.py

@@ -6,7 +6,7 @@ from typing import Any, Optional
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs, create_attention_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -64,8 +64,8 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        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)

llms/mlx_lm/requirements.txt

@@ -1,4 +1,4 @@
-mlx>=0.17.0
+mlx>=0.19.2
 numpy
 transformers[sentencepiece]>=4.39.3
 protobuf

llms/mlx_lm/sample_utils.py

@@ -1,13 +1,87 @@
 # Copyright © 2023-2024 Apple Inc.
 
+import math
 from functools import partial
+from typing import Callable, Dict, Optional
 
 import mlx.core as mx
 
 
+def make_sampler(
+    temp: float = 0.0,
+    top_p: float = 0.0,
+    min_p: float = 0.0,
+    min_tokens_to_keep: int = 1,
+) -> Callable[mx.array, mx.array]:
+    """
+    Make a sampler function for use with ``generate_step``.
+
+    Args:
+        temp (float): The temperature for sampling, if 0 the argmax is used.
+          Default: ``0``.
+        top_p (float, optional): Nulceus sampling, higher means model considers
+          more less likely words.
+        min_p (float, optional): The minimum value (scaled by the top token's
+          probability) that a token probability must have to be considered.
+        min_tokens_to_keep (int, optional): Minimum number of tokens that cannot
+          be filtered by min_p sampling.
+
+    Returns:
+        Callable[mx.array, mx.array]:
+            A sampler which takes log-probabilities and returns tokens.
+    """
+    if temp == 0:
+        return lambda x: mx.argmax(x, axis=-1)
+    elif top_p > 0 and top_p < 1.0:
+        return lambda x: top_p_sampling(x, top_p, temp)
+    elif min_p != 0.0:
+        return lambda x: min_p_sampling(x, min_p, min_tokens_to_keep, temp)
+    else:
+        return lambda x: categorical_sampling(x, temp)
+
+
+def make_logits_processors(
+    logit_bias: Optional[Dict[int, float]] = None,
+    repetition_penalty: Optional[float] = None,
+    repetition_context_size: Optional[int] = 20,
+):
+    """
+    Make logits processors for use with ``generate_step``.
+
+    Args:
+        repetition_penalty (float, optional): The penalty factor for repeating
+          tokens.
+        repetition_context_size (int, optional): The number of tokens to
+          consider for repetition penalty. Default: ``20``.
+        logit_bias (dictionary, optional): Additive logit bias.
+
+    Returns:
+        List[Callable[[mx.array, mx.array], mx.array]]:
+            A list of logits processors. Each processor in the list is a
+            callable which takes an array of tokens and an array of logits
+            and returns the updated logits.
+    """
+    logits_processors = []
+    if logit_bias:
+        indices = mx.array(list(logit_bias.keys()))
+        values = mx.array(list(logit_bias.values()))
+
+        def logit_bias_processor(_, logits):
+            logits[:, indices] += values
+            return logits
+
+        logits_processors.append(logit_bias_processor)
+
+    if repetition_penalty and repetition_penalty != 0.0:
+        logits_processors.append(
+            make_repetition_penalty(repetition_penalty, repetition_context_size)
+        )
+    return logits_processors
+
+
 @partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)
 def min_p_sampling(
-    logits: mx.array,
+    logprobs: mx.array,
     min_p: float,
     min_tokens_to_keep: int = 1,
     temperature=1.0,
@@ -20,7 +94,7 @@ def min_p_sampling(
     aggressive given a very high-probability token.
 
     Args:
-        logits: The logits from the model's output.
+        logprobs: A vector of log probabilities.
         min_p (float): Minimum token probability. Typical values are in the
             0.01-0.2 range, comparably selective as setting `top_p` in the
             0.99-0.8 range.
@@ -38,28 +112,27 @@ def min_p_sampling(
         )
     # reference implementation: https://github.com/huggingface/transformers/blob/main/src/transformers/generation/logits_process.py#L531-L605
 
-    # Softmax probabilities
-    probs = mx.softmax(logits * (1 / temperature), axis=-1)
+    logprobs = logprobs * (1 / temperature)
 
     # Indices sorted in decreasing order
-    sorted_indices = mx.argsort(-logits).squeeze(0)
-    sorted_probs = probs[..., sorted_indices]
+    sorted_indices = mx.argsort(-logprobs).squeeze(0)
+    sorted_logprobs = logprobs[..., sorted_indices]
 
     # Top probability
-    top_probs = probs[..., sorted_indices[0]]
+    top_logprobs = logprobs[..., sorted_indices[0]]
 
     # Calculate the min_p threshold
-    scaled_min_p = min_p * top_probs
+    scaled_min_p = top_logprobs + math.log(min_p)
 
     # Mask tokens that have a probability less than the scaled min_p
-    tokens_to_remove = sorted_probs < scaled_min_p
+    tokens_to_remove = sorted_logprobs < scaled_min_p
     tokens_to_remove[..., :min_tokens_to_keep] = False
 
     # Create pool of tokens with probability less than scaled min_p
-    selected_probs = mx.where(tokens_to_remove, 0, sorted_probs)
+    selected_logprobs = mx.where(tokens_to_remove, -float("inf"), sorted_logprobs)
 
     # Return sampled token
-    sorted_token = mx.random.categorical(mx.log(selected_probs))
+    sorted_token = mx.random.categorical(selected_logprobs)
     return sorted_indices[sorted_token]
 
 
@@ -100,3 +173,36 @@ def top_p_sampling(logits: mx.array, top_p: float, temperature: float) -> mx.arr
 @partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)
 def categorical_sampling(logits, temp):
     return mx.random.categorical(logits * (1 / temp))
+
+
+def make_repetition_penalty(penalty: float, context_size: int = 20):
+    """
+    Make repetition penalty processor.
+
+    Paper: https://arxiv.org/abs/1909.05858
+
+    Args:
+        penalty (float): The repetition penalty factor to be applied.
+        context_size (int): The number of previous tokens to use.
+            Default: ``20``.
+
+    Returns:
+        Callable[[mx.array, List[int]], mx.array]:
+            The repetition penalty processor.
+    """
+    if penalty < 0 or not isinstance(penalty, (int, float)):
+        raise ValueError(f"penalty must be a non-negative float, got {penalty}")
+
+    def repetition_penalty_processor(tokens, logits):
+        if len(tokens) > 0:
+            tokens = tokens[-context_size:]
+            selected_logits = logits[:, tokens]
+            selected_logits = mx.where(
+                selected_logits < 0,
+                selected_logits * penalty,
+                selected_logits / penalty,
+            )
+            logits[:, tokens] = selected_logits
+        return logits
+
+    return repetition_penalty_processor

llms/mlx_lm/server.py

@@ -27,7 +27,8 @@ from huggingface_hub import scan_cache_dir
 
 from ._version import __version__
 from .models.cache import make_prompt_cache
-from .utils import generate_step, load
+from .sample_utils import make_logits_processors, make_sampler
+from .utils import load, stream_generate
 
 
 def get_system_fingerprint():
@@ -64,7 +65,7 @@ def stopping_criteria(
           end if it has (`trim_length`).
     """
     if tokens and tokens[-1] == eos_token_id:
-        return StopCondition(stop_met=True, trim_length=1)
+        return StopCondition(stop_met=True, trim_length=0)
 
     for stop_ids in stop_id_sequences:
         if len(tokens) >= len(stop_ids):
@@ -253,7 +254,7 @@ class APIHandler(BaseHTTPRequestHandler):
         self.max_tokens = self.body.get("max_completion_tokens", None)
         if self.max_tokens is None:
             self.max_tokens = self.body.get("max_tokens", 512)
-        self.temperature = self.body.get("temperature", 1.0)
+        self.temperature = self.body.get("temperature", 0.0)
         self.top_p = self.body.get("top_p", 1.0)
         self.repetition_penalty = self.body.get("repetition_penalty", 1.0)
         self.repetition_context_size = self.body.get("repetition_context_size", 20)
@@ -290,10 +291,7 @@ class APIHandler(BaseHTTPRequestHandler):
 
         # Call endpoint specific method
         prompt = endpoints[self.path]()
-
-        # Call method based on response type
-        method = self.handle_stream if self.stream else self.handle_completion
-        method(prompt, stop_id_sequences)
+        self.handle_completion(prompt, stop_id_sequences)
 
     def validate_model_parameters(self):
         """
@@ -452,32 +450,39 @@ class APIHandler(BaseHTTPRequestHandler):
             stop_id_sequences (List[List[int]]): A list of stop words passed
               to the stopping_criteria function
         """
-        detokenizer = self.tokenizer.detokenizer
-        detokenizer.reset()
         tokens = []
         finish_reason = "length"
         stop_sequence_suffix = None
-        logging.debug(f"Starting completion:")
+        if self.stream:
+            self.end_headers()
+            logging.debug(f"Starting stream:")
+        else:
+            logging.debug(f"Starting completion:")
         token_logprobs = []
         top_tokens = []
 
         prompt = self.get_prompt_cache(prompt)
 
-        for _, (token, logprobs) in zip(
-            range(self.max_tokens),
-            generate_step(
-                prompt=mx.array(prompt),
-                model=self.model,
-                temp=self.temperature,
-                top_p=self.top_p,
-                repetition_penalty=self.repetition_penalty,
-                repetition_context_size=self.repetition_context_size,
-                logit_bias=self.logit_bias,
-                prompt_cache=self.prompt_cache.cache,
-            ),
+        text = ""
+        tic = time.perf_counter()
+        sampler = make_sampler(self.temperature, top_p=self.top_p)
+        logits_processors = make_logits_processors(
+            self.logit_bias, self.repetition_penalty, self.repetition_context_size
+        )
+        for gen_response in stream_generate(
+            model=self.model,
+            tokenizer=self.tokenizer,
+            prompt=prompt,
+            max_tokens=self.max_tokens,
+            sampler=sampler,
+            logits_processors=logits_processors,
+            prompt_cache=self.prompt_cache.cache,
         ):
-            detokenizer.add_token(token)
-            logging.debug(detokenizer.text)
+            segment = gen_response.text
+            text += segment
+            logging.debug(text)
+            token = gen_response.token
+            logprobs = gen_response.logprobs
             tokens.append(token)
 
             if self.logprobs > 0:
@@ -498,121 +503,59 @@ class APIHandler(BaseHTTPRequestHandler):
                     stop_sequence_suffix = self.tokenizer.decode(
                         tokens[-stop_condition.trim_length :]
                     )
+                    text = text[: -len(stop_sequence_suffix)]
                 break
 
-        self.prompt_cache.tokens.extend(tokens)
-        detokenizer.finalize()
-        text = (
-            detokenizer.text
-            if stop_sequence_suffix is None
-            else detokenizer.text[: -len(stop_sequence_suffix)]
-        )
-        response = self.generate_response(
-            text,
-            finish_reason,
-            len(prompt),
-            len(tokens),
-            token_logprobs=token_logprobs,
-            top_tokens=top_tokens,
-            tokens=tokens,
-        )
-
-        response_json = json.dumps(response).encode()
-        indent = "\t"  # Backslashes can't be inside of f-strings
-        logging.debug(f"Outgoing Response: {json.dumps(response, indent=indent)}")
-
-        # Send an additional Content-Length header when it is known
-        self.send_header("Content-Length", str(len(response_json)))
-        self.end_headers()
-
-        self.wfile.write(response_json)
-        self.wfile.flush()
-
-    def handle_stream(
-        self,
-        prompt: List[int],
-        stop_id_sequences: List[List[int]],
-    ):
-        """
-        Generate response to prompt and foward it to the client using a Server
-        Sent Events (SSE) stream.
-
-        Args:
-            prompt (mx.array): The tokenized prompt
-            stop_id_sequences (List[List[int]]): A list of stop words passed to
-              the stopping_criteria function
-        """
-        # No additional headers are needed, call end_headers
-        self.end_headers()
-
-        detokenizer = self.tokenizer.detokenizer
-        detokenizer.reset()
-        tokens = []
-
-        stop_sequence_suffix = None
-        logging.debug(f"Starting stream:")
-
-        prompt = self.get_prompt_cache(prompt)
-
-        for _, (token, _) in zip(
-            range(self.max_tokens),
-            generate_step(
-                prompt=mx.array(prompt),
-                model=self.model,
-                temp=self.temperature,
-                top_p=self.top_p,
-                repetition_penalty=self.repetition_penalty,
-                repetition_context_size=self.repetition_context_size,
-                prompt_cache=self.prompt_cache.cache,
-            ),
-        ):
-            detokenizer.add_token(token)
-            logging.debug(detokenizer.text)
-            tokens.append(token)
-
-            stop_condition = stopping_criteria(
-                tokens,
-                stop_id_sequences,
-                self.tokenizer.eos_token_id,
-            )
-            if stop_condition.stop_met:
-                if stop_condition.trim_length:
-                    stop_sequence_suffix = self.tokenizer.decode(
-                        tokens[-stop_condition.trim_length :]
+            if self.stream:
+                # If the end of tokens overlaps with a stop sequence, generate new
+                # tokens until we know if the stop sequence is hit or not
+                if any(
+                    (
+                        sequence_overlap(tokens, sequence)
+                        for sequence in stop_id_sequences
                     )
-                break
-
-            # If the end of tokens overlaps with a stop sequence, generate new
-            # tokens until we know if the stop sequence is hit or not
-            if any(
-                (sequence_overlap(tokens, sequence) for sequence in stop_id_sequences)
-            ):
-                continue
-
-            new_text = detokenizer.last_segment
-            if new_text:
-                response = self.generate_response(new_text, None)
-                self.wfile.write(f"data: {json.dumps(response)}\n\n".encode())
-                self.wfile.flush()
+                ):
+                    continue
+                elif segment:
+                    response = self.generate_response(segment, None)
+                    self.wfile.write(f"data: {json.dumps(response)}\n\n".encode())
+                    self.wfile.flush()
 
         self.prompt_cache.tokens.extend(tokens)
 
-        # check is there any remaining text to send
-        detokenizer.finalize()
-        last_segment = detokenizer.last_segment
-        if last_segment:
-            if stop_sequence_suffix is not None:
-                last_segment = last_segment[: -len(stop_sequence_suffix)]
-            response = self.generate_response(last_segment, "length")
+        logging.debug(f"Prompt: {gen_response.prompt_tps:.3f} tokens-per-sec")
+        logging.debug(f"Generation: {gen_response.generation_tps:.3f} tokens-per-sec")
+        logging.debug(f"Peak memory: {gen_response.peak_memory:.3f} GB")
+
+        if self.stream:
+            response = self.generate_response(segment, finish_reason)
             self.wfile.write(f"data: {json.dumps(response)}\n\n".encode())
             self.wfile.flush()
+            if self.stream_options is not None and self.stream_options["include_usage"]:
+                response = self.completion_usage_response(len(prompt), len(tokens))
+                self.wfile.write(f"data: {json.dumps(response)}\n\n".encode())
+                self.wfile.flush()
+            self.wfile.write("data: [DONE]\n\n".encode())
+            self.wfile.flush()
+        else:
+            response = self.generate_response(
+                text,
+                finish_reason,
+                len(prompt),
+                len(tokens),
+                token_logprobs=token_logprobs,
+                top_tokens=top_tokens,
+                tokens=tokens,
+            )
+            response_json = json.dumps(response).encode()
+            indent = "\t"  # Backslashes can't be inside of f-strings
+            logging.debug(f"Outgoing Response: {json.dumps(response, indent=indent)}")
 
-        if self.stream_options is not None and self.stream_options["include_usage"]:
-            response = self.completion_usage_response(len(prompt), len(tokens))
-            self.wfile.write(f"data: {json.dumps(response)}\n\n".encode())
-
-        self.wfile.write("data: [DONE]\n\n".encode())
-        self.wfile.flush()
+            # Send an additional Content-Length header when it is known
+            self.send_header("Content-Length", str(len(response_json)))
+            self.end_headers()
+            self.wfile.write(response_json)
+            self.wfile.flush()
 
     def completion_usage_response(
         self,
@@ -646,9 +589,7 @@ class APIHandler(BaseHTTPRequestHandler):
 
         # Determine response type
         self.request_id = f"chatcmpl-{uuid.uuid4()}"
-        self.object_type = (
-            "chat.completions.chunk" if self.stream else "chat.completions"
-        )
+        self.object_type = "chat.completion.chunk" if self.stream else "chat.completion"
         if (
             hasattr(self.tokenizer, "apply_chat_template")
             and self.tokenizer.chat_template

llms/mlx_lm/tokenizer_utils.py

@@ -3,14 +3,6 @@ from functools import partial
 
 from transformers import AutoTokenizer
 
-REPLACEMENT_CHAR = "\ufffd"
-
-
-def _remove_space(x):
-    if x and x[0] == " ":
-        return x[1:]
-    return x
-
 
 class StreamingDetokenizer:
     """The streaming detokenizer interface so that we can detokenize one token at a time.
@@ -57,11 +49,9 @@ class StreamingDetokenizer:
     def last_segment(self):
         """Return the last segment of readable text since last time this property was accessed."""
         text = self.text
-        if text and text[-1] != REPLACEMENT_CHAR:
-            segment = text[self.offset :]
-            self.offset = len(text)
-            return segment
-        return ""
+        segment = text[self.offset :]
+        self.offset = len(text)
+        return segment
 
 
 class NaiveStreamingDetokenizer(StreamingDetokenizer):
@@ -79,16 +69,16 @@ class NaiveStreamingDetokenizer(StreamingDetokenizer):
 
     def reset(self):
         self.offset = 0
-        self._tokens = []
+        self.tokens = []
         self._text = ""
         self._current_tokens = []
         self._current_text = ""
 
     def add_token(self, token):
         self._current_tokens.append(token)
+        self.tokens.append(token)
 
     def finalize(self):
-        self._tokens.extend(self._current_tokens)
         self._text += self._tokenizer.decode(self._current_tokens)
         self._current_tokens = []
         self._current_text = ""
@@ -103,16 +93,11 @@ class NaiveStreamingDetokenizer(StreamingDetokenizer):
             ):
                 self._current_text = self._current_text[:-1]
         if self._current_text and self._current_text[-1] == "\n":
-            self._tokens.extend(self._current_tokens)
             self._text += self._current_text
             self._current_tokens.clear()
             self._current_text = ""
         return self._text + self._current_text
 
-    @property
-    def tokens(self):
-        return self._tokens
-
 
 class SPMStreamingDetokenizer(StreamingDetokenizer):
     """A streaming detokenizer for SPM models.
@@ -123,42 +108,43 @@ class SPMStreamingDetokenizer(StreamingDetokenizer):
 
     def __init__(self, tokenizer, trim_space=True):
         self.trim_space = trim_space
+        self._sep = "\u2581".encode()
 
         # Extract the tokens in a list from id to text
         self.tokenmap = [""] * (max(tokenizer.vocab.values()) + 1)
         for value, tokenid in tokenizer.vocab.items():
-            self.tokenmap[tokenid] = value
-
-        # Replace bytes with their value
-        for i in range(len(self.tokenmap)):
-            if self.tokenmap[i].startswith("<0x"):
-                self.tokenmap[i] = chr(int(self.tokenmap[i][3:5], 16))
+            if value.startswith("<0x"):
+                # Replace bytes with their value
+                self.tokenmap[tokenid] = bytes([int(value[3:5], 16)])
+            else:
+                self.tokenmap[tokenid] = value.encode()
 
         self.reset()
 
     def reset(self):
         self.offset = 0
-        self._unflushed = ""
+        self._unflushed = b""
         self.text = ""
         self.tokens = []
 
+    def _try_flush(self, force=False):
+        text = self._unflushed.replace(self._sep, b" ").decode("utf-8", "replace")
+        if not force and text.endswith("\ufffd"):
+            return
+        if not self.text and self.trim_space and text and text[0] == " ":
+            text = text[1:]
+        self.text += text
+        self._unflushed = b""
+
     def add_token(self, token):
+        self.tokens.append(token)
         v = self.tokenmap[token]
-        if v[0] == "\u2581":
-            if self.text or not self.trim_space:
-                self.text += self._unflushed.replace("\u2581", " ")
-            else:
-                self.text = _remove_space(self._unflushed.replace("\u2581", " "))
-            self._unflushed = v
-        else:
-            self._unflushed += v
+        self._unflushed += v
+        self._try_flush()
 
     def finalize(self):
-        if self.text or not self.trim_space:
-            self.text += self._unflushed.replace("\u2581", " ")
-        else:
-            self.text = _remove_space(self._unflushed.replace("\u2581", " "))
-        self._unflushed = ""
+        self._try_flush(force=True)
+        self._unflushed = b""
 
 
 class BPEStreamingDetokenizer(StreamingDetokenizer):
@@ -169,10 +155,9 @@ class BPEStreamingDetokenizer(StreamingDetokenizer):
     """
 
     _byte_decoder = None
-    _space_matches = (".", "?", "!", ",", "'", "n't", "'m", "'s", "'ve", "'re")
+    _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
@@ -192,6 +177,16 @@ class BPEStreamingDetokenizer(StreamingDetokenizer):
         self.text = ""
         self.tokens = []
 
+    def _decode_bytes(self, seq):
+        barr = bytearray()
+        for c in seq:
+            res = self._byte_decoder.get(c, False)
+            if res:
+                barr.append(res)
+            else:
+                barr.extend(bytes(c, "utf-8"))
+        return barr.decode("utf-8", "replace")
+
     def _maybe_trim_space(self, current_text):
         if len(current_text) == 0:
             return current_text
@@ -204,19 +199,23 @@ class BPEStreamingDetokenizer(StreamingDetokenizer):
         return current_text
 
     def add_token(self, token):
+        self.tokens.append(token)
         v = self.tokenmap[token]
-        if self._byte_decoder[v[0]] == 32:
-            current_text = bytearray(
-                self._byte_decoder[c] for c in self._unflushed
-            ).decode("utf-8")
-            self.text += self._maybe_trim_space(current_text)
-            self._unflushed = v
-        else:
-            self._unflushed += v
+        self._unflushed += v
+        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[v[0]] == 32
+        ):
+            self.text += self._maybe_trim_space(text)
+            self._unflushed = ""
 
     def finalize(self):
         current_text = bytearray(self._byte_decoder[c] for c in self._unflushed).decode(
-            "utf-8"
+            "utf-8",
+            "replace",
         )
         self.text += self._maybe_trim_space(current_text)
         self._unflushed = ""
@@ -256,21 +255,33 @@ class TokenizerWrapper:
     huggingface tokenizer.
     """
 
-    def __init__(self, tokenizer, detokenizer_class=NaiveStreamingDetokenizer):
+    def __init__(
+        self, tokenizer, detokenizer_class=NaiveStreamingDetokenizer, eos_token_ids=None
+    ):
         self._tokenizer = tokenizer
         self._detokenizer = detokenizer_class(tokenizer)
+        self._eos_token_ids = (
+            set(eos_token_ids)
+            if eos_token_ids is not None
+            else {tokenizer.eos_token_id}
+        )
 
     def __getattr__(self, attr):
         if attr == "detokenizer":
             return self._detokenizer
+        elif attr == "eos_token_ids":
+            return self._eos_token_ids
         elif attr.startswith("_"):
             return self.__getattribute__(attr)
         else:
             return getattr(self._tokenizer, attr)
 
     def __setattr__(self, attr, value):
-        if attr == "detokenizer":
-            raise AttributeError("Cannot set the detokenizer.")
+        if attr in {"detokenizer", "eos_token_ids"}:
+            if attr == "detokenizer":
+                raise AttributeError("Cannot set the detokenizer.")
+            elif attr == "eos_token_ids":
+                self._eos_token_ids = set(value) if value is not None else set()
         elif attr.startswith("_"):
             super().__setattr__(attr, value)
         else:
@@ -317,7 +328,7 @@ def _is_bpe_decoder(decoder):
     return isinstance(decoder, dict) and decoder.get("type", None) == "ByteLevel"
 
 
-def load_tokenizer(model_path, tokenizer_config_extra={}):
+def load_tokenizer(model_path, tokenizer_config_extra={}, eos_token_ids=None):
     """Load a huggingface tokenizer and try to infer the type of streaming
     detokenizer to use.
 
@@ -338,7 +349,10 @@ def load_tokenizer(model_path, tokenizer_config_extra={}):
             elif _is_bpe_decoder(tokenizer_content["decoder"]):
                 detokenizer_class = BPEStreamingDetokenizer
 
+    if isinstance(eos_token_ids, int):
+        eos_token_ids = [eos_token_ids]
     return TokenizerWrapper(
         AutoTokenizer.from_pretrained(model_path, **tokenizer_config_extra),
         detokenizer_class,
+        eos_token_ids=eos_token_ids,
     )

llms/mlx_lm/tuner/trainer.py

@@ -10,6 +10,7 @@ from typing import Union
 import mlx.core as mx
 import mlx.nn as nn
 import numpy as np
+from mlx.nn.utils import average_gradients
 from mlx.utils import tree_flatten
 
 
@@ -84,9 +85,16 @@ def iterate_batches(dataset, tokenizer, batch_size, max_seq_length, train=False)
             f" examples but only has {len(dataset)}."
         )
 
+    # If running in distributed mode (N machines) then each one should skip N-1
+    # samples
+    step = mx.distributed.init().size()
+    if batch_size % step != 0:
+        raise ValueError("The batch size must be divisible by the number of workers")
+
     # Make the batches:
     batch_idx = [
-        idx[i : i + batch_size] for i in range(0, len(idx) - batch_size + 1, batch_size)
+        idx[i : i + batch_size : step]
+        for i in range(0, len(idx) - batch_size + 1, batch_size)
     ]
 
     while True:
@@ -112,9 +120,9 @@ def iterate_batches(dataset, tokenizer, batch_size, max_seq_length, train=False)
             max_length_in_batch = pad_to * ((max(lengths) + pad_to - 1) // pad_to)
             max_length_in_batch = min(max_length_in_batch, max_seq_length)
 
-            batch_arr = np.zeros((batch_size, max_length_in_batch), np.int32)
+            batch_arr = np.zeros((batch_size // step, max_length_in_batch), np.int32)
 
-            for j in range(batch_size):
+            for j in range(batch_size // step):
                 truncated_length = min(lengths[j], max_seq_length)
                 batch_arr[j, :truncated_length] = batch[j][:truncated_length]
                 lengths[j] = (
@@ -138,7 +146,7 @@ def evaluate(
     loss: callable = default_loss,
     iterate_batches: callable = iterate_batches,
 ):
-    all_losses = []
+    all_losses = 0
     ntokens = 0
 
     index_iterator = iter(range(num_batches)) if num_batches != -1 else iter(int, 1)
@@ -153,10 +161,14 @@ def evaluate(
         ),
     ):
         losses, toks = loss(model, *batch)
-        all_losses.append((losses * toks).item())
-        ntokens += toks.item()
+        all_losses += losses * toks
+        ntokens += toks
+        mx.eval(all_losses, ntokens)
 
-    return np.sum(all_losses) / ntokens
+    all_losses = mx.distributed.all_sum(all_losses)
+    ntokens = mx.distributed.all_sum(ntokens)
+
+    return (all_losses / ntokens).item()
 
 
 class TrainingCallback:
@@ -182,6 +194,11 @@ def train(
     training_callback: TrainingCallback = None,
 ):
     print(f"Starting training..., iters: {args.iters}")
+    world = mx.distributed.init()
+    world_size = world.size()
+    rank = world.rank()
+    if world_size > 1:
+        print(f"Node {rank} of {world_size}")
 
     if args.grad_checkpoint:
         grad_checkpoint(model.layers[0])
@@ -192,6 +209,9 @@ def train(
         # Forward and backward pass
         (lvalue, toks), grad = loss_value_and_grad(model, *batch)
 
+        # All reduce the gradients if running in distributed mode
+        grad = average_gradients(grad)
+
         # Model update
         optimizer.update(model, grad)
 
@@ -199,8 +219,9 @@ def train(
 
     loss_value_and_grad = nn.value_and_grad(model, loss)
 
-    losses = []
+    losses = 0
     n_tokens = 0
+    steps = 0
     trained_tokens = 0
     # Main training loop
     start = time.perf_counter()
@@ -229,9 +250,13 @@ def train(
                 iterate_batches=iterate_batches,
             )
             val_time = time.perf_counter() - stop
-            print(
-                f"Iter {it}: " f"Val loss {val_loss:.3f}, " f"Val took {val_time:.3f}s"
-            )
+            if rank == 0:
+                print(
+                    f"Iter {it}: "
+                    f"Val loss {val_loss:.3f}, "
+                    f"Val took {val_time:.3f}s",
+                    flush=True,
+                )
 
             if training_callback is not None:
                 val_info = {
@@ -244,30 +269,33 @@ def train(
             start = time.perf_counter()
 
         lvalue, toks = step(batch)
-        mx.eval(state, lvalue, toks)
-
-        # Record loss
-        losses.append(lvalue.item())
-        n_tokens += toks.item()
+        losses += lvalue
+        n_tokens += toks
+        steps += 1
+        mx.eval(state, losses, n_tokens)
 
         # Report training loss if needed
         if it % args.steps_per_report == 0 or it == args.iters:
             stop = time.perf_counter()
 
-            train_loss = np.mean(losses)
+            train_loss = mx.distributed.all_sum(losses).item()
+            train_loss /= steps * mx.distributed.init().size()
+            n_tokens = mx.distributed.all_sum(n_tokens).item()
             learning_rate = optimizer.learning_rate.item()
             it_sec = args.steps_per_report / (stop - start)
             tokens_sec = float(n_tokens) / (stop - start)
             trained_tokens += n_tokens
-            peak_mem = mx.metal.get_peak_memory() / 2**30
-            print(
-                f"Iter {it}: Train loss {train_loss:.3f}, "
-                f"Learning Rate {learning_rate:.3e}, "
-                f"It/sec {it_sec:.3f}, "
-                f"Tokens/sec {tokens_sec:.3f}, "
-                f"Trained Tokens {trained_tokens}, "
-                f"Peak mem {peak_mem:.3f} GB"
-            )
+            peak_mem = mx.metal.get_peak_memory() / 1e9
+            if rank == 0:
+                print(
+                    f"Iter {it}: Train loss {train_loss:.3f}, "
+                    f"Learning Rate {learning_rate:.3e}, "
+                    f"It/sec {it_sec:.3f}, "
+                    f"Tokens/sec {tokens_sec:.3f}, "
+                    f"Trained Tokens {trained_tokens}, "
+                    f"Peak mem {peak_mem:.3f} GB",
+                    flush=True,
+                )
 
             if training_callback is not None:
                 train_info = {
@@ -281,8 +309,9 @@ def train(
                 }
                 training_callback.on_train_loss_report(train_info)
 
-            losses = []
+            losses = 0
             n_tokens = 0
+            steps = 0
             start = time.perf_counter()
 
         # Save adapter weights

llms/mlx_lm/tuner/utils.py

@@ -96,8 +96,10 @@ def linear_to_lora_layers(
         "gemma2",
         "starcoder2",
         "cohere",
+        "cohere2",
         "minicpm",
         "deepseek",
+        "olmo2",
     ]:
         keys = set(["self_attn.q_proj", "self_attn.v_proj"])
         if model.model_type in ["mixtral", "phimoe"]:
@@ -143,6 +145,8 @@ def linear_to_lora_layers(
                 "mixer.out_proj",
             ]
         )
+    elif model.model_type == "exaone":
+        keys = set(["attn.attention.q_proj", "attn.attention.v_proj"])
     else:
         raise ValueError(f"Lora does not support {model.model_type}")
 
@@ -249,12 +253,14 @@ def remove_lora_layers(model: nn.Module) -> nn.Module:
     return model
 
 
-def print_trainable_parameters(model):
-    def nparams(m):
-        if isinstance(m, (nn.QuantizedLinear, nn.QuantizedEmbedding)):
-            return m.weight.size * (32 // m.bits)
-        return sum(v.size for _, v in tree_flatten(m.parameters()))
+def nparams(module):
+    if hasattr(module, "bits"):
+        n = 0 if not hasattr(module, "bias") else module.bias.size
+        return n + module.weight.size * 32 // module.bits
+    return sum(v.size for _, v in tree_flatten(module.parameters()))
 
+
+def print_trainable_parameters(model):
     leaf_modules = tree_flatten(
         model.leaf_modules(), is_leaf=lambda m: isinstance(m, nn.Module)
     )

llms/mlx_lm/utils.py

@@ -1,5 +1,6 @@
 # Copyright © 2023-2024 Apple Inc.
 
+import contextlib
 import copy
 import glob
 import importlib
@@ -7,6 +8,7 @@ import json
 import logging
 import shutil
 import time
+from dataclasses import dataclass
 from pathlib import Path
 from textwrap import dedent
 from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Type, Union
@@ -14,24 +16,28 @@ from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Type,
 import mlx.core as mx
 import mlx.nn as nn
 from huggingface_hub import snapshot_download
-from mlx.utils import tree_flatten
+from mlx.utils import tree_flatten, tree_reduce
 from transformers import PreTrainedTokenizer
 
 # Local imports
-from .models import base, cache
-from .sample_utils import categorical_sampling, min_p_sampling, top_p_sampling
+from .models import cache
+from .sample_utils import make_logits_processors, make_sampler
 from .tokenizer_utils import TokenizerWrapper, load_tokenizer
 from .tuner.utils import dequantize as dequantize_model
-from .tuner.utils import load_adapters
+from .tuner.utils import load_adapters, nparams
 
 # Constants
 MODEL_REMAPPING = {
     "mistral": "llama",  # mistral is compatible with llama
     "phi-msft": "phixtral",
+    "falcon_mamba": "mamba",
 }
 
 MAX_FILE_SIZE_GB = 5
 
+# A stream on the default device just for generation
+generation_stream = mx.new_stream(mx.default_device())
+
 
 class ModelNotFoundError(Exception):
     def __init__(self, message):
@@ -39,6 +45,68 @@ class ModelNotFoundError(Exception):
         super().__init__(self.message)
 
 
+@dataclass
+class GenerationResponse:
+    """
+    The output of :func:`stream_generate`.
+
+    Args:
+        text (str): The next segment of decoded text. This can be an empty string.
+        token (int): The next token.
+        logprobs (mx.array): A vector of log probabilities.
+        prompt_tokens (int): The number of tokens in the prompt.
+        prompt_tps (float): The prompt processing tokens-per-second.
+        generation_tokens (int): The number of generated tokens.
+        generation_tps (float): The tokens-per-second for generation.
+        peak_memory (float): The peak memory used so far in GB.
+        finish_reason (str): The reason the response is being sent: "length", "stop" or `None`
+    """
+
+    text: str
+    token: int
+    logprobs: mx.array
+    prompt_tokens: int
+    prompt_tps: float
+    generation_tokens: int
+    generation_tps: float
+    peak_memory: float
+    finish_reason: Optional[str] = None
+
+
+@contextlib.contextmanager
+def wired_limit(model: nn.Module, streams: Optional[List[mx.Stream]] = None):
+    """
+    A context manager to temporarily change the wired limit.
+
+    Note, the wired limit should not be changed during an async eval.  If an
+    async eval could be running pass in the streams to synchronize with prior
+    to exiting the context manager.
+    """
+    model_bytes = tree_reduce(
+        lambda acc, x: acc + x.nbytes if isinstance(x, mx.array) else acc, model, 0
+    )
+    max_rec_size = mx.metal.device_info()["max_recommended_working_set_size"]
+    if model_bytes > 0.9 * max_rec_size:
+        model_mb = model_bytes // 2**20
+        max_rec_mb = max_rec_size // 2**20
+        print(
+            f"[WARNING] Generating with a model that requires {model_mb} MB "
+            f"which is close to the maximum recommended size of {max_rec_mb} "
+            "MB. This can be slow. See the documentation for possible work-arounds: "
+            "https://github.com/ml-explore/mlx-examples/tree/main/llms#large-models"
+        )
+    old_limit = mx.metal.set_wired_limit(max_rec_size)
+    try:
+        yield None
+    finally:
+        if streams is not None:
+            for s in streams:
+                mx.synchronize(s)
+        else:
+            mx.synchronize()
+        mx.metal.set_wired_limit(old_limit)
+
+
 def _get_classes(config: dict):
     """
     Retrieve the model and model args classes based on the configuration.
@@ -61,6 +129,17 @@ def _get_classes(config: dict):
     return arch.Model, arch.ModelArgs
 
 
+def compute_bits_per_weight(model):
+    model_bytes = tree_reduce(
+        lambda acc, x: acc + x.nbytes if isinstance(x, mx.array) else acc, model, 0
+    )
+    leaf_modules = tree_flatten(
+        model.leaf_modules(), is_leaf=lambda m: isinstance(m, nn.Module)
+    )
+    model_params = sum(nparams(m) for _, m in leaf_modules)
+    return model_bytes * 8 / model_params
+
+
 def get_model_path(path_or_hf_repo: str, revision: Optional[str] = None) -> Path:
     """
     Ensures the model is available locally. If the path does not exist locally,
@@ -101,43 +180,39 @@ def get_model_path(path_or_hf_repo: str, revision: Optional[str] = None) -> Path
     return model_path
 
 
-def apply_repetition_penalty(logits: mx.array, tokens: mx.array, penalty: float):
-    """
-    Apply repetition penalty to specific logits based on the given context.
-
-    Paper: https://arxiv.org/abs/1909.05858
-
-    Args:
-        logits (mx.array): The logits produced by the language model.
-        tokens (mx.array): A list of N previous tokens.
-        penalty (float): The repetition penalty factor to be applied.
-
-    Returns:
-        logits (mx.array): Logits with repetition penalty applied to generated tokens.
-    """
-    if len(tokens) > 0:
-        selected_logits = logits[:, tokens]
-        selected_logits = mx.where(
-            selected_logits < 0, selected_logits * penalty, selected_logits / penalty
-        )
-        logits[:, tokens] = selected_logits
-    return logits
+def maybe_quantize_kv_cache(prompt_cache, quantized_kv_start, kv_group_size, kv_bits):
+    if (
+        kv_bits is not None
+        and not isinstance(prompt_cache[0], cache.QuantizedKVCache)
+        and prompt_cache[0].offset > quantized_kv_start
+    ):
+        for i in range(len(prompt_cache)):
+            if isinstance(prompt_cache[i], cache.KVCache):
+                prompt_cache[i] = prompt_cache[i].to_quantized(
+                    group_size=kv_group_size, bits=kv_bits
+                )
 
 
 def generate_step(
     prompt: mx.array,
     model: nn.Module,
-    temp: float = 0.0,
-    repetition_penalty: Optional[float] = None,
-    repetition_context_size: Optional[int] = 20,
-    top_p: float = 1.0,
-    min_p: float = 0.0,
-    min_tokens_to_keep: int = 1,
-    prefill_step_size: int = 512,
+    *,
+    max_tokens: int = 256,
+    sampler: Optional[Callable[mx.array, mx.array]] = None,
+    logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = None,
     max_kv_size: Optional[int] = None,
     prompt_cache: Optional[Any] = None,
-    logit_bias: Optional[Dict[int, float]] = None,
-    logits_processor: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = None,
+    prefill_step_size: int = 512,
+    kv_bits: Optional[int] = None,
+    kv_group_size: int = 64,
+    quantized_kv_start: int = 0,
+    prompt_progress_callback: Optional[Callable[int, int]] = None,
+    temp: Optional[float] = None,
+    repetition_penalty: Optional[float] = None,
+    repetition_context_size: Optional[int] = None,
+    top_p: Optional[float] = None,
+    min_p: Optional[float] = None,
+    min_tokens_to_keep: Optional[int] = None,
 ) -> Generator[Tuple[mx.array, mx.array], None, None]:
     """
     A generator producing token ids based on the given prompt from the model.
@@ -145,166 +220,192 @@ def generate_step(
     Args:
         prompt (mx.array): The input prompt.
         model (nn.Module): The model to use for generation.
-        temp (float): The temperature for sampling, if 0 the argmax is used.
-          Default: ``0``.
-        repetition_penalty (float, optional): The penalty factor for repeating
-          tokens.
-        repetition_context_size (int, optional): The number of tokens to
-          consider for repetition penalty. Default: ``20``.
-        top_p (float, optional): Nulceus sampling, higher means model considers
-          more less likely words.
-        min_p (float, optional): The minimum value (scaled by the top token's
-          probability) that a token probability must have to be considered.
-        min_tokens_to_keep (int, optional): Minimum number of tokens that cannot
-          be filtered by min_p sampling.
-        prefill_step_size (int): Step size for processing the prompt.
+        max_tokens (int): The maximum number of tokens. Use``-1`` for an infinite
+          generator. Default: ``256``.
+        sampler (Callable[mx.array, mx.array], optional): A sampler for sampling a
+          token from a vector of log probabilities. Default: ``None``.
+        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``.
         max_kv_size (int, optional): Maximum size of the key-value cache. Old
           entries (except the first 4 tokens) will be overwritten.
         prompt_cache (List[Any], optional): A pre-computed prompt cache. Note, if
           provided, the cache will be updated in place.
-        logit_bias (dictionary, optional): Additive logit bias.
-        logits_processor (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``.
+        prefill_step_size (int): Step size for processing the prompt.
+        kv_bits (int, optional): Number of bits to use for KV cache quantization.
+          None implies no cache quantization. Default: ``None``.
+        kv_group_size (int): Group size for KV cache quantization. Default: ``64``.
+        quantized_kv_start (int): Step to begin using a quantized KV cache.
+           when ``kv_bits`` is non-None. Default: ``0``.
+        prompt_prorgress_callback (Callable[int, int]): A call-back which takes the
+           prompt tokens processed so far and the total number of prompt tokens.
 
     Yields:
-        Generator[Tuple[mx.array, mx.array], None, None]: A generator producing
-          one token and a vector of log probabilities.
+        Tuple[mx.array, mx.array]: One token and a vector of log probabilities.
     """
 
-    def sample(logits: mx.array) -> Tuple[mx.array, float]:
-        logprobs = logits - mx.logsumexp(logits)
-
-        if temp == 0:
-            token = mx.argmax(logits, axis=-1)
-        else:
-            if top_p > 0 and top_p < 1.0:
-                token = top_p_sampling(logits, top_p, temp)
-            elif min_p != 0.0:
-                token = min_p_sampling(logits, min_p, min_tokens_to_keep, temp)
-            else:
-                token = categorical_sampling(logits, temp)
-
-        return token, logprobs
-
-    if repetition_penalty and (
-        repetition_penalty < 0 or not isinstance(repetition_penalty, float)
-    ):
-        raise ValueError(
-            f"repetition_penalty must be a non-negative float, got {repetition_penalty}"
-        )
-
-    logits_processor = logits_processor or []
-
-    if repetition_penalty:
-
-        def repetition_penalty_processor(tokens, logits):
-            return apply_repetition_penalty(
-                logits, tokens[-repetition_context_size:], repetition_penalty
-            )
-
-        logits_processor.append(repetition_penalty_processor)
-
-    if logit_bias:
-        indices = mx.array(list(logit_bias.keys()))
-        values = mx.array(list(logit_bias.values()))
-
-        def logit_bias_processor(_, logits):
-            logits[:, indices] += values
-            return logits
-
-        logits_processor.append(logit_bias_processor)
-
     y = prompt
     tokens = None
 
     # Create the KV cache for generation
     if prompt_cache is None:
-        prompt_cache = cache.make_prompt_cache(model, max_kv_size)
+        prompt_cache = cache.make_prompt_cache(
+            model,
+            max_kv_size=max_kv_size,
+        )
     elif len(prompt_cache) != len(model.layers):
         raise ValueError("Wrong number of layers in the prompt cache.")
 
+    if temp is not None or top_p is not None or min_tokens_to_keep is not None:
+        print(
+            "[Warning] Specifying sampling arguments to ``generate_step`` is "
+            "deprecated. Pass in a ``sampler`` instead."
+        )
+    if repetition_penalty is not None:
+        print(
+            "[Warning] Specifying ``repetition_penalty`` is deprecated. "
+            "Pass in ``logits_processors`` instead."
+        )
+
+    sampler = sampler or make_sampler(
+        temp or 0.0, top_p or 0.0, min_p or 0.0, min_tokens_to_keep or 1
+    )
+    logits_processors = logits_processors or make_logits_processors(
+        None, repetition_penalty, repetition_context_size or 20
+    )
+    prompt_progress_callback = prompt_progress_callback or (lambda *_: None)
+
     def _step(y):
-        logits = model(y[None], cache=prompt_cache)
-        logits = logits[:, -1, :]
+        with mx.stream(generation_stream):
+            logits = model(y[None], cache=prompt_cache)
+            logits = logits[:, -1, :]
 
-        if logits_processor:
-            nonlocal tokens
-            tokens = mx.concat([tokens, y]) if tokens is not None else y
+            if logits_processors:
+                nonlocal tokens
+                tokens = mx.concat([tokens, y]) if tokens is not None else y
 
-            for processor in logits_processor:
-                logits = processor(tokens, logits)
+                for processor in logits_processors:
+                    logits = processor(tokens, logits)
 
-        y, logprobs = sample(logits)
-        return y, logprobs.squeeze(0)
+            maybe_quantize_kv_cache(
+                prompt_cache, quantized_kv_start, kv_group_size, kv_bits
+            )
 
-    while y.size > prefill_step_size:
-        model(y[:prefill_step_size][None], cache=prompt_cache)
-        mx.eval([c.state for c in prompt_cache])
-        y = y[prefill_step_size:]
-        mx.metal.clear_cache()
+            logprobs = logits - mx.logsumexp(logits, keepdims=True)
+            y = sampler(logprobs)
+            return y, logprobs.squeeze(0)
 
-    y, logprobs = _step(y)
+    with mx.stream(generation_stream):
+        total_prompt_tokens = y.size
+        prompt_processed_tokens = 0
+        while y.size > prefill_step_size:
+            model(y[:prefill_step_size][None], cache=prompt_cache)
+            maybe_quantize_kv_cache(
+                prompt_cache, quantized_kv_start, kv_group_size, kv_bits
+            )
+            mx.eval([c.state for c in prompt_cache])
+            prompt_progress_callback(prompt_processed_tokens, total_prompt_tokens)
+            prompt_processed_tokens += prefill_step_size
+            y = y[prefill_step_size:]
+            mx.metal.clear_cache()
+
+        y, logprobs = _step(y)
 
     mx.async_eval(y, logprobs)
+    n = 0
     while True:
-        next_y, next_logprobs = _step(y)
-        mx.async_eval(next_y, next_logprobs)
+        if n != max_tokens:
+            next_y, next_logprobs = _step(y)
+            mx.async_eval(next_y, next_logprobs)
+        if n == 0:
+            mx.eval(y)
+            prompt_progress_callback(total_prompt_tokens, total_prompt_tokens)
+        if n == max_tokens:
+            break
         yield y.item(), logprobs
+        if n % 256 == 0:
+            mx.metal.clear_cache()
         y, logprobs = next_y, next_logprobs
+        n += 1
 
 
 def stream_generate(
     model: nn.Module,
     tokenizer: Union[PreTrainedTokenizer, TokenizerWrapper],
-    prompt: str,
-    max_tokens: int = 100,
+    prompt: Union[str, mx.array, List[int]],
     **kwargs,
-) -> Union[str, Generator[str, None, None]]:
+) -> Generator[GenerationResponse, None, None]:
     """
     A generator producing text based on the given prompt from the model.
 
     Args:
-        prompt (mx.array): The input prompt.
         model (nn.Module): The model to use for generation.
-        max_tokens (int): The ma
+        tokenizer (PreTrainedTokenizer): The tokenizer.
+        prompt (Union[str, mx.array, List[int]]): The input prompt string or integer tokens.
         kwargs: The remaining options get passed to :func:`generate_step`.
           See :func:`generate_step` for more details.
 
     Yields:
-        Generator[Tuple[mx.array, mx.array]]: A generator producing text.
+        GenerationResponse: An instance containing the generated text segment and
+            associated metadata. See :class:`GenerationResponse` for details.
     """
     if not isinstance(tokenizer, TokenizerWrapper):
         tokenizer = TokenizerWrapper(tokenizer)
 
-    prompt_tokens = mx.array(tokenizer.encode(prompt))
+    if not isinstance(prompt, mx.array):
+        prompt = mx.array(
+            prompt if isinstance(prompt, list) else tokenizer.encode(prompt)
+        )
+
     detokenizer = tokenizer.detokenizer
 
-    detokenizer.reset()
-    for n, (token, _) in zip(
-        range(max_tokens),
-        generate_step(prompt_tokens, model, **kwargs),
-    ):
-        if token == tokenizer.eos_token_id:
-            break
-        detokenizer.add_token(token)
+    with wired_limit(model, [generation_stream]):
+        detokenizer.reset()
+        tic = time.perf_counter()
+        for n, (token, logprobs) in enumerate(generate_step(prompt, model, **kwargs)):
+            if n == 0:
+                prompt_time = time.perf_counter() - tic
+                prompt_tps = prompt.size / prompt_time
+                tic = time.perf_counter()
+            if token in tokenizer.eos_token_ids:
+                break
 
-        # Yield the last segment if streaming
-        yield detokenizer.last_segment
+            detokenizer.add_token(token)
 
-    detokenizer.finalize()
-    yield detokenizer.last_segment
+            yield GenerationResponse(
+                text=detokenizer.last_segment,
+                token=token,
+                logprobs=logprobs,
+                prompt_tokens=prompt.size,
+                prompt_tps=prompt_tps,
+                generation_tokens=n + 1,
+                generation_tps=(n + 1) / (time.perf_counter() - tic),
+                peak_memory=mx.metal.get_peak_memory() / 1e9,
+                finish_reason=None,
+            )
+
+        detokenizer.finalize()
+        yield GenerationResponse(
+            text=detokenizer.last_segment,
+            token=token,
+            logprobs=logprobs,
+            prompt_tokens=prompt.size,
+            prompt_tps=prompt_tps,
+            generation_tokens=n + 1,
+            generation_tps=(n + 1) / (time.perf_counter() - tic),
+            peak_memory=mx.metal.get_peak_memory() / 1e9,
+            finish_reason="stop" if token in tokenizer.eos_token_ids else "length",
+        )
 
 
 def generate(
     model: nn.Module,
     tokenizer: Union[PreTrainedTokenizer, TokenizerWrapper],
     prompt: str,
-    max_tokens: int = 100,
     verbose: bool = False,
     formatter: Optional[Callable] = None,
     **kwargs,
-) -> Union[str, Generator[str, None, None]]:
+) -> str:
     """
     Generate a complete response from the model.
 
@@ -312,66 +413,42 @@ def generate(
        model (nn.Module): The language model.
        tokenizer (PreTrainedTokenizer): The tokenizer.
        prompt (str): The string prompt.
-       max_tokens (int): The maximum number of tokens. Default: ``100``.
        verbose (bool): If ``True``, print tokens and timing information.
            Default: ``False``.
-       formatter (Optional[Callable]): A function which takes a token and a
-           probability and displays it.
-       kwargs: The remaining options get passed to :func:`generate_step`.
-          See :func:`generate_step` for more details.
+       kwargs: The remaining options get passed to :func:`stream_generate`.
+          See :func:`stream_generate` for more details.
     """
-    if not isinstance(tokenizer, TokenizerWrapper):
-        tokenizer = TokenizerWrapper(tokenizer)
-
+    if formatter is not None:
+        print(
+            "[Warning] Text formatting is deprecated and no longer used. "
+            "The argument will be removed in a future version."
+        )
     if verbose:
         print("=" * 10)
         print("Prompt:", prompt)
 
-    prompt_tokens = mx.array(tokenizer.encode(prompt))
-    detokenizer = tokenizer.detokenizer
-
-    tic = time.perf_counter()
-    detokenizer.reset()
-
-    for n, (token, logprobs) in zip(
-        range(max_tokens),
-        generate_step(prompt_tokens, model, **kwargs),
-    ):
-        if n == 0:
-            prompt_time = time.perf_counter() - tic
-            tic = time.perf_counter()
-        if token == tokenizer.eos_token_id:
-            break
-        detokenizer.add_token(token)
-
+    text = ""
+    for response in stream_generate(model, tokenizer, prompt, **kwargs):
         if verbose:
-            if formatter:
-                # We have to finalize so that the prob corresponds to the last segment
-                detokenizer.finalize()
-                with mx.stream(mx.cpu):
-                    prob = mx.exp(logprobs[token]).item()
-                formatter(detokenizer.last_segment, prob)
-            else:
-                print(detokenizer.last_segment, end="", flush=True)
-
-    token_count = n + 1
-    detokenizer.finalize()
+            print(response.text, end="", flush=True)
+        text += response.text
 
     if verbose:
-        gen_time = time.perf_counter() - tic
-        print(detokenizer.last_segment, flush=True)
+        print()
         print("=" * 10)
-        if token_count == 0:
-            print("No tokens generated for this prompt")
+        if len(text) == 0:
+            print("No text generated for this prompt")
             return
-        prompt_tps = prompt_tokens.size / prompt_time
-        gen_tps = (token_count - 1) / gen_time
-        print(f"Prompt: {prompt_tokens.size} tokens, {prompt_tps:.3f} tokens-per-sec")
-        print(f"Generation: {token_count} tokens, {gen_tps:.3f} tokens-per-sec")
-        peak_mem = mx.metal.get_peak_memory() / 2**30
-        print(f"Peak memory: {peak_mem:.3f} GB")
-
-    return detokenizer.text
+        print(
+            f"Prompt: {response.prompt_tokens} tokens, "
+            f"{response.prompt_tps:.3f} tokens-per-sec"
+        )
+        print(
+            f"Generation: {response.generation_tokens} tokens, "
+            f"{response.generation_tps:.3f} tokens-per-sec"
+        )
+        print(f"Peak memory: {response.peak_memory:.3f} GB")
+    return text
 
 
 def load_config(model_path: Path) -> dict:
@@ -398,11 +475,11 @@ def load_model(
         lazy (bool): If False eval the model parameters to make sure they are
             loaded in memory before returning, otherwise they will be loaded
             when needed. Default: ``False``
-        model_config (dict, optional): Configuration parameters for the model.
-            Defaults to an empty dictionary.
+        model_config (dict, optional): Optional configuration parameters for the
+            model. Defaults to an empty dictionary.
         get_model_classes (Callable[[dict], Tuple[Type[nn.Module], Type]], optional):
             A function that returns the model class and model args class given a config.
-            Defaults to the _get_classes function.
+            Defaults to the ``_get_classes`` function.
 
     Returns:
         nn.Module: The loaded and initialized model.
@@ -411,7 +488,6 @@ def load_model(
         FileNotFoundError: If the weight files (.safetensors) are not found.
         ValueError: If the model class or args class are not found or cannot be instantiated.
     """
-
     config = load_config(model_path)
     config.update(model_config)
 
@@ -438,15 +514,21 @@ def load_model(
         weights = model.sanitize(weights)
 
     if (quantization := config.get("quantization", None)) is not None:
-        # Handle legacy models which may not have everything quantized
+
         def class_predicate(p, m):
+            # Handle custom per layer quantizations
+            if p in config["quantization"]:
+                return config["quantization"][p]
             if not hasattr(m, "to_quantized"):
                 return False
+            # Handle legacy models which may not have everything quantized
             return f"{p}.scales" in weights
 
         nn.quantize(
             model,
-            **quantization,
+            group_size=quantization["group_size"],
+            bits=quantization["bits"],
+            quantization_type=quantization["quantization_type"],
             class_predicate=class_predicate,
         )
 
@@ -456,7 +538,7 @@ def load_model(
         mx.eval(model.parameters())
 
     model.eval()
-    return model
+    return model, config
 
 
 def load(
@@ -489,11 +571,13 @@ def load(
     """
     model_path = get_model_path(path_or_hf_repo)
 
-    model = load_model(model_path, lazy, model_config)
+    model, config = load_model(model_path, lazy)
     if adapter_path is not None:
         model = load_adapters(model, adapter_path)
         model.eval()
-    tokenizer = load_tokenizer(model_path, tokenizer_config)
+    tokenizer = load_tokenizer(
+        model_path, tokenizer_config, eos_token_ids=config.get("eos_token_id", None)
+    )
 
     return model, tokenizer
 
@@ -501,9 +585,10 @@ def load(
 def fetch_from_hub(
     model_path: Path, lazy: bool = False
 ) -> Tuple[nn.Module, dict, PreTrainedTokenizer]:
-    model = load_model(model_path, lazy)
-    config = load_config(model_path)
-    tokenizer = load_tokenizer(model_path)
+    model, config = load_model(model_path, lazy)
+    tokenizer = load_tokenizer(
+        model_path, eos_token_ids=config.get("eos_token_id", None)
+    )
     return model, config, tokenizer
 
 
@@ -553,7 +638,9 @@ def upload_to_hub(path: str, upload_repo: str, hf_path: str):
         f"""
         # {upload_repo}
 
-        The Model [{upload_repo}](https://huggingface.co/{upload_repo}) was converted to MLX format from [{hf_path}](https://huggingface.co/{hf_path}) using mlx-lm version **{__version__}**.
+        The Model [{upload_repo}](https://huggingface.co/{upload_repo}) was
+        converted to MLX format from [{hf_path}](https://huggingface.co/{hf_path})
+        using mlx-lm version **{__version__}**.
 
         ## Use with mlx
 
@@ -647,7 +734,14 @@ def save_weights(
 
 
 def quantize_model(
-    model: nn.Module, config: dict, q_group_size: int, q_bits: int
+    model: nn.Module,
+    config: dict,
+    q_group_size: int,
+    q_bits: int,
+    q_type: str,
+    quant_predicate: Optional[
+        Callable[[str, nn.Module, dict], Union[bool, dict]]
+    ] = None,
 ) -> Tuple:
     """
     Applies quantization to the model weights.
@@ -657,17 +751,53 @@ def quantize_model(
         config (dict): Model configuration.
         q_group_size (int): Group size for quantization.
         q_bits (int): Bits per weight for quantization.
+        q_type (str): Quantization type
+        quant_predicate (Callable): A callable that decides how
+            to quantize each layer based on the path.
+            Accepts the layer `path`, the `module` and the model `config`.
+            Returns either a bool to signify quantize/no quantize or
+            a dict of quantization parameters to pass to `to_quantized`.
 
     Returns:
         Tuple: Tuple containing quantized weights and config.
     """
     quantized_config = copy.deepcopy(config)
-    nn.quantize(model, q_group_size, q_bits)
-    quantized_config["quantization"] = {"group_size": q_group_size, "bits": q_bits}
+    quantized_config["quantization"] = {
+        "group_size": q_group_size,
+        "bits": q_bits,
+        "quantization_type": q_type,
+    }
+
+    # Add any custom quantization parameters to the config as we go
+    def _class_predicate(p, m):
+        if quant_predicate:
+            bool_or_params = quant_predicate(p, m, config)
+        else:
+            if isinstance(m, nn.Embedding):
+                bool_or_params = {
+                    "group_size": q_group_size,
+                    "bits": q_bits,
+                    "quantization_type": "affine",
+                }
+            else:
+                bool_or_params = hasattr(m, "to_quantized")
+        quantized_config["quantization"][p] = bool_or_params
+        return bool_or_params
+
+    nn.quantize(
+        model,
+        q_group_size,
+        q_bits,
+        quantization_type=q_type,
+        class_predicate=_class_predicate,
+    )
     # support hf model tree #957
     quantized_config["quantization_config"] = quantized_config["quantization"]
     quantized_weights = dict(tree_flatten(model.parameters()))
 
+    bpw = compute_bits_per_weight(model)
+    print(f"[INFO] Quantized model with {bpw:.3f} bits per weight.")
+
     return quantized_weights, quantized_config
 
 
@@ -700,10 +830,14 @@ def convert(
     quantize: bool = False,
     q_group_size: int = 64,
     q_bits: int = 4,
+    q_type: str = "affine",
     dtype: str = "float16",
     upload_repo: str = None,
     revision: Optional[str] = None,
     dequantize: bool = False,
+    quant_predicate: Optional[
+        Callable[[str, nn.Module, dict], Union[bool, dict]]
+    ] = None,
 ):
     # Check the save path is empty
     if isinstance(mlx_path, str):
@@ -729,7 +863,14 @@ def convert(
     if quantize:
         print("[INFO] Quantizing")
         model.load_weights(list(weights.items()))
-        weights, config = quantize_model(model, config, q_group_size, q_bits)
+        weights, config = quantize_model(
+            model,
+            config,
+            q_group_size,
+            q_bits,
+            q_type=q_type,
+            quant_predicate=quant_predicate,
+        )
 
     if dequantize:
         print("[INFO] Dequantizing")

llms/setup.py

@@ -28,12 +28,14 @@ setup(
     python_requires=">=3.8",
     extras_require={
         "testing": ["datasets"],
+        "evaluation": ["lm-eval"],
     },
     entry_points={
         "console_scripts": [
             "mlx_lm.cache_prompt = mlx_lm.cache_prompt:main",
             "mlx_lm.chat = mlx_lm.chat:main",
             "mlx_lm.convert = mlx_lm.convert:main",
+            "mlx_lm.evaluate = mlx_lm.evaluate:main",
             "mlx_lm.fuse = mlx_lm.fuse:main",
             "mlx_lm.generate = mlx_lm.generate:main",
             "mlx_lm.lora = mlx_lm.lora:main",

llms/tests/test_finetune.py

@@ -3,6 +3,7 @@
 import math
 import sys
 import unittest
+from contextlib import contextmanager
 from io import StringIO
 from unittest.mock import MagicMock
 
@@ -17,6 +18,14 @@ from mlx_lm.tuner.trainer import evaluate
 from mlx_lm.tuner.utils import build_schedule
 
 
+@contextmanager
+def swapped_with_identity(obj, func):
+    old_func = getattr(obj, func)
+    setattr(obj, func, lambda x: x)
+    yield
+    setattr(obj, func, old_func)
+
+
 class TestLora(unittest.TestCase):
     def setUp(self):
         self.capturedOutput = StringIO()
@@ -374,16 +383,17 @@ class TestScheduleConfig(unittest.TestCase):
             (MagicMock(return_value=0.4), MagicMock(return_value=180)),
             (MagicMock(return_value=0.6), MagicMock(return_value=120)),
         ]
-        evaluate(
-            model=mock_model,
-            dataset=mock_dataset,
-            tokenizer=mock_tokenizer,
-            batch_size=2,
-            num_batches=2,
-            max_seq_length=2048,
-            loss=mock_default_loss,
-            iterate_batches=mock_iterate_batches,
-        )
+        with swapped_with_identity(mx.distributed, "all_sum"):
+            evaluate(
+                model=mock_model,
+                dataset=mock_dataset,
+                tokenizer=mock_tokenizer,
+                batch_size=2,
+                num_batches=2,
+                max_seq_length=2048,
+                loss=mock_default_loss,
+                iterate_batches=mock_iterate_batches,
+            )
 
         mock_iterate_batches.assert_called_once_with(
             dataset=mock_dataset,
@@ -412,16 +422,17 @@ class TestScheduleConfig(unittest.TestCase):
             (MagicMock(return_value=0.2), MagicMock(return_value=150)),
         ]
 
-        evaluate(
-            model=mock_model,
-            dataset=mock_dataset,
-            tokenizer=mock_tokenizer,
-            batch_size=2,
-            num_batches=-1,
-            max_seq_length=2048,
-            loss=mock_default_loss,
-            iterate_batches=mock_iterate_batches,
-        )
+        with swapped_with_identity(mx.distributed, "all_sum"):
+            evaluate(
+                model=mock_model,
+                dataset=mock_dataset,
+                tokenizer=mock_tokenizer,
+                batch_size=2,
+                num_batches=-1,
+                max_seq_length=2048,
+                loss=mock_default_loss,
+                iterate_batches=mock_iterate_batches,
+            )
 
         mock_iterate_batches.assert_called_once_with(
             dataset=mock_dataset,

llms/tests/test_generate.py

@@ -2,6 +2,7 @@
 
 import unittest
 
+from mlx_lm.sample_utils import make_logits_processors
 from mlx_lm.utils import generate, load
 
 
@@ -25,8 +26,8 @@ class TestGenerate(unittest.TestCase):
             self.tokenizer,
             "hello",
             max_tokens=5,
+            logits_processors=make_logits_processors(logit_bias),
             verbose=False,
-            logit_bias=logit_bias,
         )
         self.assertEqual(text, "!!!!!")
 
@@ -46,7 +47,7 @@ class TestGenerate(unittest.TestCase):
             "hello",
             max_tokens=5,
             verbose=False,
-            logits_processor=[logits_processor],
+            logits_processors=[logits_processor],
         )
         self.assertEqual(len(all_toks), len(init_toks) + 5)
 

llms/tests/test_models.py

@@ -2,7 +2,9 @@
 import unittest
 
 import mlx.core as mx
+import mlx.nn as nn
 from mlx.utils import tree_map
+from mlx_lm.models import rope_utils
 from mlx_lm.models.cache import KVCache, RotatingKVCache, make_prompt_cache
 
 
@@ -126,6 +128,26 @@ class TestModels(unittest.TestCase):
         self.assertEqual(cache.offset, 22)
         self.assertTrue(mx.allclose(x, k[..., -2:, :]))
 
+    def test_rope(self):
+        rope = rope_utils.initialize_rope(32, base=100, traditional=False)
+        self.assertTrue(isinstance(rope, nn.RoPE))
+
+        rope = rope_utils.initialize_rope(
+            32,
+            base=100,
+            traditional=False,
+            scaling_config={"rope_type": "linear", "factor": 10.0},
+        )
+        self.assertTrue(isinstance(rope, nn.RoPE))
+
+        rope = rope_utils.initialize_rope(
+            32,
+            base=100,
+            traditional=False,
+            scaling_config={"rope_type": "llama3", "factor": 2.0},
+        )
+        self.assertTrue(isinstance(rope, rope_utils.Llama3RoPE))
+
     def model_test_runner(self, model, model_type, vocab_size, num_layers):
 
         self.assertEqual(len(model.layers), num_layers)
@@ -760,6 +782,91 @@ class TestModels(unittest.TestCase):
             model, args.model_type, args.vocab_size, args.num_hidden_layers
         )
 
+    def test_hunyuan(self):
+        from mlx_lm.models import hunyuan
+
+        args = hunyuan.ModelArgs(
+            model_type="hunyuan",
+            hidden_size=128,
+            attention_bias=False,
+            intermediate_size=256,
+            num_attention_heads=4,
+            num_hidden_layers=4,
+            num_key_value_heads=2,
+            rms_norm_eps=1e-4,
+            rope_theta=1000,
+            vocab_size=1000,
+            moe_topk=2,
+            num_experts=2,
+            num_shared_expert=1,
+            use_mixed_mlp_moe=True,
+            use_qk_norm=True,
+            rope_scaling={
+                "alpha": 1000.0,
+                "factor": 1.0,
+                "type": "dynamic",
+            },
+            use_cla=True,
+            cla_share_factor=2,
+        )
+        model = hunyuan.Model(args)
+        self.model_test_runner(
+            model, args.model_type, args.vocab_size, args.num_hidden_layers
+        )
+
+    def test_olmo2(self):
+        from mlx_lm.models import olmo2
+
+        args = olmo2.ModelArgs(
+            model_type="olmo2",
+            hidden_size=128,
+            attention_bias=False,
+            intermediate_size=256,
+            num_attention_heads=4,
+            num_hidden_layers=4,
+            num_key_value_heads=2,
+            rms_norm_eps=1e-4,
+            rope_theta=1000,
+            vocab_size=1000,
+        )
+        model = olmo2.Model(args)
+        self.model_test_runner(
+            model, args.model_type, args.vocab_size, args.num_hidden_layers
+        )
+
+    def test_exaone(self):
+        from mlx_lm.models import exaone
+
+        args = exaone.ModelArgs(
+            model_type="exaone",
+            hidden_size=128,
+            num_layers=4,
+            intermediate_size=256,
+            num_attention_heads=8,
+            num_key_value_heads=2,
+            vocab_size=1000,
+            layer_norm_epsilon=1e-4,
+            rope_theta=10000,
+        )
+        model = exaone.Model(args)
+        self.model_test_runner(model, args.model_type, args.vocab_size, args.num_layers)
+
+    def test_cohere2(self):
+        from mlx_lm.models import cohere2
+
+        args = cohere2.ModelArgs(
+            model_type="cohere2",
+            hidden_size=4096,
+            head_dim=128,
+            num_hidden_layers=40,
+            sliding_window=4096,
+            sliding_window_pattern=4,
+        )
+        model = cohere2.Model(args)
+        self.model_test_runner(
+            model, args.model_type, args.vocab_size, args.num_hidden_layers
+        )
+
 
 if __name__ == "__main__":
     unittest.main()

llms/tests/test_prompt_cache.py

@@ -9,6 +9,7 @@ import mlx.core as mx
 from mlx_lm.models.cache import (
     KVCache,
     MambaCache,
+    QuantizedKVCache,
     RotatingKVCache,
     load_prompt_cache,
     make_prompt_cache,
@@ -120,21 +121,20 @@ class TestPromptCache(unittest.TestCase):
     def test_cache_with_generate(self):
         model, tokenizer = load(HF_MODEL_PATH)
         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))
+        results = list(generate_step(prompt, model, max_tokens=4))
+        toks, all_logits = zip(*results)
 
         prompt_cache = make_prompt_cache(model)
         i = 0
-        for _, (tok, logits) in zip(
-            range(2), generate_step(prompt, model, prompt_cache=prompt_cache)
+        for tok, logits in generate_step(
+            prompt, model, prompt_cache=prompt_cache, max_tokens=2
         ):
             self.assertEqual(tok, toks[i])
             self.assertTrue(mx.allclose(logits, all_logits[i]))
             i += 1
 
-        for _, (tok, logits) in zip(
-            range(1),
-            generate_step(mx.array([toks[i]]), model, prompt_cache=prompt_cache),
+        for tok, logits in generate_step(
+            mx.array([toks[i]]), model, prompt_cache=prompt_cache, max_tokens=1
         ):
             i += 1
             self.assertEqual(tok, toks[i])
@@ -186,6 +186,18 @@ class TestPromptCache(unittest.TestCase):
         num_trimmed = trim_prompt_cache(cache, 4)
         self.assertEqual(num_trimmed, 0)
 
+        cache = [QuantizedKVCache() for _ in range(2)]
+        for c in cache:
+            x = mx.random.uniform(shape=(1, 8, 10, 64))
+            c.update_and_fetch(x, x)
+
+        num_trimmed = trim_prompt_cache(cache, 7)
+        self.assertEqual(num_trimmed, 7)
+
+        # Trim more tokens than remain
+        num_trimmed = trim_prompt_cache(cache, 4)
+        self.assertEqual(num_trimmed, 3)
+
     def test_trim_cache_with_generate(self):
         model, tokenizer = load(HF_MODEL_PATH)
         prompt = tokenizer.encode("this is a prompt", return_tensors="mlx")[0]
@@ -238,6 +250,56 @@ class TestPromptCache(unittest.TestCase):
         self.assertTrue(mx.allclose(old_cache[0].keys[..., 10:11, :], y))
         self.assertTrue(mx.allclose(cache[0].keys[..., 10:11, :], z))
 
+    def test_save_load_quantized_cache(self):
+        cache = [QuantizedKVCache(bits=4, group_size=32) for _ in range(4)]
+        for c in cache:
+            x = mx.random.uniform(shape=(1, 8, 10, 32))
+            c.update_and_fetch(x, x)
+        cache_file = os.path.join(self.test_dir, "prompt_cache.safetensors")
+        save_prompt_cache(cache_file, cache)
+        loaded_cache = load_prompt_cache(cache_file)
+        self.assertTrue(loaded_cache[0].bits == cache[0].bits)
+        self.assertTrue(loaded_cache[0].group_size == cache[0].group_size)
+        self.assertTrue(len(cache), len(loaded_cache))
+        for c, lc in zip(cache, loaded_cache):
+            self.assertEqual(c.offset, lc.offset)
+            # Loop over quantized tuple
+            for i in range(3):
+                self.assertTrue(mx.array_equal(c.state[0][i], lc.state[0][i]))
+                self.assertTrue(mx.array_equal(c.state[1][i], lc.state[1][i]))
+
+        # Test with metadata
+        cache_file = os.path.join(self.test_dir, "prompt_cache.safetensors")
+        metadata = {"a": "b", "c": "d"}
+        save_prompt_cache(cache_file, cache, metadata)
+        _, loaded_metadata = load_prompt_cache(cache_file, return_metadata=True)
+        self.assertEqual(metadata, loaded_metadata)
+
+    def test_cache_to_quantized(self):
+        model, tokenizer = load(HF_MODEL_PATH)
+        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))
+
+        prompt_cache = make_prompt_cache(model)
+        i = 0
+        for _, (tok, logits) in zip(
+            range(2), generate_step(prompt, model, prompt_cache=prompt_cache)
+        ):
+            self.assertEqual(tok, toks[i])
+            self.assertTrue(mx.allclose(logits, all_logits[i]))
+            i += 1
+
+        prompt_cache = [c.to_quantized(bits=8, group_size=32) for c in prompt_cache]
+
+        for _, (tok, logits) in zip(
+            range(1),
+            generate_step(mx.array([toks[i]]), model, prompt_cache=prompt_cache),
+        ):
+            i += 1
+            self.assertEqual(tok, toks[i])
+            self.assertTrue(mx.allclose(logits, all_logits[i], rtol=2e-2))
+
 
 if __name__ == "__main__":
     unittest.main()

llms/tests/test_sample_utils.py

@@ -1,10 +1,10 @@
 import unittest
 
 import mlx.core as mx
-from mlx_lm.sample_utils import top_p_sampling
+from mlx_lm.sample_utils import min_p_sampling, top_p_sampling
 
 
-class TestSamplingUtils(unittest.TestCase):
+class TestSampleUtils(unittest.TestCase):
     def test_top_p_sampling(self):
         probs = mx.array([0.9, 0.0, 0.0, 0.1])[None]
         logits = mx.log(probs)
@@ -28,6 +28,20 @@ class TestSamplingUtils(unittest.TestCase):
         token = top_p_sampling(logits, 0.95, temperature).item()
         self.assertTrue(token in (1, 2, 3))
 
+    def test_min_p_sampling(self):
+        probs = mx.array([0.9, 0.0, 0.0, 0.1])[None]
+        logits = mx.log(probs)
+        temperature = 1.0
+        token = min_p_sampling(logits, 0.8)
+        self.assertEqual(token, 0)
+
+        probs = mx.array([0.9, 0.0, 0.0, 0.1])[None]
+        logits = mx.log(probs)
+        temperature = 1.0
+        for _ in range(5):
+            token = min_p_sampling(logits, 0.05)
+            self.assertTrue(token in (0, 3))
+
 
 if __name__ == "__main__":
     unittest.main()

llms/tests/test_tokenizers.py

@@ -34,14 +34,18 @@ class TestTokenizers(unittest.TestCase):
             detokenizer = tokenizer.detokenizer
             detokenizer.reset()
             text = ""
-            for t in tokens:
+            for e, t in enumerate(tokens):
                 detokenizer.add_token(t)
                 seg = detokenizer.last_segment
                 text += seg
+                self.assertEqual(detokenizer.tokens, tokens[: e + 1])
             detokenizer.finalize()
             text += detokenizer.last_segment
             self.assertEqual(text, expected_text)
 
+        tokens = tokenizer.encode("こんにちは！私の名前はAI")
+        check(tokens)
+
         tokens = tokenizer.encode("a ,b")
         check(tokens)
 
@@ -51,6 +55,12 @@ class TestTokenizers(unittest.TestCase):
         tokens = tokenizer.encode("3 3")
         check(tokens)
 
+        tokens = tokenizer.encode("import 'package:flutter/material.dart';")
+        check(tokens)
+
+        tokens = tokenizer.encode("hello\nworld")
+        check(tokens)
+
     def test_tokenizers(self):
         tokenizer_repos = [
             ("mlx-community/Qwen1.5-0.5B-Chat-4bit", BPEStreamingDetokenizer),
@@ -58,6 +68,7 @@ class TestTokenizers(unittest.TestCase):
             ("mlx-community/Phi-3.5-mini-instruct-4bit", SPMStreamingDetokenizer),
             ("mlx-community/Mistral-7B-Instruct-v0.3", SPMStreamingDetokenizer),
             ("mlx-community/Llama-3.2-1B-Instruct-4bit", BPEStreamingDetokenizer),
+            ("mlx-community/Falcon3-7B-Instruct-4bit", BPEStreamingDetokenizer),
         ]
         for tokenizer_repo, expected_detokenizer in tokenizer_repos:
             with self.subTest(tokenizer=tokenizer_repo):
@@ -71,6 +82,17 @@ class TestTokenizers(unittest.TestCase):
         tokenizer._detokenizer = NaiveStreamingDetokenizer(tokenizer)
         self.check_tokenizer(tokenizer)
 
+    def test_special_tokens(self):
+        tokenizer_repo = "mlx-community/DeepSeek-Coder-V2-Lite-Instruct-4bit-mlx"
+        tokenizer = self.download_tokenizer(tokenizer_repo)
+
+        detokenizer = tokenizer.detokenizer
+        detokenizer.reset()
+        detokenizer.add_token(tokenizer.eos_token_id)
+        detokenizer.finalize()
+
+        self.assertEqual(detokenizer.last_segment, tokenizer.eos_token)
+
 
 if __name__ == "__main__":
     unittest.main()

llms/tests/test_utils_load_model.py

@@ -32,7 +32,7 @@ class TestLoadModelCustomGetClasses(unittest.TestCase):
             return CustomQwenModel, CustomQwenConfig
 
         model_path = get_model_path(HF_MODEL_PATH)
-        model = load_model(model_path, get_model_classes=custom_get_classes)
+        model, _ = load_model(model_path, get_model_classes=custom_get_classes)
 
         self.assertIsInstance(model, CustomQwenModel)
         self.assertTrue(hasattr(model, "custom_attribute"))
@@ -41,7 +41,7 @@ class TestLoadModelCustomGetClasses(unittest.TestCase):
 
     def test_load_model_with_default_get_classes(self):
         model_path = get_model_path(HF_MODEL_PATH)
-        model = load_model(model_path)
+        model, _ = load_model(model_path)
 
         self.assertIsInstance(model, Qwen2Model)
 

speechcommands/main.py

@@ -76,6 +76,7 @@ def train_epoch(model, train_iter, optimizer, epoch):
     samples_per_sec = []
 
     model.train(True)
+    train_iter.reset()
     for batch_counter, batch in enumerate(train_iter):
         x = mx.array(batch["audio"])
         y = mx.array(batch["label"])
@@ -111,6 +112,7 @@ def test_epoch(model, test_iter):
     model.train(False)
     accs = []
     throughput = []
+    test_iter.reset()
     for batch_counter, batch in enumerate(test_iter):
         x = mx.array(batch["audio"])
         y = mx.array(batch["label"])

stable_diffusion/image2image.py

@@ -30,6 +30,7 @@ if __name__ == "__main__":
     parser.add_argument("--preload-models", action="store_true")
     parser.add_argument("--output", default="out.png")
     parser.add_argument("--verbose", "-v", action="store_true")
+    parser.add_argument("--seed", type=int)
     args = parser.parse_args()
 
     # Load the models
@@ -94,6 +95,7 @@ if __name__ == "__main__":
         cfg_weight=args.cfg,
         num_steps=args.steps,
         negative_text=args.negative_prompt,
+        seed=args.seed,
     )
     for x_t in tqdm(latents, total=int(args.steps * args.strength)):
         mx.eval(x_t)

whisper/README.md

@@ -25,7 +25,7 @@ pip install mlx-whisper
 
 At its simplest:
 
-```
+```sh
 mlx_whisper audio_file.mp3
 ```
 
@@ -35,6 +35,15 @@ Use `-f` to specify the output format and `--model` to specify the model. There
 are many other supported command line options. To see them all, run
 `mlx_whisper -h`.
 
+You can also pipe the audio content of other programs via stdin:
+
+```sh
+some-process | mlx_whisper -
+```
+
+The default output file name will be `content.*`. You can specify the name with
+the `--output-name` flag.
+
 #### API
 
 Transcribe audio with:
@@ -103,7 +112,7 @@ python convert.py --help
 ```
 
 By default, the conversion script will make the directory `mlx_models`
-and save the converted `weights.npz` and `config.json` there. 
+and save the converted `weights.npz` and `config.json` there.
 
 Each time it is run, `convert.py` will overwrite any model in the provided
 path. To save different models, make sure to set `--mlx-path` to a unique

whisper/convert.py

@@ -174,14 +174,9 @@ def load_torch_weights_and_config(
                 "*.txt",
             ],
         )
-    else:
-        raise RuntimeError(
-            f"Model {name_or_path} is not found in {available_models()},"
-            "on Hugging Face or as a local path."
-        )
 
     if name_or_path.endswith(".pt"):
-        checkpoint = torch.load(name_or_path, map_location="cpu")
+        checkpoint = torch.load(name_or_path, map_location="cpu", weights_only=False)
         weights, config = checkpoint["model_state_dict"], checkpoint["dims"]
     else:
         name_or_path = Path(name_or_path)
@@ -387,7 +382,7 @@ if __name__ == "__main__":
 
     # Save weights
     print("[INFO] Saving")
-    np.savez(str(mlx_path / "weights.npz"), **weights)
+    mx.save_safetensors(str(mlx_path / "weights.safetensors"), weights)
 
     # Save config.json with model_type
     with open(str(mlx_path / "config.json"), "w") as f:

whisper/mlx_whisper/_version.py

@@ -1,3 +1,3 @@
 # Copyright © 2023-2024 Apple Inc.
 
-__version__ = "0.3.0"
+__version__ = "0.4.1"

whisper/mlx_whisper/audio.py

@@ -3,7 +3,7 @@
 import os
 from functools import lru_cache
 from subprocess import CalledProcessError, run
-from typing import Union
+from typing import Optional, Union
 
 import mlx.core as mx
 import numpy as np
@@ -21,7 +21,7 @@ FRAMES_PER_SECOND = SAMPLE_RATE // HOP_LENGTH  # 10ms per audio frame
 TOKENS_PER_SECOND = SAMPLE_RATE // N_SAMPLES_PER_TOKEN  # 20ms per audio token
 
 
-def load_audio(file: str, sr: int = SAMPLE_RATE):
+def load_audio(file: str = Optional[str], sr: int = SAMPLE_RATE, from_stdin=False):
     """
     Open an audio file and read as mono waveform, resampling as necessary
 
@@ -39,19 +39,21 @@ def load_audio(file: str, sr: int = SAMPLE_RATE):
     """
 
     # This launches a subprocess to decode audio while down-mixing
-    # and resampling as necessary.  Requires the ffmpeg CLI in PATH.
+    # and resampling as necessary. Requires the ffmpeg CLI in PATH.
+    if from_stdin:
+        cmd = ["ffmpeg", "-i", "pipe:0"]
+    else:
+        cmd = ["ffmpeg", "-nostdin", "-i", file]
+
     # fmt: off
-    cmd = [
-        "ffmpeg",
-        "-nostdin",
+    cmd.extend([
         "-threads", "0",
-        "-i", file,
         "-f", "s16le",
         "-ac", "1",
         "-acodec", "pcm_s16le",
         "-ar", str(sr),
         "-"
-    ]
+    ])
     # fmt: on
     try:
         out = run(cmd, capture_output=True, check=True).stdout

whisper/mlx_whisper/cli.py

@@ -2,9 +2,11 @@
 
 import argparse
 import os
+import pathlib
 import traceback
 import warnings
 
+from . import audio
 from .tokenizer import LANGUAGES, TO_LANGUAGE_CODE
 from .transcribe import transcribe
 from .writers import get_writer
@@ -27,15 +29,24 @@ def build_parser():
     parser = argparse.ArgumentParser(
         formatter_class=argparse.ArgumentDefaultsHelpFormatter
     )
-    parser.add_argument(
-        "audio", nargs="+", type=str, help="Audio file(s) to transcribe"
-    )
+
+    parser.add_argument("audio", nargs="+", help="Audio file(s) to transcribe")
+
     parser.add_argument(
         "--model",
         default="mlx-community/whisper-tiny",
         type=str,
         help="The model directory or hugging face repo",
     )
+    parser.add_argument(
+        "--output-name",
+        type=str,
+        default=None,
+        help=(
+            "The name of transcription/translation output files before "
+            "--output-format extensions"
+        ),
+    )
     parser.add_argument(
         "--output-dir",
         "-o",
@@ -200,6 +211,7 @@ def main():
     path_or_hf_repo: str = args.pop("model")
     output_dir: str = args.pop("output_dir")
     output_format: str = args.pop("output_format")
+    output_name: str = args.pop("output_name")
     os.makedirs(output_dir, exist_ok=True)
 
     writer = get_writer(output_format, output_dir)
@@ -219,17 +231,25 @@ def main():
         warnings.warn("--max-line-count has no effect without --max-line-width")
     if writer_args["max_words_per_line"] and writer_args["max_line_width"]:
         warnings.warn("--max-words-per-line has no effect with --max-line-width")
-    for audio_path in args.pop("audio"):
+
+    for audio_obj in args.pop("audio"):
+        if audio_obj == "-":
+            # receive the contents from stdin rather than read a file
+            audio_obj = audio.load_audio(from_stdin=True)
+
+            output_name = output_name or "content"
+        else:
+            output_name = output_name or pathlib.Path(audio_obj).stem
         try:
             result = transcribe(
-                audio_path,
+                audio_obj,
                 path_or_hf_repo=path_or_hf_repo,
                 **args,
             )
-            writer(result, audio_path, **writer_args)
+            writer(result, output_name, **writer_args)
         except Exception as e:
             traceback.print_exc()
-            print(f"Skipping {audio_path} due to {type(e).__name__}: {str(e)}")
+            print(f"Skipping {audio_obj} due to {type(e).__name__}: {str(e)}")
 
 
 if __name__ == "__main__":

whisper/mlx_whisper/decoding.py

@@ -58,11 +58,12 @@ def detect_language(
     logits = model.logits(x, mel)[:, 0]
 
     # collect detected languages; suppress all non-language tokens
-    mask = np.full(logits.shape[-1], -np.inf, dtype=np.float32)
+    mask = mx.full(logits.shape[-1], -mx.inf, dtype=mx.float32)
     mask[list(tokenizer.all_language_tokens)] = 0.0
-    logits += mx.array(mask)
+    logits += mask
     language_tokens = mx.argmax(logits, axis=-1)
     language_token_probs = mx.softmax(logits, axis=-1)
+    language_token_probs = np.array(language_token_probs)
     language_probs = [
         {
             c: language_token_probs[i, j].item()
@@ -129,17 +130,12 @@ class DecodingResult:
 
 
 class Inference:
-    def __init__(self, model: "Whisper", initial_token_length: int):
+    def __init__(self, model: "Whisper"):
         self.model: "Whisper" = model
-        self.initial_token_length = initial_token_length
         self.kv_cache = None
 
     def logits(self, tokens: mx.array, audio_features: mx.array) -> mx.array:
         """Perform a forward pass on the decoder and return per-token logits"""
-        if tokens.shape[-1] > self.initial_token_length:
-            # only need to use the last token except in the first forward pass
-            tokens = tokens[:, -1:]
-
         logits, self.kv_cache, _ = self.model.decoder(
             tokens, audio_features, kv_cache=self.kv_cache
         )
@@ -251,6 +247,11 @@ class TokenDecoder:
         raise NotImplementedError
 
 
+@mx.compile
+def categorical(logits, temp):
+    return mx.random.categorical(logits / temp)
+
+
 class GreedyDecoder(TokenDecoder):
     def __init__(self, temperature: float, eot: int):
         self.temperature = temperature
@@ -262,10 +263,8 @@ class GreedyDecoder(TokenDecoder):
         if self.temperature == 0:
             next_tokens = logits.argmax(axis=-1)
         else:
-            next_tokens = mx.random.categorical(logits=logits / self.temperature)
+            next_tokens = categorical(logits, self.temperature)
 
-        next_tokens = mx.argmax(logits, axis=-1)
-        logits = logits.astype(mx.float32)
         logprobs = logits - mx.logsumexp(logits, axis=-1)
 
         current_logprobs = logprobs[mx.arange(logprobs.shape[0]), next_tokens]
@@ -281,7 +280,7 @@ class GreedyDecoder(TokenDecoder):
     def finalize(self, tokens: mx.array, sum_logprobs: mx.array):
         # make sure each sequence has at least one EOT token at the end
         tokens = mx.pad(tokens, [(0, 0), (0, 0), (0, 1)], constant_values=self.eot)
-        return tokens, sum_logprobs.tolist()
+        return tokens, sum_logprobs
 
 
 class LogitFilter:
@@ -340,10 +339,10 @@ class ApplyTimestampRules(LogitFilter):
         if self.tokenizer.no_timestamps is not None:
             mask[:, self.tokenizer.no_timestamps] = -np.inf
 
-        # timestamps have to appear in pairs, except directly before EOT; mask logits accordingly
-        for k in range(tokens.shape[0]):
-            sampled_tokens = tokens[k, self.sample_begin :]
-            seq = sampled_tokens.tolist()
+        ## timestamps have to appear in pairs, except directly before EOT; mask logits accordingly
+        tokens = tokens.tolist()
+        for k in range(len(tokens)):
+            seq = tokens[k][self.sample_begin :]
             last_was_timestamp = (
                 len(seq) >= 1 and seq[-1] >= self.tokenizer.timestamp_begin
             )
@@ -368,7 +367,7 @@ class ApplyTimestampRules(LogitFilter):
                     last_timestamp += 1
                 mask[k, self.tokenizer.timestamp_begin : last_timestamp] = -np.inf
 
-        if tokens.shape[1] == self.sample_begin:
+        if len(tokens[0]) == self.sample_begin:
             # suppress generating non-timestamp tokens at the beginning
             mask[:, : self.tokenizer.timestamp_begin] = -np.inf
 
@@ -380,16 +379,20 @@ class ApplyTimestampRules(LogitFilter):
                 mask[:, last_allowed + 1 :] = -np.inf
 
         # if sum of probability over timestamps is above any other token, sample timestamp
+        mask = mx.array(mask)
         logprobs = logits - mx.logsumexp(logits, axis=-1)
-        for k in range(tokens.shape[0]):
-            timestamp_logprob = logprobs[k, self.tokenizer.timestamp_begin :].logsumexp(
-                axis=-1
-            )
-            max_text_token_logprob = logprobs[k, : self.tokenizer.timestamp_begin].max()
-            if timestamp_logprob > max_text_token_logprob:
-                mask[k, : self.tokenizer.timestamp_begin] = -np.inf
-
-        return logits + mx.array(mask, logits.dtype)
+        timestamp_logprob = logprobs[:, self.tokenizer.timestamp_begin :].logsumexp(
+            axis=-1, keepdims=True
+        )
+        max_text_token_logprob = logprobs[:, : self.tokenizer.timestamp_begin].max(
+            axis=-1, keepdims=True
+        )
+        mask[:, : self.tokenizer.timestamp_begin] = mx.where(
+            timestamp_logprob > max_text_token_logprob,
+            -mx.inf,
+            mask[:, : self.tokenizer.timestamp_begin],
+        )
+        return logits + mask
 
 
 class DecodingTask:
@@ -424,7 +427,7 @@ class DecodingTask:
         self.sot_index: int = self.initial_tokens.index(tokenizer.sot)
 
         # inference: implements the forward pass through the decoder, including kv caching
-        self.inference = Inference(model, len(self.initial_tokens))
+        self.inference = Inference(model)
 
         # sequence ranker: implements how to rank a group of sampled sequences
         self.sequence_ranker = MaximumLikelihoodRanker(options.length_penalty)
@@ -432,9 +435,6 @@ class DecodingTask:
         # decoder: implements how to select the next tokens, given the autoregressive distribution
         if options.beam_size is not None:
             raise NotImplementedError("Beam search decoder is not yet implemented")
-            # self.decoder = BeamSearchDecoder(
-            #    options.beam_size, tokenizer.eot, self.inference, options.patience
-            # )
         else:
             self.decoder = GreedyDecoder(options.temperature, tokenizer.eot)
 
@@ -448,6 +448,7 @@ class DecodingTask:
             self.logit_filters.append(
                 SuppressTokens(self._get_suppress_tokens(), model.dims.n_vocab)
             )
+
         if not options.without_timestamps:
             precision = CHUNK_LENGTH / model.dims.n_audio_ctx  # usually 0.02 seconds
             max_initial_timestamp_index = None
@@ -570,48 +571,59 @@ class DecodingTask:
 
     def _main_loop(self, audio_features: mx.array, tokens: mx.array):
         n_batch = tokens.shape[0]
-        sum_logprobs: mx.array = mx.zeros(n_batch)
-        no_speech_probs = [np.nan] * n_batch
+        sum_logprobs = mx.zeros(n_batch)
 
-        try:
-            for i in range(self.sample_len):
-                logits = self.inference.logits(tokens, audio_features)
+        def _step(inputs, audio_features, tokens, sum_logprobs):
+            pre_logits = self.inference.logits(inputs, audio_features)
 
-                if (
-                    i == 0 and self.tokenizer.no_speech is not None
-                ):  # save no_speech_probs
-                    probs_at_sot = mx.softmax(
-                        logits[:, self.sot_index].astype(mx.float32), axis=-1
-                    )
-                    no_speech_probs = probs_at_sot[:, self.tokenizer.no_speech].tolist()
+            # consider the logits at the last token only
+            logits = pre_logits[:, -1]
 
-                # now we need to consider the logits at the last token only
-                logits = logits[:, -1]
+            # apply the logit filters, e.g. for suppressing or applying penalty to
+            for logit_filter in self.logit_filters:
+                logits = logit_filter.apply(logits, tokens)
 
-                # apply the logit filters, e.g. for suppressing or applying penalty to
-                for logit_filter in self.logit_filters:
-                    logits = logit_filter.apply(logits, tokens)
+            # expand the tokens tensor with the selected next tokens
+            tokens, completed, sum_logprobs = self.decoder.update(
+                tokens, logits, sum_logprobs
+            )
+            return tokens, completed, sum_logprobs, pre_logits
 
-                # expand the tokens tensor with the selected next tokens
-                tokens, completed, sum_logprobs = self.decoder.update(
-                    tokens, logits, sum_logprobs
-                )
+        tokens, completed, sum_logprobs, pre_logits = _step(
+            tokens, audio_features, tokens, sum_logprobs
+        )
+        if self.tokenizer.no_speech is not None:  # compute no_speech_probs
+            probs_at_sot = mx.softmax(pre_logits[:, self.sot_index], axis=-1)
+            no_speech_probs = probs_at_sot[:, self.tokenizer.no_speech]
+        else:
+            no_speech_probs = mx.full(n_batch, mx.nan)
+        mx.async_eval(completed, tokens, sum_logprobs, no_speech_probs)
 
-                if completed or tokens.shape[-1] > self.n_ctx:
-                    break
-        finally:
-            self.inference.reset()
+        for i in range(1, self.sample_len):
+            inputs = tokens[:, -1:]
+            if tokens.shape[-1] > self.n_ctx:
+                break
+            next_tokens, next_completed, next_sum_logprobs, _ = _step(
+                inputs, audio_features, tokens, sum_logprobs
+            )
+            mx.async_eval(next_completed, next_tokens, next_sum_logprobs)
+            if completed:
+                break
+            tokens = next_tokens
+            completed = next_completed
+            sum_logprobs = next_sum_logprobs
 
         return tokens, sum_logprobs, no_speech_probs
 
     def run(self, mel: mx.array) -> List[DecodingResult]:
+        self.inference.reset()
         self.decoder.reset()
         tokenizer: Tokenizer = self.tokenizer
         n_audio: int = mel.shape[0]
 
         audio_features: mx.array = self._get_audio_features(mel)  # encoder forward pass
-        tokens: np.array = np.array(self.initial_tokens)
-        tokens = np.broadcast_to(tokens, (n_audio, len(self.initial_tokens))).copy()
+        tokens: mx.array = mx.array(self.initial_tokens)
+        tokens = mx.broadcast_to(tokens, (n_audio, len(self.initial_tokens)))
 
         # detect language if requested, overwriting the language token
         languages, language_probs = self._detect_language(audio_features, tokens)
@@ -626,7 +638,6 @@ class DecodingTask:
             ]
 
         # repeat tokens by the group size, for beam search or best-of-n sampling
-        tokens = mx.array(tokens)
         if self.n_group > 1:
             tokens = tokens[:, None, :]
             tokens = mx.broadcast_to(
@@ -649,7 +660,13 @@ class DecodingTask:
 
         # get the final candidates for each group, and slice between the first sampled token and EOT
         tokens, sum_logprobs = self.decoder.finalize(tokens, sum_logprobs)
-        tokens = tokens[..., self.sample_begin :].tolist()
+        tokens = tokens[..., self.sample_begin :]
+
+        # eval and convert to list
+        mx.eval(tokens, sum_logprobs, no_speech_probs)
+        tokens = tokens.tolist()
+        sum_logprobs = sum_logprobs.tolist()
+        no_speech_probs = no_speech_probs.tolist()
         tokens = [[t[: t.index(tokenizer.eot)] for t in s] for s in tokens]
 
         # select the top-ranked sample in each group

whisper/mlx_whisper/load_models.py

@@ -26,7 +26,10 @@ def load_model(
 
     model_args = whisper.ModelDimensions(**config)
 
-    weights = mx.load(str(model_path / "weights.npz"))
+    wf = model_path / "weights.safetensors"
+    if not wf.exists():
+        wf = model_path / "weights.npz"
+    weights = mx.load(str(wf))
 
     model = whisper.Whisper(model_args, dtype)
 

whisper/mlx_whisper/transcribe.py

@@ -293,6 +293,7 @@ def transcribe(
 
                 decode_options["prompt"] = all_tokens[prompt_reset_since:]
                 result: DecodingResult = decode_with_fallback(mel_segment)
+
                 tokens = np.array(result.tokens)
 
                 if no_speech_threshold is not None:

whisper/mlx_whisper/whisper.py

@@ -80,12 +80,11 @@ class MultiHeadAttention(nn.Module):
         qk = q @ k
         if mask is not None:
             qk = qk + mask[:n_ctx, :n_ctx]
-        qk = qk.astype(mx.float32)
 
-        w = mx.softmax(qk, axis=-1).astype(q.dtype)
+        w = mx.softmax(qk, axis=-1, precise=True)
         out = (w @ v).transpose(0, 2, 1, 3)
         out = out.reshape(n_batch, n_ctx, n_state)
-        return out, qk
+        return out, qk.astype(mx.float32)
 
 
 class ResidualAttentionBlock(nn.Module):

whisper/mlx_whisper/writers.py

@@ -1,10 +1,8 @@
 # Copyright © 2024 Apple Inc.
 
 import json
-import os
+import pathlib
 import re
-import sys
-import zlib
 from typing import Callable, List, Optional, TextIO
 
 
@@ -43,15 +41,13 @@ class ResultWriter:
         self.output_dir = output_dir
 
     def __call__(
-        self, result: dict, audio_path: str, options: Optional[dict] = None, **kwargs
+        self, result: dict, output_name: str, options: Optional[dict] = None, **kwargs
     ):
-        audio_basename = os.path.basename(audio_path)
-        audio_basename = os.path.splitext(audio_basename)[0]
-        output_path = os.path.join(
-            self.output_dir, audio_basename + "." + self.extension
+        output_path = (pathlib.Path(self.output_dir) / output_name).with_suffix(
+            f".{self.extension}"
         )
 
-        with open(output_path, "w", encoding="utf-8") as f:
+        with output_path.open("wt", encoding="utf-8") as f:
             self.write_result(result, file=f, options=options, **kwargs)
 
     def write_result(