superdecider/engine.py

"""Inference engine: the model on Metal, one encoder pass per prompt, and canvas reads.

A decision read is: encode the prompt once (with the per-question bidirectional blocks the published
deployment enables), then one denoiser pass over the answer canvas whose logits at the slot positions give the
label probabilities. ``read`` does exactly that; ``generate`` runs the block-diffusion loop for chat and for
thoughts.
"""

from __future__ import annotations

import os
import subprocess
import threading
import time
from dataclasses import replace

import mlx.core as mx
import numpy as np
from transformers import AutoTokenizer

from .config import DiffusionGemmaConfig, GenerationConfig
from .generate import (
    EntropyBoundSampler,
    LinearTemperatureScheduleLogitsProcessor,
    MLXRng,
    generate as generate_loop,
)
from .modeling import DiffusionGemma, KVCache
from .weights import load_weights

CANVAS_LENGTH = 128
# Instruction prefixes kept for reuse by later requests. A prefix costs ~216 KiB per token of KV (24 sliding
# layers at 8 KV heads plus 6 full ones at 2), so the cap is small on purpose: 2 x 1024 tokens is ~430 MiB.
MIN_PREFIX_TOKENS = 64
MAX_PREFIX_TOKENS = 1024
MAX_PREFIX_CACHES = 2
# Instruction prefixes are reused across requests only when asked: re-encoding a suffix over a shorter key
# dimension shifts the last bf16 bits, which moves label probabilities by ~0.02 on a 1.5k-token prompt
# (median, worst 0.05 in 14 draws) without changing the top label. vLLM's prefix cache has the same property.
PREFIX_CACHE = os.environ.get("SD_PREFIX_CACHE", "1") != "0"
LOG = os.environ.get("SD_LOG") == "1"  # the served canvas (vLLM --diffusion-config canvas_length); the model config says 256


def swap_free_gb() -> float | None:
    """macOS swap free in GiB, else None."""
    try:
        out = subprocess.run(["sysctl", "-n", "vm.swapusage"], capture_output=True, text=True, timeout=2).stdout
        return float(out.split("free =")[1].split("M")[0]) / 1024
    except Exception:
        return None


class Engine:
    def __init__(self, model_dir: str, canvas_length: int = CANVAS_LENGTH, gen_config: GenerationConfig | None = None):
        self.config = DiffusionGemmaConfig.from_dir(model_dir)
        self.config.canvas_length = canvas_length
        self.state = load_weights(model_dir)
        self.model = DiffusionGemma(self.config, self.state)
        self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
        self.gen_config = gen_config or GenerationConfig.from_dir(model_dir)
        # One request at a time: MLX ops are not thread safe and a single 4-bit MoE already saturates the GPU.
        self.lock = threading.Lock()
        # Encoder KV of the last read, reused by the next one when it extends the same token prefix.
        self._cached_ids: list[int] | None = None
        self._cached_blocks: list[int] | None = None
        self._cached_prefix: KVCache | None = None
        # Instruction prefixes shared by requests that differ later (a fixed text in front of a variable body).
        self._prefixes: dict[tuple[int, ...], KVCache] = {}
        self._last_hit = 0  # encoder tokens the last read reused, for the SD_LOG line

    # -- prompt ------------------------------------------------------------------

    def encode(self, input_ids, block_ids=None, cache: KVCache | None = None, start: int = 0) -> KVCache:
        """Encoder pass over ``input_ids`` (absolute positions from ``start``) into ``cache``."""
        ids = mx.array([list(input_ids)], dtype=mx.int32)
        positions = mx.arange(start, start + ids.shape[1])
        _, cache = self.model.encode(ids, positions, cache, block_ids)
        return cache

    def encode_reusing_prefix(self, prompt_ids, block_ids=None) -> KVCache:
        """Encoder KV for ``prompt_ids``, re-encoding only the part no cache already covers.

        Three sources, longest first. Sliding layers keep the last ``sliding_window - 1`` keys only, so a cache
        can only be resumed at the token it ends on: reuse is either a prompt that *extends* the previous read,
        or one of the registered instruction prefixes, which are encoded at length < window and then extended
        exactly like a first prompt would be. Resuming inside a longer prompt would need keys the cache dropped,
        and the first new tokens would attend to a window shifted by the length difference, corrupting every
        sliding layer silently. Bidirectional question blocks rewrite their members' KV, so nothing at or after
        the first such block is reused.
        """
        ids = [int(t) for t in prompt_ids]
        blocks = None if block_ids is None else [int(b) for b in block_ids]
        limit = len(ids) if blocks is None else next((i for i, b in enumerate(blocks) if b >= 0), len(ids))
        cache, hit = None, 0

        # 1. the previous read's cache, when this prompt extends the same tokens. The comparison runs from the
        # front: the *tail* of one request is often the tail of the next too (same questions), so scanning back
        # from the end would accept a prompt that differs in the middle.
        if self._cached_ids is not None and len(ids) >= len(self._cached_ids):
            h, n = 0, len(self._cached_ids)
            while (
                h < n
                and ids[h] == self._cached_ids[h]
                and self._cached_blocks[h] < 0
                and (blocks is None or (blocks[h] == self._cached_blocks[h] and blocks[h] < 0))
            ):
                h += 1
            if h == n:
                cache, hit = self._cached_prefix.copy(), h

        # 2. a registered prefix shared with other requests, which may be shorter than the last prompt
        for prefix in [t for t in self._prefixes if len(t) > hit]:
            if len(prefix) <= len(ids) and len(prefix) <= limit and tuple(ids[: len(prefix)]) == prefix:
                kept = self._prefixes.pop(prefix)
                self._prefixes[prefix] = kept  # pop/push keeps it the most recent
                cache, hit = kept.copy(), len(prefix)
                break

        # 3. cold: what this prompt shares with the previous one is the instruction prefix of the next ones
        if PREFIX_CACHE and hit == 0 and self._cached_ids is not None:
            h, room = 0, min(len(ids), len(self._cached_ids), limit, MAX_PREFIX_TOKENS)
            while h < room and ids[h] == self._cached_ids[h]:
                h += 1
            if h >= MIN_PREFIX_TOKENS:
                cache = self.encode(ids[:h], None if blocks is None else block_ids[:h])
                self._prefixes.pop(tuple(ids[:h]), None)
                self._prefixes[tuple(ids[:h])] = cache.copy()  # the copy below must not reach the kept prefix
                while len(self._prefixes) > MAX_PREFIX_CACHES:
                    self._prefixes.pop(next(iter(self._prefixes)))
                hit = h

        self._last_hit = hit
        if hit < len(ids):
            cache = self.encode(ids[hit:], None if blocks is None else block_ids[hit:], cache, hit)
        self._cached_ids = ids
        self._cached_blocks = blocks if blocks is not None else [-1] * len(ids)
        self._cached_prefix = cache.copy()
        return cache

    def chat_ids(self, messages, thinking: bool = False) -> list[int]:
        out = self.tokenizer.apply_chat_template(
            messages, tokenize=True, add_generation_prompt=True, enable_thinking=thinking
        )
        return [int(t) for t in (out["input_ids"] if hasattr(out, "keys") else out)]

    # -- reads -------------------------------------------------------------------

    def read(
        self,
        prompt_ids,
        canvas_ids,
        slot_positions,
        label_ids,
        block_ids=None,
        steps: int = 1,
        pinned: dict[int, int] | None = None,
        seed: int = 0,
        samples: int = 1,
    ) -> np.ndarray:
        """Log-probabilities of every label at every slot, shape [samples, slots, labels].

        One denoiser pass over the canvas for the served default (one step); with more steps the canvas is
        denoised first, pinned positions held at their template values, and the last step's logits are read --
        the semantics of the deployment's ``diffusion_max_steps`` reads.
        """
        canvas_len = len(canvas_ids)
        positions = mx.arange(len(prompt_ids), len(prompt_ids) + canvas_len)
        slots = mx.array(slot_positions, dtype=mx.int32)
        labels = mx.array(sorted(set(label_ids)), dtype=mx.int32)
        index = {int(t): i for i, t in enumerate(label_ids if isinstance(label_ids, list) else labels.tolist())}

        log = time.perf_counter() if LOG else 0.0
        with self.lock:
            cache = self.encode_reusing_prefix(prompt_ids, block_ids)
            if LOG:
                mx.eval([a for a in cache.keys if a is not None])
                t_enc = time.perf_counter() - log
            td = time.perf_counter()
            if samples > 1 and steps <= 1 and pinned:
                # One forward for every draw. The draws differ only in the noise re-filled at the unpinned
                # (answer) positions -- what upstream's diffusion_samples/pinned pair asks for -- so they
                # batch like any canvas. Serialising them re-streams the 18 GB of experts per draw.
                logits = self.model.denoise(
                    self._draw_canvases(canvas_ids, pinned, seed, samples), cache, positions, None, slots, labels
                )
            else:
                logits = mx.concatenate(
                    [
                        self._denoise_canvas(
                            cache, canvas_ids, positions, steps, pinned, seed + k * 7919, slots, labels
                        )
                        for k in range(samples)
                    ],
                    axis=0,
                )
            logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
            arr = np.array(logprobs.astype(mx.float32))
            draws = [
                [[arr[k, i, index[t]] for t in label_ids] for i in range(len(slot_positions))]
                for k in range(arr.shape[0])
            ]
            if LOG:
                # A start-time memory floor can pass and the run still starve later: swap nearly full
                # is the last warning before a read slows to disk speed.
                swap = swap_free_gb()
                print(
                    f"read prompt={len(prompt_ids)} reused={self._last_hit} "
                    f"canvas={canvas_len} samples={samples} steps={steps} "
                    f"encode={t_enc:.2f}s denoise={time.perf_counter() - td:.2f}s total={time.perf_counter() - log:.2f}s"
                    + (f" WARN swap free {swap:.2f} GiB" if swap is not None and swap < 0.5 else ""),
                    flush=True,
                )
        return np.array(draws)

    def _draw_canvases(self, canvas_ids, pinned, seed, samples):
        """``samples`` copies of the canvas, each with fresh noise at the unpinned answer positions."""
        base = np.array(canvas_ids, dtype=np.int32)
        free = [p for p in range(len(base)) if p not in pinned]
        rows = []
        for k in range(samples):
            row = base.copy()
            if free:
                rng = np.random.default_rng((seed + k * 7919) & 0xFFFFFFFF)
                row[free] = rng.integers(0, self.model.text.vocab_size, len(free), dtype=np.int32)
            rows.append(row)
        return mx.array(np.stack(rows), dtype=mx.int32)

    def _denoise_canvas(self, cache, canvas_ids, positions, steps, pinned, seed, slots, labels):
        """Denoise the canvas ``steps`` times; returns the last step's logits at ``slots`` over ``labels``."""
        canvas = mx.array([canvas_ids], dtype=mx.int32)
        if steps <= 1:
            return self.model.denoise(canvas, cache, positions, None, slots, labels)

        sampler = EntropyBoundSampler(self.gen_config.sampler_config, canvas.shape[1], self.model.text.vocab_size)
        processor = LinearTemperatureScheduleLogitsProcessor(
            self.gen_config.t_min, self.gen_config.t_max, steps
        )
        rng = MLXRng(seed)
        pinned_mask = pinned_values = None
        if pinned:
            mask = np.zeros(canvas.shape[1], dtype=bool)
            values = np.zeros(canvas.shape[1], dtype=np.int32)
            for pos, tok in pinned.items():
                mask[pos], values[pos] = True, tok
            pinned_mask, pinned_values = mx.array(mask), mx.array(values[None, :])
        self_conditioning, logits = None, None
        for step in reversed(range(1, steps + 1)):
            logits = processor(self.model.denoise(canvas, cache, positions, self_conditioning), step)
            accepted = sampler.renoise_canvas(
                sampler.accept_canvas(canvas, rng.tokens(logits), logits), rng
            )
            if pinned_mask is not None:
                accepted = mx.where(pinned_mask[None], pinned_values, accepted)
            self_conditioning, canvas = logits, accepted
        return mx.take(logits, slots, axis=1)[..., labels]

    # -- generation --------------------------------------------------------------

    def generate(
        self,
        prompt_ids,
        max_new_tokens: int | None = None,
        max_denoising_steps: int | None = None,
        stop_ids=None,
        block_ids=None,
        token_logprobs: bool = False,
    ) -> list[int] | tuple[list[int], list[float]]:
        """Block-diffusion generation; returns the generated token ids (prompt excluded), and with
        ``token_logprobs`` the model's log-probability for each of them."""
        gcfg = self.gen_config
        if max_new_tokens is not None or max_denoising_steps is not None or stop_ids is not None:
            gcfg = replace(
                gcfg,
                max_new_tokens=max_new_tokens or gcfg.max_new_tokens,
                max_denoising_steps=max_denoising_steps or gcfg.max_denoising_steps,
                eos_token_id=list(stop_ids) if stop_ids else gcfg.eos_token_id,
            )
        with self.lock:
            out = generate_loop(
                self.model,
                mx.array([list(prompt_ids)], dtype=mx.int32),
                gcfg,
                MLXRng(0),
                block_ids,
                token_logprobs=token_logprobs,
            )
        if token_logprobs:
            ids, lp = out
            return [int(t) for t in np.array(ids)[0, len(prompt_ids):]], [float(v) for v in np.array(lp)[0]]
        return [int(t) for t in np.array(out)[0, len(prompt_ids):]]