| """Foundry-4: sparse hybrid kernel-optimization model. |
| |
| Architecture (see DESIGN notes at bottom): |
| |
| - Embedding + KernelGram n-gram memory (hashed bigram..5-gram tables, |
| zero-init gates — proven on Foundry-2/3 to carry ~half the val loss). |
| - Hybrid trunk: GatedDeltaNet linear-attention blocks everywhere except |
| every ``attn_every``-th layer which is full causal attention (RoPE + |
| QK-norm + SDPA). Linear layers handle the stream; attention layers |
| give exact long-range recall for pointer names / block sizes. |
| - Sparse MoE FFN on every block: 1 shared expert + N fine-grained |
| routed experts, top-k sigmoid router (DeepSeekMoE-style), Switch-style |
| load-balance aux loss. |
| - PonderLoop: a recurrent block (shared weights, iteration embedding) |
| applied K times to the trunk output. Per-iteration deep supervision |
| (weighted CE on each iteration's logits) so every loop step is a |
| usable state — the substrate for latent "try many paths" reasoning |
| and later RL-driven non-lazy search. |
| - JudgmentHead (Jev-like): small MLP over the pooled latent state that |
| emits typed verdicts — p(compiles), p(correct), expected speedup, |
| halt. In pretrain it is self-supervised on a signal we already have: |
| per-iteration LM improvement. In RL it is re-pointed at real harness |
| outcomes and becomes the best-of-loop selector / idea filter. |
| - Tied 65k-token embedding + chunked CE scoring (never materialise |
| [B,T,V] logits at once), optional MTP head. |
| |
| Invariants (same discipline as Foundry-1/2 — do not break): |
| 1. Trunk is causal. Attention uses the explicit -inf mask; GDN is |
| causal by construction; the ponder block uses the same mask. |
| 2. Rewards/judgments are measured or self-supervised, never faked. |
| 3. save_checkpoint stage->verify->copy->verify on every write. |
| """ |
| from __future__ import annotations |
|
|
| import math |
| from dataclasses import dataclass |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import torch.utils.checkpoint |
|
|
| from gdn import GatedDeltaNet |
| from segmm import seg_mm, wscatter |
|
|
|
|
| |
| |
| |
|
|
| @dataclass |
| class Foundry4Config: |
| vocab_size: int = 65536 |
| n_specials: int = 16 |
| d_model: int = 1280 |
| n_layers: int = 24 |
| n_heads: int = 10 |
| head_dim: int = 128 |
| seq_len: int = 4096 |
| dropout: float = 0.0 |
| rope_theta: float = 10000.0 |
| qk_norm: bool = True |
| attn_every: int = 4 |
| |
| gdn_heads: int = 10 |
| gdn_head_k: int = 128 |
| gdn_head_v: int = 128 |
| gdn_chunk: int = 32 |
| |
| n_experts: int = 384 |
| expert_hidden: int = 128 |
| moe_topk: int = 8 |
| shared_hidden: int = 640 |
| lb_weight: float = 0.01 |
| |
| kg_buckets: int = 262144 |
| kg_dim: int = 80 |
| kg_max_n: int = 4 |
| kg_probes: int = 2 |
| |
| ponder_iters: int = 4 |
| ponder_hidden: int = 2560 |
| ponder_iter_weight_decay: float = 1.0 |
| |
| judge_hidden: int = 512 |
| |
| use_mtp: bool = True |
| mtp_weight: float = 0.05 |
| |
| moe_ckpt: bool = False |
|
|
|
|
| def beast_config() -> Foundry4Config: |
| """~5.1B total / ~570M active. The target.""" |
| return Foundry4Config() |
|
|
|
|
| def dev_config() -> Foundry4Config: |
| """~700M total / ~120M active — smoke + kernel-tuning config.""" |
| return Foundry4Config( |
| d_model=768, n_layers=12, n_heads=6, head_dim=128, |
| gdn_heads=6, gdn_head_k=128, gdn_head_v=128, |
| n_experts=32, expert_hidden=192, moe_topk=4, |
| shared_hidden=512, seq_len=2048, |
| kg_buckets=65536, kg_dim=64, kg_max_n=4, kg_probes=2, |
| ponder_iters=3, ponder_hidden=1536, |
| judge_hidden=256) |
|
|
|
|
| |
| |
| |
|
|
| class KernelGram(nn.Module): |
| _C1 = 0x9E3779B1 |
| _C2 = 0x85EBCA77 |
| _SEEDS = (0x0, 0x9E3779B1, 0x85EBCA77, 0xC2B2AE3D) |
|
|
| def __init__(self, cfg: Foundry4Config): |
| super().__init__() |
| self.cfg = cfg |
| self.n_grams = cfg.kg_max_n - 1 |
| self.n_probes = cfg.kg_probes |
| self.n_slots = self.n_grams * self.n_probes |
| self.tables = nn.ParameterList([ |
| nn.Parameter(torch.randn(cfg.kg_buckets, cfg.kg_dim) |
| / math.sqrt(cfg.kg_dim)) |
| for _ in range(self.n_slots)]) |
| self.proj = nn.ParameterList([ |
| nn.Parameter(torch.empty(cfg.kg_dim, cfg.d_model)) |
| for _ in range(self.n_slots)]) |
| for p in self.proj: |
| nn.init.normal_(p, std=0.02) |
| self.gates = nn.Parameter(torch.zeros(self.n_slots)) |
|
|
| def _hashes(self, ids: torch.Tensor) -> list[torch.Tensor]: |
| B, T = ids.shape |
| out = [] |
| for k_ in range(self.n_grams): |
| n = k_ + 2 |
| for p in range(self.n_probes): |
| h = torch.zeros(B, T, dtype=torch.long, device=ids.device) |
| if p: |
| h[:, n - 1:] = self._SEEDS[p % len(self._SEEDS)] |
| for j in range(n): |
| tok = ids[:, j:T - n + 1 + j] |
| h[:, n - 1:] = (h[:, n - 1:] * self._C1 |
| + tok * self._C2) |
| out.append(h % self.cfg.kg_buckets) |
| return out |
|
|
| def forward(self, ids, x_emb): |
| for s, (h, tab, pr) in enumerate( |
| zip(self._hashes(ids), self.tables, self.proj)): |
| g = torch.tanh(self.gates[s]) |
| x_emb = x_emb + g * (F.embedding(h, tab) @ pr) |
| return x_emb |
|
|
|
|
| |
| |
| |
| |
|
|
| class SparseMoE(nn.Module): |
| def __init__(self, cfg: Foundry4Config): |
| super().__init__() |
| d, h = cfg.d_model, cfg.expert_hidden |
| E = cfg.n_experts |
| self.cfg = cfg |
| self.E, self.k = E, cfg.moe_topk |
| self.router = nn.Linear(d, E, bias=False) |
| self.w_g = nn.Parameter(torch.empty(E, d, h)) |
| self.w_w = nn.Parameter(torch.empty(E, d, h)) |
| self.w_o = nn.Parameter(torch.empty(E, h, d)) |
| for w in (self.w_g, self.w_w, self.w_o): |
| nn.init.normal_(w, std=0.02) |
| |
| self.sh_g = nn.Linear(d, cfg.shared_hidden, bias=False) |
| self.sh_w = nn.Linear(d, cfg.shared_hidden, bias=False) |
| self.sh_o = nn.Linear(cfg.shared_hidden, d, bias=False) |
|
|
| def _swiglu(self, x, wg, ww, wo): |
| return (F.silu(x @ wg) * (x @ ww)) @ wo |
|
|
| def _dispatch(self, xf, top_i, w): |
| """Packed-segment expert compute (segmm.py Triton grouped GEMM). |
| DROPLESS (causality invariant): every routed token keeps all k |
| experts — a capacity mask would couple position i's output to |
| future tokens' routing and break KV-cache semantics. Packed |
| layout: memory stays at Nk rows regardless of routing skew, and |
| there is no host<->device sync at all.""" |
| N, d = xf.shape |
| counts = torch.bincount( |
| top_i.reshape(-1), minlength=self.E) |
| order = torch.argsort(top_i.reshape(-1)) |
| flat_t = (order // self.k) |
| flat_w = w.reshape(-1)[order] |
| offs = torch.zeros(self.E + 1, dtype=torch.int64, |
| device=xf.device) |
| offs[1:] = torch.cumsum(counts, 0) |
| xs = xf[flat_t] |
| dt = xs.dtype |
| wg = self.w_g.to(dt) |
| ww = self.w_w.to(dt) |
| wo = self.w_o.to(dt) |
| h1 = seg_mm(xs, wg, offs) |
| h2 = seg_mm(xs, ww, offs) |
| yp = seg_mm(F.silu(h1) * h2, wo, offs) |
| out = xf.new_zeros(N, d) |
| out.index_add_(0, flat_t, |
| (yp * flat_w.unsqueeze(-1)).to(out.dtype)) |
| return out |
|
|
| def forward(self, x: torch.Tensor): |
| """x [B,T,d] -> out [B,T,d], lb_loss scalar.""" |
| B, T, d = x.shape |
| xf = x.reshape(-1, d) |
| N = xf.size(0) |
|
|
| logits = self.router(xf) |
| scores = torch.sigmoid(logits) |
| top_s, top_i = scores.topk(self.k, dim=-1) |
| w = top_s / top_s.sum(-1, keepdim=True).clamp_min(1e-6) |
|
|
| if self.cfg.moe_ckpt and self.training and xf.requires_grad: |
| out = torch.utils.checkpoint.checkpoint( |
| self._dispatch, xf, top_i, w, |
| use_reentrant=False) |
| else: |
| out = self._dispatch(xf, top_i, w) |
|
|
| out = out + self.sh_o(F.silu(self.sh_g(xf)) * self.sh_w(xf)) |
| out = out.view(B, T, d) |
|
|
| |
| p = F.softmax(logits, -1) |
| f = torch.zeros_like(p).scatter_( |
| 1, top_i, 1.0 / self.k).mean(0) |
| lb = self.E * (f * p.mean(0)).sum() |
| return out, lb |
|
|
|
|
| |
| |
| |
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, d: int, eps: float = 1e-6): |
| super().__init__() |
| self.w = nn.Parameter(torch.ones(d)) |
| self.eps = eps |
|
|
| def forward(self, x): |
| return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) |
| + self.eps) * self.w |
|
|
|
|
| def _rope_cos_sin(T, hd, theta, device): |
| half = hd // 2 |
| fr = 1.0 / (theta ** (torch.arange(0, half, device=device, |
| dtype=torch.float32) / half)) |
| t = torch.arange(T, device=device, dtype=torch.float32) |
| ang = torch.outer(t, fr) |
| return ang.cos(), ang.sin() |
|
|
|
|
| def _rope(x, cos, sin): |
| half = x.size(-1) // 2 |
| x1, x2 = x[..., :half], x[..., half:] |
| c = cos[None, None] |
| s = sin[None, None] |
| return torch.cat([x1 * c - x2 * s, x1 * s + x2 * c], -1) |
|
|
|
|
| class Attention(nn.Module): |
| def __init__(self, cfg: Foundry4Config): |
| super().__init__() |
| d = cfg.d_model |
| self.h, self.hd = cfg.n_heads, cfg.head_dim |
| self.wq = nn.Linear(d, self.h * self.hd, bias=False) |
| self.wk = nn.Linear(d, self.h * self.hd, bias=False) |
| self.wv = nn.Linear(d, self.h * self.hd, bias=False) |
| self.wo = nn.Linear(self.h * self.hd, d, bias=False) |
| self.qn = RMSNorm(self.hd) if cfg.qk_norm else None |
| self.kn = RMSNorm(self.hd) if cfg.qk_norm else None |
|
|
| def forward(self, x, cos, sin, causal): |
| B, T, _ = x.shape |
| q = self.wq(x).view(B, T, self.h, self.hd).transpose(1, 2) |
| k = self.wk(x).view(B, T, self.h, self.hd).transpose(1, 2) |
| v = self.wv(x).view(B, T, self.h, self.hd).transpose(1, 2) |
| if self.qn is not None: |
| q, k = self.qn(q), self.kn(k) |
| q, k = _rope(q, cos, sin), _rope(k, cos, sin) |
| a = F.scaled_dot_product_attention(q, k, v, is_causal=True) |
| return self.wo(a.transpose(1, 2).reshape(B, T, -1)) |
|
|
| def forward_cached(self, x, cos, sin, past=None): |
| """Decode-time path. past = (k, v) post-rope from earlier calls. |
| Returns (out, (k_all, v_all)) — the returned tensors become the |
| next call's `past`. cos/sin must cover only the NEW positions.""" |
| B, T, _ = x.shape |
| q = self.wq(x).view(B, T, self.h, self.hd).transpose(1, 2) |
| k = self.wk(x).view(B, T, self.h, self.hd).transpose(1, 2) |
| v = self.wv(x).view(B, T, self.h, self.hd).transpose(1, 2) |
| if self.qn is not None: |
| q, k = self.qn(q), self.kn(k) |
| q, k = _rope(q, cos, sin), _rope(k, cos, sin) |
| if past is not None: |
| k = torch.cat([past[0], k], 2) |
| v = torch.cat([past[1], v], 2) |
| if past is None or T == k.size(2): |
| a = F.scaled_dot_product_attention(q, k, v, is_causal=True) |
| else: |
| P = k.size(2) - T |
| m = torch.ones(T, k.size(2), dtype=torch.bool, |
| device=x.device) |
| m[:, P:] = torch.tril(torch.ones(T, T, dtype=torch.bool, |
| device=x.device)) |
| a = F.scaled_dot_product_attention(q, k, v, attn_mask=m) |
| return self.wo(a.transpose(1, 2).reshape(B, T, -1)), (k, v) |
|
|
|
|
| class Block(nn.Module): |
| """mixer (GDN or attention) + sparse MoE, pre-norm.""" |
|
|
| def __init__(self, cfg: Foundry4Config, use_attn: bool): |
| super().__init__() |
| self.use_attn = use_attn |
| d = cfg.d_model |
| self.norm1 = RMSNorm(d) |
| self.norm2 = RMSNorm(d) |
| if use_attn: |
| self.mixer = Attention(cfg) |
| else: |
| self.mixer = GatedDeltaNet( |
| d, cfg.gdn_heads, cfg.gdn_head_k, cfg.gdn_head_v, |
| chunk_size=cfg.gdn_chunk) |
| self.ffn = SparseMoE(cfg) |
|
|
| def forward(self, h, cos, sin, causal): |
| if self.use_attn: |
| h = h + self.mixer(self.norm1(h), cos, sin, causal) |
| else: |
| h = h + self.mixer(self.norm1(h)) |
| ff, lb = self.ffn(self.norm2(h)) |
| h = h + ff |
| return h, lb |
|
|
| def forward_cached(self, h, cos, sin, cache): |
| """cache: dict — {'kv': (k,v)} for attn, {'S': state} for GDN.""" |
| if self.use_attn: |
| o, kv = self.mixer.forward_cached(self.norm1(h), cos, sin, |
| cache.get("kv")) |
| cache["kv"] = kv |
| h = h + o |
| else: |
| o, S, ct = self.mixer(self.norm1(h), S=cache.get("S"), |
| use_cache=True, |
| conv_prev=cache.get("conv")) |
| cache["S"] = S |
| cache["conv"] = ct |
| h = h + o |
| ff, _ = self.ffn(self.norm2(h)) |
| return h + ff |
|
|
|
|
| |
| |
| |
|
|
| class PonderBlock(nn.Module): |
| """One shared-weight recurrent step: causal attn + dense SwiGLU.""" |
|
|
| def __init__(self, cfg: Foundry4Config): |
| super().__init__() |
| d = cfg.d_model |
| self.norm1 = RMSNorm(d) |
| self.norm2 = RMSNorm(d) |
| self.attn = Attention(cfg) |
| self.ff_g = nn.Linear(d, cfg.ponder_hidden, bias=False) |
| self.ff_w = nn.Linear(d, cfg.ponder_hidden, bias=False) |
| self.ff_o = nn.Linear(cfg.ponder_hidden, d, bias=False) |
| self.iter_emb = nn.Parameter(torch.zeros(cfg.ponder_iters, d)) |
|
|
| def forward(self, h, it, cos, sin, causal): |
| h = h + self.iter_emb[it] |
| h = h + self.attn(self.norm1(h), cos, sin, causal) |
| x = self.norm2(h) |
| h = h + self.ff_o(F.silu(self.ff_g(x)) * self.ff_w(x)) |
| return h |
|
|
| def forward_cached(self, h, it, cos, sin, past): |
| o, kv = self.attn.forward_cached( |
| self.norm1(h + self.iter_emb[it]), cos, sin, past) |
| h = h + self.iter_emb[it] + o |
| x = self.norm2(h) |
| return h + self.ff_o(F.silu(self.ff_g(x)) * self.ff_w(x)), kv |
|
|
|
|
| class JudgmentHead(nn.Module): |
| """Typed verdicts over the pooled latent state. |
| |
| Outputs: p_compile, p_correct logits, log-speedup regression, |
| halt logit. Trained self-supervised in pretrain (per-iter LM |
| improvement); re-pointed at harness rewards in SFT/RL. |
| """ |
|
|
| def __init__(self, cfg: Foundry4Config): |
| super().__init__() |
| d = cfg.d_model |
| self.net = nn.Sequential( |
| nn.Linear(d, cfg.judge_hidden), nn.SiLU(), |
| nn.Linear(cfg.judge_hidden, 4)) |
|
|
| def forward(self, h): |
| pooled = h.mean(1) |
| return self.net(pooled) |
|
|
|
|
| |
| |
| |
|
|
| class Foundry4Model(nn.Module): |
| def __init__(self, cfg: Foundry4Config | None = None): |
| super().__init__() |
| self.cfg = cfg = cfg or Foundry4Config() |
| d = cfg.d_model |
| self.tok_emb = nn.Embedding(cfg.vocab_size, d) |
| self.kg = KernelGram(cfg) |
| self.blocks = nn.ModuleList([ |
| Block(cfg, use_attn=((i + 1) % cfg.attn_every == 0)) |
| for i in range(cfg.n_layers)]) |
| self.norm = RMSNorm(d) |
| self.ponder = PonderBlock(cfg) |
| self.judge = JudgmentHead(cfg) |
| self.head = nn.Linear(d, cfg.vocab_size, bias=False) |
| self.head.weight = self.tok_emb.weight |
| self.mtp = (nn.Linear(d, d, bias=False) |
| if cfg.use_mtp else None) |
| self.apply(self._init) |
|
|
| @staticmethod |
| def _init(m): |
| if isinstance(m, nn.Linear): |
| nn.init.normal_(m.weight, std=0.02) |
| elif isinstance(m, nn.Embedding): |
| nn.init.normal_(m.weight, std=0.02) |
|
|
| |
|
|
| def trunk(self, ids): |
| x = self.kg(ids, self.tok_emb(ids)) |
| dev = ids.device |
| T = ids.size(1) |
| cos, sin = _rope_cos_sin(T, self.cfg.head_dim, |
| self.cfg.rope_theta, dev) |
| h, lb = x, x.new_zeros(()) |
| for blk in self.blocks: |
| h, l = blk(h, cos, sin, None) |
| lb = lb + l |
| return self.norm(h), lb |
|
|
| |
|
|
| def forward(self, ids, return_aux=False): |
| cfg = self.cfg |
| h, lb = self.trunk(ids) |
| T = ids.size(1) |
| cos, sin = _rope_cos_sin(T, cfg.head_dim, cfg.rope_theta, |
| ids.device) |
| h_iters, judge_per_iter = [h], [] |
| for it in range(cfg.ponder_iters): |
| h = self.ponder(h, it, cos, sin, None) |
| h_iters.append(h) |
| if return_aux: |
| judge_per_iter.append(self.judge(h)) |
| if return_aux: |
| |
| |
| aux = {"lb_loss": lb * cfg.lb_weight, |
| "judge": torch.stack(judge_per_iter, 1), |
| "h_iters": h_iters} |
| if self.mtp is not None: |
| aux["mtp_h"] = self.mtp(h) |
| return None, aux |
| return self.head(h) |
|
|
| |
|
|
| def new_cache(self): |
| return {"blocks": [{} for _ in self.blocks], |
| "ponder": [None] * self.cfg.ponder_iters, |
| "tail": None} |
|
|
| def forward_cached(self, ids, pos0, cache): |
| """Decode/prefill for generation. ids [B,Tn] at absolute |
| positions pos0..pos0+Tn-1. Returns logits [B,Tn,V]; cache is |
| updated in place. Matches forward() token-for-token (causality |
| invariant), just incremental.""" |
| cfg = self.cfg |
| B, Tn = ids.shape |
| dev = ids.device |
| |
| tail = cache["tail"] |
| kg_in = ids if tail is None else torch.cat([tail, ids], 1) |
| x = self.kg(kg_in, self.tok_emb(kg_in))[:, -Tn:] |
| cache["tail"] = kg_in[:, -(cfg.kg_max_n - 1):] |
| cos, sin = _rope_cos_sin(pos0 + Tn, cfg.head_dim, |
| cfg.rope_theta, dev) |
| cos, sin = cos[pos0:], sin[pos0:] |
| h = x |
| for blk, bc in zip(self.blocks, cache["blocks"]): |
| h = blk.forward_cached(h, cos, sin, bc) |
| h = self.norm(h) |
| for it in range(cfg.ponder_iters): |
| h, cache["ponder"][it] = self.ponder.forward_cached( |
| h, it, cos, sin, cache["ponder"][it]) |
| return self.head(h) |
|
|
|
|
| |
| |
| |
|
|
| def param_report(model: Foundry4Model) -> dict: |
| cfg = model.cfg |
| total = sum(p.numel() for p in model.parameters()) |
|
|
| emb = cfg.vocab_size * cfg.d_model |
| kg = sum(p.numel() for p in model.kg.parameters()) |
| experts = cfg.n_layers * cfg.n_experts * 3 * cfg.d_model \ |
| * cfg.expert_hidden |
|
|
| |
| |
| |
| routed = cfg.n_layers * cfg.moe_topk * 3 * cfg.d_model \ |
| * cfg.expert_hidden |
| active = total - experts + routed |
| n_attn = cfg.n_layers // cfg.attn_every |
| n_gdn = cfg.n_layers - n_attn |
| return {"total": total, "active": active, |
| "experts": experts, "kg": kg, "emb": emb, |
| "n_attn": n_attn, "n_gdn": n_gdn} |
|
|
|
|
| def save_checkpoint(obj: dict, out_path) -> None: |
| """Stage-write-verify-copy-verify; survives flaky filesystems.""" |
| import io |
| import shutil |
| import tempfile |
| from pathlib import Path |
|
|
| out_path = Path(out_path) |
| out_path.parent.mkdir(parents=True, exist_ok=True) |
| buf = io.BytesIO() |
| torch.save(obj, buf) |
| data = buf.getvalue() |
| with tempfile.NamedTemporaryFile( |
| dir=out_path.parent, delete=False) as f: |
| f.write(data) |
| tmp = Path(f.name) |
| |
| torch.load(tmp, map_location="cpu", weights_only=False) |
| shutil.copy2(tmp, out_path) |
| torch.load(out_path, map_location="cpu", weights_only=False) |
| tmp.unlink() |
|
|
|
|
| |
| |
| |
|
|
| def flops_per_token(cfg: Foundry4Config) -> float: |
| """Forward+backward FLOPs per token (~6N_active + attention terms).""" |
| d = cfg.d_model |
| n_attn = cfg.n_layers // cfg.attn_every |
| n_gdn = cfg.n_layers - n_attn |
| |
| |
| mixer = n_attn * 4 * d * cfg.n_heads * cfg.head_dim \ |
| + n_gdn * d * (4 * cfg.gdn_heads * cfg.gdn_head_k) |
| ffn = cfg.n_layers * 3 * d * (cfg.shared_hidden |
| + cfg.moe_topk * cfg.expert_hidden) |
| ponder = cfg.ponder_iters * (4 * d * cfg.n_heads * cfg.head_dim |
| + 3 * d * cfg.ponder_hidden) |
| heads = cfg.vocab_size * d |
| n_active = mixer + ffn + ponder + heads |
| |
| |
| seq_term = n_attn * 12 * cfg.seq_len * d \ |
| + n_gdn * 2 * cfg.gdn_chunk * cfg.gdn_heads \ |
| * cfg.gdn_head_v |
| return 6.0 * n_active + seq_term |
|
|
