spec_decode.py

"""Speculative decoding for CohereAsr on CPU: the candidate generator, the model hook, and an
oracle harness that measures the speedup-vs-acceptance curve through the *real* loop.

Why this exists
---------------
At batch 1 the decoder is 7.4% of the parameters but ~50% of the wall time, because each token is a
handful of skinny GEMVs that never saturate the CPU -- it is dispatch-bound, not compute-bound. The
consequence measured earlier: a 32-token decoder forward costs only 1.30x a 1-token forward. So if
something cheap can guess the next K tokens, one verify forward can retire up to K+1 of them for
roughly the price of one.

transformers already ships the verification half (`GenerationMixin._assisted_decoding`) and it works
on this model unmodified, including `EncoderDecoderCache.crop` for rejected tokens. Only the drafter
is missing, and a drafter is exactly a `CandidateGenerator`. Nothing here reimplements verification.

What the drafter is
-------------------
`SequenceDraftCandidateGenerator` holds a *whole-utterance* token hypothesis produced up front (a CTC
head over the encoder output emits one, non-autoregressively, in ~25 ms) and serves slices of it. The
only real work is alignment: after a rejection the generated prefix no longer matches the hypothesis
index-for-index, so the cursor is re-anchored on the last few accepted tokens.

Correctness is not at risk from a bad draft. `_assisted_decoding` keeps only the prefix its own
argmax agrees with, and the token after the last match comes from the target model itself, so a
100%-wrong draft degrades to a plain greedy step -- never to a wrong token.

Oracle mode
-----------
Passing the model's own greedy output as the hypothesis gives acceptance ~1 without training
anything, which measures the *ceiling* of this plumbing on this CPU. Corrupting a fraction of the
hypothesis walks down the acceptance axis. That replaces the analytical simulation in
results/ctc_feasibility.json with numbers from the real loop.
"""
from __future__ import annotations

import os
import sys
import json
import time
import types
import argparse

sys.path.insert(0, r"K:\cohere-to-cpu")
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
os.environ.setdefault("COHERE_CPU_THREADS", "8")

import torch
from transformers.generation.candidate_generator import CandidateGenerator

# Everything below this line is needed only by the --oracle harness at the bottom of the file, and
# is imported inside it. Keeping module scope to torch + transformers is what lets the drafter and
# the model hook be copied into the published model repo and imported there.

torch.set_grad_enabled(False)

SR = 16000
PROMPT_LEN = 10   # the processor emits a 10-token decoder prompt


# ------------------------------------------------------------------ drafter
class SequenceDraftCandidateGenerator(CandidateGenerator):
    """Serve slices of a precomputed token hypothesis as speculative candidates.

    Args:
        hypothesis: 1-D LongTensor of draft token ids, *excluding* the decoder prompt.
        confidence: optional 1-D FloatTensor, same length, per-token drafter confidence.
        k: initial draft length.
        threshold: stop a draft at the first token whose confidence is below this (0 disables).
        adaptive: +2 to k after a fully accepted round, -1 otherwise (the transformers heuristic).
        anchor: how many trailing accepted tokens to re-align on after a rejection.
    """

    requires_model_outputs = False

    def __init__(self, hypothesis, confidence=None, k=8, k_max=24, threshold=0.0,
                 adaptive=True, anchor=3, prompt_len=PROMPT_LEN,
                 eos_token_id=None, max_length=None):
        self.hyp = hypothesis.flatten().to(torch.long)
        self.conf = None if confidence is None else confidence.flatten().float()
        self.k = int(k)
        self.k_max = int(k_max)
        self.threshold = float(threshold)
        self.adaptive = bool(adaptive)
        self.anchor = int(anchor)
        self.prompt_len = int(prompt_len)
        self.eos_token_id = eos_token_id
        self.max_length = max_length
        # telemetry, read back by the harness
        self.stats = {"rounds": 0, "drafted": 0, "accepted": 0, "empty_drafts": 0,
                      "anchor_hits": 0, "backward_anchor": 0, "index_fallbacks": 0, "off_end": 0,
                      "gate_truncations": 0, "draft_lens": [], "accepted_per_round": []}
        self._last_draft_len = 0

    # -- alignment ---------------------------------------------------------
    def _cursor(self, generated: torch.Tensor) -> int | None:
        """Where in the hypothesis does the next token belong?

        Index alignment holds only while nothing has been rejected. After a rejection the target
        model has substituted its own token, so we re-anchor: find where the last `anchor` accepted
        tokens occur in the hypothesis and continue from just after them. Among several matches,
        take the one nearest the current index -- the drafter is usually locally right, and jumping
        to a far-away repeat of a common n-gram would draft nonsense.

        Forward hits are preferred over backward ones at equal distance, and a backward hit is used
        only when there is no forward one. ASR output repeats n-grams often, and a cursor behind the
        current position re-serves hypothesis slots already spent on earlier tokens: harmless for
        correctness (the verifier rejects them) but it depresses the measured acceptance rate, which
        is the x-axis of the whole speedup curve.
        """
        g = generated.tolist()
        if not g:
            return 0
        n_hyp = len(self.hyp)
        a = min(self.anchor, len(g))
        if a > 0:
            pat = g[-a:]
            hyp = self.hyp.tolist()
            hits = [i + a for i in range(0, n_hyp - a + 1) if hyp[i:i + a] == pat]
            if hits:
                self.stats["anchor_hits"] += 1
                fwd = [p for p in hits if p >= len(g)]
                if fwd:
                    return min(fwd)
                self.stats["backward_anchor"] += 1
                return max(hits)
        if len(g) < n_hyp:
            self.stats["index_fallbacks"] += 1
            return len(g)
        self.stats["off_end"] += 1
        return None

    # -- CandidateGenerator API -------------------------------------------
    def get_candidates(self, input_ids: torch.LongTensor, **kwargs):
        cur_len = input_ids.shape[1]
        if self.max_length is not None and cur_len + 1 >= self.max_length:
            # must reset like the other empty-draft paths: update_candidate_strategy runs on every
            # iteration, so a stale _last_draft_len here would be counted as drafted-but-rejected
            self.stats["empty_drafts"] += 1
            self._last_draft_len = 0
            return input_ids, None

        p = self._cursor(input_ids[0, self.prompt_len:])
        if p is None:
            self.stats["empty_drafts"] += 1
            self._last_draft_len = 0
            return input_ids, None

        end = min(p + self.k, len(self.hyp))
        if self.max_length is not None:
            end = min(end, p + (self.max_length - cur_len - 1))
        draft = self.hyp[p:end]

        # confidence gate: cut at the first token the drafter is unsure about
        if self.threshold > 0 and self.conf is not None and len(draft):
            c = self.conf[p:end]
            bad = (c < self.threshold).nonzero()
            if bad.numel():
                cut = int(bad[0])
                if cut < len(draft):
                    self.stats["gate_truncations"] += 1
                draft = draft[:cut]

        # never draft past EOS: the verifier could otherwise accept EOS *and* what follows it
        if self.eos_token_id is not None and len(draft):
            m = torch.isin(draft, self.eos_token_id).nonzero()
            if m.numel():
                draft = draft[:int(m[0])]

        if not len(draft):
            self.stats["empty_drafts"] += 1
            self._last_draft_len = 0
            return input_ids, None

        self._last_draft_len = len(draft)
        self.stats["draft_lens"].append(len(draft))
        return torch.cat((input_ids, draft.unsqueeze(0).to(input_ids.device)), dim=1), None

    def update_candidate_strategy(self, input_ids, scores, num_matches):
        self.stats["rounds"] += 1
        self.stats["drafted"] += self._last_draft_len
        self.stats["accepted"] += int(num_matches)
        self.stats["accepted_per_round"].append(int(num_matches))
        if self.adaptive and self._last_draft_len:
            if int(num_matches) == self._last_draft_len:
                self.k = min(self.k + 2, self.k_max)
            else:
                self.k = max(self.k - 1, 1)

    # -- reporting ---------------------------------------------------------
    def summary(self):
        s = self.stats
        r = max(1, s["rounds"])
        return {
            "rounds": s["rounds"],
            "drafted": s["drafted"],
            "accepted": s["accepted"],
            "acceptance_rate": round(s["accepted"] / max(1, s["drafted"]), 4),
            "mean_draft_len": round(s["drafted"] / r, 2),
            "mean_accepted_per_round": round(s["accepted"] / r, 2),
            "tokens_per_round": round((s["accepted"] + s["rounds"]) / r, 2),
            "empty_drafts": s["empty_drafts"],
            "anchor_hits": s["anchor_hits"],
            "backward_anchor": s["backward_anchor"],
            "index_fallbacks": s["index_fallbacks"],
            "off_end": s["off_end"],
            "gate_truncations": s["gate_truncations"],
        }


# --------------------------------------------------------------------- hook
def attach(model, generator):
    """Route `_assisted_decoding` to our generator.

    `_get_candidate_generator` is looked up on the instance first, so binding a replacement there
    overrides the class method without subclassing -- which matters because the deployed model is
    unpickled from the int8 cache, not constructed. `prompt_lookup_num_tokens` is still what puts
    `generate` into assisted mode; the value is irrelevant since our generator ignores it.
    """
    model._spec_generator = generator
    model._get_candidate_generator = types.MethodType(
        lambda self, **kw: self._spec_generator, model)
    return model


def detach(model):
    for a in ("_spec_generator", "_get_candidate_generator"):
        if a in model.__dict__:
            del model.__dict__[a]
    return model


def spec_generate(model, generator, *, trigger_k=8, **gen_kwargs):
    """One assisted-decoding run with `generator` as the drafter. Restores the model afterwards."""
    attach(model, generator)
    try:
        return model.generate(prompt_lookup_num_tokens=trigger_k, **gen_kwargs)
    finally:
        detach(model)


# ------------------------------------------------------------------ harness
def _clips(tr):
    from transformers.audio_utils import load_audio
    y = load_audio(r"K:\cohere-to-cpu\samples_16k\Video_Voice.wav", sampling_rate=SR)
    e = load_audio(r"K:\cohere-to-cpu\samples_16k\ElevenLabs_2026-02-22T00_18_03_Hanafi_pvc_sp100_s50_sb75_se0.wav",
                   sampling_rate=SR)
    return {
        "egyptian_10s": np.asarray(y[120 * SR:130 * SR], dtype=np.float32),
        "tts_33s": np.asarray(e, dtype=np.float32),
    }


def _encode_once(tr, clip, repeats=2):
    """Run the encoder once and time it, so the decoder can be measured on its own afterwards.

    Speculative decoding cannot touch the encoder, so folding a fresh encoder pass into every timing
    would dilute every ratio by a constant that has nothing to do with the drafter.
    """
    inp = tr.processor(audio=clip, sampling_rate=SR, return_tensors="pt", language="ar")
    f, m = inp["input_features"], inp["attention_mask"]
    if f.shape[0] != 1:
        raise SystemExit(f"clip produced {f.shape[0]} chunks; the oracle harness is single-chunk only")
    L = int(m[0].sum())
    f, m = f[:1, :L].to(tr.dtype), m[:1, :L]
    # min over repeats, like the decoder timings. A single cold reading here would be inflated by
    # oneDNN's first-call algorithm selection and would drag e2e_speedup toward 1.0, because the
    # encoder time sits in both the numerator and the denominator of that ratio.
    ts = []
    for _ in range(max(1, repeats)):
        t = time.perf_counter()
        enc = tr.model.model.encoder(input_features=f, attention_mask=m)
        ts.append(time.perf_counter() - t)
    return f, m, enc, min(ts)


def _corrupt(ids, frac, vocab, seed):
    """Replace `frac` of the hypothesis with random tokens: a drafter that is right (1-frac) of the
    time, in the worst way -- errors scattered uniformly rather than clustered at the end."""
    if frac <= 0:
        return ids.clone()
    g = torch.Generator().manual_seed(seed)
    out = ids.clone()
    n = max(1, int(round(frac * len(ids))))
    pos = torch.randperm(len(ids), generator=g)[:n]
    out[pos] = torch.randint(0, vocab, (n,), generator=g, dtype=out.dtype)
    return out


def cmd_oracle(a):
    import numpy as np                     # noqa: F401  (used by _clips/_corrupt)
    from cohere_cpu import Transcriber
    globals().update(np=np)
    torch.set_num_threads(a.threads)
    tr = Transcriber(quant=a.quant, batch_size=1, verbose=False)
    tok = tr.processor.tokenizer
    vocab = tr.model.config.get_text_config().vocab_size
    eos = torch.tensor([tr.model.generation_config.eos_token_id]).flatten()

    out = {"quant": a.quant, "threads": a.threads, "ks": a.k, "corruptions": a.corrupt, "clips": {}}
    print(f"=== oracle harness | quant={a.quant} threads={a.threads} ===", flush=True)
    print("  warming up (oneDNN picks its algorithms on the first call)...", flush=True)
    tr.warmup()

    for name, clip in _clips(tr).items():
        f, m, enc, t_enc = _encode_once(tr, clip, repeats=a.repeats)
        base = dict(input_features=f, attention_mask=m,
                    decoder_input_ids=tr._prompt_ids("ar", True, 1),
                    max_new_tokens=a.max_new_tokens, num_beams=1, do_sample=False, use_cache=True)
        # Supplying encoder_outputs makes generate skip _prepare_encoder_decoder_kwargs_for_
        # generation entirely, so the encoder runs once here and never again. The object is
        # passed through as-is: CohereAsrModel.forward reads .last_hidden_state and
        # .attention_mask off it, and only the encoder can produce that subsampled mask.
        cached = dict(base, encoder_outputs=enc)

        # -- greedy reference, decoder only
        ts = []
        for _ in range(a.repeats):
            t = time.perf_counter()
            g = tr.model.generate(**cached)
            ts.append(time.perf_counter() - t)
        t_greedy = min(ts)
        n_tok = g.shape[1] - PROMPT_LEN
        hyp_full = g[0, PROMPT_LEN:].clone()
        gtext = tok.decode(g[0], skip_special_tokens=True)

        rec = {"encoder_s": round(t_enc, 3), "greedy_decoder_s": round(t_greedy, 3),
               "greedy_total_s": round(t_enc + t_greedy, 3), "greedy_tokens": n_tok,
               "ms_per_token": round(1000 * t_greedy / max(1, n_tok), 2), "runs": {}}
        print(f"\n[{name}] encoder {t_enc*1000:.0f} ms | greedy decoder {t_greedy*1000:.0f} ms "
              f"for {n_tok} tok ({1000*t_greedy/max(1,n_tok):.1f} ms/tok) | total {(t_enc+t_greedy):.2f} s",
              flush=True)
        print(f"    {'draft':<14}{'K':>3}{'dec s':>8}{'tot s':>8}{'dec x':>7}{'e2e x':>7}"
              f"{'acc':>7}{'tok/rnd':>9}{'exact':>7}{'nWER':>7}", flush=True)

        for frac in a.corrupt:
            hyp = _corrupt(hyp_full, frac, vocab, seed=1234)
            for K in a.k:
                gen = SequenceDraftCandidateGenerator(
                    hyp, k=K, k_max=a.k_max, adaptive=not a.no_adaptive,
                    anchor=a.anchor, eos_token_id=eos,
                    max_length=PROMPT_LEN + a.max_new_tokens)
                ts = []
                for i in range(a.repeats):
                    if i:  # fresh generator per repeat, else adaptive-K state carries over
                        gen = SequenceDraftCandidateGenerator(
                            hyp, k=K, k_max=a.k_max, adaptive=not a.no_adaptive,
                            anchor=a.anchor, eos_token_id=eos,
                            max_length=PROMPT_LEN + a.max_new_tokens)
                    t = time.perf_counter()
                    s = spec_generate(tr.model, gen, trigger_k=K, **cached)
                    ts.append(time.perf_counter() - t)
                t_spec = min(ts)
                exact = bool(torch.equal(s, g))
                stext = tok.decode(s[0], skip_special_tokens=True)
                w = 0.0 if exact else 100.0 * edit(n_ortho(gtext).split(), n_ortho(stext).split()) \
                    / max(1, len(n_ortho(gtext).split()))
                sm = gen.summary()
                key = f"corrupt{frac:g}_k{K}"
                rec["runs"][key] = dict(
                    sm, corrupt=frac, k=K, decoder_s=round(t_spec, 3),
                    total_s=round(t_enc + t_spec, 3),
                    decoder_speedup=round(t_greedy / t_spec, 3),
                    e2e_speedup=round((t_enc + t_greedy) / (t_enc + t_spec), 3),
                    tokens=s.shape[1] - PROMPT_LEN, exact=exact, norm_wer_vs_greedy=round(w, 3))
                print(f"    {f'{int(frac*100)}% wrong':<14}{K:>3}{t_spec:>8.2f}{t_enc+t_spec:>8.2f}"
                      f"{t_greedy/t_spec:>7.2f}{(t_enc+t_greedy)/(t_enc+t_spec):>7.2f}"
                      f"{sm['acceptance_rate']:>7.2f}{sm['tokens_per_round']:>9.2f}"
                      f"{str(exact):>7}{w:>6.1f}%", flush=True)

        out["clips"][name] = rec

    p = a.out or r"K:\cohere-to-cpu\results\spec_oracle.json"
    json.dump(out, open(p, "w", encoding="utf-8"), ensure_ascii=False, indent=2)
    print(f"\nwrote {p}")

    # ---------------------------------------------------------------- the Phase 2 gate
    # A speedup is only meaningful against a run that produced the same transcript. Under dynamic
    # int8 a wider verify window quantizes activations more coarsely (see spec_exactness), so a
    # speculative run can diverge from greedy and stop early -- which is FASTER, and would read as
    # a pass. Eligibility is normalized-output equivalence, the same criterion the project quotes
    # WER at; byte-identity is reported but not required.
    runs = out["clips"]["egyptian_10s"]["runs"]
    perfect = {k: v for k, v in runs.items() if v["corrupt"] == 0.0}
    eligible = {k: v for k, v in perfect.items() if v["norm_wer_vs_greedy"] == 0.0}
    excluded = {k: v for k, v in perfect.items() if v["norm_wer_vs_greedy"] != 0.0}
    print()
    print("=" * 74)
    for k, v in excluded.items():
        print(f"  excluded (diverged, not comparable): {k} -> {v['e2e_speedup']:.2f}x but "
              f"normWER {v['norm_wer_vs_greedy']:.2f}%, {v['tokens']} tokens vs greedy "
              f"{out['clips']['egyptian_10s']['greedy_tokens']}")
    if not eligible:
        print("GATE: INCONCLUSIVE - every perfect-draft K diverged from greedy on the 10 s clip.")
        print("  Re-run with smaller K, or read the excluded rows as an upper bound.")
        return
    bk, bv = max(eligible.items(), key=lambda kv: kv[1]["e2e_speedup"])
    best = bv["e2e_speedup"]
    print(f"GATE: best end-to-end speedup with a perfect draft on the 10 s clip = {best:.2f}x "
          f"({bk}, decoder {bv['decoder_speedup']:.2f}x, byte-exact={bv['exact']})")
    print("      need >= 1.5x to justify training a drafter")
    print("  -> PASS" if best >= 1.5 else "  -> FAIL: stop before Phase 3")


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--oracle", action="store_true", help="run the oracle speedup-vs-acceptance sweep")
    ap.add_argument("--k", type=int, nargs="+", default=[4, 8, 16])
    ap.add_argument("--k-max", type=int, default=24)
    ap.add_argument("--corrupt", type=float, nargs="+", default=[0.0, 0.1, 0.3])
    ap.add_argument("--anchor", type=int, default=3)
    ap.add_argument("--no-adaptive", action="store_true")
    ap.add_argument("--quant", default="int8", choices=["int8", "fp32", "bf16"])
    ap.add_argument("--threads", type=int, default=8)
    ap.add_argument("--repeats", type=int, default=2, help="take the min, to reject scheduler noise")
    ap.add_argument("--max-new-tokens", type=int, default=256)
    ap.add_argument("--out")
    a = ap.parse_args(argv)
    if not a.oracle:
        ap.error("nothing to do; pass --oracle (the drafter classes are meant to be imported)")
    cmd_oracle(a)


if __name__ == "__main__":
    main()