Spaces:
Sleeping
Sleeping
| """User API: load a published head set with its base model, generate with provenance. | |
| from dssd import DSSD | |
| m = DSSD.from_pretrained("<namespace>/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"} | |
| 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 | |
| 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 | |
| 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 | |
| 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) | |
| 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) | |
| 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 | |
| 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 | |
| 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() | |