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