From b9b6054a8fd6e79d2fcfa853c32b8eeb59fd6482 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Mon, 26 Oct 2026 05:00:41 -0500 Subject: [PATCH 1/7] llama: Nanbeige4.3-33B on Metal, the bf16 release served as published A new `llama` family for the native Metal engine: dense Llama-architecture checkpoints (heads of 128, default RoPE, no biases, an untied head) in bf16, read in place from the shards with nothing converted. It keeps HF modeling_llama's rounding points (a bf16 residual, LlamaRMSNorm, RoPE with bf16 cos and sin, SiLU), runs prompt chunks on gemm_nax's tensor-unit GEMM and decode windows on a simdgroup-a-row matvec, decode attention on sdpa_vector lanes at absolute positions, and sampled draws on tf_sample_full. tf-llama-run checks greedy tokens, every position's logits, and that windows and one-row steps agree bit for bit. Measured on an M5 Ultra against HF's bf16 run of the release: 1,500 positions teacher-forced at mean KL 0.0015 and top-1 99.1%; decode 16.4 tokens/s a stream and 28-32 with two; cold prompts 1.3k tokens/s at 8k tokens and 563 at 131k. docs/recipes/nanbeige-4.3-33b-metal.md has the rest. --- docs/recipes/README.md | 1 + docs/recipes/nanbeige-4.3-33b-metal.md | 105 +++++ zig/build/metal_targets.zig | 1 + zig/kernels/metal/llama/attention_nax.metal | 204 ++++++++++ zig/kernels/metal/llama/gemm.metal | 14 + zig/kernels/metal/llama/glue.metal | 269 +++++++++++++ zig/kernels/metal/sources.zig | 6 + zig/src/cli/hub.zig | 1 + zig/src/families/llama/backend.zig | 406 ++++++++++++++++++++ zig/src/families/llama/budget.zig | 46 +++ zig/src/families/llama/config.zig | 241 ++++++++++++ zig/src/families/llama/fixtures/config.json | 37 ++ zig/src/families/llama/forward.zig | 253 ++++++++++++ zig/src/families/llama/llama.zig | 17 + zig/src/families/llama/model.zig | 115 ++++++ zig/src/families/llama/residency.zig | 25 ++ zig/src/families/llama/snapshot.zig | 77 ++++ zig/src/families/llama/state.zig | 149 +++++++ zig/src/families/llama/weights.zig | 97 +++++ zig/src/native/llama.zig | 180 +++++++++ zig/src/native/metal.zig | 11 +- zig/src/tensorfold.zig | 2 + zig/tests/llama/run.zig | 193 ++++++++++ 23 files changed, 2444 insertions(+), 6 deletions(-) create mode 100644 docs/recipes/nanbeige-4.3-33b-metal.md create mode 100644 zig/kernels/metal/llama/attention_nax.metal create mode 100644 zig/kernels/metal/llama/gemm.metal create mode 100644 zig/kernels/metal/llama/glue.metal create mode 100644 zig/src/families/llama/backend.zig create mode 100644 zig/src/families/llama/budget.zig create mode 100644 zig/src/families/llama/config.zig create mode 100644 zig/src/families/llama/fixtures/config.json create mode 100644 zig/src/families/llama/forward.zig create mode 100644 zig/src/families/llama/llama.zig create mode 100644 zig/src/families/llama/model.zig create mode 100644 zig/src/families/llama/residency.zig create mode 100644 zig/src/families/llama/snapshot.zig create mode 100644 zig/src/families/llama/state.zig create mode 100644 zig/src/families/llama/weights.zig create mode 100644 zig/src/native/llama.zig create mode 100644 zig/tests/llama/run.zig diff --git a/docs/recipes/README.md b/docs/recipes/README.md index 8017124388..4af62b9fb0 100644 --- a/docs/recipes/README.md +++ b/docs/recipes/README.md @@ -24,6 +24,7 @@ For agent-assisted contributions, read the [TensorFold contributor skill](SKILL. | Kolibri-1 | [Native CUDA, block FP8, research checkpoint](kolibri-1.md) | | LFM2.5-24B-A2B | [Native Metal, MLX 4-bit, shared rounds](lfm2.5-24b-a2b.md) | | Spark-X2.6-4B | [Native CUDA](spark-x2.6-4b-cuda.md), [native Metal, the bf16 release, contexts to 1M](spark-x2.6-4b-metal.md) | +| Nanbeige4.3-33B | [Native Metal, the bf16 release, the Llama family](nanbeige-4.3-33b-metal.md) | ## Capability floor diff --git a/docs/recipes/nanbeige-4.3-33b-metal.md b/docs/recipes/nanbeige-4.3-33b-metal.md new file mode 100644 index 0000000000..3fa56775d4 --- /dev/null +++ b/docs/recipes/nanbeige-4.3-33b-metal.md @@ -0,0 +1,105 @@ +# Nanbeige4.3-33B, native Metal + +This recipe covers text generation with [`Nanbeige/Nanbeige4.3-33B`](https://huggingface.co/Nanbeige/Nanbeige4.3-33B) +(model_type `llama`: a dense 33B Llama-architecture model, 512K context, thinking and tool calls in its chat template) +on the Metal engine, served from the bf16 release as published: no conversion, nothing lost. It was brought up on an +M5 Ultra (256 GB); other Macs are untested. The family serves any `llama` checkpoint with heads of 128, default RoPE, +no biases and its own output head, in bf16; it refuses others by name. + +## Build and serve + +```sh +hf download Nanbeige/Nanbeige4.3-33B --local-dir models/nanbeige4.3-33b +zig build native -Doptimize=ReleaseFast -Dcpu=native +zig build grammar -Dxgrammar=PATH_TO_XGRAMMAR_v0.2.8 -Dcpu=native # structured output +LLAMA_LOAD_LIMIT_GB=150 zig-out/native/bin/tensorfold-native serve models/nanbeige4.3-33b --name nanbeige --port 8888 \ + --context 131072 --parallel 2 --prompt-cache-gib 8 +``` + +`LLAMA_LOAD_LIMIT_GB` caps weights, scratch and caches together (75% of RAM when unset). The weights take 61.7 GiB; a +stream's caches take 208 KiB a position (every one of the 52 layers keeps bf16 keys and values), so a 131,072-token +stream holds 26 GiB and the full 524,288-token context 104 GiB: 150 GB holds two 131,072-token streams, or one of +262,144; the whole context needs a Mac with the weights and 104 GiB free (192 GB, the 75% default on a 256 GB Mac, +holds one 524,288-token stream). Thinking is on by default, as in the chat template; a request sends +`"chat_template_kwargs": {"enable_thinking": false}` to turn it off. The server forwards only the thinking switch to +the template: earlier turns' reasoning stays out of the prompt (the template's `preserve_thinking` default) and tool +calls take the template's XML form (its `tool_call_format` default), which the server parses. + +## What the family is + +`zig/src/families/llama/` follows the Spark-X2.6 Metal family, with HF's `modeling_llama` as the reference: + +- The reference's rounding points, its PyTorch code in bf16: the residual stays bf16 and each add rounds to it; the + norm is `bf16(w * bf16(x * rsqrt(mean(x^2) + eps)))` with the mean in fp32; RoPE is `bf16(bf16(x * c) + bf16(o * s))` + with the fp32 angle's cos and sin rounded to bf16 as HF rounds them, pairs (d, d + 64); the MLP is + `bf16(bf16(silu(g)) * u)` (`kernels/metal/llama/glue.metal`). +- Projections over the release's bf16 rows, read in place from the mapped shards (no copies, no stacking): prompt + chunks of up to 2,048 rows on `gemm_nax`'s tensor-unit GEMM (`gemm.metal`); decode windows on `ll_matvec`, a + simdgroup an output row, fp32 sums rounded once as torch's bf16 linear rounds them. q, k and v write side by side + into one row, gate and up into another. +- Attention: prompt chunks on the tensor units (Gemma 4's prompt attention over a position-major cache, + `attention_nax.metal`); decode rows on sdpa_vector lanes over key blocks at absolute positions, so a row's bits + never depend on the rows beside it. +- The lane host's shared rounds, a lone stream's steps drawn on the GPU a step ahead (argmax, or `tf_sample_full`), + copy drafts and the prompt cache. + +The Metal glue kernels are not registered with the kernel contract, as with the other Metal family glue. + +## Measured (M5 Ultra, the bf16 release, 26 Oct 2026) + +Against Hugging Face's `modeling_llama` on the release (transformers 4.57.1, bf16 on the CPU): 1,500 positions +teacher-forced, mean KL 0.0015, top-1 99.1%; 64 positions after a 9,000-token prompt, mean KL 0.0013, top-1 100%. Six +greedy chats of 64 tokens (code, prose, Chinese, a document, with and without thinking): three equal HF's token for +token; the other three part at a position where the engine's two best logits are the same bf16 value. One-row steps +and windows of 7 or 16 rows give the same logits bit for bit after prompts of 64 and 1,500 tokens. + +Served through the recipe harness (an idle 97 GiB server of another model was loaded beside it, so memory readings +include it), `--context 131072 --parallel 2`, greedy 256-token replies: + +| streams | prose | code | structured | list | first token | a round | +|---:|---:|---:|---:|---:|---:|---:| +| 1 | 16.4 | 16.3 | 16.4 | 16.4 | 0.23-0.26 s | 61.4 ms | +| 2 | 31.8 | 28.0 | 31.7 | 31.7 | 0.34-0.38 s | 66.4 ms | + +Tokens a second, all streams together. One stream reads the 61.7 GiB of weights each step; two streams read them once +for both rows. Sampled replies (temperature, top-p, seeds) decode at 16.2-16.3 tokens/s, the same as greedy. + +| cold prompt | 2,109 | 8,254 | 32,827 | 61,502 | 131,132 tokens | +|---|---:|---:|---:|---:|---:| +| tokens a second | 1,259 | 1,313 | 1,074 | 866 | 563 | +| time to first token | 1.68 s | 6.29 s | 30.6 s | 71.0 s | 233 s | + +A fresh server, each prompt starting with words no earlier one did; the 131,132-token prompt on a +`--context 262144 --parallel 1` server, the others on the 131,072-token one. With thinking on, within 4%. A multi-turn +chat's third turn, resumed from 5,493 of its 5,649 kept tokens, starts in 0.34 s against 4.5 s cold, with the same reply. + +- Structured output (json_object, json_schema, regex, choice, grammar) 18/18, drafted == plain 12/12; tool calls + (auto, required, named, thinking, streamed) 5/5; seeded sampling repeats 9/9. +- A 60-turn soak (3 chats of 20 turns beside a decoding stream) and 12 concurrent replies each equal the request + alone; streaming 4/4, multi-turn recall 8/8, stop strings 3/3, thinking 10/10, max_tokens 3/3. +- A passphrase at 10%, 50% and 90% of a 128,550-token prompt comes back exactly on the 131,072-token server. + +## Performance options + +Each keeps the replies' bits. + +- **Streams (`--parallel`).** One stream decodes at 16.4 tokens/s, bound by reading the weights each step; two give + 28-32 tokens/s in all (66 ms a round against 61), each reply equal to its run alone. The row kernel holds up to 32 + rows a launch; the caches set the count (26 GiB a 131,072-token stream). +- **Context and memory (`--context`, `LLAMA_LOAD_LIMIT_GB`).** 208 KiB a token. Cold prompts run 1.1-1.3k tokens/s + to 33k tokens, 866 at 61k and 563 at 131k (attention grows with the prompt). +- **Prompt cache (`--prompt-cache-gib`).** A resumed turn starts in 0.34 s against 4.5 s cold; kept states give way to + streams that need the room. +- **Copy drafts.** On from the lane core, replies equal to drafts off; they accepted 0-30% of their tokens on the + harness's prompts and gave no measurable decode gain (drafted 16 tokens/s, undrafted 16). +- **Sampling.** Draws run `tf_sample_full` over the whole 166,912-token vocabulary (the 1,024-candidate kernel is + compiled for a 131,072-token vocabulary): sampled replies decode at the greedy speed. +- **Not offered.** No int8 or fp8 cache, no fast-math, no reduced-precision sums, no quantized weights. + +## Not done + +No quality suite was run. No drafter is served (the family has no draft head; copy drafts come from the lane core). +Logprobs are refused by name: the release's tokenizer decodes through a `Strip` step, which the engine's tokenizer does +not count as token-local, so the server cannot give a token's bytes on its own. +Only the M5's tensor units were checked; M1-M4 Macs and smaller RAM are untested. Tied heads, other head sizes, +Llama 3's scaled RoPE and biased projections are refused at load rather than served untested. diff --git a/zig/build/metal_targets.zig b/zig/build/metal_targets.zig index 07c2711739..8538b92b0a 100644 --- a/zig/build/metal_targets.zig +++ b/zig/build/metal_targets.zig @@ -327,6 +327,7 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s .{ .name = "tf-gemma4-bench", .path = "zig/tests/gemma4/bench.zig", .about = "Gemma 4 dense step costs: projections alone by rows against whole forwards" }, .{ .name = "tf-gemma4-no-tensor", .path = "zig/tests/gemma4/no_tensor.zig", .about = "Gemma 4's Metal kernels on a Mac without tensor units: what builds, and the row path's bits" }, .{ .name = "tf-spark-run", .path = "zig/tests/spark25/run.zig", .about = "Spark2.5 on Metal: greedy tokens, every position's logits or a prompt's last logits, for the reference's check" }, + .{ .name = "tf-llama-run", .path = "zig/tests/llama/run.zig", .about = "Llama on Metal: greedy tokens, every position's logits or a prompt's last logits, for the reference's check" }, .{ .name = "tf-qwen36-row-seam", .path = "zig/tests/qwen36_row_seam.zig", .about = "Model-free fixture: the row projection's bias sums stay fp32 at every width" }, }; for (engine_programs) |p| { diff --git a/zig/kernels/metal/llama/attention_nax.metal b/zig/kernels/metal/llama/attention_nax.metal new file mode 100644 index 0000000000..c2642c4896 --- /dev/null +++ b/zig/kernels/metal/llama/attention_nax.metal @@ -0,0 +1,204 @@ +// Llama's prompt attention on the tensor units (gemma4's prompt kernel over a [position, KVH, DH] cache). +#include "../nax.h" +using namespace tfp; + +// Defines: DH (128), NQ, G (query heads a key head), KVH, SCALE. Every layer attends to every key before. + +constant constexpr int LD = KVH * DH; +constant constexpr float kMasked = -3.402823466e+38f; +constant constexpr float kLog2e = 1.44269504089f; + +// Row max (MAX) or sum over a fragment's 32 columns into r[0], r[1] (rows home.y, home.y + 8). +template +inline void fold(thread const frag& f, thread float* r) { + TF_UNROLL + for (short h = 0; h < 2; h++) { + const short b = 4 * h; + float t = MAX ? max(max(f[b], f[b + 1]), max(f[b + 2], f[b + 3])) : (f[b] + f[b + 1]) + (f[b + 2] + f[b + 3]); + const float u = simd_shuffle_xor(t, ushort(1)); + t = MAX ? max(t, u) : t + u; + const float w = simd_shuffle_xor(t, ushort(8)); + t = MAX ? max(t, w) : t + w; + r[h] = MAX ? max(r[h], t) : r[h] + t; + } +} + +constant constexpr int RT = 32; // query rows a tile +constant constexpr int NS = DH / 64; // simdgroups a tile, each its own 64 output columns +constant constexpr int KS = 32; // keys a simdgroup a block +constant constexpr int BKT = KS * NS; // keys a block +constant constexpr int PH = BKT < 128 ? BKT : 128; // keys whose probabilities cross threadgroup memory at once +constant constexpr int PP = PH + 8; // the probabilities' row pitch +using PT = float; + +template +inline frag blk(const device bfloat* p, int ld, int r, int c, short2 home) { +#ifndef TF_SIMD_FRAGS + const device bfloat* q = p + (r + home.y) * ld + c + home.x; + return frag(*(const device bfloat4*)q, *(const device bfloat4*)(q + 8 * ld)); +#else + frag f; + if (TR) { + frag_get_t(f, p, ld, r, c, home); + } else { + frag_get(f, p, ld, r, c, home); + } + return f; +#endif +} + +inline frag prob_frag(const threadgroup PT* pt, int r, int c, short2 home) { +#ifndef TF_SIMD_FRAGS + const threadgroup PT* q = pt + (r + home.y) * PP + c + home.x; + return frag(*(const threadgroup vec*)q, *(const threadgroup vec*)(q + 8 * PP)); +#else + frag f; + frag_get(f, pt, PP, r, c, home); + return f; +#endif +} + +inline void prob_put(threadgroup PT* pt, int r, int c, short2 home, thread const frag& s) { +#ifndef TF_SIMD_FRAGS + threadgroup PT* q = pt + (r + home.y) * PP + c + home.x; + *(threadgroup vec*)q = vec(s[0], s[1], s[2], s[3]); + *(threadgroup vec*)(q + 8 * PP) = vec(s[4], s[5], s[6], s[7]); +#else + TF_UNROLL + for (short e = 0; e < 8; e++) pt[(r + home.y + (e >> 2) * 8) * PP + c + home.x + TF_COL(e)] = PT(s[e]); +#endif +} + +// Q [R, NQ, DH] at positions p0..; K, V the layer's cache; O [R, NQ, DH]. P: rows, p0. A tile is RT rows of a head. +[[max_total_threads_per_threadgroup(256)]] +[[kernel]] void ll_prompt_attention(const device bfloat* Q [[buffer(0)]], const device bfloat* K [[buffer(1)]], + const device bfloat* V [[buffer(2)]], device bfloat* O [[buffer(3)]], + constant int2& P [[buffer(4)]], + uint sg [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], uint3 tg [[threadgroup_position_in_grid]]) { + const int R = P.x, p0 = P.y; + const int qh = int(tg.y), kv = qh / G; + const int r0 = int(tg.x) * RT; + const int q_rows = min(RT, R - r0); + const int qs = NQ * DH; + const device bfloat* q = Q + r0 * qs + qh * DH; + const device bfloat* kc = K + kv * DH; + const device bfloat* vc = V + kv * DH + int(sg) * 64; + const short2 home = frag_home(ushort(lane)); + threadgroup PT pt[RT * PP]; + threadgroup float rowmax[NS][RT]; + frag acc[8]; + TF_UNROLL + for (short i = 0; i < 8; i++) acc[i] = frag(0); + float top[4] = {kMasked, kMasked, kMasked, kMasked}, total[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + const int first_pos = p0 + r0, last_pos = p0 + r0 + q_rows - 1; // the tile's last key + for (int kb = 0; kb <= last_pos / BKT; kb++) { + const int kpos = kb * BKT + int(sg) * KS; + const device bfloat* kblk = kc + long(kpos) * LD; + frag s[4]; + TF_UNROLL + for (short i = 0; i < 4; i++) s[i] = frag(0); + for (int d = 0; d < DH / 16; d++) { + const frag k0 = blk(kblk, LD, 0, 16 * d, home); + const frag k1 = blk(kblk, LD, 16, 16 * d, home); + TF_UNROLL + for (short h = 0; h < 2; h++) { // rows past q_rows read inside the chunk's buffer; their outputs stay unwritten + const frag qf = blk(q, qs, 16 * h, 16 * d, home); + mma_16x32(s[2 * h], s[2 * h + 1], qf, k0, k1); + } + } + TF_UNROLL + for (short i = 0; i < 4; i++) { + TF_UNROLL + for (short e = 0; e < 8; e++) { + const int local = 16 * (i >> 1) + home.y + (e >> 2) * 8, row = first_pos + local; + const int col = kpos + 16 * (i & 1) + home.x + TF_COL(e); + s[i][e] = col > row ? kMasked : s[i][e] * (SCALE * kLog2e); + } + } + float top_new[4] = {top[0], top[1], top[2], top[3]}; + TF_UNROLL + for (short i = 0; i < 4; i++) fold(s[i], top_new + 2 * (i >> 1)); + if (home.x == 0) { + TF_UNROLL + for (short j = 0; j < 4; j++) rowmax[sg][home.y + 8 * j] = top_new[j]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + TF_UNROLL + for (short j = 0; j < 4; j++) { + for (int t = 0; t < NS; t++) top_new[j] = max(top_new[j], rowmax[t][home.y + 8 * j]); + } + TF_UNROLL + for (short i = 0; i < 4; i++) { + TF_UNROLL + for (short e = 0; e < 8; e++) { + s[i][e] = s[i][e] == kMasked ? 0.0f : fast::exp2(s[i][e] - top_new[2 * (i >> 1) + (e >> 2)]); + } + } + float keep[4]; + TF_UNROLL + for (short j = 0; j < 4; j++) { + keep[j] = top[j] == kMasked ? 0.0f : fast::exp2(top[j] - top_new[j]); + top[j] = top_new[j]; + total[j] = total[j] * keep[j]; + } + TF_UNROLL + for (short i = 0; i < 4; i++) fold(s[i], total + 2 * (i >> 1)); + TF_UNROLL + for (short i = 0; i < 8; i++) { + TF_UNROLL + for (short e = 0; e < 8; e++) acc[i][e] = acc[i][e] * keep[2 * (i >> 2) + (e >> 2)]; + } + for (int ph = 0; ph < BKT / PH; ph++) { + if (ph > 0) threadgroup_barrier(mem_flags::mem_threadgroup); // the phase before has read the tile + if (int(sg) / (PH / KS) == ph) { + TF_UNROLL + for (short i = 0; i < 4; i++) prob_put(pt, 16 * (i >> 1), int(sg) * KS - ph * PH + 16 * (i & 1), home, s[i]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int kk = 0; kk < PH / 16; kk++) { + const int kp = kb * BKT + ph * PH + 16 * kk; + if (kp > last_pos) break; // keys past the tile's rows: masked, and their cache rows may hold anything + const int valid = min(16, last_pos + 1 - kp); // a value row never written would poison the tile as 0 x NaN + const device bfloat* vblk = vc + long(kp) * LD; + frag v0, v1; + TF_UNROLL + for (short d = 0; d < 4; d += 2) { + if (valid == 16) { + v0 = blk(vblk, LD, 0, 16 * d, home); + v1 = blk(vblk, LD, 0, 16 * d + 16, home); + } else { + frag_get_in(v0, vblk, LD, 0, 16 * d, home, valid, DH); + frag_get_in(v1, vblk, LD, 0, 16 * d + 16, home, valid, DH); + } + TF_UNROLL + for (short h = 0; h < 2; h++) { + const frag pf = prob_frag(pt, 16 * h, 16 * kk, home); + mma_16x32(acc[4 * h + d], acc[4 * h + d + 1], pf, v0, v1); + } + } + } + } + } + // the tile's totals: every simdgroup's partial sums, scaled alike by the shared maxes + threadgroup_barrier(mem_flags::mem_threadgroup); + if (home.x == 0) { + TF_UNROLL + for (short j = 0; j < 4; j++) rowmax[sg][home.y + 8 * j] = total[j]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float inv[4]; + TF_UNROLL + for (short j = 0; j < 4; j++) { + float t = 0.0f; + for (int u = 0; u < NS; u++) t += rowmax[u][home.y + 8 * j]; + inv[j] = 1.0f / t; + } + device bfloat* o = O + r0 * qs + qh * DH + int(sg) * 64; + TF_UNROLL + for (short i = 0; i < 8; i++) { + TF_UNROLL + for (short e = 0; e < 8; e++) acc[i][e] = acc[i][e] * inv[2 * (i >> 2) + (e >> 2)]; + frag_put_in(acc[i], o, qs, 16 * (i >> 2), 16 * (i & 3), home, q_rows, 64); + } +} diff --git a/zig/kernels/metal/llama/gemm.metal b/zig/kernels/metal/llama/gemm.metal new file mode 100644 index 0000000000..b392568e87 --- /dev/null +++ b/zig/kernels/metal/llama/gemm.metal @@ -0,0 +1,14 @@ +// Llama's prompt projections on the tensor units: bf16 x [M, K] by the release's W [N, K]; after gemm_nax.metal. + +#define LL_GEMM(NAME, AM, AN) \ + [[max_total_threads_per_threadgroup(256)]] [[kernel]] void NAME( \ + const device bfloat16_t* A [[buffer(0)]], const device bfloat16_t* B [[buffer(1)]], \ + const device int32_t* P [[buffer(2)]], device bfloat16_t* D [[buffer(3)]], \ + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], \ + uint3 tg [[threadgroup_position_in_grid]]) { \ + tfp::gemm_nax(A, B, D, P, tg, sg, lane); \ + } +LL_GEMM(ll_gemm_00, false, false) +LL_GEMM(ll_gemm_01, false, true) +LL_GEMM(ll_gemm_10, true, false) +LL_GEMM(ll_gemm_11, true, true) diff --git a/zig/kernels/metal/llama/glue.metal b/zig/kernels/metal/llama/glue.metal new file mode 100644 index 0000000000..6b4db9bcf1 --- /dev/null +++ b/zig/kernels/metal/llama/glue.metal @@ -0,0 +1,269 @@ +// Llama's glue at HF's rounding points (modeling_llama in bf16): bf16 residual and norm, rope with bf16 cos and sin. +#include +using namespace metal; + +constant constexpr uint DH = LL_DH; // the head dim (128), defined by the loader + +struct ll_args { uint rows, width, heads; float eps; }; + +// Row r of residual R = bf16 embedding row ids[r] (the release's table, [vocab, width]). +[[kernel]] void ll_embed(const device uint* ids [[buffer(0)]], const device bfloat* E [[buffer(1)]], + constant ll_args& a [[buffer(2)]], device bfloat* R [[buffer(3)]], uint2 p [[thread_position_in_grid]]) { + if (p.x >= a.width || p.y >= a.rows) return; + R[(size_t)p.y * a.width + p.x] = E[(size_t)ids[p.y] * a.width + p.x]; +} + +// LlamaRMSNorm on a row: out = bf16(w * bf16(x * rsqrt(mean(x^2) + eps))), the mean in fp32 over the bf16 row. +inline void ll_norm_row(const device bfloat* x, const device bfloat* w, constant ll_args& a, device bfloat* out, + float sum, uint t, uint n, threadgroup float* part) { + sum = simd_sum(sum); + if (t % 32 == 0) part[t / 32] = sum; + threadgroup_barrier(mem_flags::mem_threadgroup); + float total = 0.0f; + for (uint s = 0; s < n / 32; ++s) total += part[s]; + const float inv = precise::rsqrt(total / float(a.width) + a.eps); + for (uint i = t; i < a.width; i += n) out[i] = bfloat(float(w[i]) * float(bfloat(float(x[i]) * inv))); +} + +// The residual's rows normed by w into out. +[[kernel]] void ll_norm(const device bfloat* R [[buffer(0)]], const device bfloat* w [[buffer(1)]], + constant ll_args& a [[buffer(2)]], device bfloat* out [[buffer(3)]], uint row [[threadgroup_position_in_grid]], + uint t [[thread_position_in_threadgroup]], uint n [[threads_per_threadgroup]]) { + threadgroup float part[32]; + const device bfloat* x = R + (size_t)row * a.width; + float sum = 0.0f; + for (uint i = t; i < a.width; i += n) { + const float v = float(x[i]); + sum = fma(v, v, sum); + } + ll_norm_row(x, w, a, out + (size_t)row * a.width, sum, t, n, part); +} + +// R = bf16(R + X): the reference's bf16 residual add. +[[kernel]] void ll_add(device bfloat* R [[buffer(0)]], const device bfloat* X [[buffer(1)]], + constant ll_args& a [[buffer(2)]], uint2 p [[thread_position_in_grid]]) { + if (p.x >= a.width || p.y >= a.rows) return; + const size_t i = (size_t)p.y * a.width + p.x; + R[i] = bfloat(float(R[i]) + float(X[i])); +} + +// ll_add then ll_norm on a row, one threadgroup of 256. +[[kernel]] void ll_add_norm(device bfloat* R [[buffer(0)]], const device bfloat* X [[buffer(1)]], + const device bfloat* w [[buffer(2)]], constant ll_args& a [[buffer(3)]], device bfloat* out [[buffer(4)]], + uint row [[threadgroup_position_in_grid]], uint t [[thread_position_in_threadgroup]], + uint n [[threads_per_threadgroup]]) { + threadgroup float part[32]; + device bfloat* x = R + (size_t)row * a.width; + float sum = 0.0f; + for (uint i = t; i < a.width; i += n) { + const bfloat v = bfloat(float(x[i]) + float(X[(size_t)row * a.width + i])); + x[i] = v; + const float f = float(v); + sum = fma(f, f, sum); + } + ll_norm_row(x, w, a, out + (size_t)row * a.width, sum, t, n, part); +} + +struct ll_prep_args { uint rows, heads, kv_heads, pos, stride; }; + +// HF's apply_rotary_pos_emb, pairs (d, d + DH / 2): bf16(bf16(x * c) + bf16(o * s)), c and s the angle's cos and sin +[[kernel]] void ll_prep(const device bfloat* X [[buffer(0)]], const device float* inv [[buffer(1)]], + constant ll_prep_args& a [[buffer(2)]], device bfloat* QN [[buffer(3)]], device bfloat* KC [[buffer(4)]], + device bfloat* VC [[buffer(5)]], uint2 tg [[threadgroup_position_in_grid]], uint lane [[thread_index_in_simdgroup]]) { + // in bf16 as HF rounds them; q rotated to QN, k rotated into the cache at its position, v beside it + constexpr uint PER = DH / 32, HALF = DH / 2; + const uint unit = tg.x, row = tg.y; + if (unit >= a.heads + a.kv_heads || row >= a.rows) return; + const bool is_q = unit < a.heads; + const uint head = is_q ? unit : unit - a.heads; + const device bfloat* src = X + (size_t)row * a.stride + (is_q ? 0 : a.heads * DH) + (size_t)head * DH; + const size_t slot = a.pos + row; + const size_t kvw = (size_t)a.kv_heads * DH; + device bfloat* y = is_q ? QN + ((size_t)row * a.heads + head) * DH : KC + slot * kvw + (size_t)head * DH; + const float position = float(a.pos + row); + for (uint j = 0; j < PER; ++j) { + const uint d = lane + 32 * j; + const uint i = d % HALF; + const float angle = position * inv[i]; + const float c = float(bfloat(precise::cos(angle))), s = float(bfloat(precise::sin(angle))); + const float x = float(src[d]); + const float other = d < HALF ? -float(src[d + HALF]) : float(src[d - HALF]); + y[d] = bfloat(float(bfloat(x * c)) + float(bfloat(other * s))); + } + if (is_q) return; + const device bfloat* vs = X + (size_t)row * a.stride + (a.heads + a.kv_heads) * DH + (size_t)head * DH; + device bfloat* vd = VC + slot * kvw + (size_t)head * DH; + for (uint j = 0; j < PER; ++j) vd[lane + 32 * j] = vs[lane + 32 * j]; +} + +// torch's bf16 SiLU gate on a stacked gate|up row (gate at 0, up at a.width): A = bf16(bf16(silu(g)) * u). +[[kernel]] void ll_silu_gu(const device bfloat* GU [[buffer(0)]], constant ll_args& a [[buffer(1)]], + device bfloat* A [[buffer(2)]], uint2 p [[thread_position_in_grid]]) { + if (p.x >= a.width || p.y >= a.rows) return; + const size_t at = (size_t)p.y * 2 * a.width + p.x; + const float g = float(GU[at]), u = float(GU[at + a.width]); + const float silu = float(bfloat(g / (1.0f + precise::exp(-g)))); + A[(size_t)p.y * a.width + p.x] = bfloat(silu * u); +} + +#define TF_UNROLL_4 _Pragma("clang loop unroll(full)") + +struct ll_vec_args { uint rows, pos, heads, groups, blocks, chunk, first; float scale; }; + +// Decode attention partials (sdpa_vector lanes, DH / 32 dims a lane): a (key head, row, key block) a threadgroup. +[[kernel]] void ll_decode_vec(const device bfloat* Q [[buffer(0)]], const device bfloat* K [[buffer(1)]], + const device bfloat* V [[buffer(2)]], constant ll_vec_args& a [[buffer(3)]], device float* PM [[buffer(4)]], + device float* PL [[buffer(5)]], device float* PO [[buffer(6)]], uint3 tg [[threadgroup_position_in_grid]], + uint g [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]]) { + constexpr uint PER = DH / 32; + static_assert(PER == 4, "a lane holds one bfloat4 of a head"); + const uint kv = tg.x, row = tg.y, block = a.first + tg.z; + const uint G = a.heads / a.groups, head = kv * G + g; + const uint p = a.pos + row; + const uint from = block * a.chunk, to = min(p + 1, (block + 1) * a.chunk); + const size_t ld = (size_t)a.groups * DH; + const device bfloat* qp = Q + ((size_t)row * a.heads + head) * DH + PER * lane; + const float4 q0 = a.scale * float4(*(const device bfloat4*)qp); + float4 acc0 = 0.0f; + float top = -INFINITY, mass = 0.0f; + for (uint i0 = from; i0 < to; i0 += 4) { + float score[4]; + float4 v0[4]; + TF_UNROLL_4 + for (uint u = 0; u < 4; ++u) { + const bool live = i0 + u < to; // past the block's end: a masked key (read at the block's first, weight 0) + const uint i = live ? i0 + u : from; + const size_t at = (size_t)i * ld + kv * DH + PER * lane; + v0[u] = float4(*(const device bfloat4*)(V + at)); + const float sc = simd_sum(dot(q0, float4(*(const device bfloat4*)(K + at)))); + score[u] = live ? sc : -INFINITY; + } + const float peak = max(max(top, max(score[0], score[1])), max(score[2], score[3])); + const float rescale = fast::exp(top - peak); + mass *= rescale; + acc0 *= rescale; + TF_UNROLL_4 + for (uint u = 0; u < 4; ++u) { + const float w = fast::exp(score[u] - peak); + mass += w; + acc0 += w * v0[u]; + } + top = peak; + } + const size_t slot = ((size_t)head * a.rows + row) * a.blocks + tg.z; + if (lane == 0) { PM[slot] = top; PL[slot] = mass; } + *(device float4*)(PO + slot * DH + PER * lane) = acc0; +} + +// MLX's sdpa_vector_2pass_2: 32 simdgroups a (head, row) fold the blocks, DH / 32 output dims a simdgroup. +[[kernel]] void ll_merge_vec(const device float* PM [[buffer(0)]], const device float* PL [[buffer(1)]], + const device float* PO [[buffer(2)]], constant ll_vec_args& a [[buffer(3)]], device bfloat* O [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], uint sg [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + constexpr uint PER = DH / 32; + threadgroup float merged[32 * 32]; + const uint head = tg.x, row = tg.y; + const size_t at = (size_t)head * a.rows + row; + const device float* ms = PM + at * a.blocks; + const device float* ls = PL + at * a.blocks; + // lanes and simdgroups take blocks by absolute index (first + b): a row's sums never move with the window's first row + const uint lane0 = (lane + 32 - a.first % 32) % 32, sg0 = (sg + 32 - a.first % 32) % 32; + float top = -INFINITY; + for (uint b = lane0; b < a.blocks; b += 32) top = max(top, ms[b]); + top = simd_max(top); + float total = 0.0f; + for (uint b = lane0; b < a.blocks; b += 32) total += ls[b] > 0.0f ? ls[b] * fast::exp(ms[b] - top) : 0.0f; + total = simd_sum(total); + float acc[PER]; + for (uint j = 0; j < PER; ++j) acc[j] = 0.0f; + for (uint b = sg0; b < a.blocks; b += 32) { + if (ls[b] <= 0.0f) continue; + const float f = fast::exp(ms[b] - top); + for (uint j = 0; j < PER; ++j) acc[j] += f * PO[(at * a.blocks + b) * DH + PER * lane + j]; + } + for (uint j = 0; j < PER; ++j) { + merged[lane * 32 + sg] = acc[j]; + threadgroup_barrier(mem_flags::mem_threadgroup); + acc[j] = simd_sum(merged[sg * 32 + lane]) / total; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (lane == 0) + for (uint j = 0; j < PER; ++j) O[(row * a.heads + head) * DH + PER * sg + j] = bfloat(acc[j]); +} + +inline bool ll_better(float a, uint ia, float b, uint ib) { + if (isnan(a)) return !isnan(b) || ia < ib; + if (isnan(b)) return false; + return a > b || (a == b && ia < ib); +} + +// The greedy draw (nemotron_glue's tf_argmax_bf16): a bf16 logits row's argmax, ties to the lower id, into ids[0]. +[[kernel]] void ll_argmax(const device bfloat* logits [[buffer(0)]], device uint* ids [[buffer(1)]], + constant uint& vocab [[buffer(2)]], uint t [[thread_position_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], uint sg [[simdgroup_index_in_threadgroup]]) { + threadgroup float vals[32]; + threadgroup uint idxs[32]; + float best = -INFINITY; + uint at = 0xffffffffu; + for (uint i = t; i < vocab; i += 1024) { + const float v = float(logits[i]); + if (ll_better(v, i, best, at)) { best = v; at = i; } + } + for (ushort d = 16; d > 0; d >>= 1) { + const float ov = simd_shuffle_xor(best, d); + const uint oi = simd_shuffle_xor(at, d); + if (ll_better(ov, oi, best, at)) { best = ov; at = oi; } + } + if (lane == 0) { vals[sg] = best; idxs[sg] = at; } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg != 0) return; + best = vals[lane]; + at = idxs[lane]; + for (ushort d = 16; d > 0; d >>= 1) { + const float ov = simd_shuffle_xor(best, d); + const uint oi = simd_shuffle_xor(at, d); + if (ll_better(ov, oi, best, at)) { best = ov; at = oi; } + } + if (lane == 0) ids[0] = at; +} + +struct ll_mv_args { uint rows, n, k, ldy; }; + +// Y[r, j] = bf16(sum_i X[r, i] * W[j, i]) over bf16 W [n, k] in fp32 (torch's bf16 linear): a simdgroup a row j. +template +inline void ll_matvec_t(const device bfloat* X, const device bfloat* W, constant ll_mv_args& a, device bfloat* Y, + uint3 tg, uint sg, uint lane) { + constexpr uint SG = 8; + const uint j = tg.x * SG + sg; + if (j >= a.n) return; + const device bfloat* w = W + (size_t)j * a.k; + float acc[ROWS]; + TF_UNROLL_4 + for (uint r = 0; r < ROWS; ++r) acc[r] = 0.0f; + for (uint i = 8 * lane; i < a.k; i += 256) { + const float4 w0 = float4(*(const device bfloat4*)(w + i)), w1 = float4(*(const device bfloat4*)(w + i + 4)); + TF_UNROLL_4 + for (uint r = 0; r < ROWS; ++r) { + const device bfloat* x = X + (size_t)r * a.k + i; + acc[r] += dot(w0, float4(*(const device bfloat4*)x)) + dot(w1, float4(*(const device bfloat4*)(x + 4))); + } + } + TF_UNROLL_4 + for (uint r = 0; r < ROWS; ++r) { + const float s = simd_sum(acc[r]); + if (lane == 0 && r < a.rows) Y[(size_t)r * a.ldy + j] = bfloat(s); + } +} + +#define LL_MATVEC(ROWS) \ + [[kernel]] void ll_matvec_##ROWS(const device bfloat* X [[buffer(0)]], const device bfloat* W [[buffer(1)]], \ + constant ll_mv_args& a [[buffer(2)]], device bfloat* Y [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]], \ + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]]) { \ + ll_matvec_t(X, W, a, Y, tg, sg, lane); \ + } +LL_MATVEC(1) +LL_MATVEC(2) +LL_MATVEC(4) +LL_MATVEC(8) +LL_MATVEC(16) +LL_MATVEC(32) diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig index 2506fb4905..e8deb1979c 100644 --- a/zig/kernels/metal/sources.zig +++ b/zig/kernels/metal/sources.zig @@ -176,6 +176,12 @@ pub const spark25_glue = @embedFile("spark25/glue.metal"); pub const spark25_attention_nax = @embedFile("spark25/attention_nax.metal"); /// Spark2.5's bf16 prompt projections on gemm_nax's tensor-unit GEMM (nax.h inlined at load). pub const spark25_gemm = @embedFile("prefill/gemm_nax.metal") ++ @embedFile("spark25/gemm.metal"); +/// Llama's glue at HF's rounding points: embed, a bf16 residual and its norm, rope with bf16 cos and sin, SiLU. +pub const llama_glue = @embedFile("llama/glue.metal"); +/// Llama's prompt attention on the tensor units (gemma4's prompt kernel, every key; nax.h inlined at load). +pub const llama_attention_nax = @embedFile("llama/attention_nax.metal"); +/// Llama's bf16 prompt projections on gemm_nax's tensor-unit GEMM (nax.h inlined at load). +pub const llama_gemm = @embedFile("prefill/gemm_nax.metal") ++ @embedFile("llama/gemm.metal"); /// The kernel contract's planted controls: one kernel breaking each gate (zig/tests/contract/). pub const core_contract_controls = @embedFile("core/contract_controls.metal"); diff --git a/zig/src/cli/hub.zig b/zig/src/cli/hub.zig index fe4602429d..e5c732e01b 100644 --- a/zig/src/cli/hub.zig +++ b/zig/src/cli/hub.zig @@ -18,6 +18,7 @@ const titles = [_]struct { model_type: []const u8, title: []const u8 }{ .{ .model_type = "gemma4_unified", .title = "Gemma 4 12B" }, .{ .model_type = "step5_vl", .title = "Step-5-Preview" }, .{ .model_type = "spark2_5", .title = "Spark-X2.6-4B" }, + .{ .model_type = "llama", .title = "Llama architecture (Nanbeige4.3-33B)" }, .{ .model_type = "lfm2_moe", .title = "LFM2.5-24B-A2B" }, }; diff --git a/zig/src/families/llama/backend.zig b/zig/src/families/llama/backend.zig new file mode 100644 index 0000000000..fd1f17b9c4 --- /dev/null +++ b/zig/src/families/llama/backend.zig @@ -0,0 +1,406 @@ +//! Llama on Metal behind the lane engine: windows a stream, a cache each; lone steps run ahead, drawn on the GPU. +const std = @import("std"); +const mtl = @import("metal"); +const lanes = @import("lanes"); +const st = @import("state.zig"); +const fwd = @import("forward.zig"); +const Model = @import("model.zig").Model; +const snapshot = @import("snapshot.zig"); +const row = @import("../step5/backend.zig").row; +const gpu_full = lanes.gpu_full; +const be = lanes.backend; + +/// The longest context a stream may take, the streams at once, and the bytes every stream's caches may hold. +pub const Options = struct { capacity: usize, streams: u32 = 4, budget: u64 }; + +/// Token slots a stream's handles point at: drawn tokens and host-fed ones (no command in flight reads a fresh one). +const ring = 4096; + +/// A committed step: the token slots below `end` are written when `cb` completes. +const Flight = struct { end: u64, cb: mtl.CommandBuffer }; + +/// The prompt cache's room beside the streams: its kept states fit in `bytes`, never evicting `restoring`. +pub const Room = struct { ctx: *anyopaque, fit: *const fn (ctx: *anyopaque, bytes: u64, restoring: ?*anyopaque, prompt: []const u32) void }; + +pub const Metal = struct { + gpa: std.mem.Allocator, + model: *Model, + options: Options, + scratch: st.Scratch, // decode rounds + prompt: st.Scratch, // prompt chunks + caches: std.AutoHashMapUnmanaged(*lanes.Stream, *st.Cache) = .empty, + follow: ?mtl.ResidencySet.Follow = null, + held: u64 = 0, // bytes the streams' caches hold + room: ?Room = null, // the prompt cache, which takes what the streams leave of the budget + tokens: mtl.Buffer, // u32 [ring] + next: u64 = 0, // the next free token slot + landed: u64 = 0, // slots below it are written + flights: std.ArrayList(Flight) = .empty, + values: []f32, // a logits row as f32, for the host twin of the GPU draw + + pub fn init(gpa: std.mem.Allocator, model: *Model, options: Options) !*Metal { + const out = try gpa.create(Metal); + errdefer gpa.destroy(out); + var scratch = try st.Scratch.init(model.device, &model.config, st.batch_rows); + errdefer scratch.deinit(); + var prompt = try st.Scratch.init(model.device, &model.config, st.prompt_rows); + errdefer prompt.deinit(); + const tokens = try model.device.buffer(ring * 4, mtl.ResourceOptions.shared); + errdefer tokens.deinit(); + out.* = .{ .gpa = gpa, .model = model, .options = options, .scratch = scratch, .prompt = prompt, .tokens = tokens, .values = try gpa.alloc(f32, model.config.vocab) }; + return out; + } + + pub fn deinit(self: *Metal) void { + self.drain() catch {}; + self.flights.deinit(self.gpa); + self.tokens.deinit(); + self.gpa.free(self.values); + self.setResident(null); + var it = self.caches.valueIterator(); + while (it.next()) |cache| { + cache.*.deinit(); + self.gpa.destroy(cache.*); + } + self.caches.deinit(self.gpa); + self.scratch.deinit(); + self.prompt.deinit(); + self.gpa.destroy(self); + } + + pub fn setResident(self: *Metal, hook: ?mtl.ResidencySet.Follow) void { + if (self.follow) |f| { + mtl.resident.fields.apply(&self.scratch, f, false); + mtl.resident.fields.apply(&self.prompt, f, false); + } + self.follow = hook; + if (hook) |f| { + mtl.resident.fields.apply(&self.scratch, f, true); + mtl.resident.fields.apply(&self.prompt, f, true); + } + var it = self.caches.valueIterator(); + while (it.next()) |cache| cache.*.setResident(hook); + } + + pub fn createCache(self: *Metal, capacity: usize) !*st.Cache { + const cache = try self.gpa.create(st.Cache); + errdefer self.gpa.destroy(cache); + cache.* = try st.Cache.init(self.model.device, &self.model.config, capacity); + cache.setResident(self.follow); + return cache; + } + + pub fn facts(self: *const Metal) lanes.Model { + // rows are independent in every decode kernel, so any width, alone or beside other streams, gives one-row bits + return .{ .exact_width = st.window_rows, .gpu_tokens = true, .hidden_rows = true, .batch_rows = st.batch_rows, .mtp = false, .speculate = false, .speculate_early = false, .drafts = 1, .max_streams = self.options.streams, .streams_exact = true }; + } + + pub fn backend(self: *Metal) be.Backend { + return .{ .ptr = self, .vtable = &.{ .prefill = prefillFn, .first = firstFn, .first_masked = firstMaskedFn, .queue = queueFn, .read = readFn, .verify = verifyFn, .keep = keepFn, .draft = draftFn, .release = releaseFn, .first_row = firstRowFn, .held = heldFn } }; + } + + fn cast(ptr: *anyopaque) *Metal { + return @ptrCast(@alignCast(ptr)); + } + + fn cacheOf(self: *Metal, stream: *lanes.Stream) !*st.Cache { + return self.caches.get(stream) orelse error.UnknownLlamaStream; + } + + fn logits(self: *Metal, cache: *st.Cache) []const u16 { + return cache.logits.slice(u16, self.model.config.vocab); + } + + /// The GPU draw's host twin: argmax, or tf_sample_full's keyed race at `position` (gpu_full), a mask's tokens only. + fn draw(self: *Metal, l: []const u16, sampling: ?lanes.Sampling, position: u64, mask: []const u32) !u32 { + const t = sampled(sampling) orelse { // gpu_rule's argmax: the first largest, masked tokens at -inf + var best: ?usize = null; + for (l, 0..) |w, id| { + if (mask.len > 0 and !lanes.grammar.allows(mask, id)) continue; + if (best == null or logit(w) > logit(l[best.?])) best = id; + } + return @intCast(best orelse 0); + }; + for (l, self.values, 0..) |w, *v, id| { + v.* = logit(w); + if (mask.len > 0 and !lanes.grammar.allows(mask, id)) v.* = -std.math.inf(f32); + } + return gpu_full.sample(self.gpa, self.values, t, @intCast(position), null); + } + + fn sampled(sampling: ?lanes.Sampling) ?lanes.Sampling { + const t = sampling orelse return null; + return if (t.temperature > 0) t else null; + } + + /// `n` consecutive token slots, never wrapping the ring. + fn take(self: *Metal, n: usize) u64 { + if (self.next % ring + n > ring) self.next += ring - self.next % ring; + const h = self.next; + self.next += n; + return h; + } + + fn slotOf(self: *Metal, h: u64) fwd.B { + return .{ .buf = self.tokens, .off = @intCast(h % ring * 4) }; + } + + /// Wait for the oldest step in flight. + fn land(self: *Metal) !void { + const f = self.flights.orderedRemove(0); + defer f.cb.release(); + f.cb.wait(); + if (f.cb.failure()) |msg| { + std.log.err("llama step: {s}", .{msg}); + return error.GpuFailed; + } + self.landed = @max(self.landed, f.end); + } + + /// Every step in flight landed: before the host reads caches or logits, or the GPU runs anything else. + fn drain(self: *Metal) !void { + while (self.flights.items.len > 0) try self.land(); + } + + /// Run `ids` at the cache's end in prompt chunks; a stream's pass stops at its chunk starts and reuse marks. + fn feed(self: *Metal, cache: *st.Cache, ids: []const u32, stream: ?*lanes.Stream) !void { + const base = cache.len; + var at: usize = 0; + while (at < ids.len) { + var end = @min(ids.len, at + st.prompt_rows); + if (stream) |sm| { + if (sm.isCancelled()) return error.Cancelled; + for (sm.chunks) |x| if (x > base + at and x < base + end) { + end = x - base; + }; + for (sm.reuse.marks) |x| if (x > base + at and x < base + end) { + end = x - base; + }; + } + const s = if (end - at > st.batch_rows) &self.prompt else &self.scratch; + try fwd.run(self.model, s, cache, ids[at..end], .{}); + cache.len += end - at; + at = end; + if (stream) |sm| if (std.mem.indexOfScalar(u32, sm.reuse.marks, @intCast(cache.len)) != null) if (sm.reuse.hook) |k| k.at(k.ptr, sm, @intCast(cache.len)); + } + } + + fn prefillFn(ptr: *anyopaque, stream: *lanes.Stream) !void { + const self = cast(ptr); + try self.drain(); + const ids = stream.prompt(); + if (ids.len == 0 or ids.len + stream.max_new > self.options.capacity or self.caches.contains(stream)) return error.InvalidLlamaPrompt; + if (stream.images.len > 0) return error.LlamaHasNoImageTower; + const capacity = ids.len + stream.max_new + 1; + const need = st.cacheBytes(&self.model.config, capacity); + if (self.held + need > self.options.budget) { + std.log.warn("llama: a {d}-token stream needs {d:.1} GiB of caches; {d:.1} of {d:.1} GiB are held", .{ capacity, gib(need), gib(self.held), gib(self.options.budget) }); + return error.LaneMemory; // the lane host queues the stream until another finishes, or fails it alone + } + self.held += need; + errdefer self.held -= need; + if (stream.reuse.saved) |saved| { // a state the stream cannot hold beside its caches: freed now, no restore + const snap: *snapshot.Snap = @ptrCast(@alignCast(saved)); + if (self.held + snap.bytes > self.options.budget) snap.free(); + } + // kept states give way first + if (self.room) |r| r.fit(r.ctx, self.options.budget - self.held, stream.reuse.saved, ids); + const cache = try self.createCache(capacity); + errdefer self.gpa.destroy(cache); + errdefer cache.deinit(); + try self.caches.put(self.gpa, stream, cache); // a mark's save finds it there + errdefer _ = self.caches.remove(stream); + stream.cached = 0; + if (stream.reuse.saved) |saved| { // a kept state of this prompt's prefix: the pass starts there + const snap: *const snapshot.Snap = @ptrCast(@alignCast(saved)); + if (snap.at < ids.len) { + if (snapshot.restore(self.model.queue, snap, cache)) stream.cached = snap.at else |_| stream.reuse_failed = true; + } else stream.reuse_failed = true; + } + try self.feed(cache, ids[cache.len..], stream); + } + + /// A prompt chunk and a decode window once: the first request skips the GPU's first touch. + pub fn warm(self: *Metal) !void { + var cache = try st.Cache.init(self.model.device, &self.model.config, st.prompt_rows + st.window_rows); + defer cache.deinit(); + var ids: [st.prompt_rows]u32 = undefined; + for (&ids, 0..) |*t, i| t.* = @intCast((1000 + i * 37) % self.model.config.vocab); + try self.feed(&cache, &ids, null); + try fwd.run(self.model, &self.scratch, &cache, ids[0..st.window_rows], .{ .all = true }); + } + + /// The first token, drawn on the host from the prompt's last row into a fresh slot. + fn firstFn(ptr: *anyopaque, stream: *lanes.Stream, position: u64) !u64 { + return cast(ptr).firstSlot(stream, position, &.{}); + } + + fn firstMaskedFn(ptr: *anyopaque, stream: *lanes.Stream, position: u64, mask: []const u32) !u64 { + return cast(ptr).firstSlot(stream, position, mask); + } + + fn firstSlot(self: *Metal, stream: *lanes.Stream, position: u64, mask: []const u32) !u64 { + try self.drain(); + const cache = try self.cacheOf(stream); + if (position != cache.len) return error.InvalidLlamaPosition; + const h = self.take(1); + self.tokens.slice(u32, ring)[h % ring] = try self.draw(self.logits(cache), stream.sampling, position, mask); + return h; + } + + /// A step behind those in flight: its token from the ring or the host, the next drawn into a slot; read() lands it. + fn queueFn(ptr: *anyopaque, stream: *lanes.Stream, f: be.Feed, position: u64) !u64 { + const self = cast(ptr); + const cache = try self.cacheOf(stream); + if (position != cache.len + 1) return error.InvalidLlamaPosition; + const in = switch (f) { + .handle => |h| h, + .value => |v| blk: { + if (v >= self.model.config.vocab) return error.InvalidLlamaToken; + const h = self.take(1); + self.tokens.slice(u32, ring)[h % ring] = v; + break :blk h; + }, + }; + const out = self.take(1); + const how: @FieldType(fwd.Draw, "how") = if (sampled(stream.sampling)) |t| .{ .sampled = .{ .sampling = t, .position = position } } else .greedy; + var cb: mtl.CommandBuffer = undefined; + try fwd.run(self.model, &self.scratch, cache, &.{0}, .{ .ids = self.slotOf(in), .draw = .{ .how = how, .out = self.slotOf(out) }, .flight = &cb }); + self.flights.append(self.gpa, .{ .end = out + 1, .cb = cb }) catch |err| { + cb.wait(); + cb.release(); + return err; + }; + cache.len += 1; + return out; + } + + fn firstRowFn(ptr: *anyopaque, stream: *lanes.Stream) !lanes.LogprobRow { + const self = cast(ptr); + try self.drain(); + const l = self.logits(try self.cacheOf(stream)); + return row(l, @intCast(argmax(l)), stream.logprobs orelse 0); + } + + /// A slot's token, once the step that writes it has landed (host-fed slots are written at once). + fn readFn(ptr: *anyopaque, handle: u64) !u32 { + const self = cast(ptr); + while (handle >= self.landed and self.flights.items.len > 0) try self.land(); + const token = self.tokens.slice(u32, ring)[handle % ring]; + if (token >= self.model.config.vocab) return error.InvalidLlamaToken; + return token; + } + + /// Every window's pending token and drafts in one shared forward, every row drawn and kept until `keep` rolls back. + fn verifyFn(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) !void { + const self = cast(ptr); + const vocab = self.model.config.vocab; + const words = (vocab + 31) / 32; + if (windows.len != out.len or windows.len > st.batch_rows) return error.InvalidLlamaWindows; + try self.drain(); + var segments: [st.batch_rows]fwd.Segment = undefined; + var ids: [st.batch_rows]u32 = undefined; + var total: usize = 0; + for (windows, out, 0..) |w, o, i| { + const rows = w.rows(); + if (rows > st.window_rows or total + rows > st.batch_rows or w.parents != null or w.held != 0 or o.sampled.len != rows or o.drafts.len != rows - 1) return error.UnsupportedLlamaWindow; + const cache = try self.cacheOf(w.stream); + for (w.positions, 0..) |p, r| if (p != cache.len + r + 1) return error.InvalidLlamaPosition; + segments[i] = .{ .cache = cache, .rows = rows }; + ids[total] = w.pending; + @memcpy(ids[total + 1 ..][0..w.tokens.len], w.tokens); + total += rows; + } + try fwd.runShared(self.model, &self.scratch, segments[0..windows.len], ids[0..total], .{ .all = true }); + const all = self.scratch.logits_all.slice(u16, total * vocab); + var base: usize = 0; + for (windows, out, segments[0..windows.len]) |w, o, seg| { + const rows = seg.rows; + seg.cache.len += rows; + @memcpy(o.drafts, w.tokens); + for (0..rows) |r| { + const l = all[(base + r) * vocab ..][0..vocab]; + const mask = if (w.masks.len >= (r + 1) * words) w.masks[r * words ..][0..words] else &.{}; + o.sampled[r] = try self.draw(l, w.stream.sampling, w.positions[r], mask); + if (o.rows.len == rows) o.rows[r] = row(l, o.sampled[r], w.stream.logprobs orelse 0); + } + // the last kept row's logits are the next first draw's (the lane core keeps every row here) + @memcpy(seg.cache.logits.slice(u16, vocab), all[(base + rows - 1) * vocab ..][0..vocab]); + base += rows; + } + } + + fn keepFn(ptr: *anyopaque, windows: []const be.Window, paths: []const []const u32) !void { + const self = cast(ptr); + try self.drain(); + for (windows, paths) |w, path| { + const cache = try self.cacheOf(w.stream); + for (path, 0..) |r, i| if (r != i) return error.UnsupportedLlamaTree; + if (path.len == 0 or path.len > w.rows() or w.rows() > cache.len) return error.UnsupportedLlamaWindow; + cache.len = cache.len - w.rows() + path.len; + } + } + + /// No draft head: a request at depth 0 holds nothing; the lane core's copy drafts come as window tokens. + fn draftFn(ptr: *anyopaque, requests: []const be.DraftRequest) !void { + _ = ptr; + for (requests) |r| if (r.depth != 0) return error.LlamaHasNoDraftHead; + } + + fn heldFn(_: *anyopaque, _: *lanes.Stream, out: []u32) !void { + if (out.len > 0) return error.LlamaDraftsNotHeld; + } + + fn releaseFn(ptr: *anyopaque, stream: *lanes.Stream) void { + const self = cast(ptr); + self.drain() catch |err| std.log.err("llama: a step in flight failed at release: {s}", .{@errorName(err)}); + if (self.caches.fetchRemove(stream)) |entry| { + self.held -= st.cacheBytes(&self.model.config, entry.value.capacity); + entry.value.deinit(); + self.gpa.destroy(entry.value); + if (self.room) |r| r.fit(r.ctx, self.options.budget - self.held, null, &.{}); + } + } + + /// The prompt cache's Snapshots functions: a stream's caches at a mark, copied on the GPU. + pub fn snapBytes(ptr: *anyopaque, at: u32) u64 { + return snapshot.bytes(&cast(ptr).model.config, at); + } + + pub fn snapSave(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!*anyopaque { + const self = cast(ptr); + try self.drain(); + const cache = try self.cacheOf(@ptrCast(@alignCast(owner orelse return error.NoStream))); + if (cache.len != at) return error.SnapshotOutOfStep; + const snap = try self.gpa.create(snapshot.Snap); + errdefer self.gpa.destroy(snap); + snap.* = try snapshot.save(self.model.device, self.model.queue, &self.model.config, cache); + return snap; + } + + pub fn snapRestore(_: *anyopaque, _: ?*anyopaque, _: *anyopaque) anyerror!void { + return error.BackendRestores; // the prompt pass restores, before its first chunk + } + + pub fn snapDrop(ptr: *anyopaque, saved: *anyopaque) void { + const snap: *snapshot.Snap = @ptrCast(@alignCast(saved)); + snap.free(); + cast(ptr).gpa.destroy(snap); + } +}; + +fn gib(bytes: u64) f64 { + return @as(f64, @floatFromInt(bytes)) / (1 << 30); +} + +fn logit(word: u16) f32 { + return @bitCast(@as(u32, word) << 16); +} + +fn argmax(l: []const u16) usize { + var best: usize = 0; + for (l, 0..) |w, i| if (logit(w) > logit(l[best])) { + best = i; + }; + return best; +} diff --git a/zig/src/families/llama/budget.zig b/zig/src/families/llama/budget.zig new file mode 100644 index 0000000000..d518cb0f70 --- /dev/null +++ b/zig/src/families/llama/budget.zig @@ -0,0 +1,46 @@ +//! What a Llama load on Metal may hold: weights, scratch and every stream's caches, below the panic line. +const std = @import("std"); +const cfg = @import("config.zig"); +const st = @import("state.zig"); + +/// The default share of RAM: past about 80% wired, macOS stalls until its watchdog restarts the Mac. +pub const default_share = 0.75; +/// Kept free inside the limit for the command buffers, the keepalive and the host. +pub const margin: u64 = 2 << 30; + +/// Bytes for caches: `ram` times the share (or LLAMA_LOAD_LIMIT_GB's GB, at most the GPU's working set), less the rest. +pub fn caches(ram: u64, working_set: u64, text: ?[]const u8, resident: u64) error{BadLoadLimit}!u64 { + var limit: u64 = @intFromFloat(@as(f64, @floatFromInt(ram)) * default_share); + const t = std.mem.trim(u8, text orelse "", " \t"); + if (t.len > 0) { + const gb = std.fmt.parseFloat(f64, t) catch return error.BadLoadLimit; + if (!std.math.isFinite(gb) or gb <= 0) return error.BadLoadLimit; + limit = @min(@as(u64, @intFromFloat(gb * 1e9)), @min(ram, working_set)); + } + return if (limit > resident + margin) limit - resident - margin else 0; +} + +/// The longest window one stream's caches fit in `bytes`, at most `most`: a stream holds its window and one slot more. +pub fn longest(c: *const cfg.Config, bytes: u64, most: usize) usize { + if (st.cacheBytes(c, 2) > bytes) return 0; + var lo: usize = 1; + var hi: usize = most; + while (lo < hi) { + const mid = lo + (hi - lo + 1) / 2; + if (st.cacheBytes(c, mid + 1) <= bytes) lo = mid else hi = mid - 1; + } + return lo; +} + +test "llama budget: a share of RAM less the weights, a set limit capped at the working set, bad text refused" { + const gib: u64 = 1 << 30; + try std.testing.expectEqual(@as(u64, 192 * gib - 62 * gib - margin), try caches(256 * gib, 240 * gib, null, 62 * gib)); + try std.testing.expectEqual(@as(u64, 0), try caches(256 * gib, 240 * gib, null, 191 * gib)); + try std.testing.expectEqual(@as(u64, 128 * gib - 62 * gib - margin), try caches(256 * gib, 240 * gib, "137.438953472", 62 * gib)); + try std.testing.expectError(error.BadLoadLimit, caches(256 * gib, 240 * gib, "lots", 62 * gib)); + const c = try cfg.parse(std.testing.allocator, @embedFile("fixtures/config.json")); + try std.testing.expectEqual(@as(usize, 4095), longest(&c, st.cacheBytes(&c, 4096), 1 << 20)); // 4,096 + 1 needs a block more + try std.testing.expectEqual(@as(usize, 4159), longest(&c, st.cacheBytes(&c, 4097), 1 << 20)); + try std.testing.expectEqual(@as(usize, 4096), longest(&c, st.cacheBytes(&c, 4097), 4096)); + try std.testing.expectEqual(@as(usize, 0), longest(&c, 1000, 1 << 20)); +} diff --git a/zig/src/families/llama/config.zig b/zig/src/families/llama/config.zig new file mode 100644 index 0000000000..e7631560b6 --- /dev/null +++ b/zig/src/families/llama/config.zig @@ -0,0 +1,241 @@ +//! Llama-architecture (model_type llama) admission: a dense bf16 GQA stack, heads of 128, default RoPE, its own head. +const std = @import("std"); + +pub const head_dim = 128; +pub const max_layers = 128; +pub const max_eos = 4; +/// Query heads a key head at most: a decode attention threadgroup holds a simdgroup each. +pub const max_group = 32; + +pub const Error = error{ + BadLlamaConfig, + BadLlamaContext, + BadLlamaEos, + UnsupportedLlamaGeometry, + UnsupportedLlamaRotary, + UnsupportedLlamaPrecision, +} || std.json.ParseError(std.json.Scanner); + +pub const Config = struct { + hidden: usize, + intermediate: usize, + layers: usize, + heads: usize, + kv_heads: usize, + vocab: usize, + eps: f32, + theta: f64, + context: usize, + eos: [max_eos]u32 = @splat(0), + eos_count: usize = 0, + + pub fn read(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Config { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 22)); + defer gpa.free(bytes); + return parse(gpa, bytes); + } + + pub fn queryWidth(c: Config) usize { + return c.heads * head_dim; + } + + pub fn kvWidth(c: Config) usize { + return c.kv_heads * head_dim; + } + + /// The q, k and v projections' outputs side by side in a row. + pub fn qkvWidth(c: Config) usize { + return c.queryWidth() + 2 * c.kvWidth(); + } + + /// Query heads a key head. + pub fn group(c: Config) usize { + return c.heads / c.kv_heads; + } + + /// 1 / sqrt(head_dim), the attention scale. + pub fn scale(_: Config) f32 { + return 1.0 / @sqrt(@as(f32, head_dim)); + } + + pub fn isEos(c: Config, token: u32) bool { + for (c.eos[0..c.eos_count]) |e| if (e == token) return true; + return false; + } + + /// RoPE inverse frequencies [head_dim / 2] as fp32: HF's 1 / theta^(2i / head_dim). + pub fn inverseFrequencies(c: Config, out: []f32) void { + for (out, 0..) |*f, i| { + const e: f64 = @as(f64, @floatFromInt(2 * i)) / @as(f64, head_dim); + f.* = @floatCast(1.0 / std.math.pow(f64, c.theta, e)); + } + } +}; + +fn object(v: std.json.Value) Error!std.json.ObjectMap { + return if (v == .object) v.object else error.BadLlamaConfig; +} + +fn number(v: std.json.Value) Error!f64 { + return switch (v) { + .integer => |i| @floatFromInt(i), + .float => |f| f, + else => error.BadLlamaConfig, + }; +} + +fn integer(o: std.json.ObjectMap, key: []const u8) Error!i64 { + const v = o.get(key) orelse return error.BadLlamaConfig; + return if (v == .integer) v.integer else error.BadLlamaConfig; +} + +/// `key`'s integer, or `default` when absent. +fn integerOr(o: std.json.ObjectMap, key: []const u8, default: i64) Error!i64 { + if (o.get(key) == null) return default; + return integer(o, key); +} + +fn isString(v: ?std.json.Value, expected: []const u8) bool { + const s = v orelse return false; + return s == .string and std.mem.eql(u8, s.string, expected); +} + +/// `key` is true (absent: false). +fn isTrue(o: std.json.ObjectMap, key: []const u8) bool { + const v = o.get(key) orelse return false; + return v == .bool and v.bool; +} + +/// A positive dimension. +fn size(o: std.json.ObjectMap, key: []const u8) Error!usize { + const v = try integer(o, key); + return if (v > 0) @intCast(v) else error.UnsupportedLlamaGeometry; +} + +/// rope_scaling: absent, null, or the default type by either name. +fn rotary(o: std.json.ObjectMap) Error!void { + const v = o.get("rope_scaling") orelse return; + if (v == .null) return; + const r = try object(v); + const kind = r.get("rope_type") orelse r.get("type") orelse return error.UnsupportedLlamaRotary; + if (!isString(kind, "default")) return error.UnsupportedLlamaRotary; +} + +fn eosIds(v: std.json.Value, c: *Config) Error!void { + switch (v) { + .integer => |i| { + if (i < 0 or i >= c.vocab or c.eos_count == c.eos.len) return error.BadLlamaEos; + c.eos[c.eos_count] = @intCast(i); + c.eos_count += 1; + }, + .array => |a| for (a.items) |x| try eosIds(x, c), + else => return error.BadLlamaEos, + } +} + +pub fn parse(gpa: std.mem.Allocator, bytes: []const u8) Error!Config { + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); + defer parsed.deinit(); + const root = try object(parsed.value); + if (!isString(root.get("model_type"), "llama")) return error.BadLlamaConfig; + if (isTrue(root, "attention_bias") or isTrue(root, "mlp_bias") or isTrue(root, "tie_word_embeddings")) return error.UnsupportedLlamaGeometry; + if (root.get("hidden_act")) |act| if (!isString(act, "silu")) return error.UnsupportedLlamaGeometry; + for ([_][]const u8{ "torch_dtype", "dtype" }) |key| if (root.get(key)) |d| if (!isString(d, "bfloat16")) return error.UnsupportedLlamaPrecision; + if (root.get("quantization") != null or root.get("quantization_config") != null) return error.UnsupportedLlamaPrecision; + var c: Config = .{ + .hidden = try size(root, "hidden_size"), + .intermediate = try size(root, "intermediate_size"), + .layers = try size(root, "num_hidden_layers"), + .heads = try size(root, "num_attention_heads"), + .kv_heads = 0, + .vocab = try size(root, "vocab_size"), + .eps = @floatCast(try number(root.get("rms_norm_eps") orelse return error.BadLlamaConfig)), + .theta = try number(root.get("rope_theta") orelse .{ .float = 10000.0 }), + .context = 0, + }; + const kv_heads = try integerOr(root, "num_key_value_heads", @intCast(c.heads)); + if (kv_heads < 1 or kv_heads > c.heads) return error.UnsupportedLlamaGeometry; + c.kv_heads = @intCast(kv_heads); + if (try integerOr(root, "head_dim", @intCast(c.hidden / c.heads)) != head_dim) return error.UnsupportedLlamaGeometry; + if (c.heads % c.kv_heads != 0 or c.group() > max_group) return error.UnsupportedLlamaGeometry; + // the tensor-unit GEMM steps k by 256s; the row kernels read eight values a lane over 256 + if (c.hidden % 256 != 0 or c.intermediate % 256 != 0 or c.layers > max_layers) return error.UnsupportedLlamaGeometry; + if (!(c.eps > 0) or !(c.theta > 0)) return error.BadLlamaConfig; + try rotary(root); + const context = try integer(root, "max_position_embeddings"); + if (context < 1 or context > 1 << 20) return error.BadLlamaContext; + c.context = @intCast(context); + try eosIds(root.get("eos_token_id") orelse return error.BadLlamaEos, &c); + if (c.eos_count == 0) return error.BadLlamaEos; + return c; +} + +const fixture = @embedFile("fixtures/config.json"); + +test "Nanbeige4.3-33B's config is admitted with its geometry, context and end token" { + const c = try parse(std.testing.allocator, fixture); + try std.testing.expectEqual(@as(usize, 7168), c.hidden); + try std.testing.expectEqual(@as(usize, 22016), c.intermediate); + try std.testing.expectEqual(@as(usize, 52), c.layers); + try std.testing.expectEqual(@as(usize, 56), c.heads); + try std.testing.expectEqual(@as(usize, 8), c.kv_heads); + try std.testing.expectEqual(@as(usize, 166912), c.vocab); + try std.testing.expectEqual(@as(usize, 9216), c.qkvWidth()); + try std.testing.expectEqual(@as(usize, 7), c.group()); + try std.testing.expectEqual(@as(f32, 1e-5), c.eps); + try std.testing.expectEqual(@as(f64, 120000000), c.theta); + try std.testing.expectEqual(@as(usize, 524288), c.context); + try std.testing.expectEqualSlices(u32, &.{166101}, c.eos[0..c.eos_count]); + try std.testing.expect(c.isEos(166101) and !c.isEos(166100)); +} + +test "other layouts are refused by name" { + const gpa = std.testing.allocator; + const cases = .{ + .{ "\"hidden_act\": \"silu\"", "\"hidden_act\": \"gelu\"", error.UnsupportedLlamaGeometry }, + .{ "\"head_dim\": 128", "\"head_dim\": 64", error.UnsupportedLlamaGeometry }, + .{ "\"tie_word_embeddings\": false", "\"tie_word_embeddings\": true", error.UnsupportedLlamaGeometry }, + .{ "\"attention_bias\": false", "\"attention_bias\": true", error.UnsupportedLlamaGeometry }, + .{ "\"num_key_value_heads\": 8", "\"num_key_value_heads\": 3", error.UnsupportedLlamaGeometry }, + .{ "\"num_key_value_heads\": 8", "\"num_key_value_heads\": -1", error.UnsupportedLlamaGeometry }, + .{ "\"num_key_value_heads\": 8", "\"num_key_value_heads\": 112", error.UnsupportedLlamaGeometry }, + .{ "\"eos_token_id\": 166101", "\"eos_token_id\": [166101, 166102, 166103, 166104, 166105]", error.BadLlamaEos }, + .{ "\"rope_scaling\": null", "\"rope_scaling\": {\"rope_type\": \"llama3\", \"factor\": 8.0}", error.UnsupportedLlamaRotary }, + .{ "\"torch_dtype\": \"bfloat16\"", "\"torch_dtype\": \"float16\"", error.UnsupportedLlamaPrecision }, + .{ "\"torch_dtype\": \"bfloat16\"", "\"torch_dtype\": \"bfloat16\", \"quantization\": {\"bits\": 4, \"group_size\": 64}", error.UnsupportedLlamaPrecision }, + .{ "\"model_type\": \"llama\"", "\"model_type\": \"mistral\"", error.BadLlamaConfig }, + .{ "\"max_position_embeddings\": 524288", "\"max_position_embeddings\": 0", error.BadLlamaContext }, + }; + inline for (cases) |c| { + const bytes = try std.mem.replaceOwned(u8, gpa, fixture, c[0], c[1]); + defer gpa.free(bytes); + try std.testing.expect(!std.mem.eql(u8, bytes, fixture)); + try std.testing.expectError(c[2], parse(gpa, bytes)); + } +} + +test "absent fields take HF's defaults: head_dim from the hidden size, every head its own key head, default rope" { + const gpa = std.testing.allocator; + const bare = try std.mem.replaceOwned(u8, gpa, fixture, " \"head_dim\": 128,\n", ""); + defer gpa.free(bare); + const c = try parse(gpa, bare); + try std.testing.expectEqual(@as(usize, 56), c.heads); + const every = try std.mem.replaceOwned(u8, gpa, bare, " \"num_key_value_heads\": 8,\n", ""); + defer gpa.free(every); + const d = try parse(gpa, every); + try std.testing.expectEqual(@as(usize, 56), d.kv_heads); + const typed = try std.mem.replaceOwned(u8, gpa, bare, "\"rope_scaling\": null", "\"rope_scaling\": {\"type\": \"default\"}"); + defer gpa.free(typed); + _ = try parse(gpa, typed); +} + +test "RoPE's inverse frequencies: 1 at the first pair, falling to theta^-(1 - 1/64)" { + const c = try parse(std.testing.allocator, fixture); + var f: [head_dim / 2]f32 = undefined; + c.inverseFrequencies(&f); + try std.testing.expectEqual(@as(f32, 1), f[0]); + try std.testing.expect(f[1] < 1 and f[63] > 0 and f[63] < f[62]); + try std.testing.expectApproxEqRel(@as(f32, @floatCast(std.math.pow(f64, 120000000.0, -126.0 / 128.0))), f[63], 1e-6); +} diff --git a/zig/src/families/llama/fixtures/config.json b/zig/src/families/llama/fixtures/config.json new file mode 100644 index 0000000000..adee1d1b76 --- /dev/null +++ b/zig/src/families/llama/fixtures/config.json @@ -0,0 +1,37 @@ +{ + "architectures": [ + "LlamaForCausalLM" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "bos_token_id": 166100, + "embd_pdrop": 0.0, + "eos_token_id": 166101, + "ffn_hidden_size": 22016, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 7168, + "initializer_range": 0.02, + "intermediate_size": 22016, + "kv_channels": 128, + "max_position_embeddings": 524288, + "mlp_bias": false, + "model_type": "llama", + "num_attention_heads": 56, + "num_hidden_layers": 52, + "num_key_value_heads": 8, + "num_query_groups": 8, + "pad_token_id": 0, + "pretraining_tp": 1, + "resid_pdrop": 0.0, + "rms_norm_eps": 1e-05, + "rope_scaling": null, + "rope_theta": 120000000, + "rotary_base": 120000000, + "seq_length": 524288, + "tie_word_embeddings": false, + "torch_dtype": "bfloat16", + "transformers_version": "4.51.1", + "use_cache": false, + "vocab_size": 166912 +} \ No newline at end of file diff --git a/zig/src/families/llama/forward.zig b/zig/src/families/llama/forward.zig new file mode 100644 index 0000000000..353ba6e126 --- /dev/null +++ b/zig/src/families/llama/forward.zig @@ -0,0 +1,253 @@ +//! Llama's rows through every layer in one command buffer: row kernels to 32 rows, tensor-unit matmuls past them. +const std = @import("std"); +const mtl = @import("metal"); +const lanes = @import("lanes"); +const cfg = @import("config.zig"); +const st = @import("state.zig"); +const wts = @import("weights.zig"); +const Model = @import("model.zig").Model; +const Glue = @import("model.zig").Glue; + +/// The glue kernels' constants (glue.metal's ll_args). +const Args = extern struct { rows: u32 = 0, width: u32 = 0, heads: u32 = 0, eps: f32 = 0 }; +/// The q, k and v preparation's constants (glue.metal's ll_prep_args). +const PArgs = extern struct { rows: u32, heads: u32, kv_heads: u32, pos: u32, stride: u32 }; +/// The bf16 row projection's constants (glue.metal's ll_mv_args). +const MArgs = extern struct { rows: u32, n: u32, k: u32, ldy: u32 }; +/// The decode attention's constants (glue.metal's ll_vec_args). +const VArgs = extern struct { rows: u32, pos: u32, heads: u32, groups: u32, blocks: u32, chunk: u32, first: u32, scale: f32 }; + +const Enc = struct { + e: mtl.ComputeEncoder, + m: *const Model, + s: *st.Scratch, + + fn glue(x: Enc, g: Glue, args: anytype, bufs: []const B, grid: [3]usize, group: [3]usize, groups: bool) void { + x.e.setPipeline(x.m.pipe(g)); + for (bufs, 0..) |r, slot| { + if (r.args) x.e.setValue(args, slot) else x.e.setBuffer(r.buf, r.off, slot); + } + const g3 = mtl.Size.of(grid[0], grid[1], grid[2]); + const t3 = mtl.Size.of(group[0], group[1], group[2]); + if (groups) x.e.dispatchGroups(g3, t3) else x.e.dispatchThreads(g3, t3); + } + + /// y [rows, d.n] (rows `ldy` apart, from byte `out_off`) = x W^T: gemm_nax a prompt chunk, ll_matvec a window. + fn dense(x: Enc, d: wts.Dense, in: mtl.Buffer, out: mtl.Buffer, out_off: usize, n_rows: usize, ldy: usize) void { + const e = x.e; + if (n_rows <= st.batch_rows) { + // the variant holding the rows in registers: 1, 2, 4, 8, 16 or 32 + const variant: Glue = if (n_rows <= 1) .matvec_1 else if (n_rows <= 2) .matvec_2 else if (n_rows <= 4) .matvec_4 else if (n_rows <= 8) .matvec_8 else if (n_rows <= 16) .matvec_16 else .matvec_32; + e.setPipeline(x.m.pipe(variant)); + e.setBuffer(in, 0, 0); + e.setBuffer(d.buf, d.off, 1); + e.setValue(MArgs{ .rows = @intCast(n_rows), .n = @intCast(d.n), .k = @intCast(d.k), .ldy = @intCast(ldy) }, 2); + e.setBuffer(out, out_off, 3); + e.dispatchGroups(mtl.Size.of((d.n + 7) / 8, 1, 1), mtl.Size.of(256, 1, 1)); + return; + } + const M: usize = n_rows; + const tn = (d.n + 127) / 128; + const tm = (M + 63) / 64; + e.setPipeline(x.m.gemm[@intFromBool(M % 64 == 0)][@intFromBool(d.n % 128 == 0)]); + e.setBuffer(in, 0, 0); + e.setBuffer(d.buf, d.off, 1); + var p: [16]i32 = @splat(0); + const vals = [_]usize{ M, d.n, d.k, d.k, d.k, ldy, tn, tm, 2, d.k / 256, 0, 0, 0, 0, 0 }; + for (vals, 0..) |v, i| p[i] = @intCast(v); + e.setBytes(std.mem.asBytes(&p), 2); + e.setBuffer(out, out_off, 3); + e.dispatchThreads(mtl.Size.of(@as(usize, tn << 2) * 32, (tm + 3) / 4 * 4, 2), mtl.Size.of(32, 4, 2)); + } +}; + +pub const B = struct { buf: mtl.Buffer = undefined, off: usize = 0, args: bool = false }; +fn b(buf: mtl.Buffer) B { + return .{ .buf = buf }; +} +fn bo(buf: mtl.Buffer, off: usize) B { + return .{ .buf = buf, .off = off }; +} +fn t(x: wts.Tensor) B { + return .{ .buf = x.buffer, .off = x.offset }; +} +const A = B{ .args = true }; + +/// All rows' logits or the last's; ids on the GPU; the last row drawn on the GPU; committed unwaited into `flight`. +pub const Out = struct { all: bool = false, timing: ?*Timing = null, ids: ?B = null, draw: ?Draw = null, flight: ?*mtl.CommandBuffer = null }; + +/// A GPU draw of the last row's logits: argmax, or the keyed sampler at `position` (lanes' gpu_rule and gpu_full). +pub const Draw = struct { + how: union(enum) { greedy, sampled: struct { sampling: lanes.Sampling, position: u64 } }, + out: B, +}; + +/// A forward's host encoding and GPU seconds (tf-llama-run's bench). +pub const Timing = struct { encode: f64 = 0, gpu: f64 = 0 }; + +/// A stream's rows in a forward: the next `rows` of ids at positions cache.len.., its keys and values into its cache. +pub const Segment = struct { cache: *st.Cache, rows: usize }; + +/// One stream's `ids` at positions cache.len.. (len unmoved). +pub fn run(m: *const Model, s: *st.Scratch, cache: *st.Cache, ids: []const u32, out: Out) !void { + return runShared(m, s, &.{.{ .cache = cache, .rows = ids.len }}, ids, out); +} + +/// Several streams' rows in one forward, each segment's attention over its own cache: every row keeps its solo bits. +pub fn runShared(m: *const Model, s: *st.Scratch, segments: []const Segment, ids: []const u32, out: Out) !void { + const n = ids.len; + const c = &m.config; + if (n == 0 or n > s.rows or (out.all and n > st.batch_rows) or (!out.all and segments.len != 1)) return error.LlamaRowsExceeded; + var total: usize = 0; + for (segments) |seg| { + if (seg.rows == 0 or seg.cache.len + seg.rows > seg.cache.capacity) return error.LlamaContextFull; + if (segments.len > 1 and seg.rows > st.window_rows) return error.LlamaRowsExceeded; + total += seg.rows; + } + if (total != n) return error.InvalidLlamaSegments; + if (out.ids == null) { + for (ids) |id| if (id >= c.vocab) return error.InvalidLlamaToken; + @memcpy(s.ids.slice(u32, n), ids); + } + if (out.draw != null and (out.all or segments.len != 1)) return error.InvalidLlamaSegments; + const h = c.hidden; + const nr: u32 = @intCast(n); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const t0 = if (out.timing != null) seconds() else 0; + const cb = m.queue.commandBuffer(); + const x = Enc{ .e = cb.compute(.serial), .m = m, .s = s }; + const W = m.weights; + x.glue(.embed, Args{ .rows = nr, .width = @intCast(h) }, &.{ out.ids orelse b(s.ids), bo(W.embed.buf, W.embed.off), A, b(s.resid) }, .{ h, n, 1 }, .{ 256, 1, 1 }, false); + const normed = Args{ .rows = nr, .width = @intCast(h), .eps = c.eps }; + x.glue(.norm, normed, &.{ b(s.resid), t(W.layers[0].input_norm), A, b(s.x) }, .{ n, 1, 1 }, .{ 256, 1, 1 }, true); + for (W.layers, 0..) |l, i| try layer(x, segments, l, i, n, if (i + 1 < W.layers.len) W.layers[i + 1].input_norm else null); + if (out.all) { + x.glue(.norm, normed, &.{ b(s.resid), t(W.norm), A, b(s.x) }, .{ n, 1, 1 }, .{ 256, 1, 1 }, true); + x.dense(W.head, s.x, s.logits_all, 0, n, c.vocab); + } else { + x.glue(.norm, normed, &.{ bo(s.resid, (n - 1) * h * 2), t(W.norm), A, b(s.x) }, .{ 1, 1, 1 }, .{ 256, 1, 1 }, true); + x.dense(W.head, s.x, segments[0].cache.logits, 0, 1, c.vocab); + if (out.draw) |d| pick(x, segments[0].cache.logits, d); + } + x.e.end(); + if (out.timing) |tm| tm.encode += seconds() - t0; + cb.commit(); + if (out.flight) |f| { + f.* = cb.retain(); + return; + } + cb.wait(); + if (cb.failure()) |msg| { + std.log.err("llama forward: {s}", .{msg}); + return error.GpuFailed; + } + if (out.timing) |tm| tm.gpu += cb.gpuSeconds(); +} + +fn seconds() f64 { + var ts: std.c.timespec = undefined; + _ = std.c.clock_gettime(.MONOTONIC, &ts); + return @as(f64, @floatFromInt(ts.sec)) + @as(f64, @floatFromInt(ts.nsec)) * 1e-9; +} + +/// The draw of one logits row on the GPU into `d.out`: ll_argmax, or tf_sample_full one 1024-thread threadgroup. +fn pick(x: Enc, logits: mtl.Buffer, d: Draw) void { + const e = x.e; + const vocab: u32 = @intCast(x.m.config.vocab); + switch (d.how) { + .greedy => { + e.setPipeline(x.m.pipe(.argmax)); + e.setBuffer(logits, 0, 0); + e.setBuffer(d.out.buf, d.out.off, 1); + e.setValue(vocab, 2); + }, + .sampled => |sd| { + // the kernel indexes these by row: one row, padded as Nemotron's launches pad them + const r = lanes.gpu_rule.rule(sd.sampling); + var seeds: [16]u32 = @splat(0); + var positions: [8]u32 = @splat(0); + var cfgs: [32]f32 = @splat(0); + var caps: [8]u32 = @splat(0); + seeds[0..2].* = r.seedHalves(); + positions[0] = @intCast(sd.position); + cfgs[0..4].* = r.cfg(); + caps[0] = r.top_k; + e.setPipeline(x.m.sample_full); + e.setBuffer(logits, 0, 0); + e.setBytes(std.mem.sliceAsBytes(&seeds), 1); + e.setBytes(std.mem.sliceAsBytes(&positions), 2); + e.setBytes(std.mem.sliceAsBytes(&cfgs), 3); + e.setBytes(std.mem.sliceAsBytes(&caps), 4); + e.setBuffer(d.out.buf, d.out.off, 5); + e.setBuffer(d.out.buf, d.out.off, 6); + e.setValue(vocab, 7); + }, + } + e.dispatchGroups(mtl.Size.of(1, 1, 1), mtl.Size.of(1024, 1, 1)); +} + +/// Layer `i` on every segment's rows (s.x their input norm); its last add norms the residual by `next`. +fn layer(x: Enc, segments: []const Segment, l: wts.Layer, i: usize, n: usize, next: ?wts.Tensor) !void { + const s = x.s; + const c = &x.m.config; + const h = c.hidden; + const nr: u32 = @intCast(n); + const normed = Args{ .rows = nr, .width = @intCast(h), .eps = c.eps }; + const qkvw = c.qkvWidth(); + x.dense(l.attn.q, s.x, s.q, 0, n, qkvw); + x.dense(l.attn.k, s.x, s.q, c.queryWidth() * 2, n, qkvw); + x.dense(l.attn.v, s.x, s.q, (c.queryWidth() + c.kvWidth()) * 2, n, qkvw); + var base: usize = 0; + for (segments) |seg| { + try attention(x, seg.cache, i, base, seg.rows); + base += seg.rows; + } + x.dense(l.attn.o, s.att, s.y, 0, n, h); + x.glue(.add_norm, normed, &.{ b(s.resid), b(s.y), t(l.post_norm), A, b(s.x) }, .{ n, 1, 1 }, .{ 256, 1, 1 }, true); + x.dense(l.mlp.gate, s.x, s.gu, 0, n, 2 * c.intermediate); + x.dense(l.mlp.up, s.x, s.gu, c.intermediate * 2, n, 2 * c.intermediate); + x.glue(.silu_gu, Args{ .rows = nr, .width = @intCast(c.intermediate) }, &.{ b(s.gu), A, b(s.act) }, .{ c.intermediate, n, 1 }, .{ 256, 1, 1 }, false); + x.dense(l.mlp.down, s.act, s.y, 0, n, h); + if (next) |w| { + x.glue(.add_norm, normed, &.{ b(s.resid), b(s.y), t(w), A, b(s.x) }, .{ n, 1, 1 }, .{ 256, 1, 1 }, true); + } else x.glue(.add, Args{ .rows = nr, .width = @intCast(h) }, &.{ b(s.resid), b(s.y), A }, .{ h, n, 1 }, .{ 256, 1, 1 }, false); +} + +/// One stream's scratch rows base..: q and k rotated, k and v into its cache at their positions, attention to s.att. +fn attention(x: Enc, cache: *st.Cache, i: usize, base: usize, n: usize) !void { + const s = x.s; + const c = &x.m.config; + const start = cache.len; + const qrow: usize = c.queryWidth() * 2; // a row's qn and att bytes + const qn = bo(s.qn, base * qrow); + const o = bo(s.att, base * qrow); + const nr: u32 = @intCast(n); + const pa = PArgs{ .rows = nr, .heads = @intCast(c.heads), .kv_heads = @intCast(c.kv_heads), .pos = @intCast(start), .stride = @intCast(c.qkvWidth()) }; + const keys = bo(cache.kv, cache.keys(i)); + const values = bo(cache.kv, cache.values(i)); + x.glue(.prep, pa, &.{ bo(s.q, base * c.qkvWidth() * 2), b(x.m.inv), A, qn, keys, values }, .{ c.heads + c.kv_heads, n, 1 }, .{ 32, 1, 1 }, true); + const e = x.e; + if (n > st.window_rows) { // a prompt chunk: 32 rows of a head a tile, over every key the rows see + e.setPipeline(x.m.prompt_attn); + for ([_]B{ qn, keys, values, o }, 0..) |r, slot| e.setBuffer(r.buf, r.off, slot); + e.setValue([2]i32{ @intCast(n), @intCast(start) }, 4); + e.dispatchGroups(mtl.Size.of((n + 31) / 32, c.heads, 1), mtl.Size.of(cfg.head_dim / 64 * 32, 1, 1)); + return; + } + // decode rows: a key head's query heads a threadgroup over fixed key blocks at absolute places, then the merge + const chunk = st.chunk(cache.capacity); + const blocks = (start + n - 1) / chunk + 1; + if (blocks > st.max_blocks) return error.LlamaTooManyBlocks; + const va = VArgs{ .rows = nr, .pos = @intCast(start), .heads = @intCast(c.heads), .groups = @intCast(c.kv_heads), .blocks = @intCast(blocks), .chunk = @intCast(chunk), .first = 0, .scale = c.scale() }; + e.setPipeline(x.m.pipe(.decode_vec)); + for ([_]B{ qn, keys, values }, 0..) |r, slot| e.setBuffer(r.buf, r.off, slot); + e.setValue(va, 3); + for ([_]mtl.Buffer{ s.pm, s.pl, s.po }, 4..) |r, slot| e.setBuffer(r, 0, slot); + e.dispatchGroups(mtl.Size.of(c.kv_heads, n, blocks), mtl.Size.of(32, c.group(), 1)); + e.setPipeline(x.m.pipe(.merge_vec)); + for ([_]mtl.Buffer{ s.pm, s.pl, s.po }, 0..) |r, slot| e.setBuffer(r, 0, slot); + e.setValue(va, 3); + e.setBuffer(o.buf, o.off, 4); + e.dispatchGroups(mtl.Size.of(c.heads, n, 1), mtl.Size.of(1024, 1, 1)); +} diff --git a/zig/src/families/llama/llama.zig b/zig/src/families/llama/llama.zig new file mode 100644 index 0000000000..4eca3dc970 --- /dev/null +++ b/zig/src/families/llama/llama.zig @@ -0,0 +1,17 @@ +//! Native Metal Llama (model_type llama): GQA attention and a SwiGLU MLP in bf16 at HF's rounding points. +pub const config = @import("config.zig"); +pub const weights = @import("weights.zig"); +pub const Model = @import("model.zig").Model; +pub const state = @import("state.zig"); +pub const forward = @import("forward.zig"); +pub const backend = @import("backend.zig"); +pub const budget = @import("budget.zig"); +pub const residency = @import("residency.zig"); +pub const snapshot = @import("snapshot.zig"); + +test { + _ = config; + _ = state; + _ = budget; + _ = snapshot; +} diff --git a/zig/src/families/llama/model.zig b/zig/src/families/llama/model.zig new file mode 100644 index 0000000000..0d39ea15dd --- /dev/null +++ b/zig/src/families/llama/model.zig @@ -0,0 +1,115 @@ +//! A native Llama model on Metal: config, the checkpoint in Metal buffers, its matrices and every pipeline it runs. +const std = @import("std"); +const mtl = @import("metal"); +const ks = @import("kernel_sources"); +const cfg = @import("config.zig"); +const ck = @import("../../core/checkpoint_metal.zig"); +const shards = @import("../../core/checkpoint.zig"); +const frags = @import("../../core/frags.zig"); +const weights = @import("weights.zig"); + +pub const glue_names = [_][:0]const u8{ "ll_embed", "ll_norm", "ll_add", "ll_add_norm", "ll_prep", "ll_silu_gu", "ll_decode_vec", "ll_merge_vec", "ll_argmax", "ll_matvec_1", "ll_matvec_2", "ll_matvec_4", "ll_matvec_8", "ll_matvec_16", "ll_matvec_32" }; +pub const Glue = enum { embed, norm, add, add_norm, prep, silu_gu, decode_vec, merge_vec, argmax, matvec_1, matvec_2, matvec_4, matvec_8, matvec_16, matvec_32 }; + +pub const Model = struct { + gpa: std.mem.Allocator, + device: mtl.Device, + queue: mtl.Queue, + config: cfg.Config, + checkpoint: ck.Checkpoint, + weights: weights.Weights, + glue: [glue_names.len]mtl.Pipeline, + /// Prompt attention on the tensor units; decode attention is glue's ll_decode_vec. + prompt_attn: mtl.Pipeline, + /// A sampled step's keyed draw over the whole vocabulary (host twin: lanes' gpu_full). + sample_full: mtl.Pipeline, + /// Prompt chunks' bf16 GEMM on the tensor units, by whole M and N tiles. + gemm: [2][2]mtl.Pipeline, + /// RoPE's inverse frequencies: fp32 [head_dim / 2]. + inv: mtl.Buffer, + + pub fn load(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !*Model { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const m = try gpa.create(Model); + errdefer gpa.destroy(m); + m.gpa = gpa; + m.config = try cfg.Config.read(gpa, io, dir); + m.device = try mtl.Device.init(); + errdefer m.device.deinit(); + m.queue = try m.device.queue(); + errdefer m.queue.deinit(); + m.checkpoint = ck.Checkpoint.init(gpa); + errdefer m.checkpoint.deinit(); + const files = try shards.shardFiles(gpa, io, dir); + defer shards.freeShardFiles(gpa, files); + for (files) |path| try m.checkpoint.addFile(m.device, path, ""); + m.weights = try weights.load(gpa, &m.checkpoint, &m.config); + errdefer m.weights.deinit(gpa); + { + const src = try std.fmt.allocPrint(gpa, "#define LL_DH {d}\n{s}", .{ cfg.head_dim, ks.llama_glue }); + defer gpa.free(src); + const lib = try mtl.Library.fromSource(m.device, src, .{ .language = mtl.CompileOptions.mlx().language, .math_functions = .precise }); + defer lib.deinit(); + var made: usize = 0; + errdefer for (m.glue[0..made]) |p| p.deinit(); + for (glue_names, 0..) |n, i| { + m.glue[i] = try mtl.Pipeline.init(m.device, lib, n, false); + made += 1; + } + } + errdefer for (m.glue) |p| p.deinit(); + { + const body = try frags.source(m.device, gpa, ks.llama_attention_nax); + defer gpa.free(body); + const c = m.config; + const src = try std.fmt.allocPrint(gpa, "#define DH {d}\n#define NQ {d}\n#define G {d}\n#define KVH {d}\n#define SCALE {e}f\n{s}", .{ cfg.head_dim, c.heads, c.group(), c.kv_heads, c.scale(), body }); + defer gpa.free(src); + const lib = try mtl.Library.fromSource(m.device, src, mtl.CompileOptions.mlx()); + defer lib.deinit(); + m.prompt_attn = try mtl.Pipeline.init(m.device, lib, "ll_prompt_attention", false); + } + errdefer m.prompt_attn.deinit(); + { + const lib = try mtl.Library.fromSource(m.device, ks.nemotron_sample, mtl.CompileOptions.mlx()); + defer lib.deinit(); + m.sample_full = try mtl.Pipeline.init(m.device, lib, "tf_sample_full", false); + } + errdefer m.sample_full.deinit(); + { + const body = try frags.source(m.device, gpa, ks.llama_gemm); + defer gpa.free(body); + const lib = try mtl.Library.fromSource(m.device, body, mtl.CompileOptions.mlx()); + defer lib.deinit(); + var made: usize = 0; + errdefer for (0..made) |i| m.gemm[i / 2][i % 2].deinit(); + for (0..4) |i| { + var nm: [16]u8 = undefined; + const full = try std.fmt.bufPrintSentinel(&nm, "ll_gemm_{d}{d}", .{ i / 2, i % 2 }, 0); + m.gemm[i / 2][i % 2] = try mtl.Pipeline.init(m.device, lib, full, false); + made += 1; + } + } + errdefer for (m.gemm) |row| for (row) |p| p.deinit(); + m.inv = try m.device.buffer(cfg.head_dim / 2 * 4, mtl.ResourceOptions.shared); + m.config.inverseFrequencies(m.inv.slice(f32, cfg.head_dim / 2)); + return m; + } + + pub fn pipe(m: *const Model, g: Glue) mtl.Pipeline { + return m.glue[@backingInt(g)]; + } + + pub fn deinit(m: *Model) void { + m.inv.deinit(); + for (m.gemm) |row| for (row) |p| p.deinit(); + m.sample_full.deinit(); + m.prompt_attn.deinit(); + for (m.glue) |p| p.deinit(); + m.weights.deinit(m.gpa); + m.checkpoint.deinit(); + m.queue.deinit(); + m.device.deinit(); + m.gpa.destroy(m); + } +}; diff --git a/zig/src/families/llama/residency.zig b/zig/src/families/llama/residency.zig new file mode 100644 index 0000000000..4b37f78969 --- /dev/null +++ b/zig/src/families/llama/residency.zig @@ -0,0 +1,25 @@ +//! One initial inventory and a stable follow hook for every served Llama cache and saved state on Metal. +const std = @import("std"); +const mtl = @import("metal"); +const Resident = mtl.resident.Resident; + +/// Weights, tables and scratches in one residency set for the keepalive, if the 70% limit holds them. +pub fn hold(h: anytype) !void { + var list: std.ArrayList(mtl.Buffer) = .empty; + defer list.deinit(h.gpa); + const m = h.model; + for (m.checkpoint.shards.items) |sh| try list.append(h.gpa, sh.buffer); + try list.append(h.gpa, m.inv); + try mtl.resident.fields.append(h.gpa, &list, &h.metal.scratch); + try mtl.resident.fields.append(h.gpa, &list, &h.metal.prompt); + h.resident = try Resident.hold(h.gpa, h.io, m.device, m.queue, list.items, "llama"); + if (h.resident) |*r| h.metal.setResident(r.follow()); + h.warm = if (h.resident) |*r| r.target() else .{ .queue = m.queue }; +} + +pub fn drop(h: anytype) void { + h.metal.setResident(null); + if (h.resident) |*r| r.deinit(); + h.resident = null; + h.warm = .{ .queue = h.model.queue }; +} diff --git a/zig/src/families/llama/snapshot.zig b/zig/src/families/llama/snapshot.zig new file mode 100644 index 0000000000..6832eb9de6 --- /dev/null +++ b/zig/src/families/llama/snapshot.zig @@ -0,0 +1,77 @@ +//! A stream's prompt state kept for reuse on Metal: every layer's first `at` keys and values. +const std = @import("std"); +const mtl = @import("metal"); +const cfg = @import("config.zig"); +const st = @import("state.zig"); + +/// One kept state: its bytes on the GPU (private) and the prompt tokens it holds. +pub const Snap = struct { + buf: mtl.Buffer, + at: u32, + bytes: u64, + freed: bool = false, // its storage given back early (a stream needed the room); its entry goes at the pass's report + follow: ?mtl.ResidencySet.Follow = null, + + pub fn free(self: *Snap) void { + if (!self.freed) { + if (self.follow) |f| f.leave(f.ctx, self.buf); + self.buf.deinit(); + } + self.freed = true; + } +}; + +/// Bytes a state after `at` prompt tokens takes: the keys and values of `at` positions in every layer. +pub fn bytes(c: *const cfg.Config, at: usize) u64 { + return @as(u64, c.layers) * st.kvBytes(c) * at; +} + +/// Copy `cache`'s state at its length into a new snapshot. +pub fn save(device: mtl.Device, queue: mtl.Queue, c: *const cfg.Config, cache: *const st.Cache) !Snap { + const n = bytes(c, cache.len); + const buf = try device.buffer(@intCast(@max(n, 16)), mtl.ResourceOptions.private); + errdefer buf.deinit(); + try copy(queue, cache, buf, true); + if (cache.follow) |f| f.join(f.ctx, buf); + return .{ .buf = buf, .at = @intCast(cache.len), .bytes = n, .follow = cache.follow }; +} + +/// Make `cache` (empty, any capacity past the state) hold `snap`'s state. +pub fn restore(queue: mtl.Queue, snap: *const Snap, cache: *st.Cache) !void { + if (snap.freed or snap.at >= cache.capacity) return error.LlamaSnapshotGone; + cache.len = snap.at; + errdefer cache.len = 0; + try copy(queue, cache, snap.buf, false); +} + +/// The pieces between the cache and a snapshot buffer (a layer's keys, then its values), one blit pass, waited on. +fn copy(queue: mtl.Queue, cache: *const st.Cache, buf: mtl.Buffer, out: bool) !void { + const row = cache.kv_width * 2; + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const cb = queue.commandBuffer(); + const e = cb.blit(); + var at: usize = 0; + for (0..cache.layers) |l| { + for ([_]usize{ cache.keys(l), cache.values(l) }) |off| { + const len = cache.len * row; + if (len == 0) continue; + if (out) e.copy(cache.kv, off, buf, at, len) else e.copy(buf, at, cache.kv, off, len); + at += len; + } + } + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.log.err("llama prompt state copy: {s}", .{msg}); + return error.GpuFailed; + } +} + +test "llama snapshot: a state's bytes are every layer's keys and values to its length" { + const c = try cfg.parse(std.testing.allocator, @embedFile("fixtures/config.json")); + try std.testing.expectEqual(@as(u64, 0), bytes(&c, 0)); + try std.testing.expectEqual(@as(u64, 52 * 4096 * 1000), bytes(&c, 1000)); + try std.testing.expectEqual(st.cacheBytes(&c, 4096) - c.vocab * 2, bytes(&c, 4096)); +} diff --git a/zig/src/families/llama/state.zig b/zig/src/families/llama/state.zig new file mode 100644 index 0000000000..fd36ed9fe1 --- /dev/null +++ b/zig/src/families/llama/state.zig @@ -0,0 +1,149 @@ +//! Llama on Metal: a stream's caches (every layer holds every position), and a window's or a chunk's scratch. +const std = @import("std"); +const mtl = @import("metal"); +const cfg = @import("config.zig"); + +/// Rows one stream's decode window takes at most (a verify window: the pending token and its drafts). +pub const window_rows = 16; +/// Rows a shared decode forward takes at most (several streams' windows, row kernels); more take the prompt path. +pub const batch_rows = 32; +/// Rows a prompt chunk takes at most. +pub const prompt_rows = 2048; + +/// Slots of a layer for `capacity` positions: whole 64-key blocks, so a prompt tile's reads stay in the layer. +pub fn slots(capacity: usize) usize { + return std.mem.alignForward(usize, capacity, 64); +} + +/// Keys a decode attention block takes for a stream of `capacity`: at most 512 blocks, 32 to 2,048 keys each. +pub fn chunk(capacity: usize) usize { + return std.math.clamp(std.math.ceilPowerOfTwo(usize, std.math.divCeil(usize, capacity, 512) catch unreachable) catch 2048, 32, 2048); +} +/// Partials a decode row's head holds at most: a layer's blocks at capacity, and more. +pub const max_blocks = 544; +/// The longest stream the blocks reach (512 of 2,048 keys). +pub const max_capacity = 512 * 2048; + +/// Bytes one position takes in one layer: its keys and values. +pub fn kvBytes(c: *const cfg.Config) usize { + return c.kvWidth() * 2 * 2; +} + +/// Bytes a stream of `capacity` positions holds in caches, its logits row included. +pub fn cacheBytes(c: *const cfg.Config, capacity: usize) usize { + return c.vocab * 2 + c.layers * slots(capacity) * kvBytes(c); +} + +pub const Cache = struct { + kv: mtl.Buffer, // keys then values of each layer, at `layer_at` + logits: mtl.Buffer, + layer_at: [cfg.max_layers]usize = @splat(0), + layers: usize, + kv_width: usize, + capacity: usize, + len: usize = 0, + follow: ?mtl.ResidencySet.Follow = null, + + pub fn setResident(self: *Cache, hook: ?mtl.ResidencySet.Follow) void { + if (self.follow) |f| mtl.resident.fields.apply(self, f, false); + self.follow = hook; + if (hook) |f| mtl.resident.fields.apply(self, f, true); + } + + pub fn init(device: mtl.Device, c: *const cfg.Config, capacity: usize) !Cache { + var out = Cache{ .kv = undefined, .logits = undefined, .layers = c.layers, .kv_width = c.kvWidth(), .capacity = capacity }; + var kv: usize = 0; + for (0..c.layers) |i| { + out.layer_at[i] = kv; + kv += slots(capacity) * kvBytes(c); + } + out.kv = try device.buffer(@max(kv, 16), mtl.ResourceOptions.private); + errdefer out.kv.deinit(); + out.logits = try device.buffer(c.vocab * 2, mtl.ResourceOptions.shared); + return out; + } + + pub fn keys(self: *const Cache, l: usize) usize { + return self.layer_at[l]; + } + + pub fn values(self: *const Cache, l: usize) usize { + return self.layer_at[l] + slots(self.capacity) * self.kv_width * 2; + } + + pub fn deinit(self: *Cache) void { + self.setResident(null); + self.kv.deinit(); + self.logits.deinit(); + } +}; + +pub const fields = [_][]const u8{ "ids", "resid", "x", "q", "qn", "att", "y", "gu", "act", "logits_all", "pm", "pl", "po" }; + +/// Bytes of each Scratch buffer for `rows` rows at once, in `fields` order. +fn sizes(c: *const cfg.Config, rows: usize) [fields.len]usize { + const r = rows; + const prompt = rows > batch_rows; + const h = c.hidden; + const qw = r * c.queryWidth() * 2; + const partials = @as(usize, @min(r, window_rows)) * c.heads * max_blocks * 4; // decode attention's maxima, sums, outputs (fp32) + return .{ + r * 4, r * h * 2, r * h * 2, r * c.qkvWidth() * 2, + qw, qw, r * h * 2, r * 2 * c.intermediate * 2, + r * c.intermediate * 2, @as(usize, @min(r, batch_rows)) * c.vocab * 2, if (prompt) 16 else partials, if (prompt) 16 else partials, + if (prompt) 16 else partials * cfg.head_dim, + }; +} + +/// Bytes a Scratch for `rows` rows holds. +pub fn scratchBytes(c: *const cfg.Config, rows: usize) usize { + var total: usize = 0; + for (sizes(c, rows)) |n| total += n; + return total; +} + +pub const Scratch = struct { + ids: mtl.Buffer, + resid: mtl.Buffer, // bf16 [rows, hidden]: the reference's residual + x: mtl.Buffer, // bf16 normed rows + q: mtl.Buffer, // the projections' rows: q | k | v + qn: mtl.Buffer, // q rotated + att: mtl.Buffer, // attention's output + y: mtl.Buffer, // a sublayer's bf16 output [rows, hidden] + gu: mtl.Buffer, // gate | up rows + act: mtl.Buffer, + logits_all: mtl.Buffer, // every row's logits of a window + pm: mtl.Buffer, // decode attention's partials by block: maxima, sums, outputs + pl: mtl.Buffer, + po: mtl.Buffer, + rows: usize, + + /// Buffers for `rows` rows at once (batch_rows for decode, prompt_rows for prompt chunks). + pub fn init(device: mtl.Device, c: *const cfg.Config, rows: usize) !Scratch { + const sizes_ = sizes(c, rows); + var s: Scratch = undefined; + var made: usize = 0; + errdefer inline for (fields, 0..) |f, i| if (i < made) @field(s, f).deinit(); + inline for (fields, 0..) |f, i| { + @field(s, f) = try device.buffer(sizes_[i], mtl.ResourceOptions.shared); + made += 1; + } + s.rows = rows; + return s; + } + + pub fn deinit(s: *Scratch) void { + inline for (fields) |f| @field(s, f).deinit(); + } +}; + +test "llama state: a stream's caches on Nanbeige4.3-33B's geometry, and the decode blocks" { + const c = try cfg.parse(std.testing.allocator, @embedFile("fixtures/config.json")); + // 52 layers of 4,096 positions, 4 KiB a position a layer, and a logits row + try std.testing.expectEqual(@as(usize, 166912 * 2 + 52 * 4096 * 4096), cacheBytes(&c, 4096)); + try std.testing.expectEqual(@as(usize, 208 << 10), c.layers * kvBytes(&c)); + try std.testing.expect(std.math.divCeil(usize, max_capacity, chunk(max_capacity)) catch unreachable <= 512); + try std.testing.expectEqual(@as(usize, 32), chunk(4096)); + try std.testing.expectEqual(@as(usize, 1024), chunk(524288)); + try std.testing.expect(scratchBytes(&c, prompt_rows) > scratchBytes(&c, batch_rows)); +} diff --git a/zig/src/families/llama/weights.zig b/zig/src/families/llama/weights.zig new file mode 100644 index 0000000000..0b78e0d012 --- /dev/null +++ b/zig/src/families/llama/weights.zig @@ -0,0 +1,97 @@ +//! Llama's checkpoint on Metal: the release's bf16 matrices as stored (torch's Linear [n, k]), nothing copied. +const std = @import("std"); +const mtl = @import("metal"); +const ck = @import("../../core/checkpoint_metal.zig"); +const cfg = @import("config.zig"); + +pub const Tensor = ck.Tensor; + +/// A bf16 [n, k] matrix as the release stores it, at `off` in `buf`. +pub const Dense = struct { buf: mtl.Buffer, off: usize, n: usize, k: usize }; + +pub const Attention = struct { q: Dense, k: Dense, v: Dense, o: Dense }; +pub const Mlp = struct { gate: Dense, up: Dense, down: Dense }; +pub const Layer = struct { input_norm: Tensor, post_norm: Tensor, attn: Attention, mlp: Mlp }; + +pub const Weights = struct { + embed: Dense, + head: Dense, + norm: Tensor, + layers: []Layer, + + pub fn deinit(w: *Weights, gpa: std.mem.Allocator) void { + gpa.free(w.layers); + } +}; + +fn name(buf: []u8, a: []const u8, b: []const u8) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}{s}", .{ a, b }); +} + +/// `t` is bf16 with exactly `dims`, its bytes those dims' extent. +fn require(t: Tensor, dims: []const usize) error{UnexpectedLlamaTensor}!void { + if (t.dtype != .bf16 or t.rank != dims.len or t.rank > t.shape.len) return error.UnexpectedLlamaTensor; + if (!std.mem.eql(usize, t.shape[0..t.rank], dims)) return error.UnexpectedLlamaTensor; + var bytes: usize = 2; + for (dims) |n| bytes = std.math.mul(usize, bytes, n) catch return error.UnexpectedLlamaTensor; + if (t.bytes != bytes) return error.UnexpectedLlamaTensor; +} + +fn norm(c: *const ck.Checkpoint, hidden: usize, prefix: []const u8, suffix: []const u8) !Tensor { + var buf: [256]u8 = undefined; + const n = try name(&buf, prefix, suffix); + const t = try c.get(n); + require(t, &.{hidden}) catch |err| { + std.log.err("llama: {s} is not bf16[{d}]", .{ n, hidden }); + return err; + }; + return t; +} + +/// The release's bf16 [n, k] matrix in place: k whole 256s for the GEMM's steps, 8-byte aligned for the row kernel. +fn dense(c: *const ck.Checkpoint, prefix: []const u8, suffix: []const u8, n: usize, k: usize) !Dense { + var buf: [256]u8 = undefined; + const full = try name(&buf, prefix, suffix); + const w = try c.get(full); + require(w, &.{ n, k }) catch |err| { + std.log.err("llama: {s} is not bf16[{d}, {d}]", .{ full, n, k }); + return err; + }; + if (k % 256 != 0 or w.offset % 8 != 0) return error.UnexpectedLlamaTensor; + return .{ .buf = w.buffer, .off = w.offset, .n = n, .k = k }; +} + +pub fn load(gpa: std.mem.Allocator, c: *const ck.Checkpoint, f: *const cfg.Config) !Weights { + const h = f.hidden; + var w = Weights{ + .embed = try dense(c, "model.", "embed_tokens.weight", f.vocab, h), + .head = try dense(c, "", "lm_head.weight", f.vocab, h), + .norm = try norm(c, h, "model.", "norm.weight"), + .layers = try gpa.alloc(Layer, f.layers), + }; + errdefer w.deinit(gpa); + for (w.layers, 0..) |*l, i| { + var pb: [64]u8 = undefined; + const p = try std.fmt.bufPrint(&pb, "model.layers.{d}.", .{i}); + var ab: [96]u8 = undefined; + const a = try name(&ab, p, "self_attn."); + var mb: [96]u8 = undefined; + const m = try name(&mb, p, "mlp."); + l.* = .{ + .input_norm = try norm(c, h, p, "input_layernorm.weight"), + .post_norm = try norm(c, h, p, "post_attention_layernorm.weight"), + .attn = .{ + .q = try dense(c, a, "q_proj.weight", f.queryWidth(), h), + .k = try dense(c, a, "k_proj.weight", f.kvWidth(), h), + .v = try dense(c, a, "v_proj.weight", f.kvWidth(), h), + .o = try dense(c, a, "o_proj.weight", h, f.queryWidth()), + }, + .mlp = .{ + .gate = try dense(c, m, "gate_proj.weight", f.intermediate, h), + .up = try dense(c, m, "up_proj.weight", f.intermediate, h), + .down = try dense(c, m, "down_proj.weight", h, f.intermediate), + }, + }; + } + return w; +} diff --git a/zig/src/native/llama.zig b/zig/src/native/llama.zig new file mode 100644 index 0000000000..bbfc69c779 --- /dev/null +++ b/zig/src/native/llama.zig @@ -0,0 +1,180 @@ +//! Llama on Metal served by the lane host: one-row rounds and copy-drafted windows, contexts within the budget. +const std = @import("std"); +const mtl = @import("metal"); +const api = @import("engine_api"); +const tf = @import("tensorfold"); +const fl = tf.llama; +const lanes = tf.lanes; +const Resident = @import("resident.zig").Resident; + +const Host = struct { + gpa: std.mem.Allocator, + io: std.Io, + model: *fl.Model, + metal: *fl.backend.Metal, + warm: mtl.keepalive.Target = undefined, // the residency set's, or the bare queue's, for the lane host's idle ticker + resident: ?Resident = null, // weights and scratch kept wired between requests, when they fit the 70% limit + config: lanes.Config, + clock: lanes.backend.WallClock, + core: lanes.Engine, + host: api.LaneHost, + cache: ?api.prompt_cache.Store = null, // kept prompt states, in what the streams' caches leave of the budget + cache_cap: u64 = 0, // --prompt-cache-gib's bytes (the budget when unset) + + fn close(ctx: *anyopaque) void { + const h: *Host = @ptrCast(@alignCast(ctx)); + h.host.stop(); + if (h.cache) |*store| store.deinit(); + h.core.deinit(); + h.config.deinit(h.gpa); + fl.residency.drop(h); + h.metal.deinit(); + h.model.deinit(); + h.gpa.destroy(h); + } +}; + +pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !?api.Opened { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const model = fl.Model.load(gpa, io, o.dir) catch |err| { + problem.* = try std.fmt.allocPrint(a, "the native Llama engine cannot load this checkpoint ({s})", .{@errorName(err)}); + return null; + }; + errdefer model.deinit(); + // refusals that return null free what they hold themselves: errdefer runs only for errors + var ram: u64 = 0; + var len: usize = @sizeOf(u64); + if (std.c.sysctlbyname("hw.memsize", &ram, &len, null, 0) != 0) ram = 0; + const text: ?[]const u8 = if (std.c.getenv("LLAMA_LOAD_LIMIT_GB")) |v| std.mem.span(v) else null; + // the decode window's and the prompt chunk's buffers + const scratch: u64 = fl.state.scratchBytes(&model.config, fl.state.batch_rows) + fl.state.scratchBytes(&model.config, fl.state.prompt_rows); + const budget = fl.budget.caches(ram, model.device.maxWorkingSet(), text, model.checkpoint.residentBytes() + scratch) catch { + model.deinit(); + problem.* = "LLAMA_LOAD_LIMIT_GB is not a load limit: give a number of GB above 0, or unset it for 75% of RAM"; + return null; + }; + const fits = fl.budget.longest(&model.config, budget, @min(model.config.context, fl.state.max_capacity)); + const window: i64 = o.context orelse @intCast(fits); + if (window < 1 or window > fits) { + problem.* = try std.fmt.allocPrint(a, "--context {d}: this Mac's Llama caches fit {d} tokens under its load limit (75% of RAM, or LLAMA_LOAD_LIMIT_GB)", .{ window, fits }); + model.deinit(); // after the message, so a failed allocation leaves the model to the errdefer + return null; + } + std.log.info("llama: {d:.1} GiB for caches, context up to {d} tokens", .{ @as(f64, @floatFromInt(budget)) / (1 << 30), window }); + const h = try gpa.create(Host); + errdefer gpa.destroy(h); + h.gpa = gpa; + h.io = io; + h.model = model; + h.resident = null; + h.metal = try fl.backend.Metal.init(gpa, model, .{ .capacity = @intCast(window), .streams = @max(o.lanes, 1), .budget = budget }); + errdefer h.metal.deinit(); + try h.metal.warm(); + h.config = try lanes.Config.init(gpa, h.metal.facts(), fl.state.window_rows, fl.state.window_rows - 1); + errdefer h.config.deinit(gpa); + h.clock = .{ .io = io }; + h.core = lanes.Engine.init(gpa, &h.config, h.metal.backend(), h.clock.clock()); + errdefer h.core.deinit(); + h.host = api.LaneHost.init(gpa, io, &h.core, .{ .lanes = o.lanes, .context_window = @intCast(window), .prefill_step = fl.state.prompt_rows, .structures = true, .logprobs = true }); + h.cache = null; + h.cache_cap = if (o.prompt_cache_gib) |g| @intFromFloat(g * (1 << 30)) else budget; + if (h.cache_cap > 0) { + const M = fl.backend.Metal; + // prompt chunks change arithmetic with their rows + const rules: api.prompt_cache.Rules = .{ .lookahead = 0, .planned = true, .grid = fl.state.prompt_rows, .min_prompt = 0 }; + h.cache = api.prompt_cache.Store.init(gpa, .{ .ptr = h.metal, .vtable = &.{ .bytes = M.snapBytes, .save = M.snapSave, .restore = M.snapRestore, .drop = M.snapDrop } }, rules, @min(h.cache_cap, budget)); + h.host.cache = &h.cache.?; + h.host.info_.prompt_cache = true; + h.metal.room = .{ .ctx = h, .fit = fitCache }; + } + errdefer if (h.cache) |*store| store.deinit(); + try fl.residency.hold(h); + errdefer fl.residency.drop(h); + h.host.keepalive_target = .{ .ctx = &h.warm, .tick = mtl.keepalive.Target.tick }; + try h.host.start(); + return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h }; +} + +/// Metal.room's fit: the kept states' budget is what the streams leave, at most the cap; the least valuable go first. +fn fitCache(ctx: *anyopaque, bytes: u64, restoring: ?*anyopaque, prompt: []const u32) void { + const h: *Host = @ptrCast(@alignCast(ctx)); + const store = if (h.cache) |*owned| owned else return; + store.budget = @min(h.cache_cap, bytes); + const Left = struct { + s: *const api.prompt_cache.Store, + pub fn fits(l: @This()) bool { + return l.s.held <= l.s.budget; + } + }; + _ = store.giveBack(restoring, prompt, Left{ .s = store }); +} + +const CacheFixture = struct { + live: usize = 0, + fn bytes(_: *anyopaque, at: u32) u64 { + return 100 + at; + } + fn save(ptr: *anyopaque, _: ?*anyopaque, at: u32) !api.prompt_cache.Saved { + const self: *@This() = @ptrCast(@alignCast(ptr)); + const held = try std.testing.allocator.create(u32); + held.* = at; + self.live += 1; + return held; + } + fn restore(_: *anyopaque, _: ?*anyopaque, _: api.prompt_cache.Saved) !void {} + fn drop(ptr: *anyopaque, saved: api.prompt_cache.Saved) void { + const self: *@This() = @ptrCast(@alignCast(ptr)); + std.testing.allocator.destroy(@as(*u32, @ptrCast(@alignCast(saved)))); + self.live -= 1; + } + fn snapshots(self: *@This()) api.prompt_cache.Snapshots { + return .{ .ptr = self, .vtable = &.{ .bytes = bytes, .save = save, .restore = restore, .drop = drop } }; + } +}; + +test "llama host: cache callback updates the owned store rather than a temporary value" { + var f: CacheFixture = .{}; + var h: Host = undefined; + h.cache_cap = 500; + h.cache = api.prompt_cache.Store.init(std.testing.allocator, f.snapshots(), .{ .min_prompt = 0 }, 500); + defer h.cache.?.deinit(); + fitCache(&h, 200, null, &.{}); + try std.testing.expectEqual(@as(u64, 200), h.cache.?.budget); + fitCache(&h, 800, null, &.{}); + try std.testing.expectEqual(@as(u64, 500), h.cache.?.budget); +} + +test "llama host: cache callback evicts from the real store and preserves the restored entry" { + var f: CacheFixture = .{}; + var h: Host = undefined; + h.cache_cap = 1000; + h.cache = api.prompt_cache.Store.init(std.testing.allocator, f.snapshots(), .{ .min_prompt = 0 }, 1000); + defer h.cache.?.deinit(); + const store = &h.cache.?; + try std.testing.expect(store.keep(&.{ 1, 1, 1, 1, 9 }, 4, null, &.{}, &.{})); + try std.testing.expect(store.keep(&.{ 2, 2, 2, 2, 2, 2, 9 }, 6, null, &.{}, &.{})); + try std.testing.expectEqual(@as(u64, 210), store.held); + const kept = store.entries.items[0].saved; + fitCache(&h, 110, kept, &.{ 1, 1, 1, 1, 9 }); + try std.testing.expectEqual(@as(u64, 110), store.budget); + try std.testing.expectEqual(@as(u64, 104), store.held); + try std.testing.expectEqual(@as(usize, 1), store.entries.items.len); + try std.testing.expect(store.entries.items[0].saved == kept); + try std.testing.expectEqual(@as(usize, 1), f.live); + try std.testing.expectEqual(@as(u64, 1), store.counts.evicted); + fitCache(&h, 110, kept, &.{ 1, 1, 1, 1, 9 }); + try std.testing.expectEqual(@as(u64, 1), store.counts.evicted); + fitCache(&h, 0, null, &.{}); + try std.testing.expectEqual(@as(u64, 0), store.held); + try std.testing.expectEqual(@as(usize, 0), store.entries.items.len); + try std.testing.expectEqual(@as(usize, 0), f.live); + try std.testing.expectEqual(@as(u64, 2), store.counts.evicted); +} + +test "llama host: cache callback with no configured cache leaves no state to mutate" { + var h: Host = undefined; + h.cache = null; + fitCache(&h, 0, null, &.{}); + try std.testing.expect(h.cache == null); +} diff --git a/zig/src/native/metal.zig b/zig/src/native/metal.zig index a42272fe56..61e30041c9 100644 --- a/zig/src/native/metal.zig +++ b/zig/src/native/metal.zig @@ -51,6 +51,7 @@ pub const families: []const api.Family = &.{ .{ .model_type = "gemma4_unified", .formats = gemma4_widths }, .{ .model_type = "step5_vl", .formats = &.{"mlx-q2g128"} }, .{ .model_type = "spark2_5", .formats = &.{ "bf16", "mlx-q2g64", "mlx-q3g64", "mlx-q4g64", "mlx-q5g64", "mlx-q6g64", "mlx-q8g64" } }, + .{ .model_type = "llama", .formats = &.{"bf16"} }, }; /// Gemma 4 31B (gemma4) and 12B (gemma4_unified, its text model): one native family, gemma4.zig's host. @@ -176,10 +177,6 @@ pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]c problem.* = why; return null; } - if (o.drafter != null and !dense and !isGemma4(o.model_type)) { - problem.* = "--drafter is supported for Qwen3.8-27B-class checkpoints (qwen3_5 with its own output head) and Gemma 4 dense"; - return null; - } if (dense) return openQwen27(a, gpa, io, o, problem); if (std.mem.eql(u8, o.model_type, "qwen4_exp")) return openFlashNext(a, gpa, io, o, problem); if (std.mem.eql(u8, o.model_type, "glm5_next")) return openGlm(a, gpa, io, o, problem); @@ -192,6 +189,7 @@ pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]c if (std.mem.eql(u8, o.model_type, "qwen3_5")) return @import("qwen35.zig").open(a, gpa, io, o, problem); if (std.mem.eql(u8, o.model_type, "step5_vl")) return @import("step5.zig").open(a, gpa, io, o, problem); if (std.mem.eql(u8, o.model_type, "spark2_5")) return @import("spark25.zig").open(a, gpa, io, o, problem); + if (std.mem.eql(u8, o.model_type, "llama")) return @import("llama.zig").open(a, gpa, io, o, problem); if (std.mem.eql(u8, o.model_type, "lfm2_moe")) return @import("lfm2_moe.zig").open(a, gpa, io, o, problem); if (!std.mem.eql(u8, o.model_type, "nemotron_h")) { problem.* = try std.fmt.allocPrint(a, "the native engine has no backend for {s} checkpoints yet; the Python engine 0.6.6 may serve them: python -m pip install git+https://github.com/ashhart/TensorFold.git@v0.6.6", .{o.model_type}); @@ -334,14 +332,15 @@ fn refusal(a: Allocator, o: api.Open, dense: bool) !?[]const u8 { if (o.kv8) return "--kv-cache fp8 keeps fp8 caches on CUDA; the Metal engines keep bf16: serve without it"; if (o.prompt_int8) return prompt_int8_refusal; if (o.slide and !std.mem.eql(u8, t, "nemotron_h")) return try std.fmt.allocPrint(a, "--slide: {s} checkpoints do not learn on Macs yet", .{t}); - const one_mac: ?[]const u8 = if (std.mem.eql(u8, t, "nemotron_h")) "Nemotron" else if (std.mem.eql(u8, t, "lfm2_moe")) "LFM2.5-24B-A2B" else if (isGemma4(t)) "Gemma 4" else if (std.mem.eql(u8, t, "step5_vl")) "Step-5" else if (std.mem.eql(u8, t, "spark2_5")) "Spark-X2.6" else if (!std.mem.eql(u8, t, "qwen3_5")) null else if (dense) "Qwen3.8-27B" else "Qwen3.5-2B"; + if (o.drafter != null and !dense and !isGemma4(t)) return "--drafter is supported for Qwen3.8-27B-class checkpoints (qwen3_5 with its own output head) and Gemma 4 dense"; + const one_mac: ?[]const u8 = if (std.mem.eql(u8, t, "nemotron_h")) "Nemotron" else if (std.mem.eql(u8, t, "lfm2_moe")) "LFM2.5-24B-A2B" else if (isGemma4(t)) "Gemma 4" else if (std.mem.eql(u8, t, "step5_vl")) "Step-5" else if (std.mem.eql(u8, t, "spark2_5")) "Spark-X2.6" else if (std.mem.eql(u8, t, "llama")) "Llama" else if (!std.mem.eql(u8, t, "qwen3_5")) null else if (dense) "Qwen3.8-27B" else "Qwen3.5-2B"; if (o.speed_up != null) if (one_mac) |name| return try std.fmt.allocPrint(a, "--speed-up joins two Macs for Flash Next, GLM-5.3-Flash and GLM-5.3 only: serve {s} without it", .{name}); if (o.recall and std.mem.eql(u8, t, "qwen4_exp")) return recall_refusal; if (spill.refusal(o)) |why| return why; if (o.learn == null) return null; // the 27B writes conversation states if (o.learn_conversations and std.mem.eql(u8, t, "qwen3_5") and dense) return null; - const name: []const u8 = if (std.mem.eql(u8, t, "glm_moe_dsa")) "GLM-5.3 (EXLR8)" else if (std.mem.eql(u8, t, "lfm2_moe")) "LFM2.5-24B-A2B" else if (isGemma4(t)) "Gemma 4" else if (std.mem.eql(u8, t, "step5_vl")) "Step-5" else if (std.mem.eql(u8, t, "spark2_5")) "Spark-X2.6" else if (!std.mem.eql(u8, t, "qwen3_5")) return null else if (dense) "Qwen3.8-27B" else "Qwen3.5-2B"; + const name: []const u8 = if (std.mem.eql(u8, t, "glm_moe_dsa")) "GLM-5.3 (EXLR8)" else if (std.mem.eql(u8, t, "lfm2_moe")) "LFM2.5-24B-A2B" else if (isGemma4(t)) "Gemma 4" else if (std.mem.eql(u8, t, "step5_vl")) "Step-5" else if (std.mem.eql(u8, t, "spark2_5")) "Spark-X2.6" else if (std.mem.eql(u8, t, "llama")) "Llama" else if (!std.mem.eql(u8, t, "qwen3_5")) return null else if (dense) "Qwen3.8-27B" else "Qwen3.5-2B"; if (o.learn_conversations) return try std.fmt.allocPrint(a, "{s} keeps conversation states on disk for Qwen3.8-27B, Nemotron, Flash Next and GLM-5.3-Flash: serve {s} without it", .{ spill.flag(o), name }); const or_talks: []const u8 = if (dense) " (Qwen3.8-27B keeps conversation states on disk with --learn-conversations)" else ""; return try std.fmt.allocPrint(a, "--learn learns prompt prefixes for Nemotron, Flash Next and GLM-5.3-Flash only{s}: serve {s} without --learn, --learn-dir or --learn-gib", .{ or_talks, name }); diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index 4fa4fbe147..f581ec5934 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -56,6 +56,7 @@ pub const qwen35 = @import("families/qwen3_5/qwen3_5.zig"); pub const gemma4 = @import("families/gemma4/gemma4.zig"); pub const step5 = @import("families/step5/step5.zig"); pub const spark25_metal = @import("families/spark25/metal.zig"); +pub const llama = @import("families/llama/llama.zig"); pub const lfm2_moe = @import("families/lfm2_moe/lfm2_moe.zig"); pub const flashnext_pack = @import("families/flashnext/pack.zig"); @@ -85,6 +86,7 @@ test { _ = @import("families/qwen3_5/backend.zig"); _ = @import("families/step5/step5.zig"); _ = @import("families/spark25/metal.zig"); + _ = @import("families/llama/llama.zig"); _ = @import("families/lfm2_moe/config.zig"); _ = @import("families/lfm2_moe/backend.zig"); _ = @import("families/lfm2_moe/state.zig"); diff --git a/zig/tests/llama/run.zig b/zig/tests/llama/run.zig new file mode 100644 index 0000000000..2c70c10e32 --- /dev/null +++ b/zig/tests/llama/run.zig @@ -0,0 +1,193 @@ +//! Llama on Metal on real weights: greedy tokens, every position's logits or forced steps, for a reference. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const fl = tf.llama; + +const usage = + \\tf-llama-run greedy the prompt in chunks, then n one-row steps' tokens + \\tf-llama-run logits 16-row windows from 0; bf16 logits [tokens, vocab] + \\tf-llama-run prompt the prompt in chunks; the last row's bf16 logits + \\tf-llama-run forced , then cont's tokens a step; bf16 logits before each + \\tf-llama-run exact cont's tokens as one-row steps and as 16-row windows + \\ after the prompt: every row's logits bit for bit + \\tf-llama-run bench the prompt, then n one-row steps: wall, encode, GPU ms + \\LLAMA_LAYERS=N runs the first N layers only (a truncated reference's check) + \\ +; + +pub fn main(init: std.process.Init) !void { + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len != 5) { + std.debug.print("{s}", .{usage}); + return error.ExpectedModelTokensModeArg; + } + const gpa = init.gpa; + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const bytes = try std.Io.Dir.cwd().readFileAlloc(init.io, args[2], gpa, .limited(64 << 20)); + defer gpa.free(bytes); + const ids = try std.json.parseFromSlice([]u32, gpa, bytes, .{}); + defer ids.deinit(); + const prompt = ids.value; + if (prompt.len == 0) return error.EmptyPrompt; + const m = try fl.Model.load(gpa, init.io, args[1]); + defer m.deinit(); + // LLAMA_LAYERS=N: the first N layers only, for a layer-by-layer check against a truncated reference + if (std.c.getenv("LLAMA_LAYERS")) |v| m.weights.layers = m.weights.layers[0..@min(m.weights.layers.len, try std.fmt.parseInt(usize, std.mem.span(v), 10))]; + const mode = args[3]; + var forced: ?std.json.Parsed([]u32) = null; + defer if (forced) |f| f.deinit(); + var forced_out: []const u8 = ""; + if (std.mem.eql(u8, mode, "forced")) { + const comma = std.mem.indexOfScalar(u8, args[4], ',') orelse return error.ExpectedContCommaOut; + const text = try std.Io.Dir.cwd().readFileAlloc(init.io, args[4][0..comma], gpa, .limited(1 << 20)); + defer gpa.free(text); + forced = try std.json.parseFromSlice([]u32, gpa, text, .{}); + forced_out = args[4][comma + 1 ..]; + } + const steps: usize = if (std.mem.eql(u8, mode, "greedy") or std.mem.eql(u8, mode, "bench")) try std.fmt.parseInt(usize, args[4], 10) else if (forced) |f| f.value.len else 0; + var cache = try fl.state.Cache.init(m.device, &m.config, prompt.len + steps + 1); + defer cache.deinit(); + var small = try fl.state.Scratch.init(m.device, &m.config, fl.state.batch_rows); + defer small.deinit(); + var wide = try fl.state.Scratch.init(m.device, &m.config, fl.state.prompt_rows); + defer wide.deinit(); + const vocab: usize = m.config.vocab; + var t0 = now(); + if (std.mem.eql(u8, mode, "logits")) { + const out = try gpa.alloc(u16, prompt.len * vocab); + defer gpa.free(out); + var at: usize = 0; + while (at < prompt.len) { + const n = @min(prompt.len - at, fl.state.window_rows); + try fl.forward.run(m, &small, &cache, prompt[at..][0..n], .{ .all = true }); + @memcpy(out[at * vocab ..][0 .. n * vocab], small.logits_all.slice(u16, n * vocab)); + cache.len += n; + at += n; + } + try std.Io.Dir.cwd().writeFile(init.io, .{ .sub_path = args[4], .data = std.mem.sliceAsBytes(out) }); + std.debug.print("llama logits: {d} rows in {d} ms\n", .{ prompt.len, (now() - t0) / std.time.ns_per_ms }); + return; + } + if (std.mem.eql(u8, mode, "exact")) { + const text = try std.Io.Dir.cwd().readFileAlloc(init.io, args[4], gpa, .limited(1 << 20)); + defer gpa.free(text); + const cont = try std.json.parseFromSlice([]u32, gpa, text, .{}); + defer cont.deinit(); + const toks = cont.value; + // both caches at one capacity: the decode blocks follow it + var scache = try fl.state.Cache.init(m.device, &m.config, prompt.len + toks.len + 1); + defer scache.deinit(); + var wcache = try fl.state.Cache.init(m.device, &m.config, prompt.len + toks.len + 1); + defer wcache.deinit(); + try feed(m, &small, &wide, &scache, prompt); + try feed(m, &small, &wide, &wcache, prompt); + const one = try gpa.alloc(u16, toks.len * vocab); + defer gpa.free(one); + for (toks, 0..) |tok, i| { + try fl.forward.run(m, &small, &scache, &.{tok}, .{ .all = true }); + @memcpy(one[i * vocab ..][0..vocab], small.logits_all.slice(u16, vocab)); + scache.len += 1; + } + // LLAMA_WINDOW=W: windows of W rows (16 by default) + const width: usize = if (std.c.getenv("LLAMA_WINDOW")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else fl.state.window_rows; + if (width == 0 or width > fl.state.window_rows) return error.BadWindow; + var differ: usize = 0; + var at: usize = 0; + while (at < toks.len) { + const n = @min(toks.len - at, width); + try fl.forward.run(m, &small, &wcache, toks[at..][0..n], .{ .all = true }); + for (0..n) |r| { + if (!std.mem.eql(u16, one[(at + r) * vocab ..][0..vocab], small.logits_all.slice(u16, n * vocab)[r * vocab ..][0..vocab])) { + differ += 1; + if (differ <= 4) std.debug.print("row {d} (window row {d}) differs\n", .{ at + r, r }); + } + } + wcache.len += n; + at += n; + } + std.debug.print("llama exact at {d}: {d} of {d} rows' logits differ between one-row steps and windows\n", .{ prompt.len, differ, toks.len }); + return; + } + try feed(m, &small, &wide, &cache, prompt); + const ttft = (now() - t0); + if (std.mem.eql(u8, mode, "prompt")) { + try std.Io.Dir.cwd().writeFile(init.io, .{ .sub_path = args[4], .data = std.mem.sliceAsBytes(cache.logits.slice(u16, vocab)) }); + std.debug.print("llama prompt: {d} tokens in {d} ms ({d:.0} tok/s)\n", .{ prompt.len, ttft / std.time.ns_per_ms, @as(f64, @floatFromInt(prompt.len)) * 1e9 / @as(f64, @floatFromInt(ttft)) }); + return; + } + if (forced) |f| { + const out = try gpa.alloc(u16, f.value.len * vocab); + defer gpa.free(out); + for (f.value, 0..) |tok, i| { + @memcpy(out[i * vocab ..][0..vocab], cache.logits.slice(u16, vocab)); + if (i + 1 == f.value.len) break; + try fl.forward.run(m, &small, &cache, &.{tok}, .{}); + cache.len += 1; + } + try std.Io.Dir.cwd().writeFile(init.io, .{ .sub_path = forced_out, .data = std.mem.sliceAsBytes(out) }); + return; + } + if (std.mem.eql(u8, mode, "bench")) { + const n = try std.fmt.parseInt(usize, args[4], 10); + var tm: fl.forward.Timing = .{}; + t0 = now(); + for (0..n) |_| { + try fl.forward.run(m, &small, &cache, &.{argmax(cache.logits.slice(u16, vocab))}, .{ .timing = &tm }); + cache.len += 1; + } + const wall = @as(f64, @floatFromInt(now() - t0)) / 1e6 / @as(f64, @floatFromInt(n)); + const per = 1e3 / @as(f64, @floatFromInt(n)); + std.debug.print("llama bench at {d}: {d:.3} ms a step wall, {d:.3} encode, {d:.3} GPU\n", .{ prompt.len, wall, tm.encode * per, tm.gpu * per }); + return; + } + if (!std.mem.eql(u8, mode, "greedy")) { + std.debug.print("{s}", .{usage}); + return error.UnknownMode; + } + const picks = try gpa.alloc(u32, steps); + defer gpa.free(picks); + t0 = now(); + for (picks, 0..) |*p, i| { + p.* = argmax(cache.logits.slice(u16, vocab)); + if (i + 1 == steps) break; + try fl.forward.run(m, &small, &cache, &.{p.*}, .{}); + cache.len += 1; + } + const decode = (now() - t0); + var w: std.Io.Writer.Allocating = .init(gpa); + defer w.deinit(); + try std.json.Stringify.value(picks, .{}, &w.writer); + std.debug.print("{s}\n", .{w.written()}); + std.debug.print("llama greedy: prompt {d} tokens {d} ms, {d} steps {d:.1} tok/s\n", .{ prompt.len, ttft / std.time.ns_per_ms, steps, @as(f64, @floatFromInt(steps - 1)) * 1e9 / @as(f64, @floatFromInt(decode)) }); +} + +/// The prompt in prompt_rows chunks (the tensor-unit path past 16 rows), as the backend feeds it. +fn feed(m: *const fl.Model, small: *fl.state.Scratch, wide: *fl.state.Scratch, cache: *fl.state.Cache, ids: []const u32) !void { + var at: usize = 0; + while (at < ids.len) { + const n = @min(ids.len - at, fl.state.prompt_rows); + try fl.forward.run(m, if (n > fl.state.batch_rows) wide else small, cache, ids[at..][0..n], .{}); + cache.len += n; + at += n; + } +} + +fn now() u64 { + var ts: std.c.timespec = undefined; + _ = std.c.clock_gettime(.MONOTONIC, &ts); + return @as(u64, @intCast(ts.sec)) * std.time.ns_per_s + @as(u64, @intCast(ts.nsec)); +} + +fn argmax(l: []const u16) u32 { + var best: usize = 0; + for (l, 0..) |x, i| if (logit(x) > logit(l[best])) { + best = i; + }; + return @intCast(best); +} + +fn logit(word: u16) f32 { + return @bitCast(@as(u32, word) << 16); +} From 6f467523a7c3c2d76b9e15ba0cbaf59d10ec67ba Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Mon, 26 Oct 2026 13:27:47 -0500 Subject: [PATCH 2/7] llama: every kernel registered with the kernel contract The family's Metal kernels through gates 1-5 (zig/tests/contract/metal_llama*.zig, the roll call and sources in contracts.zig): the glue (embedding, the norms at HF's rounding points, the residual add, RoPE with bf16 cos and sin, the SiLU gate, the decode attention's partials and their merge, the argmax draw by its stated NaN rule), the six row-projection forms, the four GEMM forms and the prompt attention at both served geometries, each with its float64 reference, served shapes and tiles, launched as forward.zig launches it; gate 5 reads the pipelines on this GPU. tf-kernel-contract --only llama on the M5 Ultra: 21 kernels, 21 pass gates 1-4, 0 fail, 0 known findings, 0 limit problems. --- docs/kernel-contract.md | 4 + docs/recipes/nanbeige-4.3-33b-metal.md | 8 +- zig/src/families/llama/config.zig | 7 +- zig/src/families/llama/contracts.zig | 66 +++ zig/src/families/llama/llama.zig | 2 + zig/tests/contract/metal_llama.zig | 518 +++++++++++++++++++ zig/tests/contract/metal_llama_all.zig | 191 +++++++ zig/tests/contract/metal_llama_attention.zig | 511 ++++++++++++++++++ zig/tests/contract/metal_llama_dense.zig | 230 ++++++++ zig/tests/contract/metal_main.zig | 5 + 10 files changed, 1540 insertions(+), 2 deletions(-) create mode 100644 zig/src/families/llama/contracts.zig create mode 100644 zig/tests/contract/metal_llama.zig create mode 100644 zig/tests/contract/metal_llama_all.zig create mode 100644 zig/tests/contract/metal_llama_attention.zig create mode 100644 zig/tests/contract/metal_llama_dense.zig diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 1ed43ac146..f59d204984 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -141,6 +141,10 @@ Metal also checks Nemotron's routed experts on the engine's route tables, one ro against float64 at tf_load16's bf16 x sums (up's relu2 as an interval over its two bf16 roundings), and the rows2, rows4 and six geometry pipelines of each projection bit-equal to them (two controls: member rows read from one row, and w3's x sums without their bf16 steps). +The Llama family's own kernels (`metal_llama*.zig`): its glue (embedding, the norms at HF's rounding points, the +residual add, RoPE with bf16 cos and sin, the SiLU gate, the decode attention's partials and their merge, the argmax +draw by its stated rule), its row projections (1 to 32 rows) and, on the M5's tensor units, its four GEMM forms and the +prompt attention at both served geometries; gate 5 reads their pipelines on this GPU. CUDA: the RMS norm rows (16-byte form and the fp32-out tree), the bf16 residual's norm tails (MLX's rounded rows and the once-rounded ones, each checked on its residual and its output, with four controls each dropping a rounding point), Nemotron's router logits (and their 128-row tiles, bit-equal to them) and top-k, the latent top-k's decode form (bit-equal to `tf_super_topk`), the latent combine and diff --git a/docs/recipes/nanbeige-4.3-33b-metal.md b/docs/recipes/nanbeige-4.3-33b-metal.md index 3fa56775d4..a30a2eb0c7 100644 --- a/docs/recipes/nanbeige-4.3-33b-metal.md +++ b/docs/recipes/nanbeige-4.3-33b-metal.md @@ -43,7 +43,13 @@ calls take the template's XML form (its `tool_call_format` default), which the s - The lane host's shared rounds, a lone stream's steps drawn on the GPU a step ahead (argmax, or `tf_sample_full`), copy drafts and the prompt cache. -The Metal glue kernels are not registered with the kernel contract, as with the other Metal family glue. +Every kernel is registered with the kernel contract (`zig/tests/contract/metal_llama*.zig`, the roll call in +`contracts.zig`): the glue, the six row-projection forms, the four GEMM forms and the prompt attention at both served +geometries. `tf-kernel-contract --only llama` on the M5 Ultra: + +``` +SUMMARY 21 kernels: 21 pass gates 1-4, 0 fail, 0 known findings; 0 controls: 0 red as planted, 0 missed; 0 limit problems; gates 432063 ms, limits 105 ms +``` ## Measured (M5 Ultra, the bf16 release, 26 Oct 2026) diff --git a/zig/src/families/llama/config.zig b/zig/src/families/llama/config.zig index e7631560b6..5db0942bab 100644 --- a/zig/src/families/llama/config.zig +++ b/zig/src/families/llama/config.zig @@ -16,6 +16,11 @@ pub const Error = error{ UnsupportedLlamaPrecision, } || std.json.ParseError(std.json.Scanner); +/// Attention's score scale, 1 / sqrt(head_dim). +pub fn headScale() f32 { + return 1.0 / @sqrt(@as(f32, head_dim)); +} + pub const Config = struct { hidden: usize, intermediate: usize, @@ -57,7 +62,7 @@ pub const Config = struct { /// 1 / sqrt(head_dim), the attention scale. pub fn scale(_: Config) f32 { - return 1.0 / @sqrt(@as(f32, head_dim)); + return headScale(); } pub fn isEos(c: Config, token: u32) bool { diff --git a/zig/src/families/llama/contracts.zig b/zig/src/families/llama/contracts.zig new file mode 100644 index 0000000000..46770e40a7 --- /dev/null +++ b/zig/src/families/llama/contracts.zig @@ -0,0 +1,66 @@ +//! The family's Metal kernels under the kernel contract: their sources as the loader compiles them, and the roll call. +const std = @import("std"); +const mtl = @import("metal"); +const ks = @import("kernel_sources"); +const cfg = @import("config.zig"); +const model = @import("model.zig"); +const frags = @import("../../core/frags.zig"); + +/// A served geometry the attention compiles for (the head size is the family's 128). +pub const Geometry = struct { name: []const u8, heads: u32, kv_heads: u32, hidden: usize, intermediate: usize, vocab: usize }; +/// Nanbeige4.3-33B, the family's first model, and Nanbeige4.1-3B, its drafter. +pub const nanbeige43: Geometry = .{ .name = "33B", .heads = 56, .kv_heads = 8, .hidden = 7168, .intermediate = 22016, .vocab = 166912 }; +pub const nanbeige41: Geometry = .{ .name = "3B", .heads = 20, .kv_heads = 4, .hidden = 2560, .intermediate = 10496, .vocab = 166144 }; +pub const geometries = [_]Geometry{ nanbeige43, nanbeige41 }; + +pub const glue_names = model.glue_names; +pub const gemm_names = [_][]const u8{ "ll_gemm_00", "ll_gemm_01", "ll_gemm_10", "ll_gemm_11" }; +pub const attention_name = "ll_prompt_attention"; +/// Every pipeline the family runs: the glue, the four GEMM forms and the prompt attention. +pub const pipelines = glue_names.len + gemm_names.len + 1; + +/// The glue's text as model.zig builds it. +pub fn glueSource(a: std.mem.Allocator) ![]u8 { + return std.fmt.allocPrint(a, "#define LL_DH {d}\n{s}", .{ cfg.head_dim, ks.llama_glue }); +} + +/// The glue library as model.zig compiles it: precise math functions. +pub fn glueLibrary(device: mtl.Device, a: std.mem.Allocator) !mtl.Library { + const src = try glueSource(a); + defer a.free(src); + return mtl.Library.fromSource(device, src, .{ .language = mtl.CompileOptions.mlx().language, .math_functions = .precise }); +} + +/// The prompt attention's text for a geometry, as model.zig builds it (nax.h inlined for this GPU). +pub fn attentionSource(device: mtl.Device, a: std.mem.Allocator, g: Geometry) ![]u8 { + const body = try frags.source(device, a, ks.llama_attention_nax); + defer a.free(body); + return std.fmt.allocPrint(a, "#define DH {d}\n#define NQ {d}\n#define G {d}\n#define KVH {d}\n#define SCALE {e}f\n{s}", .{ cfg.head_dim, g.heads, g.heads / g.kv_heads, g.kv_heads, cfg.headScale(), body }); +} + +/// The GEMM forms' text as model.zig builds it. +pub fn gemmSource(device: mtl.Device, a: std.mem.Allocator) ![]u8 { + return frags.source(device, a, ks.llama_gemm); +} + +/// Every pipeline the family runs appears among the registrations' functions. +pub fn requireComplete(functions: []const []const u8) !void { + inline for (.{ glue_names, gemm_names, [_][]const u8{attention_name} }) |list| for (list) |name| { + var seen = false; + for (functions) |f| if (std.mem.eql(u8, f, name)) { + seen = true; + }; + if (!seen) return error.IncompleteLlamaContract; + }; +} + +test "llama contracts: the roll call wants every pipeline" { + var all: [pipelines][]const u8 = undefined; + var i: usize = 0; + inline for (.{ glue_names, gemm_names, [_][]const u8{attention_name} }) |list| for (list) |name| { + all[i] = name; + i += 1; + }; + try requireComplete(&all); + try std.testing.expectError(error.IncompleteLlamaContract, requireComplete(all[1..])); +} diff --git a/zig/src/families/llama/llama.zig b/zig/src/families/llama/llama.zig index 4eca3dc970..ecaa37ae83 100644 --- a/zig/src/families/llama/llama.zig +++ b/zig/src/families/llama/llama.zig @@ -1,5 +1,6 @@ //! Native Metal Llama (model_type llama): GQA attention and a SwiGLU MLP in bf16 at HF's rounding points. pub const config = @import("config.zig"); +pub const contracts = @import("contracts.zig"); pub const weights = @import("weights.zig"); pub const Model = @import("model.zig").Model; pub const state = @import("state.zig"); @@ -11,6 +12,7 @@ pub const snapshot = @import("snapshot.zig"); test { _ = config; + _ = contracts; _ = state; _ = budget; _ = snapshot; diff --git a/zig/tests/contract/metal_llama.zig b/zig/tests/contract/metal_llama.zig new file mode 100644 index 0000000000..1c5c55993d --- /dev/null +++ b/zig/tests/contract/metal_llama.zig @@ -0,0 +1,518 @@ +//! Llama's glue, row projections and GEMM under the contract, compiled as the loader compiles them, launched as served. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const c = contract.spec; +const g = contract.gates; +const rows_in = contract.rows; +pub const family = tf.llama.contracts; + +pub const DH = 128; +const G33 = family.nanbeige43; +const G3 = family.nanbeige41; + +/// glue.metal's argument structs, as forward.zig fills them. +pub const Args = extern struct { rows: u32 = 0, width: u32 = 0, heads: u32 = 0, eps: f32 = 0 }; +const PArgs = extern struct { rows: u32, heads: u32, kv_heads: u32, pos: u32, stride: u32 }; + +/// A value through fp32 and then bf16, as the kernels round. +pub fn bf(v: f64) f64 { + return c.fromBf16(c.toBf16(@floatCast(v))); +} + +fn none(_: c.Shape) usize { + return 0; +} + +/// One thread's fma chain over the row (256 threads), the simd tree and the eight partials. +fn normDepth(s: c.Shape) usize { + return s.k / 256 + 16; +} + +const no_sums: c.Sums = .{ .f16 = 0, .bf16 = 0, .why = "no sums: elementwise, a gather or a draw" }; +// the norm's inner bf16(x * inv) rounds either way under the fp32 sum's error (mids 3): fp16's overflow misses 64x +const rms_why = "a bf16 output follows the sum of squares as ss^-1/2: a bf16 sum's error is a few output ulps"; +const rms_sums: c.Sums = .{ .f16 = 50, .bf16 = 0, .why = rms_why }; +const argmax_rule = "a NaN logit wins (the lowest NaN id); else the largest logit, ties to the lower id"; + +pub fn pipelineOf(lib: mtl.Library, device: mtl.Device, name: []const u8) !mtl.Pipeline { + var buf: [48]u8 = undefined; + const z = try std.fmt.bufPrintSentinel(&buf, "{s}", .{name}, 0); + return mtl.Pipeline.init(device, lib, z, false); +} + +pub const Which = enum { embed, norm, add, add_norm, silu, prep }; + +pub fn function(w: Which) []const u8 { + return switch (w) { + .embed => "ll_embed", + .norm => "ll_norm", + .add => "ll_add", + .add_norm => "ll_add_norm", + .silu => "ll_silu_gu", + .prep => "ll_prep", + }; +} + +pub fn spec(w: Which) c.Spec { + const widths = &[_]c.Shape{ .{ .rows = 32, .n = G33.hidden, .k = G33.hidden }, .{ .rows = 2048, .n = G3.hidden, .k = G3.hidden } }; + const anyk: [3]c.Tile = .{ .{}, .{}, .{ .size = 256, .step = 1, .min = 1 } }; + return switch (w) { + .embed => .{ .name = "llama glue ll_embed", .accum = .none, .depth = none, .out = .bf16, .served = widths, .sums = no_sums }, + .norm => .{ .name = "llama glue ll_norm", .accum = .f32, .depth = normDepth, .ops = 3, .mids = 3, .out = .bf16, .tiles = anyk, .served = widths, .sums = rms_sums }, + .add => .{ .name = "llama glue ll_add", .accum = .none, .depth = none, .out = .bf16, .in_place = true, .served = widths, .sums = no_sums }, + .add_norm => .{ .name = "llama glue ll_add_norm", .accum = .f32, .depth = normDepth, .ops = 3, .mids = 3, .out = .bf16, .in_place = true, .tiles = anyk, .served = widths, .sums = rms_sums }, + .silu => .{ .name = "llama glue ll_silu_gu", .accum = .none, .depth = none, .ops = 3, .mids = 1, .out = .bf16, .served = &.{ .{ .rows = 32, .n = G33.intermediate, .k = G33.intermediate }, .{ .rows = 2048, .n = G3.intermediate, .k = G3.intermediate } }, .sums = no_sums }, + .prep => .{ .name = "llama glue ll_prep", .accum = .none, .depth = none, .ops = 5, .mids = 2, .out = .bf16, .served = &.{ .{ .rows = 32, .n = G33.heads + G33.kv_heads, .k = (G33.heads + 2 * G33.kv_heads) * DH }, .{ .rows = 2048, .n = G3.heads + G3.kv_heads, .k = (G3.heads + 2 * G3.kv_heads) * DH } }, .sums = no_sums }, + }; +} + +/// The glue kernels: the width is the shape's k (n says the same), the rows its rows. +pub const Glue = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + which: Which, + lib: mtl.Library, + pipe: mtl.Pipeline, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + geo: family.Geometry = G33, + pos: u32 = 0, + r: ?rows_in.Rows = null, + x: []u16 = &.{}, // the residual's first contents (add, add_norm), gu (silu), q|k|v (prep), the logits (argmax) + y: []u16 = &.{}, // the added rows (add, add_norm) + table: []u16 = &.{}, + ids: []u32 = &.{}, + inv: []f32 = &.{}, + in: [3]rig_.Ref = undefined, + out: [3]rig_.Ref = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, which: Which) !Glue { + const lib = try family.glueLibrary(device, a); + errdefer lib.deinit(); + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .which = which, .lib = lib, .pipe = try pipelineOf(lib, device, function(which)) }; + } + + fn drop(p: *Glue) void { + if (p.r) |r| r.deinit(); + p.r = null; + for ([_]*[]u16{ &p.x, &p.y, &p.table }) |s| { + p.a.free(s.*); + s.* = &.{}; + } + p.a.free(p.ids); + p.ids = &.{}; + p.a.free(p.inv); + p.inv = &.{}; + p.rig.reset(); + } + + pub fn deinit(p: *Glue) void { + p.drop(); + p.rig.deinit(); + p.pipe.deinit(); + p.lib.deinit(); + } + + const vt = g.hooks(Glue); + + pub fn kernel(p: *Glue) g.Kernel { + return .{ .spec = spec(p.which), .ctx = p, .vt = &vt }; + } + + fn width(p: *const Glue) usize { + return p.shape.k; + } + + fn random16(rng: std.Random, out: []u16, lo: f32, hi: f32) void { + for (out) |*v| v.* = c.toBf16(lo + rng.float(f32) * (hi - lo)); + } + + /// NaN, +Inf and -Inf into rows 0, 1 and 2 of `v` (rows `w` wide), as the input families plant them. + fn plant(v: []u16, w: usize, rows: usize) void { + const at = [_][2]usize{ .{ 0, w / 3 }, .{ 1, w / 2 }, .{ 2, w - 1 } }; + const put = [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32) }; + for (at, put) |x, f| if (x[0] < rows) { + v[x[0] * w + x[1]] = c.toBf16(f); + }; + } + + pub fn prepare(p: *Glue, s: c.Shape, fill: g.Fill, seed: u64) !bool { + if (s.rows == 0 or s.k == 0) return false; + p.drop(); + p.shape = s; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const w = s.k; + const n = s.rows * w; + switch (p.which) { + .embed => { + p.table = try p.a.alloc(u16, 8 * w); + random16(rng, p.table, -2, 2); + if (fill == .nonfinite) plant(p.table, w, 8); + p.ids = try p.a.alloc(u32, s.rows); + for (p.ids, 0..) |*id, i| id.* = @intCast(i % 8); + p.in[0] = try p.rig.input(std.mem.sliceAsBytes(p.ids)); + p.in[1] = try p.rig.input(std.mem.sliceAsBytes(p.table)); + p.out[0] = try p.rig.output(n * 2); + }, + .norm, .add_norm => { + const r = try rows_in.Rows.init(p.a, .rms, .{ .rows = s.rows, .n = w, .k = w }); + p.r = r; + r.fill(if (p.which == .add_norm and fill == .nonfinite) .random else fill, seed); + if (p.which == .norm) { + p.in[0] = try p.rig.input(std.mem.sliceAsBytes(r.x)); + p.in[1] = try p.rig.input(std.mem.sliceAsBytes(r.w)); + p.out[0] = try p.rig.output(n * 2); + } else { + // the added rows: zero for the sums gate (the residual's stress rows stay whole), else small + p.y = try p.a.alloc(u16, n); + if (fill == .stress) @memset(p.y, 0) else random16(rng, p.y, -0.5, 0.5); + if (fill == .nonfinite) plant(p.y, w, s.rows); + p.in[0] = try p.rig.input(std.mem.sliceAsBytes(p.y)); + p.in[1] = try p.rig.input(std.mem.sliceAsBytes(r.w)); + p.out[0] = try p.rig.output(n * 2); + p.out[1] = try p.rig.output(n * 2); + } + }, + .add => { + p.x = try p.a.alloc(u16, n); + p.y = try p.a.alloc(u16, n); + random16(rng, p.x, -4, 4); + random16(rng, p.y, -4, 4); + if (fill == .nonfinite) { + plant(p.y, w, s.rows); + if (s.rows > 3) { // +Inf added to -Inf: a NaN out of two infinities + p.x[3 * w] = c.toBf16(std.math.inf(f32)); + p.y[3 * w] = c.toBf16(-std.math.inf(f32)); + } + } + p.in[0] = try p.rig.input(std.mem.sliceAsBytes(p.y)); + p.out[0] = try p.rig.output(n * 2); + }, + .silu => { + p.x = try p.a.alloc(u16, 2 * n); + random16(rng, p.x, -6, 6); + if (fill == .nonfinite) plant(p.x, 2 * w, s.rows); + p.in[0] = try p.rig.input(std.mem.sliceAsBytes(p.x)); + p.out[0] = try p.rig.output(n * 2); + }, + .prep => { + p.geo = if (s.n == G3.heads + G3.kv_heads) G3 else G33; + const stride = (p.geo.heads + 2 * p.geo.kv_heads) * DH; + if (s.k != stride) return false; + p.pos = if (s.rows <= 64) 8192 else 0; // a window past a long prompt, or a prompt's first chunk + p.x = try p.a.alloc(u16, s.rows * stride); + random16(rng, p.x, -3, 3); + if (fill == .nonfinite) plant(p.x, stride, s.rows); + p.inv = try p.a.alloc(f32, DH / 2); + for (p.inv, 0..) |*f, i| f.* = @floatCast(1.0 / std.math.pow(f64, 1.2e8, @as(f64, @floatFromInt(2 * i)) / DH)); + p.in[0] = try p.rig.input(std.mem.sliceAsBytes(p.x)); + p.in[1] = try p.rig.input(std.mem.sliceAsBytes(p.inv)); + p.out[0] = try p.rig.output(s.rows * p.geo.heads * DH * 2); + p.out[1] = try p.rig.output((p.pos + s.rows) * p.geo.kv_heads * DH * 2); + p.out[2] = try p.rig.output((p.pos + s.rows) * p.geo.kv_heads * DH * 2); + }, + } + return true; + } + + pub fn count(p: *Glue) usize { + const n = p.shape.rows * p.width(); + return switch (p.which) { + .add_norm => 2 * n, + .prep => p.shape.rows * p.geo.heads * DH + 2 * (p.pos + p.shape.rows) * p.geo.kv_heads * DH, + else => n, + }; + } + + pub fn run(p: *Glue) !void { + const s = p.shape; + const w: u32 = @intCast(p.width()); + const rows: u32 = @intCast(s.rows); + const pass = p.rig.begin(); + const e = pass.e; + e.setPipeline(p.pipe); + const grid = mtl.Size.of(p.width(), s.rows, 1); + const threads = mtl.Size.of(256, 1, 1); + switch (p.which) { + .embed => { + e.setBuffer(p.in[0].buf, p.in[0].off, 0); + e.setBuffer(p.in[1].buf, p.in[1].off, 1); + e.setValue(Args{ .rows = rows, .width = w }, 2); + e.setBuffer(p.out[0].buf, p.out[0].off, 3); + e.dispatchThreads(grid, threads); + }, + .norm => { + e.setBuffer(p.in[0].buf, p.in[0].off, 0); + e.setBuffer(p.in[1].buf, p.in[1].off, 1); + e.setValue(Args{ .rows = rows, .width = w, .eps = p.r.?.eps }, 2); + e.setBuffer(p.out[0].buf, p.out[0].off, 3); + e.dispatchGroups(mtl.Size.of(s.rows, 1, 1), threads); + }, + .add => { + @memcpy(p.rig.bytes(p.out[0], p.x.len * 2), std.mem.sliceAsBytes(p.x)); // the residual before the add + e.setBuffer(p.out[0].buf, p.out[0].off, 0); + e.setBuffer(p.in[0].buf, p.in[0].off, 1); + e.setValue(Args{ .rows = rows, .width = w }, 2); + e.dispatchThreads(grid, threads); + }, + .add_norm => { + const r = p.r.?; + @memcpy(p.rig.bytes(p.out[0], r.x.len * 2), std.mem.sliceAsBytes(r.x)); + e.setBuffer(p.out[0].buf, p.out[0].off, 0); + e.setBuffer(p.in[0].buf, p.in[0].off, 1); + e.setBuffer(p.in[1].buf, p.in[1].off, 2); + e.setValue(Args{ .rows = rows, .width = w, .eps = r.eps }, 3); + e.setBuffer(p.out[1].buf, p.out[1].off, 4); + e.dispatchGroups(mtl.Size.of(s.rows, 1, 1), threads); + }, + .silu => { + e.setBuffer(p.in[0].buf, p.in[0].off, 0); + e.setValue(Args{ .rows = rows, .width = w }, 1); + e.setBuffer(p.out[0].buf, p.out[0].off, 2); + e.dispatchThreads(grid, threads); + }, + .prep => { + // the cache's slots before the window: unwritten by the kernel, zero in the reference + for (p.out[1..]) |o| @memset(p.rig.bytes(o, p.pos * p.geo.kv_heads * DH * 2), 0); + e.setBuffer(p.in[0].buf, p.in[0].off, 0); + e.setBuffer(p.in[1].buf, p.in[1].off, 1); + e.setValue(PArgs{ .rows = rows, .heads = p.geo.heads, .kv_heads = p.geo.kv_heads, .pos = p.pos, .stride = w }, 2); + for (p.out, 3..) |o, slot| e.setBuffer(o.buf, o.off, slot); + e.dispatchGroups(mtl.Size.of(p.geo.heads + p.geo.kv_heads, s.rows, 1), mtl.Size.of(32, 1, 1)); + }, + } + try p.rig.finish(pass); + } + + fn xv(v: []const u16, i: usize) f64 { + return c.fromBf16(v[i]); + } + + /// The norm of a row held in `x` (bf16 values), HF's rounding points, into want and mag. + fn normRow(p: *Glue, x: []const f64, ws: []const u16, want: []f64, mag: []f64) void { + var ss: f64 = 0; + for (x) |v| ss += v * v; + const inv = 1.0 / @sqrt(ss / @as(f64, @floatFromInt(x.len)) + @as(f64, p.r.?.eps)); + for (x, 0..) |v, i| { + want[i] = bf(c.fromBf16(ws[i]) * bf(v * inv)); + mag[i] = @abs(want[i]) / 2; + } + } + + pub fn reference(p: *Glue, want: []f64, mag: []f64) !void { + const s = p.shape; + const w = p.width(); + switch (p.which) { + .embed => for (0..s.rows) |r| for (0..w) |i| { + want[r * w + i] = xv(p.table, p.ids[r] * w + i); + mag[r * w + i] = @abs(want[r * w + i]); + }, + .norm => { + const x = try p.a.alloc(f64, w); + defer p.a.free(x); + for (0..s.rows) |r| { + for (x, 0..) |*v, i| v.* = xv(p.r.?.x, r * w + i); + p.normRow(x, p.r.?.w, want[r * w ..][0..w], mag[r * w ..][0..w]); + } + }, + .add => for (0..s.rows * w) |i| { + want[i] = bf(xv(p.x, i) + xv(p.y, i)); + mag[i] = @abs(want[i]); + }, + .add_norm => { + const x = try p.a.alloc(f64, w); + defer p.a.free(x); + const n = s.rows * w; + for (0..s.rows) |r| { + for (x, 0..) |*v, i| { + v.* = bf(xv(p.r.?.x, r * w + i) + xv(p.y, r * w + i)); + want[r * w + i] = v.*; + mag[r * w + i] = @abs(v.*); + } + p.normRow(x, p.r.?.w, want[n + r * w ..][0..w], mag[n + r * w ..][0..w]); + } + }, + .silu => for (0..s.rows) |r| for (0..w) |i| { + const gate = xv(p.x, r * 2 * w + i); + const up = xv(p.x, r * 2 * w + w + i); + const silu = bf(gate / (1.0 + @exp(-gate))); + want[r * w + i] = bf(silu * up); + mag[r * w + i] = @abs(want[r * w + i]); + }, + .prep => try p.prepReference(want, mag), + } + } + + /// HF's rotation of q and k (fp32 angle, bf16 cos and sin), v copied: QN, then caches of pos + rows slots. + fn prepReference(p: *Glue, want: []f64, mag: []f64) !void { + const s = p.shape; + const geo = p.geo; + const stride = (geo.heads + 2 * geo.kv_heads) * DH; + const kvw = geo.kv_heads * DH; + const qn = s.rows * geo.heads * DH; + const kc = (p.pos + s.rows) * kvw; + @memset(want[qn .. qn + 2 * kc], 0); + @memset(mag[qn .. qn + 2 * kc], 0); + for (0..s.rows) |r| { + const position: f32 = @floatFromInt(p.pos + r); + for (0..geo.heads + geo.kv_heads) |unit| { + const q = unit < geo.heads; + const head = if (q) unit else unit - geo.heads; + const src = r * stride + (if (q) 0 else geo.heads * DH) + head * DH; + const dst = if (q) (r * geo.heads + head) * DH else qn + (p.pos + r) * kvw + head * DH; + for (0..DH) |d| { + const i = d % (DH / 2); + const angle: f32 = position * p.inv[i]; // the kernel's fp32 angle + const cs = bf(@cos(@as(f64, angle))); + const sn = bf(@sin(@as(f64, angle))); + const x = xv(p.x, src + d); + const other = if (d < DH / 2) -xv(p.x, src + d + DH / 2) else xv(p.x, src + d - DH / 2); + want[dst + d] = bf(bf(x * cs) + bf(other * sn)); + mag[dst + d] = @abs(x * cs) + @abs(other * sn); + } + } + for (0..geo.kv_heads) |head| for (0..DH) |d| { + const at = qn + kc + (p.pos + r) * kvw + head * DH + d; + want[at] = xv(p.x, r * stride + (geo.heads + geo.kv_heads) * DH + head * DH + d); + mag[at] = 0; + }; + } + } + + /// The rotation's allowance: a bf16 ulp of each term (its rounding, its cos or sin), then the output's rounding. + pub fn bounds(p: *Glue, want: []const f64, mag: []const f64, allow: []f64) !void { + if (p.which != .prep) return; + for (allow, want, mag) |*l, w, m| l.* = 0x1p-7 * m + 0x1p-8 * @abs(w) + 0x1p-133; + } + + pub fn simulate(p: *Glue, acc: c.Accum, order: g.Order, out: []f64) !void { + const s = p.shape; + const w = p.width(); + const t = try p.a.alloc(f64, w); + defer p.a.free(t); + const r = p.r.?; + const base: usize = if (p.which == .add_norm) s.rows * w else 0; + for (0..s.rows) |row| { + for (t, 0..) |*v, i| { + const x = if (p.which == .add_norm) bf(xv(r.x, row * w + i) + xv(p.y, row * w + i)) else xv(r.x, row * w + i); + if (p.which == .add_norm) out[row * w + i] = x; + v.* = x * x; + } + const ss = rows_in.sumIn(acc, order, t); + const inv = 1.0 / @sqrt(ss / @as(f64, @floatFromInt(w)) + @as(f64, r.eps)); + for (0..w) |i| { + const x = if (p.which == .add_norm) out[row * w + i] else xv(r.x, row * w + i); + out[base + row * w + i] = bf(c.fromBf16(r.w[i]) * bf(x * inv)); + } + } + } +}; + +/// ll_argmax, a draw a logits row: the stated rule checks the ids (gates 1 and 3), the bytes repeat (gate 2). +pub const Argmax = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: mtl.Library, + pipe: mtl.Pipeline, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + x: []u16 = &.{}, + in: rig_.Ref = undefined, + out: rig_.Ref = undefined, + msg: [160]u8 = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue) !Argmax { + const lib = try family.glueLibrary(device, a); + errdefer lib.deinit(); + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .lib = lib, .pipe = try pipelineOf(lib, device, "ll_argmax") }; + } + + pub fn deinit(p: *Argmax) void { + p.a.free(p.x); + p.rig.deinit(); + p.pipe.deinit(); + p.lib.deinit(); + } + + const vt = g.hooks(Argmax); + + pub fn kernel(p: *Argmax) g.Kernel { + return .{ .spec = .{ .name = "llama glue ll_argmax", .accum = .none, .depth = none, .out = .raw, .nan = .{ .stated = argmax_rule }, .tiles = .{ .{}, .{}, .{ .size = 1024, .step = 1, .min = 1 } }, .served = &.{ .{ .rows = 1, .n = 1, .k = G33.vocab }, .{ .rows = 1, .n = 1, .k = G3.vocab } }, .sums = no_sums }, .ctx = p, .vt = &vt }; + } + + /// rows logits rows of k: NaNs and infinities as planted, two NaNs in row 3 and a tie in row 4. + pub fn prepare(p: *Argmax, s: c.Shape, fill: g.Fill, seed: u64) !bool { + if (s.rows == 0 or s.k == 0) return false; + p.a.free(p.x); + p.x = &.{}; + p.rig.reset(); + p.shape = s; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const w = s.k; + p.x = try p.a.alloc(u16, s.rows * w); + for (p.x) |*v| v.* = c.toBf16(rng.float(f32) * 16 - 8); + if (fill == .nonfinite) { + Glue.plant(p.x, w, s.rows); + if (s.rows > 3 and w > 7) { // two NaNs: the lower id wins + p.x[3 * w + 7] = c.toBf16(std.math.nan(f32)); + p.x[3 * w + 5] = c.toBf16(std.math.nan(f32)); + } + } + if (s.rows > 4 and w > 2) { // a tie: the lower id wins + p.x[4 * w + 2] = c.toBf16(9); + p.x[4 * w + w - 1] = c.toBf16(9); + } + p.in = try p.rig.input(std.mem.sliceAsBytes(p.x)); + p.out = try p.rig.output(s.rows * 4); + return true; + } + + pub fn count(p: *Argmax) usize { + return p.shape.rows * 4; + } + + pub fn run(p: *Argmax) !void { + const s = p.shape; + const pass = p.rig.begin(); + const e = pass.e; + e.setPipeline(p.pipe); + for (0..s.rows) |r| { // a draw a logits row, as the forward draws its last row + e.setBuffer(p.in.buf, p.in.off + r * s.k * 2, 0); + e.setBuffer(p.out.buf, p.out.off + r * 4, 1); + e.setValue(@as(u32, @intCast(s.k)), 2); + e.dispatchGroups(mtl.Size.of(1, 1, 1), mtl.Size.of(1024, 1, 1)); + } + try p.rig.finish(pass); + } + + pub fn reference(_: *Argmax, want: []f64, mag: []f64) !void { + @memset(want, 0); + @memset(mag, 0); + } + + /// The argmax rule on the host against the drawn ids (gates 1 and 3). + pub fn check(p: *Argmax, _: g.Gate) !?[]const u8 { + const w = p.shape.k; + const ids = std.mem.bytesAsSlice(u32, p.rig.bytes(p.out, p.shape.rows * 4)); + for (0..p.shape.rows) |r| { + var best: f64 = -std.math.inf(f64); + var at: usize = 0; + var nan: ?usize = null; + for (0..w) |i| { + const v = c.fromBf16(p.x[r * w + i]); + if (std.math.isNan(v)) { + if (nan == null) nan = i; + } else if (v > best) { + best = v; + at = i; + } + } + const want = nan orelse at; + if (ids[r] != want) return try std.fmt.bufPrint(&p.msg, "row {d}: drew {d} where the rule gives {d}", .{ r, ids[r], want }); + } + return null; + } +}; diff --git a/zig/tests/contract/metal_llama_all.zig b/zig/tests/contract/metal_llama_all.zig new file mode 100644 index 0000000000..dd5971e34c --- /dev/null +++ b/zig/tests/contract/metal_llama_all.zig @@ -0,0 +1,191 @@ +//! The llama family's registrations as one set: every pipeline the loader builds, and gate 5 on this GPU. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const limits_k = @import("metal_limits.zig"); +const llama = @import("metal_llama.zig"); +const dense = @import("metal_llama_dense.zig"); +const attention = @import("metal_llama_attention.zig"); +const g = contract.gates; +const family = llama.family; + +const matvec_rows = [_]u32{ 1, 2, 4, 8, 16, 32 }; +const matvec_names = [_][]const u8{ "ll_matvec_1", "ll_matvec_2", "ll_matvec_4", "ll_matvec_8", "ll_matvec_16", "ll_matvec_32" }; + +fn wanted(only: ?[]const u8, name: []const u8) bool { + const text = only orelse return true; + return std.mem.indexOf(u8, name, text) != null; +} + +/// The probes made for a run, freed together. +pub const Set = struct { + glue: std.ArrayList(*llama.Glue) = .empty, + matvec: std.ArrayList(*dense.Matvec) = .empty, + gemm: std.ArrayList(*dense.Gemm) = .empty, + argmax: ?*llama.Argmax = null, + decode: ?*attention.DecodeVec = null, + merge: ?*attention.MergeVec = null, + prompt: std.ArrayList(*attention.Prompt) = .empty, + + pub fn deinit(s: *Set, a: std.mem.Allocator) void { + for (s.glue.items) |p| { + p.deinit(); + a.destroy(p); + } + for (s.matvec.items) |p| { + p.deinit(); + a.destroy(p); + } + for (s.gemm.items) |p| { + p.deinit(); + a.destroy(p); + } + for (s.prompt.items) |p| { + p.deinit(); + a.destroy(p); + } + if (s.argmax) |p| { + p.deinit(); + a.destroy(p); + } + if (s.decode) |p| { + p.deinit(); + a.destroy(p); + } + if (s.merge) |p| { + p.deinit(); + a.destroy(p); + } + s.glue.deinit(a); + s.matvec.deinit(a); + s.gemm.deinit(a); + s.prompt.deinit(a); + } + + /// Every registration whose name holds `only`; the roll call checks the registrations cover every pipeline. + pub fn collect(s: *Set, a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, tensor: bool, kernels: *std.ArrayList(g.Kernel), only: ?[]const u8) !void { + var functions: [family.pipelines + 1][]const u8 = undefined; + var n: usize = 0; + for (std.enums.values(llama.Which)) |w| { + functions[n] = llama.function(w); + n += 1; + if (!wanted(only, llama.spec(w).name)) continue; + const p = try a.create(llama.Glue); + errdefer a.destroy(p); + p.* = try llama.Glue.init(a, device, queue, w); + try s.glue.append(a, p); + try kernels.append(a, p.kernel()); + } + functions[n] = "ll_argmax"; + n += 1; + if (wanted(only, "llama glue ll_argmax")) { + const p = try a.create(llama.Argmax); + errdefer a.destroy(p); + p.* = try llama.Argmax.init(a, device, queue); + s.argmax = p; + try kernels.append(a, p.kernel()); + } + for ([_][]const u8{ "ll_decode_vec", "ll_merge_vec" }) |f| { + functions[n] = f; + n += 1; + } + if (wanted(only, "llama glue ll_decode_vec")) { + const p = try a.create(attention.DecodeVec); + errdefer a.destroy(p); + p.* = try attention.DecodeVec.init(a, device, queue); + s.decode = p; + try kernels.append(a, p.kernel()); + } + if (wanted(only, "llama glue ll_merge_vec")) { + const p = try a.create(attention.MergeVec); + errdefer a.destroy(p); + p.* = try attention.MergeVec.init(a, device, queue); + s.merge = p; + try kernels.append(a, p.kernel()); + } + for (matvec_rows, matvec_names) |rows, fun| { + functions[n] = fun; + n += 1; + var name_buf: [40]u8 = undefined; + if (!wanted(only, try std.fmt.bufPrint(&name_buf, "llama glue {s}", .{fun}))) continue; + const p = try a.create(dense.Matvec); + errdefer a.destroy(p); + p.* = try dense.Matvec.init(a, device, queue, rows); + try s.matvec.append(a, p); + try kernels.append(a, p.kernel()); + } + for ([_][2]bool{ .{ false, false }, .{ false, true }, .{ true, false }, .{ true, true } }) |form| { + const fun = family.gemm_names[@as(usize, @intFromBool(form[0])) * 2 + @intFromBool(form[1])]; + functions[n] = fun; + n += 1; + var name_buf: [40]u8 = undefined; + if (!tensor or !wanted(only, try std.fmt.bufPrint(&name_buf, "llama gemm {s}", .{fun}))) continue; + const p = try a.create(dense.Gemm); + errdefer a.destroy(p); + p.* = try dense.Gemm.init(a, device, queue, form[0], form[1]); + try s.gemm.append(a, p); + try kernels.append(a, p.kernel()); + } + functions[n] = family.attention_name; + n += 1; + for (family.geometries) |geo| { + var name_buf: [64]u8 = undefined; + if (!tensor or !wanted(only, try std.fmt.bufPrint(&name_buf, "llama {s} attention ll_prompt_attention", .{geo.name}))) continue; + const p = try a.create(attention.Prompt); + errdefer a.destroy(p); + p.* = try attention.Prompt.init(a, device, queue, geo); + try s.prompt.append(a, p); + try kernels.append(a, p.kernel()); + } + try family.requireComplete(functions[0..n]); + } +}; + +/// Gate 5 for the family's pipelines as built here: the threads each launch takes and its static threadgroup bytes. +pub fn limits(a: std.mem.Allocator, device: mtl.Device, tensor: bool, only: ?[]const u8) !usize { + var problems: usize = 0; + const glue = try family.glueSource(a); + defer a.free(glue); + const Decl = struct { name: []const u8, threads: u32, tg: u32 }; + const decls = [_]Decl{ + .{ .name = "llama glue ll_embed", .threads = 256, .tg = 0 }, + .{ .name = "llama glue ll_norm", .threads = 256, .tg = 128 }, + .{ .name = "llama glue ll_add", .threads = 256, .tg = 0 }, + .{ .name = "llama glue ll_add_norm", .threads = 256, .tg = 128 }, + .{ .name = "llama glue ll_prep", .threads = 32, .tg = 0 }, + .{ .name = "llama glue ll_silu_gu", .threads = 256, .tg = 0 }, + .{ .name = "llama glue ll_decode_vec", .threads = 32 * (family.nanbeige43.heads / family.nanbeige43.kv_heads), .tg = 0 }, + .{ .name = "llama glue ll_merge_vec", .threads = 1024, .tg = 32 * 32 * 4 }, + .{ .name = "llama glue ll_argmax", .threads = 1024, .tg = 32 * 8 }, + .{ .name = "llama glue ll_matvec_1", .threads = 256, .tg = 0 }, + .{ .name = "llama glue ll_matvec_2", .threads = 256, .tg = 0 }, + .{ .name = "llama glue ll_matvec_4", .threads = 256, .tg = 0 }, + .{ .name = "llama glue ll_matvec_8", .threads = 256, .tg = 0 }, + .{ .name = "llama glue ll_matvec_16", .threads = 256, .tg = 0 }, + .{ .name = "llama glue ll_matvec_32", .threads = 256, .tg = 0 }, + }; + for (decls) |d| { + if (!wanted(only, d.name)) continue; + const fun = d.name[std.mem.lastIndexOfScalar(u8, d.name, ' ').? + 1 ..]; + problems += try limits_k.onThisGpu(a, device, .{ .decl = .{ .name = d.name, .threads = d.threads, .tg_memory = d.tg, .offline = false }, .function = fun, .here = glue, .pre_m5 = glue }, false); + } + if (!tensor) return problems; + const gemm = try family.gemmSource(device, a); + defer a.free(gemm); + for (family.gemm_names) |fun| { + const name = try std.fmt.allocPrint(a, "llama gemm {s}", .{fun}); + defer a.free(name); + if (!wanted(only, name)) continue; + problems += try limits_k.onThisGpu(a, device, .{ .decl = .{ .name = name, .threads = 256, .tg_memory = 0, .offline = false }, .function = fun, .here = gemm, .pre_m5 = gemm }, false); + } + for (family.geometries) |geo| { + const name = try std.fmt.allocPrint(a, "llama {s} attention ll_prompt_attention", .{geo.name}); + defer a.free(name); + if (!wanted(only, name)) continue; + const src = try family.attentionSource(device, a, geo); + defer a.free(src); + // pt[32 x 136] floats and rowmax[2][32] + problems += try limits_k.onThisGpu(a, device, .{ .decl = .{ .name = name, .threads = 64, .tg_memory = 32 * 136 * 4 + 2 * 32 * 4, .offline = false }, .function = family.attention_name, .here = src, .pre_m5 = src }, false); + } + return problems; +} diff --git a/zig/tests/contract/metal_llama_attention.zig b/zig/tests/contract/metal_llama_attention.zig new file mode 100644 index 0000000000..b0905e9636 --- /dev/null +++ b/zig/tests/contract/metal_llama_attention.zig @@ -0,0 +1,511 @@ +//! Llama's attention under the contract: decode partials and their merge (glue), and the tensor-unit prompt kernel. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const llama = @import("metal_llama.zig"); +const c = contract.spec; +const g = contract.gates; +const rows_in = contract.rows; +const family = tf.llama.contracts; +const DH = llama.DH; +const bf = llama.bf; + +/// glue.metal's ll_vec_args, as forward.zig fills it. +const VArgs = extern struct { rows: u32, pos: u32, heads: u32, groups: u32, blocks: u32, chunk: u32, first: u32, scale: f32 }; + +const softmax_sums: c.Sums = .{ .f16 = 0, .bf16 = 0, .exact = false, .why = "weights of exp: fp32 sums held to the float64 bound, not integers" }; + +fn geometryOf(heads: usize) ?family.Geometry { + for (family.geometries) |geo| if (geo.heads == heads) return geo; + return null; +} + +fn random16(rng: std.Random, out: []u16, lo: f32, hi: f32) void { + for (out) |*v| v.* = c.toBf16(lo + rng.float(f32) * (hi - lo)); +} + +fn xv(v: []const u16, i: usize) f64 { + return c.fromBf16(v[i]); +} + +/// A block's keys in sequence, four at a time, then the merge of the four lanes. +fn decodeDepth(s: c.Shape) usize { + return @min(s.k, 2048) + 16; +} + +/// ll_decode_vec: a (key head, row, key block) a threadgroup; its partial maximum, mass and output over the block. +pub const DecodeVec = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: mtl.Library, + pipe: mtl.Pipeline, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + geo: family.Geometry = family.nanbeige43, + chunk: usize = 32, + blocks: usize = 0, + q: []u16 = &.{}, + k: []u16 = &.{}, + v: []u16 = &.{}, + in: [3]rig_.Ref = undefined, + out: [3]rig_.Ref = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue) !DecodeVec { + const lib = try family.glueLibrary(device, a); + errdefer lib.deinit(); + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .lib = lib, .pipe = try llama.pipelineOf(lib, device, "ll_decode_vec") }; + } + + fn drop(p: *DecodeVec) void { + for ([_]*[]u16{ &p.q, &p.k, &p.v }) |s| { + p.a.free(s.*); + s.* = &.{}; + } + p.rig.reset(); + } + + pub fn deinit(p: *DecodeVec) void { + p.drop(); + p.rig.deinit(); + p.pipe.deinit(); + p.lib.deinit(); + } + + const vt = g.hooks(DecodeVec); + + pub fn kernel(p: *DecodeVec) g.Kernel { + return .{ .spec = .{ .name = "llama glue ll_decode_vec", .accum = .f32, .depth = decodeDepth, .ops = 24, .out = .f32, .tiles = .{ .{ .size = 16, .step = 1, .min = 1 }, .{}, .{ .size = 32, .step = 1, .min = 1 } }, .served = &.{ .{ .rows = 16, .n = family.nanbeige43.heads, .k = 4096 }, .{ .rows = 4, .n = family.nanbeige41.heads, .k = 1024 } }, .sums = softmax_sums }, .ctx = p, .vt = &vt }; + } + + /// rows query rows at the end of k positions, n query heads; the blocks as the served cache of k would hold. + pub fn prepare(p: *DecodeVec, s: c.Shape, fill: g.Fill, seed: u64) !bool { + const geo = geometryOf(s.n) orelse return false; + if (s.rows == 0 or s.rows > 16 or s.k < s.rows) return false; + p.drop(); + p.shape = s; + p.geo = geo; + p.chunk = tf.llama.state.chunk(s.k); + p.blocks = (s.k - 1) / p.chunk + 1; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const kvw = geo.kv_heads * DH; + p.q = try p.a.alloc(u16, s.rows * geo.heads * DH); + p.k = try p.a.alloc(u16, s.k * kvw); + p.v = try p.a.alloc(u16, s.k * kvw); + random16(rng, p.q, -2, 2); + random16(rng, p.k, -2, 2); + random16(rng, p.v, -2, 2); + if (fill == .nonfinite) { + p.k[5 * kvw + 3] = c.toBf16(std.math.nan(f32)); // the first key head's queries all see key 5 + p.v[@min(9, s.k - 1) * kvw + DH + 1] = c.toBf16(std.math.nan(f32)); // the second head's outputs only + } + p.in[0] = try p.rig.input(std.mem.sliceAsBytes(p.q)); + p.in[1] = try p.rig.input(std.mem.sliceAsBytes(p.k)); + p.in[2] = try p.rig.input(std.mem.sliceAsBytes(p.v)); + const slots = geo.heads * s.rows * p.blocks; + p.out[0] = try p.rig.output(slots * 4); + p.out[1] = try p.rig.output(slots * 4); + p.out[2] = try p.rig.output(slots * DH * 4); + return true; + } + + pub fn count(p: *DecodeVec) usize { + return p.geo.heads * p.shape.rows * p.blocks * (2 + DH); + } + + pub fn run(p: *DecodeVec) !void { + const s = p.shape; + const geo = p.geo; + const pass = p.rig.begin(); + const e = pass.e; + e.setPipeline(p.pipe); + for (p.in, 0..) |r, slot| e.setBuffer(r.buf, r.off, slot); + e.setValue(VArgs{ .rows = @intCast(s.rows), .pos = @intCast(s.k - s.rows), .heads = geo.heads, .groups = geo.kv_heads, .blocks = @intCast(p.blocks), .chunk = @intCast(p.chunk), .first = 0, .scale = tf.llama.config.headScale() }, 3); + for (p.out, 4..) |r, slot| e.setBuffer(r.buf, r.off, slot); + e.dispatchGroups(mtl.Size.of(geo.kv_heads, s.rows, p.blocks), mtl.Size.of(32, geo.heads / geo.kv_heads, 1)); + try p.rig.finish(pass); + } + + /// A slot's keys: from the block's first to the row's own position, within the block. + fn span(p: *const DecodeVec, row: usize, block: usize) [2]usize { + const pos = p.shape.k - p.shape.rows + row; + return .{ block * p.chunk, @min(pos + 1, (block + 1) * p.chunk) }; + } + + fn score(p: *const DecodeVec, head: usize, row: usize, key: usize) f64 { + return p.dot(head, row, key)[0]; + } + + /// The scaled dot and the magnitude of its terms. + fn dot(p: *const DecodeVec, head: usize, row: usize, key: usize) [2]f64 { + const geo = p.geo; + const kv = head / (geo.heads / geo.kv_heads); + var sum: f64 = 0; + var abs: f64 = 0; + for (0..DH) |d| { + const t = @as(f64, tf.llama.config.headScale()) * xv(p.q, (row * geo.heads + head) * DH + d) * xv(p.k, key * geo.kv_heads * DH + kv * DH + d); + sum += t; + abs += @abs(t); + } + return .{ sum, abs }; + } + + /// PM (the block's top score), PL (its mass) and PO (its weighted values), each in its own output. + pub fn reference(p: *DecodeVec, want: []f64, mag: []f64) !void { + const s = p.shape; + const geo = p.geo; + const slots = geo.heads * s.rows * p.blocks; + const kvw = geo.kv_heads * DH; + for (0..geo.heads) |head| for (0..s.rows) |row| for (0..p.blocks) |block| { + const slot = (head * s.rows + row) * p.blocks + block; + const kv = head / (geo.heads / geo.kv_heads); + const from, const to = p.span(row, block); + var top: f64 = -std.math.inf(f64); // Metal's max drops a NaN score; exp(NaN - top) carries it on + var top_abs: f64 = 0; // the top score's terms, the magnitude its fp32 dot rounds against + var key = from; + while (key < to) : (key += 1) { + const sc, const abs = p.dot(head, row, key); + if (!std.math.isNan(sc) and sc > top) { + top = sc; + top_abs = abs; + } + } + var mass: f64 = 0; + var mass_abs: f64 = 0; + const o = want[2 * slots + slot * DH ..][0..DH]; + const om = mag[2 * slots + slot * DH ..][0..DH]; + @memset(o, 0); + @memset(om, 0); + key = from; + while (key < to) : (key += 1) { + const w = @exp(p.score(head, row, key) - top); + mass += w; + mass_abs += @abs(w); + for (0..DH) |d| { + const t = w * xv(p.v, key * kvw + kv * DH + d); + o[d] += t; + om[d] += @abs(t); + } + } + want[slot] = top; + mag[slot] = top_abs; + want[slots + slot] = mass; + mag[slots + slot] = mass_abs; + }; + } + + pub fn simulate(p: *DecodeVec, acc: c.Accum, order: g.Order, out: []f64) !void { + const s = p.shape; + const geo = p.geo; + const slots = geo.heads * s.rows * p.blocks; + const kvw = geo.kv_heads * DH; + const t = try p.a.alloc(f64, p.chunk); + defer p.a.free(t); + for (0..geo.heads) |head| for (0..s.rows) |row| for (0..p.blocks) |block| { + const slot = (head * s.rows + row) * p.blocks + block; + const kv = head / (geo.heads / geo.kv_heads); + const from, const to = p.span(row, block); + var top: f64 = -std.math.inf(f64); + for (from..to) |key| { + const sc = p.score(head, row, key); + if (!std.math.isNan(sc)) top = @max(top, sc); + } + for (from..to, 0..) |key, i| t[i] = @exp(p.score(head, row, key) - top); + out[slot] = top; + out[slots + slot] = rows_in.sumIn(acc, order, t[0 .. to - from]); + for (0..DH) |d| { + for (from..to, 0..) |key, i| t[i] = @exp(p.score(head, row, key) - top) * xv(p.v, key * kvw + kv * DH + d); + out[2 * slots + slot * DH + d] = rows_in.sumIn(acc, order, t[0 .. to - from]); + } + }; + } +}; + +/// The simdgroups' strided sums over the blocks, the fold through threadgroup memory and the division. +fn mergeDepth(s: c.Shape) usize { + return s.k / 32 + 16; +} + +/// ll_merge_vec: the partials of every block folded into a row's bf16 attention output. +pub const MergeVec = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: mtl.Library, + pipe: mtl.Pipeline, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + pm: []f32 = &.{}, + pl: []f32 = &.{}, + po: []f32 = &.{}, + in: [3]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue) !MergeVec { + const lib = try family.glueLibrary(device, a); + errdefer lib.deinit(); + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .lib = lib, .pipe = try llama.pipelineOf(lib, device, "ll_merge_vec") }; + } + + fn drop(p: *MergeVec) void { + for ([_]*[]f32{ &p.pm, &p.pl, &p.po }) |s| { + p.a.free(s.*); + s.* = &.{}; + } + p.rig.reset(); + } + + pub fn deinit(p: *MergeVec) void { + p.drop(); + p.rig.deinit(); + p.pipe.deinit(); + p.lib.deinit(); + } + + const vt = g.hooks(MergeVec); + + pub fn kernel(p: *MergeVec) g.Kernel { + return .{ .spec = .{ .name = "llama glue ll_merge_vec", .accum = .f32, .depth = mergeDepth, .ops = 6, .out = .bf16, .tiles = .{ .{ .size = 16, .step = 1, .min = 1 }, .{}, .{ .size = 32, .step = 1, .min = 1 } }, .served = &.{ .{ .rows = 16, .n = family.nanbeige43.heads, .k = 128 }, .{ .rows = 4, .n = family.nanbeige41.heads, .k = 64 } }, .sums = softmax_sums }, .ctx = p, .vt = &vt }; + } + + /// rows rows of n heads over k blocks of partials, as decode leaves them: every fourth block past the first empty. + pub fn prepare(p: *MergeVec, s: c.Shape, fill: g.Fill, seed: u64) !bool { + if (s.rows == 0 or s.rows > 16 or s.n == 0 or s.k == 0) return false; + p.drop(); + p.shape = s; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const slots = s.n * s.rows * s.k; + p.pm = try p.a.alloc(f32, slots); + p.pl = try p.a.alloc(f32, slots); + p.po = try p.a.alloc(f32, slots * DH); + for (p.pm, p.pl, 0..) |*m, *l, i| { + const empty = i % s.k != 0 and i % 4 == 3; + m.* = if (empty) -std.math.inf(f32) else rng.float(f32) * 20 - 10; + l.* = if (empty) 0 else 0.5 + rng.float(f32) * 64; + } + for (p.po, 0..) |*o, i| o.* = (rng.float(f32) * 2 - 1) * p.pl[i / DH]; + if (fill == .nonfinite) { + p.po[3 * DH + 7] = std.math.nan(f32); // the first head's first row + if (s.k > 1) p.pm[s.k + 1] = std.math.inf(f32); // the second slot's top: exp(inf - inf) + } + p.in[0] = try p.rig.input(std.mem.sliceAsBytes(p.pm)); + p.in[1] = try p.rig.input(std.mem.sliceAsBytes(p.pl)); + p.in[2] = try p.rig.input(std.mem.sliceAsBytes(p.po)); + p.out = try p.rig.output(s.rows * s.n * DH * 2); + return true; + } + + pub fn count(p: *MergeVec) usize { + return p.shape.rows * p.shape.n * DH; + } + + pub fn run(p: *MergeVec) !void { + const s = p.shape; + const pass = p.rig.begin(); + const e = pass.e; + e.setPipeline(p.pipe); + for (p.in, 0..) |r, slot| e.setBuffer(r.buf, r.off, slot); + e.setValue(VArgs{ .rows = @intCast(s.rows), .pos = 0, .heads = @intCast(s.n), .groups = 1, .blocks = @intCast(s.k), .chunk = 32, .first = 0, .scale = 1 }, 3); + e.setBuffer(p.out.buf, p.out.off, 4); + e.dispatchGroups(mtl.Size.of(s.n, s.rows, 1), mtl.Size.of(1024, 1, 1)); + try p.rig.finish(pass); + } + + fn fold(p: *MergeVec, acc: ?struct { a: c.Accum, o: g.Order }, want: []f64, mag: ?[]f64) !void { + const s = p.shape; + const t = try p.a.alloc(f64, s.k); + defer p.a.free(t); + for (0..s.n) |head| for (0..s.rows) |row| { + const at = head * s.rows + row; + var top: f64 = -std.math.inf(f64); // Metal's max drops a NaN partial + for (0..s.k) |b| if (!std.math.isNan(p.pm[at * s.k + b])) { + top = @max(top, p.pm[at * s.k + b]); + }; + var total: f64 = 0; + var total_abs: f64 = 0; + for (0..s.k) |b| { + const l: f64 = p.pl[at * s.k + b]; + t[b] = if (l > 0) l * @exp(@as(f64, p.pm[at * s.k + b]) - top) else 0; + total += t[b]; + total_abs += @abs(t[b]); + } + if (acc) |x| total = rows_in.sumIn(x.a, x.o, t); + for (0..DH) |d| { + var sum: f64 = 0; + var abs: f64 = 0; + for (0..s.k) |b| { + const l: f64 = p.pl[at * s.k + b]; + t[b] = if (l > 0) @exp(@as(f64, p.pm[at * s.k + b]) - top) * p.po[(at * s.k + b) * DH + d] else 0; + sum += t[b]; + abs += @abs(t[b]); + } + if (acc) |x| sum = rows_in.sumIn(x.a, x.o, t); + const o = (row * s.n + head) * DH + d; + want[o] = if (acc == null) sum / total else bf(sum / total); + if (mag) |m| m[o] = abs / total + @abs(sum) * total_abs / (total * total); + } + }; + } + + pub fn reference(p: *MergeVec, want: []f64, mag: []f64) !void { + return p.fold(null, want, mag); + } + + pub fn simulate(p: *MergeVec, acc: c.Accum, order: g.Order, out: []f64) !void { + return p.fold(.{ .a = acc, .o = order }, out, null); + } +}; + +/// The tensor-op accumulation over the keys, the rescales and the final division. +fn promptDepth(s: c.Shape) usize { + return s.k / 16 + 32; +} + +/// ll_prompt_attention for one geometry: rows query rows at positions k - rows.. over a cache of k positions. +pub const Prompt = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + geo: family.Geometry, + name: []const u8, + served: [2]c.Shape, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + q: []u16 = &.{}, + k: []u16 = &.{}, + v: []u16 = &.{}, + in: [3]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, geo: family.Geometry) !Prompt { + const src = try family.attentionSource(device, a, geo); + defer a.free(src); + const name = try std.fmt.allocPrint(a, "llama {s} attention ll_prompt_attention", .{geo.name}); + errdefer a.free(name); + const big = geo.heads == family.nanbeige43.heads; + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .pipe = try rig_.pipeline(device, src, family.attention_name), .geo = geo, .name = name, .served = .{ .{ .rows = if (big) 256 else 64, .n = geo.heads, .k = if (big) 1280 else 64 }, .{ .rows = 33, .n = geo.heads, .k = 33 } } }; + } + + fn drop(p: *Prompt) void { + for ([_]*[]u16{ &p.q, &p.k, &p.v }) |s| { + p.a.free(s.*); + s.* = &.{}; + } + p.rig.reset(); + } + + pub fn deinit(p: *Prompt) void { + p.drop(); + p.rig.deinit(); + p.pipe.deinit(); + p.a.free(p.name); + } + + const vt = g.hooks(Prompt); + + pub fn kernel(p: *Prompt) g.Kernel { + return .{ .spec = .{ .name = p.name, .accum = .f32, .depth = promptDepth, .ops = 8, .out = .bf16, .tiles = .{ .{ .size = 32, .step = 1, .min = 1 }, .{}, .{ .size = 64, .step = 1, .min = 1 } }, .served = &p.served, .sums = softmax_sums }, .ctx = p, .vt = &vt }; + } + + /// Q padded to whole 32-row tiles and the cache to whole 64-key blocks, as the served buffers hold them. + pub fn prepare(p: *Prompt, s: c.Shape, fill: g.Fill, seed: u64) !bool { + if (s.n != p.geo.heads or s.rows == 0 or s.rows > 2048 or s.k < s.rows) return false; + p.drop(); + p.shape = s; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const geo = p.geo; + const kvw = geo.kv_heads * DH; + const slots = std.mem.alignForward(usize, s.k, 64); + p.q = try p.a.alloc(u16, std.mem.alignForward(usize, s.rows, 32) * geo.heads * DH); + p.k = try p.a.alloc(u16, slots * kvw); + p.v = try p.a.alloc(u16, slots * kvw); + random16(rng, p.q, -2, 2); + random16(rng, p.k, -2, 2); + random16(rng, p.v, -2, 2); + if (fill == .nonfinite) { + p.k[3 * kvw + 2] = c.toBf16(std.math.nan(f32)); // key 3: every row of the first key head's queries + if (geo.kv_heads > 1) p.v[@min(10, s.k - 1) * kvw + DH + 4] = c.toBf16(std.math.nan(f32)); + } + p.in[0] = try p.rig.input(std.mem.sliceAsBytes(p.q)); + p.in[1] = try p.rig.input(std.mem.sliceAsBytes(p.k)); + p.in[2] = try p.rig.input(std.mem.sliceAsBytes(p.v)); + p.out = try p.rig.output(s.rows * geo.heads * DH * 2); + return true; + } + + pub fn count(p: *Prompt) usize { + return p.shape.rows * p.geo.heads * DH; + } + + pub fn run(p: *Prompt) !void { + const s = p.shape; + const pass = p.rig.begin(); + const e = pass.e; + e.setPipeline(p.pipe); + for (p.in, 0..) |r, slot| e.setBuffer(r.buf, r.off, slot); + e.setBuffer(p.out.buf, p.out.off, 3); + e.setValue([2]i32{ @intCast(s.rows), @intCast(s.k - s.rows) }, 4); + e.dispatchGroups(mtl.Size.of((s.rows + 31) / 32, p.geo.heads, 1), mtl.Size.of(DH / 64 * 32, 1, 1)); + try p.rig.finish(pass); + } + + fn score(p: *const Prompt, head: usize, row: usize, key: usize) f64 { + const geo = p.geo; + const kv = head / (geo.heads / geo.kv_heads); + var dot: f64 = 0; + for (0..DH) |d| dot += xv(p.q, (row * geo.heads + head) * DH + d) * xv(p.k, key * geo.kv_heads * DH + kv * DH + d); + return dot * @as(f64, tf.llama.config.headScale()); + } + + fn fold(p: *Prompt, acc: ?struct { a: c.Accum, o: g.Order }, want: []f64, mag: ?[]f64) !void { + const s = p.shape; + const geo = p.geo; + const kvw = geo.kv_heads * DH; + const p0 = s.k - s.rows; + const t = try p.a.alloc(f64, s.k); + defer p.a.free(t); + for (0..s.rows) |row| for (0..geo.heads) |head| { + const kv = head / (geo.heads / geo.kv_heads); + const keys = p0 + row + 1; + var top: f64 = -std.math.inf(f64); // the tile's max drops a NaN score; exp2(NaN - top) carries it on + for (0..keys) |key| { + const sc = p.score(head, row, key); + if (!std.math.isNan(sc)) top = @max(top, sc); + } + var total: f64 = 0; + var total_abs: f64 = 0; + for (0..keys) |key| { + t[key] = @exp(p.score(head, row, key) - top); + total += t[key]; + total_abs += @abs(t[key]); + } + if (acc) |x| total = rows_in.sumIn(x.a, x.o, t[0..keys]); + for (0..DH) |d| { + var sum: f64 = 0; + var abs: f64 = 0; + for (0..keys) |key| { + t[key] = @exp(p.score(head, row, key) - top) * xv(p.v, key * kvw + kv * DH + d); + sum += t[key]; + abs += @abs(t[key]); + } + if (acc) |x| sum = rows_in.sumIn(x.a, x.o, t[0..keys]); + const o = (row * geo.heads + head) * DH + d; + want[o] = if (acc == null) sum / total else bf(sum / total); + if (mag) |m| m[o] = abs / total + @abs(sum) * total_abs / (total * total); + } + }; + } + + pub fn reference(p: *Prompt, want: []f64, mag: []f64) !void { + return p.fold(null, want, mag); + } + + /// The tensor op takes the fp32 probabilities at relaxed precision: a bf16 ulp of each weighted term more. + pub fn bounds(_: *Prompt, want: []const f64, mag: []const f64, allow: []f64) !void { + for (allow, want, mag) |*l, w, m| l.* += 0x1p-7 * m + 0x1p-8 * @abs(w); + } + + pub fn simulate(p: *Prompt, acc: c.Accum, order: g.Order, out: []f64) !void { + return p.fold(.{ .a = acc, .o = order }, out, null); + } +}; diff --git a/zig/tests/contract/metal_llama_dense.zig b/zig/tests/contract/metal_llama_dense.zig new file mode 100644 index 0000000000..6b1973917c --- /dev/null +++ b/zig/tests/contract/metal_llama_dense.zig @@ -0,0 +1,230 @@ +//! Llama's bf16 projections under the contract: the row kernels of a window and the tensor-unit GEMM of a prompt. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const llama = @import("metal_llama.zig"); +const c = contract.spec; +const g = contract.gates; +const dense_in = contract.dense; +const family = llama.family; +const G33 = family.nanbeige43; +const G3 = family.nanbeige41; + +/// glue.metal's ll_mv_args, as forward.zig fills it. +const MArgs = extern struct { rows: u32, n: u32, k: u32, ldy: u32 }; + +/// A lane's dots over its 8-wide slices of the row, then the simd tree. +fn matvecDepth(s: c.Shape) usize { + return s.k / 32 + 8; +} + +/// The row projection ll_matvec_ROWS: Y[r, j] = bf16(sum X[r, i] W[j, i]) in fp32, ROWS rows a launch. +pub const Matvec = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: mtl.Library, + pipe: mtl.Pipeline, + per: u32, + name: []const u8, + served: [2]c.Shape, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + d: ?dense_in.Dense = null, + w16: []u16 = &.{}, + x: rig_.Ref = undefined, + w: rig_.Ref = undefined, + y: rig_.Ref = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, per: u32) !Matvec { + const lib = try family.glueLibrary(device, a); + errdefer lib.deinit(); + const name = try std.fmt.allocPrint(a, "llama glue ll_matvec_{d}", .{per}); + errdefer a.free(name); + const fun = name[std.mem.lastIndexOfScalar(u8, name, ' ').? + 1 ..]; + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .lib = lib, .pipe = try llama.pipelineOf(lib, device, fun), .per = per, .name = name, .served = .{ .{ .rows = per, .n = 1024, .k = G33.intermediate }, .{ .rows = per, .n = G33.hidden, .k = G3.hidden } } }; + } + + fn drop(p: *Matvec) void { + if (p.d) |d| d.deinit(); + p.d = null; + p.a.free(p.w16); + p.w16 = &.{}; + p.rig.reset(); + } + + pub fn deinit(p: *Matvec) void { + p.drop(); + p.rig.deinit(); + p.pipe.deinit(); + p.lib.deinit(); + p.a.free(p.name); + } + + const vt = g.hooks(Matvec); + + pub fn kernel(p: *Matvec) g.Kernel { + return .{ .spec = .{ .name = p.name, .accum = .f32, .depth = matvecDepth, .out = .bf16, .tiles = .{ .{ .size = p.per, .step = 1, .min = 1 }, .{ .size = 8, .step = 1, .min = 1 }, .{ .size = 256, .step = 8, .min = 8 } }, .served = &p.served }, .ctx = p, .vt = &vt }; + } + + pub fn prepare(p: *Matvec, s: c.Shape, fill: g.Fill, seed: u64) !bool { + if (s.rows == 0 or s.n == 0 or s.k == 0 or s.k % 8 != 0) return false; + p.drop(); + p.shape = s; + const d = try dense_in.Dense.init(p.a, s, 1); + p.d = d; + d.fillX(fill, seed, 2); + p.w16 = try p.a.alloc(u16, s.n * s.k); + try denseWeights(p.w16, d, s, fill, seed); + p.x = try p.rig.input(std.mem.sliceAsBytes(d.x)); + p.w = try p.rig.input(std.mem.sliceAsBytes(p.w16)); + p.y = try p.rig.output(d.outputs() * 2); + return true; + } + + pub fn count(p: *Matvec) usize { + return p.d.?.outputs(); + } + + pub fn run(p: *Matvec) !void { + const s = p.shape; + const pass = p.rig.begin(); + const e = pass.e; + e.setPipeline(p.pipe); + var at: usize = 0; + while (at < s.rows) : (at += p.per) { + const here: u32 = @intCast(@min(p.per, s.rows - at)); + e.setBuffer(p.x.buf, p.x.off + at * s.k * 2, 0); + e.setBuffer(p.w.buf, p.w.off, 1); + e.setValue(MArgs{ .rows = here, .n = @intCast(s.n), .k = @intCast(s.k), .ldy = @intCast(s.n) }, 2); + e.setBuffer(p.y.buf, p.y.off + at * s.n * 2, 3); + e.dispatchGroups(mtl.Size.of((s.n + 7) / 8, 1, 1), mtl.Size.of(256, 1, 1)); + } + try p.rig.finish(pass); + } + + pub fn reference(p: *Matvec, want: []f64, mag: []f64) !void { + return p.d.?.reference(want, mag); + } + + pub fn simulate(p: *Matvec, acc: c.Accum, order: g.Order, out: []f64) !void { + return p.d.?.simulate(acc, order, out); + } +}; + +/// Bf16 weights for a dense family: random in [-1, 1), or gate 4's integers, mirrored and negated over the halves. +pub fn denseWeights(w16: []u16, d: dense_in.Dense, s: c.Shape, fill: g.Fill, seed: u64) !void { + var prng = std.Random.DefaultPrng.init(seed ^ 0x5151); + const rng = prng.random(); + const h = s.k / 2; + for (0..s.n) |j| for (0..s.k) |i| { + const v: f32 = if (fill != .stress) rng.float(f32) * 2 - 1 else if (i < h) @floatFromInt(1 + (i * 7 + j) % 6) else -@as(f32, @floatFromInt(1 + ((i - h) * 7 + j) % 6)); + w16[j * s.k + i] = c.toBf16(v); + d.w[j * s.k + i] = c.fromBf16(w16[j * s.k + i]); + }; +} + +/// The tensor-op accumulation over the K tiles, then the output rounding. +fn gemmDepth(s: c.Shape) usize { + return s.k / 16 + 16; +} + +/// The prompt GEMM ll_gemm_AB: A whole 64-row tiles of M, B whole 128-column tiles of N; a probe pads to its form. +pub const Gemm = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + am: bool, + an: bool, + name: []const u8, + served: [2]c.Shape, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + d: ?dense_in.Dense = null, + w16: []u16 = &.{}, + x: rig_.Ref = undefined, + w: rig_.Ref = undefined, + y: rig_.Ref = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, am: bool, an: bool) !Gemm { + const src = try family.gemmSource(device, a); + defer a.free(src); + const name = try std.fmt.allocPrint(a, "llama gemm ll_gemm_{d}{d}", .{ @intFromBool(am), @intFromBool(an) }); + errdefer a.free(name); + const fun = name[std.mem.lastIndexOfScalar(u8, name, ' ').? + 1 ..]; + const rows: usize = if (am) 256 else 100; + const n: usize = if (an) 1024 else 1088; + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .pipe = try rig_.pipeline(device, src, fun), .am = am, .an = an, .name = name, .served = .{ .{ .rows = rows, .n = n, .k = G33.hidden }, .{ .rows = if (am) 64 else 33, .n = if (an) G3.hidden else G3.hidden + 64, .k = G3.hidden } } }; + } + + fn drop(p: *Gemm) void { + if (p.d) |d| d.deinit(); + p.d = null; + p.a.free(p.w16); + p.w16 = &.{}; + p.rig.reset(); + } + + pub fn deinit(p: *Gemm) void { + p.drop(); + p.rig.deinit(); + p.pipe.deinit(); + p.a.free(p.name); + } + + const vt = g.hooks(Gemm); + + pub fn kernel(p: *Gemm) g.Kernel { + return .{ .spec = .{ .name = p.name, .accum = .f32, .depth = gemmDepth, .out = .bf16, .tiles = .{ .{ .size = 64, .step = 1, .min = 1 }, .{ .size = 128, .step = 1, .min = 1 }, .{ .size = 256, .step = 256, .min = 256 } }, .served = &p.served }, .ctx = p, .vt = &vt }; + } + + /// The shape the kernel runs at: whole tiles where its form assumes them (the padded rows are the probe's too). + fn padded(p: *const Gemm, s: c.Shape) c.Shape { + return .{ .rows = if (p.am) std.mem.alignForward(usize, s.rows, 64) else s.rows, .n = if (p.an) std.mem.alignForward(usize, s.n, 128) else s.n, .k = s.k }; + } + + pub fn prepare(p: *Gemm, s0: c.Shape, fill: g.Fill, seed: u64) !bool { + if (s0.rows == 0 or s0.n == 0 or s0.k == 0 or s0.k % 256 != 0) return false; + const s = p.padded(s0); + if (s.rows > 2048) return false; + p.drop(); + p.shape = s; + const d = try dense_in.Dense.init(p.a, s, 1); + p.d = d; + d.fillX(fill, seed, 2); + p.w16 = try p.a.alloc(u16, s.n * s.k); + try denseWeights(p.w16, d, s, fill, seed); + p.x = try p.rig.input(std.mem.sliceAsBytes(d.x)); + p.w = try p.rig.input(std.mem.sliceAsBytes(p.w16)); + p.y = try p.rig.output(d.outputs() * 2); + return true; + } + + pub fn count(p: *Gemm) usize { + return p.d.?.outputs(); + } + + pub fn run(p: *Gemm) !void { + const s = p.shape; + const pass = p.rig.begin(); + const e = pass.e; + e.setPipeline(p.pipe); + e.setBuffer(p.x.buf, p.x.off, 0); + e.setBuffer(p.w.buf, p.w.off, 1); + const tn = (s.n + 127) / 128; + const tm = (s.rows + 63) / 64; + var prm: [16]i32 = @splat(0); + const vals = [_]usize{ s.rows, s.n, s.k, s.k, s.k, s.n, tn, tm, 2, s.k / 256 }; + for (vals, 0..) |v, i| prm[i] = @intCast(v); + e.setBytes(std.mem.asBytes(&prm), 2); + e.setBuffer(p.y.buf, p.y.off, 3); + e.dispatchThreads(mtl.Size.of((tn << 2) * 32, (tm + 3) / 4 * 4, 2), mtl.Size.of(32, 4, 2)); + try p.rig.finish(pass); + } + + pub fn reference(p: *Gemm, want: []f64, mag: []f64) !void { + return p.d.?.reference(want, mag); + } + + pub fn simulate(p: *Gemm, acc: c.Accum, order: g.Order, out: []f64) !void { + return p.d.?.simulate(acc, order, out); + } +}; diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index 46801e171c..6df3de8ad0 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -27,6 +27,7 @@ const gemma_mixed = @import("metal_gemma4_mixed.zig"); const lfm2_argmax = @import("metal_lfm2_argmax.zig"); const lfm2_streams = @import("metal_lfm2_streams.zig"); const lfm2_controls = @import("metal_lfm2_controls.zig"); +const llama_k = @import("metal_llama_all.zig"); const g = contract.gates; const usage = "usage: tf-kernel-contract [--only TEXT] [--no-gates] [--no-limits] [--offline DIR] [--seed N] [--repro] [--rotary-long] [--rotary-accuracy FILE]"; @@ -298,6 +299,9 @@ pub fn main(init: std.process.Init) !u8 { }; var h3_list: std.ArrayList(*h3_k.H3) = .empty; defer for (h3_list.items) |x| x.deinit(); + var llama_set = llama_k.Set{}; + defer llama_set.deinit(a); + try llama_set.collect(a, device, queue, tensor, &kernels, o.only); if (tensor) for (std.enums.values(h3_k.Which)) |w| { if (!wanted(o, h3_k.spec(w).name)) continue; const x = try a.create(h3_k.H3); @@ -382,6 +386,7 @@ pub fn main(init: std.process.Init) !u8 { lim += try gemma.limits(a, device, o.only); lim += try gemma_embed.limits(a, device, o.only); lim += try gemma_mixed.limits(a, device, o.only); + lim += try llama_k.limits(a, device, tensor, o.only); } const t2 = std.Io.Clock.awake.now(init.io); const ms = struct { From e020a705facb3171132c6cfa894db4910360a4e4 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Tue, 27 Oct 2026 05:19:51 -0500 Subject: [PATCH 3/7] llama: admit only the checkpoint checked against HF, refuse Macs without tensor units by name; the teacher-forced misses and the licence in the doc tf-llama-run takes exact's window width as an argument and loses its environment switches (LLAMA_LAYERS, LLAMA_WINDOW). --- docs/recipes/nanbeige-4.3-33b-metal.md | 40 ++++++++++++++++--- zig/src/families/llama/config.zig | 54 +++++++++++++++++++++++++- zig/src/families/llama/model.zig | 2 + zig/src/native/llama.zig | 13 ++++++- zig/tests/llama/run.zig | 13 +++---- 5 files changed, 106 insertions(+), 16 deletions(-) diff --git a/docs/recipes/nanbeige-4.3-33b-metal.md b/docs/recipes/nanbeige-4.3-33b-metal.md index a30a2eb0c7..a2f4a1ff1d 100644 --- a/docs/recipes/nanbeige-4.3-33b-metal.md +++ b/docs/recipes/nanbeige-4.3-33b-metal.md @@ -3,13 +3,16 @@ This recipe covers text generation with [`Nanbeige/Nanbeige4.3-33B`](https://huggingface.co/Nanbeige/Nanbeige4.3-33B) (model_type `llama`: a dense 33B Llama-architecture model, 512K context, thinking and tool calls in its chat template) on the Metal engine, served from the bf16 release as published: no conversion, nothing lost. It was brought up on an -M5 Ultra (256 GB); other Macs are untested. The family serves any `llama` checkpoint with heads of 128, default RoPE, -no biases and its own output head, in bf16; it refuses others by name. +M5 Ultra (256 GB); other Macs are untested. The family admits only the checkpoints checked against Hugging Face, by the +geometry their configs name (`qualified` in `zig/src/families/llama/config.zig`: Nanbeige4.3-33B); another `llama` +checkpoint is refused by name until it has been checked the same way. Prompts run on Metal 4 tensor units, so it needs +an M5 or later: M1 to M4 Macs are refused by name at load. ## Build and serve ```sh -hf download Nanbeige/Nanbeige4.3-33B --local-dir models/nanbeige4.3-33b +hf download Nanbeige/Nanbeige4.3-33B --revision c1093c501cc6460427c0ac634db7223d0450c0ee \ + --local-dir models/nanbeige4.3-33b zig build native -Doptimize=ReleaseFast -Dcpu=native zig build grammar -Dxgrammar=PATH_TO_XGRAMMAR_v0.2.8 -Dcpu=native # structured output LLAMA_LOAD_LIMIT_GB=150 zig-out/native/bin/tensorfold-native serve models/nanbeige4.3-33b --name nanbeige --port 8888 \ @@ -59,6 +62,27 @@ greedy chats of 64 tokens (code, prose, Chinese, a document, with and without th token; the other three part at a position where the engine's two best logits are the same bf16 value. One-row steps and windows of 7 or 16 rows give the same logits bit for bit after prompts of 64 and 1,500 tokens. +The 13 teacher-forced positions of the 1,500 whose top tokens differ (row n predicts token n + 1; a tie goes to the +lower id on both sides). At each, HF's and the engine's two contenders are within one bf16 step of each other on both +sides (0.125 below 32, 0.25 from 32 to 64): HF's two tie at 3 of them, the engine's at 8, neither at 22 and 166, and +at 1370 a third token, `no`, ties at the top on both sides. + +| position | HF's top-1 (a) | engine's top-1 (b) | HF's logits a, b | engine's logits a, b | KL | +|---:|---|---|---|---|---:| +| 22 | `en` 271 | `see` 3574 | 20.75, 20.625 | 20.625, 20.75 | 0.0047 | +| 83 | `the` 269 | `Python` 9120 | 26.25, 26.25 | 26.25, 26.375 | 0.0022 | +| 85 | `developed` 3260 | `-` 152369 | 35.5, 35.5 | 35.5, 35.75 | 0.0134 | +| 166 | `";` 2651 | `"))` 9441 | 36.25, 36 | 36.25, 36.5 | 0.0234 | +| 224 | `(` 152374 | `('` 1091 | 39.75, 39.5 | 39.5, 39.5 | 0.0078 | +| 401 | `(` 152374 | `({'` 50468 | 39.5, 39.25 | 39.25, 39.25 | 0.0076 | +| 425 | `{` 152405 | `1` 152367 | 40, 39.75 | 40, 40 | 0.0163 | +| 586 | `representation` 7443 | `encoding` 13571 | 30, 30 | 29.875, 30 | 0.0027 | +| 1302 | `If` 3934 | ``` `` ``` 2283 | 32, 31.875 | 32, 32 | 0.0022 | +| 1304 | `indent` 28799 | `allow` 13471 | 37.25, 37 | 37, 37 | 0.0078 | +| 1370 | `it` 388 | `the` 269 | 38.25, 38 (`no` 754: 38.25) | 38, 38 (`no` 754: 38) | 0.0076 | +| 1384 | ``` `` ``` 16411 | `circular` 13380 | 30.875, 30.75 | 30.75, 30.75 | 0.0025 | +| 1447 | `instead` 62222 | `with` 3629 | 35, 34.75 | 34.5, 34.5 | 0.0129 | + Served through the recipe harness (an idle 97 GiB server of another model was loaded beside it, so memory readings include it), `--context 131072 --parallel 2`, greedy 256-token replies: @@ -107,5 +131,11 @@ Each keeps the replies' bits. No quality suite was run. No drafter is served (the family has no draft head; copy drafts come from the lane core). Logprobs are refused by name: the release's tokenizer decodes through a `Strip` step, which the engine's tokenizer does not count as token-local, so the server cannot give a token's bytes on its own. -Only the M5's tensor units were checked; M1-M4 Macs and smaller RAM are untested. Tied heads, other head sizes, -Llama 3's scaled RoPE and biased projections are refused at load rather than served untested. +Only the M5's tensor units were checked; M1-M4 Macs are refused at load, and smaller RAM is untested. Other `llama` +checkpoints, tied heads, other head sizes, Llama 3's scaled RoPE and biased projections are refused at load rather than +served untested. + +## Licence + +Nanbeige4.3-33B is published under Apache-2.0 at the revision measured here +(`c1093c501cc6460427c0ac634db7223d0450c0ee`, the licence in its model card). diff --git a/zig/src/families/llama/config.zig b/zig/src/families/llama/config.zig index 5db0942bab..b7a3609170 100644 --- a/zig/src/families/llama/config.zig +++ b/zig/src/families/llama/config.zig @@ -1,4 +1,4 @@ -//! Llama-architecture (model_type llama) admission: a dense bf16 GQA stack, heads of 128, default RoPE, its own head. +//! Llama (model_type llama) admission: dense bf16 GQA, 128-wide heads, default RoPE; `qualified` checkpoints only. const std = @import("std"); pub const head_dim = 128; @@ -14,8 +14,27 @@ pub const Error = error{ UnsupportedLlamaGeometry, UnsupportedLlamaRotary, UnsupportedLlamaPrecision, + UnqualifiedLlamaCheckpoint, } || std.json.ParseError(std.json.Scanner); +/// A checkpoint checked against HF on the device, by the geometry its config names. +pub const Qualified = struct { + name: []const u8, + hidden: usize, + intermediate: usize, + layers: usize, + heads: usize, + kv_heads: usize, + vocab: usize, + theta: f64, + context: usize, +}; + +/// The checkpoints the family serves; another `llama` checkpoint is refused until it has been checked the same way. +pub const qualified = [_]Qualified{ + .{ .name = "Nanbeige4.3-33B", .hidden = 7168, .intermediate = 22016, .layers = 52, .heads = 56, .kv_heads = 8, .vocab = 166912, .theta = 120000000, .context = 524288 }, +}; + /// Attention's score scale, 1 / sqrt(head_dim). pub fn headScale() f32 { return 1.0 / @sqrt(@as(f32, head_dim)); @@ -39,7 +58,18 @@ pub const Config = struct { defer gpa.free(path); const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 22)); defer gpa.free(bytes); - return parse(gpa, bytes); + const c = try parse(gpa, bytes); + _ = try c.qualification(); + return c; + } + + /// The qualified checkpoint this config names. + pub fn qualification(c: Config) Error!*const Qualified { + for (&qualified) |*q| { + if (c.hidden == q.hidden and c.intermediate == q.intermediate and c.layers == q.layers and c.heads == q.heads and + c.kv_heads == q.kv_heads and c.vocab == q.vocab and c.theta == q.theta and c.context == q.context) return q; + } + return error.UnqualifiedLlamaCheckpoint; } pub fn queryWidth(c: Config) usize { @@ -196,6 +226,26 @@ test "Nanbeige4.3-33B's config is admitted with its geometry, context and end to try std.testing.expect(c.isEos(166101) and !c.isEos(166100)); } +test "only the qualified checkpoint is admitted: a llama of another geometry parses but is refused" { + const gpa = std.testing.allocator; + const c = try parse(gpa, fixture); + try std.testing.expectEqualStrings("Nanbeige4.3-33B", (try c.qualification()).name); + const cases = .{ + .{ "\"hidden_size\": 7168", "\"hidden_size\": 4096" }, + .{ "\"num_hidden_layers\": 52", "\"num_hidden_layers\": 32" }, + .{ "\"num_key_value_heads\": 8", "\"num_key_value_heads\": 7" }, + .{ "\"rope_theta\": 120000000", "\"rope_theta\": 500000" }, + .{ "\"max_position_embeddings\": 524288", "\"max_position_embeddings\": 131072" }, + }; + inline for (cases) |k| { + const bytes = try std.mem.replaceOwned(u8, gpa, fixture, k[0], k[1]); + defer gpa.free(bytes); + try std.testing.expect(!std.mem.eql(u8, bytes, fixture)); + const d = try parse(gpa, bytes); + try std.testing.expectError(error.UnqualifiedLlamaCheckpoint, d.qualification()); + } +} + test "other layouts are refused by name" { const gpa = std.testing.allocator; const cases = .{ diff --git a/zig/src/families/llama/model.zig b/zig/src/families/llama/model.zig index 0d39ea15dd..4bdf970516 100644 --- a/zig/src/families/llama/model.zig +++ b/zig/src/families/llama/model.zig @@ -37,6 +37,8 @@ pub const Model = struct { m.config = try cfg.Config.read(gpa, io, dir); m.device = try mtl.Device.init(); errdefer m.device.deinit(); + // prompt chunks' GEMM and attention run on Metal 4 tensor units: an M5 or later + if (!m.device.tensorUnits()) return error.LlamaNeedsTensorUnits; m.queue = try m.device.queue(); errdefer m.queue.deinit(); m.checkpoint = ck.Checkpoint.init(gpa); diff --git a/zig/src/native/llama.zig b/zig/src/native/llama.zig index bbfc69c779..3451bf9c1f 100644 --- a/zig/src/native/llama.zig +++ b/zig/src/native/llama.zig @@ -38,7 +38,11 @@ pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Ope const pool = mtl.objc.Pool.push(); defer pool.pop(); const model = fl.Model.load(gpa, io, o.dir) catch |err| { - problem.* = try std.fmt.allocPrint(a, "the native Llama engine cannot load this checkpoint ({s})", .{@errorName(err)}); + problem.* = switch (err) { + error.LlamaNeedsTensorUnits => "the native Llama engine needs Metal 4 tensor units (an M5 or later) for its prompt GEMM and attention: this Mac's GPU (M1 to M4) has none", + error.UnqualifiedLlamaCheckpoint => try std.fmt.allocPrint(a, "the native Llama engine serves only the checkpoints checked against Hugging Face ({s}): this llama checkpoint's geometry is not one of them", .{qualifiedNames()}), + else => try std.fmt.allocPrint(a, "the native Llama engine cannot load this checkpoint ({s})", .{@errorName(err)}), + }; return null; }; errdefer model.deinit(); @@ -96,6 +100,13 @@ pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Ope return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h }; } +/// The qualified checkpoints' names, comma-separated. +fn qualifiedNames() []const u8 { + comptime var names: []const u8 = ""; + inline for (fl.config.qualified, 0..) |q, i| names = names ++ (if (i == 0) "" else ", ") ++ q.name; + return names; +} + /// Metal.room's fit: the kept states' budget is what the streams leave, at most the cap; the least valuable go first. fn fitCache(ctx: *anyopaque, bytes: u64, restoring: ?*anyopaque, prompt: []const u32) void { const h: *Host = @ptrCast(@alignCast(ctx)); diff --git a/zig/tests/llama/run.zig b/zig/tests/llama/run.zig index 2c70c10e32..962faf3844 100644 --- a/zig/tests/llama/run.zig +++ b/zig/tests/llama/run.zig @@ -9,10 +9,9 @@ const usage = \\tf-llama-run logits 16-row windows from 0; bf16 logits [tokens, vocab] \\tf-llama-run prompt the prompt in chunks; the last row's bf16 logits \\tf-llama-run forced , then cont's tokens a step; bf16 logits before each - \\tf-llama-run exact cont's tokens as one-row steps and as 16-row windows - \\ after the prompt: every row's logits bit for bit + \\tf-llama-run exact [,] cont's tokens as one-row steps and as w-row windows + \\ (16) after the prompt: every row's logits bit for bit \\tf-llama-run bench the prompt, then n one-row steps: wall, encode, GPU ms - \\LLAMA_LAYERS=N runs the first N layers only (a truncated reference's check) \\ ; @@ -33,8 +32,6 @@ pub fn main(init: std.process.Init) !void { if (prompt.len == 0) return error.EmptyPrompt; const m = try fl.Model.load(gpa, init.io, args[1]); defer m.deinit(); - // LLAMA_LAYERS=N: the first N layers only, for a layer-by-layer check against a truncated reference - if (std.c.getenv("LLAMA_LAYERS")) |v| m.weights.layers = m.weights.layers[0..@min(m.weights.layers.len, try std.fmt.parseInt(usize, std.mem.span(v), 10))]; const mode = args[3]; var forced: ?std.json.Parsed([]u32) = null; defer if (forced) |f| f.deinit(); @@ -71,11 +68,13 @@ pub fn main(init: std.process.Init) !void { return; } if (std.mem.eql(u8, mode, "exact")) { - const text = try std.Io.Dir.cwd().readFileAlloc(init.io, args[4], gpa, .limited(1 << 20)); + var parts = std.mem.splitScalar(u8, args[4], ','); + const text = try std.Io.Dir.cwd().readFileAlloc(init.io, parts.first(), gpa, .limited(1 << 20)); defer gpa.free(text); const cont = try std.json.parseFromSlice([]u32, gpa, text, .{}); defer cont.deinit(); const toks = cont.value; + const width: usize = if (parts.next()) |w| try std.fmt.parseInt(usize, w, 10) else fl.state.window_rows; // both caches at one capacity: the decode blocks follow it var scache = try fl.state.Cache.init(m.device, &m.config, prompt.len + toks.len + 1); defer scache.deinit(); @@ -90,8 +89,6 @@ pub fn main(init: std.process.Init) !void { @memcpy(one[i * vocab ..][0..vocab], small.logits_all.slice(u16, vocab)); scache.len += 1; } - // LLAMA_WINDOW=W: windows of W rows (16 by default) - const width: usize = if (std.c.getenv("LLAMA_WINDOW")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else fl.state.window_rows; if (width == 0 or width > fl.state.window_rows) return error.BadWindow; var differ: usize = 0; var at: usize = 0; From a4d143fd90569667e0d8c8052db9eca93b0089ce Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Wed, 28 Oct 2026 01:33:21 -0500 Subject: [PATCH 4/7] llama: the exactness mode fails on any differing row, and every run checks a planted row is reported --- zig/src/families/llama/exact.zig | 33 ++++++++++++++++++++++++++++++++ zig/src/families/llama/llama.zig | 2 ++ zig/tests/llama/run.zig | 23 ++++++++++++++-------- 3 files changed, 50 insertions(+), 8 deletions(-) create mode 100644 zig/src/families/llama/exact.zig diff --git a/zig/src/families/llama/exact.zig b/zig/src/families/llama/exact.zig new file mode 100644 index 0000000000..2958a96413 --- /dev/null +++ b/zig/src/families/llama/exact.zig @@ -0,0 +1,33 @@ +//! The exactness check's comparison: one-row steps' logits against the same rows' from windows, bit for bit. +const std = @import("std"); + +/// How many rows of `vocab` bf16 logits differ in any bit between `a` and `b`; the first ones' indices go to `first`. +pub fn differing(a: []const u16, b: []const u16, vocab: usize, first: []usize) usize { + std.debug.assert(vocab > 0 and a.len == b.len and a.len % vocab == 0); + var n: usize = 0; + for (0..a.len / vocab) |r| { + if (std.mem.eql(u16, a[r * vocab ..][0..vocab], b[r * vocab ..][0..vocab])) continue; + if (n < first.len) first[n] = r; + n += 1; + } + return n; +} + +/// Flips the lowest bit of one logit in row `row`: the planted difference the check must report. +pub fn plant(rows: []u16, vocab: usize, row: usize) void { + rows[row * vocab + vocab / 2] ^= 1; +} + +test "llama exact: equal rows pass, a planted row is reported by index, and only it" { + const vocab = 8; + var a: [4 * vocab]u16 = undefined; + for (&a, 0..) |*x, i| x.* = @intCast(0x3f80 + i); + var b = a; + var first: [4]usize = undefined; + try std.testing.expectEqual(@as(usize, 0), differing(&a, &b, vocab, &first)); + plant(&b, vocab, 2); + try std.testing.expectEqual(@as(usize, 1), differing(&a, &b, vocab, &first)); + try std.testing.expectEqual(@as(usize, 2), first[0]); + plant(&b, vocab, 2); + try std.testing.expectEqual(@as(usize, 0), differing(&a, &b, vocab, &first)); +} diff --git a/zig/src/families/llama/llama.zig b/zig/src/families/llama/llama.zig index ecaa37ae83..492d2c327f 100644 --- a/zig/src/families/llama/llama.zig +++ b/zig/src/families/llama/llama.zig @@ -9,8 +9,10 @@ pub const backend = @import("backend.zig"); pub const budget = @import("budget.zig"); pub const residency = @import("residency.zig"); pub const snapshot = @import("snapshot.zig"); +pub const exact = @import("exact.zig"); test { + _ = exact; _ = config; _ = contracts; _ = state; diff --git a/zig/tests/llama/run.zig b/zig/tests/llama/run.zig index 962faf3844..f4b1bc46b0 100644 --- a/zig/tests/llama/run.zig +++ b/zig/tests/llama/run.zig @@ -10,7 +10,8 @@ const usage = \\tf-llama-run prompt the prompt in chunks; the last row's bf16 logits \\tf-llama-run forced , then cont's tokens a step; bf16 logits before each \\tf-llama-run exact [,] cont's tokens as one-row steps and as w-row windows - \\ (16) after the prompt: every row's logits bit for bit + \\ (16) after the prompt: every row's logits bit for bit, + \\ failing on any difference or a missed planted row \\tf-llama-run bench the prompt, then n one-row steps: wall, encode, GPU ms \\ ; @@ -90,21 +91,27 @@ pub fn main(init: std.process.Init) !void { scache.len += 1; } if (width == 0 or width > fl.state.window_rows) return error.BadWindow; - var differ: usize = 0; + const wins = try gpa.alloc(u16, toks.len * vocab); + defer gpa.free(wins); var at: usize = 0; while (at < toks.len) { const n = @min(toks.len - at, width); try fl.forward.run(m, &small, &wcache, toks[at..][0..n], .{ .all = true }); - for (0..n) |r| { - if (!std.mem.eql(u16, one[(at + r) * vocab ..][0..vocab], small.logits_all.slice(u16, n * vocab)[r * vocab ..][0..vocab])) { - differ += 1; - if (differ <= 4) std.debug.print("row {d} (window row {d}) differs\n", .{ at + r, r }); - } - } + @memcpy(wins[at * vocab ..][0 .. n * vocab], small.logits_all.slice(u16, n * vocab)); wcache.len += n; at += n; } + var first: [4]usize = undefined; + const differ = fl.exact.differing(one, wins, vocab, &first); + for (first[0..@min(differ, first.len)]) |r| std.debug.print("row {d} (window row {d}) differs\n", .{ r, r % width }); std.debug.print("llama exact at {d}: {d} of {d} rows' logits differ between one-row steps and windows\n", .{ prompt.len, differ, toks.len }); + // the control: a row planted in a copy of the steps' logits must be the one difference reported + const planted = try gpa.dupe(u16, one); + defer gpa.free(planted); + fl.exact.plant(planted, vocab, toks.len / 2); + if (fl.exact.differing(one, planted, vocab, &first) != 1 or first[0] != toks.len / 2) return error.PlantedRowMissed; + std.debug.print("llama exact: the planted row {d} is reported\n", .{toks.len / 2}); + if (differ > 0) return error.RowsDiffer; return; } try feed(m, &small, &wide, &cache, prompt); From 24e98179e151dc2d817e25184ea2688b596f364c Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Wed, 28 Oct 2026 02:00:17 -0500 Subject: [PATCH 5/7] llama: a judge failing on any teacher-forced difference that is not a near-tie, a float64 reference, the doc's misses and a one-process receipt --- docs/recipes/nanbeige-4.3-33b-metal.md | 73 ++++++++++++++++------ tools/zig/llama_f64_reference.py | 86 ++++++++++++++++++++++++++ tools/zig/llama_teacher_judge.py | 65 +++++++++++++++++++ tools/zig/llama_teacher_judge_test.py | 44 +++++++++++++ 4 files changed, 248 insertions(+), 20 deletions(-) create mode 100644 tools/zig/llama_f64_reference.py create mode 100644 tools/zig/llama_teacher_judge.py create mode 100644 tools/zig/llama_teacher_judge_test.py diff --git a/docs/recipes/nanbeige-4.3-33b-metal.md b/docs/recipes/nanbeige-4.3-33b-metal.md index a2f4a1ff1d..c490c860eb 100644 --- a/docs/recipes/nanbeige-4.3-33b-metal.md +++ b/docs/recipes/nanbeige-4.3-33b-metal.md @@ -62,26 +62,59 @@ greedy chats of 64 tokens (code, prose, Chinese, a document, with and without th token; the other three part at a position where the engine's two best logits are the same bf16 value. One-row steps and windows of 7 or 16 rows give the same logits bit for bit after prompts of 64 and 1,500 tokens. -The 13 teacher-forced positions of the 1,500 whose top tokens differ (row n predicts token n + 1; a tie goes to the -lower id on both sides). At each, HF's and the engine's two contenders are within one bf16 step of each other on both -sides (0.125 below 32, 0.25 from 32 to 64): HF's two tie at 3 of them, the engine's at 8, neither at 22 and 166, and -at 1370 a third token, `no`, ties at the top on both sides. - -| position | HF's top-1 (a) | engine's top-1 (b) | HF's logits a, b | engine's logits a, b | KL | -|---:|---|---|---|---|---:| -| 22 | `en` 271 | `see` 3574 | 20.75, 20.625 | 20.625, 20.75 | 0.0047 | -| 83 | `the` 269 | `Python` 9120 | 26.25, 26.25 | 26.25, 26.375 | 0.0022 | -| 85 | `developed` 3260 | `-` 152369 | 35.5, 35.5 | 35.5, 35.75 | 0.0134 | -| 166 | `";` 2651 | `"))` 9441 | 36.25, 36 | 36.25, 36.5 | 0.0234 | -| 224 | `(` 152374 | `('` 1091 | 39.75, 39.5 | 39.5, 39.5 | 0.0078 | -| 401 | `(` 152374 | `({'` 50468 | 39.5, 39.25 | 39.25, 39.25 | 0.0076 | -| 425 | `{` 152405 | `1` 152367 | 40, 39.75 | 40, 40 | 0.0163 | -| 586 | `representation` 7443 | `encoding` 13571 | 30, 30 | 29.875, 30 | 0.0027 | -| 1302 | `If` 3934 | ``` `` ``` 2283 | 32, 31.875 | 32, 32 | 0.0022 | -| 1304 | `indent` 28799 | `allow` 13471 | 37.25, 37 | 37, 37 | 0.0078 | -| 1370 | `it` 388 | `the` 269 | 38.25, 38 (`no` 754: 38.25) | 38, 38 (`no` 754: 38) | 0.0076 | -| 1384 | ``` `` ``` 16411 | `circular` 13380 | 30.875, 30.75 | 30.75, 30.75 | 0.0025 | -| 1447 | `instead` 62222 | `with` 3629 | 35, 34.75 | 34.5, 34.5 | 0.0129 | +`tools/zig/llama_teacher_judge.py` judges every teacher-forced position by the repository's near-tie rule +(`qwen27_margins.py`: the same two contenders on both sides, each margin within one bf16 step) and exits 1 on anything +else. Against HF's bf16 run, 13 of the 1,500 positions differ and two are red: at 425 and 1370 three tokens sit within +one bf16 step, so the two sides' contenders differ. `tools/zig/llama_f64_reference.py` runs the release in float64 +(its bf16 weights cast up exactly, HF's arithmetic, one rounding at the end) and settles them: + +- Against the float64 pass, the engine passes the rule: 1,483 of 1,500 positions equal, all 17 differences near-ties. + HF's bf16 run does not: 1,488 equal, and at 425 it rounds two contenders a step high (red). +- At 1370 the exact logits are `it` 388 38.1266, `no` 754 38.1037 and `the` 269 38.0106: within 0.116 of each other, + under half a bf16 step (0.25 at 32-64). 388 lies 0.0016 above the rounding midpoint 38.125, so rounded once it is + 38.25 and the others 38.0. HF's bf16 run gives 754 a step too high (38.25); the engine gives 388 a step too low + (38.0). +- Layer by layer at 1370 (the engine's forward stopped after each layer count by a local probe, not in this PR; HF's + `output_hidden_states`), the engine's residual and HF's drift from the float64 one alike (0.3% after the first layer, + 1.8% after the 51st; 0.80% and 0.82% on average over the layers). The engine's largest excess, 1.5-1.7 times HF's at + layers 20-22, is channel 4086, a massive activation near 427 where a bf16 step is 2: both hold 430 after layer 19, + and at layer 20 HF's add lands on 428 while the engine's stays at 430. + +Every position where the engine's top token differs from HF's or from the float64 pass's (row n predicts token +n + 1; a tie goes to the lower id; logits listed for each side's top-1 tokens in the order of the first three columns): + +| position | float64's top-1 | HF's | engine's | float64 (rounded once) | HF | engine | against HF | against float64 | +|---:|---|---|---|---|---|---|---|---| +| 22 | `en` 271 | `en` 271 | `see` 3574 | 20.625, 20.625 | 20.75, 20.625 | 20.625, 20.75 | near-tie | near-tie | +| 64 | `pickle` 82922 | `json` 5174 | `json` 5174 | 31.625, 31.5 | 31.625, 31.625 | 31.75, 31.75 | equal | near-tie | +| 83 | `the` 269 | `the` 269 | `Python` 9120 | 26.25, 26.25 | 26.25, 26.25 | 26.25, 26.375 | near-tie | near-tie | +| 85 | `developed` 3260 | `developed` 3260 | `-` 152369 | 35.5, 35.5 | 35.5, 35.5 | 35.5, 35.75 | near-tie | near-tie | +| 166 | `"))` 9441 | `";` 2651 | `"))` 9441 | 36.5, 36.25 | 36, 36.25 | 36.5, 36.25 | near-tie | equal | +| 171 | `\"` 12028 | `[\` 11752 | `[\` 11752 | 31.75, 31.625 | 32, 32.25 | 32.25, 32.25 | equal | near-tie | +| 224 | `('` 1091 | `(` 152374 | `('` 1091 | 39.75, 39.75 | 39.5, 39.75 | 39.5, 39.5 | near-tie | equal | +| 235 | `.` 152361 | `,` 152363 | `,` 152363 | 45.5, 45.5 | 45.25, 45.5 | 45.25, 45.5 | equal | near-tie | +| 357 | `separators` 118153 | `indent` 34489 | `indent` 34489 | 37, 36.75 | 36.75, 36.75 | 36.75, 36.75 | equal | near-tie | +| 401 | `(` 152374 | `(` 152374 | `({'` 50468 | 39.5, 39.25 | 39.5, 39.25 | 39.25, 39.25 | near-tie | near-tie | +| 425 | `1` 152367 | `{` 152405 | `1` 152367 | 39.75, 39.75 | 39.75, 40 | 40, 40 | bug | equal | +| 586 | `encoding` 13571 | `representation` 7443 | `encoding` 13571 | 29.75, 29.625 | 30, 30 | 30, 29.875 | near-tie | equal | +| 1234 | `file` 2408 | `a` 152346 | `a` 152346 | 35, 35 | 35, 35.25 | 35, 35.25 | equal | near-tie | +| 1295 | `an` 284 | `a` 261 | `a` 261 | 44, 43.75 | 43.75, 43.75 | 43.75, 43.75 | equal | near-tie | +| 1302 | `If` 3934 | `If` 3934 | `` `` `` 2283 | 32.25, 32 | 32, 31.875 | 32, 32 | near-tie | near-tie | +| 1304 | `indent` 28799 | `indent` 28799 | `allow` 13471 | 37.25, 37 | 37.25, 37 | 37, 37 | near-tie | near-tie | +| 1352 | `the` 269 | `\n` 13 | `\n` 13 | 46.25, 46 | 46, 46 | 46, 46 | equal | near-tie | +| 1370 | `it` 388 | `it` 388 | `the` 269 | 38.25, 38 | 38.25, 38 | 38, 38 | bug | near-tie | +| 1381 | `skipped` 61065 | `disabled` 18191 | `disabled` 18191 | 41.5, 41.25 | 41.25, 41.25 | 41.25, 41.25 | equal | near-tie | +| 1384 | `` `` `` 16411 | `` `` `` 16411 | `circular` 13380 | 30.875, 30.75 | 30.875, 30.75 | 30.75, 30.75 | near-tie | near-tie | +| 1447 | `instead` 62222 | `instead` 62222 | `with` 3629 | 35.5, 35.25 | 35, 34.75 | 34.5, 34.5 | near-tie | near-tie | + +A device receipt with no other model server or GPU job on the Mac (28 Oct 2026, this branch's build; a CPU float64 +reference run, one layer's weights at a time, overlapped its first minutes; greedy, 128-token replies, +`return_token_ids`, `--context 32768 --parallel 2`): the 1,500-position logits are bit-identical to the ones judged +above, and the six greedy chats to the earlier run; `tf-llama-run ... exact` finds 0 differing rows after the 64- and +1,500-token prompts at windows of 16 and 7 rows, and reports its planted row each time. Six prompts sent two at a time +return the same token ids as each one alone. Three chats whose first turn carries an 8,500-token log: the second turn +resumed from the prompt cache (8,521-8,528 of 8,554-8,568 prompt tokens kept) returns the same ids as on a server +without the cache. No choice differed, so there are no margins to list. Served through the recipe harness (an idle 97 GiB server of another model was loaded beside it, so memory readings include it), `--context 131072 --parallel 2`, greedy 256-token replies: diff --git a/tools/zig/llama_f64_reference.py b/tools/zig/llama_f64_reference.py new file mode 100644 index 0000000000..6d83bfa5f2 --- /dev/null +++ b/tools/zig/llama_f64_reference.py @@ -0,0 +1,86 @@ +"""A Llama checkpoint's teacher-forced logits in float64 on the CPU, rounded once to bf16: the exact reference.""" +# usage: python3 tools/zig/llama_f64_reference.py MODEL_DIR tokens.json out.bin (raw bf16 [positions, vocab]) +import json +import sys +import time + +import numpy as np +import torch +from safetensors import safe_open + + +def bf16_once(x: np.ndarray) -> np.ndarray: + """float64 to bf16 bits, rounded once to nearest (ties to even) at bf16's spacing.""" + spacing = np.ldexp(1.0, np.frexp(np.abs(x))[1] - 1 - 7) + y = (np.round(x / spacing) * spacing).astype(np.float32) # exact: y has at most 8 significant bits + return (y.view(np.uint32) >> 16).astype(np.uint16) + + +def main(model_dir: str, tokens: str, out: str) -> None: + torch.set_num_threads(24) + cfg = json.load(open(f"{model_dir}/config.json")) + index = json.load(open(f"{model_dir}/model.safetensors.index.json"))["weight_map"] + if cfg.get("rope_scaling") or cfg.get("attention_bias") or cfg.get("mlp_bias") or cfg.get("tie_word_embeddings"): + raise SystemExit("only the plain Llama layout: no scaled RoPE, biases or tied head") + layers, heads, kv = cfg["num_hidden_layers"], cfg["num_attention_heads"], cfg["num_key_value_heads"] + dim = cfg.get("head_dim") or cfg["hidden_size"] // heads + eps, theta = cfg["rms_norm_eps"], float(cfg["rope_theta"]) + f64 = torch.float64 + quant = cfg.get("quantization") + if quant and (quant.get("mode", "affine") != "affine" or quant["bits"] not in (4, 8)): + raise SystemExit("only MLX affine checkpoints at 4 or 8 bits") + + def tensor(name): + with safe_open(f"{model_dir}/{index[name]}", framework="pt") as f: + return f.get_tensor(name) + + def w(name): + if not name.endswith(".weight") or name[: -len("weight")] + "scales" not in index: + return tensor(name).to(f64) + # mlx_lm's layout: u32 words of `bits`-bit codes, lowest bits first; a scale and a bias a group of codes + bits, group = quant["bits"], quant["group_size"] + words = tensor(name).to(torch.int64) & 0xFFFFFFFF + shifts = torch.arange(0, 32, bits, dtype=torch.int64) + q = ((words.unsqueeze(-1) >> shifts) & ((1 << bits) - 1)).reshape(words.shape[0], -1).to(f64) + scales = tensor(name[: -len("weight")] + "scales").to(f64).repeat_interleave(group, -1) + biases = tensor(name[: -len("weight")] + "biases").to(f64).repeat_interleave(group, -1) + return scales * q + biases + + def norm(x, g): + return g * (x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)) + + def rotate(x): + return torch.cat((-x[..., dim // 2 :], x[..., : dim // 2]), -1) + + ids = json.load(open(tokens)) + n = len(ids) + t0 = time.time() + with torch.no_grad(): + inv = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=f64) / dim)) + angle = torch.outer(torch.arange(n, dtype=f64), inv) + cos, sin = torch.cat((angle, angle), -1).cos(), torch.cat((angle, angle), -1).sin() + x = w("model.embed_tokens.weight")[torch.tensor(ids)] + mask = torch.full((n, n), float("-inf"), dtype=f64).triu(1) + for i in range(layers): + p = f"model.layers.{i}." + h = norm(x, w(p + "input_layernorm.weight")) + q = (h @ w(p + "self_attn.q_proj.weight").T).view(n, heads, dim).transpose(0, 1) + k = (h @ w(p + "self_attn.k_proj.weight").T).view(n, kv, dim).transpose(0, 1) + v = (h @ w(p + "self_attn.v_proj.weight").T).view(n, kv, dim).transpose(0, 1) + q, k = q * cos + rotate(q) * sin, k * cos + rotate(k) * sin + k, v = k.repeat_interleave(heads // kv, 0), v.repeat_interleave(heads // kv, 0) + a = torch.softmax(q @ k.transpose(1, 2) / dim**0.5 + mask, -1) @ v + x = x + a.transpose(0, 1).reshape(n, heads * dim) @ w(p + "self_attn.o_proj.weight").T + h = norm(x, w(p + "post_attention_layernorm.weight")) + x = x + (torch.nn.functional.silu(h @ w(p + "mlp.gate_proj.weight").T) * (h @ w(p + "mlp.up_proj.weight").T)) @ w(p + "mlp.down_proj.weight").T + print(f"layer {i + 1} of {layers}, {time.time() - t0:.0f} s", flush=True) + h = norm(x, w("model.norm.weight")) + head = w("lm_head.weight") + with open(out, "wb") as f: + for s in range(0, n, 64): + bf16_once((h[s : s + 64] @ head.T).numpy()).tofile(f) + print(f"{n} positions' logits in float64, rounded once to bf16, in {time.time() - t0:.0f} s") + + +if __name__ == "__main__": + main(*sys.argv[1:4]) diff --git a/tools/zig/llama_teacher_judge.py b/tools/zig/llama_teacher_judge.py new file mode 100644 index 0000000000..e0fe624b60 --- /dev/null +++ b/tools/zig/llama_teacher_judge.py @@ -0,0 +1,65 @@ +"""Judge a Llama checkpoint's teacher-forced logits against a reference's at every position.""" +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path + +from qwen27_margins import verdict + + +def top2(row) -> dict: + """A row's best two ids and scores, ties to the lower id (the engine's argmax and numpy's).""" + import numpy as np + + first = int(np.argmax(row)) + rest = row.copy() + rest[first] = -np.inf + second = int(np.argmax(rest)) + return {"ids": [first, second], "logits": [float(row[first]), float(row[second])]} + + +def judge(native: list[dict], reference: list[dict]) -> dict: + if len(native) != len(reference) or not native: + raise ValueError("expected the same positive number of positions on both sides") + exceptions = [] + for position, (a, b) in enumerate(zip(native, reference)): + actual, wanted = a["ids"][0], b["ids"][0] + if actual != wanted: + exceptions.append({"position": position, "expected": wanted, "actual": actual, **verdict(a, b, wanted, actual)}) + return {"positions": len(native), "strict_equal": len(native) - len(exceptions), "exceptions": exceptions, + "policy_pass": all(item["verdict"] == "near-tie" for item in exceptions)} + + +def load(path: Path, vocab: int): + import numpy as np + + words = np.fromfile(path, dtype=np.uint16) + if words.size % vocab: + raise ValueError(f"{path}: not whole rows of {vocab} bf16 logits") + return (words.astype(np.uint32) << 16).view(np.float32).reshape(-1, vocab) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--reference", type=Path, required=True) + parser.add_argument("--native", type=Path, required=True) + parser.add_argument("--vocab", type=int, required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + ref, nat = load(args.reference, args.vocab), load(args.native, args.vocab) + result = judge([top2(row) for row in nat], [top2(row) for row in ref]) + result["sha256"] = {key: hashlib.sha256(path.read_bytes()).hexdigest() for key, path in (("reference", args.reference), ("native", args.native))} + with args.output.open("x") as out: + out.write(json.dumps(result, indent=2) + "\n") + for item in result["exceptions"]: + print(f"position {item['position']}: {item['verdict']}: reference {item['python_top2']}, native {item['native_top2']}") + print(f"{result['strict_equal']} of {result['positions']} positions equal; {len(result['exceptions'])} differ; " + f"policy {'pass' if result['policy_pass'] else 'FAIL'}") + if not result["policy_pass"]: + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/tools/zig/llama_teacher_judge_test.py b/tools/zig/llama_teacher_judge_test.py new file mode 100644 index 0000000000..3617f67c87 --- /dev/null +++ b/tools/zig/llama_teacher_judge_test.py @@ -0,0 +1,44 @@ +"""The Llama teacher judge passes equal rows and one-step ties, and is red on wider margins and other contenders.""" +import copy +import unittest + +from llama_teacher_judge import judge + + +def rows(n=8): + return [{"ids": [1, 2], "logits": [38.25, 38.0]} for _ in range(n)] + + +class LlamaTeacherJudgeTest(unittest.TestCase): + def test_equal(self): + result = judge(rows(), rows()) + self.assertEqual(result["strict_equal"], 8) + self.assertTrue(result["policy_pass"]) + + def test_one_step_tie(self): + native = rows() + native[3] = {"ids": [2, 1], "logits": [38.0, 38.0]} + result = judge(native, rows()) + self.assertTrue(result["policy_pass"]) + self.assertEqual([item["position"] for item in result["exceptions"]], [3]) + + def test_wider_margin_is_red(self): + native = rows() + native[0] = {"ids": [2, 1], "logits": [38.5, 38.0]} + self.assertFalse(judge(native, rows())["policy_pass"]) + + def test_third_contender_is_red(self): + # three tokens within a step (#1415's position 1370): the two sides' contenders differ + native = rows() + native[5] = {"ids": [0, 1], "logits": [38.0, 38.0]} + reference = copy.deepcopy(rows()) + reference[5] = {"ids": [1, 2], "logits": [38.25, 38.25]} + self.assertFalse(judge(native, reference)["policy_pass"]) + + def test_lengths_must_match(self): + with self.assertRaises(ValueError): + judge(rows(7), rows(8)) + + +if __name__ == "__main__": + unittest.main() From 7609b48d90412099c74c362b261e1811d40eb5ea Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Wed, 28 Oct 2026 08:11:54 -0500 Subject: [PATCH 6/7] llama: the teacher-forced checker refuses a row holding NaN or infinity, on either side, before comparing --- tools/zig/llama_teacher_judge.py | 23 ++++++++- tools/zig/llama_teacher_judge_test.py | 71 ++++++++++++++++++++++++++- 2 files changed, 91 insertions(+), 3 deletions(-) diff --git a/tools/zig/llama_teacher_judge.py b/tools/zig/llama_teacher_judge.py index e0fe624b60..ff94cfb41b 100644 --- a/tools/zig/llama_teacher_judge.py +++ b/tools/zig/llama_teacher_judge.py @@ -4,15 +4,18 @@ import argparse import hashlib import json +import math from pathlib import Path from qwen27_margins import verdict def top2(row) -> dict: - """A row's best two ids and scores, ties to the lower id (the engine's argmax and numpy's).""" + """A row's best two ids and scores, ties to the lower id (as the engine's argmax); non-finite rows are refused.""" import numpy as np + if not np.isfinite(row).all(): + raise ValueError(f"non-finite logits at id {int(np.flatnonzero(~np.isfinite(row))[0])}") first = int(np.argmax(row)) rest = row.copy() rest[first] = -np.inf @@ -25,6 +28,9 @@ def judge(native: list[dict], reference: list[dict]) -> dict: raise ValueError("expected the same positive number of positions on both sides") exceptions = [] for position, (a, b) in enumerate(zip(native, reference)): + for side, item in (("native", a), ("reference", b)): + if not all(math.isfinite(x) for x in item["logits"]): + raise ValueError(f"position {position}: non-finite {side} logits {item['logits']}") actual, wanted = a["ids"][0], b["ids"][0] if actual != wanted: exceptions.append({"position": position, "expected": wanted, "actual": actual, **verdict(a, b, wanted, actual)}) @@ -49,7 +55,20 @@ def main(): parser.add_argument("--output", type=Path, required=True) args = parser.parse_args() ref, nat = load(args.reference, args.vocab), load(args.native, args.vocab) - result = judge([top2(row) for row in nat], [top2(row) for row in ref]) + try: + sides = [] + for name, rows in (("native", nat), ("reference", ref)): + records = [] + for position, row in enumerate(rows): + try: + records.append(top2(row)) + except ValueError as err: + raise ValueError(f"{name} position {position}: {err}") from None + sides.append(records) + result = judge(*sides) + except ValueError as err: + print(f"refused: {err}") + raise SystemExit(1) result["sha256"] = {key: hashlib.sha256(path.read_bytes()).hexdigest() for key, path in (("reference", args.reference), ("native", args.native))} with args.output.open("x") as out: out.write(json.dumps(result, indent=2) + "\n") diff --git a/tools/zig/llama_teacher_judge_test.py b/tools/zig/llama_teacher_judge_test.py index 3617f67c87..1b1ba8128d 100644 --- a/tools/zig/llama_teacher_judge_test.py +++ b/tools/zig/llama_teacher_judge_test.py @@ -1,8 +1,19 @@ """The Llama teacher judge passes equal rows and one-step ties, and is red on wider margins and other contenders.""" import copy +import math +import os +import struct +import subprocess +import sys +import tempfile import unittest -from llama_teacher_judge import judge +from llama_teacher_judge import judge, top2 + +try: + import numpy +except ImportError: + numpy = None def rows(n=8): @@ -35,6 +46,64 @@ def test_third_contender_is_red(self): reference[5] = {"ids": [1, 2], "logits": [38.25, 38.25]} self.assertFalse(judge(native, reference)["policy_pass"]) + def test_non_finite_rows_are_refused_on_either_side(self): + for bad in (math.nan, math.inf, -math.inf): + for side in ("native", "reference"): + native, reference = rows(), rows() + row = {"ids": [1, 2], "logits": [bad, 38.0]} + (native if side == "native" else reference)[4] = row + with self.assertRaises(ValueError, msg=f"{bad} on the {side} side"): + judge(native, reference) + + def test_non_finite_with_the_same_winner_is_refused(self): + # the winning id unchanged on both sides must not hide the score + native, reference = rows(), rows() + native[2] = {"ids": [1, 2], "logits": [math.nan, 0.0]} + reference[2] = {"ids": [1, 2], "logits": [math.nan, 0.0]} + with self.assertRaises(ValueError): + judge(native, reference) + native[2] = {"ids": [1, 2], "logits": [38.25, -math.inf]} + with self.assertRaises(ValueError): + judge(native, rows()) + + @unittest.skipUnless(numpy, "numpy") + def test_top2_refuses_a_non_finite_row_anywhere(self): + for bad in (numpy.nan, numpy.inf, -numpy.inf): + row = numpy.array([1.0, 3.0, 2.0, 0.5], dtype=numpy.float32) + row[3] = bad # not the winner: the winning id would be unchanged + with self.assertRaises(ValueError): + top2(row) + + @unittest.skipUnless(numpy, "numpy") + def test_the_checker_exits_1_on_a_non_finite_row(self): + def bf16(values): + return b"".join(struct.pack("> 16) for v in values) + good = [1.0, 3.0, 2.0, 0.5] * 2 + here = os.path.dirname(os.path.abspath(__file__)) + for side in ("native", "reference"): + for bad in (math.nan, math.inf, -math.inf): + with tempfile.TemporaryDirectory() as d: + planted = list(good) + planted[4 + 3] = bad # row 1, not its winner + files = {"native": good, "reference": good} + files[side] = planted + for name, values in files.items(): + with open(os.path.join(d, name + ".bin"), "wb") as f: + f.write(bf16(values)) + run = subprocess.run([sys.executable, os.path.join(here, "llama_teacher_judge.py"), "--native", os.path.join(d, "native.bin"), + "--reference", os.path.join(d, "reference.bin"), "--vocab", "4", "--output", os.path.join(d, "out.json")], + capture_output=True, text=True, cwd=here) + self.assertEqual(run.returncode, 1, f"{bad} on the {side} side: {run.stdout}{run.stderr}") + self.assertIn("refused", run.stdout) + with tempfile.TemporaryDirectory() as d: # the control: the same files, all finite, pass + for name in ("native", "reference"): + with open(os.path.join(d, name + ".bin"), "wb") as f: + f.write(bf16(good)) + run = subprocess.run([sys.executable, os.path.join(here, "llama_teacher_judge.py"), "--native", os.path.join(d, "native.bin"), + "--reference", os.path.join(d, "reference.bin"), "--vocab", "4", "--output", os.path.join(d, "out.json")], + capture_output=True, text=True, cwd=here) + self.assertEqual(run.returncode, 0, run.stdout + run.stderr) + def test_lengths_must_match(self): with self.assertRaises(ValueError): judge(rows(7), rows(8)) From 567379e260396be8c178ce23dcdedb8385003523 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Wed, 28 Oct 2026 09:52:37 -0500 Subject: [PATCH 7/7] llama: the float64 reference can round at every one of the reference's rounding points (--bf16-points) --- tools/zig/llama_f64_reference.py | 37 +++++++++++++++++++++----------- 1 file changed, 24 insertions(+), 13 deletions(-) diff --git a/tools/zig/llama_f64_reference.py b/tools/zig/llama_f64_reference.py index 6d83bfa5f2..360d1210c7 100644 --- a/tools/zig/llama_f64_reference.py +++ b/tools/zig/llama_f64_reference.py @@ -1,5 +1,5 @@ """A Llama checkpoint's teacher-forced logits in float64 on the CPU, rounded once to bf16: the exact reference.""" -# usage: python3 tools/zig/llama_f64_reference.py MODEL_DIR tokens.json out.bin (raw bf16 [positions, vocab]) +# usage: llama_f64_reference.py MODEL_DIR tokens.json out.bin [--bf16-points]: raw bf16 [positions, vocab] import json import sys import time @@ -16,7 +16,8 @@ def bf16_once(x: np.ndarray) -> np.ndarray: return (y.view(np.uint32) >> 16).astype(np.uint16) -def main(model_dir: str, tokens: str, out: str) -> None: +def main(model_dir: str, tokens: str, out: str, points: bool = False) -> None: + """`points`: round to bf16 at every one of the reference's rounding points, each operation exact before it.""" torch.set_num_threads(24) cfg = json.load(open(f"{model_dir}/config.json")) index = json.load(open(f"{model_dir}/model.safetensors.index.json"))["weight_map"] @@ -34,7 +35,7 @@ def tensor(name): with safe_open(f"{model_dir}/{index[name]}", framework="pt") as f: return f.get_tensor(name) - def w(name): + def w(name, mlx=False): if not name.endswith(".weight") or name[: -len("weight")] + "scales" not in index: return tensor(name).to(f64) # mlx_lm's layout: u32 words of `bits`-bit codes, lowest bits first; a scale and a bias a group of codes @@ -44,10 +45,16 @@ def w(name): q = ((words.unsqueeze(-1) >> shifts) & ((1 << bits) - 1)).reshape(words.shape[0], -1).to(f64) scales = tensor(name[: -len("weight")] + "scales").to(f64).repeat_interleave(group, -1) biases = tensor(name[: -len("weight")] + "biases").to(f64).repeat_interleave(group, -1) + if mlx: # mlx's dequantize, as an affine embedding's rows are read: bf16(bf16(scale * q) + bias) + b16 = lambda t: t.to(torch.bfloat16).to(f64) + return b16(b16(scales * q) + biases) return scales * q + biases + def r(x): # a rounding point: bf16 when `points`, else exact + return x.to(torch.bfloat16).to(f64) if points else x + def norm(x, g): - return g * (x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)) + return r(g * r(x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps))) def rotate(x): return torch.cat((-x[..., dim // 2 :], x[..., : dim // 2]), -1) @@ -59,20 +66,24 @@ def rotate(x): inv = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=f64) / dim)) angle = torch.outer(torch.arange(n, dtype=f64), inv) cos, sin = torch.cat((angle, angle), -1).cos(), torch.cat((angle, angle), -1).sin() - x = w("model.embed_tokens.weight")[torch.tensor(ids)] + if points: # HF's fp32 angle, its cos and sin rounded to bf16 + a32 = torch.outer(torch.arange(n, dtype=torch.float32), 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))) + cos, sin = r(torch.cat((a32, a32), -1).cos().to(f64)), r(torch.cat((a32, a32), -1).sin().to(f64)) + x = w("model.embed_tokens.weight", mlx=points)[torch.tensor(ids)] mask = torch.full((n, n), float("-inf"), dtype=f64).triu(1) for i in range(layers): p = f"model.layers.{i}." h = norm(x, w(p + "input_layernorm.weight")) - q = (h @ w(p + "self_attn.q_proj.weight").T).view(n, heads, dim).transpose(0, 1) - k = (h @ w(p + "self_attn.k_proj.weight").T).view(n, kv, dim).transpose(0, 1) - v = (h @ w(p + "self_attn.v_proj.weight").T).view(n, kv, dim).transpose(0, 1) - q, k = q * cos + rotate(q) * sin, k * cos + rotate(k) * sin + q = r(h @ w(p + "self_attn.q_proj.weight").T).view(n, heads, dim).transpose(0, 1) + k = r(h @ w(p + "self_attn.k_proj.weight").T).view(n, kv, dim).transpose(0, 1) + v = r(h @ w(p + "self_attn.v_proj.weight").T).view(n, kv, dim).transpose(0, 1) + q, k = r(r(q * cos) + r(rotate(q) * sin)), r(r(k * cos) + r(rotate(k) * sin)) k, v = k.repeat_interleave(heads // kv, 0), v.repeat_interleave(heads // kv, 0) - a = torch.softmax(q @ k.transpose(1, 2) / dim**0.5 + mask, -1) @ v - x = x + a.transpose(0, 1).reshape(n, heads * dim) @ w(p + "self_attn.o_proj.weight").T + a = r(torch.softmax(q @ k.transpose(1, 2) / dim**0.5 + mask, -1) @ v) + x = r(x + r(a.transpose(0, 1).reshape(n, heads * dim) @ w(p + "self_attn.o_proj.weight").T)) h = norm(x, w(p + "post_attention_layernorm.weight")) - x = x + (torch.nn.functional.silu(h @ w(p + "mlp.gate_proj.weight").T) * (h @ w(p + "mlp.up_proj.weight").T)) @ w(p + "mlp.down_proj.weight").T + act = r(r(torch.nn.functional.silu(r(h @ w(p + "mlp.gate_proj.weight").T))) * r(h @ w(p + "mlp.up_proj.weight").T)) + x = r(x + r(act @ w(p + "mlp.down_proj.weight").T)) print(f"layer {i + 1} of {layers}, {time.time() - t0:.0f} s", flush=True) h = norm(x, w("model.norm.weight")) head = w("lm_head.weight") @@ -83,4 +94,4 @@ def rotate(x): if __name__ == "__main__": - main(*sys.argv[1:4]) + main(*sys.argv[1:4], points="--bf16-points" in sys.argv[4:])