Accelerate Decider 0.8B, 2B, and 2B Vision on Apple Silicon

README.md

@@ -35,6 +35,8 @@ The v8 weights stay available under the Hub tag `v8`. `docs/HISTORY.md` describe
 pip install git+https://github.com/Mapika/decider          # or: git clone ... && pip install -e ".[serve]"
 ```
 
+On Apple Silicon, install the optional MLX/Metal kernel with `pip install -e ".[metal]"` from a clone. Without it, MPS inference uses the PyTorch implementation.
+
 ```python
 from decider.infer import Decider
 d = Decider("Mapika/decider-2b")                             # one CUDA GPU, bf16, about 4 GB; downloads the weights on first use

decider/bench/mps.py

@@ -0,0 +1,110 @@
+"""Offline full-model MPS check.
+
+Usage: python -m decider.bench.mps MODEL reference|conv|optimized
+
+Runs one mode per process without concurrent GPU work. JSON includes per-request
+probabilities and synchronized timings; model loading is excluded. ``reference``
+is the pure Transformers PyTorch path, ``conv`` adds only Decider's fused
+causal convolution, and ``optimized`` adds the MPS attention patch as well.
+"""
+import inspect
+import json
+import platform
+import statistics
+import sys
+import time
+from pathlib import Path
+
+import torch
+import transformers
+from huggingface_hub import snapshot_download
+
+from decider import mps_ops
+from decider.engine import patch_conv
+from decider.infer import Decider, Example, Q
+
+
+def configure(mode):
+    import transformers.models.qwen3_5.modeling_qwen3_5 as mq
+
+    if mode not in ("reference", "conv", "optimized"):
+        raise ValueError("mode must be reference, conv, or optimized")
+    # Avoid FLA/Triton in the reference arm: this is the pure Transformers path.
+    mq.torch_chunk_gated_delta_rule = inspect.unwrap(mq.torch_chunk_gated_delta_rule)
+    mq.causal_conv1d_fn = inspect.unwrap(mq.causal_conv1d_fn)
+    if mode == "reference":
+        mps_ops.patch_mps = lambda: False
+        return False, False
+    if mode == "conv":
+        mps_ops.patch_mps = lambda: False
+        patch_conv()
+        return False, True
+    return mps_ops.patch_mps(), True
+
+
+def main():
+    if len(sys.argv) != 3:
+        raise SystemExit(__doc__)
+    name, mode = sys.argv[1:]
+    assert torch.backends.mps.is_available(), "MPS required"
+    path = snapshot_download(name, local_files_only=True)
+    patch_result, conv_result = configure(mode)
+    if name.endswith("vision"):
+        from PIL import Image
+        from decider.vision.model import VisionDecisionModel
+
+        model = VisionDecisionModel(path, grad_ckpt=False).to("mps").eval()
+        cases = [(Image.new("RGB", (224, 224), color), Example(
+            "Identify the dominant color in the image.",
+            [Q("What color is shown?", ["red", "green", "blue"])]))
+            for color in ("red", "green", "blue")]
+
+        def run(case):
+            logits = model.slot_logits(model.prepare([case]))
+            return torch.softmax(logits, -1)[0, :3].cpu().tolist()
+    else:
+        model = Decider(path, device="mps", use_graphs=False)
+        cases = ["My card was charged twice for the same purchase.",
+                 "I cannot log in after resetting my password.",
+                 "I would like pricing for fifty licenses."]
+
+        def run(case):
+            return model.decide(case, [{"question": "Which department should handle this?",
+                "options": ["billing", "technical", "sales"]}])[0]["probs_list"]
+
+    results = []
+    with torch.inference_mode():
+        for index, case in enumerate(cases):
+            for _ in range(2):
+                run(case)
+                torch.mps.synchronize()
+            times = []
+            for _ in range(5):
+                torch.mps.synchronize()
+                start = time.perf_counter()
+                probs = run(case)
+                torch.mps.synchronize()
+                times.append((time.perf_counter() - start) * 1000)
+            assert all(torch.isfinite(torch.tensor(probs)))
+            results.append(dict(case=index, probs=probs, times_ms=times,
+                                median_ms=statistics.median(times)))
+    weights = model.lm if name.endswith("vision") else model.m.lm
+    print(json.dumps({
+        "model": name,
+        "snapshot_revision": Path(path).name,
+        "mode": mode,
+        "patch_mps": patch_result,
+        "patch_conv": conv_result,
+        "torch": torch.__version__,
+        "transformers": transformers.__version__,
+        "dtype": str(next(weights.parameters()).dtype),
+        "hardware": platform.machine(),
+        "macOS": platform.mac_ver()[0],
+        "warmups": 2,
+        "measurements_per_case": 5,
+        "results": results,
+    }, indent=2))
+
+
+if __name__ == "__main__":
+    main()

decider/engine.py

@@ -68,8 +68,13 @@ class Engine:
     def __init__(self, path, device="cuda", dtype=torch.bfloat16, use_graphs=True, max_ctx_tokens=1536,
                  compile=True, fp8=False, conv_patch=True):
         if conv_patch:
-            patch_conv()
+            if str(device).startswith("mps"):
+                from decider.mps_ops import patch_mps
+                patch_mps()
+            else:
+                patch_conv()
         self.m = DecisionModel(path, dtype=dtype, grad_ckpt=False).to(device).eval()
+        use_graphs = use_graphs and torch.device(device).type == "cuda"
         self.tok = self.m.tok; self.dev = device; self.use_graphs = use_graphs; self.max_ctx = max_ctx_tokens
         self.core, self.W = self.m.lm.model, self.m.lm.lm_head.weight[self.m.letters].detach().clone()
         self.cfg = dict(compile=compile, fp8=fp8, conv_patch=conv_patch, graphs=use_graphs)
@@ -82,7 +87,7 @@ def __init__(self, path, device="cuda", dtype=torch.bfloat16, use_graphs=True, m
         else:
             self._fwd_impl = self._fwd_eager
         self.graphs = {}                       # (B, T) -> (static_ids, static_out, graph)
-        self.pool = torch.cuda.graph_pool_handle() if use_graphs else None
+        self.pool = torch.cuda.graph_pool_handle() if (use_graphs and str(device).startswith("cuda")) else None
         self.stats = dict(graph_captures=0, forwards=0)
 
     def _fwd_eager(self, ids):

decider/evaluate.py

@@ -70,6 +70,7 @@ def agg(keys):
     ap.add_argument("--temperature", type=float, default=1.0)
     ap.add_argument("--limit", type=int, default=0)
     ap.add_argument("--engine", default="eager", help="eager | graph | compile | fp8")
+    ap.add_argument("--device", default=None, help="cuda | mps | cpu (auto-detected when omitted)")
     ap.add_argument("--max_options", type=int, default=0, help="0 = sub-sample large label sets to 10 (the original protocol); 255 = offer the full label set")
     ap.add_argument("--max_ctx", type=int, default=1536)
     ap.add_argument("--layout", default="state_first", help="state_first | schema_first")
@@ -80,16 +81,21 @@ def agg(keys):
         evals = {k: v for k, v in evals.items() if k in a.tasks.split(",")}
     if a.limit:
         evals = {k: v[:a.limit] for k, v in evals.items()}
+    device = a.device or ("cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu"))
+    dtype = torch.float16 if device == "mps" else torch.bfloat16
     eng = None
     if a.engine:
         from decider.engine import Engine
-        eng = Engine(a.model, compile=a.engine in ("compile", "fp8"), fp8=a.engine == "fp8", conv_patch=a.engine in ("compile", "fp8"))
+        eng = Engine(a.model, device=device, dtype=dtype, compile=a.engine in ("compile", "fp8"), fp8=a.engine == "fp8", conv_patch=a.engine in ("compile", "fp8"))
         m = eng.m
     else:
-        m = DecisionModel(a.model, grad_ckpt=False).cuda()
+        if device == "mps":
+            from decider.mps_ops import patch_mps
+            patch_mps()
+        m = DecisionModel(a.model, dtype=dtype, grad_ckpt=False).to(device).eval()
     os.makedirs(a.out, exist_ok=True)
     res, dump = run_eval(m, evals, bs=a.bs, temperature=a.temperature, engine=eng, max_options=a.max_options or None, max_ctx=a.max_ctx, layout=a.layout)
     agg = aggregate(res)
     print("[agg]", json.dumps(agg, indent=1))
-    json.dump(dict(results=res, agg=agg, model=a.model, engine=a.engine), open(f"{a.out}/eval.json", "w"), indent=1)
+    json.dump(dict(results=res, agg=agg, model=a.model, engine=a.engine, device=device, dtype=str(dtype)), open(f"{a.out}/eval.json", "w"), indent=1)
     pickle.dump(dump, open(f"{a.out}/preds.pkl", "wb"))

decider/infer.py

@@ -7,6 +7,8 @@
                     {"question": "How urgent is this?", "options": ["low", "medium", "high"]}])
     # -> [{'choice': 'billing', 'confidence': 0.97, 'probs': {...}}, {...}]
 """
+import logging
+
 import torch
 from decider.model import DecisionModel, collate
 from decider.prompt import build, MAX_OPTIONS
@@ -23,6 +25,7 @@ class Example:
     context: str; qs: list; task: str = "infer"; image: bytes = None
 
 
+logger = logging.getLogger(__name__)
 NEUTRAL_NONE = "not listed here"
 
 
@@ -53,9 +56,18 @@ def __call__(self, state, max_state_tokens=32768):
 
 
 class Decider:
-    """use_graphs=True (default on CUDA) routes scoring through decider.engine.Engine: shape-bucketed
-    CUDA graphs, ~7x lower single-request latency than eager. Set False for CPU or debugging."""
-    def __init__(self, path, device="cuda", dtype=torch.bfloat16, temperature=None, abstain_below=0.0, use_graphs=None):
+    """One-pass decisions with automatic CUDA, MPS, or CPU device selection.
+
+    CUDA uses shape-bucketed graphs by default. MPS defaults to float16 and uses
+    the optional MPS patch; CPU defaults to bfloat16. Set ``use_graphs=False``
+    for eager execution or debugging.
+    """
+    def __init__(self, path, device=None, dtype=None, temperature=None, abstain_below=0.0, use_graphs=None):
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
+        if dtype is None:
+            dtype = torch.float16 if str(device).startswith("mps") else torch.bfloat16
+        logger.info("Decider device=%s dtype=%s", device, dtype)
         import json, os
         cfg = {}
         try:                                          # model folder may carry decider_config.json (temperature, flags)
@@ -73,6 +85,9 @@ def __init__(self, path, device="cuda", dtype=torch.bfloat16, temperature=None,
             from decider.engine import Engine
             self.eng = Engine(path, device=device, dtype=dtype); self.m = self.eng.m
         else:
+            if str(device).startswith("mps"):
+                from decider.mps_ops import patch_mps
+                patch_mps()
             self.eng = None; self.m = DecisionModel(path, dtype=dtype, grad_ckpt=False).to(device).eval()
         self.dev = device; self.T = temperature; self.abstain_below = abstain_below
         self.name = "decider-" + str(cfg.get("version", "dev"))