tensor_roll/model.py

# Tensor Roll — Recursive CUDA-Q Model Quantizer
# Copyright (C) 2026 SnapKitty Collective
# SPDX-License-Identifier: AGPL-3.0-or-later
"""Tiny transformer implemented from scratch in NumPy (forward + backward).

This is the empirical validation vehicle: a small teacher is trained on
synthetic data, then the ENTIRE Tensor Roll pipeline (roll -> compress ->
distill -> finetune -> quantize -> .trq -> sandboxed inference) runs on it.
Same code path as the 30B scale-up; only the config and data change.
No torch / transformers / llama.cpp anywhere in this file.
"""

from __future__ import annotations

import time
from dataclasses import dataclass

import numpy as np

DTYPE = np.float32
EPS = 1e-5


@dataclass
class Config:
    vocab: int = 32
    d: int = 48
    heads: int = 3
    layers: int = 2
    dff: int = 96
    seq: int = 16

    @property
    def dh(self) -> int:
        assert self.d % self.heads == 0
        return self.d // self.heads

    def to_dict(self):
        return {"vocab": self.vocab, "d": self.d, "heads": self.heads,
                "layers": self.layers, "dff": self.dff, "seq": self.seq}

    @staticmethod
    def from_dict(d):
        return Config(**d)


def count_params(params: dict) -> int:
    return int(sum(v.size for v in params.values()))


def init_params(rng: np.random.Generator, cfg: Config) -> dict:
    def xavier(shape):
        fan = shape[0] + shape[1] if len(shape) == 2 else shape[0]
        lim = np.sqrt(6.0 / fan)
        return rng.uniform(-lim, lim, size=shape).astype(DTYPE)

    p = {"emb": xavier((cfg.vocab, cfg.d))}
    for l in range(cfg.layers):
        p[f"L{l}.ln1g"] = np.ones(cfg.d, dtype=DTYPE)
        p[f"L{l}.ln1b"] = np.zeros(cfg.d, dtype=DTYPE)
        for nm in ("Wq", "Wk", "Wv", "Wo"):
            p[f"L{l}.{nm}"] = xavier((cfg.d, cfg.d))
            p[f"L{l}.b{nm[1:]}"] = np.zeros(cfg.d, dtype=DTYPE)
        p[f"L{l}.ln2g"] = np.ones(cfg.d, dtype=DTYPE)
        p[f"L{l}.ln2b"] = np.zeros(cfg.d, dtype=DTYPE)
        p[f"L{l}.W1"] = xavier((cfg.d, cfg.dff))
        p[f"L{l}.b1"] = np.zeros(cfg.dff, dtype=DTYPE)
        p[f"L{l}.W2"] = xavier((cfg.dff, cfg.d))
        p[f"L{l}.b2"] = np.zeros(cfg.d, dtype=DTYPE)
    p["lnfg"] = np.ones(cfg.d, dtype=DTYPE)
    p["lnfb"] = np.zeros(cfg.d, dtype=DTYPE)
    p["Whead"] = xavier((cfg.d, cfg.vocab))
    p["bhead"] = np.zeros(cfg.vocab, dtype=DTYPE)
    return p


# ---------------------------------------------------------------------------
# Forward
# ---------------------------------------------------------------------------

def _ln_forward(x, g, b):
    mean = x.mean(axis=-1, keepdims=True)
    var = x.var(axis=-1, keepdims=True)
    std = np.sqrt(var + EPS)
    xh = (x - mean) / std
    return g * xh + b, (x, xh, std, g)


def _ln_backward(dy, cache):
    x, xh, std, g = cache
    n_feat = x.shape[-1]
    dxh = dy * g
    dvar = np.sum(dxh * (x - x.mean(axis=-1, keepdims=True))
                  * -0.5 * (std ** -3), axis=-1, keepdims=True)
    dmean = (np.sum(dxh * -1.0 / std, axis=-1, keepdims=True)
             + dvar * np.mean(-2.0 * (x - x.mean(axis=-1, keepdims=True)), axis=-1, keepdims=True))
    dx = dxh / std + dvar * 2.0 * (x - x.mean(axis=-1, keepdims=True)) / n_feat + dmean / n_feat
    dg = np.sum(dy * xh, axis=(0, 1))
    db = np.sum(dy, axis=(0, 1))
    return dx, dg, db


def _softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True)
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)


def forward(params: dict, cfg: Config, ids: np.ndarray):
    """Returns (logits, cache). ids: (B,T) int."""
    B, T = ids.shape
    H, dh = cfg.heads, cfg.dh
    cache: dict = {"ids": ids}
    x = params["emb"][ids]  # (B,T,d)

    cmask = np.zeros((T, T), dtype=DTYPE)
    cmask[np.triu_indices(T, k=1)] = -1e9

    for l in range(cfg.layers):
        c: dict = {"x_prev": x}
        h, c["ln1"] = _ln_forward(x, params[f"L{l}.ln1g"], params[f"L{l}.ln1b"])
        Q = (h @ params[f"L{l}.Wq"] + params[f"L{l}.bq"]).reshape(B, T, H, dh)
        K = (h @ params[f"L{l}.Wk"] + params[f"L{l}.bk"]).reshape(B, T, H, dh)
        V = (h @ params[f"L{l}.Wv"] + params[f"L{l}.bv"]).reshape(B, T, H, dh)
        s = np.einsum("bthd,bshd->bhts", Q, K) / np.sqrt(dh) + cmask[None, None, :, :]
        P = _softmax(s)
        A = np.einsum("bhts,bshd->bthd", P, V).reshape(B, T, cfg.d)
        O = A @ params[f"L{l}.Wo"] + params[f"L{l}.bo"]
        y = x + O
        z, c["ln2"] = _ln_forward(y, params[f"L{l}.ln2g"], params[f"L{l}.ln2b"])
        u = z @ params[f"L{l}.W1"] + params[f"L{l}.b1"]
        a = np.maximum(u, 0)
        m_ = a @ params[f"L{l}.W2"] + params[f"L{l}.b2"]
        x = y + m_
        c.update({"h": h, "Q": Q, "K": K, "V": V, "P": P, "A": A,
                  "y": y, "z": z, "u": u, "a": a})
        cache[f"L{l}"] = c
    f, cache["lnf"] = _ln_forward(x, params["lnfg"], params["lnfb"])
    cache["f"] = f
    cache["x_pre_lnf"] = x
    logits = f @ params["Whead"] + params["bhead"]
    return logits, cache


def backward(params: dict, cfg: Config, cache: dict, dlogits: np.ndarray,
             d_hiddens: dict | None = None, d_attn: dict | None = None) -> dict:
    """dlogits: (B,T,V). d_hiddens: {layer_idx: (B,T,d) upstream grad added
    at that layer's output} for hidden-state matching. d_attn: {layer_idx:
    (B,H,T,T) upstream grad added to attention probs} for attention matching."""
    B, T = cache["ids"].shape
    grads = {}
    f = cache["f"]
    grads["Whead"] = np.einsum("btd,btv->dv", f, dlogits)
    grads["bhead"] = dlogits.sum(axis=(0, 1))
    df = dlogits @ params["Whead"].T
    dx, dg, db = _ln_backward(df, cache["lnf"])
    grads["lnfg"], grads["lnfb"] = dg, db

    for l in reversed(range(cfg.layers)):
        c = cache[f"L{l}"]
        if d_hiddens and l in d_hiddens:
            dx = dx + d_hiddens[l]
        # x = y + m ; m = a@W2+b2
        dy = dx.copy()
        dm = dx
        A = c["A"]
        grads[f"L{l}.W2"] = np.einsum("bti,btj->ij", c["a"], dm)
        grads[f"L{l}.b2"] = dm.sum(axis=(0, 1))
        da = dm @ params[f"L{l}.W2"].T
        du = da * (c["u"] > 0)
        grads[f"L{l}.W1"] = np.einsum("bti,btj->ij", c["z"], du)
        grads[f"L{l}.b1"] = du.sum(axis=(0, 1))
        dz = du @ params[f"L{l}.W1"].T
        dy2, dg2, db2 = _ln_backward(dz, c["ln2"])
        grads[f"L{l}.ln2g"], grads[f"L{l}.ln2b"] = dg2, db2
        dy = dy + dy2
        # y = x_prev + O ; O = A@Wo+bo
        dx_prev = dy.copy()
        dO = dy
        grads[f"L{l}.Wo"] = np.einsum("bti,btj->ij", A, dO)
        grads[f"L{l}.bo"] = dO.sum(axis=(0, 1))
        dA = (dO @ params[f"L{l}.Wo"].T).reshape(B, T, cfg.heads, cfg.dh)
        Q, K, V, P = c["Q"], c["K"], c["V"], c["P"]
        dP = np.einsum("bthd,bshd->bhts", dA, V)
        if d_attn is not None and l in d_attn:
            dP = dP + d_attn[l]
        dV = np.einsum("bhts,bthd->bshd", P, dA)
        dS = P * (dP - (dP * P).sum(axis=-1, keepdims=True))
        inv = 1.0 / np.sqrt(cfg.dh)
        dQ = np.einsum("bhts,bshd->bthd", dS, K) * inv
        dK = np.einsum("bhts,bthd->bshd", dS, Q) * inv
        h = c["h"]
        for nm, dM in (("Wq", dQ), ("Wk", dK), ("Wv", dV)):
            dM2 = dM.reshape(B, T, cfg.d)
            grads[f"L{l}.{nm}"] = np.einsum("bti,btj->ij", h, dM2)
            grads[f"L{l}.b{nm[1:]}"] = dM2.sum(axis=(0, 1))
            dx_prev = dx_prev + dM2 @ params[f"L{l}.{nm}"].T
        dh, dg1, db1 = _ln_backward(dx_prev, c["ln1"])
        grads[f"L{l}.ln1g"], grads[f"L{l}.ln1b"] = dg1, db1
        dx = dh  # grad w.r.t. this layer's input = next (previous) layer's output grad
    dEmb = np.zeros_like(params["emb"])
    np.add.at(dEmb, cache["ids"], dx)
    grads["emb"] = dEmb
    return grads


# ---------------------------------------------------------------------------
# Losses
# ---------------------------------------------------------------------------

def ce_loss(logits: np.ndarray, targets: np.ndarray):
    """Mean softmax cross-entropy. Returns (loss, dlogits)."""
    B, T, V = logits.shape
    x = logits.reshape(-1, V).astype(np.float64)
    t = targets.reshape(-1)
    x = x - x.max(axis=1, keepdims=True)
    logp = x - np.log(np.exp(x).sum(axis=1, keepdims=True))
    loss = float(-logp[np.arange(t.size), t].mean())
    d = np.exp(logp)
    d[np.arange(t.size), t] -= 1.0
    d /= t.size
    return loss, d.reshape(B, T, V).astype(DTYPE)


def kl_loss(s_logits: np.ndarray, t_logits: np.ndarray, tau: float = 2.0):
    """tau^2 * mean KL(softmax(t/tau) || softmax(s/tau)). Returns (loss, dlogits_s)."""
    B, T, V = s_logits.shape
    s = (s_logits / tau).reshape(-1, V).astype(np.float64)
    t = (t_logits / tau).reshape(-1, V).astype(np.float64)
    s = s - s.max(1, keepdims=True)
    t = t - t.max(1, keepdims=True)
    log_q = s - np.log(np.exp(s).sum(1, keepdims=True))
    log_p = t - np.log(np.exp(t).sum(1, keepdims=True))
    p = np.exp(log_p)
    kl = (p * (log_p - log_q)).sum(1).mean()
    loss = float(tau * tau * kl)
    d = (np.exp(log_q) - p) * tau / p.shape[0]
    return loss, d.reshape(B, T, V).astype(DTYPE)


def kl_divergence_mean(t_logits: np.ndarray, s_logits: np.ndarray) -> float:
    """Mean KL(teacher || student) at tau=1 — teacher-divergence metric."""
    t = t_logits.reshape(-1, t_logits.shape[-1]).astype(np.float64)
    s = s_logits.reshape(-1, s_logits.shape[-1]).astype(np.float64)
    t = t - t.max(1, keepdims=True)
    s = s - s.max(1, keepdims=True)
    log_p = t - np.log(np.exp(t).sum(1, keepdims=True))
    log_q = s - np.log(np.exp(s).sum(1, keepdims=True))
    p = np.exp(log_p)
    return float((p * (log_p - log_q)).sum(1).mean())


# ---------------------------------------------------------------------------
# Optimizer
# ---------------------------------------------------------------------------

class Adam:
    def __init__(self, params: dict, lr: float = 3e-3,
                 beta1: float = 0.9, beta2: float = 0.999, eps: float = 1e-8):
        self.lr = lr
        self.b1, self.b2, self.eps = beta1, beta2, eps
        self.m = {k: np.zeros_like(v) for k, v in params.items()}
        self.v = {k: np.zeros_like(v) for k, v in params.items()}
        self.t = 0

    def step(self, params: dict, grads: dict):
        self.t += 1
        for k in params:
            g = grads[k].astype(np.float64)
            self.m[k] = self.b1 * self.m[k] + (1 - self.b1) * g
            self.v[k] = self.b2 * self.v[k] + (1 - self.b2) * g * g
            mh = self.m[k] / (1 - self.b1 ** self.t)
            vh = self.v[k] / (1 - self.b2 ** self.t)
            params[k] = (params[k].astype(np.float64)
                         - self.lr * mh / (np.sqrt(vh) + self.eps)).astype(DTYPE)


# ---------------------------------------------------------------------------
# Synthetic data + training
# ---------------------------------------------------------------------------

def gen_data(rng: np.random.Generator, n: int, cfg: Config) -> np.ndarray:
    """Deterministic grammar: x[t] = (3*x[t-1] + 5*x[t-2] + 7) mod V."""
    ids = np.zeros((n, cfg.seq), dtype=np.int64)
    ids[:, 0] = rng.integers(0, cfg.vocab, size=n)
    ids[:, 1] = rng.integers(0, cfg.vocab, size=n)
    for t in range(2, cfg.seq):
        ids[:, t] = (3 * ids[:, t - 1] + 5 * ids[:, t - 2] + 7) % cfg.vocab
    return ids


def train_teacher(cfg: Config, *, steps: int, batch: int, lr: float,
                  seed: int, log_every: int = 100):
    rng = np.random.default_rng(seed)
    params = init_params(rng, cfg)
    opt = Adam(params, lr=lr)
    history = []
    t0 = time.perf_counter()
    for step in range(1, steps + 1):
        ids = gen_data(rng, batch, cfg)
        logits, cache = forward(params, cfg, ids[:, :-1])
        loss, dl = ce_loss(logits, ids[:, 1:])
        grads = backward(params, cfg, cache, dl)
        opt.step(params, grads)
        if step % log_every == 0 or step == steps:
            history.append({"step": step, "loss": loss,
                            "elapsed_s": time.perf_counter() - t0})
    return params, history


def eval_metrics(params: dict, cfg: Config, *, batches: int = 8,
                 batch: int = 256, seed: int = 999):
    """Real held-out accuracy + loss on the synthetic grammar."""
    rng = np.random.default_rng(seed)
    tot_loss, tot_correct, tot_n = 0.0, 0, 0
    for _ in range(batches):
        ids = gen_data(rng, batch, cfg)
        logits, _ = forward(params, cfg, ids[:, :-1])
        loss, _ = ce_loss(logits, ids[:, 1:])
        pred = logits.argmax(axis=-1)
        tot_loss += loss
        tot_correct += int((pred == ids[:, 1:]).sum())
        tot_n += batch
    return {"loss": tot_loss / batches, "accuracy": tot_correct / (tot_n * (cfg.seq - 1))}


def _act_input(name: str, cache: dict, cfg: Config):
    """Input activation feeding the op that owns tensor `name`, taken from
    a real forward cache. Returns None when not attributable."""
    if name == "emb":
        return cache["L0"]["x_prev"]
    if name in ("Whead", "bhead"):
        return cache["f"]
    if name in ("lnfg", "lnfb"):
        return cache["x_pre_lnf"]
    if name.startswith("L"):
        _, rest = name.split(".", 1)
        l = name[1:].split(".")[0]
        c = cache[f"L{l}"]
        if rest in ("Wq", "Wk", "Wv", "bq", "bk", "bv"):
            return c["h"]
        if rest in ("Wo", "bo"):
            return c["A"]
        if rest in ("W1", "b1"):
            return c["z"]
        if rest in ("W2", "b2"):
            return c["a"]
        if rest in ("ln1g", "ln1b"):
            return c["x_prev"]
        if rest in ("ln2g", "ln2b"):
            return c["y"]
    return None


def measure_tensor_stats(params: dict, cfg: Config, batches: int = 4,
                         batch: int = 64, seed: int = 0) -> dict:
    """Measured per-tensor contributions over real forward/backward passes.

    activation contribution: RMS of the input activation feeding each
        tensor's op, averaged over batches (from the forward cache).
    gradient contribution: mean Frobenius norm of the tensor's gradient,
        plus the relative ||grad|| / ||W||.
    Returns {tensor_name: {"act_rms": float, "grad_norm": float,
                           "grad_rel": float}}.
    """
    rng = np.random.default_rng(seed)
    acc_a, acc_g, n = {}, {}, 0
    for _ in range(batches):
        ids = gen_data(rng, batch, cfg)
        logits, cache = forward(params, cfg, ids[:, :-1])
        _, dlogits = ce_loss(logits, ids[:, 1:])
        grads = backward(params, cfg, cache, dlogits)
        n += 1
        for k, W in params.items():
            ai = _act_input(k, cache, cfg)
            if ai is not None:
                v = float(np.sqrt(np.mean(np.asarray(ai, dtype=np.float64) ** 2)))
                acc_a[k] = acc_a.get(k, 0.0) + v
            g = grads.get(k)
            if g is not None:
                acc_g[k] = acc_g.get(k, 0.0) + float(np.linalg.norm(g))
    stats = {}
    for k, W in params.items():
        wn = float(np.linalg.norm(W))
        gn = acc_g.get(k, 0.0) / n
        stats[k] = {
            "act_rms": acc_a.get(k, 0.0) / n if k in acc_a else None,
            "grad_norm": gn,
            "grad_rel": gn / (wn + 1e-12),
        }
    return stats


# ---------------------------------------------------------------------------
# Generation (greedy; streaming-capable)
# ---------------------------------------------------------------------------

def generate(params: dict, cfg: Config, prompt, max_tokens: int,
             seed: int = 0):
    """Greedy generation. Returns (tokens, tokens_per_sec, elapsed_s)."""
    return _generate_impl(params, cfg, prompt, max_tokens, on_token=None)


def generate_stream(params: dict, cfg: Config, prompt, max_tokens: int,
                    seed: int = 0, on_token=None):
    """Greedy generation, yielding each token as produced (true streaming).
    Returns (tokens, tokens_per_sec, elapsed_s)."""
    return _generate_impl(params, cfg, prompt, max_tokens, on_token=on_token)


def _generate_impl(params: dict, cfg: Config, prompt, max_tokens: int,
                   on_token):
    ids = np.array([list(prompt)], dtype=np.int64)
    out = []
    t0 = time.perf_counter()
    for _ in range(max_tokens):
        ctx = ids[:, -cfg.seq:]
        logits, _ = forward(params, cfg, ctx)
        nxt = int(logits[0, -1].argmax())
        out.append(nxt)
        if on_token is not None:
            on_token(nxt)
        ids = np.concatenate([ids, [[nxt]]], axis=1)
    dt = time.perf_counter() - t0
    tps = len(out) / dt if dt > 0 else 0.0
    return out, tps, dt