Add optional MLX backend for Apple Silicon (backend='mlx' or 'auto'; ONNX stays the default)

README.md

@@ -25,6 +25,11 @@ tags:
 - distillation
 - knowledge-distillation
 - onnxruntime
+- mlx
+- apple-silicon
+- macos
+- metal
+- on-device
 - zero-shot
 - nlu
 - intent-router
@@ -70,11 +75,11 @@ model-index:
       name: Routing accuracy, 6 unseen domains
 ---
 
-# Laya Intent Router 150M: zero-shot intent classification on CPU (ONNX, int8)
+# Laya Intent Router 150M: zero-shot intent classification on CPU (ONNX, int8) and Apple Silicon (MLX)
 
 **A 150M-parameter zero-shot intent detection model with out-of-scope detection, 160 ms p95 on 4 CPU threads.** You give it a user message and a list of intents written in plain English. It tells you which intent the message belongs to, or that it belongs to none of them.
 
-No training. No labelled data. You write the intents when you call it, and you can change them on every request. It's a small ONNX file you run with `onnxruntime` and `numpy`, no GPU and no torch.
+No training. No labelled data. You write the intents when you call it, and you can change them on every request. It's a small ONNX file you run with `onnxruntime` and `numpy`, no GPU and no torch. On an Apple Silicon Mac, an optional MLX backend runs the same model on the GPU in about 15 to 25 ms.
 
 On a held-out set of 1,636 messages it gets **94.9% routing accuracy and 98.6% out-of-scope recall**, up from 72.8% and 68.2% for the original [Laya](https://huggingface.co/convaiinnovations/laya).
 
@@ -131,6 +136,37 @@ Each intent is a key plus 1 or more example phrasings or a short description. `m
 
 Or clone [the GitHub repo](https://github.com/vrajnotviraj/laya-intent-router) and use `from laya_intent_router import LayaIntentRouter`, or run `python laya_intent_router.py "your message here"`.
 
+## Faster on a Mac: MLX backend
+
+On Apple Silicon the same model can run on the GPU through [MLX](https://github.com/ml-explore/mlx). Install `mlx` next to the usual packages and pass `backend="mlx"` or `backend="auto"`:
+
+```bash
+pip install onnxruntime tokenizers numpy huggingface_hub mlx
+```
+
+```python
+r = lir.LayaIntentRouter.from_pretrained(backend="auto")   # MLX if it can, ONNX if not
+r.backend                                              # 'mlx'
+lir.LayaIntentRouter.from_pretrained(backend="mlx")      # force MLX, errors out if mlx is missing
+lir.LayaIntentRouter.from_pretrained()                   # default: ONNX on CPU, never downloads mlx/
+```
+
+The default is ONNX, so nothing changes unless you ask for MLX. `auto` only picks MLX on an Apple Silicon Mac with `mlx` importable and the `mlx/` weights in the repo. Everywhere else (Linux, Windows, Intel Macs, or a Mac without `mlx`) it runs ONNX on CPU, same as before. Each backend downloads only its own weights, so ONNX users never pull the 328 MB `mlx/` folder and MLX users never pull `laya.onnx`. `mlx` is optional and never required.
+
+One catch: a copy of `laya_intent_router.py` from before September 29, 2026 downloads the whole repo, `mlx/` included. Grab the current file to skip that.
+
+The MLX weights are converted from the same trained checkpoint as `laya.onnx` (fp16, with the Linear layers quantized to 8 bits at load). On 5,000 dev routing episodes they made the same match-or-none decision as the ONNX model 99.7% of the time, and accuracy and out-of-scope recall were within 0.001. The few that differ sit right on the 0.625 threshold, where ONNX's own int8 rounding tips them one way or the other.
+
+Latency on an M2 Pro, 1 message at a time, warm, whole `route()` call including the shortlist (p50 / p95, 300 calls each):
+
+| intents | ONNX, 4 CPU threads | MLX (M2 Pro GPU) |
+|---|---|---|
+| 2 | 66 / 107 ms | 15 / 22 ms |
+| 7 | 107 / 213 ms | 24 / 37 ms |
+| 20 | 150 / 219 ms | 28 / 38 ms |
+
+So about 4 to 5x faster. The machine was busy during these runs (ONNX p95 is 160 ms on a quiet one), so treat the p95s as upper bounds and compare the columns with each other. It doesn't get under 10 ms: the cost is 22 encoder layers over a ~225-token prompt, and dropping from fp32 to fp16 or 8-bit saves at most about 8 ms. I've only measured an M2 Pro.
+
 ## More examples
 
 A SaaS support bot, 4 intents, nothing fine-tuned:
@@ -205,6 +241,16 @@ Everything trained locally on a 32 GB M2 Pro. Training code: [github.com/vrajnot
 | `tokenizer.json`, `rl_agent_config.json` | tokenizer, prompt format, calibrated temperatures, threshold |
 | `shortlist/` | bge-small-en-v1.5 ONNX embedder for long intent lists |
 | `laya_intent_router.py` | the whole inference code, one file, no torch |
+| `mlx/` | the same model as MLX weights for Apple Silicon (fp16, 8-bit Linear at load, 328 MB), optional |
+| `laya_intent_router_mlx.py` | the MLX model code, only loaded when you pass `backend="mlx"` or `"auto"` |
+
+## FAQ
+
+**Can I run it in Ollama or LM Studio?**
+
+No, and I don't think it's possible without breaking the model. Ollama and LM Studio run GGUF files through llama.cpp, which is built for text generation and plain embeddings. This model isn't either of those. After the encoder it has a small decision head: 2 extra transformer layers, a scorer that reads one position per intent (the `[MASK]` in front of each one) and a temperature per menu size. GGUF has nowhere to put that head, and llama.cpp has no way to run it. You'd get an embedding model and lose the part that says "none of these".
+
+If you want it local and fast on a Mac, use the MLX backend above. It runs the whole graph, head included, in-process.
 
 ## License and credits
 

laya_intent_router.py

@@ -9,12 +9,17 @@
     # {'match': 'cancel_order', 'score': 0.97, 'probabilities': {...}}
 
 Needs: pip install onnxruntime tokenizers numpy huggingface_hub
+On Apple Silicon, `pip install mlx` as well and the model runs on the GPU with backend="mlx" or "auto";
+laya_intent_router_mlx.py has the MLX code). The default is backend="onnx" (CPU), which never downloads mlx/.
 """
 
 import hashlib
 import json
+import importlib.util
 import os
+import platform
 import re
+import sys
 from collections import OrderedDict
 
 import numpy as np
@@ -30,6 +35,9 @@ HISTORY_HINT = (
 )
 NO_MATCH_DESCRIPTION = "Gibberish, filler words, or a message unrelated to every other path."
 NONE = "__none__"
+BACKENDS = ("auto", "onnx", "mlx")
+SHARED_FILES = ["tokenizer.json", "rl_agent_config.json", "shortlist/*"]
+BACKEND_FILES = {"onnx": ["laya.onnx"], "mlx": ["mlx/*", "laya_intent_router_mlx.py"]}
 OPTION_CAP, MIN_HEAD_TOKENS, PHRASING_SEP = 96, 16, '" | "'
 
 
@@ -53,6 +61,32 @@ def _humanise(path_id):
     return re.sub(r"[_\-.:/]+", " ", s).strip().lower()
 
 
+def _mlx_available():
+    if sys.platform != "darwin" or platform.machine() != "arm64":
+        return False
+    try:
+        import mlx.core as mx
+
+        return bool(mx.metal.is_available())  # False in macOS VMs / CI runners without a Metal GPU
+    except Exception:
+        return False
+
+
+def _has_mlx_weights(model_dir):
+    return os.path.exists(os.path.join(model_dir, "mlx", "model.safetensors"))
+
+
+def _mlx_session(model_dir):
+    for d in (os.path.dirname(os.path.abspath(__file__)), model_dir):  # next to this file, else the downloaded copy
+        f = os.path.join(d, "laya_intent_router_mlx.py")
+        if os.path.exists(f):
+            spec = importlib.util.spec_from_file_location("laya_intent_router_mlx", f)
+            mod = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(mod)
+            return mod.Session(os.path.join(model_dir, "mlx"))
+    raise FileNotFoundError("laya_intent_router_mlx.py not found next to laya_intent_router.py or in " + model_dir)
+
+
 def _session(path, threads):
     import onnxruntime as ort
 
@@ -100,13 +134,21 @@ class Shortlist:
 
 
 class LayaIntentRouter:
-    def __init__(self, model_dir, threads=4, shortlist_k=4):
+    def __init__(self, model_dir, threads=4, shortlist_k=4, backend="onnx"):
+        """backend: "onnx" (CPU, default), "mlx" (Apple GPU) or "auto" (mlx on Apple Silicon when mlx is installed
+        and mlx/ weights are present, onnx otherwise). The chosen one is in self.backend."""
         from tokenizers import Tokenizer
 
+        if backend not in BACKENDS:
+            raise ValueError(f"backend must be one of {BACKENDS}, got {backend!r}")
+        if backend == "auto":
+            backend = "mlx" if _has_mlx_weights(model_dir) and _mlx_available() else "onnx"
+
         self.cfg = json.load(open(os.path.join(model_dir, "rl_agent_config.json")))
         self.tok = Tokenizer.from_file(os.path.join(model_dir, "tokenizer.json"))
         self.cls, self.sep, self.mask = (self.tok.token_to_id(t) for t in ("[CLS]", "[SEP]", "[MASK]"))
-        self.sess = _session(os.path.join(model_dir, "laya.onnx"), threads)
+        self.backend = backend
+        self.sess = _mlx_session(model_dir) if backend == "mlx" else _session(os.path.join(model_dir, "laya.onnx"), threads)
         self.threshold = self.cfg.get("match_threshold", 0.625)
         self.temps = self.cfg["temperature_by_options"]
         sl_dir = os.path.join(model_dir, "shortlist")
@@ -114,10 +156,21 @@ class LayaIntentRouter:
         self.shortlist = Shortlist(sl_dir, threads) if self.k else None
 
     @classmethod
-    def from_pretrained(cls, repo_id=REPO_ID, **kw):
+    def from_pretrained(cls, repo_id=REPO_ID, backend="onnx", **kw):
+        """Downloads only what the backend needs: ONNX never fetches mlx/, MLX never fetches laya.onnx."""
         from huggingface_hub import snapshot_download
 
-        return cls(snapshot_download(repo_id), **kw)
+        if backend not in BACKENDS:
+            raise ValueError(f"backend must be one of {BACKENDS}, got {backend!r}")
+        if backend == "auto":
+            if _mlx_available():
+                d = snapshot_download(repo_id, allow_patterns=SHARED_FILES + BACKEND_FILES["mlx"])
+                if _has_mlx_weights(d):
+                    return cls(d, backend="mlx", **kw)
+            backend = "onnx"
+        if backend == "mlx" and not _mlx_available():  # fail before downloading 328 MB we can't use
+            raise ImportError('backend="mlx" needs Apple Silicon and `pip install mlx`')
+        return cls(snapshot_download(repo_id, allow_patterns=SHARED_FILES + BACKEND_FILES[backend]), backend=backend, **kw)
 
     def _tokens(self, text):
         return self.tok.encode(text.replace("[MASK]", " "), add_special_tokens=False)
@@ -189,8 +242,6 @@ class LayaIntentRouter:
 
 
 if __name__ == "__main__":
-    import sys
-
     here = os.path.dirname(os.path.abspath(__file__))
     r = LayaIntentRouter(here) if os.path.exists(os.path.join(here, "laya.onnx")) else LayaIntentRouter.from_pretrained()
     paths = {
@@ -199,6 +250,7 @@ if __name__ == "__main__":
         "return_order": ["Customer wants to return an order"],
         "wrong_item": ["Customer received an item different from what they ordered"],
     }
+    print(f"backend: {r.backend}")
     for msg in sys.argv[1:] or ["where is my parcel #A-7721", "qwewqeqw", "i got blue shoes but ordered black"]:
         out = r.route(msg, paths)
         print(f"{msg!r:45} -> {out['match']} ({out['score']:.2f})")

laya_intent_router_mlx.py

@@ -0,0 +1,174 @@
+"""MLX backend for laya_intent_router.py: the same model running on the Apple Silicon GPU.
+
+Used by LayaIntentRouter(backend="mlx" or "auto"). Needs `pip install mlx` and the weights in `mlx/`.
+Encoder: ModernBERT (Ettin-150M). Pre-norm layers (layer 0 has no attn_norm), bias-free Linear and
+LayerNorm, GeGLU MLP with exact GELU, rotate-half RoPE with a theta per layer type, global attention
+on every 3rd layer and a |i - j| <= local_attention // 2 window on the rest. Head: 2 norm-first
+transformer layers (ReLU), then a scorer on each [MASK] marker. Weights are stored in fp16; the Linear
+layers are quantized to 8 bits at load when mlx/config.json asks for it.
+"""
+
+import json
+import os
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+
+
+def _mask(attention_mask, window=None, L=None):
+    m = None
+    if window is not None:
+        i = mx.arange(L)
+        m = (mx.abs(i[:, None] - i[None, :]) <= window)[None, None]
+    if attention_mask is not None:
+        pad = attention_mask.astype(mx.bool_)[:, None, None, :]
+        m = pad if m is None else m & pad
+    return m
+
+
+class Attention(nn.Module):
+    def __init__(self, d, heads, theta):
+        super().__init__()
+        self.heads, self.theta = heads, theta
+        self.Wqkv = nn.Linear(d, 3 * d, bias=False)
+        self.Wo = nn.Linear(d, d, bias=False)
+
+    def __call__(self, x, mask):
+        B, L, d = x.shape
+        q, k, v = self.Wqkv(x).reshape(B, L, 3, self.heads, -1).transpose(2, 0, 3, 1, 4)
+        q = mx.fast.rope(q, q.shape[-1], traditional=False, base=self.theta, scale=1.0, offset=0)
+        k = mx.fast.rope(k, k.shape[-1], traditional=False, base=self.theta, scale=1.0, offset=0)
+        o = mx.fast.scaled_dot_product_attention(q, k, v, scale=q.shape[-1] ** -0.5, mask=mask)
+        return self.Wo(o.transpose(0, 2, 1, 3).reshape(B, L, d))
+
+
+class MLP(nn.Module):
+    def __init__(self, d, inter):
+        super().__init__()
+        self.Wi = nn.Linear(d, 2 * inter, bias=False)
+        self.Wo = nn.Linear(inter, d, bias=False)
+
+    def __call__(self, x):
+        a, g = mx.split(self.Wi(x), 2, axis=-1)
+        return self.Wo(nn.gelu(a) * g)
+
+
+class EncoderLayer(nn.Module):
+    def __init__(self, c, idx):
+        super().__init__()
+        d, eps = c["hidden_size"], c["norm_eps"]
+        self.sliding = c["layer_types"][idx] == "sliding_attention"
+        theta = c["rope_parameters"]["sliding_attention" if self.sliding else "full_attention"]["rope_theta"]
+        if idx:
+            self.attn_norm = nn.LayerNorm(d, eps=eps, bias=False)
+        self.attn = Attention(d, c["num_attention_heads"], theta)
+        self.mlp_norm = nn.LayerNorm(d, eps=eps, bias=False)
+        self.mlp = MLP(d, c["intermediate_size"])
+
+    def __call__(self, x, mask):
+        x = x + self.attn(self.attn_norm(x) if "attn_norm" in self else x, mask)
+        return x + self.mlp(self.mlp_norm(x))
+
+
+class Embeddings(nn.Module):
+    def __init__(self, c):
+        super().__init__()
+        self.tok_embeddings = nn.Embedding(c["vocab_size"], c["hidden_size"])
+        self.norm = nn.LayerNorm(c["hidden_size"], eps=c["norm_eps"], bias=False)
+
+    def __call__(self, ids):
+        return self.norm(self.tok_embeddings(ids))
+
+
+class ModernBert(nn.Module):
+    def __init__(self, c):
+        super().__init__()
+        self.window = c["local_attention"] // 2
+        self.embeddings = Embeddings(c)
+        self.layers = [EncoderLayer(c, i) for i in range(c["num_hidden_layers"])]
+        self.final_norm = nn.LayerNorm(c["hidden_size"], eps=c["norm_eps"], bias=False)
+
+    def __call__(self, ids, attention_mask=None):
+        x = self.embeddings(ids)
+        full, local = _mask(attention_mask), _mask(attention_mask, self.window, ids.shape[1])
+        for layer in self.layers:
+            x = layer(x, local if layer.sliding else full)
+        return self.final_norm(x)
+
+
+class MultiheadAttention(nn.Module):
+    def __init__(self, d, heads):
+        super().__init__()
+        self.heads = heads
+        self.in_proj_weight, self.in_proj_bias = mx.zeros((3 * d, d)), mx.zeros((3 * d,))
+        self.out_proj = nn.Linear(d, d)
+
+    def __call__(self, x, mask):
+        B, L, d = x.shape
+        qkv = x @ self.in_proj_weight.T + self.in_proj_bias
+        q, k, v = qkv.reshape(B, L, 3, self.heads, -1).transpose(2, 0, 3, 1, 4)
+        o = mx.fast.scaled_dot_product_attention(q, k, v, scale=q.shape[-1] ** -0.5, mask=mask)
+        return self.out_proj(o.transpose(0, 2, 1, 3).reshape(B, L, d))
+
+
+class HeadLayer(nn.Module):
+    def __init__(self, d, heads):
+        super().__init__()
+        self.self_attn = MultiheadAttention(d, heads)
+        self.linear1, self.linear2 = nn.Linear(d, 4 * d), nn.Linear(4 * d, d)
+        self.norm1, self.norm2 = nn.LayerNorm(d), nn.LayerNorm(d)
+
+    def __call__(self, x, mask):
+        x = x + self.self_attn(self.norm1(x), mask)
+        return x + self.linear2(nn.relu(self.linear1(self.norm2(x))))
+
+
+class Head(nn.Module):
+    def __init__(self, d, n):
+        super().__init__()
+        self.layers = [HeadLayer(d, max(1, d // 64)) for _ in range(n)]
+
+
+class DecisionModel(nn.Module):
+    def __init__(self, c):
+        super().__init__()
+        d = c["hidden_size"]
+        self.encoder = ModernBert(c)
+        self.head = Head(d, c["head_layers"])
+        self.type_emb = nn.Embedding(3, d)
+        self.scorer = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, d), nn.GELU(), nn.Linear(d, 1))
+        # the act head ships with the weights but routing never reads it
+        self.act_head = nn.Sequential(nn.Linear(d + 4, 256), nn.GELU(), nn.Linear(256, c["n_act"]))
+        self.temperature = mx.ones((3,))
+
+    def __call__(self, input_ids, attention_mask, marker_pos, marker_mask, qtype):
+        h = self.encoder(input_ids, attention_mask) + self.type_emb(qtype)[:, None, :]
+        mask = _mask(attention_mask)
+        for layer in self.head.layers:
+            h = layer(h, mask)
+        m = mx.take_along_axis(h, mx.maximum(marker_pos, 0)[:, :, None], axis=1)
+        return mx.where(marker_mask, self.scorer(m).squeeze(-1).astype(mx.float32), -1e4)
+
+
+class Session:
+    """Drop-in for the onnxruntime session the router calls: run(["logits"], feeds) -> [np.ndarray]."""
+
+    def __init__(self, weights_dir):
+        c = json.load(open(os.path.join(weights_dir, "config.json")))
+        self.model = DecisionModel(c)
+        self.model.load_weights(os.path.join(weights_dir, "model.safetensors"), strict=True)
+        q = c.get("quantization")
+        if q:
+            nn.quantize(self.model, group_size=q["group_size"], bits=q["bits"],
+                        class_predicate=lambda _, m: isinstance(m, nn.Linear) and m.weight.shape[-1] % q["group_size"] == 0)
+        self.model.eval()
+        mx.eval(self.model.parameters())
+
+    def run(self, names, feeds):
+        assert list(names) == ["logits"], names
+        am = feeds["attention_mask"]
+        f = {k: mx.array(v) for k, v in feeds.items()}
+        logits = self.model(f["input_ids"], None if am.all() else f["attention_mask"], f["marker_pos"],
+                            f["marker_mask"], f["qtype"])
+        return [np.array(logits)]

mlx/config.json

@@ -0,0 +1,50 @@
+{
+  "hidden_size": 768,
+  "num_attention_heads": 12,
+  "num_hidden_layers": 22,
+  "intermediate_size": 1152,
+  "layer_types": [
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention"
+  ],
+  "local_attention": 128,
+  "rope_parameters": {
+    "full_attention": {
+      "rope_theta": 160000.0,
+      "rope_type": "default"
+    },
+    "sliding_attention": {
+      "rope_theta": 160000.0,
+      "rope_type": "default"
+    }
+  },
+  "norm_eps": 1e-05,
+  "vocab_size": 50368,
+  "head_layers": 2,
+  "n_act": 2,
+  "model_type": "laya-decision-modernbert",
+  "quantization": {
+    "bits": 8,
+    "group_size": 64
+  }
+}

mlx/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b8b637753dbc51c3df33ce165444df95d692309b770cbee0759c346e4a9be0c1
+size 327984812