"""Graph-captured greedy decoding: AR baseline and self-speculative decoding.

All compute runs in captured CUDA/HIP graphs with static buffers:
  ar             token -> all layers -> LM head argmax                       (1 row)
  seg[a, b, m]   m hidden rows (buffer) -> layers [a, b) at positions pos[:m] (m rows)
  head[k]        one hidden row -> head k -> (confidence, token id)
  fill[a, b]     CALM K/V copy for layers [a, b) from one hidden row          (regime B)
  verify[c, n]   n hidden rows at layer c -> layers [c, L) -> argmax           (n rows)
Positions are read from a device buffer, so graphs replay anywhere. Attention
uses the row-invariant Triton kernel and GEMMs use rocBLAS, so AR and every
batched path compute each row bitwise identically (greedy parity).

Regime A (default, kv_fill="lazy"): before a draft row climbs to a deeper head,
the rows of the block that are shallower are caught up exactly in one batched
segment pass; head scores depend only on the exact prefix and verification runs
only [c_max, L). Regime B (kv_fill="copy"): K/V of skipped layers are copied from
the exit hidden state; verification runs [c_min, L).
Prefill (shared by all methods) runs eagerly.
"""

from __future__ import annotations

import math
import time

import numpy as np

import torch
import torch.nn.functional as F

from .decode import GenStats, StaticPolicy
from .heads import ExitHead
try:
    from .kernels import row_attention
except ImportError:  # no Triton (CPU tests)
    row_attention = None
from .layered import LayeredDecoder, _rotate_half


def _lam(controller) -> float:
    """Scalar lambda of an online controller (NoSlowdown wraps it in .aci)."""
    c = getattr(controller, "aci", controller)
    return float(np.min(getattr(c, "lam", float("nan"))))


class DraftModel:
    """Classical speculative decoding: a small model of the same tokenizer drafts for the target
    (P1 baseline). Own FastDecoder (own KV, AR graph). Per cycle: catch up the committed tokens it
    has not processed, then draft `gamma` tokens greedily, or stop early when its max probability
    falls below `conf` (adaptive). The drafts are verified exactly by the target like copies."""

    def __init__(self, fast: "FastDecoder", gamma: int = 4, conf: float | None = None):
        self.fast, self.gamma, self.conf = fast, gamma, conf

    def start(self, prompt: list[int]) -> None:
        self.fast._check_capacity(prompt, 0)
        self.fast._prefill(prompt)
        self.d_len = len(prompt)  # positions whose draft KV is valid

    def _step(self, tok: int, pos: int):
        f = self.fast
        f._set_pos([pos])
        f.ids.fill_(tok)
        a, c = f._g_ar_conf().tolist()
        return int(a), float(c)

    def _catchup(self, seq: list[int], start: int, end: int) -> None:
        """Committed tokens [start, end) the draft has not processed (e.g. after prompt-lookup
        cycles): processed in blocks of max_rows through the draft's layer-0 verify graphs (KV
        only; a draft-side error could only lower acceptance, never exactness)."""
        f = self.fast
        W = f.max_rows
        ids = torch.tensor(seq[start:end], device=f.dec.device)
        for s0 in range(0, end - start, W):
            m = min(W, end - start - s0)
            f._set_pos(list(range(start + s0, start + s0 + m)))
            f.buf[:m].copy_(f.dec.embed(ids[s0:s0 + m]))
            f._g_verify(0, m)

    def propose(self, seq: list[int], budget: int) -> list[int]:
        pos0 = len(seq) - 1
        if pos0 - self.d_len > 1:
            self._catchup(seq, self.d_len, pos0)
        else:
            for p in range(self.d_len, pos0):  # committed tokens the draft has not processed yet
                self._step(seq[p], p)
        x, drafts, self.n_inputs = seq[pos0], [], 0
        for j in range(min(self.gamma, budget)):
            t, c = self._step(x, pos0 + j)
            self.n_inputs += 1
            if self.conf is not None and j > 0 and c < self.conf:
                break
            drafts.append(t)
            x = t
        self.pos0 = pos0
        return drafts

    def commit(self, accepted: int) -> None:
        # inputs at pos0 + i (i >= 1) were d_i: valid iff i <= accepted
        self.d_len = self.pos0 + 1 + min(accepted, self.n_inputs - 1)


class FastDecoder:
    def __init__(self, dec: LayeredDecoder, heads: list[ExitHead], head_layers: list[int],
                 max_rows: int = 16, use_graphs: bool | None = None, shared_pool: bool = False,
                 prefill: str = "engine", prefill_chunk: int = 2048):
        assert prefill in ("engine", "full"), prefill
        self.prefill_mode, self.prefill_chunk = prefill, prefill_chunk
        self.dec = dec
        self.heads = heads
        self.head_layers = head_layers
        self.max_rows = max_rows
        dev, d = dec.device, dec.spec.hidden_size
        # Row metadata in ONE device buffer [pos | slot | cl | anc], filled by one copy from a
        # host staging buffer (pinned on GPU): audit T5, piste 1 (3-4 pageable copies before).
        W = max_rows
        self._meta_d = torch.zeros(3 * W + W * W, dtype=torch.int64, device=dev)
        self._meta_h = torch.zeros(3 * W + W * W, dtype=torch.int64)
        if dev.type == "cuda":
            self._meta_h = self._meta_h.pin_memory()
        self.pos = self._meta_d[0:W]
        # Tree drafts: slot = KV write index of each row; keys t < cl are read at index t,
        # keys t >= cl at anc[row, t - cl] (the row's path). Chains: slot = pos, cl = pos + 1.
        self.slot = self._meta_d[W:2 * W]
        self.cl = self._meta_d[2 * W:3 * W]
        self.anc = self._meta_d[3 * W:].view(W, W)
        self.cl.fill_(1)
        self.idx = torch.zeros(max_rows, dtype=torch.int64, device=dev)
        self.ids = torch.zeros(1, dtype=torch.long, device=dev)
        self.buf = torch.zeros(max_rows, d, device=dev, dtype=dec.dtype)    # segment rows
        self.hrows = torch.zeros(max_rows, d, device=dev, dtype=dec.dtype)  # block rows
        self.hb = torch.zeros(1, d, device=dev, dtype=dec.dtype)            # head input
        # sampling (T > 0): Gumbel noise filled eagerly before each head replay (no RNG inside
        # graphs), per-head renormalised V_k distribution of the last sample, saved per draft row
        vmax = max([h.n_vocab for h in heads] + [1])
        self.gnoise = torch.zeros(vmax, device=dev, dtype=torch.float32)
        self.qout = torch.zeros(vmax, device=dev, dtype=torch.float32)
        self.qsave = torch.zeros(max_rows, vmax, device=dev, dtype=torch.float32)
        self.temperature = 1.0
        self._sampling, self._gen, self._rng = False, None, None
        self.graphs: dict = {}
        self._trace = None  # debug: per-layer tensors (eager mode only)
        # ROCm 7.0 (CLR graph "packet capture"): a graph replay is not reliably ordered
        # with later work on the same stream; this faulted on MI210 (jobs 893662-893901).
        # Synchronising after every replay removes it (job 893903) and is the fastest
        # fix measured (96 tok/s AR vs 72 with DEBUG_CLR_GRAPH_PACKET_CAPTURE=0).
        self.sync_each = bool(torch.version.hip)
        self.use_graphs = dev.type == "cuda" if use_graphs is None else use_graphs
        # One memory pool for all graphs: safe because graphs never run concurrently and
        # every graph output is consumed (or copied into a persistent buffer) right away.
        self.pool = torch.cuda.graph_pool_handle() if (shared_pool and self.use_graphs) else None
        self.capture_seconds = 0.0
        self.stats = {"seg_replays": 0, "seg_rows": 0, "noncontig": 0}

    def _sync(self) -> None:
        if self.dec.device.type == "cuda":
            torch.cuda.synchronize()

    # ------------------------------------------------------------------ compute
    def _rope_tables(self, n):
        pos = self.pos[:n]
        return pos, self.dec.cos.index_select(0, pos)[:, None], self.dec.sin.index_select(0, pos)[:, None]

    def _pad(self, x: torch.Tensor) -> torch.Tensor:
        n = x.shape[0]
        return x if n == self.max_rows else F.pad(x, (0, 0) * (x.dim() - 1) + (0, self.max_rows - n))

    def _lin(self, mod, x: torch.Tensor) -> torch.Tensor:
        """GEMM on a fixed-width (max_rows) input; returns all max_rows rows."""
        return F.linear(self._pad(x), mod.weight, mod.bias)

    def _mlp(self, mlp, x: torch.Tensor) -> torch.Tensor:
        return self._lin(mlp.down_proj, F.silu(self._lin(mlp.gate_proj, x)) * self._lin(mlp.up_proj, x))

    def _final_logits(self, h: torch.Tensor) -> torch.Tensor:
        n = h.shape[0]
        return self._lin(self.dec.lm_head, self.dec.norm(self._pad(h))).float()[:n]

    def _layers(self, h: torch.Tensor, first: int, last: int) -> torch.Tensor:
        """Layers [first, last) on n real rows, at FIXED WIDTH.

        Every elementwise op, reduction (RMSNorm, q/k norm) and GEMM runs on exactly
        `max_rows` rows (real rows first, zero padding after), so each kernel sees the
        same shape whatever n is and a row is computed bitwise identically in AR (n=1),
        catch-up (m rows) and verification (n rows). Library kernels pick different
        code paths by shape: hipBLASLt/rocBLAS GEMMs (jobs 893616, 893848) and the
        MLP-block elementwise/reduction kernels (job 893872) were not row-invariant.
        Only the K/V cache writes (exact copies) and attention (row-invariant Triton
        kernel) touch the n real rows alone. Pad rows use position pos[n:] (kept valid).
        """
        dec = self.dec
        n = h.shape[0]
        W = self.max_rows
        hd = dec.spec.head_dim
        pos, cos, sin = self._rope_tables(W)
        h = self._pad(h)
        for li in range(first, last):
            layer = dec.layers[li]
            attn = layer.self_attn
            x = layer.input_layernorm(h)
            q = self._lin(attn.q_proj, x).view(W, -1, hd)
            k = self._lin(attn.k_proj, x).view(W, -1, hd)
            v = self._lin(attn.v_proj, x).view(W, -1, hd)
            if dec.spec.has_qk_norm:
                q, k = attn.q_norm(q), attn.k_norm(k)
            q = q * cos + _rotate_half(q) * sin
            k = k * cos + _rotate_half(k) * sin
            dec.k_cache[li].index_copy_(1, self.slot[:n], k[:n].transpose(0, 1))
            dec.v_cache[li].index_copy_(1, self.slot[:n], v[:n].transpose(0, 1))
            o = self._attention(q[:n].contiguous(), li, n)
            h = h + self._lin(attn.o_proj, o.reshape(n, -1))
            if self._trace is not None:
                self._trace.append(("attn_out", li, o.clone(), h[:n].clone()))
            h = h + self._mlp(layer.mlp, layer.post_attention_layernorm(h))
            if self._trace is not None:
                self._trace.append(("layer_out", li, h[:n].clone()))
        return h[:n]

    def _attention(self, q: torch.Tensor, li: int, n: int) -> torch.Tensor:
        dec = self.dec
        if row_attention is not None and q.is_cuda:
            return row_attention(q, dec.k_cache[li], dec.v_cache[li], self.pos, dec.scaling,
                                 cl=self.cl, anc=self.anc)
        outs = []  # reference path (CPU tests): one query per row over its exact path
        cls, ancs = self.cl[:n].tolist(), self.anc[:n].tolist()
        for i, p in enumerate(self.pos[:n].tolist()):
            idx = torch.tensor(list(range(min(cls[i], p + 1))) + ancs[i][: max(p + 1 - cls[i], 0)])
            outs.append(F.scaled_dot_product_attention(
                q[i:i + 1].transpose(0, 1)[None], dec.k_cache[li, :, idx][None],
                dec.v_cache[li, :, idx][None], scale=dec.scaling, enable_gqa=True)[0].transpose(0, 1))
        return torch.cat(outs)

    def _fill(self, h: torch.Tensor, first: int, last: int) -> None:
        dec = self.dec
        hd = dec.spec.head_dim
        W = self.max_rows
        pos, cos, sin = self._rope_tables(W)
        x0 = self._pad(h)
        for li in range(first, last):
            layer = dec.layers[li]
            attn = layer.self_attn
            x = layer.input_layernorm(x0)
            k = self._lin(attn.k_proj, x).view(W, -1, hd)
            if dec.spec.has_qk_norm:
                k = attn.k_norm(k)
            k = k * cos + _rotate_half(k) * sin
            v = self._lin(attn.v_proj, x).view(W, -1, hd)
            dec.k_cache[li].index_copy_(1, pos[:1], k[:1].transpose(0, 1))
            dec.v_cache[li].index_copy_(1, pos[:1], v[:1].transpose(0, 1))

    def _head_out(self, k: int) -> torch.Tensor:
        head = self.heads[k]
        logits = head(self.hb)[0]
        probs = F.softmax(logits, -1)
        conf, arg = probs.max(-1)
        if head.reject:
            body = logits[:-1].argmax(-1)
            conf = torch.where(arg == logits.shape[-1] - 1, torch.zeros_like(conf), conf)
        else:
            body = arg
        return torch.stack([conf.float(), head.to_token_ids(body).float()])

    def _head_out4(self, k: int) -> torch.Tensor:
        """Head k on hb: [conf, id_1..id_4, p_1..p_4] over V_k ∪ ⊥ (⊥ reported as id -1).
        Iterated argmax instead of torch.topk (ROCm 7.0 graph bug, see _g_verify_topk)."""
        head = self.heads[k]
        logits = head(self.hb)[0].float()
        probs = F.softmax(logits, -1)
        conf, arg = probs.max(-1)
        rej = logits.shape[-1] - 1
        if head.reject:
            conf = torch.where(arg == rej, torch.zeros_like(conf), conf)
        ids, ps, pr = [], [], probs
        for _ in range(4):
            i = pr.argmax(-1)
            ps.append(pr.gather(0, i[None])[0])
            tok = head.to_token_ids(i.clamp_max(head.n_vocab - 1)[None])[0].float()
            ids.append(torch.where(i == rej, torch.full_like(tok, -1.0), tok) if head.reject else tok)
            pr = pr.scatter(0, i[None], float("-inf"))
        return torch.stack([conf] + ids + ps)

    def _g_head4(self, k):
        return self._run(("head4", k), lambda: self._head_out4(k))

    # ------------------------------------------------------------------ graphs
    def _run(self, key, fn):
        if not self.use_graphs:
            return fn()
        if key not in self.graphs:
            t0 = time.perf_counter()
            s = torch.cuda.Stream()
            s.wait_stream(torch.cuda.current_stream())
            with torch.cuda.stream(s):
                for _ in range(2):
                    fn()
            torch.cuda.current_stream().wait_stream(s)
            g = torch.cuda.CUDAGraph()
            with torch.cuda.graph(g, pool=self.pool):
                out = fn()
            self.graphs[key] = (g, out)
            self._sync()
            self.capture_seconds += time.perf_counter() - t0
        g, out = self.graphs[key]
        if self.sync_each:
            # ROCm 7.0: a replay is not reliably ordered with the eager work queued just
            # before it either (e.g. the buf <- hrows copy): prompt-lookup cycles verified
            # from layer 0 read a stale buf (job 899287, parity 0.67; exact without graphs or
            # with packet capture off, 899290-899291). A bare sync costs ~4 us (job 899254).
            torch.cuda.synchronize()
        g.replay()
        if self.sync_each:
            torch.cuda.synchronize()
        return out

    def _g_ar(self):
        L = self.dec.spec.num_layers
        return self._run(("ar",), lambda: self._final_logits(
            self._layers(self.dec.embed(self.ids), 0, L)).argmax(-1))

    def _head_sample_out(self, k: int) -> torch.Tensor:
        """Head k on hb at temperature self.temperature: [conf (exit score, 0 if ⊥ wins),
        sampled token id from q' = softmax(V_k logits / T) via Gumbel-max on gnoise, q'(sample)];
        q' is left in qout[:n_vocab]."""
        head = self.heads[k]
        lg = head(self.hb)[0].float()
        probs = F.softmax(lg, -1)
        conf, arg = probs.max(-1)
        if head.reject:
            conf = torch.where(arg == lg.shape[-1] - 1, torch.zeros_like(conf), conf)
        nv = head.n_vocab
        qb = F.softmax(lg[:nv] / self.temperature, -1)
        x = (qb.clamp_min(1e-30).log() + self.gnoise[:nv]).argmax(-1)
        self.qout[:nv].copy_(qb)
        tok = head.to_token_ids(x[None])[0].float()
        return torch.stack([conf, tok, qb.gather(0, x[None])[0]])

    def _g_head_sample(self, k):
        return self._run(("head_sample", k), lambda: self._head_sample_out(k))

    def _g_ar_logits(self):
        L = self.dec.spec.num_layers
        return self._run(("ar_logits",), lambda: self._final_logits(
            self._layers(self.dec.embed(self.ids), 0, L)).float()[0])

    def _g_verify_logits(self, c, n):
        L = self.dec.spec.num_layers
        return self._run(("verify_logits", c, n), lambda: self._final_logits(
            self._layers(self.buf[:n], c, L)).float())

    def _g_ar_conf(self):
        """AR step returning [argmax, max prob] (draft model with adaptive stop)."""
        L = self.dec.spec.num_layers

        def fn():
            pr = F.softmax(self._final_logits(self._layers(self.dec.embed(self.ids), 0, L)).float(), -1)[0]
            conf, arg = pr.max(-1)
            return torch.stack([arg.float(), conf])
        return self._run(("ar_conf",), fn)

    def _g_seg(self, a, b, m):
        def fn():
            self.buf[:m].copy_(self._layers(self.buf[:m], a, b))
            return self.buf
        self.stats["seg_replays"] += 1
        self.stats["seg_rows"] += m
        return self._run(("seg", a, b, m), fn)

    def _g_head(self, k):
        return self._run(("head", k), lambda: self._head_out(k))

    def _g_fill(self, a, b):
        def fn():
            self._fill(self.hb, a, b)
            return self.hb
        return self._run(("fill", a, b), fn)

    TOPK = 8

    def _g_verify_topk(self, c, n):
        """Verification returning, per row, the argmax then the top-TOPK token ids (lossy rules)."""
        L = self.dec.spec.num_layers

        def fn():
            # Iterated argmax instead of torch.topk: topk on a 151k vocabulary crashed inside a
            # HIP graph (illegal address, ROCm 7.0); argmax is the kernel the lossless path uses.
            lg = self._final_logits(self._layers(self.buf[:n], c, L)).float()
            cols = []
            for _ in range(self.TOPK):
                i = lg.argmax(-1, keepdim=True)
                cols.append(i)
                lg = lg.scatter(-1, i, float("-inf"))
            return torch.cat([cols[0]] + cols, -1)
        return self._run(("verify_topk", c, n), fn)

    def _g_verify(self, c, n):
        L = self.dec.spec.num_layers
        return self._run(("verify", c, n), lambda: self._final_logits(
            self._layers(self.buf[:n], c, L)).argmax(-1))

    # ------------------------------------------------------------------ helpers
    def _set_pos(self, plist: list[int], slots: list[int] | None = None) -> None:
        # Pad rows (fixed width) reuse the last real position: always a valid index.
        # slots: KV indices of the rows (default = positions; differ only for star siblings).
        W = self.max_rows
        full = list(plist) + [plist[-1]] * (W - len(plist))
        sl = full if slots is None else list(slots) + [slots[-1]] * (W - len(slots))
        h = self._meta_h
        h[0:W] = torch.tensor(full)
        h[W:2 * W] = torch.tensor(sl)
        # a row whose slot differs from its position reads its own key from the slot
        h[2 * W:3 * W] = torch.tensor([p + 1 if s_ == p else p for p, s_ in zip(full, sl)])
        if slots is None:  # chains never read anc (cl = pos + 1): skip its transfer
            self._meta_d[:3 * W].copy_(h[:3 * W])
        else:
            a = h[3 * W:].view(W, W)
            a.zero_()
            a[:, 0] = h[W:2 * W]
            self._meta_d.copy_(h)

    def _set_rows(self, pos: list[int], slots: list[int], cls: list[int], ancs: list[list[int]]) -> None:
        """Tree rows: positions, KV slots, contiguous prefix lengths and path slots."""
        W = self.max_rows
        assert len(set(slots)) == len(slots), "two rows would write the same KV slot"
        pad = W - len(pos)
        h = self._meta_h
        h[0:W] = torch.tensor(pos + [pos[-1]] * pad)
        h[W:2 * W] = torch.tensor(slots + [slots[-1]] * pad)
        h[2 * W:3 * W] = torch.tensor(cls + [cls[-1]] * pad)
        a = h[3 * W:].view(W, W)
        a.zero_()
        for i, row in enumerate(ancs + [ancs[-1]] * pad):
            if row:
                a[i, :len(row)] = torch.tensor(row)
        self._meta_d.copy_(h)

    def _gather(self, rows: list[int]) -> None:
        m = len(rows)
        if rows == list(range(rows[0], rows[0] + m)):
            self.buf[:m].copy_(self.hrows[rows[0]:rows[0] + m])
        else:
            self.stats["noncontig"] = self.stats.get("noncontig", 0) + 1
            self.idx[:m].copy_(torch.tensor(rows, dtype=torch.int64))
            torch.index_select(self.hrows, 0, self.idx[:m], out=self.buf[:m])

    def _scatter(self, rows: list[int]) -> None:
        m = len(rows)
        if rows == list(range(rows[0], rows[0] + m)):
            self.hrows[rows[0]:rows[0] + m].copy_(self.buf[:m])
        else:
            self.hrows.index_copy_(0, self.idx[:m], self.buf[:m])

    def _embed_row(self, j: int, tok: int) -> None:
        self.ids.fill_(tok)
        self.hrows[j:j + 1].copy_(self.dec.embed(self.ids))

    def _prefill(self, prompt: list[int]) -> int:
        """Prompt processed by the engine itself, in causal chunks of max_rows rows
        through the fixed-width graphs (verify graphs from layer 0). No third-party
        attention kernel runs between graph replays, and the prompt KV is computed
        with exactly the same operators as generated tokens. prefill="full": wide GEMMs
        and standard causal attention instead (_prefill_full)."""
        self._prefix = None  # a greedy prefill overwrites the cached sampling prefix
        if self.prefill_mode == "full":
            return self._prefill_full(prompt)
        W = self.max_rows
        ids = torch.tensor(prompt, device=self.dec.device)
        last = None
        for s0 in range(0, len(prompt), W):
            m = min(W, len(prompt) - s0)
            self._set_pos(list(range(s0, s0 + m)))
            self.buf[:m].copy_(self.dec.embed(ids[s0:s0 + m]))
            last = self._g_verify(0, m)
        return int(last[m - 1])

    @torch.no_grad()
    def _prefill_full(self, prompt: list[int]) -> int:
        """Standard prefill: eager wide GEMMs over the whole prompt (chunks of prefill_chunk
        rows) and causal scaled-dot-product attention, writing the prompt KV in place.

        Decoding exactness is unchanged: AR and speculative decoding both start from this
        same KV and first token, and every decode-time row still goes through the
        fixed-width graphs. Only the prompt KV differs (bf16 rounding of other kernels)
        from the in-engine prefill, so outputs may differ from an engine-prefill run."""
        dec = self.dec
        hd = dec.spec.head_dim
        ids = torch.tensor(prompt, device=dec.device)
        n = len(prompt)
        h = None
        for s0 in range(0, n, self.prefill_chunk):
            m = min(self.prefill_chunk, n - s0)
            pos = torch.arange(s0, s0 + m, device=dec.device)
            cos, sin = dec.cos.index_select(0, pos)[:, None], dec.sin.index_select(0, pos)[:, None]
            h = dec.embed(ids[s0:s0 + m])
            if s0 == 0:
                mask, causal = None, True
            else:  # query i (position s0 + i) sees keys t <= s0 + i
                mask = (torch.arange(s0 + m, device=dec.device)[None] <= pos[:, None])[None, None]
                causal = False
            for li in range(dec.spec.num_layers):
                layer = dec.layers[li]
                attn = layer.self_attn
                x = layer.input_layernorm(h)
                q = attn.q_proj(x).view(m, -1, hd)
                k = attn.k_proj(x).view(m, -1, hd)
                v = attn.v_proj(x).view(m, -1, hd)
                if dec.spec.has_qk_norm:
                    q, k = attn.q_norm(q), attn.k_norm(k)
                q = q * cos + _rotate_half(q) * sin
                k = k * cos + _rotate_half(k) * sin
                dec.k_cache[li, :, s0:s0 + m].copy_(k.transpose(0, 1))
                dec.v_cache[li, :, s0:s0 + m].copy_(v.transpose(0, 1))
                o = F.scaled_dot_product_attention(
                    q.transpose(0, 1)[None], dec.k_cache[li, :, :s0 + m][None], dec.v_cache[li, :, :s0 + m][None],
                    attn_mask=mask, is_causal=causal, scale=dec.scaling, enable_gqa=True)[0]
                h = h + attn.o_proj(o.transpose(0, 1).reshape(m, -1))
                h = h + layer.mlp(layer.post_attention_layernorm(h))
        tok = int(dec.lm_head(dec.norm(h[-1:])).float().argmax(-1)[0])
        self._sync()  # eager kernels must finish before the next graph replay (ROCm ordering)
        return tok

    @staticmethod
    def _boundaries(policy, head_layers):
        # extra_bounds: segment boundaries without a head. Prompt lookup alone uses the h0
        # layer: copies go through seg(0, L_h0) + verify(L_h0) like every other policy.
        # Decode-time verification straight from layer 0 gave a wrong first replay after
        # the prefill on MI210 (ROCm 7.0, jobs 899287/899321: graph token 0, eager correct).
        return sorted({0} | {head_layers[k] for k in policy.head_order}
                      | set(getattr(policy, "extra_bounds", ())))

    def warm_blas(self, max_prefill: int = 0) -> None:
        """Kept for compatibility: with the in-engine prefill, every GEMM has the fixed
        width max_rows, so no larger BLAS call happens after capture."""
        self._warm = self.dec.max_len

    def precapture(self, policies: list, kv_fill: str = "lazy") -> None:
        """Capture every graph the given policies can use, at scratch positions."""
        self._precapture_body(policies, kv_fill)

    def _precapture_body(self, policies: list, kv_fill: str = "lazy") -> None:
        if not getattr(self, "_warm", 0):
            self.warm_blas()
        scratch = self.dec.max_len - self.max_rows - 2
        self._set_pos(list(range(scratch, scratch + self.max_rows)))
        self.ids.fill_(0)
        self._g_ar()
        for m in range(1, self.max_rows + 1):  # in-engine prefill chunks
            self._g_verify(0, m)
        for pol in policies:
            b = self._boundaries(pol, self.head_layers)
            c_min, c_max = (b[1] if len(b) > 1 else 0), b[-1]  # head-less (prompt lookup): [0]
            n_rows = self.max_rows if getattr(pol, "tree", None) else min(pol.max_draft + 1, self.max_rows)
            for k in pol.head_order:
                self._g_head(k)
            for u, v in zip(b, b[1:]):
                for m in range(1, (n_rows if kv_fill == "lazy" else 1) + 1):
                    self._g_seg(u, v, m)
            if kv_fill == "copy":
                for k in pol.head_order:
                    if self.head_layers[k] < c_max:
                        self._g_fill(self.head_layers[k], c_max)
            for n in range(1, n_rows + 1):
                self._g_verify(c_max if kv_fill == "lazy" else c_min, n)
                if getattr(pol, "lossy", False):
                    self._g_verify_topk(c_max if kv_fill == "lazy" else c_min, n)
        for pol in policies:  # tree drafts (star, or trie with star siblings): the top-4 head graph
            st_cfg = getattr(pol, "star", None) or (getattr(pol, "tree", None) or {}).get("star")
            if st_cfg is not None:
                self.hb.zero_()
                self._g_head4(st_cfg.get("head", 0))
        self.stats = {"seg_replays": 0, "seg_rows": 0, "noncontig": 0}
        self._sync()

    # ------------------------------------------------------------------ decoding
    @torch.no_grad()
    def _check_capacity(self, prompt, max_new) -> None:
        # drafts write KV up to max_rows slots beyond the last committed position
        if len(prompt) + max_new + self.max_rows > self.dec.max_len:
            raise ValueError(f"prompt {len(prompt)} + max_new {max_new} + {self.max_rows} draft rows "
                             f"exceeds the KV capacity {self.dec.max_len}")

    def _sample(self, probs: torch.Tensor) -> int:
        return int(torch.multinomial(probs, 1, generator=self._gen))

    def _first_token_sampled(self, prompt: list[int]) -> int:
        """Prefill prompt[:-1] (KV only), then the last prompt token through the AR graph: the
        first new token is SAMPLED from its logits (same arithmetic for AR and speculation)."""
        # prompt[:-1] KV is never rewritten by generation (writes start at len(prompt) - 1): reuse it
        # across samples of the same prompt (distribution tests draw thousands)
        if len(prompt) > 1 and not (getattr(self, "reuse_prefix", False)
                                    and getattr(self, "_prefix", None) == tuple(prompt[:-1])):
            self._prefill(prompt[:-1])
            self._prefix = tuple(prompt[:-1])
        self._set_pos([len(prompt) - 1])
        self.ids.fill_(prompt[-1])
        return self._sample(F.softmax(self._g_ar_logits() / self.temperature, -1))

    def _sample_verify(self, c_verify: int, drafts: list[int], st) -> list[int]:
        """Exact speculative sampling on a chain (Leviathan et al.): draft i is accepted with
        probability min(1, p_i(x_i) / q_i(x_i)); at the first rejection the token is resampled
        from norm(max(p_i - q_i, 0)); if all are accepted the bonus is sampled from p_n. q_i is
        the head's renormalised V_k distribution at T (saved at drafting) or a point mass (copy)."""
        n = len(drafts) + 1
        P = F.softmax(self._g_verify_logits(c_verify, n)[:n] / self.temperature, -1)
        st.cycles += 1
        st.drafted += len(drafts)
        st.empty_drafts += int(not drafts)
        committed = []
        if drafts:
            pv = P[torch.arange(len(drafts), device=P.device), torch.tensor(drafts, device=P.device)].tolist()
        for i, x in enumerate(drafts):
            k, qx = self._draft_q[i]
            if self._rng.random() < min(1.0, pv[i] / max(qx, 1e-30)):
                committed.append(x)
                continue
            res = P[i].clone()
            if k < 0:
                res[x] = 0.0
            else:
                h = self.heads[k]
                q = torch.zeros_like(res)
                q[h.vocab_ids] = self.qsave[i, :h.n_vocab]
                res = (res - q).clamp_min(0.0)
            tot = float(res.sum())
            committed.append(self._sample(res / tot if tot > 0 else P[i]))
            st.rejected_cycles += 1
            break
        else:
            committed.append(self._sample(P[len(drafts)]))
        st.accepted += min(len(committed) - 1, len(drafts))
        return committed

    @staticmethod
    def _emit(st: GenStats, on_commit, toks: list[int], srcs: list[str]) -> None:
        st.sources.extend(srcs)
        if on_commit is not None:
            on_commit(list(toks), list(srcs))

    def _commit(self, out, committed, srcs, index, stop, eos_ids, max_new, st, on_commit) -> bool:
        """Append a cycle's committed tokens to out (truncated at stop / EOS / max_new), feed the
        prompt-lookup index, record their sources. Returns True if the stop rule fired."""
        n0, stopped = len(out), False
        for t in committed:
            out.append(t)
            if index is not None:
                index.extend([t])
            if stop is not None and stop(t):
                stopped = True
                break
            if t in eos_ids or len(out) >= max_new:
                break
        self._emit(st, on_commit, out[n0:], srcs[:len(out) - n0])
        return stopped

    def generate_ar(self, prompt: list[int], max_new: int, eos_ids=frozenset(), stop=None,
                    temperature: float = 0.0, gen: torch.Generator | None = None,
                    on_commit=None) -> tuple[list[int], GenStats]:
        """stop: optional callable fed every new token, True = stop after it (same rule as
        generate_spec, so both stop at the same token). on_commit(tokens, sources): optional
        callable fed the tokens appended to the output (and their sources) after the prefill
        and after every step."""
        self._check_capacity(prompt, max_new)
        st = GenStats()
        self._sync()
        t0 = time.perf_counter()
        if temperature > 0:
            self.temperature, self._gen = temperature, gen
            out = [self._first_token_sampled(prompt)]
            self._sync()
            st.prefill_seconds = time.perf_counter() - t0
            self._emit(st, on_commit, out[:1], ["model"])
            while len(out) < max_new and out[-1] not in eos_ids and not (stop is not None and stop(out[-1])):
                self._set_pos([len(prompt) + len(out) - 1])
                self.ids.fill_(out[-1])
                out.append(self._sample(F.softmax(self._g_ar_logits() / temperature, -1)))
                self._emit(st, on_commit, out[-1:], ["model"])
                st.cycles += 1
            self._sync()
            st.seconds = time.perf_counter() - t0
            st.new_tokens = len(out)
            return out, st
        out = [self._prefill(prompt)]
        self._sync()
        st.prefill_seconds = time.perf_counter() - t0
        self._emit(st, on_commit, out[:1], ["model"])
        stopped = stop is not None and stop(out[0])
        steps, t_prev = [], time.perf_counter()
        while not stopped and len(out) < max_new and out[-1] not in eos_ids:
            self._set_pos([len(prompt) + len(out) - 1])
            self.ids.fill_(out[-1])
            out.append(int(self._g_ar()))       # .item() synchronises: per-step wall time
            self._emit(st, on_commit, out[-1:], ["model"])
            if stop is not None and stop(out[-1]):
                break
            t_now = time.perf_counter()
            steps.append(t_now - t_prev)
            t_prev = t_now
            st.cycles += 1
        if steps:  # robust unit for cost controllers (outliers: syncs, GC)
            st.ar_step_median = float(np.median(steps))
        self._sync()
        st.seconds = time.perf_counter() - t0
        st.new_tokens = len(out)
        return out, st

    def _draft_lazy(self, policy, first_tok, pos0, budget, eos_ids, st, proposal=None, pld_mode=None):
        """Regime A. Returns drafts; on exit all block rows sit in hrows at c_max.

        proposal / pld_mode (prompt lookup): "first" drafts the copied tokens without
        evaluating any head; "fallback" uses them only if no head exits at the first
        position (an empty cycle). Copied drafts enter at depth 0 and are caught up
        lazily with the other rows (head index -1)."""
        hl = self.head_layers
        b = self._boundaries(policy, hl)
        c_max = b[-1]
        nxt = {u: v for u, v in zip(b, b[1:])}
        depth: list[int] = []

        def advance(target: int) -> None:
            while True:
                lag = [j for j, d in enumerate(depth) if d < target]
                if not lag:
                    return
                u = min(depth[j] for j in lag)
                batch = [j for j in lag if depth[j] == u]
                self._set_pos([pos0 + j for j in batch])
                self._gather(batch)
                self._g_seg(u, nxt[u], len(batch))
                self._scatter(batch)
                for j in batch:
                    depth[j] = nxt[u]

        drafts: list[int] = []
        self._draft_heads, self._stop_info, self._draft_q = [], None, []
        tok = first_tok

        def copy_in() -> None:
            for t in proposal[: budget - len(drafts)]:
                drafts.append(t)
                self._draft_heads.append(-1)
                if self._sampling:
                    self._draft_q.append((-1, 1.0))  # copy: deterministic draft, q = delta
                st.exits["pld"] = st.exits.get("pld", 0) + 1
                self._embed_row(len(depth), t)
                depth.append(0)
                if t in eos_ids:
                    break

        while True:
            j = len(depth)
            self._embed_row(j, tok)
            depth.append(0)
            if len(drafts) >= budget or (drafts and drafts[-1] in eos_ids):
                break
            if j == 0 and proposal and pld_mode == "first":
                copy_in()
                break
            exited, row = None, {}
            for k in policy.head_order:
                advance(hl[k])
                self.hb.copy_(self.hrows[j:j + 1])
                if self._sampling:
                    self.gnoise.uniform_(1e-10, 1.0, generator=self._gen).log_().neg_().log_().neg_()
                    conf, tok_f, qx = self._g_head_sample(k).tolist()
                else:
                    conf, tok_f = self._g_head(k).tolist()
                row[k] = (conf, int(tok_f))
                st.head_evals += 1
                if policy.decide(k, conf, len(drafts)):
                    exited = k
                    if self._sampling:  # keep q'_k of this draft row for the residual
                        self.qsave[len(drafts), :self.heads[k].n_vocab].copy_(self.qout[:self.heads[k].n_vocab])
                        self._draft_q.append((k, qx))
                    break
            if exited is None:
                self._stop_info = row  # every head was evaluated at the stop position
                if j == 0 and proposal and pld_mode == "fallback":
                    self._stop_info = None
                    copy_in()
                break
            tok = int(tok_f)
            st.exits[exited] = st.exits.get(exited, 0) + 1
            drafts.append(tok)
            self._draft_heads.append(exited)
        advance(c_max)
        return drafts, c_max

    def _draft_copy(self, policy, first_tok, pos0, budget, eos_ids, st):
        """Regime B (CALM copy). Returns drafts; rows at c_min are left in hrows."""
        hl = self.head_layers
        b = self._boundaries(policy, hl)
        c_min, c_max = b[1], b[-1]
        nxt = {u: v for u, v in zip(b, b[1:])}
        drafts: list[int] = []
        self._draft_heads = []  # head of each draft (provenance)
        tok = first_tok
        while True:
            j = len(drafts)
            self._set_pos([pos0 + j])
            self._embed_row(j, tok)
            self.buf[:1].copy_(self.hrows[j:j + 1])
            self._g_seg(0, c_min, 1)
            self.hrows[j:j + 1].copy_(self.buf[:1])  # hidden at c_min (verified later)
            if len(drafts) >= budget or (drafts and drafts[-1] in eos_ids):
                break
            depth, exited = c_min, None
            for k in policy.head_order:
                while depth < hl[k]:
                    self._g_seg(depth, nxt[depth], 1)
                    depth = nxt[depth]
                self.hb.copy_(self.buf[:1])
                conf, tok_f = self._g_head(k).tolist()
                st.head_evals += 1
                if policy.decide(k, conf, len(drafts)):
                    exited = k
                    break
            if exited is None:
                break
            if depth < c_max:
                self._g_fill(depth, c_max)
            tok = int(tok_f)
            st.exits[exited] = st.exits.get(exited, 0) + 1
            drafts.append(tok)
            self._draft_heads.append(exited)
        return drafts, c_min

    def _star_cycle(self, star: dict, first_tok: int, pos0: int, budget: int, st) -> list[int]:
        """Depth-1 tree draft (star): root = last committed token; the shallow head (index
        star["head"]) proposes up to C children from its top-4 at the root; children are
        computed on their own KV slots pos0+1+i at position pos0+1, verified together with
        the root; the accepted child's KV is compacted to slot pos0+1. Exact (lossless).
        Admission: q_r = p_r (head probability of rank r, ⊥ excluded) >= tau[r], and the
        survival sum of the admitted q_r >= gamma. Returns the committed tokens."""
        k = star.get("head", 0)
        c = self.head_layers[k]
        self._embed_row(0, first_tok)
        self._set_pos([pos0])
        self._gather([0])
        self._g_seg(0, c, 1)
        self._scatter([0])
        self.hb.copy_(self.hrows[0:1])
        r = self._g_head4(k).tolist()
        st.head_evals += 1
        ids, ps = [int(x) for x in r[1:5]], r[5:9]
        tau = star.get("tau", [0.05, 0.02, 0.02, 0.02])
        kids = [t for t, p, th in zip(ids, ps, tau) if t >= 0 and p >= th][: min(star.get("C", 4), budget)]
        qs = [p for t, p, th in zip(ids, ps, tau) if t >= 0 and p >= th][: len(kids)]
        if kids and sum(qs) < star.get("gamma", 0.3):
            kids = []
        m = len(kids)
        if m:
            for i, t in enumerate(kids):
                self._embed_row(1 + i, t)
            self._set_pos([pos0 + 1] * m, [pos0 + 1 + i for i in range(m)])
            self._gather(list(range(1, 1 + m)))
            self._g_seg(0, c, m)
            self._scatter(list(range(1, 1 + m)))
        n = 1 + m
        self._set_pos([pos0] + [pos0 + 1] * m, [pos0] + [pos0 + 1 + i for i in range(m)])
        self.buf[:n].copy_(self.hrows[:n])
        target = self._g_verify(c, n).tolist()
        st.cycles += 1
        st.drafted += m
        st.empty_drafts += int(m == 0)
        committed = [target[0]]
        if m:
            st.exits["star"] = st.exits.get("star", 0) + m
            j = next((i for i, t in enumerate(kids) if t == target[0]), None)
            if j is None:
                st.rejected_cycles += 1
            else:
                st.accepted += 1
                committed.append(target[1 + j])
                if j > 0:  # compaction: the accepted child's KV (all layers) to slot pos0+1
                    self.stats["star_compactions"] = self.stats.get("star_compactions", 0) + 1
                    dk, dv = self.dec.k_cache, self.dec.v_cache
                    dk[:, :, pos0 + 1].copy_(dk[:, :, pos0 + 1 + j])
                    dv[:, :, pos0 + 1].copy_(dv[:, :, pos0 + 1 + j])
        return committed

    def _tree_cycle(self, policy, cfg: dict, root_tok: int, pos0: int, eos_ids, st, proposal,
                    lam: float = float("-inf"), match_n: int = 0, ext=None):
        """General tree draft (trie), one verification. Rows = nodes in creation order
        (row 0 = root at pos0); node i sits at position pos0 + depth, writes its KV at slot
        pos0 + i and attends to [0, pos0] plus the slots of its own path.
          1. prompt-lookup copy (proposal, <= cfg["m"]) inserted under the root first, so it
             rides in the root's first segment;
          2. head chain (lazy, policy.decide, <= cfg["head_budget"] drafts) that reuses a
             copy row when its token matches; a head is never evaluated on a row already
             computed beyond its layer;
          3. optional star siblings at the root from the top-4 of cfg["star"]["head"].
        Verification walks the trie from the root; the accepted path's KV is compacted to
        positions pos0+1.. Exact (lossless).

        Safety gate (chantier 3b §7.6, alpha guarantee on trees): every node v carries a
        predicted acceptance q_v (head: a_hat_k; copy: cfg["gate"]["q_copy"][n][d] for match
        length n and copy depth d; star: isotone map of the rank probability). After the
        rows are computed, the children of u are KEPT (drafted, verified) only if
        p(u) = min(1, sum_v q_v) >= lam; lam > 1 keeps none (pure probe). Returns
        (committed tokens, info) with info = path origins, failed child origin, stop (p, a)
        and the number of drafted nodes, for online.cycle_from_tree."""
        gate = cfg.get("gate") or {}
        qc = gate.get("q_copy")
        qs = gate.get("q_star")
        hl = self.head_layers
        W = self.max_rows
        b = self._boundaries(policy, hl)
        c_max = b[-1]
        nxt = dict(zip(b, b[1:]))
        tok, par, dep, ld, kids, origin, q = [], [], [], [], [], [], []
        shown = []  # provenance reported to the user (st.sources); origin feeds the statistics

        def add(t, parent, how, qv=0.0):
            if len(tok) >= W:
                return None
            i = len(tok)
            tok.append(t); par.append(parent); dep.append(0 if parent is None else dep[parent] + 1)
            ld.append(0); kids.append({}); origin.append(how); shown.append(how); q.append(float(qv))
            if parent is not None:
                kids[parent][t] = i
            self._embed_row(i, t)
            return i

        def path(i):
            out = []
            while par[i] is not None:
                out.append(pos0 + i)
                i = par[i]
            return out[::-1]

        def set_rows(rows):
            self._set_rows([pos0 + dep[i] for i in rows], [pos0 + i for i in rows],
                           [pos0 + 1] * len(rows), [path(i) for i in rows])

        def advance(target):
            while True:
                lag = [i for i in range(len(tok)) if ld[i] < target]
                if not lag:
                    return
                u = min(ld[i] for i in lag)
                batch = [i for i in lag if ld[i] == u]
                set_rows(batch)
                self._gather(batch)
                self._g_seg(u, nxt[u], len(batch))
                self._scatter(batch)
                for i in batch:
                    ld[i] = nxt[u]

        root = add(root_tok, None, "root")
        c = root
        for t in (ext or []):  # external draft model chain (P1: SD + PLD in one verification)
            n = kids[c].get(t)
            if n is None:
                n = add(t, c, "draft")
                if n is None:
                    break
                st.exits["draft"] = st.exits.get("draft", 0) + 1
            c = n
            if t in eos_ids:
                break
        c = root
        for d, t in enumerate((proposal or [])[: cfg.get("m", 15)], start=1):
            n = kids[c].get(t)
            if n is None:
                qv = 0.0
                if qc is not None:
                    row = qc.get(str(match_n)) or qc.get(str(max(int(k) for k in qc)))
                    qv = row[min(d, len(row)) - 1]
                n = add(t, c, "copy", qv)
                if n is None:
                    break
                st.pld_drafted += 1
            c = n
            if t in eos_ids:
                break
        star = cfg.get("star")
        if hasattr(policy, "start_block"):
            policy.start_block()
        c, drafted = root, 0
        while drafted < cfg.get("head_budget", 8) and len(tok) < W and (c == root or tok[c] not in eos_ids):
            exited = None
            for k in policy.head_order:
                if hl[k] < ld[c]:
                    continue  # row already beyond this head's layer (shared copy row)
                advance(hl[k])
                self.hb.copy_(self.hrows[c:c + 1])
                if star is not None and c == root and k == star.get("head", 0):
                    r = self._g_head4(k).tolist()
                    conf, t_k = r[0], int(r[1])
                    tau = star.get("tau", [0.05, 0.02, 0.02, 0.02])
                    adm = [(int(x), p, rk) for rk, (x, p, th) in enumerate(zip(r[1:5], r[5:9], tau))
                           if x >= 0 and p >= th]
                    if adm and sum(p for _, p, _ in adm) >= star.get("gamma", 0.3):
                        for x, p, rk in adm[: star.get("C", 4)]:
                            qv = float(np.interp(p, *qs[rk])) if qs is not None else 0.0
                            if x in kids[root]:
                                q[kids[root][x]] = max(q[kids[root][x]], qv)
                            elif add(x, root, "star", qv) is not None:
                                st.exits["star"] = st.exits.get("star", 0) + 1
                else:
                    conf, t_k = self._g_head(k).tolist()
                    t_k = int(t_k)
                st.head_evals += 1
                if conf > 0 and policy.decide(k, conf, drafted):  # conf == 0: ⊥ won
                    exited = k
                    break
            if exited is None:
                break
            if gate.get("q_head") is not None:  # head map refitted for the gate (pool A)
                q_head = float(np.interp(conf, *gate["q_head"]))
            else:
                q_head = float(policy._ahat(exited, conf)) if hasattr(policy, "_ahat") else float(conf)
            n = kids[c].get(t_k)
            if n is None:
                n = add(t_k, c, "head", q_head)
                if n is None:
                    break
                st.exits[exited] = st.exits.get(exited, 0) + 1
            else:
                q[n] = max(q[n], q_head)
                if origin[n] == "star":  # the star's rank-0 candidate is the chain's own token
                    shown[n] = "head"
            c = n
            drafted += 1
            if t_k in eos_ids:
                break
        advance(c_max)
        # safety gate: keep the children of u only if p(u) = min(1, sum_v q_v) >= lam. Children
        # are distinct tokens and at most one equals the greedy target: disjoint events, so the
        # sum is the calibrated probability that one is accepted (chantier 3b §7.6)
        gate_p = [min(1.0, sum(q[v] for v in kids[u].values())) for u in range(len(tok))]
        if gate.get("g") is not None:  # final isotone recalibration g(P) -> HIT (pool A)
            gate_p = [float(np.interp(x, *gate["g"])) if kids[u] else 0.0 for u, x in enumerate(gate_p)]
        keep, frontier = [root], [root]
        while frontier:
            u = frontier.pop(0)
            if kids[u] and gate_p[u] >= lam:
                ch = sorted(kids[u].values())
                keep += ch
                frontier += ch
        keep.sort()
        kept = set(keep)
        n_rows = len(keep)
        set_rows(keep)
        if keep == list(range(n_rows)):
            self.buf[:n_rows].copy_(self.hrows[:n_rows])
        else:
            self.buf[:n_rows].copy_(self.hrows.index_select(0, torch.tensor(keep, device=self.hrows.device)))
        tgt_rows = self._g_verify(c_max, n_rows).tolist()
        target = {u: tgt_rows[j] for j, u in enumerate(keep)}
        st.cycles += 1
        st.drafted += n_rows - 1
        st.empty_drafts += int(n_rows == 1)
        acc, node = [], root
        while True:
            nxt_node = kids[node].get(target[node])
            if nxt_node is None or nxt_node not in kept:
                break
            acc.append(nxt_node)
            node = nxt_node
        st.accepted += len(acc)
        st.pld_accepted += sum(1 for i in acc if origin[i] == "copy")
        kept_kids = [v for v in kids[node].values() if v in kept]
        st.rejected_cycles += int(bool(kept_kids))  # some draft below the last accepted node failed
        info = {"path": [origin[i] for i in acc], "sources": [shown[i] for i in acc], "n_nodes": n_rows - 1,
                "failed": origin[max(kept_kids, key=lambda v: q[v])] if kept_kids else None,
                "stop": None if kept_kids else (gate_p[node], float(target[node] in kids[node]))}
        src = [pos0 + i for i in acc]
        dst = [pos0 + 1 + d for d in range(len(acc))]
        if src != dst:  # compaction of the accepted path (all layers, K and V)
            self.stats["tree_compactions"] = self.stats.get("tree_compactions", 0) + 1
            si = torch.tensor(src, device=self.dec.k_cache.device)
            di = torch.tensor(dst, device=self.dec.k_cache.device)
            for cache in (self.dec.k_cache, self.dec.v_cache):
                cache.index_copy_(2, di, cache.index_select(2, si))
        return [tok[i] for i in acc] + [target[node]], info

    @torch.no_grad()
    def generate_spec(self, policy, prompt: list[int], max_new: int, eos_ids=frozenset(),
                      kv_fill: str = "lazy", controller=None, predictor=None,
                      ar_step_seconds: float | None = None, accept_rule=None, pld=None,
                      stop=None, star=None, drafter=None, temperature: float = 0.0,
                      gen: torch.Generator | None = None, on_commit=None) -> tuple[list[int], GenStats]:
        """controller/predictor: online calibration (chantier 3b). Before each cycle the
        controller's lambda is turned into confidence thresholds (exit at the first head
        with a_hat_k >= lambda); lambda > 1 means a plain AR step (or, if the controller
        probes, a cycle that evaluates every head without drafting). After verification
        the cycle feedback (labels on the verified prefix, stop-position head outputs) is
        sent back. Requires kv_fill="lazy".

        accept_rule: None (exact, lossless) or a LOSSY relaxation of greedy verification:
          ("topk", k)       accept a draft that is among the target's top-k tokens (r3)
          ("equiv", cls)    accept a draft in the same formatting class as the target
                            argmax, cls = list mapping token id -> class id (r5)
        An accepted draft is committed as is; rejections and bonus use the target argmax.

        pld: prompt lookup drafts, dict(mode="first"|"fallback", n_max=4, n_min=1, m=15); with
        a policy without heads, "fallback" is plain PLD. stop, on_commit: as in generate_ar.
        st.sources[i] is where output token i came from: "copy" (prompt lookup), "head" (head
        chain), "star" (top-k sibling of h0), "draft" (draft model) or "model" (prefill token,
        correction or bonus token of each cycle)."""
        from .calibration import online as onl
        from .pld import NgramIndex
        self._check_capacity(prompt, max_new)
        assert policy.max_draft + 1 <= self.max_rows
        assert controller is None or kv_fill == "lazy"
        st = GenStats()
        K = len(self.heads)
        gamma_only = bool(getattr(controller, "gamma_only", False))
        xi = float(getattr(controller, "xi", 0.0))

        def account(cost, emitted, drafting):
            g = cost - (1 + xi) * emitted
            st.cost_units += cost
            st.excess_end += g
            st.excess_max = max(st.excess_max, st.excess_end)
            if drafting:
                st.h_max = max(st.h_max, g)
        if controller is not None:
            st.lam_start = _lam(controller)
            st.gamma_start = float(getattr(controller, "gamma", float("nan")))
        draft = self._draft_lazy if kv_fill == "lazy" else self._draft_copy
        seg_rows0 = self.stats["seg_rows"]  # engine counters are cumulative: report this request only
        self._sync()
        t0 = time.perf_counter()
        self._sampling = temperature > 0
        if self._sampling:  # exact speculative sampling on chains (heads + copies)
            assert controller is None and accept_rule is None and star is None and drafter is None
            assert getattr(policy, "tree", None) is None and kv_fill == "lazy"
            import random
            self.temperature, self._gen = temperature, gen
            self._rng = random.Random(gen.initial_seed() if gen is not None else 0)
            out = [self._first_token_sampled(prompt)]
        else:
            out = [self._prefill(prompt)]
        self._sync()
        st.prefill_seconds = time.perf_counter() - t0
        self._emit(st, on_commit, out[:1], ["model"])
        assert pld is None or (controller is None and kv_fill == "lazy")
        index = None
        if pld is not None:
            index = NgramIndex(pld.get("n_max", 4), pld.get("n_min", 1))
            index.extend(prompt)
            index.extend(out)
        stopped = stop is not None and stop(out[0])
        assert star is None or (controller is None and pld is None and accept_rule is None and kv_fill == "lazy")
        if drafter is not None:  # optional prompt lookup FIRST (copy if a match exists, else the draft model)
            assert controller is None and kv_fill == "lazy" and (pld is None or pld.get("mode") == "first")
            drafter.start(prompt)
        tree = getattr(policy, "tree", None)
        if tree is not None:
            # controllers on a tree: a gamma-only cost safety (NoSlowdown, Prop. 4), or an ACI
            # on the safety gate lambda (virtual head, chantier 3b §7.6), possibly wrapped in it
            assert accept_rule is None and kv_fill == "lazy"
            if tree.get("copy") and index is None:
                index = NgramIndex(tree["copy"].get("n_max", 4), tree["copy"].get("n_min", 1))
                index.extend(prompt)
                index.extend(out)
        while tree is not None and not stopped and len(out) < max_new and out[-1] not in eos_ids:
            pos0 = len(prompt) + len(out) - 1
            lam, probe = float("-inf"), False
            if controller is not None:  # gamma floor / gate lambda; plain AR step when asked
                lamv, gamma, ar = controller.policy()
                if not gamma_only:
                    lam = float(np.atleast_1d(lamv)[0])
                if ar and controller.probe:  # counterfactual probe: compute the tree, draft nothing
                    probe, lam, ar = True, float("inf"), False
                if ar:
                    self._set_pos([pos0])
                    self.ids.fill_(out[-1])
                    out.append(int(self._g_ar()))
                    self._emit(st, on_commit, out[-1:], ["model"])
                    st.cycles += 1
                    st.ar_steps += 1
                    controller.update(onl.Cycle([], [], 0, 0, 1, 1.0, ar=True))
                    account(1.0, 1, False)
                    if index is not None:
                        index.extend([out[-1]])
                    if stop is not None and stop(out[-1]):
                        stopped = True
                    continue
                base = getattr(policy, "_gamma_base", None)
                if base is None:
                    base = policy._gamma_base = policy.cum_gamma
                policy.cum_gamma = ({k: max(v, float(gamma)) for k, v in base.items()} if isinstance(base, dict)
                                    else max(float(base), float(gamma)))
                t_cycle = time.perf_counter()
            prop = index.propose(tree["copy"].get("m", 15)) if tree.get("copy") else None
            ext = drafter.propose(prompt + out, min(max_new - len(out), self.max_rows - 1)) if drafter is not None else None
            committed, info = self._tree_cycle(policy, tree, out[-1], pos0, eos_ids, st, prop, lam=lam,
                                               match_n=getattr(index, "last_n", 0) if index is not None else 0,
                                               ext=ext)
            if drafter is not None:  # draft KV valid up to the common prefix with the accepted path
                k = 0
                while k < min(len(ext), len(committed) - 1) and ext[k] == committed[k]:
                    k += 1
                drafter.commit(k)
            if controller is not None:
                cost = ((time.perf_counter() - t_cycle) / ar_step_seconds) if ar_step_seconds else 1.0
                if gamma_only:
                    cyc = onl.Cycle([], [], len(committed) - 1, 0, len(committed), cost)
                else:
                    cyc = onl.cycle_from_tree(info["path"], info["failed"], emitted=len(committed), cost=cost,
                                              stop=info["stop"], n_nodes=info["n_nodes"])
                    cyc.ar = probe
                controller.update(cyc)
                st.lab_n += len(cyc.labels)
                st.lab_s += float(sum(cyc.labels))
                account(cost, len(committed), True)
            stopped = self._commit(out, committed, info["sources"] + ["model"], index, stop, eos_ids,
                                   max_new, st, on_commit) or stopped
        while star is not None and not stopped and len(out) < max_new and out[-1] not in eos_ids:
            pos0 = len(prompt) + len(out) - 1
            committed = self._star_cycle(star, out[-1], pos0, max_new - len(out), st)
            stopped = self._commit(out, committed, ["star"] * (len(committed) - 1) + ["model"], None, stop,
                                   eos_ids, max_new, st, on_commit) or stopped
        while not stopped and len(out) < max_new and out[-1] not in eos_ids:
            pos0 = len(prompt) + len(out) - 1
            budget = min(policy.max_draft, max_new - len(out))
            t_cycle = time.perf_counter()
            probe = False
            if controller is not None:
                taus, gamma, ar = controller.policy()
                if ar and not controller.probe:
                    self._set_pos([pos0])
                    self.ids.fill_(out[-1])
                    out.append(int(self._g_ar()))
                    self._emit(st, on_commit, out[-1:], ["model"])
                    st.cycles += 1
                    st.ar_steps += 1
                    controller.update(onl.Cycle([], [], 0, 0, 1, 1.0, ar=True))
                    account(1.0, 1, False)
                    if stop is not None and stop(out[-1]):  # same rule as the tree branch
                        stopped = True
                    continue
                probe = bool(ar)
                if not gamma_only:  # gamma_only: the policy keeps its heads and a_hat thresholds
                    tau_h = predictor.thresholds(taus)          # exit iff H = -conf <= tau
                    policy.thresholds = {k: (float("inf") if probe else float(-tau_h[k])) for k in range(K)}
                    policy.head_order = list(range(K))
                if hasattr(policy, "cum_gamma"):
                    base = getattr(policy, "_gamma_base", None)
                    if base is None:
                        base = policy._gamma_base = policy.cum_gamma
                    if isinstance(base, dict):  # per-head gamma_k: gamma_k <- max(gamma_k^pol, gamma_safety)
                        policy.cum_gamma = {k: max(v, float(gamma)) for k, v in base.items()}
                    else:
                        policy.cum_gamma = float(gamma)
            if hasattr(policy, "start_block"):
                policy.start_block()
            used_drafter = False
            if drafter is not None:
                prop = index.propose(pld.get("m", 15)) if index is not None else []
                if not prop:
                    prop, used_drafter = drafter.propose(prompt + out, budget), True
                drafts, c_verify = draft(policy, out[-1], pos0, budget, eos_ids, st, proposal=prop, pld_mode="fallback")
                if used_drafter:  # draft-model tokens went through the copy path: count them apart
                    nd = sum(1 for k in self._draft_heads if k < 0)
                    st.exits["pld"] = st.exits.get("pld", 0) - nd
                    st.exits["draft"] = st.exits.get("draft", 0) + nd
            elif index is not None:
                drafts, c_verify = draft(policy, out[-1], pos0, budget, eos_ids, st,
                                         proposal=index.propose(pld.get("m", 15)), pld_mode=pld.get("mode", "fallback"))
            else:
                drafts, c_verify = draft(policy, out[-1], pos0, budget, eos_ids, st)
            n = len(drafts) + 1
            self._set_pos(list(range(pos0, pos0 + n)))
            self.buf[:n].copy_(self.hrows[:n])
            if self._sampling:
                committed = self._sample_verify(c_verify, drafts, st)
            else:
                if accept_rule is None:
                    target = self._g_verify(c_verify, n).tolist()
                    tk = None
                else:
                    rows = self._g_verify_topk(c_verify, n).tolist()
                    target, tk = [r[0] for r in rows], [r[1:] for r in rows]
                st.cycles += 1
                st.drafted += len(drafts)
                st.empty_drafts += int(not drafts)
                committed = []
                for i, d in enumerate(drafts):
                    ok = target[i] == d
                    if not ok and accept_rule is not None:
                        if accept_rule[0] == "topk":
                            ok = d in tk[i][: accept_rule[1]]
                        elif accept_rule[0] == "equiv":
                            ok = accept_rule[1][d] == accept_rule[1][target[i]]
                    if not ok:
                        break
                    committed.append(d)
                st.accepted += len(committed)
                if index is not None and not (drafter is not None and used_drafter):
                    st.pld_drafted += sum(1 for k in self._draft_heads if k < 0)
                    st.pld_accepted += sum(1 for k in self._draft_heads[:len(committed)] if k < 0)
                st.rejected_cycles += int(len(committed) < len(drafts))
                if hasattr(policy, "update"):
                    policy.update(len(drafts), len(committed))
                if drafter is not None and used_drafter:
                    drafter.commit(len(committed))
                committed.append(target[len(committed)])  # correction or bonus token
            if controller is not None:
                n_acc = st_acc = len(committed) - 1
                per_token = [(k, i < n_acc) for i, k in enumerate(self._draft_heads)]
                stop_pos = None  # (a_hat, labels) at the stop position (not the `stop` rule)
                if (predictor is not None and n_acc == len(drafts) and self._stop_info is not None
                        and len(self._stop_info) == K):
                    confs = np.array([self._stop_info[k][0] for k in range(K)])
                    toks = [self._stop_info[k][1] for k in range(K)]
                    ahat = predictor.predict(-confs[None, :])[0]
                    stop_pos = (ahat, np.array([float(t == target[len(drafts)]) for t in toks]))
                cost = ((time.perf_counter() - t_cycle) / ar_step_seconds) if ar_step_seconds else 1.0
                cyc = onl.cycle_from_hook(per_token, emitted=len(committed), cost=cost, stop=stop_pos)
                cyc.ar = probe
                controller.update(cyc)
                st.lab_n += len(cyc.labels)
                st.lab_s += float(sum(cyc.labels))
                account(cost, len(committed), True)
            srcs = [("draft" if used_drafter else "copy") if k < 0 else "head"
                    for k in self._draft_heads[:len(committed) - 1]] + ["model"]
            stopped = self._commit(out, committed, srcs, index, stop, eos_ids, max_new, st, on_commit) or stopped
        self._sync()
        st.seconds = time.perf_counter() - t0
        st.new_tokens = len(out)
        st.layer_rows = self.stats["seg_rows"] - seg_rows0
        self._sampling = False
        if controller is not None:
            st.lam_end = _lam(controller)
            st.n_jumps = int(getattr(getattr(controller, "aci", controller), "n_jumps", -1))
            st.gamma_end = float(getattr(controller, "gamma", float("nan")))
        return out, st