"""TinyLLama-v0 int4 for the 20-bit WROM, with per-group bases, row scales and tables.""" import os import struct import sys import numpy as np from safetensors import safe_open HERE = os.path.dirname(os.path.abspath(__file__)) SRC = os.environ.get("TL_SRC", os.path.join(os.path.dirname(os.path.abspath(__file__)), "model.safetensors")) WORDS_PER_UNIT = 15 * 64 DIM, NL, NH, HID, VOCAB, MAXPOS = 64, 8, 16, 256, 32000, 512 HS, KVD = DIM // NH, DIM def rope_perm(w): idx = [(m // 2) + (m % 2) * (HS // 2) for m in range(HS)] return w.reshape(NH, HS, -1)[:, idx, :].reshape(NH * HS, -1) def packed(unit, off): return (unit << 10) | off class Placer: def __init__(self): self.unit = 0 self.off = 0 self.cells = bytearray() def _grow(self, wp): need = 8 * (wp + 1) if len(self.cells) < need: self.cells.extend(bytes(need - len(self.cells))) def add_nibbles(self, nibs): first = None for i in range(0, len(nibs), 8): if self.off + 1 > WORDS_PER_UNIT: self.unit += 1 self.off = 0 if first is None: first = packed(self.unit, self.off) wp = packed(self.unit, self.off) self._grow(wp) for k in range(8): self.cells[8 * wp + k] = int(nibs[i + k]) & 15 self.off += 1 return first def add(self, q): rows, n = q.shape wpr = n // 8 gw = 4 * wpr first = None bases = [] for g0 in range(0, rows, 4): if self.off + gw > WORDS_PER_UNIT: self.unit += 1 self.off = 0 if first is None: first = packed(self.unit, self.off) wp0 = packed(self.unit, self.off) bases.append(wp0) for r in range(4): row = q[g0 + r] if g0 + r < rows else np.zeros(n, np.int8) for w in range(wpr): wp = wp0 + r * wpr + w self._grow(wp) for k in range(8): self.cells[8 * wp + k] = int(row[8 * w + k] + 8) & 15 self.off += gw return first def quantize(w, grid=13): w = w.astype(np.float32) best = None for c in np.linspace(0.5, 1.0, grid): s = np.abs(w).max(-1, keepdims=True) * c / 7.0 s[s == 0] = 1e-12 q = np.clip(np.round(w / s), -8, 7) e = ((q * s - w) ** 2).sum(-1, keepdims=True) if best is None: best = (q, s, e) else: m = e < best[2] best = (np.where(m, q, best[0]), np.where(m, s, best[1]), np.minimum(e, best[2])) return best[0].astype(np.int8), best[1][..., 0].astype(np.float64) def carr(name, ctype, a): a = np.asarray(a).ravel() return f"static const {ctype} {name}[{len(a)}] = {{{','.join(str(int(v)) for v in a)}}};\n" def load(): t = {} with safe_open(SRC, "pt") as f: for k in f.keys(): t[k] = f.get_tensor(k).float().numpy() return t def build(): t = load() P = Placer() Q, S, W = {}, {}, {} SC = {} def pack_scales(v): rows = len(v) nibs = [] for g in range(0, rows, 4): for pair in range(2): for r in range(2): idx = g + pair * 2 + r x = int(v[idx]) if idx < rows else 0 nibs += [(x >> (4 * j)) & 15 for j in range(4)] return nibs def put(name, w): q, s = quantize(w) Q[name], S[name] = q, np.round(s * 2**15).astype(np.int64) assert S[name].max() < 65536, (name, S[name].max()) if name.startswith("down"): W[name] = [P.add(q[:, 64 * c:64 * c + 64]) for c in range(4)] else: W[name] = P.add(q) SC[name] = P.add_nibbles(pack_scales(S[name])) fin = t["model.norm.weight"] put("lm_head", t["lm_head.weight"] * fin[None, :]) put("embed", t["model.embed_tokens.weight"]) for l in range(NL): ga = t[f"model.layers.{l}.input_layernorm.weight"] gf = t[f"model.layers.{l}.post_attention_layernorm.weight"] put(f"q{l}", rope_perm(t[f"model.layers.{l}.self_attn.q_proj.weight"] * ga[None, :])) put(f"k{l}", rope_perm(t[f"model.layers.{l}.self_attn.k_proj.weight"] * ga[None, :])) put(f"v{l}", t[f"model.layers.{l}.self_attn.v_proj.weight"] * ga[None, :]) put(f"o{l}", t[f"model.layers.{l}.self_attn.o_proj.weight"]) put(f"gate{l}", t[f"model.layers.{l}.mlp.gate_proj.weight"] * gf[None, :]) put(f"up{l}", t[f"model.layers.{l}.mlp.up_proj.weight"] * gf[None, :]) put(f"down{l}", t[f"model.layers.{l}.mlp.down_proj.weight"]) return P, Q, S, W, SC def emit(P, Q, S, W, SC): h = ["#include \n"] h.append(f"#define DIM {DIM}\n#define HID {HID}\n#define NL {NL}\n#define NH {NH}\n" f"#define HS {HS}\n#define KVD {KVD}\n#define VOCAB {VOCAB}\n#define MAXPOS {MAXPOS}\n") h.append(f"#define ATT_C {round(2 ** 31 / (127 ** 2 * np.sqrt(HS)))}\n") h.append(f"#define G_HEAD {W['lm_head']}\n#define G_EMB {W['embed']}\n") for nm in ("q", "k", "v", "o", "gate", "up"): h.append(carr(f"G_{nm}", "int32_t", [W[f"{nm}{l}"] for l in range(NL)])) h.append(carr("G_down", "int32_t", [W[f"down{l}"][c] for l in range(NL) for c in range(4)])) h.append(f"#define SC_HEAD {SC['lm_head']}\n#define SC_EMB {SC['embed']}\n") for nm in ("q", "k", "v", "o", "gate", "up", "down"): h.append(carr(f"SC_{nm}", "int32_t", [SC[f"{nm}{l}"] for l in range(NL)])) h.append(carr("SILU", "int32_t", [round((x := (i - 256) / 32) / (1 + np.exp(-x)) * 4096) for i in range(512)])) cs = [] for pos in range(MAXPOS): for i in range(0, HS, 2): fr = 1.0 / 10000 ** (i / HS) cs += [round(np.cos(pos * fr) * 16384), round(np.sin(pos * fr) * 16384)] h.append(carr("ROPE", "int16_t", cs)) h.append(carr("EXP2", "uint16_t", [round(2 ** (-i / 256) * 32768) for i in range(256)])) import json kv = json.load(open(HERE + "/kvmax.json")) h.append(carr("KMAX", "int32_t", kv["kmax"])) h.append(carr("VMAX", "int32_t", kv["vmax"])) h.append(carr("P8", "uint8_t", [round(127 * np.exp(-i / 32)) for i in range(256)])) open(HERE + "/weights_tl.h", "w").write("".join(h)) open(HERE + "/wrom_tl.c", "w").write( "#include \n__attribute__((section(\".wrom\"), used)) const uint8_t " f"WROM_DATA[{len(P.cells)}] = {{{','.join(map(str, P.cells))}}};\n") print(f"header y wrom escritos: {len(P.cells)} celdas, {P.unit + 1} unidades") if __name__ == "__main__": P, Q, S, W, SC = build() emit(P, Q, S, W, SC)