Something went wrong. Try again.
A general purpose .. programming language written in zig ijadux2.tngl.sh/Rinnescript
zig language
Something went wrong. Try again.
22 kB · 581 lines
Zig
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582//! Minimal LLM inference for RinneScript: a llama2.c-scale decoder-only//! transformer loaded from GGUF (F32 weights, zero-copy over the mmap) with//! a CPU forward pass (RMSNorm -> RoPE -> GQA attention -> SwiGLU MLP).//!//! Self-defined GGUF convention (documented in docs/libs/llm.md)://! metadata (integer kvs): rinne.llm.dim, .hidden_dim, .n_layers,//! .n_heads, .n_kv_heads, .vocab_size, .seq_len//! tensors (all F32)://! tok_embd [vocab_size, dim]//! output_norm [dim]//! output [vocab_size, dim]//! blk.{i}.attn_norm [dim]//! blk.{i}.attn_q [dim, dim]//! blk.{i}.attn_k [kv_dim, dim]//! blk.{i}.attn_v [kv_dim, dim]//! blk.{i}.attn_o [dim, dim]//! blk.{i}.ffn_norm [dim]//! blk.{i}.ffn_gate [hidden_dim, dim]//! blk.{i}.ffn_down [dim, hidden_dim]//! blk.{i}.ffn_up [hidden_dim, dim]//!//! Text-level generation needs a tokenizer, which is future work; v1 exposes//! token-id-level forward passes (`model.forward(ids) -> last-token logits`).const std = @import("std");const value = @import("value.zig");const gguf = @import("ai_gguf.zig");
pub const Error = error{ InvalidModel, OutOfMemory,};
const rope_theta: f32 = 10000.0;const norm_eps: f32 = 1e-5;
/// Reinterprets a mapped byte window as an F32 slice, validating dtype and/// element count against `expect_numel`.fn f32View(v: value.TensorView, file_bytes: []u8, expect_numel: u64) Error![]f32 { if (v.dtype != .f32) return Error.InvalidModel; var n: u64 = 1; for (v.shape) |d| { n = std.math.mul(u64, n, d) catch return Error.InvalidModel; } if (n != expect_numel or @as(u64, v.len) != expect_numel * 4) return Error.InvalidModel; const start: usize = @intCast(v.off); const end: usize = start + @as(usize, @intCast(v.len)); if (end > file_bytes.len) return Error.InvalidModel; const aligned: []align(4) u8 = @alignCast(file_bytes[start..end]); return std.mem.bytesAsSlice(f32, aligned);}
fn findByPrefix(views: []const value.TensorView, name_buf: []const u8) ?value.TensorView { for (views) |v| { if (std.mem.eql(u8, v.name, name_buf)) return v; } return null;}
/// Builds a model from parsed GGUF views over `file_bytes`. Weight slices/// borrow `file_bytes`; keep its mapping alive until `deinit`. When/// `owner` is provided it is released by `deinit` (refcount pattern shared/// with the ai library loaders).pub fn load( gpa: std.mem.Allocator, cfg: value.LlmConfig, nums: []const gguf.KvNum, views: []const value.TensorView, file_bytes: []u8, owner: ?*value.MapOwner,) Error!*value.LlmModel { // Config sanity. if (cfg.dim == 0 or cfg.n_heads == 0 or cfg.dim % cfg.n_heads != 0) return Error.InvalidModel; if (cfg.n_kv_heads == 0 or cfg.n_heads % cfg.n_kv_heads != 0) return Error.InvalidModel; if (cfg.n_layers == 0 or cfg.hidden_dim == 0 or cfg.vocab_size == 0 or cfg.seq_len == 0) return Error.InvalidModel; _ = nums; // config was already extracted by the caller
const dim: usize = @intCast(cfg.dim); const hidden: usize = @intCast(cfg.hidden_dim); const kv_dim: usize = @intCast(cfg.kvDim()); const vocab: usize = @intCast(cfg.vocab_size);
const m = gpa.create(value.LlmModel) catch return Error.OutOfMemory; errdefer gpa.destroy(m); m.* = .{ .cfg = cfg, .tok_embd = &.{}, .output_norm = &.{}, .output = &.{}, .layers = &.{}, .owner = owner, };
m.tok_embd = try f32View(findByPrefix(views, "tok_embd") orelse return Error.InvalidModel, file_bytes, vocab * dim); m.output_norm = try f32View(findByPrefix(views, "output_norm") orelse return Error.InvalidModel, file_bytes, dim); m.output = try f32View(findByPrefix(views, "output") orelse return Error.InvalidModel, file_bytes, vocab * dim);
const layers = gpa.alloc(value.LlmLayerWeights, @intCast(cfg.n_layers)) catch return Error.OutOfMemory; errdefer { // Individual slices borrow mapped memory; nothing to free per-layer. gpa.free(layers); } m.layers = layers;
for (layers, 0..) |*lw, li| { lw.* = .{ .attn_norm = try layerTensor(views, file_bytes, li, "attn_norm", dim), .attn_q = try layerTensor(views, file_bytes, li, "attn_q", dim * dim), .attn_k = try layerTensor(views, file_bytes, li, "attn_k", kv_dim * dim), .attn_v = try layerTensor(views, file_bytes, li, "attn_v", kv_dim * dim), .attn_o = try layerTensor(views, file_bytes, li, "attn_o", dim * dim), .ffn_norm = try layerTensor(views, file_bytes, li, "ffn_norm", dim), .ffn_gate = try layerTensor(views, file_bytes, li, "ffn_gate", hidden * dim), .ffn_down = try layerTensor(views, file_bytes, li, "ffn_down", dim * hidden), .ffn_up = try layerTensor(views, file_bytes, li, "ffn_up", hidden * dim), }; }
// KV caches are allocated lazily by `forward`. return m;}
fn layerTensor( views: []const value.TensorView, file_bytes: []u8, layer: usize, suffix: []const u8, expect_numel: usize,) Error![]f32 { var name_buf: [128]u8 = undefined; const name = std.fmt.bufPrint(&name_buf, "blk.{d}.{s}", .{ layer, suffix }) catch return Error.InvalidModel; const v = findByPrefix(views, name) orelse return Error.InvalidModel; return f32View(v, file_bytes, expect_numel);}
/// Releases the model: KV caches, layer table, and the shared mapping./// Idempotent via `dead`.pub fn deinit(gpa: std.mem.Allocator, m: *value.LlmModel) void { if (m.dead) return; m.dead = true; for (m.key_cache) |kc| gpa.free(kc); for (m.value_cache) |vc| gpa.free(vc); if (m.key_cache.len != 0) gpa.free(m.key_cache); if (m.value_cache.len != 0) gpa.free(m.value_cache); gpa.free(m.layers); if (m.owner) |o| o.release(); gpa.destroy(m);}
// ---------------------------------------------------------------------------// Forward pass (llama2.c-style math)// ---------------------------------------------------------------------------
fn rmsnorm(out: []f32, x: []const f32, w: []const f32) void { var ss: f32 = 0; for (x) |v| ss += v * v; const inv = 1.0 / @sqrt(ss / @as(f32, @floatFromInt(x.len)) + norm_eps); for (x, 0..) |v, i| out[i] = w[i] * (v * inv);}
fn matmul(out: []f32, x: []const f32, w: []const f32) void { const n_out = out.len; const n_in = x.len; for (0..n_out) |o| { var sum: f32 = 0; const row = w[o * n_in ..][0..n_in]; for (row, x) |wv, xv| sum += wv * xv; out[o] = sum; }}
fn softmax(s: []f32) void { var max = -std.math.inf(f32); for (s) |v| max = @max(max, v); var sum: f32 = 0; for (s) |*v| { v.* = @exp(v.* - max); sum += v.*; } for (s) |*v| v.* /= sum;}
fn silu(x: f32) f32 { return x / (1.0 + @exp(-x));}
/// Applies rotary embeddings in-place to q ([dim]) and k ([kv_dim]).fn rope(q: []f32, k: []f32, head_size: usize, pos: usize) void { rotateSlice(q, head_size, pos); rotateSlice(k, head_size, pos);}
fn rotateSlice(s: []f32, head_size: usize, pos: usize) void { var off: usize = 0; while (off < s.len) : (off += head_size) { rotateHead(s[off .. off + head_size], pos); }}
fn rotateHead(h: []f32, pos: usize) void { const hs = h.len; var i: usize = 0; while (i < hs) : (i += 2) { const freq = 1.0 / std.math.pow(f32, rope_theta, @as(f32, @floatFromInt(i)) / @as(f32, @floatFromInt(hs))); const ang: f32 = @floatFromInt(pos); const val = ang * freq; const fcr = @cos(val); const fci = @sin(val); const v0 = h[i]; const v1 = h[i + 1]; h[i] = v0 * fcr - v1 * fci; h[i + 1] = v0 * fci + v1 * fcr; }}
fn ensureCaches(gpa: std.mem.Allocator, m: *value.LlmModel) Error!void { const nl: usize = @intCast(m.cfg.n_layers); if (m.key_cache.len == nl and m.value_cache.len == nl) return; // (Re)allocate for the current architecture. for (m.key_cache) |kc| gpa.free(kc); for (m.value_cache) |vc| gpa.free(vc); if (m.key_cache.len != 0) gpa.free(m.key_cache); if (m.value_cache.len != 0) gpa.free(m.value_cache); m.key_cache = &.{}; m.value_cache = &.{};
const kv_dim: usize = @intCast(m.cfg.kvDim()); const seq: usize = @intCast(m.cfg.seq_len); m.key_cache = gpa.alloc([]f32, nl) catch return Error.OutOfMemory; m.value_cache = gpa.alloc([]f32, nl) catch return Error.OutOfMemory; var filled: usize = 0; errdefer { for (m.key_cache[0..filled]) |kc| gpa.free(kc); for (m.value_cache[0..filled]) |vc| gpa.free(vc); gpa.free(m.key_cache); gpa.free(m.value_cache); m.key_cache = &.{}; m.value_cache = &.{}; } for (0..nl) |i| { m.key_cache[i] = gpa.alloc(f32, seq * kv_dim) catch return Error.OutOfMemory; @memset(m.key_cache[i], 0); m.value_cache[i] = gpa.alloc(f32, seq * kv_dim) catch return Error.OutOfMemory; @memset(m.value_cache[i], 0); filled += 1; }}
/// Runs the full sequence through the transformer and returns the LAST/// token's logits ([vocab_size], owned by the caller). The KV cache is reset/// at the start of every call.pub fn forward(gpa: std.mem.Allocator, m: *value.LlmModel, ids: []const i64) Error![]f32 { if (ids.len == 0) return Error.InvalidModel; if (ids.len > m.cfg.seq_len) return Error.InvalidModel; for (ids) |id| { if (id < 0 or @as(u64, @intCast(id)) >= m.cfg.vocab_size) return Error.InvalidModel; }
try ensureCaches(gpa, m); defer m.pos = 0;
const dim: usize = @intCast(m.cfg.dim); const hidden: usize = @intCast(m.cfg.hidden_dim); const kv_dim: usize = @intCast(m.cfg.kvDim()); const hs: usize = @intCast(m.cfg.headSize()); const n_heads: usize = @intCast(m.cfg.n_heads); const n_kv: usize = @intCast(m.cfg.n_kv_heads); const kv_mul: usize = n_heads / n_kv; const vocab: usize = @intCast(m.cfg.vocab_size);
// Scratch buffers. const xt = gpa.alloc(f32, dim) catch return Error.OutOfMemory; defer gpa.free(xt); const xb = gpa.alloc(f32, dim) catch return Error.OutOfMemory; defer gpa.free(xb); const hb = gpa.alloc(f32, hidden) catch return Error.OutOfMemory; defer gpa.free(hb); const xb2 = gpa.alloc(f32, hidden) catch return Error.OutOfMemory; defer gpa.free(xb2); const q = gpa.alloc(f32, dim) catch return Error.OutOfMemory; defer gpa.free(q); const k = gpa.alloc(f32, kv_dim) catch return Error.OutOfMemory; defer gpa.free(k); const v = gpa.alloc(f32, kv_dim) catch return Error.OutOfMemory; defer gpa.free(v); const att = gpa.alloc(f32, m.cfg.seq_len) catch return Error.OutOfMemory; defer gpa.free(att);
var x = gpa.alloc(f32, dim) catch return Error.OutOfMemory; defer gpa.free(x);
for (ids, 0..) |tok, t| { // Embed. const emb_row = m.tok_embd[@as(usize, @intCast(tok)) * dim ..][0..dim]; @memcpy(x, emb_row);
for (m.layers, 0..) |lw, li| { // --- Attention block --- rmsnorm(xt, x, lw.attn_norm); matmul(q, xt, lw.attn_q); matmul(k, xt, lw.attn_k); matmul(v, xt, lw.attn_v); rope(q, k, hs, t);
const kc = m.key_cache[li]; const vc = m.value_cache[li]; @memcpy(kc[t * kv_dim ..][0..kv_dim], k); @memcpy(vc[t * kv_dim ..][0..kv_dim], v);
// Multi-head attention with GQA grouping. for (0..n_heads) |h| { const qh = q[h * hs ..][0..hs]; const kvh = (h / kv_mul) * hs; // Scores over positions 0..=t. var valid: usize = 0; for (0..t + 1) |tp| { const kh = kc[tp * kv_dim + kvh ..][0..hs]; var dot: f32 = 0; for (qh, kh) |qv, kvv| dot += qv * kvv; att[valid] = dot / @sqrt(@as(f32, @floatFromInt(hs))); valid += 1; } softmax(att[0..valid]); @memset(xb[h * hs ..][0..hs], 0); for (0..valid) |tp| { const vh = vc[tp * kv_dim + kvh ..][0..hs]; const wgt = att[tp]; for (vh, 0..) |vv, j| xb[h * hs + j] += wgt * vv; } } matmul(hb[0..dim], xb, lw.attn_o); // reuse hb head as [dim] buffer for (hb[0..dim], 0..) |av, i| x[i] += av;
// --- MLP block (SwiGLU) --- rmsnorm(xt, x, lw.ffn_norm); matmul(hb, xt, lw.ffn_gate); matmul(xb2, xt, lw.ffn_up); for (hb, xb2) |*gv, uv| gv.* = silu(gv.*) * uv; matmul(xb[0..dim], hb, lw.ffn_down); for (xb[0..dim], 0..) |mv, i| x[i] += mv; } }
// Final norm + classifier; return last-position logits. rmsnorm(xt, x, m.output_norm); const logits = gpa.alloc(f32, vocab) catch return Error.OutOfMemory; errdefer gpa.free(logits); matmul(logits, xt, m.output); return logits;}
// ---------------------------------------------------------------------------// Tests: tiny fixtures built in-memory.// ---------------------------------------------------------------------------
const testing = std.testing;
/// Test/demo support: serializes a tiny deterministic 1-layer model/// (dim=4, hidden=8, vocab=3, seq_len=8; attention/MLP zeroed so logits are/// easy to reason about). Caller owns the returned bytes.pub fn buildMicroModelGguf(a: std.mem.Allocator) ![]u8 { var list: std.ArrayList(u8) = .empty; errdefer list.deinit(a); _ = try buildMicroModel(&list, .{}); return list.toOwnedSlice(a);}
fn wU32(list: *std.ArrayList(u8), n: u32) !void { var b: [4]u8 = undefined; std.mem.writeInt(u32, &b, n, .little); try list.appendSlice(testing.allocator, &b);}
fn wU64(list: *std.ArrayList(u8), n: u64) !void { var b: [8]u8 = undefined; std.mem.writeInt(u64, &b, n, .little); try list.appendSlice(testing.allocator, &b);}
fn wStr(list: *std.ArrayList(u8), s: []const u8) !void { try wU64(list, s.len); try list.appendSlice(testing.allocator, s);}
fn wF32(list: *std.ArrayList(u8), xs: []const f32) !void { for (xs) |x| { var b: [4]u8 = undefined; std.mem.writeInt(u32, &b, @bitCast(x), .little); try list.appendSlice(testing.allocator, &b); }}
const FixtureTensor = struct { name: []const u8, /// Shape as [outer, inner] (numel only matters to the loader). outer: u64, inner: u64, data: []const f32,};
const MicroCfg = struct { const dim: u64 = 4; const hidden: u64 = 8; const vocab: u64 = 3; const seq: u64 = 8;};
/// Builds a tiny 1-layer model GGUF: attention and MLP weights zeroed so the/// residual stream equals the token embedding; final norm ones; output rows/// distinct so different tokens yield different logits.fn buildMicroModel(list: *std.ArrayList(u8), opts: struct { drop_tensor: ?[]const u8 = null }) !value.LlmConfig { const a = testing.allocator; const dim = MicroCfg.dim; const hidden = MicroCfg.hidden; const vocab = MicroCfg.vocab;
var tensors = std.ArrayList(FixtureTensor).empty; defer tensors.deinit(a); const zeros_h_d = [_]f32{0} ** (hidden * dim); const zeros_d_d = [_]f32{0} ** (dim * dim); try tensors.append(a, .{ .name = "tok_embd", .outer = vocab, .inner = dim, .data = &[_]f32{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12 } }); try tensors.append(a, .{ .name = "output_norm", .outer = 1, .inner = dim, .data = &[_]f32{ 1, 1, 1, 1 } }); try tensors.append(a, .{ .name = "output", .outer = vocab, .inner = dim, .data = &[_]f32{ 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0 } }); try tensors.append(a, .{ .name = "blk.0.attn_norm", .outer = 1, .inner = dim, .data = &[_]f32{ 1, 1, 1, 1 } }); try tensors.append(a, .{ .name = "blk.0.attn_q", .outer = dim, .inner = dim, .data = &zeros_d_d }); try tensors.append(a, .{ .name = "blk.0.attn_k", .outer = dim, .inner = dim, .data = &zeros_d_d }); try tensors.append(a, .{ .name = "blk.0.attn_v", .outer = dim, .inner = dim, .data = &zeros_d_d }); try tensors.append(a, .{ .name = "blk.0.attn_o", .outer = dim, .inner = dim, .data = &zeros_d_d }); try tensors.append(a, .{ .name = "blk.0.ffn_norm", .outer = 1, .inner = dim, .data = &[_]f32{ 1, 1, 1, 1 } }); try tensors.append(a, .{ .name = "blk.0.ffn_gate", .outer = hidden, .inner = dim, .data = &zeros_h_d }); try tensors.append(a, .{ .name = "blk.0.ffn_down", .outer = dim, .inner = hidden, .data = &[_]f32{0} ** (dim * hidden) }); try tensors.append(a, .{ .name = "blk.0.ffn_up", .outer = hidden, .inner = dim, .data = &zeros_h_d });
try list.appendSlice(a, "GGUF"); try wU32(list, 3); // version var kept: u64 = tensors.items.len; if (opts.drop_tensor != null) kept -= 1; try wU64(list, kept); // kvs: general.alignment + 7 rinne.llm.* try wU64(list, 8); try wStr(list, "general.alignment"); try wU32(list, 4); try wU32(list, 32); const cfg_vals = [_]struct { name: []const u8, v: u64 }{ .{ .name = "rinne.llm.dim", .v = dim }, .{ .name = "rinne.llm.hidden_dim", .v = hidden }, .{ .name = "rinne.llm.n_layers", .v = 1 }, .{ .name = "rinne.llm.n_heads", .v = 1 }, .{ .name = "rinne.llm.n_kv_heads", .v = 1 }, .{ .name = "rinne.llm.vocab_size", .v = vocab }, .{ .name = "rinne.llm.seq_len", .v = MicroCfg.seq }, }; for (cfg_vals) |cv| { try wStr(list, cv.name); try wU32(list, 10); // UINT64 try wU64(list, cv.v); }
// Tensor infos (dims written inner-first; loader only checks numel). for (tensors.items) |t| { if (opts.drop_tensor) |dropped| { if (std.mem.eql(u8, dropped, t.name)) continue; } try wStr(list, t.name); try wU32(list, 2); // n_dims try wU64(list, t.inner); try wU64(list, t.outer); try wU32(list, 0); // F32 try wU64(list, 0); // rel_off patched below }
// Rebuild cleanly: infos carry real running offsets this time. list.clearRetainingCapacity(); try list.appendSlice(a, "GGUF"); try wU32(list, 3); try wU64(list, kept); try wU64(list, 8); try wStr(list, "general.alignment"); try wU32(list, 4); try wU32(list, 32); for (cfg_vals) |cv| { try wStr(list, cv.name); try wU32(list, 10); try wU64(list, cv.v); } var running: u64 = 0; for (tensors.items) |t| { if (opts.drop_tensor) |dropped| { if (std.mem.eql(u8, dropped, t.name)) continue; } try wStr(list, t.name); try wU32(list, 2); try wU64(list, t.inner); try wU64(list, t.outer); try wU32(list, 0); try wU64(list, running); running += @as(u64, @intCast(t.data.len)) * 4; } while (list.items.len % 32 != 0) try list.append(a, 0); for (tensors.items) |t| { if (opts.drop_tensor) |dropped| { if (std.mem.eql(u8, dropped, t.name)) continue; } try wF32(list, t.data); }
return .{ .dim = dim, .hidden_dim = hidden, .n_layers = 1, .n_heads = 1, .n_kv_heads = 1, .vocab_size = vocab, .seq_len = MicroCfg.seq, };}
test "micro-model loads and forwards deterministically" { const a = testing.allocator; var bytes_list: std.ArrayList(u8) = .empty; defer bytes_list.deinit(a); const cfg = try buildMicroModel(&bytes_list, .{});
const r = try gguf.parseWithKv(a, bytes_list.items); defer { for (r.views) |v| { a.free(v.name); a.free(v.shape); } a.free(r.views); for (r.nums) |n| a.free(n.name); a.free(r.nums); }
const m = try load(a, cfg, r.nums, r.views, bytes_list.items, null); defer deinit(a, m);
const ids = [_]i64{ 0, 1 }; const lg1 = try forward(a, m, &ids); defer a.free(lg1); try testing.expectEqual(@as(usize, 3), lg1.len); for (lg1) |x| try testing.expect(std.math.isFinite(x));
const lg2 = try forward(a, m, &ids); defer a.free(lg2); for (lg1, lg2) |p, q| try testing.expectEqual(p, q);
// Different tokens differ (output rows are distinct basis vectors). const ids_b = [_]i64{0}; const lgb = try forward(a, m, &ids_b); defer a.free(lgb); // Token 0 embeds to an ascending vector, so the third output row // (e2 basis) yields the largest logit. try testing.expect(lgb[2] > lgb[0]);
// Token id range is enforced. const bad = [_]i64{99}; try testing.expectError(Error.InvalidModel, forward(a, m, &bad)); const empty_ids = [_]i64{}; try testing.expectError(Error.InvalidModel, forward(a, m, &empty_ids));}
test "load rejects missing kv and missing/wrong tensors" { const a = testing.allocator;
// Build a config-only fixture missing one tensor. var bytes_list: std.ArrayList(u8) = .empty; defer bytes_list.deinit(a); const cfg = try buildMicroModel(&bytes_list, .{ .drop_tensor = "blk.0.attn_o" });
const r = try gguf.parseWithKv(a, bytes_list.items); defer { for (r.views) |v| { a.free(v.name); a.free(v.shape); } a.free(r.views); for (r.nums) |n| a.free(n.name); a.free(r.nums); } try testing.expectError(Error.InvalidModel, load(a, cfg, r.nums, r.views, bytes_list.items, null));}