diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 7fa49b6419..98fc94513d 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -150,6 +150,10 @@ finite (its id word carries bit 31, its value word a NaN) and takes the lower id rank's word into the id sentinel, so both Macs refuse a non-finite round where one Mac does. Six controls edit the source: a stage reading the peer's half, a gather writing its own half twice, and an argmax or pick tie to the higher id must go red at `shapes`, and an argmax or pick that does not flag at `nonfinite`. +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/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..f99a084331 --- /dev/null +++ b/docs/recipes/nanbeige-4.3-33b-metal.md @@ -0,0 +1,228 @@ +# 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 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 --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 \ + --context 131072 --parallel 2 --prompt-cache-gib 8 +# with Nanbeige4.1-3B drafting from its bf16 release (the same replies, faster: see Drafter below) +hf download Nanbeige/Nanbeige4.1-3B --revision 9c4555c37921c0982af8fffeee1e0a6e6b953c4f --local-dir models/nanbeige4.1-3b +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 --drafter models/nanbeige4.1-3b +``` + +`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. +- `--drafter`: a smaller Llama on the same tokenizer (`drafter.zig`) as the lane core's priced copy proposer, its + drafts verified by the model's windows like any copy, so replies keep the model's bits. + +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) + +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. + +`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: + +| 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. + +## Drafter: Nanbeige4.1-3B + +[`Nanbeige/Nanbeige4.1-3B`](https://huggingface.co/Nanbeige/Nanbeige4.1-3B) (revision +`9c4555c37921c0982af8fffeee1e0a6e6b953c4f`, Apache-2.0; model_type `llama`, 32 layers, hidden 2,560, 20 query and 4 key +heads of 128) shares the 33B's tokenizer: the same 166,100 ids, merges and added tokens, the same bos and eos (its +`vocab_size` pads to 166,144, the 33B's to 166,912). The family admits it as a drafter only, served from its bf16 +release as published (7.33 GiB), so the plain `--drafter ` serves it: `--drafter-bits` (which picks Qwen's +DFlash width) is not read for Llama, as for Gemma 4, and quantized drafter checkpoints are refused by name. + +- **The proposer.** `--drafter` loads the smaller Llama on the served model's queue and registers it as the lane + host's copy proposer (`LaneHost.proposers`), priced, so a round takes its drafts as given. The model's window + verifies them row by row and keeps the prefix its own draws agree with, as it keeps any copy draft: replies are the + model's whatever the drafter says. +- **One budget.** The streams' caches and the drafter's (64 KiB a token) share one ledger (`budget.Ledger`). A + stream's admission (`budget.admitStream`) checks its peak first (its caches and its drafter cache's growth, the old + state's copy included), trims the kept prompt states to what that peak leaves, and only then allocates; a freed + drafter cache keeps its tokens for the next request sharing its prefix, counted, and gives way first. +- **Drafting.** Before each round the drafter's cache is synced to the stream's context by their longest common + prefix, the new tokens are fed, and the drafts are greedy steps chained on the GPU. The depth comes from measured + rates: the stream's decayed acceptance, a chained step's time (8.5-8.9 ms) and a line fitted to the round's time by + its window's rows. With several streams in a round, each stream drafts the depth with the most tokens a millisecond + for the whole round (`pricing.depthShared`), nothing only when drafting costs the round more than it lands. + +Checked with `tf-llama-run <33B> draft ,`, which decodes greedily through the drafter's windows: +on eight chat prompts (128 tokens, thinking on and off) and after a 1,500-token prompt (256 tokens), the drafted +tokens equal the plain greedy tokens exactly (28 Oct 2026, one model process on the Mac): + +| prompt | plain | drafted | tokens a round | drafts kept | +|---|---:|---:|---:|---:| +| Fibonacci function, thinking off | 15.8 tok/s | 40.7 | 4.54 | 99 of 127 | +| Fibonacci function, thinking on | 15.8 | 22.0 | 1.81 | 57 of 142 | +| why the sky is blue, thinking off | 15.8 | 26.8 | 2.44 | 75 of 151 | +| why the sky is blue, thinking on | 15.9 | 22.8 | 1.88 | 60 of 125 | +| Hangzhou in two sentences (Chinese), thinking off | 15.9 | 21.1 | 1.73 | 54 of 137 | +| Hangzhou in two sentences (Chinese), thinking on | 15.9 | 24.1 | 2.02 | 64 of 130 | +| 9.11 or 9.8, thinking off | 15.8 | 31.1 | 3.31 | 90 of 167 | +| 9.11 or 9.8, thinking on | 15.8 | 25.9 | 2.42 | 75 of 158 | +| 1,500 tokens of Python's json docs, 256 tokens | 16.1 | 27.8 | 2.68 | 160 of 332 | + +Served (`--context 32768 --parallel 2 --prompt-cache-gib 8 --drafter ... --drafter-bits 0`, fresh servers): eight new +prompts sent two at a time (greedy, 256 tokens, thinking off) decode together at 30.1-40.9 tokens/s (36.1 on average) +against 26.6-30.1 (29.0) without `--drafter`, every reply's token ids equal to the undrafted one. Streams joining, +leaving, queued, refused (HTTP 400) and cancelled each return the token ids of their plain solo run, and a reply alone +after a cancel is unchanged; in rounds shared from the first token to the last, 199 drafts were offered and 125 kept. +The drafter's cache is fed the prompt on the first round, after the model's own prompt pass, so the first token keeps +its time. + +## Not done + +No quality suite was run. The family has no draft head of its own: drafts come from `--drafter`'s smaller Llama or +the lane core's copies. Quantized drafters are refused by name: the drafter is served from its bf16 release. +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 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/tools/zig/llama_f64_reference.py b/tools/zig/llama_f64_reference.py new file mode 100644 index 0000000000..360d1210c7 --- /dev/null +++ b/tools/zig/llama_f64_reference.py @@ -0,0 +1,97 @@ +"""A Llama checkpoint's teacher-forced logits in float64 on the CPU, rounded once to bf16: the exact reference.""" +# usage: llama_f64_reference.py MODEL_DIR tokens.json out.bin [--bf16-points]: 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, 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"] + 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, 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 + 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) + 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 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) + + 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() + 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 = 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 = 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")) + 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") + 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], points="--bf16-points" in sys.argv[4:]) diff --git a/tools/zig/llama_teacher_judge.py b/tools/zig/llama_teacher_judge.py new file mode 100644 index 0000000000..ff94cfb41b --- /dev/null +++ b/tools/zig/llama_teacher_judge.py @@ -0,0 +1,84 @@ +"""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 +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 (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 + 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)): + 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)}) + 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) + 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") + 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..1b1ba8128d --- /dev/null +++ b/tools/zig/llama_teacher_judge_test.py @@ -0,0 +1,113 @@ +"""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, top2 + +try: + import numpy +except ImportError: + numpy = None + + +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_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)) + + +if __name__ == "__main__": + unittest.main() diff --git a/zig/build/metal_targets.zig b/zig/build/metal_targets.zig index efa248606d..907830a28f 100644 --- a/zig/build/metal_targets.zig +++ b/zig/build/metal_targets.zig @@ -329,6 +329,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 6e06cb42f9..aa15844f4b 100644 --- a/zig/kernels/metal/sources.zig +++ b/zig/kernels/metal/sources.zig @@ -180,6 +180,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..6c465fc4ca --- /dev/null +++ b/zig/src/families/llama/backend.zig @@ -0,0 +1,440 @@ +//! 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 budget = @import("budget.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 }; + +/// The drafter's slots in the same ledger: a stream's slot growth (with its copy), made, and kept slots shed to fit. +pub const Slots = struct { + ctx: *anyopaque, + growth: *const fn (ctx: *anyopaque, stream: *lanes.Stream, capacity: usize) u64, + prepare: *const fn (ctx: *anyopaque, stream: *lanes.Stream, capacity: usize) anyerror!void, + shed: *const fn (ctx: *anyopaque, bytes: u64) 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, + ledger: budget.Ledger, // every cache's bytes against the one budget: the streams' and the drafter's slots + room: ?Room = null, // the prompt cache, which takes what the ledger leaves + slots: ?Slots = null, // --drafter's slots, admitted with their streams + 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), .ledger = .{ .limit = options.budget } }; + 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); + // the stream's peak (caches, drafter growth) fits before anything is made; kept states trim first + const Admit = struct { + m: *Metal, + stream: *lanes.Stream, + capacity: usize, + ids: []const u32, + pub fn growth(a: *const @This()) u64 { + return if (a.m.slots) |d| d.growth(d.ctx, a.stream, a.capacity) else 0; + } + pub fn shed(a: *const @This(), bytes: u64) void { + if (a.m.slots) |d| d.shed(d.ctx, bytes); + } + pub fn restoreBytes(a: *const @This()) u64 { + const saved = a.stream.reuse.saved orelse return 0; + const snap: *const snapshot.Snap = @ptrCast(@alignCast(saved)); + return if (snap.freed) 0 else snap.bytes; + } + pub fn dropRestore(a: *const @This()) void { // a state the stream cannot hold beside its peak: freed now, no restore + const snap: *snapshot.Snap = @ptrCast(@alignCast(a.stream.reuse.saved.?)); + snap.free(); + } + pub fn trim(a: *const @This(), bytes: u64) void { + if (a.m.room) |r| r.fit(r.ctx, bytes, a.stream.reuse.saved, a.ids); + } + pub fn prepare(a: *const @This()) !void { + if (a.m.slots) |d| try d.prepare(d.ctx, a.stream, a.capacity); + } + }; + const admit: Admit = .{ .m = self, .stream = stream, .capacity = capacity, .ids = ids }; + budget.admitStream(&self.ledger, need, &admit) catch |err| { + if (err == error.LaneMemory) 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 + admit.growth()), gib(self.ledger.held), gib(self.ledger.limit) }); + return err; // LaneMemory: the lane host queues the stream until another finishes, or fails it alone + }; + errdefer self.ledger.give(need); + 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); + // the kept states may use what the stream's growth copy and the slot's old cache gave back + if (self.room) |r| r.fit(r.ctx, self.ledger.left(), stream.reuse.saved, ids); + 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.ledger.give(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.ledger.left(), 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..9ad889d512 --- /dev/null +++ b/zig/src/families/llama/budget.zig @@ -0,0 +1,204 @@ +//! 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; +} + +/// One budget for every cache a load holds: each stream's, the drafter's slots (taken or kept) and their growth copies. +pub const Ledger = struct { + limit: u64, + held: u64 = 0, + + pub fn fits(l: *const Ledger, bytes: u64) bool { + return l.held + bytes <= l.limit; + } + + pub fn take(l: *Ledger, bytes: u64) void { + std.debug.assert(l.fits(bytes)); + l.held += bytes; + } + + pub fn give(l: *Ledger, bytes: u64) void { + l.held -= bytes; + } + + pub fn left(l: *const Ledger) u64 { + return l.limit - l.held; + } +}; + +/// A stream's admission inside the limit: peak checked, kept states (and an unfitting restore) trimmed, then made. +pub fn admitStream(l: *Ledger, need: u64, h: anytype) !void { + var peak = h.growth(); + if (!l.fits(need + peak)) { + h.shed(need + peak); + peak = h.growth(); + } + if (!l.fits(need + peak)) return error.LaneMemory; + const room = l.left() - need - peak; + if (h.restoreBytes() > room) h.dropRestore(); + h.trim(room); + l.take(need); + errdefer l.give(need); + try h.prepare(); +} + +/// A drafter's caches beside a stream's: `streams` of them, each the window and a decode window of rows more. +pub const Drafted = struct { c: *const cfg.Config, streams: u32 }; + +/// 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 { + return longestWith(c, null, bytes, most); +} + +/// As `longest`, with the drafter's caches for every stream taken from the same bytes. +pub fn longestWith(c: *const cfg.Config, drafted: ?Drafted, bytes: u64, most: usize) usize { + if (needed(c, drafted, 1) > bytes) return 0; + var lo: usize = 1; + var hi: usize = most; + while (lo < hi) { + const mid = lo + (hi - lo + 1) / 2; + if (needed(c, drafted, mid) <= bytes) lo = mid else hi = mid - 1; + } + return lo; +} + +/// Bytes a window of `n` tokens takes: one stream's caches a slot longer, and the drafter's caches when it drafts. +pub fn needed(c: *const cfg.Config, drafted: ?Drafted, n: usize) u64 { + var total: u64 = st.cacheBytes(c, n + 1); + if (drafted) |d| total += @as(u64, @max(d.streams, 1)) * st.cacheBytes(d.c, n + st.window_rows); + return total; +} + +test "llama budget: two drafted requests that each fit alone are refused together, and pass in turn" { + const gpa = std.testing.allocator; + const gib: u64 = 1 << 30; + const c = try cfg.parse(gpa, @embedFile("fixtures/config.json")); + const dc = try cfg.parse(gpa, @embedFile("fixtures/drafter_config.json")); + // a 100 GiB budget, a 262,144-token context and two streams: startup admits the window + const drafted: Drafted = .{ .c = &dc, .streams = 2 }; + try std.testing.expect(longestWith(&c, drafted, 100 * gib, 262144) == 262144); + // a 240,000-token request: its caches and its drafter slot's at the same capacity (backend.prefillFn's need) + const capacity = 240_000 + 1; + const one = st.cacheBytes(&c, capacity) + st.cacheBytes(&dc, capacity + st.window_rows); + var l: Ledger = .{ .limit = 100 * gib }; + try std.testing.expect(l.fits(one)); + l.take(one); + try std.testing.expect(!l.fits(one)); // the second waits: together they would pass the budget by 24.5 GiB + try std.testing.expect(2 * one - l.limit > 24 * gib); + // the first ends: its stream's caches are given back, its drafter slot kept (and still counted) + l.give(st.cacheBytes(&c, capacity)); + try std.testing.expect(l.fits(one)); +} + +/// A fake prompt store, drafter slot and allocation recorder: the most the ledger and the kept states held at once. +const Fake = struct { + const gib: u64 = 1 << 30; + l: *Ledger, + kept: u64, // kept prompt states but the restore entry + restore: u64, // the entry this stream resumes from (0: none) + old: u64, // the drafter slot's cache before it grows + grown: u64, // its cache after + copy: u64, // the old state's copy while it moves + fail: bool = false, + peak: u64 = 0, + + fn note(f: *Fake) void { + f.peak = @max(f.peak, f.l.held + f.kept + f.restore); + } + fn growth(f: *Fake) u64 { + return f.grown + f.copy; + } + fn shed(_: *Fake, _: u64) void {} + fn restoreBytes(f: *Fake) u64 { + return f.restore; + } + fn dropRestore(f: *Fake) void { + f.restore = 0; + } + fn trim(f: *Fake, bytes: u64) void { + f.kept = @min(f.kept, bytes -| f.restore); + f.note(); + } + fn prepare(f: *Fake) !void { + f.l.take(f.grown); + f.note(); + if (f.fail) { + f.l.give(f.grown); + return error.OutOfMemory; + } + f.l.take(f.copy); + f.note(); + f.l.give(f.copy); + f.l.give(f.old); + } +}; + +test "llama budget: admission trims kept prompt states before it allocates, so its peak stays inside the budget" { + const g = Fake.gib; + // 8 GiB, slot 1 and kept states 3 held; allocating the stream's 3 + 2.5 before the trim would reach 9.5 GiB + var l: Ledger = .{ .limit = 8 * g }; + l.take(g); + var f: Fake = .{ .l = &l, .kept = 2 * g, .restore = g, .old = g, .grown = 2 * g, .copy = g / 2 }; + try admitStream(&l, 3 * g, &f); + try std.testing.expect(f.peak <= l.limit); + try std.testing.expectEqual(g, f.restore); // the protected entry stayed + try std.testing.expectEqual(5 * g, l.held); // the stream's 3 and the grown slot's 2 + // a restore entry that cannot stay beside the stream's peak is dropped, not kept past the budget + var l2: Ledger = .{ .limit = 8 * g }; + l2.take(g); + var f2: Fake = .{ .l = &l2, .kept = g, .restore = 2 * g, .old = g, .grown = 2 * g, .copy = g / 2 }; + try admitStream(&l2, 3 * g, &f2); + try std.testing.expect(f2.peak <= l2.limit); + try std.testing.expectEqual(@as(u64, 0), f2.restore); +} + +test "llama budget: a failed drafter growth gives every byte back, and its peak stays inside the budget" { + const g = Fake.gib; + var l: Ledger = .{ .limit = 8 * g }; + l.take(g); + var f: Fake = .{ .l = &l, .kept = 2 * g, .restore = g, .old = g, .grown = 2 * g, .copy = g / 2, .fail = true }; + try std.testing.expectError(error.OutOfMemory, admitStream(&l, 3 * g, &f)); + try std.testing.expect(f.peak <= l.limit); + try std.testing.expectEqual(g, l.held); // the stream's caches and the growth given back; the old slot still held + // a peak that cannot fit at all waits before anything is trimmed or made + var f2: Fake = .{ .l = &l, .kept = 2 * g, .restore = g, .old = g, .grown = 6 * g, .copy = g }; + try std.testing.expectError(error.LaneMemory, admitStream(&l, 3 * g, &f2)); + try std.testing.expectEqual(@as(u64, 2 * g), f2.kept); + try std.testing.expectEqual(@as(u64, 0), f2.peak); +} + +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)); + // a drafter of the same geometry for two streams: its caches come out of the same bytes + const two: Drafted = .{ .c = &c, .streams = 2 }; + try std.testing.expectEqual(st.cacheBytes(&c, 4081) + 2 * st.cacheBytes(&c, 4096), needed(&c, two, 4080)); + try std.testing.expectEqual(@as(usize, 4080), longestWith(&c, two, needed(&c, two, 4080), 1 << 20)); // 4,081 needs a block more + try std.testing.expect(longestWith(&c, two, st.cacheBytes(&c, 4097), 1 << 20) < 1500); + try std.testing.expectEqual(@as(usize, 0), longestWith(&c, two, 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..a53ff4217c --- /dev/null +++ b/zig/src/families/llama/config.zig @@ -0,0 +1,327 @@ +//! 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; +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, + UnqualifiedLlamaCheckpoint, +} || std.json.ParseError(std.json.Scanner); + +/// What a qualified checkpoint was checked for: served on its own, or drafting for a served one. +pub const Role = enum { served, drafter }; + +/// A checkpoint checked on the device, by the geometry its config names. +pub const Qualified = struct { + name: []const u8, + role: Role, + 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", .role = .served, .hidden = 7168, .intermediate = 22016, .layers = 52, .heads = 56, .kv_heads = 8, .vocab = 166912, .theta = 120000000, .context = 524288 }, + // its drafts are checked by the served model's windows, its own logits not against HF + .{ .name = "Nanbeige4.1-3B", .role = .drafter, .hidden = 2560, .intermediate = 10496, .layers = 32, .heads = 20, .kv_heads = 4, .vocab = 166144, .theta = 70000000, .context = 262144 }, +}; + +/// 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, + 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); + 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 { + 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 headScale(); + } + + 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 "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 "Nanbeige4.1-3B's geometry is admitted as a drafter, not as a served model" { + const gpa = std.testing.allocator; + var bytes = try gpa.dupe(u8, fixture); + defer gpa.free(bytes); + const geometry = .{ + .{ "\"hidden_size\": 7168", "\"hidden_size\": 2560" }, + .{ "\"intermediate_size\": 22016", "\"intermediate_size\": 10496" }, + .{ "\"num_hidden_layers\": 52", "\"num_hidden_layers\": 32" }, + .{ "\"num_attention_heads\": 56", "\"num_attention_heads\": 20" }, + .{ "\"num_key_value_heads\": 8", "\"num_key_value_heads\": 4" }, + .{ "\"vocab_size\": 166912", "\"vocab_size\": 166144" }, + .{ "\"rope_theta\": 120000000", "\"rope_theta\": 70000000" }, + .{ "\"max_position_embeddings\": 524288", "\"max_position_embeddings\": 262144" }, + }; + inline for (geometry) |k| { + const next = try std.mem.replaceOwned(u8, gpa, bytes, k[0], k[1]); + try std.testing.expect(!std.mem.eql(u8, next, bytes)); + gpa.free(bytes); + bytes = next; + } + const q = try (try parse(gpa, bytes)).qualification(); + try std.testing.expectEqualStrings("Nanbeige4.1-3B", q.name); + try std.testing.expectEqual(Role.drafter, q.role); +} + +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/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/drafter.zig b/zig/src/families/llama/drafter.zig new file mode 100644 index 0000000000..fb7c160a89 --- /dev/null +++ b/zig/src/families/llama/drafter.zig @@ -0,0 +1,508 @@ +//! A smaller Llama drafting for the served one as the lane core's priced proposer; the target's windows verify it. +const std = @import("std"); +const mtl = @import("metal"); +const lanes = @import("lanes"); +const cfg = @import("config.zig"); +const st = @import("state.zig"); +const fwd = @import("forward.zig"); +const snapshot = @import("snapshot.zig"); +const budget = @import("budget.zig"); +const backend = @import("backend.zig"); +const Model = @import("model.zig").Model; +const Proposer = lanes.proposer.Proposer; +const pricing = @import("pricing.zig"); +const Fit = pricing.Fit; +const depth = pricing.depth; +const depthShared = pricing.depthShared; +const ema = pricing.ema; + +/// Up to `capacity` tokens for each of `streams` caches (one a request, kept for its prefix), in the target's ledger. +pub const Options = struct { capacity: usize, streams: u32, ledger: *budget.Ledger }; + +/// Token slots the chained steps draw into; a round takes at most window_rows of them. +const ring = 4096; +/// Drafts a round until the rates are measured. +const first_depth: usize = 6; +/// The acceptance prior (two drafts of three) and how much of the history a round keeps. +const prior_kept = 2.0; +const prior_tried = 3.0; +const decay = 0.9; + +/// A cache a stream synced to its context by longest common prefix; greedy steps chained on the GPU; a measured depth. +pub const Drafter = struct { + gpa: std.mem.Allocator, + model: *Model, // owned: freed with the drafter + options: Options, + scratch: st.Scratch, // chained steps and short deltas + prompt: st.Scratch, // prompt-sized deltas + tokens: mtl.Buffer, // u32 [ring] + next: u64 = 0, // the next free token slot + slots: []Slot, + follow: ?mtl.ResidencySet.Follow = null, + step_ms: ?f64 = null, // one chained step, wall (null: not measured yet) + rounds: Fit = .{}, // the target's round by its window's rows, from a proposal's return to its outcome + + /// A stream's cache and what it holds; a freed slot keeps both for the next request sharing its prefix. + const Slot = struct { + cache: ?st.Cache = null, // made or grown when its stream is admitted, to that stream's capacity + seen: std.ArrayList(u32) = .empty, // the tokens whose keys and values the cache holds, in order + next: ?u32 = null, // the token drawn after `seen` (the last fed row's argmax) + taken: bool = false, + }; + + pub fn init(gpa: std.mem.Allocator, model: *Model, options: Options) !*Drafter { + const d = try gpa.create(Drafter); + errdefer gpa.destroy(d); + 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(); + const slots = try gpa.alloc(Slot, @max(options.streams, 1)); + for (slots) |*s| s.* = .{}; + d.* = .{ .gpa = gpa, .model = model, .options = options, .scratch = scratch, .prompt = prompt, .tokens = tokens, .slots = slots }; + return d; + } + + /// A drafter for `model`, warmed: it owns the model from the call on, freed with it if the warm-up fails. + pub fn start(gpa: std.mem.Allocator, model: *Model, options: Options) !*Drafter { + const d = init(gpa, model, options) catch |err| { + model.deinit(); + return err; + }; + return started(Drafter, d); + } + + pub fn deinit(d: *Drafter) void { + for (d.slots) |*s| { + d.drop(s); + s.seen.deinit(d.gpa); + } + d.gpa.free(d.slots); + d.tokens.deinit(); + d.scratch.deinit(); + d.prompt.deinit(); + d.model.deinit(); + d.gpa.destroy(d); + } + + /// Bytes the drafter holds beside its weights and caches: its scratches and token slots. + pub fn fixedBytes(c: *const cfg.Config) u64 { + return st.scratchBytes(c, st.batch_rows) + st.scratchBytes(c, st.prompt_rows) + ring * 4; + } + + pub fn setResident(d: *Drafter, hook: ?mtl.ResidencySet.Follow) void { + if (d.follow) |f| { + mtl.resident.fields.apply(&d.scratch, f, false); + mtl.resident.fields.apply(&d.prompt, f, false); + } + d.follow = hook; + if (hook) |f| { + mtl.resident.fields.apply(&d.scratch, f, true); + mtl.resident.fields.apply(&d.prompt, f, true); + } + for (d.slots) |*s| if (s.cache) |*c| c.setResident(hook); + } + + /// A prompt chunk and a chained step once, so the first request skips the GPU's first touch. + pub fn warm(d: *Drafter) !void { + var cache = try st.Cache.init(d.model.device, &d.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) % d.model.config.vocab); + try d.feed(&cache, &ids, null); + var flights: [1]mtl.CommandBuffer = undefined; + try d.step(&cache, ids[0], &flights[0]); + try land(&flights); + } + + /// The lane host's proposer factory: a request takes the free slot sharing the longest prefix with its prompt. + pub fn make(ctx: *anyopaque, gpa: std.mem.Allocator, prompt: []const u32, eos: []const u32) anyerror!Proposer { + const d: *Drafter = @ptrCast(@alignCast(ctx)); + var best: ?usize = null; + var best_len: usize = 0; + for (d.slots, 0..) |*s, i| { + if (s.taken) continue; + const shared = commonPrefix(s.seen.items, prompt); + if (best == null or shared > best_len) { + best = i; + best_len = shared; + } + } + const i = best orelse return error.LlamaDrafterBusy; + const slot = &d.slots[i]; + if (prompt.len + 1 > d.options.capacity) return error.LlamaDrafterPromptTooLong; + const r = try gpa.create(Request); + r.* = .{ .d = d, .gpa = gpa, .slot = i, .eos = eos }; + slot.taken = true; + return r.proposer(); + } + + pub fn free(ctx: *anyopaque, p: Proposer) void { + const d: *Drafter = @ptrCast(@alignCast(ctx)); + const r: *Request = @ptrCast(@alignCast(p.ptr)); + d.slots[r.slot].taken = false; + if (r.rounds > 0) std.log.debug("llama drafter: {d} rounds, {d} drafts offered, {d} kept, {d:.2} tokens a round", .{ r.rounds, r.offered, r.landed, @as(f64, @floatFromInt(r.rounds + r.landed)) / @as(f64, @floatFromInt(r.rounds)) }); + r.out.deinit(r.gpa); + r.gpa.destroy(r); + } + + /// The streams holding a slot: more than one shares each round, every stream's step waiting on each one's drafts. + fn sharers(d: *const Drafter) usize { + var n: usize = 0; + for (d.slots) |s| n += @intFromBool(s.taken); + return n; + } + + /// The target's admission hooks: a stream's slot grows with it, in its ledger. + pub fn slotsHook(d: *Drafter) backend.Slots { + return .{ .ctx = d, .growth = growthFn, .prepare = prepareFn, .shed = shedFn }; + } + + /// The slot `proposer` drafts from (null: something else's proposer, or none). + fn slotFor(d: *Drafter, proposer: ?Proposer) ?*Slot { + const p = proposer orelse return null; + if (p.vtable != &Request.vtable) return null; + const r: *Request = @ptrCast(@alignCast(p.ptr)); + return if (r.d == d) &d.slots[r.slot] else null; + } + + /// Rows a stream of `capacity` tokens drafts within: its context, a window of drafts and the draw after them. + fn rowsFor(d: *const Drafter, capacity: usize) usize { + return @min(capacity + st.window_rows, d.options.capacity); + } + + fn growthFn(ctx: *anyopaque, stream: *lanes.Stream, capacity: usize) u64 { + const d: *Drafter = @ptrCast(@alignCast(ctx)); + const slot = d.slotFor(stream.proposer) orelse return 0; + return growth(&d.model.config, if (slot.cache) |c| .{ .capacity = c.capacity, .len = c.len } else null, d.rowsFor(capacity)); + } + + fn prepareFn(ctx: *anyopaque, stream: *lanes.Stream, capacity: usize) anyerror!void { + const d: *Drafter = @ptrCast(@alignCast(ctx)); + try d.prepare(stream.proposer, capacity); + } + + /// The proposer's slot grown for a request of `capacity` tokens, as the target's admission grows a stream's. + pub fn prepare(d: *Drafter, proposer: ?Proposer, capacity: usize) !void { + const slot = d.slotFor(proposer) orelse return; + try d.grow(slot, d.rowsFor(capacity)); + } + + fn shedFn(ctx: *anyopaque, bytes: u64) void { + const d: *Drafter = @ptrCast(@alignCast(ctx)); + for (d.slots) |*s| { + if (d.options.ledger.fits(bytes)) return; + if (!s.taken) d.drop(s); + } + } + + /// A slot's cache and what it knew, freed and given back to the ledger. + fn drop(d: *Drafter, slot: *Slot) void { + if (slot.cache) |*c| { + d.options.ledger.give(st.cacheBytes(&d.model.config, c.capacity)); + c.deinit(); + } + slot.cache = null; + slot.seen.clearRetainingCapacity(); + slot.next = null; + } + + /// The slot's cache made or grown to `rows` (its state copied), every byte taken from the ledger as it is held. + fn grow(d: *Drafter, slot: *Slot, rows: usize) !void { + if (slot.cache) |c| if (c.capacity >= rows) return; + const c = &d.model.config; + const ledger = d.options.ledger; + const made = st.cacheBytes(c, rows); + ledger.take(made); + errdefer ledger.give(made); + var bigger = try st.Cache.init(d.model.device, c, rows); + errdefer bigger.deinit(); + if (slot.cache) |*old| { + const copy = snapshot.bytes(c, old.len); + ledger.take(copy); + defer ledger.give(copy); + var snap = try snapshot.save(d.model.device, d.model.queue, c, old); + defer snap.free(); + try snapshot.restore(d.model.queue, &snap, &bigger); + ledger.give(st.cacheBytes(c, old.capacity)); + old.deinit(); + } + bigger.setResident(d.follow); + slot.cache = bigger; + } + + /// `ids` at the cache's end in prompt chunks, waited on, the last chunk's row drawn into `draw` when given. + fn feed(d: *Drafter, cache: *st.Cache, ids: []const u32, draw: ?*mtl.CommandBuffer) !void { + var at: usize = 0; + while (at < ids.len) { + const end = @min(ids.len, at + st.prompt_rows); + const s = if (end - at > st.batch_rows) &d.prompt else &d.scratch; + const last = end == ids.len; + if (last and draw != null) { + const out = d.take(); + try fwd.run(d.model, s, cache, ids[at..end], .{ .draw = .{ .how = .greedy, .out = d.slotOf(out) }, .flight = draw }); + } else try fwd.run(d.model, s, cache, ids[at..end], .{}); + cache.len += end - at; + at = end; + } + } + + /// One row from the host into a fresh slot, its argmax drawn into the slot after it; committed into `flight`. + fn step(d: *Drafter, cache: *st.Cache, row: u32, flight: *mtl.CommandBuffer) !void { + const in = d.take(); + d.tokens.slice(u32, ring)[in % ring] = row; + const out = d.take(); + try fwd.run(d.model, &d.scratch, cache, &.{0}, .{ .ids = d.slotOf(in), .draw = .{ .how = .greedy, .out = d.slotOf(out) }, .flight = flight }); + cache.len += 1; + } + + /// A step whose row is a drawn slot (the previous step's), the next drawn into a fresh slot. + fn chain(d: *Drafter, cache: *st.Cache, in: u64, flight: *mtl.CommandBuffer) !void { + const out = d.take(); + try fwd.run(d.model, &d.scratch, cache, &.{0}, .{ .ids = d.slotOf(in), .draw = .{ .how = .greedy, .out = d.slotOf(out) }, .flight = flight }); + cache.len += 1; + } + + fn take(d: *Drafter) u64 { + const h = d.next; + d.next += 1; + return h; + } + + fn slotOf(d: *Drafter, h: u64) fwd.B { + return .{ .buf = d.tokens, .off = @intCast(h % ring * 4) }; + } + + fn token(d: *Drafter, h: u64) u32 { + return d.tokens.slice(u32, ring)[h % ring]; + } +}; + +/// What a slot's cache holds now: its capacity and the rows written. +pub const Held = struct { capacity: usize, len: usize }; + +/// Bytes a slot holding `have` takes to grow to `rows`: the new cache, and the old state's copy while it moves. +pub fn growth(c: *const cfg.Config, have: ?Held, rows: usize) u64 { + if (have) |h| if (h.capacity >= rows) return 0; + return st.cacheBytes(c, rows) + if (have) |h| snapshot.bytes(c, h.len) else 0; +} + +/// `d` warmed, or freed when its warm-up fails: nothing made before the warm-up outlives a failed start. +fn started(comptime D: type, d: *D) !*D { + errdefer d.deinit(); + try d.warm(); + return d; +} + +/// Wait for every step in flight, oldest first, and release them. +fn land(flights: []mtl.CommandBuffer) !void { + var failed = false; + for (flights) |cb| { + cb.wait(); + if (cb.failure()) |msg| { + std.log.err("llama drafter step: {s}", .{msg}); + failed = true; + } + cb.release(); + } + if (failed) return error.GpuFailed; +} + +fn commonPrefix(a: []const u32, b: []const u32) usize { + const n = @min(a.len, b.len); + var i: usize = 0; + while (i < n and a[i] == b[i]) i += 1; + return i; +} + +fn now() f64 { + var ts: std.c.timespec = undefined; + _ = std.c.clock_gettime(.MONOTONIC, &ts); + return @as(f64, @floatFromInt(ts.sec)) * 1e3 + @as(f64, @floatFromInt(ts.nsec)) * 1e-6; +} + +/// One stream's proposer: its slot, its acceptance and what the last round asked. +pub const Request = struct { + d: *Drafter, + gpa: std.mem.Allocator, + slot: usize, + eos: []const u32, + out: std.ArrayList(u32) = .empty, + asked: ?struct { len: usize, max_draft: i64 } = null, + returned: ?f64 = null, // when the last proposal returned (ms), for the round time + kept: f64 = prior_kept, // decayed drafts landed + tried: f64 = prior_tried, // decayed drafts judged: the landed ones and the first to break + rounds: u64 = 0, + offered: u64 = 0, + landed: u64 = 0, + unfit: bool = false, // a context token past the drafter's vocabulary: no proposals for this request + alone: bool = true, // the last proposal's round had this stream alone: its time fits the rounds' line + + pub fn proposer(r: *Request) Proposer { + return .{ .ptr = r, .vtable = &vtable }; + } + + const vtable: Proposer.VTable = .{ .propose = proposeFn, .last_match = lastMatch, .priced = true, .round = roundFn }; + + fn self(p: *anyopaque) *Request { + return @ptrCast(@alignCast(p)); + } + + fn proposeFn(p: *anyopaque, context: []const u32, max_draft: i64) anyerror![]const u32 { + return self(p).propose(context, max_draft); + } + + fn lastMatch(_: *anyopaque) i64 { + return 0; + } + + fn roundFn(p: *anyopaque, rows: u32, kept: u32) void { + self(p).round(rows, kept); + } + + /// Drafts a round at the measured rates (first_depth before), within the ask and window; shared rounds priced. + pub fn want(r: *const Request, max_draft: i64) usize { + const most: usize = @intCast(@min(max_draft, @as(i64, st.window_rows - 1))); + const step = r.d.step_ms orelse return @min(first_depth, most); + if (r.d.rounds.samples == 0) return @min(first_depth, most); + const streams = r.d.sharers(); + if (streams > 1) return depthShared(r.kept / r.tried, r.d.rounds, step, most, streams); + return depth(r.kept / r.tried, r.d.rounds, step, most); + } + + /// The drafter's continuation of `context`, up to `max_draft` tokens, cut after an eos; borrowed until the next. + pub fn propose(r: *Request, context: []const u32, max_draft: i64) ![]const u32 { + if (r.asked) |a| if (a.len == context.len and a.max_draft == max_draft) return r.out.items; + r.out.clearRetainingCapacity(); + r.asked = .{ .len = context.len, .max_draft = max_draft }; + r.returned = null; // an empty answer times no round + if (max_draft <= 0 or r.unfit or context.len == 0) return r.out.items; + const d = r.d; + const slot = &d.slots[r.slot]; + const n = r.want(max_draft); + if (n == 0) return r.out.items; // the shared round's drafts cost more than they land + r.alone = d.sharers() <= 1; + // the cache keeps the context's prefix it agrees with (an earlier round's drafts past the context included) + const shared = commonPrefix(slot.seen.items, context); + if (shared < context.len and shared < slot.seen.items.len) { + slot.seen.shrinkRetainingCapacity(shared); + slot.cache.?.len = shared; + slot.next = null; + } + // room for the context, the drafts and the draw after them: the slot grew to its stream's capacity at admission + const cache = if (slot.cache) |*c| c else return r.out.items; + if (@max(context.len, slot.seen.items.len) + n + 1 > cache.capacity) return r.out.items; + var flights: [st.window_rows]mtl.CommandBuffer = undefined; + var draws: [st.window_rows]u64 = undefined; // each pass's drawn slot, in order + var count: usize = 0; + if (slot.seen.items.len < context.len) { + // the rest of the context, its last row drawn; waited on, so the chain's time is its own + const delta = context[slot.seen.items.len..]; + for (delta) |t| if (t >= d.model.config.vocab) { + r.unfit = true; + return r.out.items; + }; + if (slot.seen.items.len + delta.len + n > cache.capacity) return r.out.items; + try d.feed(cache, delta, &flights[0]); + draws[0] = d.next - 1; + count = 1; + try slot.seen.appendSlice(d.gpa, delta); + try land(flights[0..1]); + } else if (slot.next == null) return error.LlamaDrafterUndrawn; + const fed = slot.seen.items.len - context.len; // an earlier round's drafts the cache holds past the context + var more = n -| (fed + 1); // chained steps past the drawn token + if (slot.seen.items.len + more > cache.capacity) more = 0; + if (more > 0) { + const t0 = now(); + const first = count; + if (count == 0) { + // the drawn token is the host's copy: its step starts the chain + try d.step(cache, slot.next.?, &flights[0]); + draws[0] = d.next - 1; + try slot.seen.append(d.gpa, slot.next.?); + count = 1; + more -= 1; + } + for (0..more) |_| { + try d.chain(cache, draws[count - 1], &flights[count]); + draws[count] = d.next - 1; + count += 1; + } + try land(flights[first..count]); + d.step_ms = ema(d.step_ms, (now() - t0) / @as(f64, @floatFromInt(count - first))); + } + if (count > 0) { + // every pass but the last fed its draw to the next; the last draw is the token after `seen` + for (draws[0 .. count - 1]) |h| try slot.seen.append(d.gpa, d.token(h)); + slot.next = d.token(draws[count - 1]); + } + const available = slot.seen.items.len - context.len; + try r.out.appendSlice(r.gpa, slot.seen.items[context.len..][0..@min(available, n)]); + if (r.out.items.len < n) try r.out.append(r.gpa, slot.next.?); + for (r.out.items, 0..) |t, i| if (std.mem.indexOfScalar(u32, r.eos, t) != null) { + r.out.shrinkRetainingCapacity(i + 1); + break; + }; + r.returned = now(); + return r.out.items; + } + + /// The round's outcome: `rows` the window's rows, `kept` its drafts landed (no ask this round: a plain step). + pub fn round(r: *Request, rows: u32, kept: u32) void { + r.rounds += 1; + defer r.asked = null; + if (r.asked == null or rows < 2) return; + if (r.returned) |t| if (r.alone) r.d.rounds.add(@floatFromInt(rows), now() - t); + const offered = rows - 1; + r.offered += offered; + r.landed += kept; + const k: f64 = @floatFromInt(kept); + r.kept = decay * r.kept + k; + r.tried = decay * r.tried + k + @as(f64, if (kept < offered) 1.0 else 0.0); + } +}; + +test "llama drafter: a warm-up that fails frees the drafter it was made for" { + const Fake = struct { + freed: *bool, + fn warm(_: *@This()) !void { + return error.WarmUpFailed; + } + fn deinit(f: *@This()) void { + f.freed.* = true; + } + }; + var freed = false; + var f: Fake = .{ .freed = &freed }; + try std.testing.expectError(error.WarmUpFailed, started(Fake, &f)); + try std.testing.expect(freed); +} + +test "llama drafter: a slot's growth counts the new cache and the old state's copy, nothing when it holds enough" { + const c = try cfg.parse(std.testing.allocator, @embedFile("fixtures/drafter_config.json")); + try std.testing.expectEqual(@as(u64, st.cacheBytes(&c, 9000)), growth(&c, null, 9000)); + try std.testing.expectEqual(@as(u64, st.cacheBytes(&c, 9000) + snapshot.bytes(&c, 3000)), growth(&c, .{ .capacity = 4096, .len = 3000 }, 9000)); + try std.testing.expectEqual(@as(u64, 0), growth(&c, .{ .capacity = 9000, .len = 3000 }, 9000)); +} + +test "llama drafter: the common prefix and the acceptance history" { + try std.testing.expectEqual(@as(usize, 3), commonPrefix(&.{ 1, 2, 3, 4 }, &.{ 1, 2, 3, 9, 9 })); + try std.testing.expectEqual(@as(usize, 2), commonPrefix(&.{ 1, 2 }, &.{ 1, 2, 3 })); + try std.testing.expectEqual(@as(usize, 0), commonPrefix(&.{}, &.{1})); + var d: Drafter = undefined; + var r: Request = .{ .d = &d, .gpa = std.testing.allocator, .slot = 0, .eos = &.{} }; + r.asked = .{ .len = 10, .max_draft = 15 }; + r.round(7, 6); // six drafts, all landed: no break counted + try std.testing.expectApproxEqAbs(decay * prior_kept + 6, r.kept, 1e-9); + try std.testing.expectApproxEqAbs(decay * prior_tried + 6, r.tried, 1e-9); + try std.testing.expect(r.asked == null); + r.asked = .{ .len = 17, .max_draft = 15 }; + r.round(7, 2); // two of six: the third broke + try std.testing.expectEqual(@as(u64, 12), r.offered); + try std.testing.expectEqual(@as(u64, 8), r.landed); + r.round(1, 0); // a plain step: no history + try std.testing.expectEqual(@as(u64, 3), r.rounds); + try std.testing.expectEqual(@as(u64, 12), r.offered); +} 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/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/fixtures/drafter_config.json b/zig/src/families/llama/fixtures/drafter_config.json new file mode 100644 index 0000000000..0b3dd5089f --- /dev/null +++ b/zig/src/families/llama/fixtures/drafter_config.json @@ -0,0 +1,33 @@ +{ + "architectures": [ + "LlamaForCausalLM" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "bos_token_id": 166100, + "dtype": "bfloat16", + "embd_pdrop": 0.0, + "eos_token_id": 166101, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 2560, + "initializer_range": 0.02, + "intermediate_size": 10496, + "max_position_embeddings": 262144, + "mlp_bias": false, + "model_type": "llama", + "num_attention_heads": 20, + "num_hidden_layers": 32, + "num_key_value_heads": 4, + "pad_token_id": 0, + "pretraining_tp": 1, + "resid_pdrop": 0.0, + "rms_norm_eps": 1e-05, + "rope_scaling": null, + "rope_theta": 70000000, + "tie_word_embeddings": false, + "torch_dtype": "bfloat16", + "transformers_version": "4.51.3", + "use_cache": true, + "vocab_size": 166144 +} 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..2b8b131310 --- /dev/null +++ b/zig/src/families/llama/llama.zig @@ -0,0 +1,25 @@ +//! Native Metal Llama (model_type llama): GQA attention and a SwiGLU MLP at HF's rounding points; a smaller one drafts. +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"); +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"); +pub const exact = @import("exact.zig"); +pub const drafter = @import("drafter.zig"); +pub const pricing = @import("pricing.zig"); + +test { + _ = exact; + _ = config; + _ = contracts; + _ = state; + _ = budget; + _ = snapshot; + _ = drafter; + _ = pricing; +} diff --git a/zig/src/families/llama/model.zig b/zig/src/families/llama/model.zig new file mode 100644 index 0000000000..86e2a4ea2d --- /dev/null +++ b/zig/src/families/llama/model.zig @@ -0,0 +1,125 @@ +//! 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, + /// The queue is another model's (a drafter's commands behind the served model's, in its residency set). + shared_queue: bool, + + pub fn load(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !*Model { + return loadOn(gpa, io, dir, null); + } + + /// As `load`, on `queue` when given (the model then never releases it). + pub fn loadOn(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, queue: ?mtl.Queue) !*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(); + // prompt chunks' GEMM and attention run on Metal 4 tensor units: an M5 or later + if (!m.device.tensorUnits()) return error.LlamaNeedsTensorUnits; + m.shared_queue = queue != null; + m.queue = queue orelse try m.device.queue(); + errdefer if (!m.shared_queue) 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(); + if (!m.shared_queue) m.queue.deinit(); + m.device.deinit(); + m.gpa.destroy(m); + } +}; diff --git a/zig/src/families/llama/pricing.zig b/zig/src/families/llama/pricing.zig new file mode 100644 index 0000000000..5cd5491c4f --- /dev/null +++ b/zig/src/families/llama/pricing.zig @@ -0,0 +1,112 @@ +//! What a drafted round is worth: rounds' times by their rows, and the depth landing the most tokens a millisecond. +const std = @import("std"); + +/// The newest sample's weight in the running times. +const rate = 0.1; + +pub fn ema(kept: ?f64, sample: f64) f64 { + return if (kept) |k| k + rate * (sample - k) else sample; +} + +/// A round's milliseconds by its window's rows, a line fitted to the recent rounds (a window costs more with rows). +pub const Fit = struct { + samples: u32 = 0, + x: f64 = 0, // rows + y: f64 = 0, // ms + xx: f64 = 0, + xy: f64 = 0, + + pub fn add(f: *Fit, rows: f64, ms: f64) void { + if (f.samples == 0) { + f.* = .{ .samples = 1, .x = rows, .y = ms, .xx = rows * rows, .xy = rows * ms }; + return; + } + f.samples += 1; + f.x = ema(f.x, rows); + f.y = ema(f.y, ms); + f.xx = ema(f.xx, rows * rows); + f.xy = ema(f.xy, rows * ms); + } + + /// Milliseconds a window of `rows` takes: the fitted line, or the mean while the rows seen hardly vary. + pub fn at(f: Fit, rows: f64) f64 { + const spread = f.xx - f.x * f.x; + if (spread < 0.25) return f.y; + const slope = @max(0, (f.xy - f.x * f.y) / spread); + return @max(f.y + slope * (rows - f.x), 0.5 * f.y); + } +}; + +/// The depth with the most tokens a millisecond: a round lands 1 + p + ... + p^k tokens for k + 1 rows and k steps. +pub fn depth(p: f64, rounds: Fit, step_ms: f64, most: usize) usize { + var best: usize = 1; + var best_rate: f64 = 0; + var expect: f64 = 1; + var pk: f64 = 1; + for (1..most + 1) |k| { + pk *= p; + expect += pk; + const kf: f64 = @floatFromInt(k); + const r = expect / (rounds.at(kf + 1) + kf * step_ms); + if (r > best_rate) { + best_rate = r; + best = k; + } + } + return best; +} + +/// A shared round's depth for each of `streams` streams: most tokens a millisecond for the round, 0 if drafting costs. +pub fn depthShared(p: f64, rounds: Fit, step_ms: f64, most: usize, streams: usize) usize { + const s: f64 = @floatFromInt(streams); + var best: usize = 0; + var best_rate = s / rounds.at(s); + var expect: f64 = 1; + var pk: f64 = 1; + for (1..most + 1) |k| { + pk *= p; + expect += pk; + const kf: f64 = @floatFromInt(k); + const r = s * expect / (rounds.at(s * (kf + 1)) + s * kf * step_ms); + if (r > best_rate) { + best_rate = r; + best = k; + } + } + return best; +} + +test "llama pricing: the depth pays for its steps and rows, more at a higher acceptance, none past the ask" { + // flat 60 ms rounds and 3 ms steps: at p 0.7 drafting about 5-6 pays best; at 0.95 the window's 15; at 0.1 one + var flat: Fit = .{}; + flat.add(4, 60); + try std.testing.expect(depth(0.7, flat, 3, 15) >= 4 and depth(0.7, flat, 3, 15) <= 7); + try std.testing.expectEqual(@as(usize, 15), depth(0.95, flat, 3, 15)); + try std.testing.expectEqual(@as(usize, 1), depth(0.1, flat, 3, 15)); + try std.testing.expectEqual(@as(usize, 3), depth(0.95, flat, 3, 3)); + // dear steps pull the depth down + try std.testing.expect(depth(0.7, flat, 30, 15) <= 2); + // rounds that grow 5 ms a row (as the M5's windows do) pull it down too + var sloped: Fit = .{}; + for (0..40) |i| sloped.add(@floatFromInt(2 + i % 8), 50 + 5 * @as(f64, @floatFromInt(2 + i % 8))); + try std.testing.expectApproxEqAbs(@as(f64, 100), sloped.at(10), 3); + try std.testing.expect(depth(0.7, sloped, 3, 15) < depth(0.7, flat, 3, 15)); + try std.testing.expect(depth(0.95, sloped, 3, 15) >= 8); + // one row count seen: the mean stands in for the line + var one: Fit = .{}; + one.add(4, 60); + one.add(4, 70); + try std.testing.expectApproxEqAbs(one.y, one.at(12), 1e-9); +} + +test "llama pricing: a shared round drafts where it pays for every stream, and not where it costs more" { + // 60 ms rounds growing 5 ms a row, 3 ms steps: two streams at a high acceptance still draft + var sloped: Fit = .{}; + for (0..40) |i| sloped.add(@floatFromInt(2 + i % 8), 50 + 5 * @as(f64, @floatFromInt(2 + i % 8))); + try std.testing.expect(depthShared(0.9, sloped, 3, 15, 2) >= 2); + // a low acceptance: the round runs undrafted + try std.testing.expectEqual(@as(usize, 0), depthShared(0.2, sloped, 3, 15, 2)); + // more streams make a row dearer for the round: never deeper than for fewer + try std.testing.expect(depthShared(0.7, sloped, 3, 15, 4) <= depthShared(0.7, sloped, 3, 15, 2)); + try std.testing.expect(depthShared(0.99, sloped, 3, 3, 2) <= 3); +} diff --git a/zig/src/families/llama/residency.zig b/zig/src/families/llama/residency.zig new file mode 100644 index 0000000000..557ea0fdb6 --- /dev/null +++ b/zig/src/families/llama/residency.zig @@ -0,0 +1,36 @@ +//! 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, the drafter's too, in one residency set for the keepalive, under the 70% limit. +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); + if (h.drafter) |d| { + for (d.model.checkpoint.shards.items) |sh| try list.append(h.gpa, sh.buffer); + try list.append(h.gpa, d.model.inv); + try list.append(h.gpa, d.tokens); + try mtl.resident.fields.append(h.gpa, &list, &d.scratch); + try mtl.resident.fields.append(h.gpa, &list, &d.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()); + if (h.drafter) |d| d.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.drafter) |d| d.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..eefe9fe172 --- /dev/null +++ b/zig/src/native/llama.zig @@ -0,0 +1,268 @@ +//! Llama on Metal served by the lane host: one-row rounds, windows of --drafter's or copied drafts, contexts budgeted. +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, + drafter: ?*fl.drafter.Drafter = null, // --drafter: a smaller Llama proposing each stream's next tokens + 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); + if (h.drafter) |d| d.deinit(); // its slots give their bytes back to the backend's ledger + 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.* = 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(.served)}), + else => try std.fmt.allocPrint(a, "the native Llama engine cannot load this checkpoint ({s})", .{@errorName(err)}), + }; + return null; + }; + errdefer model.deinit(); + const q = try model.config.qualification(); + if (q.role != .served) { + problem.* = try std.fmt.allocPrint(a, "{s} is checked only as a drafter (--drafter): the native Llama engine serves {s}", .{ q.name, qualifiedNames(.served) }); + model.deinit(); + return null; + } + // refusals that return null free what they hold themselves: errdefer runs only for errors + var drafter_model: ?*fl.Model = null; + errdefer if (drafter_model) |dm| dm.deinit(); + if (drafterDir(o)) |dir| { + // a smaller Llama on the model's tokens, its commands on the model's queue (one residency set for both) + const dm = fl.Model.loadOn(gpa, io, dir, model.queue) catch |err| { + problem.* = if (err == error.UnsupportedLlamaPrecision) + try std.fmt.allocPrint(a, "--drafter: {s} is a quantized checkpoint; a Llama drafter is served from its bf16 release as published", .{dir}) + else + try std.fmt.allocPrint(a, "--drafter: the native Llama engine cannot load {s} ({s}); it takes a smaller Llama checkpoint on the served model's tokenizer", .{ dir, @errorName(err) }); + model.deinit(); + return null; + }; + drafter_model = dm; // before the refusal's message is built, so a failed allocation leaves it to the errdefer + if (try drafterRefusal(a, dir, &model.config, &dm.config)) |why| { + problem.* = why; + dm.deinit(); + model.deinit(); + return null; + } + } + 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 + var scratch: u64 = fl.state.scratchBytes(&model.config, fl.state.batch_rows) + fl.state.scratchBytes(&model.config, fl.state.prompt_rows); + var most: usize = @min(model.config.context, fl.state.max_capacity); + if (drafter_model) |dm| { + scratch += dm.checkpoint.residentBytes() + fl.drafter.Drafter.fixedBytes(&dm.config); + most = @min(most, dm.config.context); // the drafter drafts within its own trained window + } + const budget = fl.budget.caches(ram, model.device.maxWorkingSet(), text, model.checkpoint.residentBytes() + scratch) catch { + if (drafter_model) |dm| dm.deinit(); + 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 streams: u32 = @max(o.lanes, 1); + const drafted: ?fl.budget.Drafted = if (drafter_model) |dm| .{ .c = &dm.config, .streams = streams } else null; + const fits = fl.budget.longestWith(&model.config, drafted, budget, most); + 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){s}", .{ window, fits, if (drafter_model != null) ", the drafter's caches and trained window counted" else "" }); + if (drafter_model) |dm| dm.deinit(); + 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.drafter = null; + h.metal = try fl.backend.Metal.init(gpa, model, .{ .capacity = @intCast(window), .streams = streams, .budget = budget }); + errdefer h.metal.deinit(); + try h.metal.warm(); + if (drafter_model) |dm| { + const capacity: usize = @as(usize, @intCast(window)) + fl.state.window_rows; + drafter_model = null; // the drafter's now: freed with it, even when its warm-up fails + h.drafter = try fl.drafter.Drafter.start(gpa, dm, .{ .capacity = capacity, .streams = streams, .ledger = &h.metal.ledger }); + h.metal.slots = h.drafter.?.slotsHook(); // its slots grow with their streams' admission, in the same budget + const dc = &dm.config; + std.log.info("llama drafter: {s}: {d} layers, bf16, {d:.2} GiB of weights, {d} caches of up to {d} tokens ({d:.2} GiB each), each grown with its stream in the same budget", .{ o.drafter.?, dc.layers, @as(f64, @floatFromInt(dm.checkpoint.residentBytes())) / (1 << 30), streams, capacity, @as(f64, @floatFromInt(fl.state.cacheBytes(dc, capacity))) / (1 << 30) }); + } + errdefer if (h.drafter) |d| d.deinit(); + 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); + if (h.drafter) |d| h.host.proposers = .{ .ctx = d, .make = fl.drafter.Drafter.make, .free = fl.drafter.Drafter.free }; + 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 }; +} + +/// The names of the checkpoints qualified for `role`, comma-separated. +fn qualifiedNames(comptime role: fl.config.Role) []const u8 { + comptime var names: []const u8 = ""; + inline for (fl.config.qualified) |q| if (q.role == role) { + names = names ++ (if (names.len == 0) "" else ", ") ++ q.name; + }; + return names; +} + +/// --drafter's directory when drafts are on; one width (bf16 as published), so --drafter-bits is unread, as in Gemma 4. +fn drafterDir(o: api.Open) ?[]const u8 { + return if (o.drafts) o.drafter else null; +} + +/// Why a drafter cannot serve: tokens that are not the model's. +fn drafterRefusal(a: std.mem.Allocator, dir: []const u8, target: *const fl.config.Config, drafter: *const fl.config.Config) !?[]const u8 { + var same = drafter.vocab <= target.vocab and drafter.eos_count == target.eos_count; + if (same) for (drafter.eos[0..drafter.eos_count]) |e| { + if (!target.isEos(e)) same = false; + }; + if (!same) return try std.fmt.allocPrint(a, "--drafter: {s} does not share the model's tokens (vocabulary {d} against {d}, or other end tokens); a drafter is a smaller Llama on the served model's tokenizer", .{ dir, drafter.vocab, target.vocab }); + return null; +} + +/// 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: the default options serve the drafter, whatever --drafter-bits holds" { + const plain: api.Open = .{ .dir = "m", .model_type = "llama", .drafter = "d" }; + try std.testing.expectEqual(@as(u8, 4), plain.drafter_bits); // the shared default, Qwen's DFlash width + try std.testing.expectEqualStrings("d", drafterDir(plain).?); + var bf16 = plain; + bf16.drafter_bits = 0; + try std.testing.expectEqualStrings("d", drafterDir(bf16).?); + var off = plain; + off.drafts = false; // --no-drafts + try std.testing.expectEqual(@as(?[]const u8, null), drafterDir(off)); +} + +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 d3ed284021..d2175e2238 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}); @@ -337,14 +335,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) and !std.mem.eql(u8, t, "llama")) return "--drafter is supported for Qwen3.8-27B-class checkpoints (qwen3_5 with its own output head), Gemma 4 dense and Llama (a smaller Llama on the same tokenizer)"; + 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 59afec1320..cce3ec1cfe 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -58,6 +58,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"); @@ -87,6 +88,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/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 1cf2a29ab8..561c178eea 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -30,6 +30,7 @@ const lfm2_argmax = @import("metal_lfm2_argmax.zig"); const lfm2_streams = @import("metal_lfm2_streams.zig"); const lfm2_tap = @import("metal_lfm2_tap.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]"; @@ -343,6 +344,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); @@ -433,6 +437,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 { diff --git a/zig/tests/llama/run.zig b/zig/tests/llama/run.zig new file mode 100644 index 0000000000..7c39c84d99 --- /dev/null +++ b/zig/tests/llama/run.zig @@ -0,0 +1,259 @@ +//! 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 w-row windows + \\ (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 + \\tf-llama-run draft , greedy's n tokens through the drafter's windows: the same + \\ tokens, with the rounds, drafts kept and tok/s + \\ +; + +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(); + 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 ..]; + } + var drafter_dir: []const u8 = ""; + if (std.mem.eql(u8, mode, "draft")) { + const comma = std.mem.indexOfScalar(u8, args[4], ',') orelse return error.ExpectedDrafterCommaSteps; + drafter_dir = args[4][0..comma]; + } + 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 (drafter_dir.len > 0) try std.fmt.parseInt(usize, args[4][drafter_dir.len + 1 ..], 10) else if (forced) |f| f.value.len else 0; + var cache = try fl.state.Cache.init(m.device, &m.config, prompt.len + steps + fl.state.window_rows + 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")) { + 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(); + 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; + } + if (width == 0 or width > fl.state.window_rows) return error.BadWindow; + 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 }); + @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); + 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 (drafter_dir.len > 0) return draft(gpa, init.io, m, &small, &cache, drafter_dir, prompt, steps); + 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)) }); +} + +/// Greedy decoding through the drafter's windows after the prompt: drafts kept while the target's rows agree. +fn draft(gpa: std.mem.Allocator, io: std.Io, m: *fl.Model, small: *fl.state.Scratch, cache: *fl.state.Cache, dir: []const u8, prompt: []const u32, steps: usize) !void { + const vocab = m.config.vocab; + const dm = try fl.Model.loadOn(gpa, io, dir, m.queue); + var ledger: fl.budget.Ledger = .{ .limit = std.math.maxInt(u64) }; // the tool's one stream: no budget to share + const capacity = prompt.len + steps + 2 * fl.state.window_rows; + const d = try fl.drafter.Drafter.start(gpa, dm, .{ .capacity = capacity, .streams = 1, .ledger = &ledger }); + defer d.deinit(); + const p = try fl.drafter.Drafter.make(d, gpa, prompt, m.config.eos[0..m.config.eos_count]); + defer fl.drafter.Drafter.free(d, p); + try d.prepare(p, capacity); + var context: std.ArrayList(u32) = .empty; + defer context.deinit(gpa); + try context.appendSlice(gpa, prompt); + var picks: std.ArrayList(u32) = .empty; + defer picks.deinit(gpa); + var pending = argmax(cache.logits.slice(u16, vocab)); + try picks.append(gpa, pending); + try context.append(gpa, pending); + var rounds: usize = 0; + var offered: usize = 0; + var landed: usize = 0; + var ids: [fl.state.window_rows]u32 = undefined; + const t0 = now(); + while (picks.items.len < steps) { + const drafts = try p.propose(context.items, fl.state.window_rows - 1); + ids[0] = pending; + @memcpy(ids[1..][0..drafts.len], drafts); + const rows = 1 + drafts.len; + try fl.forward.run(m, small, cache, ids[0..rows], .{ .all = true }); + const all = small.logits_all.slice(u16, rows * vocab); + var keep: usize = 0; + while (keep < drafts.len and argmax(all[keep * vocab ..][0..vocab]) == drafts[keep]) keep += 1; + pending = argmax(all[keep * vocab ..][0..vocab]); + cache.len += keep + 1; + try picks.appendSlice(gpa, drafts[0..keep]); + try picks.append(gpa, pending); + try context.appendSlice(gpa, drafts[0..keep]); + try context.append(gpa, pending); + p.round(@intCast(rows), @intCast(keep)); + rounds += 1; + offered += drafts.len; + landed += keep; + } + const decode = now() - t0; + picks.shrinkRetainingCapacity(steps); + var w: std.Io.Writer.Allocating = .init(gpa); + defer w.deinit(); + try std.json.Stringify.value(picks.items, .{}, &w.writer); + std.debug.print("{s}\n", .{w.written()}); + const per_round = @as(f64, @floatFromInt(rounds + landed)) / @as(f64, @floatFromInt(@max(rounds, 1))); + std.debug.print("llama draft: prompt {d} tokens, {d} tokens in {d} rounds ({d:.2} a round), {d} of {d} drafts kept, {d:.1} tok/s; step {d:.2} ms, rounds {d:.1} ms at 4 rows, {d:.1} at 8\n", .{ prompt.len, steps, rounds, per_round, landed, offered, @as(f64, @floatFromInt(steps - 1)) * 1e9 / @as(f64, @floatFromInt(decode)), d.step_ms orelse 0, d.rounds.at(4), d.rounds.at(8) }); +} + +/// 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); +}