## MLX ## mlx.MlxEngine Runs batched MLX inference on a dedicated worker thread. Usage ``` python 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](mlx.md#spawnllm.mlx.MLXPatches), and precomputes a prompt cache for `prefix_messages`. Await [ensure_loaded](mlx.md#spawnllm.mlx.MlxEngine.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](mlx.md#spawnllm.mlx.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](mlx.md#spawnllm.mlx.MlxEngine.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()](#spawnllm.mlx.MlxEngine.close) | Signal the worker thread to exit after completing already-queued jobs. | | [ensure_loaded()](#spawnllm.mlx.MlxEngine.ensure_loaded) | Wait for the worker thread to finish loading, re-raising the error if loading failed. | | [generate()](#spawnllm.mlx.MlxEngine.generate) | Generate a completion per conversation, defaulting to a single token. | | [peak_memory_gb()](#spawnllm.mlx.MlxEngine.peak_memory_gb) | Return the process's peak resident set size in GiB. | | [submit()](#spawnllm.mlx.MlxEngine.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 ``` python close() ``` ##### ensure_loaded() Wait for the worker thread to finish loading, re-raising the error if loading failed. Usage ``` python ensure_loaded() ``` ##### generate() Generate a completion per conversation, defaulting to a single token. Usage ``` python 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 ``` python peak_memory_gb() ``` ##### submit() Run `fn(*args)` on the worker thread and await its result. Usage ``` python 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 ``` python 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 ``` python >>> 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()](#spawnllm.mlx.AdapterCodec.decode) | Decompress `ZST` and write the restored safetensors adapter to `dst`. | | [digest()](#spawnllm.mlx.AdapterCodec.digest) | Return the first 16 hex digits of the SHA-256 of the compressed adapter `ZST`. | | [dtype()](#spawnllm.mlx.AdapterCodec.dtype) | Return the safetensors dtype shared by the compressed adapter's tensors, e.g. `"BF16"`. | | [encode()](#spawnllm.mlx.AdapterCodec.encode) | Compress the safetensors adapter at `src` into `ZST`. | ##### decode() Decompress `ZST` and write the restored safetensors adapter to `dst`. Usage ``` python 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 ``` python digest() ``` ##### dtype() Return the safetensors dtype shared by the compressed adapter's tensors, e.g. `"BF16"`. Usage ``` python dtype() ``` ##### encode() Compress the safetensors adapter at `src` into `ZST`. Usage ``` python 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 ``` python 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()](#spawnllm.mlx.AdapterFuser.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 ``` python 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 ``` python mlx.MLXPatches() ``` [MlxEngine](mlx.md#spawnllm.mlx.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()](#spawnllm.mlx.MLXPatches.apply) | Apply all patches once; subsequent calls are no-ops. | ##### apply() Apply all patches once; subsequent calls are no-ops. Usage ``` python apply() ```