Blame for dssd/api.py

"""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"}


@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()