MLX

mlx.MlxEngine

Runs batched MLX inference on a dedicated worker thread.

Usage

mlx.MlxEngine(
    fused_dir,
    *,
    logits_processor_factory,
    prefix_messages,
    batch_size,
    worker_name="mlx"
)

The constructor starts a daemon worker thread that loads the model from fused_dir, applies MLXPatches, and precomputes a prompt cache for prefix_messages. Await ensure_loaded to surface load failures before generating.

Parameters

fused_dir: Path: Directory of the MLX model to load, such as the path returned by AdapterFuser.ensure_fused.
logits_processor_factory: Callable[[Any], Callable[…, Any]]: Builds the logits processor from the loaded tokenizer; the processor is applied to every generation call.
prefix_messages: list[dict[str, str]]: Chat messages shared by every conversation; their prompt cache is computed once and reused across batches. Each conversation passed to generate must start with these messages.
batch_size: int: Number of conversations generated per batch.
worker_name: str = "mlx": Name of the worker thread.

Raises

RuntimeError: When not running on macOS with Apple Silicon.

Methods

Name	Description
close()	Signal the worker thread to exit after completing already-queued jobs.
ensure_loaded()	Wait for the worker thread to finish loading, re-raising the error if loading failed.
generate()	Generate a completion per conversation, defaulting to a single token.
peak_memory_gb()	Return the process’s peak resident set size in GiB.
submit()	Run `fn(*args)` on the worker thread and await its result.

close()

Signal the worker thread to exit after completing already-queued jobs.

Usage

Source

close()

ensure_loaded()

Wait for the worker thread to finish loading, re-raising the error if loading failed.

Usage

Source

ensure_loaded()

generate()

Generate a completion per conversation, defaulting to a single token.

Usage

Source

generate(message_lists, on_progress, *, max_tokens=1)

Conversations are processed in chunks of batch_size, ordered by the length of each conversation’s final message so similar-length prompts batch together. Every conversation must start with the engine’s prefix_messages; the shared prefix is served from the precomputed prompt cache.

Parameters

message_lists: list[list[dict[str, str]]]: Conversations, each a list of chat messages with role and content keys.
on_progress: Callable[[int], None]: Called after each chunk with the number of conversations completed in that chunk.
max_tokens: int = 1: Maximum tokens generated per conversation; the default of 1 yields single-token classification output.

Returns

list[str]: Generated texts, in the same order as message_lists.

peak_memory_gb()

Return the process’s peak resident set size in GiB.

Usage

Source

peak_memory_gb()

submit()

Run fn(*args) on the worker thread and await its result.

Usage

Source

submit(fn, *args)

Parameters

fn: Callable[…, R]: Callable to execute on the worker thread.
*args: Any: Positional arguments passed to fn.

Returns

R: The value returned by fn; an exception raised by fn propagates to the awaiter.

mlx.AdapterCodec

Compresses and decompresses a homogeneous-dtype safetensors LoRA adapter.

Usage

Source

mlx.AdapterCodec()

Each tensor’s bytes are grouped by byte position within the dtype (byte-shuffle) before zstd compression. Subclasses override DIR, ZST, and CONFIG to point at their shipped package data.

Example

>>> class ShippedCodec(AdapterCodec):
...     DIR = Path(__file__).parent
...     ZST = DIR / "adapters.safetensors.zst"
...     CONFIG = DIR / "adapter_config.json"
>>> ShippedCodec.encode(Path("trained/adapters.safetensors"))

Methods

Name	Description
decode()	Decompress `ZST` and write the restored safetensors adapter to `dst`.
digest()	Return the first 16 hex digits of the SHA-256 of the compressed adapter `ZST`.
dtype()	Return the safetensors dtype shared by the compressed adapter’s tensors, e.g. `"BF16"`.
encode()	Compress the safetensors adapter at `src` into `ZST`.

decode()

Decompress ZST and write the restored safetensors adapter to dst.

Usage

Source

decode(dst)

Parameters

dst: Path: Destination path for the safetensors file.

digest()

Return the first 16 hex digits of the SHA-256 of the compressed adapter ZST.

Usage

Source

digest()

dtype()

Return the safetensors dtype shared by the compressed adapter’s tensors, e.g. "BF16".

Usage

Source

dtype()

encode()

Compress the safetensors adapter at src into ZST.

Usage

Source

encode(src)

Byte-shuffles each tensor’s data, then compresses with zstd at COMPRESSION_LEVEL. All tensors must share one dtype drawn from TYPESIZES.

Parameters

src: Path: Path to the adapter’s safetensors file.

mlx.AdapterFuser

Fuses a shipped LoRA adapter into a base MLX model.

Usage

Source

mlx.AdapterFuser()

The fused model is stored in the Hugging Face hub cache using the standard models--*/snapshots/* layout, keyed by the codec’s digest, so repeat calls reuse the cached result.

Methods

Name	Description
ensure_fused()	Fuse the codec’s adapter into `model_repo`, returning the cached result when present.

ensure_fused()

Fuse the codec’s adapter into model_repo, returning the cached result when present.

Usage

Source

ensure_fused(model_repo, *, codec, cache_namespace, tqdm_class=None)

On a cache miss, downloads the base model from the Hugging Face hub, decodes the compressed adapter, applies and fuses the LoRA layers, and saves the merged model into the hub cache.

Parameters

model_repo: str: Hugging Face repo id of the base MLX model.
codec: AdapterCodec: Codec providing the compressed adapter and its config.
cache_namespace: str: Names the cache entry models--{cache_namespace}-{digest}.
tqdm_class: type | None = None: Progress-bar class forwarded to snapshot_download.

Returns

Path: Path to the fused model’s snapshot directory.

mlx.MLXPatches

Idempotent runtime patches for mlx_lm.

Usage

Source

mlx.MLXPatches()

MlxEngine applies these on its worker thread before the first batch_generate call. The single current patch replaces BatchGenerator.stats with a version that clamps the elapsed-time denominators when computing tokens-per-second, preventing a ZeroDivisionError.

Methods

Name	Description
apply()	Apply all patches once; subsequent calls are no-ops.

apply()

Apply all patches once; subsequent calls are no-ops.

Usage

Source

apply()