Blame for code/offload/dspark/worker/model.py

"""The DSpark drafter of the served checkpoint on one RTX 3090 (SM86): weights, per-request window state, draft step.

Semantics: NOTES.md (the engine's DSpark round, traced from the running image). Tensors: the exact checkpoint bytes
(checkpoint.Checkpoint refuses anything else). Representation changes are draft-only and listed in NOTES.md.

Public surface (all device tensors, no host syncs inside, so the draft step and the insert capture into CUDA graphs):
  * ``main_x(taps)``                      [n, 15360] bf16 target taps -> [n, 5120] bf16 conditioning features
  * ``insert(slot_rows, pos_rows, mx)``  write the context KV of accepted positions into request windows
  * ``draft(slots, anchors, last_pos, k)`` -> (draft ids [B, k] int64, confidences [B, k] fp32)
  * ``reset_slots(slots)``                empty request windows
Graphs: ``capture(batch_sizes, ks, n_outs)`` records one graph per (B, k, n_out) (and a margin twin);
``draft_graphed`` replays it.
"""

from __future__ import annotations

import json
import os
import time
from dataclasses import dataclass, field

import torch

from . import env  # noqa: F401
from . import fused as FU
from .checkpoint import Checkpoint, DrafterShape
from .exl3 import kernels as K
from .exl3.bank import Exl3Bank
from .exl3.plan import PROJ_DOWN, PROJ_GATE, PROJ_UP, Exl3MoEPlan
from .ops import (MarkovBias, apply_rope, block_attention, hc_collapse, hc_mixes, hc_post, kv_record_qdq, rms_norm,
                  route, rope_cos_sin)

MUL1 = 0x83DCD12D - (1 << 32)       # the mul1 codebook marker as int32


@dataclass
class Window:
    """What a draft block attends to: the ring rows of ``slots`` whose position lies in [lo, last] (int64 [B] each).
    ``kv`` [B, n_stages, W, D] / ``valid`` [B, W] are the gathered rows, filled only for the unfused (ops) path; the
    fused path reads the ring in place (fused.block_attention_ring)."""
    slots: torch.Tensor
    lo: torch.Tensor
    last: torch.Tensor
    kv: torch.Tensor | None = None
    valid: torch.Tensor | None = None


def facc_ext():
    """exllamav3_ext_facc: the same exllamav3 1.6.0 source built with fp32 MMA accumulation on sm_86 (stock 1.6.0
    accumulates exl3_gemm in fp16 on sm_86 only, EXL3_GEMM_H_ACC; ENV.md). Raises if it is not installed."""
    import exllamav3_ext_facc  # type: ignore[import-not-found]
    return exllamav3_ext_facc


class Exl3Linear:
    """One EXL3 matrix on exllamav3 1.6.0: fp16 input, fp32 output (the plugin's prompt-size path; NOTES.md).

    ``prescale``: divide the input by the power of two that brings max |x| under 2^14 and multiply the output back
    (exact; the plugin does this for main_proj, whose input is the target's raw residual stream).
    ``ext`` / ``force_shape``: call ``ext.exl3_gemm`` directly (``ext`` defaults to the stock exllamav3_ext) with a
    fixed kernel shape instead of LinearEXL3's timing-based autotune: the kernel, and so each row's result, is then
    the same for every row count the shape covers (1..16 for the 16-row shapes) and in every process."""

    def __init__(self, trellis: torch.Tensor, suh: torch.Tensor, svh: torch.Tensor, k: int, dev: torch.device,
                 prescale: bool = False, ext=None, force_shape: int = 0):
        from exllamav3.modules.quant.exl3 import LinearEXL3

        self.in_features, self.out_features, self.k = trellis.shape[0] * 16, trellis.shape[1] * 16, k
        self.trellis = trellis.to(dev).contiguous()
        self.suh = suh.to(dev).contiguous()
        self.svh = svh.to(dev).contiguous()
        self.mul1 = torch.tensor(MUL1, dtype=torch.int32, device=dev)
        self.lin = LinearEXL3(None, self.in_features, self.out_features, suh=self.suh, svh=self.svh,
                              trellis=self.trellis, mul1=self.mul1)
        self.prescale = prescale
        self.direct = ext is not None or force_shape > 0
        if self.direct:
            if ext is None:
                import exllamav3_ext as ext
            self.ext = ext
        self.force_shape = force_shape

    def _gemm(self, x16: torch.Tensor, out: torch.Tensor | None) -> torch.Tensor:
        if not self.direct:
            if out is not None:
                raise ValueError("out= needs the direct exl3_gemm path")
            return self.lin.forward(x16, {}, out_dtype=torch.float32)
        if out is None:
            out = torch.empty(x16.shape[0], self.out_features, dtype=torch.float32, device=x16.device)
        self.ext.exl3_gemm(x16, self.trellis, out, self.suh, torch.empty_like(x16), self.svh, self.force_shape, False,
                           True, 0)
        return out

    def __call__(self, x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
        """x [m, in] (bf16 / fp32 are converted to fp16; fp16 is used as is) -> fp32 [m, out]."""
        if self.prescale:
            sc = torch.exp2(torch.ceil(torch.log2(x.abs().amax().float().clamp_min(1.0) / 16384.0)).clamp_min(0.0))
            y = self._gemm((x.float() / sc).to(torch.float16).contiguous(), out)
            return y.mul_(sc)
        if x.dtype != torch.float16:
            x = x.to(torch.float16)
        return self._gemm(x.contiguous(), out)

    def nbytes(self) -> int:
        return sum(t.numel() * t.element_size() for t in (self.trellis, self.suh, self.svh))


@dataclass
class Stage:
    wq_a: Exl3Linear
    wkv: Exl3Linear
    wq_b: Exl3Linear
    wo_b: Exl3Linear
    wo_a8: torch.Tensor           # uint8 [groups * 1024, 4096]: wo_a's e4m3fn bytes (fused path: fused.wo_a_fp8)
    wo_a_sc: torch.Tensor         # fp32 [256, 128]: its UE8M0 32 x 32 block scales
    wo_a_t: torch.Tensor | None   # bf16 [groups, 4096, 1024]: wo_a dequantized exactly, transposed per group (unfused
                                  # path and tools; None when fused, see DSparkDrafter.wo_a_bf16)
    attn_sink: torch.Tensor       # fp32 [heads]
    q_norm: torch.Tensor
    kv_norm: torch.Tensor
    attn_norm: torch.Tensor
    ffn_norm: torch.Tensor
    hc_attn: tuple[torch.Tensor, torch.Tensor, torch.Tensor]   # fn fp32 [24, 4H], scale [3], base [24]
    hc_ffn: tuple[torch.Tensor, torch.Tensor, torch.Tensor]
    gate_w: torch.Tensor          # fp32 [E, H]
    gate_bias: torch.Tensor       # fp32 [E]
    sh_w1: Exl3Linear
    sh_w3: Exl3Linear
    sh_w2: Exl3Linear
    bank: Exl3Bank
    groups: list
    perm: torch.Tensor            # int64 [E] logical -> physical slot


def dequant_wo_a(weight: torch.Tensor, scale: torch.Tensor, groups: int) -> torch.Tensor:
    """FP8 e4m3 [G*R, C] x UE8M0 block scales [G*R/bs, C/bs] -> bf16 [G, C, R] (exact: e4m3 x 2^k is a bf16 value)."""
    rows, cols = weight.shape
    br, bc = rows // scale.shape[0], cols // scale.shape[1]
    s = scale.float().repeat_interleave(br, 0).repeat_interleave(bc, 1)
    w = (weight.float() * s).to(torch.bfloat16)
    if not torch.equal(w.float(), weight.float() * s):
        raise RuntimeError("wo_a dequantization is not exact in bf16")
    return w.view(groups, rows // groups, cols).transpose(1, 2).contiguous()


class DSparkDrafter:
    """All weights resident on one GPU; per-request windows in ``slots`` fixed slots (slot ``slots`` = scratch)."""

    HEAD_SHAPE = 3      # exl3_gemm shape (16, 32, 256): what the autotuner picks for the head at 1-16 rows (689 us)
    HEAD_CHUNK = 16     # head rows per GEMM call: every call stays on the 16-row shape (n_out bit-identity)

    def __init__(self, ck: Checkpoint, dev: torch.device, slots: int = 16, max_model_len: int = 327680,
                 max_rows: int = 32, kv_format: str = "v4", fused: bool = True, main_proj_fp32acc: bool = True,
                 markov_cache_rows: int = 8192):
        if not ck.verified:
            raise RuntimeError("checkpoint not identity-verified")
        t0 = time.time()
        self.dev = dev
        self.s: DrafterShape = ck.shape
        s = self.s
        self.kv_format = kv_format          # "v4": the SM12x fp8_ds_mla record values; "bf16": no quantization
        self.fused = fused                  # Triton kernels of fused.py (False: the plain torch ops of ops.py)
        torch.backends.cuda.matmul.allow_tf32 = False
        torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
        self.max_model_len = max_model_len
        self.max_pos = max_model_len + 2048          # RoPE rows, like the engine's cap (patch 0033a)
        self.max_rows = max_rows
        self.eps = s.rms_eps
        self.scale = s.head_dim ** -0.5
        k_dense = ck.drafter_k
        plan = Exl3MoEPlan.from_dict(ck.plan_json)

        def lin(name: str, want_k: int | None, prescale: bool = False, ext=None, force_shape: int = 0) -> Exl3Linear:
            tr, suh, svh, k = ck.exl3(name)
            if want_k is not None and k != want_k:
                raise RuntimeError(f"{name}: trellis K{k}, config says K{want_k}")
            return Exl3Linear(tr, suh, svh, k, dev, prescale, ext, force_shape)

        def g(name: str, dtype: torch.dtype | None = None) -> torch.Tensor:
            t = ck.get(name)
            return t.to(dev, dtype=dtype or t.dtype).contiguous()

        self.embed = g("embed.weight")                                     # bf16 [V, H]
        self.head = lin("head", ck.head_k, force_shape=self.HEAD_SHAPE)
        # main_proj with fp32 accumulation: its 15360-wide input (three nearly equal target residual streams) cancels
        # heavily, and exllamav3's sm_86 fp16 accumulators put main_x 0.37 % from fp32 (the engine: 0.11 %; D3b)
        self.main_proj_fp32acc = main_proj_fp32acc
        self.main_proj = lin("mtp.0.main_proj", k_dense["main_proj"], prescale=True,
                             ext=facc_ext() if main_proj_fp32acc else None)
        self.main_norm = g("mtp.0.main_norm.weight")
        last = f"mtp.{s.n_stages - 1}."
        self.norm = g(last + "norm.weight")
        self.conf_w = g(last + "confidence_head.proj.weight", torch.float32).view(-1)
        tokens = None
        tok_path = os.path.join(ck.root, "markov_top_tokens.json")
        if os.path.exists(tok_path):
            with open(tok_path) as f:
                tokens = json.load(f)
        # cached Markov bias rows: the serve caches the first 4,096 ids of markov_top_tokens.json; the worker caches all
        # 8,192 (same kernel, same bits; 80 % instead of 74 % of the live captures' Markov inputs hit the cache, each
        # miss costs a 66 MB GEMV, ~85 us on the 3090). +1 GB.
        self.markov = MarkovBias(g(last + "markov_head.embed.weight"), g(last + "markov_head.head.weight"), tokens,
                                 rows=markov_cache_rows)

        self.stages: list[Stage] = []
        for i in range(s.n_stages):
            p = f"mtp.{i}."
            L = s.n_target_layers + i
            lp = plan.drafter_layers[L]
            if lp.num_experts != s.n_experts:
                raise RuntimeError(f"plan layer {L}: {lp.num_experts} experts, config {s.n_experts}")
            bank = Exl3Bank(lp, s.hidden, s.inter, 1, 0)
            bank.allocate(dev)
            for e in range(s.n_experts):
                for proj, wn in ((PROJ_GATE, "w1"), (PROJ_UP, "w3"), (PROJ_DOWN, "w2")):
                    tr, suh, svh, _ = ck.exl3(f"{p}ffn.experts.{e}.{wn}")
                    bank.load_piece("trellis", proj, e, tr)
                    bank.load_piece("suh", proj, e, suh)
                    bank.load_piece("svh", proj, e, svh)
            groups, _ = bank.tables()
            st = Stage(
                wq_a=lin(p + "attn.wq_a", k_dense["attn.wq_a"]), wkv=lin(p + "attn.wkv", k_dense["attn.wq_a"]),
                wq_b=lin(p + "attn.wq_b", k_dense["attn.wq_b"]), wo_b=lin(p + "attn.wo_b", k_dense["attn.wo_b"]),
                wo_a8=g(p + "attn.wo_a.weight").view(torch.uint8), wo_a_sc=g(p + "attn.wo_a.scale", torch.float32),
                wo_a_t=None if fused else dequant_wo_a(ck.get(p + "attn.wo_a.weight").to(dev),
                                                       ck.get(p + "attn.wo_a.scale").to(dev), s.o_groups),
                attn_sink=g(p + "attn.attn_sink", torch.float32), q_norm=g(p + "attn.q_norm.weight"),
                kv_norm=g(p + "attn.kv_norm.weight"), attn_norm=g(p + "attn_norm.weight"),
                ffn_norm=g(p + "ffn_norm.weight"),
                hc_attn=(g(p + "hc_attn_fn"), g(p + "hc_attn_scale"), g(p + "hc_attn_base")),
                hc_ffn=(g(p + "hc_ffn_fn"), g(p + "hc_ffn_scale"), g(p + "hc_ffn_base")),
                gate_w=g(p + "ffn.gate.weight", torch.float32), gate_bias=g(p + "ffn.gate.bias", torch.float32),
                sh_w1=lin(p + "ffn.shared_experts.w1", k_dense["ffn.shared_experts.w1"]),
                sh_w3=lin(p + "ffn.shared_experts.w3", k_dense["ffn.shared_experts.w1"]),
                sh_w2=lin(p + "ffn.shared_experts.w2", k_dense["ffn.shared_experts.w2"]),
                bank=bank, groups=groups, perm=torch.tensor(lp.perm, dtype=torch.int64, device=dev),
            )
            self.stages.append(st)
        from .exl3.hostcopy import release_stage
        release_stage()

        self.coop = K.get_coop_scratch(dev, max_rows, s.top_k, s.hidden, s.inter)
        K.prime(dev, self.stages[0].groups, self.coop, None, s.hidden, s.top_k, s.swiglu_limit, False, True)
        self.cos_sin = rope_cos_sin(self.max_pos, s.rope_dim, s.rope_theta).to(dev)
        # query positions are clamped to max_model_len - 1 (the engine's clamp); a device scalar so CUDA graphs read
        # the current value (set_query_clamp)
        self.q_clamp = torch.tensor(max_model_len - 1, dtype=torch.int64, device=dev)

        # per-request windows: slot `slots` is a scratch slot that padded rows write to
        self.n_slots = slots
        self.ring_kv = torch.zeros(slots + 1, s.n_stages, s.window, s.head_dim, dtype=torch.bfloat16, device=dev)
        self.ring_pos = torch.full((slots + 1, s.window), -1, dtype=torch.int64, device=dev)
        self.arange_block = torch.arange(s.block + 1, dtype=torch.int64, device=dev)
        self.graphs: dict[tuple[int, int, int], dict] = {}
        self.margin_graphs: dict[tuple[int, int, int], dict] = {}
        self.reply_graphs: dict[tuple[int, int, int, bool], dict] = {}
        self.insert_graphs: dict[int, dict] = {}
        self.pool = None
        torch.cuda.synchronize(dev)
        self.load_seconds = time.time() - t0

    def wo_a_bf16(self, i: int) -> torch.Tensor:
        """Stage i's wo_a as bf16 [groups, 4096, 1024] (exact dequantization; for tools on the fused model)."""
        st = self.stages[i]
        if st.wo_a_t is not None:
            return st.wo_a_t
        return dequant_wo_a(st.wo_a8.view(torch.float8_e4m3fn), st.wo_a_sc, self.s.o_groups)

    # ------------------------------------------------------------------------------------------ op selection
    def _norm(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
        return FU.rms_norm(x, w, self.eps) if self.fused else rms_norm(x.to(torch.bfloat16), w, self.eps)

    def _rope(self, x: torch.Tensor, pos: torch.Tensor, inverse: bool = False) -> torch.Tensor:
        return FU.rope(x, pos, self.cos_sin, inverse) if self.fused else apply_rope(x, pos, self.cos_sin, inverse)

    def _kv(self, kv: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
        """rope + (V4 record) of normed KV rows."""
        if self.kv_format != "v4":
            return self._rope(kv, pos)
        if self.fused:
            return FU.kv_rope_qdq(kv, pos, self.cos_sin)
        return kv_record_qdq(apply_rope(kv, pos, self.cos_sin))

    def _mixes(self, stream, fn, scale, base):
        f = FU.hc_mixes if self.fused else hc_mixes
        return f(stream, fn, scale, base, self.eps, self.s.hc_eps, self.s.hc_iters)

    def _collapse_norm(self, stream, pre, w):
        if self.fused:
            return FU.collapse_norm(stream, pre, w, self.eps)
        return rms_norm(hc_collapse(stream, pre), w, self.eps)

    def _post(self, y, stream, post, comb):
        return (FU.hc_post if self.fused else hc_post)(y, stream, post, comb)

    def _collapse(self, stream, pre):
        return FU.collapse(stream, pre) if self.fused else hc_collapse(stream, pre)

    # ------------------------------------------------------------------------------------------ context side
    def main_x(self, taps: torch.Tensor) -> torch.Tensor:
        """[n, 3H] bf16 taps of target layers 37/38/39 -> main_x [n, H] bf16."""
        if self.fused and isinstance(self.main_proj, Exl3Linear):
            x16, sc = FU.prescale_fp16(taps)
            return FU.scaled_rms_norm(self.main_proj._gemm(x16, None), sc, self.main_norm, self.eps)
        return self._norm(self.main_proj(taps), self.main_norm)

    def context_kv(self, mx: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
        """Stage context KV [n_stages, n, D] (rope at pos, record-quantized) of main_x rows."""
        out = []
        for st in self.stages:
            out.append(self._kv(self._norm(st.wkv(mx), st.kv_norm), pos))
        return torch.stack(out)

    def insert(self, slot_rows: torch.Tensor, pos_rows: torch.Tensor, mx: torch.Tensor) -> None:
        """Write the context KV of rows (slot, position, main_x) into the windows. Callers pass at most the last
        `window` positions of a request per call (older ones would be overwritten anyway; duplicates are refused
        upstream). Padded rows use slot == n_slots (scratch)."""
        if self.fused and self.kv_format == "v4":
            mx16 = mx.to(torch.float16)
            for i, st in enumerate(self.stages):
                FU.kv_store(st.wkv(mx16), st.kv_norm, self.eps, pos_rows, self.cos_sin, self.ring_kv[:, i], slot_rows,
                            self.ring_pos if i == 0 else None)
            return
        kv = self.context_kv(mx, pos_rows)                          # [S, n, D]
        idx = pos_rows % self.s.window
        self.ring_kv[slot_rows, :, idx] = kv.transpose(0, 1)
        self.ring_pos[slot_rows, idx] = pos_rows

    def reset_slots(self, slots: torch.Tensor) -> None:
        self.ring_pos[slots] = -1
        self.ring_kv[slots] = 0

    # ------------------------------------------------------------------------------------------ draft step
    def _attention(self, st: Stage, i: int, x: torch.Tensor, pos: torch.Tensor, win: Window, B: int, k: int
                   ) -> torch.Tensor:
        s = self.s
        T = B * k
        x16 = x.to(torch.float16)
        qr = self._norm(st.wq_a(x16), st.q_norm)
        kvb = self._kv(self._norm(st.wkv(x16), st.kv_norm), pos)
        q = self._rope(st.wq_b(qr).to(torch.bfloat16).view(T, s.n_heads, s.head_dim), pos)
        qb, kb = q.view(B, k, s.n_heads, s.head_dim), kvb.view(B, k, s.head_dim)
        if win.kv is None:
            o = FU.block_attention_ring(qb, self.ring_kv[:, i], self.ring_pos, win.slots, win.lo, win.last, kb,
                                        st.attn_sink, self.scale)
        else:
            o = block_attention(qb, win.kv[:, i], win.valid, kb, st.attn_sink, self.scale)
        o = self._rope(o.view(T, s.n_heads, s.head_dim), pos, inverse=True)
        if st.wo_a_t is None:
            z = FU.wo_a_fp8(o.view(T, -1), st.wo_a8, st.wo_a_sc, s.o_groups)  # [T, G*1024] bf16
        else:
            og = o.view(T, s.o_groups, -1).transpose(0, 1)                   # [G, T, 4096]
            z = torch.bmm(og, st.wo_a_t).transpose(0, 1).reshape(T, -1)      # [T, G*1024] bf16
        return st.wo_b(z).to(torch.bfloat16)

    def _moe(self, st: Stage, x: torch.Tensor, trace: dict | None = None, tag: str = "") -> torch.Tensor:
        s = self.s
        ids, w = route(x, st.gate_w, st.gate_bias, s.top_k, s.route_scale)
        if trace is not None:
            trace[tag + ".experts"] = ids
            trace[tag + ".route_w"] = w
        x16 = x.to(torch.float16).contiguous()
        sel = st.perm[ids].contiguous()
        routed = K.run_decode_coop(x16, sel, w.to(torch.float16).contiguous(), st.groups, self.coop,
                                   s.swiglu_limit, False, True)
        g = st.sh_w1(x16)
        u = st.sh_w3(x16)
        lim = s.swiglu_limit
        a = torch.nn.functional.silu(g.clamp(max=lim)) * u.clamp(min=-lim, max=lim)
        sh = st.sh_w2(a.to(torch.float16))
        return (routed + sh).to(torch.bfloat16)

    def forward_block(self, ids: torch.Tensor, pos: torch.Tensor, win: Window, trace: dict | None = None
                      ) -> torch.Tensor:
        """The three stages over a query block. ids/pos [B, k]; win: the context window (block_inputs).
        Returns head_hidden [B*k, H] bf16 (pre-norm). ``trace`` (eager only) collects per-stage tensors."""
        s = self.s
        B, k = ids.shape
        T = B * k
        pos_f = pos.reshape(T)
        x = self.embed[ids.reshape(T)]
        stream = x.unsqueeze(1).expand(T, s.hc_mult, s.hidden).contiguous()
        pre_in = None
        for i, st in enumerate(self.stages):
            pre, post, comb = self._mixes(stream, *st.hc_attn)
            xin = self._collapse_norm(stream, pre_in, st.attn_norm)
            a = self._attention(st, i, xin, pos_f, win, B, k)
            stream = self._post(a, stream, post, comb)
            fpre, fpost, fcomb = self._mixes(stream, *st.hc_ffn)
            xin = self._collapse_norm(stream, pre, st.ffn_norm)
            y = self._moe(st, xin, trace, f"stage{i}")
            stream = self._post(y, stream, fpost, fcomb)
            pre_in = fpre
            if trace is not None:
                trace[f"stage{i}.attn_out"] = a
                trace[f"stage{i}.ffn_in"] = xin
                trace[f"stage{i}.ffn_out"] = y
                trace[f"stage{i}.stream"] = stream
        return self._collapse(stream, pre_in)

    def heads(self, hh: torch.Tensor, anchors: torch.Tensor, B: int, k: int, trace: dict | None = None,
              n_out: int | None = None) -> tuple[torch.Tensor, torch.Tensor]:
        """Greedy sequential Markov drafting + confidence over head_hidden [B*k, H], for the first ``n_out`` (default
        k) positions of every request; later positions return draft 0 and confidence 0. Positions < n_out are
        bit-identical for every n_out (CONTRACT 1.3): head rows go through the GEMM in chunks of <= HEAD_CHUNK rows on
        one fixed kernel shape, Markov and argmax are per row, and the confidence product keeps its [B, k] shape
        (zero Markov embeddings past n_out)."""
        s = self.s
        n = k if n_out is None else n_out
        if not 1 <= n <= k:
            raise ValueError(f"n_out {n} outside 1..{k}")
        rows = self._norm(hh, self.norm).view(B, k, s.hidden)[:, :n].reshape(B * n, s.hidden).to(torch.float16)
        l32 = torch.empty(B * n, s.vocab, dtype=torch.float32, device=self.dev)
        for c in range(0, B * n, self.HEAD_CHUNK):
            self.head(rows[c:c + self.HEAD_CHUNK], out=l32[c:c + self.HEAD_CHUNK])
        logits = l32.to(torch.bfloat16).view(B, n, s.vocab)
        out = torch.zeros(B, k, dtype=torch.int64, device=self.dev)
        if self.fused and trace is None:
            for i in range(n):
                FU.markov_argmax_into(logits[:, i], anchors if i == 0 else out[:, i - 1], self.markov, out, i)
            emb = self.markov.embed(torch.cat((anchors.unsqueeze(1), out[:, :n - 1]), 1))      # [B, n, R]
        else:
            prev = anchors
            embs = []
            for i in range(n):
                e = self.markov.embed(prev)
                li = logits[:, i] + self.markov(prev, e)
                d = li.argmax(dim=-1)
                if trace is not None:
                    trace[f"markov{i}.logits"] = li
                out[:, i] = d
                embs.append(e)
                prev = d
            emb = torch.stack(embs, 1)
        if n < k:
            emb = torch.cat((emb, emb.new_zeros(B, k - n, emb.shape[-1])), 1)
        feat = torch.cat((hh.view(B, k, s.hidden), emb), dim=-1).float()
        conf = torch.sigmoid(feat @ self.conf_w)
        if n < k:
            conf[:, n:] = 0.0
        if trace is not None:
            trace["base_logits"] = logits
        return out, conf

    def block_inputs(self, slots: torch.Tensor, anchors: torch.Tensor, last_pos: torch.Tensor, k: int,
                     ctx_lo: torch.Tensor | None = None, gather: bool | None = None
                     ) -> tuple[torch.Tensor, torch.Tensor, Window]:
        """Query ids/positions and the window: ring rows whose position is in [ctx_lo, last_pos]. ctx_lo defaults to
        max(last_pos - W + 1, 0); the server passes max(lo, anchor_pos - W) with the slot's contiguous lower bound.
        ``gather`` (default: not fused) also materializes the window rows and their mask."""
        s = self.s
        ar = self.arange_block[:k]
        ids = torch.where(ar.unsqueeze(0) == 0, anchors.unsqueeze(1), torch.full_like(ar, s.noise_token).unsqueeze(0))
        pos = torch.minimum(last_pos.unsqueeze(1) + 1 + ar.unsqueeze(0), self.q_clamp)
        if ctx_lo is None:
            ctx_lo = (last_pos - (s.window - 1)).clamp_min(0)
        win = Window(slots, ctx_lo, last_pos)
        if gather is None:
            gather = not self.fused
        if gather:
            cpos = self.ring_pos[slots]
            win.valid = (cpos >= ctx_lo.unsqueeze(1)) & (cpos <= last_pos.unsqueeze(1))
            win.kv = self.ring_kv[slots]
        return ids, pos, win

    def draft(self, slots: torch.Tensor, anchors: torch.Tensor, last_pos: torch.Tensor, k: int,
              trace: dict | None = None, ctx_lo: torch.Tensor | None = None, n_out: int | None = None
              ) -> tuple[torch.Tensor, torch.Tensor]:
        """Eager draft step. slots/anchors/last_pos (/ctx_lo) int64 [B] on the device. The body always runs the
        whole k-token block; ``n_out`` limits the head / Markov / confidence to the first n_out positions."""
        B = slots.shape[0]
        ids, pos, win = self.block_inputs(slots, anchors, last_pos, k, ctx_lo)
        hh = self.forward_block(ids, pos, win, trace)
        if trace is not None:
            trace["head_hidden"] = hh
        return self.heads(hh, anchors, B, k, trace, n_out)

    def set_query_clamp(self, max_model_len: int) -> None:
        if not 0 < max_model_len <= self.max_model_len:
            raise ValueError(f"max_model_len {max_model_len} outside 1..{self.max_model_len}")
        self.q_clamp.fill_(max_model_len - 1)

    def draft_with_margin(self, slots: torch.Tensor, anchors: torch.Tensor, last_pos: torch.Tensor, k: int,
                          ctx_lo: torch.Tensor | None = None, n_out: int | None = None
                          ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Eager draft that also returns top1 - top2 of (base + Markov) logits per position (CONTRACT WANT_MARGIN;
        0 past n_out). Same drafts as draft(): the argmax runs on the same bf16 logits."""
        tr: dict = {}
        ids, conf = self.draft(slots, anchors, last_pos, k, trace=tr, ctx_lo=ctx_lo, n_out=n_out)
        n = k if n_out is None else n_out
        mg = torch.zeros(slots.shape[0], k, dtype=torch.float32, device=self.dev)
        for i in range(n):
            v = tr[f"markov{i}.logits"].float().topk(2, dim=-1).values
            mg[:, i] = v[:, 0] - v[:, 1]
        return ids, conf, mg

    # ------------------------------------------------------------------------------------------ CUDA graphs
    def capture(self, batch_sizes=(1, 2, 3, 4), ks=(4, 5), insert_rows=(8, 16, 32, 64, 128), n_outs=(4, 5),
                margins: bool = True, pack_reply: bool = False) -> float:
        """One graph per (B, k, n_out) (n_out in n_outs and <= k; n_out = k always) for the draft step, the same with
        top1 - top2 margins (CONTRACT WANT_MARGIN; shadow mode asks for them) when ``margins``, and one per row
        bucket for the insert. ``pack_reply`` adds optional float64 reply twins;
        default False, so normal capture count and outputs remain unchanged. Returns seconds."""
        t0 = time.time()
        dev = self.dev
        if self.pool is None:
            self.pool = torch.cuda.graph_pool_handle()
        bmax = max(batch_sizes)
        self.g_slots = torch.full((bmax,), self.n_slots, dtype=torch.int64, device=dev)
        self.g_anchor = torch.zeros(bmax, dtype=torch.int64, device=dev)
        self.g_last = torch.zeros(bmax, dtype=torch.int64, device=dev)
        self.g_lo = torch.zeros(bmax, dtype=torch.int64, device=dev)
        rmax = max(insert_rows)
        self.gi_slot = torch.full((rmax,), self.n_slots, dtype=torch.int64, device=dev)
        self.gi_pos = torch.zeros(rmax, dtype=torch.int64, device=dev)
        self.gi_taps = torch.zeros(rmax, self.s.hidden * len(self.s.taps), dtype=torch.bfloat16, device=dev)
        s1 = torch.cuda.Stream(dev)
        for B in batch_sizes:
            for k in ks:
                for n in sorted({x for x in n_outs if x <= k} | {k}):
                    args = (self.g_slots[:B], self.g_anchor[:B], self.g_last[:B], k)
                    lo = self.g_lo[:B]
                    s1.wait_stream(torch.cuda.current_stream(dev))
                    with torch.cuda.stream(s1):
                        for _ in range(2):
                            self.draft(*args, ctx_lo=lo, n_out=n)   # warm-up: kernel attributes, buffers, Triton JIT
                    torch.cuda.current_stream(dev).wait_stream(s1)
                    torch.cuda.synchronize(dev)
                    gr = torch.cuda.CUDAGraph()
                    # A process-wide default capture stream can belong to another device (e.g. TP2 rank 0).
                    # Keep the pool on this model's stream, regardless of which model captured first.
                    with torch.cuda.graph(gr, pool=self.pool, stream=s1):
                        ids, conf = self.draft(*args, ctx_lo=lo, n_out=n)
                    self.graphs[(B, k, n)] = {"graph": gr, "ids": ids, "conf": conf}
                    if margins:
                        s1.wait_stream(torch.cuda.current_stream(dev))
                        with torch.cuda.stream(s1):
                            for _ in range(2):
                                self.draft_with_margin(*args, ctx_lo=lo, n_out=n)
                        torch.cuda.current_stream(dev).wait_stream(s1)
                        torch.cuda.synchronize(dev)
                        gr = torch.cuda.CUDAGraph()
                        with torch.cuda.graph(gr, pool=self.pool, stream=s1):
                            ids, conf, mg = self.draft_with_margin(*args, ctx_lo=lo, n_out=n)
                        self.margin_graphs[(B, k, n)] = {"graph": gr, "ids": ids, "conf": conf, "mg": mg}
                    if pack_reply:
                        for margin in ((False, True) if margins else (False,)):
                            gr = torch.cuda.CUDAGraph()
                            with torch.cuda.graph(gr, pool=self.pool, stream=s1):
                                values = (self.draft_with_margin(*args, ctx_lo=lo, n_out=n) if margin
                                          else self.draft(*args, ctx_lo=lo, n_out=n))
                                packed = torch.cat([v.to(torch.float64) for v in values], 1)
                            self.reply_graphs[(B, k, n, margin)] = {"graph": gr, "packed": packed}
        for n in insert_rows:
            args = (self.gi_slot[:n], self.gi_pos[:n])
            s1.wait_stream(torch.cuda.current_stream(dev))
            with torch.cuda.stream(s1):
                for _ in range(2):
                    self.insert(*args, self.main_x(self.gi_taps[:n]))
            torch.cuda.current_stream(dev).wait_stream(s1)
            torch.cuda.synchronize(dev)
            gr = torch.cuda.CUDAGraph()
            with torch.cuda.graph(gr, pool=self.pool, stream=s1):
                self.insert(*args, self.main_x(self.gi_taps[:n]))
            self.insert_graphs[n] = {"graph": gr}
        # the warm-ups wrote into the scratch slot only; keep it empty
        self.reset_slots(torch.tensor([self.n_slots], device=dev))
        torch.cuda.synchronize(dev)
        return time.time() - t0

    def has_graph(self, B: int, k: int, n_out: int | None = None, margin: bool = False) -> bool:
        return (B, k, k if n_out is None else n_out) in (self.margin_graphs if margin else self.graphs)

    def draft_graphed(self, slots: torch.Tensor, anchors: torch.Tensor, last_pos: torch.Tensor, k: int,
                      ctx_lo: torch.Tensor | None = None, n_out: int | None = None, margin: bool = False,
                      packed_reply: bool = False
                      ) -> tuple[torch.Tensor, ...]:
        """Replay the (B, k, n_out) graph; returns views of the graph's output buffers (ids, conf[, margin]; copy
        before the next replay). ``margin``: the draft_with_margin graph (same ids and confidences).
        ``packed_reply`` returns a one-tuple holding the captured float64 wire staging tensor;
        requires capture(pack_reply=True), including a margin twin when margin=True."""
        B = slots.shape[0]
        n = k if n_out is None else n_out
        g = (self.reply_graphs.get((B, k, n, margin)) if packed_reply
             else (self.margin_graphs if margin else self.graphs).get((B, k, n)))
        if g is None:
            raise KeyError(f"no CUDA graph for batch {B}, depth {k}, n_out {n}, margin {margin}, packed {packed_reply}")
        self.g_slots[:B].copy_(slots, non_blocking=True)
        self.g_anchor[:B].copy_(anchors, non_blocking=True)
        self.g_last[:B].copy_(last_pos, non_blocking=True)
        if ctx_lo is None:
            ctx_lo = (last_pos - (self.s.window - 1)).clamp_min(0)
        self.g_lo[:B].copy_(ctx_lo, non_blocking=True)
        g["graph"].replay()
        if packed_reply:
            return (g["packed"],)
        return (g["ids"], g["conf"], g["mg"]) if margin else (g["ids"], g["conf"])

    def insert_taps_graphed(self, slot_rows: torch.Tensor, pos_rows: torch.Tensor, taps: torch.Tensor) -> None:
        """main_x + insert for n rows through the smallest captured bucket >= n (padding rows hit the scratch slot)."""
        n = slot_rows.shape[0]
        sizes = sorted(self.insert_graphs)
        b = next((x for x in sizes if x >= n), None)
        if b is None:
            raise ValueError(f"{n} insert rows exceed the largest graph bucket {sizes[-1]}")
        self.gi_slot[:n].copy_(slot_rows, non_blocking=True)
        self.gi_slot[n:b].fill_(self.n_slots)
        self.gi_pos[:n].copy_(pos_rows, non_blocking=True)
        self.gi_pos[n:b].zero_()
        self.gi_taps[:n].copy_(taps, non_blocking=True)
        self.gi_taps[n:b].zero_()           # padding rows: no stale taps in main_proj's batch-wide prescale
        self.insert_graphs[b]["graph"].replay()

    # ------------------------------------------------------------------------------------------ accounting
    def memory_table(self) -> dict:
        def nb(t: torch.Tensor) -> int:
            return t.numel() * t.element_size()

        st_dense = sum(sum(l.nbytes() for l in (st.wq_a, st.wkv, st.wq_b, st.wo_b, st.sh_w1, st.sh_w3, st.sh_w2))
                       for st in self.stages)
        st_woa = sum(nb(st.wo_a8) + nb(st.wo_a_sc) + (nb(st.wo_a_t) if st.wo_a_t is not None else 0)
                     for st in self.stages)
        st_exp = sum(sum(nb(t) for t in st.bank.t.values()) for st in self.stages)
        st_small = sum(nb(t) for st in self.stages for t in (st.attn_sink, st.q_norm, st.kv_norm, st.attn_norm,
                                                               st.ffn_norm, st.gate_w, st.gate_bias, *st.hc_attn,
                                                               *st.hc_ffn))
        return {
            "embed_bf16": nb(self.embed),
            "head_exl3_k6": self.head.nbytes(),
            "main_proj_exl3_k5": self.main_proj.nbytes(),
            "stages_dense_exl3_k5": st_dense,
            "stages_wo_a_fp8": st_woa,
            "stages_experts_exl3_k2": st_exp,
            "stages_norms_hc_gate_fp32": st_small,
            "markov_w1_w2_bf16": nb(self.markov.w1) + nb(self.markov.w2),
            "markov_cache_rows_bf16": nb(self.markov.cache) + nb(self.markov.slot_of_token),
            "confidence_and_norms": nb(self.conf_w) + nb(self.norm) + nb(self.main_norm),
            "rope_cos_sin_fp32": nb(self.cos_sin),
            "windows_per_slot": nb(self.ring_kv[0]) + nb(self.ring_pos[0]),
            "windows_total": nb(self.ring_kv) + nb(self.ring_pos),
            "coop_scratch": sum(nb(t) for t in (self.coop.had_g, self.coop.had_u, self.coop.gu_g, self.coop.gu_u,
                                                self.coop.act_out, self.coop.d_out, self.coop.ctr, *self.coop.out)),
            "torch_allocated": torch.cuda.memory_allocated(self.dev),
            "torch_reserved": torch.cuda.memory_reserved(self.dev),
            "torch_max_allocated": torch.cuda.max_memory_allocated(self.dev),
        }