mlx/src/solomon_mlx/engine.py

"""BF16 Metal execution with instance-owned adaptation and isolated question caches."""

import copy
import json
import threading
import time
from pathlib import Path

import mlx.core as mx
import numpy as np
from mlx import nn
from PIL import Image

from ._vendor.prompts import PAGE, SYSTEM
from .artifacts import ADAPTER_SHA, BASE_REVISION, HEADS_SHA, SOLOMON_REVISION, runtime_identity, sha256


class SwitchLoRA(nn.Module):
    def __init__(self, linear, a, b, context):
        super().__init__()
        self.linear, self.lora_a, self.lora_b = linear, a, b
        self._context = context

    def __call__(self, x):
        y = self.linear(x)
        start = self._context["start"]
        if start is None or start >= x.shape[1]:
            return y
        delta = (2.0 * ((x[:, start:].astype(mx.float32) @ self.lora_a) @ self.lora_b)).astype(y.dtype)
        return y + delta if start == 0 else mx.concatenate([y[:, :start], y[:, start:] + delta], axis=1)


def fork_cache(caches):
    """New cache containers and array handles; MLX owns copy-on-write storage.

    mx.array creates a distinct handle, so slice updates cannot change a prefix's
    Python array. Recurrent/window updates replace the branch's private slots.
    """
    from mlx_vlm.models.cache import ArraysCache, KVCache

    result = []
    for original in caches:
        if isinstance(original, ArraysCache):
            branch = ArraysCache(len(original.cache))
            branch.cache = [None if x is None else mx.array(x) for x in original.cache]
        elif isinstance(original, KVCache):
            branch = KVCache()
            branch.state = tuple(None if x is None else mx.array(x) for x in original.state)
        else:
            raise TypeError(f"Unsupported prefix cache: {type(original).__name__}")
        result.append(branch)
    return result


class Engine:
    def __init__(self, directory, *, chunk_size=2048, max_tokens=40960):
        from mlx_vlm.models.qwen3_vl.processing_qwen3_vl import Qwen3VLProcessor
        from mlx_vlm.utils import load_model

        self.directory = Path(directory).resolve()
        self.binding = json.loads((self.directory / "binding.json").read_text())
        if (
            self.binding.get("schema") != "solomon-mlx-binding-v1"
            or self.binding.get("base_revision") != BASE_REVISION
            or self.binding.get("solomon_revision") != SOLOMON_REVISION
        ):
            raise ValueError("Unrecognized or unpinned Solomon MLX binding")
        if self.binding["profile"] != "quality" or self.binding["dtype"] != "bfloat16":
            raise ValueError("This runtime currently accepts only the BF16 quality profile")
        for name, expected in self.binding["files"].items():
            path = (self.directory / name).resolve()
            if not path.is_relative_to(self.directory) or sha256(path) != expected:
                raise ValueError(f"Model artifact checksum mismatch: {name}")
        adapter, heads = self.directory / "adapter.safetensors", self.directory / "heads.npz"
        if sha256(adapter) != ADAPTER_SHA or sha256(heads) != HEADS_SHA:
            raise ValueError("Solomon checkpoint identity mismatch")
        if not 1 <= chunk_size <= 2048 or not 1 <= max_tokens <= 40960:
            raise ValueError("Invalid chunk size or context ceiling")
        weight_bytes = sum(
            (self.directory / name).stat().st_size
            for name in self.binding["files"]
            if name.endswith((".safetensors", ".npz"))
        )
        if weight_bytes + 4 * 2**30 > mx.device_info()["max_recommended_working_set_size"]:
            raise MemoryError(
                "Full BF16 weights and minimum workspace exceed this Mac’s recommended Metal working set"
            )
        self.chunk_size, self.max_tokens = chunk_size, max_tokens
        self.lock = threading.RLock()
        self.context = {"start": None}
        self.model = load_model(self.directory / "backbone", lazy=True, strict=True)
        self.processor = Qwen3VLProcessor.from_pretrained(
            str(self.directory / "backbone"), trust_remote_code=False
        )
        self.lm, self.t = self.model.language_model, self.processor.tokenizer
        self.pad = self.t.convert_tokens_to_ids("<|image_pad|>")
        weights = mx.load(str(adapter))
        for name in sorted({key.rsplit(".", 1)[0] for key in weights}):
            parts = name.split(".")
            if parts[:2] != ["model", "layers"]:
                raise ValueError(f"Unexpected adapter target: {name}")
            owner = self.lm.model.layers[int(parts[2])]
            for part in parts[3:-1]:
                owner = getattr(owner, part)
            linear = getattr(owner, parts[-1])
            a, b = weights[name + ".lora_a"].astype(mx.float32), weights[name + ".lora_b"].astype(mx.float32)
            if a.shape != (linear.weight.shape[1], 64) or b.shape != (64, linear.weight.shape[0]):
                raise ValueError(f"Adapter orientation/shape mismatch: {name}")
            setattr(owner, parts[-1], SwitchLoRA(linear, a, b, self.context))
        with np.load(heads, allow_pickle=False) as archive:
            keys = {k[:-7] for k in archive.files if k.endswith("/weight")}
            required_heads = {
                "boolean/state4",
                "entity/state4",
                "multilabel/state4",
                "ordered/threshold4",
                "single/choiceR",
                "single/choiceS",
                "single/sufficiency3",
                "ordered/choiceR",
                "ordered/choiceS",
                "ordered/sufficiency3",
            }
            if keys != required_heads:
                raise ValueError("All ten semantic heads are required")
            self.heads = {}
            for key in keys:
                w, b = archive[key + "/weight"], archive[key + "/bias"]
                if (
                    w.shape != (10, 5120)
                    or b.shape != (10,)
                    or not np.isfinite(w).all()
                    or not np.isfinite(b).all()
                ):
                    raise ValueError("Invalid semantic head")
                self.heads[key] = (mx.array(w, mx.float32), mx.array(b, mx.float32))
        self.model.freeze()
        self.model.eval()
        mx.eval(self.model.parameters(), self.heads)
        self.identity = runtime_identity(self.binding, chunk_size=self.chunk_size, max_tokens=self.max_tokens)

    def render(self, parts, block):
        content = ""
        for i, p in enumerate(parts):
            if "text" in p:
                content += ("\n" if i and "image" in parts[i - 1] else "") + p["text"]
            else:
                content += ("\n" if i and "text" in parts[i - 1] else "") + PAGE
        return self.t.apply_chat_template(
            [
                {"role": "system", "content": SYSTEM},
                {"role": "user", "content": "Document:\n" + content + "\n\n" + block},
            ],
            tokenize=False,
            add_generation_prompt=True,
            enable_thinking=False,
        )

    def expand(self, ids, counts):
        out, index = [], 0
        for token in ids:
            if token == self.pad:
                if index >= len(counts):
                    raise ValueError("Unexpected image placeholder in document text")
                out.extend([token] * counts[index])
                index += 1
            else:
                out.append(token)
        if index != len(counts):
            raise ValueError("Image placeholder count mismatch")
        return out

    def positions(self, start, count):
        return mx.broadcast_to(mx.arange(start, start + count)[None, None, :], (3, 1, count))

    def admit(self, count):
        if count < 1 or count > self.max_tokens:
            raise ValueError(f"{count} tokens exceeds the {self.max_tokens}-token scope ceiling")
        # Conservative allowance: BF16 attention KV + FP32 recurrent states and
        # chunk intermediates. This supplements the token ceiling, not a promise
        # of availability in the presence of other processes.
        temporary = 4 * 2**30 + count * 16 * 2 * 4 * 256 * 2
        limit = mx.device_info()["max_recommended_working_set_size"]
        if mx.get_active_memory() + temporary > limit:
            raise MemoryError("Insufficient recommended Metal working set for this request")

    def forward(self, ids, positions, cache, *, embeds=None, adapter_from=None, taps=()):
        hidden, captured = None, {}
        try:
            for start in range(0, len(ids), self.chunk_size):
                end = min(start + self.chunk_size, len(ids))
                self.context["start"] = None if adapter_from is None else max(0, adapter_from - start)
                last = end == len(ids)
                out = self.lm(
                    mx.array([ids[start:end]]),
                    cache=cache,
                    position_ids=positions[:, :, start:end],
                    inputs_embeds=None if embeds is None else embeds[:, start:end],
                    skip_logits=True,
                    return_hidden=last,
                    capture_layer_ids=list(taps) if last else None,
                )
                if last:
                    hidden = out.hidden_states[-1][0, -1].astype(mx.float32)
                    captured = {
                        str(i): self.lm.model.norm(h[:, -1:])[0, -1].astype(mx.float32)
                        for i, h in zip(sorted(set(taps)), out.hidden_states[:-1])
                    }
                    mx.eval(hidden, captured)
                mx.eval([c.state for c in cache])
            return hidden, captured
        finally:
            self.context["start"] = None

    def prefill(self, parts):
        with self.lock:
            prefill_started = time.perf_counter()
            text = self.render(parts, "X")
            boundary = text.rfind("\n\nX")
            if boundary < 0:
                raise ValueError("Missing document boundary")
            raw = self.t.encode(text[:boundary], add_special_tokens=False)
            if "text" in parts[-1]:
                raw = raw[:-1]
            vision_started = time.perf_counter()
            counts, grids, features = [], [], []
            for part in parts:
                if "image" not in part:
                    continue
                with Image.open(part["image"]) as image:
                    processed = self.processor.image_processor(images=[image.convert("RGB")])
                grid_np = np.asarray(processed["image_grid_thw"])
                count = int(grid_np.prod()) // self.model.config.vision_config.spatial_merge_size**2
                self.admit(len(raw) + sum(counts) + count - len(counts) - 1)
                grid = mx.array(grid_np)
                pixels = mx.array(np.asarray(processed["pixel_values"])).astype(
                    self.model.vision_tower.patch_embed.proj.weight.dtype
                )
                feature, _ = self.model.vision_tower(pixels, grid)
                mx.eval(feature)
                counts.append(count)
                grids.append(grid)
                features.append(feature)
            vision_seconds = time.perf_counter() - vision_started if counts else 0.0
            ids = self.expand(raw, counts)
            self.admit(len(ids))
            embeds, delta, feats, grid = None, 0, None, None
            if counts:
                feats, grid = mx.concatenate(features), mx.concatenate(grids)
                f = self.model.get_input_embeddings(
                    mx.array([ids]), mx.zeros((1,)), image_grid_thw=grid, cached_image_features=feats
                )
                embeds, positions = f.inputs_embeds, f.position_ids
                delta = int(np.asarray(f.rope_deltas).reshape(-1)[0])
                if delta != int(mx.max(positions).item()) + 1 - len(ids):
                    raise ValueError("Multimodal RoPE offset mismatch")
            else:
                positions = self.positions(0, len(ids))
            cache = self.lm.make_cache()
            started = time.perf_counter()
            self.forward(ids, positions, cache, embeds=embeds)
            return {
                "parts": copy.deepcopy(parts),
                "prefix_ids": ids,
                "cache": cache,
                "counts": counts,
                "rope_delta": delta,
                "features": feats,
                "grid": grid,
                "positions": positions,
                "prefill_seconds": time.perf_counter() - prefill_started,
                "language_prefill_seconds": time.perf_counter() - started,
                "vision_seconds": vision_seconds,
            }

    def ask(self, state, block, width, head, *, execution="cached", taps=()):
        with self.lock:
            if head not in self.heads or not 2 <= width <= 10:
                raise ValueError("Unknown semantic head or invalid width")
            ids = self.expand(
                self.t.encode(self.render(state["parts"], block), add_special_tokens=False), state["counts"]
            )
            self.admit(len(ids))
            p = len(state["prefix_ids"])
            if ids[:p] != state["prefix_ids"]:
                raise ValueError("Question token prefix differs from cached document")
            started = time.perf_counter()
            if execution == "cached":
                hidden, captured = self.forward(
                    ids[p:],
                    self.positions(p + state["rope_delta"], len(ids) - p),
                    fork_cache(state["cache"]),
                    adapter_from=0,
                    taps=taps,
                )
            elif execution == "full":
                embeds = None
                if state["counts"]:
                    f = self.model.get_input_embeddings(
                        mx.array([ids]),
                        mx.zeros((1,)),
                        image_grid_thw=state["grid"],
                        cached_image_features=state["features"],
                    )
                    embeds, positions = f.inputs_embeds, f.position_ids
                else:
                    positions = self.positions(0, len(ids))
                hidden, captured = self.forward(
                    ids, positions, self.lm.make_cache(), embeds=embeds, adapter_from=p, taps=taps
                )
            else:
                raise ValueError("Execution must be cached or full")
            w, b = self.heads[head]
            logits = (w @ hidden + b)[:width]
            mx.eval(logits)
            values = np.asarray(logits)
            if not np.isfinite(values).all():
                raise ValueError("Nonfinite trained-head output")
            result = {
                "letter_logits": values.tolist(),
                "head_key": head,
                "prompt_tokens": len(ids),
                "branch_tokens": len(ids) - p,
                "reused_prefix_tokens": p if execution == "cached" else 0,
                "seconds": time.perf_counter() - started,
            }
            if taps:
                result.update(
                    hidden=np.asarray(hidden).tolist(),
                    taps={k: np.asarray(v).tolist() for k, v in captured.items()},
                    token_ids=ids,
                )
            return result