diff --git a/CHANGELOG.md b/CHANGELOG.md index c580e697e5..a796271390 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,13 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Ripple-65B-Preview runs on native Metal (M5 or later) from the publisher's NVFP4 and bf16 releases as published, + with exact bf16 caches, thinking, tool calls and structured output; `--context` is capped at 262,144, the longest + checked (#2231, fixes #2171). Thanks to @BobClawblaw. + +- Ripple-65B-Preview runs on native CUDA (GB10 checked) from the publisher's NVFP4 release as published, with the + same bf16 layer math as Metal (#2302, fixes #2300). + - Internal: 42 more lean findings fixed in the tests' comments and the tools' docstrings, with no change in behaviour (#2315). - HIP: the decode round's residual-add and norm tail issues each row's loads before its branches, and a MoE tail loads its slots' rows together, with output bytes unchanged. On a gfx1151, one-stream decode is 2.7% faster on Qwen3.6-35B-A3B and 1.1% on the 27B. `tf_tail` is now under the kernel contract (#2247). Thanks to @AbrahamCete. diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index e9889018a6..0c47433d80 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -306,6 +306,7 @@ These follow a model's release code or a library's model code for their math, ro - `tools/lfm3/ref64.py` and `zig/kernels/xpu/lfm3.cl`: `modeling_lfm3.py` in [LiquidAI/LFM3-1.7B](https://huggingface.co/LiquidAI/LFM3-1.7B). Copyright Liquid AI, LFM Open License v1.0; see [the license text](LICENSES/LFM-Open-1.0.txt). - `zig/kernels/metal/minimax_m3/glue.metal`: the rounding points of `modeling_minimax_m3_vl.py` in [transformers](https://github.com/huggingface/transformers). +- `zig/src/families/ricursive/` (its layer order and rounding points), `zig/kernels/metal/ricursive/glue.metal`, `zig/kernels/cuda/ricursive.cu`, `ricursive.cuh`, `ricursive_prep.cu`, `ricursive_prep.cuh` and `tools/zig/ricursive_reference.py`: `modeling_ricursive.py` and the SGLang model `inference/server/model.py` in [ricursiveAI/Ripple-65B-Preview](https://huggingface.co/ricursiveAI/Ripple-65B-Preview) at revision `240d5966cbe00f0627c699e59ef58560b60303cf`. Copyright ricursiveAI, Apache License 2.0; see [the license text](LICENSES/Apache-2.0.txt). The reference's `--release` and `--release-f64` modes import the release's own file from the checkpoint at run time; TensorFold does not include it. - `zig/src/vision/qwen_cuda.zig`: `Qwen3_5VisionModel` in transformers. - `zig/src/core/picture/patches_test.zig`: recorded image-processor output data from `Lfm3VlImageProcessor` and `Qwen2VLImageProcessor` in transformers; the processing implementation is described below. - `zig/src/families/flashnext/rotary_config.zig`: `_compute_yarn_parameters` in transformers, for the 64 rotary dims. diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index b19db41e0a..6eaa3bf8bc 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -253,6 +253,22 @@ never rescales, a merge dropping each block's scale, a row projection dropping i untransposed and a prompt attention seeing the next key go red at `shapes`, and an argmax never picking a NaN at `nonfinite`. The INT4 release's row projections and dequantization have two more: every group taking its row's first zero point, and a word's nibbles read in reverse, both red at `shapes`. +Ripple-65B's own kernels (`metal_ricursive*.zig`): its glue, each output inside float64's range at the release's bf16 +rounding points over the spread its fp32 steps allow (the scaled embedding; the add-norm in the forms the forward runs, +the value residual's lerp, the sublayer's add and both; q and k's norms with RoPE and a ring slot on sliding layers, +without RoPE on global ones; the output gate's fp32 sigmoid; the clamped SwiGLU), the decode attention's partials over a +1,024-key window in a 3,072-slot ring and over a global layer's every key, and on the M5's tensor units the prompt +attention for both layer kinds; the NVFP4 MLP (E2M1 nibbles low first, an unsigned E4M3 scale a block of 16, the +tensor's fp32 global): its row kernels, bf16 of the global times the fp32 sum, the prompt's exact dequantization and the +GEMM taking the global before its one rounding. Llama's row projections, merge, argmax and GEMM forms register again at +Ripple's shapes (k to 40,960, 80 query heads, 100,352 logits), and Nemotron's `tf_sample_full`, Ripple's sampled draw, +is held to a stated rule: each draw equals lanes' `gpu_full` twin. Fifteen controls edit the source, each red at +`shapes`: an embedding reading its own row, an add-norm not writing its residual back, a lerp dropping its blend, a prep +rotating a global layer or writing one ring slot late, a gate rounding its sigmoid's exp to bf16, a clamp missing its +lower bound, a decode seeing the key before its window, a prompt attention seeing the 64 keys before it, one whose keys +wrap 64 slots early, nibbles read high first, a scale's exponent bias off by one, the global applied after the sum +rounds, a dequantization leaving out a block's scale, and a GEMM never taking the global. Their launches are in +`ricursive_records.zig`. 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 63e476b72d..63fb0a70d0 100644 --- a/docs/recipes/README.md +++ b/docs/recipes/README.md @@ -28,6 +28,7 @@ For agent-assisted contributions, read the [TensorFold contributor skill](SKILL. | 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 and INT4 releases, the Llama family](nanbeige-4.3-33b-metal.md) | +| Ripple-65B-Preview | [Native CUDA, the publisher's NVFP4 release](ripple-65b-cuda.md), [native Metal, the publisher's NVFP4 and bf16 releases](ripple-65b-metal.md) | ## Capability floor diff --git a/docs/recipes/ripple-65b-cuda.md b/docs/recipes/ripple-65b-cuda.md new file mode 100644 index 0000000000..68f0fff1a4 --- /dev/null +++ b/docs/recipes/ripple-65b-cuda.md @@ -0,0 +1,111 @@ +# Ripple-65B-Preview, CUDA + +This recipe covers text generation with ricursiveAI's Ripple-65B-Preview (model_type `ricursive`; see +[the Metal recipe](ripple-65b-metal.md) for the model) on the CUDA engine. One release is served as published, +read in place from its shards with nothing converted: + +- [`ricursiveAI/Ripple-65B-Preview-NVFP4`](https://huggingface.co/ricursiveAI/Ripple-65B-Preview-NVFP4) (revision + `99eea4477a001a41b86914e41feb4b2ed86105e5`, 53 GiB): ModelOpt NVFP4 MLPs, attention, embeddings and head in bf16. + +The bf16 release is refused by name (`RicursiveCudaNeedsNvfp4`): at 121 GiB it does not fit a GB10, and the CUDA +family reads only the NVFP4 MLPs. + +It was brought up on a DGX Spark (GB10, sm_121, 121.7 GiB shared with the host, CUDA 13). The kernels build for +sm_80, 86, 89, 120 and 121 and their launches fit each, but only GB10 was run. + +## Build and serve + +```sh +hf download ricursiveAI/Ripple-65B-Preview-NVFP4 --revision 99eea4477a001a41b86914e41feb4b2ed86105e5 \ + --local-dir models/ripple-65b-nvfp4 +zig build native -Dnvcc=/usr/local/cuda/bin/nvcc -Dsm=121 +zig-out/native/bin/tensorfold-native serve models/ripple-65b-nvfp4 --name ripple --port 8888 --context 32768 +``` + +At `--context 32768` admission gives eight streams of 2.78 GiB and a 16.0 GiB prompt cache, after a 12.2 GiB +reserve. A stream's caches: 64 of the 80 layers slide over 1,024 keys and keep a 3,072-row ring each; every fifth +layer is global and keeps every key, in 512-row pages. + +Thinking follows the chat template, as on Metal: Ripple writes `` itself when it thinks and answers bare +otherwise, and the server files a reply without the block as answer. Its think markers are spelled from several +tokens, so a thinking budget and logprobs are refused by name. + +## What the family is + +`zig/src/families/ricursive/cuda*.zig` and `zig/kernels/cuda/ricursive*.cu`, with the same layer math as Metal and +the release's `modeling_ricursive.py` as the reference: + +- **Projections.** NVFP4 MLPs on the `.nvfp4` W4A16 lane (nibbles and E4M3 scales repacked at load, the fp32 global + per tensor; `input_scale` is not used); attention projections, embeddings and head in bf16 through `qmmf`. Q, K and + the gate share one projection; V reads the value-residual blend. +- **Glue** (`ricursive.cu`): the embedding scale in fp32, the add-and-norm with the value-residual lerp, the gate + (bf16 sigmoid before the multiply) and the clamped SwiGLU. `ricursive_prep.cu` does the Q/K norm, RoPE (bf16 cos + and sin tables; none on the global layers) and the K/V write into a ring or pages in one kernel. +- **Attention.** Sliding decode on `tf_window_decode_tc` over each layer's ring; global decode on the paged + `tree_attention` kernels; prompt chunks on the window kernel (paged for the global layers) with two query heads a + block, for Ripple's group of 10. +- **Serving.** The lane host's shared rounds as CUDA graphs, greedy or sampled draws, copy drafts, verify windows, + kept states and the prompt cache. + +**Activations stay bf16**, as on Metal: the engine serves the NVFP4 weights with bf16 activations and fp32 sums, not +the publisher's W4A4 kernels. + +Every kernel is registered with the kernel contract, with planted controls that each must turn its gate red. On +GB10 (nvcc 13.3, sm_121): + +``` +SUMMARY ricursive: 20 kernels, 20 pass gates 1-4, 0 fail; 24 controls: 24 red as planted, 0 missed +``` + +## Correctness (GB10) + +`tools/zig/ricursive_reference.py` (from the Metal PR) on the Spark's CPU, teacher-forced over the same 2,586-token +chat as the Metal recipe, against the NVFP4 release's own passes: + +| engine logits | against | mean KL | top-1 | +|---|---|---:|---:| +| 2,048-row prompt chunks | the float64 pass | 0.00069 | 98.72% | +| 2,048-row prompt chunks | the `--bf16-points` pass | 0.00078 | 98.84% | +| one-row decode | the float64 pass | 0.00076 | 98.92% | +| one-row decode | the `--bf16-points` pass | 0.00080 | 99.07% | +| `--bf16-points` pass (no engine) | the float64 pass | 0.00073 | 98.96% | + +The engine sits as close to the float64 pass as the bf16-rounded reference does (Metal: 0.00069 prompt, 0.00107 +decode). The float64 pass equals the release's own `modeling_ricursive.py` run in float64 on the same weights +(`--release-f64`: top-1 100%, largest KL 0.0005). Every engine miss is a near-tie: the reference leads the engine's +pick by at most 0.25 nats. + +- Prompt chunks of 2,048, 1,000 and 333 rows give bit-identical logits on all 2,586 rows; no non-finite logit. +- A greedy 128-token reply equals the float64 pass's argmax token for token. +- With graphs on and with `--eager`, every dump and reply hash equals the first run's bit for bit. +- Served replies are bit-identical alone, beside seven other streams and one at a time (greedy and sampled). A + conversation sent again, reusing 79 to 81 cached tokens, equals its first reply, and 60 repeats of a reply equal the + first. Seeded replies equal their repeats, and streamed text equals the reply sent whole. + +## Measured (GB10, NVFP4 release, `--context 32768`, no other model busy) + +One stream decodes at 248.3 ms a token (4.03 tok/s); each step reads all 53 GiB of weights. Prompts run +at 518-520 tok/s at 2,586 rows. Greedy 256-token essays, thinking off, all streams together: + +| streams | tok/s | +|---:|---:| +| 1 | 4.03 | +| 4 | 15.54 | +| 8 | 29.93 | + +At `--context 131072` admission gives four streams of 8.78 GiB and a 7.1 GiB prompt cache. Cold prompts there +(`tools/bench_cache.py --require-cold`), time to first token: 3.76 s at 2,048 tokens (545 tok/s), 15.6 s at 8,192 +(524), 68.4 s at 32,768 (479) and 150.7 s at 65,536 (435). Each reply equals its repeat and itself beside a second. +A 114,528-token prompt with a random four-word passphrase a third of the way in: the reply names it, first token +after 300.5 s (381 tok/s). + +## Not done + +Structured output (`response_format`) and forced tool calls (`tool_choice` `"required"` or a named function) are not +supported on CUDA yet: the server refuses them, because the CUDA family has no masked sampling. Tool calls with +`tool_choice` `"auto"` work. No quality suite was run. Contexts past 131,072 tokens were not served on CUDA. Only +GB10 was run. The publisher's DFlash2 drafter (`ricursiveAI/Ripple-65B-Preview-DFlash2`) is not served on CUDA yet. + +## Licence + +Ripple-65B-Preview and its NVFP4 release are published under Apache-2.0 at the revisions measured here. diff --git a/docs/recipes/ripple-65b-metal.md b/docs/recipes/ripple-65b-metal.md new file mode 100644 index 0000000000..be8a909428 --- /dev/null +++ b/docs/recipes/ripple-65b-metal.md @@ -0,0 +1,170 @@ +# Ripple-65B-Preview, native Metal + +This recipe covers text generation with ricursiveAI's Ripple-65B-Preview (model_type `ricursive`: a dense 65B model, +80 layers, contexts up to 1M tokens by its model card, thinking and tool calls in its chat template) on the Metal +engine. Two of the publisher's releases are served as published, read in place from their shards with nothing +converted: + +- [`ricursiveAI/Ripple-65B-Preview-NVFP4`](https://huggingface.co/ricursiveAI/Ripple-65B-Preview-NVFP4) (revision + `99eea4477a001a41b86914e41feb4b2ed86105e5`, 53 GiB): ModelOpt NVFP4 MLPs, attention, embeddings and head in bf16. + The recommended one on a 256 GB Mac: it leaves 135 GiB for caches. +- [`ricursiveAI/Ripple-65B-Preview`](https://huggingface.co/ricursiveAI/Ripple-65B-Preview) (revision + `240d5966cbe00f0627c699e59ef58560b60303cf`, 121 GiB): the bf16 release, the exact model. + +It was brought up on an M5 Ultra (256 GB); other Macs are untested. Prompts run on Metal 4 tensor units, so it needs +an M5 or later: M1 to M4 Macs are refused by name at load. Other `ricursive` checkpoints, and either release quantized +any other way, are refused by name. + +## Build and serve + +```sh +hf download ricursiveAI/Ripple-65B-Preview-NVFP4 --revision 99eea4477a001a41b86914e41feb4b2ed86105e5 \ + --local-dir models/ripple-65b-nvfp4 +zig build native -Doptimize=ReleaseFast -Dcpu=native +zig build grammar -Dxgrammar=PATH_TO_XGRAMMAR_v0.2.8 -Dcpu=native # structured output +zig-out/native/bin/tensorfold-native serve models/ripple-65b-nvfp4 --name ripple --port 8888 \ + --context 262144 --parallel 2 +``` + +The bf16 release serves the same way from its own directory. A stream's caches take 16.75 GiB at 262,144 +tokens: 64 of the 80 layers slide over 1,024 keys and keep a 3,072-slot ring each (the window and a 2,048-row prompt +chunk); every fifth layer is global and keeps every key, in bf16. Under the default load limit (75% of RAM) a 256 GB +Mac holds the NVFP4 weights (54.8 GiB, wired between requests) and 135 GiB of caches: eight 262,144-token streams. +Thinking follows the chat template: Ripple writes `` itself when it thinks and answers bare otherwise, and the +server files a reply without the block as answer, not reasoning, whatever `enable_thinking` says (the template has no +switch). `--context` takes up to 262,144 tokens, the default and the longest checked (below); the card's 1M is +refused by name until a run at that length qualifies it. + +## What the family is + +`zig/src/families/ricursive/`, with the release's `modeling_ricursive.py` and ricursiveAI's SGLang model as the +reference: + +- **Layers.** PeriNorm (a norm before and after each sublayer, bf16 residual, each add rounded to it), per-head q + and k norms, RoPE (theta 10,000) on sliding layers and none on global ones, a sigmoid output gate in fp32, the value + residual (each layer's value input blends the first layer's scaled embeddings with its own normed input), and the + clamped SwiGLU (gate at most 20, up within ±20; a NaN stays NaN, as `torch.clamp` keeps it). The glue is + `kernels/metal/ricursive/glue.metal`. +- **Projections.** bf16 rows read in place from the mapped shards. Prompt chunks of up to 2,048 rows run on llama's + tensor-unit GEMM; decode windows (up to 32 rows) run llama's `ll_matvec`. Every chunk is planned on the same + 2,048-row grid, so a resumed prompt equals the prompt served fresh. +- **NVFP4 MLPs** (`fp4.metal`, `fp4_gemm.metal`). E2M1 nibbles (low first), an unsigned E4M3 scale per 16 values, + one fp32 global per tensor. A nibble times its scale is exact in bf16. Decode windows sum in fp32 and apply the global + before the one bf16 rounding; prompt chunks first write the exact fp4 × scale values to a scratch and run the GEMM + over them, each fp32 sum times the global before it rounds (CUTLASS's alpha epilogue). +- **Attention.** Prompt chunks on the tensor units over the sliding ring or every key (`attention_nax.metal`); decode + rows over key blocks at absolute positions, so a row's bits never depend on the rows beside it. +- **Serving.** The lane host's shared rounds, greedy or `tf_sample_full` draws, copy drafts, the prompt cache. + +**Activations stay bf16.** The publisher serves the NVFP4 release with SGLang's `flashinfer_cutlass`, which also +quantizes activations to fp4 (W4A4). TensorFold does not run reduced-precision activations, so the engine serves the +release's NVFP4 weights with bf16 activations and fp32 sums: what is measured below is that model, not the +publisher's W4A4 kernels. + +Every kernel is registered with the kernel contract (`zig/tests/contract/metal_ricursive*.zig`, the roll call in +`contracts.zig`): the glue, decode and prompt attention for both layer kinds, the NVFP4 row kernels, dequantization and +GEMM, and llama's row projections, merge, argmax and GEMM again at Ripple's shapes (every reduction it runs: k 5,120, +10,240 for the attention output and 40,960). Fifteen planted controls each edit one kernel's source, and each must turn +its gate red. `tf-kernel-contract --only ricursive` on the M5 Ultra: + +``` +SUMMARY 46 kernels: 46 pass gates 1-4, 0 fail, 0 known findings; 15 controls: 15 red as planted, 0 missed; 0 limit problems; gates 645430 ms, limits 140 ms +``` + +## Correctness (M5 Ultra, 4 Nov 2026) + +`tools/zig/ricursive_reference.py` runs a checkpoint on the CPU, one layer's weights at a time: in float64 (the +weights cast up exactly; the NVFP4 MLPs dequantized exactly), and with `--bf16-points` rounded to bf16 at every point +the release's bf16 forward rounds; with `--release`, the checkpoint's own `modeling_ricursive.py` layers in bf16 (the +torch CPU forward, its weights loaded a layer at a time). Teacher-forced over a 2,586-token chat: + +| engine logits | against | mean KL | notes | +|---|---|---:|---| +| bf16 release, 2,048-row prompt chunks | its float64 pass | 0.00086 | | +| bf16 release, 16-row decode windows | its float64 pass | 0.00073 | | +| bf16 release, 2,048-row prompt chunks | `modeling_ricursive.py` in bf16 | 0.00073 | top-1 98.7% | +| bf16 release, 16-row decode windows | `modeling_ricursive.py` in bf16 | 0.00069 | top-1 98.8% | +| `modeling_ricursive.py` in bf16 (no engine) | the float64 pass | 0.00079 | top-1 98.3% | +| NVFP4 release, 2,048-row prompt chunks | its float64 pass | 0.00069 | | +| NVFP4 release, 16-row decode windows | its float64 pass | 0.00107 | | +| NVFP4 release, prompt chunks | the bf16 release's float64 pass | 0.0060 | top-1 96.1% | +| NVFP4 float64 pass (no engine) | the bf16 release's float64 pass | 0.0059 | top-1 95.8% | + +Against the release's own bf16 code the engine sits as close as the float64 pass does (0.00073 and 0.00069 against +0.00076 for float64 measured the same way): the differences are bf16 rounding order, on either side. The last two rows +are **what the NVFP4 release costs**: the publisher's quantization, not the engine's, about eight +times the engine's own rounding. The engine adds nothing measurable to it. + +- Prompt chunks of 2,048, 1,000 and 333 rows give bit-identical logits; decode windows of 16 and 7 rows equal one-row + steps bit for bit (both releases). +- A greedy bf16 reply of 77 tokens equals the float64 pass's argmax token for token; on a short NVFP4 prompt, 21 of + 21 positions take the float64 pass's top token (mean KL 0.0005). +- One NVFP4 position of 5,200 in a longer prompt flips. After "every segment gets 4," the float64 pass gives `000` + (0.645), decode windows `000` (0.621), and the prompt path `096` (0.855). Moving the NVFP4 projections to the row + kernels still flips it; moving every projection there does not. The flip comes from the shared bf16 tensor-unit + GEMM on the attention projections, which matches MLX bit for bit and passes the contract at these shapes, at a + position where the NVFP4 model is nearly torn between two answers. bf16 does not flip there. Leaving that position + out, the prompt path's mean KL equals the decode path's. + +## Measured + +Served through the recipe harness (4 Nov 2026, the NVFP4 release, no other model busy; speed and prompt runs on +build 588e7d2a2, the serving checks repeated on d26cd59ad with the same results and speeds), `--context 262144 +--parallel 2`, greedy 256-token replies: + +| streams | prose | code | structured | list | first token | a round | +|---:|---:|---:|---:|---:|---:|---:| +| 1 | 10.1 | 10.1 | 10.2 | 10.1 | 0.61-0.74 s | 99.2 ms | +| 2 | 18.8 | 18.6 | 19.1 | 18.8 | 0.93-1.10 s | 107.0 ms | + +Tokens a second, all streams together. One stream reads the 53 GiB of weights each step; two read them once for both +rows. The same ruler after the serving checks gives the same speeds. + +| cold prompt | 2,220 | 8,776 | 35,015 | 139,723 tokens | +|---|---:|---:|---:|---:| +| tokens a second | 569 | 605 | 587 | 434 | +| time to first token | 3.9 s | 14.5 s | 59.6 s | 322 s | + +A fresh server, each prompt starting with words no earlier one did. With thinking requested (new prompts of the same +lengths), 488, 458, 503 and 410 tokens a second. +A multi-turn chat's third turn, resumed from 4,513 of its 4,597 kept tokens, starts in 0.67 s against 8.2 s cold, +with the same reply as on a restarted server. + +The bf16 release through the same harness (5 Nov 2026, built from the context-cap commit before its rebase; no other model busy, `--context 262144 +--parallel 2`, greedy 256-token replies): + +| streams | prose | code | structured | list | first token | a round | +|---:|---:|---:|---:|---:|---:|---:| +| 1 | 5.8 | 5.0 | 4.9 | 5.0 | 0.51-1.30 s | 194.7 ms | +| 2 | 11.5 | 11.0 | 10.5 | 11.8 | 0.86-1.81 s | 179.6 ms | + +One stream reads the 121 GiB of weights each step. The rows spread more than NVFP4's (the ruler after the serving +checks gives 4.4 to 5.8 at 197.2 ms a round); the rounds are the steady figure. Cold prompts on a fresh server, with +thinking requested: 2,206 tokens 519 tok/s (4.3 s), 8,830 596 (14.8 s), 34,931 586 (59.6 s), 139,870 459 (305 s). +Without thinking requested the harness timed nothing: on all four prompts Ripple opened a think block of its own, +and the 8-token replies held no answer token to time. A resumed third turn (4,513 of 4,597 kept) starts in 0.59 s +against 7.5 s cold, with the same reply as on a restarted server. + +- 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 9/9 equal on repeat and undrafted. +- A 60-turn soak and 12 concurrent replies each equal to the request alone; streaming 4/4, multi-turn recall 8/8, + stop strings 3/3, verbatim copy 2/2, max_tokens 3/3. +- Thinking 5/10 by the suite's rule, which wants reasoning in every reply: Ripple decides for itself whether to think + (its template has no switch). Sent again, five replies thought and five answered directly, nine of ten answers right + (the miss, a clock angle, answered directly); no think tags leaked into a reply. +- A passphrase a third of the way into a 120,041-token and a 258,039-token prompt (greedy, on the same server) comes + back exactly; the cold prompts take 273 s and 790 s to the first token. The bf16 release found it too, at 120k and + 258k tokens. + +## Not done + +No quality suite was run. Contexts past 258,039 tokens were not run, so `--context` stops at 262,144 and refuses +the card's 1M by name. Ripple's think markers are spelled from several tokens, not one each, so a thinking budget +(which forces a close) and logprobs (which cut the reply at its close token) are refused by name, and a reply's +`reasoning_tokens` reads 0. The family has no draft head: drafts come only from the lane core's copies. No images: +Ripple has no vision tower. Only the M5's tensor units were checked; M1 to M4 Macs are refused at load, and smaller +RAM is untested. The publisher's DFlash2 drafter (`ricursiveAI/Ripple-65B-Preview-DFlash2`) is not served yet. + +## Licence + +Ripple-65B-Preview and its NVFP4 release are published under Apache-2.0 at the revisions measured here. diff --git a/tools/zig/ricursive_gate_compare.py b/tools/zig/ricursive_gate_compare.py new file mode 100644 index 0000000000..d8b14fba2a --- /dev/null +++ b/tools/zig/ricursive_gate_compare.py @@ -0,0 +1,266 @@ +"""Ripple's engine logits against ricursive_reference.py's: KL, top-1, worst rows, identical dumps, a greedy reply.""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +from pathlib import Path + +import numpy as np + +USAGE = """ +Logits files are bf16 bits [rows][vocab]. The reference writes them raw, with no header. The engine's `tensorfold kl +--dump` writes u32 rows, u32 vocab, then a record a row: the token fed next (u32) and the row's logits. Either is read +(the format is told from the file size); row r is the distribution of token r + 1 given tokens 0..r. + + ricursive_gate_compare.py finite FILE... [--ids ids.json] (non-finite logits, distinct top-1 tokens, hits on the next token) + ricursive_gate_compare.py kl REF ENGINE --ids ids.json [--max-kl X --min-top1 Y] [--out report.json] + (a REF with more rows than ids, as one run over prompt + reply gives, is read for its first rows) + ricursive_gate_compare.py same FILE FILE... [--vocab N] (every row's logits bit for bit equal to the first file's) + ricursive_gate_compare.py reply REF REPLY.json --prompt-len P (REF over prompt + reply: argmax of row P-1+i is reply[i]) + ricursive_gate_compare.py join IDS.json REPLY.json OUT.json (the prompt's ids followed by the reply's) + ricursive_gate_compare.py split IDS.json PROMPT.json --prompt N (the first N ids to PROMPT.json; the rest on stdout, comma + separated, for `tensorfold kl --gold`) +""" + +HEADER = 8 # a kl dump's u32 rows and u32 vocab + + +def bf16_to_f32(bits: np.ndarray) -> np.ndarray: + return (bits.astype(np.uint32) << 16).view(np.float32) + + +def log_softmax(x: np.ndarray) -> np.ndarray: + x = x.astype(np.float64) + m = x.max(-1, keepdims=True) + return x - m - np.log(np.exp(x - m).sum(-1, keepdims=True)) + + +def read_ids(path: str) -> list[int]: + data = json.loads(Path(path).read_text()) + if isinstance(data, dict): + data = data.get("ids", data.get("tokens")) + return [int(t) for t in data] + + +def open_logits(path: str, vocab: int) -> tuple[np.ndarray, np.ndarray | None]: + """The file's logits as a [rows, vocab] uint16 view, and the fed tokens when it is a kl dump.""" + size = os.path.getsize(path) + rec = 4 + 2 * vocab + raw = size > 0 and size % (2 * vocab) == 0 + kl = size > HEADER and (size - HEADER) % rec == 0 + if raw == kl: + raise SystemExit(f"{path}: {size} bytes is neither [rows, {vocab}] raw bf16 nor a kl dump of that vocabulary") + if raw: + rows = size // (2 * vocab) + return np.memmap(path, dtype=" dict: + """Per row KL(ref || engine) over the whole vocabulary in float64, both greedy picks, and the worst rows.""" + rows = ref.shape[0] + kl = np.zeros(rows) + agree = np.zeros(rows, dtype=bool) + ref_top = np.zeros((rows, 3), dtype=np.int64) + eng_top = np.zeros((rows, 3), dtype=np.int64) + ref_lp = np.zeros((rows, 3)) + eng_lp = np.zeros((rows, 3)) + finite = np.ones(rows, dtype=bool) + lead = np.zeros(rows) # the reference's logprob lead of its top-1 over the engine's top-1 (0 where they agree) + for lo in range(0, rows, chunk): + hi = min(rows, lo + chunk) + r, e = bf16_to_f32(np.asarray(ref[lo:hi])), bf16_to_f32(np.asarray(eng[lo:hi])) + finite[lo:hi] = np.isfinite(e).all(-1) + lp, lq = log_softmax(r), log_softmax(e) + kl[lo:hi] = (np.exp(lp) * (lp - lq)).sum(-1) + ref_top[lo:hi] = np.argsort(-lp, axis=-1, kind="stable")[:, :3] + eng_top[lo:hi] = np.argsort(-lq, axis=-1, kind="stable")[:, :3] + ref_lp[lo:hi] = np.take_along_axis(lp, ref_top[lo:hi], -1) + eng_lp[lo:hi] = np.take_along_axis(lq, eng_top[lo:hi], -1) + agree[lo:hi] = ref_top[lo:hi, 0] == eng_top[lo:hi, 0] + at = np.arange(hi - lo) + lead[lo:hi] = lp[at, ref_top[lo:hi, 0]] - lp[at, eng_top[lo:hi, 0]] + worst = [] + for r in np.argsort(-np.nan_to_num(kl, nan=np.inf), kind="stable")[:top]: + worst.append({ + "row": int(r), + "next_token": ids[r + 1] if r + 1 < len(ids) else None, + "kl": float(kl[r]), + "same_top1": bool(agree[r]), + "reference": [[int(t), round(float(p), 4)] for t, p in zip(ref_top[r], ref_lp[r])], + "engine": [[int(t), round(float(p), 4)] for t, p in zip(eng_top[r], eng_lp[r])], + "reference_margin": round(float(ref_lp[r, 0] - ref_lp[r, 1]), 4), + }) + return { + "rows": rows, + "kl_mean": float(kl.mean()), + "kl_median": float(np.median(kl)), + "kl_p99": float(np.percentile(kl, 99)), + "kl_max": float(kl.max()), + "top1_agreement": float(agree.mean()), + "top1_misses": int((~agree).sum()), + "nonfinite_rows": int((~finite).sum()), + "miss_leads": sorted(round(float(m), 4) for m in lead[~agree]), + "worst": worst, + } + + +def run_kl(a: argparse.Namespace) -> int: + ids = read_ids(a.ids) + ref, ref_fed = open_logits(a.ref, a.vocab) + eng, fed = open_logits(a.engine, a.vocab) + if ref.shape[0] > len(ids): + ref = ref[: len(ids)] # a causal model: the first rows of a longer run are the shorter run's + if fed is None and eng.shape[0] > len(ids): + eng = eng[: len(ids)] # the same for a raw file, so two reference passes compare over the chat's rows + if ref.shape != eng.shape or ref.shape[0] != len(ids): + raise SystemExit(f"shapes differ: reference {ref.shape}, engine {eng.shape}, {len(ids)} ids") + if fed is not None: # a kl dump says which token it fed after each row: the chat's next one, but for the last row + wrong = np.flatnonzero(np.asarray(fed[:-1]) != np.asarray(ids[1:], dtype=np.uint32)) + if wrong.size: + raise SystemExit(f"the engine's file fed other tokens than the ids at {wrong.size} rows, first row {wrong[0]}") + rep = compare(ref, eng, ids, a.top) + print(f"{rep['rows']} rows: mean KL {rep['kl_mean']:.5f}, median {rep['kl_median']:.5f}, p99 {rep['kl_p99']:.4f}, " + f"max {rep['kl_max']:.4f}; top-1 {100 * rep['top1_agreement']:.2f}% ({rep['top1_misses']} misses); " + f"{rep['nonfinite_rows']} rows with a non-finite engine logit") + print(f"the {len(rep['worst'])} worst rows (row, next token, KL, reference and engine top 3 as token:logprob, margin):") + for w in rep["worst"]: + ref3 = " ".join(f"{t}:{p}" for t, p in w["reference"]) + eng3 = " ".join(f"{t}:{p}" for t, p in w["engine"]) + print(f" row {w['row']}, next {w['next_token']}, KL {w['kl']:.4f}, reference [{ref3}], engine [{eng3}], " + f"margin {w['reference_margin']}{'' if w['same_top1'] else ', TOP-1 DIFFERS'}") + leads = np.array(rep["miss_leads"]) + if leads.size: # a miss where the reference itself barely prefers its pick is a near tie, not an error + print(f"top-1 misses by the reference's lead of its pick over the engine's: {int((leads < 0.1).sum())} under " + f"0.1, {int((leads < 0.5).sum())} under 0.5, the largest {leads.max():.4f} (nats)") + if a.out: + Path(a.out).write_text(json.dumps(rep, indent=1)) + bad = rep["nonfinite_rows"] > 0 or not rep["kl_mean"] == rep["kl_mean"] + if a.max_kl is not None and not rep["kl_mean"] <= a.max_kl: + print(f"FAIL: mean KL {rep['kl_mean']:.5f} is above {a.max_kl}") + bad = True + if a.min_top1 is not None and not rep["top1_agreement"] >= a.min_top1: + print(f"FAIL: top-1 {100 * rep['top1_agreement']:.2f}% is below {100 * a.min_top1:.2f}%") + bad = True + return 1 if bad else 0 + + +def run_finite(a: argparse.Namespace) -> int: + """A first look at a dump before any reference exists: a broken engine shows as NaNs, one token, or no hits.""" + ids = read_ids(a.ids) if a.ids else None + status = 0 + for path in a.files: + logits, _ = open_logits(path, a.vocab) + bad, top = 0, [] + for lo in range(0, logits.shape[0], 64): + x = bf16_to_f32(np.asarray(logits[lo:lo + 64])) + bad += int((~np.isfinite(x).all(-1)).sum()) + top.append(x.argmax(-1)) + top = np.concatenate(top) + text = f"{path}: {logits.shape[0]} rows, {bad} with a non-finite logit, {len(set(top.tolist()))} distinct top-1 tokens" + if ids is not None and len(ids) == len(top): + text += f", top-1 is the next token at {100 * np.mean(top[:-1] == np.array(ids[1:])):.1f}% of rows" + print(text) + status |= int(bad > 0) + return status + + +def run_same(a: argparse.Namespace) -> int: + base, _ = open_logits(a.files[0], a.vocab) + status = 0 + for path in a.files[1:]: + other, _ = open_logits(path, a.vocab) + if other.shape != base.shape: + print(f"{path}: {other.shape[0]} rows, not {base.shape[0]}") + status = 1 + continue + differ = [r for r in range(base.shape[0]) if not np.array_equal(base[r], other[r])] + if not differ: + print(f"{path}: all {base.shape[0]} rows bit-identical to {a.files[0]}") + continue + status = 1 + first = differ[0] + print(f"{path}: {len(differ)} of {base.shape[0]} rows differ from {a.files[0]}; first row {first} " + f"({int((base[first] != other[first]).sum())} of {a.vocab} logits)") + return status + + +def read_reply(path: str) -> list[int]: + return read_ids(path) # a plain list, or the run command's --report (its "tokens") + + +def run_reply(a: argparse.Namespace) -> int: + reply = read_reply(a.reply) + ref, _ = open_logits(a.ref, a.vocab) + need = a.prompt_len - 1 + len(reply) + if ref.shape[0] < need: + raise SystemExit(f"the reference has {ref.shape[0]} rows, the prompt and reply need {need}") + lp = log_softmax(bf16_to_f32(np.asarray(ref[a.prompt_len - 1:need]))) + pick = lp.argmax(-1) + miss = [i for i, t in enumerate(reply) if int(pick[i]) != t] + print(f"greedy reply of {len(reply)} tokens: {len(reply) - len(miss)} equal the float64 pass's argmax") + for i in miss[:5]: + print(f" reply token {i}: engine {reply[i]} (reference logprob {lp[i, reply[i]]:.4f}), " + f"reference {int(pick[i])} ({lp[i, pick[i]]:.4f}), margin {lp[i, pick[i]] - lp[i, reply[i]]:.4f}") + return 1 if miss else 0 + + +def run_join(a: argparse.Namespace) -> int: + ids = read_ids(a.ids) + read_reply(a.reply) + Path(a.out).write_text(json.dumps(ids)) + print(f"{len(ids)} ids to {a.out}") + return 0 + + +def run_split(a: argparse.Namespace) -> int: + ids = read_ids(a.ids) + Path(a.prompt_out).write_text(json.dumps(ids[: a.prompt])) + print(",".join(str(t) for t in ids[a.prompt:])) + return 0 + + +def main(argv: list[str] | None = None) -> int: + ap = argparse.ArgumentParser(description=__doc__, epilog=USAGE, formatter_class=argparse.RawDescriptionHelpFormatter) + sub = ap.add_subparsers(dest="cmd", required=True) + kl = sub.add_parser("kl") + kl.add_argument("ref") + kl.add_argument("engine") + kl.add_argument("--ids", required=True) + kl.add_argument("--top", type=int, default=10) + kl.add_argument("--max-kl", type=float, help="exit 1 if the mean KL is above this") + kl.add_argument("--min-top1", type=float, help="exit 1 if the top-1 agreement (0 to 1) is below this") + kl.add_argument("--out", default="") + fin = sub.add_parser("finite") + fin.add_argument("files", nargs="+") + fin.add_argument("--ids", default="") + same = sub.add_parser("same") + same.add_argument("files", nargs="+") + reply = sub.add_parser("reply") + reply.add_argument("ref") + reply.add_argument("reply") + reply.add_argument("--prompt-len", type=int, required=True) + join = sub.add_parser("join") + join.add_argument("ids") + join.add_argument("reply") + join.add_argument("out") + split = sub.add_parser("split") + split.add_argument("ids") + split.add_argument("prompt_out") + split.add_argument("--prompt", type=int, required=True) + for p in (kl, fin, same, reply): + p.add_argument("--vocab", type=int, default=100352) + a = ap.parse_args(argv) + return {"kl": run_kl, "finite": run_finite, "same": run_same, "reply": run_reply, "join": run_join, "split": run_split}[a.cmd](a) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/zig/ricursive_gate_ids.py b/tools/zig/ricursive_gate_ids.py new file mode 100644 index 0000000000..9f34d97d8e --- /dev/null +++ b/tools/zig/ricursive_gate_ids.py @@ -0,0 +1,100 @@ +"""The Ripple gate's chat as token ids: the release's chat template over a fixed text, cut to an exact length.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Callable + +USAGE = """ +The Metal check's 2,586-token chat is not in the repository, so this builds one of the same length: a system turn, a +user turn holding the text's first 300 words, a fixed assistant turn, and a last user turn holding as many of the next +words as make the whole prompt, generation prompt included, exactly --length tokens. The tokens come from the release's +tokenizer.json and chat_template.jinja (needs `tokenizers` and `jinja2`); nothing else is read. + + ricursive_gate_ids.py MODEL_DIR --text FILE --out ids.json [--length 2586] +""" + +SYSTEM = "You are a careful technical assistant." +FIRST = "Summarise this passage in one sentence:\n\n" +ANSWER = "The passage describes how a model is served and checked." +LAST = "Now list every number in this passage, with what it measures:\n\n" +FILLERS = (" a", " the", " of", " and") + + +def messages(head: str, tail: str) -> list[dict]: + return [ + {"role": "system", "content": SYSTEM}, + {"role": "user", "content": FIRST + head}, + {"role": "assistant", "content": ANSWER}, + {"role": "user", "content": LAST + tail}, + ] + + +def fit(count: Callable[[list[dict]], int], text: str, length: int) -> list[dict]: + """The chat whose `count` is exactly `length`: the last user turn takes the most words that fit, then filler.""" + words = text.split() + head, rest = " ".join(words[:300]), words[300:] + if count(messages(head, "")) > length: + raise SystemExit(f"the fixed turns alone are over {length} tokens") + if count(messages(head, " ".join(rest))) < length: + raise SystemExit(f"the text holds too few words for {length} tokens") + lo, hi = 0, len(rest) # the most words k with count <= length; count never falls as words are added + while lo < hi: + mid = (lo + hi + 1) // 2 + if count(messages(head, " ".join(rest[:mid]))) <= length: + lo = mid + else: + hi = mid - 1 + tail = " ".join(rest[:lo]) + for filler in FILLERS: + for extra in range(5): + chat = messages(head, tail + filler * extra) + if count(chat) == length: + return chat + raise SystemExit(f"no filler reached exactly {length} tokens after {lo} words") + + +def load_renderer(model_dir: Path) -> Callable[[list[dict]], str]: + """The release's chat template as a function of the messages, with the generation prompt added.""" + from jinja2.sandbox import ImmutableSandboxedEnvironment + + def refuse(message: str) -> None: + raise ValueError(message) + + config = json.loads((model_dir / "tokenizer_config.json").read_text()) if (model_dir / "tokenizer_config.json").exists() else {} + special = {k: (v.get("content") if isinstance(v, dict) else v) or "" for k, v in config.items() if k in ("bos_token", "eos_token")} + env = ImmutableSandboxedEnvironment(trim_blocks=True, lstrip_blocks=True) + env.globals["raise_exception"] = refuse + template = env.from_string((model_dir / "chat_template.jinja").read_text()) + return lambda chat: template.render(messages=chat, add_generation_prompt=True, **special) + + +def main(argv: list[str] | None = None) -> int: + ap = argparse.ArgumentParser(description=__doc__, epilog=USAGE, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("model_dir") + ap.add_argument("--text", required=True) + ap.add_argument("--out", required=True) + ap.add_argument("--length", type=int, default=2586) + a = ap.parse_args(argv) + from tokenizers import Tokenizer + + model = Path(a.model_dir) + tok = Tokenizer.from_file(str(model / "tokenizer.json")) + render = load_renderer(model) + + def encode(chat: list[dict]) -> list[int]: + return tok.encode(render(chat), add_special_tokens=False).ids + + chat = fit(lambda c: len(encode(c)), Path(a.text).read_text(), a.length) + ids = encode(chat) + Path(a.out).write_text(json.dumps(ids)) + print(f"{len(ids)} ids to {a.out}; first {ids[:4]}, last {ids[-4:]}") + print("the prompt ends: " + repr(render(chat)[-60:])) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/zig/ricursive_gate_test.py b/tools/zig/ricursive_gate_test.py new file mode 100644 index 0000000000..1726727e04 --- /dev/null +++ b/tools/zig/ricursive_gate_test.py @@ -0,0 +1,183 @@ +"""The Ripple gate tools on synthetic files: KL, top-1, worst rows, identical dumps, replies, an exact-length chat.""" + +import contextlib +import io +import json +import os +import tempfile +import unittest +from pathlib import Path + +import numpy as np + +import ricursive_gate_compare as cmp +import ricursive_gate_ids as ids_tool + +VOCAB, ROWS = 50, 12 + + +def bf16(x: np.ndarray) -> np.ndarray: + return (x.astype(np.float32).view(np.uint32) >> 16).astype(np.uint16) + + +def f32(bits: np.ndarray) -> np.ndarray: + return (bits.astype(np.uint32) << 16).view(np.float32) + + +def run(*argv: str) -> tuple[int, str]: + out = io.StringIO() + with contextlib.redirect_stdout(out): + code = cmp.main(list(argv)) + return code, out.getvalue() + + +class CompareTests(unittest.TestCase): + def setUp(self): + folder = tempfile.TemporaryDirectory() + self.addCleanup(folder.cleanup) + self.dir = Path(folder.name) + rng = np.random.default_rng(1) + self.ref = bf16(rng.normal(0, 3, (ROWS, VOCAB))) + self.eng = self.ref.copy() + self.eng[3] = bf16(rng.normal(0, 3, VOCAB)) # row 3 is another distribution + self.ids = [int(t) for t in rng.integers(0, VOCAB, ROWS)] + (self.dir / "ids.json").write_text(json.dumps(self.ids)) + + def raw(self, name: str, bits: np.ndarray) -> str: + path = self.dir / name + bits.astype(" str: + path = self.dir / name + with open(path, "wb") as f: + f.write(np.array([bits.shape[0], bits.shape[1]], dtype=" float: + p, q = f32(self.ref[r]).astype(np.float64), f32(self.eng[r]).astype(np.float64) + p, q = p - np.log(np.exp(p).sum()), q - np.log(np.exp(q).sum()) + return float((np.exp(p) * (p - q)).sum()) + + def test_kl_matches_the_formula_for_both_formats(self): + fed = self.ids[1:] + [0] + for eng in (self.raw("e.bin", self.eng), self.kl_dump("e.kl", self.eng, fed)): + code, text = run("kl", self.raw("r.bin", self.ref), eng, "--ids", str(self.dir / "ids.json"), + "--vocab", str(VOCAB), "--out", str(self.dir / "rep.json")) + rep = json.loads((self.dir / "rep.json").read_text()) + self.assertEqual(code, 0) + self.assertAlmostEqual(rep["kl_mean"], self.expected_kl(3) / ROWS, places=9) + self.assertEqual(rep["worst"][0]["row"], 3) + self.assertAlmostEqual(rep["worst"][0]["kl"], self.expected_kl(3), places=9) + self.assertEqual(rep["worst"][0]["next_token"], self.ids[4]) + self.assertIn("worst rows", text) + same = float(np.mean(f32(self.ref).argmax(-1) == f32(self.eng).argmax(-1))) + self.assertAlmostEqual(rep["top1_agreement"], same) + + def test_pass_rule_flags_kl_and_nonfinite_logits(self): + args = ("--ids", str(self.dir / "ids.json"), "--vocab", str(VOCAB)) + ref = self.raw("r.bin", self.ref) + eng = self.raw("e.bin", self.eng) + self.assertEqual(run("kl", ref, eng, *args, "--max-kl", "10")[0], 0) + code, text = run("kl", ref, eng, *args, "--max-kl", "0.0001") + self.assertEqual(code, 1) + self.assertIn("FAIL", text) + bad = self.eng.copy() + bad[5, 2] = 0x7FC0 # a bf16 NaN + self.assertEqual(run("kl", ref, self.raw("n.bin", bad), *args)[0], 1) + + def test_a_dump_that_fed_other_tokens_is_refused(self): + dump = self.kl_dump("e.kl", self.eng, [t + 1 for t in self.ids]) + with self.assertRaises(SystemExit): + run("kl", self.raw("r.bin", self.ref), dump, "--ids", str(self.dir / "ids.json"), "--vocab", str(VOCAB)) + + def test_same_finds_the_first_differing_row(self): + a = self.raw("a.bin", self.ref) + self.assertEqual(run("same", a, self.raw("b.bin", self.ref), "--vocab", str(VOCAB))[0], 0) + code, text = run("same", a, self.raw("c.bin", self.eng), "--vocab", str(VOCAB)) + self.assertEqual(code, 1) + self.assertIn("1 of 12 rows differ", text) + self.assertIn("first row 3", text) + + def test_reply_equal_and_not(self): + prompt = 5 + ref = self.ref + reply = [int(f32(ref[prompt - 1 + i]).argmax()) for i in range(4)] + path = self.dir / "reply.json" + path.write_text(json.dumps({"tokens": reply})) # the run command's report + self.assertEqual(run("reply", self.raw("r.bin", ref), str(path), "--prompt-len", str(prompt), "--vocab", str(VOCAB))[0], 0) + reply[2] = (reply[2] + 1) % VOCAB + path.write_text(json.dumps(reply)) + code, text = run("reply", self.raw("r.bin", ref), str(path), "--prompt-len", str(prompt), "--vocab", str(VOCAB)) + self.assertEqual(code, 1) + self.assertIn("3 equal", text) + self.assertIn("reply token 2", text) + + def test_a_longer_reference_is_read_for_its_first_rows(self): + longer = np.concatenate([self.ref, self.ref[:3]]) + code, _ = run("kl", self.raw("r.bin", longer), self.raw("e.bin", self.eng), "--ids", str(self.dir / "ids.json"), + "--vocab", str(VOCAB), "--out", str(self.dir / "rep.json")) + self.assertEqual(json.loads((self.dir / "rep.json").read_text())["rows"], ROWS) + self.assertEqual(code, 0) + + def test_finite_reports_nan_rows_and_next_token_hits(self): + good = self.raw("g.bin", self.ref) + code, text = run("finite", good, "--vocab", str(VOCAB), "--ids", str(self.dir / "ids.json")) + self.assertEqual(code, 0) + self.assertIn("0 with a non-finite logit", text) + hits = 100 * np.mean(f32(self.ref).argmax(-1)[:-1] == np.array(self.ids[1:])) + self.assertIn(f"{hits:.1f}%", text) + bad = self.ref.copy() + bad[4, 0] = 0x7F80 # bf16 +inf + code, text = run("finite", self.raw("b.bin", bad), "--vocab", str(VOCAB)) + self.assertEqual(code, 1) + self.assertIn("1 with a non-finite logit", text) + + def test_split_writes_the_prompt_and_prints_the_rest(self): + code, text = run("split", str(self.dir / "ids.json"), str(self.dir / "p.json"), "--prompt", "2") + self.assertEqual(code, 0) + self.assertEqual(json.loads((self.dir / "p.json").read_text()), self.ids[:2]) + self.assertEqual(text.strip(), ",".join(map(str, self.ids[2:]))) + + def test_join(self): + (self.dir / "reply.json").write_text(json.dumps({"tokens": [7, 8]})) + run("join", str(self.dir / "ids.json"), str(self.dir / "reply.json"), str(self.dir / "all.json")) + self.assertEqual(json.loads((self.dir / "all.json").read_text()), self.ids + [7, 8]) + + +class IdsTests(unittest.TestCase): + def count(self, chat): # a stand-in tokenizer: a token a word, two for a long word, and four a turn + return sum(sum(2 if len(w) > 3 else 1 for w in m["content"].split()) + 4 for m in chat) + 3 + + def test_fit_is_exact_and_deterministic(self): + text = " ".join(f"w{i}" for i in range(2000)) + for length in (1500, 1501, 1502, 1503): + chat = ids_tool.fit(self.count, text, length) + self.assertEqual(self.count(chat), length) + self.assertEqual(chat, ids_tool.fit(self.count, text, length)) + + def test_fit_refuses_a_short_text_and_long_fixed_turns(self): + with self.assertRaises(SystemExit): + ids_tool.fit(self.count, " ".join(["w"] * 400), 2586) + with self.assertRaises(SystemExit): + ids_tool.fit(self.count, " ".join(["w"] * 2000), 100) + + def test_the_release_template_renders_the_chat(self): + try: + import jinja2 # noqa: F401 + except ImportError: + self.skipTest("jinja2 is not installed") + folder = tempfile.TemporaryDirectory() + self.addCleanup(folder.cleanup) + (Path(folder.name) / "chat_template.jinja").write_text( + "{% for m in messages %}<|im_start|>{{ m.role }}\n{{ m.content }}<|im_end|>\n{% endfor %}" + "{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}") + text = ids_tool.load_renderer(Path(folder.name))(ids_tool.messages("a b", "c d")) + self.assertTrue(text.startswith("<|im_start|>system\n")) + self.assertTrue(text.endswith("<|im_start|>assistant\n")) + + +if __name__ == "__main__": + unittest.main() diff --git a/tools/zig/ricursive_reference.py b/tools/zig/ricursive_reference.py new file mode 100644 index 0000000000..91dab02f3f --- /dev/null +++ b/tools/zig/ricursive_reference.py @@ -0,0 +1,190 @@ +"""A Ripple (ricursive) checkpoint's teacher-forced logits on the CPU, a layer at a time from its safetensors.""" +import importlib.util +import json +import sys +import time + +import numpy as np +import torch +from safetensors import safe_open + +USAGE = """usage: ricursive_reference.py MODEL_DIR tokens.json out.bin [--bf16-points | --release | --release-f64] +writes raw bf16 logits [positions, vocab]; an NVFP4 release's MLPs are its exact dequantized values + (default) float64: bf16 weights cast up exactly, the release's arithmetic, one rounding at the end + --bf16-points float64 rounded to bf16 at every point the release's bf16 forward rounds + --release the checkpoint's own modeling_ricursive.py layers in bf16 (bf16 releases only) + --release-f64 those layers in float64, an NVFP4 release's MLPs as their exact dequantized values""" + + +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 checked(cfg: dict) -> None: + """Only the release's layer: gated (fp32 gate), value residuals, sliding and global layers, no scaled RoPE.""" + want = {"model_type": "ricursive", "use_gated_attention": True, "gate_full_precision": True, "use_valres": True, + "is_hybrid_swa": True, "tie_word_embeddings": False, "hidden_act": "silu"} + for key, value in want.items(): + if cfg.get(key) != value: + raise SystemExit(f"config {key} is {cfg.get(key)!r}, not the release's {value!r}") + q = cfg.get("quantization_config") + if q and (q.get("quant_method") != "modelopt" or q.get("quant_algo") != "NVFP4"): + raise SystemExit("only the publisher's ModelOpt NVFP4 MLPs are read") + if cfg.get("rope_scaling"): + raise SystemExit("scaled RoPE is not the release's") + + +def main(model_dir: str, tokens: str, out: str, mode: str = "") -> None: + torch.set_num_threads(24) + cfg = json.load(open(f"{model_dir}/config.json")) + checked(cfg) + index = json.load(open(f"{model_dir}/model.safetensors.index.json"))["weight_map"] + ids = json.load(open(tokens)) + n = len(ids) + + def tensor(name): + with safe_open(f"{model_dir}/{index[name]}", framework="pt") as f: + return f.get_tensor(name) + + t0 = time.time() + with torch.no_grad(): + h = release(model_dir, cfg, tensor, ids, mode == "--release-f64") if mode.startswith("--release") else exact(cfg, tensor, ids, mode == "--bf16-points", index) + head = tensor("lm_head.weight").to(h.dtype) + with open(out, "wb") as f: + for s in range(0, n, 64): + rows = (h[s : s + 64] @ head.T).to(torch.float64).numpy() + bf16_once(rows).tofile(f) + print(f"{n} positions' logits ({mode or 'float64'}) in {time.time() - t0:.0f} s") + + +E2M1 = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], dtype=torch.float64) + + +def nvfp4(packed, scale, global_): + """ModelOpt NVFP4 [n, k] in float64: E2M1 nibbles (low first) x each block's E4M3 scale x the global, exact.""" + n = packed.shape[0] + codes = torch.stack((packed & 15, packed >> 4), -1).reshape(n, -1).long() + blocks = scale.to(torch.float64).repeat_interleave(16, 1) + return E2M1[codes] * blocks * global_.to(torch.float64) + + +def exact(cfg, tensor, ids, points, index): + """The release's forward in float64; with `points`, rounded to bf16 where its bf16 code rounds.""" + f64 = torch.float64 + layers, heads, kv, dim = cfg["num_hidden_layers"], cfg["num_attention_heads"], cfg["num_key_value_heads"], cfg["head_dim"] + eps, theta, window, clamp = cfg["rms_norm_eps"], float(cfg["rope_theta"]), cfg["sliding_window"], cfg["ffn_clamp_threshold"] + n = len(ids) + + def w(name): + if name.endswith(".weight") and f"{name}_scale" in index: # an NVFP4 projection, its exact values + return nvfp4(tensor(name), tensor(f"{name}_scale"), tensor(f"{name}_scale_2")) + return tensor(name).to(f64) + + 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) + + pos = torch.arange(n, dtype=f64) + inv = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=f64) / dim)) + angle = torch.outer(pos, inv) + cos, sin = torch.cat((angle, angle), -1).cos(), torch.cat((angle, angle), -1).sin() + if points: # the release's fp32 angle, its cos and sin rounded to bf16 + inv32 = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + a32 = torch.outer(torch.arange(n, dtype=torch.float32), inv32) + cos, sin = r(torch.cat((a32, a32), -1).cos().to(f64)), r(torch.cat((a32, a32), -1).sin().to(f64)) + key, query = torch.arange(n)[None, :], torch.arange(n)[:, None] + causal = torch.where(key <= query, 0.0, float("-inf")).to(f64) + sliding = torch.where((key <= query) & (key >= query - (window - 1)), 0.0, float("-inf")).to(f64) + scale = cfg["embed_scale"] + x = r(w("model.embed_tokens.weight")[torch.tensor(ids)] * (np.float32(scale) if points else scale)) + v0 = x + t0 = time.time() + for i in range(layers): + p, global_ = f"model.layers.{i}.", cfg["layer_types"][i] == "full_attention" + a = p + "self_attn." + hs = norm(x, w(p + "input_layernorm.weight")) + blend = r(torch.sigmoid(w(a + "value_res_lambda")[0])) + vin = r(torch.lerp(v0, hs, blend)) + q = r(hs @ w(a + "q_proj.weight").T).view(n, heads, dim) + k = r(hs @ w(a + "k_proj.weight").T).view(n, kv, dim) + v = r(vin @ w(a + "v_proj.weight").T).view(n, kv, dim) + q, k = norm(q, w(a + "q_norm.weight")), norm(k, w(a + "k_norm.weight")) + q, k, v = q.transpose(0, 1), k.transpose(0, 1), v.transpose(0, 1) + if not global_: + q, k = r(q * cos + rotate(q) * sin), r(k * cos + rotate(k) * sin) + k, v = k.repeat_interleave(heads // kv, 0), v.repeat_interleave(heads // kv, 0) + att = r(torch.softmax(q @ k.transpose(1, 2) / dim**0.5 + (causal if global_ else sliding), -1) @ v) + att = att.transpose(0, 1).reshape(n, heads * dim) + att = r(att * r(torch.sigmoid(r(hs @ w(a + "g_proj.weight").T)))) # the gate's sigmoid in fp32, then bf16 + x = r(x + norm(r(att @ w(a + "o_proj.weight").T), w(p + "post_attention_layernorm.weight"))) + hs = norm(x, w(p + "pre_feedforward_layernorm.weight")) + gate = r(hs @ w(p + "mlp.gate_proj.weight").T).clamp(max=clamp) + up = r(hs @ w(p + "mlp.up_proj.weight").T).clamp(-clamp, clamp) + act = r(r(torch.nn.functional.silu(gate)) * up) + x = r(x + norm(r(act @ w(p + "mlp.down_proj.weight").T), w(p + "post_feedforward_layernorm.weight"))) + print(f"layer {i + 1} of {layers}, {time.time() - t0:.0f} s", flush=True) + return norm(x, w("model.norm.weight")) + + +def release(model_dir, cfg, tensor, ids, f64=False): + """The checkpoint's own decoder layers and norms in bf16 (or float64), each layer's weights loaded alone.""" + sys.path.insert(0, model_dir) # its modeling file imports configuration_ricursive as a package sibling + pkg = importlib.util.spec_from_file_location("ricursive_release", f"{model_dir}/modeling_ricursive.py", + submodule_search_locations=[model_dir]) + spec_c = importlib.util.spec_from_file_location("ricursive_release.configuration_ricursive", + f"{model_dir}/configuration_ricursive.py") + conf = importlib.util.module_from_spec(spec_c) + sys.modules[spec_c.name] = conf + spec_c.loader.exec_module(conf) + spec_m = importlib.util.spec_from_file_location("ricursive_release.modeling_ricursive", f"{model_dir}/modeling_ricursive.py") + mod = importlib.util.module_from_spec(spec_m) + sys.modules["ricursive_release"] = importlib.util.module_from_spec(pkg) + sys.modules[spec_m.name] = mod + spec_m.loader.exec_module(mod) + config = conf.RicursiveConfig.from_pretrained(model_dir) + bf = torch.float64 if f64 else torch.bfloat16 + n = len(ids) + x = tensor("model.embed_tokens.weight")[torch.tensor(ids)].unsqueeze(0).to(bf) * config.embed_scale + v0 = x + position_ids = torch.arange(n).unsqueeze(0) + t0 = time.time() + for i in range(config.num_hidden_layers): + layer = mod.RicursiveDecoderLayer(config, i).to(bf) + prefix = f"model.layers.{i}." + state = layer_state(tensor, index_names(model_dir, prefix), prefix, f64) + layer.load_state_dict(state, strict=True) + x = layer(x, position_ids=position_ids, value_residual=v0)[0] + print(f"layer {i + 1} of {config.num_hidden_layers}, {time.time() - t0:.0f} s", flush=True) + final = mod.RicursiveRMSNorm(config.hidden_size, eps=config.rms_norm_eps).to(bf) + final.load_state_dict({"weight": tensor("model.norm.weight")}) + return final(x)[0] + + +def layer_state(tensor, names, prefix, f64): + """A layer's weights by module name; in float64 an NVFP4 weight is its exact values and its scales are dropped.""" + have = set(names) + state = {} + for k in names: + if f64 and k.endswith((".weight_scale", ".weight_scale_2", ".input_scale")): + continue + nv = f64 and f"{k}_scale" in have + state[k[len(prefix):]] = nvfp4(tensor(k), tensor(f"{k}_scale"), tensor(f"{k}_scale_2")) if nv else tensor(k) + return state + + +def index_names(model_dir, prefix): + return [k for k in json.load(open(f"{model_dir}/model.safetensors.index.json"))["weight_map"] if k.startswith(prefix)] + + +if __name__ == "__main__": + if len(sys.argv) not in (4, 5) or sys.argv[4:] not in ([], ["--bf16-points"], ["--release"], ["--release-f64"]): + sys.exit(USAGE) + main(*sys.argv[1:4], mode=sys.argv[4] if len(sys.argv) > 4 else "") diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 732a9fda95..14952c7992 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -78,10 +78,12 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu const mods = family(b, target, optimize, cuda, draft_ids); const cli = b.createModule(.{ .root_source_file = b.path("zig/src/cli/cuda_main.zig"), .target = target, .optimize = optimize, .link_libc = true, .strip = strip }); const native = engines(b, target, optimize, cuda, mods.core, mods.lanes, mods.nemotron, mods.heat, mods.qwen35, mods.qwen_vision); - for ([_]struct { []const u8, *std.Build.Module }{ .{ "cuda", cuda }, .{ "core", mods.core }, .{ "lanes", mods.lanes }, .{ "nemotron", mods.nemotron }, .{ "dory", native.engines.import_table.get("dory").? } }) |m| cli.addImport(m[0], m[1]); + for ([_]struct { []const u8, *std.Build.Module }{ .{ "cuda", cuda }, .{ "core", mods.core }, .{ "lanes", mods.lanes }, .{ "nemotron", mods.nemotron }, .{ "dory", native.engines.import_table.get("dory").? }, .{ "ricursive_cuda", native.engines.import_table.get("ricursive_cuda").? } }) |m| cli.addImport(m[0], m[1]); cli.addOptions("build_options", build_options); cli.addImport("checkpoint_cli", checkpointCli(b, target, optimize, strip, native.engines, mods.core)); - b.installArtifact(b.addExecutable(.{ .name = "tensorfold", .root_module = cli })); + const cli_exe = b.addExecutable(.{ .name = "tensorfold", .root_module = cli }); + b.installArtifact(cli_exe); + b.step("compile-cuda-cli", "Compile the CUDA tensorfold CLI without executing a GPU").dependOn(&cli_exe.step); const runner = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/main.zig"), .target = target, .optimize = optimize, .link_libc = true, .strip = strip }); const dump_options = b.addOptions(); dump_options.addOption(bool, "enabled", opts.test_dump); @@ -105,7 +107,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu for (exlr8_build.denseModules(b, target, optimize, exlr8, exlr8_tests[0])) |m| runner.addImport(m.name, m.module); runner.addImport("glm_cuda", glmCuda(b, target, optimize, cuda, mods.core, mods.lanes, mods.heat, mods.api)); @import("minimax_cuda.zig").targets(b, target, optimize, strip, cuda, mods.core, mods.lanes, runner); - for ([_]struct { []const u8, *std.Build.Module }{ .{ "kolibri1", kolibri1Module(b, target, optimize, cuda, mods.core, mods.lanes) }, .{ "dory", @import("dory.zig").module(b, target, optimize, cuda, mods.core, mods.lanes) } }) |m| runner.addImport(m[0], m[1]); + for ([_]struct { []const u8, *std.Build.Module }{ .{ "kolibri1", kolibri1Module(b, target, optimize, cuda, mods.core, mods.lanes) }, .{ "dory", @import("dory.zig").module(b, target, optimize, cuda, mods.core, mods.lanes) }, .{ "ricursive_cuda", @import("ricursive_cuda.zig").module(b, target, optimize, cuda, mods.core, mods.lanes) } }) |m| runner.addImport(m[0], m[1]); runner.addImport("gemma4_cuda", @import("gemma4_cuda.zig").module(b, target, optimize, cuda, mods.core, mods.lanes)); runner.addImport("spark25", @import("spark25.zig").module(b, target, optimize, cuda, mods.core, mods.lanes)); runner.addImport("moss_image_cuda", b.createModule(.{ .root_source_file = b.path("zig/src/families/moss_image/cuda_kernels.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "cuda", .module = cuda }} })); diff --git a/zig/build/cuda_engines.zig b/zig/build/cuda_engines.zig index 0bc47cfb82..c7696389a5 100644 --- a/zig/build/cuda_engines.zig +++ b/zig/build/cuda_engines.zig @@ -56,6 +56,7 @@ pub fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu .{ .name = "gemma4_cuda", .module = @import("gemma4_cuda.zig").module(b, target, optimize, cuda, core, lanes) }, .{ .name = "spark25", .module = @import("spark25.zig").module(b, target, optimize, cuda, core, lanes) }, .{ .name = "dory", .module = @import("dory.zig").module(b, target, optimize, cuda, core, lanes) }, + .{ .name = "ricursive_cuda", .module = @import("ricursive_cuda.zig").module(b, target, optimize, cuda, core, lanes) }, }, }); mod.addImport("qwen4", @import("qwen4.zig").module(b, target, optimize, core)); diff --git a/zig/build/cuda_family_host.zig b/zig/build/cuda_family_host.zig index a0406a1e2e..2f1cd6734a 100644 --- a/zig/build/cuda_family_host.zig +++ b/zig/build/cuda_family_host.zig @@ -7,6 +7,7 @@ pub fn steps(b: *std.Build, cuda: *std.Build.Module, mods: anytype, step: *std.B b.step("test-qwen35-metal-scope-host", "Metal 2B format admission without device or Metal compiler").dependOn(&metal_scope.step); @import("load_prefetch_host.zig").tests(b, mods.core, step); @import("minimax_cuda.zig").hostTests(b, host, cuda, mods.core, mods.lanes, step); + @import("ricursive_cuda.zig").hostTests(b, host, cuda, mods.core, mods.lanes, step); @import("cuda_groups_host.zig").host(b, cuda, mods, step); const prefill = b.addRunArtifact(b.addTest(.{ .root_module = mods.qwen35, .filters = &.{"Qwen stepped prefill:"} })); b.step("test-qwen35-prefill-host", "Interrupted stepped prompt lifecycle without devices").dependOn(&prefill.step); diff --git a/zig/build/cuda_kernels.zig b/zig/build/cuda_kernels.zig index 20de0fa767..fd0cb34125 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -143,6 +143,8 @@ pub const kernels = [_]Kernel{ .{ .name = "glm_index", .flags = &.{"-O3"} }, // a prompt chunk's DSA scores and top-512 pools .{ .name = "minimax_m3", .flags = glue }, // ours: MiniMax-M3.1-Flash's glue (norms, RoPE prep, routing, DSpark, tower RoPE) .{ .name = "minimax_m3_attn", .flags = glue }, // ours: MiniMax-M3.1-Flash's indexer and block-sparse attention + .{ .name = "ricursive", .flags = glue }, // ours: Ripple's glue (embed, add-norms, output gate, FFN activation), glue.metal step for step + .{ .name = "ricursive_prep", .flags = glue }, // ours: Ripple's attention prep (q/k head norms, RoPE, K and V into rings or pages) .{ .name = "glm_quant", .flags = torch_ops }, // ours: GLM's draft-only weight quantizers, torch at load in the Python family .{ .name = "exl3_linear", .flags = &.{ "-O3", "--expt-relaxed-constexpr" } }, // exl3/linear.cu's device code, tensorfold_exl3_linear_v5 .{ .name = "exl3_experts", .flags = &.{ "-O3", "--expt-relaxed-constexpr" } }, // exl3/experts.cu's device code, tensorfold_exl3_experts diff --git a/zig/build/metal_targets.zig b/zig/build/metal_targets.zig index 35352c3e9d..8fb36a2d0c 100644 --- a/zig/build/metal_targets.zig +++ b/zig/build/metal_targets.zig @@ -344,6 +344,7 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s .{ .name = "tf-minimax-run", .path = "zig/tests/minimax_m3/run.zig", .about = "MiniMax-M3.1-Flash greedy tokens, every position's logits or a prompt's last logits, for the reference's check" }, .{ .name = "tf-minimax-vision", .path = "zig/tests/minimax_m3/vision.zig", .about = "MiniMax-M3.1-Flash's image tower alone: an image's patches and text rows, for tools/vision_ref.py'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-ricursive-run", .path = "zig/tests/ricursive/run.zig", .about = "Ripple on Metal: greedy tokens, every position's logits, a prompt's last logits or forced steps, 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" }, }; // the tokenizer for the programs that tokenize prompts themselves (the engine module has none) diff --git a/zig/build/ricursive_cuda.zig b/zig/build/ricursive_cuda.zig new file mode 100644 index 0000000000..d2a7344698 --- /dev/null +++ b/zig/build/ricursive_cuda.zig @@ -0,0 +1,18 @@ +//! Ripple's CUDA backend in the build: its module for tf-cuda-test and the native server, and its host regressions. +const std = @import("std"); + +/// The Ripple CUDA family over the runtime, the core and the lane core. +pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module) *std.Build.Module { + const m = b.createModule(.{ .root_source_file = b.path("zig/src/families/ricursive/cuda.zig"), .target = target, .optimize = optimize, .link_libc = true }); + m.addImport("cuda", cuda); + m.addImport("core", core); + m.addImport("lanes", lanes); + return m; +} + +/// The family's host regressions (launch geometry, layout, names, packing, admission) in `all`. +pub fn hostTests(b: *std.Build, host: std.Build.ResolvedTarget, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, all: *std.Build.Step) void { + const run = b.addRunArtifact(b.addTest(.{ .root_module = module(b, host, .debug, cuda, core, lanes) })); + b.step("test-ricursive-cuda-host", "Ripple's CUDA backend host regressions (no GPU work)").dependOn(&run.step); + all.dependOn(&run.step); +} diff --git a/zig/kernels/cuda/contract_controls.cu b/zig/kernels/cuda/contract_controls.cu index 2970dfa23c..5aacff3bcc 100644 --- a/zig/kernels/cuda/contract_controls.cu +++ b/zig/kernels/cuda/contract_controls.cu @@ -74,6 +74,27 @@ extern "C" __global__ void __launch_bounds__(tf_qk_round::WIDE) tf_ctl_qk_round( tf_qk_round::rows(in, nw, cosv, sinv, inv, rot, pos, slot, dst, kc, vc, len, h, hk, eps, keep); } +// Gate 3: ricursive_prep.cu's tf_rc_prep keeping only the rounding points in `keep`; the shapes gate must go red. +#include "ricursive_prep.cuh" + +extern "C" __global__ void __launch_bounds__(tf_rcp::NT) tf_ctl_rc_prep( + const tf_rcp::Rows XQ, const tf_rcp::Rows XK, const tf_rcp::Rows XV, const float* inv, const __nv_bfloat16* qw, + const __nv_bfloat16* kw, __nv_bfloat16* Q, __nv_bfloat16* KC, __nv_bfloat16* VC, int rows, int heads, int kv_heads, int pos, + int ring, int rope, float eps, const int* tab, long long page, int tok, int keep) { + tf_rcp::prep(XQ, XK, XV, inv, qw, kw, Q, KC, VC, rows, heads, kv_heads, pos, ring, rope, eps, tab, page, tok, keep); +} + +// Gate 3: tf_rc_prep_rows with each row's slot dropped (all rows in stream 0's cache), `keep` as above; shapes go red. +extern "C" __global__ void __launch_bounds__(tf_rcp::NT) tf_ctl_rc_prep_rows( + const tf_rcp::Rows XQ, const tf_rcp::Rows XK, const tf_rcp::Rows XV, const float* inv, const __nv_bfloat16* qw, + const __nv_bfloat16* kw, __nv_bfloat16* Q, __nv_bfloat16* KC, __nv_bfloat16* VC, int rows, int heads, int kv_heads, + const int* pos, const int* slot, int ring, int rope, float eps, const int* tab, long long page, int tok, int stride, + int keep) { + (void)slot; // the planted fault: no row's own stream + tf_rcp::prep_rows(XQ, XK, XV, inv, qw, kw, Q, KC, VC, rows, heads, kv_heads, pos, nullptr, ring, rope, eps, tab, page, + tok, stride, keep); +} + // Gate 3: act_rows.cu's tf_softcap_rows keeping only the rounding points in `keep`; the exact check must go red. #include "act_round.cuh" diff --git a/zig/kernels/cuda/ricursive.cu b/zig/kernels/cuda/ricursive.cu new file mode 100644 index 0000000000..c6b35dc64f --- /dev/null +++ b/zig/kernels/cuda/ricursive.cu @@ -0,0 +1,61 @@ +// Ripple's glue on CUDA: glue.metal step for step (bf16 residual, fp32 sums), and its planted controls. + +#include "ricursive.cuh" + +namespace tf_rc { + +// Embed: grid (ceil(width / 256), rows), 256 threads; ids are u32. +#define TF_RC_EMBED(NAME, BF_SCALE) \ + extern "C" __global__ void __launch_bounds__(NT) NAME(const unsigned* __restrict__ ids, const bf16* __restrict__ E, \ + bf16* __restrict__ R, bf16* __restrict__ V0, int rows, int width, \ + float scale) { \ + embed(ids, E, R, V0, rows, width, scale); \ + } + +// Add-norm: grid rows, 256 threads; Y null skips the add, V0 null the lerp; width % 8 == 0, <= 8,192, 16-byte rows. +#define TF_RC_ADD_NORM(NAME, K) \ + extern "C" __global__ void __launch_bounds__(NT) \ + NAME(bf16* __restrict__ R, const bf16* __restrict__ Y, const bf16* __restrict__ w_post, \ + const bf16* __restrict__ w_next, const bf16* __restrict__ V0, bf16* __restrict__ X, bf16* __restrict__ VX, \ + int rows, int width, float eps, float blend) { \ + add_norm(R, Y, w_post, w_next, V0, X, VX, rows, width, eps, blend); \ + } + +// Gate: grid (ceil(width / 256), rows), 256 threads, A in place. +#define TF_RC_GATE(NAME, K) \ + extern "C" __global__ void __launch_bounds__(NT) NAME(bf16* __restrict__ A, const bf16* __restrict__ G, int rows, \ + int width) { \ + gate(A, G, rows, width); \ + } + +// SiLU-clamp: grid (ceil(width / 256), rows), 256 threads; GU rows are 2 x width (gate | up). +#define TF_RC_SILU_CLAMP(NAME, K) \ + extern "C" __global__ void __launch_bounds__(NT) NAME(const bf16* __restrict__ GU, bf16* __restrict__ A, int rows, \ + int width, float limit) { \ + silu_clamp(GU, A, rows, width, limit); \ + } + +TF_RC_EMBED(tf_rc_embed, false) +TF_RC_ADD_NORM(tf_rc_add_norm, NORM_ALL) +TF_RC_GATE(tf_rc_gate, SIG) +TF_RC_SILU_CLAMP(tf_rc_silu_clamp, SILU | LOW) + +// Controls, one a rounding point (each must fail the contract's shapes gate): the served kernel without it. +TF_RC_EMBED(tf_rc_ctl_embed_bf16_scale, true) +TF_RC_ADD_NORM(tf_rc_ctl_add_norm_y_row, NORM_ALL & ~Y_ROW) +TF_RC_ADD_NORM(tf_rc_ctl_add_norm_y_out, NORM_ALL & ~Y_OUT) +TF_RC_ADD_NORM(tf_rc_ctl_add_norm_resid, NORM_ALL & ~RESID) +TF_RC_ADD_NORM(tf_rc_ctl_add_norm_x_row, NORM_ALL & ~X_ROW) +TF_RC_ADD_NORM(tf_rc_ctl_add_norm_x_lerp, NORM_ALL & ~X_LERP) +TF_RC_GATE(tf_rc_ctl_gate_sigmoid, 0) +TF_RC_SILU_CLAMP(tf_rc_ctl_silu_round, LOW) +TF_RC_SILU_CLAMP(tf_rc_ctl_silu_no_lower, SILU) + +#undef TF_RC_EMBED +#undef TF_RC_ADD_NORM +#undef TF_RC_GATE +#undef TF_RC_SILU_CLAMP + +// Plug-in point: the prep (per-head norm, RoPE, KV-ring write) and attention kernels go below, named tf_rc_*. + +} // namespace tf_rc diff --git a/zig/kernels/cuda/ricursive.cuh b/zig/kernels/cuda/ricursive.cuh new file mode 100644 index 0000000000..8ac92f8030 --- /dev/null +++ b/zig/kernels/cuda/ricursive.cuh @@ -0,0 +1,163 @@ +// Ripple's CUDA glue: glue.metal's rc_embed, rc_add_norm, rc_gate and rc_silu_clamp at the same rounding points. +#pragma once + +#include +#include +#include + +#include "common/act.cuh" +#include "common/reduce.cuh" + +namespace tf_rc { + +typedef __nv_bfloat16 bf16; + +constexpr int NT = 256; // threads a block in every glue launch +constexpr int MOST = 4; // 8-value chunks a thread holds in add_norm: rows up to 8,192 wide + +__device__ __forceinline__ float bf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } +__device__ __forceinline__ float ld(const bf16* p) { return __bfloat162float(*p); } + +// A rounding point: bf16 when K keeps P; a control's K drops it. +template +__device__ __forceinline__ float point(float x) { return (K & P) ? bf(x) : x; } + +// add_norm's rounding points, a bit each; a control clears one. +constexpr int Y_ROW = 1; // bf16(y * inv): Y's norm before w_post +constexpr int Y_OUT = 2; // bf16(w_post * that): the normed branch +constexpr int RESID = 4; // bf16(r + normed): the residual the next norm reads +constexpr int X_ROW = 8; // bf16(r * inv): R's norm before w_next +constexpr int X_LERP = 16; // the lerp reads X as stored (bf16), not the fp32 product +constexpr int NORM_ALL = Y_ROW | Y_OUT | RESID | X_ROW | X_LERP; + +constexpr int SIG = 1; // gate: bf16(sigmoid(g)), the release's fp32 sigmoid then bf16 +constexpr int SILU = 1; // silu_clamp: bf16(silu(g)) +constexpr int LOW = 2; // silu_clamp: up's lower clamp (not a rounding point; the Metal family's control drops it) + +// R and V0 row r = bf16(E[ids[r]] * scale), the fp32 scale; BF_SCALE (a control) rounds the scale to bf16 first. +template +__device__ __forceinline__ void embed(const unsigned* __restrict__ ids, const bf16* __restrict__ E, bf16* __restrict__ R, + bf16* __restrict__ V0, int rows, int width, float scale) { + const int x = blockIdx.x * blockDim.x + threadIdx.x, row = blockIdx.y; + if (x >= width || row >= rows) return; + const size_t at = (size_t)row * width + x; + const bf16 v = __float2bfloat16_rn(__fmul_rn(ld(E + (size_t)ids[row] * width + x), BF_SCALE ? bf(scale) : scale)); + R[at] = v; + V0[at] = v; +} + +// torch.lerp on bf16 (its vectorized form): fma(w, e - s, s) below one half, else fma(w - 1, e - s, e); unrounded. +__device__ __forceinline__ float lerp(float s, float e, float w) { + const float d = __fsub_rn(e, s); + return fabsf(w) < 0.5f ? __fmaf_rn(w, d, s) : __fmaf_rn(__fsub_rn(w, 1.0f), d, e); +} + +// A thread's chunks t + j NT of a row's 8-value chunks, zeros past the row. +__device__ __forceinline__ void load_row(const bf16* __restrict__ p, int chunks, float (&v)[MOST][8]) { + const uint4* q = reinterpret_cast(p); +#pragma unroll + for (int j = 0; j < MOST; ++j) { + const int c = threadIdx.x + j * NT; + if (c < chunks) { + tf_common::unpack8(q[c], v[j]); + } else { +#pragma unroll + for (int e = 0; e < 8; ++e) v[j][e] = 0.0f; + } + } +} + +// A row's fp32 sum of squares (each thread's fma chain, then the block's tree) into 1 / sqrt(mean + eps). +__device__ __forceinline__ float inv_rms(const float (&v)[MOST][8], float* part, int width, float eps) { + float s = 0.0f; +#pragma unroll + for (int j = 0; j < MOST; ++j) +#pragma unroll + for (int e = 0; e < 8; ++e) s = __fmaf_rn(v[j][e], v[j][e], s); + return tf_common::rms_scale(tf_common::block_sum_tree(s, part, NT / 32), width, eps); +} + +// A row a block, R in place: R = bf16(R + bf16(w_post bf16(Y inv))); X = bf16(w_next bf16(R inv)); VX = lerp(V0, X). +template +__device__ __forceinline__ void add_norm(bf16* __restrict__ R, const bf16* __restrict__ Y, const bf16* __restrict__ w_post, + const bf16* __restrict__ w_next, const bf16* __restrict__ V0, bf16* __restrict__ X, + bf16* __restrict__ VX, int rows, int width, float eps, float blend) { + __shared__ float part[NT / 32]; + const int row = blockIdx.x, chunks = width / 8; + if (row >= rows) return; + const size_t base = (size_t)row * width; + float r[MOST][8]; + load_row(R + base, chunks, r); + if (Y != nullptr) { + float y[MOST][8]; + load_row(Y + base, chunks, y); + const float inv = inv_rms(y, part, width, eps); + const uint4* w4 = reinterpret_cast(w_post); + uint4* out = reinterpret_cast(R + base); +#pragma unroll + for (int j = 0; j < MOST; ++j) { + const int c = threadIdx.x + j * NT; + if (c >= chunks) continue; + float w[8]; + tf_common::unpack8(w4[c], w); +#pragma unroll + for (int e = 0; e < 8; ++e) { + const float normed = point(__fmul_rn(w[e], point(__fmul_rn(y[j][e], inv)))); + r[j][e] = point(__fadd_rn(r[j][e], normed)); + } + out[c] = tf_common::pack8(r[j]); + } + } + const float inv = inv_rms(r, part, width, eps); + const uint4* w4 = reinterpret_cast(w_next); +#pragma unroll + for (int j = 0; j < MOST; ++j) { + const int c = threadIdx.x + j * NT; + if (c >= chunks) continue; + float w[8], x[8]; + tf_common::unpack8(w4[c], w); +#pragma unroll + for (int e = 0; e < 8; ++e) x[e] = __fmul_rn(w[e], point(__fmul_rn(r[j][e], inv))); + const uint4 xs = tf_common::pack8(x); + reinterpret_cast(X + base)[c] = xs; + if (V0 == nullptr) continue; + float s[8], e8[8]; + tf_common::unpack8(reinterpret_cast(V0 + base)[c], s); + if (K & X_LERP) { + tf_common::unpack8(xs, e8); + } else { +#pragma unroll + for (int e = 0; e < 8; ++e) e8[e] = x[e]; + } +#pragma unroll + for (int e = 0; e < 8; ++e) s[e] = lerp(s[e], e8[e], blend); + reinterpret_cast(VX + base)[c] = tf_common::pack8(s); + } +} + +// The attention's output gate in place: A = bf16(A bf16(sigmoid(G))), G the bf16 g_proj rows. +template +__device__ __forceinline__ void gate(bf16* __restrict__ A, const bf16* __restrict__ G, int rows, int width) { + const int x = blockIdx.x * blockDim.x + threadIdx.x, row = blockIdx.y; + if (x >= width || row >= rows) return; + const size_t i = (size_t)row * width + x; + const float s = point(tf_common::sigmoid_rn(ld(G + i))); + A[i] = __float2bfloat16_rn(__fmul_rn(ld(A + i), s)); +} + +// The FFN on a stacked gate | up row: gate at most limit, up within +-limit, then bf16(bf16(silu(g)) u). +template +__device__ __forceinline__ void silu_clamp(const bf16* __restrict__ GU, bf16* __restrict__ A, int rows, int width, float limit) { + const int x = blockIdx.x * blockDim.x + threadIdx.x, row = blockIdx.y; + if (x >= width || row >= rows) return; + const size_t at = (size_t)row * 2 * width + x; + const float g0 = ld(GU + at), u0 = ld(GU + at + width); + // a comparison is false for a NaN, so it passes both clamps as torch.clamp keeps it + const float g = g0 > limit ? limit : g0; + const float hi = u0 > limit ? limit : u0; + const float u = (K & LOW) && u0 < -limit ? -limit : hi; + const float silu = point(tf_common::silu_rn(g)); + A[(size_t)row * width + x] = __float2bfloat16_rn(__fmul_rn(silu, u)); +} + +} // namespace tf_rc diff --git a/zig/kernels/cuda/ricursive_prep.cu b/zig/kernels/cuda/ricursive_prep.cu new file mode 100644 index 0000000000..b15f6366f5 --- /dev/null +++ b/zig/kernels/cuda/ricursive_prep.cu @@ -0,0 +1,23 @@ +// Ripple's attention prep on CUDA: the Metal family's rc_prep step for step, every rounding point kept. + +#include "ricursive_prep.cuh" + +// q and k heads normed (w, not 1 + w), NeoX-rotated if rope != 0; k and v at ring slot (pos + row) % ring, or in pages. +extern "C" __global__ void __launch_bounds__(tf_rcp::NT) + tf_rc_prep(const tf_rcp::Rows XQ, const tf_rcp::Rows XK, const tf_rcp::Rows XV, const float* __restrict__ inv, + const __nv_bfloat16* __restrict__ qw, const __nv_bfloat16* __restrict__ kw, __nv_bfloat16* __restrict__ Q, + __nv_bfloat16* __restrict__ KC, __nv_bfloat16* __restrict__ VC, int rows, int heads, int kv_heads, int pos, + int ring, int rope, float eps, const int* __restrict__ tab, long long page, int tok) { + tf_rcp::prep(XQ, XK, XV, inv, qw, kw, Q, KC, VC, rows, heads, kv_heads, pos, ring, rope, eps, tab, page, tok, tf_rcp::ALL); +} + +// tf_rc_prep with row r at pos[r] of stream slot[r] (device int32, so graphs replay new values); tf_rc_prep's bits. +extern "C" __global__ void __launch_bounds__(tf_rcp::NT) + tf_rc_prep_rows(const tf_rcp::Rows XQ, const tf_rcp::Rows XK, const tf_rcp::Rows XV, const float* __restrict__ inv, + const __nv_bfloat16* __restrict__ qw, const __nv_bfloat16* __restrict__ kw, + __nv_bfloat16* __restrict__ Q, __nv_bfloat16* __restrict__ KC, __nv_bfloat16* __restrict__ VC, int rows, + int heads, int kv_heads, const int* __restrict__ pos, const int* __restrict__ slot, int ring, int rope, + float eps, const int* __restrict__ tab, long long page, int tok, int stride) { + tf_rcp::prep_rows(XQ, XK, XV, inv, qw, kw, Q, KC, VC, rows, heads, kv_heads, pos, slot, ring, rope, eps, tab, page, tok, + stride, tf_rcp::ALL); +} diff --git a/zig/kernels/cuda/ricursive_prep.cuh b/zig/kernels/cuda/ricursive_prep.cuh new file mode 100644 index 0000000000..f60864a9e5 --- /dev/null +++ b/zig/kernels/cuda/ricursive_prep.cuh @@ -0,0 +1,119 @@ +// Ripple's attention prep (Metal ricursive/glue.metal's rc_prep) step for step; `keep`: its bf16 rounding points. +#pragma once + +#include +#include +#include "common/reduce.cuh" + +// tf_rcp, not tf_rc_prep: the kernel's extern "C" symbol takes that name. +namespace tf_rcp { + +typedef __nv_bfloat16 bf16; + +// Head dim 128, a lane's dims lane + 32 j; a warp a (head, row), 8 warps a block. +constexpr int D = 128, HALF = 64, PER = D / 32, WARPS = 8, NT = 32 * WARPS; + +// rc_prep's rounding points: x r, w times that, cos, sin, and the stored output (dropped: truncated, not rounded). +enum : int { NORM_ROW = 1, NORM_OUT = 2, COS = 4, SIN = 8, OUT = 16 }; +constexpr int ALL = 31; + +__device__ __forceinline__ float bf(float v) { return __bfloat162float(__float2bfloat16_rn(v)); } + +// v rounded to bf16 where `keep` holds the point p (a constant in the served kernel, so the test folds away). +__device__ __forceinline__ float point(float v, int keep, int p) { return (keep & p) != 0 ? bf(v) : v; } + +// The output as stored: to nearest even, or its fp32 bits' top half where `keep` drops OUT. +__device__ __forceinline__ bf16 store(float v, int keep) { + return (keep & OUT) != 0 ? __float2bfloat16_rn(v) : __float2bfloat16_rz(v); +} + +// One projection's output rows: its first row and the elements between rows (V is its own projection). +struct Rows { + const bf16* x; + int stride; + __device__ __forceinline__ const bf16* at(int row, int head) const { return x + (size_t)row * stride + (size_t)head * D; } +}; + +// A warp's unit of row `row` at position p: query head `unit` into Q, else key head unit - heads and its value at `at`. +__device__ __forceinline__ void unit_at(const Rows XQ, const Rows XK, const Rows XV, const float* __restrict__ inv, + const bf16* __restrict__ qw, const bf16* __restrict__ kw, bf16* __restrict__ Q, + bf16* __restrict__ KC, bf16* __restrict__ VC, int heads, int unit, int lane, int row, + int p, size_t at, int rope, float eps, int keep) { + const bool q = unit < heads; + const int head = q ? unit : unit - heads; + const bf16* src = q ? XQ.at(row, head) : XK.at(row, head); + const bf16* w = q ? qw : kw; + bf16* y = q ? Q + ((size_t)row * heads + head) * D : KC + at + (size_t)head * D; + float x[PER]; + float sum = 0.0f; // bf16(w bf16(x r)), r = rsqrt(sum x^2 / 128 + eps); intrinsics and fma: the same under any flags +#pragma unroll + for (int j = 0; j < PER; ++j) { + x[j] = __bfloat162float(src[lane + 32 * j]); + sum = fmaf(x[j], x[j], sum); + } + const float r = rsqrtf(__fadd_rn(__fdiv_rn(tf_common::warp_sum(sum), (float)D), eps)); +#pragma unroll + for (int j = 0; j < PER; ++j) { + const float wv = __bfloat162float(w[lane + 32 * j]); + x[j] = point(__fmul_rn(wv, point(__fmul_rn(x[j], r), keep, NORM_ROW)), keep, NORM_OUT); + } + const float position = (float)p; +#pragma unroll + for (int j = 0; j < PER; ++j) { + const int d = lane + 32 * j; + float v = x[j]; + if (rope != 0) { + const float angle = __fmul_rn(position, inv[d % HALF]); + const float c = point(cosf(angle), keep, COS), s = point(sinf(angle), keep, SIN); + const float other = j < 2 ? -x[j ^ 2] : x[j ^ 2]; // d < 64 pairs with d + 64 (negated), else with d - 64 + v = __fadd_rn(__fmul_rn(x[j], c), __fmul_rn(other, s)); + } + y[d] = store(v, keep); + } + if (q) return; + const bf16* vs = XV.at(row, head); + bf16* vd = VC + at + (size_t)head * D; +#pragma unroll + for (int j = 0; j < PER; ++j) vd[lane + 32 * j] = vs[lane + 32 * j]; // the value as it is +} + +// One stream's rows at pos + row: keys and values at ring slot (pos + row) % ring, or row p % tok of page tab[p / tok]. +__device__ __forceinline__ void prep(const Rows XQ, const Rows XK, const Rows XV, const float* __restrict__ inv, + const bf16* __restrict__ qw, const bf16* __restrict__ kw, bf16* __restrict__ Q, + bf16* __restrict__ KC, bf16* __restrict__ VC, int rows, int heads, int kv_heads, + int pos, int ring, int rope, float eps, const int* __restrict__ tab, long long page, + int tok, int keep) { + const int unit = blockIdx.x * (blockDim.x >> 5) + (threadIdx.x >> 5); + const int lane = threadIdx.x & 31; + const int row = blockIdx.y; + if (unit >= heads + kv_heads || row >= rows) return; // whole warps: unit_at's shuffles stay converged + const size_t kvw = (size_t)kv_heads * D; + const int p = pos + row; + // the key's cache row: its ring slot, or its page's row (an address only: the numerics are the ring's) + const size_t at = tab != nullptr ? (size_t)tab[p / tok] * (size_t)page + (size_t)(p % tok) * kvw + : (size_t)(p % ring) * kvw; + unit_at(XQ, XK, XV, inv, qw, kw, Q, KC, VC, heads, unit, lane, row, p, at, rope, eps, keep); +} + +// Row r at pos[r] of stream slot[r] (null: stream 0): ring row slot * ring + p % ring, or tab[slot * stride + p / tok]. +__device__ __forceinline__ void prep_rows(const Rows XQ, const Rows XK, const Rows XV, const float* __restrict__ inv, + const bf16* __restrict__ qw, const bf16* __restrict__ kw, bf16* __restrict__ Q, + bf16* __restrict__ KC, bf16* __restrict__ VC, int rows, int heads, int kv_heads, + const int* __restrict__ pos, const int* __restrict__ slot, int ring, int rope, + float eps, const int* __restrict__ tab, long long page, int tok, int stride, + int keep) { + const int unit = blockIdx.x * (blockDim.x >> 5) + (threadIdx.x >> 5); + const int lane = threadIdx.x & 31; + const int row = blockIdx.y; + if (unit >= heads + kv_heads || row >= rows) return; // whole warps: unit_at's shuffles stay converged + const size_t kvw = (size_t)kv_heads * D; + const int p = pos[row]; + const size_t s = slot != nullptr ? (size_t)slot[row] : 0; + // the key's cache row: its stream's ring slot, or its page's row by the stream's row of the table + const size_t at = tab != nullptr + ? (size_t)tab[s * (size_t)stride + (size_t)(p / tok)] * (size_t)page + (size_t)(p % tok) * kvw + : (s * (size_t)ring + (size_t)(p % ring)) * kvw; + unit_at(XQ, XK, XV, inv, qw, kw, Q, KC, VC, heads, unit, lane, row, p, at, rope, eps, keep); +} + +} // namespace tf_rcp diff --git a/zig/kernels/cuda/window_attention.cu b/zig/kernels/cuda/window_attention.cu index c90859146d..95764803eb 100644 --- a/zig/kernels/cuda/window_attention.cu +++ b/zig/kernels/cuda/window_attention.cu @@ -574,6 +574,10 @@ template __global__ void tf_window_prompt::wattn_kernel<128, 8, 4, 8>(const __nv template __global__ void tf_window_prompt::wattn_kernel<128, 8, 8, 8>(const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, __nv_bfloat16* __restrict__, int, int, int, int, int, float, int, int); +// Ripple's group 10 (no multiple of 4): two query heads a block, so 64 rows a block; each warp's arithmetic as above. +template __global__ void tf_window_prompt::wattn_kernel<128, 8, 2, 8>(const __nv_bfloat16* __restrict__, + const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, __nv_bfloat16* __restrict__, int, int, int, + int, int, float, int, int); // Kolibri-1's full layers' prompt rows over pages: head dim 128, eight warps, four or eight query heads a block. #define TF_WATTN_PG(NAME, HPC) \ @@ -586,3 +590,5 @@ template __global__ void tf_window_prompt::wattn_kernel<128, 8, 8, 8>(const __nv TF_WATTN_PG(tf_window_prompt_h4_pg, 4) TF_WATTN_PG(tf_window_prompt_h8_pg, 8) +// Ripple's global layers over pages: two query heads a block. +TF_WATTN_PG(tf_window_prompt_h2_pg, 2) diff --git a/zig/kernels/metal/ricursive/attention_nax.metal b/zig/kernels/metal/ricursive/attention_nax.metal new file mode 100644 index 0000000000..529e200283 --- /dev/null +++ b/zig/kernels/metal/ricursive/attention_nax.metal @@ -0,0 +1,204 @@ +// Ripple's prompt attention on the tensor units: llama's kernel over a row's window, keys at slot position % ring. +#include "../nax.h" +using namespace tfp; + +// Defines: DH (128), NQ, G (query heads a key head), KVH, SCALE. P: rows, p0, window, ring (a multiple of BKT). + +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, O [R, NQ, DH] from p0; K, V the cache; a tile RT rows of a head, each seeing its window (global: every key). +[[max_total_threads_per_threadgroup(256)]] +[[kernel]] void rc_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 int4& 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, window = P.z, ring = P.w; + 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 = max(0, first_pos - (window - 1)) / BKT; kb <= last_pos / BKT; kb++) { + const int kpos = kb * BKT + int(sg) * KS; + const device bfloat* kblk = kc + long(kpos % ring) * 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 || col < row - (window - 1) ? 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 % ring) * 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/ricursive/fp4.metal b/zig/kernels/metal/ricursive/fp4.metal new file mode 100644 index 0000000000..2b03aeae94 --- /dev/null +++ b/zig/kernels/metal/ricursive/fp4.metal @@ -0,0 +1,86 @@ +// Ripple's NVFP4 MLPs (ModelOpt): E2M1 nibbles low first, E4M3 scale per 16, fp32 global; fp4 * scale exact in bf16. +#include +using namespace metal; + +#define TF_UNROLL_4 _Pragma("clang loop unroll(full)") + +struct rc_fp4_args { uint rows, n, k, ldy; float global; }; + +// E2M1: sign, two exponent bits, one mantissa bit. +inline float rc_e2m1(uint code) { + const uint m = code & 7u; + const float v = m < 4u ? float(m) * 0.5f : (m == 4u ? 2.0f : m == 5u ? 3.0f : m == 6u ? 4.0f : 6.0f); + return (code & 8u) ? -v : v; +} + +// Unsigned E4M3 (ModelOpt's scales are positive): exponent bias 7, three mantissa bits, 0x7f NaN as torch reads it. +inline float rc_e4m3(uint code) { + const uint e = (code >> 3) & 15u, m = code & 7u; + if ((code & 0x7fu) == 0x7fu) return NAN; + return e ? ldexp(1.0f + float(m) * 0.125f, int(e) - 7) : float(m) * 0.001953125f; +} + +// The 16 weights of block `b` in row j: two words of nibbles, each times the block's scale, exact in fp32 and bf16. +inline void rc_fp4_block(const device uint2* W, const device uchar* S, size_t at, thread float4 (&w)[4]) { + const uint2 q = W[at]; + const float s = rc_e4m3(S[at]); + TF_UNROLL_4 + for (uint c = 0; c < 4; ++c) { + const uint word = c < 2 ? q.x : q.y, sh = 16u * (c & 1u); + w[c] = float4(rc_e2m1(word >> sh), rc_e2m1(word >> (sh + 4u)), rc_e2m1(word >> (sh + 8u)), rc_e2m1(word >> (sh + 12u))) * s; + } +} + +// Y[r, j] = bf16(global * sum_i X[r, i] * W[j, i]) in fp32: a simdgroup an output j, a lane a block of 16 at a time. +template +inline void rc_fp4_matvec_t(const device bfloat* X, const device uint2* W, const device uchar* S, + constant rc_fp4_args& a, device bfloat* Y, uint j, uint lane) { + const uint blocks = a.k / 16; + float acc[ROWS]; + TF_UNROLL_4 + for (uint r = 0; r < ROWS; ++r) acc[r] = 0.0f; + for (uint b = lane; b < blocks; b += 32) { + float4 w[4]; + rc_fp4_block(W, S, (size_t)j * blocks + b, w); + TF_UNROLL_4 + for (uint r = 0; r < ROWS; ++r) { + const device bfloat4* x = (const device bfloat4*)(X + (size_t)r * a.k + b * 16); + float t = 0.0f; + TF_UNROLL_4 + for (uint c = 0; c < 4; ++c) t += dot(w[c], float4(x[c])); + acc[r] += t; + } + } + TF_UNROLL_4 + for (uint r = 0; r < ROWS; ++r) { + const float sum = simd_sum(acc[r]); + if (lane == 0 && r < a.rows) Y[(size_t)r * a.ldy + j] = bfloat(sum * a.global); + } +} + +#define RC_FP4_MATVEC(ROWS) \ + [[kernel]] void rc_fp4_matvec_##ROWS(const device bfloat* X [[buffer(0)]], const device uint2* W [[buffer(1)]], \ + const device uchar* S [[buffer(2)]], constant rc_fp4_args& a [[buffer(3)]], device bfloat* Y [[buffer(4)]], \ + uint3 tg [[threadgroup_position_in_grid]], uint sg [[simdgroup_index_in_threadgroup]], \ + uint lane [[thread_index_in_simdgroup]]) { \ + const uint j = tg.x * 8 + sg; \ + if (j < a.n) rc_fp4_matvec_t(X, W, S, a, Y, j, lane); \ + } +RC_FP4_MATVEC(1) +RC_FP4_MATVEC(2) +RC_FP4_MATVEC(4) +RC_FP4_MATVEC(8) +RC_FP4_MATVEC(16) +RC_FP4_MATVEC(32) + +// W as bf16 [n, k] without its global, for a prompt chunk's GEMM: a thread a block of 16, each value exact. +[[kernel]] void rc_fp4_dequant(const device uint2* W [[buffer(0)]], const device uchar* S [[buffer(1)]], + constant rc_fp4_args& a [[buffer(2)]], device bfloat4* D [[buffer(3)]], uint2 p [[thread_position_in_grid]]) { + const uint blocks = a.k / 16; + if (p.x >= blocks || p.y >= a.n) return; + const size_t at = (size_t)p.y * blocks + p.x; + float4 w[4]; + rc_fp4_block(W, S, at, w); + TF_UNROLL_4 + for (uint c = 0; c < 4; ++c) D[at * 4 + c] = bfloat4(w[c]); +} diff --git a/zig/kernels/metal/ricursive/fp4_gemm.metal b/zig/kernels/metal/ricursive/fp4_gemm.metal new file mode 100644 index 0000000000..409c9fadc7 --- /dev/null +++ b/zig/kernels/metal/ricursive/fp4_gemm.metal @@ -0,0 +1,53 @@ +// NVFP4 prompt projections after gemm_nax.metal: x by the exact dequantized W, fp32 sums times the global, then bf16. + +namespace tfp { +template +inline void rc_fp4_gemm_t(const device bfloat16_t* A, const device bfloat16_t* B, device bfloat16_t* D, + const device int* P, float global, uint3 tg, uint sg, uint lane) { + constexpr int BM = 64, BN = 128, BK = 256, WN = 4; + const int swz = P[8]; + const int ty = (int(tg.y) << swz) + (int(tg.x) & ((1 << swz) - 1)), tx = int(tg.x) >> swz; + if (tx >= P[6] || ty >= P[7]) { + return; + } + const int lda = P[3], ldb = P[4], ldd = P[5]; + const int row = ty * BM + 32 * int(sg / WN), col = tx * BN + 32 * int(sg % WN); + A += long(P[13]) + long(row) * lda; + B += long(P[14]) + long(col) * ldb; + D += long(row) * ldd + col; + const int rows = AM ? 32 : min(32, P[0] - row), cols = AN ? 32 : min(32, P[1] - col); + const bool full = (AM || rows == 32) && (AN || cols == 32); + const short2 home = frag_home(ushort(lane)); + frag acc[2][2]; + if (full) { + gemm_block(acc, A, B, lda, ldb, P[2], P[9], rows, cols, home); + } else { + gemm_block(acc, A, B, lda, ldb, P[2], P[9], rows, cols, home); + } + TF_UNROLL + for (short i = 0; i < 2; i++) { + TF_UNROLL + for (short j = 0; j < 2; j++) { + acc[i][j] *= global; + } + } + if (full) { + gemm_store(acc, D, ldd, rows, cols, home); + } else if (rows > 0 && cols > 0) { + gemm_store(acc, D, ldd, rows, cols, home); + } +} +} // namespace tfp + +#define RC_FP4_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)]], constant float& g [[buffer(4)]], \ + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], \ + uint3 tg [[threadgroup_position_in_grid]]) { \ + tfp::rc_fp4_gemm_t(A, B, D, P, g, tg, sg, lane); \ + } +RC_FP4_GEMM(rc_fp4_gemm_00, false, false) +RC_FP4_GEMM(rc_fp4_gemm_01, false, true) +RC_FP4_GEMM(rc_fp4_gemm_10, true, false) +RC_FP4_GEMM(rc_fp4_gemm_11, true, true) diff --git a/zig/kernels/metal/ricursive/glue.metal b/zig/kernels/metal/ricursive/glue.metal new file mode 100644 index 0000000000..203a3d77d5 --- /dev/null +++ b/zig/kernels/metal/ricursive/glue.metal @@ -0,0 +1,197 @@ +// Ripple's glue at the release's rounding points (modeling_ricursive in bf16): bf16 residual, four norms a layer. +#include +using namespace metal; + +#define TF_UNROLL_4 _Pragma("clang loop unroll(full)") + +constant constexpr uint DH = RC_DH; // the head dim (128), defined by the loader + +// Up to RC_NP of a thread's row elements (t, t + n, ...) held in registers. +constant constexpr uint RC_NP = 32; + +struct rc_args { uint rows, width, flags; float eps, scale, blend; }; +constant constexpr uint RC_ADD = 1; // rc_add_norm: R += norm(Y) by w_post before the norm +constant constexpr uint RC_LERP = 2; // rc_add_norm: also VX = lerp(V0, X, blend), v's input + +// R and V0 row r = bf16(E[ids[r]] * scale): the scaled embedding, also the value residual's source. +[[kernel]] void rc_embed(const device uint* ids [[buffer(0)]], const device bfloat* E [[buffer(1)]], + constant rc_args& a [[buffer(2)]], device bfloat* R [[buffer(3)]], device bfloat* V0 [[buffer(4)]], + uint2 p [[thread_position_in_grid]]) { + if (p.x >= a.width || p.y >= a.rows) return; + const size_t at = (size_t)p.y * a.width + p.x; + const bfloat v = bfloat(float(E[(size_t)ids[p.y] * a.width + p.x]) * a.scale); + R[at] = v; + V0[at] = v; +} + +// A row's fp32 sum of squares over the threadgroup into 1 / rms: rsqrt(mean + eps), the parts summed in order. +inline float rc_inv(float sum, constant rc_args& a, 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]; + return precise::rsqrt(total / float(a.width) + a.eps); +} + +// torch.lerp on bf16 (its vectorized form): fma(w, e - s, s) below one half, else fma(w - 1, e - s, e); rounded once. +inline bfloat rc_lerp(float s, float e, float w) { + return bfloat(fabs(w) < 0.5f ? fma(w, e - s, s) : fma(w - 1.0f, e - s, e)); +} + +// A row a group: R = bf16(R + norm(Y, w_post)) if RC_ADD; X = norm(R, w_next); VX = lerp(V0, X, blend) if RC_LERP. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void rc_add_norm(device bfloat* R [[buffer(0)]], const device bfloat* Y [[buffer(1)]], + const device bfloat* w_post [[buffer(2)]], const device bfloat* w_next [[buffer(3)]], + const device bfloat* V0 [[buffer(4)]], constant rc_args& a [[buffer(5)]], device bfloat* X [[buffer(6)]], + device bfloat* VX [[buffer(7)]], uint row [[threadgroup_position_in_grid]], + uint t [[thread_position_in_threadgroup]], uint n [[threads_per_threadgroup]]) { + threadgroup float part_y[32], part_r[32]; + const size_t base = (size_t)row * a.width; + float r[RC_NP]; + TF_UNROLL_4 + for (uint p = 0; p < RC_NP; ++p) r[p] = t + p * n < a.width ? float(R[base + t + p * n]) : 0.0f; + if (a.flags & RC_ADD) { + float y[RC_NP], sum = 0.0f; + TF_UNROLL_4 + for (uint p = 0; p < RC_NP; ++p) { + y[p] = t + p * n < a.width ? float(Y[base + t + p * n]) : 0.0f; + sum = fma(y[p], y[p], sum); + } + const float inv = rc_inv(sum, a, t, n, part_y); + TF_UNROLL_4 + for (uint p = 0; p < RC_NP; ++p) { + if (t + p * n >= a.width) continue; + const float normed = float(bfloat(float(w_post[t + p * n]) * float(bfloat(y[p] * inv)))); + r[p] = float(bfloat(r[p] + normed)); + R[base + t + p * n] = bfloat(r[p]); + } + } + float sum = 0.0f; + TF_UNROLL_4 + for (uint p = 0; p < RC_NP; ++p) sum = fma(r[p], r[p], sum); + const float inv = rc_inv(sum, a, t, n, part_r); + TF_UNROLL_4 + for (uint p = 0; p < RC_NP; ++p) { + const uint i = t + p * n; + if (i >= a.width) continue; + const bfloat x = bfloat(float(w_next[i]) * float(bfloat(r[p] * inv))); + X[base + i] = x; + if (a.flags & RC_LERP) VX[base + i] = rc_lerp(float(V0[base + i]), float(x), a.blend); + } +} + +struct rc_prep_args { uint rows, heads, kv_heads, pos, stride, ring, rope; float eps; }; + +// Per head: q_norm or k_norm, then on sliding layers fp32 RoPE (bf16 cos, sin) rounded once; k, v at (pos+row) % ring. +[[kernel]] void rc_prep(const device bfloat* X [[buffer(0)]], const device float* inv [[buffer(1)]], + const device bfloat* qw [[buffer(2)]], const device bfloat* kw [[buffer(3)]], constant rc_prep_args& a [[buffer(4)]], + device bfloat* QN [[buffer(5)]], device bfloat* KC [[buffer(6)]], device bfloat* VC [[buffer(7)]], + uint2 tg [[threadgroup_position_in_grid]], uint lane [[thread_index_in_simdgroup]]) { + constexpr uint PER = DH / 32, HALF = DH / 2; + static_assert(PER == 4, "a lane holds dims lane + 32j: a dim's RoPE partner is in the same lane"); + 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 device bfloat* w = is_q ? qw : kw; + float x[PER], sum = 0.0f; + TF_UNROLL_4 + for (uint j = 0; j < PER; ++j) { + x[j] = float(src[lane + 32 * j]); + sum = fma(x[j], x[j], sum); + } + const float r = precise::rsqrt(simd_sum(sum) / float(DH) + a.eps); + TF_UNROLL_4 + for (uint j = 0; j < PER; ++j) x[j] = float(bfloat(float(w[lane + 32 * j]) * float(bfloat(x[j] * r)))); + const size_t slot = (a.pos + row) % a.ring; + 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); + TF_UNROLL_4 + for (uint j = 0; j < PER; ++j) { + const uint d = lane + 32 * j; + if (a.rope == 0) { + y[d] = bfloat(x[j]); + continue; + } + const float angle = position * inv[d % HALF]; + const float c = float(bfloat(precise::cos(angle))), s = float(bfloat(precise::sin(angle))); + const float other = d < HALF ? -x[j + 2] : x[j - 2]; + y[d] = bfloat(x[j] * c + 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]; +} + +// The output gate: A = bf16(A * bf16(sigmoid(fp32 G))), G the bf16 g_proj rows (the release's gate_full_precision). +[[kernel]] void rc_gate(device bfloat* A [[buffer(0)]], const device bfloat* G [[buffer(1)]], + constant rc_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; + const float g = float(G[i]); + const float s = float(bfloat(1.0f / (1.0f + precise::exp(-g)))); + A[i] = bfloat(float(A[i]) * s); +} + +// The release's FFN on a stacked gate|up row: gate at most scale, up within +-scale, then bf16(bf16(silu(g)) * u). +[[kernel]] void rc_silu_clamp(const device bfloat* GU [[buffer(0)]], constant rc_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; + // torch.clamp keeps a NaN where Metal's min and clamp return the other operand + const float g0 = float(GU[at]), u0 = float(GU[at + a.width]); + const float g = isnan(g0) ? g0 : min(g0, a.scale), u = isnan(u0) ? u0 : clamp(u0, -a.scale, a.scale); + const float silu = float(bfloat(g / (1.0f + precise::exp(-g)))); + A[(size_t)p.y * a.width + p.x] = bfloat(silu * u); +} + +struct rc_vec_args { uint rows, pos, heads, groups, blocks, chunk, first; float scale; uint window, ring; }; + +// Decode partials a (key head, row, block) a group: ll_decode_vec over the window at slot i % ring, blocks fixed. +[[kernel]] void rc_decode_vec(const device bfloat* Q [[buffer(0)]], const device bfloat* K [[buffer(1)]], + const device bfloat* V [[buffer(2)]], constant rc_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 low = p + 1 > a.window ? p + 1 - a.window : 0; + const uint from = max(block * a.chunk, low), 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 % a.ring) * 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; +} diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig index 26a4ffe619..1ab74007fb 100644 --- a/zig/kernels/metal/sources.zig +++ b/zig/kernels/metal/sources.zig @@ -203,6 +203,13 @@ pub const llama_glue = @embedFile("llama/glue.metal"); 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"); +/// Ripple's glue at the release's rounding points: embedding, four norms a layer, q/k norm, RoPE, gate, clamp. +pub const ricursive_glue = @embedFile("ricursive/glue.metal"); +/// The NVFP4 release's MLP row kernels and dequantization; its tensor-unit GEMM after gemm_nax.metal. +pub const ricursive_fp4 = @embedFile("ricursive/fp4.metal"); +pub const ricursive_fp4_gemm = @embedFile("prefill/gemm_nax.metal") ++ @embedFile("ricursive/fp4_gemm.metal"); +/// Ripple's prompt attention on the tensor units (llama's kernel over a window, keys in a ring; nax.h inlined at load). +pub const ricursive_attention_nax = @embedFile("ricursive/attention_nax.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/cuda_kl.zig b/zig/src/cli/cuda_kl.zig index 4d3a98bb8b..25fcb13575 100644 --- a/zig/src/cli/cuda_kl.zig +++ b/zig/src/cli/cuda_kl.zig @@ -35,39 +35,61 @@ fn logSoftmax(row: []const u8, out: []f64) void { for (out) |*o| o.* -= z; } -fn download(gpa: std.mem.Allocator, e: *nemotron.Engine, ptr: u64, len: usize) ![]u8 { - const host = try gpa.alloc(u8, len); - errdefer gpa.free(host); - try e.ops().download(host, ptr); - try e.stream.synchronize(); - return host; +/// An engine's logits rows in order: row(i) gives row i's bf16 logits, step(i, t) feeds token t after it. +pub const Source = struct { + ctx: *anyopaque, + vocab: usize, + row: *const fn (ctx: *anyopaque, i: usize, out: []u8) anyerror!void, + step: *const fn (ctx: *anyopaque, i: usize, t: u32) anyerror!void, +}; + +/// The usage refusal before any GPU work: a prompt and exactly one of --dump and --against. +pub fn refused(prompt: []const u32, dump: ?[]const u8, against: ?[]const u8) bool { + if ((dump == null) != (against == null) and prompt.len != 0) return false; + std.debug.print("kl: give --tokens and one of --dump FILE or --against FILE\n", .{}); + return true; } +const Nemotron = struct { + fn row(ctx: *anyopaque, i: usize, out: []u8) anyerror!void { + const e: *nemotron.Engine = @ptrCast(@alignCast(ctx)); + try e.ops().download(out, if (i == 0) e.b.p_logits else e.b.logits); + try e.stream.synchronize(); + } + fn step(ctx: *anyopaque, _: usize, t: u32) anyerror!void { + const e: *nemotron.Engine = @ptrCast(@alignCast(ctx)); + _ = try e.step(t, null); + } +}; + /// `dump`: `n` rows after `prompt` into the file, fed `feed` then greedy. `against`: the file's tokens fed, row by row. pub fn run(gpa: std.mem.Allocator, io: std.Io, e: *nemotron.Engine, prompt: []const u32, feed: []const u32, n: usize, dump: ?[]const u8, against: ?[]const u8) !u8 { - if ((dump == null) == (against == null) or prompt.len == 0) { - std.debug.print("kl: give --tokens and one of --dump FILE or --against FILE\n", .{}); - return 2; - } - const v = e.c.vocab; - const rec = recordBytes(v); + if (refused(prompt, dump, against)) return 2; _ = try e.prefill(prompt, null, null); + return compare(gpa, io, .{ .ctx = e, .vocab = e.c.vocab, .row = Nemotron.row, .step = Nemotron.step }, prompt.len, feed, n, dump, against); +} + +/// `src`'s rows after a `prompt_len`-token prompt: `dump` writes `n`, fed `feed` then greedy; `against` feeds the file. +pub fn compare(gpa: std.mem.Allocator, io: std.Io, src: Source, prompt_len: usize, feed: []const u32, n: usize, dump: ?[]const u8, against: ?[]const u8) !u8 { + const v = src.vocab; + const rec = recordBytes(v); + const row = try gpa.alloc(u8, v * 2); + defer gpa.free(row); if (dump) |path| { const buf = try gpa.alloc(u8, header_bytes + n * rec); defer gpa.free(buf); std.mem.writeInt(u32, buf[0..4], @intCast(n), .little); std.mem.writeInt(u32, buf[4..8], @intCast(v), .little); for (0..n) |i| { - const row = try download(gpa, e, if (i == 0) e.b.p_logits else e.b.logits, v * 2); - defer gpa.free(row); + try src.row(src.ctx, i, row); const t: u32 = if (i < feed.len) feed[i] else @intCast(argmax(row, v)); const at = header_bytes + i * rec; std.mem.writeInt(u32, buf[at..][0..4], t, .little); @memcpy(buf[at + 4 ..][0 .. v * 2], row); - if (i + 1 < n) _ = try e.step(t, null); + if (i + 1 < n) try src.step(src.ctx, i, t); } try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = path, .data = buf }); - std.debug.print("RESULT kl dump: {d} rows ({d} fed) after a {d}-token prompt to {s}\n", .{ n, @min(n, feed.len), prompt.len, path }); + std.debug.print("RESULT kl dump: {d} rows ({d} fed) after a {d}-token prompt to {s}\n", .{ n, @min(n, feed.len), prompt_len, path }); return 0; } const ref = try std.Io.Dir.cwd().readFileAlloc(io, against.?, gpa, .limited(1 << 34)); @@ -89,8 +111,7 @@ pub fn run(gpa: std.mem.Allocator, io: std.Io, e: *nemotron.Engine, prompt: []co const at = header_bytes + i * rec; const t = std.mem.readInt(u32, ref[at..][0..4], .little); const want = ref[at + 4 ..][0 .. v * 2]; - const row = try download(gpa, e, if (i == 0) e.b.p_logits else e.b.logits, v * 2); - defer gpa.free(row); + try src.row(src.ctx, i, row); logSoftmax(want, lp); logSoftmax(row, lq); var kl: f64 = 0; @@ -105,7 +126,7 @@ pub fn run(gpa: std.mem.Allocator, io: std.Io, e: *nemotron.Engine, prompt: []co if (!same and first_diff == null) first_diff = i; nll[0] -= lp[t]; nll[1] -= lq[t]; - if (i + 1 < rows) _ = try e.step(t, null); + if (i + 1 < rows) try src.step(src.ctx, i, t); } const r: f64 = @floatFromInt(rows); std.debug.print("RESULT kl: {d} rows, mean KL {e:.3} nats, max {e:.3} at row {d}; top-1 {d} of {d}; first differs at {?d}; mean NLL of the fed tokens {d:.4} (reference) {d:.4} (this run)\n", .{ rows, kl_sum / r, kl_max, kl_at, top1, rows, first_diff, nll[0] / r, nll[1] / r }); diff --git a/zig/src/cli/cuda_main.zig b/zig/src/cli/cuda_main.zig index 6fdea9e952..5f8f88dde3 100644 --- a/zig/src/cli/cuda_main.zig +++ b/zig/src/cli/cuda_main.zig @@ -15,6 +15,7 @@ const lanes_cli = @import("cuda_lanes.zig"); const segments_cli = @import("cuda_segments.zig"); const train_check = @import("cuda_train_check.zig"); const dory_cli = @import("cuda_dory.zig"); +const ricursive_cli = @import("cuda_ricursive.zig"); const decode = nemotron.decode; const usage = @@ -42,6 +43,7 @@ const usage = \\ tensorfold digest MODEL (every byte of the checkpoint hashed, as a tensor-parallel rank checks it at load) \\ tensorfold kl MODEL --tokens ID,... [--gold ID,...] [--max-tokens N] (--dump FILE | --against FILE) [--precision P] \\ (logits rows to FILE, fed --gold then greedy; or FILE's tokens fed: KL, top-1 and NLL against its rows) + \\ [--gold-file PATH] [--chunk ROWS] (a Ripple checkpoint: every prompt row's logits, the prompt in chunks of ROWS) \\ tensorfold widths MODEL --tokens ID,... [--max-tokens N] [--eager] [sampling as run] \\ (every verify width, 1-32 for a latent MoE else 1-16, against serial decoding, a wrong draft every third window) \\ tensorfold models @@ -91,6 +93,13 @@ pub fn main(init: std.process.Init) !u8 { } else if (std.mem.eql(u8, a, "--gold")) { opts.gold = try parseIds(gpa, value); i += 1; + } else if (std.mem.eql(u8, a, "--gold-file")) { + gpa.free(opts.gold); + opts.gold = try segments_cli.readIds(gpa, init.io, value); + i += 1; + } else if (std.mem.eql(u8, a, "--chunk")) { + opts.chunk = try std.fmt.parseInt(usize, value, 10); + i += 1; } else if (std.mem.eql(u8, a, "--precision")) { opts.math = std.meta.stringToEnum(nemotron.Precision, value) orelse return error.BadPrecision; i += 1; @@ -192,10 +201,25 @@ pub fn main(init: std.process.Init) !u8 { const s: ?lanes.Sampling = if (opts.sampling.temperature > 0) opts.sampling else null; return dory_cli.run(gpa, init.io, &ctx, opts.model, .{ .tokens = opts.tokens, .max_tokens = opts.max_tokens, .loops = opts.loops, .context = opts.context, .drafts = opts.drafts, .graphs = opts.graphs, .stop_eos = opts.stop_eos, .sampling = s }); } + if (ricursive_cli.wants(gpa, init.io, opts.model)) { + if (kernels != null or opts.sampling.temperature > 0 or opts.loops != null) { + std.debug.print("Ripple on CUDA runs greedily on its own kernels: --kernels native, temperature 0, no --loops\n", .{}); + return 2; + } + const ro: ricursive_cli.Options = .{ .tokens = opts.tokens, .gold = opts.gold, .max_tokens = opts.max_tokens, .context = opts.context, .stop_eos = opts.stop_eos, .report = opts.report, .chunk = opts.chunk, .dump = opts.dump, .against = opts.against, .graphs = opts.graphs }; + if (std.mem.eql(u8, args[1], "run")) return ricursive_cli.run(gpa, init.io, &ctx, opts.model, ro); + if (std.mem.eql(u8, args[1], "kl")) return ricursive_cli.klRun(gpa, init.io, &ctx, opts.model, ro); + std.debug.print("Ripple on CUDA takes run and kl so far\n", .{}); + return 2; + } if (opts.loops != null) { std.debug.print("--loops sets a Dory checkpoint's recurrent passes: {s} is not one\n", .{opts.model}); return 2; } + if (opts.chunk != null) { + std.debug.print("--chunk forwards a Ripple prompt in chunks: {s} is not a Ripple checkpoint\n", .{opts.model}); + return 2; + } const env = init.environ_map; var carve: ?cuda.Carveout = null; if (try cuda.carveout.requested(opts.carveout, env.get("TF_CUDA_CARVEOUT"), env.get("TF_CUDA_CARVEOUT_MIB"))) |bytes| { @@ -288,6 +312,7 @@ const Options = struct { costs: ?[]f64 = null, segments: ?usize = null, loops: ?usize = null, // a Dory checkpoint's recurrent passes (run only) + chunk: ?usize = null, // kl on a Ripple checkpoint: every prompt row's logits, the prompt in chunks of this many rows counts: Counts = .{ .items = .{ 1, 2, 3, 4 }, .len = 4 }, repeat: usize = 3, streams: usize = 4, diff --git a/zig/src/cli/cuda_ricursive.zig b/zig/src/cli/cuda_ricursive.zig new file mode 100644 index 0000000000..aaace2eda1 --- /dev/null +++ b/zig/src/cli/cuda_ricursive.zig @@ -0,0 +1,195 @@ +//! `tensorfold run` and `kl` on a Ripple checkpoint: the CUDA forward driven directly, one stream; --eager: no graphs. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rc = @import("ricursive_cuda"); +const kl = @import("cuda_kl.zig"); + +const fwd = rc.model; + +/// The checkpoint at `dir` is a Ripple one: its config.json names model_type ricursive. +pub fn wants(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) bool { + const path = std.fs.path.join(gpa, &.{ dir, "config.json" }) catch return false; + defer gpa.free(path); + const text = std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 22)) catch return false; + defer gpa.free(text); + const Head = struct { model_type: []const u8 = "" }; + const head = std.json.parseFromSlice(Head, gpa, text, .{ .ignore_unknown_fields = true }) catch return false; + defer head.deinit(); + return std.mem.eql(u8, head.value.model_type, rc.native.model_type); +} + +pub const Options = struct { + tokens: []const u32, + gold: []const u32 = &.{}, + max_tokens: usize, + context: ?usize, + stop_eos: bool = true, + report: ?[]const u8 = null, + chunk: ?usize = null, // kl: every prompt row's logits, the prompt in chunks of this many rows + dump: ?[]const u8 = null, + against: ?[]const u8 = null, + graphs: bool = true, // decode rounds as CUDA graphs; false: eager launches +}; + +fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, context: ?usize) !*rc.native.Owned { + const own = try rc.native.openOwned(gpa, io, ctx, dir, context orelse @intCast(rc.native.default_context), null); + std.debug.print("Ripple on CUDA: {d}-token context, prompt chunks of {d} rows\n", .{ own.m.context, fwd.prompt_rows }); + return own; +} + +/// The prompt from position 0 in chunks of up to 2,048 rows, its last row through the head. +fn prefill(m: *const fwd.Model, st: *fwd.Seq, ids: []const u32) !void { + st.len = 0; + var at: usize = 0; + while (at < ids.len) { + const end = @min(ids.len, at + fwd.prompt_rows); + try m.forward(st, ids[at..end], .prompt); + if (end == ids.len) try m.head(end - at - 1); + at = end; + } +} + +/// The greedy token of the head's last row. +fn greedy(own: *rc.native.Owned) !u32 { + try own.m.draw(own.rb.tok.ptr); + var out: [1]u32 = undefined; + try own.m.download(own.rb.tok.ptr, std.mem.sliceAsBytes(&out)); + return out[0]; +} + +fn seconds(io: std.Io, since: std.Io.Timestamp) f64 { + return @as(f64, @floatFromInt(since.durationTo(std.Io.Timestamp.now(io, .awake)).nanoseconds)) / 1e9; +} + +/// A greedy reply: the prompt's chunks, then one-row decode steps; the output line and --report as Nemotron's run. +pub fn run(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, o: Options) !u8 { + if (o.tokens.len == 0) return error.NoPromptTokens; + const own = try open(gpa, io, ctx, dir, o.context); + defer rc.native.closeOwned(own); + const m = &own.m; + m.graphs = o.graphs; + const st = &own.rb.st; + const c = own.w.config; + const count = @max(1, @min(o.max_tokens, m.context -| o.tokens.len)); + var out: std.ArrayList(u32) = .empty; + defer out.deinit(gpa); + const t0 = std.Io.Timestamp.now(io, .awake); + try prefill(m, st, o.tokens); + try out.append(gpa, try greedy(own)); + const prefill_seconds = seconds(io, t0); + const t1 = std.Io.Timestamp.now(io, .awake); + while (out.items.len < count) { + const t = out.items[out.items.len - 1]; + if (o.stop_eos and c.isEos(t)) break; + try m.forward(st, &.{t}, .decode); + try m.head(0); + try out.append(gpa, try greedy(own)); + } + const decode_seconds = seconds(io, t1); + var digest: [32]u8 = undefined; + const text = try core.ids_json.write(gpa, out.items); + defer gpa.free(text); + std.crypto.hash.sha2.Sha256.hash(text, &digest, .{}); + const hex = std.fmt.bytesToHex(digest, .lower); + const steps: f64 = @floatFromInt(@max(1, out.items.len - 1)); + const ms = decode_seconds * 1e3 / steps; + std.debug.print("tokens {d} sha {s} prefill {d:.4}s decode {d:.4}s {d:.3} ms/token {d:.2} tok/s prompt {d:.1} tok/s\n", .{ out.items.len, hex[0..12], prefill_seconds, decode_seconds, ms, 1e3 / ms, @as(f64, @floatFromInt(o.tokens.len)) / prefill_seconds }); + if (o.report) |path| { + const report = .{ + .engine = "zig-cuda", + .model_type = rc.native.model_type, + .prompt_tokens = o.tokens, + .tokens = out.items, + .token_sha256 = hex, + .prefill_seconds = prefill_seconds, + .decode_seconds = decode_seconds, + .ms_per_token = ms, + .tokens_per_second = 1e3 / ms, + .prompt_tokens_per_second = @as(f64, @floatFromInt(o.tokens.len)) / prefill_seconds, + .drafts = false, + .graphs = m.captured() > 0, // a decode round replayed a graph + .max_len = m.context, + }; + const json = try std.json.Stringify.valueAlloc(gpa, report, .{}); + defer gpa.free(json); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = path, .data = json }); + } + return 0; +} + +/// kl's rows after a prefill: row 0 the prompt's last, then one-row decode steps. +const Decode = struct { + m: *const fwd.Model, + st: *fwd.Seq, + + fn row(p: *anyopaque, _: usize, out: []u8) anyerror!void { + const x: *Decode = @ptrCast(@alignCast(p)); + try x.m.download(x.m.b.logits, out); + } + + fn step(p: *anyopaque, _: usize, t: u32) anyerror!void { + const x: *Decode = @ptrCast(@alignCast(p)); + try x.m.forward(x.st, &.{t}, .decode); + try x.m.head(0); + } +}; + +/// kl's rows of a whole prompt teacher-forced: row i after ids[0..i+1], the prompt forwarded in chunks of `chunk` rows. +const Chunks = struct { + m: *const fwd.Model, + st: *fwd.Seq, + ids: []const u32, + chunk: usize, + base: usize = 0, // the current forward's first row + done: usize = 0, // rows forwarded + head_at: ?usize = null, // the forward's row in `rows`' first slot + + fn row(p: *anyopaque, i: usize, out: []u8) anyerror!void { + const x: *Chunks = @ptrCast(@alignCast(p)); + if (i >= x.done) { + const end = @min(x.ids.len, x.done + x.chunk); + try x.m.forward(x.st, x.ids[x.done..end], .prompt); + x.base = x.done; + x.done = end; + x.head_at = null; + } + const r = i - x.base; + const v = x.m.c.vocab; + if (x.head_at == null or r < x.head_at.? or r >= x.head_at.? + fwd.head_rows) { + const r0 = r - r % fwd.head_rows; + try x.m.headRows(r0, @min(fwd.head_rows, x.done - x.base - r0)); + x.head_at = r0; + } + try x.m.download(x.m.b.rows + (r - x.head_at.?) * v * 2, out); + } + + /// The prompt's own tokens only: a reference whose tokens differ is another prompt. + fn step(p: *anyopaque, i: usize, t: u32) anyerror!void { + const x: *Chunks = @ptrCast(@alignCast(p)); + if (i + 1 >= x.ids.len or x.ids[i + 1] != t) return error.TokenMismatch; + } +}; + +/// --chunk ROWS: every prompt row's logits; else the last prompt row and `max_tokens - 1` decode rows (--gold, greedy). +pub fn klRun(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, o: Options) !u8 { + if (kl.refused(o.tokens, o.dump, o.against)) return 2; + if (o.chunk) |c| if (c == 0 or c > fwd.prompt_rows) { + std.debug.print("--chunk {d}: prompt rows a forward, 1 to {d}\n", .{ c, fwd.prompt_rows }); + return 2; + }; + const own = try open(gpa, io, ctx, dir, o.context); + defer rc.native.closeOwned(own); + const m = &own.m; + m.graphs = o.graphs; + const st = &own.rb.st; + const vocab = own.w.config.vocab; + if (o.chunk) |c| { + st.len = 0; + var src: Chunks = .{ .m = m, .st = st, .ids = o.tokens, .chunk = c }; + return kl.compare(gpa, io, .{ .ctx = &src, .vocab = vocab, .row = Chunks.row, .step = Chunks.step }, 1, o.tokens[1..], o.tokens.len, o.dump, o.against); + } + try prefill(m, st, o.tokens); + var src: Decode = .{ .m = m, .st = st }; + return kl.compare(gpa, io, .{ .ctx = &src, .vocab = vocab, .row = Decode.row, .step = Decode.step }, o.tokens.len, o.gold, o.max_tokens, o.dump, o.against); +} diff --git a/zig/src/cli/hub.zig b/zig/src/cli/hub.zig index a13e2e4fed..2c8737eb9b 100644 --- a/zig/src/cli/hub.zig +++ b/zig/src/cli/hub.zig @@ -21,6 +21,7 @@ const titles = [_]struct { model_type: []const u8, title: []const u8 }{ .{ .model_type = "spark2_5", .title = "Spark-X2.6-4B" }, .{ .model_type = "minimax_m3_vl", .title = "MiniMax-M3.1-Flash" }, .{ .model_type = "llama", .title = "Llama architecture (Nanbeige4.3-33B)" }, + .{ .model_type = "ricursive", .title = "Ripple-65B (ricursive)" }, .{ .model_type = "lfm2_moe", .title = "LFM2.5-24B-A2B" }, }; diff --git a/zig/src/core/contract/measurements/pending_133.zig b/zig/src/core/contract/measurements/pending_133.zig index bd21f7a0c3..136e4a4637 100644 --- a/zig/src/core/contract/measurements/pending_133.zig +++ b/zig/src/core/contract/measurements/pending_133.zig @@ -59,4 +59,12 @@ pub const records = [_]Override{ .{ .name = "w4a8_prefill tf_w4a8_quantize", .usage = .{ .{ .regs = 27, .static = 0, .bound = 128 }, .{ .regs = 27, .static = 0, .bound = 128 }, .{ .regs = 27, .static = 0, .bound = 128 }, .{ .regs = 27, .static = 0, .bound = 128 }, .{ .regs = 27, .static = 0, .bound = 128 } } }, .{ .name = "qmm_group MLP128", .usage = .{ .{ .regs = 0, .static = 0 }, .{ .regs = 0, .static = 0 }, .{ .regs = 0, .static = 0 }, .{ .regs = 54, .static = 0, .bound = 256 }, .{ .regs = 54, .static = 0, .bound = 256 } } }, .{ .name = "qwen35_keep tf_q35k_apply_pend", .usage = .{ .{ .regs = 32, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 } } }, + .{ .name = "ricursive tf_rc_embed", .usage = .{ .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 } } }, + .{ .name = "ricursive tf_rc_add_norm", .usage = .{ .{ .regs = 80, .static = 32, .bound = 256 }, .{ .regs = 79, .static = 32, .bound = 256 }, .{ .regs = 79, .static = 32, .bound = 256 }, .{ .regs = 79, .static = 32, .bound = 256 }, .{ .regs = 79, .static = 32, .bound = 256 } } }, + .{ .name = "ricursive tf_rc_gate", .usage = .{ .{ .regs = 15, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 } } }, + .{ .name = "ricursive tf_rc_silu_clamp", .usage = .{ .{ .regs = 17, .static = 0, .bound = 256 }, .{ .regs = 18, .static = 0, .bound = 256 }, .{ .regs = 18, .static = 0, .bound = 256 }, .{ .regs = 19, .static = 0, .bound = 256 }, .{ .regs = 19, .static = 0, .bound = 256 } } }, + .{ .name = "ricursive_prep tf_rc_prep", .usage = .{ .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 38, .static = 0, .bound = 256 }, .{ .regs = 38, .static = 0, .bound = 256 }, .{ .regs = 36, .static = 0, .bound = 256 }, .{ .regs = 36, .static = 0, .bound = 256 } } }, + .{ .name = "ricursive_prep tf_rc_prep_rows", .usage = .{ .{ .regs = 38, .static = 0, .bound = 256 }, .{ .regs = 38, .static = 0, .bound = 256 }, .{ .regs = 38, .static = 0, .bound = 256 }, .{ .regs = 38, .static = 0, .bound = 256 }, .{ .regs = 38, .static = 0, .bound = 256 } } }, + .{ .name = "window attention wattn_kernel<128, 8, 2, 8>", .usage = .{ .{ .regs = 222, .static = 0, .bound = 256 }, .{ .regs = 222, .static = 0, .bound = 256 }, .{ .regs = 222, .static = 0, .bound = 256 }, .{ .regs = 218, .static = 0, .bound = 256 }, .{ .regs = 218, .static = 0, .bound = 256 } } }, + .{ .name = "window attention tf_window_prompt_h2_pg", .usage = .{ .{ .regs = 222, .static = 0, .bound = 256 }, .{ .regs = 222, .static = 0, .bound = 256 }, .{ .regs = 222, .static = 0, .bound = 256 }, .{ .regs = 231, .static = 0, .bound = 256 }, .{ .regs = 231, .static = 0, .bound = 256 } } }, }; diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 552a9243d1..5b083ac106 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -350,8 +350,8 @@ pub fn usageFor(r: CudaRecord, release: ?limits.Nvcc) [limits.sms.len]U { } return u; } -/// Restored Qwen launches retain a pending entry wherever a compiler receipt is absent. -pub const cuda_pending = @import("qwen_whole_records.zig").records ++ @import("cuda_lane_pending.zig").records ++ @import("weight_only_records.zig").cuda ++ @import("hc_activation_records.zig").cuda ++ @import("qwen4_records.zig").cuda; +/// Restored Qwen launches, and Ripple's unmeasured ones, retain a pending entry wherever a compiler receipt is absent. +pub const cuda_pending = @import("qwen_whole_records.zig").records ++ @import("cuda_lane_pending.zig").records ++ @import("weight_only_records.zig").cuda ++ @import("hc_activation_records.zig").cuda ++ @import("qwen4_records.zig").cuda ++ @import("ricursive_cuda_records.zig").cuda ++ @import("ricursive_prep_records.zig").cuda ++ @import("ricursive_attn_records.zig").cuda; /// Gate 5's planted CUDA controls: a launch past its __launch_bounds__, one needing more registers than a block has. pub const cuda_controls = [_]CudaRecord{ .{ .launch = .{ .name = "control: past its bounds", .symbol = "tf_ctl_bounded", .image = "contract_controls", .threads = 512 }, .usage = .{ .{ .regs = 10, .static = 0, .bound = 256 }, .{ .regs = 10, .static = 0, .bound = 256 }, .{ .regs = 10, .static = 0, .bound = 256 }, .{ .regs = 10, .static = 0, .bound = 256 }, .{ .regs = 10, .static = 0, .bound = 256 } } }, @@ -570,7 +570,7 @@ test { const paging = @import("paging_pending.zig"); /// The Metal collector keeps measured slots and the restored vision declarations. pub const all_metal = @import("attention_pv_records.zig").records ++ metal ++ @import("delta_slots_records.zig").records ++ @import("flashnext_pending.zig").records ++ @import("simd_pairs_records.zig").records ++ @import("metal_pair_records.zig").records ++ @import("glm_rows_records.zig").records ++ @import("weight_only_records.zig").metal ++ @import("hc_activation_records.zig").metal ++ @import("qwen4_records.zig").metal ++ @import("lfm2_records.zig").records ++ @import("nemotron_experts_records.zig").records ++ @import("nemotron_pages_records.zig").records ++ @import("nemotron_attention_records.zig").records ++ - @import("nemotron_ssd_records.zig").records ++ @import("flashnext_attention_records.zig").records ++ @import("spark25_records.zig").records ++ @import("gemma_attention_records.zig").records ++ @import("moss_image_records.zig").records ++ @import("dsv41_metal_records.zig").records ++ @import("minimax_records.zig").records ++ @import("ling_route_records.zig").records ++ @import("qwen27_moe_records.zig").records ++ @import("affine_gather_records.zig").records; + @import("nemotron_ssd_records.zig").records ++ @import("flashnext_attention_records.zig").records ++ @import("spark25_records.zig").records ++ @import("gemma_attention_records.zig").records ++ @import("moss_image_records.zig").records ++ @import("dsv41_metal_records.zig").records ++ @import("minimax_records.zig").records ++ @import("ling_route_records.zig").records ++ @import("qwen27_moe_records.zig").records ++ @import("affine_gather_records.zig").records ++ @import("ricursive_records.zig").records; const learning = @import("registry_learning.zig"); pub const gradient_pending = learning.gradient_pending; pub const qwen_learning_pending = learning.qwen_learning_pending; diff --git a/zig/src/core/contract/ricursive_attn_records.zig b/zig/src/core/contract/ricursive_attn_records.zig new file mode 100644 index 0000000000..92886739d0 --- /dev/null +++ b/zig/src/core/contract/ricursive_attn_records.zig @@ -0,0 +1,16 @@ +//! Ripple's prompt attention on CUDA: window_attention.cu's two-heads-a-block instances, for kernel-contract-limits. +const registry = @import("registry.zig"); +const limits = @import("limits.zig"); + +/// nvcc 13.0's usage, the compiler of the release fatbins: not yet measured (every value 0). +pub const cuda = [_]registry.CudaRecord{ + // TODO(spark): fill from kernel-contract-limits + .{ .launch = .{ .name = "window attention wattn_kernel<128, 8, 2, 8>", .symbol = "_ZN16tf_window_prompt12wattn_kernelILi128ELi8ELi2ELi8EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifii", .image = "window_attention", .threads = 256, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 } } }, + // TODO(spark): fill from kernel-contract-limits + .{ .launch = .{ .name = "window attention tf_window_prompt_h2_pg", .symbol = "tf_window_prompt_h2_pg", .image = "window_attention", .threads = 256, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 } } }, +}; + +test "Ripple's prompt attention: unmeasured records cannot pass a compiled resource check" { + const t = @import("std").testing; + for (cuda) |r| try t.expectEqual(limits.sms.len, try registry.checkCuda(t.allocator, &.{r}, true)); +} diff --git a/zig/src/core/contract/ricursive_cuda_records.zig b/zig/src/core/contract/ricursive_cuda_records.zig new file mode 100644 index 0000000000..bcce96a81e --- /dev/null +++ b/zig/src/core/contract/ricursive_cuda_records.zig @@ -0,0 +1,19 @@ +//! Ripple's CUDA glue (ricursive.cu): its served launches, pending in registry.cuda_pending until nvcc's records land. +const std = @import("std"); +const registry = @import("registry.zig"); + +/// The glue's launches, 256 threads each; usage zero (no compiled record yet): gate 5 refuses them until filled. +pub const cuda = [_]registry.CudaRecord{ + .{ .launch = .{ .name = "ricursive tf_rc_embed", .symbol = "tf_rc_embed", .image = "ricursive", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 } } }, // TODO(spark): fill from kernel-contract-limits + .{ .launch = .{ .name = "ricursive tf_rc_add_norm", .symbol = "tf_rc_add_norm", .image = "ricursive", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 } } }, // TODO(spark): fill from kernel-contract-limits + .{ .launch = .{ .name = "ricursive tf_rc_gate", .symbol = "tf_rc_gate", .image = "ricursive", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 } } }, // TODO(spark): fill from kernel-contract-limits + .{ .launch = .{ .name = "ricursive tf_rc_silu_clamp", .symbol = "tf_rc_silu_clamp", .image = "ricursive", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 } } }, // TODO(spark): fill from kernel-contract-limits +}; + +test "Ripple's CUDA glue records: one served launch a symbol, on the ricursive image" { + for (cuda, 0..) |r, i| { + try std.testing.expect(std.mem.startsWith(u8, r.launch.symbol, "tf_rc_")); + try std.testing.expectEqualStrings("ricursive", r.launch.image); + for (cuda[i + 1 ..]) |o| try std.testing.expect(!std.mem.eql(u8, r.launch.symbol, o.launch.symbol)); + } +} diff --git a/zig/src/core/contract/ricursive_prep_records.zig b/zig/src/core/contract/ricursive_prep_records.zig new file mode 100644 index 0000000000..cc65a0f4f4 --- /dev/null +++ b/zig/src/core/contract/ricursive_prep_records.zig @@ -0,0 +1,21 @@ +//! Ripple's CUDA attention prep: one stream's rows, or each row at its own position and stream; a warp a head. +const registry = @import("registry.zig"); + +/// Usage pending: no nvcc has compiled the image yet, so every value is zero until the Spark's receipts land. +pub const cuda = [_]registry.CudaRecord{ + // TODO(spark): fill from kernel-contract-limits + .{ .launch = .{ .name = "ricursive_prep tf_rc_prep", .symbol = "tf_rc_prep", .image = "ricursive_prep", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 } } }, + // TODO(spark): fill from kernel-contract-limits + .{ .launch = .{ .name = "ricursive_prep tf_rc_prep_rows", .symbol = "tf_rc_prep_rows", .image = "ricursive_prep", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 }, .{ .regs = 0, .static = 0, .bound = 0 } } }, +}; + +test "ricursive prep CUDA: its launch is in the live registry once, at the kernel's 256 threads" { + const std = @import("std"); + const t = std.testing; + for (cuda) |r| { + var found: usize = 0; + for (registry.all_cuda) |active| found += @intFromBool(std.mem.eql(u8, active.launch.name, r.launch.name)); + try t.expectEqual(@as(usize, 1), found); + try t.expectEqual(@as(u32, 256), r.launch.threads); + } +} diff --git a/zig/src/core/contract/ricursive_records.zig b/zig/src/core/contract/ricursive_records.zig new file mode 100644 index 0000000000..6dd4c1e650 --- /dev/null +++ b/zig/src/core/contract/ricursive_records.zig @@ -0,0 +1,27 @@ +//! Ripple-65B's Metal launches (metal_ricursive*.zig): its glue, its prompt attention and the draw it runs. +const M = @import("limits.zig").MetalLaunch; +pub const records = [_]M{ + .{ .name = "ricursive glue rc_embed", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + // part_y[32] and part_r[32] floats; the source reserves its 256 threads + .{ .name = "ricursive glue rc_add_norm", .threads = 256, .tg_memory = 256, .reserved = 256 }, + .{ .name = "ricursive glue rc_prep", .threads = 32, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "ricursive glue rc_gate", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "ricursive glue rc_silu_clamp", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + // a simdgroup for each of a key head's 10 query heads + .{ .name = "ricursive glue rc_decode_vec", .threads = 320, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + // tensor ops before the M5 too (nax.h's simdgroup fragments still multiply through mpp): no offline build + .{ .name = "ricursive attention rc_prompt_attention", .threads = 64, .tg_memory = 9472, .reserved = 256, .offline = false }, + .{ .name = "ricursive fp4 rc_fp4_matvec_1", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "ricursive fp4 rc_fp4_matvec_2", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "ricursive fp4 rc_fp4_matvec_4", .threads = 256, .tg_memory = 0, .m1 = 768, .m2 = 768 }, + .{ .name = "ricursive fp4 rc_fp4_matvec_8", .threads = 256, .tg_memory = 0, .m1 = 512, .m2 = 512 }, + .{ .name = "ricursive fp4 rc_fp4_matvec_16", .threads = 256, .tg_memory = 0, .m1 = 448, .m2 = 448 }, + .{ .name = "ricursive fp4 rc_fp4_matvec_32", .threads = 256, .tg_memory = 0, .m1 = 448, .m2 = 448 }, + .{ .name = "ricursive fp4 rc_fp4_dequant", .threads = 256, .tg_memory = 0, .m1 = 832, .m2 = 832 }, + // gemm_nax's tiles: tensor ops, no offline build; the source reserves its 256 threads + .{ .name = "ricursive fp4 gemm rc_fp4_gemm_00", .threads = 256, .tg_memory = 0, .reserved = 256, .offline = false }, + .{ .name = "ricursive fp4 gemm rc_fp4_gemm_01", .threads = 256, .tg_memory = 0, .reserved = 256, .offline = false }, + .{ .name = "ricursive fp4 gemm rc_fp4_gemm_10", .threads = 256, .tg_memory = 0, .reserved = 256, .offline = false }, + .{ .name = "ricursive fp4 gemm rc_fp4_gemm_11", .threads = 256, .tg_memory = 0, .reserved = 256, .offline = false }, + .{ .name = "ricursive sample tf_sample_full", .threads = 1024, .tg_memory = 9760, .reserved = 1024, .m1 = 1024, .m2 = 1024 }, +}; diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index 08825ae3f2..154046fb92 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -179,6 +179,8 @@ pub const Info = struct { thinking_budget: bool = true, /// The model can emit a reasoning block; false: all reply text is answer content, whatever enable_thinking renders. reasoning: bool = true, + /// The reply opens its own think block or answers without one: text that does not open a block is all answer. + reasoning_optional: bool = false, /// The engine stops the reply after the think block at Request.answer_max (every lane engine: Info.ofLanes). answer_limit: bool = false, structures: bool = false, diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 9ce511329a..7e2c646519 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -76,6 +76,8 @@ pub const glm_ops: []const u8 = if (available) &Blob("fatbin_glm_ops").bytes els pub const glm_index: []const u8 = if (available) &Blob("fatbin_glm_index").bytes else &.{}; pub const minimax_m3: []const u8 = if (available) &Blob("fatbin_minimax_m3").bytes else &.{}; pub const minimax_m3_attn: []const u8 = if (available) &Blob("fatbin_minimax_m3_attn").bytes else &.{}; +pub const ricursive: []const u8 = if (available) &Blob("fatbin_ricursive").bytes else &.{}; +pub const ricursive_prep: []const u8 = if (available) &Blob("fatbin_ricursive_prep").bytes else &.{}; pub const glm_quant: []const u8 = if (available) &Blob("fatbin_glm_quant").bytes else &.{}; pub const glm_vision: []const u8 = if (available) &Blob("fatbin_glm_vision").bytes else &.{}; pub const glm_mm: []const u8 = if (available) &Blob("fatbin_glm_mm").bytes else &.{}; diff --git a/zig/src/families/ricursive/backend.zig b/zig/src/families/ricursive/backend.zig new file mode 100644 index 0000000000..d062b17a5d --- /dev/null +++ b/zig/src/families/ricursive/backend.zig @@ -0,0 +1,426 @@ +//! Ripple 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 }; + +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 stream's caches against the one budget + room: ?Room = null, // the prompt cache, which takes what the ledger leaves + 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.UnknownRicursiveStream; + } + + 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("ricursive 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.InvalidRicursivePrompt; + if (stream.images.len > 0) return error.RicursiveHasNoImageTower; + const capacity = ids.len + stream.max_new + 1; + const need = st.cacheBytes(&self.model.config, capacity); + // the stream's caches fit before anything is made; kept states trim first + const Admit = struct { + m: *Metal, + stream: *lanes.Stream, + ids: []const u32, + pub fn growth(_: *const @This()) u64 { + return 0; // no drafter: a stream's caches are all it takes + } + pub fn shed(_: *const @This(), _: u64) void {} + 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(_: *const @This()) !void {} + }; + const admit: Admit = .{ .m = self, .stream = stream, .ids = ids }; + budget.admitStream(&self.ledger, need, &admit) catch |err| { + if (err == error.LaneMemory) std.log.warn("ricursive: 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, &self.model.config, 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.InvalidRicursivePosition; + 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.InvalidRicursivePosition; + const in = switch (f) { + .handle => |h| h, + .value => |v| blk: { + if (v >= self.model.config.vocab) return error.InvalidRicursiveToken; + 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.InvalidRicursiveToken; + 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.InvalidRicursiveWindows; + 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.UnsupportedRicursiveWindow; + const cache = try self.cacheOf(w.stream); + for (w.positions, 0..) |p, r| if (p != cache.len + r + 1) return error.InvalidRicursivePosition; + 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.UnsupportedRicursiveTree; + if (path.len == 0 or path.len > w.rows() or w.rows() > cache.len) return error.UnsupportedRicursiveWindow; + 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.RicursiveHasNoDraftHead; + } + + fn heldFn(_: *anyopaque, _: *lanes.Stream, out: []u32) !void { + if (out.len > 0) return error.RicursiveDraftsNotHeld; + } + + fn releaseFn(ptr: *anyopaque, stream: *lanes.Stream) void { + const self = cast(ptr); + self.drain() catch |err| std.log.err("ricursive: 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/ricursive/budget.zig b/zig/src/families/ricursive/budget.zig new file mode 100644 index 0000000000..a7b51364c7 --- /dev/null +++ b/zig/src/families/ricursive/budget.zig @@ -0,0 +1,29 @@ +//! What a Ripple load on Metal may hold: llama's ledger and limit, and the longest context Ripple's caches fit. +const std = @import("std"); +const cfg = @import("config.zig"); +const st = @import("state.zig"); +const llama = @import("../llama/budget.zig"); + +pub const Ledger = llama.Ledger; +pub const caches = llama.caches; +pub const admitStream = llama.admitStream; + +/// The longest context whose stream (its caches and logits row) fits in `bytes`, at most `most`. +pub fn longest(c: *const cfg.Config, bytes: u64, most: usize) usize { + if (st.cacheBytes(c, 2) > bytes) return 0; + var lo: usize = 1; + var hi: usize = most; + while (lo < hi) { + const mid = lo + (hi - lo + 1) / 2; + if (st.cacheBytes(c, mid + 1) <= bytes) lo = mid else hi = mid - 1; + } + return lo; +} + +test "ricursive budget: a 262,144-token stream's caches fit 17 GiB, not 16" { + const c = try cfg.parse(std.testing.allocator, @embedFile("fixtures/config.json")); + const gib: u64 = 1 << 30; + try std.testing.expect(longest(&c, 17 * gib, 1 << 20) >= 262144); + try std.testing.expect(longest(&c, 16 * gib, 1 << 20) < 262144); + try std.testing.expectEqual(@as(usize, 0), longest(&c, 1 << 20, 1 << 20)); +} diff --git a/zig/src/families/ricursive/config.zig b/zig/src/families/ricursive/config.zig new file mode 100644 index 0000000000..87cd732932 --- /dev/null +++ b/zig/src/families/ricursive/config.zig @@ -0,0 +1,347 @@ +//! Ricursive (model_type ricursive, Ripple-65B) admission: sliding and global attention, gated, with value residuals. +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; +/// The longest context taken: the longest checked (a needle at 258,039 tokens); longer waits for a qualifying run. +pub const max_context = 262144; + +pub const Error = error{ + BadRicursiveConfig, + BadRicursiveEos, + UnsupportedRicursiveGeometry, + UnsupportedRicursiveRotary, + UnsupportedRicursivePrecision, + UnqualifiedRicursiveCheckpoint, +} || std.json.ParseError(std.json.Scanner); + +/// A checkpoint checked on the device, by the geometry its config names. +pub const Qualified = struct { + name: []const u8, + hidden: usize, + intermediate: usize, + layers: usize, + heads: usize, + kv_heads: usize, + vocab: usize, + window: usize, + theta: f64, +}; + +/// The checkpoints the family serves; another `ricursive` checkpoint is refused until it has been checked the same way. +pub const qualified = [_]Qualified{ + .{ .name = "Ripple-65B-Preview", .hidden = 5120, .intermediate = 40960, .layers = 80, .heads = 80, .kv_heads = 8, .vocab = 100352, .window = 1024, .theta = 10000 }, +}; + +/// 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, + /// Keys a sliding layer's row sees: itself and the window - 1 before it. + window: usize, + /// The embeddings' scale, as the model multiplies its bf16 rows (fp32 opmath). + embed_scale: f32, + /// The FFN's clamp: gate at most this, up within plus or minus it. + clamp: f32, + context: usize, + /// The MLP's projections are the publisher's NVFP4 (ModelOpt: E2M1 in blocks of 16, E4M3 scales, an fp32 global). + fp4_mlp: bool = false, + /// A global layer: full attention without RoPE; the others slide with RoPE. + global: [max_layers]bool = @splat(false), + 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.window == q.window and c.theta == q.theta and + c.fifthGlobal()) return q; + } + return error.UnqualifiedRicursiveCheckpoint; + } + + /// The release's layer kinds: every fifth layer global, the others sliding. + fn fifthGlobal(c: Config) bool { + for (c.global[0..c.layers], 0..) |g, i| if (g != (i % 5 == 4)) return false; + return true; + } + + 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(); + } + + /// Global layers: every position's keys kept. + pub fn globals(c: Config) usize { + var n: usize = 0; + for (c.global[0..c.layers]) |g| n += @intFromBool(g); + return n; + } + + 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 the release makes them: 1 / theta^(2i / head_dim), both steps fp32. + 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); + const p: f32 = @floatCast(std.math.pow(f64, c.theta, e)); // powf, correctly rounded + f.* = 1.0 / p; + } + } +}; + +fn object(v: std.json.Value) Error!std.json.ObjectMap { + return if (v == .object) v.object else error.BadRicursiveConfig; +} + +fn number(v: std.json.Value) Error!f64 { + return switch (v) { + .integer => |i| @floatFromInt(i), + .float => |f| f, + else => error.BadRicursiveConfig, + }; +} + +fn integer(o: std.json.ObjectMap, key: []const u8) Error!i64 { + const v = o.get(key) orelse return error.BadRicursiveConfig; + return if (v == .integer) v.integer else error.BadRicursiveConfig; +} + +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 the JSON boolean `want`. +fn isBool(o: std.json.ObjectMap, key: []const u8, want: bool) bool { + const v = o.get(key) orelse return false; + return v == .bool and v.bool == want; +} + +/// 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.UnsupportedRicursiveGeometry; +} + +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.BadRicursiveEos; + c.eos[c.eos_count] = @intCast(i); + c.eos_count += 1; + }, + .array => |a| for (a.items) |x| try eosIds(x, c), + else => return error.BadRicursiveEos, + } +} + +/// layer_types and hybrid_layer_pattern (1 sliding, 0 global) name the same layers, one a layer. +fn layerKinds(root: std.json.ObjectMap, c: *Config) Error!void { + const types = root.get("layer_types") orelse return error.BadRicursiveConfig; + const pattern = root.get("hybrid_layer_pattern") orelse return error.BadRicursiveConfig; + if (types != .array or pattern != .array) return error.BadRicursiveConfig; + if (types.array.items.len != c.layers or pattern.array.items.len != c.layers) return error.UnsupportedRicursiveGeometry; + for (types.array.items, pattern.array.items, 0..) |t, p, i| { + const global = isString(t, "full_attention"); + if (!global and !isString(t, "sliding_attention")) return error.UnsupportedRicursiveGeometry; + if (p != .integer or p.integer != @intFromBool(!global)) return error.UnsupportedRicursiveGeometry; + c.global[i] = global; + } +} + +/// ModelOpt NVFP4, groups of 16, no cache scheme, ignoring exactly the embeddings, head and every layer's attention. +fn nvfp4Mlp(q: std.json.Value, layers: usize) Error!bool { + const o = try object(q); + if (!isString(o.get("quant_method"), "modelopt") or !isString(o.get("quant_algo"), "NVFP4")) return error.UnsupportedRicursivePrecision; + for ([_][]const u8{ "kv_cache_quant_algo", "kv_cache_scheme" }) |key| if (o.get(key)) |v| if (v != .null) return error.UnsupportedRicursivePrecision; + const groups = try object(o.get("config_groups") orelse return error.UnsupportedRicursivePrecision); + if (groups.count() != 1) return error.UnsupportedRicursivePrecision; + const g = try object(groups.values()[0]); + const w = try object(g.get("weights") orelse return error.UnsupportedRicursivePrecision); + if (!isString(w.get("type"), "float") or (number(w.get("num_bits") orelse return error.UnsupportedRicursivePrecision) catch 0) != 4 or + (number(w.get("group_size") orelse return error.UnsupportedRicursivePrecision) catch 0) != 16) return error.UnsupportedRicursivePrecision; + const ignore = switch (o.get("ignore") orelse return error.UnsupportedRicursivePrecision) { + .array => |a| a.items, + else => return error.UnsupportedRicursivePrecision, + }; + var attention: [max_layers]bool = @splat(false); + for (ignore) |item| { + const name = switch (item) { + .string => |t| t, + else => return error.UnsupportedRicursivePrecision, + }; + if (std.mem.eql(u8, name, "lm_head") or std.mem.eql(u8, name, "model.embed_tokens")) continue; + const prefix = "model.layers."; + const suffix = ".self_attn*"; + if (!std.mem.startsWith(u8, name, prefix) or !std.mem.endsWith(u8, name, suffix) or name.len <= prefix.len + suffix.len) return error.UnsupportedRicursivePrecision; + const l = std.fmt.parseInt(usize, name[prefix.len .. name.len - suffix.len], 10) catch return error.UnsupportedRicursivePrecision; + if (l >= layers or l >= max_layers) return error.UnsupportedRicursivePrecision; + attention[l] = true; + } + // every layer's attention stays bf16, as the kernels read it + for (attention[0..@min(layers, max_layers)]) |kept| if (!kept) return error.UnsupportedRicursivePrecision; + return true; +} + +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"), "ricursive")) return error.BadRicursiveConfig; + // the release's layer: gated attention (an fp32 gate), value residuals, sliding and global layers, an untied head + for ([_][]const u8{ "use_gated_attention", "gate_full_precision", "use_valres", "is_hybrid_swa" }) |key| { + if (!isBool(root, key, true)) return error.UnsupportedRicursiveGeometry; + } + if (!isBool(root, "tie_word_embeddings", false)) return error.UnsupportedRicursiveGeometry; + if (!isString(root.get("hidden_act"), "silu")) return error.UnsupportedRicursiveGeometry; + for ([_][]const u8{ "torch_dtype", "dtype" }) |key| if (root.get(key)) |d| if (!isString(d, "bfloat16")) return error.UnsupportedRicursivePrecision; + if (root.get("quantization") != null) return error.UnsupportedRicursivePrecision; + if (root.get("rope_scaling")) |r| if (r != .null) return error.UnsupportedRicursiveRotary; + 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 = try size(root, "num_key_value_heads"), + .vocab = try size(root, "vocab_size"), + .eps = @floatCast(try number(root.get("rms_norm_eps") orelse return error.BadRicursiveConfig)), + .theta = try number(root.get("rope_theta") orelse return error.BadRicursiveConfig), + .window = try size(root, "sliding_window"), + .embed_scale = @floatCast(try number(root.get("embed_scale") orelse return error.BadRicursiveConfig)), + .clamp = @floatCast(try number(root.get("ffn_clamp_threshold") orelse return error.BadRicursiveConfig)), + .context = 0, + }; + if (root.get("quantization_config")) |q| c.fp4_mlp = try nvfp4Mlp(q, c.layers); + if (try size(root, "head_dim") != head_dim) return error.UnsupportedRicursiveGeometry; + if (c.layers > max_layers or c.heads % c.kv_heads != 0 or c.group() > max_group) return error.UnsupportedRicursiveGeometry; + // 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.queryWidth() % 256 != 0) return error.UnsupportedRicursiveGeometry; + if (!(c.eps > 0) or !(c.theta > 0) or !(c.embed_scale > 0) or !(c.clamp > 0)) return error.BadRicursiveConfig; + try layerKinds(root, &c); + c.context = @min(try size(root, "max_position_embeddings"), max_context); + try eosIds(root.get("eos_token_id") orelse return error.BadRicursiveEos, &c); + if (c.eos_count == 0) return error.BadRicursiveEos; + return c; +} + +const fixture = @embedFile("fixtures/config.json"); + +test "Ripple-65B-Preview's config is admitted with its geometry, layer kinds and end tokens" { + const c = try parse(std.testing.allocator, fixture); + try std.testing.expectEqualStrings("Ripple-65B-Preview", (try c.qualification()).name); + try std.testing.expectEqual(@as(usize, 5120), c.hidden); + try std.testing.expectEqual(@as(usize, 40960), c.intermediate); + try std.testing.expectEqual(@as(usize, 80), c.layers); + try std.testing.expectEqual(@as(usize, 10240), c.queryWidth()); + try std.testing.expectEqual(@as(usize, 12288), c.qkvWidth()); + try std.testing.expectEqual(@as(usize, 10), c.group()); + try std.testing.expectEqual(@as(usize, 1024), c.window); + try std.testing.expectEqual(@as(usize, 16), c.globals()); + for (0..c.layers) |i| try std.testing.expectEqual(i % 5 == 4, c.global[i]); + try std.testing.expectEqual(@as(f32, 20), c.clamp); + try std.testing.expectEqual(@as(f32, @floatCast(71.55417527999327)), c.embed_scale); + try std.testing.expectEqual(@as(usize, max_context), c.context); + try std.testing.expectEqualSlices(u32, &.{ 100265, 100257 }, c.eos[0..c.eos_count]); +} + +test "ricursive config: another geometry or layer is refused" { + const gpa = std.testing.allocator; + const refused = .{ + .{ "\"use_valres\": true", "\"use_valres\": false", error.UnsupportedRicursiveGeometry }, + .{ "\"gate_full_precision\": true", "\"gate_full_precision\": false", error.UnsupportedRicursiveGeometry }, + .{ "\"tie_word_embeddings\": false", "\"tie_word_embeddings\": true", error.UnsupportedRicursiveGeometry }, + .{ "\"head_dim\": 128", "\"head_dim\": 64", error.UnsupportedRicursiveGeometry }, + .{ "\"hidden_act\": \"silu\"", "\"hidden_act\": \"gelu\"", error.UnsupportedRicursiveGeometry }, + .{ "\"model_type\": \"ricursive\"", "\"model_type\": \"llama\"", error.BadRicursiveConfig }, + .{ "\"rope_theta\": 10000.0", "\"rope_theta\": 10000.0, \"rope_scaling\": {\"rope_type\": \"yarn\"}", error.UnsupportedRicursiveRotary }, + .{ "\"layer_types\": [\n \"sliding_attention\"", "\"layer_types\": [\n \"full_attention\"", error.UnsupportedRicursiveGeometry }, + }; + inline for (refused) |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)); + try std.testing.expectError(k[2], parse(gpa, bytes)); + } + const other = try std.mem.replaceOwned(u8, gpa, fixture, "\"sliding_window\": 1024", "\"sliding_window\": 4096"); + defer gpa.free(other); + try std.testing.expectError(error.UnqualifiedRicursiveCheckpoint, (try parse(gpa, other)).qualification()); +} + +/// The fixture with an NVFP4 quantization_config: `extra` in it, every layer's attention ignored but `skip`'s. +fn withNvfp4(gpa: std.mem.Allocator, extra: []const u8, skip: ?usize) ![]u8 { + var ignore: std.ArrayList(u8) = .empty; + defer ignore.deinit(gpa); + try ignore.appendSlice(gpa, "\"lm_head\", \"model.embed_tokens\""); + for (0..80) |l| if (l != skip) try ignore.print(gpa, ", \"model.layers.{d}.self_attn*\"", .{l}); + const q = try std.fmt.allocPrint(gpa, "\"model_type\": \"ricursive\", \"quantization_config\": {{{s}\"quant_method\": \"modelopt\", \"quant_algo\": \"NVFP4\", \"config_groups\": {{\"group_0\": {{\"weights\": {{\"type\": \"float\", \"num_bits\": 4, \"group_size\": 16}}}}}}, \"ignore\": [{s}]}}", .{ extra, ignore.items }); + defer gpa.free(q); + return std.mem.replaceOwned(u8, gpa, fixture, "\"model_type\": \"ricursive\"", q); +} + +test "ricursive config: the NVFP4 release's scheme is admitted, a cache scheme or a quantized attention refused" { + const gpa = std.testing.allocator; + const ok = try withNvfp4(gpa, "", null); + defer gpa.free(ok); + try std.testing.expect((try parse(gpa, ok)).fp4_mlp); + const cache = try withNvfp4(gpa, "\"kv_cache_quant_algo\": \"FP8\", ", null); + defer gpa.free(cache); + try std.testing.expectError(error.UnsupportedRicursivePrecision, parse(gpa, cache)); + const attention = try withNvfp4(gpa, "", 41); + defer gpa.free(attention); + try std.testing.expectError(error.UnsupportedRicursivePrecision, parse(gpa, attention)); +} + +test "ricursive config: another layer pattern with the release's geometry is not qualified" { + const gpa = std.testing.allocator; + const types = try std.mem.replaceOwned(u8, gpa, fixture, "\"layer_types\": [\n \"sliding_attention\"", "\"layer_types\": [\n \"full_attention\""); + defer gpa.free(types); + const both = try std.mem.replaceOwned(u8, gpa, types, "\"hybrid_layer_pattern\": [\n 1,", "\"hybrid_layer_pattern\": [\n 0,"); + defer gpa.free(both); + try std.testing.expectError(error.UnqualifiedRicursiveCheckpoint, (try parse(gpa, both)).qualification()); +} diff --git a/zig/src/families/ricursive/contracts.zig b/zig/src/families/ricursive/contracts.zig new file mode 100644 index 0000000000..5c4ebe4e27 --- /dev/null +++ b/zig/src/families/ricursive/contracts.zig @@ -0,0 +1,84 @@ +//! The family's Metal kernels under the kernel contract: their sources as model.zig 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 st = @import("state.zig"); +const frags = @import("../../core/frags.zig"); + +/// Ripple-65B, the family's one qualified checkpoint. +pub const ripple = cfg.qualified[0]; +/// The engine's sliding ring for a stream past it (state.ring at any long capacity). +pub const sliding_ring = std.mem.alignForward(usize, ripple.window - 1 + st.prompt_rows, st.ring_chunk); + +pub const glue_names = model.glue_names; +pub const llama_names = model.llama_names; +pub const gemm_names = [_][]const u8{ "ll_gemm_00", "ll_gemm_01", "ll_gemm_10", "ll_gemm_11" }; +pub const attention_name = "rc_prompt_attention"; +pub const sample_name = "tf_sample_full"; +pub const fp4_names = model.fp4_names; +pub const fp4_gemm_names = [_][]const u8{ "rc_fp4_gemm_00", "rc_fp4_gemm_01", "rc_fp4_gemm_10", "rc_fp4_gemm_11" }; +/// Every pipeline model.zig builds: glue, llama's row kernels and GEMMs, attention, the draw and the NVFP4 MLP's. +pub const pipelines = glue_names.len + llama_names.len + gemm_names.len + 2 + fp4_names.len + fp4_gemm_names.len; + +/// The glue's text as model.zig builds it. +pub fn glueSource(a: std.mem.Allocator) ![]u8 { + return model.glueSource(a, "#define RC_DH {d}\n{s}", ks.ricursive_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 header as model.zig writes it for Ripple. +pub fn attentionHeader(a: std.mem.Allocator) ![]u8 { + const fmt = "#define DH {d}\n#define NQ {d}\n#define G {d}\n#define KVH {d}\n#define SCALE {e}f\n"; + return std.fmt.allocPrint(a, fmt, .{ cfg.head_dim, ripple.heads, ripple.heads / ripple.kv_heads, ripple.kv_heads, cfg.headScale() }); +} + +/// The prompt attention's text as model.zig builds it (nax.h inlined for this GPU). +pub fn attentionSource(device: mtl.Device, a: std.mem.Allocator) ![]u8 { + const head = try attentionHeader(a); + defer a.free(head); + const body = try frags.source(device, a, ks.ricursive_attention_nax); + defer a.free(body); + return std.mem.concat(a, u8, &.{ head, body }); +} + +/// The NVFP4 row kernels' library as model.zig compiles it: precise math functions. +pub fn fp4Library(device: mtl.Device, src: []const u8) !mtl.Library { + return mtl.Library.fromSource(device, src, .{ .language = mtl.CompileOptions.mlx().language, .math_functions = .precise }); +} + +/// The NVFP4 GEMM's text: nax.h inlined for this GPU (model.zig's build), or as a GPU without tensor units takes it. +pub fn fp4GemmSource(device: mtl.Device, a: std.mem.Allocator, units: bool) ![]u8 { + if (units) return frags.source(device, a, ks.ricursive_fp4_gemm); + return frags.precision.source(a, false, ks.ricursive_fp4_gemm, ks.nax, false); +} + +/// Every pipeline the family runs appears among the registrations' functions. +pub fn requireComplete(functions: []const []const u8) !void { + inline for (.{ glue_names, llama_names, gemm_names, [_][]const u8{ attention_name, sample_name }, fp4_names, fp4_gemm_names }) |list| for (list) |name| { + var seen = false; + for (functions) |f| if (std.mem.eql(u8, f, name)) { + seen = true; + }; + if (!seen) return error.IncompleteRicursiveContract; + }; +} + +test "ricursive contracts: the roll call wants every pipeline" { + var all: [pipelines][]const u8 = undefined; + var i: usize = 0; + inline for (.{ glue_names, llama_names, gemm_names, [_][]const u8{ attention_name, sample_name }, fp4_names, fp4_gemm_names }) |list| for (list) |name| { + all[i] = name; + i += 1; + }; + try requireComplete(&all); + try std.testing.expectError(error.IncompleteRicursiveContract, requireComplete(all[1..])); + try std.testing.expectEqual(@as(usize, 3072), sliding_ring); +} diff --git a/zig/src/families/ricursive/cuda.zig b/zig/src/families/ricursive/cuda.zig new file mode 100644 index 0000000000..98b3e09714 --- /dev/null +++ b/zig/src/families/ricursive/cuda.zig @@ -0,0 +1,19 @@ +//! Ripple on CUDA: the family's kernels (glue, attention prep, attention), weights and forward, beside Metal's. + +pub const kernels = @import("cuda_kernels.zig"); +pub const prep = @import("cuda_prep.zig"); +pub const attn = @import("cuda_attn.zig"); +pub const weights = @import("cuda_weights.zig"); +pub const model = @import("cuda_model.zig"); +pub const backend = @import("cuda_backend.zig"); +pub const native = @import("cuda_native.zig"); + +test { + _ = kernels; + _ = prep; + _ = attn; + _ = weights; + _ = model; + _ = backend; + _ = native; +} diff --git a/zig/src/families/ricursive/cuda_attn.zig b/zig/src/families/ricursive/cuda_attn.zig new file mode 100644 index 0000000000..d05d73627c --- /dev/null +++ b/zig/src/families/ricursive/cuda_attn.zig @@ -0,0 +1,397 @@ +//! Ripple's attention on CUDA (group 10, dim 128): window_attention's prompt and decode, tree_attention's global. +const std = @import("std"); +const cuda = @import("cuda"); + +const wa = cuda.window_attention; +const ta = cuda.tree_attention; +const Stream = cuda.Stream; +const Function = cuda.Function; + +pub const head_dim = 128; +/// Query heads a block: Ripple's group 10 is no multiple of the shared prompt's 4 (WindowAttention.prompt refuses it). +pub const heads_a_block = 2; +/// Eight warps, each 16 rows of one head: 64 rows a block. +pub const rows_a_block = 16 * (8 / heads_a_block); + +/// wattn_kernel<128, 8, 2, 8>: keys in a ring (sliding layers) or linear (global layers). +pub const prompt_symbol = "_ZN16tf_window_prompt12wattn_kernelILi128ELi8ELi2ELi8EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifii"; +/// The same over a stream's pages (global layers only: no window). +pub const prompt_pg_symbol = "tf_window_prompt_h2_pg"; +/// window_attention.zig's prompt_smem: eight staging slots of 32 keys at head dim 128. +const prompt_smem: u32 = 8 * 32 * 16 * 16; + +pub const Kernels = struct { + prompt: Function, + prompt_pg: Function, + + /// The two instances from the window_attention image WindowAttention.load loaded (it owns the module). + pub fn init(w: *const wa.WindowAttention) !Kernels { + const p = try w.mod.function(prompt_symbol); + try p.allowDynamicShared(prompt_smem); + const pg = try w.mod.function(prompt_pg_symbol); + try pg.allowDynamicShared(prompt_smem); + return .{ .prompt = p, .prompt_pg = pg }; + } +}; + +/// A sliding layer: keys (pos - window, pos] of a ring holding key p at row p % ring (prompt and decode). +pub fn sliding(h: usize, hk: usize, window: usize, ring: usize, scale: f32) wa.Window { + return .{ .h = h, .hk = hk, .d = head_dim, .win = window, .ring = ring, .scale = scale }; +} + +/// A global layer's prompt rows: every key up to the row's own, linear (key p at row p) or in pages. +pub fn global(h: usize, hk: usize, scale: f32, pg: ?wa.Pages) wa.Window { + return .{ .h = h, .hk = hk, .d = head_dim, .win = 0, .ring = 0, .scale = scale, .pg = pg }; +} + +/// A ring holds every key a call's rows see: it never wraps in the call, or it takes the window and the rows. +fn ringHolds(w: wa.Window, p0: usize, n: usize) bool { + if (w.ring == 0 or p0 + n <= w.ring) return true; + return w.win > 0 and w.win - 1 + n <= w.ring; +} + +/// The shapes the two-heads-a-block kernels take; the shared prompt's checks plus the ring's. +pub fn check(w: wa.Window, p0: usize, n: usize) error{WindowShape}!void { + if (w.d != head_dim or w.hk == 0 or w.h % w.hk != 0 or (w.h / w.hk) % heads_a_block != 0) return error.WindowShape; + if (p0 + n > std.math.maxInt(i32) or w.win > std.math.maxInt(i32) or w.ring > std.math.maxInt(i32)) return error.WindowShape; + if (w.pg) |pg| if (w.win != 0 or w.ring != 0 or pg.tokens == 0 or pg.page_elems > std.math.maxInt(i63)) return error.WindowShape; + if (!ringHolds(w, p0, n)) return error.WindowShape; +} + +/// The launch: (rows / 64) x (key heads x head pairs) blocks of 256 threads, 64 KiB of staged keys. +pub fn config(w: wa.Window, n: usize) cuda.Config { + const pairs = w.h / w.hk / heads_a_block; + return .{ .grid = .{ .x = @intCast((n + rows_a_block - 1) / rows_a_block), .y = @intCast(w.hk * pairs), .z = 1 }, .block = .{ .x = 256 }, .shared = prompt_smem }; +} + +/// n prompt rows of one stream at p0, p0 + 1, ...: q [n, h, 128] -> out [n, h, 128]; kc/vc as WindowAttention.prompt. +pub fn prompt(k: *const Kernels, s: Stream, q: u64, kc: u64, vc: u64, out: u64, p0: usize, n: usize, w: wa.Window) !void { + try check(w, p0, n); + if (n == 0) return; + var a: cuda.Args = .{}; + for ([_]u64{ q, kc, vc, out }) |v| a.add(v); + for ([_]usize{ p0, n, w.h, w.hk, w.h / w.hk }) |v| a.add(@as(i32, @intCast(v))); + a.add(w.scale); + a.add(@as(i32, @intCast(w.win))); + a.add(@as(i32, @intCast(w.ring))); + if (w.pg) |pg| { + a.add(pg.table); + a.add(@as(i64, @intCast(pg.page_elems))); + a.add(@as(i32, @intCast(pg.tokens))); + } + try cuda.launch.launch(if (w.pg != null) k.prompt_pg else k.prompt, config(w, n), s, &a); +} + +/// Global layers' heads: 80 query heads over 8 key heads (group 10). +pub const global_heads = 80; +pub const global_kv_heads = 8; + +/// A call's plan for tree_attention, int32 row_stream [m], streams [m, 4], depths [m], then int64 offs [m, 2]. +fn planBytes(max_rows: usize) usize { + return max_rows * (6 * 4 + 2 * 8); +} + +/// Row r its own stream: committed keys 0..pos + r from the pages, depth 0 (no path keys); K, V from `origin`. +fn plan(buf: []u8, max_rows: usize, pos: usize, rows: usize, k_off: i64, v_off: i64) void { + for (0..rows) |r| planRow(buf, max_rows, r, pos + r + 1, k_off, v_off); +} + +/// Row r its own stream (and its call table row), `committed` keys from the pages, depth 0; K, V from `origin`. +fn planRow(buf: []u8, max_rows: usize, r: usize, committed: usize, k_off: i64, v_off: i64) void { + const ints = std.mem.bytesAsSlice(i32, buf[0 .. 24 * max_rows]); + const offs = std.mem.bytesAsSlice(i64, buf[24 * max_rows ..][0 .. 16 * max_rows]); + ints[r] = @intCast(r); + ints[max_rows + 4 * r ..][0..4].* = .{ @intCast(r), 1, @intCast(committed), 0 }; + ints[5 * max_rows + r] = 0; + offs[2 * r] = k_off; + offs[2 * r + 1] = v_off; +} + +/// Every row's fixed fields, staged once: its own stream and table row, K at the pool, V `v_off` past it, committed 0. +fn stagePlan(buf: []u8, max_rows: usize, v_off: i64) void { + for (0..max_rows) |r| planRow(buf, max_rows, r, 0, 0, v_off); +} + +/// Row r's committed keys (pos[r] + 1) into the staged plan; returns the streams' bytes to upload, 16 a row. +fn commitRows(buf: []u8, max_rows: usize, pos: []const usize) []const u8 { + const streams = std.mem.bytesAsSlice(i32, buf[4 * max_rows .. 20 * max_rows]); + for (pos, 0..) |p, r| streams[4 * r + 2] = @intCast(p + 1); + return buf[4 * max_rows ..][0 .. 16 * pos.len]; +} + +/// tree_attention's fp32 partials for `max_rows` rows over keys to `context`: o, max, sum by group, then slices. +pub fn scratchBytes(max_rows: usize, context: usize) usize { + return ta.TreeAttention.scratchBytes(max_rows, global_heads, global_kv_heads, ta.TreeAttention.groups(context), head_dim); +} + +/// The device's SMs: tree_attention takes one-warp blocks for small grids (the same bits as four-warp blocks). +fn smCount(d: *const cuda.Driver) !usize { + var dev: cuda.abi.Device = 0; + try d.check(d.api.cuCtxGetDevice(&dev), "cuCtxGetDevice"); + var sms: c_int = 0; + try d.check(d.api.cuDeviceGetAttribute(&sms, .multiprocessor_count, dev), "cuDeviceGetAttribute"); + return @intCast(@max(sms, 1)); +} + +/// Global decode: tree_attention's paged kernels and merge, each row a stream of its own over the layer's pages. +pub const GlobalDecode = struct { + tree: ta.TreeAttention, + plan: cuda.DeviceBuffer, + scratch: cuda.DeviceBuffer, + host: []u8, + gpa: std.mem.Allocator, + max_rows: usize, + /// The most keys a row sees (pos + rows): the partials are strided by its groups. + context: usize, + /// V's offset from K that `stage` uploaded (null: not staged); decodeGlobal's plan overwrites the staged fields. + v_off: ?usize = null, + + /// The device bytes init allocates: the partials and the plan. + pub fn bytes(max_rows: usize, context: usize) usize { + return scratchBytes(max_rows, context) + std.mem.alignForward(usize, planBytes(max_rows), 256); + } + + pub fn init(d: *const cuda.Driver, gpa: std.mem.Allocator, max_rows: usize, context: usize) !GlobalDecode { + if (max_rows == 0 or context < max_rows or context > std.math.maxInt(i32)) return error.GlobalShape; + var tree = try ta.TreeAttention.load(d); + errdefer tree.unload(); + tree.sms = try smCount(d); + const host = try gpa.alloc(u8, planBytes(max_rows)); + errdefer gpa.free(host); + @memset(host, 0); + var p = try cuda.DeviceBuffer.alloc(d, host.len); + errdefer p.free(); + const sc = try cuda.DeviceBuffer.alloc(d, scratchBytes(max_rows, context)); + return .{ .tree = tree, .plan = p, .scratch = sc, .host = host, .gpa = gpa, .max_rows = max_rows, .context = context }; + } + + pub fn deinit(g: *GlobalDecode) void { + g.scratch.free(); + g.plan.free(); + g.gpa.free(g.host); + g.tree.unload(); + } + + /// The plan's fixed fields for every row, uploaded once (stagePlan): a round then uploads only its committed keys. + pub fn stage(g: *GlobalDecode, s: Stream, v_off: usize) !void { + if (v_off > std.math.maxInt(i63)) return error.GlobalShape; + stagePlan(g.host, g.max_rows, @intCast(v_off)); + try g.plan.uploadAsync(0, g.host, s.handle); + g.v_off = v_off; + } + + /// The key groups a graph's rounds take: those of the context, so its grid and partials' strides serve every round. + pub fn groups(g: *const GlobalDecode) usize { + return ta.TreeAttention.groups(g.context); + } +}; + +/// q [rows, 80, 128], row r at pos + r over keys 0..pos + r of `pages` (kc, vc: tf_rc_prep's pools) -> out. +pub const DecodeArgs = struct { q: u64, kc: u64, vc: u64, out: u64, pos: usize, rows: usize, pages: wa.Pages, scale: f32 }; + +/// The calls the plan takes: at most init's rows, keys to its context, pages of whole int32 positions. +pub fn admit(max_rows: usize, context: usize, a: DecodeArgs) error{GlobalShape}!void { + if (a.rows > max_rows or a.pos > context or a.rows > context - a.pos) return error.GlobalShape; + const pg = a.pages; + if (pg.tokens == 0 or pg.tokens > std.math.maxInt(i32) or pg.page_elems == 0 or pg.page_elems > std.math.maxInt(i63)) return error.GlobalShape; +} + +/// One global layer's decode rows. The plan is uploaded each call (not inside a CUDA graph capture). +pub fn decodeGlobal(g: *const GlobalDecode, s: Stream, a: DecodeArgs) !void { + return decodeHeads(g, s, a, global_heads); +} + +/// decodeGlobal with `h` query heads over the 8 key heads: Ripple's 80, or a control's other group. +pub fn decodeHeads(g: *const GlobalDecode, s: Stream, a: DecodeArgs, h: usize) !void { + try admit(g.max_rows, g.context, a); + if (h == 0 or h > global_heads or h % global_kv_heads != 0) return error.GlobalShape; + if (a.rows == 0) return; + const origin = @min(a.kc, a.vc); + plan(g.host, g.max_rows, a.pos, a.rows, @intCast((a.kc - origin) / 2), @intCast((a.vc - origin) / 2)); + try g.plan.uploadAsync(0, g.host, s.handle); + const pg: ta.Pages = .{ .table = a.pages.table, .page = a.pages.page_elems, .tok = a.pages.tokens, .stride = 0 }; + try run(g, s, a.q, a.kc, a.vc, a.out, a.rows, h, pg, ta.TreeAttention.groups(a.pos + a.rows), a.scale); +} + +/// A round's global-layer plan: row r its own tree stream at pos[r], K at `kc`, V `v_off` past it, `ng` key groups. +pub const Planned = struct { rows: usize, ng: usize, v_off: usize }; + +/// The rounds the plan takes: 1 to init's rows, each position within its context, V offsets of int64 elements. +pub fn admitRows(max_rows: usize, context: usize, pos: []const usize, v_off: usize) error{GlobalShape}!void { + if (pos.len == 0 or pos.len > max_rows or v_off > std.math.maxInt(i63)) return error.GlobalShape; + for (pos) |p| if (p >= context) return error.GlobalShape; +} + +/// The groups a call's partials are strided by: those of its row with the most keys. +pub fn groupsOf(pos: []const usize) usize { + var most: usize = 0; + for (pos) |p| most = @max(most, p + 1); + return ta.TreeAttention.groups(most); +} + +/// tree_attention.zig's choice for a call: one-warp blocks (true) or four-warp blocks, by rows, groups and key heads. +pub fn oneWarp(g: *const GlobalDecode, rows: usize, ng: usize) bool { + return g.tree.sms > 0 and rows * ng * global_kv_heads <= ta.slice_items; +} + +/// A round's rows into the staged plan, each a one-row call's stream; `fixed`: the context's key groups, same bits. +pub fn planRows(g: *const GlobalDecode, s: Stream, pos: []const usize, fixed: bool) !Planned { + const v_off = g.v_off orelse return error.PlanNotStaged; + try admitRows(g.max_rows, g.context, pos, v_off); + try g.plan.uploadAsync(4 * g.max_rows, commitRows(g.host, g.max_rows, pos), s.handle); + return .{ .rows = pos.len, .ng = if (fixed) g.groups() else groupsOf(pos), .v_off = v_off }; +} + +/// One global layer's planned rows, `h` query heads over the 8 key heads: q [rows, h, 128] -> out [rows, h, 128]. +pub fn runRows(g: *const GlobalDecode, s: Stream, p: Planned, q: u64, kc: u64, out: u64, h: usize, pages: wa.Pages, scale: f32) !void { + if (h == 0 or h > global_heads or h % global_kv_heads != 0) return error.GlobalShape; + if (pages.tokens == 0 or pages.tokens > std.math.maxInt(i32) or pages.page_elems == 0 or pages.page_elems > std.math.maxInt(i63) or pages.stride > std.math.maxInt(i32)) return error.GlobalShape; + const pg: ta.Pages = .{ .table = pages.table, .page = pages.page_elems, .tok = pages.tokens, .stride = pages.stride }; + try run(g, s, q, kc, kc + 2 * p.v_off, out, p.rows, h, pg, p.ng, scale); +} + +/// tree_attention over the uploaded plan: `n` rows, `ng` groups of partials, the pages `pg`. +fn run(g: *const GlobalDecode, s: Stream, q: u64, kc: u64, vc: u64, out: u64, n: usize, h: usize, pg: ta.Pages, ng: usize, scale: f32) !void { + const m = g.max_rows; + const base = g.plan.ptr; + const depths = base + 20 * m; // also `paths`: depth 0 reads none + const rows: ta.Rows = .{ .row_stream = base, .streams = base + 4 * m, .offs = base + 24 * m, .origin = @min(kc, vc), .paths = depths, .depths = depths, .pages = pg }; + const sc = ta.TreeAttention.scratch(g.scratch.ptr, n, h, global_kv_heads, ng, head_dim); + // kn, vn: path keys, of which there are none (each row's own key is committed in the pages) + try g.tree.run(s, q, kc, vc, rows, sc, out, n, h, global_kv_heads, head_dim, ng, scale); +} + +test "Ripple's prompt grid: 64 rows a block, five head pairs of each of 8 key heads" { + const t = std.testing; + const w = sliding(80, 8, 1024, 3072, 0.08838835); + for ([_][2]u32{ .{ 2048, 32 }, .{ 1000, 16 }, .{ 333, 6 }, .{ 1, 1 } }) |c| { + const cfg = config(w, c[0]); + try t.expectEqual(c[1], @as(u32, cfg.grid.x)); + try t.expectEqual(@as(u32, 40), @as(u32, cfg.grid.y)); + try t.expectEqual(@as(u32, 256), @as(u32, cfg.block.x)); + try t.expectEqual(@as(u32, 65536), cfg.shared); + } +} + +test "the ring takes the window and a 2,048-row chunk past its wrap; a longer chunk or a short ring is refused" { + const t = std.testing; + const w = sliding(80, 8, 1024, 3072, 0.08838835); + try check(w, 8216, 2048); + try check(w, 0, 3072); // no wrap in the call + try t.expectError(error.WindowShape, check(w, 8216, 2050)); + try t.expectError(error.WindowShape, check(sliding(80, 8, 1024, 3000, 0.08838835), 8216, 2048)); + try check(global(80, 8, 0.08838835, null), 6000, 2048); // linear keys: no ring + var every = global(80, 8, 0.08838835, null); + every.ring = 4096; + try t.expectError(error.WindowShape, check(every, 4000, 200)); // every key, from a ring that wraps +} + +test "shapes the kernel does not take are refused" { + const t = std.testing; + try t.expectError(error.WindowShape, check(sliding(72, 8, 1024, 3072, 1), 0, 16)); // group 9 + var wide = sliding(80, 8, 1024, 3072, 1); + wide.d = 256; + try t.expectError(error.WindowShape, check(wide, 0, 16)); + var pg = sliding(80, 8, 1024, 3072, 1); + pg.pg = .{ .table = 0, .page_elems = 2 * 512 * 8 * 128, .tokens = 512 }; + try t.expectError(error.WindowShape, check(pg, 0, 16)); // pages hold global layers' keys only + try check(global(80, 8, 1, .{ .table = 0, .page_elems = 2 * 512 * 8 * 128, .tokens = 512 }), 6000, 2048); +} + +test "global decode plan: each row its own stream, committed to its own key, no path keys" { + const t = std.testing; + var buf: [planBytes(4)]u8 = @splat(0xee); + plan(&buf, 4, 510, 3, 0, 2048); + const ints = std.mem.bytesAsSlice(i32, buf[0..96]); + const offs = std.mem.bytesAsSlice(i64, buf[96..]); + for (0..3) |r| { + const ri: i32 = @intCast(r); + try t.expectEqual(ri, ints[r]); + try t.expectEqualSlices(i32, &.{ ri, 1, 511 + ri, 0 }, &[4]i32{ ints[4 + 4 * r], ints[5 + 4 * r], ints[6 + 4 * r], ints[7 + 4 * r] }); + try t.expectEqual(@as(i32, 0), ints[20 + r]); + try t.expectEqual(@as(i64, 0), offs[2 * r]); + try t.expectEqual(@as(i64, 2048), offs[2 * r + 1]); + } + try t.expectEqual(@as(i32, @bitCast(@as(u32, 0xeeeeeeee))), ints[3]); // row 3: past the call's rows, untouched +} + +test "global decode takes 1 to 16 rows to the context, and pages of whole positions" { + const t = std.testing; + const pg: wa.Pages = .{ .table = 0, .page_elems = 2 * 512 * 8 * 128, .tokens = 512 }; + const at = struct { + fn f(pos: usize, rows: usize, p: wa.Pages) DecodeArgs { + return .{ .q = 0, .kc = 0, .vc = 0, .out = 0, .pos = pos, .rows = rows, .pages = p, .scale = 0.08838835 }; + } + }.f; + try admit(16, 131072, at(131056, 16, pg)); + try admit(16, 131072, at(0, 1, pg)); + try t.expectError(error.GlobalShape, admit(16, 131072, at(131057, 16, pg))); // past the context + try t.expectError(error.GlobalShape, admit(16, 131072, at(200000, 1, pg))); + try t.expectError(error.GlobalShape, admit(16, 131072, at(0, 17, pg))); + var none = pg; + none.tokens = 0; + try t.expectError(error.GlobalShape, admit(16, 131072, at(0, 1, none))); +} + +test "global decode rows: each row its own stream and call table row, committed to its own position" { + const t = std.testing; + var buf: [planBytes(4)]u8 = @splat(0xee); + const pos = [_]usize{ 4095, 511, 3071 }; + for (pos, 0..) |p, r| planRow(&buf, 4, r, p + 1, 0, 4608); + const ints = std.mem.bytesAsSlice(i32, buf[0..96]); + const offs = std.mem.bytesAsSlice(i64, buf[96..]); + for (pos, 0..) |p, r| { + const ri: i32 = @intCast(r); + try t.expectEqual(ri, ints[r]); + try t.expectEqualSlices(i32, &.{ ri, 1, @as(i32, @intCast(p + 1)), 0 }, &[4]i32{ ints[4 + 4 * r], ints[5 + 4 * r], ints[6 + 4 * r], ints[7 + 4 * r] }); + try t.expectEqual(@as(i64, 4608), offs[2 * r + 1]); + } + var one: [planBytes(4)]u8 = @splat(0xee); + plan(&one, 4, 510, 3, 0, 4608); + var rows: [planBytes(4)]u8 = @splat(0xee); + for (0..3) |r| planRow(&rows, 4, r, 510 + r + 1, 0, 4608); + try t.expectEqualSlices(u8, &one, &rows); // a stream's consecutive rows: M1's plan + try t.expectEqual(@as(usize, 8), groupsOf(&pos)); // 4,096 keys: 8 groups of 512 + try t.expectEqual(@as(usize, 9), groupsOf(&.{4096})); +} + +test "staged plan: fixed fields once, a round's committed keys its only upload, B1's per-row plan for its rows" { + const t = std.testing; + var staged: [planBytes(4)]u8 = @splat(0xee); + stagePlan(&staged, 4, 4608); + const first = [_]usize{ 4095, 511, 3071 }; + _ = commitRows(&staged, 4, &first); + const pos = [_]usize{ 7, 3072 }; // a later round of fewer rows: rows past it keep the last round's keys + const up = commitRows(&staged, 4, &pos); + try t.expectEqual(@as(usize, 2 * 16), up.len); + try t.expectEqual(@intFromPtr(&staged) + 16, @intFromPtr(up.ptr)); // streams' rows, after row_stream [4] + var b1: [planBytes(4)]u8 = @splat(0xee); + for (0..4) |r| planRow(&b1, 4, r, 0, 0, 4608); + for (pos, 0..) |p, r| planRow(&b1, 4, r, p + 1, 0, 4608); + planRow(&b1, 4, 2, 3072, 0, 4608); + try t.expectEqualSlices(u8, &b1, &staged); + const ints = std.mem.bytesAsSlice(i32, staged[0..96]); + try t.expectEqualSlices(i32, &.{ 1, 1, 3073, 0 }, &[4]i32{ ints[8], ints[9], ints[10], ints[11] }); +} + +test "a graph's key groups: the context's, at least any round's; their partials fit init's scratch" { + const t = std.testing; + try t.expectEqual(@as(usize, 16), ta.TreeAttention.groups(8192)); // the forward test's context + try t.expectEqual(@as(usize, 64), ta.TreeAttention.groups(32768)); // the default context + try t.expect(groupsOf(&.{ 8191, 0 }) <= ta.TreeAttention.groups(8192)); + const w = ta.TreeAttention.scratch(0, 48, global_heads, global_kv_heads, ta.TreeAttention.groups(8192), head_dim); + try t.expectEqual(@as(u64, scratchBytes(48, 8192)), w.sl + ta.slice_items * 4 * 10 * 4); // the last slice sum's end +} + +test "global decode rows: 1 to init's rows, each position within the context" { + const t = std.testing; + try admitRows(16, 131072, &.{ 131071, 0, 70000 }, 512 * 1024 * 256); + try t.expectError(error.GlobalShape, admitRows(16, 131072, &.{131072}, 0)); + try t.expectError(error.GlobalShape, admitRows(2, 131072, &.{ 1, 2, 3 }, 0)); + try t.expectError(error.GlobalShape, admitRows(2, 131072, &.{}, 0)); +} + +test "global decode scratch: 16 rows' partials over 131,072 keys (256 groups), then the slices" { + try std.testing.expectEqual(@as(usize, (16 * 80 * 256 + 1024 * 4 * 10) * 130 * 4), scratchBytes(16, 131072)); + try std.testing.expectEqual(@as(usize, 191_692_800), scratchBytes(16, 131072)); +} diff --git a/zig/src/families/ricursive/cuda_backend.zig b/zig/src/families/ricursive/cuda_backend.zig new file mode 100644 index 0000000000..43a24abcdd --- /dev/null +++ b/zig/src/families/ricursive/cuda_backend.zig @@ -0,0 +1,466 @@ +//! Ripple behind lanes.Backend on CUDA: a slot a stream, a chunk a step, a round all windows; cuda.sample draws. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const cfg = @import("config.zig"); +const fwd = @import("cuda_model.zig"); +/// The prompt cache's states (cuda_snap.zig), reached through the backend. +pub const snap = @import("cuda_snap.zig"); + +const be = lanes.backend; +const Stream = lanes.Stream; +const LogRow = lanes.LogprobRow; +const row_words: usize = lanes.logprob.words; + +pub const prompt_rows: u32 = fwd.prompt_rows; +/// Rows a round takes over every stream: the model's scratch for rounds is opened at this many. +pub const round_rows: u32 = 128; +/// The widest window a stream verifies (B1: a window's rows are the one-row steps' rows). +pub const max_window: u32 = 16; +/// Positions a slot keeps past a reply for a verify window's rows. +pub const reserve: usize = max_window; +/// Drawn tokens a handle names (a ring; read before the round loop could wrap it). +const ring = 1024; +/// Rows of logprob words at once: a round's, or a prompt pass's head rows. +const word_rows = @max(round_rows, fwd.head_rows); + +/// The positions opened for an asked context: it and the verify reserve, within the maximum, at least a round's rows. +pub fn physical(context: usize) usize { + return std.math.clamp(context +| reserve, @as(usize, round_rows), cfg.max_context); +} + +/// A prompt and its reply fit `positions` with the verify reserve left over. +pub fn fits(positions: usize, prompt: usize, reply: usize) bool { + return prompt <= positions and reply <= positions - prompt and reserve <= positions - prompt - reply; +} + +const Lane = struct { + slot: usize, + len: u64 = 0, // the cache length: positions [0, len) hold this stream's keys + base: u64 = 0, // the cache length the last verify started from + hi: usize = 0, // positions written (rejected rows too): release zeroes [0, hi) + rows: std.ArrayList(u32) = .empty, // the last verify's row tokens + last: u64, // the prompt's last row's bf16 logits [V] (the first draw reads them) + rule: u64, // sample.cu's Rule, when the stream samples + sampled: bool = false, + fill: ?usize = null, // a prompt pass going a chunk at a time: where its next chunk starts + first: ?LogRow = null, // the first token's logprob row, with `logprobs` +}; + +pub const Ripple = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + m: *const fwd.Model, + st: fwd.Seq = .{}, // slot 0's positions when the CLI drives the forward itself (cli/cuda_ricursive.zig) + tok: cuda.DeviceBuffer, // int32 drawn tokens: the handle ring, then a round's rows + meta: cuda.DeviceBuffer, // a draw's first position minus one, a window each + words: cuda.DeviceBuffer, // logprob words, word_rows rows of row_words + picks: cuda.DeviceBuffer, // a prompt pass's next tokens (head_rows) + slot_bufs: cuda.DeviceBuffer, // a slot's last-row logits, then its rule, a slot each + sample: cuda.Module, + draws: cuda.sample.Draws, + lanes: std.AutoHashMapUnmanaged(*const Stream, Lane) = .empty, + free: []bool, // slot i is free + next: u64 = 0, + + /// A slot's share of slot_bufs: its last row's logits, then its rule. + fn slotStride(vocab: usize) usize { + return std.mem.alignForward(usize, vocab * 2, 256) + 256; + } + + /// The device bytes `init` allocates for `slots` slots (the caches are the model's, fwd.Model.openBytes). + pub fn openBytes(vocab: usize, slots: usize) usize { + return (ring + round_rows) * 4 + round_rows * 16 + word_rows * row_words * 4 + fwd.head_rows * 4 + slots * slotStride(vocab); + } + + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, m: *const fwd.Model) !Ripple { + const slots = m.arena.slots; + const free = try gpa.alloc(bool, slots); + errdefer gpa.free(free); + @memset(free, true); + var tok = try cuda.DeviceBuffer.alloc(d, (ring + round_rows) * 4); + errdefer tok.free(); + var meta = try cuda.DeviceBuffer.alloc(d, round_rows * 16); + errdefer meta.free(); + var words = try cuda.DeviceBuffer.alloc(d, word_rows * row_words * 4); + errdefer words.free(); + var picks = try cuda.DeviceBuffer.alloc(d, fwd.head_rows * 4); + errdefer picks.free(); + var slot_bufs = try cuda.DeviceBuffer.alloc(d, slots * slotStride(m.c.vocab)); + errdefer slot_bufs.free(); + var sample = try cuda.Module.load(d, cuda.kernels.sample); + errdefer sample.unload(); + const draws = try cuda.sample.Draws.resolve(sample); + return .{ .gpa = gpa, .d = d, .m = m, .tok = tok, .meta = meta, .words = words, .picks = picks, .slot_bufs = slot_bufs, .sample = sample, .draws = draws, .free = free }; + } + + pub fn deinit(x: *Ripple) void { + var it = x.lanes.valueIterator(); + while (it.next()) |l| l.rows.deinit(x.gpa); + x.lanes.deinit(x.gpa); + x.gpa.free(x.free); + x.sample.unload(); + x.slot_bufs.free(); + x.picks.free(); + x.words.free(); + x.meta.free(); + x.tok.free(); + x.* = undefined; + } + + pub fn backend(x: *Ripple) be.Backend { + return .{ .ptr = x, .vtable = &vtable, .chunk_invariant = true }; // M1: chunk cuts give the same bits + } + + const vtable: be.Backend.VTable = .{ .prefill = prefill, .prefill_step = prefillStep, .cancel_prefill = cancelPrefill, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .release = release, .first_row = firstRow, .prompt_rows = true, .reports_fill = true }; + + /// The round loop's facts: no draft head; host copies verified in windows up to 16 rows, streams sharing a round. + pub fn facts(x: *const Ripple) lanes.Model { + const rows: u32 = @intCast(x.m.round_rows); + return .{ .copy_match = 8, .exact_width = @min(max_window, rows), .gpu_tokens = false, .mtp = false, .speculate = false, .hidden_rows = true, .batch_rows = rows, .max_streams = @intCast(@min(x.m.arena.slots, rows)) }; + } + + fn self(ptr: *anyopaque) *Ripple { + return @ptrCast(@alignCast(ptr)); + } + + fn upload(x: *Ripple, dst: u64, bytes: []const u8) !void { + try cuda.copy_source.enqueue(x.d, dst, bytes, x.m.s); + } + + fn tokAt(x: *const Ripple, i: usize) u64 { + return x.tok.ptr + i * 4; + } + + /// `rows` bf16 logits rows drawn into `out`, row r at `position0 + r` (argmax when greedy). + fn drawRows(x: *Ripple, l: *const Lane, logits: u64, rows: usize, position0: u64, out: u64, meta_at: usize) !void { + const v = x.m.c.vocab; + if (!l.sampled) return x.m.k.argmaxRows(x.m.s, logits, rows, v, v, out); + const meta = [4]i32{ @intCast(position0 - 1), 0, 0, 0 }; + const mp = x.meta.ptr + meta_at * 16; + try x.upload(mp, std.mem.asBytes(&meta)); + try x.draws.draw(x.m.s, logits, v, l.rule, mp, 0, out, rows, null, null); + } + + /// The stream's caches at `at` for the prompt cache, copied behind the chunk that ends there. + pub fn saveState(x: *Ripple, owner: ?*anyopaque, at: u32) !*snap.Saved { + const s: *const Stream = @ptrCast(@alignCast(owner orelse return error.NoOwner)); + const l = x.lanes.getPtr(s) orelse return error.UnknownStream; + if ((l.fill orelse l.len) != at) return error.NotAtMark; // a state is copied only where the pass stands + return snap.save(x.gpa, x.m, l.slot, at); + } + + /// The stream's slot made to hold `k`'s state (a restore outside its prompt pass). + pub fn restoreState(x: *Ripple, owner: ?*anyopaque, k: *const snap.Saved) !void { + const s: *const Stream = @ptrCast(@alignCast(owner orelse return error.NoOwner)); + const l = x.lanes.getPtr(s) orelse return error.UnknownStream; + try snap.restore(x.m, l.slot, k); + l.len = k.at; + l.hi = @max(l.hi, k.at); + } + + pub fn dropState(x: *Ripple, k: *snap.Saved) void { + snap.drop(x.gpa, x.m.s, k); + } + + fn freeSlot(x: *const Ripple) ?usize { + return std.mem.indexOfScalar(bool, x.free, true); + } + + /// Whether a request could start now: a free slot; one that never fits passes (its prefill refuses it by name). + pub fn admits(x: *const Ripple, prompt: usize, max_new: usize) bool { + if (!fits(x.m.context, prompt, max_new)) return true; + return x.freeSlot() != null; + } + + /// A prompt pass begins: a slot, the stream's rule, a kept state restored; where its first chunk starts. + fn opening(x: *Ripple, s: *Stream) !*Lane { + const p = s.prompt(); + if (p.len == 0) return error.EmptyPrompt; + if (!fits(x.m.context, p.len, @as(usize, s.max_new))) return error.PastContext; + if (s.grammar != null) return error.StructuresUnsupported; + const got = try x.lanes.getOrPut(x.gpa, s); + if (!got.found_existing) { + const slot = x.freeSlot() orelse { + x.lanes.removeByPtr(got.key_ptr); + return error.NoFreeSlot; + }; + x.free[slot] = false; + const at = x.slot_bufs.ptr + slot * slotStride(x.m.c.vocab); + got.value_ptr.* = .{ .slot = slot, .last = at, .rule = at + slotStride(x.m.c.vocab) - 256 }; + } + const l = got.value_ptr; + l.sampled = false; + l.first = null; + if (s.sampling) |sp| if (sp.temperature > 0) { + const r = cuda.sample.Rule.from(lanes.gpu_rule.rule(sp)); + try x.upload(l.rule, std.mem.asBytes(&r)); + l.sampled = true; + }; + s.cached = 0; + s.prompt_rows.clearRetainingCapacity(); + // a pass with prompt logprobs starts at 0: every prompt row is scored + l.fill = if (s.prompt_logprobs != null) 0 else x.resumeAt(s, l, p.len); + s.fill_at.store(@intCast(l.fill.?), .monotonic); + return l; + } + + /// The kept state's position, restored into the slot (0: none, not this prompt's, or another shape's). + fn resumeAt(x: *Ripple, s: *Stream, l: *Lane, len: usize) usize { + const saved = s.reuse.saved orelse return 0; + const k: *const snap.Saved = @ptrCast(@alignCast(saved)); + if (k.at != s.reuse.at or k.at == 0 or k.at >= len or !snap.sameKey(k.key, snap.keyOf(x.m.c))) { + s.reuse_failed = true; + return 0; + } + snap.restore(x.m, l.slot, k) catch { + s.reuse_failed = true; + return 0; + }; + l.hi = @max(l.hi, k.at); + s.cached = k.at; + return k.at; + } + + /// The chunk's rows [at, end) that predict a prompt token, `head_rows` at a time: words and rank to prompt rows. + fn promptRows(x: *Ripple, s: *Stream, at: usize, end: usize, k: u8) !void { + const p = s.prompt(); + const v = x.m.c.vocab; + const scored = @min(end, p.len - 1) -| at; + var r0: usize = 0; + while (r0 < scored) : (r0 += fwd.head_rows) { + const n = @min(fwd.head_rows, scored - r0); + const next = p[at + r0 + 1 ..][0..n]; + try x.m.headRows(r0, n); + try x.upload(x.picks.ptr, std.mem.sliceAsBytes(next)); + try x.draws.logprobRows(x.m.s, x.m.b.rows, v, x.picks.ptr, k, true, x.words.ptr, n); + var words: [fwd.head_rows * row_words]u32 = undefined; + try x.m.download(x.words.ptr, std.mem.sliceAsBytes(words[0 .. n * row_words])); + for (next, 0..) |t, i| try s.prompt_rows.append(x.gpa, LogRow.fromWords(words[i * row_words ..][0..row_words], t, k)); + } + } + + /// One chunk of `l.fill`, at most `max_rows` (0: full), cut at the next mark; true once the last prompt row is in. + fn chunk(x: *Ripple, s: *Stream, l: *Lane, max_rows: u32) !bool { + const p = s.prompt(); + const at = l.fill orelse return true; + if (s.isCancelled()) { + s.keepInterruptedPrefill(@intCast(at)); + return error.Cancelled; + } + const rows: usize = if (max_rows == 0) prompt_rows else @min(max_rows, prompt_rows); + var end = @min(p.len, at + rows); + for (s.reuse.marks) |mk| if (mk > at and mk < end) { + end = mk; + }; + const done = end == p.len; + try x.m.prompt(l.slot, at, p[at..end]); + if (s.prompt_logprobs) |k| try x.promptRows(s, at, end, k); + if (done) { + try x.m.head(end - at - 1); + try x.d.check(x.d.api.cuMemcpyDtoDAsync_v2(l.last, x.m.b.logits, x.m.c.vocab * 2, x.m.s.handle), "cuMemcpyDtoDAsync"); + } + s.fill_at.store(@intCast(end), .monotonic); + l.hi = @max(l.hi, end); + l.fill = if (done) null else end; + if (done) l.len = p.len; + if (s.reuse.hook) |hk| if (std.mem.indexOfScalar(u32, s.reuse.marks, @intCast(end)) != null) hk.at(hk.ptr, s, @intCast(end)); + return done; + } + + fn prefill(ptr: *anyopaque, s: *Stream) anyerror!void { + const x = self(ptr); + const l = try x.opening(s); + errdefer x.dropLane(s); + while (!try x.chunk(s, l, 0)) {} + } + + fn prefillStep(ptr: *anyopaque, s: *Stream, max_rows: u32) anyerror!bool { + const x = self(ptr); + const going = if (x.lanes.getPtr(s)) |l| l.fill != null else false; + const l = if (going) x.lanes.getPtr(s).? else try x.opening(s); + errdefer x.dropLane(s); + return x.chunk(s, l, max_rows); + } + + fn cancelPrefill(ptr: *anyopaque, s: *Stream) !void { + const l = self(ptr).lanes.getPtr(s) orelse return; + if (l.fill) |at| s.keepInterruptedPrefill(@intCast(at)); + } + + fn laneOf(x: *Ripple, s: *const Stream) !*Lane { + return x.lanes.getPtr(s) orelse error.UnknownStream; + } + + fn handle(x: *Ripple) u64 { + const h = x.next; + x.next += 1; + return h; + } + + /// The prompt's last row drawn at `position`; with logprobs, its row kept for first_row. + fn first(ptr: *anyopaque, s: *Stream, position: u64) anyerror!u64 { + const x = self(ptr); + const l = try x.laneOf(s); + if (position != l.len or l.fill != null) return error.PositionMismatch; + const h = x.handle(); + try x.drawRows(l, l.last, 1, position, x.tokAt(h % ring), 0); + if (s.logprobs) |k| { + try x.draws.logprobRows(x.m.s, l.last, x.m.c.vocab, x.tokAt(h % ring), k, false, x.words.ptr, 1); + var w: [row_words]u32 = undefined; + try x.m.download(x.words.ptr, std.mem.sliceAsBytes(&w)); + l.first = LogRow.fromWords(&w, try read(x, h), k); + } + return h; + } + + fn firstRow(ptr: *anyopaque, s: *Stream) anyerror!LogRow { + const l = try self(ptr).laneOf(s); + return l.first orelse error.NoFirstRow; + } + + fn value(x: *Ripple, feed: be.Feed) !u32 { + return switch (feed) { + .value => |v| v, + .handle => |h| read(x, h), + }; + } + + /// One token at the cache's end: a round of one row, its row through the head, drawn at `position`. + fn queue(ptr: *anyopaque, s: *Stream, feed: be.Feed, position: u64) anyerror!u64 { + const x = self(ptr); + const l = try x.laneOf(s); + const t = try x.value(feed); + if (l.fill != null or position != l.len + 1) return error.PositionMismatch; + try x.m.round(&.{.{ .slot = l.slot, .p0 = l.len, .tokens = 1 }}, &.{t}); + try x.m.rowLogits(1); + l.len += 1; + l.hi = @max(l.hi, l.len); + const h = x.handle(); + try x.drawRows(l, x.m.b.rows, 1, position, x.tokAt(h % ring), 0); + return h; + } + + fn read(ptr: *anyopaque, h: u64) anyerror!u32 { + const x = self(ptr); + if (h >= x.next or x.next - h > ring) return error.NoSuchToken; + var out: [1]u32 = undefined; + try x.m.download(x.tokAt(h % ring), std.mem.sliceAsBytes(&out)); + return out[0]; + } + + /// Every window's rows in one round (a chain a stream, from its cache length), each window drawn at its positions. + fn verify(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { + const x = self(ptr); + const v = x.m.c.vocab; + const chains = try x.gpa.alloc(fwd.Chain, windows.len); + defer x.gpa.free(chains); + var ids: [round_rows]u32 = undefined; + var total: usize = 0; + for (windows, chains) |w, *c| { + if (w.held != 0) return error.NoHeldDrafts; // no draft head: every draft comes from the host + if (w.parents != null) return error.ChainsOnly; + if (w.masks.len != 0) return error.StructuresUnsupported; + const l = try x.laneOf(w.stream); + if (w.rows() > max_window) return error.WindowTooWide; + if (total + w.rows() > @min(round_rows, x.m.round_rows)) return error.RoundTooWide; + if (l.fill != null) return error.PositionMismatch; + l.base = l.len; + l.rows.clearRetainingCapacity(); + try l.rows.append(x.gpa, w.pending); + try l.rows.appendSlice(x.gpa, w.tokens); + for (w.positions, 0..) |pos, r| if (pos != l.base + r + 1) return error.PositionMismatch; + c.* = .{ .slot = l.slot, .p0 = l.base, .tokens = l.rows.items.len }; + @memcpy(ids[total..][0..l.rows.items.len], l.rows.items); + total += l.rows.items.len; + } + try x.m.round(chains, ids[0..total]); + try x.m.rowLogits(total); + var at: usize = 0; + var scored = false; + for (windows, out, 0..) |w, o, i| { + const l = x.lanes.getPtr(w.stream).?; + const n = l.rows.items.len; + const logits = x.m.b.rows + at * v * 2; + try x.drawRows(l, logits, n, w.positions[0], x.tokAt(ring + at), i); + if (o.rows.len > 0) { + try x.draws.logprobRows(x.m.s, logits, v, x.tokAt(ring + at), w.stream.logprobs.?, false, x.words.ptr + at * row_words * 4, n); + scored = true; + } + at += n; + } + var got: [round_rows]u32 = undefined; + try x.m.download(x.tokAt(ring), std.mem.sliceAsBytes(got[0..total])); + var words: [round_rows * row_words]u32 = undefined; + if (scored) try x.m.download(x.words.ptr, std.mem.sliceAsBytes(words[0 .. total * row_words])); + at = 0; + for (windows, out) |w, o| { + const l = x.lanes.getPtr(w.stream).?; + const n = l.rows.items.len; + @memcpy(o.sampled[0..n], got[at..][0..n]); + @memcpy(o.drafts[0 .. n - 1], l.rows.items[1..]); + for (o.rows, o.sampled[0..o.rows.len], at..) |*r, pick, row| r.* = LogRow.fromWords(words[row * row_words ..][0..row_words], pick, w.stream.logprobs.?); + l.len = l.base + n; // a chain keeps every row unless rolled back + l.hi = @max(l.hi, l.len); + at += n; + } + } + + fn keep(ptr: *anyopaque, windows: []const be.Window, paths: []const []const u32) anyerror!void { + const x = self(ptr); + for (windows, paths) |w, path| { + const l = try x.laneOf(w.stream); + if (path.len == 0 or path.len > l.rows.items.len) return error.EmptyPath; + for (path, 0..) |r, i| if (r != i) return error.ChainsOnly; // a chain keeps a prefix of its rows + l.len = l.base + path.len; // caches are positional: the next round overwrites what the rejected rows wrote + } + } + + /// No draft head: absorbing is a no-op, and a request for drafts is refused by name. + fn draft(_: *anyopaque, requests: []const be.DraftRequest) anyerror!void { + for (requests) |r| if (r.depth != 0 or r.lanes != null) return error.NoDraftHead; + } + + /// The stream's slot zeroed where it wrote (a fresh arena's bytes for the next stream) and freed. + fn dropLane(x: *Ripple, s: *const Stream) void { + if (x.lanes.fetchRemove(s)) |entry| { + var l = entry.value; + x.m.reset(l.slot, l.hi) catch |e| std.log.warn("ripple: slot {d} not zeroed on release ({s})", .{ l.slot, @errorName(e) }); + x.free[l.slot] = true; + l.rows.deinit(x.gpa); + } + } + + fn release(ptr: *anyopaque, s: *Stream) void { + self(ptr).dropLane(s); + } +}; + +test "ricursive cuda backend: the verify reserve past a reply; positions opened for a context" { + try std.testing.expect(fits(32784, 30000, 2768)); + try std.testing.expect(fits(32784, 32784 - reserve - 1, 1)); + try std.testing.expect(!fits(32784, 32784 - reserve, 1)); + try std.testing.expect(!fits(32784, 32785, 0)); + try std.testing.expectEqual(@as(usize, 32768 + reserve), physical(32768)); + try std.testing.expect(fits(physical(32768), 32000, 768)); // the asked context holds a prompt and its reply + try std.testing.expectEqual(@as(usize, 262144), physical(262144)); // the model's maximum caps it + try std.testing.expectEqual(@as(usize, round_rows), physical(1)); +} + +test "ricursive cuda backend: windows of 16 rows, a stream a slot, rows a round" { + var m: fwd.Model = undefined; + m.round_rows = round_rows; + m.arena.slots = 6; + const x: Ripple = .{ .gpa = std.testing.allocator, .d = undefined, .m = &m, .tok = undefined, .meta = undefined, .words = undefined, .picks = undefined, .slot_bufs = undefined, .sample = undefined, .draws = undefined, .free = &.{} }; + const f = x.facts(); + try std.testing.expectEqual(@as(u32, 16), f.exact_width); + try std.testing.expectEqual(@as(u32, 6), f.max_streams); + try std.testing.expectEqual(@as(u32, 128), f.batch_rows); + try std.testing.expect(f.hidden_rows and f.copy_drafts and !f.mtp and !f.gpu_tokens); + try std.testing.expect(Ripple.slotStride(100352) >= 100352 * 2 + @sizeOf(cuda.sample.Rule)); + std.testing.refAllDecls(Ripple); +} + +test { + _ = snap; +} diff --git a/zig/src/families/ricursive/cuda_kernels.zig b/zig/src/families/ricursive/cuda_kernels.zig new file mode 100644 index 0000000000..7f2bb50f27 --- /dev/null +++ b/zig/src/families/ricursive/cuda_kernels.zig @@ -0,0 +1,171 @@ +//! Ripple's glue kernels on CUDA (ricursive.cu): the Metal family's glue.metal launches, one by one. + +const std = @import("std"); +const cuda = @import("cuda"); + +const Stream = cuda.Stream; +const Function = cuda.Function; + +fn int(x: usize) c_int { + return @intCast(x); +} + +fn ceil(n: usize, d: usize) u32 { + return @intCast((n + d - 1) / d); +} + +/// Threads a block in every glue launch (ricursive.cuh's NT). +pub const threads = 256; +/// The widest row tf_rc_add_norm holds: each thread four 8-value chunks (MOST). +pub const most_width = 8 * threads * 4; + +pub const Kernels = struct { + mod: cuda.Module, + embed: Function, + add_norm: Function, + gate: Function, + silu_clamp: Function, + // plug-in point: the prep and attention launches, a Function each, tf_rc_ in ricursive.cu + + pub fn load(d: *const cuda.Driver) !Kernels { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var mod = try cuda.Module.load(d, cuda.kernels.ricursive); + errdefer mod.unload(); + var k: Kernels = undefined; + k.mod = mod; + inline for (@typeInfo(Kernels).@"struct".field_names, @typeInfo(Kernels).@"struct".field_types) |name, T| if (T == Function) { + const sym = "tf_rc_" ++ name; + @field(k, name) = try mod.function(sym); + }; + return k; + } + + pub fn unload(k: *Kernels) void { + k.mod.unload(); + k.* = undefined; + } +}; + +fn go(f: Function, s: Stream, grid: cuda.Dim3, a: *cuda.Args) !void { + try cuda.launch.launch(f, .{ .grid = grid, .block = .{ .x = threads }, .shared = 0 }, s, a); +} + +/// An elementwise launch's grid: 256 of a row's values a block, a grid row a row (at most 65,535). +fn plane(rows: usize, width: usize) !cuda.Dim3 { + if (rows == 0 or rows > 65535 or width == 0) return error.UnsupportedShape; + return .{ .x = ceil(width, threads), .y = @intCast(rows) }; +} + +/// R and V0 = bf16(E[ids] * scale), the fp32 scale: the residual's rows and the value residual's source (u32 ids). +pub fn embed(k: *const Kernels, s: Stream, ids: u64, table: u64, resid: u64, v0: u64, rows: usize, width: usize, scale: f32) !void { + return embedBy(k.embed, s, ids, table, resid, v0, rows, width, scale); +} + +/// `embed` through `f`: the served kernel or a control of it. +pub fn embedBy(f: Function, s: Stream, ids: u64, table: u64, resid: u64, v0: u64, rows: usize, width: usize, scale: f32) !void { + const grid = try plane(rows, width); + var a: cuda.Args = .{}; + for ([_]u64{ ids, table, resid, v0 }) |v| a.add(v); + a.add(int(rows)); + a.add(int(width)); + a.add(scale); + try go(f, s, grid, &a); +} + +/// One tf_rc_add_norm launch: `branch` 0 skips the add (and `post`), `v0` 0 skips the lerp (and `vx`). +pub const NormRows = struct { + resid: u64, // R, read and written in place + branch: u64 = 0, // Y: the attention's or the FFN's output rows + post: u64 = 0, // w_post, Y's norm weight + next: u64, // w_next, the next norm's weight + v0: u64 = 0, // the value residual's source rows + x: u64, // X: the next norm's rows + vx: u64 = 0, // lerp(V0, X, blend): v's input + rows: usize, + width: usize, + eps: f32, + blend: f32 = 0, +}; + +/// tf_rc_add_norm's layout: 8-value chunks, at most most_width a row, every pointer 16-byte aligned. +pub fn normLayout(p: NormRows) !void { + if (p.width == 0 or p.width % 8 != 0 or p.width > most_width) return error.UnsupportedWidth; + if (p.rows == 0 or p.rows > std.math.maxInt(i32)) return error.UnsupportedShape; + if ((p.branch == 0) != (p.post == 0) or (p.v0 == 0) != (p.vx == 0)) return error.MissingOperand; + for ([_]u64{ p.resid, p.branch, p.post, p.next, p.v0, p.x, p.vx }) |v| if (v % 16 != 0) return error.Misaligned; +} + +/// R = bf16(R + norm(Y, w_post)) with Y, then X = norm(R, w_next), and VX = lerp(V0, X, blend) with V0: a row a block. +pub fn addNorm(k: *const Kernels, s: Stream, p: NormRows) !void { + return addNormBy(k.add_norm, s, p); +} + +/// `addNorm` through `f`: the served kernel or a control of it. +pub fn addNormBy(f: Function, s: Stream, p: NormRows) !void { + try normLayout(p); + var a: cuda.Args = .{}; + for ([_]u64{ p.resid, p.branch, p.post, p.next, p.v0, p.x, p.vx }) |v| a.add(v); + a.add(int(p.rows)); + a.add(int(p.width)); + a.add(p.eps); + a.add(p.blend); + try go(f, s, .{ .x = @intCast(p.rows) }, &a); +} + +/// The attention's output gate in place: A = bf16(A * bf16(sigmoid(G))) over rows of `width` (query heads x 128). +pub fn gate(k: *const Kernels, s: Stream, act: u64, g: u64, rows: usize, width: usize) !void { + return gateBy(k.gate, s, act, g, rows, width); +} + +/// `gate` through `f`: the served kernel or a control of it. +pub fn gateBy(f: Function, s: Stream, act: u64, g: u64, rows: usize, width: usize) !void { + const grid = try plane(rows, width); + var a: cuda.Args = .{}; + a.add(act); + a.add(g); + a.add(int(rows)); + a.add(int(width)); + try go(f, s, grid, &a); +} + +/// A = bf16(bf16(silu(min(g, limit))) * clamp(u, -limit, limit)) from stacked gate | up rows, 2 x `width` each. +pub fn siluClamp(k: *const Kernels, s: Stream, gu: u64, act: u64, rows: usize, width: usize, limit: f32) !void { + return siluClampBy(k.silu_clamp, s, gu, act, rows, width, limit); +} + +/// `siluClamp` through `f`: the served kernel or a control of it. +pub fn siluClampBy(f: Function, s: Stream, gu: u64, act: u64, rows: usize, width: usize, limit: f32) !void { + const grid = try plane(rows, width); + var a: cuda.Args = .{}; + a.add(gu); + a.add(act); + a.add(int(rows)); + a.add(int(width)); + a.add(limit); + try go(f, s, grid, &a); +} + +test "launch geometry: ceilings and the add-norm's layout" { + try std.testing.expectEqual(@as(u32, 20), ceil(5120, threads)); + try std.testing.expectEqual(@as(u32, 40), ceil(10240, threads)); + try std.testing.expectEqual(@as(u32, 160), ceil(40960, threads)); + try std.testing.expectEqual(@as(usize, 8192), most_width); + try std.testing.expectError(error.UnsupportedShape, plane(65536, 5120)); + try std.testing.expectError(error.UnsupportedShape, plane(1, 0)); + const p: NormRows = .{ .resid = 256, .next = 512, .x = 768, .rows = 2048, .width = 5120, .eps = 1e-6 }; + try normLayout(p); + for ([_]usize{ 5124, 8200, 0 }) |w| { + var bad = p; + bad.width = w; + try std.testing.expectError(error.UnsupportedWidth, normLayout(bad)); + } + var bad = p; + bad.branch = 1024; + try std.testing.expectError(error.MissingOperand, normLayout(bad)); + bad = p; + bad.v0 = 1024; + try std.testing.expectError(error.MissingOperand, normLayout(bad)); + bad = p; + bad.x = 776; + try std.testing.expectError(error.Misaligned, normLayout(bad)); +} diff --git a/zig/src/families/ricursive/cuda_model.zig b/zig/src/families/ricursive/cuda_model.zig new file mode 100644 index 0000000000..5e5367164f --- /dev/null +++ b/zig/src/families/ricursive/cuda_model.zig @@ -0,0 +1,589 @@ +//! Ripple's forward on CUDA: Metal forward.zig's layer order; prompt chunks of one slot, and rounds of rows over slots. +const std = @import("std"); +const cuda = @import("cuda"); +const cfg = @import("config.zig"); +const glue = @import("cuda_kernels.zig"); +const rp = @import("cuda_prep.zig"); +const attn = @import("cuda_attn.zig"); +const wts = @import("cuda_weights.zig"); + +const wa = cuda.window_attention; +const Stream = cuda.Stream; +const DeviceBuffer = cuda.DeviceBuffer; +const Config = cfg.Config; +const Weight = cuda.qmmf.Weight; + +/// A prompt forward's rows at most (a chunk). +pub const prompt_rows = 2048; +/// A sliding layer's ring: the window's keys before a chunk and the chunk's own (window - 1 + 2,048, to 128s). +pub const ring_rows = 3072; +/// Positions a global layer's page holds. +pub const page_tokens = 512; +/// Rows the head writes at once for a chunk's every-row logits. +pub const head_rows = 128; +/// A one-row round's rows: the round maximum of the CLI and the one-stream backend (M1's one-row chain decode). +pub const decode_rows = 1; +/// Device offsets start on 256 bytes (the row kernels read 16-byte chunks). +const buffer_align = 256; + +/// prompt: a chunk by the prompt kernels; decode: a round of one chain by the lane matmul and the decode attention. +pub const Mode = enum { prompt, decode }; + +/// One stream's rows in a round: `tokens` consecutive positions from p0 in its slot. +pub const Chain = struct { slot: usize, p0: usize, tokens: usize }; + +pub const Kernels = @import("cuda_model_kernels.zig").Kernels; + +/// The forward's scratch for `prompt_rows` rows (a round's at most), each at its own 256-byte boundary. +const Bufs = struct { + resid: u64 = 0, // bf16 [R, H]: the residual + v0: u64 = 0, // bf16 [R, H]: the value residual's source (the scaled embedding) + x: u64 = 0, // bf16 [R, H]: the next norm's rows + vx: u64 = 0, // bf16 [R, H]: lerp(v0, x, blend), v's input + y: u64 = 0, // bf16 [R, H]: o's and down's output + qkg: u64 = 0, // bf16 [R, q + kv + q]: the fused projection + v: u64 = 0, // bf16 [R, kv] + q: u64 = 0, // bf16 [R, heads, 128]: prep's normed, rotated queries + att: u64 = 0, // bf16 [R, q]: the attention, then gated in place + g: u64 = 0, // bf16 [R, q]: the gate's rows contiguous (several rows) + gu: u64 = 0, // bf16 [R, 2I]: gate then up a row, as tf_rc_silu_clamp reads them + act: u64 = 0, // bf16 [R, I]: the planar gate, then silu_clamp's output + up: u64 = 0, // bf16 [R, I]: the planar up + ids: u64 = 0, // u32 [R] + pos: u64 = 0, // int32 [round]: a round's rows' positions + slot: u64 = 0, // int32 [round]: a round's rows' slots + inv: u64 = 0, // f32 [64]: RoPE's inverse frequencies + wscr: u64 = 0, // the sliding decode's partials (16 bytes at head dim 128) + part: u64 = 0, // the lane matmul's slice partials at a round's rows + logits: u64 = 0, // bf16 [V]: the last row's + rows: u64 = 0, // bf16 [max(head_rows, round), V]: a chunk's or a round's rows + ctab: u64 = 0, // int32 [round, pages a slot]: a round's call table, row r row r's slot's page ids +}; + +/// The scratch for rounds of at most `round` rows over slots of `per_slot` pages. +fn bufSizes(c: Config, part: usize, round: usize, per_slot: usize) Bufs { + const r = prompt_rows; + const h = c.hidden; + const q = c.queryWidth(); + const f = c.intermediate; + return .{ + .resid = r * h * 2, + .v0 = r * h * 2, + .x = r * h * 2, + .vx = r * h * 2, + .y = r * h * 2, + .qkg = r * wts.qkgWidth(c) * 2, + .v = r * c.kvWidth() * 2, + .q = r * q * 2, + .att = r * q * 2, + .g = r * q * 2, + .gu = r * 2 * f * 2, + .act = r * f * 2, + .up = r * f * 2, + .ids = r * 4, + .pos = round * 4, + .slot = round * 4, + .inv = cfg.head_dim / 2 * 4, + .wscr = wa.WindowAttention.scratchBytes(slidingWindow(c), round), + .part = part, + .logits = c.vocab * 2, + .rows = @max(head_rows, round) * c.vocab * 2, + .ctab = round * per_slot * 4, + }; +} + +/// Each buffer's offset from `base` in field order; the bytes they span. +fn place(sizes: Bufs, base: u64, out: *Bufs) usize { + var at: usize = 0; + inline for (@typeInfo(Bufs).@"struct".field_names) |name| { + @field(out, name) = base + at; + at = std.mem.alignForward(usize, at + @field(sizes, name), buffer_align); + } + return at; +} + +fn slidingWindow(c: Config) wa.Window { + return attn.sliding(c.heads, c.kv_heads, c.window, ring_rows, c.scale()); +} + +/// The pages a slot's global layer holds: the context in whole pages. +pub fn pagesOf(context: usize) usize { + return (context + page_tokens - 1) / page_tokens; +} + +/// A layer's caches from slot 0: a sliding layer's rings [slots][ring][kv_heads][128], or (`pg`) a global's pages. +pub const Cache = struct { kc: u64, vc: u64, pg: ?wa.Pages = null }; + +/// One slot's rows of a layer as a snapshot copies them: position p at ring row p % ring_rows, or global row p. +pub const Span = struct { k: u64, v: u64, rows: usize, row_bytes: usize, ring: bool }; + +/// One slot's caches: a sliding layer's K and V rings, a global layer's K and V pages for `context`. +pub fn slotBytes(c: Config, context: usize) usize { + const kvw = c.kvWidth(); + const g = c.globals(); + return (c.layers - g) * 2 * ring_rows * kvw * 2 + g * 2 * pagesOf(context) * page_tokens * kvw * 2; +} + +/// Every slot's caches in one buffer, a layer at a time (each slot's rows, K then V), then the page table. +pub const Arena = struct { + gpa: std.mem.Allocator, + kv: DeviceBuffer, + caches: []Cache, // a layer each, at slot 0 + ids: []i32, // the page table's host copy [slots][per_slot]: the call tables gather its rows + slots: usize, + per_slot: usize, // pages a slot's global layer holds + kvw: usize, + + /// `slots` slots' caches for `context` positions and the int32 table (one slot: M1's State.bytes). + pub fn bytes(c: Config, context: usize, slots: usize) usize { + return slots * slotBytes(c, context) + std.mem.alignForward(usize, slots * pagesOf(context) * 4, buffer_align); + } + + /// Zeroed caches (no stale bits under a masked key), the identity table: slot s's page j is page s * per_slot + j. + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, s: Stream, c: Config, context: usize, slots: usize) !Arena { + const kvw = c.kvWidth(); + const per = pagesOf(context); + var kv = try DeviceBuffer.alloc(d, bytes(c, context, slots)); + errdefer kv.free(); + try d.check(d.api.cuMemsetD8Async(kv.ptr, 0, kv.len, s.handle), "cuMemsetD8Async"); + const caches = try gpa.alloc(Cache, c.layers); + errdefer gpa.free(caches); + const ids = try gpa.alloc(i32, slots * per); + errdefer gpa.free(ids); + for (ids, 0..) |*x, i| x.* = @intCast(i); + const table_at = kv.len - std.mem.alignForward(usize, ids.len * 4, buffer_align); + try cuda.copy_source.enqueue(d, kv.ptr + table_at, std.mem.sliceAsBytes(ids), s); + var at: u64 = kv.ptr; + for (caches, 0..) |*cache, i| { + const half: usize = slots * if (c.global[i]) per * page_tokens * kvw * 2 else ring_rows * kvw * 2; + cache.* = .{ .kc = at, .vc = at + half }; + if (c.global[i]) cache.pg = .{ .table = kv.ptr + table_at, .page_elems = page_tokens * kvw, .tokens = page_tokens, .stride = per }; + at += 2 * half; + } + return .{ .gpa = gpa, .kv = kv, .caches = caches, .ids = ids, .slots = slots, .per_slot = per, .kvw = kvw }; + } + + pub fn deinit(a: *Arena) void { + a.gpa.free(a.ids); + a.gpa.free(a.caches); + a.kv.free(); + a.* = undefined; + } + + /// Layer i's caches of one slot as a prompt chunk reads them: its rings, or the pools and its table row (stride 0). + pub fn slotCache(a: *const Arena, i: usize, slot: usize) Cache { + const c = a.caches[i]; + if (c.pg) |pg| { + var row = pg; + row.table += slot * a.per_slot * 4; + row.stride = 0; + return .{ .kc = c.kc, .vc = c.vc, .pg = row }; + } + const ring = ring_rows * a.kvw * 2; + return .{ .kc = c.kc + slot * ring, .vc = c.vc + slot * ring }; + } + + /// Layer i's rows of one slot. + pub fn span(a: *const Arena, i: usize, slot: usize) Span { + const c = a.caches[i]; + const row = a.kvw * 2; + const rows: usize = if (c.pg != null) a.per_slot * page_tokens else ring_rows; + return .{ .k = c.kc + slot * rows * row, .v = c.vc + slot * rows * row, .rows = rows, .row_bytes = row, .ring = c.pg == null }; + } + + /// The V pools' offset from the K pools in bf16 elements: the same in every global layer. + fn vOff(a: *const Arena) usize { + for (a.caches) |c| if (c.pg != null) return (c.vc - c.kc) / 2; + return 0; + } +}; + +/// A round's rows on the host: positions and slots (int32 for the kernels, usize for the plan), the call table. +const Host = struct { pos: []i32, slot: []i32, at: []usize, ctab: []i32 }; + +/// Checks a round's chains against `ids` and fills the host rows: chain j's rows in order, each slot at most once. +fn fill(h: Host, table: []const i32, per_slot: usize, slots: usize, context: usize, chains: []const Chain, n: usize) !void { + if (n == 0 or n > h.pos.len) return error.ForwardRows; + var r: usize = 0; + for (chains, 0..) |ch, j| { + if (ch.slot >= slots) return error.NoSuchSlot; + if (ch.tokens == 0 or ch.tokens > n - r) return error.ForwardRows; + if (ch.p0 > context or ch.tokens > context - ch.p0) return error.PastContext; + for (chains[0..j]) |o| if (o.slot == ch.slot) return error.SlotTwice; // its rows would race in the caches + for (r..r + ch.tokens, ch.p0..) |row, p| { + h.pos[row] = @intCast(p); + h.slot[row] = @intCast(ch.slot); + h.at[row] = p; + @memcpy(h.ctab[row * per_slot ..][0..per_slot], table[ch.slot * per_slot ..][0..per_slot]); + } + r += ch.tokens; + } + if (r != n) return error.ForwardRows; +} + +/// One slot's stream as the CLI and the one-stream backend drive it: its slot and the positions its caches hold. +pub const Seq = struct { slot: usize = 0, len: usize = 0 }; + +/// A forward's rows: a prompt chunk of one chain, or a round (its global layers' plan, null without global layers). +const Pass = union(enum) { chunk: Chain, round: ?attn.Planned }; + +pub const Model = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + s: Stream, + w: *const wts.Weights, + k: *const Kernels, + c: Config, + context: usize, + /// A round's rows at most (a caller's choice, for example 16 rows for each of its streams; at most prompt_rows). + round_rows: usize, + scratch: DeviceBuffer, + b: Bufs = .{}, + gd: attn.GlobalDecode, + arena: Arena, + host: Host, + /// Rounds replay CUDA graphs (round's `.graph`); false: eager launches with the round's own key groups (B1's). + graphs: bool = true, + /// Captured rounds by row count ([n], n to round_rows), each the layers' launches over the staged rows. + execs: []?cuda.graph.Exec, + + /// The lane's slice partials of every projection and the head at up to `rows` rows (they reduce below fused_rows). + fn partBytes(w: *const wts.Weights, k: *const Kernels, rows: usize) usize { + var most: usize = 0; + for ([_]usize{ rows, @min(rows, cuda.qmmf.fused_rows - 1) }) |m| { + most = @max(most, k.lane.partBytes(m, w.head)); + for (w.layers) |l| for ([_]Weight{ l.qkg, l.v, l.o, l.gate, l.up, l.down }) |x| { + most = @max(most, k.lane.partBytes(m, x)); + }; + } + return most; + } + + /// The device bytes init allocates beside the weights: scratch, `slots` slots' caches, global decode for `round`. + pub fn openBytes(w: *const wts.Weights, k: *const Kernels, context: usize, slots: usize, round_rows: usize) usize { + const c = w.config; + var b: Bufs = .{}; + const scratch = place(bufSizes(c, partBytes(w, k, round_rows), round_rows, pagesOf(context)), 0, &b); + return scratch + Arena.bytes(c, context, slots) + attn.GlobalDecode.bytes(round_rows, context); + } + + /// Scratch for `round_rows`-row rounds and `slots` slots of `context` positions; refuses a shape no ring holds. + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, s: Stream, w: *const wts.Weights, k: *const Kernels, context: usize, slots: usize, round_rows: usize) !Model { + const c = w.config; + if (c.window == 0 or c.window - 1 + prompt_rows > ring_rows) return error.UnsupportedRicursiveCudaShape; + if (context == 0 or context > cfg.max_context) return error.PastContext; + // a round's rows fit the chunk's scratch (and a chain's rows never lap the ring); page ids are int32 + if (round_rows == 0 or round_rows > prompt_rows or round_rows > context or slots == 0 or slots * pagesOf(context) > std.math.maxInt(i32)) return error.ForwardRows; + const sizes = bufSizes(c, partBytes(w, k, round_rows), round_rows, pagesOf(context)); + var b: Bufs = .{}; + var scratch = try DeviceBuffer.alloc(d, place(sizes, 0, &b)); + errdefer scratch.free(); + _ = place(sizes, scratch.ptr, &b); + var inv: [cfg.head_dim / 2]f32 = undefined; + c.inverseFrequencies(&inv); + try cuda.copy_source.enqueue(d, b.inv, std.mem.sliceAsBytes(&inv), s); + var gd = try attn.GlobalDecode.init(d, gpa, round_rows, context); + errdefer gd.deinit(); + var arena = try Arena.init(gpa, d, s, c, context, slots); + errdefer arena.deinit(); + try gd.stage(s, arena.vOff()); + const execs = try gpa.alloc(?cuda.graph.Exec, round_rows + 1); + errdefer gpa.free(execs); + @memset(execs, null); + const h: Host = .{ + .pos = try gpa.alloc(i32, round_rows), + .slot = try gpa.alloc(i32, round_rows), + .at = try gpa.alloc(usize, round_rows), + .ctab = try gpa.alloc(i32, round_rows * arena.per_slot), + }; + return .{ .gpa = gpa, .d = d, .s = s, .w = w, .k = k, .c = c, .context = context, .round_rows = round_rows, .scratch = scratch, .b = b, .gd = gd, .arena = arena, .host = h, .execs = execs }; + } + + pub fn deinit(m: *Model) void { + m.dropGraphs(); + m.gpa.free(m.execs); + m.gpa.free(m.host.ctab); + m.gpa.free(m.host.at); + m.gpa.free(m.host.slot); + m.gpa.free(m.host.pos); + m.arena.deinit(); + m.gd.deinit(); + m.scratch.free(); + m.* = undefined; + } + + /// Zeroes a slot's positions [0, written) (a ring past its rows whole): a fresh arena's bytes, for another stream. + pub fn reset(m: *const Model, slot: usize, written: usize) !void { + if (slot >= m.arena.slots) return error.NoSuchSlot; + for (0..m.c.layers) |i| { + const sp = m.arena.span(i, slot); + const n = @min(written, sp.rows) * sp.row_bytes; + for ([_]u64{ sp.k, sp.v }) |x| try m.d.check(m.d.api.cuMemsetD8Async(x, 0, n, m.s.handle), "cuMemsetD8Async"); + } + } + + fn linear(m: *const Model, x: u64, ldx: usize, n: usize, w: Weight, y: u64, mode: Mode) !void { + if (mode == .decode) return m.k.lane.matmulOut(m.s, x, ldx, n, w, y, .bf16, m.b.part); + return m.k.prompt.matmulOut(m.s, x, ldx, n, w, y, .bf16); + } + + fn copy2D(m: *const Model, src: u64, src_pitch: usize, dst: u64, dst_pitch: usize, width: usize, height: usize) !void { + const p: cuda.abi.Memcpy2D = .{ .src = src, .src_pitch = src_pitch, .dst = dst, .dst_pitch = dst_pitch, .width = width, .height = height }; + try m.d.check(m.d.api.cuMemcpy2DAsync_v2(&p, m.s.handle), "cuMemcpy2DAsync"); + } + + /// Gate and up into gu's rows (gate at column 0, up at column I): one row straight in, else planar then copied. + fn mlpIn(m: *const Model, l: *const wts.Layer, n: usize, mode: Mode) !void { + const h = m.c.hidden; + const f2 = m.c.intermediate * 2; + const b = m.b; + if (n == 1) { + try m.linear(b.x, h, 1, l.gate, b.gu, mode); + return m.linear(b.x, h, 1, l.up, b.gu + f2, mode); + } + try m.linear(b.x, h, n, l.gate, b.act, mode); + try m.linear(b.x, h, n, l.up, b.up, mode); + try m.copy2D(b.act, f2, b.gu, 2 * f2, f2, n); + try m.copy2D(b.up, f2, b.gu + f2, 2 * f2, f2, n); + } + + /// Layer i's q, k, v prep and attention: a chunk's prep and prompt kernels, or a round's prep_rows and decode. + fn attention(m: *const Model, l: *const wts.Layer, i: usize, n: usize, pass: Pass) !void { + const c = m.c; + const b = m.b; + const k = m.k; + const qw = c.queryWidth(); + const kvw = c.kvWidth(); + const wq = wts.qkgWidth(c); + const scale = c.scale(); + const xq = rp.Rows.of(b.qkg, wq); + const xk = rp.Rows.of(b.qkg + qw * 2, wq); + const xv = rp.Rows.of(b.v, kvw); + switch (pass) { + .chunk => |ch| { + const cache = m.arena.slotCache(i, ch.slot); + try rp.prep(&k.prep, m.s, .{ .xq = xq, .xk = xk, .xv = xv, .inv = b.inv, .qw = l.q_norm, .kw = l.k_norm, .q = b.q, .kc = cache.kc, .vc = cache.vc, .rows = n, .heads = c.heads, .kv_heads = c.kv_heads, .pos = ch.p0, .ring = if (cache.pg == null) ring_rows else 0, .rope = !c.global[i], .eps = c.eps, .pg = cache.pg }); + const w = if (cache.pg) |pg| attn.global(c.heads, c.kv_heads, scale, pg) else slidingWindow(c); + return attn.prompt(&k.attn, m.s, b.q, cache.kc, cache.vc, b.att, ch.p0, n, w); + }, + .round => |planned| { + const cache = m.arena.caches[i]; + try rp.prepRows(&k.prep, m.s, .{ .xq = xq, .xk = xk, .xv = xv, .inv = b.inv, .qw = l.q_norm, .kw = l.k_norm, .q = b.q, .kc = cache.kc, .vc = cache.vc, .rows = n, .heads = c.heads, .kv_heads = c.kv_heads, .pos = b.pos, .slot = b.slot, .ring = if (cache.pg == null) ring_rows else 0, .rope = !c.global[i], .eps = c.eps, .pg = cache.pg }); + if (cache.pg) |pg| { + var call = pg; // tree_attention reads its streams' rows: the round's gathered table + call.table = b.ctab; + return attn.runRows(&m.gd, m.s, planned.?, b.q, cache.kc, b.att, c.heads, call, scale); + } + return k.win.decode(m.s, b.q, cache.kc, cache.vc, b.att, b.pos, b.slot, slidingWindow(c), b.wscr, n); + }, + } + } + + /// The forward's n rows (ids uploaded): every layer, then the final norm into x (the heads read it). + fn layers(m: *const Model, n: usize, pass: Pass) !void { + const c = m.c; + const s = m.s; + const b = m.b; + const k = m.k; + const h = c.hidden; + const qw = c.queryWidth(); + const kvw = c.kvWidth(); + const wq = wts.qkgWidth(c); + const mode: Mode = if (pass == .chunk) .prompt else .decode; + const ls = m.w.layers; + try glue.embed(&k.glue, s, b.ids, m.w.embed, b.resid, b.v0, n, h, c.embed_scale); + try glue.addNorm(&k.glue, s, .{ .resid = b.resid, .next = ls[0].input_norm, .v0 = b.v0, .x = b.x, .vx = b.vx, .rows = n, .width = h, .eps = c.eps, .blend = ls[0].blend }); + for (ls, 0..) |*l, i| { + try m.linear(b.x, h, n, l.qkg, b.qkg, mode); + try m.linear(b.vx, h, n, l.v, b.v, mode); + try m.attention(l, i, n, pass); + const g_at = b.qkg + (qw + kvw) * 2; + const g = if (n == 1) g_at else blk: { // tf_rc_gate reads G rows contiguous + try m.copy2D(g_at, wq * 2, b.g, qw * 2, qw * 2, n); + break :blk b.g; + }; + try glue.gate(&k.glue, s, b.att, g, n, qw); + try m.linear(b.att, qw, n, l.o, b.y, mode); + try glue.addNorm(&k.glue, s, .{ .resid = b.resid, .branch = b.y, .post = l.post_attn, .next = l.pre_ffn, .x = b.x, .rows = n, .width = h, .eps = c.eps }); + try m.mlpIn(l, n, mode); + try glue.siluClamp(&k.glue, s, b.gu, b.act, n, c.intermediate, c.clamp); + try m.linear(b.act, c.intermediate, n, l.down, b.y, mode); + if (i + 1 < ls.len) { + const next = &ls[i + 1]; + try glue.addNorm(&k.glue, s, .{ .resid = b.resid, .branch = b.y, .post = l.post_ffn, .next = next.input_norm, .v0 = b.v0, .x = b.x, .vx = b.vx, .rows = n, .width = h, .eps = c.eps, .blend = next.blend }); + } else try glue.addNorm(&k.glue, s, .{ .resid = b.resid, .branch = b.y, .post = l.post_ffn, .next = m.w.norm, .x = b.x, .rows = n, .width = h, .eps = c.eps }); + } + } + + /// A prompt chunk of one slot: `ids` at p0, p0 + 1, ... (at most prompt_rows), M1's prompt path and bits. + pub fn prompt(m: *const Model, slot: usize, p0: usize, ids: []const u32) !void { + const n = ids.len; + if (n == 0 or n > prompt_rows) return error.ForwardRows; + if (slot >= m.arena.slots) return error.NoSuchSlot; + if (p0 > m.context or n > m.context - p0) return error.PastContext; + try cuda.copy_source.enqueue(m.d, m.b.ids, std.mem.sliceAsBytes(ids), m.s); + try m.layers(n, .{ .chunk = .{ .slot = slot, .p0 = p0, .tokens = n } }); + } + + /// One round: chain j's `tokens` rows at p0, p0 + 1, ... of its slot, each seeing its slot's keys to its position. + pub fn round(m: *const Model, chains: []const Chain, ids: []const u32) !void { + return m.roundAs(chains, ids, if (m.graphs) .graph else .exact); + } + + /// Launches: graph (captured, the context's key groups), exact (eager, own key groups), fixed (eager, graph's). + pub const Run = enum { graph, exact, fixed }; + + /// A round's rows launched as `run` says; ids, positions, slots, table rows and committed keys upload first. + pub fn roundAs(m: *const Model, chains: []const Chain, ids: []const u32, run: Run) !void { + const n = ids.len; + const h = m.host; + try fill(h, m.arena.ids, m.arena.per_slot, m.arena.slots, m.context, chains, n); + const b = m.b; + try cuda.copy_source.enqueue(m.d, b.ids, std.mem.sliceAsBytes(ids), m.s); + try cuda.copy_source.enqueue(m.d, b.pos, std.mem.sliceAsBytes(h.pos[0..n]), m.s); + try cuda.copy_source.enqueue(m.d, b.slot, std.mem.sliceAsBytes(h.slot[0..n]), m.s); + const planned: ?attn.Planned = if (m.c.globals() == 0) null else blk: { + try cuda.copy_source.enqueue(m.d, b.ctab, std.mem.sliceAsBytes(h.ctab[0 .. n * m.arena.per_slot]), m.s); + break :blk try attn.planRows(&m.gd, m.s, h.at[0..n], run != .exact); + }; + if (run != .graph) return m.layers(n, .{ .round = planned }); + const g = &m.execs[n]; + if (g.* == null) g.* = try m.capture(n, planned); + try g.*.?.launchOn(m.s); + } + + /// The layers' launches for a round of n rows as a graph: they read every per-round value from the staged buffers. + fn capture(m: *const Model, n: usize, planned: ?attn.Planned) !cuda.graph.Exec { + try m.s.synchronize(); + try cuda.graph.beginCapture(m.s, .thread_local); + m.layers(n, .{ .round = planned }) catch |e| { + if (cuda.graph.endCapture(m.s)) |g| { + var gg = g; + gg.deinit(); + } else |_| {} + return e; + }; + var g = try cuda.graph.endCapture(m.s); + defer g.deinit(); + return g.instantiate(); + } + + /// Every captured graph dropped (the next round of each row count captures again). + pub fn dropGraphs(m: *const Model) void { + for (m.execs) |*g| if (g.*) |*e| { + e.deinit(); + g.* = null; + }; + } + + /// The row counts whose rounds have a graph. + pub fn captured(m: *const Model) usize { + var n: usize = 0; + for (m.execs) |g| n += @intFromBool(g != null); + return n; + } + + /// `ids` at seq.len, seq.len + 1, ...: a prompt chunk, or (decode) a round of one chain; seq.len moves past them. + pub fn forward(m: *const Model, seq: *Seq, ids: []const u32, mode: Mode) !void { + switch (mode) { + .prompt => try m.prompt(seq.slot, seq.len, ids), + .decode => try m.round(&.{.{ .slot = seq.slot, .p0 = seq.len, .tokens = ids.len }}, ids), + } + seq.len += ids.len; + } + + /// The last forward's row `row` through the bf16 head (the lane matmul) into `logits` [V]. + pub fn head(m: *const Model, row: usize) !void { + const h = m.c.hidden; + try m.k.lane.matmulOut(m.s, m.b.x + row * h * 2, h, 1, m.w.head, m.b.logits, .bf16, m.b.part); + } + + /// The last forward's rows [r0, r0 + rows) (at most head_rows) through the head (the prompt matmul) into `rows`. + pub fn headRows(m: *const Model, r0: usize, rows: usize) !void { + if (rows == 0 or rows > head_rows) return error.ForwardRows; + const h = m.c.hidden; + try m.k.prompt.matmulOut(m.s, m.b.x + r0 * h * 2, h, rows, m.w.head, m.b.rows, .bf16); + } + + /// The last round's rows [0, rows) through the head (the lane matmul, as head's) into `rows` [rows, V]. + pub fn rowLogits(m: *const Model, rows: usize) !void { + if (rows == 0 or rows > m.round_rows) return error.ForwardRows; + try m.k.lane.matmulOut(m.s, m.b.x, m.c.hidden, rows, m.w.head, m.b.rows, .bf16, m.b.part); + } + + /// The greedy draw of `logits` into `out` (int32). + pub fn draw(m: *const Model, out: u64) !void { + try m.k.argmaxRows(m.s, m.b.logits, 1, m.c.vocab, m.c.vocab, out); + } + + /// The greedy draws of rowLogits' rows into `out` (int32 [rows]). + pub fn drawRows(m: *const Model, rows: usize, out: u64) !void { + try m.k.argmaxRows(m.s, m.b.rows, rows, m.c.vocab, m.c.vocab, out); + } + + /// Device bytes at `ptr` into `out`, after the stream's work. + pub fn download(m: *const Model, ptr: u64, out: []u8) !void { + try m.s.synchronize(); + try m.d.check(m.d.api.cuMemcpyDtoH_v2(out.ptr, ptr, out.len), "cuMemcpyDtoH"); + } +}; + +const testing = std.testing; + +fn release() !Config { + var c = try cfg.parse(testing.allocator, @embedFile("fixtures/config.json")); + c.fp4_mlp = true; + return c; +} + +test "ricursive cuda caches: 64 rings of 3,072 rows, 16 layers' pages for the context; slots side by side" { + const c = try release(); + try testing.expectEqual(@as(usize, 16), c.globals()); + const ring = 64 * 2 * 3072 * 1024 * 2; + try testing.expectEqual(ring + 16 * 2 * 64 * 512 * 1024 * 2 + 256, Arena.bytes(c, 32768, 1)); // M1's State.bytes + try testing.expectEqual(ring + 16 * 2 * 9 * 512 * 1024 * 2 + 256, Arena.bytes(c, 4096 + 1, 1)); // a part page is a page + try testing.expectEqual(4 * slotBytes(c, 32768) + 4 * 64 * 4, Arena.bytes(c, 32768, 4)); // the table: 4 rows of 64 ids + try testing.expect(c.window - 1 + prompt_rows <= ring_rows); +} + +test "ricursive cuda scratch: every buffer on 256 bytes, about 1 GiB for 2,048 rows" { + const c = try release(); + var b: Bufs = .{}; + const total = place(bufSizes(c, 1 << 20, 1, 64), 4096, &b); + inline for (@typeInfo(Bufs).@"struct".field_names) |name| try testing.expect(@field(b, name) % buffer_align == 0); + try testing.expectEqual(@as(u64, 4096), b.resid); + try testing.expect(b.rows > b.logits and b.logits > b.part); + try testing.expect(total > 1_000_000_000 and total < 1_200_000_000); + var wide: Bufs = .{}; + const sizes = bufSizes(c, 0, 256, 64); + try testing.expectEqual(@as(usize, 256 * c.vocab * 2), sizes.rows); // a round past head_rows: its rows' logits + try testing.expectEqual(@as(usize, 256 * 64 * 4), sizes.ctab); + _ = place(sizes, 0, &wide); +} + +test "ricursive cuda round rows: chains in order, each row its slot's table row; a slot twice or a short round refused" { + const per = 3; + var table: [4 * per]i32 = undefined; + for (&table, 0..) |*x, i| x.* = @intCast(i); + var pos: [8]i32 = undefined; + var slot: [8]i32 = undefined; + var at: [8]usize = undefined; + var ctab: [8 * per]i32 = undefined; + const h: Host = .{ .pos = &pos, .slot = &slot, .at = &at, .ctab = &ctab }; + try fill(h, &table, per, 4, 1536, &.{ .{ .slot = 2, .p0 = 510, .tokens = 3 }, .{ .slot = 0, .p0 = 7, .tokens = 1 } }, 4); + try testing.expectEqualSlices(i32, &.{ 510, 511, 512, 7 }, pos[0..4]); + try testing.expectEqualSlices(i32, &.{ 2, 2, 2, 0 }, slot[0..4]); + try testing.expectEqualSlices(usize, &.{ 510, 511, 512, 7 }, at[0..4]); + try testing.expectEqualSlices(i32, &.{ 6, 7, 8, 6, 7, 8, 6, 7, 8, 0, 1, 2 }, ctab[0 .. 4 * per]); + try testing.expectError(error.SlotTwice, fill(h, &table, per, 4, 1536, &.{ .{ .slot = 1, .p0 = 0, .tokens = 1 }, .{ .slot = 1, .p0 = 1, .tokens = 1 } }, 2)); + try testing.expectError(error.ForwardRows, fill(h, &table, per, 4, 1536, &.{.{ .slot = 1, .p0 = 0, .tokens = 2 }}, 3)); + try testing.expectError(error.ForwardRows, fill(h, &table, per, 4, 1536, &.{.{ .slot = 1, .p0 = 0, .tokens = 9 }}, 9)); + try testing.expectError(error.NoSuchSlot, fill(h, &table, per, 4, 1536, &.{.{ .slot = 4, .p0 = 0, .tokens = 1 }}, 1)); + try testing.expectError(error.PastContext, fill(h, &table, per, 4, 1536, &.{.{ .slot = 3, .p0 = 1535, .tokens = 2 }}, 2)); +} + +test "ricursive cuda forward type-checks against the kernels' wrappers" { + testing.refAllDecls(Model); + testing.refAllDecls(Kernels); + testing.refAllDecls(Arena); +} diff --git a/zig/src/families/ricursive/cuda_model_kernels.zig b/zig/src/families/ricursive/cuda_model_kernels.zig new file mode 100644 index 0000000000..e91e22cec6 --- /dev/null +++ b/zig/src/families/ricursive/cuda_model_kernels.zig @@ -0,0 +1,60 @@ +//! The kernels Ripple's CUDA forward loads together: glue, prep, attention, the lane and prompt matmuls, argmax. +const cuda = @import("cuda"); +const glue = @import("cuda_kernels.zig"); +const rp = @import("cuda_prep.zig"); +const attn = @import("cuda_attn.zig"); + +const wa = cuda.window_attention; +const Stream = cuda.Stream; + +pub const Kernels = struct { + glue: glue.Kernels, + prep: rp.Kernels, + win: wa.WindowAttention, // sliding decode (tf_window_decode_tc); owns the prompt attention's module + attn: attn.Kernels, + lane: cuda.qmmf.Lane, + prompt: cuda.nvfp4.Prompt, + argmax_mod: cuda.Module, + argmax: cuda.Function, // torch.argmax's first maximum, int32 a row + + pub fn load(ctx: *const cuda.Context) !Kernels { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + const d = ctx.d; + var k: Kernels = undefined; + k.glue = try glue.Kernels.load(d); + errdefer k.glue.unload(); + k.prep = try rp.Kernels.load(d); + errdefer k.prep.unload(); + k.win = try wa.WindowAttention.load(d); + errdefer k.win.unload(); + k.attn = try attn.Kernels.init(&k.win); + k.lane = try cuda.qmmf.Lane.load(d, (try ctx.capability()) / 10); + errdefer k.lane.unload(); + k.prompt = try cuda.nvfp4.Prompt.load(d); + errdefer k.prompt.unload(); + k.argmax_mod = try cuda.Module.load(d, cuda.kernels.torch_argmax); + errdefer k.argmax_mod.unload(); + k.argmax = try k.argmax_mod.function("tf_argmax_rows_i32_kernel"); + return k; + } + + pub fn deinit(k: *Kernels) void { + k.argmax_mod.unload(); + k.prompt.unload(); + k.lane.unload(); + k.win.unload(); + k.prep.unload(); + k.glue.unload(); + k.* = undefined; + } + + /// int32 argmax of `rows` bf16 rows of `vocab` logits (`stride` apart) into `out`. + pub fn argmaxRows(k: *const Kernels, s: Stream, logits: u64, rows: usize, vocab: usize, stride: usize, out: u64) !void { + var a: cuda.Args = .{}; + a.add(logits); + a.add(out); + for ([_]usize{ rows, vocab, stride }) |x| a.add(@as(i64, @intCast(x))); + a.add(@as(i32, 0)); // dtype 0: bf16 + try cuda.launch.launch(k.argmax, .{ .grid = .{ .x = @intCast(rows), .y = 1, .z = 1 }, .block = .{ .x = 1024 }, .shared = 0 }, s, &a); + } +}; diff --git a/zig/src/families/ricursive/cuda_native.zig b/zig/src/families/ricursive/cuda_native.zig new file mode 100644 index 0000000000..2b17bd165e --- /dev/null +++ b/zig/src/families/ricursive/cuda_native.zig @@ -0,0 +1,302 @@ +//! Ripple behind the native CUDA registry: several streams, an arena slot each; the prompt cache's states; refusals. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const wts = @import("cuda_weights.zig"); +const fwd = @import("cuda_model.zig"); +const be = @import("cuda_backend.zig"); +const cfg = @import("config.zig"); + +const snap = be.snap; + +pub const model_type = "ricursive"; +pub const formats: []const []const u8 = &.{"modelopt-nvfp4"}; +pub const default_context: i64 = 32768; +pub const max_segments: u32 = 1; +pub const prompt_rows: u32 = be.prompt_rows; +/// A kept state is the caches through its position and nothing past it (no draft head reads a token on). +pub const snap_lookahead: u32 = 0; +pub const takes_drafter = false; +pub const drafter_refusal = "--drafter: Ripple on CUDA loads no drafter yet; serve it without --drafter"; +/// `explain` words this family's open refusals too. +pub const explains_open = true; +/// The server's think block rules: Ripple opens `` itself or answers bare, and its markers take no budget. +pub fn thinking(info: anytype) void { + info.reasoning_optional = true; + info.thinking_budget = false; +} + +/// GB10's memory (121.7 GiB), the budget when the server gives no reading of its room. +pub const gb10_bytes: u64 = @intFromFloat(121.7 * (1 << 30)); + +/// The server's reading of the room left after the weights; null if unreadable. +pub const Room = struct { ctx: *anyopaque, read: *const fn (*anyopaque) ?u64 }; +/// --parallel: the streams made at open (auto: the most that fit, at most a round's rows over 16-row windows). +pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false, lanes: u32 = 8, lanes_fixed: bool = false, room: ?Room = null }; + +/// What the native server drives: every stream's caches, allocated at open. +pub const Loaded = struct { + backend: lanes.backend.Backend, + facts: lanes.Model, + rows: u32, + stream_bytes: usize, + free_streams: u32, + ctx: *anyopaque, + deinit: *const fn (*anyopaque) void, + target: *be.Ripple, // the prompt cache keeps this backend's states (snap_*) +}; + +/// The weights, kernels, stream, forward and backend one open owns. +pub const Owned = struct { gpa: std.mem.Allocator, w: wts.Weights, k: fwd.Kernels, s: cuda.Stream, m: fwd.Model, rb: be.Ripple }; + +/// Device bytes past the weights for the CLI's one stream at `context` positions: scratch, caches, the backend's. +pub fn openBytes(w: *const wts.Weights, k: *const fwd.Kernels, context: usize) u64 { + return fwd.Model.openBytes(w, k, context, 1, fwd.decode_rows) + be.Ripple.openBytes(w.config.vocab, 1); +} + +/// `need` bytes in `room`, or the refusal: the caches of the asked context do not fit. +fn admit(need: u64, room: u64, context: usize) !void { + if (need <= room) return; + std.log.err("ripple: a {d}-token context needs {d:.2} GiB of caches and scratch past the weights; the CUDA memory budget leaves {d:.2} GiB", .{ context, gib(need), gib(room) }); + return error.NoStreamFits; +} + +fn gib(b: u64) f64 { + return @as(f64, @floatFromInt(b)) / (1 << 30); +} + +/// The request's context option, refused past the model's or empty. +fn contextOf(context: usize) !usize { + if (context == 0 or context > cfg.max_context) return error.PastContext; + return context; +} + +/// Streams to allocate in `room`: as asked, else (auto) the most that fit; a fixed count past them or a round refuses. +fn fit(room: u64, asked: u32, fixed: bool, plan: anytype) !usize { + const most = be.round_rows / be.max_window; + if (fixed and asked > most) return error.ParallelPastRound; + var n = std.math.clamp(@as(usize, asked), 1, most); + while (try plan.bytes(n) > room) : (n -= 1) { + if (fixed) return error.ParallelPastMemory; + if (n == 1) return error.NoStreamFits; + } + return n; +} + +/// The device bytes `n` streams take at `positions`: the forward's scratch for a round, their slots, the backend's. +const Plan = struct { + w: *const wts.Weights, + k: *const fwd.Kernels, + positions: usize, + + fn bytes(p: Plan, n: usize) !u64 { + return fwd.Model.openBytes(p.w, p.k, p.positions, n, be.round_rows) + be.Ripple.openBytes(p.w.config.vocab, n); + } +}; + +/// The checkpoint through the CUDA gate, the kernels and a stream; the forward is made by the caller. +fn begin(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8) !*Owned { + const own = try gpa.create(Owned); + errdefer gpa.destroy(own); + own.gpa = gpa; + own.w = try wts.load(gpa, io, ctx.d, dir); + errdefer own.w.deinit(); + own.k = try fwd.Kernels.load(ctx); + errdefer own.k.deinit(); + own.s = try cuda.Stream.init(ctx.d, false); + return own; +} + +/// begin's parts, freed when an open stops before its forward is made. +fn abandon(own: *Owned) void { + own.s.deinit(); + own.k.deinit(); + own.w.deinit(); + own.gpa.destroy(own); +} + +fn roomOf(room_reading: ?Room) u64 { + // the server's budget; a CLI run has none and lets the allocation refuse, as Dory's does + return if (room_reading) |r| r.read(r.ctx) orelse 0 else std.math.maxInt(u64); +} + +/// The CLI's open: one slot of `context` positions and one-row rounds, if they fit the room (null: unchecked). +pub fn openOwned(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, context_asked: usize, room_reading: ?Room) !*Owned { + const context = try contextOf(context_asked); + const own = try begin(gpa, io, ctx, dir); + errdefer abandon(own); + try admit(openBytes(&own.w, &own.k, context), roomOf(room_reading), context); + own.m = try fwd.Model.init(gpa, ctx.d, own.s, &own.w, &own.k, context, 1, fwd.decode_rows); + errdefer own.m.deinit(); + own.rb = try be.Ripple.init(gpa, ctx.d, &own.m); + return own; +} + +pub fn closeOwned(own: *Owned) void { + own.rb.deinit(); + own.m.deinit(); + abandon(own); +} + +/// The server's open: as many streams as asked and fit (auto: the most), each a slot of context and verify reserve. +pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { + if (kernels != null) return error.CapturedKernels; + if (o.lanes_fixed and o.lanes > be.round_rows / be.max_window) return error.ParallelPastRound; // before the weights load + if (o.segments > 1) return error.SegmentsUnsupported; + const positions = be.physical(try contextOf(o.context)); + const own = try begin(gpa, io, ctx, dir); + errdefer abandon(own); + const plan: Plan = .{ .w = &own.w, .k = &own.k, .positions = positions }; + const room = roomOf(o.room); + const n = fit(room, o.lanes, o.lanes_fixed, plan) catch |e| { + if (e == error.ParallelPastMemory) { + const most = fit(room, o.lanes, false, plan) catch 0; + std.log.err("--parallel {d} needs {d:.2} GiB at a {d}-token context, and the CUDA memory budget leaves {d:.2} GiB after the weights: {d} stream{s} fit", .{ o.lanes, gib(try plan.bytes(o.lanes)), o.context, gib(room), most, if (most == 1) "" else "s" }); + } + if (e == error.NoStreamFits) std.log.err("ripple: one stream at a {d}-token context needs {d:.2} GiB of caches and scratch past the weights; the CUDA memory budget leaves {d:.2} GiB", .{ o.context, gib(try plan.bytes(1)), gib(room) }); + return e; + }; + own.m = try fwd.Model.init(gpa, ctx.d, own.s, &own.w, &own.k, positions, n, be.round_rows); + errdefer own.m.deinit(); + own.rb = try be.Ripple.init(gpa, ctx.d, &own.m); + std.log.info("Ripple on CUDA: {d} stream{s} of {d} positions (a {d}-token context and a {d}-row verify reserve), rounds of {d} rows, prompt chunks of {d}", .{ n, if (n == 1) "" else "s", positions, o.context, be.reserve, be.round_rows, fwd.prompt_rows }); + return .{ + .backend = own.rb.backend(), + .facts = own.rb.facts(), + .rows = be.max_window, + .stream_bytes = fwd.Arena.bytes(own.w.config, positions, 1), + .free_streams = @intCast(n), + .ctx = own, + .deinit = release, + .target = &own.rb, + }; +} + +fn release(p: *anyopaque) void { + closeOwned(@ptrCast(@alignCast(p))); +} + +/// Whether a request's caches could start now (the lane host's admission, on its engine thread). +pub fn fits(p: *anyopaque, prompt: u32, max_tokens: u32) bool { + const own: *const Owned = @ptrCast(@alignCast(p)); + return own.rb.admits(prompt, max_tokens); +} + +fn ripple(ptr: *anyopaque) *be.Ripple { + return @ptrCast(@alignCast(ptr)); +} + +/// Bytes a state at `at` prompt tokens takes. +pub fn snap_bytes(ptr: *anyopaque, at: u32) u64 { + return snap.bytes(ripple(ptr).m.c, at); +} + +/// The prompt pass stands at `at` (a chunk start or a mark): copy the stream's caches there. +pub fn snap_save(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!*anyopaque { + return @ptrCast(try ripple(ptr).saveState(owner, at)); +} + +pub fn snap_restore(ptr: *anyopaque, owner: ?*anyopaque, saved: *anyopaque) anyerror!void { + return ripple(ptr).restoreState(owner, @ptrCast(@alignCast(saved))); +} + +pub fn snap_drop(ptr: *anyopaque, saved: *anyopaque) void { + ripple(ptr).dropState(@ptrCast(@alignCast(saved))); +} + +/// A request or load this engine refuses, in words; null: none of its own. +pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { + return switch (err) { + error.PastContext => "the prompt and its reply exceed this server's context window: shorten it, lower max_tokens or raise --context (at most 262,144)", + error.NoFreeSlot => "every stream's caches are in use: retry once another request ends", + error.NoStreamFits => "the CUDA memory budget fits no Ripple stream at this context after the weights: lower --context, free device memory, or adjust TENSORFOLD_MEMORY_RESERVE_GIB / TENSORFOLD_CUDA_MEMORY_LIMIT_GB", + error.ParallelPastRound => std.fmt.comptimePrint("--parallel: Ripple on CUDA holds at most {d} streams (a round's rows over 16-row windows): serve --parallel {d} or fewer, or auto", .{ be.round_rows / be.max_window, be.round_rows / be.max_window }), + error.ParallelPastMemory => "--parallel: the CUDA memory budget fits fewer Ripple streams than asked (each stream's caches are allocated at load): lower --parallel or --context, serve --parallel auto, or adjust TENSORFOLD_MEMORY_RESERVE_GIB / TENSORFOLD_CUDA_MEMORY_LIMIT_GB", + error.SegmentsUnsupported => "--segments: Ripple on CUDA runs its prompt chunks one after another: serve without --segments", + error.NoDraftHead => "Ripple on CUDA has no draft head: drafts come from the prompt's own text only", + error.StructuresUnsupported => "structured output is not built for Ripple on CUDA yet", + error.CapturedKernels => "Ripple runs on the engine's own kernels: unset TENSORFOLD_CUDA_KERNELS", + error.BuiltWithoutKernels => "this build carries no CUDA kernel images: build with -Dnvcc", + error.RicursiveCudaNeedsNvfp4 => "Ripple on CUDA serves the NVFP4 release (ModelOpt W4A16 MLPs, bf16 attention): the BF16 checkpoint's 130 GB do not fit GB10", + error.UnsupportedRicursiveCudaShape => "this Ripple checkpoint's shape is not one the CUDA kernels take (group of 2 to 16 query heads in pairs, hidden size at most 8,192, widths in whole 128s)", + error.Unimplemented => "Ripple's CUDA global-layer decode is not in this build yet", + else => null, + }; +} + +test "ricursive cuda native: a reply thinks or answers bare, with no thinking budget, as on Metal" { + var info: struct { reasoning_optional: bool = false, thinking_budget: bool = true } = .{}; + thinking(&info); + try std.testing.expect(info.reasoning_optional and !info.thinking_budget); +} + +test "ricursive cuda native: every refusal of its own is worded, the rest left to the server" { + for ([_]anyerror{ error.PastContext, error.NoFreeSlot, error.NoStreamFits, error.ParallelPastRound, error.ParallelPastMemory, error.SegmentsUnsupported, error.NoDraftHead, error.StructuresUnsupported, error.CapturedKernels, error.BuiltWithoutKernels, error.RicursiveCudaNeedsNvfp4, error.UnsupportedRicursiveCudaShape, error.Unimplemented }) |e| { + try std.testing.expect(explain(null, e) != null); + } + for ([_]anyerror{ error.OutOfMemory, error.SamplingUnsupported, error.LogprobsUnsupported, error.WindowTooWide }) |e| { + try std.testing.expect(explain(null, e) == null); + } +} + +test "ricursive cuda native: the release at 32,768 and 262,144 positions on GB10; a context past the room refused" { + var c = try cfg.parse(std.testing.allocator, @embedFile("fixtures/config.json")); + c.fp4_mlp = true; + const weights: u64 = wts.deviceBytes(c); + const scratch: u64 = 1_200_000_000; // the forward's 2,048 rows (cuda_model's test bounds it) + for ([_]usize{ 32768, 262144 }) |context| { + const need = fwd.Arena.bytes(c, be.physical(context), 1) + scratch; + try admit(need, gb10_bytes - weights, context); + } + try std.testing.expectEqual(@as(usize, 32768), try contextOf(32768)); + try std.testing.expectError(error.PastContext, contextOf(cfg.max_context + 1)); + try std.testing.expectError(error.PastContext, contextOf(0)); +} + +/// A synthetic plan: `stream` bytes a stream plus `fixed` bytes of scratch. +const Fake = struct { + stream: u64, + fixed: u64, + fn bytes(f: Fake, n: usize) !u64 { + return f.fixed + n * f.stream; + } +}; + +test "ricursive cuda native: the streams made at open fit the room; auto takes what fits, a fixed count refuses" { + const g: u64 = 1 << 30; + const one: Fake = .{ .stream = g, .fixed = 0 }; + try std.testing.expectEqual(@as(usize, 2), try fit(2 * g, 8, false, one)); + try std.testing.expectError(error.ParallelPastMemory, fit(2 * g, 8, true, one)); + try std.testing.expectEqual(@as(usize, 1), try fit(5 * g / 2, 8, false, Fake{ .stream = g, .fixed = g })); + try std.testing.expectError(error.NoStreamFits, fit(g - 1, 8, false, one)); + try std.testing.expectEqual(@as(usize, 8), try fit(std.math.maxInt(u64), 8, false, one)); + try std.testing.expectEqual(@as(usize, 8), try fit(std.math.maxInt(u64), 64, false, one)); + try std.testing.expectError(error.ParallelPastRound, fit(1024 * g, 9, true, one)); +} + +test "ricursive cuda native: streams at a context on GB10 by the release's slot bytes" { + var c = try cfg.parse(std.testing.allocator, @embedFile("fixtures/config.json")); + c.fp4_mlp = true; + const room = gb10_bytes - wts.deviceBytes(c); + const scratch: u64 = 1_400_000_000; // a round's and a chunk's scratch, about (cuda_model's test bounds the chunk's) + // a slot: 64 rings of 3,072 rows (805 MB) and 16 global layers' pages for the positions + const Slots = struct { + c: cfg.Config, + positions: usize, + fixed: u64, + fn bytes(p: @This(), n: usize) !u64 { + return p.fixed + fwd.Arena.bytes(p.c, p.positions, n); + } + }; + try std.testing.expectEqual(@as(usize, 8), try fit(room, 8, false, Slots{ .c = c, .positions = be.physical(32768), .fixed = scratch })); + try std.testing.expect(try fit(room, 8, false, Slots{ .c = c, .positions = be.physical(262144), .fixed = scratch }) < 8); +} + +test "the prompt cache's hooks compile against the store's vtable" { + const pc: struct { bytes: *const fn (*anyopaque, u32) u64, save: *const fn (*anyopaque, ?*anyopaque, u32) anyerror!*anyopaque, restore: *const fn (*anyopaque, ?*anyopaque, *anyopaque) anyerror!void, drop: *const fn (*anyopaque, *anyopaque) void } = .{ .bytes = snap_bytes, .save = snap_save, .restore = snap_restore, .drop = snap_drop }; + try std.testing.expect(@intFromPtr(pc.save) != 0); +} + +test "ricursive cuda native: the open path type-checks against the registry's shape" { + std.testing.refAllDecls(@This()); +} diff --git a/zig/src/families/ricursive/cuda_prep.zig b/zig/src/families/ricursive/cuda_prep.zig new file mode 100644 index 0000000000..9c78409c31 --- /dev/null +++ b/zig/src/families/ricursive/cuda_prep.zig @@ -0,0 +1,170 @@ +//! Ripple's attention prep on CUDA (ricursive_prep.cu): glue.metal's rc_prep, q dense and k, v into a ring or pages. + +const std = @import("std"); +const cuda = @import("cuda"); + +const Stream = cuda.Stream; +const Function = cuda.Function; +/// A stream's pages as window_attention's paged prompt reads them; `stride` (a slot's table entries) is prep_rows'. +pub const Pages = cuda.window_attention.Pages; + +const head_dim = 128; + +/// A warp a (head, row), 8 warps a block: tf_rc_prep's and tf_rc_prep_rows' __launch_bounds__. +const warps = 8; +const threads = 32 * warps; + +fn int(x: usize) c_int { + return @intCast(x); +} + +fn ceil(n: usize, d: usize) u32 { + return @intCast((n + d - 1) / d); +} + +pub const Kernels = struct { + mod: cuda.Module, + prep: Function, + prep_rows: Function, + + pub fn load(d: *const cuda.Driver) !Kernels { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var mod = try cuda.Module.load(d, cuda.kernels.ricursive_prep); + errdefer mod.unload(); + return .{ .mod = mod, .prep = try mod.function("tf_rc_prep"), .prep_rows = try mod.function("tf_rc_prep_rows") }; + } + + pub fn unload(k: *Kernels) void { + k.mod.unload(); + k.* = undefined; + } +}; + +/// One projection's bf16 output rows (ricursive_prep.cuh Rows): the first row and the elements between rows. +pub const Rows = extern struct { + x: u64, + stride: i32, + pad: i32 = 0, + + pub fn of(x: u64, stride: usize) Rows { + return .{ .x = x, .stride = int(stride) }; + } +}; + +/// `rows` rows at pos + row: q, and k and v at ring slot (pos + row) % ring or in pages; `rope` false: NoPE, normed. +pub const Prep = struct { xq: Rows, xk: Rows, xv: Rows, inv: u64, qw: u64, kw: u64, q: u64, kc: u64, vc: u64, rows: usize, heads: usize, kv_heads: usize, pos: usize, ring: usize, rope: bool, eps: f32, pg: ?Pages = null }; + +/// The writes the kernel takes: a ring at least the call's rows, or pages each holding `tokens` whole rows. +pub fn admit(p: Prep) !void { + if (p.pos + p.rows > std.math.maxInt(i32)) return error.PrepShape; + if (p.pg) |pg| { + if (pg.tokens == 0 or pg.tokens > std.math.maxInt(i32) or pg.page_elems < pg.tokens * p.kv_heads * head_dim or pg.page_elems > std.math.maxInt(i63)) return error.PageShape; + } else if (p.rows > p.ring) return error.RingTooSmall; // a call's rows would overwrite their own slots +} + +/// A block of 8 heads across, a row down. +pub fn grid(p: Prep) cuda.Dim3 { + return .{ .x = ceil(p.heads + p.kv_heads, warps), .y = @intCast(p.rows), .z = 1 }; +} + +pub fn prep(k: *const Kernels, s: Stream, p: Prep) !void { + return prepVia(k.prep, null, s, p); +} + +/// The launch on `f`: the served tf_rc_prep, or (keep not null) contract_controls' tf_ctl_rc_prep keeping those points. +pub fn prepVia(f: Function, keep: ?i32, s: Stream, p: Prep) !void { + if (p.rows == 0) return; + try admit(p); + var a: cuda.Args = .{}; + for ([_]Rows{ p.xq, p.xk, p.xv }) |r| a.add(r); + for ([_]u64{ p.inv, p.qw, p.kw, p.q, p.kc, p.vc }) |v| a.add(v); + for ([_]usize{ p.rows, p.heads, p.kv_heads, p.pos, p.ring }) |v| a.add(int(v)); + a.add(@as(c_int, @intFromBool(p.rope))); + a.add(p.eps); + a.add(if (p.pg) |pg| pg.table else @as(u64, 0)); // null: the ring + a.add(@as(i64, if (p.pg) |pg| @intCast(pg.page_elems) else 0)); + a.add(@as(c_int, if (p.pg) |pg| @intCast(pg.tokens) else 0)); + if (keep) |v| a.add(@as(c_int, v)); + try cuda.launch.launch(f, .{ .grid = grid(p), .block = .{ .x = threads }, .shared = 0 }, s, &a); +} + +/// tf_rc_prep_rows: row r at pos[r] of stream slot[r] (device int32, replayable); rows a ring apart must not share. +pub const PrepRows = struct { xq: Rows, xk: Rows, xv: Rows, inv: u64, qw: u64, kw: u64, q: u64, kc: u64, vc: u64, rows: usize, heads: usize, kv_heads: usize, pos: u64, slot: ?u64 = null, ring: usize, rope: bool, eps: f32, pg: ?Pages = null }; + +/// The writes tf_rc_prep_rows takes: a ring (positions on the device, unchecked), or pages of `tokens` whole rows. +pub fn admitRows(p: PrepRows) !void { + if (p.pg) |pg| { + if (pg.tokens == 0 or pg.tokens > std.math.maxInt(i32) or pg.page_elems < pg.tokens * p.kv_heads * head_dim or pg.page_elems > std.math.maxInt(i63) or pg.stride > std.math.maxInt(i32)) return error.PageShape; + } else if (p.ring == 0 or p.ring > std.math.maxInt(i32)) return error.RingTooSmall; +} + +/// tf_rc_prep's grid: a block of 8 heads across, a row down. +pub fn gridRows(p: PrepRows) cuda.Dim3 { + return .{ .x = ceil(p.heads + p.kv_heads, warps), .y = @intCast(p.rows), .z = 1 }; +} + +pub fn prepRows(k: *const Kernels, s: Stream, p: PrepRows) !void { + return prepRowsVia(k.prep_rows, null, s, p); +} + +/// The launch on `f`: the served tf_rc_prep_rows, or (keep not null) contract_controls' tf_ctl_rc_prep_rows. +pub fn prepRowsVia(f: Function, keep: ?i32, s: Stream, p: PrepRows) !void { + if (p.rows == 0) return; + try admitRows(p); + var a: cuda.Args = .{}; + for ([_]Rows{ p.xq, p.xk, p.xv }) |r| a.add(r); + for ([_]u64{ p.inv, p.qw, p.kw, p.q, p.kc, p.vc }) |v| a.add(v); + for ([_]usize{ p.rows, p.heads, p.kv_heads }) |v| a.add(int(v)); + a.add(p.pos); + a.add(p.slot orelse @as(u64, 0)); // null: stream 0 + a.add(int(p.ring)); + a.add(@as(c_int, @intFromBool(p.rope))); + a.add(p.eps); + a.add(if (p.pg) |pg| pg.table else @as(u64, 0)); // null: the rings + a.add(@as(i64, if (p.pg) |pg| @intCast(pg.page_elems) else 0)); + a.add(@as(c_int, if (p.pg) |pg| @intCast(pg.tokens) else 0)); + a.add(@as(c_int, if (p.pg) |pg| @intCast(pg.stride) else 0)); + if (keep) |v| a.add(@as(c_int, v)); + try cuda.launch.launch(f, .{ .grid = gridRows(p), .block = .{ .x = threads }, .shared = 0 }, s, &a); +} + +test "launch geometry: 88 heads in 11 blocks a row, Rows as the kernel's struct" { + const p: Prep = .{ .xq = Rows.of(0, 0), .xk = Rows.of(0, 0), .xv = Rows.of(0, 0), .inv = 0, .qw = 0, .kw = 0, .q = 0, .kc = 0, .vc = 0, .rows = 2048, .heads = 80, .kv_heads = 8, .pos = 0, .ring = 3072, .rope = true, .eps = 1e-6 }; + const g = grid(p); + try std.testing.expectEqual(@as(c_uint, 11), g.x); + try std.testing.expectEqual(@as(c_uint, 2048), g.y); + try std.testing.expectEqual(@as(usize, 16), @sizeOf(Rows)); +} + +test "admit: a ring holds the call's rows; pages hold 512 whole rows of 8 heads" { + const t = std.testing; + var p: Prep = .{ .xq = Rows.of(0, 0), .xk = Rows.of(0, 0), .xv = Rows.of(0, 0), .inv = 0, .qw = 0, .kw = 0, .q = 0, .kc = 0, .vc = 0, .rows = 2048, .heads = 80, .kv_heads = 8, .pos = 6000, .ring = 3072, .rope = true, .eps = 1e-6 }; + try admit(p); + p.ring = 2047; + try t.expectError(error.RingTooSmall, admit(p)); + p.pg = .{ .table = 0, .page_elems = 512 * 8 * 128, .tokens = 512 }; + try admit(p); // pages: no ring + p.pg.?.page_elems -= 1; + try t.expectError(error.PageShape, admit(p)); + p.pg = .{ .table = 0, .page_elems = 0, .tokens = 0 }; + try t.expectError(error.PageShape, admit(p)); +} + +test "rows: tf_rc_prep's grid; a ring of some rows, or pages of whole rows and an int32 table stride" { + const t = std.testing; + var p: PrepRows = .{ .xq = Rows.of(0, 0), .xk = Rows.of(0, 0), .xv = Rows.of(0, 0), .inv = 0, .qw = 0, .kw = 0, .q = 0, .kc = 0, .vc = 0, .rows = 64, .heads = 80, .kv_heads = 8, .pos = 0, .ring = 3072, .rope = true, .eps = 1e-6 }; + const g = gridRows(p); + try t.expectEqual(@as(c_uint, 11), g.x); + try t.expectEqual(@as(c_uint, 64), g.y); + try admitRows(p); + p.rows = 4096; // several streams' rows: more than one ring holds, each in its own + try admitRows(p); + p.ring = 0; + try t.expectError(error.RingTooSmall, admitRows(p)); + p.pg = .{ .table = 0, .page_elems = 512 * 8 * 128, .tokens = 512, .stride = 256 }; + try admitRows(p); // pages: no ring + p.pg.?.stride = 1 << 31; + try t.expectError(error.PageShape, admitRows(p)); + p.pg = .{ .table = 0, .page_elems = 512 * 8 * 128 - 1, .tokens = 512, .stride = 256 }; + try t.expectError(error.PageShape, admitRows(p)); +} diff --git a/zig/src/families/ricursive/cuda_snap.zig b/zig/src/families/ricursive/cuda_snap.zig new file mode 100644 index 0000000000..9a8c6fc180 --- /dev/null +++ b/zig/src/families/ricursive/cuda_snap.zig @@ -0,0 +1,169 @@ +//! A stream's caches for the prompt cache: each sliding layer's last window - 1 ring rows, each global's first `at`. +const std = @import("std"); +const cuda = @import("cuda"); +const cfg = @import("config.zig"); +const fwd = @import("cuda_model.zig"); + +/// What a state was made under: its layers, which of them are global, the window and the KV width. +pub const Key = struct { layers: u32, global: [cfg.max_layers]bool, window: u32, kv_width: u32 }; + +/// The key every state of a model at `c` carries. +pub fn keyOf(c: cfg.Config) Key { + return .{ .layers = @intCast(c.layers), .global = c.global, .window = @intCast(c.window), .kv_width = @intCast(c.kvWidth()) }; +} + +pub fn sameKey(a: Key, b: Key) bool { + return std.meta.eql(a, b); +} + +/// A kept state: a layer at a time, K then V, global layers' positions [0, at), sliding rings' last positions in order. +pub const Saved = struct { at: u32, key: Key, buf: cuda.DeviceBuffer }; + +/// Positions a sliding ring keeps at `at`: a later row at `at` or past reads keys from at - window + 1 on. +fn tail(window: usize, at: usize) usize { + return @min(at, window - 1); +} + +/// Bytes a state at `at` takes. +pub fn bytes(c: cfg.Config, at: usize) u64 { + const row = c.kvWidth() * 2; + var n: usize = 0; + for (0..c.layers) |i| n += 2 * (if (c.global[i]) at else tail(c.window, at)) * row; + return n; +} + +/// Each layer's K then V rows of `slot` for positions [first, at): one call a run of consecutive rows, in order. +fn each(a: *const fwd.Arena, window: usize, slot: usize, at: usize, ctx: anytype, comptime f: fn (@TypeOf(ctx), dev: u64, off: usize, len: usize) anyerror!void) !void { + var off: usize = 0; + for (0..a.caches.len) |i| { + const sp = a.span(i, slot); + for ([_]u64{ sp.k, sp.v }) |base| { + var p = if (sp.ring) at - tail(window, at) else 0; + while (p < at) { + const row = if (sp.ring) p % sp.rows else p; // a global layer's table is the identity + const run = @min(at - p, sp.rows - row); + try f(ctx, base + row * sp.row_bytes, off, run * sp.row_bytes); + off += run * sp.row_bytes; + p += run; + } + } + } +} + +const Copy = struct { d: *const cuda.Driver, s: cuda.Stream, buf: u64, out: bool }; + +fn copy(c: Copy, dev: u64, off: usize, len: usize) anyerror!void { + const src = if (c.out) dev else c.buf + off; + const dst = if (c.out) c.buf + off else dev; + try c.d.check(c.d.api.cuMemcpyDtoDAsync_v2(dst, src, len, c.s.handle), "cuMemcpyDtoDAsync"); +} + +/// Copy `slot`'s state at `at` into new storage, queued on the model's stream behind the chunk that ends there. +pub fn save(gpa: std.mem.Allocator, m: *const fwd.Model, slot: usize, at: u32) !*Saved { + if (at == 0 or at > m.context) return error.NotAtMark; + const k = try gpa.create(Saved); + errdefer gpa.destroy(k); + k.* = .{ .at = at, .key = keyOf(m.c), .buf = try cuda.DeviceBuffer.alloc(m.d, @max(bytes(m.c, at), 256)) }; + errdefer k.buf.free(); + try each(&m.arena, m.c.window, slot, at, Copy{ .d = m.d, .s = m.s, .buf = k.buf.ptr, .out = true }, copy); + return k; +} + +/// Make `slot`'s caches hold `k`'s state, queued before its next chunk; another shape's state is refused. +pub fn restore(m: *const fwd.Model, slot: usize, k: *const Saved) !void { + if (!sameKey(k.key, keyOf(m.c))) return error.StateOfAnotherShape; + if (k.at == 0 or k.at > m.context) return error.PastContext; + try each(&m.arena, m.c.window, slot, k.at, Copy{ .d = m.d, .s = m.s, .buf = k.buf.ptr, .out = false }, copy); +} + +/// Free a state once the queued copies of it are done. +pub fn drop(gpa: std.mem.Allocator, s: cuda.Stream, k: *Saved) void { + s.synchronize() catch {}; + k.buf.free(); + gpa.destroy(k); +} + +const testing = std.testing; + +const Count = struct { + runs: *std.ArrayList([3]usize), + fn add(c: Count, dev: u64, off: usize, len: usize) anyerror!void { + try c.runs.append(testing.allocator, .{ dev, off, len }); + } +}; + +fn release() !cfg.Config { + var c = try cfg.parse(testing.allocator, @embedFile("fixtures/config.json")); + c.fp4_mlp = true; + return c; +} + +/// An arena of `slots` slots at `context` positions with fake device addresses (no device memory): layer i at i << 40. +fn fakeArena(c: cfg.Config, caches: []fwd.Cache, slots: usize, context: usize) fwd.Arena { + const per = fwd.pagesOf(context); + const kvw = c.kvWidth(); + for (caches, 0..) |*x, i| { + const half: usize = slots * if (c.global[i]) per * fwd.page_tokens * kvw * 2 else fwd.ring_rows * kvw * 2; + x.* = .{ .kc = @as(u64, i + 1) << 40, .vc = (@as(u64, i + 1) << 40) + half }; + if (c.global[i]) x.pg = .{ .table = 0, .page_elems = fwd.page_tokens * kvw, .tokens = fwd.page_tokens, .stride = per }; + } + return .{ .gpa = testing.allocator, .kv = undefined, .caches = caches, .ids = &.{}, .slots = slots, .per_slot = per, .kvw = kvw }; +} + +test "ricursive cuda state: each ring's last 1,023 rows at their ring rows, wrapped in two runs; globals' first at rows" { + const c = try release(); + var caches: [cfg.max_layers]fwd.Cache = undefined; + const a = fakeArena(c, caches[0..c.layers], 3, 8192); + const rb = c.kvWidth() * 2; + const at = 3500; // a ring of 3,072: positions 2,477 to 3,071 then 3,072 to 3,499 at ring rows 0 to 427 + try testing.expectEqual(@as(usize, 2048), rb); + try testing.expectEqual(@as(u64, 2 * (16 * at + 64 * 1023) * rb), bytes(c, at)); + var runs: std.ArrayList([3]usize) = .empty; + defer runs.deinit(testing.allocator); + try each(&a, c.window, 2, at, Count{ .runs = &runs }, Count.add); + try testing.expectEqual(@as(usize, 2 * (16 + 64 * 2)), runs.items.len); + // layer 0 is sliding: slot 2's ring starts 2 * 3,072 rows in; its K rows 2,477.. then 0.. + const k0 = caches[0].kc + 2 * fwd.ring_rows * rb; + try testing.expectEqual([3]usize{ k0 + (at - 1023) * rb, 0, (fwd.ring_rows - (at - 1023)) * rb }, runs.items[0]); + try testing.expectEqual([3]usize{ k0, runs.items[0][2], (at % fwd.ring_rows) * rb }, runs.items[1]); + try testing.expectEqual(caches[0].vc + 2 * fwd.ring_rows * rb + (at - 1023) * rb, runs.items[2][0]); + // layer 4 is global: one run of K and one of V from the slot's first page (the identity table) + try testing.expect(c.global[4]); + const g = runs.items[4 * 4 ..][0..2]; + try testing.expectEqual([3]usize{ caches[4].kc + 2 * fwd.pagesOf(8192) * fwd.page_tokens * rb, g[0][1], at * rb }, g[0]); + try testing.expectEqual(caches[4].vc + 2 * fwd.pagesOf(8192) * fwd.page_tokens * rb, g[1][0]); + var total: usize = 0; + for (runs.items, 0..) |r, i| { + if (i > 0) try testing.expectEqual(runs.items[i - 1][1] + runs.items[i - 1][2], r[1]); // packed in order + total += r[2]; + } + try testing.expectEqual(bytes(c, at), total); +} + +test "ricursive cuda state: before the window fills, a ring keeps every position from 0 in one run" { + const c = try release(); + var caches: [cfg.max_layers]fwd.Cache = undefined; + const a = fakeArena(c, caches[0..c.layers], 1, 4096); + const at = 510; // a page edge two rows on + var runs: std.ArrayList([3]usize) = .empty; + defer runs.deinit(testing.allocator); + try each(&a, c.window, 0, at, Count{ .runs = &runs }, Count.add); + try testing.expectEqual(@as(usize, 2 * c.layers), runs.items.len); + for (runs.items) |r| try testing.expectEqual(@as(usize, at * 2048), r[2]); + try testing.expectEqual(@as(u64, 2 * c.layers * at * 2048), bytes(c, at)); +} + +test "ricursive cuda state: the key names the shape; another shape's state is refused" { + const c = try release(); + var other = c; + other.global[0] = true; + try testing.expect(sameKey(keyOf(c), keyOf(c))); + try testing.expect(!sameKey(keyOf(c), keyOf(other))); + var m: fwd.Model = undefined; + m.c = c; + m.context = 8192; + var buf: cuda.DeviceBuffer = undefined; + buf.ptr = 0; + const k: Saved = .{ .at = 10, .key = keyOf(other), .buf = buf }; + try testing.expectError(error.StateOfAnotherShape, restore(&m, 0, &k)); +} diff --git a/zig/src/families/ricursive/cuda_weights.zig b/zig/src/families/ricursive/cuda_weights.zig new file mode 100644 index 0000000000..1b755ddec5 --- /dev/null +++ b/zig/src/families/ricursive/cuda_weights.zig @@ -0,0 +1,431 @@ +//! Ripple's weights on CUDA: NVFP4 MLPs for the lane matmul, bf16 attention (q, k, g as one) and head; gate refusals. +const std = @import("std"); +const core = @import("core"); +const cuda = @import("cuda"); +const cfg = @import("config.zig"); +const glue = @import("cuda_kernels.zig"); +const attn = @import("cuda_attn.zig"); + +const Config = cfg.Config; +const Checkpoint = core.checkpoint.Checkpoint; +const Tensor = core.checkpoint.Tensor; +const Weight = cuda.qmmf.Weight; +const DeviceBuffer = cuda.DeviceBuffer; + +/// Bytes a small tensor's device copy starts on: tf_rc_add_norm reads 16-byte chunks. +const small_align = 256; + +/// What CUDA serves: the NVFP4 release's MLPs, and shapes the kernels and packers take. +pub fn gate(c: *const Config) !void { + if (!c.fp4_mlp) return error.RicursiveCudaNeedsNvfp4; + const g = c.group(); + // tf_window_decode_tc holds a key head's queries in one tile (at most 16); the prompt kernel pairs them + if (g % attn.heads_a_block != 0 or g > 16) return error.UnsupportedRicursiveCudaShape; + if (c.hidden % 8 != 0 or c.hidden > glue.most_width) return error.UnsupportedRicursiveCudaShape; + // packBf16 and packWords take k in whole 64s; every output a whole number of 128-row blocks (no padding) + for ([_]usize{ c.hidden, c.queryWidth(), c.intermediate }) |k| if (k % cuda.qmmf.group != 0) return error.UnsupportedRicursiveCudaShape; + for ([_]usize{ c.queryWidth(), c.kvWidth(), c.hidden, c.intermediate, c.vocab }) |n| if (n % 128 != 0) return error.UnsupportedRicursiveCudaShape; +} + +/// The fused projection's rows: q, then k, then g; v stays apart (its input is the value-residual blend). +pub fn qkgWidth(c: Config) usize { + return 2 * c.queryWidth() + c.kvWidth(); +} + +/// An NVFP4 projection as ModelOpt stores it: nibbles [n, k / 2], E4M3 scales [n, k / 16], the fp32 global. +pub const Fp4 = struct { w: Tensor, scales: Tensor, global: f32, n: usize, k: usize }; + +pub const LayerMap = struct { + q: Tensor, + k: Tensor, + v: Tensor, + g: Tensor, + o: Tensor, + q_norm: Tensor, + k_norm: Tensor, + lambda: u16, // value_res_lambda's bf16 bits + input_norm: Tensor, + post_attn: Tensor, + pre_ffn: Tensor, + post_ffn: Tensor, + gate: Fp4, + up: Fp4, + down: Fp4, +}; + +/// Every tensor the CUDA engine reads, checked by dtype and shape; input_scale is taken and unread, as on Metal. +pub const Map = struct { + gpa: std.mem.Allocator, + embed: Tensor, + head: Tensor, + norm: Tensor, + layers: []LayerMap, + + pub fn deinit(m: *Map) void { + m.gpa.free(m.layers); + m.* = undefined; + } +}; + +const Names = struct { + ck: *Checkpoint, + buf: [160]u8 = undefined, + + fn bf16(n: *Names, shape: []const usize, comptime fmt: []const u8, args: anytype) !Tensor { + return n.ck.expect(try std.fmt.bufPrint(&n.buf, fmt, args), .bf16, shape); + } + + fn fp4(n: *Names, rows: usize, k: usize, comptime fmt: []const u8, args: anytype) !Fp4 { + const w = try n.ck.expect(try std.fmt.bufPrint(&n.buf, fmt ++ ".weight", args), .u8, &.{ rows, k / 2 }); + const s = try n.ck.expect(try std.fmt.bufPrint(&n.buf, fmt ++ ".weight_scale", args), .f8_e4m3, &.{ rows, k / 16 }); + const t = try n.ck.expect(try std.fmt.bufPrint(&n.buf, fmt ++ ".weight_scale_2", args), .f32, &.{}); + const v: f32 = @bitCast(std.mem.readInt(u32, t.bytes[0..4], .little)); + if (!(v > 0) or !std.math.isFinite(v)) { + std.log.err("ricursive: " ++ fmt ++ ".weight_scale_2 {e} is not a positive finite scale", args ++ .{v}); + return error.UnexpectedTensor; + } + const input = try std.fmt.bufPrint(&n.buf, fmt ++ ".input_scale", args); + if (n.ck.has(input)) _ = try n.ck.get(input); // the activations stay bf16 (W4A16): taken, never read + return .{ .w = w, .scales = s, .global = v, .n = rows, .k = k }; + } +}; + +/// The release's names (PR 2231's weights.zig); a missing, mistyped or leftover tensor refuses the checkpoint. +pub fn map(gpa: std.mem.Allocator, ck: *Checkpoint, c: Config) !Map { + var nm: Names = .{ .ck = ck }; + const h = c.hidden; + const q = c.queryWidth(); + const kv = c.kvWidth(); + const f = c.intermediate; + var m: Map = .{ .gpa = gpa, .embed = undefined, .head = undefined, .norm = undefined, .layers = try gpa.alloc(LayerMap, c.layers) }; + errdefer gpa.free(m.layers); + m.embed = try nm.bf16(&.{ c.vocab, h }, "model.embed_tokens.weight", .{}); + m.head = try nm.bf16(&.{ c.vocab, h }, "lm_head.weight", .{}); + m.norm = try nm.bf16(&.{h}, "model.norm.weight", .{}); + for (m.layers, 0..) |*l, i| { + const p = "model.layers.{d}."; + const a = p ++ "self_attn."; + const lambda = try nm.bf16(&.{1}, a ++ "value_res_lambda", .{i}); + l.* = .{ + .q = try nm.bf16(&.{ q, h }, a ++ "q_proj.weight", .{i}), + .k = try nm.bf16(&.{ kv, h }, a ++ "k_proj.weight", .{i}), + .v = try nm.bf16(&.{ kv, h }, a ++ "v_proj.weight", .{i}), + .g = try nm.bf16(&.{ q, h }, a ++ "g_proj.weight", .{i}), + .o = try nm.bf16(&.{ h, q }, a ++ "o_proj.weight", .{i}), + .q_norm = try nm.bf16(&.{cfg.head_dim}, a ++ "q_norm.weight", .{i}), + .k_norm = try nm.bf16(&.{cfg.head_dim}, a ++ "k_norm.weight", .{i}), + .lambda = std.mem.readInt(u16, lambda.bytes[0..2], .little), + .input_norm = try nm.bf16(&.{h}, p ++ "input_layernorm.weight", .{i}), + .post_attn = try nm.bf16(&.{h}, p ++ "post_attention_layernorm.weight", .{i}), + .pre_ffn = try nm.bf16(&.{h}, p ++ "pre_feedforward_layernorm.weight", .{i}), + .post_ffn = try nm.bf16(&.{h}, p ++ "post_feedforward_layernorm.weight", .{i}), + .gate = try nm.fp4(f, h, p ++ "mlp.gate_proj", .{i}), + .up = try nm.fp4(f, h, p ++ "mlp.up_proj", .{i}), + .down = try nm.fp4(h, f, p ++ "mlp.down_proj", .{i}), + }; + } + if (ck.unused() != 0) return error.UnusedTensors; + return m; +} + +/// torch's bf16 sigmoid of a bf16 scalar: fp32 1 / (1 + exp(-x)) to bf16, nearest even (Metal's weights.zig's copy). +pub fn blendOf(lambda: u16) f32 { + const x: f32 = @bitCast(@as(u32, lambda) << 16); + const s: f32 = 1.0 / (1.0 + @exp(-x)); + const bits: u32 = @bitCast(s); + const rounded = (bits + 0x7fff + ((bits >> 16) & 1)) & 0xffff0000; + return @bitCast(rounded); +} + +pub const Layer = struct { + qkg: Weight, // bf16 [q + kv + q, H]: q rows, k rows, g rows + v: Weight, // bf16 [kv, H] + o: Weight, // bf16 [H, q] + gate: Weight, // NVFP4 [I, H] + up: Weight, + down: Weight, // NVFP4 [H, I] + input_norm: u64, // bf16 [H] + post_attn: u64, + pre_ffn: u64, + post_ffn: u64, + q_norm: u64, // bf16 [128] + k_norm: u64, + blend: f32, // bf16(sigmoid(value_res_lambda)) +}; + +/// Device bytes `deviceBytes` counts beside the small tensors: every projection, the table and the head. +pub fn deviceBytes(c: Config) usize { + const h = c.hidden; + const fp4 = c.intermediate * h / 2 + c.intermediate * h / 16; // words and scale tiles (no padding: gate) + const per_layer = (qkgWidth(c) + c.kvWidth()) * h * 2 + h * c.queryWidth() * 2 + 3 * fp4 + smallLayerBytes(c); + return 2 * c.vocab * h * 2 + c.layers * per_layer + std.mem.alignForward(usize, h * 2, small_align); +} + +/// A layer's norm weights, each at its own 256-byte boundary. +fn smallLayerBytes(c: Config) usize { + return 4 * std.mem.alignForward(usize, c.hidden * 2, small_align) + 2 * std.mem.alignForward(usize, cfg.head_dim * 2, small_align); +} + +pub const Weights = struct { + gpa: std.mem.Allocator, + config: Config, + embed: u64 = 0, // bf16 [V, H] as stored (tf_rc_embed gathers rows) + head: Weight = undefined, // bf16 [V, H] packed + norm: u64 = 0, // bf16 [H] + layers: []Layer = &.{}, + buffers: std.ArrayList(DeviceBuffer) = .empty, + bytes: usize = 0, + + pub fn deinit(w: *Weights) void { + for (w.buffers.items) |*b| b.free(); + w.buffers.deinit(w.gpa); + w.gpa.free(w.layers); + w.* = undefined; + } + + fn upload(w: *Weights, d: *const cuda.Driver, bytes: []const u8) !u64 { + var b = try DeviceBuffer.fromHost(d, bytes); + errdefer b.free(); + try w.buffers.append(w.gpa, b); + w.bytes += bytes.len; + return b.ptr; + } + + /// bf16 matrices one under another ([sum n, k]), packed as one projection for the lane and prompt matmuls. + fn dense(w: *Weights, d: *const cuda.Driver, parts: []const Tensor, k: usize) !Weight { + var n: usize = 0; + for (parts) |p| n += p.bytes.len / 2 / k; + const src = try w.gpa.alloc(u16, n * k); + defer w.gpa.free(src); + var at: usize = 0; + for (parts) |p| { + @memcpy(std.mem.sliceAsBytes(src)[at..][0..p.bytes.len], p.bytes); + at += p.bytes.len; + } + const out = try w.gpa.alloc(u16, cuda.modelopt.padded(n) * k); + defer w.gpa.free(out); + cuda.qmmf.packBf16(out, src, n, k); + return cuda.modelopt.bf16Weight(try w.upload(d, std.mem.sliceAsBytes(out)), n, k); + } + + fn fp4(w: *Weights, d: *const cuda.Driver, p: Fp4) !Weight { + const words = try w.gpa.alloc(u32, cuda.nvfp4.padded(p.n) / 64 * (p.k / 64) * 512); + defer w.gpa.free(words); + const tiles = try w.gpa.alloc(u8, cuda.nvfp4.padded(p.n) * (p.k / 16)); + defer w.gpa.free(tiles); + cuda.nvfp4.packWords(words, p.w.bytes, p.n, p.k); + cuda.nvfp4.packScales(tiles, p.scales.bytes, p.n, p.k); + const wp = try w.upload(d, std.mem.sliceAsBytes(words)); + const tp = try w.upload(d, tiles); + return cuda.nvfp4.weight(wp, tp, p.global, p.n, p.k); + } +}; + +/// The small tensors in one host image, each at a 256-byte boundary; their offsets into it. +const Small = struct { + bytes: std.ArrayList(u8) = .empty, + + fn put(s: *Small, gpa: std.mem.Allocator, t: Tensor) !usize { + const at = std.mem.alignForward(usize, s.bytes.items.len, small_align); + try s.bytes.appendNTimes(gpa, 0, at - s.bytes.items.len); + try s.bytes.appendSlice(gpa, t.bytes); + return at; + } +}; + +/// config.json through the CUDA gate, every tensor mapped and checked, then layer by layer onto the GPU. +pub fn load(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, dir: []const u8) !Weights { + const c = try Config.read(gpa, io, dir); + try gate(&c); + var ck = try Checkpoint.openModel(gpa, io, dir); + defer core.checkpoint.cache.close(&ck); + var m = try map(gpa, &ck, c); + defer m.deinit(); + var w: Weights = .{ .gpa = gpa, .config = c }; + errdefer w.deinit(); + const h = c.hidden; + w.layers = try gpa.alloc(Layer, c.layers); + w.embed = try w.upload(d, m.embed.bytes); + w.head = try w.dense(d, &.{m.head}, h); + var small: Small = .{}; + defer small.bytes.deinit(gpa); + const at_norm = try small.put(gpa, m.norm); + const offsets = try gpa.alloc([6]usize, c.layers); + defer gpa.free(offsets); + for (w.layers, m.layers, offsets) |*l, t, *o| { + l.* = .{ + .qkg = try w.dense(d, &.{ t.q, t.k, t.g }, h), + .v = try w.dense(d, &.{t.v}, h), + .o = try w.dense(d, &.{t.o}, c.queryWidth()), + .gate = try w.fp4(d, t.gate), + .up = try w.fp4(d, t.up), + .down = try w.fp4(d, t.down), + .input_norm = 0, + .post_attn = 0, + .pre_ffn = 0, + .post_ffn = 0, + .q_norm = 0, + .k_norm = 0, + .blend = blendOf(t.lambda), + }; + for ([_]Tensor{ t.input_norm, t.post_attn, t.pre_ffn, t.post_ffn, t.q_norm, t.k_norm }, o) |x, *at| at.* = try small.put(gpa, x); + } + const base = try w.upload(d, small.bytes.items); + w.norm = base + at_norm; + for (w.layers, offsets) |*l, o| { + l.input_norm = base + o[0]; + l.post_attn = base + o[1]; + l.pre_ffn = base + o[2]; + l.post_ffn = base + o[3]; + l.q_norm = base + o[4]; + l.k_norm = base + o[5]; + } + return w; +} + +const testing = std.testing; +const tio = testing.io; + +/// The release's config (the fixture), as its NVFP4 quantization_config sets it. +fn release() !Config { + var c = try cfg.parse(testing.allocator, @embedFile("fixtures/config.json")); + c.fp4_mlp = true; + return c; +} + +test "ricursive cuda gate: the NVFP4 release passes, the BF16 checkpoint and unbuilt shapes are refused" { + var c = try release(); + try gate(&c); + try testing.expectEqual(@as(usize, 21504), qkgWidth(c)); + c.fp4_mlp = false; + try testing.expectError(error.RicursiveCudaNeedsNvfp4, gate(&c)); + c = try release(); + c.kv_heads = 16; // group 5: the prompt kernel pairs query heads + try testing.expectError(error.UnsupportedRicursiveCudaShape, gate(&c)); + c = try release(); + c.hidden = 10240; // past tf_rc_add_norm's 8,192 + try testing.expectError(error.UnsupportedRicursiveCudaShape, gate(&c)); +} + +test "ricursive cuda weights: the release's device bytes, about 57 GB, the planned 53.3 GiB" { + const c = try release(); + const b = deviceBytes(c); + try testing.expect(b > 57_000_000_000 and b < 57_400_000_000); +} + +test "ricursive cuda blendOf: torch's bf16 sigmoid" { + try testing.expectEqual(@as(f32, 0.5), blendOf(0x0000)); + try testing.expectEqual(@as(f32, 0.73046875), blendOf(0x3f80)); + try testing.expectEqual(@as(f32, 0.26953125), blendOf(0xbf80)); +} + +test "ricursive cuda fused q, k and g: one packing of the rows equals each part packed in its place" { + const gpa = testing.allocator; + const k = 128; + const ns = [_]usize{ 256, 128, 256 }; + var total: usize = 0; + for (ns) |n| total += n; + const src = try gpa.alloc(u16, total * k); + defer gpa.free(src); + for (src, 0..) |*v, i| v.* = @truncate(i *% 2654435761); + const whole = try gpa.alloc(u16, total * k); + defer gpa.free(whole); + cuda.qmmf.packBf16(whole, src, total, k); + var at: usize = 0; + const part = try gpa.alloc(u16, 256 * k); + defer gpa.free(part); + for (ns) |n| { + cuda.qmmf.packBf16(part[0 .. n * k], src[at..][0 .. n * k], n, k); + try testing.expectEqualSlices(u16, whole[at..][0 .. n * k], part[0 .. n * k]); + at += n * k; + } +} + +/// A tiny Ripple: 5 layers (the fifth global), 4 query heads over 2 key heads. +fn tiny() Config { + var c: Config = .{ .hidden = 128, .intermediate = 256, .layers = 5, .heads = 4, .kv_heads = 2, .vocab = 128, .eps = 1e-6, .theta = 10000, .window = 16, .embed_scale = 1, .clamp = 20, .context = 64, .fp4_mlp = true }; + c.global[4] = true; + return c; +} + +/// A safetensors file of every tensor `c` names (zeros, each global scale 1). +fn writeFixture(a: std.mem.Allocator, dir: std.Io.Dir, c: Config, input_scales: bool) !void { + const E = struct { name: []const u8, dtype: []const u8, shape: []const usize }; + var list: std.ArrayList(E) = .empty; + const h = c.hidden; + const Add = struct { + fn bf(l: *std.ArrayList(E), al: std.mem.Allocator, name: []const u8, shape: []const usize) !void { + try l.append(al, .{ .name = name, .dtype = "BF16", .shape = try al.dupe(usize, shape) }); + } + fn fp4(l: *std.ArrayList(E), al: std.mem.Allocator, pre: []const u8, n: usize, k: usize, input: bool) !void { + try l.append(al, .{ .name = try std.fmt.allocPrint(al, "{s}.weight", .{pre}), .dtype = "U8", .shape = try al.dupe(usize, &.{ n, k / 2 }) }); + try l.append(al, .{ .name = try std.fmt.allocPrint(al, "{s}.weight_scale", .{pre}), .dtype = "F8_E4M3", .shape = try al.dupe(usize, &.{ n, k / 16 }) }); + try l.append(al, .{ .name = try std.fmt.allocPrint(al, "{s}.weight_scale_2", .{pre}), .dtype = "F32", .shape = &.{} }); + if (input) try l.append(al, .{ .name = try std.fmt.allocPrint(al, "{s}.input_scale", .{pre}), .dtype = "F32", .shape = &.{} }); + } + }; + try Add.bf(&list, a, "model.embed_tokens.weight", &.{ c.vocab, h }); + try Add.bf(&list, a, "lm_head.weight", &.{ c.vocab, h }); + try Add.bf(&list, a, "model.norm.weight", &.{h}); + for (0..c.layers) |i| { + const p = try std.fmt.allocPrint(a, "model.layers.{d}.", .{i}); + for ([_][]const u8{ "input_layernorm", "post_attention_layernorm", "pre_feedforward_layernorm", "post_feedforward_layernorm" }) |n| try Add.bf(&list, a, try std.fmt.allocPrint(a, "{s}{s}.weight", .{ p, n }), &.{h}); + for ([_][]const u8{ "q_proj", "g_proj" }) |n| try Add.bf(&list, a, try std.fmt.allocPrint(a, "{s}self_attn.{s}.weight", .{ p, n }), &.{ c.queryWidth(), h }); + for ([_][]const u8{ "k_proj", "v_proj" }) |n| try Add.bf(&list, a, try std.fmt.allocPrint(a, "{s}self_attn.{s}.weight", .{ p, n }), &.{ c.kvWidth(), h }); + try Add.bf(&list, a, try std.fmt.allocPrint(a, "{s}self_attn.o_proj.weight", .{p}), &.{ h, c.queryWidth() }); + for ([_][]const u8{ "q_norm", "k_norm" }) |n| try Add.bf(&list, a, try std.fmt.allocPrint(a, "{s}self_attn.{s}.weight", .{ p, n }), &.{cfg.head_dim}); + try Add.bf(&list, a, try std.fmt.allocPrint(a, "{s}self_attn.value_res_lambda", .{p}), &.{1}); + for ([_][]const u8{ "gate_proj", "up_proj" }) |n| try Add.fp4(&list, a, try std.fmt.allocPrint(a, "{s}mlp.{s}", .{ p, n }), c.intermediate, h, input_scales); + try Add.fp4(&list, a, try std.fmt.allocPrint(a, "{s}mlp.down_proj", .{p}), h, c.intermediate, input_scales); + } + var header: std.Io.Writer.Allocating = .init(a); + try header.writer.writeAll("{"); + var at: usize = 0; + const sizes = try a.alloc(usize, list.items.len); + for (list.items, sizes, 0..) |e, *size, j| { + var n: usize = if (e.dtype[0] == 'B') 2 else if (e.dtype[0] == 'F' and e.dtype[1] == '3') 4 else 1; + for (e.shape) |x| n *= x; + size.* = n; + try header.writer.print("{s}\"{s}\":{{\"dtype\":\"{s}\",\"shape\":[", .{ if (j == 0) "" else ",", e.name, e.dtype }); + for (e.shape, 0..) |x, i| try header.writer.print("{s}{d}", .{ if (i == 0) "" else ",", x }); + try header.writer.print("],\"data_offsets\":[{d},{d}]}}", .{ at, at + n }); + at += n; + } + try header.writer.writeAll("}"); + const hd = header.written(); + const file = try a.alloc(u8, 8 + hd.len + at); + @memset(file, 0); + std.mem.writeInt(u64, file[0..8], hd.len, .little); + @memcpy(file[8..][0..hd.len], hd); + var off: usize = 0; + for (list.items, sizes) |e, n| { + if (std.mem.endsWith(u8, e.name, "weight_scale_2")) std.mem.writeInt(u32, file[8 + hd.len + off ..][0..4], @bitCast(@as(f32, 1)), .little); + off += n; + } + try dir.writeFile(tio, .{ .sub_path = "model.safetensors", .data = file }); +} + +fn tmpPath(a: std.mem.Allocator, tmp: testing.TmpDir) ![]u8 { + const cwd = try std.process.currentPathAlloc(tio, a); + return std.fmt.allocPrint(a, "{s}/.zig-cache/tmp/{s}", .{ cwd, tmp.sub_path }); +} + +test "ricursive cuda names: every tensor of a tiny NVFP4 release maps, input_scale taken and unread" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const c = tiny(); + for ([_]bool{ true, false }) |input| { + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + try writeFixture(a, tmp.dir, c, input); + var ck = try Checkpoint.openModel(testing.allocator, tio, try tmpPath(a, tmp)); + defer ck.close(); + var m = try map(testing.allocator, &ck, c); + defer m.deinit(); + try testing.expectEqual(@as(usize, 5), m.layers.len); + try testing.expectEqual(@as(f32, 1), m.layers[3].down.global); + try testing.expectEqual(@as(usize, 128), m.layers[0].gate.k); + try testing.expectEqual(@as(f32, 0.5), blendOf(m.layers[4].lambda)); + } +} diff --git a/zig/src/families/ricursive/fixtures/config.json b/zig/src/families/ricursive/fixtures/config.json new file mode 100644 index 0000000000..180aab5d39 --- /dev/null +++ b/zig/src/families/ricursive/fixtures/config.json @@ -0,0 +1,202 @@ +{ + "architectures": [ + "RicursiveForCausalLM" + ], + "auto_map": { + "AutoConfig": "configuration_ricursive.RicursiveConfig", + "AutoModel": "modeling_ricursive.RicursiveModel", + "AutoModelForCausalLM": "modeling_ricursive.RicursiveForCausalLM" + }, + "bos_token_id": 100257, + "embed_scale": 71.55417527999327, + "eos_token_id": [ + 100265, + 100257 + ], + "ffn_clamp_threshold": 20.0, + "gate_full_precision": true, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 5120, + "hybrid_layer_pattern": [ + 1, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 1, + 0 + ], + "intermediate_size": 40960, + "is_hybrid_swa": true, + "layer_types": [ + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention" + ], + "max_position_embeddings": 2097152, + "model_type": "ricursive", + "num_attention_heads": 80, + "num_hidden_layers": 80, + "num_key_value_heads": 8, + "pad_token_id": 100277, + "rms_norm_eps": 1e-06, + "rope_theta": 10000.0, + "sliding_window": 1024, + "tie_word_embeddings": false, + "transformers_version": "4.57.6", + "use_cache": true, + "use_gated_attention": true, + "use_valres": true, + "vocab_size": 100352 +} diff --git a/zig/src/families/ricursive/forward.zig b/zig/src/families/ricursive/forward.zig new file mode 100644 index 0000000000..dc77e1370a --- /dev/null +++ b/zig/src/families/ricursive/forward.zig @@ -0,0 +1,327 @@ +//! Ripple'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; + +/// The glue kernels' constants (glue.metal's rc_args). +const Args = extern struct { rows: u32 = 0, width: u32 = 0, flags: u32 = 0, eps: f32 = 0, scale: f32 = 0, blend: f32 = 0 }; +const add_flag: u32 = 1; // rc_add_norm: the sublayer's output normed and added first +const lerp_flag: u32 = 2; // rc_add_norm: v's input blended from the embeddings too +/// q and k's norm and RoPE, into the cache (glue.metal's rc_prep_args). +const PArgs = extern struct { rows: u32, heads: u32, kv_heads: u32, pos: u32, stride: u32, ring: u32, rope: u32, eps: f32 }; +/// The NVFP4 row kernels' and dequantization's constants (fp4.metal's rc_fp4_args). +const Fp4Args = extern struct { rows: u32, n: u32, k: u32, ldy: u32, global: f32 }; +/// The bf16 row projection's constants (llama's ll_mv_args). +const MArgs = extern struct { rows: u32, n: u32, k: u32, ldy: u32 }; +/// The decode attention's constants (rc_vec_args); its first eight are llama's ll_vec_args, which the merge reads. +const VArgs = extern struct { rows: u32, pos: u32, heads: u32, groups: u32, blocks: u32, chunk: u32, first: u32, scale: f32, window: u32 = 0, ring: u32 = 0 }; +/// A global layer's window: every key before. +const every_key = std.math.maxInt(i32); + +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 }; + +const Enc = struct { + e: mtl.ComputeEncoder, + m: *const Model, + s: *st.Scratch, + + fn glue(x: Enc, p: mtl.Pipeline, args: anytype, bufs: []const B, grid: [3]usize, group: [3]usize, groups: bool) void { + x.e.setPipeline(p); + 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, at byte `out_off`) = x W^T: the GEMM a prompt chunk, the row kernels a window. + fn dense(x: Enc, d: wts.Dense, in: B, 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 v: @import("model.zig").Llama = 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.ll(v)); + e.setBuffer(in.buf, in.off, 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 tn = (d.n + 127) / 128; + const tm = (n_rows + 63) / 64; + e.setPipeline(x.m.gemm[@intFromBool(n_rows % 64 == 0)][@intFromBool(d.n % 128 == 0)]); + e.setBuffer(in.buf, in.off, 0); + e.setBuffer(d.buf, d.off, 1); + var p: [16]i32 = @splat(0); + const vals = [_]usize{ n_rows, 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)); + } + + /// An MLP projection: bf16, or NVFP4 rows; a prompt chunk dequantizes exactly, then the GEMM applies the global. + fn proj(x: Enc, p: wts.Proj, in: B, out: mtl.Buffer, out_off: usize, n_rows: usize, ldy: usize) void { + const d = switch (p) { + .bf16 => |d| return x.dense(d, in, out, out_off, n_rows, ldy), + .fp4 => |d| d, + }; + const f = &(x.m.fp4 orelse unreachable); + const e = x.e; + const fa = Fp4Args{ .rows = @intCast(n_rows), .n = @intCast(d.n), .k = @intCast(d.k), .ldy = @intCast(ldy), .global = d.global }; + if (n_rows <= st.batch_rows) { + const v: usize = if (n_rows <= 1) 0 else if (n_rows <= 2) 1 else if (n_rows <= 4) 2 else if (n_rows <= 8) 3 else if (n_rows <= 16) 4 else 5; + e.setPipeline(f.rows[v]); + e.setBuffer(in.buf, in.off, 0); + e.setBuffer(d.w, d.w_off, 1); + e.setBuffer(d.s, d.s_off, 2); + e.setValue(fa, 3); + e.setBuffer(out, out_off, 4); + e.dispatchGroups(mtl.Size.of((d.n + 7) / 8, 1, 1), mtl.Size.of(256, 1, 1)); + return; + } + e.setPipeline(f.rows[6]); + e.setBuffer(d.w, d.w_off, 0); + e.setBuffer(d.s, d.s_off, 1); + e.setValue(fa, 2); + e.setBuffer(x.s.wq, 0, 3); + e.dispatchThreads(mtl.Size.of(d.k / 16, d.n, 1), mtl.Size.of(64, 4, 1)); + const tn = (d.n + 127) / 128; + const tm = (n_rows + 63) / 64; + e.setPipeline(f.gemm[@intFromBool(n_rows % 64 == 0)][@intFromBool(d.n % 128 == 0)]); + e.setBuffer(in.buf, in.off, 0); + e.setBuffer(x.s.wq, 0, 1); + var pa: [16]i32 = @splat(0); + const vals = [_]usize{ n_rows, 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| pa[i] = @intCast(v); + e.setBytes(std.mem.asBytes(&pa), 2); + e.setBuffer(out, out_off, 3); + e.setValue(d.global, 4); + e.dispatchThreads(mtl.Size.of(@as(usize, tn << 2) * 32, (tm + 3) / 4 * 4, 2), mtl.Size.of(32, 4, 2)); + } + + /// rc_add_norm over n rows: with `post`, R += norm(y, post); X = norm(R, next); with `blend`, VX too. + fn addNorm(x: Enc, n: usize, post: ?wts.Tensor, next: wts.Tensor, blend: ?f32) void { + const s = x.s; + const c = &x.m.config; + const flags = (if (post != null) add_flag else 0) | (if (blend != null) lerp_flag else 0); + const a = Args{ .rows = @intCast(n), .width = @intCast(c.hidden), .flags = flags, .eps = c.eps, .blend = blend orelse 0 }; + const p = if (post) |w| t(w) else t(next); // unread without the add + x.glue(x.m.pipe(.add_norm), a, &.{ b(s.resid), b(s.y), p, t(next), b(s.v0), A, b(s.x), b(s.vx) }, .{ n, 1, 1 }, .{ 256, 1, 1 }, 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-ricursive-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.RicursiveRowsExceeded; + var total: usize = 0; + for (segments) |seg| { + if (seg.rows == 0 or seg.cache.len + seg.rows > seg.cache.capacity) return error.RicursiveContextFull; + if (segments.len > 1 and seg.rows > st.window_rows) return error.RicursiveRowsExceeded; + total += seg.rows; + } + if (total != n) return error.InvalidRicursiveSegments; + if (out.ids == null) { + for (ids) |id| if (id >= c.vocab) return error.InvalidRicursiveToken; + @memcpy(s.ids.slice(u32, n), ids); + } + if (out.draw != null and (out.all or segments.len != 1)) return error.InvalidRicursiveSegments; + const h = c.hidden; + 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; + const ea = Args{ .rows = @intCast(n), .width = @intCast(h), .scale = c.embed_scale }; + x.glue(m.pipe(.embed), ea, &.{ out.ids orelse b(s.ids), bo(W.embed.buf, W.embed.off), A, b(s.resid), b(s.v0) }, .{ h, n, 1 }, .{ 256, 1, 1 }, false); + x.addNorm(n, null, W.layers[0].input_norm, W.layers[0].attn.blend); + for (W.layers, 0..) |l, i| try layer(x, segments, l, i, n, if (i + 1 < W.layers.len) W.layers[i + 1] else null); + if (out.all) { + x.dense(W.head, b(s.x), s.logits_all, 0, n, c.vocab); + } else { + x.dense(W.head, bo(s.x, (n - 1) * h * 2), 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("ricursive forward: {s}", .{msg}); + return error.GpuFailed; + } + if (out.timing) |tm| tm.gpu += cb.gpuSeconds(); +} + +/// The head alone on `rows` final-normed rows in s.x, into s.logits_all (a prompt chunk's every row, for the check). +pub fn head(m: *const Model, s: *st.Scratch, rows: usize) !void { + if (rows == 0 or rows > st.batch_rows or rows > s.rows) return error.RicursiveRowsExceeded; + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const cb = m.queue.commandBuffer(); + const x = Enc{ .e = cb.compute(.serial), .m = m, .s = s }; + x.dense(m.weights.head, b(s.x), s.logits_all, 0, rows, m.config.vocab); + x.e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.log.err("ricursive head: {s}", .{msg}); + return error.GpuFailed; + } +} + +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.ll(.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 and s.vx its inputs); its last add norms the residual for `next`. +fn layer(x: Enc, segments: []const Segment, l: wts.Layer, i: usize, n: usize, next: ?wts.Layer) !void { + const s = x.s; + const c = &x.m.config; + const h = c.hidden; + const qkvw = c.qkvWidth(); + x.dense(l.attn.q, b(s.x), s.q, 0, n, qkvw); + x.dense(l.attn.k, b(s.x), s.q, c.queryWidth() * 2, n, qkvw); + x.dense(l.attn.v, b(s.vx), s.q, (c.queryWidth() + c.kvWidth()) * 2, n, qkvw); + x.dense(l.attn.g, b(s.x), s.g, 0, n, c.queryWidth()); + var base: usize = 0; + for (segments) |seg| { + try attention(x, seg.cache, l, i, base, seg.rows); + base += seg.rows; + } + const ga = Args{ .rows = @intCast(n), .width = @intCast(c.queryWidth()) }; + x.glue(x.m.pipe(.gate), ga, &.{ b(s.att), b(s.g), A }, .{ c.queryWidth(), n, 1 }, .{ 256, 1, 1 }, false); + x.dense(l.attn.o, b(s.att), s.y, 0, n, h); + x.addNorm(n, l.post_attn_norm, l.pre_ffn_norm, null); + x.proj(l.mlp.gate, b(s.x), s.gu, 0, n, 2 * c.intermediate); + x.proj(l.mlp.up, b(s.x), s.gu, c.intermediate * 2, n, 2 * c.intermediate); + const fa = Args{ .rows = @intCast(n), .width = @intCast(c.intermediate), .scale = c.clamp }; + x.glue(x.m.pipe(.silu_clamp), fa, &.{ b(s.gu), A, b(s.act) }, .{ c.intermediate, n, 1 }, .{ 256, 1, 1 }, false); + x.proj(l.mlp.down, b(s.act), s.y, 0, n, h); + if (next) |nl| { + x.addNorm(n, l.post_ffn_norm, nl.input_norm, nl.attn.blend); + } else x.addNorm(n, l.post_ffn_norm, x.m.weights.norm, null); +} + +/// One stream's scratch rows base..: q and k normed (rotated on sliding layers), k and v into its cache, attention. +fn attention(x: Enc, cache: *st.Cache, l: wts.Layer, i: usize, base: usize, n: usize) !void { + const s = x.s; + const c = &x.m.config; + const start = cache.len; + const global = c.global[i]; + const ring = cache.layer_slots[i]; + 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 pa = PArgs{ .rows = @intCast(n), .heads = @intCast(c.heads), .kv_heads = @intCast(c.kv_heads), .pos = @intCast(start), .stride = @intCast(c.qkvWidth()), .ring = @intCast(ring), .rope = @intFromBool(!global), .eps = c.eps }; + const keys = bo(cache.kv, cache.keys(i)); + const values = bo(cache.kv, cache.values(i)); + x.glue(x.m.pipe(.prep), pa, &.{ bo(s.q, base * c.qkvWidth() * 2), b(x.m.inv), t(l.attn.q_norm), t(l.attn.k_norm), A, qn, keys, values }, .{ c.heads + c.kv_heads, n, 1 }, .{ 32, 1, 1 }, true); + const e = x.e; + const window: usize = if (global) every_key else c.window; + if (n > st.window_rows) { // a prompt chunk: 32 rows of a head a tile, over every key in the rows' windows + e.setPipeline(x.m.prompt_attn); + for ([_]B{ qn, keys, values, o }, 0..) |r, slot| e.setBuffer(r.buf, r.off, slot); + e.setValue([4]i32{ @intCast(n), @intCast(start), @intCast(window), @intCast(ring) }, 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 = if (global) st.chunk(cache.capacity) else st.ring_chunk; + const low = if (start + 1 > window) start + 1 - window else 0; // the first row's first key + const first = low / chunk; + const blocks = (start + n - 1) / chunk + 1 - first; + if (blocks > st.max_blocks) return error.RicursiveTooManyBlocks; + const va = VArgs{ .rows = @intCast(n), .pos = @intCast(start), .heads = @intCast(c.heads), .groups = @intCast(c.kv_heads), .blocks = @intCast(blocks), .chunk = @intCast(chunk), .first = @intCast(first), .scale = c.scale(), .window = @intCast(window), .ring = @intCast(ring) }; + 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.ll(.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/ricursive/model.zig b/zig/src/families/ricursive/model.zig new file mode 100644 index 0000000000..4109af6e93 --- /dev/null +++ b/zig/src/families/ricursive/model.zig @@ -0,0 +1,193 @@ +//! A native Ripple 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"); + +/// The family's own glue (glue.metal in kernels/metal/ricursive). +pub const glue_names = [_][:0]const u8{ "rc_embed", "rc_add_norm", "rc_prep", "rc_gate", "rc_silu_clamp", "rc_decode_vec" }; +pub const Glue = enum { embed, add_norm, prep, gate, silu_clamp, decode_vec }; +/// Llama's kernels it runs as they are: bf16 row projections, the decode attention's merge and the greedy draw. +pub const llama_names = [_][:0]const u8{ "ll_matvec_1", "ll_matvec_2", "ll_matvec_4", "ll_matvec_8", "ll_matvec_16", "ll_matvec_32", "ll_merge_vec", "ll_argmax" }; +pub const Llama = enum { matvec_1, matvec_2, matvec_4, matvec_8, matvec_16, matvec_32, merge_vec, argmax }; + +/// The NVFP4 MLP's row kernels (1 to 32 rows) and its prompt dequantization (fp4.metal). +pub const fp4_names = [_][:0]const u8{ "rc_fp4_matvec_1", "rc_fp4_matvec_2", "rc_fp4_matvec_4", "rc_fp4_matvec_8", "rc_fp4_matvec_16", "rc_fp4_matvec_32", "rc_fp4_dequant" }; + +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, + llama: [llama_names.len]mtl.Pipeline, + /// Prompt attention on the tensor units over a window and a ring; decode attention is rc_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 (llama's). + gemm: [2][2]mtl.Pipeline, + /// The NVFP4 release's MLP kernels; null for the bf16 release. + fp4: ?Fp4 = null, + /// RoPE's inverse frequencies: fp32 [head_dim / 2]. + inv: mtl.Buffer, + + pub fn load(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !*Model { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const m = try gpa.create(Model); + errdefer gpa.destroy(m); + m.gpa = gpa; + m.config = try cfg.Config.read(gpa, io, dir); + m.device = try mtl.Device.init(); + errdefer m.device.deinit(); + // prompt chunks' GEMM and attention run on Metal 4 tensor units: an M5 or later + if (!m.device.tensorUnits()) return error.RicursiveNeedsTensorUnits; + m.queue = try m.device.queue(); + errdefer m.queue.deinit(); + m.checkpoint = ck.Checkpoint.init(gpa); + errdefer m.checkpoint.deinit(); + const files = try shards.shardFiles(gpa, io, dir); + defer shards.freeShardFiles(gpa, files); + // the release's 2-byte value_res_lambda scalars leave later tensors 2-14 bytes off 8: each placed 256-aligned + for (files) |path| try m.checkpoint.addFileSplit(m.device, path, "", null, keepAll); + m.weights = try weights.load(gpa, &m.checkpoint, &m.config); + errdefer m.weights.deinit(gpa); + try pipelines(m.device, gpa, "#define RC_DH {d}\n{s}", ks.ricursive_glue, &glue_names, &m.glue); + errdefer for (m.glue) |p| p.deinit(); + try pipelines(m.device, gpa, "#define LL_DH {d}\n{s}", ks.llama_glue, &llama_names, &m.llama); + errdefer for (m.llama) |p| p.deinit(); + { + const src = try attentionSource(m.device, gpa, &m.config); + 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, "rc_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.fp4 = if (m.config.fp4_mlp) try Fp4.init(m.device, gpa) else null; + errdefer if (m.fp4) |*f| f.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 ll(m: *const Model, g: Llama) mtl.Pipeline { + return m.llama[@backingInt(g)]; + } + + pub fn deinit(m: *Model) void { + m.inv.deinit(); + if (m.fp4) |*f| f.deinit(); + for (m.gemm) |row| for (row) |p| p.deinit(); + m.sample_full.deinit(); + m.prompt_attn.deinit(); + for (m.llama) |p| p.deinit(); + for (m.glue) |p| p.deinit(); + m.weights.deinit(m.gpa); + m.checkpoint.deinit(); + m.queue.deinit(); + m.device.deinit(); + m.gpa.destroy(m); + } +}; + +/// The NVFP4 pipelines: the row kernels and dequantization (precise math) and the tensor-unit GEMM with its global. +pub const Fp4 = struct { + rows: [fp4_names.len]mtl.Pipeline, + gemm: [2][2]mtl.Pipeline, + + fn init(device: mtl.Device, gpa: std.mem.Allocator) !Fp4 { + var f: Fp4 = undefined; + { + const lib = try mtl.Library.fromSource(device, ks.ricursive_fp4, .{ .language = mtl.CompileOptions.mlx().language, .math_functions = .precise }); + defer lib.deinit(); + var made: usize = 0; + errdefer for (f.rows[0..made]) |p| p.deinit(); + for (fp4_names, 0..) |n, i| { + f.rows[i] = try mtl.Pipeline.init(device, lib, n, false); + made += 1; + } + } + errdefer for (f.rows) |p| p.deinit(); + const body = try frags.source(device, gpa, ks.ricursive_fp4_gemm); + defer gpa.free(body); + const lib = try mtl.Library.fromSource(device, body, mtl.CompileOptions.mlx()); + defer lib.deinit(); + var made: usize = 0; + errdefer for (0..made) |i| f.gemm[i / 2][i % 2].deinit(); + for (0..4) |i| { + var nm: [24]u8 = undefined; + const full = try std.fmt.bufPrintSentinel(&nm, "rc_fp4_gemm_{d}{d}", .{ i / 2, i % 2 }, 0); + f.gemm[i / 2][i % 2] = try mtl.Pipeline.init(device, lib, full, false); + made += 1; + } + return f; + } + + fn deinit(f: *Fp4) void { + for (f.gemm) |row| for (row) |p| p.deinit(); + for (f.rows) |p| p.deinit(); + } +}; + +/// No tensor stays in its file: every one is read into the shard's buffer at its own aligned place. +fn keepAll(_: []const u8, _: usize) bool { + return false; +} + +/// The glue text with its head-size define, as the loader compiles it (precise math functions). +pub fn glueSource(a: std.mem.Allocator, comptime fmt: []const u8, body: []const u8) ![]u8 { + return std.fmt.allocPrint(a, fmt, .{ cfg.head_dim, body }); +} + +/// The prompt attention's text for a geometry (nax.h inlined for this GPU). +pub fn attentionSource(device: mtl.Device, a: std.mem.Allocator, c: *const cfg.Config) ![]u8 { + const body = try frags.source(device, a, ks.ricursive_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, c.heads, c.group(), c.kv_heads, c.scale(), body }); +} + +fn pipelines(device: mtl.Device, gpa: std.mem.Allocator, comptime fmt: []const u8, body: []const u8, names: []const [:0]const u8, out: []mtl.Pipeline) !void { + const src = try glueSource(gpa, fmt, body); + defer gpa.free(src); + const lib = try mtl.Library.fromSource(device, src, .{ .language = mtl.CompileOptions.mlx().language, .math_functions = .precise }); + defer lib.deinit(); + var made: usize = 0; + errdefer for (out[0..made]) |p| p.deinit(); + for (names, 0..) |n, i| { + out[i] = try mtl.Pipeline.init(device, lib, n, false); + made += 1; + } +} diff --git a/zig/src/families/ricursive/ricursive.zig b/zig/src/families/ricursive/ricursive.zig new file mode 100644 index 0000000000..2f00b7e37a --- /dev/null +++ b/zig/src/families/ricursive/ricursive.zig @@ -0,0 +1,19 @@ +//! Native Metal Ripple (ricursive): sliding and global gated attention, value residuals, PeriNorm layers. +pub const config = @import("config.zig"); +pub const weights = @import("weights.zig"); +pub const Model = @import("model.zig").Model; +pub const state = @import("state.zig"); +pub const forward = @import("forward.zig"); +pub const backend = @import("backend.zig"); +pub const budget = @import("budget.zig"); +pub const snapshot = @import("snapshot.zig"); +pub const contracts = @import("contracts.zig"); + +test { + _ = config; + _ = weights; + _ = state; + _ = budget; + _ = snapshot; + _ = contracts; +} diff --git a/zig/src/families/ricursive/snapshot.zig b/zig/src/families/ricursive/snapshot.zig new file mode 100644 index 0000000000..04bc31040d --- /dev/null +++ b/zig/src/families/ricursive/snapshot.zig @@ -0,0 +1,85 @@ +//! A stream's prompt state kept for reuse on Metal: each layer's held keys and values (a sliding layer's ring). +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; + } +}; + +/// Slots of layer `l` a state after `at` tokens holds: every position, or a sliding layer's full ring past it. +fn held(c: *const cfg.Config, l: usize, at: usize) usize { + return if (c.global[l]) at else @min(at, st.ring(c, st.max_capacity)); +} + +/// Bytes a state after `at` prompt tokens takes. +pub fn bytes(c: *const cfg.Config, at: usize) u64 { + var total: u64 = 0; + for (0..c.layers) |l| total += @as(u64, held(c, l, at)) * st.kvBytes(c); + return total; +} + +/// 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, c, cache, cache.len, 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: a ring past its size is a full ring on both. +pub fn restore(queue: mtl.Queue, c: *const cfg.Config, snap: *const Snap, cache: *st.Cache) !void { + if (snap.freed or snap.at >= cache.capacity) return error.RicursiveSnapshotGone; + for (0..c.layers) |l| if (held(c, l, snap.at) > cache.layer_slots[l]) return error.RicursiveSnapshotGone; + try copy(queue, c, cache, snap.at, snap.buf, false); + cache.len = snap.at; +} + +/// 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, c: *const cfg.Config, cache: *const st.Cache, at: usize, 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 off: usize = 0; + for (0..cache.layers) |l| { + const len = held(c, l, at) * row; + if (len == 0) continue; + for ([_]usize{ cache.keys(l), cache.values(l) }) |from| { + if (out) e.copy(cache.kv, from, buf, off, len) else e.copy(buf, off, cache.kv, from, len); + off += len; + } + } + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.log.err("ricursive prompt state copy: {s}", .{msg}); + return error.GpuFailed; + } +} + +test "ricursive snapshot: global layers keep every position, sliding layers at most their ring" { + 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, 80 * 1000 * 4096), bytes(&c, 1000)); + try std.testing.expectEqual(@as(u64, (16 * 100000 + 64 * 3072) * 4096), bytes(&c, 100000)); + try std.testing.expectEqual(st.cacheBytes(&c, 262144) - c.vocab * 2, bytes(&c, 262144)); +} diff --git a/zig/src/families/ricursive/state.zig b/zig/src/families/ricursive/state.zig new file mode 100644 index 0000000000..357656a838 --- /dev/null +++ b/zig/src/families/ricursive/state.zig @@ -0,0 +1,170 @@ +//! Ripple on Metal: a stream's caches (global layers hold every position, sliding layers a ring) and the 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; +/// Keys a sliding layer's decode attention block takes: its window spans a handful of blocks. +pub const ring_chunk = 128; + +/// Slots of a global 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); +} + +/// Slots of a sliding layer: a chunk's rows and the window behind its first row, in whole 128-key blocks. +pub fn ring(c: *const cfg.Config, capacity: usize) usize { + return @min(slots(capacity), std.mem.alignForward(usize, c.window - 1 + prompt_rows, ring_chunk)); +} + +/// Keys a global layer's decode attention block takes for a stream of `capacity`: at most 512 blocks, 32 to 2,048 keys. +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; +} + +/// Slots layer `l` holds for a stream of `capacity`. +pub fn layerSlots(c: *const cfg.Config, l: usize, capacity: usize) usize { + return if (c.global[l]) slots(capacity) else ring(c, capacity); +} + +/// Bytes a stream of `capacity` positions holds in caches, its logits row included. +pub fn cacheBytes(c: *const cfg.Config, capacity: usize) usize { + var total: usize = c.vocab * 2; + for (0..c.layers) |l| total += layerSlots(c, l, capacity) * kvBytes(c); + return total; +} + +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), + layer_slots: [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; + out.layer_slots[i] = layerSlots(c, i, capacity); + kv += out.layer_slots[i] * 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] + self.layer_slots[l] * 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", "v0", "x", "vx", "q", "g", "qn", "att", "y", "gu", "act", "logits_all", "pm", "pl", "po", "wq" }; + +/// 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 hb = r * c.hidden * 2; + const qw = r * c.queryWidth() * 2; + const partials = @as(usize, @min(r, window_rows)) * c.heads * max_blocks * 4; // a segment's decode attention partials + return .{ + r * 4, hb, hb, hb, hb, + r * c.qkvWidth() * 2, qw, qw, qw, hb, + 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, if (prompt and c.fp4_mlp) c.intermediate * c.hidden * 2 else 16, + }; +} + +/// 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 release's residual + v0: mtl.Buffer, // the rows' scaled embeddings: the value residual's source in every layer + x: mtl.Buffer, // bf16 normed rows + vx: mtl.Buffer, // v's input: lerp(v0, x, the layer's blend) + q: mtl.Buffer, // the projections' rows: q | k | v + g: mtl.Buffer, // the gate projection's rows + qn: mtl.Buffer, // q normed (and rotated on sliding layers) + att: mtl.Buffer, // attention's output, then gated + 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, + wq: mtl.Buffer, // a prompt chunk's NVFP4 projection as bf16 (its global left to the GEMM) + 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 "ricursive state: global layers hold every position, sliding layers a ring of 3,072" { + const c = try cfg.parse(std.testing.allocator, @embedFile("fixtures/config.json")); + try std.testing.expectEqual(@as(usize, 3072), ring(&c, 262144)); + try std.testing.expectEqual(@as(usize, 128), ring(&c, 100)); // a short stream: its own whole blocks + try std.testing.expectEqual(@as(usize, 4096), kvBytes(&c)); + // 16 global layers of 262,144 positions, 64 sliding layers of 3,072, 4 KiB a position a layer, and a logits row + try std.testing.expectEqual(@as(usize, 100352 * 2 + (16 * 262144 + 64 * 3072) * 4096), cacheBytes(&c, 262144)); + try std.testing.expect(ring(&c, max_capacity) % ring_chunk == 0 and ring(&c, max_capacity) >= c.window - 1 + prompt_rows); + try std.testing.expect(std.math.divCeil(usize, max_capacity, chunk(max_capacity)) catch unreachable <= 512); + try std.testing.expect(scratchBytes(&c, prompt_rows) > scratchBytes(&c, batch_rows)); +} diff --git a/zig/src/families/ricursive/weights.zig b/zig/src/families/ricursive/weights.zig new file mode 100644 index 0000000000..636acf6f3b --- /dev/null +++ b/zig/src/families/ricursive/weights.zig @@ -0,0 +1,156 @@ +//! Ripple's checkpoint on Metal: tensors as stored (bf16, or the NVFP4 release's MLPs), nothing converted. +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, + g: Dense, // the output gate's projection + o: Dense, + q_norm: Tensor, + k_norm: Tensor, + /// sigmoid(value_res_lambda) rounded to bf16, as the release blends the layer-0 embeddings into v's input. + blend: f32, +}; +/// An NVFP4 [n, k] matrix as ModelOpt stores it: nibbles [n, k / 2], E4M3 scales [n, k / 16] and the fp32 global. +pub const Fp4 = struct { w: mtl.Buffer, w_off: usize, s: mtl.Buffer, s_off: usize, n: usize, k: usize, global: f32 }; +/// An MLP projection: the release's bf16 rows, or the publisher's NVFP4. +pub const Proj = union(enum) { bf16: Dense, fp4: Fp4 }; +pub const Mlp = struct { gate: Proj, up: Proj, down: Proj }; +pub const Layer = struct { + input_norm: Tensor, + post_attn_norm: Tensor, + pre_ffn_norm: Tensor, + post_ffn_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); + } +}; + +/// `t` is `dtype` with exactly `dims`, its bytes those dims' extent. +fn require(t: Tensor, dtype: @TypeOf(t.dtype), dims: []const usize) error{UnexpectedRicursiveTensor}!void { + if (t.dtype != dtype or t.rank != dims.len or t.rank > t.shape.len) return error.UnexpectedRicursiveTensor; + if (!std.mem.eql(usize, t.shape[0..t.rank], dims)) return error.UnexpectedRicursiveTensor; + var bytes: usize = switch (dtype) { + .bf16 => 2, + .f32 => 4, + else => 1, + }; + for (dims) |n| bytes = std.math.mul(usize, bytes, n) catch return error.UnexpectedRicursiveTensor; + if (t.bytes != bytes) return error.UnexpectedRicursiveTensor; +} + +/// The bf16 tensor `prefix ++ suffix` of exactly `dims`. +fn get(c: *const ck.Checkpoint, prefix: []const u8, suffix: []const u8, dims: []const usize) !Tensor { + return getAs(c, prefix, suffix, .bf16, dims); +} + +/// The `dtype` tensor `prefix ++ suffix` of exactly `dims`. +fn getAs(c: *const ck.Checkpoint, prefix: []const u8, suffix: []const u8, dtype: anytype, dims: []const usize) !Tensor { + var buf: [256]u8 = undefined; + const full = try std.fmt.bufPrint(&buf, "{s}{s}", .{ prefix, suffix }); + const t = try c.get(full); + require(t, dtype, dims) catch |err| { + std.log.err("ricursive: {s} is not {s}{any}", .{ full, @tagName(dtype), dims }); + return err; + }; + return t; +} + +/// The MLP projection `prefix ++ name`: bf16 rows, or NVFP4 (the row kernels read 8-byte blocks of 16 nibbles). +fn proj(c: *const ck.Checkpoint, f: *const cfg.Config, prefix: []const u8, comptime name: []const u8, n: usize, k: usize) !Proj { + if (!f.fp4_mlp) return .{ .bf16 = try dense(c, prefix, name ++ ".weight", n, k) }; + if (k % 256 != 0) return error.UnexpectedRicursiveTensor; + const w = try getAs(c, prefix, name ++ ".weight", .u8, &.{ n, k / 2 }); + const s = try getAs(c, prefix, name ++ ".weight_scale", .f8_e4m3, &.{ n, k / 16 }); + const g = try getAs(c, prefix, name ++ ".weight_scale_2", .f32, &.{}); + if (w.offset % 8 != 0) return error.UnexpectedRicursiveTensor; + const global = g.host(f32)[0]; + if (!std.math.isFinite(global) or global <= 0) return error.UnexpectedRicursiveTensor; + return .{ .fp4 = .{ .w = w.buffer, .w_off = w.offset, .s = s.buffer, .s_off = s.offset, .n = n, .k = k, .global = global } }; +} + +/// 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 { + const w = try get(c, prefix, suffix, &.{ n, k }); + if (k % 256 != 0 or w.offset % 8 != 0) { + std.log.err("ricursive: {s}{s} at byte {d}: the row kernels need 8-byte rows of whole 256s", .{ prefix, suffix, w.offset }); + return error.UnexpectedRicursiveTensor; + } + return .{ .buf = w.buffer, .off = w.offset, .n = n, .k = k }; +} + +/// torch's bf16 sigmoid of a bf16 scalar: fp32 1 / (1 + exp(-x)), rounded to bf16 (to nearest, ties to even). +pub fn blendOf(lambda: u16) f32 { + const x: f32 = @bitCast(@as(u32, lambda) << 16); + const s: f32 = 1.0 / (1.0 + @exp(-x)); + const bits: u32 = @bitCast(s); + const rounded = (bits + 0x7fff + ((bits >> 16) & 1)) & 0xffff0000; + return @bitCast(rounded); +} + +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 get(c, "model.", "norm.weight", &.{h}), + .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 std.fmt.bufPrint(&ab, "{s}self_attn.", .{p}); + var mb: [96]u8 = undefined; + const m = try std.fmt.bufPrint(&mb, "{s}mlp.", .{p}); + const lambda = try get(c, a, "value_res_lambda", &.{1}); + l.* = .{ + .input_norm = try get(c, p, "input_layernorm.weight", &.{h}), + .post_attn_norm = try get(c, p, "post_attention_layernorm.weight", &.{h}), + .pre_ffn_norm = try get(c, p, "pre_feedforward_layernorm.weight", &.{h}), + .post_ffn_norm = try get(c, p, "post_feedforward_layernorm.weight", &.{h}), + .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), + .g = try dense(c, a, "g_proj.weight", f.queryWidth(), h), + .o = try dense(c, a, "o_proj.weight", h, f.queryWidth()), + .q_norm = try get(c, a, "q_norm.weight", &.{cfg.head_dim}), + .k_norm = try get(c, a, "k_norm.weight", &.{cfg.head_dim}), + .blend = blendOf(lambda.host(u16)[0]), + }, + .mlp = .{ + .gate = try proj(c, f, m, "gate_proj", f.intermediate, h), + .up = try proj(c, f, m, "up_proj", f.intermediate, h), + .down = try proj(c, f, m, "down_proj", h, f.intermediate), + }, + }; + } + return w; +} + +test "ricursive weights: the blend is torch's bf16 sigmoid of the stored lambda" { + try std.testing.expectEqual(@as(f32, 0.5), blendOf(0x0000)); // lambda 0: one half + try std.testing.expectEqual(@as(f32, 0.73046875), blendOf(0x3f80)); // lambda 1: 0.7310586 rounds to 0.73046875 + try std.testing.expectEqual(@as(f32, 0.26953125), blendOf(0xbf80)); // lambda -1: 0.2689414 rounds to 0.26953125 +} diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index e7e02669cd..fc23ffe974 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -30,7 +30,7 @@ const vision_offer = @import("cuda_vision_offer.zig"); const drafted = @import("cuda_drafted.zig"); const kv_pages = @import("cuda_kv_pool.zig"); /// CUDA families provide metadata, open their lane backend and explain their refusals. -const registry = .{ nemotron.native, nemotron.omni, glm.native, minimax.native, qwen35.native, qwen35.native_moe, flashnext.native, kolibri1.native, gemma4.native, gemma4.native31, spark25.native, dory.native }; +const registry = .{ nemotron.native, nemotron.omni, glm.native, minimax.native, qwen35.native, qwen35.native_moe, flashnext.native, kolibri1.native, gemma4.native, gemma4.native31, spark25.native, dory.native, @import("ricursive_cuda").native }; /// --record opens here (engines.zig refuses it for a backend without this). pub const records = true; pub const backends: []const []const u8 = &.{"cuda"}; @@ -494,6 +494,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O .slide_conversations = conversations, }, .{ .ctx = loaded.ctx, .text = F.explain }); opened = true; + if (@hasDecl(F, "thinking")) F.thinking(&h.host.info_); // a family's own think block rules kv_pages.warnCapacity(cache_capacity, o.keep == 0 or o.prompt_cache_gib == 0, if (paged) |p| p.taken else 0, std.log); if (@hasField(@TypeOf(loaded), "follow")) if (loaded.follow) |f| { h.follow_fn = f; @@ -513,7 +514,6 @@ test { _ = snaps; _ = prompt_pages; _ = cuda_device; // the device and variable tests - _ = windows; } test "CUDA host preserves prefill_step, first-before-draft and chunk-invariance capability" { diff --git a/zig/src/native/metal.zig b/zig/src/native/metal.zig index f702f70523..ab1a8abcaa 100644 --- a/zig/src/native/metal.zig +++ b/zig/src/native/metal.zig @@ -56,6 +56,7 @@ pub const families: []const api.Family = &.{ .{ .model_type = "spark2_5", .formats = &.{ "bf16", "mlx-q2g64", "mlx-q3g64", "mlx-q4g64", "mlx-q5g64", "mlx-q6g64", "mlx-q8g64" } }, .{ .model_type = "minimax_m3_vl", .formats = &.{ "mlx-q3g64", "mlx-q8g64" } }, .{ .model_type = "llama", .formats = &.{ "bf16", "compressed-tensors-int4-g32" } }, + .{ .model_type = "ricursive", .formats = &.{ "bf16", "modelopt-nvfp4-mlp" } }, }; /// Gemma 4 31B (gemma4) and 12B (gemma4_unified, its text model): one native family, gemma4.zig's host. @@ -194,6 +195,7 @@ pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]c 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, "minimax_m3_vl")) return @import("minimax_m3.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, "ricursive")) return @import("ricursive.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}); @@ -322,14 +324,14 @@ fn refusal(a: Allocator, o: api.Open, dense: bool) !?[]const u8 { if (o.prompt_int8) return prompt_int8_refusal; if (!learners.supports(t)) if (try o.slideRefusal(a, "Macs")) |why| return why; if (o.drafter != null and !dense and !isGemma4(t) and !std.mem.eql(u8, t, "llama") and !std.mem.eql(u8, t, "qwen3_5_moe")) return "--drafter is supported for Qwen3.8-27B-class checkpoints (qwen3_5 with its own output head), Qwen3.5-MoE (qwen3_5_moe), 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, "minimax_m3_vl")) "MiniMax-M3.1-Flash" else if (std.mem.eql(u8, t, "qwen3_5_moe")) "Qwen3.5-MoE" 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"; + 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, "minimax_m3_vl")) "MiniMax-M3.1-Flash" else if (std.mem.eql(u8, t, "qwen3_5_moe")) "Qwen3.5-MoE" else if (std.mem.eql(u8, t, "llama")) "Llama" else if (std.mem.eql(u8, t, "ricursive")) "Ripple" 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, "minimax_m3_vl")) "MiniMax-M3.1-Flash" else if (std.mem.eql(u8, t, "qwen3_5_moe")) "Qwen3.5-MoE" 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"; + 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, "minimax_m3_vl")) "MiniMax-M3.1-Flash" else if (std.mem.eql(u8, t, "qwen3_5_moe")) "Qwen3.5-MoE" else if (std.mem.eql(u8, t, "llama")) "Llama" else if (std.mem.eql(u8, t, "ricursive")) "Ripple" 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 and std.mem.eql(u8, t, "qwen3_5")) " (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/native/ricursive.zig b/zig/src/native/ricursive.zig new file mode 100644 index 0000000000..4b95b112a7 --- /dev/null +++ b/zig/src/native/ricursive.zig @@ -0,0 +1,143 @@ +//! Ripple on Metal served by the lane host: one-row rounds, windows of copied drafts, contexts budgeted. +const std = @import("std"); +const mtl = @import("metal"); +const api = @import("engine_api"); +const tf = @import("tensorfold"); +const fr = tf.ricursive; +const lanes = tf.lanes; +const Resident = @import("resident.zig").Resident; + +const Host = struct { + gpa: std.mem.Allocator, + io: std.Io, + model: *fr.Model, + metal: *fr.backend.Metal, + warm: mtl.keepalive.Target = undefined, // the residency set's, or the bare queue's, for the lane host's idle ticker + resident: ?Resident = null, // weights and scratch kept wired between requests, when they fit the 70% limit + config: lanes.Config, + clock: lanes.backend.WallClock, + core: lanes.Engine, + host: api.LaneHost, + cache: ?api.prompt_cache.Store = null, // kept prompt states, in what the streams' caches leave of the budget + cache_cap: u64 = 0, // --prompt-cache-gib's bytes (the budget when unset) + + fn close(ctx: *anyopaque) void { + const h: *Host = @ptrCast(@alignCast(ctx)); + h.host.stop(); + if (h.cache) |*store| store.deinit(); + h.core.deinit(); + h.config.deinit(h.gpa); + drop(h); + h.metal.deinit(); + h.model.deinit(); + h.gpa.destroy(h); + } +}; + +pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !?api.Opened { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const model = fr.Model.load(gpa, io, o.dir) catch |err| { + problem.* = switch (err) { + error.RicursiveNeedsTensorUnits => "the native Ripple 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.UnqualifiedRicursiveCheckpoint => "the native Ripple engine serves only the checkpoints checked against the release (Ripple-65B-Preview): this ricursive checkpoint's geometry is not one of them", + error.UnsupportedRicursivePrecision => "the native Ripple engine serves the publisher's releases as published (ricursiveAI/Ripple-65B-Preview in bf16, or its -NVFP4 MLPs); this checkpoint is quantized otherwise", + else => try std.fmt.allocPrint(a, "the native Ripple engine cannot load this checkpoint ({s})", .{@errorName(err)}), + }; + return null; + }; + errdefer model.deinit(); + var ram: u64 = 0; + var len: usize = @sizeOf(u64); + if (std.c.sysctlbyname("hw.memsize", &ram, &len, null, 0) != 0) ram = 0; + // the decode window's and the prompt chunk's buffers + const scratch: u64 = fr.state.scratchBytes(&model.config, fr.state.batch_rows) + fr.state.scratchBytes(&model.config, fr.state.prompt_rows); + const budget = try fr.budget.caches(ram, model.device.maxWorkingSet(), null, model.checkpoint.residentBytes() + scratch); + const most: usize = @min(model.config.context, fr.state.max_capacity); + const fits = fr.budget.longest(&model.config, budget, most); + if (o.context) |asked| if (asked > fr.config.max_context) { + problem.* = try std.fmt.allocPrint(a, "--context {d}: the native Ripple engine is checked to {d} tokens; longer contexts are refused until a run at that length qualifies them", .{ asked, fr.config.max_context }); + model.deinit(); + return null; + }; + const window: i64 = o.context orelse @intCast(@min(fits, fr.config.max_context)); + if (window < 1 or window > fits) { + problem.* = if (o.context) |asked| + try std.fmt.allocPrint(a, "--context {d}: this Mac's Ripple caches fit {d} tokens beside the weights under its load limit (75% of RAM)", .{ asked, fits }) + else + try a.dupe(u8, "this Mac's load limit (75% of RAM) leaves no room for a Ripple stream's caches beside the weights"); + model.deinit(); // after the message, so a failed allocation leaves the model to the errdefer + return null; + } + std.log.info("ricursive: {d:.1} GiB for caches, context up to {d} tokens ({d:.2} GiB a full stream)", .{ @as(f64, @floatFromInt(budget)) / (1 << 30), window, @as(f64, @floatFromInt(fr.state.cacheBytes(&model.config, @intCast(window + 1)))) / (1 << 30) }); + const streams: u32 = @max(o.lanes, 1); + const h = try gpa.create(Host); + errdefer gpa.destroy(h); + h.gpa = gpa; + h.io = io; + h.model = model; + h.resident = null; + h.metal = try fr.backend.Metal.init(gpa, model, .{ .capacity = @intCast(window), .streams = streams, .budget = budget }); + errdefer h.metal.deinit(); + try h.metal.warm(); + h.config = try lanes.Config.init(gpa, h.metal.facts(), fr.state.window_rows, fr.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 = fr.state.prompt_rows, .structures = true, .logprobs = true }); + h.host.info_.reasoning_optional = true; // Ripple writes `` itself when it thinks, and answers bare otherwise + h.host.info_.thinking_budget = false; // its think markers are spelled, not one token each: no close can be forced + 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 = fr.backend.Metal; + // prompt chunks change arithmetic with their rows + const rules: api.prompt_cache.Rules = .{ .lookahead = 0, .planned = true, .grid = fr.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 hold(h); + errdefer drop(h); + h.host.keepalive_target = .{ .ctx = &h.warm, .tick = mtl.keepalive.Target.tick }; + try h.host.start(); + return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h }; +} + +/// Weights, tables and scratches in one residency set for the keepalive, under the 70% limit. +fn hold(h: *Host) !void { + var list: std.ArrayList(mtl.Buffer) = .empty; + defer list.deinit(h.gpa); + const m = h.model; + for (m.checkpoint.shards.items) |sh| try list.append(h.gpa, sh.buffer); + try list.append(h.gpa, m.inv); + try mtl.resident.fields.append(h.gpa, &list, &h.metal.scratch); + try mtl.resident.fields.append(h.gpa, &list, &h.metal.prompt); + h.resident = try Resident.hold(h.gpa, h.io, m.device, m.queue, list.items, "ricursive"); + if (h.resident) |*r| h.metal.setResident(r.follow()); + h.warm = if (h.resident) |*r| r.target() else .{ .queue = m.queue }; +} + +fn drop(h: *Host) void { + h.metal.setResident(null); + if (h.resident) |*r| r.deinit(); + h.resident = null; + h.warm = .{ .queue = h.model.queue }; +} + +/// 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 }); +} diff --git a/zig/src/server/chat.zig b/zig/src/server/chat.zig index 85143559f0..8e4e235568 100644 --- a/zig/src/server/chat.zig +++ b/zig/src/server/chat.zig @@ -143,7 +143,8 @@ pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepa errdefer srv.release(preparing, null); if (gone.check()) return error.Cancelled; const thinking = input.prompt == null and (flag(f, "enable_thinking") orelse srv.config.enable_thinking); - const reasoning = thinking and srv.info.reasoning and (srv.text.vtable.format != .kimi or input.prompt == null); // the prompt renders `thinking`; the reply reads `reasoning` + // the prompt renders `thinking`; the reply reads `reasoning` (always, when the model decides for itself) + const reasoning = (thinking or (srv.markers.optional and input.prompt == null)) and srv.info.reasoning and (srv.text.vtable.format != .kimi or input.prompt == null); const effort = srv.effortFor(if (f.get("reasoning_effort")) |e| (if (e == .string) e.string else null) else null); cx.clear_thinking = flag(f, "clear_thinking"); cx.template_fields = f; diff --git a/zig/src/server/reply_text.zig b/zig/src/server/reply_text.zig index 6f942d1e4a..0f44d8b492 100644 --- a/zig/src/server/reply_text.zig +++ b/zig/src/server/reply_text.zig @@ -4,7 +4,8 @@ const json = @import("json"); const Value = json.Value; const Allocator = std.mem.Allocator; -pub const Markers = struct { open: []const u8, close: []const u8 }; +/// `optional`: the reply writes its own opener when it thinks, and a reply that does not is all answer (Ripple's). +pub const Markers = struct { open: []const u8, close: []const u8, optional: bool = false }; /// Qwen's prompt opens the think block; Gemma 4's reply opens its thought channel. pub const think_markers: Markers = .{ .open = "", .close = "" }; /// MiniMax-M3's adaptive thinking: the reply writes `...` itself, or `` alone. @@ -155,6 +156,10 @@ pub fn splitThinkingAt(a: Allocator, full: []const u8, finished: bool, m: Marker text = trimmed; } else if (!finished and std.mem.startsWith(u8, openerOf(m), text)) { return .{ .reasoning = "", .answer = "" }; // held while the reply may still be writing that opener + } else if (m.optional) { + var held: usize = 0; // a reply without the opener answers: only a partial call opener is held while it streams + for (o.calls) |c| held = @max(held, if (finished) 0 else partialTag(text, c[0])); + return .{ .reasoning = "", .answer = text[0 .. text.len - held] }; } if (at) |v| at = if (v >= shift) v - shift else null; const end: ?usize = switch (close) { @@ -561,6 +566,19 @@ test "reply text: MiniMax-M3's adaptive think block, written by the reply or ski try std.testing.expectEqualStrings("", held.answer); } +test "reply text: an optional think block, a reply that opens one splits, a bare reply is all answer" { + const m: Markers = .{ .open = "", .close = "", .optional = true }; + for ([_]struct { []const u8, bool, []const u8, []const u8 }{ + .{ "\nAdd them.\n\n\n42", true, "Add them.\n", "42" }, + .{ "17 * 23 = 391.", true, "", "17 * 23 = 391." }, + .{ "Sure. = 0) 0 else @min(r[0] * r[0], r[1] * r[1]), hi }; +} + +/// rsqrt(ss / width + eps) over a sum of squares in `ss`: the fp32 sum's gamma(depth), the mean, the eps and rsqrt. +pub fn invRms(ss: Range, width: usize, depth: usize) Range { + const gm = c.gamma(@floatFromInt(depth), u); + const w: f64 = @floatFromInt(width); + return .{ (1 - 6 * u) / @sqrt(ss[1] * (1 + gm) / w + eps), (1 + 6 * u) / @sqrt(@max(ss[0] * (1 - gm), 0) / w + eps) }; +} + +/// Each output's allowance: its distance to the far end of its range, so the range alone decides. +pub fn allowRange(lo: []const f64, hi: []const f64, want: []const f64, allow: []f64) void { + for (allow, lo, hi, want) |*l, a, b, w| l.* = @max(w - a, b - w, 0) + 0x1p-133; +} + +pub fn none(_: c.Shape) usize { + return 0; +} + +pub const no_sums: c.Sums = .{ .f16 = 0, .bf16 = 0, .why = "no sums: elementwise, a gather or a copy" }; + +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 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. +pub 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 xv(v: []const u16, i: usize) f64 { + return c.fromBf16(v[i]); +} + +pub const Which = enum { embed, gate, silu_clamp }; + +pub fn function(w: Which) []const u8 { + return switch (w) { + .embed => "rc_embed", + .gate => "rc_gate", + .silu_clamp => "rc_silu_clamp", + }; +} + +/// Prompt chunks to 2,048 rows and decode windows of 16; the FFN's 40,960 wide rows to 512 (the grid's rows repeat). +pub fn spec(w: Which) c.Spec { + const q = R.heads * DH; + return switch (w) { + .embed => .{ .name = "ricursive glue rc_embed", .accum = .none, .depth = none, .ops = 1, .out = .bf16, .served = &.{ .{ .rows = 16, .n = R.hidden, .k = R.hidden }, .{ .rows = 2048, .n = R.hidden, .k = R.hidden } }, .sums = no_sums }, + .gate => .{ .name = "ricursive glue rc_gate", .accum = .none, .depth = none, .ops = 8, .mids = 1, .out = .bf16, .in_place = true, .served = &.{ .{ .rows = 16, .n = q, .k = q }, .{ .rows = 2048, .n = q, .k = q } }, .sums = no_sums }, + .silu_clamp => .{ .name = "ricursive glue rc_silu_clamp", .accum = .none, .depth = none, .ops = 8, .mids = 1, .out = .bf16, .served = &.{ .{ .rows = 16, .n = R.intermediate, .k = R.intermediate }, .{ .rows = 512, .n = R.intermediate, .k = R.intermediate } }, .sums = no_sums }, + }; +} + +/// The elementwise glue: the width is the shape's k, the rows its rows. +pub const Pointwise = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + which: Which, + mutation: controls.Mutation, + lib: mtl.Library, + pipe: mtl.Pipeline, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + x: []u16 = &.{}, // the table (embed), the attention rows (gate), the stacked gate|up rows (silu_clamp) + y: []u16 = &.{}, // the gate rows (gate) + ids: []u32 = &.{}, + lo: []f64 = &.{}, + hi: []f64 = &.{}, + in: [2]rig_.Ref = undefined, + out: [2]rig_.Ref = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, which: Which, mutation: controls.Mutation) !Pointwise { + const lib = try controls.glueLibrary(device, a, mutation); + errdefer lib.deinit(); + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .which = which, .mutation = mutation, .lib = lib, .pipe = try pipelineOf(lib, device, function(which)) }; + } + + fn drop(p: *Pointwise) void { + for ([_]*[]u16{ &p.x, &p.y }) |s| { + p.a.free(s.*); + s.* = &.{}; + } + for ([_]*[]f64{ &p.lo, &p.hi }) |s| { + p.a.free(s.*); + s.* = &.{}; + } + p.a.free(p.ids); + p.ids = &.{}; + p.rig.reset(); + } + + pub fn deinit(p: *Pointwise) void { + p.drop(); + p.rig.deinit(); + p.pipe.deinit(); + p.lib.deinit(); + } + + const vt = g.hooks(Pointwise); + + pub fn kernel(p: *Pointwise) g.Kernel { + return controls.mark(.{ .spec = spec(p.which), .ctx = p, .vt = &vt }, p.mutation); + } + + pub fn prepare(p: *Pointwise, 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.x = try p.a.alloc(u16, 8 * w); + random16(rng, p.x, -2, 2); + if (fill == .nonfinite) plant(p.x, w, 8); + p.ids = try p.a.alloc(u32, s.rows); + for (p.ids, 0..) |*id, i| id.* = @intCast((i * 5 + 3) % 8); + p.in[0] = try p.rig.input(std.mem.sliceAsBytes(p.ids)); + p.in[1] = try p.rig.input(std.mem.sliceAsBytes(p.x)); + p.out[0] = try p.rig.output(n * 2); + p.out[1] = try p.rig.output(n * 2); + }, + .gate => { + 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, -8, 8); + if (fill == .nonfinite) { + plant(p.y, w, s.rows); + if (s.rows > 4) { // an infinite attention value gated to 0: NaN; a NaN value + p.x[3 * w] = c.toBf16(std.math.inf(f32)); + p.y[3 * w] = c.toBf16(-std.math.inf(f32)); + p.x[4 * w + 1] = c.toBf16(std.math.nan(f32)); + } + } + p.in[0] = try p.rig.input(std.mem.sliceAsBytes(p.y)); + p.out[0] = try p.rig.output(n * 2); + }, + .silu_clamp => { + p.x = try p.a.alloc(u16, 2 * n); + random16(rng, p.x, -40, 40); // both sides of the clamp + if (fill == .nonfinite) { + plant(p.x, 2 * w, s.rows); + // up's NaN and infinities in rows 3 to 5, where each row's gate is finite + if (s.rows > 5) for ([_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32) }, 3..) |f, r| { + p.x[r * 2 * w + w + 7] = c.toBf16(f); + }; + } + p.in[0] = try p.rig.input(std.mem.sliceAsBytes(p.x)); + p.out[0] = try p.rig.output(n * 2); + }, + } + p.lo = try p.a.alloc(f64, p.count()); + p.hi = try p.a.alloc(f64, p.count()); + return true; + } + + pub fn count(p: *Pointwise) usize { + const n = p.shape.rows * p.shape.k; + return if (p.which == .embed) 2 * n else n; + } + + pub fn run(p: *Pointwise) !void { + const s = p.shape; + const w: u32 = @intCast(s.k); + const rows: u32 = @intCast(s.rows); + const pass = p.rig.begin(); + const e = pass.e; + e.setPipeline(p.pipe); + 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, .scale = embed_scale }, 2); + e.setBuffer(p.out[0].buf, p.out[0].off, 3); + e.setBuffer(p.out[1].buf, p.out[1].off, 4); + }, + .gate => { + @memcpy(p.rig.bytes(p.out[0], p.x.len * 2), std.mem.sliceAsBytes(p.x)); // the rows before the gate + 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); + }, + .silu_clamp => { + e.setBuffer(p.in[0].buf, p.in[0].off, 0); + e.setValue(Args{ .rows = rows, .width = w, .scale = clamp }, 1); + e.setBuffer(p.out[0].buf, p.out[0].off, 2); + }, + } + e.dispatchThreads(mtl.Size.of(s.k, s.rows, 1), mtl.Size.of(256, 1, 1)); + try p.rig.finish(pass); + } + + /// The release's formula in float64 at its rounding points (want), and the range its fp32 steps allow (lo, hi). + pub fn reference(p: *Pointwise, want: []f64, mag: []f64) !void { + const s = p.shape; + const w = s.k; + const n = s.rows * w; + for (0..s.rows) |r| for (0..w) |i| { + const at = r * w + i; + switch (p.which) { + .embed => { + const v = xv(p.x, p.ids[r] * w + i) * @as(f64, embed_scale); + const e = round(widen(point(v), u)); + for ([_]usize{ at, n + at }) |o| { + want[o] = bf(v); + p.lo[o] = e[0]; + p.hi[o] = e[1]; + } + }, + .gate => { + const gv = xv(p.y, at); + const sig = 1.0 / (1.0 + @exp(-gv)); // exp, 1 + e and the division: 8 fp32 ulps + const e = round(mul(point(xv(p.x, at)), round(widen(point(sig), 8 * u)))); + want[at] = bf(xv(p.x, at) * bf(sig)); + p.lo[at] = e[0]; + p.hi[at] = e[1]; + }, + .silu_clamp => { + const g0 = xv(p.x, r * 2 * w + i); + const up0 = xv(p.x, r * 2 * w + w + i); + // torch's clamp keeps a NaN + const gv = if (std.math.isNan(g0)) g0 else @min(g0, clamp); + const up = if (std.math.isNan(up0)) up0 else std.math.clamp(up0, -clamp, clamp); + const silu = gv / (1.0 + @exp(-gv)); + const e = round(mul(round(widen(point(silu), 8 * u)), point(up))); + want[at] = bf(bf(silu) * up); + p.lo[at] = e[0]; + p.hi[at] = e[1]; + }, + } + }; + for (mag, want) |*m, v| m.* = @abs(v); + } + + pub fn bounds(p: *Pointwise, want: []const f64, _: []const f64, allow: []f64) !void { + allowRange(p.lo, p.hi, want, allow); + } + + pub fn intervals(p: *Pointwise, low: []f64, high: []f64) !void { + @memcpy(low, p.lo); + @memcpy(high, p.hi); + } +}; diff --git a/zig/tests/contract/metal_ricursive_all.zig b/zig/tests/contract/metal_ricursive_all.zig new file mode 100644 index 0000000000..41c47bbde3 --- /dev/null +++ b/zig/tests/contract/metal_ricursive_all.zig @@ -0,0 +1,259 @@ +//! Ripple's registrations as one set: its kernels (NVFP4 too) and controls, llama's, the draw; and the roll call. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const sources = @import("kernel_sources"); +const rc = @import("metal_ricursive.zig"); +const norm = @import("metal_ricursive_norm.zig"); +const attention = @import("metal_ricursive_attention.zig"); +const sample = @import("metal_ricursive_sample.zig"); +const fp4 = @import("metal_ricursive_fp4.zig"); +const fp4_gemm = @import("metal_ricursive_fp4_gemm.zig"); +const controls = @import("metal_ricursive_controls.zig"); +const llama = @import("metal_llama.zig"); +const dense = @import("metal_llama_dense.zig"); +const llama_attention = @import("metal_llama_attention.zig"); +const c = contract.spec; +const g = contract.gates; +const family = rc.family; +const R = rc.R; + +/// The label Ripple's registrations of llama's pipelines carry, and those at o_proj's reduction. +const label = "ricursive llama"; +const o_label = label ++ " o"; +const matvec_rows = [_]u32{ 1, 2, 4, 8, 16, 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 row projections at Ripple's reductions: the FFN's down (k 40,960) and those over the hidden width (k 5,120). +fn matvecShapes(per: u32) [2]c.Shape { + return .{ .{ .rows = per, .n = 1024, .k = R.intermediate }, .{ .rows = per, .n = 1024, .k = R.hidden } }; +} + +/// A prompt chunk's GEMM at k 5,120 and the down's 40,960, whole or ragged tiles as the form takes them. +fn gemmShapes(am: bool, an: bool) [2]c.Shape { + const n: usize = if (an) 1024 else 1088; + return .{ .{ .rows = if (am) 64 else 33, .n = n, .k = R.hidden }, .{ .rows = if (am) 128 else 100, .n = n, .k = R.intermediate } }; +} + +/// o_proj's reduction over the heads (k 10,240): its 5,120 outputs, and a slice of them. +fn oMatvecShapes(per: u32) [2]c.Shape { + const k = R.heads * rc.DH; + return .{ .{ .rows = per, .n = 1024, .k = k }, .{ .rows = per, .n = R.hidden, .k = k } }; +} + +/// o_proj's GEMM over a prompt chunk, whole or ragged tiles as the form takes them. +fn oGemmShapes(am: bool, an: bool) [2]c.Shape { + const n: usize = if (an) 1024 else 1088; + const k = R.heads * rc.DH; + return .{ .{ .rows = if (am) 64 else 33, .n = n, .k = k }, .{ .rows = if (am) 128 else 100, .n = n, .k = k } }; +} + +/// A decode window's merge over a sliding layer's handful of blocks, and over a global layer's most (state.max_blocks). +const merge_shapes = [_]c.Shape{ .{ .rows = 16, .n = R.heads, .k = 544 }, .{ .rows = 16, .n = R.heads, .k = 9 } }; +const argmax_shapes = [_]c.Shape{.{ .rows = 1, .n = 1, .k = R.vocab }}; + +/// The probes made for a run, freed together. +pub const Set = struct { + point: std.ArrayList(*rc.Pointwise) = .empty, + add_norm: std.ArrayList(*norm.AddNorm) = .empty, + prep: std.ArrayList(*norm.Prep) = .empty, + decode: std.ArrayList(*attention.DecodeVec) = .empty, + prompt: std.ArrayList(*attention.Prompt) = .empty, + matvec: std.ArrayList(*dense.Matvec) = .empty, + gemm: std.ArrayList(*dense.Gemm) = .empty, + merge: std.ArrayList(*llama_attention.MergeVec) = .empty, + argmax: std.ArrayList(*llama.Argmax) = .empty, + sample: std.ArrayList(*sample.Sample) = .empty, + fp4_matvec: std.ArrayList(*fp4.Matvec) = .empty, + fp4_dequant: std.ArrayList(*fp4.Dequant) = .empty, + fp4_gemm: std.ArrayList(*fp4_gemm.Gemm) = .empty, + + pub fn deinit(s: *Set, a: std.mem.Allocator) void { + inline for (@typeInfo(Set).@"struct".field_names) |f| { + const list = &@field(s, f); + for (list.items) |p| { + p.deinit(); + a.destroy(p); + } + list.deinit(a); + } + } + + /// `v` kept in `list` (freed with the set) and registered; a failure frees it. + fn keep(a: std.mem.Allocator, list: anytype, v: anytype, kernels: *std.ArrayList(g.Kernel)) !void { + var x = v; + const p = a.create(@TypeOf(v)) catch |err| { + x.deinit(); + return err; + }; + p.* = x; + list.append(a, p) catch |err| { + p.deinit(); + a.destroy(p); + return err; + }; + try kernels.append(a, p.kernel()); + } + + /// 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: std.ArrayList([]const u8) = .empty; + defer functions.deinit(a); + for (std.enums.values(rc.Which)) |w| { + try functions.append(a, rc.function(w)); + if (wanted(only, rc.spec(w).name)) try keep(a, &s.point, try rc.Pointwise.init(a, device, queue, w, .none), kernels); + } + try functions.append(a, "rc_add_norm"); + for (std.enums.values(norm.Form)) |f| if (wanted(only, norm.addNormSpec(f).name)) { + try keep(a, &s.add_norm, try norm.AddNorm.init(a, device, queue, f, .none), kernels); + }; + try functions.append(a, "rc_prep"); + for (std.enums.values(norm.Layer)) |l| if (wanted(only, norm.prepSpec(l).name)) { + try keep(a, &s.prep, try norm.Prep.init(a, device, queue, l, .none), kernels); + }; + try functions.append(a, "rc_decode_vec"); + for (std.enums.values(attention.Layer)) |l| if (wanted(only, attention.decodeSpec(l).name)) { + try keep(a, &s.decode, try attention.DecodeVec.init(a, device, queue, l, .none), kernels); + }; + // the prompt path's GEMM and attention run on the tensor units only (model.zig refuses a GPU without them) + try functions.append(a, family.attention_name); + for (std.enums.values(attention.Layer)) |l| if (tensor and wanted(only, attention.promptSpec(l).name)) { + try keep(a, &s.prompt, try attention.Prompt.init(a, device, queue, l, .none), kernels); + }; + for (matvec_rows, family.llama_names[0..matvec_rows.len]) |per, fun| { + try functions.append(a, fun); + var buf: [48]u8 = undefined; + if (!wanted(only, try std.fmt.bufPrint(&buf, "{s} {s}", .{ label, fun }))) continue; + var m = try dense.Matvec.init(a, device, queue, per, .none); + m.retarget(label, matvecShapes(per)) catch |err| { + m.deinit(); + return err; + }; + try keep(a, &s.matvec, m, kernels); + } + for (family.gemm_names, 0..) |fun, i| { + try functions.append(a, fun); + var buf: [48]u8 = undefined; + if (!tensor or !wanted(only, try std.fmt.bufPrint(&buf, "{s} {s}", .{ label, fun }))) continue; + const am = i / 2 == 1; + const an = i % 2 == 1; + var m = try dense.Gemm.init(a, device, queue, am, an, .none); + m.retarget(label, gemmShapes(am, an)) catch |err| { + m.deinit(); + return err; + }; + try keep(a, &s.gemm, m, kernels); + } + try functions.append(a, "ll_merge_vec"); + if (wanted(only, label ++ " ll_merge_vec")) { + var m = try llama_attention.MergeVec.init(a, device, queue, .none); + m.name = label ++ " ll_merge_vec"; + m.served = &merge_shapes; + try keep(a, &s.merge, m, kernels); + } + try functions.append(a, "ll_argmax"); + if (wanted(only, label ++ " ll_argmax")) { + var m = try llama.Argmax.init(a, device, queue, .none); + m.name = label ++ " ll_argmax"; + m.served = &argmax_shapes; + try keep(a, &s.argmax, m, kernels); + } + try functions.append(a, family.sample_name); + if (wanted(only, sample.spec().name)) try keep(a, &s.sample, try sample.Sample.init(a, device, queue), kernels); + try s.collectFp4(a, device, queue, tensor, kernels, only, &functions); + try s.collectO(a, device, queue, tensor, kernels, only); + try family.requireComplete(functions.items); + try s.plant(a, device, queue, tensor, kernels, only); + } + + /// The NVFP4 MLP's row kernels, its dequantization and (tensor units only) its four GEMM forms. + fn collectFp4(s: *Set, a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, tensor: bool, kernels: *std.ArrayList(g.Kernel), only: ?[]const u8, functions: *std.ArrayList([]const u8)) !void { + for (matvec_rows, family.fp4_names[0..matvec_rows.len]) |per, fun| { + try functions.append(a, fun); + var buf: [48]u8 = undefined; + if (wanted(only, try std.fmt.bufPrint(&buf, "ricursive fp4 {s}", .{fun}))) try keep(a, &s.fp4_matvec, try fp4.Matvec.init(a, device, queue, per, .none), kernels); + } + try functions.append(a, "rc_fp4_dequant"); + if (wanted(only, "ricursive fp4 rc_fp4_dequant")) try keep(a, &s.fp4_dequant, try fp4.Dequant.init(a, device, queue, .none), kernels); + for (family.fp4_gemm_names, 0..) |fun, i| { + try functions.append(a, fun); + var buf: [48]u8 = undefined; + if (tensor and wanted(only, try std.fmt.bufPrint(&buf, "ricursive fp4 gemm {s}", .{fun}))) { + try keep(a, &s.fp4_gemm, try fp4_gemm.Gemm.init(a, device, queue, i / 2 == 1, i % 2 == 1, .none), kernels); + } + } + } + + /// llama's row projections and GEMMs at o_proj's reduction (already on the roll call above). + fn collectO(s: *Set, a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, tensor: bool, kernels: *std.ArrayList(g.Kernel), only: ?[]const u8) !void { + var buf: [48]u8 = undefined; + for (matvec_rows, family.llama_names[0..matvec_rows.len]) |per, fun| { + if (!wanted(only, try std.fmt.bufPrint(&buf, "{s} {s}", .{ o_label, fun }))) continue; + var m = try dense.Matvec.init(a, device, queue, per, .none); + m.retarget(o_label, oMatvecShapes(per)) catch |err| { + m.deinit(); + return err; + }; + try keep(a, &s.matvec, m, kernels); + } + for (family.gemm_names, 0..) |fun, i| { + if (!tensor or !wanted(only, try std.fmt.bufPrint(&buf, "{s} {s}", .{ o_label, fun }))) continue; + const am = i / 2 == 1; + const an = i % 2 == 1; + var m = try dense.Gemm.init(a, device, queue, am, an, .none); + m.retarget(o_label, oGemmShapes(am, an)) catch |err| { + m.deinit(); + return err; + }; + try keep(a, &s.gemm, m, kernels); + } + } + + /// Each control: its kernel built with one planted fault, registered beside the kernels. + fn plant(s: *Set, a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, tensor: bool, kernels: *std.ArrayList(g.Kernel), only: ?[]const u8) !void { + for (std.enums.values(controls.Mutation)) |m| { + if (m == .none or !wanted(only, m.name())) continue; + switch (m) { + .embed_own_row => try keep(a, &s.point, try rc.Pointwise.init(a, device, queue, .embed, m), kernels), + .gate_early_round => try keep(a, &s.point, try rc.Pointwise.init(a, device, queue, .gate, m), kernels), + .clamp_no_lower => try keep(a, &s.point, try rc.Pointwise.init(a, device, queue, .silu_clamp, m), kernels), + .add_norm_no_residual => try keep(a, &s.add_norm, try norm.AddNorm.init(a, device, queue, .add, m), kernels), + .lerp_no_blend => try keep(a, &s.add_norm, try norm.AddNorm.init(a, device, queue, .lerp, m), kernels), + .prep_ropes_global => try keep(a, &s.prep, try norm.Prep.init(a, device, queue, .global, m), kernels), + .prep_next_slot => try keep(a, &s.prep, try norm.Prep.init(a, device, queue, .sliding, m), kernels), + .decode_past_window => try keep(a, &s.decode, try attention.DecodeVec.init(a, device, queue, .sliding, m), kernels), + .prompt_past_window, .prompt_short_ring => if (tensor) { + try keep(a, &s.prompt, try attention.Prompt.init(a, device, queue, .sliding, m), kernels); + }, + .fp4_high_first, .fp4_scale_bias, .fp4_global_late => try keep(a, &s.fp4_matvec, try fp4.Matvec.init(a, device, queue, 4, m), kernels), + .fp4_dequant_unscaled => try keep(a, &s.fp4_dequant, try fp4.Dequant.init(a, device, queue, m), kernels), + .fp4_gemm_no_global => if (tensor) { + try keep(a, &s.fp4_gemm, try fp4_gemm.Gemm.init(a, device, queue, false, false, m), kernels); + }, + .none => unreachable, + } + } + } +}; + +/// Gate 5's texts for `ricursive ` launches: the glue, the draw, or the attention (nax.h inlined; fragments pre-M5). +pub fn limitsText(a: std.mem.Allocator, device: mtl.Device, name: []const u8) !struct { here: []const u8, pre_m5: []const u8 } { + if (std.mem.indexOf(u8, name, " attention ") != null) { + const head = try family.attentionHeader(a); + const body = try std.mem.replaceOwned(u8, a, sources.ricursive_attention_nax, "#include \"../nax.h\"", sources.nax); + const pre = try std.mem.concat(a, u8, &.{ head, "#define TF_SIMD_FRAGS 1\n", body }); + return .{ .here = if (device.tensorUnits()) try family.attentionSource(device, a) else pre, .pre_m5 = pre }; + } + if (std.mem.indexOf(u8, name, " fp4 gemm ") != null) { + const pre = try family.fp4GemmSource(device, a, false); + return .{ .here = if (device.tensorUnits()) try family.fp4GemmSource(device, a, true) else pre, .pre_m5 = pre }; + } + if (std.mem.indexOf(u8, name, " fp4 ") != null) return .{ .here = sources.ricursive_fp4, .pre_m5 = sources.ricursive_fp4 }; + if (std.mem.indexOf(u8, name, " sample ") != null) return .{ .here = sources.nemotron_sample, .pre_m5 = sources.nemotron_sample }; + const glue = try family.glueSource(a); + return .{ .here = glue, .pre_m5 = glue }; +} diff --git a/zig/tests/contract/metal_ricursive_attention.zig b/zig/tests/contract/metal_ricursive_attention.zig new file mode 100644 index 0000000000..68d8a42bd6 --- /dev/null +++ b/zig/tests/contract/metal_ricursive_attention.zig @@ -0,0 +1,421 @@ +//! Ripple's attention under the contract: rc_decode_vec and rc_prompt_attention, sliding (a ring) and global layers. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const controls = @import("metal_ricursive_controls.zig"); +const rc = @import("metal_ricursive.zig"); +const c = contract.spec; +const g = contract.gates; +const rows_in = contract.rows; +const st = tf.ricursive.state; +const R = rc.R; +const DH = rc.DH; +const xv = rc.xv; +const kvw = R.kv_heads * DH; +const grp = R.heads / R.kv_heads; +const scale: f64 = tf.ricursive.config.headScale(); +const nan16 = 0x7fc0; +/// A global layer's window: every key before (forward.zig's every_key). +const every_key: usize = std.math.maxInt(i32); + +/// Sliding layers (RoPE'd keys in a ring, a 1,024-key window) and global layers (every key, a slot a position). +pub const Layer = enum { sliding, global }; + +fn windowOf(l: Layer) usize { + return if (l == .sliding) R.window else every_key; +} + +/// The cache's slots for `k` positions: the sliding ring, or whole 64-key blocks (state.slots). +fn ringOf(l: Layer, k: usize) usize { + return if (l == .sliding) rc.family.sliding_ring else st.slots(k); +} + +/// The first key position row `pos` sees. +fn firstKey(l: Layer, pos: usize) usize { + return (pos + 1) -| windowOf(l); +} + +const softmax_sums: c.Sums = .{ .f16 = 0, .bf16 = 0, .exact = false, .why = "weights of exp: fp32 sums held to the float64 bound, not integers" }; + +/// rc_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, window: u32, ring: u32 }; + +/// Random K and V in slots holding positions to `last`; NaN in slots none reached or holding one before `from`. +fn fillCache(rng: std.Random, k: []u16, v: []u16, ring: usize, last: usize, from: ?usize) void { + rc.random16(rng, k, -2, 2); + rc.random16(rng, v, -2, 2); + for (0..ring) |s| { + const newest: ?usize = if (s > last) null else last - (last - s) % ring; + const poison = newest == null or (from != null and newest.? < from.?); + if (poison) { + @memset(k[s * kvw ..][0..kvw], nan16); + @memset(v[s * kvw ..][0..kvw], nan16); + } + } +} + +/// The scaled dot of query `qrow` (a [heads, DH] row) head `h` with the key at position `key`, and its terms' size. +fn dot(q: []const u16, qrow: usize, h: usize, k: []const u16, slot: usize) [2]f64 { + const kv = h / grp; + var sum: f64 = 0; + var abs: f64 = 0; + for (0..DH) |d| { + const t = scale * xv(q, (qrow * R.heads + h) * DH + d) * xv(k, slot * kvw + kv * DH + d); + sum += t; + abs += @abs(t); + } + return .{ sum, abs }; +} + +/// Every key of a block in sequence, four at a time, then the merge of the four lanes. +fn slidingDepth(_: c.Shape) usize { + return st.ring_chunk + 16; +} + +fn globalDepth(s: c.Shape) usize { + return @min(s.k, 2048) + 16; +} + +pub fn decodeSpec(l: Layer) c.Spec { + const served: []const c.Shape = switch (l) { + // a verify window deep in a stream (its window across the ring's wrap), a short one early + .sliding => &.{ .{ .rows = 16, .n = R.heads, .k = 100016 }, .{ .rows = 4, .n = R.heads, .k = 1500 } }, + .global => &.{ .{ .rows = 16, .n = R.heads, .k = 6000 }, .{ .rows = 3, .n = R.heads, .k = 700 } }, + }; + const chunk: usize = if (l == .sliding) st.ring_chunk else 32; + return .{ .name = if (l == .sliding) "ricursive sliding rc_decode_vec" else "ricursive global rc_decode_vec", .accum = .f32, .depth = if (l == .sliding) slidingDepth else globalDepth, .ops = 24, .out = .f32, .tiles = .{ .{ .size = 16, .step = 1, .min = 1 }, .{}, .{ .size = chunk, .step = 1, .min = 1 } }, .served = served, .sums = softmax_sums }; +} + +/// rc_decode_vec: a (key head, row, key block) a threadgroup over the blocks forward.zig dispatches for the rows. +pub const DecodeVec = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + layer: Layer, + mutation: controls.Mutation, + lib: mtl.Library, + pipe: mtl.Pipeline, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + ring: usize = 0, + chunk: usize = 0, + first: usize = 0, + 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, layer: Layer, mutation: controls.Mutation) !DecodeVec { + const lib = try controls.glueLibrary(device, a, mutation); + errdefer lib.deinit(); + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .layer = layer, .mutation = mutation, .lib = lib, .pipe = try rc.pipelineOf(lib, device, "rc_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 controls.mark(.{ .spec = decodeSpec(p.layer), .ctx = p, .vt = &vt }, p.mutation); + } + + fn p0(p: *const DecodeVec) usize { + return p.shape.k - p.shape.rows; + } + + /// rows query rows at the end of k positions; the blocks forward.zig dispatches for a stream of capacity k. + pub fn prepare(p: *DecodeVec, s: c.Shape, fill: g.Fill, seed: u64) !bool { + if (s.n != R.heads or s.rows == 0 or s.rows > st.window_rows or s.k < s.rows) return false; + p.drop(); + p.shape = s; + p.ring = ringOf(p.layer, s.k); + p.chunk = if (p.layer == .sliding) st.ring_chunk else st.chunk(s.k); + const start = p.p0(); + p.first = firstKey(p.layer, start) / p.chunk; + p.blocks = (start + s.rows - 1) / p.chunk + 1 - p.first; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + p.q = try p.a.alloc(u16, s.rows * R.heads * DH); + p.k = try p.a.alloc(u16, p.ring * kvw); + p.v = try p.a.alloc(u16, p.ring * kvw); + rc.random16(rng, p.q, -2, 2); + fillCache(rng, p.k, p.v, p.ring, s.k - 1, firstKey(p.layer, start)); + if (fill == .nonfinite and start > 1) { // keys every row sees: the first key head's at p0, the second's values + p.k[(start % p.ring) * kvw + 3] = c.toBf16(std.math.nan(f32)); + p.v[((start - 1) % p.ring) * kvw + DH + 1] = c.toBf16(std.math.nan(f32)); + p.v[((start - 2) % p.ring) * kvw + 2 * DH + 5] = c.toBf16(std.math.inf(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)); + const slots = R.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 R.heads * p.shape.rows * p.blocks * (2 + DH); + } + + pub fn run(p: *DecodeVec) !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); + const a = VArgs{ .rows = @intCast(s.rows), .pos = @intCast(p.p0()), .heads = R.heads, .groups = R.kv_heads, .blocks = @intCast(p.blocks), .chunk = @intCast(p.chunk), .first = @intCast(p.first), .scale = @floatCast(scale), .window = @intCast(windowOf(p.layer)), .ring = @intCast(p.ring) }; + e.setValue(a, 3); + for (p.out, 4..) |r, slot| e.setBuffer(r.buf, r.off, slot); + e.dispatchGroups(mtl.Size.of(R.kv_heads, s.rows, p.blocks), mtl.Size.of(32, grp, 1)); + try p.rig.finish(pass); + } + + /// A partial's keys: the block's within the row's window, up to the row's own position. + fn span(p: *const DecodeVec, row: usize, z: usize) [2]usize { + const pos = p.p0() + row; + const block = p.first + z; + const from = @max(block * p.chunk, firstKey(p.layer, pos)); + return .{ from, @max(from, @min(pos + 1, (block + 1) * p.chunk)) }; + } + + /// PM (the block's top score), PL (its mass) and PO (its weighted values); `acc` replays the sums. + fn partials(p: *DecodeVec, acc: ?struct { a: c.Accum, o: g.Order }, want: []f64, mag: ?[]f64) !void { + const s = p.shape; + const slots = R.heads * s.rows * p.blocks; + const sc = try p.a.alloc(f64, 2 * p.chunk); + defer p.a.free(sc); + const t = sc[p.chunk..]; + var abs_of: [2048]f64 = undefined; + for (0..R.heads) |h| for (0..s.rows) |row| for (0..p.blocks) |z| { + const slot = (h * s.rows + row) * p.blocks + z; + const kv = h / grp; + const from, const to = p.span(row, z); + 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; + for (from..to, 0..) |key, i| { + sc[i], abs_of[i] = dot(p.q, row, h, p.k, key % p.ring); + if (!std.math.isNan(sc[i]) and sc[i] > top) { + top = sc[i]; + top_abs = abs_of[i]; + } + } + const n = to - from; + var mass: f64 = 0; + for (0..n) |i| { + t[i] = @exp(sc[i] - top); + mass += t[i]; + } + want[slot] = top; + want[slots + slot] = if (acc) |x| rows_in.sumIn(x.a, x.o, t[0..n]) else mass; + if (mag) |m| { + m[slot] = top_abs; + m[slots + slot] = mass; + } + for (0..DH) |d| { + var sum: f64 = 0; + var abs: f64 = 0; + for (from..to, 0..) |key, i| { + t[i] = @exp(sc[i] - top) * xv(p.v, (key % p.ring) * kvw + kv * DH + d); + sum += t[i]; + abs += @abs(t[i]); + } + const o = 2 * slots + slot * DH + d; + want[o] = if (acc) |x| rows_in.sumIn(x.a, x.o, t[0..n]) else sum; + if (mag) |m| m[o] = abs; + } + }; + } + + pub fn reference(p: *DecodeVec, want: []f64, mag: []f64) !void { + return p.partials(null, want, mag); + } + + pub fn simulate(p: *DecodeVec, acc: c.Accum, order: g.Order, out: []f64) !void { + return p.partials(.{ .a = acc, .o = order }, out, null); + } +}; + +/// The tensor-op accumulation over a row's keys (a window's at most), the rescales and the final division. +fn promptDepth(s: c.Shape) usize { + return s.k / 16 + 32; +} + +fn slidingPromptDepth(s: c.Shape) usize { + return @min(s.k, R.window + 64) / 16 + 32; +} + +pub fn promptSpec(l: Layer) c.Spec { + const served: []const c.Shape = switch (l) { + // a chunk mid-block past a full window, the shortest from the start, one whose window wraps the ring + .sliding => &.{ .{ .rows = 33, .n = R.heads, .k = 1233 }, .{ .rows = 17, .n = R.heads, .k = 17 }, .{ .rows = 300, .n = R.heads, .k = 3 * 3072 + 500 } }, + // a chunk mid-block, the shortest from the start, a longer one deeper in + .global => &.{ .{ .rows = 33, .n = R.heads, .k = 333 }, .{ .rows = 17, .n = R.heads, .k = 17 }, .{ .rows = 100, .n = R.heads, .k = 1300 } }, + }; + const name = if (l == .sliding) "ricursive sliding attention rc_prompt_attention" else "ricursive global attention rc_prompt_attention"; + return .{ .name = name, .accum = .f32, .depth = if (l == .sliding) slidingPromptDepth else promptDepth, .ops = 8, .out = .bf16, .tiles = .{ .{ .size = 32, .step = 1, .min = st.window_rows + 1 }, .{}, .{ .size = 64, .step = 1, .min = 1 } }, .served = served, .sums = softmax_sums }; +} + +/// rc_prompt_attention: rows query rows at positions k - rows.. over the layer's cache (a ring, or a slot a position). +pub const Prompt = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + layer: Layer, + mutation: controls.Mutation, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + ring: usize = 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, layer: Layer, mutation: controls.Mutation) !Prompt { + const plain = try rc.family.attentionSource(device, a); + defer a.free(plain); + const src = try controls.apply(a, plain, mutation); + defer a.free(src); + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .pipe = try rig_.pipeline(device, src, rc.family.attention_name), .layer = layer, .mutation = mutation }; + } + + 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(); + } + + const vt = g.hooks(Prompt); + + pub fn kernel(p: *Prompt) g.Kernel { + return controls.mark(.{ .spec = promptSpec(p.layer), .ctx = p, .vt = &vt }, p.mutation); + } + + fn p0(p: *const Prompt) usize { + return p.shape.k - p.shape.rows; + } + + /// Q padded to whole 32-row tiles (the padding NaN), the cache as the engine leaves it before the attention. + pub fn prepare(p: *Prompt, s: c.Shape, fill: g.Fill, seed: u64) !bool { + if (s.n != R.heads or s.rows == 0 or s.rows > st.prompt_rows or s.k < s.rows) return false; + p.drop(); + p.shape = s; + p.ring = ringOf(p.layer, s.k); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const pad = std.mem.alignForward(usize, s.rows, 32); + p.q = try p.a.alloc(u16, pad * R.heads * DH); + p.k = try p.a.alloc(u16, p.ring * kvw); + p.v = try p.a.alloc(u16, p.ring * kvw); + rc.random16(rng, p.q, -2, 2); + @memset(p.q[s.rows * R.heads * DH ..], nan16); + fillCache(rng, p.k, p.v, p.ring, s.k - 1, null); + const start = p.p0(); + if (fill == .nonfinite and start > 1) { // keys every row sees + p.k[(start % p.ring) * kvw + 2] = c.toBf16(std.math.nan(f32)); + p.v[((start - 1) % p.ring) * kvw + DH + 4] = c.toBf16(std.math.nan(f32)); + p.v[((start - 2) % p.ring) * kvw + 2 * DH + 9] = c.toBf16(std.math.inf(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 * R.heads * DH * 2); + return true; + } + + pub fn count(p: *Prompt) usize { + return p.shape.rows * R.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([4]i32{ @intCast(s.rows), @intCast(p.p0()), @intCast(windowOf(p.layer)), @intCast(p.ring) }, 4); + e.dispatchGroups(mtl.Size.of((s.rows + 31) / 32, R.heads, 1), mtl.Size.of(DH / 64 * 32, 1, 1)); + try p.rig.finish(pass); + } + + /// Each row's heads over the keys of its window, in float64 (or with `acc`, the sums replayed). + fn fold(p: *Prompt, acc: ?struct { a: c.Accum, o: g.Order }, want: []f64, mag: ?[]f64) !void { + const s = p.shape; + const most = @min(s.k, windowOf(p.layer)); + const buf = try p.a.alloc(f64, 3 * most); + defer p.a.free(buf); + const sc = buf[0..most]; + const w = buf[most .. 2 * most]; + const t = buf[2 * most ..]; + for (0..s.rows) |row| for (0..R.heads) |h| { + const kv = h / grp; + const pos = p.p0() + row; + const from = firstKey(p.layer, pos); + const n = pos + 1 - from; + var top: f64 = -std.math.inf(f64); // the tile's max drops a NaN score; exp2(NaN - top) carries it on + for (0..n) |i| { + sc[i] = dot(p.q, row, h, p.k, (from + i) % p.ring)[0]; + if (!std.math.isNan(sc[i])) top = @max(top, sc[i]); + } + var total: f64 = 0; + var total_abs: f64 = 0; + for (0..n) |i| { + w[i] = @exp(sc[i] - top); + total += w[i]; + total_abs += @abs(w[i]); + } + if (acc) |x| total = rows_in.sumIn(x.a, x.o, w[0..n]); + for (0..DH) |d| { + var sum: f64 = 0; + var abs: f64 = 0; + for (0..n) |i| { + t[i] = w[i] * xv(p.v, ((from + i) % p.ring) * kvw + kv * DH + d); + sum += t[i]; + abs += @abs(t[i]); + } + if (acc) |x| sum = rows_in.sumIn(x.a, x.o, t[0..n]); + const o = (row * R.heads + h) * DH + d; + want[o] = if (acc == null) sum / total else rc.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_ricursive_controls.zig b/zig/tests/contract/metal_ricursive_controls.zig new file mode 100644 index 0000000000..5061837c07 --- /dev/null +++ b/zig/tests/contract/metal_ricursive_controls.zig @@ -0,0 +1,102 @@ +//! Planted controls for Ripple's kernels: a source as the loader builds it with one fault, which one gate must report. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const tf = @import("tensorfold"); +const family = tf.ricursive.contracts; + +pub const Mutation = enum { + none, + embed_own_row, + add_norm_no_residual, + lerp_no_blend, + prep_ropes_global, + prep_next_slot, + gate_early_round, + clamp_no_lower, + decode_past_window, + prompt_past_window, + prompt_short_ring, + fp4_high_first, + fp4_scale_bias, + fp4_global_late, + fp4_dequant_unscaled, + fp4_gemm_no_global, + + pub fn name(m: Mutation) []const u8 { + return switch (m) { + .none => "", + .embed_own_row => "control: metal ricursive embed, row r reads table row r, not row ids[r]", + .add_norm_no_residual => "control: metal ricursive add_norm, the sum is not written back to the residual", + .lerp_no_blend => "control: metal ricursive add_norm, v's input takes the normed row whole, not its blend", + .prep_ropes_global => "control: metal ricursive prep, a global layer's q and k are rotated too", + .prep_next_slot => "control: metal ricursive prep, k and v land one ring slot late", + .gate_early_round => "control: metal ricursive gate, the sigmoid's exp is rounded to bf16 before the division", + .clamp_no_lower => "control: metal ricursive silu_clamp, up is clamped above only", + .decode_past_window => "control: metal ricursive decode, each row also sees the key before its window", + .prompt_past_window => "control: metal ricursive prompt attention, each row also sees the 64 keys before its window", + .prompt_short_ring => "control: metal ricursive prompt attention, keys wrap 64 slots before the ring's end", + .fp4_high_first => "control: metal ricursive fp4 matvec, each byte's nibbles read high first", + .fp4_scale_bias => "control: metal ricursive fp4 matvec, the E4M3 scale's exponent bias off by one", + .fp4_global_late => "control: metal ricursive fp4 matvec, the global applied after the sum rounds to bf16", + .fp4_dequant_unscaled => "control: metal ricursive fp4 dequant, a block's values left without its scale", + .fp4_gemm_no_global => "control: metal ricursive fp4 gemm, the sums never take the global", + }; + } + + /// The gate that must go red. + pub fn gate(m: Mutation) ?contract.gates.Gate { + return if (m == .none) null else .shapes; + } + + const Edit = struct { from: []const u8, to: []const u8 }; + + fn edit(m: Mutation) Edit { + return switch (m) { + .none => unreachable, + .embed_own_row => .{ .from = "E[(size_t)ids[p.y] * a.width + p.x]", .to = "E[(size_t)p.y * a.width + p.x]" }, + .add_norm_no_residual => .{ .from = " R[base + t + p * n] = bfloat(r[p]);\n", .to = "" }, + .lerp_no_blend => .{ .from = "rc_lerp(float(V0[base + i]), float(x), a.blend)", .to = "rc_lerp(float(V0[base + i]), float(x), 1.0f)" }, + .prep_ropes_global => .{ .from = "if (a.rope == 0) {", .to = "if (false) {" }, + .prep_next_slot => .{ .from = "const size_t slot = (a.pos + row) % a.ring;", .to = "const size_t slot = (a.pos + row + 1) % a.ring;" }, + .gate_early_round => .{ .from = "1.0f / (1.0f + precise::exp(-g))", .to = "1.0f / (1.0f + float(bfloat(precise::exp(-g))))" }, + .clamp_no_lower => .{ .from = "clamp(u0, -a.scale, a.scale)", .to = "min(u0, a.scale)" }, + .decode_past_window => .{ .from = "p + 1 > a.window ? p + 1 - a.window : 0;", .to = "p + 1 > a.window ? p - a.window : 0;" }, + .prompt_past_window => .{ .from = "col < row - (window - 1)", .to = "col < row - (window + 63)" }, + .prompt_short_ring => .{ .from = "kc + long(kpos % ring) * LD", .to = "kc + long(kpos % (ring - BKT)) * LD" }, + .fp4_high_first => .{ .from = "rc_e2m1(word >> sh), rc_e2m1(word >> (sh + 4u)), rc_e2m1(word >> (sh + 8u)), rc_e2m1(word >> (sh + 12u))", .to = "rc_e2m1(word >> (sh + 4u)), rc_e2m1(word >> sh), rc_e2m1(word >> (sh + 12u)), rc_e2m1(word >> (sh + 8u))" }, + .fp4_scale_bias => .{ .from = "int(e) - 7", .to = "int(e) - 6" }, + .fp4_global_late => .{ .from = "bfloat(sum * a.global)", .to = "bfloat(float(bfloat(sum)) * a.global)" }, + .fp4_dequant_unscaled => .{ .from = "D[at * 4 + c] = bfloat4(w[c]);", .to = "D[at * 4 + c] = bfloat4(w[c] / rc_e4m3(S[at]));" }, + .fp4_gemm_no_global => .{ .from = "acc[i][j] *= global;", .to = "acc[i][j] *= 1.0f;" }, + }; + } +}; + +/// `text` with the control's one edit (a copy either way); an edit that finds nothing proves nothing, so it errors. +pub fn apply(a: std.mem.Allocator, text: []const u8, m: Mutation) ![]u8 { + if (m == .none) return a.dupe(u8, text); + const e = m.edit(); + if (std.mem.count(u8, text, e.from) != 1) return error.ControlEditMissing; + return std.mem.replaceOwned(u8, a, text, e.from, e.to); +} + +/// The glue library as the loader compiles it (precise math functions), with the control's edit. +pub fn glueLibrary(device: mtl.Device, a: std.mem.Allocator, m: Mutation) !mtl.Library { + if (m == .none) return family.glueLibrary(device, a); + const plain = try family.glueSource(a); + defer a.free(plain); + const src = try apply(a, plain, m); + defer a.free(src); + return mtl.Library.fromSource(device, src, .{ .language = mtl.CompileOptions.mlx().language, .math_functions = .precise }); +} + +/// A spec's name and control for a probe built with `m`. +pub fn mark(k: contract.gates.Kernel, m: Mutation) contract.gates.Kernel { + var out = k; + if (m != .none) { + out.spec.name = m.name(); + out.control = m.gate(); + } + return out; +} diff --git a/zig/tests/contract/metal_ricursive_fp4.zig b/zig/tests/contract/metal_ricursive_fp4.zig new file mode 100644 index 0000000000..d86de80077 --- /dev/null +++ b/zig/tests/contract/metal_ricursive_fp4.zig @@ -0,0 +1,274 @@ +//! Ripple's NVFP4 MLP under the contract: rc_fp4_matvec's rows (bf16 of global x fp32 sum) and rc_fp4_dequant (exact). +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const sources = @import("kernel_sources"); +const rig_ = @import("rig.zig"); +const rc = @import("metal_ricursive.zig"); +const controls = @import("metal_ricursive_controls.zig"); +const c = contract.spec; +const g = contract.gates; +const Dense = contract.dense.Dense; +const family = rc.family; +const R = rc.R; + +/// fp4.metal's rc_fp4_args, as forward.zig fills it. +pub const Args = extern struct { rows: u32, n: u32, k: u32, ldy: u32, global: f32 }; + +/// A served tensor's global (ModelOpt's, 5e-4 to 7e-4); gate 4's is a power of two, so its sums stay exact. +pub const global: f32 = 6.1e-4; +pub const stress_global: f32 = 0x1p-11; + +/// E2M1's eight magnitudes; bit 3 is the sign. +const e2m1 = [8]f64{ 0, 0.5, 1, 1.5, 2, 3, 4, 6 }; + +pub fn fp4(code: u4) f64 { + const v = e2m1[code & 7]; + return if (code & 8 != 0) -v else v; +} + +/// Unsigned E4M3 (OCP's e4m3fn without its sign): bias 7, three mantissa bits, subnormals m x 2^-9, 0x7f NaN. +pub fn e4m3(code: u8) f64 { + if (code & 0x7f == 0x7f) return std.math.nan(f64); + const e: i32 = (code >> 3) & 15; + const m: f64 = @floatFromInt(code & 7); + return if (e == 0) m * 0x1p-9 else std.math.ldexp(1 + m / 8, e - 7); +} + +/// A matrix [n, k] in the release's layout: nibbles low first, 16 a block (two u32 words), one E4M3 code a block. +pub const Packed = struct { + a: std.mem.Allocator, + words: []u32, + scales: []u8, + + pub fn deinit(p: Packed) void { + p.a.free(p.words); + p.a.free(p.scales); + } +}; + +/// Random codes or gate 4's mirrored integers at scale 1, into `w` as fp4 x scale; edge scales and gate 1's NaN code. +pub fn pack(a: std.mem.Allocator, n: usize, k: usize, fill: g.Fill, seed: u64, w: []f64) !Packed { + const words = try a.alloc(u32, n * k / 8); + errdefer a.free(words); + const scales = try a.alloc(u8, n * k / 16); + @memset(words, 0); + var prng = std.Random.DefaultPrng.init(seed ^ 0xf4f4); + const rng = prng.random(); + const ints = [5]u4{ 2, 4, 5, 6, 7 }; // 1, 2, 3, 4, 6 + // every fifth row's first block: the largest scale (448), a subnormal one, zero; gate 1's rows 3, 8, ... a NaN + const h = k / 2; + for (0..n) |j| for (0..k / 16) |b| { + const special = [3]u8{ 0x7e, 0x03, 0x00 }; + const nan_code = fill == .nonfinite and b == 1 and j % 5 == 3; + scales[j * k / 16 + b] = if (fill == .stress) 0x38 else if (nan_code) 0x7f else if (b == 0 and j % 5 < 3) special[j % 5] else 0x30 + rng.uintLessThan(u8, 16); + }; + for (0..n) |j| for (0..k) |i| { + const code: u4 = if (fill != .stress) rng.int(u4) else if (i < h) ints[(i * 7 + j) % 5] else 8 | ints[((i - h) * 7 + j) % 5]; + words[(j * k + i) / 8] |= @as(u32, code) << @intCast(4 * (i % 8)); + w[j * k + i] = fp4(code) * e4m3(scales[(j * k + i) / 16]); + }; + return .{ .a = a, .words = words, .scales = scales }; +} + +/// fp4.metal as model.zig compiles it, with the control's edit. +pub fn library(device: mtl.Device, a: std.mem.Allocator, m: controls.Mutation) !mtl.Library { + const src = try controls.apply(a, sources.ricursive_fp4, m); + defer a.free(src); + return family.fp4Library(device, src); +} + +/// A lane's dots over its blocks (four dot4s each), its running sum, the simd tree, then the global. +fn matvecDepth(s: c.Shape) usize { + return s.k / 512 + 16; +} + +/// rc_fp4_matvec_ROWS: the window's MLP projection, every launch at most ROWS rows (the rows past it NaN, unread). +pub const Matvec = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: mtl.Library, + pipe: mtl.Pipeline, + per: u32, + mutation: controls.Mutation, + name: []const u8, + served: [3]c.Shape, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + global: f32 = global, + d: ?Dense = null, + w: ?Packed = null, + in: [3]rig_.Ref = undefined, // x padded to whole launches, words, scales + y: rig_.Ref = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, per: u32, mutation: controls.Mutation) !Matvec { + const lib = try library(device, a, mutation); + errdefer lib.deinit(); + const name = try std.fmt.allocPrint(a, "ricursive fp4 rc_fp4_matvec_{d}", .{per}); + errdefer a.free(name); + const fun = name[std.mem.lastIndexOfScalar(u8, name, ' ').? + 1 ..]; + const served = [3]c.Shape{ .{ .rows = per, .n = 1024, .k = R.hidden }, .{ .rows = per, .n = 1024, .k = R.intermediate }, .{ .rows = per / 2 + 1, .n = 1032, .k = R.hidden } }; + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .lib = lib, .pipe = try rc.pipelineOf(lib, device, fun), .per = per, .mutation = mutation, .name = name, .served = served }; + } + + fn drop(p: *Matvec) void { + if (p.d) |d| d.deinit(); + if (p.w) |w| w.deinit(); + p.d = null; + p.w = null; + 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 { + const tiles: [3]c.Tile = .{ .{ .size = p.per, .step = 1, .min = 1 }, .{ .size = 8, .step = 1, .min = 1 }, .{ .size = 512, .step = 16, .min = 16 } }; + return controls.mark(.{ .spec = .{ .name = p.name, .accum = .f32, .depth = matvecDepth, .ops = 1, .out = .bf16, .tiles = tiles, .served = &p.served }, .ctx = p, .vt = &vt }, p.mutation); + } + + 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 % 16 != 0) return false; + p.drop(); + p.shape = s; + p.global = if (fill == .stress) stress_global else global; + const d = try Dense.init(p.a, s, 1); + p.d = d; + d.fillX(fill, seed, 2); + p.w = try pack(p.a, s.n, s.k, fill, seed, d.w); + const padded = try p.a.alloc(u16, std.mem.alignForward(usize, s.rows, p.per) * s.k); + defer p.a.free(padded); + @memset(padded, c.toBf16(std.math.nan(f32))); + @memcpy(padded[0..d.x.len], d.x); + p.in[0] = try p.rig.input(std.mem.sliceAsBytes(padded)); + p.in[1] = try p.rig.input(std.mem.sliceAsBytes(p.w.?.words)); + p.in[2] = try p.rig.input(p.w.?.scales); + 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.in[0].buf, p.in[0].off + at * s.k * 2, 0); + e.setBuffer(p.in[1].buf, p.in[1].off, 1); + e.setBuffer(p.in[2].buf, p.in[2].off, 2); + e.setValue(Args{ .rows = here, .n = @intCast(s.n), .k = @intCast(s.k), .ldy = @intCast(s.n), .global = p.global }, 3); + e.setBuffer(p.y.buf, p.y.off + at * s.n * 2, 4); + 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 { + try p.d.?.reference(want, mag); + for (want, mag) |*w, *m| { + w.* *= p.global; + m.* *= p.global; + } + } + + pub fn simulate(p: *Matvec, acc: c.Accum, order: g.Order, out: []f64) !void { + try p.d.?.simulate(acc, order, out); + for (out) |*o| o.* *= p.global; + } +}; + +/// rc_fp4_dequant: a prompt's W as bf16, each value fp4 x scale, exact (no global: the GEMM applies it). +pub const Dequant = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: mtl.Library, + pipe: mtl.Pipeline, + mutation: controls.Mutation, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + w: ?Packed = null, + values: []f64 = &.{}, + in: [2]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, mutation: controls.Mutation) !Dequant { + const lib = try library(device, a, mutation); + errdefer lib.deinit(); + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .lib = lib, .pipe = try rc.pipelineOf(lib, device, "rc_fp4_dequant"), .mutation = mutation }; + } + + fn drop(p: *Dequant) void { + if (p.w) |w| w.deinit(); + p.w = null; + p.a.free(p.values); + p.values = &.{}; + p.rig.reset(); + } + + pub fn deinit(p: *Dequant) void { + p.drop(); + p.rig.deinit(); + p.pipe.deinit(); + p.lib.deinit(); + } + + const vt = g.hooks(Dequant); + + /// A shape's rows are W's rows (forward's n), k its columns. + pub fn kernel(p: *Dequant) g.Kernel { + const served: []const c.Shape = &.{ .{ .rows = 64, .n = 1, .k = R.hidden }, .{ .rows = 16, .n = 1, .k = R.intermediate } }; + const tiles: [3]c.Tile = .{ .{ .size = 4, .step = 1, .min = 1 }, .{}, .{ .size = 1024, .step = 16, .min = 16 } }; + return controls.mark(.{ .spec = .{ .name = "ricursive fp4 rc_fp4_dequant", .accum = .none, .depth = rc.none, .out = .bf16, .tiles = tiles, .served = served, .sums = rc.no_sums }, .ctx = p, .vt = &vt }, p.mutation); + } + + pub fn prepare(p: *Dequant, s: c.Shape, fill: g.Fill, seed: u64) !bool { + if (s.rows == 0 or s.n != 1 or s.k == 0 or s.k % 16 != 0) return false; + p.drop(); + p.shape = s; + p.values = try p.a.alloc(f64, s.rows * s.k); + p.w = try pack(p.a, s.rows, s.k, fill, seed, p.values); + p.in[0] = try p.rig.input(std.mem.sliceAsBytes(p.w.?.words)); + p.in[1] = try p.rig.input(p.w.?.scales); + p.out = try p.rig.output(s.rows * s.k * 2); + return true; + } + + pub fn count(p: *Dequant) usize { + return p.shape.rows * p.shape.k; + } + + /// forward.zig's launch: the args carry the chunk's rows and the tensor's real global, which the kernel ignores. + pub fn run(p: *Dequant) !void { + const s = p.shape; + const pass = p.rig.begin(); + const e = pass.e; + e.setPipeline(p.pipe); + 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 = 100, .n = @intCast(s.rows), .k = @intCast(s.k), .ldy = @intCast(s.rows), .global = global }, 2); + e.setBuffer(p.out.buf, p.out.off, 3); + e.dispatchThreads(mtl.Size.of(s.k / 16, s.rows, 1), mtl.Size.of(64, 4, 1)); + try p.rig.finish(pass); + } + + pub fn reference(p: *Dequant, want: []f64, mag: []f64) !void { + @memcpy(want, p.values); + for (mag, p.values) |*m, v| m.* = @abs(v); + } + + /// Every value is exact in bf16 (at most six significant bits): no error allowed. + pub fn bounds(_: *Dequant, _: []const f64, _: []const f64, allow: []f64) !void { + @memset(allow, 0); + } +}; diff --git a/zig/tests/contract/metal_ricursive_fp4_gemm.zig b/zig/tests/contract/metal_ricursive_fp4_gemm.zig new file mode 100644 index 0000000000..4829826e85 --- /dev/null +++ b/zig/tests/contract/metal_ricursive_fp4_gemm.zig @@ -0,0 +1,136 @@ +//! rc_fp4_gemm_AB under the contract: a prompt chunk by the dequantized W on the tensor units, global before rounding. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const rc = @import("metal_ricursive.zig"); +const fp4 = @import("metal_ricursive_fp4.zig"); +const controls = @import("metal_ricursive_controls.zig"); +const c = contract.spec; +const g = contract.gates; +const Dense = contract.dense.Dense; +const family = rc.family; +const R = rc.R; + +/// The tensor-op accumulation over the K tiles, the global, then the output rounding. +fn depth(s: c.Shape) usize { + return s.k / 16 + 16; +} + +/// A chunk at k 5,120 (gate and up) and the down's 40,960, whole or ragged tiles as the form takes them. +fn served(am: bool, an: bool) [2]c.Shape { + const n: usize = if (an) 1024 else 1088; + return .{ .{ .rows = if (am) 64 else 33, .n = n, .k = R.hidden }, .{ .rows = if (am) 128 else 100, .n = n, .k = R.intermediate } }; +} + +pub const Gemm = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + am: bool, + an: bool, + mutation: controls.Mutation, + name: []const u8, + served: [2]c.Shape, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + global: f32 = fp4.global, + d: ?Dense = null, + 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, mutation: controls.Mutation) !Gemm { + const plain = try family.fp4GemmSource(device, a, true); + defer a.free(plain); + const src = try controls.apply(a, plain, mutation); + defer a.free(src); + const name = try std.fmt.allocPrint(a, "ricursive fp4 gemm rc_fp4_gemm_{d}{d}", .{ @intFromBool(am), @intFromBool(an) }); + errdefer a.free(name); + const fun = name[std.mem.lastIndexOfScalar(u8, name, ' ').? + 1 ..]; + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .pipe = try rig_.pipeline(device, src, fun), .am = am, .an = an, .mutation = mutation, .name = name, .served = served(am, an) }; + } + + fn drop(p: *Gemm) void { + if (p.d) |d| d.deinit(); + p.d = null; + 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 { + const tiles: [3]c.Tile = .{ .{ .size = 64, .step = 1, .min = 1 }, .{ .size = 128, .step = 1, .min = 1 }, .{ .size = 256, .step = 256, .min = 256 } }; + return controls.mark(.{ .spec = .{ .name = p.name, .accum = .f32, .depth = depth, .ops = 1, .out = .bf16, .tiles = tiles, .served = &p.served }, .ctx = p, .vt = &vt }, p.mutation); + } + + /// 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 }; + } + + /// x as llama's GEMM probe fills it; W as rc_fp4_dequant leaves it (fp4 x scale, exact in bf16). + 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; + p.global = if (fill == .stress) fp4.stress_global else fp4.global; + const d = try Dense.init(p.a, s, 1); + p.d = d; + d.fillX(fill, seed, 2); + const w = try fp4.pack(p.a, s.n, s.k, fill, seed, d.w); + w.deinit(); + const w16 = try p.a.alloc(u16, s.n * s.k); + defer p.a.free(w16); + for (w16, d.w) |*o, v| o.* = c.toBf16(@floatCast(v)); + p.x = try p.rig.input(std.mem.sliceAsBytes(d.x)); + p.w = try p.rig.input(std.mem.sliceAsBytes(w16)); + p.y = try p.rig.output(d.outputs() * 2); + return true; + } + + pub fn count(p: *Gemm) usize { + return p.d.?.outputs(); + } + + /// forward.zig's launch of the chunk's projection. + 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.setValue(p.global, 4); + 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 { + try p.d.?.reference(want, mag); + for (want, mag) |*w, *m| { + w.* *= p.global; + m.* *= p.global; + } + } + + pub fn simulate(p: *Gemm, acc: c.Accum, order: g.Order, out: []f64) !void { + try p.d.?.simulate(acc, order, out); + for (out) |*o| o.* *= p.global; + } +}; diff --git a/zig/tests/contract/metal_ricursive_norm.zig b/zig/tests/contract/metal_ricursive_norm.zig new file mode 100644 index 0000000000..4113254625 --- /dev/null +++ b/zig/tests/contract/metal_ricursive_norm.zig @@ -0,0 +1,499 @@ +//! Ripple's norms under the contract: rc_add_norm in the forward's three forms, and rc_prep on both layer kinds. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const controls = @import("metal_ricursive_controls.zig"); +const rc = @import("metal_ricursive.zig"); +const c = contract.spec; +const g = contract.gates; +const rows_in = contract.rows; +const Range = rc.Range; +const R = rc.R; +const DH = rc.DH; +const u = rc.u; +const bf = rc.bf; +const xv = rc.xv; + +const add_flag: u32 = 1; +const lerp_flag: u32 = 2; + +/// rc_add_norm's forms: a layer's input norm with v's lerp, the attention's add, the FFN's add into the next layer. +pub const Form = enum { lerp, add, add_lerp }; + +/// The blend each form runs at: below one half and above it, torch.lerp's two branches. +fn blendOf(f: Form) f32 { + return if (f == .add_lerp) 0.69921875 else 0.30078125; +} + +fn flagsOf(f: Form) u32 { + return switch (f) { + .lerp => lerp_flag, + .add => add_flag, + .add_lerp => add_flag | lerp_flag, + }; +} + +/// 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 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 }; + +pub fn addNormSpec(f: Form) c.Spec { + const name = switch (f) { + .lerp => "ricursive glue rc_add_norm lerp", + .add => "ricursive glue rc_add_norm add", + .add_lerp => "ricursive glue rc_add_norm add lerp", + }; + const anyk: [3]c.Tile = .{ .{}, .{}, .{ .size = 256, .step = 1, .min = 1 } }; + return .{ .name = name, .accum = .f32, .depth = normDepth, .ops = 3, .mids = 3, .out = .bf16, .in_place = f != .lerp, .tiles = anyk, .served = &.{ .{ .rows = 16, .n = R.hidden, .k = R.hidden }, .{ .rows = 2048, .n = R.hidden, .k = R.hidden } }, .sums = rms_sums }; +} + +const Acc = ?struct { a: c.Accum, o: g.Order }; + +/// rc_add_norm on rows of the hidden width: R (and Y, w_post) with the add, V0 and the blend with the lerp. +pub const AddNorm = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + form: Form, + mutation: controls.Mutation, + lib: mtl.Library, + pipe: mtl.Pipeline, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + r: ?rows_in.Rows = null, // the residual (x) and w_next (w) + y: []u16 = &.{}, + post: []u16 = &.{}, + v0: []u16 = &.{}, + lo: []f64 = &.{}, + hi: []f64 = &.{}, + refs: [5]rig_.Ref = undefined, // R, Y, w_post, w_next, V0 + outs: [2]rig_.Ref = undefined, // X, VX + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, form: Form, mutation: controls.Mutation) !AddNorm { + const lib = try controls.glueLibrary(device, a, mutation); + errdefer lib.deinit(); + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .form = form, .mutation = mutation, .lib = lib, .pipe = try rc.pipelineOf(lib, device, "rc_add_norm") }; + } + + fn drop(p: *AddNorm) void { + if (p.r) |r| r.deinit(); + p.r = null; + for ([_]*[]u16{ &p.y, &p.post, &p.v0 }) |s| { + p.a.free(s.*); + s.* = &.{}; + } + for ([_]*[]f64{ &p.lo, &p.hi }) |s| { + p.a.free(s.*); + s.* = &.{}; + } + p.rig.reset(); + } + + pub fn deinit(p: *AddNorm) void { + p.drop(); + p.rig.deinit(); + p.pipe.deinit(); + p.lib.deinit(); + } + + const vt = g.hooks(AddNorm); + + pub fn kernel(p: *AddNorm) g.Kernel { + return controls.mark(.{ .spec = addNormSpec(p.form), .ctx = p, .vt = &vt }, p.mutation); + } + + fn adds(p: *const AddNorm) bool { + return p.form != .lerp; + } + + fn lerps(p: *const AddNorm) bool { + return p.form != .add; + } + + pub fn prepare(p: *AddNorm, s: c.Shape, fill: g.Fill, seed: u64) !bool { + if (s.rows == 0 or s.k == 0 or s.k > 32 * 256) return false; + p.drop(); + p.shape = s; + var prng = std.Random.DefaultPrng.init(seed ^ 0x6e6f); + const rng = prng.random(); + const w = s.k; + const n = s.rows * w; + const r = try rows_in.Rows.init(p.a, .rms, .{ .rows = s.rows, .n = w, .k = w }); + p.r = r; + // with the add, the planted values go into Y (the residual stays finite); without it, into R + r.fill(if (p.adds() and fill == .nonfinite) .random else fill, seed); + p.y = try p.a.alloc(u16, n); + p.post = try p.a.alloc(u16, w); + p.v0 = try p.a.alloc(u16, n); + // Y is zero for the sums gate (the residual's stress rows stay whole), else small + if (fill == .stress) @memset(p.y, 0) else rc.random16(rng, p.y, -0.5, 0.5); + rc.random16(rng, p.post, 0.5, 1.5); + rc.random16(rng, p.v0, -2, 2); + if (fill == .nonfinite) { + if (p.adds()) rc.plant(p.y, w, s.rows); + if (p.lerps() and s.rows > 4) for ([_]f32{ std.math.nan(f32), std.math.inf(f32) }, 4..) |f, at| { + p.v0[at * w + w / 4] = c.toBf16(f); + }; + } + p.refs[1] = try p.rig.input(std.mem.sliceAsBytes(p.y)); + p.refs[2] = try p.rig.input(std.mem.sliceAsBytes(p.post)); + p.refs[3] = try p.rig.input(std.mem.sliceAsBytes(r.w)); + p.refs[4] = try p.rig.input(std.mem.sliceAsBytes(p.v0)); + p.refs[0] = if (p.adds()) try p.rig.output(n * 2) else try p.rig.input(std.mem.sliceAsBytes(r.x)); + p.outs[0] = try p.rig.output(n * 2); + p.outs[1] = if (p.lerps()) try p.rig.output(n * 2) else try p.rig.scratch(n * 2); + p.lo = try p.a.alloc(f64, p.count()); + p.hi = try p.a.alloc(f64, p.count()); + return true; + } + + /// R (with the add), X, then VX (with the lerp). + pub fn count(p: *AddNorm) usize { + const n = p.shape.rows * p.shape.k; + return n * (1 + @as(usize, @intFromBool(p.adds())) + @intFromBool(p.lerps())); + } + + pub fn run(p: *AddNorm) !void { + const s = p.shape; + const r = p.r.?; + if (p.adds()) @memcpy(p.rig.bytes(p.refs[0], r.x.len * 2), std.mem.sliceAsBytes(r.x)); // the residual before + const pass = p.rig.begin(); + const e = pass.e; + e.setPipeline(p.pipe); + for (p.refs, 0..) |x, slot| e.setBuffer(x.buf, x.off, slot); + e.setValue(rc.Args{ .rows = @intCast(s.rows), .width = @intCast(s.k), .flags = flagsOf(p.form), .eps = rc.eps, .blend = blendOf(p.form) }, 5); + e.setBuffer(p.outs[0].buf, p.outs[0].off, 6); + e.setBuffer(p.outs[1].buf, p.outs[1].off, 7); + e.dispatchGroups(mtl.Size.of(s.rows, 1, 1), mtl.Size.of(256, 1, 1)); + try p.rig.finish(pass); + } + + /// Row `row` at the release's rounding points into want, its range into lo and hi; `acc` replays the two sums. + fn row(p: *AddNorm, at: usize, acc: Acc, want: []f64, t: []f64, rr: []Range) void { + const s = p.shape; + const w = s.k; + const n = s.rows * w; + const base = at * w; + const r = p.r.?; + const depth = normDepth(s); + var o: usize = 0; // where X's outputs start + for (rr, 0..) |*x, i| x.* = rc.point(xv(r.x, base + i)); + if (p.adds()) { + o = n; + for (t, 0..) |*x, i| x.* = xv(p.y, base + i) * xv(p.y, base + i); + const ss = if (acc) |x| rows_in.sumIn(x.a, x.o, t) else exactSum(t); + const inv = 1.0 / @sqrt(ss / @as(f64, @floatFromInt(w)) + rc.eps); + const ir = rc.invRms(rc.point(ss), w, depth); + for (0..w) |i| { + const y = xv(p.y, base + i); + const pw = xv(p.post, i); + const res = xv(r.x, base + i); + want[base + i] = bf(res + bf(pw * bf(y * inv))); + const normed = rc.round(rc.mul(rc.point(pw), rc.round(rc.widen(rc.mul(rc.point(y), ir), u)))); + rr[i] = rc.round(rc.widen(rc.add(rc.point(res), normed), u)); + } + } + var ss: f64 = 0; + var sr: Range = .{ 0, 0 }; + for (0..w) |i| { + const v = if (p.adds()) want[base + i] else xv(r.x, base + i); + t[i] = v * v; + sr = rc.add(sr, rc.square(rr[i])); + } + ss = if (acc) |x| rows_in.sumIn(x.a, x.o, t) else exactSum(t); + const inv = 1.0 / @sqrt(ss / @as(f64, @floatFromInt(w)) + rc.eps); + const ir = rc.invRms(sr, w, depth); + const blend: f64 = blendOf(p.form); + for (0..w) |i| { + const v = if (p.adds()) want[base + i] else xv(r.x, base + i); + const nw = xv(r.w, i); + const x = bf(nw * bf(v * inv)); + want[o + base + i] = x; + const xr = rc.round(rc.mul(rc.point(nw), rc.round(rc.widen(rc.mul(rr[i], ir), u)))); + if (p.adds()) { + p.lo[base + i] = rr[i][0]; + p.hi[base + i] = rr[i][1]; + } + p.lo[o + base + i] = xr[0]; + p.hi[o + base + i] = xr[1]; + if (!p.lerps()) continue; + const s0 = xv(p.v0, base + i); + const vo = o + n + base + i; + want[vo] = bf(lerp(s0, x, blend)); + // lerp rises with its end: the range's ends, each with the difference's rounding and the fma's + const ends = Range{ lerp(s0, xr[0], blend), lerp(s0, xr[1], blend) }; + const err = u * (blend * @max(@abs(xr[0] - s0), @abs(xr[1] - s0)) + @max(@abs(ends[0]), @abs(ends[1]))); + const vr = rc.round(.{ ends[0] - err, ends[1] + err }); + p.lo[vo] = vr[0]; + p.hi[vo] = vr[1]; + } + } + + fn rows(p: *AddNorm, acc: Acc, want: []f64) !void { + const w = p.shape.k; + const t = try p.a.alloc(f64, w); + defer p.a.free(t); + const rr = try p.a.alloc(Range, w); + defer p.a.free(rr); + for (0..p.shape.rows) |at| p.row(at, acc, want, t, rr); + } + + pub fn reference(p: *AddNorm, want: []f64, mag: []f64) !void { + try p.rows(null, want); + for (mag, want) |*m, v| m.* = @abs(v); + } + + pub fn simulate(p: *AddNorm, acc: c.Accum, order: g.Order, out: []f64) !void { + const lo = try p.a.dupe(f64, p.lo); + defer p.a.free(lo); + const hi = try p.a.dupe(f64, p.hi); + defer p.a.free(hi); + try p.rows(.{ .a = acc, .o = order }, out); + @memcpy(p.lo, lo); + @memcpy(p.hi, hi); + } + + pub fn bounds(p: *AddNorm, want: []const f64, _: []const f64, allow: []f64) !void { + rc.allowRange(p.lo, p.hi, want, allow); + } + + pub fn intervals(p: *AddNorm, low: []f64, high: []f64) !void { + @memcpy(low, p.lo); + @memcpy(high, p.hi); + } +}; + +/// torch.lerp's vectorized form (the kernel's rc_lerp) in float64. +fn lerp(s: f64, e: f64, w: f64) f64 { + return if (@abs(w) < 0.5) s + w * (e - s) else e - (1 - w) * (e - s); +} + +fn exactSum(t: []const f64) f64 { + var s: f64 = 0; + for (t) |x| s += x; + return s; +} + +/// rc_prep's layers: sliding (RoPE, a ring) and global (no RoPE, a slot a position). +pub const Layer = enum { sliding, global }; + +/// A lane's four squares, then the simd tree. +fn prepDepth(_: c.Shape) usize { + return 12; +} + +pub fn prepSpec(l: Layer) c.Spec { + const units = R.heads + R.kv_heads; + const stride = (R.heads + 2 * R.kv_heads) * DH; + const name = if (l == .sliding) "ricursive glue rc_prep sliding" else "ricursive glue rc_prep global"; + return .{ .name = name, .accum = .f32, .depth = prepDepth, .ops = 6, .mids = 3, .out = .bf16, .served = &.{ .{ .rows = 16, .n = units, .k = stride }, .{ .rows = 2048, .n = units, .k = stride } }, .sums = rms_sums }; +} + +/// glue.metal's rc_prep_args, as forward.zig fills it. +const PArgs = extern struct { rows: u32, heads: u32, kv_heads: u32, pos: u32, stride: u32, ring: u32, rope: u32, eps: f32 }; + +/// rc_prep on a window (deep in a stream, its slots wrapping the ring) or a prompt chunk. +pub const Prep = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + layer: Layer, + mutation: controls.Mutation, + lib: mtl.Library, + pipe: mtl.Pipeline, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + pos: usize = 0, + ring: usize = 0, + x: []u16 = &.{}, + qw: []u16 = &.{}, + kw: []u16 = &.{}, + inv: []f32 = &.{}, + lo: []f64 = &.{}, + hi: []f64 = &.{}, + in: [4]rig_.Ref = undefined, + out: [3]rig_.Ref = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, layer: Layer, mutation: controls.Mutation) !Prep { + const lib = try controls.glueLibrary(device, a, mutation); + errdefer lib.deinit(); + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .layer = layer, .mutation = mutation, .lib = lib, .pipe = try rc.pipelineOf(lib, device, "rc_prep") }; + } + + fn drop(p: *Prep) void { + for ([_]*[]u16{ &p.x, &p.qw, &p.kw }) |s| { + p.a.free(s.*); + s.* = &.{}; + } + for ([_]*[]f64{ &p.lo, &p.hi }) |s| { + p.a.free(s.*); + s.* = &.{}; + } + p.a.free(p.inv); + p.inv = &.{}; + p.rig.reset(); + } + + pub fn deinit(p: *Prep) void { + p.drop(); + p.rig.deinit(); + p.pipe.deinit(); + p.lib.deinit(); + } + + const vt = g.hooks(Prep); + + pub fn kernel(p: *Prep) g.Kernel { + return controls.mark(.{ .spec = prepSpec(p.layer), .ctx = p, .vt = &vt }, p.mutation); + } + + const kvw = R.kv_heads * DH; + + pub fn prepare(p: *Prep, s: c.Shape, fill: g.Fill, seed: u64) !bool { + const stride = (R.heads + 2 * R.kv_heads) * DH; + if (s.rows == 0 or s.k != stride or s.n != R.heads + R.kv_heads) return false; + p.drop(); + p.shape = s; + const window = s.rows <= 16; + // sliding: a window across the ring's wrap, a chunk deep in a stream; global: a window at 700, a first chunk + if (p.layer == .sliding) { + p.ring = rc.family.sliding_ring; + p.pos = if (window) 3 * p.ring - 7 else 100003; + } else { + p.pos = if (window) 700 else 0; + p.ring = std.mem.alignForward(usize, p.pos + s.rows, 64); + } + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + p.x = try p.a.alloc(u16, s.rows * stride); + p.qw = try p.a.alloc(u16, DH); + p.kw = try p.a.alloc(u16, DH); + if (fill == .stress) { // integers whose 128 squares overflow fp16 and lose bits in bf16 + for (p.x) |*v| v.* = c.toBf16(@floatFromInt(@as(i32, @intCast(rng.uintLessThan(u32, 129))) - 64)); + } else rc.random16(rng, p.x, -3, 3); + rc.random16(rng, p.qw, 0.5, 1.5); + rc.random16(rng, p.kw, 0.5, 1.5); + if (fill == .nonfinite) rc.plant(p.x, stride, s.rows); + p.inv = try p.a.alloc(f32, DH / 2); + for (p.inv, 0..) |*f, i| { + const pw: f32 = @floatCast(std.math.pow(f64, 10000, @as(f64, @floatFromInt(2 * i)) / DH)); + f.* = 1.0 / pw; // config.zig's inverseFrequencies + } + 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.in[2] = try p.rig.input(std.mem.sliceAsBytes(p.qw)); + p.in[3] = try p.rig.input(std.mem.sliceAsBytes(p.kw)); + p.out[0] = try p.rig.output(s.rows * R.heads * DH * 2); + p.out[1] = try p.rig.output(p.ring * kvw * 2); + p.out[2] = try p.rig.output(p.ring * kvw * 2); + p.lo = try p.a.alloc(f64, p.count()); + p.hi = try p.a.alloc(f64, p.count()); + return true; + } + + /// QN, then the whole key and value rings. + pub fn count(p: *Prep) usize { + return p.shape.rows * R.heads * DH + 2 * p.ring * kvw; + } + + pub fn run(p: *Prep) !void { + const s = p.shape; + // slots no row of this call reaches: unwritten by the kernel, zero in the reference + for (p.out[1..]) |o| @memset(p.rig.bytes(o, p.ring * kvw * 2), 0); + const pass = p.rig.begin(); + const e = pass.e; + e.setPipeline(p.pipe); + for (p.in[0..4], 0..) |r, slot| e.setBuffer(r.buf, r.off, slot); + const a = PArgs{ .rows = @intCast(s.rows), .heads = R.heads, .kv_heads = R.kv_heads, .pos = @intCast(p.pos), .stride = @intCast(s.k), .ring = @intCast(p.ring), .rope = @intFromBool(p.layer == .sliding), .eps = rc.eps }; + e.setValue(a, 4); + for (p.out, 5..) |o, slot| e.setBuffer(o.buf, o.off, slot); + e.dispatchGroups(mtl.Size.of(R.heads + R.kv_heads, s.rows, 1), mtl.Size.of(32, 1, 1)); + try p.rig.finish(pass); + } + + /// One head's norm and rotation at row `r`: want at the rounding points, its range; `acc` replays the sum. + fn head(p: *Prep, r: usize, unit: usize, acc: Acc, want: []f64) void { + const q = unit < R.heads; + const h = if (q) unit else unit - R.heads; + const src = r * p.shape.k + (if (q) 0 else R.heads * DH) + h * DH; + const qn = p.shape.rows * R.heads * DH; + const slot = (p.pos + r) % p.ring; + const dst = if (q) (r * R.heads + h) * DH else qn + slot * kvw + h * DH; + const w = if (q) p.qw else p.kw; + var t: [DH]f64 = undefined; + for (&t, 0..) |*x, d| x.* = xv(p.x, src + d) * xv(p.x, src + d); + const ss = if (acc) |x| rows_in.sumIn(x.a, x.o, &t) else exactSum(&t); + const inv = 1.0 / @sqrt(ss / DH + rc.eps); + const ir = rc.invRms(rc.point(exactSum(&t)), DH, 12); + var n: [DH]f64 = undefined; + var nr: [DH]Range = undefined; + for (0..DH) |d| { + const x = xv(p.x, src + d); + const wd = xv(w, d); + n[d] = bf(wd * bf(x * inv)); + nr[d] = rc.round(rc.mul(rc.point(wd), rc.round(rc.widen(rc.mul(rc.point(x), ir), u)))); + } + const position: f32 = @floatFromInt(p.pos + r); + for (0..DH) |d| { + var v = n[d]; + var e = nr[d]; + if (p.layer == .sliding) { + const angle: f64 = position * p.inv[d % (DH / 2)]; // the kernel's fp32 angle + const cs = @cos(angle); + const sn = @sin(angle); + const cr = rc.round(rc.widen(rc.point(cs), 4 * u)); // precise cos and sin: 4 ulps + const sr = rc.round(rc.widen(rc.point(sn), 4 * u)); + const other = if (d < DH / 2) -n[d + DH / 2] else n[d - DH / 2]; + const or_ = if (d < DH / 2) Range{ -nr[d + DH / 2][1], -nr[d + DH / 2][0] } else nr[d - DH / 2]; + v = bf(n[d] * bf(cs) + other * bf(sn)); + const t1 = rc.mul(nr[d], cr); + const t2 = rc.mul(or_, sr); + const err = 2 * u * (@max(@abs(t1[0]), @abs(t1[1])) + @max(@abs(t2[0]), @abs(t2[1]))); + e = rc.round(.{ t1[0] + t2[0] - err, t1[1] + t2[1] + err }); + } + want[dst + d] = v; + p.lo[dst + d] = e[0]; + p.hi[dst + d] = e[1]; + } + if (q) return; + for (0..DH) |d| { + const at = qn + p.ring * kvw + slot * kvw + h * DH + d; + want[at] = xv(p.x, r * p.shape.k + (R.heads + R.kv_heads) * DH + h * DH + d); + p.lo[at] = want[at]; + p.hi[at] = want[at]; + } + } + + fn all(p: *Prep, acc: Acc, want: []f64) void { + const qn = p.shape.rows * R.heads * DH; + @memset(want[qn..], 0); + @memset(p.lo[qn..], 0); + @memset(p.hi[qn..], 0); + for (0..p.shape.rows) |r| for (0..R.heads + R.kv_heads) |unit| p.head(r, unit, acc, want); + } + + pub fn reference(p: *Prep, want: []f64, mag: []f64) !void { + p.all(null, want); + for (mag, want) |*m, v| m.* = @abs(v); + } + + pub fn simulate(p: *Prep, acc: c.Accum, order: g.Order, out: []f64) !void { + const lo = try p.a.dupe(f64, p.lo); + defer p.a.free(lo); + const hi = try p.a.dupe(f64, p.hi); + defer p.a.free(hi); + p.all(.{ .a = acc, .o = order }, out); + @memcpy(p.lo, lo); + @memcpy(p.hi, hi); + } + + pub fn bounds(p: *Prep, want: []const f64, _: []const f64, allow: []f64) !void { + rc.allowRange(p.lo, p.hi, want, allow); + } + + pub fn intervals(p: *Prep, low: []f64, high: []f64) !void { + @memcpy(low, p.lo); + @memcpy(high, p.hi); + } +}; diff --git a/zig/tests/contract/metal_ricursive_sample.zig b/zig/tests/contract/metal_ricursive_sample.zig new file mode 100644 index 0000000000..6f7ed3dcbf --- /dev/null +++ b/zig/tests/contract/metal_ricursive_sample.zig @@ -0,0 +1,128 @@ +//! Nemotron's tf_sample_full as Ripple draws with it, held to its stated rule: each draw equals lanes' gpu_full twin. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const contract = @import("contract"); +const sources = @import("kernel_sources"); +const rig_ = @import("rig.zig"); +const rc = @import("metal_ricursive.zig"); +const c = contract.spec; +const g = contract.gates; +const full = tf.lanes.gpu_full; +const Sampling = tf.lanes.Sampling; + +const rule = "each row's draw equals lanes' gpu_full twin: the plan (top_k, top_p, min_p) then the keyed race"; + +/// The settings rows cycle through: top_k off with a nucleus, a wide top_k, a narrow one at T 0.7, and min_p. +const settings = [_]Sampling{ + .{ .seed = 11, .temperature = 1.0, .top_k = 0, .top_p = 0.9 }, + .{ .seed = 12, .temperature = 1.0, .top_k = 3000, .top_p = 0.5 }, + .{ .seed = 13, .temperature = 0.7, .top_k = 20, .top_p = 0.95 }, + .{ .seed = 14, .temperature = 1.0, .top_k = 0, .top_p = 1.0, .min_p = 0.01 }, +}; +const position: u32 = 40000; + +pub fn spec() c.Spec { + return .{ .name = "ricursive sample tf_sample_full", .accum = .none, .depth = rc.none, .out = .raw, .nan = .{ .stated = rule }, .tiles = .{ .{}, .{}, .{ .size = 1024, .step = 1, .min = 1 } }, .served = &.{.{ .rows = 8, .n = 1, .k = rc.R.vocab }}, .sums = rc.no_sums }; +} + +pub const Sample = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + logits: []f32 = &.{}, + in: [5]rig_.Ref = undefined, // logits, seeds, positions, cfgs, caps + out: rig_.Ref = undefined, + msg: [160]u8 = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue) !Sample { + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .pipe = try rig_.pipeline(device, sources.nemotron_sample, "tf_sample_full") }; + } + + pub fn deinit(p: *Sample) void { + p.a.free(p.logits); + p.rig.deinit(); + p.pipe.deinit(); + } + + const vt = g.hooks(Sample); + + pub fn kernel(p: *Sample) g.Kernel { + return .{ .spec = spec(), .ctx = p, .vt = &vt }; + } + + /// A row of logits in [-3, 0] in bf16 (every fourth row a peak over 2,000 tokens); planted NaN and infinities. + pub fn prepare(p: *Sample, s: c.Shape, fill: g.Fill, seed: u64) !bool { + if (s.rows == 0 or s.k == 0 or s.n != 1) return false; + p.a.free(p.logits); + p.logits = &.{}; + p.rig.reset(); + p.shape = s; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + p.logits = try p.a.alloc(f32, s.rows * s.k); + const words = try p.a.alloc(u16, s.rows * s.k); + defer p.a.free(words); + for (words, 0..) |*w, i| { + const peak = (i / s.k) % 4 == 3 and i % s.k < 2000; + w.* = c.toBf16(if (peak) 0 else -3 * rng.float(f32)); + } + if (fill == .nonfinite) rc.plant(words, s.k, s.rows); + for (p.logits, words) |*l, w| l.* = c.fromBf16(w); + var seeds: [64]u32 = undefined; + var positions: [32]u32 = undefined; + var cfgs: [128]f32 = undefined; + var caps: [32]u32 = undefined; + if (s.rows > positions.len) return false; + for (0..s.rows) |r| { + const ru = tf.lanes.gpu_rule.rule(settings[r % settings.len]); + seeds[2 * r ..][0..2].* = ru.seedHalves(); + positions[r] = position + @as(u32, @intCast(r)); + cfgs[4 * r ..][0..4].* = ru.cfg(); + caps[r] = ru.top_k; + } + p.in[0] = try p.rig.input(std.mem.sliceAsBytes(words)); + p.in[1] = try p.rig.input(std.mem.sliceAsBytes(seeds[0 .. 2 * s.rows])); + p.in[2] = try p.rig.input(std.mem.sliceAsBytes(positions[0..s.rows])); + p.in[3] = try p.rig.input(std.mem.sliceAsBytes(cfgs[0 .. 4 * s.rows])); + p.in[4] = try p.rig.input(std.mem.sliceAsBytes(caps[0..s.rows])); + p.out = try p.rig.output(s.rows * 4); + return true; + } + + pub fn count(p: *Sample) usize { + return p.shape.rows * 4; + } + + pub fn run(p: *Sample) !void { + 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, 5); + e.setBuffer(p.out.buf, p.out.off, 6); // the id map, unread without MAP (forward.zig binds the same) + e.setValue(@as(u32, @intCast(p.shape.k)), 7); + e.dispatchGroups(mtl.Size.of(p.shape.rows, 1, 1), mtl.Size.of(1024, 1, 1)); + try p.rig.finish(pass); + } + + pub fn reference(_: *Sample, want: []f64, mag: []f64) !void { + @memset(want, 0); + @memset(mag, 0); + } + + /// The twin's draw for every row (gates 1 and 3). + pub fn check(p: *Sample, _: g.Gate) !?[]const u8 { + const k = p.shape.k; + const ids = std.mem.bytesAsSlice(u32, p.rig.bytes(p.out, p.shape.rows * 4)); + for (0..p.shape.rows) |r| { + const s = settings[r % settings.len]; + const row = p.logits[r * k ..][0..k]; + const plan = try full.plan(p.a, row, s); + const want = full.race(row, plan, s.seed, position + @as(u32, @intCast(r)), null); + if (ids[r] != want) return try std.fmt.bufPrint(&p.msg, "row {d}: drew {d} where the twin draws {d}", .{ r, ids[r], want }); + } + return null; + } +}; diff --git a/zig/tests/cuda/contract_limits.zig b/zig/tests/cuda/contract_limits.zig index a1778cb2b0..eb7ea623ff 100644 --- a/zig/tests/cuda/contract_limits.zig +++ b/zig/tests/cuda/contract_limits.zig @@ -26,7 +26,7 @@ pub fn image(name: []const u8) error{UnknownContractImage}![]const u8 { if (std.mem.eql(u8, name, "h3")) return &h3_image.bytes; inline for (.{ "qwen4_projection", "qwen4_hc", "qwen4_gdn", "qwen4_ple", "gemma4_ple" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); - inline for (.{ "tree_round", "act_rows", "gdn_prep", "qk_rows", "norms", "moe", "experts", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "qmm_prefill", "lane_gemv", "prompt16", "qmmf", "ck_experts", "nvfp4_experts", "nvfp4_experts_tile", "nvfp4_experts32", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "glm_glue", "glm_ops", "glm_index", "nemotron_mamba", "sample", "qwen35", "qwen35_bf16_rows", "qwen35_keep", "qwen35_kv8", "w4a8_prefill", "qwen35_draft", "ck_lane", "gdn_train", "train_q4", "train_attn", "flashnext_train", "train", "attn_train", "qwen27_train", "tree_attention", "block_attention", "window_attention", "window_pages", "moe_glue", "dense_rows", "qmm", "flashnext_fact", "flashnext_hcf", "flashnext_fexp", "experts_rows", "nemotron_ops", "minimax_m3", "minimax_m3_attn", "moss_image" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); + inline for (.{ "tree_round", "act_rows", "gdn_prep", "qk_rows", "norms", "moe", "experts", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "qmm_prefill", "lane_gemv", "prompt16", "qmmf", "ck_experts", "nvfp4_experts", "nvfp4_experts_tile", "nvfp4_experts32", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "glm_glue", "glm_ops", "glm_index", "nemotron_mamba", "sample", "qwen35", "qwen35_bf16_rows", "qwen35_keep", "qwen35_kv8", "w4a8_prefill", "qwen35_draft", "ck_lane", "gdn_train", "train_q4", "train_attn", "flashnext_train", "train", "attn_train", "qwen27_train", "tree_attention", "block_attention", "window_attention", "window_pages", "moe_glue", "dense_rows", "qmm", "flashnext_fact", "flashnext_hcf", "flashnext_fexp", "experts_rows", "nemotron_ops", "minimax_m3", "minimax_m3_attn", "moss_image", "ricursive", "ricursive_prep" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); if (std.mem.eql(u8, name, "contract_controls")) return if (options.controls) &controls.bytes else &.{}; return error.UnknownContractImage; } diff --git a/zig/tests/cuda/contract_ricursive.zig b/zig/tests/cuda/contract_ricursive.zig new file mode 100644 index 0000000000..4888f49060 --- /dev/null +++ b/zig/tests/cuda/contract_ricursive.zig @@ -0,0 +1,57 @@ +//! Ripple's CUDA kernels under the contract, with controls (ricursive-contract): glue, attention prep, attention. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rc = @import("ricursive_cuda"); +const check = @import("check.zig"); +const glue = @import("contract_ricursive_glue.zig"); +const prep = @import("contract_ricursive_prep.zig"); +const attn = @import("contract_ricursive_attn.zig"); + +const g = core.contract.gates; + +/// Every registration, on the heap: the gates hold pointers to them. +pub const Set = struct { + gpa: std.mem.Allocator, + k: rc.kernels.Kernels, + glue: glue.Set = .{}, + prep: *prep.Set = undefined, // its own image, and contract_controls' tf_ctl_rc_prep + attn: *attn.Set = undefined, // window_attention's image (WindowAttention owns it), tree_attention's (GlobalDecode) + + /// The set, its kernels appended to `list`. + pub fn init(gpu: check.Gpu, s: cuda.Stream, a: std.mem.Allocator, list: *std.ArrayList(g.Kernel)) !*Set { + const set = try gpu.gpa.create(Set); + errdefer gpu.gpa.destroy(set); + set.* = .{ .gpa = gpu.gpa, .k = try rc.kernels.Kernels.load(gpu.d) }; + errdefer set.k.unload(); + try set.glue.init(gpu, s, &set.k, a, list); + errdefer set.glue.deinit(); + set.prep = try prep.Set.init(gpu, s, a, list); + errdefer set.prep.deinit(); + set.attn = try attn.Set.init(gpu, s, a, list); + return set; + } + + pub fn deinit(set: *Set) void { + set.attn.deinit(); + set.prep.deinit(); + set.glue.deinit(); + set.k.unload(); + set.gpa.destroy(set); + } +}; + +/// Ripple's kernels' gates 1-4 and their controls alone. +pub fn run(gpu: check.Gpu) !void { + const a = gpu.gpa; + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var list: std.ArrayList(g.Kernel) = .empty; + defer list.deinit(a); + const set = try Set.init(gpu, s, a, &list); + defer set.deinit(); + var t: g.Tally = .{}; + try g.runAll(a, list.items, 0x726970, &t); + std.debug.print("SUMMARY ricursive: {d} kernels, {d} pass gates 1-4, {d} fail; {d} controls: {d} red as planted, {d} missed\n", .{ t.kernels, t.passed, t.failed, t.controls, t.caught, t.missed }); + try check.expect(t.problems() == 0, "ricursive-contract: {d} gate failures, {d} controls missed", .{ t.failed, t.missed }); +} diff --git a/zig/tests/cuda/contract_ricursive_attn.zig b/zig/tests/cuda/contract_ricursive_attn.zig new file mode 100644 index 0000000000..cfb1f4d925 --- /dev/null +++ b/zig/tests/cuda/contract_ricursive_attn.zig @@ -0,0 +1,304 @@ +//! Ripple's attention under the contract on CUDA (80/8 heads): prompt rows, sliding decode, global decode over pages. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rc = @import("ricursive_cuda").attn; +const check = @import("check.zig"); +const rig_ = @import("contract_rig.zig"); +const attn = @import("contract_attn.zig"); +const wk = @import("contract_window.zig"); + +const c = core.contract.spec; +const g = core.contract.gates; +const wa = cuda.window_attention; + +/// Ripple-65B's attention: 80 query heads over 8 key heads, head dim 128, a 1,024-key window in a ring of 3,072. +const h = 80; +const hk = 8; +const window = 1024; +const ring = 3072; +/// Global layers' pages: tree_attention's span, as its decode reads them. +const page = 512; + +/// Sliding chunks: before the window fills, across its edge, and past the ring's wrap (key 9,216 at slot 0). +const sliding_served = [_]c.Shape{ .{ .rows = 333, .n = h, .k = 0 }, .{ .rows = 1000, .n = h, .k = 700 }, .{ .rows = 2048, .n = h, .k = 8216 } }; +/// Global chunks (every key): from the start, past 1,500 keys, and a whole 2,048-row pass. +const global_served = [_]c.Shape{ .{ .rows = 333, .n = h, .k = 0 }, .{ .rows = 1000, .n = h, .k = 1500 }, .{ .rows = 2048, .n = h, .k = 0 } }; +/// A control's one shape: past the window and the ring's wrap, where each fault binds. +const control_served = [_]c.Shape{.{ .rows = 333, .n = h, .k = 8216 }}; +/// The chunk control's: chunks of 333 after the first. +const chunk_served = [_]c.Shape{.{ .rows = 1000, .n = h, .k = 700 }}; + +/// A control's fault: score scale doubled, window doubled, ring one slot short, or later chunks' scale one ulp up. +pub const Fault = enum { none, scale, window, ring, late_scale }; + +/// A prompt registration: the window test's call, cut into `chunk` rows a launch (0: one launch), and its fault. +pub const Cfg = struct { base: wk.Cfg, chunk: usize = 0, fault: Fault = .none }; + +fn slidingCfg(name: []const u8, served: []const c.Shape) wk.Cfg { + return .{ .name = name, .h = h, .hk = hk, .d = 128, .win = window, .ring = ring, .prompt = true, .served = served }; +} + +fn globalCfg(name: []const u8, pages: usize, served: []const c.Shape) wk.Cfg { + return .{ .name = name, .h = h, .hk = hk, .d = 128, .win = 0, .ring = 0, .prompt = true, .pages = pages, .served = served }; +} + +pub const prompts = [_]Cfg{ + .{ .base = slidingCfg("cuda ricursive prompt wattn_kernel<128, 8, 2, 8>, 80/8 heads, window 1024, ring 3072", &sliding_served) }, + .{ .base = globalCfg("cuda ricursive prompt wattn_kernel<128, 8, 2, 8>, 80/8 heads, every key (global layers)", 0, &global_served) }, + .{ .base = globalCfg("cuda ricursive prompt tf_window_prompt_h2_pg, 80/8 heads, every key, pages of 512", page, &global_served) }, + // a row's bits do not depend on its chunk: cuts of the same rows bit-equal to one call + .{ .base = slidingCfg("cuda ricursive prompt in chunks of 1,000, window 1024: one call's bits", &sliding_served), .chunk = 1000 }, + .{ .base = slidingCfg("cuda ricursive prompt in chunks of 333, window 1024: one call's bits", &sliding_served), .chunk = 333 }, + .{ .base = globalCfg("cuda ricursive prompt in chunks of 333, every key: one call's bits", 0, &global_served), .chunk = 333 }, + .{ .base = slidingCfg("control: Ripple's prompt rows with the score scaled twice", &control_served), .fault = .scale }, + .{ .base = slidingCfg("control: Ripple's prompt rows with a window of 2,048 keys", &control_served), .fault = .window }, + .{ .base = slidingCfg("control: Ripple's prompt rows with a ring one slot short (3,071)", &control_served), .fault = .ring }, + .{ .base = .{ .name = "control: Ripple's paged prompt rows whose first page id is a spare page's (all NaN)", .h = h, .hk = hk, .d = 128, .win = 0, .ring = 0, .prompt = true, .pages = page, .stale = true, .served = &control_served } }, + .{ .base = slidingCfg("control: Ripple's prompt in chunks of 333, the later chunks' scale one ulp up", &chunk_served), .chunk = 333, .fault = .late_scale }, +}; + +/// Sliding decode through tf_window_decode_tc as it is: one-row steps past the ring's wrap, and 16-row windows. +pub const decodes = [_]wk.Cfg{ + .{ .name = "cuda ricursive decode tf_window_decode_tc, 80/8 heads, window 1024, ring 3072", .h = h, .hk = hk, .d = 128, .win = window, .ring = ring, .streams = 2, .served = &.{ .{ .rows = 1, .n = h, .k = 9500 }, .{ .rows = 16, .n = h, .k = 3000 } } }, +}; + +/// The P x V chain of a sliding prompt's last row: at most its window's keys, a tile at a time, then the fold. +fn slidingChain(s: c.Shape) usize { + return @min(s.k + s.rows, window) + 64 + 64; +} + +/// A global prompt's last row: every key to its own. +fn globalChain(s: c.Shape) usize { + return s.k + s.rows + 64 + 64; +} + +/// Global decode: two pages' rows across a page's edge, rows 1 and 16 from the first key, a page's first key, deep. +const gdec_served = [_]c.Shape{ .{ .rows = 16, .n = h, .k = 505 }, .{ .rows = 1, .n = h, .k = 0 }, .{ .rows = 16, .n = h, .k = 0 }, .{ .rows = 1, .n = h, .k = 512 }, .{ .rows = 1, .n = h, .k = 30000 }, .{ .rows = 16, .n = h, .k = 30000 } }; +/// A control's shape: 16 rows across a page's edge. +const gdec_control = [_]c.Shape{.{ .rows = 16, .n = h, .k = 505 }}; +/// The short control's: rows over their first keys, where one key is much of a row's weight. +const gdec_short = [_]c.Shape{.{ .rows = 16, .n = h, .k = 1 }}; +/// GlobalDecode's rows and keys in the contract: the served shapes' most. +const global_rows = 16; +const global_context = 30000 + 16; + +/// Gate 4 at 30,000 keys: past about 8,000 keys the stress values' sums leave fp32's 24 bits, so no order is exact. +const deep_sums: c.Sums = .{ .bf16 = 0, .exact = false, .why = "past 8,000 keys the stress values' sums leave fp32's 24 bits: gate 4 holds the float64 bound, which fp16 sums overflow" }; + +/// A global decode control's fault: score scale doubled, each row one key short, or 64 query heads (group 8). +pub const GlobalFault = enum { none, scale, short, heads }; + +/// A global decode registration: the window test's paged prompt rows (row r at k + r over keys 0..k + r). +pub const GlobalCfg = struct { base: wk.Cfg, fault: GlobalFault = .none }; + +fn pagedCfg(name: []const u8, served: []const c.Shape, stale: bool) wk.Cfg { + return .{ .name = name, .h = h, .hk = hk, .d = 128, .win = 0, .ring = 0, .prompt = true, .pages = page, .stale = stale, .served = served }; +} + +pub const globals = [_]GlobalCfg{ + .{ .base = pagedCfg("cuda ricursive global decode tf_tree_*_pg + tf_tree_merge, 80/8 heads, every key, pages of 512", &gdec_served, false) }, + .{ .base = pagedCfg("control: Ripple's global decode with the score scaled twice", &gdec_control, false), .fault = .scale }, + .{ .base = pagedCfg("control: Ripple's global decode whose first page id is a spare page's (all NaN)", &gdec_control, true) }, + .{ .base = pagedCfg("control: Ripple's global decode with each row one key short", &gdec_short, false), .fault = .short }, + .{ .base = pagedCfg("control: Ripple's global decode with 64 query heads over the 8 key heads (group 8)", &gdec_control, false), .fault = .heads }, +}; + +/// Decode rows: rows, query heads and the first row's position; tree_attention's 16-pair tiles and 512-key groups. +pub fn globalSpec(f: GlobalCfg) c.Spec { + const b = f.base; + return .{ .name = b.name, .accum = .f32, .depth = &globalChain, .ops = 1, .out = .bf16, .tiles = .{ .{ .size = 16 }, .{}, .{ .size = page, .min = 0 } }, .served = b.served, .sums = deep_sums }; +} + +/// Prompt rows: chunk rows, query heads and the chunk's start; 64-row blocks, 64-key tiles. +pub fn spec(f: Cfg) c.Spec { + const b = f.base; + return .{ .name = b.name, .accum = .f32, .depth = if (b.win > 0) &slidingChain else &globalChain, .ops = 1, .out = .bf16, .replaces = if (f.chunk > 0) "wattn_kernel<128, 8, 2, 8> in one call" else null, .tiles = .{ .{ .size = rc.rows_a_block }, .{}, .{ .size = 64, .min = 0 } }, .served = b.served, .sums = .{ .why = attn.why } }; +} + +pub const Prompt = struct { + rig: *rig_.Rig, // the window test's own + w: wk.Window, + k: *const rc.Kernels, + f: Cfg, + base: u64 = 0, // chunked: the one call's output (scratch) + + pub fn deinit(self: *Prompt) void { + self.w.deinit(); + } + + pub fn kernel(self: *Prompt) g.Kernel { + const control = self.f.fault != .none or self.f.base.stale; + return .{ .spec = spec(self.f), .ctx = self, .vt = &vt, .control = if (control) .shapes else null }; + } + + const vt = g.hooks(Prompt); + + pub fn prepare(self: *Prompt, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (!try self.w.prepare(shape, fill, seed)) return false; + if (self.f.chunk > 0) self.base = try self.rig.scratch(self.w.count() * 2); + return true; + } + + /// The launch's window: the registration's, with its control's fault (`later`: a chunk after the first). + fn window(self: *const Prompt, later: bool) wa.Window { + const f = self.f.base; + const scale: f32 = @floatCast(1 / @sqrt(@as(f64, @floatFromInt(f.d)))); // the window test's, so its reference's + const pg: ?wa.Pages = if (f.pages > 0) .{ .table = self.w.table, .page_elems = 2 * f.pages * f.hk * f.d, .tokens = f.pages } else null; + var w = if (f.win > 0) rc.sliding(f.h, f.hk, f.win, f.ring, scale) else rc.global(f.h, f.hk, scale, pg); + switch (self.f.fault) { + .none => {}, + .scale => w.scale *= 2, + .window => w.win *= 2, + .ring => w.ring -= 1, + .late_scale => if (later) { + w.scale = std.math.nextAfter(f32, w.scale, 1); + }, + } + return w; + } + + /// Rows [r, r + m) of the prepared call into `out`, at positions k + r, ... + fn launch(self: *Prompt, r: usize, m: usize, out: u64) !void { + const at = r * self.f.base.h * self.f.base.d * 2; // a row's heads, bf16, in q and out + const w = &self.w; + try rc.prompt(self.k, self.rig.s, w.dev[0] + at, w.dev[1], w.dev[2], out + at, w.shape.k + r, m, self.window(r > 0)); + } + + pub fn run(self: *Prompt) !void { + const n = self.w.shape.rows; + const step = if (self.f.chunk > 0) self.f.chunk else n; + var r: usize = 0; + while (r < n) : (r += step) try self.launch(r, @min(step, n - r), self.w.out); + try self.rig.s.synchronize(); + } + + /// The chunks' baseline: the prepared rows in one call. + pub fn replaced(self: *Prompt) ![]const u8 { + const n = self.w.shape.rows; + try rc.prompt(self.k, self.rig.s, self.w.dev[0], self.w.dev[1], self.w.dev[2], self.base, self.w.shape.k, n, self.window(false)); + return try self.rig.bytes(self.base, self.w.count() * 2); + } + + pub fn count(self: *Prompt) usize { + return self.w.count(); + } + + pub fn check(self: *Prompt, gate: g.Gate) !?[]const u8 { + return self.w.check(gate); + } + + pub fn reference(self: *Prompt, want: []f64, mag: []f64) !void { + return self.w.reference(want, mag); + } + + pub fn simulate(self: *Prompt, acc: c.Accum, order: g.Order, out: []f64) !void { + return self.w.simulate(acc, order, out); + } +}; + +/// Global decode through rc.decodeGlobal, over the window test's paged prompt inputs and its float64 reference. +pub const Global = struct { + rig: *rig_.Rig, // the window test's own + w: wk.Window, + gd: *const rc.GlobalDecode, + f: GlobalCfg, + + pub fn deinit(self: *Global) void { + self.w.deinit(); + } + + pub fn kernel(self: *Global) g.Kernel { + const control = self.f.fault != .none or self.f.base.stale; + return .{ .spec = globalSpec(self.f), .ctx = self, .vt = &vt, .control = if (control) .shapes else null }; + } + + const vt = g.hooks(Global); + + pub fn prepare(self: *Global, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows > self.gd.max_rows or shape.k + shape.rows > self.gd.context) return false; + return self.w.prepare(shape, fill, seed); + } + + pub fn run(self: *Global) !void { + const f = self.f.base; + const w = &self.w; + var a: rc.DecodeArgs = .{ + .q = w.dev[0], + .kc = w.dev[1], + .vc = w.dev[2], + .out = w.out, + .pos = w.shape.k, + .rows = w.shape.rows, + .pages = .{ .table = w.table, .page_elems = 2 * f.pages * f.hk * f.d, .tokens = f.pages }, + .scale = @floatCast(1 / @sqrt(@as(f64, @floatFromInt(f.d)))), // the window test's, so its reference's + }; + switch (self.f.fault) { + .none, .heads => {}, + .scale => a.scale *= 2, + .short => a.pos -|= 1, + } + if (self.f.fault == .heads) try rc.decodeHeads(self.gd, w.rig.s, a, 64) else try rc.decodeGlobal(self.gd, w.rig.s, a); + try w.rig.s.synchronize(); + } + + pub fn count(self: *Global) usize { + return self.w.count(); + } + + pub fn check(self: *Global, gate: g.Gate) !?[]const u8 { + return self.w.check(gate); + } + + pub fn reference(self: *Global, want: []f64, mag: []f64) !void { + return self.w.reference(want, mag); + } + + pub fn simulate(self: *Global, acc: c.Accum, order: g.Order, out: []f64) !void { + return self.w.simulate(acc, order, out); + } +}; + +/// Every registration, on the heap: the gates hold pointers to them, and each prompt's rig pointer into itself. +pub const Set = struct { + gpa: std.mem.Allocator, + win: wa.WindowAttention, + k: rc.Kernels = undefined, + gd: rc.GlobalDecode = undefined, + prompts: [prompts.len]Prompt = undefined, + decodes: [decodes.len]wk.Window = undefined, + globals: [globals.len]Global = undefined, + + /// The set, its kernels appended to `list`. + pub fn init(gpu: check.Gpu, s: cuda.Stream, a: std.mem.Allocator, list: *std.ArrayList(g.Kernel)) !*Set { + const set = try gpu.gpa.create(Set); + errdefer gpu.gpa.destroy(set); + set.* = .{ .gpa = gpu.gpa, .win = try wa.WindowAttention.load(gpu.d) }; + errdefer set.win.unload(); + set.k = try rc.Kernels.init(&set.win); + set.gd = try rc.GlobalDecode.init(gpu.d, gpu.gpa, global_rows, global_context); + errdefer set.gd.deinit(); + for (&set.prompts, prompts) |*x, f| { + x.* = .{ .rig = undefined, .w = .{ .rig = .{ .gpu = gpu, .s = s }, .f = f.base, .g = &set.win }, .k = &set.k, .f = f }; + x.rig = &x.w.rig; + } + for (&set.decodes, decodes) |*x, f| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .f = f, .g = &set.win }; + for (&set.globals, globals) |*x, f| { + x.* = .{ .rig = undefined, .w = .{ .rig = .{ .gpu = gpu, .s = s }, .f = f.base, .g = &set.win }, .gd = &set.gd, .f = f }; + x.rig = &x.w.rig; + } + for (&set.prompts) |*x| try list.append(a, x.kernel()); + for (&set.decodes) |*x| try list.append(a, x.kernel()); + for (&set.globals) |*x| try list.append(a, x.kernel()); + return set; + } + + pub fn deinit(set: *Set) void { + for (&set.prompts) |*x| x.deinit(); + for (&set.decodes) |*x| x.deinit(); + for (&set.globals) |*x| x.deinit(); + set.gd.deinit(); + set.win.unload(); + set.gpa.destroy(set); + } +}; diff --git a/zig/tests/cuda/contract_ricursive_glue.zig b/zig/tests/cuda/contract_ricursive_glue.zig new file mode 100644 index 0000000000..e51c8fd883 --- /dev/null +++ b/zig/tests/cuda/contract_ricursive_glue.zig @@ -0,0 +1,492 @@ +//! Ripple's CUDA glue under the contract: tf_rc_embed, tf_rc_add_norm (four forms), tf_rc_gate, tf_rc_silu_clamp. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rc = @import("ricursive_cuda"); +const check = @import("check.zig"); +const rig_ = @import("contract_rig.zig"); +const ref = @import("contract_minimax_glue_ref.zig"); // the family-neutral Gated wrapper and fills + +const c = core.contract.spec; +const g = core.contract.gates; +const kern = rc.kernels; +const Iv = ref.Iv; +const Sums = ref.Sums; +const Kernels = kern.Kernels; +const Function = cuda.Function; + +const hidden = 5120; +const query = 80 * 128; // the output gate's rows: query heads x head dim +const intermediate = 40960; +const vocab = 100352; +const eps: f32 = 1e-6; +const embed_scale: f32 = 71.55417527999327; +const limit: f32 = 20; // the FFN's clamp + +/// The forward's row counts at width n: decode 1 to 16, odd tails and prompt chunks up to 2,048. +fn rowsAt(n: usize, k: usize) [7]c.Shape { + var out: [7]c.Shape = undefined; + for ([_]usize{ 16, 1, 3, 7, 333, 1000, 2048 }, &out) |r, *s| s.* = .{ .rows = r, .n = n, .k = k }; + return out; +} + +const embed_rows = rowsAt(hidden, 1); +const norm_rows = rowsAt(hidden, hidden); +const gate_rows = rowsAt(query, 1); +const silu_rows = rowsAt(intermediate, 2 * intermediate); +const elementwise: c.Tile = .{ .size = 256, .step = 1, .min = 1 }; + +pub fn embedSpec(control: bool) c.Spec { + const name = if (control) "control: cuda ricursive tf_rc_embed, the scale rounded to bf16 (71.5)" else "cuda ricursive tf_rc_embed (fp32 scale, rounded once; R and V0)"; + return .{ .name = name, .accum = .none, .depth = ref.noSums, .out = .bf16, .tiles = .{ .{}, elementwise, .{} }, .served = &embed_rows }; +} + +/// R and V0 = bf16(E[ids] * scale) over the whole vocabulary's table: row 0 reads its last row. +pub const Embed = struct { + kern: *const Kernels, + f: Function, + /// control: the scale rounded to bf16 before the product + ctl: bool = false, + shape: c.Shape = undefined, + table: []u16 = &.{}, + ids: []u32 = &.{}, + dev: [2]u64 = @splat(0), // the table, the ids + out: [2]u64 = @splat(0), // R, V0 + + pub fn spec(self: *const Embed) c.Spec { + return embedSpec(self.ctl); + } + + pub fn control(self: *const Embed) ?g.Gate { + return if (self.ctl) .shapes else null; + } + + pub fn free(self: *Embed, a: std.mem.Allocator) void { + a.free(self.table); + a.free(self.ids); + self.table = &.{}; + self.ids = &.{}; + } + + /// The rows the ids read in [-1, 1), the rest zero (never read); non-finite: NaN, +Inf, -Inf in the first rows. + pub fn prepare(self: *Embed, a: std.mem.Allocator, rig: *rig_.Rig, shape: c.Shape, fill: g.Fill, rng: std.Random) !bool { + if (shape.rows == 0 or shape.n == 0) return false; + self.free(a); + self.shape = shape; + const n = shape.n; + self.ids = try a.alloc(u32, shape.rows); + for (self.ids, 0..) |*id, r| id.* = if (r == 0) vocab - 1 else rng.uintLessThan(u32, vocab); + self.table = try a.alloc(u16, vocab * n); + @memset(self.table, 0); + for (self.ids) |id| for (self.table[@as(usize, id) * n ..][0..n]) |*v| { + v.* = c.toBf16(rng.float(f32) * 2 - 1); + }; + if (fill == .nonfinite) for ([_]f32{ ref.nan, ref.inf, -ref.inf }, 0..) |x, r| if (r < shape.rows) { + self.table[@as(usize, self.ids[r]) * n + (7 * r + 1) % n] = c.toBf16(x); + }; + self.dev = .{ try ref.input(rig, self.table), try ref.input(rig, self.ids) }; + self.out = .{ try rig.output(shape.rows * n * 2), try rig.output(shape.rows * n * 2) }; + return true; + } + + pub fn run(self: *Embed, rig: *rig_.Rig) !void { + try kern.embedBy(self.f, rig.s, self.dev[1], self.dev[0], self.out[0], self.out[1], self.shape.rows, self.shape.n, embed_scale); + } + + pub fn count(self: *const Embed) usize { + return 2 * self.shape.rows * self.shape.n; + } + + pub fn solve(self: *const Embed, _: std.mem.Allocator, _: Sums, out: []Iv) !void { + const n = self.shape.n; + const half = self.shape.rows * n; + for (0..self.shape.rows) |r| for (0..n) |i| { + const v = Iv.of(ref.bf(self.table[@as(usize, self.ids[r]) * n + i])).mul(Iv.of(embed_scale)).bf(); + out[r * n + i] = v; + out[half + r * n + i] = v; + }; + } +}; + +/// tf_rc_add_norm's forms in the forward: a first norm (plain or with v's lerp), an add, an add into the next lerp. +pub const Form = enum { plain, lerp, add, add_lerp }; + +/// A control's dropped rounding point (ricursive.cuh's add_norm bits). +pub const Point = enum { y_row, y_out, resid, x_row, x_lerp }; + +fn adds(f: Form) bool { + return f == .add or f == .add_lerp; +} + +fn lerps(f: Form) bool { + return f == .lerp or f == .add_lerp; +} + +/// The blend each lerp form runs at: below one half and above it, torch.lerp's two branches. +fn blendOf(f: Form) f32 { + return if (f == .add_lerp) 0.69921875 else 0.30078125; +} + +/// A thread's fma chain (8 values a chunk, a chunk per 2,048 of the row), the warp's 5 shuffles and the block's 5. +fn normDepth(s: c.Shape) usize { + return 8 * ((s.n + 2047) / 2048) + 10; +} + +pub fn normSpec(f: Form, drop: ?Point) c.Spec { + const name = if (drop) |p| switch (p) { + .y_row => "control: cuda ricursive tf_rc_add_norm without bf16(y / rms)", + .y_out => "control: cuda ricursive tf_rc_add_norm without bf16(w_post y / rms)", + .resid => "control: cuda ricursive tf_rc_add_norm, the next norm reads the residual before its bf16", + .x_row => "control: cuda ricursive tf_rc_add_norm without bf16(r / rms)", + .x_lerp => "control: cuda ricursive tf_rc_add_norm, the lerp reads X before its bf16", + } else switch (f) { + .plain => "cuda ricursive tf_rc_add_norm (X = norm(R))", + .lerp => "cuda ricursive tf_rc_add_norm lerp (X and lerp(V0, X), blend below one half)", + .add => "cuda ricursive tf_rc_add_norm add (R += norm(Y), in place, then X)", + .add_lerp => "cuda ricursive tf_rc_add_norm add lerp (R += norm(Y), X, lerp(V0, X) above one half)", + }; + const chunks: c.Tile = .{ .size = 2048, .step = 8, .min = 8 }; // a thread's chunk per 2,048 of the row + return .{ .name = name, .accum = .f32, .depth = normDepth, .out = .bf16, .in_place = adds(f), .tiles = .{ .{}, chunks, .{} }, .served = &norm_rows, .sums = .{} }; +} + +/// Residual rows as bf16: [-1, 1), one in 64 twenty times larger; stress: integers in [-32, 32]. +fn residual(a: std.mem.Allocator, rng: std.Random, n: usize, fill: g.Fill) ![]u16 { + const x = try a.alloc(u16, n); + for (x, 0..) |*v, i| v.* = c.toBf16(if (fill == .stress) ref.int(rng, -32, 32) else (rng.float(f32) * 2 - 1) * (if (i % 64 == 7) @as(f32, 20) else 1)); + return x; +} + +/// NaN, +Inf and -Inf at known places: a row each, row 3 holding both infinities. +fn plant(x: []u16, rows: usize, n: usize) void { + const at = [_][2]usize{ .{ 0, n / 3 }, .{ 1, n / 2 }, .{ 2, n - 1 }, .{ 3, 0 }, .{ 3, n - 1 } }; + const put = [_]f32{ ref.nan, ref.inf, -ref.inf, ref.inf, -ref.inf }; + for (at, put) |p, v| if (p[0] < rows) { + x[p[0] * n + p[1]] = c.toBf16(v); + }; +} + +/// 1 / rms of a row's values (ranges): an fp32 sum of squares over `depth` roundings, then 1 / sqrt(mean + eps). +fn invRms(s: Sums, x: []const Iv, depth: usize, scratch: []f64, sq: []Iv) Iv { + for (sq, x) |*q, v| { + const lo = @min(@abs(v.lo), @abs(v.hi)); + const hi = @max(@abs(v.lo), @abs(v.hi)); + q.* = .{ .m = v.m * v.m, .lo = if (v.lo <= 0 and v.hi >= 0) 0 else lo * lo, .hi = hi * hi }; + } + const total = s.sum(sq, depth, scratch); + return total.div(Iv.of(@floatFromInt(x.len))).add(Iv.of(eps)).rsqrt(); +} + +/// torch.lerp as the kernel's fma: w (e - s) + s below one half, (w - 1)(e - s) + e above it; rounded to bf16 once. +fn lerp(s: Iv, e: Iv, w: f32) Iv { + const d = e.sub(s); + return (if (@abs(w) < 0.5) Iv.of(w).mul(d).add(s) else Iv.of(w - 1).mul(d).add(e)).bf(); +} + +/// tf_rc_add_norm on hidden-wide rows: R, w_next; Y and w_post with the add; V0 and the blend with the lerp. +pub const AddNorm = struct { + kern: *const Kernels, + f: Function, + form: Form, + drop: ?Point = null, + shape: c.Shape = undefined, + r: []u16 = &.{}, // R before the launch + y: []u16 = &.{}, + post: []u16 = &.{}, + next: []u16 = &.{}, + v0: []u16 = &.{}, + p: kern.NormRows = undefined, + + pub const sums = {}; + + pub fn spec(self: *const AddNorm) c.Spec { + return normSpec(self.form, self.drop); + } + + pub fn control(self: *const AddNorm) ?g.Gate { + return if (self.drop != null) .shapes else null; + } + + pub fn free(self: *AddNorm, a: std.mem.Allocator) void { + for ([_]*[]u16{ &self.r, &self.y, &self.post, &self.next, &self.v0 }) |s| { + a.free(s.*); + s.* = &.{}; + } + } + + /// Y in [-0.5, 0.5), weights [0.5, 1.5), V0 [-2, 2); stress: Y zero; non-finite: in Y (add) or R, V0 rows 4-6. + pub fn prepare(self: *AddNorm, a: std.mem.Allocator, rig: *rig_.Rig, shape: c.Shape, fill: g.Fill, rng: std.Random) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.n % 8 != 0 or shape.n > kern.most_width) return false; + self.free(a); + self.shape = shape; + const rows = shape.rows; + const n = shape.n; + const add = adds(self.form); + const mix = lerps(self.form); + self.r = try residual(a, rng, rows * n, fill); + self.y = try ref.bfs(a, rng, if (add) rows * n else 0, -0.5, 0.5); + if (fill == .stress) @memset(self.y, 0); + self.post = try ref.bfs(a, rng, if (add) n else 0, 0.5, 1.5); + self.next = try ref.bfs(a, rng, n, 0.5, 1.5); + self.v0 = try ref.bfs(a, rng, if (mix) rows * n else 0, -2, 2); + if (fill == .nonfinite) { + plant(if (add) self.y else self.r, rows, n); + if (mix) for ([_]f32{ ref.nan, ref.inf, -ref.inf }, 4..) |x, row| if (row < rows) { + self.v0[row * n + n / 4] = c.toBf16(x); + }; + } + const bytes = rows * n * 2; + self.p = .{ .resid = 0, .next = try ref.input(rig, self.next), .x = 0, .rows = rows, .width = n, .eps = eps, .blend = if (mix) blendOf(self.form) else 0 }; + if (add) { + self.p.branch = try ref.input(rig, self.y); + self.p.post = try ref.input(rig, self.post); + self.p.resid = try rig.output(bytes); // in place: restored before each launch + } else self.p.resid = try ref.input(rig, self.r); + if (mix) self.p.v0 = try ref.input(rig, self.v0); + self.p.x = try rig.output(bytes); + if (mix) self.p.vx = try rig.output(bytes); + return true; + } + + pub fn run(self: *AddNorm, rig: *rig_.Rig) !void { + if (adds(self.form)) try ref.restore(rig, self.p.resid, std.mem.sliceAsBytes(self.r)); // read and written in place: every run the same pass + try kern.addNormBy(self.f, rig.s, self.p); + } + + pub fn count(self: *const AddNorm) usize { + const outs = 1 + @as(usize, @intFromBool(adds(self.form))) + @intFromBool(lerps(self.form)); + return outs * self.shape.rows * self.shape.n; + } + + /// The outputs in the order prepare allocates them: R (with the add), X, VX (with the lerp). + pub fn solve(self: *const AddNorm, a: std.mem.Allocator, s: Sums, out: []Iv) !void { + const n = self.shape.n; + const total = self.shape.rows * n; + const buf = try a.alloc(Iv, 3 * n); + defer a.free(buf); + const scratch = try a.alloc(f64, n); + defer a.free(scratch); + const depth = normDepth(self.shape); + const r = buf[0..n]; + const y = buf[n .. 2 * n]; + const sq = buf[2 * n ..]; + const lead: usize = if (adds(self.form)) total else 0; + const out_r = out[0..lead]; + const out_x = out[lead..][0..total]; + const out_vx = out[lead + total ..]; + for (0..self.shape.rows) |row| { + const o = row * n; + for (r, self.r[o..][0..n]) |*v, h| v.* = Iv.of(ref.bf(h)); + if (adds(self.form)) { + for (y, self.y[o..][0..n]) |*v, h| v.* = Iv.of(ref.bf(h)); + const inv_y = invRms(s, y, depth, scratch, sq); + for (r, y, self.post) |*v, yv, w| v.* = v.add(Iv.of(ref.bf(w)).mul(yv.mul(inv_y).bf()).bf()).bf(); + @memcpy(out_r[o..][0..n], r); + } + const inv = invRms(s, r, depth, scratch, sq); + for (out_x[o..][0..n], r, self.next) |*x, v, w| x.* = Iv.of(ref.bf(w)).mul(v.mul(inv).bf()).bf(); + if (lerps(self.form)) for (out_vx[o..][0..n], out_x[o..][0..n], self.v0[o..][0..n]) |*vx, x, h| { + vx.* = lerp(Iv.of(ref.bf(h)), x, blendOf(self.form)); + }; + } + } +}; + +/// The control kernel without point `p` (ricursive.cu's tf_rc_ctl_add_norm_*). +fn dropped(k: *const Kernels, p: Point) !Function { + return switch (p) { + .y_row => k.mod.function("tf_rc_ctl_add_norm_y_row"), + .y_out => k.mod.function("tf_rc_ctl_add_norm_y_out"), + .resid => k.mod.function("tf_rc_ctl_add_norm_resid"), + .x_row => k.mod.function("tf_rc_ctl_add_norm_x_row"), + .x_lerp => k.mod.function("tf_rc_ctl_add_norm_x_lerp"), + }; +} + +pub fn gateSpec(control: bool) c.Spec { + const name = if (control) "control: cuda ricursive tf_rc_gate without bf16(sigmoid(g))" else "cuda ricursive tf_rc_gate (fp32 sigmoid then bf16, in place)"; + return .{ .name = name, .accum = .none, .depth = ref.noSums, .out = .bf16, .in_place = true, .tiles = .{ .{}, elementwise, .{} }, .served = &gate_rows }; +} + +/// A = bf16(A * bf16(sigmoid(G))) in place: A in [-4, 4), G in [-12, 12); non-finite: G's NaN and infinities, then A's. +pub const Gate = struct { + kern: *const Kernels, + f: Function, + /// control: the sigmoid left in fp32 + ctl: bool = false, + shape: c.Shape = undefined, + act: []u16 = &.{}, + gv: []u16 = &.{}, + dev_a: u64 = 0, + dev_g: u64 = 0, + + pub fn spec(self: *const Gate) c.Spec { + return gateSpec(self.ctl); + } + + pub fn control(self: *const Gate) ?g.Gate { + return if (self.ctl) .shapes else null; + } + + pub fn free(self: *Gate, a: std.mem.Allocator) void { + a.free(self.act); + a.free(self.gv); + self.act = &.{}; + self.gv = &.{}; + } + + pub fn prepare(self: *Gate, a: std.mem.Allocator, rig: *rig_.Rig, shape: c.Shape, fill: g.Fill, rng: std.Random) !bool { + if (shape.rows == 0 or shape.n == 0) return false; + self.free(a); + self.shape = shape; + const n = shape.n; + self.act = try ref.bfs(a, rng, shape.rows * n, -4, 4); + self.gv = try ref.bfs(a, rng, shape.rows * n, -12, 12); + if (fill == .nonfinite) { + const put = [_]struct { usize, bool, f32 }{ .{ 0, true, ref.nan }, .{ 1, true, ref.inf }, .{ 2, true, -ref.inf }, .{ 3, false, ref.nan }, .{ 4, false, ref.inf }, .{ 5, false, -ref.inf } }; + for (put) |p| if (p[0] < shape.rows) { + (if (p[1]) self.gv else self.act)[p[0] * n + p[0] + 1] = c.toBf16(p[2]); + }; + } + self.dev_g = try ref.input(rig, self.gv); + self.dev_a = try rig.output(shape.rows * n * 2); // in place: restored before each launch + return true; + } + + pub fn run(self: *Gate, rig: *rig_.Rig) !void { + try ref.restore(rig, self.dev_a, std.mem.sliceAsBytes(self.act)); + try kern.gateBy(self.f, rig.s, self.dev_a, self.dev_g, self.shape.rows, self.shape.n); + } + + pub fn count(self: *const Gate) usize { + return self.shape.rows * self.shape.n; + } + + pub fn solve(self: *const Gate, _: std.mem.Allocator, _: Sums, out: []Iv) !void { + for (out, self.act, self.gv) |*o, x, gh| o.* = Iv.of(ref.bf(x)).mul(Iv.of(ref.bf(gh)).sigmoid().bf()).bf(); + } +}; + +/// A SiLU-clamp control: the silu left in fp32, or up clamped above only. +pub const SiluDrop = enum { round, lower }; + +const silu_nan = "torch.clamp's bounds: a NaN gate or up gives NaN, +Inf reads as the limit, a -Inf up as -limit; a -Inf gate gives NaN (-Inf / Inf)"; + +pub fn siluSpec(drop: ?SiluDrop) c.Spec { + const name = if (drop) |d| switch (d) { + .round => "control: cuda ricursive tf_rc_silu_clamp without bf16(silu(g))", + .lower => "control: cuda ricursive tf_rc_silu_clamp, up clamped above only", + } else "cuda ricursive tf_rc_silu_clamp (stacked gate | up, clamp 20)"; + return .{ .name = name, .accum = .none, .depth = ref.noSums, .out = .bf16, .nan = .{ .stated = silu_nan }, .tiles = .{ .{}, elementwise, .{} }, .served = &silu_rows }; +} + +/// bf16(bf16(silu(min(g, 20))) * clamp(u, -20, 20)) from stacked gate | up rows, as the up-gate projection writes them. +pub const SiluClamp = struct { + kern: *const Kernels, + f: Function, + drop: ?SiluDrop = null, + shape: c.Shape = undefined, + gate: []u16 = &.{}, + up: []u16 = &.{}, + dev: u64 = 0, + out: u64 = 0, + + pub fn spec(self: *const SiluClamp) c.Spec { + return siluSpec(self.drop); + } + + pub fn control(self: *const SiluClamp) ?g.Gate { + return if (self.drop != null) .shapes else null; + } + + pub fn free(self: *SiluClamp, a: std.mem.Allocator) void { + a.free(self.gate); + a.free(self.up); + self.gate = &.{}; + self.up = &.{}; + } + + /// Gate and up in [-30, 30) (past the clamp both ways); non-finite: NaN, +Inf and -Inf in each of gate and up. + pub fn prepare(self: *SiluClamp, a: std.mem.Allocator, rig: *rig_.Rig, shape: c.Shape, fill: g.Fill, rng: std.Random) !bool { + if (shape.rows == 0 or shape.n == 0) return false; + self.free(a); + self.shape = shape; + const rows = shape.rows; + const n = shape.n; + self.gate = try ref.bfs(a, rng, rows * n, -30, 30); + self.up = try ref.bfs(a, rng, rows * n, -30, 30); + if (fill == .nonfinite) { + const put = [_]struct { usize, bool, f32 }{ .{ 0, true, ref.nan }, .{ 1, true, ref.inf }, .{ 2, true, -ref.inf }, .{ 3, false, ref.nan }, .{ 4, false, ref.inf }, .{ 5, false, -ref.inf } }; + for (put) |p| if (p[0] < rows) { + (if (p[1]) self.gate else self.up)[p[0] * n + p[0] + 1] = c.toBf16(p[2]); + }; + } + const both = try a.alloc(u16, 2 * rows * n); + defer a.free(both); + for (0..rows) |r| { + @memcpy(both[2 * r * n ..][0..n], self.gate[r * n ..][0..n]); + @memcpy(both[(2 * r + 1) * n ..][0..n], self.up[r * n ..][0..n]); + } + self.dev = try ref.input(rig, both); + self.out = try rig.output(rows * n * 2); + return true; + } + + pub fn run(self: *SiluClamp, rig: *rig_.Rig) !void { + try kern.siluClampBy(self.f, rig.s, self.dev, self.out, self.shape.rows, self.shape.n, limit); + } + + pub fn count(self: *const SiluClamp) usize { + return self.shape.rows * self.shape.n; + } + + pub fn solve(self: *const SiluClamp, _: std.mem.Allocator, _: Sums, out: []Iv) !void { + for (out, self.gate, self.up) |*o, gh, uh| { + const gv = Iv.of(ref.bf(gh)).atMost(limit); + const u = Iv.of(ref.bf(uh)).atMost(limit).atLeast(-limit); + const silu = gv.div(Iv.of(1).add(gv.neg().exp())).bf(); + o.* = silu.mul(u).bf(); + } + } +}; + +/// The glue's registrations, held where the gates point at them (the caller's heap set). +pub const Set = struct { + embeds: [2]ref.Gated(Embed) = undefined, + norms: [4]ref.Gated(AddNorm) = undefined, + norm_ctls: [5]ref.Gated(AddNorm) = undefined, + gates: [2]ref.Gated(Gate) = undefined, + silus: [3]ref.Gated(SiluClamp) = undefined, + + /// Every registration over `k`, appended to `list` (each control after the kernel it breaks). + pub fn init(set: *Set, gpu: check.Gpu, s: cuda.Stream, k: *const Kernels, a: std.mem.Allocator, list: *std.ArrayList(g.Kernel)) !void { + const rig: rig_.Rig = .{ .gpu = gpu, .s = s }; + set.embeds = .{ + .{ .rig = rig, .k = .{ .kern = k, .f = k.embed } }, + .{ .rig = rig, .k = .{ .kern = k, .f = try k.mod.function("tf_rc_ctl_embed_bf16_scale"), .ctl = true } }, + }; + for (&set.norms, std.enums.values(Form)) |*x, f| x.* = .{ .rig = rig, .k = .{ .kern = k, .f = k.add_norm, .form = f } }; + for (&set.norm_ctls, std.enums.values(Point)) |*x, p| x.* = .{ .rig = rig, .k = .{ .kern = k, .f = try dropped(k, p), .form = .add_lerp, .drop = p } }; + set.gates = .{ + .{ .rig = rig, .k = .{ .kern = k, .f = k.gate } }, + .{ .rig = rig, .k = .{ .kern = k, .f = try k.mod.function("tf_rc_ctl_gate_sigmoid"), .ctl = true } }, + }; + set.silus = .{ + .{ .rig = rig, .k = .{ .kern = k, .f = k.silu_clamp } }, + .{ .rig = rig, .k = .{ .kern = k, .f = try k.mod.function("tf_rc_ctl_silu_round"), .drop = .round } }, + .{ .rig = rig, .k = .{ .kern = k, .f = try k.mod.function("tf_rc_ctl_silu_no_lower"), .drop = .lower } }, + }; + for (&set.embeds) |*x| try list.append(a, x.kernel()); + for (&set.norms) |*x| try list.append(a, x.kernel()); + for (&set.norm_ctls) |*x| try list.append(a, x.kernel()); + for (&set.gates) |*x| try list.append(a, x.kernel()); + for (&set.silus) |*x| try list.append(a, x.kernel()); + } + + pub fn deinit(set: *Set) void { + for (&set.embeds) |*x| x.deinit(); + for (&set.norms) |*x| x.deinit(); + for (&set.norm_ctls) |*x| x.deinit(); + for (&set.gates) |*x| x.deinit(); + for (&set.silus) |*x| x.deinit(); + } +}; diff --git a/zig/tests/cuda/contract_ricursive_prep.zig b/zig/tests/cuda/contract_ricursive_prep.zig new file mode 100644 index 0000000000..34e68d2c7d --- /dev/null +++ b/zig/tests/cuda/contract_ricursive_prep.zig @@ -0,0 +1,560 @@ +//! Ripple's CUDA attention prep under the contract: tf_rc_prep with its controls, tf_rc_prep_rows bit-equal to it. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rp = @import("ricursive_cuda").prep; +const check = @import("check.zig"); +const rig_ = @import("contract_rig.zig"); +const ref = @import("contract_minimax_glue_ref.zig"); + +const c = core.contract.spec; +const g = core.contract.gates; +const Iv = ref.Iv; +const Sums = ref.Sums; + +const controls = struct { + pub const bytes align(16) = @embedFile("fatbin_contract_controls").*; +}; + +const D = 128; +const heads = 80; +const kv_heads = 8; +const kvw = kv_heads * D; +const qk_stride = (heads + kv_heads) * D; // a fused q | k projection's row; v is its own projection +const width = (heads + 2 * kv_heads) * D; // a row's outputs: its q, k and v heads +const eps: f32 = 1e-6; +const theta: f64 = 10000; +const sliding_ring = 3072; // state.zig: the window less 1 plus a 2,048-row chunk, to 128s +const page = 512; // positions a global layer's page holds: tree_attention's span, as Kolibri-1's pages + +/// ricursive_prep.cuh's rounding points. +const norm_row: i32 = 1; +const norm_out: i32 = 2; +const cos_pt: i32 = 4; +const sin_pt: i32 = 8; +const out_pt: i32 = 16; +const all: i32 = 31; + +/// Layer kinds: sliding (RoPE, the 3,072-slot ring), global (NoPE, a ring that never wraps), global in 512-row pages. +pub const Form = enum { sliding, global, paged }; + +/// A registration: the served kernel (keep null), or a control keeping every rounding point but one. +pub const Entry = struct { name: []const u8, form: Form, keep: ?i32 = null }; + +/// NoPE stores the normed value as it is, so the norm's last rounding and the output's show only on sliding layers. +pub const entries = [_]Entry{ + .{ .name = "cuda ricursive tf_rc_prep sliding (RoPE, the 3,072-slot ring)", .form = .sliding }, + .{ .name = "cuda ricursive tf_rc_prep global (NoPE, a slot a position)", .form = .global }, + .{ .name = "cuda ricursive tf_rc_prep global over pages of 512 (NoPE, a shuffled page table)", .form = .paged }, + .{ .name = "control: ricursive tf_rc_prep drops bf16(x r)", .form = .sliding, .keep = all & ~norm_row }, + .{ .name = "control: ricursive tf_rc_prep drops bf16(w bf16(x r))", .form = .sliding, .keep = all & ~norm_out }, + .{ .name = "control: ricursive tf_rc_prep drops bf16(cos)", .form = .sliding, .keep = all & ~cos_pt }, + .{ .name = "control: ricursive tf_rc_prep drops bf16(sin)", .form = .sliding, .keep = all & ~sin_pt }, + .{ .name = "control: ricursive tf_rc_prep truncates its output, not rounding it", .form = .sliding, .keep = all & ~out_pt }, +}; + +/// Rows x a row's outputs x the call's end position (pos + rows): windows across the wrap, deep and last; chunks. +const sliding_served = [_]c.Shape{ + .{ .rows = 16, .n = width, .k = 3 * sliding_ring + 9 }, // a decode window across the ring's wrap + .{ .rows = 1, .n = width, .k = sliding_ring }, // the ring's last slot + .{ .rows = 1, .n = width, .k = sliding_ring + 1 }, // slot 0 again + .{ .rows = 7, .n = width, .k = 100_007 }, + .{ .rows = 16, .n = width, .k = 262_144 }, // the longest context checked + .{ .rows = 2048, .n = width, .k = 2048 }, // the first prompt chunk + .{ .rows = 2048, .n = width, .k = 4096 }, // the second, across the wrap + .{ .rows = 1000, .n = width, .k = 2 * sliding_ring + 500 }, // half before the wrap, half after + .{ .rows = 333, .n = width, .k = 333 }, + .{ .rows = 333, .n = width, .k = sliding_ring + 100 }, +}; + +/// The ring is slots(capacity) for a capacity of the end position: k a multiple of 64 ends on its last slot. +const global_served = [_]c.Shape{ + .{ .rows = 16, .n = width, .k = 100_016 }, + .{ .rows = 1, .n = width, .k = 262_144 }, // the last position of the longest context + .{ .rows = 5, .n = width, .k = 704 }, + .{ .rows = 2048, .n = width, .k = 2048 }, + .{ .rows = 1000, .n = width, .k = 3048 }, + .{ .rows = 333, .n = width, .k = 65_869 }, +}; + +/// The global shapes over pages, and a decode window across a page's edge (positions 1,016 to 1,031). +const paged_served = global_served ++ [_]c.Shape{.{ .rows = 16, .n = width, .k = 1032 }}; + +/// A lane's 4 fma steps, the warp's 5 shuffles. +fn depth(_: c.Shape) usize { + return 4 + 5; +} + +pub fn prepSpec(e: Entry) c.Spec { + const served: []const c.Shape = switch (e.form) { + .sliding => &sliding_served, + .global => &global_served, + .paged => &paged_served, + }; + return .{ .name = e.name, .accum = .f32, .depth = depth, .out = .bf16, .served = served, .sums = .{} }; +} + +/// q [rows][80][128], then the k and v rings' held slots, or (paged) every row of the call's pages and two spares. +pub const Prep = struct { + kern: *const rp.Kernels, + entry: Entry, + ctl: ?cuda.Function = null, + shape: c.Shape = undefined, + pos: usize = 0, + ring: usize = 0, + first: usize = 0, // the lowest slot the outputs hold + held: usize = 0, // the slots they hold + qk: []u16 = &.{}, + v: []u16 = &.{}, + qw: []u16 = &.{}, + kw: []u16 = &.{}, + inv: []f32 = &.{}, + table: []i32 = &.{}, // paged: the stream's page ids by position / 512 + dev: [6]u64 = @splat(0), // qk, v, inv, qw, kw, table + outs: [3]u64 = @splat(0), // q, and the k and v rings' bases (slot 0) + + pub const sums = {}; + + pub fn spec(self: *const Prep) c.Spec { + return prepSpec(self.entry); + } + + pub fn control(self: *const Prep) ?g.Gate { + return if (self.entry.keep != null) .shapes else null; + } + + pub fn free(self: *Prep, a: std.mem.Allocator) void { + for ([_]*[]u16{ &self.qk, &self.v, &self.qw, &self.kw }) |x| { + a.free(x.*); + x.* = &.{}; + } + a.free(self.inv); + self.inv = &.{}; + a.free(self.table); + self.table = &.{}; + } + + /// The call at `shape`: its position, ring or pages, inputs and outputs. + pub fn prepare(self: *Prep, a: std.mem.Allocator, rig: *rig_.Rig, shape: c.Shape, fill: g.Fill, rng: std.Random) !bool { + if (shape.n != width or shape.rows == 0 or shape.k < shape.rows) return false; + self.free(a); + self.shape = shape; + const rows = shape.rows; + self.pos = shape.k - rows; + const sliding = self.entry.form == .sliding; + const paged = self.entry.form == .paged; + self.ring = if (sliding) sliding_ring else if (paged) 0 else std.mem.alignForward(usize, shape.k, 64); + if (!paged and rows > self.ring) return false; + self.first = if (sliding or paged) 0 else self.pos; + self.held = if (sliding) self.ring else rows; + if (paged) { + // the call's pages drawn from a pool two larger, in shuffled order; pages below them are never read (-1) + const lo = self.pos / page; + const hi = (shape.k - 1) / page; + const pool = hi - lo + 1 + 2; + const ids = try a.alloc(i32, pool); + defer a.free(ids); + for (ids, 0..) |*x, i| x.* = @intCast(i); + rng.shuffle(i32, ids); + self.table = try a.alloc(i32, hi + 1); + @memset(self.table, -1); + @memcpy(self.table[lo..], ids[0 .. hi - lo + 1]); + self.held = pool * page; + } + try self.inputs(a, rig, rows, fill, rng); + const off = self.first * kvw * 2; // the rings' base, so slot `first` lands on the buffer's first row + self.outs[0] = try rig.output(rows * heads * D * 2); // allocated in the order the outputs gather: q, k, v + self.outs[1] = try rig.output(self.held * kvw * 2) - off; + self.outs[2] = try rig.output(self.held * kvw * 2) - off; + return true; + } + + /// Heads in [-3, 3) (one in 9 a tenth the size), weights in [0.5, 1.5); stress: integers in [-64, 64]; uploaded. + fn inputs(self: *Prep, a: std.mem.Allocator, rig: *rig_.Rig, rows: usize, fill: g.Fill, rng: std.Random) !void { + self.qk = try a.alloc(u16, rows * qk_stride); + self.v = try a.alloc(u16, rows * kvw); + for ([_][]u16{ self.qk, self.v }, [_]usize{ qk_stride, kvw }) |x, stride| for (x, 0..) |*e, i| { + e.* = c.toBf16(if (fill == .stress) ref.int(rng, -64, 64) else (rng.float(f32) * 6 - 3) * (if (i % stride / D % 9 == 4) @as(f32, 0.1) else 1)); + }; + if (fill == .nonfinite and rows >= 4) { + self.qk[3] = c.toBf16(ref.nan); // row 0's first q head: its every output NaN + self.qk[qk_stride + D + 9] = c.toBf16(ref.inf); // row 1's second q head: r = 0, so dim 9 (and 73 rotated) NaN + self.qk[2 * qk_stride + heads * D + 70] = c.toBf16(-ref.inf); // row 2's first k head, in the second half + self.v[3 * kvw + 2 * D + 5] = c.toBf16(ref.inf); // row 3's third v head: copied as it is + } + self.qw = try ref.bfs(a, rng, D, 0.5, 1.5); + self.kw = try ref.bfs(a, rng, D, 0.5, 1.5); + self.inv = try a.alloc(f32, D / 2); + for (self.inv, 0..) |*f, i| f.* = 1.0 / @as(f32, @floatCast(std.math.pow(f64, theta, @as(f64, @floatFromInt(2 * i)) / D))); // config.zig's inverseFrequencies + self.dev = .{ try ref.input(rig, self.qk), try ref.input(rig, self.v), try ref.input(rig, self.inv), try ref.input(rig, self.qw), try ref.input(rig, self.kw), if (self.table.len > 0) try ref.input(rig, self.table) else 0 }; + } + + pub fn run(self: *Prep, rig: *rig_.Rig) !void { + const off = self.first * kvw * 2; + if (self.held > self.shape.rows) for (self.outs[1..]) |o| try rig.zero(o + off, self.held * kvw * 2); // slots no row reaches: zero, as in the reference + const d = self.dev; + const p: rp.Prep = .{ .xq = rp.Rows.of(d[0], qk_stride), .xk = rp.Rows.of(d[0] + heads * D * 2, qk_stride), .xv = rp.Rows.of(d[1], kvw), .inv = d[2], .qw = d[3], .kw = d[4], .q = self.outs[0], .kc = self.outs[1], .vc = self.outs[2], .rows = self.shape.rows, .heads = heads, .kv_heads = kv_heads, .pos = self.pos, .ring = self.ring, .rope = self.entry.form == .sliding, .eps = eps, .pg = if (self.entry.form == .paged) .{ .table = d[5], .page_elems = page * kvw, .tokens = page } else null }; + if (self.ctl) |f| try rp.prepVia(f, self.entry.keep.?, rig.s, p) else try rp.prep(self.kern, rig.s, p); + } + + pub fn count(self: *const Prep) usize { + return self.shape.rows * heads * D + 2 * self.held * kvw; + } + + /// One head's 128 outputs: bf16(w bf16(x r)), then on sliding layers bf16(x c + other s), c, s the bf16 cos, sin. + fn head(self: *const Prep, s: Sums, x: []const u16, w: []const u16, cs: []const Iv, sn: []const Iv, out: []Iv, scratch: []f64) void { + var sq: [D]Iv = undefined; + var n: [D]Iv = undefined; + for (&sq, x) |*q, v| q.* = Iv.of(ref.bf(v) * ref.bf(v)); + const r = s.sum(&sq, depth(self.shape), scratch).div(Iv.of(D)).add(Iv.of(eps)).rsqrt(); + for (&n, x, w) |*o, v, wv| o.* = Iv.of(ref.bf(wv)).mul(Iv.of(ref.bf(v)).mul(r).bf()).bf(); + for (out, 0..) |*o, d| { + if (self.entry.form != .sliding) { + o.* = n[d]; + continue; + } + const other = if (d < D / 2) n[d + D / 2].neg() else n[d - D / 2]; + o.* = n[d].mul(cs[d % (D / 2)]).add(other.mul(sn[d % (D / 2)])).bf(); + } + } + + /// Position p's row in the outputs' k and v: its ring slot, or its page's row in the pool. + fn slot(self: *const Prep, p: usize) usize { + if (self.entry.form == .paged) return @as(usize, @intCast(self.table[p / page])) * page + p % page; + return p % self.ring - self.first; + } + + /// Row r's position and its row in the outputs' k and v: the call's pos + r. + fn rowPos(self: *const Prep, r: usize) usize { + return self.pos + r; + } + + fn rowSlot(self: *const Prep, r: usize) usize { + return self.slot(self.pos + r); + } + + /// Every rounding point kept, whatever the entry: a control's kernel drops one and must leave these ranges. + pub fn solve(self: *const Prep, _: std.mem.Allocator, s: Sums, out: []Iv) !void { + self.solveAt(s, out, self); + } + + /// The outputs with row r at by.rowPos(r), its k and v at row by.rowSlot(r) of the held ones (the rest zero). + fn solveAt(self: *const Prep, s: Sums, out: []Iv, by: anytype) void { + const rows = self.shape.rows; + const qn = rows * heads * D; + const ks = out[qn..][0 .. self.held * kvw]; + const vs = out[qn + self.held * kvw ..]; + @memset(ks, Iv.of(0)); + @memset(vs, Iv.of(0)); + var scratch: [D]f64 = undefined; + var cs: [D / 2]Iv = undefined; + var sn: [D / 2]Iv = undefined; + for (0..rows) |r| { + const position: f32 = @floatFromInt(by.rowPos(r)); + for (&cs, &sn, self.inv) |*co, *si, f| { + const angle = Iv.of(position * f); // the kernel's fp32 angle, exactly + co.* = angle.trig(false).bf(); + si.* = angle.trig(true).bf(); + } + const t = by.rowSlot(r); // the row's slot, as the outputs hold it + const xr = self.qk[r * qk_stride ..][0..qk_stride]; + for (0..heads) |h| self.head(s, xr[h * D ..][0..D], self.qw, &cs, &sn, out[(r * heads + h) * D ..][0..D], &scratch); + for (0..kv_heads) |h| self.head(s, xr[(heads + h) * D ..][0..D], self.kw, &cs, &sn, ks[t * kvw + h * D ..][0..D], &scratch); + for (vs[t * kvw ..][0..kvw], self.v[r * kvw ..][0..kvw]) |*o, x| o.* = Iv.of(ref.bf(x)); + } + } +}; + +/// tf_rc_prep_rows' calls: 4 streams' rows round-robin in 5 shuffled slots (one unused), last rows 3,073 apart. +const streams = 4; +const slots = streams + 1; +const apart = sliding_ring + 1; + +/// A registration of tf_rc_prep_rows (form sliding or paged): the served kernel, or the control naming no row's stream. +pub const RowsEntry = struct { name: []const u8, form: Form, keep: ?i32 = null }; + +pub const rows_entries = [_]RowsEntry{ + .{ .name = "cuda ricursive tf_rc_prep_rows sliding (4 streams' rows in their 3,072-slot rings)", .form = .sliding }, + .{ .name = "cuda ricursive tf_rc_prep_rows global over pages of 512 (4 streams, each its row of the page table)", .form = .paged }, + .{ .name = "control: ricursive tf_rc_prep_rows drops the slot (every row into slot 0's ring)", .form = .sliding, .keep = all }, +}; + +/// Rows x a row's outputs x the first stream's last position + 1 (the others end 3,073, 6,146 and 9,219 below it). +const rows_served = [_]c.Shape{ + .{ .rows = 16, .n = width, .k = 4 * sliding_ring + 2 }, // verify windows of 4 ending at ring slots 1, 0, 3,071, 3,070 + .{ .rows = 4, .n = width, .k = 4 * sliding_ring + 2 }, // a decode round: a row a stream, at those slots + .{ .rows = 7, .n = width, .k = 3 * apart + 3 }, // windows of 2, 2, 2 and 1, the last stream's row at position 2 + .{ .rows = 64, .n = width, .k = 100_007 }, + .{ .rows = 16, .n = width, .k = 262_144 }, // the longest context checked +}; + +/// Gate 3 holds tf_rc_prep_rows to tf_rc_prep's bits; gates 1 and 4 to the float64 reference, as tf_rc_prep. +pub fn rowsSpec(e: RowsEntry) c.Spec { + return .{ .name = e.name, .accum = .f32, .depth = depth, .out = .bf16, .replaces = "tf_rc_prep (a launch a row, at its position on its stream's cache)", .served = &rows_served, .sums = .{} }; +} + +/// Rows stream i holds in a call of `rows`: rows i, i + 4, ... +fn streamRows(rows: usize, i: usize) usize { + return (rows + streams - 1 - i) / streams; +} + +/// Stream i's last position. +fn lastPos(shape: c.Shape, i: usize) usize { + return shape.k - 1 - i * apart; +} + +/// Stream i's pages: its first and one past its last (none without rows). +fn pageSpan(shape: c.Shape, i: usize) [2]usize { + const n = streamRows(shape.rows, i); + if (n == 0) return .{ 0, 0 }; + return .{ (lastPos(shape, i) + 1 - n) / page, lastPos(shape, i) / page + 1 }; +} + +/// tf_rc_prep_rows on Prep's inputs, against tf_rc_prep a launch a row into scratch alike (every slot's ring or pages). +pub const PrepRows = struct { + kern: *const rp.Kernels, + entry: RowsEntry, + ctl: ?cuda.Function = null, + p: Prep, // inputs, outputs and arithmetic: the form's one-stream registration's + at: []i32 = &.{}, // row r's position + of: []i32 = &.{}, // row r's slot + stride: usize = 0, // paged: a slot's entries in the table + dev: [2]u64 = @splat(0), // at, of + base: [3]u64 = @splat(0), // tf_rc_prep's q and k and v caches + host: std.ArrayList(u8) = .empty, // their bytes, in the outputs' order + + pub const sums = {}; + + pub fn spec(self: *const PrepRows) c.Spec { + return rowsSpec(self.entry); + } + + pub fn control(self: *const PrepRows) ?g.Gate { + return if (self.entry.keep != null) .shapes else null; + } + + pub fn free(self: *PrepRows, a: std.mem.Allocator) void { + self.p.free(a); + a.free(self.at); + self.at = &.{}; + a.free(self.of); + self.of = &.{}; + self.host.deinit(a); + self.host = .empty; + } + + pub fn prepare(self: *PrepRows, a: std.mem.Allocator, rig: *rig_.Rig, shape: c.Shape, fill: g.Fill, rng: std.Random) !bool { + const rows = shape.rows; + const most = streamRows(rows, 0); + if (shape.n != width or rows == 0 or most > sliding_ring or shape.k < (streams - 1) * apart + most) return false; + self.free(a); + const p = &self.p; + p.shape = shape; + const paged = self.entry.form == .paged; + var ids: [slots]i32 = undefined; + for (&ids, 0..) |*x, i| x.* = @intCast(i); + rng.shuffle(i32, &ids); // stream i in slot ids[i] + self.at = try a.alloc(i32, rows); + self.of = try a.alloc(i32, rows); + for (self.at, self.of, 0..) |*x, *t, r| { + const i = r % streams; + x.* = @intCast(lastPos(shape, i) + 1 - streamRows(rows, i) + r / streams); + t.* = ids[i]; + } + p.ring = if (paged) 0 else sliding_ring; + p.first = 0; + p.held = slots * sliding_ring; + if (paged) { + // each stream's pages from a pool two larger, in shuffled order; a slot's other entries are never read (-1) + self.stride = (shape.k - 1) / page + 1; + var need: usize = 0; + for (0..streams) |i| need += pageSpan(shape, i)[1] - pageSpan(shape, i)[0]; + const pool = try a.alloc(i32, need + 2); + defer a.free(pool); + for (pool, 0..) |*x, i| x.* = @intCast(i); + rng.shuffle(i32, pool); + p.table = try a.alloc(i32, slots * self.stride); + @memset(p.table, -1); + var next: usize = 0; + for (0..streams) |i| { + const span = pageSpan(shape, i); + const row = p.table[@as(usize, @intCast(ids[i])) * self.stride ..]; + for (span[0]..span[1]) |pg| { + row[pg] = pool[next]; + next += 1; + } + } + p.held = pool.len * page; + } + try p.inputs(a, rig, rows, fill, rng); + self.dev = .{ try ref.input(rig, self.at), try ref.input(rig, self.of) }; + const qb = rows * heads * D * 2; + const kb = p.held * kvw * 2; + p.outs = .{ try rig.output(qb), try rig.output(kb), try rig.output(kb) }; // gathered in this order: q, k, v + self.base = .{ try rig.scratch(qb), try rig.scratch(kb), try rig.scratch(kb) }; + return true; + } + + fn call(self: *const PrepRows) rp.PrepRows { + const p = &self.p; + const d = p.dev; + return .{ .xq = rp.Rows.of(d[0], qk_stride), .xk = rp.Rows.of(d[0] + heads * D * 2, qk_stride), .xv = rp.Rows.of(d[1], kvw), .inv = d[2], .qw = d[3], .kw = d[4], .q = p.outs[0], .kc = p.outs[1], .vc = p.outs[2], .rows = p.shape.rows, .heads = heads, .kv_heads = kv_heads, .pos = self.dev[0], .slot = self.dev[1], .ring = p.ring, .rope = self.entry.form == .sliding, .eps = eps, .pg = if (self.entry.form == .paged) .{ .table = d[5], .page_elems = page * kvw, .tokens = page, .stride = self.stride } else null }; + } + + pub fn run(self: *PrepRows, rig: *rig_.Rig) !void { + for (self.p.outs[1..]) |o| try rig.zero(o, self.p.held * kvw * 2); // slots no row reaches: zero, as in the reference + if (self.ctl) |f| try rp.prepRowsVia(f, self.entry.keep.?, rig.s, self.call()) else try rp.prepRows(self.kern, rig.s, self.call()); + } + + /// Gate 3's baseline: tf_rc_prep, a launch a row at the row's position, on its slot's ring or its slot's table row. + pub fn baseline(self: *PrepRows, rig: *rig_.Rig) ![]const u8 { + const p = &self.p; + const d = p.dev; + const b = self.base; + const kb = p.held * kvw * 2; + for (b[1..]) |o| try rig.zero(o, kb); + const paged = self.entry.form == .paged; + for (self.at, self.of, 0..) |x, t, r| { + const ts: usize = @intCast(t); + const ring_at = ts * p.ring * kvw * 2; // the slot's ring (pages: ring 0, the pools themselves) + try rp.prep(self.kern, rig.s, .{ + .xq = rp.Rows.of(d[0] + r * qk_stride * 2, qk_stride), + .xk = rp.Rows.of(d[0] + (r * qk_stride + heads * D) * 2, qk_stride), + .xv = rp.Rows.of(d[1] + r * kvw * 2, kvw), + .inv = d[2], + .qw = d[3], + .kw = d[4], + .q = b[0] + r * heads * D * 2, + .kc = b[1] + ring_at, + .vc = b[2] + ring_at, + .rows = 1, + .heads = heads, + .kv_heads = kv_heads, + .pos = @intCast(x), + .ring = p.ring, + .rope = !paged, + .eps = eps, + .pg = if (paged) .{ .table = d[5] + ts * self.stride * 4, .page_elems = page * kvw, .tokens = page } else null, + }); + } + const a = rig.gpu.gpa; + self.host.clearRetainingCapacity(); + for (b, [_]usize{ p.shape.rows * heads * D * 2, kb, kb }) |o, n| try self.host.appendSlice(a, try rig.bytes(o, n)); + return self.host.items; + } + + pub fn count(self: *const PrepRows) usize { + return self.p.count(); + } + + /// Row r's position, and its row in the outputs' k and v: its slot's ring slot, or its page's row in the pool. + fn rowPos(self: *const PrepRows, r: usize) usize { + return @intCast(self.at[r]); + } + + fn rowSlot(self: *const PrepRows, r: usize) usize { + const x: usize = @intCast(self.at[r]); + const t: usize = @intCast(self.of[r]); + if (self.entry.form == .paged) return @as(usize, @intCast(self.p.table[t * self.stride + x / page])) * page + x % page; + return t * sliding_ring + x % sliding_ring; + } + + /// Every rounding point kept: Prep's reference at each row's own position and slot. + pub fn solve(self: *const PrepRows, _: std.mem.Allocator, s: Sums, out: []Iv) !void { + self.p.solveAt(s, out, self); + } +}; + +/// ref.Gated(PrepRows) and the spec's baseline: tf_rc_prep's bytes on the same inputs (ref.Gated has no `replaced`). +pub const RowsGated = struct { + in: ref.Gated(PrepRows), + rig: *rig_.Rig, // in.rig: the hooks read the outputs there + + const vt = g.hooks(RowsGated); + + pub fn kernel(self: *RowsGated) g.Kernel { + return .{ .spec = self.in.k.spec(), .ctx = self, .vt = &vt, .control = self.in.k.control() }; + } + + pub fn prepare(self: *RowsGated, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + return self.in.prepare(shape, fill, seed); + } + + pub fn run(self: *RowsGated) !void { + return self.in.run(); + } + + pub fn count(self: *RowsGated) usize { + return self.in.count(); + } + + pub fn reference(self: *RowsGated, want: []f64, mag: []f64) !void { + return self.in.reference(want, mag); + } + + pub fn bounds(self: *RowsGated, want: []const f64, mag: []const f64, allow: []f64) !void { + return self.in.bounds(want, mag, allow); + } + + pub fn intervals(self: *RowsGated, low: []f64, high: []f64) !void { + return self.in.intervals(low, high); + } + + pub fn simulate(self: *RowsGated, acc: c.Accum, order: g.Order, out: []f64) !void { + return self.in.simulate(acc, order, out); + } + + pub fn replaced(self: *RowsGated) ![]const u8 { + return self.in.k.baseline(self.rig); + } +}; + +/// The registrations, on the heap: the gates hold pointers to them. +pub const Set = struct { + gpa: std.mem.Allocator, + k: rp.Kernels, + ctl: cuda.Module, + preps: [entries.len]ref.Gated(Prep) = undefined, + rows: [rows_entries.len]RowsGated = undefined, + + /// The set, its kernels appended to `list` (the controls after the kernel they break). + pub fn init(gpu: check.Gpu, s: cuda.Stream, a: std.mem.Allocator, list: *std.ArrayList(g.Kernel)) !*Set { + const set = try gpu.gpa.create(Set); + errdefer gpu.gpa.destroy(set); + var k = try rp.Kernels.load(gpu.d); + errdefer k.unload(); + var ctl = try cuda.Module.load(gpu.d, &controls.bytes); + errdefer ctl.unload(); + const f = try ctl.function("tf_ctl_rc_prep"); + const fr = try ctl.function("tf_ctl_rc_prep_rows"); + set.* = .{ .gpa = gpu.gpa, .k = k, .ctl = ctl }; + for (&set.preps, entries) |*x, e| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .k = .{ .kern = &set.k, .entry = e, .ctl = if (e.keep != null) f else null } }; + for (&set.rows, rows_entries) |*x, e| { + x.* = .{ .in = .{ .rig = .{ .gpu = gpu, .s = s }, .k = .{ .kern = &set.k, .entry = e, .ctl = if (e.keep != null) fr else null, .p = .{ .kern = &set.k, .entry = .{ .name = e.name, .form = e.form } } } }, .rig = undefined }; + x.rig = &x.in.rig; + } + for (&set.preps) |*x| try list.append(a, x.kernel()); + for (&set.rows) |*x| try list.append(a, x.kernel()); + return set; + } + + pub fn deinit(set: *Set) void { + for (&set.rows) |*x| x.in.deinit(); + for (&set.preps) |*x| x.deinit(); + set.ctl.unload(); + set.k.unload(); + set.gpa.destroy(set); + } +}; + +test "each control keeps all of tf_rc_prep's rounding points but one, and each point has its control" { + var dropped: i32 = 0; + for (entries) |e| if (e.keep) |k| { + const gone = all & ~k; + try std.testing.expectEqual(@as(u32, 1), @popCount(gone)); + try std.testing.expect((dropped & gone) == 0); + try std.testing.expect(e.form == .sliding); + dropped |= gone; + }; + try std.testing.expectEqual(all, dropped); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index f816738b4d..c4c773f983 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -261,11 +261,14 @@ const usage = \\ pocket-contract Nemotron's pocket kernels (learned neurons) through gates 1-4, controls red \\ flashnext-modelopt-contract Flash Next NVFP4's BF16 dense and FP8 MTP-expert downs through gates 1-4 \\ minimax-contract MiniMax-M3.1-Flash's indexer and block-sparse attention through gates 1-4, controls red + \\ ricursive-contract Ripple's glue, attention prep and attention (group 10) through gates 1-4, controls red + \\ ripple-forward Ripple's forward rounds (streams, verify windows, mixes) == one-row steps, bit for bit (synthetic) \\ moss-contract MOSS-Image's own kernels and its full attention through gates 1-4, controls red \\ glm-exl3 GLM's EXL3 routed experts, decode and prompt, against float64 (synthetic) \\ glm-kda-pad-split TF_GLM_KDA_PAD's padding of the bf16 and split KDA stacks at TP=3 (synthetic) \\ glm-vision MODEL F32 BF16 IMG GLM's image tower and preparation against vision_ref.py's torch runs \\ glm-mm-fixed Prompt K policies: FP64, every row alone/batched, offsets and guarded tails + \\ ripple-serve Ripple's lane backend: shared rounds, copied drafts, kept states, logprob rows (synthetic) \\ glm-mm-bench GLM's bf16 linears: b16.matmul against glm_mm, by rows (one GPU) \\ glm-index-check GLM's indexer with glm_index against tf_glm_scores and tf_glm_select_rows, byte for byte (one GPU) \\ glm-kv-host-slots TF_GLM_KV_HOST's gather past its slots is refused as error.GlmKvHostSlots, then one within (one GPU) @@ -513,9 +516,12 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "qwen-pages")) return @import("contract_qwen_pages.zig").run(gpu); if (std.mem.eql(u8, cmd, "qwen-pending")) return @import("contract_qwen_pending.zig").run(gpu); if (std.mem.eql(u8, cmd, "minimax-contract")) return @import("contract_minimax.zig").run(gpu); + if (std.mem.eql(u8, cmd, "ricursive-contract")) return @import("contract_ricursive.zig").run(gpu); + if (std.mem.eql(u8, cmd, "ripple-forward")) return @import("ripple_forward_tests.zig").run(gpu); if (std.mem.eql(u8, cmd, "rope-yarn")) return @import("contract_qwen_rope_yarn.zig").run(gpu); if (std.mem.eql(u8, cmd, "kernel-contract-limits")) return contract_tests.fatbins(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "pocket-contract")) return @import("contract_pocket.zig").run(gpu); + if (std.mem.eql(u8, cmd, "ripple-serve")) return @import("ripple_serve_tests.zig").run(gpu); if (std.mem.eql(u8, cmd, "flashnext-modelopt-contract")) return @import("contract_flashnext_modelopt.zig").run(gpu); if (std.mem.eql(u8, cmd, "moss-contract")) return @import("contract_moss_image.zig").run(gpu); if (std.mem.eql(u8, cmd, "tree")) return tree_tests.run(gpu); diff --git a/zig/tests/cuda/ripple_forward_tests.zig b/zig/tests/cuda/ripple_forward_tests.zig new file mode 100644 index 0000000000..695664a0b6 --- /dev/null +++ b/zig/tests/cuda/ripple_forward_tests.zig @@ -0,0 +1,489 @@ +//! Ripple's forward rounds == one-row steps, bit for bit (ripple-forward): streams, verify windows, graphs, reuse. +const std = @import("std"); +const cuda = @import("cuda"); +const rc = @import("ricursive_cuda"); +const check = @import("check.zig"); + +const fwd = rc.model; +const wts = rc.weights; +const Config = @FieldType(wts.Weights, "config"); +const Weight = cuda.qmmf.Weight; + +/// The release's attention and widths, so every lane matmul takes the release's plan (tile, K slices, cluster). +pub fn config() Config { + var c: Config = .{ .hidden = 5120, .intermediate = 12288, .layers = 10, .heads = 80, .kv_heads = 8, .vocab = 16384, .eps = 1e-6, .theta = 10000, .window = 1024, .embed_scale = 71.55417527999327, .clamp = 20, .context = 8192, .fp4_mlp = true }; + for (0..c.layers) |i| c.global[i] = i % 5 == 4; + return c; +} + +/// bf16 bits of x, rounded to nearest even. +fn bf16(x: f32) u16 { + const b: u32 = @bitCast(x); + return @intCast((b + 0x7fff + ((b >> 16) & 1)) >> 16); +} + +const Build = struct { + w: *wts.Weights, + d: *const cuda.Driver, + r: std.Random, + + fn put(b: Build, bytes: []const u8) !u64 { + var buf = try cuda.DeviceBuffer.fromHost(b.d, bytes); + errdefer buf.free(); + try b.w.buffers.append(b.w.gpa, buf); + b.w.bytes += bytes.len; + return buf.ptr; + } + + /// `n` bf16 values of magnitude in [2^e, 2^(e+1)), random signs and mantissas (fast: the release's sizes). + fn values(b: Build, n: usize, e: i32) ![]u16 { + const v = try b.w.gpa.alloc(u16, n); + const exp: u16 = @intCast((e + 127) << 7); + var bits: u64 = 0; + for (v, 0..) |*x, i| { + if (i % 4 == 0) bits = b.r.int(u64); + x.* = exp | @as(u16, @truncate(bits & 0x807f)); + bits >>= 16; + } + return v; + } + + /// A norm's or the embedding's bf16 values, normal with deviation `sd` around `mean`. + fn small(b: Build, n: usize, mean: f32, sd: f32) !u64 { + const v = try b.w.gpa.alloc(u16, n); + defer b.w.gpa.free(v); + for (v) |*x| x.* = bf16(mean + sd * b.r.floatNorm(f32)); + return b.put(std.mem.sliceAsBytes(v)); + } + + /// A bf16 [n, k] projection packed for the lane and prompt matmuls, its values about 1 / sqrt(k). + fn dense(b: Build, n: usize, k: usize) !Weight { + const src = try b.values(n * k, -@as(i32, std.math.log2_int(usize, k) / 2) - 1); + defer b.w.gpa.free(src); + const out = try b.w.gpa.alloc(u16, cuda.modelopt.padded(n) * k); + defer b.w.gpa.free(out); + cuda.qmmf.packBf16(out, src, n, k); + return cuda.modelopt.bf16Weight(try b.put(std.mem.sliceAsBytes(out)), n, k); + } + + /// An NVFP4 [n, k] projection: random nibbles, E4M3 scales of 1 to 3, a global of 1/1024 (about 1 / sqrt(k)). + fn fp4(b: Build, n: usize, k: usize) !Weight { + const gpa = b.w.gpa; + const codes = try gpa.alloc(u8, n * k / 2); + defer gpa.free(codes); + b.r.bytes(codes); + const scales = try gpa.alloc(u8, n * k / 16); + defer gpa.free(scales); + const e4m3 = [_]u8{ 0x38, 0x3c, 0x40, 0x44 }; + b.r.bytes(scales); + for (scales) |*x| x.* = e4m3[x.* & 3]; + const words = try gpa.alloc(u32, cuda.nvfp4.padded(n) / 64 * (k / 64) * 512); + defer gpa.free(words); + const tiles = try gpa.alloc(u8, cuda.nvfp4.padded(n) * (k / 16)); + defer gpa.free(tiles); + cuda.nvfp4.packWords(words, codes, n, k); + cuda.nvfp4.packScales(tiles, scales, n, k); + return cuda.nvfp4.weight(try b.put(std.mem.sliceAsBytes(words)), try b.put(tiles), 1.0 / 1024.0, n, k); + } +}; + +/// Random weights in the loader's layout (cuda_weights.load): q, k and g one projection, NVFP4 MLPs. +pub fn build(gpu: check.Gpu, c: Config, seed: u64) !wts.Weights { + var w: wts.Weights = .{ .gpa = gpu.gpa, .config = c }; + errdefer w.deinit(); + var prng = std.Random.DefaultPrng.init(seed); + const b: Build = .{ .w = &w, .d = gpu.d, .r = prng.random() }; + const h = c.hidden; + const q = c.queryWidth(); + const kv = c.kvWidth(); + w.embed = try b.small(c.vocab * h, 0, 0.02); + w.head = try b.dense(c.vocab, h); + w.norm = try b.small(h, 1, 0.1); + w.layers = try gpu.gpa.alloc(wts.Layer, c.layers); + for (w.layers) |*l| l.* = .{ + .qkg = try b.dense(2 * q + kv, h), + .v = try b.dense(kv, h), + .o = try b.dense(h, q), + .gate = try b.fp4(c.intermediate, h), + .up = try b.fp4(c.intermediate, h), + .down = try b.fp4(h, c.intermediate), + .input_norm = try b.small(h, 1, 0.1), + .post_attn = try b.small(h, 1, 0.1), + .pre_ffn = try b.small(h, 1, 0.1), + .post_ffn = try b.small(h, 1, 0.1), + .q_norm = try b.small(128, 1, 0.1), + .k_norm = try b.small(128, 1, 0.1), + .blend = wts.blendOf(bf16(b.r.float(f32) * 2 - 1)), + }; + return w; +} + +/// A stream: its prompt's length and decode tokens (teacher-forced, so every run feeds the same rows). +const Spec = struct { name: []const u8, prompt: usize }; + +const specs = [_]Spec{ + .{ .name = "A", .prompt = 507 }, // rows across the 512-row page edge + .{ .name = "B", .prompt = 3066 }, // across the 3,072-row ring's wrap + .{ .name = "C", .prompt = 1019 }, // across 1,024: a page edge and the window's + .{ .name = "D", .prompt = 37 }, + .{ .name = "E", .prompt = 4093 }, // a wrapped ring, rows across 4,096 keys (8 to 9 key groups) +}; + +/// Rounds of (stream, rows): several streams a row each, a verify window, mixes of 20 and 48 rows, one and k rows. +const rounds = [_][]const [2]usize{ + &.{ .{ 0, 1 }, .{ 1, 1 }, .{ 2, 1 }, .{ 3, 1 }, .{ 4, 1 } }, + &.{.{ 1, 8 }}, + &.{ .{ 0, 5 }, .{ 2, 5 }, .{ 3, 5 }, .{ 4, 5 } }, + &.{ .{ 0, 10 }, .{ 1, 10 }, .{ 2, 10 }, .{ 3, 10 }, .{ 4, 8 } }, + &.{ .{ 3, 1 }, .{ 1, 3 } }, +}; +const round_rows = 48; + +fn decodeTokens(j: usize) usize { + var n: usize = 0; + for (rounds) |r| for (r) |x| if (x[0] == j) { + n += x[1]; + }; + return n; +} + +const Streams = struct { + gpa: std.mem.Allocator, + ids: [specs.len][]u32 = undefined, // the prompt, then the decode tokens + want: [specs.len][][]u8 = undefined, // one-row steps' logits a decode token + prompt: [specs.len][]u8 = undefined, // the prompt's last row's logits + + fn init(gpa: std.mem.Allocator, vocab: usize, seed: u64) !Streams { + var x: Streams = .{ .gpa = gpa }; + var prng = std.Random.DefaultPrng.init(seed); + for (specs, 0..) |sp, j| { + x.ids[j] = try gpa.alloc(u32, sp.prompt + decodeTokens(j)); + for (x.ids[j]) |*t| t.* = prng.random().uintLessThan(u32, @intCast(vocab)); + x.want[j] = try gpa.alloc([]u8, decodeTokens(j)); + for (x.want[j]) |*l| l.* = try gpa.alloc(u8, vocab * 2); + x.prompt[j] = try gpa.alloc(u8, vocab * 2); + } + return x; + } + + fn deinit(x: *Streams) void { + for (0..specs.len) |j| { + for (x.want[j]) |l| x.gpa.free(l); + x.gpa.free(x.want[j]); + x.gpa.free(x.ids[j]); + x.gpa.free(x.prompt[j]); + } + } +}; + +/// Equal logits rows, or the first differing logit and how many differ. +fn same(comptime what: []const u8, args: anytype, got: []const u8, want: []const u8) bool { + if (std.mem.eql(u8, got, want)) return true; + const g = std.mem.bytesAsSlice(u16, got); + const w = std.mem.bytesAsSlice(u16, want); + var diff: usize = 0; + var first: usize = g.len; + for (g, w, 0..) |a, b, i| if (a != b) { + diff += 1; + first = @min(first, i); + }; + std.debug.print(" DIFF ", .{}); + std.debug.print(what, args); + std.debug.print(": {d} of {d} logits differ, first {d}: {x:0>4} for {x:0>4}\n", .{ diff, g.len, first, g[first], w[first] }); + return false; +} + +/// A prompt into `slot` from 0 in chunks of up to 2,048 rows (M1's path), its last row through the head. +fn prefill(m: *const fwd.Model, slot: usize, ids: []const u32) !fwd.Seq { + var seq: fwd.Seq = .{ .slot = slot }; + while (seq.len < ids.len) { + const end = @min(ids.len, seq.len + fwd.prompt_rows); + const at = seq.len; + try m.forward(&seq, ids[at..end], .prompt); + if (end == ids.len) try m.head(end - at - 1); + } + return seq; +} + +/// One-row steps of stream j in its reference slot j: the logits each round must give. +fn reference(m: *fwd.Model, x: *Streams, j: usize, logits: []u8) !usize { + const ids = x.ids[j]; + const n = specs[j].prompt; + var seq = try prefill(m, j, ids[0..n]); + try m.download(m.b.logits, x.prompt[j]); + var fails: usize = 0; + for (x.want[j], n..) |want, p| { + try m.forward(&seq, ids[p..][0..1], .decode); + try m.head(0); + try m.download(m.b.logits, want); + if (p != n and p != n + 5) continue; + // the same row again on four-warp tree blocks (sms 0): its keys are rewritten with the same bits + const sms = m.gd.tree.sms; + m.gd.tree.sms = 0; + seq.len -= 1; + try m.forward(&seq, ids[p..][0..1], .decode); + try m.head(0); + try m.download(m.b.logits, logits); + m.gd.tree.sms = sms; + if (same("tree one-warp vs four-warp, stream {s} at {d}", .{ specs[j].name, p }, logits, want)) { + check.pass("ripple-forward tree one-warp == four-warp blocks: stream {s} at {d} ({d} key groups)", .{ specs[j].name, p, rc.attn.groupsOf(&.{p}) }); + } else { + std.debug.print("FAIL ripple-forward tree one-warp != four-warp blocks: stream {s} at {d}\n", .{ specs[j].name, p }); + fails += 1; + } + } + return fails; +} + +/// The driver's free device bytes (GB10: of its unified memory) after the stream's work. +fn freeBytes(m: *const fwd.Model) !usize { + try m.s.synchronize(); + var free: usize = 0; + var total: usize = 0; + try m.d.check(m.d.api.cuMemGetInfo_v2(&free, &total), "cuMemGetInfo"); + return free; +} + +/// The last round's rows' logits by rowLogits into `out`. +fn roundRows(m: *const fwd.Model, n: usize, out: []u8) !void { + try m.rowLogits(n); + try m.download(m.b.rows, out[0 .. n * m.c.vocab * 2]); +} + +/// Each of n rows of `got` the same as `want`'s. +fn sameRows(what: []const u8, n: usize, v: usize, got: []const u8, want: []const u8) bool { + var ok = true; + for (0..n) |r| ok = same("{s} row {d}", .{ what, r }, got[r * v ..][0..v], want[r * v ..][0..v]) and ok; + return ok; +} + +/// A round from its graph into `got`, then eager over the same caches into `want`; true when every row's logits match. +fn graphVsEager(m: *const fwd.Model, what: []const u8, chains: []const fwd.Chain, ids: []const u32, got: []u8, want: []u8) !bool { + try m.roundAs(chains, ids, .graph); + try roundRows(m, ids.len, got); + try m.roundAs(chains, ids, .exact); + try roundRows(m, ids.len, want); + return sameRows(what, ids.len, m.c.vocab * 2, got, want); +} + +/// Rounds over the test slots against the references, by rowLogits and by head: a graph, fixed eager, then B1's eager. +fn runRounds(m: *const fwd.Model, x: *Streams, base: usize, done: *[specs.len]usize, logits: []u8, rows: []u8, by_graph: []u8, by_fixed: []u8) !usize { + const v = m.c.vocab; + var fails: usize = 0; + var chains: [specs.len]fwd.Chain = undefined; + var ids: [round_rows]u32 = undefined; + var who: [round_rows][2]usize = undefined; // (stream, decode token) a row + var pos: [round_rows]usize = undefined; + for (rounds, 0..) |rd, ri| { + var n: usize = 0; + for (rd, 0..) |c, ci| { + const j = c[0]; + const p0 = specs[j].prompt + done[j]; + chains[ci] = .{ .slot = base + j, .p0 = p0, .tokens = c[1] }; + for (0..c[1]) |i| { + ids[n] = x.ids[j][p0 + i]; + who[n] = .{ j, done[j] + i }; + pos[n] = p0 + i; + n += 1; + } + done[j] += c[1]; + } + const ch = chains[0..rd.len]; + const free0 = try freeBytes(m); + try m.roundAs(ch, ids[0..n], .graph); + try roundRows(m, n, by_graph); + const graph_mib = (@as(f64, @floatFromInt(free0)) - @as(f64, @floatFromInt(try freeBytes(m)))) / (1 << 20); + std.debug.print("INFO ripple-forward graph of {d} rows: {d:.1} MiB of device memory (the driver's free bytes before and after its capture), {d} graphs\n", .{ n, graph_mib, m.captured() }); + try m.roundAs(ch, ids[0..n], .fixed); + try roundRows(m, n, by_fixed); + try m.roundAs(ch, ids[0..n], .exact); + try roundRows(m, n, rows); + var bad: usize = 0; + for (0..n) |r| { + const want = x.want[who[r][0]][who[r][1]]; + const name = specs[who[r][0]].name; + if (!same("round {d} row {d} (stream {s} at {d}) by rowLogits", .{ ri, r, name, pos[r] }, rows[r * v * 2 ..][0 .. v * 2], want)) bad += 1; + try m.head(r); + try m.download(m.b.logits, logits); + if (!same("round {d} row {d} (stream {s} at {d}) by head", .{ ri, r, name, pos[r] }, logits, want)) bad += 1; + } + const ng = rc.attn.groupsOf(pos[0..n]); + const path = if (rc.attn.oneWarp(&m.gd, n, ng)) "one-warp" else "four-warp"; + if (bad == 0) { + check.pass("ripple-forward round {d}: {d} chains, {d} rows ({s} tree blocks, {d} key groups) == one-row steps", .{ ri, rd.len, n, path, ng }); + } else { + std.debug.print("FAIL ripple-forward round {d}: {d} chains, {d} rows ({s} tree blocks, {d} key groups): {d} row comparisons differ\n", .{ ri, rd.len, n, path, ng, bad }); + fails += 1; + } + const fixed = m.gd.groups(); + const fixed_path = if (rc.attn.oneWarp(&m.gd, n, fixed)) "one-warp" else "four-warp"; + if (sameRows("fixed key groups", n, v * 2, by_fixed, rows)) { + check.pass("ripple-forward fixed key groups, round {d}: {d} groups ({s}) == the round's own {d} ({s}), eager", .{ ri, fixed, fixed_path, ng, path }); + } else { + std.debug.print("FAIL ripple-forward fixed key groups, round {d}: {d} groups ({s}) != the round's own {d} ({s})\n", .{ ri, fixed, fixed_path, ng, path }); + fails += 1; + } + if (sameRows("graph", n, v * 2, by_graph, rows)) { + check.pass("ripple-forward graph, round {d}: {d} rows replayed from a graph == the eager round", .{ ri, n }); + } else { + std.debug.print("FAIL ripple-forward graph, round {d}: {d} rows from a graph != the eager round\n", .{ ri, n }); + fails += 1; + } + } + return fails; +} + +/// A replayed round: (slot, p0, rows) a chain. +const Replay = struct { what: []const u8, chains: []const [3]usize }; + +/// Graphs captured by the rounds above, replayed at other positions and slots (each against eager). +const elsewhere = [_]Replay{ + .{ .what = "the 5-row graph (five one-row chains) as one window in slot 1 at 3,070 (the ring's wrap, a page edge)", .chains = &.{.{ 1, 3070, 5 }} }, + .{ .what = "the 8-row graph (one window) as four 2-row chains: page edges at 512 and 1,024, 4,096 keys, the reused slot", .chains = &.{ .{ 0, 511, 2 }, .{ 2, 1023, 2 }, .{ 4, 4095, 2 }, .{ 6, 52, 2 } } }, + .{ .what = "the one-row graph in slot 8 at 54", .chains = &.{.{ 8, 54, 1 }} }, + .{ .what = "the 4-row graph (two chains) as one window over 4,096 keys of a wrapped ring", .chains = &.{.{ 9, 4094, 4 }} }, + .{ .what = "the 20-row graph (four chains) as one window in slot 4 at 3,060 (the ring's wrap, a page edge)", .chains = &.{.{ 4, 3060, 20 }} }, +}; + +/// Graphs of several row counts replayed in another order than their captures (5, 8, 20, 48, 4, 17, 1). +const any_order = [_]Replay{ + .{ .what = "48 rows", .chains = &.{ .{ 5, 3040, 24 }, .{ 7, 1000, 24 } } }, + .{ .what = "1 row", .chains = &.{.{ 3, 53, 1 }} }, + .{ .what = "17 rows", .chains = &.{.{ 9, 4080, 17 }} }, + .{ .what = "4 rows", .chains = &.{ .{ 0, 500, 1 }, .{ 1, 3080, 3 } } }, + .{ .what = "1 row again", .chains = &.{.{ 2, 1030, 1 }} }, + .{ .what = "8 rows", .chains = &.{.{ 8, 40, 8 }} }, +}; + +/// A replay's chains and `toks`' first rows as its ids; true when the graph's rows equal eager's. +fn replay(m: *const fwd.Model, r: Replay, toks: []const u32, got: []u8, want: []u8) !bool { + var chains: [specs.len]fwd.Chain = undefined; + var n: usize = 0; + for (r.chains, 0..) |c, i| { + chains[i] = .{ .slot = c[0], .p0 = c[1], .tokens = c[2] }; + n += c[2]; + } + return graphVsEager(m, r.what, chains[0..r.chains.len], toks[0..n], got, want); +} + +pub fn run(gpu: check.Gpu) !void { + const a = gpu.gpa; + const c = config(); + try wts.gate(&c); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var w = try build(gpu, c, 0x7269_7070); + defer w.deinit(); + var k = try fwd.Kernels.load(gpu.ctx); + defer k.deinit(); + // slots 0-4 the references', 5-9 the rounds' + var m = try fwd.Model.init(a, gpu.d, s, &w, &k, c.context, 2 * specs.len, round_rows); + defer m.deinit(); + var x = try Streams.init(a, c.vocab, 0x6677); + defer x.deinit(); + const logits = try a.alloc(u8, c.vocab * 2); + defer a.free(logits); + const rows = try a.alloc(u8, round_rows * c.vocab * 2); + defer a.free(rows); + const by_graph = try a.alloc(u8, round_rows * c.vocab * 2); + defer a.free(by_graph); + const by_fixed = try a.alloc(u8, round_rows * c.vocab * 2); + defer a.free(by_fixed); + var fails: usize = 0; + + // the references and B1's checks: eager one-row steps and rounds (graphs only where a check names them) + m.graphs = false; + for (0..specs.len) |j| fails += try reference(&m, &x, j, logits); + + // 1. a prompt chunk in another slot: the same bits + const base = specs.len; + var prompts: usize = 0; + for (specs, 0..) |sp, j| { + _ = try prefill(&m, base + j, x.ids[j][0..sp.prompt]); + try m.download(m.b.logits, logits); + if (!same("prompt of stream {s} in slot {d}", .{ sp.name, base + j }, logits, x.prompt[j])) { + std.debug.print("FAIL ripple-forward prompt of stream {s}: slot {d} != slot {d}\n", .{ sp.name, base + j, j }); + prompts += 1; + } + } + if (prompts == 0) check.pass("ripple-forward prompts: {d} streams' last rows the same in their test slots", .{specs.len}); + fails += prompts; + + // 2. rounds of several streams, verify windows and mixes + var done: [specs.len]usize = @splat(0); + fails += try runRounds(&m, &x, base, &done, logits, rows, by_graph, by_fixed); + + // 3. a slot reused after reset: B's (a wrapped ring, 7 pages) takes D, all its decode rows one verify window + const slot = base + 1; + try m.reset(slot, specs[1].prompt + decodeTokens(1)); + const d = 3; + const n = specs[d].prompt; + const t = decodeTokens(d); + _ = try prefill(&m, slot, x.ids[d][0..n]); + try m.download(m.b.logits, logits); + var bad: usize = @intFromBool(!same("reused slot's prompt", .{}, logits, x.prompt[d])); + // the window from a graph first, then B1's eager window into `rows` + const window_by_graph = try graphVsEager(&m, "reused slot window", &.{.{ .slot = slot, .p0 = n, .tokens = t }}, x.ids[d][n..][0..t], by_graph, rows); + for (0..t) |r| { + if (!same("reused slot row {d}", .{r}, rows[r * c.vocab * 2 ..][0 .. c.vocab * 2], x.want[d][r])) bad += 1; + } + if (bad == 0) { + check.pass("ripple-forward reused slot: stream D after B's reset, prompt and a {d}-row window == one-row steps", .{t}); + } else { + std.debug.print("FAIL ripple-forward reused slot: {d} rows differ\n", .{bad}); + fails += 1; + } + if (window_by_graph) { + check.pass("ripple-forward graph, reused slot: a {d}-row window replayed from a graph == the eager window", .{t}); + } else { + std.debug.print("FAIL ripple-forward graph, reused slot: a {d}-row window from a graph != the eager window\n", .{t}); + fails += 1; + } + + // 4. the one-row graph (the CLI's decode) by `forward`: captured once, replayed across the wrap and a page edge + const reused = base; + try m.reset(reused, specs[0].prompt + decodeTokens(0)); + var seq = try prefill(&m, reused, x.ids[1][0..specs[1].prompt]); + try m.download(m.b.logits, logits); + var steps: usize = @intFromBool(!same("B's prompt in slot {d}", .{reused}, logits, x.prompt[1])); + m.graphs = true; + const first = seq.len; + for (x.want[1], first..) |want, p| { + try m.forward(&seq, x.ids[1][p..][0..1], .decode); + try m.head(0); + try m.download(m.b.logits, logits); + if (!same("one-row graph step at {d}", .{p}, logits, want)) steps += 1; + } + m.graphs = false; + if (steps == 0) { + check.pass("ripple-forward one-row graph: captured in slot {d} at {d}, replayed to {d} (the ring's wrap, a page edge) == one-row eager steps", .{ reused, first, seq.len - 1 }); + } else { + std.debug.print("FAIL ripple-forward one-row graph: {d} of {d} steps differ from one-row eager steps\n", .{ steps, x.want[1].len }); + fails += 1; + } + + // 5. graphs captured above, replayed at other positions and slots + var toks: [round_rows]u32 = undefined; + var prng = std.Random.DefaultPrng.init(0x6772_6170); + for (&toks) |*tk| tk.* = prng.random().uintLessThan(u32, @intCast(c.vocab)); + for (elsewhere) |r| { + if (try replay(&m, r, &toks, by_graph, rows)) { + check.pass("ripple-forward graph elsewhere: {s} == eager", .{r.what}); + } else { + std.debug.print("FAIL ripple-forward graph elsewhere: {s} != eager\n", .{r.what}); + fails += 1; + } + } + + // 6. graphs of several row counts coexist: replayed in another order, none captured again + const graphs = m.captured(); + var order: usize = 0; + for (any_order) |r| order += @intFromBool(!try replay(&m, r, &toks, by_graph, rows)); + if (order == 0 and m.captured() == graphs) { + check.pass("ripple-forward graphs coexist: {d} row counts' graphs, replayed as 48, 1, 17, 4, 1 and 8 rows == eager, none captured again", .{graphs}); + } else { + std.debug.print("FAIL ripple-forward graphs coexist: {d} of {d} replays differ, {d} graphs before and {d} after\n", .{ order, any_order.len, graphs, m.captured() }); + fails += 1; + } + try check.expect(fails == 0, "ripple-forward: {d} checks differ", .{fails}); +} diff --git a/zig/tests/cuda/ripple_serve_tests.zig b/zig/tests/cuda/ripple_serve_tests.zig new file mode 100644 index 0000000000..55014f9f94 --- /dev/null +++ b/zig/tests/cuda/ripple_serve_tests.zig @@ -0,0 +1,484 @@ +//! Ripple served through its lane backend, bit for bit (ripple-serve): shared rounds, drafts, kept states, logprobs. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const rc = @import("ricursive_cuda"); +const check = @import("check.zig"); + +const fwd = rc.model; +const wts = rc.weights; +const be = rc.backend; +const snap = be.snap; +const Stream = lanes.Stream; +const Backend = lanes.backend.Backend; +const Window = lanes.backend.Window; +const Verified = lanes.backend.Verified; +const LogRow = lanes.LogprobRow; + +const fixture = @import("ripple_forward_tests.zig"); // the release-width weights on 10 layers + +const Spec = struct { name: []const u8, prompt: usize }; + +const specs = [_]Spec{ + .{ .name = "D", .prompt = 37 }, + .{ .name = "A", .prompt = 507 }, // decode rows across the 512-row page edge + .{ .name = "E", .prompt = 4093 }, // a wrapped ring; decode rows across 4,096 keys + .{ .name = "C", .prompt = 1019 }, // kept states at 510 and 514, about the first page edge +}; +const D = 0; +const A = 1; +const E = 2; +const C = 3; +/// The three streams decoding together (checks 1 and 2). +const shared = [_]usize{ D, A, E }; + +/// Draws a stream makes: the first (the prompt's last row), then one a round. +const steps = 8; +const max_new = 64; +const slots = 4; + +/// Each draw of a stream (the token at prompt length + i) and the bf16 logits it was drawn from. +const Trace = struct { + tokens: [steps]u32 = undefined, + logits: [steps][]u8 = undefined, + + fn init(gpa: std.mem.Allocator, vocab: usize) !Trace { + var t: Trace = .{}; + var made: usize = 0; + errdefer for (t.logits[0..made]) |l| gpa.free(l); + for (&t.logits) |*l| { + l.* = try gpa.alloc(u8, vocab * 2); + made += 1; + } + return t; + } + + fn deinit(t: *Trace, gpa: std.mem.Allocator) void { + for (t.logits) |l| gpa.free(l); + } +}; + +/// Draws and logits equal from step `from` on, or the first difference printed. +fn sameTrace(comptime what: []const u8, args: anytype, got: *const Trace, want: *const Trace, from: usize) bool { + for (from..steps) |i| { + if (got.tokens[i] == want.tokens[i] and std.mem.eql(u8, got.logits[i], want.logits[i])) continue; + const g = std.mem.bytesAsSlice(u16, got.logits[i]); + const w = std.mem.bytesAsSlice(u16, want.logits[i]); + var diff: usize = 0; + for (g, w) |a, b| diff += @intFromBool(a != b); + std.debug.print(" DIFF ", .{}); + std.debug.print(what, args); + std.debug.print(": step {d}: token {d} for {d}, {d} of {d} logits differ\n", .{ i, got.tokens[i], want.tokens[i], diff, g.len }); + return false; + } + return true; +} + +const Rig = struct { + gpa: std.mem.Allocator, + m: *fwd.Model, + x: *be.Ripple, + b: Backend, + + fn stream(r: Rig, id: []const u8, p: []const u32, sampling: ?lanes.Sampling) !Stream { + return Stream.init(r.gpa, .{ .id = id, .prompt = p, .max_new = max_new, .sampling = sampling, .drafts = false }); + } + + /// The prompt pass; the slot the stream took. + fn open(r: Rig, s: *Stream) !usize { + try r.b.prefill(s); + return r.x.lanes.getPtr(s).?.slot; + } + + /// The first draw and its logits (the prompt's last row, kept in the stream's slot). + fn first(r: Rig, s: *Stream, t: *Trace) !u32 { + const tok = try r.b.read(try r.b.first(s, s.prompt().len)); + t.tokens[0] = tok; + try r.m.download(r.x.lanes.getPtr(s).?.last, t.logits[0]); + return tok; + } + + /// One-row rounds by queue from step `from` (fed `tok`, the draw before it), each draw and its logits into `t`. + fn decode(r: Rig, s: *Stream, t: *Trace, from: usize, tok0: u32) !void { + var tok = tok0; + for (from..steps) |i| { + tok = try r.b.read(try r.b.queue(s, .{ .value = tok }, s.prompt().len + i)); + t.tokens[i] = tok; + try r.m.download(r.m.b.rows, t.logits[i]); + } + } + + /// Stream `s` alone from its prompt (or its kept state); the slot it took. + fn alone(r: Rig, s: *Stream, t: *Trace) !usize { + const slot = try r.open(s); + defer r.b.release(s); + try r.decode(s, t, 1, try r.first(s, t)); + return slot; + } + + /// Streams `ss` together: their prompt passes, first draws, then one-row windows by verify and keep (no drafts). + fn together(r: Rig, ss: []const *Stream, ts: []Trace) !void { + const v2 = r.m.c.vocab * 2; + var opened: usize = 0; + defer for (ss[0..opened]) |s| r.b.release(s); + for (ss) |s| { + _ = try r.open(s); + opened += 1; + } + var pending: [shared.len]u32 = undefined; + for (ss, ts, 0..) |s, *t, j| pending[j] = try r.first(s, t); + var windows: [shared.len]Window = undefined; + var out: [shared.len]Verified = undefined; + var pos: [shared.len][1]u64 = undefined; + var sampled: [shared.len][1]u32 = undefined; + var drafts: [shared.len][1]u32 = undefined; + const row0 = [_]u32{0}; + const paths: [shared.len][]const u32 = @splat(&row0); + const n = ss.len; + for (1..steps) |i| { + for (ss, 0..) |s, j| { + pos[j] = .{s.prompt().len + i}; + windows[j] = .{ .stream = s, .pending = pending[j], .held = 0, .tokens = &.{}, .parents = null, .positions = &pos[j] }; + out[j] = .{ .sampled = &sampled[j], .drafts = drafts[j][0..0] }; + } + try r.b.verify(windows[0..n], out[0..n]); + for (ts[0..n], 0..) |*t, j| { + try r.m.download(r.m.b.rows + j * v2, t.logits[i]); + t.tokens[i] = sampled[j][0]; + pending[j] = sampled[j][0]; + } + try r.b.keep(windows[0..n], paths[0..n]); + } + } +}; + +/// Check 1 and 2: three streams together == each alone (`sp` null: greedy), their slots 0 to 2 against slot 0 alone. +fn sharedRounds(r: Rig, prompts: []const []u32, alone: []const Trace, sp: ?lanes.Sampling, label: []const u8) !usize { + var ss: [shared.len]Stream = undefined; + var made: usize = 0; + defer for (ss[0..made]) |*s| s.deinit(r.gpa); + var ptrs: [shared.len]*Stream = undefined; + for (shared, &ss, &ptrs) |j, *s, *p| { + var each = sp; + if (each) |*q| q.seed +%= j; + s.* = try r.stream(specs[j].name, prompts[j], each); + made += 1; + p.* = s; + } + var ts: [shared.len]Trace = undefined; + var traced: usize = 0; + defer for (ts[0..traced]) |*t| t.deinit(r.gpa); + for (&ts) |*t| { + t.* = try Trace.init(r.gpa, r.m.c.vocab); + traced += 1; + } + try r.together(&ptrs, &ts); + var bad: usize = 0; + for (shared, &ts) |j, *t| { + if (!sameTrace("{s} stream {s} together vs alone", .{ label, specs[j].name }, t, &alone[j], 0)) bad += 1; + } + if (bad == 0) { + check.pass("ripple-serve {s} shared rounds: streams D, A, E ({d}, {d}, {d} prompt tokens) together == each alone, {d} draws and their logits each", .{ label, specs[D].prompt, specs[A].prompt, specs[E].prompt, steps }); + return 0; + } + std.debug.print("FAIL ripple-serve {s} shared rounds: {d} of {d} streams differ from alone\n", .{ label, bad, shared.len }); + return 1; +} + +/// Check 3: a 7-row window of copied drafts with a wrong fourth keeps 4 rows; draws then == plain decode. +fn copiedDrafts(r: Rig, p: []const u32, plain: *const Trace) !usize { + const v = r.m.c.vocab; + var s = try r.stream("copied", p, null); + defer s.deinit(r.gpa); + var t = try Trace.init(r.gpa, v); + defer t.deinit(r.gpa); + _ = try r.open(&s); + defer r.b.release(&s); + const n = p.len; + const t0 = try r.first(&s, &t); + const wrong = 4; // the draft for prompt length + 4 + var drafts: [6]u32 = undefined; + for (&drafts, 1..) |*d, i| d.* = plain.tokens[i]; + drafts[wrong - 1] = @intCast((plain.tokens[wrong] + 1) % v); + var pos: [7]u64 = undefined; + for (&pos, 1..) |*q, i| q.* = n + i; + var sampled: [7]u32 = undefined; + var dr: [6]u32 = undefined; + const w: Window = .{ .stream = &s, .pending = t0, .held = 0, .tokens = &drafts, .parents = null, .positions = &pos }; + var out = [_]Verified{.{ .sampled = &sampled, .drafts = &dr }}; + try r.b.verify(&.{w}, &out); + // rows 0 to 3 read plain tokens: their draws are plain steps 1 to 4; row 3's is the bonus past the wrong draft + for (0..wrong) |row| { + t.tokens[row + 1] = sampled[row]; + try r.m.download(r.m.b.rows + row * v * 2, t.logits[row + 1]); + } + const kept = [_]u32{ 0, 1, 2, 3 }; + try r.b.keep(&.{w}, &.{&kept}); + try r.decode(&s, &t, wrong + 1, sampled[wrong - 1]); + const ok = std.mem.eql(u32, &dr, &drafts) and sameTrace("copied drafts", .{}, &t, plain, 0); + if (ok) { + check.pass("ripple-serve copied drafts: stream A, a 7-row window (positions {d} to {d}) with a wrong 4th draft keeps 4 rows; draws and logits == plain decode, and the next 3 steps too", .{ n + 1, n + 7 }); + return 0; + } + std.debug.print("FAIL ripple-serve copied drafts: stream A's window or the steps after it differ from plain decode\n", .{}); + return 1; +} + +/// The prompt cache's hook: a state at each mark the pass stands at. +const Keep = struct { + x: *be.Ripple, + saved: [2]?*snap.Saved = .{ null, null }, + n: usize = 0, + + fn hook(ptr: *anyopaque, s: *Stream, at: u32) void { + const k: *Keep = @ptrCast(@alignCast(ptr)); + if (k.n == k.saved.len) return; + k.saved[k.n] = k.x.saveState(s, at) catch |e| blk: { + std.debug.print(" saveState at {d}: {s}\n", .{ at, @errorName(e) }); + break :blk null; + }; + k.n += 1; + } +}; + +/// Check 4: states kept at `ats`, each restored into another slot and continued == the uninterrupted `plain`. +fn keptStates(r: Rig, name: []const u8, p: []const u32, ats: []const u32, notes: []const []const u8, plain: *const Trace) !usize { + var keep: Keep = .{ .x = r.x }; + var saver = try Stream.init(r.gpa, .{ .id = "saver", .prompt = p, .max_new = max_new, .drafts = false, .reuse = .{ .marks = ats, .hook = .{ .ptr = &keep, .at = Keep.hook } } }); + defer saver.deinit(r.gpa); + const from = try r.open(&saver); + defer r.b.release(&saver); // held while the restores run, so each lands in another slot + defer for (keep.saved[0..keep.n]) |k| if (k) |saved| r.x.dropState(saved); + var bad: usize = 0; + for (ats, notes, 0..) |at, note, i| { + const saved = (if (i < keep.n) keep.saved[i] else null) orelse { + std.debug.print("FAIL ripple-serve kept state at {d} of stream {s}: none kept\n", .{ at, name }); + bad += 1; + continue; + }; + var s = try Stream.init(r.gpa, .{ .id = "restored", .prompt = p, .max_new = max_new, .drafts = false, .reuse = .{ .saved = saved, .at = at } }); + defer s.deinit(r.gpa); + var t = try Trace.init(r.gpa, r.m.c.vocab); + defer t.deinit(r.gpa); + const into = try r.alone(&s, &t); + const ok = !s.reuse_failed and s.cached == at and into != from and sameTrace("state at {d}", .{at}, &t, plain, 0); + if (ok) { + check.pass("ripple-serve kept state at {d} of stream {s} ({s}, {d:.1} MiB): saved from slot {d}, restored into slot {d}, {d} more prompt rows and {d} draws == uninterrupted", .{ at, name, note, @as(f64, @floatFromInt(snap.bytes(r.m.c, at))) / (1 << 20), from, into, p.len - at, steps }); + } else { + std.debug.print("FAIL ripple-serve kept state at {d} of stream {s}: reuse_failed {s}, cached {d}, slots {d} -> {d}, or draws differ\n", .{ at, name, if (s.reuse_failed) "yes" else "no", s.cached, from, into }); + bad += 1; + } + } + return bad; +} + +/// The Nemotron CUDA logprobs' bound against a float64 log-softmax of the same logits (CHANGELOG, #572). +const tol: f64 = 2e-6; + +/// A float64 log-softmax of one bf16 logits row. +const Host = struct { + vals: []align(1) const u16, + lse: f64, + + fn init(row: []const u8) Host { + var h: Host = .{ .vals = std.mem.bytesAsSlice(u16, row), .lse = 0 }; + var most = -std.math.inf(f64); + for (0..h.vals.len) |i| most = @max(most, h.at(i)); + var sum: f64 = 0; + for (0..h.vals.len) |i| sum += @exp(h.at(i) - most); + h.lse = most + @log(sum); + return h; + } + + fn at(h: Host, i: usize) f64 { + const f: f32 = @bitCast(@as(u32, h.vals[i]) << 16); + return f; + } + + fn lp(h: Host, i: usize) f64 { + return h.at(i) - h.lse; + } + + /// The best ids, logit descending, ties by the lower id. + fn best(h: Host, out: []u32) void { + var n: usize = 0; + for (0..h.vals.len) |i| { + const v = h.at(i); + var j = n; + while (j > 0 and h.at(out[j - 1]) < v) j -= 1; + if (j == out.len) continue; + var q = @min(n, out.len - 1); + while (q > j) : (q -= 1) out[q] = out[q - 1]; + out[j] = @intCast(i); + n = @min(n + 1, out.len); + } + } + + /// vLLM's rank: the ids whose logit is at or above the pick's, itself included. + fn rank(h: Host, pick: u32) u32 { + var n: u32 = 0; + for (0..h.vals.len) |i| n += @intFromBool(h.at(i) >= h.at(pick)); + return n; + } +}; + +/// A backend row against the host's: the pick, its best ids exactly, their log probabilities within `tol`. +fn rowMatches(h: Host, row: LogRow, pick: u32, k: u8, ranked: bool, worst: *f64) bool { + var ok = row.token == pick and row.count == k; + const e0 = @abs(@as(f64, row.logprob) - h.lp(pick)); + worst.* = @max(worst.*, e0); + ok = ok and e0 <= tol; + var ids: [lanes.logprob.max_top]u32 = undefined; + h.best(ids[0..k]); + for (0..k) |i| { + const e = @abs(@as(f64, row.logprobs[i]) - h.lp(ids[i])); + worst.* = @max(worst.*, e); + ok = ok and row.ids[i] == ids[i] and e <= tol; + } + if (ranked) ok = ok and row.rank == h.rank(pick); + return ok; +} + +fn verdict(ok: bool, comptime what: []const u8, args: anytype, worst: f64) usize { + std.debug.print("{s} ripple-serve logprob rows: ", .{if (ok) "PASS" else "FAIL"}); + std.debug.print(what, args); + if (ok) { + std.debug.print(" == a float64 log-softmax (ids exact, worst {e:.2} <= {e:.0})\n", .{ worst, tol }); + return 0; + } + std.debug.print(" differ from a float64 log-softmax (worst {e:.2}, bound {e:.0})\n", .{ worst, tol }); + return 1; +} + +/// Check 5: the first draw's row and a verify window's rows (20 best), and prompt rows with ranks (5 best). +fn logprobRows(r: Rig, pa: []const u32, plain: *const Trace, pe: []const u32) !usize { + const v = r.m.c.vocab; + const row_bytes = v * 2; + const k: u8 = lanes.logprob.max_top; + var bad: usize = 0; + const buf = try r.gpa.alloc(u8, fwd.head_rows * row_bytes); + defer r.gpa.free(buf); + { + var s = try Stream.init(r.gpa, .{ .id = "logprobs", .prompt = pa, .max_new = max_new, .drafts = false, .logprobs = k }); + defer s.deinit(r.gpa); + var t = try Trace.init(r.gpa, v); + defer t.deinit(r.gpa); + _ = try r.open(&s); + defer r.b.release(&s); + const t0 = try r.first(&s, &t); + var worst: f64 = 0; + const row = try r.b.firstRow(&s); + bad += verdict(rowMatches(Host.init(t.logits[0]), row, t0, k, false, &worst), "the first draw's row (stream A, {d} best)", .{k}, worst); + + // a 4-row window of plain drafts, every row scored + const n = pa.len; + const drafts = plain.tokens[1..4]; + var pos: [4]u64 = undefined; + for (&pos, 1..) |*q, i| q.* = n + i; + var sampled: [4]u32 = undefined; + var dr: [3]u32 = undefined; + var rows: [4]LogRow = undefined; + const w: Window = .{ .stream = &s, .pending = t0, .held = 0, .tokens = drafts, .parents = null, .positions = &pos }; + var out = [_]Verified{.{ .sampled = &sampled, .drafts = &dr, .rows = &rows }}; + try r.b.verify(&.{w}, &out); + try r.m.download(r.m.b.rows, buf[0 .. 4 * row_bytes]); + worst = 0; + var ok = true; + for (rows, sampled, 0..) |lr, pick, i| ok = rowMatches(Host.init(buf[i * row_bytes ..][0..row_bytes]), lr, pick, k, false, &worst) and ok; + bad += verdict(ok, "a 4-row verify window's rows (stream A at {d} to {d}, {d} best)", .{ n + 1, n + 4, k }, worst); + } + { + // prompt rows: the reference logits from the same prompt in a slot no stream holds, head_rows at a time + const kp: u8 = 5; + const slot = slots - 1; + const scored = pe.len - 1; + const want = try r.gpa.alloc(u8, scored * row_bytes); + defer r.gpa.free(want); + try r.m.prompt(slot, 0, pe); + var r0: usize = 0; + while (r0 < scored) : (r0 += fwd.head_rows) { + const rows = @min(fwd.head_rows, scored - r0); + try r.m.headRows(r0, rows); + try r.m.download(r.m.b.rows, want[r0 * row_bytes ..][0 .. rows * row_bytes]); + } + try r.m.reset(slot, pe.len); + var s = try Stream.init(r.gpa, .{ .id = "prompt-logprobs", .prompt = pe, .max_new = max_new, .drafts = false, .prompt_logprobs = kp }); + defer s.deinit(r.gpa); + try r.b.prefill(&s); + r.b.release(&s); + var worst: f64 = 0; + var ok = s.prompt_rows.items.len == scored; + if (ok) { + for (s.prompt_rows.items, 0..) |lr, i| ok = rowMatches(Host.init(want[i * row_bytes ..][0..row_bytes]), lr, pe[i + 1], kp, true, &worst) and ok; + } + bad += verdict(ok, "{d} prompt rows with ranks ({d} best, {d}-row head batches)", .{ scored, kp, fwd.head_rows }, worst); + } + return bad; +} + +pub fn run(gpu: check.Gpu) !void { + const a = gpu.gpa; + const c = fixture.config(); + try wts.gate(&c); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + const t0 = check.now(gpu.io); + var w = try fixture.build(gpu, c, 0x7269_7070); + defer w.deinit(); + std.debug.print("ripple-serve: weights built in {d:.1} s\n", .{@as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e9}); + var k = try fwd.Kernels.load(gpu.ctx); + defer k.deinit(); + var m = try fwd.Model.init(a, gpu.d, s, &w, &k, c.context, slots, be.round_rows); + defer m.deinit(); + var x = try be.Ripple.init(a, gpu.d, &m); + defer x.deinit(); + const r: Rig = .{ .gpa = a, .m = &m, .x = &x, .b = x.backend() }; + + var prompts: [specs.len][]u32 = undefined; + var made: usize = 0; + defer for (prompts[0..made]) |p| a.free(p); + var prng = std.Random.DefaultPrng.init(0x7365_7276); + for (specs, &prompts) |sp, *p| { + p.* = try a.alloc(u32, sp.prompt); + made += 1; + for (p.*) |*t| t.* = prng.random().uintLessThan(u32, @intCast(c.vocab)); + } + const sampling: lanes.Sampling = .{ .seed = 0x5eed, .temperature = 1.0, .top_p = 0.95, .top_k = 20 }; + + // each stream alone in a fresh slot, greedy (all four) and sampled (the shared three, a seed each) + var greedy: [specs.len]Trace = undefined; + var sampled: [specs.len]Trace = undefined; + var traced: usize = 0; + defer for (greedy[0..traced], sampled[0..traced]) |*g, *q| { + g.deinit(a); + q.deinit(a); + }; + for (&greedy, &sampled) |*g, *q| { + g.* = try Trace.init(a, c.vocab); + q.* = Trace.init(a, c.vocab) catch |e| { + g.deinit(a); + return e; + }; + traced += 1; + } + for (specs, prompts, &greedy, 0..) |sp, p, *g, j| { + var st = try r.stream(sp.name, p, null); + defer st.deinit(a); + _ = try r.alone(&st, g); + if (std.mem.indexOfScalar(usize, &shared, j) == null) continue; + var each = sampling; + each.seed +%= j; + var sq = try r.stream(sp.name, p, each); + defer sq.deinit(a); + _ = try r.alone(&sq, &sampled[j]); + } + + var fails: usize = 0; + fails += try sharedRounds(r, &prompts, &greedy, null, "greedy"); + fails += try sharedRounds(r, &prompts, &sampled, sampling, "sampled (temperature 1, top-p 0.95, top-k 20)"); + fails += try copiedDrafts(r, prompts[A], &greedy[A]); + fails += try keptStates(r, "E", prompts[E], &.{3500}, &.{"past the 3,072-row ring's wrap"}, &greedy[E]); + fails += try keptStates(r, "C", prompts[C], &.{ 510, 514 }, &.{ "two rows before a 512-row page edge", "two rows across a 512-row page edge" }, &greedy[C]); + fails += try logprobRows(r, prompts[A], &greedy[A], prompts[E][0..300]); + try check.expect(fails == 0, "ripple-serve: {d} checks differ", .{fails}); +} diff --git a/zig/tests/ricursive/run.zig b/zig/tests/ricursive/run.zig new file mode 100644 index 0000000000..89d892c954 --- /dev/null +++ b/zig/tests/ricursive/run.zig @@ -0,0 +1,236 @@ +//! Ripple on Metal on real weights: greedy tokens, every position's logits or forced steps, for the reference's check. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const fr = tf.ricursive; + +const usage = + \\tf-ricursive-run greedy the prompt in chunks, then n one-row steps' tokens + \\tf-ricursive-run logits 16-row windows from 0; bf16 logits [tokens, vocab] + \\tf-ricursive-run chunked [,] every row's logits through prompt chunks of rows + \\ (2048): the tensor-unit path's, [tokens, vocab] + \\tf-ricursive-run prompt the prompt in chunks; the last row's bf16 logits + \\tf-ricursive-run forced , then cont's tokens a step; bf16 logits before each + \\tf-ricursive-run exact [,] cont's tokens as one-row steps and as w-row windows + \\ (16) after the prompt: every row's logits bit for bit + \\tf-ricursive-run bench the prompt, then n one-row steps: wall, encode, GPU ms + \\ +; + +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 io = init.io; + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const prompt_json = try readIds(gpa, io, args[2]); + defer prompt_json.deinit(); + const prompt = prompt_json.value; + if (prompt.len == 0) return error.EmptyPrompt; + const mode = args[3]; + const arg = args[4]; + const m = try fr.Model.load(gpa, io, args[1]); + defer m.deinit(); + var small = try fr.state.Scratch.init(m.device, &m.config, fr.state.batch_rows); + defer small.deinit(); + var wide = try fr.state.Scratch.init(m.device, &m.config, fr.state.prompt_rows); + defer wide.deinit(); + const vocab: usize = m.config.vocab; + var t0 = now(); + if (eql(mode, "logits") or eql(mode, "chunked")) { + var rows: usize = fr.state.window_rows; + var path: []const u8 = arg; + if (eql(mode, "chunked")) { + rows = fr.state.prompt_rows; + if (std.mem.indexOfScalar(u8, arg, ',')) |comma| { + path = arg[0..comma]; + rows = try std.fmt.parseInt(usize, arg[comma + 1 ..], 10); + } + if (rows <= fr.state.batch_rows or rows > fr.state.prompt_rows) return error.ChunkRowsOutOfRange; + } + var cache = try fr.state.Cache.init(m.device, &m.config, prompt.len + 1); + defer cache.deinit(); + 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, rows); + if (n <= fr.state.batch_rows) { + try fr.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)); + } else { + // a chunk's every row: the chunk on the prompt path, its rows' logits one 32-row head pass at a time + try fr.forward.run(m, &wide, &cache, prompt[at..][0..n], .{}); + try headRows(m, &wide, &small, out[at * vocab ..][0 .. n * vocab]); + } + cache.len += n; + at += n; + } + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = path, .data = std.mem.sliceAsBytes(out) }); + std.debug.print("ricursive {s}: {d} rows in {d} ms\n", .{ mode, prompt.len, (now() - t0) / std.time.ns_per_ms }); + return; + } + if (eql(mode, "exact")) return exact(gpa, io, m, &small, &wide, prompt, arg); + var forced: ?std.json.Parsed([]u32) = null; + defer if (forced) |f| f.deinit(); + var forced_out: []const u8 = ""; + if (eql(mode, "forced")) { + const comma = std.mem.indexOfScalar(u8, arg, ',') orelse return error.ExpectedContCommaOut; + forced = try readIds(gpa, io, arg[0..comma]); + forced_out = arg[comma + 1 ..]; + } + const steps: usize = if (eql(mode, "greedy") or eql(mode, "bench")) try std.fmt.parseInt(usize, arg, 10) else if (forced) |f| f.value.len else 0; + var cache = try fr.state.Cache.init(m.device, &m.config, prompt.len + steps + fr.state.window_rows + 1); + defer cache.deinit(); + try feed(m, &small, &wide, &cache, prompt); + const ttft = now() - t0; + if (eql(mode, "prompt")) { + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = arg, .data = std.mem.sliceAsBytes(cache.logits.slice(u16, vocab)) }); + std.debug.print("ricursive 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 fr.forward.run(m, &small, &cache, &.{tok}, .{}); + cache.len += 1; + } + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = forced_out, .data = std.mem.sliceAsBytes(out) }); + return; + } + if (eql(mode, "bench")) { + var tm: fr.forward.Timing = .{}; + t0 = now(); + for (0..steps) |_| { + try fr.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(steps)); + const per = 1e3 / @as(f64, @floatFromInt(steps)); + std.debug.print("ricursive 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 (!eql(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 fr.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("ricursive 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)) }); +} + +/// cont's tokens after the prompt as one-row steps and as w-row windows: every row's logits must agree bit for bit. +fn exact(gpa: std.mem.Allocator, io: std.Io, m: *fr.Model, small: *fr.state.Scratch, wide: *fr.state.Scratch, prompt: []const u32, arg: []const u8) !void { + const vocab = m.config.vocab; + var parts = std.mem.splitScalar(u8, arg, ','); + const cont = try readIds(gpa, io, parts.first()); + defer cont.deinit(); + const toks = cont.value; + const width: usize = if (parts.next()) |w| try std.fmt.parseInt(usize, w, 10) else fr.state.window_rows; + if (width == 0 or width > fr.state.window_rows) return error.BadWindow; + // both caches at one capacity: the decode blocks follow it + var scache = try fr.state.Cache.init(m.device, &m.config, prompt.len + toks.len + 1); + defer scache.deinit(); + var wcache = try fr.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 fr.forward.run(m, small, &scache, &.{tok}, .{ .all = true }); + @memcpy(one[i * vocab ..][0..vocab], small.logits_all.slice(u16, vocab)); + scache.len += 1; + } + 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 fr.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 differ: usize = 0; + for (0..toks.len) |r| { + if (std.mem.eql(u16, one[r * vocab ..][0..vocab], wins[r * vocab ..][0..vocab])) continue; + if (differ < 4) std.debug.print("row {d} (window row {d}) differs\n", .{ r, r % width }); + differ += 1; + } + std.debug.print("ricursive exact at {d}: {d} of {d} rows' logits differ between one-row steps and {d}-row windows\n", .{ prompt.len, differ, toks.len, width }); + if (differ > 0) return error.RowsDiffer; +} + +/// Every row of the chunk just run on `wide`: its final-normed rows through the head, 32 at a time on `small`'s logits. +fn headRows(m: *const fr.Model, wide: *fr.state.Scratch, small: *fr.state.Scratch, out: []u16) !void { + const vocab = m.config.vocab; + const h = m.config.hidden; + const n = out.len / vocab; + var at: usize = 0; + while (at < n) { + const r = @min(n - at, fr.state.batch_rows); + @memcpy(small.x.slice(u16, r * h), wide.x.slice(u16, (at + r) * h)[at * h ..]); + try fr.forward.head(m, small, r); + @memcpy(out[at * vocab ..][0 .. r * vocab], small.logits_all.slice(u16, r * vocab)); + at += r; + } +} + +/// The prompt in prompt_rows chunks (the tensor-unit path past 32 rows), as the backend feeds it. +fn feed(m: *const fr.Model, small: *fr.state.Scratch, wide: *fr.state.Scratch, cache: *fr.state.Cache, ids: []const u32) !void { + var at: usize = 0; + while (at < ids.len) { + const n = @min(ids.len - at, fr.state.prompt_rows); + try fr.forward.run(m, if (n > fr.state.batch_rows) wide else small, cache, ids[at..][0..n], .{}); + cache.len += n; + at += n; + } +} + +fn readIds(gpa: std.mem.Allocator, io: std.Io, path: []const u8) !std.json.Parsed([]u32) { + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(64 << 20)); + defer gpa.free(bytes); + return std.json.parseFromSlice([]u32, gpa, bytes, .{}); +} + +fn eql(a: []const u8, b: []const u8) bool { + return std.mem.eql(u8, a, b); +} + +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); +}