kvgen.py

"""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)