mlx

AsyncMlxLM

Bases: AsyncLM

Asynchronous MLX-based language model wrapper.

This class provides an async interface to MLX language models with automatic batching, caching, and KV cache management. It extends AsyncLM to provide efficient batched inference with prefix caching.

The model automatically batches concurrent requests, continues the live KV cache when a batch extends the previous one, and optionally caches computed log probabilities for reuse.

Source code in genlm/backend/llm/mlx.py

class AsyncMlxLM(AsyncLM):
    """Asynchronous MLX-based language model wrapper.

    This class provides an async interface to MLX language models with
    automatic batching, caching, and KV cache management. It extends
    AsyncLM to provide efficient batched inference with prefix caching.

    The model automatically batches concurrent requests, continues the live
    KV cache when a batch extends the previous one, and optionally caches
    computed log probabilities for reuse.
    """

    def __init__(
        self,
        mlx_lm_model,
        tokenizer,
        timeout=0.0,
        prefill_step_size=2048,
        cache_size=0,
        cache_opts=None,
    ):
        """Initialize an `AsyncMlxLM` instance.

        Args:
            mlx_lm_model: The MLX language model instance.
            tokenizer: The tokenizer for encoding/decoding text.
            timeout (float, optional): Seconds to wait once for late concurrent
                requests before processing a batch. Defaults to 0.
            prefill_step_size (int, optional): Number of tokens to process
                per step during prompt prefilling.
            cache_size (int, optional): Maximum size of the output cache. If 0,
                caching is disabled. Defaults to 0.
            cache_opts (dict, optional): Additional options to pass to the
                [`OutputCache`][genlm.backend.cache.OutputCache] constructor.
                Defaults to None (no extra options).
        """
        self.mlx_lm_model = mlx_lm_model
        self.timeout = timeout
        self.prefill_step_size = prefill_step_size
        self._adapters = _Adapters(mlx_lm_model)
        # KV rows never cross adapters, so each adapter keeps its own pool.
        self._slots = defaultdict(
            partial(_SlotPool, mlx_lm_model, prefill_step_size)
        )
        self.cache = (
            OutputCache(maxsize=cache_size, **(cache_opts or {}))
            if cache_size > 0
            else None
        )
        # mlx_lm.load returns a TokenizerWrapper; AsyncLM needs the HuggingFace
        # tokenizer it wraps.
        super().__init__(tokenizer=getattr(tokenizer, "_tokenizer", tokenizer))

    @classmethod
    def from_name(cls, model_name, **kwargs):
        """Create an `AsyncMlxLM` instance from a model name.

        Args:
            model_name (str): Name of the model to load. Can be a Hugging Face
                model identifier or local path.
            **kwargs: Additional arguments passed to `AsyncMlxLM` constructor,
                such as `timeout`, `prefill_step_size`, `cache_size`.

        Returns:
            AsyncMlxLM: An `AsyncMlxLM` instance with the loaded model and tokenizer.
        """
        model, tokenizer = mlx_lm.load(model_name)
        return cls(model, tokenizer, **kwargs)

    def add_new_lora(self, lora_path, lora_name="lora_1"):
        """Register the adapter at ``lora_path`` under ``lora_name``.

        Re-registering a name rebinds it. Every adapter on a model must share the
        layout the first one wrapped it with.

        Args:
            lora_path (str): Directory holding ``adapter_config.json`` and
                ``adapters.safetensors``.
            lora_name (str, optional): Name to select the adapter by.
        """
        self._adapters.add(lora_name, lora_path)
        self.clear_cache()

    def remove_lora(self, lora_name):
        """Unregister ``lora_name``.

        Args:
            lora_name (str): Name of the adapter to remove.
        """
        self._adapters.remove(lora_name)
        self.clear_cache()

    def clear_cache(self):
        """Clear the output cache, the live KV rows, and the MLX device cache."""
        if self.cache is not None:
            self.cache.clear()
        self._slots.clear()
        mx.clear_cache()

    def _forward(self, prompts, lora_name):
        """Next-token log-probs under one adapter, ``[len(prompts), vocab]``."""
        with _wired(self.mlx_lm_model):
            self._adapters.select(lora_name)
            logits = self._slots[lora_name].logits(prompts).astype(mx.float32)
            logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
            mx.eval(logprobs)
        return _to_torch(logprobs)

    def _resolve(self, keys):
        """``{key: logprobs}`` for distinct ``(context, lora_name)`` keys.

        Uncached keys are forwarded one batch per adapter, then memoized.
        """
        out, todo = {}, defaultdict(list)
        for key in keys:
            if self.cache is not None and key in self.cache:
                out[key] = self.cache[key]
            else:
                todo[key[1]].append(key)
        for lora_name, batch in todo.items():
            logprobs = self._forward([list(k[0]) for k in batch], lora_name)
            for row, key in enumerate(batch):
                out[key] = logprobs[row]
                if self.cache is not None:
                    self.cache[key] = logprobs[row]
        return out

    def _batch_evaluate(self, queries):
        """Resolve a batch of queries, deduplicating equal keys.

        Every future gets its row or a failure; a failed forward would
        otherwise leave the co-callers awaiting.
        """
        if not queries:
            return
        futures = defaultdict(list)
        for key, future in queries:
            futures[key].append(future)
        try:
            logprobs = self._resolve(list(futures))
        except Exception as exc:
            fail_futures(queries, exc)
            raise
        except BaseException as exc:
            fail_futures(queries, batch_abandoned(exc))
            raise
        for key, waiting in futures.items():
            for future in waiting:
                future.set_result(logprobs[key])

    async def next_token_logprobs(self, token_ids, lora_name=None):
        """Request log probabilities of next token. This version is asynchronous because it automatically batches concurrent requests; use with `await`.

        Args:
            token_ids (list[int]): a list of token ids, representing a prompt to the language model.
            lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

        Returns:
            logprobs (torch.Tensor): a tensor of with the language model's log (normalized) probabilities for the next token following the prompt.
        """
        if not token_ids:
            raise ValueError("Token ids must not be empty")
        key = (tuple(token_ids), lora_name)
        if self.cache is not None and key in self.cache:
            return self.cache[key]
        future = asyncio.get_running_loop().create_future()
        batch = await self._join_batch([(key, future)], timeout=self.timeout)
        if batch is not None:
            self._batch_evaluate(batch)
        return await future

    def next_token_logprobs_sync(self, token_ids, lora_name=None):
        """Request log probabilities of next token synchronously.

        Args:
            token_ids (list[int]): A list of token IDs, representing a prompt to the language model.
            lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

        Returns:
            (torch.Tensor): Normalized log probability tensor.
        """
        if not token_ids:
            raise ValueError("Token ids must not be empty")
        key = (tuple(token_ids), lora_name)
        return self._resolve([key])[key]

    def batch_next_token_logprobs_sync(self, token_ids_list, lora_name=None):
        """Next-token log-probs for each sequence, in one batched forward.

        Args:
            token_ids_list (list[list[int]]): Prompts' token ids.
            lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

        Returns:
            (torch.Tensor): ``[len(token_ids_list), vocab]`` log-probabilities.
        """
        if any(not ids for ids in token_ids_list):
            raise ValueError("Token ids must not be empty")
        keys = [(tuple(ids), lora_name) for ids in token_ids_list]
        logprobs = self._resolve(list(dict.fromkeys(keys)))
        return torch.stack([logprobs[k] for k in keys])

    async def sample(
        self,
        prompt_token_ids,
        max_tokens,
        eos_token_ids,
        temperature=1.0,
        seed=None,
        lora_name=None,
    ):
        """Sample from the language model.

        Args:
            prompt_token_ids (list[int]): The token IDs of the prompt to
                start generation from.
            max_tokens (int): The maximum number of tokens to generate.
            eos_token_ids (list[int]): The token IDs that signal
                end-of-sequence. Generation stops when one of these is
                sampled.
            temperature (float, optional): The temperature to use for
                sampling. Higher values make the distribution more uniform,
                lower values make it more peaked. Defaults to 1.0.
            seed (int, optional): The seed for the random number generator.
                If provided, sets the random seed before sampling.
                Defaults to None.
            lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

        Returns:
            (list[int]): The sampled token IDs.
        """
        self._adapters.select(lora_name)
        if seed is not None:
            mx.random.seed(seed)

        generated = []
        for token, _ in generate_step(
            mx.array(prompt_token_ids),
            self.mlx_lm_model,
            max_tokens=max_tokens,
            sampler=make_sampler(temp=temperature),
        ):
            if token in eos_token_ids:
                break
            generated.append(token)
        return generated

__init__(mlx_lm_model, tokenizer, timeout=0.0, prefill_step_size=2048, cache_size=0, cache_opts=None)

Initialize an AsyncMlxLM instance.

Parameters:

Name	Type	Description	Default
mlx_lm_model		The MLX language model instance.	required
tokenizer		The tokenizer for encoding/decoding text.	required
timeout	float	Seconds to wait once for late concurrent requests before processing a batch. Defaults to 0.	0.0
prefill_step_size	int	Number of tokens to process per step during prompt prefilling.	2048
cache_size	int	Maximum size of the output cache. If 0, caching is disabled. Defaults to 0.	0
cache_opts	dict	Additional options to pass to the OutputCache constructor. Defaults to None (no extra options).	None

Source code in genlm/backend/llm/mlx.py

def __init__(
    self,
    mlx_lm_model,
    tokenizer,
    timeout=0.0,
    prefill_step_size=2048,
    cache_size=0,
    cache_opts=None,
):
    """Initialize an `AsyncMlxLM` instance.

    Args:
        mlx_lm_model: The MLX language model instance.
        tokenizer: The tokenizer for encoding/decoding text.
        timeout (float, optional): Seconds to wait once for late concurrent
            requests before processing a batch. Defaults to 0.
        prefill_step_size (int, optional): Number of tokens to process
            per step during prompt prefilling.
        cache_size (int, optional): Maximum size of the output cache. If 0,
            caching is disabled. Defaults to 0.
        cache_opts (dict, optional): Additional options to pass to the
            [`OutputCache`][genlm.backend.cache.OutputCache] constructor.
            Defaults to None (no extra options).
    """
    self.mlx_lm_model = mlx_lm_model
    self.timeout = timeout
    self.prefill_step_size = prefill_step_size
    self._adapters = _Adapters(mlx_lm_model)
    # KV rows never cross adapters, so each adapter keeps its own pool.
    self._slots = defaultdict(
        partial(_SlotPool, mlx_lm_model, prefill_step_size)
    )
    self.cache = (
        OutputCache(maxsize=cache_size, **(cache_opts or {}))
        if cache_size > 0
        else None
    )
    # mlx_lm.load returns a TokenizerWrapper; AsyncLM needs the HuggingFace
    # tokenizer it wraps.
    super().__init__(tokenizer=getattr(tokenizer, "_tokenizer", tokenizer))

from_name(model_name, **kwargs) classmethod

Create an AsyncMlxLM instance from a model name.

Parameters:

Name	Type	Description	Default
model_name	str	Name of the model to load. Can be a Hugging Face model identifier or local path.	required
**kwargs		Additional arguments passed to AsyncMlxLM constructor, such as timeout, prefill_step_size, cache_size.	{}

Returns:

Name	Type	Description
AsyncMlxLM		An AsyncMlxLM instance with the loaded model and tokenizer.

Source code in genlm/backend/llm/mlx.py

@classmethod
def from_name(cls, model_name, **kwargs):
    """Create an `AsyncMlxLM` instance from a model name.

    Args:
        model_name (str): Name of the model to load. Can be a Hugging Face
            model identifier or local path.
        **kwargs: Additional arguments passed to `AsyncMlxLM` constructor,
            such as `timeout`, `prefill_step_size`, `cache_size`.

    Returns:
        AsyncMlxLM: An `AsyncMlxLM` instance with the loaded model and tokenizer.
    """
    model, tokenizer = mlx_lm.load(model_name)
    return cls(model, tokenizer, **kwargs)

add_new_lora(lora_path, lora_name='lora_1')

Register the adapter at lora_path under lora_name.

Re-registering a name rebinds it. Every adapter on a model must share the layout the first one wrapped it with.

Parameters:

Name	Type	Description	Default
lora_path	str	Directory holding adapter_config.json and adapters.safetensors.	required
lora_name	str	Name to select the adapter by.	'lora_1'

Source code in genlm/backend/llm/mlx.py

def add_new_lora(self, lora_path, lora_name="lora_1"):
    """Register the adapter at ``lora_path`` under ``lora_name``.

    Re-registering a name rebinds it. Every adapter on a model must share the
    layout the first one wrapped it with.

    Args:
        lora_path (str): Directory holding ``adapter_config.json`` and
            ``adapters.safetensors``.
        lora_name (str, optional): Name to select the adapter by.
    """
    self._adapters.add(lora_name, lora_path)
    self.clear_cache()

remove_lora(lora_name)

Unregister lora_name.

Parameters:

Name	Type	Description	Default
lora_name	str	Name of the adapter to remove.	required

Source code in genlm/backend/llm/mlx.py

def remove_lora(self, lora_name):
    """Unregister ``lora_name``.

    Args:
        lora_name (str): Name of the adapter to remove.
    """
    self._adapters.remove(lora_name)
    self.clear_cache()

clear_cache()

Clear the output cache, the live KV rows, and the MLX device cache.

Source code in genlm/backend/llm/mlx.py

def clear_cache(self):
    """Clear the output cache, the live KV rows, and the MLX device cache."""
    if self.cache is not None:
        self.cache.clear()
    self._slots.clear()
    mx.clear_cache()

next_token_logprobs(token_ids, lora_name=None) async

Request log probabilities of next token. This version is asynchronous because it automatically batches concurrent requests; use with await.

Parameters:

Name	Type	Description	Default
token_ids	list[int]	a list of token ids, representing a prompt to the language model.	required
lora_name	str	Name of the LoRA adapter to use. Defaults to None (the base model).	None

Returns:

Name	Type	Description
logprobs	Tensor	a tensor of with the language model's log (normalized) probabilities for the next token following the prompt.

Source code in genlm/backend/llm/mlx.py

async def next_token_logprobs(self, token_ids, lora_name=None):
    """Request log probabilities of next token. This version is asynchronous because it automatically batches concurrent requests; use with `await`.

    Args:
        token_ids (list[int]): a list of token ids, representing a prompt to the language model.
        lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

    Returns:
        logprobs (torch.Tensor): a tensor of with the language model's log (normalized) probabilities for the next token following the prompt.
    """
    if not token_ids:
        raise ValueError("Token ids must not be empty")
    key = (tuple(token_ids), lora_name)
    if self.cache is not None and key in self.cache:
        return self.cache[key]
    future = asyncio.get_running_loop().create_future()
    batch = await self._join_batch([(key, future)], timeout=self.timeout)
    if batch is not None:
        self._batch_evaluate(batch)
    return await future

next_token_logprobs_sync(token_ids, lora_name=None)

Request log probabilities of next token synchronously.

Parameters:

Name	Type	Description	Default
token_ids	list[int]	A list of token IDs, representing a prompt to the language model.	required
lora_name	str	Name of the LoRA adapter to use. Defaults to None (the base model).	None

Returns:

Type	Description
Tensor	Normalized log probability tensor.

Source code in genlm/backend/llm/mlx.py

def next_token_logprobs_sync(self, token_ids, lora_name=None):
    """Request log probabilities of next token synchronously.

    Args:
        token_ids (list[int]): A list of token IDs, representing a prompt to the language model.
        lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

    Returns:
        (torch.Tensor): Normalized log probability tensor.
    """
    if not token_ids:
        raise ValueError("Token ids must not be empty")
    key = (tuple(token_ids), lora_name)
    return self._resolve([key])[key]

batch_next_token_logprobs_sync(token_ids_list, lora_name=None)

Next-token log-probs for each sequence, in one batched forward.

Parameters:

Name	Type	Description	Default
token_ids_list	list[list[int]]	Prompts' token ids.	required
lora_name	str	Name of the LoRA adapter to use. Defaults to None (the base model).	None

Returns:

Type	Description
Tensor	[len(token_ids_list), vocab] log-probabilities.

Source code in genlm/backend/llm/mlx.py

def batch_next_token_logprobs_sync(self, token_ids_list, lora_name=None):
    """Next-token log-probs for each sequence, in one batched forward.

    Args:
        token_ids_list (list[list[int]]): Prompts' token ids.
        lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

    Returns:
        (torch.Tensor): ``[len(token_ids_list), vocab]`` log-probabilities.
    """
    if any(not ids for ids in token_ids_list):
        raise ValueError("Token ids must not be empty")
    keys = [(tuple(ids), lora_name) for ids in token_ids_list]
    logprobs = self._resolve(list(dict.fromkeys(keys)))
    return torch.stack([logprobs[k] for k in keys])

sample(prompt_token_ids, max_tokens, eos_token_ids, temperature=1.0, seed=None, lora_name=None) async

Sample from the language model.

Parameters:

Name	Type	Description	Default
prompt_token_ids	list[int]	The token IDs of the prompt to start generation from.	required
max_tokens	int	The maximum number of tokens to generate.	required
eos_token_ids	list[int]	The token IDs that signal end-of-sequence. Generation stops when one of these is sampled.	required
temperature	float	The temperature to use for sampling. Higher values make the distribution more uniform, lower values make it more peaked. Defaults to 1.0.	1.0
seed	int	The seed for the random number generator. If provided, sets the random seed before sampling. Defaults to None.	None
lora_name	str	Name of the LoRA adapter to use. Defaults to None (the base model).	None

Returns:

Type	Description
list[int]	The sampled token IDs.

Source code in genlm/backend/llm/mlx.py

async def sample(
    self,
    prompt_token_ids,
    max_tokens,
    eos_token_ids,
    temperature=1.0,
    seed=None,
    lora_name=None,
):
    """Sample from the language model.

    Args:
        prompt_token_ids (list[int]): The token IDs of the prompt to
            start generation from.
        max_tokens (int): The maximum number of tokens to generate.
        eos_token_ids (list[int]): The token IDs that signal
            end-of-sequence. Generation stops when one of these is
            sampled.
        temperature (float, optional): The temperature to use for
            sampling. Higher values make the distribution more uniform,
            lower values make it more peaked. Defaults to 1.0.
        seed (int, optional): The seed for the random number generator.
            If provided, sets the random seed before sampling.
            Defaults to None.
        lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

    Returns:
        (list[int]): The sampled token IDs.
    """
    self._adapters.select(lora_name)
    if seed is not None:
        mx.random.seed(seed)

    generated = []
    for token, _ in generate_step(
        mx.array(prompt_token_ids),
        self.mlx_lm_model,
        max_tokens=max_tokens,
        sampler=make_sampler(temp=temperature),
    ):
        if token in eos_token_ids:
            break
        generated.append(token)
    return generated