| """Decoding strategies implemented on top of ``model.forward()``. |
| |
| ``model.generate()`` is never called. Each strategy drives the forward pass |
| itself and maintains its own KV cache, which is what keeps the cost linear in |
| the number of generated tokens instead of quadratic: without a cache every new |
| token would re-encode the entire prefix. |
| |
| Prompts in a batch have different lengths, so they are **left**-padded. That |
| puts every real last token at index -1, which is what the decoder needs, but |
| it also means the absolute positions are wrong unless ``position_ids`` is |
| derived from the attention mask rather than from ``arange`` -- which is what |
| ``_initial_forward`` does below. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
|
|
| |
| |
| |
| def _initial_forward(model, input_ids, attention_mask): |
| """Encode the prompt once and return (last-step logits, cache, position).""" |
| pos = attention_mask.long().cumsum(-1) - 1 |
| pos.clamp_(min=0) |
| out = model(input_ids=input_ids, attention_mask=attention_mask, |
| position_ids=pos, use_cache=True) |
| return out.logits[:, -1, :].float(), out.past_key_values, pos[:, -1] |
|
|
|
|
| def _step(model, token, attention_mask, past, pos): |
| """One incremental decoding step through the cache.""" |
| out = model(input_ids=token.unsqueeze(-1), attention_mask=attention_mask, |
| position_ids=pos.unsqueeze(-1), past_key_values=past, |
| use_cache=True) |
| return out.logits[:, -1, :].float(), out.past_key_values |
|
|
|
|
| def _grow_mask(attention_mask): |
| ones = attention_mask.new_ones((attention_mask.size(0), 1)) |
| return torch.cat([attention_mask, ones], dim=1) |
|
|
|
|
| def _finish(seqs, eos_id, prompt_len): |
| """Trim each generated sequence at its first EOS.""" |
| out = [] |
| for s in seqs: |
| if eos_id is not None and eos_id in s: |
| s = s[:s.index(eos_id)] |
| out.append(s) |
| return out |
|
|
|
|
| |
| |
| |
| @torch.no_grad() |
| def greedy_decode(model, input_ids, attention_mask, max_new=64, eos_id=None): |
| """Always take the arg-max. Deterministic, and the width-1 limit of beam |
| search *only* when no length normalisation is applied.""" |
| B = input_ids.size(0) |
| logits, past, pos = _initial_forward(model, input_ids, attention_mask) |
| gen = [[] for _ in range(B)] |
| done = torch.zeros(B, dtype=torch.bool, device=input_ids.device) |
| scores = torch.zeros(B, device=input_ids.device) |
|
|
| for _ in range(max_new): |
| lp = F.log_softmax(logits, dim=-1) |
| nxt = lp.argmax(-1) |
| scores += lp.gather(1, nxt.unsqueeze(1)).squeeze(1) * (~done) |
| for i, t in enumerate(nxt.tolist()): |
| if not done[i]: |
| gen[i].append(t) |
| if eos_id is not None: |
| done |= nxt == eos_id |
| if bool(done.all()): |
| break |
| attention_mask = _grow_mask(attention_mask) |
| pos = pos + 1 |
| logits, past = _step(model, nxt, attention_mask, past, pos) |
| return _finish(gen, eos_id, input_ids.size(1)), scores.tolist() |
|
|
|
|
| |
| |
| |
| def _top_k_filter(logits, k): |
| """Keep the k highest logits, mask the rest to -inf.""" |
| if k is None or k <= 0 or k >= logits.size(-1): |
| return logits |
| kth = logits.topk(k, dim=-1).values[..., -1, None] |
| return logits.masked_fill(logits < kth, float("-inf")) |
|
|
|
|
| def _top_p_filter(logits, p): |
| """Nucleus filtering: keep the smallest prefix of the sorted distribution |
| whose cumulative probability first reaches p. |
| |
| The shift below is what guarantees the top-1 token always survives: a |
| token is removed only if the mass *strictly before* it already reached p, |
| so even a single token with probability > p is kept. |
| """ |
| if p is None or p >= 1.0: |
| return logits |
| srt, idx = logits.sort(descending=True, dim=-1) |
| cum = srt.softmax(-1).cumsum(-1) |
| remove = cum - srt.softmax(-1) >= p |
| srt = srt.masked_fill(remove, float("-inf")) |
| return srt.scatter(-1, idx, srt) |
|
|
|
|
| @torch.no_grad() |
| def sample_decode(model, input_ids, attention_mask, max_new=64, eos_id=None, |
| top_k=None, top_p=None, temperature=1.0, generator=None): |
| """Ancestral sampling from a truncated, temperature-scaled distribution.""" |
| B = input_ids.size(0) |
| logits, past, pos = _initial_forward(model, input_ids, attention_mask) |
| gen = [[] for _ in range(B)] |
| done = torch.zeros(B, dtype=torch.bool, device=input_ids.device) |
| scores = torch.zeros(B, device=input_ids.device) |
|
|
| for _ in range(max_new): |
| full_lp = F.log_softmax(logits, dim=-1) |
| z = logits / max(temperature, 1e-6) |
| z = _top_k_filter(z, top_k) |
| z = _top_p_filter(z, top_p) |
| probs = z.softmax(-1) |
| nxt = torch.multinomial(probs, 1, generator=generator).squeeze(1) |
| scores += full_lp.gather(1, nxt.unsqueeze(1)).squeeze(1) * (~done) |
| for i, t in enumerate(nxt.tolist()): |
| if not done[i]: |
| gen[i].append(t) |
| if eos_id is not None: |
| done |= nxt == eos_id |
| if bool(done.all()): |
| break |
| attention_mask = _grow_mask(attention_mask) |
| pos = pos + 1 |
| logits, past = _step(model, nxt, attention_mask, past, pos) |
| return _finish(gen, eos_id, input_ids.size(1)), scores.tolist() |
|
|
|
|
| |
| |
| |
| @torch.no_grad() |
| def beam_search(model, input_ids, attention_mask, beam_width=4, max_new=64, |
| eos_id=None, length_penalty=1.0): |
| """Keep the ``beam_width`` highest-scoring prefixes at every step. |
| |
| Cost per step is ``O(W * |V|)`` for the scoring plus a ``top-W`` over the |
| ``W * |V|`` candidates, and the whole batch is run as one ``B*W`` forward, |
| so wall time grows roughly linearly in W (memory does too -- the KV cache |
| is W copies of the prompt). Finished hypotheses are moved aside and their |
| beam slot is neutralised with a -inf mask so it cannot be extended again. |
| |
| Hypotheses are *selected* by summed log-probability divided by |
| ``len ** length_penalty``; without that normalisation beam search |
| systematically prefers short outputs, since every extra token adds a |
| negative number. The score *returned* is the raw summed log-probability, |
| so that it means the same thing as the score the other strategies return |
| and the caller can derive a per-token figure the same way for all of them. |
| """ |
| B, L = input_ids.shape |
| dev = input_ids.device |
| W, V = beam_width, model.config.vocab_size |
|
|
| logits, past, pos = _initial_forward(model, input_ids, attention_mask) |
| lp = F.log_softmax(logits, dim=-1) |
|
|
| |
| top_lp, top_tok = lp.topk(W, dim=-1) |
| beam_scores = top_lp.reshape(B * W) |
| beams = [[int(t)] for t in top_tok.reshape(-1).tolist()] |
|
|
| past.batch_repeat_interleave(W) |
| attention_mask = attention_mask.repeat_interleave(W, dim=0) |
| pos = pos.repeat_interleave(W, dim=0) |
| cur = top_tok.reshape(-1) |
|
|
| |
| finished = [[] for _ in range(B)] |
| alive = torch.ones(B * W, dtype=torch.bool, device=dev) |
|
|
| for i in range(B * W): |
| if eos_id is not None and beams[i][0] == eos_id: |
| finished[i // W].append((beam_scores[i].item(), |
| beam_scores[i].item(), [])) |
| alive[i] = False |
|
|
| for step in range(1, max_new): |
| if not bool(alive.any()): |
| break |
| attention_mask = _grow_mask(attention_mask) |
| pos = pos + 1 |
| logits, past = _step(model, cur, attention_mask, past, pos) |
| lp = F.log_softmax(logits, dim=-1) |
|
|
| |
| |
| lp = lp.masked_fill(~alive.unsqueeze(1), float("-inf")) |
| cand = (beam_scores.unsqueeze(1) + lp).view(B, W * V) |
| top_sc, top_ix = cand.topk(W, dim=-1) |
| src_beam = top_ix // V |
| tok = top_ix % V |
|
|
| flat_src = (torch.arange(B, device=dev).unsqueeze(1) * W |
| + src_beam).reshape(-1) |
| past.reorder_cache(flat_src) |
| new_beams = [beams[i] + [int(t)] |
| for i, t in zip(flat_src.tolist(), tok.reshape(-1).tolist())] |
| beams = new_beams |
| beam_scores = top_sc.reshape(-1) |
| alive = alive[flat_src].clone() |
| cur = tok.reshape(-1) |
|
|
| if eos_id is not None: |
| hit = (cur == eos_id) & alive |
| for i in hit.nonzero(as_tuple=True)[0].tolist(): |
| seq = beams[i][:-1] |
| norm = len(seq) ** length_penalty if seq else 1.0 |
| raw = beam_scores[i].item() |
| finished[i // W].append((raw / norm, raw, seq)) |
| alive[i] = False |
| beam_scores[i] = float("-inf") |
|
|
| |
| for i in range(B * W): |
| if alive[i]: |
| seq = beams[i] |
| norm = len(seq) ** length_penalty if seq else 1.0 |
| raw = beam_scores[i].item() |
| finished[i // W].append((raw / norm, raw, seq)) |
|
|
| out, out_scores = [], [] |
| for b in range(B): |
| if finished[b]: |
| _, raw, seq = max(finished[b], key=lambda t: t[0]) |
| else: |
| raw, seq = 0.0, [] |
| out.append(seq) |
| out_scores.append(raw) |
| return out, out_scores |
|
|
