Laya-MLX: a decision model that answers in 13 ms on a Mac

Laya-MLX does not produce sentences. It takes a state and a question together and returns a probability for each named option, which is why output tokens² stay at zero.

The structure splits into an encoder and decision heads. A bidirectional encoder reads state and question together and the heads compute the probabilities. With no token-by-token loop the latency stays short.

Asking a generative model to choose means getting a sentence back and parsing it. A broken format adds a retry and the latency that comes with it. Receiving the probability directly removes the parsing step.

The models come in two sizes, 421 million and 322 million parameters. Compared with running a generative model on the same Mac, memory use stays under a gigabyte.

A choice question returns a probability for each named option. A score question returns probabilities over ordered levels and their expected score, and a true/false question returns a single probability.

It does not fit writing prose or multi-step reasoning. It is for choosing among named options or rating on fixed levels.

Install the package, load a model and pass a state and a question. The example below routes a support message into one of three departments.

pip install laya-mlx

The first load downloads the checkpoint³, and after that inference⁴ stays local. It needs Python 3.11 or later and macOS 14 or later.

import laya_mlx as laya

agent = laya.load("aac6fef/laya-mlx")
result = agent.predict(
    "I was billed twice. Please refund the duplicate.",
    {"department": {
        "type": "choice",
        "instructions": "Who should handle this?",
        "criteria": ["billing", "technical", "sales"],
    }},
)
print(result["answers"]["department"])

It can also run on the CPU⁵ instead of the GPU⁶. Changing only the backend moves the same code between the two.

If you only handle English, the 421M model is enough; with mixed languages the multilingual checkpoint is faster. Their context limits differ, so check that number before passing long instructions.

Pre-converted FP16⁷ checkpoints are published, so you can load one directly instead of converting the original weights yourself.

The measurement ran on an M3 Max with a 40-core GPU and 128 GB of memory. The timing includes prompt⁸ preparation, tokenization, synchronized inference, calibration and formatting, and excludes model loading.

The numbers move with question length and count. The table measures one short question and a batch of 50, so longer instructions add latency.

The MLX port was checked against the original on 63 validation questions. All 63 matched in both FP32⁹ and FP16, so all 378 comparisons passed.

Repeating the same call 100 times returned finite, identical results without memory growth. That holds for those fixtures, not for every possible question.

The game demo serves the same check. A Snake run handled 2,400 moves with two safety interventions and no deaths, about 6.5% faster than its eager control.

Fidelity came before speed. Where a probability is used, a one-point difference changes the decision, so matching the original is the starting point.

Entry points for an HTTP server and an MCP¹⁰ server ship with the package. Calling it from an app or exposing it as an agent tool needs no extra code.

pip install 'laya-mlx[serve]'
laya-mlx-serve

Extras for frameworks such as LangChain and LlamaIndex ship separately, so you can swap only the decision step in a pipeline¹¹ you already run.

When sending many questions at once, tune the batch size¹². The default is 16, and the 50-question throughput¹³ above used a batch of 64.

This package is an independent MLX port, not an official release from the original authors. It keeps the original weights, question formatting, calibration and output schema, but its support surface differs.

CUDA¹⁴ and ONNX execution and training are not included. Training stays with the upstream project; this package covers inference and conversion.

The repository records which upstream revision it follows. Checking that record against the upstream documentation shows which features are still missing.

The measurement ran on macOS 27.2, and older versions cannot be assumed to match. Measure the latency yourself with a short question before relying on it.