"""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):]]