Slice Desk

#!/usr/bin/env python3
"""Slice Desk — LittleBit init on one synthetic matrix.

LOCKED
- Expert up-projection reference rank is 243. Do not paste it onto this tensor.
- Init is not QAT. Do not add an optimizer.
- No Grok perplexity exists. Do not print one.
- This file is the whole loop. Do not download a checkpoint.
- Closed weights stay closed.
"""
import argparse
import json
import time
import torch


def bits(dout, din, r, paths=2):
    span = dout + din
    return paths * (r * span + 16 * (span + r))


def rank_for(bpw, dout, din, paths=2):
    n = dout * din
    span = dout + din
    exact = (bpw * n - paths * 16 * span) / (paths * (span + 16))
    floor_bpw = bits(dout, din, 1, paths) / n
    feasible = exact >= 1
    r = max(1, int(round(exact))) if feasible else 1
    return exact, r, floor_bpw, feasible


def dual_svid(W, r):
    # Truncated SVD, sign factors, rank-1 magnitude scales. Eq. 6-8.
    U, S, Vh = torch.linalg.svd(W.float(), full_matrices=False)
    r = min(int(r), int(S.numel()))
    root = S[:r].sqrt()
    Uf = U[:, :r] * root
    Vf = Vh[:r, :].T * root
    Us = torch.sign(Uf)
    Vs = torch.sign(Vf)
    Us = torch.where(Us == 0, torch.ones_like(Us), Us)
    Vs = torch.where(Vs == 0, torch.ones_like(Vs), Vs)

    def rank1(M):
        u, s, vh = torch.linalg.svd(M.abs(), full_matrices=False)
        return u[:, 0] * s[0].sqrt(), vh[0] * s[0].sqrt()

    h, lu = rank1(Uf)
    g, lv = rank1(Vf)
    return Us, Vs, h, g, (lu * lv)


def reconstruct(Us, Vs, h, g, l):
    # diag(h) Us diag(l) Vs^T diag(g), without building the diagonal matrices.
    mid = (Vs * l) @ Us.T
    return h[:, None] * mid.T * g[None, :]


def both_paths(W, r):
    pri_f = dual_svid(W, r)
    pri = reconstruct(*pri_f)
    res_f = dual_svid(W - pri, r)
    hat = pri + reconstruct(*res_f)
    return hat, pri_f, res_f


def forward(X, Us, Vs, h, g, l):
    # Y = (((X * g) @ V_sign) * l) @ U_sign.T * h
    z = (X * g) @ Vs
    z = z * l
    return (z @ Us.T) * h


def quant4(W, group=128):
    # Group absmax proxy. 16 code points stand in for NF4. Not GPTQ.
    flat = W.reshape(-1)
    pad = (group - flat.numel() % group) % group
    x = torch.nn.functional.pad(flat, (0, pad)).view(-1, group)
    scale = x.abs().amax(1, keepdim=True).clamp_min(1e-8)
    q = (x / scale * 7).round().clamp(-8, 7)
    hat = (q / 7) * scale
    return hat.reshape(-1)[: flat.numel()].reshape(W.shape)


def rel(a, b):
    return (a - b).norm() / b.norm().clamp_min(1e-12)


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--rows", type=int, default=256)
    p.add_argument("--cols", type=int, default=256)
    p.add_argument("--bpw", type=float, default=0.5)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--seq", type=int, default=64)
    args = p.parse_args()

    ref_exact, ref_r, _, _ = rank_for(0.1, 32768, 6144)
    ref = {
        "expert_up_shape": [32768, 6144],
        "expert_up_r_exact": round(ref_exact, 4),
        "expert_up_r": ref_r,
        "expert_up_mb": round(bits(32768, 6144, ref_r) / 8 / 1e6, 3),
        "expert_up_note": "Locked reference. r is 243. This run does not load that matrix.",
    }

    g = torch.Generator().manual_seed(args.seed)
    W = torch.randn(args.rows, args.cols, generator=g)
    src = "synthetic %dx%d seed %d" % (args.rows, args.cols, args.seed)

    dout, din = int(W.shape[0]), int(W.shape[1])
    exact, r, floor_bpw, ok = rank_for(args.bpw, dout, din)
    report = {"source": src, "shape": [dout, din], **ref}
    if not ok:
        report.update({
            "stopped": True,
            "reason": "target BPW is below the fp16-scale floor for this shape",
            "floor_bpw": round(floor_bpw, 4),
            "target_bpw": args.bpw,
        })
        print(json.dumps(report, indent=2))
        return

    try:
        t0 = time.perf_counter()
        hat, pri, res = both_paths(W, r)
        fit_s = time.perf_counter() - t0
    except RuntimeError as exc:
        report.update({"stopped": True, "reason": "svd failed: %s" % exc})
        print(json.dumps(report, indent=2))
        return

    q = quant4(W)
    X = torch.randn(args.seq, din)

    def timed(fn, n=8):
        for _ in range(2):
            fn()
        t = time.perf_counter()
        for _ in range(n):
            fn()
        return (time.perf_counter() - t) / n

    little = timed(lambda: forward(X, *pri) + forward(X, *res))
    dense = timed(lambda: X @ W.T)
    report.update({
        "stopped": False,
        "bpw_target": args.bpw,
        "r_exact": round(exact, 4),
        "r": r,
        "bpw_achieved": round(bits(dout, din, r) / W.numel(), 6),
        "floor_bpw": round(floor_bpw, 4),
        "littlebit_mb": round(bits(dout, din, r) / 8 / 1e6, 4),
        "fp16_mb": round(W.numel() * 2 / 1e6, 3),
        "g4_group128_mb": round(W.numel() * 4.125 / 8 / 1e6, 3),
        "littlebit_rel_l2": round(rel(hat, W).item(), 6),
        "g4_rel_l2": round(rel(q, W).item(), 6),
        "init_seconds": round(fit_s, 3),
        "seq": args.seq,
        "littlebit_forward_ms": round(little * 1e3, 2),
        "fp32_forward_ms": round(dense * 1e3, 2),
        "note": "Init only, both paths, dense sign matmul. Not QAT. Not the 11.6x kernel claim.",
    })
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()