Hugging Face's logo

GenbaLabs
/
super-decider-mlx

Text Generation
MLX
English
diffusion_gemma
nvfp4
diffusion-gemma
diffusion-lm
structured-output
decision-making
conversational
modelopt
Model card Files
xet
 Community
super-decider-mlx / superdecider /engine.py
GenbaLabs's picture
GenbaLabs
card, MLX package and runtime
88a6753 verified
Raw History Blame Contribute Delete
14.6 kB
"""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):]]