# Copyright © 2023-2024 Apple Inc.
import argparse
import contextlib
import copy
import functools
import json
import sys
import time
from collections import deque
from dataclasses import dataclass
from functools import partial
from typing import Any, Callable, Generator, List, Optional, Sequence, Tuple, Union
import mlx.core as mx
import mlx.nn as nn
from mlx.utils import tree_reduce
from transformers import PreTrainedTokenizer
from .models import cache
from .models.cache import (
ArraysCache,
BatchKVCache,
BatchRotatingKVCache,
CacheList,
KVCache,
QuantizedKVCache,
RotatingKVCache,
TokenBuffer,
load_prompt_cache,
)
from .sample_utils import make_sampler
from .tokenizer_utils import TokenizerWrapper
from .utils import does_model_support_input_embeddings, load
DEFAULT_PROMPT = "hello"
DEFAULT_MAX_TOKENS = 100
DEFAULT_TEMP = 0.0
DEFAULT_TOP_P = 1.0
DEFAULT_MIN_P = 0.0
DEFAULT_TOP_K = 0
DEFAULT_XTC_PROBABILITY = 0.0
DEFAULT_XTC_THRESHOLD = 0.0
DEFAULT_MIN_TOKENS_TO_KEEP = 1
DEFAULT_SEED = None
DEFAULT_MODEL = "mlx-community/Llama-3.2-3B-Instruct-4bit"
DEFAULT_QUANTIZED_KV_START = 5000
def str2bool(string):
return string.lower() not in ["false", "f"]
def setup_arg_parser():
"""Set up and return the argument parser."""
parser = argparse.ArgumentParser(description="LLM inference script")
parser.add_argument(
"--model",
type=str,
help=(
"The path to the local model directory or Hugging Face repo. "
f"If no model is specified, then {DEFAULT_MODEL} is used."
),
default=None,
)
parser.add_argument(
"--trust-remote-code",
action="store_true",
help="Enable trusting remote code for tokenizer",
)
parser.add_argument(
"--adapter-path",
type=str,
help="Optional path for the trained adapter weights and config.",
)
parser.add_argument(
"--extra-eos-token",
type=str,
default=(),
nargs="+",
help="Add tokens in the list of eos tokens that stop generation.",
)
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)",
)
parser.add_argument(
"--prefill-response",
default=None,
help="Prefill response to be used for the chat template",
)
parser.add_argument(
"--max-tokens",
"-m",
type=int,
default=DEFAULT_MAX_TOKENS,
help="Maximum number of tokens to generate",
)
parser.add_argument(
"--temp", type=float, default=DEFAULT_TEMP, help="Sampling temperature"
)
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(
"--top-k", type=int, default=DEFAULT_TOP_K, help="Sampling top-k"
)
parser.add_argument(
"--xtc-probability",
type=float,
default=DEFAULT_XTC_PROBABILITY,
help="Probability of XTC sampling to happen each next token",
)
parser.add_argument(
"--xtc-threshold",
type=float,
default=0.0,
help="Thresold the probs of each next token candidate to be sampled by XTC",
)
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",
action="store_true",
help="Use the raw prompt without the tokenizer's chat template.",
)
parser.add_argument(
"--use-default-chat-template",
action="store_true",
help="Use the default chat template",
)
parser.add_argument(
"--chat-template-config",
help="Additional config for `apply_chat_template`. Should be a dictionary of"
" string keys to values represented as a JSON decodable string.",
default=None,
)
parser.add_argument(
"--verbose",
type=str2bool,
default=True,
help="Log verbose output when 'True' or 'T' or only print the response when 'False' or 'F'",
)
parser.add_argument(
"--max-kv-size",
type=int,
help="Set the maximum key-value cache size",
default=None,
)
parser.add_argument(
"--prompt-cache-file",
type=str,
default=None,
help="A file containing saved KV caches to avoid recomputing them",
)
parser.add_argument(
"--quantize-activations",
"-qa",
action="store_true",
help="Quantize activations using the same quantization config as the corresponding layer.",
)
parser.add_argument(
"--kv-bits",
type=int,
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,
)
parser.add_argument(
"--draft-model",
type=str,
help="A model to be used for speculative decoding.",
default=None,
)
parser.add_argument(
"--num-draft-tokens",
type=int,
help="Number of tokens to draft when using speculative decoding.",
default=3,
)
parser.add_argument(
"--flash-head",
action=argparse.BooleanOptionalAction,
default=None,
help="Maple only: use the approximate FlashHead output layer. Faster "
"decode, approximate token stream. Omit to follow the checkpoint config.",
)
return parser
# A stream on the default device just for generation
generation_stream = mx.new_thread_local_stream(mx.default_device())
@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.
"""
if not mx.metal.is_available():
try:
yield
finally:
pass
else:
model_bytes = tree_reduce(
lambda acc, x: acc + x.nbytes if isinstance(x, mx.array) else acc, model, 0
)
max_rec_size = mx.device_info()["max_recommended_working_set_size"]
if model_bytes > 0.9 * max_rec_size:
model_mb = model_bytes // 2**20
max_rec_mb = max_rec_size // 2**20
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-lm/tree/main#large-models"
)
old_limit = mx.set_wired_limit(max_rec_size)
try:
yield
finally:
if streams is not None:
for s in streams:
mx.synchronize(s)
else:
mx.synchronize()
mx.set_wired_limit(old_limit)
@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.
from_draft (bool): Whether the token was generated by the draft model.
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
from_draft: bool
prompt_tokens: int
prompt_tps: float
generation_tokens: int
generation_tps: float
peak_memory: float
finish_reason: Optional[str] = None
def maybe_quantize_kv_cache(prompt_cache, quantized_kv_start, kv_group_size, kv_bits):
if kv_bits is None:
return
for e, c in enumerate(prompt_cache):
if hasattr(c, "to_quantized") and c.offset >= quantized_kv_start:
prompt_cache[e] = c.to_quantized(group_size=kv_group_size, bits=kv_bits)
def generate_step(
prompt: mx.array,
model: nn.Module,
*,
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,
prefill_step_size: int = 2048,
kv_bits: Optional[int] = None,
kv_group_size: int = 64,
quantized_kv_start: int = 0,
prompt_progress_callback: Optional[Callable[[int, int], None]] = None,
input_embeddings: Optional[mx.array] = None,
) -> Generator[Tuple[mx.array, mx.array], None, None]:
"""
A generator producing token ids 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 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.
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_progress_callback (Callable[[int, int], None]): A call-back which takes the
prompt tokens processed so far and the total number of prompt tokens.
input_embeddings (mx.array, optional): Input embeddings to use instead of or in
conjunction with prompt tokens. Default: ``None``.
Yields:
Tuple[mx.array, mx.array]: One token and a vector of log probabilities.
"""
if input_embeddings is not None:
if not does_model_support_input_embeddings(model):
raise ValueError("Model does not support input embeddings.")
elif len(prompt) > 0 and len(prompt) != len(input_embeddings):
raise ValueError(
f"When providing input_embeddings, their sequence length ({len(input_embeddings)}) "
f"must match the sequence length of the prompt ({len(prompt)}), or the "
"prompt must be empty."
)
elif len(prompt) == 0:
raise ValueError(
"Either input_embeddings or prompt (or both) must be provided."
)
tokens = None
# Create the KV cache for generation
if prompt_cache is None:
prompt_cache = cache.make_prompt_cache(
model,
max_kv_size=max_kv_size,
)
prompt_progress_callback = prompt_progress_callback or (lambda *_: None)
quantize_cache_fn = functools.partial(
maybe_quantize_kv_cache,
quantized_kv_start=quantized_kv_start,
kv_group_size=kv_group_size,
kv_bits=kv_bits,
)
sampler = sampler or (lambda x: mx.argmax(x, axis=-1))
def _model_call(input_tokens: mx.array, input_embeddings: Optional[mx.array]):
if input_embeddings is not None:
return model(
input_tokens, cache=prompt_cache, input_embeddings=input_embeddings
)
else:
return model(input_tokens, cache=prompt_cache)
def _step(input_tokens: mx.array, input_embeddings: Optional[mx.array] = None):
nonlocal tokens
with mx.stream(generation_stream):
logits = _model_call(
input_tokens=input_tokens[None],
input_embeddings=(
input_embeddings[None] if input_embeddings is not None else None
),
)
logits = logits[:, -1, :]
if logits_processors and len(input_tokens) > 0:
tokens = (
mx.concat([tokens, input_tokens])
if tokens is not None
else input_tokens
)
for processor in logits_processors:
logits = processor(tokens, logits)
quantize_cache_fn(prompt_cache)
logprobs = logits - mx.logsumexp(logits, keepdims=True)
sampled = sampler(logprobs)
return sampled, logprobs.squeeze(0)
with mx.stream(generation_stream):
total_prompt_tokens = (
len(input_embeddings) if input_embeddings is not None else len(prompt)
)
prompt_processed_tokens = 0
prompt_progress_callback(prompt_processed_tokens, total_prompt_tokens)
while total_prompt_tokens - prompt_processed_tokens > 1:
remaining = (total_prompt_tokens - prompt_processed_tokens) - 1
n_to_process = min(prefill_step_size, remaining)
_model_call(
input_tokens=prompt[:n_to_process][None],
input_embeddings=(
input_embeddings[:n_to_process][None]
if input_embeddings is not None
else None
),
)
quantize_cache_fn(prompt_cache)
mx.eval([c.state for c in prompt_cache])
prompt_processed_tokens += n_to_process
prompt_progress_callback(prompt_processed_tokens, total_prompt_tokens)
prompt = prompt[n_to_process:]
input_embeddings = (
input_embeddings[n_to_process:]
if input_embeddings is not None
else input_embeddings
)
mx.clear_cache()
y, logprobs = _step(input_tokens=prompt, input_embeddings=input_embeddings)
mx.async_eval(y, logprobs)
n = 0
while True:
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.clear_cache()
y, logprobs = next_y, next_logprobs
n += 1
def speculative_generate_step(
prompt: mx.array,
model: nn.Module,
draft_model: nn.Module,
*,
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,
prompt_cache: Optional[Any] = None,
prefill_step_size: int = 512,
kv_bits: Optional[int] = None,
kv_group_size: int = 64,
quantized_kv_start: int = 0,
) -> Generator[Tuple[mx.array, mx.array, bool], None, None]:
"""
A generator producing token ids based on the given prompt from the model.
Args:
prompt (mx.array): The input prompt.
model (nn.Module): The model to use for generation.
draft_model (nn.Module): The draft model for speculative decoding.
num_draft_tokens (int, optional): The number of draft tokens for
speculative decoding. Default: ``2``.
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``.
prompt_cache (List[Any], optional): A pre-computed prompt cache. Note, if
provided, the cache will be updated in place. The cache must be trimmable.
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``.
Yields:
Tuple[mx.array, mx.array, bool]: One token, a vector of log probabilities,
and a bool indicating if the token was generated by the draft model
"""
y = prompt.astype(mx.uint32)
prev_tokens = None
# Create the KV cache for generation
if prompt_cache is None:
model_cache = cache.make_prompt_cache(model)
draft_cache = cache.make_prompt_cache(draft_model)
else:
model_cache = prompt_cache[: len(model.layers)]
draft_cache = prompt_cache[len(model.layers) :]
if not cache.can_trim_prompt_cache(model_cache):
types = {type(c).__name__ for c in model_cache if not c.is_trimmable()}
raise ValueError(
f"Speculative decoding requires a trimmable prompt cache " f"(got {types})."
)
sampler = sampler or (lambda x: mx.argmax(x, axis=-1))
quantize_cache_fn = functools.partial(
maybe_quantize_kv_cache,
quantized_kv_start=quantized_kv_start,
kv_group_size=kv_group_size,
kv_bits=kv_bits,
)
def _process_and_sample(tokens, logits):
if logits_processors:
for processor in logits_processors:
logits = processor(tokens, logits)
logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
y = sampler(logprobs)
return y, logprobs
def _step(model, cache, y, n_predict=1):
with mx.stream(generation_stream):
logits = model(y[None], cache=cache)
logits = logits[:, -n_predict:, :]
quantize_cache_fn(cache)
if logits_processors:
nonlocal prev_tokens
out_y, out_logprobs = [], []
if n_predict > 1:
y = y[: -(n_predict - 1)]
for i in range(n_predict):
prev_tokens = (
mx.concatenate([prev_tokens, y])
if prev_tokens is not None
else y
)
y, logprobs = _process_and_sample(prev_tokens, logits[:, i, :])
out_y.append(y)
out_logprobs.append(logprobs)
return mx.concatenate(out_y, axis=0), mx.concatenate(
out_logprobs, axis=0
)
else:
return _process_and_sample(None, logits.squeeze(0))
def _prefill(model, cache, y):
while y.size > 1:
n_to_process = min(prefill_step_size, y.size - 1)
model(y[:n_to_process][None], cache=cache)
quantize_cache_fn(cache)
mx.eval([c.state for c in cache])
y = y[n_to_process:]
mx.clear_cache()
return y
def _rewind_cache(num_draft, num_accept):
cache.trim_prompt_cache(model_cache, num_draft - num_accept)
cache.trim_prompt_cache(draft_cache, max(num_draft - num_accept - 1, 0))
def _draft_generate(y, num_draft):
if num_draft == 0:
return mx.array([], mx.uint32)
ys = []
for _ in range(num_draft):
y, _ = _step(draft_model, draft_cache, y)
mx.async_eval(y)
ys.append(y)
return mx.concatenate(ys)
with mx.stream(generation_stream):
draft_y = _prefill(draft_model, draft_cache, y)
y = _prefill(model, model_cache, y)
ntoks = 0
# Set these so the finally block doesn't raise
num_draft = 0
n = 0
try:
while True:
num_draft = min(max_tokens - ntoks, num_draft_tokens)
draft_tokens = _draft_generate(draft_y, num_draft)
if prev_tokens is not None:
prev_tokens = prev_tokens[: prev_tokens.size - y.size - num_draft + 1]
y = mx.concatenate([y, draft_tokens])
tokens, logprobs = _step(model, model_cache, y, num_draft + 1)
mx.eval(tokens, draft_tokens)
draft_tokens = draft_tokens.tolist()
tokens = tokens.tolist()
n = 0
while n < num_draft:
tn, dtn, lpn = tokens[n], draft_tokens[n], logprobs[n]
if tn != dtn:
break
n += 1
ntoks += 1
yield tn, lpn, True
if ntoks == max_tokens:
break
if ntoks < max_tokens:
ntoks += 1
yield tokens[n], logprobs[n], False
if ntoks == max_tokens:
break
y = mx.array([tokens[n]], mx.uint32)
draft_y = y
# If we accepted all the draft tokens, include the last
# draft token in the next draft step since it hasn't been
# processed yet by the draft model
if n == num_draft:
draft_y = mx.concatenate(
[mx.array(draft_tokens[-1:], mx.uint32), draft_y]
)
if prev_tokens is not None:
prev_tokens = prev_tokens[: -max(num_draft - n, 1)]
_rewind_cache(num_draft, n)
finally:
_rewind_cache(num_draft, n)
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]:
"""
A generator producing text based on the given prompt from the model.
Args:
model (nn.Module): The model to use for generation.
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``.
kwargs: The remaining options get passed to :func:`generate_step`.
See :func:`generate_step` for more details.
Yields:
GenerationResponse: An instance containing the generated text segment and
associated metadata. See :class:`GenerationResponse` for details.
"""
if not isinstance(tokenizer, TokenizerWrapper):
tokenizer = TokenizerWrapper(tokenizer)
if not isinstance(prompt, mx.array):
if isinstance(prompt, str):
# Try to infer if special tokens are needed
add_special_tokens = tokenizer.bos_token is None or not prompt.startswith(
tokenizer.bos_token
)
prompt = tokenizer.encode(prompt, add_special_tokens=add_special_tokens)
prompt = mx.array(prompt)
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)
# from_draft always false for non-speculative generation
token_generator = (
(token, logprobs, False) for token, logprobs in token_generator
)
else:
kwargs.pop("max_kv_size", None)
kwargs.pop("prompt_progress_callback", None)
token_generator = speculative_generate_step(
prompt, model, draft_model, **kwargs
)
with wired_limit(model, [generation_stream]):
tic = time.perf_counter()
for n, (token, logprobs, from_draft) in enumerate(token_generator):
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
detokenizer.add_token(token)
if (n + 1) == max_tokens:
break
yield GenerationResponse(
text=detokenizer.last_segment,
token=token,
logprobs=logprobs,
from_draft=from_draft,
prompt_tokens=prompt.size,
prompt_tps=prompt_tps,
generation_tokens=n + 1,
generation_tps=(n + 1) / (time.perf_counter() - tic),
peak_memory=mx.get_peak_memory() / 1e9,
finish_reason=None,
)
detokenizer.finalize()
yield GenerationResponse(
text=detokenizer.last_segment,
token=token,
logprobs=logprobs,
from_draft=from_draft,
prompt_tokens=prompt.size,
prompt_tps=prompt_tps,
generation_tokens=n + 1,
generation_tps=(n + 1) / (time.perf_counter() - tic),
peak_memory=mx.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: Union[str, List[int]],
verbose: bool = False,
**kwargs,
) -> str:
"""
Generate a complete response from the model.
Args:
model (nn.Module): The language model.
tokenizer (PreTrainedTokenizer): The tokenizer.
prompt (Union[str, List[int]]): The input prompt string or integer tokens.
verbose (bool): If ``True``, print tokens and timing information.
Default: ``False``.
kwargs: The remaining options get passed to :func:`stream_generate`.
See :func:`stream_generate` for more details.
"""
if verbose:
print("=" * 10)
text = ""
for response in stream_generate(model, tokenizer, prompt, **kwargs):
if verbose:
print(response.text, end="", flush=True)
text += response.text
if verbose:
print()
print("=" * 10)
if len(text) == 0:
print("No text generated for this prompt")
return
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 _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])
def _right_pad_prompts(prompts, max_length=None):
if max_length is None:
max_length = max(len(p) for p in prompts)
return mx.array([p + [0] * (max_length - len(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
def _make_cache(model, left_padding, max_kv_size):
"""
Convert a list of regular caches into their corresponding
batch-aware caches.
"""
def to_batch_cache(c):
if type(c) is KVCache:
return BatchKVCache(left_padding)
elif isinstance(c, ArraysCache):
c.left_padding = mx.array(left_padding)
return c
elif isinstance(c, RotatingKVCache):
if c.keep > 0:
raise ValueError("RotatingKVCache with keep tokens is not supported.")
return BatchRotatingKVCache(c.max_size, left_padding)
elif isinstance(c, CacheList):
return CacheList(*(to_batch_cache(sub_c) for sub_c in c.caches))
else:
raise ValueError(f"{type(c)} does not yet support batching")
if hasattr(model, "make_cache"):
cache = model.make_cache()
return [to_batch_cache(c) for c in cache]
else:
if max_kv_size is not None:
return [
BatchRotatingKVCache(max_kv_size, left_padding) for _ in model.layers
]
return [BatchKVCache(left_padding) for _ in model.layers]
def _merge_caches(caches):
batch_cache = []
if not caches:
return batch_cache
for i in range(len(caches[0])):
if hasattr(caches[0][i], "merge"):
batch_cache.append(caches[0][i].merge([c[i] for c in caches]))
else:
raise ValueError(
f"{type(caches[0][i])} does not yet support batching with history"
)
return batch_cache
def _extend_cache(cache_a, cache_b):
if not cache_a:
return cache_b
if not cache_b:
return cache_a
for ca, cb in zip(cache_a, cache_b):
ca.extend(cb)
return cache_a
def _build_trie(sequences):
"""Build an Aho-Corasick trie from the provided sequences
See https://en.wikipedia.org/wiki/Aho–Corasick_algorithm .
"""
trie = {}
for idx, seq in enumerate(sequences):
node = trie
try:
for tok in seq:
node = node.setdefault(tok, {})
node["__match__"] = (tuple(seq), idx)
except TypeError:
node = node.setdefault(seq, {})
node["__match__"] = ((seq,), idx)
# BFS to set failure links and propagate matches.
queue = deque()
for key, child in trie.items():
if key == "__match__":
continue
child["__fail__"] = trie
queue.append(child)
while queue:
parent = queue.popleft()
for key, child in parent.items():
if key in ("__fail__", "__match__"):
continue
queue.append(child)
fail = parent["__fail__"]
while key not in fail and fail is not trie:
fail = fail["__fail__"]
child["__fail__"] = fail[key] if key in fail else trie
if "__match__" not in child and "__match__" in child["__fail__"]:
child["__match__"] = child["__fail__"]["__match__"]
return trie
def _step_trie(node, trie, x):
"""One step in the Aho-Corasick trie."""
while x not in node and node is not trie:
node = node["__fail__"]
if x in node:
node = node[x]
return node
class StopSequenceMatcher:
"""Detect stop sequences in a stream of tokens using an Aho-Corasick trie.
Any matched sequence signals stop. Used by the batch generator for EOS and
stop word detection.
"""
def __init__(self, stop_sequences=None):
self._trie = _build_trie(stop_sequences) if stop_sequences else {}
def __deepcopy__(self, memo):
new = object.__new__(StopSequenceMatcher)
new._trie = self._trie
return new
def make_state(self):
return self._trie
@staticmethod
def match(state, trie, x):
"""Advance by one token. Returns (new_state, matched)."""
node = _step_trie(state, trie, x)
return node, node.get("__match__") is not None
class TextStateMachine:
"""A state machine that matches decoded text to track state transitions
(reasoning, tool calling) and strip the matched control sequences from the
output.
Transitions are provided as state -> [(text, new_state)]. Matching on text
rather than token ids is robust to tokenization differences (e.g. a
marker's trailing ``>`` being merged with the following byte).
The runtime state carries a buffer holding text that might be part of a
control sequence. Text is only emitted once it is known not to be part of
any match.
Example:
sm = TextStateMachine(
transitions={
"normal": [("", "reasoning"), ("", "tool")],
"reasoning": [("", "normal")],
"tool": [("", "normal")],
},
)
state = sm.make_state(initial="normal")
"""
def __init__(self, transitions=None):
self._states = {}
for src, edges in (transitions or {}).items():
strings, dst = zip(*edges) if edges else ([], [])
self._states[src] = (_build_trie(strings), dst)
def make_state(self, initial="normal"):
"""Create a fresh runtime state (state_name, trie_node, states, buffer)."""
if initial not in self._states:
self._states[initial] = (_build_trie([]), [])
return (initial, self._states[initial][0], self._states, "")
@staticmethod
def step(state, text):
"""Consume a chunk of decoded text.
Returns (new_state, emittable_text, current_state_name) where
emittable_text is the text safe to show (control sequences stripped,
possible partial matches held back in the buffer).
"""
s, n, states, buf = state
buf += text
trie = states[s][0]
emittable = ""
# buf[:consumed] has been emitted or discarded; buf[consumed:] pending.
consumed = 0
for i in range(len(buf)):
ch = buf[i]
while ch not in n and n is not trie:
n = n["__fail__"]
if ch in n:
n = n[ch]
match = n.get("__match__")
if match is not None:
match_start = i + 1 - len(match[0])
emittable += buf[consumed:match_start]
consumed = i + 1
s = states[s][1][match[1]]
if s is None:
return (s, None, states, buf[consumed:]), emittable, s
trie = states[s][0]
n = trie
elif n is trie:
# At the root: no partial match in progress, everything is safe.
emittable += buf[consumed : i + 1]
consumed = i + 1
return (s, n, states, buf[consumed:]), emittable, s
@staticmethod
def flush(state):
"""Emit the remaining buffer (use on finish_reason="length")."""
s, n, states, buf = state
trie = states[s][0] if s is not None else None
return (s, trie, states, ""), buf, s
@staticmethod
def discard(state):
"""Drop the remaining buffer (use on finish_reason="stop")."""
s, n, states, buf = state
trie = states[s][0] if s is not None else None
return (s, trie, states, ""), s
def make_stop_matcher(tokenizer, stop_words=None):
"""Build a StopSequenceMatcher from EOS tokens and stop words."""
stop_sequences = [(t,) for t in tokenizer.eos_token_ids]
for w in stop_words or []:
stop_sequences.append(tuple(tokenizer.encode(w, add_special_tokens=False)))
return StopSequenceMatcher(stop_sequences)
def make_text_state_machine(tokenizer, stop_words=None):
"""Build a TextStateMachine with reasoning/tool transitions and stop words.
Stop words are added as self-transitions in every state so they are
stripped from the output without changing state.
"""
transitions = {}
if tokenizer.has_thinking:
transitions.setdefault("normal", []).append(
(tokenizer.think_start, "reasoning")
)
transitions["reasoning"] = [(tokenizer.think_end, "normal")]
if tokenizer.has_tool_calling:
transitions.setdefault("normal", []).append((tokenizer.tool_call_start, "tool"))
if tokenizer.has_thinking:
transitions["reasoning"].append((tokenizer.tool_call_start, "tool"))
transitions["tool"] = (
[(tokenizer.tool_call_end, "normal")] if tokenizer.tool_call_end else []
)
if stop_words:
for state_name in set(transitions) | {"normal"}:
for w in stop_words:
transitions.setdefault(state_name, []).append((w, state_name))
return TextStateMachine(transitions or None)
class PromptProcessingBatch:
"""
A batch processor for prompt tokens with support for incremental processing.
This class handles batched prompt processing, managing KV caches and preparing
tokens for generation. It supports extending, filtering, and splitting batches.
"""
@dataclass
class Response:
uid: int
progress: tuple
end_of_segment: bool
end_of_prompt: bool
def __init__(
self,
model: nn.Module,
uids: List[int],
caches: List[List[Any]],
tokens: Optional[List[List[int]]] = None,
prefill_step_size: int = 2048,
samplers: Optional[List[Callable[[mx.array], mx.array]]] = None,
fallback_sampler: Optional[Callable[[mx.array], mx.array]] = None,
logits_processors: Optional[
List[List[Callable[[mx.array, mx.array], mx.array]]]
] = None,
stop_matchers: Optional[List[StopSequenceMatcher]] = None,
max_tokens: Optional[List[int]] = None,
):
self.model = model
self.uids = uids
self.prompt_cache = _merge_caches(caches)
self.tokens = tokens if tokens is not None else [[] for _ in uids]
self.prefill_step_size = prefill_step_size
self.samplers = samplers if samplers is not None else []
self.fallback_sampler = fallback_sampler or (lambda x: mx.argmax(x, axis=-1))
self.logits_processors = (
logits_processors if logits_processors is not None else []
)
self.stop_matchers = (
stop_matchers
if stop_matchers is not None
else [StopSequenceMatcher()] * len(uids)
)
self.max_tokens = (
max_tokens
if max_tokens is not None
else [DEFAULT_MAX_TOKENS] * len(self.uids)
)
def __len__(self):
return len(self.uids)
def extract_cache(self, idx: int) -> List[Any]:
return [c.extract(idx) for c in self.prompt_cache]
def extend(self, batch):
if not any(self.samplers):
self.samplers = [None] * len(self.uids)
if not any(self.logits_processors):
self.logits_processors = [None] * len(self.uids)
samplers = batch.samplers if any(batch.samplers) else [None] * len(batch.uids)
logits_processors = (
batch.logits_processors
if any(batch.logits_processors)
else [None] * len(batch.uids)
)
self.uids.extend(batch.uids)
self.prompt_cache = _extend_cache(self.prompt_cache, batch.prompt_cache)
self.tokens.extend(batch.tokens)
self.samplers.extend(samplers)
self.logits_processors.extend(logits_processors)
self.max_tokens.extend(batch.max_tokens)
self.stop_matchers.extend(batch.stop_matchers)
def _copy(self):
new_batch = self.__class__.__new__(self.__class__)
new_batch.model = self.model
new_batch.uids = list(self.uids)
new_batch.prompt_cache = copy.deepcopy(self.prompt_cache)
new_batch.tokens = list(self.tokens)
new_batch.prefill_step_size = self.prefill_step_size
new_batch.samplers = list(self.samplers)
new_batch.fallback_sampler = self.fallback_sampler
new_batch.logits_processors = list(self.logits_processors)
new_batch.stop_matchers = list(self.stop_matchers)
new_batch.max_tokens = list(self.max_tokens)
return new_batch
def split(self, indices: List[int]):
indices = sorted(indices)
indices_left = sorted(set(range(len(self.uids))) - set(indices))
new_batch = self._copy()
self.filter(indices_left)
new_batch.filter(indices)
return new_batch
def filter(self, keep: List[int]):
self.uids = [self.uids[idx] for idx in keep]
if not keep:
self.prompt_cache.clear()
else:
for c in self.prompt_cache:
c.filter(keep)
self.tokens = [self.tokens[idx] for idx in keep]
if any(self.samplers):
self.samplers = [self.samplers[idx] for idx in keep]
else:
self.samplers = [None] * len(keep)
if any(self.logits_processors):
self.logits_processors = [self.logits_processors[idx] for idx in keep]
else:
self.logits_processors = [[]] * len(keep)
self.max_tokens = [self.max_tokens[idx] for idx in keep]
self.stop_matchers = [self.stop_matchers[idx] for idx in keep]
def prompt(self, tokens: List[List[int]]):
"""
Process prompt tokens through the model.
Args:
tokens: List of token sequences to process.
"""
if len(self.uids) != len(tokens):
raise ValueError("The batch length doesn't match the number of inputs")
if not tokens:
return
# Add the tokens to the self.tokens so they represent the tokens
# contained in the KV Cache.
for sti, ti in zip(self.tokens, tokens):
sti += ti
# Calculate if we need to pad
lengths = [len(p) for p in tokens]
max_length = max(lengths)
padding = [max_length - l for l in lengths]
max_padding = max(padding)
# Prepare the caches and inputs. Right pad if needed otherwise just
# cast to array.
if max_padding > 0:
tokens = _right_pad_prompts(tokens, max_length=max_length)
for c in self.prompt_cache:
c.prepare(lengths=lengths, right_padding=padding)
else:
tokens = mx.array(tokens)
# Actual prompt processing loop
while tokens.shape[1] > 0:
n_to_process = min(self.prefill_step_size, tokens.shape[1])
self.model(tokens[:, :n_to_process], cache=self.prompt_cache)
mx.eval([c.state for c in self.prompt_cache])
mx.clear_cache()
tokens = tokens[:, n_to_process:]
# Finalize the cache if there was any padding
if max_padding > 0:
for c in self.prompt_cache:
c.finalize()
mx.eval([c.state for c in self.prompt_cache])
mx.clear_cache()
def generate(self, tokens: List[List[int]]):
"""
Transition from prompt processing to generation.
Args:
tokens: Final tokens for each sequence to start generation.
Returns:
A GenerationBatch ready for token generation.
"""
if any(len(t) > 1 for t in tokens):
self.prompt([t[:-1] for t in tokens])
last_token = mx.array([t[-1] for t in tokens])
generation = GenerationBatch(
self.model,
self.uids,
last_token,
self.prompt_cache,
self.tokens,
self.samplers,
self.fallback_sampler,
self.logits_processors,
self.stop_matchers,
self.max_tokens,
)
self.uids = []
self.prompt_cache = []
self.tokens = []
self.samplers = []
self.logits_processors = []
self.max_tokens = []
return generation
@classmethod
def empty(
cls,
model: nn.Module,
fallback_sampler: Callable[[mx.array], mx.array],
prefill_step_size: int = 2048,
):
return cls(
model=model,
fallback_sampler=fallback_sampler,
prefill_step_size=prefill_step_size,
uids=[],
caches=[],
tokens=[],
samplers=[],
logits_processors=[],
max_tokens=[],
stop_matchers=[],
)
class GenerationBatch:
"""
A batched token generator that manages multiple sequences in parallel.
This class handles the generation phase after prompt processing, managing
KV caches, sampling, and stop sequence detection for multiple sequences.
"""
@dataclass
class Response:
uid: int
token: int
logprobs: mx.array
finish_reason: Optional[str]
prompt_cache: Optional[List[Any]]
all_tokens: Optional[List[int]]
def __init__(
self,
model: nn.Module,
uids: List[int],
inputs: mx.array,
prompt_cache: List[Any],
tokens: List[List[int]],
samplers: Optional[List[Callable[[mx.array], mx.array]]],
fallback_sampler: Callable[[mx.array], mx.array],
logits_processors: Optional[
List[List[Callable[[mx.array, mx.array], mx.array]]]
],
stop_matchers: List[StopSequenceMatcher],
max_tokens: List[int],
):
self.model = model
self.uids = uids
self.prompt_cache = prompt_cache
self.tokens = tokens
self.samplers = samplers
self.fallback_sampler = fallback_sampler
self.logits_processors = logits_processors
self.stop_matchers = stop_matchers
self.max_tokens = max_tokens
if self.samplers and len(self.samplers) != len(self.uids):
raise ValueError("Insufficient number of samplers provided")
if self.logits_processors and len(self.logits_processors) != len(self.uids):
raise ValueError("Insufficient number of logits_processors provided")
self._current_tokens = None
self._current_logprobs = []
self._next_tokens = inputs
self._next_logprobs = []
self._token_context = [TokenBuffer(t) for t in tokens]
self._num_tokens = [0] * len(self.uids)
self._matcher_states = [m.make_state() for m in stop_matchers]
if self.uids:
self._step()
def __len__(self):
return len(self.uids)
def extend(self, batch):
"""Extend this batch with another generation batch."""
self.uids.extend(batch.uids)
self.prompt_cache = _extend_cache(self.prompt_cache, batch.prompt_cache)
self.tokens.extend(batch.tokens)
self.samplers.extend(batch.samplers)
self.logits_processors.extend(batch.logits_processors)
self.max_tokens.extend(batch.max_tokens)
self.stop_matchers.extend(batch.stop_matchers)
if self._current_tokens is None:
self._current_tokens = batch._current_tokens
self._current_logprobs = batch._current_logprobs
elif batch._current_tokens is not None:
self._current_tokens = mx.concatenate(
[self._current_tokens, batch._current_tokens]
)
self._current_logprobs.extend(batch._current_logprobs)
if self._next_tokens is None:
self._next_tokens = batch._next_tokens
self._next_logprobs = batch._next_logprobs
elif batch._next_tokens is not None:
self._next_tokens = mx.concatenate([self._next_tokens, batch._next_tokens])
self._next_logprobs.extend(batch._next_logprobs)
self._token_context.extend(batch._token_context)
self._num_tokens.extend(batch._num_tokens)
self._matcher_states.extend(batch._matcher_states)
def _step(self) -> Tuple[List[int], List[mx.array]]:
"""
Perform a single generation step.
Returns:
Tuple of token list and logprobs list.
"""
self._current_tokens = self._next_tokens
self._current_logprobs = self._next_logprobs
inputs = self._current_tokens
# Forward pass
logits = self.model(inputs[:, None], cache=self.prompt_cache)
logits = logits[:, -1, :]
# Logits processors
token_context = []
if any(self.logits_processors):
# Update the token context that will be used by the logits processors
token_context = [
tc.update_and_fetch(inputs[i : i + 1])
for i, tc in enumerate(self._token_context)
]
processed_logits = []
for e in range(len(self.uids)):
sample_logits = logits[e : e + 1]
for processor in self.logits_processors[e]:
sample_logits = processor(token_context[e], sample_logits)
processed_logits.append(sample_logits)
logits = mx.concatenate(processed_logits, axis=0)
# Normalize the logits
logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
# Sample
if any(self.samplers):
all_samples = []
for e in range(len(self.uids)):
sample_sampler = self.samplers[e] or self.fallback_sampler
sampled = sample_sampler(logprobs[e : e + 1])
all_samples.append(sampled)
sampled = mx.concatenate(all_samples, axis=0)
else:
sampled = self.fallback_sampler(logprobs)
# Assign the next step to member variables and start computing it
# asynchronously
self._next_tokens = sampled
self._next_logprobs = list(logprobs)
mx.async_eval(self._next_tokens, self._next_logprobs, token_context)
# Eval the current tokens and current logprobs. After that also add
# them to self.tokens so that it always represents the tokens contained
# in the KV Cache.
mx.eval(inputs, self._current_logprobs)
inputs = inputs.tolist()
for sti, ti in zip(self.tokens, inputs):
sti.append(ti)
return inputs, self._current_logprobs
def extract_cache(self, idx: int) -> List[Any]:
return [c.extract(idx) for c in self.prompt_cache]
def filter(self, keep: List[int]):
"""Filter the batch to keep only the specified indices."""
self.uids = [self.uids[idx] for idx in keep]
if not keep:
self.prompt_cache.clear()
else:
for c in self.prompt_cache:
c.filter(keep)
self.tokens = [self.tokens[idx] for idx in keep]
if any(self.samplers):
self.samplers = [self.samplers[idx] for idx in keep]
if any(self.logits_processors):
self.logits_processors = [self.logits_processors[idx] for idx in keep]
self.max_tokens = [self.max_tokens[idx] for idx in keep]
self.stop_matchers = [self.stop_matchers[idx] for idx in keep]
self._next_tokens = self._next_tokens[keep] if keep else None
self._next_logprobs = [self._next_logprobs[idx] for idx in keep]
self._token_context = [self._token_context[idx] for idx in keep]
self._num_tokens = [self._num_tokens[idx] for idx in keep]
self._matcher_states = [self._matcher_states[idx] for idx in keep]
def next(self) -> List[Response]:
"""
Generate the next batch of tokens.
Returns:
List of Response objects for each sequence in the batch.
"""
if not self.uids:
return []
tokens, logprobs = self._step()
keep = []
responses = []
for i in range(len(self.uids)):
finish_reason = None
self._num_tokens[i] += 1
if self._num_tokens[i] >= self.max_tokens[i]:
finish_reason = "length"
self._matcher_states[i], matched = StopSequenceMatcher.match(
self._matcher_states[i],
self.stop_matchers[i]._trie,
tokens[i],
)
if matched:
finish_reason = "stop"
if finish_reason is not None:
responses.append(
self.Response(
uid=self.uids[i],
token=tokens[i],
logprobs=logprobs[i],
finish_reason=finish_reason,
prompt_cache=self.extract_cache(i),
all_tokens=self.tokens[i],
)
)
else:
keep.append(i)
responses.append(
self.Response(
uid=self.uids[i],
token=tokens[i],
logprobs=logprobs[i],
finish_reason=None,
prompt_cache=None,
all_tokens=None,
)
)
if len(keep) < len(self.uids):
self.filter(keep)
return responses
@classmethod
def empty(
cls,
model: nn.Module,
fallback_sampler: Callable[[mx.array], mx.array],
):
return cls(
model=model,
fallback_sampler=fallback_sampler,
uids=[],
inputs=mx.array([], dtype=mx.uint32),
prompt_cache=[],
tokens=[],
samplers=[],
logits_processors=[],
max_tokens=[],
stop_matchers=[],
)
class BatchGenerator:
"""
A batch generator implements continuous batching.
This class provides automatic management of prompt processing and generation
batches, handling the transition between the two.
It also allows for segmented prompt processing which guarantees that the
generator will stop at these boundaries when processing an input.
"""
def __init__(
self,
model: nn.Module,
*,
max_tokens: int = 128,
stop_tokens: Optional[Sequence[Sequence[int]]] = None,
sampler: Optional[Callable[[mx.array], mx.array]] = None,
logits_processors: Optional[
List[Callable[[mx.array, mx.array], mx.array]]
] = None,
completion_batch_size: int = 32,
prefill_batch_size: int = 8,
prefill_step_size: int = 2048,
max_kv_size: Optional[int] = None,
stream=None,
):
self.model = model
self.max_tokens = max_tokens
self.sampler = sampler or (lambda x: mx.argmax(x, axis=-1))
self.logits_processors = logits_processors or []
self.uid_count = 0
self.prefill_step_size = prefill_step_size
self.prefill_batch_size = prefill_batch_size
self.completion_batch_size = max(completion_batch_size, prefill_batch_size)
self.max_kv_size = max_kv_size
self._stream = stream or generation_stream
self._default_stop_matcher = StopSequenceMatcher(
stop_tokens if stop_tokens else None,
)
self._uid_count = 0
self._prompt_batch = PromptProcessingBatch.empty(
self.model,
self.sampler,
prefill_step_size=prefill_step_size,
)
self._generation_batch = GenerationBatch.empty(self.model, self.sampler)
self._unprocessed_sequences = deque()
self._currently_processing = []
self._prompt_tokens_counter = 0
self._prompt_time_counter = 0
self._gen_tokens_counter = 0
self._steps_counter = 0
if mx.metal.is_available():
self._old_wired_limit = mx.set_wired_limit(
mx.device_info()["max_recommended_working_set_size"]
)
else:
self._old_wired_limit = None
@property
def stream(self):
return self._stream
def close(self):
if self._old_wired_limit is not None:
mx.synchronize(self._stream)
mx.set_wired_limit(self._old_wired_limit)
self._old_wired_limit = None
def __del__(self):
self.close()
@contextlib.contextmanager
def stats(self, stats=None):
stats = stats or BatchStats()
self._prompt_tokens_counter = 0
self._prompt_time_counter = 0
self._gen_tokens_counter = 0
tic = time.perf_counter()
try:
yield stats
finally:
toc = time.perf_counter()
total_time = toc - tic
gen_time = total_time - self._prompt_time_counter
stats.prompt_tokens += self._prompt_tokens_counter
stats.prompt_time += self._prompt_time_counter
stats.prompt_tps = stats.prompt_tokens / stats.prompt_time
stats.generation_tokens += self._gen_tokens_counter
stats.generation_time += gen_time
stats.generation_tps = stats.generation_tokens / stats.generation_time
stats.peak_memory = max(stats.peak_memory, mx.get_peak_memory() / 1e9)
def insert(
self,
prompts: List[List[int]],
max_tokens: Optional[List[int]] = None,
caches: Optional[List[List[Any]]] = None,
all_tokens: Optional[List[List[int]]] = None,
samplers: Optional[List[Callable[[mx.array], mx.array]]] = None,
logits_processors: Optional[
List[List[Callable[[mx.array, mx.array], mx.array]]]
] = None,
stop_matchers: Optional[List[StopSequenceMatcher]] = None,
):
return self.insert_segments(
[[p] for p in prompts],
max_tokens,
caches,
all_tokens,
samplers,
logits_processors,
stop_matchers,
)
def insert_segments(
self,
segments: List[List[List[int]]],
max_tokens: Optional[List[int]] = None,
caches: Optional[List[List[Any]]] = None,
all_tokens: Optional[List[List[int]]] = None,
samplers: Optional[List[Callable[[mx.array], mx.array]]] = None,
logits_processors: Optional[
List[List[Callable[[mx.array, mx.array], mx.array]]]
] = None,
stop_matchers: Optional[List[StopSequenceMatcher]] = None,
):
uids = []
max_tokens = max_tokens or [self.max_tokens] * len(segments)
all_tokens = all_tokens or [[] for _ in segments]
samplers = samplers or [None] * len(segments)
logits_processors = logits_processors or (
[self.logits_processors] * len(segments)
)
stop_matchers = stop_matchers or ([self._default_stop_matcher] * len(segments))
caches = caches or [None] * len(segments)
for i in range(len(segments)):
if caches[i] is None:
caches[i] = self._make_new_cache()
for seq, m, c, at, s, lp, sm in zip(
segments,
max_tokens,
caches,
all_tokens,
samplers,
logits_processors,
stop_matchers,
):
seq = list(seq)
if len(seq[-1]) != 1:
seq.append(seq[-1][-1:])
seq[-2] = seq[-2][:-1]
self._unprocessed_sequences.append(
(self._uid_count, seq, m, c, at, s, lp, sm)
)
uids.append(self._uid_count)
self._uid_count += 1
return uids
def _make_new_cache(self):
if self.max_kv_size is None:
return cache.make_prompt_cache(self.model)
return [
(
RotatingKVCache(max_size=self.max_kv_size)
if isinstance(ci, KVCache)
else ci
)
for ci in cache.make_prompt_cache(self.model)
]
def _find_uids(self, uids):
uids = set(uids)
results = {}
for i, uid_i in enumerate(self._generation_batch.uids):
if uid_i in uids:
results[uid_i] = (2, i)
for i, uid_i in enumerate(self._prompt_batch.uids):
if uid_i in uids:
results[uid_i] = (1, i)
for i, seq in enumerate(self._unprocessed_sequences):
if seq[0] in uids:
results[seq[0]] = (0, i)
return results
def extract_cache(self, uids):
results = {}
for uid, (stage, idx) in self._find_uids(uids).items():
if stage == 0:
results[uid] = self._unprocessed_sequences[idx][3:5]
elif stage == 1:
results[uid] = (
self._prompt_batch.extract_cache(idx),
self._prompt_batch.tokens[idx],
)
else:
results[uid] = (
self._generation_batch.extract_cache(idx),
self._generation_batch.tokens[idx],
)
return results
def remove(self, uids, return_prompt_caches=False):
caches = {}
if return_prompt_caches:
caches = self.extract_cache(uids)
keep = (
set(range(len(self._unprocessed_sequences))),
set(range(len(self._prompt_batch))),
set(range(len(self._generation_batch))),
)
for stage, idx in self._find_uids(uids).values():
keep[stage].remove(idx)
if len(keep[0]) < len(self._unprocessed_sequences):
self._unprocessed_sequences = deque(
x for i, x in enumerate(self._unprocessed_sequences) if i in keep[0]
)
if len(keep[1]) < len(self._prompt_batch):
self._prompt_batch.filter(sorted(keep[1]))
self._currently_processing = [
x for i, x in enumerate(self._currently_processing) if i in keep[1]
]
if len(keep[2]) < len(self._generation_batch):
self._generation_batch.filter(sorted(keep[2]))
return caches
@property
def prompt_cache_nbytes(self):
total = sum(c.nbytes for p in self._unprocessed_sequences for c in p[3])
total += sum(c.nbytes for c in self._prompt_batch.prompt_cache)
total += sum(c.nbytes for c in self._generation_batch.prompt_cache)
return total
def _make_batch(self, n: int):
uids = []
caches = []
tokens = []
samplers = []
logits_processors = []
max_tokens = []
stop_matchers = []
for _ in range(n):
sequence = self._unprocessed_sequences.popleft()
uids.append(sequence[0])
caches.append(sequence[3])
tokens.append(sequence[4])
samplers.append(sequence[5])
logits_processors.append(sequence[6])
max_tokens.append(sequence[2])
stop_matchers.append(sequence[7])
self._currently_processing.append(
[sequence[1], 0, sum(len(s) for s in sequence[1])]
)
return PromptProcessingBatch(
model=self.model,
uids=uids,
caches=caches,
tokens=tokens,
prefill_step_size=self.prefill_step_size,
samplers=samplers,
fallback_sampler=self.sampler,
logits_processors=logits_processors,
stop_matchers=stop_matchers,
max_tokens=max_tokens,
)
def _next(self):
generation_responses = []
prompt_responses = []
# Generate tokens first
if len(self._generation_batch) > 0:
generation_responses = self._generation_batch.next()
self._gen_tokens_counter += len(generation_responses)
self._steps_counter += 1
if self._steps_counter % 512 == 0:
mx.clear_cache()
# Exit early because we already have our hands full with decoding
if len(self._generation_batch) >= self.completion_batch_size:
return prompt_responses, generation_responses
# Check if we have sequences and add them to the prompt batch
n = min(
self.prefill_batch_size - len(self._prompt_batch),
self.completion_batch_size - len(self._generation_batch),
len(self._unprocessed_sequences),
)
if n > 0:
self._prompt_batch.extend(self._make_batch(n))
# Split the prompt sequences to the ones moving to generation and the rest
keep = []
split = []
for i, seq in enumerate(self._currently_processing):
segments = seq[0]
if len(segments) == 1 and len(segments[0]) == 1:
split.append(i)
else:
keep.append(i)
# Actually split off part of the prompt batch and start generation
if split:
last_inputs = [self._currently_processing[i][0][0] for i in split]
progress = [(self._currently_processing[i][2],) * 2 for i in split]
self._currently_processing = [self._currently_processing[i] for i in keep]
gen_batch = self._prompt_batch.split(split).generate(last_inputs)
for i, p in enumerate(progress):
prompt_responses.append(
PromptProcessingBatch.Response(
gen_batch.uids[i],
p,
True,
True,
)
)
self._generation_batch.extend(gen_batch)
# Extract the next prompts input
prompts = []
for i, seq in enumerate(self._currently_processing):
response = PromptProcessingBatch.Response(
self._prompt_batch.uids[i], 0, False, False
)
segments = seq[0]
n = min(len(segments[0]), self.prefill_step_size)
prompts.append(segments[0][:n])
segments[0] = segments[0][n:]
if len(segments[0]) == 0:
segments.pop(0)
response.end_of_segment = True
seq[1] += len(prompts[-1])
response.progress = (seq[1], seq[2])
prompt_responses.append(response)
# Process the prompts
self._prompt_tokens_counter += sum(len(p) for p in prompts)
tic = time.perf_counter()
self._prompt_batch.prompt(prompts)
toc = time.perf_counter()
self._prompt_time_counter += toc - tic
return prompt_responses, generation_responses
def next(self):
"""
Get the next batch of responses.
Returns:
Tuple of prompt processing responses and generation responses.
"""
with mx.stream(self._stream):
return self._next()
def next_generated(self):
"""
Return only generated tokens ignoring batch generation responses.
Returns:
List of GenerationBatch.Response objects
"""
with mx.stream(self._stream):
while True:
prompt_responses, generation_responses = self._next()
if not generation_responses and prompt_responses:
continue
return generation_responses
@dataclass
class BatchResponse:
"""
A data object to hold a batch generation response.
Args:
texts: (List[str]): The generated text for each prompt.
stats (BatchStats): Statistics about the generation.
caches: Optional prompt caches for each sequence.
token_ids (Optional[List[List[int]]]): The generated token IDs for each
prompt. Only present when ``return_token_ids=True``.
logprobs (Optional[List[List[float]]]): The per-token log-probabilities
of the sampled tokens for each prompt. Only present when
``return_logprobs=True``.
"""
texts: List[str]
stats: BatchStats
caches: Optional[List[List[Any]]]
token_ids: Optional[List[List[int]]] = None
logprobs: Optional[List[List[float]]] = None
def batch_generate(
model,
tokenizer,
prompts: List[List[int]],
prompt_caches: Optional[List[List[Any]]] = None,
max_tokens: Union[int, List[int]] = 128,
verbose: bool = False,
return_prompt_caches: bool = False,
return_token_ids: bool = False,
return_logprobs: bool = False,
**kwargs,
) -> BatchResponse:
"""
Generate responses for the given batch of prompts.
Args:
model (nn.Module): The language model.
tokenizer (PreTrainedTokenizer): The tokenizer.
prompts (List[List[int]]): The input prompts.
prompt_caches (List[List[Any]], optional): Pre-computed prompt-caches
for each input prompt. Note, unlike ``generate_step``, the caches
won't be updated in-place.
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.
return_prompt_caches (bool): Return the prompt caches in the batch
responses. Default: ``False``.
return_token_ids (bool): Return the generated token IDs in the batch
responses. Default: ``False``.
return_logprobs (bool): Return the per-token log-probability of the
sampled token for each generated token. Useful for reinforcement
learning (e.g. RLOO, PPO) where behavior log-probabilities are needed
for importance weighting. Default: ``False``.
kwargs: The remaining options get passed to :obj:`BatchGenerator`.
See :obj:`BatchGenerator` for more details.
"""
gen = BatchGenerator(
model,
stop_tokens=[[t] for t in tokenizer.eos_token_ids],
**kwargs,
)
num_samples = len(prompts)
fin = 0
if verbose:
print(f"[batch_generate] Finished processing 0/{num_samples} ...", end="\r")
if isinstance(max_tokens, int):
max_tokens = [max_tokens] * len(prompts)
uids = gen.insert(prompts, max_tokens, caches=prompt_caches)
results = {uid: [] for uid in uids}
logprob_results = {uid: [] for uid in uids} if return_logprobs else None
prompt_caches = {}
with gen.stats() as stats:
while responses := gen.next_generated():
for r in responses:
if r.finish_reason is not None:
if return_prompt_caches:
prompt_caches[r.uid] = r.prompt_cache
if verbose:
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 return_logprobs:
logprob_results[r.uid].append(r.logprobs[r.token].item())
gen.close()
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]
caches = [prompt_caches[uid] for uid in uids] if return_prompt_caches else None
token_ids = [results[uid] for uid in uids] if return_token_ids else None
logprobs = [logprob_results[uid] for uid in uids] if return_logprobs else None
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, caches, token_ids, logprobs)
def main():
parser = setup_arg_parser()
args = parser.parse_args()
if args.seed is not None:
mx.random.seed(args.seed)
# 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,
)
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"])
)
tokenizer_config["trust_remote_code"] = args.trust_remote_code
model_path = args.model
if using_cache:
if model_path is None:
model_path = metadata["model"]
elif model_path != metadata["model"]:
raise ValueError(
f"Providing a different model ({model_path}) than that "
f"used to create the prompt cache ({metadata['model']}) "
"is an error."
)
model_path = model_path or DEFAULT_MODEL
model_config = {"quantize_activations": args.quantize_activations}
if args.flash_head is not None:
model_config["use_flash_head"] = args.flash_head
model, tokenizer = load(
model_path,
adapter_path=args.adapter_path,
tokenizer_config=tokenizer_config,
model_config=model_config,
trust_remote_code=args.trust_remote_code,
)
for eos_token in args.extra_eos_token:
tokenizer.add_eos_token(eos_token)
template_kwargs = {}
if args.chat_template_config is not None:
template_kwargs = json.loads(args.chat_template_config)
prompt = args.prompt.replace("\\n", "\n").replace("\\t", "\t")
prompt = sys.stdin.read() if prompt == "-" else prompt
if not args.ignore_chat_template and tokenizer.has_chat_template:
if args.system_prompt is not None:
messages = [{"role": "system", "content": args.system_prompt}]
else:
messages = []
messages.append({"role": "user", "content": prompt})
has_prefill = args.prefill_response is not None
if has_prefill:
messages.append({"role": "assistant", "content": args.prefill_response})
prompt = tokenizer.apply_chat_template(
messages,
tokenize=False,
continue_final_message=has_prefill,
add_generation_prompt=not has_prefill,
**template_kwargs,
)
# Treat the prompt as a suffix assuming that the prefix is in the
# stored kv cache.
if using_cache:
messages[-1]["content"] = ""
test_prompt = tokenizer.apply_chat_template(
messages,
tokenize=False,
continue_final_message=has_prefill,
add_generation_prompt=not has_prefill,
)
prompt = prompt[test_prompt.index("") :]
prompt = tokenizer.encode(prompt, add_special_tokens=False)
else:
prompt = tokenizer.encode(prompt)
if args.draft_model is not None:
draft_model, draft_tokenizer = load(args.draft_model)
if draft_tokenizer.vocab_size != tokenizer.vocab_size:
raise ValueError("Draft model tokenizer does not match model tokenizer.")
else:
draft_model = None
sampler = make_sampler(
args.temp,
args.top_p,
args.min_p,
args.min_tokens_to_keep,
top_k=args.top_k,
xtc_probability=args.xtc_probability,
xtc_threshold=args.xtc_threshold,
xtc_special_tokens=tokenizer.encode("\n") + list(tokenizer.eos_token_ids),
)
response = generate(
model,
tokenizer,
prompt,
max_tokens=args.max_tokens,
verbose=args.verbose,
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,
draft_model=draft_model,
num_draft_tokens=args.num_draft_tokens,
)
if not args.verbose:
print(response)
if __name__ == "__main__":
print(
"Calling `python -m mlx_lm.generate...` directly is deprecated."
" Use `mlx_lm.generate...` or `python -m mlx_lm generate ...` instead."
)
main()