import json import sys import numpy as np sys.path.insert(0, "/tmp/opencode/tl") import export_full as E P, Q, Sc, W, SC = E.build() DIM, HID, NL, NH, HS, KVD = 64, 256, 8, 16, 4, 64 KV = json.load(open("/tmp/opencode/tl/kvmax.json")) KMAX = np.array(KV["kmax"], np.int64) VMAX = np.array(KV["vmax"], np.int64) ATT_C = 66572 PKSH = 0 SILU_SH = 11 SILU_RND = False ROPE = {} LAST = {} STATS = {"ca": 0, "prod": 0, "quant": 0, "x": 0} FLOAT_ATT = False FLOAT_SILU = False FLOAT_ROPE = False FLOAT_KV = False def build_tables(): global P8, SILU, SRND, SG8 pk = 32 << PKSH P8 = np.array([round(127 * np.exp(-i / pk)) for i in range(8 * pk)], np.int64) g = 1 << (16 - SILU_SH) SG8 = 8 * g SRND = (1 << (SILU_SH - 1)) if SILU_RND else 0 SILU = np.array([round((x := (i - SG8) / g) / (1 + np.exp(-x)) * 4096) for i in range(16 * g + 1)], np.int64) build_tables() def w32(x): return ((np.asarray(x, np.int64) + (1 << 31)) & 0xFFFFFFFF) - (1 << 31) def i8(x): return ((np.asarray(x, np.int64) + 128) & 255) - 128 def mulh(a, b): return (np.asarray(a, np.int64) * np.int64(b)) >> 32 def quant(v, amax): amax = max(int(amax), 256) r = 0x7F000000 // amax h = amax // 254 STATS["quant"] = max(STATS["quant"], int((np.abs(np.asarray(v, np.int64)) + h).max()) << 8) return i8(mulh(w32((np.asarray(v, np.int64) + h) << 8), r)) def norm_f(xq): ss = int((xq.astype(np.int64) ** 2).sum()) m = (ss << 10) & 0xFFFFFFFF r, b = 0, 1 << 30 while b > m: b >>= 2 while b: if m >= r + b: m -= r + b r = (r >> 1) + b else: r >>= 1 b >>= 2 return int(w32((0x80000000 // r) << 10)) def plain_f(a): return int(w32(((int(a) << 9) // 127) << 8)) def mv(xq, name, f): acc = Q[name].astype(np.int64) @ xq.astype(np.int64) return mulh(w32(acc * Sc[name].astype(np.int64)), f) def sb(x): x &= 0xFF return x - 256 if x >= 128 else x def pack4(b): return sum((int(x) & 0xFF) << (8 * i) for i, x in enumerate(b)) def smaqa(a, b): return sum(sb(a >> (8 * k)) * sb(b >> (8 * k)) for k in range(4)) def q8o(v, r, h): t = int(mulh(w32((int(v) + (h if v >= 0 else -h)) << 8), r)) return max(-127, min(127, t)) def rope(v, pos): if FLOAT_ROPE: out = v.astype(np.float64).copy() for i in range(0, len(v), 2): fr = 1.0 / 10000 ** ((i % HS) / HS) c, s_ = np.cos(pos * fr), np.sin(pos * fr) a, b = float(v[i]), float(v[i + 1]) out[i] = a * c - b * s_ out[i + 1] = a * s_ + b * c return np.round(out).astype(np.int64) if pos not in ROPE: t = [] for i in range(0, HS, 2): fr = 1.0 / 10000 ** (i / HS) t += [round(np.cos(pos * fr) * 16384), round(np.sin(pos * fr) * 16384)] ROPE[pos] = t t = ROPE[pos] v = v.copy() for i in range(0, len(v), 2): c, s = t[i % HS], t[i % HS + 1] a, b = int(v[i]) >> 6, int(v[i + 1]) >> 6 v[i] = (a * c - b * s) >> 8 v[i + 1] = (a * s + b * c) >> 8 return v class State: def __init__(self): self.kq = np.zeros((NL, 32, KVD), np.int64) self.vt = np.zeros((NL, KVD, 32), np.int64) self.kf = np.zeros((NL, 32, KVD), np.float64) self.vf = np.zeros((NL, 32, KVD), np.float64) def step(tok, pos, st): x = Q["embed"][tok].astype(np.int64) * int(Sc["embed"][tok]) * 2 LAST["xs"] = [x.copy()] slot = 0 if pos == 0 else 1 + (pos - 1) % 31 n = pos + 1 if pos + 1 < 32 else 32 for l in range(NL): xq = quant(x, np.abs(x).max()) f = norm_f(xq) qv, kk, vv = mv(xq, f"q{l}", f), mv(xq, f"k{l}", f), mv(xq, f"v{l}", f) qv, kk = rope(qv, pos), rope(kk, pos) kr, kh = 0x7F000000 // int(KMAX[l]), int(KMAX[l]) // 254 vr, vh = 0x7F000000 // int(VMAX[l]), int(VMAX[l]) // 254 for i in range(KVD): st.kq[l][slot][i] = q8o(kk[i], kr, kh) st.vt[l][i][slot] = q8o(vv[i], vr, vh) st.kf[l][slot] = kk st.vf[l][slot] = vv xb = np.zeros(DIM, np.int64) for h in range(NH): if FLOAT_ATT or FLOAT_KV: qr = qv[h * HS:(h + 1) * HS].astype(np.float64) ks = st.kf[l][:n, h * HS:(h + 1) * HS] vs = st.vf[l][:n, h * HS:(h + 1) * HS] if FLOAT_KV: ks = np.round(ks / (int(KMAX[l]) / 127)).clip(-127, 127) * (int(KMAX[l]) / 127) vs = np.round(vs / (int(VMAX[l]) / 127)).clip(-127, 127) * (int(VMAX[l]) / 127) z = (ks @ qr) / 2.0 ** 32 * 0.5 pr = np.exp(z - z.max()) pr /= pr.sum() xb[h * HS:(h + 1) * HS] = np.round(pr @ vs).astype(np.int64) continue qh = qv[h * HS:(h + 1) * HS] am = max(256, int(np.abs(qh).max())) r, hf = 0x7F000000 // am, am // 254 b = [q8o(v, r, hf) for v in qh] q0, q1 = pack4(b[:4]), pack4(b[4:]) ca = min((((am * int(KMAX[l])) >> 10) * ATT_C) >> 16, 0x7FFFFFFF) STATS["ca"] = max(STATS["ca"], ca) sc = [smaqa(q0, pack4(st.kq[l][t][h * HS:h * HS + 4])) + smaqa(q1, pack4(st.kq[l][t][h * HS + 4:h * HS + 8])) for t in range(n)] m = max(sc) p = [] for s in sc: i = (int(m - s) * ca) >> (32 - PKSH) p.append(int(P8[i]) if 0 <= i < len(P8) else 0) sp = sum(p) rr = (int(VMAX[l]) << 12) // (127 * sp) for j in range(HS): o = sum(p[t] * int(st.vt[l][h * HS + j][t]) for t in range(n)) xb[h * HS + j] = (o * rr) >> 12 am = int(np.abs(xb).max()) xq = quant(xb, am) x = x + mv(xq, f"o{l}", plain_f(max(am, 256))) if l == 0: LAST["xmid"] = x.copy() xq = quant(x, np.abs(x).max()) f = norm_f(xq) hb, hb2 = mv(xq, f"gate{l}", f), mv(xq, f"up{l}", f) for i in range(HID): a = int(hb[i]) if FLOAT_SILU: hb[i] = int(round(a / (1 + np.exp(-a / 65536.0)) * int(hb2[i]) / 65536.0)) continue s = a >> 4 if a >= 8 << 16 else (0 if a < -(8 << 16) else int(SILU[((a + SRND) >> SILU_SH) + SG8])) STATS["prod"] = max(STATS["prod"], abs(s * (int(hb2[i]) >> 8))) hb[i] = (s * (int(hb2[i]) >> 8)) >> 4 am = int(np.abs(hb).max()) xq = quant(hb, am) y = Q[f"down{l}"].astype(np.int64) @ xq.astype(np.int64) x = x + mulh(w32(y * Sc[f"down{l}"].astype(np.int64)), plain_f(max(am, 256))) LAST["xs"].append(x.copy()) LAST["amax"] = max(int(np.abs(x).max()), 256) STATS["x"] = max(STATS["x"], int(np.abs(x).max())) LAST["rms"] = float(np.sqrt(np.mean((x.astype(np.float64) / 65536) ** 2))) xq = quant(x, np.abs(x).max()) acc = Q["lm_head"].astype(np.int64) @ xq.astype(np.int64) return mulh(w32(acc * Sc["lm_head"].astype(np.int64)), 1 << 30) def run_free(tok, ntok): st, out = State(), [] for pos in range(ntok): lg = step(tok, pos, st) tok = int(lg.argmax()) out.append(tok) return out