Laya‑MLX – Fast, typed‑decision inference on Apple Silicon

What it is – A Python package that lets you run the Laya family of decision‑making language models completely locally on macOS Apple Silicon (M‑series) GPUs. It ports the original Convai Innovations checkpoints to MLX, Apple’s metal‑accelerated tensor library, so you get sub‑15 ms latency for a single short question without any PyTorch, 🤗 Transformers, or cloud calls.

---

Core ideas

Concept	How Laya‑MLX implements it
Typed decisions	Instead of free‑form text generation, the model returns structured answers (choice, score, or binary “noul”) in a single forward pass. This avoids token‑by‑token decoding and yields deterministic, low‑latency outputs.
Bidirectional encoder	The input (state + question) is encoded with a ModernBERT‑large or mmBERT‑base backbone, then specialised heads produce probabilities for the requested type.
Local, no‑runtime dependencies	All inference runs in MLX; tokenisation uses the Hugging Face Rust tokenizer compiled into the wheel. No PyTorch/Transformers binaries are required.
Apple‑silicon optimisation	Optional compile=True, prefix‑caching and padding tricks give a ~6 % speed boost on M3 Max; the library also ships pre‑converted FP16 checkpoints for the GPU.

---

Quick start (30 s)

pip install laya-mlx            # core library
pip install 'laya-mlx[demo]'    # optional demo utilities

import laya_mlx as laya
agent = laya.load("aac6fef/laya-mlx")   # downloads the FP16 checkpoint on first use
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"])   # → "billing"

Works on macOS 14+, Python 3.11+, and any Apple‑silicon GPU (M1‑M3). The first call downloads the model; subsequent calls are fully offline.

---

Available checkpoints

Model ID (load)	Encoder	Params	Context	Language
aac6fef/laya-mlx	ModernBERT‑large	421 M	512	English
aac6fef/laya-multilingual-mlx	mmBERT‑base	322 M	1 024	Multilingual
aac6fef/laya-typed-decisions-mlx	ModernBERT‑large	421 M	1 024	English (typed‑decision workflow)

All three are exact FP16 conversions of the upstream Convai Innovations weights and are hosted on Hugging Face. You can also point laya.load at the original hub IDs (convaiinnovations/laya, etc.) – the library will download and convert them on the fly.

---

Performance (M3 Max, FP16)

Metric	English (Laya 421M)	Multilingual (Laya‑multilingual 322M)
Median latency, one short question	13.4 ms	7.4 ms
95‑th‑percentile latency	13.9 ms	7.8 ms
Throughput (50‑question batch, batch‑size 64)	146 q/s	395 q/s
Peak GPU memory per request	944 MiB	688 MiB
Optimised laya‑snake --optimize --max‑speed	75.4 moves / s (≈6.5 % faster than eager)

Numbers include tokenisation, tensor preparation, inference, calibration and result formatting; model loading is excluded.

---

Main API surface

agent = laya.load(
    checkpoint,          # HF repo ID or local path
    dtype="float16",   # or "float32", "bfloat16"
    batch_size=16,       # max questions per forward pass
    device="gpu",       # "cpu" also works (much slower)
    compile=False,       # enable MLX compilation for speed
    cache_prompts=False, # keep tokenised prompts for reuse
)

# Predict – `system_one` is an alias
answers = agent.predict(state, questions)

*state can be a plain string, a JSON dict, or a list of prior messages. questions is a dict where each entry describes the desired answer type (choice, score, noul). The return value contains:

• answers (structured results)
• action.act_probability (raw head probability)
• token‑usage statistics
• rounding to four decimal places (as in the upstream model).

Router helper

For applications that need to auto‑select the language‑specific checkpoint:

router = laya.Router(dtype="float16", max_loaded=2)
out = router.predict(state, triage_questions())
print(out["routing"])   # e.g. "multilingual"

The router keeps up to max_loaded models resident and can be pre‑loaded with Router(preload=True).

---

Command‑line utilities

Command	Purpose
laya-mlx predict …	Run a single inference from JSON files or inline strings.
laya-mlx convert …	Convert a Hugging Face checkpoint to an MLX‑compatible directory (safetensors + config).
laya-snake	Interactive terminal demo that plays the classic Snake game, calling the model for each move. Use --optimize for the compiled fast path.

All CLIs accept the same --model argument used by laya.load.

---

Development & testing

• Dependencies: the project uses uv for reproducible environments. uv sync --extra dev pulls testing, benchmarking, and reference extras.
• Tests: unit tests compare the MLX implementation against the original Transformers heads on small random models; full‑checkpoint validation checks tokenisation, calibrated probabilities, deterministic repeats, and memory growth.
• Benchmarks: benchmarks/run measures latency/throughput; results are stored under benchmarks/results and summarised in BENCHMARKS.md.
• Research: the docs/ folder contains deep‑dive reports on performance bottlenecks and speculative 10× speed‑up ideas (mathematical, engineering, and implementation‑level). Scripts are in experiments/.

---

License & attribution

• Code – Apache‑2.0 (see LICENSE).
• Weights – Original Laya weights are © Convai Innovations; they are redistributed under the same license via the Hugging Face repos linked above.
• Port – The MLX re‑implementation and surrounding utilities are authored by mizorewww and adapt parts of the upstream NandhaKishorM/laya repository (MIT‑style attribution in NOTICE).

---

Who might use this?

• Product teams that need deterministic, low‑latency routing or classification on‑device (e.g., ticket triage, urgency scoring, binary policy checks).
• Developers building macOS‑only AI assistants or edge services where sending data to the cloud is undesirable.
• Researchers interested in benchmarking MLX against PyTorch/Transformers on Apple GPUs or exploring further speed‑up techniques.

---

TL;DR – Laya‑MLX gives you a ready‑to‑run, high‑performance inference library for the Laya decision‑making models on Apple Silicon, with a clean Python API, a terminal Snake demo, and full‑stack tooling for conversion, benchmarking, and model‑card publishing.