| 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 |
|
|
