File size: 4,113 Bytes
91fca1d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | """
runner_patch.py — swap HyperExpertRunner onto FastGenerator WITHOUT editing
runtime_adapters.py. Reversible, A/B-able.
from runner_patch import install, uninstall
install(runner) # runner._generate now uses cached decode (sliding)
uninstall(runner) # back to stock
GATED: only install after divergence_quality.py says benign AND
help/shatter + smoke re-run matches native within noise.
Preserves every stock control byte-for-byte (verified identical output with
rep_pen=1.3 + no_repeat_3gram + pad masking on a real runner instance): the
loop body is stock _generate's; only the logits source changes (full forward
per step -> FastGenerator prefill/step). Deltas set before the call are
picked up automatically (snapshot-at-prefill).
One stated semantic difference: stock _generate rolls a sliding window when
len(prompt)+new exceeds block_size; a cache cannot roll, so we pre-crop the
prompt to (block_size - max_new). Your harness already budgets context so
this should never trigger; if it does, the crop is logged.
"""
import types
from time import perf_counter
from fast_generate import FastGenerator
def _generate_fast(self, prompt_text: str, max_new_tokens: int) -> str:
torch = self.torch
g = self.gen
from runtime_adapters import _banned_ngram_tokens, log
ids = self.tok.encode(prompt_text)[: self.block_size]
if len(ids) + max_new_tokens > self.block_size:
log(f"fast decode: pre-cropping prompt {len(ids)} -> "
f"{self.block_size - max_new_tokens} (cache cannot roll)",
tag="GENERATE")
ids = ids[-(self.block_size - max_new_tokens):]
idx = torch.tensor([ids], dtype=torch.long, device=self.device)
t0 = perf_counter()
fg = FastGenerator(self.adapted.model) # sliding; sees live deltas
logits_next = fg.prefill(idx)
new_ids: list[int] = []
n_blocked = 0
for _ in range(max_new_tokens):
logits = logits_next.float() # (1, V)
if g.mask_pad_vocab and self.tok_vocab and self.tok_vocab < logits.size(-1):
logits[:, self.tok_vocab:] = -float("inf")
if g.repetition_penalty and g.repetition_penalty != 1.0:
seen = torch.unique(idx[0])
vals = logits[0, seen]
logits[0, seen] = torch.where(vals > 0, vals / g.repetition_penalty,
vals * g.repetition_penalty)
if g.no_repeat_ngram_size and g.no_repeat_ngram_size >= 2:
banned = _banned_ngram_tokens(idx[0].tolist(), g.no_repeat_ngram_size)
if banned:
n_blocked += len(banned)
logits[0, list(banned)] = -float("inf")
if g.temperature <= 0:
nxt = int(logits.argmax(-1))
else:
lg = logits / g.temperature
if g.top_k:
v, _ = torch.topk(lg, min(g.top_k, lg.size(-1)))
lg[lg < v[:, [-1]]] = -float("inf")
nxt = int(torch.multinomial(torch.softmax(lg, -1), 1))
new_ids.append(nxt)
idx = torch.cat([idx, torch.tensor([[nxt]], device=self.device)], dim=1)
if g.stop_on_eos and nxt in self.stop_ids:
new_ids.pop()
break
logits_next = fg.step(torch.tensor([[nxt]], device=self.device))
extra = (f", rep_pen={g.repetition_penalty}" if g.repetition_penalty != 1.0 else "")
extra += (f", no_repeat_{g.no_repeat_ngram_size}gram(blocked {n_blocked})"
if g.no_repeat_ngram_size >= 2 else "")
log(f"generate[FAST]: prompt {len(ids)} tok -> +{len(new_ids)} tok "
f"in {perf_counter()-t0:.2f}s "
f"({'greedy' if g.temperature <= 0 else f'T={g.temperature}'}{extra})",
tag="GENERATE")
return self.tok.decode(new_ids)
def install(runner):
if not hasattr(runner, "_generate_stock"):
runner._generate_stock = runner._generate
runner._generate = types.MethodType(_generate_fast, runner)
return runner
def uninstall(runner):
if hasattr(runner, "_generate_stock"):
runner._generate = runner._generate_stock
return runner
|