simit 0.1.0__py3-none-any.whl

simit/backends/mot/engine.py

@@ -0,0 +1,844 @@
+"""Continuous-batching inference engine for MoT unified multimodal models.
+
+One background thread owns the GPU. Every step it packs, into a single
+forward pass:
+
+* one decode token for every running text generation,
+* prefill chunks (text: causal, image: bidirectional) of newly admitted
+  requests, within a token budget,
+* one denoising step (conditional and unconditional CFG branch) for every
+  running image generation.
+
+KV lives in a contiguous arena; each request owns a region. Long text
+prefixes (skill and router prompts) are kept after a request finishes and
+copied into new requests that share them (automatic prefix caching).
+"""
+
+from __future__ import annotations
+
+import hashlib
+import itertools
+import math
+import queue
+import threading
+import time
+import traceback
+from collections import OrderedDict
+from concurrent.futures import Future
+from dataclasses import dataclass, field
+from typing import Any, Callable, Optional
+
+import torch
+import torch.nn.functional as F
+
+from .modeling import AttnBatch, KVArena
+
+
+# ------------------------------------------------------------------ prompts
+@dataclass
+class TextChunk:
+    ids: list[int]
+
+
+@dataclass
+class ImageChunk:
+    """An understanding image: ``n`` tokens (start marker, ViT tokens, end
+    marker) sharing one rope position. ``encode`` returns their embeddings."""
+
+    n: int
+    key: str                      # content hash, for prefix-cache keys
+    payload: Any                  # backend-specific preprocessed image
+
+
+@dataclass
+class TextJob:
+    chunks: list                  # TextChunk / ImageChunk
+    max_new_tokens: int = 0       # 0 for scoring
+    temperature: float = 0.0
+    top_p: float = 1.0
+    repetition_penalty: float = 1.0
+    logprobs: bool = False
+    score_len: int = 0            # >0: teacher-forced scoring of the last `score_len` tokens
+    stop_ids: tuple = ()
+    on_tokens: Optional[Callable[[list], None]] = None   # streaming: called with all output ids so far
+    future: Future = field(default_factory=Future)
+
+
+@dataclass
+class ImageJob:
+    chunks: list                  # conditioning prompt (text chunks)
+    height: int
+    width: int
+    seed: Optional[int] = None
+    future: Future = field(default_factory=Future)
+
+
+# ------------------------------------------------------------- arena regions
+class RegionAllocator:
+    def __init__(self, capacity: int):
+        self.free = [(0, capacity)]  # sorted (start, size)
+
+    def alloc(self, size: int) -> Optional[int]:
+        for i, (start, n) in enumerate(self.free):
+            if n >= size:
+                if n == size:
+                    self.free.pop(i)
+                else:
+                    self.free[i] = (start + size, n - size)
+                return start
+        return None
+
+    def release(self, start: int, size: int):
+        if size <= 0:
+            return
+        self.free.append((start, size))
+        self.free.sort()
+        merged = []
+        for s, n in self.free:
+            if merged and merged[-1][0] + merged[-1][1] == s:
+                merged[-1] = (merged[-1][0], merged[-1][1] + n)
+            else:
+                merged.append((s, n))
+        self.free = merged
+
+
+_G = 16  # prefix-cache granularity (tokens)
+
+
+def _prefix_hashes(keys: list) -> list[tuple[int, str]]:
+    """Chained hashes of every _G-token prefix: [(length, hash), ...]."""
+    out, h = [], hashlib.sha1()
+    for i, k in enumerate(keys, 1):
+        h.update(str(k).encode() + b",")
+        if i % _G == 0:
+            out.append((i, h.copy().hexdigest()))
+    return out
+
+
+@dataclass
+class _CacheEntry:
+    start: int
+    length: int
+    owner: Any = None             # live request still using the region
+    ready: bool = False
+
+
+class PrefixCache:
+    def __init__(self, max_tokens: int):
+        self.max_tokens = max_tokens
+        self.entries: "OrderedDict[int, _CacheEntry]" = OrderedDict()   # id -> entry (LRU order)
+        self.index: dict[str, tuple[int, int]] = {}                       # hash -> (entry id, length)
+        self._ids = itertools.count()
+        self.tokens = 0
+        self.hits = 0
+        self.hit_tokens = 0
+
+    def lookup(self, hashes):
+        for length, h in reversed(hashes):
+            ref = self.index.get(h)
+            if ref and ref[0] in self.entries:
+                e = self.entries[ref[0]]
+                self.entries.move_to_end(ref[0])
+                return e, length
+        return None, 0
+
+    def add(self, start: int, length: int, hashes, owner) -> int:
+        eid = next(self._ids)
+        self.entries[eid] = _CacheEntry(start, length, owner=owner)
+        for n, h in hashes:
+            if n <= length:
+                self.index[h] = (eid, n)
+        self.tokens += length
+        return eid
+
+    def evictable(self):
+        for eid, e in self.entries.items():
+            if e.owner is None:
+                yield eid, e
+
+    def remove(self, eid):
+        e = self.entries.pop(eid)
+        self.tokens -= e.length
+        return e
+
+
+# --------------------------------------------------------------- running state
+class _Text:
+    def __init__(self, job: TextJob, region: int, cap: int, keys: list, n_prompt: int):
+        self.job, self.region, self.cap = job, region, cap
+        self.keys, self.n_prompt = keys, n_prompt
+        self.written = 0           # tokens whose KV is in the region
+        self.pos = 0               # next rope position (1-D)
+        self.chunk_i, self.chunk_off = 0, 0
+        self.out_ids: list[int] = []
+        self.out_lp: list[float] = []
+        self.score_lp: list[float] = []
+        self.last_token: Optional[int] = None
+        self.cache_eid: Optional[int] = None
+        self.done = False
+        if job.score_len:  # rows (key indices) whose next-token distribution scores the targets
+            self.score_lo = n_prompt - job.score_len - 1
+            self.score_hi = n_prompt - 2
+
+    @property
+    def prefilled(self):
+        return self.chunk_i >= len(self.job.chunks)
+
+
+class _Img:
+    def __init__(self, job: ImageJob, region, cap, uncond_region, n_lat, h, w, pos_ids, noise, steps):
+        self.job, self.region, self.cap = job, region, cap
+        self.uncond_region = uncond_region
+        self.n_lat, self.h, self.w = n_lat, h, w
+        self.lat_pos_ids = pos_ids
+        self.x = noise
+        self.steps = steps            # list of (t, dt)
+        self.step_i = 0
+        self.written = 0
+        self.pos = 0
+        self.chunk_i, self.chunk_off = 0, 0
+        self.n_prompt = sum(len(c.ids) for c in job.chunks)
+
+    @property
+    def prefilled(self):
+        return self.chunk_i >= len(self.job.chunks)
+
+
+def _positions_tensor(positions, device):
+    """[T] for 1-D rope, [3, T] for multimodal rope (positions given as triples)."""
+    t = torch.tensor(positions, dtype=torch.long)
+    if t.dim() == 2:
+        t = t.T.contiguous()
+    return t.to(device, non_blocking=True)
+
+
+class _DecodeGraphs:
+    """CUDA graphs for pure-decode steps (one new token per sequence), bucketed
+    by batch size. Padding rows write into a reserved dummy region."""
+
+    BUCKETS = (1, 2, 4, 8, 16, 24, 32, 48, 64)
+
+    def __init__(self, engine, max_k: int):
+        self.engine = engine
+        self.max_k = max_k
+        self.graphs = {}
+        self.pool = None
+        self.dummy_slot = engine.arena.capacity - 1  # last arena slot, never allocated
+        self.failed = False
+
+    def bucket(self, b):
+        for n in self.BUCKETS:
+            if n >= b:
+                return n
+        return None
+
+    def _capture(self, n):
+        eng = self.engine
+        lm, dev = eng.model.lm, eng.device
+        head_dim = lm.cfg["hidden_size"] // lm.cfg["num_attention_heads"]
+        st = {
+            "ids": torch.zeros(n, dtype=torch.long, device=dev),
+            "cos": torch.zeros(n, head_dim, dtype=eng.model.dtype, device=dev),
+            "sin": torch.zeros(n, head_dim, dtype=eng.model.dtype, device=dev),
+            "slots": torch.full((n,), self.dummy_slot, dtype=torch.long, device=dev),
+            "cu": torch.arange(n + 1, dtype=torch.int32, device=dev),
+            "kstart": torch.full((n + 1,), self.dummy_slot, dtype=torch.int32, device=dev),
+            "kused": torch.ones(n, dtype=torch.int32, device=dev),
+        }
+        batch = AttnBatch(slots=st["slots"], n_causal=n,
+                          causal=(st["cu"], st["kstart"], st["kused"], 1, self.max_k))
+
+        def run():
+            x = lm.model.embed_tokens(st["ids"])
+            h = lm(x, None, batch, eng.arena, cos_sin=(st["cos"], st["sin"]))
+            return lm.lm_head(h).float()
+
+        s = torch.cuda.Stream(device=dev)
+        s.wait_stream(torch.cuda.current_stream(dev))
+        with torch.cuda.stream(s):
+            for _ in range(2):
+                run()
+        torch.cuda.current_stream(dev).wait_stream(s)
+        g = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(g, pool=self.pool, stream=torch.cuda.current_stream(dev),
+                              capture_error_mode="thread_local"):
+            out = run()
+        self.pool = g.pool()
+        st["out"] = out
+        self.graphs[n] = (g, st)
+
+    def run(self, ids, positions, slots, kstart, kused):
+        b = len(ids)
+        n = self.bucket(b)
+        if n is None or self.failed:
+            return None
+        try:
+            if n not in self.graphs:
+                self._capture(n)
+        except Exception:
+            self.failed = True
+            return None
+        g, st = self.graphs[n]
+        dev = self.engine.device
+        lm = self.engine.model.lm
+        cos, sin = lm.rotary(_positions_tensor(positions, dev), self.engine.model.dtype)
+        st["ids"][:b].copy_(torch.tensor(ids, dtype=torch.long), non_blocking=True)
+        st["cos"][:b].copy_(cos)
+        st["sin"][:b].copy_(sin)
+        st["slots"][:b].copy_(torch.tensor(slots, dtype=torch.long), non_blocking=True)
+        st["slots"][b:].fill_(self.dummy_slot)
+        ks = torch.tensor(kstart + [self.dummy_slot] * (n - b + 1), dtype=torch.int32)
+        st["kstart"].copy_(ks, non_blocking=True)
+        ku = torch.tensor(kused + [1] * (n - b), dtype=torch.int32)
+        st["kused"].copy_(ku, non_blocking=True)
+        g.replay()
+        return st["out"][:b]
+
+
+def auto_kv_tokens(model, headroom_gib: float = 10.0) -> int:
+    """KV arena size from the GPU memory that is free -- including memory PyTorch
+    still caches from an earlier, closed engine -- minus headroom for activations."""
+    device = model.device
+    free, _ = torch.cuda.mem_get_info(device)
+    free += torch.cuda.memory_reserved(device) - torch.cuda.memory_allocated(device)
+    cfg = model.lm.cfg
+    per_token = 2 * cfg["num_hidden_layers"] * cfg["num_key_value_heads"] * (
+        cfg["hidden_size"] // cfg["num_attention_heads"]) * 2
+    tokens = int((free - headroom_gib * 2 ** 30) * 0.8 / per_token)
+    if tokens < 16384:
+        raise RuntimeError(f"not enough free GPU memory for the KV cache ({free / 2 ** 30:.1f} GiB free); "
+                           "close other engines on this GPU or pass kv_cache_tokens")
+    return tokens
+
+
+class MoTEngine:
+    """Owns the model on one device and serves TextJob / ImageJob requests."""
+
+    def __init__(self, model, *, kv_tokens: int, max_step_tokens: int = 16384,
+                 max_image_step_tokens: int = 12288, prefix_cache_fraction: float = 0.3,
+                 min_cache_prefix: int = 256, use_cuda_graphs: bool = True):
+        self.model = model          # adapter: see BagelModel for the interface
+        self.device = model.device
+        lm = model.lm
+        cfg = lm.cfg
+        self.arena = KVArena(lm.num_layers, kv_tokens, cfg["num_key_value_heads"],
+                             cfg["hidden_size"] // cfg["num_attention_heads"], model.dtype, self.device)
+        self.alloc = RegionAllocator(kv_tokens - 1)  # last slot: scratch for padded CUDA-graph rows
+        self.graphs = _DecodeGraphs(self, max_k=min(kv_tokens, 32768)) if use_cuda_graphs else None
+        self.cache = PrefixCache(int(kv_tokens * prefix_cache_fraction))
+        self.max_step_tokens = max_step_tokens
+        self.max_image_step_tokens = max_image_step_tokens
+        self.min_cache_prefix = min_cache_prefix
+        self.lock = threading.RLock()          # allocator + prefix cache
+        self.inbox = {"text": queue.Queue(), "image": queue.Queue()}
+        self.pending = {"text": [], "image": []}
+        self.texts: list[_Text] = []
+        self.images: list[_Img] = []
+        self.stats = {"steps": 0, "tokens": 0, "prefix_hit_tokens": 0, "image_steps": 0}
+        self._stop = False
+        # Text (latency-critical decoding) and image generation run in separate
+        # threads on separate CUDA streams so that they overlap on the GPU.
+        self._threads = [threading.Thread(target=self._run, args=(kind,), daemon=True, name=f"simit-mot-{kind}")
+                         for kind in ("text", "image")]
+        for t in self._threads:
+            t.start()
+
+    # -------------------------------------------------------------- submit
+    def submit(self, job) -> Future:
+        self.inbox["text" if isinstance(job, TextJob) else "image"].put(job)
+        return job.future
+
+    def close(self):
+        """Stop the engine threads and give the KV arena and graph memory back."""
+        self._stop = True
+        for q in self.inbox.values():
+            q.put(None)
+        for t in self._threads:
+            if t is not threading.current_thread():
+                t.join(timeout=10)
+        self.arena = None
+        self.graphs = None
+        self.cache = None
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    # ---------------------------------------------------------------- loop
+    def _active(self, kind):
+        return self.texts if kind == "text" else self.images
+
+    def _run(self, kind):
+        torch.set_grad_enabled(False)
+        # Decode is on every request's critical path: its kernels get scheduling
+        # priority over the (compute-bound) diffusion kernels.
+        lo, hi = torch.cuda.Stream.priority_range()
+        stream = torch.cuda.Stream(device=self.device, priority=hi if kind == "text" else lo)
+        inbox, pending = self.inbox[kind], self.pending[kind]
+        with torch.cuda.stream(stream):
+            while not self._stop:
+                idle = not (pending or self._active(kind))
+                try:
+                    job = inbox.get(timeout=0.2 if idle else 0)
+                    if job is not None:
+                        pending.append(job)
+                    while True:
+                        job = inbox.get_nowait()
+                        if job is not None:
+                            pending.append(job)
+                except queue.Empty:
+                    pass
+                if not (pending or self._active(kind)):
+                    continue
+                try:
+                    self._admit(kind)
+                    if self._active(kind):
+                        with torch.inference_mode():
+                            self._step(kind)
+                except Exception as e:  # never kill the engine thread; fail the affected requests
+                    tb = traceback.format_exc()
+                    active = self._active(kind)
+                    for r in active:
+                        if not r.job.future.done():
+                            r.job.future.set_exception(RuntimeError(f"engine step failed: {e}\n{tb}"))
+                        self._release(r)
+                    active.clear()
+                    torch.cuda.empty_cache()
+
+    # ------------------------------------------------------------ admission
+    def _region(self, size: int) -> Optional[int]:
+        start = self.alloc.alloc(size)
+        while start is None:
+            victim = next(self.cache.evictable(), None)
+            if victim is None:
+                return None
+            e = self.cache.remove(victim[0])
+            self.alloc.release(e.start, e.length)
+            start = self.alloc.alloc(size)
+        return start
+
+    def _keys(self, chunks):
+        keys, text_prefix = [], 0
+        seen_image = False
+        for c in chunks:
+            if isinstance(c, TextChunk):
+                keys.extend(c.ids)
+                if not seen_image:
+                    text_prefix += len(c.ids)
+            else:
+                seen_image = True
+                keys.extend(f"{c.key}:{i}" for i in range(c.n))
+        return keys, text_prefix
+
+    def _admit(self, kind):
+        still = []
+        for job in self.pending[kind]:
+            if job.future.cancelled():
+                continue
+            with self.lock:
+                ok = self._admit_text(job) if kind == "text" else self._admit_image(job)
+            if not ok:
+                still.append(job)
+        self.pending[kind][:] = still
+
+    def _admit_text(self, job: TextJob) -> bool:
+        keys, text_prefix = self._keys(job.chunks)
+        n_prompt = len(keys)
+        cap = n_prompt + job.max_new_tokens + 1
+        hashes = _prefix_hashes(keys[:text_prefix]) if text_prefix >= _G else []
+        entry, hit = self.cache.lookup([(n, h) for n, h in hashes if n < n_prompt])
+        if entry is not None and not entry.ready:
+            return False  # an identical long prefix is being computed right now: wait and reuse it
+        start = self._region(cap)
+        if start is None:
+            return False
+        seq = _Text(job, start, cap, keys, n_prompt)
+        if entry is not None and hit >= _G:
+            self.arena.copy(entry.start, start, hit)
+            self._skip(seq, hit)
+            self.stats["prefix_hit_tokens"] += hit
+        if text_prefix >= self.min_cache_prefix:
+            seq.cache_eid = self.cache.add(start, text_prefix, hashes, owner=seq)
+        self.texts.append(seq)
+        return True
+
+    def _skip(self, seq, n: int):
+        """Mark the first ``n`` prompt tokens (text only) as already written."""
+        seq.written = n
+        left = n
+        while left > 0:
+            c = seq.job.chunks[seq.chunk_i]
+            take = min(left, len(c.ids) - seq.chunk_off)
+            seq.chunk_off += take
+            left -= take
+            if seq.chunk_off == len(c.ids):
+                seq.chunk_i, seq.chunk_off = seq.chunk_i + 1, 0
+        seq.pos = self.model.position_after(seq.job.chunks, n)
+
+    def _admit_image(self, job: ImageJob) -> bool:
+        m = self.model
+        h, w = job.height // m.latent_downsample, job.width // m.latent_downsample
+        n_lat = h * w
+        n_prompt = sum(len(c.ids) for c in job.chunks)
+        cap = n_prompt + n_lat + 2
+        start = self._region(cap)
+        if start is None:
+            return False
+        uncond = self._region(n_lat + 2)
+        if uncond is None:
+            self.alloc.release(start, cap)
+            return False
+        gen = torch.Generator(device="cpu")
+        gen.manual_seed(job.seed if job.seed is not None else torch.seed() % (2 ** 31))
+        noise = torch.randn(n_lat, m.latent_dim, generator=gen).to(self.device, torch.float32)
+        self.images.append(_Img(job, start, cap, uncond, n_lat, h, w, m.latent_pos_ids(h, w).to(self.device),
+                                noise, m.denoise_schedule()))
+        return True
+
+    # ---------------------------------------------------------------- step
+    def _drop_cancelled(self, kind):
+        active = self._active(kind)
+        live = []
+        for r in active:
+            if r.job.future.cancelled():
+                self._release(r)
+            else:
+                live.append(r)
+        active[:] = live
+
+    def _step(self, kind):
+        self._drop_cancelled(kind)
+        m, dev = self.model, self.device
+        budget = self.max_step_tokens if kind == "text" else self.max_image_step_tokens
+        texts = self.texts if kind == "text" else []
+        images = self.images if kind == "image" else []
+        causal, full = [], []      # pieces: dict(kind, ...)
+        # 1) decode tokens
+        for s in texts:
+            if s.prefilled and s.job.score_len == 0 and s.last_token is not None:
+                causal.append(dict(kind="decode", seq=s, ids=[s.last_token], pos=m.text_positions(s.pos, 1),
+                                   logits=True))
+                budget -= 1
+        # 2) denoising steps (as many images as fit the step budget; the rest wait a step)
+        for g in images:
+            if g.prefilled:
+                t = g.steps[g.step_i][0]
+                need = (g.n_lat + 2) * (2 if m.cfg_active(t) else 1)
+                if need > budget and full:
+                    continue
+                full.append(dict(kind="latent", job=g, cond=True))
+                if m.cfg_active(t):
+                    full.append(dict(kind="latent", job=g, cond=False))
+                budget -= need
+        # 3) prefill chunks within the remaining budget
+        for r in texts + images:
+            if r.prefilled:
+                continue
+            while not r.prefilled and budget > 0:
+                c = r.job.chunks[r.chunk_i]
+                if isinstance(c, ImageChunk):
+                    if c.n > budget and (causal or full):
+                        break
+                    full.append(dict(kind="image", req=r, chunk=c, pos=r.pos, start_written=r.written))
+                    r.written += c.n
+                    r.pos = m.position_after_image(r.pos, c)
+                    r.chunk_i += 1
+                    budget -= c.n
+                else:
+                    take = min(len(c.ids) - r.chunk_off, max(budget, 1))
+                    ids = c.ids[r.chunk_off:r.chunk_off + take]
+                    pos = m.text_positions(r.pos, take)
+                    r.chunk_off += take
+                    if r.chunk_off == len(c.ids):
+                        r.chunk_i, r.chunk_off = r.chunk_i + 1, 0
+                    final = r.prefilled
+                    causal.append(dict(kind="prefill", req=r, ids=ids, pos=pos, start_written=r.written, final=final))
+                    r.written += take
+                    r.pos += take
+                    budget -= take
+            if budget <= 0:
+                break
+        if not causal and not full:
+            return
+        self._forward(causal, full, kind)
+
+    def _forward(self, causal, full, kind):
+        m, dev = self.model, self.device
+        lm = m.lm
+        if (self.graphs is not None and not full and causal and all(p["kind"] == "decode" for p in causal)
+                and max(p["seq"].written + 1 for p in causal) <= self.graphs.max_k):
+            seqs = [p["seq"] for p in causal]
+            logits = self.graphs.run([p["ids"][0] for p in causal], [p["pos"][0] for p in causal],
+                                     [s.region + s.written for s in seqs], [s.region for s in seqs],
+                                     [s.written + 1 for s in seqs])
+            if logits is not None:
+                self.stats["steps"] += 1
+                self.stats["tokens"] += len(seqs)
+                self._consume_logits(logits, [(i, 1, p) for i, p in enumerate(causal)])
+                self._finish(kind)
+                return
+        embeds, positions, slots = [], [], []
+        logit_rows, latent_rows = [], []    # (row_start, n, piece)
+        q_lens_c, kstart_c, kused_c = [], [], []
+        q_lens_f, kstart_f, kused_f = [], [], []
+        token_ids, token_rows = [], []
+        image_pieces = []
+        row = 0
+
+        def add_text(ids, pos, region, written, n_used):
+            nonlocal row
+            token_ids.extend(ids)
+            token_rows.extend(range(row, row + len(ids)))
+            embeds.append(None)
+            positions.extend(pos)
+            slots.extend(range(region + written, region + written + len(ids)))
+            row += len(ids)
+
+        for p in causal:
+            if p["kind"] == "decode":
+                s = p["seq"]
+                start_row = row
+                add_text(p["ids"], p["pos"], s.region, s.written, s.written + 1)
+                q_lens_c.append(1); kstart_c.append(s.region); kused_c.append(s.written + 1)
+                logit_rows.append((start_row, 1, p))
+            else:
+                r = p["req"]
+                start_row = row
+                n = len(p["ids"])
+                add_text(p["ids"], p["pos"], r.region, p["start_written"], p["start_written"] + n)
+                q_lens_c.append(n); kstart_c.append(r.region); kused_c.append(p["start_written"] + n)
+                if isinstance(r, _Text):
+                    a = p["start_written"]
+                    if r.job.score_len:
+                        lo, hi = max(a, r.score_lo), min(a + n - 1, r.score_hi)
+                        if lo <= hi:
+                            p["score_from"] = lo
+                            logit_rows.append((start_row + lo - a, hi - lo + 1, p))
+                    elif p["final"]:
+                        logit_rows.append((start_row + n - 1, 1, p))
+        n_causal = row
+        for p in full:
+            if p["kind"] == "image":
+                r, c = p["req"], p["chunk"]
+                start_row = row
+                image_pieces.append((start_row, p))
+                positions.extend(m.image_positions(p["pos"], c))
+                slots.extend(range(r.region + p["start_written"], r.region + p["start_written"] + c.n))
+                embeds.append(("image", start_row, c))
+                row += c.n
+                q_lens_f.append(c.n); kstart_f.append(r.region); kused_f.append(p["start_written"] + c.n)
+        # A latent block "[start_of_image] latents [end_of_image]" is attended
+        # bidirectionally, so its queries can be split into two pseudo-sequences
+        # over the same keys: the two markers (understanding expert) here and the
+        # latents (generation expert) last, giving each expert contiguous rows.
+        blocks = []
+        for p in full:
+            if p["kind"] == "latent":
+                g = p["job"]
+                region = g.region if p["cond"] else g.uncond_region
+                prefix = g.n_prompt if p["cond"] else 0
+                lpos = m.latent_positions(g.pos if p["cond"] else 0, g.h, g.w)
+                blocks.append((p, g, region, prefix, lpos))
+        for p, g, region, prefix, lpos in blocks:
+            n = g.n_lat + 2
+            token_ids.extend([m.start_of_image, m.end_of_image]); token_rows.extend([row, row + 1])
+            positions.extend([lpos[0], lpos[-1]])
+            slots.extend([region + prefix, region + prefix + n - 1])
+            row += 2
+            q_lens_f.append(2); kstart_f.append(region); kused_f.append(prefix + n)
+        n_und = row
+        for p, g, region, prefix, lpos in blocks:
+            n = g.n_lat + 2
+            embeds.append(("latent", row, g))
+            positions.extend(lpos[1:-1])
+            slots.extend(range(region + prefix + 1, region + prefix + 1 + g.n_lat))
+            latent_rows.append((row, g.n_lat, p))
+            row += g.n_lat
+            q_lens_f.append(g.n_lat); kstart_f.append(region); kused_f.append(prefix + n)
+
+        T = row
+        hidden = lm.cfg["hidden_size"]
+        x = torch.empty(T, hidden, dtype=m.dtype, device=dev)
+        if token_ids:
+            ids_t = torch.tensor(token_ids, dtype=torch.long, device=dev)
+            rows_t = torch.tensor(token_rows, dtype=torch.long, device=dev)
+            x[rows_t] = lm.model.embed_tokens(ids_t)
+        imgs = [e for e in embeds if e is not None and e[0] == "image"]
+        if imgs:
+            vit = m.encode_images([c.payload for _, _, c in imgs])
+            for (_, start_row, c), emb in zip(imgs, vit):
+                # start marker / ViT tokens / end marker
+                x[start_row] = lm.model.embed_tokens.weight[m.start_of_image]
+                x[start_row + 1:start_row + 1 + emb.shape[0]] = emb
+                x[start_row + c.n - 1] = lm.model.embed_tokens.weight[m.end_of_image]
+        for e in embeds:
+            if e is not None and e[0] == "latent":
+                _, start_row, g = e
+                t = g.steps[g.step_i][0]
+                x[start_row:start_row + g.n_lat] = m.embed_latents(g.x, t, g.lat_pos_ids)
+
+        pos_t = _positions_tensor(positions, dev)
+        slots_t = torch.tensor(slots, dtype=torch.long, device=dev)
+
+        def meta(q_lens, kstart, kused):
+            if not q_lens:
+                return None
+            cu = torch.tensor([0] + list(itertools.accumulate(q_lens)), dtype=torch.int32, device=dev)
+            ks = torch.tensor(kstart + [kstart[-1] + kused[-1]], dtype=torch.int32, device=dev)
+            ku = torch.tensor(kused, dtype=torch.int32, device=dev)
+            return (cu, ks, ku, max(q_lens), max(kused))
+
+        batch = AttnBatch(slots=slots_t, n_causal=n_causal, causal=meta(q_lens_c, kstart_c, kused_c),
+                          full=meta(q_lens_f, kstart_f, kused_f), n_und=n_und if blocks else None)
+        h = lm(x, pos_t, batch, self.arena)
+        self.stats["steps" if kind == "text" else "image_steps"] += 1
+        self.stats["tokens"] += T
+
+        # ---- text outputs
+        if logit_rows:
+            rows = torch.cat([torch.arange(r0, r0 + n, device=dev) for r0, n, _ in logit_rows])
+            logits = lm.lm_head(h[rows]).float()
+            self._consume_logits(logits, logit_rows)
+        # ---- latent outputs
+        if latent_rows:
+            self._consume_latents(h, latent_rows)
+        with self.lock:
+            for s in self.texts if kind == "text" else []:
+                if s.cache_eid is not None:
+                    e = self.cache.entries.get(s.cache_eid)
+                    if e is not None and not e.ready and s.written >= e.length:
+                        e.ready = True
+        self._finish(kind)
+
+    # ------------------------------------------------------------- outputs
+    def _consume_logits(self, logits, logit_rows):
+        off = 0
+        sample_rows, sample_seqs = [], []
+        for r0, n, p in logit_rows:
+            seq = p["seq"] if p["kind"] == "decode" else p["req"]
+            if seq.job.score_len:
+                # the row at key index i predicts keys[i + 1]
+                i0 = p["score_from"]
+                targets = torch.tensor(seq.keys[i0 + 1:i0 + 1 + n], device=logits.device)
+                lp = torch.log_softmax(logits[off:off + n], dim=-1)
+                seq.score_lp.extend(lp[torch.arange(n, device=logits.device), targets].tolist())
+                if i0 + n - 1 >= seq.score_hi:
+                    seq.done = True
+            else:
+                sample_rows.append(off)
+                sample_seqs.append(seq)
+            off += n
+        if not sample_rows:
+            return
+        lg = logits[sample_rows]
+        tokens, lps = self._sample(lg, sample_seqs)
+        for seq, tok, lp in zip(sample_seqs, tokens, lps):
+            if seq.last_token is not None:
+                seq.pos += 1  # the decode token just consumed was written at seq.pos
+                seq.written += 1
+            if tok in seq.job.stop_ids or len(seq.out_ids) >= seq.job.max_new_tokens:
+                seq.done = True
+                continue
+            seq.out_ids.append(tok)
+            if seq.job.logprobs:
+                seq.out_lp.append(lp)
+            seq.last_token = tok
+            if len(seq.out_ids) >= seq.job.max_new_tokens:
+                seq.done = True
+            if seq.job.on_tokens is not None and len(seq.out_ids) % 4 == 0:
+                try:
+                    seq.job.on_tokens(seq.out_ids)
+                except Exception:
+                    pass
+
+    def _sample(self, logits, seqs):
+        raw_lp = torch.log_softmax(logits, dim=-1)
+        lg = logits.clone()
+        for i, s in enumerate(seqs):
+            rp = s.job.repetition_penalty
+            if rp != 1.0 and s.out_ids:
+                seen = torch.tensor(sorted(set(s.out_ids)), device=lg.device)
+                vals = lg[i, seen]
+                lg[i, seen] = torch.where(vals < 0, vals * rp, vals / rp)
+        temps = torch.tensor([max(s.job.temperature, 0.0) for s in seqs], device=lg.device)
+        greedy = temps == 0
+        tokens = lg.argmax(-1)
+        if not bool(greedy.all()):
+            probs = torch.softmax(lg / temps.clamp(min=1e-5)[:, None], dim=-1)
+            for i, s in enumerate(seqs):
+                if s.job.top_p < 1.0 and not greedy[i]:
+                    sp, si = probs[i].sort(descending=True)
+                    keep = (sp.cumsum(0) - sp) <= s.job.top_p
+                    filt = torch.zeros_like(probs[i]).scatter(0, si[keep], sp[keep])
+                    probs[i] = filt / filt.sum()
+            sampled = torch.multinomial(probs, 1).squeeze(1)
+            tokens = torch.where(greedy, tokens, sampled)
+        chosen_lp = raw_lp.gather(1, tokens[:, None]).squeeze(1)
+        return tokens.tolist(), chosen_lp.tolist()
+
+    def _consume_latents(self, h, latent_rows):
+        m = self.model
+        by_job: dict[int, dict] = {}
+        r_first = latent_rows[0][0]   # latent rows are contiguous and in order
+        vel = m.latents_to_velocity(h[r_first:r_first + sum(n for _, n, _ in latent_rows)])
+        for r0, n, p in latent_rows:
+            v = vel[r0 - r_first:r0 - r_first + n]
+            by_job.setdefault(id(p["job"]), {"job": p["job"]})["cond" if p["cond"] else "uncond"] = v
+        for d in by_job.values():
+            g = d["job"]
+            t, dt = g.steps[g.step_i]
+            v = d["cond"]
+            if "uncond" in d:
+                v = m.apply_cfg(v, d["uncond"])
+            g.x = g.x - v.float() * dt
+            g.step_i += 1
+
+    def _finish(self, kind):
+        m = self.model
+        if kind == "text":
+            keep = []
+            for s in self.texts:
+                if s.done:
+                    self._release(s)
+                    if not s.job.future.done():
+                        s.job.future.set_result(s)
+                else:
+                    keep.append(s)
+            self.texts[:] = keep
+            return
+        done_imgs, keep = [], []
+        for g in self.images:
+            (done_imgs if g.step_i >= len(g.steps) else keep).append(g)
+        self.images[:] = keep
+        if done_imgs:
+            pics = m.decode_latents([(g.x, g.h, g.w) for g in done_imgs])
+            for g, pic in zip(done_imgs, pics):
+                self._release(g)
+                if not g.job.future.done():
+                    g.job.future.set_result(pic)
+
+    def _release(self, r):
+        with self.lock:
+            self._release_locked(r)
+
+    def _release_locked(self, r):
+        if isinstance(r, _Text):
+            if r.cache_eid is not None and r.cache_eid in self.cache.entries and r.written >= self.cache.entries[r.cache_eid].length:
+                e = self.cache.entries[r.cache_eid]
+                e.owner, e.ready = None, True
+                # keep only the cached prefix; give back the rest of the region
+                self.alloc.release(r.region + e.length, r.cap - e.length)
+                self._trim_cache()
+            else:
+                if r.cache_eid is not None and r.cache_eid in self.cache.entries:
+                    self.cache.remove(r.cache_eid)
+                self.alloc.release(r.region, r.cap)
+        else:
+            self.alloc.release(r.region, r.cap)
+            self.alloc.release(r.uncond_region, r.n_lat + 2)
+
+    def _trim_cache(self):
+        while self.cache.tokens > self.cache.max_tokens:
+            victim = next(self.cache.evictable(), None)
+            if victim is None:
+                return
+            e = self.cache.remove(victim[0])
+            self.alloc.release(e.start, e.length)