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