Spaces:
Sleeping
Sleeping
| """Autoregressive and self-speculative decoding on a LayeredDecoder. | |
| One cycle of self-speculative decoding: | |
| 1. Drafting. The last committed token (not yet cached) and then each draft | |
| token climb the layers; at every head allowed by the policy we ask the | |
| policy whether to exit. An exit yields a draft token from that head; the | |
| K/V of the deeper drafting layers are filled from the exit hidden state | |
| (CALM-style copy) so later draft tokens can attend to it. A token with no | |
| exit, the draft budget, or a drafted EOS ends the draft. | |
| 2. Verification. Every row of the cycle has passed through layer c_min (the | |
| shallowest exit the policy can use) with exact inputs, so its hidden state | |
| there is exact. One pass runs layers [c_min, L) over all rows, overwriting | |
| the approximate K/V, and yields the full-model distribution at each row. | |
| 3. Acceptance. Standard speculative sampling (Leviathan et al. 2023): accept | |
| draft x with prob min(1, p(x)/q(x)), else resample from norm(max(p-q, 0)); | |
| if all drafts are accepted, sample a bonus token from the last row. At T=0 | |
| this is exact argmax matching. The output law (greedy: the output) equals | |
| the full model's. | |
| Cache invariant between cycles: positions [0, len(seq)-1) are exact in every | |
| layer; the last committed token is processed at the next cycle. | |
| """ | |
| from __future__ import annotations | |
| import bisect | |
| import time | |
| from dataclasses import dataclass, field | |
| import torch | |
| import torch.nn.functional as F | |
| from .heads import ExitHead, entropy, head_draft_distribution | |
| from .layered import LayeredDecoder | |
| # ---------------------------------------------------------------------- policies | |
| class ExitPolicy: | |
| """Which heads may be tried, in which order, and when to exit.""" | |
| head_order: list[int] # indices into `heads`, shallow to deep | |
| max_draft: int | |
| def exit(self, k: int, logits: torch.Tensor, n_drafted: int) -> bool: | |
| raise NotImplementedError | |
| def decide(self, k: int, conf: float, n_drafted: int) -> bool: | |
| """Exit decision from the head confidence (graph engine).""" | |
| raise NotImplementedError | |
| class StaticPolicy(ExitPolicy): | |
| """Always exit at head k; draft exactly `length` tokens (LayerSkip-like).""" | |
| def __init__(self, k: int, length: int): | |
| self.head_order = [k] | |
| self.max_draft = length | |
| def exit(self, k, logits, n_drafted): | |
| return True | |
| def decide(self, k, conf, n_drafted): | |
| return True | |
| class ThresholdPolicy(ExitPolicy): | |
| """Exit at the first head (in `head_order`) whose score passes its threshold. | |
| score: "maxprob" (max over V_k ∪ ⊥, 0 if ⊥ wins), "neg_entropy", or | |
| "one_minus_reject". With a single head this is the Kangaroo rule; with all | |
| heads it is DSSD. | |
| """ | |
| def __init__(self, heads: list[ExitHead], thresholds: dict[int, float], | |
| max_draft: int = 32, score: str = "maxprob", cum_gamma: float = 0.0, | |
| acc_maps: dict | None = None): | |
| self.heads = heads | |
| # Optional per-head calibrated acceptance a_hat_k(conf) (isotonic, piecewise linear | |
| # knots/values): thresholds and the cumulative product then apply to a_hat. | |
| self.acc_maps = {int(k): (list(v[0]), list(v[1])) for k, v in (acc_maps or {}).items()} | |
| self.thresholds = thresholds | |
| self.head_order = sorted(thresholds) | |
| self.max_draft = max_draft | |
| self.score = score | |
| # Cumulative stop (g): a draft token is only emitted while the product of the | |
| # exit confidences of the block stays >= cum_gamma (proxy for the probability | |
| # that the whole block survives verification, cf. SpecDec++). | |
| self.cum_gamma = cum_gamma | |
| self._cum = 1.0 | |
| def start_block(self) -> None: | |
| self._cum = 1.0 | |
| def _cum_ok(self, conf: float, k: int | None = None) -> bool: | |
| # cum_gamma: a scalar, or a per-head dict {k: gamma_k} (F7 option B: the exit test | |
| # at head k is cum * a_hat_k >= gamma_k) | |
| g = self.cum_gamma.get(k, 0.0) if isinstance(self.cum_gamma, dict) else self.cum_gamma | |
| if g <= 0.0: | |
| self._cum *= conf | |
| return True | |
| if self._cum * conf < g: | |
| return False | |
| self._cum *= conf | |
| return True | |
| def exit(self, k, logits, n_drafted): | |
| head = self.heads[k] | |
| if self.score == "maxprob": | |
| _, _, s = head_draft_distribution(logits, head.reject, 1.0) | |
| elif self.score == "neg_entropy": | |
| s = -entropy(logits) | |
| if head.reject and int(logits.argmax(-1)) == logits.shape[-1] - 1: | |
| return False | |
| elif self.score == "one_minus_reject": | |
| _, p_rej, s = head_draft_distribution(logits, head.reject, 1.0) | |
| s = 1.0 - p_rej | |
| else: | |
| raise ValueError(self.score) | |
| return float(s) >= self.thresholds[k] and self._cum_ok(float(s), k) | |
| def _ahat(self, k, conf): | |
| if k not in self.acc_maps: | |
| return conf | |
| xs, ys = self.acc_maps[k] | |
| if conf <= xs[0]: | |
| return ys[0] | |
| if conf >= xs[-1]: | |
| return ys[-1] | |
| j = bisect.bisect_right(xs, conf) | |
| x0, x1, y0, y1 = xs[j - 1], xs[j], ys[j - 1], ys[j] | |
| return y0 + (y1 - y0) * (conf - x0) / (x1 - x0 if x1 > x0 else 1.0) | |
| def decide(self, k, conf, n_drafted): | |
| if conf <= 0.0 and self.heads[k].reject: # ⊥ won: never draft (T4: gamma_k > a_hat(⊥)) | |
| return False | |
| a = self._ahat(k, conf) | |
| return a >= self.thresholds[k] and self._cum_ok(a, k) | |
| class ScriptedPolicy(ExitPolicy): | |
| """Testing: draft token i exits at head schedule[i % len] (None = no exit).""" | |
| def __init__(self, n_heads: int, schedule: list, max_draft: int = 15): | |
| self.head_order = list(range(n_heads)) | |
| self.schedule = schedule | |
| self.max_draft = max_draft | |
| def _want(self, n_drafted): | |
| return self.schedule[n_drafted % len(self.schedule)] | |
| def exit(self, k, logits, n_drafted): | |
| return self._want(n_drafted) == k | |
| def decide(self, k, conf, n_drafted): | |
| return self._want(n_drafted) == k | |
| class AdaptiveThresholdPolicy(ThresholdPolicy): | |
| """Single head, threshold adapted online toward a target acceptance rate | |
| (Draft & Verify / AdaEDL style): after each cycle, raise the threshold if the | |
| cycle's acceptance fell below `target`, lower it otherwise.""" | |
| def __init__(self, heads, k: int, gamma0: float, target: float, step: float = 0.01, | |
| max_draft: int = 15): | |
| super().__init__(heads, {k: gamma0}, max_draft=max_draft, score="maxprob") | |
| self.k, self.target, self.step = k, target, step | |
| def update(self, n_drafted: int, n_accepted: int) -> None: | |
| if n_drafted == 0: | |
| return | |
| g = self.thresholds[self.k] + (self.step if n_accepted / n_drafted < self.target else -self.step) | |
| self.thresholds[self.k] = min(0.999, max(0.0, g)) | |
| # ---------------------------------------------------------------------- stats | |
| class GenStats: | |
| new_tokens: int = 0 | |
| cycles: int = 0 # full-model passes after prefill | |
| drafted: int = 0 | |
| accepted: int = 0 | |
| empty_drafts: int = 0 # cycles whose draft had length 0 (fallback) | |
| rejected_cycles: int = 0 # cycles ended by a rejected draft | |
| ar_steps: int = 0 # online control: plain AR steps (lambda > 1, no probe) | |
| lam_start: float = float("nan") # online control: lambda at the start / end of the request | |
| lam_end: float = float("nan") | |
| gamma_start: float = float("nan") # cost safety (NoSlowdown): gamma at the start / end | |
| gamma_end: float = float("nan") | |
| ar_step_median: float = float("nan") # AR runs: median per-token step time | |
| lab_n: int = 0 # online control: labels sent to the controller (and their sum) | |
| lab_s: float = 0.0 | |
| n_jumps: int = -1 # online control: re-anchoring jumps of the ACI (cumulative) | |
| pld_drafted: int = 0 # prompt lookup: copied drafts / accepted copied drafts | |
| pld_accepted: int = 0 | |
| cost_units: float = 0.0 # online control: sum of cycle costs (1 = one AR step) | |
| h_max: float = float("-inf") # max over drafting cycles of cost - (1 + xi) * emitted | |
| excess_end: float = 0.0 # sum over the request of cost - (1 + xi) * emitted | |
| excess_max: float = float("-inf") # running maximum of that sum within the request | |
| exits: dict = field(default_factory=dict) # head index -> count | |
| sources: list = field(default_factory=list) # per output token: copy/head/star/draft/model | |
| head_evals: int = 0 | |
| layer_rows: int = 0 | |
| kv_fill_rows: int = 0 | |
| vocab_proj_rows: int = 0 | |
| seconds: float = 0.0 | |
| prefill_seconds: float = 0.0 | |
| def tau(self) -> float: | |
| """Tokens committed per full-model pass (incl. correction/bonus).""" | |
| return self.new_tokens / max(self.cycles, 1) | |
| def as_dict(self) -> dict: | |
| d = dict(self.__dict__) | |
| d.pop("sources") # per-token list: kept out of the bench jsonl | |
| d["tau"] = self.tau | |
| d["acceptance"] = self.accepted / max(self.drafted, 1) | |
| # drafts emitted on the output prefix = accepted + one rejected per rejecting cycle | |
| d["acceptance_prefix"] = self.accepted / max(self.accepted + self.rejected_cycles, 1) | |
| d["tokens_per_s"] = self.new_tokens / max(self.seconds, 1e-9) | |
| return d | |
| def _sync(device) -> None: | |
| if device.type == "cuda": | |
| torch.cuda.synchronize(device) | |
| def _sample(probs: torch.Tensor, gen: torch.Generator | None) -> int: | |
| return int(torch.multinomial(probs, 1, generator=gen)) | |
| def _target_token(logits: torch.Tensor, temperature: float, gen) -> int: | |
| if temperature == 0.0: | |
| return int(logits.argmax()) | |
| return _sample(F.softmax(logits / temperature, dim=-1), gen) | |
| # ---------------------------------------------------------------------- decoders | |
| def generate_ar(dec: LayeredDecoder, prompt: list[int], max_new: int, | |
| temperature: float = 0.0, eos_ids: set[int] | None = None, | |
| gen: torch.Generator | None = None) -> tuple[list[int], GenStats]: | |
| eos_ids = eos_ids or set() | |
| st = GenStats() | |
| dec.reset_counters() | |
| ids = torch.tensor(prompt, device=dec.device) | |
| _sync(dec.device) | |
| t0 = time.perf_counter() | |
| logits = dec.prefill(ids) | |
| _sync(dec.device) | |
| st.prefill_seconds = time.perf_counter() - t0 | |
| out = [_target_token(logits, temperature, gen)] | |
| t1 = time.perf_counter() | |
| while len(out) < max_new and out[-1] not in eos_ids: | |
| pos = len(prompt) + len(out) - 1 | |
| logits = dec.full_step(torch.tensor([out[-1]], device=dec.device), pos)[0] | |
| out.append(_target_token(logits, temperature, gen)) | |
| st.cycles += 1 | |
| _sync(dec.device) | |
| st.seconds = time.perf_counter() - t1 + st.prefill_seconds | |
| st.new_tokens = len(out) | |
| st.layer_rows = dec.counters["layer_rows"] | |
| st.vocab_proj_rows = dec.counters["vocab_proj_rows"] | |
| return out, st | |
| def generate_spec(dec: LayeredDecoder, heads: list[ExitHead], head_layers: list[int], | |
| policy: ExitPolicy, prompt: list[int], max_new: int, | |
| temperature: float = 0.0, eos_ids: set[int] | None = None, | |
| gen: torch.Generator | None = None, | |
| kv_fill: str = "lazy", trace: list | None = None) -> tuple[list[int], GenStats]: | |
| """Self-speculative decoding; head k reads the hidden state after head_layers[k] layers. | |
| kv_fill: | |
| "lazy" regime A (reference): before a draft row climbs past the depth of | |
| earlier rows of the block, those rows are caught up exactly (batched), | |
| so every head score depends only on the exact prefix. Verification | |
| then runs only layers [c_max, L). | |
| "copy" regime B: CALM state copy of the K/V of skipped drafting layers; | |
| verification runs [c_min, L). | |
| "recompute" regime B variant: run the skipped drafting layers for exited tokens. | |
| trace: if a list, receives (position, head, logits) for every head evaluation. | |
| """ | |
| eos_ids = eos_ids or set() | |
| st = GenStats() | |
| dec.reset_counters() | |
| L = dec.spec.num_layers | |
| order = policy.head_order | |
| c_min = head_layers[order[0]] | |
| c_max = head_layers[order[-1]] | |
| ids = torch.tensor(prompt, device=dec.device) | |
| _sync(dec.device) | |
| t0 = time.perf_counter() | |
| logits = dec.prefill(ids) | |
| _sync(dec.device) | |
| st.prefill_seconds = time.perf_counter() - t0 | |
| out = [_target_token(logits, temperature, gen)] | |
| t1 = time.perf_counter() | |
| while len(out) < max_new and out[-1] not in eos_ids: | |
| pos0 = len(prompt) + len(out) - 1 | |
| inputs = [out[-1]] | |
| h_min: list[torch.Tensor] = [] # hidden after c_min layers, per row | |
| drafts: list[int] = [] | |
| q_rows: list[torch.Tensor | None] = [] | |
| budget = min(policy.max_draft, max_new - len(out)) | |
| if hasattr(policy, "start_block"): | |
| policy.start_block() | |
| if kv_fill == "lazy": | |
| drafts, q_rows, H, verify_from = _draft_lazy( | |
| dec, heads, head_layers, policy, out[-1], pos0, budget, temperature, gen, | |
| eos_ids, st, trace, context=(prompt + out[:-1]) if trace is not None else None) | |
| inputs = [out[-1]] + drafts | |
| H = dec.run_layers(H, pos0, verify_from, L) | |
| target = dec.final_logits(H) | |
| st.cycles += 1 | |
| st.drafted += len(drafts) | |
| if not drafts: | |
| st.empty_drafts += 1 | |
| n_acc, committed = _accept(drafts, q_rows, target, temperature, gen) | |
| st.accepted += n_acc | |
| for t in committed: | |
| out.append(t) | |
| if t in eos_ids or len(out) >= max_new: | |
| break | |
| continue | |
| # ---------------- drafting (regime B) | |
| while True: | |
| j = len(inputs) - 1 | |
| if len(drafts) >= budget or (drafts and drafts[-1] in eos_ids): | |
| break # last input is processed by the verification pass | |
| pos = pos0 + j | |
| h = dec.run_layers(dec.embed(torch.tensor([inputs[j]], device=dec.device)), | |
| pos, 0, c_min) | |
| h_min.append(h) | |
| depth = c_min | |
| exited = None | |
| for k in order: | |
| if head_layers[k] > depth: | |
| h = dec.run_layers(h, pos, depth, head_layers[k]) | |
| depth = head_layers[k] | |
| logits_k = heads[k](h)[0] | |
| st.head_evals += 1 | |
| if trace is not None: | |
| trace.append((pos, k, logits_k.detach().clone())) | |
| if policy.exit(k, logits_k, len(drafts)): | |
| exited = k | |
| break | |
| if exited is None: | |
| break | |
| head = heads[exited] | |
| q_k, _, _ = head_draft_distribution(logits_k, head.reject, temperature) | |
| if temperature == 0.0: | |
| tok_local = int(q_k.argmax()) | |
| q_full = None | |
| else: | |
| tok_local = _sample(q_k, gen) | |
| q_full = head.scatter_full(q_k) | |
| tok = int(head.to_token_ids(torch.tensor(tok_local, device=dec.device))) | |
| if depth < c_max: | |
| if kv_fill == "copy": | |
| dec.fill_kv_from_hidden(h, pos, depth, c_max) | |
| elif kv_fill == "recompute": | |
| dec.run_layers(h, pos, depth, c_max) | |
| else: | |
| raise ValueError(kv_fill) | |
| st.exits[exited] = st.exits.get(exited, 0) + 1 | |
| drafts.append(tok) | |
| q_rows.append(q_full) | |
| inputs.append(tok) | |
| # ---------------- verification (one full-model pass) | |
| if len(h_min) < len(inputs): # last draft token never entered the stack | |
| j = len(inputs) - 1 | |
| h_min.append(dec.run_layers( | |
| dec.embed(torch.tensor([inputs[j]], device=dec.device)), pos0 + j, 0, c_min)) | |
| H = torch.cat(h_min, dim=0) | |
| H = dec.run_layers(H, pos0, c_min, L) | |
| target = dec.final_logits(H) # (len(inputs), V) | |
| st.cycles += 1 | |
| st.drafted += len(drafts) | |
| if not drafts: | |
| st.empty_drafts += 1 | |
| # ---------------- acceptance | |
| n_acc, committed = _accept(drafts, q_rows, target, temperature, gen) | |
| st.accepted += n_acc | |
| for t in committed: | |
| out.append(t) | |
| if t in eos_ids or len(out) >= max_new: | |
| break | |
| _sync(dec.device) | |
| st.seconds = time.perf_counter() - t1 + st.prefill_seconds | |
| st.new_tokens = len(out) | |
| st.layer_rows = dec.counters["layer_rows"] | |
| st.kv_fill_rows = dec.counters["kv_fill_rows"] | |
| st.vocab_proj_rows = dec.counters["vocab_proj_rows"] | |
| return out, st | |
| def _accept(drafts, q_rows, target, temperature, gen): | |
| """Speculative acceptance (Leviathan et al.): returns (#accepted, committed tokens).""" | |
| committed: list[int] = [] | |
| for i, d in enumerate(drafts): | |
| if temperature == 0.0: | |
| if int(target[i].argmax()) != d: | |
| committed.append(int(target[i].argmax())) | |
| return len(committed) - 1, committed | |
| else: | |
| p = F.softmax(target[i] / temperature, dim=-1) | |
| q = q_rows[i] | |
| ratio = (p[d] / q[d].clamp_min(1e-20)).clamp(max=1.0) | |
| u = torch.rand((), device=p.device, generator=gen) | |
| if not bool(u < ratio): | |
| resid = (p - q).clamp_min(0.0) | |
| z = resid.sum() | |
| committed.append(_sample(resid / z if z > 0 else p, gen)) | |
| return len(committed) - 1, committed | |
| committed.append(d) | |
| n_acc = len(committed) | |
| committed.append(_target_token(target[len(drafts)], temperature, gen)) | |
| return n_acc, committed | |
| def _draft_lazy(dec, heads, head_layers, policy, first_token, pos0, budget, temperature, gen, | |
| eos_ids, st, trace, context=None): | |
| """Regime-A drafting. Returns (drafts, q_rows, hidden rows at c_max, c_max). | |
| With `trace`, records (position, head, logits, exact input context) per head evaluation. | |
| """ | |
| order = policy.head_order | |
| c_min, c_max = head_layers[order[0]], head_layers[order[-1]] | |
| rows: list[dict] = [] # {"pos", "h", "depth"} for every input row of the block | |
| def advance(target: int) -> None: | |
| while True: | |
| lag = [r for r in rows if r["depth"] < target] | |
| if not lag: | |
| return | |
| u = min(r["depth"] for r in lag) | |
| v = min([r["depth"] for r in lag if r["depth"] > u] + [target]) | |
| batch = [r for r in lag if r["depth"] == u] | |
| out = dec.run_layers(torch.cat([r["h"] for r in batch]), 0, u, v, | |
| positions=[r["pos"] for r in batch]) | |
| for i, r in enumerate(batch): | |
| r["h"], r["depth"] = out[i:i + 1], v | |
| def new_row(tok: int, pos: int) -> dict: | |
| r = {"pos": pos, "h": dec.embed(torch.tensor([tok], device=dec.device)), "depth": 0} | |
| rows.append(r) | |
| return r | |
| drafts: list[int] = [] | |
| q_rows: list = [] | |
| tok = first_token | |
| while True: | |
| row = new_row(tok, pos0 + len(drafts)) | |
| if len(drafts) >= budget or (drafts and drafts[-1] in eos_ids): | |
| break # trailing row: only verified | |
| exited = None | |
| for k in order: | |
| advance(head_layers[k]) | |
| logits_k = heads[k](row["h"])[0] | |
| st.head_evals += 1 | |
| if trace is not None: | |
| ctx = context + [first_token] + drafts | |
| trace.append((row["pos"], k, logits_k.detach().clone(), ctx)) | |
| if policy.exit(k, logits_k, len(drafts)): | |
| exited = k | |
| break | |
| if exited is None: | |
| break | |
| head = heads[exited] | |
| q_k, _, _ = head_draft_distribution(logits_k, head.reject, temperature) | |
| if temperature == 0.0: | |
| tok_local, q_full = int(q_k.argmax()), None | |
| else: | |
| tok_local = _sample(q_k, gen) | |
| q_full = head.scatter_full(q_k) | |
| tok = int(head.to_token_ids(torch.tensor(tok_local, device=dec.device))) | |
| st.exits[exited] = st.exits.get(exited, 0) + 1 | |
| drafts.append(tok) | |
| q_rows.append(q_full) | |
| if hasattr(policy, "update_draft"): | |
| policy.update_draft() | |
| advance(c_max) | |
| return drafts, q_rows, torch.cat([r["h"] for r in rows]), c_max | |