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