"""User API: load a published head set with its base model, generate with provenance. from dssd import DSSD m = DSSD.from_pretrained("/DSSD-Qwen3-8B-Base") g = m.generate("def fib(n):", max_new_tokens=256, domain="code", compare_ar=True) g.text, g.sources, g.speedup, g.identical Greedy decoding only; the output is identical to greedy autoregressive decoding of the same engine (exact verification, row-invariant computation). """ from __future__ import annotations import queue import threading import time from collections import Counter from dataclasses import dataclass, field from typing import Iterator import numpy as np import torch from .decode import ThresholdPolicy from .fast import FastDecoder from .hub import load_heads, read_config from .layered import LayeredDecoder, use_row_invariant_blas from .stops import make_stop ALIASES = {"math": "gsm8k", "humaneval": "code"} @dataclass class Generation: text: str token_ids: list[int] sources: list[str] stats: dict decode_tokens_per_s: float # (new_tokens - 1) / (seconds - prefill_seconds) ar: "Generation | None" = None speedup: float | None = None # decode_tokens_per_s / ar.decode_tokens_per_s identical: bool | None = None # token_ids == ar.token_ids @dataclass class StreamEvent: token_ids: list[int] # tokens appended by this event sources: list[str] text: str # text of the phase so far done: bool phase: str = "dssd" # "dssd" or "ar" result: Generation | None = None # on the final event t: float = 0.0 # seconds since the phase's first token (end of prefill), taken # when the engine committed these tokens class _GammaOnly: """Inner controller of a gamma-only cost safety: never asks for AR, no lambda.""" lam, probe, K = 0.0, False, 4 def policy(self): return np.zeros(self.K), 0.0, False def update(self, c): pass class _Cancelled(Exception): pass class DSSD: MAX_ROWS = 16 # fixed engine width (row-invariant GEMMs) def __init__(self, model, tokenizer, heads, layers, config: dict, graphs: bool | str = "auto", max_len: int = 4096, prefill: str = "engine"): """No engine work here (ZeroGPU: CUDA is only emulated at module level): the engine is built and its graphs captured on first use, or by warmup().""" self.model, self.tokenizer, self.config = model.eval().requires_grad_(False), tokenizer, config self.heads = [h.eval().requires_grad_(False) for h in heads] self.layers, self.graphs, self.max_len, self.prefill = list(layers), graphs, max_len, prefill self.domains = list(config["domains"]) eos = getattr(tokenizer, "eos_token_id", None) self.eos_ids = frozenset() if eos is None else frozenset([eos]) self._lock = threading.Lock() # one generation at a time on the engine buffers self._ar_step = None # seconds per AR step (cost safety), measured once self._fast, self._ready = None, False @property def fast(self) -> FastDecoder: if self._fast is None: dev = next(self.model.parameters()).device use_graphs = (dev.type == "cuda") if self.graphs == "auto" else bool(self.graphs) self._fast = FastDecoder(LayeredDecoder(self.model, max_len=self.max_len), self.heads, self.layers, max_rows=self.MAX_ROWS, use_graphs=use_graphs, prefill=self.prefill) return self._fast def _warmup(self) -> None: if not self._ready: with torch.no_grad(): self.fast.precapture([self._policy(d) for d in self.domains] + [ThresholdPolicy(self.heads, {k: 1.0 for k in range(len(self.heads))}, max_draft=15)]) self._ready = True def warmup(self) -> None: """Build the engine and capture its graphs now (otherwise done by the first generation).""" with self._lock: self._warmup() # ------------------------------------------------------------------ loading @classmethod def from_components(cls, model, tokenizer, heads, layers, config: dict, graphs: bool | str = "auto", max_len: int = 4096, prefill: str = "engine") -> "DSSD": use_row_invariant_blas() # rocBLAS on ROCm (row-invariant GEMMs); no-op elsewhere return cls(model, tokenizer, heads, layers, config, graphs, max_len, prefill) @classmethod def from_pretrained(cls, repo_or_dir: str, device: str | None = None, dtype=torch.bfloat16, graphs: bool | str = "auto", max_len: int = 4096, base_model: str | None = None, revision: str | None = None, token: str | None = None, prefill: str = "engine") -> "DSSD": """Heads from a Hub repo or directory; base model and tokenizer from the config's base_model at its pinned revision (or `base_model`, unpinned).""" from transformers import AutoModelForCausalLM, AutoTokenizer device = device or ("cuda" if torch.cuda.is_available() else "cpu") if device == "cpu": dtype = torch.float32 conf = read_config(repo_or_dir, revision, token) base = base_model or conf["base_model"] base_rev = None if base_model else conf.get("base_revision") tok = AutoTokenizer.from_pretrained(base, revision=base_rev, token=token) model = AutoModelForCausalLM.from_pretrained(base, revision=base_rev, dtype=dtype, token=token) model = model.to(device).eval() heads, layers, conf = load_heads(repo_or_dir, model.config, device, dtype, revision, token) return cls.from_components(model, tok, heads, layers, conf, graphs, max_len, prefill) # ------------------------------------------------------------------ helpers def _domain(self, domain: str | None) -> str: d = ALIASES.get(domain, domain) if domain is not None else self.config.get("default_domain", self.domains[0]) if d not in self.config["domains"]: raise ValueError(f"unknown domain {domain!r}; available: {self.domains}") return d def _policy(self, domain: str) -> ThresholdPolicy: spec = self.config["domains"][domain]["policy"] if spec.get("type") != "threshold": raise ValueError(f"unsupported policy type {spec.get('type')!r}") g = spec.get("cum_gamma", 0.0) pol = ThresholdPolicy(self.heads, {int(k): float(v) for k, v in spec["thresholds"].items()}, max_draft=spec.get("max_draft", 15), cum_gamma={int(k): float(v) for k, v in g.items()} if isinstance(g, dict) else float(g), acc_maps=spec.get("acc_maps")) pol.tree = spec.get("tree") pol.star = spec.get("star") return pol def _ids(self, prompt) -> list[int]: if not isinstance(prompt, str): return list(prompt) # as for head training and the benchmark prompts: no BOS / special tokens return list(self.tokenizer.encode(prompt, add_special_tokens=False)) def _stop(self, domain: str, stop): kind = self.config["domains"][domain].get("stop") if stop == "auto" else stop return make_stop(self.tokenizer, kind) def _check(self, ids, max_new): if len(ids) + max_new > self.max_len - self.MAX_ROWS: raise ValueError(f"prompt ({len(ids)} tokens) + max_new_tokens ({max_new}) exceeds " f"max_len={self.max_len} minus {self.MAX_ROWS} draft rows") def _result(self, out, st) -> Generation: d = st.as_dict() d["source_counts"] = dict(Counter(st.sources)) dec_s = st.seconds - st.prefill_seconds return Generation(self.tokenizer.decode(out), list(out), list(st.sources), d, (st.new_tokens - 1) / dec_s if dec_s > 0 else float("nan")) def _ar_step_seconds(self, ids) -> float: if self._ar_step is None: n = min(33, self.fast.dec.max_len - self.fast.max_rows - len(ids)) _, st = self.fast.generate_ar(ids, n) self._ar_step = st.ar_step_median return self._ar_step def _run_dssd(self, ids, max_new, domain, stop, no_slowdown, on_commit=None): from .calibration import online pol = self._policy(domain) kw = {} if no_slowdown: # cost safety (chantier 3b, NoSlowdown, gamma only): never slower than AR (+xi) ctl = online.NoSlowdown(0.02, 0.02, 0.0, _GammaOnly(), gamma_min=0.0) ctl.K, ctl.gamma_only = len(self.heads), True kw = {"controller": ctl, "ar_step_seconds": self._ar_step_seconds(ids)} out, st = self.fast.generate_spec(pol, ids, max_new, eos_ids=self.eos_ids, stop=self._stop(domain, stop), pld=self.config["domains"][domain]["policy"].get("pld"), on_commit=on_commit, **kw) return self._result(out, st) def _run_ar(self, ids, max_new, domain, stop, on_commit=None): out, st = self.fast.generate_ar(ids, max_new, eos_ids=self.eos_ids, stop=self._stop(domain, stop), on_commit=on_commit) return self._result(out, st) @staticmethod def _compare(g: Generation, ar: Generation) -> Generation: g.ar, g.identical = ar, g.token_ids == ar.token_ids g.speedup = g.decode_tokens_per_s / ar.decode_tokens_per_s if ar.decode_tokens_per_s > 0 else None return g # ------------------------------------------------------------------ API @torch.no_grad() def generate(self, prompt, max_new_tokens: int = 256, domain: str | None = None, stop: str | None = "auto", compare_ar: bool = False, no_slowdown: bool = False) -> Generation: """Greedy DSSD generation. domain: "web", "gsm8k" (alias "math") or "code" selects the calibrated policy and, with stop="auto", the stop rule. compare_ar: also run greedy AR on the same engine and report speedup (decoding, prefill excluded) and identity.""" d, ids = self._domain(domain), self._ids(prompt) self._check(ids, max_new_tokens) with self._lock: self._warmup() g = self._run_dssd(ids, max_new_tokens, d, stop, no_slowdown) if compare_ar: g = self._compare(g, self._run_ar(ids, max_new_tokens, d, stop)) return g @torch.no_grad() def generate_ar(self, prompt, max_new_tokens: int = 256, domain: str | None = None, stop: str | None = "auto") -> Generation: d, ids = self._domain(domain), self._ids(prompt) self._check(ids, max_new_tokens) with self._lock: self._warmup() return self._run_ar(ids, max_new_tokens, d, stop) def stream(self, prompt, max_new_tokens: int = 256, domain: str | None = None, stop: str | None = "auto", no_slowdown: bool = False, ar: bool = False) -> Iterator[StreamEvent]: """One event per verification cycle (DSSD), then, if ar=True, one per AR step; the final event (done=True) carries the Generation (with .ar, speedup, identical).""" d, ids = self._domain(domain), self._ids(prompt) self._check(ids, max_new_tokens) q: queue.Queue = queue.Queue() cancel = threading.Event() def sink(phase): t0 = [] def cb(toks, srcs): # runs inside the timed region: enqueue only, no decoding if cancel.is_set(): raise _Cancelled now = time.perf_counter() if not t0: t0.append(now) q.put((phase, toks, srcs, now - t0[0])) return cb def work(): try: with self._lock, torch.no_grad(): self._warmup() g = self._run_dssd(ids, max_new_tokens, d, stop, no_slowdown, sink("dssd")) if ar: g = self._compare(g, self._run_ar(ids, max_new_tokens, d, stop, sink("ar"))) q.put(StreamEvent([], [], g.text, True, "ar" if ar else "dssd", g)) except _Cancelled: pass except BaseException as e: # re-raised in the consumer q.put(e) th = threading.Thread(target=work, daemon=True) th.start() acc = {"dssd": [], "ar": []} try: while True: ev = q.get() if isinstance(ev, BaseException): raise ev if isinstance(ev, tuple): # (phase, tokens, sources, t) from the worker phase, toks, srcs, t = ev acc[phase].extend(toks) ev = StreamEvent(toks, srcs, self.tokenizer.decode(acc[phase]), False, phase, t=t) yield ev if ev.done: return finally: cancel.set() th.join()