| """Text generation. |
| |
| - sample(): naive, a full forward pass over the context per new token (no cache). |
| Used for training samples, and the reference that every faster path must match (4.1). |
| - generate(): KV-cached, one prefill pass then one new token per forward (4.3). |
| cuda_graph=True replays each decode step as one CUDA graph (4.11, tinyllm.cuda_graph). |
| Both choose tokens with sample_next (4.2). |
| """ |
|
|
| import torch |
| from torch import Tensor |
|
|
| from tinyllm.kv_cache import KVCache |
| from tinyllm.model import Transformer |
|
|
|
|
| def sample_next( |
| logits: Tensor, |
| temperature: float = 1.0, |
| top_k: int | None = None, |
| top_p: float | None = None, |
| generator: torch.Generator | None = None, |
| ) -> Tensor: |
| """Pick the next token from last-position logits (B, V) → (B, 1). Step 4.2. |
| |
| Order (same as HF's warpers): temperature, then top-k, then top-p. |
| - temperature 0 means greedy (argmax); top_k / top_p are then irrelevant. |
| - top_k keeps the k most likely tokens. |
| - top_p (nucleus) keeps the smallest set of most likely tokens whose total |
| probability is >= top_p. The most likely token is always kept. |
| """ |
| if top_k is not None and top_k < 1: |
| raise ValueError(f"top_k must be >= 1, got {top_k}") |
| if top_p is not None and not 0.0 < top_p <= 1.0: |
| raise ValueError(f"top_p must be in (0, 1], got {top_p}") |
| logits = logits.float() |
| if temperature == 0: |
| return logits.argmax(-1, keepdim=True) |
| logits = logits / temperature |
| if top_k is not None: |
| kth = torch.topk(logits, min(top_k, logits.size(-1))).values[:, -1:] |
| logits = logits.masked_fill(logits < kth, float("-inf")) |
| if top_p is not None and top_p < 1.0: |
| sorted_logits, order = logits.sort(-1, descending=True) |
| probs = sorted_logits.softmax(-1) |
| mass_before = probs.cumsum(-1) - probs |
| drop_sorted = mass_before >= top_p |
| drop = drop_sorted.scatter(-1, order, drop_sorted) |
| logits = logits.masked_fill(drop, float("-inf")) |
| return torch.multinomial(logits.softmax(-1), 1, generator=generator) |
|
|
|
|
| @torch.no_grad() |
| def sample( |
| model: Transformer, |
| idx: Tensor, |
| max_new_tokens: int, |
| temperature: float = 1.0, |
| top_k: int | None = None, |
| eos_id: int | None = None, |
| generator: torch.Generator | None = None, |
| top_p: float | None = None, |
| ) -> Tensor: |
| """Append up to max_new_tokens to idx (1, T). Stops early when eos_id is sampled. |
| |
| Token choice is sample_next (temperature, top_k, top_p). |
| """ |
| was_training = model.training |
| model.eval() |
| max_len = model.cfg.max_seq_len |
| for _ in range(max_new_tokens): |
| context = idx[:, -max_len:] |
| with torch.autocast(idx.device.type, dtype=torch.bfloat16, enabled=idx.device.type == "cuda"): |
| logits, _ = model(context) |
| nxt = sample_next(logits[:, -1, :], temperature, top_k, top_p, generator) |
| idx = torch.cat([idx, nxt], dim=1) |
| if eos_id is not None and nxt.item() == eos_id: |
| break |
| model.train(was_training) |
| return idx |
|
|
|
|
| @torch.no_grad() |
| def generate( |
| model: Transformer, |
| idx: Tensor, |
| max_new_tokens: int, |
| temperature: float = 1.0, |
| top_k: int | None = None, |
| top_p: float | None = None, |
| eos_id: int | None = None, |
| generator: torch.Generator | None = None, |
| cuda_graph=False, |
| ) -> Tensor: |
| """KV-cached generation. Append up to max_new_tokens to idx (B, T). |
| |
| Prefill runs the prompt once and fills the cache; each decode step feeds only the |
| newest token. Stops early once every row has produced eos_id (rows that finish |
| first keep generating; the caller trims after EOS). |
| |
| Unlike sample(), the context can't slide: RoPE positions are baked into the cached |
| keys. So prompt + max_new_tokens must fit in max_seq_len. |
| |
| cuda_graph: False (eager), True (build a GraphDecoder), or a GraphDecoder to reuse, |
| so the graphs it captured serve several calls with the same batch size (4.11). |
| """ |
| B, T = idx.shape |
| max_len = model.cfg.max_seq_len |
| if T + max_new_tokens > max_len: |
| raise ValueError(f"prompt {T} + new {max_new_tokens} exceeds max_seq_len {max_len}") |
| was_training = model.training |
| model.eval() |
| if cuda_graph: |
| from tinyllm.cuda_graph import GraphDecoder |
| dec = cuda_graph if isinstance(cuda_graph, GraphDecoder) else GraphDecoder(model, B) |
| prefill, decode = dec.prefill, dec.step |
| else: |
| cache = KVCache(model.cfg.n_layers, max_len) |
| autocast = dict(device_type=idx.device.type, dtype=torch.bfloat16, |
| enabled=idx.device.type == "cuda") |
|
|
| def prefill(x: Tensor) -> Tensor: |
| with torch.autocast(**autocast): |
| return model(x, cache=cache)[0] |
| decode = prefill |
| done = torch.zeros(B, dtype=torch.bool, device=idx.device) |
| out = [idx] |
| logits = prefill(idx) |
| for i in range(max_new_tokens): |
| nxt = sample_next(logits[:, -1, :], temperature, top_k, top_p, generator) |
| out.append(nxt) |
| if eos_id is not None: |
| done |= nxt.squeeze(-1) == eos_id |
| if done.all(): |
| break |
| if i + 1 == max_new_tokens: |
| break |
| logits = decode(nxt) |
| model.train(was_training) |
| return torch.cat(out, dim=1) |
|
|
