Batch generation (#443)

* initial batch generation

* more in batch generate

* concatenation

* use batch API in eval

* unique max tokens per prompt

* basic continuous batching

* simplify

* better perf by ensuring everything in same stream

* use data class for response

* check cache type

README.md

@@ -85,7 +85,9 @@ To see a description of all the arguments you can do:
 
 Check out the [generation
 example](https://github.com/ml-explore/mlx-lm/tree/main/mlx_lm/examples/generate_response.py)
-to see how to use the API in more detail.
+to see how to use the API in more detail. Check out the [batch generation
+example](https://github.com/ml-explore/mlx-lm/tree/main/mlx_lm/examples/batch_generate_response.py)
+to see how to efficiently generate continuations for a batch of prompts.
 
 The `mlx-lm` package also comes with functionality to quantize and optionally
 upload models to the Hugging Face Hub.

mlx_lm/__init__.py

@@ -7,5 +7,5 @@ from ._version import __version__
 os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "1"
 
 from .convert import convert
-from .generate import generate, stream_generate
+from .generate import batch_generate, generate, stream_generate
 from .utils import load

mlx_lm/benchmark.py

@@ -4,7 +4,7 @@ import argparse
 
 import mlx.core as mx
 
-from mlx_lm import stream_generate
+from mlx_lm import batch_generate, stream_generate
 from mlx_lm.generate import DEFAULT_MODEL
 from mlx_lm.tokenizer_utils import load_tokenizer
 from mlx_lm.utils import (
@@ -44,6 +44,13 @@ def setup_arg_parser():
         help="Length of completion",
         type=int,
     )
+    parser.add_argument(
+        "--batch-size",
+        "-b",
+        default=1,
+        help="Batch size",
+        type=int,
+    )
     parser.add_argument(
         "--num-trials",
         "-n",
@@ -71,20 +78,31 @@ def main():
 
     prompt_tokens = args.prompt_tokens
     generation_tokens = args.generation_tokens
-    prompt = mx.random.randint(0, config["vocab_size"], (prompt_tokens,))
+    batch_size = args.batch_size
+    prompts = mx.random.randint(
+        0, config["vocab_size"], (batch_size, prompt_tokens)
+    ).tolist()
+    prompt = prompts[0]
 
-    def _bench():
+    def single_bench():
         for response in stream_generate(
             model, tokenizer, prompt, max_tokens=generation_tokens
         ):
             pass
         return response
 
+    def batch_bench():
+        return batch_generate(
+            model, tokenizer, prompts, max_tokens=generation_tokens
+        ).stats
+
+    _bench = batch_bench
+
     print("Running warmup..")
     _bench()
 
     report_keys = ["prompt_tps", "generation_tps", "peak_memory"]
-    print(f"Timing with {prompt_tokens=} and {generation_tokens=}.")
+    print(f"Timing with {prompt_tokens=}, {generation_tokens=}, {batch_size=}.")
     responses = []
     for i in range(args.num_trials):
         response = _bench()

mlx_lm/evaluate.py

@@ -23,8 +23,7 @@ from lm_eval.api.registry import register_model
 from lm_eval.models import huggingface
 from tqdm import tqdm
 
-from .generate import stream_generate
-from .models.base import create_causal_mask
+from .generate import batch_generate
 from .models.cache import make_prompt_cache
 from .utils import common_prefix_len, load
 
@@ -69,7 +68,7 @@ class MLXLM(LM):
     def __init__(
         self,
         path_or_hf_repo: str,
-        max_tokens: Optional[int] = None,
+        max_tokens: int,
         use_chat_template: Optional[bool] = None,
         trust_remote_code: bool = False,
     ) -> None:
@@ -78,7 +77,7 @@ class MLXLM(LM):
         self._model, self.tokenizer = load(
             path_or_hf_repo, tokenizer_config=tokenizer_config
         )
-        self._max_tokens = max_tokens or self.tokenizer.model_max_length
+        self._max_tokens = max_tokens
         self._batch_size = 8
         self.use_chat_template = use_chat_template
         if use_chat_template is None:
@@ -307,30 +306,33 @@ class MLXLM(LM):
         """
         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
+        # The second element of the tuple contains:
         # {'do_sample': False, 'until': ['\n\n'], 'temperature': 0}
         completions = []
 
-        for context, opt in tqdm(zip(contexts, options), total=len(contexts)):
-            until = opt["until"]
-            context = self.tokenizer.encode(
+        # Tokenize all contexts
+        contexts = [
+            self.tokenizer.encode(
                 context, add_special_tokens=not self.use_chat_template
             )
-            max_tokens = min(
-                opt.get("max_gen_tokens", self._max_tokens),
-                self.tokenizer.model_max_length - len(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)
+            for context in contexts
+        ]
+
+        # TODO consider multi-token, per-prompt stop conditions
+        max_tokens = [opt.get("max_gen_toks", self._max_tokens) for opt in options]
+
+        completions = batch_generate(
+            model=self._model,
+            tokenizer=self.tokenizer,
+            prompts=contexts,
+            max_tokens=max_tokens,
+            verbose=True,
+        ).texts
+
+        for e, (text, opt) in enumerate(zip(completions, options)):
+            until = opt["until"]
+            if any(u in text for u in until):
+                completions[e] = _rstrip_until(text, until)
         return completions
 
 
@@ -348,7 +350,8 @@ def main():
     parser.add_argument(
         "--max-tokens",
         type=int,
-        help="Maximum nunber of tokens to generate. Defaults to the model's max context length.",
+        help="Maximum number of tokens to generate.",
+        default=8912,
     )
     parser.add_argument(
         "--limit",

mlx_lm/examples/batch_generate_response.py

@@ -0,0 +1,32 @@
+# Copyright © 2025 Apple Inc.
+
+from mlx_lm import batch_generate, load
+
+# Specify the checkpoint
+checkpoint = "mlx-community/Llama-3.2-3B-Instruct-4bit"
+
+# Load the corresponding model and tokenizer
+model, tokenizer = load(path_or_hf_repo=checkpoint)
+
+# A batch of prompts
+prompts = [
+    "Write a story about Einstein.",
+    "Why is the sky blue?",
+    "What time is it?",
+    "How tall is Mt Everest?",
+]
+
+# Apply the chat template and encode to tokens
+prompts = [
+    tokenizer.apply_chat_template(
+        [{"role": "user", "content": p}],
+        add_generation_prompt=True,
+    )
+    for p in prompts
+]
+
+# Set `verbose=True` to see generation statistics
+result = batch_generate(model, tokenizer, prompts, verbose=False, max_tokens=128)
+
+# The returned result contains texts completions in the same order as prompts
+print(result.texts[0])

mlx_lm/generate.py

@@ -24,6 +24,8 @@ from transformers import PreTrainedTokenizer
 
 from .models import cache
 from .models.cache import (
+    BatchKVCache,
+    KVCache,
     QuantizedKVCache,
     load_prompt_cache,
 )
@@ -463,7 +465,7 @@ def speculative_generate_step(
     model: nn.Module,
     draft_model: nn.Module,
     *,
-    num_draft_tokens=2,
+    num_draft_tokens: int = 2,
     max_tokens: int = 256,
     sampler: Optional[Callable[mx.array, mx.array]] = None,
     logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = None,
@@ -637,6 +639,7 @@ def stream_generate(
     model: nn.Module,
     tokenizer: Union[PreTrainedTokenizer, TokenizerWrapper],
     prompt: Union[str, mx.array, List[int]],
+    max_tokens: int = 256,
     draft_model: Optional[nn.Module] = None,
     **kwargs,
 ) -> Generator[GenerationResponse, None, None]:
@@ -648,6 +651,8 @@ def stream_generate(
         tokenizer (PreTrainedTokenizer): The tokenizer.
         prompt (Union[str, mx.array, List[int]]): The input prompt string or
           integer tokens.
+        max_tokens (int): The maximum number of tokens to generate.
+          Default: ``256``.
         draft_model (Optional[nn.Module]): An optional draft model. If provided
           then speculative decoding is used. The draft model must use the same
           tokenizer as the main model. Default: ``None``.
@@ -672,6 +677,8 @@ def stream_generate(
 
     detokenizer = tokenizer.detokenizer
 
+    kwargs["max_tokens"] = max_tokens
+
     if draft_model is None:
         kwargs.pop("num_draft_tokens", None)
         token_generator = generate_step(prompt, model, **kwargs)
@@ -696,6 +703,8 @@ def stream_generate(
                 break
 
             detokenizer.add_token(token)
+            if (n + 1) == max_tokens:
+                break
 
             yield GenerationResponse(
                 text=detokenizer.last_segment,
@@ -771,6 +780,329 @@ def generate(
     return text
 
 
+def _left_pad_prompts(prompts, max_length=None):
+    if max_length is None:
+        max_length = max(len(p) for p in prompts)
+    return mx.array([[0] * (max_length - len(p)) + p for p in prompts])
+
+
+@dataclass
+class BatchStats:
+    """
+    An data object to hold generation stats.
+
+    Args:
+        prompt_tokens (int): The number of prompt tokens processed.
+        prompt_tps (float): The prompt processing tokens-per-second.
+        prompt_time (float): The time in seconds spent in prompt processing.
+        generation_tokens (int): The number of generated tokens.
+        generation_tps (float): The tokens-per-second for generation.
+        generation_time (float): The time in seconds spent in generation .
+        peak_memory (float): The peak memory used so far in GB.
+    """
+
+    prompt_tokens: int = 0
+    prompt_tps: float = 0
+    prompt_time: float = 0
+    generation_tokens: int = 0
+    generation_tps: float = 0
+    generation_time: float = 0
+    peak_memory: float = 0
+
+
+@dataclass
+class BatchResponse:
+    """
+    An data object to hold a batch generation response.
+
+    Args:
+        texts: (List[str]): The generated text for each prompt.
+        stats (BatchStats): Statistics about the generation.
+    """
+
+    texts: List[str]
+    stats: BatchStats
+
+
+@dataclass
+class Batch:
+    uids: List[int]
+    y: mx.array
+    logprobs: mx.array
+    max_tokens: List[int]
+    num_tokens: List[int]
+    cache: List[Any]
+
+    def __len__(self):
+        return len(self.uids)
+
+    def filter(self, keep_idx: List[int]):
+        self.uids = [self.uids[k] for k in keep_idx]
+        self.max_tokens = [self.max_tokens[k] for k in keep_idx]
+        self.num_tokens = [self.num_tokens[k] for k in keep_idx]
+        keep_idx = mx.array(keep_idx, mx.int32)
+        self.y = self.y[keep_idx]
+        self.logprobs = self.logprobs[keep_idx]
+        for c in self.cache:
+            c.filter(keep_idx)
+
+    def extend(self, other):
+        self.uids.extend(other.uids)
+        self.y = mx.concatenate([self.y, other.y])
+        self.logprobs = mx.concatenate([self.logprobs, other.logprobs])
+        self.num_tokens.extend(other.num_tokens)
+        self.max_tokens.extend(other.max_tokens)
+        for c, o in zip(self.cache, other.cache):
+            c.extend(o)
+
+
+def _make_cache(model, left_padding):
+    """
+    Convert a list of regular caches into their corresponding
+    batch-aware caches.
+    """
+    if hasattr(model, "make_cache"):
+        cache = model.make_cache()
+        batch_cache = []
+        for c in cache:
+            if not isinstance(c, KVCache):
+                raise ValueError(f"{type(c)} does not yet support batching")
+            # Convert cache to batched cache
+            batch_cache.append(BatchKVCache(left_padding))
+    else:
+        return [BatchKVCache(left_padding) for _ in model.layers]
+
+
+class BatchGenerator:
+
+    @dataclass
+    class Response:
+        uid: int
+        token: int
+        logprobs: mx.array
+        finish_reason: Optional[str]
+
+    def __init__(
+        self,
+        model,
+        max_tokens: int = 128,
+        stop_tokens: Optional[set] = None,
+        sampler: Optional[Callable[mx.array, mx.array]] = None,
+        completion_batch_size: int = 32,
+        prefill_batch_size: int = 8,
+        prefill_step_size: int = 2048,
+    ):
+        self.model = model
+        self.unprocessed_prompts = []
+        self.max_tokens = max_tokens
+        self.stop_tokens = stop_tokens or set()
+        self.sampler = sampler or (lambda x: mx.argmax(x, keepdims=True, axis=-1))
+        self.uid_count = 0
+        self.prefill_step_size = prefill_step_size
+        self.prefill_batch_size = prefill_batch_size
+        self.completion_batch_size = completion_batch_size
+        self._stats = BatchStats()
+
+        self.active_batch = None
+
+    def insert(self, prompts, max_tokens: Union[List[int], int, None] = None):
+        uids = []
+
+        if max_tokens is None or isinstance(max_tokens, int):
+            max_tokens = [max_tokens or self.max_tokens] * len(prompts)
+
+        for p, m in zip(prompts, max_tokens):
+            self.unprocessed_prompts.append((self.uid_count, p, m))
+            uids.append(self.uid_count)
+            self.uid_count += 1
+        # Sort in ascending order of length
+        self.unprocessed_prompts = sorted(
+            self.unprocessed_prompts, key=lambda x: len(x[1])
+        )
+        return uids
+
+    def _process_prompts(self, prompts):
+        uids, inputs, max_tokens = zip(*prompts)
+        lengths = [len(p) for p in inputs]
+        max_length = max(lengths)
+        batch_size = self.prefill_batch_size
+        self._stats.prompt_tokens += sum(lengths)
+        left_padding = [max_length - l for l in lengths]
+        inputs = _left_pad_prompts(inputs, max_length=max_length)
+
+        prompt_cache = _make_cache(self.model, left_padding)
+
+        while inputs.shape[1] > 1:
+            n_to_process = min(self.prefill_step_size, inputs.shape[1] - 1)
+            self.model(inputs[:, :n_to_process], cache=prompt_cache)
+            mx.eval([c.state for c in prompt_cache])
+            inputs = inputs[:, n_to_process:]
+            mx.clear_cache()
+
+        y, logprobs = self._step(inputs, prompt_cache)
+        mx.async_eval(y, logprobs)
+        return Batch(
+            list(uids), y, logprobs, list(max_tokens), [0] * len(uids), prompt_cache
+        )
+
+    def _step(self, input_tokens: mx.array, prompt_cache: List[Any]):
+        logits = self.model(input_tokens, cache=prompt_cache)
+        logits = logits[:, -1, :]
+        logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
+        sampled = self.sampler(logprobs).squeeze(1)
+        return sampled, logprobs
+
+    def stats(self):
+        self._stats.prompt_tps = self._stats.prompt_tokens / self._stats.prompt_time
+        self._stats.generation_tps = (
+            self._stats.generation_tokens / self._stats.generation_time
+        )
+        self._stats.peak_memory = mx.get_peak_memory() / 1e9
+        return self._stats
+
+    def _next(self):
+        tic = time.perf_counter()
+
+        prompt_processing = False
+        batch = self.active_batch
+        num_active = len(batch) if batch else 0
+        num_to_add = self.completion_batch_size - num_active
+        while num_to_add >= self.prefill_batch_size:
+            prompts = self.unprocessed_prompts[: self.prefill_batch_size]
+            # Finish processing the last examples of the last batch
+            if len(prompts) == 0 and num_active > 0:
+                break
+            # No more prompts and no more completions, all done
+            elif len(prompts) == 0:
+                self.active_batch = None
+                return []
+            # Process prompts
+            if batch is not None and not prompt_processing:
+                # Finish any active completion tokens
+                mx.eval(batch.y, batch.logprobs)
+                self._stats.generation_time += time.perf_counter() - tic
+                tic = time.perf_counter()
+
+            batch = self._process_prompts(prompts)
+            self.unprocessed_prompts = self.unprocessed_prompts[
+                self.prefill_batch_size :
+            ]
+            prompt_processing = True
+            # If there was no active batch, set it
+            if self.active_batch is None:
+                self.active_batch = batch
+            else:
+                self.active_batch.extend(batch)
+
+            num_active = len(self.active_batch)
+            num_to_add -= len(batch)
+
+        batch = self.active_batch
+        y, logprobs = batch.y, batch.logprobs
+        batch.y, batch.logprobs = self._step(y[:, None], batch.cache)
+        mx.async_eval(batch.y, batch.logprobs)
+
+        y = y.tolist()
+        toc = time.perf_counter()
+        if prompt_processing:
+            self._stats.prompt_time += toc - tic
+        else:
+            self._stats.generation_time += toc - tic
+        keep_idx = []
+        end_idx = []
+        responses = []
+
+        for e, (t, uid, num_tok, max_tok) in enumerate(
+            zip(y, batch.uids, batch.num_tokens, batch.max_tokens)
+        ):
+            num_tok += 1
+            batch.num_tokens[e] = num_tok
+            if t in self.stop_tokens:
+                finish_reason = "stop"
+                end_idx.append(e)
+            elif num_tok >= max_tok:
+                finish_reason = "length"
+                end_idx.append(e)
+            else:
+                finish_reason = None
+                keep_idx.append(e)
+            responses.append(self.Response(uid, t, logprobs[e], finish_reason))
+
+        # Remove any finished completions
+        if len(end_idx):
+            if len(keep_idx) > 0:
+                batch.filter(keep_idx)
+            else:
+                self.active_batch = None
+
+        self._stats.generation_tokens += len(responses)
+        return responses
+
+    def next(self):
+        with mx.stream(generation_stream):
+            return self._next()
+
+
+def batch_generate(
+    model,
+    tokenizer,
+    prompts: List[int],
+    max_tokens: Union[int, List[int]] = 128,
+    verbose: bool = False,
+    **kwargs,
+) -> BatchResponse:
+    """
+    Generate responses for the given batch of prompts.
+
+    Args:
+       model (nn.Module): The language model.
+       tokenizer (PreTrainedTokenizer): The tokenizer.
+       prompt (List[List[int]]): The input prompts.
+       verbose (bool): If ``True``, print tokens and timing information.
+          Default: ``False``.
+       max_tokens (Union[int, List[int]): Maximum number of output tokens. This
+          can be per prompt if a list is provided.
+       kwargs: The remaining options get passed to :obj:`BatchGenerator`.
+          See :obj:`BatchGenerator` for more details.
+    """
+
+    gen = BatchGenerator(model, stop_tokens=tokenizer.eos_token_ids, **kwargs)
+    num_samples = len(prompts)
+    fin = 0
+    if verbose:
+        print(f"[batch_generate] Finished processing 0/{num_samples} ...", end="\r")
+
+    with wired_limit(model, [generation_stream]):
+        uids = gen.insert(prompts, max_tokens)
+        results = {uid: [] for uid in uids}
+        while responses := gen.next():
+            for r in responses:
+                if verbose and r.finish_reason != None:
+                    fin += 1
+                    print(
+                        f"[batch_generate] Finished processing {fin}/{num_samples} ...",
+                        end="\r",
+                    )
+                if r.finish_reason != "stop":
+                    results[r.uid].append(r.token)
+    if verbose:
+        print(f"[batch_generate] Finished processing {fin}/{num_samples}")
+
+    # Return results in correct order
+    texts = [tokenizer.decode(results[uid]) for uid in uids]
+    stats = gen.stats()
+    if verbose:
+        print(
+            f"[batch_generate] Prompt: {stats.prompt_tokens} tokens, {stats.prompt_tps:.3f} tokens-per-sec"
+        )
+        print(
+            f"[batch_generate] Generation: {stats.generation_tokens} tokens, "
+            f"{stats.generation_tps:.3f} tokens-per-sec"
+        )
+        print(f"[batch_generate] Peak memory: {stats.peak_memory:.3f} GB")
+    return BatchResponse(texts, stats)
+
+
 def main():
     parser = setup_arg_parser()
     args = parser.parse_args()

mlx_lm/models/base.py

@@ -25,7 +25,8 @@ def create_causal_mask(
     N: int,
     offset: int = 0,
     window_size: Optional[int] = None,
-    lengths: Optional[mx.array] = None,
+    right_padding: Optional[mx.array] = None,
+    left_padding: Optional[mx.array] = None,
 ):
     rinds = mx.arange(offset + N)
     linds = mx.arange(offset, offset + N) if offset else rinds
@@ -34,9 +35,10 @@ def create_causal_mask(
     mask = linds >= rinds
     if window_size is not None:
         mask = mask & (linds < rinds + window_size)
-    if lengths is not None:
-        lengths = lengths[:, None, None, None]
-        mask = mask & (rinds < lengths)
+    if right_padding is not None:
+        mask = mask & (rinds < mx.expand_dims((offset + N) - right_padding, (1, 2, 3)))
+    if left_padding is not None:
+        mask = mask & (mx.expand_dims(left_padding, (1, 2, 3)) <= rinds)
     return mask
 
 

mlx_lm/models/cache.py

@@ -609,3 +609,130 @@ class CacheList(KVCache):
             l = len(c.state)
             c.state = v[start : start + l]
             start += l
+
+
+class BatchKVCache(_BaseCache):
+    def __init__(self, left_padding: List[int]):
+        """
+        The BatchKV cache expects inputs to be left-padded.
+
+        E.g. the following prompts:
+
+            [1, 3, 5]
+            [7]
+            [2, 6, 8, 9]
+
+        Should be padded like so:
+
+            [0, 1, 3, 5]
+            [0, 0, 0, 7]
+            [2, 6, 8, 9]
+
+        And ``left_padding`` specifies the amount of padding for each.
+        In this case, ``left_padding = [1, 3, 0]``.
+        """
+        self.keys = None
+        self.values = None
+        self.left_padding = mx.array(left_padding)
+        self.offset = mx.array([-l for l in left_padding])
+        self._idx = 0
+        self.step = 256
+
+    def update_and_fetch(self, keys, values):
+        prev = self._idx
+        if self.keys is None or (prev + keys.shape[2]) > self.keys.shape[2]:
+            B, n_kv_heads, _, k_head_dim = keys.shape
+            v_head_dim = values.shape[3]
+            n_steps = (self.step + keys.shape[2] - 1) // self.step
+            k_shape = (B, n_kv_heads, n_steps * self.step, k_head_dim)
+            v_shape = (B, n_kv_heads, n_steps * self.step, v_head_dim)
+            new_k = mx.zeros(k_shape, keys.dtype)
+            new_v = mx.zeros(v_shape, values.dtype)
+            if self.keys is not None:
+                if prev % self.step != 0:
+                    self.keys = self.keys[..., :prev, :]
+                    self.values = self.values[..., :prev, :]
+                self.keys = mx.concatenate([self.keys, new_k], axis=2)
+                self.values = mx.concatenate([self.values, new_v], axis=2)
+            else:
+                self.keys, self.values = new_k, new_v
+
+        self.offset += keys.shape[2]
+        self._idx += keys.shape[2]
+        self.keys[..., prev : self._idx, :] = keys
+        self.values[..., prev : self._idx, :] = values
+        return self.keys[..., : self._idx, :], self.values[..., : self._idx, :]
+
+    @property
+    def state(self):
+        k, v = self.keys, self.values
+        if self._idx < k.shape[2]:
+            k = k[..., : self._idx, :]
+            v = v[..., : self._idx, :]
+        return k, v, self.offset, self.left_padding
+
+    @state.setter
+    def state(self, v):
+        self.keys, self.values, self.offset, self.left_padding = v
+        self._idx = self.keys.shape[2]
+
+    def is_trimmable(self):
+        return True
+
+    def trim(self, n):
+        n = min(self._idx, n)
+        self._idx -= n
+        self.offset -= n
+        return n
+
+    def make_mask(self, N: int, return_array: bool = False, **kwargs):
+        return create_causal_mask(
+            N, offset=self._idx, left_padding=self.left_padding, **kwargs
+        )
+
+    def filter(self, batch_indices):
+        """
+        In-place filter to keep just the given indices in the cache.
+        """
+        self.keys = self.keys[batch_indices]
+        self.values = self.values[batch_indices]
+        if isinstance(self.offset, mx.array):
+            self.offset = self.offset[batch_indices]
+            self.left_padding = self.left_padding[batch_indices]
+
+        # Shift left to reduce padding
+        min_left_pad = self.left_padding.min().item()
+        if min_left_pad > 0:
+            self.keys = self.keys[..., min_left_pad:, :]
+            self.values = self.values[..., min_left_pad:, :]
+            self._idx -= min_left_pad
+            self.left_padding -= min_left_pad
+
+    def extend(self, other):
+        """
+        In-place extend this cache with the other cache.
+        """
+        max_idx = max(self._idx, other._idx)
+        max_size = max(self.keys.shape[2], other.keys.shape[2])
+
+        # Pad the keys and values so they are right-justified
+        # with the index and the same size
+        def pad(c):
+            left = max_idx - c._idx
+            right = max_size - c.keys.shape[2] - left
+            k, v = c.keys, c.values
+            if right < 0:
+                k = k[..., :right, :]
+                v = v[..., :right, :]
+                right = 0
+            if left != 0 or right != 0:
+                pad = [(0, 0), (0, 0), (left, right), (0, 0)]
+                k = mx.pad(k, pad)
+                v = mx.pad(v, pad)
+            left_padding = c.left_padding + left
+            return k, v, c.offset, left_padding
+
+        self.keys, self.values, self.offset, self.left_padding = map(
+            mx.concatenate, zip(*(pad(self), pad(other)))
+        )
+        self._idx = max_idx

tests/test_evaluate.py

@@ -19,7 +19,7 @@ class TestMLXLM(unittest.TestCase):
 
         with patch("mlx_lm.evaluate.load") as mock_load:
             mock_load.return_value = (self.mock_model, self.mock_tokenizer)
-            self.mlx_lm = MLXLM("test_model")
+            self.mlx_lm = MLXLM("test_model", max_tokens=128)
 
     def test_loglikelihood_rolling_processes_all_inputs(self):
         """Test that loglikelihood_rolling processes all inputs correctly when batching."""

tests/test_generate.py

@@ -3,7 +3,10 @@
 import unittest
 from typing import List
 
+import mlx.core as mx
+
 from mlx_lm.generate import (
+    BatchGenerator,
     GenerationResponse,
     generate,
     stream_generate,
@@ -18,6 +21,7 @@ class TestGenerate(unittest.TestCase):
     def setUpClass(cls):
         cls.HF_MODEL_PATH = "mlx-community/Qwen1.5-0.5B-Chat-4bit"
         cls.model, cls.tokenizer = load(cls.HF_MODEL_PATH)
+        cls.model.set_dtype(mx.float32)
 
     def test_generate(self):
         # Simple test that generation runs
@@ -37,6 +41,23 @@ class TestGenerate(unittest.TestCase):
         )
         self.assertEqual(text, "!!!!!")
 
+    def test_stream_generate_max_tokens(self):
+        prompt = self.tokenizer.apply_chat_template(
+            [{"role": "user", "content": "Write a story about Einstein"}],
+            tokenize=True,
+            add_generation_prompt=True,
+        )
+
+        tokens = []
+        for response in stream_generate(
+            self.model,
+            self.tokenizer,
+            prompt,
+            max_tokens=4,
+        ):
+            tokens.append(response.token)
+        self.assertEqual(len(tokens), 4)
+
     def test_generate_with_processor(self):
         init_toks = self.tokenizer.encode("hello")
 
@@ -83,8 +104,7 @@ class TestGenerate(unittest.TestCase):
             drafted.append(generation_result.from_draft)
             results.append(generation_result)
 
-        self.assertEqual(len(results), 6)
-        drafted.pop()
+        self.assertEqual(len(results), 5)
         # since num_draft_tokens is 2 and draft model is the same, the
         # first 2 generations should be drafts, the third should come
         # from the target model, and last two should be drafts
@@ -151,6 +171,136 @@ class TestGenerate(unittest.TestCase):
             num_embeddings / prefill_step_size < num_prompt_processing_callbacks
         )
 
+    def test_batch_matches_single(self):
+
+        prompts = [
+            "Write a story about Einstein",
+            "Hi",
+            "What time is it?",
+            "How tall is Mt Everest?",
+        ]
+        prompts = [
+            self.tokenizer.apply_chat_template(
+                [{"role": "user", "content": p}],
+                tokenize=True,
+                add_generation_prompt=True,
+            )
+            for p in prompts
+        ]
+
+        gen = BatchGenerator(
+            self.model, stop_tokens=self.tokenizer.eos_token_ids, max_tokens=1
+        )
+        uids = gen.insert(prompts)
+        batch_responses = {r.uid: r for r in gen.next()}
+
+        # Do a test for each prompt the logits are close
+        for e, prompt in enumerate(prompts):
+
+            for response in stream_generate(
+                self.model, self.tokenizer, prompt, max_tokens=1
+            ):
+                blp = batch_responses[uids[e]].logprobs
+                lp = response.logprobs
+                self.assertTrue(mx.allclose(blp, lp))
+                break
+
+    def test_many_batches(self):
+
+        prompts = [
+            "Write a story about Einstein",
+            "Hi",
+            "What time is it?",
+            "How tall is Mt Everest?",
+        ]
+        prompts = [
+            self.tokenizer.apply_chat_template(
+                [{"role": "user", "content": p}],
+                tokenize=True,
+                add_generation_prompt=True,
+            )
+            for p in prompts
+        ]
+
+        gen = BatchGenerator(
+            self.model,
+            stop_tokens=self.tokenizer.eos_token_ids,
+            max_tokens=1,
+            prefill_batch_size=2,
+            prefill_step_size=8,
+            completion_batch_size=3,
+        )
+        uids = gen.insert(prompts)
+        batch_responses = {}
+        not_in = True
+        iters = 0
+        while responses := gen.next():
+            for r in responses:
+                not_in &= r.uid not in batch_responses
+                batch_responses[r.uid] = r
+            iters += 1
+        # only one token per prompt means only one response per prompt
+        self.assertTrue(not_in)
+
+        # completion batch size is too small for a single iteration
+        self.assertTrue(iters > 1)
+
+        # Do a test for each prompt the logits are close
+        for e, prompt in enumerate(prompts):
+
+            for response in stream_generate(
+                self.model, self.tokenizer, prompt, max_tokens=1
+            ):
+                blp = batch_responses[uids[e]].logprobs
+                lp = response.logprobs
+                self.assertTrue(mx.allclose(blp, lp))
+                break
+
+    def test_batch_unique_max_toks(self):
+        prompts = [
+            "Write a story about Einstein",
+            "Hi",
+            "What time is it?",
+            "How tall is Mt Everest?",
+        ]
+        prompts = [
+            self.tokenizer.apply_chat_template(
+                [{"role": "user", "content": p}],
+                tokenize=True,
+                add_generation_prompt=True,
+            )
+            for p in prompts
+        ]
+
+        gen = BatchGenerator(
+            self.model,
+            stop_tokens=self.tokenizer.eos_token_ids,
+            prefill_batch_size=2,
+            prefill_step_size=8,
+            completion_batch_size=3,
+        )
+        num_toks = [2, 3, 4, 5]
+        uids = gen.insert(prompts, max_tokens=num_toks)
+        batch_responses = {uid: [] for uid in uids}
+        while responses := gen.next():
+            for r in responses:
+                batch_responses[r.uid].append(r.token)
+
+        # Do a test for each prompt the logits are close
+        for e, prompt in enumerate(prompts):
+
+            tokens = []
+            for response in stream_generate(
+                self.model,
+                self.tokenizer,
+                prompt,
+                max_tokens=num_toks[e],
+            ):
+                tokens.append(response.token)
+
+            batch_tokens = batch_responses[uids[e]]
+            self.assertEqual(tokens, batch_tokens)
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_models.py

@@ -132,21 +132,29 @@ class TestModels(unittest.TestCase):
         self.assertEqual(cache.offset, 22)
         self.assertTrue(mx.allclose(x, k[..., -2:, :]))
 
-    def test_causal_mask_lengths(self):
-        mx.random.seed(8)
-        B, N_q, T_q, N_kv, T_kv, D = (4, 8, 3, 2, 3, 2)
-        lengths = mx.array([1, 2, 3, 1])
-        q = mx.random.uniform(shape=(B, N_q, T_q, D))
-        k = mx.random.uniform(shape=(B, N_kv, T_kv, D))
-        v = k
-        mask = create_causal_mask(T_q, 0, lengths=lengths)
+    def test_causal_mask_padding(self):
+        right_padding = mx.array([2, 1, 0])
+        mask = create_causal_mask(3, right_padding=right_padding)
 
-        out1 = mx.fast.scaled_dot_product_attention(q, k, v, scale=1.0, mask=mask)
-        q[1, :, 2:] = mx.ones_like(q[1, :, 2:])
-        k[1, :, 2:] = mx.ones_like(k[1, :, 2:])
-        v[1, :, 2:] = mx.ones_like(v[1, :, 2:])
-        out2 = mx.fast.scaled_dot_product_attention(q, k, v, scale=1.0, mask=mask)
-        self.assertTrue(mx.allclose(out1[1, :, :2], out2[1, :, :2]))
+        causal_mask = create_causal_mask(3)
+        self.assertTrue(
+            mx.array_equal(mask[0, 0], causal_mask & mx.array([True, False, False]))
+        )
+        self.assertTrue(
+            mx.array_equal(mask[1, 0], causal_mask & mx.array([True, True, False]))
+        )
+        self.assertTrue(mx.array_equal(mask[2, 0], causal_mask))
+
+        left_padding = mx.array([2, 1, 0])
+        mask = create_causal_mask(3, left_padding=left_padding)
+
+        self.assertTrue(
+            mx.array_equal(mask[0, 0], causal_mask & mx.array([False, False, True]))
+        )
+        self.assertTrue(
+            mx.array_equal(mask[1, 0], causal_mask & mx.array([False, True, True]))
+        )
+        self.assertTrue(mx.array_equal(mask[2, 0], causal_mask))
 
     def test_mask_with_window(self):
         mask = create_causal_mask(5, 0, window_size=3)

tests/test_prompt_cache.py

@@ -10,6 +10,7 @@ import mlx.core as mx
 from mlx_lm.generate import generate_step
 from mlx_lm.models.base import create_attention_mask, create_causal_mask
 from mlx_lm.models.cache import (
+    BatchKVCache,
     CacheList,
     ChunkedKVCache,
     KVCache,
@@ -418,6 +419,47 @@ class TestPromptCache(unittest.TestCase):
         cmask = create_attention_mask(mx.zeros((1, 1, 32)), cache, window_size=5)
         self.assertTrue(mx.array_equal(cmask, mask))
 
+    def test_batch_kv_cache(self):
+        cache = BatchKVCache(left_padding=[2, 3, 4])
+        k, v = mx.zeros((3, 1, 4, 8)), mx.zeros((3, 1, 4, 8))
+        # Update works
+        k, v = cache.update_and_fetch(k, v)
+        self.assertEqual(k.shape, (3, 1, 4, 8))
+
+        # State can be evaluated
+        mx.eval(cache.state)
+
+        # State can be set
+        cache.state = cache.state
+
+        # Test filtering
+        cache.filter([0, 1])
+
+        # In this case filtering left shifts the cache so it has zero padding
+        self.assertEqual(cache.state[0].shape, (2, 1, 2, 8))
+
+        mask = cache.make_mask(1)
+        self.assertEqual(mask[0].squeeze().tolist(), [True, True, True])
+        self.assertEqual(mask[1].squeeze().tolist(), [False, True, True])
+
+        # Test extension
+        cache_a = BatchKVCache(left_padding=[2, 1, 2])
+        cache_b = BatchKVCache(left_padding=[3, 0])
+
+        k = mx.zeros((3, 1, 8, 1))
+        v = mx.zeros((3, 1, 8, 1))
+        cache_a.update_and_fetch(k, v)
+
+        k = mx.zeros((2, 1, 4, 1))
+        v = mx.zeros((2, 1, 4, 1))
+        cache_b.update_and_fetch(k, v)
+
+        cache_a.extend(cache_b)
+        self.assertEqual(cache_a.keys.shape[0], 5)
+        self.assertEqual(cache_a.values.shape[0], 5)
+        self.assertEqual(cache_a.offset.tolist(), [6, 7, 6, 1, 4])
+        self.assertEqual(cache_a.left_padding.tolist(), [2, 1, 2, 7, 4])
+
 
 if __name__ == "__main__":
     unittest.main()