Read images and prompts up to 262,144 tokens

README.md

@@ -12,7 +12,8 @@ tags:
 
 <p align="center"><img src="assets/hero.png" alt="StartLux-Decision: a probability for every option" width="100%"></p>
 
-StartLux-Decision-27B answers typed questions about a state: pick one of several options, yes or no, or a rating on a scale. Every
+StartLux-Decision-27B answers typed questions about a state: pick one of several options, yes or no, or a rating on a scale. The
+state can be text, JSON or images, up to 262,144 tokens (256K). Every
 question comes back with a probability for each option. Requests and responses use the TypeSafe `/v1/systemone`
 format, so clients written for Jev work unchanged.
 
@@ -30,6 +31,7 @@ Sizes: [0.8B](https://huggingface.co/startlux-models/StartLux-Decision-0.8B) ·
 | Intern-Decision, average accuracy over seven suites | 91.82 |
 | Latency, one request with three questions | 102.3 ms |
 | Latency, one yes/no question | 50.7 ms |
+| Input | text, JSON or images, up to 262,144 tokens |
 
 </div>
 
@@ -84,6 +86,15 @@ The folder ships its own inference package, `startlux_decision/`, which is the f
 - For bulk work, `decide_batch` batches the questions of many requests together, which is several times faster than
   sending them one at a time.
 
+## Images and long inputs
+
+The weights include a vision tower, and the inference package uses it: a request can carry images as part of its
+evidence, `images=[...]` in Python (PIL images, file paths, encoded bytes, base64 strings or data URIs) or
+`"images": [...]` over HTTP (base64 strings or data URIs). `<image>` in a string state marks where each image goes.
+Prompts can run to 262,144 tokens (256K), the model's native context: a long state is read once, in chunks, and every
+question of the request branches off it. The MLX backend and the GGUF files read text only. `confidence` follows
+TypeSafe's definitions (for a choice, (p_max − 1/n) / (1 − 1/n)); the top probability is in `probabilities`.
+
 ## Usage
 
 ```bash
@@ -119,6 +130,14 @@ answers, usage = m.decide(state, questions)
 many = m.decide_batch([(state, questions), ...])
 ```
 
+With images, from the same folder:
+
+```python
+answers, usage = m.decide("Photo taken at delivery: <image>",
+                          {"damaged": {"type": "noul", "instructions": "Is the parcel damaged?"}},
+                          images=["parcel.jpg"])
+```
+
 ## License
 
 The model weights are released under [CC BY-NC 4.0](https://creativecommons.org/licenses/by-nc/4.0/): free for

requirements.txt

@@ -3,9 +3,14 @@ torch>=2.8
 transformers>=5.8
 huggingface_hub>=0.34
 safetensors>=0.4
+# image input: the checkpoints' image processor uses torchvision; install the build that matches your torch
+pillow>=10
+torchvision>=0.23
 # required on CUDA: without them transformers silently runs a plain torch path (>10x slower)
-flash-linear-attention>=0.5.2
-causal-conv1d>=1.7          # CUDA build; if pip compiles it, add --no-build-isolation
+flash-linear-attention>=0.5.2; sys_platform == "linux"
+causal-conv1d>=1.7; sys_platform == "linux"          # CUDA build; if pip compiles it, add --no-build-isolation
+# Apple Silicon: the forward pass runs in MLX instead (startlux_decision/mlx_model.py)
+mlx-lm>=0.31.3; sys_platform == "darwin" and platform_machine == "arm64"
 # fine-tuning
 peft>=0.20
 # evaluation (Typed Decisions parquet); the Decision Index kit is installed from its own repository

startlux_decision/check.py

@@ -3,18 +3,34 @@
     python -m startlux_decision.check /path/to/StartLux-Decision-4B
 
 Exits with status 1 when flash-linear-attention or causal-conv1d is missing or not importable; transformers would then
-fall back to a plain torch path that is more than ten times slower.
+fall back to a plain torch path that is more than ten times slower.  On Apple Silicon with mlx-lm installed the server
+runs the model with MLX instead, which needs neither.  Also says whether image input will work (torch backend only).
 """
+import importlib.util
+import os
 import sys
 
 from .model import fast_kernels_active
 
 
+def images_ready(path):
+    """"yes", or what image input still needs for the model in `path`."""
+    if not os.path.exists(os.path.join(path, "preprocessor_config.json")):
+        return "no, the model folder has no preprocessor_config.json"
+    missing = [m for m in ("PIL", "torchvision") if importlib.util.find_spec(m) is None]
+    return f"no, pip install {' '.join('pillow' if m == 'PIL' else m for m in missing)}" if missing else "yes"
+
+
 def main():
     if len(sys.argv) != 2:
         sys.exit(__doc__.strip())
+    from .server import mlx_available
+    if mlx_available():                              # Apple Silicon: mlx-lm has its own kernels for these layers
+        print("Apple Silicon: the server runs the model with MLX (text only); nothing else to install")
+        sys.exit(0)
     ok = fast_kernels_active(sys.argv[1])
     print("fast kernels: " + ("active" if ok else "NOT active, pip install flash-linear-attention causal-conv1d"))
+    print("images: " + images_ready(sys.argv[1]))
     sys.exit(0 if ok else 1)
 
 

startlux_decision/gguf_server.py

@@ -0,0 +1,115 @@
+"""A /v1/systemone server for a StartLux-Decision GGUF file run by llama.cpp.
+
+    llama-server -m StartLux-Decision-4B-Q8_0.gguf -ngl 99 -c 16384 --parallel 4 --port 8081
+    python -m startlux_decision.gguf_server --model-dir StartLux-Decision-4B-GGUF --llama http://127.0.0.1:8081 --port 8090
+
+It is StartLuxDecision with only the forward pass replaced by a llama-server call: prompt rendering, the option-letter
+readout, the per-type temperatures and wide choices are the package's own code, so the answers match the original
+model question by question up to the numerics of the GGUF file.  MODEL_DIR holds the tokenizer files, config.json and
+decision_config.json (the GGUF repositories ship them next to the .gguf files).  llama-server returns the
+log-probabilities of the whole vocabulary at the answer position; dividing them by the temperature and normalising over
+the listed options is the same as doing it on the logits."""
+import argparse
+import json
+import os
+import urllib.request
+import uuid
+from concurrent.futures import ThreadPoolExecutor
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+
+import torch
+
+from . import jevfmt as J
+from .model import StartLuxDecision
+
+
+class GGUFDecision(StartLuxDecision):
+    def __init__(self, path, llama, workers):                     # no torch model: llama-server runs the weights
+        from transformers import AutoTokenizer
+        cfg = json.load(open(os.path.join(path, "decision_config.json")))
+        self.tok = AutoTokenizer.from_pretrained(path)
+        self.letters = J.check_tokenizer(self.tok)
+        if self.letters != cfg["letter_token_ids"]:
+            raise ValueError("tokenizer letter ids differ from decision_config.json")
+        self.temperature = {k: float(v) for k, v in cfg["temperature_by_type"].items()}
+        wide = cfg.get("wide_choice", {})
+        self.group, self.keep, self.residual = wide.get("group", 25), wide.get("keep", 3), wide.get("residual", 1e-3)
+        self.max_length, self.graphs = 65536, {}
+        mc = json.load(open(os.path.join(path, "config.json")))
+        self.vocab = int(mc.get("text_config", mc).get("vocab_size", 262144))
+        self.url = llama.rstrip("/") + "/completion"
+        self.pool = ThreadPoolExecutor(workers)
+
+    def _letter_logprobs(self, ids, count):
+        """Next-token log-probabilities of the option letters at the last prompt position (full-vocabulary softmax;
+        dividing by the temperature and normalising over the options afterwards equals doing it on the logits)."""
+        targets = self.letters[:count]
+        for n in (256, self.vocab):
+            body = json.dumps({"prompt": ids, "n_predict": 1, "temperature": -1, "n_probs": n, "cache_prompt": False}).encode()
+            with urllib.request.urlopen(urllib.request.Request(self.url, body, {"Content-Type": "application/json"}), timeout=900) as r:
+                out = json.loads(r.read())
+            rows = out.get("completion_probabilities", out.get("probs"))
+            got = {t["id"]: t["logprob"] for t in rows[0]["top_logprobs"]}
+            if all(t in got for t in targets):
+                return [got[t] for t in targets]
+        raise RuntimeError("option letters missing from llama-server's log-probabilities")
+
+    def _logits(self, rows):
+        jobs = []
+        for r in rows:
+            order = [o["id"] for o in r["options"]]
+            ids, _ = J.render_ids(r, self.tok, order, max_length=self.max_length)
+            jobs.append((ids, len(order)))
+        outs = list(self.pool.map(lambda j: self._letter_logprobs(*j), jobs))
+        return [torch.tensor(o, dtype=torch.float32) for o in outs], sum(len(i) for i, _ in jobs)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model-dir", required=True)
+    ap.add_argument("--llama", required=True)
+    ap.add_argument("--port", type=int, required=True)
+    ap.add_argument("--workers", type=int, default=16)
+    ap.add_argument("--name", help="model name reported in responses (default: the GGUF directory name)")
+    a = ap.parse_args()
+    engine = GGUFDecision(a.model_dir, a.llama, a.workers)
+    a.name = a.name or os.path.basename(os.path.abspath(a.model_dir))
+
+    class Handler(BaseHTTPRequestHandler):
+        def _send(self, code, obj):
+            data = json.dumps(obj).encode()
+            self.send_response(code)
+            self.send_header("Content-Type", "application/json")
+            self.send_header("Content-Length", str(len(data)))
+            self.send_header("x-typesafe-request-id", uuid.uuid4().hex)
+            self.end_headers()
+            self.wfile.write(data)
+
+        def do_GET(self):
+            if self.path.rstrip("/") == "/v1/models":
+                return self._send(200, {"models": [{"name": a.name, "description": "StartLux-Decision typed decision model "
+                                                    "(GGUF)", "release_date": "2026-10-01"}]})
+            self._send(200, {"status": "ok", "model": a.name})
+
+        def do_POST(self):
+            try:
+                req = json.loads(self.rfile.read(int(self.headers.get("Content-Length", 0))))
+                answers, usage = engine.decide(req.get("state"), req["questions"])
+                self._send(200, {"answers": answers, "usage": usage, "model": a.name})
+            except (ValueError, KeyError) as e:                    # malformed request or one the model cannot answer
+                self._send(422, {"error": str(e), "detail": [{"loc": ["body"], "msg": str(e), "type": "value_error"}]})
+            except Exception as e:                                 # surfaced to the client; suites.py stops on it
+                self._send(500, {"error": repr(e)})
+
+        def log_message(self, *args):
+            pass
+
+    class Server(ThreadingHTTPServer):
+        request_queue_size = 1024                           # the default backlog of 5 resets connections under load
+        daemon_threads = True
+
+    Server(("127.0.0.1", a.port), Handler).serve_forever()
+
+
+if __name__ == "__main__":
+    main()

startlux_decision/mlx_int8.py

@@ -0,0 +1,232 @@
+"""int8 matmuls for MLXDecision(int8=True) on the neural accelerators of M5 and later GPUs.
+
+The large projections (MLP, attention and linear-attention in/out projections) run as int8 x int8 -> int32 through Metal
+4's MetalPerformancePrimitives matmul2d on the M5's neural accelerators.  Activations are quantized per token on every
+call, weights per output channel once at load.  SmoothQuant (alpha 0.75) first moves
+activation outliers into the weights, using activation ranges from the built-in CALIBRATION requests, which contain no
+benchmark items.  The embedding, the output head and projections too small for the kernel stay as they are.
+
+The quantize and matmul kernels are adapted from JetBrains/mlx-vlm (branch feature/int8-prefill, mlx_vlm/int8_prefill.py),
+released under the MIT License:
+
+    Copyright © 2025 Prince Canuma
+
+    Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
+    documentation files (the "Software"), to deal in the Software without restriction, including without limitation the
+    rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit
+    persons to whom the Software is furnished to do so, subject to the following conditions:
+
+    The above copyright notice and this permission notice shall be included in all copies or substantial portions of the
+    Software.
+
+    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
+    WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+    COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
+    OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+"""
+import re
+
+import mlx.core as mx
+import mlx.nn as nn
+
+ALPHA = 0.75                # SmoothQuant migration strength
+TILE = (64, 128, 4)         # output tile rows, columns and simdgroups; requests are tens to hundreds of tokens
+MIN_DIM = 1024              # projections smaller than this in either dimension stay as they are
+
+_HEADER = """
+#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
+using namespace mpp::tensor_ops;
+using namespace metal;
+"""
+
+# One threadgroup per row: absmax of x * inv_s over the row, then symmetric int8 with that row's scale.
+_QUANT_SRC = """
+    constexpr int K = {K};
+    constexpr int NTH = 256;
+    uint row = threadgroup_position_in_grid.x;
+    uint tid = thread_position_in_threadgroup.x;
+    uint lane = tid % 32;
+    uint sg = tid / 32;
+    const device {T}* xrow = x + size_t(row) * K;
+    float amax = 0.0f;
+    for (int i = tid; i < K; i += NTH) amax = max(amax, fabs(float(xrow[i]) * inv_s[i]));
+    amax = simd_max(amax);
+    threadgroup float tg_max[NTH / 32];
+    if (lane == 0) tg_max[sg] = amax;
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    amax = tg_max[lane % (NTH / 32)];
+    amax = simd_max(amax);
+    float scale = max(amax, 1e-8f) / 127.0f;
+    float inv = 1.0f / scale;
+    if (tid == 0) xs[row] = scale;
+    device int8_t* qrow = xq + size_t(row) * K;
+    for (int i = tid; i < K; i += NTH) qrow[i] = int8_t(clamp(rint(float(xrow[i]) * inv_s[i] * inv), -127.0f, 127.0f));
+"""
+
+# One threadgroup per TM x TN output tile; matmul2d loops over K.  Edge rows are bounds-checked, so any M works.
+_GEMM_SRC = """
+    constexpr int N = {N};
+    constexpr int K = {K};
+    constexpr int TM = {TM};
+    constexpr int TN = {TN};
+    uint2 tgid = threadgroup_position_in_grid.xy;
+    const int M = m_dim[0];
+    constexpr auto desc = matmul2d_descriptor(TM, TN, static_cast<int>(dynamic_extent), false, true, true,
+                                              matmul2d_descriptor::mode::multiply_accumulate);
+    matmul2d<desc, execution_simdgroups<{NS}>> op;
+    auto A = tensor<device int8_t, dextents<int32_t, 2>, tensor_inline>((device int8_t*)xq, dextents<int32_t, 2>(K, M));
+    auto B = tensor<device int8_t, dextents<int32_t, 2>, tensor_inline>((device int8_t*)wq, dextents<int32_t, 2>(K, N));
+    auto tA = A.slice(0, int(tgid.y) * TM);
+    auto tB = B.slice(0, int(tgid.x) * TN);
+    auto cT = op.get_destination_cooperative_tensor<decltype(tA), decltype(tB), int32_t>();
+#pragma unroll
+    for (uint16_t i = 0; i < cT.get_capacity(); ++i) if (cT.is_valid_element(i)) cT[i] = 0;
+    op.run(tA, tB, cT);
+#pragma unroll
+    for (uint16_t i = 0; i < cT.get_capacity(); ++i) {{
+        if (cT.is_valid_element(i)) {{
+            auto idx = cT.get_multidimensional_index(i);
+            int n = int(tgid.x) * TN + idx[0];
+            int m = int(tgid.y) * TM + idx[1];
+            if (m < M && n < N) out[size_t(m) * N + n] = bfloat(float(cT[i]) * xs[m] * ws[n]);
+        }}
+    }}
+"""
+
+_kernels = {}
+
+
+def supported():
+    """True on GPUs with neural accelerators (Apple GPU family 17, the M5, and later)."""
+    info = mx.device_info() if hasattr(mx, "device_info") else mx.metal.device_info()
+    m = re.match(r"applegpu_g(\d+)", info.get("architecture", ""))
+    return bool(m) and int(m.group(1)) >= 17
+
+
+def _quantize_rows(x, inv_s):
+    M, K = x.shape
+    key = ("q", K)
+    if key not in _kernels:
+        _kernels[key] = mx.fast.metal_kernel(name=f"sld_quant_{K}", input_names=["x", "inv_s"], output_names=["xq", "xs"],
+                                             source=_QUANT_SRC.format(K=K, T="bfloat"), header=_HEADER)
+    return _kernels[key](inputs=[x, inv_s], output_shapes=[(M, K), (M,)], output_dtypes=[mx.int8, mx.float32],
+                         grid=(M * 256, 1, 1), threadgroup=(256, 1, 1))
+
+
+def _matmul(xq, xs, wq, ws):
+    TM, TN, NS = TILE
+    M, K = xq.shape
+    N = wq.shape[0]
+    key = ("g", N, K)
+    if key not in _kernels:
+        _kernels[key] = mx.fast.metal_kernel(name=f"sld_i8mm_{N}_{K}", input_names=["xq", "xs", "wq", "ws", "m_dim"],
+                                             output_names=["out"], header=_HEADER,
+                                             source=_GEMM_SRC.format(N=N, K=K, TM=TM, TN=TN, NS=NS))
+    out, = _kernels[key](inputs=[xq, xs, wq, ws, mx.array([M], dtype=mx.int32)], output_shapes=[(M, N)],
+                         output_dtypes=[mx.bfloat16], grid=((N // TN) * 32 * NS, (M + TM - 1) // TM, 1),
+                         threadgroup=(32 * NS, 1, 1))
+    return out
+
+
+class Int8Linear(nn.Module):
+    """nn.Linear (or an mlx-lm QuantizedLinear) as a W8A8 matmul with SmoothQuant scales."""
+
+    def __init__(self, linear, act_max):
+        super().__init__()
+        w = linear.weight
+        if hasattr(linear, "scales"):
+            w = mx.dequantize(w, linear.scales, linear.biases, linear.group_size, linear.bits)
+        w = w.astype(mx.float32)
+        s = mx.clip(mx.power(mx.maximum(act_max, 1e-5), ALPHA) / mx.power(mx.maximum(mx.abs(w).max(axis=0), 1e-5), 1 - ALPHA),
+                    1e-3, 1e3)
+        w = w * s[None, :]
+        self.inv_s = 1.0 / s
+        self.ws = mx.maximum(mx.abs(w).max(axis=1), 1e-8) / 127.0
+        self.wq = mx.clip(mx.round(w / self.ws[:, None]), -127, 127).astype(mx.int8)
+        self.bias = linear.bias if "bias" in linear else None
+        self.freeze()
+
+    def __call__(self, x):
+        shape = x.shape
+        x2 = x.reshape(-1, shape[-1]).astype(mx.bfloat16)
+        y = _matmul(*_quantize_rows(x2, self.inv_s), self.wq, self.ws)
+        if self.bias is not None:
+            y = y + self.bias
+        return y.reshape(*shape[:-1], self.wq.shape[0])
+
+
+def _eligible(layer):
+    n, k = layer.weight.shape
+    k = k * 32 // layer.bits if hasattr(layer, "scales") else k
+    return n % TILE[1] == 0 and k % 32 == 0 and min(n, k) >= MIN_DIM
+
+
+def _projections(body):
+    return [(name, mod, cname, child) for name, mod in body.named_modules() if not isinstance(mod, _Recorder)
+            for cname, child in mod.children().items()
+            if isinstance(child, (nn.Linear, nn.QuantizedLinear)) and _eligible(child)]
+
+
+class _Recorder(nn.Module):
+    def __init__(self, inner):
+        super().__init__()
+        self.inner, self.amax = inner, None
+
+    def __call__(self, x):
+        m = mx.abs(x.reshape(-1, x.shape[-1]).astype(mx.float32)).max(axis=0)
+        self.amax = m if self.amax is None else mx.maximum(self.amax, m)
+        return self.inner(x)
+
+
+def convert(body, sequences):
+    """Replace the eligible projections of `body` by Int8Linear, calibrated on `sequences` (token id lists); returns
+    how many were replaced.  Layers are converted one at a time, so a lazily loaded model never holds two copies."""
+    slots = [(mod, cname, _Recorder(child)) for _, mod, cname, child in _projections(body)]
+    for mod, cname, rec in slots:
+        setattr(mod, cname, rec)
+    for ids in sequences:
+        body(mx.array([ids]))
+        mx.eval([rec.amax for _, _, rec in slots])
+    for i, (mod, cname, rec) in enumerate(slots):
+        new = Int8Linear(rec.inner, rec.amax)
+        mx.eval(new.parameters())
+        setattr(mod, cname, new)
+        if i % 16 == 15:
+            mx.clear_cache()
+    mx.clear_cache()
+    return len(slots)
+
+
+CALIBRATION = []            # (state, questions) pairs of the kinds of requests the models serve
+for _t in ("My parcel was marked delivered but never arrived.", "The app crashes when I open settings on Android 15.",
+           "I want a refund for the duplicate charge on invoice 7781.", "Can I change the shipping address after ordering?",
+           "Password reset emails never arrive.", "The discount code SPRING20 is rejected at checkout."):
+    CALIBRATION.append(({"ticket": _t}, {
+        "team": {"type": "choice", "instructions": "Which team should handle this ticket?",
+                 "criteria": {"billing": "Payments, refunds and invoices", "shipping": "Delivery and tracking",
+                              "technical": "App, login and account problems"}},
+        "urgent": {"type": "noul", "instructions": "Should this ticket be answered today?"},
+        "severity": {"type": "score", "instructions": "How severe is the impact?",
+                     "criteria": ["cosmetic", "annoying", "blocks the customer"]}}))
+for _n in (5, 20, 60):
+    CALIBRATION.append(({"page": " ".join(f"Item {i}: {['laptop', 'desk lamp', 'usb cable', 'office chair'][i % 4]} "
+                                          f"${(i * 37) % 900}.99, {'in stock' if i % 3 else 'backorder'}." for i in range(_n)),
+                         "goal": "Buy the cheapest in-stock chair."},
+                        {"next": {"type": "choice", "instructions": "Which item should be clicked next?",
+                                  "criteria": {f"item{i}": f"Item {i}" for i in range(min(_n, 12))}},
+                         "done": {"type": "noul", "instructions": "Is the goal already met?"}}))
+for _r in ("The hotel was spotless and the staff helpful.", "Food arrived cold and two items were missing.",
+           "Average experience, nothing special.", "Absolutely terrible support, never again!"):
+    CALIBRATION.append(({"review": _r}, {
+        "stars": {"type": "score", "instructions": "How satisfied is the customer?", "criteria": ["1", "2", "3", "4", "5"]},
+        "topic": {"type": "choice", "instructions": "What is the review mainly about?",
+                  "criteria": {"service": "staff and support", "product": "food or room", "price": "cost",
+                               "delivery": "delivery"}}}))
+for _q, _o in (("Which planet is known as the Red Planet?", ["Venus", "Mars", "Jupiter", "Saturn"]),
+               ("What is 17 * 3?", ["41", "51", "57", "61"]), ("Who wrote Hamlet?", ["Dickens", "Shakespeare", "Austen", "Twain"])):
+    CALIBRATION.append(({"quiz": _q}, {"answer": {"type": "choice", "instructions": _q, "criteria": {o.lower(): o for o in _o}}}))
+CALIBRATION.append(({"log": "\n".join(f"2026-09-30 12:{i:02d} {'ERROR' if i % 7 == 0 else 'INFO'} worker-{i % 3} step {i}"
+                                      for i in range(40))},
+                    {"failing": {"type": "noul", "instructions": "Is any worker failing repeatedly?"},
+                     "worker": {"type": "choice", "instructions": "Which worker logs the most errors?",
+                                "criteria": {"w0": "worker-0", "w1": "worker-1", "w2": "worker-2"}}}))

startlux_decision/mlx_model.py

@@ -0,0 +1,166 @@
+"""StartLux-Decision on Apple Silicon: the forward pass in MLX, everything else unchanged.
+
+    python -m startlux_decision.server --model StartLux-Decision-4B              # --backend auto picks MLX on a Mac
+    python -m startlux_decision.server --model StartLux-Decision-4B --int8       # M5 and later: int8 matmuls
+
+MLXDecision is StartLuxDecision with the forward pass run by mlx-lm's Qwen3.5 implementation, which has Metal kernels
+for the linear-attention layers (PyTorch on MPS runs reference code there).  As in GGUFDecision, prompt rendering, the
+option-letter readout, the per-type temperatures and wide choices are the package's own code.
+
+All questions of a request run in one forward pass, one right-padded row per question.  Every prompt starts with the
+same system text: it runs once at start-up and its cache (attention keys and values, linear-attention conv and recurrent
+states) is copied to every row.  When the questions of a request share a long piece of evidence, the evidence runs once
+too and is kept, so asking about the same evidence again runs only the questions.  int8=True moves the large
+projections to int8 matmuls on the GPU's neural accelerators (mlx_int8.py).
+"""
+import json
+import os
+
+import mlx.core as mx
+import numpy as np
+import torch
+
+from . import jevfmt as J
+from .model import StartLuxDecision
+
+CACHE_LIMIT = 1 << 30       # MLX keeps freed buffers for reuse; unbounded, the cache grows with every new request shape
+SHARE_MIN = 192             # evidence tokens a shared prefix saves, above which running it separately pays off
+MAX_BATCH_TOKENS = 16384    # padded tokens per forward pass when many rows arrive at once (decide_batch)
+
+
+class MLXDecision(StartLuxDecision):
+    backend = "mlx"
+
+    def __init__(self, path, int8=False, max_length=65536):            # no torch model: MLX runs the weights
+        from mlx_lm import load
+        from transformers import AutoTokenizer
+
+        config = json.load(open(os.path.join(path, "config.json")))
+        if config.get("text_config", config).get("num_experts"):
+            raise NotImplementedError("mixture-of-experts checkpoints such as StartLux-Decision-35B-A3B run on the torch "
+                                      "backend: start the server with --backend torch")
+        cfg = json.load(open(os.path.join(path, "decision_config.json")))
+        self.tok = AutoTokenizer.from_pretrained(path)
+        self.letters = J.check_tokenizer(self.tok)
+        if self.letters != cfg["letter_token_ids"]:
+            raise ValueError("tokenizer letter ids differ from decision_config.json")
+        self.temperature = {k: float(v) for k, v in cfg["temperature_by_type"].items()}
+        wide = cfg.get("wide_choice", {})
+        self.group, self.keep, self.residual = wide.get("group", 25), wide.get("keep", 3), wide.get("residual", 1e-3)
+        self.max_length, self.graphs, self.fast_kernels = int(max_length), {}, True
+        self.pad = self.tok.pad_token_id if self.tok.pad_token_id is not None else 0
+        mx.set_cache_limit(CACHE_LIMIT)
+        self.model, _ = load(path, lazy=int8)
+        lm = getattr(self.model, "language_model", self.model)
+        self.body = lm.model
+        self.int8 = 0
+        if int8:
+            from . import mlx_int8
+            if not mlx_int8.supported():
+                raise RuntimeError("int8 needs the neural accelerators of an M5 or later GPU; run without --int8")
+            rows = [r for s, q in mlx_int8.CALIBRATION for _, r in self._split(s, q)[1]]
+            self.int8 = mlx_int8.convert(self.body, [t for t, _ in self._encode(rows)])
+        mx.eval(self.model.parameters())
+        head = lm.lm_head if hasattr(lm, "lm_head") else self.body.embed_tokens
+        idx = mx.array(self.letters)
+        w = head.weight[idx]
+        if hasattr(head, "scales"):                        # a quantized build: dequantize the letter rows only
+            w = mx.dequantize(w, head.scales[idx], head.biases[idx], head.group_size, head.bits)
+        self.letter_rows = w.astype(mx.float32)
+        self.snapshots = {}                                            # token prefix -> prompt cache right after it
+        probes = (({"x": "a"}, {"q": {"type": "noul", "instructions": "?"}}),
+                  ({"y": 1}, {"q": {"type": "noul", "instructions": "!"}}))
+        a, b = (self._encode([r for _, r in self._split(s, q)[1]])[0][0] for s, q in probes)
+        n = 0
+        while n < min(len(a), len(b)) and a[n] == b[n]:
+            n += 1
+        self.preamble = tuple(a[:n])                                   # the system text every prompt starts with
+        self._prefill(self.preamble)
+
+    def warm_up(self):
+        """Run a few request shapes once, so Metal compiles their kernels before the first real request."""
+        page = lambda n: {"page": " ".join(f"Row {i}: item {i}, qty {i % 7}." for i in range(n))}
+        yes_no = {"q": {"type": "noul", "instructions": "Is the item in stock?"}}
+        three = dict(yes_no,
+                     c={"type": "choice", "instructions": "Which item?", "criteria": {"a": "first", "b": "second"}},
+                     s={"type": "score", "instructions": "How full is the page?", "criteria": ["empty", "some", "full"]})
+        for n in (4, 40, 160):
+            self.decide(page(n), yes_no)
+            self.decide(page(n), three)
+
+    def _encode(self, rows):
+        out = []
+        for r in rows:
+            order = [o["id"] for o in r["options"]]
+            ids, _ = J.render_ids(r, self.tok, order, max_length=self.max_length)
+            out.append((ids, len(order)))
+        return out
+
+    def _forward(self, ids, last, cache=None):
+        h = self.body(ids, cache=cache)
+        return h[mx.arange(ids.shape[0]), last].astype(mx.float32) @ self.letter_rows.T
+
+    @staticmethod
+    def _copy(cache, rows):
+        """`cache` repeated to `rows` rows; the stored cache itself is never written to."""
+        from mlx_lm.models.cache import ArraysCache, KVCache
+        out = []
+        for c in cache:
+            if isinstance(c, KVCache):
+                n = KVCache()
+                if c.keys is not None:
+                    n.keys = mx.repeat(c.keys[..., :c.offset, :], rows, axis=0)
+                    n.values = mx.repeat(c.values[..., :c.offset, :], rows, axis=0)
+                n.offset = c.offset
+            else:
+                n = ArraysCache(len(c.cache))
+                n.cache = [None if a is None else mx.repeat(a, rows, axis=0) for a in c.cache]
+            out.append(n)
+        return out
+
+    def _prefill(self, prefix):
+        """The prompt cache right after `prefix`, resumed from the longest stored prefix it starts with, and stored."""
+        from mlx_lm.models.cache import make_prompt_cache
+        best = max((s for s in self.snapshots if prefix[:len(s)] == s), key=len, default=())
+        cache = self._copy(self.snapshots[best], 1) if best else make_prompt_cache(self.model)
+        if len(prefix) > len(best):
+            self.body(mx.array([prefix[len(best):]]), cache=cache)
+        mx.eval([c.state for c in cache])
+        self.snapshots[prefix] = cache
+        return cache
+
+    def _logits(self, rows):
+        enc = self._encode(rows)
+        toks = [t for t, _ in enc]
+        c = 0                                                          # common prefix, leaving every row one token
+        while c < min(len(t) for t in toks) - 1 and all(t[c] == toks[0][c] for t in toks):
+            c += 1
+        prefix = tuple(toks[0][:c])
+        pre = list(self.preamble)
+        if prefix != self.preamble and prefix in self.snapshots:       # the same evidence as before: questions only
+            groups = [(list(range(len(toks))), c, self.snapshots[prefix])]
+        elif (len(toks) - 1) * (c - len(pre)) >= SHARE_MIN:            # long shared evidence: run it once, keep it
+            for s in [s for s in self.snapshots if s != self.preamble]:
+                del self.snapshots[s]
+            groups = [(list(range(len(toks))), c, self._prefill(prefix))]
+        else:                                                          # one pass from the stored system text
+            start = len(pre) if all(t[:len(pre)] == pre for t in toks) else 0
+            base = self.snapshots[self.preamble] if start else None
+            groups, batch = [], []
+            for j in sorted(range(len(toks)), key=lambda j: -len(toks[j])):
+                if batch and (len(batch) + 1) * (len(toks[batch[0]]) - start) > MAX_BATCH_TOKENS:
+                    groups.append((batch, start, base))
+                    batch = []
+                batch.append(j)
+            groups.append((batch, start, base))
+        out = [None] * len(toks)
+        for idx, start, cache in groups:
+            n = max(len(toks[j]) - start for j in idx)
+            ids = np.full((len(idx), n), self.pad, dtype=np.int32)
+            for r, j in enumerate(idx):
+                ids[r, :len(toks[j]) - start] = toks[j][start:]
+            z = np.array(self._forward(mx.array(ids), mx.array([len(toks[j]) - start - 1 for j in idx]),
+                                       self._copy(cache, len(idx)) if cache is not None else None))
+            for r, j in enumerate(idx):
+                out[j] = torch.from_numpy(z[r, :enc[j][1]].copy())
+        return out, sum(len(t) for t in toks)

startlux_decision/model.py

@@ -10,17 +10,41 @@ to four questions of up to 1024 tokens each.  Replaying a graph removes the per-
 the JevK5 runtime, Apache-2.0).  Inputs are right-padded; every layer is causal and rows never mix, so padding after the
 last prompt token never reaches the position that is read.  Other requests run eagerly as one padded batch (inputs
 padded to one of PAD_LENGTHS); decide_batch() batches questions across many requests.
+
+Long prompts.  The questions of one request repeat the same evidence, so their prompts begin with the same tokens.  When
+that shared beginning is long it runs once: the decoder reads it in chunks of prefill_chunk tokens and keeps what each
+layer carries forward (the keys and values of the attention layers, the convolution and recurrent state of the
+linear-attention layers), then the rest of every question runs as one batch on top of it.  Chunks bound the memory for
+activations, so a single question over 128K tokens runs the same way and needs little more than the keys and values of
+its prompt.  A chunk attends to the keys before it without a mask and to its own keys causally, two fused attention
+calls whose outputs are combined by their log-sum-exp; the questions attend to the shared keys the same way, all rows
+against the one stored copy.  Prompts can be as long as the checkpoints' native context, 262,144 tokens.
+
+Images.  A request can carry images as part of its evidence.  The checkpoint's image processor resizes each to at most
+max_pixels pixels (1M by default) and cuts it into patches, the vision tower turns every 2x2 patches into one
+embedding, and these take the places of the image tokens in each question's prompt, with the multimodal positions the
+checkpoints use.  The vision tower runs once per request; the questions then take the same paths as text questions
+(one padded batch, or the shared beginning in chunks), without the graphs.  In a string state "<image>" marks where each
+image goes; without marks the images come first, then the state.
 """
+import base64
+import binascii
+import contextlib
+import functools
 import importlib
+import io
 import json
 import os
+import re
+import threading
 import warnings
 
 import torch
+import torch.nn.functional as F
 
 from . import jevfmt as J
 
-__all__ = ["StartLuxDecision", "load_model", "fast_kernels_active"]
+__all__ = ["StartLuxDecision", "load_model", "load_image", "fast_kernels_active", "choice_confidence", "score_confidence"]
 
 GRAPH_LENGTHS = (128, 192, 256, 320, 384, 512, 640, 768, 1024, 1536, 2048, 3072, 4096)
 GRAPH_ROWS = (1, 2, 3, 4)          # questions per request replayed as one graph
@@ -32,6 +56,64 @@ PAD_LENGTHS = (128, 192, 256, 320, 384, 448, 512, 640, 768, 896, 1024, 1280, 153
                5120, 6144, 7168, 8192, 10240, 12288, 14336, 16384, 20480, 24576, 28672, 32768, 40960, 49152, 57344, 65536)
 
 
+SHARE_MIN = 4096          # a shared beginning this long runs once for all questions
+SHARE_STEP = 1024         # the shared part is cut to a multiple of this; the tokens after it run with each question
+PREFILL_CHUNK = 32768     # default chunk: activations stay a few GiB, so the 35B-A3B reads 128K tokens on one 80 GB GPU
+CAUSAL = "startlux_causal"
+_SHARED = threading.local()   # .kv while questions run on a shared part: {attention layer index: (keys, values)}
+_LSE_KERNEL = []          # the fused attention here that also returns the log-sum-exp: "cudnn", "flash" or None
+IMAGE = "<image>"         # in a string state: where the next image goes
+_SLOT = "\x00IMAGE\x00"   # an image's place in a rendered prompt
+
+
+def load_image(x, paths=True):
+    """A PIL image, encoded image bytes, a base64 string or data URI, or (paths=True) a file path -> RGB PIL image."""
+    from PIL import Image, UnidentifiedImageError
+
+    if isinstance(x, Image.Image):
+        return x.convert("RGB")
+    raw = x
+    if isinstance(x, str):
+        if paths and not x.startswith("data:") and os.path.isfile(x):
+            with open(x, "rb") as f:
+                raw = f.read()
+        else:
+            try:
+                raw = base64.b64decode(re.sub(r"\s+", "", x.split(",", 1)[1] if x.startswith("data:") else x),
+                                      validate=True)
+            except (binascii.Error, IndexError):
+                raise ValueError("an image must be base64 or a data URI" + (", a file path" if paths else "") +
+                                 " or bytes") from None
+    if not isinstance(raw, (bytes, bytearray)):
+        raise ValueError(f"unsupported image value of type {type(x).__name__}")
+    try:
+        return Image.open(io.BytesIO(raw)).convert("RGB")
+    except (UnidentifiedImageError, OSError) as e:
+        raise ValueError(f"not a readable image: {e}") from None
+
+
+class _Row(list):
+    """The tokens of one prompt in a request with images, with each token's multimodal position (.pos, 3 x n) and the
+    index of the vision embedding each image token takes (.feature, -1 for text tokens)."""
+
+
+def choice_confidence(p):
+    """TypeSafe's confidence of a Choice answer: how far the top probability sits above an even split, from 0 (even) to
+    1 (all on one option), (p_max - 1/n) / (1 - 1/n)."""
+    n = len(p)
+    return 1.0 if n < 2 else min(1.0, max(0.0, (n * max(p) - 1) / (n - 1)))
+
+
+def score_confidence(p):
+    """TypeSafe's confidence of a Score answer: 1 minus the probability-weighted distance in levels from the most likely
+    level, relative to the same distance for an even spread measured from the middle level, floored at 0."""
+    n = len(p)
+    m = max(range(n), key=p.__getitem__)
+    spread = sum(v * abs(i - m) for i, v in enumerate(p))
+    even = sum(abs(i - (n - 1) / 2) for i in range(n)) / n
+    return max(0.0, 1 - spread / even) if even else 1.0
+
+
 def padded_length(n):
     """The length a batch whose longest input has n tokens is padded to."""
     for length in PAD_LENGTHS:
@@ -40,6 +122,125 @@ def padded_length(n):
     return -(-n // 8192) * 8192
 
 
+def _attend(query, key, value, causal, scale):
+    """softmax(q k^T) v and the log-sum-exp of each query's scores, from cuDNN or else the flash kernel; None when
+    neither runs here."""
+    if not _LSE_KERNEL:
+        for kind in ("cudnn", "flash", None):
+            try:
+                out = _attend_with(kind, query, key, value, causal, scale)
+            except RuntimeError:
+                continue
+            _LSE_KERNEL.append(kind)
+            return out
+    return _attend_with(_LSE_KERNEL[0], query, key, value, causal, scale)
+
+
+def _attend_with(kind, query, key, value, causal, scale):
+    if kind is None:
+        return None
+    if kind == "cudnn":
+        out = torch.ops.aten._scaled_dot_product_cudnn_attention(query, key, value, None, True, 0.0, causal, False,
+                                                                  scale=scale)
+    else:
+        groups = query.shape[1] // key.shape[1]
+        if groups > 1:
+            key, value = key.repeat_interleave(groups, 1), value.repeat_interleave(groups, 1)
+        out = torch.ops.aten._scaled_dot_product_flash_attention(query, key, value, 0.0, causal, False, scale=scale)
+    return out[0], out[1].reshape(query.shape[:3])
+
+
+def _merge(a, b):
+    """Attention over two disjoint key sets, from each part's output and log-sum-exp."""
+    (out_a, lse_a), (out_b, lse_b) = a, b
+    top = torch.maximum(lse_a, lse_b)
+    wa, wb = (lse_a - top).exp().unsqueeze(-1), (lse_b - top).exp().unsqueeze(-1)
+    return ((out_a.float() * wa + out_b.float() * wb) / (wa + wb)).to(out_a.dtype)
+
+
+def _causal_attention(module, query, key, value, attention_mask, dropout=0.0, scaling=None, **kwargs):
+    """Causal attention for queries at the end of the keys, without a mask: inputs on this path are right-padded, so
+    causality is the only mask needed.  The first chunk is one causal call.  A later chunk attends to the keys before
+    it without a mask and to its own keys causally; questions on a shared part attend to its stored keys (all rows
+    against the one copy) and to their own keys causally."""
+    from torch.nn.attention.bias import causal_lower_right
+
+    rows, heads, q, dim = query.shape
+    shared = getattr(_SHARED, "kv", None)
+    shared = shared.get(module.layer_idx) if shared else None
+    if shared is not None:
+        folded = query.transpose(0, 1).reshape(1, heads, rows * q, dim)
+        part = _attend(folded, shared[0], shared[1], False, scaling)
+        if part is not None:
+            out_a = part[0].reshape(heads, rows, q, dim).transpose(0, 1)
+            lse_a = part[1].reshape(heads, rows, q).transpose(0, 1)
+            out = _merge((out_a, lse_a), _attend(query, key, value, True, scaling))
+            return out.transpose(1, 2).contiguous(), None
+        key = torch.cat([shared[0].expand(rows, -1, -1, -1), key], dim=-2)
+        value = torch.cat([shared[1].expand(rows, -1, -1, -1), value], dim=-2)
+    k = key.shape[2]
+    if q == k:
+        out = F.scaled_dot_product_attention(query, key, value, is_causal=q > 1, scale=scaling, enable_gqa=True)
+    else:
+        before = _attend(query, key[:, :, :k - q], value[:, :, :k - q], False, scaling)
+        if before is not None:
+            out = _merge(before, _attend(query, key[:, :, k - q:], value[:, :, k - q:], True, scaling))
+        else:
+            out = F.scaled_dot_product_attention(query, key, value, attn_mask=causal_lower_right(q, k), scale=scaling,
+                                                 enable_gqa=True)
+    return out.transpose(1, 2).contiguous(), None
+
+
+def _register_causal_attention():
+    from transformers import AttentionInterface
+    from transformers.masking_utils import AttentionMaskInterface
+
+    AttentionInterface.register(CAUSAL, _causal_attention)
+    AttentionMaskInterface.register(CAUSAL, lambda *args, **kwargs: None)
+
+
+def _common_length(a, b):
+    n = min(len(a), len(b))
+    diff = (torch.tensor(a[:n]) != torch.tensor(b[:n])).nonzero()
+    return int(diff[0, 0]) if len(diff) else n
+
+
+@functools.cache
+def _shared_kv_layer():
+    from transformers.cache_utils import DynamicLayer
+
+    class SharedKV(DynamicLayer):
+        """Keys and values of one attention layer over the shared tokens.  The chunks of the shared part fill a buffer
+        sized for all of them; once frozen, nothing more is stored and a pass over the questions reads them from that
+        one copy."""
+
+        def __init__(self, length):
+            super().__init__()
+            self.length, self.end, self.frozen = length, 0, False
+
+        def update(self, key_states, value_states, *args, **kwargs):
+            if self.frozen:                     # the attention reads the shared keys through _SHARED
+                return key_states, value_states
+            if not self.is_initialized:
+                b, h, _, d = key_states.shape
+                self.keys = key_states.new_empty((b, h, self.length, d))
+                self.values = value_states.new_empty((b, h, self.length, d))
+                self.dtype, self.device, self.is_initialized = key_states.dtype, key_states.device, True
+            n = key_states.shape[-2]
+            self.keys[:, :, self.end:self.end + n] = key_states
+            self.values[:, :, self.end:self.end + n] = value_states
+            self.end += n
+            return self.keys[:, :, :self.end], self.values[:, :, :self.end]
+
+        def get_seq_length(self):
+            return self.end
+
+        def get_mask_sizes(self, query_length):
+            return self.end + query_length, 0
+
+    return SharedKV
+
+
 def fast_kernels_active(path):
     """True when transformers will use the fla / causal-conv1d kernels for the linear-attention layers of the model
     in `path`."""
@@ -55,18 +256,26 @@ def fast_kernels_active(path):
 
 def load_model(path, device):
     """The checkpoint's own transformers class in bf16 on `device`, and its text decoder (the stack without the
-    output head; checkpoints that also carry other towers keep the text decoder under .language_model)."""
+    output head; checkpoints that also carry other towers keep the text decoder under .language_model).  The experts
+    of a mixture-of-experts checkpoint run as grouped matrix multiplications: one kernel per projection for all
+    experts and no host synchronisation, so the CUDA graphs can record them."""
     import transformers
 
-    arch = transformers.AutoConfig.from_pretrained(path).architectures[0]
-    model = getattr(transformers, arch).from_pretrained(path, dtype=torch.bfloat16, device_map={"": device})
+    config = transformers.AutoConfig.from_pretrained(path)
+    extra = {"experts_implementation": "grouped_mm"} if getattr(config.get_text_config(), "num_experts", 0) else {}
+    model = getattr(transformers, config.architectures[0]).from_pretrained(path, dtype=torch.bfloat16,
+                                                                           device_map={"": device}, **extra)
     return model, getattr(model.model, "language_model", model.model)
 
 
 class StartLuxDecision:
-    """decide(state, questions) -> (answers, usage), answers in the TypeSafe /v1/systemone format."""
+    """decide(state, questions, images=None) -> (answers, usage), answers in the TypeSafe /v1/systemone format.
 
-    def __init__(self, path, device=None, max_length=65536, max_batch_tokens=65536, graphs=True):
+    images=False leaves the vision tower out (about 0.2 to 0.9 GB less GPU memory); max_pixels bounds the size of
+    each image after resizing."""
+
+    def __init__(self, path, device=None, max_length=262144, max_batch_tokens=65536, graphs=True,
+                 prefill_chunk=PREFILL_CHUNK, images=True, max_pixels=1 << 20, min_pixels=65536):
         from transformers import AutoTokenizer
 
         if not os.path.isdir(path):
@@ -92,9 +301,26 @@ class StartLuxDecision:
         self.body = body.eval()
         head = model.get_output_embeddings().weight
         self.letter_rows = head.index_select(0, torch.tensor(self.letters, device=head.device)).float()
-        del model                       # only the decoder and the letter rows of the output head are used
+        self.vision, self.image_processor, self.vision_note = None, None, "the checkpoint has no vision tower"
+        if not images:
+            self.vision_note = "the model was loaded with images=False"
+        elif getattr(model.model, "visual", None) is not None:
+            try:
+                from transformers import AutoImageProcessor
+
+                self.image_processor = AutoImageProcessor.from_pretrained(path)
+                self.image_processor.size = {"longest_edge": int(max_pixels), "shortest_edge": int(min_pixels)}
+                self.vision = model.model.eval()        # vision tower, multimodal positions; its decoder is self.body
+                self.image_token = model.config.image_token_id
+            except (OSError, ImportError, ValueError) as e:
+                self.vision_note = f"no image processor ({str(e).splitlines()[0][:200]})"
+        del model                       # the decoder, the letter rows of the output head and the vision tower are used
+        self._media = None              # the vision embeddings of the request being decided, when it has images
         self.pad = self.tok.pad_token_id
         self.max_length, self.max_batch_tokens = int(max_length), int(max_batch_tokens)
+        self.share, self.prefill_chunk = True, int(prefill_chunk)
+        self._kept = None                   # (tokens, cache) of the last shared part, for the rest of one request
+        _register_causal_attention()
         self.graphs = {}
         if graphs and self.device.type == "cuda" and os.environ.get("STARTLUX_GRAPHS", "1") != "0":
             self._capture()
@@ -143,6 +369,87 @@ class StartLuxDecision:
         res = out.cpu()
         return [res[j, :c].clone() for j, (_, _, c) in enumerate(enc)]
 
+    # ---- images
+    def _media_state(self, state, n):
+        """The state as text with the places of n images: at its "<image>" marks, else first."""
+        text = J.state_text(state)
+        marks = text.count(IMAGE)
+        if marks == n:
+            return text.replace(IMAGE, _SLOT)
+        if marks:
+            raise ValueError(f"the state has {marks} {IMAGE} marks for {n} images")
+        head = _SLOT if n == 1 else "\n".join(f"Image {k + 1}: {_SLOT}" for k in range(n))
+        return head + ("\n\n" + text if text.strip() else "")
+
+    @torch.no_grad()
+    def _encode_images(self, images):
+        """-> the vision embeddings of a request's images, and per image its (t, h, w) patch grid and token count."""
+        if getattr(self, "vision", None) is None:          # also the MLX and GGUF backends, which read text only
+            raise ValueError(f"this model takes no images: {getattr(self, 'vision_note', 'this backend reads text only')}")
+        enc = self.image_processor(images=[load_image(x) for x in images], return_tensors="pt")
+        grid = enc["image_grid_thw"]
+        with torch.autocast(self.device.type, dtype=torch.bfloat16):
+            out = self.vision.get_image_features(enc["pixel_values"].to(self.device), grid.to(self.device),
+                                                 return_dict=True)
+        merge = self.vision.visual.spatial_merge_size
+        return {"features": torch.cat(out.pooler_output, 0), "grid": grid, "tokens": (grid.prod(-1) // merge ** 2).tolist()}
+
+    def _render(self, row, order):
+        """A rendered record -> its prompt tokens; in a request with images, each image's slot becomes its tokens."""
+        if self._media is None:
+            return J.render_ids(row, self.tok, order, max_length=self.max_length)[0]
+        msgs, _ = J.messages(row, order)
+        content, parts = [], msgs[1]["content"].split(_SLOT)
+        for k, part in enumerate(parts):
+            if part:
+                content.append({"type": "text", "text": part})
+            if k < len(parts) - 1:
+                content.append({"type": "image"})
+        chat = [{"role": "system", "content": [{"type": "text", "text": msgs[0]["content"]}]},
+                {"role": "user", "content": content}]
+        text = self.tok.apply_chat_template(chat, tokenize=False, add_generation_prompt=True, enable_thinking=False)
+        if not text.endswith(J.THINK_OFF_SUFFIX):
+            raise ValueError("chat template lacks the thinking-off assistant prefix")
+        counts, ids, k = self._media["tokens"], _Row(), 0
+        for t in self.tok.encode(text, add_special_tokens=False):
+            if t == self.image_token:
+                if k == len(counts):
+                    raise ValueError("more image places than images")
+                ids.extend([t] * counts[k])
+                k += 1
+            else:
+                ids.append(t)
+        if k != len(counts):
+            raise ValueError("fewer image places than images")
+        if len(ids) > self.max_length:
+            raise ValueError(f"length {len(ids)} > {self.max_length}")
+        t = torch.tensor(ids)
+        image = t == self.image_token
+        pos, _ = self.vision.get_rope_index(t[None], mm_token_type_ids=image.int()[None],
+                                            image_grid_thw=self._media["grid"])
+        ids.pos = pos[:, 0]
+        ids.feature = torch.where(image, image.long().cumsum(0) - 1, -1)
+        return ids
+
+    def _forward(self, ids, rows, offset=0, **kwargs):
+        """The decoder over ids (right-padded rows: tokens offset.. of each prompt in rows).  In a request with images
+        the image tokens take their vision embeddings and every token its multimodal position."""
+        if self._media is None:
+            return self.body(input_ids=ids.to(self.device), **kwargs)
+        ids = ids.to(self.device)
+        embeds = self.body.get_input_embeddings()(ids)
+        features, n = self._media["features"], ids.shape[1]
+        pos = torch.empty((3, ids.shape[0], n), dtype=torch.long)
+        for j, t in enumerate(rows):
+            p = t.pos[:, offset:offset + n]
+            pos[:, j, :p.shape[1]] = p
+            pos[:, j, p.shape[1]:] = p[:, -1:] + torch.arange(1, n - p.shape[1] + 1)
+            f = t.feature[offset:offset + n]
+            at = (f >= 0).nonzero().squeeze(1)
+            if len(at):
+                embeds[j, at.to(self.device)] = features[f[at].to(self.device)].to(embeds.dtype)
+        return self.body(inputs_embeds=embeds, position_ids=pos.to(self.device), **kwargs)
+
     # ---- readout
     @torch.no_grad()
     def _logits(self, rows):
@@ -150,15 +457,19 @@ class StartLuxDecision:
         enc = []
         for i, r in enumerate(rows):
             order = [o["id"] for o in r["options"]]
-            ids, _ = J.render_ids(r, self.tok, order, max_length=self.max_length)
-            enc.append((i, ids, len(order)))
+            enc.append((i, self._render(r, order), len(order)))
         out, tokens = [None] * len(rows), 0
-        n = self._graph_length(enc) if self.graphs and enc else None
+        n = self._graph_length(enc) if self.graphs and enc and self._media is None else None
         if n is not None:                           # every question of the request in one replay
             for (i, t, _), z in zip(enc, self._graph_logits(enc, n)):
                 out[i] = z
                 tokens += len(t)
             return out, tokens
+        groups, enc = self._shared_groups(enc) if self.share else ([], enc)
+        for c, group in groups:                     # long prompts with a shared beginning
+            for (i, t, _), z in zip(group, self._shared_logits(group, c)):
+                out[i] = z
+                tokens += len(t)
         enc.sort(key=lambda x: -len(x[1]))
         start = 0
         while start < len(enc):
@@ -167,21 +478,125 @@ class StartLuxDecision:
             batch = enc[start:start + n]
             start += n
             ids = torch.full((len(batch), longest), self.pad, dtype=torch.long)
-            attn = torch.zeros_like(ids)
             for j, (_, t, _) in enumerate(batch):
                 ids[j, :len(t)] = torch.tensor(t)
-                attn[j, :len(t)] = 1
                 tokens += len(t)
-            ids, attn = ids.to(self.device), attn.to(self.device)
-            with torch.autocast(self.device.type, dtype=torch.bfloat16):
-                hidden = self.body(input_ids=ids, attention_mask=attn, use_cache=False, return_dict=True).last_hidden_state
-            pos = attn.sum(1) - 1
+            with torch.autocast(self.device.type, dtype=torch.bfloat16), self._causal():
+                hidden = self._forward(ids, [t for _, t, _ in batch], use_cache=False, return_dict=True).last_hidden_state
+            pos = torch.tensor([len(t) - 1 for _, t, _ in batch], device=self.device)
             h = hidden[torch.arange(len(batch), device=self.device), pos].float()
             logits = h @ self.letter_rows.T
             for j, (i, _, c) in enumerate(batch):
                 out[i] = logits[j, :c].cpu()
         return out, tokens
 
+    @contextlib.contextmanager
+    def _causal(self):
+        """Attention through _causal_attention: right-padded rows, causality as the only mask, no mask tensor."""
+        config = self.body.config
+        before, config._attn_implementation = config._attn_implementation, CAUSAL
+        try:
+            yield
+        finally:
+            config._attn_implementation = before
+
+    # ---- long prompts: the shared beginning once, in chunks, then the questions on top of it
+    def _shared_groups(self, enc):
+        """-> ([(c, rows that share their first c tokens)], other rows).  Rows are grouped by their first SHARE_MIN
+        tokens; c is cut to a multiple of SHARE_STEP and leaves every row at least one token of its own.  One row
+        alone is a group when it is longer than one chunk."""
+        by_start, rest = {}, []
+        for e in enc:
+            if len(e[1]) > SHARE_MIN:
+                by_start.setdefault(tuple(e[1][:SHARE_MIN]), []).append(e)
+            else:
+                rest.append(e)
+        groups = []
+        for group in by_start.values():
+            first = group[0][1]
+            c = min(len(t) for _, t, _ in group) - 1
+            for _, t, _ in group[1:]:
+                c = min(c, _common_length(first, t))
+            c = c // SHARE_STEP * SHARE_STEP
+            kept = self._kept[0] if self._kept else ()
+            if (kept and all(len(t) > len(kept) and tuple(t[:len(kept)]) == kept for _, t, _ in group)
+                    and max(len(t) for _, t, _ in group) - len(kept) <= self.prefill_chunk):
+                c = len(kept)                       # a later round of the same request (wide choices) reuses it
+            if c >= SHARE_MIN and (len(group) > 1 or len(first) > self.prefill_chunk):
+                groups.append((c, group))
+            else:
+                rest.extend(group)
+        return groups, rest
+
+    def _rows_cache(self, shared, rows):
+        """A cache for one pass of `rows` questions on top of the shared part: the attention layers read the shared
+        keys and values (one copy); each linear-attention layer starts every row from a copy of the shared state."""
+        from transformers import DynamicCache
+        from transformers.cache_utils import LinearAttentionLayer
+
+        cache = DynamicCache(config=self.body.config)
+        for i, layer in enumerate(shared.layers):
+            if not isinstance(layer, LinearAttentionLayer):
+                cache.layers[i] = layer
+                continue
+            copy = LinearAttentionLayer()
+            for name in ("conv_states", "recurrent_states"):
+                x = getattr(layer, name)
+                setattr(copy, name, x.repeat(rows, *[1] * (x.dim() - 1)))
+            copy.dtype, copy.device = layer.dtype, layer.device
+            copy.max_batch_size, copy.conv_kernel_size = rows, layer.conv_kernel_size
+            copy.is_conv_states_initialized = copy.is_recurrent_states_initialized = copy.has_previous_state = True
+            cache.layers[i] = copy
+        return cache
+
+    @torch.no_grad()
+    def _shared_logits(self, group, c):
+        """Letter logits for rows (rendered prompts) that share their first c tokens."""
+        from transformers import DynamicCache
+        from transformers.cache_utils import LinearAttentionLayer
+
+        SharedKV = _shared_kv_layer()
+        prefix = tuple(group[0][1][:c])
+        out = []
+        try:
+            with torch.autocast(self.device.type, dtype=torch.bfloat16), self._causal():
+                if self._kept and self._kept[0] == prefix:
+                    shared = self._kept[1]
+                else:
+                    self._kept = None
+                    shared = DynamicCache(config=self.body.config)
+                    for i, layer in enumerate(shared.layers):
+                        if not isinstance(layer, LinearAttentionLayer):
+                            shared.layers[i] = SharedKV(c)
+                    ids = torch.tensor([prefix])
+                    for start in range(0, c, self.prefill_chunk):
+                        self._forward(ids[:, start:start + self.prefill_chunk], [group[0][1]], start,
+                                      past_key_values=shared, use_cache=True)
+                    del ids
+                    self._kept = (prefix, shared)
+                _SHARED.kv = {}
+                for i, layer in enumerate(shared.layers):
+                    if isinstance(layer, SharedKV):
+                        layer.frozen = True
+                        _SHARED.kv[i] = (layer.keys[:, :, :layer.end], layer.values[:, :, :layer.end])
+                tails = [(t, k) for _, t, k in group]          # whole prompts; their tokens from c on run here
+                per_pass = max(1, self.max_batch_tokens // padded_length(max(len(t) - c for t, _ in tails)))
+                for start in range(0, len(tails), per_pass):
+                    batch = tails[start:start + per_pass]
+                    n = padded_length(max(len(t) - c for t, _ in batch))
+                    ids = torch.full((len(batch), n), self.pad, dtype=torch.long)
+                    for j, (t, _) in enumerate(batch):
+                        ids[j, :len(t) - c] = torch.tensor(t[c:])
+                    cache = self._rows_cache(shared, len(batch))
+                    hidden = self._forward(ids, [t for t, _ in batch], c, past_key_values=cache,
+                                           use_cache=True).last_hidden_state
+                    last = torch.tensor([len(t) - c - 1 for t, _ in batch], device=self.device)
+                    logits = hidden[torch.arange(len(batch), device=self.device), last].float() @ self.letter_rows.T
+                    out.extend(logits[j, :k].cpu() for j, (_, k) in enumerate(batch))
+        finally:
+            _SHARED.kv = None
+        return out
+
     def _probs(self, rows):
         logits, tokens = self._logits(rows)
         return [torch.softmax(z / self.temperature.get(r["type"], 1.0), -1).tolist() for r, z in zip(rows, logits)], tokens
@@ -196,7 +611,7 @@ class StartLuxDecision:
         try:
             eager, _ = self._probs(rows)
         finally:
-            self.graphs = graphs
+            self.graphs, self._kept = graphs, None
         return max(abs(a - b) for pg, pe in zip(graphed, eager) for a, b in zip(pg, pe))
 
     def _wide(self, state, spec):
@@ -234,13 +649,13 @@ class StartLuxDecision:
             return {"type": "noul", "noul": p[ids.index("true")]}
         if row["type"] == "score":
             levels = question.get("criteria") or []
-            levels = list(levels.values()) if isinstance(levels, dict) else list(levels)
-            return {"type": "score", "score": sum(i * v for i, v in enumerate(p)), "confidence": max(p),
+            levels = [levels[k] for k in J.score_keys(levels)] if isinstance(levels, dict) else list(levels)
+            return {"type": "score", "score": sum(i * v for i, v in enumerate(p)), "confidence": score_confidence(p),
                     "legend": {str(i): (levels[i] if i < len(levels) else str(i)) for i in range(len(p))},
                     "probabilities": {str(i): v for i, v in enumerate(p)}}
         dist = dict(zip(ids, p))
         best = max(dist, key=dist.get)
-        return {"type": "choice", "choice": best, "confidence": dist[best], "probabilities": dist}
+        return {"type": "choice", "choice": best, "confidence": choice_confidence(p), "probabilities": dist}
 
     def _split(self, state, questions):
         """-> (answers decided without the model, [(key, rendered row)], tokens spent on wide lists)"""
@@ -255,27 +670,42 @@ class StartLuxDecision:
                 p, n = self._wide(state, q)
                 tokens += n
                 best = max(p, key=p.get)
-                answers[k] = {"type": "choice", "choice": best, "confidence": p[best], "probabilities": p}
+                answers[k] = {"type": "choice", "choice": best, "confidence": choice_confidence(list(p.values())),
+                              "probabilities": p}
             else:
                 rows.append((k, J.from_systemone(state, q)))
         return answers, rows, tokens
 
-    def decide(self, state, questions):
-        """One request -> (answers, usage), answers in the TypeSafe /v1/systemone format."""
-        answers, rows, tokens = self._split(state, questions)
-        if rows:
-            probs, n = self._probs([r for _, r in rows])
-            tokens += n
-            for (k, row), p in zip(rows, probs):
-                answers[k] = self._answer(row, p, questions[k])
+    def decide(self, state, questions, images=None):
+        """One request -> (answers, usage), answers in the TypeSafe /v1/systemone format.  images: the request's
+        evidence images (PIL images, file paths, encoded bytes, base64 strings or data URIs), placed at the "<image>"
+        marks of a string state, or before the state."""
+        if images is not None and not isinstance(images, (list, tuple)):
+            images = [images]
+        try:
+            if images:
+                self._media = self._encode_images(images)
+                state = self._media_state(state, len(images))
+            answers, rows, tokens = self._split(state, questions)
+            if rows:
+                probs, n = self._probs([r for _, r in rows])
+                tokens += n
+                for (k, row), p in zip(rows, probs):
+                    answers[k] = self._answer(row, p, questions[k])
+        finally:
+            self._kept = self._media = None
         return answers, {"input_tokens": tokens, "output_tokens": 0}
 
     def decide_batch(self, requests):
-        """[(state, questions), ...] -> [answers, ...].  Every question of every request goes through one length-sorted
-        set of padded forward passes (up to max_batch_tokens each), which is much faster than calling decide() in a loop
-        when the requests are short.  Answers are the same as decide() up to bf16 rounding."""
+        """[(state, questions) or (state, questions, images), ...] -> [answers, ...].  Every question of every text
+        request goes through one length-sorted set of padded forward passes (up to max_batch_tokens each), which is
+        much faster than calling decide() in a loop when the requests are short; requests with images run one at a time.
+        Answers are the same as decide() up to bf16 rounding."""
         parts, flat = [], []
-        for state, questions in requests:
+        for state, questions, *images in requests:
+            if images and images[0]:
+                parts.append((self.decide(state, questions, images[0])[0], [], questions))
+                continue
             answers, rows, _ = self._split(state, questions)
             parts.append((answers, rows, questions))
             flat.extend(r for _, r in rows)
@@ -283,7 +713,7 @@ class StartLuxDecision:
         try:
             probs, _ = self._probs(flat) if flat else ([], 0)
         finally:
-            self.graphs = graphs
+            self.graphs, self._kept = graphs, None
         out, i = [], 0
         for answers, rows, questions in parts:
             for k, row in rows:

startlux_decision/server.py

@@ -2,32 +2,64 @@
 
     python -m startlux_decision.server --model /path/to/StartLux-Decision-4B --port 8090
 
-POST /v1/systemone  {"state": ..., "questions": {key: {"type", "instructions", "criteria"}}}
+POST /v1/systemone  {"state": ..., "questions": {key: {"type", "instructions", "criteria"}}, "images": [...]}
   -> {"answers": {key: answer}, "usage": {"input_tokens", "output_tokens"}, "model": ...}
-GET  /health        {"status", "model", "fast_kernels", "cuda_graphs"}
-GET  /v1/models
+     "images" is optional: base64 strings or data URIs, part of the evidence ("<image>" in a string state marks where
+     each goes, otherwise they come first); the torch backend only.
+GET  /health        {"status", "model", "backend", "fast_kernels", "cuda_graphs", "images", "max_length"}
+GET  /v1/models     {"models": [{"name", "description", "release_date"}]}
+Every response carries an x-typesafe-request-id header; errors are {"error": message, "detail": [{"loc", "msg", "type"}]}
+(400 for a malformed body, 422 for a request the model cannot answer).  confidence follows TypeSafe's definitions.
 Requests are served one at a time on one GPU.  At start-up one request runs through both the CUDA-graph path and the
 eager path; if their probabilities differ by more than 0.02 the graphs are dropped.
+
+--backend auto (the default) runs the model with MLX on Apple Silicon when mlx-lm is installed (mlx_model.py) and with
+PyTorch everywhere else; --int8 adds int8 matmuls on M5 and later Macs (mlx_int8.py).
 """
 import argparse
 import json
 import os
 import threading
 import time
+import uuid
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 
 
+def mlx_available():
+    """True on Apple Silicon with mlx-lm installed (requirements.txt installs it there)."""
+    try:
+        import mlx.core as mx
+        import mlx_lm  # noqa: F401
+        return mx.metal.is_available()
+    except ImportError:
+        return False
+
+
 def main():
     ap = argparse.ArgumentParser()
     ap.add_argument("--model", required=True, help="local StartLux-Decision directory")
     ap.add_argument("--host", default="127.0.0.1")
     ap.add_argument("--port", type=int, default=8090)
-    ap.add_argument("--device", help="cuda or cpu (default: cuda when available)")
+    ap.add_argument("--device", help="cuda or cpu for the torch backend (default: cuda when available)")
+    ap.add_argument("--backend", choices=("auto", "torch", "mlx"), default="auto",
+                    help="auto: MLX on Apple Silicon when mlx-lm is installed, else torch")
+    ap.add_argument("--int8", action="store_true", help="MLX on M5 and later: int8 matmuls on the neural accelerators")
     ap.add_argument("--name", help="model name reported in responses (default: the directory name)")
+    ap.add_argument("--max-length", type=int, default=262144, help="longest prompt in tokens (torch backend)")
+    ap.add_argument("--max-pixels", type=int, default=1 << 20, help="pixels per image after resizing (torch backend)")
+    ap.add_argument("--no-images", action="store_true", help="leave the vision tower out (torch backend)")
     a = ap.parse_args()
-    from .model import StartLuxDecision
-
-    engine = StartLuxDecision(a.model, device=a.device)
+    backend = a.backend
+    if backend == "auto":
+        backend = "mlx" if a.device is None and mlx_available() else "torch"
+    if backend == "mlx":
+        from .mlx_model import MLXDecision
+        engine = MLXDecision(a.model, int8=a.int8)
+        engine.warm_up()
+    else:
+        from .model import StartLuxDecision
+        engine = StartLuxDecision(a.model, device=a.device, max_length=a.max_length, images=not a.no_images,
+                                  max_pixels=a.max_pixels)
     demo_state = {"ticket": "I was charged twice for order #4411 and the app still shows it as unpaid."}
     demo_questions = {
         "team": {"type": "choice", "instructions": "Which team should handle this ticket?",
@@ -46,6 +78,10 @@ def main():
             engine.graphs = {}
     lock = threading.Lock()
     name = a.name or os.path.basename(os.path.normpath(a.model))
+    try:
+        released = json.load(open(os.path.join(a.model, "decision_config.json"))).get("release_date", "2026-10-01")
+    except (OSError, ValueError):
+        released = "2026-10-01"
 
     class Handler(BaseHTTPRequestHandler):
         def _send(self, code, obj):
@@ -53,42 +89,60 @@ def main():
             self.send_response(code)
             self.send_header("Content-Type", "application/json")
             self.send_header("Content-Length", str(len(data)))
+            self.send_header("x-typesafe-request-id", uuid.uuid4().hex)
             self.end_headers()
             self.wfile.write(data)
 
+        def _error(self, code, message):
+            self._send(code, {"error": message, "detail": [{"loc": ["body"], "msg": message, "type": "value_error"}]})
+
         def do_GET(self):
             path = self.path.rstrip("/")
             if path in ("/health", "/v1/health"):
-                return self._send(200, {"status": "ok", "model": name, "fast_kernels": engine.fast_kernels,
-                                        "cuda_graphs": len(engine.graphs)})
+                return self._send(200, {"status": "ok", "model": name, "backend": backend, "fast_kernels": engine.fast_kernels,
+                                        "cuda_graphs": len(engine.graphs), "images": getattr(engine, "vision", None) is not None,
+                                        "max_length": getattr(engine, "max_length", None)})
             if path == "/v1/models":
-                return self._send(200, {"models": [{"name": name, "description": "StartLux-Decision typed decision model"}]})
-            self._send(404, {"error": "not found"})
+                return self._send(200, {"models": [{"name": name, "description": "StartLux-Decision typed decision model",
+                                                    "release_date": released}]})
+            self._error(404, "not found")
 
         def do_POST(self):
             if self.path.rstrip("/") != "/v1/systemone":
-                return self._send(404, {"error": "not found"})
+                return self._error(404, "not found")
             try:
                 body = json.loads(self.rfile.read(int(self.headers.get("Content-Length") or 0)) or b"{}")
                 questions = body.get("questions")
                 if not isinstance(questions, dict) or not questions:
                     raise ValueError("questions must be a non-empty object")
+                images = body.get("images") or None
+                if images is not None:
+                    if not isinstance(images, list) or not all(isinstance(x, str) for x in images):
+                        raise ValueError("images must be a list of base64 strings or data URIs")
+                    if backend != "torch":
+                        raise ValueError("images need the torch backend")
+                    from .model import load_image
+                    images = [load_image(x, paths=False) for x in images]
             except ValueError as e:
-                return self._send(400, {"error": str(e)})
+                return self._error(400, str(e))
             t = time.perf_counter()
             try:
                 with lock:
-                    answers, usage = engine.decide(body.get("state"), questions)
+                    answers, usage = (engine.decide(body.get("state"), questions, images=images) if images else
+                                      engine.decide(body.get("state"), questions))
             except ValueError as e:          # unknown question type, prompt over the context limit
-                return self._send(422, {"error": str(e)})
+                return self._error(422, str(e))
             self._send(200, {"answers": answers, "usage": usage, "model": name,
                              "latency_ms": round(1000 * (time.perf_counter() - t), 2)})
 
         def log_message(self, *args):
             pass
 
-    print(f"{name} serving on http://{a.host}:{a.port}/v1/systemone (fast kernels: {engine.fast_kernels}, "
-          f"cuda graphs: {len(engine.graphs)})", flush=True)
+    detail = (f"MLX, int8 projections: {engine.int8}" if backend == "mlx" else
+              f"fast kernels: {engine.fast_kernels}, cuda graphs: {len(engine.graphs)}, "
+              f"images: {'yes' if engine.vision is not None else 'no, ' + engine.vision_note}, "
+              f"max length: {engine.max_length}")
+    print(f"{name} serving on http://{a.host}:{a.port}/v1/systemone ({detail})", flush=True)
     ThreadingHTTPServer((a.host, a.port), Handler).serve_forever()