Laya-MLX — Native Apple Silicon port of Laya typed decision models

TL;DR

Laya-MLX is an independent MLX port of Convai Innovations’ Laya (see Laya — Non-Autoregressive System One Decision Model (open-source RLCD)) that runs all three published Laya checkpoints natively on Apple Silicon — no PyTorch, Transformers runtime, or cloud API. Author @mizorewww ports the ModernBERT-large backbone, decision transformer, scoring and action heads to MLX; keeps the original weights, prompt formatting, calibration, and output schema; and publishes pre-converted FP16 checkpoints on Hugging Face under aac6fef/laya-mlx*. On an M3 Max the port reports 13.42 ms P50 for a short English typed decision (7.39 ms for the multilingual checkpoint), 146.8–395 questions/sec batched throughput, and ≤944 MiB peak MLX allocation — with 378/378 selected-answer matches against upstream on the validation fixtures across FP32 and FP16.

Key claims

Native MLX runtime with no PyTorch/Transformers dependency; encoder, decision transformer, scoring head, and action head all run in MLX, tokenization via Hugging Face’s Rust tokenizer [§Why typed decisions?, §Attribution and license].
M3 Max latency (FP16, end-to-end, model loading excluded): 13.42 ms P50 / 13.92 ms P95 single-question for Laya 421M; 7.39 ms P50 / 7.79 ms P95 for the multilingual 322M checkpoint [Performance on M3 Max table].
M3 Max batched throughput (batch_size=64, 50-question measurement): 146.8 questions/sec (Laya 421M), 395.0 questions/sec (multilingual) [Performance on M3 Max table].
Peak MLX allocation on a single short question: 943.6 MiB (Laya 421M), 687.6 MiB (multilingual) — fits comfortably in ≤1 GB [Performance on M3 Max table].
Port fidelity: all three checkpoints match upstream’s selected answer on 63/63 validation questions in both FP32 and FP16 (378/378 comparisons); 100 repeated finite deterministic calls per configuration with zero measured active-memory growth [§Performance on M3 Max, port fidelity paragraph].
Snake demo peaks at 75.40 moves/s across 2,400 moves with --optimize --max-speed, zero deaths and 2 visible safety interventions — ~6.5% faster than its same-run eager control on the same M3 Max [§Quick start].
Independent-port disclaimer: this is not an official Convai Innovations release; RLCD training and fine-tuning remain in the upstream project, and the port covers inference and MLX conversion only [§Supported checkpoints, §Attribution and license].
Performance-research investigation reports that the current engineering does not support a further universal 10× speedup on the same checkpoints; selected cases show 1.03–1.08× paired median speedups from compilation, prefix-cache, and custom Metal kernels [§Performance research].
Optional runtime knobs (compile=True, pad_to_multiple=16, cache_prompts=True) are opt-in with first-use compilation cost and shape specialization; the prefix cache is bounded to 128 questions and shares state-tokenization but not encoder computation [§Python API].
License: Apache-2.0; upstream attribution pinned to Laya commit 6a5819129eb220570792e417e49723d697efd76f [§Attribution and license].

Method

The port reimplements the Laya architecture — bidirectional ModernBERT-large encoder + 2-layer decision transformer + option-marker scoring head + act/escalate head — in MLX, targeting the Apple Silicon unified memory + Metal Performance Shaders stack. ModernBERT-specific structure is preserved: the global/local attention pattern, inclusive sliding-window boundary, distinct local vs global RoPE bases, and first-layer normalization behavior all carry over. Unsupported encoders and non-default RoPE scaling fail explicitly at load time rather than degrading silently.

Weight conversion (laya-mlx convert) reads the original safetensors, remaps parameter names into MLX’s convention, and writes model.safetensors plus encoder/agent configs, tokenizer files, and mlx_config.json. It is a name/dtype conversion, not quantization or retraining. Source checkpoints already store FP16 weights; FP32 mode increases arithmetic precision only. BF16 is available on request but excluded from the published validation matrix.

At runtime, the Python API mirrors upstream Laya’s predict(state, questions) contract, including four-decimal rounding, action.act_probability, and token usage fields. Router(dtype=..., max_loaded=..., preload=...) reproduces upstream’s language routing and typed-decisions workflow detection. Load-bearing detail: question rows are batched independently, and encoder representations depend jointly on state and question — the README explicitly disclaims any state-only encoder reuse across questions.

Results

Reported on M3 Max (40 GPU cores, 128 GiB memory), timing inclusive of prompt preparation, tokenization, tensor build, synchronized inference, calibration, and result formatting; model load excluded. Single-question end-to-end P50 lands at 13.42 ms (Laya 421M) and 7.39 ms (multilingual 322M). At batch_size=64 the multilingual checkpoint reaches 395 questions/sec. Peak MLX allocation stays under 1 GB for either checkpoint on a single short question.

Fidelity to upstream: 378/378 selected-answer matches across three checkpoints × 63 validation questions × 2 precisions (FP32, FP16). Numerical parity in probabilities is not exact across precisions but selected labels agree; per-question probability errors are documented in the repo’s BENCHMARKS.md.

Snake benchmark under the tested optimization flag combination (--optimize --max-speed): 75.40 moves/s across 2,400 moves, zero deaths, 2 visible cycle-safety interventions, ≈6.5% faster than the same-run eager control on the same hardware.

Set against Jev’s published third-party numbers (236–276 ms P50, cited in Laya — Non-Autoregressive System One Decision Model (open-source RLCD)): 276 ms / 7.4 ms ≈ 37× for the multilingual checkpoint, 276 ms / 13.4 ms ≈ 20× for Laya 421M — the launch tweet’s “50×” framing is directionally in the same ballpark but does not correspond exactly to any single benchmark cell in the repo README.

Why it’s interesting

This is the load-bearing test of the Laya — Non-Autoregressive System One Decision Model (open-source RLCD) argument that a 421M non-autoregressive decision model with a bidirectional encoder can be moved off cloud GPU inference onto commodity Apple Silicon and stay in the sub-15 ms tail — the concrete counter to Jev’s proprietary-cloud-endpoint framing. It shows that once the architectural bet (parallel-decode + short-output + specialised arch, three-lever stack tracked on LLM Inference Efficiency) is public and Apache-2.0, the last mile is a runtime-engineering exercise a single contributor can complete in weeks — with published FP16 checkpoints on Hugging Face and 378/378 upstream-parity on validation fixtures. Contrasts with the APXinf-robo — Edge Inference Engine for VLA/WAM Models on Jetson Thor / Orin Jetson-class edge-inference pattern in the VLA/WAM cluster: same “batch-1 latency-first serving on a device the user already owns” regime, but on a laptop rather than a robot compute box, and for a decision model rather than a robot policy. Also sharpens the Tool-Use Agents output-space-specialist pole: the fully-local sub-15 ms Snake demo is the first filed evidence that a Laya-class typed-decision agent can run inside a real-time interactive loop with no network hop, which is what the Jev launch post was pointing at but couldn’t demonstrate publicly.

Caveats worth flagging: benchmarks are single-machine self-reported on one M3 Max configuration; the “50×” framing in the launch tweet is the author’s rhetorical framing rather than a repo-reported number and doesn’t correspond to a specific benchmark cell; the performance-research thread explicitly concludes the current port does not support a further universal 10× speedup on the same checkpoints — the 1.03–1.08× medians from compilation and custom Metal kernels are the honest headroom. All accuracy claims are inherited from upstream Laya’s own numbers (which themselves have the “base and multilingual checkpoints regress below majority-class baseline zero-shot” caveat noted on the Laya page).