| """ |
| 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) |
| logits_next = fg.prefill(idx) |
| new_ids: list[int] = [] |
| n_blocked = 0 |
| for _ in range(max_new_tokens): |
| logits = logits_next.float() |
| 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 |
|
|
