ml/model.py

"""Shared tiny Transformer backbone.

One class, two modes. `causal=False` → bidirectional attention, the diffusion
denoiser (all positions see all positions, no KV cache; DESIGN.md §1).
`causal=True` → causal mask, the AR baseline, WITH a KV cache for generation —
AR's natural advantage, which a fair iso-latency comparison must grant it.
Diffusion inherently has no cache: each denoising step reprocesses the whole
sequence (EVALUATION.md is explicit that this makes a diffusion step heavier).

Same depth/width/heads in both modes, so the comparison isolates the attention
pattern + objective, not capacity. No recurrent depth, no MoE — Stage 0 is the
bare substrate (EXPERIMENTS.md).
"""

from __future__ import annotations

from contextlib import nullcontext

import torch
import torch.nn as nn
import torch.nn.functional as F

from .config import ModelConfig


def amp_ctx(device: str):
    """fp16 autocast on GPU backends — ~5x faster attention on MPS, big memory
    savings. CPU stays fp32."""
    if device in ("mps", "cuda"):
        return torch.autocast(device_type=device, dtype=torch.float16)
    return nullcontext()


class MHA(nn.Module):
    def __init__(self, cfg: ModelConfig):
        super().__init__()
        self.h = cfg.n_heads
        self.hd = cfg.d_model // cfg.n_heads
        self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model)
        self.proj = nn.Linear(cfg.d_model, cfg.d_model)
        self.dropout = cfg.dropout

    def forward(self, x, attn_bias=None, cache=None, return_kv=False):
        """x: (B,T,C). attn_bias: precomputed additive mask (B,1,T,Tk) or None
        (built once per forward by the Transformer). cache: (k,v) past or None.
        Returns (out, kv_or_None)."""
        B, T, C = x.shape
        qkv = self.qkv(x).view(B, T, 3, self.h, self.hd).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]  # each (B,h,T,hd)

        incremental = cache is not None
        if incremental:
            pk, pv = cache
            if pk is not None:
                k = torch.cat([pk, k], dim=2)
                v = torch.cat([pv, v], dim=2)
        new_kv = (k, v) if (return_kv or incremental) else None

        # Incremental single-step decode attends all cached keys → no mask needed
        # (causal-correct because tokens are fed in order).
        mask = None if incremental else attn_bias
        if mask is not None and mask.dtype != q.dtype:
            mask = mask.to(q.dtype)

        out = F.scaled_dot_product_attention(
            q, k, v,
            attn_mask=mask,
            dropout_p=self.dropout if self.training else 0.0,
        )
        out = out.transpose(1, 2).reshape(B, T, C)
        return self.proj(out), new_kv


class Block(nn.Module):
    def __init__(self, cfg: ModelConfig):
        super().__init__()
        self.ln1 = nn.LayerNorm(cfg.d_model)
        self.attn = MHA(cfg)
        self.ln2 = nn.LayerNorm(cfg.d_model)
        self.mlp = nn.Sequential(
            nn.Linear(cfg.d_model, cfg.d_ff),
            nn.GELU(),
            nn.Linear(cfg.d_ff, cfg.d_model),
            nn.Dropout(cfg.dropout),
        )

    def forward(self, x, attn_bias=None, cache=None, return_kv=False):
        a, kv = self.attn(self.ln1(x), attn_bias, cache, return_kv)
        x = x + a
        x = x + self.mlp(self.ln2(x))
        return x, kv


class Transformer(nn.Module):
    def __init__(self, cfg: ModelConfig, causal: bool):
        super().__init__()
        self.cfg = cfg
        self.causal = causal
        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.pos_emb = nn.Embedding(cfg.max_len, cfg.d_model)
        self.drop = nn.Dropout(cfg.dropout)
        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layers)])
        self.ln_f = nn.LayerNorm(cfg.d_model)
        self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        self.head.weight = self.tok_emb.weight  # weight tying
        self.apply(self._init)

    def _init(self, m):
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, std=0.02)
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, std=0.02)

    def _build_bias(self, attn_keep, causal):
        """Additive attention mask (B,1,T,T), built ONCE per forward. -inf where a
        query may not attend (padding, and future positions when causal)."""
        B, T = attn_keep.shape
        device = attn_keep.device
        bias = torch.zeros(B, 1, T, T, device=device)
        bias = bias.masked_fill(~attn_keep[:, None, None, :], float("-inf"))
        if causal:
            cmask = torch.triu(torch.ones(T, T, device=device, dtype=torch.bool), 1)
            bias = bias.masked_fill(cmask, float("-inf"))
        return bias

    def forward(self, ids, attn_keep):
        """Full pass. ids:(B,T); attn_keep:(B,T) True=real token. Returns (B,T,V)."""
        B, T = ids.shape
        pos = torch.arange(T, device=ids.device).unsqueeze(0)
        x = self.drop(self.tok_emb(ids) + self.pos_emb(pos))
        bias = self._build_bias(attn_keep, self.causal)
        for blk in self.blocks:
            x, _ = blk(x, attn_bias=bias, cache=None)
        return self.head(self.ln_f(x))

    # ---- block diffusion (Nemotron-style, KV-cacheable) ----

    def _block_bias(self, attn_keep, region, block_id):
        """Block-causal additive mask (B,1,T,T) for training/teacher-forcing.
        Clean context attends only to context (so its representation is
        independent of the generated region -> cacheable). A region query in
        block b attends to all context plus region blocks <= b, bidirectional
        within its block."""
        keepj = attn_keep[:, None, None, :]            # (B,1,1,T)
        reg_i = region[:, None, :, None]               # query is region
        reg_j = region[:, None, None, :]               # key is region
        bid_i = block_id[:, None, :, None]
        bid_j = block_id[:, None, None, :]
        allow_region_q = (~reg_j) | (reg_j & (bid_j <= bid_i))
        allow = torch.where(reg_i, allow_region_q, ~reg_j) & keepj
        return torch.zeros_like(allow, dtype=torch.float32).masked_fill(~allow, float("-inf"))

    def forward_blocks(self, ids, attn_keep, region, block_id):
        """Full teacher-forcing pass under the block-causal mask. Returns (B,T,V)."""
        B, T = ids.shape
        pos = torch.arange(T, device=ids.device).unsqueeze(0)
        x = self.drop(self.tok_emb(ids) + self.pos_emb(pos))
        bias = self._block_bias(attn_keep, region, block_id)
        for blk in self.blocks:
            x, _ = blk(x, attn_bias=bias, cache=None)
        return self.head(self.ln_f(x))

    @torch.no_grad()
    def encode_context(self, ctx_ids, ctx_pos):
        """Per-layer K/V for the clean context (prefix+suffix), attending only
        among itself. ctx_ids,ctx_pos: (1,Lc). Returns list of (k,v)."""
        x = self.tok_emb(ctx_ids) + self.pos_emb(ctx_pos)
        caches = []
        for blk in self.blocks:
            x, kv = blk(x, attn_bias=None, cache=None, return_kv=True)
            caches.append(kv)
        return caches

    @torch.no_grad()
    def decode_block(self, blk_ids, blk_pos, caches):
        """Forward one block's positions against cached context+committed K/V.
        blk_ids,blk_pos: (1,Lb). Returns (logits (1,Lb,V), new_caches) where
        new_caches has the block's K/V appended (use after committing)."""
        x = self.tok_emb(blk_ids) + self.pos_emb(blk_pos)
        new_caches = []
        for blk, c in zip(self.blocks, caches):
            x, kv = blk(x, attn_bias=None, cache=c, return_kv=True)
            new_caches.append(kv)
        return self.head(self.ln_f(x)), new_caches

    @torch.no_grad()
    def generate(self, head_ids, max_new, eos_id):
        """KV-cached greedy decode (causal only). head_ids: 1-D prompt tensor.
        Returns the list of generated token ids (excluding eos)."""
        assert self.causal
        device = head_ids.device
        L = head_ids.size(0)
        pos = torch.arange(L, device=device).unsqueeze(0)
        x = self.tok_emb(head_ids.unsqueeze(0)) + self.pos_emb(pos)

        # Prompt pass: causal among the prompt, seed the caches.
        cbias = torch.triu(
            torch.full((1, 1, L, L), float("-inf"), device=device), diagonal=1
        )
        caches = []
        for blk in self.blocks:
            x, kv = blk(x, attn_bias=cbias, cache=None, return_kv=True)
            caches.append(kv)
        logits = self.head(self.ln_f(x))
        nxt = int(logits[0, -1].argmax().item())

        out = []
        cur_len = L
        for _ in range(max_new):
            if nxt == eos_id:
                break
            out.append(nxt)
            cur_len += 1
            if cur_len >= self.cfg.max_len:
                break
            tok = torch.tensor([[nxt]], device=device)
            p = torch.tensor([[cur_len - 1]], device=device)
            x = self.tok_emb(tok) + self.pos_emb(p)
            new_caches = []
            for blk, c in zip(self.blocks, caches):
                x, kv = blk(x, attn_bias=None, cache=c, return_kv=True)
                new_caches.append(kv)
            caches = new_caches
            logits = self.head(self.ln_f(x))
            nxt = int(logits[0, -1].argmax().item())
        return out

    def num_params(self) -> int:
        return sum(p.numel() for p in self.parameters())