"""KV-cached batched greedy generation for ByteGPT — the wall-clock fix. The naive loop re-forwards the full prefix for every generated token (O(T·L²) attention work); harmless at tier-2's ~20-token CoT, fatal for tier-3+ chains-of-thought (hundreds-thousands of tokens) under the challenge's 5-min/1100-problem inference budget. This module reuses the model's OWN weights and norm functions — it is a faster schedule for the same computation, not a different model. Prefill captures each block's K/V via a forward hook on `attn.c_attn` (one ordinary forward over the prompt, which also yields the first generated token); each subsequent token does one single-position pass per block against the cache. Supported config (asserted): attn_kind=mha, attn_norm in {entmax15, softmax}, position_encoding in {alibi, learned, none}, no block gates / loop blocks / attention residuals. Validated token-identical against the naive loop in __main__ (random-init smoke + optional real checkpoint). """ from __future__ import annotations import sys from pathlib import Path import torch _HERE = Path(__file__).resolve().parent sys.path.insert(0, str(_HERE)) # submission layout (vendored) for _anc in Path(__file__).resolve().parents: # dev layout (repo experiments/) if (_anc / "experiments" / "coppola_pretrain_tiny.py").exists(): sys.path.insert(0, str(_anc / "experiments")) break import coppola_pretrain_tiny as cpt # noqa: E402 def _norm_row(scores: torch.Tensor, mode: str) -> torch.Tensor: """Attention norm over one query row (all cached keys are valid: no mask).""" if mode == "softmax": return torch.softmax(scores, dim=-1) if mode == "entmax15": return (cpt._pkg_entmax15(scores, dim=-1) if cpt._HAVE_ENTMAX_PKG else cpt._entmax15(scores, dim=-1)) raise ValueError(f"kvgen: unsupported attn_norm {mode!r}") def _assert_supported(model) -> None: cfg = model.config assert cfg.attn_kind == "mha", f"kvgen: attn_kind={cfg.attn_kind!r}" assert cfg.attn_norm in ("entmax15", "softmax"), f"kvgen: attn_norm={cfg.attn_norm!r}" assert cfg.position_encoding in ("alibi", "learned", "none"), \ f"kvgen: position_encoding={cfg.position_encoding!r}" assert cfg.res_attn == "none", "kvgen: res_attn unsupported" assert not cfg.loop_block_indices and not cfg.loop_unit_indices, "kvgen: loop blocks unsupported" assert model._gate_values() is None, "kvgen: block gates unsupported" assert cfg.dropout == 0.0 or not model.training, "kvgen: eval mode required" @torch.no_grad() def generate_kv(model, ids: torch.Tensor, n_gen: int) -> torch.Tensor: """Append n_gen greedy tokens to each row of ids ([B, T0] long). Returns the [B, n_gen] generated tokens. Total length must fit cfg.seq_len.""" _assert_supported(model) cfg = model.config B, T0 = ids.shape assert T0 + n_gen <= cfg.seq_len, f"kvgen: {T0}+{n_gen} exceeds seq_len {cfg.seq_len}" device = ids.device blocks = model.transformer.h L = len(blocks) nh = [blk.attn.n_head for blk in blocks] hd = [blk.attn.head_dim for blk in blocks] # ---- prefill: one ordinary forward over the prompt, K/V captured by hook kv: list[tuple[torch.Tensor, torch.Tensor] | None] = [None] * L def _mk_hook(i): def hook(_mod, _inp, qkv): _q, k, v = qkv.chunk(3, dim=-1) kv[i] = (k.view(B, -1, nh[i], hd[i]).transpose(1, 2).contiguous(), v.view(B, -1, nh[i], hd[i]).transpose(1, 2).contiguous()) return hook handles = [blk.attn.c_attn.register_forward_hook(_mk_hook(i)) for i, blk in enumerate(blocks)] try: logits, _ = model(ids) finally: for h in handles: h.remove() out = [logits[:, -1].argmax(dim=-1)] # first generated token slopes = (model._alibi_slopes.to(device) if cfg.position_encoding == "alibi" else None) # ---- incremental steps: one single-position pass per block per token for step in range(1, n_gen): t = T0 + step - 1 # position of the token we feed in x = model.transformer.wte(out[-1]).unsqueeze(1) # [B, 1, E] if cfg.position_encoding == "learned": x = x + model.transformer.wpe(torch.tensor([t], device=device))[None, :, :] for i, blk in enumerate(blocks): h = blk.ln_1(x) q, k, v = blk.attn.c_attn(h).chunk(3, dim=-1) q = q.view(B, 1, nh[i], hd[i]).transpose(1, 2) # [B, nh, 1, hd] k = k.view(B, 1, nh[i], hd[i]).transpose(1, 2) v = v.view(B, 1, nh[i], hd[i]).transpose(1, 2) K = torch.cat([kv[i][0], k], dim=2) # [B, nh, t+1, hd] V = torch.cat([kv[i][1], v], dim=2) kv[i] = (K, V) # [B, nh, 1, t+1]: dot of the single query with every cached key scores = torch.einsum("bhqd,bhjd->bhqj", q, K) / (hd[i] ** 0.5) if slopes is not None: # ALiBi row: slope * -(t - j) j = torch.arange(K.size(2), device=device, dtype=scores.dtype) scores = scores + slopes.view(1, -1, 1, 1) * (-(t - j).abs()).view(1, 1, 1, -1) attn = _norm_row(scores, cfg.attn_norm) y = torch.einsum("bhqj,bhjd->bhqd", attn, V) y = y.transpose(1, 2).contiguous().view(B, 1, -1) x = x + blk.attn.c_proj(y) x = x + blk.mlp(blk.ln_2(x)) logits1 = model.lm_head(model.transformer.ln_f(x)) # [B, 1, vocab] out.append(logits1[:, -1].argmax(dim=-1)) return torch.stack(out, dim=1) # [B, n_gen] @torch.no_grad() def generate_naive(model, ids: torch.Tensor, n_gen: int) -> torch.Tensor: """Reference loop (full re-forward per token), for validation/timing.""" cap = model.config.seq_len for _ in range(n_gen): logits, _ = model(ids[:, -cap:]) ids = torch.cat([ids, logits[:, -1].argmax(dim=-1, keepdim=True)], dim=1) return ids[:, -n_gen:] if __name__ == "__main__": import argparse import time ap = argparse.ArgumentParser(description="validate kvgen vs naive loop") ap.add_argument("--ckpt", default=None, help="optional trained checkpoint (.pt)") ap.add_argument("--n-gen", type=int, default=40) ap.add_argument("--batch", type=int, default=8) args = ap.parse_args() device = "cuda" if torch.cuda.is_available() else "cpu" if args.ckpt: sys.path.insert(0, str(_HERE)) from train_arith_bp_supervised import TrainConfig, build_model ck = torch.load(args.ckpt, map_location=device, weights_only=False) model = build_model(TrainConfig(**ck["config"]), device) model.load_state_dict(ck["state_dict"]) else: cfg = cpt.GPTConfig(vocab_size=256, n_layer=4, n_head=4, n_embd=128, seq_len=256, dropout=0.0, attn_kind="mha", attn_norm="entmax15", position_encoding="alibi", loss_kind="bce") torch.manual_seed(0) model = cpt.ByteGPT(cfg).to(device) model.eval() torch.manual_seed(1) ids = torch.randint(40, 70, (args.batch, 24), device=device) t0 = time.time() a = generate_naive(model, ids, args.n_gen) t_naive = time.time() - t0 t0 = time.time() b = generate_kv(model, ids, args.n_gen) t_kv = time.time() - t0 same = (a == b).all().item() n_diff = (a != b).sum().item() print(f"identical={same} diff_tokens={n_diff}/{a.numel()} " f"naive={t_naive:.2f}s kv={t_kv:.2f}s speedup={t_naive / max(t_kv, 1e-9):.1f}x") if not same: sys.exit(1)