diff --git a/CHANGELOG.md b/CHANGELOG.md index 3ded82e8c44..46cc113ffba 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -72,10 +72,15 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Qwen3.8-27B on Metal: prompt passes no longer write the GDN tails buffer, which nothing read. Replies are unchanged; on an M5 Ultra, time to first token is about 0.5% shorter at 2K, 8K and 32K prompts, one and four streams (#2915, for #2905). - Flash Next on CUDA at `--tp 2`: a prompt pass wider than 64 rows reduce-scatters its routed experts' FP32 slots by row halves, sums each rank's rows and gathers the sums, where it all-reduced every slot; decode rounds and verify windows keep the all-reduce. Every reduced value is the same single add and the slots are summed in the writeback's order, so replies keep their tokens. Over two GB10s' RoCE link rank 0 sends 3.43 MiB a prompt token instead of 5.81, and cold prefill of 2k to 32k tokens is 26.5% to 26.8% faster there and 30.9% to 33.9% faster on two RTX PRO 6000s (an 8,192-token prompt: 7.12 to 5.62 s and 3.61 to 2.70 s) (#2977). Thanks to @swbiggart. +- Vulkan: the native dense Qwen3.5-2B greedy engine and its operator gates, the second part of the opt-in Vulkan backend built only with `-Dvulkan=true`: checked BF16 tensor operations, the norm/MLP, attention (RoPE, causal GQA SDPA, output gate) and GatedDeltaNet gates, one GDN path for prefill and decode, prefill chunked within the context ceiling, and the kernel contract extended to every operator's SPIR-V through gates 1 to 4 with planted controls and recorded device findings (residual denormals, the preserve-add subnormal flush, the gated norm at the production per-head geometry). `zig build vulkan-kernel-contract -Dvulkan=true` runs it; default builds are unchanged (#1544). Thanks to @joshuaswarren. + +- The Vulkan engine reads Qwen3.5's 4-bit projection weights and activations 16 bytes at a time (`qmm_q4_vec4`) when both are 16-byte aligned, and keeps the scalar kernel otherwise. Each output keeps the same sum in the same order, so the logits are byte-identical. On a 13-inch M1 with Qwen3.5-2B 4-bit, a run (prefill, 63 decode steps and the per-step logits readback) took 34.8% to 43.0% less time than the same build with the scalar kernel selected, on three inputs: a 24-token prompt, and documents of 256 and 1,024 tokens (Linux 7.3.0-2-20261110.1-lab-ARCH; the first measurement, against #1544's head on an earlier kernel and Mesa, gave 40% to 52%). The kernel contract checks the new kernel with the same gates as `qmm_bf16`. Receipts: [docs/receipts/vulkan-qmm-vec4-m1](docs/receipts/vulkan-qmm-vec4-m1/README.md). - Tests: `tf-cuda-test flashnext-rank RANK PORT timeout` passes again on two GPUs. NCCL start-up keeps its 30 s bound and only the stalled collective gets 200 ms, and both ranks meet once their communicators are up, so `timeout` and `eof` fail the collective under test, not the start-up (#2976, fixes #2975). Thanks to @swbiggart. +- The Vulkan engine stores each Qwen3.5 4-bit projection in tiles of 32 columns, repacked once at load, and reads them with `qmm_q4_tile`: a group of 32 lanes then loads one contiguous 512-byte run instead of 32 rows a kilobyte apart. The loop, the sums and their order are `qmm_q4_vec4`'s, so the logits are byte-identical. A projection is padded to whole tiles, and the tied output head is a second, tiled copy of the embedding table (about 320 MB more device memory). On a 13-inch M1 with Qwen3.5-2B 4-bit, a run took 10.9% less time than with `qmm_q4_vec4` on a 24-token prompt, 35.6% less on a 256-token document and 44.9% less on a 1,024-token document (Linux 7.3.0-2-20261110.1-lab-ARCH; earlier builds on an earlier kernel and Mesa gave 19%, 50% and 63%). The kernel contract checks `qmm_q4_tile` with the same gates as `qmm_bf16`. Receipts: [docs/receipts/vulkan-qmm-tile-m1](docs/receipts/vulkan-qmm-tile-m1/README.md). - Tokenizer: a split pattern is held to a work allowance per piece (2,048 instructions a byte, counted, not timed) and refused at load when too deep or too large, so a crafted tokenizer or prompt can't hang the pre-tokenizer; every split of the patterns we serve is unchanged (#3071). +- The Vulkan engine records its dispatches into one command buffer and submits when 128 are waiting, or when the next upload, download or fill needs the queue, instead of one submission and fence wait per dispatch. Each pipeline holds a ring of 256 descriptor sets. No kernel changes, so the logits are byte-identical. On a 13-inch M1 with Qwen3.5-2B 4-bit and tiled projections, a run took 56.0% less time on a 24-token prompt, 42.0% less on a 256-token document and 30.1% less on a 1,024-token document. On a 16-inch M1 Max, it took 84.1%, 82.5% and 61.8% less time on those inputs; the 256-token cell was noisy. Receipts: [docs/receipts/vulkan-submit-batch-m1](docs/receipts/vulkan-submit-batch-m1/README.md). - Cluster pull: a pull that would leave a file held only in part is refused unless the caller allows unverified bytes (off by default); a pull that ends whole still resumes and is checked against the manifest's hash (#3074). @@ -134,6 +139,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Metal: kernels nothing launches are removed: Qwen3.8-27B's `prompt_glue.metal` and `prompt_state.metal` (six `q27_prompt_*` kernels) and core's `tf_lane_rope` with `Attention.rotate`; saved Qwen3.8-27B prompt states on Metal are recomputed once after upgrading, since their identity covers the kernel sources (#3065, part of #2251). Thanks to @jschmied. - Tests: main's nightly is green again. The Flash Next learn, prompt and compare tools build again on Linux (#2585 had put one file in two modules), and the dense serving and serving-identity pins take the changes #2579, #2585, #2285, #2205 and #2789 made, each named beside its pin. Serving is unchanged (#3002). - HIP: Qwen3.8-Flash-Next's MTP head takes a prompt chunk on the prompt tiles, not a decode window at a time, so cold prompts finish about 5% sooner at 8K and 32K on a Radeon 8060S (gfx1151). Replies are unchanged; the head's caches and drafts are not bit-identical to before (#1878). Thanks to @gilby. +- The Vulkan engine stores each Qwen3.5 4-bit projection in tiles of 32 columns, repacked once at load, and reads them with `qmm_q4_tile`: a group of 32 lanes then loads one contiguous 512-byte run instead of 32 rows a kilobyte apart. The loop, the sums and their order are `qmm_q4_vec4`'s, so the logits are byte-identical. A projection is padded to whole tiles, and the tied output head is a second, tiled copy of the embedding table (about 320 MB more device memory). On a 13-inch M1 with Qwen3.5-2B 4-bit, a run took 19% less time than with `qmm_q4_vec4` on a 24-token prompt, 50% less on a 256-token document and 63% less on a 1,024-token document. The kernel contract checks `qmm_q4_tile` with the same gates as `qmm_bf16`. Receipts: [docs/receipts/vulkan-qmm-tile-m1](docs/receipts/vulkan-qmm-tile-m1/README.md). +- The Vulkan engine records its dispatches into one command buffer and submits when 128 are waiting, or when the next upload, download or fill needs the queue, instead of one submission and fence wait per dispatch. Each pipeline holds a ring of 256 descriptor sets. No kernel changes, so the logits are byte-identical. On a 13-inch M1 with Qwen3.5-2B 4-bit and tiled projections, a run took 68% less time on a 24-token prompt, 49% less on a 256-token document and 30% less on a 1,024-token document. Receipts: [docs/receipts/vulkan-submit-batch-m1](docs/receipts/vulkan-submit-batch-m1/README.md). - Cortheon (`--cortheon`): DuckDuckGo results are kept, and its ads still skipped, when only the spelling differs (either quote style, a `class` attribute in any letter case, an unquoted `href`, spaces around its `=`, or a line break before it); a duplicate whose scheme or host differs only in letter case is dropped while paths and queries keep their case, and the PyPI, npm and crates.io allow-list takes the host in any case with the path exact (#2750). Thanks to @ppardus. @@ -1062,6 +1069,9 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## 1.2.2 +- The Nemotron XPU loader and the XPU safetensors reader read shard headers through core, so a malformed or oversized header is refused before anything is read (#1919). Thanks to @jschmied. + + - CUDA lanes: Gemma 4, MiniMax M3, Kolibri-1, Spark 2.5 and Dory can verify root and context copies as Nemotron does, opt-in: set `TF_GEMMA4_COPIES`, `TF_MINIMAX_COPIES`, `TF_KOLIBRI_COPIES`, `TF_SPARK_COPIES` or `TF_DORY_COPIES` to 1. Replies equal `"draft": false`, and unset is the old policy (#2022, part of #1910). - The dense Qwen3.5 Metal loader takes its shape from the checkpoint's config rather than the 27B's fixed geometry, checked against the 27B, 9B and 2B with identical tokens (#1817). @@ -1090,8 +1100,6 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - `zig build test-info-config` runs info's issue-934 test again: it was built from the CUDA CLI root, which never reached it (#2006). Thanks to @jishnuvenugopal. -- The Nemotron XPU loader and the XPU safetensors reader read shard headers through core, so a malformed or oversized header is refused before anything is read (#1919). Thanks to @jschmied. - - Windows: a cost-cache read that lands mid-replace is counted as a miss (safe: the window is priced again), so its concurrency test no longer flakes (#2001). - GLM on CUDA with TF_GLM_KV_SHARE no longer reports a balanced idle KV pool as (UNEVEN) at close; nothing else changes (#1993). Thanks to @Bizuayeu. diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index 7a72107791f..46a61b0ec87 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -71,12 +71,31 @@ The native engine carries that adaptation. `zig/kernels/metal/qwen3_5/gdn_pre.me `zig/kernels/metal/core/deltanet.metal` and Flash Next's `zig/kernels/metal/flashnext/q4_gdn_*.metal` run the same model math. -## Vulkan affine projections +## Vulkan shaders from omarchy-mlx + +These files in `zig/kernels/vulkan/` come from [omarchy-mlx](https://github.com/joshuaswarren/omarchy-mlx), +MIT License, Copyright © 2026 Joshua Warren and mlx-omarchy contributors. The first lines of each file name its +source path and commit. The native Vulkan path does not link MLX. + +- `argreduce_chunk.comp` +- `argreduce_combine.comp` +- `binary_vec.comp` +- `cast.comp` +- `conv_dw1d.comp` +- `copy_general.comp` +- `dequant.comp` +- `fast_norm.comp` +- `fast_norm_gated.comp` +- `fast_rope.comp` +- `gated_delta_decode.comp` +- `gated_delta_prefill_ash.comp` +- `gather_take.comp` +- `gdn_conv_decode.comp` +- `omarchy_trig.h` +- `qmm.comp` +- `sdpa_decode_native.comp` +- `swiglu.comp` -`zig/kernels/vulkan/qmm.comp` adapts the GLSL shader from -[mlx-omarchy](https://github.com/joshuaswarren/omarchy-mlx), MIT License, -Copyright © 2026 Joshua Warren and mlx-omarchy contributors. -Its arithmetic is unchanged; the native Vulkan path does not link MLX. See [the license text](LICENSES/mlx-omarchy-MIT.txt). ## Qwen Flash Next diff --git a/build.zig b/build.zig index c871eac8561..8dd47fe2606 100644 --- a/build.zig +++ b/build.zig @@ -218,7 +218,10 @@ pub fn build(b: *std.Build) void { if (with_hip) hip_build.steps(b, target, optimize, build_options, test_step, hip_options); const with_vulkan = b.option(bool, "vulkan", "Enable the opt-in Vulkan backend and its host/device targets") orelse false; const vulkan_options = vulkan_build.options(b); // declared on every target, used only with -Dvulkan=true - if (with_vulkan) vulkan_build.steps(b, target, optimize, test_step, vulkan_options); + if (with_vulkan) { + vulkan_build.steps(b, target, optimize, test_step, vulkan_options); + @import("zig/build/vulkan_model.zig").steps(b, target, optimize, test_step, vulkan_options, build_options); + } exlr8_build.hostTests(b, target, test_step); xpu_build.steps(b, target, optimize, build_options); video_build.steps(b, target, test_step); diff --git a/docs/credits.md b/docs/credits.md index a8a94cae80c..ae329811ff3 100644 --- a/docs/credits.md +++ b/docs/credits.md @@ -120,6 +120,14 @@ contributions are in it. ([lane_qmm.py](https://github.com/ashhart/TensorFold/blob/v0.6.6/src/tensorfold/kernels/qwen/dense/v1/lane_qmm.py)), widening each group to nibbles for the 4-bit tensor op. The native `zig/kernels/metal/core/lane_projection.metal` stages 2- and 3-bit codes the same way. +- The native Vulkan dense Qwen3.5-2B engine (`zig/src/families/qwen3_5/backend/vulkan/`) follows the same frozen + dense-lane design ([lane_qmm.py](https://github.com/ashhart/TensorFold/blob/v0.6.6/src/tensorfold/kernels/qwen/dense/v1/lane_qmm.py), + [#34](https://github.com/ashhart/TensorFold/pull/34)) for its per-lane qmm decomposition. Fifteen of its compute + shaders and `omarchy_trig.h` come from Joshua Warren's MIT-licensed omarchy-mlx unchanged (`argreduce_chunk`, + `argreduce_combine`, `binary_vec`, `cast`, `conv_dw1d`, `copy_general`, `dequant`, `fast_norm`, `fast_norm_gated`, + `fast_rope`, `gather_take`, `gated_delta_decode`, `gated_delta_prefill_oracle`, `gdn_conv_decode`, `sdpa_decode_native` + and `swiglu`); `qmm.comp` is derived from it. The rest (`attention_gate`, `finite_logits`, `gdn_split`, `qmm_q4_vec4` + and the contract controls) are new work, written for this engine. - @di37 made the engine read the GPU generation once from `applegpu_gN`, and a virtual machine's GPU as none, in [#209](https://github.com/ashhart/TensorFold/pull/209) ([device.py](https://github.com/ashhart/TensorFold/blob/v0.6.6/src/tensorfold/kernels/device.py)), as diff --git a/docs/receipts/vulkan-qmm-tile-m1/PREREG.md b/docs/receipts/vulkan-qmm-tile-m1/PREREG.md new file mode 100644 index 00000000000..71c13ad6f3f --- /dev/null +++ b/docs/receipts/vulkan-qmm-tile-m1/PREREG.md @@ -0,0 +1,48 @@ +# qmm_q4_tile pre-registration (written 2026-11-03 before any GPU run of this change) + +## Question + +Does storing each projection's weights, scales and biases in 32-column tiles (a load-time repack) make a Qwen3.5-2B run shorter +on the 13-inch M1, with byte-identical logits? The kernel is `qmm_q4_vec4` with only its address formulas changed, so every sum is the same. + +## Arms (two ELFs of `tf-qwen35-model`, ReleaseSafe, aarch64) + +- old: the vec4 build, commit `e742e023`, ELF `683fcf46...`. +- new: the same plus the tiled image: weights of every projection with n a multiple of 32 are repacked at load; the lm head is a second, + tiled copy of the embedding table (the embedding keeps the row layout). Projections with n = 16 keep the vec4 image. + +## Gates (every process, both arms) + +1. Exit 0 and a complete receipt (the house receipt check). +2. Generated ids equal the pinned ids. +3. Logits SHA-256 of every process equals one digest per case, the same digest the earlier vec4 runs gave (unicode-24 `e5447d9a...`, + document-256 `147a2d2e...`, document-1024 `7f9d8e65...`). +4. Before the pairs: the stage 1 probe on the GPU runs the tiled image on every tileable projection case and compares bytes to the default image. + +## Metric and design + +Run 2 duration from the probe's stderr milestones (prefill plus 63 decode steps plus the per-step logits readback), as before. +Three pairs per case, order old,new,new,old,old,new. Cases: unicode-24, document-256, document-1024. + +## Decision rule + +Speedup: median pair gain of at least 3 percent and all three pair gains positive and all gates hold. Otherwise reported as no speedup and not proposed. + +## Forecast (a prediction; I do not know the answer) + +Gain of 20 to 45 percent of run 2 on unicode-24 and similar on the documents. Why not more: qmm is about 63 percent of the run net of the submit +cost, the multi-row result showed qmm is not simply load-count limited, and the hardware cache may already hold part of the 1 KiB stride. +I give 20 percent to a gain under 3 percent, and 10 percent to a regression. + +## Not claimed + +Other chips, other models, other quantizations, the lm head on devices whose storage range is below the table size (they keep the row layout). +The software-device pass is not GPU evidence. Memory: the tied lm head adds about 320 MB of device memory. + +## Addendum, written after the final pairs ran (2026-11-03) + +The first contract run on the GPU refused the served shape 1x16x64: a registered kernel must serve every shape the contract serves, and the build +above kept the n = 16 projections on the vec4 image. The design changed before the final runs: the repack pads the last tile with zero columns, so the tiled +image serves every n, the n = 16 projections are tiled too, and the arms read "old: the vec4 build, new: the padded tile build". The gates, metric and +decision rule did not change. This addendum was written after the final pairs; the three earlier pairs on the unpadded build are kept in the receipts as an +earlier measurement and are not the result. The forecast above was for the unpadded build; it was not revised. diff --git a/docs/receipts/vulkan-qmm-tile-m1/README.md b/docs/receipts/vulkan-qmm-tile-m1/README.md new file mode 100644 index 00000000000..7300716e3f2 --- /dev/null +++ b/docs/receipts/vulkan-qmm-tile-m1/README.md @@ -0,0 +1,50 @@ +# qmm_q4_tile receipts, Apple M1 (13-inch) + +Qwen3.5-2B 4-bit, 64 greedy steps, two runs per process. Old build is #1990's head (the vec4 build, ELF `46571f11`), new build is #2015's head (ELF `6e6f47f2`). Both include the float32 fix. The ELFs were built before the comment and name commits that sit above them; those commits change no code. +Each case ran three old and three new processes, interleaved in both orders (old,new,new,old,old,new). Every process ran under `taskset` on the four P cores (CPU 4-7). Run 2 is the prefill plus 63 decode steps plus the per-step logits readback. Linux 7.3.0-2-20261110.1-lab-ARCH, started 2026-11-12 16:16Z, GPU timeouts 0. All 18 processes returned 0 with the pinned ids and the pinned logits digest. + +| Case | Old run 2 (median) | New run 2 (median) | Gain: median (pair range) | Logits digest, all six processes | +|---|---:|---:|---|---| +| unicode-24 | 11.8 s | 8.4 s | 28.6% (27.9 to 31.2) | `6bfbb1a55a3e` | +| document-256 | 26.9 s | 12.2 s | 54.4% (53.9 to 61.0) | `8901e17ccef7` | +| document-1024 | 76.9 s | 25.8 s | 66.8% (66.1 to 66.9) | `a472a62c7e22` | + +`pinned-heads-summary.json` holds every process row of this run. It is generated from the raw run output. + +An earlier run on the same chip did not pin the cores (old build `1e1b7d8c`, new build `b6d9a35b`, started 2026-11-11 19:15Z). Its gains were lower and its logits digests differ from the ones above; the two runs do not say which of the changes caused that: + +| Case | Old run 2 (median) | New run 2 (median) | Gain: median (pair range) | Logits digest, all six processes | +|---|---:|---:|---|---| +| unicode-24 | 10.2 s | 9.1 s | 10.9% (6.7 to 22.2) | `e946185d64ae` | +| document-256 | 23.8 s | 15.3 s | 35.6% (32.4 to 37.1) | `3188aaae7d11` | +| document-1024 | 70.2 s | 38.7 s | 44.9% (34.8 to 45.1) | `82c353787e89` | + +The gain depends on the environment. The first version of this receipt measured 19.3%, 49.9% and 62.6% on earlier builds with Linux 7.1.12-2-20261013.18-sep-ARCH and Mesa git-9323171c12 (digests `e5447d9a3ab5`, `147a2d2e0232`, `7f9d8e65e90e`). These runs do not separate the cause. + +First version's table (earlier builds, earlier environment; old build the vec4 build `e742e023`, new build `618ed699`): + +| Case | Old run 2 (median) | New run 2 (median) | Gain: median (pair range) | Logits digest, all six processes | +|---|---:|---:|---|---| +| unicode-24 | 9.2 s | 7.3 s | 19.3% (18.3 to 24.3) | `e5447d9a3ab5` | +| document-256 | 20.5 s | 10.2 s | 49.9% (48.5 to 50.6) | `147a2d2e0232` | +| document-1024 | 59.7 s | 22.3 s | 62.6% (61.4 to 64.0) | `7f9d8e65e90e` | + +An earlier measurement, on the build before the last tile was padded (`a2dcaa1d`; the n = 16 projections kept the vec4 image): + +| Case | Old run 2 (median) | New run 2 (median) | Gain: median (pair range) | Logits digest, all six processes | +|---|---:|---:|---|---| +| unicode-24 | 9.0 s | 6.9 s | 24.7% (19.5 to 27.6) | `e5447d9a3ab5` | +| document-256 | 20.6 s | 10.5 s | 48.8% (48.0 to 50.6) | `147a2d2e0232` | +| document-1024 | 59.6 s | 22.2 s | 62.8% (62.7 to 63.0) | `7f9d8e65e90e` | + +Files: + +- `pinned-heads-summary.json` holds every process row, the pairs, the ELF hashes and the setup of the pinned run in the first table. It is generated from the raw run output. +- `fixed-heads-summary.json` is the same for the unpinned run of 2026-11-11 in the second table. +- `summary.json` is the first version's summary (earlier builds, earlier environment, both tile builds). It stays as the record of those runs. +- `runners-2a5a06c2/` holds the output of the ten contract and stage runners at #2015's head (`2a5a06c2`, 2026-11-12, GPU timeouts 0), one file per runner, with `log.txt` listing the exit codes. `tf-vulkan-kernel-contract.txt` includes `qmm_q4_tile_f32`: four kernels pass gates 1 to 4, eight controls caught. `tf-stage1-gpu.txt` holds the projection cases with `tile=identical` (the tiled image against the default image, byte for byte, on the bf16 projections). No byte-level comparison of the f32 tile image against `qmm_f32` exists; the equal logits digests of the two heads are the identity evidence for it. +- `PREREG.md` was written before the first GPU run. Its addendum, written after the final pairs, says what changed and why. +- `raw-files-sha256.txt` lists the SHA-256 of every raw run file the summary was generated from (`final/` and `interim/`) for `summary.json`. + +Byte identity: each process writes the SHA-256 of its full logits stream. A case passes only if all six digests are equal and the ids match the pinned ids. +The digests of #1990's head and #2015's head are equal on every input. diff --git a/docs/receipts/vulkan-qmm-tile-m1/fixed-heads-summary.json b/docs/receipts/vulkan-qmm-tile-m1/fixed-heads-summary.json new file mode 100644 index 00000000000..604b4eb8fe3 --- /dev/null +++ b/docs/receipts/vulkan-qmm-tile-m1/fixed-heads-summary.json @@ -0,0 +1,262 @@ +{ + "change": "qmm_q4_tile (#2015) measured on the float32-fixed heads", + "setup": { + "chip": "Apple M1 (G13G B1)", + "kernel_release": "7.3.0-2-20261110.1-lab-ARCH", + "mesa": "Honeykrisp, Omarchy Mesa build, library sha256 d75f4cbc", + "checkpoint": "Qwen3.5-2B-MLX-4bit, revision 93760be4", + "settings": "free-running greedy, 64 steps, two runs per process, vendor id 0x10005, one model process at a time, fresh process per run", + "decision_rule": "median run-2 gain of at least 3 percent and every pair positive", + "old_build": "#1990 head (vec4 build), ELF sha256 prefix 1e1b7d8c", + "new_build": "#2015 head, ELF sha256 prefix b6d9a35b", + "note": "Both include the float32 fix." + }, + "cases": [ + { + "case": "unicode-24", + "status": [ + "PASS_WITH_NOT_VERIFIED" + ], + "start_utc": [ + "2026-11-11T19:15:24+00:00" + ], + "kernel_release": "7.3.0-2-20261110.1-lab-ARCH", + "elfs": { + "old": { + "sha256": "1e1b7d8cdddf57645c2090207eef24812e4cbc6e6b90a11dc9d43492ce90ac5a", + "bytes": 6087520 + }, + "new": { + "sha256": "b6d9a35b97786770f9424b6e0089c2daead0625e0e21597f8abb96c6f6cba3cb", + "bytes": 6116824 + } + }, + "order": "old,new,new,old,old,new", + "pairs": 3, + "logits_sha256_all_processes": "e946185d64aea5942265276e7b76f7d4c4560b543145e1139e9c2a09432ca00c", + "run2_ms_median": { + "old": 10194, + "new": 9080 + }, + "run2_gain_median": 0.10927996860898567, + "run2_gain_range": [ + 0.06705035971223021, + 0.22169437846397466 + ], + "processes": [ + { + "label": "pr0-old", + "arm": "old", + "run1_ms": 9942, + "run2_ms": 10104, + "peak_rss_bytes": 1233780736, + "logits_sha256": "e946185d64aea5942265276e7b76f7d4c4560b543145e1139e9c2a09432ca00c" + }, + { + "label": "pr1-new", + "arm": "new", + "run1_ms": 8311, + "run2_ms": 7864, + "peak_rss_bytes": 1503756288, + "logits_sha256": "e946185d64aea5942265276e7b76f7d4c4560b543145e1139e9c2a09432ca00c" + }, + { + "label": "pr2-new", + "arm": "new", + "run1_ms": 7935, + "run2_ms": 9080, + "peak_rss_bytes": 1503297536, + "logits_sha256": "e946185d64aea5942265276e7b76f7d4c4560b543145e1139e9c2a09432ca00c" + }, + { + "label": "pr3-old", + "arm": "old", + "run1_ms": 10025, + "run2_ms": 10194, + "peak_rss_bytes": 1215594496, + "logits_sha256": "e946185d64aea5942265276e7b76f7d4c4560b543145e1139e9c2a09432ca00c" + }, + { + "label": "pr4-old", + "arm": "old", + "run1_ms": 10147, + "run2_ms": 10425, + "peak_rss_bytes": 1246560256, + "logits_sha256": "e946185d64aea5942265276e7b76f7d4c4560b543145e1139e9c2a09432ca00c" + }, + { + "label": "pr5-new", + "arm": "new", + "run1_ms": 8464, + "run2_ms": 9726, + "peak_rss_bytes": 1504264192, + "logits_sha256": "e946185d64aea5942265276e7b76f7d4c4560b543145e1139e9c2a09432ca00c" + } + ] + }, + { + "case": "document-256", + "status": [ + "PASS_WITH_NOT_VERIFIED" + ], + "start_utc": [ + "2026-11-11T19:22:57+00:00" + ], + "kernel_release": "7.3.0-2-20261110.1-lab-ARCH", + "elfs": { + "old": { + "sha256": "1e1b7d8cdddf57645c2090207eef24812e4cbc6e6b90a11dc9d43492ce90ac5a", + "bytes": 6087520 + }, + "new": { + "sha256": "b6d9a35b97786770f9424b6e0089c2daead0625e0e21597f8abb96c6f6cba3cb", + "bytes": 6116824 + } + }, + "order": "old,new,new,old,old,new", + "pairs": 3, + "logits_sha256_all_processes": "3188aaae7d1182c133a89fde103722d23f3d81b57df259498c9b778ebe2d0d7a", + "run2_ms_median": { + "old": 23785, + "new": 15322 + }, + "run2_gain_median": 0.3561914366149838, + "run2_gain_range": [ + 0.3235652722303973, + 0.3711087645195354 + ], + "processes": [ + { + "label": "pr0-old", + "arm": "old", + "run1_ms": 23758, + "run2_ms": 23785, + "peak_rss_bytes": 1215184896, + "logits_sha256": "3188aaae7d1182c133a89fde103722d23f3d81b57df259498c9b778ebe2d0d7a" + }, + { + "label": "pr1-new", + "arm": "new", + "run1_ms": 14613, + "run2_ms": 16089, + "peak_rss_bytes": 1534902272, + "logits_sha256": "3188aaae7d1182c133a89fde103722d23f3d81b57df259498c9b778ebe2d0d7a" + }, + { + "label": "pr2-new", + "arm": "new", + "run1_ms": 15119, + "run2_ms": 14889, + "peak_rss_bytes": 1504591872, + "logits_sha256": "3188aaae7d1182c133a89fde103722d23f3d81b57df259498c9b778ebe2d0d7a" + }, + { + "label": "pr3-old", + "arm": "old", + "run1_ms": 23755, + "run2_ms": 23675, + "peak_rss_bytes": 1215578112, + "logits_sha256": "3188aaae7d1182c133a89fde103722d23f3d81b57df259498c9b778ebe2d0d7a" + }, + { + "label": "pr4-old", + "arm": "old", + "run1_ms": 23943, + "run2_ms": 23799, + "peak_rss_bytes": 1246576640, + "logits_sha256": "3188aaae7d1182c133a89fde103722d23f3d81b57df259498c9b778ebe2d0d7a" + }, + { + "label": "pr5-new", + "arm": "new", + "run1_ms": 14885, + "run2_ms": 15322, + "peak_rss_bytes": 1503232000, + "logits_sha256": "3188aaae7d1182c133a89fde103722d23f3d81b57df259498c9b778ebe2d0d7a" + } + ] + }, + { + "case": "document-1024", + "status": [ + "PASS_WITH_NOT_VERIFIED" + ], + "start_utc": [ + "2026-11-11T20:07:45+00:00" + ], + "kernel_release": "7.3.0-2-20261110.1-lab-ARCH", + "elfs": { + "old": { + "sha256": "1e1b7d8cdddf57645c2090207eef24812e4cbc6e6b90a11dc9d43492ce90ac5a", + "bytes": 6087520 + }, + "new": { + "sha256": "b6d9a35b97786770f9424b6e0089c2daead0625e0e21597f8abb96c6f6cba3cb", + "bytes": 6116824 + } + }, + "order": "old,new,new,old,old,new", + "pairs": 3, + "logits_sha256_all_processes": "82c353787e896b52c8def7f4cebbc087d42f71c4837505db8100df67b2235356", + "run2_ms_median": { + "old": 70186, + "new": 38670 + }, + "run2_gain_median": 0.4493346251389166, + "run2_gain_range": [ + 0.34842693744830144, + 0.4505229055359782 + ], + "processes": [ + { + "label": "pr0-old", + "arm": "old", + "run1_ms": 70216, + "run2_ms": 70186, + "peak_rss_bytes": 1228144640, + "logits_sha256": "82c353787e896b52c8def7f4cebbc087d42f71c4837505db8100df67b2235356" + }, + { + "label": "pr1-new", + "arm": "new", + "run1_ms": 39480, + "run2_ms": 38649, + "peak_rss_bytes": 1549025280, + "logits_sha256": "82c353787e896b52c8def7f4cebbc087d42f71c4837505db8100df67b2235356" + }, + { + "label": "pr2-new", + "arm": "new", + "run1_ms": 39377, + "run2_ms": 38670, + "peak_rss_bytes": 1516339200, + "logits_sha256": "82c353787e896b52c8def7f4cebbc087d42f71c4837505db8100df67b2235356" + }, + { + "label": "pr3-old", + "arm": "old", + "run1_ms": 70308, + "run2_ms": 70376, + "peak_rss_bytes": 1228161024, + "logits_sha256": "82c353787e896b52c8def7f4cebbc087d42f71c4837505db8100df67b2235356" + }, + { + "label": "pr4-old", + "arm": "old", + "run1_ms": 70105, + "run2_ms": 70118, + "peak_rss_bytes": 1228111872, + "logits_sha256": "82c353787e896b52c8def7f4cebbc087d42f71c4837505db8100df67b2235356" + }, + { + "label": "pr5-new", + "arm": "new", + "run1_ms": 39922, + "run2_ms": 45687, + "peak_rss_bytes": 1515782144, + "logits_sha256": "82c353787e896b52c8def7f4cebbc087d42f71c4837505db8100df67b2235356" + } + ] + } + ] +} diff --git a/docs/receipts/vulkan-qmm-tile-m1/pinned-heads-summary.json b/docs/receipts/vulkan-qmm-tile-m1/pinned-heads-summary.json new file mode 100644 index 00000000000..40fa4db8c91 --- /dev/null +++ b/docs/receipts/vulkan-qmm-tile-m1/pinned-heads-summary.json @@ -0,0 +1,210 @@ +{ + "change": "qmm_q4_tile_f32 (#2015) against qmm_q4_vec4 (#1990), P-core pinned", + "setup": { + "chip": "Apple M1 (G13G B1)", + "kernel_release": "7.3.0-2-20261110.1-lab-ARCH", + "settings": "free-running greedy, 64 steps, two runs per process, fresh process per run, one model process at a time, every process under taskset on the four P cores (CPU 4-7)", + "start_utc": "2026-11-12T16:16:03Z", + "old_build": "#1990 head (vec4 build), ELF sha256 prefix 46571f11", + "new_build": "#2015 head (tile build), ELF sha256 prefix 6e6f47f2", + "note": "Both include the float32 fix. The ELFs were built before the comment and name only commits that sit above them; those commits change no code.", + "gpu_timeouts": 0 + }, + "cases": [ + { + "case": "unicode-24", + "order": "h1990-1,h2015-1,h2015-2,h1990-2,h1990-3,h2015-3", + "pairs": 3, + "logits_sha256_all_processes": "6bfbb1a55a3e7d8d95bfdf4fac7b2d5a46df28eac5a83f56bf6e55eb965c26c9", + "run2_ms_median": { + "old": 11803, + "new": 8416 + }, + "run2_gain_median": 0.28569003564759804, + "run2_gain_range": [ + 0.2785484689561834, + 0.3115309667033805 + ], + "processes": [ + { + "label": "old-1", + "rc": 0, + "run1_ms": 12449, + "run2_ms": 11782, + "ids_match": true, + "logits_sha256": "6bfbb1a55a3e7d8d95bfdf4fac7b2d5a46df28eac5a83f56bf6e55eb965c26c9" + }, + { + "label": "new-1", + "rc": 0, + "run1_ms": 9192, + "run2_ms": 8416, + "ids_match": true, + "logits_sha256": "6bfbb1a55a3e7d8d95bfdf4fac7b2d5a46df28eac5a83f56bf6e55eb965c26c9" + }, + { + "label": "old-2", + "rc": 0, + "run1_ms": 12410, + "run2_ms": 11822, + "ids_match": true, + "logits_sha256": "6bfbb1a55a3e7d8d95bfdf4fac7b2d5a46df28eac5a83f56bf6e55eb965c26c9" + }, + { + "label": "new-2", + "rc": 0, + "run1_ms": 9255, + "run2_ms": 8529, + "ids_match": true, + "logits_sha256": "6bfbb1a55a3e7d8d95bfdf4fac7b2d5a46df28eac5a83f56bf6e55eb965c26c9" + }, + { + "label": "old-3", + "rc": 0, + "run1_ms": 12364, + "run2_ms": 11803, + "ids_match": true, + "logits_sha256": "6bfbb1a55a3e7d8d95bfdf4fac7b2d5a46df28eac5a83f56bf6e55eb965c26c9" + }, + { + "label": "new-3", + "rc": 0, + "run1_ms": 9573, + "run2_ms": 8126, + "ids_match": true, + "logits_sha256": "6bfbb1a55a3e7d8d95bfdf4fac7b2d5a46df28eac5a83f56bf6e55eb965c26c9" + } + ] + }, + { + "case": "document-256", + "order": "h1990-1,h2015-1,h2015-2,h1990-2,h1990-3,h2015-3", + "pairs": 3, + "logits_sha256_all_processes": "8901e17ccef7bffde7332519f5c70cb2389f47aaff31ddc7c91b4fe9099b709b", + "run2_ms_median": { + "old": 26853, + "new": 12247 + }, + "run2_gain_median": 0.5442976744186047, + "run2_gain_range": [ + 0.5386565272496832, + 0.6095408334264328 + ], + "processes": [ + { + "label": "old-1", + "rc": 0, + "run1_ms": 27453, + "run2_ms": 26875, + "ids_match": true, + "logits_sha256": "8901e17ccef7bffde7332519f5c70cb2389f47aaff31ddc7c91b4fe9099b709b" + }, + { + "label": "new-1", + "rc": 0, + "run1_ms": 13154, + "run2_ms": 12247, + "ids_match": true, + "logits_sha256": "8901e17ccef7bffde7332519f5c70cb2389f47aaff31ddc7c91b4fe9099b709b" + }, + { + "label": "old-2", + "rc": 0, + "run1_ms": 27295, + "run2_ms": 26853, + "ids_match": true, + "logits_sha256": "8901e17ccef7bffde7332519f5c70cb2389f47aaff31ddc7c91b4fe9099b709b" + }, + { + "label": "new-2", + "rc": 0, + "run1_ms": 12509, + "run2_ms": 10485, + "ids_match": true, + "logits_sha256": "8901e17ccef7bffde7332519f5c70cb2389f47aaff31ddc7c91b4fe9099b709b" + }, + { + "label": "old-3", + "rc": 0, + "run1_ms": 27483, + "run2_ms": 26826, + "ids_match": true, + "logits_sha256": "8901e17ccef7bffde7332519f5c70cb2389f47aaff31ddc7c91b4fe9099b709b" + }, + { + "label": "new-3", + "rc": 0, + "run1_ms": 13112, + "run2_ms": 12376, + "ids_match": true, + "logits_sha256": "8901e17ccef7bffde7332519f5c70cb2389f47aaff31ddc7c91b4fe9099b709b" + } + ] + }, + { + "case": "document-1024", + "order": "h1990-1,h2015-1,h2015-2,h1990-2,h1990-3,h2015-3", + "pairs": 3, + "logits_sha256_all_processes": "a472a62c7e221de2b9ec38e07b72011f29abc6f23adc393cf0f7b7cae2259c37", + "run2_ms_median": { + "old": 76868, + "new": 25765 + }, + "run2_gain_median": 0.6675827011405274, + "run2_gain_range": [ + 0.6607732634338139, + 0.66934224905032 + ], + "processes": [ + { + "label": "old-1", + "rc": 0, + "run1_ms": 78056, + "run2_ms": 77508, + "ids_match": true, + "logits_sha256": "a472a62c7e221de2b9ec38e07b72011f29abc6f23adc393cf0f7b7cae2259c37" + }, + { + "label": "new-1", + "rc": 0, + "run1_ms": 27080, + "run2_ms": 25765, + "ids_match": true, + "logits_sha256": "a472a62c7e221de2b9ec38e07b72011f29abc6f23adc393cf0f7b7cae2259c37" + }, + { + "label": "old-2", + "rc": 0, + "run1_ms": 77068, + "run2_ms": 76868, + "ids_match": true, + "logits_sha256": "a472a62c7e221de2b9ec38e07b72011f29abc6f23adc393cf0f7b7cae2259c37" + }, + { + "label": "new-2", + "rc": 0, + "run1_ms": 26849, + "run2_ms": 25417, + "ids_match": true, + "logits_sha256": "a472a62c7e221de2b9ec38e07b72011f29abc6f23adc393cf0f7b7cae2259c37" + }, + { + "label": "old-3", + "rc": 0, + "run1_ms": 77571, + "run2_ms": 76300, + "ids_match": true, + "logits_sha256": "a472a62c7e221de2b9ec38e07b72011f29abc6f23adc393cf0f7b7cae2259c37" + }, + { + "label": "new-3", + "rc": 0, + "run1_ms": 26738, + "run2_ms": 25883, + "ids_match": true, + "logits_sha256": "a472a62c7e221de2b9ec38e07b72011f29abc6f23adc393cf0f7b7cae2259c37" + } + ] + } + ] +} \ No newline at end of file diff --git a/docs/receipts/vulkan-qmm-tile-m1/raw-files-sha256.txt b/docs/receipts/vulkan-qmm-tile-m1/raw-files-sha256.txt new file mode 100644 index 00000000000..3c68ea37316 --- /dev/null +++ b/docs/receipts/vulkan-qmm-tile-m1/raw-files-sha256.txt @@ -0,0 +1,135 @@ +0112e8e1d578a77a796fc4ce49e2ee161a4f5b7605631d38e0e9e56ee419e24f final/pair-document256/result.json +041cd74df7b9c2ea7e25de69066152110bb0f9cd2b76f2f7869f81f8b9876a99 final/pair-unicode24/timing.tsv +06c56fb46cb8790443160a8d945f708d8ef01ee3143c466213e4cb100edae262 interim/pair-unicode24/out/unicode-24--p0-free/result.json +07629cfec834975af4acb84c4b50f0c28ab3cb6433cd4729e2f279a4a3a4e386 final/pair-document1024/result.json +0a8312e907ecf4185ba34aa0a6899e8cc0830221e5713a5d129cb9b6c1b5ca05 final/pair-unicode24/log--pr2-new.err +0c0c1525c29136df0a96bb15925d7fbc36023e17e64e25b800eefe22bd5c6f6e interim/pair-document256/log--pr1-new.err +10836d84335714907c116ac56c7e70809fa747bce08a1159048b2ce17bc811d8 final/pair-document1024/out/document-1024--p0-free/result.json +10ce11a6fc836ea364d8065458e21647ae164a56bf945f61ded011240e468dad interim/pair-document256/log--pr3-old.err +186f6c514d6f61d53107819aa6380c6acf529cdd0a1a62d009751348ab4fe9b5 interim/pair-document1024/out/document-1024--p5-free/result.json +1964ce86d8f2b0bf3bcb556cc595a142fb7f2f483f2e3f097f3a144a71bde11d final/pair-document1024/log--pr4-old.err +1b60af3d291c0bc9a430527d26b371b1b0b6dda5a1c89f22572f14710b11b1b8 interim/pair-unicode24/log--pr5-new.err +1bb8521e60ac2471f179b2cf67566315625591e571effb8a5798c4fd746b89cd final/pair-document1024/log--pr5-new.err +210256e95b35cd0fd63c25b35d8c3ffd89897a7fd03eedef0be4390549d07786 final/pair-document1024/log--pr3-old.err +27cac8197a46b96d42900e1574cafc2b73756d96edeee57432a183d38eb4afce interim/pair-document256/out/document-256--p5-free/result.json +2aed068bcfe9dc8cc4439bfb88058b75bafd231b9162b35ee0265d5b8b72d53c final/contract-m1/run.err +2f722cee84bf32d87db55a09b7f9dcff661685205627871bf79ec979f448b6b5 interim/pair-document1024/log--pr0-old.err +3128c9fd9feb31fcae42873b105a88b6efaed403dcde6e8bb236c744d3db52ea interim/pair-unicode24/log--pr2-new.err +36f2f22b846d4b11a99122a0c4f8e97a24ac391324a0962b660c49c660f4b62b final/pair-unicode24/result.json +37c1b2a7abb9467191cb9164a6e484744fc0484aa355b537d3951938c59c0875 final/pair-document256/log--pr0-old.err +3dc54b731aa58ef74b824d0eabab283adc7457cd0bf1435f0b07c543dc3d06df final/pair-document1024/out/document-1024--p5-free/result.json +401a8ef94f32f9cda722192bcdb42a85a2ca50b0d6cfe11e3b8ac1adbbfe303d final/pair-unicode24/log--pr0-old.err +41a189c56406b51ec0f549db554f7136b95dc10df98d5e34ce0461012ffacba3 interim/pair-document1024/log--pr1-new.err +4741adb372eb9d0ebbff269905acbce8bed57538cebba6a84057e425c793a8da final/pair-document256/out/document-256--p1-free/result.json +47d653d098ec7a7658c172ab3435f149207398466fe09322d8a9b0fd1876d7aa final/pair-unicode24/out/unicode-24--p3-free/result.json +4c98f154b6c1642191871148f7063558242da200f454f262deb7bf96bbc4197c final/pair-unicode24/log--pr5-new.err +5011041e4bf80318251f07cfa8d25a00a24c846f4e48467b634dd596efb88518 interim/pair-document1024/result.json +507fb2afbb59d245f76c95a0d6f68993fbc79fa2cb1eb14774ca4a0addfe2fef interim/pair-document1024/log--pr5-new.err +537445c7d17c972ab266c723bd448c887c1a3a72b6236da6ffd3d49333b7cc46 final/pair-unicode24/log--pr3-old.err +5540fcbcd5bd6a0e30763282398e8f379b32df0988b936baaebec1c9e1f62cb3 interim/pair-document256/out/document-256--p2-free/result.json +582239964e403a7f707e8e5e574be532cb7ab60067a08010a9395137285f9656 final/pair-document256/out/document-256--p2-free/result.json +5c5dcf7f3cb37199375541475a98b7ff95744b7bde95f8224ae4684825c7af65 final/pair-document1024/out/document-1024--p1-free/result.json +600ee4652d8dbde91e4565e77655e089f34b8d7350ef092962a9e5c4f1c6fb1e interim/pair-document256/out/document-256--p0-free/result.json +66285671a068368f7905c43e7c55a686f7fa948056037e37a00fa098b02f6872 final/pair-document1024/log--pr0-old.err +67a7b4d75f65b4422b61772a90d1d9d807f226fc49bafb1072db234e321b4792 interim/pair-document256/result.json +6a8bde588dc7027e3f67f274fe336e505058f979072bf337bed04875e9c5d920 final/pair-unicode24/log--pr4-old.err +6c4a57612696c8058f56cc5666df6da42be2399b28e67e01abfa97f2a62c881e interim/id-m1/env.txt +6e017bc58596a8245abfb6b7cc5bc48967ebe8a19a7159a05c9a792578720946 interim/pair-document1024/out/document-1024--p3-free/result.json +6fabe7a032b4bc75ba96bea8741bdcaee624753096e82b402152be3d1a0610ec final/pair-document1024/log--pr2-new.err +70be688a55d59c4ed80567a1fb88d09be7ca897785a9f32d5be8549e28ed41f6 final/pair-document256/log--pr3-old.err +724a9d0c7b69eaba365545112c1e44928bfefef92401717281b8942abcfab4d5 final/pair-document256/out/document-256--p5-free/result.json +75154068bbe0bf2888fa42d214d7134fd21deb6f6cfacc7b18dade0d633df5d4 interim/pair-document1024/out/document-1024--p1-free/result.json +7dd036a67cc1a1a4bf094f10b96495305a305af78182415679c6677f512e9ee5 interim/pair-document256/out/document-256--p4-free/result.json +7fac308c4fd59e705616dda63c2c963fabb562ecddf119e95a6c5155e854d361 interim/pair-document1024/timing.tsv +7fee6fe9ed08c45aeff03970fc0461675ecc86ee68f14f4bc8133b153a3c5011 interim/pair-document1024/log--pr2-new.err +85b2c8f6c83314ead09cd7e7bab23b23f7f7aab33579a7e38e015ca4ff77768c final/contract-m1/env.txt +879b6179fe9897b628b17ee5bae33e28d64e04dcaf334d980879dc38435d0edd interim/pair-document256/out/document-256--p3-free/result.json +87a3dce51955649c33d2798d5a1b7aa8eea64a7ed8ba15ae99aa5acc35c2e952 final/id-m1/SHA256SUMS +92e64de27e2025582305a10a512924cb858f0363256f92a3f3b35b668e850bfc interim/pair-unicode24/out/unicode-24--p3-free/result.json +941c7bf552ff35788e55421aab709c4e5d0706e54f0997368fa47850fd6857cc interim/pair-document256/out/document-256--p1-free/result.json +95a5543cb1ec4a75f90bcc57099522e273f6cbcd53aad901cf20bd114edbefbe final/pair-document256/log--pr1-new.err +960bf5e65d527541129a745fe452e0dab5e6204de627c05c1b265f79cf496b7f interim/pair-document256/timing.tsv +9a271f2a916b0b6ee6cecb2426f0b3206ef074578be55d9bc94f6f3fe3ab86aa final/contract-m1/run.rc +9a271f2a916b0b6ee6cecb2426f0b3206ef074578be55d9bc94f6f3fe3ab86aa final/id-m1/run.rc +9a271f2a916b0b6ee6cecb2426f0b3206ef074578be55d9bc94f6f3fe3ab86aa interim/id-m1/run.rc +9f48284b2e0a4bd126e41f18aed9d468c6dd848265a7c4f95aecda1fe975aa77 final/pair-document256/out/document-256--p3-free/result.json +9f48284b2e0a4bd126e41f18aed9d468c6dd848265a7c4f95aecda1fe975aa77 final/pair-document256/out/document-256--p4-free/result.json +9fb7614fbe7317a29e92b9eaa5cf6225ad1d8a67b6d79058ab9ac10a4a5f48bd interim/pair-document1024/out/document-1024--p2-free/result.json +ac012b673b78ae365287eb7fa34b5c1d7abeeb79af89cb41922366a26d6e819d final/pair-unicode24/out/unicode-24--p4-free/result.json +ad2a571d2334a3e29e7f01e90e8cd2bc0d59391f78aab2153de89c29d56dbe7c final/pair-unicode24/log--pr1-new.err +af2c4d4a3d13506bd482bfef62d825b4326e2c3cae19d6608cfcad45bb1c2644 interim/pair-document256/log--pr5-new.err +b0ecd463abe4919b5b08696bd4316565c58822105f8b7ec120c3796822ff2564 interim/pair-document1024/out/document-1024--p4-free/result.json +b2c1460bd9ad655cd5ff41dd5ee1d04045a30aff7801d97f3226acb9e6302657 interim/pair-document256/log--pr2-new.err +b48c8445e9583fb975063248ca22562cbf1e9eace2d72138807c2290436f1363 interim/id-m1/SHA256SUMS +b8afdcea51547849cf165ac610345f934eb345f159e617a50c90e5040d7a809e interim/id-m1/run.err +bd9739318bcad863ef2b750b794f8ade8e0a958a08148d7ef76d178cd2b2ccef interim/pair-unicode24/out/unicode-24--p4-free/result.json +c0dcf94ad6937d0a395cbdb94c625973f870a0aecb38eda98aeebaa55934552f final/contract-m1/SHA256SUMS +c275dc78591ba0518e6fc4292063ab7ec825fd62e3632a826567d9120eac94e4 interim/pair-document256/log--pr0-old.err +c4e71cca1d84e1317dc2c59c974f92567b07bab9535f4f7b79b82d6c71f10210 final/pair-document1024/out/document-1024--p2-free/result.json +c8ac1562a39a08ef330f32b928fce5282eae165081f42331c377a46f43e482c1 interim/pair-document1024/log--pr3-old.err +cbcc039001a45904bdfb53c1494ebf1c3c60b281d6e6dc652c6cc501411dbcee interim/pair-document1024/log--pr4-old.err +cd674021ce2c7555055327284610e76517fa61d64f825d7f24f5501e36ed0f25 interim/pair-document1024/out/document-1024--p0-free/result.json +d2237bb5f8bb55a844ccb51977152a00778961533dd3444dafc568ef7be99f7b interim/pair-unicode24/out/unicode-24--p2-free/result.json +d417991d9d5b6b470b4674fb2e556de3610e5d25303e28d374f50bc049f097c9 final/pair-document1024/out/document-1024--p4-free/result.json +d584318cd52d60af1e0a1006868bb31bb6ff07c8fb712f96a2f9af82e5c402e1 final/pair-unicode24/out/unicode-24--p0-free/result.json +d6c68e98d5bbf9848c482a51bc8112d32c06c8370faf2b427b8b1b4d24d48dcb interim/pair-unicode24/log--pr0-old.err +d7f4e1e7a90e25f10c065879392ff7eb7c328da4eb90b744593bf160e57dd210 final/pair-document256/timing.tsv +d91a08f63ec2bfbef6d077f5189a3b1f81d8790c7eb2929c15a19e1f28627489 interim/pair-document256/log--pr4-old.err +d941b38021f0746705caf914644759fab8c30779b3f70caacf54e7eae7d711ba final/id-m1/env.txt +dd73b55699b04abb00c3a8585d9afaa9efc13fdfb85847c5a23eead677b54669 interim/pair-unicode24/log--pr1-new.err +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 final/contract-m1/run.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 final/id-m1/run.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 final/pair-document1024/log--pr0-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 final/pair-document1024/log--pr1-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 final/pair-document1024/log--pr2-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 final/pair-document1024/log--pr3-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 final/pair-document1024/log--pr4-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 final/pair-document1024/log--pr5-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 final/pair-document256/log--pr0-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 final/pair-document256/log--pr1-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 final/pair-document256/log--pr2-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 final/pair-document256/log--pr3-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 final/pair-document256/log--pr4-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 final/pair-document256/log--pr5-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 final/pair-unicode24/log--pr0-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 final/pair-unicode24/log--pr1-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 final/pair-unicode24/log--pr2-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 final/pair-unicode24/log--pr3-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 final/pair-unicode24/log--pr4-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 final/pair-unicode24/log--pr5-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 interim/id-m1/run.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 interim/pair-document1024/log--pr0-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 interim/pair-document1024/log--pr1-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 interim/pair-document1024/log--pr2-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 interim/pair-document1024/log--pr3-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 interim/pair-document1024/log--pr4-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 interim/pair-document1024/log--pr5-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 interim/pair-document256/log--pr0-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 interim/pair-document256/log--pr1-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 interim/pair-document256/log--pr2-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 interim/pair-document256/log--pr3-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 interim/pair-document256/log--pr4-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 interim/pair-document256/log--pr5-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 interim/pair-unicode24/log--pr0-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 interim/pair-unicode24/log--pr1-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 interim/pair-unicode24/log--pr2-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 interim/pair-unicode24/log--pr3-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 interim/pair-unicode24/log--pr4-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 interim/pair-unicode24/log--pr5-new.out +e54945448af0546fbe2ffcf08b367f5502fc9d1e82ed8e8e44d5f8f127ab24b2 final/pair-unicode24/out/unicode-24--p1-free/result.json +e7368fd827063ebef4d57f7f380b50d7c2bb5005cffd84955761be0d3204a3d3 final/pair-document256/log--pr5-new.err +eaefc453eed30e7d28414cf766bc36855852e458295b634ed6df8a9db1930da3 interim/pair-unicode24/result.json +ec980ea786ccbf531ae6666c2997745695c2a8dc9c57643b0a319b08f52ccec2 final/pair-document1024/out/document-1024--p3-free/result.json +edee7152422fe1a5921c2fcd36a43ae5fc70115f0a61ea30543ea881e384791c final/pair-unicode24/out/unicode-24--p5-free/result.json +ee4331b413d8b8981f0ce84285cf43adf8a9e856ea56514bf8de28bcaa6eb021 final/pair-document256/log--pr2-new.err +efabfef0ee8f8243a92caf34bd3b1a491f2aaa19dad9217944d578a86c603837 final/pair-document1024/timing.tsv +f0f995c62b3e054db35bb43283e18a3e612a8b7feb21db9703a4e6429a2095b5 interim/pair-unicode24/log--pr3-old.err +f1292aa3599151bbc49559c953dbacc6b5aa8f2efec3c06414d3686ba18155d3 interim/pair-unicode24/out/unicode-24--p1-free/result.json +f2bb0317c7b31e33b3f46d53913d6981dc68c9f2e6247809470e0bfa1e19b312 final/pair-document256/out/document-256--p0-free/result.json +f57d1d76acf4c8aef3a9b02d655b15482da72fb7441149ea61584e4d0e674a8d final/pair-document256/log--pr4-old.err +f596dbcc023f9cb50562e14c82e0a36c4520f7d9bd494ddb24763bef01dc469d interim/pair-unicode24/out/unicode-24--p5-free/result.json +f990e1b4053e3b831f1176b434d9d69bc1ba485ffd882fbda597363c8065260c final/pair-document1024/log--pr1-new.err +f9b852ef11bfc0a55cb7087e9f09f08b7918790158e371a932ea893533b65fe0 interim/pair-unicode24/timing.tsv +fc8da160de81188672396ba74dc1aaae5798d7859df37fe37fbaf112c6a02fe2 interim/pair-unicode24/log--pr4-old.err +fe4a2722b8b9007c487a34d7745291b7bb09f62c71c44c781505c6c65518023f final/pair-unicode24/out/unicode-24--p2-free/result.json +ff86fd5bca1d4af44191fa64b94ef4148be34afa8d9dbed16553b55a573b2476 final/id-m1/run.err diff --git a/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/log.txt b/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/log.txt new file mode 100644 index 00000000000..dafd7f2c0a1 --- /dev/null +++ b/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/log.txt @@ -0,0 +1,22 @@ +tenrun 2026-11-12T16:33:33Z kernel 7.3.0-2-20261110.1-lab-ARCH load 1.27 gpu-timeouts=0 +tenrun tf-attention-contract rc 0 +RESULT attention-contract elapsed_ms=6692 on "Apple M1 (G13G B1)": passed=15 total=15 +tenrun tf-contract-pointwise rc 0 +RESULT contract-pointwise elapsed_ms=72 on "Apple M1 (G13G B1)": passed=16 known=1 total=17 +tenrun tf-stage1-datamove rc 0 +RESULT datamove on "Apple M1 (G13G B1)": kernels=7 passed=7 failed=0 controls=0 known=0 +tenrun tf-stage1-gpu rc 0 +RESULT stage1 elapsed_ms=7304 on "Apple M1 (G13G B1)": passed=21 total=21 +tenrun tf-stage2-gpu rc 0 +RESULT stage2 elapsed_ms=36 on "Apple M1 (G13G B1)": passed=9 total=9 +tenrun tf-stage3-gpu rc 0 +RESULT stage3 elapsed_ms=173 on "Apple M1 (G13G B1)": passed=4 total=4 +tenrun tf-stage4-gpu rc 0 +RESULT stage4 elapsed_ms=195 on "Apple M1 (G13G B1)": passed=6 total=6 +tenrun tf-stage5-gpu rc 0 +RESULT stage5 elapsed_ms=109 on "Apple M1 (G13G B1)": passed=7 total=7 +tenrun tf-vulkan-greedy-contract rc 0 +SUMMARY vulkan greedy contract on "Apple M1 (G13G B1)": 1 kernel, 1 pass gates 1-4, 0 fail; 1 control: 1 red as planted, 0 missed +tenrun tf-vulkan-kernel-contract rc 0 +SUMMARY Apple M1 (G13G B1) 4 kernels: 4 pass gates 1-4, 0 fail, 0 known findings; 8 controls: 8 caught, 0 missed; 0 limit failures; gates 41802 ms +tenrun-complete gpu-timeouts=0 diff --git a/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-attention-contract.txt b/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-attention-contract.txt new file mode 100644 index 00000000000..f357faee5e9 --- /dev/null +++ b/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-attention-contract.txt @@ -0,0 +1,17 @@ +MILESTONE device-found elapsed_ms=4 name="Apple M1 (G13G B1)" vendor=0x10005 device=0x0000 api=1.4.362 storage16=true int16=true shared=32768 fence-timeout=10000ms +MILESTONE rope elapsed_ms=161 case="queries rows4 base0" rows=4 base=0 worst=4.964008589780876e-1 of bound, repeats byte-identical, passthrough tail bit-exact, input unchanged, planted outputs refused +MILESTONE rope elapsed_ms=163 case="queries rows2 base5" rows=2 base=5 worst=4.8822881879460084e-1 of bound, repeats byte-identical, passthrough tail bit-exact, input unchanged, planted outputs refused +MILESTONE sdpa elapsed_ms=165 case="decode n=1" rows=1 prefix=0 visible=[1..1] worst=0e0 of bound, repeats byte-identical, planted outputs refused, inputs unchanged +MILESTONE sdpa elapsed_ms=168 case="decode n=16 tile" rows=1 prefix=15 visible=[16..16] worst=3.5347427802135684e-1 of bound, repeats byte-identical, planted outputs refused, inputs unchanged +MILESTONE sdpa elapsed_ms=171 case="decode n=17 tile+1" rows=1 prefix=16 visible=[17..17] worst=3.9242733174609323e-1 of bound, repeats byte-identical, planted outputs refused, inputs unchanged +MILESTONE sdpa elapsed_ms=187 case="decode n=255 scan" rows=1 prefix=254 visible=[255..255] worst=3.2373066331518274e-1 of bound, repeats byte-identical, planted outputs refused, inputs unchanged +MILESTONE sdpa elapsed_ms=201 case="decode n=256 scan" rows=1 prefix=255 visible=[256..256] worst=3.5630466330042704e-1 of bound, repeats byte-identical, planted outputs refused, inputs unchanged +MILESTONE sdpa elapsed_ms=216 case="decode n=257 scan+1" rows=1 prefix=256 visible=[257..257] worst=3.44416466623321e-1 of bound, repeats byte-identical, planted outputs refused, inputs unchanged +MILESTONE sdpa elapsed_ms=268 case="decode n=1024 stride" rows=1 prefix=1023 visible=[1024..1024] worst=3.371824275163813e-1 of bound, repeats byte-identical, planted outputs refused, inputs unchanged +MILESTONE sdpa elapsed_ms=320 case="decode n=1025 stride+1" rows=1 prefix=1024 visible=[1025..1025] worst=3.1211365844085626e-1 of bound, repeats byte-identical, planted outputs refused, inputs unchanged +MILESTONE sdpa elapsed_ms=669 case="decode n=7167 stream-1" rows=1 prefix=7166 visible=[7167..7167] worst=3.099820299106263e-1 of bound, repeats byte-identical, planted outputs refused, inputs unchanged +MILESTONE sdpa elapsed_ms=1014 case="decode n=7168 stream" rows=1 prefix=7167 visible=[7168..7168] worst=3.0423502310343786e-1 of bound, repeats byte-identical, planted outputs refused, inputs unchanged +MILESTONE sdpa elapsed_ms=1047 case="prefill rows32 n1..32" rows=32 prefix=0 visible=[1..32] worst=3.8573441919364143e-1 of bound, repeats byte-identical, rows decode alone bit-identical, planted outputs refused, inputs unchanged +MILESTONE sdpa elapsed_ms=1459 case="prefill rows17 across 1024" rows=17 prefix=1007 visible=[1008..1024] worst=3.454415083464824e-1 of bound, repeats byte-identical, rows decode alone bit-identical, planted outputs refused, inputs unchanged +MILESTONE sdpa elapsed_ms=6692 case="prefill rows32 to 7168" rows=32 prefix=7136 visible=[7137..7168] worst=3.1143891982294897e-1 of bound, repeats byte-identical, rows decode alone bit-identical, planted outputs refused, inputs unchanged +RESULT attention-contract elapsed_ms=6692 on "Apple M1 (G13G B1)": passed=15 total=15 diff --git a/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-contract-pointwise.txt b/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-contract-pointwise.txt new file mode 100644 index 00000000000..10e05a79826 --- /dev/null +++ b/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-contract-pointwise.txt @@ -0,0 +1,19 @@ +MILESTONE device-found elapsed_ms=8 name="Apple M1 (G13G B1)" vendor=0x10005 device=0x0000 type=integrated api=1.4.362 storage16=true int16=true preserve32=false flush32=true fence-timeout=10000ms +MILESTONE residual-add elapsed_ms=42 case="random 512" elements=512 matches the f32-add + RNE-store pipeline bit for bit, subnormal-class-flush-allowed=false, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE residual-add elapsed_ms=43 case="cancel 512" elements=512 matches the f32-add + RNE-store pipeline bit for bit, subnormal-class-flush-allowed=false, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE residual-add elapsed_ms=44 case="classes 20" elements=20 matches the f32-add + RNE-store pipeline bit for bit, subnormal-class-flush-allowed=true, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE swiglu elapsed_ms=46 case="random 512" elements=512 worst=0.9668 of bound at element 471, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE swiglu elapsed_ms=47 case="controls 16" elements=16 worst=0.5616 of bound at element 9, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE rms-norm elapsed_ms=49 case="1x2048" rows=1 width=2048 worst=0.9742 of bound at row 0 col 669, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE rms-norm elapsed_ms=50 case="4x2048 offsets zero row" rows=4 width=2048 worst=0.9861 of bound at row 2 col 478, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused, zero row exactly zero +MILESTONE rms-norm elapsed_ms=52 case="2x2048 dynamic range" rows=2 width=2048 worst=0.9722 of bound at row 1 col 1079, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE rms-norm elapsed_ms=54 case="2x256 head_norm offsets" rows=2 width=256 worst=0.9734 of bound at row 0 col 213, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE preserve-add elapsed_ms=55 case="random 512" elements=512 matches the stage 2 add pipeline bits over the DenormPreserve image, strict on subnormal sums, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE preserve-add elapsed_ms=56 case="cancel 512" elements=512 matches the stage 2 add pipeline bits over the DenormPreserve image, strict on subnormal sums, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +KNOWN preserve-add case="classes 20": 1 subnormal sums read as zero, honest by design — the device reports preserve32=false (shaderDenormPreserveFloat32 = 0) and flushes fp32 denormals, so the DenormPreserve image cannot keep them: recorded, not fitted +MILESTONE output-gate elapsed_ms=59 case="random 2x2048" rows=2 worst=0.9081 of bound at row 1 col 983, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE output-gate elapsed_ms=61 case="controls 2x2048" rows=2 worst=0.3234 of bound at row 0 col 8, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE gated-norm elapsed_ms=63 case="1x16x128" tokens=1 head-rows=16 width=128 worst=0.0750 of bound at row 0 col 1016, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE gated-norm elapsed_ms=70 case="32x16x128 dynamic zero token" tokens=32 head-rows=512 width=128 worst=0.0763 of bound at row 2 col 191, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused, zero token exactly zero +MILESTONE gated-norm elapsed_ms=72 case="controls 3x16x128" tokens=3 head-rows=48 width=128 worst=0.0760 of bound at row 2 col 1342, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +RESULT contract-pointwise elapsed_ms=72 on "Apple M1 (G13G B1)": passed=16 known=1 total=17 diff --git a/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-stage1-datamove.txt b/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-stage1-datamove.txt new file mode 100644 index 00000000000..056bf123b3b --- /dev/null +++ b/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-stage1-datamove.txt @@ -0,0 +1,30 @@ +MILESTONE device-found name="Apple M1 (G13G B1)" vendor=0x10005 storage16=true int16=true timeout=10000ms +PASS nonfinite vulkan cast_f32_bf16 (RNE): the stated rule holds: outputs equal the contract's RNE bytes: ties to even, bf16 subnormals included, a NaN keeps its class +PASS repeat vulkan cast_f32_bf16 (RNE): five runs byte-identical (292 bytes), inputs and guards untouched, no overlap +PASS shapes vulkan cast_f32_bf16 (RNE): 4 shapes exact by its own check (worst 0.000 of the bound), 3 edges refused by the kernel +N/A sums vulkan cast_f32_bf16 (RNE): declares none sums +PASS nonfinite vulkan cast_bf16_f32 (exact widening): the stated rule holds: outputs equal the reference bits: bf16 widening is exact, every value and NaN byte copies +PASS repeat vulkan cast_bf16_f32 (exact widening): five runs byte-identical (584 bytes), inputs and guards untouched, no overlap +PASS shapes vulkan cast_bf16_f32 (exact widening): 4 shapes exact by its own check (worst 0.000 of the bound), 3 edges refused by the kernel +N/A sums vulkan cast_bf16_f32 (exact widening): declares none sums +PASS nonfinite vulkan copy_u32 (bit copy): the stated rule holds: outputs equal the source bytes: a raw bit copy, every value and NaN pattern included +PASS repeat vulkan copy_u32 (bit copy): five runs byte-identical (1200 bytes), inputs and guards untouched, no overlap +PASS shapes vulkan copy_u32 (bit copy): 7 shapes exact by its own check (worst 0.000 of the bound), 0 edges refused by the kernel +N/A sums vulkan copy_u32 (bit copy): declares none sums +PASS nonfinite vulkan copy_u16 (bit copy): the stated rule holds: outputs equal the source bytes: a raw bit copy, every value and NaN pattern included +PASS repeat vulkan copy_u16 (bit copy): five runs byte-identical (600 bytes), inputs and guards untouched, no overlap +PASS shapes vulkan copy_u16 (bit copy): 4 shapes exact by its own check (worst 0.000 of the bound), 3 edges refused by the kernel +N/A sums vulkan copy_u16 (bit copy): declares none sums +PASS nonfinite vulkan gather_u32 (bit rows): the stated rule holds: outputs equal the addressed row's bytes; an out-of-range id's row is zeros, never another row (ops_q4.zig) +PASS repeat vulkan gather_u32 (bit rows): five runs byte-identical (160 bytes), inputs and guards untouched, no overlap +PASS shapes vulkan gather_u32 (bit rows): 7 shapes exact by its own check (worst 0.000 of the bound), 0 edges refused by the kernel +N/A sums vulkan gather_u32 (bit rows): declares none sums +PASS nonfinite vulkan gather_u16 (bit rows): the stated rule holds: outputs equal the addressed row's bytes; an out-of-range id's row is zeros, never another row (ops_q4.zig) +PASS repeat vulkan gather_u16 (bit rows): five runs byte-identical (80 bytes), inputs and guards untouched, no overlap +PASS shapes vulkan gather_u16 (bit rows): 7 shapes exact by its own check (worst 0.000 of the bound), 0 edges refused by the kernel +N/A sums vulkan gather_u16 (bit rows): declares none sums +PASS nonfinite vulkan dequant_bf16 (q4 g64 RNE): the stated rule holds: one fp32 expression, then the contract's RNE: a non-finite scale or bias makes the output the fp32 expression's class +PASS repeat vulkan dequant_bf16 (q4 g64 RNE): five runs byte-identical (1536 bytes), inputs and guards untouched, no overlap +PASS shapes vulkan dequant_bf16 (q4 g64 RNE): 7 shapes exact by its own check (worst 0.000 of the bound), 0 edges refused by the kernel +N/A sums vulkan dequant_bf16 (q4 g64 RNE): declares none sums +RESULT datamove on "Apple M1 (G13G B1)": kernels=7 passed=7 failed=0 controls=0 known=0 diff --git a/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-stage1-gpu.txt b/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-stage1-gpu.txt new file mode 100644 index 00000000000..1f882137a1d --- /dev/null +++ b/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-stage1-gpu.txt @@ -0,0 +1,23 @@ +MILESTONE device-found elapsed_ms=9 name="Apple M1 (G13G B1)" vendor=0x10005 device=0x0000 type=integrated api=1.4.362 driver="Honeykrisp" info="Mesa 26.3.0-devel (git-439dc032d7)" subgroup=32 storage16=true int16=true max-storage-range=2147483647 max-buffer-size=137438953472 fence-timeout=10000ms +MILESTONE cast-f32-bf16 elapsed_ms=75 elements=146 vectors=18 randoms=128 exact bits, tails unchanged +MILESTONE cast-bf16-f32 elapsed_ms=75 elements=146 vectors=18 randoms=128 exact bits, tails unchanged +MILESTONE copy-u32 elapsed_ms=76 elements=9 src_off=64 dst_off=128 exact bytes, guards and source unchanged +MILESTONE copy-u16 elapsed_ms=76 elements=22 src_off=128 dst_off=256 exact bytes, guards and source unchanged +MILESTONE gather-u32 elapsed_ms=77 rows=6 cols=10 tokens=4 rows exact incl repeat, out-of-range id gathers zeros, tails unchanged +MILESTONE gather-bf16 elapsed_ms=78 rows=6 cols=10 tokens=4 rows exact incl repeat, out-of-range id gathers zeros, tails unchanged +MILESTONE check-ids elapsed_ms=78 valid ids pass, boundary id refused before upload +MILESTONE embedding elapsed_ms=79 case="6 tok 8 rows k64" tokens=6 rows=8 k=64 within bf16 bound, ids include a repeat, table and tails unchanged +MILESTONE embedding elapsed_ms=81 case="3 tok 4 rows k128 repeats" tokens=3 rows=4 k=128 within bf16 bound, ids include a repeat, table and tails unchanged +MILESTONE project elapsed_ms=84 case="m1 n8 k64" worst=0.7485 of bound at row 0 col 1, max_abs=2.522643655538559e-3, tile=identical, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE project elapsed_ms=89 case="m2 n257 k128" worst=0.8880 of bound at row 0 col 131, max_abs=7.288329303264618e-3, tile=identical, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE project elapsed_ms=93 case="m2 n256 k128 offsets" worst=0.9205 of bound at row 0 col 35, max_abs=1.5103191137313843e-2, tile=identical, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE project elapsed_ms=96 case="m3 n8 k64 cancel" worst=0.6477 of bound at row 0 col 4, max_abs=4.5108795166015625e-4, tile=identical, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE project elapsed_ms=102 case="m1 n16 k2048" worst=0.2208 of bound at row 0 col 12, max_abs=1.1380931362509727e-2, tile=identical, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE project elapsed_ms=113 case="m2 n16 k6144 offsets" worst=0.0988 of bound at row 1 col 4, max_abs=2.9790184809826314e-2, tile=identical, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE project elapsed_ms=597 case="m32 n2048 k6144" worst=0.3003 of bound at row 16 col 745, max_abs=1.2159601086750627e-1, tile=identical, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE project elapsed_ms=763 case="m1 n4096 k2048" worst=0.5109 of bound at row 0 col 2175, max_abs=3.1081650871783495e-2, tile=identical, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE project elapsed_ms=1007 case="m1 n6144 k2048" worst=0.5713 of bound at row 0 col 3328, max_abs=5.346063122851774e-2, tile=identical, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE project elapsed_ms=1254 case="m1 n2048 k6144 offsets" worst=0.1835 of bound at row 0 col 422, max_abs=6.236894195899367e-2, tile=identical, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE project elapsed_ms=1782 case="m32 n6144 k2048" worst=0.6675 of bound at row 21 col 5456, max_abs=6.222089286893606e-2, tile=identical, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE project elapsed_ms=7281 case="m1 n140000 k2048 chunked on a 128 MiB range" worst=0.6567 of bound at row 0 col 62267, max_abs=6.247375486418605e-2, tile=identical, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +RESULT stage1 elapsed_ms=7304 on "Apple M1 (G13G B1)": passed=21 total=21 diff --git a/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-stage2-gpu.txt b/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-stage2-gpu.txt new file mode 100644 index 00000000000..f74fac40a6c --- /dev/null +++ b/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-stage2-gpu.txt @@ -0,0 +1,11 @@ +MILESTONE device-found elapsed_ms=5 name="Apple M1 (G13G B1)" vendor=0x10005 device=0x0000 type=integrated api=1.4.362 driver="Honeykrisp" info="Mesa 26.3.0-devel (git-439dc032d7)" subgroup=32 storage16=true int16=true max-storage-range=2147483647 max-buffer-size=137438953472 fence-timeout=10000ms preserve32=false flush32=true denorm-independence=1 +MILESTONE residual-add elapsed_ms=23 case="random 512" elements=512 matches f32-add + RNE-store with subnormal-class-flush-allowed=false, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE residual-add elapsed_ms=25 case="cancel 512" elements=512 matches f32-add + RNE-store with subnormal-class-flush-allowed=false, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE residual-add elapsed_ms=27 case="classes 20" elements=20 matches f32-add + RNE-store with subnormal-class-flush-allowed=true, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE swiglu elapsed_ms=28 case="random 512" elements=512 worst=0.9787 of bound at element 124, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE swiglu elapsed_ms=30 case="controls 16" elements=16 worst=0.5616 of bound at element 9, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE rms-norm elapsed_ms=31 case="1x2048" rows=1 width=2048 worst=0.9873 of bound at row 0 col 1339, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE rms-norm elapsed_ms=33 case="4x2048 offsets zero row" rows=4 width=2048 worst=0.9873 of bound at row 0 col 1339, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused, zero row exactly zero +MILESTONE rms-norm elapsed_ms=35 case="2x2048 dynamic range" rows=2 width=2048 worst=0.9791 of bound at row 1 col 220, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE rms-norm elapsed_ms=36 case="2x256 head_norm offsets" rows=2 width=256 worst=0.9778 of bound at row 1 col 67, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +RESULT stage2 elapsed_ms=36 on "Apple M1 (G13G B1)": passed=9 total=9 diff --git a/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-stage3-gpu.txt b/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-stage3-gpu.txt new file mode 100644 index 00000000000..6e908bd281b --- /dev/null +++ b/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-stage3-gpu.txt @@ -0,0 +1,7 @@ +MILESTONE device-found elapsed_ms=4 name="Apple M1 (G13G B1)" vendor=0x10005 device=0x0000 api=1.4.362 storage16=true int16=true shared=32768 fence-timeout=10000ms +MILESTONE rope elapsed_ms=159 case="queries rows4 base0" rows=4 base=0 worst=4.964008589780876e-1 of bound, repeats byte-identical, passthrough tail bit-exact, input unchanged, planted outputs refused +MILESTONE rope elapsed_ms=161 case="queries rows2 base5" rows=2 base=5 worst=4.8822881879460084e-1 of bound, repeats byte-identical, passthrough tail bit-exact, input unchanged, planted outputs refused +MILESTONE cache elapsed_ms=163 case="cache prefix2 base2 rows3 cap8" base=2 rows=3 capacity=8 keys worst=4.954321215226423e-1 of bound, V transport bit-exact, poison rows untouched, inputs unchanged +MILESTONE sdpa elapsed_ms=170 case="sdpa prefix2 rows3 cap8" prefix=2 rows=3 visible-counts={ 3, 4, 5 } worst=5.687629560222356e-2 of bound, repeats byte-identical, nan/large-finite/zero future keys bit-identical, q unchanged, planted outputs refused +MILESTONE gate elapsed_ms=173 case="gate rows4" rows=4 worst=4.8945424413889344e-1 of bound, repeats byte-identical, inputs unchanged, planted outputs refused +RESULT stage3 elapsed_ms=173 on "Apple M1 (G13G B1)": passed=4 total=4 diff --git a/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-stage4-gpu.txt b/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-stage4-gpu.txt new file mode 100644 index 00000000000..0ded2127598 --- /dev/null +++ b/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-stage4-gpu.txt @@ -0,0 +1,10 @@ +MILESTONE device-found elapsed_ms=8 name="Apple M1 (G13G B1)" vendor=0x10005 device=0x0000 api=1.4.362 storage16=true int16=true subgroups=63 shuffle-relative=true subgroup-size=32 fence-timeout=10000ms +MILESTONE conv-split elapsed_ms=106 tokens=8 split worst=1.1318588482563632e-2 conv worst=9.960479179851554e-1 of bounds, conv input and state carry-out bit-exact, repeats byte-identical, inputs unchanged, planted outputs refused +MILESTONE split-standalone worst=3.689453278087952e-2 of bound +MILESTONE prefill elapsed_ms=132 tokens=8 chunk=8 chunks=null worst=6.315034419925415e-1 of bound, hf matches independent f64 reference from poison, repeats byte-identical, h0 and soa unchanged, planted outputs refused +MILESTONE split-standalone worst=3.689453278087952e-2 of bound +MILESTONE prefill elapsed_ms=172 tokens=8 chunk=4 chunks=1 worst=6.315034419925415e-1 of bound, hf matches independent f64 reference from poison, repeats byte-identical, h0 and soa unchanged, planted outputs refused +MILESTONE decode elapsed_ms=186 worst=9.852585125692864e-1 of bound, repeats byte-identical, inputs unchanged, planted outputs refused +MILESTONE conv-decode elapsed_ms=193 conv worst=9.960479179851554e-1 split worst=7.421311625377545e-3 of bounds, carry-out state bit-exact, split ran on the ACTUAL activated bits, repeats byte-identical, inputs unchanged, planted outputs refused +MILESTONE norm elapsed_ms=195 rows=4 worst=2.750397409250156e-3 of bound, repeats byte-identical, inputs unchanged, planted outputs refused +RESULT stage4 elapsed_ms=195 on "Apple M1 (G13G B1)": passed=6 total=6 diff --git a/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-stage5-gpu.txt b/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-stage5-gpu.txt new file mode 100644 index 00000000000..5d6b58bd7dc --- /dev/null +++ b/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-stage5-gpu.txt @@ -0,0 +1,12 @@ +MILESTONE device-found elapsed_ms=8 name="Apple M1 (G13G B1)" vendor=0x10005 device=0x0000 api=1.4.362 storage16=true int16=true subgroups=63 shuffle-relative=true subgroup-size=32 fence-timeout=10000ms +MILESTONE greedy elapsed_ms=36 gate="vocab-248320" count=248320 token=124160 reference=124160 repeats=3 identical inputs-guards-unchanged +MILESTONE greedy elapsed_ms=49 gate="ties-across-chunks" count=248320 token=1 reference=1 repeats=3 identical inputs-guards-unchanged +MILESTONE greedy elapsed_ms=63 gate="all-negative" count=248320 token=124160 reference=124160 repeats=3 identical inputs-guards-unchanged +MILESTONE greedy elapsed_ms=77 gate="signed-zero-both-orders" count=248320 plus-first-token=5 minus-first-token=5 reference=lowest-index repeats=3 identical inputs-guards-unchanged +error: greedy: non-finite bf16 logit at index 82773 of 248320; row refused +error: greedy: non-finite bf16 logit at index 124160 of 248320; row refused +error: greedy: non-finite bf16 logit at index 165546 of 248320; row refused +MILESTONE greedy elapsed_ms=98 gate="nonfinite-refused" count=248320 nan-first-bad=82773 pinf-first-bad=124160 ninf-first-bad=165546 verdict=NonFinite decode-refuses inputs-guards-unchanged +MILESTONE greedy elapsed_ms=109 gate="min-count-2" count=2 token=1 reference=1 repeats=3 identical inputs-guards-unchanged +MILESTONE greedy elapsed_ms=109 gate="odd-count-refused" count=7 refused=Misaligned before-any-dispatch +RESULT stage5 elapsed_ms=109 on "Apple M1 (G13G B1)": passed=7 total=7 diff --git a/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-vulkan-greedy-contract.txt b/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-vulkan-greedy-contract.txt new file mode 100644 index 00000000000..77b272bb33b --- /dev/null +++ b/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-vulkan-greedy-contract.txt @@ -0,0 +1,13 @@ +MILESTONE greedy-contract device-found name="Apple M1 (G13G B1)" vendor=0x10005 device=0x0000 api=1.4.362 storage16=true int16=true fence-timeout=10000ms +error: greedy: non-finite bf16 logit at index 0 of 248320; row refused +error: greedy: non-finite bf16 logit at index 0 of 248320; row refused +error: greedy: non-finite bf16 logit at index 0 of 248320; row refused +error: greedy: non-finite bf16 logit at index 0 of 248320; row refused +error: greedy: non-finite bf16 logit at index 0 of 248320; row refused +PASS nonfinite vulkan greedy draw: finite_logits, argreduce_chunk, argreduce_combine (bf16): the stated rule holds: any non-finite anywhere in the row refuses it, winner or not; the largest value wins; exact ties (including +0/-0) go to the lowest index +PASS repeat vulkan greedy draw: finite_logits, argreduce_chunk, argreduce_combine (bf16): five runs byte-identical (4 bytes), inputs and guards untouched, no overlap +PASS shapes vulkan greedy draw: finite_logits, argreduce_chunk, argreduce_combine (bf16): 7 shapes exact by its own check (worst 0.000 of the bound), 0 edges refused by the kernel +N/A sums vulkan greedy draw: finite_logits, argreduce_chunk, argreduce_combine (bf16): declares none sums +FAIL nonfinite control: fmax row max (drops a NaN): row 0 did not refuse: it reported 2139095040 where the first non-finite sits at 0 +CONTROL control: fmax row max (drops a NaN): red at nonfinite, as planted +SUMMARY vulkan greedy contract on "Apple M1 (G13G B1)": 1 kernel, 1 pass gates 1-4, 0 fail; 1 control: 1 red as planted, 0 missed diff --git a/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-vulkan-kernel-contract.txt b/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-vulkan-kernel-contract.txt new file mode 100644 index 00000000000..632fe0ea072 --- /dev/null +++ b/docs/receipts/vulkan-qmm-tile-m1/runners-2a5a06c2/tf-vulkan-kernel-contract.txt @@ -0,0 +1,45 @@ +LIMITS deviceName="Apple M1 (G13G B1)" +LIMITS driverName="Honeykrisp" driverInfo="Mesa 26.3.0-devel (git-439dc032d7)" subgroupSize=32 +LIMITS maxComputeWorkGroupInvocations=1024 maxComputeWorkGroupSize=1024,1024,64 maxComputeSharedMemorySize=32768 +LIMITS maxStorageBufferRange=2147483647 minStorageBufferOffsetAlignment=16 +PASS limits qmm_bf16: local=256x1x1 invocations=256 shared=0 +PASS limits qmm_q4_vec4: local=256x1x1 invocations=256 shared=0 +PASS limits qmm_q4_tile: local=256x1x1 invocations=256 shared=0 +PASS limits qmm_q4_tile_f32: local=256x1x1 invocations=256 shared=0 +PASS limits control:qmm_fmax_nonfinite: local=256x1x1 invocations=256 shared=0 +PASS limits control:atomic_sum: local=256x1x1 invocations=256 shared=0 +PASS limits control:scribble_input: local=256x1x1 invocations=256 shared=0 +PASS limits control:drop_tail: local=256x1x1 invocations=256 shared=0 +PASS limits control:fp16_accum: local=256x1x1 invocations=256 shared=0 +PASS limits control:over_limit: local=2048x1x1 invocations=2048 shared=0 +PASS nonfinite vulkan qmm_bf16 affine q4 g64 transposed: 64 non-finite and 16 finite outputs match the float64 classes +PASS repeat vulkan qmm_bf16 affine q4 g64 transposed: five runs byte-identical (64 bytes), inputs and guards untouched, no overlap +PASS shapes vulkan qmm_bf16 affine q4 g64 transposed: 15 shapes within the float64 bound (worst 0.990 of the bound), 0 edges refused by the kernel +PASS sums vulkan qmm_bf16 affine q4 g64 transposed: exact sums at k 6144, where fp16 sums would miss by 5268x and bf16 by 256x +PASS nonfinite vulkan qmm_q4_vec4 affine q4 g64 transposed: 64 non-finite and 16 finite outputs match the float64 classes +PASS repeat vulkan qmm_q4_vec4 affine q4 g64 transposed: five runs byte-identical (64 bytes), inputs and guards untouched, no overlap +PASS shapes vulkan qmm_q4_vec4 affine q4 g64 transposed: 15 shapes within the float64 bound (worst 0.990 of the bound), 0 edges refused by the kernel +PASS sums vulkan qmm_q4_vec4 affine q4 g64 transposed: exact sums at k 6144, where fp16 sums would miss by 4902x and bf16 by 256x +PASS nonfinite vulkan qmm_q4_tile affine q4 g64 transposed tiled 32: 64 non-finite and 16 finite outputs match the float64 classes +PASS repeat vulkan qmm_q4_tile affine q4 g64 transposed tiled 32: five runs byte-identical (64 bytes), inputs and guards untouched, no overlap +PASS shapes vulkan qmm_q4_tile affine q4 g64 transposed tiled 32: 15 shapes within the float64 bound (worst 0.987 of the bound), 0 edges refused by the kernel +PASS sums vulkan qmm_q4_tile affine q4 g64 transposed tiled 32: exact sums at k 6144, where fp16 sums would miss by 5756x and bf16 by 256x +FAIL nonfinite control: qmm fmax drops NaN: 1 of 5 outputs differ in class; first at 0: finite where float64 gives nan +CONTROL control: qmm fmax drops NaN: red at nonfinite, as planted +FAIL repeat control: atomic add row sum: 4 of 4 repeats differ from the first run's bytes, first at byte 0 +CONTROL control: atomic add row sum: red at repeat, as planted +FAIL repeat control: row sum scribbles input: 4 of 4 repeats differ from the first run's bytes, first at byte 0 +CONTROL control: row sum scribbles input: red at repeat, as planted +FAIL shapes control: row sum drops K tail: at 4x1x4097 output 0 is 8.023634e1 where float64 gives 8.094338e1: 224.57 of its bound +CONTROL control: row sum drops K tail: red at shapes, as planted +FAIL sums control: dot accumulates in fp16: infx over its exact-sum bound at k 4096: not fp32 sums (fp16 would miss 33554432x, bf16 16777216x) +CONTROL control: dot accumulates in fp16: red at sums, as planted +FAIL repeat control: qmm writes its weights: a run changed its inputs +CONTROL control: qmm writes its weights: red at repeat, as planted +FAIL repeat control: qmm writes its biases: a run changed its inputs +CONTROL control: qmm writes its biases: red at repeat, as planted +PASS nonfinite vulkan qmm_q4_tile_f32 affine q4 g64 transposed tiled 32 f32 activations: 64 non-finite and 16 finite outputs match the float64 classes +PASS repeat vulkan qmm_q4_tile_f32 affine q4 g64 transposed tiled 32 f32 activations: five runs byte-identical (96 bytes), inputs and guards untouched, no overlap +PASS shapes vulkan qmm_q4_tile_f32 affine q4 g64 transposed tiled 32 f32 activations: 15 shapes within the float64 bound (worst 0.034 of the bound), 0 edges refused by the kernel +PASS sums vulkan qmm_q4_tile_f32 affine q4 g64 transposed tiled 32 f32 activations: exact sums at k 6144, where fp16 sums would miss by 882128357x and bf16 by 16777216x +SUMMARY Apple M1 (G13G B1) 4 kernels: 4 pass gates 1-4, 0 fail, 0 known findings; 8 controls: 8 caught, 0 missed; 0 limit failures; gates 41802 ms diff --git a/docs/receipts/vulkan-qmm-tile-m1/summary.json b/docs/receipts/vulkan-qmm-tile-m1/summary.json new file mode 100644 index 00000000000..3c09cb79da9 --- /dev/null +++ b/docs/receipts/vulkan-qmm-tile-m1/summary.json @@ -0,0 +1,844 @@ +{ + "change": { + "old_build_commit": "e742e02320bcddb3969d80b04f5cd1a11591625b", + "old_build_is": "the qmm_q4_vec4 build", + "new_build_commit": "618ed699d0c62a7114afb758c2037974079a4559", + "new_build_elf_built_from_commit": "6a158edcf2a02050181b85c3b2eb6a69d9f1cf31 (identical tree; only the first commit's message was reworded)", + "new_build_adds": "qmm_q4_tile: weights, scales and biases repacked into zero-padded 32-column tiles at load; same arithmetic and order as qmm_q4_vec4", + "interim_build_commit": "a2dcaa1d4a0148d75f967fe3a4848264f07f31b8" + }, + "setup": { + "chip": "Apple M1 (G13G B1)", + "gpu_driver": "Honeykrisp, Mesa 26.3.0-devel (git-9323171c12)", + "os": "Linux", + "kernel_release": "7.1.12-2-20261013.18-sep-ARCH", + "checkpoint": "Qwen3.5-2B-MLX-4bit, revision 93760be4", + "quantization": "4-bit, group 64, bf16 activations", + "settings": "free-running greedy, 64 steps, two runs per process, vendor id 0x10005, one model process at a time, fresh process per run", + "metric": "run 2 duration: prefill of the prompt plus 63 decode steps plus the per-step logits readback; from the probe's own stderr milestones, 1 ms resolution", + "order": "old,new,new,old,old,new (three pairs, both orders)", + "decision_rule": "median run-2 gain of at least 3 percent and every pair positive" + }, + "identity_probe": { + "exit_code": 0, + "result": [ + "RESULT stage1 elapsed_ms=7550 on \"Apple M1 (G13G B1)\": passed=21 total=21" + ], + "tiled_image_vs_default_image": [ + "m1 n8 k64: identical", + "m2 n257 k128: identical", + "m2 n256 k128 offsets: identical", + "m3 n8 k64 cancel: identical", + "m1 n16 k2048: identical", + "m2 n16 k6144 offsets: identical", + "m32 n2048 k6144: identical", + "m1 n4096 k2048: identical", + "m1 n6144 k2048: identical", + "m1 n2048 k6144 offsets: identical", + "m32 n6144 k2048: identical", + "m1 n140000 k2048 chunked on a 128 MiB range: identical" + ] + }, + "cases": [ + { + "case": "unicode-24", + "status": "PASS_WITH_NOT_VERIFIED", + "kernel_release": "7.1.12-2-20261013.18-sep-ARCH", + "start_utc": "2026-11-03T08:36:08+00:00", + "end_utc": "2026-11-03T08:37:53+00:00", + "elfs": { + "old": { + "sha256": "683fcf46a89c3d99d6db94c3f9fda614a9340d7193dfdc301451225f5db20ed8", + "bytes": 5778952 + }, + "new": { + "sha256": "d4949bef20e3b175365d26fab943cafe6f28befdd060185366ff2713368530a3", + "bytes": 5809824 + } + }, + "logits_sha256_all_six_processes": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f", + "processes": [ + { + "label": "pr0-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 401, + "run1_ms": 9204, + "run2_ms": 9155, + "peak_rss_bytes": 118636544, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + }, + { + "label": "pr1-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 714, + "run1_ms": 7429, + "run2_ms": 7387, + "peak_rss_bytes": 327843840, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + }, + { + "label": "pr2-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 714, + "run1_ms": 7237, + "run2_ms": 7253, + "peak_rss_bytes": 325926912, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + }, + { + "label": "pr3-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 396, + "run1_ms": 8924, + "run2_ms": 8878, + "peak_rss_bytes": 86360064, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + }, + { + "label": "pr4-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 397, + "run1_ms": 9001, + "run2_ms": 9177, + "peak_rss_bytes": 95977472, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + }, + { + "label": "pr5-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 697, + "run1_ms": 7145, + "run2_ms": 6950, + "peak_rss_bytes": 327876608, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + } + ], + "pairs": [ + { + "pair": 1, + "order": [ + "old", + "new" + ], + "run2_old_ms": 9155, + "run2_new_ms": 7387, + "run2_gain": 0.193119, + "run1_old_ms": 9204, + "run1_new_ms": 7429, + "run1_gain": 0.192851 + }, + { + "pair": 2, + "order": [ + "new", + "old" + ], + "run2_old_ms": 8878, + "run2_new_ms": 7253, + "run2_gain": 0.183037, + "run1_old_ms": 8924, + "run1_new_ms": 7237, + "run1_gain": 0.189041 + }, + { + "pair": 3, + "order": [ + "old", + "new" + ], + "run2_old_ms": 9177, + "run2_new_ms": 6950, + "run2_gain": 0.242672, + "run1_old_ms": 9001, + "run1_new_ms": 7145, + "run1_gain": 0.206199 + } + ], + "run2_gain_median": 0.193119, + "run2_gain_min": 0.183037, + "run2_gain_max": 0.242672, + "run1_gain_median": 0.192851, + "speedup_rule": "MET", + "raw_result_json_sha256": "36f2f22b846d4b11a99122a0c4f8e97a24ac391324a0962b660c49c660f4b62b" + }, + { + "case": "document-256", + "status": "PASS_WITH_NOT_VERIFIED", + "kernel_release": "7.1.12-2-20261013.18-sep-ARCH", + "start_utc": "2026-11-03T08:42:19+00:00", + "end_utc": "2026-11-03T08:45:33+00:00", + "elfs": { + "old": { + "sha256": "683fcf46a89c3d99d6db94c3f9fda614a9340d7193dfdc301451225f5db20ed8", + "bytes": 5778952 + }, + "new": { + "sha256": "d4949bef20e3b175365d26fab943cafe6f28befdd060185366ff2713368530a3", + "bytes": 5809824 + } + }, + "logits_sha256_all_six_processes": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb", + "processes": [ + { + "label": "pr0-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 402, + "run1_ms": 20331, + "run2_ms": 20567, + "peak_rss_bytes": 87228416, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + }, + { + "label": "pr1-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 712, + "run1_ms": 10943, + "run2_ms": 10593, + "peak_rss_bytes": 327270400, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + }, + { + "label": "pr2-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 712, + "run1_ms": 10668, + "run2_ms": 10154, + "peak_rss_bytes": 326615040, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + }, + { + "label": "pr3-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 400, + "run1_ms": 20796, + "run2_ms": 20535, + "peak_rss_bytes": 86999040, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + }, + { + "label": "pr4-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 402, + "run1_ms": 20886, + "run2_ms": 20444, + "peak_rss_bytes": 87588864, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + }, + { + "label": "pr5-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 713, + "run1_ms": 10822, + "run2_ms": 10240, + "peak_rss_bytes": 327450624, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + } + ], + "pairs": [ + { + "pair": 1, + "order": [ + "old", + "new" + ], + "run2_old_ms": 20567, + "run2_new_ms": 10593, + "run2_gain": 0.484952, + "run1_old_ms": 20331, + "run1_new_ms": 10943, + "run1_gain": 0.461758 + }, + { + "pair": 2, + "order": [ + "new", + "old" + ], + "run2_old_ms": 20535, + "run2_new_ms": 10154, + "run2_gain": 0.505527, + "run1_old_ms": 20796, + "run1_new_ms": 10668, + "run1_gain": 0.487017 + }, + { + "pair": 3, + "order": [ + "old", + "new" + ], + "run2_old_ms": 20444, + "run2_new_ms": 10240, + "run2_gain": 0.49912, + "run1_old_ms": 20886, + "run1_new_ms": 10822, + "run1_gain": 0.481854 + } + ], + "run2_gain_median": 0.49912, + "run2_gain_min": 0.484952, + "run2_gain_max": 0.505527, + "run1_gain_median": 0.481854, + "speedup_rule": "MET", + "raw_result_json_sha256": "0112e8e1d578a77a796fc4ce49e2ee161a4f5b7605631d38e0e9e56ee419e24f" + }, + { + "case": "document-1024", + "status": "PASS_WITH_NOT_VERIFIED", + "kernel_release": "7.1.12-2-20261013.18-sep-ARCH", + "start_utc": "2026-11-03T08:49:52+00:00", + "end_utc": "2026-11-03T08:58:11+00:00", + "elfs": { + "old": { + "sha256": "683fcf46a89c3d99d6db94c3f9fda614a9340d7193dfdc301451225f5db20ed8", + "bytes": 5778952 + }, + "new": { + "sha256": "d4949bef20e3b175365d26fab943cafe6f28befdd060185366ff2713368530a3", + "bytes": 5809824 + } + }, + "logits_sha256_all_six_processes": "7f9d8e65e90eb3751edc41aed99c9ac5d37ce3f7d99dae1e626497632661a820", + "processes": [ + { + "label": "pr0-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 406, + "run1_ms": 59589, + "run2_ms": 59651, + "peak_rss_bytes": 87539712, + "logits_sha256": "7f9d8e65e90eb3751edc41aed99c9ac5d37ce3f7d99dae1e626497632661a820" + }, + { + "label": "pr1-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 705, + "run1_ms": 22098, + "run2_ms": 23047, + "peak_rss_bytes": 326664192, + "logits_sha256": "7f9d8e65e90eb3751edc41aed99c9ac5d37ce3f7d99dae1e626497632661a820" + }, + { + "label": "pr2-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 714, + "run1_ms": 22025, + "run2_ms": 22310, + "peak_rss_bytes": 326696960, + "logits_sha256": "7f9d8e65e90eb3751edc41aed99c9ac5d37ce3f7d99dae1e626497632661a820" + }, + { + "label": "pr3-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 413, + "run1_ms": 59403, + "run2_ms": 59711, + "peak_rss_bytes": 87572480, + "logits_sha256": "7f9d8e65e90eb3751edc41aed99c9ac5d37ce3f7d99dae1e626497632661a820" + }, + { + "label": "pr4-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 411, + "run1_ms": 59789, + "run2_ms": 59652, + "peak_rss_bytes": 87556096, + "logits_sha256": "7f9d8e65e90eb3751edc41aed99c9ac5d37ce3f7d99dae1e626497632661a820" + }, + { + "label": "pr5-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 696, + "run1_ms": 22667, + "run2_ms": 21499, + "peak_rss_bytes": 326877184, + "logits_sha256": "7f9d8e65e90eb3751edc41aed99c9ac5d37ce3f7d99dae1e626497632661a820" + } + ], + "pairs": [ + { + "pair": 1, + "order": [ + "old", + "new" + ], + "run2_old_ms": 59651, + "run2_new_ms": 23047, + "run2_gain": 0.613636, + "run1_old_ms": 59589, + "run1_new_ms": 22098, + "run1_gain": 0.62916 + }, + { + "pair": 2, + "order": [ + "new", + "old" + ], + "run2_old_ms": 59711, + "run2_new_ms": 22310, + "run2_gain": 0.626367, + "run1_old_ms": 59403, + "run1_new_ms": 22025, + "run1_gain": 0.629227 + }, + { + "pair": 3, + "order": [ + "old", + "new" + ], + "run2_old_ms": 59652, + "run2_new_ms": 21499, + "run2_gain": 0.639593, + "run1_old_ms": 59789, + "run1_new_ms": 22667, + "run1_gain": 0.620883 + } + ], + "run2_gain_median": 0.626367, + "run2_gain_min": 0.613636, + "run2_gain_max": 0.639593, + "run1_gain_median": 0.62916, + "speedup_rule": "MET", + "raw_result_json_sha256": "07629cfec834975af4acb84c4b50f0c28ab3cb6433cd4729e2f279a4a3a4e386" + } + ], + "interim_cases_before_padding": [ + { + "case": "unicode-24", + "status": "PASS_WITH_NOT_VERIFIED", + "kernel_release": "7.1.12-2-20261013.18-sep-ARCH", + "start_utc": "2026-11-03T07:14:43+00:00", + "end_utc": "2026-11-03T07:16:26+00:00", + "elfs": { + "old": { + "sha256": "683fcf46a89c3d99d6db94c3f9fda614a9340d7193dfdc301451225f5db20ed8", + "bytes": 5778952 + }, + "new": { + "sha256": "3a22c95969c1b05c607f61c11f53628eb267fe88bc79fbb2caf6cefef47b2fce", + "bytes": 5812152 + } + }, + "logits_sha256_all_six_processes": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f", + "processes": [ + { + "label": "pr0-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 404, + "run1_ms": 9151, + "run2_ms": 8805, + "peak_rss_bytes": 87605248, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + }, + { + "label": "pr1-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 677, + "run1_ms": 7175, + "run2_ms": 7091, + "peak_rss_bytes": 311967744, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + }, + { + "label": "pr2-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 689, + "run1_ms": 6852, + "run2_ms": 6926, + "peak_rss_bytes": 311984128, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + }, + { + "label": "pr3-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 403, + "run1_ms": 9196, + "run2_ms": 9199, + "peak_rss_bytes": 86769664, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + }, + { + "label": "pr4-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 405, + "run1_ms": 9257, + "run2_ms": 9030, + "peak_rss_bytes": 87605248, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + }, + { + "label": "pr5-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 693, + "run1_ms": 7255, + "run2_ms": 6539, + "peak_rss_bytes": 311820288, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + } + ], + "pairs": [ + { + "pair": 1, + "order": [ + "old", + "new" + ], + "run2_old_ms": 8805, + "run2_new_ms": 7091, + "run2_gain": 0.194662, + "run1_old_ms": 9151, + "run1_new_ms": 7175, + "run1_gain": 0.215933 + }, + { + "pair": 2, + "order": [ + "new", + "old" + ], + "run2_old_ms": 9199, + "run2_new_ms": 6926, + "run2_gain": 0.247092, + "run1_old_ms": 9196, + "run1_new_ms": 6852, + "run1_gain": 0.254893 + }, + { + "pair": 3, + "order": [ + "old", + "new" + ], + "run2_old_ms": 9030, + "run2_new_ms": 6539, + "run2_gain": 0.275858, + "run1_old_ms": 9257, + "run1_new_ms": 7255, + "run1_gain": 0.216269 + } + ], + "run2_gain_median": 0.247092, + "run2_gain_min": 0.194662, + "run2_gain_max": 0.275858, + "run1_gain_median": 0.216269, + "speedup_rule": "MET", + "raw_result_json_sha256": "eaefc453eed30e7d28414cf766bc36855852e458295b634ed6df8a9db1930da3" + }, + { + "case": "document-256", + "status": "PASS_WITH_NOT_VERIFIED", + "kernel_release": "7.1.12-2-20261013.18-sep-ARCH", + "start_utc": "2026-11-03T07:32:44+00:00", + "end_utc": "2026-11-03T07:35:57+00:00", + "elfs": { + "old": { + "sha256": "683fcf46a89c3d99d6db94c3f9fda614a9340d7193dfdc301451225f5db20ed8", + "bytes": 5778952 + }, + "new": { + "sha256": "3a22c95969c1b05c607f61c11f53628eb267fe88bc79fbb2caf6cefef47b2fce", + "bytes": 5812152 + } + }, + "logits_sha256_all_six_processes": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb", + "processes": [ + { + "label": "pr0-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 396, + "run1_ms": 20458, + "run2_ms": 20466, + "peak_rss_bytes": 87572480, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + }, + { + "label": "pr1-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 710, + "run1_ms": 10482, + "run2_ms": 10481, + "peak_rss_bytes": 311885824, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + }, + { + "label": "pr2-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 692, + "run1_ms": 10155, + "run2_ms": 10684, + "peak_rss_bytes": 311574528, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + }, + { + "label": "pr3-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 402, + "run1_ms": 20661, + "run2_ms": 20552, + "peak_rss_bytes": 87588864, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + }, + { + "label": "pr4-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 404, + "run1_ms": 20679, + "run2_ms": 20759, + "peak_rss_bytes": 95256576, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + }, + { + "label": "pr5-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 684, + "run1_ms": 10021, + "run2_ms": 10246, + "peak_rss_bytes": 311197696, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + } + ], + "pairs": [ + { + "pair": 1, + "order": [ + "old", + "new" + ], + "run2_old_ms": 20466, + "run2_new_ms": 10481, + "run2_gain": 0.487882, + "run1_old_ms": 20458, + "run1_new_ms": 10482, + "run1_gain": 0.487633 + }, + { + "pair": 2, + "order": [ + "new", + "old" + ], + "run2_old_ms": 20552, + "run2_new_ms": 10684, + "run2_gain": 0.480148, + "run1_old_ms": 20661, + "run1_new_ms": 10155, + "run1_gain": 0.508494 + }, + { + "pair": 3, + "order": [ + "old", + "new" + ], + "run2_old_ms": 20759, + "run2_new_ms": 10246, + "run2_gain": 0.506431, + "run1_old_ms": 20679, + "run1_new_ms": 10021, + "run1_gain": 0.515402 + } + ], + "run2_gain_median": 0.487882, + "run2_gain_min": 0.480148, + "run2_gain_max": 0.506431, + "run1_gain_median": 0.508494, + "speedup_rule": "MET", + "raw_result_json_sha256": "67a7b4d75f65b4422b61772a90d1d9d807f226fc49bafb1072db234e321b4792" + }, + { + "case": "document-1024", + "status": "PASS_WITH_NOT_VERIFIED", + "kernel_release": "7.1.12-2-20261013.18-sep-ARCH", + "start_utc": "2026-11-03T07:40:25+00:00", + "end_utc": "2026-11-03T07:48:45+00:00", + "elfs": { + "old": { + "sha256": "683fcf46a89c3d99d6db94c3f9fda614a9340d7193dfdc301451225f5db20ed8", + "bytes": 5778952 + }, + "new": { + "sha256": "3a22c95969c1b05c607f61c11f53628eb267fe88bc79fbb2caf6cefef47b2fce", + "bytes": 5812152 + } + }, + "logits_sha256_all_six_processes": "7f9d8e65e90eb3751edc41aed99c9ac5d37ce3f7d99dae1e626497632661a820", + "processes": [ + { + "label": "pr0-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 426, + "run1_ms": 59518, + "run2_ms": 59445, + "peak_rss_bytes": 94896128, + "logits_sha256": "7f9d8e65e90eb3751edc41aed99c9ac5d37ce3f7d99dae1e626497632661a820" + }, + { + "label": "pr1-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 684, + "run1_ms": 23027, + "run2_ms": 22199, + "peak_rss_bytes": 311066624, + "logits_sha256": "7f9d8e65e90eb3751edc41aed99c9ac5d37ce3f7d99dae1e626497632661a820" + }, + { + "label": "pr2-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 691, + "run1_ms": 23627, + "run2_ms": 22058, + "peak_rss_bytes": 311984128, + "logits_sha256": "7f9d8e65e90eb3751edc41aed99c9ac5d37ce3f7d99dae1e626497632661a820" + }, + { + "label": "pr3-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 415, + "run1_ms": 59663, + "run2_ms": 59557, + "peak_rss_bytes": 96944128, + "logits_sha256": "7f9d8e65e90eb3751edc41aed99c9ac5d37ce3f7d99dae1e626497632661a820" + }, + { + "label": "pr4-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 400, + "run1_ms": 59568, + "run2_ms": 59689, + "peak_rss_bytes": 119373824, + "logits_sha256": "7f9d8e65e90eb3751edc41aed99c9ac5d37ce3f7d99dae1e626497632661a820" + }, + { + "label": "pr5-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 688, + "run1_ms": 22476, + "run2_ms": 22188, + "peak_rss_bytes": 311296000, + "logits_sha256": "7f9d8e65e90eb3751edc41aed99c9ac5d37ce3f7d99dae1e626497632661a820" + } + ], + "pairs": [ + { + "pair": 1, + "order": [ + "old", + "new" + ], + "run2_old_ms": 59445, + "run2_new_ms": 22199, + "run2_gain": 0.626562, + "run1_old_ms": 59518, + "run1_new_ms": 23027, + "run1_gain": 0.613109 + }, + { + "pair": 2, + "order": [ + "new", + "old" + ], + "run2_old_ms": 59557, + "run2_new_ms": 22058, + "run2_gain": 0.629632, + "run1_old_ms": 59663, + "run1_new_ms": 23627, + "run1_gain": 0.603992 + }, + { + "pair": 3, + "order": [ + "old", + "new" + ], + "run2_old_ms": 59689, + "run2_new_ms": 22188, + "run2_gain": 0.628273, + "run1_old_ms": 59568, + "run1_new_ms": 22476, + "run1_gain": 0.622683 + } + ], + "run2_gain_median": 0.628273, + "run2_gain_min": 0.626562, + "run2_gain_max": 0.629632, + "run1_gain_median": 0.613109, + "speedup_rule": "MET", + "raw_result_json_sha256": "5011041e4bf80318251f07cfa8d25a00a24c846f4e48467b634dd596efb88518" + } + ] +} diff --git a/docs/receipts/vulkan-qmm-vec4-m1/PREREG.md b/docs/receipts/vulkan-qmm-vec4-m1/PREREG.md new file mode 100644 index 00000000000..5ec5a4126e6 --- /dev/null +++ b/docs/receipts/vulkan-qmm-vec4-m1/PREREG.md @@ -0,0 +1,67 @@ +# qmm_q4_vec4 pre-registration (written 2026-11-03 before any GPU run of this change) + + +## Question + +Does reading the q4 projection's weights and activations with 16-byte loads (`qmm_q4_vec4`) make +Qwen3.5-2B decode faster on the 13-inch M1, with byte-identical outputs? + +## Arms (two ELFs of `tf-qwen35-model`, ReleaseSafe, aarch64) + +- old: TensorFold PR #1544 tip `b3ab9ba2` (scalar `qmm_bf16` only). +- new: the same plus commit `e742e023` (`qmm_q4_vec4` chosen when x and w windows are 16-byte aligned). +- Nothing else differs. Same checkpoint, same input (unicode-24, 24 prompt tokens), context 512, 64 free-running + greedy steps, two runs per process. + +## Gates (every process, both arms; any failure stops the unit) + +1. Exit 0 and a complete receipt (house receipt check: status ok, oracle ids and hashes, byte-equal repeats). +2. Generated ids equal the pinned ids (hash `6bfbb1a5...`). +3. Full logits bytes of every process equal the first valid old-arm process's logits (one digest for all six). +4. Disk, ICD, checkpoint, input and ELF identity pins hold before the first process starts. + +Identity across arms is a gate. A speed number from a run that fails a gate is not reported as a result. + +## Metric + +Per process, from the probe's own stderr milestones (1 ms resolution): + +- primary: run 2 duration = `run2-done` minus `run1-done` (prefill of 24 tokens plus 63 decode steps, + warm, with the per-step logits readback the probe always does); +- secondary: run 1 duration = `run1-done` minus `engine-ready`. + +Order old,new,new,old,old,new (three pairs, both orders). Per pair: gain = (old - new) / old on run 2. +Report the median of the three pairs, the range, and every raw value. + +## Decision rule + +- Speedup: median pair gain of at least 3 percent AND all three pair gains positive AND all gates hold. +- No speedup: median under 3 percent or any pair not positive. Reported as such; the change stays unmerged. +- Any gate failure: the candidate dies, no speed claim. + +## Forecast (a prediction, written before the run; I do not know the answer) + +I expect run 2 to shrink by 10 to 35 percent: the profile of the earlier build puts quantized matmul near 80 percent of +decode and each thread issues about six times fewer loads. I give a regression (new slower) a 10 percent chance, +because a wider load can raise register pressure on this compiler. + +## Not claimed + +No claim about other chips, other quantizations (3, 5, 6, 8 bit and fp modes keep the scalar image), other models, +or the Metal/MLX path. The software-device (llvmpipe) identity pass does not count as GPU evidence. +Two runs of one build on one host are not a significance test. + +## Addendum 2026-11-03 (written after the first pair, before the follow-up rows) + +First pair result (unicode-24, context 512): all gates held, six logits identical, median run-2 gain 40.09 percent +(pairs 41.0, 40.0, 40.1), rule MET. That was above my 10 to 35 percent forecast. + +Follow-up rows use the same two ELFs, the same gates and the same decision rule on three more cases: +a fresh replicate of unicode-24 (host-state variance), document-256 at context 512, and document-1024 at context 1536. +The documents prefill 8 and 32 chunks of 32 rows, so the projections run with m up to 32, a different regime from m = 1 decode. + +Forecast, written before these rows: every gate holds (the identity probe already passed on this GPU for m up to 32). +Gains on the document inputs are positive but smaller than 40 percent, because prefill and attention take a larger share of +run time. I give 15 percent to a document-1024 gain below 3 percent. The replicate lands within 5 points of 40 percent. +If any gate fails on any case, the candidate is not mergeable on that evidence, whatever the other cases show. +Each case is judged by the rule on its own; the cases are not pooled. diff --git a/docs/receipts/vulkan-qmm-vec4-m1/README.md b/docs/receipts/vulkan-qmm-vec4-m1/README.md new file mode 100644 index 00000000000..e2305d5b951 --- /dev/null +++ b/docs/receipts/vulkan-qmm-vec4-m1/README.md @@ -0,0 +1,32 @@ +# qmm_q4_vec4 receipts, Apple M1 (13-inch) + +Qwen3.5-2B 4-bit, 64 greedy steps, two runs per process. Old build is #1990's head built with the vec4 selection forced to the scalar kernel (`qmmVariant = .scalar` in `ops_q4.zig`, not committed; ELF `4d1c6ba6`), new build is #1990's head (`1e1b7d8c`). The float32 fix sits above the vec4 commits, so an older #1544 build would not carry it. +Unicode-24 and document-256 ran three old and three new processes (old,new,new,old,old,new). Document-1024 ran two pairs, one old,new and one new,old, because the scalar arm takes about 250 s a process. Run 2 is the prefill plus 63 decode steps plus the per-step logits readback. Linux 7.3.0-2-20261110.1-lab-ARCH, Honeykrisp from the Omarchy Mesa build (library sha256 `d75f4cbc`). + +| Case | Old run 2 (median) | New run 2 (median) | Gain: median (pair range) | Logits digest, all processes | +|---|---:|---:|---|---| +| unicode-24 | 15.6 s | 10.2 s | 34.8% (34.6 to 35.2) | `e946185d64ae` | +| document-256 | 40.4 s | 23.8 s | 41.0% (41.0 to 41.0) | `3188aaae7d11` | +| document-1024 (two pairs) | 123.0 s | 70.1 s | 43.0% (43.0 to 43.0) | `82c353787e89` | + +The gain depends on the environment. The first version of this receipt measured 40.1%, 48.5% and 51.5% against #1544's head (`b3ab9ba2`; new build `e742e023`) on earlier builds with Linux 7.1.12-2-20261013.18-sep-ARCH and Mesa git-9323171c12 (digests `e5447d9a3ab5`, `147a2d2e0232`, `7f9d8e65e90e`). These runs do not separate the cause. + +First version's table (earlier builds, earlier environment): + +| Case | Old run 2 (median) | New run 2 (median) | Gain: median (pair range) | Logits digest, all six processes | +|---|---:|---:|---|---| +| unicode-24, first pair | 15.6 s | 9.4 s | 40.1% (40.0 to 41.0) | `e5447d9a3ab5` | +| unicode-24, replicate | 15.3 s | 9.0 s | 41.0% (39.8 to 41.3) | `e5447d9a3ab5` | +| document-256 | 40.0 s | 20.7 s | 48.5% (48.2 to 48.5) | `147a2d2e0232` | +| document-1024 | 122.8 s | 59.5 s | 51.5% (51.5 to 51.5) | `7f9d8e65e90e` | + +Files: + +- `fixed-heads-summary.json` holds every process row (with peak resident memory), the pairs, the ELF hashes and the setup of the fixed-head run. It is generated from the raw run files. +- `summary.json` is the first version's summary (earlier builds, earlier environment). It stays as the record of those runs. +- `identity-probe.txt` holds the probe that runs both images on 17 projection cases and compares the output bytes. It ran on the earlier builds and was not re-run on the fixed head. +- `PREREG.md` was written before any GPU run. It holds the forecast, the gates and the decision rule. +- `raw-files-sha256.txt` lists the SHA-256 of every raw run file the summary was generated from. +- `kernel-contract-*.txt` holds the contract runner's output, one file per GPU, from the earlier builds. It was not re-run on the fixed head. + +Byte identity: each process writes the SHA-256 of its full logits stream. A case passes only if all six digests are equal and the ids match the pinned ids. diff --git a/docs/receipts/vulkan-qmm-vec4-m1/fixed-heads-summary.json b/docs/receipts/vulkan-qmm-vec4-m1/fixed-heads-summary.json new file mode 100644 index 00000000000..911fa8363c7 --- /dev/null +++ b/docs/receipts/vulkan-qmm-vec4-m1/fixed-heads-summary.json @@ -0,0 +1,247 @@ +{ + "change": "qmm_q4_vec4 (#1990) measured on the float32-fixed head", + "setup": { + "chip": "Apple M1 (G13G B1)", + "kernel_release": "7.3.0-2-20261110.1-lab-ARCH", + "mesa": "Honeykrisp, Omarchy Mesa build, library sha256 d75f4cbc", + "checkpoint": "Qwen3.5-2B-MLX-4bit, revision 93760be4", + "settings": "free-running greedy, 64 steps, two runs per process, vendor id 0x10005, one model process at a time, fresh process per run", + "decision_rule": "median run-2 gain of at least 3 percent and every pair positive", + "old_build": "#1990 head built with the vec4 selection forced to the scalar kernel (qmmVariant = .scalar in ops_q4.zig, not committed), ELF sha256 prefix 4d1c6ba6", + "new_build": "#1990 head, ELF sha256 prefix 1e1b7d8c", + "note": "The float32 fix sits above the vec4 commits, so an older #1544 build would not carry it. The old arm is the same head with the scalar kernel selected." + }, + "cases": [ + { + "case": "unicode-24", + "status": [ + "PASS_WITH_NOT_VERIFIED" + ], + "start_utc": [ + "2026-11-11T19:30:52+00:00" + ], + "kernel_release": "7.3.0-2-20261110.1-lab-ARCH", + "elfs": { + "old": { + "sha256": "4d1c6ba6c702220c44303ba674f4c3247e671b3e585d3c83bcd6960d65c2cb31", + "bytes": 6087144 + }, + "new": { + "sha256": "1e1b7d8cdddf57645c2090207eef24812e4cbc6e6b90a11dc9d43492ce90ac5a", + "bytes": 6087520 + } + }, + "order": "old,new,new,old,old,new", + "pairs": 3, + "logits_sha256_all_processes": "e946185d64aea5942265276e7b76f7d4c4560b543145e1139e9c2a09432ca00c", + "run2_ms_median": { + "old": 15588, + "new": 10157 + }, + "run2_gain_median": 0.34840903258917116, + "run2_gain_range": [ + 0.34632530887784996, + 0.35175297523319393 + ], + "processes": [ + { + "label": "pr0-old", + "arm": "old", + "run1_ms": 15336, + "run2_ms": 15588, + "peak_rss_bytes": 1214251008, + "logits_sha256": "e946185d64aea5942265276e7b76f7d4c4560b543145e1139e9c2a09432ca00c" + }, + { + "label": "pr1-new", + "arm": "new", + "run1_ms": 10106, + "run2_ms": 10157, + "peak_rss_bytes": 1215561728, + "logits_sha256": "e946185d64aea5942265276e7b76f7d4c4560b543145e1139e9c2a09432ca00c" + }, + { + "label": "pr2-new", + "arm": "new", + "run1_ms": 10255, + "run2_ms": 10264, + "peak_rss_bytes": 1214332928, + "logits_sha256": "e946185d64aea5942265276e7b76f7d4c4560b543145e1139e9c2a09432ca00c" + }, + { + "label": "pr3-old", + "arm": "old", + "run1_ms": 15606, + "run2_ms": 15702, + "peak_rss_bytes": 1214431232, + "logits_sha256": "e946185d64aea5942265276e7b76f7d4c4560b543145e1139e9c2a09432ca00c" + }, + { + "label": "pr4-old", + "arm": "old", + "run1_ms": 15738, + "run2_ms": 15545, + "peak_rss_bytes": 1215594496, + "logits_sha256": "e946185d64aea5942265276e7b76f7d4c4560b543145e1139e9c2a09432ca00c" + }, + { + "label": "pr5-new", + "arm": "new", + "run1_ms": 9984, + "run2_ms": 10077, + "peak_rss_bytes": 1246920704, + "logits_sha256": "e946185d64aea5942265276e7b76f7d4c4560b543145e1139e9c2a09432ca00c" + } + ] + }, + { + "case": "document-256", + "status": [ + "PASS_WITH_NOT_VERIFIED" + ], + "start_utc": [ + "2026-11-11T19:39:53+00:00" + ], + "kernel_release": "7.3.0-2-20261110.1-lab-ARCH", + "elfs": { + "old": { + "sha256": "4d1c6ba6c702220c44303ba674f4c3247e671b3e585d3c83bcd6960d65c2cb31", + "bytes": 6087144 + }, + "new": { + "sha256": "1e1b7d8cdddf57645c2090207eef24812e4cbc6e6b90a11dc9d43492ce90ac5a", + "bytes": 6087520 + } + }, + "order": "old,new,new,old,old,new", + "pairs": 3, + "logits_sha256_all_processes": "3188aaae7d1182c133a89fde103722d23f3d81b57df259498c9b778ebe2d0d7a", + "run2_ms_median": { + "old": 40352, + "new": 23796 + }, + "run2_gain_median": 0.410289452815226, + "run2_gain_range": [ + 0.4101222506013341, + 0.41031867968478963 + ], + "processes": [ + { + "label": "pr0-old", + "arm": "old", + "run1_ms": 40217, + "run2_ms": 40354, + "peak_rss_bytes": 1215201280, + "logits_sha256": "3188aaae7d1182c133a89fde103722d23f3d81b57df259498c9b778ebe2d0d7a" + }, + { + "label": "pr1-new", + "arm": "new", + "run1_ms": 23852, + "run2_ms": 23796, + "peak_rss_bytes": 1214218240, + "logits_sha256": "3188aaae7d1182c133a89fde103722d23f3d81b57df259498c9b778ebe2d0d7a" + }, + { + "label": "pr2-new", + "arm": "new", + "run1_ms": 23972, + "run2_ms": 23788, + "peak_rss_bytes": 1214545920, + "logits_sha256": "3188aaae7d1182c133a89fde103722d23f3d81b57df259498c9b778ebe2d0d7a" + }, + { + "label": "pr3-old", + "arm": "old", + "run1_ms": 40231, + "run2_ms": 40327, + "peak_rss_bytes": 1215545344, + "logits_sha256": "3188aaae7d1182c133a89fde103722d23f3d81b57df259498c9b778ebe2d0d7a" + }, + { + "label": "pr4-old", + "arm": "old", + "run1_ms": 40383, + "run2_ms": 40352, + "peak_rss_bytes": 1243185152, + "logits_sha256": "3188aaae7d1182c133a89fde103722d23f3d81b57df259498c9b778ebe2d0d7a" + }, + { + "label": "pr5-new", + "arm": "new", + "run1_ms": 23658, + "run2_ms": 23796, + "peak_rss_bytes": 1247002624, + "logits_sha256": "3188aaae7d1182c133a89fde103722d23f3d81b57df259498c9b778ebe2d0d7a" + } + ] + }, + { + "case": "document-1024", + "status": [ + "PASS_WITH_NOT_VERIFIED" + ], + "start_utc": [ + "2026-11-11T20:23:08+00:00", + "2026-11-11T20:35:25+00:00" + ], + "kernel_release": "7.3.0-2-20261110.1-lab-ARCH", + "elfs": { + "old": { + "sha256": "4d1c6ba6c702220c44303ba674f4c3247e671b3e585d3c83bcd6960d65c2cb31", + "bytes": 6087144 + }, + "new": { + "sha256": "1e1b7d8cdddf57645c2090207eef24812e4cbc6e6b90a11dc9d43492ce90ac5a", + "bytes": 6087520 + } + }, + "order": "pair 1 old,new; pair 2 new,old (two pairs, not three: the scalar arm takes about 250 s per process)", + "pairs": 2, + "logits_sha256_all_processes": "82c353787e896b52c8def7f4cebbc087d42f71c4837505db8100df67b2235356", + "run2_ms_median": { + "old": 122959.5, + "new": 70124.5 + }, + "run2_gain_median": 0.4296944653090763, + "run2_gain_range": [ + 0.42951296616849927, + 0.4298759644496533 + ], + "processes": [ + { + "label": "pr0-old", + "arm": "old", + "run1_ms": 123115, + "run2_ms": 123051, + "peak_rss_bytes": 1227128832, + "logits_sha256": "82c353787e896b52c8def7f4cebbc087d42f71c4837505db8100df67b2235356" + }, + { + "label": "pr1-new", + "arm": "new", + "run1_ms": 70313, + "run2_ms": 70199, + "peak_rss_bytes": 1227112448, + "logits_sha256": "82c353787e896b52c8def7f4cebbc087d42f71c4837505db8100df67b2235356" + }, + { + "label": "pr0-new", + "arm": "new", + "run1_ms": 69977, + "run2_ms": 70050, + "peak_rss_bytes": 1259487232, + "logits_sha256": "82c353787e896b52c8def7f4cebbc087d42f71c4837505db8100df67b2235356" + }, + { + "label": "pr1-old", + "arm": "old", + "run1_ms": 122909, + "run2_ms": 122868, + "peak_rss_bytes": 1228144640, + "logits_sha256": "82c353787e896b52c8def7f4cebbc087d42f71c4837505db8100df67b2235356" + } + ] + } + ] +} diff --git a/docs/receipts/vulkan-qmm-vec4-m1/identity-probe.txt b/docs/receipts/vulkan-qmm-vec4-m1/identity-probe.txt new file mode 100644 index 00000000000..d5b47c51f93 --- /dev/null +++ b/docs/receipts/vulkan-qmm-vec4-m1/identity-probe.txt @@ -0,0 +1,11 @@ +MILESTONE device-found elapsed_ms=25 name="Apple M1 (G13G B1)" vendor=0x10005 device=0x0000 type=integrated api=1.4.362 driver="Honeykrisp" info="Mesa 26.3.0-devel (git-9323171c12)" subgroup=32 storage16=true int16=true max-storage-range=2147483647 max-buffer-size=137438953472 fence-timeout=60000ms +MILESTONE project elapsed_ms=61 case="m1 n8 k64" worst=0.7485 of bound at row 0 col 1, max_abs=2.522643655538559e-3, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE project elapsed_ms=65 case="m2 n257 k128" worst=0.8880 of bound at row 0 col 131, max_abs=7.288329303264618e-3, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE project elapsed_ms=68 case="m2 n256 k128 offsets" worst=0.9205 of bound at row 0 col 35, max_abs=1.5103191137313843e-2, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE project elapsed_ms=70 case="m3 n8 k64 cancel" worst=0.6477 of bound at row 0 col 4, max_abs=4.5108795166015625e-4, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE project elapsed_ms=74 case="m1 n16 k2048" worst=0.2208 of bound at row 0 col 12, max_abs=1.1380931362509727e-2, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE project elapsed_ms=83 case="m2 n16 k6144 offsets" worst=0.0988 of bound at row 1 col 4, max_abs=2.9790184809826314e-2, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE project elapsed_ms=884 case="m32 n2048 k6144" worst=0.3003 of bound at row 16 col 745, max_abs=1.2159601086750627e-1, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +MILESTONE project elapsed_ms=5856 case="m1 n140000 k2048 chunked on a 128 MiB range" worst=0.6758 of bound at row 0 col 86073, max_abs=6.206762744113803e-2, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused +RESULT stage1 elapsed_ms=5875 on "Apple M1 (G13G B1)": passed=17 total=17 +exit code 0 diff --git a/docs/receipts/vulkan-qmm-vec4-m1/kernel-contract-m1-13-inch.txt b/docs/receipts/vulkan-qmm-vec4-m1/kernel-contract-m1-13-inch.txt new file mode 100644 index 00000000000..e394bda61da --- /dev/null +++ b/docs/receipts/vulkan-qmm-vec4-m1/kernel-contract-m1-13-inch.txt @@ -0,0 +1,35 @@ +LIMITS deviceName="Apple M1 (G13G B1)" +LIMITS driverName="Honeykrisp" driverInfo="Mesa 26.3.0-devel (git-9323171c12)" subgroupSize=32 +LIMITS maxComputeWorkGroupInvocations=1024 maxComputeWorkGroupSize=1024,1024,64 maxComputeSharedMemorySize=32768 +LIMITS maxStorageBufferRange=2147483647 minStorageBufferOffsetAlignment=16 +PASS limits qmm_bf16: local=256x1x1 invocations=256 shared=0 +PASS limits qmm_q4_vec4: local=256x1x1 invocations=256 shared=0 +PASS limits control:qmm_fmax_nonfinite: local=256x1x1 invocations=256 shared=0 +PASS limits control:atomic_sum: local=256x1x1 invocations=256 shared=0 +PASS limits control:scribble_input: local=256x1x1 invocations=256 shared=0 +PASS limits control:drop_tail: local=256x1x1 invocations=256 shared=0 +PASS limits control:fp16_accum: local=256x1x1 invocations=256 shared=0 +PASS limits control:over_limit: local=2048x1x1 invocations=2048 shared=0 +PASS nonfinite vulkan qmm_bf16 affine q4 g64 transposed: 64 non-finite and 16 finite outputs match the float64 classes +PASS repeat vulkan qmm_bf16 affine q4 g64 transposed: five runs byte-identical (64 bytes), inputs and guards untouched, no overlap +PASS shapes vulkan qmm_bf16 affine q4 g64 transposed: 15 shapes within the float64 bound (worst 0.990 of the bound), 0 edges refused by the kernel +PASS sums vulkan qmm_bf16 affine q4 g64 transposed: exact sums at k 6144, where fp16 sums would miss by 5268x and bf16 by 256x +PASS nonfinite vulkan qmm_q4_vec4 affine q4 g64 transposed: 64 non-finite and 16 finite outputs match the float64 classes +PASS repeat vulkan qmm_q4_vec4 affine q4 g64 transposed: five runs byte-identical (64 bytes), inputs and guards untouched, no overlap +PASS shapes vulkan qmm_q4_vec4 affine q4 g64 transposed: 15 shapes within the float64 bound (worst 0.990 of the bound), 0 edges refused by the kernel +PASS sums vulkan qmm_q4_vec4 affine q4 g64 transposed: exact sums at k 6144, where fp16 sums would miss by 4902x and bf16 by 256x +FAIL nonfinite control: qmm fmax drops NaN: 1 of 5 outputs differ in class; first at 0: finite where float64 gives nan +CONTROL control: qmm fmax drops NaN: red at nonfinite, as planted +FAIL repeat control: atomic add row sum: 4 of 4 repeats differ from the first run's bytes, first at byte 0 +CONTROL control: atomic add row sum: red at repeat, as planted +FAIL repeat control: row sum scribbles input: 4 of 4 repeats differ from the first run's bytes, first at byte 0 +CONTROL control: row sum scribbles input: red at repeat, as planted +FAIL shapes control: row sum drops K tail: at 4x1x4097 output 0 is 8.023634e1 where float64 gives 8.094338e1: 224.57 of its bound +CONTROL control: row sum drops K tail: red at shapes, as planted +FAIL sums control: dot accumulates in fp16: infx over its exact-sum bound at k 4096: not fp32 sums (fp16 would miss 33554432x, bf16 16777216x) +CONTROL control: dot accumulates in fp16: red at sums, as planted +FAIL repeat control: qmm writes its weights: a run changed its inputs +CONTROL control: qmm writes its weights: red at repeat, as planted +FAIL repeat control: qmm writes its biases: a run changed its inputs +CONTROL control: qmm writes its biases: red at repeat, as planted +SUMMARY Apple M1 (G13G B1) 2 kernels: 2 pass gates 1-4, 0 fail, 0 known findings; 8 controls: 8 caught, 0 missed; 0 limit failures; gates 24786 ms diff --git a/docs/receipts/vulkan-qmm-vec4-m1/raw-files-sha256.txt b/docs/receipts/vulkan-qmm-vec4-m1/raw-files-sha256.txt new file mode 100644 index 00000000000..0d2ef247838 --- /dev/null +++ b/docs/receipts/vulkan-qmm-vec4-m1/raw-files-sha256.txt @@ -0,0 +1,85 @@ +00de9332b242f2083ed2512031b0d44ea2f321f1f6238f24c79855be8c95bc38 pair-unicode24-first/out/unicode-24--p2-free/result.json +0343c456bd1ae948a266a571ad709cda3b91234565ee8310176c2ac24c88bd36 pair-unicode24-replicate/log--pr1-new.err +04ebc43b28980dd65de8614ba8100a0622ad69c2e6651f0587762696f07cec13 pair-unicode24-first/log--pr4-old.err +081bc7e682e750f211a27152f4be7b681f117df3b672d654a0f11b5a1a4f39d4 pair-unicode24-replicate/log--pr5-new.err +0e891794913e6d510faf56bfd97e09eca7a741b7ac5007cfee1eb3e65850768f pair-unicode24-first/log--pr0-old.err +1457e826c0dad08b914159a5cee7540503e3a9126f6e4df547d4b74ae21e80e3 id-m1/probe.err +1a92ee3b69c1b34617f8df742422ffad4d78b41a8ea7d72fb99cde647d04393a pair-document1024/log--pr5-new.err +20fee10364d35387eb693f7fa56f04eb5a762d4f37bee2549614926485c8c0ce pair-document1024/out/document-1024--p5-free/result.json +2635fb9aeecdd509080281d9b6561b02aec4c4f161088a7da47ed7d0d5652c1f pair-document256/out/document-256--p5-free/result.json +372f70b6e59dafa32bfc79b7ecea1448cf34df31eddb7358c9a05184593d25cc pair-document256/log--pr3-old.err +3798affa950cf7b748d3e41ef0d8a4f4229db922f395cf565854c409653f1429 pair-document256/out/document-256--p3-free/result.json +428d09c66811b0160b4ae47429ecb21b7ddbd977ecaed6579ef28a9ba78f7155 pair-document1024/log--pr4-old.err +43d6902b39d5a591a1e98caf6eb8be7e0cec511a5ffc606780a02638f3673488 pair-unicode24-first/log--pr3-old.err +46e0114ae80a12f8f92256607df6e4cec990e3062407d3cadfd0e70340ff0c44 pair-unicode24-replicate/log--pr3-old.err +4ef58a5c82708475b000c19cdf680032a2b42fa333a4f3b66626f142f9fc4845 pair-unicode24-replicate/out/unicode-24--p1-free/result.json +58d196dac92aeed0955dc69a6d2e0901ce316a29a9cb6abb6cde38887e0104e6 pair-unicode24-first/log--pr1-new.err +5a979292886be194f6df2b24dae2beb44de837ebeb9f7aedb2159745e5e8f573 pair-unicode24-replicate/log--pr2-new.err +5e0048a8fc2e99792fe382523a3fd8653d53eaf4c053e7db4a07f3ff34fc85ed pair-document1024/out/document-1024--p0-free/result.json +61c7e214a685cfd385f4dc4f64ef7a3de1e0e4e61143b0deec7665cf5be74bc0 pair-document256/out/document-256--p0-free/result.json +6374582883f99c2c7657480fa6430677efddd7c6be1750c0086e3c64a579b738 pair-document256/log--pr1-new.err +6645cdbcb504957c06ab23c907e1baa4daed264fa0d99ce500458ed878d9b906 pair-document1024/result.json +6f7d5ef63e38b965cd987c10d898a3260680b5196b9ccf19adfc5f63a9657f74 pair-document1024/out/document-1024--p2-free/result.json +753684fd21cfed17592e03ef4494635779d32c445c95d278393d26381821bbc8 pair-document256/log--pr0-old.err +781e735b3e7dbef578f8d4fb9328769a7b88a29e8f3310ee7f33efd1bed171a9 pair-unicode24-replicate/timing.tsv +7962c64466d2412c79a5d6126c19d6bd0c1313bbc54c7aeb218dea7a7a9c315d pair-unicode24-first/log--pr2-new.err +7a9f928f4d55c8d31823c52f94d9d33fa08a1827dc3f4194b40a4ac03df57fea pair-document256/result.json +7b03f71a1186c57f18723a8aded5043fbf75ab03635a7c00083d06d42767788e pair-unicode24-replicate/result.json +7d3b924ae929a275ac7b626082d81dbaf6138354a343960b36138a62e8a398d6 pair-document256/log--pr4-old.err +7eec78d59444ffeee6b88d1f51f569770060d4c7000056b62a0e1833695a0b25 pair-document1024/out/document-1024--p3-free/result.json +7f5755405e47b6b4d905cf38e02d44249e144292f9c6990cb21b38b0ed097c7d pair-unicode24-first/timing.tsv +84eec989b51b9c1e9c1f9f48d519d1058bb0c3c28c0a9979eab8629ed7d865b8 pair-document1024/log--pr2-new.err +98eef4ab838f9b00de1010952aa4afc32a28b1955090f89a8db3903aa520dc6d pair-unicode24-first/out/unicode-24--p0-free/result.json +9a271f2a916b0b6ee6cecb2426f0b3206ef074578be55d9bc94f6f3fe3ab86aa id-m1/probe.rc +9d42bcad73c381104821dc953e5c30ab334997157496ddc445f9d7d8c694e440 pair-document256/out/document-256--p4-free/result.json +9f666589be57881ed1f8fe50a187e312ca63521a848f9285b48563846daa5828 pair-unicode24-replicate/out/unicode-24--p5-free/result.json +a3ea80e2fb893686b8bd125cef9a48756a61cbbae5cf7f49758e121fa7fa5b03 pair-document256/out/document-256--p1-free/result.json +ab75e0fb627d64d9b378907b8dfb01c7a64a8b50dcc816cd977c5c0e1e32204e pair-document256/log--pr2-new.err +b0d6dfa3c9f115b4c64cd63774c3b26ce76fe4647bca8afeb4faae2a7e290d84 pair-unicode24-replicate/out/unicode-24--p4-free/result.json +b66c1c1b37ee4ee4f1e8731db46700a5243cd8587f5d9fa5883da28774c03f8b pair-unicode24-first/out/unicode-24--p5-free/result.json +b998ffacb4d9b67f125a71ec05c89b3cbb0af0ede48e34004c95c56ebc10bc3b pair-unicode24-replicate/out/unicode-24--p0-free/result.json +be394002cdd07e9efc36372c777ed5751ae776b3dee72014bee52d097ba26f0f pair-document256/out/document-256--p2-free/result.json +c15864ac328c11383275d576a95ab0c3d0eb834edf67725002e7a7f7dee7bc9e pair-unicode24-first/log--pr5-new.err +c2317714a5c735940e369f68ad85d2f269d948cc253cbc4560619f2000fe60b0 pair-unicode24-replicate/log--pr4-old.err +c2b5d61fd967ced7b154d781dec6606bfb307ad810d3a12c05c33adc3b41c9fe pair-unicode24-first/out/unicode-24--p3-free/result.json +c6b6792326226e2f3d40954f522c7410366ffdd127c67f6d89c156be466fc402 id-m1/SHA256SUMS +d0f3ae6be8a26fd84d15365276e181feacceaef62e13b1d3a4a6d598db5868b6 pair-document1024/log--pr3-old.err +d14e129b5ad3278076f750a6ef058be641e957aa5de5e2e3228e0986640709be pair-unicode24-replicate/log--pr0-old.err +d2a0749aad1830e165a850ecfcfddcaf35e80679696f8c38582f3cda6541647a pair-document1024/out/document-1024--p1-free/result.json +d584318cd52d60af1e0a1006868bb31bb6ff07c8fb712f96a2f9af82e5c402e1 pair-unicode24-replicate/out/unicode-24--p2-free/result.json +dbc6895cb8fd134a072d0a23e4547505d10d780262a035eaa1deffb9c4830ed8 pair-unicode24-first/out/unicode-24--p4-free/result.json +deea0f6bcaaa540b1e79f4f1baf69d30be4642b6199bb0d54db4d03491b876f2 pair-document1024/out/document-1024--p4-free/result.json +df884bf9565b1d98c7e67907cac4c09ddb8e92e0e2765c001b95479da884241c pair-document256/log--pr5-new.err +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 id-m1/probe.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-document1024/log--pr0-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-document1024/log--pr1-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-document1024/log--pr2-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-document1024/log--pr3-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-document1024/log--pr4-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-document1024/log--pr5-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-document256/log--pr0-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-document256/log--pr1-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-document256/log--pr2-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-document256/log--pr3-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-document256/log--pr4-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-document256/log--pr5-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-unicode24-first/log--pr0-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-unicode24-first/log--pr1-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-unicode24-first/log--pr2-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-unicode24-first/log--pr3-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-unicode24-first/log--pr4-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-unicode24-first/log--pr5-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-unicode24-replicate/log--pr0-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-unicode24-replicate/log--pr1-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-unicode24-replicate/log--pr2-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-unicode24-replicate/log--pr3-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-unicode24-replicate/log--pr4-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-unicode24-replicate/log--pr5-new.out +e3dd9a01835e1bf1e50531944b0fd3c67031186d717c0c0b5734ad039b55b948 pair-unicode24-first/out/unicode-24--p1-free/result.json +e69d2b2eb1fc902b199f3153a974fe12c6ad69db44f56087e6ba7dcda1d6eebc pair-unicode24-replicate/out/unicode-24--p3-free/result.json +e712f71034441f07a206e7b1cb295647eea4a18e4dff2188bf344fb6a4f8e328 pair-document1024/timing.tsv +ea56b4bcb7e3d8c3e79ee461914d2df1f0379bb791816edc17eb21a7352985d3 pair-document1024/log--pr1-new.err +ecebe13ee377d2554ff0bac45c6df9d02c0d5c54c2ae06f3b4d30cb31fb27197 pair-unicode24-first/result.json +f07cfd8bd585cc86e3b8ba784204250047b1e982604994b3134ccdea35c0b678 id-m1/env.txt +f25780901c2dec1f8f4b3f7511769b094ab384dfe6334284359dc9755a3f1d42 pair-document1024/log--pr0-old.err +fa0053d4243e8e9e22eb7afc05ecfdd4d8c7682f5a0047edde16bd4822409f73 pair-document256/timing.tsv diff --git a/docs/receipts/vulkan-qmm-vec4-m1/summary.json b/docs/receipts/vulkan-qmm-vec4-m1/summary.json new file mode 100644 index 00000000000..ab027262556 --- /dev/null +++ b/docs/receipts/vulkan-qmm-vec4-m1/summary.json @@ -0,0 +1,569 @@ +{ + "change": { + "base_commit": "b3ab9ba217a80c63fcc48f7643ddfe339a8f4bdc", + "base_is": "head of #1544", + "candidate_commit": "e742e02320bcddb3969d80b04f5cd1a11591625b", + "candidate_adds": "qmm_q4_vec4, the q4 projection with 16-byte weight and activation loads" + }, + "setup": { + "chip": "Apple M1 (G13G B1)", + "gpu_driver": "Honeykrisp, Mesa 26.3.0-devel (git-9323171c12)", + "os": "Linux", + "kernel_release": "7.1.12-2-20261013.18-sep-ARCH", + "checkpoint": "Qwen3.5-2B-MLX-4bit, revision 93760be4", + "quantization": "4-bit, group 64, bf16 activations", + "settings": "free-running greedy, 64 steps, two runs per process, vendor id 0x10005, one model process at a time, fresh process per run", + "metric": "run 2 duration: prefill of the prompt plus 63 decode steps plus the per-step logits readback; from the probe's own stderr milestones, 1 ms resolution", + "order": "old,new,new,old,old,new (three pairs, both orders)", + "decision_rule": "median run-2 gain of at least 3 percent and every pair positive" + }, + "identity_probe": { + "exit_code": 0, + "stderr_lines": [ + "MILESTONE device-found elapsed_ms=25 name=\"Apple M1 (G13G B1)\" vendor=0x10005 device=0x0000 type=integrated api=1.4.362 driver=\"Honeykrisp\" info=\"Mesa 26.3.0-devel (git-9323171c12)\" subgroup=32 storage16=true int16=true max-storage-range=2147483647 max-buffer-size=137438953472 fence-timeout=60000ms", + "MILESTONE project elapsed_ms=61 case=\"m1 n8 k64\" worst=0.7485 of bound at row 0 col 1, max_abs=2.522643655538559e-3, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused", + "MILESTONE project elapsed_ms=65 case=\"m2 n257 k128\" worst=0.8880 of bound at row 0 col 131, max_abs=7.288329303264618e-3, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused", + "MILESTONE project elapsed_ms=68 case=\"m2 n256 k128 offsets\" worst=0.9205 of bound at row 0 col 35, max_abs=1.5103191137313843e-2, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused", + "MILESTONE project elapsed_ms=70 case=\"m3 n8 k64 cancel\" worst=0.6477 of bound at row 0 col 4, max_abs=4.5108795166015625e-4, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused", + "MILESTONE project elapsed_ms=74 case=\"m1 n16 k2048\" worst=0.2208 of bound at row 0 col 12, max_abs=1.1380931362509727e-2, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused", + "MILESTONE project elapsed_ms=83 case=\"m2 n16 k6144 offsets\" worst=0.0988 of bound at row 1 col 4, max_abs=2.9790184809826314e-2, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused", + "MILESTONE project elapsed_ms=884 case=\"m32 n2048 k6144\" worst=0.3003 of bound at row 16 col 745, max_abs=1.2159601086750627e-1, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused", + "MILESTONE project elapsed_ms=5856 case=\"m1 n140000 k2048 chunked on a 128 MiB range\" worst=0.6758 of bound at row 0 col 86073, max_abs=6.206762744113803e-2, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused", + "RESULT stage1 elapsed_ms=5875 on \"Apple M1 (G13G B1)\": passed=17 total=17" + ] + }, + "cases": [ + { + "case": "unicode-24, first pair", + "status": "PASS_WITH_NOT_VERIFIED", + "kernel_release": "7.1.12-2-20261013.18-sep-ARCH", + "start_utc": "2026-11-03T02:31:39+00:00", + "end_utc": "2026-11-03T02:34:14+00:00", + "elfs": { + "old": { + "sha256": "ba69208ac5442e81fd682a37328791ef713bdc78ba7b2212ce29eecbb13f16c1", + "bytes": 5766584 + }, + "new": { + "sha256": "683fcf46a89c3d99d6db94c3f9fda614a9340d7193dfdc301451225f5db20ed8", + "bytes": 5778952 + } + }, + "logits_sha256_all_six_processes": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f", + "processes": [ + { + "label": "pr0-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 402, + "run1_ms": 15478, + "run2_ms": 15516, + "peak_rss_bytes": 86638592, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + }, + { + "label": "pr1-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 417, + "run1_ms": 9398, + "run2_ms": 9155, + "peak_rss_bytes": 87539712, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + }, + { + "label": "pr2-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 412, + "run1_ms": 9287, + "run2_ms": 9359, + "peak_rss_bytes": 87261184, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + }, + { + "label": "pr3-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 393, + "run1_ms": 15375, + "run2_ms": 15611, + "peak_rss_bytes": 86671360, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + }, + { + "label": "pr4-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 407, + "run1_ms": 15562, + "run2_ms": 15612, + "peak_rss_bytes": 86425600, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + }, + { + "label": "pr5-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 415, + "run1_ms": 9280, + "run2_ms": 9353, + "peak_rss_bytes": 86589440, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + } + ], + "pairs": [ + { + "pair": 1, + "order": [ + "old", + "new" + ], + "run2_old_ms": 15516, + "run2_new_ms": 9155, + "run2_gain": 0.409964, + "run1_old_ms": 15478, + "run1_new_ms": 9398, + "run1_gain": 0.392816 + }, + { + "pair": 2, + "order": [ + "new", + "old" + ], + "run2_old_ms": 15611, + "run2_new_ms": 9359, + "run2_gain": 0.400487, + "run1_old_ms": 15375, + "run1_new_ms": 9287, + "run1_gain": 0.395967 + }, + { + "pair": 3, + "order": [ + "old", + "new" + ], + "run2_old_ms": 15612, + "run2_new_ms": 9353, + "run2_gain": 0.40091, + "run1_old_ms": 15562, + "run1_new_ms": 9280, + "run1_gain": 0.403676 + } + ], + "run2_gain_median": 0.40091, + "run2_gain_min": 0.400487, + "run2_gain_max": 0.409964, + "run1_gain_median": 0.395967, + "speedup_rule": "MET", + "raw_result_json_sha256": "ecebe13ee377d2554ff0bac45c6df9d02c0d5c54c2ae06f3b4d30cb31fb27197" + }, + { + "case": "unicode-24, replicate", + "status": "PASS_WITH_NOT_VERIFIED", + "kernel_release": "7.1.12-2-20261013.18-sep-ARCH", + "start_utc": "2026-11-03T03:41:11+00:00", + "end_utc": "2026-11-03T03:43:43+00:00", + "elfs": { + "old": { + "sha256": "ba69208ac5442e81fd682a37328791ef713bdc78ba7b2212ce29eecbb13f16c1", + "bytes": 5766584 + }, + "new": { + "sha256": "683fcf46a89c3d99d6db94c3f9fda614a9340d7193dfdc301451225f5db20ed8", + "bytes": 5778952 + } + }, + "logits_sha256_all_six_processes": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f", + "processes": [ + { + "label": "pr0-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 402, + "run1_ms": 15465, + "run2_ms": 15386, + "peak_rss_bytes": 87506944, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + }, + { + "label": "pr1-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 405, + "run1_ms": 9180, + "run2_ms": 9036, + "peak_rss_bytes": 87588864, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + }, + { + "label": "pr2-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 403, + "run1_ms": 9169, + "run2_ms": 9188, + "peak_rss_bytes": 87244800, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + }, + { + "label": "pr3-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 400, + "run1_ms": 15530, + "run2_ms": 15264, + "peak_rss_bytes": 86605824, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + }, + { + "label": "pr4-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 402, + "run1_ms": 15442, + "run2_ms": 15220, + "peak_rss_bytes": 87195648, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + }, + { + "label": "pr5-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 411, + "run1_ms": 8768, + "run2_ms": 8980, + "peak_rss_bytes": 87588864, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + } + ], + "pairs": [ + { + "pair": 1, + "order": [ + "old", + "new" + ], + "run2_old_ms": 15386, + "run2_new_ms": 9036, + "run2_gain": 0.412713, + "run1_old_ms": 15465, + "run1_new_ms": 9180, + "run1_gain": 0.406402 + }, + { + "pair": 2, + "order": [ + "new", + "old" + ], + "run2_old_ms": 15264, + "run2_new_ms": 9188, + "run2_gain": 0.398061, + "run1_old_ms": 15530, + "run1_new_ms": 9169, + "run1_gain": 0.409594 + }, + { + "pair": 3, + "order": [ + "old", + "new" + ], + "run2_old_ms": 15220, + "run2_new_ms": 8980, + "run2_gain": 0.409987, + "run1_old_ms": 15442, + "run1_new_ms": 8768, + "run1_gain": 0.432198 + } + ], + "run2_gain_median": 0.409987, + "run2_gain_min": 0.398061, + "run2_gain_max": 0.412713, + "run1_gain_median": 0.409594, + "speedup_rule": "MET", + "raw_result_json_sha256": "7b03f71a1186c57f18723a8aded5043fbf75ab03635a7c00083d06d42767788e" + }, + { + "case": "document-256", + "status": "PASS_WITH_NOT_VERIFIED", + "kernel_release": "7.1.12-2-20261013.18-sep-ARCH", + "start_utc": "2026-11-03T02:45:29+00:00", + "end_utc": "2026-11-03T02:51:41+00:00", + "elfs": { + "old": { + "sha256": "ba69208ac5442e81fd682a37328791ef713bdc78ba7b2212ce29eecbb13f16c1", + "bytes": 5766584 + }, + "new": { + "sha256": "683fcf46a89c3d99d6db94c3f9fda614a9340d7193dfdc301451225f5db20ed8", + "bytes": 5778952 + } + }, + "logits_sha256_all_six_processes": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb", + "processes": [ + { + "label": "pr0-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 416, + "run1_ms": 40306, + "run2_ms": 40543, + "peak_rss_bytes": 87392256, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + }, + { + "label": "pr1-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 410, + "run1_ms": 20917, + "run2_ms": 20865, + "peak_rss_bytes": 119373824, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + }, + { + "label": "pr2-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 395, + "run1_ms": 20908, + "run2_ms": 20596, + "peak_rss_bytes": 87572480, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + }, + { + "label": "pr3-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 399, + "run1_ms": 40129, + "run2_ms": 39985, + "peak_rss_bytes": 119308288, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + }, + { + "label": "pr4-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 400, + "run1_ms": 40019, + "run2_ms": 40011, + "peak_rss_bytes": 87408640, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + }, + { + "label": "pr5-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 399, + "run1_ms": 20678, + "run2_ms": 20714, + "peak_rss_bytes": 86261760, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + } + ], + "pairs": [ + { + "pair": 1, + "order": [ + "old", + "new" + ], + "run2_old_ms": 40543, + "run2_new_ms": 20865, + "run2_gain": 0.485361, + "run1_old_ms": 40306, + "run1_new_ms": 20917, + "run1_gain": 0.481045 + }, + { + "pair": 2, + "order": [ + "new", + "old" + ], + "run2_old_ms": 39985, + "run2_new_ms": 20596, + "run2_gain": 0.484907, + "run1_old_ms": 40129, + "run1_new_ms": 20908, + "run1_gain": 0.47898 + }, + { + "pair": 3, + "order": [ + "old", + "new" + ], + "run2_old_ms": 40011, + "run2_new_ms": 20714, + "run2_gain": 0.482292, + "run1_old_ms": 40019, + "run1_new_ms": 20678, + "run1_gain": 0.483295 + } + ], + "run2_gain_median": 0.484907, + "run2_gain_min": 0.482292, + "run2_gain_max": 0.485361, + "run1_gain_median": 0.481045, + "speedup_rule": "MET", + "raw_result_json_sha256": "7a9f928f4d55c8d31823c52f94d9d33fa08a1827dc3f4194b40a4ac03df57fea" + }, + { + "case": "document-1024", + "status": "PASS_WITH_NOT_VERIFIED", + "kernel_release": "7.1.12-2-20261013.18-sep-ARCH", + "start_utc": "2026-11-03T03:16:36+00:00", + "end_utc": "2026-11-03T03:34:56+00:00", + "elfs": { + "old": { + "sha256": "ba69208ac5442e81fd682a37328791ef713bdc78ba7b2212ce29eecbb13f16c1", + "bytes": 5766584 + }, + "new": { + "sha256": "683fcf46a89c3d99d6db94c3f9fda614a9340d7193dfdc301451225f5db20ed8", + "bytes": 5778952 + } + }, + "logits_sha256_all_six_processes": "7f9d8e65e90eb3751edc41aed99c9ac5d37ce3f7d99dae1e626497632661a820", + "processes": [ + { + "label": "pr0-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 404, + "run1_ms": 122612, + "run2_ms": 122756, + "peak_rss_bytes": 119308288, + "logits_sha256": "7f9d8e65e90eb3751edc41aed99c9ac5d37ce3f7d99dae1e626497632661a820" + }, + { + "label": "pr1-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 406, + "run1_ms": 59575, + "run2_ms": 59497, + "peak_rss_bytes": 86409216, + "logits_sha256": "7f9d8e65e90eb3751edc41aed99c9ac5d37ce3f7d99dae1e626497632661a820" + }, + { + "label": "pr2-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 408, + "run1_ms": 59572, + "run2_ms": 59529, + "peak_rss_bytes": 86671360, + "logits_sha256": "7f9d8e65e90eb3751edc41aed99c9ac5d37ce3f7d99dae1e626497632661a820" + }, + { + "label": "pr3-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 404, + "run1_ms": 122713, + "run2_ms": 122768, + "peak_rss_bytes": 86622208, + "logits_sha256": "7f9d8e65e90eb3751edc41aed99c9ac5d37ce3f7d99dae1e626497632661a820" + }, + { + "label": "pr4-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 420, + "run1_ms": 122805, + "run2_ms": 122704, + "peak_rss_bytes": 87523328, + "logits_sha256": "7f9d8e65e90eb3751edc41aed99c9ac5d37ce3f7d99dae1e626497632661a820" + }, + { + "label": "pr5-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 407, + "run1_ms": 59706, + "run2_ms": 59498, + "peak_rss_bytes": 86540288, + "logits_sha256": "7f9d8e65e90eb3751edc41aed99c9ac5d37ce3f7d99dae1e626497632661a820" + } + ], + "pairs": [ + { + "pair": 1, + "order": [ + "old", + "new" + ], + "run2_old_ms": 122756, + "run2_new_ms": 59497, + "run2_gain": 0.515323, + "run1_old_ms": 122612, + "run1_new_ms": 59575, + "run1_gain": 0.514118 + }, + { + "pair": 2, + "order": [ + "new", + "old" + ], + "run2_old_ms": 122768, + "run2_new_ms": 59529, + "run2_gain": 0.51511, + "run1_old_ms": 122713, + "run1_new_ms": 59572, + "run1_gain": 0.514542 + }, + { + "pair": 3, + "order": [ + "old", + "new" + ], + "run2_old_ms": 122704, + "run2_new_ms": 59498, + "run2_gain": 0.51511, + "run1_old_ms": 122805, + "run1_new_ms": 59706, + "run1_gain": 0.513815 + } + ], + "run2_gain_median": 0.51511, + "run2_gain_min": 0.51511, + "run2_gain_max": 0.515323, + "run1_gain_median": 0.514118, + "speedup_rule": "MET", + "raw_result_json_sha256": "6645cdbcb504957c06ab23c907e1baa4daed264fa0d99ce500458ed878d9b906" + } + ] +} diff --git a/docs/receipts/vulkan-submit-batch-m1/PREREG.md b/docs/receipts/vulkan-submit-batch-m1/PREREG.md new file mode 100644 index 00000000000..ad70a6a7024 --- /dev/null +++ b/docs/receipts/vulkan-submit-batch-m1/PREREG.md @@ -0,0 +1,34 @@ +# One submission per batch of dispatches: pre-registration (written 2026-11-03 before any GPU run of this change) + +## Question + +Every dispatch of the Vulkan engine is its own submission with a fence wait. The profile of the tile build (unicode-24) shows 31,104 dispatches +in a 64-step run, about 19,000 of them small kernels at about 100 us each, and the tiny ones (copy, gather, dequant) at 89 to 111 us with a +minimum of 66 to 79 us: mostly submit and wait. Does recording many dispatches into one command buffer and submitting once per batch shorten a +run, with byte-identical logits? No kernel and no sum changes. + +## Arms (two ELFs of `tf-qwen35-model`, ReleaseSafe, aarch64) + +- old: the tile build, commit `6a158edc`, ELF `d4949bef...`. +- new: commit `0d73b0fb`, ELF `d8ee787b...`: dispatches are recorded into the open command buffer; it is submitted with 128 waiting dispatches, or when + an upload, download or fill needs the queue. Each pipeline has a ring of 256 descriptor sets. + +## Gates and metric + +As before: exit 0, pinned ids, one logits digest per case (unicode-24 `e5447d9a...`, document-256 `147a2d2e...`, document-1024 `7f9d8e65...`), run 2 duration +from the probe's milestones, three pairs in the order old,new,new,old,old,new. Also: the stage 1 to 5 probes pass on a software device (done), and a hang +or a validation error in any process fails the unit. + +## Decision rule + +Median gain of at least 3 percent and all pairs positive and all gates hold. Otherwise not proposed. + +## Forecast (a prediction) + +unicode-24 gains 15 to 35 percent: the floor is about 40 percent of the run, and I expect to keep half to four fifths of it (each dispatch still costs a +record, a barrier and the GPU's launch gap). document-256 and document-1024 gain less (a smaller share of their time is dispatch cost): 5 to 15 percent. +I give 15 percent to a gain under 3 percent on unicode-24, and 10 percent to a failure (a hang, a wrong digest or a crash) that needs a fix before any number counts. + +## Not claimed + +Other chips and models. The software-device pass is correctness evidence only. diff --git a/docs/receipts/vulkan-submit-batch-m1/README.md b/docs/receipts/vulkan-submit-batch-m1/README.md new file mode 100644 index 00000000000..e9bfaae1b83 --- /dev/null +++ b/docs/receipts/vulkan-submit-batch-m1/README.md @@ -0,0 +1,23 @@ +# Batched submission receipts, Apple M1 (13-inch) + +Qwen3.5-2B 4-bit, 64 greedy steps, two runs per process. Old build is #2015's head (the tile build, `b6d9a35b`), new build is #2016's head (adds batched submission, `d9b28417`). Both include the float32 fix. +Each case ran three old and three new processes, interleaved in both orders. Run 2 is the prefill plus 63 decode steps plus the per-step logits readback. Linux 7.3.0-2-20261110.1-lab-ARCH, Honeykrisp from the Omarchy Mesa build (library sha256 `d75f4cbc`). + +| Case | Old run 2 (median) | New run 2 (median) | Gain: median (pair range) | Logits digest, all six processes | +|---|---:|---:|---|---| +| unicode-24 | 8.0 s | 3.4 s | 57.8% (57.0 to 59.5) | `e946185d64ae` | +| document-256 | 15.1 s | 9.9 s | 34.1% (34.0 to 36.8) | `3188aaae7d11` | +| document-1024 | 38.7 s | 32.3 s | 16.6% (16.2 to 17.3) | `82c353787e89` | + +The gain depends on the environment. The first version of this receipt measured 68.2%, 48.8% and 30.5% on earlier builds with Linux 7.1.12-2-20261013.18-sep-ARCH and Mesa git-9323171c12 (digests `e5447d9a3ab5`, `147a2d2e0232`, `7f9d8e65e90e`). Running the earlier document-256 builds (old `151319df`, new `e84e79dd`) in the current environment gave 20.8 s and 16.6 s, a 20.5% gain, with the earlier digest `147a2d2e0232`. Every run time here is slower than in the first version, and these runs do not separate the cause. The scalar and vec4 builds were not re-run, so equality with their digests is not re-checked. + +Files: + +- `fixed-heads-summary.json` holds every process row (with peak resident memory), the pairs, the ELF hashes, the setup, and the earlier-builds run described above. It is generated from the raw run files. +- `summary.json` is the first version's summary (earlier builds, earlier environment). It stays as the record of those runs. +- `identity-probe.txt` holds the stage 1 probe run on the GPU with the earlier builds: 21 of 21, and the tiled image equal to the default image on every projection case. It was not re-run on the fixed heads. +- `kernel-contract-m1-13-inch.txt` holds the contract runner's output on the GPU with the earlier builds. It was not re-run on the fixed heads. +- `PREREG.md` was written before any GPU run. It holds the forecast, the gates and the decision rule. +- `raw-files-sha256.txt` lists the SHA-256 of every raw run file `summary.json` was generated from. + +Byte identity: each process writes the SHA-256 of its full logits stream. A case passes only if all six digests are equal and the ids match the pinned ids. diff --git a/docs/receipts/vulkan-submit-batch-m1/fixed-heads-summary.json b/docs/receipts/vulkan-submit-batch-m1/fixed-heads-summary.json new file mode 100644 index 00000000000..3b816c4a9f6 --- /dev/null +++ b/docs/receipts/vulkan-submit-batch-m1/fixed-heads-summary.json @@ -0,0 +1,322 @@ +{ + "change": "Vulkan batched submission (#2016) measured on the float32-fixed heads", + "setup": { + "chip": "Apple M1 (G13G B1)", + "kernel_release": "7.3.0-2-20261110.1-lab-ARCH", + "mesa": "Honeykrisp, Omarchy Mesa build, library sha256 d75f4cbc", + "checkpoint": "Qwen3.5-2B-MLX-4bit, revision 93760be4", + "settings": "free-running greedy, 64 steps, two runs per process, vendor id 0x10005, one model process at a time, fresh process per run", + "order": "old,new,new,old,old,new (three pairs, both orders)", + "decision_rule": "median run-2 gain of at least 3 percent and every pair positive" + }, + "cases": [ + { + "case": "unicode-24", + "status": "PASS_WITH_NOT_VERIFIED", + "start_utc": "2026-11-11T17:04:25+00:00", + "kernel_release": "7.3.0-2-20261110.1-lab-ARCH", + "elfs": { + "old": { + "sha256": "b6d9a35b97786770f9424b6e0089c2daead0625e0e21597f8abb96c6f6cba3cb", + "bytes": 6116824 + }, + "new": { + "sha256": "d9b2841744a03c80160d92959dd27ff40cb53110b0cd33f188ea24c9d729fa5d", + "bytes": 6109056 + } + }, + "logits_sha256_all_six_processes": "e946185d64aea5942265276e7b76f7d4c4560b543145e1139e9c2a09432ca00c", + "run2_ms_median": { + "old": 7994, + "new": 3371 + }, + "run2_gain_median": 0.5783087315486615, + "run2_gain_range": [ + 0.570097037793667, + 0.5948388631228446 + ], + "processes": [ + { + "label": "pr0-old", + "arm": "old", + "run1_ms": 8242, + "run2_ms": 7832, + "peak_rss_bytes": 1503100928, + "logits_sha256": "e946185d64aea5942265276e7b76f7d4c4560b543145e1139e9c2a09432ca00c" + }, + { + "label": "pr1-new", + "arm": "new", + "run1_ms": 3360, + "run2_ms": 3367, + "peak_rss_bytes": 1517191168, + "logits_sha256": "e946185d64aea5942265276e7b76f7d4c4560b543145e1139e9c2a09432ca00c" + }, + { + "label": "pr2-new", + "arm": "new", + "run1_ms": 3361, + "run2_ms": 3371, + "peak_rss_bytes": 1517191168, + "logits_sha256": "e946185d64aea5942265276e7b76f7d4c4560b543145e1139e9c2a09432ca00c" + }, + { + "label": "pr3-old", + "arm": "old", + "run1_ms": 7994, + "run2_ms": 7994, + "peak_rss_bytes": 1536081920, + "logits_sha256": "e946185d64aea5942265276e7b76f7d4c4560b543145e1139e9c2a09432ca00c" + }, + { + "label": "pr4-old", + "arm": "old", + "run1_ms": 8579, + "run2_ms": 8409, + "peak_rss_bytes": 1535164416, + "logits_sha256": "e946185d64aea5942265276e7b76f7d4c4560b543145e1139e9c2a09432ca00c" + }, + { + "label": "pr5-new", + "arm": "new", + "run1_ms": 3373, + "run2_ms": 3407, + "peak_rss_bytes": 1548992512, + "logits_sha256": "e946185d64aea5942265276e7b76f7d4c4560b543145e1139e9c2a09432ca00c" + } + ] + }, + { + "case": "document-256", + "status": "PASS_WITH_NOT_VERIFIED", + "start_utc": "2026-11-11T17:40:46+00:00", + "kernel_release": "7.3.0-2-20261110.1-lab-ARCH", + "elfs": { + "old": { + "sha256": "b6d9a35b97786770f9424b6e0089c2daead0625e0e21597f8abb96c6f6cba3cb", + "bytes": 6116824 + }, + "new": { + "sha256": "d9b2841744a03c80160d92959dd27ff40cb53110b0cd33f188ea24c9d729fa5d", + "bytes": 6109056 + } + }, + "logits_sha256_all_six_processes": "3188aaae7d1182c133a89fde103722d23f3d81b57df259498c9b778ebe2d0d7a", + "run2_ms_median": { + "old": 15070, + "new": 9906 + }, + "run2_gain_median": 0.34100891787568216, + "run2_gain_range": [ + 0.3396814863968149, + 0.367876970199732 + ], + "processes": [ + { + "label": "pr0-old", + "arm": "old", + "run1_ms": 15175, + "run2_ms": 15026, + "peak_rss_bytes": 1523728384, + "logits_sha256": "3188aaae7d1182c133a89fde103722d23f3d81b57df259498c9b778ebe2d0d7a" + }, + { + "label": "pr1-new", + "arm": "new", + "run1_ms": 9901, + "run2_ms": 9902, + "peak_rss_bytes": 1517191168, + "logits_sha256": "3188aaae7d1182c133a89fde103722d23f3d81b57df259498c9b778ebe2d0d7a" + }, + { + "label": "pr2-new", + "arm": "new", + "run1_ms": 9936, + "run2_ms": 9951, + "peak_rss_bytes": 1517174784, + "logits_sha256": "3188aaae7d1182c133a89fde103722d23f3d81b57df259498c9b778ebe2d0d7a" + }, + { + "label": "pr3-old", + "arm": "old", + "run1_ms": 14738, + "run2_ms": 15070, + "peak_rss_bytes": 1504526336, + "logits_sha256": "3188aaae7d1182c133a89fde103722d23f3d81b57df259498c9b778ebe2d0d7a" + }, + { + "label": "pr4-old", + "arm": "old", + "run1_ms": 15781, + "run2_ms": 15671, + "peak_rss_bytes": 1529544704, + "logits_sha256": "3188aaae7d1182c133a89fde103722d23f3d81b57df259498c9b778ebe2d0d7a" + }, + { + "label": "pr5-new", + "arm": "new", + "run1_ms": 9903, + "run2_ms": 9906, + "peak_rss_bytes": 1517142016, + "logits_sha256": "3188aaae7d1182c133a89fde103722d23f3d81b57df259498c9b778ebe2d0d7a" + } + ] + }, + { + "case": "document-1024", + "status": "PASS_WITH_NOT_VERIFIED", + "start_utc": "2026-11-11T18:20:53+00:00", + "kernel_release": "7.3.0-2-20261110.1-lab-ARCH", + "elfs": { + "old": { + "sha256": "b6d9a35b97786770f9424b6e0089c2daead0625e0e21597f8abb96c6f6cba3cb", + "bytes": 6116824 + }, + "new": { + "sha256": "d9b2841744a03c80160d92959dd27ff40cb53110b0cd33f188ea24c9d729fa5d", + "bytes": 6109056 + } + }, + "logits_sha256_all_six_processes": "82c353787e896b52c8def7f4cebbc087d42f71c4837505db8100df67b2235356", + "run2_ms_median": { + "old": 38745, + "new": 32299 + }, + "run2_gain_median": 0.1663698541747322, + "run2_gain_range": [ + 0.1620646701614159, + 0.17282750505231384 + ], + "processes": [ + { + "label": "pr0-old", + "arm": "old", + "run1_ms": 38501, + "run2_ms": 38745, + "peak_rss_bytes": 1516552192, + "logits_sha256": "82c353787e896b52c8def7f4cebbc087d42f71c4837505db8100df67b2235356" + }, + { + "label": "pr1-new", + "arm": "new", + "run1_ms": 32292, + "run2_ms": 32299, + "peak_rss_bytes": 1529741312, + "logits_sha256": "82c353787e896b52c8def7f4cebbc087d42f71c4837505db8100df67b2235356" + }, + { + "label": "pr2-new", + "arm": "new", + "run1_ms": 32293, + "run2_ms": 32335, + "peak_rss_bytes": 1529757696, + "logits_sha256": "82c353787e896b52c8def7f4cebbc087d42f71c4837505db8100df67b2235356" + }, + { + "label": "pr3-old", + "arm": "old", + "run1_ms": 39385, + "run2_ms": 39091, + "peak_rss_bytes": 1517207552, + "logits_sha256": "82c353787e896b52c8def7f4cebbc087d42f71c4837505db8100df67b2235356" + }, + { + "label": "pr4-old", + "arm": "old", + "run1_ms": 38865, + "run2_ms": 38534, + "peak_rss_bytes": 1517223936, + "logits_sha256": "82c353787e896b52c8def7f4cebbc087d42f71c4837505db8100df67b2235356" + }, + { + "label": "pr5-new", + "arm": "new", + "run1_ms": 32254, + "run2_ms": 32289, + "peak_rss_bytes": 1529462784, + "logits_sha256": "82c353787e896b52c8def7f4cebbc087d42f71c4837505db8100df67b2235356" + } + ] + } + ], + "isolation": { + "case": "document-256, earlier builds in the current environment", + "status": "PASS_WITH_NOT_VERIFIED", + "start_utc": "2026-11-11T18:08:43+00:00", + "kernel_release": "7.3.0-2-20261110.1-lab-ARCH", + "elfs": { + "old": { + "sha256": "151319df8f011fbb23c753a9398836e28215582c26ec8d7e6b4ad7aa0854c7b7", + "bytes": 5784808 + }, + "new": { + "sha256": "e84e79dd1fefed9889f833d92bed5c19d36033a189851b878232248dd7a6a7ea", + "bytes": 5765768 + } + }, + "logits_sha256_all_six_processes": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb", + "run2_ms_median": { + "old": 20838, + "new": 16567 + }, + "run2_gain_median": 0.20539399174584894, + "run2_gain_range": [ + 0.19968135953266064, + 0.21128302785051178 + ], + "processes": [ + { + "label": "pr0-old", + "arm": "old", + "run1_ms": 20810, + "run2_ms": 20713, + "peak_rss_bytes": 1212760064, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + }, + { + "label": "pr1-new", + "arm": "new", + "run1_ms": 16554, + "run2_ms": 16577, + "peak_rss_bytes": 1225080832, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + }, + { + "label": "pr2-new", + "arm": "new", + "run1_ms": 16556, + "run2_ms": 16558, + "peak_rss_bytes": 1224343552, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + }, + { + "label": "pr3-old", + "arm": "old", + "run1_ms": 20883, + "run2_ms": 20838, + "peak_rss_bytes": 1244938240, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + }, + { + "label": "pr4-old", + "arm": "old", + "run1_ms": 20919, + "run2_ms": 21005, + "peak_rss_bytes": 1222115328, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + }, + { + "label": "pr5-new", + "arm": "new", + "run1_ms": 16564, + "run2_ms": 16567, + "peak_rss_bytes": 1224294400, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + } + ] + }, + "isolation_note": "Same case, same host and kernel, with the earlier builds (old 151319df, new e84e79dd). The earlier README table was measured on kernel 7.1.12-2-20261013.18-sep-ARCH with Mesa git-9323171c12. Run times and gains differ between the environments, and these runs do not separate the cause.", + "excluded_attempts": [ + "A first document-1024 attempt failed the packet's own prompt-token pin (256 instead of 1024) after one process. It is not counted." + ] +} diff --git a/docs/receipts/vulkan-submit-batch-m1/identity-probe.txt b/docs/receipts/vulkan-submit-batch-m1/identity-probe.txt new file mode 100644 index 00000000000..0016db5202e --- /dev/null +++ b/docs/receipts/vulkan-submit-batch-m1/identity-probe.txt @@ -0,0 +1,14 @@ +RESULT stage1 elapsed_ms=7652 on "Apple M1 (G13G B1)": passed=21 total=21 +m1 n8 k64: identical +m2 n257 k128: identical +m2 n256 k128 offsets: identical +m3 n8 k64 cancel: identical +m1 n16 k2048: identical +m2 n16 k6144 offsets: identical +m32 n2048 k6144: identical +m1 n4096 k2048: identical +m1 n6144 k2048: identical +m1 n2048 k6144 offsets: identical +m32 n6144 k2048: identical +m1 n140000 k2048 chunked on a 128 MiB range: identical +exit code 0 diff --git a/docs/receipts/vulkan-submit-batch-m1/kernel-contract-m1-13-inch.txt b/docs/receipts/vulkan-submit-batch-m1/kernel-contract-m1-13-inch.txt new file mode 100644 index 00000000000..64bdebe817c --- /dev/null +++ b/docs/receipts/vulkan-submit-batch-m1/kernel-contract-m1-13-inch.txt @@ -0,0 +1,40 @@ +LIMITS deviceName="Apple M1 (G13G B1)" +LIMITS driverName="Honeykrisp" driverInfo="Mesa 26.3.0-devel (git-9323171c12)" subgroupSize=32 +LIMITS maxComputeWorkGroupInvocations=1024 maxComputeWorkGroupSize=1024,1024,64 maxComputeSharedMemorySize=32768 +LIMITS maxStorageBufferRange=2147483647 minStorageBufferOffsetAlignment=16 +PASS limits qmm_bf16: local=256x1x1 invocations=256 shared=0 +PASS limits qmm_q4_vec4: local=256x1x1 invocations=256 shared=0 +PASS limits qmm_q4_tile: local=256x1x1 invocations=256 shared=0 +PASS limits control:qmm_fmax_nonfinite: local=256x1x1 invocations=256 shared=0 +PASS limits control:atomic_sum: local=256x1x1 invocations=256 shared=0 +PASS limits control:scribble_input: local=256x1x1 invocations=256 shared=0 +PASS limits control:drop_tail: local=256x1x1 invocations=256 shared=0 +PASS limits control:fp16_accum: local=256x1x1 invocations=256 shared=0 +PASS limits control:over_limit: local=2048x1x1 invocations=2048 shared=0 +PASS nonfinite vulkan qmm_bf16 affine q4 g64 transposed: 64 non-finite and 16 finite outputs match the float64 classes +PASS repeat vulkan qmm_bf16 affine q4 g64 transposed: five runs byte-identical (64 bytes), inputs and guards untouched, no overlap +PASS shapes vulkan qmm_bf16 affine q4 g64 transposed: 15 shapes within the float64 bound (worst 0.990 of the bound), 0 edges refused by the kernel +PASS sums vulkan qmm_bf16 affine q4 g64 transposed: exact sums at k 6144, where fp16 sums would miss by 5268x and bf16 by 256x +PASS nonfinite vulkan qmm_q4_vec4 affine q4 g64 transposed: 64 non-finite and 16 finite outputs match the float64 classes +PASS repeat vulkan qmm_q4_vec4 affine q4 g64 transposed: five runs byte-identical (64 bytes), inputs and guards untouched, no overlap +PASS shapes vulkan qmm_q4_vec4 affine q4 g64 transposed: 15 shapes within the float64 bound (worst 0.990 of the bound), 0 edges refused by the kernel +PASS sums vulkan qmm_q4_vec4 affine q4 g64 transposed: exact sums at k 6144, where fp16 sums would miss by 4902x and bf16 by 256x +PASS nonfinite vulkan qmm_q4_tile affine q4 g64 transposed tiled 32: 64 non-finite and 16 finite outputs match the float64 classes +PASS repeat vulkan qmm_q4_tile affine q4 g64 transposed tiled 32: five runs byte-identical (64 bytes), inputs and guards untouched, no overlap +PASS shapes vulkan qmm_q4_tile affine q4 g64 transposed tiled 32: 15 shapes within the float64 bound (worst 0.987 of the bound), 0 edges refused by the kernel +PASS sums vulkan qmm_q4_tile affine q4 g64 transposed tiled 32: exact sums at k 6144, where fp16 sums would miss by 5756x and bf16 by 256x +FAIL nonfinite control: qmm fmax drops NaN: 1 of 5 outputs differ in class; first at 0: finite where float64 gives nan +CONTROL control: qmm fmax drops NaN: red at nonfinite, as planted +FAIL repeat control: atomic add row sum: 4 of 4 repeats differ from the first run's bytes, first at byte 0 +CONTROL control: atomic add row sum: red at repeat, as planted +FAIL repeat control: row sum scribbles input: 4 of 4 repeats differ from the first run's bytes, first at byte 0 +CONTROL control: row sum scribbles input: red at repeat, as planted +FAIL shapes control: row sum drops K tail: at 4x1x4097 output 0 is 8.023634e1 where float64 gives 8.094338e1: 224.57 of its bound +CONTROL control: row sum drops K tail: red at shapes, as planted +FAIL sums control: dot accumulates in fp16: infx over its exact-sum bound at k 4096: not fp32 sums (fp16 would miss 33554432x, bf16 16777216x) +CONTROL control: dot accumulates in fp16: red at sums, as planted +FAIL repeat control: qmm writes its weights: a run changed its inputs +CONTROL control: qmm writes its weights: red at repeat, as planted +FAIL repeat control: qmm writes its biases: a run changed its inputs +CONTROL control: qmm writes its biases: red at repeat, as planted +SUMMARY Apple M1 (G13G B1) 3 kernels: 3 pass gates 1-4, 0 fail, 0 known findings; 8 controls: 8 caught, 0 missed; 0 limit failures; gates 32345 ms diff --git a/docs/receipts/vulkan-submit-batch-m1/raw-files-sha256.txt b/docs/receipts/vulkan-submit-batch-m1/raw-files-sha256.txt new file mode 100644 index 00000000000..cabbf3789fe --- /dev/null +++ b/docs/receipts/vulkan-submit-batch-m1/raw-files-sha256.txt @@ -0,0 +1,70 @@ +037279260ffd2b5540aac5ed7e43db2af9c7e554ecc86e54a63e90369d0b9eff pair-document1024/out/document-1024--p4-free/result.json +0fd3cd673700e6bcd2a1ccbed9046a7b8f37803b285e1434ad04eabcd9aafc15 pair-unicode24/out/unicode-24--p1-free/result.json +155956133c9e2aceb250674d0a5c6ef722b2239b8c3a86bec70ea261e14ba3c6 pair-document1024/out/document-1024--p3-free/result.json +16137f0840cb06a715bd295520faea3958680924e950b1c5a2de90b91bbe06fa pair-unicode24/log--pr3-old.err +1d18e90f0d98f1694e928736bad522064f1d93191608717082031cd75f828203 pair-unicode24/log--pr0-old.err +1d3f971eeb69ada6e0e26f0829686d52c04e8233a98e6235f52763b986e699f8 pair-document256/log--pr4-old.err +1f99302956b7c994ba95bfc7fd5e335deb1d85ee3080f6a11732c8e753b68160 pair-document256/out/document-256--p0-free/result.json +27d1b1843f0ea180978c084e7717aee9349e0e996cbad406812bcf4cebb2ee7c pair-document1024/timing.tsv +2adc1e5b2fd3a2bd39eadd2e05e74c573413a24e576f81c852b4bd1a159760b4 pair-document256/out/document-256--p1-free/result.json +2cbabd9fdbe4291189c7ee99f8ab3c1302a1107000887c8f2c09f1f7777d53bd id-m1/env.txt +2cf91bc926a02ade4ec4a9c4a5b0bc231121fb2aa5d8716ad74be3eadfa6be67 pair-document256/out/document-256--p2-free/result.json +34983c20030757a2564aaebeac834c51000b75fa801292e664fc64a6931fb4df pair-document256/log--pr5-new.err +3714cd4fa99d4dc3e38e04dd069457dc8843037407d5a605c4eb292daabc83b0 pair-unicode24/out/unicode-24--p4-free/result.json +395b44200cb4fcfbf26a3222e763caa8dd34c35be4059bcc488e50218b62c2ff id-m1/run.err +4b8b192bb8ce18c1636b7797ab0daef39da9e9759c1960940498c54ee2bd8d57 pair-unicode24/log--pr2-new.err +52d467d803f456fa8d11412866033a6f40dc443101bf051618fff315086d0832 pair-unicode24/log--pr1-new.err +52f61b395aa76d16ca0aa5041c62cbf41161cb23c2cee8c60deb04b7c42e0add pair-unicode24/result.json +538d859027c43f300cc56393b0005a361a7e468d01b9f2a4c1ceecc8da5b8527 pair-document256/result.json +5d47d2e174771f6d0a41f51f3930b478e89e6cee7fc272d0955bb72cdaf9eebb pair-unicode24/out/unicode-24--p3-free/result.json +5fa86448b889034ddfb7d9f343f2d3d5d09c92e622ba0dfea17110dbe0458879 pair-document256/out/document-256--p5-free/result.json +6d03440b1ecc73c2c1f886e3096d827389e8cef1bdd11e9b917c448094f2b97a pair-document1024/log--pr1-new.err +70dd7bc7c3fec3f6a069cfa6c17214856b8f980645d66ad87fae37eb74819894 pair-document1024/log--pr4-old.err +74210119a6dc14db39d7ddfacb5e68382f0a9737e6797ee3bc8ae2d0be2bfce7 pair-document1024/log--pr3-old.err +7aa14a731447c2bc7efa69a43db3e75e0a6b1a25c4407087c198401375d6749f pair-document256/out/document-256--p4-free/result.json +80f19cef8e73966fb1d3ae4d1e27cc51e0cb86bd02cd272bbec169ea846ec199 contract-m1/SHA256SUMS +8e834cfbb6a275d31216bf7da1dbda59739e7ca2a986774ee59589ed64eabdc6 contract-m1/env.txt +9a271f2a916b0b6ee6cecb2426f0b3206ef074578be55d9bc94f6f3fe3ab86aa contract-m1/run.rc +9a271f2a916b0b6ee6cecb2426f0b3206ef074578be55d9bc94f6f3fe3ab86aa id-m1/run.rc +9e911bdc3ec2e0ae13acea07e6246f7d0a0baff1339eb5b3e9b5b44a0d19213d pair-document1024/out/document-1024--p2-free/result.json +9ff1107c61fd519cd2409387caa2e1f894027d6add831cde5656037f4aeb88bb pair-unicode24/log--pr5-new.err +a900a94ee456df4d3eb8e131400519e07c75eb3415d918589c98c72615776040 pair-document1024/log--pr5-new.err +ac6b5e9f8a19f78125e58151152eb9e76a2f081fdd3e03447268bdf566118ec1 pair-document1024/out/document-1024--p0-free/result.json +b2a8087f2dc71e9144d9437177e008bdb350b6a38b5f9b612a72d67aa8f6bda3 id-m1/SHA256SUMS +c21cd18f7fdc5fa0654f7e3d6bd3dc714833937fcd2180da98367e9b3c56c218 pair-document256/log--pr1-new.err +c4d82d4add20a4f908e7841f1ea1a940837ad13a3d4cd2a61f22e0d819a1966b pair-unicode24/out/unicode-24--p2-free/result.json +c55cd4f04a086296e6c702a612d7e90b0e2c45da1d532d32fc4f27d3dafabb64 pair-document256/timing.tsv +c750d2b47ce643bf4c63cf2d7e67c3e64367e734f5548075f792a03b215ce1aa pair-document1024/result.json +c84c8782b28451b010d52e771873d67de94021ad02dada05bbadb220b598c98e contract-m1/run.err +ccf7fea1b5800785199010e2243fe2664882ab06da3e90fcb883077fa945c16a pair-document1024/log--pr0-old.err +d39380dab8699f53f013df679301e5f27bfb74d54733a410e50e85bf39833628 pair-unicode24/out/unicode-24--p5-free/result.json +d4cad957d0eed8b5f4c703fb2638b279565fc5ce5ac03d8bea1851fcbc8c0d07 pair-unicode24/timing.tsv +dd082d3e9a433bc58e1a60743e8e8a3774be6f5588bbfdaa182d2104aa62afed pair-document256/out/document-256--p3-free/result.json +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 contract-m1/run.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 id-m1/run.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-document1024/log--pr0-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-document1024/log--pr1-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-document1024/log--pr2-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-document1024/log--pr3-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-document1024/log--pr4-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-document1024/log--pr5-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-document256/log--pr0-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-document256/log--pr1-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-document256/log--pr2-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-document256/log--pr3-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-document256/log--pr4-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-document256/log--pr5-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-unicode24/log--pr0-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-unicode24/log--pr1-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-unicode24/log--pr2-new.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-unicode24/log--pr3-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-unicode24/log--pr4-old.out +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 pair-unicode24/log--pr5-new.out +e6ae0d1f2ec67b6c8eff60c81d3906adb0f95cdd9e466e65253a2c0ab8292c3d pair-unicode24/out/unicode-24--p0-free/result.json +ecf8dfc820ff4bf9030f4e65b90eb98fb5d61d0df1ee3ceb0178a17cabf32fb7 pair-document256/log--pr3-old.err +efaa4939678aa40d6ad38ea4ba543c66ed447848102eb51607a2e461a3b0940c pair-document1024/log--pr2-new.err +f414799e85ff84ead7c7334ed4ab1dc28cb4f0f52e8355af928c2294bae2495c pair-document256/log--pr2-new.err +f51117a4c9950d4c5d705f03da05271857d7d64a32117bd64e522b7c42260efd pair-document256/log--pr0-old.err +f53d4968b749a7fd9832eda73dd2134722e21cf2ece41ec0547283161931796b pair-document1024/out/document-1024--p5-free/result.json +f7800b8bdfdfebe02f63bc73c693c5f8274f38343d259aea8082cc11a79402bf pair-document1024/out/document-1024--p1-free/result.json +fd2e821a527bfb516f2102abc4f9940c2b5575da3c859c2543ece39b1c57f8f0 pair-unicode24/log--pr4-old.err diff --git a/docs/receipts/vulkan-submit-batch-m1/summary.json b/docs/receipts/vulkan-submit-batch-m1/summary.json new file mode 100644 index 00000000000..c00408d06c0 --- /dev/null +++ b/docs/receipts/vulkan-submit-batch-m1/summary.json @@ -0,0 +1,443 @@ +{ + "change": { + "old_build_commit": "618ed699d0c62a7114afb758c2037974079a4559", + "old_build_elf_built_from_tree_of": "6a158edcf2a02050181b85c3b2eb6a69d9f1cf31", + "old_build_is": "the qmm_q4_tile build", + "new_build_commit": "5f8d4c3bcdd4d84b4977d1ea6c17d5ccf24cb707", + "new_build_elf_built_from_commit": "0d73b0fbdf06ced6643b89476342629703d5697f (the same change on the tile build before main was merged in; the two files of the change are identical)", + "new_build_adds": "Pipeline.run records into the open command buffer; Commands submits at 128 waiting dispatches or at the next upload, download or fill; each pipeline has a ring of 256 descriptor sets" + }, + "setup": { + "chip": "Apple M1 (G13G B1)", + "gpu_driver": "Honeykrisp, Mesa 26.3.0-devel (git-9323171c12)", + "os": "Linux", + "kernel_release": "7.1.12-2-20261013.18-sep-ARCH", + "checkpoint": "Qwen3.5-2B-MLX-4bit, revision 93760be4", + "quantization": "4-bit, group 64, bf16 activations", + "settings": "free-running greedy, 64 steps, two runs per process, vendor id 0x10005, one model process at a time, fresh process per run", + "metric": "run 2 duration: prefill of the prompt plus 63 decode steps plus the per-step logits readback; from the probe's own stderr milestones, 1 ms resolution", + "order": "old,new,new,old,old,new (three pairs, both orders)", + "decision_rule": "median run-2 gain of at least 3 percent and every pair positive" + }, + "identity_probe": { + "exit_code": 0, + "result": [ + "RESULT stage1 elapsed_ms=7652 on \"Apple M1 (G13G B1)\": passed=21 total=21" + ], + "tiled_image_vs_default_image": [ + "m1 n8 k64: identical", + "m2 n257 k128: identical", + "m2 n256 k128 offsets: identical", + "m3 n8 k64 cancel: identical", + "m1 n16 k2048: identical", + "m2 n16 k6144 offsets: identical", + "m32 n2048 k6144: identical", + "m1 n4096 k2048: identical", + "m1 n6144 k2048: identical", + "m1 n2048 k6144 offsets: identical", + "m32 n6144 k2048: identical", + "m1 n140000 k2048 chunked on a 128 MiB range: identical" + ] + }, + "cases": [ + { + "case": "unicode-24", + "status": "PASS_WITH_NOT_VERIFIED", + "kernel_release": "7.1.12-2-20261013.18-sep-ARCH", + "start_utc": "2026-11-03T09:02:25+00:00", + "end_utc": "2026-11-03T09:03:29+00:00", + "elfs": { + "old": { + "sha256": "d4949bef20e3b175365d26fab943cafe6f28befdd060185366ff2713368530a3", + "bytes": 5809824 + }, + "new": { + "sha256": "d8ee787b9aefd5f7461ebe3ed40d0dbbacd06f2891822cde476c1b054abfe85c", + "bytes": 5789144 + } + }, + "logits_sha256_all_six_processes": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f", + "processes": [ + { + "label": "pr0-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 706, + "run1_ms": 7277, + "run2_ms": 6977, + "peak_rss_bytes": 325550080, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + }, + { + "label": "pr1-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 711, + "run1_ms": 2263, + "run2_ms": 2277, + "peak_rss_bytes": 325533696, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + }, + { + "label": "pr2-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 759, + "run1_ms": 2236, + "run2_ms": 2237, + "peak_rss_bytes": 295583744, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + }, + { + "label": "pr3-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 694, + "run1_ms": 6758, + "run2_ms": 7038, + "peak_rss_bytes": 326549504, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + }, + { + "label": "pr4-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 700, + "run1_ms": 7297, + "run2_ms": 7450, + "peak_rss_bytes": 327860224, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + }, + { + "label": "pr5-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 724, + "run1_ms": 2256, + "run2_ms": 2262, + "peak_rss_bytes": 327073792, + "logits_sha256": "e5447d9a3ab50c35e9b6e9b41dc3670040bb163080484d6aef41a1afab8b669f" + } + ], + "pairs": [ + { + "pair": 1, + "order": [ + "old", + "new" + ], + "run2_old_ms": 6977, + "run2_new_ms": 2277, + "run2_gain": 0.673642, + "run1_old_ms": 7277, + "run1_new_ms": 2263, + "run1_gain": 0.68902 + }, + { + "pair": 2, + "order": [ + "new", + "old" + ], + "run2_old_ms": 7038, + "run2_new_ms": 2237, + "run2_gain": 0.682154, + "run1_old_ms": 6758, + "run1_new_ms": 2236, + "run1_gain": 0.669133 + }, + { + "pair": 3, + "order": [ + "old", + "new" + ], + "run2_old_ms": 7450, + "run2_new_ms": 2262, + "run2_gain": 0.696376, + "run1_old_ms": 7297, + "run1_new_ms": 2256, + "run1_gain": 0.690832 + } + ], + "run2_gain_median": 0.682154, + "run2_gain_min": 0.673642, + "run2_gain_max": 0.696376, + "run1_gain_median": 0.68902, + "speedup_rule": "MET", + "raw_result_json_sha256": "52f61b395aa76d16ca0aa5041c62cbf41161cb23c2cee8c60deb04b7c42e0add" + }, + { + "case": "document-256", + "status": "PASS_WITH_NOT_VERIFIED", + "kernel_release": "7.1.12-2-20261013.18-sep-ARCH", + "start_utc": "2026-11-03T09:07:47+00:00", + "end_utc": "2026-11-03T09:09:28+00:00", + "elfs": { + "old": { + "sha256": "d4949bef20e3b175365d26fab943cafe6f28befdd060185366ff2713368530a3", + "bytes": 5809824 + }, + "new": { + "sha256": "d8ee787b9aefd5f7461ebe3ed40d0dbbacd06f2891822cde476c1b054abfe85c", + "bytes": 5789144 + } + }, + "logits_sha256_all_six_processes": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb", + "processes": [ + { + "label": "pr0-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 697, + "run1_ms": 9751, + "run2_ms": 10155, + "peak_rss_bytes": 311312384, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + }, + { + "label": "pr1-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 707, + "run1_ms": 5201, + "run2_ms": 5201, + "peak_rss_bytes": 324845568, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + }, + { + "label": "pr2-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 710, + "run1_ms": 5232, + "run2_ms": 5202, + "peak_rss_bytes": 327745536, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + }, + { + "label": "pr3-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 706, + "run1_ms": 10797, + "run2_ms": 10961, + "peak_rss_bytes": 327614464, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + }, + { + "label": "pr4-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 711, + "run1_ms": 10134, + "run2_ms": 9921, + "peak_rss_bytes": 327630848, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + }, + { + "label": "pr5-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 697, + "run1_ms": 5205, + "run2_ms": 5206, + "peak_rss_bytes": 327925760, + "logits_sha256": "147a2d2e023222c0fd8f12c88ac72bfc8888faed858c2939252149f88e6bd8cb" + } + ], + "pairs": [ + { + "pair": 1, + "order": [ + "old", + "new" + ], + "run2_old_ms": 10155, + "run2_new_ms": 5201, + "run2_gain": 0.487839, + "run1_old_ms": 9751, + "run1_new_ms": 5201, + "run1_gain": 0.466619 + }, + { + "pair": 2, + "order": [ + "new", + "old" + ], + "run2_old_ms": 10961, + "run2_new_ms": 5202, + "run2_gain": 0.525408, + "run1_old_ms": 10797, + "run1_new_ms": 5232, + "run1_gain": 0.515421 + }, + { + "pair": 3, + "order": [ + "old", + "new" + ], + "run2_old_ms": 9921, + "run2_new_ms": 5206, + "run2_gain": 0.475255, + "run1_old_ms": 10134, + "run1_new_ms": 5205, + "run1_gain": 0.486382 + } + ], + "run2_gain_median": 0.487839, + "run2_gain_min": 0.475255, + "run2_gain_max": 0.525408, + "run1_gain_median": 0.486382, + "speedup_rule": "MET", + "raw_result_json_sha256": "538d859027c43f300cc56393b0005a361a7e468d01b9f2a4c1ceecc8da5b8527" + }, + { + "case": "document-1024", + "status": "PASS_WITH_NOT_VERIFIED", + "kernel_release": "7.1.12-2-20261013.18-sep-ARCH", + "start_utc": "2026-11-03T09:15:13+00:00", + "end_utc": "2026-11-03T09:19:10+00:00", + "elfs": { + "old": { + "sha256": "d4949bef20e3b175365d26fab943cafe6f28befdd060185366ff2713368530a3", + "bytes": 5809824 + }, + "new": { + "sha256": "d8ee787b9aefd5f7461ebe3ed40d0dbbacd06f2891822cde476c1b054abfe85c", + "bytes": 5789144 + } + }, + "logits_sha256_all_six_processes": "7f9d8e65e90eb3751edc41aed99c9ac5d37ce3f7d99dae1e626497632661a820", + "processes": [ + { + "label": "pr0-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 701, + "run1_ms": 22647, + "run2_ms": 22320, + "peak_rss_bytes": 311967744, + "logits_sha256": "7f9d8e65e90eb3751edc41aed99c9ac5d37ce3f7d99dae1e626497632661a820" + }, + { + "label": "pr1-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 722, + "run1_ms": 15790, + "run2_ms": 15730, + "peak_rss_bytes": 327598080, + "logits_sha256": "7f9d8e65e90eb3751edc41aed99c9ac5d37ce3f7d99dae1e626497632661a820" + }, + { + "label": "pr2-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 734, + "run1_ms": 15738, + "run2_ms": 15749, + "peak_rss_bytes": 308314112, + "logits_sha256": "7f9d8e65e90eb3751edc41aed99c9ac5d37ce3f7d99dae1e626497632661a820" + }, + { + "label": "pr3-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 720, + "run1_ms": 22652, + "run2_ms": 22651, + "peak_rss_bytes": 326664192, + "logits_sha256": "7f9d8e65e90eb3751edc41aed99c9ac5d37ce3f7d99dae1e626497632661a820" + }, + { + "label": "pr4-old", + "arm": "old", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 730, + "run1_ms": 21861, + "run2_ms": 22934, + "peak_rss_bytes": 309542912, + "logits_sha256": "7f9d8e65e90eb3751edc41aed99c9ac5d37ce3f7d99dae1e626497632661a820" + }, + { + "label": "pr5-new", + "arm": "new", + "exit_code": 0, + "receipt": "COMPLETE_OK", + "engine_ready_ms": 700, + "run1_ms": 15800, + "run2_ms": 15797, + "peak_rss_bytes": 327483392, + "logits_sha256": "7f9d8e65e90eb3751edc41aed99c9ac5d37ce3f7d99dae1e626497632661a820" + } + ], + "pairs": [ + { + "pair": 1, + "order": [ + "old", + "new" + ], + "run2_old_ms": 22320, + "run2_new_ms": 15730, + "run2_gain": 0.295251, + "run1_old_ms": 22647, + "run1_new_ms": 15790, + "run1_gain": 0.302777 + }, + { + "pair": 2, + "order": [ + "new", + "old" + ], + "run2_old_ms": 22651, + "run2_new_ms": 15749, + "run2_gain": 0.304711, + "run1_old_ms": 22652, + "run1_new_ms": 15738, + "run1_gain": 0.305227 + }, + { + "pair": 3, + "order": [ + "old", + "new" + ], + "run2_old_ms": 22934, + "run2_new_ms": 15797, + "run2_gain": 0.311197, + "run1_old_ms": 21861, + "run1_new_ms": 15800, + "run1_gain": 0.277252 + } + ], + "run2_gain_median": 0.304711, + "run2_gain_min": 0.295251, + "run2_gain_max": 0.311197, + "run1_gain_median": 0.302777, + "speedup_rule": "MET", + "raw_result_json_sha256": "c750d2b47ce643bf4c63cf2d7e67c3e64367e734f5548075f792a03b215ce1aa" + } + ] +} diff --git a/zig/build/vulkan.zig b/zig/build/vulkan.zig index 283ce2edae6..b1df3e10190 100644 --- a/zig/build/vulkan.zig +++ b/zig/build/vulkan.zig @@ -1,131 +1,257 @@ //! Vulkan, built only with -Dvulkan=true: host tests in `test`, named GPU steps; options on every target (#1165). const std = @import("std"); +const vulkan_stages = @import("vulkan_stages.zig"); /// Vulkan's options as given, declared whether or not -Dvulkan=true builds anything with them. -pub const Options = struct { include: ?[]const u8, glslang: ?[]const u8 }; +pub const Options = struct { include: ?[]const u8, glslang: ?[]const u8, stage1_root: ?[]const u8, contract_root: ?[]const u8 }; pub fn options(b: *std.Build) Options { return .{ .include = b.option([]const u8, "vulkan-include", "Directory holding vulkan/vulkan_core.h and vk_video/ (default /usr/include)"), .glslang = b.option([]const u8, "glslang", "glslang for the Vulkan shaders (default: glslang on PATH)"), + .stage1_root = b.option([]const u8, "stage1-test-root", "Root .zig of the stage 1 GPU probe (default zig/tests/vulkan/stage1_probe.zig)"), + .contract_root = b.option([]const u8, "contract-root", "Root .zig of the Vulkan kernel-contract runner (default zig/tests/vulkan/contract_main.zig)"), }; } /// vulkan-host-test (part of `test`), vulkan-gpu-build and vulkan-gpu-test, vulkan-probe-build and vulkan-probe. pub fn steps(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.Optimize, test_step: *std.Build.Step, o: Options) void { - const host = b.step("vulkan-host-test", "Vulkan loader admission, device selection, buffer geometry and the qmm fixture, without GPU work"); - host.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = module(b, b.graph.host, .debug, o) })).step); - const fixture = b.createModule(.{ - .root_source_file = b.path("zig/tests/vulkan/affine_fixture.zig"), - .target = b.graph.host, - .optimize = .debug, - .imports = &.{.{ .name = "contract", .module = @import("kernel_contract.zig").module(b, b.graph.host, .debug) }}, - }); - host.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = fixture })).step); - test_step.dependOn(host); - const vulkan = module(b, target, optimize, o); - const gpu = b.addTest(.{ .root_module = b.createModule(.{ - .root_source_file = b.path("zig/src/vulkan/gpu_tests.zig"), + var ctx = vulkan_stages.Context{ + .b = b, .target = target, .optimize = optimize, - .link_libc = true, - .imports = &.{.{ .name = "vulkan", .module = vulkan }}, - }) }); - b.step("vulkan-gpu-build", "Compile the Vulkan device tests without running GPU work").dependOn(&gpu.step); - b.step("vulkan-gpu-test", "Real Vulkan device selection and staged device-local buffer round trips").dependOn(&b.addRunArtifact(gpu).step); - const glslang = o.glslang orelse "glslang"; - // vulkan1.2 emits SPIR-V 1.5: 1.6's LocalSizeId needs maintenance4, which Device does not enable - const spirv = b.addSystemCommand(&.{ glslang, "-V", "--target-env", "vulkan1.2", "-S", "comp", "-o" }); - const qmm = spirv.addOutputFileArg("qmm.spv"); - spirv.addFileArg(b.path("zig/kernels/vulkan/qmm.comp")); - const files = b.addWriteFiles(); - _ = files.addCopyFile(qmm, "qmm.spv"); - const shaders = b.createModule(.{ .root_source_file = files.add("shaders.zig", "pub const qmm align(4) = @embedFile(\"qmm.spv\").*;\n") }); - const probe = b.addExecutable(.{ .name = "tf-vulkan-probe", .root_module = b.createModule(.{ - .root_source_file = b.path("zig/tests/vulkan/probe.zig"), - .target = target, - .optimize = optimize, - .link_libc = true, - .imports = &.{ - .{ .name = "vulkan", .module = vulkan }, - .{ .name = "vulkan_shaders", .module = shaders }, - .{ .name = "contract", .module = @import("kernel_contract.zig").module(b, target, .debug) }, - }, - }) }); - const install = b.addInstallArtifact(probe, .{}); - b.step("vulkan-probe-build", "Build zig-out/bin/tf-vulkan-probe and its SPIR-V without running GPU work").dependOn(&install.step); - const run = b.addRunArtifact(probe); - run.addPassthruArgs(); - b.step("vulkan-probe", "Run the Vulkan qmm probe on the chosen GPU (args after --)").dependOn(&run.step); - const contract_fixture = b.createModule(.{ - .root_source_file = b.path("zig/tests/vulkan/contract_qmm.zig"), + .o = o, + .glslang = o.glslang orelse "glslang", + .kernel_dir = b.path("zig/kernels/vulkan"), + .vulkan = module(b, target, optimize, o), + }; + vulkan_stages.hostAdmission(&ctx); + test_step.dependOn(ctx.host); + vulkan_stages.probeSteps(&ctx); + vulkan_stages.stage1(&ctx); + vulkan_stages.stage2(&ctx); + vulkan_stages.stage3(&ctx); + vulkan_stages.stage4(&ctx); + vulkan_stages.stage5(&ctx); + vulkan_stages.contracts(&ctx); + vulkan_stages.kernelContract(&ctx); + // Family host admission: refAllDeclsRecursive over the family modules; house wires attach in build.zig. + const model_host_tests = b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/src/vulkan_model_host_tests.zig"), .target = b.graph.host, .optimize = .debug, - .imports = &.{ - .{ .name = "vulkan", .module = module(b, b.graph.host, .debug, o) }, - .{ .name = "contract", .module = @import("kernel_contract.zig").module(b, b.graph.host, .debug) }, - }, - }); - host.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = contract_fixture })).step); - const contract_glslang = o.glslang orelse "glslang"; - const qmm_bf16_cmd = b.addSystemCommand(&.{ contract_glslang, "-V", "--target-env", "vulkan1.2", "-S", "comp", "-DUSE_BF16=1", "-o" }); - const qmm_bf16 = qmm_bf16_cmd.addOutputFileArg("qmm_bf16.spv"); - qmm_bf16_cmd.addFileArg(b.path("zig/kernels/vulkan/qmm.comp")); - const variants = .{ - .{ "qmm_fmax", "-DTF_CTL_QMM_FMAX=1" }, - .{ "atomic_sum", "-DTF_CTL_ATOMIC_SUM=1" }, - .{ "scribble_sum", "-DTF_CTL_SCRIBBLE=1" }, - .{ "drop_tail", "-DTF_CTL_DROP_TAIL=1" }, - .{ "fp16_accum", "-DTF_CTL_FP16=1" }, - .{ "over_limit", "-DTF_CTL_OVER_LIMIT=1" }, - }; - const controls: [6]std.Build.LazyPath = blk: { - var out: [6]std.Build.LazyPath = undefined; - inline for (variants, 0..) |variant, i| { - const command = b.addSystemCommand(&.{ contract_glslang, "-V", "--target-env", "vulkan1.2", "-S", "comp", variant[1], "-o" }); - out[i] = command.addOutputFileArg(variant[0] ++ ".spv"); - command.addFileArg(b.path("zig/kernels/vulkan/contract_controls.comp")); - } - break :blk out; - }; - const contract_files = b.addWriteFiles(); - _ = contract_files.addCopyFile(qmm_bf16, "qmm_bf16.spv"); - inline for (variants, 0..) |variant, i| { - _ = contract_files.addCopyFile(controls[i], variant[0] ++ ".spv"); - } - const contract_shaders = b.createModule(.{ .root_source_file = contract_files.add("vulkan_contract_shaders.zig", "pub const qmm_bf16 align(4) = @embedFile(\"qmm_bf16.spv\").*;\n" ++ - "pub const qmm_fmax align(4) = @embedFile(\"qmm_fmax.spv\").*;\n" ++ - "pub const atomic_sum align(4) = @embedFile(\"atomic_sum.spv\").*;\n" ++ - "pub const scribble_sum align(4) = @embedFile(\"scribble_sum.spv\").*;\n" ++ - "pub const drop_tail align(4) = @embedFile(\"drop_tail.spv\").*;\n" ++ - "pub const fp16_accum align(4) = @embedFile(\"fp16_accum.spv\").*;\n" ++ - "pub const over_limit align(4) = @embedFile(\"over_limit.spv\").*;\n") }); - const contract_module = @import("kernel_contract.zig").module(b, target, .debug); - const contract_exe = b.addExecutable(.{ .name = "tf-vulkan-kernel-contract", .root_module = b.createModule(.{ - .root_source_file = b.path("zig/tests/vulkan/contract_main.zig"), - .target = target, - .optimize = optimize, .link_libc = true, .imports = &.{ - .{ .name = "vulkan", .module = vulkan }, - .{ .name = "vulkan_shaders", .module = contract_shaders }, - .{ .name = "core", .module = contract_module }, - .{ .name = "contract", .module = contract_module }, + .{ .name = "vulkan", .module = module(b, b.graph.host, .debug, ctx.o) }, + .{ .name = "core", .module = coreModule(b, b.graph.host, .debug) }, }, }) }); - const contract_install = b.addInstallArtifact(contract_exe, .{}); - b.step("tf-vulkan-kernel-contract", "Build Vulkan QMM contract runner and planted control shaders without running them").dependOn(&contract_install.step); - b.step("vulkan-kernel-contract", "Run Vulkan QMM contract gates and planted controls on the selected device").dependOn(&b.addRunArtifact(contract_exe).step); + ctx.host.dependOn(&b.addRunArtifact(model_host_tests).step); + b.step("model-host-test", "Qwen3.5 Vulkan family weights/state host tests (TF_QWEN_DIR real-load test stays skipped without the env var)").dependOn(&b.addRunArtifact(model_host_tests).step); +} + +const ShaderVariant = struct { name: []const u8, src: []const u8, defines: []const []const u8 }; + +/// The stage 1 closure's shader images: one variant per op path the host code names in ops.Shaders. +pub const shader_variants = [_]ShaderVariant{ + .{ .name = "qmm_bf16", .src = "zig/kernels/vulkan/qmm.comp", .defines = &.{"-DUSE_BF16"} }, + .{ .name = "qmm_q4_vec4", .src = "zig/kernels/vulkan/qmm_q4_vec4.comp", .defines = &.{} }, + .{ .name = "qmm_q4_tile", .src = "zig/kernels/vulkan/qmm_q4_tile.comp", .defines = &.{} }, + .{ .name = "qmm_q4_tile_f32", .src = "zig/kernels/vulkan/qmm_q4_tile.comp", .defines = &.{"-DACT_F32"} }, + .{ .name = "dequant_bf16", .src = "zig/kernels/vulkan/dequant.comp", .defines = &.{"-DUSE_BF16"} }, + .{ .name = "gather_u32", .src = "zig/kernels/vulkan/gather_take.comp", .defines = &.{"-DUSE_U32"} }, + .{ .name = "gather_u16", .src = "zig/kernels/vulkan/gather_take.comp", .defines = &.{"-DUSE_U16"} }, + .{ .name = "cast_bf16_f32", .src = "zig/kernels/vulkan/cast.comp", .defines = &.{"-DSOURCE_BF16"} }, + .{ .name = "cast_f32_bf16", .src = "zig/kernels/vulkan/cast.comp", .defines = &.{"-DDEST_BF16"} }, + .{ .name = "copy_u32", .src = "zig/kernels/vulkan/copy_general.comp", .defines = &.{"-DUSE_U32"} }, + .{ .name = "copy_u16", .src = "zig/kernels/vulkan/copy_general.comp", .defines = &.{"-DUSE_U16"} }, + .{ .name = "copy_f32_bf16", .src = "zig/kernels/vulkan/copy_general.comp", .defines = &.{"-DUSE_F32_TO_BF16"} }, + .{ .name = "qmm_f32", .src = "zig/kernels/vulkan/qmm.comp", .defines = &.{ "-DUSE_BF16", "-DACT_F32" } }, +}; + +pub const stage1_shaders_src = + \\pub const qmm_bf16 align(4) = @embedFile("qmm_bf16.spv").*; + \\pub const qmm_q4_vec4 align(4) = @embedFile("qmm_q4_vec4.spv").*; + \\pub const qmm_q4_tile align(4) = @embedFile("qmm_q4_tile.spv").*; + \\pub const qmm_q4_tile_f32 align(4) = @embedFile("qmm_q4_tile_f32.spv").*; + \\pub const dequant_bf16 align(4) = @embedFile("dequant_bf16.spv").*; + \\pub const gather_u32 align(4) = @embedFile("gather_u32.spv").*; + \\pub const gather_u16 align(4) = @embedFile("gather_u16.spv").*; + \\pub const cast_bf16_f32 align(4) = @embedFile("cast_bf16_f32.spv").*; + \\pub const cast_f32_bf16 align(4) = @embedFile("cast_f32_bf16.spv").*; + \\pub const copy_u32 align(4) = @embedFile("copy_u32.spv").*; + \\pub const copy_u16 align(4) = @embedFile("copy_u16.spv").*; + \\pub const copy_f32_bf16 align(4) = @embedFile("copy_f32_bf16.spv").*; + \\pub const qmm_f32 align(4) = @embedFile("qmm_f32.spv").*; +; + +/// Stage 2's images: norm/swiglu in the ORACLE_QWEN35 oracle form; add_preserve_bf16 is add plus DenormPreserve (f32). +pub const stage2_variants = [_]ShaderVariant{ + .{ .name = "norm_bf16", .src = "zig/kernels/vulkan/fast_norm.comp", .defines = &.{ "-DUSE_BF16", "-DORACLE_QWEN35" } }, + .{ .name = "add_bf16", .src = "zig/kernels/vulkan/binary_vec.comp", .defines = &.{} }, + .{ .name = "add_preserve_bf16", .src = "zig/kernels/vulkan/binary_vec.comp", .defines = &.{"-DPRESERVE_F32"} }, + .{ .name = "swiglu_bf16", .src = "zig/kernels/vulkan/swiglu.comp", .defines = &.{ "-DUSE_BF16", "-DORACLE_QWEN35" } }, + .{ .name = "norm_f32", .src = "zig/kernels/vulkan/fast_norm.comp", .defines = &.{ "-DUSE_BF16", "-DORACLE_QWEN35", "-DACT_IN_F32", "-DACT_OUT_F32" } }, + .{ .name = "norm_f32_in", .src = "zig/kernels/vulkan/fast_norm.comp", .defines = &.{ "-DUSE_BF16", "-DORACLE_QWEN35", "-DACT_IN_F32", "-DACT_OUT_BF16" } }, + .{ .name = "norm_bf16_in", .src = "zig/kernels/vulkan/fast_norm.comp", .defines = &.{ "-DUSE_BF16", "-DORACLE_QWEN35", "-DACT_OUT_F32" } }, + .{ .name = "add_f32", .src = "zig/kernels/vulkan/binary_vec.comp", .defines = &.{"-DACT_F32"} }, + .{ .name = "add_wide_bf16", .src = "zig/kernels/vulkan/binary_vec.comp", .defines = &.{"-DADD_WIDE"} }, + .{ .name = "add_bf16_f32", .src = "zig/kernels/vulkan/binary_vec.comp", .defines = &.{"-DADD_WIDE_F32"} }, + .{ .name = "swiglu_f32", .src = "zig/kernels/vulkan/swiglu.comp", .defines = &.{ "-DUSE_BF16", "-DORACLE_QWEN35", "-DACT_F32" } }, +}; + +pub const stage2_shaders_src = + \\pub const norm_bf16 align(4) = @embedFile("norm_bf16.spv").*; + \\pub const add_bf16 align(4) = @embedFile("add_bf16.spv").*; + \\pub const add_preserve_bf16 align(4) = @embedFile("add_preserve_bf16.spv").*; + \\pub const swiglu_bf16 align(4) = @embedFile("swiglu_bf16.spv").*; + \\pub const norm_f32 align(4) = @embedFile("norm_f32.spv").*; + \\pub const norm_f32_in align(4) = @embedFile("norm_f32_in.spv").*; + \\pub const norm_bf16_in align(4) = @embedFile("norm_bf16_in.spv").*; + \\pub const add_f32 align(4) = @embedFile("add_f32.spv").*; + \\pub const add_wide_bf16 align(4) = @embedFile("add_wide_bf16.spv").*; + \\pub const add_bf16_f32 align(4) = @embedFile("add_bf16_f32.spv").*; + \\pub const swiglu_f32 align(4) = @embedFile("swiglu_f32.spv").*; +; + +/// Stage 3's images: the ORACLE_QWEN35 rope variant, the causal multi-row SDPA decode arm, the output gate. +pub const stage3_variants = [_]ShaderVariant{ + .{ .name = "rope_bf16", .src = "zig/kernels/vulkan/fast_rope.comp", .defines = &.{ "-DUSE_BF16", "-DORACLE_QWEN35" } }, + .{ .name = "sdpa_rows_bf16", .src = "zig/kernels/vulkan/sdpa_decode_native.comp", .defines = &.{ "-DBF16_IO", "-DSDPA_DIM=256", "-DSDPA_ROWS" } }, + .{ .name = "gate_bf16", .src = "zig/kernels/vulkan/attention_gate.comp", .defines = &.{} }, + .{ .name = "rope_f32", .src = "zig/kernels/vulkan/fast_rope.comp", .defines = &.{"-DORACLE_QWEN35"} }, + .{ .name = "rope_f32_kcache", .src = "zig/kernels/vulkan/fast_rope.comp", .defines = &.{ "-DORACLE_QWEN35", "-DACT_OUT_BF16" } }, + .{ .name = "sdpa_rows_f32q", .src = "zig/kernels/vulkan/sdpa_decode_native.comp", .defines = &.{ "-DBF16_IO", "-DF32_Q", "-DF32_OUT", "-DSDPA_DIM=256", "-DSDPA_ROWS" } }, + .{ .name = "gate_f32", .src = "zig/kernels/vulkan/attention_gate.comp", .defines = &.{"-DACT_F32"} }, +}; + +pub const stage3_shaders_src = + \\pub const rope_bf16 align(4) = @embedFile("rope_bf16.spv").*; + \\pub const sdpa_rows_bf16 align(4) = @embedFile("sdpa_rows_bf16.spv").*; + \\pub const gate_bf16 align(4) = @embedFile("gate_bf16.spv").*; + \\pub const rope_f32 align(4) = @embedFile("rope_f32.spv").*; + \\pub const rope_f32_kcache align(4) = @embedFile("rope_f32_kcache.spv").*; + \\pub const sdpa_rows_f32q align(4) = @embedFile("sdpa_rows_f32q.spv").*; + \\pub const gate_f32 align(4) = @embedFile("gate_f32.spv").*; +; + +/// Stage 4's images: GDN conv, gate split, prefill scan, decode, RMSNorm epilogue and the bf16 copy, all -DUSE_BF16. +pub const stage4_variants = [_]ShaderVariant{ + .{ .name = "conv_dw1d_bf16", .src = "zig/kernels/vulkan/conv_dw1d.comp", .defines = &.{"-DUSE_BF16"} }, + .{ .name = "gdn_conv_decode_bf16", .src = "zig/kernels/vulkan/gdn_conv_decode.comp", .defines = &.{"-DUSE_BF16"} }, + .{ .name = "gdn_split_bf16", .src = "zig/kernels/vulkan/gdn_split.comp", .defines = &.{"-DUSE_BF16"} }, + .{ .name = "prefill_recur_bf16", .src = "zig/kernels/vulkan/gated_delta_prefill_oracle.comp", .defines = &.{"-DUSE_BF16"} }, + .{ .name = "gated_delta_decode_bf16", .src = "zig/kernels/vulkan/gated_delta_decode.comp", .defines = &.{ "-DUSE_BF16", "-DORACLE_QWEN35" } }, + .{ .name = "fast_norm_gated_bf16", .src = "zig/kernels/vulkan/fast_norm_gated.comp", .defines = &.{ "-DUSE_BF16", "-DORACLE_QWEN35" } }, + .{ .name = "conv_dw1d_f32", .src = "zig/kernels/vulkan/conv_dw1d.comp", .defines = &.{ "-DUSE_BF16", "-DACT_F32" } }, + .{ .name = "gdn_split_f32", .src = "zig/kernels/vulkan/gdn_split.comp", .defines = &.{"-DACT_F32"} }, + .{ .name = "prefill_recur_f32", .src = "zig/kernels/vulkan/gated_delta_prefill_oracle.comp", .defines = &.{"-DACT_F32"} }, + .{ .name = "fast_norm_gated_f32", .src = "zig/kernels/vulkan/fast_norm_gated.comp", .defines = &.{ "-DUSE_BF16", "-DORACLE_QWEN35", "-DACT_F32" } }, + .{ .name = "gdn_cast_bf16_f32", .src = "zig/kernels/vulkan/cast.comp", .defines = &.{"-DSOURCE_BF16"} }, + .{ .name = "gdn_cast_f32_bf16", .src = "zig/kernels/vulkan/cast.comp", .defines = &.{"-DDEST_BF16"} }, + .{ .name = "gdn_copy_u32", .src = "zig/kernels/vulkan/copy_general.comp", .defines = &.{"-DUSE_U32"} }, + .{ .name = "copy_u16", .src = "zig/kernels/vulkan/copy_general.comp", .defines = &.{"-DUSE_U16"} }, +}; + +pub const stage4_shaders_src = + \\pub const conv_dw1d_bf16 align(4) = @embedFile("conv_dw1d_bf16.spv").*; + \\pub const gdn_conv_decode_bf16 align(4) = @embedFile("gdn_conv_decode_bf16.spv").*; + \\pub const gdn_split_bf16 align(4) = @embedFile("gdn_split_bf16.spv").*; + \\pub const prefill_recur_bf16 align(4) = @embedFile("prefill_recur_bf16.spv").*; + \\pub const gated_delta_decode_bf16 align(4) = @embedFile("gated_delta_decode_bf16.spv").*; + \\pub const fast_norm_gated_bf16 align(4) = @embedFile("fast_norm_gated_bf16.spv").*; + \\pub const conv_dw1d_f32 align(4) = @embedFile("conv_dw1d_f32.spv").*; + \\pub const gdn_split_f32 align(4) = @embedFile("gdn_split_f32.spv").*; + \\pub const prefill_recur_f32 align(4) = @embedFile("prefill_recur_f32.spv").*; + \\pub const fast_norm_gated_f32 align(4) = @embedFile("fast_norm_gated_f32.spv").*; + \\pub const gdn_cast_bf16_f32 align(4) = @embedFile("gdn_cast_bf16_f32.spv").*; + \\pub const gdn_cast_f32_bf16 align(4) = @embedFile("gdn_cast_f32_bf16.spv").*; + \\pub const gdn_copy_u32 align(4) = @embedFile("gdn_copy_u32.spv").*; + \\pub const copy_u16 align(4) = @embedFile("copy_u16.spv").*; +; + +/// Stage 5's images: the greedy gate's three passes; argreduce_chunk is -DUSE_BF16, combine reduces the f32 partials. +pub const stage5_variants = [_]ShaderVariant{ + .{ .name = "finite_logits", .src = "zig/kernels/vulkan/finite_logits.comp", .defines = &.{} }, + .{ .name = "argreduce_chunk", .src = "zig/kernels/vulkan/argreduce_chunk.comp", .defines = &.{"-DUSE_BF16"} }, + .{ .name = "argreduce_combine", .src = "zig/kernels/vulkan/argreduce_combine.comp", .defines = &.{} }, +}; + +pub const stage5_shaders_src = + \\pub const finite_logits align(4) = @embedFile("finite_logits.spv").*; + \\pub const argreduce_chunk align(4) = @embedFile("argreduce_chunk.spv").*; + \\pub const argreduce_combine align(4) = @embedFile("argreduce_combine.spv").*; +; + +/// The qwen3.5 Vulkan server's modules: the family root, the embedded SPIR-V bundle, the native_engines root. +pub fn serverModules(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.Optimize, o: Options, lanes: *std.Build.Module) struct { qwen35_vulkan: *std.Build.Module, qwen35_shaders: *std.Build.Module, native_engines: *std.Build.Module, engine_api: *std.Build.Module, core: *std.Build.Module } { + const vulkan = module(b, target, optimize, o); + const core = coreModule(b, target, optimize); + const qwen35_vulkan = b.createModule(.{ .root_source_file = b.path("zig/src/families/qwen3_5/vulkan.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ + .{ .name = "vulkan", .module = vulkan }, + .{ .name = "lanes", .module = lanes }, + .{ .name = "core", .module = core }, + } }); + // one glslang invocation per image, all embedded under their bundle names + const glslang = o.glslang orelse "glslang"; + const kernel_dir = b.path("zig/kernels/vulkan"); + const files = b.addWriteFiles(); + var src: std.ArrayListUnmanaged(u8) = .empty; + for (server_shader_images) |img| { + const cmd = b.addSystemCommand(&.{ glslang, "-V", "--target-env", "vulkan1.2", "-S", "comp", "-I." }); + cmd.setCwd(kernel_dir); + for (img.defines) |d| cmd.addArg(d); + cmd.addArg("-o"); + const spv = cmd.addOutputFileArg(b.fmt("{s}.spv", .{img.name})); + cmd.addFileArg(b.path(img.src)); + _ = files.addCopyFile(spv, b.fmt("{s}.spv", .{img.name})); + src.appendSlice(b.allocator, b.fmt("pub const {s} align(4) = @embedFile(\"{s}.spv\").*;\n", .{ img.name, img.name })) catch @panic("OOM"); + } + const qwen35_shaders = b.createModule(.{ .root_source_file = files.add("qwen35_shaders.zig", src.items) }); + const engine_api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = lanes }} }); + const native_engines = b.createModule(.{ .root_source_file = b.path("zig/src/native/vulkan.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ + .{ .name = "engine_api", .module = engine_api }, + .{ .name = "lanes", .module = lanes }, + .{ .name = "qwen35_vulkan", .module = qwen35_vulkan }, + .{ .name = "qwen35_shaders", .module = qwen35_shaders }, + } }); + return .{ .qwen35_vulkan = qwen35_vulkan, .qwen35_shaders = qwen35_shaders, .native_engines = native_engines, .engine_api = engine_api, .core = core }; +} + +/// Every image the bundle embeds: the real stage arrays, concatenated; names are what the adapter's bundle() reads. +const server_shader_images = shader_variants ++ stage2_variants ++ stage3_variants ++ stage4_variants[0 .. stage4_variants.len - 1] ++ stage5_variants; + +/// The house core module the family loader's config/shard imports resolve through, built exactly as cuda.zig builds it. +fn coreModule(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.Optimize) *std.Build.Module { + const tokenizer = b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = target, .optimize = optimize, .link_libc = true }); + const core = b.createModule(.{ .root_source_file = b.path("zig/src/core/root.zig"), .target = target, .optimize = optimize, .link_libc = true }); + core.addImport("gpu_memory", @import("gpu_memory.zig").module(b, target)); + core.addImport("tokenizer", tokenizer); + const living = b.createModule(.{ .root_source_file = b.path("zig/src/living/living.zig"), .target = target, .optimize = optimize, .link_libc = true }); + core.addImport("living", living); + return core; } /// The backend over `vk`, vulkan_core.h translated without prototypes: every entry point comes from the loader. -fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.Optimize, o: Options) *std.Build.Module { +pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.Optimize, o: Options) *std.Build.Module { const header = b.addWriteFiles().add("vulkan_c.h", "#define VK_NO_PROTOTYPES\n#include \n"); const vk = b.addTranslateC(.{ .root_source_file = header, .target = target, .optimize = optimize }); // after the target's libc headers, so a cross build never takes the host's stdint.h vk.addAfterIncludePath(.{ .cwd_relative = o.include orelse "/usr/include" }); const m = b.createModule(.{ .root_source_file = b.path("zig/src/vulkan/root.zig"), .target = target, .optimize = optimize, .link_libc = true }); m.addImport("vk", vk.createModule()); + m.addImport("stage1_params", b.createModule(.{ .root_source_file = b.path("zig/src/vulkan/params.zig"), .target = target, .optimize = optimize })); + // Host-test-only .comp sources: flag-pin tests match push flags against literal masks; shipped ELFs embed SPIR-V. + const comp_files = b.addWriteFiles(); + inline for (.{ "gated_delta_decode", "gated_delta_prefill_oracle", "fast_rope", "gdn_conv_decode" }) |name| + _ = comp_files.addCopyFile(b.path("zig/kernels/vulkan/" ++ name ++ ".comp"), name ++ ".comp"); + m.addImport("vulkan_comp_sources", b.createModule(.{ .root_source_file = comp_files.add("comp_sources.zig", + \\pub const gated_delta_decode = @embedFile("gated_delta_decode.comp"); + \\pub const gated_delta_prefill_oracle = @embedFile("gated_delta_prefill_oracle.comp"); + \\pub const fast_rope = @embedFile("fast_rope.comp"); + \\pub const gdn_conv_decode = @embedFile("gdn_conv_decode.comp"); + \\ + ) })); return m; } diff --git a/zig/build/vulkan_model.zig b/zig/build/vulkan_model.zig new file mode 100644 index 00000000000..fc6c91c45cc --- /dev/null +++ b/zig/build/vulkan_model.zig @@ -0,0 +1,106 @@ +//! The native Vulkan model build, opt-in -Dvulkan=true: the tf-qwen35-model probe, the native server, GPU-free tests. + +const std = @import("std"); +const vulkan_build = @import("vulkan.zig"); +const cuda_build = @import("cuda.zig"); +const server_build = @import("server.zig"); + +/// vulkan-model-probe-build, compile-native-vulkan-server, native-vulkan, qwen35-model-host-test. +pub fn steps( + b: *std.Build, + target: std.Build.ResolvedTarget, + optimize: std.builtin.Optimize, + test_step: *std.Build.Step, + o: vulkan_build.Options, + build_options: *std.Build.Step.Options, +) void { + const m = vulkan_build.serverModules(b, target, optimize, o, lanesModule(b, target, optimize)); + + // (a) the real qwen3.5 model probe over the named family and shader modules; compile only. + const probe = b.addExecutable(.{ .name = "tf-qwen35-model", .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/vulkan/qwen35_model_probe.zig"), + .target = target, + .optimize = optimize, + .link_libc = true, + .imports = &.{ + .{ .name = "qwen35_vulkan", .module = m.qwen35_vulkan }, + .{ .name = "qwen35_shaders", .module = m.qwen35_shaders }, + .{ .name = "core", .module = m.core }, + }, + }) }); + b.step("vulkan-model-probe-build", "Build zig-out/bin/tf-qwen35-model without running GPU work").dependOn(&b.addInstallArtifact(probe, .{}).step); + + // (b) Server entry point, checkpoint CLI and factory engine_api over Vulkan engines. + const template = b.createModule(.{ .root_source_file = b.path("zig/src/core/template/template.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true }); + const server = b.addExecutable(.{ .name = "tensorfold-native-vulkan", .root_module = b.createModule(.{ + .root_source_file = b.path("zig/src/server/main.zig"), + .target = target, + .optimize = .ReleaseSafe, + .link_libc = true, + .imports = &.{ + .{ .name = "core", .module = m.core }, + .{ .name = "engine_api", .module = m.engine_api }, + .{ .name = "living", .module = m.core.import_table.get("living").? }, + .{ .name = "tokenizer", .module = m.core.import_table.get("tokenizer").? }, + .{ .name = "template", .module = template }, + .{ .name = "native_engines", .module = m.native_engines }, + .{ .name = "checkpoint_cli", .module = cuda_build.checkpointCli(b, target, .ReleaseSafe, null, m.native_engines, m.core) }, + }, + }) }); + server.root_module.addOptions("build_options", build_options); + _ = server_build.addTo(b, server.root_module, target, .ReleaseSafe); + b.step("compile-native-vulkan-server", "Compile the actual server entry point over the Vulkan engines without installing or running").dependOn(&server.step); + b.step("native-vulkan", "tensorfold-native with the Vulkan engines into zig-out/native/bin").dependOn(&b.addInstallArtifact(server, .{ .dest_dir = .{ .override = .{ .custom = "native/bin" } } }).step); + + // (c) GPU-free host admission: the adapter's tests, the family backend/engine/forward decls, the probe's parse. + const host = b.graph.host; + const h = vulkan_build.serverModules(b, host, .debug, o, lanesModule(b, host, .debug)); + const all = b.step("qwen35-model-host-test", "Qwen3.5 Vulkan engine/forward/native-backend and probe argument parser host tests, without GPU work"); + const native_tests = b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/src/native/vulkan.zig"), + .target = host, + .optimize = .debug, + .link_libc = true, + .imports = &.{ + .{ .name = "engine_api", .module = h.engine_api }, + .{ .name = "lanes", .module = h.engine_api.import_table.get("lanes").? }, + .{ .name = "qwen35_vulkan", .module = h.qwen35_vulkan }, + .{ .name = "qwen35_shaders", .module = h.qwen35_shaders }, + }, + }) }); + all.dependOn(&b.addRunArtifact(native_tests).step); + // A file-local root discovers tests; a separate imported module does not. + const family_tests = b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/src/families/qwen3_5/vulkan_test.zig"), + .target = host, + .optimize = .debug, + .link_libc = true, + .imports = &.{ + .{ .name = "vulkan", .module = h.qwen35_vulkan.import_table.get("vulkan").? }, + .{ .name = "lanes", .module = h.engine_api.import_table.get("lanes").? }, + .{ .name = "core", .module = h.core }, + }, + }) }); + all.dependOn(&b.addRunArtifact(family_tests).step); + const probe_tests = b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/vulkan/qwen35_model_probe.zig"), + .target = host, + .optimize = .debug, + .link_libc = true, + .imports = &.{ + .{ .name = "qwen35_vulkan", .module = h.qwen35_vulkan }, + .{ .name = "qwen35_shaders", .module = h.qwen35_shaders }, + .{ .name = "core", .module = h.core }, + }, + }) }); + // TF_VULKAN_LOAD_DIR is not a cached input: the engine load reruns when the env changes, like TF_QWEN_DIR's. + const probe_run = b.addRunArtifact(probe_tests); + probe_run.has_side_effects = true; + all.dependOn(&probe_run.step); + test_step.dependOn(all); +} + +/// The lane core the factory takes, built as build.zig builds it for the other seams. +fn lanesModule(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.Optimize) *std.Build.Module { + return b.createModule(.{ .root_source_file = b.path("zig/src/core/lanes/lanes.zig"), .target = target, .optimize = optimize, .link_libc = true }); +} diff --git a/zig/build/vulkan_stages.zig b/zig/build/vulkan_stages.zig new file mode 100644 index 00000000000..1f21ed69f45 --- /dev/null +++ b/zig/build/vulkan_stages.zig @@ -0,0 +1,557 @@ +//! The Vulkan stages' build steps, split out of vulkan.zig: one function per stage, called from it in step order. + +const std = @import("std"); +const vulkan_build = @import("vulkan.zig"); + +/// What a stage's build needs: the graph, the target, the backend module, the host step and the shader compiler. +pub const Context = struct { + b: *std.Build, + target: std.Build.ResolvedTarget, + optimize: std.builtin.Optimize, + o: vulkan_build.Options, + glslang: []const u8, + kernel_dir: std.Build.LazyPath, + vulkan: *std.Build.Module, + host: *std.Build.Step = undefined, + stage3_host: *std.Build.Step = undefined, + stage5_host: *std.Build.Step = undefined, + stage1_shaders_mod: *std.Build.Module = undefined, + stage2_shaders_mod: *std.Build.Module = undefined, + stage3_shaders_mod: *std.Build.Module = undefined, + stage4_shaders_mod: *std.Build.Module = undefined, + stage5_shaders_mod: *std.Build.Module = undefined, + stage3_build: *std.Build.Step = undefined, + stage4_build: *std.Build.Step = undefined, +}; + +/// vulkan-host-test and the five stage host suites: the host admission, without GPU work. +pub fn hostAdmission(ctx: *Context) void { + const b = ctx.b; + const vk_host_tests = b.addTest(.{ .root_module = vulkan_build.module(b, b.graph.host, .debug, ctx.o) }); + ctx.host = b.step("vulkan-host-test", "Vulkan loader admission, device selection, buffer geometry and the qmm fixture, without GPU work"); + ctx.host.dependOn(&b.addRunArtifact(vk_host_tests).step); + const fixture = b.createModule(.{ + .root_source_file = b.path("zig/tests/vulkan/affine_fixture.zig"), + .target = b.graph.host, + .optimize = .debug, + .imports = &.{ + .{ .name = "contract", .module = @import("kernel_contract.zig").module(b, b.graph.host, .debug) }, + .{ .name = "stage1_params", .module = b.createModule(.{ .root_source_file = b.path("zig/src/vulkan/params.zig"), .target = b.graph.host, .optimize = .debug }) }, + }, + }); + ctx.host.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = fixture })).step); + const stage1_fixture = b.createModule(.{ + .root_source_file = b.path("zig/tests/vulkan/stage1_fixture.zig"), + .target = b.graph.host, + .optimize = .debug, + .imports = &.{.{ .name = "contract", .module = @import("kernel_contract.zig").module(b, b.graph.host, .debug) }}, + }); + ctx.host.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = stage1_fixture })).step); + b.step("stage1-host-test", "Stage 1 tensor, arena, params and operator-plan host tests, plus the qmm fixture suite").dependOn(ctx.host); + + // Stage 2 host admission: the vulkan module's own tests plus the stage 2 fixture suite, without GPU work. + const stage2_fixture = b.createModule(.{ + .root_source_file = b.path("zig/tests/vulkan/stage2_fixture.zig"), + .target = b.graph.host, + .optimize = .debug, + .imports = &.{.{ .name = "contract", .module = @import("kernel_contract.zig").module(b, b.graph.host, .debug) }}, + }); + const stage2_fixture_tests = b.addTest(.{ .root_module = stage2_fixture }); + ctx.host.dependOn(&b.addRunArtifact(stage2_fixture_tests).step); + const stage2_host = b.step("stage2-host-test", "Stage 2 norm/MLP operator-plan host tests and the stage 2 fixture suite, without GPU work"); + stage2_host.dependOn(&b.addRunArtifact(vk_host_tests).step); + stage2_host.dependOn(&b.addRunArtifact(stage2_fixture_tests).step); + // The stage 2 probe's host tests (the Bound16 view regression), pure host; `shaders` and `main` stay unanalyzed. + const stage2_probe_host_tests = b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/vulkan/stage2_probe.zig"), + .target = b.graph.host, + .optimize = .debug, + .imports = &.{ + .{ .name = "vulkan", .module = vulkan_build.module(b, b.graph.host, .debug, ctx.o) }, + .{ .name = "contract", .module = @import("kernel_contract.zig").module(b, b.graph.host, .debug) }, + }, + }) }); + stage2_host.dependOn(&b.addRunArtifact(stage2_probe_host_tests).step); + ctx.host.dependOn(&b.addRunArtifact(stage2_probe_host_tests).step); + + // Stages 3 and 4 host admission: the vulkan module's own tests plus each stage's fixture suite, without GPU work. + const stage3_fixture_tests = b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/vulkan/stage3_fixture.zig"), + .target = b.graph.host, + .optimize = .debug, + .imports = &.{.{ .name = "contract", .module = @import("kernel_contract.zig").module(b, b.graph.host, .debug) }}, + }) }); + ctx.host.dependOn(&b.addRunArtifact(stage3_fixture_tests).step); + ctx.stage3_host = b.step("stage3-host-test", "Stage 3 attention plan host tests and the stage 3 fixture suite, without GPU work"); + ctx.stage3_host.dependOn(&b.addRunArtifact(vk_host_tests).step); + ctx.stage3_host.dependOn(&b.addRunArtifact(stage3_fixture_tests).step); + const stage4_fixture_tests = b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/vulkan/stage4_fixture.zig"), + .target = b.graph.host, + .optimize = .debug, + .link_libc = true, + .imports = &.{ + .{ .name = "contract", .module = @import("kernel_contract.zig").module(b, b.graph.host, .debug) }, + .{ .name = "vulkan", .module = vulkan_build.module(b, b.graph.host, .debug, ctx.o) }, + }, + }) }); + ctx.host.dependOn(&b.addRunArtifact(stage4_fixture_tests).step); + const stage4_host = b.step("stage4-host-test", "Stage 4 GDN plan host tests and the stage 4 fixture suite, without GPU work"); + stage4_host.dependOn(&b.addRunArtifact(vk_host_tests).step); + stage4_host.dependOn(&b.addRunArtifact(stage4_fixture_tests).step); + + // Stage 5 host admission: the vulkan module's tests plus its fixture suite; honest failure until the fixture lands. + const stage5_fixture_tests = b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/vulkan/stage5_fixture.zig"), + .target = b.graph.host, + .optimize = .debug, + .imports = &.{.{ .name = "contract", .module = @import("kernel_contract.zig").module(b, b.graph.host, .debug) }}, + }) }); + ctx.host.dependOn(&b.addRunArtifact(stage5_fixture_tests).step); + ctx.stage5_host = b.step("stage5-host-test", "Stage 5 greedy plan host tests and the stage 5 fixture suite, without GPU work"); + ctx.stage5_host.dependOn(&b.addRunArtifact(vk_host_tests).step); + ctx.stage5_host.dependOn(&b.addRunArtifact(stage5_fixture_tests).step); +} + +/// vulkan-gpu-build, vulkan-gpu-test, vulkan-probe-build and vulkan-probe: the device tests and the qmm probe. +pub fn probeSteps(ctx: *Context) void { + const b = ctx.b; + const gpu = b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/src/vulkan/gpu_tests.zig"), + .target = ctx.target, + .optimize = ctx.optimize, + .link_libc = true, + .imports = &.{.{ .name = "vulkan", .module = ctx.vulkan }}, + }) }); + b.step("vulkan-gpu-build", "Compile the Vulkan device tests without running GPU work").dependOn(&gpu.step); + b.step("vulkan-gpu-test", "Real Vulkan device selection and staged device-local buffer round trips").dependOn(&b.addRunArtifact(gpu).step); + // vulkan1.2 emits SPIR-V 1.5: 1.6's LocalSizeId needs maintenance4, which Device does not enable + const spirv = b.addSystemCommand(&.{ ctx.glslang, "-V", "--target-env", "vulkan1.2", "-S", "comp", "-o" }); + const qmm = spirv.addOutputFileArg("qmm.spv"); + spirv.addFileArg(b.path("zig/kernels/vulkan/qmm.comp")); + const files = b.addWriteFiles(); + _ = files.addCopyFile(qmm, "qmm.spv"); + const shaders = b.createModule(.{ .root_source_file = files.add("shaders.zig", "pub const qmm align(4) = @embedFile(\"qmm.spv\").*;\n") }); + const probe = b.addExecutable(.{ .name = "tf-vulkan-probe", .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/vulkan/probe.zig"), + .target = ctx.target, + .optimize = ctx.optimize, + .link_libc = true, + .imports = &.{ + .{ .name = "vulkan", .module = ctx.vulkan }, + .{ .name = "vulkan_shaders", .module = shaders }, + .{ .name = "contract", .module = @import("kernel_contract.zig").module(b, ctx.target, .debug) }, + .{ .name = "stage1_params", .module = ctx.vulkan.import_table.get("stage1_params").? }, + }, + }) }); + const install = b.addInstallArtifact(probe, .{}); + b.step("vulkan-probe-build", "Build zig-out/bin/tf-vulkan-probe and its SPIR-V without running GPU work").dependOn(&install.step); + const run = b.addRunArtifact(probe); + run.addPassthruArgs(); + b.step("vulkan-probe", "Run the Vulkan qmm probe on the chosen GPU (args after --)").dependOn(&run.step); +} + +/// stage1-gpu-build, stage1-gpu-run and stage1-datamove: the stage 1 images and the datamove contract. +pub fn stage1(ctx: *Context) void { + const b = ctx.b; + // Stage 1: the operator variants as SPIR-V assets, and the GPU probe from the fixture owner's root. + const stage1_files = b.addWriteFiles(); + ctx.stage1_shaders_mod = b.createModule(.{ .root_source_file = stage1_files.add("stage1_shaders.zig", vulkan_build.stage1_shaders_src) }); + const stage1_build = b.step("stage1-gpu-build", "Compile the stage 1 SPIR-V variants, install them under stage1/, and build the stage 1 GPU probe without running it"); + const stage1_run = b.step("stage1-gpu-run", "Run the stage 1 GPU probe on the chosen device (args after --)"); + for (vulkan_build.shader_variants) |v| { + const cmd = b.addSystemCommand(&.{ ctx.glslang, "-V", "--target-env", "vulkan1.2", "-S", "comp" }); + for (v.defines) |d| cmd.addArg(d); + cmd.addArg("-o"); + const spv = cmd.addOutputFileArg(b.fmt("{s}.spv", .{v.name})); + cmd.addFileArg(b.path(v.src)); + stage1_build.dependOn(&b.addInstallFile(spv, b.fmt("stage1/{s}.spv", .{v.name})).step); + _ = stage1_files.addCopyFile(spv, b.fmt("{s}.spv", .{v.name})); + } + const stage1_imports: []const std.Build.Module.Import = &.{ + .{ .name = "vulkan", .module = ctx.vulkan }, + .{ .name = "stage1_shaders", .module = ctx.stage1_shaders_mod }, + .{ .name = "contract", .module = @import("kernel_contract.zig").module(b, ctx.target, .debug) }, + }; + const probe_root = b.path(ctx.o.stage1_root orelse "zig/tests/vulkan/stage1_probe.zig"); + const stage1_probe = b.addExecutable(.{ .name = "tf-stage1-gpu", .root_module = b.createModule(.{ + .root_source_file = probe_root, + .target = ctx.target, + .optimize = ctx.optimize, + .link_libc = true, + .imports = stage1_imports, + }) }); + stage1_build.dependOn(&b.addInstallArtifact(stage1_probe, .{}).step); + const stage1_probe_run = b.addRunArtifact(stage1_probe); + stage1_probe_run.addPassthruArgs(); + stage1_run.dependOn(&stage1_probe_run.step); + + // The datamove group's contract: the seven exact kernels through gates 1-4. + const datamove_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/vulkan/contract_datamove_main.zig"), + .target = ctx.target, + .optimize = ctx.optimize, + .link_libc = true, + .imports = stage1_imports, + }); + const datamove = b.addExecutable(.{ .name = "tf-stage1-datamove", .root_module = datamove_module }); + stage1_build.dependOn(&b.addInstallArtifact(datamove, .{}).step); + const datamove_run = b.addRunArtifact(datamove); + datamove_run.addPassthruArgs(); + b.step("stage1-datamove", "Run the datamove group's contract gates on the chosen device (args after --)").dependOn(&datamove_run.step); + // The datamove fixtures' host suite runs on the build host, whatever -Dtarget the runner is built for. + const datamove_host_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/vulkan/contract_datamove_main.zig"), + .target = b.graph.host, + .optimize = .debug, + .link_libc = true, + .imports = &.{ + .{ .name = "vulkan", .module = vulkan_build.module(b, b.graph.host, .debug, ctx.o) }, + .{ .name = "stage1_shaders", .module = ctx.stage1_shaders_mod }, + .{ .name = "contract", .module = @import("kernel_contract.zig").module(b, b.graph.host, .debug) }, + }, + }); + ctx.host.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = datamove_host_module })).step); +} + +/// stage2-gpu-build and stage2-gpu-run: the stage 2 norm/swiglu images and the NormMlp probe. +pub fn stage2(ctx: *Context) void { + const b = ctx.b; + // Stage 2: norm/swiglu as the ORACLE_QWEN35 oracle form, plus plain bf16 and -DUSE_BF16; the NormMlp probe. + const stage2_files = b.addWriteFiles(); + const stage2_build = b.step("stage2-gpu-build", "Compile the stage 2 SPIR-V variants, install them under stage2/, and build the stage 2 GPU probe without running it"); + const stage2_run = b.step("stage2-gpu-run", "Run the stage 2 GPU probe on the chosen device (args after --)"); + for (vulkan_build.stage2_variants) |v| { + const cmd = b.addSystemCommand(&.{ ctx.glslang, "-V", "--target-env", "vulkan1.2", "-S", "comp" }); + for (v.defines) |d| cmd.addArg(d); + cmd.addArg("-o"); + const spv = cmd.addOutputFileArg(b.fmt("{s}.spv", .{v.name})); + cmd.addFileArg(b.path(v.src)); + stage2_build.dependOn(&b.addInstallFile(spv, b.fmt("stage2/{s}.spv", .{v.name})).step); + _ = stage2_files.addCopyFile(spv, b.fmt("{s}.spv", .{v.name})); + } + ctx.stage2_shaders_mod = b.createModule(.{ .root_source_file = stage2_files.add("stage2_shaders.zig", vulkan_build.stage2_shaders_src) }); + const stage2_probe = b.addExecutable(.{ .name = "tf-stage2-gpu", .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/vulkan/stage2_probe.zig"), + .target = ctx.target, + .optimize = ctx.optimize, + .link_libc = true, + .imports = &.{ + .{ .name = "vulkan", .module = ctx.vulkan }, + .{ .name = "stage2_shaders", .module = ctx.stage2_shaders_mod }, + .{ .name = "contract", .module = @import("kernel_contract.zig").module(b, ctx.target, .debug) }, + }, + }) }); + stage2_build.dependOn(&b.addInstallArtifact(stage2_probe, .{}).step); + const stage2_probe_run = b.addRunArtifact(stage2_probe); + stage2_probe_run.addPassthruArgs(); + stage2_run.dependOn(&stage2_probe_run.step); +} + +/// stage3-gpu-build, stage3-gpu-run and attention-contract-run: the attention images and the attention contract. +pub fn stage3(ctx: *Context) void { + const b = ctx.b; + // Stages 3 and 4: attention and GatedDeltaNet; compiles run from the kernel dir with -I. for fast_rope's #include. + const stage3_files = b.addWriteFiles(); + ctx.stage3_build = b.step("stage3-gpu-build", "Compile the stage 3 attention SPIR-V variants, install them under stage3/, and build the stage 3 GPU probe without running it"); + const stage3_run = b.step("stage3-gpu-run", "Run the stage 3 GPU probe on the chosen device (args after --)"); + for (vulkan_build.stage3_variants) |v| { + const cmd = b.addSystemCommand(&.{ ctx.glslang, "-V", "--target-env", "vulkan1.2", "-S", "comp", "-I." }); + cmd.setCwd(ctx.kernel_dir); + for (v.defines) |d| cmd.addArg(d); + cmd.addArg("-o"); + const spv = cmd.addOutputFileArg(b.fmt("{s}.spv", .{v.name})); + cmd.addFileArg(b.path(v.src)); + ctx.stage3_build.dependOn(&b.addInstallFile(spv, b.fmt("stage3/{s}.spv", .{v.name})).step); + _ = stage3_files.addCopyFile(spv, b.fmt("{s}.spv", .{v.name})); + } + ctx.stage3_shaders_mod = b.createModule(.{ .root_source_file = stage3_files.add("stage3_shaders.zig", vulkan_build.stage3_shaders_src) }); + const stage3_probe = b.addExecutable(.{ .name = "tf-stage3-gpu", .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/vulkan/stage3_probe.zig"), + .target = ctx.target, + .optimize = ctx.optimize, + .link_libc = true, + .imports = &.{ + .{ .name = "vulkan", .module = ctx.vulkan }, + .{ .name = "stage1_shaders", .module = ctx.stage1_shaders_mod }, + .{ .name = "stage3_shaders", .module = ctx.stage3_shaders_mod }, + .{ .name = "contract", .module = @import("kernel_contract.zig").module(b, ctx.target, .debug) }, + }, + }) }); + ctx.stage3_build.dependOn(&b.addInstallArtifact(stage3_probe, .{}).step); + const stage3_probe_run = b.addRunArtifact(stage3_probe); + stage3_probe_run.addPassthruArgs(); + stage3_run.dependOn(&stage3_probe_run.step); + + // The attention contract: rope and sdpa_rows against the fixture's f64 refs; the GPU probe reuses stage 3 images. + const attention_files = b.addWriteFiles(); + _ = attention_files.addCopyFile(b.path("zig/kernels/vulkan/sdpa_decode_native.comp"), "sdpa_decode_native.comp"); + const attention_sources = b.createModule(.{ .root_source_file = attention_files.add("attention_sources.zig", + \\pub const sdpa_decode_native = @embedFile("sdpa_decode_native.comp"); + \\ + ) }); + const attention_contract_tests = b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/vulkan/contract_attention.zig"), + .target = b.graph.host, + .optimize = .debug, + .imports = &.{ + .{ .name = "contract", .module = @import("kernel_contract.zig").module(b, b.graph.host, .debug) }, + .{ .name = "attention_sources", .module = attention_sources }, + }, + }) }); + ctx.host.dependOn(&b.addRunArtifact(attention_contract_tests).step); + ctx.stage3_host.dependOn(&b.addRunArtifact(attention_contract_tests).step); + const attention_contract = b.addExecutable(.{ .name = "tf-attention-contract", .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/vulkan/contract_attention_main.zig"), + .target = ctx.target, + .optimize = ctx.optimize, + .link_libc = true, + .imports = &.{ + .{ .name = "vulkan", .module = ctx.vulkan }, + .{ .name = "stage1_shaders", .module = ctx.stage1_shaders_mod }, + .{ .name = "stage3_shaders", .module = ctx.stage3_shaders_mod }, + .{ .name = "contract", .module = @import("kernel_contract.zig").module(b, ctx.target, .debug) }, + .{ .name = "attention_sources", .module = attention_sources }, + }, + }) }); + ctx.stage3_build.dependOn(&b.addInstallArtifact(attention_contract, .{}).step); + const attention_contract_run = b.addRunArtifact(attention_contract); + attention_contract_run.addPassthruArgs(); + b.step("attention-contract-run", "Run the attention contract (rope_bf16 + sdpa_rows_bf16, served shapes) on the chosen device (args after --)").dependOn(&attention_contract_run.step); +} + +/// stage4-gpu-build and stage4-gpu-run: the stage 4 GDN images and the probe. +pub fn stage4(ctx: *Context) void { + const b = ctx.b; + const stage4_files = b.addWriteFiles(); + ctx.stage4_build = b.step("stage4-gpu-build", "Compile the stage 4 GDN SPIR-V variants, install them under stage4/, and build the stage 4 GPU probe without running it"); + const stage4_run = b.step("stage4-gpu-run", "Run the stage 4 GPU probe on the chosen device (args after --)"); + for (vulkan_build.stage4_variants) |v| { + const cmd = b.addSystemCommand(&.{ ctx.glslang, "-V", "--target-env", "vulkan1.2", "-S", "comp", "-I." }); + cmd.setCwd(ctx.kernel_dir); + for (v.defines) |d| cmd.addArg(d); + cmd.addArg("-o"); + const spv = cmd.addOutputFileArg(b.fmt("{s}.spv", .{v.name})); + cmd.addFileArg(b.path(v.src)); + ctx.stage4_build.dependOn(&b.addInstallFile(spv, b.fmt("stage4/{s}.spv", .{v.name})).step); + _ = stage4_files.addCopyFile(spv, b.fmt("{s}.spv", .{v.name})); + } + ctx.stage4_shaders_mod = b.createModule(.{ .root_source_file = stage4_files.add("stage4_shaders.zig", vulkan_build.stage4_shaders_src) }); + const stage4_probe = b.addExecutable(.{ .name = "tf-stage4-gpu", .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/vulkan/stage4_probe.zig"), + .target = ctx.target, + .optimize = ctx.optimize, + .link_libc = true, + .imports = &.{ + .{ .name = "vulkan", .module = ctx.vulkan }, + .{ .name = "stage4_shaders", .module = ctx.stage4_shaders_mod }, + .{ .name = "contract", .module = @import("kernel_contract.zig").module(b, ctx.target, .debug) }, + }, + }) }); + ctx.stage4_build.dependOn(&b.addInstallArtifact(stage4_probe, .{}).step); + const stage4_probe_run = b.addRunArtifact(stage4_probe); + stage4_probe_run.addPassthruArgs(); + stage4_run.dependOn(&stage4_probe_run.step); +} + +/// stage5-gpu-build, stage5-gpu-run and greedy-contract-run: the greedy images and the greedy contract. +pub fn stage5(ctx: *Context) void { + const b = ctx.b; + // Stage 5: the greedy gate (finite logits, chunked argreduce, combine); fails honestly until its probe root lands. + const stage5_files = b.addWriteFiles(); + const stage5_build = b.step("stage5-gpu-build", "Compile the stage 5 greedy SPIR-V variants, install them under stage5/, and build the stage 5 GPU probe without running it"); + const stage5_run = b.step("stage5-gpu-run", "Run the stage 5 GPU probe on the chosen device (args after --)"); + for (vulkan_build.stage5_variants) |v| { + const cmd = b.addSystemCommand(&.{ ctx.glslang, "-V", "--target-env", "vulkan1.2", "-S", "comp", "-I." }); + cmd.setCwd(ctx.kernel_dir); + for (v.defines) |d| cmd.addArg(d); + cmd.addArg("-o"); + const spv = cmd.addOutputFileArg(b.fmt("{s}.spv", .{v.name})); + cmd.addFileArg(b.path(v.src)); + stage5_build.dependOn(&b.addInstallFile(spv, b.fmt("stage5/{s}.spv", .{v.name})).step); + _ = stage5_files.addCopyFile(spv, b.fmt("{s}.spv", .{v.name})); + } + ctx.stage5_shaders_mod = b.createModule(.{ .root_source_file = stage5_files.add("stage5_shaders.zig", vulkan_build.stage5_shaders_src) }); + const stage5_probe = b.addExecutable(.{ .name = "tf-stage5-gpu", .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/vulkan/stage5_probe.zig"), + .target = ctx.target, + .optimize = ctx.optimize, + .link_libc = true, + .imports = &.{ + .{ .name = "vulkan", .module = ctx.vulkan }, + .{ .name = "stage5_shaders", .module = ctx.stage5_shaders_mod }, + .{ .name = "contract", .module = @import("kernel_contract.zig").module(b, ctx.target, .debug) }, + }, + }) }); + stage5_build.dependOn(&b.addInstallArtifact(stage5_probe, .{}).step); + const stage5_probe_run = b.addRunArtifact(stage5_probe); + stage5_probe_run.addPassthruArgs(); + stage5_run.dependOn(&stage5_probe_run.step); + + // The greedy contract: the composed draw through gates 1-4 plus the NaN control; the probe reuses stage 5 images. + const greedy_files = b.addWriteFiles(); + const greedy_ctl_cmd = b.addSystemCommand(&.{ ctx.glslang, "-V", "--target-env", "vulkan1.2", "-S", "comp", "-o" }); + const greedy_ctl_spv = greedy_ctl_cmd.addOutputFileArg("contract_controls_greedy.spv"); + greedy_ctl_cmd.addFileArg(b.path("zig/kernels/vulkan/contract_controls_greedy.comp")); + _ = greedy_files.addCopyFile(greedy_ctl_spv, "contract_controls_greedy.spv"); + const greedy_controls = b.createModule(.{ .root_source_file = greedy_files.add("greedy_controls.zig", "pub const fmax_max align(4) = @embedFile(\"contract_controls_greedy.spv\").*;\n") }); + const greedy_contract_tests = b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/vulkan/contract_greedy.zig"), + .target = b.graph.host, + .optimize = .debug, + .imports = &.{ + .{ .name = "vulkan", .module = vulkan_build.module(b, b.graph.host, .debug, ctx.o) }, + .{ .name = "contract", .module = @import("kernel_contract.zig").module(b, b.graph.host, .debug) }, + }, + }) }); + ctx.host.dependOn(&b.addRunArtifact(greedy_contract_tests).step); + ctx.stage5_host.dependOn(&b.addRunArtifact(greedy_contract_tests).step); + const greedy_contract = b.addExecutable(.{ .name = "tf-vulkan-greedy-contract", .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/vulkan/contract_greedy_main.zig"), + .target = ctx.target, + .optimize = ctx.optimize, + .link_libc = true, + .imports = &.{ + .{ .name = "vulkan", .module = ctx.vulkan }, + .{ .name = "stage5_shaders", .module = ctx.stage5_shaders_mod }, + .{ .name = "greedy_controls", .module = greedy_controls }, + .{ .name = "contract", .module = @import("kernel_contract.zig").module(b, ctx.target, .debug) }, + }, + }) }); + stage5_build.dependOn(&b.addInstallArtifact(greedy_contract, .{}).step); + const greedy_contract_run = b.addRunArtifact(greedy_contract); + greedy_contract_run.addPassthruArgs(); + b.step("greedy-contract-run", "Run the greedy contract (the composed draw and its planted control) on the chosen device (args after --)").dependOn(&greedy_contract_run.step); +} + +/// contract-pointwise-host-test, contract-pointwise-build and contract-pointwise: the pointwise contract. +pub fn contracts(ctx: *Context) void { + const b = ctx.b; + const contract_fixture = b.createModule(.{ + .root_source_file = b.path("zig/tests/vulkan/contract_qmm.zig"), + .target = b.graph.host, + .optimize = .debug, + .imports = &.{ + .{ .name = "vulkan", .module = vulkan_build.module(b, b.graph.host, .debug, ctx.o) }, + .{ .name = "contract", .module = @import("kernel_contract.zig").module(b, b.graph.host, .debug) }, + }, + }); + ctx.host.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = contract_fixture })).step); + // The gdn host proofs: walker vs the stage 4 fixture's, depth shape honesty, stress sums' f32 and fp16 margins + const contract_gdn_fixture = b.createModule(.{ + .root_source_file = b.path("zig/tests/vulkan/contract_gdn.zig"), + .target = b.graph.host, + .optimize = .debug, + .imports = &.{ + .{ .name = "vulkan", .module = vulkan_build.module(b, b.graph.host, .debug, ctx.o) }, + .{ .name = "contract", .module = @import("kernel_contract.zig").module(b, b.graph.host, .debug) }, + }, + }); + ctx.host.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = contract_gdn_fixture })).step); + // The pointwise contract: all six closures against one GPU probe, plus the fixtures' host suite in `test`. + const pointwise_fixture_tests = b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/vulkan/contract_pointwise.zig"), + .target = b.graph.host, + .optimize = .debug, + .imports = &.{ + .{ .name = "contract", .module = @import("kernel_contract.zig").module(b, b.graph.host, .debug) }, + .{ .name = "stage2_shaders", .module = ctx.stage2_shaders_mod }, + .{ .name = "vulkan", .module = vulkan_build.module(b, b.graph.host, .debug, ctx.o) }, + }, + }) }); + ctx.host.dependOn(&b.addRunArtifact(pointwise_fixture_tests).step); + b.step("contract-pointwise-host-test", "The pointwise contract fixtures' host suite (references and checkers), without GPU work").dependOn(ctx.host); + const pointwise = b.addExecutable(.{ .name = "tf-contract-pointwise", .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/vulkan/contract_pointwise_main.zig"), + .target = ctx.target, + .optimize = ctx.optimize, + .link_libc = true, + .imports = &.{ + .{ .name = "vulkan", .module = ctx.vulkan }, + .{ .name = "stage2_shaders", .module = ctx.stage2_shaders_mod }, + .{ .name = "stage3_shaders", .module = ctx.stage3_shaders_mod }, + .{ .name = "stage4_shaders", .module = ctx.stage4_shaders_mod }, + .{ .name = "contract", .module = @import("kernel_contract.zig").module(b, ctx.target, .debug) }, + }, + }) }); + const pointwise_install = b.addInstallArtifact(pointwise, .{}); + ctx.stage3_build.dependOn(&pointwise_install.step); + ctx.stage4_build.dependOn(&pointwise_install.step); + b.step("contract-pointwise-build", "Compile the six pointwise SPIR-V images and the pointwise contract probe without running GPU work").dependOn(&pointwise_install.step); + const pointwise_run = b.addRunArtifact(pointwise); + pointwise_run.addPassthruArgs(); + b.step("contract-pointwise", "Run the pointwise shader contract on the chosen GPU (args after --)").dependOn(&pointwise_run.step); +} + +/// tf-vulkan-kernel-contract and vulkan-kernel-contract: the QMM contract runner and its planted controls. +pub fn kernelContract(ctx: *Context) void { + const b = ctx.b; + const qmm_bf16_cmd = b.addSystemCommand(&.{ ctx.glslang, "-V", "--target-env", "vulkan1.2", "-S", "comp", "-DUSE_BF16=1", "-o" }); + const qmm_bf16 = qmm_bf16_cmd.addOutputFileArg("qmm_bf16.spv"); + qmm_bf16_cmd.addFileArg(b.path("zig/kernels/vulkan/qmm.comp")); + const qmm_vec4_cmd = b.addSystemCommand(&.{ ctx.glslang, "-V", "--target-env", "vulkan1.2", "-S", "comp", "-o" }); + const qmm_vec4 = qmm_vec4_cmd.addOutputFileArg("qmm_q4_vec4.spv"); + qmm_vec4_cmd.addFileArg(b.path("zig/kernels/vulkan/qmm_q4_vec4.comp")); + const qmm_tile_cmd = b.addSystemCommand(&.{ ctx.glslang, "-V", "--target-env", "vulkan1.2", "-S", "comp", "-o" }); + const qmm_tile = qmm_tile_cmd.addOutputFileArg("qmm_q4_tile.spv"); + qmm_tile_cmd.addFileArg(b.path("zig/kernels/vulkan/qmm_q4_tile.comp")); + const qmm_tile_f32_cmd = b.addSystemCommand(&.{ ctx.glslang, "-V", "--target-env", "vulkan1.2", "-S", "comp", "-DACT_F32=1", "-o" }); + const qmm_tile_f32 = qmm_tile_f32_cmd.addOutputFileArg("qmm_q4_tile_f32.spv"); + qmm_tile_f32_cmd.addFileArg(b.path("zig/kernels/vulkan/qmm_q4_tile.comp")); + const variants = .{ + .{ "qmm_fmax", "-DTF_CTL_QMM_FMAX=1" }, + .{ "atomic_sum", "-DTF_CTL_ATOMIC_SUM=1" }, + .{ "scribble_sum", "-DTF_CTL_SCRIBBLE=1" }, + .{ "drop_tail", "-DTF_CTL_DROP_TAIL=1" }, + .{ "fp16_accum", "-DTF_CTL_FP16=1" }, + .{ "over_limit", "-DTF_CTL_OVER_LIMIT=1" }, + }; + const controls: [6]std.Build.LazyPath = blk: { + var out: [6]std.Build.LazyPath = undefined; + inline for (variants, 0..) |variant, i| { + const command = b.addSystemCommand(&.{ ctx.glslang, "-V", "--target-env", "vulkan1.2", "-S", "comp", variant[1], "-o" }); + out[i] = command.addOutputFileArg(variant[0] ++ ".spv"); + command.addFileArg(b.path("zig/kernels/vulkan/contract_controls.comp")); + } + break :blk out; + }; + const contract_files = b.addWriteFiles(); + _ = contract_files.addCopyFile(qmm_bf16, "qmm_bf16.spv"); + _ = contract_files.addCopyFile(qmm_vec4, "qmm_q4_vec4.spv"); + _ = contract_files.addCopyFile(qmm_tile, "qmm_q4_tile.spv"); + _ = contract_files.addCopyFile(qmm_tile_f32, "qmm_q4_tile_f32.spv"); + inline for (variants, 0..) |variant, i| { + _ = contract_files.addCopyFile(controls[i], variant[0] ++ ".spv"); + } + const contract_shaders = b.createModule(.{ .root_source_file = contract_files.add("vulkan_contract_shaders.zig", "pub const qmm_bf16 align(4) = @embedFile(\"qmm_bf16.spv\").*;\n" ++ + "pub const qmm_q4_vec4 align(4) = @embedFile(\"qmm_q4_vec4.spv\").*;\n" ++ + "pub const qmm_q4_tile align(4) = @embedFile(\"qmm_q4_tile.spv\").*;\n" ++ + "pub const qmm_q4_tile_f32 align(4) = @embedFile(\"qmm_q4_tile_f32.spv\").*;\n" ++ + "pub const qmm_fmax align(4) = @embedFile(\"qmm_fmax.spv\").*;\n" ++ + "pub const atomic_sum align(4) = @embedFile(\"atomic_sum.spv\").*;\n" ++ + "pub const scribble_sum align(4) = @embedFile(\"scribble_sum.spv\").*;\n" ++ + "pub const drop_tail align(4) = @embedFile(\"drop_tail.spv\").*;\n" ++ + "pub const fp16_accum align(4) = @embedFile(\"fp16_accum.spv\").*;\n" ++ + "pub const over_limit align(4) = @embedFile(\"over_limit.spv\").*;\n") }); + const contract_module = @import("kernel_contract.zig").module(b, ctx.target, .debug); + const contract_exe = b.addExecutable(.{ .name = "tf-vulkan-kernel-contract", .root_module = b.createModule(.{ + .root_source_file = b.path(ctx.o.contract_root orelse "zig/tests/vulkan/contract_main.zig"), + .target = ctx.target, + .optimize = ctx.optimize, + .link_libc = true, + .imports = &.{ + .{ .name = "vulkan", .module = ctx.vulkan }, + .{ .name = "vulkan_shaders", .module = contract_shaders }, + .{ .name = "stage1_shaders", .module = ctx.stage1_shaders_mod }, + .{ .name = "stage2_shaders", .module = ctx.stage2_shaders_mod }, + .{ .name = "stage3_shaders", .module = ctx.stage3_shaders_mod }, + .{ .name = "stage4_shaders", .module = ctx.stage4_shaders_mod }, + .{ .name = "stage5_shaders", .module = ctx.stage5_shaders_mod }, + .{ .name = "core", .module = contract_module }, + .{ .name = "contract", .module = contract_module }, + }, + }) }); + const contract_install = b.addInstallArtifact(contract_exe, .{}); + b.step("tf-vulkan-kernel-contract", "Build Vulkan QMM contract runner and planted control shaders without running them").dependOn(&contract_install.step); + b.step("vulkan-kernel-contract", "Run Vulkan QMM contract gates and planted controls on the selected device").dependOn(&b.addRunArtifact(contract_exe).step); +} diff --git a/zig/kernels/vulkan/argreduce_chunk.comp b/zig/kernels/vulkan/argreduce_chunk.comp new file mode 100644 index 00000000000..c963d8c8b26 --- /dev/null +++ b/zig/kernels/vulkan/argreduce_chunk.comp @@ -0,0 +1,120 @@ +// From omarchy-mlx overlay/mlx/backend/omarchy/shaders/argreduce_chunk.comp at 6f0de59de6d0. +// Code unchanged; used here for operation = 1 (max), one row, matrix_m = 4096. +// Copyright © 2026 Joshua Warren and mlx-omarchy contributors; MIT License, full text in LICENSES/mlx-omarchy-MIT.txt. + +#version 460 +#ifdef USE_FP16 +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#define STORAGE_TYPE float16_t +#define LOAD_VALUE(x) float(x) +#elif defined(USE_BF16) +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require +#define STORAGE_TYPE uint16_t +#define LOAD_VALUE(x) uintBitsToFloat(uint(x) << 16u) +#else +#define STORAGE_TYPE float +#define LOAD_VALUE(x) float(x) +#endif + +#define ACC_LOWEST uintBitsToFloat(0xff800000u) +#define ACC_HIGHEST uintBitsToFloat(0x7f800000u) + +layout(local_size_x = 256) in; + +struct Partial { + float value; + uint index; +}; + +layout(set = 0, binding = 0, std430) readonly buffer Input { + STORAGE_TYPE values[]; +} input_data; +layout(set = 0, binding = 1, std430) writeonly buffer Partials { + Partial values[]; +} partial_data; + +layout(push_constant) uniform Params { + uint count; + uint operation; + uint lhs_size; + uint rhs_size; + uint reduce_size; + uint output_size; + uint lhs_offset; + uint rhs_offset; + uint output_offset; + uint aux_size; + uint aux_offset; + uint matrix_m; + uint matrix_n; + uint matrix_k; + uint flags; +} params; + +shared float block_values[256]; +shared uint block_indices[256]; + +bool challenges(float next, float incumbent, bool is_max) { + return isnan(next) ? !isnan(incumbent) + : (!isnan(incumbent) && + (is_max ? next > incumbent : next < incumbent)); +} + +bool challenges_at( + float next, float incumbent, uint next_index, uint incumbent_index, + bool is_max) { + return isnan(next) + ? (!isnan(incumbent) || next_index < incumbent_index) + : (!isnan(incumbent) && + ((is_max ? next > incumbent : next < incumbent) || + (next == incumbent && next_index < incumbent_index))); +} + +void main() { + uint thread = gl_LocalInvocationID.x; + uint n = params.reduce_size; + uint chunks = params.matrix_n; + uint chunk_length = params.matrix_m; + bool is_max = params.operation == 1u; + uint id = gl_WorkGroupID.x; + uint row = id / chunks; + uint chunk = id - row * chunks; + uint begin = chunk * chunk_length; + uint end = min(begin + chunk_length, n); + uint start = params.lhs_offset + row * n; + float never = is_max ? ACC_LOWEST : ACC_HIGHEST; + float v[16]; + for (uint k = 0u; k < 16u; ++k) { + uint index = begin + thread + k * 256u; + v[k] = index < end ? LOAD_VALUE(input_data.values[start + index]) : never; + } + float local_value = never; + uint local_index = 0u; + for (uint k = 0u; k < 16u; ++k) { + if (challenges(v[k], local_value, is_max)) { + local_value = v[k]; + local_index = begin + thread + k * 256u; + } + } + block_values[thread] = local_value; + block_indices[thread] = local_index; + barrier(); + for (uint stride = 128u; stride > 0u; stride >>= 1u) { + if (thread < stride) { + float mine = block_values[thread]; + float theirs = block_values[thread + stride]; + if (challenges_at( + theirs, mine, block_indices[thread + stride], + block_indices[thread], is_max)) { + block_values[thread] = theirs; + block_indices[thread] = block_indices[thread + stride]; + } + } + barrier(); + } + if (thread == 0u) { + partial_data.values[id] = Partial(block_values[0], block_indices[0]); + } +} diff --git a/zig/kernels/vulkan/argreduce_combine.comp b/zig/kernels/vulkan/argreduce_combine.comp new file mode 100644 index 00000000000..fa7ed9cbab2 --- /dev/null +++ b/zig/kernels/vulkan/argreduce_combine.comp @@ -0,0 +1,86 @@ +// From omarchy-mlx overlay/mlx/backend/omarchy/shaders/argreduce_combine.comp at 6f0de59de6d0. +// Code unchanged; used here for operation = 1 (max), one row, output_offset = 2. +// Copyright © 2026 Joshua Warren and mlx-omarchy contributors; MIT License, full text in LICENSES/mlx-omarchy-MIT.txt. + +#version 460 +#define ACC_LOWEST uintBitsToFloat(0xff800000u) +#define ACC_HIGHEST uintBitsToFloat(0x7f800000u) + +layout(local_size_x = 256) in; + +struct Partial { + float value; + uint index; +}; + +layout(set = 0, binding = 0, std430) readonly buffer Partials { + Partial values[]; +} partial_data; +layout(set = 0, binding = 1, std430) writeonly buffer Output { + uint values[]; +} output_data; + +layout(push_constant) uniform Params { + uint count; + uint operation; + uint lhs_size; + uint rhs_size; + uint reduce_size; + uint output_size; + uint lhs_offset; + uint rhs_offset; + uint output_offset; + uint aux_size; + uint aux_offset; + uint matrix_m; + uint matrix_n; + uint matrix_k; + uint flags; +} params; + +shared float block_values[256]; +shared uint block_indices[256]; + +bool challenges_at( + float next, float incumbent, uint next_index, uint incumbent_index, + bool is_max) { + return isnan(next) + ? (!isnan(incumbent) || next_index < incumbent_index) + : (!isnan(incumbent) && + ((is_max ? next > incumbent : next < incumbent) || + (next == incumbent && next_index < incumbent_index))); +} + +void main() { + uint thread = gl_LocalInvocationID.x; + uint chunks = params.matrix_n; + bool is_max = params.operation == 1u; + float never = is_max ? ACC_LOWEST : ACC_HIGHEST; + uint row = gl_WorkGroupID.x; + float value = never; + uint index = 0u; + if (thread < chunks) { + Partial p = partial_data.values[row * chunks + thread]; + value = p.value; + index = p.index; + } + block_values[thread] = value; + block_indices[thread] = index; + barrier(); + for (uint stride = 128u; stride > 0u; stride >>= 1u) { + if (thread < stride) { + float mine = block_values[thread]; + float theirs = block_values[thread + stride]; + if (challenges_at( + theirs, mine, block_indices[thread + stride], + block_indices[thread], is_max)) { + block_values[thread] = theirs; + block_indices[thread] = block_indices[thread + stride]; + } + } + barrier(); + } + if (thread == 0u) { + output_data.values[params.output_offset + row] = block_indices[0]; + } +} diff --git a/zig/kernels/vulkan/attention_gate.comp b/zig/kernels/vulkan/attention_gate.comp new file mode 100644 index 00000000000..0de1f35a36c --- /dev/null +++ b/zig/kernels/vulkan/attention_gate.comp @@ -0,0 +1,92 @@ +#version 460 +#extension GL_EXT_shader_16bit_storage : require + +layout(local_size_x = 256) in; + +#ifdef ACT_F32 +// F32 activations: ctx, gate halves and output address f32 elements directly. +layout(set = 0, binding = 0, std430) readonly buffer Context { + float values[]; +} ctx_data; +layout(set = 0, binding = 1, std430) readonly buffer Gate { + float values[]; +} gate_data; +layout(set = 0, binding = 2, std430) writeonly buffer Output { + float values[]; +} out_data; + +layout(push_constant) uniform Push { + uint words; // rows * 2048 f32 elements +} push; + +void main() { + uint index = gl_GlobalInvocationID.x; + uint stride = gl_NumWorkGroups.x * gl_WorkGroupSize.x; + while (index < push.words) { + const uint row = index >> 11u; // elements per row: 2048 + const uint d = index & 2047u; + const uint head = d >> 8u; // 256 elements per head block + const float gate = gate_data.values[row * 4096u + head * 512u + 256u + (d & 255u)]; + const float s = 1.0f / (1.0f + exp(-gate)); + out_data.values[index] = ctx_data.values[index] * s; + if (push.words - index <= stride) { + break; + } + index += stride; + } +} +#else + +layout(set = 0, binding = 0, std430) readonly buffer Context { + uint values[]; +} ctx_data; +layout(set = 0, binding = 1, std430) readonly buffer Gate { + uint values[]; +} gate_data; +layout(set = 0, binding = 2, std430) writeonly buffer Output { + uint values[]; +} out_data; + +layout(push_constant) uniform Push { + uint words; // rows * 1024 +} push; + +uint bf16_bits(float value) { + uint bits = floatBitsToUint(value); + return isnan(value) + ? ((bits >> 16u) | 0x40u) + : ((bits + 0x7fffu + ((bits >> 16u) & 1u)) >> 16u); +} + +float bf16_lane(uint word, uint lane) { + return uintBitsToFloat(lane == 0u ? (word << 16u) : (word & 0xFFFF0000u)); +} + +void main() { + uint index = gl_GlobalInvocationID.x; + uint stride = gl_NumWorkGroups.x * gl_WorkGroupSize.x; + while (index < push.words) { + const uint row = index >> 10u; // words per row: 2048 elements / 2 + const uint w = index & 1023u; + const uint head = w >> 7u; // 128 words per 256-wide block + const uint d0 = (w & 127u) << 1u; // element within the head block + // Gate word: qraw[r*4096 + head*512 + 256 + d0 .. +1], all even. + const uint gate_word = gate_data.values[ + ((row * 4096u + head * 512u + 256u + d0) >> 1u)]; + const uint ctx_word = ctx_data.values[index]; + uint out_word = 0u; + for (uint lane = 0u; lane < 2u; ++lane) { + float gate = bf16_lane(gate_word, lane); + float s = 1.0f / (1.0f + exp(-gate)); + s = bf16_lane(bf16_bits(s), 0u); // Metal's bfloat(s) rounding + float v = bf16_lane(ctx_word, lane) * s; + out_word |= bf16_bits(v) << (16u * lane); + } + out_data.values[index] = out_word; + if (push.words - index <= stride) { + break; + } + index += stride; + } +} +#endif diff --git a/zig/kernels/vulkan/binary_vec.comp b/zig/kernels/vulkan/binary_vec.comp new file mode 100644 index 00000000000..cf3aede7274 --- /dev/null +++ b/zig/kernels/vulkan/binary_vec.comp @@ -0,0 +1,138 @@ +// From omarchy-mlx overlay/mlx/backend/omarchy/shaders/binary_vec.comp at 94e1f535bc3f. +// Code adapted: 3 lines added or changed, 0 removed. +// Copyright © 2026 Joshua Warren and mlx-omarchy contributors; MIT License, full text in LICENSES/mlx-omarchy-MIT.txt. + +#version 460 + +// -DPRESERVE_F32: DenormPreserve 4459/32; SPV_KHR_float_controls; cap 4464; shaderDenormPreserveFloat32; signed zero. +#ifdef PRESERVE_F32 +#extension GL_EXT_spirv_intrinsics : require +spirv_execution_mode(extensions = ["SPV_KHR_float_controls"], capabilities = [4464], 4459, 32); +#endif + +#ifdef ACT_F32 +// F32 activations: all three buffers are plain vec4. +#define VEC_A vec4 +#define VEC_B vec4 +#define VEC_OUT vec4 +#define LOAD_A(x) (x) +#define LOAD_B(x) (x) +#define STORE_OUT(v) (v) +#elif defined(ADD_WIDE) +// Mixed add: the f32 wide stream adds a bf16 residual and rounds the sum to bf16. +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require +#define VEC_A vec4 +#define VEC_B u16vec4 +#define VEC_OUT u16vec4 +#define LOAD_A(x) (x) +#define LOAD_B(x) uintBitsToFloat(uvec4(x) << 16u) +uint16_t bf16_store(float value) { + uint bits = floatBitsToUint(value); + if (isnan(value)) { + return uint16_t((bits >> 16) | 0x40u); + } + return uint16_t((bits + 0x7fffu + ((bits >> 16) & 1u)) >> 16); +} +#define STORE_OUT(v) \ + u16vec4(bf16_store(v.x), bf16_store(v.y), bf16_store(v.z), bf16_store(v.w)) +#elif defined(ADD_WIDE_F32) +// Mixed add: the bf16 narrow stream joins the f32 wide stream, keeping the f32 sum. +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require +#define VEC_A u16vec4 +#define VEC_B vec4 +#define VEC_OUT vec4 +#define LOAD_A(x) uintBitsToFloat(uvec4(x) << 16u) +#define LOAD_B(x) (x) +#define STORE_OUT(v) (v) +#else +#ifdef USE_FP16 +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#define STORAGE_VEC f16vec4 +#define LOAD_VEC(x) vec4(x) +#define STORE_VEC(v) f16vec4(v) +#else +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require +#define STORAGE_VEC u16vec4 +#define LOAD_VEC(x) uintBitsToFloat(uvec4(x) << 16u) +uint16_t bf16_store(float value) { + uint bits = floatBitsToUint(value); + if (isnan(value)) { + return uint16_t((bits >> 16) | 0x40u); + } + return uint16_t((bits + 0x7fffu + ((bits >> 16) & 1u)) >> 16); +} +#define STORE_VEC(v) \ + u16vec4(bf16_store(v.x), bf16_store(v.y), bf16_store(v.z), bf16_store(v.w)) +#endif +#define VEC_A STORAGE_VEC +#define VEC_B STORAGE_VEC +#define VEC_OUT STORAGE_VEC +#define LOAD_A(x) LOAD_VEC(x) +#define LOAD_B(x) LOAD_VEC(x) +#define STORE_OUT(v) STORE_VEC(v) +#endif + +layout(local_size_x = 256) in; + +layout(set = 0, binding = 0, std430) readonly buffer InputA { + VEC_A values[]; +} input_a; +layout(set = 0, binding = 1, std430) readonly buffer InputB { + VEC_B values[]; +} input_b; +layout(set = 0, binding = 2, std430) writeonly buffer Output { + VEC_OUT values[]; +} output_data; + +// Push constants match elementwise.comp: first nine words are read; count, sizes, offsets are elements divided by four. +layout(push_constant) uniform Params { + uint count; + uint operation; + uint lhs_size; + uint rhs_size; + uint reduce_size; + uint output_size; + uint lhs_offset; + uint rhs_offset; + uint output_offset; +} params; + +void main() { + uint vec_count = params.count / 4u; + uint lhs_vecs = params.lhs_size / 4u; + uint rhs_vecs = params.rhs_size / 4u; + uint index = gl_GlobalInvocationID.x; + uint stride = gl_NumWorkGroups.x * gl_WorkGroupSize.x; + while (index < vec_count) { + uint lhs_index = params.lhs_offset / 4u + + (lhs_vecs == vec_count ? index : index % lhs_vecs); + uint rhs_index = params.rhs_offset / 4u + + (rhs_vecs == vec_count ? index : index % rhs_vecs); + vec4 lhs = LOAD_A(input_a.values[lhs_index]); + vec4 rhs = LOAD_B(input_b.values[rhs_index]); + vec4 value; + switch (params.operation) { + case 0u: + value = lhs + rhs; + break; + case 1u: + value = lhs * rhs; + break; + case 2u: + value = lhs / rhs; + break; + default: + value = lhs - rhs; + break; + } + output_data.values[params.output_offset / 4u + index] = STORE_OUT(value); + if (vec_count - index <= stride) { + break; + } + index += stride; + } +} diff --git a/zig/kernels/vulkan/cast.comp b/zig/kernels/vulkan/cast.comp new file mode 100644 index 00000000000..eb2e7701b2e --- /dev/null +++ b/zig/kernels/vulkan/cast.comp @@ -0,0 +1,148 @@ +#version 460 + +// From omarchy-mlx overlay/mlx/backend/omarchy/shaders/cast.comp at fde15c1ece69: + +// code unchanged, comments condensed. + +// Copyright © 2026 Joshua Warren; MIT License, full text in LICENSES/mlx-omarchy-MIT.txt. + +// Element cast across SOURCE_*/DEST_* macros; complex64 real-part load drops imag. + +// Real->complex64 is vec2(x, 0); BF16 uint16 RNE store with NaN quiet bit. + +// SOURCE_BOOL: one word load per stride iter outside lane loop (Honeykrisp bool load quirk). + +#if defined(SOURCE_FP16) || defined(DEST_FP16) +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#endif +#if defined(SOURCE_BF16) || defined(DEST_BF16) +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require +#endif + +#ifdef SOURCE_FP16 +#define SOURCE_TYPE float16_t +#elif defined(SOURCE_BF16) +#define SOURCE_TYPE uint16_t +#elif defined(SOURCE_C64) +#define SOURCE_TYPE vec2 +#elif defined(SOURCE_I32) +#define SOURCE_TYPE int +#elif defined(SOURCE_U32) +#define SOURCE_TYPE uint +#elif defined(SOURCE_BOOL) +#define SOURCE_TYPE uint +#else +#define SOURCE_TYPE float +#endif + +#ifndef SOURCE_BOOL +#ifdef SOURCE_C64 +// Real-part projection: complex64->real load drops imag (upstream operator float). +#define SOURCE_LOAD(i) input_data.values[(i)].x +#else +#define SOURCE_LOAD(i) input_data.values[(i)] +#endif +#endif + +#ifdef DEST_FP16 +#define DEST_TYPE float16_t +#elif defined(DEST_BF16) +#define DEST_TYPE uint16_t +#elif defined(DEST_C64) +#define DEST_TYPE vec2 +#elif defined(DEST_I32) +#define DEST_TYPE int +#else +#define DEST_TYPE float +#endif + +#if defined(SOURCE_BF16) || defined(DEST_BF16) +float bf16_load(uint16_t bits) { + return uintBitsToFloat(uint(bits) << 16u); +} +uint16_t bf16_store(float value) { + uint bits = floatBitsToUint(value); + if (isnan(value)) { + return uint16_t((bits >> 16) | 0x40u); + } + return uint16_t((bits + 0x7fffu + ((bits >> 16) & 1u)) >> 16); +} +#endif + +#if defined(DEST_C64) +// Real->complex64 imag=0; BF16 widens via bf16_load; CONVERT arg appears once. +#if defined(SOURCE_BF16) +#define CONVERT(x) vec2(bf16_load(x), 0.0) +#else +#define CONVERT(x) vec2(float(x), 0.0) +#endif +#elif defined(SOURCE_BF16) && defined(DEST_BF16) +#define CONVERT(x) (x) +#elif defined(SOURCE_BF16) +#define CONVERT(x) DEST_TYPE(bf16_load(x)) +#elif defined(DEST_BF16) +#define CONVERT(x) bf16_store(float(x)) +#else +#define CONVERT(x) DEST_TYPE(x) +#endif + +layout(local_size_x = 256) in; + +layout(set = 0, binding = 0, std430) readonly buffer Input { + SOURCE_TYPE values[]; +} input_data; +layout(set = 0, binding = 2, std430) writeonly buffer Output { + DEST_TYPE values[]; +} output_data; + +layout(push_constant) uniform Params { + uint count; + uint operation; + uint lhs_size; + uint rhs_size; + uint reduce_size; + uint output_size; + uint lhs_offset; + uint rhs_offset; + uint output_offset; +} params; + +void main() { +#ifdef SOURCE_BOOL +// SOURCE_BOOL: one word load per stride iter outside lane loop (Honeykrisp bool load quirk). + +// mis/word_count cover unaligned start; grid-stride for counts above 65535*256*4. + uint mis = params.lhs_offset & 3u; + uint word_count = (params.count + mis + 3u) >> 2u; + uint stride = gl_NumWorkGroups.x * gl_WorkGroupSize.x; + for (uint word_index = gl_GlobalInvocationID.x; word_index < word_count; + word_index += stride) { + uint word = input_data.values[(params.lhs_offset >> 2u) + word_index]; + uint first = 4u * word_index - mis; + for (uint lane = 0u; lane < 4u; ++lane) { + if (word_index == 0u && lane < mis) { + continue; + } + uint index = first + lane; + if (index >= params.count) { + break; + } + output_data.values[params.output_offset + index] = + CONVERT((word >> (lane * 8u)) & 0xFFu); + } + } +#else + uint index = gl_GlobalInvocationID.x; + uint stride = gl_NumWorkGroups.x * gl_WorkGroupSize.x; + while (index < params.count) { + output_data.values[params.output_offset + index] = + CONVERT(SOURCE_LOAD(params.lhs_offset + index)); + if (params.count - index <= stride) { + break; + } + index += stride; + } +#endif +} diff --git a/zig/kernels/vulkan/contract_controls_greedy.comp b/zig/kernels/vulkan/contract_controls_greedy.comp new file mode 100644 index 00000000000..451dde15bfc --- /dev/null +++ b/zig/kernels/vulkan/contract_controls_greedy.comp @@ -0,0 +1,37 @@ +// Greedy Vulkan control: OpFMax lets NaN lose; it accepts nonfinite rows, so the nonfinite gate must fail. +#version 460 + +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require + +layout(local_size_x = 256) in; + +layout(set = 0, binding = 0, std430) readonly buffer Input { + uint16_t values[]; +} input_data; +layout(set = 0, binding = 1, std430) writeonly buffer Maxima { + float values[]; +} maxima_data; + +layout(push_constant) uniform Params { + uint count; +} params; + +shared float block_max[256]; + +void main() { + uint thread = gl_LocalInvocationID.x; + uint row = gl_WorkGroupID.x; + float m = uintBitsToFloat(0xff800000u); // -inf + for (uint i = thread; i < params.count; i += 256u) + m = max(m, uintBitsToFloat(uint(input_data.values[row * params.count + i]) << 16u)); + block_max[thread] = m; + barrier(); + for (uint stride = 128u; stride > 0u; stride >>= 1u) { + if (thread < stride) + block_max[thread] = max(block_max[thread], block_max[thread + stride]); + barrier(); + } + if (thread == 0u) + maxima_data.values[row] = block_max[0]; +} diff --git a/zig/kernels/vulkan/conv_dw1d.comp b/zig/kernels/vulkan/conv_dw1d.comp new file mode 100644 index 00000000000..9c3f09b1a7b --- /dev/null +++ b/zig/kernels/vulkan/conv_dw1d.comp @@ -0,0 +1,88 @@ +// From omarchy-mlx overlay/mlx/backend/omarchy/shaders/conv_dw1d.comp at 7addbbe9d2ff. +// Code unchanged, comments condensed. +// Copyright © 2026 Joshua Warren and mlx-omarchy contributors; MIT License, full text in LICENSES/mlx-omarchy-MIT.txt. +#version 460 + + +#ifdef USE_FP16 +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#define STORAGE_TYPE float16_t +#define LOAD_VALUE(x) float(x) +#define STORE_VALUE(f) STORAGE_TYPE(f) +#elif defined(USE_BF16) +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require +#define STORAGE_TYPE uint16_t +#define LOAD_VALUE(x) uintBitsToFloat(uint(x) << 16u) +uint16_t bf16_store(float value) { + uint bits = floatBitsToUint(value); + if (isnan(value)) { + return uint16_t((bits >> 16) | 0x40u); + } + return uint16_t((bits + 0x7fffu + ((bits >> 16) & 1u)) >> 16); +} +#define STORE_VALUE(f) bf16_store(f) +#else +#define STORAGE_TYPE float +#define LOAD_VALUE(x) float(x) +#define STORE_VALUE(f) STORAGE_TYPE(f) +#endif + +// ACT_F32 (with USE_BF16): input and output are f32 while the taps stay bf16. +#ifdef ACT_F32 +#define ACT_TYPE float +#define ACT_LOAD(x) (x) +#define ACT_STORE(f) (f) +#else +#define ACT_TYPE STORAGE_TYPE +#define ACT_LOAD(x) LOAD_VALUE(x) +#define ACT_STORE(f) STORE_VALUE(f) +#endif + +layout(local_size_x = 256) in; + +layout(set = 0, binding = 0, std430) readonly buffer InputX { + ACT_TYPE values[]; +} input_x; +layout(set = 0, binding = 1, std430) readonly buffer Weight { + STORAGE_TYPE values[]; +} weight; +layout(set = 0, binding = 2, std430) writeonly buffer Output { + ACT_TYPE values[]; +} output_data; + +layout(push_constant) uniform Params { + uint count; + uint operation; + uint lhs_size; + uint rhs_size; + uint reduce_size; + uint output_size; + uint lhs_offset; + uint rhs_offset; + uint output_offset; +} params; + +void main() { + const uint channels = params.reduce_size; + const uint taps = params.operation; + const uint l_in = params.lhs_size; + const uint l_out = params.rhs_size; + uint stride_threads = gl_NumWorkGroups.x * 256u; + for (uint index = gl_GlobalInvocationID.x; index < params.count; + index += stride_threads) { + uint channel = index % channels; + uint row = index / channels; + uint t = row % l_out; + uint batch = row / l_out; + uint in_base = params.lhs_offset + (batch * l_in + t) * channels + channel; + uint w_base = params.rhs_offset + channel * taps; + float accumulator = 0.0; + for (uint k = 0u; k < taps; ++k) { + accumulator += ACT_LOAD(input_x.values[in_base + k * channels]) * + LOAD_VALUE(weight.values[w_base + k]); + } + output_data.values[params.output_offset + index] = ACT_STORE(accumulator); + } +} diff --git a/zig/kernels/vulkan/copy_general.comp b/zig/kernels/vulkan/copy_general.comp new file mode 100644 index 00000000000..1aca4e70ebb --- /dev/null +++ b/zig/kernels/vulkan/copy_general.comp @@ -0,0 +1,198 @@ +#version 460 + +// From omarchy-mlx overlay/mlx/backend/omarchy/shaders/copy_general.comp at fde15c1ece69: + +// code unchanged, comments condensed. + +// Copyright © 2026 Joshua Warren; MIT License, full text in LICENSES/mlx-omarchy-MIT.txt. + +// Strided copy; USE_U32/U16/U64 bit-exact; USE_BOOL/U8 packed-byte atomicAnd/Or (scatter.comp). + +// USE_COMPLEX64: vec2 item copy; signed index math for negative-stride views (int32 span). + +// Honeykrisp: one straight-line per-thread load at a linear address. + +#ifdef USE_FP16 +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#define STORAGE_TYPE float16_t +#define LOAD_VALUE(x) float(x) +#define STORE_VALUE(f) STORAGE_TYPE(f) +#elif defined(USE_U32) +#define STORAGE_TYPE uint +#define LOAD_VALUE(x) (x) +#define STORE_VALUE(x) (x) +#elif defined(USE_BF16) +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require +#define STORAGE_TYPE uint16_t +#define LOAD_VALUE(x) uintBitsToFloat(uint(x) << 16u) +uint16_t bf16_store(float value) { + uint bits = floatBitsToUint(value); + if (isnan(value)) { + return uint16_t((bits >> 16) | 0x40u); + } + return uint16_t((bits + 0x7fffu + ((bits >> 16) & 1u)) >> 16); +} +#define STORE_VALUE(f) bf16_store(f) +#elif defined(USE_COMPLEX64) +// Complex64: vec2 (re,im) bit-exact like USE_U32; strides/offsets are item offsets. +#define STORAGE_TYPE vec2 +#define LOAD_VALUE(x) (x) +#define STORE_VALUE(f) STORAGE_TYPE(f) +#elif defined(USE_BOOL) || defined(USE_U8) +#define STORAGE_TYPE uint +#define LOAD_VALUE(x) (x) +#define STORE_VALUE(x) (x) +#elif defined(USE_U16) +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require +#define STORAGE_TYPE uint16_t +#define LOAD_VALUE(x) (x) +#define STORE_VALUE(x) (x) +#elif defined(USE_U64) +#extension GL_EXT_shader_explicit_arithmetic_types_int64 : require +#define STORAGE_TYPE uint64_t +#define LOAD_VALUE(x) (x) +#define STORE_VALUE(x) (x) +#elif defined(USE_F32_TO_BF16) +// F32 input, bf16 output with RNE store: the V-cache write rounds f32 activations on landing. +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require +#define STORAGE_TYPE float +#define LOAD_VALUE(x) (x) +uint16_t bf16_store(float value) { + uint bits = floatBitsToUint(value); + if (isnan(value)) { + return uint16_t((bits >> 16) | 0x40u); + } + return uint16_t((bits + 0x7fffu + ((bits >> 16) & 1u)) >> 16); +} +#define STORE_VALUE(f) bf16_store(f) +#define OUT_TYPE uint16_t +#else +#define STORAGE_TYPE float +#define LOAD_VALUE(x) float(x) +#define STORE_VALUE(f) STORAGE_TYPE(f) +#endif + +#ifndef OUT_TYPE +#define OUT_TYPE STORAGE_TYPE +#endif + +layout(local_size_x = 256) in; + +layout(set = 0, binding = 0, std430) readonly buffer InputA { + STORAGE_TYPE values[]; +} input_a; +layout(set = 0, binding = 1, std430) readonly buffer InputB { + STORAGE_TYPE values[]; +} input_b; +#if defined(USE_BOOL) || defined(USE_U8) +// Packed-byte dest uses atomics so binding is read-write. +layout(set = 0, binding = 2, std430) buffer Output { + uint values[]; +} output_data; +#else +layout(set = 0, binding = 2, std430) writeonly buffer Output { + OUT_TYPE values[]; +} output_data; +#endif + +layout(push_constant) uniform Params { + uint count; + uint operation; + uint lhs_size; + uint rhs_size; + uint reduce_size; + uint output_size; + uint lhs_offset; + uint rhs_offset; + uint output_offset; + uint aux_size; + uint aux_offset; + uint matrix_m; + uint matrix_n; + uint matrix_k; + uint flags; + float alpha; + float beta; + uint dims; + uint shape[4]; + uint in_strides[4]; + uint out_strides[4]; +} params; + +layout(set = 0, binding = 3, std430) readonly buffer AxisMetadata { + uint values[]; +} axis_data; + +uint axis_extent(uint index) { + return params.matrix_k == 0u ? params.shape[index] : axis_data.values[index]; +} + +uint in_stride_at(uint index) { + return params.matrix_k == 0u + ? params.in_strides[index] + : axis_data.values[params.matrix_k + index]; +} + +uint out_stride_at(uint index) { + return params.matrix_k == 0u + ? params.out_strides[index] + : axis_data.values[2u * params.matrix_k + index]; +} + +void main() { +#if defined(USE_BOOL) || defined(USE_U8) + uint index = gl_GlobalInvocationID.x; + uint stride = gl_NumWorkGroups.x * gl_WorkGroupSize.x; + while (index < params.count) { + uint rem = index; + int in_index = int(params.lhs_offset); + int out_index = int(params.output_offset); + for (int dim = int(params.dims) - 1; dim >= 0; --dim) { + uint coord = rem % axis_extent(uint(dim)); + rem /= axis_extent(uint(dim)); + in_index += int(coord) * int(in_stride_at(uint(dim))); + out_index += int(coord) * int(out_stride_at(uint(dim))); + } + uint word = input_a.values[in_index >> 2]; + uint lane = uint(in_index & 3); + uint b = (lane == 0u) + ? (word & 0xFFu) + : (lane == 1u) + ? ((word >> 8u) & 0xFFu) + : (lane == 2u) + ? ((word >> 16u) & 0xFFu) + : ((word >> 24u) & 0xFFu); + uint shift = uint(out_index & 3) * 8u; + atomicAnd(output_data.values[out_index >> 2], ~(0xFFu << shift)); + atomicOr(output_data.values[out_index >> 2], b << shift); + if (params.count - index <= stride) { + break; + } + index += stride; + } +#else + uint index = gl_GlobalInvocationID.x; + uint stride = gl_NumWorkGroups.x * gl_WorkGroupSize.x; + while (index < params.count) { + uint rem = index; + int in_index = int(params.lhs_offset); + int out_index = int(params.output_offset); + for (int dim = int(params.dims) - 1; dim >= 0; --dim) { + uint coord = rem % axis_extent(uint(dim)); + rem /= axis_extent(uint(dim)); + in_index += int(coord) * int(in_stride_at(uint(dim))); + out_index += int(coord) * int(out_stride_at(uint(dim))); + } + output_data.values[out_index] = + STORE_VALUE(LOAD_VALUE(input_a.values[in_index])); + if (params.count - index <= stride) { + break; + } + index += stride; + } +#endif +} diff --git a/zig/kernels/vulkan/dequant.comp b/zig/kernels/vulkan/dequant.comp new file mode 100644 index 00000000000..372b909a045 --- /dev/null +++ b/zig/kernels/vulkan/dequant.comp @@ -0,0 +1,260 @@ +#version 460 + +// From omarchy-mlx overlay/mlx/backend/omarchy/shaders/dequant.comp at fde15c1ece69: + +// code unchanged, comments condensed. + +// Copyright © 2026 Joshua Warren; MIT License, full text in LICENSES/mlx-omarchy-MIT.txt. + +// Affine dequant: out = q * scale + bias; one thread owns one packed unit along the last axis. + +// bits 2/4/8: uint32 words, 32/bits values low-bits-first; bits 3/5/6: byte packs via BYTE_AT. + +// FP_MODE: e8m0/e4m3 scale bytes; fp4 e2m1 nibbles or fp8 e4m3 bytes (Metal fp4.h/fp8.h). + +// Honeykrisp: straight-line per-thread unit load; no lane-divergent word loads. + +#ifdef USE_FP16 +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#define STORAGE_TYPE float16_t +#define LOAD_VALUE(x) float(x) +#define STORE_VALUE(f) STORAGE_TYPE(f) +#elif defined(USE_BF16) +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require +#define STORAGE_TYPE uint16_t +#define LOAD_VALUE(x) uintBitsToFloat(uint(x) << 16u) +uint16_t bf16_store(float value) { + uint bits = floatBitsToUint(value); + if (isnan(value)) { + return uint16_t((bits >> 16) | 0x40u); + } + return uint16_t((bits + 0x7fffu + ((bits >> 16) & 1u)) >> 16); +} +#define STORE_VALUE(f) bf16_store(f) +#else +#define STORAGE_TYPE float +#define LOAD_VALUE(x) float(x) +#define STORE_VALUE(f) STORAGE_TYPE(f) +#endif + +layout(local_size_x = 256) in; + +#ifdef FP_MODE +layout(set = 0, binding = 0, std430) readonly buffer InputW { + uint values[]; +} input_w; +// Scale bytes addressed through the word view. +layout(set = 0, binding = 1, std430) readonly buffer Scales { + uint values[]; +} scales_data; +// Optional nvfp4 global scale float32 word; dummy when unbound (flags bit 1). +layout(set = 0, binding = 2, std430) readonly buffer GlobalScale { + uint values[]; +} global_scale_data; +layout(set = 0, binding = 3, std430) writeonly buffer Output { + STORAGE_TYPE values[]; +} output_data; +#else +layout(set = 0, binding = 0, std430) readonly buffer InputW { + uint values[]; +} input_w; +layout(set = 0, binding = 1, std430) readonly buffer Scales { + STORAGE_TYPE values[]; +} scales_data; +layout(set = 0, binding = 2, std430) readonly buffer Biases { + STORAGE_TYPE values[]; +} biases_data; +layout(set = 0, binding = 3, std430) writeonly buffer Output { + STORAGE_TYPE values[]; +} output_data; +#endif + +layout(push_constant) uniform Params { + uint count; + uint operation; + uint lhs_size; + uint rhs_size; + uint reduce_size; + uint output_size; + uint lhs_offset; + uint rhs_offset; + uint output_offset; + uint aux_size; + uint aux_offset; + uint matrix_m; + uint matrix_n; + uint matrix_k; + uint flags; + float alpha; + float beta; + uint dims; + uint shape[4]; + uint in_strides[4]; + uint out_strides[4]; +} params; + +// FP_MODE push: operation=bits, reduce_size=group, matrix_n=packs/row, matrix_k=groups/row. + +// aux_offset=scale-byte start, rhs_size=global-scale word, flags bit1=has global scale. + +#ifdef FP_MODE +float fp8_decode(uint bits) { + float magnitude = unpackHalf2x16((bits & 127u) << 7u).x * 256.0; + return ((bits & 128u) != 0u) ? -magnitude : magnitude; +} + +float e8m0_decode(uint bits) { + return uintBitsToFloat(bits == 0u ? 0x400000u : bits << 23u); +} + +// FP4 E2M1 via Metal fp4_e2m1 fp16 trick: (bits&7)<<9 subnormal, *16384 for 2^14. +float fp4_decode(uint code) { + float magnitude = unpackHalf2x16((code & 7u) << 9u).x * 16384.0; + return ((code & 8u) != 0u) ? -magnitude : magnitude; +} +#endif + +// Constant-shift select chain from scatter.comp: runtime word/lane, constant shift per lane. +uint byte_at(uint word, uint lane) { + return (lane == 0u) + ? (word & 0xFFu) + : (lane == 1u) + ? ((word >> 8u) & 0xFFu) + : (lane == 2u) + ? ((word >> 16u) & 0xFFu) + : ((word >> 24u) & 0xFFu); +} + +void main() { + uint bits = params.operation; + uint group_size = params.reduce_size; + uint stride = gl_NumWorkGroups.x * 256u; +#ifdef FP_MODE + uint pack_factor = (bits == 8u) ? 1u : 2u; + bool has_global_scale = (params.flags & 2u) != 0u; + float global_dec = 1.0; + if (has_global_scale) { + global_dec = + uintBitsToFloat(global_scale_data.values[params.rhs_size]) / 2688.0; + } + for (uint unit = gl_GlobalInvocationID.x; unit < params.count; + unit += stride) { + uint row = unit / params.matrix_n; + uint unit_column = unit % params.matrix_n; + uint group = (unit_column * pack_factor) / group_size; + uint scale_addr = params.aux_offset + row * params.matrix_k + group; + uint scale_byte = byte_at( + scales_data.values[scale_addr >> 2u], scale_addr & 3u); + float scale = (group_size == 16u) ? fp8_decode(scale_byte) + : e8m0_decode(scale_byte); + if (has_global_scale) { + scale *= global_dec; + } + uint packed_addr = params.lhs_offset + unit; + uint packed_byte = byte_at( + input_w.values[packed_addr >> 2u], packed_addr & 3u); + uint value_base = params.output_offset + unit * pack_factor; + if (bits == 8u) { + output_data.values[value_base] = + STORE_VALUE(scale * fp8_decode(packed_byte)); + } else { + output_data.values[value_base] = STORE_VALUE( + scale * fp4_decode(packed_byte & 0xFu)); + output_data.values[value_base + 1u] = STORE_VALUE( + scale * fp4_decode(packed_byte >> 4u)); + } + } +#else + uint is_byte_bits = (bits == 3u || bits == 5u || bits == 6u) ? 1u : 0u; + uint pack_factor = is_byte_bits == 1u + ? (bits == 6u ? 4u : 8u) + : (32u / bits); + uint bytes_per_pack = is_byte_bits == 1u ? (bits == 5u ? 5u : 3u) : 1u; + uint value_mask = (1u << bits) - 1u; + uint units_per_group = group_size / pack_factor; + + for (uint unit = gl_GlobalInvocationID.x; unit < params.count; + unit += stride) { + uint row = unit / params.matrix_n; + uint unit_column = unit % params.matrix_n; + uint group = unit_column / units_per_group; + float scale = LOAD_VALUE( + scales_data.values[params.rhs_offset + row * params.matrix_k + + group]); + float bias = LOAD_VALUE( + biases_data.values[params.aux_offset + row * params.matrix_k + + group]); + uint value_base = params.output_offset + unit * pack_factor; + if (is_byte_bits == 1u) { +// Byte pack: BYTE_AT bytes then upstream per-element extract masks (compile-time constants). + uint pack_base = params.lhs_offset + unit * bytes_per_pack; + uint c0 = byte_at(input_w.values[pack_base >> 2u], pack_base & 3u); + uint c1 = byte_at( + input_w.values[(pack_base + 1u) >> 2u], (pack_base + 1u) & 3u); + uint c2 = byte_at( + input_w.values[(pack_base + 2u) >> 2u], (pack_base + 2u) & 3u); + uint c3 = uint(0u); + uint c4 = uint(0u); + if (bytes_per_pack == 5u) { + c3 = byte_at( + input_w.values[(pack_base + 3u) >> 2u], (pack_base + 3u) & 3u); + c4 = byte_at( + input_w.values[(pack_base + 4u) >> 2u], (pack_base + 4u) & 3u); + } + if (bits == 3u) { + output_data.values[value_base] = STORE_VALUE( + scale * float(c0 & 0x7u) + bias); + output_data.values[value_base + 1u] = STORE_VALUE( + scale * float((c0 & 0x38u) >> 3) + bias); + output_data.values[value_base + 2u] = STORE_VALUE( + scale * float(((c0 & 0xc0u) >> 6) + ((c1 & 0x1u) << 2)) + bias); + output_data.values[value_base + 3u] = STORE_VALUE( + scale * float((c1 & 0xeu) >> 1) + bias); + output_data.values[value_base + 4u] = STORE_VALUE( + scale * float((c1 & 0x70u) >> 4) + bias); + output_data.values[value_base + 5u] = STORE_VALUE( + scale * float(((c1 & 0x80u) >> 7) + ((c2 & 0x3u) << 1)) + bias); + output_data.values[value_base + 6u] = STORE_VALUE( + scale * float((c2 & 0x1cu) >> 2) + bias); + output_data.values[value_base + 7u] = STORE_VALUE( + scale * float((c2 & 0xe0u) >> 5) + bias); + } else if (bits == 5u) { + output_data.values[value_base] = STORE_VALUE( + scale * float(c0 & 0x1fu) + bias); + output_data.values[value_base + 1u] = STORE_VALUE( + scale * float(((c0 & 0xe0u) >> 5) + ((c1 & 0x3u) << 3)) + bias); + output_data.values[value_base + 2u] = STORE_VALUE( + scale * float((c1 & 0x7cu) >> 2) + bias); + output_data.values[value_base + 3u] = STORE_VALUE( + scale * float(((c1 & 0x80u) >> 7) + ((c2 & 0xfu) << 1)) + bias); + output_data.values[value_base + 4u] = STORE_VALUE( + scale * float(((c2 & 0xf0u) >> 4) + ((c3 & 0x1u) << 4)) + bias); + output_data.values[value_base + 5u] = STORE_VALUE( + scale * float((c3 & 0x3eu) >> 1) + bias); + output_data.values[value_base + 6u] = STORE_VALUE( + scale * float(((c3 & 0xc0u) >> 6) + ((c4 & 0x7u) << 2)) + bias); + output_data.values[value_base + 7u] = STORE_VALUE( + scale * float((c4 & 0xf8u) >> 3) + bias); + } else { + output_data.values[value_base] = STORE_VALUE( + scale * float(c0 & 0x3fu) + bias); + output_data.values[value_base + 1u] = STORE_VALUE( + scale * float(((c0 >> 6) & 0x3u) + ((c1 & 0xfu) << 2)) + bias); + output_data.values[value_base + 2u] = STORE_VALUE( + scale * float(((c1 >> 4) & 0xfu) + ((c2 & 0x3u) << 4)) + bias); + output_data.values[value_base + 3u] = STORE_VALUE( + scale * float((c2 >> 2) & 0x3fu) + bias); + } + } else { + uint word = input_w.values[params.lhs_offset + unit]; + for (uint lane = 0u; lane < pack_factor; ++lane) { + uint q = (word >> (lane * bits)) & value_mask; + output_data.values[value_base + lane] = + STORE_VALUE(scale * float(q) + bias); + } + } + } +#endif +} diff --git a/zig/kernels/vulkan/fast_norm.comp b/zig/kernels/vulkan/fast_norm.comp new file mode 100644 index 00000000000..88a90fd6e15 --- /dev/null +++ b/zig/kernels/vulkan/fast_norm.comp @@ -0,0 +1,164 @@ +// From omarchy-mlx overlay/mlx/backend/omarchy/shaders/fast_norm.comp at 8d9aaa2db04a. +// Code adapted: 2 lines added or changed, 0 removed. +// Copyright © 2026 Joshua Warren and mlx-omarchy contributors; MIT License, full text in LICENSES/mlx-omarchy-MIT.txt. + +#version 460 + +#ifdef USE_FP16 +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#define STORAGE_TYPE float16_t +#define LOAD_VALUE(x) float(x) +#define STORE_VALUE(f) STORAGE_TYPE(f) +#elif defined(USE_BF16) +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require +#define STORAGE_TYPE uint16_t +#define LOAD_VALUE(x) uintBitsToFloat(uint(x) << 16u) +uint16_t bf16_store(float value) { + uint bits = floatBitsToUint(value); + if (isnan(value)) { + return uint16_t((bits >> 16) | 0x40u); + } + return uint16_t((bits + 0x7fffu + ((bits >> 16) & 1u)) >> 16); +} +#define STORE_VALUE(f) bf16_store(f) +#else +#define STORAGE_TYPE float +#define LOAD_VALUE(x) float(x) +#define STORE_VALUE(f) STORAGE_TYPE(f) +#endif + +// ACT_IN_F32/ACT_OUT_F32/ACT_OUT_BF16: the norm's activation in/out while the gain stays bf16. +#ifdef ACT_IN_F32 +#define ACT_IN float +#define ACT_IN_LOAD(x) (x) +#else +#define ACT_IN STORAGE_TYPE +#define ACT_IN_LOAD(x) LOAD_VALUE(x) +#endif +#ifdef ACT_OUT_F32 +#define ACT_OUT float +#define ACT_OUT_STORE(f) (f) +#elif defined(ACT_OUT_BF16) +#define ACT_OUT uint16_t +#define ACT_OUT_STORE(f) bf16_store(f) +#else +#define ACT_OUT STORAGE_TYPE +#define ACT_OUT_STORE(f) STORE_VALUE(f) +#endif + +layout(local_size_x = 256) in; + +layout(set = 0, binding = 0, std430) readonly buffer Input { + ACT_IN values[]; +} input_data; +layout(set = 0, binding = 1, std430) readonly buffer Weight { + STORAGE_TYPE values[]; +} weight_data; +#ifdef LAYER_NORM +layout(set = 0, binding = 2, std430) readonly buffer Bias { + STORAGE_TYPE values[]; +} bias_data; +#endif +layout(set = 0, binding = 3, std430) writeonly buffer Output { + ACT_OUT values[]; +} output_data; + +layout(push_constant) uniform Params { + uint count; + uint operation; + uint lhs_size; + uint rhs_size; + uint reduce_size; + uint output_size; + uint lhs_offset; + uint rhs_offset; + uint output_offset; + uint aux_size; + uint aux_offset; + uint matrix_m; + uint matrix_n; + uint matrix_k; + uint flags; + float alpha; + float beta; +} params; + +shared float block_values[256]; + +void main() { + uint thread = gl_LocalInvocationID.x; + uint row_length = params.reduce_size; + float eps = params.alpha; + uint threads = 256u; + + for (uint row = gl_WorkGroupID.x; row < params.output_size; + row += gl_NumWorkGroups.x) { + uint start = params.lhs_offset + row * row_length; + +#ifdef LAYER_NORM + float local_sum = 0.0; + for (uint index = thread; index < row_length; index += threads) { + local_sum += ACT_IN_LOAD(input_data.values[start + index]); + } + block_values[thread] = local_sum; + barrier(); + for (uint stride = 128u; stride > 0u; stride >>= 1u) { + if (thread < stride) { + block_values[thread] += block_values[thread + stride]; + } + barrier(); + } + float mu = block_values[0] / float(row_length); + barrier(); +#endif + + float local_sq = 0.0; + for (uint index = thread; index < row_length; index += threads) { + float value = ACT_IN_LOAD(input_data.values[start + index]); + local_sq += value * value; + } + block_values[thread] = local_sq; + barrier(); + for (uint stride = 128u; stride > 0u; stride >>= 1u) { + if (thread < stride) { + block_values[thread] += block_values[thread + stride]; + } + barrier(); + } + float mean_square = block_values[0] / float(row_length); + barrier(); +#ifdef LAYER_NORM + float variance = mean_square - mu * mu; +#else + float variance = mean_square; +#endif + float norm = inversesqrt(variance + eps); + + for (uint index = thread; index < row_length; index += threads) { + float value = ACT_IN_LOAD(input_data.values[start + index]); + float weight = + (params.lhs_size == 1u) + ? LOAD_VALUE(weight_data.values[params.rhs_offset]) + : LOAD_VALUE(weight_data.values[params.rhs_offset + index]); +#ifdef LAYER_NORM + float bias = + (params.aux_size == 1u) + ? LOAD_VALUE(bias_data.values[params.aux_offset]) + : LOAD_VALUE(bias_data.values[params.aux_offset + index]); + float result = (value - mu) * norm * weight + bias; +#else +#ifdef ORACLE_QWEN35 + // Ash8c Metal: bfloat(Wt*(hv*inv)); weight multiplies the product; one store rounding. + float result = weight * (value * norm); +#else + float result = value * norm * weight; +#endif +#endif + output_data.values[params.output_offset + row * row_length + index] = + ACT_OUT_STORE(result); + } + barrier(); + } +} diff --git a/zig/kernels/vulkan/fast_norm_gated.comp b/zig/kernels/vulkan/fast_norm_gated.comp new file mode 100644 index 00000000000..4c55d071103 --- /dev/null +++ b/zig/kernels/vulkan/fast_norm_gated.comp @@ -0,0 +1,277 @@ +// From omarchy-mlx overlay/mlx/backend/omarchy/shaders/fast_norm_gated.comp at 69ab4cfc3f6c. +// Code adapted: 4 lines added or changed, 0 removed. +// Copyright © 2026 Joshua Warren and mlx-omarchy contributors; MIT License, full text in LICENSES/mlx-omarchy-MIT.txt. +#version 460 +#ifdef APPLE_REDUCTION +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_shader_subgroup_arithmetic : require +#endif + +#ifdef USE_FP16 +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#define STORAGE_TYPE float16_t +#define LOAD_VALUE(x) float(x) +#define STORE_VALUE(f) STORAGE_TYPE(f) +#define WIDEN(x) float(x) +#elif defined(USE_BF16) +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require +#define STORAGE_TYPE uint16_t +#define LOAD_VALUE(x) uintBitsToFloat(uint(x) << 16u) +uint16_t bf16_store(float value) { + uint bits = floatBitsToUint(value); + if (isnan(value)) { + return uint16_t((bits >> 16) | 0x40u); + } + return uint16_t((bits + 0x7fffu + ((bits >> 16) & 1u)) >> 16); +} +// Exact bf16 -> float32 widening of a stored value. +#define WIDEN(x) (uintBitsToFloat(uint(x) << 16u)) +#define STORE_VALUE(f) bf16_store(f) +#else +#define STORAGE_TYPE float +#define LOAD_VALUE(x) float(x) +#define STORE_VALUE(f) STORAGE_TYPE(f) +#define WIDEN(x) float(x) +#endif + +// ACT_F32: the gated norm's y/z input and mix output are f32 while the gain stays bf16. +#ifdef ACT_F32 +#define ACT_TYPE float +#define ACT_LOAD(x) (x) +#define ACT_STORE(f) (f) +#define ACT_WIDEN(x) (x) +#define ACT_NORMED(f) (f) +#else +#define ACT_TYPE STORAGE_TYPE +#define ACT_LOAD(x) LOAD_VALUE(x) +#define ACT_STORE(f) STORE_VALUE(f) +#define ACT_WIDEN(x) WIDEN(x) +#define ACT_NORMED(f) WIDEN(STORE_VALUE(f)) +#endif + +// Preserve the omitted contract, precision, and host-interface facts from the original comment. +#ifdef GATED_ONLY +#define IS_SCALED (false) +#define IS_GATED (true) +#else +#define IS_SCALED (params.operation == 1u) +#define IS_GATED (params.operation == 0u) +#endif + +layout(local_size_x = 256) in; + +layout(set = 0, binding = 0, std430) readonly buffer Input { + ACT_TYPE values[]; +} input_data; +layout(set = 0, binding = 1, std430) readonly buffer Weight { + STORAGE_TYPE values[]; +} weight_data; +layout(set = 0, binding = 2, std430) readonly buffer Gate { + ACT_TYPE values[]; +} gate_data; +layout(set = 0, binding = 3, std430) writeonly buffer Output { + ACT_TYPE values[]; +} output_data; + +layout(push_constant) uniform Params { + uint count; + uint operation; + uint lhs_size; + uint rhs_size; + uint reduce_size; + uint output_size; + uint lhs_offset; + uint rhs_offset; + uint output_offset; + uint aux_size; + uint aux_offset; + uint matrix_m; + uint matrix_n; + uint matrix_k; + uint flags; + float alpha; + float beta; +} params; + +#ifdef APPLE_REDUCTION +shared float subgroup_values[8]; +shared float apple_norm; + +void store_normed(uint row, uint row_length, uint index, float value, + float norm, float scale) { + float weight = + (params.lhs_size == 1u) + ? LOAD_VALUE(weight_data.values[params.rhs_offset]) + : LOAD_VALUE(weight_data.values[params.rhs_offset + index]); +#ifdef ORACLE_QWEN35 + // Ash8c Metal oracle (gdn_post.metal): BF16(weight*(value*norm)); gate is BF16((gate/(1+exp(-gate)))*normed). + precise float scaled = value * norm; + precise float result = weight * scaled; + float normed = ACT_NORMED(result); + float out_value; + if (IS_GATED) { + uint gate_index = params.aux_offset + row * row_length + index; + float gate = ACT_WIDEN(gate_data.values[gate_index]); + precise float silu = gate / (1.0f + exp(-gate)); + precise float prod = silu * normed; + out_value = prod; + } else { + out_value = normed * scale; + } +#else + precise float norm_weight = norm * weight; + precise float result = value * norm_weight; + float normed = ACT_NORMED(result); + float out_value; + if (IS_GATED) { + uint gate_index = params.aux_offset + row * row_length + index; + float gate = ACT_WIDEN(gate_data.values[gate_index]); + float sigmoid_value = 1.0 / (1.0 + exp(-gate)); + precise float silu = gate * sigmoid_value; + precise float prod = silu * normed; + out_value = prod; + } else { + out_value = normed * scale; + } +#endif + output_data.values[params.output_offset + row * row_length + index] = + ACT_STORE(out_value); +} +#else +shared float block_values[256]; +#endif + +void main() { + uint thread = gl_LocalInvocationID.x; + uint row_length = params.reduce_size; + float eps = params.alpha; + uint threads = 256u; + + // Mode-1 scalar: round the f32 scale to the storage dtype exactly as the composed promote cast does, then widen. + float scale = 0.0; + if (IS_SCALED) { +#ifdef USE_BF16 + scale = WIDEN(bf16_store(params.beta)); +#else + scale = params.beta; +#endif + } + + for (uint row = gl_WorkGroupID.x; row < params.output_size; + row += gl_NumWorkGroups.x) { + uint start = params.lhs_offset + row * row_length; + +#ifdef APPLE_REDUCTION + bool contiguous = row_length == 8u * threads; + float cached_x[8]; + float local_sq = 0.0; + if (contiguous) { + for (uint read = 0u; read < 8u; ++read) { + float value = ACT_LOAD(input_data.values[start + thread * 8u + read]); + cached_x[read] = value; + precise float square = value * value; + local_sq += square; + } + } else { + for (uint index = thread; index < row_length; index += threads) { + float value = ACT_LOAD(input_data.values[start + index]); + precise float square = value * value; + local_sq += square; + } + } + float subgroup_sum = subgroupAdd(local_sq); + if (gl_SubgroupInvocationID == 0u) { + subgroup_values[gl_SubgroupID] = subgroup_sum; + } + barrier(); + float partial = gl_SubgroupID == 0u && gl_SubgroupInvocationID < 8u + ? subgroup_values[gl_SubgroupInvocationID] + : 0.0; + float row_sum = subgroupAdd(partial); + if (gl_LocalInvocationIndex == 0u) { + apple_norm = inversesqrt(row_sum / float(row_length) + eps); + } + barrier(); + float norm = apple_norm; + barrier(); + if (contiguous) { + for (uint read = 0u; read < 8u; ++read) { + store_normed(row, row_length, thread * 8u + read, cached_x[read], + norm, scale); + } + } else { + for (uint index = thread; index < row_length; index += threads) { + store_normed(row, row_length, index, + ACT_LOAD(input_data.values[start + index]), norm, + scale); + } + } +#else + float local_sq = 0.0; + for (uint index = thread; index < row_length; index += threads) { + float value = ACT_LOAD(input_data.values[start + index]); + local_sq += value * value; + } + block_values[thread] = local_sq; + barrier(); + for (uint stride = 128u; stride > 0u; stride >>= 1u) { + if (thread < stride) { + block_values[thread] += block_values[thread + stride]; + } + barrier(); + } + float mean_square = block_values[0] / float(row_length); + barrier(); + float norm = inversesqrt(mean_square + eps); + + for (uint index = thread; index < row_length; index += threads) { + float value = ACT_LOAD(input_data.values[start + index]); + float weight = + (params.lhs_size == 1u) + ? LOAD_VALUE(weight_data.values[params.rhs_offset]) + : LOAD_VALUE(weight_data.values[params.rhs_offset + index]); +#ifdef ORACLE_QWEN35 + // Ash8c Metal gdn_post keeps value*norm then weight; its gate uses division, as in store_normed. + precise float scaled = value * norm; + precise float result = weight * scaled; + // FastRmsNormBF16 -> FusedChainF32(sigmoid,mul,mul) -> CastF32BF16 rounds normed once, then widens exactly. + float normed = ACT_NORMED(result); + float out_value; + if (IS_GATED) { + uint gate_index = params.aux_offset + row * row_length + index; + float gate = ACT_WIDEN(gate_data.values[gate_index]); + precise float silu = gate / (1.0f + exp(-gate)); + precise float prod = silu * normed; + out_value = prod; + } else { + out_value = normed * scale; + } +#else + // Reference uses value*(norm*weight); ~5e-6 outputs differed 1 bf16 ulp; precise: 0/5.4M mismatches; source 25. + precise float norm_weight = norm * weight; + precise float result = value * norm_weight; + // The composed path rounds the normed value once to storage, then widens it exactly. + float normed = ACT_NORMED(result); + float out_value; + if (IS_GATED) { + uint gate_index = params.aux_offset + row * row_length + index; + float gate = ACT_WIDEN(gate_data.values[gate_index]); + // Sigmoid and silu stay f32 until the final multiply; rounding them to bf16 caused ~36% one-ulp differences. + float sigmoid_value = 1.0 / (1.0 + exp(-gate)); + // Keep (gate*sigmoid)*normed as the oracle does; reassociating it left outputs 7e-6 off. + precise float silu = gate * sigmoid_value; + precise float prod = silu * normed; + out_value = prod; + } else { + out_value = normed * scale; + } +#endif + output_data.values[params.output_offset + row * row_length + index] = + ACT_STORE(out_value); + } + barrier(); +#endif + } +} diff --git a/zig/kernels/vulkan/fast_rope.comp b/zig/kernels/vulkan/fast_rope.comp new file mode 100644 index 00000000000..9026d9b283e --- /dev/null +++ b/zig/kernels/vulkan/fast_rope.comp @@ -0,0 +1,438 @@ +// From omarchy-mlx overlay/mlx/backend/omarchy/shaders/fast_rope.comp at de34407c136e. +// Code adapted: 7 lines added or changed, 0 removed. +// Copyright © 2026 Joshua Warren and mlx-omarchy contributors; MIT License, full text in LICENSES/mlx-omarchy-MIT.txt. +#version 460 + +#extension GL_GOOGLE_include_directive : require +#include "omarchy_trig.h" + +#ifdef USE_FP16 +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#define STORAGE_TYPE float16_t +layout(set = 0, binding = 0, std430) readonly buffer Input { + STORAGE_TYPE values[]; +} input_data; +layout(set = 0, binding = 1, std430) writeonly buffer Output { + STORAGE_TYPE values[]; +} output_data; +float rope_load(uint abs_index) { + return float(input_data.values[abs_index]); +} +void rope_store(uint abs_index, float value) { + output_data.values[abs_index] = float16_t(value); +} +float rope_round(float value) { + // Match eager storage rounding after each operation; preserve float16 output bits. + return float(float16_t(value)); +} +#elif defined(USE_BF16) +#extension GL_EXT_shader_16bit_storage : require +layout(set = 0, binding = 0, std430) readonly buffer Input { + uint values[]; +} input_data; +layout(set = 0, binding = 1, std430) buffer Output { + uint values[]; +} output_data; +float rope_load(uint abs_index) { + uint word = input_data.values[abs_index >> 1u]; + uint bits = ((abs_index & 1u) == 0u) ? (word & 0xFFFFu) : (word >> 16u); + return uintBitsToFloat(bits << 16u); +} +uint bf16_bits(float value) { + uint bits = floatBitsToUint(value); + return isnan(value) + ? ((bits >> 16u) | 0x40u) + : ((bits + 0x7fffu + ((bits >> 16u) & 1u)) >> 16u); +} +// ORACLE_QWEN35 matches Metal's float32 angle/products and performs one final bfloat16 RNE store. +float rope_round(float value) { +#ifdef ORACLE_QWEN35 + return value; +#else + uint bits = floatBitsToUint(value); + uint rounded = isnan(value) + ? ((bits >> 16u) | 0x40u) + : ((bits + 0x7fffu + ((bits >> 16u) & 1u)) >> 16u); + return uintBitsToFloat(rounded << 16u); +#endif +} +#else +#define STORAGE_TYPE float +layout(set = 0, binding = 0, std430) readonly buffer Input { + STORAGE_TYPE values[]; +} input_data; +float rope_load(uint abs_index) { + return input_data.values[abs_index]; +} +#ifdef ACT_OUT_BF16 +// F32 rope with a bf16 cache write: element addressing, RNE store. +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require +layout(set = 0, binding = 1, std430) writeonly buffer Output { + uint16_t values[]; +} output_data; +uint16_t bf16_store(float value) { + uint bits = floatBitsToUint(value); + return isnan(value) + ? uint16_t((bits >> 16u) | 0x40u) + : uint16_t((bits + 0x7fffu + ((bits >> 16u) & 1u)) >> 16u); +} +void rope_store(uint abs_index, float value) { + output_data.values[abs_index] = bf16_store(value); +} +#else +layout(set = 0, binding = 1, std430) writeonly buffer Output { + STORAGE_TYPE values[]; +} output_data; +void rope_store(uint abs_index, float value) { + output_data.values[abs_index] = value; +} +#endif +float rope_round(float value) { + return value; +} +#endif + +layout(set = 0, binding = 2, std430) readonly buffer Offsets { + int values[]; +} offset_data; +#ifdef USE_FREQS +layout(set = 0, binding = 3, std430) readonly buffer Freqs { + float values[]; +} freqs_data; +#endif + +// ORACLE_QWEN35 uses Metal's pow(10000000.0f, -i/32) expression with i < 32. +float rope_inv_freq(uint i) { +#if defined(ORACLE_QWEN35) + return pow(1.0e7f, -float(i) / 32.0f); +#elif defined(USE_FREQS) + return 1.0f / freqs_data.values[params.aux_offset + i]; +#else + return exp(-float(i) * params.beta); +#endif +} + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout(push_constant) uniform Params { + uint count; + uint operation; + uint lhs_size; + uint rhs_size; + uint reduce_size; + uint output_size; + uint lhs_offset; + uint rhs_offset; + uint output_offset; + uint aux_size; + uint aux_offset; + uint matrix_m; + uint matrix_n; + uint matrix_k; + uint flags; + float alpha; + float beta; + uint dims; + uint shape[4]; + uint in_strides[4]; + uint out_strides[4]; +} params; + +void main() { +#ifdef USE_BF16 + uint index = gl_GlobalInvocationID.x; + uint stride = gl_NumWorkGroups.x * gl_WorkGroupSize.x; + uint half_dims = params.dims; + uint time = params.shape[1]; + uint heads = params.shape[0]; + bool forward = (params.flags & 1u) != 0u; + bool traditional = (params.flags & 2u) != 0u; + bool head_seq_transpose = (params.flags & 4u) != 0u; + bool passthrough = (params.flags & 8u) != 0u; + while (index < params.count) { + uint k; + uint t; + uint mat; + // Half-split rotation assigns D/2 words to each row. + if (!passthrough && !traditional) { + { + uint row_words = half_dims; + k = index % row_words; + t = (index / row_words) % time; + mat = index / (row_words * time); + uint b = mat / heads; + uint head = mat - b * heads; + float offset = float( + offset_data.values[params.rhs_offset + b * params.matrix_m]); + float position = params.alpha * (float(t) + offset); + // Each k maps two adjacent r1 elements or the mirrored r2 pair. + bool r1_side = k < (half_dims >> 1u); + // On r2, theta and operands use the mirrored i = element - half_dims. + uint j0 = r1_side ? (2u * k) + : (half_dims + 2u * (k - (half_dims >> 1u))); + uint j1 = j0 + 1u; + uint i0 = r1_side ? j0 : (j0 - half_dims); + uint i1 = i0 + 1u; +#ifdef USE_FREQS + float inv0 = 1.0f / freqs_data.values[params.aux_offset + i0]; + float inv1 = 1.0f / freqs_data.values[params.aux_offset + i1]; +#else + float inv0 = exp(-float(i0) * params.beta); + float inv1 = exp(-float(i1) * params.beta); +#endif + float cos0 = rope_round(omarchy_trig_cos(position * inv0)); + float sin0 = rope_round(omarchy_trig_sin(position * inv0)); + float cos1 = rope_round(omarchy_trig_cos(position * inv1)); + float sin1 = rope_round(omarchy_trig_sin(position * inv1)); + uint in_base; + if (head_seq_transpose) { + in_base = b * (time * params.in_strides[1]) + + t * params.in_strides[1] + head * params.in_strides[0]; + } else { + in_base = t * params.in_strides[1] + mat * params.in_strides[0]; + } + uint out_row = + t * params.out_strides[1] + mat * params.out_strides[0]; + uint x1a = in_base + i0 * params.in_strides[2]; + uint x2a = x1a + half_dims * params.in_strides[2]; + uint x1b = in_base + i1 * params.in_strides[2]; + uint x2b = x1b + half_dims * params.in_strides[2]; + float xa = rope_load(params.lhs_offset + x1a); + float xb = rope_load(params.lhs_offset + x2a); + float xc = rope_load(params.lhs_offset + x1b); + float xd = rope_load(params.lhs_offset + x2b); + float v0; + float v1; + if (r1_side) { + v0 = forward ? (rope_round(xa * cos0) - rope_round(xb * sin0)) + : (rope_round(xb * sin0) + rope_round(xa * cos0)); + v1 = forward ? (rope_round(xc * cos1) - rope_round(xd * sin1)) + : (rope_round(xd * sin1) + rope_round(xc * cos1)); + } else { + v0 = forward ? (rope_round(xa * sin0) + rope_round(xb * cos0)) + : (rope_round(xb * cos0) - rope_round(xa * sin0)); + v1 = forward ? (rope_round(xc * sin1) + rope_round(xd * cos1)) + : (rope_round(xd * cos1) - rope_round(xc * sin1)); + } + uint word_base = (params.output_offset + out_row) >> 1u; + output_data.values[word_base + k] = + bf16_bits(v0) | (bf16_bits(v1) << 16u); + } + } else { + // Traditional and passthrough pairs each occupy one output word. + uint row_words = passthrough ? (params.shape[2] >> 1u) : half_dims; + uint k = index % row_words; + t = (index / row_words) % time; + mat = index / (row_words * time); + uint b = mat / heads; + uint head = mat - b * heads; + float offset = float( + offset_data.values[params.rhs_offset + b * params.matrix_m]); + float position = params.alpha * (float(t) + offset); + uint in_base; + if (head_seq_transpose) { + in_base = b * (time * params.in_strides[1]) + + t * params.in_strides[1] + head * params.in_strides[0]; + } else { + in_base = t * params.in_strides[1] + mat * params.in_strides[0]; + } + uint out_row = t * params.out_strides[1] + mat * params.out_strides[0]; + uint word_base = (params.output_offset + out_row) >> 1u; + if (!passthrough) { + // Traditional: element 2i = r1, element 2i+1 = r2, one theta. + uint i = k; +#ifdef USE_FREQS + float inv_freq = 1.0f / freqs_data.values[params.aux_offset + i]; +#else + float inv_freq = exp(-float(i) * params.beta); +#endif + float theta = position * inv_freq; + float cos_t = rope_round(omarchy_trig_cos(theta)); + float sin_t = rope_round(omarchy_trig_sin(theta)); + uint in_row = in_base; + float x1 = rope_load( + params.lhs_offset + in_row + (2u * i) * params.in_strides[2]); + float x2 = rope_load( + params.lhs_offset + in_row + (2u * i + 1u) * params.in_strides[2]); + float xa = rope_round(x1 * cos_t); + float xb = rope_round(x2 * sin_t); + float xc = rope_round(x1 * sin_t); + float xd = rope_round(x2 * cos_t); + float r1 = forward ? (xa - xb) : (xb + xa); + float r2 = forward ? (xc + xd) : (xd - xa); + output_data.values[word_base + k] = + bf16_bits(r1) | (bf16_bits(r2) << 16u); + } else { + // Passthrough pair: two adjacent d, both rotated (d < dims) or both tail. + uint d0 = 2u * k; + uint d1 = d0 + 1u; + float v0; + float v1; + if (d0 >= 2u * half_dims) { + v0 = rope_load( + params.lhs_offset + in_base + d0 * params.in_strides[2]); + v1 = rope_load( + params.lhs_offset + in_base + d1 * params.in_strides[2]); + } else { + float vals[2]; + for (uint side = 0u; side < 2u; ++side) { + uint d = d0 + side; + uint i = d % half_dims; +#ifdef USE_FREQS + float inv_freq = 1.0f / freqs_data.values[params.aux_offset + i]; +#else + float inv_freq = exp(-float(i) * params.beta); +#endif + float theta = position * inv_freq; + float cos_t = rope_round(omarchy_trig_cos(theta)); + float sin_t = rope_round(omarchy_trig_sin(theta)); + uint x1_index = in_base + i * params.in_strides[2]; + uint x2_index = x1_index + half_dims * params.in_strides[2]; + float x1 = rope_load(params.lhs_offset + x1_index); + float x2 = rope_load(params.lhs_offset + x2_index); + float xa = rope_round(x1 * cos_t); + float xb = rope_round(x2 * sin_t); + float xc = rope_round(x1 * sin_t); + float xd = rope_round(x2 * cos_t); + vals[side] = (d < half_dims) ? (forward ? (xa - xb) : (xb + xa)) + : (forward ? (xc + xd) : (xd - xa)); + } + v0 = vals[0]; + v1 = vals[1]; + } + output_data.values[word_base + k] = + bf16_bits(v0) | (bf16_bits(v1) << 16u); + } + } + if (params.count - index <= stride) { + break; + } + index += stride; + } + return; +#else + uint index = gl_GlobalInvocationID.x; + uint stride = gl_NumWorkGroups.x * gl_WorkGroupSize.x; + uint half_dims = params.dims; + uint time = params.shape[1]; + uint heads = params.shape[0]; + bool forward = (params.flags & 1u) != 0u; + bool traditional = (params.flags & 2u) != 0u; + bool head_seq_transpose = (params.flags & 4u) != 0u; + bool passthrough = (params.flags & 8u) != 0u; + while (index < params.count) { + uint i; + uint t; + uint mat; + if (passthrough) { + uint d = index % params.shape[2]; + t = (index / params.shape[2]) % time; + mat = index / (params.shape[2] * time); + uint b = mat / heads; + uint head = mat - b * heads; + i = traditional ? (d >> 1u) : (d - (d / half_dims) * half_dims); + float offset = + float(offset_data.values[params.rhs_offset + b * params.matrix_m]); + float position = params.alpha * (float(t) + offset); +#ifdef USE_FREQS + float inv_freq = 1.0f / freqs_data.values[params.aux_offset + i]; +#else + float inv_freq = exp(-float(i) * params.beta); +#endif + float theta = position * inv_freq; + float cos_t = rope_round(omarchy_trig_cos(theta)); + float sin_t = rope_round(omarchy_trig_sin(theta)); + uint in_base; + if (head_seq_transpose) { + in_base = b * (time * params.in_strides[1]) + + t * params.in_strides[1] + head * params.in_strides[0]; + } else { + in_base = t * params.in_strides[1] + mat * params.in_strides[0]; + } + uint out_base = + t * params.out_strides[1] + mat * params.out_strides[0]; + uint x1_index = in_base + + (traditional ? (d & ~1u) : i) * params.in_strides[2]; + uint x2_index = traditional + ? (x1_index + params.in_strides[2]) + : (x1_index + half_dims * params.in_strides[2]); + uint out_index = out_base + d * params.out_strides[2]; + if (d >= 2u * half_dims) { + // Passthrough tail: the rotated range ends at dims; the rest of the axis carries the input through untouched. + rope_store( + params.output_offset + out_index, + rope_load(params.lhs_offset + in_base + d * params.in_strides[2])); + } else { + float x1 = rope_load(params.lhs_offset + x1_index); + float x2 = rope_load(params.lhs_offset + x2_index); + float xa = rope_round(x1 * cos_t); + float xb = rope_round(x2 * sin_t); + float xc = rope_round(x1 * sin_t); + float xd = rope_round(x2 * cos_t); + bool first = traditional ? ((d & 1u) == 0u) : (d < half_dims); + float r1 = forward ? (xa - xb) : (xb + xa); + float r2 = forward ? (xc + xd) : (xd - xa); + rope_store(params.output_offset + out_index, first ? r1 : r2); + } + } else { + i = index % half_dims; + t = (index / half_dims) % time; + mat = index / (half_dims * time); + uint b = mat / heads; + uint head = mat - b * heads; + float offset = + float(offset_data.values[params.rhs_offset + b * params.matrix_m]); + float position = params.alpha * (float(t) + offset); +#ifdef USE_FREQS + float inv_freq = 1.0f / freqs_data.values[params.aux_offset + i]; +#else + float inv_freq = exp(-float(i) * params.beta); +#endif + float theta = position * inv_freq; + float cos_t = rope_round(omarchy_trig_cos(theta)); + float sin_t = rope_round(omarchy_trig_sin(theta)); + uint in_base; + if (head_seq_transpose) { + in_base = b * (time * params.in_strides[1]) + + t * params.in_strides[1] + head * params.in_strides[0]; + } else { + in_base = t * params.in_strides[1] + mat * params.in_strides[0]; + } + uint out_base = + t * params.out_strides[1] + mat * params.out_strides[0]; + uint in_index_1; + uint in_index_2; + uint out_index_1; + // Match eager intermediate rounding; float32 storage is transparent. + uint out_index_2; + if (traditional) { + in_index_1 = in_base + (2u * i) * params.in_strides[2]; + in_index_2 = in_index_1 + params.in_strides[2]; + out_index_1 = out_base + (2u * i) * params.out_strides[2]; + out_index_2 = out_index_1 + params.out_strides[2]; + } else { + in_index_1 = in_base + i * params.in_strides[2]; + in_index_2 = in_index_1 + half_dims * params.in_strides[2]; + out_index_1 = out_base + i * params.out_strides[2]; + out_index_2 = out_index_1 + half_dims * params.out_strides[2]; + } + float x1 = rope_load(params.lhs_offset + in_index_1); + float x2 = rope_load(params.lhs_offset + in_index_2); + // Match eager storage-precision rounding after each operation; f32 storage is transparent. + float xa = rope_round(x1 * cos_t); + float xb = rope_round(x2 * sin_t); + float xc = rope_round(x1 * sin_t); + float xd = rope_round(x2 * cos_t); + float r1 = forward ? (xa - xb) : (xb + xa); + float r2 = forward ? (xc + xd) : (xd - xa); + rope_store(params.output_offset + out_index_1, r1); + rope_store(params.output_offset + out_index_2, r2); + } + if (params.count - index <= stride) { + break; + } + index += stride; + } +#endif +} diff --git a/zig/kernels/vulkan/finite_logits.comp b/zig/kernels/vulkan/finite_logits.comp new file mode 100644 index 00000000000..42dcd3732be --- /dev/null +++ b/zig/kernels/vulkan/finite_logits.comp @@ -0,0 +1,56 @@ +// Set a flag for any nonfinite BF16 input and record the first bad index. +#version 460 + +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require + +layout(local_size_x = 256) in; + +layout(set = 0, binding = 0, std430) readonly buffer Input { + uint16_t values[]; +} input_data; +layout(set = 0, binding = 1, std430) buffer Results { + uint finite_flag; + uint first_bad; +} results; + +layout(push_constant) uniform Params { + uint count; +} params; + +shared uint block_bad[256]; +shared uint block_flags[256]; + +void main() { + uint thread = gl_LocalInvocationID.x; + uint begin = gl_WorkGroupID.x * 4096u; + uint end = min(begin + 4096u, params.count); + // Bindings: 0 input BF16[count], 1 results 8B; push=count; grid=ceil(count/4096); local=256; shared=2x256 uints. + uint local_bad = 0xffffffffu; + uint local_flag = 0u; + for (uint k = 0u; k < 16u; ++k) { + uint index = begin + thread + k * 256u; + if (index < end) { + uint bits = uint(input_data.values[index]); + if ((bits & 0x7fffu) >= 0x7f80u) { + local_flag = 1u; + local_bad = min(local_bad, index); + } + } + } + block_bad[thread] = local_bad; + block_flags[thread] = local_flag; + barrier(); + for (uint stride = 128u; stride > 0u; stride >>= 1u) { + if (thread < stride) { + block_bad[thread] = min(block_bad[thread], block_bad[thread + stride]); + block_flags[thread] = block_flags[thread] | block_flags[thread + stride]; + } + barrier(); + } + if (thread == 0u) { + // Clean workgroups contribute (0, 0xffffffff): no-ops on both words. + atomicOr(results.finite_flag, block_flags[0]); + atomicMin(results.first_bad, block_bad[0]); + } +} diff --git a/zig/kernels/vulkan/gated_delta_decode.comp b/zig/kernels/vulkan/gated_delta_decode.comp new file mode 100644 index 00000000000..87fa374e4cf --- /dev/null +++ b/zig/kernels/vulkan/gated_delta_decode.comp @@ -0,0 +1,256 @@ +// From omarchy-mlx overlay/mlx/backend/omarchy/shaders/gated_delta_decode.comp at 16c800c7e89d. +// Code adapted: 5 lines added or changed, 3 removed. +// Copyright © 2026 Joshua Warren and mlx-omarchy contributors; MIT License, full text in LICENSES/mlx-omarchy-MIT.txt. +#version 460 + +#extension GL_EXT_shader_16bit_storage : require +// int16 is required: raw-gates slots read uint16_t A_log words in every build. +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require + +layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in; + +layout(set = 0, binding = 0, std430) readonly buffer QBuf { uint16_t v[]; } q_buf; +layout(set = 0, binding = 1, std430) readonly buffer KBuf { uint16_t v[]; } k_buf; +layout(set = 0, binding = 2, std430) readonly buffer VBuf { uint16_t v[]; } v_buf; +layout(set = 0, binding = 3, std430) readonly buffer GBuf { uint16_t v[]; } g_buf; +layout(set = 0, binding = 4, std430) readonly buffer BBuf { uint16_t v[]; } b_buf; +layout(set = 0, binding = 5, std430) readonly buffer SIn { float v[]; } s_in; +layout(set = 0, binding = 6, std430) writeonly buffer YBuf { uint16_t v[]; } y_buf; +layout(set = 0, binding = 7, std430) buffer SOut { float v[]; } s_out; +// Flag 3 raw-gates mode uses slots 3/4 for a/b and slots 8/9 for A_log/dt_bias (bf16). +layout(set = 0, binding = 8, std430) readonly buffer ALogBuf { uint16_t v[]; } alog_buf; +layout(set = 0, binding = 9, std430) readonly buffer DtBuf { uint16_t v[]; } dt_buf; + +layout(push_constant) uniform Params { + uint count; + uint operation; + uint lhs_size; + uint rhs_size; + uint reduce_size; + uint output_size; + uint lhs_offset; + uint rhs_offset; + uint output_offset; + uint aux_size; + uint aux_offset; + uint matrix_m; // Dk + uint matrix_n; // Dv + uint matrix_k; // Hv + uint flags; + float alpha; + float beta_f; + uint dims; + uint shape[4]; + uint in_strides[4]; + uint out_strides[4]; +} params; + +float bf16_load(uint16_t bits) { + return uintBitsToFloat(uint(bits) << 16u); +} + +uint16_t bf16_store(float value) { + uint bits = floatBitsToUint(value); + if (isnan(value)) { + return uint16_t((bits >> 16) | 0x40u); + } + return uint16_t((bits + 0x7fffu + ((bits >> 16) & 1u)) >> 16); +} + +const uint ROWS = 32u; +const uint TILE_STRIDE = 129u; // padded: row r, column i -> bank (r + i) % 32 + +shared float sk[128]; +shared float sq[128]; +shared float skg[128]; +shared float tile[ROWS * TILE_STRIDE]; + +// One decayed state row's k-weighted sum; mlx_order selects MLX's rounding order. +float gdn_kv_walk(uint trow, uint dk, float g, bool mlx_order) { + float kv_acc = 0.0f; + for (uint i = 0u; i < dk; ++i) { + if (mlx_order) { + precise float s_g = tile[trow + i] * g; + precise float term = s_g * sk[i]; + kv_acc += term; + } else { + kv_acc += tile[trow + i] * skg[i]; + } + } + return kv_acc; +} + +// ns = s * g + dk_term, contracted (dk_term arrives rounded). +float gdn_decay_add(float s, float g, float dk_term) { + return s * g + dk_term; +} + +#ifdef GDN_COMPOSED_ORDER +// Match elementwise.comp LogAddExp(0,x) without `precise`, preserving compiler choices. +float gdn_log1p_ew(float v) { + float u = 1.0f + v; + if (u == 1.0f) { + return v; + } + return log(u) * (v / (u - 1.0f)); +} + +float gdn_softplus_ew(float x) { + float value = max(0.0f, x); + if (!isnan(value) && !isinf(value)) { + value += gdn_log1p_ew(exp(-abs(0.0f - x))); + } + return value; +} + +// Decay gate g = exp(-exp(a_log) * softplus(x)), each step rounded separately. +float gdn_g_composed(float x, float a_log) { + float sp = gdn_softplus_ew(x); + precise float e = exp(a_log); + precise float p = e * sp; + precise float n = -p; + precise float g = exp(n); + return g; +} + + +#endif /* GDN_COMPOSED_ORDER */ +// Gate prologue (g, beta) for one head. +vec2 gdn_gates(uint head) { + float g; + float beta; + // g stays f32; eager-chain bf16 roundings apply to x, softplus, and beta. + if ((params.flags & 8u) != 0u) { + // Raw gates: beta=sigmoid(b) f32->bf16; g=exp(-exp(A_log_f32)*softplus(a+dt_bias)) f32->bf16. + float b_in = bf16_load(b_buf.v[params.aux_offset + head]); + float beta_f = 1.0f / (1.0f + exp(-b_in)); + beta = bf16_load(bf16_store(beta_f)); + float a_in = bf16_load(g_buf.v[params.shape[0] + head]); + // ph: head index within one batch row (in_strides[2] heads per row). + uint ph = head % params.in_strides[2]; + float a_log; + if ((params.flags & 512u) != 0u) { + uint base = params.in_strides[0] + 2u * ph; + a_log = uintBitsToFloat(uint(alog_buf.v[base]) | + (uint(alog_buf.v[base + 1u]) << 16u)); + } else { + a_log = bf16_load(alog_buf.v[params.in_strides[0] + ph]); + } + // dt_bias rides bf16 in the Qwen3.8 checkpoint; widening is exact. + float dt = bf16_load(dt_buf.v[params.in_strides[1] + ph]); + // The eager chain rounds x = a + dt_bias to bf16 (bf16 + bf16 stays +#ifndef GDN_COMPOSED_ORDER + // bf16) and softplus(x) to bf16; replicate those roundings exactly. +#else /* GDN_COMPOSED_ORDER */ + // bf16); the widen is exact. +#endif /* GDN_COMPOSED_ORDER */ + float x = bf16_load(bf16_store(a_in + dt)); +#ifdef ORACLE_QWEN35 + // Preserve the omitted contract, precision, and host-interface facts from the original comment. + float sp = bf16_load(bf16_store(max(x, 0.0f) + log(1.0f + exp(-abs(x))))); + g = exp(-(exp(a_log) * sp)); +#else /* ORACLE_QWEN35 */ +#ifndef GDN_COMPOSED_ORDER + // Compute elementwise LogAddExp softplus with corrected log1p ratio, then round to bf16. + float u = 1.0f + exp(-abs(x)); + float l1p = (u == 1.0f) ? exp(-abs(x)) : log(u) * (exp(-abs(x)) / (u - 1.0f)); + // C++ fallback keeps softplus f32; flag 9 selects it. Other models round softplus to bf16. + float sp_raw = max(x, 0.0f) + l1p; + float sp = ((params.flags & 512u) != 0u) ? sp_raw + : bf16_load(bf16_store(sp_raw)); + // g stays f32; eager-chain bf16 roundings apply to x, softplus, and beta. + g = exp(-(exp(a_log) * sp)); +#else /* GDN_COMPOSED_ORDER */ + if ((params.flags & 512u) != 0u) { + // Flag 9 (f32 A_log) selects the C++ raw fallback; softplus output stays f32. + g = gdn_g_composed(x, a_log); + } else { + // softplus=logaddexp(0,x)=max(a,b)+log1p(exp(-|a-b|)); corrected log1p=log(u)*x/(u-1), then bf16. + float u = 1.0f + exp(-abs(x)); + float l1p = (u == 1.0f) ? exp(-abs(x)) : log(u) * (exp(-abs(x)) / (u - 1.0f)); + float sp = bf16_load(bf16_store(max(x, 0.0f) + l1p)); + g = exp(-(exp(a_log) * sp)); + } +#endif /* GDN_COMPOSED_ORDER */ +#endif /* ORACLE_QWEN35 */ + } else { + g = bf16_load(g_buf.v[params.shape[0] + head]); + beta = bf16_load(b_buf.v[params.aux_offset + head]); + } + return vec2(g, beta); +} + +void main() { + uint dk = params.matrix_m; + uint dv = params.matrix_n; + uint hv = params.matrix_k; + uint head = gl_WorkGroupID.x; + uint row0 = gl_WorkGroupID.y * ROWS; + uint lane = gl_LocalInvocationID.x; + + if (head >= hv) { + return; + } + vec2 gates = gdn_gates(head); + float g = gates.x; + float beta = gates.y; +#ifdef GDN_COMPOSED_ORDER + // khead: the key head shared by this value head (in_strides[2] value and in_strides[3] key heads per row). + uint khead = head; + if (params.in_strides[3] != 0u) { + uint hv_b = params.in_strides[2]; + uint hk_b = params.in_strides[3]; + khead = (head / hv_b) * hk_b + (head % hv_b) / (hv_b / hk_b); + } +#elif defined(GDN_NATIVE_GQA_INDEX) + // Native GQA key index: the key-head order BF16 A_log models use. + uint khead = head; + if (params.in_strides[3] != 0u) { + uint hv_b = params.in_strides[2]; + uint hk_b = params.in_strides[3]; + khead = (head / hv_b) * hk_b + (head % hv_b) / (hv_b / hk_b); + } +#else /* GDN_COMPOSED_ORDER */ + uint khead = head; +#endif /* GDN_COMPOSED_ORDER */ + sk[lane] = bf16_load(k_buf.v[params.rhs_offset + khead * dk + lane]); + sq[lane] = bf16_load(q_buf.v[params.lhs_offset + khead * dk + lane]); + skg[lane] = sk[lane] * g; + uint head_base = head * dv * dk; + uint tile_base = params.shape[1] + head_base + row0 * dk; + for (uint r = 0u; r < ROWS; ++r) { + tile[r * TILE_STRIDE + lane] = s_in.v[tile_base + r * dk + lane]; + } + barrier(); + + if (lane < ROWS) { + uint row = row0 + lane; + uint trow = lane * TILE_STRIDE; + bool mlx_order = (params.flags & 512u) != 0u; + float kv_acc = gdn_kv_walk(trow, dk, g, mlx_order); + precise float delta = + (bf16_load(v_buf.v[params.aux_size + head * dv + row]) - kv_acc) * beta; + + float o_acc = 0.0f; + for (uint i = 0u; i < dk; ++i) { + precise float dk_term = delta * sk[i]; + precise float ns; + if (mlx_order) { + // Composed chain rounds state*g and delta*k separately, then adds; no FMA contraction. + precise float s_g = tile[trow + i] * g; + ns = s_g + dk_term; + } else { + ns = gdn_decay_add(tile[trow + i], g, dk_term); + } + tile[trow + i] = ns; + precise float o_term = ns * sq[i]; + o_acc += o_term; + } + y_buf.v[params.output_offset + head * dv + row] = bf16_store(o_acc); + } + barrier(); + uint out_base = params.shape[2] + head_base + row0 * dk; + for (uint r = 0u; r < ROWS; ++r) { + s_out.v[out_base + r * dk + lane] = tile[r * TILE_STRIDE + lane]; + } +} diff --git a/zig/kernels/vulkan/gated_delta_prefill_oracle.comp b/zig/kernels/vulkan/gated_delta_prefill_oracle.comp new file mode 100644 index 00000000000..275ade77cae --- /dev/null +++ b/zig/kernels/vulkan/gated_delta_prefill_oracle.comp @@ -0,0 +1,238 @@ +// From omarchy-mlx overlay/mlx/backend/omarchy/shaders/gated_delta_prefill.comp at 7d3f69ff2010. +// Code adapted: 3 lines added or changed, 7 removed. +// Copyright © 2026 Joshua Warren and mlx-omarchy contributors; MIT License, full text in LICENSES/mlx-omarchy-MIT.txt. +#version 460 + +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require + +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_shader_subgroup_shuffle_relative : require + +// Dk is 128 by the dispatch contract (GatedDeltaUpdate::use_fallback). +#define DK 128u +#define LANES 4u +#define SLICE (DK / LANES) +#define WG_THREADS 128u +#define ROWS_PER_WG (WG_THREADS / LANES) + +layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in; + +// ACT_F32: the soa (q/k/v), beta and y activations are f32; the bf16 g binding stays for the non-f32 g path. +#ifdef ACT_F32 +#define ACT_IO float +#define ACT_LOADV(x) (x) +#define ACT_STOREV(x) (x) +#else +#define ACT_IO uint16_t +#define ACT_LOADV(x) bf16_load(x) +#define ACT_STOREV(x) bf16_store(x) +#endif + +layout(set = 0, binding = 0, std430) readonly buffer QBuf { ACT_IO v[]; } q_buf; +layout(set = 0, binding = 1, std430) readonly buffer KBuf { ACT_IO v[]; } k_buf; +layout(set = 0, binding = 2, std430) readonly buffer VBuf { ACT_IO v[]; } v_buf; +layout(set = 0, binding = 3, std430) readonly buffer GBuf { uint16_t v[]; } g_buf; +layout(set = 0, binding = 4, std430) readonly buffer BBuf { ACT_IO v[]; } b_buf; +layout(set = 0, binding = 5, std430) readonly buffer SIn { float v[]; } s_in; +layout(set = 0, binding = 6, std430) writeonly buffer YBuf { ACT_IO v[]; } y_buf; +layout(set = 0, binding = 7, std430) buffer SOut { float v[]; } s_out; +layout(set = 0, binding = 8, std430) readonly buffer GBufF { float v[]; } g_buf_f; +layout(set = 0, binding = 9, std430) buffer Snap { float v[]; } snap; + +layout(push_constant) uniform Params { + uint count; + uint operation; + uint lhs_size; + uint rhs_size; + uint reduce_size; + uint output_size; + uint lhs_offset; + uint rhs_offset; + uint output_offset; + uint aux_size; + uint aux_offset; + uint matrix_m; // Dk + uint matrix_n; // Dv + uint matrix_k; // Hv + uint flags; // bit0: per-Dk decay, bit1: mask, bit2: g is f32, + // bit3: output pass + float alpha; + float beta_f; + uint dims; // T + uint shape[4]; // [0]=g offset, [1]=h0 offset, [2]=hf offset, [3]=mask offset + uint in_strides[4]; + uint out_strides[4]; + uint chunk; // chunk size C + uint chunks; // ceil(T / C) +} params; + +float bf16_load(uint16_t bits) { + return uintBitsToFloat(uint(bits) << 16u); +} + +uint16_t bf16_store(float value) { + uint bits = floatBitsToUint(value); + if (isnan(value)) { + return uint16_t((bits >> 16) | 0x40u); + } + return uint16_t((bits + 0x7fffu + ((bits >> 16) & 1u)) >> 16); +} + + +shared float sk[2][DK]; +shared float sq[2][DK]; +shared float sg[2][DK]; + +float load_g_scalar(uint t, uint head, uint hv) { + if ((params.flags & 4u) != 0u) { + return g_buf_f.v[params.shape[0] + t * hv + head]; + } + return bf16_load(g_buf.v[params.shape[0] + t * hv + head]); +} + +// Broadcast lane LANES-1's value to the LANES lanes of its row (two uniform relative shuffles; LANES == 4). +float row_total(float acc, uint sub) { + float d1 = subgroupShuffleDown(acc, 1u); + acc = (sub == 2u) ? d1 : acc; + float d2 = subgroupShuffleDown(acc, 2u); + return (sub < 2u) ? d2 : acc; +} + +// Each token updates one register state slice. Dot products chain lanes in ascending Dk order with ops-path rounding. +void token_step( + inout float st[SLICE], + uint sub, + bool tok, + float beta, + float g, + float v_row, + bool compute_o, + uint buf, + out float o_acc) { + o_acc = 0.0f; + if (!tok) { + return; + } + uint k0 = sub * SLICE; + // Round each product and add separately; NoContraction prevents FMA fusion. + precise float acc = 0.0f; + for (uint j = 0u; j < LANES; ++j) { + precise float prev = subgroupShuffleUp(acc, 1u); + if (sub == j) { + acc = (j == 0u) ? 0.0f : prev; + for (uint i = 0u; i < SLICE; ++i) { + precise float gi = (params.flags & 1u) != 0u ? sg[buf][k0 + i] : g; + precise float p = st[i] * gi; + p = p * sk[buf][k0 + i]; + acc += p; + } + } + } + precise float kv_acc = row_total(acc, sub); + precise float delta = (v_row - kv_acc) * beta; + for (uint i = 0u; i < SLICE; ++i) { + precise float gi = (params.flags & 1u) != 0u ? sg[buf][k0 + i] : g; + precise float dec = st[i] * gi; + // Match the eager fused state add with an explicit fma after the separately rounded decay. + precise float ns = fma(delta, sk[buf][k0 + i], dec); + st[i] = ns; + } + if (compute_o) { + // Prefix scan skips the ns*q output dot; state arithmetic is unchanged. + acc = 0.0f; + for (uint j = 0u; j < LANES; ++j) { + precise float prev = subgroupShuffleUp(acc, 1u); + if (sub == j) { + acc = (j == 0u) ? 0.0f : prev; + for (uint i = 0u; i < SLICE; ++i) { + precise float o = st[i] * sq[buf][k0 + i]; + acc += o; + } + } + } + o_acc = row_total(acc, sub); + } +} + +void main() { + uint dv = params.matrix_n; + uint hv = params.matrix_k; + uint T = params.dims; + uint C = params.chunk; + uint head = gl_WorkGroupID.x / LANES; + uint row_block = gl_WorkGroupID.x % LANES; + uint tid = gl_LocalInvocationID.x; + uint sub = tid % LANES; // slice index within the row + uint row = row_block * ROWS_PER_WG + tid / LANES; + bool out_pass = (params.flags & 8u) != 0u; + uint chunk = out_pass ? gl_WorkGroupID.y : 0u; + + if (head >= hv) { + return; + } + uint head_base = head * dv * DK; + uint slice_base = head_base + row * DK + sub * SLICE; + + float st[SLICE]; + + // Restore chunk initial state: chunk 0 reads h0; later chunks read snapshot slot chunk-1. + if (chunk == 0u) { + for (uint i = 0u; i < SLICE; ++i) { + st[i] = s_in.v[params.shape[1] + slice_base + i]; + } + } else { + uint snap_base = ((chunk - 1u) * hv) * dv * DK + slice_base; + for (uint i = 0u; i < SLICE; ++i) { + st[i] = snap.v[snap_base + i]; + } + } + + // Pass 0 (prefix scan) walks every token; pass 1 replays one chunk. + uint t_begin = out_pass ? chunk * C : 0u; + uint t_end = out_pass ? min(t_begin + C, T) : T; + + for (uint t = t_begin; t < t_end; ++t) { + uint t_head = t * hv + head; + // Oracle arm: no mask binding; the host rejects flag bit 1, so every token is valid. + bool tok = true; + // Double buffer: even and odd tokens use alternate shared slots. + uint buf = t & 1u; + sk[buf][tid] = ACT_LOADV(k_buf.v[params.rhs_offset + t_head * DK + tid]); + sq[buf][tid] = ACT_LOADV(q_buf.v[params.lhs_offset + t_head * DK + tid]); + if ((params.flags & 1u) != 0u) { + sg[buf][tid] = (params.flags & 4u) != 0u + ? g_buf_f.v[params.shape[0] + (t * hv + head) * DK + tid] + : bf16_load(g_buf.v[params.shape[0] + (t * hv + head) * DK + tid]); + } + barrier(); + + float beta = ACT_LOADV(b_buf.v[params.aux_offset + t * hv + head]); + float g = (params.flags & 1u) != 0u ? 0.0f : load_g_scalar(t, head, hv); + float v_row = ACT_LOADV(v_buf.v[params.aux_size + t_head * dv + row]); + + float o_acc; + token_step(st, sub, tok, beta, g, v_row, out_pass, buf, o_acc); + + if (out_pass) { + if (sub == 0u) { + y_buf.v[params.output_offset + t_head * dv + row] = + tok ? ACT_STOREV(o_acc) : ACT_STOREV(0.0f); + } + } else if ((t + 1u) % C == 0u && t + 1u < T) { + // Chunk boundary: snapshot the slice for pass 1's chunk t/C + 1 (slot index t/C). + uint snap_base = ((t / C) * hv) * dv * DK + slice_base; + for (uint i = 0u; i < SLICE; ++i) { + snap.v[snap_base + i] = st[i]; + } + } + // No trailing barrier: alternating slots and next load barrier order reuse after prior reads. + } + + if (out_pass && chunk + 1u == params.chunks) { + // Last chunk's terminal state is hf. + for (uint i = 0u; i < SLICE; ++i) { + s_out.v[params.shape[2] + slice_base + i] = st[i]; + } + } +} diff --git a/zig/kernels/vulkan/gather_take.comp b/zig/kernels/vulkan/gather_take.comp new file mode 100644 index 00000000000..cffdafbe06d --- /dev/null +++ b/zig/kernels/vulkan/gather_take.comp @@ -0,0 +1,251 @@ +#version 460 + +// From omarchy-mlx overlay/mlx/backend/omarchy/shaders/gather_take.comp at fde15c1ece69: + +// code unchanged, comments condensed. + +// Copyright © 2026 Joshua Warren; MIT License, full text in LICENSES/mlx-omarchy-MIT.txt. + +// General gather (take / take_along): t -> idx_lin and window s; table window + gathered axes. + +// operation: 0 i32, 1 u32, 2 i64 LE, 3 u8, 4 i8, 5 u16, 6 i16; signed neg wraps like offset_neg_idx. + +// OOR index (or i64 high not 0/0xFFFFFFFF) writes zero; upstream leaves such reads undefined. + +// USE_U32/U16/I64 raw copy; USE_BOOL/U8 packed bytes via BYTE_AT + atomicOr (host pre-zeros). + +// USE_MULTI_INDEX: indices bind 1, axis meta bind 2, output bind 3; push tables rank-cap 4. + +#ifdef USE_COMPLEX +#define STORAGE_TYPE vec2 +#define LOAD_VALUE(x) (x) +#define STORE_VALUE(x) (x) +#elif defined(USE_FP16) +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#define STORAGE_TYPE float16_t +#define LOAD_VALUE(x) float(x) +#define STORE_VALUE(f) STORAGE_TYPE(f) +#elif defined(USE_BF16) +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require +#define STORAGE_TYPE uint16_t +#define LOAD_VALUE(x) uintBitsToFloat(uint(x) << 16u) +uint16_t bf16_store(float value) { + uint bits = floatBitsToUint(value); + if (isnan(value)) { + return uint16_t((bits >> 16) | 0x40u); + } + return uint16_t((bits + 0x7fffu + ((bits >> 16) & 1u)) >> 16); +} +#define STORE_VALUE(f) bf16_store(f) +#elif defined(USE_U32) +#define STORAGE_TYPE uint +#define LOAD_VALUE(x) (x) +#define STORE_VALUE(x) (x) +#elif defined(USE_BOOL) +// Packed-bool: 1 byte/elem via uint words; BYTE_AT load, atomicOr store; host densifies tables. +#define STORAGE_TYPE uint +#define LOAD_BOOL_LANE(word, lane) (((word) >> ((lane) * 8u)) & 0xFFu) +#define STORE_BOOL(b) ((b) ? 1u : 0u) +#elif defined(USE_U8) +// Raw u8/i8 via uint words + BYTE_AT; no 0/1 normalize; no 8-bit storage feature. +#define STORAGE_TYPE uint +#elif defined(USE_U16) +// Raw u16/i16 bit copy; needs 16-bit storage; host gates dispatch. +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require +#define STORAGE_TYPE uint16_t +#define LOAD_VALUE(x) (x) +#define STORE_VALUE(x) (x) +#elif defined(USE_I64) +#define STORAGE_TYPE uvec2 +#define LOAD_VALUE(x) (x) +#define STORE_VALUE(x) (x) +#else +#define STORAGE_TYPE float +#define LOAD_VALUE(x) float(x) +#define STORE_VALUE(f) STORAGE_TYPE(f) +#endif + +layout(local_size_x = 256) in; + +layout(set = 0, binding = 0, std430) readonly buffer Table { + STORAGE_TYPE values[]; +} input_data; +layout(set = 0, binding = 1, std430) readonly buffer Indices { + int values[]; +} indices_data; +#ifdef USE_MULTI_INDEX +layout(set = 0, binding = 2, std430) readonly buffer AxisMetadata { + uint values[]; +} axis_metadata; +layout(set = 0, binding = 3, std430) writeonly buffer Output { + STORAGE_TYPE values[]; +} output_data; +#else +layout(set = 0, binding = 2, std430) writeonly buffer Output { + STORAGE_TYPE values[]; +} output_data; +#endif + +layout(push_constant) uniform Params { + uint count; + uint operation; + uint lhs_size; + uint rhs_size; + uint reduce_size; + uint output_size; + uint lhs_offset; + uint rhs_offset; + uint output_offset; + uint aux_size; + uint aux_offset; + uint matrix_m; + uint matrix_n; + uint matrix_k; + uint flags; + float alpha; + float beta; + uint dims; + uint shape[4]; + uint in_strides[4]; + uint out_strides[4]; +} params; + +// Honeykrisp: shift-mask byte extract miscompiles; BYTE_AT constant-shift select is proven. +#define BYTE_AT(word, lane) \ + ((lane) == 0u ? ((word) & 0xFFu) \ + : (lane) == 1u ? (((word) >> 8u) & 0xFFu) \ + : (lane) == 2u ? (((word) >> 16u) & 0xFFu) \ + : (((word) >> 24u) & 0xFFu)) + +// DECODE_INDEX: width-scaled slot on word transport; BYTE_AT narrow lanes; OOR => out_index==dim. + +// Expands as a braced block (GLSL bool conditions; do/while(0) does not compile). +#define DECODE_INDEX(buf, base, slot, op, dim, out_index) \ + { \ + uint decoded_index = (dim); \ + if ((op) == 3u || (op) == 4u) { \ + uint bword = uint(buf.values[(base) + ((slot) >> 2u)]); \ + uint b = BYTE_AT(bword, (slot) & 3u); \ + if ((op) == 4u && (b & 0x80u) != 0u) { \ + b |= 0xFFFFFF00u; \ + } \ + decoded_index = b < 0x80000000u ? b : (dim) + b; \ + } else if ((op) == 5u || (op) == 6u) { \ + uint hword = uint(buf.values[(base) + ((slot) >> 1u)]); \ + uint h = ((slot) & 1u) == 0u ? (hword & 0xFFFFu) : (hword >> 16u); \ + if ((op) == 6u && (h & 0x8000u) != 0u) { \ + h |= 0xFFFF0000u; \ + } \ + decoded_index = h < 0x80000000u ? h : (dim) + h; \ + } else { \ + uint word = uint(buf.values[(base) + (slot)]); \ + if ((op) == 0u) { \ + decoded_index = word < 0x80000000u ? word : (dim) + word; \ + } else if ((op) == 1u) { \ + decoded_index = word; \ + } else { \ + uint high = uint(buf.values[(base) + (slot) + 1u]); \ + if (high == 0u) { \ + decoded_index = word; \ + } else if (high == 0xFFFFFFFFu) { \ + decoded_index = (dim) + word; \ + } \ + } \ + } \ + (out_index) = decoded_index; \ + } \ + +void main() { + uint t = gl_GlobalInvocationID.x; + uint stride = gl_NumWorkGroups.x * gl_WorkGroupSize.x; + while (t < params.count) { +// Split output element into index position and window position. + uint idx_count = 1u; + for (uint i = 0u; i < params.dims; ++i) { + idx_count *= params.out_strides[i]; + } + uint slice_total = params.count / idx_count; + uint idx_lin = t / slice_total; + uint s = t - idx_lin * slice_total; + +// Decode window position across the table dims. + uint win = 0u; + uint scale = 1u; + for (uint i = params.flags; i > 0u; --i) { + uint dim = params.shape[i - 1u]; + win += ((s / scale) % dim) * params.in_strides[i - 1u]; + scale *= dim; + } + + uint word_slot = + params.operation == 2u ? idx_lin + idx_lin : idx_lin; + bool valid = true; + uint src = 0u; +#ifdef USE_MULTI_INDEX + for (uint axis = 0u; axis < params.aux_size; ++axis) { + uint dim = axis_metadata.values[axis]; + uint index; + DECODE_INDEX(indices_data, + axis_metadata.values[2u * params.aux_size + axis], + word_slot, + params.operation, + dim, + index); + valid = valid && index < dim; + src += index * axis_metadata.values[params.aux_size + axis]; + } +#else + { + uint dim = params.reduce_size; + uint index; + DECODE_INDEX(indices_data, + params.rhs_offset, + word_slot, + params.operation, + dim, + index); + valid = index < dim; + src = index * params.matrix_m; + } +#endif +#if defined(USE_COMPLEX) || defined(USE_U32) || defined(USE_U16) || defined(USE_I64) + STORAGE_TYPE value = STORAGE_TYPE(0); +#else + float value = 0.0; +#endif +#if defined(USE_BOOL) || defined(USE_U8) +// Bool/U8 byte-lane merge: unique lane in output word via atomicOr (four threads, four lanes). + +// Invalid index ORs 0 into host-prezeroed output (documented zero read). + uint in_elem = params.lhs_offset + src + win; +#ifdef USE_BOOL + bool bvalue = valid + ? (LOAD_BOOL_LANE(input_data.values[in_elem >> 2u], + in_elem & 3u) != 0u) + : false; + uint bbit = STORE_BOOL(bvalue); +#else + uint bbit = valid + ? BYTE_AT(input_data.values[in_elem >> 2u], in_elem & 3u) + : 0u; +#endif + uint out_elem = params.output_offset + t; + uint out_byte = out_elem & 3u; + uint out_word_index = out_elem >> 2u; + atomicOr(output_data.values[out_word_index], bbit << (out_byte * 8u)); +#else + if (valid) { + value = LOAD_VALUE( + input_data.values[params.lhs_offset + src + win]); + } + output_data.values[params.output_offset + t] = STORE_VALUE(value); +#endif + if (params.count - t <= stride) { + break; + } + t += stride; + } +} diff --git a/zig/kernels/vulkan/gdn_conv_decode.comp b/zig/kernels/vulkan/gdn_conv_decode.comp new file mode 100644 index 00000000000..aea3fa6c6b9 --- /dev/null +++ b/zig/kernels/vulkan/gdn_conv_decode.comp @@ -0,0 +1,178 @@ +// From omarchy-mlx overlay/mlx/backend/omarchy/shaders/gdn_conv_decode.comp at 69ab4cfc3f6c. +// Code unchanged, comments condensed. +// Copyright © 2026 Joshua Warren and mlx-omarchy contributors; MIT License, full text in LICENSES/mlx-omarchy-MIT.txt. +#version 460 +#ifdef APPLE_REDUCTION +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_shader_subgroup_arithmetic : require +#endif + + +#ifdef USE_BF16 +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require +#define STORAGE_TYPE uint16_t +#define LOAD_VALUE(x) uintBitsToFloat(uint(x) << 16u) +uint16_t bf16_store(float value) { + uint bits = floatBitsToUint(value); + if (isnan(value)) { + return uint16_t((bits >> 16) | 0x40u); + } + return uint16_t((bits + 0x7fffu + ((bits >> 16) & 1u)) >> 16); +} +#define STORE_VALUE(f) bf16_store(f) +#else +#define STORAGE_TYPE float +#define LOAD_VALUE(x) float(x) +#define STORE_VALUE(f) STORAGE_TYPE(f) +#endif + +layout(local_size_x = 256) in; + +layout(set = 0, binding = 0, std430) readonly buffer InState { + STORAGE_TYPE values[]; +} in_state; +layout(set = 0, binding = 1, std430) readonly buffer InX { + STORAGE_TYPE values[]; +} in_x; +layout(set = 0, binding = 2, std430) readonly buffer Weight { + STORAGE_TYPE values[]; +} weight; +layout(set = 0, binding = 3, std430) writeonly buffer Output { + STORAGE_TYPE values[]; +} output_data; +layout(set = 0, binding = 4, std430) buffer StateOut { + STORAGE_TYPE values[]; +} state_out; + +// Partial sums for the sum-of-squares reduction. +shared float norm_block[256]; + +// Field order matches our fixed push-constant layout: count, operation, lhs_size, rhs_size, reduce_size, ... +layout(push_constant) uniform Params { + uint count; + uint operation; + uint lhs_size; + uint rhs_size; + uint reduce_size; + uint output_size; + uint lhs_offset; + uint rhs_offset; + uint output_offset; + uint aux_size; + uint aux_offset; + uint matrix_m; + uint matrix_n; + uint matrix_k; + uint flags; +} params; + +void main() { + const uint channels = params.reduce_size; + const uint km1 = params.lhs_size; + uint stride_threads = gl_NumWorkGroups.x * 256u; + for (uint work_index = gl_GlobalInvocationID.x; + work_index < params.count; + work_index += stride_threads) { + uint batch = work_index / channels; + uint channel = work_index % channels; + uint state_base = batch * km1 * channels; + uint x_base = batch * channels; + float accumulator = 0.0; + for (uint tap = 0u; tap < params.operation; ++tap) { + float value = tap < km1 + ? LOAD_VALUE( + in_state.values[state_base + tap * channels + channel]) + : LOAD_VALUE(in_x.values[x_base + channel]); + accumulator += value * + LOAD_VALUE(weight.values[channel * params.operation + tap]); + } + // Flag 0 (`activate`): match swiglu.comp SILU_ONLY with storage rounding after conv, sigmoid, and multiply. + STORAGE_TYPE conv_out = STORE_VALUE(accumulator); + STORAGE_TYPE activated = conv_out; + if ((params.flags & 1u) != 0u) { + // Flag 0 (`activate`): match swiglu.comp SILU_ONLY with storage rounding after conv, sigmoid, and multiply. + float value = LOAD_VALUE(conv_out); + float sigmoid_value = 1.0 / (1.0 + exp(-value)); + STORAGE_TYPE sig_b = STORE_VALUE(sigmoid_value); + float silu = value * LOAD_VALUE(sig_b); + activated = STORE_VALUE(silu); + } + if ((params.flags & 2u) != 0u) { + // flags bit 2: RMS-normalize the activated values with a sum-of-squares reduction. +#ifdef APPLE_REDUCTION + uint thread = gl_LocalInvocationID.x; + uint half_id = thread / 128u; + float norm_value = LOAD_VALUE(activated); + float subgroup_sum = subgroupAdd(norm_value * norm_value); + if (gl_SubgroupInvocationID == 0u) { + norm_block[gl_SubgroupID] = subgroup_sum; + } + barrier(); + float partial = gl_SubgroupID == half_id * 4u && + gl_SubgroupInvocationID < 4u + ? norm_block[half_id * 4u + gl_SubgroupInvocationID] + : 0.0; + float row_sum = subgroupAdd(partial); + if (gl_SubgroupID == half_id * 4u && + gl_SubgroupInvocationID == 0u) { + float eps = uintBitsToFloat(params.matrix_k); + norm_block[8u + half_id] = + inversesqrt(row_sum / 128.0 + eps); + } + barrier(); + float norm = norm_block[8u + half_id]; + barrier(); + uint channel_head = channel / 128u; + if (channel_head < params.rhs_size / 128u * 2u) { + uint scale_bits = (channel_head < params.rhs_size / 128u) + ? params.matrix_m + : params.matrix_n; + float scale = + LOAD_VALUE(bf16_store(uintBitsToFloat(scale_bits))); + precise float result = norm_value * (norm * 1.0); + float normed = LOAD_VALUE(bf16_store(result)); + activated = bf16_store(normed * scale); + } +#else + uint thread = gl_LocalInvocationID.x; + uint half_id = thread / 128u; + uint slot = thread % 128u; + float norm_value = LOAD_VALUE(activated); + norm_block[thread] = norm_value * norm_value; + barrier(); + for (uint stride = 64u; stride > 0u; stride >>= 1u) { + if (slot < stride) { + norm_block[half_id * 128u + slot] += + norm_block[half_id * 128u + slot + stride]; + } + barrier(); + } + uint channel_head = channel / 128u; + if (channel_head < params.rhs_size / 128u * 2u) { + float eps = uintBitsToFloat(params.matrix_k); + float mean_square = norm_block[half_id * 128u] / 128.0; + float norm = inversesqrt(mean_square + eps); + uint scale_bits = (channel_head < params.rhs_size / 128u) + ? params.matrix_m + : params.matrix_n; + // Round f32 scale to storage, then widen, as the composed promote cast does. + float scale = + LOAD_VALUE(bf16_store(uintBitsToFloat(scale_bits))); + // Match FastNormGatedBF16: value * (norm * dummy weight 1.0). + precise float result = norm_value * (norm * 1.0); + // Round the normed result to storage, then widen as the composed path does. + float normed = LOAD_VALUE(bf16_store(result)); + activated = bf16_store(normed * scale); + } +#endif + } + output_data.values[work_index] = activated; + for (uint row = 0u; row + 1u < km1; ++row) { + state_out.values[state_base + row * channels + channel] = + in_state.values[state_base + (row + 1u) * channels + channel]; + } + state_out.values[state_base + (km1 - 1u) * channels + channel] = + in_x.values[x_base + channel]; + } +} diff --git a/zig/kernels/vulkan/gdn_split.comp b/zig/kernels/vulkan/gdn_split.comp new file mode 100644 index 00000000000..a569786e001 --- /dev/null +++ b/zig/kernels/vulkan/gdn_split.comp @@ -0,0 +1,133 @@ +// Splits the conv output into q, k and v (SoA by head-row) and computes the gates g and beta from a, b, A_log and dt. + +#version 460 + +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require + +layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in; + +// ACT_F32: split activations (conv, a, b, q/k/v, beta) are f32; dt stays bf16 and every Metal-matching +// intermediate bf16 rounding becomes transparent, as the validated f32 store policy rounds them to f32. +#ifdef ACT_F32 +#define SPLIT_IO float +#define SPLIT_LOAD(x) (x) +#define SPLIT_RND(x) (x) +#define SPLIT_STORE(x) (x) +#else +#define SPLIT_IO uint16_t +#define SPLIT_LOAD(x) bf16_load(x) +#define SPLIT_RND(x) bf16_load(bf16_store(x)) +#define SPLIT_STORE(x) bf16_store(x) +#endif + +layout(set = 0, binding = 0, std430) readonly buffer ConvOut { SPLIT_IO v[]; } conv_out; +layout(set = 0, binding = 1, std430) readonly buffer Ain { SPLIT_IO v[]; } ain; +layout(set = 0, binding = 2, std430) readonly buffer Bin { SPLIT_IO v[]; } bin; +layout(set = 0, binding = 3, std430) readonly buffer ALog { float v[]; } alog; +layout(set = 0, binding = 4, std430) readonly buffer Dt { uint16_t v[]; } dt; +layout(set = 0, binding = 5, std430) writeonly buffer Qkv { SPLIT_IO v[]; } qkv; +layout(set = 0, binding = 6, std430) writeonly buffer Gout { float v[]; } g_out; +layout(set = 0, binding = 7, std430) writeonly buffer Beta { SPLIT_IO v[]; } beta; + +layout(push_constant) uniform Params { + uint count; // T + uint operation; + uint lhs_size; // C + uint rhs_size; // T (SoA section stride in head-rows) + uint reduce_size; + uint output_size; // heads + uint lhs_offset; + uint rhs_offset; + uint output_offset; + uint aux_size; + uint aux_offset; + uint matrix_m; + uint matrix_n; + uint matrix_k; + uint flags; + float alpha; + float beta_f; + uint dims; + uint shape[4]; + uint in_strides[4]; + uint out_strides[4]; +} params; + +float bf16_load(uint16_t bits) { + return uintBitsToFloat(uint(bits) << 16u); +} + +uint16_t bf16_store(float value) { + uint bits = floatBitsToUint(value); + if (isnan(value)) { + return uint16_t((bits >> 16) | 0x40u); + } + return uint16_t((bits + 0x7fffu + ((bits >> 16) & 1u)) >> 16); +} + + +// Match Ash8c Metal gdn_pre lane 0: bf16-round s and softplus; keep g f32. +vec2 oracle_gate(float a_in, float b_in, float a_log, float dt_in) { + float s = SPLIT_RND(a_in + dt_in); + float sp = SPLIT_RND(max(s, 0.0f) + log(1.0f + exp(-abs(s)))); + float g = exp(-(exp(a_log) * sp)); + float beta = SPLIT_RND(1.0f / (1.0f + exp(-b_in))); + return vec2(g, beta); +} + +shared float row_reduce[32]; + +void main() { + const uint C = params.lhs_size; + const uint T = params.count; + const uint t = gl_WorkGroupID.x; + const uint head = gl_WorkGroupID.y; + const uint lane = gl_LocalInvocationID.x; + const uint hv = head % 16u; + const bool isq = head < 16u; + const bool isk = head >= 16u && head < 32u; + + float vals[4]; + uint base = t * C + head * 128u + lane * 4u; + for (uint j = 0u; j < 4u; ++j) { + float conv = SPLIT_LOAD(conv_out.v[base + j]); + float sig = SPLIT_RND(1.0f / (1.0f + exp(-conv))); + vals[j] = SPLIT_RND(conv * sig); + } + + uint out_head_row = t * 16u + hv; + uint section = isq ? 0u : (isk ? params.rhs_size * 2048u : 2u * params.rhs_size * 2048u); // T*16*128 per section + if (isq || isk) { + float ss = 0.0f; + for (uint j = 0u; j < 4u; ++j) ss += vals[j] * vals[j]; + row_reduce[lane] = ss; + barrier(); + for (uint stride = 16u; stride > 0u; stride >>= 1u) { + float other = lane < stride ? row_reduce[lane + stride] : 0.0f; + barrier(); + if (lane < stride) row_reduce[lane] += other; + barrier(); + } + float inv = inversesqrt(row_reduce[0] / 128.0f + (1e-6f / 128.0f)); + float scale = SPLIT_RND(isq ? (1.0f / 128.0f) : inversesqrt(128.0f)); + uint dst = section + out_head_row * 128u + lane * 4u; + for (uint j = 0u; j < 4u; ++j) { + qkv.v[dst + j] = SPLIT_STORE(scale * SPLIT_RND(vals[j] * inv)); + } + } else { + uint dst = section + out_head_row * 128u + lane * 4u; + for (uint j = 0u; j < 4u; ++j) { + qkv.v[dst + j] = SPLIT_STORE(vals[j]); + } + } + + if (lane == 0u) { + vec2 gates = oracle_gate(SPLIT_LOAD(ain.v[t * 16u + hv]), + SPLIT_LOAD(bin.v[t * 16u + hv]), + alog.v[hv], + bf16_load(dt.v[hv])); + g_out.v[t * 16u + hv] = gates.x; + beta.v[t * 16u + hv] = SPLIT_STORE(gates.y); + } +} diff --git a/zig/kernels/vulkan/omarchy_trig.h b/zig/kernels/vulkan/omarchy_trig.h new file mode 100644 index 00000000000..1db944bab8f --- /dev/null +++ b/zig/kernels/vulkan/omarchy_trig.h @@ -0,0 +1,68 @@ +// From omarchy-mlx overlay/mlx/backend/omarchy/shaders/omarchy_trig.h: code unchanged, comments condensed. + +// Copyright © 2026 Joshua Warren; MIT License, full text in LICENSES/mlx-omarchy-MIT.txt. + +// Shared trig range reduction for omarchy compute shaders. + +// The built-in sin/cos/tan lose upstream accuracy above ~1e3 and collapse past ~1e6 (llvmpipe past ~1e7, measured). + +// Every path that evaluates a trigonometric function of an argument not bounded by construction uses these wrappers. + +// The wrappers stay accurate to OMARCHY_TRIG_LIMIT and return NaN above it; a finite wrong value is never produced. + +// Below OMARCHY_TRIG_BUILTIN_MAX the raw built-in stays bit-identical. + +// The reduction is Cody-Waite with C1 = 6.28125 = 201 * 2^-5 (8 significant bits). + +// So k * C1 is exact in float32 for every |k| <= 83601, that is for every argument up to 525371.9. + +// The remaining 2*pi pieces are a dyadic 2^-9 (k times it is exact for all k) plus two f32 correction terms. + +// The three-subtract chain is marked precise, so a driver may not reassociate it (Honeykrisp has a documented defect). + +// Above the limit the argument leaves the exact-product zone, so the wrappers return NaN, not a wrong finite value. + +// Do not copy this logic into another shader; include this header. + +const float OMARCHY_TRIG_LIMIT = 500000.0; +const float OMARCHY_TRIG_BUILTIN_MAX = 10000.0; + +const float OMARCHY_TRIG_C1 = 6.28125; // 201 * 2^-5 +const float OMARCHY_TRIG_C2HI = 0.001953125; // 2^-9 +const float OMARCHY_TRIG_C2LO = -1.781781975296326e-05; // 2*pi - C1 - C2HI +const float OMARCHY_TRIG_C3 = -6.608047442568932e-13; // 2*pi - C1 - C2HI - C2LO +const float OMARCHY_TRIG_INV_2PI = 0.15915494309189535; // f32(1 / 2*pi) + +float omarchy_trig_nan() { return uintBitsToFloat(0x7fc00000u); } + +precise float omarchy_trig_reduce(float x) { + float k = floor(x * OMARCHY_TRIG_INV_2PI); + float r = x - k * OMARCHY_TRIG_C1; + r -= k * OMARCHY_TRIG_C2HI; + r -= k * OMARCHY_TRIG_C2LO; + r -= k * OMARCHY_TRIG_C3; + return r; +} + +float omarchy_trig_sin(float x) { + float ax = abs(x); + if (ax <= OMARCHY_TRIG_BUILTIN_MAX) return sin(x); + if (ax > OMARCHY_TRIG_LIMIT) return omarchy_trig_nan(); + return sin(omarchy_trig_reduce(x)); +} + +float omarchy_trig_cos(float x) { + float ax = abs(x); + if (ax <= OMARCHY_TRIG_BUILTIN_MAX) return cos(x); + if (ax > OMARCHY_TRIG_LIMIT) return omarchy_trig_nan(); + return cos(omarchy_trig_reduce(ax)); +} + +// tan reduces through the 2*pi-periodic sin/cos forms; the error grows near the poles as at the built-in. + +float omarchy_trig_tan(float x) { + float ax = abs(x); + if (ax <= OMARCHY_TRIG_BUILTIN_MAX) return tan(x); + if (ax > OMARCHY_TRIG_LIMIT) return omarchy_trig_nan(); + return sin(omarchy_trig_reduce(x)) / cos(omarchy_trig_reduce(ax)); +} diff --git a/zig/kernels/vulkan/qmm.comp b/zig/kernels/vulkan/qmm.comp index 5a58fa90c32..0d01be552ac 100644 --- a/zig/kernels/vulkan/qmm.comp +++ b/zig/kernels/vulkan/qmm.comp @@ -2,7 +2,7 @@ // From omarchy-mlx overlay/mlx/backend/omarchy/shaders/qmm.comp at 388b41f509fb: code unchanged, comments condensed. -// Copyright © 2026 Joshua Warren and mlx-omarchy contributors; MIT License, full text in LICENSES/mlx-omarchy-MIT.txt. +// Copyright © 2026 Joshua Warren; MIT License, full text in LICENSES/mlx-omarchy-MIT.txt. // QuantizedMatmul: out[m, n] = x row m . dequantized w row n (transposed) or column n (flags bit 0 set). @@ -37,10 +37,21 @@ uint16_t bf16_store(float value) { #define STORE_VALUE(f) STORAGE_TYPE(f) #endif +// ACT_F32: activations (x rows and output) are f32 while scales and biases stay bf16. +#ifdef ACT_F32 +#define ACT_TYPE float +#define ACT_LOAD(x) (x) +#define ACT_STORE(f) (f) +#else +#define ACT_TYPE STORAGE_TYPE +#define ACT_LOAD(x) LOAD_VALUE(x) +#define ACT_STORE(f) STORE_VALUE(f) +#endif + layout(local_size_x = 256) in; layout(set = 0, binding = 0, std430) readonly buffer InputX { - STORAGE_TYPE values[]; + ACT_TYPE values[]; } input_x; layout(set = 0, binding = 1, std430) readonly buffer InputW { uint values[]; @@ -62,7 +73,7 @@ layout(set = 0, binding = 3, std430) readonly buffer Biases { STORAGE_TYPE values[]; } biases; layout(set = 0, binding = 4, std430) writeonly buffer Output { - STORAGE_TYPE values[]; + ACT_TYPE values[]; } output_data; #endif @@ -210,7 +221,7 @@ void main() { : ((nibble == 0u) ? (packed_byte & 0xFu) : (packed_byte >> 4u)); float value = (bits == 8u) ? fp8_decode(code) : fp4_decode(code); - accumulator += LOAD_VALUE(input_x.values[x_base + k]) * + accumulator += ACT_LOAD(input_x.values[x_base + k]) * (scale * value); } } else { @@ -233,7 +244,7 @@ void main() { : (((k & 1u) == 0u) ? (packed_byte & 0xFu) : (packed_byte >> 4u)); float value = (bits == 8u) ? fp8_decode(code) : fp4_decode(code); - accumulator += LOAD_VALUE(input_x.values[x_base + k]) * + accumulator += ACT_LOAD(input_x.values[x_base + k]) * (scale * value); } } @@ -284,7 +295,7 @@ void main() { b * w_batch_words + k * n_words_per_row + col_word_off]; q = (word >> lane_shift) & value_mask; } - accumulator += LOAD_VALUE(input_x.values[x_base + k]) * + accumulator += ACT_LOAD(input_x.values[x_base + k]) * (scale * float(q) + bias); } } else { @@ -324,7 +335,7 @@ void main() { (byte_addr + 4u) & 3u)) : q6_at(lane, c0, c1, c2); accumulator += - LOAD_VALUE(input_x.values[x_base + k_base + lane]) * + ACT_LOAD(input_x.values[x_base + k_base + lane]) * (scale * float(q) + bias); } } else { @@ -332,7 +343,7 @@ void main() { for (uint lane = 0u; lane < pack_factor; ++lane) { uint q = (word >> (lane * bits)) & value_mask; accumulator += - LOAD_VALUE(input_x.values[x_base + k_base + lane]) * + ACT_LOAD(input_x.values[x_base + k_base + lane]) * (scale * float(q) + bias); } } @@ -341,6 +352,6 @@ void main() { } #endif output_data.values[params.output_offset + index] = - STORE_VALUE(accumulator); + ACT_STORE(accumulator); } } diff --git a/zig/kernels/vulkan/qmm_q4_tile.comp b/zig/kernels/vulkan/qmm_q4_tile.comp new file mode 100644 index 00000000000..8b1993cb60a --- /dev/null +++ b/zig/kernels/vulkan/qmm_q4_tile.comp @@ -0,0 +1,120 @@ +#version 460 +// qmm_q4_vec4 over tiled weights, 32 columns a tile: words [tile][k/32][lane], scales [tile][group][lane]. +// ACT_F32: x rows and the output are f32 (two vec4 a group of eight), scales and biases stay bf16. +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require + +layout(local_size_x = 256) in; + +#ifdef ACT_F32 +layout(set = 0, binding = 0, std430) readonly buffer InputX { + vec4 values[]; +} input_x; +#else +layout(set = 0, binding = 0, std430) readonly buffer InputX { + uvec4 values[]; +} input_x; +#endif +layout(set = 0, binding = 1, std430) readonly buffer InputW { + uvec4 values[]; +} input_w; +layout(set = 0, binding = 2, std430) readonly buffer Scales { + uint16_t values[]; +} scales; +layout(set = 0, binding = 3, std430) readonly buffer Biases { + uint16_t values[]; +} biases; +#ifdef ACT_F32 +layout(set = 0, binding = 4, std430) writeonly buffer Output { + float values[]; +} output_data; +#else +layout(set = 0, binding = 4, std430) writeonly buffer Output { + uint16_t values[]; +} output_data; +#endif + +layout(push_constant) uniform Params { + uint count; + uint operation; + uint lhs_size; + uint rhs_size; + uint reduce_size; + uint output_size; + uint lhs_offset; + uint rhs_offset; + uint output_offset; + uint aux_size; + uint aux_offset; + uint matrix_m; + uint matrix_n; + uint matrix_k; + uint flags; + float alpha; + float beta; + uint dims; + uint shape[4]; + uint in_strides[4]; + uint out_strides[4]; +} params; + +uint16_t bf16_store(float value) { + uint bits = floatBitsToUint(value); + if (isnan(value)) { + return uint16_t((bits >> 16) | 0x40u); + } + return uint16_t((bits + 0x7fffu + ((bits >> 16) & 1u)) >> 16); +} + +float bf16_value(uint16_t x) { + return uintBitsToFloat(uint(x) << 16u); +} + +void main() { + uint groups_per_row = params.matrix_k / 64u; + uint total = params.matrix_m * params.matrix_n; + uint stride = gl_NumWorkGroups.x * 256u; + for (uint index = gl_GlobalInvocationID.x; index < total; index += stride) { + uint row = index / params.matrix_n; + uint column = index - row * params.matrix_n; + uint tile = column >> 5u; + uint lane = column & 31u; + uint w_base = tile * params.matrix_k + lane; + uint x_vec = row * (params.matrix_k / 8u); + uint param_base = tile * (groups_per_row * 32u) + lane; + float accumulator = 0.0; + for (uint group = 0u; group < groups_per_row; ++group) { + float scale = bf16_value(scales.values[param_base + group * 32u]); + float bias = bf16_value(biases.values[param_base + group * 32u]); + for (uint half_group = 0u; half_group < 2u; ++half_group) { + uvec4 w4 = input_w.values[w_base + (group * 2u + half_group) * 32u]; + for (uint t = 0u; t < 4u; ++t) { + uint word = w4[t]; + uint x_at = x_vec + group * 8u + half_group * 4u + t; +#ifdef ACT_F32 + vec4 x_lo = input_x.values[x_at * 2u]; + vec4 x_hi = input_x.values[x_at * 2u + 1u]; + for (uint e = 0u; e < 8u; ++e) { + float xv = e < 4u ? x_lo[e] : x_hi[e - 4u]; + accumulator += xv * (scale * float((word >> (e * 4u)) & 0xfu) + bias); + } +#else + uvec4 x8 = input_x.values[x_at]; + for (uint pair = 0u; pair < 4u; ++pair) { + uint x_pair = x8[pair]; + accumulator += uintBitsToFloat(x_pair << 16u) * + (scale * float((word >> (pair * 8u)) & 0xfu) + bias); + accumulator += uintBitsToFloat(x_pair & 0xffff0000u) * + (scale * float((word >> (pair * 8u + 4u)) & 0xfu) + bias); + } +#endif + } + } + } +#ifdef ACT_F32 + output_data.values[params.output_offset + index] = accumulator; +#else + output_data.values[params.output_offset + index] = bf16_store(accumulator); +#endif + } +} diff --git a/zig/kernels/vulkan/qmm_q4_vec4.comp b/zig/kernels/vulkan/qmm_q4_vec4.comp new file mode 100644 index 00000000000..8cbd47d05c4 --- /dev/null +++ b/zig/kernels/vulkan/qmm_q4_vec4.comp @@ -0,0 +1,91 @@ +#version 460 +// qmm.comp q4 g64 transposed bf16, 16-byte loads, one batch, zero offsets; same arithmetic; see ops_q4.qmmVariant. +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require + +layout(local_size_x = 256) in; + +layout(set = 0, binding = 0, std430) readonly buffer InputX { + uvec4 values[]; +} input_x; +layout(set = 0, binding = 1, std430) readonly buffer InputW { + uvec4 values[]; +} input_w; +layout(set = 0, binding = 2, std430) readonly buffer Scales { + uint16_t values[]; +} scales; +layout(set = 0, binding = 3, std430) readonly buffer Biases { + uint16_t values[]; +} biases; +layout(set = 0, binding = 4, std430) writeonly buffer Output { + uint16_t values[]; +} output_data; + +layout(push_constant) uniform Params { + uint count; + uint operation; + uint lhs_size; + uint rhs_size; + uint reduce_size; + uint output_size; + uint lhs_offset; + uint rhs_offset; + uint output_offset; + uint aux_size; + uint aux_offset; + uint matrix_m; + uint matrix_n; + uint matrix_k; + uint flags; + float alpha; + float beta; + uint dims; + uint shape[4]; + uint in_strides[4]; + uint out_strides[4]; +} params; + +uint16_t bf16_store(float value) { + uint bits = floatBitsToUint(value); + if (isnan(value)) { + return uint16_t((bits >> 16) | 0x40u); + } + return uint16_t((bits + 0x7fffu + ((bits >> 16) & 1u)) >> 16); +} + +float bf16_value(uint16_t x) { + return uintBitsToFloat(uint(x) << 16u); +} + +void main() { + uint groups_per_row = params.matrix_k / 64u; + uint total = params.matrix_m * params.matrix_n; + uint stride = gl_NumWorkGroups.x * 256u; + for (uint index = gl_GlobalInvocationID.x; index < total; index += stride) { + uint row = index / params.matrix_n; + uint column = index - row * params.matrix_n; + uint w_vec = column * (params.matrix_k / 32u); + uint x_vec = row * (params.matrix_k / 8u); + uint param_base = column * groups_per_row; + float accumulator = 0.0; + for (uint group = 0u; group < groups_per_row; ++group) { + float scale = bf16_value(scales.values[param_base + group]); + float bias = bf16_value(biases.values[param_base + group]); + for (uint half_group = 0u; half_group < 2u; ++half_group) { + uvec4 w4 = input_w.values[w_vec + group * 2u + half_group]; + for (uint t = 0u; t < 4u; ++t) { + uint word = w4[t]; + uvec4 x8 = input_x.values[x_vec + group * 8u + half_group * 4u + t]; + for (uint pair = 0u; pair < 4u; ++pair) { + uint x_pair = x8[pair]; + accumulator += uintBitsToFloat(x_pair << 16u) * + (scale * float((word >> (pair * 8u)) & 0xfu) + bias); + accumulator += uintBitsToFloat(x_pair & 0xffff0000u) * + (scale * float((word >> (pair * 8u + 4u)) & 0xfu) + bias); + } + } + } + } + output_data.values[params.output_offset + index] = bf16_store(accumulator); + } +} diff --git a/zig/kernels/vulkan/sdpa_decode_native.comp b/zig/kernels/vulkan/sdpa_decode_native.comp new file mode 100644 index 00000000000..1c131fda724 --- /dev/null +++ b/zig/kernels/vulkan/sdpa_decode_native.comp @@ -0,0 +1,431 @@ +// From omarchy-mlx overlay/mlx/backend/omarchy/shaders/sdpa_decode_native.comp at 2dcd24ba9c50. +// Code adapted: 1 lines added or changed, 1 removed. +// Copyright © 2026 Joshua Warren and mlx-omarchy contributors; MIT License, full text in LICENSES/mlx-omarchy-MIT.txt. + +// One 1024-thread workgroup handles each query head. + +// Short and mid decode use 32 interleaved online streams and the final SIMD reduction. + +// Long contexts switch to sdpa_decode_native_p1.comp + _p2.comp. + +// That pair replaces the former in-workgroup 64/128-block two-pass. + +// The -DBF16_IO arm matches f32-score composition and matmul.comp per-dot accumulation order. + +// softmax_suffix.comp computes max/exp/sum/normalize over stored f32 scores. + +// One f32 q*scale and one RNE bf16 store make it bit-identical to the ten-dispatch path. + +#version 460 + +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_shader_subgroup_arithmetic : require +#extension GL_KHR_shader_subgroup_shuffle : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#ifdef BF16_IO +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require +#endif + +layout(local_size_x = 1024, local_size_y = 1, local_size_z = 1) in; +#ifdef BF16_IO +// bf16 uses a uint16 word view; loads widen from the f32 top half. + +// RNE stores match every other USE_BF16 leg; see cast.comp. + +// F32_Q/F32_OUT keep the cache bf16 while the query and y activations are f32. +layout(set = 0, binding = 0, std430) readonly buffer Query { +#ifdef F32_Q + float values[]; +#else + uint16_t values[]; +#endif +} query_data; +layout(set = 0, binding = 1, std430) readonly buffer Key { + uint16_t values[]; +} key_data; +layout(set = 0, binding = 2, std430) readonly buffer Value { + uint16_t values[]; +} value_data; +layout(set = 0, binding = 3, std430) writeonly buffer Output { +#ifdef F32_OUT + float values[]; +#else + uint16_t values[]; +#endif +} output_data; +#else +layout(set = 0, binding = 0, std430) readonly buffer Query { + float16_t values[]; +} query_data; +layout(set = 0, binding = 1, std430) readonly buffer Key { + float16_t values[]; +} key_data; +layout(set = 0, binding = 2, std430) readonly buffer Value { + float16_t values[]; +} value_data; +layout(set = 0, binding = 3, std430) writeonly buffer Output { + float16_t values[]; +} output_data; +#endif + +layout(push_constant) uniform Params { + uint count; + uint operation; + uint lhs_size; + uint rhs_size; + uint reduce_size; + uint output_size; + uint lhs_offset; + uint rhs_offset; + uint output_offset; + uint aux_size; + uint aux_offset; + uint matrix_m; + uint matrix_n; + uint matrix_k; + uint flags; + float alpha; + float beta; + uint dims; + uint shape[4]; + uint in_strides[4]; + uint out_strides[4]; +} params; + +const uint SIMD = 32u; +const float FINITE_MIN = -3.4028234663852886e38; + +// The default SDPA_DIM=64 keeps original blobs; -DSDPA_DIM=N compiles wider twins. + +// In bf16, N must be a multiple of 32 up to 256. + +// Each bf16 width uses ascending-d dots, 256-lane softmax, and ascending-key 16-wide output tiles. + +// Results match the composed path at every width; only per-dot length and output thread count scale. + +// In f16, each lane covers SDPA_DIM/64 consecutive dim pairs; the 64-wide simdgroup mapping is applied PAIRS times. + +// Therefore, f16 N must be a multiple of 64. + +// Native 32-interleaved-stream order and merge stay unchanged per pair. +#ifndef SDPA_DIM +#define SDPA_DIM 64u +#endif +const uint PAIRS = SDPA_DIM / 64u; + +#ifndef BF16_IO +shared float s_outputs[SIMD * SIMD]; +shared float s_max[SIMD]; +shared float s_sum[SIMD]; +#else +// This composition-exact arm replaces online softmax with the f32-score composition's arithmetic. + +// Probs overwrite scores in one f32 stream; normalize reads and writes only its own strided slots. + +// The host guard enforces kStreamKeys; larger contexts fall through to the same bit-exact composition. +const uint kStreamKeys = 7168u; +shared float s_qrow[SDPA_DIM]; +shared float s_stream[kStreamKeys]; +shared float s_tree[256]; +#endif + +float native_exp(float x) { + return exp(x); +} + +#ifdef BF16_IO +// The cast.comp conversion pair, verbatim: exact widening in, RNE out. +float bf16_load(uint16_t bits) { + return uintBitsToFloat(uint(bits) << 16u); +} +uint16_t bf16_store(float value) { + uint bits = floatBitsToUint(value); + if (isnan(value)) { + return uint16_t((bits >> 16) | 0x40u); + } + return uint16_t((bits + 0x7fffu + ((bits >> 16) & 1u)) >> 16); +} +#define LD16(x) bf16_load(x) +#else +#define LD16(x) float(x) +#endif + +// F32_Q/F32_OUT: the query and y activations are f32; the K/V cache stays bf16. +#ifdef BF16_IO +#ifdef F32_Q +#define Q_LOAD(x) (x) +#else +#define Q_LOAD(x) bf16_load(x) +#endif +#ifdef F32_OUT +#define OUT_STORE(x) (x) +#else +#define OUT_STORE(x) bf16_store(x) +#endif +#else +#define Q_LOAD(x) LD16(x) +#define OUT_STORE(x) float16_t(x) +#endif + +void store_pair( + uint head, + uint lane, + uint pair, + float x, + float y, + float denominator) { + if (lane == 0u) { + uint out_base = params.output_offset + head * SDPA_DIM + + pair * 64u + 2u * gl_SubgroupID; + output_data.values[out_base] = OUT_STORE(x / denominator); + output_data.values[out_base + 1u] = OUT_STORE(y / denominator); + } +} + +void main() { +#ifdef BF16_IO + uint thread = gl_LocalInvocationID.x; + uint head = gl_WorkGroupID.x; + uint repeats = params.matrix_m / params.matrix_n; + uint kv_head = head / repeats; + uint k_len = params.matrix_k; + // Workgroup z selects the batch row; offsets are q out_strides[2], k lhs_size, v rhs_size, and out out_strides[3]. + + // All offsets are zero at batch 0. + uint batch = gl_WorkGroupID.z; + uint q_batch = params.lhs_offset + batch * params.out_strides[2]; + uint k_batch = params.rhs_offset + batch * params.lhs_size; + uint v_batch = params.aux_offset + batch * params.rhs_size; + uint o_batch = params.output_offset + batch * params.out_strides[3]; +#ifdef SDPA_ROWS + // Multi-token verification uses workgroup (head, row) and the row's causal bound. + + // Scanning keys 0..key_end-1 is bit-identical to decoding that row alone at q_len=1 with the same cache prefix. + + // params.dims carries q_len; out_strides[0] is the q row stride. + + // Output and q use shape[0] as their shared head stride. + uint row = gl_WorkGroupID.y; + uint key_end = params.matrix_k - params.dims + row + 1u; +#else + uint row = 0u; + uint key_end = params.matrix_k; +#endif + + if (thread < SDPA_DIM) { +#ifdef SDPA_ROWS + s_qrow[thread] = params.alpha * Q_LOAD( + query_data.values[q_batch + head * params.shape[0] + + row * params.out_strides[0] + thread * params.shape[3]]); +#else + s_qrow[thread] = params.alpha * Q_LOAD( + query_data.values[q_batch + + head * params.shape[0] + thread * params.shape[3]]); +#endif + } + barrier(); + + // Scores over every key: ascending-d accumulation, matmul.comp order. + for (uint key = thread; key < key_end; key += 1024u) { + uint k_base = k_batch + kv_head * params.shape[1] + + key * params.shape[2]; + float accumulator = 0.0; + // Load 16 K elements before 16 in-order FMAs, independent of the accumulator chain and bit-identical in order. + for (uint d = 0u; d < SDPA_DIM; d += 16u) { + float kv[16]; + for (uint i = 0u; i < 16u; ++i) { + kv[i] = bf16_load(key_data.values[k_base + (d + i) * params.in_strides[0]]); + } + for (uint i = 0u; i < 16u; ++i) { + accumulator += s_qrow[d + i] * kv[i]; + } + } + s_stream[key] = 1.0 * accumulator; + } + barrier(); + + // Row max: strided scan then the 256-lane binary tree. + if (thread < 256u) { + float local_max = uintBitsToFloat(0xff800000u); + for (uint index = thread; index < key_end; index += 256u) { + float next = s_stream[index]; + if (isnan(next) || next > local_max) { + local_max = next; + } + } + s_tree[thread] = local_max; + } + barrier(); + for (uint stride = 128u; stride > 0u; stride >>= 1u) { + if (thread < stride) { + s_tree[thread] = max(s_tree[thread], s_tree[thread + stride]); + } + barrier(); + } + float row_max = s_tree[0]; + barrier(); + + // Denominator: strided sequential exp-sum per lane, then the tree. + if (thread < 256u) { + float local_sum = 0.0; + for (uint index = thread; index < key_end; index += 256u) { + local_sum += exp(s_stream[index] - row_max); + } + s_tree[thread] = local_sum; + } + barrier(); + for (uint stride = 128u; stride > 0u; stride >>= 1u) { + if (thread < stride) { + s_tree[thread] += s_tree[thread + stride]; + } + barrier(); + } + float normalizer = 1.0 / s_tree[0]; + barrier(); + + // Recompute probs from stored f32 scores, as softmax_suffix.comp's normalize pass does. + if (thread < 256u) { + for (uint index = thread; index < key_end; index += 256u) { + s_stream[index] = exp(s_stream[index] - row_max) * normalizer; + } + } + barrier(); + + // Output dimension d accumulates keys in ascending order over 16-wide tiles with exact-zero tail padding. + + // Drain with 1.0f and store RNE bf16. + if (thread < SDPA_DIM) { + uint d = thread; + float accumulator = 0.0; + uint key = 0u; + uint v_base = v_batch + kv_head * params.in_strides[1] + + d * params.in_strides[3]; + // Load 16 independent V elements before the 16 in-order FMAs. + + // A 64-key block spilled load registers and ran 9x slower. + for (; key + 16u <= key_end; key += 16u) { + float vv[16]; + for (uint index = 0u; index < 16u; ++index) { + vv[index] = bf16_load( + value_data.values[v_base + (key + index) * params.in_strides[2]]); + } + for (uint index = 0u; index < 16u; ++index) { + accumulator += s_stream[key + index] * vv[index]; + } + } + for (uint index = 0u; index < 16u; ++index) { + float prob = key < key_end ? s_stream[key] : 0.0; + float value16 = key < key_end ? bf16_load( + value_data.values[v_batch + + kv_head * params.in_strides[1] + + key * params.in_strides[2] + + d * params.in_strides[3]]) : 0.0; + accumulator += prob * value16; + ++key; + } +#ifdef SDPA_ROWS + // Row stride is a push constant, not SDPA_DIM. + + // The dense Qwen3.5 oracle stores y row-major as (rows, 8, 256), with row stride heads*SDPA_DIM (2048). + + // This matches the q read's use of out_strides[0]. + output_data.values[o_batch + head * params.shape[0] + + row * params.out_strides[0] + d] = OUT_STORE(1.0 * accumulator); +#else + output_data.values[o_batch + head * SDPA_DIM + d] = + OUT_STORE(1.0 * accumulator); +#endif + } + return; +#else + uint lane = gl_SubgroupInvocationID; + uint sgid = gl_SubgroupID; + uint head = gl_WorkGroupID.x; + uint repeats = params.matrix_m / params.matrix_n; + uint kv_head = head / repeats; + uint q_base = params.lhs_offset + head * params.shape[0]; + uint q_dim = 2u * lane; + float q0[PAIRS]; + float q1[PAIRS]; + for (uint p = 0u; p < PAIRS; ++p) { + q0[p] = params.alpha * + LD16(query_data.values[q_base + (p * 64u + q_dim) * params.shape[3]]); + q1[p] = params.alpha * LD16( + query_data.values[q_base + (p * 64u + q_dim + 1u) * params.shape[3]]); + } + + // Native sdpa_vector assigns each SIMD keys sgid, sgid+32, ... with independent online-softmax state. + + // It merges 32 states at the end; each lane's score dot accumulates ascending-d across its dim pairs. + + // PAIRS=1 preserves the original two-fma sequence verbatim. + float max_score = FINITE_MIN; + float sum_exp_score = 0.0; + float o0[PAIRS]; + float o1[PAIRS]; + for (uint p = 0u; p < PAIRS; ++p) { + o0[p] = 0.0; + o1[p] = 0.0; + } + for (uint key = sgid; key < params.matrix_k; key += SIMD) { + uint k_base = params.rhs_offset + kv_head * params.shape[1] + + key * params.shape[2] + q_dim * params.in_strides[0]; + float score = 0.0; + for (uint p = 0u; p < PAIRS; ++p) { + score = fma( + q0[p], + LD16(key_data.values[k_base + p * 64u * params.in_strides[0]]), + score); + score = fma( + q1[p], + LD16(key_data.values[ + k_base + (p * 64u + 1u) * params.in_strides[0]]), + score); + } + score = subgroupAdd(score); + float new_max = max(max_score, score); + float factor = native_exp(max_score - new_max); + float exp_score = native_exp(score - new_max); + max_score = new_max; + sum_exp_score = fma(sum_exp_score, factor, exp_score); + uint v_base = params.aux_offset + kv_head * params.in_strides[1] + + key * params.in_strides[2] + q_dim * params.in_strides[3]; + for (uint p = 0u; p < PAIRS; ++p) { + o0[p] = fma( + o0[p], + factor, + exp_score * LD16(value_data.values[ + v_base + p * 64u * params.in_strides[3]])); + o1[p] = fma( + o1[p], + factor, + exp_score * LD16(value_data.values[ + v_base + (p * 64u + 1u) * params.in_strides[3]])); + } + } + + if (lane == 0u) { + s_max[sgid] = max_score; + s_sum[sgid] = sum_exp_score; + } + barrier(); + max_score = s_max[lane]; + float merged_max = subgroupMax(max_score); + float factor = native_exp(max_score - merged_max); + float denominator = subgroupAdd(s_sum[lane] * factor); + + for (uint p = 0u; p < PAIRS; ++p) { + s_outputs[lane * SIMD + sgid] = o0[p]; + barrier(); + float merged0 = subgroupAdd(s_outputs[sgid * SIMD + lane] * factor); + barrier(); + s_outputs[lane * SIMD + sgid] = o1[p]; + barrier(); + float merged1 = subgroupAdd(s_outputs[sgid * SIMD + lane] * factor); + store_pair(head, lane, p, merged0, merged1, denominator); + // The next pair's write must not start before this pair's read ends. + barrier(); + } +#endif +} diff --git a/zig/kernels/vulkan/swiglu.comp b/zig/kernels/vulkan/swiglu.comp new file mode 100644 index 00000000000..c1869e78ed5 --- /dev/null +++ b/zig/kernels/vulkan/swiglu.comp @@ -0,0 +1,112 @@ +// From omarchy-mlx overlay/mlx/backend/omarchy/shaders/swiglu.comp at 98650249f97f. +// Code adapted: 5 lines added or changed, 0 removed. +// Copyright © 2026 Joshua Warren and mlx-omarchy contributors; MIT License, full text in LICENSES/mlx-omarchy-MIT.txt. + +// Four-wide SwiGLU rounds each op; aligned host inputs; optional materialize writes sig/silu buffers. + +#version 460 + +#ifdef ACT_F32 +// F32 activations: every buffer is plain vec4 and the store rounding is transparent. +#define STORAGE_VEC vec4 +#define LOAD_VEC(x) (x) +#define STORE_VEC(v) (v) +float round_intermediate(float f) { + return f; +} +#else +#ifdef USE_FP16 +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#define STORAGE_VEC f16vec4 +#define LOAD_VEC(x) vec4(x) +#define STORE_VEC(v) f16vec4(v) +float round_intermediate(float f) { + return unpackHalf2x16(packHalf2x16(vec2(f, 0.0))).x; +} +#else +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require +#define STORAGE_VEC u16vec4 +#define LOAD_VEC(x) uintBitsToFloat(uvec4(x) << 16u) +uint16_t bf16_store(float value) { + uint bits = floatBitsToUint(value); + if (isnan(value)) { + return uint16_t((bits >> 16) | 0x40u); + } + return uint16_t((bits + 0x7fffu + ((bits >> 16) & 1u)) >> 16); +} +#define STORE_VEC(v) \ + u16vec4(bf16_store(v.x), bf16_store(v.y), bf16_store(v.z), bf16_store(v.w)) +float round_intermediate(float f) { + return uintBitsToFloat(uint(bf16_store(f)) << 16u); +} +#endif +#endif + +layout(local_size_x = 256) in; + +layout(set = 0, binding = 0, std430) readonly buffer Gate { + STORAGE_VEC values[]; +} gate_buf; +layout(set = 0, binding = 1, std430) readonly buffer Up { + STORAGE_VEC values[]; +} up_buf; +layout(set = 0, binding = 2, std430) writeonly buffer Output { + STORAGE_VEC values[]; +} output_data; +layout(set = 0, binding = 3, std430) writeonly buffer SigmoidOut { + STORAGE_VEC values[]; +} sigmoid_out; +layout(set = 0, binding = 4, std430) writeonly buffer SiluOut { + STORAGE_VEC values[]; +} silu_out; + +layout(push_constant) uniform Params { + uint count; + uint gate_offset; + uint up_offset; + uint materialize; +} params; + +vec4 round4(vec4 v) { + return vec4( + round_intermediate(v.x), + round_intermediate(v.y), + round_intermediate(v.z), + round_intermediate(v.w)); +} + +void main() { + uint vec_count = params.count / 4u; + uint index = gl_GlobalInvocationID.x; + uint stride = gl_NumWorkGroups.x * gl_WorkGroupSize.x; + while (index < vec_count) { + vec4 gate = LOAD_VEC(gate_buf.values[params.gate_offset / 4u + index]); +#ifdef SILU_ONLY + vec4 up = vec4(1.0); // standalone silu: silu * 1.0 is exact +#else + vec4 up = LOAD_VEC(up_buf.values[params.up_offset / 4u + index]); +#endif +#ifdef ORACLE_QWEN35 + // Match Metal's g/(1+exp(-g)): one f32 division, no intermediate rounding, one output-store round. + vec4 denom = 1.0 + exp(-gate); + vec4 sig = 1.0 / denom; // exact only for the materialize path + vec4 silu = gate / denom; + vec4 value = silu * up; +#else + vec4 sig = round4(1.0 / (1.0 + exp(-gate))); + vec4 silu = round4(gate * sig); + vec4 value = round4(silu * up); +#endif + if (params.materialize != 0u) { + sigmoid_out.values[index] = STORE_VEC(sig); + silu_out.values[index] = STORE_VEC(silu); + } + output_data.values[index] = STORE_VEC(value); + if (vec_count - index <= stride) { + break; + } + index += stride; + } +} diff --git a/zig/src/core/contract/vulkan_records.zig b/zig/src/core/contract/vulkan_records.zig index 8d1969ce65e..7bbc52fae92 100644 --- a/zig/src/core/contract/vulkan_records.zig +++ b/zig/src/core/contract/vulkan_records.zig @@ -25,6 +25,9 @@ pub const devices = [_]limits.VulkanDeviceLimits{ pub const records = [_]Record{ .{ .launch = .{ .name = "qmm_bf16", .local_size = .{ 256, 1, 1 } } }, + .{ .launch = .{ .name = "qmm_q4_vec4", .local_size = .{ 256, 1, 1 } } }, + .{ .launch = .{ .name = "qmm_q4_tile", .local_size = .{ 256, 1, 1 } } }, + .{ .launch = .{ .name = "qmm_q4_tile_f32", .local_size = .{ 256, 1, 1 } } }, .{ .launch = .{ .name = "control:qmm_fmax_nonfinite", .local_size = .{ 256, 1, 1 } }, .control_gate = 1 }, .{ .launch = .{ .name = "control:atomic_sum", .local_size = .{ 256, 1, 1 } }, .control_gate = 2 }, .{ .launch = .{ .name = "control:scribble_input", .local_size = .{ 256, 1, 1 } }, .control_gate = 2 }, diff --git a/zig/src/families/qwen3_5/backend/vulkan/backend.zig b/zig/src/families/qwen3_5/backend/vulkan/backend.zig new file mode 100644 index 00000000000..3e406817f18 --- /dev/null +++ b/zig/src/families/qwen3_5/backend/vulkan/backend.zig @@ -0,0 +1,256 @@ +// Vulkan Qwen3.5-2B backend owns one stream and sequence; greedy decode has no rewind, so reset clears state. +const std = @import("std"); +const lanes = @import("lanes"); +const engine_mod = @import("engine.zig"); +const pinned = @import("../../config.zig"); + +const Engine = engine_mod.Engine; +const be = lanes.backend; + +// Each prompt pass uses chunks of at most 32 rows. +pub const max_prompt_rows: usize = 32; + +/// The SDPA arm's visible-key limit, the state's capacity ceiling (state.zig). +pub const max_capacity = @import("state.zig").max_capacity; + +// A prompt accepts at most max_capacity real token ids; reply decodes are admitted separately. +pub fn validPrompt(ids: []const u32) bool { + if (ids.len == 0 or ids.len > max_capacity) return false; + for (ids) |id| if (id >= pinned.vocab) return false; + return true; +} + +/// A stream's request must decode greedily; the host's `greedy_only` gate refuses the rest first. +fn samplingRefused(s: *const lanes.Stream) bool { + const sampling = s.sampling orelse return false; + return sampling.temperature > 0; +} + +pub const Vulkan = struct { + gpa: std.mem.Allocator, + engine: *Engine, + owner: ?*lanes.Stream = null, // the one active stream; the engine holds a single sequence + + pub fn init(gpa: std.mem.Allocator, engine: *Engine) !*Vulkan { + const out = try gpa.create(Vulkan); + errdefer gpa.destroy(out); + out.* = .{ .gpa = gpa, .engine = engine }; + return out; + } + + // The engine (device, weights, sequence) outlives this backend; its owner deinits it once quiescent. + pub fn deinit(self: *Vulkan) void { + self.gpa.destroy(self); + } + + pub fn facts(self: *const Vulkan) lanes.Model { + _ = self; + return .{ + .exact_width = 1, // serial one-row exactness + .copy_drafts = false, + .gpu_tokens = false, // no pipelined rounds + .mtp = false, + .speculate = false, // no draft head + .speculate_early = false, + .drafts = 0, + .hidden_rows = false, // no multi-stream shared forwards + .streams_exact = true, + .max_streams = 1, + .batch_rows = 1, + .head_trees = false, + .draft_probabilities = false, + }; + } + + pub fn backend(self: *Vulkan) be.Backend { + return .{ .ptr = self, .vtable = &.{ + .prefill = prefillFn, + .prefill_step = prefillStepFn, + .cancel_prefill = cancelPrefillFn, + .first = firstFn, + .queue = queueFn, + .read = readFn, + .verify = verifyFn, + .keep = keepFn, + .draft = draftFn, + .release = releaseFn, + .reports_fill = true, + } }; + } + + fn cast(ptr: *anyopaque) *Vulkan { + return @ptrCast(@alignCast(ptr)); + } + + fn prefillFn(ptr: *anyopaque, stream: *lanes.Stream) !void { + const self = cast(ptr); + if (self.owner != null) return error.QwenStreamBusy; + if (samplingRefused(stream)) return error.UnsupportedQwenSampling; + const ids = stream.prompt(); + if (!validPrompt(ids)) return error.InvalidQwenPrompt; + // steps floor 1: a prompt-only warm pass (max_new 0) still reads one row past the prompt + try engine_mod.planFor(self.engine.capacity(), 0, ids.len, @max(stream.max_new, 1)); + try self.engine.reset(); // a new prompt starts from a really zeroed state (also clears a poisoned one) + try self.engine.prefill(ids); // the same canonical 32-row schedule the stepped path runs + stream.fill_at.store(@intCast(ids.len), .monotonic); + self.owner = stream; + } + + // Each 32-row call checks ids and capacity before reset; visible() tracks progress. + fn prefillStepFn(ptr: *anyopaque, stream: *lanes.Stream, max_rows: u32) !bool { + const self = cast(ptr); + if (self.owner) |o| { + if (o != stream) return error.QwenStreamBusy; + } else { + if (samplingRefused(stream)) return error.UnsupportedQwenSampling; + const ids = stream.prompt(); + if (!validPrompt(ids)) return error.InvalidQwenPrompt; + // steps floor 1: same footprint as the whole-pass admission (a prompt-only request reads one row past) + try engine_mod.planFor(self.engine.capacity(), 0, ids.len, @max(stream.max_new, 1)); + try self.engine.reset(); // full request admitted: now the real zeroing + self.owner = stream; + } + const ids = stream.prompt(); + const done: usize = @intCast(self.engine.visible()); + if (done >= ids.len) return true; + var end = engine_mod.chunkEnd(done, ids.len); + if (max_rows > 0) end = @min(end, done + @as(usize, max_rows)); // 0: uncapped, per the lanes contract + try self.engine.prefillChunk(ids[done..end]); + const filled = self.engine.visible(); + stream.fill_at.store(@intCast(filled), .monotonic); + return filled == ids.len; + } + + // Cancellation cannot keep a partial prefix; reset clears state and ownership, so later release is a no-op. + fn cancelPrefillFn(ptr: *anyopaque, stream: *lanes.Stream) !void { + const self = cast(ptr); + if (self.owner != stream) return error.UnknownQwenStream; + self.owner = null; + try self.engine.reset(); + } + + fn firstFn(ptr: *anyopaque, stream: *lanes.Stream, position: u64) !u64 { + const self = cast(ptr); + if (self.owner != stream) return error.UnknownQwenStream; + if (position != self.engine.visible()) return error.InvalidQwenPosition; + if (samplingRefused(stream)) return error.UnsupportedQwenSampling; + return try self.engine.draw(); + } + + fn queueFn(ptr: *anyopaque, stream: *lanes.Stream, feed: be.Feed, position: u64) !u64 { + const self = cast(ptr); + if (self.owner != stream) return error.UnknownQwenStream; + if (position != self.engine.visible() + 1) return error.InvalidQwenPosition; + const token: u32 = switch (feed) { + .value => |v| v, + .handle => |h| std.math.cast(u32, h) orelse return error.InvalidQwenToken, + }; + if (token >= pinned.vocab) return error.InvalidQwenToken; + try self.engine.decode(token); + return try self.engine.draw(); + } + + fn readFn(_: *anyopaque, handle: u64) !u32 { + if (handle >= pinned.vocab) return error.InvalidQwenToken; + return @intCast(handle); + } + + fn verifyFn(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) !void { + const self = cast(ptr); + if (windows.len != 1 or out.len != 1) return error.UnsupportedQwenWindow; + const w = windows[0]; + const o = &out[0]; + if (w.parents != null or w.held != 0 or w.tokens.len != 0 or w.masks.len != 0 or w.early) return error.UnsupportedQwenWindow; + if (w.rows() != 1) return error.UnsupportedQwenWindow; + if (o.sampled.len != 1 or o.drafts.len != 0) return error.InvalidQwenOutputs; + if (self.owner != w.stream) return error.UnknownQwenStream; + if (w.positions.len != 1 or w.positions[0] != self.engine.visible() + 1) return error.InvalidQwenPosition; + if (samplingRefused(w.stream)) return error.UnsupportedQwenSampling; + try self.engine.decode(w.pending); + o.sampled[0] = try self.engine.draw(); + } + + fn keepFn(ptr: *anyopaque, windows: []const be.Window, paths: []const []const u32) !void { + const self = cast(ptr); + if (windows.len != 1 or paths.len != 1) return error.UnsupportedQwenWindow; + const w = windows[0]; + if (self.owner != w.stream) return error.UnknownQwenStream; + if (w.parents != null or w.held != 0) return error.UnsupportedQwenWindow; + // The state is append-only: the one committed row is already in, and no rewind exists. + if (paths[0].len != 1) return error.UnsupportedQwenRewind; + } + + fn draftFn(_: *anyopaque, requests: []const be.DraftRequest) !void { + for (requests) |r| if (r.depth != 0) return error.QwenCheckpointHasNoDraftHead; + } + + fn releaseFn(ptr: *anyopaque, stream: *lanes.Stream) void { + const self = cast(ptr); + if (self.owner == stream) { + self.owner = null; + self.engine.reset() catch |err| + std.log.err("qwen3.5 vulkan: the release reset failed ({s}); the next prompt resets again", .{@errorName(err)}); + } + } +}; + +test "facts pin the serial one-row greedy engine" { + const v: Vulkan = .{ .gpa = undefined, .engine = undefined }; + var cfg = try lanes.Config.init(std.testing.allocator, v.facts(), 1, 0); + defer cfg.deinit(std.testing.allocator); + try std.testing.expectEqual(@as(u32, 1), cfg.family_width); + try std.testing.expectEqual(@as(u32, 0), cfg.max_copy); + try std.testing.expectEqual(@as(u32, 0), cfg.most_drafts); + try std.testing.expectEqual(@as(u32, 1), cfg.batch_streams); + try std.testing.expect(!cfg.family_streams); + try std.testing.expect(!cfg.pipelined); +} + +test "prompt validation bounds: whole requests to the ceiling, never past, real ids only" { + try std.testing.expect(!validPrompt(&.{})); + try std.testing.expect(validPrompt(&.{1})); + try std.testing.expect(!validPrompt(&.{pinned.vocab})); // OOV + var oov_last: [65]u32 = @splat(7); + oov_last[oov_last.len - 1] = pinned.vocab; // valid length, OOV beyond the first chunk + try std.testing.expect(!validPrompt(&oov_last)); + const over: [max_capacity + 1]u32 = @splat(7); // 7169 ids: one past the ceiling + try std.testing.expect(!validPrompt(&over)); + const at_cap: [max_capacity]u32 = @splat(7); // 7168 ids: exactly the ceiling + try std.testing.expect(validPrompt(&at_cap)); + const buf: [7168]u32 = @splat(7); + for ([_]usize{ 1, 31, 32, 33, 63, 64, 65, 7105 }) |len| { + try std.testing.expect(validPrompt(buf[0..len])); + } +} + +test "the whole-request plan gates exact fit and one past, with checked arithmetic" { + const cap: usize = max_capacity; + for ([_]usize{ 1, 31, 32, 33, 63, 64, 65, 7105, 7168 }) |len| { + try engine_mod.planFor(cap, 0, len, 1); // prompt-only: last draw reads a written row + } + try engine_mod.planFor(cap, 0, 7167, 2); // 7167 rows + 1 decode = exact fit + try std.testing.expectError(error.ContextTooLong, engine_mod.planFor(cap, 0, 7168, 2)); // one past + try std.testing.expectError(error.ContextTooLong, engine_mod.planFor(cap, 0, 7169, 1)); + try std.testing.expectError(error.UnsupportedPrompt, engine_mod.planFor(cap, 0, 0, 1)); // empty + try std.testing.expectError(error.UnsupportedPrompt, engine_mod.planFor(cap, 0, 1, 0)); + try std.testing.expectError(error.ContextTooLong, engine_mod.planFor(cap, std.math.maxInt(u64) - 2, 2, 2)); // checked add refuses, never wraps +} + +test "stepped and cancel callbacks exist and their ownership guards fire before any engine access" { + var v: Vulkan = .{ .gpa = undefined, .engine = undefined }; + const table = v.backend().vtable; + try std.testing.expect(table.prefill_step != null); // the advertised step is wired + try std.testing.expect(table.cancel_prefill != null); + try std.testing.expect(table.reports_fill); + // busy guard: a foreign stream while another owns the pass refuses before the engine is touched + var owned: Vulkan = .{ .gpa = undefined, .engine = undefined }; + var other = try lanes.Stream.init(std.testing.allocator, .{ .id = "owned", .prompt = &.{1}, .max_new = 4 }); + defer other.deinit(std.testing.allocator); + owned.owner = &other; + var s = try lanes.Stream.init(std.testing.allocator, .{ .id = "t", .prompt = &.{ 1, 2 }, .max_new = 4 }); + defer s.deinit(std.testing.allocator); + try std.testing.expectError(error.QwenStreamBusy, table.prefill_step.?(@ptrCast(&owned), &s, 0)); + // cancel guard: no owner, or a foreign stream, refuses without touching the engine + var fresh: Vulkan = .{ .gpa = undefined, .engine = undefined }; + try std.testing.expectError(error.UnknownQwenStream, table.cancel_prefill.?(@ptrCast(&fresh), &s)); +} diff --git a/zig/src/families/qwen3_5/backend/vulkan/engine.zig b/zig/src/families/qwen3_5/backend/vulkan/engine.zig new file mode 100644 index 00000000000..5df5e92d8ee --- /dev/null +++ b/zig/src/families/qwen3_5/backend/vulkan/engine.zig @@ -0,0 +1,393 @@ +// Admission precedes weights, state allocates last, and post-write failures poison until reset; no rollback. +const std = @import("std"); +const vk = @import("vulkan"); +const pinned = @import("../../config.zig"); +const weights = @import("weights.zig"); +const state = @import("state.zig"); +const fwd = @import("forward.zig"); + +const Allocator = std.mem.Allocator; +const View = vk.tensor.View; + +pub const max_prefill_rows = fwd.max_prefill_rows; + +// Chunk cut = min(done+32,total), pure, host-tested; totals 32/33/63/64/65 cut at 32 / 32,33 / 32,63 / 32,64 / 32,64,65 +pub fn chunkEnd(done: usize, total: usize) usize { + return @min(done +| max_prefill_rows, total); +} + +// The staged-transfer buffer: 8 MiB so the load's ~27k tensors stream in large chunks instead of one 64 KiB fence each. +pub const staging_bytes: u64 = 8 << 20; +// The 8 MiB scratch arena holds persistent windows and worst-layer prefill transients (~5.6 MiB), with headroom. +pub const arena_bytes: u64 = 8 << 20; + +// Step errors include dispatch.Error; Timeout exits 124 and InFlight exits 1 inside `vk.orExit`, before teardown. +pub const Error = vk.dispatch.Error || + error{ ContextTooLong, NonFinite, UnsupportedPrompt, Poisoned }; +// Engine.init has no explicit error set: inferred platform I/O/config errors pass through; steps do no filesystem I/O. + +// Build-supplied SPIR-V variants (norm/swiglu/rope use ORACLE_QWEN35 Metal oracle); engine checks and pipelines them. +pub const ShaderBundle = struct { + ops: vk.ops.Shaders, + norm: vk.norm_mlp.Shaders, + attention: vk.attention.Shaders, + gdn: vk.gdn.Shaders, + greedy: vk.greedy.Shaders, +}; + +// The plan is checked before GPU work: `steps` draws need `rows + steps - 1` positions; checked sums refuse overflow. +pub fn planFor(capacity: usize, visible: u64, rows: usize, steps: usize) Error!void { + if (rows == 0 or steps == 0) return error.UnsupportedPrompt; + if (capacity > state.max_capacity) return error.ContextTooLong; + const need = std.math.add(u64, visible, rows) catch return error.ContextTooLong; + const total = std.math.add(u64, need, steps) catch return error.ContextTooLong; + if (total - 1 > capacity) return error.ContextTooLong; +} + +// Shared logits readiness: poison or zero visible prefix refuses draw/readback before GPU work; pure and host-tested. +pub fn logitsReady(poisoned: bool, visible: u64) Error!void { + if (poisoned) return error.Poisoned; + if (visible == 0) return error.Invalid; +} + +// The arena stores four f32 hidden-row buffers, ids, and full logits; it also fits the worst transient layer. +pub fn persistentBytes() u64 { + const rows: u64 = max_prefill_rows; + return 4 * rows * pinned.hidden * 4 + rows * 4 + pinned.vocab * 2; +} + +pub const Engine = struct { + gpa: Allocator, + timeout_ns: u64, + rt: *vk.Runtime, + dev: *vk.Device, + cmds: *vk.Commands, + staging: vk.Buffer, + ops: vk.ops.Ops, + norm: vk.norm_mlp.NormMlp, + attn: vk.Attention, + gdn: vk.Gdn, + greedy: vk.Greedy, + arena: vk.arena.Arena, + model: weights.Model, + state: state.State, + // The residual ring: h[0]/h[1] alternate as the add's input and output, so no add ever writes the buffer it reads. + h: [2]View, + r: View, + x: View, + ids: View, + logits: View, + base_mark: u64, + poisoned: bool, + + // Pinned-model admission runs inside the weights load before GPU work; SDPA arm limit is 7168; probe uses 256. + pub fn init( + gpa: Allocator, + io: std.Io, + bundle: ShaderBundle, + model_dir: []const u8, + context: usize, + timeout_ns: u64, + ) !*Engine { + const cap = try state.capacityFor(context); + const e = try gpa.create(Engine); + errdefer gpa.destroy(e); + e.gpa = gpa; + e.timeout_ns = timeout_ns; + e.poisoned = false; + + e.rt = try gpa.create(vk.Runtime); + errdefer gpa.destroy(e.rt); + e.rt.* = try vk.Runtime.open(); + errdefer e.rt.close(); + e.dev = try gpa.create(vk.Device); + errdefer gpa.destroy(e.dev); + e.dev.* = try vk.Device.open(e.rt, .{ .features = .{ .storage16 = true, .int16 = true } }); + errdefer e.dev.close(); + e.cmds = try gpa.create(vk.Commands); + errdefer gpa.destroy(e.cmds); + e.cmds.* = try vk.Commands.init(e.dev); + errdefer e.cmds.deinit(); + e.staging = try vk.Buffer.init(e.dev, .staging, staging_bytes); + errdefer e.staging.deinit(); + + // device gates first: every operator refuses a lacking device before a single weight byte is read + e.ops = try vk.ops.Ops.init(e.dev, bundle.ops); + errdefer e.ops.deinit(); + e.norm = try vk.norm_mlp.NormMlp.init(e.dev, bundle.norm); + errdefer e.norm.deinit(); + e.attn = try vk.Attention.init(e.dev, bundle.attention, &e.ops); + errdefer e.attn.deinit(); + e.gdn = try vk.Gdn.init(e.dev, bundle.gdn); + errdefer e.gdn.deinit(); + e.greedy = try vk.Greedy.init(e.dev, bundle.greedy, pinned.vocab); + errdefer e.greedy.deinit(); + + e.arena = try vk.arena.Arena.init(e.dev, arena_bytes); + errdefer e.arena.deinit(); + // h/r/x hold f32 bytes at their widest; per-layer views re-kind the same span (see ringView). + const rows: u64 = max_prefill_rows; + e.h[0] = try e.arena.alloc(.f32, rows * pinned.hidden); + e.h[1] = try e.arena.alloc(.f32, rows * pinned.hidden); + e.r = try e.arena.alloc(.f32, rows * pinned.hidden); + e.x = try e.arena.alloc(.f32, rows * pinned.hidden); + e.ids = try e.arena.alloc(.u32, rows); + e.logits = try e.arena.alloc(.bf16, pinned.vocab); + e.base_mark = e.arena.mark(); + + e.model = try vk.orExit(weights.Model.load(gpa, e.dev, e.cmds, io, model_dir, e.staging, timeout_ns)); + errdefer e.model.deinit(); + e.state = try vk.orExit(state.State.init(gpa, e.dev, e.cmds, cap, e.staging, timeout_ns)); + return e; + } + + // Reverse teardown is legal only with no pending submission; `orExit` ensures Timeout/InFlight never return here. + pub fn deinit(e: *Engine) void { + e.state.deinit(); + e.model.deinit(); + e.arena.deinit(); + e.greedy.deinit(); + e.gdn.deinit(); + e.attn.deinit(); + e.norm.deinit(); + e.ops.deinit(); + e.staging.deinit(); + e.cmds.deinit(); + e.dev.close(); + e.rt.close(); + const gpa = e.gpa; + gpa.destroy(e.cmds); + gpa.destroy(e.dev); + gpa.destroy(e.rt); + gpa.destroy(e); + } + + // Reset zeroes all state on GPU; commits length/parities after fences; clears poison and rewinds scratch to base. + pub fn reset(e: *Engine) Error!void { + try vk.orExit(e.state.reset(e.cmds, e.staging, e.timeout_ns)); + e.arena.reset(e.base_mark) catch {}; + e.poisoned = false; + } + + // Check ids/bounds, commit 32-row chunks. Failure poisons (partial RNN state never reused); refusals touch nothing. + pub fn prefill(e: *Engine, tokens: []const u32) Error!void { + if (e.poisoned) return error.Poisoned; + if (tokens.len == 0) return error.UnsupportedPrompt; + try vk.ops_q4.checkIds(tokens, e.model.embed.n); + try planFor(@intCast(e.state.capacity), e.state.visible(), tokens.len, 1); + var done: usize = 0; + while (done < tokens.len) { + const end = chunkEnd(done, tokens.len); + try e.prefillChunk(tokens[done..end]); + done = end; + } + } + + // One forward chunk has 1..32 checked rows; its logits give the next-token distribution after the chunk. + pub fn prefillChunk(e: *Engine, tokens: []const u32) Error!void { + if (e.poisoned) return error.Poisoned; + if (tokens.len == 0 or tokens.len > max_prefill_rows) return error.UnsupportedPrompt; + try vk.ops_q4.checkIds(tokens, e.model.embed.n); + try planFor(@intCast(e.state.capacity), e.state.visible(), tokens.len, 1); + e.poisoned = true; + try vk.orExit(vk.upload(e.cmds, e.ids.buffer.*, e.ids.byte_offset, e.staging, std.mem.sliceAsBytes(tokens), e.timeout_ns)); + try vk.orExit(embed(e, ringView(try e.h[0].slice(0, tokens.len * pinned.hidden), .bf16), tokens.len)); + try runLayers(e, tokens.len); + try vk.orExit(e.state.append(tokens.len)); + e.poisoned = false; + } + + // Decode embeds the fed token, runs all 24 layers and final norm/head; resulting logits cover the next token. + pub fn decode(e: *Engine, token: u32) Error!void { + if (e.poisoned) return error.Poisoned; + try vk.ops_q4.checkIds(&.{token}, e.model.embed.n); + try planFor(@intCast(e.state.capacity), e.state.visible(), 1, 1); + e.poisoned = true; + try vk.orExit(vk.upload(e.cmds, e.ids.buffer.*, e.ids.byte_offset, e.staging, std.mem.sliceAsBytes(&[1]u32{token}), e.timeout_ns)); + try vk.orExit(embed(e, ringView(try e.h[0].slice(0, pinned.hidden), .bf16), 1)); + try runLayers(e, 1); + try vk.orExit(e.state.append(1)); + e.poisoned = false; + } + + // The oracle draw picks the largest finite logit; ties (+0/-0) go lowest index. Any NaN/Inf refuses and poisons. + pub fn draw(e: *Engine) Error!u32 { + try logitsReady(e.poisoned, e.state.visible()); + return vk.orExit(e.greedy.select(e.cmds, e.staging, e.logits, e.timeout_ns)) catch |err| { + if (err == error.NonFinite) e.poisoned = true; + return err; + }; + } + + // Logits are bytes for receipt readback, never CPU argmax; wrong `out.len` is refused without poisoning. + pub fn logitsBytes(e: *Engine, out: []u8) Error!void { + try logitsReady(e.poisoned, e.state.visible()); + if (out.len != e.logits.byteLen()) return error.Invalid; + try vk.orExit(vk.download(e.cmds, e.logits.buffer.*, e.logits.byte_offset, e.staging, out, e.timeout_ns)); + } + + pub fn visible(e: *Engine) u64 { + return e.state.visible(); + } + + /// The logits row's byte length, for sizing caller buffers. + pub fn logitsLen(e: *Engine) u64 { + return e.logits.byteLen(); + } + + /// The capacity actually allocated (the context request, SDPA-capped). + pub fn capacity(e: *Engine) usize { + return @intCast(e.state.capacity); + } + + /// config.json's max_position_embeddings, as the weights admission read it. + pub fn modelContext(e: *Engine) usize { + return e.model.context; + } + + /// The chosen device's Vulkan vendorID, for the probe's optional guard. + pub fn vendorId(e: *Engine) u32 { + return e.dev.properties.vendorID; + } + + fn embed(e: *Engine, out: View, rows: usize) !void { + const cx = e.core(); + const mark = e.arena.mark(); + errdefer e.arena.reset(mark) catch {}; + try vk.orExit(vk.ops_q4.embedding(cx.acc.ops, cx.cmds, out, try e.ids.slice(0, rows), e.model.embed.w, e.model.embed.scales, e.model.embed.biases, e.model.embed.k, &e.arena, e.timeout_ns)); + e.arena.reset(mark) catch {}; + } + + fn core(e: *Engine) fwd.Core { + return .{ + .cmds = e.cmds, + .staging = e.staging, + .acc = .{ .ops = &e.ops, .norm = &e.norm, .attn = &e.attn, .gdn = &e.gdn }, + .scratch = &e.arena, + .timeout_ns = e.timeout_ns, + }; + } + + // 24 layers, store kinds per fwd.actKind: norm (plain at 0, else add), mixer, post norm, MLP. + fn runLayers(e: *Engine, rows: usize) Error!void { + const cx = e.core(); + const span = rows * pinned.hidden; + const base = e.state.visible(); + var ring: usize = 0; + for (e.model.layers, 0..) |layer, i| { + const act = fwd.actKind(i); + const prev: vk.tensor.Kind = if (i == 0) .bf16 else fwd.actKind(i - 1); + const x = ringView(try e.x.slice(0, span), act); + const input_norm, const post_norm, const mlp_q = switch (layer) { + .full => |fl| .{ fl.input_norm, fl.post_norm, fl.mlp }, + .linear => |ll| .{ ll.input_norm, ll.post_norm, ll.mlp }, + }; + if (i == 0) { + try cx.plainNorm(x, ringView(try e.h[0].slice(0, span), prev), input_norm, rows); + } else { + // the residual keeps the previous kind; layer 4's add rounds the f32 sum to bf16 on the store. + try cx.addNorm(x, ringView(try e.h[ring ^ 1].slice(0, span), prev), ringView(try e.h[ring].slice(0, span), prev), ringView(try e.r.slice(0, span), prev), input_norm, rows); + ring ^= 1; + } + const r = ringView(try e.r.slice(0, span), act); + switch (layer) { + .full => |fl| try cx.fullLayer(x, fl, &e.state, state.fullSlot(i).?, rows, r, base), + .linear => |ll| try cx.linearLayer(x, ll, &e.state, state.linearSlot(i).?, rows, r), + } + // the sum lands in the current layer's kind; its `a` side still reads the previous sum + try cx.addNorm(ringView(try e.x.slice(0, span), act), ringView(try e.h[ring ^ 1].slice(0, span), act), ringView(try e.h[ring].slice(0, span), prev), r, post_norm, rows); + ring ^= 1; + try cx.mlp(ringView(try e.x.slice(0, span), act), rows, mlp_q, ringView(try e.r.slice(0, span), act)); + } + const bf16: vk.tensor.Kind = .bf16; + try cx.finish(ringView(try e.h[ring ^ 1].slice(0, span), bf16), ringView(try e.h[ring].slice(0, span), bf16), ringView(try e.r.slice(0, span), bf16), ringView(try e.x.slice(0, span), bf16), e.model.final_norm, e.model.outputHead(), e.logits, rows); + } +}; + +/// The same span over the arena's f32-backed window, viewed in another element kind. +fn ringView(v: View, kind: vk.tensor.Kind) View { + return .{ .buffer = v.buffer, .byte_offset = v.byte_offset, .kind = kind, .elements = v.elements }; +} + +const testing = std.testing; + +test "logitsReady admits only a clean engine with a visible prefix" { + try logitsReady(false, 1); + try logitsReady(false, 256); + try testing.expectError(error.Invalid, logitsReady(false, 0)); + try testing.expectError(error.Poisoned, logitsReady(true, 7)); + try testing.expectError(error.Poisoned, logitsReady(true, 0)); // poison wins +} + +test "the plan refuses empty, stepless and overflowing sequences before any GPU work" { + const cap = 256; + try planFor(cap, 0, 16, 16); + try planFor(cap, 240, 16, 1); // exactly fills: 240 + 16 + 1 - 1 = 256 + try testing.expectError(error.UnsupportedPrompt, planFor(cap, 0, 0, 1)); + try testing.expectError(error.UnsupportedPrompt, planFor(cap, 0, 1, 0)); + try testing.expectError(error.ContextTooLong, planFor(cap, 241, 16, 1)); + try testing.expectError(error.ContextTooLong, planFor(0, 0, 1, 1)); + // checked sums: extreme caller inputs refuse instead of wrapping + try testing.expectError(error.ContextTooLong, planFor(cap, std.math.maxInt(u64), 1, 1)); + try testing.expectError(error.ContextTooLong, planFor(cap, 0, 1, std.math.maxInt(usize))); + try testing.expectError(error.ContextTooLong, planFor(cap, std.math.maxInt(u64) - 1, 1, 3)); +} + +test "the canonical cut ends every whole-request length at 32-row boundaries" { + for ([_]struct { len: usize, cuts: []const usize }{ + .{ .len = 1, .cuts = &.{1} }, + .{ .len = 31, .cuts = &.{31} }, + .{ .len = 32, .cuts = &.{32} }, + .{ .len = 33, .cuts = &.{ 32, 33 } }, + .{ .len = 63, .cuts = &.{ 32, 63 } }, + .{ .len = 64, .cuts = &.{ 32, 64 } }, + .{ .len = 65, .cuts = &.{ 32, 64, 65 } }, + }) |case| { + var done: usize = 0; + var at: usize = 0; + while (done < case.len) : (at += 1) { + const end = chunkEnd(done, case.len); + try testing.expect(at < case.cuts.len and end == case.cuts[at]); // exactly the expected schedule + done = end; + } + try testing.expectEqual(case.cuts.len, at); + } +} + +test "the plan admits whole requests up to the capacity ceiling and refuses one past" { + const cap: usize = 7168; + for ([_]usize{ 1, 31, 32, 33, 63, 64, 65, 7105, 7168 }) |len| { + try planFor(cap, 0, len, 1); // a prompt-only request always fits at or under the ceiling + } + try testing.expectError(error.ContextTooLong, planFor(cap, 0, 7169, 1)); // one past the ceiling + try planFor(cap, 0, 7168, 1); // exact fit prompt + one draw reads a written row + try testing.expectError(error.ContextTooLong, planFor(cap, 0, 7168, 2)); // one past with a decode + try planFor(cap, 0, 7105, 64); // the frozen long case ends at exactly 7168 writes + try testing.expectError(error.ContextTooLong, planFor(cap, 0, 7106, 64)); + try testing.expectError(error.ContextTooLong, planFor(cap + 1, 0, 1, 1)); + try planFor(cap, 0, 7167, 2); // 7167 + 2 - 1 = 7168: exact fit +} + +test "the wide store policy covers exactly layers 0-3 of the 24-layer model" { + for (0..24) |i| { + const want: vk.tensor.Kind = if (i < 4) .f32 else .bf16; + try testing.expectEqual(want, fwd.actKind(i)); + } +} + +test "the arena holds the persistent windows plus the worst transient layer" { + const rows: u64 = max_prefill_rows; + const conv: u64 = pinned.conv_dim; + const dim: u64 = pinned.linear_heads * pinned.linear_dim; + // one linear layer at the cap in f32 activations: qkv, z, a, b, soa, g, beta, conv_out, conv_in, y, mix + const linear = 4 * (rows * (5 * conv + 3 * dim + 4 * pinned.linear_heads) + (rows + state.conv_rows) * conv); + // one full-attention layer at the cap in f32: qraw, q, qn, k, kn, v, queries, y, mix + const full = 4 * rows * (2 * 8 * 256 + 5 * 8 * 256 + 3 * 2 * 256); + // the MLP in f32: gate + up + act + const mlp = 12 * pinned.intermediate * rows; + const worst = @max(linear, @max(full, mlp)); + try testing.expect(worst > 0); + try testing.expectEqual(fwd.actKind(0), vk.tensor.Kind.f32); // the wide layers the buffers size for + try testing.expectEqual(fwd.actKind(fwd.wide_layers), vk.tensor.Kind.bf16); // the first narrow layer + try testing.expect(arena_bytes >= persistentBytes() + worst + (1 << 20)); +} diff --git a/zig/src/families/qwen3_5/backend/vulkan/forward.zig b/zig/src/families/qwen3_5/backend/vulkan/forward.zig new file mode 100644 index 00000000000..9c82a4a775c --- /dev/null +++ b/zig/src/families/qwen3_5/backend/vulkan/forward.zig @@ -0,0 +1,196 @@ +// Single-sequence Qwen3.5-2B Vulkan matches the Metal oracle (24 layers: 18 GDN, 6 attention); GPU calls via vk.orExit. +const std = @import("std"); +const vk = @import("vulkan"); +const pinned = @import("../../config.zig"); +const weights = @import("weights.zig"); +const state = @import("state.zig"); + +const Arena = vk.arena.Arena; +const View = vk.tensor.View; + +// Whole-prompt rows are split into dispatch chunks capped here; persistent windows use this cap. +pub const max_prefill_rows: usize = 32; + +// Layers 0-3 keep every activation store in f32; 4+ round stores to bf16 (the validated store policy). +pub const wide_layers: usize = state.wide_layers; + +// The activation kind a layer's stores use; pure, host-tested alongside the ring kinds. +pub fn actKind(layer: usize) vk.tensor.Kind { + return if (layer < wide_layers) .f32 else .bf16; +} +/// The GDN prefill scan's chunk cap (chunks of 1 when a prefill fits one). +pub const max_gdn_chunk: u64 = 32; + +comptime { + if (2 * pinned.query_heads * pinned.head_dim != 4096) + @compileError("the interleaved Q layout drifted from [q 256 | gate 256] x 8"); + if (pinned.conv_dim != 2 * pinned.linear_heads * pinned.linear_dim + pinned.linear_heads * pinned.linear_dim) + @compileError("GDN channel order drifted from q|k|v"); +} + +/// The admitted operators the composition drives. +pub const Acc = struct { + ops: *vk.ops.Ops, + norm: *vk.norm_mlp.NormMlp, + attn: *vk.Attention, + gdn: *vk.Gdn, +}; + +// `Core` borrows commands, staging, operators, scratch and fence timeout; the owner owns the arena and every view. +pub const Core = struct { + cmds: *vk.Commands, + staging: vk.Buffer, + acc: Acc, + scratch: *Arena, + timeout_ns: u64, + + /// One affine projection through the admitted Q4 closure; weights repacked at load run the tiled image. + fn project(c: *const Core, out: View, x: View, q: weights.Quant) !void { + if (q.tiled) return vk.orExit(vk.ops_q4.projectTiled(c.acc.ops, c.cmds, out, x, q.w, q.scales, q.biases, q.k, c.timeout_ns)); + try vk.orExit(vk.ops_q4.project(c.acc.ops, c.cmds, out, x, q.w, q.scales, q.biases, q.k, c.scratch, c.timeout_ns)); + } + + // Fused norm: bf16 residual add (`h = bf16(H + R)`) out of place into `sum`, then raw-gain RMSNorm into `out`. + pub fn addNorm(c: *const Core, out: View, sum: View, h: View, r: View, gain: View, rows: usize) !void { + try vk.orExit(c.acc.norm.add(c.cmds, sum, h, r, c.timeout_ns)); + try vk.orExit(c.acc.norm.rmsNorm(c.cmds, out, sum, gain, rows, pinned.hidden, pinned.eps, c.timeout_ns)); + } + + /// The layer-0 norm: no residual, the gain over the raw embedding. + pub fn plainNorm(c: *const Core, out: View, h: View, gain: View, rows: usize) !void { + try vk.orExit(c.acc.norm.rmsNorm(c.cmds, out, h, gain, rows, pinned.hidden, pinned.eps, c.timeout_ns)); + } + + // `extractQueries` reads `qraw[d + k * 512]` (k = row * 8 + head); the transport rides the views' kind. + pub fn extractQueries(c: *const Core, q_out: View, qraw: View, rows: usize) !void { + const plan = extractQPlan(rows, q_out.elements, qraw.elements) orelse return error.Invalid; + if (q_out.kind != qraw.kind or (q_out.kind != .bf16 and q_out.kind != .f32)) return error.Invalid; + const pipe = c.acc.ops.copyPipeline(q_out.kind); + try pipe.bind(&.{ qraw.bind(), c.acc.ops.dummyBinding(), q_out.bind(), c.acc.ops.dummyBinding() }); + var p: vk.params.Params = .{ + .count = plan.count, + .dims = 2, + .shape = .{ plan.dim, plan.heads_rows, 0, 0 }, + .in_strides = .{ 1, plan.in_head_stride, 0, 0 }, + .out_strides = .{ 1, plan.dim, 0, 0 }, + }; + try vk.orExit(pipe.run(c.cmds, std.mem.asBytes(&p), vk.ops.dispatchGroups(plan.count, pipe.local[0]), c.timeout_ns)); + } + + // Attention: Q/K/V, gamma, RoPE, KV at `base`, GQA SDPA, gate/O; every activation store takes the layer's kind. + pub fn fullLayer(c: *const Core, x: View, layer: weights.FullLayer, s: *state.State, slot: usize, rows: usize, r_out: View, base: u64) !void { + const mark = c.scratch.mark(); + errdefer c.scratch.reset(mark) catch {}; + const hidden = pinned.hidden; + const act = x.kind; + const qraw = try c.scratch.alloc(act, rows * 2 * pinned.query_heads * pinned.head_dim); + try c.project(qraw, x, layer.q); + const q = try c.scratch.alloc(act, rows * pinned.query_heads * pinned.head_dim); + try c.extractQueries(q, qraw, rows); + const qn = try c.scratch.alloc(act, q.elements); + try vk.orExit(c.acc.norm.rmsNorm(c.cmds, qn, q, layer.q_norm, rows * pinned.query_heads, pinned.head_dim, pinned.eps, c.timeout_ns)); + const k = try c.scratch.alloc(act, rows * pinned.kv_heads * pinned.head_dim); + try c.project(k, x, layer.k); + const kn = try c.scratch.alloc(act, k.elements); + try vk.orExit(c.acc.norm.rmsNorm(c.cmds, kn, k, layer.k_norm, rows * pinned.kv_heads, pinned.head_dim, pinned.eps, c.timeout_ns)); + const v = try c.scratch.alloc(act, k.elements); + try c.project(v, x, layer.v); + const queries = try c.scratch.alloc(act, q.elements); + try vk.orExit(c.acc.attn.setBasePosition(c.cmds, c.staging, @intCast(base), c.timeout_ns)); + try vk.orExit(c.acc.attn.ropeQueries(c.cmds, queries, qn, rows, c.timeout_ns)); + try vk.orExit(c.acc.attn.cacheKV(c.cmds, s.keysFor(slot), s.valuesFor(slot), kn, v, rows, base, s.capacity, c.timeout_ns)); + const y = try c.scratch.alloc(act, q.elements); + try vk.orExit(c.acc.attn.attend(c.cmds, y, queries, s.keysFor(slot), s.valuesFor(slot), base, rows, s.capacity, c.timeout_ns)); + const mix = try c.scratch.alloc(act, rows * hidden); + try vk.orExit(c.acc.attn.gateOutput(c.cmds, mix, y, qraw, rows, c.timeout_ns)); + try c.project(r_out, mix, layer.o); + c.scratch.reset(mark) catch {}; + } + + // GDN: qkv/z/a/b, conv, gates/recurrence; every store takes the layer's kind; g stays f32. + pub fn linearLayer(c: *const Core, x: View, layer: weights.LinearLayer, s: *state.State, slot: usize, rows: usize, r_out: View) !void { + const mark = c.scratch.mark(); + errdefer c.scratch.reset(mark) catch {}; + const dim = pinned.linear_heads * pinned.linear_dim; + const act = x.kind; + const qkv = try c.scratch.alloc(act, rows * pinned.conv_dim); + try c.project(qkv, x, layer.qkv); + const z = try c.scratch.alloc(act, rows * dim); + try c.project(z, x, layer.z); + const a = try c.scratch.alloc(act, rows * pinned.linear_heads); + try c.project(a, x, layer.a); + const b = try c.scratch.alloc(act, rows * pinned.linear_heads); + try c.project(b, x, layer.b); + const soa = try c.scratch.alloc(act, rows * 3 * dim); + const g = try c.scratch.alloc(.f32, rows * pinned.linear_heads); + const beta = try c.scratch.alloc(act, rows * pinned.linear_heads); + const conv_out = try c.scratch.alloc(act, rows * pinned.conv_dim); + const conv_in = try c.scratch.alloc(act, (rows + state.conv_rows) * pinned.conv_dim); + try vk.orExit(c.acc.gdn.convPrefill(c.cmds, soa, g, beta, conv_out, conv_in, s.convInput(slot), qkv, layer.conv, a, b, layer.a_log, layer.dt_bias, s.convOutput(slot), rows, c.timeout_ns)); + const y = try c.scratch.alloc(act, rows * dim); + try vk.orExit(c.acc.gdn.prefillRecur(c.cmds, soa, g, beta, s.gdnInput(slot), y, s.gdnOutput(slot), null, rows, @min(rows, max_gdn_chunk), c.timeout_ns)); + const mix = try c.scratch.alloc(act, rows * dim); + try vk.orExit(c.acc.gdn.gatedNorm(c.cmds, mix, y, z, layer.gnorm, rows, pinned.eps, c.timeout_ns)); + try c.project(r_out, mix, layer.out); + s.flip(slot); + c.scratch.reset(mark) catch {}; + } + + // The dense MLP: Q4 gate and up projections, the one-round oracle SwiGLU, and the Q4 down projection into `r_out`. + pub fn mlp(c: *const Core, x: View, rows: usize, m: weights.DenseMlp, r_out: View) !void { + const mark = c.scratch.mark(); + errdefer c.scratch.reset(mark) catch {}; + const act = x.kind; + const gate = try c.scratch.alloc(act, rows * pinned.intermediate); + try c.project(gate, x, m.gate); + const up = try c.scratch.alloc(act, rows * pinned.intermediate); + try c.project(up, x, m.up); + const activated = try c.scratch.alloc(act, rows * pinned.intermediate); + try vk.orExit(c.acc.norm.swiglu(c.cmds, activated, gate, up, c.timeout_ns)); + try c.project(r_out, activated, m.down); + c.scratch.reset(mark) catch {}; + } + + // The last row gets residual add, raw-gain norm, and tied embedding or separate `lm_head` for greedy draw. + pub fn finish(c: *const Core, sum: View, h: View, r: View, x: View, final_norm: View, head: weights.Quant, logits: View, rows: usize) !void { + try vk.orExit(c.acc.norm.add(c.cmds, sum, h, r, c.timeout_ns)); + const last = (rows - 1) * pinned.hidden; + const x_last = try x.slice(last, pinned.hidden); + try vk.orExit(c.acc.norm.rmsNorm(c.cmds, x_last, try sum.slice(last, pinned.hidden), final_norm, 1, pinned.hidden, pinned.eps, c.timeout_ns)); + const mark = c.scratch.mark(); + defer c.scratch.reset(mark) catch {}; + try c.project(logits, x_last, head); + } +}; + +// `extractQPlan` is null on torn views; pure, host-tested; `heads_rows=rows*query_heads`, `in_head_stride=2*head_dim`. +pub const ExtractQPlan = struct { count: u32, dim: u32, heads_rows: u32, in_head_stride: u32 }; + +pub fn extractQPlan(rows: usize, out_elements: u64, in_elements: u64) ?ExtractQPlan { + if (rows == 0 or rows > max_prefill_rows) return null; + const dim: u64 = pinned.head_dim; + const heads: u64 = pinned.query_heads; + const per_row_out = heads * dim; + const per_row_in = 2 * heads * dim; + if (out_elements != rows * per_row_out or in_elements != rows * per_row_in) return null; + return .{ + .count = std.math.cast(u32, rows * per_row_out) orelse return null, + .dim = @intCast(dim), + .heads_rows = std.math.cast(u32, rows * heads) orelse return null, + .in_head_stride = @intCast(2 * dim), + }; +} + +const testing = std.testing; + +test "extractQPlan derives the interleaved geometry and refuses torn views" { + const plan = extractQPlan(3, 3 * 8 * 256, 3 * 4096).?; + try testing.expectEqual(@as(u32, 3 * 2048), plan.count); + try testing.expectEqual(@as(u32, 256), plan.dim); + try testing.expectEqual(@as(u32, 24), plan.heads_rows); + try testing.expectEqual(@as(u32, 512), plan.in_head_stride); + try testing.expectEqual(@as(?ExtractQPlan, null), extractQPlan(0, 0, 0)); + try testing.expectEqual(@as(?ExtractQPlan, null), extractQPlan(33, 33 * 2048, 33 * 4096)); // past the prefill cap + try testing.expectEqual(@as(?ExtractQPlan, null), extractQPlan(2, 2 * 2048, 2 * 4095)); + try testing.expectEqual(@as(?ExtractQPlan, null), extractQPlan(2, 2 * 2047, 2 * 4096)); +} diff --git a/zig/src/families/qwen3_5/backend/vulkan/qwen35_2b_config.json b/zig/src/families/qwen3_5/backend/vulkan/qwen35_2b_config.json new file mode 100644 index 00000000000..6d901903a28 --- /dev/null +++ b/zig/src/families/qwen3_5/backend/vulkan/qwen35_2b_config.json @@ -0,0 +1,106 @@ +{ + "architectures": [ + "Qwen3_5ForConditionalGeneration" + ], + "image_token_id": 248056, + "model_type": "qwen3_5", + "quantization": { + "group_size": 64, + "bits": 4, + "mode": "affine" + }, + "quantization_config": { + "group_size": 64, + "bits": 4, + "mode": "affine" + }, + "text_config": { + "attention_bias": false, + "attention_dropout": 0.0, + "attn_output_gate": true, + "dtype": "bfloat16", + "eos_token_id": 248044, + "full_attention_interval": 4, + "head_dim": 256, + "hidden_act": "silu", + "hidden_size": 2048, + "initializer_range": 0.02, + "intermediate_size": 6144, + "layer_types": [ + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention" + ], + "linear_conv_kernel_dim": 4, + "linear_key_head_dim": 128, + "linear_num_key_heads": 16, + "linear_num_value_heads": 16, + "linear_value_head_dim": 128, + "max_position_embeddings": 262144, + "mlp_only_layers": [], + "model_type": "qwen3_5_text", + "mtp_num_hidden_layers": 1, + "mtp_use_dedicated_embeddings": false, + "num_attention_heads": 8, + "num_hidden_layers": 24, + "num_key_value_heads": 2, + "rms_norm_eps": 1e-06, + "tie_word_embeddings": true, + "use_cache": true, + "vocab_size": 248320, + "mamba_ssm_dtype": "float32", + "rope_parameters": { + "mrope_interleaved": true, + "mrope_section": [ + 11, + 11, + 10 + ], + "rope_type": "default", + "rope_theta": 10000000, + "partial_rotary_factor": 0.25 + } + }, + "tie_word_embeddings": true, + "transformers_version": "4.57.0.dev0", + "video_token_id": 248057, + "vision_config": { + "deepstack_visual_indexes": [], + "depth": 24, + "hidden_act": "gelu_pytorch_tanh", + "hidden_size": 1024, + "in_channels": 3, + "initializer_range": 0.02, + "intermediate_size": 4096, + "model_type": "qwen3_5", + "num_heads": 16, + "num_position_embeddings": 2304, + "out_hidden_size": 2048, + "patch_size": 16, + "spatial_merge_size": 2, + "temporal_patch_size": 2 + }, + "vision_end_token_id": 248054, + "vision_start_token_id": 248053 +} \ No newline at end of file diff --git a/zig/src/families/qwen3_5/backend/vulkan/state.zig b/zig/src/families/qwen3_5/backend/vulkan/state.zig new file mode 100644 index 00000000000..79728af59ed --- /dev/null +++ b/zig/src/families/qwen3_5/backend/vulkan/state.zig @@ -0,0 +1,351 @@ +// Qwen3.5-2B state stores six full-attention KV caches and eighteen linear-layer conv/GDN states in typed device views. +const std = @import("std"); +const vk = @import("vulkan"); +const pinned = @import("../../config.zig"); + +const Allocator = std.mem.Allocator; +const Buffer = vk.Buffer; +const Commands = vk.Commands; +const Device = vk.Device; +const View = vk.tensor.View; + +pub const Error = vk.runtime.Error || error{ContextTooLong}; + +// Full-attention layers are indices 3, 7, ...; divide by four after checking the pinned interval. +pub const full_layers: usize = pinned.layers / 4; +pub const linear_layers: usize = pinned.layers - full_layers; + +// SDPA shared memory admits visible prefixes up to 7168 keys. +pub const max_capacity: usize = 7168; + +pub const kv_heads: usize = pinned.kv_heads; // 2 +pub const head_dim: usize = pinned.head_dim; // 256 +pub const conv_rows: usize = pinned.conv_taps - 1; // 3 +pub const conv_channels: usize = pinned.conv_dim; // 6144 +pub const gdn_heads: usize = pinned.linear_heads; // 16 +pub const gdn_head_dim: usize = pinned.linear_dim; // 128 + +/// Layers 0..wide_layers-1 keep f32 activation stores; their conv windows are f32 too, every other layer's stay bf16. +pub const wide_layers: usize = 4; + +/// The conv window kind of linear slot `slot` (layers 0-2 are slots 0-2; layer 3 is the first full layer). +pub fn convKind(slot: usize) vk.tensor.Kind { + return if (slot < wide_layers - wide_layers / 4) .f32 else .bf16; +} + +/// The whole-model capacity: the context pin admits, the SDPA limit bounds. +pub fn capacityFor(context: usize) Error!usize { + if (context == 0 or context > max_capacity) return error.ContextTooLong; + return context; +} + +/// Layer `index`'s full-attention cache slot (indices 3, 7, ... map to 0..5). +pub fn fullSlot(index: usize) ?usize { + if (index % 4 != 3 or index >= pinned.layers) return null; + return (index + 1) / 4 - 1; +} + +/// Layer `index`'s linear-layer slot (the other eighteen indices, in order). +pub fn linearSlot(index: usize) ?usize { + if (index >= pinned.layers or index % 4 == 3) return null; + return index - index / 4; +} + +// Heap-owned buffers keep views stable; parity selects each linear layer's ping-pong input half. +pub const State = struct { + gpa: Allocator, + dev: *const Device, + capacity: u64, + /// The committed visible prefix; only `append` and `reset` change it. + committed: u64 = 0, + buffers: []*Buffer, + keys: []View, + values: []View, + /// Two halves per linear layer, even = half 0; `parity` names the input. + conv: []View, + gdn: []View, + parity: []u1, + // How many of `buffers` hold real created buffers; init's error path frees exactly this prefix. + created: usize, + + // Allocate and zero min(context, max_capacity) positions; staged transfers chunk through caller storage. + pub fn init( + gpa: Allocator, + dev: *const Device, + cmds: *Commands, + context: usize, + staging: Buffer, + timeout_ns: u64, + ) (Allocator.Error || Error)!State { + const capacity = try capacityFor(context); + + var s: State = .{ + .gpa = gpa, + .dev = dev, + .capacity = capacity, + .buffers = &.{}, + .keys = &.{}, + .values = &.{}, + .conv = &.{}, + .gdn = &.{}, + .parity = &.{}, + .created = 0, + }; + errdefer s.deinit(); + + // Record each completed allocation so the error path frees exactly the created prefix. + s.buffers = try gpa.alloc(*Buffer, 2 * full_layers + 4 * linear_layers); + s.keys = try gpa.alloc(View, full_layers); + s.values = try gpa.alloc(View, full_layers); + s.conv = try gpa.alloc(View, 2 * linear_layers); + s.gdn = try gpa.alloc(View, 2 * linear_layers); + s.parity = try gpa.alloc(u1, linear_layers); + @memset(s.parity, 0); + + const kv_elements: u64 = kv_heads * capacity * head_dim; + const conv_elements: u64 = conv_rows * conv_channels; + const gdn_elements: u64 = gdn_heads * gdn_head_dim * gdn_head_dim; + var next: usize = 0; + for (0..full_layers) |slot| { + s.keys[slot] = try s.region(next, .bf16, kv_elements); + next += 1; + s.values[slot] = try s.region(next, .bf16, kv_elements); + next += 1; + } + for (0..linear_layers) |slot| { + s.conv[2 * slot] = try s.region(next, convKind(slot), conv_elements); + next += 1; + s.conv[2 * slot + 1] = try s.region(next, convKind(slot), conv_elements); + next += 1; + } + for (0..linear_layers) |slot| { + s.gdn[2 * slot] = try s.region(next, .f32, gdn_elements); + next += 1; + s.gdn[2 * slot + 1] = try s.region(next, .f32, gdn_elements); + next += 1; + } + try s.reset(cmds, staging, timeout_ns); + s.created = s.buffers.len; + return s; + } + + // Register the buffer as soon as Vulkan allocates it; State.deinit remains its sole owner. + fn region(s: *State, at: usize, kind: vk.tensor.Kind, elements: u64) (Allocator.Error || vk.runtime.Error)!View { + const bytes = try vk.tensor.extent(kind, elements); + const ptr = try s.gpa.create(Buffer); + ptr.* = Buffer.init(s.dev, .device, bytes) catch |e| { + s.gpa.destroy(ptr); + return e; + }; + s.buffers[at] = ptr; + s.created = at + 1; + return View.init(ptr, kind, 0, elements); + } + + // Zero every device region through staged transfer using at most a 16 MiB host buffer. + pub fn reset(s: *State, cmds: *Commands, staging: Buffer, timeout_ns: u64) Error!void { + const zeros = try s.gpa.alloc(u8, @min(1 << 24, staging.len)); + defer s.gpa.free(zeros); + @memset(zeros, 0); + for (s.buffers) |ptr| { + var done: u64 = 0; + while (done < ptr.len) { + const n = @min(zeros.len, ptr.len - done); + // in-flight failures exit here, never unwinding into frees + try vk.orExit(vk.upload(cmds, ptr.*, done, staging, zeros[0..n], timeout_ns)); + done += n; + } + } + s.committed = 0; + @memset(s.parity, 0); + } + + /// The full-attention cache slot's key rows, `[2, capacity, 256]` BF16. + pub fn keysFor(s: *const State, slot: usize) View { + return s.keys[slot]; + } + + /// The full-attention cache slot's value rows, same shape. + pub fn valuesFor(s: *const State, slot: usize) View { + return s.values[slot]; + } + + /// The linear slot's conv window the next step reads, `[3, 6144]` in `convKind(slot)`. + pub fn convInput(s: *const State, slot: usize) View { + return s.conv[2 * slot + s.parity[slot]]; + } + + /// The disjoint half the next step writes, same shape. + pub fn convOutput(s: *const State, slot: usize) View { + return s.conv[2 * slot + (s.parity[slot] ^ 1)]; + } + + /// The linear slot's GDN recurrence the next step reads, `[16, 128, 128]` F32. + pub fn gdnInput(s: *const State, slot: usize) View { + return s.gdn[2 * slot + s.parity[slot]]; + } + + /// The disjoint half the next step writes, same shape. + pub fn gdnOutput(s: *const State, slot: usize) View { + return s.gdn[2 * slot + (s.parity[slot] ^ 1)]; + } + + /// After a step's writes completed: the written half becomes the input. + pub fn flip(s: *State, slot: usize) void { + s.parity[slot] ^= 1; + } + + /// The committed visible prefix (what attention may attend over). + pub fn visible(s: *const State) u64 { + return s.committed; + } + + // Append refuses capacity overflow; recurrence makes rewind invalid, so speculation requires restoring state. + pub fn append(s: *State, rows: u64) Error!void { + if (rows > s.capacity - s.committed) return error.ContextTooLong; + s.committed += rows; + } + + // Frees every device buffer. Call only when no submission touching any view can still be pending. + pub fn deinit(s: *State) void { + for (s.buffers[0..s.created]) |ptr| { + var buf = ptr.*; + buf.deinit(); + s.gpa.destroy(ptr); + } + s.gpa.free(s.buffers); + s.gpa.free(s.keys); + s.gpa.free(s.values); + s.gpa.free(s.conv); + s.gpa.free(s.gdn); + s.gpa.free(s.parity); + s.* = undefined; + } +}; + +const testing = std.testing; + +fn fakeState(gpa: Allocator, capacity: usize) !State { + // content-free state for layout tests: fake buffers, no GPU behind them + const kv_elements: u64 = kv_heads * capacity * head_dim; + const conv_elements: u64 = conv_rows * conv_channels; + const gdn_elements: u64 = gdn_heads * gdn_head_dim * gdn_head_dim; + const s: State = .{ + .gpa = gpa, + .dev = undefined, + .capacity = capacity, + .buffers = try gpa.alloc(*Buffer, 2 * full_layers + 4 * linear_layers), + .keys = try gpa.alloc(View, full_layers), + .values = try gpa.alloc(View, full_layers), + .conv = try gpa.alloc(View, 2 * linear_layers), + .gdn = try gpa.alloc(View, 2 * linear_layers), + .parity = try gpa.alloc(u1, linear_layers), + .created = 0, + }; + @memset(s.parity, 0); + const conv_bytes = 2 * conv_elements; + const gdn_bytes = 4 * gdn_elements; + // Fake views use two buffers sized for the largest F32 GDN region, preserving ping-pong separation. + const fakes: [2]*Buffer = .{ try gpa.create(Buffer), try gpa.create(Buffer) }; + defer for (fakes) |ptr| gpa.destroy(ptr); + for (fakes) |ptr| ptr.* = .{ .dev = undefined, .handle = null, .memory = null, .len = gdn_bytes, .kind = .device, .mapped = null, .coherent = false }; + for (s.buffers, 0..) |*ptr, i| { + ptr.* = fakes[i % 2]; + } + const kb: *Buffer = fakes[0]; + const kb2: *Buffer = fakes[1]; + for (s.keys) |*v| v.* = try View.init(kb, .bf16, 0, kv_elements); + for (s.values) |*v| v.* = try View.init(kb, .bf16, 0, kv_elements); + for (s.conv, 0..) |*v, i| v.* = try View.init(if (i % 2 == 0) kb else kb2, convKind(i / 2), 0, conv_elements); + for (s.gdn, 0..) |*v, i| v.* = try View.init(if (i % 2 == 0) kb else kb2, .f32, 0, gdn_elements); + _ = conv_bytes; + return s; +} + +test "capacity is the context pinned, or refused past the SDPA limit" { + try testing.expectEqual(@as(usize, 4096), try capacityFor(4096)); + try testing.expectEqual(@as(usize, max_capacity), try capacityFor(max_capacity)); + // a context beyond the SDPA arm's limit is refused, never truncated + try testing.expectError(error.ContextTooLong, capacityFor(0)); + try testing.expectError(error.ContextTooLong, capacityFor(max_capacity + 1)); + try testing.expectError(error.ContextTooLong, capacityFor(262144)); +} + +test "layer slots partition the twenty-four layers exactly" { + var fulls: usize = 0; + var linears: usize = 0; + for (0..pinned.layers) |i| { + if (fullSlot(i)) |slot| { + try testing.expectEqual(fulls, slot); + fulls += 1; + try testing.expectEqual(@as(?usize, null), linearSlot(i)); + } else { + try testing.expectEqual(linears, linearSlot(i).?); + linears += 1; + } + } + try testing.expectEqual(@as(usize, full_layers), fulls); + try testing.expectEqual(@as(usize, linear_layers), linears); + try testing.expectEqual(@as(usize, 18), linear_layers); + try testing.expectEqual(@as(?usize, null), fullSlot(pinned.layers)); + try testing.expectEqual(@as(?usize, null), linearSlot(pinned.layers)); +} + +test "cache regions carry the pinned shapes and stay inside their buffers" { + const gpa = testing.allocator; + const capacity = 64; + var s = try fakeState(gpa, capacity); + defer { + gpa.free(s.buffers); + gpa.free(s.keys); + gpa.free(s.values); + gpa.free(s.conv); + gpa.free(s.gdn); + gpa.free(s.parity); + } + try testing.expectEqual(@as(u64, kv_heads * capacity * head_dim), s.keysFor(0).elements); + try testing.expectEqual(vk.tensor.Kind.bf16, s.keysFor(5).kind); + try testing.expectEqual(vk.tensor.Kind.bf16, s.valuesFor(5).kind); + try testing.expectEqual(@as(u64, conv_rows * conv_channels), s.convInput(17).elements); + try testing.expectEqual(vk.tensor.Kind.f32, s.gdnInput(0).kind); + try testing.expectEqual(@as(u64, gdn_heads * gdn_head_dim * gdn_head_dim), s.gdnInput(0).elements); + // ping-pong halves are disjoint views of different buffers + try testing.expect(!s.convInput(3).overlaps(s.convOutput(3))); + try testing.expect(!s.gdnInput(9).overlaps(s.gdnOutput(9))); +} + +test "convKind is f32 exactly for the linear slots of the wide layers" { + // layers 0-2 are linear slots 0-2 (layer 3 is full attention); layer 4 is slot 3 + try testing.expectEqual(vk.tensor.Kind.f32, convKind(0)); + try testing.expectEqual(vk.tensor.Kind.f32, convKind(2)); + try testing.expectEqual(vk.tensor.Kind.bf16, convKind(3)); +} + +test "flip swaps the halves and append is the only length commit" { + const gpa = testing.allocator; + var s = try fakeState(gpa, 32); + defer { + gpa.free(s.buffers); + gpa.free(s.keys); + gpa.free(s.values); + gpa.free(s.conv); + gpa.free(s.gdn); + gpa.free(s.parity); + } + const in0 = s.convInput(0); + const out0 = s.convOutput(0); + try s.append(7); + try testing.expectEqual(@as(u64, 7), s.visible()); + s.flip(0); + try testing.expectEqual(out0.byte_offset, s.convInput(0).byte_offset); + try testing.expectEqual(in0.byte_offset, s.convOutput(0).byte_offset); + s.flip(0); + try testing.expectEqual(in0.byte_offset, s.convInput(0).byte_offset); + // the capacity is a hard commit boundary + try testing.expectError(error.ContextTooLong, s.append(26)); + try s.append(25); + try testing.expectEqual(@as(u64, 32), s.visible()); + try testing.expectError(error.ContextTooLong, s.append(1)); + // no rewind exists: visible length only ever grows until a real reset + try testing.expectEqual(@as(u64, 32), s.visible()); +} diff --git a/zig/src/families/qwen3_5/backend/vulkan/weights.zig b/zig/src/families/qwen3_5/backend/vulkan/weights.zig new file mode 100644 index 00000000000..f8183c6825a --- /dev/null +++ b/zig/src/families/qwen3_5/backend/vulkan/weights.zig @@ -0,0 +1,423 @@ +// Qwen3.5-2B Vulkan: fail-closed Q4 admission before reads; raw BF16 norms; heap views; staged uploads use orExit. +const std = @import("std"); +const vk = @import("vulkan"); +const pinned = @import("../../config.zig"); +const loader = @import("../../loader.zig"); +const Allocator = std.mem.Allocator; +const Buffer = vk.Buffer; +const Commands = vk.Commands; +const Device = vk.Device; +const Table = loader.table.Table; +const Tensor = loader.table.Tensor; +const View = vk.tensor.View; + +pub const Error = loader.table.Error || pinned.Error || vk.runtime.Error || + error{UnsupportedModel}; + +/// The text tower's key prefix, as the shared loader names it. +const prefix = loader.checkpoint.prefix; + +// A `Quant` stores packed affine words, per-group scales and biases, plus row/column geometry checked by qmm. +pub const Quant = struct { + w: View, + scales: View, + biases: View, + n: u64, + k: u64, + /// Words, scales and biases are repacked by `vk.tile` for the tiled qmm image. + tiled: bool = false, +}; + +pub const DenseMlp = struct { gate: Quant, up: Quant, down: Quant }; + +// Full attention uses Q4/raw BF16; `q` has 4096 interleaved query/gate rows (8×256), split by forward to match Metal. +pub const FullLayer = struct { + input_norm: View, + post_norm: View, + q: Quant, + k: Quant, + v: Quant, + o: Quant, + q_norm: View, + k_norm: View, + mlp: DenseMlp, +}; + +// GatedDeltaNet keeps raw BF16 taps/norms/`dt_bias`; `conv_dw1d` widens uint16, and `A_log` stays raw F32. +pub const LinearLayer = struct { + input_norm: View, + post_norm: View, + qkv: Quant, + z: Quant, + a: Quant, + b: Quant, + out: Quant, + gnorm: View, + conv: View, + a_log: View, + dt_bias: View, + mlp: DenseMlp, +}; + +pub const Layer = union(enum) { full: FullLayer, linear: LinearLayer }; + +// The one supported model, checked field by field against the pinned family constants before anything is trusted. +pub fn expectPinnedModel(s: loader.config.Spec) Error!void { + const expected = .{ + .hidden = pinned.hidden, + .vocab = pinned.vocab, + .n_layers = pinned.layers, + .intermediate = pinned.intermediate, + .heads = pinned.query_heads, + .kv_heads = pinned.kv_heads, + .head_dim = pinned.head_dim, + .key_heads = pinned.linear_heads, + .value_heads = pinned.linear_heads, + // Spec stores per-head dimensions and conv taps, not channel widths. + .key_dim = pinned.linear_dim, + .value_dim = pinned.linear_dim, + .conv = pinned.conv_taps, + .rotary_dim = pinned.rotary_dim, + .bits = 4, + .group = pinned.group, + .experts = 0, + }; + inline for (@typeInfo(@TypeOf(expected)).@"struct".field_names) |field| { + if (@field(s, field) != @field(expected, field)) { + std.debug.print("Vulkan weights admission: Spec.{s}: expected {d}, found {d}\n", .{ field, @field(expected, field), @field(s, field) }); + return error.UnsupportedModel; + } + } +} + +comptime { + if (pinned.layers % 4 != 0) @compileError("the pinned layer count must keep the full-attention interval exact"); + if (pinned.conv_dim != 2 * pinned.linear_heads * pinned.linear_dim + pinned.linear_heads * pinned.linear_dim) + @compileError("conv channels drifted from 2*NK*DK + NV*DV"); + if (pinned.query_heads * pinned.head_dim * 2 != 4096) @compileError("qraw row count drifted from the interleaved q/gate layout"); +} + +fn kindOf(dt: loader.table.DType) ?vk.tensor.Kind { + return switch (dt) { + .f32 => .f32, + .bf16 => .bf16, + .u32 => .u32, + else => null, + }; +} + +/// A checked exact dtype/shape, logged by key when it refuses. +fn expectTensor(key: []const u8, t: Tensor, dt: loader.table.DType, shape: []const usize) Error!void { + if (!t.is(dt, shape)) { + std.log.err("{s}: expected {s} {any}, found {s} {any}", .{ key, @tagName(dt), shape, @tagName(t.dtype), t.shape[0..t.rank] }); + return error.UnsupportedModel; + } +} + +// `Builder` records heap-owned buffers; `deinit` frees created buffers on errors, preventing leaks and double-frees. +const Builder = struct { + gpa: Allocator, + dev: *const Device, + cmds: *Commands, + staging: Buffer, + timeout_ns: u64, + table: *const Table, + keep: std.ArrayList(*Buffer) = .empty, + + fn deinit(b: *Builder) void { + for (b.keep.items) |ptr| { + var buf = ptr.*; + buf.deinit(); + b.gpa.destroy(ptr); + } + b.keep.deinit(b.gpa); + b.* = undefined; + } + + fn key(buf: *[192]u8, comptime fmt: []const u8, args: anytype) Error![]const u8 { + // Layer prefixes can point into buf: format elsewhere before copying. + var scratch: [192]u8 = undefined; + const name = std.fmt.bufPrint(&scratch, fmt, args) catch return error.UnsupportedModel; + @memcpy(buf[0..name.len], name); + return buf[0..name.len]; + } + + /// The named tensor, dtype-checked before its mapped bytes are touched. + fn named(b: *Builder, name: []const u8) Error!Tensor { + const t = b.table.get(name) catch |err| { + std.debug.print("Vulkan weights tensor {s}: {s}\n", .{ name, @errorName(err) }); + return err; + }; + switch (t.dtype) { + .f32, .bf16, .u32 => {}, + else => { + std.log.err("{s}: unsupported tensor dtype {s}", .{ name, @tagName(t.dtype) }); + return error.UnsupportedModel; + }, + } + return t; + } + + /// Uploads go through staging; `View` pads words; `Builder.deinit` owns slots after `keep`. + fn tensor(b: *Builder, t: Tensor) !View { + return b.tensorBytes(t, t.bytes, @intCast(t.numel())); + } + + /// The tensor's dtype with `bytes` as its device contents, viewed as `elements` of that dtype. + fn tensorBytes(b: *Builder, t: Tensor, bytes: []const u8, elements: u64) !View { + const kind = kindOf(t.dtype) orelse return error.UnsupportedModel; + const ptr = try b.gpa.create(Buffer); + ptr.* = Buffer.init(b.dev, .device, bytes.len) catch |e| { + b.gpa.destroy(ptr); + return e; + }; + b.keep.append(b.gpa, ptr) catch |e| { + var buf = ptr.*; + buf.deinit(); + b.gpa.destroy(ptr); + return e; + }; + // Registered slots belong to `Builder.deinit`; `orExit` prevents in-flight failures from unwinding into frees. + try vk.orExit(vk.upload(b.cmds, ptr.*, 0, b.staging, bytes, b.timeout_ns)); + return View.init(ptr, kind, 0, elements); + } + + /// A raw tensor of the exact given dtype and shape, uploaded unchanged. + fn raw(b: *Builder, name: []const u8, dt: loader.table.DType, shape: []const usize) !View { + const t = try b.named(name); + try expectTensor(name, t, dt, shape); + return b.tensor(t); + } + + // For `n` rows, `k` inputs, group 64: checked U32 `[n,k/8]` words and BF16 scales/biases `[n,k/64]`. + fn quant(b: *Builder, name: []const u8, n: usize, k: usize) !Quant { + var buf: [192]u8 = undefined; + const w = try b.raw(try key(&buf, "{s}.weight", .{name}), .u32, &.{ n, k / 8 }); + const scales = try b.raw(try key(&buf, "{s}.scales", .{name}), .bf16, &.{ n, k / 64 }); + const biases = try b.raw(try key(&buf, "{s}.biases", .{name}), .bf16, &.{ n, k / 64 }); + return .{ .w = w, .scales = scales, .biases = biases, .n = n, .k = k }; + } + + // Tiled when k fills whole groups, the padded words fit one binding, and the arena keeps x 16-byte aligned. + fn tilable(b: *const Builder, n: usize, k: usize) bool { + const limits = b.dev.properties.limits; + return vk.tile.fits(n, k) and vk.tile.padded(n) * k / 2 <= limits.maxStorageBufferRange and limits.minStorageBufferOffsetAlignment % 16 == 0; + } + + // One tensor of a tiled projection: checked like `raw`, then repacked on the host into zero-padded 32-column tiles. + fn repacked(b: *Builder, name: []const u8, dt: loader.table.DType, shape: []const usize, n: usize, k: usize, words: bool) !View { + const t = try b.named(name); + try expectTensor(name, t, dt, shape); + const moved = try b.gpa.alloc(u8, vk.tile.padded(n) * k / @as(usize, if (words) 2 else 32)); + defer b.gpa.free(moved); + if (words) vk.tile.tileWords(moved, t.bytes, n, k) else vk.tile.tileParams(moved, t.bytes, n, k); + return b.tensorBytes(t, moved, moved.len / @as(usize, if (words) 4 else 2)); + } + + // As `quant`, tiled when the shape fits; the tiled image reads bf16 and f32 activations alike. + fn quantTiled(b: *Builder, name: []const u8, n: usize, k: usize) !Quant { + if (!b.tilable(n, k)) return b.quant(name, n, k); + var buf: [192]u8 = undefined; + const w = try b.repacked(try key(&buf, "{s}.weight", .{name}), .u32, &.{ n, k / 8 }, n, k, true); + const scales = try b.repacked(try key(&buf, "{s}.scales", .{name}), .bf16, &.{ n, k / 64 }, n, k, false); + const biases = try b.repacked(try key(&buf, "{s}.biases", .{name}), .bf16, &.{ n, k / 64 }, n, k, false); + return .{ .w = w, .scales = scales, .biases = biases, .n = n, .k = k, .tiled = true }; + } + + // Raw BF16 norm gains pass byte-for-byte; the Metal kernel multiplies gain directly, never `1 + w`. + fn norm(b: *Builder, name: []const u8, elements: usize) !View { + return b.raw(name, .bf16, &.{elements}); + } + + fn mlp(b: *Builder, base: []const u8) !DenseMlp { + var buf: [192]u8 = undefined; + return .{ + .gate = try b.quantTiled(try key(&buf, "{s}mlp.gate_proj", .{base}), pinned.intermediate, pinned.hidden), + .up = try b.quantTiled(try key(&buf, "{s}mlp.up_proj", .{base}), pinned.intermediate, pinned.hidden), + .down = try b.quantTiled(try key(&buf, "{s}mlp.down_proj", .{base}), pinned.hidden, pinned.intermediate), + }; + } + + fn layer(b: *Builder, index: usize) !Layer { + var buf: [192]u8 = undefined; + const base = try key(&buf, prefix ++ "layers.{d}.", .{index}); + const input_norm = try b.norm(try key(&buf, "{s}input_layernorm.weight", .{base}), pinned.hidden); + const post_norm = try b.norm(try key(&buf, "{s}post_attention_layernorm.weight", .{base}), pinned.hidden); + const m = try b.mlp(base); + if (pinned.linear(index)) { + const lin = try key(&buf, "{s}linear_attn.", .{base}); + return .{ .linear = .{ + .input_norm = input_norm, + .post_norm = post_norm, + .qkv = try b.quantTiled(try key(&buf, "{s}in_proj_qkv", .{lin}), 3 * pinned.linear_heads * pinned.linear_dim, pinned.hidden), + .z = try b.quantTiled(try key(&buf, "{s}in_proj_z", .{lin}), pinned.linear_heads * pinned.linear_dim, pinned.hidden), + .a = try b.quantTiled(try key(&buf, "{s}in_proj_a", .{lin}), pinned.linear_heads, pinned.hidden), + .b = try b.quantTiled(try key(&buf, "{s}in_proj_b", .{lin}), pinned.linear_heads, pinned.hidden), + .out = try b.quantTiled(try key(&buf, "{s}out_proj", .{lin}), pinned.hidden, pinned.linear_heads * pinned.linear_dim), + .gnorm = try b.norm(try key(&buf, "{s}norm.weight", .{lin}), pinned.linear_dim), + .conv = try b.raw(try key(&buf, "{s}conv1d.weight", .{lin}), .bf16, &.{ pinned.conv_dim, pinned.conv_taps, 1 }), + .a_log = try b.raw(try key(&buf, "{s}A_log", .{lin}), .f32, &.{pinned.linear_heads}), + .dt_bias = try b.raw(try key(&buf, "{s}dt_bias", .{lin}), .bf16, &.{pinned.linear_heads}), + .mlp = m, + } }; + } + const attn = try key(&buf, "{s}self_attn.", .{base}); + return .{ .full = .{ + .input_norm = input_norm, + .post_norm = post_norm, + .q = try b.quantTiled(try key(&buf, "{s}q_proj", .{attn}), 2 * pinned.query_heads * pinned.head_dim, pinned.hidden), + .k = try b.quantTiled(try key(&buf, "{s}k_proj", .{attn}), pinned.kv_heads * pinned.head_dim, pinned.hidden), + .v = try b.quantTiled(try key(&buf, "{s}v_proj", .{attn}), pinned.kv_heads * pinned.head_dim, pinned.hidden), + .o = try b.quantTiled(try key(&buf, "{s}o_proj", .{attn}), pinned.hidden, pinned.query_heads * pinned.head_dim), + .q_norm = try b.norm(try key(&buf, "{s}q_norm.weight", .{attn}), pinned.head_dim), + .k_norm = try b.norm(try key(&buf, "{s}k_norm.weight", .{attn}), pinned.head_dim), + .mlp = m, + } }; + } +}; + +// Typed text-tower views use owned heap buffers; checkpoint config decides whether the head ties to the embedding. +pub const Model = struct { + gpa: Allocator, + spec: loader.config.Spec, + /// config.json's max_position_embeddings, as admission pinned it. + context: usize, + /// The head is the packed input embedding (`tie_word_embeddings`). + tied: bool, + buffers: []*Buffer, + embed: Quant, + final_norm: View, + // When tied, `outputHead` returns `embed`; untied checkpoints require a separate head or loading is refused. + head: ?Quant, + layers: []Layer, + + // Loads `dir` via `cmds` with nonempty `staging`; it checks admission first and keeps heap-backed views valid. + pub fn load( + gpa: Allocator, + dev: *const Device, + cmds: *Commands, + io: std.Io, + dir: []const u8, + staging: Buffer, + timeout_ns: u64, + ) !Model { + // Admission pins model type, geometry, quantization, and tying from `config.json` before other reads. + const cfg = pinned.Config.read(gpa, io, dir) catch |err| { + std.debug.print("Vulkan weights config.json admission: {s}\n", .{@errorName(err)}); + return err; + }; + var ck = loader.checkpoint.Checkpoint.open(gpa, io, dir) catch |err| { + std.debug.print("Vulkan weights Checkpoint.open: {s}\n", .{@errorName(err)}); + return err; + }; + defer ck.close(); + try expectPinnedModel(ck.spec()); + + var b: Builder = .{ + .gpa = gpa, + .dev = dev, + .cmds = cmds, + .staging = staging, + .timeout_ns = timeout_ns, + .table = &ck.main, + }; + defer b.deinit(); + + var layers: std.ArrayList(Layer) = .empty; + errdefer layers.deinit(gpa); + + // The head is tiled; tied, it is a second upload of the embedding table (the embedding keeps its row layout). + const head: ?Quant = if (!b.tilable(pinned.vocab, pinned.hidden)) + (if (cfg.tied) null else try b.quant("language_model.lm_head", pinned.vocab, pinned.hidden)) + else + try b.quantTiled(if (cfg.tied) prefix ++ "embed_tokens" else "language_model.lm_head", pinned.vocab, pinned.hidden); + + const embed = try b.quant(prefix ++ "embed_tokens", pinned.vocab, pinned.hidden); + const final_norm = try b.norm(prefix ++ "norm.weight", pinned.hidden); + // Mapped pages are released per part, so host RSS holds one checkpoint layer, not the full file set. + ck.release(null); + for (0..pinned.layers) |i| { + try layers.append(gpa, try b.layer(i)); + ck.release(i); + } + + const owned_layers = try layers.toOwnedSlice(gpa); + errdefer gpa.free(owned_layers); + // ownership moves to the model; the builder's defer is now a no-op + const buffers = try b.keep.toOwnedSlice(gpa); + return .{ + .gpa = gpa, + .spec = ck.spec(), + .context = cfg.context, + .tied = cfg.tied, + .buffers = buffers, + .embed = embed, + .final_norm = final_norm, + .head = head, + .layers = owned_layers, + }; + } + + /// The output head: the embedding itself when the checkpoint ties them. + pub fn outputHead(m: *const Model) Quant { + return m.head orelse m.embed; + } + + // Free device buffers only when no submission can touch a view; in-flight errors exit through `orExit`. + pub fn deinit(m: *Model) void { + for (m.buffers) |ptr| { + var buf = ptr.*; + buf.deinit(); + m.gpa.destroy(ptr); + } + m.gpa.free(m.buffers); + m.gpa.free(m.layers); + m.* = undefined; + } +}; + +const testing = std.testing; + +test "kindOf maps exactly the dtypes the views carry" { + try testing.expectEqual(@as(?vk.tensor.Kind, .f32), kindOf(.f32)); + try testing.expectEqual(@as(?vk.tensor.Kind, .bf16), kindOf(.bf16)); + try testing.expectEqual(@as(?vk.tensor.Kind, .u32), kindOf(.u32)); + try testing.expectEqual(@as(?vk.tensor.Kind, null), kindOf(.f16)); +} + +test "tensor keys can extend a prefix in the destination buffer" { + var buf: [192]u8 = undefined; + const base = try Builder.key(&buf, prefix ++ "layers.{d}.", .{0}); + const name = try Builder.key(&buf, "{s}input_layernorm.weight", .{base}); + try testing.expectEqualStrings(prefix ++ "layers.0.input_layernorm.weight", name); + const oversized: [193]u8 = @splat('x'); + try testing.expectError(error.UnsupportedModel, Builder.key(&buf, "{s}", .{oversized[0..]})); +} + +test "the real pinned config geometry is admitted and field drift is refused" { + const bytes = @embedFile("qwen35_2b_config.json"); + var digest: [32]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(bytes, &digest, .{}); + try testing.expectEqualStrings("beb7fc5a6e0405fe332821cf1a8ef7b69bb390a8c8933171647de5579debf949", &std.fmt.bytesToHex(digest, .lower)); + var arena = std.heap.ArenaAllocator.init(testing.allocator); + defer arena.deinit(); + const doc = try pinned.declared(arena.allocator(), bytes); + const s = try pinned.geometry(try pinned.textTower(doc), 4, 64); + try testing.expectEqual(@as(usize, 128), s.key_dim); + try testing.expectEqual(@as(usize, 128), s.value_dim); + try testing.expectEqual(@as(usize, 4), s.conv); + try expectPinnedModel(s); + inline for (.{ "hidden", "vocab", "n_layers", "intermediate", "heads", "kv_heads", "head_dim", "key_heads", "value_heads", "key_dim", "value_dim", "conv", "rotary_dim", "bits", "group", "experts" }) |field| { + var drift = s; + @field(drift, field) += 1; + try testing.expectError(error.UnsupportedModel, expectPinnedModel(drift)); + } +} + +test "tensor shape checks refuse rank and dim drift by name" { + const t: Tensor = .{ .dtype = .bf16, .rank = 1, .shape = .{ 2048, 0, 0, 0, 0 }, .bytes = &.{} }; + try testing.expect(t.is(.bf16, &.{2048})); + try testing.expect(!t.is(.f32, &.{2048})); + try testing.expect(!t.is(.bf16, &.{2049})); + try testing.expect(!t.is(.bf16, &.{ 2048, 1 })); +} diff --git a/zig/src/families/qwen3_5/vulkan.zig b/zig/src/families/qwen3_5/vulkan.zig new file mode 100644 index 00000000000..0236df9e3bd --- /dev/null +++ b/zig/src/families/qwen3_5/vulkan.zig @@ -0,0 +1,4 @@ +// Qwen3.5 root resolves backend ../../ imports; backend.zig cannot resolve them alone. +pub const backend = @import("backend/vulkan/backend.zig"); +pub const engine = @import("backend/vulkan/engine.zig"); +pub const config = @import("config.zig"); diff --git a/zig/src/families/qwen3_5/vulkan_test.zig b/zig/src/families/qwen3_5/vulkan_test.zig new file mode 100644 index 00000000000..7c4859dd562 --- /dev/null +++ b/zig/src/families/qwen3_5/vulkan_test.zig @@ -0,0 +1,5 @@ +test { + _ = @import("backend/vulkan/backend.zig"); + _ = @import("backend/vulkan/engine.zig"); + _ = @import("backend/vulkan/forward.zig"); +} diff --git a/zig/src/native/vulkan.zig b/zig/src/native/vulkan.zig new file mode 100644 index 00000000000..f246a39e3e6 --- /dev/null +++ b/zig/src/native/vulkan.zig @@ -0,0 +1,278 @@ +// A native server opens one Qwen3.5-2B dense model on Vulkan behind the lane host, one stream at a time, greedy only. +const std = @import("std"); +const api = @import("engine_api"); +const lanes = @import("lanes"); +const fam = @import("qwen35_vulkan"); // the family's thin vulkan root: backend and engine +const sh = @import("qwen35_shaders"); // the build's compiled SPIR-V images, one decl per shader +const Allocator = std.mem.Allocator; + +/// Every fence wait inside the engine is bound by this; the probes serve with the same bound. +const timeout_ns: u64 = 10 * std.time.ns_per_s; + +pub const backends: []const []const u8 = &.{"vulkan"}; +pub const serves: []const []const u8 = &.{"--backend=vulkan"}; +pub const families: []const api.Family = &.{ + .{ .model_type = "qwen3_5", .formats = &.{"mlx-q4g64"} }, +}; + +// GPU family gates name a chip class; Vulkan names no chip; `capabilities` reports none; all eligible devices open. +pub fn chip(a: Allocator) ?[]const u8 { + _ = a; + return null; +} + +/// A flag or a resource the engine cannot take: `problem` says why. +const Refused = error{Refused}; + +fn refuse(a: Allocator, problem: *[]const u8, comptime fmt: []const u8, args: anytype) Refused { + problem.* = std.fmt.allocPrint(a, fmt, args) catch fmt; + return error.Refused; +} + +/// Flags no family on this engine serves, refused before any device work; null when `o` may open. +pub fn refused(o: api.Open) ?[]const u8 { + if (o.tp != 1) return "the Vulkan engine has no tensor-parallel ranks: serve without --tp"; + if (o.drafter != null) return "--drafter loads a DFlash2 drafter; the Vulkan engine drafts with no head: serve without --drafter"; + if (o.speed_up != null) return "--speed-up links two Macs: serve without --speed-up on Vulkan"; + if (o.kv8) return "--kv-cache fp8: the Vulkan engine keeps bf16 caches; serve without it"; + if (o.prompt_int8) return "--prompt-precision int8: the Vulkan engine's prompt matmuls take bf16 inputs; serve without it"; + if (o.learn != null) return "--learn keeps learned prompt states for the Mac engines; the Vulkan engine keeps none: serve without --learn"; + if (o.slide) return "--slide learns into the served weights on Metal; the Vulkan engine has no learner: serve without --slide"; + if (o.vision != null) return "--vision: the Vulkan engine reads no images; serve it without --vision"; + if (o.loops != null) return "--loops sets a looped model's recurrent passes (Dory's): the Vulkan engine has none; serve without it"; + if (o.mtp_drafts != null) return "--mtp-drafts: the Vulkan engine drafts with no head; serve without --mtp-drafts"; + if (o.mtp_confidence != null) return "--mtp-confidence tunes a draft head, and the Vulkan engine has none: serve without --mtp-confidence"; + if ((o.prompt_cache_gib orelse 0) > 0) return "--prompt-cache-gib keeps prompt states between requests; the Vulkan engine resets its state for every prompt: serve without it"; + if (o.prompt_cache_over_cap) return "--prompt-cache-over-cap lifts the Mac prompt-cache cap; the Vulkan engine keeps no prompt cache: serve without it"; + if (o.keep) |n| if (n > 0) return "--checkpoint-slots keeps prompt states; the Vulkan engine resets its state for every prompt: serve without --checkpoint-slots"; + if (o.device != null) return "--device picks among several GPUs; the Vulkan engine opens its one eligible device: serve without --device"; + if (o.cfg_enabled or o.cfg_batched) return "--cfg-* guidance belongs to the Drama 3 speech model; the Vulkan engine reads no speech model"; + return null; +} + +/// The model's window (config.json's max_position_embeddings), 0 when it names none. +fn modelContext(a: Allocator, io: std.Io, dir: []const u8) i64 { + const path = std.fs.path.join(a, &.{ dir, "config.json" }) catch return 0; + const bytes = std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(16 << 20)) catch return 0; + const doc = std.json.parseFromSliceLeaky(std.json.Value, a, bytes, .{}) catch return 0; + if (doc != .object) return 0; + const text = if (doc.object.get("text_config")) |t| (if (t == .object) t else doc) else doc; + const limit = text.object.get("max_position_embeddings") orelse return 0; + return if (limit == .integer and limit.integer > 0) limit.integer else 0; +} + +// Explicit `--context` must fit model window and SDPA visible keys; by default, use model window capped at arm keys. +fn contextFor(a: Allocator, o: api.Open, native: i64, problem: *[]const u8) error{Refused}!usize { + if (o.context) |asked| { + if (asked < 1) return refuse(a, problem, "--context {d}: give a positive token count", .{asked}); + if (native > 0 and asked > native) return refuse(a, problem, "--context {d} exceeds this model's {d}-token window", .{ asked, native }); + if (asked > fam.backend.max_capacity) return refuse(a, problem, "--context {d}: the Vulkan attention arm sees at most {d} visible keys; serve with --context {d} or less", .{ asked, fam.backend.max_capacity, fam.backend.max_capacity }); + return @intCast(asked); + } + if (native < 1) return refuse(a, problem, "the checkpoint's config.json names no context window; give --context as a positive token count", .{}); + return @intCast(@min(native, fam.backend.max_capacity)); +} + +/// The compiled shader images the engine's ops run, one decl per field, `copy_u16` shared by ops and GDN. +fn bundle() fam.engine.ShaderBundle { + return .{ + .ops = .{ + .qmm_bf16 = &sh.qmm_bf16, + .qmm_q4_vec4 = &sh.qmm_q4_vec4, + .qmm_f32 = &sh.qmm_f32, + .qmm_q4_tile = &sh.qmm_q4_tile, + .qmm_q4_tile_f32 = &sh.qmm_q4_tile_f32, + .dequant_bf16 = &sh.dequant_bf16, + .gather_u32 = &sh.gather_u32, + .gather_u16 = &sh.gather_u16, + .cast_bf16_f32 = &sh.cast_bf16_f32, + .cast_f32_bf16 = &sh.cast_f32_bf16, + .copy_u32 = &sh.copy_u32, + .copy_u16 = &sh.copy_u16, + .copy_f32_bf16 = &sh.copy_f32_bf16, + }, + .norm = .{ + .norm_bf16 = &sh.norm_bf16, + .add_bf16 = &sh.add_bf16, + .add_preserve_bf16 = &sh.add_preserve_bf16, + .swiglu_bf16 = &sh.swiglu_bf16, + .norm_f32 = &sh.norm_f32, + .norm_f32_in = &sh.norm_f32_in, + .norm_bf16_in = &sh.norm_bf16_in, + .add_f32 = &sh.add_f32, + .add_wide_bf16 = &sh.add_wide_bf16, + .add_bf16_f32 = &sh.add_bf16_f32, + .swiglu_f32 = &sh.swiglu_f32, + }, + .attention = .{ + .rope_bf16 = &sh.rope_bf16, + .sdpa_rows_bf16 = &sh.sdpa_rows_bf16, + .gate_bf16 = &sh.gate_bf16, + .rope_f32 = &sh.rope_f32, + .rope_f32_kcache = &sh.rope_f32_kcache, + .sdpa_rows_f32q = &sh.sdpa_rows_f32q, + .gate_f32 = &sh.gate_f32, + }, + .gdn = .{ + .conv_dw1d_bf16 = &sh.conv_dw1d_bf16, + .gdn_conv_decode_bf16 = &sh.gdn_conv_decode_bf16, + .gdn_split_bf16 = &sh.gdn_split_bf16, + .prefill_recur_bf16 = &sh.prefill_recur_bf16, + .gated_delta_decode_bf16 = &sh.gated_delta_decode_bf16, + .fast_norm_gated_bf16 = &sh.fast_norm_gated_bf16, + .copy_u16 = &sh.copy_u16, + .conv_dw1d_f32 = &sh.conv_dw1d_f32, + .gdn_split_f32 = &sh.gdn_split_f32, + .prefill_recur_f32 = &sh.prefill_recur_f32, + .fast_norm_gated_f32 = &sh.fast_norm_gated_f32, + .copy_u32 = &sh.gdn_copy_u32, + }, + .greedy = .{ + .finite_logits = &sh.finite_logits, + .argreduce_chunk = &sh.argreduce_chunk, + .argreduce_combine = &sh.argreduce_combine, + }, + }; +} + +/// One loaded model behind the lane host: everything the engine thread reads lives here. +const Host = struct { + gpa: Allocator, + engine: *fam.engine.Engine, + vk: *fam.backend.Vulkan, + config: lanes.Config, + clock: lanes.backend.WallClock, + core: lanes.Engine, + host: api.LaneHost, + + fn close(ctx: *anyopaque) void { + const h: *Host = @ptrCast(@alignCast(ctx)); + h.host.stop(); + h.core.deinit(); + h.config.deinit(h.gpa); + h.vk.deinit(); + h.engine.deinit(); + h.gpa.destroy(h); + } +}; + +/// The engine for `o.dir`, or null with `problem` set when the checkpoint or the flags are not Vulkan's. +pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !?api.Opened { + if (!std.mem.eql(u8, o.model_type, "qwen3_5")) { + problem.* = try std.fmt.allocPrint(a, "the native Vulkan engine has no backend for {s} checkpoints yet; serve with --engine python", .{o.model_type}); + return null; + } + if (refused(o)) |why| { + problem.* = why; + return null; + } + if (o.lanes > 1 and o.lanes_fixed) { + problem.* = try std.fmt.allocPrint(a, "--parallel {d}: the Vulkan engine serves one stream; serve with --parallel 1", .{o.lanes}); + return null; + } + const window = contextFor(a, o, modelContext(a, io, o.dir), problem) catch |err| switch (err) { + error.Refused => return null, + else => |x| return x, + }; + const engine = fam.engine.Engine.init(gpa, io, bundle(), o.dir, window, timeout_ns) catch |err| switch (err) { + error.OutOfMemory => return error.OutOfMemory, + error.UnsupportedModel => { + problem.* = "the checkpoint is not the pinned Qwen3.5-2B mlx-q4g64 layout the Vulkan engine reads"; + return null; + }, + else => { + problem.* = try std.fmt.allocPrint(a, "the Vulkan engine cannot open a device or its weights ({s})", .{@errorName(err)}); + return null; + }, + }; + errdefer engine.deinit(); + const vk = try fam.backend.Vulkan.init(gpa, engine); + errdefer vk.deinit(); + const h = try gpa.create(Host); + errdefer gpa.destroy(h); + h.* = .{ + .gpa = gpa, + .engine = engine, + .vk = vk, + .config = try lanes.Config.init(gpa, vk.facts(), 1, 0), + .clock = undefined, + .core = undefined, + .host = undefined, + }; + errdefer h.config.deinit(gpa); + h.clock = .{ .io = io }; + h.core = lanes.Engine.init(gpa, &h.config, h.vk.backend(), h.clock.clock()); + errdefer h.core.deinit(); + h.host = api.LaneHost.init(gpa, io, &h.core, .{ + .lanes = 1, + .context_window = @intCast(engine.capacity()), + .prefill_step = fam.backend.max_prompt_rows, + .greedy_only = true, + }); + try h.host.start(); + return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h }; +} + +test "the registry names qwen3_5 over mlx-q4g64 on the vulkan backend" { + try std.testing.expectEqual(@as(usize, 1), backends.len); + try std.testing.expectEqualStrings("vulkan", backends[0]); + try std.testing.expectEqualStrings("qwen3_5", families[0].model_type); + try std.testing.expectEqual(@as(usize, 1), families[0].formats.len); + try std.testing.expectEqualStrings("mlx-q4g64", families[0].formats[0]); + for (serves) |s| try std.testing.expect(!std.mem.eql(u8, s, "--backend=rocm")); + var has_vulkan = false; + for (serves) |s| has_vulkan = has_vulkan or std.mem.eql(u8, s, "--backend=vulkan"); + try std.testing.expect(has_vulkan); +} + +test "flags nothing on this engine serves are refused, not ignored" { + const base = api.Open{ .dir = "d", .model_type = "qwen3_5" }; + try std.testing.expect(refused(base) == null); + try std.testing.expectStringStartsWith(refused(.{ .dir = "d", .model_type = "qwen3_5", .tp = 2 }) orelse "", "the Vulkan engine has no tensor-parallel ranks"); + try std.testing.expectStringStartsWith(refused(.{ .dir = "d", .model_type = "qwen3_5", .drafter = "x" }) orelse "", "--drafter"); + try std.testing.expectStringStartsWith(refused(.{ .dir = "d", .model_type = "qwen3_5", .speed_up = "r.json" }) orelse "", "--speed-up"); + try std.testing.expectStringStartsWith(refused(.{ .dir = "d", .model_type = "qwen3_5", .kv8 = true }) orelse "", "--kv-cache fp8"); + try std.testing.expectStringStartsWith(refused(.{ .dir = "d", .model_type = "qwen3_5", .prompt_int8 = true }) orelse "", "--prompt-precision int8"); + try std.testing.expectStringStartsWith(refused(.{ .dir = "d", .model_type = "qwen3_5", .learn = "/x" }) orelse "", "--learn"); + try std.testing.expectStringStartsWith(refused(.{ .dir = "d", .model_type = "qwen3_5", .slide = true }) orelse "", "--slide"); + try std.testing.expectStringStartsWith(refused(.{ .dir = "d", .model_type = "qwen3_5", .vision = .{} }) orelse "", "--vision"); + try std.testing.expectStringStartsWith(refused(.{ .dir = "d", .model_type = "qwen3_5", .loops = 2 }) orelse "", "--loops"); + try std.testing.expectStringStartsWith(refused(.{ .dir = "d", .model_type = "qwen3_5", .mtp_drafts = 2 }) orelse "", "--mtp-drafts"); + try std.testing.expectStringStartsWith(refused(.{ .dir = "d", .model_type = "qwen3_5", .mtp_confidence = 0.5 }) orelse "", "--mtp-confidence"); + try std.testing.expectStringStartsWith(refused(.{ .dir = "d", .model_type = "qwen3_5", .prompt_cache_gib = 4 }) orelse "", "--prompt-cache-gib"); + try std.testing.expectStringStartsWith(refused(.{ .dir = "d", .model_type = "qwen3_5", .prompt_cache_over_cap = true }) orelse "", "--prompt-cache-over-cap"); + try std.testing.expectStringStartsWith(refused(.{ .dir = "d", .model_type = "qwen3_5", .keep = 2 }) orelse "", "--checkpoint-slots"); + try std.testing.expectStringStartsWith(refused(.{ .dir = "d", .model_type = "qwen3_5", .device = 1 }) orelse "", "--device"); + try std.testing.expectStringStartsWith(refused(.{ .dir = "d", .model_type = "qwen3_5", .cfg_enabled = true }) orelse "", "--cfg-*"); + try std.testing.expect(refused(.{ .dir = "d", .model_type = "qwen3_5", .prompt_cache_gib = 0 }) == null); +} + +test "the context window clamps to the model's own and the arm's visible keys" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var problem: []const u8 = ""; + try std.testing.expectEqual(@as(usize, 256), try contextFor(a, .{ .dir = "d", .model_type = "qwen3_5", .context = 256 }, 4096, &problem)); + try std.testing.expectEqual(@as(usize, 4096), try contextFor(a, .{ .dir = "d", .model_type = "qwen3_5" }, 4096, &problem)); + try std.testing.expectEqual(@as(usize, fam.backend.max_capacity), try contextFor(a, .{ .dir = "d", .model_type = "qwen3_5", .context = fam.backend.max_capacity }, 0, &problem)); + try std.testing.expectError(error.Refused, contextFor(a, .{ .dir = "d", .model_type = "qwen3_5", .context = 8000 }, 4096, &problem)); + try std.testing.expectError(error.Refused, contextFor(a, .{ .dir = "d", .model_type = "qwen3_5" }, 0, &problem)); + try std.testing.expectError(error.Refused, contextFor(a, .{ .dir = "d", .model_type = "qwen3_5", .context = fam.backend.max_capacity + 1 }, 0, &problem)); +} + +test "the pinned checkpoint's wide native window clamps by default, and an explicit over-cap ask is refused" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var problem: []const u8 = ""; + // config.json names 262144 (the pinned checkpoint's): a default open serves the arm's 7168. + try std.testing.expectEqual(@as(usize, fam.backend.max_capacity), try contextFor(a, .{ .dir = "d", .model_type = "qwen3_5" }, 262144, &problem)); + try std.testing.expectError(error.Refused, contextFor(a, .{ .dir = "d", .model_type = "qwen3_5", .context = 262144 }, 262144, &problem)); +} + +test "the core's production bounds helpers gate prompt and capacity" { + try std.testing.expect(fam.backend.validPrompt(&.{1})); + try std.testing.expect(!fam.backend.validPrompt(&.{})); + try fam.engine.planFor(11, 0, 4, 8); // four prompt rows and seven decodes + try std.testing.expectError(error.ContextTooLong, fam.engine.planFor(11, 0, 4, 9)); // one past +} diff --git a/zig/src/server/cli.zig b/zig/src/server/cli.zig index 82975c5f7e5..9679fab435b 100644 --- a/zig/src/server/cli.zig +++ b/zig/src/server/cli.zig @@ -19,7 +19,7 @@ pub const Flag = struct { /// The values it honours when only some of the choices. native_values: ?[]const []const u8 = null, }; -const backend_values: []const []const u8 = if (builtin.os.tag == .macos) &.{ "auto", "mlx" } else if (served("--backend=rocm")) &.{ "auto", "rocm" } else &.{ "auto", "cuda" }; +const backend_values: []const []const u8 = if (builtin.os.tag == .macos) &.{ "auto", "mlx" } else if (served("--backend=vulkan")) &.{ "auto", "vulkan" } else if (served("--backend=rocm")) &.{ "auto", "rocm" } else &.{ "auto", "cuda" }; /// The flags past the common ones the binary's backend serves, from its root's `serves` (test roots: none). const native_serves: []const []const u8 = if (@hasDecl(@import("root"), "serves")) @import("root").serves else &.{}; @@ -129,7 +129,7 @@ pub const flags = [_]Flag{ .{ .name = "--ple-on-ssd", .kind = .store_true, .native = builtin.os.tag == .macos }, .{ .name = "--fsr-route-trace", .native = builtin.os.tag == .macos }, .{ .name = "--no-update-check", .kind = .store_true, .native = true }, - .{ .name = "--backend", .choices = &.{ "auto", "mlx", "cuda", "rocm" }, .native = true, .native_values = backend_values }, + .{ .name = "--backend", .choices = &.{ "auto", "mlx", "cuda", "rocm", "vulkan" }, .native = true, .native_values = backend_values }, .{ .name = "--tp", .choices = &.{ "1", "2", "4", "8" }, .native = served("--tp") }, .{ .name = "--rank", .native = served("--tp") }, .{ .name = "--master", .native = served("--tp") }, @@ -585,6 +585,7 @@ const cli = @This(); /// The private flag helpers cli_test.zig checks. pub const Test = if (builtin.is_test) struct { pub const cuda_build = cli.cuda_build; + pub const backend_values = cli.backend_values; pub const native_serves = cli.native_serves; pub const servedBy = cli.servedBy; pub const cudaFlag = cli.cudaFlag; diff --git a/zig/src/server/cli_backend_test.zig b/zig/src/server/cli_backend_test.zig new file mode 100644 index 00000000000..0f5cc538cff --- /dev/null +++ b/zig/src/server/cli_backend_test.zig @@ -0,0 +1,21 @@ +//! The --backend flag's choice list versus the build's served gate, apart from cli_test.zig's 600-line cap. +const std = @import("std"); +const cli = @import("cli.zig"); +const Usage = cli.Usage; +const parse = cli.parse; +const servedBy = cli.Test.servedBy; + +test "flag deps: --backend vulkan parses where served and gets the gated refusal elsewhere" { + const backend_values = cli.Test.backend_values; + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var u: Usage = .{}; + if (servedBy(backend_values, "vulkan")) { + const parsed = try parse(a, &.{ "m", "--backend", "vulkan" }, &u); + try std.testing.expectEqualStrings("vulkan", parsed.backend); + } else { + try std.testing.expectError(error.Usage, parse(a, &.{ "m", "--backend", "vulkan" }, &u)); + try std.testing.expect(std.mem.indexOf(u8, u.message, "not served by the native engine") != null); + } +} diff --git a/zig/src/server/root.zig b/zig/src/server/root.zig index 1afe9ea063f..d57954d825a 100644 --- a/zig/src/server/root.zig +++ b/zig/src/server/root.zig @@ -77,6 +77,7 @@ test { _ = @import("cli_test.zig"); _ = @import("cli_profile_test.zig"); _ = @import("compact_at_cli_test.zig"); + _ = @import("cli_backend_test.zig"); _ = @import("unix_socket_cli_test.zig"); _ = @import("effort_cli_audit_test.zig"); _ = @import("live.zig"); diff --git a/zig/src/vulkan/arena.zig b/zig/src/vulkan/arena.zig new file mode 100644 index 00000000000..d1ed9c97a5f --- /dev/null +++ b/zig/src/vulkan/arena.zig @@ -0,0 +1,109 @@ +// Arena suballocates aligned device-local views; do not move it after views exist, as they point to its buffer. +const std = @import("std"); +const memory = @import("memory.zig"); +const tensor = @import("tensor.zig"); +const Device = @import("device.zig").Device; +const Buffer = memory.Buffer; +const Error = @import("runtime.zig").Error; + +pub const Arena = struct { + buffer: Buffer, + /// Next free byte; only grows through alloc and shrinks through reset. + next: u64, + /// Every allocation's alignment: the device's storage offset alignment, at least a word. + alignment: u64, + + /// A device-local arena of `bytes`; its size must fit the device's buffer and allocation limits. + pub fn init(dev: *const Device, bytes: u64) Error!Arena { + try memory.fits(bytes, dev.v13.maxBufferSize, dev.v11.maxMemoryAllocationSize); + return .{ + .buffer = try Buffer.init(dev, .device, bytes), + .next = 0, + .alignment = @max(dev.properties.limits.minStorageBufferOffsetAlignment, 4), + }; + } + + /// Frees the arena's buffer; call only after the last submission using any of its views returned success. + pub fn deinit(a: *Arena) void { + a.buffer.deinit(); + } + + // alloc returns an aligned view or named error; cursor advances only after full fit, so failure wastes nothing. + pub fn alloc(a: *Arena, kind: tensor.Kind, elements: u64) Error!tensor.View { + if (elements == 0) return error.Invalid; + const start = std.mem.alignForward(u64, a.next, a.alignment); + const v = tensor.View.init(&a.buffer, kind, start, elements) catch |e| return if (e == error.OutOfRange) error.OutOfMemory else e; + const end = std.math.add(u64, v.byte_offset, v.byteLen()) catch return error.OutOfMemory; + // Round the committed cursor up so every mark/reset boundary is an allocation boundary. + a.next = (std.math.add(u64, end, a.alignment - 1) catch return error.OutOfMemory) / a.alignment * a.alignment; + return v; + } + + /// The cursor to reset to; taken before the allocations it unwinds. + pub fn mark(a: *const Arena) u64 { + return a.next; + } + + /// Winds the cursor back to an earlier mark; call only after the serial commands using those views completed. + pub fn reset(a: *Arena, to: u64) Error!void { + if (to > a.next) return error.OutOfRange; + if (to % a.alignment != 0) return error.Misaligned; + a.next = to; + } + + pub fn used(a: *const Arena) u64 { + return a.next; + } + + pub fn capacity(a: *const Arena) u64 { + return a.buffer.len; + } +}; + +const testing = std.testing; +const tensorKind = tensor.Kind; + +fn arena(len: u64, alignment: u64) Arena { + return .{ .buffer = .{ .dev = undefined, .handle = null, .memory = null, .len = len, .kind = .device, .mapped = null, .coherent = false }, .next = 0, .alignment = alignment }; +} + +test "allocations are aligned, bounded and hold their geometry" { + var a = arena(1024, 256); + const f = try a.alloc(.f32, 4); // 16 bytes, cursor committed at the next 256 boundary + try testing.expectEqual(@as(u64, 0), f.byte_offset); + const b = try a.alloc(tensorKind.bf16, 130); // 260 bytes from 256, cursor rounds to 768 + try testing.expectEqual(@as(u64, 256), b.byte_offset); + try testing.expectEqual(@as(u64, 768), a.used()); + try testing.expectError(error.OutOfMemory, a.alloc(.u32, 65)); // 260 bytes past 768 cannot round-trip 1024 + try testing.expectEqual(@as(u64, 768), a.used()); // the failed allocation left the cursor alone + try testing.expect(f.buffer == &a.buffer); +} + +test "the arena never returns empty or ragged views" { + var a = arena(512, 64); + try testing.expectError(error.Invalid, a.alloc(.u32, 0)); + try testing.expectError(error.Misaligned, a.alloc(.bf16, 1)); // a two-byte tail drops its element + try testing.expectEqual(@as(u64, 0), a.used()); + const ok = try a.alloc(.bf16, 2); // whole words only + try testing.expectEqual(@as(u64, 4), ok.byteLen()); +} + +test "reset returns to an earlier mark and refuses later or misaligned marks" { + var a = arena(512, 64); + const first = try a.alloc(.u32, 8); // bytes 0..32, cursor committed at 64 + try testing.expectEqual(@as(u64, 0), first.byte_offset); + const m = a.mark(); + try testing.expectEqual(@as(u64, 64), m); + _ = try a.alloc(.u32, 8); // bytes 64..96 + try a.reset(m); + try testing.expectEqual(@as(u64, m), a.used()); + const again = try a.alloc(.u32, 8); // lands where the reset-away allocation was + try testing.expectEqual(@as(u64, 64), again.byte_offset); + try testing.expectError(error.OutOfRange, a.reset(a.used() + 64)); + try testing.expectError(error.Misaligned, a.reset(17)); +} + +test "a zero-capacity arena allocates nothing" { + var a = arena(0, 16); + try testing.expectError(error.OutOfMemory, a.alloc(.f32, 1)); +} diff --git a/zig/src/vulkan/attention.zig b/zig/src/vulkan/attention.zig new file mode 100644 index 00000000000..709db3f2175 --- /dev/null +++ b/zig/src/vulkan/attention.zig @@ -0,0 +1,355 @@ +// Qwen3.5-2B uses half-split RoPE: head_dim 256, rotary_dim 64, theta 1e7; 8 query heads, 2 KV heads; scale 1/16. + +// fast_rope: -DUSE_BF16 -DORACLE_QWEN35 (no mid bf16 rounding); sdpa: -DBF16_IO -DSDPA_DIM=256 -DSDPA_ROWS; gate: none + +// zig/src/families/qwen3_5/forward.zig attention(); zig/kernels/metal/qwen3_5/{layout,attn_partial,attn_merge}.metal. +const std = @import("std"); +const dispatch = @import("dispatch.zig"); +const memory = @import("memory.zig"); +const ops_mod = @import("ops.zig"); +const params = @import("stage1_params"); +const tensor = @import("tensor.zig"); +const commands = @import("commands.zig"); +const Device = @import("device.zig").Device; +const Commands = commands.Commands; +const Buffer = memory.Buffer; +// dispatch.Error adds ContextTooLong; long prefixes refuse, geometry never falls through, and there is no CPU fallback. +pub const Error = dispatch.Error || error{ContextTooLong}; +const View = tensor.View; +const Ops = ops_mod.Ops; + +pub const query_heads: u64 = 8; +pub const kv_heads: u64 = 2; +pub const head_dim: u64 = 256; +// Metal rotates i and i+32 for i < 32 (rotary_dim 64) and passes through d >= 64. +pub const rotary_pairs: u64 = 32; +pub const theta: f32 = 1e7; +/// head_dim^-1/2 as the oracle passes it: e.setValue(1.0 / 16.0, 3). +pub const scale: f32 = 1.0 / 16.0; +// bf16 stages kStreamKeys f32 slots; long prefixes refuse, not truncate; p1/p2 is not wired. +pub const max_visible_keys: u64 = 7168; +/// s_qrow (256) + s_stream (7168) + s_tree (256) f32 words. +const sdpa_shared_bytes: u64 = 30720; + +// The compiled shader images Attention runs; every field is a real SPIR-V image of the named variant. +pub const Shaders = struct { + rope_bf16: []align(4) const u8, + sdpa_rows_bf16: []align(4) const u8, + gate_bf16: []align(4) const u8, + rope_f32: []align(4) const u8, + rope_f32_kcache: []align(4) const u8, + sdpa_rows_f32q: []align(4) const u8, + gate_f32: []align(4) const u8, +}; + +/// attention_gate.comp's own one-word push block. +const GateParams = extern struct { words: u32 = 0 }; +const gate_push_bytes: u32 = @sizeOf(GateParams); + +const gate_row_elems: u64 = query_heads * head_dim; // 2048 +const gate_row_words: u64 = gate_row_elems / 2; + +comptime { + if (gate_push_bytes != 4) @compileError("attention_gate push block must stay one word"); +} + +pub const Attention = struct { + dev: *const Device, + /// Mutable: cacheKV rebinds the shared Ops copy_u16 pipeline (bind needs *Pipeline). + ops: *Ops, + rope_bf16: dispatch.Pipeline, + sdpa_bf16: dispatch.Pipeline, + gate_bf16: dispatch.Pipeline, + rope_f32: dispatch.Pipeline, + rope_f32_kcache: dispatch.Pipeline, + sdpa_f32q: dispatch.Pipeline, + gate_f32: dispatch.Pipeline, + // The i32 base position is read by each rope dispatch; setBasePosition writes it between operations. + offsets: Buffer, + + // Builds three pipelines; checks bf16 features, shared memory, and Pipeline.init’s 1024-thread workgroup limit. + pub fn init(dev: *const Device, shaders: Shaders, ops: *Ops) Error!Attention { + if (!dev.features.storage16 or !dev.features.int16) return error.MissingFeature; + if (dev.properties.limits.maxComputeSharedMemorySize < sdpa_shared_bytes) return error.MissingFeature; + var a: Attention = .{ + .dev = dev, + .ops = ops, + .rope_bf16 = undefined, + .sdpa_bf16 = undefined, + .gate_bf16 = undefined, + .rope_f32 = undefined, + .rope_f32_kcache = undefined, + .sdpa_f32q = undefined, + .gate_f32 = undefined, + .offsets = try Buffer.init(dev, .device, 4), + }; + errdefer a.offsets.deinit(); + // fast_rope reads input, output and the offsets slot; output is binding 1. + a.rope_bf16 = try dispatch.Pipeline.init(dev, shaders.rope_bf16, .{ .bindings = 3, .push_bytes = params.push_bytes, .writes = 1 << 1 }); + errdefer a.rope_bf16.deinit(); + a.sdpa_bf16 = try dispatch.Pipeline.init(dev, shaders.sdpa_rows_bf16, .{ .bindings = 4, .push_bytes = params.push_bytes, .writes = 1 << 3 }); + errdefer a.sdpa_bf16.deinit(); + a.gate_bf16 = try dispatch.Pipeline.init(dev, shaders.gate_bf16, .{ .bindings = 3, .push_bytes = gate_push_bytes, .writes = 1 << 2 }); + errdefer a.gate_bf16.deinit(); + a.rope_f32 = try dispatch.Pipeline.init(dev, shaders.rope_f32, .{ .bindings = 3, .push_bytes = params.push_bytes, .writes = 1 << 1 }); + errdefer a.rope_f32.deinit(); + a.rope_f32_kcache = try dispatch.Pipeline.init(dev, shaders.rope_f32_kcache, .{ .bindings = 3, .push_bytes = params.push_bytes, .writes = 1 << 1 }); + errdefer a.rope_f32_kcache.deinit(); + a.sdpa_f32q = try dispatch.Pipeline.init(dev, shaders.sdpa_rows_f32q, .{ .bindings = 4, .push_bytes = params.push_bytes, .writes = 1 << 3 }); + errdefer a.sdpa_f32q.deinit(); + a.gate_f32 = try dispatch.Pipeline.init(dev, shaders.gate_f32, .{ .bindings = 3, .push_bytes = gate_push_bytes, .writes = 1 << 2 }); + return a; + } + + // Deinit after successful completion only; Timeout/InFlight may leave GPU-owned resources, so use commands.orExit. + pub fn deinit(a: *Attention) void { + a.rope_bf16.deinit(); + a.sdpa_bf16.deinit(); + a.gate_bf16.deinit(); + a.rope_f32.deinit(); + a.rope_f32_kcache.deinit(); + a.sdpa_f32q.deinit(); + a.gate_f32.deinit(); + a.offsets.deinit(); + } + + // Uploads the segment base (cache length); rope reads it, and queue order requires upload to finish first. + pub fn setBasePosition(a: *Attention, cmds: *Commands, staging: Buffer, position: i64, timeout_ns: u64) Error!void { + if (position < 0 or position > std.math.maxInt(i32)) return error.Invalid; + const bytes = std.mem.toBytes(@as(i32, @intCast(position))); + try commands.upload(cmds, a.offsets, 0, staging, &bytes, timeout_ns); + } + + // Rotates normalized q (rows,8,256); bf16 stays word-packed, the wide f32 path addresses elements. + pub fn ropeQueries(a: *Attention, cmds: *Commands, out: View, q: View, rows: u64, timeout_ns: u64) Error!void { + const plan = try ropePlan(out.elements, q.elements, rows, query_heads); + if (out.overlaps(q)) return error.Aliased; + const f32io = out.kind == .f32; + if (q.kind != (if (f32io) tensor.Kind.f32 else tensor.Kind.bf16) or out.kind != q.kind) return error.Invalid; + const pipe = if (f32io) &a.rope_f32 else &a.rope_bf16; + try pipe.bind(&.{ q.bind(), out.bind(), a.offsetBinding() }); + try a.ropeRun(pipe, cmds, plan, queryStrides(), 0, if (f32io) plan.elements else plan.words, timeout_ns); + } + + // K rotates into cache (rounded to bf16 on landing); V rides the f32->bf16 rounding copy in wide layers. + pub fn cacheKV(a: *Attention, cmds: *Commands, cache_keys: View, cache_values: View, knorm: View, v: View, rows: u64, base: u64, capacity: u64, timeout_ns: u64) Error!void { + if (cache_keys.kind != .bf16 or cache_values.kind != .bf16) return error.Invalid; + if (knorm.kind != .bf16 and knorm.kind != .f32) return error.Invalid; + if (v.kind != knorm.kind) return error.Invalid; + const cache_elems = kv_heads *| capacity *| head_dim; + if (cache_keys.elements != cache_elems or cache_values.elements != cache_elems) return error.Invalid; + const end = std.math.add(u64, base, rows) catch return error.TooLarge; + if (end > capacity) return error.OutOfRange; + if (cache_keys.overlaps(knorm) or cache_values.overlaps(v)) return error.Aliased; + const plan = try ropePlan(knorm.elements, knorm.elements, rows, kv_heads); + const wide = knorm.kind == .f32; + // K: passthrough tail copies d >= 64; rotation covers d < 64; f32 in rounds to bf16 on the cache store. + const rope = if (wide) &a.rope_f32_kcache else &a.rope_bf16; + try rope.bind(&.{ knorm.bind(), cache_keys.bind(), a.offsetBinding() }); + try a.ropeRun(rope, cmds, plan, .{ + .in_head = head_dim, + .in_time = kv_heads * head_dim, + .out_head = capacity * head_dim, // cache head stride + .out_time = head_dim, // cache time stride + }, base * head_dim, if (wide) plan.elements else plan.words, timeout_ns); + // V: copy_general, dims 3; bf16 is a bit-exact uint16 transport, f32 rounds to bf16 on landing. + const copy = try vCopyPlan(rows, capacity, v.elements); + const vpipe = if (wide) &a.ops.copy_f32_bf16 else &a.ops.copy_u16; + try vpipe.bind(&.{ v.bind(), a.ops.dummyBinding(), cache_values.bind(), a.ops.dummyBinding() }); + var p: params.Params = .{ + .count = copy.elements, + .dims = 3, + .shape = .{ @intCast(head_dim), @intCast(kv_heads), try ops_mod.cast32(rows), 0 }, + .in_strides = .{ 1, @intCast(head_dim), @intCast(kv_heads * head_dim), 0 }, + .out_strides = .{ 1, try ops_mod.cast32(capacity * head_dim), @intCast(head_dim), 0 }, + }; + p.output_offset = try ops_mod.cast32(base * head_dim); + try vpipe.run(cmds, std.mem.asBytes(&p), try groups(copy.elements, vpipe.local[0], a.dev.properties.limits.maxComputeWorkGroupCount[0]), timeout_ns); + } + + // SDPA scales q once, streams f32 scores; Metal scales per-score in 128-key chunks. No bit-exact/performance claim. + pub fn attend(a: *Attention, cmds: *Commands, y: View, queries: View, cache_keys: View, cache_values: View, prefix: u64, rows: u64, capacity: u64, timeout_ns: u64) Error!void { + const plan = try attendPlan(prefix, rows, capacity, queries.elements, y.elements, cache_keys.elements, cache_values.elements); + if (y.kind != queries.kind or (queries.kind != .bf16 and queries.kind != .f32)) return error.Invalid; + if (y.overlaps(queries) or y.overlaps(cache_keys) or y.overlaps(cache_values)) return error.Aliased; + const pipe = if (queries.kind == .f32) &a.sdpa_f32q else &a.sdpa_bf16; + try pipe.bind(&.{ queries.bind(), cache_keys.bind(), cache_values.bind(), y.bind() }); + var p: params.Params = .{ + .matrix_m = @intCast(query_heads), // repeats = 8/2 = 4 + .matrix_n = @intCast(kv_heads), + .matrix_k = plan.total_keys, + .flags = 0, + .alpha = scale, + .dims = plan.rows, // q_len + .shape = .{ @intCast(head_dim), try ops_mod.cast32(capacity * head_dim), @intCast(head_dim), 1 }, + .in_strides = .{ 1, try ops_mod.cast32(capacity * head_dim), @intCast(head_dim), 1 }, + .out_strides = .{ @intCast(query_heads * head_dim), 0, 0, 0 }, + }; + const groups_x = std.math.cast(u32, plan.rows) orelse return error.TooLarge; + try pipe.run(cmds, std.mem.asBytes(&p), .{ @intCast(query_heads), groups_x, 1 }, timeout_ns); + } + + // Gate is bf16(ctx*bf16(sigmoid(gate))) narrow or f32(ctx*sigmoid(gate)) wide; qraw's second halves hold the gate. + pub fn gateOutput(a: *Attention, cmds: *Commands, out: View, ctx: View, qraw: View, rows: u64, timeout_ns: u64) Error!void { + const plan = try gatePlan(out.elements, ctx.elements, qraw.elements, rows); + const wide = out.kind == .f32; + if (ctx.kind != out.kind or qraw.kind != out.kind) return error.Invalid; + if (out.overlaps(ctx) or out.overlaps(qraw)) return error.Aliased; + const pipe = if (wide) &a.gate_f32 else &a.gate_bf16; + try pipe.bind(&.{ ctx.bind(), qraw.bind(), out.bind() }); + const words = if (wide) plan.elements else plan.words; + const p: GateParams = .{ .words = words }; + try pipe.run(cmds, std.mem.asBytes(&p), try groups(words, pipe.local[0], a.dev.properties.limits.maxComputeWorkGroupCount[0]), timeout_ns); + } + + fn offsetBinding(a: *const Attention) dispatch.Binding { + return .{ .buffer = &a.offsets, .len = 4 }; + } + + // out_base offsets cache; trig uses omarchy_trig.h wrappers, not metal::cos/sin; exp/pow/cos may differ in ulps. + fn ropeRun(a: *Attention, pipe: *dispatch.Pipeline, cmds: *Commands, plan: RopePlan, strides: RopeStrides, out_base: u64, items: u32, timeout_ns: u64) Error!void { + // Passthrough inv_freq is exp(-i*beta), with beta=ln(1e7)/32, yielding 1e7^(-i/32) up to exp/pow ulps. + var p: params.Params = .{ + .count = items, + .flags = rope_flags, + .alpha = 1.0, // position = 1 * (t + base) + .beta = @log(theta) / @as(f32, rotary_pairs), + .dims = @intCast(rotary_pairs), // half_dims + .shape = .{ plan.heads, plan.rows, @intCast(head_dim), 0 }, // heads, time, rotated-axis extent + .in_strides = .{ @intCast(strides.in_head), try ops_mod.cast32(strides.in_time), 1, 0 }, + .out_strides = .{ try ops_mod.cast32(strides.out_head), try ops_mod.cast32(strides.out_time), 1, 0 }, + }; + p.output_offset = try ops_mod.cast32(out_base); + try pipe.run(cmds, std.mem.asBytes(&p), try groups(items, pipe.local[0], a.dev.properties.limits.maxComputeWorkGroupCount[0]), timeout_ns); + } +}; + +/// ropeRun's push flags: bit0 forward, bit3 passthrough tail (fast_rope.comp). +const rope_flags: u32 = 1 | 8; + +const RopeStrides = struct { in_head: u64, in_time: u64, out_head: u64, out_time: u64 }; + +/// (rows, heads, 256) bf16, contiguous on both sides. +fn queryStrides() RopeStrides { + return .{ + .in_head = head_dim, + .in_time = query_heads * head_dim, + .out_head = head_dim, + .out_time = query_heads * head_dim, + }; +} + +fn groups(count: u64, local: u32, max_x: u32) Error![3]u32 { + const n = std.math.cast(u32, (count + local - 1) / local) orelse return error.TooLarge; + return .{ @min(n, max_x), 1, 1 }; +} + +// Rope geometry: bf16 addresses packed words (two halves), f32 addresses elements; pure and host-tested. +pub const RopePlan = struct { rows: u32, heads: u32, words: u32, elements: u32 }; + +pub fn ropePlan(out_elements: u64, in_elements: u64, rows: u64, heads: u64) Error!RopePlan { + if (rows == 0 or rows > std.math.maxInt(u32)) return error.Invalid; + const row_elems = heads *| head_dim; + const total = rows *| row_elems; + if (in_elements != total or out_elements != total) return error.Invalid; + if (heads != query_heads and heads != kv_heads) return error.Invalid; + return .{ + .rows = @intCast(rows), + .heads = @intCast(heads), + .words = try ops_mod.cast32(total / 2), // one word per bf16 pair + .elements = try ops_mod.cast32(total), + }; +} + +/// V scatter geometry (rows of 2*256 elements into a capacity cache). +pub const VCopyPlan = struct { elements: u32 }; + +pub fn vCopyPlan(rows: u64, capacity: u64, v_elements: u64) Error!VCopyPlan { + if (rows == 0 or capacity == 0 or rows > capacity) return error.Invalid; + const row_elems = kv_heads * head_dim; + const total = rows *| row_elems; + if (v_elements != total) return error.Invalid; + if (capacity *| row_elems > std.math.maxInt(u32)) return error.TooLarge; // output offsets stay u32 + return .{ .elements = try ops_mod.cast32(total) }; +} + +// Causal SDPA row r sees prefix+r+1 keys, including its own; later keys are unread. Pure and host-tested. +pub const AttendPlan = struct { rows: u32, total_keys: u32 }; + +pub fn attendPlan(prefix: u64, rows: u64, capacity: u64, q_elements: u64, y_elements: u64, k_elements: u64, v_elements: u64) Error!AttendPlan { + if (rows == 0 or rows > std.math.maxInt(u32)) return error.Invalid; + const total = std.math.add(u64, prefix, rows) catch return error.TooLarge; + if (capacity == 0 or total > capacity) return error.OutOfRange; + if (total > max_visible_keys) return error.ContextTooLong; + const row_elems = query_heads * head_dim; + if (q_elements != rows *| row_elems or y_elements != rows *| row_elems) return error.Invalid; + const cache_elems = kv_heads *| capacity *| head_dim; + if (k_elements != cache_elems or v_elements != cache_elems) return error.Invalid; + if (capacity *| head_dim > std.math.maxInt(u32)) return error.TooLarge; + return .{ .rows = @intCast(rows), .total_keys = try ops_mod.cast32(total) }; +} + +/// Gate geometry; words covers the bf16 pair count, elements the f32 count; pure and host-tested. +pub const GatePlan = struct { rows: u32, words: u32, elements: u32 }; + +pub fn gatePlan(out_elements: u64, ctx_elements: u64, qraw_elements: u64, rows: u64) Error!GatePlan { + if (rows == 0 or rows > std.math.maxInt(u32)) return error.Invalid; + if (ctx_elements != rows *| gate_row_elems or out_elements != rows *| gate_row_elems) return error.Invalid; + if (qraw_elements != rows *| (2 * gate_row_elems)) return error.Invalid; + return .{ + .rows = @intCast(rows), + .words = try ops_mod.cast32(rows * gate_row_words), + .elements = try ops_mod.cast32(rows * gate_row_elems), + }; +} + +const testing = std.testing; + +test "ropePlan derives the word count and refuses torn or foreign shapes" { + const plan = try ropePlan(4 * 8 * 256, 4 * 8 * 256, 4, query_heads); + try testing.expectEqual(@as(u32, 4 * 8 * 128), plan.words); + try testing.expectEqual(@as(u32, 4 * 8 * 256), plan.elements); + const kv = try ropePlan(3 * 2 * 256, 3 * 2 * 256, 3, kv_heads); + try testing.expectEqual(@as(u32, 3 * 2 * 128), kv.words); + try testing.expectEqual(@as(u32, 3 * 2 * 256), kv.elements); + try testing.expectError(error.Invalid, ropePlan(4 * 8 * 256 + 1, 4 * 8 * 256, 4, query_heads)); + try testing.expectError(error.Invalid, ropePlan(4 * 8 * 256, 4 * 8 * 256, 4, 4)); // not the model's head count + try testing.expectError(error.Invalid, ropePlan(0, 0, 0, query_heads)); + try testing.expectError(error.TooLarge, ropePlan(1 << 35, 1 << 35, 1 << 24, query_heads)); // words past u32 +} + +test "attendPlan keeps the causal prefix and refuses long or torn contexts" { + // rows > 1: each row's visible count is prefix + row + 1. + const plan = try attendPlan(100, 3, 128, 3 * 8 * 256, 3 * 8 * 256, 2 * 128 * 256, 2 * 128 * 256); + try testing.expectEqual(@as(u32, 103), plan.total_keys); + try testing.expectEqual(@as(u32, 3), plan.rows); + try testing.expectError(error.OutOfRange, attendPlan(126, 3, 128, 3 * 2048, 3 * 2048, 2 * 128 * 256, 2 * 128 * 256)); + try testing.expectError(error.ContextTooLong, attendPlan(max_visible_keys, 1, max_visible_keys + 1, 2048, 2048, 2 * (max_visible_keys + 1) * 256, 2 * (max_visible_keys + 1) * 256)); + try testing.expectError(error.Invalid, attendPlan(0, 3, 128, 2 * 2048, 3 * 2048, 2 * 128 * 256, 2 * 128 * 256)); + try testing.expectError(error.Invalid, attendPlan(0, 3, 128, 3 * 2048, 3 * 2048, 2 * 127 * 256, 2 * 128 * 256)); + try testing.expectError(error.TooLarge, attendPlan(0, 1, 1 << 25, 2048, 2048, 2 * (1 << 25) * 256, 2 * (1 << 25) * 256)); +} + +test "vCopyPlan refuses rows past capacity and caches past u32 offsets" { + const plan = try vCopyPlan(4, 16, 4 * 2 * 256); + try testing.expectEqual(@as(u32, 4 * 2 * 256), plan.elements); + try testing.expectError(error.Invalid, vCopyPlan(17, 16, 17 * 2 * 256)); + try testing.expectError(error.Invalid, vCopyPlan(4, 16, 5 * 2 * 256)); + try testing.expectError(error.TooLarge, vCopyPlan(1, 1 << 24, 2 * 256)); +} + +test "gatePlan derives the word count and refuses torn rows" { + const plan = try gatePlan(4 * 2048, 4 * 2048, 4 * 4096, 4); + try testing.expectEqual(@as(u32, 4 * 1024), plan.words); + try testing.expectEqual(@as(u32, 4 * 2048), plan.elements); + try testing.expectError(error.Invalid, gatePlan(4 * 2048, 4 * 2048 + 1, 4 * 4096, 4)); + try testing.expectError(error.Invalid, gatePlan(4 * 2048, 4 * 2048, 4 * 4095, 4)); + try testing.expectError(error.Invalid, gatePlan(0, 0, 0, 0)); +} + +test "rope push flags match the masks the embedded rope shader actually tests" { + const src = @import("vulkan_comp_sources").fast_rope; + try testing.expectEqual(@as(u32, 1 | 8), rope_flags); + try testing.expect(std.mem.indexOf(u8, src, "(params.flags & 1u)") != null); + try testing.expect(std.mem.indexOf(u8, src, "(params.flags & 8u)") != null); +} diff --git a/zig/src/vulkan/commands.zig b/zig/src/vulkan/commands.zig index e0f0b9f1048..aeafcac9c1d 100644 --- a/zig/src/vulkan/commands.zig +++ b/zig/src/vulkan/commands.zig @@ -1,5 +1,6 @@ //! One command buffer and fence on the device's queue, and staged uploads and readbacks through it. const std = @import("std"); +const builtin = @import("builtin"); const c = @import("vk"); const runtime = @import("runtime.zig"); const memory = @import("memory.zig"); @@ -13,6 +14,12 @@ pub const Commands = struct { pool: c.VkCommandPool, cmd: c.VkCommandBuffer, fence: c.VkFence, + /// A command buffer is begun and holds the dispatches and copies recorded since the last submission. + recording: bool = false, + /// Dispatches in the open command buffer; Pipeline.run submits when it reaches half a batch. + pending: u32 = 0, + /// Set by a state-corrupting failure (signaled fence or invalid buffer): later begin and submit refuse. + poison: ?Error = null, pub fn init(dev: *const Device) Error!Commands { const api = &dev.api; @@ -38,28 +45,50 @@ pub const Commands = struct { return .{ .dev = dev, .pool = pool, .cmd = cmd, .fence = fence }; } - /// Frees the pool, its command buffer and the fence; call only after the last submit returned success. + /// Frees the pool, its command buffer and the fence after the last submit; dispatches never submitted are dropped. pub fn deinit(self: *Commands) void { self.dev.api.vkDestroyFence(self.dev.handle, self.fence, null); self.dev.api.vkDestroyCommandPool(self.dev.handle, self.pool, null); self.* = undefined; } - /// The command buffer, reset and begun for one submission. + /// Drops unsubmitted records unrun. Call before freeing any buffer an open command buffer may reference. + pub fn discard(self: *Commands) void { + if (!self.recording) return; + _ = self.dev.api.vkResetCommandBuffer(self.cmd, 0); + self.recording = false; + self.pending = 0; + } + + /// The open command buffer; begun (and reset) only when nothing recorded earlier still waits for a submission. pub fn begin(self: *Commands) Error!c.VkCommandBuffer { + if (self.poison) |p| return p; + if (self.recording) return self.cmd; try check(self.dev.api.vkResetCommandBuffer(self.cmd, 0)); const info: c.VkCommandBufferBeginInfo = .{ .sType = c.VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO, .flags = c.VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT, }; try check(self.dev.api.vkBeginCommandBuffer(self.cmd, &info)); + self.recording = true; return self.cmd; } - /// Finite fence wait; Timeout/InFlight require orExit to avoid freeing resources the GPU may still use. + /// Submits what is recorded (nothing if none); timeout_ns covers the whole batch; Timeout and InFlight need orExit. pub fn submit(self: *Commands, timeout_ns: u64) Error!void { + if (self.poison) |p| return p; + if (!self.recording) return; const api = &self.dev.api; - try check(api.vkEndCommandBuffer(self.cmd)); + const ended = api.vkEndCommandBuffer(self.cmd); + if (ended != c.VK_SUCCESS) { + // Invalid after a failed end: poison the stream so no recording begin reuses it. + self.poison = error.CommandBufferInvalid; + self.recording = false; + if (!builtin.is_test) std.log.err("vulkan: vkEndCommandBuffer returned {d}; the command stream is poisoned", .{ended}); + return error.CommandBufferInvalid; + } + self.recording = false; + self.pending = 0; const info: c.VkSubmitInfo = .{ .sType = c.VK_STRUCTURE_TYPE_SUBMIT_INFO, .commandBufferCount = 1, @@ -75,7 +104,13 @@ pub const Commands = struct { std.log.err("vulkan: vkWaitForFences returned {d} within {d} ns of a submission", .{ waited_on, timeout_ns }); return e; }; - try check(api.vkResetFences(self.dev.handle, 1, &self.fence)); + const reset_done = api.vkResetFences(self.dev.handle, 1, &self.fence); + if (reset_done != c.VK_SUCCESS) { + // The fence stays signaled after a good wait: poison the stream; orExit ends the process. + self.poison = error.FencePoisoned; + if (!builtin.is_test) std.log.err("vulkan: vkResetFences returned {d} after a signaled wait; the command stream is poisoned", .{reset_done}); + return error.FencePoisoned; + } } }; @@ -101,6 +136,8 @@ pub fn orExit(result: anytype) @TypeOf(result) { _ = result catch |e| { if (e == error.Timeout) std.process.exit(124); if (e == error.InFlight) std.process.exit(1); + if (e == error.FencePoisoned) std.process.exit(1); + if (e == error.CommandBufferInvalid) std.process.exit(1); return e; }; return result; @@ -116,6 +153,35 @@ pub fn barrier(dev: *const Device, cmd: c.VkCommandBuffer, src_stage: u32, src_a dev.api.vkCmdPipelineBarrier(cmd, src_stage, dst_stage, 0, 1, &b, 0, null, 0, null); } +// Fill geometry requires nonzero 4-byte-aligned offset/length within the buffer; host-tested. +pub fn fillRange(buf_len: u64, offset: u64, len: u64) Error!void { + if (len == 0 or len % 4 != 0 or offset % 4 != 0) return error.Misaligned; + try memory.span(buf_len, offset, len); +} + +// Record a 4-byte little-endian fill in a caller-begun command buffer; dst must be a device buffer from dev. +pub fn recordFill(dev: *const Device, cmd: c.VkCommandBuffer, dst: Buffer, offset: u64, len: u64, value: u32) Error!void { + if (dst.dev != dev or dst.kind != .device) return error.Invalid; + try fillRange(dst.len, offset, len); + dev.api.vkCmdFillBuffer(cmd, dst.handle, offset, len, value); +} + +/// Synchronous vkCmdFillBuffer of a device buffer range, with barriers; submits and waits. +pub fn fill(cmds: *Commands, dst: Buffer, offset: u64, len: u64, value: u32, timeout_ns: u64) Error!void { + if (dst.dev != cmds.dev or dst.kind != .device) return error.Invalid; + try fillRange(dst.len, offset, len); + const cmd = try cmds.begin(); + barrier(cmds.dev, cmd, gpu_stages, gpu_writes, transfer, transfer_write); + cmds.dev.api.vkCmdFillBuffer(cmd, dst.handle, offset, len, value); + barrier(cmds.dev, cmd, transfer, transfer_write, gpu_stages, gpu_access); + try cmds.submit(timeout_ns); +} + +/// fill(..., value = 0): the real device-side zero every family state uses. +pub fn zero(cmds: *Commands, dst: Buffer, offset: u64, len: u64, timeout_ns: u64) Error!void { + try fill(cmds, dst, offset, len, 0, timeout_ns); +} + const gpu_stages: u32 = c.VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | c.VK_PIPELINE_STAGE_TRANSFER_BIT; const gpu_writes: u32 = c.VK_ACCESS_SHADER_WRITE_BIT | c.VK_ACCESS_TRANSFER_WRITE_BIT; const gpu_access: u32 = gpu_writes | c.VK_ACCESS_SHADER_READ_BIT | c.VK_ACCESS_TRANSFER_READ_BIT; @@ -181,6 +247,164 @@ test "only a successful submit and a signalled fence let the caller free a submi try std.testing.expectError(error.InFlight, waited(c.VK_NOT_READY)); } +test "discard drops the open command buffer's records and clears the batch state" { + const Calls = struct { + var resets: usize = 0; + fn reset(_: c.VkCommandBuffer, _: c.VkCommandBufferResetFlags) callconv(.c) c.VkResult { + resets += 1; + return c.VK_SUCCESS; + } + }; + var api: @import("device.zig").DeviceApi = undefined; + api.vkResetCommandBuffer = Calls.reset; + var dev: Device = undefined; + dev.api = api; + var cmds: Commands = .{ .dev = &dev, .pool = undefined, .cmd = @ptrFromInt(0x1000), .fence = undefined, .recording = true, .pending = 7 }; + cmds.discard(); + try std.testing.expectEqual(@as(usize, 1), Calls.resets); + try std.testing.expectEqual(false, cmds.recording); + try std.testing.expectEqual(@as(u32, 0), cmds.pending); + cmds.discard(); // nothing recording: no second reset call + try std.testing.expectEqual(@as(usize, 1), Calls.resets); +} + +test "a failed reset still clears the batch state, and reuse fails without recording or submitting" { + const Calls = struct { + var resets: usize = 0; + var begins: usize = 0; + var submits: usize = 0; + fn reset(_: c.VkCommandBuffer, _: c.VkCommandBufferResetFlags) callconv(.c) c.VkResult { + resets += 1; + return c.VK_ERROR_OUT_OF_DEVICE_MEMORY; + } + fn begin(_: c.VkCommandBuffer, _: [*c]const c.VkCommandBufferBeginInfo) callconv(.c) c.VkResult { + begins += 1; + return c.VK_SUCCESS; + } + fn submit(_: c.VkQueue, _: c_uint, _: [*c]const c.VkSubmitInfo, _: c.VkFence) callconv(.c) c.VkResult { + submits += 1; + return c.VK_SUCCESS; + } + }; + var api: @import("device.zig").DeviceApi = undefined; + api.vkResetCommandBuffer = Calls.reset; + api.vkBeginCommandBuffer = Calls.begin; + api.vkQueueSubmit = Calls.submit; + var dev: Device = undefined; + dev.api = api; + var cmds: Commands = .{ .dev = &dev, .pool = undefined, .cmd = @ptrFromInt(0x1000), .fence = undefined, .recording = true, .pending = 7 }; + // Discard ignores the reset result on purpose: nothing recorded was ever submitted. + cmds.discard(); + try std.testing.expectEqual(@as(usize, 1), Calls.resets); + try std.testing.expectEqual(false, cmds.recording); + try std.testing.expectEqual(@as(u32, 0), cmds.pending); + // Submit is a no-op while not recording, and a failed begin reset adds no record. + try cmds.submit(0); + try std.testing.expectEqual(@as(usize, 0), Calls.submits); + try std.testing.expectError(error.OutOfMemory, cmds.begin()); + try std.testing.expectEqual(@as(usize, 2), Calls.resets); + try std.testing.expectEqual(@as(usize, 0), Calls.begins); + try std.testing.expectEqual(false, cmds.recording); +} + +test "a vkResetFences failure after a signaled wait poisons the stream terminally" { + const Calls = struct { + var ends: usize = 0; + var queues: usize = 0; + var waits: usize = 0; + var resets: usize = 0; + fn end(_: c.VkCommandBuffer) callconv(.c) c.VkResult { + ends += 1; + return c.VK_SUCCESS; + } + fn queueSubmit(_: c.VkQueue, _: c_uint, _: [*c]const c.VkSubmitInfo, _: c.VkFence) callconv(.c) c.VkResult { + queues += 1; + return c.VK_SUCCESS; + } + fn wait(_: c.VkDevice, _: c_uint, _: [*c]const c.VkFence, _: c.VkBool32, _: c_ulonglong) callconv(.c) c.VkResult { + waits += 1; + return c.VK_SUCCESS; + } + fn reset(_: c.VkDevice, _: c_uint, _: [*c]const c.VkFence) callconv(.c) c.VkResult { + resets += 1; + return c.VK_ERROR_OUT_OF_DEVICE_MEMORY; + } + }; + var api: @import("device.zig").DeviceApi = undefined; + api.vkEndCommandBuffer = Calls.end; + api.vkQueueSubmit = Calls.queueSubmit; + api.vkWaitForFences = Calls.wait; + api.vkResetFences = Calls.reset; + var dev: Device = undefined; + dev.api = api; + dev.handle = @ptrFromInt(0x2000); + dev.queue = @ptrFromInt(0x3000); + var cmds: Commands = .{ .dev = &dev, .pool = undefined, .cmd = @ptrFromInt(0x1000), .fence = @ptrFromInt(0x4000), .recording = true, .pending = 3 }; + std.testing.log_level = .err; // the poison path logs the reset failure, as it is meant to + defer std.testing.log_level = .warn; + // The fence is signaled after a good wait: a failed reset is terminal, and no submit may reuse it. + try std.testing.expectError(error.FencePoisoned, cmds.submit(0)); + try std.testing.expectEqual(@as(usize, 1), Calls.queues); + try std.testing.expectEqual(@as(usize, 1), Calls.waits); + try std.testing.expectEqual(@as(usize, 1), Calls.resets); + try std.testing.expectEqual(false, cmds.recording); + cmds.pending = 3; // pretend records arrived after the poisoning + try std.testing.expectError(error.FencePoisoned, cmds.submit(0)); + try std.testing.expectEqual(@as(usize, 1), Calls.queues); // no second submission ever reached the queue +} + +test "a failed vkEndCommandBuffer poisons the stream: a caught error cannot reach a recording begin" { + const Calls = struct { + var ends: usize = 0; + var begins: usize = 0; + var queues: usize = 0; + fn end(_: c.VkCommandBuffer) callconv(.c) c.VkResult { + ends += 1; + return c.VK_ERROR_OUT_OF_DEVICE_MEMORY; + } + fn begin(_: c.VkCommandBuffer, _: [*c]const c.VkCommandBufferBeginInfo) callconv(.c) c.VkResult { + begins += 1; + return c.VK_SUCCESS; + } + fn queueSubmit(_: c.VkQueue, _: c_uint, _: [*c]const c.VkSubmitInfo, _: c.VkFence) callconv(.c) c.VkResult { + queues += 1; + return c.VK_SUCCESS; + } + }; + var api: @import("device.zig").DeviceApi = undefined; + api.vkEndCommandBuffer = Calls.end; + api.vkBeginCommandBuffer = Calls.begin; + api.vkQueueSubmit = Calls.queueSubmit; + var dev: Device = undefined; + dev.api = api; + dev.handle = @ptrFromInt(0x2000); + dev.queue = @ptrFromInt(0x3000); + var cmds: Commands = .{ .dev = &dev, .pool = undefined, .cmd = @ptrFromInt(0x1000), .fence = @ptrFromInt(0x4000), .recording = true, .pending = 5 }; + std.testing.log_level = .err; // the poison path logs the end failure, as it is meant to + defer std.testing.log_level = .warn; + // A gate wrapper catches this and continues: the buffer is invalid and must never be handed out again. + try std.testing.expectError(error.CommandBufferInvalid, cmds.submit(0)); + try std.testing.expectEqual(@as(usize, 1), Calls.ends); + try std.testing.expectEqual(@as(usize, 0), Calls.queues); + try std.testing.expectEqual(false, cmds.recording); + // The caught-and-continued path: begin must refuse, not return the invalid open buffer. + try std.testing.expectError(error.CommandBufferInvalid, cmds.begin()); + try std.testing.expectEqual(@as(usize, 0), Calls.begins); + // And submit refuses without queueing the poisoned records. + try std.testing.expectError(error.CommandBufferInvalid, cmds.submit(0)); + try std.testing.expectEqual(@as(usize, 0), Calls.queues); +} + +test "fill geometry refuses torn, empty and out-of-range fills" { + try fillRange(4096, 0, 4096); + try fillRange(4096, 12, 1024); + try std.testing.expectError(error.Misaligned, fillRange(4096, 0, 0)); // empty + try std.testing.expectError(error.Misaligned, fillRange(4096, 2, 1024)); + try std.testing.expectError(error.Misaligned, fillRange(4096, 0, 6)); + try std.testing.expectError(error.OutOfRange, fillRange(4096, 8, 4096)); + try std.testing.expectError(error.OutOfRange, fillRange(0, 0, 4)); +} + test "staged transfers refuse foreign devices, wrong kinds, aliasing and ranges past the buffer" { // operand identities only: nothing here is dereferenced or reaches Vulkan const dev: *const Device = @ptrFromInt(0x1000); diff --git a/zig/src/vulkan/device.zig b/zig/src/vulkan/device.zig index 06354b2b460..08a5abb9e2b 100644 --- a/zig/src/vulkan/device.zig +++ b/zig/src/vulkan/device.zig @@ -32,6 +32,7 @@ pub const DeviceApi = struct { vkWaitForFences: Fn(c.PFN_vkWaitForFences), vkResetFences: Fn(c.PFN_vkResetFences), vkCmdPipelineBarrier: Fn(c.PFN_vkCmdPipelineBarrier), + vkCmdFillBuffer: Fn(c.PFN_vkCmdFillBuffer), vkCmdCopyBuffer: Fn(c.PFN_vkCmdCopyBuffer), vkCmdBindPipeline: Fn(c.PFN_vkCmdBindPipeline), vkCmdBindDescriptorSets: Fn(c.PFN_vkCmdBindDescriptorSets), diff --git a/zig/src/vulkan/dispatch.zig b/zig/src/vulkan/dispatch.zig index a74e2954350..598e98ca5c8 100644 --- a/zig/src/vulkan/dispatch.zig +++ b/zig/src/vulkan/dispatch.zig @@ -10,7 +10,11 @@ const check = runtime.check; pub const Error = runtime.Error || error{ BadSpirv, NoLocalSize, TooManyBindings, BadPushSize, BadWorkgroup, BadGroupCount, BadBindings, Aliased }; -pub const max_bindings = 8; +pub const max_bindings = 10; + +/// Dispatches recorded into one submission, so a descriptor set per dispatch; run submits at half of it. +pub const batch_max: u32 = 256; +const batch_flush: u32 = batch_max / 2; /// The shader's set 0: storage buffers 0..bindings-1, its push bytes, and a mask of the bindings it writes. pub const Layout = struct { bindings: u32, push_bytes: u32, writes: u32 }; @@ -27,7 +31,9 @@ pub const Pipeline = struct { pipeline_layout: c.VkPipelineLayout = null, pipeline: c.VkPipeline = null, pool: c.VkDescriptorPool = null, - set: c.VkDescriptorSet = null, + /// One set per dispatch recorded since the last submission; bind fills the next one, record binds the current one. + sets: [batch_max]c.VkDescriptorSet = @splat(null), + slot: usize = 0, bound: bool = false, /// Checks the module and the device's limits before creating anything; a failure destroys what was made. @@ -98,22 +104,23 @@ pub const Pipeline = struct { return e; }; errdefer api.vkDestroyPipeline(h, p.pipeline, null); - const size: c.VkDescriptorPoolSize = .{ .type = c.VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, .descriptorCount = layout.bindings }; + const size: c.VkDescriptorPoolSize = .{ .type = c.VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, .descriptorCount = layout.bindings * batch_max }; const pool_info: c.VkDescriptorPoolCreateInfo = .{ .sType = c.VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO, - .maxSets = 1, + .maxSets = batch_max, .poolSizeCount = 1, .pPoolSizes = &size, }; try check(api.vkCreateDescriptorPool(h, &pool_info, null, &p.pool)); errdefer api.vkDestroyDescriptorPool(h, p.pool, null); + const set_layouts: [batch_max]c.VkDescriptorSetLayout = @splat(p.set_layout); const alloc_info: c.VkDescriptorSetAllocateInfo = .{ .sType = c.VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO, .descriptorPool = p.pool, - .descriptorSetCount = 1, - .pSetLayouts = &p.set_layout, + .descriptorSetCount = batch_max, + .pSetLayouts = &set_layouts, }; - try check(api.vkAllocateDescriptorSets(h, &alloc_info, &p.set)); + try check(api.vkAllocateDescriptorSets(h, &alloc_info, &p.sets)); return p; } @@ -127,7 +134,7 @@ pub const Pipeline = struct { api.vkDestroyDescriptorSetLayout(h, p.set_layout, null); } - /// Updates checked device-local bindings; call only when no submission using this descriptor set is pending. + /// Fills the next descriptor set with checked device-local bindings; earlier recorded dispatches keep their sets. pub fn bind(p: *Pipeline, bindings: []const Binding) Error!void { if (bindings.len != p.layout.bindings) return error.BadBindings; var infos: [max_bindings]c.VkDescriptorBufferInfo = undefined; @@ -138,7 +145,7 @@ pub const Pipeline = struct { infos[i] = try b.buffer.descriptor(b.offset, b.len); writes[i] = .{ .sType = c.VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET, - .dstSet = p.set, + .dstSet = p.sets[(p.slot + 1) % batch_max], .dstBinding = @intCast(i), .descriptorCount = 1, .descriptorType = c.VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, @@ -147,6 +154,7 @@ pub const Pipeline = struct { } try disjoint(bindings, p.layout.writes); p.dev.api.vkUpdateDescriptorSets(p.dev.handle, p.layout.bindings, &writes, 0, null); + p.slot = (p.slot + 1) % batch_max; p.bound = true; } @@ -161,7 +169,7 @@ pub const Pipeline = struct { }; const api = &p.dev.api; api.vkCmdBindPipeline(cmd, c.VK_PIPELINE_BIND_POINT_COMPUTE, p.pipeline); - api.vkCmdBindDescriptorSets(cmd, c.VK_PIPELINE_BIND_POINT_COMPUTE, p.pipeline_layout, 0, 1, &p.set, 0, null); + api.vkCmdBindDescriptorSets(cmd, c.VK_PIPELINE_BIND_POINT_COMPUTE, p.pipeline_layout, 0, 1, &p.sets[p.slot], 0, null); if (push.len > 0) api.vkCmdPushConstants(cmd, p.pipeline_layout, c.VK_SHADER_STAGE_COMPUTE_BIT, 0, p.layout.push_bytes, push.ptr); api.vkCmdDispatch(cmd, groups[0], groups[1], groups[2]); const reads = c.VK_ACCESS_SHADER_READ_BIT | c.VK_ACCESS_SHADER_WRITE_BIT | c.VK_ACCESS_TRANSFER_READ_BIT | c.VK_ACCESS_TRANSFER_WRITE_BIT; @@ -169,10 +177,11 @@ pub const Pipeline = struct { commands.barrier(p.dev, cmd, c.VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, c.VK_ACCESS_SHADER_WRITE_BIT, stages, reads); } - /// One dispatch in its own submission; wrap it in commands.orExit, which ends the process on error.Timeout. + /// Records one dispatch into the open command buffer; submitted with half a batch, or with the next copy or fill. pub fn run(p: *const Pipeline, cmds: *commands.Commands, push: []const u8, groups: [3]u32, timeout_ns: u64) Error!void { try p.record(try cmds.begin(), push, groups); - try cmds.submit(timeout_ns); + cmds.pending += 1; + if (cmds.pending >= batch_flush) try cmds.submit(timeout_ns); } }; @@ -314,7 +323,17 @@ test "device limits refuse bindings, push bytes, workgroups and group counts pas try std.testing.expectError(error.BadWorkgroup, fits(l, qmm, .{ 512, 1, 1 })); l.maxComputeWorkGroupInvocations = 128; try std.testing.expectError(error.BadWorkgroup, fits(l, qmm, .{ 256, 1, 1 })); + // restore the limit the negative case reduced, before the positive cases + l.maxComputeWorkGroupInvocations = 256; try std.testing.expectError(error.BadGroupCount, counts(.{ 65535, 65535, 65535 }, .{ 0, 1, 1 })); try std.testing.expectError(error.BadGroupCount, counts(.{ 65535, 65535, 65535 }, .{ 65536, 1, 1 })); try counts(.{ 65535, 65535, 65535 }, .{ 2, 1, 1 }); + // The 10-binding GDN prefill/decode closures fit; 11 do not; fewer offered still refuses (real capacity). + l.maxPerStageDescriptorStorageBuffers = 10; + l.maxDescriptorSetStorageBuffers = 10; + l.maxPerStageResources = 10; + try fits(l, .{ .bindings = max_bindings, .push_bytes = 120, .writes = 1 << 9 }, .{ 256, 1, 1 }); + try std.testing.expectError(error.TooManyBindings, fits(l, .{ .bindings = max_bindings + 1, .push_bytes = 120, .writes = 1 << 10 }, .{ 256, 1, 1 })); + l.maxPerStageDescriptorStorageBuffers = 8; + try std.testing.expectError(error.TooManyBindings, fits(l, .{ .bindings = max_bindings, .push_bytes = 120, .writes = 1 << 9 }, .{ 256, 1, 1 })); } diff --git a/zig/src/vulkan/gdn.zig b/zig/src/vulkan/gdn.zig new file mode 100644 index 00000000000..e41094dc482 --- /dev/null +++ b/zig/src/vulkan/gdn.zig @@ -0,0 +1,555 @@ +// Vulkan host dispatch for the Gated DeltaNet kernels. + +const std = @import("std"); +const dispatch = @import("dispatch.zig"); +const memory = @import("memory.zig"); +const params = @import("stage1_params"); +const tensor = @import("tensor.zig"); +const c = @import("vk"); +const Device = @import("device.zig").Device; +const Commands = @import("commands.zig").Commands; +const Buffer = memory.Buffer; +const Error = dispatch.Error; +const View = tensor.View; + +/// The fixed Qwen3.5 linear-layer geometry these primitives serve. +pub const Model = struct { + pub const linear_layers = 18; + pub const heads = 16; + pub const head_dim = 128; + pub const taps = 4; + pub const conv_channels = 6144; + pub const conv_state_rows = taps - 1; + pub const linear_dim = heads * head_dim; // 2048 + pub const state_elements = heads * head_dim * head_dim; // 262144 f32 +}; + +/// qk l2 norm epsilon term, as the Metal oracle writes it: 1e-6 / DK. +const qk_norm_eps: f32 = 1e-6 / @as(f32, Model.head_dim); + +const dummy_bytes: u64 = 256; + +comptime { + if (Model.conv_channels != 2 * Model.heads * Model.head_dim + Model.heads * Model.head_dim) + @compileError("conv channels drifted from 2*NK*DK + NV*DV"); +} + +pub const Shaders = struct { + /// conv_dw1d.comp -DUSE_BF16 + conv_dw1d_bf16: []align(4) const u8, + // gdn_conv_decode.comp -DUSE_BF16 (no APPLE_REDUCTION: the oracle arm uses flag bit 0 only) + gdn_conv_decode_bf16: []align(4) const u8, + /// gdn_split.comp -DUSE_BF16 + gdn_split_bf16: []align(4) const u8, + // Both prefill passes share this image; a push-constant flag selects the pass. + prefill_recur_bf16: []align(4) const u8, + /// gated_delta_decode.comp -DUSE_BF16 (raw gates ride flags bit3|bit9) + gated_delta_decode_bf16: []align(4) const u8, + /// fast_norm_gated.comp -DUSE_BF16 (default build, mode 0 at runtime) + fast_norm_gated_bf16: []align(4) const u8, + /// copy_general.comp -DUSE_U16 + copy_u16: []align(4) const u8, + /// conv_dw1d.comp -DUSE_BF16 -DACT_F32: f32 activations over the bf16 taps for the wide layers + conv_dw1d_f32: []align(4) const u8, + /// gdn_split.comp -DACT_F32 + gdn_split_f32: []align(4) const u8, + /// gated_delta_prefill_oracle.comp -DACT_F32 + prefill_recur_f32: []align(4) const u8, + /// fast_norm_gated.comp -DACT_F32 + fast_norm_gated_f32: []align(4) const u8, + /// copy_general.comp -DUSE_U32: the wide prefill conv-state copies + copy_u32: []align(4) const u8, +}; + +/// Push block sizes, each equal to the GLSL block glslang emits. +const conv_push_bytes: u32 = 36; // count..output_offset (9 words) +const conv_decode_push_bytes: u32 = 60; // count..flags (15 words) +const full_push_bytes: u32 = params.push_bytes; // 120 +// gated_delta_prefill_oracle.comp appends `chunk` and `chunks` after the standard Params: 32 words. +const prefill_push_bytes: u32 = 128; +const decode_raw_gates: u32 = 8; +const decode_f32_alog: u32 = 512; +// Prefill push flags (gated_delta_prefill_oracle.comp): bit2 g is f32, and the output pass adds bit3. +const scan_flags: u32 = 4; +const out_flags: u32 = 4 | 8; +// Decode SiLU/norm flags stay off; the split owns Metal's single SiLU. +const conv_decode_flags: u32 = 0; + +// Params plus the prefill scan's trailing chunk fields, matching the shader's push block field for field. +const PrefillPush = extern struct { + base: params.Params = .{}, + chunk: u32 = 0, + chunks: u32 = 0, +}; + +comptime { + if (@sizeOf(PrefillPush) != prefill_push_bytes) @compileError("prefill push block drifted"); + if (@offsetOf(PrefillPush, "chunk") != params.push_bytes) @compileError("prefill chunk field misaligned"); +} +const norm_push_bytes: u32 = @offsetOf(params.Params, "beta") + 4; // 68 + +fn groups(count: u64, local: u32, max_x: u32) Error![3]u32 { + const n = std.math.cast(u32, (count + local - 1) / local) orelse return error.TooLarge; + return .{ @min(n, max_x), 1, 1 }; +} + +fn wantKind(kind: tensor.Kind, views: []const View) Error!void { + for (views) |v| if (v.kind != kind) return error.Invalid; +} + +fn wantBf16(views: []const View) Error!void { + try wantKind(.bf16, views); +} + +fn wantF32(views: []const View) Error!void { + for (views) |v| if (v.kind != .f32) return error.Invalid; +} + +// Standalone and composite split outputs are bf16 soa/beta and f32 g (narrow) or f32 soa/beta/g (wide). +pub fn wantSplitOutputs(out_soa: View, g_out: View, beta_out: View) Error!void { + if (out_soa.kind != .bf16 and out_soa.kind != .f32) return error.Invalid; + if (beta_out.kind != out_soa.kind) return error.Invalid; + try wantF32(&.{g_out}); +} + +/// The split's activation kinds: conv_out and out_soa share bf16 (narrow) or f32 (wide layers); pure and host-tested. +pub const SplitVariant = enum { bf16, f32 }; + +pub fn splitVariant(conv_out: View, out_soa: View) Error!SplitVariant { + if (conv_out.kind != out_soa.kind) return error.Invalid; + if (conv_out.kind != .bf16 and conv_out.kind != .f32) return error.Invalid; + return if (conv_out.kind == .f32) .f32 else .bf16; +} + +/// The gated norm's kinds: bf16 everywhere, or f32 y/z/out over a bf16 weight; pure and host-tested. +pub const GatedNormVariant = enum { bf16, f32 }; + +pub fn gatedNormVariant(out: View, y: View, z: View, weight: View) Error!GatedNormVariant { + if (weight.kind != .bf16) return error.Invalid; + if (out.kind != y.kind or y.kind != z.kind) return error.Invalid; + if (out.kind != .bf16 and out.kind != .f32) return error.Invalid; + return if (out.kind == .f32) .f32 else .bf16; +} + +// Reject cross-dispatch read/write overlap; bind checks same-dispatch overlap. +fn refuseOverlap(writes: []const View, reads: []const View) Error!void { + for (writes) |w| for (reads) |r| { + if (w.overlaps(r)) return error.Aliased; + }; +} + +pub const Gdn = struct { + dev: *const Device, + copy_u16: dispatch.Pipeline, + conv_bf16: dispatch.Pipeline, + conv_decode_bf16: dispatch.Pipeline, + split_bf16: dispatch.Pipeline, + prefill_bf16: dispatch.Pipeline, + decode_bf16: dispatch.Pipeline, + norm_bf16: dispatch.Pipeline, + split_f32: dispatch.Pipeline, + prefill_f32: dispatch.Pipeline, + norm_f32: dispatch.Pipeline, + conv_f32: dispatch.Pipeline, + copy_u32: dispatch.Pipeline, + /// Stand-in for declared bindings a dispatch never touches. + dummy: Buffer, + + pub fn init(dev: *const Device, shaders: Shaders) Error!Gdn { + if (!dev.features.storage16 or !dev.features.int16) return error.MissingFeature; + // Row sums need relative shuffles and four row lanes in a subgroup divisible by four. + if (dev.v11.subgroupSupportedStages & c.VK_SHADER_STAGE_COMPUTE_BIT == 0) return error.MissingFeature; + if (dev.v11.subgroupSupportedOperations & c.VK_SUBGROUP_FEATURE_SHUFFLE_RELATIVE_BIT == 0) return error.MissingFeature; + if (dev.v11.subgroupSize == 0 or dev.v11.subgroupSize % 4 != 0) return error.MissingFeature; + var g: Gdn = .{ + .dev = dev, + .copy_u16 = undefined, + .conv_bf16 = undefined, + .conv_decode_bf16 = undefined, + .split_bf16 = undefined, + .prefill_bf16 = undefined, + .decode_bf16 = undefined, + .norm_bf16 = undefined, + .split_f32 = undefined, + .prefill_f32 = undefined, + .norm_f32 = undefined, + .conv_f32 = undefined, + .copy_u32 = undefined, + .dummy = try Buffer.init(dev, .device, 2 * dummy_bytes), + }; + errdefer g.dummy.deinit(); + g.copy_u16 = try dispatch.Pipeline.init(dev, shaders.copy_u16, .{ .bindings = 4, .push_bytes = full_push_bytes, .writes = 1 << 2 }); + errdefer g.copy_u16.deinit(); + g.conv_bf16 = try dispatch.Pipeline.init(dev, shaders.conv_dw1d_bf16, .{ .bindings = 3, .push_bytes = conv_push_bytes, .writes = 1 << 2 }); + errdefer g.conv_bf16.deinit(); + g.conv_decode_bf16 = try dispatch.Pipeline.init(dev, shaders.gdn_conv_decode_bf16, .{ .bindings = 5, .push_bytes = conv_decode_push_bytes, .writes = (1 << 3) | (1 << 4) }); + errdefer g.conv_decode_bf16.deinit(); + g.split_bf16 = try dispatch.Pipeline.init(dev, shaders.gdn_split_bf16, .{ .bindings = 8, .push_bytes = full_push_bytes, .writes = 0b1110_0000 }); + errdefer g.split_bf16.deinit(); + // 10 bindings each: refused with TooManyBindings while dispatch.max_bindings is 8. See the module doc block. + g.prefill_bf16 = try dispatch.Pipeline.init(dev, shaders.prefill_recur_bf16, .{ .bindings = 10, .push_bytes = prefill_push_bytes, .writes = (1 << 6) | (1 << 7) | (1 << 9) }); + errdefer g.prefill_bf16.deinit(); + g.decode_bf16 = try dispatch.Pipeline.init(dev, shaders.gated_delta_decode_bf16, .{ .bindings = 10, .push_bytes = full_push_bytes, .writes = (1 << 6) | (1 << 7) }); + errdefer g.decode_bf16.deinit(); + g.norm_bf16 = try dispatch.Pipeline.init(dev, shaders.fast_norm_gated_bf16, .{ .bindings = 4, .push_bytes = norm_push_bytes, .writes = 1 << 3 }); + errdefer g.norm_bf16.deinit(); + g.split_f32 = try dispatch.Pipeline.init(dev, shaders.gdn_split_f32, .{ .bindings = 8, .push_bytes = full_push_bytes, .writes = 0b1110_0000 }); + errdefer g.split_f32.deinit(); + g.prefill_f32 = try dispatch.Pipeline.init(dev, shaders.prefill_recur_f32, .{ .bindings = 10, .push_bytes = prefill_push_bytes, .writes = (1 << 6) | (1 << 7) | (1 << 9) }); + errdefer g.prefill_f32.deinit(); + g.norm_f32 = try dispatch.Pipeline.init(dev, shaders.fast_norm_gated_f32, .{ .bindings = 4, .push_bytes = norm_push_bytes, .writes = 1 << 3 }); + errdefer g.norm_f32.deinit(); + g.conv_f32 = try dispatch.Pipeline.init(dev, shaders.conv_dw1d_f32, .{ .bindings = 3, .push_bytes = conv_push_bytes, .writes = 1 << 2 }); + errdefer g.conv_f32.deinit(); + g.copy_u32 = try dispatch.Pipeline.init(dev, shaders.copy_u32, .{ .bindings = 4, .push_bytes = full_push_bytes, .writes = 1 << 2 }); + return g; + } + + pub fn deinit(m: *Gdn) void { + m.copy_u16.deinit(); + m.conv_bf16.deinit(); + m.conv_decode_bf16.deinit(); + m.split_bf16.deinit(); + m.prefill_bf16.deinit(); + m.decode_bf16.deinit(); + m.norm_bf16.deinit(); + m.split_f32.deinit(); + m.prefill_f32.deinit(); + m.norm_f32.deinit(); + m.conv_f32.deinit(); + m.copy_u32.deinit(); + m.dummy.deinit(); + } + + fn dummyView(m: *const Gdn, kind: tensor.Kind) Error!View { + return View.init(&m.dummy, kind, 0, dummy_bytes / kind.size()); + } + + fn snapshotDummyView(m: *const Gdn) Error!View { + // Binding 9 is declared writable even when the single-chunk shader never touches it. + return View.init(&m.dummy, .f32, dummy_bytes, dummy_bytes / tensor.Kind.f32.size()); + } + + fn bindPrefill(pipe: *dispatch.Pipeline, bindings: []const dispatch.Binding) Error!void { + pipe.bind(bindings) catch |e| { + if (e == error.Aliased) { + const names = [_][]const u8{ "q", "k", "v", "g_bf16_dummy", "beta", "h0", "y", "hf", "g_f32", "snapshot" }; + for (bindings, 0..) |w, i| if (pipe.layout.writes >> @intCast(i) & 1 != 0) { + for (bindings, 0..) |r, j| { + if (i != j and w.buffer.handle == r.buffer.handle and w.offset < r.offset + r.len and r.offset < w.offset + w.len) + std.debug.print("GDN prefill alias: {s} binding={d} buffer={any} bytes=[{d},{d}) overlaps {s} binding={d} buffer={any} bytes=[{d},{d})\n", .{ + names[i], i, w.buffer.handle, w.offset, w.offset + w.len, names[j], j, r.buffer.handle, r.offset, r.offset + r.len, + }); + } + }; + } + return e; + }; + } + + // Same-shape 1-D copy; narrow rides raw u16 words, wide f32 rides raw u32 words. + fn copyAny(m: *Gdn, cmds: *Commands, wide: bool, out: View, src: View, timeout_ns: u64) Error!void { + const d = try m.dummyView(.u32); + const pipe = if (wide) &m.copy_u32 else &m.copy_u16; + try pipe.bind(&.{ src.bind(), d.bind(), out.bind(), d.bind() }); + const count = try src.elemCount32(); + var p: params.Params = .{ + .count = count, + .dims = 1, + .shape = .{ count, 0, 0, 0 }, + .in_strides = .{ 1, 0, 0, 0 }, + .out_strides = .{ 1, 0, 0, 0 }, + }; + try pipe.run(cmds, std.mem.asBytes(&p), try groups(count, pipe.local[0], m.dev.properties.limits.maxComputeWorkGroupCount[0]), timeout_ns); + } + + // Split applies SiLU and Q/K L2 norm; V passes through as SoA, gates are f32 g and beta in the layer's kind. + pub fn split(m: *Gdn, cmds: *Commands, out_soa: View, g_out: View, beta_out: View, conv_out: View, a: View, b: View, a_log: View, dt: View, tokens: u64, timeout_ns: u64) Error!void { + const variant = try splitVariant(conv_out, out_soa); + try wantKind(conv_out.kind, &.{ a, b, beta_out }); + try wantBf16(&.{dt}); + try wantF32(&.{ a_log, g_out }); + const plan = try splitPlan(conv_out.elements, out_soa.elements, g_out.elements, beta_out.elements, a.elements, b.elements, a_log.elements, dt.elements, tokens); + try refuseOverlap(&.{ out_soa, g_out, beta_out }, &.{ conv_out, a, b, a_log, dt }); + const pipe = if (variant == .f32) &m.split_f32 else &m.split_bf16; + try pipe.bind(&.{ + conv_out.bind(), a.bind(), b.bind(), a_log.bind(), dt.bind(), + out_soa.bind(), g_out.bind(), beta_out.bind(), + }); + var p: params.Params = .{ + .count = plan.tokens, + .lhs_size = plan.conv_channels, + .rhs_size = plan.tokens, + .output_size = 3 * Model.heads, + }; + const most = m.dev.properties.limits.maxComputeWorkGroupCount; + if (plan.tokens > most[0] or 3 * Model.heads > most[1]) return error.BadGroupCount; + try pipe.run(cmds, std.mem.asBytes(&p), .{ plan.tokens, 3 * Model.heads, 1 }, timeout_ns); + } + + // Depthwise conv over the stored conv state and the new tokens, then the SoA split to q/k/v, g and beta. + pub fn convPrefill(m: *Gdn, cmds: *Commands, out_soa: View, g_out: View, beta_out: View, conv_out: View, conv_input: View, conv_state_in: View, qkv: View, weight: View, a: View, b: View, a_log: View, dt: View, conv_state_out: View, tokens: u64, timeout_ns: u64) Error!void { + const variant = try splitVariant(conv_out, out_soa); + const wide = variant == .f32; + try wantKind(conv_out.kind, &.{ conv_input, conv_state_in, qkv, conv_state_out, beta_out, a, b }); + try wantBf16(&.{ weight, dt }); + try wantF32(&.{ a_log, g_out }); + const plan = try convPrefillPlan(conv_input.elements, conv_state_in.elements, qkv.elements, weight.elements, conv_state_out.elements, tokens); + try refuseOverlap(&.{ conv_input, conv_state_out }, &.{ conv_state_in, qkv, weight }); + try refuseOverlap(&.{ out_soa, g_out, beta_out, conv_out }, &.{ conv_input, a, b, a_log, dt }); + // rows 0..2 of the conv input are the stored conv state + try m.copyAny(cmds, wide, try conv_input.slice(0, plan.state_elements), conv_state_in, timeout_ns); + // rows 3.. are the projection output + try m.copyAny(cmds, wide, try conv_input.slice(plan.state_elements, plan.tokens_elements), qkv, timeout_ns); + const conv_pipe = if (wide) &m.conv_f32 else &m.conv_bf16; + try conv_pipe.bind(&.{ conv_input.bind(), weight.bind(), conv_out.bind() }); + var p: params.Params = .{ + .count = plan.count, + .operation = Model.taps, + .lhs_size = plan.l_in, + .rhs_size = plan.tokens, + .reduce_size = Model.conv_channels, + .output_size = 1, + }; + try conv_pipe.run(cmds, std.mem.asBytes(&p)[0..conv_push_bytes], try groups(plan.count, conv_pipe.local[0], m.dev.properties.limits.maxComputeWorkGroupCount[0]), timeout_ns); + try m.split(cmds, out_soa, g_out, beta_out, conv_out, a, b, a_log, dt, tokens, timeout_ns); + // the next conv state: conv input rows tokens..tokens+3 + try m.copyAny(cmds, wide, conv_state_out, try conv_input.slice(plan.tokens_elements, plan.state_elements), timeout_ns); + } + + // Decode fuses state+x, 4-tap conv and shifted state; raw bf16 output lets split apply the sole SiLU. + pub fn convDecodeStep(m: *Gdn, cmds: *Commands, out_soa: View, g_out: View, beta_out: View, activated: View, conv_state_in: View, x: View, weight: View, a: View, b: View, a_log: View, dt: View, conv_state_out: View, timeout_ns: u64) Error!void { + try wantBf16(&.{ activated, conv_state_in, x, weight, a, b, dt, conv_state_out }); + try wantSplitOutputs(out_soa, g_out, beta_out); + try wantF32(&.{a_log}); + const plan = try convDecodePlan(conv_state_in.elements, x.elements, weight.elements, activated.elements, conv_state_out.elements, out_soa.elements, g_out.elements, beta_out.elements, a.elements, b.elements, a_log.elements, dt.elements); + try refuseOverlap(&.{ activated, conv_state_out }, &.{ conv_state_in, x, weight }); + try refuseOverlap(&.{ out_soa, g_out, beta_out }, &.{ activated, a, b, a_log, dt }); + try m.conv_decode_bf16.bind(&.{ conv_state_in.bind(), x.bind(), weight.bind(), activated.bind(), conv_state_out.bind() }); + var p: params.Params = .{ + .count = plan.count, + .operation = Model.taps, + .lhs_size = Model.conv_state_rows, + .reduce_size = plan.conv_channels, + .flags = conv_decode_flags, + }; + try m.conv_decode_bf16.run(cmds, std.mem.asBytes(&p)[0..conv_decode_push_bytes], try groups(plan.count, m.conv_decode_bf16.local[0], m.dev.properties.limits.maxComputeWorkGroupCount[0]), timeout_ns); + try m.split(cmds, out_soa, g_out, beta_out, activated, a, b, a_log, dt, 1, timeout_ns); + } + + // Delta-net prefill recurrence: h0 in, hf out, y written; snap is an optional f32 side output. + pub fn prefillRecur(m: *Gdn, cmds: *Commands, soa: View, g: View, beta: View, h0: View, y: View, hf: View, snap: ?View, tokens: u64, chunk: u64, timeout_ns: u64) Error!void { + if (beta.kind != soa.kind) return error.Invalid; + try wantF32(&.{ g, h0, hf }); + if (snap) |s| try wantF32(&.{s}); + // the scan's activations: soa and y share bf16 (narrow) or f32 (wide layers) + if (soa.kind != y.kind or (soa.kind != .bf16 and soa.kind != .f32)) return error.Invalid; + const pipe = if (soa.kind == .f32) &m.prefill_f32 else &m.prefill_bf16; + const plan = try prefillPlan(soa.elements, g.elements, beta.elements, h0.elements, y.elements, hf.elements, if (snap) |s| s.elements else 0, tokens, chunk); + const q = try soa.slice(0, plan.section_elements); + const k = try soa.slice(plan.section_elements, plan.section_elements); + const v = try soa.slice(2 * plan.section_elements, plan.section_elements); + try refuseOverlap(&.{ y, hf }, &.{ soa, g, beta, h0 }); + if (snap) |s| try refuseOverlap(&.{s}, &.{ h0, hf, y }); + if (plan.chunks > 1) { + const snap_v = snap orelse return error.Invalid; + const db = try m.dummyView(.bf16); + try bindPrefill(pipe, &.{ + q.bind(), k.bind(), v.bind(), db.bind(), beta.bind(), + h0.bind(), y.bind(), hf.bind(), g.bind(), snap_v.bind(), + }); + try m.runPrefillPass(pipe, cmds, scan_flags, plan, timeout_ns); + } + const snap_binding = snap orelse try m.snapshotDummyView(); + const db = try m.dummyView(.bf16); + try bindPrefill(pipe, &.{ + q.bind(), k.bind(), v.bind(), db.bind(), beta.bind(), + h0.bind(), y.bind(), hf.bind(), g.bind(), snap_binding.bind(), + }); + try m.runPrefillPass(pipe, cmds, out_flags, plan, timeout_ns); + } + + fn runPrefillPass(m: *Gdn, pipe: *dispatch.Pipeline, cmds: *Commands, flags: u32, plan: PrefillPlan, timeout_ns: u64) Error!void { + const base: params.Params = .{ + .dims = plan.tokens, + .matrix_m = Model.head_dim, // Dk + .matrix_n = Model.head_dim, // Dv + .matrix_k = Model.heads, // Hv + .flags = flags, + .shape = .{ 0, 0, 0, 0 }, // views bind at their base + }; + const p: PrefillPush = .{ .base = base, .chunk = plan.chunk, .chunks = plan.chunks }; + const grid_x = Model.heads * 4; // LANES == 4 row blocks per head + const most = m.dev.properties.limits.maxComputeWorkGroupCount; + if (grid_x > most[0] or plan.chunks > most[1]) return error.BadGroupCount; + try pipe.run(cmds, std.mem.asBytes(&p), .{ grid_x, plan.chunks, 1 }, timeout_ns); + } + + // Decode T=1 computes Metal-rounded raw gates and f32 A_log; state writes out of place. + pub fn decodeRecur(m: *Gdn, cmds: *Commands, soa: View, a: View, b: View, a_log_u16: View, dt: View, state_in: View, state_out: View, y: View, timeout_ns: u64) Error!void { + try wantBf16(&.{ soa, a, b, dt, y }); + try wantF32(&.{ state_in, state_out }); + if (a_log_u16.kind != .u16) return error.Invalid; + _ = try decodePlan(soa.elements, a.elements, b.elements, a_log_u16.elements, dt.elements, state_in.elements, state_out.elements, y.elements); + const q = try soa.slice(0, Model.linear_dim); + const k = try soa.slice(Model.linear_dim, Model.linear_dim); + const v = try soa.slice(2 * Model.linear_dim, Model.linear_dim); + try refuseOverlap(&.{ state_out, y }, &.{ soa, a, b, a_log_u16, dt, state_in }); + try m.decode_bf16.bind(&.{ + q.bind(), k.bind(), v.bind(), a.bind(), b.bind(), + state_in.bind(), y.bind(), state_out.bind(), a_log_u16.bind(), dt.bind(), + }); + var p: params.Params = .{ + .matrix_m = Model.head_dim, + .matrix_n = Model.head_dim, + .matrix_k = Model.heads, + .flags = decode_raw_gates | decode_f32_alog, + .shape = .{ 0, 0, 0, 0 }, + .in_strides = .{ 0, 0, Model.heads, 0 }, // [2] = Hv for head % Hv + }; + const most = m.dev.properties.limits.maxComputeWorkGroupCount; + if (Model.heads > most[0] or (Model.head_dim / 32) > most[1]) return error.BadGroupCount; + try m.decode_bf16.run(cmds, std.mem.asBytes(&p), .{ Model.heads, Model.head_dim / 32, 1 }, timeout_ns); + } + + // Gated RMSNorm of y with weight, gated by z. + pub fn gatedNorm(m: *Gdn, cmds: *Commands, out: View, y: View, z: View, weight: View, rows: u64, eps: f32, timeout_ns: u64) Error!void { + const variant = try gatedNormVariant(out, y, z, weight); + const plan = try gatedNormPlan(out.elements, y.elements, z.elements, weight.elements, rows, eps); + try refuseOverlap(&.{out}, &.{ y, z, weight }); + const pipe = if (variant == .f32) &m.norm_f32 else &m.norm_bf16; + try pipe.bind(&.{ y.bind(), weight.bind(), z.bind(), out.bind() }); + var p: params.Params = .{ + .operation = 0, // gated mode + .lhs_size = Model.head_dim, // per-column weight within the head (!= 1) + .reduce_size = Model.head_dim, // the 128-wide per-head row + .output_size = plan.rows, // tokens * 16 head-rows + }; + p.alpha = eps; + const g = try groups(plan.rows, 1, m.dev.properties.limits.maxComputeWorkGroupCount[0]); + try pipe.run(cmds, std.mem.asBytes(&p)[0..norm_push_bytes], g, timeout_ns); + } +}; + +/// The GDN geometry plans live in gdn_plan.zig; every plan name stays reachable from here. +const plans = @import("gdn_plan.zig"); +pub const SplitPlan = plans.SplitPlan; +pub const splitPlan = plans.splitPlan; +pub const ConvPrefillPlan = plans.ConvPrefillPlan; +pub const convPrefillPlan = plans.convPrefillPlan; +pub const ConvDecodePlan = plans.ConvDecodePlan; +pub const convDecodePlan = plans.convDecodePlan; +pub const PrefillPlan = plans.PrefillPlan; +pub const prefillPlan = plans.prefillPlan; +pub const DecodePlan = plans.DecodePlan; +pub const decodePlan = plans.decodePlan; +pub const GatedNormPlan = plans.GatedNormPlan; +pub const gatedNormPlan = plans.gatedNormPlan; + +// Independent f64 reference for split/decode gates; host assertions only, never a tensor fallback. +pub const GateF64 = struct { g: f64, beta: f64 }; + +pub fn gateReferenceF64(a: f64, b: f64, a_log: f64, dt: f64) GateF64 { + const s = a + dt; + const sp = @max(s, 0.0) + std.math.log1p(std.math.exp(-@abs(s))); + return .{ + .g = @exp(-@exp(a_log) * sp), + .beta = 1.0 / (1.0 + @exp(-b)), + }; +} + +const testing = std.testing; + +test "splitVariant and gatedNormVariant accept only the narrow and wide combinations" { + const dev: *const Device = @ptrFromInt(0x5000); + const words: usize = 4096; + const buf = Buffer{ .dev = dev, .handle = @ptrFromInt(0x10), .memory = @ptrFromInt(0x20), .len = words * 4, .kind = .device, .mapped = null, .coherent = false }; + const bf = try View.init(&buf, .bf16, 0, words / 2); + const f3 = try View.init(&buf, .f32, 0, words / 4); + try testing.expectEqual(SplitVariant.bf16, try splitVariant(bf, bf)); + try testing.expectEqual(SplitVariant.f32, try splitVariant(f3, f3)); + try testing.expectError(error.Invalid, splitVariant(bf, f3)); // mixed activation kinds + try testing.expectError(error.Invalid, splitVariant(try View.init(&buf, .u16, 0, words / 2), bf)); + const gain = try View.init(&buf, .bf16, 256, 8); + try testing.expectEqual(GatedNormVariant.bf16, try gatedNormVariant(bf, bf, bf, gain)); + try testing.expectEqual(GatedNormVariant.f32, try gatedNormVariant(f3, f3, f3, gain)); + try testing.expectError(error.Invalid, gatedNormVariant(f3, f3, f3, f3)); // a f32 weight is never accepted + try testing.expectError(error.Invalid, gatedNormVariant(bf, f3, bf, gain)); + try testing.expectError(error.Invalid, gatedNormVariant(bf, bf, bf, try View.init(&buf, .u16, 512, 8))); +} + +test "wantSplitOutputs requires bf16 soa/beta and f32 g on real View kinds" { + // Fake buffers test view kinds/extents only, as in commands.zig; no device is needed. + const dev: *const Device = @ptrFromInt(0x5000); + const words: usize = 4096; + const buf = Buffer{ .dev = dev, .handle = @ptrFromInt(0x10), .memory = @ptrFromInt(0x20), .len = words * 4, .kind = .device, .mapped = null, .coherent = false }; + const soa = try View.init(&buf, .bf16, 0, words / 2); + const g = try View.init(&buf, .f32, 0, words / 4); + const beta = try View.init(&buf, .bf16, 256, words / 2 - 128); + try wantSplitOutputs(soa, g, beta); + // a bf16 g_out — the regression this guards — must refuse + const g_bf16 = try View.init(&buf, .bf16, 0, words / 2); + try testing.expectError(error.Invalid, wantSplitOutputs(soa, g_bf16, beta)); + // so does an f32 beta or a u16 soa (u16 has the same extent as bf16, so the refusal is genuinely the kind check) + const beta_f32 = try View.init(&buf, .f32, 0, words / 4); + try testing.expectError(error.Invalid, wantSplitOutputs(soa, g, beta_f32)); + const soa_u16 = try View.init(&buf, .u16, 0, words / 2); + try testing.expectError(error.Invalid, wantSplitOutputs(soa_u16, g, beta)); + // the wide layers: f32 soa takes an f32 beta (the shader stores floats) and refuses a bf16 one + const soa_f32 = try View.init(&buf, .f32, 0, words / 4); + try wantSplitOutputs(soa_f32, g, beta_f32); + try testing.expectError(error.Invalid, wantSplitOutputs(soa_f32, g, beta)); +} + +test "gateReferenceF64 behaves like exp(-exp(A_log) softplus) and a sigmoid" { + const g0 = gateReferenceF64(0, 0, 0, 0); + try testing.expectApproxEqAbs(@as(f64, @exp(-@exp(0.0) * std.math.log1p(@as(f64, 1.0)))), g0.g, 1e-12); + try testing.expectApproxEqAbs(@as(f64, 0.5), g0.beta, 1e-12); + // larger a (larger softplus) decays g; larger a_log decays it harder + const up = gateReferenceF64(2, 0, 0, 0); + try testing.expect(g0.g > up.g and up.g > 0 and up.g < 1); + const loud = gateReferenceF64(2, 0, 2, 0); + try testing.expect(loud.g < up.g); + try testing.expect(gateReferenceF64(0, 4, 0, 0).beta > g0.beta); + try testing.expect(gateReferenceF64(-40, 0, 0, 0).g > 0.999); // softplus saturates to 0 +} + +test "the fixed model constants agree with the Metal kernel constexprs" { + try testing.expectEqual(@as(usize, 16), Model.heads); + try testing.expectEqual(@as(usize, 128), Model.head_dim); + try testing.expectEqual(@as(usize, 6144), Model.conv_channels); + try testing.expectEqual(@as(usize, 4), Model.taps); + try testing.expectEqual(@as(usize, 18), Model.linear_layers); + try testing.expectEqual(@as(usize, 262144), Model.state_elements); +} + +test "single-chunk prefill read and written snapshot dummies are disjoint" { + var m: Gdn = undefined; + m.dummy = .{ .dev = undefined, .handle = null, .memory = null, .len = 2 * dummy_bytes, .kind = .device, .mapped = null, .coherent = false }; + const read = try m.dummyView(.bf16); + const snapshot = try m.snapshotDummyView(); + try testing.expect(read.buffer == snapshot.buffer); + try testing.expect(!read.overlaps(snapshot)); + try testing.expectEqual(@as(u64, 0), read.byte_offset); + try testing.expectEqual(dummy_bytes, snapshot.byte_offset); + try testing.expectEqual(dummy_bytes, read.byteLen()); + try testing.expectEqual(dummy_bytes, snapshot.byteLen()); +} + +test "gdn push flags match the masks the embedded shaders actually test" { + const src = @import("vulkan_comp_sources"); + // Decode uses bit3 for raw gates and bit9 for f32 A_log; 0x300 missed the shader's &8u test. + try testing.expectEqual(@as(u32, 8), decode_raw_gates); + try testing.expectEqual(@as(u32, 512), decode_f32_alog); + try testing.expect(std.mem.indexOf(u8, src.gated_delta_decode, "(params.flags & 8u)") != null); + try testing.expect(std.mem.indexOf(u8, src.gated_delta_decode, "(params.flags & 512u)") != null); + // Prefill: scan pass keeps g f32 (bit2); the output pass adds bit3. + try testing.expectEqual(@as(u32, 4), scan_flags); + try testing.expectEqual(@as(u32, 4 | 8), out_flags); + try testing.expect(std.mem.indexOf(u8, src.gated_delta_prefill_oracle, "(params.flags & 4u)") != null); + try testing.expect(std.mem.indexOf(u8, src.gated_delta_prefill_oracle, "(params.flags & 8u)") != null); + // Conv decode SiLU/norm masks exist but stay off; split owns the sole SiLU. + try testing.expectEqual(@as(u32, 0), conv_decode_flags); + try testing.expect(std.mem.indexOf(u8, src.gdn_conv_decode, "(params.flags & 1u)") != null); + try testing.expect(std.mem.indexOf(u8, src.gdn_conv_decode, "(params.flags & 2u)") != null); +} diff --git a/zig/src/vulkan/gdn_plan.zig b/zig/src/vulkan/gdn_plan.zig new file mode 100644 index 00000000000..111c99e125b --- /dev/null +++ b/zig/src/vulkan/gdn_plan.zig @@ -0,0 +1,197 @@ +// The GDN geometry plans: split, conv prefill and decode, the prefill recurrence, decode and the gated norm. +const std = @import("std"); +const dispatch = @import("dispatch.zig"); +const gdn = @import("gdn.zig"); + +const Model = gdn.Model; +const Error = dispatch.Error; + +/// The gate split's derived geometry, or the named reason it cannot run. +pub const SplitPlan = struct { tokens: u32, conv_channels: u32, section_elements: u32 }; + +pub fn splitPlan(conv_out_elements: u64, soa_elements: u64, g_elements: u64, beta_elements: u64, a_elements: u64, b_elements: u64, a_log_elements: u64, dt_elements: u64, tokens: u64) Error!SplitPlan { + if (tokens == 0) return error.Invalid; + const section = std.math.mul(u64, tokens, Model.linear_dim) catch return error.TooLarge; + const rows = std.math.mul(u64, tokens, Model.heads) catch return error.TooLarge; + const soa = std.math.mul(u64, section, 3) catch return error.TooLarge; + const conv = std.math.mul(u64, tokens, Model.conv_channels) catch return error.TooLarge; + if (conv_out_elements != conv or soa_elements != soa) return error.Invalid; + if (g_elements != rows or beta_elements != rows or a_elements != rows or b_elements != rows) return error.Invalid; + if (a_log_elements != Model.heads or dt_elements != Model.heads) return error.Invalid; + return .{ + .tokens = try cast32(tokens), + .conv_channels = try cast32(Model.conv_channels), + .section_elements = try cast32(section), + }; +} + +pub const ConvPrefillPlan = struct { tokens: u32, count: u32, l_in: u32, tokens_elements: u64, state_elements: u64 }; + +pub fn convPrefillPlan(conv_input_elements: u64, conv_state_in_elements: u64, qkv_elements: u64, weight_elements: u64, conv_state_out_elements: u64, tokens: u64) Error!ConvPrefillPlan { + if (tokens == 0) return error.Invalid; + const rows = std.math.mul(u64, tokens, Model.conv_channels) catch return error.TooLarge; + const input = std.math.add(u64, rows, std.math.mul(u64, Model.conv_state_rows, Model.conv_channels) catch return error.TooLarge) catch return error.TooLarge; + const state = std.math.mul(u64, Model.conv_state_rows, Model.conv_channels) catch return error.TooLarge; + if (conv_input_elements != input) return error.Invalid; + if (conv_state_in_elements != state or conv_state_out_elements != state) return error.Invalid; + if (qkv_elements != rows) return error.Invalid; + if (weight_elements != Model.conv_channels * Model.taps) return error.Invalid; + return .{ + .tokens = try cast32(tokens), + .count = try cast32(rows), + .l_in = try cast32(tokens + Model.conv_state_rows), + .tokens_elements = rows, + .state_elements = state, + }; +} + +pub const ConvDecodePlan = struct { count: u32, conv_channels: u32 }; + +pub fn convDecodePlan(conv_state_in_elements: u64, x_elements: u64, weight_elements: u64, activated_elements: u64, conv_state_out_elements: u64, soa_elements: u64, g_elements: u64, beta_elements: u64, a_elements: u64, b_elements: u64, a_log_elements: u64, dt_elements: u64) Error!ConvDecodePlan { + const state = std.math.mul(u64, Model.conv_state_rows, Model.conv_channels) catch return error.TooLarge; + const channels: u64 = Model.conv_channels; + if (conv_state_in_elements != state or conv_state_out_elements != state) return error.Invalid; + if (x_elements != channels or activated_elements != channels) return error.Invalid; + if (weight_elements != channels * Model.taps) return error.Invalid; + if (soa_elements != 3 * Model.linear_dim) return error.Invalid; + if (g_elements != Model.heads or beta_elements != Model.heads) return error.Invalid; + if (a_elements != Model.heads or b_elements != Model.heads) return error.Invalid; + if (a_log_elements != Model.heads or dt_elements != Model.heads) return error.Invalid; + return .{ .count = try cast32(channels), .conv_channels = try cast32(channels) }; +} + +pub const PrefillPlan = struct { tokens: u32, chunk: u32, chunks: u32, section_elements: u32 }; + +pub fn prefillPlan(soa_elements: u64, g_elements: u64, beta_elements: u64, h0_elements: u64, y_elements: u64, hf_elements: u64, snap_elements: u64, tokens: u64, chunk: u64) Error!PrefillPlan { + if (tokens == 0 or chunk == 0) return error.Invalid; + const section = std.math.mul(u64, tokens, Model.linear_dim) catch return error.TooLarge; + const soa = std.math.mul(u64, section, 3) catch return error.TooLarge; + const rows = std.math.mul(u64, tokens, Model.heads) catch return error.TooLarge; + if (soa_elements != soa or g_elements != rows or beta_elements != rows) return error.Invalid; + if (h0_elements != Model.state_elements or hf_elements != Model.state_elements) return error.Invalid; + if (y_elements != section) return error.Invalid; + const chunks = (tokens + chunk - 1) / chunk; + const snap = std.math.mul(u64, chunks - 1, Model.state_elements) catch return error.TooLarge; + if (snap_elements != snap) return error.Invalid; + return .{ + .tokens = try cast32(tokens), + .chunk = try cast32(chunk), + .chunks = try cast32(chunks), + .section_elements = try cast32(section), + }; +} + +pub const DecodePlan = struct {}; + +pub fn decodePlan(soa_elements: u64, a_elements: u64, b_elements: u64, a_log_u16_elements: u64, dt_elements: u64, state_in_elements: u64, state_out_elements: u64, y_elements: u64) Error!DecodePlan { + if (soa_elements != 3 * Model.linear_dim) return error.Invalid; + if (a_elements != Model.heads or b_elements != Model.heads or dt_elements != Model.heads) return error.Invalid; + // f32 A_log viewed as u16 pairs: heads * 2 + if (a_log_u16_elements != 2 * Model.heads) return error.Invalid; + if (state_in_elements != Model.state_elements or state_out_elements != Model.state_elements) return error.Invalid; + if (y_elements != Model.linear_dim) return error.Invalid; + return .{}; +} + +pub const GatedNormPlan = struct { rows: u32 }; + +// `rows` counts tokens; dispatch rows*16 width-128 head rows with per-head weight [128], not [2048]. +pub fn gatedNormPlan(out_elements: u64, y_elements: u64, z_elements: u64, weight_elements: u64, rows: u64, eps: f32) Error!GatedNormPlan { + if (!std.math.isFinite(eps) or eps <= 0) return error.Invalid; + if (rows == 0) return error.Invalid; + const total = std.math.mul(u64, rows, Model.linear_dim) catch return error.TooLarge; + if (out_elements != total or y_elements != total or z_elements != total) return error.Invalid; + if (weight_elements != Model.head_dim) return error.Invalid; + const head_rows = std.math.mul(u64, rows, Model.heads) catch return error.TooLarge; + return .{ .rows = try cast32(head_rows) }; +} + +fn cast32(v: u64) Error!u32 { + return std.math.cast(u32, v) orelse error.TooLarge; +} + +const testing = std.testing; + +test "splitPlan derives the geometry and refuses wrong shapes and overflow" { + const t: u64 = 7; + const section = t * Model.linear_dim; + try testing.expectEqual(@as(u32, 7), (try splitPlan(t * Model.conv_channels, 3 * section, t * Model.heads, t * Model.heads, t * Model.heads, t * Model.heads, Model.heads, Model.heads, t)).tokens); + try testing.expectError(error.Invalid, splitPlan(t * Model.conv_channels + 1, 3 * section, t * Model.heads, t * Model.heads, t * Model.heads, t * Model.heads, Model.heads, Model.heads, t)); + try testing.expectError(error.Invalid, splitPlan(t * Model.conv_channels, 3 * section, t * Model.heads, t * Model.heads, t * Model.heads, t * Model.heads, Model.heads + 1, Model.heads, t)); + try testing.expectError(error.Invalid, splitPlan(t * Model.conv_channels, 3 * section, t * Model.heads, t * Model.heads, t * Model.heads, t * Model.heads, Model.heads, Model.heads, 0)); + // tokens=2^44 yields section=2^55; u64 geometry fits, but u32 casts return TooLarge. + const big: u64 = 1 << 44; + const big_section = big * Model.linear_dim; + try testing.expectError(error.TooLarge, splitPlan(big * Model.conv_channels, 3 * big_section, big * Model.heads, big * Model.heads, big * Model.heads, big * Model.heads, Model.heads, Model.heads, big)); +} + +test "convPrefillPlan checks the conv input, state and weight geometry" { + const t: u64 = 5; + const ch: u64 = Model.conv_channels; + const state = Model.conv_state_rows * ch; + try testing.expectEqual(@as(u32, 8), (try convPrefillPlan((t + 3) * ch, state, t * ch, ch * 4, state, t)).l_in); + try testing.expectError(error.Invalid, convPrefillPlan((t + 2) * ch, state, t * ch, ch * 4, state, t)); + try testing.expectError(error.Invalid, convPrefillPlan((t + 3) * ch, state, t * ch, ch * 5, state, t)); + try testing.expectError(error.Invalid, convPrefillPlan((t + 3) * ch, state + ch, t * ch, ch * 4, state, t)); +} + +test "convDecodePlan refuses any torn decode geometry" { + const ch: u64 = Model.conv_channels; + const state = Model.conv_state_rows * ch; + const heads: u64 = Model.heads; + const ok = try convDecodePlan(state, ch, ch * 4, ch, state, 3 * Model.linear_dim, heads, heads, heads, heads, heads, heads); + try testing.expect(ok.count == ch); + try testing.expectError(error.Invalid, convDecodePlan(state, ch + 2, ch * 4, ch, state, 3 * Model.linear_dim, heads, heads, heads, heads, heads, heads)); + try testing.expectError(error.Invalid, convDecodePlan(state, ch, ch * 4, ch, state, 3 * Model.linear_dim, heads, heads, heads, heads, heads, heads + 1)); +} + +test "prefillPlan derives chunks and the snapshot size, refusing snap mismatch" { + const t: u64 = 130; + const chunk: u64 = 64; + const section = t * Model.linear_dim; + const state: u64 = Model.state_elements; + // Three chunks need two interior snapshots; chunk c restores slot c-1, so size is (chunks-1)*state. + const plan = try prefillPlan(3 * section, t * Model.heads, t * Model.heads, state, section, state, 2 * state, t, chunk); + try testing.expectEqual(@as(u32, 3), plan.chunks); + try testing.expectEqual(@as(u32, 64), plan.chunk); + try testing.expectError(error.Invalid, prefillPlan(3 * section, t * Model.heads, t * Model.heads, state, section, state, 3 * state, t, chunk)); // too many slots + try testing.expectError(error.Invalid, prefillPlan(3 * section, t * Model.heads, t * Model.heads, state, section, state, 1 * state, t, chunk)); // too few + try testing.expectError(error.Invalid, prefillPlan(3 * section, t * Model.heads, t * Model.heads, state, section, state, 0, t, chunk)); // chunks>1 needs snap + try testing.expectError(error.Invalid, prefillPlan(3 * section, t * Model.heads, t * Model.heads, state, section, state, 0, t, 0)); + const single = try prefillPlan(3 * section, t * Model.heads, t * Model.heads, state, section, state, 0, t, t); + try testing.expectEqual(@as(u32, 1), single.chunks); +} + +test "the common prefill graph admits decode and full prompt chunks" { + const ch: u64 = Model.conv_channels; + const conv_state = Model.conv_state_rows * ch; + const state: u64 = Model.state_elements; + for ([_]u64{ 1, 31, 32 }) |tokens| { + const conv = try convPrefillPlan((tokens + Model.conv_state_rows) * ch, conv_state, tokens * ch, ch * Model.taps, conv_state, tokens); + try testing.expectEqual(@as(u32, @intCast(tokens)), conv.tokens); + try testing.expectEqual(@as(u32, @intCast(tokens + Model.conv_state_rows)), conv.l_in); + const section = tokens * Model.linear_dim; + const recurrence = try prefillPlan(3 * section, tokens * Model.heads, tokens * Model.heads, state, section, state, 0, tokens, tokens); + try testing.expectEqual(@as(u32, 1), recurrence.chunks); + try testing.expectEqual(@as(u32, @intCast(tokens)), recurrence.chunk); + } +} + +test "decodePlan checks the u16 A_log window and the fixed decode shapes" { + const heads: u64 = Model.heads; + _ = try decodePlan(3 * Model.linear_dim, heads, heads, 2 * heads, heads, Model.state_elements, Model.state_elements, Model.linear_dim); + try testing.expectError(error.Invalid, decodePlan(3 * Model.linear_dim, heads, heads, heads, heads, Model.state_elements, Model.state_elements, Model.linear_dim)); + try testing.expectError(error.Invalid, decodePlan(3 * Model.linear_dim, heads, heads, 2 * heads, heads, Model.state_elements + 4, Model.state_elements, Model.linear_dim)); +} + +test "gatedNormPlan is the per-head 128 norm and refuses torn geometry" { + const rows: u64 = 3; + const total = rows * Model.linear_dim; + // rows = tokens; the plan dispatches tokens*16 head-rows, weight [128] + try testing.expectEqual(@as(u32, 3 * Model.heads), (try gatedNormPlan(total, total, total, Model.head_dim, rows, 1e-6)).rows); + // the old 2048 global-row weight is REFUSED + try testing.expectError(error.Invalid, gatedNormPlan(total, total, total, Model.linear_dim, rows, 1e-6)); + try testing.expectError(error.Invalid, gatedNormPlan(total + 2, total, total, Model.head_dim, rows, 1e-6)); + try testing.expectError(error.Invalid, gatedNormPlan(total, total, total, Model.head_dim, rows, 0)); + try testing.expectError(error.Invalid, gatedNormPlan(total, total, total, Model.head_dim, 0, 1e-6)); +} diff --git a/zig/src/vulkan/greedy.zig b/zig/src/vulkan/greedy.zig new file mode 100644 index 00000000000..89ef47e4e8d --- /dev/null +++ b/zig/src/vulkan/greedy.zig @@ -0,0 +1,208 @@ +// Greedy: max finite, +0/-0 ties to lowest index, any NaN/Inf refuses; argreduce = mlx-omarchy copies (max, one row) +const std = @import("std"); +const c = @import("vk"); +const commands = @import("commands.zig"); +const dispatch = @import("dispatch.zig"); +const memory = @import("memory.zig"); +const params = @import("stage1_params"); +const tensor = @import("tensor.zig"); +const Device = @import("device.zig").Device; +const Commands = commands.Commands; +const Buffer = memory.Buffer; +const View = tensor.View; + +pub const Error = dispatch.Error || error{NonFinite}; + +// finite_logits scans; argreduce_chunk and argreduce_combine reduce; chunk/combine use -DUSE_BF16; scan has no defines. +pub const Shaders = struct { + finite_logits: []align(4) const u8, + argreduce_chunk: []align(4) const u8, + argreduce_combine: []align(4) const u8, +}; + +// Each workgroup covers 256 lanes x 16 prefetched values, matching the argreduce chunk kernel. +const block_elements: u64 = 4096; +/// The combine kernel folds at most 256 pairs in its single workgroup. +const max_chunks: u64 = 256; + +// 12-byte readback stores verdict, first bad index, token; sentinels prevent unwritten indices passing as token 0. +pub const Result = extern struct { + finite_flag: u32, + first_bad: u32, + index: u32, +}; + +comptime { + if (@sizeOf(Result) != 12) @compileError("greedy readback drifted"); + if (@offsetOf(params.Params, "flags") + 4 != 60) @compileError("argreduce push prefix drifted"); +} + +/// The argreduce kernels' push block ends at `flags`: the standard 15-word prefix. +const greedy_push_bytes: u32 = 60; + +/// finite_logits.comp reads only the row length. +const FiniteParams = extern struct { count: u32 = 0 }; + +// Pure plan: geometry only (Invalid empty, TooLarge huge); odd bf16 counts fail View word-padding at the caller +pub const Plan = struct { count: u32, chunk_len: u32, chunks: u32 }; + +pub fn greedyPlan(count: u64) Error!Plan { + if (count == 0) return error.Invalid; + const n = std.math.cast(u32, count) orelse return error.TooLarge; + const chunks = std.math.cast(u32, (count + block_elements - 1) / block_elements) orelse return error.TooLarge; + if (chunks > max_chunks) return error.TooLarge; // combine folds one row in a single workgroup + return .{ .count = n, .chunk_len = @intCast(block_elements), .chunks = chunks }; +} + +// Decode returns NonFinite on dirty flags and VulkanFailed on missing/bad indices; otherwise it returns the token. +pub fn decode(count: u32, out: Result) Error!u32 { + if (out.finite_flag != 0) return error.NonFinite; + if (out.index >= count) return error.VulkanFailed; + return out.index; +} + +// The host widens bf16 exactly, rejects non-finites, and keeps the first max; ascending order resolves +0/-0 ties. +pub fn greedyRef(logits: []const f32) Error!u32 { + if (logits.len == 0) return error.Invalid; + var best = -std.math.inf(f32); + var index: u32 = 0; + for (logits, 0..) |v, i| { + if (!std.math.isFinite(v)) return error.NonFinite; + if (v > best) { + best = v; + index = std.math.cast(u32, i) orelse return error.TooLarge; + } + } + return index; +} + +pub const Greedy = struct { + dev: *const Device, + plan: Plan, + finite_logits: dispatch.Pipeline, + chunk: dispatch.Pipeline, + combine: dispatch.Pipeline, + /// Per-row (value, index) pairs the chunk kernel writes, one per chunk. + partials: Buffer, + /// The three result words: finite flag, first bad index, selected token. + results: Buffer, + + // Builds pipelines and scratch; the device must support storage16 and int16 for bf16. + pub fn init(dev: *const Device, shaders: Shaders, count: u64) Error!Greedy { + if (!dev.features.storage16 or !dev.features.int16) return error.MissingFeature; + const plan = try greedyPlan(count); + var partials = try Buffer.init(dev, .device, @as(u64, plan.chunks) * 8); + errdefer partials.deinit(); + var results = try Buffer.init(dev, .device, @sizeOf(Result)); + errdefer results.deinit(); + var g: Greedy = .{ + .dev = dev, + .plan = plan, + .finite_logits = undefined, + .chunk = undefined, + .combine = undefined, + .partials = partials, + .results = results, + }; + g.finite_logits = try dispatch.Pipeline.init(dev, shaders.finite_logits, .{ .bindings = 2, .push_bytes = @sizeOf(FiniteParams), .writes = 1 << 1 }); + errdefer g.finite_logits.deinit(); + g.chunk = try dispatch.Pipeline.init(dev, shaders.argreduce_chunk, .{ .bindings = 2, .push_bytes = greedy_push_bytes, .writes = 1 << 1 }); + errdefer g.chunk.deinit(); + g.combine = try dispatch.Pipeline.init(dev, shaders.argreduce_combine, .{ .bindings = 2, .push_bytes = greedy_push_bytes, .writes = 1 << 1 }); + return g; + } + + /// Destroys the pipelines and the scratch; call only after the last submission using them returned success. + pub fn deinit(g: *Greedy) void { + g.combine.deinit(); + g.chunk.deinit(); + g.finite_logits.deinit(); + g.partials.deinit(); + g.results.deinit(); + } + + // One draw resets, scans, reduces, combines, then fences readback; wrap in commands.orExit for Timeout/InFlight. + pub fn select(g: *Greedy, cmds: *Commands, staging: Buffer, logits: View, timeout_ns: u64) Error!u32 { + if (logits.kind != .bf16 or logits.elements != g.plan.count) return error.Invalid; + try g.finite_logits.bind(&.{ logits.bind(), .{ .buffer = &g.results, .len = 8 } }); + try g.chunk.bind(&.{ logits.bind(), .{ .buffer = &g.partials, .len = g.partials.len } }); + try g.combine.bind(&.{ .{ .buffer = &g.partials, .len = g.partials.len }, .{ .buffer = &g.results, .len = g.results.len } }); + const cmd = try cmds.begin(); + // Each call resets flag=0 and bad/index=0xffffffff; a barrier orders fills before compute. + try commands.recordFill(g.dev, cmd, g.results, 0, 4, 0); + try commands.recordFill(g.dev, cmd, g.results, 4, 8, std.math.maxInt(u32)); + commands.barrier(g.dev, cmd, c.VK_PIPELINE_STAGE_TRANSFER_BIT, c.VK_ACCESS_TRANSFER_WRITE_BIT, fill_stages, fill_access); + var fp: FiniteParams = .{ .count = g.plan.count }; + try g.finite_logits.record(cmd, std.mem.asBytes(&fp), .{ g.plan.chunks, 1, 1 }); + var p: params.Params = .{ + .operation = 1, // max + .reduce_size = g.plan.count, + .matrix_m = g.plan.chunk_len, + .matrix_n = g.plan.chunks, + }; + try g.chunk.record(cmd, std.mem.asBytes(&p)[0..greedy_push_bytes], .{ g.plan.chunks, 1, 1 }); + var o: params.Params = .{ + .operation = 1, // max + .matrix_n = g.plan.chunks, + .output_offset = 2, // Result.index: word 2, after flag and first_bad + }; + try g.combine.record(cmd, std.mem.asBytes(&o)[0..greedy_push_bytes], .{ 1, 1, 1 }); + try cmds.submit(timeout_ns); + var out: Result = undefined; + try commands.download(cmds, g.results, 0, staging, std.mem.asBytes(&out), timeout_ns); + return decode(g.plan.count, out) catch |e| { + if (e == error.NonFinite) std.log.err("greedy: non-finite bf16 logit at index {d} of {d}; row refused", .{ out.first_bad, g.plan.count }); + return e; + }; + } +}; + +const fill_stages: u32 = c.VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | c.VK_PIPELINE_STAGE_TRANSFER_BIT; +const fill_access: u32 = c.VK_ACCESS_SHADER_READ_BIT | c.VK_ACCESS_SHADER_WRITE_BIT | + c.VK_ACCESS_TRANSFER_READ_BIT | c.VK_ACCESS_TRANSFER_WRITE_BIT; + +/// f32 exactly as the kernels widen a bf16 halfword. +fn widen(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +const testing = std.testing; + +test "greedyPlan derives 4096-wide chunks and refuses empty, oversized and over-partitioned rows" { + const plan = try greedyPlan(248320); + try testing.expectEqual(@as(u32, 248320), plan.count); + try testing.expectEqual(@as(u32, 4096), plan.chunk_len); + try testing.expectEqual(@as(u32, 61), plan.chunks); + try testing.expectEqual(@as(u32, 2), (try greedyPlan(4097)).chunks); + try testing.expectError(error.Invalid, greedyPlan(0)); + try testing.expectError(error.TooLarge, greedyPlan(1 << 32)); + try testing.expectError(error.TooLarge, greedyPlan(max_chunks * block_elements + 1)); +} + +test "greedyRef applies the draw rule: max wins, exact ties keep the lowest index, any non-finite refuses" { + try testing.expectEqual(@as(u32, 2), try greedyRef(&.{ -1.5, -1.5, -0.5, -2 })); // all negative + try testing.expectEqual(@as(u32, 0), try greedyRef(&.{ 0.5, 0.5 })); + try testing.expectEqual(@as(u32, 0), try greedyRef(&.{ 0.0, -0.0, -1.0 })); // +0/-0 compare equal: first of the tie + try testing.expectEqual(@as(u32, 0), try greedyRef(&.{ -0.0, 0.0, -1.0 })); // whichever sign comes first + try testing.expectEqual(@as(u32, 1), try greedyRef(&.{ widen(0x3f80), widen(0x4000), widen(0x4000), widen(0x4000) })); + try testing.expectError(error.NonFinite, greedyRef(&.{ 1.0, std.math.nan(f32) })); + try testing.expectError(error.NonFinite, greedyRef(&.{ 1.0, std.math.inf(f32) })); + try testing.expectError(error.NonFinite, greedyRef(&.{ -std.math.inf(f32), -1.0 })); // losing -inf still refuses + try testing.expectError(error.Invalid, greedyRef(&.{})); +} + +test "bf16 logits ride word-padded views: an odd count has no view and is refused before any dispatch" { + const buf: Buffer = .{ .dev = undefined, .handle = null, .memory = null, .len = 248322 * 2, .kind = .device, .mapped = null, .coherent = false }; + try testing.expectError(error.Misaligned, View.init(&buf, .bf16, 0, 1)); + // Both rows fit the fake buffer; the odd row fails View word-padding rather than exceeding the buffer. + try testing.expectError(error.Misaligned, View.init(&buf, .bf16, 0, 248321)); + _ = try View.init(&buf, .bf16, 0, 248320); +} + +test "decode refuses a dirty flag and an unwritten or out-of-row index and passes a real token" { + const clean: Result = .{ .finite_flag = 0, .first_bad = std.math.maxInt(u32), .index = 262 }; + try testing.expectEqual(@as(u32, 262), try decode(248320, clean)); + try testing.expectError(error.NonFinite, decode(248320, .{ .finite_flag = 1, .first_bad = 7, .index = 0 })); + try testing.expectError(error.VulkanFailed, decode(248320, .{ .finite_flag = 0, .first_bad = std.math.maxInt(u32), .index = std.math.maxInt(u32) })); + try testing.expectError(error.VulkanFailed, decode(248320, .{ .finite_flag = 0, .first_bad = std.math.maxInt(u32), .index = 248320 })); +} diff --git a/zig/src/vulkan/norm_mlp.zig b/zig/src/vulkan/norm_mlp.zig new file mode 100644 index 00000000000..a5c9c2207b4 --- /dev/null +++ b/zig/src/vulkan/norm_mlp.zig @@ -0,0 +1,309 @@ +// Norm/SwiGLU build with -DUSE_BF16; binary_vec is bf16 by default (no defines); runtime has no CPU fallback. +const std = @import("std"); +const dispatch = @import("dispatch.zig"); +const memory = @import("memory.zig"); +const params = @import("stage1_params"); +const tensor = @import("tensor.zig"); +const Device = @import("device.zig").Device; +const Commands = @import("commands.zig").Commands; +const Buffer = memory.Buffer; +const Error = dispatch.Error; +const View = tensor.View; + +// norm/SwiGLU: -DUSE_BF16; -DORACLE_QWEN35=Ash8c (build-time); -DPRESERVE_F32 needs shaderDenormPreserveFloat32. +pub const Shaders = struct { + norm_bf16: []align(4) const u8, + add_bf16: []align(4) const u8, + add_preserve_bf16: []align(4) const u8, + swiglu_bf16: []align(4) const u8, + norm_f32: []align(4) const u8, + norm_f32_in: []align(4) const u8, + norm_bf16_in: []align(4) const u8, + add_f32: []align(4) const u8, + add_wide_bf16: []align(4) const u8, + add_bf16_f32: []align(4) const u8, + swiglu_f32: []align(4) const u8, +}; + +const dummy_bytes: u64 = 256; + +/// fast_norm.comp's push block ends at `beta`: 17 words. +const norm_push_bytes: u32 = 68; +/// binary_vec.comp reads only the first nine words, the cast prefix. +const add_push_bytes: u32 = 36; +/// swiglu.comp has its own four-word block. +const SwigluParams = extern struct { count: u32 = 0, gate_offset: u32 = 0, up_offset: u32 = 0, materialize: u32 = 0 }; +const swiglu_push_bytes: u32 = @sizeOf(SwigluParams); + +comptime { + if (norm_push_bytes != @offsetOf(params.Params, "beta") + 4) @compileError("fast_norm push block drifted from Params"); + if (add_push_bytes != params.cast_push_bytes) @compileError("binary_vec push block drifted from the cast prefix"); +} + +fn groups(count: u64, local: u32, max_x: u32) Error![3]u32 { + const n = std.math.cast(u32, (count + local - 1) / local) orelse return error.TooLarge; + return .{ @min(n, max_x), 1, 1 }; +} + +/// The rmsNorm kind combinations: gains stay bf16 in all four. +pub const NormVariant = enum { bf16, f32, f32_in, bf16_in }; + +/// (out, x, gamma) kinds -> norm pipeline, or Invalid; pure and host-tested. +pub fn normVariant(out: View, x: View, gamma: View) Error!NormVariant { + if (gamma.kind != .bf16) return error.Invalid; + if (x.kind == .bf16 and out.kind == .bf16) return .bf16; + if (x.kind == .f32 and out.kind == .f32) return .f32; + if (x.kind == .f32 and out.kind == .bf16) return .f32_in; + if (x.kind == .bf16 and out.kind == .f32) return .bf16_in; + return error.Invalid; +} + +/// The residual add kind combinations: bf16 h+r, f32 h+r, and the two mixed transitions. +pub const AddVariant = enum { bf16, f32, wide_bf16, wide_f32 }; + +/// (out, a, b) kinds -> add pipeline, or Invalid; pure and host-tested. +pub fn addVariant(out: View, a: View, b: View) Error!AddVariant { + if (out.elements != a.elements or a.elements != b.elements) return error.Invalid; + if (a.kind == .bf16 and b.kind == .bf16 and out.kind == .bf16) return .bf16; + if (a.kind == .f32 and b.kind == .f32 and out.kind == .f32) return .f32; + if (a.kind == .f32 and b.kind == .bf16 and out.kind == .bf16) return .wide_bf16; + if (a.kind == .bf16 and b.kind == .f32 and out.kind == .f32) return .wide_f32; + return error.Invalid; +} + +/// SwiGLU takes all-bf16 or all-f32 operands; mixed kinds refuse. +pub fn swigluVariant(out: View, gate: View, up: View) Error!tensor.Kind { + if (out.elements != gate.elements or gate.elements != up.elements) return error.Invalid; + if (out.kind != gate.kind or gate.kind != up.kind) return error.Invalid; + if (out.kind != .bf16 and out.kind != .f32) return error.Invalid; + return out.kind; +} + +pub const NormMlp = struct { + dev: *const Device, + norm_bf16: dispatch.Pipeline, + add_bf16: dispatch.Pipeline, + swiglu_bf16: dispatch.Pipeline, + norm_f32: dispatch.Pipeline, + norm_f32_in: dispatch.Pipeline, + norm_bf16_in: dispatch.Pipeline, + add_f32: dispatch.Pipeline, + add_wide_bf16: dispatch.Pipeline, + add_wide_f32: dispatch.Pipeline, + swiglu_f32: dispatch.Pipeline, + // Dummy bindings cover unused fast_norm bias 2 and unmaterialized SwiGLU sigmoid/silu 3,4. + dummy: Buffer, + + /// Builds the nine pipelines; the device must have storage16 and int16 for the bf16 paths. + pub fn init(dev: *const Device, shaders: Shaders) Error!NormMlp { + if (!dev.features.storage16 or !dev.features.int16) return error.MissingFeature; + var m: NormMlp = .{ + .dev = dev, + .norm_bf16 = undefined, + .add_bf16 = undefined, + .swiglu_bf16 = undefined, + .norm_f32 = undefined, + .norm_f32_in = undefined, + .norm_bf16_in = undefined, + .add_f32 = undefined, + .add_wide_bf16 = undefined, + .add_wide_f32 = undefined, + .swiglu_f32 = undefined, + .dummy = try Buffer.init(dev, .device, dummy_bytes), + }; + errdefer m.dummy.deinit(); + m.norm_bf16 = try dispatch.Pipeline.init(dev, shaders.norm_bf16, .{ .bindings = 4, .push_bytes = norm_push_bytes, .writes = 1 << 3 }); + errdefer m.norm_bf16.deinit(); + m.add_bf16 = try dispatch.Pipeline.init(dev, if (dev.v12.shaderDenormPreserveFloat32 != 0) shaders.add_preserve_bf16 else shaders.add_bf16, .{ .bindings = 3, .push_bytes = add_push_bytes, .writes = 1 << 2 }); + errdefer m.add_bf16.deinit(); + m.swiglu_bf16 = try dispatch.Pipeline.init(dev, shaders.swiglu_bf16, .{ .bindings = 5, .push_bytes = swiglu_push_bytes, .writes = 1 << 2 }); + errdefer m.swiglu_bf16.deinit(); + m.norm_f32 = try dispatch.Pipeline.init(dev, shaders.norm_f32, .{ .bindings = 4, .push_bytes = norm_push_bytes, .writes = 1 << 3 }); + errdefer m.norm_f32.deinit(); + m.norm_f32_in = try dispatch.Pipeline.init(dev, shaders.norm_f32_in, .{ .bindings = 4, .push_bytes = norm_push_bytes, .writes = 1 << 3 }); + errdefer m.norm_f32_in.deinit(); + m.norm_bf16_in = try dispatch.Pipeline.init(dev, shaders.norm_bf16_in, .{ .bindings = 4, .push_bytes = norm_push_bytes, .writes = 1 << 3 }); + errdefer m.norm_bf16_in.deinit(); + m.add_f32 = try dispatch.Pipeline.init(dev, shaders.add_f32, .{ .bindings = 3, .push_bytes = add_push_bytes, .writes = 1 << 2 }); + errdefer m.add_f32.deinit(); + m.add_wide_bf16 = try dispatch.Pipeline.init(dev, shaders.add_wide_bf16, .{ .bindings = 3, .push_bytes = add_push_bytes, .writes = 1 << 2 }); + errdefer m.add_wide_bf16.deinit(); + m.add_wide_f32 = try dispatch.Pipeline.init(dev, shaders.add_bf16_f32, .{ .bindings = 3, .push_bytes = add_push_bytes, .writes = 1 << 2 }); + errdefer m.add_wide_f32.deinit(); + m.swiglu_f32 = try dispatch.Pipeline.init(dev, shaders.swiglu_f32, .{ .bindings = 5, .push_bytes = swiglu_push_bytes, .writes = 1 << 2 }); + return m; + } + + /// Destroys the pipelines and the dummy; call only after the last submission using them returned success. + pub fn deinit(m: *NormMlp) void { + m.norm_bf16.deinit(); + m.add_bf16.deinit(); + m.swiglu_bf16.deinit(); + m.norm_f32.deinit(); + m.norm_f32_in.deinit(); + m.norm_bf16_in.deinit(); + m.add_f32.deinit(); + m.add_wide_bf16.deinit(); + m.add_wide_f32.deinit(); + m.swiglu_f32.deinit(); + m.dummy.deinit(); + } + + fn dummyBinding(m: *const NormMlp) dispatch.Binding { + return .{ .buffer = &m.dummy, .len = dummy_bytes }; + } + + fn normPipeline(m: *NormMlp, variant: NormVariant) *dispatch.Pipeline { + return switch (variant) { + .bf16 => &m.norm_bf16, + .f32 => &m.norm_f32, + .f32_in => &m.norm_f32_in, + .bf16_in => &m.norm_bf16_in, + }; + } + + /// Preserve the omitted contract, precision, and host-interface facts from the original comment. + pub fn rmsNorm(m: *NormMlp, cmds: *Commands, out: View, x: View, gamma: View, rows: u64, width: u64, eps: f32, timeout_ns: u64) Error!void { + const variant = try normVariant(out, x, gamma); + const plan = try normPlan(x.elements, out.elements, gamma.elements, rows, width, eps); + // grid-stride row loop: a group count past the device limit still covers every row + const g = try groups(plan.rows, 1, m.dev.properties.limits.maxComputeWorkGroupCount[0]); + const pipe = m.normPipeline(variant); + try pipe.bind(&.{ x.bind(), gamma.bind(), m.dummyBinding(), out.bind() }); + var p: params.Params = .{ + .count = plan.count, + .lhs_size = plan.width, // not 1: the row-gamma read, per column + .reduce_size = plan.width, + .output_size = plan.rows, + }; + p.alpha = eps; + try pipe.run(cmds, std.mem.asBytes(&p)[0..norm_push_bytes], g, timeout_ns); + } + + pub fn addPipeline(m: *NormMlp, variant: AddVariant) *dispatch.Pipeline { + return switch (variant) { + .bf16 => &m.add_bf16, + .f32 => &m.add_f32, + .wide_bf16 => &m.add_wide_bf16, + .wide_f32 => &m.add_wide_f32, + }; + } + + // Add is out-of-place; binary_vec loads 4 elements per run, so count and operands must be divisible by 4. + pub fn add(m: *NormMlp, cmds: *Commands, out: View, a: View, b: View, timeout_ns: u64) Error!void { + const variant = try addVariant(out, a, b); + try vec4Guard(out.elements); + const pipe = m.addPipeline(variant); + try pipe.bind(&.{ a.bind(), b.bind(), out.bind() }); + var p: params.Params = .{ + .count = try out.elemCount32(), + .operation = 0, // add + .lhs_size = try a.elemCount32(), + .rhs_size = try b.elemCount32(), + }; + try pipe.run(cmds, std.mem.asBytes(&p)[0..add_push_bytes], try groups(p.count, pipe.local[0], m.dev.properties.limits.maxComputeWorkGroupCount[0]), timeout_ns); + } + + // SwiGLU is 4-wide, out-of-place; bindings 3/4 materialize=0; ORACLE_QWEN35 rounds once, MLX fused_chain 3 times. + pub fn swiglu(m: *NormMlp, cmds: *Commands, out: View, gate: View, up: View, timeout_ns: u64) Error!void { + const kind = try swigluVariant(out, gate, up); + try vec4Guard(out.elements); + const pipe = if (kind == .f32) &m.swiglu_f32 else &m.swiglu_bf16; + try pipe.bind(&.{ gate.bind(), up.bind(), out.bind(), m.dummyBinding(), m.dummyBinding() }); + const p: SwigluParams = .{ .count = try out.elemCount32() }; + try pipe.run(cmds, std.mem.asBytes(&p), try groups(p.count, pipe.local[0], m.dev.properties.limits.maxComputeWorkGroupCount[0]), timeout_ns); + } +}; + +fn cast32(v: u64) Error!u32 { + return std.math.cast(u32, v) orelse error.TooLarge; +} + +fn vec4Guard(elements: u64) Error!void { + if (elements % 4 != 0) return error.Invalid; +} + +/// The RMSNorm's derived geometry, or the named reason it cannot run; pure and host-tested. +pub const NormPlan = struct { rows: u32, width: u32, count: u32 }; + +pub fn normPlan(x_elements: u64, out_elements: u64, gamma_elements: u64, rows: u64, width: u64, eps: f32) Error!NormPlan { + if (!std.math.isFinite(eps) or eps <= 0) return error.Invalid; + if (rows == 0 or width == 0) return error.Invalid; + const total = std.math.mul(u64, rows, width) catch return error.TooLarge; + if (x_elements != total or out_elements != total) return error.Invalid; + // A one-element bf16 gamma cannot pass View.init's word-padding rule, so no scalar form. + if (gamma_elements != width) return error.Invalid; + return .{ + .rows = try cast32(rows), + .width = try cast32(width), + .count = try cast32(total), + }; +} + +const testing = std.testing; + +fn fakeBuffer(len: u64) Buffer { + return .{ .dev = undefined, .handle = null, .memory = null, .len = len, .kind = .device, .mapped = null, .coherent = false }; +} + +fn view(buf: *const Buffer, kind: tensor.Kind, offset: u64, elements: u64) View { + return View.init(buf, kind, offset, elements) catch unreachable; +} + +test "normVariant maps the four combinations and refuses mixed or foreign gains" { + const buf = fakeBuffer(1 << 16); + const bf = view(&buf, .bf16, 0, 64); + const f3 = view(&buf, .f32, 256, 64); + const gain = view(&buf, .bf16, 1024, 8); + try testing.expectEqual(NormVariant.bf16, try normVariant(bf, bf, gain)); + try testing.expectEqual(NormVariant.f32, try normVariant(f3, f3, gain)); + try testing.expectEqual(NormVariant.f32_in, try normVariant(bf, f3, gain)); + try testing.expectEqual(NormVariant.bf16_in, try normVariant(f3, bf, gain)); + try testing.expectError(error.Invalid, normVariant(bf, bf, f3)); // a f32 gain is never accepted + try testing.expectError(error.Invalid, normVariant(f3, view(&buf, .u16, 0, 64), gain)); +} + +test "addVariant maps the four combinations and refuses mixed kinds" { + const buf = fakeBuffer(1 << 16); + const bf = view(&buf, .bf16, 0, 64); + const f3 = view(&buf, .f32, 256, 64); + try testing.expectEqual(AddVariant.bf16, try addVariant(bf, bf, bf)); + try testing.expectEqual(AddVariant.f32, try addVariant(f3, f3, f3)); + try testing.expectEqual(AddVariant.wide_bf16, try addVariant(bf, f3, bf)); + try testing.expectEqual(AddVariant.wide_f32, try addVariant(f3, bf, f3)); + try testing.expectError(error.Invalid, addVariant(f3, bf, bf)); // bf16 stream under an f32 output + try testing.expectError(error.Invalid, addVariant(bf, bf, f3)); + try testing.expectError(error.Invalid, addVariant(bf, f3, f3)); + try testing.expectError(error.Invalid, addVariant(bf, bf, view(&buf, .bf16, 512, 32))); // torn shapes +} + +test "swigluVariant takes all-bf16 or all-f32 and refuses the rest" { + const buf = fakeBuffer(1 << 16); + const bf = view(&buf, .bf16, 0, 64); + const f3 = view(&buf, .f32, 256, 64); + try testing.expectEqual(tensor.Kind.bf16, try swigluVariant(bf, bf, bf)); + try testing.expectEqual(tensor.Kind.f32, try swigluVariant(f3, f3, f3)); + try testing.expectError(error.Invalid, swigluVariant(f3, bf, bf)); + try testing.expectError(error.Invalid, swigluVariant(bf, f3, f3)); + try testing.expectError(error.Invalid, swigluVariant(bf, bf, f3)); + try testing.expectError(error.Invalid, swigluVariant(view(&buf, .u16, 0, 64), bf, bf)); +} + +test "normPlan derives the geometry and refuses torn rows, wrong gamma, bad eps and overflow" { + const plan = try normPlan(4 * 8, 4 * 8, 8, 4, 8, 1e-6); + try testing.expectEqual(@as(u32, 4), plan.rows); + try testing.expectEqual(@as(u32, 8), plan.width); + try testing.expectEqual(@as(u32, 32), plan.count); + try testing.expectError(error.Invalid, normPlan(4 * 8, 4 * 8 + 1, 8, 4, 8, 1e-6)); + try testing.expectError(error.Invalid, normPlan(4 * 8, 4 * 8, 7, 4, 8, 1e-6)); + try testing.expectError(error.Invalid, normPlan(4 * 8, 4 * 8, 8, 4, 8, 0)); + try testing.expectError(error.Invalid, normPlan(4 * 8, 4 * 8, 8, 4, 8, std.math.inf(f32))); + try testing.expectError(error.Invalid, normPlan(0, 0, 8, 0, 8, 1e-6)); + try testing.expectError(error.TooLarge, normPlan(1 << 32, 1 << 32, 1 << 32, 1, 1 << 32, 1e-6)); // width past u32 + try testing.expectError(error.TooLarge, normPlan(1 << 32, 1 << 32, 1 << 16, 1 << 16, 1 << 16, 1e-6)); // rows*width past u32 +} + +test "the four-wide gate refuses ragged element counts" { + try testing.expectError(error.Invalid, vec4Guard(6)); + try vec4Guard(8); +} diff --git a/zig/src/vulkan/ops.zig b/zig/src/vulkan/ops.zig new file mode 100644 index 00000000000..cbac018bfd9 --- /dev/null +++ b/zig/src/vulkan/ops.zig @@ -0,0 +1,292 @@ +//! Stage 1 Vulkan operators: cast, raw copy, gather and affine dequant, each validated before any dispatch. +const std = @import("std"); +const memory = @import("memory.zig"); +const device = @import("device.zig"); +const commands = @import("commands.zig"); +const dispatch = @import("dispatch.zig"); +const params = @import("stage1_params"); +const tensor = @import("tensor.zig"); +const Device = device.Device; +const Commands = commands.Commands; +const Buffer = memory.Buffer; +const Error = dispatch.Error; + +/// The compiled shader images Ops runs; every field is a real SPIR-V image of the named variant. +pub const Shaders = struct { + qmm_bf16: []align(4) const u8, + qmm_q4_vec4: []align(4) const u8, + qmm_f32: []align(4) const u8, + qmm_q4_tile: []align(4) const u8, + qmm_q4_tile_f32: []align(4) const u8, + dequant_bf16: []align(4) const u8, + gather_u32: []align(4) const u8, + gather_u16: []align(4) const u8, + cast_bf16_f32: []align(4) const u8, + cast_f32_bf16: []align(4) const u8, + copy_u32: []align(4) const u8, + copy_u16: []align(4) const u8, + copy_f32_bf16: []align(4) const u8, +}; + +const dummy_bytes: u64 = 256; + +pub fn cast32(v: u64) Error!u32 { + return std.math.cast(u32, v) orelse error.TooLarge; +} + +pub fn dispatchGroups(count: u32, local: u32) [3]u32 { + return .{ @intCast((@as(u64, count) + local - 1) / local), 1, 1 }; +} + +/// gather_take geometry: one index array over a row-major table, window elements per row. +pub const GatherPlan = struct { window: u64, count: u32, ids: u32, rows: u32, stride: u32 }; + +/// The gather's derived geometry, or the named reason it cannot run; pure and host-tested. +pub fn gatherPlan(out_elements: u64, table_elements: u64, ids_elements: u64, rows: u64) Error!GatherPlan { + if (ids_elements == 0 or rows == 0) return error.Invalid; + if (out_elements != ids_elements *| (table_elements / rows)) return error.Invalid; + const window = table_elements / rows; + if (window * rows != table_elements or window == 0) return error.Invalid; + return .{ + .window = window, + .count = try cast32(out_elements), + .ids = try cast32(ids_elements), + .rows = try cast32(rows), + .stride = try cast32(window), + }; +} + +/// dequant geometry: rows of k elements, k/group packed words and scale groups per row. +pub const DequantPlan = struct { rows: u64, words: u64, groups: u32, count: u32, group: u32 }; + +/// The dequant's derived geometry, or the named reason it cannot run; pure and host-tested. +pub fn dequantPlan(out_elements: u64, w_elements: u64, scales_elements: u64, biases_elements: u64, k: u64, group: u64) Error!DequantPlan { + if (k == 0 or group == 0 or k % group != 0 or group % 8 != 0) return error.Invalid; + const words = k / 8; + const groups = k / group; + if (out_elements == 0 or out_elements % k != 0) return error.Invalid; + const rows = out_elements / k; + if (w_elements != rows * words or scales_elements != rows * groups or biases_elements != rows * groups) return error.Invalid; + return .{ + .rows = rows, + .words = words, + .groups = try cast32(groups), + .count = try cast32(rows * words), + .group = try cast32(group), + }; +} + +pub const Ops = struct { + dev: *const Device, + qmm_bf16: dispatch.Pipeline, + qmm_q4_vec4: dispatch.Pipeline, + qmm_f32: dispatch.Pipeline, + qmm_q4_tile: dispatch.Pipeline, + qmm_q4_tile_f32: dispatch.Pipeline, + dequant_bf16: dispatch.Pipeline, + gather_u32: dispatch.Pipeline, + gather_u16: dispatch.Pipeline, + cast_bf16_f32: dispatch.Pipeline, + cast_f32_bf16: dispatch.Pipeline, + copy_u32: dispatch.Pipeline, + copy_u16: dispatch.Pipeline, + copy_f32_bf16: dispatch.Pipeline, + /// Read-only stand-in for declared bindings a kernel never touches (cast 1, copy 1 and 3). + dummy: Buffer, + + /// Builds the ten pipelines; the device must have been opened with storage16 and int16 for the bf16 paths. + pub fn init(dev: *const Device, shaders: Shaders) Error!Ops { + if (!dev.features.storage16 or !dev.features.int16) return error.MissingFeature; + var o: Ops = .{ + .dev = dev, + .qmm_bf16 = undefined, + .qmm_q4_vec4 = undefined, + .qmm_f32 = undefined, + .qmm_q4_tile = undefined, + .qmm_q4_tile_f32 = undefined, + .dequant_bf16 = undefined, + .gather_u32 = undefined, + .gather_u16 = undefined, + .cast_bf16_f32 = undefined, + .cast_f32_bf16 = undefined, + .copy_u32 = undefined, + .copy_u16 = undefined, + .copy_f32_bf16 = undefined, + .dummy = try Buffer.init(dev, .device, dummy_bytes), + }; + errdefer o.dummy.deinit(); + o.qmm_bf16 = try dispatch.Pipeline.init(dev, shaders.qmm_bf16, .{ .bindings = 5, .push_bytes = params.push_bytes, .writes = 1 << 4 }); + errdefer o.qmm_bf16.deinit(); + o.qmm_q4_vec4 = try dispatch.Pipeline.init(dev, shaders.qmm_q4_vec4, .{ .bindings = 5, .push_bytes = params.push_bytes, .writes = 1 << 4 }); + errdefer o.qmm_q4_vec4.deinit(); + o.qmm_f32 = try dispatch.Pipeline.init(dev, shaders.qmm_f32, .{ .bindings = 5, .push_bytes = params.push_bytes, .writes = 1 << 4 }); + errdefer o.qmm_f32.deinit(); + o.qmm_q4_tile = try dispatch.Pipeline.init(dev, shaders.qmm_q4_tile, .{ .bindings = 5, .push_bytes = params.push_bytes, .writes = 1 << 4 }); + errdefer o.qmm_q4_tile.deinit(); + o.qmm_q4_tile_f32 = try dispatch.Pipeline.init(dev, shaders.qmm_q4_tile_f32, .{ .bindings = 5, .push_bytes = params.push_bytes, .writes = 1 << 4 }); + errdefer o.qmm_q4_tile_f32.deinit(); + o.dequant_bf16 = try dispatch.Pipeline.init(dev, shaders.dequant_bf16, .{ .bindings = 4, .push_bytes = params.push_bytes, .writes = 1 << 3 }); + errdefer o.dequant_bf16.deinit(); + o.gather_u32 = try dispatch.Pipeline.init(dev, shaders.gather_u32, .{ .bindings = 3, .push_bytes = params.push_bytes, .writes = 1 << 2 }); + errdefer o.gather_u32.deinit(); + o.gather_u16 = try dispatch.Pipeline.init(dev, shaders.gather_u16, .{ .bindings = 3, .push_bytes = params.push_bytes, .writes = 1 << 2 }); + errdefer o.gather_u16.deinit(); + o.cast_bf16_f32 = try dispatch.Pipeline.init(dev, shaders.cast_bf16_f32, .{ .bindings = 3, .push_bytes = params.cast_push_bytes, .writes = 1 << 2 }); + errdefer o.cast_bf16_f32.deinit(); + o.cast_f32_bf16 = try dispatch.Pipeline.init(dev, shaders.cast_f32_bf16, .{ .bindings = 3, .push_bytes = params.cast_push_bytes, .writes = 1 << 2 }); + errdefer o.cast_f32_bf16.deinit(); + o.copy_u32 = try dispatch.Pipeline.init(dev, shaders.copy_u32, .{ .bindings = 4, .push_bytes = params.push_bytes, .writes = 1 << 2 }); + errdefer o.copy_u32.deinit(); + o.copy_u16 = try dispatch.Pipeline.init(dev, shaders.copy_u16, .{ .bindings = 4, .push_bytes = params.push_bytes, .writes = 1 << 2 }); + errdefer o.copy_u16.deinit(); + o.copy_f32_bf16 = try dispatch.Pipeline.init(dev, shaders.copy_f32_bf16, .{ .bindings = 4, .push_bytes = params.push_bytes, .writes = 1 << 2 }); + return o; + } + + /// Destroys the pipelines and the dummy; call only after the last submission using them returned success. + pub fn deinit(o: *Ops) void { + o.qmm_bf16.deinit(); + o.qmm_q4_vec4.deinit(); + o.qmm_q4_tile.deinit(); + o.qmm_q4_tile_f32.deinit(); + o.dequant_bf16.deinit(); + o.gather_u32.deinit(); + o.gather_u16.deinit(); + o.cast_bf16_f32.deinit(); + o.cast_f32_bf16.deinit(); + o.copy_u32.deinit(); + o.copy_u16.deinit(); + o.copy_f32_bf16.deinit(); + o.dummy.deinit(); + } + + pub fn dummyBinding(o: *const Ops) dispatch.Binding { + return .{ .buffer = &o.dummy, .len = dummy_bytes }; + } + + /// bf16 -> f32 (RNE-free widening) or f32 -> bf16 (RNE with NaN quiet bit), one dispatch. + pub fn cast(o: *Ops, cmds: *Commands, comptime to_bf16: bool, out: tensor.View, src: tensor.View, timeout_ns: u64) Error!void { + const src_kind: tensor.Kind = if (to_bf16) .f32 else .bf16; + const out_kind: tensor.Kind = if (to_bf16) .bf16 else .f32; + if (src.kind != src_kind or out.kind != out_kind or src.elements != out.elements) return error.Invalid; + const pipe = if (to_bf16) &o.cast_f32_bf16 else &o.cast_bf16_f32; + try pipe.bind(&.{ src.bind(), o.dummyBinding(), out.bind() }); + var p: params.Params = .{ .count = try src.elemCount32() }; + try pipe.run(cmds, std.mem.asBytes(&p)[0..params.cast_push_bytes], dispatchGroups(p.count, pipe.local[0]), timeout_ns); + } + + pub fn castF32ToBf16(o: *Ops, cmds: *Commands, out: tensor.View, src: tensor.View, timeout_ns: u64) Error!void { + return o.cast(cmds, true, out, src, timeout_ns); + } + + pub fn castBf16ToF32(o: *Ops, cmds: *Commands, out: tensor.View, src: tensor.View, timeout_ns: u64) Error!void { + return o.cast(cmds, false, out, src, timeout_ns); + } + + /// Contiguous same-kind raw copy through copy_general (dims 1, unit strides); strided copies are refused. + pub fn copyPipeline(o: *Ops, kind: tensor.Kind) *dispatch.Pipeline { + return switch (kind) { + .f32, .u32 => &o.copy_u32, + .bf16, .u16 => &o.copy_u16, + }; + } + + pub fn copy(o: *Ops, cmds: *Commands, out: tensor.View, src: tensor.View, timeout_ns: u64) Error!void { + if (out.kind != src.kind or out.elements != src.elements) return error.Invalid; + const pipe = switch (out.kind) { + .f32, .u32 => &o.copy_u32, + .bf16, .u16 => &o.copy_u16, + }; + try pipe.bind(&.{ src.bind(), o.dummyBinding(), out.bind(), o.dummyBinding() }); + const count = try src.elemCount32(); + var p: params.Params = .{ + .count = count, + .dims = 1, + .shape = .{ count, 0, 0, 0 }, + .in_strides = .{ 1, 0, 0, 0 }, + .out_strides = .{ 1, 0, 0, 0 }, + }; + try pipe.run(cmds, std.mem.asBytes(&p), dispatchGroups(count, pipe.local[0]), timeout_ns); + } + + // gather_take uses u32 ids: out[t] reads table[id*window+s]; out-of-range ids yield zero; see ops_q4.checkIds. + pub fn gather(o: *Ops, cmds: *Commands, out: tensor.View, table: tensor.View, ids: tensor.View, rows: u64, timeout_ns: u64) Error!void { + if (ids.kind != .u32 or table.kind != out.kind) return error.Invalid; + const pipe = switch (out.kind) { + .u32 => &o.gather_u32, + .bf16, .u16 => &o.gather_u16, + .f32 => return error.Invalid, + }; + const plan = try gatherPlan(out.elements, table.elements, ids.elements, rows); + try pipe.bind(&.{ table.bind(), ids.bind(), out.bind() }); + var p: params.Params = .{ + .count = plan.count, + .operation = 1, // raw u32 indices + .reduce_size = plan.rows, + .matrix_m = plan.stride, + .dims = 1, + .shape = .{ try cast32(plan.window), 0, 0, 0 }, + .in_strides = .{ 1, 0, 0, 0 }, + .out_strides = .{ plan.ids, 0, 0, 0 }, + .flags = 1, // one window axis after the gathered one + }; + try pipe.run(cmds, std.mem.asBytes(&p), dispatchGroups(plan.count, pipe.local[0]), timeout_ns); + } + + /// Affine dequant: out = q * scale + bias per packed unit, bf16 round-to-nearest-even output. + pub fn dequant(o: *Ops, cmds: *Commands, out: tensor.View, w: tensor.View, scales: tensor.View, biases: tensor.View, k: u64, group: u64, timeout_ns: u64) Error!void { + if (out.kind != .bf16 or w.kind != .u32 or scales.kind != .bf16 or biases.kind != .bf16) return error.Invalid; + const plan = try dequantPlan(out.elements, w.elements, scales.elements, biases.elements, k, group); + try o.dequant_bf16.bind(&.{ w.bind(), scales.bind(), biases.bind(), out.bind() }); + var p: params.Params = .{ + .count = plan.count, + .operation = 4, // 4-bit codes + .reduce_size = plan.group, + .matrix_n = try cast32(plan.words), + .matrix_k = plan.groups, + }; + try o.dequant_bf16.run(cmds, std.mem.asBytes(&p), dispatchGroups(plan.count, o.dequant_bf16.local[0]), timeout_ns); + } +}; + +const testing = std.testing; + +test "gather plans derive the window and refuse torn shapes" { + const plan = try gatherPlan(6 * 128, 512 * 128, 6, 512); + try testing.expectEqual(@as(u64, 128), plan.window); + try testing.expectEqual(@as(u32, 768), plan.count); + try testing.expectEqual(@as(u32, 512), plan.rows); + try testing.expectEqual(@as(u32, 128), plan.stride); + try testing.expectError(error.Invalid, gatherPlan(0, 128, 1, 1)); + try testing.expectError(error.Invalid, gatherPlan(128, 128, 1, 0)); + try testing.expectError(error.Invalid, gatherPlan(768, 512, 6, 512)); // torn window + try testing.expectError(error.Invalid, gatherPlan(768, 510 * 128, 6, 511)); + try testing.expectError(error.TooLarge, gatherPlan((1 << 32) * 4, (1 << 32) * 4, 1 << 32, 1 << 32)); +} + +test "stage1 out-of-range gather uses one output row for one id" { + _ = try gatherPlan(40, 60, 4, 6); + try testing.expectError(error.Invalid, gatherPlan(40, 60, 1, 6)); + const one = try gatherPlan(10, 60, 1, 6); + try testing.expectEqual(@as(u32, 10), one.count); +} + +test "dequant plans derive rows, words and groups and refuse torn geometry" { + const plan = try dequantPlan(4 * 4096, 4 * 512, 4 * 64, 4 * 64, 4096, 64); + try testing.expectEqual(@as(u64, 4), plan.rows); + try testing.expectEqual(@as(u64, 512), plan.words); + try testing.expectEqual(@as(u32, 64), plan.groups); + try testing.expectEqual(@as(u32, 2048), plan.count); + try testing.expectError(error.Invalid, dequantPlan(4096, 512, 64, 64, 4096, 65)); // k % group + try testing.expectError(error.Invalid, dequantPlan(4096, 511, 64, 64, 4096, 64)); // torn words + try testing.expectError(error.Invalid, dequantPlan(4096, 512, 63, 64, 4096, 64)); // torn scales + try testing.expectError(error.Invalid, dequantPlan(4095, 511, 64, 64, 4096, 64)); // torn rows + try testing.expectError(error.TooLarge, dequantPlan(8 << 35, 8 << 32, 8 << 29, 8 << 29, 1 << 35, 64)); +} + +test "dispatch groups round up without wrapping" { + try testing.expectEqual([3]u32{ 1, 1, 1 }, dispatchGroups(1, 256)); + try testing.expectEqual([3]u32{ 1, 1, 1 }, dispatchGroups(256, 256)); + try testing.expectEqual([3]u32{ 2, 1, 1 }, dispatchGroups(257, 256)); + try testing.expectEqual([3]u32{ (1 << 32) / 256, 1, 1 }, dispatchGroups(std.math.maxInt(u32), 256)); +} diff --git a/zig/src/vulkan/ops_q4.zig b/zig/src/vulkan/ops_q4.zig new file mode 100644 index 00000000000..66f2dbe4cae --- /dev/null +++ b/zig/src/vulkan/ops_q4.zig @@ -0,0 +1,312 @@ +//! The composed stage 1 closures: Q4 affine projection and Q4-table embedding over gathered, dequantized rows. +const std = @import("std"); +const commands = @import("commands.zig"); +const dispatch = @import("dispatch.zig"); +const params = @import("stage1_params"); +const tensor = @import("tensor.zig"); +const arena_mod = @import("arena.zig"); +const ops_mod = @import("ops.zig"); +const tile = @import("tile.zig"); +const Commands = commands.Commands; +const Arena = arena_mod.Arena; +const Ops = ops_mod.Ops; +const View = tensor.View; +const Error = dispatch.Error; + +/// Stage 1 serves Q4, group 64, transposed weights only; wider coverage is later work, not a silent pass. +pub const group_size: u64 = 64; +pub const bits: u64 = 4; + +/// The device numbers the chunk planner shapes itself around. +pub const Limits = struct { max_range: u64, alignment: u64 }; + +// n rows per qmm chunk (0 if all of w fits one binding); k % 64 == 0; w bytes within max_range; w/scale starts aligned. +pub fn chunkRows(l: Limits, k: u64, n: u64) Error!u64 { + const row_w = k / 2; // k/8 u32 words per row + const row_s = k / 32; // k/64 bf16 scales (or biases) per row + if (n *| row_w <= l.max_range) return 0; + const step_w = l.alignment / std.math.gcd(l.alignment, row_w); + const step_s = l.alignment / std.math.gcd(l.alignment, row_s); + const step = std.math.mul(u64, step_w / std.math.gcd(step_w, step_s), step_s) catch return error.TooLarge; + const group_bytes = std.math.mul(u64, step, row_w) catch return error.TooLarge; + if (group_bytes > l.max_range) return error.TooLarge; + return (l.max_range / group_bytes) * step; +} + +const Geometry = struct { m: u64, n: u64 }; + +/// How w, scales and biases are stored: one row a column, or tiles of 32 columns with the last tile padded. +const Storage = enum { rows, tiles }; + +/// The projection's m and n from the views and k, or the named reason it cannot run; x/out bf16 or f32; weights u32. +fn projectGeometry(out: View, x: View, w: View, scales: View, biases: View, k: u64, storage: Storage) Error!Geometry { + if (out.kind != x.kind or (out.kind != .bf16 and out.kind != .f32)) return error.Invalid; + if (w.kind != .u32 or scales.kind != .bf16 or biases.kind != .bf16) return error.Invalid; + if (k < group_size or k % group_size != 0) return error.Invalid; + if (x.elements == 0 or x.elements % k != 0) return error.Invalid; + const m = x.elements / k; + if (out.elements == 0 or out.elements % m != 0) return error.Invalid; + const n = out.elements / m; + const stored: u64 = if (storage == .tiles) tile.padded(@intCast(n)) else n; + const per_row_w = k / 8; + const per_row_p = k / 64; + const want_w = std.math.mul(u64, stored, per_row_w) catch return error.Invalid; + const want_p = std.math.mul(u64, stored, per_row_p) catch return error.Invalid; + if (w.elements != want_w or scales.elements != want_p or biases.elements != want_p) return error.Invalid; + return .{ .m = m, .n = n }; +} + +fn q4Push(m: u64, n: u64, k: u64) Error!params.Params { + const count = try ops_mod.cast32(m * n); + return .{ + .count = count, + .operation = @intCast(bits), + .reduce_size = @intCast(group_size), + .output_size = count, + .matrix_m = try ops_mod.cast32(m), + .matrix_n = try ops_mod.cast32(n), + .matrix_k = try ops_mod.cast32(k), + .shape = .{ 1, 0, 0, 0 }, // one x batch, shared w and per-row params + }; +} + +/// Which qmm image a dispatch runs: `scalar` is qmm.comp, `vec4` is qmm_q4_vec4.comp (same arithmetic, 16-byte loads). +pub const Qmm = enum { scalar, vec4 }; + +// vec4 loads uvec4 from x and w: both windows must start 16-byte aligned (k % 64 keeps strides and ranges aligned). +pub fn qmmVariant(x: View, w: View) Qmm { + return if (x.byte_offset % 16 == 0 and w.byte_offset % 16 == 0) .vec4 else .scalar; +} + +fn qmmPipeline(o: *Ops, force: ?Qmm, act_f32: bool, x: View, w: View) Error!*dispatch.Pipeline { + if (act_f32) return &o.qmm_f32; // the wide layers 0-3 read f32 activations + const fit = qmmVariant(x, w); + const want = force orelse fit; + if (want == .vec4 and fit == .scalar) return error.Misaligned; + return if (want == .vec4) &o.qmm_q4_vec4 else &o.qmm_bf16; +} + +// Q4 projection out[m,n] = x[m,k] . dequant(w,scales,biases)[n,k]; scratch + copy_general above maxStorageBufferRange. +pub fn project(o: *Ops, cmds: *Commands, out: View, x: View, w: View, scales: View, biases: View, k: u64, scratch: *Arena, timeout_ns: u64) Error!void { + return projectAs(o, null, cmds, out, x, w, scales, biases, k, scratch, timeout_ns); +} + +/// `project` over weights repacked by `tile`: one dispatch; bf16 or f32 (both x and out) activations. +pub fn tiledActivation(out: View, x: View) Error!void { + if (out.kind != x.kind or (x.kind != .bf16 and x.kind != .f32)) return error.Invalid; +} + +pub fn projectTiled(o: *Ops, cmds: *Commands, out: View, x: View, w: View, scales: View, biases: View, k: u64, timeout_ns: u64) Error!void { + try tiledActivation(out, x); + const g = try projectGeometry(out, x, w, scales, biases, k, .tiles); + if (!tile.fits(g.n, k)) return error.Invalid; + const range = o.dev.properties.limits.maxStorageBufferRange; + if (x.byteLen() > range or out.byteLen() > range or w.byteLen() > range) return error.TooLarge; + if (qmmVariant(x, w) != .vec4) return error.Misaligned; + const pipe = if (out.kind == .f32) &o.qmm_q4_tile_f32 else &o.qmm_q4_tile; + try pipe.bind(&.{ x.bind(), w.bind(), scales.bind(), biases.bind(), out.bind() }); + const p = try q4Push(g.m, g.n, k); + try pipe.run(cmds, std.mem.asBytes(&p), ops_mod.dispatchGroups(p.count, pipe.local[0]), timeout_ns); +} + +/// `project` with the qmm image forced (null picks by alignment); a forced vec4 over unaligned windows is refused. +pub fn projectAs(o: *Ops, force: ?Qmm, cmds: *Commands, out: View, x: View, w: View, scales: View, biases: View, k: u64, scratch: *Arena, timeout_ns: u64) Error!void { + const g = try projectGeometry(out, x, w, scales, biases, k, .rows); + const limits = o.dev.properties.limits; + const lim = Limits{ .max_range = limits.maxStorageBufferRange, .alignment = limits.minStorageBufferOffsetAlignment }; + if (x.byteLen() > lim.max_range or out.byteLen() > lim.max_range) return error.TooLarge; + const rows = try chunkRows(lim, k, g.n); + const act_f32 = out.kind == .f32; + if (rows == 0) { + const pipe = try qmmPipeline(o, force, act_f32, x, w); + try pipe.bind(&.{ x.bind(), w.bind(), scales.bind(), biases.bind(), out.bind() }); + const p = try q4Push(g.m, g.n, k); + try pipe.run(cmds, std.mem.asBytes(&p), ops_mod.dispatchGroups(p.count, pipe.local[0]), timeout_ns); + return; + } + const block = try scratch.alloc(out.kind, g.m * rows); + if (block.byteLen() > lim.max_range) return error.TooLarge; + // bf16 scatters ride raw u16 columns, f32 scatters raw u32 words; both keep every RNE byte. + const scatter = if (act_f32) &o.copy_u32 else &o.copy_u16; + const row_words = k / 8; + const row_groups = k / 64; + var first: u64 = 0; + while (first < g.n) { + const chunk = @min(rows, g.n - first); + // chunkRows made every chunk start an aligned byte of the w and scale streams. + const w_chunk = try w.slice(first * row_words, chunk * row_words); + const s_chunk = try scales.slice(first * row_groups, chunk * row_groups); + const b_chunk = try biases.slice(first * row_groups, chunk * row_groups); + const block_chunk = try block.slice(0, g.m * chunk); + const pipe = try qmmPipeline(o, force, act_f32, x, w_chunk); + try pipe.bind(&.{ x.bind(), w_chunk.bind(), s_chunk.bind(), b_chunk.bind(), block_chunk.bind() }); + const p = try q4Push(g.m, chunk, k); + try pipe.run(cmds, std.mem.asBytes(&p), ops_mod.dispatchGroups(p.count, pipe.local[0]), timeout_ns); + // Scatter the packed block into the full row-major output. + try scatter.bind(&.{ block_chunk.bind(), o.dummyBinding(), out.bind(), o.dummyBinding() }); + var cp: params.Params = .{ + .count = try ops_mod.cast32(g.m * chunk), + .dims = 2, + .shape = .{ try ops_mod.cast32(g.m), try ops_mod.cast32(chunk), 0, 0 }, + .in_strides = .{ try ops_mod.cast32(chunk), 1, 0, 0 }, + .out_strides = .{ try ops_mod.cast32(g.n), 1, 0, 0 }, + .output_offset = try ops_mod.cast32(first), + }; + try scatter.run(cmds, std.mem.asBytes(&cp), ops_mod.dispatchGroups(cp.count, scatter.local[0]), timeout_ns); + first += chunk; + } +} + +// Host token validation: every id must index a table row, checked before upload; out-of-range ids gather to zero. +pub fn checkIds(ids: []const u32, rows: u64) Error!void { + for (ids) |id| if (id >= rows) return error.OutOfRange; +} + +// Embedding rows/tokens or why not; scales/biases ride the 32-bit transport, so tokens*(k/64) fills whole words. +fn embeddingGeometry(out: View, ids: View, table_w: View, table_s: View, table_b: View, k: u64) Error!struct { rows: u64, tokens: u64 } { + if (out.kind != .bf16 or ids.kind != .u32 or table_w.kind != .u32 or table_s.kind != .bf16 or table_b.kind != .bf16) return error.Invalid; + if (k < group_size or k % group_size != 0 or ids.elements == 0) return error.Invalid; + const tokens = ids.elements; + const want_out = std.math.mul(u64, tokens, k) catch return error.Invalid; + if (out.elements != want_out) return error.Invalid; + const row_words = k / 8; + const row_groups = k / 64; + if ((tokens * row_groups) % 2 != 0) return error.Invalid; + if (table_w.elements == 0 or table_w.elements % row_words != 0) return error.Invalid; + const rows = table_w.elements / row_words; + const want_p = std.math.mul(u64, rows, row_groups) catch return error.Invalid; + if (table_s.elements != want_p or table_b.elements != want_p) return error.Invalid; + return .{ .rows = rows, .tokens = tokens }; +} + +// Q4 embedding: rows to `scratch`, bf16 out; tables are single bindings — larger refused, never chunked. +pub fn embedding(o: *Ops, cmds: *Commands, out: View, ids: View, table_w: View, table_s: View, table_b: View, k: u64, scratch: *Arena, timeout_ns: u64) Error!void { + const g = try embeddingGeometry(out, ids, table_w, table_s, table_b, k); + const range = o.dev.properties.limits.maxStorageBufferRange; + if (table_w.byteLen() > range or table_s.byteLen() > range or table_b.byteLen() > range) return error.TooLarge; + const row_words = k / 8; + const row_groups = k / 64; + // Allocations happen before any dispatch, so a full scratch is a clean refusal, not a torn result. + const gathered_w = try scratch.alloc(.u32, g.tokens * row_words); + const gathered_s = try scratch.alloc(.bf16, g.tokens * row_groups); + const gathered_b = try scratch.alloc(.bf16, g.tokens * row_groups); + try o.gather(cmds, gathered_w, table_w, ids, g.rows, timeout_ns); + try o.gather(cmds, gathered_s, table_s, ids, g.rows, timeout_ns); + try o.gather(cmds, gathered_b, table_b, ids, g.rows, timeout_ns); + try o.dequant(cmds, out, gathered_w, gathered_s, gathered_b, k, group_size, timeout_ns); +} + +const testing = std.testing; + +test "chunkRows returns zero when one binding fits and aligned groups when it does not" { + const fits: Limits = .{ .max_range = 1 << 30, .alignment = 256 }; + try testing.expectEqual(@as(u64, 0), try chunkRows(fits, 4096, 1000)); + const tight: Limits = .{ .max_range = 1 << 27, .alignment = 256 }; + try testing.expectEqual(@as(u64, 65536), try chunkRows(tight, 4096, 100_000)); // 2-row steps of 4096 bytes fill 128 MiB + const small: Limits = .{ .max_range = 1 << 16, .alignment = 256 }; + const rows = try chunkRows(small, 64, 1 << 20); // row 32B, scale row 2B: 128-row steps of 4096 bytes, 16 steps fill 64 KiB + try testing.expectEqual(@as(u64, 2048), rows); + try testing.expectEqual(@as(u64, 0), rows % 128); + try testing.expect(rows * 32 <= small.max_range); +} + +test "chunkRows refuses an aligned group larger than the range" { + const absurd: Limits = .{ .max_range = 1 << 12, .alignment = 1 << 12 }; + try testing.expectError(error.TooLarge, chunkRows(absurd, 64, 1 << 20)); +} + +test "project geometry derives m and n and refuses torn or wrong-kind views" { + const dev: @import("device.zig").Device = undefined; + const buf = @import("memory.zig").Buffer{ .dev = &dev, .handle = null, .memory = null, .len = 1 << 20, .kind = .device, .mapped = null, .coherent = false }; + const x = try View.init(&buf, .bf16, 0, 3 * 128); + const out = try View.init(&buf, .bf16, 1024, 3 * 100); + const w = try View.init(&buf, .u32, 2048, 100 * 16); + const scales = try View.init(&buf, .bf16, 4096, 100 * 2); + const biases = try View.init(&buf, .bf16, 4096, 100 * 2); + const g = try projectGeometry(out, x, w, scales, biases, 128, .rows); + try testing.expectEqual(@as(u64, 3), g.m); + try testing.expectEqual(@as(u64, 100), g.n); + try testing.expectError(error.Invalid, projectGeometry(out, x, w, scales, biases, 100, .rows)); + try testing.expectError(error.Invalid, projectGeometry(out, x, w, scales, biases, 256, .rows)); + const f32x = try View.init(&buf, .f32, 1024, 3 * 128); + try testing.expectError(error.Invalid, projectGeometry(out, f32x, w, scales, biases, 128, .rows)); // mixed kinds + // the wide combination: f32 x and f32 out over the same bf16 scales and biases + const f32out = try View.init(&buf, .f32, 8192, 3 * 100); + const g32 = try projectGeometry(f32out, f32x, w, scales, biases, 128, .rows); + try testing.expectEqual(@as(u64, 3), g32.m); + try testing.expectEqual(@as(u64, 100), g32.n); + const torn = try View.init(&buf, .u32, 2048, 100 * 16 + 1); + try testing.expectError(error.Invalid, projectGeometry(out, x, torn, scales, biases, 128, .rows)); +} + +test "tiled storage holds n rounded up to whole tiles of columns" { + const dev: @import("device.zig").Device = undefined; + const buf = @import("memory.zig").Buffer{ .dev = &dev, .handle = null, .memory = null, .len = 1 << 20, .kind = .device, .mapped = null, .coherent = false }; + const out = try View.init(&buf, .bf16, 0, 3 * 100); + const x = try View.init(&buf, .bf16, 1024, 3 * 128); + const w = try View.init(&buf, .u32, 2048, 128 * 16); + const scales = try View.init(&buf, .bf16, 16384, 128 * 2); + const biases = try View.init(&buf, .bf16, 32768, 128 * 2); + const g = try projectGeometry(out, x, w, scales, biases, 128, .tiles); + try testing.expectEqual(@as(u64, 100), g.n); + try testing.expectError(error.Invalid, projectGeometry(out, x, w, scales, biases, 128, .rows)); +} + +test "the tiled image takes matching bf16 or f32 activations and refuses a mix" { + const dev: @import("device.zig").Device = undefined; + const buf = @import("memory.zig").Buffer{ .dev = &dev, .handle = null, .memory = null, .len = 1 << 20, .kind = .device, .mapped = null, .coherent = false }; + const out = try View.init(&buf, .bf16, 0, 3 * 100); + const x = try View.init(&buf, .bf16, 1024, 3 * 128); + try tiledActivation(out, x); + const f32x = try View.init(&buf, .f32, 8192, 3 * 128); + const f32out = try View.init(&buf, .f32, 16384, 3 * 100); + try tiledActivation(f32out, f32x); + try testing.expectError(error.Invalid, tiledActivation(f32out, x)); + try testing.expectError(error.Invalid, tiledActivation(out, f32x)); +} + +test "the vec4 qmm image is chosen only when every vector load is 16-byte aligned" { + const dev: @import("device.zig").Device = undefined; + const buf = @import("memory.zig").Buffer{ .dev = &dev, .handle = null, .memory = null, .len = 1 << 20, .kind = .device, .mapped = null, .coherent = false }; + const w = try View.init(&buf, .u32, 2048, 100 * 16); + const x = try View.init(&buf, .bf16, 64, 3 * 128); + try testing.expectEqual(Qmm.vec4, qmmVariant(x, w)); + // 4-byte windows are legal views, but a uvec4 load through them would be misaligned. + try testing.expectEqual(Qmm.scalar, qmmVariant(try View.init(&buf, .bf16, 68, 3 * 128), w)); + try testing.expectEqual(Qmm.scalar, qmmVariant(try View.init(&buf, .bf16, 72, 3 * 128), w)); + try testing.expectEqual(Qmm.scalar, qmmVariant(x, try View.init(&buf, .u32, 2052, 100 * 16))); + try testing.expectEqual(Qmm.scalar, qmmVariant(x, try View.init(&buf, .u32, 2056, 100 * 16))); + try testing.expectEqual(Qmm.vec4, qmmVariant(try View.init(&buf, .bf16, 16, 3 * 128), try View.init(&buf, .u32, 32, 100 * 16))); +} + +test "embedding geometry derives rows and tokens and refuses torn tables" { + const dev: @import("device.zig").Device = undefined; + const buf = @import("memory.zig").Buffer{ .dev = &dev, .handle = null, .memory = null, .len = 1 << 20, .kind = .device, .mapped = null, .coherent = false }; + const ids = try View.init(&buf, .u32, 0, 5); + const out = try View.init(&buf, .bf16, 64, 5 * 128); + const w = try View.init(&buf, .u32, 2048, 100 * 16); + const scales = try View.init(&buf, .bf16, 4096, 100 * 2); + const biases = try View.init(&buf, .bf16, 4096, 100 * 2); + const g = try embeddingGeometry(out, ids, w, scales, biases, 128); + try testing.expectEqual(@as(u64, 100), g.rows); + try testing.expectEqual(@as(u64, 5), g.tokens); + try testing.expectError(error.Invalid, embeddingGeometry(out, ids, w, scales, biases, 64)); // out rows torn + const torn_s = try View.init(&buf, .bf16, 4096, 99 * 2); + try testing.expectError(error.Invalid, embeddingGeometry(out, ids, w, torn_s, biases, 128)); + // k 64 gathers one scale a token: five tokens leave a ragged bf16 pair, six fill whole words. + const k64_out5 = try View.init(&buf, .bf16, 8192, 5 * 64); + const ids6 = try View.init(&buf, .u32, 0, 6); + const k64_out6 = try View.init(&buf, .bf16, 8192, 6 * 64); + const w64 = try View.init(&buf, .u32, 8704, 100 * 8); + const s64 = try View.init(&buf, .bf16, 8704, 100); + try testing.expectError(error.Invalid, embeddingGeometry(k64_out5, ids, w64, s64, s64, 64)); + const g64 = try embeddingGeometry(k64_out6, ids6, w64, s64, s64, 64); + try testing.expectEqual(@as(u64, 6), g64.tokens); + try testing.expectEqual(@as(u64, 100), g64.rows); +} + +test "checkIds refuses any id past the table" { + try checkIds(&.{ 0, 3, 99 }, 100); + try checkIds(&.{}, 1); + try testing.expectError(error.OutOfRange, checkIds(&.{ 0, 100 }, 100)); +} diff --git a/zig/src/vulkan/params.zig b/zig/src/vulkan/params.zig new file mode 100644 index 00000000000..38f3afad6a3 --- /dev/null +++ b/zig/src/vulkan/params.zig @@ -0,0 +1,44 @@ +//! The 120-byte push-constant block shared by every stage 1 Vulkan shader. +const std = @import("std"); + +/// qmm.comp's push constants (std430); shaders read the fields each op sets, the rest keep the defaults below. +pub const Params = extern struct { + count: u32 = 0, + operation: u32 = 0, + lhs_size: u32 = 0, + rhs_size: u32 = 0, + reduce_size: u32 = 0, + output_size: u32 = 0, + lhs_offset: u32 = 0, + rhs_offset: u32 = 0, + output_offset: u32 = 0, + aux_size: u32 = 0, + aux_offset: u32 = 0, + matrix_m: u32 = 0, + matrix_n: u32 = 0, + matrix_k: u32 = 0, + flags: u32 = 0, + alpha: f32 = 1, + beta: f32 = 0, + dims: u32 = 0, + shape: [4]u32 = .{ 0, 0, 0, 0 }, + in_strides: [4]u32 = .{ 0, 0, 0, 0 }, + out_strides: [4]u32 = .{ 0, 0, 0, 0 }, +}; + +/// Push bytes a full-Params pipeline (qmm, dequant, gather_take, copy_general) records. +pub const push_bytes: u32 = 120; +/// cast.comp reads only these first fields (count through output_offset). +pub const cast_push_bytes: u32 = 36; + +comptime { + const names = [_][]const u8{ "count", "operation", "lhs_size", "rhs_size", "reduce_size", "output_size", "lhs_offset", "rhs_offset", "output_offset", "aux_size", "aux_offset", "matrix_m", "matrix_n", "matrix_k", "flags", "alpha", "beta", "dims" }; + for (names, 0..) |name, i| if (@offsetOf(Params, name) != 4 * i) @compileError("qmm.comp push constant moved: " ++ name); + if (@offsetOf(Params, "shape") != 72 or @offsetOf(Params, "in_strides") != 88 or @offsetOf(Params, "out_strides") != 104) @compileError("qmm.comp push constant arrays moved"); + if (@sizeOf(Params) != push_bytes) @compileError("qmm.comp push constants are 120 bytes"); +} + +test "cast reads only the prefix the cast pipeline records" { + try std.testing.expect(cast_push_bytes <= push_bytes); + try std.testing.expectEqual(@as(u32, 36), cast_push_bytes); +} diff --git a/zig/src/vulkan/root.zig b/zig/src/vulkan/root.zig index 11fece45fcf..478fb965dbb 100644 --- a/zig/src/vulkan/root.zig +++ b/zig/src/vulkan/root.zig @@ -16,11 +16,27 @@ pub const Commands = commands.Commands; pub const barrier = commands.barrier; pub const upload = commands.upload; pub const download = commands.download; +pub const fill = commands.fill; +pub const zero = commands.zero; +pub const recordFill = commands.recordFill; pub const orExit = commands.orExit; pub const dispatch = @import("dispatch.zig"); +pub const params = @import("stage1_params"); +pub const tensor = @import("tensor.zig"); +pub const arena = @import("arena.zig"); +pub const ops = @import("ops.zig"); +pub const ops_q4 = @import("ops_q4.zig"); +pub const tile = @import("tile.zig"); +pub const norm_mlp = @import("norm_mlp.zig"); +pub const attention = @import("attention.zig"); +pub const Attention = attention.Attention; +pub const gdn = @import("gdn.zig"); +pub const Gdn = gdn.Gdn; +pub const greedy = @import("greedy.zig"); +pub const Greedy = greedy.Greedy; test { const std = @import("std"); - inline for (.{ runtime, Runtime, device, Device, memory, Buffer, commands, Commands, dispatch, dispatch.Pipeline }) |T| + inline for (.{ runtime, Runtime, device, Device, memory, Buffer, commands, Commands, dispatch, dispatch.Pipeline, params, tensor, tensor.View, arena, ops, ops_q4, tile, norm_mlp, norm_mlp.NormMlp, attention, Attention, gdn, Gdn, greedy, Greedy }) |T| std.testing.refAllDecls(T); } diff --git a/zig/src/vulkan/runtime.zig b/zig/src/vulkan/runtime.zig index 93ca6574d04..c51864de4c0 100644 --- a/zig/src/vulkan/runtime.zig +++ b/zig/src/vulkan/runtime.zig @@ -14,6 +14,10 @@ pub const Error = error{ Timeout, /// A submit or fence wait failed with the work's state unknown: the GPU may still use it; the caller exits. InFlight, + /// A post-wait fence reset failed and left the fence signaled: the stream is poisoned; the caller exits. + FencePoisoned, + /// A command-buffer end failed and left the buffer invalid: the stream is poisoned; the caller exits. + CommandBufferInvalid, DeviceLost, NoDevice, SoftwareDevice, diff --git a/zig/src/vulkan/tensor.zig b/zig/src/vulkan/tensor.zig new file mode 100644 index 00000000000..a53cae047d1 --- /dev/null +++ b/zig/src/vulkan/tensor.zig @@ -0,0 +1,134 @@ +//! Honest device tensor views: a byte range of one buffer with a fixed element kind, checked before any shader sees it. +const std = @import("std"); +const memory = @import("memory.zig"); +const dispatch = @import("dispatch.zig"); +const Buffer = memory.Buffer; +const Error = @import("runtime.zig").Error; + +/// The element kinds stage 1 moves and computes; bf16 rides storage as raw uint16 bits. +pub const Kind = enum { + f32, + bf16, + u32, + u16, + + /// Bytes per element; every kind is 2 or 4, so 4-aligned byte offsets are element-aligned. + pub fn size(kind: Kind) u64 { + return switch (kind) { + .f32, .u32 => 4, + .bf16, .u16 => 2, + }; + } +}; + +/// `elements` of `kind` in bytes, refusing overflow. +pub fn extent(kind: Kind, elements: u64) Error!u64 { + const bytes = elements *| kind.size(); + if (elements != 0 and bytes / kind.size() != elements) return error.TooLarge; + return bytes; +} + +/// A checked `{ byte_offset, elements, kind }` window over one buffer; never a raw device address. +pub const View = struct { + buffer: *const Buffer, + byte_offset: u64, + elements: u64, + kind: Kind, + + /// A nonempty, 4-byte-aligned, word-padded window inside `buffer`; bf16/u16 windows hold whole words. + pub fn init(buffer: *const Buffer, kind: Kind, byte_offset: u64, elements: u64) Error!View { + if (elements == 0) return error.OutOfRange; + if (byte_offset % 4 != 0) return error.Misaligned; + const bytes = try extent(kind, elements); + try memory.span(buffer.len, byte_offset, bytes); + if (bytes % 4 != 0) return error.Misaligned; // std430 word transport: a ragged tail drops its last element + return .{ .buffer = buffer, .byte_offset = byte_offset, .elements = elements, .kind = kind }; + } + + pub fn byteLen(v: View) u64 { + return v.elements * v.kind.size(); // init proved this fits and is word-padded + } + + /// A subwindow starting `first` elements in; it cannot escape the window it slices, only the buffer. + pub fn slice(v: View, first: u64, elements: u64) Error!View { + if (first > v.elements or elements > v.elements - first) return error.OutOfRange; + return init(v.buffer, v.kind, v.byte_offset + first * v.kind.size(), elements); + } + + /// The window's byte offset as an element index, as the shaders' push constants carry offsets. + pub fn elemOffset32(v: View) Error!u32 { + return std.math.cast(u32, v.byte_offset / v.kind.size()) orelse error.TooLarge; + } + + pub fn elemCount32(v: View) Error!u32 { + return std.math.cast(u32, v.elements) orelse error.TooLarge; + } + + /// The descriptor range this window binds: its own base, so shader element offsets start at zero. + pub fn bind(v: View) dispatch.Binding { + return .{ .buffer = v.buffer, .offset = v.byte_offset, .len = v.byteLen() }; + } + + /// Whether two windows overlap (same buffer only); used to refuse aliasing before Pipeline.bind does. + pub fn overlaps(a: View, b: View) bool { + return a.buffer == b.buffer and a.byte_offset < b.byte_offset + b.byteLen() and b.byte_offset < a.byte_offset + a.byteLen(); + } +}; + +const testing = std.testing; + +fn fakeBuffer(len: u64) Buffer { + return .{ .dev = undefined, .handle = null, .memory = null, .len = len, .kind = .device, .mapped = null, .coherent = false }; +} + +test "views hold their geometry and refuse empty, misaligned and out-of-range windows" { + const buf = fakeBuffer(1024); + const v = try View.init(&buf, .f32, 16, 128); + try testing.expectEqual(@as(u64, 16), v.byte_offset); + try testing.expectEqual(@as(u64, 512), v.byteLen()); + try testing.expectEqual(@as(u32, 4), try v.elemOffset32()); + try testing.expectEqual(@as(u32, 128), try v.elemCount32()); + try testing.expectError(error.OutOfRange, View.init(&buf, .f32, 0, 0)); + try testing.expectError(error.Misaligned, View.init(&buf, .f32, 2, 8)); + try testing.expectError(error.OutOfRange, View.init(&buf, .f32, 512, 129)); + try testing.expectError(error.OutOfRange, View.init(&buf, .f32, 1024, 1)); + try testing.expectError(error.TooLarge, extent(.bf16, std.math.maxInt(u64) / 2 + 1)); +} + +test "bf16 and u16 windows must end on a word boundary" { + const buf = fakeBuffer(64); + try testing.expectError(error.Misaligned, View.init(&buf, .bf16, 0, 3)); + _ = try View.init(&buf, .bf16, 8, 4); + try testing.expectError(error.Misaligned, View.init(&buf, .u16, 0, 1)); +} + +test "slices stay inside their window and keep the base offset" { + const buf = fakeBuffer(256); + const v = try View.init(&buf, .u32, 32, 32); + const s = try v.slice(4, 8); + try testing.expectEqual(@as(u64, 48), s.byte_offset); + try testing.expectEqual(@as(u64, 8), s.elements); + try testing.expectError(error.OutOfRange, v.slice(33, 1)); + try testing.expectError(error.OutOfRange, v.slice(0, 33)); + try testing.expectError(error.OutOfRange, v.slice(31, 2)); // may not escape past the window's end + try testing.expectError(error.OutOfRange, v.slice(5, std.math.maxInt(u64))); +} + +test "element offsets past 32 bits refuse instead of wrapping" { + const buf = fakeBuffer((1 << 34) + 64); + const v = try View.init(&buf, .u32, 4, (1 << 32) + 1); + try testing.expectError(error.TooLarge, v.elemCount32()); + try testing.expectEqual(@as(u32, 1), try v.elemOffset32()); +} + +test "overlap is same-buffer range intersection" { + const buf = fakeBuffer(256); + const other = fakeBuffer(256); + const a = try View.init(&buf, .u32, 0, 16); + const b = try View.init(&buf, .u32, 64, 16); + const c = try View.init(&buf, .u32, 60, 4); + const d = try View.init(&other, .u32, 0, 16); + try testing.expect(a.overlaps(c)); + try testing.expect(!a.overlaps(b)); + try testing.expect(!a.overlaps(d)); +} diff --git a/zig/src/vulkan/tile.zig b/zig/src/vulkan/tile.zig new file mode 100644 index 00000000000..062523b3996 --- /dev/null +++ b/zig/src/vulkan/tile.zig @@ -0,0 +1,104 @@ +//! Host repack of Q4 group-64 weights into 32-column tiles for qmm_q4_tile; whole vectors and scales move unchanged. +const std = @import("std"); + +/// Columns (outputs) in one tile; one Honeykrisp subgroup, so a lane's 16-byte load sits next to its neighbours'. +pub const columns: usize = 32; +const vector: usize = 16; +const group: usize = 64; + +/// A projection can be tiled when k fills whole groups; the last tile gets zero columns the kernel never computes. +pub fn fits(n: u64, k: u64) bool { + return n != 0 and k != 0 and k % group == 0; +} + +/// Columns stored for n outputs: n rounded up to whole tiles. +pub fn padded(n: usize) usize { + return (n + columns - 1) / columns * columns; +} + +/// Words (rows of k/32 vectors) into [tiles][k/32 vectors][32 lanes]; src is n * k / 2 bytes, dst padded(n) * k / 2. +pub fn tileWords(dst: []u8, src: []const u8, n: usize, k: usize) void { + std.debug.assert(fits(n, k) and src.len == n * k / 2 and dst.len == padded(n) * k / 2); + @memset(dst, 0); + const vectors = k / 32; + for (0..n) |column| { + for (0..vectors) |j| { + const from = (column * vectors + j) * vector; + const to = (((column / columns) * vectors + j) * columns + column % columns) * vector; + @memcpy(dst[to..][0..vector], src[from..][0..vector]); + } + } +} + +/// Scales or biases (rows of k/64) into [tiles][k/64 groups][32 lanes]; src is n*k/32 bytes, dst padded(n)*k/32. +pub fn tileParams(dst: []u8, src: []const u8, n: usize, k: usize) void { + std.debug.assert(fits(n, k) and src.len == n * k / 32 and dst.len == padded(n) * k / 32); + @memset(dst, 0); + const groups = k / group; + for (0..n) |column| { + for (0..groups) |g| { + const from = (column * groups + g) * 2; + const to = (((column / columns) * groups + g) * columns + column % columns) * 2; + @memcpy(dst[to..][0..2], src[from..][0..2]); + } + } +} + +const testing = std.testing; + +fn byteSum(bytes: []const u8) u64 { + var sum: u64 = 0; + for (bytes) |b| sum += b; + return sum; +} + +test "fits needs whole groups of k, and padded rounds n up to whole tiles" { + try testing.expect(fits(32, 64)); + try testing.expect(fits(16, 2048)); + try testing.expect(fits(248320, 2048)); + try testing.expect(!fits(32, 96)); + try testing.expect(!fits(0, 64)); + try testing.expectEqual(@as(usize, 32), padded(1)); + try testing.expectEqual(@as(usize, 32), padded(32)); + try testing.expectEqual(@as(usize, 64), padded(33)); + try testing.expectEqual(@as(usize, 248320), padded(248320)); +} + +// The shader reads column c, vector j at ((c / 32) * (k / 32) + j) * 32 + c % 32 vectors; scales likewise per group. +test "a tiled vector is the row's vector at the address the shader computes, padding is zero, and nothing is lost" { + const n = 70; + const k = 192; + var src: [n * k / 2]u8 = undefined; + for (&src, 0..) |*byte, i| byte.* = @truncate((i *% 2654435761 >> 7) | 1); + var dst: [padded(n) * k / 2]u8 = undefined; + tileWords(&dst, &src, n, k); + for (0..n) |c| { + for (0..k / 32) |j| { + const at = (((c / 32) * (k / 32) + j) * 32 + c % 32) * 16; + try testing.expectEqualSlices(u8, src[(c * (k / 32) + j) * 16 ..][0..16], dst[at..][0..16]); + } + } + try testing.expectEqual(byteSum(&src), byteSum(&dst)); + for (n..padded(n)) |c| { + for (0..k / 32) |j| { + const at = (((c / 32) * (k / 32) + j) * 32 + c % 32) * 16; + try testing.expectEqual(@as(u64, 0), byteSum(dst[at..][0..16])); + } + } +} + +test "tiled scales are the row's scale at the address the shader computes, padding is zero, and nothing is lost" { + const n = 40; + const k = 256; + var src: [n * k / 32]u8 = undefined; + for (&src, 0..) |*byte, i| byte.* = @truncate((i *% 40503 >> 3) | 1); + var dst: [padded(n) * k / 32]u8 = undefined; + tileParams(&dst, &src, n, k); + for (0..n) |c| { + for (0..k / 64) |g| { + const at = (((c / 32) * (k / 64) + g) * 32 + c % 32) * 2; + try testing.expectEqualSlices(u8, src[(c * (k / 64) + g) * 2 ..][0..2], dst[at..][0..2]); + } + } + try testing.expectEqual(byteSum(&src), byteSum(&dst)); +} diff --git a/zig/src/vulkan_model_host_tests.zig b/zig/src/vulkan_model_host_tests.zig new file mode 100644 index 00000000000..c22c615a0bf --- /dev/null +++ b/zig/src/vulkan_model_host_tests.zig @@ -0,0 +1,10 @@ +const std = @import("std"); +const weights = @import("families/qwen3_5/backend/vulkan/weights.zig"); +const state = @import("families/qwen3_5/backend/vulkan/state.zig"); + +test { + std.testing.refAllDecls(weights); + std.testing.refAllDecls(weights.Model); + std.testing.refAllDecls(state); + std.testing.refAllDecls(state.State); +} diff --git a/zig/tests/vulkan/affine_fixture.zig b/zig/tests/vulkan/affine_fixture.zig index 0ed78b6ee38..6f241d5356b 100644 --- a/zig/tests/vulkan/affine_fixture.zig +++ b/zig/tests/vulkan/affine_fixture.zig @@ -1,5 +1,6 @@ //! MLX affine 4-bit, group-64, transposed inputs for Vulkan qmm.comp's qmm_f32 build, checked against float64. const std = @import("std"); +const stage1_params = @import("stage1_params"); const contract = @import("contract"); const c = contract.spec; const Affine = contract.affine.Affine; @@ -13,37 +14,8 @@ pub const poison: u32 = 0x7fc5_a5a5; /// Poisoned elements after each binding's data, so an overrun reads poison rather than past the buffer. pub const tail = 16; -/// qmm.comp's push constants (std430); qmm reads the fields params() sets, the rest keep MLX's ComputeParams defaults. -pub const Params = extern struct { - count: u32 = 0, - operation: u32 = 0, - lhs_size: u32 = 0, - rhs_size: u32 = 0, - reduce_size: u32 = 0, - output_size: u32 = 0, - lhs_offset: u32 = 0, - rhs_offset: u32 = 0, - output_offset: u32 = 0, - aux_size: u32 = 0, - aux_offset: u32 = 0, - matrix_m: u32 = 0, - matrix_n: u32 = 0, - matrix_k: u32 = 0, - flags: u32 = 0, - alpha: f32 = 1, - beta: f32 = 0, - dims: u32 = 0, - shape: [4]u32 = .{ 0, 0, 0, 0 }, - in_strides: [4]u32 = .{ 0, 0, 0, 0 }, - out_strides: [4]u32 = .{ 0, 0, 0, 0 }, -}; - -comptime { - const names = [_][]const u8{ "count", "operation", "lhs_size", "rhs_size", "reduce_size", "output_size", "lhs_offset", "rhs_offset", "output_offset", "aux_size", "aux_offset", "matrix_m", "matrix_n", "matrix_k", "flags", "alpha", "beta", "dims" }; - for (names, 0..) |name, i| if (@offsetOf(Params, name) != 4 * i) @compileError("qmm.comp push constant moved: " ++ name); - if (@offsetOf(Params, "shape") != 72 or @offsetOf(Params, "in_strides") != 88 or @offsetOf(Params, "out_strides") != 104) @compileError("qmm.comp push constant arrays moved"); - if (@sizeOf(Params) != 120) @compileError("qmm.comp push constants are 120 bytes"); -} +/// qmm.comp's push constants (std430), owned by the Vulkan params module; the layout asserts live there. +pub const Params = stage1_params.Params; /// Base items of each binding's data: lhs_offset, rhs_offset (words), aux_offset, aux_size and output_offset. pub const Offsets = struct { x: usize = 0, w: usize = 0, scales: usize = 0, biases: usize = 0, out: usize = 0 }; diff --git a/zig/tests/vulkan/contract_attention.zig b/zig/tests/vulkan/contract_attention.zig new file mode 100644 index 00000000000..04d51ede22b --- /dev/null +++ b/zig/tests/vulkan/contract_attention.zig @@ -0,0 +1,337 @@ +// Host-only rope/SDPA contract vs f64; fixture caps at 8 keys; shader SDPA uses 256 lanes and 7168-key limit. +const std = @import("std"); +const contract = @import("contract"); +const c = contract.spec; +pub const fx = @import("stage3_fixture.zig"); +const sources = @import("attention_sources"); + +pub const Report = fx.Report; + +// The reduction geometry the depth derivation names, pinned to the shader source by the tests below. +pub const lanes: u64 = 256; // softmax threads: s_tree's width +pub const tree_levels: u64 = 8; // log2(lanes) binary-tree levels +pub const dot_terms: u64 = fx.head_dim; // 256 ascending-d score FMAs +pub const key_tile: u64 = 16; // 16-key output tiles / load batches +pub const key_stride: u64 = 1024; // threads striding the score scan +pub const max_visible_keys: u64 = 7168; // kStreamKeys: the shared-memory stream + +// For n keys, denominator depth is ceil(n/256)+8: 256 lanes, 8 tree levels, and 36 at n=7168. +pub fn softmaxDepth(n: u64) u64 { + return (n + lanes - 1) / lanes + tree_levels; +} + +/// The output accumulation's depth: sequential over the n visible keys. +pub fn outputDepth(n: u64) u64 { + return n; +} + +fn rel(depth: f64) f64 { + const u: f64 = 0x1p-24; + return depth * u / (1 - depth * u); +} + +/// The reference walk's magnitudes for one case, at the shader's own scales. +pub const Mags = struct { + /// max over keys of sum_d |alpha*q_d*k_d|: the score dot's term scale. + terms: f64 = 0, + /// max |score|: exp's argument is |s - max| <= 2*score. + score: f64 = 0, + /// max |value|: probs sum to 1, so the output term scale is this. + value: f64 = 0, +}; + +// Bound each output by per-stage depths/scales, two f32 exp ulps, and one RNE bf16 store (2^-8). +pub fn sdpaBound(n: u64, m: Mags, want: f64) f64 { + const u: f64 = 0x1p-24; + const dot = rel(dot_terms) * m.terms; // score dot depth 256 + const arg = dot + u * 2 * m.score; // s - max: subtraction at |arg| <= 2S + const e_rel = arg + 2 * u; // exp's own f32 error, two ops + const den_rel = rel(@floatFromInt(softmaxDepth(n))) + e_rel; // chain + tree, NOT n + const prob_rel = e_rel + den_rel + u; // p = e * (1/D): one multiply + const work = prob_rel * m.value + rel(@floatFromInt(outputDepth(n))) * m.value; + return work + 0x1p-8 * (@abs(want) + work) + 0x1p-133; +} + +pub const SdpaCase = struct { name: []const u8, rows: u64, prefix: u64 }; + +// Capacity covers prefix + rows and rejects totals above max_visible_keys. +pub fn capacityOf(case: SdpaCase) u64 { + return case.prefix + case.rows; +} + +/// One row's visible key count: the causal bound, prefix + r + 1. +pub fn visibleKeys(case: SdpaCase, r: u64) u64 { + return case.prefix + r + 1; +} + +// Cases test n=1,15/16/17,255/256/257,1023/1024/1025,7167/7168 and rows 17/32 at 1024/7168 limits. +pub const served_cases = [_]SdpaCase{ + .{ .name = "decode n=1", .rows = 1, .prefix = 0 }, + .{ .name = "decode n=16 tile", .rows = 1, .prefix = 15 }, + .{ .name = "decode n=17 tile+1", .rows = 1, .prefix = 16 }, + .{ .name = "decode n=255 scan", .rows = 1, .prefix = 254 }, + .{ .name = "decode n=256 scan", .rows = 1, .prefix = 255 }, + .{ .name = "decode n=257 scan+1", .rows = 1, .prefix = 256 }, + .{ .name = "decode n=1024 stride", .rows = 1, .prefix = 1023 }, + .{ .name = "decode n=1025 stride+1", .rows = 1, .prefix = 1024 }, + .{ .name = "decode n=7167 stream-1", .rows = 1, .prefix = 7166 }, + .{ .name = "decode n=7168 stream", .rows = 1, .prefix = 7167 }, + .{ .name = "prefill rows32 n1..32", .rows = 32, .prefix = 0 }, + .{ .name = "prefill rows17 across 1024", .rows = 17, .prefix = 1007 }, + .{ .name = "prefill rows32 to 7168", .rows = 32, .prefix = 7136 }, +}; + +// Attend f64 dots/max/softmax/output for n=prefix+r+1 in (kv_head,capacity,head_dim); matches fx.Attend.init order. +pub fn attendRef( + a: std.mem.Allocator, + q_bits: []const u16, + keys: []const u16, + values: []const u16, + case: SdpaCase, + capacity: u64, + mags: *Mags, + want: []f64, + visible: []usize, +) !void { + const row_elems: usize = @intCast(fx.heads * fx.head_dim); + const cap: usize = @intCast(capacity); + const kh_span: usize = @intCast(fx.heads / fx.kv_heads); // repeats: 8/2 + // Per-(row, head) score scratch; capacity <= 7168 (attendPlan's bound). + const scores = try a.alloc(f64, cap); + defer a.free(scores); + for (0..@intCast(case.rows)) |r| { + const n: usize = @intCast(visibleKeys(case, @intCast(r))); + visible[r] = n; + for (0..@intCast(fx.heads)) |h| { + const kh = h / kh_span; + const qrow = q_bits[r * row_elems + h * @as(usize, @intCast(fx.head_dim)) ..][0..@intCast(fx.head_dim)]; + var m: f64 = -std.math.floatMax(f64); + for (0..n) |j| { + const kbase = (kh * cap + j) * @as(usize, @intCast(fx.head_dim)); + var dot: f64 = 0; + var terms: f64 = 0; + for (0..@intCast(fx.head_dim)) |d| { + const qv = c.fromBf16(qrow[d]); + const kv = c.fromBf16(keys[kbase + d]); + dot += qv * kv; + terms += @abs(qv) * @abs(kv); + } + scores[j] = dot * fx.scale; + mags.terms = @max(mags.terms, terms * fx.scale); + m = @max(m, scores[j]); + mags.score = @max(mags.score, @abs(scores[j])); + } + var denom: f64 = 0; + for (scores[0..n]) |s| denom += @exp(s - m); + for (0..@intCast(fx.head_dim)) |d| { + var acc: f64 = 0; + for (0..n) |j| { + const vbase = (kh * cap + j) * @as(usize, @intCast(fx.head_dim)); + acc += (@exp(scores[j] - m) / denom) * c.fromBf16(values[vbase + d]); + } + want[r * row_elems + h * @as(usize, @intCast(fx.head_dim)) + d] = acc; + } + } + } +} + +// Causal SDPA compares GPU bits to f64 references; n=1 preserves V bits, and finite wants require finite outputs. +pub const Sdpa = struct { + a: std.mem.Allocator, + case: SdpaCase, + capacity: u64, + q_bits: []u16, + k_bits: []u16, + v_bits: []u16, + want: []f64, + visible: []usize, + mags: Mags, + + pub fn init(a: std.mem.Allocator, case: SdpaCase, seed: u64) !Sdpa { + const capacity = capacityOf(case); + const q_elems: usize = @intCast(case.rows * fx.heads * fx.head_dim); + const cache_elems: usize = @intCast(fx.kv_heads * capacity * fx.head_dim); + const q_bits = try a.alloc(u16, q_elems); + errdefer a.free(q_bits); + const k_bits = try a.alloc(u16, cache_elems); + errdefer a.free(k_bits); + const v_bits = try a.alloc(u16, cache_elems); + errdefer a.free(v_bits); + const want = try a.alloc(f64, q_elems); + errdefer a.free(want); + const visible = try a.alloc(usize, @intCast(case.rows)); + errdefer a.free(visible); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (q_bits) |*x| x.* = fx.sampleDyadic(rng, -5, 3); + for (k_bits) |*x| x.* = fx.sampleDyadic(rng, -5, 3); + for (v_bits) |*x| x.* = fx.sampleDyadic(rng, -5, 3); + var f: Sdpa = .{ + .a = a, + .case = case, + .capacity = capacity, + .q_bits = q_bits, + .k_bits = k_bits, + .v_bits = v_bits, + .want = want, + .visible = visible, + .mags = .{}, + }; + for (f.v_bits) |bits| f.mags.value = @max(f.mags.value, @abs(c.fromBf16(bits))); + try attendRef(a, f.q_bits, f.k_bits, f.v_bits, case, capacity, &f.mags, f.want, f.visible); + return f; + } + + pub fn deinit(f: Sdpa) void { + f.a.free(f.q_bits); + f.a.free(f.k_bits); + f.a.free(f.v_bits); + f.a.free(f.want); + f.a.free(f.visible); + } + + // Check outputs against row bounds; single-key rows preserve V bits, and finite references require finite outputs. + pub fn check(f: Sdpa, bits: []const u16) !Report { + var r: Report = .{}; + const row_elems: usize = @intCast(fx.heads * fx.head_dim); + const kh_span: usize = @intCast(fx.heads / fx.kv_heads); + for (0..@intCast(f.case.rows)) |row| { + const n = f.visible[row]; + for (0..@intCast(fx.heads)) |h| { + for (0..@intCast(fx.head_dim)) |d| { + const i = row * row_elems + h * @as(usize, @intCast(fx.head_dim)) + d; + const got = c.fromBf16(bits[i]); + if (!std.math.isFinite(got)) return error.NonFiniteOutput; + const w = f.want[i]; + const bound = sdpaBound(n, f.mags, w); + const dist = @abs(got - w); + if (dist > bound) return error.OverBound; + const ratio = dist / bound; + if (ratio > r.worst) r = .{ .worst = ratio, .where = i }; + if (n == 1) { + const vb = f.v_bits[(h / kh_span) * @as(usize, @intCast(f.capacity)) * @as(usize, @intCast(fx.head_dim)) + d]; + if (bits[i] != vb) return error.SingleKeyNotTransport; + } + } + } + } + return r; + } +}; + +const testing = std.testing; + +test "softmax depth is the shader's chain plus tree, never the plain key sum" { + // One strided pass per lane + 8 tree levels at and under one full scan. + try testing.expectEqual(@as(u64, 9), softmaxDepth(1)); + try testing.expectEqual(@as(u64, 9), softmaxDepth(255)); + try testing.expectEqual(@as(u64, 9), softmaxDepth(256)); + try testing.expectEqual(@as(u64, 10), softmaxDepth(257)); + try testing.expectEqual(@as(u64, 12), softmaxDepth(1024)); + try testing.expectEqual(@as(u64, 13), softmaxDepth(1025)); + try testing.expectEqual(@as(u64, 36), softmaxDepth(7168)); + // The point of the derivation: 36, not 7168. + try testing.expect(softmaxDepth(max_visible_keys) < max_visible_keys / 100); +} + +test "sdpaBound covers a planted wrong finite and grows with the honest depths" { + const m: Mags = .{ .terms = 128, .score = 64, .value = 8 }; + const b1 = sdpaBound(1, m, 1); + try testing.expect(b1 > 0); + try testing.expect(sdpaBound(257, m, 1) > sdpaBound(256, m, 1)); // chain edge + try testing.expect(sdpaBound(7168, m, 1) > sdpaBound(1025, m, 1)); // output depth + try testing.expect(sdpaBound(7168, m, 1) < 1 << 20); // 2^31 plants are refused + try testing.expect(sdpaBound(1, .{}, 0) > 0); // finite on zeros +} + +test "generalized reference equals the fixture's Attend element for element" { + const a = testing.allocator; + // Same shape family the fixture's [8]f64 score buffer allows. + const case = SdpaCase{ .name = "fixture parity", .rows = 2, .prefix = 1 }; + const capacity = capacityOf(case); + const row_elems: usize = @intCast(fx.heads * fx.head_dim); + const cache_elems: usize = @intCast(fx.kv_heads * capacity * fx.head_dim); + var prng = std.Random.DefaultPrng.init(0xc0ffee); + const rng = prng.random(); + const q = try a.alloc(u16, case.rows * row_elems); + defer a.free(q); + const k = try a.alloc(u16, cache_elems); + defer a.free(k); + const v = try a.alloc(u16, cache_elems); + defer a.free(v); + for (q) |*x| x.* = fx.sampleDyadic(rng, -5, 3); + for (k) |*x| x.* = fx.sampleDyadic(rng, -5, 3); + for (v) |*x| x.* = fx.sampleDyadic(rng, -5, 3); + const fixture = try fx.Attend.init(a, .{ .name = case.name, .rows = case.rows, .prefix = case.prefix, .capacity = capacity }, q, k, v, 0); + defer fixture.deinit(); + const want = try a.alloc(f64, case.rows * row_elems); + defer a.free(want); + const visible = try a.alloc(usize, @intCast(case.rows)); + defer a.free(visible); + var mags: Mags = .{}; + try attendRef(a, q, k, v, case, capacity, &mags, want, visible); + try testing.expectEqualSlices(f64, fixture.want, want); + for (0..case.rows) |r| try testing.expectEqual(@as(usize, @intCast(visibleKeys(case, @intCast(r)))), visible[r]); +} + +test "served cases stay inside the engine's shapes and hit every reduction edge" { + for (served_cases) |case| { + try testing.expect(case.rows == 1 or case.rows <= 32); // decode rows 1, prefill <= 32 + const capacity = capacityOf(case); + try testing.expect(capacity <= max_visible_keys); // kStreamKeys + try testing.expectEqual(visibleKeys(case, case.rows - 1), capacity); // exact fit + try testing.expect(visibleKeys(case, 0) >= 1); + } + // The named edges all appear as some decode case's single-row count. + const decode_edges = [_]u64{ 1, 16, 17, 256, 257, 1024, 1025, 7168 }; + for (decode_edges) |n| { + var found = false; + for (served_cases) |case| { + if (case.rows == 1 and visibleKeys(case, 0) == n) found = true; + } + try testing.expect(found); // every claimed edge is actually served + } + // Multi-row prefill crosses the 1024 stride and fills 7168. + try testing.expect(served_cases[11].prefix < key_stride and capacityOf(served_cases[11]) >= key_stride); + try testing.expectEqual(max_visible_keys, capacityOf(served_cases[served_cases.len - 1])); +} + +test "sdpa checkers accept the honest bf16 output and refuse planted ones" { + const a = testing.allocator; + const case = SdpaCase{ .name = "plants", .rows = 2, .prefix = 0 }; // row 0 sees exactly 1 key + const f = try Sdpa.init(a, case, 0x5d1); + defer f.deinit(); + const good = try a.alloc(u16, f.want.len); + defer a.free(good); + for (f.want, 0..) |w, i| good[i] = c.toBf16(@floatCast(w)); + const rep = try f.check(good); // the honest rounded output passes + try testing.expect(rep.worst <= 1); + // single-key row: bit-exact V transport, and tampering with it is caught + try testing.expectEqual(f.v_bits[0], good[0]); // row 0, head 0, d 0: n=1 + good[0] ^= 1; + try testing.expectError(error.SingleKeyNotTransport, f.check(good)); + good[0] = f.v_bits[0]; + good[1] = 0x4f00; // 2^31: an absolutely wrong finite past any bound + try testing.expectError(error.OverBound, f.check(good)); + good[1] = 0x7fc0; + try testing.expectError(error.NonFiniteOutput, f.check(good)); + // an all-poison output (dead kernel) is refused + const dead = try a.alloc(u16, f.want.len); + defer a.free(dead); + @memset(dead, fx.poison16); + try testing.expectError(error.NonFiniteOutput, f.check(dead)); +} + +test "depth derivation is pinned to the shader source this contract reads" { + const src = sources.sdpa_decode_native; + // 256-lane scan and the 8-level binary tree: the denominator's depth. + try testing.expect(std.mem.indexOf(u8, src, "if (thread < 256u)") != null); + try testing.expect(std.mem.indexOf(u8, src, "index = thread; index < key_end; index += 256u") != null); + try testing.expect(std.mem.indexOf(u8, src, "for (uint stride = 128u; stride > 0u; stride >>= 1u)") != null); + // The output accumulation's depth: sequential over the visible keys. + try testing.expect(std.mem.indexOf(u8, src, "key + 16u <= key_end; key += 16u") != null); + // The causal multi-row bound the served shapes ride. + try testing.expect(std.mem.indexOf(u8, src, "params.matrix_k - params.dims + row + 1u") != null); + // The shared-memory bound that caps the served context. + try testing.expect(std.mem.indexOf(u8, src, "kStreamKeys = 7168u") != null); +} diff --git a/zig/tests/vulkan/contract_attention_main.zig b/zig/tests/vulkan/contract_attention_main.zig new file mode 100644 index 00000000000..398db75af6b --- /dev/null +++ b/zig/tests/vulkan/contract_attention_main.zig @@ -0,0 +1,305 @@ +// The attention gate checks rope_bf16 and sdpa_rows_bf16: decode=1, prefill<=32, keys=1..7168; exits 0/1/2/124. +const std = @import("std"); +const vk = @import("vulkan"); +const shaders = @import("stage3_shaders"); +const ops_shaders = @import("stage1_shaders"); +const fx = @import("stage3_fixture.zig"); +const ca = @import("contract_attention.zig"); +const contract = @import("contract"); +const cx = contract.spec; + +const usage = "usage: tf-attention-contract [--vendor ID] [--timeout-ms MS] | --help\n" ++ + " --help show this text and exit without opening a loader or device\n" ++ + " --vendor ID refuse a device whose Vulkan vendorID differs (decimal or 0x hex; no default)\n" ++ + " --timeout-ms MS bound on every fence wait, 1 to 600000 (default 10000); a timeout ends the process\n"; + +const Elapsed = struct { + io: std.Io, + start: std.Io.Timestamp, + + fn ms(e: Elapsed) i64 { + return e.start.durationTo(std.Io.Clock.awake.now(e.io)).toMilliseconds(); + } +}; + +const repeats = 3; +// Large FINITE bf16 planted over non-visible keys: an uncausal read turns the output huge and fails every check. +const finite_poison: u16 = 0x7f7f; +/// 1 MiB: the 7168-key cache buffers (7.3 MB each) upload in few chunks. +const staging_bytes = 1 << 20; + +pub fn main(init: std.process.Init) !u8 { + const clock: Elapsed = .{ .io = init.io, .start = std.Io.Clock.awake.now(init.io) }; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + for (args[1..]) |arg| if (std.mem.eql(u8, arg, "--help")) { + std.debug.print(usage, .{}); + return 0; + }; + var vendor: ?u32 = null; + var timeout_ms: u64 = 10_000; + var i: usize = 1; + while (i < args.len) : (i += 2) { + const value = if (i + 1 < args.len) args[i + 1] else ""; + if (std.mem.eql(u8, args[i], "--vendor")) { + vendor = std.fmt.parseInt(u32, value, 0) catch null; + if (vendor == null) break; + } else if (std.mem.eql(u8, args[i], "--timeout-ms")) { + timeout_ms = std.fmt.parseInt(u64, value, 10) catch 0; + if (timeout_ms == 0 or timeout_ms > 600_000) break; + } else break; + } else return if (try probe(init.gpa, clock, vendor, timeout_ms * std.time.ns_per_ms)) 0 else 1; + std.debug.print(usage, .{}); + return 2; +} + +fn probe(a: std.mem.Allocator, clock: Elapsed, vendor: ?u32, timeout: u64) !bool { + var rt = try vk.Runtime.open(); + defer rt.close(); + var dev = try vk.Device.open(&rt, .{ .features = .{ .storage16 = true, .int16 = true } }); + defer dev.close(); + const pr = dev.properties; + const name = std.mem.sliceTo(&pr.deviceName, 0); + const v = pr.apiVersion; + std.debug.print("MILESTONE device-found elapsed_ms={d} name=\"{s}\" vendor=0x{x:0>4} device=0x{x:0>4} " ++ + "api={d}.{d}.{d} storage16={} int16={} shared={d} fence-timeout={d}ms\n", .{ + clock.ms(), name, pr.vendorID, pr.deviceID, + v >> 22 & 0x7f, v >> 12 & 0x3ff, v & 0xfff, dev.features.storage16, + dev.features.int16, pr.limits.maxComputeSharedMemorySize, timeout / std.time.ns_per_ms, + }); + if (pr.deviceType == vk.c.VK_PHYSICAL_DEVICE_TYPE_CPU) { + std.debug.print("FAIL device: a CPU (software) Vulkan device is never a GPU result\n", .{}); + return false; + } + if (vendor) |want| if (pr.vendorID != want) { + std.debug.print("FAIL device: vendor 0x{x:0>4} is not the requested 0x{x:0>4}\n", .{ pr.vendorID, want }); + return false; + }; + + var cmds = try vk.Commands.init(&dev); + defer cmds.deinit(); + var staging = try vk.Buffer.init(&dev, .staging, staging_bytes); + defer staging.deinit(); + var ops = try vk.ops.Ops.init(&dev, .{ + .qmm_bf16 = &ops_shaders.qmm_bf16, + .qmm_q4_vec4 = &ops_shaders.qmm_q4_vec4, + .qmm_f32 = &ops_shaders.qmm_f32, + .copy_f32_bf16 = &ops_shaders.copy_f32_bf16, + .qmm_q4_tile = &ops_shaders.qmm_q4_tile, + .qmm_q4_tile_f32 = &ops_shaders.qmm_q4_tile_f32, + .dequant_bf16 = &ops_shaders.dequant_bf16, + .gather_u32 = &ops_shaders.gather_u32, + .gather_u16 = &ops_shaders.gather_u16, + .cast_bf16_f32 = &ops_shaders.cast_bf16_f32, + .cast_f32_bf16 = &ops_shaders.cast_f32_bf16, + .copy_u32 = &ops_shaders.copy_u32, + .copy_u16 = &ops_shaders.copy_u16, + }); + defer ops.deinit(); + var attn = try vk.Attention.init(&dev, .{ + .rope_bf16 = &shaders.rope_bf16, + .sdpa_rows_bf16 = &shaders.sdpa_rows_bf16, + .gate_bf16 = &shaders.gate_bf16, + .rope_f32 = &shaders.rope_f32, + .rope_f32_kcache = &shaders.rope_f32_kcache, + .sdpa_rows_f32q = &shaders.sdpa_rows_f32q, + .gate_f32 = &shaders.gate_f32, + }, &ops); + defer attn.deinit(); + + var passed: usize = 0; + var total: usize = 0; + for (fx.rope_cases) |case| { + total += 1; + passed += @intFromBool(try ropeGate(a, clock, &dev, &cmds, staging, &attn, case, timeout)); + } + for (ca.served_cases, 0..) |case, i| { + total += 1; + passed += @intFromBool(try sdpaGate(a, clock, &dev, &cmds, staging, &attn, case, 0x5dcafe + @as(u64, i), timeout)); + } + std.debug.print("RESULT attention-contract elapsed_ms={d} on \"{s}\": passed={d} total={d}\n", .{ clock.ms(), name, passed, total }); + return passed == total; +} + +fn upload16(dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, shell: []const u16, timeout: u64) !vk.Buffer { + var b = try vk.Buffer.init(dev, .device, std.mem.sliceAsBytes(shell).len); + errdefer b.deinit(); + try vk.orExit(vk.upload(cmds, b, 0, staging, std.mem.sliceAsBytes(shell), timeout)); + return b; +} + +fn fetch(cmds: *vk.Commands, staging: vk.Buffer, b: vk.Buffer, out: []u8, timeout: u64) !void { + try vk.orExit(vk.download(cmds, b, 0, staging, out, timeout)); +} + +fn unchanged(a: std.mem.Allocator, cmds: *vk.Commands, staging: vk.Buffer, shell: []const u16, buf: vk.Buffer, timeout: u64) !bool { + const back = try a.alloc(u16, shell.len); + defer a.free(back); + try fetch(cmds, staging, buf, std.mem.sliceAsBytes(back), timeout); + return std.mem.indexOfDiff(u16, shell, back) == null; +} + +// Every plant must fail: downloaded GOOD output's first finite-want element gets wrong finite, then nan/+inf/-inf. +fn plantRefuse(f: anytype, good: []const u16, name: []const u8, gate: []const u8) !bool { + const mutated = try f.a.dupe(u16, good); + defer f.a.free(mutated); + var first_finite: ?usize = null; + for (good, 0..) |bits, i| { + if (std.math.isFinite(cx.fromBf16(bits))) { + first_finite = i; + break; + } + } + const at = first_finite orelse return error.NoFiniteReference; + mutated[at] = 0x4f00; // 2^31: an absolutely wrong finite the bounds must refuse + if (f.check(mutated)) |_| { + std.debug.print("FAIL {s} case=\"{s}\": a planted wrong finite output passed\n", .{ gate, name }); + return false; + } else |_| {} + for ([_]u16{ 0x7fc0, 0x7f80, 0xff80 }, [_][]const u8{ "nan", "+inf", "-inf" }) |pb, pname| { + mutated[at] = pb; + if (f.check(mutated)) |_| { + std.debug.print("FAIL {s} case=\"{s}\": a planted {s} output passed\n", .{ gate, name, pname }); + return false; + } else |_| {} + } + return true; +} + +// Rope gates dispatch three byte-identical repeats, check stage3_fixture f64, refuse plants, and leave input unchanged. +fn ropeGate(a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, attn: *vk.Attention, case: fx.RopeCase, timeout: u64) !bool { + const f = try fx.Rope.init(a, case, 0x57a3); + defer f.deinit(); + var q = try upload16(dev, cmds, staging, f.q_bits, timeout); + defer q.deinit(); + // Poison every output word so missing writes cannot inherit a reference. + const out_shell = try a.alloc(u16, f.want.len + fx.tail); + defer a.free(out_shell); + @memset(out_shell, fx.poison16); + var out = try upload16(dev, cmds, staging, out_shell, timeout); + defer out.deinit(); + const out_view = try vk.tensor.View.init(&out, .bf16, 0, f.want.len); + const q_view = try vk.tensor.View.init(&q, .bf16, 0, f.want.len); + try vk.orExit(attn.setBasePosition(cmds, staging, case.base, timeout)); + const rows: u64 = case.rows; + var first: ?[]u16 = null; + defer if (first) |x| a.free(x); + const full_bits = try a.alloc(u16, f.want.len + fx.tail); + defer a.free(full_bits); + const bits = full_bits[0..f.want.len]; + for (0..repeats) |rep| { + // re-upload the all-poison shell before every dispatch: the kernel must write every output word itself. + try vk.orExit(vk.upload(cmds, out, 0, staging, std.mem.sliceAsBytes(out_shell), timeout)); + try vk.orExit(attn.ropeQueries(cmds, out_view, q_view, rows, timeout)); + try fetch(cmds, staging, out, std.mem.sliceAsBytes(full_bits), timeout); + for (full_bits[f.want.len..]) |g_elem| { + if (g_elem != fx.poison16) { + std.debug.print("FAIL rope case=\"{s}\": output tail overwritten at element {d}\n", .{ case.name, f.want.len + (std.mem.indexOf(u16, full_bits[f.want.len..], &[_]u16{g_elem}) orelse 0) }); + return false; + } + } + if (first) |x| { + if (std.mem.indexOfDiff(u16, x, bits)) |at| { + std.debug.print("FAIL rope case=\"{s}\": repeat {d} differs at element {d}\n", .{ case.name, rep, at }); + return false; + } + } else first = try a.dupe(u16, bits); + } + const rep = f.check(bits) catch |e| { + std.debug.print("FAIL rope case=\"{s}\": {s}\n", .{ case.name, @errorName(e) }); + return false; + }; + if (!try plantRefuse(f, bits, case.name, "rope")) return false; + if (!try unchanged(a, cmds, staging, f.q_bits, q, timeout)) { + std.debug.print("FAIL rope case=\"{s}\": the input changed\n", .{case.name}); + return false; + } + std.debug.print("MILESTONE rope elapsed_ms={d} case=\"{s}\" rows={d} base={d} worst={e} of bound, repeats byte-identical, " ++ + "passthrough tail bit-exact, input unchanged, planted outputs refused\n", .{ clock.ms(), case.name, rows, case.base, rep.worst }); + return true; +} + +// SDPA: poison shell, 3 identical repeats, uploaded-bit f64 bounds, single-key exactness, plants refused, unchanged +fn sdpaGate(a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, attn: *vk.Attention, case: ca.SdpaCase, seed: u64, timeout: u64) !bool { + const row_elems: u64 = fx.heads * fx.head_dim; + const f = try ca.Sdpa.init(a, case, seed); + defer f.deinit(); + var q = try upload16(dev, cmds, staging, f.q_bits, timeout); + defer q.deinit(); + var keys = try upload16(dev, cmds, staging, f.k_bits, timeout); + defer keys.deinit(); + var values = try upload16(dev, cmds, staging, f.v_bits, timeout); + defer values.deinit(); + // Poison every output word so missing writes cannot inherit a reference. + const out_len: usize = @intCast(case.rows * row_elems); + const out_shell = try a.alloc(u16, out_len + fx.tail); + defer a.free(out_shell); + @memset(out_shell, fx.poison16); + var out = try upload16(dev, cmds, staging, out_shell, timeout); + defer out.deinit(); + const q_view = try vk.tensor.View.init(&q, .bf16, 0, out_len); + const keys_view = try vk.tensor.View.init(&keys, .bf16, 0, @intCast(f.k_bits.len)); + const values_view = try vk.tensor.View.init(&values, .bf16, 0, @intCast(f.v_bits.len)); + const y_view = try vk.tensor.View.init(&out, .bf16, 0, out_len); + var first: ?[]u16 = null; + defer if (first) |x| a.free(x); + const full_bits = try a.alloc(u16, out_len + fx.tail); + defer a.free(full_bits); + const bits = full_bits[0..out_len]; + for (0..repeats) |rep| { + try vk.orExit(vk.upload(cmds, out, 0, staging, std.mem.sliceAsBytes(out_shell), timeout)); + try vk.orExit(attn.attend(cmds, y_view, q_view, keys_view, values_view, case.prefix, case.rows, f.capacity, timeout)); + try fetch(cmds, staging, out, std.mem.sliceAsBytes(full_bits), timeout); + for (full_bits[out_len..]) |g_elem| { + if (g_elem != fx.poison16) { + std.debug.print("FAIL sdpa case=\"{s}\": output tail overwritten at element {d}\n", .{ case.name, out_len + (std.mem.indexOf(u16, full_bits[out_len..], &[_]u16{g_elem}) orelse 0) }); + return false; + } + } + if (first) |x| { + if (std.mem.indexOfDiff(u16, x, bits)) |at| { + std.debug.print("FAIL sdpa case=\"{s}\": repeat {d} differs at element {d}\n", .{ case.name, rep, at }); + return false; + } + } else first = try a.dupe(u16, bits); + } + const rep = f.check(bits) catch |e| { + std.debug.print("FAIL sdpa case=\"{s}\": {s}\n", .{ case.name, @errorName(e) }); + return false; + }; + if (!try plantRefuse(f, bits, case.name, "sdpa")) return false; + if (!try unchanged(a, cmds, staging, f.q_bits, q, timeout) or + !try unchanged(a, cmds, staging, f.k_bits, keys, timeout) or + !try unchanged(a, cmds, staging, f.v_bits, values, timeout)) + { + std.debug.print("FAIL sdpa case=\"{s}\": an input changed\n", .{case.name}); + return false; + } + // Each row decoded alone at q_len 1 and prefix + r must bit-match its multi-row result (decode-rows-1 shape). + if (case.rows > 1) { + const one_len: usize = @intCast(row_elems); + const one_shell = try a.alloc(u16, one_len + fx.tail); + defer a.free(one_shell); + @memset(one_shell, fx.poison16); + var one = try upload16(dev, cmds, staging, one_shell, timeout); + defer one.deinit(); + const one_view = try vk.tensor.View.init(&one, .bf16, 0, one_len); + const one_bits = try a.alloc(u16, one_len); + defer a.free(one_bits); + for (0..@intCast(case.rows)) |r| { + const q_row = try q_view.slice(@intCast(r * row_elems), @intCast(row_elems)); + try vk.orExit(vk.upload(cmds, one, 0, staging, std.mem.sliceAsBytes(one_shell), timeout)); + try vk.orExit(attn.attend(cmds, one_view, q_row, keys_view, values_view, case.prefix + @as(u64, r), 1, f.capacity, timeout)); + try fetch(cmds, staging, one, std.mem.sliceAsBytes(one_bits), timeout); + if (std.mem.indexOfDiff(u16, one_bits, bits[@intCast(r * row_elems)..][0..one_len])) |at| { + std.debug.print("FAIL sdpa case=\"{s}\": row {d} decoded alone differs at element {d}\n", .{ case.name, r, at }); + return false; + } + } + } + std.debug.print("MILESTONE sdpa elapsed_ms={d} case=\"{s}\" rows={d} prefix={d} visible=[{d}..{d}] worst={e} of bound, " ++ + "repeats byte-identical, {s}planted outputs refused, inputs unchanged\n", .{ + clock.ms(), case.name, case.rows, case.prefix, f.visible[0], f.visible[f.visible.len - 1], rep.worst, + if (case.rows > 1) "rows decode alone bit-identical, " else "", + }); + return true; +} diff --git a/zig/tests/vulkan/contract_datamove.zig b/zig/tests/vulkan/contract_datamove.zig new file mode 100644 index 00000000000..d2818042f56 --- /dev/null +++ b/zig/tests/vulkan/contract_datamove.zig @@ -0,0 +1,582 @@ +// Stage 1: cast_f32_bf16/cast_bf16_f32, copy_u32/copy_u16, gather_u32/gather_u16, dequant_bf16 match reference bytes. +const std = @import("std"); +const vk = @import("vulkan"); +const contract = @import("contract"); +const fx = @import("stage1_fixture.zig"); + +const c = contract.spec; +const g = contract.gates; + +const rig = @import("contract_datamove_rig.zig"); +const Rig = rig.Rig; +const read16 = rig.read16; +const write16 = rig.write16; +const read32 = rig.read32; +const write32 = rig.write32; +const buf = rig.buf; +pub const guard = rig.guard; + +fn none(_: c.Shape) usize { + return 0; +} + +/// The contract's gather id scheme: [0] out of range, [1] == [2] (a repeated row), 0, then cycling rows. +fn gatherIds(ids: []u32, table_rows: usize) void { + for (ids, 0..) |*id, i| { + id.* = switch (i) { + 0 => @intCast(table_rows), // out of range: the shader's documented zero row + 1, 2 => @intCast(table_rows - 1), + 3 => 0, + else => @intCast(i % table_rows), + }; + } +} + +// Dequant widens bf16 and uses one fp32 multiply-add (q4 product exact, so one rounding either way), then contract RNE. +fn dequantBits(scale_bits: u16, q: u4, bias_bits: u16) u16 { + const s: f32 = c.fromBf16(scale_bits); + const b: f32 = c.fromBf16(bias_bits); + const v = s * @as(f32, @floatFromInt(q)) + b; + return c.toBf16(v); +} + +// Random bf16 scale/bias begins at 2^-113, preventing fp32 subnormals that FTZ devices flush but hosts retain. +fn randomBf16(rng: std.Random) u16 { + const sign: u16 = rng.int(u1); + const exp: u16 = 14 + rng.uintLessThan(u8, 113); + const mant: u16 = rng.int(u7); + return sign << 15 | exp << 7 | mant; +} + +// Casts use contract RNE for f32 -> bf16 (NaNs keep class) and exact bf16 -> f32 widening; both match reference bytes. +pub fn Cast(comptime to_bf16: bool) type { + return struct { + const Self = @This(); + const esz: usize = if (to_bf16) 2 else 4; + const in_kind: vk.tensor.Kind = if (to_bf16) .f32 else .bf16; + const out_kind: vk.tensor.Kind = if (to_bf16) .bf16 else .f32; + const out_elem: c.Elem = if (to_bf16) .bf16 else .f32; + const rule = if (to_bf16) + "outputs equal the contract's RNE bytes: ties to even, bf16 subnormals included, a NaN keeps its class" + else + "outputs equal the reference bits: bf16 widening is exact, every value and NaN byte copies"; + + pub const spec: c.Spec = .{ + .name = if (to_bf16) "vulkan cast_f32_bf16 (RNE)" else "vulkan cast_bf16_f32 (exact widening)", + .accum = .none, + .depth = none, + .out = out_elem, + .nan = .{ .stated = rule }, + .tiles = .{ .{}, .{ .size = 256, .step = 1, .min = 1 }, .{} }, + .served = &.{ .{ .rows = 1, .n = 146, .k = 0 }, .{ .rows = 3, .n = 700, .k = 0 } }, + }; + + rig: Rig, + o: *vk.ops.Ops, + arena: std.heap.ArenaAllocator, + in_at: usize = 0, + out_at: usize = 0, + elements: usize = 0, + want: []u8 = &.{}, + msg: [160]u8 = undefined, + + pub fn init(a: std.mem.Allocator, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, o: *vk.ops.Ops, timeout_ns: u64) Self { + return .{ .rig = .{ .a = a, .dev = dev, .cmds = cmds, .staging = staging, .timeout_ns = timeout_ns }, .o = o, .arena = .init(a) }; + } + + pub fn deinit(self: *Self) void { + self.rig.deinit(); + self.arena.deinit(); + } + + pub fn kernel(self: *Self) g.Kernel { + return .{ .spec = spec, .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Self); + + pub fn prepare(self: *Self, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0) return false; + const elements = std.math.mul(usize, shape.rows, shape.n) catch return false; + // Odd 16-bit counts are unserved because dispatch.bind needs 32-bit words and View.init rejects the tail. + if (elements % 2 != 0 and (in_kind.size() == 2 or esz == 2)) return false; + self.rig.reset(); + _ = self.arena.reset(.retain_capacity); + const a = self.arena.allocator(); + const in = try a.alloc(u8, elements * in_kind.size()); + const want = try a.alloc(u8, elements * esz); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (0..elements) |i| { + if (to_bf16) { + const x: f32 = if (i < fx.bf16_vectors.len) fx.bf16_vectors[i].in else (rng.float(f32) * 2 - 1) * std.math.pow(f32, 2, @floatFromInt(@as(i32, @intCast(i % 41)) - 20)); + write32(in, i, @bitCast(x)); + write16(want, i, c.toBf16(x)); + } else { + const bits: u16 = if (i < fx.bf16_vectors.len) fx.bf16_vectors[i].want else rng.int(u16); + write16(in, i, bits); + write32(want, i, @as(u32, bits) << 16); + } + } + if (fill == .nonfinite) { + const planted = [_]u32{ 0x7fc00000, 0x7f800000, 0xff800000 }; + for (planted, 0..) |bits, k| { + const i = (k * 37 + 11) % elements; + if (to_bf16) { + write32(in, i, bits); + write16(want, i, c.toBf16(@bitCast(bits))); + } else { + write16(in, i, @truncate(bits >> 16)); + write32(want, i, bits & 0xffff0000); + } + } + } + self.in_at = try self.rig.input(in); + self.out_at = try self.rig.output(want.len); + self.elements = elements; + self.want = want; + return true; + } + + pub fn run(self: *Self) !void { + const src = try self.rig.view(self.in_at, in_kind, self.elements); + const out = try self.rig.view(self.out_at, out_kind, self.elements); + if (to_bf16) { + try vk.orExit(self.o.castF32ToBf16(self.rig.cmds, out, src, self.rig.timeout_ns)); + } else { + try vk.orExit(self.o.castBf16ToF32(self.rig.cmds, out, src, self.rig.timeout_ns)); + } + } + + pub fn count(self: *Self) usize { + return self.elements; + } + + pub fn reference(self: *Self, want: []f64, mag: []f64) !void { + _ = self; + @memset(want, 0); + @memset(mag, 0); // never called: the check hook owns every comparison + } + + pub fn check(self: *Self, at: g.Gate) !?[]const u8 { + _ = at; + const got = try self.rig.outputBytes(); + for (0..self.elements) |i| { + const want_bits: u32 = if (to_bf16) read16(self.want, i) else read32(self.want, i); + const got_bits: u32 = if (to_bf16) read16(got, i) else read32(got, i); + if (got_bits != want_bits) return buf(&self.msg, "element {d} read 0x{x:0>8}, want 0x{x:0>8}", .{ i, got_bits, want_bits }); + } + return null; + } + }; +} + +/// A contiguous same-kind raw copy: exact bytes out, guards and source untouched. +pub fn Copy(comptime kind: vk.tensor.Kind) type { + return struct { + const Self = @This(); + const esz = kind.size(); + + pub const spec: c.Spec = .{ + .name = "vulkan copy_" ++ (if (kind == .u32) "u32" else "u16") ++ " (bit copy)", + .accum = .none, + .depth = none, + .out = .raw, + .nan = .{ .stated = "outputs equal the source bytes: a raw bit copy, every value and NaN pattern included" }, + .tiles = .{ .{}, .{ .size = 256, .step = 1, .min = 1 }, .{} }, + .served = &.{ .{ .rows = 1, .n = 300, .k = 0 }, .{ .rows = 2, .n = 520, .k = 0 } }, + }; + + rig: Rig, + o: *vk.ops.Ops, + arena: std.heap.ArenaAllocator, + in_at: usize = 0, + out_at: usize = 0, + elements: usize = 0, + want: []u8 = &.{}, + msg: [160]u8 = undefined, + + pub fn init(a: std.mem.Allocator, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, o: *vk.ops.Ops, timeout_ns: u64) Self { + return .{ .rig = .{ .a = a, .dev = dev, .cmds = cmds, .staging = staging, .timeout_ns = timeout_ns }, .o = o, .arena = .init(a) }; + } + + pub fn deinit(self: *Self) void { + self.rig.deinit(); + self.arena.deinit(); + } + + pub fn kernel(self: *Self) g.Kernel { + return .{ .spec = spec, .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Self); + + pub fn prepare(self: *Self, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0) return false; + const elements = std.math.mul(usize, shape.rows, shape.n) catch return false; + // See the casts: an odd 16-bit element count cannot be bound as whole 32-bit words. + if (elements % 2 != 0 and esz == 2) return false; + self.rig.reset(); + _ = self.arena.reset(.retain_capacity); + const a = self.arena.allocator(); + const data = try a.alloc(u8, elements * esz); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + rng.bytes(data); + if (fill == .nonfinite) { + const planted32 = [_]u32{ 0x7fc00000, 0x7f800000, 0xff800000 }; + const planted16 = [_]u16{ 0x7fc0, 0x7f80, 0xff80 }; + for (0..3) |k| { + const i = (k * 53 + 7) % elements; + if (esz == 4) write32(data, i, planted32[k]) else write16(data, i, planted16[k]); + } + } + self.in_at = try self.rig.input(data); + self.out_at = try self.rig.output(data.len); + self.elements = elements; + self.want = data; + return true; + } + + pub fn run(self: *Self) !void { + const src = try self.rig.view(self.in_at, kind, self.elements); + const dst = try self.rig.view(self.out_at, kind, self.elements); + try vk.orExit(self.o.copy(self.rig.cmds, dst, src, self.rig.timeout_ns)); + } + + pub fn count(self: *Self) usize { + return self.elements; + } + + pub fn reference(self: *Self, want: []f64, mag: []f64) !void { + _ = self; + @memset(want, 0); + @memset(mag, 0); // never called: the check hook owns every comparison + } + + pub fn check(self: *Self, at: g.Gate) !?[]const u8 { + _ = at; + const got = try self.rig.outputBytes(); + if (std.mem.indexOfDiff(u8, got, self.want)) |byte| { + return buf(&self.msg, "element {d} differs from the source bytes", .{byte / esz}); + } + return null; + } + }; +} + +// gather_take copies u32/u16 rows bit for bit; out-of-range ids gather zeros, never another row (ops_q4.zig). +pub fn Gather(comptime kind: vk.tensor.Kind) type { + return struct { + const Self = @This(); + const esz = kind.size(); + + pub const spec: c.Spec = .{ + .name = "vulkan gather_" ++ (if (kind == .u32) "u32" else "u16") ++ " (bit rows)", + .accum = .none, + .depth = none, + .out = .raw, + .nan = .{ .stated = "outputs equal the addressed row's bytes; an out-of-range id's row is zeros, never another row (ops_q4.zig)" }, + .tiles = .{ .{}, .{ .size = 256, .step = 1, .min = 1 }, .{} }, + .served = &.{ .{ .rows = 4, .n = 10, .k = 0 }, .{ .rows = 8, .n = 300, .k = 0 } }, + }; + + rig: Rig, + o: *vk.ops.Ops, + arena: std.heap.ArenaAllocator, + tokens: usize = 0, + cols: usize = 0, + table_rows: usize = 0, + table_at: usize = 0, + ids_at: usize = 0, + out_at: usize = 0, + ids: []u32 = &.{}, + want: []u8 = &.{}, + msg: [160]u8 = undefined, + + pub fn init(a: std.mem.Allocator, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, o: *vk.ops.Ops, timeout_ns: u64) Self { + return .{ .rig = .{ .a = a, .dev = dev, .cmds = cmds, .staging = staging, .timeout_ns = timeout_ns }, .o = o, .arena = .init(a) }; + } + + pub fn deinit(self: *Self) void { + self.rig.deinit(); + self.arena.deinit(); + } + + pub fn kernel(self: *Self) g.Kernel { + return .{ .spec = spec, .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Self); + + pub fn prepare(self: *Self, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows < 4 or shape.n == 0) return false; + const tokens = shape.rows; + const cols = shape.n; + const cells = std.math.mul(usize, tokens, cols) catch return false; + const table_rows = tokens + 2; + self.rig.reset(); + _ = self.arena.reset(.retain_capacity); + const a = self.arena.allocator(); + const table = try a.alloc(u8, table_rows * cells * esz); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + rng.bytes(table); + if (fill == .nonfinite) { + if (esz == 4) { + const planted = [_]u32{ 0x7fc00000, 0x7f800000, 0xff800000 }; + for (planted, [_]usize{ 0, 1, cols }) |bits, cell| write32(table, cell, bits); + } else { + const planted = [_]u16{ 0x7fc0, 0x7f80, 0xff80 }; + for (planted, [_]usize{ 0, 1, cols }) |bits, cell| write16(table, cell, bits); + } + } + const ids = try a.alloc(u32, tokens); + gatherIds(ids, table_rows); + const want = try a.alloc(u8, cells * esz); + for (ids, 0..) |id, tok| { + const row = want[tok * cols * esz ..][0 .. cols * esz]; + if (id < table_rows) { + @memcpy(row, table[id * cols * esz ..][0 .. cols * esz]); + } else { + @memset(row, 0); // the documented out-of-range read + } + } + self.table_at = try self.rig.input(table); + self.ids_at = try self.rig.input(std.mem.sliceAsBytes(ids)); + self.out_at = try self.rig.output(want.len); + self.tokens = tokens; + self.cols = cols; + self.table_rows = table_rows; + self.ids = ids; + self.want = want; + return true; + } + + pub fn run(self: *Self) !void { + const out = try self.rig.view(self.out_at, kind, self.tokens * self.cols); + const table = try self.rig.view(self.table_at, kind, self.table_rows * self.cols); + const ids = try self.rig.view(self.ids_at, .u32, self.tokens); + try vk.orExit(self.o.gather(self.rig.cmds, out, table, ids, self.table_rows, self.rig.timeout_ns)); + } + + pub fn count(self: *Self) usize { + return self.tokens * self.cols; + } + + pub fn reference(self: *Self, want: []f64, mag: []f64) !void { + _ = self; + @memset(want, 0); + @memset(mag, 0); // never called: the check hook owns every comparison + } + + pub fn check(self: *Self, at: g.Gate) !?[]const u8 { + _ = at; + const got = try self.rig.outputBytes(); + if (std.mem.indexOfDiff(u8, got, self.want)) |byte| { + const tok = byte / (self.cols * esz); + const id = self.ids[tok]; + if (id >= self.table_rows) return buf(&self.msg, "out-of-range id {d} gathered nonzero bytes at {d}: it must gather zeros (ops_q4.zig)", .{ id, byte }); + return buf(&self.msg, "token {d} (id {d}) differs at byte {d}", .{ tok, id, byte }); + } + return null; + } + }; +} + +// Q4 g64 outputs match fp32 scale*q+bias model bits per code after contract RNE; non-finite scale/bias keeps its class. +pub const Dequant = struct { + const group: u64 = 64; + + pub const spec: c.Spec = .{ + .name = "vulkan dequant_bf16 (q4 g64 RNE)", + .accum = .none, + .depth = none, + .out = .bf16, + .nan = .{ .stated = "one fp32 expression, then the contract's RNE: a non-finite scale or bias makes the output the fp32 expression's class" }, + .tiles = .{ .{}, .{ .size = 256, .step = 64, .min = 64 }, .{} }, + .served = &.{ .{ .rows = 6, .n = 128, .k = 0 }, .{ .rows = 2, .n = 2048, .k = 0 } }, + }; + + rig: Rig, + o: *vk.ops.Ops, + arena: std.heap.ArenaAllocator, + rows: usize = 0, + elements: usize = 0, + words: usize = 0, + groups: usize = 0, + w_at: usize = 0, + s_at: usize = 0, + b_at: usize = 0, + out_at: usize = 0, + want: []u8 = &.{}, + msg: [160]u8 = undefined, + + pub fn init(a: std.mem.Allocator, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, o: *vk.ops.Ops, timeout_ns: u64) Dequant { + return .{ .rig = .{ .a = a, .dev = dev, .cmds = cmds, .staging = staging, .timeout_ns = timeout_ns }, .o = o, .arena = .init(a) }; + } + + pub fn deinit(self: *Dequant) void { + self.rig.deinit(); + self.arena.deinit(); + } + + pub fn kernel(self: *Dequant) g.Kernel { + return .{ .spec = spec, .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Dequant); + + pub fn prepare(self: *Dequant, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + // two rows minimum: the nonfinite plants need a second row's first group + if (shape.rows < 2 or shape.n < group or shape.n % group != 0) return false; + const rows = shape.rows; + const per_row = shape.n; + const elements = std.math.mul(usize, rows, per_row) catch return false; + const words = per_row / 8; + const groups = per_row / @as(usize, @intCast(group)); + self.rig.reset(); + _ = self.arena.reset(.retain_capacity); + const a = self.arena.allocator(); + const w = try a.alloc(u8, rows * words * 4); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + rng.bytes(w); + const scales = try a.alloc(u16, rows * groups); + const biases = try a.alloc(u16, rows * groups); + for (scales, biases) |*s, *b| { + s.* = randomBf16(rng); + b.* = randomBf16(rng); + } + if (fill == .nonfinite) { + scales[0] = 0x7fc0; // a NaN scale: its group's outputs are NaN + scales[groups] = 0x7f80; // +inf scale with -inf bias on the next row: inf + -inf is NaN + biases[0] = 0x0000; + biases[groups] = 0xff80; + biases[rows * groups - 1] = 0xff80; // -inf bias on a finite scale: -inf out + } + const want = try a.alloc(u8, elements * 2); + for (0..rows * words) |unit| { + const row = unit / words; + const col = unit % words; + const grp = col / (@as(usize, @intCast(group)) / 8); + const sbits = scales[row * groups + grp]; + const bbits = biases[row * groups + grp]; + const word = read32(w, unit); + for (0..8) |lane| { + const q: u4 = @truncate(word >> @intCast(4 * lane)); + write16(want, unit * 8 + lane, dequantBits(sbits, q, bbits)); + } + } + self.w_at = try self.rig.input(w); + self.s_at = try self.rig.input(std.mem.sliceAsBytes(scales)); + self.b_at = try self.rig.input(std.mem.sliceAsBytes(biases)); + self.out_at = try self.rig.output(want.len); + self.rows = rows; + self.elements = elements; + self.words = words; + self.groups = groups; + self.want = want; + return true; + } + + pub fn run(self: *Dequant) !void { + const out = try self.rig.view(self.out_at, .bf16, self.elements); + const w = try self.rig.view(self.w_at, .u32, self.rows * self.words); + const scales = try self.rig.view(self.s_at, .bf16, self.rows * self.groups); + const biases = try self.rig.view(self.b_at, .bf16, self.rows * self.groups); + try vk.orExit(self.o.dequant(self.rig.cmds, out, w, scales, biases, self.words * 8, group, self.rig.timeout_ns)); + } + + pub fn count(self: *Dequant) usize { + return self.elements; + } + + pub fn reference(self: *Dequant, want: []f64, mag: []f64) !void { + _ = self; + @memset(want, 0); + @memset(mag, 0); // never called: the check hook owns every comparison + } + + pub fn check(self: *Dequant, at: g.Gate) !?[]const u8 { + _ = at; + const got = try self.rig.outputBytes(); + for (0..self.elements) |i| { + const wbits = read16(self.want, i); + const gbits = read16(got, i); + const wv = c.fromBf16(wbits); + if (c.classOf(wv) == .finite) { + if (gbits != wbits) return buf(&self.msg, "element {d} read 0x{x:0>4}, want the model's 0x{x:0>4}", .{ i, gbits, wbits }); + } else if (c.classOf(c.fromBf16(gbits)) != c.classOf(wv)) { + return buf(&self.msg, "element {d} is {t} where the fp32 expression is {t}", .{ i, c.classOf(c.fromBf16(gbits)), c.classOf(wv) }); + } + } + return null; + } +}; + +const testing = std.testing; + +test "the dequant model is one fp32 multiply-add then the contract's RNE" { + try testing.expectEqual(@as(u16, 0x3f80), dequantBits(0x3f80, 1, 0x0000)); + try testing.expectEqual(@as(u16, 0x4170), dequantBits(0x3f80, 15, 0x0000)); // 15.0 is a grid point + try testing.expectEqual(@as(u16, 0x3f80), dequantBits(0x3f80, 1, 0x3b80)); // 1 + 2^-8: tie to even, down + try testing.expectEqual(@as(u16, 0x3f82), dequantBits(0x3f80, 1, 0x3dc0)); // 1 + 3*2^-8: tie with lsb 1, up + try testing.expectEqual(@as(u16, 0x3f81), dequantBits(0x3f80, 1, 0x3b81)); // a hair above half rounds up + try testing.expectEqual(@as(u16, 0x7f80), dequantBits(0x7f80, 1, 0x0000)); // +inf scale stays +inf +} + +test "gather ids include an out-of-range id and a repeat, the rest are valid rows" { + const ids = try testing.allocator.alloc(u32, 6); + defer testing.allocator.free(ids); + gatherIds(ids, 5); + try testing.expectEqual(@as(u32, 5), ids[0]); // out of range: gathers zeros, never a table row + try testing.expectEqual(ids[1], ids[2]); // a repeated row + for (ids[1..]) |id| try testing.expect(id < 5); +} + +test "runAll swallows a kernel run's fence Timeout, so every fixture GPU call must exit before it escapes" { + // runAll turns a fence Timeout into a FAIL verdict, so vk.orExit in the fixture's GPU calls is the only exit. + const F = struct { + fn prepare(_: *anyopaque, _: c.Shape, _: g.Fill, _: u64) anyerror!bool { + return true; + } + fn run(_: *anyopaque) anyerror!void { + return error.Timeout; + } + fn output(_: *anyopaque) anyerror![]const u8 { + return ""; + } + fn inputs(_: *anyopaque) anyerror![]const u8 { + return ""; + } + fn intact(_: *anyopaque) anyerror!bool { + return true; + } + fn overlap(_: *anyopaque) bool { + return false; + } + fn count(_: *anyopaque) usize { + return 0; + } + fn reference(_: *anyopaque, want: []f64, mag: []f64) anyerror!void { + @memset(want, 0); + @memset(mag, 0); + } + }; + const vt = g.VTable{ + .prepare = F.prepare, + .run = F.run, + .output = F.output, + .inputs = F.inputs, + .intact = F.intact, + .overlap = F.overlap, + .count = F.count, + .reference = F.reference, + }; + const k = g.Kernel{ .spec = Cast(true).spec, .ctx = undefined, .vt = &vt }; + var t: g.Tally = .{}; + try g.runAll(testing.allocator, &.{k}, 0, &t); + try testing.expectEqual(@as(usize, 1), t.failed); // the Timeout became a verdict, not an error return + try testing.expectEqual(@as(usize, 0), t.passed); +} diff --git a/zig/tests/vulkan/contract_datamove_main.zig b/zig/tests/vulkan/contract_datamove_main.zig new file mode 100644 index 00000000000..3a8a5a25015 --- /dev/null +++ b/zig/tests/vulkan/contract_datamove_main.zig @@ -0,0 +1,105 @@ +// GPU runner for stage 1 datamove gates 1-4; each gate prints a line. Exit 0/1/2 = pass/fail/usage. +const std = @import("std"); +const vk = @import("vulkan"); +const shaders = @import("stage1_shaders"); +const contract = @import("contract"); +const dm = @import("contract_datamove.zig"); + +const g = contract.gates; + +const seed: u64 = 0xda7a; + +const usage = "usage: tf-stage1-datamove [--vendor ID] [--timeout-ms MS] | --help\n" ++ + " --help show this text and exit without opening a loader or device\n" ++ + " --vendor ID refuse a device whose Vulkan vendorID differs (decimal or 0x hex; no default)\n" ++ + " --timeout-ms MS bound on every submission and transfer, 1 to 600000 (default 10000)\n"; + +pub fn main(init: std.process.Init) !u8 { + const args = try init.minimal.args.toSlice(init.arena.allocator()); + var vendor: ?u32 = null; + var timeout_ms: u64 = 10_000; + var i: usize = 1; + while (i < args.len) : (i += 2) { + const value = if (i + 1 < args.len) args[i + 1] else ""; + if (std.mem.eql(u8, args[i], "--help")) { + std.debug.print(usage, .{}); + return 0; + } else if (std.mem.eql(u8, args[i], "--vendor")) { + vendor = std.fmt.parseInt(u32, value, 0) catch null; + if (vendor == null) break; + } else if (std.mem.eql(u8, args[i], "--timeout-ms")) { + timeout_ms = std.fmt.parseInt(u64, value, 10) catch 0; + if (timeout_ms == 0 or timeout_ms > 600_000) break; + } else break; + } else return run(init.gpa, vendor, timeout_ms * std.time.ns_per_ms); + std.debug.print(usage, .{}); + return 2; +} + +fn run(a: std.mem.Allocator, vendor: ?u32, timeout_ns: u64) !u8 { + var rt = try vk.Runtime.open(); + defer rt.close(); + var dev = try vk.Device.open(&rt, .{ .features = .{ .storage16 = true, .int16 = true } }); + defer dev.close(); + const pr = dev.properties; + const name = std.mem.sliceTo(&pr.deviceName, 0); + if (pr.deviceType == vk.c.VK_PHYSICAL_DEVICE_TYPE_CPU) { + std.debug.print("FAIL datamove: a CPU (software) Vulkan device is never a GPU result\n", .{}); + return 1; + } + if (vendor) |want| if (pr.vendorID != want) { + std.debug.print("FAIL datamove: vendor 0x{x:0>4} is not the requested 0x{x:0>4}\n", .{ pr.vendorID, want }); + return 1; + }; + std.debug.print("MILESTONE device-found name=\"{s}\" vendor=0x{x:0>4} storage16={} int16={} timeout={d}ms\n", .{ name, pr.vendorID, dev.features.storage16, dev.features.int16, timeout_ns / std.time.ns_per_ms }); + var cmds = try vk.Commands.init(&dev); + defer cmds.deinit(); + var staging = try vk.Buffer.init(&dev, .staging, 1 << 20); + defer staging.deinit(); + const sh = vk.ops.Shaders{ + .qmm_bf16 = &shaders.qmm_bf16, + .qmm_q4_vec4 = &shaders.qmm_q4_vec4, + .qmm_f32 = &shaders.qmm_f32, + .copy_f32_bf16 = &shaders.copy_f32_bf16, + .qmm_q4_tile = &shaders.qmm_q4_tile, + .qmm_q4_tile_f32 = &shaders.qmm_q4_tile_f32, + .dequant_bf16 = &shaders.dequant_bf16, + .gather_u32 = &shaders.gather_u32, + .gather_u16 = &shaders.gather_u16, + .cast_bf16_f32 = &shaders.cast_bf16_f32, + .cast_f32_bf16 = &shaders.cast_f32_bf16, + .copy_u32 = &shaders.copy_u32, + .copy_u16 = &shaders.copy_u16, + }; + var o = try vk.ops.Ops.init(&dev, sh); + defer o.deinit(); + + var cast_to = dm.Cast(true).init(a, &dev, &cmds, staging, &o, timeout_ns); + defer cast_to.deinit(); + var cast_from = dm.Cast(false).init(a, &dev, &cmds, staging, &o, timeout_ns); + defer cast_from.deinit(); + var copy_u32 = dm.Copy(.u32).init(a, &dev, &cmds, staging, &o, timeout_ns); + defer copy_u32.deinit(); + var copy_u16 = dm.Copy(.u16).init(a, &dev, &cmds, staging, &o, timeout_ns); + defer copy_u16.deinit(); + var gather_u32 = dm.Gather(.u32).init(a, &dev, &cmds, staging, &o, timeout_ns); + defer gather_u32.deinit(); + var gather_u16 = dm.Gather(.u16).init(a, &dev, &cmds, staging, &o, timeout_ns); + defer gather_u16.deinit(); + var dequant = dm.Dequant.init(a, &dev, &cmds, staging, &o, timeout_ns); + defer dequant.deinit(); + + var kernels = [_]g.Kernel{ + cast_to.kernel(), + cast_from.kernel(), + copy_u32.kernel(), + copy_u16.kernel(), + gather_u32.kernel(), + gather_u16.kernel(), + dequant.kernel(), + }; + var t: g.Tally = .{}; + try g.runAll(a, &kernels, seed, &t); + std.debug.print("RESULT datamove on \"{s}\": kernels={d} passed={d} failed={d} controls={d} known={d}\n", .{ name, t.kernels, t.passed, t.failed, t.controls, t.known }); + return if (t.problems() == 0) 0 else 1; +} diff --git a/zig/tests/vulkan/contract_datamove_rig.zig b/zig/tests/vulkan/contract_datamove_rig.zig new file mode 100644 index 00000000000..5dd119e0302 --- /dev/null +++ b/zig/tests/vulkan/contract_datamove_rig.zig @@ -0,0 +1,125 @@ +// The datamove contract's rig: guarded device buffers, the little-endian byte helpers and the failure message. +const std = @import("std"); +const vk = @import("vulkan"); + +/// Guard bytes on both sides of every rig payload (the largest minStorageBufferOffsetAlignment). +pub const guard = 256; +const pattern = 0xa5; + +// Each call uploads inputs, poisons outputs with 0xff, and guards every payload. +pub const Rig = struct { + a: std.mem.Allocator, + dev: *const vk.Device, + cmds: *vk.Commands, + staging: vk.Buffer, + timeout_ns: u64, + slots: std.ArrayList(Slot) = .empty, + seen: std.ArrayList(u8) = .empty, + outs: std.ArrayList(u8) = .empty, + + const Kind = enum { input, output }; + const Slot = struct { buf: vk.Buffer, len: usize, kind: Kind }; + + pub fn deinit(r: *Rig) void { + r.reset(); + r.slots.deinit(r.a); + r.seen.deinit(r.a); + r.outs.deinit(r.a); + } + + pub fn reset(r: *Rig) void { + // A hook error after a recorded dispatch leaves the buffer open: drop those records before freeing slots. + r.cmds.discard(); + for (r.slots.items) |*s| s.buf.deinit(); + r.slots.clearRetainingCapacity(); + } + + fn slot(r: *Rig, data: ?[]const u8, len: usize, kind: Kind) !usize { + const all = try r.a.alloc(u8, len + 2 * guard); + defer r.a.free(all); + @memset(all, pattern); + if (data) |d| { + @memcpy(all[guard..][0..len], d); + } else { + @memset(all[guard..][0..len], 0xff); // an output element no launch writes reads as NaN + } + var b = try vk.Buffer.init(r.dev, .device, all.len); + errdefer b.deinit(); + // Timeout or InFlight leaves completion unknown: exit before any teardown can free GPU-owned resources. + try vk.orExit(vk.upload(r.cmds, b, 0, r.staging, all, r.timeout_ns)); + try r.slots.append(r.a, .{ .buf = b, .len = len, .kind = kind }); + return r.slots.items.len - 1; + } + + pub fn input(r: *Rig, data: []const u8) !usize { + return r.slot(data, data.len, .input); + } + + /// An output of `len` bytes, all 0xff until the launch writes it. + pub fn output(r: *Rig, len: usize) !usize { + return r.slot(null, len, .output); + } + + /// The payload of slot `i` as a typed view for a binding; call between prepare and the next prepare. + pub fn view(r: *Rig, i: usize, kind: vk.tensor.Kind, elements: usize) !vk.tensor.View { + const s = &r.slots.items[i]; + return vk.tensor.View.init(&s.buf, kind, guard, elements); + } + + fn gatherBytes(r: *Rig, kind: Kind, out: *std.ArrayList(u8), whole: bool) ![]const u8 { + out.clearRetainingCapacity(); + for (r.slots.items) |s| if (s.kind == kind or whole) { + const at = out.items.len; + const n = if (whole) s.len + 2 * guard else s.len; + try out.resize(r.a, at + n); + try vk.orExit(vk.download(r.cmds, s.buf, if (whole) 0 else guard, r.staging, out.items[at..][0..n], r.timeout_ns)); + }; + return out.items; + } + + pub fn outputBytes(r: *Rig) ![]const u8 { + return r.gatherBytes(.output, &r.outs, false); + } + + pub fn inputBytes(r: *Rig) ![]const u8 { + return r.gatherBytes(.input, &r.seen, false); + } + + /// Every guard of every buffer still holds its pattern. + pub fn intact(r: *Rig) !bool { + const all = try r.gatherBytes(.input, &r.seen, true); + var at: usize = 0; + for (r.slots.items) |s| { + for (all[at..][0..guard]) |v| if (v != pattern) return false; + for (all[at + guard + s.len ..][0..guard]) |v| if (v != pattern) return false; + at += s.len + 2 * guard; + } + return true; + } + + /// Every payload is its own buffer, so no output range overlaps an input range. + pub fn overlap(r: *Rig) bool { + _ = r; + return false; + } +}; + +pub fn read16(bytes: []const u8, i: usize) u16 { + return std.mem.readInt(u16, bytes[2 * i ..][0..2], .little); +} + +pub fn write16(bytes: []u8, i: usize, bits: u16) void { + std.mem.writeInt(u16, bytes[2 * i ..][0..2], bits, .little); +} + +pub fn read32(bytes: []const u8, i: usize) u32 { + return std.mem.readInt(u32, bytes[4 * i ..][0..4], .little); +} + +pub fn write32(bytes: []u8, i: usize, bits: u32) void { + std.mem.writeInt(u32, bytes[4 * i ..][0..4], bits, .little); +} + +pub fn buf(msg: *[160]u8, comptime fmt: []const u8, args: anytype) ?[]const u8 { + return std.fmt.bufPrint(msg, fmt, args) catch "message too long"; +} diff --git a/zig/tests/vulkan/contract_gdn.zig b/zig/tests/vulkan/contract_gdn.zig new file mode 100644 index 00000000000..d8e6480cd90 --- /dev/null +++ b/zig/tests/vulkan/contract_gdn.zig @@ -0,0 +1,399 @@ +// The GDN contract's host suite: the fixtures, the references and the uploads the gates check. +const std = @import("std"); +const contract = @import("contract"); +const fx = @import("stage4_fixture.zig"); +const refs = @import("contract_gdn_ref.zig"); +const ops = @import("contract_gdn_ops.zig"); + +const c = contract.spec; +const g = contract.gates; +const rows = contract.rows; + +pub const Gdn = ops.Gdn; +pub const Which = refs.Which; +pub const chunk_fixed = refs.chunk_fixed; +pub const spec = refs.spec; +pub const walk = refs.walk; +pub const splitReference = refs.splitReference; +pub const convReference = refs.convReference; + +const heads = refs.heads; +const dk = refs.dk; +const linear_dim = refs.linear_dim; +const channels = refs.channels; +const state_elems = refs.state_elems; +const poison = refs.poison; +const decodeGates = refs.decodeGates; +const f32Of = refs.f32Of; +const bf = refs.bf; +const putF32 = refs.putF32; +const splitDepth = refs.splitDepth; +const tapsDepth = refs.tapsDepth; +const recurYDepth = refs.recurYDepth; +const recurStateDepth = refs.recurStateDepth; +const decodeYDepth = refs.decodeYDepth; +const decodeStateDepth = refs.decodeStateDepth; +const snapDepth = refs.snapDepth; + +const testing = std.testing; + +// Host fixtures need a real allocator; device fields stay untouched because host tests never call `uploadBuffers`. +fn hostFixture(which: Which) Gdn { + return Gdn.init(testing.allocator, undefined, undefined, undefined, undefined, which); +} + +test "gdn contract prepareHost refuses torn geometry and serves the fixed one" { + var f = hostFixture(.split_soa); + defer f.deinit(); + try testing.expectEqual(@as(?usize, null), try f.prepareHost(.{ .rows = 0, .n = heads, .k = dk }, .random, 1)); + try testing.expectEqual(@as(?usize, null), try f.prepareHost(.{ .rows = 8, .n = 8, .k = dk }, .random, 1)); + try testing.expectEqual(@as(?usize, null), try f.prepareHost(.{ .rows = 8, .n = heads, .k = 64 }, .random, 1)); + for (std.enums.values(Which)) |which| { + var g_ = hostFixture(which); + defer g_.deinit(); + const out_len = (try g_.prepareHost(.{ .rows = 1, .n = heads, .k = dk }, .random, 3)) orelse return error.Refused; + try testing.expect(out_len > 0); + } +} + +test "gdn contract prepareHost random fill: finite references within fixture budgets" { + for (std.enums.values(Which)) |which| { + var f = hostFixture(which); + defer f.deinit(); + for ([_]usize{ 1, 8, 32 }) |T| { + _ = (try f.prepareHost(.{ .rows = T, .n = heads, .k = dk }, .random, 0x6764 + T)) orelse return error.Refused; + try testing.expectEqual(f.want.len, f.count()); + var buf: [8][]u16 = undefined; + try testing.expect(f.inputList(&buf).len > 0); + for (f.want, f.mag) |w, m| { + try testing.expect(std.math.isFinite(w)); + try testing.expect(std.math.isFinite(m) and m >= 0); + } + const allow = try f.arena.allocator().alloc(f64, f.want.len); + try f.bounds(f.want, f.mag, allow); + for (allow) |l| try testing.expect(std.math.isFinite(l) and l > 0); + // Before run, outputs stay poison; `want` is the GPU target, exposing a wrong reference length. + for (f.out[f.want.len..]) |w| try testing.expectEqual(poison, w); + } + } +} + +test "gdn contract prepareHost stress fill: bounded state, planted mirror, fp16-proving partials" { + for ([_]Which{ .prefill_y, .prefill_state, .decode_y, .decode_state }) |which| { + var f = hostFixture(which); + defer f.deinit(); + _ = (try f.prepareHost(.{ .rows = 4, .n = heads, .k = dk }, .stress, 0x6764)) orelse return error.Refused; + // Stress g=1.0; beta=2^-10 (prefill) or rounded sigmoid(-10) (decode), so state feedback converges. + const state_src = if (which == .prefill_y or which == .prefill_state) f.h0_bits else f.state_init_bits; + for (0..state_elems) |i| { + const v = f32Of(state_src, i); + try testing.expect(v >= 4096 and v < 4104); + } + // `v` is bf16 of pinned kv dot, keeping delta tiny; scan soa words (prefill v contiguous, decode per token). + var mirror_ok = false; + for (f.soa_bits) |bits| { + const v = bf(bits); + if (v > 500000 and v < 550000) mirror_ok = true; + } + try testing.expect(mirror_ok); + // the f64 reference stays finite and bounded at stress scale + for (f.want) |w| try testing.expect(std.math.isFinite(w)); + try testing.expect(f.max_state < 1e6); + } +} + +test "gdn contract prepareHost nonfinite fill plants classes the walkers propagate" { + for (std.enums.values(Which)) |which| { + var f = hostFixture(which); + defer f.deinit(); + _ = (try f.prepareHost(.{ .rows = 8, .n = heads, .k = dk }, .nonfinite, 0x6765)) orelse return error.Refused; + var planted: usize = 0; + for (f.want) |w| { + if (c.classOf(w) != .finite) planted += 1; + } + try testing.expect(planted > 0); + try testing.expect(planted < f.want.len); + } +} + +test "gdn contract conv prefill reference reproduces the fixture Conv checker's want" { + // the assembled input taps and the state carry, at the fixture's shape + var f = hostFixture(.conv_prefill); + defer f.deinit(); + _ = (try f.prepareHost(.{ .rows = @intCast(fx.tokens), .n = heads, .k = dk }, .random, 0x57a5)) orelse return error.Refused; + const cf = try fx.Conv.init(testing.allocator, 0x57a4 + 1); + defer cf.deinit(); + // same seed family as the fixture: check MY reference math against a directly computed conv over MY generated bits + const C: usize = channels; + for (0..@intCast(fx.tokens)) |t| { + for (0..C) |ch| { + var acc: f64 = 0; + for (0..4) |k| { + const row = t + k; + const in_bit = if (row < 3) f.state_bits[row * C + ch] else f.qkv_bits[(row - 3) * C + ch]; + acc += bf(in_bit) * bf(f.weight_bits[ch * 4 + k]); + } + try testing.expectEqual(f.want[t * C + ch], acc); + } + } +} + +test "gdn contract depth functions grow with tokens, never constant" { + try testing.expect(recurYDepth(.{ .rows = 130, .n = 16, .k = 128 }) > recurYDepth(.{ .rows = 8, .n = 16, .k = 128 })); + try testing.expectEqual(@as(usize, 130 * 130 + 130), recurYDepth(.{ .rows = 130, .n = 16, .k = 128 })); + try testing.expectEqual(@as(usize, 4 * 130 + 2), recurStateDepth(.{ .rows = 4, .n = 16, .k = 128 })); + try testing.expectEqual(@as(usize, 3 * 260 + 130), decodeYDepth(.{ .rows = 3, .n = 16, .k = 128 })); + try testing.expectEqual(@as(usize, 132 + 2), decodeStateDepth(.{ .rows = 1, .n = 16, .k = 128 })); + try testing.expectEqual(@as(f64, 8 * 130 + 2), snapDepth()); + // the conv and split depths are genuinely shape-independent chains + try testing.expectEqual(@as(usize, 4), tapsDepth(.{ .rows = 99, .n = 16, .k = 128 })); + try testing.expectEqual(@as(usize, 40), splitDepth(.{ .rows = 99, .n = 16, .k = 128 })); +} + +test "gdn contract walker matches the stage 4 fixture walker bit for bit at the fixture shape" { + const a = testing.allocator; + const T: usize = @intCast(fx.tokens); // 8, the fixture's comptime shape + const soa_bits = try a.alloc(u16, 3 * T * linear_dim); + defer a.free(soa_bits); + const g_pairs = try a.alloc(u16, 2 * T * heads); + defer a.free(g_pairs); + const beta_bits = try a.alloc(u16, T * heads); + defer a.free(beta_bits); + const h0_bits = try a.alloc(u16, 2 * state_elems); + defer a.free(h0_bits); + var prng = std.Random.DefaultPrng.init(0xc7); + const rng = prng.random(); + for (soa_bits) |*x| x.* = fx.sampleDyadic(rng, -3, 1); + for (g_pairs) |*x| { + const bits: u32 = @bitCast(fx.sampleDyadicF32(rng, -4, 0)); + x.* = @truncate(bits); + } + for (beta_bits) |*x| x.* = fx.sampleDyadic(rng, -3, 1); + for (0..state_elems) |i| putF32(h0_bits, i, fx.sampleDyadicF32(rng, -6, -2)); + + const mine_y = try a.alloc(f64, T * linear_dim); + defer a.free(mine_y); + const mine_state = try a.alloc(f64, state_elems); + defer a.free(mine_state); + const mine_st = try a.alloc(f64, state_elems); + defer a.free(mine_st); + var mine_max: f64 = 0; + walk(T, soa_bits, g_pairs, beta_bits, h0_bits, mine_st, mine_y, mine_state, &mine_max, null, null, null, null); + + const fx_y = try a.alloc(f64, T * linear_dim); + defer a.free(fx_y); + const fx_state = try a.alloc(f64, state_elems); + defer a.free(fx_state); + var fx_max: f64 = 0; + fx.RecurWalk.run(soa_bits, g_pairs, beta_bits, h0_bits, fx_y, fx_state, &fx_max, null, null); + try testing.expectEqualSlices(f64, fx_y, mine_y); + try testing.expectEqualSlices(f64, fx_state, mine_state); + try testing.expectEqual(fx_max, mine_max); + + // the chunked variant's interior snapshot: boundary at token 4 + const mine_snap = try a.alloc(f64, state_elems); + defer a.free(mine_snap); + const fx_snap = try a.alloc(f64, state_elems); + defer a.free(fx_snap); + var mine_max2: f64 = 0; + var fx_max2: f64 = 0; + walk(T, soa_bits, g_pairs, beta_bits, h0_bits, mine_st, mine_y, mine_state, &mine_max2, 4, mine_snap, null, null); + fx.RecurWalk.run(soa_bits, g_pairs, beta_bits, h0_bits, fx_y, fx_state, &fx_max2, 4, fx_snap); + try testing.expectEqualSlices(f64, fx_snap, mine_snap); + try testing.expectEqualSlices(f64, fx_state, mine_state); +} + +test "gdn contract walker matches the fixture decode reference exactly" { + const a = testing.allocator; + const f = try fx.Decode.init(a, 0x53); + defer f.deinit(); + // pack the raw-gates chain the fixture derived into the walk's inputs + const g_pairs = try a.alloc(u16, 2 * heads); + defer a.free(g_pairs); + const beta_bits = try a.alloc(u16, heads); + defer a.free(beta_bits); + decodeGates(f.a_bits, f.b_bits, f.a_log_bits, f.dt_bits, 0, g_pairs, beta_bits); + const y = try a.alloc(f64, linear_dim); + defer a.free(y); + const st_out = try a.alloc(f64, state_elems); + defer a.free(st_out); + const st = try a.alloc(f64, state_elems); + defer a.free(st); + var max_state: f64 = 0; + walk(1, f.soa_bits, g_pairs, beta_bits, f.state_in_bits, st, y, st_out, &max_state, null, null, null, null); + for (f.want_y, y) |w, m| try testing.expectEqual(w, m); + for (f.want_state, st_out) |w, m| try testing.expectEqual(w, m); + try testing.expectEqual(f.max_state, max_state); +} + +test "gdn contract split reference replays clean through the fixture Split checker" { + const a = testing.allocator; + const T8: usize = @intCast(fx.tokens); // the fixture checker's comptime shape + var f = try fx.Split.init(a, 0x9e); // the fixture generates its own inputs + defer f.deinit(); + // my reference over the SAME bits, rendered to bits, through the PRODUCTION fixture checker + const want8 = try a.alloc(f64, 3 * T8 * linear_dim + T8 * heads); + defer a.free(want8); + const mag8 = try a.alloc(f64, want8.len); + defer a.free(mag8); + splitReference(T8, f.conv_out_bits, f.a_bits, f.b_bits, f.a_log_bits, f.dt_bits, want8, mag8); + const soa_bits = try a.alloc(u16, 3 * T8 * linear_dim); + defer a.free(soa_bits); + const beta_bits = try a.alloc(u16, T8 * heads); + defer a.free(beta_bits); + const g_pairs = try a.alloc(u16, 2 * T8 * heads); + defer a.free(g_pairs); + for (0..3 * T8 * linear_dim) |i| soa_bits[i] = c.toBf16(@floatCast(want8[i])); + for (0..T8 * heads) |i| beta_bits[i] = c.toBf16(@floatCast(want8[3 * T8 * linear_dim + i])); + for (0..T8) |t| for (0..heads) |h| { + const ref = fx.gateRefOracle(f.a_bits[t * heads + h], f.b_bits[t * heads + h], f32Of(f.a_log_bits, h), f.dt_bits[h]); + putF32(g_pairs, t * heads + h, @floatCast(ref.g)); + }; + const rep = try f.check(soa_bits, g_pairs, beta_bits, null); + try testing.expect(rep.worst <= 1); + // and the same want/mag pass my own budget hook's allowances + for (want8, mag8) |w, m| { + const allow = fx.budget(40, 5, 3, w, m); + try testing.expect(allow > 0 and std.math.isFinite(allow)); + } +} + +test "gdn contract split g reference equals the fixture's oracle chain" { + const a = testing.allocator; + const f = try fx.SplitOne.init(a, 0x57ac); + defer f.deinit(); + // `splitReference` omits g; `split_g` must match fixture `gateRefOracle`, so replay it here. + for (0..heads) |h| { + const ref = fx.gateRefOracle(f.a_bits[h], f.b_bits[h], f32Of(f.a_log_bits, h), f.dt_bits[h]); + try testing.expect(ref.g > 0 and ref.g < 1); + try testing.expect(ref.beta > 0 and ref.beta < 1); + } +} + +test "gdn contract stress keeps f32 sums exact and makes fp16 miss by over 100x" { + const a = testing.allocator; + var arena = std.heap.ArenaAllocator.init(a); + defer arena.deinit(); + const ar = arena.allocator(); + const T: usize = 4; + const soa_bits = try ar.alloc(u16, 3 * T * linear_dim); + const g_pairs = try ar.alloc(u16, 2 * T * heads); + const beta_bits = try ar.alloc(u16, T * heads); + const h0_bits = try ar.alloc(u16, 2 * state_elems); + for (0..state_elems) |i| putF32(h0_bits, i, @floatFromInt(4096 + i % 7)); + for (soa_bits) |*x| x.* = c.toBf16(1); + for (0..T * linear_dim) |i| { + if (i % dk % 2 == 1) soa_bits[i] = c.toBf16(-1); + } + for (0..heads) |h| for (0..dk) |row| { + var kv: f64 = 0; + for (0..dk) |i| kv += @as(f64, @floatFromInt(4096 + (h * dk * dk + row * dk + i) % 7)); + for (0..T) |t| soa_bits[2 * T * linear_dim + (t * heads + h) * dk + row] = c.toBf16(@floatCast(kv)); + }; + for (0..T * heads) |i| { + putF32(g_pairs, i, 1.0); + beta_bits[i] = c.toBf16(0x1p-10); + } + // f32 sums fit every-order state budget; dyadic denominators exceed 24 bits after 2 tokens, covered by depth. + const want_y = try ar.alloc(f64, T * linear_dim); + const want_state = try ar.alloc(f64, state_elems); + const st = try ar.alloc(f64, state_elems); + var max_state: f64 = 0; + walk(T, soa_bits, g_pairs, beta_bits, h0_bits, st, want_y, want_state, &max_state, null, null, null, null); + const sim_y = try ar.alloc(f64, T * linear_dim); + const sim_state = try ar.alloc(f64, state_elems); + var sim_max: f64 = 0; + for ([_]g.Order{ .sequential, .pairwise, .strided }) |order| { + walk(T, soa_bits, g_pairs, beta_bits, h0_bits, st, sim_y, sim_state, &sim_max, null, null, .f32, order); + for (sim_state, want_state) |gv, wv| { + const allow = fx.budget(@floatFromInt(T * 130 + 2), 2, 0, wv, max_state); + try testing.expect(@abs(gv - wv) <= allow); + } + } + // fp16: the kv dot's partial sums pass 65504, so the state and y blow past the budget by orders of magnitude + const allow = fx.budget(@floatFromInt(T * 130 + 2), 2, 0, want_state[0], max_state); + var worst: f64 = 0; + walk(T, soa_bits, g_pairs, beta_bits, h0_bits, st, sim_y, sim_state, &sim_max, null, null, .f16, .sequential); + for (sim_state, want_state) |gv, wv| { + if (!std.math.isFinite(gv)) { + worst = std.math.inf(f64); + break; + } + worst = @max(worst, @abs(gv - wv) / allow); + } + try testing.expect(worst >= 100); +} + +// Buffer `i` takes `out` words iff a checked output is fetched there. +fn expectOut(which: Which, i: usize, chunks: usize) bool { + return switch (which) { + .split_soa => i == 5 or i == 7, + .split_g => i == 6, + .conv_prefill => i == 8 or i == 12, + .conv_decode => i == 9, + .prefill_y => i == 4, + .prefill_state => i == 5 or (i == 6 and chunks > 1), + .decode_y => i == 8, + .decode_state => false, + }; +} + +// The device buffers every which uploads, so the fixed indices run() binds stay the ones uploadWords hands the device. +fn expectCount(which: Which, chunks: usize) usize { + return switch (which) { + .split_soa, .split_g => 8, + .conv_prefill, .conv_decode => 13, + .prefill_y, .prefill_state => if (chunks > 1) 7 else 6, + .decode_y, .decode_state => 9, + }; +} + +test "gdn contract upload words: every upload inside the region prepare sized" { + for (std.enums.values(Which)) |which| { + var f = hostFixture(which); + defer f.deinit(); + const s = spec(which); + // the shapes the gates drive: gate 1 and 2's small(), gate 3's served and edge, gate 4's stress + var steps: [16]struct { shape: c.Shape, fill: rows.Fill } = undefined; + var ns: usize = 0; + var one = s.served[0]; + one.rows = std.math.clamp(one.rows, 5, 8); + steps[ns] = .{ .shape = one, .fill = .nonfinite }; + ns += 1; + one = s.served[0]; + one.rows = std.math.clamp(one.rows, 1, 64); + steps[ns] = .{ .shape = one, .fill = .random }; + ns += 1; + const list = try c.shapes(testing.allocator, s); + defer testing.allocator.free(list); + for (list) |sh| { + steps[ns] = .{ .shape = sh, .fill = .random }; + ns += 1; + } + one = s.served[0]; + one.rows = @min(one.rows, 4); + steps[ns] = .{ .shape = one, .fill = .stress }; + ns += 1; + + for (steps[0..ns]) |step| { + const prepared = f.prepareHost(step.shape, step.fill, 0x6764) catch |err| switch (err) { + // the fixture refuses this fill before any upload; the gate reports the error and never uploads + error.NonFiniteReference => continue, + else => return err, + }; + const out_len = prepared orelse continue; + try testing.expect(out_len > 0); + // the upload path uploadBuffers takes, without a device + var words: [16][]u16 = undefined; + const n = try f.uploadWords(&words); + try testing.expectEqual(expectCount(which, f.chunks), n); + const lo = @intFromPtr(f.out.ptr); + const hi = lo + f.out.len * @sizeOf(u16); + for (words[0..n], 0..) |w, i| { + const at = @intFromPtr(w.ptr); + const in_out = at >= lo and at + w.len * @sizeOf(u16) <= hi; + try testing.expectEqual(expectOut(which, i, f.chunks), in_out); + } + } + } +} diff --git a/zig/tests/vulkan/contract_gdn_host.zig b/zig/tests/vulkan/contract_gdn_host.zig new file mode 100644 index 00000000000..787b54a6a64 --- /dev/null +++ b/zig/tests/vulkan/contract_gdn_host.zig @@ -0,0 +1,530 @@ +// The GDN contract's host layer: the fills, the references and the upload order, none of it on a device. +const std = @import("std"); +const contract = @import("contract"); +const fx = @import("stage4_fixture.zig"); +const refs = @import("contract_gdn_ref.zig"); +const ops = @import("contract_gdn_ops.zig"); + +const c = contract.spec; +const g = contract.gates; +const rows = contract.rows; + +const Gdn = ops.Gdn; +const Which = refs.Which; +const heads = refs.heads; +const dk = refs.dk; +const linear_dim = refs.linear_dim; +const channels = refs.channels; +const taps = refs.taps; +const state_elems = refs.state_elems; +const chunk_fixed = refs.chunk_fixed; +const poison = refs.poison; +const putF32 = refs.putF32; +const guard_words = refs.guard_words; +const max_tokens = refs.max_tokens; +const splitReference = refs.splitReference; +const convReference = refs.convReference; +const decodeGates = refs.decodeGates; +const f32Of = refs.f32Of; +const bf = refs.bf; +const walk = refs.walk; +const PrefillInput = refs.PrefillInput; +const DecodeInput = refs.DecodeInput; + +// The fixture's dyadic sampler so gate 3's inputs sit inside the ranges the stated budgets were derived over. +fn fillDyadic(rng: std.Random, v: []u16, lo: i32, hi: i32) void { + for (v) |*x| x.* = fx.sampleDyadic(rng, lo, hi); +} + +fn fillDyadicF32(rng: std.Random, pairs: []u16, lo: i32, hi: i32) void { + for (0..pairs.len / 2) |i| { + const bits: u32 = @bitCast(fx.sampleDyadicF32(rng, lo, hi)); + pairs[2 * i] = @truncate(bits); + pairs[2 * i + 1] = @truncate(bits >> 16); + } +} + +// Gate 1 plants NaN and ±inf in token row 1; finite per-head `a_log`/`dt`/`weight` stop one class masking outputs. +pub fn plantNonfinite(f: *Gdn, rng: std.Random) !void { + const T = f.tokens; + const word = [_]u16{ 0x7fc0, 0x7f80, 0xff80 }; // NaN, +inf, -inf + switch (f.which) { + .split_soa, .split_g => { + f.conv_out_bits[1 * channels + 5] = word[0]; + f.conv_out_bits[1 * channels + 6] = word[1]; + f.conv_out_bits[1 * channels + 7] = word[2]; + f.a_bits[1 * heads + 3] = word[0]; + f.a_bits[1 * heads + 11] = word[1]; + f.b_bits[1 * heads + 9] = word[2]; + }, + .conv_prefill => { + f.qkv_bits[5] = word[0]; // qkv row 0 = assembled input row 3 + f.state_bits[3] = word[1]; // state row 0: +inf and -inf one row + f.state_bits[9] = word[2]; + }, + .conv_decode => { + f.qkv_bits[3 * channels + 5] = word[0]; // token 3's x + f.state_bits[3] = word[1]; + f.state_bits[9] = word[2]; + }, + .prefill_y, .prefill_state => { + const section = T * linear_dim; + f.soa_bits[0 * section + (1 * heads + 3) * dk + 5] = word[0]; // q + f.soa_bits[1 * section + (1 * heads + 3) * dk + 7] = word[1]; // k + f.soa_bits[2 * section + (1 * heads + 3) * dk + 9] = word[2]; // v + putF32(f.g_pairs, 1 * heads + 3, std.math.inf(f32)); + f.beta_bits[1 * heads + 7] = word[0]; + f.h0_bits[2 * (11 * dk * dk)] = word[0]; + f.h0_bits[2 * (11 * dk * dk) + 1] = word[0]; + }, + .decode_y, .decode_state => { + f.soa_bits[0 * linear_dim + 1] = word[0]; // q head 0 + f.soa_bits[1 * linear_dim + 2] = word[1]; // k head 1 + f.soa_bits[2 * linear_dim + 3] = word[2]; // v head 2 + if (T > 1) { + f.soa_bits[(1 * 3 + 0) * linear_dim + 130] = word[0]; + f.soa_bits[(1 * 3 + 1) * linear_dim + 131] = word[1]; + f.soa_bits[(1 * 3 + 2) * linear_dim + 132] = word[2]; + f.a_bits[1 * heads + 5] = word[1]; + f.b_bits[1 * heads + 6] = word[2]; + } + f.state_bits[2 * (4 * dk * dk)] = word[0]; + f.state_bits[2 * (4 * dk * dk) + 1] = word[0]; + }, + } + _ = rng; +} + +// F32 dyadic sums are order-exact below 2^24; fp16 misses by thousands, while bf16 `v` pins state with tiny delta. +pub fn fillStress(f: *Gdn) void { + const T = f.tokens; + switch (f.which) { + .split_soa, .split_g => { + for (0..T * channels) |i| { + const mag: i64 = 128 + @as(i64, @intCast(i % 8)); + f.conv_out_bits[i] = c.toBf16(@floatFromInt(if (i % 2 == 0) mag else -mag)); + } + for (0..T * heads) |i| { + f.a_bits[i] = c.toBf16(0); + f.b_bits[i] = c.toBf16(0); + } + for (0..heads) |h| { + f.dt_bits[h] = c.toBf16(0); + putF32(f.a_log_bits, h, 0); + } + }, + .conv_prefill => { + for (f.state_bits, 0..) |*x, i| { + const sgn: f32 = if (i % 2 == 0) 1 else -1; + x.* = c.toBf16(sgn * @as(f32, @floatFromInt(16384 + i % 64))); + } + for (f.qkv_bits, 0..) |*x, i| { + const sgn: f32 = if (i % 2 == 0) 1 else -1; + x.* = c.toBf16(sgn * @as(f32, @floatFromInt(16384 + i % 64))); + } + for (f.weight_bits, 0..) |*x, i| { + const sgn: f32 = if (i % 3 == 0) 1 else -1; + x.* = c.toBf16(sgn * @as(f32, @floatFromInt(2 + i % 3))); + } + }, + .conv_decode => { + for (f.state_bits, 0..) |*x, i| { + const sgn: f32 = if (i % 2 == 0) 1 else -1; + x.* = c.toBf16(sgn * @as(f32, @floatFromInt(16384 + i % 64))); + } + for (f.qkv_bits, 0..) |*x, i| { + const sgn: f32 = if (i % 2 == 0) 1 else -1; + x.* = c.toBf16(sgn * @as(f32, @floatFromInt(16384 + i % 64))); + } + for (f.weight_bits, 0..) |*x, i| { + const sgn: f32 = if (i % 3 == 0) 1 else -1; + x.* = c.toBf16(sgn * @as(f32, @floatFromInt(2 + i % 3))); + } + }, + .prefill_y, .prefill_state => { + // State 2^12+small; k=1, q=±1, g=1, beta=2^-10. Gain 1.0 blows 2^24; bf16 v keeps delta tiny. + for (0..state_elems) |i| putF32(f.h0_bits, i, @floatFromInt(4096 + i % 7)); + for (0..3 * T * linear_dim) |i| f.soa_bits[i] = c.toBf16(1); // q, k + for (0..T * linear_dim) |i| { + if (i % dk % 2 == 1) f.soa_bits[i] = c.toBf16(-1); // q alternates + } + // With k=1 and state varying only by i%7, every head/row kv dot is the same ~524667 integer. + var kv: f64 = 0; + for (0..dk) |i| kv += @as(f64, @floatFromInt(4096 + i % 7)); + for (0..T * linear_dim) |i| f.soa_bits[2 * T * linear_dim + i] = c.toBf16(@floatCast(kv)); + for (0..T * heads) |i| { + putF32(f.g_pairs, i, 1.0); + f.beta_bits[i] = c.toBf16(0x1p-10); + } + }, + .decode_y, .decode_state => { + for (0..state_elems) |i| putF32(f.state_bits, i, @floatFromInt(4096 + i % 7)); + for (f.soa_bits) |*x| x.* = c.toBf16(1); // q, k = 1 + for (0..T * linear_dim) |i| { + if (i % dk % 2 == 1) f.soa_bits[i] = c.toBf16(-1); + } + // v mirrors the kv dot (state pinned, k = 1) + for (0..heads) |h| for (0..dk) |row| { + var kv: f64 = 0; + for (0..dk) |i| kv += @as(f64, @floatFromInt(4096 + @as(i64, @intCast((h * dk * dk + row * dk + i) % 7)))); + for (0..T) |t| f.soa_bits[(t * 3 + 2) * linear_dim + h * dk + row] = c.toBf16(@floatCast(kv)); + }; + for (0..T * heads) |i| { + f.a_bits[i] = c.toBf16(0); + f.b_bits[i] = c.toBf16(-10); // sigmoid(-10): a small convergent beta gain + } + for (0..heads) |h| { + f.dt_bits[h] = c.toBf16(0); + putF32(f.a_log_bits, h, -88.0); // exp(-88) softplus gates g to f32 1.0 + } + }, + } +} + +// `prepareHost` refuses shapes, generates fills, and computes references without a device for host tests. +pub fn prepareHost(f: *Gdn, shape: c.Shape, fill: g.Fill, seed: u64) !?usize { + if (shape.rows == 0 or shape.n != heads or shape.k != dk or shape.rows > max_tokens) return null; + f.reset(); + _ = f.arena.reset(.retain_capacity); + f.shape = shape; + f.tokens = shape.rows; + const T = f.tokens; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + f.chunks = @intCast((T + @as(usize, @intCast(chunk_fixed)) - 1) / @as(usize, @intCast(chunk_fixed))); + + // host inputs per which + switch (f.which) { + .split_soa, .split_g => { + f.conv_out_bits = try f.allocU16(T * channels, poison); + f.a_bits = try f.allocU16(T * heads, poison); + f.b_bits = try f.allocU16(T * heads, poison); + f.dt_bits = try f.allocU16(heads, poison); + f.a_log_bits = try f.allocU16(2 * heads, poison); + }, + .conv_prefill => { + f.state_bits = try f.allocU16(3 * channels, poison); + f.qkv_bits = try f.allocU16(T * channels, poison); + f.weight_bits = try f.allocU16(taps * channels, poison); + f.a_bits = try f.allocU16(T * heads, poison); + f.b_bits = try f.allocU16(T * heads, poison); + f.dt_bits = try f.allocU16(heads, poison); + f.a_log_bits = try f.allocU16(2 * heads, poison); + }, + .conv_decode => { + f.state_bits = try f.allocU16(3 * channels, poison); + f.state_init_bits = try f.allocU16(3 * channels, poison); + f.state_swap_bits = try f.allocU16(3 * channels, poison); + f.qkv_bits = try f.allocU16(T * channels, poison); // x rows + f.weight_bits = try f.allocU16(taps * channels, poison); + f.a_bits = try f.allocU16(T * heads, poison); + f.b_bits = try f.allocU16(T * heads, poison); + f.dt_bits = try f.allocU16(heads, poison); + f.a_log_bits = try f.allocU16(2 * heads, poison); + }, + .prefill_y, .prefill_state => { + f.soa_bits = try f.allocU16(3 * T * linear_dim, poison); + f.g_pairs = try f.allocU16(2 * T * heads, poison); + f.beta_bits = try f.allocU16(T * heads, poison); + f.h0_bits = try f.allocU16(2 * state_elems, poison); + f.walk_st = try f.allocF64(state_elems); + }, + .decode_y, .decode_state => { + f.soa_bits = try f.allocU16(3 * T * linear_dim, poison); + f.a_bits = try f.allocU16(T * heads, poison); + f.b_bits = try f.allocU16(T * heads, poison); + f.a_log_bits = try f.allocU16(2 * heads, poison); + f.dt_bits = try f.allocU16(heads, poison); + f.state_bits = try f.allocU16(2 * state_elems, poison); + f.state_init_bits = try f.allocU16(2 * state_elems, poison); + f.state_swap_bits = try f.allocU16(2 * state_elems, poison); + f.walk_st = try f.allocF64(state_elems); + @memcpy(f.state_init_bits, f.state_bits); + }, + } + + // fills: the fixture's dyadic ranges for gate 3, engineered integers for gate 4, planted classes for gate 1 + switch (fill) { + .random => switch (f.which) { + .split_soa, .split_g => { + fillDyadic(rng, f.conv_out_bits, -4, 2); + fillDyadic(rng, f.a_bits, -3, 2); + fillDyadic(rng, f.b_bits, -3, 2); + fillDyadic(rng, f.dt_bits, -4, 1); + // Served `a_log` keeps g normal; above ~2.4, devices flush subnormal g (census, not fill failure). + fillDyadicF32(rng, f.a_log_bits, -6, 0); + }, + .conv_prefill, .conv_decode => { + fillDyadic(rng, f.state_bits, -4, 2); + fillDyadic(rng, f.qkv_bits, -4, 2); + fillDyadic(rng, f.weight_bits, -3, 1); + fillDyadic(rng, f.a_bits, -3, 2); + fillDyadic(rng, f.b_bits, -3, 2); + fillDyadic(rng, f.dt_bits, -4, 1); + fillDyadicF32(rng, f.a_log_bits, -6, 0); + }, + .prefill_y, .prefill_state => { + // For 130 tokens, g, beta, soa <=2^-3 keep feedback <1; f32 state decays to the delta floor. + fillDyadic(rng, f.soa_bits, -6, -3); + fillDyadicF32(rng, f.g_pairs, -6, -3); + fillDyadic(rng, f.beta_bits, -6, -4); + fillDyadicF32(rng, f.h0_bits, -6, -2); + }, + .decode_y, .decode_state => { + fillDyadic(rng, f.soa_bits, -4, 2); + fillDyadic(rng, f.a_bits, -3, 2); + fillDyadic(rng, f.b_bits, -3, 2); + fillDyadic(rng, f.dt_bits, -4, 1); + fillDyadicF32(rng, f.a_log_bits, -4, 1); + fillDyadicF32(rng, f.state_bits, -6, -2); + }, + }, + .stress => fillStress(f), + .nonfinite => switch (f.which) { + .split_soa, .split_g => { + fillDyadic(rng, f.conv_out_bits, -4, 2); + fillDyadic(rng, f.a_bits, -3, 2); + fillDyadic(rng, f.b_bits, -3, 2); + fillDyadic(rng, f.dt_bits, -4, 1); + // Served `a_log` keeps g normal; above ~2.4, devices flush subnormal g (census, not fill failure). + fillDyadicF32(rng, f.a_log_bits, -6, 0); + }, + .conv_prefill, .conv_decode => { + fillDyadic(rng, f.state_bits, -4, 2); + fillDyadic(rng, f.qkv_bits, -4, 2); + fillDyadic(rng, f.weight_bits, -3, 1); + fillDyadic(rng, f.a_bits, -3, 2); + fillDyadic(rng, f.b_bits, -3, 2); + fillDyadic(rng, f.dt_bits, -4, 1); + fillDyadicF32(rng, f.a_log_bits, -6, 0); + }, + .prefill_y, .prefill_state => { + // For 130 tokens, g, beta, soa <=2^-3 keep feedback <1; f32 state decays to the delta floor. + fillDyadic(rng, f.soa_bits, -6, -3); + fillDyadicF32(rng, f.g_pairs, -6, -3); + fillDyadic(rng, f.beta_bits, -6, -4); + fillDyadicF32(rng, f.h0_bits, -6, -2); + }, + .decode_y, .decode_state => { + fillDyadic(rng, f.soa_bits, -4, 2); + fillDyadic(rng, f.a_bits, -3, 2); + fillDyadic(rng, f.b_bits, -3, 2); + fillDyadic(rng, f.dt_bits, -4, 1); + fillDyadicF32(rng, f.a_log_bits, -4, 1); + fillDyadicF32(rng, f.state_bits, -6, -2); + }, + }, + } + if (fill == .nonfinite) try plantNonfinite(f, rng); + // the decode loops' uploaded initial state: the FINAL generated bits + if (f.state_init_bits.len != 0) @memcpy(f.state_init_bits, f.state_bits); + + // `out` holds guarded u16 poison (2 words/f32); `spec.out` reads `count` output elements. + const out_len: usize = switch (f.which) { + .split_soa => 3 * T * linear_dim + T * heads, + .split_g => 2 * T * heads, + .conv_prefill, .conv_decode => T * channels + 3 * channels, + .prefill_y => T * linear_dim, + .prefill_state => 2 * state_elems * f.chunks, + .decode_y => T * linear_dim, + .decode_state => 2 * state_elems, + }; + const count_len: usize = switch (f.which) { + .split_g => T * heads, + .prefill_state => state_elems * f.chunks, + .decode_state => state_elems, + else => out_len, + }; + f.out = try f.allocU16(out_len + guard_words, poison); + f.out_words = out_len; + f.want = try f.allocF64(count_len); + f.mag = try f.allocF64(count_len); + try computeWant(f); + if (fill != .nonfinite) for (f.want, f.mag) |w, m| { + if (!std.math.isFinite(w) or !std.math.isFinite(m)) return error.NonFiniteReference; + }; + return out_len; +} + +/// The uploaded inputs in the fixed order the unchanged-check compares. +pub fn inputList(f: *Gdn, buf: *[8][]u16) []const []u16 { + switch (f.which) { + .split_soa, .split_g => { + buf[0..5].* = .{ f.conv_out_bits, f.a_bits, f.b_bits, f.dt_bits, f.a_log_bits }; + return buf[0..5]; + }, + .conv_prefill => { + buf[0..7].* = .{ f.state_bits, f.qkv_bits, f.weight_bits, f.a_bits, f.b_bits, f.dt_bits, f.a_log_bits }; + return buf[0..7]; + }, + // Decode initial state uses an unwritten buffer so unchanged checks see the exact first-step input. + .conv_decode => { + buf[0..7].* = .{ f.state_init_bits, f.qkv_bits, f.weight_bits, f.a_bits, f.b_bits, f.dt_bits, f.a_log_bits }; + return buf[0..7]; + }, + .prefill_y, .prefill_state => { + buf[0..4].* = .{ f.soa_bits, f.g_pairs, f.beta_bits, f.h0_bits }; + return buf[0..4]; + }, + .decode_y, .decode_state => { + buf[0..6].* = .{ f.soa_bits, f.a_bits, f.b_bits, f.a_log_bits, f.dt_bits, f.state_init_bits }; + return buf[0..6]; + }, + } +} + +// `outWords` uses fetched outputs' `out` slice or poison scratch; bad bounds report gate errors, not panics. +pub fn outWords(f: *Gdn, checked: bool, start: usize, words: usize) ![]u16 { + if (!checked) return f.allocU16(words, poison); + if (start > f.out_words or words > f.out_words - start) return error.OutputRegionOutOfRange; + return f.out[start .. start + words]; +} + +// `uploadWords` orders input bits/output poison; checked outputs use sized `out`, others fresh scratch. +pub fn uploadWords(f: *Gdn, buf: *[16][]u16) !usize { + const T = f.tokens; + var n: usize = 0; + switch (f.which) { + .split_soa, .split_g => { + // 0 conv_out, 1 a, 2 b, 3 dt, 4 a_log, 5 soa, 6 g, 7 beta + buf[0..5].* = .{ f.conv_out_bits, f.a_bits, f.b_bits, f.dt_bits, f.a_log_bits }; + n = 5; + buf[n] = try outWords(f, f.which == .split_soa, 0, 3 * T * linear_dim); + buf[n + 1] = try outWords(f, f.which == .split_g, 0, 2 * T * heads); + buf[n + 2] = try outWords(f, f.which == .split_soa, 3 * T * linear_dim, T * heads); + n += 3; + }, + .conv_prefill => { + // Slots 0-12 are state_in, qkv, weight, a, b, dt, a_log, conv_input, conv_out, soa, g, beta, state_out. + buf[0..7].* = .{ f.state_bits, f.qkv_bits, f.weight_bits, f.a_bits, f.b_bits, f.dt_bits, f.a_log_bits }; + n = 7; + buf[n] = try outWords(f, false, 0, (T + 3) * channels); + buf[n + 1] = try outWords(f, true, 0, T * channels); + buf[n + 2] = try outWords(f, false, 0, 3 * T * linear_dim); + buf[n + 3] = try outWords(f, false, 0, 2 * T * heads); + buf[n + 4] = try outWords(f, false, 0, T * heads); + buf[n + 5] = try outWords(f, true, T * channels, 3 * channels); + n += 6; + }, + .conv_decode => { + // Slots 0-12 are state_init, state_a, state_b, x, weight, a, b, dt, a_log, activated, soa, g, beta. + buf[0..9].* = .{ f.state_init_bits, f.state_swap_bits, f.state_swap_bits, f.qkv_bits, f.weight_bits, f.a_bits, f.b_bits, f.dt_bits, f.a_log_bits }; + n = 9; + buf[n] = try outWords(f, true, 0, T * channels); + buf[n + 1] = try outWords(f, false, 0, 3 * T * linear_dim); + buf[n + 2] = try outWords(f, false, 0, 2 * T * heads); + buf[n + 3] = try outWords(f, false, 0, T * heads); + n += 4; + }, + .prefill_y, .prefill_state => { + // 0 soa, 1 g, 2 beta, 3 h0, 4 y, 5 hf, 6 snap? + buf[0..4].* = .{ f.soa_bits, f.g_pairs, f.beta_bits, f.h0_bits }; + n = 4; + buf[n] = try outWords(f, f.which == .prefill_y, 0, T * linear_dim); + buf[n + 1] = try outWords(f, f.which == .prefill_state, 0, 2 * state_elems); + n += 2; + if (f.chunks > 1) { + buf[n] = try outWords(f, f.which == .prefill_state, 2 * state_elems, 2 * state_elems * (f.chunks - 1)); + n += 1; + } + }, + .decode_y, .decode_state => { + // 0 soa, 1 a, 2 b, 3 a_log, 4 dt, 5 state_init, 6 state_a, 7 state_b, 8 y + buf[0..8].* = .{ f.soa_bits, f.a_bits, f.b_bits, f.a_log_bits, f.dt_bits, f.state_init_bits, f.state_swap_bits, f.state_swap_bits }; + n = 8; + buf[n] = try outWords(f, f.which == .decode_y, 0, T * linear_dim); + n += 1; + }, + } + return n; +} + +// `computeWant` builds f64 `want/mag` from bits before device work; `reference()` returns the prepared copy. +pub fn computeWant(f: *Gdn) !void { + const T = f.tokens; + switch (f.which) { + .split_soa => splitReference(T, f.conv_out_bits, f.a_bits, f.b_bits, f.a_log_bits, f.dt_bits, f.want, f.mag), + .split_g => { + for (0..T * heads) |i| { + const ref = fx.gateRefOracle(f.a_bits[i], f.b_bits[i], f32Of(f.a_log_bits, i % heads), f.dt_bits[i % heads]); + f.want[i] = ref.g; + f.mag[i] = @abs(ref.g); + } + }, + .conv_prefill => { + const ctx: PrefillInput = .{ .state_bits = f.state_bits, .qkv_bits = f.qkv_bits }; + convReference(T, f.weight_bits, ctx, f.want[0 .. T * channels], f.mag[0 .. T * channels], f.want[T * channels ..]); + @memset(f.mag[T * channels ..], 0); + }, + .conv_decode => { + // the chain: each step taps the carried state then shifts it + for (0..T) |t| { + const ctx: DecodeInput = .{ .state_bits = f.state_bits, .x_bits = f.qkv_bits, .t = t }; + convReference(1, f.weight_bits, ctx, f.want[t * channels ..][0..channels], f.mag[t * channels ..][0..channels], null); + for (0..2) |r| @memcpy(f.state_bits[r * channels ..][0..channels], f.state_bits[(r + 1) * channels ..][0..channels]); + @memcpy(f.state_bits[2 * channels ..][0..channels], f.qkv_bits[t * channels ..][0..channels]); + } + // the final carry is the input rows T..T+2 of the walked chain + for (0..3 * channels) |i| { + f.want[T * channels + i] = bf(f.state_bits[i]); + f.mag[T * channels + i] = 0; + } + }, + .prefill_y, .prefill_state => { + f.max_state = 0; + const snap_chunk: ?usize = if (f.chunks > 1) @intCast(chunk_fixed) else null; + const y: ?[]f64 = if (f.which == .prefill_y) f.want else null; + const hf: ?[]f64 = if (f.which == .prefill_state) f.want[0..state_elems] else null; + const snap: ?[]f64 = if (f.which == .prefill_state and f.chunks > 1) f.want[state_elems..] else null; + walk(T, f.soa_bits, f.g_pairs, f.beta_bits, f.h0_bits, f.walk_st, y, hf, &f.max_state, snap_chunk, snap, null, null); + @memset(f.mag, f.max_state); + }, + .decode_y, .decode_state => { + f.max_state = 0; + const g_t = try f.allocU16(2 * heads, 0); + const beta_t = try f.allocU16(heads, 0); + const y_t = try f.allocF64(linear_dim); + const st_scratch = try f.allocF64(state_elems); + const st_out = try f.allocF64(state_elems); + for (0..T) |t| { + decodeGates(f.a_bits, f.b_bits, f.a_log_bits, f.dt_bits, t, g_t, beta_t); + walk(1, f.soa_bits[t * 3 * linear_dim ..][0 .. 3 * linear_dim], g_t, beta_t, f.state_bits, st_scratch, y_t, st_out, &f.max_state, null, null, null, null); + if (f.which == .decode_y) { + @memcpy(f.want[t * linear_dim ..][0..linear_dim], y_t); + } else { + // only the final carry is checked; the chain still walks every step + @memcpy(f.want[0..state_elems], st_out); + } + // the next step reads the f32 state this one left + for (0..state_elems) |i| putF32(f.state_bits, i, @floatCast(st_out[i])); + } + @memset(f.mag, f.max_state); + }, + } +} + +/// The device buffer index of every uploaded input, in inputList's order. +pub fn inputBufIndices(f: *Gdn, buf: *[8]usize) []const usize { + switch (f.which) { + .split_soa, .split_g => { + buf[0..5].* = .{ 0, 1, 2, 3, 4 }; + return buf[0..5]; + }, + .conv_prefill => { + buf[0..7].* = .{ 0, 1, 2, 3, 4, 5, 6 }; + return buf[0..7]; + }, + .conv_decode => { + buf[0..7].* = .{ 0, 3, 4, 5, 6, 7, 8 }; + return buf[0..7]; + }, + .prefill_y, .prefill_state => { + buf[0..4].* = .{ 0, 1, 2, 3 }; + return buf[0..4]; + }, + .decode_y, .decode_state => { + buf[0..6].* = .{ 0, 1, 2, 3, 4, 5 }; + return buf[0..6]; + }, + } +} diff --git a/zig/tests/vulkan/contract_gdn_main.zig b/zig/tests/vulkan/contract_gdn_main.zig new file mode 100644 index 00000000000..ed09d5eb6c6 --- /dev/null +++ b/zig/tests/vulkan/contract_gdn_main.zig @@ -0,0 +1,133 @@ +// Stage 4 GDN contract uses production Gdn and planted controls; gate 5 checks live device limits and records. +const std = @import("std"); +const vk = @import("vulkan"); +const shaders = @import("vulkan_shaders"); +const shaders4 = @import("stage4_shaders"); +const core = @import("core"); +const ctl = @import("contract_qmm.zig"); +const fx = @import("contract_gdn.zig"); +const g = core.gates; +const records = core.vulkan_records; + +// Shared bytes as each shader source declares (stage 4 doc block + prefill [2][DK] staging); values on the next line +const shared_bytes = [7]u64{ 0, 1024, 128, 3072, 3 * 128 * 4 + 32 * 129 * 4, 1024, 0 }; +const shared_names = [7][]const u8{ "copy_u16", "conv_dw1d_bf16", "gdn_conv_decode_bf16", "gdn_split_bf16", "prefill_recur_bf16", "gated_delta_decode_bf16", "fast_norm_gated_bf16" }; + +pub fn main(init: std.process.Init) !u8 { + var rt = try vk.Runtime.open(); + defer rt.close(); + var dev = try vk.Device.open(&rt, .{ .features = .{ .storage16 = true, .int16 = true } }); + defer dev.close(); + const props = dev.properties; + const lim = props.limits; + const name = std.mem.sliceTo(&props.deviceName, 0); + std.debug.print("LIMITS deviceName=\"{s}\"\n", .{name}); + std.debug.print("LIMITS driverName=\"{s}\" driverInfo=\"{s}\" subgroupSize={d} shuffle-relative={} storage16={} int16={}\n", .{ std.mem.sliceTo(&dev.v12.driverName, 0), std.mem.sliceTo(&dev.v12.driverInfo, 0), dev.v11.subgroupSize, dev.v11.subgroupSupportedOperations & vk.c.VK_SUBGROUP_FEATURE_SHUFFLE_RELATIVE_BIT != 0, dev.features.storage16, dev.features.int16 }); + std.debug.print("LIMITS maxComputeWorkGroupInvocations={d} maxComputeWorkGroupSize={d},{d},{d} maxComputeSharedMemorySize={d}\n", .{ lim.maxComputeWorkGroupInvocations, lim.maxComputeWorkGroupSize[0], lim.maxComputeWorkGroupSize[1], lim.maxComputeWorkGroupSize[2], lim.maxComputeSharedMemorySize }); + std.debug.print("LIMITS maxStorageBufferRange={d} minStorageBufferOffsetAlignment={d}\n", .{ lim.maxStorageBufferRange, lim.minStorageBufferOffsetAlignment }); + if (props.deviceType == vk.c.VK_PHYSICAL_DEVICE_TYPE_CPU) { + std.debug.print("FAIL gdn-contract: a CPU (software) Vulkan device is never a GPU result\n", .{}); + return 1; + } + var problems: usize = 0; + var gate5_controls: usize = 0; + var tally: g.Tally = .{}; + for (records.records) |record| { + const local = record.launch.local_size; + const invocations: u64 = @as(u64, local[0]) * local[1] * local[2]; + const fits = invocations <= lim.maxComputeWorkGroupInvocations and local[0] <= lim.maxComputeWorkGroupSize[0] and local[1] <= lim.maxComputeWorkGroupSize[1] and local[2] <= lim.maxComputeWorkGroupSize[2] and record.launch.shared_bytes <= lim.maxComputeSharedMemorySize; + const expected_fail = record.control_gate == 5; + const caught = fits != expected_fail; + std.debug.print("{s} limits {s}: local={d}x{d}x{d} invocations={d} shared={d}\n", .{ if (caught) "PASS" else "FAIL", record.launch.name, local[0], local[1], local[2], invocations, record.launch.shared_bytes }); + if (!caught) problems += 1; + if (expected_fail) { + if (caught) gate5_controls += 1 else { + tally.missed += 1; + std.debug.print("CONTROL MISSED {s}: gate 5 did not reject the oversized declaration\n", .{record.launch.name}); + } + } + } + var cmds = try vk.Commands.init(&dev); + defer cmds.deinit(); + var staging = try vk.Buffer.init(&dev, .staging, 32 << 20); // the T=130 prefill's snapshot slot + defer staging.deinit(); + var m = try vk.Gdn.init(&dev, .{ + .conv_dw1d_bf16 = &shaders4.conv_dw1d_bf16, + .gdn_conv_decode_bf16 = &shaders4.gdn_conv_decode_bf16, + .gdn_split_bf16 = &shaders4.gdn_split_bf16, + .prefill_recur_bf16 = &shaders4.prefill_recur_bf16, + .gated_delta_decode_bf16 = &shaders4.gated_delta_decode_bf16, + .fast_norm_gated_bf16 = &shaders4.fast_norm_gated_bf16, + .copy_u16 = &shaders4.copy_u16, + .conv_dw1d_f32 = &shaders4.conv_dw1d_f32, + .gdn_split_f32 = &shaders4.gdn_split_f32, + .prefill_recur_f32 = &shaders4.prefill_recur_f32, + .fast_norm_gated_f32 = &shaders4.fast_norm_gated_f32, + .copy_u32 = &shaders4.gdn_copy_u32, + }); + defer m.deinit(); + // Gate 5 checks every reflected local size and each shader's declared shared bytes against this device's limits. + const locals = [7][3]u32{ m.copy_u16.local, m.conv_bf16.local, m.conv_decode_bf16.local, m.split_bf16.local, m.prefill_bf16.local, m.decode_bf16.local, m.norm_bf16.local }; + for (locals, shared_bytes, shared_names) |local, shared, sname| { + const invocations: u64 = @as(u64, local[0]) * local[1] * local[2]; + const fits = invocations <= lim.maxComputeWorkGroupInvocations and local[0] <= lim.maxComputeWorkGroupSize[0] and local[1] <= lim.maxComputeWorkGroupSize[1] and local[2] <= lim.maxComputeWorkGroupSize[2] and shared <= lim.maxComputeSharedMemorySize; + std.debug.print("{s} limits gdn {s}: local={d}x{d}x{d} invocations={d} shared={d}\n", .{ if (fits) "PASS" else "FAIL", sname, local[0], local[1], local[2], invocations, shared }); + if (!fits) problems += 1; + } + + var split_soa = fx.Gdn.init(init.gpa, &dev, &cmds, staging, &m, .split_soa); + defer split_soa.deinit(); + var split_g = fx.Gdn.init(init.gpa, &dev, &cmds, staging, &m, .split_g); + defer split_g.deinit(); + var conv_pre = fx.Gdn.init(init.gpa, &dev, &cmds, staging, &m, .conv_prefill); + defer conv_pre.deinit(); + var conv_dec = fx.Gdn.init(init.gpa, &dev, &cmds, staging, &m, .conv_decode); + defer conv_dec.deinit(); + var pre_y = fx.Gdn.init(init.gpa, &dev, &cmds, staging, &m, .prefill_y); + defer pre_y.deinit(); + var pre_state = fx.Gdn.init(init.gpa, &dev, &cmds, staging, &m, .prefill_state); + defer pre_state.deinit(); + var dec_y = fx.Gdn.init(init.gpa, &dev, &cmds, staging, &m, .decode_y); + defer dec_y.deinit(); + var dec_state = fx.Gdn.init(init.gpa, &dev, &cmds, staging, &m, .decode_state); + defer dec_state.deinit(); + + const control_layout: vk.dispatch.Layout = .{ .bindings = 3, .push_bytes = 8, .writes = 1 << 2 }; + var p_fmax = try vk.dispatch.Pipeline.init(&dev, &shaders.qmm_fmax, control_layout); + defer p_fmax.deinit(); + var p_atomic = try vk.dispatch.Pipeline.init(&dev, &shaders.atomic_sum, control_layout); + defer p_atomic.deinit(); + var p_scribble = try vk.dispatch.Pipeline.init(&dev, &shaders.scribble_sum, control_layout); + defer p_scribble.deinit(); + var p_tail = try vk.dispatch.Pipeline.init(&dev, &shaders.drop_tail, control_layout); + defer p_tail.deinit(); + var p_fp16 = try vk.dispatch.Pipeline.init(&dev, &shaders.fp16_accum, control_layout); + defer p_fp16.deinit(); + var controls = [_]ctl.Control{ + ctl.Control.init(init.gpa, &dev, &cmds, staging, &p_fmax, .qmm_fmax), + ctl.Control.init(init.gpa, &dev, &cmds, staging, &p_atomic, .atomic_sum), + ctl.Control.init(init.gpa, &dev, &cmds, staging, &p_scribble, .scribble_sum), + ctl.Control.init(init.gpa, &dev, &cmds, staging, &p_tail, .drop_tail), + ctl.Control.init(init.gpa, &dev, &cmds, staging, &p_fp16, .fp16_accum), + }; + defer for (&controls) |*control| control.deinit(); + + var kernels: [13]g.Kernel = undefined; + kernels[0] = split_soa.kernel(); + kernels[1] = split_g.kernel(); + kernels[2] = conv_pre.kernel(); + kernels[3] = conv_dec.kernel(); + kernels[4] = pre_y.kernel(); + kernels[5] = pre_state.kernel(); + kernels[6] = dec_y.kernel(); + kernels[7] = dec_state.kernel(); + for (&controls, 8..) |*control, i| kernels[i] = control.kernel(); + const start = std.Io.Clock.awake.now(init.io); + tally.controls += gate5_controls; + tally.caught += gate5_controls; + try g.runAll(init.gpa, &kernels, 0x6764, &tally); + const elapsed = start.durationTo(std.Io.Clock.awake.now(init.io)).toMilliseconds(); + std.debug.print("SUMMARY {s} gdn group: {d} kernels: {d} pass gates 1-4, {d} fail, {d} known findings; {d} controls: {d} caught, {d} missed; {d} limit failures; gates {d} ms\n", .{ name, tally.kernels, tally.passed, tally.failed, tally.known, tally.controls, tally.caught, tally.missed, problems, elapsed }); + if (tally.problems() != 0 or tally.missed != 0 or problems != 0) return 1; + return 0; +} diff --git a/zig/tests/vulkan/contract_gdn_ops.zig b/zig/tests/vulkan/contract_gdn_ops.zig new file mode 100644 index 00000000000..5b3cc58c820 --- /dev/null +++ b/zig/tests/vulkan/contract_gdn_ops.zig @@ -0,0 +1,451 @@ +// The GDN contract's device layer: the fixture's buffers, the uploads and the launches. +const std = @import("std"); +const vk = @import("vulkan"); +const contract = @import("contract"); +const fx = @import("stage4_fixture.zig"); +const refs = @import("contract_gdn_ref.zig"); +const host = @import("contract_gdn_host.zig"); + +const c = contract.spec; +const g = contract.gates; +const rows = contract.rows; + +const Which = refs.Which; +const heads = refs.heads; +const dk = refs.dk; +const linear_dim = refs.linear_dim; +const channels = refs.channels; +const taps = refs.taps; +const state_elems = refs.state_elems; +const chunk_fixed = refs.chunk_fixed; +const timeout_ns = refs.timeout_ns; +const spec = refs.spec; +const walk = refs.walk; +const decodeGates = refs.decodeGates; +const f32Of = refs.f32Of; +const bf = refs.bf; +const recurYDepth = refs.recurYDepth; +const recurStateDepth = refs.recurStateDepth; +const decodeYDepth = refs.decodeYDepth; +const decodeStateDepth = refs.decodeStateDepth; +const snapDepth = refs.snapDepth; +const Rig = refs.Rig; +const poison = refs.poison; +const putF32 = refs.putF32; +const PrefillInput = refs.PrefillInput; +const DecodeInput = refs.DecodeInput; + +pub const Gdn = struct { + a: std.mem.Allocator, + dev: *const vk.Device, + cmds: *vk.Commands, + staging: vk.Buffer, + m: *vk.Gdn, + which: Which, + arena: std.heap.ArenaAllocator, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + tokens: usize = 0, + chunks: usize = 1, + max_state: f64 = 0, + + // host inputs (the uploaded bits, kept for the references) + conv_out_bits: []u16 = &.{}, + a_bits: []u16 = &.{}, + b_bits: []u16 = &.{}, + dt_bits: []u16 = &.{}, + a_log_bits: []u16 = &.{}, + state_bits: []u16 = &.{}, // conv state (bf16) or recurrence state (f32 pairs) + state_init_bits: []u16 = &.{}, // the decode loops' uploaded initial state, never written + state_swap_bits: []u16 = &.{}, // poison for the decode loops' scratch carry buffers + qkv_bits: []u16 = &.{}, // conv_prefill projection rows; conv_decode/decode x rows + weight_bits: []u16 = &.{}, + soa_bits: []u16 = &.{}, + g_pairs: []u16 = &.{}, + beta_bits: []u16 = &.{}, + h0_bits: []u16 = &.{}, + walk_st: []f64 = &.{}, // the walker's state scratch + + // host outputs and the contract bookkeeping + out: []u16 = &.{}, + out_words: usize = 0, // out's data width in u16 words (f32 outputs take two a word) + want: []f64 = &.{}, + mag: []f64 = &.{}, + input_snapshot: []u8 = &.{}, + input_now: []u8 = &.{}, + bufs: [16]vk.Buffer = undefined, + made: usize = 0, + /// The decode loops' carry buffer holding the final state after run(). + final_state_idx: usize = 0, + rig: Rig = .{}, + + pub fn init(a: std.mem.Allocator, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, m: *vk.Gdn, which: Which) Gdn { + return .{ .a = a, .dev = dev, .cmds = cmds, .staging = staging, .m = m, .which = which, .arena = .init(a) }; + } + + pub fn deinit(f: *Gdn) void { + f.reset(); + f.arena.deinit(); + } + + pub fn reset(f: *Gdn) void { + while (f.made > 0) { + f.made -= 1; + f.bufs[f.made].deinit(); + } + _ = f.arena.reset(.retain_capacity); + f.conv_out_bits = &.{}; + f.a_bits = &.{}; + f.b_bits = &.{}; + f.dt_bits = &.{}; + f.a_log_bits = &.{}; + f.state_bits = &.{}; + f.state_init_bits = &.{}; + f.state_swap_bits = &.{}; + f.qkv_bits = &.{}; + f.weight_bits = &.{}; + f.soa_bits = &.{}; + f.g_pairs = &.{}; + f.beta_bits = &.{}; + f.h0_bits = &.{}; + f.walk_st = &.{}; + f.out = &.{}; + f.want = &.{}; + f.mag = &.{}; + f.input_snapshot = &.{}; + f.input_now = &.{}; + f.rig = .{}; + f.shape = .{ .rows = 0, .n = 0, .k = 0 }; + f.tokens = 0; + f.chunks = 1; + f.max_state = 0; + } + + const vt = g.hooks(Gdn); + pub fn kernel(f: *Gdn) g.Kernel { + return .{ .spec = spec(f.which), .ctx = f, .vt = &vt }; + } + + pub fn allocU16(f: *Gdn, n: usize, fill: ?u16) ![]u16 { + const v = try f.arena.allocator().alloc(u16, n); + if (fill) |p| @memset(v, p) else for (v) |*x| x.* = 0; + return v; + } + + pub fn allocF64(f: *Gdn, n: usize) ![]f64 { + return f.arena.allocator().alloc(f64, n); + } + + // `devBuf` uploads input `bytes` or poisons output `poison_bytes`, registers cleanup, and returns its index. + fn devBuf(f: *Gdn, bytes: []const u8) !usize { + const b = try vk.Buffer.init(f.dev, .device, bytes.len); + f.bufs[f.made] = b; + f.made += 1; + try vk.orExit(vk.upload(f.cmds, f.bufs[f.made - 1], 0, f.staging, bytes, timeout_ns)); + return f.made - 1; + } + + fn view(f: *Gdn, idx: usize, kind: vk.tensor.Kind, elements: u64) !vk.tensor.View { + return vk.tensor.View.init(&f.bufs[idx], kind, 0, elements); + } + + pub fn prepare(f: *Gdn, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + const out_len = (try f.prepareHost(shape, fill, seed)) orelse return false; + const a = f.arena.allocator(); + + // Fixed per-`which` buffers upload input bits and poisoned output regions or scratch. + var input_slices: [8][]u16 = undefined; + const inputs = f.inputList(&input_slices); + var input_total: usize = 0; + for (inputs) |s| input_total += std.mem.sliceAsBytes(s).len; + try f.uploadBuffers(); + f.input_snapshot = try a.alloc(u8, input_total); + f.input_now = try a.alloc(u8, input_total); + var pos: usize = 0; + for (inputs) |s| { + const bytes = std.mem.sliceAsBytes(s); + @memcpy(f.input_snapshot[pos..][0..bytes.len], bytes); + pos += bytes.len; + } + f.rig = .{ .output = std.mem.sliceAsBytes(f.out[0..out_len]), .inputs = f.input_snapshot, .guards_ok = true }; + return true; + } + + // The host fills and the upload order live in contract_gdn_host.zig; these keep the fixture's own call sites. + pub fn prepareHost(f: *Gdn, shape: c.Shape, fill: g.Fill, seed: u64) !?usize { + return host.prepareHost(f, shape, fill, seed); + } + + pub fn inputList(f: *Gdn, buf: *[8][]u16) []const []u16 { + return host.inputList(f, buf); + } + + pub fn uploadWords(f: *Gdn, buf: *[16][]u16) !usize { + return host.uploadWords(f, buf); + } + + // `uploadWords`/`run` share indices; `out` poison catches dead stores, and other outputs use fresh scratch. + fn uploadBuffers(f: *Gdn) !void { + var words: [16][]u16 = undefined; + const n = try f.uploadWords(&words); + for (words[0..n]) |w| _ = try f.devBuf(std.mem.sliceAsBytes(w)); + } + + pub fn run(f: *Gdn) !void { + const T = f.tokens; + switch (f.which) { + .split_soa, .split_g => { + const conv_out = try f.view(0, .bf16, @intCast(T * channels)); + const a_v = try f.view(1, .bf16, @intCast(T * heads)); + const b_v = try f.view(2, .bf16, @intCast(T * heads)); + const dt_v = try f.view(3, .bf16, heads); + const a_log_v = try f.view(4, .f32, heads); + const soa_v = try f.view(5, .bf16, @intCast(3 * T * linear_dim)); + const g_v = try f.view(6, .f32, @intCast(T * heads)); + const beta_v = try f.view(7, .bf16, @intCast(T * heads)); + try vk.orExit(f.m.split(f.cmds, soa_v, g_v, beta_v, conv_out, a_v, b_v, a_log_v, dt_v, T, timeout_ns)); + }, + .conv_prefill => { + const state_v = try f.view(0, .bf16, @intCast(3 * channels)); + const qkv_v = try f.view(1, .bf16, @intCast(T * channels)); + const weight_v = try f.view(2, .bf16, @intCast(taps * channels)); + const a_v = try f.view(3, .bf16, @intCast(T * heads)); + const b_v = try f.view(4, .bf16, @intCast(T * heads)); + const dt_v = try f.view(5, .bf16, heads); + const a_log_v = try f.view(6, .f32, heads); + const conv_input_v = try f.view(7, .bf16, @intCast((T + 3) * channels)); + const conv_out_v = try f.view(8, .bf16, @intCast(T * channels)); + const soa_v = try f.view(9, .bf16, @intCast(3 * T * linear_dim)); + const g_v = try f.view(10, .f32, @intCast(T * heads)); + const beta_v = try f.view(11, .bf16, @intCast(T * heads)); + const state_out_v = try f.view(12, .bf16, @intCast(3 * channels)); + try vk.orExit(f.m.convPrefill(f.cmds, soa_v, g_v, beta_v, conv_out_v, conv_input_v, state_v, qkv_v, weight_v, a_v, b_v, a_log_v, dt_v, state_out_v, T, timeout_ns)); + }, + .conv_decode => { + const x_v = try f.view(3, .bf16, @intCast(T * channels)); + const weight_v = try f.view(4, .bf16, @intCast(taps * channels)); + const a_v = try f.view(5, .bf16, @intCast(T * heads)); + const b_v = try f.view(6, .bf16, @intCast(T * heads)); + const dt_v = try f.view(7, .bf16, heads); + const a_log_v = try f.view(8, .f32, heads); + const act_v = try f.view(9, .bf16, @intCast(T * channels)); + const soa_v = try f.view(10, .bf16, @intCast(3 * T * linear_dim)); + const g_v = try f.view(11, .f32, @intCast(T * heads)); + const beta_v = try f.view(12, .bf16, @intCast(T * heads)); + var in_idx: usize = 0; // state_init; carry buffers are 1 and 2 + for (0..T) |t| { + const out_idx: usize = if (in_idx == 1) 2 else 1; + const st_in = try f.view(in_idx, .bf16, @intCast(3 * channels)); + const st_out = try f.view(out_idx, .bf16, @intCast(3 * channels)); + try vk.orExit(f.m.convDecodeStep(f.cmds, try soa_v.slice(t * 3 * linear_dim, 3 * linear_dim), try g_v.slice(t * heads, heads), try beta_v.slice(t * heads, heads), try act_v.slice(t * channels, channels), st_in, try x_v.slice(t * channels, channels), weight_v, try a_v.slice(t * heads, heads), try b_v.slice(t * heads, heads), a_log_v, dt_v, st_out, timeout_ns)); + in_idx = out_idx; + } + f.final_state_idx = in_idx; // the initial buffer at 0 was never written + }, + .prefill_y, .prefill_state => { + const soa_v = try f.view(0, .bf16, @intCast(3 * T * linear_dim)); + const g_v = try f.view(1, .f32, @intCast(T * heads)); + const beta_v = try f.view(2, .bf16, @intCast(T * heads)); + const h0_v = try f.view(3, .f32, state_elems); + const y_v = try f.view(4, .bf16, @intCast(T * linear_dim)); + const hf_v = try f.view(5, .f32, state_elems); + const snap_v: ?vk.tensor.View = if (f.chunks > 1) try f.view(6, .f32, state_elems * (f.chunks - 1)) else null; + try vk.orExit(f.m.prefillRecur(f.cmds, soa_v, g_v, beta_v, h0_v, y_v, hf_v, snap_v, T, chunk_fixed, timeout_ns)); + }, + .decode_y, .decode_state => { + const a_v = try f.view(1, .bf16, @intCast(T * heads)); + const b_v = try f.view(2, .bf16, @intCast(T * heads)); + const a_log_v = try f.view(3, .u16, 2 * heads); + const dt_v = try f.view(4, .bf16, heads); + const y_v = try f.view(8, .bf16, @intCast(T * linear_dim)); + var in_idx: usize = 5; // state_init + for (0..T) |t| { + const st_in = try f.view(in_idx, .f32, state_elems); + const out_idx: usize = if (in_idx == 5) 6 else if (in_idx == 6) 7 else 6; + const st_out = try f.view(out_idx, .f32, state_elems); + try vk.orExit(f.m.decodeRecur(f.cmds, try f.soaView(t), try a_v.slice(t * heads, heads), try b_v.slice(t * heads, heads), a_log_v, dt_v, st_in, st_out, try y_v.slice(t * linear_dim, linear_dim), timeout_ns)); + in_idx = out_idx; + } + f.final_state_idx = in_idx; + }, + } + try f.fetchOutputs(); + } + + fn soaView(f: *Gdn, t: usize) !vk.tensor.View { + const soa_v = try f.view(0, .bf16, @intCast(3 * f.tokens * linear_dim)); + return soa_v.slice(t * 3 * linear_dim, 3 * linear_dim); + } + + // Fetch checked outputs, reread inputs for unchanged checks, then verify guard words. + fn fetchOutputs(f: *Gdn) !void { + const T = f.tokens; + switch (f.which) { + .split_soa => { + try vk.download(f.cmds, f.bufs[5], 0, f.staging, std.mem.sliceAsBytes(f.out[0 .. 3 * T * linear_dim]), timeout_ns); + try vk.download(f.cmds, f.bufs[7], 0, f.staging, std.mem.sliceAsBytes(f.out[3 * T * linear_dim .. 3 * T * linear_dim + T * heads]), timeout_ns); + }, + .split_g => { + try vk.download(f.cmds, f.bufs[6], 0, f.staging, std.mem.sliceAsBytes(f.out[0 .. 2 * T * heads]), timeout_ns); + }, + .conv_prefill => { + try vk.download(f.cmds, f.bufs[8], 0, f.staging, std.mem.sliceAsBytes(f.out[0 .. T * channels]), timeout_ns); + try vk.download(f.cmds, f.bufs[12], 0, f.staging, std.mem.sliceAsBytes(f.out[T * channels .. T * channels + 3 * channels]), timeout_ns); + }, + .conv_decode => { + try vk.download(f.cmds, f.bufs[9], 0, f.staging, std.mem.sliceAsBytes(f.out[0 .. T * channels]), timeout_ns); + try vk.download(f.cmds, f.bufs[f.final_state_idx], 0, f.staging, std.mem.sliceAsBytes(f.out[T * channels .. T * channels + 3 * channels]), timeout_ns); + }, + .prefill_y => { + try vk.download(f.cmds, f.bufs[4], 0, f.staging, std.mem.sliceAsBytes(f.out[0 .. T * linear_dim]), timeout_ns); + }, + .prefill_state => { + try vk.download(f.cmds, f.bufs[5], 0, f.staging, std.mem.sliceAsBytes(f.out[0 .. 2 * state_elems]), timeout_ns); + if (f.chunks > 1) try vk.download(f.cmds, f.bufs[6], 0, f.staging, std.mem.sliceAsBytes(f.out[2 * state_elems .. 2 * state_elems * f.chunks]), timeout_ns); + }, + .decode_y => { + try vk.download(f.cmds, f.bufs[8], 0, f.staging, std.mem.sliceAsBytes(f.out[0 .. T * linear_dim]), timeout_ns); + }, + .decode_state => { + try vk.download(f.cmds, f.bufs[f.final_state_idx], 0, f.staging, std.mem.sliceAsBytes(f.out[0 .. 2 * state_elems]), timeout_ns); + }, + } + // inputs, post-run, in the snapshot's order + var slices: [8][]u16 = undefined; + const inputs = f.inputList(&slices); + var idxs: [8]usize = undefined; + const dev_idx = host.inputBufIndices(f, &idxs); + var pos: usize = 0; + for (inputs, dev_idx) |s, di| { + try vk.download(f.cmds, f.bufs[di], 0, f.staging, std.mem.sliceAsBytes(s), timeout_ns); + const bytes = std.mem.sliceAsBytes(s); + @memcpy(f.input_now[pos..][0..bytes.len], bytes); + pos += bytes.len; + } + f.rig.inputs = f.input_now; + f.rig.guards_ok = true; + for (f.out[f.out_words..]) |w| if (w != poison) { + f.rig.guards_ok = false; + break; + }; + } + + // Gate 4 replays each dot in `accum`/`order`, sums other ops in f64, and writes `out` in `reference()` layout. + pub fn simulate(f: *Gdn, accum: c.Accum, order: g.Order, out: []f64) !void { + const T = f.tokens; + switch (f.which) { + .split_soa => { + var terms: [dk]f64 = undefined; + for (0..T) |t| { + for (0..3 * heads) |head| { + const hv: usize = head % heads; + var vals: [dk]f64 = undefined; + var vmax: f64 = 0; + for (0..dk) |i| { + const cv = bf(f.conv_out_bits[t * channels + head * dk + i]); + vals[i] = cv / (1.0 + @exp(-cv)); + vmax = @max(vmax, @abs(vals[i])); + } + for (0..dk) |i| terms[i] = vals[i] * vals[i]; + const ss = rows.sumIn(accum, order, &terms); + const inv: f64 = 1.0 / @sqrt(ss / 128.0 + 1e-6 / 128.0); + const scale_f: f64 = if (head < heads) 1.0 / 128.0 else @sqrt(1.0 / 128.0); + const dst = (head / heads) * (T * linear_dim) + (t * heads + hv) * dk; + for (0..dk) |i| out[dst + i] = if (head < 2 * heads) scale_f * (vals[i] * inv) else vals[i]; + } + for (0..heads) |hv| { + const ref = fx.gateRefOracle(f.a_bits[t * heads + hv], f.b_bits[t * heads + hv], f32Of(f.a_log_bits, hv), f.dt_bits[hv]); + out[3 * T * linear_dim + t * heads + hv] = ref.beta; + } + } + }, + .split_g => return error.Unimplemented, // accum none; gate 4 is N/A + .conv_prefill => { + const ctx: PrefillInput = .{ .state_bits = f.state_bits, .qkv_bits = f.qkv_bits }; + var terms: [taps]f64 = undefined; + for (0..channels) |ch| { + for (0..T) |t| { + for (0..taps) |k| terms[k] = bf(ctx.inWord(t + k, ch)) * bf(f.weight_bits[ch * taps + k]); + out[t * channels + ch] = rows.sumIn(accum, order, &terms); + } + for (0..3) |j| out[T * channels + j * channels + ch] = bf(ctx.inWord(T + j, ch)); + } + }, + .conv_decode => { + // the chain replays from the UPLOADED initial state + var terms: [taps]f64 = undefined; + const chain_a = try f.allocU16(3 * channels, 0); + const chain_b = try f.allocU16(3 * channels, 0); + @memcpy(chain_a, f.state_init_bits); + var chain: []u16 = chain_a; + for (0..T) |t| { + const ctx: DecodeInput = .{ .state_bits = chain, .x_bits = f.qkv_bits, .t = t }; + for (0..channels) |ch| { + for (0..taps) |k| terms[k] = bf(ctx.inWord(t + k, ch)) * bf(f.weight_bits[ch * taps + k]); + out[t * channels + ch] = rows.sumIn(accum, order, &terms); + } + const nxt: []u16 = if (chain.ptr == chain_a.ptr) chain_b else chain_a; + for (0..2) |r| @memcpy(nxt[r * channels ..][0..channels], chain[(r + 1) * channels ..][0..channels]); + @memcpy(nxt[2 * channels ..][0..channels], f.qkv_bits[t * channels ..][0..channels]); + chain = nxt; + } + for (0..3 * channels) |i| out[T * channels + i] = bf(chain[i]); + }, + .prefill_y => walk(T, f.soa_bits, f.g_pairs, f.beta_bits, f.h0_bits, f.walk_st, out, null, &f.max_state, null, null, accum, order), + .prefill_state => { + const snap_chunk: ?usize = if (f.chunks > 1) @intCast(chunk_fixed) else null; + walk(T, f.soa_bits, f.g_pairs, f.beta_bits, f.h0_bits, f.walk_st, null, out[0..state_elems], &f.max_state, snap_chunk, if (f.chunks > 1) out[state_elems..] else null, accum, order); + }, + .decode_y, .decode_state => { + const g_t = try f.allocU16(2 * heads, 0); + const beta_t = try f.allocU16(heads, 0); + const y_t: ?[]f64 = if (f.which == .decode_y) out else null; + const st_scratch = try f.allocF64(state_elems); + const st_out = try f.allocF64(state_elems); + const bits = try f.allocU16(2 * state_elems, 0); + @memcpy(bits, f.state_init_bits); + for (0..T) |t| { + decodeGates(f.a_bits, f.b_bits, f.a_log_bits, f.dt_bits, t, g_t, beta_t); + walk(1, f.soa_bits[t * 3 * linear_dim ..][0 .. 3 * linear_dim], g_t, beta_t, bits, st_scratch, if (y_t) |y| y[t * linear_dim ..][0..linear_dim] else null, st_out, &f.max_state, null, null, accum, order); + if (f.which == .decode_state) @memcpy(out[0..state_elems], st_out); + for (0..state_elems) |i| putF32(bits, i, @floatCast(st_out[i])); + } + }, + } + } + + // Bounds use stage 4 depth/ops/mids and magnitude scales: `max_state`, row `vmax^2`, `|acc|`, and `|want|`. + pub fn bounds(f: *Gdn, want: []const f64, mag: []const f64, allow: []f64) !void { + switch (f.which) { + .split_soa => for (want, mag, allow) |w, m, *l| { + l.* = fx.budget(40, 5, 3, w, m); + }, + .split_g => for (want, mag, allow) |w, m, *l| { + l.* = fx.budget(1, 4, 2, w, m); + }, + .conv_prefill, .conv_decode => for (want, mag, allow) |w, m, *l| { + l.* = fx.budget(taps, 0, 0, w, m); + }, + .prefill_y => for (want, mag, allow) |w, m, *l| { + l.* = fx.budget(@floatFromInt(recurYDepth(f.shape)), 4, 0, w, m); + }, + .prefill_state => for (want, mag, 0..) |w, m, i| { + allow[i] = fx.budget(if (i < state_elems) @floatFromInt(recurStateDepth(f.shape)) else snapDepth(), 2, 0, w, m); + }, + .decode_y => for (want, mag, allow) |w, m, *l| { + l.* = fx.budget(@floatFromInt(decodeYDepth(f.shape)), 4, 0, w, m); + }, + .decode_state => for (want, mag, allow) |w, m, *l| { + l.* = fx.budget(@floatFromInt(decodeStateDepth(f.shape)), 2, 0, w, m); + }, + } + } + + pub fn count(f: *Gdn) usize { + return f.want.len; + } + + pub fn reference(f: *Gdn, want: []f64, mag: []f64) !void { + @memcpy(want, f.want); + @memcpy(mag, f.mag); + } +}; diff --git a/zig/tests/vulkan/contract_gdn_ref.zig b/zig/tests/vulkan/contract_gdn_ref.zig new file mode 100644 index 00000000000..efa22b66194 --- /dev/null +++ b/zig/tests/vulkan/contract_gdn_ref.zig @@ -0,0 +1,347 @@ +// The GDN contract's reference layer: the served shapes, the specs and the f64 references the gates check. +const std = @import("std"); +const contract = @import("contract"); +const fx = @import("stage4_fixture.zig"); + +const c = contract.spec; +const g = contract.gates; +const rows = contract.rows; + +const Model = fx.Model; +pub const heads: usize = @intCast(Model.heads); // 16 +pub const dk: usize = @intCast(Model.head_dim); // 128 +pub const linear_dim: usize = @intCast(Model.linear_dim); // 2048 +pub const channels: usize = @intCast(Model.conv_channels); // 6144 +pub const taps: usize = @intCast(Model.taps); // 4 +pub const state_elems: usize = @intCast(Model.state_elements); // 262144 f32 + +// The served prefill chunk: snapshots appear from 9 tokens on, one interior boundary per additional chunk. +pub const chunk_fixed: u64 = 8; +pub const timeout_ns: u64 = 30 * std.time.ns_per_s; +pub const poison: u16 = 0xffff; +pub const guard_words: usize = 8; +pub const max_tokens: usize = 4096; // host allocation cap on generated edges + +pub const Which = enum { split_soa, split_g, conv_prefill, conv_decode, prefill_y, prefill_state, decode_y, decode_state }; + +pub fn splitDepth(_: c.Shape) usize { + return 40; // the 32-lane square sum, the stage 4 fixture's stated depth +} +pub fn gateDepth(_: c.Shape) usize { + return 1; +} +pub fn tapsDepth(_: c.Shape) usize { + return taps; +} +// Recurrence depth adds about 130 f32 roundings per token plus one output-dot depth; stage 4 derives 130 per update. +pub fn recurYDepth(s: c.Shape) usize { + return s.rows * 130 + 130; +} +pub fn recurStateDepth(s: c.Shape) usize { + return s.rows * 130 + 2; +} +// Decode repeats prefill per-token depth; each 130-rounding step includes the 4-op raw-gates prologue. +pub fn decodeYDepth(s: c.Shape) usize { + return s.rows * 260 + 130; +} +pub fn decodeStateDepth(s: c.Shape) usize { + return s.rows * 132 + 2; +} +// A snapshot slot's depth: `chunk_fixed` tokens from the chunk's start, plus the store. +pub fn snapDepth() f64 { + return @floatFromInt(@as(usize, @intCast(chunk_fixed)) * 130 + 2); +} + +const token_tile = c.Tile{ .size = 8, .step = 1, .min = 1 }; + +fn servedOf(which: Which) []const c.Shape { + return switch (which) { + .split_soa, .split_g, .conv_prefill, .conv_decode, .decode_y, .decode_state => &.{ + .{ .rows = 8, .n = heads, .k = dk }, + .{ .rows = 1, .n = heads, .k = dk }, + .{ .rows = 32, .n = heads, .k = dk }, + }, + .prefill_y, .prefill_state => &.{ + .{ .rows = 4, .n = heads, .k = dk }, + .{ .rows = 1, .n = heads, .k = dk }, + .{ .rows = 32, .n = heads, .k = dk }, + .{ .rows = 130, .n = heads, .k = dk }, + }, + }; +} + +const soa_sums_why = "the 32-lane square sum reaches the output only through its inverse square root, and the stated budget's magnitude term at the output's scale is wider than any fp16 miss the stress can engineer: gate 3's shapes bound the sum to f32 quality instead (a bf16-summed ss misses by ~2^-8 relative, far past the observed worst of bound)"; +const conv_sums_why = "these stress inputs do not separate bf16 tap sums beyond the rounded output bound; fp16 overflows"; +const recur_sums_why = "the kv dot's bf16 error enters the state only scaled by beta and the budget's magnitude term already carries the f32 chain: the fp16 margin proves fp32 sums"; + +pub fn spec(which: Which) c.Spec { + return switch (which) { + .split_soa => .{ + .name = "vulkan gdn split soa+beta (bf16)", + .accum = .f32, + .depth = splitDepth, + .ops = 5, + .mids = 3, + .out = .bf16, + .tiles = .{ token_tile, .{}, .{} }, + .served = servedOf(which), + .sums = .{ .f16 = 0, .bf16 = 0, .why = soa_sums_why }, + }, + .split_g => .{ + .name = "vulkan gdn split g (f32)", + .accum = .none, // pointwise gate chain: no sums (CUDA gdn_prep's gb precedent) + .depth = gateDepth, + .ops = 4, + .mids = 2, + .out = .f32, + .tiles = .{ token_tile, .{}, .{} }, + .served = servedOf(which), + }, + .conv_prefill => .{ + .name = "vulkan gdn conv prefill out+state (bf16)", + .accum = .f32, + .depth = tapsDepth, + .out = .bf16, + .tiles = .{ token_tile, .{}, .{} }, + .served = servedOf(which), + .sums = .{ .bf16 = 0, .why = conv_sums_why }, + }, + .conv_decode => .{ + .name = "vulkan gdn conv decode out+state (bf16)", + .accum = .f32, + .depth = tapsDepth, + .out = .bf16, + .tiles = .{ token_tile, .{}, .{} }, + .served = servedOf(which), + .sums = .{ .bf16 = 0, .why = conv_sums_why }, + }, + .prefill_y => .{ + .name = "vulkan gdn prefill y (bf16)", + .accum = .f32, + .depth = recurYDepth, + .ops = 4, + .out = .bf16, + .tiles = .{ token_tile, .{}, .{} }, + .served = servedOf(which), + .sums = .{ .bf16 = 0, .why = recur_sums_why }, + }, + .prefill_state => .{ + .name = "vulkan gdn prefill hf+snap (f32)", + .accum = .f32, + .depth = recurStateDepth, + .ops = 2, + .out = .f32, + .tiles = .{ token_tile, .{}, .{} }, + .served = servedOf(which), + .sums = .{ .bf16 = 0, .why = recur_sums_why }, + }, + .decode_y => .{ + .name = "vulkan gdn decode y (bf16)", + .accum = .f32, + .depth = decodeYDepth, + .ops = 4, + .out = .bf16, + .tiles = .{ token_tile, .{}, .{} }, + .served = servedOf(which), + .sums = .{ .bf16 = 0, .why = recur_sums_why }, + }, + .decode_state => .{ + .name = "vulkan gdn decode state (f32)", + .accum = .f32, + .depth = decodeStateDepth, + .ops = 2, + .out = .f32, + .tiles = .{ token_tile, .{}, .{} }, + .served = servedOf(which), + .sums = .{ .bf16 = 0, .why = recur_sums_why }, + }, + }; +} + +pub fn f32Of(pairs: []const u16, word: usize) f64 { + return @as(f32, @bitCast(@as(u32, pairs[2 * word]) | (@as(u32, pairs[2 * word + 1]) << 16))); +} + +pub fn bf(bits: u16) f64 { + return c.fromBf16(bits); +} + +pub fn putF32(pairs: []u16, word: usize, v: f32) void { + const bits: u32 = @bitCast(v); + pairs[2 * word] = @truncate(bits); + pairs[2 * word + 1] = @truncate(bits >> 16); +} + +// `walk` matches fx.RecurWalk/fx.Decode bitwise (null `acc`=f64); writes y/state, snapshots, pre-update `max_state`. +pub fn walk( + tokens: usize, + soa_bits: []const u16, + g_pairs: []const u16, + beta_bits: []const u16, + h0_bits: []const u16, + st: []f64, + want_y: ?[]f64, + want_state: ?[]f64, + max_state: *f64, + snap_chunk: ?usize, + want_snap: ?[]f64, + acc: ?c.Accum, + order: ?g.Order, +) void { + const section = tokens * linear_dim; + for (0..state_elems) |i| st[i] = f32Of(h0_bits, i); + var terms: [dk]f64 = undefined; + for (0..tokens) |t| { + for (0..heads) |h| { + const gv = f32Of(g_pairs, t * heads + h); + const beta = bf(beta_bits[t * heads + h]); + const qsec = 0 * section + (t * heads + h) * dk; + const ksec = 1 * section + (t * heads + h) * dk; + const vsec = 2 * section + (t * heads + h) * dk; + for (0..dk) |row| { + for (0..dk) |i| { + const s = st[h * dk * dk + row * dk + i]; + max_state.* = @max(max_state.*, @abs(s)); + terms[i] = (s * gv) * bf(soa_bits[ksec + i]); + } + const kv = if (acc) |a| rows.sumIn(a, order.?, &terms) else seq(&terms); + const delta = (bf(soa_bits[vsec + row]) - kv) * beta; + for (0..dk) |i| { + const idx = h * dk * dk + row * dk + i; + const ns = (st[idx] * gv) + delta * bf(soa_bits[ksec + i]); + st[idx] = ns; + terms[i] = ns * bf(soa_bits[qsec + i]); + } + const o = if (acc) |a| rows.sumIn(a, order.?, &terms) else seq(&terms); + if (want_y) |wy| wy[(t * heads + h) * dk + row] = o; + } + } + if (want_snap) |ws| if (snap_chunk) |chk| { + if ((t + 1) % chk == 0 and t + 1 < tokens) { + const slot = (t + 1) / chk - 1; + @memcpy(ws[slot * state_elems ..][0..state_elems], st[0..state_elems]); + } + }; + } + if (want_state) |ws| @memcpy(ws[0..state_elems], st[0..state_elems]); +} + +fn seq(terms: []const f64) f64 { + var sum: f64 = 0; + for (terms) |v| sum += v; + return sum; +} + +// Split bounds use `want/mag`: soa then beta; mag is row `vmax^2` for soa and `|want|` for beta. +pub fn splitReference( + tokens: usize, + conv_out_bits: []const u16, + a_bits: []const u16, + b_bits: []const u16, + a_log_bits: []const u16, + dt_bits: []const u16, + want: []f64, + mag: []f64, +) void { + const soa_elems = 3 * tokens * linear_dim; + for (0..tokens) |t| { + for (0..3 * heads) |head| { + const hv: usize = head % heads; + var vals: [dk]f64 = undefined; + var ss: f64 = 0; + var vmax: f64 = 0; + for (0..dk) |i| { + const cv = bf(conv_out_bits[t * channels + head * dk + i]); + const sig = 1.0 / (1.0 + @exp(-cv)); + vals[i] = cv * sig; + ss += vals[i] * vals[i]; + vmax = @max(vmax, @abs(vals[i])); + } + const is_q = head < heads; + const is_k = head >= heads and head < 2 * heads; + const inv: f64 = 1.0 / @sqrt(ss / 128.0 + 1e-6 / 128.0); + const scale_f: f64 = if (is_q) 1.0 / 128.0 else @sqrt(1.0 / 128.0); + const dst_base = (head / heads) * (tokens * linear_dim) + (t * heads + hv) * dk; + for (0..dk) |i| { + want[dst_base + i] = if (is_q or is_k) scale_f * (vals[i] * inv) else vals[i]; + mag[dst_base + i] = vmax * vmax; + } + } + for (0..heads) |hv| { + const ref = fx.gateRefOracle(a_bits[t * heads + hv], b_bits[t * heads + hv], f32Of(a_log_bits, hv), dt_bits[hv]); + want[soa_elems + t * heads + hv] = ref.beta; + mag[soa_elems + t * heads + hv] = @abs(ref.beta); + } + } +} + +// Conv uses `in_word` on state then projection rows; it writes T rows and shifted carry when `want_state` is set. +pub fn convReference( + tokens: usize, + weight_bits: []const u16, + ctx: anytype, + want: []f64, + mag: []f64, + want_state: ?[]f64, +) void { + for (0..channels) |ch| { + for (0..tokens) |t| { + var acc: f64 = 0; + for (0..taps) |k| { + acc += bf(ctx.inWord(t + k, ch)) * bf(weight_bits[ch * taps + k]); + } + want[t * channels + ch] = acc; + mag[t * channels + ch] = @abs(acc); + } + if (want_state) |ws| for (0..3) |j| { + ws[j * channels + ch] = bf(ctx.inWord(tokens + j, ch)); + }; + } +} + +// The prefill conv's assembled-input reader: rows 0..2 are the stored conv state, rows 3.. the projection output. +pub const PrefillInput = struct { + state_bits: []const u16, + qkv_bits: []const u16, + pub fn inWord(self: PrefillInput, row: usize, ch: usize) u16 { + return if (row < 3) self.state_bits[row * channels + ch] else self.qkv_bits[(row - 3) * channels + ch]; + } +}; + +// The decode conv's assembled-input reader at token `t`: rows 0..2 of that step's carried state, row 3 the token's x. +pub const DecodeInput = struct { + state_bits: []const u16, + x_bits: []const u16, + t: usize, + pub fn inWord(self: DecodeInput, row: usize, ch: usize) u16 { + return if (row < 3) self.state_bits[row * channels + ch] else self.x_bits[self.t * channels + ch]; + } +}; + +// The g/beta pair one decode token's raw-gates prologue produces (the oracle chain), as the walk reads them. +pub fn decodeGates(a_bits: []const u16, b_bits: []const u16, a_log_bits: []const u16, dt_bits: []const u16, t: usize, g_pairs: []u16, beta_bits: []u16) void { + for (0..heads) |h| { + const ref = fx.gateRefOracle(a_bits[t * heads + h], b_bits[t * heads + h], f32Of(a_log_bits, h), dt_bits[h]); + putF32(g_pairs, h, @floatCast(ref.g)); + beta_bits[h] = c.toBf16(@floatCast(ref.beta)); + } +} + +// The shared gate harness expects a rig object; as in HIP, it exposes bytes while the fixture owns device buffers. +pub const Rig = struct { + output: []const u8 = &.{}, + inputs: []const u8 = &.{}, + guards_ok: bool = false, + pub fn outputBytes(r: Rig) []const u8 { + return r.output; + } + pub fn inputBytes(r: Rig) []const u8 { + return r.inputs; + } + pub fn intact(r: Rig) bool { + return r.guards_ok; + } + pub fn overlap(_: Rig) bool { + return false; // every view here is out of place; the wrapper's refuseOverlap refuses aliasing at bind + } +}; diff --git a/zig/tests/vulkan/contract_greedy.zig b/zig/tests/vulkan/contract_greedy.zig new file mode 100644 index 00000000000..24cc5a8c97a --- /dev/null +++ b/zig/tests/vulkan/contract_greedy.zig @@ -0,0 +1,458 @@ +// Vulkan greedy contract: one composed draw checks non-finite refusal and lowest-index ties against host greedyRef. +const std = @import("std"); +const vk = @import("vulkan"); +const contract = @import("contract"); +const fx = @import("stage5_fixture.zig"); + +const c = contract.spec; +const g = contract.gates; + +// The shell reserves 256 bytes before data, matching Vulkan's maximum minStorageBufferOffsetAlignment. +pub const guard: usize = fx.at * 2; +const pattern: u8 = 0xa5; + +// One workgroup folds a row: max_chunks mirrors greedy.zig's max_chunks; fx.chunk mirrors its block_elements. +const max_chunks: usize = 256; + +/// A row the combine folds in its single workgroup cannot exceed. +const max_row: usize = max_chunks * fx.chunk; + +// Refused rows retain this reset index, so unwritten output cannot pass as token 0. +pub const refused: u32 = std.math.maxInt(u32); + +pub const rule = "any non-finite anywhere in the row refuses it, winner or not; the largest value wins; " ++ + "exact ties (including +0/-0) go to the lowest index"; + +fn flat(_: c.Shape) usize { + return 0; +} + +// Served vocab and draft-window rows cover one row at 2 elements and chunk edges 4094, 4096, 4098. +pub fn spec() c.Spec { + return .{ + .name = "vulkan greedy draw: finite_logits, argreduce_chunk, argreduce_combine (bf16)", + .accum = .none, + .depth = flat, + .out = .raw, + .nan = .{ .stated = rule }, + .tiles = .{ .{}, .{ .size = fx.chunk, .step = 2, .min = 2 }, .{} }, + .served = &.{ + .{ .rows = 1, .n = fx.vocab, .k = 0 }, + .{ .rows = 4, .n = fx.vocab, .k = 0 }, + .{ .rows = 1, .n = max_row, .k = 0 }, + }, + }; +} + +/// The control's registration: it claims the same rule and must fail it. +pub fn controlSpec() c.Spec { + return .{ + .name = "control: fmax row max (drops a NaN)", + .accum = .none, + .depth = flat, + .out = .raw, + .nan = .{ .stated = rule }, + .tiles = spec().tiles, + .served = spec().served, + }; +} + +// Input data sits between downloaded guard spans for repeat checks of unchanged inputs and intact guards. +pub const Rig = struct { + a: std.mem.Allocator, + dev: *const vk.Device, + cmds: *vk.Commands, + staging: vk.Buffer, + timeout: u64, + // prepare owns the guarded shell; views copy its handle but do not deinitialize it. + shell: ?vk.Buffer = null, + data_bytes: usize = 0, + seen: std.ArrayList(u8) = .empty, + whole: std.ArrayList(u8) = .empty, + outs: std.ArrayList(u8) = .empty, + + pub fn deinit(r: *Rig) void { + r.reset(); + r.seen.deinit(r.a); + r.whole.deinit(r.a); + r.outs.deinit(r.a); + } + + pub fn reset(r: *Rig) void { + // No freed buffer may stay referenced by an open command buffer's records. + r.cmds.discard(); + if (r.shell) |b| { + var s = b; + s.deinit(); + r.shell = null; + } + r.outs.clearRetainingCapacity(); + } + + // Upload the pattern with data; fresh device memory is zero, not the pattern required by intact(). + pub fn shellFor(r: *Rig, data: []const u8) ![]u8 { + const all = try r.a.alloc(u8, data.len + 2 * guard); + errdefer r.a.free(all); + @memset(all, pattern); + @memcpy(all[guard..][0..data.len], data); + return all; + } + + /// Uploads `data` between two guard spans; the shell's data window starts at `guard`. + pub fn input(r: *Rig, data: []const u8) !vk.Buffer { + const all = try r.shellFor(data); + defer r.a.free(all); + var b = try vk.Buffer.init(r.dev, .device, all.len); + errdefer b.deinit(); + try vk.orExit(vk.upload(r.cmds, b, 0, r.staging, all, r.timeout)); + r.shell = b; + r.data_bytes = data.len; + return b; + } + + /// Every input's data window now, on the host: a run must leave it as prepared. + pub fn inputBytes(r: *Rig) ![]const u8 { + const b = r.shell orelse return error.Invalid; + try r.seen.resize(r.a, r.data_bytes); + try vk.orExit(vk.download(r.cmds, b, guard, r.staging, r.seen.items, r.timeout)); + return r.seen.items; + } + + /// Every guard of the shell still holds its pattern. + pub fn intact(r: *Rig) !bool { + const b = r.shell orelse return error.Invalid; + try r.whole.resize(r.a, r.data_bytes + 2 * guard); + try vk.orExit(vk.download(r.cmds, b, 0, r.staging, r.whole.items, r.timeout)); + for (r.whole.items, 0..) |v, i| { + if (i >= guard and i < guard + r.data_bytes) continue; + if (v != pattern) return false; + } + return true; + } + + /// The bytes the last run wrote, on the host. + pub fn outputBytes(r: *Rig) ![]const u8 { + return r.outs.items; + } + + // Draw outputs use separate internal result and partial buffers; the control owns its maxima allocation. + pub fn overlap(_: *Rig) bool { + return false; + } +}; + +// Same shapes, fill, and check for draw and control; only run differs, so the control fails the named gate. +pub const Greedy = struct { + a: std.mem.Allocator, + arena: std.heap.ArenaAllocator, + rig: Rig, + dev: *const vk.Device, + cmds: *vk.Commands, + staging: vk.Buffer, + timeout: u64, + shaders: vk.greedy.Shaders, + /// Present turns the registration into the control. + control_spirv: ?[]align(4) const u8, + greedy: ?vk.Greedy = null, + pipe: ?vk.dispatch.Pipeline = null, + maxima: ?vk.Buffer = null, + shape: c.Shape = undefined, + shell: vk.Buffer = undefined, + views: []vk.tensor.View = &.{}, + bits: []u16 = &.{}, + msg: [192]u8 = undefined, + + pub fn init( + a: std.mem.Allocator, + dev: *const vk.Device, + cmds: *vk.Commands, + staging: vk.Buffer, + timeout: u64, + shaders: vk.greedy.Shaders, + control_spirv: ?[]align(4) const u8, + ) Greedy { + return .{ + .a = a, + .arena = .init(a), + .rig = .{ .a = a, .dev = dev, .cmds = cmds, .staging = staging, .timeout = timeout }, + .dev = dev, + .cmds = cmds, + .staging = staging, + .timeout = timeout, + .shaders = shaders, + .control_spirv = control_spirv, + }; + } + + pub fn deinit(self: *Greedy) void { + self.dropDevices(); + self.rig.deinit(); + self.arena.deinit(); + } + + fn dropDevices(self: *Greedy) void { + // No freed pipeline or buffer may stay referenced by an open command buffer's records. + self.cmds.discard(); + if (self.greedy) |*x| x.deinit(); + self.greedy = null; + if (self.pipe) |*p| p.deinit(); + self.pipe = null; + if (self.maxima) |*m| m.deinit(); + self.maxima = null; + } + + const Arm = enum { draw, control }; + + fn arm(self: *const Greedy) Arm { + return if (self.control_spirv == null) .draw else .control; + } + + pub fn kernel(self: *Greedy) g.Kernel { + return .{ + .spec = if (self.arm() == .draw) spec() else controlSpec(), + .ctx = self, + .vt = &vt, + .control = if (self.arm() == .draw) null else .nonfinite, + }; + } + + const vt = g.hooks(Greedy); + + pub fn prepare(self: *Greedy, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + const n = shape.n; + const rows = shape.rows; + if (rows == 0 or n == 0) return false; + if (n % 2 != 0) return false; // an odd bf16 count has no legal view: the word-padding rule refuses + if (rows > 64 or n > max_row) return false; + self.dropDevices(); + self.rig.reset(); + _ = self.arena.reset(.retain_capacity); + self.shape = shape; + const a = self.arena.allocator(); + self.bits = try a.alloc(u16, rows * n); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (0..rows) |r| fillRow(self.bits[r * n ..][0..n], rng); + if (fill == .nonfinite) plantNonFinite(self.bits, rows, n); + const data = try a.alloc(u8, rows * n * 2); + for (self.bits, 0..) |b, i| std.mem.writeInt(u16, data[2 * i ..][0..2], b, .little); + self.shell = try self.rig.input(data); + self.views = try a.alloc(vk.tensor.View, rows); + for (self.views, 0..) |*v, r| + v.* = try vk.tensor.View.init(&self.shell, .bf16, guard + r * n * 2, n); + switch (self.arm()) { + .draw => self.greedy = try vk.Greedy.init(self.dev, self.shaders, n), + .control => { + self.pipe = try vk.dispatch.Pipeline.init(self.dev, self.control_spirv.?, .{ + .bindings = 2, + .push_bytes = 4, // the control's push block: one uint, the row length + .writes = 1 << 1, + }); + self.maxima = try vk.Buffer.init(self.dev, .device, rows * 4); + }, + } + try self.rig.outs.resize(self.a, rows * 4); + @memset(self.rig.outs.items, 0xff); // an unwritten row reads as the refused marker + return true; + } + + pub fn run(self: *Greedy) !void { + const n = self.shape.n; + const rows = self.shape.rows; + switch (self.arm()) { + .draw => { + const greedy = &self.greedy.?; + for (0..rows) |r| { + // Timeout or InFlight leaves completion unknown: exit here, never unwind into the rig's frees. + const tok: u32 = vk.orExit(greedy.select(self.cmds, self.staging, self.views[r], self.timeout)) catch |e| switch (e) { + // The draw rule's refusal; check names the row and its evidence. + error.NonFinite => refused, + else => return e, + }; + std.mem.writeInt(u32, self.rig.outs.items[4 * r ..][0..4], tok, .little); + } + }, + .control => { + const pipe = &self.pipe.?; + const maxima = &self.maxima.?; + try pipe.bind(&.{ + .{ .buffer = &self.shell, .offset = guard, .len = rows * n * 2 }, + .{ .buffer = maxima, .len = rows * 4 }, + }); + var push: [4]u8 align(4) = undefined; + std.mem.writeInt(u32, &push, @intCast(n), .little); + const cmd = try self.cmds.begin(); + try pipe.record(cmd, &push, .{ @intCast(rows), 1, 1 }); + // Timeout or InFlight leaves completion unknown: exit here, never unwind into the rig's frees. + try vk.orExit(self.cmds.submit(self.timeout)); + try vk.orExit(vk.download(self.cmds, maxima.*, 0, self.staging, self.rig.outs.items, self.timeout)); + }, + } + } + + pub fn count(self: *Greedy) usize { + return self.shape.rows * 4; + } + + /// Raw outputs check themselves; the float64 slots stay empty. + pub fn reference(_: *Greedy, want: []f64, mag: []f64) !void { + @memset(want, 0); + @memset(mag, 0); + } + + // A planted row refuses at its first offender; a finite row matches greedyRef's token. + pub fn check(self: *Greedy, at: g.Gate) !?[]const u8 { + const n = self.shape.n; + const rows = self.shape.rows; + switch (at) { + .nonfinite => for (0..rows) |r| { + const row = self.bits[r * n ..][0..n]; + if (!rowHasNonFinite(row)) + return try std.fmt.bufPrint(&self.msg, "row {d}: the planted fill holds no non-finite, so the gate proves nothing", .{r}); + const bad = fx.firstBadBits(row); + const got = std.mem.readInt(u32, self.rig.outs.items[4 * r ..][0..4], .little); + if (got != refused) + return try std.fmt.bufPrint(&self.msg, "row {d} did not refuse: it reported {d} where the first non-finite sits at {d}", .{ r, got, bad }); + }, + .shapes => { + const wide = try self.arena.allocator().alloc(f32, n); + for (0..rows) |r| { + const row = self.bits[r * n ..][0..n]; + for (row, 0..) |b, i| wide[i] = c.fromBf16(b); + const want = vk.greedy.greedyRef(wide) catch |e| + return try std.fmt.bufPrint(&self.msg, "row {d}: the reference refused a finite row: {s}", .{ r, @errorName(e) }); + const got = std.mem.readInt(u32, self.rig.outs.items[4 * r ..][0..4], .little); + if (got == refused) + return try std.fmt.bufPrint(&self.msg, "row {d} refused a finite row", .{r}); + if (got != want) + return try std.fmt.bufPrint(&self.msg, "row {d} picks {d}, greedyRef gives {d}", .{ r, got, want }); + } + }, + else => {}, + } + return null; + } +}; + +// Eighth-step logits provide 256 values and common ties, testing lowest-index selection across chunk folds. +fn fillRow(data: []u16, rng: std.Random) void { + for (data) |*d| d.* = c.toBf16(@as(f32, @floatFromInt(rng.uintLessThan(u32, 256))) / 8 - 16); +} + +// Every row: a finite would-be winner, then a losing NaN, +inf and -inf; every third row loses its winner to a NaN too. +fn plantNonFinite(bits: []u16, rows: usize, n: usize) void { + for (0..rows) |r| { + const row = bits[r * n ..][0..n]; + row[n / 2] = c.toBf16(5.0); + row[0] = 0xff80; // -inf + row[n / 3] = 0x7fc0; // NaN + row[n - 1] = 0x7f80; // +inf + if (r % 3 == 0) row[n / 2] = 0x7fc0; + } +} + +/// The finite scan's own predicate: exponent all ones is non-finite bf16. +pub fn rowHasNonFinite(row: []const u16) bool { + for (row) |b| if ((b & 0x7fff) >= 0x7f80) return true; + return false; +} + +const testing = std.testing; + +test "the spec serves the vocab, a draft window and the max_chunks edge; the chunk edges ride one row" { + const s = spec(); + try testing.expect(s.nan.stated.len > 0); + try testing.expectEqual(fx.vocab, s.served[0].n); + try testing.expectEqual(@as(usize, 4), s.served[1].rows); + try testing.expectEqual(max_row, s.served[2].n); + try testing.expectEqual(@as(usize, 1_048_576), max_row); // 256 chunks of 4096 + const list = try c.shapes(testing.allocator, s); + defer testing.allocator.free(list); + try testing.expectEqual(@as(usize, 7), list.len); + try testing.expectEqual(fx.vocab, list[0].n); + try testing.expectEqual(@as(usize, 4), list[1].rows); + try testing.expectEqual(max_row, list[2].n); + for (list[3..], [_]usize{ 2, fx.chunk - 2, fx.chunk, fx.chunk + 2 }) |x, want| { + try testing.expectEqual(@as(usize, 1), x.rows); + try testing.expectEqual(want, x.n); + } +} + +test "the eighths fill is bf16-exact and tie-rich, and greedyRef and the fixture reference agree" { + const a = testing.allocator; + const n = 512; // more values than the fill's 256 classes, so ties must exist + const bits = try a.alloc(u16, n); + defer a.free(bits); + var prng = std.Random.DefaultPrng.init(0x6b63); + fillRow(bits, prng.random()); + var distinct: usize = 0; + for (bits, 0..) |b, at| { + const v = c.fromBf16(b); + const k = (v + 16) * 8; // the fill's eighth class: an exact small integer + try testing.expect(k == @trunc(k) and k >= 0 and k < 256); + if (std.mem.indexOfScalar(u16, bits[0..at], b) == null) distinct += 1; + } + try testing.expect(distinct < n); + for (0..256) |k| // every drawn class round-trips bf16 exactly + try testing.expectEqual(@as(f32, @as(f32, @floatFromInt(k)) / 8 - 16), c.fromBf16(c.toBf16(@as(f32, @floatFromInt(k)) / 8 - 16))); + const wide = try a.alloc(f32, n); + defer a.free(wide); + for (bits, 0..) |b, i| wide[i] = c.fromBf16(b); + const want = try vk.greedy.greedyRef(wide); + try testing.expectEqual(try fx.referenceBits(bits), want); + var best = -std.math.inf(f32); + for (wide) |v| best = @max(best, v); + try testing.expectEqual(@as(u32, @intCast(std.mem.indexOfScalar(f32, wide, best).?)), want); // the tie's lowest index +} + +test "the nonfinite plant covers every row, keeps the loser rows' winner, and names the first offender" { + const a = testing.allocator; + const rows = 3; + const n = 4096; + const bits = try a.alloc(u16, rows * n); + defer a.free(bits); + plantNonFinite(bits, rows, n); + for (0..rows) |r| { + const row = bits[r * n ..][0..n]; + try testing.expect(rowHasNonFinite(row)); + try testing.expectEqual(@as(usize, 0), fx.firstBadBits(row)); // the losing -inf at index 0 + try testing.expectEqual(@as(u16, 0xff80), row[0]); + try testing.expectEqual(@as(u16, 0x7fc0), row[n / 3]); + try testing.expectEqual(@as(u16, 0x7f80), row[n - 1]); + if (r % 3 != 0) try testing.expectEqual(@as(f32, 5.0), c.fromBf16(row[n / 2])); + } +} + +test "the refused marker is the word the draw's own decode refuses to pass" { + try testing.expectEqual(std.math.maxInt(u32), refused); + try testing.expectError(error.VulkanFailed, vk.greedy.decode(4, .{ .finite_flag = 0, .first_bad = refused, .index = refused })); + try testing.expectError(error.NonFinite, vk.greedy.decode(4, .{ .finite_flag = 1, .first_bad = 1, .index = 0 })); +} + +test "the draw arm's fence guard passes a refusal through: only Timeout and InFlight exit" { + const nonFinite = struct { + fn f() anyerror!u32 { + return error.NonFinite; + } + }.f; + // The guard wraps select; a rule refusal must still reach the draw's refused marker, not exit. + const tok = vk.orExit(nonFinite()) catch |e| switch (e) { + error.NonFinite => refused, + else => return e, + }; + try testing.expectEqual(refused, tok); +} + +test "the input shell uploads the pattern intact() checks; fresh device memory reads zeros" { + var rig = Rig{ .a = testing.allocator, .dev = undefined, .cmds = undefined, .staging = undefined, .timeout = 0 }; + // shellFor is host-only and the rig owns nothing here: no deinit, since reset discards through cmds. + const data = [_]u8{ 1, 2, 3, 4 }; + const all = try rig.shellFor(&data); + defer testing.allocator.free(all); + try testing.expectEqual(data.len + 2 * guard, all.len); + for (all[0..guard]) |v| try testing.expectEqual(pattern, v); + for (all[guard + data.len ..]) |v| try testing.expectEqual(pattern, v); + try testing.expectEqualSlices(u8, &data, all[guard..][0..data.len]); + // intact() skips data between patterned spans; without uploading the pattern, the repeat gate would fail. + try testing.expect(all[0] == pattern and all[guard - 1] == pattern); +} diff --git a/zig/tests/vulkan/contract_greedy_main.zig b/zig/tests/vulkan/contract_greedy_main.zig new file mode 100644 index 00000000000..33851eba19c --- /dev/null +++ b/zig/tests/vulkan/contract_greedy_main.zig @@ -0,0 +1,86 @@ +// Greedy: gates 1-4 on the composed draw; .nonfinite control (HIP tf_ctl_fmax_max mirror) must go red; exit 0/1/2/124. +const std = @import("std"); +const vk = @import("vulkan"); +const contract = @import("contract"); +const g = contract.gates; +const shaders = @import("stage5_shaders"); +const controls = @import("greedy_controls"); +const greedy_k = @import("contract_greedy.zig"); + +const usage = "usage: tf-vulkan-greedy-contract [--vendor ID] [--timeout-ms MS] | --help\n" ++ + " --help show this text and exit without opening a loader or device\n" ++ + " --vendor ID refuse a device whose Vulkan vendorID differs (decimal or 0x hex; no default)\n" ++ + " --timeout-ms MS bound on every fence wait, 1 to 600000 (default 10000); a timeout ends the process\n"; + +pub fn main(init: std.process.Init) !u8 { + const a = init.gpa; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + for (args[1..]) |arg| if (std.mem.eql(u8, arg, "--help")) { + std.debug.print(usage, .{}); + return 0; + }; + var vendor: ?u32 = null; + var timeout_ms: u64 = 10_000; + var i: usize = 1; + while (i < args.len) : (i += 2) { + const value = if (i + 1 < args.len) args[i + 1] else ""; + if (std.mem.eql(u8, args[i], "--vendor")) { + vendor = std.fmt.parseInt(u32, value, 0) catch null; + if (vendor == null) break; + } else if (std.mem.eql(u8, args[i], "--timeout-ms")) { + timeout_ms = std.fmt.parseInt(u64, value, 10) catch 0; + if (timeout_ms == 0 or timeout_ms > 600_000) break; + } else break; + } else return run(a, vendor, timeout_ms * std.time.ns_per_ms); + std.debug.print(usage, .{}); + return 2; +} + +fn shaderSet() vk.greedy.Shaders { + return .{ + .finite_logits = &shaders.finite_logits, + .argreduce_chunk = &shaders.argreduce_chunk, + .argreduce_combine = &shaders.argreduce_combine, + }; +} + +fn run(a: std.mem.Allocator, vendor: ?u32, timeout: u64) !u8 { + var rt = try vk.Runtime.open(); + defer rt.close(); + var dev = try vk.Device.open(&rt, .{ .features = .{ .storage16 = true, .int16 = true } }); + defer dev.close(); + const pr = dev.properties; + const name = std.mem.sliceTo(&pr.deviceName, 0); + const v = pr.apiVersion; + std.debug.print("MILESTONE greedy-contract device-found name=\"{s}\" vendor=0x{x:0>4} device=0x{x:0>4} " ++ + "api={d}.{d}.{d} storage16={} int16={} fence-timeout={d}ms\n", .{ + name, pr.vendorID, + pr.deviceID, v >> 22 & 0x7f, + v >> 12 & 0x3ff, v & 0xfff, + dev.features.storage16, dev.features.int16, + timeout / std.time.ns_per_ms, + }); + if (pr.deviceType == vk.c.VK_PHYSICAL_DEVICE_TYPE_CPU) { + std.debug.print("FAIL greedy-contract: a CPU (software) Vulkan device is never a GPU result\n", .{}); + return 1; + } + if (vendor) |want| if (pr.vendorID != want) { + std.debug.print("FAIL greedy-contract: vendor 0x{x:0>4} is not the requested 0x{x:0>4}\n", .{ pr.vendorID, want }); + return 1; + }; + var cmds = try vk.Commands.init(&dev); + defer cmds.deinit(); + // The largest contract shell (max_chunks of 4096 bf16) rides staging: 2 MiB. + var staging = try vk.Buffer.init(&dev, .staging, 4 << 20); + defer staging.deinit(); + var draw = greedy_k.Greedy.init(a, &dev, &cmds, staging, timeout, shaderSet(), null); + defer draw.deinit(); + var ctl = greedy_k.Greedy.init(a, &dev, &cmds, staging, timeout, shaderSet(), &controls.fmax_max); + defer ctl.deinit(); + var kernels = [_]g.Kernel{ draw.kernel(), ctl.kernel() }; + var t: g.Tally = .{}; + try g.runAll(a, &kernels, 0x6b63, &t); + std.debug.print("SUMMARY vulkan greedy contract on \"{s}\": {d} kernel, {d} pass gates 1-4, {d} fail; " ++ + "{d} control: {d} red as planted, {d} missed\n", .{ name, t.kernels, t.passed, t.failed, t.controls, t.caught, t.missed }); + return if (t.problems() > 0) 1 else 0; +} diff --git a/zig/tests/vulkan/contract_main.zig b/zig/tests/vulkan/contract_main.zig index 5ea34f7c3fb..8b6b2c05a3a 100644 --- a/zig/tests/vulkan/contract_main.zig +++ b/zig/tests/vulkan/contract_main.zig @@ -47,6 +47,27 @@ pub fn main(init: std.process.Init) !u8 { if (!std.mem.eql(u32, &pipeline.local, &.{ fx.local_size_x, 1, 1 })) return error.LocalSizeMismatch; var fixture = fx.Fixture.init(init.gpa, &dev, &cmds, staging, &pipeline); defer fixture.deinit(); + var vec4_pipeline = try vk.dispatch.Pipeline.init(&dev, &shaders.qmm_q4_vec4, layout); + defer vec4_pipeline.deinit(); + if (!std.mem.eql(u32, &vec4_pipeline.local, &.{ fx.local_size_x, 1, 1 })) return error.LocalSizeMismatch; + var vec4_fixture = fx.Fixture.init(init.gpa, &dev, &cmds, staging, &vec4_pipeline); + defer vec4_fixture.deinit(); + vec4_fixture.name = "vulkan qmm_q4_vec4 affine q4 g64 transposed"; + var tile_pipeline = try vk.dispatch.Pipeline.init(&dev, &shaders.qmm_q4_tile, layout); + defer tile_pipeline.deinit(); + if (!std.mem.eql(u32, &tile_pipeline.local, &.{ fx.local_size_x, 1, 1 })) return error.LocalSizeMismatch; + var tile_fixture = fx.Fixture.init(init.gpa, &dev, &cmds, staging, &tile_pipeline); + defer tile_fixture.deinit(); + tile_fixture.name = "vulkan qmm_q4_tile affine q4 g64 transposed tiled 32"; + tile_fixture.tiled = true; + var tile_f32_pipeline = try vk.dispatch.Pipeline.init(&dev, &shaders.qmm_q4_tile_f32, layout); + defer tile_f32_pipeline.deinit(); + if (!std.mem.eql(u32, &tile_f32_pipeline.local, &.{ fx.local_size_x, 1, 1 })) return error.LocalSizeMismatch; + var tile_f32_fixture = fx.Fixture.init(init.gpa, &dev, &cmds, staging, &tile_f32_pipeline); + defer tile_f32_fixture.deinit(); + tile_f32_fixture.name = "vulkan qmm_q4_tile_f32 affine q4 g64 transposed tiled 32 f32 activations"; + tile_f32_fixture.tiled = true; + tile_f32_fixture.act_f32 = true; const control_layout: vk.dispatch.Layout = .{ .bindings = 3, .push_bytes = 8, .writes = 1 << 2 }; var p_fmax = try vk.dispatch.Pipeline.init(&dev, &shaders.qmm_fmax, control_layout); defer p_fmax.deinit(); @@ -77,11 +98,14 @@ pub fn main(init: std.process.Init) !u8 { fx.Control.init(init.gpa, &dev, &cmds, staging, &p_fp16, .fp16_accum), }; defer for (&controls) |*control| control.deinit(); - var kernels: [8]g.Kernel = undefined; + var kernels: [11]g.Kernel = undefined; kernels[0] = fixture.kernel(); - for (&controls, 1..) |*control, i| kernels[i] = control.kernel(); - kernels[6] = writes_w.kernel(); - kernels[7] = writes_b.kernel(); + kernels[1] = vec4_fixture.kernel(); + kernels[2] = tile_fixture.kernel(); + for (&controls, 3..) |*control, i| kernels[i] = control.kernel(); + kernels[8] = writes_w.kernel(); + kernels[9] = writes_b.kernel(); + kernels[10] = tile_f32_fixture.kernel(); const start = std.Io.Clock.awake.now(init.io); tally.controls += gate5_controls; tally.caught += gate5_controls; diff --git a/zig/tests/vulkan/contract_pointwise.zig b/zig/tests/vulkan/contract_pointwise.zig new file mode 100644 index 00000000000..03b705cffc9 --- /dev/null +++ b/zig/tests/vulkan/contract_pointwise.zig @@ -0,0 +1,584 @@ +// Pointwise contract: stages 2-4 closures as host fixtures; preserve add checked STRICTLY (a flushing device fails). +const std = @import("std"); +const contract = @import("contract"); +const stage2 = @import("stage2_fixture.zig"); +const shaders2 = @import("stage2_shaders"); +const vk = @import("vulkan"); + +const c = contract.spec; +pub const Report = stage2.Report; +pub const classify = stage2.classify; + +/// Element offsets (elements of the binding's kind; 256-byte alignment) of each binding's data. +pub const Offsets = stage2.Offsets; + +pub const tail = stage2.tail; +pub const poison16 = stage2.poison16; + +fn poisonedShell16(a: std.mem.Allocator, at: usize, len: usize) ![]u16 { + const buf = try a.alloc(u16, at + len + tail); + @memset(buf, poison16); + return buf; +} + +fn guardIntact(shell: []const u16, at: usize, len: usize, poison: u16) !void { + for (shell, 0..) |g, i| { + if (i >= at and i < at + len) continue; + if (g != poison) return error.GuardClobbered; + } +} + +/// A bf16-representable f32 value near `1` scaled by `scale` (a power of two). +fn sampleBits(rng: std.Random, scale: f32) u16 { + const v = (rng.float(f32) + 0.5) * scale * (if (rng.boolean()) @as(f32, 1) else -1); + return c.toBf16(v); +} + +// ---- Residual add, preserved: add_preserve_bf16 (binary_vec.comp + -DPRESERVE_F32), checked STRICTLY ---- + +// Preserved add: stage 2's three shapes, 256B-aligned; norm_mlp selects it where the device declares DenormPreserve. +pub const PreserveCase = stage2.AddCase; +pub const preserve_cases = [_]PreserveCase{ + .{ .name = "random 512", .elems = 512, .at = .{ .a = 128, .b = 256, .out = 384 } }, + .{ .name = "cancel 512", .elems = 512, .at = .{ .a = 128, .b = 256, .out = 384 }, .fill = .cancel }, + .{ .name = "classes 20", .elems = 20, .at = .{ .a = 128, .b = 256, .out = 384 }, .fill = .classes }, +}; + +// Preserve gate reuses stage 2's add fixture; allow_subnormal_flush false: a flushing device fails the subnormal rows. +pub const PreserveAdd = stage2.Add; + +// ---- Attention output gate: one mid bf16 rounding of the sigmoid, one bf16 store ---- + +// ops = 4 (exp, 1+, divide, ctx multiply); mids = 1: sigmoid bf16 rounding, error <= 2^-9|s|; a budget, not a proof. +const gate_spec: c.Spec = .{ .name = "vulkan attention_gate bf16 (bf16(sigmoid(gate))*ctx)", .accum = .none, .depth = gateDepth, .ops = 4, .mids = 1, .out = .bf16, .served = &[_]c.Shape{} }; +fn gateDepth(_: c.Shape) usize { + return 0; // elementwise +} + +pub const gate_row_elems: usize = 2048; // 8 heads * 256, the gate's output row +pub const gate_qraw_row_elems: usize = 2 * gate_row_elems; // the interleaved q/gate projection row + +pub const GateCase = struct { name: []const u8, rows: usize, at: Offsets = .{}, controls: bool = false }; + +pub const gate_cases = [_]GateCase{ + .{ .name = "random 2x2048", .rows = 2, .at = .{ .a = 128, .b = 256, .out = 384 } }, + .{ .name = "controls 2x2048", .rows = 2, .at = .{ .a = 128, .b = 256, .out = 384 }, .controls = true }, +}; + +// Gate controls: planted nan/±inf/±0, model-range ends, a bf16 subnormal; |gate| <= 100 keeps exp(-gate) in f32 range. +const gate_controls = [_]u16{ 0x7fc0, 0x7f80, 0xff80, 0x0000, 0x8000, 0x3f80, 0xbf80, 0x41a0, 0xc1a0, 0x42c8, 0xc2c8, 0x3e80, 0xbd80, 0x0001, 0x3f00, 0xbf00 }; + +/// |gate| <= 16 keeps exp(-gate) at most e^16 ~ 8.9e6, safely inside f32. +fn gateBits(rng: std.Random) u16 { + return c.toBf16((rng.float(f32) * 2 - 1) * 16); +} + +pub const Gate = struct { + a: std.mem.Allocator, + case: GateCase, + /// ctx bits: binding 0, rows * 2048 elements. + ctx: []u16, + // qraw bits: binding 1; element d of row r at r*4096 + (d/256)*512 + 256 + d%256 (interleaved layout). + qraw: []u16, + out: []u16, + /// f64(ctx) * sigmoid64(gate) per element; nonfinite only where the gate is planted so. + want: []f64, + + pub fn init(a: std.mem.Allocator, case: GateCase, seed: u64) !Gate { + const elems = case.rows * gate_row_elems; + const ctx = try poisonedShell16(a, case.at.a, elems); + errdefer a.free(ctx); + const qraw = try poisonedShell16(a, case.at.b, case.rows * gate_qraw_row_elems); + errdefer a.free(qraw); + const out = try poisonedShell16(a, case.at.out, elems); + errdefer a.free(out); + const want = try a.alloc(f64, elems); + errdefer a.free(want); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (0..elems) |i| { + const row = i / gate_row_elems; + const d = i % gate_row_elems; + const gate: u16 = if (case.controls) gate_controls[i % gate_controls.len] else gateBits(rng); + // ctx is 1, except -1 under the -inf gate (signed zero) and zero at 511: -0, 1023: +0; zeros keep the sign. + const ctx_bits: u16 = if (!case.controls) sampleBits(rng, 1) else if (i == 2) 0xbf80 else if (i == 511) 0x8000 else if (i == 1023) 0x0000 else 0x3f80; + ctx[case.at.a + i] = ctx_bits; + // The shader's read: head = d/256, gate element = row*4096 + head*512 + 256 + d%256. + qraw[case.at.b + row * gate_qraw_row_elems + (d / 256) * 512 + 256 + d % 256] = gate; + const g64: f64 = c.fromBf16(gate); + const e = @exp(-g64); + const s = 1 / (1 + e); + want[i] = c.fromBf16(ctx_bits) * s; + } + return .{ .a = a, .case = case, .ctx = ctx, .qraw = qraw, .out = out, .want = want }; + } + + pub fn deinit(f: Gate) void { + f.a.free(f.ctx); + f.a.free(f.qraw); + f.a.free(f.out); + f.a.free(f.want); + } + + // Finite wants within boundExact; a zero want (ctx zero, gate -inf) demands an exactly-zero signed output. + pub fn check(f: Gate, got: []const u16) !Report { + if (got.len != f.out.len) return error.LengthMismatch; + const at = f.case.at.out; + try guardIntact(got, at, f.case.rows * gate_row_elems, poison16); + var r: Report = .{}; + for (f.want, 0..) |w, i| { + const g = got[at + i]; + const gf = c.fromBf16(g); + if (!std.math.isFinite(w) or !std.math.isFinite(gf)) { + if (c.classOf(gf) != c.classOf(w)) return error.NonFiniteMismatch; + continue; + } + if (w == 0) { + if (gf != 0 or std.math.signbit(gf) != std.math.signbit(w)) return error.OutOfBound; + continue; + } + const err = @abs(gf - w); + const bound = c.boundExact(gate_spec, w); + if (!(std.math.isFinite(bound) and bound > 0)) return error.BadBound; + const ratio = err / bound; + if (ratio > r.worst) r = .{ .worst = ratio, .row = i / gate_row_elems, .col = i % gate_row_elems }; + } + if (r.worst > 1) return error.OutOfBound; + return r; + } +}; + +// ---- Gated RMSNorm: tokens x 16 head rows of 128, the [128] gain shared by every head, plus one mid rounding ---- + +// depth = 12 (1 + 8 levels + 1 divide = 10, 2 margin); ops = 10; mids = 12 (2^-8 of |normed|; |z| >= 0.25 floor). +const gated_depth = 12; +const gated_spec: c.Spec = .{ + .name = "vulkan fast_norm_gated bf16 per-head gated rms (silu(z)*bf16(x*(inv*w)), head rows of 128)", + .accum = .f32, + .depth = gatedDepth, + .ops = 10, + .mids = 12, + .out = .bf16, + .served = &[_]c.Shape{ .{ .rows = 1, .n = 1, .k = gated_head_dim }, .{ .rows = 2, .n = 1, .k = gated_head_dim } }, +}; +fn gatedDepth(_: c.Shape) usize { + return gated_depth; +} + +// gdn.Model geometry: 16 head rows of 128 a token (2048 bf16 elements); one [128] gain shared by every head. +pub const gated_heads: usize = 16; +pub const gated_head_dim: usize = 128; +pub const gated_width: usize = gated_heads * gated_head_dim; +pub const gated_eps: f32 = 1e-6; + +pub const GatedOffsets = struct { x: usize = 0, weight: usize = 0, gate: usize = 0, out: usize = 0 }; + +pub const GatedCase = struct { + name: []const u8, + tokens: usize, + fill: stage2.RmsFill = .flat, + // The last token's 16 head rows are all zeros: mean square 0, norm = rsqrt(eps), outputs exactly signed zero. + zero_token: bool = false, + // Token 0 plants z controls; tokens 1 and 2 poison x[head 0][0] with +inf and nan; the per-head norm must not leak. + controls: bool = false, + at: GatedOffsets = .{}, +}; + +// Served tokens: one decode row and a prefill chunk (rows = 1, max_prefill_rows = 32), plus controls on three tokens. +pub const gated_cases = [_]GatedCase{ + .{ .name = "1x16x128", .tokens = 1, .at = .{ .x = 128, .weight = 0, .gate = 256, .out = 384 } }, + .{ .name = "32x16x128 dynamic zero token", .tokens = 32, .fill = .dynamic, .zero_token = true, .at = .{ .x = 128, .weight = 0, .gate = 256, .out = 384 } }, + .{ .name = "controls 3x16x128", .tokens = 3, .controls = true, .at = .{ .x = 128, .weight = 0, .gate = 256, .out = 384 } }, +}; + +// z controls at token 0, head 0: |silu| >= 0.109, zeros from z = ±0, classes from nan/±inf; z = -inf: silu = nan. +const gated_z_controls = [_]u16{ 0x7fc0, 0x7f80, 0xff80, 0x0000, 0x8000, 0x3f80, 0xbf80, 0x3e80, 0xbe80, 0x4000, 0x4100, 0xc000 }; + +/// |z| in [0.25, 8] on one side of zero (resampled): the mids budget's stated gate floor. +fn gatedGateBits(rng: std.Random) u16 { + var v = rng.float(f32) * 10 - 2; + while (@abs(v) < 0.25) v = rng.float(f32) * 10 - 2; + return c.toBf16(v); +} + +fn sigmoid64(g: f64) f64 { + return 1 / (1 + @exp(-g)); +} + +fn silu64(z: f64) f64 { + return z * sigmoid64(z); +} + +pub const GatedNorm = struct { + a: std.mem.Allocator, + case: GatedCase, + x: []u16, + weight: []u16, + z: []u16, + out: []u16, + // silu64(z) * widen(bf16(x*(n*w))), f64, mid modeled; w the shared [128] gain; nonfinite only at planted controls. + want: []f64, + + pub fn init(a: std.mem.Allocator, case: GatedCase, seed: u64) !GatedNorm { + const elems = case.tokens * gated_width; + const x = try poisonedShell16(a, case.at.x, elems); + errdefer a.free(x); + const weight = try poisonedShell16(a, case.at.weight, gated_head_dim); + errdefer a.free(weight); + const z = try poisonedShell16(a, case.at.gate, elems); + errdefer a.free(z); + const out = try poisonedShell16(a, case.at.out, elems); + errdefer a.free(out); + const want = try a.alloc(f64, elems); + errdefer a.free(want); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (0..gated_head_dim) |i| weight[case.at.weight + i] = sampleBits(rng, 1); + for (0..case.tokens) |token| { + const base = case.at.x + token * gated_width; + const zbase = case.at.gate + token * gated_width; + for (0..gated_heads) |head| { + const head_base = base + head * gated_head_dim; + const zhead = zbase + head * gated_head_dim; + var sum_sq: f64 = 0; + for (0..gated_head_dim) |i| { + var bits: u16 = if (case.zero_token and token + 1 == case.tokens) 0x0000 else switch (case.fill) { + .flat => sampleBits(rng, 1), + .dynamic => sampleBits(rng, if (i % 4 == 0) 0x1p10 else if (i % 4 == 2) 0x1p-10 else 1), + }; + if (case.controls and token == 1 and head == 0 and i == 0) bits = 0x7f80; // +inf: head 0's sum goes infinite, norm 0 + if (case.controls and token == 2 and head == 0 and i == 0) bits = 0x7fc0; // nan: head 0's sums go nan + x[head_base + i] = bits; + sum_sq += @as(f64, c.fromBf16(bits)) * @as(f64, c.fromBf16(bits)); + } + const n: f64 = 1 / @sqrt(sum_sq / @as(f64, @floatFromInt(gated_head_dim)) + gated_eps); + for (0..gated_head_dim) |i| { + const zb: u16 = if (case.controls and token == 0 and head == 0 and i < gated_z_controls.len) gated_z_controls[i] else if (case.controls and token == 1 and head == 0 and i == 1) 0x7fc0 else gatedGateBits(rng); + z[zhead + i] = zb; + // Shader's association: value * (norm * weight), one mid bf16 rounding; weight rides the column. + const normed = c.toBf16(@floatCast(c.fromBf16(x[head_base + i]) * (n * c.fromBf16(weight[case.at.weight + i])))); + want[token * gated_width + head * gated_head_dim + i] = silu64(c.fromBf16(zb)) * c.fromBf16(normed); + } + } + } + return .{ .a = a, .case = case, .x = x, .weight = weight, .z = z, .out = out, .want = want }; + } + + pub fn deinit(f: GatedNorm) void { + f.a.free(f.x); + f.a.free(f.weight); + f.a.free(f.z); + f.a.free(f.out); + f.a.free(f.want); + } + + // Rows-family budget (gamma(depth) of the output's magnitude) plus mids/ops; a zero want demands a signed zero. + pub fn check(f: GatedNorm, got: []const u16) !Report { + if (got.len != f.out.len) return error.LengthMismatch; + const at = f.case.at.out; + try guardIntact(got, at, f.case.tokens * gated_width, poison16); + var r: Report = .{}; + const shape: c.Shape = .{ .rows = f.case.tokens * gated_heads, .n = 1, .k = gated_head_dim }; + for (f.want, 0..) |w, i| { + const g = got[at + i]; + const gf = c.fromBf16(g); + if (!std.math.isFinite(w) or !std.math.isFinite(gf)) { + if (c.classOf(gf) != c.classOf(w)) return error.NonFiniteMismatch; + continue; + } + if (w == 0) { + if (gf != 0 or std.math.signbit(gf) != std.math.signbit(w)) return error.OutOfBound; + continue; + } + const err = @abs(gf - w); + const bound = c.bound(gated_spec, shape, w, @abs(w)); + if (!(std.math.isFinite(bound) and bound > 0)) return error.BadBound; + const ratio = err / bound; + if (ratio > r.worst) r = .{ .worst = ratio, .row = i / gated_width, .col = i % gated_width }; + } + if (r.worst > 1) return error.OutOfBound; + return r; + } +}; + +const testing = std.testing; + +// ---- Host tests: the references and the checkers, no GPU ---- + +test "preserve wants are the stage 2 pipeline bits, and the flushed subnormal rows stay refused" { + for (preserve_cases) |case| { + const f = try PreserveAdd.init(testing.allocator, case, 0x58a2); + defer f.deinit(); + try testing.expect(!f.allow_subnormal_flush); // the preserve gate never qualifies a flush + for (0..f.case.elems) |i| { + const sum = c.fromBf16(f.a_bits[f.case.at.a + i]) + c.fromBf16(f.b_bits[f.case.at.b + i]); + try testing.expectEqual(c.toBf16(sum), f.want[i]); // f32 add, one RNE bf16 store + } + } + // The classes case carries the subnormal rows (0x0001 + 0x0001 wants 0x0002); a signed zero is a BitMismatch. + const g = try PreserveAdd.init(testing.allocator, preserve_cases[2], 0x58a3); + defer g.deinit(); + try testing.expectEqual(@as(u16, 0x0002), g.want[12]); + const bits = try testing.allocator.dupe(u16, g.out); + defer testing.allocator.free(bits); + @memcpy(bits[g.case.at.out..][0..g.want.len], g.want); + try testing.expectEqual(@as(?PreserveAdd.Mismatch, null), try g.firstMismatch(bits)); + bits[g.case.at.out + 12] = 0x0000; + const m = (try g.firstMismatch(bits)).?; + try testing.expectEqual(error.BitMismatch, m.kind); + try testing.expectEqual(@as(u16, 0x0002), m.want); + try testing.expectError(error.BitMismatch, g.check(bits)); +} + +test "preserve classes plants every nonfinite class and the cancel case produces exact tiny sums" { + const g = try PreserveAdd.init(testing.allocator, preserve_cases[2], 0x58a4); + defer g.deinit(); + var classes: std.EnumSet(c.Class) = .empty; + for (g.want) |w| classes.insert(c.classOf(c.fromBf16(w))); + try testing.expect(classes.contains(.nan)); + try testing.expect(classes.contains(.pos_inf)); + try testing.expect(classes.contains(.neg_inf)); + try testing.expect(classes.contains(.finite)); + const f = try PreserveAdd.init(testing.allocator, preserve_cases[1], 0x58a5); + defer f.deinit(); + var tiny: usize = 0; + for (0..f.case.elems) |i| { + const sum = c.fromBf16(f.a_bits[f.case.at.a + i]) + c.fromBf16(f.b_bits[f.case.at.b + i]); + if (@abs(sum) < 0x1p-20) tiny += 1; + } + try testing.expect(tiny > 0); // v + -(v - 0*ulp) pairs cancel to exact zeros +} + +test "preserve checker reports the first bad element and keeps named failures" { + const f = try PreserveAdd.init(testing.allocator, preserve_cases[0], 0x58a6); + defer f.deinit(); + const bits = try testing.allocator.dupe(u16, f.out); + defer testing.allocator.free(bits); + @memcpy(bits[f.case.at.out..][0..f.want.len], f.want); + try testing.expectEqual(@as(?PreserveAdd.Mismatch, null), try f.firstMismatch(bits)); + bits[f.case.at.out + 3] +%= 1; // one bf16 ulp off: a finite bit mismatch + const m = (try f.firstMismatch(bits)).?; + try testing.expectEqual(error.BitMismatch, m.kind); + try testing.expectEqual(@as(usize, 3), m.index); + try testing.expectEqual(f.want[3], m.want); + try testing.expectError(error.BitMismatch, f.check(bits)); + try testing.expectError(error.LengthMismatch, f.firstMismatch(f.out[0 .. f.out.len - 1])); + bits[f.case.at.out + 3] = f.want[3]; + bits[0] = 0x1234; // guard clobbered + try testing.expectError(error.GuardClobbered, f.firstMismatch(bits)); +} + +test "gate references stay finite on random cases and the controls plant nan and signed zeros" { + for (gate_cases) |case| { + const f = try Gate.init(testing.allocator, case, 0x58a7); + defer f.deinit(); + if (!case.controls) for (f.want) |w| try testing.expect(std.math.isFinite(w)); + } + const g = try Gate.init(testing.allocator, gate_cases[1], 0x58a8); + defer g.deinit(); + try testing.expect(std.math.isNan(g.want[0])); // nan gate + try testing.expectEqual(@as(f64, 1), g.want[1]); // +inf gate: sigmoid 1, ctx 1 + try testing.expect(std.math.signbit(g.want[2])); // -inf gate with ctx -1: zero carries ctx's sign + try testing.expectEqual(@as(f64, 0), @abs(g.want[2])); +} + +test "gate qraw layout matches the shader's interleaved read" { + const f = try Gate.init(testing.allocator, gate_cases[0], 0x58a9); + defer f.deinit(); + // Gate element d of row r sits at r*4096 + (d/256)*512 + 256 + d%256; nothing else in the qraw window is written. + for (f.qraw, 0..) |g, i| { + if (i < f.case.at.b) continue; // the poison before the data window + const w = i - f.case.at.b; + const row = w / gate_qraw_row_elems; + const off = w % gate_qraw_row_elems; + const within = off % 512; + const planted = within >= 256; + if (planted != (g != poison16)) return error.TestUnexpectedResult; + if (row >= f.case.rows and g != poison16) return error.TestUnexpectedResult; + } +} + +test "gate checker refuses wrong finite outputs, wrong-signed zeros and class breaks" { + const f = try Gate.init(testing.allocator, gate_cases[1], 0x58aa); + defer f.deinit(); + const bits = try testing.allocator.dupe(u16, f.out); + defer testing.allocator.free(bits); + for (f.want, 0..) |w, i| bits[f.case.at.out + i] = c.toBf16(@floatCast(w)); + _ = try f.check(bits); + // A wrong finite output is out of bound. + const finite_at = blk: { + for (f.want, 0..) |w, i| if (w != 0 and std.math.isFinite(w)) break :blk i; + unreachable; + }; + bits[f.case.at.out + finite_at] = c.toBf16(@floatCast(f.want[finite_at] * 1.1)); + try testing.expectError(error.OutOfBound, f.check(bits)); + // A sign-flipped zero is refused: the product's sign is deterministic. + for (f.want, 0..) |w, i| bits[f.case.at.out + i] = c.toBf16(@floatCast(w)); + const zero_at = blk: { + for (f.want, 0..) |w, i| if (w == 0) break :blk i; + unreachable; + }; + bits[f.case.at.out + zero_at] ^= 0x8000; + try testing.expectError(error.OutOfBound, f.check(bits)); + // A class break at the nan gate. + for (f.want, 0..) |w, i| bits[f.case.at.out + i] = c.toBf16(@floatCast(w)); + bits[f.case.at.out] = 0x3f80; // nan wanted, 1 got + try testing.expectError(error.NonFiniteMismatch, f.check(bits)); +} + +test "gated references are finite off the controls, and the zero token wants exact signed zeros" { + for (gated_cases) |case| { + const f = try GatedNorm.init(testing.allocator, case, 0x58ab); + defer f.deinit(); + if (case.controls) continue; + for (f.want) |w| try testing.expect(std.math.isFinite(w)); + if (case.zero_token) { + const zero_at = (case.tokens - 1) * gated_width; + for (f.want[zero_at..]) |w| try testing.expectEqual(@as(f64, 0), @abs(w)); + } + } + // The controls: token 0's z give nan, both infinities and signed zeros; tokens 1 and 2 poison only head 0's row. + const g = try GatedNorm.init(testing.allocator, gated_cases[2], 0x58ac); + defer g.deinit(); + try testing.expect(std.math.isNan(g.want[0])); // z = nan + try testing.expect(std.math.isInf(g.want[1])); // z = +inf, silu +inf + try testing.expect(std.math.isNan(g.want[2])); // z = -inf: silu = -inf*0 = nan + try testing.expectEqual(@as(f64, 0), @abs(g.want[3])); // z = +0 + for (0..gated_head_dim) |i| { // token 1, head 0 + const w = g.want[gated_width + i]; + if (i == 0 or i == 1) { + try testing.expect(std.math.isNan(w)); + } else { + try testing.expectEqual(@as(f64, 0), @abs(w)); + } + } + for (gated_width + gated_head_dim..2 * gated_width) |i| // token 1, heads 1..15 + try testing.expect(std.math.isFinite(g.want[i])); + for (2 * gated_width..2 * gated_width + gated_head_dim) |i| // token 2, head 0: all nan + try testing.expect(std.math.isNan(g.want[i])); + for (2 * gated_width + gated_head_dim..3 * gated_width) |i| // token 2, heads 1..15 + try testing.expect(std.math.isFinite(g.want[i])); +} + +test "gated checker refuses wrong finite outputs, sign-flipped zeros and class breaks" { + // Random case: one wrong finite output is out of bound. + { + const f = try GatedNorm.init(testing.allocator, gated_cases[0], 0x58ad); + defer f.deinit(); + const bits = try testing.allocator.dupe(u16, f.out); + defer testing.allocator.free(bits); + for (f.want, 0..) |w, i| bits[f.case.at.out + i] = c.toBf16(@floatCast(w)); + _ = try f.check(bits); + // Doubling is many bf16 ulps off at EVERY magnitude; a fixed point (w*1.1+0.05 at -0.5) bf16 collapses onto. + bits[f.case.at.out + gated_width / 2] = c.toBf16(@floatCast(f.want[gated_width / 2] * 2)); + try testing.expectError(error.OutOfBound, f.check(bits)); + } + // Zero token: one bf16 subnormal unit where an exact zero belongs is refused. + { + const f = try GatedNorm.init(testing.allocator, gated_cases[1], 0x58ae); + defer f.deinit(); + const bits = try testing.allocator.dupe(u16, f.out); + defer testing.allocator.free(bits); + for (f.want, 0..) |w, i| bits[f.case.at.out + i] = c.toBf16(@floatCast(w)); + _ = try f.check(bits); + bits[f.case.at.out + (f.case.tokens - 1) * gated_width] = 0x0001; + try testing.expectError(error.OutOfBound, f.check(bits)); + } + // Controls: a sign-flipped zero and a class break are both refused. + { + const f = try GatedNorm.init(testing.allocator, gated_cases[2], 0x58af); + defer f.deinit(); + const bits = try testing.allocator.dupe(u16, f.out); + defer testing.allocator.free(bits); + for (f.want, 0..) |w, i| bits[f.case.at.out + i] = c.toBf16(@floatCast(w)); + _ = try f.check(bits); + const zero_at = blk: { + for (f.want[0..gated_width], 0..) |w, i| if (w == 0) break :blk i; + unreachable; + }; + bits[f.case.at.out + zero_at] ^= 0x8000; + try testing.expectError(error.OutOfBound, f.check(bits)); + for (f.want, 0..) |w, i| bits[f.case.at.out + i] = c.toBf16(@floatCast(w)); + bits[f.case.at.out] = 0x3f80; // nan wanted, 1 got + try testing.expectError(error.NonFiniteMismatch, f.check(bits)); + } +} + +test "the production plan accepts every gated case, and the old 2048-gain shape is refused" { + for (gated_cases) |case| { + const elems = case.tokens * gated_width; + const plan = try vk.gdn.gatedNormPlan(elems, elems, elems, gated_head_dim, case.tokens, gated_eps); + try testing.expectEqual(case.tokens * gated_heads, plan.rows); // one head row a workgroup + } + // The old runner's [2048] weight is what the plan refuses: the production gain is the [128] head row. + try testing.expectError(error.Invalid, vk.gdn.gatedNormPlan(gated_width, gated_width, gated_width, gated_width, 1, gated_eps)); + // A wrong element count and a bad eps are refused too. + try testing.expectError(error.Invalid, vk.gdn.gatedNormPlan(gated_width + 2, gated_width, gated_width, gated_head_dim, 1, gated_eps)); + try testing.expectError(error.Invalid, vk.gdn.gatedNormPlan(gated_width, gated_width, gated_width, gated_head_dim, 1, 0)); +} + +test "gated fixture element counts are the production geometry, guards outside" { + for (gated_cases) |case| { + const f = try GatedNorm.init(testing.allocator, case, 0x58b6); + defer f.deinit(); + const elems = case.tokens * gated_width; + try testing.expectEqual(case.at.x + elems + tail, f.x.len); + try testing.expectEqual(case.at.gate + elems + tail, f.z.len); + try testing.expectEqual(case.at.out + elems + tail, f.out.len); + try testing.expectEqual(case.at.weight + gated_head_dim + tail, f.weight.len); + try testing.expectEqual(elems, f.want.len); + } +} + +fn preserveAllocRun(a: std.mem.Allocator, case: PreserveCase) !void { + (try PreserveAdd.init(a, case, 0x58b0)).deinit(); +} + +// OpCapability (2,17), DenormPreserve (4464): -DPRESERVE_F32 preserves; catches a no-ifdef define (the old -DUSE_BF16). +fn hasCapability(spv: []align(4) const u8, cap: u32) bool { + const words = std.mem.bytesAsSlice(u32, spv); + if (words.len < 5 or words[0] != 0x07230203) return false; + for (words[1 .. words.len - 1], 1..) |w, i| { + if (w == ((2 << 16) | 17) and words[i + 1] == cap) return true; + } + return false; +} + +test "the preserve build requests DenormPreserve and the plain add does not" { + try testing.expect(hasCapability(&shaders2.add_preserve_bf16, 4464)); + try testing.expect(!hasCapability(&shaders2.add_bf16, 4464)); +} + +fn gateAllocRun(a: std.mem.Allocator, case: GateCase) !void { + (try Gate.init(a, case, 0x58b1)).deinit(); +} + +fn gatedAllocRun(a: std.mem.Allocator, case: GatedCase) !void { + (try GatedNorm.init(a, case, 0x58b2)).deinit(); +} + +test "every fixture init survives each single allocation failure without leaking" { + try std.testing.checkAllAllocationFailures(testing.allocator, preserveAllocRun, .{preserve_cases[0]}); + try std.testing.checkAllAllocationFailures(testing.allocator, preserveAllocRun, .{preserve_cases[2]}); + for (gate_cases) |case| try std.testing.checkAllAllocationFailures(testing.allocator, gateAllocRun, .{case}); + for (gated_cases) |case| try std.testing.checkAllAllocationFailures(testing.allocator, gatedAllocRun, .{case}); +} + +test "guards sit poisoned outside every data window" { + const f = try GatedNorm.init(testing.allocator, gated_cases[1], 0x58b3); + defer f.deinit(); + for (f.x, 0..) |g, i| { + if (i >= f.case.at.x and i < f.case.at.x + f.case.tokens * gated_width) continue; + try testing.expectEqual(poison16, g); + } + for (f.z, 0..) |g, i| { + if (i >= f.case.at.gate and i < f.case.at.gate + f.case.tokens * gated_width) continue; + try testing.expectEqual(poison16, g); + } + const p = try PreserveAdd.init(testing.allocator, preserve_cases[0], 0x58b4); + defer p.deinit(); + for (p.a_bits, 0..) |g, i| { + if (i >= p.case.at.a and i < p.case.at.a + p.case.elems) continue; + try testing.expectEqual(poison16, g); + } +} diff --git a/zig/tests/vulkan/contract_pointwise_main.zig b/zig/tests/vulkan/contract_pointwise_main.zig new file mode 100644 index 00000000000..cb77eed17ed --- /dev/null +++ b/zig/tests/vulkan/contract_pointwise_main.zig @@ -0,0 +1,559 @@ +// Vulkan pointwise contract tests six stage 2-4 operators; exit codes: 0 pass, 1 fail, 2 usage, 124 timeout. +const std = @import("std"); +const fx = @import("contract_pointwise.zig"); +const fx2 = @import("stage2_fixture.zig"); +const contract = @import("contract"); +const vk = @import("vulkan"); +const shaders2 = @import("stage2_shaders"); +const shaders3 = @import("stage3_shaders"); +const shaders4 = @import("stage4_shaders"); + +const cx = contract.spec; + +const usage = "usage: tf-contract-pointwise [--vendor ID] [--timeout-ms MS] | --help\n" ++ + " --help show this text and exit without opening a loader or device\n" ++ + " --vendor ID refuse a device whose Vulkan vendorID differs (decimal or 0x hex; no default)\n" ++ + " --timeout-ms MS bound on every fence wait, 1 to 600000 (default 10000); a timeout ends the process\n"; + +/// Milliseconds on the awake clock since main started, printed on every milestone. +const Elapsed = struct { + io: std.Io, + start: std.Io.Timestamp, + + fn ms(e: Elapsed) i64 { + return e.start.durationTo(std.Io.Clock.awake.now(e.io)).toMilliseconds(); + } +}; + +const repeats = 3; +const planted_bits = [_]u16{ 0x7fc0, 0x7f80, 0xff80 }; +const planted = [_][]const u8{ "nan", "+inf", "-inf" }; + +// The first output whose reference is finite (u16 bits, u32 bits or f64): where a planted nonfinite must be refused. +fn firstFinite(want: anytype) ?usize { + for (want, 0..) |w, i| { + const v: f64 = switch (@TypeOf(w)) { + f64 => w, + u16 => cx.fromBf16(w), + u32 => @as(f32, @bitCast(w)), + else => @compileError("want elements are f64, u16 or u32"), + }; + if (std.math.isFinite(v)) return i; + } + return null; +} + +// attention_gate.comp's one-word push block (the wrapper's GateParams, local because the wrapper keeps it private). +const GateParams = extern struct { words: u32 = 0 }; +comptime { + if (@sizeOf(GateParams) != 4) @compileError("attention_gate push block must stay one word"); + if (gated_push_bytes != @offsetOf(vk.params.Params, "beta") + 4) @compileError("fast_norm_gated push block drifted from Params"); +} +/// fast_norm_gated.comp's push block ends at `beta`, exactly the wrapper's norm_push_bytes. +const gated_push_bytes: u32 = 68; + +pub fn main(init: std.process.Init) !u8 { + const clock: Elapsed = .{ .io = init.io, .start = std.Io.Clock.awake.now(init.io) }; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + for (args[1..]) |arg| if (std.mem.eql(u8, arg, "--help")) { + std.debug.print(usage, .{}); + return 0; + }; + var vendor: ?u32 = null; + var timeout_ms: u64 = 10_000; + var i: usize = 1; + while (i < args.len) : (i += 2) { + const value = if (i + 1 < args.len) args[i + 1] else ""; + if (std.mem.eql(u8, args[i], "--vendor")) { + vendor = std.fmt.parseInt(u32, value, 0) catch null; + if (vendor == null) break; + } else if (std.mem.eql(u8, args[i], "--timeout-ms")) { + timeout_ms = std.fmt.parseInt(u64, value, 10) catch 0; + if (timeout_ms == 0 or timeout_ms > 600_000) break; + } else break; + } else return if (try probe(init.gpa, clock, vendor, timeout_ms * std.time.ns_per_ms)) 0 else 1; + std.debug.print(usage, .{}); + return 2; +} + +fn probe(a: std.mem.Allocator, clock: Elapsed, vendor: ?u32, timeout: u64) !bool { + var rt = try vk.Runtime.open(); + defer rt.close(); + var dev = try vk.Device.open(&rt, .{ .features = .{ .storage16 = true, .int16 = true } }); + defer dev.close(); + const pr = dev.properties; + const name = std.mem.sliceTo(&pr.deviceName, 0); + const types = [_][]const u8{ "other", "integrated", "discrete", "virtual", "cpu" }; + const device_kind = if (pr.deviceType < types.len) types[pr.deviceType] else "unknown"; + const v = pr.apiVersion; + std.debug.print("MILESTONE device-found elapsed_ms={d} name=\"{s}\" vendor=0x{x:0>4} device=0x{x:0>4} type={s} api={d}.{d}.{d} " ++ + "storage16={} int16={} preserve32={} flush32={} fence-timeout={d}ms\n", .{ clock.ms(), name, pr.vendorID, pr.deviceID, device_kind, v >> 22 & 0x7f, v >> 12 & 0x3ff, v & 0xfff, dev.features.storage16, dev.features.int16, dev.v12.shaderDenormPreserveFloat32 != 0, dev.v12.shaderDenormFlushToZeroFloat32 != 0, timeout / std.time.ns_per_ms }); + if (pr.deviceType == vk.c.VK_PHYSICAL_DEVICE_TYPE_CPU) { + std.debug.print("FAIL device: a CPU (software) Vulkan device is never a GPU result\n", .{}); + return false; + } + if (vendor) |want| if (pr.vendorID != want) { + std.debug.print("FAIL device: vendor 0x{x:0>4} is not the requested 0x{x:0>4}\n", .{ pr.vendorID, want }); + return false; + }; + + var cmds = try vk.Commands.init(&dev); + defer cmds.deinit(); + var staging = try vk.Buffer.init(&dev, .staging, 64 << 10); + defer staging.deinit(); + // Stage 2 tests production RMSNorm, residual add, SwiGLU, preserve-add, output gate and gated norm pipelines. + var m = try vk.norm_mlp.NormMlp.init(&dev, .{ + .norm_bf16 = &shaders2.norm_bf16, + .add_bf16 = &shaders2.add_bf16, + .add_preserve_bf16 = &shaders2.add_preserve_bf16, + .swiglu_bf16 = &shaders2.swiglu_bf16, + .norm_f32 = &shaders2.norm_f32, + .norm_f32_in = &shaders2.norm_f32_in, + .norm_bf16_in = &shaders2.norm_bf16_in, + .add_f32 = &shaders2.add_f32, + .add_wide_bf16 = &shaders2.add_wide_bf16, + .add_bf16_f32 = &shaders2.add_bf16_f32, + .swiglu_f32 = &shaders2.swiglu_f32, + }); + defer m.deinit(); + var preserve = try vk.dispatch.Pipeline.init(&dev, &shaders2.add_preserve_bf16, .{ .bindings = 3, .push_bytes = vk.params.cast_push_bytes, .writes = 1 << 2 }); + defer preserve.deinit(); + var gate = try vk.dispatch.Pipeline.init(&dev, &shaders3.gate_bf16, .{ .bindings = 3, .push_bytes = @sizeOf(GateParams), .writes = 1 << 2 }); + defer gate.deinit(); + var gated = try vk.dispatch.Pipeline.init(&dev, &shaders4.fast_norm_gated_bf16, .{ .bindings = 4, .push_bytes = gated_push_bytes, .writes = 1 << 3 }); + defer gated.deinit(); + + var passed: usize = 0; + var total: usize = 0; + for (fx2.add_cases) |case| { + total += 1; + passed += @intFromBool(guardedGate("residual-add", case.name, addGate, .{ a, clock, &dev, &cmds, staging, &m, case, timeout })); + } + for (fx2.swiglu_cases) |case| { + total += 1; + passed += @intFromBool(guardedGate("swiglu", case.name, swigluGate, .{ a, clock, &dev, &cmds, staging, &m, case, timeout })); + } + for (fx2.rms_cases) |case| { + total += 1; + passed += @intFromBool(guardedGate("rms-norm", case.name, rmsGate, .{ a, clock, &dev, &cmds, staging, &m, case, timeout })); + } + for (fx.preserve_cases) |case| { + total += 1; + const known_before = known_findings; + const ok = guardedGate("preserve-add", case.name, preserveGate, .{ a, clock, &dev, &cmds, staging, &preserve, case, timeout }); + passed += @intFromBool(ok and known_findings == known_before); + } + for (fx.gate_cases) |case| { + total += 1; + passed += @intFromBool(guardedGate("output-gate", case.name, outputGate, .{ a, clock, &dev, &cmds, staging, &gate, case, timeout })); + } + for (fx.gated_cases) |case| { + total += 1; + passed += @intFromBool(guardedGate("gated-norm", case.name, gatedNormGate, .{ a, clock, &dev, &cmds, staging, &gated, case, timeout })); + } + // Each refused shape or plan error counts as that case's failure, not process termination. + std.debug.print("RESULT contract-pointwise elapsed_ms={d} on \"{s}\": passed={d} known={d} total={d}\n", .{ clock.ms(), name, passed, known_findings, total }); + return passed + known_findings == total; +} + +// A gate error prints that case's failure and continues to the final result line. +fn guardedGate(gate: []const u8, case_name: []const u8, comptime f: anytype, args: anytype) bool { + return @call(.auto, f, args) catch |e| blk: { + std.debug.print("FAIL {s} case=\"{s}\": {s}\n", .{ gate, case_name, @errorName(e) }); + break :blk false; + }; +} + +/// Cases that ended as a KNOWN device finding (counted apart from passes). +var known_findings: usize = 0; + +fn uploaded(_: std.mem.Allocator, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, bytes: []const u8, timeout: u64) !vk.Buffer { + var b = try vk.Buffer.init(dev, .device, bytes.len); + errdefer b.deinit(); + try vk.orExit(vk.upload(cmds, b, 0, staging, bytes, timeout)); + return b; +} + +fn fetch(cmds: *vk.Commands, staging: vk.Buffer, b: vk.Buffer, out: []u8, timeout: u64) !void { + try vk.orExit(vk.download(cmds, b, 0, staging, out, timeout)); +} + +fn checkUnchanged(a: std.mem.Allocator, cmds: *vk.Commands, staging: vk.Buffer, shell: anytype, buf: vk.Buffer, timeout: u64) !bool { + const bytes = std.mem.sliceAsBytes(shell); + const back = try a.alloc(u8, bytes.len); + defer a.free(back); + try fetch(cmds, staging, buf, back, timeout); + return std.mem.indexOfDiff(u8, bytes, back) == null; +} + +// Bind the shell through caller-owned Bound16; storing its view would reference boundShell's dead stack frame. +const Bound16 = struct { + buf: vk.Buffer, + byte_offset: u64, + elements: u64, + + fn view(b: *const Bound16) !vk.tensor.View { + return vk.tensor.View.init(&b.buf, .bf16, b.byte_offset, b.elements); + } +}; + +fn boundShell(a: std.mem.Allocator, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, shell: []const u16, at_elems: usize, elements: u64, timeout: u64) !Bound16 { + var buf = try uploaded(a, dev, cmds, staging, std.mem.sliceAsBytes(shell), timeout); + errdefer buf.deinit(); + return .{ .buf = buf, .byte_offset = at_elems * 2, .elements = elements }; +} + +fn reportFail(comptime what: []const u8, name: []const u8, e: anyerror) bool { + std.debug.print("FAIL " ++ what ++ " case=\"" ++ "{s}" ++ "\": " ++ "{s}" ++ "\n", .{ name, @errorName(e) }); + return false; +} + +// Stage 2 gates share fixtures and production operators with its probe, so this contract runs independently. +fn addGate(a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, m: *vk.norm_mlp.NormMlp, case: fx2.AddCase, timeout: u64) !bool { + var f = try fx2.Add.init(a, case, 0x59a9); + defer f.deinit(); + // Only devices lacking f32 denormal preservation may flush subnormal classes; report the property and finding. + f.allow_subnormal_flush = dev.v12.shaderDenormPreserveFloat32 == 0; + var in_a = try boundShell(a, dev, cmds, staging, f.a_bits, case.at.a, case.elems, timeout); + defer in_a.buf.deinit(); + var in_b = try boundShell(a, dev, cmds, staging, f.b_bits, case.at.b, case.elems, timeout); + defer in_b.buf.deinit(); + var out = try boundShell(a, dev, cmds, staging, f.out, case.at.out, case.elems, timeout); + defer out.buf.deinit(); + const bits = try a.alloc(u16, f.out.len); + defer a.free(bits); + var first: ?[]u16 = null; + defer if (first) |x| a.free(x); + for (0..repeats) |rep| { + try vk.orExit(vk.upload(cmds, out.buf, 0, staging, std.mem.sliceAsBytes(f.out), timeout)); + try vk.orExit(m.add(cmds, try out.view(), try in_a.view(), try in_b.view(), timeout)); + try fetch(cmds, staging, out.buf, std.mem.sliceAsBytes(bits), timeout); + if (first) |x| { + if (std.mem.indexOfDiff(u16, x, bits)) |at| { + std.debug.print("FAIL residual-add case=\"{s}\": repeat {d} differs at element {d}\n", .{ case.name, rep, at }); + return false; + } + } else first = try a.dupe(u16, bits); + } + const mm = f.firstMismatch(bits) catch |e| return reportFail("residual-add", case.name, e); + if (mm) |bad| { + std.debug.print("FAIL residual-add case=\"{s}\": {s} at element {d}: a=0x{x:0>4} ({s}) b=0x{x:0>4} ({s}) want=0x{x:0>4} ({s}) got=0x{x:0>4} ({s})\n", .{ case.name, @errorName(bad.kind), bad.index, bad.a, fx.classify(bad.a), bad.b, fx.classify(bad.b), bad.want, fx.classify(bad.want), bad.got, fx.classify(bad.got) }); + return false; + } + const data_at = case.at.out + firstFinite(f.want).?; + var mutated = try a.dupe(u16, bits); + defer a.free(mutated); + for (planted_bits, planted) |pb, pname| { + mutated[data_at] = pb; + if (f.check(mutated)) |_| { + std.debug.print("FAIL residual-add case=\"{s}\": a planted {s} output passed the check\n", .{ case.name, pname }); + return false; + } else |_| {} + } + if (!try checkUnchanged(a, cmds, staging, f.a_bits, in_a.buf, timeout) or !try checkUnchanged(a, cmds, staging, f.b_bits, in_b.buf, timeout)) { + std.debug.print("FAIL residual-add case=\"{s}\": an input binding changed\n", .{case.name}); + return false; + } + std.debug.print("MILESTONE residual-add elapsed_ms={d} case=\"{s}\" elements={d} matches the f32-add + RNE-store pipeline bit for bit, subnormal-class-flush-allowed={}, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused\n", .{ clock.ms(), case.name, case.elems, f.allow_subnormal_flush and case.fill == .classes }); + return true; +} + +fn swigluGate(a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, m: *vk.norm_mlp.NormMlp, case: fx2.SwigluCase, timeout: u64) !bool { + const f = try fx2.Swiglu.init(a, case, 0x59aa); + defer f.deinit(); + var gate = try boundShell(a, dev, cmds, staging, f.gate, case.at.a, case.elems, timeout); + defer gate.buf.deinit(); + var up = try boundShell(a, dev, cmds, staging, f.up, case.at.b, case.elems, timeout); + defer up.buf.deinit(); + var out = try boundShell(a, dev, cmds, staging, f.out, case.at.out, case.elems, timeout); + defer out.buf.deinit(); + const bits = try a.alloc(u16, f.out.len); + defer a.free(bits); + var first: ?[]u16 = null; + defer if (first) |x| a.free(x); + var r: fx.Report = .{}; + for (0..repeats) |rep| { + try vk.orExit(vk.upload(cmds, out.buf, 0, staging, std.mem.sliceAsBytes(f.out), timeout)); + try vk.orExit(m.swiglu(cmds, try out.view(), try gate.view(), try up.view(), timeout)); + try fetch(cmds, staging, out.buf, std.mem.sliceAsBytes(bits), timeout); + r = f.check(bits) catch |e| return reportFail("swiglu", case.name, e); + if (first) |x| { + if (std.mem.indexOfDiff(u16, x, bits)) |at| { + std.debug.print("FAIL swiglu case=\"{s}\": repeat {d} differs at element {d}\n", .{ case.name, rep, at }); + return false; + } + } else first = try a.dupe(u16, bits); + } + var mutated = try a.dupe(u16, bits); + defer a.free(mutated); + const data_at = case.at.out + firstFinite(f.want).?; + for (planted_bits, planted) |pb, pname| { + mutated[data_at] = pb; + if (f.check(mutated)) |_| { + std.debug.print("FAIL swiglu case=\"{s}\": a planted {s} output passed the check\n", .{ case.name, pname }); + return false; + } else |_| {} + } + if (!try checkUnchanged(a, cmds, staging, f.gate, gate.buf, timeout) or !try checkUnchanged(a, cmds, staging, f.up, up.buf, timeout)) { + std.debug.print("FAIL swiglu case=\"{s}\": an input binding changed\n", .{case.name}); + return false; + } + std.debug.print("MILESTONE swiglu elapsed_ms={d} case=\"{s}\" elements={d} worst={d:.4} of bound at element {d}, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused\n", .{ clock.ms(), case.name, case.elems, r.worst, r.col }); + return true; +} + +fn rmsGate(a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, m: *vk.norm_mlp.NormMlp, case: fx2.RmsCase, timeout: u64) !bool { + const f = try fx2.Rms.init(a, case, 0x59ab); + defer f.deinit(); + var x = try boundShell(a, dev, cmds, staging, f.x, case.at.x, case.rows * case.width, timeout); + defer x.buf.deinit(); + var gamma = try boundShell(a, dev, cmds, staging, f.gamma, case.at.gamma, case.width, timeout); + defer gamma.buf.deinit(); + var out = try boundShell(a, dev, cmds, staging, f.out, case.at.out, case.rows * case.width, timeout); + defer out.buf.deinit(); + const bits = try a.alloc(u16, f.out.len); + defer a.free(bits); + var first: ?[]u16 = null; + defer if (first) |xx| a.free(xx); + var r: fx.Report = .{}; + for (0..repeats) |rep| { + try vk.orExit(vk.upload(cmds, out.buf, 0, staging, std.mem.sliceAsBytes(f.out), timeout)); + try vk.orExit(m.rmsNorm(cmds, try out.view(), try x.view(), try gamma.view(), case.rows, case.width, fx2.rms_eps, timeout)); + try fetch(cmds, staging, out.buf, std.mem.sliceAsBytes(bits), timeout); + r = f.check(bits) catch |e| return reportFail("rms-norm", case.name, e); + if (first) |xx| { + if (std.mem.indexOfDiff(u16, xx, bits)) |at| { + std.debug.print("FAIL rms-norm case=\"{s}\": repeat {d} differs at element {d}\n", .{ case.name, rep, at }); + return false; + } + } else first = try a.dupe(u16, bits); + } + var mutated = try a.dupe(u16, bits); + defer a.free(mutated); + const data_at = case.at.out + r.row * case.width + r.col; + for (planted_bits, planted) |pb, pname| { + mutated[data_at] = pb; + if (f.check(mutated)) |_| { + std.debug.print("FAIL rms-norm case=\"{s}\": a planted {s} output passed the check\n", .{ case.name, pname }); + return false; + } else |_| {} + } + if (!try checkUnchanged(a, cmds, staging, f.x, x.buf, timeout) or !try checkUnchanged(a, cmds, staging, f.gamma, gamma.buf, timeout)) { + std.debug.print("FAIL rms-norm case=\"{s}\": an input binding changed\n", .{case.name}); + return false; + } + if (case.zero_row) { + const zero_at = case.at.out + (case.rows - 1) * case.width; + for (bits[zero_at .. zero_at + case.width]) |g| if (g != 0x0000 and g != 0x8000) { + std.debug.print("FAIL rms-norm case=\"{s}\": the zero row produced 0x{x:0>4}\n", .{ case.name, g }); + return false; + }; + } + std.debug.print("MILESTONE rms-norm elapsed_ms={d} case=\"{s}\" rows={d} width={d} worst={d:.4} of bound at row {d} col {d}, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused{s}\n", .{ clock.ms(), case.name, case.rows, case.width, r.worst, r.row, r.col, if (case.zero_row) ", zero row exactly zero" else "" }); + return true; +} + +// Preserve-add is strict: denormal-flushing devices fail classes rows; the reference never adapts to the flush. +fn preserveGate(a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, p: *vk.dispatch.Pipeline, case: fx.PreserveCase, timeout: u64) !bool { + const f = try fx.PreserveAdd.init(a, case, 0x59ac); + defer f.deinit(); + var in_a = try boundShell(a, dev, cmds, staging, f.a_bits, case.at.a, case.elems, timeout); + defer in_a.buf.deinit(); + var in_b = try boundShell(a, dev, cmds, staging, f.b_bits, case.at.b, case.elems, timeout); + defer in_b.buf.deinit(); + var out = try boundShell(a, dev, cmds, staging, f.out, case.at.out, case.elems, timeout); + defer out.buf.deinit(); + const bits = try a.alloc(u16, f.out.len); + defer a.free(bits); + var first: ?[]u16 = null; + defer if (first) |x| a.free(x); + for (0..repeats) |rep| { + try vk.orExit(vk.upload(cmds, out.buf, 0, staging, std.mem.sliceAsBytes(f.out), timeout)); + try vk.orExit(p.bind(&.{ (try in_a.view()).bind(), (try in_b.view()).bind(), (try out.view()).bind() })); + const push: vk.params.Params = .{ .count = try cast32(case.elems), .operation = 0, .lhs_size = try cast32(case.elems), .rhs_size = try cast32(case.elems) }; + try vk.orExit(p.run(cmds, std.mem.asBytes(&push)[0..vk.params.cast_push_bytes], try groups(case.elems, p.local[0], dev.properties.limits.maxComputeWorkGroupCount[0]), timeout)); + try fetch(cmds, staging, out.buf, std.mem.sliceAsBytes(bits), timeout); + if (first) |x| { + if (std.mem.indexOfDiff(u16, x, bits)) |at| { + std.debug.print("FAIL preserve-add case=\"{s}\": repeat {d} differs at element {d}\n", .{ case.name, rep, at }); + return false; + } + } else first = try a.dupe(u16, bits); + } + const mm = f.firstMismatch(bits) catch |e| return reportFail("preserve-add", case.name, e); + if (mm) |bad| { + // Same-sign subnormal-to-zero mismatches may be known on non-preserving devices; all others fail. + if (case.fill == .classes and dev.v12.shaderDenormPreserveFloat32 == 0) { + var flushes: usize = 0; + var cur: ?@TypeOf(bad) = bad; + while (cur) |c| { + const flush = c.kind == error.BitMismatch and isBf16Subnormal(c.want) and c.got & 0x7fff == 0 and (c.got ^ c.want) & 0x8000 == 0; + if (!flush) break; + flushes += 1; + bits[c.index] = c.want; + cur = f.firstMismatch(bits) catch null; + if (cur != null and cur.?.index <= c.index) break; + } + if (cur == null) { + known_findings += 1; + std.debug.print("KNOWN preserve-add case=\"{s}\": {d} subnormal sums read as zero, honest by design — the device reports preserve32=false (shaderDenormPreserveFloat32 = 0) and flushes fp32 denormals, so the DenormPreserve image cannot keep them: recorded, not fitted\n", .{ case.name, flushes }); + return true; + } + } + std.debug.print("FAIL preserve-add case=\"{s}\": {s} at element {d}: a=0x{x:0>4} ({s}) b=0x{x:0>4} ({s}) want=0x{x:0>4} ({s}) got=0x{x:0>4} ({s})\n", .{ case.name, @errorName(bad.kind), bad.index, bad.a, fx.classify(bad.a), bad.b, fx.classify(bad.b), bad.want, fx.classify(bad.want), bad.got, fx.classify(bad.got) }); + if (case.fill == .classes and isBf16Subnormal(bad.want)) { + std.debug.print("NOTE preserve-add: want is a bf16 subnormal whose f32 sum the pipeline must carry through fp32 denormals, and got differs — this device flushes fp32 denormals despite the requested DenormPreserve; the preserve gate fails here honestly by design\n", .{}); + } + return false; + } + // A planted nonfinite output at a finite-want element must be refused by the same checker. + const data_at = case.at.out + firstFinite(f.want).?; + var mutated = try a.dupe(u16, bits); + defer a.free(mutated); + for (planted_bits, planted) |pb, pname| { + mutated[data_at] = pb; + if (f.check(mutated)) |_| { + std.debug.print("FAIL preserve-add case=\"{s}\": a planted {s} output passed the check\n", .{ case.name, pname }); + return false; + } else |_| {} + } + if (!try checkUnchanged(a, cmds, staging, f.a_bits, in_a.buf, timeout) or !try checkUnchanged(a, cmds, staging, f.b_bits, in_b.buf, timeout)) { + std.debug.print("FAIL preserve-add case=\"{s}\": an input binding changed\n", .{case.name}); + return false; + } + std.debug.print("MILESTONE preserve-add elapsed_ms={d} case=\"{s}\" elements={d} matches the stage 2 add pipeline bits over the DenormPreserve image, strict on subnormal sums, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused\n", .{ clock.ms(), case.name, case.elems }); + return true; +} + +/// bf16 bits whose value is a subnormal: exponent 0 with a nonzero mantissa. +fn isBf16Subnormal(bits: u16) bool { + return bits & 0x7f80 == 0 and bits & 0x7f != 0; +} + +/// The attention output gate: bf16(sigmoid(gate)) once, times ctx, one bf16 store. +fn outputGate(a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, gate: *vk.dispatch.Pipeline, case: fx.GateCase, timeout: u64) !bool { + const f = try fx.Gate.init(a, case, 0x59ad); + defer f.deinit(); + const elems = case.rows * fx.gate_row_elems; + var ctx = try boundShell(a, dev, cmds, staging, f.ctx, case.at.a, elems, timeout); + defer ctx.buf.deinit(); + var qraw = try boundShell(a, dev, cmds, staging, f.qraw, case.at.b, case.rows * fx.gate_qraw_row_elems, timeout); + defer qraw.buf.deinit(); + var out = try boundShell(a, dev, cmds, staging, f.out, case.at.out, elems, timeout); + defer out.buf.deinit(); + const bits = try a.alloc(u16, f.out.len); + defer a.free(bits); + var first: ?[]u16 = null; + defer if (first) |x| a.free(x); + var r: fx.Report = .{}; + for (0..repeats) |rep| { + try vk.orExit(vk.upload(cmds, out.buf, 0, staging, std.mem.sliceAsBytes(f.out), timeout)); + try vk.orExit(gate.bind(&.{ (try ctx.view()).bind(), (try qraw.view()).bind(), (try out.view()).bind() })); + const plan = try vk.attention.gatePlan(elems, elems, case.rows * fx.gate_qraw_row_elems, case.rows); + const push: GateParams = .{ .words = plan.words }; + try vk.orExit(gate.run(cmds, std.mem.asBytes(&push), try groups(plan.words, gate.local[0], dev.properties.limits.maxComputeWorkGroupCount[0]), timeout)); + try fetch(cmds, staging, out.buf, std.mem.sliceAsBytes(bits), timeout); + r = f.check(bits) catch |e| return reportFail("output-gate", case.name, e); + if (first) |x| { + if (std.mem.indexOfDiff(u16, x, bits)) |at| { + std.debug.print("FAIL output-gate case=\"{s}\": repeat {d} differs at element {d}\n", .{ case.name, rep, at }); + return false; + } + } else first = try a.dupe(u16, bits); + } + var mutated = try a.dupe(u16, bits); + defer a.free(mutated); + const data_at = case.at.out + firstFinite(f.want).?; + for (planted_bits, planted) |pb, pname| { + mutated[data_at] = pb; + if (f.check(mutated)) |_| { + std.debug.print("FAIL output-gate case=\"{s}\": a planted {s} output passed the check\n", .{ case.name, pname }); + return false; + } else |_| {} + } + if (!try checkUnchanged(a, cmds, staging, f.ctx, ctx.buf, timeout) or !try checkUnchanged(a, cmds, staging, f.qraw, qraw.buf, timeout)) { + std.debug.print("FAIL output-gate case=\"{s}\": an input binding changed\n", .{case.name}); + return false; + } + std.debug.print("MILESTONE output-gate elapsed_ms={d} case=\"{s}\" rows={d} worst={d:.4} of bound at row {d} col {d}, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused\n", .{ clock.ms(), case.name, case.rows, r.worst, r.row, r.col }); + return true; +} + +// Gated RMSNorm uses 16 heads of width 128 and a shared 128-value gain; bound covers reduction and mid budget. +fn gatedNormGate(a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, g: *vk.dispatch.Pipeline, case: fx.GatedCase, timeout: u64) !bool { + const f = try fx.GatedNorm.init(a, case, 0x59ae); + defer f.deinit(); + const elems = case.tokens * fx.gated_width; + var x = try boundShell(a, dev, cmds, staging, f.x, case.at.x, elems, timeout); + defer x.buf.deinit(); + var weight = try boundShell(a, dev, cmds, staging, f.weight, case.at.weight, fx.gated_head_dim, timeout); + defer weight.buf.deinit(); + var z = try boundShell(a, dev, cmds, staging, f.z, case.at.gate, elems, timeout); + defer z.buf.deinit(); + var out = try boundShell(a, dev, cmds, staging, f.out, case.at.out, elems, timeout); + defer out.buf.deinit(); + const bits = try a.alloc(u16, f.out.len); + defer a.free(bits); + var first: ?[]u16 = null; + defer if (first) |xx| a.free(xx); + var r: fx.Report = .{}; + for (0..repeats) |rep| { + try vk.orExit(vk.upload(cmds, out.buf, 0, staging, std.mem.sliceAsBytes(f.out), timeout)); + try vk.orExit(g.bind(&.{ (try x.view()).bind(), (try weight.view()).bind(), (try z.view()).bind(), (try out.view()).bind() })); + // The wrapper plan uses token rows and 16 heads per token; weights contain one 128-value head row. + const plan = try vk.gdn.gatedNormPlan(elems, elems, elems, fx.gated_head_dim, case.tokens, fx.gated_eps); + var push: vk.params.Params = .{ + .operation = 0, // gated mode + .lhs_size = fx.gated_head_dim, // per-column weight within the head + .reduce_size = fx.gated_head_dim, // the 128-wide per-head row + .output_size = plan.rows, + }; + push.alpha = fx.gated_eps; + // One 256-thread workgroup per head row, the wrapper's grid-stride row loop. + try vk.orExit(g.run(cmds, std.mem.asBytes(&push)[0..gated_push_bytes], try groups(plan.rows, 1, dev.properties.limits.maxComputeWorkGroupCount[0]), timeout)); + try fetch(cmds, staging, out.buf, std.mem.sliceAsBytes(bits), timeout); + r = f.check(bits) catch |e| return reportFail("gated-norm", case.name, e); + if (first) |xx| { + if (std.mem.indexOfDiff(u16, xx, bits)) |at| { + std.debug.print("FAIL gated-norm case=\"{s}\": repeat {d} differs at element {d}\n", .{ case.name, rep, at }); + return false; + } + } else first = try a.dupe(u16, bits); + } + var mutated = try a.dupe(u16, bits); + defer a.free(mutated); + const data_at = case.at.out + r.row * fx.gated_width + r.col; + for (planted_bits, planted) |pb, pname| { + mutated[data_at] = pb; + if (f.check(mutated)) |_| { + std.debug.print("FAIL gated-norm case=\"{s}\": a planted {s} output passed the check\n", .{ case.name, pname }); + return false; + } else |_| {} + } + if (!try checkUnchanged(a, cmds, staging, f.x, x.buf, timeout) or !try checkUnchanged(a, cmds, staging, f.weight, weight.buf, timeout) or !try checkUnchanged(a, cmds, staging, f.z, z.buf, timeout)) { + std.debug.print("FAIL gated-norm case=\"{s}\": an input binding changed\n", .{case.name}); + return false; + } + if (case.zero_token) { + const zero_at = case.at.out + (case.tokens - 1) * fx.gated_width; + for (bits[zero_at .. zero_at + fx.gated_width]) |gg| if (gg != 0x0000 and gg != 0x8000) { + std.debug.print("FAIL gated-norm case=\"{s}\": the zero token produced 0x{x:0>4}\n", .{ case.name, gg }); + return false; + }; + } + std.debug.print("MILESTONE gated-norm elapsed_ms={d} case=\"{s}\" tokens={d} head-rows={d} width={d} worst={d:.4} of bound at row {d} col {d}, repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused{s}\n", .{ clock.ms(), case.name, case.tokens, case.tokens * fx.gated_heads, fx.gated_head_dim, r.worst, r.row, r.col, if (case.zero_token) ", zero token exactly zero" else "" }); + return true; +} + +fn cast32(v: u64) !u32 { + return std.math.cast(u32, v) orelse error.TooLarge; +} + +fn groups(count: u64, local: u32, max_x: u32) ![3]u32 { + const n = try cast32((count + local - 1) / local); + return .{ @min(n, max_x), 1, 1 }; +} + +test { + std.testing.refAllDecls(@This()); +} diff --git a/zig/tests/vulkan/contract_qmm.zig b/zig/tests/vulkan/contract_qmm.zig index c1a451ec4e6..94a63c4228d 100644 --- a/zig/tests/vulkan/contract_qmm.zig +++ b/zig/tests/vulkan/contract_qmm.zig @@ -56,6 +56,13 @@ pub const spec: c.Spec = .{ // gate 4 keeps the harness default: fp16 and bf16 sums must miss by 100x (no claim that output rounding hides them) }; +/// The f32-activation image: x and the output are f32, so gate 3's rounding model is the f32 one. +pub const spec_f32: c.Spec = blk: { + var s = spec; + s.out = .f32; + break :blk s; +}; + /// The shared gate harness expects a rig object; as in HIP, it exposes bytes while the fixture owns device buffers. const Rig = struct { output: []const u8 = &.{}, @@ -101,6 +108,12 @@ pub const Fixture = struct { bindings: [5]dispatch.Binding = undefined, /// Planted control: overwrite input buffer i (x, w, scales, biases) after dispatch; gate 2 must catch it. corrupt: ?usize = null, + /// The kernel's reported name when it is not the default qmm_bf16 spec's. + name: ?[]const u8 = null, + /// Tiled image: inputs are repacked into 32-column tiles before upload; shapes that cannot be tiled are refused. + tiled: bool = false, + /// f32 activations: x is stored as f32 (the bf16 values widened exactly) and the output is f32. + act_f32: bool = false, pub fn init(a: std.mem.Allocator, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, pipeline: *dispatch.Pipeline) Fixture { return .{ .a = a, .dev = dev, .cmds = cmds, .staging = staging, .pipeline = pipeline }; @@ -108,7 +121,12 @@ pub const Fixture = struct { pub fn deinit(f: *Fixture) void { f.reset(); } + fn actUnits(f: *const Fixture) usize { + return if (f.act_f32) 2 else 1; + } fn reset(f: *Fixture) void { + // A hook error after a recorded dispatch leaves the buffer open: drop those records before freeing buffers. + f.cmds.discard(); while (f.made > 0) { f.made -= 1; f.buffers[f.made].deinit(); @@ -138,14 +156,30 @@ pub const Fixture = struct { planted.name = corrupt_names[i]; return .{ .spec = planted, .ctx = f, .vt = &vt, .control = .repeat }; } - return .{ .spec = spec, .ctx = f, .vt = &vt }; + var named = if (f.act_f32) spec_f32 else spec; + if (f.name) |n| named.name = n; + return .{ .spec = named, .ctx = f, .vt = &vt }; } + fn allocPoison(a: std.mem.Allocator, data_len: usize) ![]u16 { const extra: usize = if (data_len % 2 == 0) 16 else 17; const p = try a.alloc(u16, data_len + extra); @memset(p, poison); return p; } + // The tiled image reads host-repacked inputs sized for whole tiles; the float64 reference came from the row layout. + fn tileInputs(a: std.mem.Allocator, words: []u8, scales: []u8, biases: []u8, shape: c.Shape) !void { + const real_words = shape.n * shape.k / 2; + const real_params = shape.n * (shape.k / group) * 2; + const copy = try a.alloc(u8, real_words); + defer a.free(copy); + @memcpy(copy, words[0..real_words]); + vk.tile.tileWords(words, copy, shape.n, shape.k); + @memcpy(copy[0..real_params], scales[0..real_params]); + vk.tile.tileParams(scales, copy[0..real_params], shape.n, shape.k); + @memcpy(copy[0..real_params], biases[0..real_params]); + vk.tile.tileParams(biases, copy[0..real_params], shape.n, shape.k); + } pub fn prepare(f: *Fixture, shape: c.Shape, fill: g.Fill, seed: u64) !bool { if (shape.k % group != 0 or shape.n * shape.k > max_weights or shape.rows * shape.n > max_outputs) return false; f.reset(); @@ -182,22 +216,29 @@ pub const Fixture = struct { .nonfinite => .nonfinite, .stress => .stress, }, seed); - const x = try allocPoison(f.a, shape.rows * shape.k); + const x = try allocPoison(f.a, shape.rows * shape.k * f.actUnits()); errdefer f.a.free(x); - @memcpy(x[0..q.x.len], q.x); + if (f.act_f32) { + for (q.x, 0..) |v, i| { + x[2 * i] = 0; + x[2 * i + 1] = v; + } + } else @memcpy(x[0..q.x.len], q.x); + const stored: usize = if (f.tiled) vk.tile.padded(shape.n) else shape.n; const words_len = shape.n * shape.k / 8; - const w = try f.a.alloc(u32, words_len + 4); + const w = try f.a.alloc(u32, stored * shape.k / 8 + 4); errdefer f.a.free(w); @memset(w, 0xffff_ffff); q.words(w[0..words_len]); const nparam = shape.n * (shape.k / group); - const scales = try allocPoison(f.a, nparam); + const nstored = stored * (shape.k / group); + const scales = try allocPoison(f.a, nstored); errdefer f.a.free(scales); @memcpy(scales[0..nparam], q.scales); - const biases = try allocPoison(f.a, nparam); + const biases = try allocPoison(f.a, nstored); errdefer f.a.free(biases); @memcpy(biases[0..nparam], q.biases); - const out = try allocPoison(f.a, shape.rows * shape.n); + const out = try allocPoison(f.a, shape.rows * shape.n * f.actUnits()); errdefer f.a.free(out); const want = try f.a.alloc(f64, shape.rows * shape.n); errdefer f.a.free(want); @@ -205,6 +246,7 @@ pub const Fixture = struct { errdefer f.a.free(mag); try q.reference(want, mag); if (fill != .nonfinite) for (want, mag) |v, m| if (!std.math.isFinite(v) or !std.math.isFinite(m)) return error.NonFiniteReference; + if (f.tiled) try tileInputs(f.a, std.mem.sliceAsBytes(w[0 .. stored * shape.k / 8]), std.mem.sliceAsBytes(scales[0..nstored]), std.mem.sliceAsBytes(biases[0..nstored]), shape); f.x = x; f.w = w; f.scales = scales; @@ -235,7 +277,7 @@ pub const Fixture = struct { try vk.orExit(vk.download(f.cmds, b.buffer.*, 0, f.staging, f.input_snapshot[at..][0..b.len], timeout_ns)); at += b.len; } - const data_count = s.rows * s.n; + const data_count = s.rows * s.n * f.actUnits(); f.rig.guards_ok = true; for (f.out[data_count..]) |v| if (v != poison) { f.rig.guards_ok = false; @@ -255,6 +297,18 @@ pub const Fixture = struct { } }; +test "a fixture name replaces the default kernel name" { + var f: Fixture = undefined; + f.corrupt = null; + f.act_f32 = false; + f.name = "vulkan qmm_q4_vec4 affine q4 g64 transposed"; + try std.testing.expectEqualStrings(f.name.?, f.kernel().spec.name); + f.name = "vulkan qmm_q4_tile affine q4 g64 transposed tiled 32"; + try std.testing.expectEqualStrings(f.name.?, f.kernel().spec.name); + f.name = null; + try std.testing.expectEqualStrings(spec.name, f.kernel().spec.name); +} + /// One source shader variant, exercised through the same hook contract as HIP's planted controls. pub const ControlKind = enum { qmm_fmax, atomic_sum, scribble_sum, drop_tail, fp16_accum }; fn controlGate(kind: ControlKind) g.Gate { @@ -311,6 +365,8 @@ pub const Control = struct { f.reset(); } fn reset(f: *Control) void { + // Same rule as the fixture reset: no freed buffer may stay referenced by open records. + f.cmds.discard(); if (f.x_device) |*b| b.deinit(); if (f.w_device) |*b| b.deinit(); if (f.y_device) |*b| b.deinit(); @@ -428,5 +484,51 @@ test "Vulkan bf16 fixture uses raw u16 storage and guards are aligned for storag try std.testing.expectEqual(@as(usize, 2), @sizeOf(u16)); } +test "a failed host allocation after a recorded dispatch is dropped before the fixture and control free their buffers" { + const Calls = struct { + var resets: usize = 0; + var submits: usize = 0; + fn reset(_: vk.c.VkCommandBuffer, _: vk.c.VkCommandBufferResetFlags) callconv(.c) vk.c.VkResult { + resets += 1; + return vk.c.VK_SUCCESS; + } + fn submit(_: vk.c.VkQueue, _: c_uint, _: [*c]const vk.c.VkSubmitInfo, _: vk.c.VkFence) callconv(.c) vk.c.VkResult { + submits += 1; + return vk.c.VK_SUCCESS; + } + }; + var api: @TypeOf(@as(vk.Device, undefined).api) = undefined; + api.vkResetCommandBuffer = Calls.reset; + api.vkQueueSubmit = Calls.submit; + var dev: vk.Device = undefined; + dev.api = api; + var pipeline: dispatch.Pipeline = undefined; + const staging: vk.Buffer = undefined; + var failing = std.testing.FailingAllocator.init(std.testing.allocator, .{ .fail_index = 0 }); + + // Fixture.run records its dispatch, then its corrupt arm allocates a tampered copy; that allocation fails. + var cmds: vk.Commands = .{ .dev = &dev, .pool = undefined, .cmd = @ptrFromInt(0x1000), .fence = undefined, .recording = true, .pending = 2 }; + var fixture = Fixture.init(failing.allocator(), &dev, &cmds, staging, &pipeline); + try std.testing.expectError(error.OutOfMemory, fixture.a.dupe(u8, "x")); + fixture.reset(); + try std.testing.expectEqual(false, cmds.recording); + try std.testing.expectEqual(@as(u32, 0), cmds.pending); + try std.testing.expectEqual(@as(usize, 1), Calls.resets); + try cmds.submit(0); + try std.testing.expectEqual(@as(usize, 0), Calls.submits); + + // A control's reset drops open records the same way; Control.run allocates nothing after its dispatch. + cmds.recording = true; + cmds.pending = 2; + var control = Control.init(failing.allocator(), &dev, &cmds, staging, &pipeline, .qmm_fmax); + try std.testing.expectError(error.OutOfMemory, control.a.dupe(u8, "x")); + control.reset(); + try std.testing.expectEqual(false, cmds.recording); + try std.testing.expectEqual(@as(u32, 0), cmds.pending); + try std.testing.expectEqual(@as(usize, 2), Calls.resets); + try cmds.submit(0); + try std.testing.expectEqual(@as(usize, 0), Calls.submits); +} + // Comptime constant for the same reason as Fixture.vt: the Kernel keeps a pointer to it. const control_vt = g.hooks(Control); diff --git a/zig/tests/vulkan/qwen35_model_probe.zig b/zig/tests/vulkan/qwen35_model_probe.zig new file mode 100644 index 00000000000..84c57e2f16f --- /dev/null +++ b/zig/tests/vulkan/qwen35_model_probe.zig @@ -0,0 +1,588 @@ +// Qwen3.5-2B: context 1-7168 (256 default); 2 exact-ID runs, lowest ties; exits 0/1/2/124; teacher-forced per-step. +const std = @import("std"); +const q35v = @import("qwen35_vulkan"); +const engine_mod = q35v.engine; +const pinned = q35v.config; +const q35s = @import("qwen35_shaders"); +const sys = @import("sys"); + +const receipt = @import("qwen35_model_receipt.zig"); +const Receipt = receipt.Receipt; +const writeResult = receipt.writeResult; +const free_running_mode = receipt.free_running_mode; +const teacher_forced_mode = receipt.teacher_forced_mode; + +const Engine = engine_mod.Engine; + +const usage = "usage: tf-qwen35-model [--vendor ID] [--timeout-ms MS] [--context N] [--teacher-forced] [--sequential-prefill] MODEL_DIR INPUT_IDS.npy ORACLE_DIR OUTPUT_DIR\n" ++ + " --help show this text and exit without opening a loader or device\n" ++ + " --vendor ID refuse a device whose Vulkan vendorID differs (decimal or 0x hex; no default)\n" ++ + " --timeout-ms MS bound on every fence wait, 1 to 300000 (default 10000); a timeout ends the process\n" ++ + " --context N KV context capacity to request, 1 to 7168 (default 256)\n" ++ + " --teacher-forced feed each scored step's oracle id as the next input (labelled per-step mode; default off)\n" ++ + " --sequential-prefill prefill one prompt token per Engine.prefill call (test-only; default is the canonical whole-prompt 32-row chunking)\n"; + +const Elapsed = struct { + io: std.Io, + start: std.Io.Timestamp, + + fn ms(e: Elapsed) i64 { + return e.start.durationTo(std.Io.Clock.awake.now(e.io)).toMilliseconds(); + } +}; + +const Options = struct { + vendor: ?u32 = null, + timeout_ms: u64 = 10_000, + context: usize = 256, + teacher_forced: bool = false, + sequential_prefill: bool = false, + positional: [4][]const u8 = undefined, +}; + +// The CLI's shape: flags then exactly four positional paths, or null; pure and host-tested. +pub fn parseArgs(args: []const []const u8) ?Options { + var o: Options = .{}; + var i: usize = 0; + while (i < args.len) { + const arg = args[i]; + if (std.mem.eql(u8, arg, "--vendor")) { + if (i + 1 >= args.len) return null; + o.vendor = std.fmt.parseInt(u32, args[i + 1], 0) catch return null; + i += 2; + } else if (std.mem.eql(u8, arg, "--timeout-ms")) { + if (i + 1 >= args.len) return null; + o.timeout_ms = std.fmt.parseInt(u64, args[i + 1], 10) catch return null; + if (o.timeout_ms == 0 or o.timeout_ms > 300_000) return null; + i += 2; + } else if (std.mem.eql(u8, arg, "--context")) { + if (i + 1 >= args.len) return null; + o.context = std.fmt.parseInt(usize, args[i + 1], 10) catch return null; + if (o.context == 0 or o.context > q35v.backend.max_capacity) return null; + i += 2; + } else if (std.mem.eql(u8, arg, "--teacher-forced")) { + o.teacher_forced = true; + i += 1; + } else if (std.mem.eql(u8, arg, "--sequential-prefill")) { + o.sequential_prefill = true; + i += 1; + } else if (std.mem.startsWith(u8, arg, "--")) { + return null; + } else break; + } + const rest = args[i..]; + if (rest.len != 4) return null; + o.positional = .{ rest[0], rest[1], rest[2], rest[3] }; + return o; +} + +// Canonical prefill chunking: how many 32-row chunks `rows` spans; pure, host-tested; the chunk size is the engine's. +pub fn chunkCount(rows: usize) usize { + return (rows + engine_mod.max_prefill_rows - 1) / engine_mod.max_prefill_rows; +} + +// The .npy header's verdict, or null unless a v1 little-endian u32 rank-1 array of `count` words; pure, host-tested. +pub const Npy = struct { count: usize, data_offset: usize }; + +pub fn parseNpy(bytes: []const u8) ?Npy { + if (bytes.len < 8 or bytes[6] != 1) return null; + const array = @import("core").npy.parse(bytes) catch return null; + if (array.rank != 1 or array.shape[0] == 0 or !std.mem.eql(u8, array.descr, " std.math.maxInt(u32)) return error.BadOracleIds; + out[i] = @intCast(item.integer); + } + return out; +} + +// Oracle result.json fields the receipt echoes; missing = unusable oracle, not a pass; arena-allocated, nothing freed. +const OracleMeta = struct { + prompt_token_hash: []const u8, + generated_ids_hash: []const u8, + logits_bf16_hash: []const u8, + logits_bytes: u64, + steps: u64, + note: []const u8, +}; + +fn readOracleMeta(a: std.mem.Allocator, io: std.Io, dir: std.Io.Dir, oracle_dir: []const u8) !OracleMeta { + const path = try std.fmt.allocPrint(a, "{s}/result.json", .{oracle_dir}); + defer a.free(path); + const bytes = try dir.readFileAlloc(io, path, a, .limited(1 << 22)); + defer a.free(bytes); + var parsed = try std.json.parseFromSlice(std.json.Value, a, bytes, .{}); + defer parsed.deinit(); + const root = if (parsed.value == .object) parsed.value.object else return error.BadOracleMeta; + const str = struct { + fn get(alloc: std.mem.Allocator, o: std.json.ObjectMap, key: []const u8) ![]const u8 { + const v = o.get(key) orelse return error.BadOracleMeta; + if (v != .string) return error.BadOracleMeta; + return alloc.dupe(u8, v.string); + } + }; + const num = struct { + fn get(o: std.json.ObjectMap, key: []const u8) !u64 { + const v = o.get(key) orelse return error.BadOracleMeta; + if (v != .integer or v.integer < 0) return error.BadOracleMeta; + return @intCast(v.integer); + } + }; + var note: []const u8 = ""; + if (root.get("reference_limits")) |rl| { + if (rl == .object) { + if (rl.object.get("checkpoint_revision")) |v| { + if (v == .string) note = try a.dupe(u8, v.string); + } + } + } + return .{ + .prompt_token_hash = try str.get(a, root, "prompt_token_hash"), + .generated_ids_hash = try str.get(a, root, "generated_ids_hash"), + .logits_bf16_hash = try str.get(a, root, "logits_bf16_hash"), + .logits_bytes = try num.get(root, "logits_bytes"), + .steps = try num.get(root, "steps"), + .note = note, + }; +} + +// The bundle: qwen35_shaders SPIR-V per ShaderBundle field, shared by probe and load test — the set the server serves. +fn bundle() engine_mod.ShaderBundle { + return .{ + .ops = .{ + .qmm_bf16 = &q35s.qmm_bf16, + .qmm_q4_vec4 = &q35s.qmm_q4_vec4, + .qmm_f32 = &q35s.qmm_f32, + .qmm_q4_tile = &q35s.qmm_q4_tile, + .qmm_q4_tile_f32 = &q35s.qmm_q4_tile_f32, + .dequant_bf16 = &q35s.dequant_bf16, + .gather_u32 = &q35s.gather_u32, + .gather_u16 = &q35s.gather_u16, + .cast_bf16_f32 = &q35s.cast_bf16_f32, + .cast_f32_bf16 = &q35s.cast_f32_bf16, + .copy_u32 = &q35s.copy_u32, + .copy_u16 = &q35s.copy_u16, + .copy_f32_bf16 = &q35s.copy_f32_bf16, + }, + .norm = .{ + .norm_bf16 = &q35s.norm_bf16, + .add_bf16 = &q35s.add_bf16, + .add_preserve_bf16 = &q35s.add_preserve_bf16, + .swiglu_bf16 = &q35s.swiglu_bf16, + .norm_f32 = &q35s.norm_f32, + .norm_f32_in = &q35s.norm_f32_in, + .norm_bf16_in = &q35s.norm_bf16_in, + .add_f32 = &q35s.add_f32, + .add_wide_bf16 = &q35s.add_wide_bf16, + .add_bf16_f32 = &q35s.add_bf16_f32, + .swiglu_f32 = &q35s.swiglu_f32, + }, + .attention = .{ + .rope_bf16 = &q35s.rope_bf16, + .sdpa_rows_bf16 = &q35s.sdpa_rows_bf16, + .gate_bf16 = &q35s.gate_bf16, + .rope_f32 = &q35s.rope_f32, + .rope_f32_kcache = &q35s.rope_f32_kcache, + .sdpa_rows_f32q = &q35s.sdpa_rows_f32q, + .gate_f32 = &q35s.gate_f32, + }, + .gdn = .{ + .conv_dw1d_bf16 = &q35s.conv_dw1d_bf16, + .gdn_conv_decode_bf16 = &q35s.gdn_conv_decode_bf16, + .gdn_split_bf16 = &q35s.gdn_split_bf16, + .prefill_recur_bf16 = &q35s.prefill_recur_bf16, + .gated_delta_decode_bf16 = &q35s.gated_delta_decode_bf16, + .fast_norm_gated_bf16 = &q35s.fast_norm_gated_bf16, + .copy_u16 = &q35s.copy_u16, + .conv_dw1d_f32 = &q35s.conv_dw1d_f32, + .gdn_split_f32 = &q35s.gdn_split_f32, + .prefill_recur_f32 = &q35s.prefill_recur_f32, + .fast_norm_gated_f32 = &q35s.fast_norm_gated_f32, + .copy_u32 = &q35s.gdn_copy_u32, + }, + .greedy = .{ .finite_logits = &q35s.finite_logits, .argreduce_chunk = &q35s.argreduce_chunk, .argreduce_combine = &q35s.argreduce_combine }, + }; +} + +pub fn main(init: std.process.Init) !u8 { + const clock: Elapsed = .{ .io = init.io, .start = std.Io.Clock.awake.now(init.io) }; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + for (args[1..]) |arg| if (std.mem.eql(u8, arg, "--help")) { + std.debug.print(usage, .{}); + return 0; + }; + const opts = parseArgs(args[1..]) orelse { + std.debug.print(usage, .{}); + return 2; + }; + const a = init.gpa; + const io = init.io; + const cwd = std.Io.Dir.cwd(); + // Immutable hash/diagnostic strings ride the process arena: allocated once, freed only at exit, never on gpa. + const diag = init.arena.allocator(); + + // every input validated before any GPU work exists + const npy_bytes = cwd.readFileAlloc(io, opts.positional[1], a, .limited(1 << 26)) catch |e| { + std.debug.print("FAIL input: cannot read {s}: {s}\n", .{ opts.positional[1], @errorName(e) }); + return 2; + }; + defer a.free(npy_bytes); + const npy = parseNpy(npy_bytes) orelse { + std.debug.print("FAIL input: {s} is not a v1 little-endian u32 rank-1 NPY with a nonempty shape\n", .{opts.positional[1]}); + return 2; + }; + if (npy.count > q35v.backend.max_capacity) { + std.debug.print("FAIL input: {d} prompt ids past the {d}-row context cap\n", .{ npy.count, q35v.backend.max_capacity }); + return 2; + } + const prompt = try readPromptIds(a, npy_bytes[npy.data_offset..]); + defer a.free(prompt); + for (prompt) |id| if (id >= pinned.vocab) { + std.debug.print("FAIL input: token id {d} past the vocab\n", .{id}); + return 2; + }; + + const expected = readExpectedIds(a, io, cwd, opts.positional[2]) catch |e| { + std.debug.print("FAIL input: oracle ids unreadable from {s}: {s}\n", .{ opts.positional[2], @errorName(e) }); + return 2; + }; + defer a.free(expected); + // immutable metadata rides the arena: readOracleMeta allocates, the process arena frees, no per-field cleanup + const oracle = readOracleMeta(diag, io, cwd, opts.positional[2]) catch |e| { + std.debug.print("FAIL input: oracle result.json unreadable from {s}: {s}\n", .{ opts.positional[2], @errorName(e) }); + return 2; + }; + if (oracle.steps != expected.len) { + std.debug.print("FAIL input: oracle steps {d} and ids {d} disagree\n", .{ oracle.steps, expected.len }); + return 2; + } + // same-prompt and expected-ids proof before any GPU work: the oracle's hashes must name this prompt and these ids + for (expected) |id| if (id >= pinned.vocab) { + std.debug.print("FAIL input: oracle id {d} past the vocab\n", .{id}); + return 2; + }; + const prompt_hash = try sha256HexA(diag, std.mem.sliceAsBytes(prompt)); + const expected_hash = try sha256HexA(diag, std.mem.sliceAsBytes(expected)); + if (!sameHex(prompt_hash, oracle.prompt_token_hash)) { + std.debug.print("FAIL input: prompt hash {s} is not the oracle's {s}\n", .{ prompt_hash, oracle.prompt_token_hash }); + return 2; + } + if (!sameHex(expected_hash, oracle.generated_ids_hash)) { + std.debug.print("FAIL input: expected ids hash {s} is not the oracle's {s}\n", .{ expected_hash, oracle.generated_ids_hash }); + return 2; + } + if (!isShaHex(oracle.logits_bf16_hash)) { + std.debug.print("FAIL input: oracle logits hash is not 64 hex characters\n", .{}); + return 2; + } + + // the prompt and every step must fit the requested context — checked before the device or weights exist + engine_mod.planFor(opts.context, 0, prompt.len, expected.len) catch |e| { + std.debug.print("FAIL plan: {s} (prompt {d} rows, {d} steps, context {d})\n", .{ @errorName(e), prompt.len, expected.len, opts.context }); + return 1; + }; + + // the output directory must be new before the GPU exists + cwd.createDir(io, opts.positional[3], .default_dir) catch |e| switch (e) { + error.PathAlreadyExists => { + std.debug.print("FAIL input: output dir {s} already exists\n", .{opts.positional[3]}); + return 2; + }, + else => { + std.debug.print("FAIL input: cannot create {s}: {s}\n", .{ opts.positional[3], @errorName(e) }); + return 2; + }, + }; + const out = cwd.openDir(io, opts.positional[3], .{}) catch |e| { + std.debug.print("FAIL input: cannot open {s}: {s}\n", .{ opts.positional[3], @errorName(e) }); + return 2; + }; + defer out.close(io); + + const engine = Engine.init(a, io, bundle(), opts.positional[0], opts.context, opts.timeout_ms * std.time.ns_per_ms) catch |e| { + std.debug.print("FAIL engine: {s}\n", .{@errorName(e)}); + return 1; + }; + defer engine.deinit(); + std.debug.print("MILESTONE engine-ready elapsed_ms={d} capacity={d} model_context={d} logits_bytes={d}\n", .{ clock.ms(), engine.capacity(), engine.modelContext(), engine.logitsLen() }); + if (opts.vendor) |want| if (engine.vendorId() != want) { + std.debug.print("FAIL device: vendor 0x{x:0>4} is not the requested 0x{x:0>4}\n", .{ engine.vendorId(), want }); + return 1; + }; + + const row_len: usize = @intCast(engine.logitsLen()); + const rows_a = try a.alloc(u8, expected.len * row_len); + defer a.free(rows_a); + const rows_b = try a.alloc(u8, expected.len * row_len); + defer a.free(rows_b); + const ids_a = try a.alloc(u32, expected.len); + defer a.free(ids_a); + const ids_b = try a.alloc(u32, expected.len); + defer a.free(ids_b); + + greedyRun(engine, opts.teacher_forced, opts.sequential_prefill, prompt, expected, ids_a, rows_a, row_len) catch |e| { + std.debug.print("FAIL run1: {s}\n", .{@errorName(e)}); + return 1; + }; + std.debug.print("MILESTONE run1-done elapsed_ms={d}\n", .{clock.ms()}); + try engine.reset(); + greedyRun(engine, opts.teacher_forced, opts.sequential_prefill, prompt, expected, ids_b, rows_b, row_len) catch |e| { + std.debug.print("FAIL run2: {s}\n", .{@errorName(e)}); + return 1; + }; + std.debug.print("MILESTONE run2-done elapsed_ms={d}\n", .{clock.ms()}); + + // within-Vulkan determinism: fresh cache, real reset, identical bytes + const ids_equal = std.mem.eql(u32, ids_a, ids_b); + const logits_equal = std.mem.eql(u8, rows_a, rows_b); + // the pass criterion: the exact oracle ids, both runs, no epsilon + const match_a = std.mem.eql(u32, ids_a, expected); + const match_b = std.mem.eql(u32, ids_b, expected); + + // diagnostic only: element-wise raw logits against the oracle's file + var mismatches: ?u64 = null; + const oracle_logits_path = try std.fmt.allocPrint(a, "{s}/logits_bf16.bin", .{opts.positional[2]}); + defer a.free(oracle_logits_path); + if (cwd.readFileAlloc(io, oracle_logits_path, a, .limited(rows_a.len + (1 << 20)))) |oracle_bytes| { + defer a.free(oracle_bytes); + mismatches = countElementMismatches(oracle_bytes, rows_a); + if (mismatches == null) std.debug.print("NOTE oracle logits {d} bytes vs run {d}; diagnostic skipped\n", .{ oracle_bytes.len, rows_a.len }); + } else |_| std.debug.print("NOTE oracle logits file unreadable; diagnostic skipped\n", .{}); + + // receipts: the full BF16 logit rows of run A, then the result record + out.writeFile(io, .{ .sub_path = "logits_bf16.bin", .data = rows_a }) catch |e| { + std.debug.print("FAIL receipt: cannot write logits_bf16.bin: {s}\n", .{@errorName(e)}); + return 1; + }; + const ids_hash = try sha256HexA(diag, std.mem.sliceAsBytes(ids_a)); + writeResult(a, io, out, .{ + .status = if (match_a and match_b and ids_equal and logits_equal) "ok" else "mismatch", + .steps = expected.len, + .execution_mode = if (opts.teacher_forced) teacher_forced_mode else free_running_mode, + .prefill_schedule = if (opts.sequential_prefill) "rowwise" else "chunk32", + .context_cap = opts.context, + .prompt_chunks = if (opts.sequential_prefill) prompt.len else chunkCount(prompt.len), + .prefill_chunk_rows = if (opts.sequential_prefill) 1 else engine_mod.max_prefill_rows, + .capacity = engine.capacity(), + .model_context = engine.modelContext(), + .prompt_tokens = prompt.len, + .input_npy_sha256 = try sha256HexA(diag, npy_bytes), + .prompt_u32le_sha256 = prompt_hash, + .generated_ids = ids_a, + .generated_ids_hash = ids_hash, + .logits_bytes = rows_a.len, + .logits_bf16_sha256 = try sha256HexA(diag, rows_a), + .repeat_checks = .{ .generated_ids_equal = ids_equal, .logits_bytes_equal = logits_equal }, + .oracle_match = .{ .ids_run_a = match_a, .ids_run_b = match_b }, + .oracle = .{ + .generated_ids_hash = oracle.generated_ids_hash, + .logits_bf16_hash = oracle.logits_bf16_hash, + .logits_bytes = oracle.logits_bytes, + .checkpoint_note = oracle.note, + }, + .oracle_logits_element_mismatches = mismatches, + .elapsed_ms = clock.ms(), + }) catch |e| { + std.debug.print("FAIL receipt: cannot write result.json: {s}\n", .{@errorName(e)}); + return 1; + }; + + std.debug.print("RESULT qwen35-model elapsed_ms={d}: ids_match_a={} ids_match_b={} repeats(ids={},logits={}) " ++ + "oracle_logits_mismatches={s} expected_hash={s} got_hash={s}\n", .{ + clock.ms(), match_a, match_b, ids_equal, logits_equal, + if (mismatches) |n| (std.fmt.allocPrint(diag, "{d}", .{n}) catch "?") else "skipped", oracle.generated_ids_hash, ids_hash, + }); + // mode labels ride only non-default modes, so the anchor default's RESULT line stays byte-identical + if (opts.teacher_forced) std.debug.print("RESULT execution_mode=teacher-forced: oracle ids fed as previous tokens; per-step agreement, not a free-running chain proof\n", .{}); + if (opts.sequential_prefill) std.debug.print("RESULT prefill_schedule=rowwise: one Engine.prefill call per prompt token\n", .{}); + return if (match_a and match_b and ids_equal and logits_equal) 0 else 1; +} + +// Greedy: prefill (32-row chunks; --sequential-prefill per token = boundary evidence); ids_out.len argreduce draws. +fn greedyRun(e: *Engine, teacher_forced: bool, sequential: bool, prompt: []const u32, expected: []const u32, ids_out: []u32, rows_out: []u8, row_len: usize) !void { + if (sequential) { + for (prompt) |tok| try e.prefill(&.{tok}); + } else { + try e.prefill(prompt); + } + for (ids_out, 0..) |*slot, step| { + slot.* = try e.draw(); + try e.logitsBytes(rows_out[step * row_len ..][0..row_len]); + if (step + 1 < ids_out.len) try e.decode(if (teacher_forced) expected[step] else slot.*); + } +} + +const testing = std.testing; + +test "isShaHex and sameHex gate the oracle hash cross-checks" { + const h = "dcd774d847e199c7a8ad93ef794e597e5a6bf672ec448afb7a2e507578b261b3"; + try testing.expect(isShaHex(h)); + try testing.expect(sameHex(h, h)); + try testing.expect(sameHex(h, "DCD774D847E199C7A8AD93EF794E597E5A6BF672EC448AFB7A2E507578B261B3")); + try testing.expect(!isShaHex(h[0..63])); // short + try testing.expect(!isShaHex(&@as([64]u8, @splat('g')))); // not hex + try testing.expect(!isShaHex("dcd774")); // far too short + try testing.expect(!sameHex(h, "d1b45ec858151fe74125ce020f637fcf1963b399a3cd32564c7ca65966624589")); // different digest +} + +test "countElementMismatches counts 2-byte pairs, not bytes" { + const a = [_]u8{ 0x11, 0x22, 0x33, 0x44, 0x55, 0x66 }; + try testing.expectEqual(@as(?u64, 0), countElementMismatches(&a, &a)); + var b = a; + b[1] ^= 0xff; // one byte of the first element flips: one element differs + try testing.expectEqual(@as(?u64, 1), countElementMismatches(&a, &b)); + var c = a; + c[2] ^= 0xff; // first element intact again, second element flipped + try testing.expectEqual(@as(?u64, 1), countElementMismatches(&a, &c)); + try testing.expectEqual(@as(?u64, null), countElementMismatches(&a, a[0..4])); // length mismatch + try testing.expectEqual(@as(?u64, null), countElementMismatches(&a, a[0..5])); // torn element +} + +test "parseArgs takes the two flags and exactly four paths" { + const ok = parseArgs(&.{ "--vendor", "0x1002", "--timeout-ms", "90000", "/m", "/i.npy", "/o", "/out" }).?; + try testing.expectEqual(@as(?u32, 0x1002), ok.vendor); + try testing.expectEqual(@as(u64, 90000), ok.timeout_ms); + try testing.expectEqualStrings("/out", ok.positional[3]); + try testing.expect(parseArgs(&.{"/m"}) == null); + try testing.expect(parseArgs(&.{"--vendor"}) == null); + try testing.expect(parseArgs(&.{ "--vendor", "x", "/m", "/i", "/o", "/out" }) == null); + try testing.expect(parseArgs(&.{ "--timeout-ms", "0", "/m", "/i", "/o", "/out" }) == null); + try testing.expect(parseArgs(&.{ "--wat", "/m", "/i", "/o", "/out" }) == null); + try testing.expect(parseArgs(&.{ "/m", "/i", "/o", "/out", "/extra" }) == null); +} + +test "parseArgs context and mode flags keep the anchor default and bound the context" { + // the anchor default: no new flags, the historical 256 free-running run + const plain = parseArgs(&.{ "/m", "/i", "/o", "/out" }).?; + try testing.expectEqual(@as(usize, 256), plain.context); + try testing.expect(!plain.teacher_forced); + try testing.expect(!plain.sequential_prefill); + + // the boundary sizes the corpus runs: 33, 65, 7105 all parse + for ([_][]const u8{ "33", "65", "7105", "7168" }) |n| { + const o = parseArgs(&.{ "--context", n, "--teacher-forced", "--sequential-prefill", "/m", "/i", "/o", "/out" }).?; + try testing.expect(o.teacher_forced); + try testing.expect(o.sequential_prefill); + } + // one past the backend's capacity is refused at the CLI, no device + try testing.expectEqual(@as(usize, 7168), q35v.backend.max_capacity); + try testing.expect(parseArgs(&.{ "--context", "7169", "/m", "/i", "/o", "/out" }) == null); + try testing.expect(parseArgs(&.{ "--context", "0", "/m", "/i", "/o", "/out" }) == null); + try testing.expect(parseArgs(&.{ "--context", "-1", "/m", "/i", "/o", "/out" }) == null); + try testing.expect(parseArgs(&.{ "--context", "abc", "/m", "/i", "/o", "/out" }) == null); + try testing.expect(parseArgs(&.{"--context"}) == null); + try testing.expect(parseArgs(&.{ "--context", "256", "/m", "/i", "/o", "/out", "/extra" }) == null); +} + +test "chunkCount spans the canonical 32-row chunks without a configurable size" { + try testing.expectEqual(@as(usize, engine_mod.max_prefill_rows), 32); + try testing.expectEqual(@as(usize, 1), chunkCount(1)); + try testing.expectEqual(@as(usize, 1), chunkCount(32)); + try testing.expectEqual(@as(usize, 2), chunkCount(33)); // boundary: two chunks + try testing.expectEqual(@as(usize, 2), chunkCount(64)); + try testing.expectEqual(@as(usize, 3), chunkCount(65)); // boundary: three chunks + try testing.expectEqual(@as(usize, 224), chunkCount(7168)); // full context +} + +test "parseNpy takes exactly a v1 LE u32 rank-1 array" { + const dict = "{'descr': ' try q.fill(.random, seed), + .cancel => try cancel(q, seed), + } + const x = try poisoned16(a, at.x, q.x.len); + errdefer a.free(x); + @memcpy(x[at.x..][0..q.x.len], q.x); + const w = try poisoned32(a, at.w, s.n * s.k / 8); + errdefer a.free(w); + q.words(w[at.w..][0 .. s.n * s.k / 8]); + const scales = try poisoned16(a, at.scales, q.scales.len); + errdefer a.free(scales); + @memcpy(scales[at.scales..][0..q.scales.len], q.scales); + const biases = try poisoned16(a, at.biases, q.biases.len); + errdefer a.free(biases); + @memcpy(biases[at.biases..][0..q.biases.len], q.biases); + const out = try poisoned16(a, at.out, s.rows * s.n); + errdefer a.free(out); + const want = try a.alloc(f64, s.rows * s.n); + errdefer a.free(want); + const mag = try a.alloc(f64, want.len); + errdefer a.free(mag); + try q.reference(want, mag); + for (want, mag) |expected, magnitude| if (!std.math.isFinite(expected) or !std.math.isFinite(magnitude) or magnitude < 0) return error.NonFiniteReference; + return .{ .a = a, .case = case, .q = q, .x = x, .w = w, .scales = scales, .biases = biases, .out = out, .want = want, .mag = mag }; + } + + pub fn deinit(f: Fixture) void { + f.q.deinit(); + f.a.free(f.x); + f.a.free(f.w); + f.a.free(f.scales); + f.a.free(f.biases); + f.a.free(f.out); + f.a.free(f.want); + f.a.free(f.mag); + } + + /// got: bf16 bits read back whole; guards keep their bits and every output meets its bound. + pub fn check(f: Fixture, got: []const u16) !Report { + if (got.len != f.out.len) return error.LengthMismatch; + const at = f.case.at.out; + for (got, f.out, 0..) |g, o, i| { + if (i >= at and i < at + f.want.len) continue; + if (g != o) return error.GuardClobbered; + } + const r = f.measure(got); + if (!(r.worst <= 1)) return error.OutOfBound; + return r; + } + + /// Each output's distance from f64 over its bound; call after check fails to name the worst output. + pub fn measure(f: Fixture, got: []const u16) Report { + std.debug.assert(got.len == f.out.len); + var r: Report = .{}; + for (got[f.case.at.out..][0..f.want.len], f.want, f.mag, 0..) |bits, w, m, i| { + const g: f64 = c.fromBf16(bits); + const d = @abs(g - w); + const e = if (std.math.isNan(d)) std.math.inf(f64) else d; + const bound = c.bound(spec, f.case.shape, w, m); + const valid = std.math.isFinite(m) and m >= 0 and std.math.isFinite(bound) and bound > 0; + const ratio = if (valid) e / bound else std.math.inf(f64); + r.max_abs = @max(r.max_abs, e); + if (ratio > r.worst) { + r.worst = ratio; + r.row = i / f.case.shape.n; + r.col = i % f.case.shape.n; + } + } + return r; + } +}; + +fn poisoned16(a: std.mem.Allocator, base: usize, len: usize) ![]u16 { + const buf = try a.alloc(u16, base + len + tail); + @memset(buf, poison16); + return buf; +} + +fn poisoned32(a: std.mem.Allocator, base: usize, len: usize) ![]u32 { + const buf = try a.alloc(u32, base + len + tail); + @memset(buf, poison32); + return buf; +} + +/// Bias -8 s (exact) and codes mirrored about 8 across a group's halves; x repeats across them but once a group. +fn cancel(q: Affine, seed: u64) !void { + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const group = q.group; + for (q.scales, q.biases, 0..) |*s, *b, i| { + s.* = c.toBf16((rng.float(f32) * 0.02 + 0.001) * @as(f32, if (i % 2 == 1) -1 else 1)); + b.* = c.toBf16(-8 * c.fromBf16(s.*)); + } + const h = group / 2; + for (0..q.codes.len / group) |g| for (0..h) |i| { + const v = 1 + rng.uintLessThan(u8, 15); + q.codes[g * group + i] = v; + q.codes[g * group + h + i] = 16 - v; + }; + for (0..q.x.len / group) |g| for (0..h) |i| { + const v = c.toBf16(rng.float(f32) * 4 - 2); + q.x[g * group + i] = v; + q.x[g * group + h + i] = if (i == g % h) c.toBf16(rng.float(f32) * 4 - 2) else v; + }; +} + +const testing = std.testing; + +test "bf16 vectors agree with the contract's RNE cast at every stated geometry" { + for (bf16_vectors) |v| try testing.expectEqual(v.want, c.toBf16(v.in)); + try testing.expectEqual(@as(u16, 0x3f81), c.toBf16(1.0 + 0x1p-8 + 0x1p-20)); // a hair above half +} + +test "projection fixture bindings meet the largest Vulkan storage alignment" { + for (cases) |case| { + for ([_]usize{ case.at.x * 2, case.at.w * 4, case.at.scales * 2, case.at.biases * 2, case.at.out * 2 }) |bytes| { + try testing.expectEqual(@as(usize, 0), bytes % 256); + } + } +} + +test "both cast widths check the actual tail after a whole output readback" { + const elements = bf16_vectors.len + 128; + inline for (.{ 1, 2 }) |words| { + var shell: [(elements + tail) * words]u16 = undefined; + @memset(shell[0 .. elements * words], 0x3f80); + @memset(shell[elements * words ..], 0xa5a5); + try castGuards(&shell, elements, words); + try testing.expectError(error.LengthMismatch, castGuards(shell[0 .. elements * words], elements, words)); + try testing.expectError(error.LengthMismatch, castGuards(shell[0 .. tail * words], elements, words)); + shell[elements * words] ^= 1; + try testing.expectError(error.GuardClobbered, castGuards(&shell, elements, words)); + } +} + +test "embedding tables choose valid, once-repeated ids and finite f64 wants over poison tails" { + var t = try EmbedTable.init(testing.allocator, embed_cases[1], 7); + defer t.deinit(); + try testing.expect(t.ids[1] == t.ids[0]); + for (t.ids) |id| try testing.expect(id < t.case.rows); + const words = t.case.rows * t.case.k / 8; + for (t.w[words..]) |g| try testing.expectEqual(poison32, g); + var nonzero: usize = 0; + for (t.want) |w| { + try testing.expect(std.math.isFinite(w)); + nonzero += @intFromBool(w != 0); + } + try testing.expect(nonzero > t.want.len / 2); +} + +test "q4 fixture: the checker refuses a planted nonfinite output and a clobbered guard" { + var f = try Fixture.init(testing.allocator, cases[0], 3); + defer f.deinit(); + const got = try testing.allocator.dupe(u16, f.out); + defer testing.allocator.free(got); + for (f.want, 0..) |w, i| got[f.case.at.out + i] = c.toBf16(@floatCast(w)); + _ = try f.check(got); + got[f.case.at.out] = 0x7f80; // +inf at output 0 + try testing.expectError(error.OutOfBound, f.check(got)); + got[f.case.at.out] = c.toBf16(@floatCast(f.want[0])); + got[f.case.at.out + f.want.len] ^= 1; + try testing.expectError(error.GuardClobbered, f.check(got)); +} + +test "q4 fixture: reference-rounded outputs meet the contract bound, including the cancel fill" { + for (cases) |case| { + var f = try Fixture.init(testing.allocator, case, 11); + defer f.deinit(); + const got = try testing.allocator.dupe(u16, f.out); + defer testing.allocator.free(got); + for (f.want, 0..) |w, i| got[f.case.at.out + i] = c.toBf16(@floatCast(w)); + const r = try f.check(got); + try testing.expect(r.worst <= 1); + } +} diff --git a/zig/tests/vulkan/stage1_probe.zig b/zig/tests/vulkan/stage1_probe.zig new file mode 100644 index 00000000000..2ebdf494a23 --- /dev/null +++ b/zig/tests/vulkan/stage1_probe.zig @@ -0,0 +1,531 @@ +// Stage 1 GPU checks typed views, arena, BF16 cast/copy and Q4 embed/affine; no model claim. Exit 0/1/2/124. +const std = @import("std"); +const vk = @import("vulkan"); +const shaders = @import("stage1_shaders"); +const fx = @import("stage1_fixture.zig"); +const c = vk.c; +const contract = @import("contract"); +const cx = contract.spec; // toBf16, fromBf16, boundExact, the C constants + +const usage = "usage: tf-stage1-probe [--vendor ID] [--timeout-ms MS] | --help\n" ++ + " --help show this text and exit without opening a loader or device\n" ++ + " --vendor ID refuse a device whose Vulkan vendorID differs (decimal or 0x hex; no default)\n" ++ + " --timeout-ms MS bound on every fence wait, 1 to 600000 (default 10000); a timeout ends the process\n"; + +/// Milliseconds on the awake clock since main started, printed on every milestone. +const Elapsed = struct { + io: std.Io, + start: std.Io.Timestamp, + + fn ms(e: Elapsed) i64 { + return e.start.durationTo(std.Io.Clock.awake.now(e.io)).toMilliseconds(); + } +}; + +const repeats = 3; + +pub fn main(init: std.process.Init) !u8 { + const clock: Elapsed = .{ .io = init.io, .start = std.Io.Clock.awake.now(init.io) }; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + for (args[1..]) |arg| if (std.mem.eql(u8, arg, "--help")) { + std.debug.print(usage, .{}); + return 0; + }; + var vendor: ?u32 = null; + var timeout_ms: u64 = 10_000; + var i: usize = 1; + while (i < args.len) : (i += 2) { + const value = if (i + 1 < args.len) args[i + 1] else ""; + if (std.mem.eql(u8, args[i], "--vendor")) { + vendor = std.fmt.parseInt(u32, value, 0) catch null; + if (vendor == null) break; + } else if (std.mem.eql(u8, args[i], "--timeout-ms")) { + timeout_ms = std.fmt.parseInt(u64, value, 10) catch 0; + if (timeout_ms == 0 or timeout_ms > 600_000) break; + } else break; + } else return if (try probe(init.gpa, clock, vendor, timeout_ms * std.time.ns_per_ms)) 0 else 1; + std.debug.print(usage, .{}); + return 2; +} + +fn probe(a: std.mem.Allocator, clock: Elapsed, vendor: ?u32, timeout: u64) !bool { + var rt = try vk.Runtime.open(); + defer rt.close(); + var dev = try vk.Device.open(&rt, .{ .features = .{ .storage16 = true, .int16 = true } }); + defer dev.close(); + const pr = dev.properties; + const name = std.mem.sliceTo(&pr.deviceName, 0); + const types = [_][]const u8{ "other", "integrated", "discrete", "virtual", "cpu" }; + const device_kind = if (pr.deviceType < types.len) types[pr.deviceType] else "unknown"; + const v = pr.apiVersion; + const driver = .{ std.mem.sliceTo(&dev.v12.driverName, 0), std.mem.sliceTo(&dev.v12.driverInfo, 0) }; + std.debug.print("MILESTONE device-found elapsed_ms={d} name=\"{s}\" vendor=0x{x:0>4} device=0x{x:0>4} type={s} " ++ + "api={d}.{d}.{d} driver=\"{s}\" info=\"{s}\" subgroup={d} storage16={} int16={} " ++ + "max-storage-range={d} max-buffer-size={d} fence-timeout={d}ms\n", .{ + clock.ms(), + name, + pr.vendorID, + pr.deviceID, + device_kind, + v >> 22 & 0x7f, + v >> 12 & 0x3ff, + v & 0xfff, + } ++ driver ++ .{ dev.v11.subgroupSize, dev.features.storage16, dev.features.int16, dev.properties.limits.maxStorageBufferRange, dev.v13.maxBufferSize, timeout / std.time.ns_per_ms }); + if (pr.deviceType == c.VK_PHYSICAL_DEVICE_TYPE_CPU) { + std.debug.print("FAIL device: a CPU (software) Vulkan device is never a GPU result\n", .{}); + return false; + } + if (vendor) |want| if (pr.vendorID != want) { + std.debug.print("FAIL device: vendor 0x{x:0>4} is not the requested 0x{x:0>4}\n", .{ pr.vendorID, want }); + return false; + }; + + var cmds = try vk.Commands.init(&dev); + defer cmds.deinit(); + var staging = try vk.Buffer.init(&dev, .staging, 64 << 10); + defer staging.deinit(); + const sh = vk.ops.Shaders{ + .qmm_bf16 = &shaders.qmm_bf16, + .qmm_q4_vec4 = &shaders.qmm_q4_vec4, + .qmm_f32 = &shaders.qmm_f32, + .copy_f32_bf16 = &shaders.copy_f32_bf16, + .qmm_q4_tile = &shaders.qmm_q4_tile, + .qmm_q4_tile_f32 = &shaders.qmm_q4_tile_f32, + .dequant_bf16 = &shaders.dequant_bf16, + .gather_u32 = &shaders.gather_u32, + .gather_u16 = &shaders.gather_u16, + .cast_bf16_f32 = &shaders.cast_bf16_f32, + .cast_f32_bf16 = &shaders.cast_f32_bf16, + .copy_u32 = &shaders.copy_u32, + .copy_u16 = &shaders.copy_u16, + }; + var o = try vk.ops.Ops.init(&dev, sh); + defer o.deinit(); + var scratch = try vk.arena.Arena.init(&dev, 1 << 20); + defer scratch.deinit(); + + var passed: usize = 0; + var total: usize = 0; + inline for (.{ "cast-f32-bf16", "cast-bf16-f32" }) |gate| { + total += 1; + passed += @intFromBool(try castGate(gate, a, clock, &dev, &cmds, staging, &o, timeout)); + } + inline for (.{ vk.tensor.Kind.u32, .u16 }) |kind| { + total += 1; + passed += @intFromBool(try copyGate(kind, a, clock, &dev, &cmds, staging, &o, timeout)); + } + inline for (.{ vk.tensor.Kind.u32, .bf16 }) |kind| { + total += 1; + passed += @intFromBool(try gatherGate(kind, a, clock, &dev, &cmds, staging, &o, timeout)); + } + total += 1; + passed += @intFromBool(try checkIdsGate(clock)); + for (fx.embed_cases) |case| { + total += 1; + passed += @intFromBool(try embedGate(a, clock, &dev, &cmds, staging, &o, &scratch, case, timeout)); + } + for (fx.cases, 0..) |case, n| { + total += 1; + passed += @intFromBool(try projectGate(a, clock, &dev, &cmds, staging, &o, &scratch, case, 0x57a6 + n, timeout)); + } + std.debug.print("RESULT stage1 elapsed_ms={d} on \"{s}\": passed={d} total={d}\n", .{ clock.ms(), name, passed, total }); + return passed == total; +} + +fn uploaded(_: std.mem.Allocator, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, bytes: []const u8, timeout: u64) !vk.Buffer { + var b = try vk.Buffer.init(dev, .device, bytes.len); + errdefer b.deinit(); + try vk.orExit(vk.upload(cmds, b, 0, staging, bytes, timeout)); + return b; +} + +fn fetch(cmds: *vk.Commands, staging: vk.Buffer, b: vk.Buffer, out: []u8, timeout: u64) !void { + try vk.orExit(vk.download(cmds, b, 0, staging, out, timeout)); +} + +/// Every stage 1 vector plus randoms: exact bf16 bits out, exact f32 back, guard tails untouched. +fn castGate(comptime gate: []const u8, a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, o: *vk.ops.Ops, timeout: u64) !bool { + const to_bf16 = std.mem.eql(u8, gate, "cast-f32-bf16"); + var prng = std.Random.DefaultPrng.init(0x57a1); + const rng = prng.random(); + const randoms = 128; + const n = fx.bf16_vectors.len + randoms; + var ins: std.ArrayList(u8) = .empty; + defer ins.deinit(a); + const wants = try a.alloc(u16, n); // the bf16 bits each element must produce or come from + defer a.free(wants); + for (0..n) |i| { + const x: f32 = if (i < fx.bf16_vectors.len) fx.bf16_vectors[i].in else (rng.float(f32) * 2 - 1) * std.math.pow(f32, 2, @floatFromInt(@as(i32, @intCast(i % 41)) - 20)); + wants[i] = if (i < fx.bf16_vectors.len) fx.bf16_vectors[i].want else cx.toBf16(x); + if (to_bf16) { + try ins.appendSlice(a, &std.mem.toBytes(x)); + } else { + try ins.appendSlice(a, &std.mem.toBytes(wants[i])); + } + } + var in = try uploaded(a, dev, cmds, staging, ins.items, timeout); + defer in.deinit(); + const out_esz: usize = if (to_bf16) 2 else 4; + var out_buf = try vk.Buffer.init(dev, .device, (n + fx.tail) * out_esz); + defer out_buf.deinit(); + const fill = try a.alloc(u8, (n + fx.tail) * out_esz); + defer a.free(fill); + @memset(fill, 0xa5); + try vk.orExit(vk.upload(cmds, out_buf, 0, staging, fill, timeout)); + const in_kind: vk.tensor.Kind = if (to_bf16) .f32 else .bf16; + const out_kind: vk.tensor.Kind = if (to_bf16) .bf16 else .f32; + const in_view = try vk.tensor.View.init(&in, in_kind, 0, n); + const out_view = try vk.tensor.View.init(&out_buf, out_kind, 0, n); + if (to_bf16) { + try vk.orExit(o.castF32ToBf16(cmds, out_view, in_view, timeout)); + } else { + try vk.orExit(o.castBf16ToF32(cmds, out_view, in_view, timeout)); + } + const words_out: usize = if (to_bf16) 1 else 2; // u16 words per output element + const back = try a.alloc(u16, (n + fx.tail) * words_out); + defer a.free(back); + try fetch(cmds, staging, out_buf, std.mem.sliceAsBytes(back), timeout); + for (0..n) |i| { + const want_bits: u32 = @as(u32, wants[i]) << (if (to_bf16) @as(u5, 0) else 16); + const got_bits: u32 = if (to_bf16) back[i] else @as(u32, back[2 * i]) | @as(u32, back[2 * i + 1]) << 16; + if (got_bits != want_bits) { + std.debug.print("FAIL {s}: element {d} read 0x{x:0>8}, want 0x{x:0>8}\n", .{ gate, i, got_bits, want_bits }); + return false; + } + } + fx.castGuards(back, n, words_out) catch { + std.debug.print("FAIL {s}: guard tail clobbered\n", .{gate}); + return false; + }; + std.debug.print("MILESTONE {s} elapsed_ms={d} elements={d} vectors={d} randoms={d} exact bits, tails unchanged\n", .{ gate, clock.ms(), n, fx.bf16_vectors.len, randoms }); + return true; +} + +/// A contiguous same-kind copy at element offsets: exact bytes, both guard tails and the source unchanged. +fn copyGate(comptime kind: vk.tensor.Kind, a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, o: *vk.ops.Ops, timeout: u64) !bool { + const gate = if (kind == .u32) "copy-u32" else "copy-u16"; + const elems: usize = if (kind == .u32) fx.copy_words else fx.copy_half_elements; + const src_at: usize = if (kind == .u32) fx.copy_u32_src else fx.copy_u16_src; + const dst_at: usize = if (kind == .u32) fx.copy_u32_dst else fx.copy_u16_dst; + const esz = kind.size(); + var prng = std.Random.DefaultPrng.init(0x57a2 + esz); + const rng = prng.random(); + const data = try a.alloc(u8, elems * esz); + defer a.free(data); + rng.bytes(data); + const poison: u8 = 0xa5; + var src = try poisonedBuf(a, dev, cmds, staging, src_at * esz + data.len, poison, timeout); + defer src.buf.deinit(); + try vk.orExit(vk.upload(cmds, src.buf, src_at * esz, staging, data, timeout)); + var dst = try poisonedBuf(a, dev, cmds, staging, dst_at * esz + data.len, poison, timeout); + defer dst.buf.deinit(); + const src_view = try vk.tensor.View.init(&src.buf, kind, src_at * esz, elems); + const dst_view = try vk.tensor.View.init(&dst.buf, kind, dst_at * esz, elems); + try vk.orExit(o.copy(cmds, dst_view, src_view, timeout)); + const back = try a.alloc(u8, dst.buf.len); + defer a.free(back); + try fetch(cmds, staging, dst.buf, back, timeout); + if (std.mem.indexOfDiff(u8, back[dst_at * esz ..][0..data.len], data)) |at| { + std.debug.print("FAIL {s}: first differing byte at element offset {d}\n", .{ gate, at / esz }); + return false; + } + for (back, 0..) |byte, i| { + if (i >= dst_at * esz and i < dst_at * esz + data.len) continue; + if (byte != poison) { + std.debug.print("FAIL {s}: guard clobbered at byte {d}\n", .{ gate, i }); + return false; + } + } + const src_back = try a.alloc(u8, src.buf.len); + defer a.free(src_back); + try fetch(cmds, staging, src.buf, src_back, timeout); + if (std.mem.indexOfDiff(u8, src_back[src_at * esz ..][0..data.len], data) != null or + std.mem.indexOfNone(u8, src_back[0 .. src_at * esz], &.{poison}) != null or + std.mem.indexOfNone(u8, src_back[src_at * esz + data.len ..], &.{poison}) != null) + { + std.debug.print("FAIL {s}: the source view changed\n", .{gate}); + return false; + } + std.debug.print("MILESTONE {s} elapsed_ms={d} elements={d} src_off={d} dst_off={d} exact bytes, guards and source unchanged\n", .{ gate, clock.ms(), elems, src_at, dst_at }); + return true; +} + +const Poisoned = struct { buf: vk.Buffer }; + +fn poisonedBuf(a: std.mem.Allocator, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, len: usize, poison: u8, timeout: u64) !Poisoned { + var b = try vk.Buffer.init(dev, .device, len); + errdefer b.deinit(); + const fill = try a.alloc(u8, len); + defer a.free(fill); + @memset(fill, poison); + try vk.orExit(vk.upload(cmds, b, 0, staging, fill, timeout)); + return .{ .buf = b }; +} + +/// Selected and repeated valid ids gather exact table rows; an out-of-range id gathers to zero rows. +fn gatherGate(comptime kind: vk.tensor.Kind, a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, o: *vk.ops.Ops, timeout: u64) !bool { + const gate = if (kind == .u32) "gather-u32" else "gather-bf16"; + const esz = kind.size(); + const rows: u32 = 6; + const cols: usize = 10; + const tokens: usize = 4; + const ids = [_]u32{ 0, 5, 5, 2 }; // selected, with a repeat + var prng = std.Random.DefaultPrng.init(0x57a3 + esz); + const rng = prng.random(); + const cell = try a.alloc(u8, cols * esz); + defer a.free(cell); + const table = try a.alloc(u8, rows * cols * esz); + defer a.free(table); + for (0..rows) |r| { + rng.bytes(cell); + @memcpy(table[r * cols * esz ..][0 .. cols * esz], cell); + } + var table_buf = try uploaded(a, dev, cmds, staging, table, timeout); + defer table_buf.deinit(); + const oor = [_]u32{rows}; // out of range: the host check refuses it + var oor_buf = try uploaded(a, dev, cmds, staging, std.mem.sliceAsBytes(&oor), timeout); + defer oor_buf.deinit(); + var ids_buf = try uploaded(a, dev, cmds, staging, std.mem.sliceAsBytes(ids[0..]), timeout); + defer ids_buf.deinit(); + const out_bytes = try a.alloc(u8, tokens * cols * esz + fx.tail * esz); + defer a.free(out_bytes); + @memset(out_bytes, 0xa5); + var out_buf = try uploaded(a, dev, cmds, staging, out_bytes, timeout); + defer out_buf.deinit(); + const out_view = try vk.tensor.View.init(&out_buf, kind, 0, tokens * cols); + const table_view = try vk.tensor.View.init(&table_buf, kind, 0, rows * cols); + const ids_view = try vk.tensor.View.init(&ids_buf, .u32, 0, tokens); + try vk.orExit(o.gather(cmds, out_view, table_view, ids_view, rows, timeout)); + try fetch(cmds, staging, out_buf, out_bytes, timeout); + for (ids, 0..) |id, t| { + const want = table[@as(usize, id) * cols * esz ..][0 .. cols * esz]; + const got = out_bytes[t * cols * esz ..][0 .. cols * esz]; + if (std.mem.indexOfDiff(u8, got, want)) |at| { + std.debug.print("FAIL {s}: token {d} (id {d}) differs at byte {d}\n", .{ gate, t, id, at }); + return false; + } + } + for (out_bytes[tokens * cols * esz ..]) |byte| if (byte != 0xa5) { + std.debug.print("FAIL {s}: guard tail clobbered\n", .{gate}); + return false; + }; + // the out-of-range id must gather zeros, never table row 0 + const oor_view = try vk.tensor.View.init(&oor_buf, .u32, 0, 1); + try vk.orExit(o.gather(cmds, try out_view.slice(0, cols), table_view, oor_view, rows, timeout)); + try fetch(cmds, staging, out_buf, out_bytes, timeout); + for (out_bytes[0 .. cols * esz]) |byte| if (byte != 0) { + std.debug.print("FAIL {s}: an out-of-range id gathered nonzero bytes\n", .{gate}); + return false; + }; + std.debug.print("MILESTONE {s} elapsed_ms={d} rows={d} cols={d} tokens={d} rows exact incl repeat, out-of-range id gathers zeros, tails unchanged\n", .{ gate, clock.ms(), rows, cols, tokens }); + return true; +} + +/// Host-side token validation: ids inside rows pass, a boundary id is refused before any upload. +fn checkIdsGate(clock: Elapsed) !bool { + const rows: u64 = 8; + vk.ops_q4.checkIds(&.{ 0, 7, 3, 7 }, rows) catch |e| { + std.debug.print("FAIL check-ids: valid ids refused with {s}\n", .{@errorName(e)}); + return false; + }; + const bad = [_]u32{ 0, 8 }; // rows itself is the first out-of-range id + if (vk.ops_q4.checkIds(&bad, rows)) |_| { + std.debug.print("FAIL check-ids: id {d} of {d} rows passed\n", .{ bad[1], rows }); + return false; + } else |_| {} + std.debug.print("MILESTONE check-ids elapsed_ms={d} valid ids pass, boundary id refused before upload\n", .{clock.ms()}); + return true; +} + +// Embedding gathers packed rows by IDs into bounded bf16 output without changing the table. +fn embedGate(a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, o: *vk.ops.Ops, scratch: *vk.arena.Arena, case: fx.EmbedCase, timeout: u64) !bool { + var t = try fx.EmbedTable.init(a, case, 0x57a4); + defer t.deinit(); + var w_buf = try uploaded(a, dev, cmds, staging, std.mem.sliceAsBytes(t.w), timeout); + defer w_buf.deinit(); + var s_buf = try uploaded(a, dev, cmds, staging, std.mem.sliceAsBytes(t.scales), timeout); + defer s_buf.deinit(); + var b_buf = try uploaded(a, dev, cmds, staging, std.mem.sliceAsBytes(t.biases), timeout); + defer b_buf.deinit(); + var id_buf = try uploaded(a, dev, cmds, staging, std.mem.sliceAsBytes(t.ids), timeout); + defer id_buf.deinit(); + const out_elems = case.tokens * case.k; + var out = try poisoned16Buf(a, dev, cmds, staging, out_elems, timeout); + defer out.buf.deinit(); + const out_view = try vk.tensor.View.init(&out.buf, .bf16, 0, out_elems); + const id_view = try vk.tensor.View.init(&id_buf, .u32, 0, case.tokens); + const w_view = try vk.tensor.View.init(&w_buf, .u32, 0, t.w.len - fx.tail); + const s_view = try vk.tensor.View.init(&s_buf, .bf16, 0, t.scales.len - fx.tail); + const b_view = try vk.tensor.View.init(&b_buf, .bf16, 0, t.biases.len - fx.tail); + const mark = scratch.mark(); + defer scratch.reset(mark) catch {}; + try vk.orExit(vk.ops_q4.embedding(o, cmds, out_view, id_view, w_view, s_view, b_view, case.k, scratch, timeout)); + const bits = try a.alloc(u16, out_elems + fx.tail); + defer a.free(bits); + try fetch(cmds, staging, out.buf, std.mem.sliceAsBytes(bits), timeout); + for (t.want, 0..) |want, i| { + const got: f64 = cx.fromBf16(bits[i]); + const err = @abs(got - want); + const bound = cx.boundExact(fx.spec, want); + if (!(std.math.isFinite(err) and err <= bound)) { + std.debug.print("FAIL embedding case=\"{s}\": token {d} element {d} got {e}, want {e}, err {e} of bound {e}\n", .{ case.name, i / case.k, i % case.k, got, want, err, bound }); + return false; + } + } + for (bits[out_elems..]) |g| if (g != fx.poison16) { + std.debug.print("FAIL embedding case=\"{s}\": guard tail clobbered\n", .{case.name}); + return false; + }; + const w_back = try a.alloc(u8, std.mem.sliceAsBytes(t.w).len); + defer a.free(w_back); + try fetch(cmds, staging, w_buf, w_back, timeout); + if (std.mem.indexOfDiff(u8, w_back, std.mem.sliceAsBytes(t.w)) != null) { + std.debug.print("FAIL embedding case=\"{s}\": the table changed\n", .{case.name}); + return false; + } + std.debug.print("MILESTONE embedding elapsed_ms={d} case=\"{s}\" tokens={d} rows={d} k={d} within bf16 bound, ids include a repeat, table and tails unchanged\n", .{ clock.ms(), case.name, case.tokens, case.rows, case.k }); + return true; +} + +// BF16 Q4 projection checks FP64 bounds, three identical repeats, unchanged inputs/guards and nonfinite rejection. +fn projectGate(a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, o: *vk.ops.Ops, scratch: *vk.arena.Arena, case: fx.Case, seed: u64, timeout: u64) !bool { + const f = try fx.Fixture.init(a, case, seed); + defer f.deinit(); + var x_buf = try uploaded(a, dev, cmds, staging, std.mem.sliceAsBytes(f.x), timeout); + defer x_buf.deinit(); + var w_buf = try uploaded(a, dev, cmds, staging, std.mem.sliceAsBytes(f.w), timeout); + defer w_buf.deinit(); + var s_buf = try uploaded(a, dev, cmds, staging, std.mem.sliceAsBytes(f.scales), timeout); + defer s_buf.deinit(); + var b_buf = try uploaded(a, dev, cmds, staging, std.mem.sliceAsBytes(f.biases), timeout); + defer b_buf.deinit(); + var out = try poisoned16Buf(a, dev, cmds, staging, f.out.len, timeout); + defer out.buf.deinit(); + const x_view = try vk.tensor.View.init(&x_buf, .bf16, f.case.at.x * 2, f.q.x.len); + const w_view = try vk.tensor.View.init(&w_buf, .u32, f.case.at.w * 4, f.w.len - f.case.at.w - fx.tail); + const s_view = try vk.tensor.View.init(&s_buf, .bf16, f.case.at.scales * 2, f.q.scales.len); + const b_view = try vk.tensor.View.init(&b_buf, .bf16, f.case.at.biases * 2, f.q.biases.len); + const out_view = try vk.tensor.View.init(&out.buf, .bf16, f.case.at.out * 2, f.want.len); + const bits = try a.alloc(u16, f.out.len); + defer a.free(bits); + var r: fx.Report = .{}; + const first = try a.dupe(u16, bits); + defer a.free(first); + const mark = scratch.mark(); + defer scratch.reset(mark) catch {}; + for (0..repeats) |rep| { + try vk.orExit(vk.upload(cmds, out.buf, 0, staging, std.mem.sliceAsBytes(f.out), timeout)); + try vk.orExit(vk.ops_q4.project(o, cmds, out_view, x_view, w_view, s_view, b_view, case.shape.k, scratch, timeout)); + try scratch.reset(mark); + try fetch(cmds, staging, out.buf, std.mem.sliceAsBytes(bits), timeout); + r = f.check(bits) catch |e| { + const m = f.measure(bits); + std.debug.print("FAIL project case=\"{s}\" repeat={d}: {s}; worst {d:.3} of bound at row {d} col {d}: " ++ + "got {e} want {e}\n", .{ case.name, rep, @errorName(e), m.worst, m.row, m.col, cx.fromBf16(bits[f.case.at.out + m.row * case.shape.n + m.col]), f.want[m.row * case.shape.n + m.col] }); + return false; + }; + if (rep == 0) { + @memcpy(first, bits); + } else if (std.mem.indexOfDiff(u16, first, bits)) |at| { + std.debug.print("FAIL project case=\"{s}\": repeat {d} differs at element {d}\n", .{ case.name, rep, at }); + return false; + } + } + // The vec4 image is the scalar image with wider loads: its bytes must equal the scalar bytes exactly. + if (vk.ops_q4.qmmVariant(x_view, w_view) != .vec4) { + std.debug.print("FAIL project case=\"{s}\": its windows are not 16-byte aligned, so the vec4 image is not exercised\n", .{case.name}); + return false; + } + const scalar_bits = try a.alloc(u16, f.out.len); + defer a.free(scalar_bits); + inline for (.{ vk.ops_q4.Qmm.scalar, vk.ops_q4.Qmm.vec4 }) |force| { + try vk.orExit(vk.upload(cmds, out.buf, 0, staging, std.mem.sliceAsBytes(f.out), timeout)); + try vk.orExit(vk.ops_q4.projectAs(o, force, cmds, out_view, x_view, w_view, s_view, b_view, case.shape.k, scratch, timeout)); + try scratch.reset(mark); + try fetch(cmds, staging, out.buf, std.mem.sliceAsBytes(if (force == .scalar) scalar_bits else bits), timeout); + } + if (std.mem.indexOfDiff(u16, scalar_bits, bits)) |at| { + std.debug.print("FAIL project case=\"{s}\": the vec4 image differs from the scalar image at element {d}\n", .{ case.name, at }); + return false; + } + if (std.mem.indexOfDiff(u16, first, bits)) |at| { + std.debug.print("FAIL project case=\"{s}\": the default pick differs from the forced vec4 image at element {d}\n", .{ case.name, at }); + return false; + } + const tiled = try tileProject(a, dev, cmds, staging, o, f, case, x_view, out, out_view, bits, timeout); + if (tiled == false) return false; + const back = try a.alloc(u8, @max(std.mem.sliceAsBytes(f.x).len, std.mem.sliceAsBytes(f.w).len, std.mem.sliceAsBytes(f.scales).len, std.mem.sliceAsBytes(f.biases).len)); + defer a.free(back); + for ([_][]const u8{ std.mem.sliceAsBytes(f.x), std.mem.sliceAsBytes(f.w), std.mem.sliceAsBytes(f.scales), std.mem.sliceAsBytes(f.biases) }, [_]vk.Buffer{ x_buf, w_buf, s_buf, b_buf }, 0..) |bytes, buf, n| { + try fetch(cmds, staging, buf, back[0..bytes.len], timeout); + if (std.mem.indexOfDiff(u8, bytes, back[0..bytes.len])) |at| { + std.debug.print("FAIL project case=\"{s}\": binding {d} changed at byte {d}\n", .{ case.name, n, at }); + return false; + } + } + const planted_bits = [_]u16{ 0x7fc0, 0x7f80, 0xff80 }; + const planted = [_][]const u8{ "nan", "+inf", "-inf" }; + const at = f.case.at.out + r.row * case.shape.n + r.col; + for (planted_bits, planted) |pb, pname| { + bits[at] = pb; + if (f.check(bits)) |_| { + std.debug.print("FAIL project case=\"{s}\": a planted {s} output passed the check\n", .{ case.name, pname }); + return false; + } else |e| if (e != error.OutOfBound) { + std.debug.print("FAIL project case=\"{s}\": planted {s} gave {s}, not OutOfBound\n", .{ case.name, pname, @errorName(e) }); + return false; + } + } + std.debug.print("MILESTONE project elapsed_ms={d} case=\"{s}\" worst={d:.4} of bound at row {d} col {d}, max_abs={e}, tile={s}, " ++ + "repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused\n", .{ clock.ms(), case.name, r.worst, r.row, r.col, r.max_abs, if (tiled == null) "n/a" else "identical" }); + return true; +} + +// The tiled image reads a host repack of the same weights; its bytes must equal the default image's. Null: too large. +fn tileProject(a: std.mem.Allocator, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, o: *vk.ops.Ops, f: fx.Fixture, case: fx.Case, x_view: vk.tensor.View, out: Poisoned16, out_view: vk.tensor.View, want: []const u16, timeout: u64) !?bool { + const n: usize = case.shape.n; + const k: usize = case.shape.k; + const stored = vk.tile.padded(n); + if (!vk.tile.fits(n, k) or stored * k / 2 > dev.properties.limits.maxStorageBufferRange) return null; + const groups = k / 64; + const words = std.mem.sliceAsBytes(f.w)[f.case.at.w * 4 ..][0 .. n * k / 2]; + const scales = std.mem.sliceAsBytes(f.scales)[f.case.at.scales * 2 ..][0 .. n * groups * 2]; + const biases = std.mem.sliceAsBytes(f.biases)[f.case.at.biases * 2 ..][0 .. n * groups * 2]; + const tiled = try a.alloc(u8, stored * k / 2); + defer a.free(tiled); + vk.tile.tileWords(tiled, words, n, k); + var w_buf = try uploaded(a, dev, cmds, staging, tiled, timeout); + defer w_buf.deinit(); + const params = tiled[0 .. stored * groups * 2]; + vk.tile.tileParams(params, scales, n, k); + var s_buf = try uploaded(a, dev, cmds, staging, params, timeout); + defer s_buf.deinit(); + vk.tile.tileParams(params, biases, n, k); + var b_buf = try uploaded(a, dev, cmds, staging, params, timeout); + defer b_buf.deinit(); + const w_view = try vk.tensor.View.init(&w_buf, .u32, 0, tiled.len / 4); + const s_view = try vk.tensor.View.init(&s_buf, .bf16, 0, stored * groups); + const b_view = try vk.tensor.View.init(&b_buf, .bf16, 0, stored * groups); + const got = try a.alloc(u16, want.len); + defer a.free(got); + try vk.orExit(vk.upload(cmds, out.buf, 0, staging, std.mem.sliceAsBytes(f.out), timeout)); + try vk.orExit(vk.ops_q4.projectTiled(o, cmds, out_view, x_view, w_view, s_view, b_view, case.shape.k, timeout)); + try fetch(cmds, staging, out.buf, std.mem.sliceAsBytes(got), timeout); + if (std.mem.indexOfDiff(u16, want, got)) |at| { + std.debug.print("FAIL project case=\"{s}\": the tiled image differs from the default image at element {d}\n", .{ case.name, at }); + return false; + } + return true; +} + +const Poisoned16 = struct { buf: vk.Buffer }; + +fn poisoned16Buf(a: std.mem.Allocator, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, elems: usize, timeout: u64) !Poisoned16 { + var out: Poisoned16 = .{ .buf = try vk.Buffer.init(dev, .device, (elems + fx.tail) * 2) }; + errdefer out.buf.deinit(); + const fill = try a.alloc(u16, elems + fx.tail); + defer a.free(fill); + @memset(fill, fx.poison16); + try vk.orExit(vk.upload(cmds, out.buf, 0, staging, std.mem.sliceAsBytes(fill), timeout)); + return out; +} diff --git a/zig/tests/vulkan/stage2_fixture.zig b/zig/tests/vulkan/stage2_fixture.zig new file mode 100644 index 00000000000..39630200636 --- /dev/null +++ b/zig/tests/vulkan/stage2_fixture.zig @@ -0,0 +1,276 @@ +// ORACLE_QWEN35 add/SwiGLU/RMSNorm fixtures under fixed budgets; nonfinite by class (drivers vary NaN payloads). +const std = @import("std"); +const contract = @import("contract"); +const c = contract.spec; +// The cases live in stage2_fixture_ops.zig; the tests stay here so every importer of this file runs them. +const ops = @import("stage2_fixture_ops.zig"); +const f32ops = @import("stage2_fixture_f32.zig"); +comptime { + _ = f32ops; +} + +pub const tail = ops.tail; +pub const poison16 = ops.poison16; +pub const classOfBits = ops.classOfBits; +pub const classify = ops.classify; +pub const Report = ops.Report; +pub const Offsets = ops.Offsets; +pub const AddFill = ops.AddFill; +pub const AddCase = ops.AddCase; +pub const add_cases = ops.add_cases; +pub const Add = ops.Add; +pub const SwigluCase = ops.SwigluCase; +pub const swiglu_cases = ops.swiglu_cases; +pub const swiglu_gate_controls = ops.swiglu_gate_controls; +pub const Swiglu = ops.Swiglu; +pub const RmsFill = ops.RmsFill; +pub const RmsCase = ops.RmsCase; +pub const RmsOffsets = ops.RmsOffsets; +pub const rms_cases = ops.rms_cases; +pub const rms_eps = ops.rms_eps; +pub const Rms = ops.Rms; + +const testing = std.testing; + +test "add wants are the shader pipeline: f32 add then one RNE bf16 store" { + for (add_cases) |case| { + const f = try Add.init(testing.allocator, case, 0x5702); + defer f.deinit(); + for (0..f.case.elems) |i| { + const want = c.toBf16(c.fromBf16(f.a_bits[f.case.at.a + i]) + c.fromBf16(f.b_bits[f.case.at.b + i])); + try testing.expectEqual(want, f.want[i]); + } + } +} + +test "add cancel case produces exact tiny sums and the classes case plants every nonfinite class" { + const f = try Add.init(testing.allocator, add_cases[1], 0x5703); + defer f.deinit(); + var zeros: usize = 0; + for (0..f.case.elems) |i| { + const sum = c.fromBf16(f.a_bits[f.case.at.a + i]) + c.fromBf16(f.b_bits[f.case.at.b + i]); + if (@abs(c.fromBf16(f.want[i]) - sum) == 0) zeros += 1; + } + try testing.expect(zeros > 0); // v + -(v - 0*ulp) pairs cancel exactly + const g = try Add.init(testing.allocator, add_cases[2], 0x5704); + defer g.deinit(); + var classes: std.EnumSet(c.Class) = .empty; + for (g.want) |w| classes.insert(classOfBits(w)); + try testing.expect(classes.contains(.nan)); + try testing.expect(classes.contains(.pos_inf)); + try testing.expect(classes.contains(.neg_inf)); + try testing.expect(classes.contains(.finite)); +} + +test "swiglu controls plant nan, both infinities and signed zeros, and references stay finite elsewhere" { + for (swiglu_cases) |case| { + const f = try Swiglu.init(testing.allocator, case, 0x5705); + defer f.deinit(); + for (f.want, 0..) |w, i| { + if (case.controls and i < swiglu_gate_controls.len) { + const gate = c.fromBf16(f.gate[f.case.at.a + i]); + if (!std.math.isFinite(gate) or gate == 0) continue; // class/plant entries + } + if (!case.controls) try testing.expect(std.math.isFinite(w)); + } + } + const g = try Swiglu.init(testing.allocator, swiglu_cases[1], 0x5706); + defer g.deinit(); + try testing.expect(std.math.isNan(g.want[0])); // nan gate + try testing.expect(std.math.isInf(g.want[1])); // +inf gate, up 1 + try testing.expect(std.math.isNan(g.want[2])); // -inf gate -> -inf/inf + try testing.expect(g.want[3] == 0); // +0 gate -> exactly zero + try testing.expect(std.math.signbit(g.want[4])); // -0 gate keeps its zero sign +} + +test "rms references are finite and the zero row wants exact zeros" { + for (rms_cases) |case| { + const f = try Rms.init(testing.allocator, case, 0x5707); + defer f.deinit(); + for (f.want) |w| try testing.expect(std.math.isFinite(w)); + if (case.zero_row) { + const zero_at = (case.rows - 1) * case.width; + for (f.want[zero_at..]) |w| try testing.expectEqual(@as(f64, 0), w); + } + } +} + +test "classify names the class of raw bf16 bits with the sign visible" { + try testing.expectEqualStrings("+zero", classify(0x0000)); + try testing.expectEqualStrings("-zero", classify(0x8000)); + try testing.expectEqualStrings("+subnormal", classify(0x0001)); + try testing.expectEqualStrings("-subnormal", classify(0x807f)); + try testing.expectEqualStrings("+normal", classify(0x0080)); + try testing.expectEqualStrings("-normal", classify(0xbf80)); + try testing.expectEqualStrings("+inf", classify(0x7f80)); + try testing.expectEqualStrings("-inf", classify(0xff80)); + try testing.expectEqualStrings("+nan", classify(0x7fc0)); + try testing.expectEqualStrings("-nan", classify(0xffc1)); +} + +test "add firstMismatch reports the first bad element with both inputs and want/got words" { + const f = try Add.init(testing.allocator, add_cases[2], 0x5717); + defer f.deinit(); + const bits = try testing.allocator.dupe(u16, f.out); + defer testing.allocator.free(bits); + @memcpy(bits[f.case.at.out..][0..f.want.len], f.want); + try testing.expectEqual(@as(?Add.Mismatch, null), try f.firstMismatch(bits)); + // Two corruptions at finite-want elements: the helper must report only the FIRST, zero-based. + bits[f.case.at.out + 10] = 0x0000; + bits[f.case.at.out + 11] = 0x0000; + const m = (try f.firstMismatch(bits)).?; + try testing.expectEqual(error.BitMismatch, m.kind); + try testing.expectEqual(@as(usize, 10), m.index); + try testing.expectEqual(f.a_bits[f.case.at.a + 10], m.a); + try testing.expectEqual(f.b_bits[f.case.at.b + 10], m.b); + try testing.expectEqual(f.want[10], m.want); + try testing.expectEqual(@as(u16, 0x0000), m.got); + try testing.expectError(error.BitMismatch, f.check(bits)); + // A class break at a nonfinite-want element: NonFiniteMismatch, still first index. + @memcpy(bits[f.case.at.out..][0..f.want.len], f.want); + const inf_at: usize = blk: { + for (f.want, 0..) |w, i| if (classOfBits(w) == .pos_inf) break :blk i; + unreachable; + }; + bits[f.case.at.out + inf_at] = 0x3f80; // +inf wanted, +1 got: class break + const nf = (try f.firstMismatch(bits)).?; + try testing.expectEqual(error.NonFiniteMismatch, nf.kind); + try testing.expectEqual(inf_at, nf.index); + try testing.expectError(error.NonFiniteMismatch, f.check(bits)); +} + +test "add firstMismatch keeps the named length and guard failures" { + const f = try Add.init(testing.allocator, add_cases[2], 0x5718); + defer f.deinit(); + try testing.expectError(error.LengthMismatch, f.firstMismatch(f.out[0 .. f.out.len - 1])); + const bits = try testing.allocator.dupe(u16, f.out); + defer testing.allocator.free(bits); + @memcpy(bits[f.case.at.out..][0..f.want.len], f.want); + bits[0] = 0x1234; // guard poison clobbered before the data window + try testing.expectError(error.GuardClobbered, f.firstMismatch(bits)); +} + +test "only unsupported-device subnormal classes may flush; preserving devices stay strict" { + const f_base = try Add.init(testing.allocator, add_cases[2], 0x5719); + defer f_base.deinit(); + const bits = try testing.allocator.dupe(u16, f_base.out); + defer testing.allocator.free(bits); + for ([_]bool{ true, false }) |preserve32| { + var f = f_base; + f.allow_subnormal_flush = !preserve32; + @memcpy(bits[f.case.at.out..][0..f.want.len], f.want); + try f.check(bits); + for ([_]usize{ 12, 13, 14 }) |i| bits[f.case.at.out + i] = f.want[i] & 0x8000; + if (preserve32) { + try testing.expectError(error.BitMismatch, f.check(bits)); + } else { + try f.check(bits); + bits[f.case.at.out + 12] = 0x8000; + try testing.expectError(error.BitMismatch, f.check(bits)); + bits[f.case.at.out + 12] = 0; + for ([_]usize{ 4, 5, 10, 15, 16 }) |i| { + const old = bits[f.case.at.out + i]; + bits[f.case.at.out + i] = old ^ 0x8000; + try testing.expectError(error.BitMismatch, f.check(bits)); + bits[f.case.at.out + i] = old; + } + bits[f.case.at.out] = 0; + try testing.expectError(error.NonFiniteMismatch, f.check(bits)); + } + } + // The policy is scoped to the classes gate, not arbitrary finite outputs. + var normal = try Add.init(testing.allocator, add_cases[0], 0x571a); + defer normal.deinit(); + normal.allow_subnormal_flush = true; + const normal_bits = try testing.allocator.dupe(u16, normal.out); + defer testing.allocator.free(normal_bits); + @memcpy(normal_bits[normal.case.at.out..][0..normal.want.len], normal.want); + normal.want[0] = 0x0001; + normal_bits[normal.case.at.out] = 0; + try testing.expectError(error.BitMismatch, normal.check(normal_bits)); +} + +test "the checkers reject a deliberately wrong finite output, not only accept good ones" { + // RMS dynamic range: want + 0.1 is far outside the output-relative budget. + { + const f = try Rms.init(testing.allocator, rms_cases[2], 0x5710); + defer f.deinit(); + const bits = try testing.allocator.dupe(u16, f.out); + defer testing.allocator.free(bits); + for (f.want, 0..) |w, i| bits[f.case.at.out + i] = c.toBf16(@floatCast(w)); + _ = try f.check(bits); + const at = f.case.width / 2; + bits[f.case.at.out + at] = c.toBf16(@floatCast(f.want[at] + 0.1)); + try testing.expectError(error.OutOfBound, f.check(bits)); + } + // RMS zero row: one bf16 subnormal unit of output must be refused outright. + { + const f = try Rms.init(testing.allocator, rms_cases[1], 0x5711); + defer f.deinit(); + const bits = try testing.allocator.dupe(u16, f.out); + defer testing.allocator.free(bits); + for (f.want, 0..) |w, i| bits[f.case.at.out + i] = c.toBf16(@floatCast(w)); + _ = try f.check(bits); + bits[f.case.at.out + (f.case.rows - 1) * f.case.width] = 0x0001; + try testing.expectError(error.OutOfBound, f.check(bits)); + } + // SwiGLU: a wrong finite output must be refused too. + { + const f = try Swiglu.init(testing.allocator, swiglu_cases[0], 0x5712); + defer f.deinit(); + const bits = try testing.allocator.dupe(u16, f.out); + defer testing.allocator.free(bits); + for (f.want, 0..) |w, i| bits[f.case.at.out + i] = c.toBf16(@floatCast(w)); + _ = try f.check(bits); + bits[f.case.at.out] = c.toBf16(@floatCast(f.want[0] + 0.1)); + try testing.expectError(error.OutOfBound, f.check(bits)); + } + // Residual add: a wrong finite output is a bit mismatch. + { + const f = try Add.init(testing.allocator, add_cases[0], 0x5713); + defer f.deinit(); + const bits = try testing.allocator.dupe(u16, f.out); + defer testing.allocator.free(bits); + @memcpy(bits[f.case.at.out..][0..f.want.len], f.want); + try f.check(bits); + bits[f.case.at.out] = c.toBf16(@floatCast(c.fromBf16(f.want[0]) + 0.1)); + try testing.expectError(error.BitMismatch, f.check(bits)); + } +} + +fn addAllocRun(a: std.mem.Allocator, case: AddCase) !void { + const f = try Add.init(a, case, 0x5714); + f.deinit(); +} + +fn swigluAllocRun(a: std.mem.Allocator, case: SwigluCase) !void { + const f = try Swiglu.init(a, case, 0x5715); + f.deinit(); +} + +fn rmsAllocRun(a: std.mem.Allocator, case: RmsCase) !void { + const f = try Rms.init(a, case, 0x5716); + f.deinit(); +} + +test "every fixture init survives each single allocation failure without leaking" { + // checkAllAllocationFailures fails each init allocation; error.OutOfMemory required; errdefers free earlier bytes. + try std.testing.checkAllAllocationFailures(testing.allocator, addAllocRun, .{add_cases[0]}); + try std.testing.checkAllAllocationFailures(testing.allocator, addAllocRun, .{add_cases[2]}); + try std.testing.checkAllAllocationFailures(testing.allocator, swigluAllocRun, .{swiglu_cases[0]}); + try std.testing.checkAllAllocationFailures(testing.allocator, swigluAllocRun, .{swiglu_cases[1]}); + for (rms_cases) |case| try std.testing.checkAllAllocationFailures(testing.allocator, rmsAllocRun, .{case}); +} + +test "guards sit poisoned outside every data window" { + const f = try Rms.init(testing.allocator, rms_cases[1], 0x5708); + defer f.deinit(); + for (f.x, 0..) |g, i| { + if (i >= f.case.at.x and i < f.case.at.x + f.case.rows * f.case.width) continue; + try testing.expectEqual(poison16, g); + } + for (f.out, 0..) |g, i| { + if (i >= f.case.at.out and i < f.case.at.out + f.case.rows * f.case.width) continue; + try testing.expectEqual(poison16, g); + } +} diff --git a/zig/tests/vulkan/stage2_fixture_f32.zig b/zig/tests/vulkan/stage2_fixture_f32.zig new file mode 100644 index 00000000000..d2e809221c4 --- /dev/null +++ b/zig/tests/vulkan/stage2_fixture_f32.zig @@ -0,0 +1,94 @@ +// The f32 activation pipelines' want formulas: no store rounding inside the wide layers, one RNE store at a transition. +const std = @import("std"); +const contract = @import("contract"); +const c = contract.spec; +const ops = @import("stage2_fixture_ops.zig"); + +const testing = std.testing; + +/// add_f32: the exact f32 sum, stored unrounded. +fn addF32Want(a_bits: u16, b_bits: u16) f32 { + return c.fromBf16(a_bits) + c.fromBf16(b_bits); +} + +/// add_wide_bf16: the f32 stream plus the widened bf16 residual, rounded once to bf16 on the store. +fn addWideWant(a_f32: f32, b_bits: u16) u16 { + return c.toBf16(a_f32 + c.fromBf16(b_bits)); +} + +/// add_bf16_f32: the widened bf16 stream joins the f32 residual, the f32 sum stays unrounded. +fn addWideF32Want(a_bits: u16, b_f32: f32) f32 { + return c.fromBf16(a_bits) + b_f32; +} + +test "add f32 stores the f32 sum; the narrow pipeline is its one-store rounding" { + for (ops.class_pairs) |pair| { + const f64_sum = @as(f64, c.fromBf16(pair[0])) + @as(f64, c.fromBf16(pair[1])); + const wide = addF32Want(pair[0], pair[1]); + if (!std.math.isFinite(wide)) { + // inf - inf and nan classes agree by class, payloads vary + try testing.expect(!std.math.isFinite(@as(f32, @floatCast(f64_sum)))); + continue; + } + // the f32 add is the f64-exact sum rounded once to f32 (both operands are bf16, hence f32-exact) + try testing.expectEqual(@as(f32, @floatCast(f64_sum)), wide); + // the bf16 pipeline's want is exactly the f32 sum's store rounding + try testing.expectEqual(c.toBf16(wide), c.toBf16(@floatCast(f64_sum))); + } +} + +test "cancel pairs stay exact under the f32 store" { + var prng = std.Random.DefaultPrng.init(0x5702); + const rng = prng.random(); + for (0..32) |i| { + const vbits = ops.sampleBits(rng, 1); + const v = c.fromBf16(vbits); + const k: f32 = @floatFromInt(i % 3); + const ulp = std.math.pow(f32, 2, @floor(std.math.log2(@abs(v))) - 7); + const nbits = c.toBf16(v - k * ulp); + // the near-cancelling sum is bf16-representable, so both pipelines are bit-exact on it + try testing.expectEqual(c.fromBf16(nbits) + v, addF32Want(vbits, nbits)); + } +} + +test "the mixed transitions round exactly once, on the side whose store holds bf16" { + const a: f32 = 1.00390625; // not bf16-representable + const b_bits = c.toBf16(0.0625); + // wide_bf16: the f32 sum rounds on the bf16 store + try testing.expectEqual(c.toBf16(a + 0.0625), addWideWant(a, b_bits)); + // wide_f32: the bf16 side widens losslessly and the f32 store keeps the sum + const a_bits = c.toBf16(1.0078125); + const b: f32 = 1.0 / 3.0; + try testing.expectEqual(c.fromBf16(a_bits) + b, addWideF32Want(a_bits, b)); +} + +test "swiglu f32 keeps the ASH single-round formula in f32 without a store rounding" { + for (ops.swiglu_gate_controls, 0..) |g_bits, i| { + const u_bits = ops.swiglu_up_controls[i]; + const g32: f32 = c.fromBf16(g_bits); + const u32v: f32 = c.fromBf16(u_bits); + if (!std.math.isFinite(g32) or !std.math.isFinite(u32v)) continue; + const want32 = g32 / (1 + @exp(-g32)) * u32v; + // the f32 result re-rounds onto the bf16 pipeline's want within one bf16 ulp of the value + try testing.expectApproxEqAbs(@as(f64, c.fromBf16(c.toBf16(want32))), @as(f64, want32), 0.004 * @abs(want32)); + } +} + +test "rms f32 stores the unrounded norm and f32_in rounds once at the store" { + var prng = std.Random.DefaultPrng.init(0x5708); + const rng = prng.random(); + const width = 64; + var x: [width]f32 = undefined; + var sum_sq: f32 = 0; + for (&x) |*v| { + v.* = c.fromBf16(ops.sampleBits(rng, 1)); + sum_sq += v.* * v.*; + } + const inv = 1.0 / std.math.sqrt(sum_sq / @as(f32, width) + ops.rms_eps); + const gamma: f32 = c.fromBf16(c.toBf16(0.75)); + for (x) |v| { + const wide = gamma * (v * inv); + // the f32 store keeps the value; the f32_in variant's bf16 store rounds it once + try testing.expectApproxEqAbs(@as(f64, c.fromBf16(c.toBf16(wide))), @as(f64, wide), 0.004 * @abs(wide)); + } +} diff --git a/zig/tests/vulkan/stage2_fixture_ops.zig b/zig/tests/vulkan/stage2_fixture_ops.zig new file mode 100644 index 00000000000..6b23c4b394f --- /dev/null +++ b/zig/tests/vulkan/stage2_fixture_ops.zig @@ -0,0 +1,372 @@ +// The stage 2 fixture's cases and checkers: residual add, SwiGLU and RMSNorm against the Metal-oracle semantics. +const std = @import("std"); +const contract = @import("contract"); +const c = contract.spec; + +/// Poisoned elements after each binding's data, so an overrun reads poison rather than past the buffer. +pub const tail = 16; +pub const poison16: u16 = 0xffff; + +pub fn classOfBits(bits: u16) c.Class { + return c.classOf(c.fromBf16(bits)); +} + +/// Human name for the value class of raw bf16 bits: sign, then zero/subnormal/normal/inf/nan. +pub fn classify(bits: u16) []const u8 { + const neg = bits & 0x8000 != 0; + const e = (bits >> 7) & 0xff; + const m = bits & 0x7f; + return switch (e) { + 0xff => if (m == 0) (if (neg) "-inf" else "+inf") else (if (neg) "-nan" else "+nan"), + 0 => if (m == 0) (if (neg) "-zero" else "+zero") else (if (neg) "-subnormal" else "+subnormal"), + else => if (neg) "-normal" else "+normal", + }; +} + +/// The worst output: |got - want| over its bound (nonfinite counts as infinitely far). +pub const Report = struct { worst: f64 = 0, row: usize = 0, col: usize = 0 }; + +/// Element offsets (bf16 elements; 256-byte alignment = multiples of 128) of each binding's data. +pub const Offsets = struct { a: usize = 0, b: usize = 0, out: usize = 0 }; + +fn poisonedShell(a: std.mem.Allocator, at: usize, len: usize) ![]u16 { + const buf = try a.alloc(u16, at + len + tail); + @memset(buf, poison16); + return buf; +} + +fn guardIntact(shell: []const u16, at: usize, len: usize) !void { + for (shell, 0..) |g, i| { + if (i >= at and i < at + len) continue; + if (g != poison16) return error.GuardClobbered; + } +} + +/// A bf16-representable f32 value near `1` scaled by `scale` (a power of two). +pub fn sampleBits(rng: std.Random, scale: f32) u16 { + const v = (rng.float(f32) + 0.5) * scale * (if (rng.boolean()) @as(f32, 1) else -1); + return c.toBf16(v); +} + +// ---- Residual add: the shader pipeline (f32 add + one RNE bf16 store), bit for bit ---- + +pub const AddFill = enum { random, cancel, classes }; + +pub const AddCase = struct { name: []const u8, elems: usize, at: Offsets = .{}, fill: AddFill = .random }; + +// Counts %4==0 (binary_vec moves 4/8-byte load); offsets 256B aligned, the largest minStorageBufferOffsetAlignment. + +// 0x0001..0x007f SUBNORMALs (0x0080 smallest NORMAL); measured properties, never vendor names, allow signed zero. +pub const add_cases = [_]AddCase{ + .{ .name = "random 512", .elems = 512, .at = .{ .a = 128, .b = 256, .out = 384 } }, + .{ .name = "cancel 512", .elems = 512, .at = .{ .a = 128, .b = 256, .out = 384 }, .fill = .cancel }, + .{ .name = "classes 20", .elems = 20, .at = .{ .a = 128, .b = 256, .out = 384 }, .fill = .classes }, +}; + +/// Planted classes-case inputs; every pair is exercised (elems == len). +pub const class_pairs = [_][2]u16{ + .{ 0x7f80, 0x3f80 }, // +inf + 1 -> +inf + .{ 0xff80, 0x3f80 }, // -inf + 1 -> -inf + .{ 0x7f80, 0xff80 }, // +inf + -inf -> nan + .{ 0x7fc0, 0x3f80 }, // nan + 1 -> nan + .{ 0x0000, 0x8000 }, // +0 + -0 -> +0 (round to nearest) + .{ 0x8000, 0x8000 }, // -0 + -0 -> -0 + .{ 0x7f7f, 0x7f7f }, // bf16 max + bf16 max -> f32 overflow -> +inf + .{ 0xff7f, 0xff7f }, // both-min -> -inf + .{ 0x3f80, 0xbf80 }, // 1 + -1 -> +0 + .{ 0x4000, 0xc000 }, // 2 + -2 -> +0 + .{ 0x3f80, 0x3f80 }, // 1 + 1 -> 2 + .{ 0x3f81, 0xbd80 }, // 1.0078125 + -0.0625, exact in f32 + .{ 0x0001, 0x0001 }, // smallest bf16 SUBNORMAL doubled -> 0x0002 (honest fail on flush-to-zero) + .{ 0x0001, 0x0002 }, // 1 + 2 subnormal units -> 0x0003 + .{ 0x0001, 0x8000 }, // subnormal + -0 -> unchanged + .{ 0x0080, 0x0080 }, // smallest NORMAL doubled -> 0x0100 + .{ 0x0080, 0x8000 }, // smallest normal + -0 -> unchanged + .{ 0x4f00, 0xb080 }, // 2^31 - 2^-29, huge spread: f32 rounds to 2^31 + .{ 0x42f6, 0xc2f6 }, // ±123 -> +0 + .{ 0x3f00, 0xc000 }, // 0.5 + -2 -> -1.5 exact +}; + +pub const Add = struct { + a: std.mem.Allocator, + case: AddCase, + /// Bindings as uploaded: shells poisoned on both sides of the data window. + a_bits: []u16, + b_bits: []u16, + out: []u16, + // The exact bits the shader pipeline (f32 add, one RNE bf16 store) must produce; nonfinite by class only. + want: []u16, + allow_subnormal_flush: bool = false, + + pub fn init(a: std.mem.Allocator, case: AddCase, seed: u64) !Add { + const a_bits = try poisonedShell(a, case.at.a, case.elems); + errdefer a.free(a_bits); + const b_bits = try poisonedShell(a, case.at.b, case.elems); + errdefer a.free(b_bits); + const out = try poisonedShell(a, case.at.out, case.elems); + errdefer a.free(out); + const want = try a.alloc(u16, case.elems); + errdefer a.free(want); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (0..case.elems) |i| { + const pair = switch (case.fill) { + .random => [2]u16{ sampleBits(rng, std.math.pow(f32, 2, @floatFromInt(@as(i32, @intCast(i % 5)) - 2))), sampleBits(rng, std.math.pow(f32, 2, @floatFromInt(@as(i32, @intCast(i % 3)) - 1))) }, + // v and -(v - k*ulp(v)): a bf16-exact difference, sum tiny but exact + .cancel => blk: { + const vbits = sampleBits(rng, 1); + const v = c.fromBf16(vbits); + const k: f32 = @floatFromInt(i % 3); + const ulp = std.math.pow(f32, 2, @floor(std.math.log2(@abs(v))) - 7); + break :blk [2]u16{ vbits, c.toBf16(-(v - k * ulp)) }; + }, + .classes => class_pairs[i % class_pairs.len], + }; + a_bits[case.at.a + i] = pair[0]; + b_bits[case.at.b + i] = pair[1]; + // The shader's own pipeline, evaluated on the host: f32 add, one RNE bf16 store. + want[i] = c.toBf16(c.fromBf16(pair[0]) + c.fromBf16(pair[1])); + } + return .{ .a = a, .case = case, .a_bits = a_bits, .b_bits = b_bits, .out = out, .want = want }; + } + + pub fn deinit(f: Add) void { + f.a.free(f.a_bits); + f.a.free(f.b_bits); + f.a.free(f.out); + f.a.free(f.want); + } + + // First mismatch under the exact check (class-vs-bit), or null when the window matches; guard failures are errors. + pub const Mismatch = struct { kind: error{ NonFiniteMismatch, BitMismatch }, index: usize, a: u16, b: u16, want: u16, got: u16 }; + + pub fn firstMismatch(f: Add, got: []const u16) !?Mismatch { + if (got.len != f.out.len) return error.LengthMismatch; + const at = f.case.at.out; + try guardIntact(got, at, f.case.elems); + for (f.want, 0..) |w, i| { + const g = got[at + i]; + if (classOfBits(g) != classOfBits(w)) return .{ .kind = error.NonFiniteMismatch, .index = i, .a = f.a_bits[f.case.at.a + i], .b = f.b_bits[f.case.at.b + i], .want = w, .got = g }; + const flushed = f.allow_subnormal_flush and f.case.fill == .classes and + (w & 0x7f80) == 0 and (w & 0x007f) != 0 and g == (w & 0x8000); + if (std.math.isFinite(c.fromBf16(w)) and g != w and !flushed) return .{ .kind = error.BitMismatch, .index = i, .a = f.a_bits[f.case.at.a + i], .b = f.b_bits[f.case.at.b + i], .want = w, .got = g }; + } + return null; + } + + /// got: the whole out shell read back; guards keep their poison, data bits equal want exactly. + pub fn check(f: Add, got: []const u16) !void { + if (try f.firstMismatch(got)) |m| return m.kind; + } +}; + +// ---- SwiGLU: one bf16 store plus a documented f32 budget ---- + +// The f32 exp/1+/divide/multiply chain uses 8 ops and a fixed 8*2^-24 budget; exp has no Vulkan-spec per-op bound. +const swiglu_spec: c.Spec = .{ .name = "vulkan swiglu bf16 (ORACLE_QWEN35 one-store)", .accum = .none, .depth = no_depth, .ops = 8, .out = .bf16, .served = &[_]c.Shape{} }; +fn no_depth(_: c.Shape) usize { + return 0; // elementwise: boundExact never calls depth +} + +pub const SwigluCase = struct { name: []const u8, elems: usize, at: Offsets = .{}, controls: bool = false }; + +pub const swiglu_cases = [_]SwigluCase{ + .{ .name = "random 512", .elems = 512, .at = .{ .a = 128, .b = 256, .out = 384 } }, + .{ .name = "controls 16", .elems = 16, .at = .{ .a = 128, .b = 256, .out = 384 }, .controls = true }, +}; + +/// The controls: planted nan/+inf/-inf gates and signed zeros, up mostly 1 with two sign flips. +pub const swiglu_gate_controls = [_]u16{ 0x7fc0, 0x7f80, 0xff80, 0x0000, 0x8000, 0x3f80, 0xbf80, 0x41a0, 0xc1a0, 0x3f00, 0xbf00, 0x4000, 0xc000, 0x3e80, 0x3d00, 0x4800 }; +pub const swiglu_up_controls = [_]u16{ 0x3f80, 0x3f80, 0x3f80, 0x3f80, 0x3f80, 0x3f80, 0x3f80, 0x3f80, 0x3f80, 0xbf80, 0x3f80, 0xc000, 0x3f80, 0xbf80, 0x3f80, 0x3f80 }; + +pub const Swiglu = struct { + a: std.mem.Allocator, + case: SwigluCase, + gate: []u16, + up: []u16, + out: []u16, + /// The f64 reference g / (1 + exp(-g)) * up of each output (finite: bound-checked; else class). + want: []f64, + + pub fn init(a: std.mem.Allocator, case: SwigluCase, seed: u64) !Swiglu { + const gate = try poisonedShell(a, case.at.a, case.elems); + errdefer a.free(gate); + const up = try poisonedShell(a, case.at.b, case.elems); + errdefer a.free(up); + const out = try poisonedShell(a, case.at.out, case.elems); + errdefer a.free(out); + const want = try a.alloc(f64, case.elems); + errdefer a.free(want); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (0..case.elems) |i| { + const g: u16 = if (case.controls) swiglu_gate_controls[i % swiglu_gate_controls.len] else sampleBits(rng, std.math.pow(f32, 2, @floatFromInt(@as(i32, @intCast(i % 4))))); + const u: u16 = if (case.controls) swiglu_up_controls[i % swiglu_up_controls.len] else sampleBits(rng, 1); + gate[case.at.a + i] = g; + up[case.at.b + i] = u; + const g64: f64 = c.fromBf16(g); + const u64v: f64 = c.fromBf16(u); + want[i] = g64 / (1 + @exp(-g64)) * u64v; + } + return .{ .a = a, .case = case, .gate = gate, .up = up, .out = out, .want = want }; + } + + pub fn deinit(f: Swiglu) void { + f.a.free(f.gate); + f.a.free(f.up); + f.a.free(f.out); + f.a.free(f.want); + } + + // Finite wants: within boundExact (bf16 rounding + f32 budget); want == 0 demands exactly-zero; nonfinite: classes. + pub fn check(f: Swiglu, got: []const u16) !Report { + if (got.len != f.out.len) return error.LengthMismatch; + const at = f.case.at.out; + try guardIntact(got, at, f.case.elems); + var r: Report = .{}; + for (f.want, 0..) |w, i| { + const g = got[at + i]; + const gf = c.fromBf16(g); + if (!std.math.isFinite(w) or !std.math.isFinite(gf)) { + if (c.classOf(gf) != c.classOf(w)) return error.NonFiniteMismatch; + continue; + } + if (w == 0) { + if (gf != 0) return error.OutOfBound; // signed zeros: never a small nonzero + continue; + } + const err = @abs(gf - w); + const bound = c.boundExact(swiglu_spec, w); + if (!(std.math.isFinite(bound) and bound > 0)) return error.BadBound; + const ratio = err / bound; + if (ratio > r.worst) { + r = .{ .worst = ratio, .row = i, .col = i }; + } + } + if (r.worst > 1) return error.OutOfBound; + return r; + } +}; + +// ---- RMSNorm: contract bound with the reduction budget named ---- + +// depth = 20: 2048 takes 8 adds + 8 levels (16), 256 1 + 8; the sum is positive, so gamma(20) is relative in any order. +const rms_depth = 20; +const rms_spec: c.Spec = .{ + .name = "vulkan fast_norm bf16 rms (ORACLE_QWEN35 weight*(x*inv))", + .accum = .f32, + .depth = rmsDepth, + .ops = 4, + .out = .bf16, + .served = &rms_served, +}; +fn rmsDepth(_: c.Shape) usize { + return rms_depth; +} +const rms_served = [_]c.Shape{ + .{ .rows = 1, .n = 1, .k = 2048 }, + .{ .rows = 4, .n = 1, .k = 2048 }, + .{ .rows = 2, .n = 1, .k = 2048 }, + .{ .rows = 2, .n = 1, .k = 256 }, +}; + +pub const RmsFill = enum { flat, dynamic }; + +pub const RmsCase = struct { + name: []const u8, + rows: usize, + width: usize, + fill: RmsFill = .flat, + /// The last row is all zeros: mean square 0, norm = rsqrt(eps), output exactly zero. + zero_row: bool = false, + at: RmsOffsets = .{}, +}; + +/// 256-byte-aligned element offsets (multiples of 128). +pub const RmsOffsets = struct { x: usize = 0, gamma: usize = 0, out: usize = 0 }; + +pub const rms_cases = [_]RmsCase{ + .{ .name = "1x2048", .rows = 1, .width = 2048 }, + .{ .name = "4x2048 offsets zero row", .rows = 4, .width = 2048, .zero_row = true, .at = .{ .x = 128, .gamma = 0, .out = 384 } }, + .{ .name = "2x2048 dynamic range", .rows = 2, .width = 2048, .fill = .dynamic }, + // The query/key head_norm width Qwen3.5-2B actually serves, at the same aligned offsets. + .{ .name = "2x256 head_norm offsets", .rows = 2, .width = 256, .at = .{ .x = 128, .gamma = 0, .out = 384 } }, +}; + +pub const rms_eps: f32 = 1e-6; + +pub const Rms = struct { + a: std.mem.Allocator, + case: RmsCase, + x: []u16, + gamma: []u16, + out: []u16, + /// gamma64 * (x64 * rsqrt(mean(x^2) + eps)) per element, in f64, independent of the kernel. + want: []f64, + + pub fn init(a: std.mem.Allocator, case: RmsCase, seed: u64) !Rms { + const x = try poisonedShell(a, case.at.x, case.rows * case.width); + errdefer a.free(x); + const gamma = try poisonedShell(a, case.at.gamma, case.width); + errdefer a.free(gamma); + const out = try poisonedShell(a, case.at.out, case.rows * case.width); + errdefer a.free(out); + const want = try a.alloc(f64, case.rows * case.width); + errdefer a.free(want); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (0..case.width) |i| gamma[case.at.gamma + i] = sampleBits(rng, 1); + for (0..case.rows) |row| { + const base = case.at.x + row * case.width; + var sum_sq: f64 = 0; + for (0..case.width) |i| { + const bits: u16 = if (case.zero_row and row + 1 == case.rows) 0x0000 else switch (case.fill) { + .flat => sampleBits(rng, 1), + // Large and tiny magnitudes interleaved: large squares dominate, small ones round on their scale. + .dynamic => sampleBits(rng, if (i % 4 == 0) 0x1p10 else if (i % 4 == 2) 0x1p-10 else 1), + }; + x[base + i] = bits; + sum_sq += @as(f64, c.fromBf16(bits)) * @as(f64, c.fromBf16(bits)); + } + const n: f64 = 1 / @sqrt(sum_sq / @as(f64, @floatFromInt(case.width)) + rms_eps); + for (0..case.width) |i| { + const x64: f64 = c.fromBf16(x[base + i]); + const g64: f64 = c.fromBf16(gamma[case.at.gamma + i]); + want[row * case.width + i] = g64 * (x64 * n); + } + } + for (want) |w| if (!std.math.isFinite(w)) return error.NonFiniteReference; + return .{ .a = a, .case = case, .x = x, .gamma = gamma, .out = out, .want = want }; + } + + pub fn deinit(f: Rms) void { + f.a.free(f.x); + f.a.free(f.gamma); + f.a.free(f.out); + f.a.free(f.want); + } + + // Fixed budget, NOT a spec proof: |got - want| <= c.bound with the output's own magnitude, gamma(depth) relative. + pub fn check(f: Rms, got: []const u16) !Report { + if (got.len != f.out.len) return error.LengthMismatch; + const at = f.case.at.out; + try guardIntact(got, at, f.case.rows * f.case.width); + var r: Report = .{}; + for (f.want, 0..) |w, i| { + const g = got[at + i]; + const gf = c.fromBf16(g); + if (!std.math.isFinite(gf)) return error.NonFiniteMismatch; + if (w == 0) { + if (gf != 0) return error.OutOfBound; // the zero row: ±0 only, never small nonzero + continue; + } + const err = @abs(gf - w); + const bound = c.bound(rms_spec, .{ .rows = f.case.rows, .n = 1, .k = f.case.width }, w, @abs(w)); + if (!(std.math.isFinite(bound) and bound > 0)) return error.BadBound; + const ratio = err / bound; + if (ratio > r.worst) r = .{ .worst = ratio, .row = i / f.case.width, .col = i % f.case.width }; + } + if (r.worst > 1) return error.OutOfBound; + return r; + } +}; diff --git a/zig/tests/vulkan/stage2_probe.zig b/zig/tests/vulkan/stage2_probe.zig new file mode 100644 index 00000000000..941a7aa5412 --- /dev/null +++ b/zig/tests/vulkan/stage2_probe.zig @@ -0,0 +1,363 @@ +// Vulkan stage 2 bring-up: BF16 add, SwiGLU and RMSNorm (Metal-oracle ORACLE_QWEN35) through NormMlp; exit 0/1/2/124. +const std = @import("std"); +const vk = @import("vulkan"); +const shaders = @import("stage2_shaders"); +const fx = @import("stage2_fixture.zig"); +const contract = @import("contract"); +const cx = contract.spec; // toBf16, fromBf16, the C constants + +const usage = "usage: tf-stage2-gpu [--vendor ID] [--timeout-ms MS] | --help\n" ++ + " --help show this text and exit without opening a loader or device\n" ++ + " --vendor ID refuse a device whose Vulkan vendorID differs (decimal or 0x hex; no default)\n" ++ + " --timeout-ms MS bound on every fence wait, 1 to 600000 (default 10000); a timeout ends the process\n"; + +/// Milliseconds on the awake clock since main started, printed on every milestone. +const Elapsed = struct { + io: std.Io, + start: std.Io.Timestamp, + + fn ms(e: Elapsed) i64 { + return e.start.durationTo(std.Io.Clock.awake.now(e.io)).toMilliseconds(); + } +}; + +const repeats = 3; +const planted_bits = [_]u16{ 0x7fc0, 0x7f80, 0xff80 }; +const planted = [_][]const u8{ "nan", "+inf", "-inf" }; + +/// The first output whose reference is finite (u16 bits or f64): where a planted nonfinite must be refused. +fn firstFinite(want: anytype) ?usize { + for (want, 0..) |w, i| { + const v = if (@TypeOf(w) == u16) cx.fromBf16(w) else w; + if (std.math.isFinite(v)) return i; + } + return null; +} + +pub fn main(init: std.process.Init) !u8 { + const clock: Elapsed = .{ .io = init.io, .start = std.Io.Clock.awake.now(init.io) }; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + for (args[1..]) |arg| if (std.mem.eql(u8, arg, "--help")) { + std.debug.print(usage, .{}); + return 0; + }; + var vendor: ?u32 = null; + var timeout_ms: u64 = 10_000; + var i: usize = 1; + while (i < args.len) : (i += 2) { + const value = if (i + 1 < args.len) args[i + 1] else ""; + if (std.mem.eql(u8, args[i], "--vendor")) { + vendor = std.fmt.parseInt(u32, value, 0) catch null; + if (vendor == null) break; + } else if (std.mem.eql(u8, args[i], "--timeout-ms")) { + timeout_ms = std.fmt.parseInt(u64, value, 10) catch 0; + if (timeout_ms == 0 or timeout_ms > 600_000) break; + } else break; + } else return if (try probe(init.gpa, clock, vendor, timeout_ms * std.time.ns_per_ms)) 0 else 1; + std.debug.print(usage, .{}); + return 2; +} + +fn probe(a: std.mem.Allocator, clock: Elapsed, vendor: ?u32, timeout: u64) !bool { + var rt = try vk.Runtime.open(); + defer rt.close(); + var dev = try vk.Device.open(&rt, .{ .features = .{ .storage16 = true, .int16 = true } }); + defer dev.close(); + const pr = dev.properties; + const name = std.mem.sliceTo(&pr.deviceName, 0); + const types = [_][]const u8{ "other", "integrated", "discrete", "virtual", "cpu" }; + const device_kind = if (pr.deviceType < types.len) types[pr.deviceType] else "unknown"; + const v = pr.apiVersion; + const driver = .{ std.mem.sliceTo(&dev.v12.driverName, 0), std.mem.sliceTo(&dev.v12.driverInfo, 0) }; + std.debug.print("MILESTONE device-found elapsed_ms={d} name=\"{s}\" vendor=0x{x:0>4} device=0x{x:0>4} type={s} " ++ + "api={d}.{d}.{d} driver=\"{s}\" info=\"{s}\" subgroup={d} storage16={} int16={} " ++ + "max-storage-range={d} max-buffer-size={d} fence-timeout={d}ms preserve32={} flush32={} denorm-independence={d}\n", .{ + clock.ms(), + name, + pr.vendorID, + pr.deviceID, + device_kind, + v >> 22 & 0x7f, + v >> 12 & 0x3ff, + v & 0xfff, + } ++ driver ++ .{ dev.v11.subgroupSize, dev.features.storage16, dev.features.int16, dev.properties.limits.maxStorageBufferRange, dev.v13.maxBufferSize, timeout / std.time.ns_per_ms, dev.v12.shaderDenormPreserveFloat32 != 0, dev.v12.shaderDenormFlushToZeroFloat32 != 0, dev.v12.denormBehaviorIndependence }); + if (pr.deviceType == vk.c.VK_PHYSICAL_DEVICE_TYPE_CPU) { + std.debug.print("FAIL device: a CPU (software) Vulkan device is never a GPU result\n", .{}); + return false; + } + if (vendor) |want| if (pr.vendorID != want) { + std.debug.print("FAIL device: vendor 0x{x:0>4} is not the requested 0x{x:0>4}\n", .{ pr.vendorID, want }); + return false; + }; + + var cmds = try vk.Commands.init(&dev); + defer cmds.deinit(); + var staging = try vk.Buffer.init(&dev, .staging, 64 << 10); + defer staging.deinit(); + // The add image requests preservation only when the device advertises it. + var m = try vk.norm_mlp.NormMlp.init(&dev, .{ + .norm_bf16 = &shaders.norm_bf16, + .add_bf16 = &shaders.add_bf16, + .add_preserve_bf16 = &shaders.add_preserve_bf16, + .swiglu_bf16 = &shaders.swiglu_bf16, + .norm_f32 = &shaders.norm_f32, + .norm_f32_in = &shaders.norm_f32_in, + .norm_bf16_in = &shaders.norm_bf16_in, + .add_f32 = &shaders.add_f32, + .add_wide_bf16 = &shaders.add_wide_bf16, + .add_bf16_f32 = &shaders.add_bf16_f32, + .swiglu_f32 = &shaders.swiglu_f32, + }); + defer m.deinit(); + + var passed: usize = 0; + var total: usize = 0; + for (fx.add_cases) |case| { + total += 1; + passed += @intFromBool(try addGate(a, clock, &dev, &cmds, staging, &m, case, timeout)); + } + for (fx.swiglu_cases) |case| { + total += 1; + passed += @intFromBool(try swigluGate(a, clock, &dev, &cmds, staging, &m, case, timeout)); + } + for (fx.rms_cases) |case| { + total += 1; + passed += @intFromBool(try rmsGate(a, clock, &dev, &cmds, staging, &m, case, timeout)); + } + std.debug.print("RESULT stage2 elapsed_ms={d} on \"{s}\": passed={d} total={d}\n", .{ clock.ms(), name, passed, total }); + return passed == total; +} + +fn uploaded(_: std.mem.Allocator, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, bytes: []const u8, timeout: u64) !vk.Buffer { + var b = try vk.Buffer.init(dev, .device, bytes.len); + errdefer b.deinit(); + try vk.orExit(vk.upload(cmds, b, 0, staging, bytes, timeout)); + return b; +} + +fn fetch(cmds: *vk.Commands, staging: vk.Buffer, b: vk.Buffer, out: []u8, timeout: u64) !void { + try vk.orExit(vk.download(cmds, b, 0, staging, out, timeout)); +} + +fn bitsBuffer(a: std.mem.Allocator, shell: []const u16) ![]u16 { + return a.alloc(u16, shell.len); +} + +// One bf16 binding's shell uploaded at its fixture offsets; the view is derived against the caller-resident Bound16. +const Bound16 = struct { + buf: vk.Buffer, + byte_offset: u64, + elements: u64, + + // The checked window over this binding's own buffer, derived against caller storage (a stored View would dangle). + fn view(b: *const Bound16) !vk.tensor.View { + return vk.tensor.View.init(&b.buf, .bf16, b.byte_offset, b.elements); + } +}; + +fn boundShell(a: std.mem.Allocator, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, shell: []const u16, at_elems: usize, elements: u64, timeout: u64) !Bound16 { + var buf = try uploaded(a, dev, cmds, staging, std.mem.sliceAsBytes(shell), timeout); + errdefer buf.deinit(); + return .{ .buf = buf, .byte_offset = at_elems * 2, .elements = elements }; +} + +fn checkUnchanged(a: std.mem.Allocator, cmds: *vk.Commands, staging: vk.Buffer, shell: []const u16, buf: vk.Buffer, timeout: u64) !bool { + const back = try a.alloc(u8, std.mem.sliceAsBytes(shell).len); + defer a.free(back); + try fetch(cmds, staging, buf, back, timeout); + return std.mem.indexOfDiff(u8, std.mem.sliceAsBytes(shell), back) == null; +} + +// Out = a + b bit for bit with the shader's pipeline (f32 add, one RNE bf16 store): three repeats, nonfinite refused. +fn addGate(a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, m: *vk.norm_mlp.NormMlp, case: fx.AddCase, timeout: u64) !bool { + var f = try fx.Add.init(a, case, 0x57a9); + defer f.deinit(); + f.allow_subnormal_flush = dev.v12.shaderDenormPreserveFloat32 == 0; + var in_a = try boundShell(a, dev, cmds, staging, f.a_bits, case.at.a, case.elems, timeout); + defer in_a.buf.deinit(); + var in_b = try boundShell(a, dev, cmds, staging, f.b_bits, case.at.b, case.elems, timeout); + defer in_b.buf.deinit(); + var out = try boundShell(a, dev, cmds, staging, f.out, case.at.out, case.elems, timeout); + defer out.buf.deinit(); + const bits = try bitsBuffer(a, f.out); + defer a.free(bits); + var first: ?[]u16 = null; + defer if (first) |x| a.free(x); + for (0..repeats) |rep| { + try vk.orExit(vk.upload(cmds, out.buf, 0, staging, std.mem.sliceAsBytes(f.out), timeout)); + try vk.orExit(m.add(cmds, try out.view(), try in_a.view(), try in_b.view(), timeout)); + try fetch(cmds, staging, out.buf, std.mem.sliceAsBytes(bits), timeout); + if (first) |x| { + if (std.mem.indexOfDiff(u16, x, bits)) |at| { + std.debug.print("FAIL residual-add case=\"{s}\": repeat {d} differs at element {d}\n", .{ case.name, rep, at }); + return false; + } + } else { + first = try a.dupe(u16, bits); + } + } + const mm = f.firstMismatch(bits) catch |e| { + std.debug.print("FAIL residual-add case=\"{s}\": {s}\n", .{ case.name, @errorName(e) }); + return false; + }; + if (mm) |bad| { + std.debug.print("FAIL residual-add case=\"{s}\": {s} at element {d}: a=0x{x:0>4} ({s}) b=0x{x:0>4} ({s}) want=0x{x:0>4} ({s}) got=0x{x:0>4} ({s})\n", .{ case.name, @errorName(bad.kind), bad.index, bad.a, fx.classify(bad.a), bad.b, fx.classify(bad.b), bad.want, fx.classify(bad.want), bad.got, fx.classify(bad.got) }); + return false; + } + // A planted nonfinite output at a finite-want element must be refused by the same checker. + const data_at = case.at.out + firstFinite(f.want).?; + var mutated = try a.dupe(u16, bits); + defer a.free(mutated); + for (planted_bits, planted) |pb, pname| { + mutated[data_at] = pb; + if (f.check(mutated)) |_| { + std.debug.print("FAIL residual-add case=\"{s}\": a planted {s} output passed the check\n", .{ case.name, pname }); + return false; + } else |_| {} + } + if (!try checkUnchanged(a, cmds, staging, f.a_bits, in_a.buf, timeout) or !try checkUnchanged(a, cmds, staging, f.b_bits, in_b.buf, timeout)) { + std.debug.print("FAIL residual-add case=\"{s}\": an input binding changed\n", .{case.name}); + return false; + } + std.debug.print("MILESTONE residual-add elapsed_ms={d} case=\"{s}\" elements={d} matches f32-add + RNE-store with subnormal-class-flush-allowed={}, repeats byte-identical, " ++ + "inputs and guards unchanged, planted nan/+inf/-inf refused\n", .{ clock.ms(), case.name, case.elems, f.allow_subnormal_flush and case.fill == .classes }); + return true; +} + +// Out = silu(gate) * up within the bound: three byte-identical repeats, unchanged inputs and guards, nonfinite refused. +fn swigluGate(a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, m: *vk.norm_mlp.NormMlp, case: fx.SwigluCase, timeout: u64) !bool { + const f = try fx.Swiglu.init(a, case, 0x57aa); + defer f.deinit(); + var gate = try boundShell(a, dev, cmds, staging, f.gate, case.at.a, case.elems, timeout); + defer gate.buf.deinit(); + var up = try boundShell(a, dev, cmds, staging, f.up, case.at.b, case.elems, timeout); + defer up.buf.deinit(); + var out = try boundShell(a, dev, cmds, staging, f.out, case.at.out, case.elems, timeout); + defer out.buf.deinit(); + const bits = try bitsBuffer(a, f.out); + defer a.free(bits); + var first: ?[]u16 = null; + defer if (first) |x| a.free(x); + var r: fx.Report = .{}; + for (0..repeats) |rep| { + try vk.orExit(vk.upload(cmds, out.buf, 0, staging, std.mem.sliceAsBytes(f.out), timeout)); + try vk.orExit(m.swiglu(cmds, try out.view(), try gate.view(), try up.view(), timeout)); + try fetch(cmds, staging, out.buf, std.mem.sliceAsBytes(bits), timeout); + r = f.check(bits) catch |e| { + // No cached worst/position: r belongs to the last PASSING repeat; printing it after a failure misreports. + std.debug.print("FAIL swiglu case=\"{s}\" repeat={d}: {s}\n", .{ case.name, rep, @errorName(e) }); + return false; + }; + if (first) |x| { + if (std.mem.indexOfDiff(u16, x, bits)) |at| { + std.debug.print("FAIL swiglu case=\"{s}\": repeat {d} differs at element {d}\n", .{ case.name, rep, at }); + return false; + } + } else { + first = try a.dupe(u16, bits); + } + } + var mutated = try a.dupe(u16, bits); + defer a.free(mutated); + const data_at = case.at.out + firstFinite(f.want).?; + for (planted_bits, planted) |pb, pname| { + mutated[data_at] = pb; + if (f.check(mutated)) |_| { + std.debug.print("FAIL swiglu case=\"{s}\": a planted {s} output passed the check\n", .{ case.name, pname }); + return false; + } else |_| {} + } + if (!try checkUnchanged(a, cmds, staging, f.gate, gate.buf, timeout) or !try checkUnchanged(a, cmds, staging, f.up, up.buf, timeout)) { + std.debug.print("FAIL swiglu case=\"{s}\": an input binding changed\n", .{case.name}); + return false; + } + std.debug.print("MILESTONE swiglu elapsed_ms={d} case=\"{s}\" elements={d} worst={d:.4} of bound at element {d}, " ++ + "repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused\n", .{ clock.ms(), case.name, case.elems, r.worst, r.col }); + return true; +} + +// RMSNorm through the production rmsNorm (normPlan): contract bound with the named f32 budget; the zero row is zero. +fn rmsGate(a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, m: *vk.norm_mlp.NormMlp, case: fx.RmsCase, timeout: u64) !bool { + const f = try fx.Rms.init(a, case, 0x57ab); + defer f.deinit(); + var x = try boundShell(a, dev, cmds, staging, f.x, case.at.x, case.rows * case.width, timeout); + defer x.buf.deinit(); + var gamma = try boundShell(a, dev, cmds, staging, f.gamma, case.at.gamma, case.width, timeout); + defer gamma.buf.deinit(); + var out = try boundShell(a, dev, cmds, staging, f.out, case.at.out, case.rows * case.width, timeout); + defer out.buf.deinit(); + const bits = try bitsBuffer(a, f.out); + defer a.free(bits); + var first: ?[]u16 = null; + defer if (first) |xx| a.free(xx); + var r: fx.Report = .{}; + for (0..repeats) |rep| { + try vk.orExit(vk.upload(cmds, out.buf, 0, staging, std.mem.sliceAsBytes(f.out), timeout)); + try vk.orExit(m.rmsNorm(cmds, try out.view(), try x.view(), try gamma.view(), case.rows, case.width, fx.rms_eps, timeout)); + try fetch(cmds, staging, out.buf, std.mem.sliceAsBytes(bits), timeout); + r = f.check(bits) catch |e| { + // No cached worst/position: r belongs to the last PASSING repeat; printing it after a failure misreports. + std.debug.print("FAIL rms-norm case=\"{s}\" repeat={d}: {s}\n", .{ case.name, rep, @errorName(e) }); + return false; + }; + if (first) |xx| { + if (std.mem.indexOfDiff(u16, xx, bits)) |at| { + std.debug.print("FAIL rms-norm case=\"{s}\": repeat {d} differs at element {d}\n", .{ case.name, rep, at }); + return false; + } + } else { + first = try a.dupe(u16, bits); + } + } + var mutated = try a.dupe(u16, bits); + defer a.free(mutated); + const data_at = case.at.out + r.row * case.width + r.col; + for (planted_bits, planted) |pb, pname| { + mutated[data_at] = pb; + if (f.check(mutated)) |_| { + std.debug.print("FAIL rms-norm case=\"{s}\": a planted {s} output passed the check\n", .{ case.name, pname }); + return false; + } else |_| {} + } + if (!try checkUnchanged(a, cmds, staging, f.x, x.buf, timeout) or !try checkUnchanged(a, cmds, staging, f.gamma, gamma.buf, timeout)) { + std.debug.print("FAIL rms-norm case=\"{s}\": an input binding changed\n", .{case.name}); + return false; + } + if (case.zero_row) { + const zero_at = case.at.out + (case.rows - 1) * case.width; + for (bits[zero_at .. zero_at + case.width]) |g| if (g != 0x0000 and g != 0x8000) { + std.debug.print("FAIL rms-norm case=\"{s}\": the zero row produced 0x{x:0>4}\n", .{ case.name, g }); + return false; + }; + } + std.debug.print("MILESTONE rms-norm elapsed_ms={d} case=\"{s}\" rows={d} width={d} worst={d:.4} of bound at row {d} col {d}, " ++ + "repeats byte-identical, inputs and guards unchanged, planted nan/+inf/-inf refused{s}\n", .{ clock.ms(), case.name, case.rows, case.width, r.worst, r.row, r.col, if (case.zero_row) ", zero row exactly zero" else "" }); + return true; +} + +// A host-only stand-in for a device buffer: View.init reads only `len`; the view derivation stays on caller storage. +fn hostBuffer(len: u64) vk.Buffer { + return .{ .dev = undefined, .handle = null, .memory = null, .len = len, .kind = .device, .mapped = null, .coherent = false }; +} + +test "Bound16.view binds the caller-owned buffer, never a factory stack frame" { + const b1: Bound16 = .{ .buf = hostBuffer(64), .byte_offset = 8, .elements = 28 }; + const b2: Bound16 = .{ .buf = hostBuffer(128), .byte_offset = 16, .elements = 56 }; + const v1 = try b1.view(); + const v2 = try b2.view(); + try std.testing.expect(v1.buffer == &b1.buf); + try std.testing.expect(v2.buffer == &b2.buf); + try std.testing.expect(v1.buffer != v2.buffer); + try std.testing.expectEqual(@as(u64, 8), v1.byte_offset); + try std.testing.expectEqual(@as(u64, 28), v1.elements); + try std.testing.expectEqual(vk.tensor.Kind.bf16, v1.kind); +} + +test "bound window keeps the offset, range and alignment checks" { + const buf = hostBuffer(64); + try std.testing.expectError(error.OutOfRange, vk.tensor.View.init(&buf, .bf16, 0, 0)); + try std.testing.expectError(error.Misaligned, vk.tensor.View.init(&buf, .bf16, 2, 4)); + try std.testing.expectError(error.OutOfRange, vk.tensor.View.init(&buf, .bf16, 8, 30)); // past the 64-byte end + const b: Bound16 = .{ .buf = hostBuffer(64), .byte_offset = 8, .elements = 28 }; + _ = try b.view(); // the derived window is the checked one +} diff --git a/zig/tests/vulkan/stage3_fixture.zig b/zig/tests/vulkan/stage3_fixture.zig new file mode 100644 index 00000000000..e6e4d5f1071 --- /dev/null +++ b/zig/tests/vulkan/stage3_fixture.zig @@ -0,0 +1,311 @@ +const std = @import("std"); +const contract = @import("contract"); +const c = contract.spec; +// The cases live in stage3_fixture_cases.zig; the tests stay here so every importer of this file runs them. +const cases = @import("stage3_fixture_cases.zig"); + +pub const tail = cases.tail; +pub const poison16 = cases.poison16; +pub const heads = cases.heads; +pub const kv_heads = cases.kv_heads; +pub const head_dim = cases.head_dim; +pub const rotary = cases.rotary; +pub const theta = cases.theta; +pub const scale = cases.scale; +pub const Report = cases.Report; +pub const budget = cases.budget; +pub const sampleDyadic = cases.sampleDyadic; +pub const ropeRef = cases.ropeRef; +pub const RopeCase = cases.RopeCase; +pub const rope_cases = cases.rope_cases; +pub const Rope = cases.Rope; +pub const CacheCase = cases.CacheCase; +pub const cache_case = cases.cache_case; +pub const Cache = cases.Cache; +pub const AttendCase = cases.AttendCase; +pub const attend_case = cases.attend_case; +pub const visibleFrom = cases.visibleFrom; +pub const Attend = cases.Attend; +pub const GateCase = cases.GateCase; +pub const gate_case = cases.gate_case; +pub const GateDiag = cases.GateDiag; +pub const Gate = cases.Gate; + +const testing = std.testing; + +test "sdpa per-row visible counts and per-row poison windows" { + const ac = attend_case; + // prefix=2, rows=3: each row's visible key count is prefix+r+1 -> 3,4,5 + try testing.expectEqual([3]u64{ 3, 4, 5 }, .{ visibleFrom(ac, 0), visibleFrom(ac, 1), visibleFrom(ac, 2) }); + for (0..ac.rows) |r| { + const from = visibleFrom(ac, @intCast(r)); + // Each row's visible key is required; its poison window is fresh and nonempty. + try testing.expect(from > ac.prefix); + try testing.expect(from <= ac.capacity); + } +} + +test "rope reference matches the oracle half-split form on a hand vector" { + // Unit input yields cos/sin in rotated halves and zero passthrough tail. + var row: [head_dim]u16 = @splat(0); + for (0..32) |i| row[i] = c.toBf16(1.0); + var out: [head_dim]f64 = undefined; + ropeRef(&row, 0, &out); + for (0..32) |i| { + try testing.expectEqual(@as(f64, 1), out[i]); // cos(0) = 1 + try testing.expectEqual(@as(f64, 0), out[i + 32]); // sin(0) = 0 + } + try testing.expectEqual(@as(f64, 0), out[128]); // passthrough tail + // pos 0.5, theta^(-1/32) at i=1: out[1] = cos(0.5*inv), out[33] = sin(0.5*inv). + const pos: f64 = 0.5; + ropeRef(&row, pos, &out); + const inv = std.math.pow(f64, theta, -1.0 / 32.0); + try testing.expectApproxEqAbs(@cos(pos * inv), out[1], 1e-12); + try testing.expectApproxEqAbs(@sin(pos * inv), out[33], 1e-12); +} + +test "budget grows with depth and stays finite on zeros" { + try testing.expect(budget(300, 6, 2, 1, 1) > budget(2, 6, 1, 1, 1)); + try testing.expect(budget(2, 6, 1, 0, 0) > 0); +} + +test "sampleDyadic is exactly bf16" { + var prng = std.Random.DefaultPrng.init(7); + const rng = prng.random(); + for (0..64) |_| { + const bits = sampleDyadic(rng, -5, 3); + const v = c.fromBf16(bits); + try testing.expectEqual(v, c.fromBf16(c.toBf16(v))); + } +} + +// Negative tests require rejection of wrong finite/nonfinite outputs and detect tail/row bit tampering. +test "rope, cache and gate checkers refuse planted outputs" { + const a = testing.allocator; + { + const f = try Rope.init(a, rope_cases[0], 0x11); + defer f.deinit(); + const good = try a.alloc(u16, f.want.len); + defer a.free(good); + for (f.want, 0..) |w, i| good[i] = c.toBf16(@floatCast(w)); + const rep = try f.check(good); // the honest output passes + try testing.expect(rep.worst <= 1); + good[0] = 0x4f00; // 2^31: an absolutely wrong finite, far past any budget + try testing.expectError(error.OverBound, f.check(good)); + good[0] = 0x7fc0; + try testing.expectError(error.NonFiniteOutput, f.check(good)); + // passthrough tail tamper + const f2 = try Rope.init(a, rope_cases[0], 0x12); + defer f2.deinit(); + const good2 = try a.alloc(u16, f2.want.len); + defer a.free(good2); + for (f2.want, 0..) |w, i| good2[i] = c.toBf16(@floatCast(w)); + good2[@intCast(rotary * 2)] ^= 1; + try testing.expectError(error.PassthroughTailChanged, f2.check(good2)); + } + { + const f = try Cache.init(a, cache_case, 0x13); + defer f.deinit(); + const kbits = try a.alloc(u16, f.keys_shell.len - tail); + defer a.free(kbits); + for (kbits) |*x| x.* = poison16; + for (f.active, 0..) |act, i| if (act) { + kbits[i] = c.toBf16(@floatCast(f.want_keys[i])); + }; + const krep = try f.checkKeys(kbits); + try testing.expect(krep.worst <= 1); + // a POISON row tamper must fail: + var inactive: ?usize = null; + for (f.active, 0..) |act, i| if (!act) { + inactive = i; + break; + }; + kbits[inactive.?] = 0; + try testing.expectError(error.PoisonRowTouched, f.checkKeys(kbits)); + // V transport tamper + const vbits = try a.dupe(u16, f.want_value_bits); + defer a.free(vbits); + var v_active: ?usize = null; + for (f.want_value_bits, 0..) |w, i| if (w != poison16) { + v_active = i; + break; + }; + vbits[v_active.?] ^= 1; + try testing.expectError(error.ValueTransportChanged, f.checkValueBits(vbits)); + } + { + const f = try Gate.init(a, gate_case, 0x14); + defer f.deinit(); + const good = try a.alloc(u16, @intCast(f.case.rows * 2048)); + defer a.free(good); + @memset(good, 0); + // the gate reference is derived from the shells via the shared element walk + for (0..@intCast(f.case.rows)) |row| { + for (0..2048) |d| { + good[row * 2048 + d] = c.toBf16(@floatCast(f.element(good, row, d).want)); + } + } + const grep2 = try f.check(good); + try testing.expect(grep2.worst <= 1); + good[0] = 0xff80; + try testing.expectError(error.NonFiniteOutput, f.check(good)); + } +} + +test "all-poison output (dead kernel on poisoned init) is rejected by checkers" { + const a = testing.allocator; + // Poison16 fills and uploads output before each dispatch; absent or partial writes must fail. + { + const f = try Rope.init(a, rope_cases[0], 0x21); + defer f.deinit(); + const shell = try a.alloc(u16, f.want.len + tail); + defer a.free(shell); + @memset(shell, poison16); + try testing.expectError(error.NonFiniteOutput, f.check(shell[0..f.want.len])); // dead kernel + // partial write: first row correct, rest poison + for (0..f.row_elems) |i| shell[i] = c.toBf16(@floatCast(f.want[i])); + try testing.expectError(error.NonFiniteOutput, f.check(shell[0..f.want.len])); + } + { + const f = try Gate.init(a, gate_case, 0x22); + defer f.deinit(); + const out = try a.alloc(u16, @intCast(f.case.rows * 2048)); + defer a.free(out); + @memset(out, poison16); + try testing.expectError(error.NonFiniteOutput, f.check(out)); + } +} + +test "fixture initialization fails cleanly under allocation failure" { + var failing = std.testing.FailingAllocator.init(testing.allocator, .{ .fail_index = 1 }); + try testing.expectError(error.OutOfMemory, Rope.init(failing.allocator(), rope_cases[0], 0x15)); + try testing.expectError(error.OutOfMemory, Cache.init(failing.allocator(), cache_case, 0x16)); + try testing.expectError(error.OutOfMemory, Gate.init(failing.allocator(), gate_case, 0x17)); +} + +// Prefix=2, rows=3, cap=8: hidden-key poisoning preserves f64 output; a one-key leak is rejected. +test "causal sdpa: poisoned future keys are invisible; a one-key leak is caught" { + const a = testing.allocator; + const ac = attend_case; + const row_elems: usize = @intCast(heads * head_dim); + const cache_elems: usize = @intCast(kv_heads * ac.capacity * head_dim); + const cap: usize = @intCast(ac.capacity); + var prng = std.Random.DefaultPrng.init(0x31); + const rng = prng.random(); + for (0..ac.rows) |r| { + // Fresh, finite, nontrivial data each row (no cross-row accumulation). + const q = try a.alloc(u16, row_elems); + defer a.free(q); + for (q) |*x| x.* = sampleDyadic(rng, -5, 3); + const keys = try a.alloc(u16, cache_elems); + defer a.free(keys); + const values = try a.alloc(u16, cache_elems); + defer a.free(values); + for (keys) |*x| x.* = sampleDyadic(rng, -5, 3); + for (values) |*x| x.* = sampleDyadic(rng, -5, 3); + + // Copy with every key at or after row r's visibility window poisoned. + const from: usize = @intCast(visibleFrom(ac, @intCast(r))); + const leaked_keys = try a.dupe(u16, keys); + defer a.free(leaked_keys); + for (0..@intCast(kv_heads)) |h| { + const head_base = @as(usize, @intCast(h)) * cap; + @memset(leaked_keys[(head_base + from) * @as(usize, @intCast(head_dim)) ..][0 .. (cap - from) * @as(usize, @intCast(head_dim))], poison16); + } + + // Baseline: this row as a single-row query over the clean cache. + const one_row = AttendCase{ .name = "baseline row", .rows = 1, .prefix = ac.prefix + @as(u64, @intCast(r)), .capacity = ac.capacity }; + const base = try Attend.init(a, one_row, q, keys, values, 0); + defer base.deinit(); + // Same single-row query over the poisoned cache: identical reference. + const pbase = try Attend.init(a, one_row, q, leaked_keys, values, 0); + defer pbase.deinit(); + try testing.expectEqualSlices(f64, base.want, pbase.want); + + // The honest bf16 output passes BOTH production checkers. + const out = try a.alloc(u16, base.want.len); + defer a.free(out); + for (base.want, 0..) |w, i| out[i] = c.toBf16(@floatCast(w)); + const rep = try pbase.check(out); + try testing.expect(rep.worst <= 1); + + // The same Attend.init leaks one poisoned key; its nonfinite reference must differ and fail. + const leak_case = AttendCase{ .name = "leak row", .rows = 1, .prefix = ac.prefix + @as(u64, @intCast(r)) + 1, .capacity = ac.capacity }; + const leak = try Attend.init(a, leak_case, q, leaked_keys, values, 0); + defer leak.deinit(); + var has_nonfinite = false; + for (leak.want) |w| { + if (!std.math.isFinite(w)) has_nonfinite = true; + } + try testing.expect(has_nonfinite); // the leak truly saw the poison + try testing.expect(!std.meta.eql(base.want[0], leak.want[0])); // differs from baseline + const leak_out = try a.alloc(u16, leak.want.len); + defer a.free(leak_out); + for (leak.want, 0..) |w, i| leak_out[i] = c.toBf16(@floatCast(w)); + try testing.expectError(error.NonFiniteOutput, leak.check(leak_out)); + try testing.expectError(error.NonFiniteOutput, base.check(leak_out)); // causal checker rejects the leaked output + } +} + +// Zeroed lane 0 was the GPU bug: all-zero output must fail OverBound and report its first element. +test "all-zero gate output is refused OverBound and the diagnostic names element 0" { + const a = testing.allocator; + const f = try Gate.init(a, gate_case, 0x25); + defer f.deinit(); + // Element 0 uses gate=1 (sigmoid ~0.731, nonzero bf16) and ctx=1; dyadic shells stay fixed. + f.qraw_bits[256] = c.toBf16(1.0); + f.ctx_bits[0] = c.toBf16(1.0); + const zeros = try a.alloc(u16, @intCast(f.case.rows * 2048)); + defer a.free(zeros); + @memset(zeros, 0); + try testing.expectError(error.OverBound, f.check(zeros)); + const diag = f.firstViolation(zeros).?; + try testing.expectEqual(@as(usize, 0), diag.index); + try testing.expectEqual(@as(usize, 0), diag.row); + try testing.expectEqual(@as(usize, 0), diag.d); + try testing.expectEqual(@as(usize, 0), diag.head); + try testing.expectEqual(@as(f64, 0), diag.got); + try testing.expectEqual(c.toBf16(1.0), diag.ctx_bits); + try testing.expectEqual(f.qraw_bits[256], diag.gate_bits); + try testing.expectEqual(@as(f64, 1), diag.ctx); + try testing.expect(diag.sigmoid_bf16 > 0.5 and diag.sigmoid_bf16 < 1); + try testing.expectEqual(diag.ctx * diag.sigmoid_bf16, diag.want); + try testing.expectEqual(diag.want, @abs(diag.got - diag.want)); + try testing.expect(diag.ratio > 1); +} + +// Diagnostics report the planted first bad element without changing OverBound acceptance. +test "gate diagnostic reports a planted bad first element with unchanged OverBound acceptance" { + const a = testing.allocator; + const f = try Gate.init(a, gate_case, 0x24); + defer f.deinit(); + f.qraw_bits[256] = c.toBf16(1.0); + f.ctx_bits[0] = c.toBf16(1.0); + const good = try a.alloc(u16, @intCast(f.case.rows * 2048)); + defer a.free(good); + @memset(good, 0); + for (0..@intCast(f.case.rows)) |row| { + for (0..2048) |d| { + good[row * 2048 + d] = c.toBf16(@floatCast(f.element(good, row, d).want)); + } + } + const rep = try f.check(good); // the honest output passes + try testing.expect(rep.worst <= 1); + good[0] = 0x4f00; // 2^31 at the first element: absolutely wrong + try testing.expectError(error.OverBound, f.check(good)); + const diag = f.firstViolation(good).?; + try testing.expectEqual(@as(usize, 0), diag.index); + try testing.expectEqual(@as(f64, 2147483648.0), diag.got); + try testing.expect(diag.ratio > 1); + try testing.expectEqual(f.qraw_bits[256], diag.gate_bits); + try testing.expectEqual(c.toBf16(1.0), diag.ctx_bits); +} + +test "gate sigmoid widen: bf16 0x3f00 widens once to 0.5; the old double shift gives 0" { + const bits: u32 = 0x3f00; // what bf16_bits(0.5) produces + const once: f32 = @bitCast(bits << 16); // bf16_lane(word, 0): the correct widen + try testing.expectEqual(@as(f32, 0.5), once); + const annihilated: u64 = @as(u64, bits) << 32; // (bits << 16) << 16 as the GPU takes it + const twice: f32 = @bitCast(@as(u32, @truncate(annihilated))); + try testing.expectEqual(@as(f32, 0.0), twice); +} diff --git a/zig/tests/vulkan/stage3_fixture_cases.zig b/zig/tests/vulkan/stage3_fixture_cases.zig new file mode 100644 index 00000000000..4695f22ccc5 --- /dev/null +++ b/zig/tests/vulkan/stage3_fixture_cases.zig @@ -0,0 +1,459 @@ +// The stage 3 fixture's cases and checkers: RoPE, the paged KV cache, causal SDPA and the output gate. +const std = @import("std"); +const contract = @import("contract"); +const c = contract.spec; + +pub const tail = 16; +pub const poison16: u16 = 0xffff; + +pub const heads: u64 = 8; +pub const kv_heads: u64 = 2; +pub const head_dim: u64 = 256; +pub const rotary: u64 = 32; +pub const theta: f64 = 1e7; +pub const scale: f64 = 1.0 / 16.0; + +/// The worst output: |got - want| OVER its budget (a ratio; >1 fails), and where. +pub const Report = struct { worst: f64 = 0, where: usize = 0 }; + +// Budget: depth f32 sum, then `ops` f32 and `mids` bf16 roundings, with one final bf16 store; never sum output squares. +pub fn budget(depth: f64, ops: f64, mids: f64, want: f64, mag: f64) f64 { + const u: f64 = 0x1p-24; + const g = if (depth > 1) depth * u / (1 - depth * u) else 0; + const value_scale = @max(@abs(want), mag); + const work = g * mag + (ops * u + mids * 0x1p-8) * value_scale; + return work + 0x1p-8 * (@abs(want) + work) + 0x1p-133; +} + +/// ±m * 2^e, m in {1, 1.5}, e in the stated range: exactly bf16-representable. +pub fn sampleDyadic(rng: std.Random, lo: i32, hi: i32) u16 { + const e = rng.intRangeAtMost(i32, lo, hi); + const m: f32 = if (rng.boolean()) 1.0 else 1.5; + return c.toBf16(m * std.math.pow(f32, 2, @floatFromInt(e)) * (if (rng.boolean()) @as(f32, 1) else -1)); +} + +fn poisonedShell(a: std.mem.Allocator, len: usize) ![]u16 { + const buf = try a.alloc(u16, len + tail); + @memset(buf, poison16); + return buf; +} + +/// The oracle rope in f64 for one 256-wide row at absolute position `pos`. +pub fn ropeRef(row: []const u16, pos: f64, out: []f64) void { + for (0..rotary) |i| { + const inv = std.math.pow(f64, theta, -@as(f64, @floatFromInt(i)) / 32.0); + const angle = pos * inv; + const cos_a = @cos(angle); + const sin_a = @sin(angle); + const x1 = c.fromBf16(row[i]); + const x2 = c.fromBf16(row[i + 32]); + out[i] = x1 * cos_a - x2 * sin_a; + out[i + 32] = x1 * sin_a + x2 * cos_a; + } + for (rotary * 2..head_dim) |d| out[d] = c.fromBf16(row[d]); +} + +pub const RopeCase = struct { name: []const u8, rows: u64, base: i64 }; + +pub const rope_cases = [_]RopeCase{ + .{ .name = "queries rows4 base0", .rows = 4, .base = 0 }, + .{ .name = "queries rows2 base5", .rows = 2, .base = 5 }, +}; + +// RoPE gate checks rotated values against f64; uploaded poison tails remain, and passthrough bits match exactly. +pub const Rope = struct { + a: std.mem.Allocator, + case: RopeCase, + q_bits: []u16, + want: []f64, + /// Elements per row (heads * head_dim). + row_elems: usize, + + pub fn init(a: std.mem.Allocator, case: RopeCase, seed: u64) !Rope { + const row_elems: usize = @intCast(heads * head_dim); + const q_bits = try poisonedShell(a, case.rows * row_elems); + errdefer a.free(q_bits); + const want = try a.alloc(f64, case.rows * row_elems); + errdefer a.free(want); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + var ref: [head_dim]f64 = undefined; + for (0..@intCast(case.rows)) |r| { + const row = q_bits[r * row_elems ..][0..row_elems]; + for (0..heads) |h| { + const head_row = row[@intCast(h * head_dim)..][0..@intCast(head_dim)]; + for (head_row) |*x| x.* = sampleDyadic(rng, -5, 3); + ropeRef(head_row, @floatFromInt(case.base + @as(i64, @intCast(r))), &ref); + @memcpy(want[r * row_elems + @as(usize, @intCast(h * head_dim)) ..][0..@intCast(head_dim)], &ref); + } + } + return .{ .a = a, .case = case, .q_bits = q_bits, .want = want, .row_elems = row_elems }; + } + + pub fn deinit(f: Rope) void { + f.a.free(f.q_bits); + f.a.free(f.want); + } + + pub fn check(f: Rope, bits: []const u16) !Report { + var r: Report = .{}; + for (f.want, 0..) |w, i| { + const got = c.fromBf16(bits[i]); + const d = @abs(got - w); + const bound = budget(2, 6, 1, w, @abs(w)); + if (!std.math.isFinite(got)) return error.NonFiniteOutput; + if (d > bound) return error.OverBound; + const ratio = d / bound; + if (ratio > r.worst) r = .{ .worst = ratio, .where = i }; + // passthrough tail: bit-exact + if (i % @as(usize, @intCast(head_dim)) >= @as(usize, @intCast(rotary * 2))) { + if (bits[i] != f.q_bits[i]) return error.PassthroughTailChanged; + } + } + return r; + } +}; + +pub const CacheCase = struct { name: []const u8, rows: u64, base: u64, capacity: u64 }; + +pub const cache_case = CacheCase{ .name = "cache prefix2 base2 rows3 cap8", .rows = 3, .base = 2, .capacity = 8 }; + +// CacheKV plants prefix rows in [kv_head,capacity,head_dim]; K matches f64 RoPE, inactive bits persist, V is exact. +pub const Cache = struct { + a: std.mem.Allocator, + case: CacheCase, + knorm_bits: []u16, + v_bits: []u16, + keys_shell: []u16, + values_shell: []u16, + want_keys: []f64, + // Exact bits required for every position: planted prefix bits, the V transport bits, or poison for untouched rows. + want_key_bits: []u16, + want_value_bits: []u16, + /// True where the operator writes (or the prefix was planted): budget-checked. + active: []bool, + + /// (h, pos, d) -> element index in the [kv_head, capacity, head_dim] cache. + fn cacheIndex(case: CacheCase, h: usize, pos: usize) usize { + return (h * @as(usize, @intCast(case.capacity)) + pos) * @as(usize, @intCast(head_dim)); + } + + pub fn init(a: std.mem.Allocator, case: CacheCase, seed: u64) !Cache { + const row_elems: usize = @intCast(kv_heads * head_dim); + const cache_elems: usize = @intCast(kv_heads * case.capacity * head_dim); + const knorm_bits = try poisonedShell(a, case.rows * row_elems); + errdefer a.free(knorm_bits); + const v_bits = try poisonedShell(a, case.rows * row_elems); + errdefer a.free(v_bits); + const keys_shell = try poisonedShell(a, cache_elems); + errdefer a.free(keys_shell); + const values_shell = try poisonedShell(a, cache_elems); + errdefer a.free(values_shell); + const want_keys = try a.alloc(f64, cache_elems); + errdefer a.free(want_keys); + const want_key_bits = try a.alloc(u16, cache_elems); + errdefer a.free(want_key_bits); + const want_value_bits = try a.alloc(u16, cache_elems); + errdefer a.free(want_value_bits); + const active = try a.alloc(bool, cache_elems); + errdefer a.free(active); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + // Everything starts inactive and poison; the loops below claim slots. + for (0..cache_elems) |i| { + want_keys[i] = 0; + want_key_bits[i] = poison16; + want_value_bits[i] = poison16; + active[i] = false; + } + // Prefix positions (both kv heads): planted data, bit-exact want. + for (0..@intCast(case.base)) |pos| { + for (0..@intCast(kv_heads)) |h| { + const base_idx = cacheIndex(case, h, pos); + for (0..@intCast(head_dim)) |d| { + const kb = sampleDyadic(rng, -5, 3); + const vb = sampleDyadic(rng, -5, 3); + keys_shell[base_idx + d] = kb; + values_shell[base_idx + d] = vb; + want_keys[base_idx + d] = c.fromBf16(kb); + want_key_bits[base_idx + d] = kb; + want_value_bits[base_idx + d] = vb; + active[base_idx + d] = true; + } + } + } + // K rows match f64 RoPE at base+r; V rows transport exact bits. + var ref: [head_dim]f64 = undefined; + for (0..@intCast(case.rows)) |r| { + for (0..@intCast(kv_heads)) |h| { + const row = knorm_bits[r * row_elems + h * @as(usize, @intCast(head_dim)) ..][0..@intCast(head_dim)]; + for (row) |*x| x.* = sampleDyadic(rng, -5, 3); + const vrow = v_bits[r * row_elems + h * @as(usize, @intCast(head_dim)) ..][0..@intCast(head_dim)]; + for (vrow) |*x| x.* = sampleDyadic(rng, -5, 3); + ropeRef(row, @floatFromInt(@as(i64, @intCast(case.base + r))), &ref); + const dst = cacheIndex(case, h, @intCast(case.base + r)); + @memcpy(want_keys[dst..][0..@intCast(head_dim)], &ref); + for (0..@intCast(head_dim)) |d| { + want_value_bits[dst + d] = vrow[d]; + active[dst + d] = true; + } + } + } + return .{ + .a = a, + .case = case, + .knorm_bits = knorm_bits, + .v_bits = v_bits, + .keys_shell = keys_shell, + .values_shell = values_shell, + .want_keys = want_keys, + .want_key_bits = want_key_bits, + .want_value_bits = want_value_bits, + .active = active, + }; + } + + pub fn deinit(f: Cache) void { + f.a.free(f.knorm_bits); + f.a.free(f.v_bits); + f.a.free(f.keys_shell); + f.a.free(f.values_shell); + f.a.free(f.want_keys); + f.a.free(f.want_key_bits); + f.a.free(f.want_value_bits); + f.a.free(f.active); + } + + // Active rows match f64 RoPE; inactive poison rows preserve uploaded bits. + pub fn checkKeys(f: Cache, bits: []const u16) !Report { + var r: Report = .{}; + for (f.active, 0..) |is_active, i| { + if (!is_active) { + if (bits[i] != f.want_key_bits[i]) return error.PoisonRowTouched; + continue; + } + const w = f.want_keys[i]; + const got = c.fromBf16(bits[i]); + if (!std.math.isFinite(got)) return error.NonFiniteOutput; + const d = @abs(got - w); + const bound = budget(2, 6, 1, w, @abs(w)); + if (d > bound) return error.OverBound; + const ratio = d / bound; + if (ratio > r.worst) r = .{ .worst = ratio, .where = i }; + } + return r; + } + + /// V rows and every untouched cache row must be bit-exact (transport or poison). + pub fn checkValueBits(f: Cache, bits: []const u16) !void { + for (f.want_value_bits, 0..) |w, i| { + if (w == poison16) { + if (bits[i] != poison16) return error.PoisonRowTouched; + } else if (bits[i] != w) return error.ValueTransportChanged; + } + } +}; + +pub const AttendCase = struct { name: []const u8, rows: u64, prefix: u64, capacity: u64 }; + +pub const attend_case = AttendCase{ .name = "sdpa prefix2 rows3 cap8", .rows = 3, .prefix = 2, .capacity = 8 }; + +// Each query row must ignore keys from its first non-visible index onward. +pub fn visibleFrom(case: AttendCase, r: u64) u64 { + return case.prefix + r + 1; +} + +// Causal SDPA uses uploaded cache and rotated query bits; its f64 reference adds no RoPE error. +pub const Attend = struct { + a: std.mem.Allocator, + case: AttendCase, + q_bits: []u16, + want: []f64, + /// Visible key count per row (prefix + row + 1), for the known-counts report. + visible: []usize, + max_score: f64, + max_value: f64, + + pub fn init(a: std.mem.Allocator, case: AttendCase, q_bits: []const u16, keys: []const u16, values: []const u16, seed: u64) !Attend { + _ = seed; + const row_elems: usize = @intCast(heads * head_dim); + const q_bits_copy = try a.dupe(u16, q_bits[0 .. case.rows * row_elems]); + errdefer a.free(q_bits_copy); + const want = try a.alloc(f64, case.rows * row_elems); + errdefer a.free(want); + const visible = try a.alloc(usize, @intCast(case.rows)); + errdefer a.free(visible); + var max_score: f64 = 0; + var max_value: f64 = 0; + for (0..@intCast(case.rows)) |r| { + const vis: usize = @intCast(case.prefix + @as(u64, @intCast(r)) + 1); + visible[r] = vis; + for (0..heads) |h| { + const kh = h / 4; // repeats = heads / kv_heads + const qrow = q_bits[r * row_elems + h * @as(usize, @intCast(head_dim)) ..][0..@intCast(head_dim)]; + var scores: [8]f64 = undefined; // vis <= prefix+rows <= 8 in the fixture case + var m: f64 = -std.math.floatMax(f64); + for (0..vis) |j| { + const kbase = (@as(usize, @intCast(kh)) * @as(usize, @intCast(case.capacity)) + j) * @as(usize, @intCast(head_dim)); + var dot: f64 = 0; + for (0..@intCast(head_dim)) |d| { + dot += @as(f64, c.fromBf16(qrow[d])) * c.fromBf16(keys[kbase + d]); + } + scores[j] = dot * scale; + m = @max(m, scores[j]); + max_score = @max(max_score, @abs(scores[j])); + } + var denom: f64 = 0; + for (scores[0..vis]) |s| denom += @exp(s - m); + for (0..@intCast(head_dim)) |d| { + var acc: f64 = 0; + for (0..vis) |j| { + const vbase = (@as(usize, @intCast(kh)) * @as(usize, @intCast(case.capacity)) + j) * @as(usize, @intCast(head_dim)); + acc += (@exp(scores[j] - m) / denom) * c.fromBf16(values[vbase + d]); + max_value = @max(max_value, @abs(c.fromBf16(values[vbase + d]))); + } + want[r * row_elems + h * @as(usize, @intCast(head_dim)) + d] = acc; + } + } + } + return .{ .a = a, .case = case, .q_bits = q_bits_copy, .want = want, .visible = visible, .max_score = max_score, .max_value = max_value }; + } + + pub fn deinit(f: Attend) void { + f.a.free(f.q_bits); + f.a.free(f.want); + f.a.free(f.visible); + } + + pub fn check(f: Attend, bits: []const u16) !Report { + var r: Report = .{}; + for (f.want, 0..) |w, i| { + const got = c.fromBf16(bits[i]); + if (!std.math.isFinite(got)) return error.NonFiniteOutput; + const d = @abs(got - w); + const bound = budget(300, 6, 2, w, f.max_score + f.max_value); + if (d > bound) return error.OverBound; + const ratio = d / bound; + if (ratio > r.worst) r = .{ .worst = ratio, .where = i }; + } + return r; + } +}; + +pub const GateCase = struct { name: []const u8, rows: u64 }; + +pub const gate_case = GateCase{ .name = "gate rows4", .rows = 4 }; + +// Gate elements retain inputs, f64 reference/intermediates, budget and output for shared checks and diagnostics. +pub const GateDiag = struct { + index: usize, + row: usize, + d: usize, + head: usize, + /// Uploaded input words: ctx at ctx_bits, the raw qraw gate at gate_bits. + ctx_bits: u16, + gate_bits: u16, + ctx: f64, + gate: f64, + // Sigmoid intermediates are f64, f32 before one RNE bf16 round, and the resulting bf16 value. + sigmoid_f64: f64, + sigmoid_f32: f64, + sigmoid_bf16_bits: u16, + sigmoid_bf16: f64, + want: f64, + got: f64, + bound: f64, + distance: f64, + ratio: f64, +}; + +// The interleaved output gate: qraw row stride 4096, head h's gate at h*512 + 256 + d within the row. +pub const Gate = struct { + a: std.mem.Allocator, + case: GateCase, + qraw_bits: []u16, + ctx_bits: []u16, + out_bits: []u16, + + pub fn init(a: std.mem.Allocator, case: GateCase, seed: u64) !Gate { + const qraw = try poisonedShell(a, case.rows * 4096); + errdefer a.free(qraw); + const ctx = try poisonedShell(a, case.rows * 2048); + errdefer a.free(ctx); + const out = try poisonedShell(a, case.rows * 2048); + errdefer a.free(out); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + // Data region only: the poison tail must survive as the overrun guard. + for (qraw[0 .. case.rows * 4096]) |*x| x.* = sampleDyadic(rng, -8, 8); + for (ctx[0 .. case.rows * 2048]) |*x| x.* = sampleDyadic(rng, -5, 3); + return .{ .a = a, .case = case, .qraw_bits = qraw, .ctx_bits = ctx, .out_bits = out }; + } + + pub fn deinit(f: Gate) void { + f.a.free(f.qraw_bits); + f.a.free(f.ctx_bits); + f.a.free(f.out_bits); + } + + // One element's reference, budget and output value; the checker and the diagnostic both walk this. + pub fn element(f: Gate, bits: []const u16, row: usize, d: usize) GateDiag { + const head = d / 256; + const gate_bits = f.qraw_bits[row * 4096 + head * 512 + 256 + (d % 256)]; + const gate = c.fromBf16(gate_bits); + // Metal rounds sigmoid to bf16 then widens; budget covers f32/f64 exp differences. + const sigmoid_f64 = 1.0 / (1.0 + @exp(-@as(f64, gate))); + const sigmoid_f32: f32 = @floatCast(sigmoid_f64); + const sigmoid_bf16_bits = c.toBf16(sigmoid_f32); + const sigmoid_bf16 = c.fromBf16(sigmoid_bf16_bits); + const ctx = c.fromBf16(f.ctx_bits[row * 2048 + d]); + const want = ctx * sigmoid_bf16; + const got = c.fromBf16(bits[row * 2048 + d]); + const distance: f32 = @abs(got - want); + const bound = budget(1, 3, 1, want, @abs(want)); + return .{ + .index = row * 2048 + d, + .row = row, + .d = d, + .head = head, + .ctx_bits = f.ctx_bits[row * 2048 + d], + .gate_bits = gate_bits, + .ctx = ctx, + .gate = gate, + .sigmoid_f64 = sigmoid_f64, + .sigmoid_f32 = sigmoid_f32, + .sigmoid_bf16_bits = sigmoid_bf16_bits, + .sigmoid_bf16 = sigmoid_bf16, + .want = want, + .got = got, + .bound = bound, + .distance = distance, + .ratio = @as(f64, distance) / bound, + }; + } + + /// The checker recomputes the reference from the uploaded shells. + pub fn check(f: Gate, bits: []const u16) !Report { + var r: Report = .{}; + for (0..@intCast(f.case.rows)) |row| { + for (0..2048) |d| { + const e = f.element(bits, row, d); + if (!std.math.isFinite(e.got)) return error.NonFiniteOutput; + if (e.distance > e.bound) return error.OverBound; + if (e.ratio > r.worst) r = .{ .worst = e.ratio, .where = e.index }; + } + } + return r; + } + + // Return first failing element with probe context; null means accepted. + pub fn firstViolation(f: Gate, bits: []const u16) ?GateDiag { + for (0..@intCast(f.case.rows)) |row| { + for (0..2048) |d| { + const e = f.element(bits, row, d); + if (!std.math.isFinite(e.got) or e.distance > e.bound) return e; + } + } + return null; + } +}; diff --git a/zig/tests/vulkan/stage3_probe.zig b/zig/tests/vulkan/stage3_probe.zig new file mode 100644 index 00000000000..cb80e958188 --- /dev/null +++ b/zig/tests/vulkan/stage3_probe.zig @@ -0,0 +1,409 @@ +// Stage 3 GPU bring-up: RoPE/KV/GQA SDPA/gate; exits 0/1/2/124; no model, perf or cross-backend bitexact claim. +const std = @import("std"); +const vk = @import("vulkan"); +const shaders = @import("stage3_shaders"); +const ops_shaders = @import("stage1_shaders"); +const fx = @import("stage3_fixture.zig"); +const contract = @import("contract"); +const cx = contract.spec; + +const usage = "usage: tf-stage3-gpu [--vendor ID] [--timeout-ms MS] | --help\n" ++ + " --help show this text and exit without opening a loader or device\n" ++ + " --vendor ID refuse a device whose Vulkan vendorID differs (decimal or 0x hex; no default)\n" ++ + " --timeout-ms MS bound on every fence wait, 1 to 600000 (default 10000); a timeout ends the process\n"; + +const Elapsed = struct { + io: std.Io, + start: std.Io.Timestamp, + + fn ms(e: Elapsed) i64 { + return e.start.durationTo(std.Io.Clock.awake.now(e.io)).toMilliseconds(); + } +}; + +const repeats = 3; +// Large FINITE bf16 (≈ max) poisons non-visible keys; uncausal reads produce ±inf/huge and fail checks. +const finite_poison: u16 = 0x7f7f; + +pub fn main(init: std.process.Init) !u8 { + const clock: Elapsed = .{ .io = init.io, .start = std.Io.Clock.awake.now(init.io) }; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + for (args[1..]) |arg| if (std.mem.eql(u8, arg, "--help")) { + std.debug.print(usage, .{}); + return 0; + }; + var vendor: ?u32 = null; + var timeout_ms: u64 = 10_000; + var i: usize = 1; + while (i < args.len) : (i += 2) { + const value = if (i + 1 < args.len) args[i + 1] else ""; + if (std.mem.eql(u8, args[i], "--vendor")) { + vendor = std.fmt.parseInt(u32, value, 0) catch null; + if (vendor == null) break; + } else if (std.mem.eql(u8, args[i], "--timeout-ms")) { + timeout_ms = std.fmt.parseInt(u64, value, 10) catch 0; + if (timeout_ms == 0 or timeout_ms > 600_000) break; + } else break; + } else return if (try probe(init.gpa, clock, vendor, timeout_ms * std.time.ns_per_ms)) 0 else 1; + std.debug.print(usage, .{}); + return 2; +} + +fn probe(a: std.mem.Allocator, clock: Elapsed, vendor: ?u32, timeout: u64) !bool { + var rt = try vk.Runtime.open(); + defer rt.close(); + var dev = try vk.Device.open(&rt, .{ .features = .{ .storage16 = true, .int16 = true } }); + defer dev.close(); + const pr = dev.properties; + const name = std.mem.sliceTo(&pr.deviceName, 0); + const v = pr.apiVersion; + std.debug.print("MILESTONE device-found elapsed_ms={d} name=\"{s}\" vendor=0x{x:0>4} device=0x{x:0>4} " ++ + "api={d}.{d}.{d} storage16={} int16={} shared={d} fence-timeout={d}ms\n", .{ + clock.ms(), name, pr.vendorID, pr.deviceID, + v >> 22 & 0x7f, v >> 12 & 0x3ff, v & 0xfff, dev.features.storage16, + dev.features.int16, pr.limits.maxComputeSharedMemorySize, timeout / std.time.ns_per_ms, + }); + if (pr.deviceType == vk.c.VK_PHYSICAL_DEVICE_TYPE_CPU) { + std.debug.print("FAIL device: a CPU (software) Vulkan device is never a GPU result\n", .{}); + return false; + } + if (vendor) |want| if (pr.vendorID != want) { + std.debug.print("FAIL device: vendor 0x{x:0>4} is not the requested 0x{x:0>4}\n", .{ pr.vendorID, want }); + return false; + }; + + var cmds = try vk.Commands.init(&dev); + defer cmds.deinit(); + var staging = try vk.Buffer.init(&dev, .staging, 256 << 10); + defer staging.deinit(); + var ops = try vk.ops.Ops.init(&dev, .{ + .qmm_bf16 = &ops_shaders.qmm_bf16, + .qmm_q4_vec4 = &ops_shaders.qmm_q4_vec4, + .qmm_f32 = &ops_shaders.qmm_f32, + .copy_f32_bf16 = &ops_shaders.copy_f32_bf16, + .qmm_q4_tile = &ops_shaders.qmm_q4_tile, + .qmm_q4_tile_f32 = &ops_shaders.qmm_q4_tile_f32, + .dequant_bf16 = &ops_shaders.dequant_bf16, + .gather_u32 = &ops_shaders.gather_u32, + .gather_u16 = &ops_shaders.gather_u16, + .cast_bf16_f32 = &ops_shaders.cast_bf16_f32, + .cast_f32_bf16 = &ops_shaders.cast_f32_bf16, + .copy_u32 = &ops_shaders.copy_u32, + .copy_u16 = &ops_shaders.copy_u16, + }); + defer ops.deinit(); + var attn = try vk.Attention.init(&dev, .{ + .rope_bf16 = &shaders.rope_bf16, + .sdpa_rows_bf16 = &shaders.sdpa_rows_bf16, + .gate_bf16 = &shaders.gate_bf16, + .rope_f32 = &shaders.rope_f32, + .rope_f32_kcache = &shaders.rope_f32_kcache, + .sdpa_rows_f32q = &shaders.sdpa_rows_f32q, + .gate_f32 = &shaders.gate_f32, + }, &ops); + defer attn.deinit(); + + var passed: usize = 0; + var total: usize = 0; + for (fx.rope_cases) |case| { + total += 1; + passed += @intFromBool(try ropeGate(a, clock, &dev, &cmds, staging, &attn, case, timeout)); + } + total += 1; + passed += @intFromBool(try attendGate(a, clock, &dev, &cmds, staging, &attn, timeout)); + total += 1; + passed += @intFromBool(try gateGate(a, clock, &dev, &cmds, staging, &attn, timeout)); + std.debug.print("RESULT stage3 elapsed_ms={d} on \"{s}\": passed={d} total={d}\n", .{ clock.ms(), name, passed, total }); + return passed == total; +} + +fn upload16(dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, shell: []const u16, timeout: u64) !vk.Buffer { + var b = try vk.Buffer.init(dev, .device, std.mem.sliceAsBytes(shell).len); + errdefer b.deinit(); + try vk.orExit(vk.upload(cmds, b, 0, staging, std.mem.sliceAsBytes(shell), timeout)); + return b; +} + +fn fetch(cmds: *vk.Commands, staging: vk.Buffer, b: vk.Buffer, out: []u8, timeout: u64) !void { + try vk.orExit(vk.download(cmds, b, 0, staging, out, timeout)); +} + +fn unchanged(a: std.mem.Allocator, cmds: *vk.Commands, staging: vk.Buffer, shell: []const u16, buf: vk.Buffer, timeout: u64) !bool { + const back = try a.alloc(u16, shell.len); + defer a.free(back); + try fetch(cmds, staging, buf, std.mem.sliceAsBytes(back), timeout); + return std.mem.indexOfDiff(u16, shell, back) == null; +} + +// Rope checks dispatch, 3 bit-identical repeats, f64 reference, wrong-finite/nonfinite refusals, and unchanged input. +fn ropeGate(a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, attn: *vk.Attention, case: fx.RopeCase, timeout: u64) !bool { + const f = try fx.Rope.init(a, case, 0x57a3); + defer f.deinit(); + var q = try upload16(dev, cmds, staging, f.q_bits, timeout); + defer q.deinit(); + // Poison every output word so missing writes cannot inherit a reference. + const out_shell = try a.alloc(u16, f.want.len + fx.tail); + defer a.free(out_shell); + @memset(out_shell, fx.poison16); + var out = try upload16(dev, cmds, staging, out_shell, timeout); + defer out.deinit(); + const out_view = try vk.tensor.View.init(&out, .bf16, 0, f.want.len); + const q_view = try vk.tensor.View.init(&q, .bf16, 0, f.want.len); + try vk.orExit(attn.setBasePosition(cmds, staging, case.base, timeout)); + const rows: u64 = case.rows; + var first: ?[]u16 = null; + defer if (first) |x| a.free(x); + // readback covers data + tail + const full_bits = try a.alloc(u16, f.want.len + fx.tail); + defer a.free(full_bits); + const bits = full_bits[0..f.want.len]; + for (0..repeats) |rep| { + // Re-upload poison before each dispatch; unwritten output words stay NaN and fail finite-output checks. + try vk.orExit(vk.upload(cmds, out, 0, staging, std.mem.sliceAsBytes(out_shell), timeout)); + try vk.orExit(attn.ropeQueries(cmds, out_view, q_view, rows, timeout)); + try fetch(cmds, staging, out, std.mem.sliceAsBytes(full_bits), timeout); + // poison tail must survive untouched + for (full_bits[f.want.len..]) |g_elem| { + if (g_elem != fx.poison16) { + std.debug.print("FAIL rope case=\"{s}\": output tail overwritten at element {d}\n", .{ case.name, f.want.len + (std.mem.indexOf(u16, full_bits[f.want.len..], &[_]u16{g_elem}) orelse 0) }); + return false; + } + } + if (first) |x| { + if (std.mem.indexOfDiff(u16, x, bits)) |at| { + std.debug.print("FAIL rope case=\"{s}\": repeat {d} differs at element {d}\n", .{ case.name, rep, at }); + return false; + } + } else first = try a.dupe(u16, bits); + } + const rep = f.check(bits) catch |e| { + std.debug.print("FAIL rope case=\"{s}\": {s}\n", .{ case.name, @errorName(e) }); + return false; + }; + if (!try plantRefuse(f, bits, case.name, "rope")) return false; + if (!try unchanged(a, cmds, staging, f.q_bits, q, timeout)) { + std.debug.print("FAIL rope case=\"{s}\": the input changed\n", .{case.name}); + return false; + } + std.debug.print("MILESTONE rope elapsed_ms={d} case=\"{s}\" rows={d} base={d} worst={e} of bound, repeats byte-identical, " ++ + "passthrough tail bit-exact, input unchanged, planted outputs refused\n", .{ clock.ms(), case.name, rows, case.base, rep.worst }); + return true; +} + +// Poison a downloaded GOOD copy with wrong finite then nan/+inf/-inf at the first finite-want element; refuse all. +fn plantRefuse(f: anytype, good: []const u16, name: []const u8, gate: []const u8) !bool { + const mutated = try f.a.dupe(u16, good); + defer f.a.free(mutated); + var first_finite: ?usize = null; + for (good, 0..) |bits, i| { + if (std.math.isFinite(cx.fromBf16(bits))) { + first_finite = i; + break; + } + } + const at = first_finite orelse return error.NoFiniteReference; + mutated[at] = 0x4f00; // 2^31: an absolutely wrong finite the budgets must refuse + if (f.check(mutated)) |_| { + std.debug.print("FAIL {s} case=\"{s}\": a planted wrong finite output passed\n", .{ gate, name }); + return false; + } else |_| {} + for ([_]u16{ 0x7fc0, 0x7f80, 0xff80 }, [_][]const u8{ "nan", "+inf", "-inf" }) |pb, pname| { + mutated[at] = pb; + if (f.check(mutated)) |_| { + std.debug.print("FAIL {s} case=\"{s}\": a planted {s} output passed\n", .{ gate, name, pname }); + return false; + } else |_| {} + } + return true; +} + +// CacheKV checks f64 rope/V bits; SDPA checks counts, NaN/finite/zero causality, repeats, refusals, unchanged inputs. +fn attendGate(a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, attn: *vk.Attention, timeout: u64) !bool { + const ac = fx.attend_case; + const cc = fx.cache_case; + const rows = ac.rows; + const qrow_elems: u64 = fx.heads * fx.head_dim; + const crow_elems: u64 = fx.kv_heads * fx.head_dim; + const cache_elems: u64 = fx.kv_heads * cc.capacity * fx.head_dim; + + const f = try fx.Cache.init(a, cc, 0xcacbe); + defer f.deinit(); + // rotated-query fixture data (rope of these rows is checked by the rope gates) + var qf = try fx.Rope.init(a, .{ .name = "attend q", .rows = rows, .base = @intCast(cc.base) }, 0x57a5); + defer qf.deinit(); + + var knorm = try upload16(dev, cmds, staging, f.knorm_bits, timeout); + defer knorm.deinit(); + var vseg = try upload16(dev, cmds, staging, f.v_bits, timeout); + defer vseg.deinit(); + var keys = try upload16(dev, cmds, staging, f.keys_shell, timeout); + defer keys.deinit(); + var values = try upload16(dev, cmds, staging, f.values_shell, timeout); + defer values.deinit(); + var qraw = try upload16(dev, cmds, staging, qf.q_bits, timeout); + defer qraw.deinit(); + var rotated = try vk.Buffer.init(dev, .device, std.mem.sliceAsBytes(qf.q_bits).len); + defer rotated.deinit(); + var ybuf = try vk.Buffer.init(dev, .device, @intCast(rows * qrow_elems * 2)); + defer ybuf.deinit(); + + const knorm_view = try vk.tensor.View.init(&knorm, .bf16, 0, rows * crow_elems); + const vseg_view = try vk.tensor.View.init(&vseg, .bf16, 0, rows * crow_elems); + const keys_view = try vk.tensor.View.init(&keys, .bf16, 0, cache_elems); + const values_view = try vk.tensor.View.init(&values, .bf16, 0, cache_elems); + const qraw_view = try vk.tensor.View.init(&qraw, .bf16, 0, rows * qrow_elems); + const rotated_view = try vk.tensor.View.init(&rotated, .bf16, 0, rows * qrow_elems); + const y_view = try vk.tensor.View.init(&ybuf, .bf16, 0, rows * qrow_elems); + + // positions of the segment rows are base..base+rows-1 + try vk.orExit(attn.setBasePosition(cmds, staging, @intCast(cc.base), timeout)); + // initialize y with poison; a causal kernel overwrites every visible row + const y_init = try a.alloc(u16, @intCast(rows * qrow_elems)); + defer a.free(y_init); + @memset(y_init, fx.poison16); + var first: ?[]u16 = null; + defer if (first) |x| a.free(x); + const ybits = try a.alloc(u16, @intCast(rows * qrow_elems)); + defer a.free(ybits); + + // --- rotate + write the cache --- + try vk.orExit(attn.ropeQueries(cmds, rotated_view, qraw_view, rows, timeout)); + try vk.orExit(attn.cacheKV(cmds, keys_view, values_view, knorm_view, vseg_view, rows, cc.base, cc.capacity, timeout)); + // --- cache checks --- + const keys_back = try a.alloc(u16, @intCast(cache_elems)); + defer a.free(keys_back); + const values_back = try a.alloc(u16, @intCast(cache_elems)); + defer a.free(values_back); + try fetch(cmds, staging, keys, std.mem.sliceAsBytes(keys_back), timeout); + try fetch(cmds, staging, values, std.mem.sliceAsBytes(values_back), timeout); + const kr = fx.Cache.checkKeys(f, keys_back) catch |e| { + std.debug.print("FAIL cache case=\"{s}\": keys {s}\n", .{ cc.name, @errorName(e) }); + return false; + }; + fx.Cache.checkValueBits(f, values_back) catch |e| { + std.debug.print("FAIL cache case=\"{s}\": values {s}\n", .{ cc.name, @errorName(e) }); + return false; + }; + if (!try unchanged(a, cmds, staging, f.knorm_bits, knorm, timeout) or !try unchanged(a, cmds, staging, f.v_bits, vseg, timeout)) { + std.debug.print("FAIL cache case=\"{s}\": a segment input changed\n", .{cc.name}); + return false; + } + std.debug.print("MILESTONE cache elapsed_ms={d} case=\"{s}\" base={d} rows={d} capacity={d} keys worst={e} of bound, " ++ + "V transport bit-exact, poison rows untouched, inputs unchanged\n", .{ clock.ms(), cc.name, cc.base, rows, cc.capacity, kr.worst }); + + // --- SDPA over the ACTUAL bits the GPU reads (rotated q, written caches) --- + const rotated_back = try a.alloc(u16, @intCast(rows * qrow_elems)); + defer a.free(rotated_back); + try fetch(cmds, staging, rotated, std.mem.sliceAsBytes(rotated_back), timeout); + const g = try fx.Attend.init(a, ac, rotated_back, keys_back, values_back, 0); + defer g.deinit(); + + for (0..repeats) |rep| { + try vk.orExit(vk.upload(cmds, ybuf, 0, staging, std.mem.sliceAsBytes(y_init), timeout)); + try vk.orExit(attn.attend(cmds, y_view, rotated_view, keys_view, values_view, ac.prefix, rows, cc.capacity, timeout)); + try fetch(cmds, staging, ybuf, std.mem.sliceAsBytes(ybits), timeout); + if (first) |x| { + if (std.mem.indexOfDiff(u16, x, ybits)) |at| { + std.debug.print("FAIL sdpa: repeat {d} differs at element {d}\n", .{ rep, at }); + return false; + } + } else first = try a.dupe(u16, ybits); + } + const rep = g.check(ybits) catch |e| { + std.debug.print("FAIL sdpa: {s}\n", .{@errorName(e)}); + return false; + }; + + // Fresh cache per row poisons keys >= prefix+r+1; rows=1/prefix=prefix+r must bit-match baseline without budget. + const alt_keys = try a.alloc(u16, @intCast(cache_elems)); + defer a.free(alt_keys); + const y_row_bits = try a.alloc(u16, @intCast(qrow_elems)); + defer a.free(y_row_bits); + const variants = [_]u16{ fx.poison16, finite_poison, 0x0000 }; + const vnames = [_][]const u8{ "nan", "large-finite", "zero" }; + for (variants, vnames) |fill, vname| { + for (0..@intCast(rows)) |r| { + @memcpy(alt_keys, keys_back); + const visible = fx.visibleFrom(ac, @intCast(r)); + for (0..@intCast(fx.kv_heads)) |kh| { + for (visible..cc.capacity) |j| { + const base = (@as(usize, @intCast(kh)) * @as(usize, @intCast(cc.capacity)) + j) * @as(usize, @intCast(fx.head_dim)); + @memset(alt_keys[base..][0..@intCast(fx.head_dim)], fill); + } + } + var alt_buf = try upload16(dev, cmds, staging, alt_keys, timeout); + defer alt_buf.deinit(); + const alt_view = try vk.tensor.View.init(&alt_buf, .bf16, 0, cache_elems); + // row-slice views: q row r of `rotated`, output row r of `ybuf` + const row_byte_offset = r * qrow_elems * 2; + const q_row_view = try vk.tensor.View.init(&rotated, .bf16, row_byte_offset, qrow_elems); + const y_row_view = try vk.tensor.View.init(&ybuf, .bf16, row_byte_offset, qrow_elems); + try vk.orExit(vk.upload(cmds, ybuf, 0, staging, std.mem.sliceAsBytes(y_init), timeout)); + try vk.orExit(attn.attend(cmds, y_row_view, q_row_view, alt_view, values_view, ac.prefix + @as(u64, @intCast(r)), 1, cc.capacity, timeout)); + try vk.orExit(vk.download(cmds, ybuf, row_byte_offset, staging, std.mem.sliceAsBytes(y_row_bits), timeout)); + const baseline_row = first.?[r * @as(usize, @intCast(qrow_elems)) ..][0..@intCast(qrow_elems)]; + if (std.mem.indexOfDiff(u16, baseline_row, y_row_bits)) |at| { + std.debug.print("FAIL sdpa causality: {s} future keys changed row {d} output at element {d}\n", .{ vname, r, at }); + return false; + } + } + } + if (!try plantRefuse(g, first.?, ac.name, "sdpa")) return false; + if (!try unchanged(a, cmds, staging, qf.q_bits, qraw, timeout)) { + std.debug.print("FAIL sdpa: the raw q input changed\n", .{}); + return false; + } + std.debug.print("MILESTONE sdpa elapsed_ms={d} case=\"{s}\" prefix={d} rows={d} visible-counts={any} worst={e} of bound, " ++ + "repeats byte-identical, nan/large-finite/zero future keys bit-identical, q unchanged, planted outputs refused\n", .{ clock.ms(), ac.name, ac.prefix, rows, g.visible, rep.worst }); + return true; +} + +// Interleaved output repeats 3 times, matches f64, refuses plants, and leaves both inputs unchanged. +fn gateGate(a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, attn: *vk.Attention, timeout: u64) !bool { + const case = fx.gate_case; + const f = try fx.Gate.init(a, case, 0x57a6); + defer f.deinit(); + var qraw = try upload16(dev, cmds, staging, f.qraw_bits, timeout); + defer qraw.deinit(); + var ctx = try upload16(dev, cmds, staging, f.ctx_bits, timeout); + defer ctx.deinit(); + // out starts poisoned: every row must be written by the kernel + var out = try upload16(dev, cmds, staging, f.out_bits, timeout); + defer out.deinit(); + const elems: u64 = case.rows * 2048; + const out_view = try vk.tensor.View.init(&out, .bf16, 0, elems); + const ctx_view = try vk.tensor.View.init(&ctx, .bf16, 0, elems); + const qraw_view = try vk.tensor.View.init(&qraw, .bf16, 0, case.rows * 4096); + const bits = try a.alloc(u16, @intCast(elems)); + defer a.free(bits); + var first: ?[]u16 = null; + defer if (first) |x| a.free(x); + for (0..repeats) |rep| { + try vk.orExit(vk.upload(cmds, out, 0, staging, std.mem.sliceAsBytes(f.out_bits), timeout)); + try vk.orExit(attn.gateOutput(cmds, out_view, ctx_view, qraw_view, case.rows, timeout)); + try fetch(cmds, staging, out, std.mem.sliceAsBytes(bits), timeout); + if (first) |x| { + if (std.mem.indexOfDiff(u16, x, bits)) |at| { + std.debug.print("FAIL gate case=\"{s}\": repeat {d} differs at element {d}\n", .{ case.name, rep, at }); + return false; + } + } else first = try a.dupe(u16, bits); + } + const rep = f.check(first.?) catch |e| { + std.debug.print("FAIL gate case=\"{s}\": {s}\n", .{ case.name, @errorName(e) }); + if (f.firstViolation(first.?)) |d| { + std.debug.print("DIAG gate index={d} row={d} col={d} head={d} ctx=0x{x:0>4} ({e}) raw_gate=0x{x:0>4} ({e}) got_word=0x{x:0>4} got={e} want={e} bound={e} ratio={e} HOST_sigmoid_f64={e} HOST_sigmoid_f32={e} HOST_sigmoid_bf16=0x{x:0>4} ({e})\n", .{ + d.index, d.row, d.d, d.head, d.ctx_bits, d.ctx, d.gate_bits, d.gate, first.?[d.index], d.got, d.want, d.bound, d.ratio, d.sigmoid_f64, d.sigmoid_f32, d.sigmoid_bf16_bits, d.sigmoid_bf16, + }); + } + return false; + }; + if (!try plantRefuse(f, first.?, case.name, "gate")) return false; + if (!try unchanged(a, cmds, staging, f.qraw_bits, qraw, timeout) or !try unchanged(a, cmds, staging, f.ctx_bits, ctx, timeout)) { + std.debug.print("FAIL gate case=\"{s}\": an input changed\n", .{case.name}); + return false; + } + std.debug.print("MILESTONE gate elapsed_ms={d} case=\"{s}\" rows={d} worst={e} of bound, repeats byte-identical, " ++ + "inputs unchanged, planted outputs refused\n", .{ clock.ms(), case.name, case.rows, rep.worst }); + return true; +} diff --git a/zig/tests/vulkan/stage4_fixture.zig b/zig/tests/vulkan/stage4_fixture.zig new file mode 100644 index 00000000000..279c558d632 --- /dev/null +++ b/zig/tests/vulkan/stage4_fixture.zig @@ -0,0 +1,283 @@ +// Stage 4 fixture: the GatedDeltaNet closure, 16 heads, head_dim 128, 6144 conv channels, 4 taps, vs f64 host refs. + +// depth/ops/mids budgets, mag: split max|vals|^2, gate |g|, recur max state, norm max|x|^2*len; chunk=4/8 one ref + +const std = @import("std"); +const core = @import("stage4_fixture_core.zig"); +const split_mod = @import("stage4_fixture_split.zig"); +const state_mod = @import("stage4_fixture_state.zig"); + +pub const c = core.c; +pub const tail = core.tail; +pub const poison16 = core.poison16; +pub const Model = core.Model; +pub const tokens = core.tokens; +pub const Report = core.Report; +pub const Output = core.Output; +pub const Violation = core.Violation; +pub const budget = core.budget; +pub const sampleDyadic = core.sampleDyadic; +pub const sampleDyadicF32 = core.sampleDyadicF32; +pub const gateRef = core.gateRef; +pub const gateRefOracle = core.gateRefOracle; +pub const Split = split_mod.Split; +pub const SplitOne = split_mod.SplitOne; +pub const Conv = split_mod.Conv; +pub const RecurWalk = state_mod.RecurWalk; +pub const Recur = state_mod.Recur; +pub const Decode = state_mod.Decode; +pub const Norm = state_mod.Norm; +const f32Of = core.f32Of; + +// Replay the decode seed through the production gate checker; report old-reference drift, not a GPU-failure cause. +test "gateRefOracle replays the fixture seed and diagnoses g and beta separately" { + const a = testing.allocator; + const f = try SplitOne.init(a, 0x57ac); + defer f.deinit(); + // what the shader computes, on the host: g f32 words, beta bf16 + const g_pairs = try a.alloc(u16, 2 * @as(usize, @intCast(Model.heads))); + defer a.free(g_pairs); + const beta_bits = try a.alloc(u16, @intCast(Model.heads)); + defer a.free(beta_bits); + var old_worst: f64 = 0; + for (0..@intCast(Model.heads)) |hv| { + const ash = gateRefOracle(f.a_bits[hv], f.b_bits[hv], f32Of(f.a_log_bits, hv), f.dt_bits[hv]); + const gbits: u32 = @bitCast(@as(f32, @floatCast(ash.g))); + g_pairs[2 * hv] = @truncate(gbits); + g_pairs[2 * hv + 1] = @truncate(gbits >> 16); + beta_bits[hv] = c.toBf16(@floatCast(ash.beta)); + // evidence: the OLD unrounded f64 reference vs the shader chain + const old = gateRef( + c.fromBf16(f.a_bits[hv]), + c.fromBf16(f.b_bits[hv]), + f32Of(f.a_log_bits, hv), + c.fromBf16(f.dt_bits[hv]), + ); + const gd = @abs(old.g - ash.g); + const gbound = budget(1, 4, 2, ash.g, @abs(ash.g)); + old_worst = @max(old_worst, gd / gbound); + } + std.debug.print("HOST gate old-reference drift seed=0x57ac max_ratio={e} (not a GPU diagnosis)\n", .{old_worst}); + // the exact chain replays CLEAN through the same checker + const rep = try f.checkGates(g_pairs, beta_bits, null); + try testing.expect(rep.worst <= 1); + var diag: Violation = undefined; + const g0 = g_pairs[0]; + const g1 = g_pairs[1]; + g_pairs[0] = 0; + g_pairs[1] = 0x4f00; + try testing.expectError(error.GateOverBound, f.checkGates(g_pairs, beta_bits, &diag)); + try testing.expectEqual(Output.g, diag.output); + try testing.expectEqual(@as(usize, 0), diag.index); + g_pairs[0] = g0; + g_pairs[1] = g1; + const beta0 = beta_bits[0]; + beta_bits[0] = 0x4f00; + try testing.expectError(error.GateOverBound, f.checkGates(g_pairs, beta_bits, &diag)); + try testing.expectEqual(Output.beta, diag.output); + try testing.expectEqual(@as(usize, 0), diag.index); + beta_bits[0] = beta0; + // subnormal census (log-only): a_log<=3, sp<=~4.5 => g~exp(-20.085*4.5)~5.6e-40; FTZ=>g=0; host ref:no FTZ. + var min_g: f32 = std.math.floatMax(f32); + var subnormals: usize = 0; + for (0..@intCast(Model.heads)) |hv| { + const ash = gateRefOracle(f.a_bits[hv], f.b_bits[hv], f32Of(f.a_log_bits, hv), f.dt_bits[hv]); + const g: f32 = @floatCast(ash.g); + if (g > 0) min_g = @min(min_g, g); + const bits: u32 = @bitCast(g); + const exponent = (bits >> 23) & 0xff; + if (exponent == 0 and (bits & 0x7fffff) != 0) subnormals += 1; // f32 subnormal + if (g == 0) subnormals += 1; // flushed/underflowed to zero + } + std.debug.print("HOST gate g census (exact seed): min_positive_g={e} subnormal_or_zero_g={d}/16\n", .{ min_g, subnormals }); +} + +const testing = std.testing; + +test "budget grows with depth and stays finite on zeros" { + try testing.expect(budget(300, 6, 2, 1, 1) > budget(2, 6, 1, 1, 1)); + try testing.expect(budget(2, 6, 1, 0, 0) > 0); +} + +test "gateRef behaves like exp(-exp(A_log) softplus) and a sigmoid" { + const g0 = gateRef(0, 0, 0, 0); + try testing.expectApproxEqAbs(@as(f64, @exp(-@exp(0.0) * std.math.log1p(@as(f64, 1.0)))), g0.g, 1e-12); + try testing.expectApproxEqAbs(@as(f64, 0.5), g0.beta, 1e-12); +} + +test "sampleDyadic is exactly bf16" { + var prng = std.Random.DefaultPrng.init(9); + const rng = prng.random(); + for (0..64) |_| { + const bits = sampleDyadic(rng, -5, 3); + const v = c.fromBf16(bits); + try testing.expectEqual(v, c.fromBf16(c.toBf16(v))); + } +} + +// Negative plants vs PRODUCTION Norm checker: honest output passes; a wrong finite and every nonfinite class refused. +test "norm checker refuses planted outputs" { + const a = testing.allocator; + const f = try Norm.init(a, 0x18); + defer f.deinit(); + const width: usize = @intCast(Model.head_dim); + const head_rows: usize = 4 * @as(usize, @intCast(Model.heads)); + const good = try a.alloc(u16, head_rows * width); + defer a.free(good); + for (0..head_rows) |row| { + var ss: f64 = 0; + for (0..width) |i| { + const v = c.fromBf16(f.y_bits[row * width + i]); + ss += v * v; + } + const inv: f64 = 1.0 / @sqrt(ss / @as(f64, width) + f.eps); + for (0..width) |i| { + const x = c.fromBf16(f.y_bits[row * width + i]); + const z = c.fromBf16(f.z_bits[row * width + i]); + // the same Metal gdn_post association the checker references + const normed = c.fromBf16(c.toBf16(@floatCast(c.fromBf16(f.w_bits[i]) * (x * inv)))); + const silu = z / (1.0 + @exp(-z)); + good[row * width + i] = c.toBf16(@floatCast(silu * normed)); + } + } + const nrep = try f.check(good); + try testing.expect(nrep.worst <= 1); + good[0] = 0x4f00; // 2^31: an absolutely wrong finite, far past any budget + try testing.expectError(error.OverBound, f.check(good)); + good[0] = 0x7f80; + try testing.expectError(error.NonFiniteOutput, f.check(good)); +} + +// Carry-out self-tests: rounded state PASSES; poisoned (0xffff=NaN) carry-out REJECTED; dead store (all-poison) fails. +test "recur checker accepts rounded state and refuses poisoned carry-out" { + const a = testing.allocator; + // fabricate the split's outputs as dyadic bits (the walker reads them) + const soa_len: usize = @intCast(3 * tokens * Model.linear_dim); + const soa_bits = try a.alloc(u16, soa_len); + defer a.free(soa_bits); + var prng = std.Random.DefaultPrng.init(0x51); + const rng = prng.random(); + for (soa_bits) |*x| x.* = sampleDyadic(rng, -3, 1); + const g_len: usize = @intCast(tokens * Model.heads * 2); + const g_pairs = try a.alloc(u16, g_len); + defer a.free(g_pairs); + for (g_pairs) |*x| { + const bits: u32 = @bitCast(sampleDyadicF32(rng, -4, 0)); + x.* = @truncate(bits); + } + const beta_len: usize = @intCast(tokens * Model.heads); + const beta_bits = try a.alloc(u16, beta_len); + defer a.free(beta_bits); + for (beta_bits) |*x| x.* = sampleDyadic(rng, -3, 1); + const f = try Recur.init(a, soa_bits, g_pairs, beta_bits, 0x52, tokens); // single chunk: no interior snapshot + defer f.deinit(); + const state_elems: usize = @intCast(Model.state_elements); + // correctly reference-rounded state must pass + const hf_good = try a.alloc(u16, 2 * state_elems); + defer a.free(hf_good); + for (f.want_state, 0..) |w, i| { + const bits: u32 = @bitCast(@as(f32, @floatCast(w))); + hf_good[2 * i] = @truncate(bits); + hf_good[2 * i + 1] = @truncate(bits >> 16); + } + const y_good = try a.alloc(u16, @intCast(tokens * Model.linear_dim)); + defer a.free(y_good); + for (f.want_y, 0..) |w, i| y_good[i] = c.toBf16(@floatCast(w)); + const rep = try f.check(y_good, hf_good, null); + try testing.expect(rep.worst <= 1); + // poisoned carry-out (NaN word) must be rejected + const hf_bad = try a.dupe(u16, hf_good); + defer a.free(hf_bad); + hf_bad[0] = 0xffff; + hf_bad[1] = 0xffff; + try testing.expectError(error.NonFiniteOutput, f.check(y_good, hf_bad, null)); + // a dead store: all-poison state must be rejected + @memset(hf_bad, 0xffff); + try testing.expectError(error.NonFiniteOutput, f.check(y_good, hf_bad, null)); + // single-chunk Recur expects NO snapshot + try testing.expect(f.want_snap == null); + // a chunked Recur has exactly one interior boundary reference; the correctly-rounded slot passes, poison is refused + const fc = try Recur.init(a, soa_bits, g_pairs, beta_bits, 0x54, 4); + defer fc.deinit(); + const ws = fc.want_snap.?; + try testing.expectEqual(@as(usize, @intCast(Model.state_elements)), ws.len); + const snap_good = try a.alloc(u16, 2 * @as(usize, @intCast(Model.state_elements))); + defer a.free(snap_good); + for (ws, 0..) |w, i| { + const bits: u32 = @bitCast(@as(f32, @floatCast(w))); + snap_good[2 * i] = @truncate(bits); + snap_good[2 * i + 1] = @truncate(bits >> 16); + } + const snap_rep = try fc.checkSnap(snap_good, null); + try testing.expect(snap_rep.worst <= 1); + snap_good[2] = 0xffff; + snap_good[3] = 0xffff; + try testing.expectError(error.NonFiniteOutput, fc.checkSnap(snap_good, null)); + @memset(snap_good, 0xffff); + try testing.expectError(error.NonFiniteOutput, fc.checkSnap(snap_good, null)); +} + +test "decode checker refuses poisoned carry-out state" { + const a = testing.allocator; + const f = try Decode.init(a, 0x53); + defer f.deinit(); + const state_elems: usize = @intCast(Model.state_elements); + const y_good = try a.alloc(u16, @intCast(Model.linear_dim)); + defer a.free(y_good); + for (f.want_y, 0..) |w, i| y_good[i] = c.toBf16(@floatCast(w)); + const souts_good = try a.alloc(u16, 2 * state_elems); + defer a.free(souts_good); + for (f.want_state, 0..) |w, i| { + const bits: u32 = @bitCast(@as(f32, @floatCast(w))); + souts_good[2 * i] = @truncate(bits); + souts_good[2 * i + 1] = @truncate(bits >> 16); + } + const rep = try f.check(y_good, souts_good, null); + try testing.expect(rep.worst <= 1); + const souts_bad = try a.dupe(u16, souts_good); + defer a.free(souts_bad); + souts_bad[100] = 0xffff; + souts_bad[101] = 0xffff; + try testing.expectError(error.NonFiniteOutput, f.check(y_good, souts_bad, null)); +} + +test "fixture initialization fails cleanly under allocation failure" { + var failing = std.testing.FailingAllocator.init(testing.allocator, .{ .fail_index = 1 }); + try testing.expectError(error.OutOfMemory, Norm.init(failing.allocator(), 0x19)); + try testing.expectError(error.OutOfMemory, Decode.init(failing.allocator(), 0x1a)); +} + +test "decode seed replays the measured wrong raw-gates branch at y zero" { + const f = try Decode.init(testing.allocator, 0x57a8); + defer f.deinit(); + const replay = struct { + fn first(d: Decode, g: f32, beta: f32) u16 { + var kv: f32 = 0; + for (0..128) |i| { + const state: f32 = @floatCast(f32Of(d.state_in_bits, i)); + const decayed: f32 = state * g; + const term: f32 = decayed * c.fromBf16(d.soa_bits[2048 + i]); + kv += term; + } + const delta: f32 = (c.fromBf16(d.soa_bits[4096]) - kv) * beta; + var out: f32 = 0; + for (0..128) |i| { + const state: f32 = @floatCast(f32Of(d.state_in_bits, i)); + const decayed: f32 = state * g; + const update: f32 = delta * c.fromBf16(d.soa_bits[2048 + i]); + const next: f32 = decayed + update; + const term: f32 = next * c.fromBf16(d.soa_bits[i]); + out += term; + } + return c.toBf16(out); + } + }.first; + const wrong = replay(f, c.fromBf16(f.a_bits[0]), c.fromBf16(f.b_bits[0])); + try testing.expectEqual(c.toBf16(892), wrong); + const gates = gateRefOracle(f.a_bits[0], f.b_bits[0], f32Of(f.a_log_bits, 0), f.dt_bits[0]); + const corrected = replay(f, @floatCast(gates.g), @floatCast(gates.beta)); + const bound = budget(260, 4, 0, f.want_y[0], f.max_state); + try testing.expect(@abs(c.fromBf16(corrected) - f.want_y[0]) <= bound); + try testing.expect(@abs(c.fromBf16(wrong) - f.want_y[0]) > bound); + std.debug.print("HOST decode first-element replay seed=0x57a8 skipped-raw={e} corrected={e} want={e} bound={e}\n", .{ c.fromBf16(wrong), c.fromBf16(corrected), f.want_y[0], bound }); +} diff --git a/zig/tests/vulkan/stage4_fixture_core.zig b/zig/tests/vulkan/stage4_fixture_core.zig new file mode 100644 index 00000000000..8c53d23617a --- /dev/null +++ b/zig/tests/vulkan/stage4_fixture_core.zig @@ -0,0 +1,90 @@ +// Stage 4 fixture core: the consts, the f64 budget and the bf16 gate references the fixture parts share. + +const std = @import("std"); +const contract = @import("contract"); +pub const c = contract.spec; +const gdn_mod = @import("vulkan").gdn; + +pub const tail = 16; +pub const poison16: u16 = 0xffff; + +pub const Model = gdn_mod.Model; +pub const tokens: u64 = 8; + +/// `worst` is |got - want| OVER its budget (a ratio; >1 fails), with `where`. +pub const Report = struct { worst: f64 = 0, where: usize = 0 }; + +// First failing element, probe print. `want` = HOST oracle-chain ref; sp/exp never stored; negative plants pass null. +pub const Output = enum { soa, g, beta, y, state, snapshot }; +pub const Violation = struct { + index: usize, + output: Output = .soa, + want: f64 = 0, + got: f64 = 0, + bound: f64 = 0, + ratio: f64 = 0, + // The want value's f32 bits where the reference is an f32 scalar (g): shows f32 class (subnormal/underflow) only. + want_f32_bits: u32 = 0, +}; + +pub fn violated(v: ?*Violation, output: Output, index: usize, want: f64, got: f64, bound: f64) void { + if (v) |d| d.* = .{ + .output = output, + .index = index, + .want = want, + .got = got, + .bound = bound, + .ratio = if (bound > 0) @abs(got - want) / bound else std.math.inf(f64), + .want_f32_bits = @bitCast(@as(f32, @floatCast(want))), + }; +} + +// budget: depth f32 sum@mag, ops f32 + mids bf16 rounds, 1 store; not sum-of-squares; never enlarge after GPU runs +pub fn budget(depth: f64, ops: f64, mids: f64, want: f64, mag: f64) f64 { + const u: f64 = 0x1p-24; + const g = if (depth > 1) depth * u / (1 - depth * u) else 0; + const value_scale = @max(@abs(want), mag); + const work = g * mag + (ops * u + mids * 0x1p-8) * value_scale; + return work + 0x1p-8 * (@abs(want) + work) + 0x1p-133; +} + +pub fn sampleDyadic(rng: std.Random, lo: i32, hi: i32) u16 { + const e = rng.intRangeAtMost(i32, lo, hi); + const m: f32 = if (rng.boolean()) 1.0 else 1.5; + return c.toBf16(m * std.math.pow(f32, 2, @floatFromInt(e)) * (if (rng.boolean()) @as(f32, 1) else -1)); +} + +pub fn sampleDyadicF32(rng: std.Random, lo: i32, hi: i32) f32 { + const e = rng.intRangeAtMost(i32, lo, hi); + const m: f32 = if (rng.boolean()) 1.0 else 1.5; + return m * std.math.pow(f32, 2, @floatFromInt(e)) * (if (rng.boolean()) @as(f32, 1) else -1); +} + +pub fn poisonedShell(a: std.mem.Allocator, len: usize) ![]u16 { + const buf = try a.alloc(u16, len + tail); + @memset(buf, poison16); + return buf; +} + +pub fn f32Of(pairs: []const u16, word: usize) f64 { + return @as(f32, @bitCast(@as(u32, pairs[2 * word]) | (@as(u32, pairs[2 * word + 1]) << 16))); +} + +// f64 chain shared by split lane-0 and decode raw-gates; omits bf16 rounds (~20x); kept for amplification test +pub fn gateRef(a_in: f64, b_in: f64, a_log: f64, dt: f64) gdn_mod.GateF64 { + return gdn_mod.gateReferenceF64(a_in, b_in, a_log, dt); +} + +// Oracle: s=BF16(a+dt); sp=BF16(max(s,0)+log(1+exp(-|s|))); g f32; beta=BF16(sigmoid(b)); residual = exp/log ulps +pub fn gateRefOracle(a_in: u16, b_in: u16, a_log: f64, dt_in: u16) gdn_mod.GateF64 { + const a = c.fromBf16(a_in); + const dt = c.fromBf16(dt_in); + const s = c.fromBf16(c.toBf16(a + dt)); + const sp_raw: f32 = @max(s, 0.0) + @log(@as(f32, 1.0) + @exp(-@abs(s))); + const sp = c.fromBf16(c.toBf16(sp_raw)); + // the shader's exp reads the f32 A_log: cast BEFORE the exp + const ea: f32 = @exp(@as(f32, @floatCast(a_log))); + const g: f32 = @exp(-ea * sp); + const beta_raw: f32 = 1.0 / (1.0 + @exp(-c.fromBf16(b_in))); + return .{ .g = @floatCast(g), .beta = @floatCast(c.fromBf16(c.toBf16(beta_raw))) }; +} diff --git a/zig/tests/vulkan/stage4_fixture_split.zig b/zig/tests/vulkan/stage4_fixture_split.zig new file mode 100644 index 00000000000..cda15f96724 --- /dev/null +++ b/zig/tests/vulkan/stage4_fixture_split.zig @@ -0,0 +1,337 @@ +// Stage 4 fixture: the split at T tokens, the split at one token and the conv prefill fixture. + +const std = @import("std"); +const core = @import("stage4_fixture_core.zig"); + +const c = core.c; +const Model = core.Model; +const tokens = core.tokens; +const budget = core.budget; +const gateRefOracle = core.gateRefOracle; +const f32Of = core.f32Of; +const poisonedShell = core.poisonedShell; +const violated = core.violated; +const Report = core.Report; +const Output = core.Output; +const Violation = core.Violation; +const sampleDyadic = core.sampleDyadic; +const sampleDyadicF32 = core.sampleDyadicF32; + +// ---- Split ---- + +// Gate split at T=tokens: conv_out/a/b/dt/a_log uploaded; check reads downloaded soa (bf16), g (f32), beta (bf16) bits. +pub const Split = struct { + a: std.mem.Allocator, + conv_out_bits: []u16, + a_bits: []u16, + b_bits: []u16, + a_log_bits: []u16, // f32 bits as u16 pairs, 16 words + dt_bits: []u16, + + pub fn init(a: std.mem.Allocator, seed: u64) !Split { + const conv_len: usize = @intCast(tokens * Model.conv_channels); + const gate_len: usize = @intCast(tokens * Model.heads); + const conv_out_bits = try poisonedShell(a, conv_len); + errdefer a.free(conv_out_bits); + const a_bits = try poisonedShell(a, gate_len); + errdefer a.free(a_bits); + const b_bits = try poisonedShell(a, gate_len); + errdefer a.free(b_bits); + const dt_bits = try poisonedShell(a, @intCast(Model.heads)); + errdefer a.free(dt_bits); + const a_log_bits = try poisonedShell(a, @intCast(2 * Model.heads)); + errdefer a.free(a_log_bits); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (conv_out_bits[0..conv_len]) |*x| x.* = sampleDyadic(rng, -4, 2); + for (a_bits[0..gate_len]) |*x| x.* = sampleDyadic(rng, -3, 2); + for (b_bits[0..gate_len]) |*x| x.* = sampleDyadic(rng, -3, 2); + for (dt_bits[0..@intCast(Model.heads)]) |*x| x.* = sampleDyadic(rng, -4, 1); + for (0..@intCast(Model.heads)) |h| { + const bits: u32 = @bitCast(sampleDyadicF32(rng, -4, 1)); + a_log_bits[2 * h] = @truncate(bits); + a_log_bits[2 * h + 1] = @truncate(bits >> 16); + } + return .{ .a = a, .conv_out_bits = conv_out_bits, .a_bits = a_bits, .b_bits = b_bits, .a_log_bits = a_log_bits, .dt_bits = dt_bits }; + } + + pub fn deinit(f: Split) void { + f.a.free(f.conv_out_bits); + f.a.free(f.a_bits); + f.a.free(f.b_bits); + f.a.free(f.dt_bits); + f.a.free(f.a_log_bits); + } + + pub fn check(f: Split, soa_bits: []const u16, g_pairs: []const u16, beta_bits: []const u16, violation: ?*Violation) !Report { + var r: Report = .{}; + const C: usize = @intCast(Model.conv_channels); + const section: usize = @intCast(tokens * Model.linear_dim); + for (0..@intCast(tokens)) |t| { + for (0..@intCast(3 * Model.heads)) |head| { + const hv: usize = @intCast(head % 16); + var vals: [128]f64 = undefined; + var ss: f64 = 0; + var vmax: f64 = 0; + for (0..128) |i| { + const cv = c.fromBf16(f.conv_out_bits[t * C + head * 128 + i]); + const sig = 1.0 / (1.0 + @exp(-cv)); + vals[i] = cv * sig; + ss += vals[i] * vals[i]; + vmax = @max(vmax, @abs(vals[i])); + } + const is_q = head < 16; + const is_k = head >= 16 and head < 32; + const inv: f64 = 1.0 / @sqrt(ss / 128.0 + 1e-6 / 128.0); + const scale_f: f64 = if (is_q) 1.0 / 128.0 else @sqrt(1.0 / 128.0); + const dst_base = (head / 16) * section + (t * 16 + hv) * 128; + for (0..128) |i| { + const want = if (is_q or is_k) scale_f * (vals[i] * inv) else vals[i]; + const got = c.fromBf16(soa_bits[dst_base + i]); + if (!std.math.isFinite(got)) { + violated(violation, .soa, dst_base + i, want, got, 0); + return error.NonFiniteOutput; + } + const d = @abs(got - want); + const bound = budget(40, 5, 3, want, vmax * vmax); + if (d > bound) { + violated(violation, .soa, dst_base + i, want, got, bound); + return error.OverBound; + } + const ratio = d / bound; + if (ratio > r.worst) r = .{ .worst = ratio, .where = dst_base + i }; + } + } + for (0..@intCast(Model.heads)) |hv| { + // HOST oracle chain: the exact bf16 rounds the shader does + const ref = gateRefOracle( + f.a_bits[t * 16 + hv], + f.b_bits[t * 16 + hv], + f32Of(f.a_log_bits, hv), + f.dt_bits[hv], + ); + const g_idx = (t * 16 + hv) * 2; + const g_got = f32Of(g_pairs, t * 16 + hv); + if (!std.math.isFinite(g_got)) { + violated(violation, .g, g_idx / 2, ref.g, g_got, 0); + return error.NonFiniteOutput; + } + const gd = @abs(g_got - ref.g); + const gbound = budget(1, 4, 2, ref.g, @abs(ref.g)); + if (gd > gbound) { + violated(violation, .g, g_idx / 2, ref.g, g_got, gbound); + return error.GateOverBound; + } + const gratio = gd / gbound; + if (gratio > r.worst) r = .{ .worst = gratio, .where = g_idx }; + const beta_want = ref.beta; + const beta_got = c.fromBf16(beta_bits[t * 16 + hv]); + if (!std.math.isFinite(beta_got)) { + violated(violation, .beta, g_idx / 2, beta_want, beta_got, 0); + return error.NonFiniteOutput; + } + const bd = @abs(beta_got - beta_want); + const bbound = budget(1, 3, 1, beta_want, @abs(beta_want)); + if (bd > bbound) { + violated(violation, .beta, g_idx / 2, beta_want, beta_got, bbound); + return error.GateOverBound; + } + const bratio = bd / bbound; + if (bratio > r.worst) r = .{ .worst = bratio, .where = g_idx + 1 }; + } + } + return r; + } +}; + +// ---- SplitOne: the gate split at exactly ONE token (the convDecodeStep tail) ---- + +// Split T=1: soa[3*2048], g/beta[16], conv_out[6144]; ACTUAL activated conv bits in `conv_out_bits`[6144] before check. +pub const SplitOne = struct { + a: std.mem.Allocator, + conv_out_bits: []u16, + a_bits: []u16, + b_bits: []u16, + a_log_bits: []u16, // f32 pairs + dt_bits: []u16, + + pub fn init(a: std.mem.Allocator, seed: u64) !SplitOne { + const C: usize = @intCast(Model.conv_channels); + const conv_out_bits = try poisonedShell(a, C); + errdefer a.free(conv_out_bits); + const a_bits = try poisonedShell(a, @intCast(Model.heads)); + errdefer a.free(a_bits); + const b_bits = try poisonedShell(a, @intCast(Model.heads)); + errdefer a.free(b_bits); + const dt_bits = try poisonedShell(a, @intCast(Model.heads)); + errdefer a.free(dt_bits); + const a_log_bits = try poisonedShell(a, @intCast(2 * Model.heads)); + errdefer a.free(a_log_bits); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (conv_out_bits[0..C]) |*x| x.* = sampleDyadic(rng, -4, 2); + for (a_bits[0..@intCast(Model.heads)]) |*x| x.* = sampleDyadic(rng, -3, 2); + for (b_bits[0..@intCast(Model.heads)]) |*x| x.* = sampleDyadic(rng, -3, 2); + for (dt_bits[0..@intCast(Model.heads)]) |*x| x.* = sampleDyadic(rng, -4, 1); + for (0..@intCast(Model.heads)) |h| { + const bits: u32 = @bitCast(sampleDyadicF32(rng, -4, 1)); + a_log_bits[2 * h] = @truncate(bits); + a_log_bits[2 * h + 1] = @truncate(bits >> 16); + } + return .{ .a = a, .conv_out_bits = conv_out_bits, .a_bits = a_bits, .b_bits = b_bits, .a_log_bits = a_log_bits, .dt_bits = dt_bits }; + } + + pub fn deinit(f: SplitOne) void { + f.a.free(f.conv_out_bits); + f.a.free(f.a_bits); + f.a.free(f.b_bits); + f.a.free(f.dt_bits); + f.a.free(f.a_log_bits); + } + + // The gate section alone (g/beta are conv_out-independent), shared by check() and the host replay test. + pub fn checkGates(f: SplitOne, g_pairs: []const u16, beta_bits: []const u16, violation: ?*Violation) !Report { + var r: Report = .{}; + for (0..@intCast(Model.heads)) |hv| { + // HOST oracle chain: the exact bf16 rounds the shader does + const ref = gateRefOracle( + f.a_bits[hv], + f.b_bits[hv], + f32Of(f.a_log_bits, hv), + f.dt_bits[hv], + ); + const g_got = f32Of(g_pairs, hv); + if (!std.math.isFinite(g_got)) { + violated(violation, .g, hv, ref.g, g_got, 0); + return error.NonFiniteOutput; + } + const gd = @abs(g_got - ref.g); + const gbound = budget(1, 4, 2, ref.g, @abs(ref.g)); + if (gd > gbound) { + violated(violation, .g, hv, ref.g, g_got, gbound); + return error.GateOverBound; + } + const gratio = gd / gbound; + if (gratio > r.worst) r = .{ .worst = gratio, .where = hv }; + const beta_got = c.fromBf16(beta_bits[hv]); + if (!std.math.isFinite(beta_got)) { + violated(violation, .beta, hv, ref.beta, beta_got, 0); + return error.NonFiniteOutput; + } + const bd = @abs(beta_got - ref.beta); + const bbound = budget(1, 3, 1, ref.beta, @abs(ref.beta)); + if (bd > bbound) { + violated(violation, .beta, hv, ref.beta, beta_got, bbound); + return error.GateOverBound; + } + const bratio = bd / bbound; + if (bratio > r.worst) r = .{ .worst = bratio, .where = 16 + hv }; + } + return r; + } + + pub fn check(f: SplitOne, soa_bits: []const u16, g_pairs: []const u16, beta_bits: []const u16, violation: ?*Violation) !Report { + var r: Report = .{}; + const section: usize = @intCast(Model.linear_dim); + for (0..@intCast(3 * Model.heads)) |head| { + const hv: usize = @intCast(head % 16); + var vals: [128]f64 = undefined; + var ss: f64 = 0; + var vmax: f64 = 0; + for (0..128) |i| { + const cv = c.fromBf16(f.conv_out_bits[head * 128 + i]); + const sig = 1.0 / (1.0 + @exp(-cv)); + vals[i] = cv * sig; + ss += vals[i] * vals[i]; + vmax = @max(vmax, @abs(vals[i])); + } + const is_q = head < 16; + const is_k = head >= 16 and head < 32; + const inv: f64 = 1.0 / @sqrt(ss / 128.0 + 1e-6 / 128.0); + const scale_f: f64 = if (is_q) 1.0 / 128.0 else @sqrt(1.0 / 128.0); + const dst_base = (head / 16) * section + hv * 128; + for (0..128) |i| { + const want = if (is_q or is_k) scale_f * (vals[i] * inv) else vals[i]; + const got = c.fromBf16(soa_bits[dst_base + i]); + if (!std.math.isFinite(got)) { + violated(violation, .soa, dst_base + i, want, got, 0); + return error.NonFiniteOutput; + } + const d = @abs(got - want); + const bound = budget(40, 5, 3, want, vmax * vmax); + if (d > bound) { + violated(violation, .soa, dst_base + i, want, got, bound); + return error.OverBound; + } + const ratio = d / bound; + if (ratio > r.worst) r = .{ .worst = ratio, .where = dst_base + i }; + } + } + const g = try f.checkGates(g_pairs, beta_bits, violation); + if (g.worst > r.worst) r = g; + return r; + } +}; + +// ---- Conv prefill ---- + +// Prefill conv: fixture states the assembled input (state rows then qkv), weight, and checked output. +pub const Conv = struct { + a: std.mem.Allocator, + state_bits: []u16, // [3, C] + qkv_bits: []u16, // [T, C] + weight_bits: []u16, // [C, 4] + conv_input_want: []u16, // [T+3, C]: what copy_u16 must assemble + + pub fn init(a: std.mem.Allocator, seed: u64) !Conv { + const C: usize = @intCast(Model.conv_channels); + const state = try poisonedShell(a, 3 * C); + errdefer a.free(state); + const qkv = try poisonedShell(a, @intCast(tokens * C)); + errdefer a.free(qkv); + const weight = try poisonedShell(a, 4 * C); + errdefer a.free(weight); + const conv_input_want = try a.alloc(u16, @intCast((tokens + 3) * C)); + errdefer a.free(conv_input_want); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (state[0 .. 3 * C]) |*x| x.* = sampleDyadic(rng, -4, 2); + for (qkv[0..@intCast(tokens * C)]) |*x| x.* = sampleDyadic(rng, -4, 2); + for (weight[0 .. 4 * C]) |*x| x.* = sampleDyadic(rng, -3, 1); + @memcpy(conv_input_want[0 .. 3 * C], state[0 .. 3 * C]); + @memcpy(conv_input_want[3 * C ..], qkv[0..@intCast(tokens * C)]); + return .{ .a = a, .state_bits = state, .qkv_bits = qkv, .weight_bits = weight, .conv_input_want = conv_input_want }; + } + + pub fn deinit(f: Conv) void { + f.a.free(f.state_bits); + f.a.free(f.qkv_bits); + f.a.free(f.weight_bits); + f.a.free(f.conv_input_want); + } + + pub fn checkInput(f: Conv, got: []const u16) !void { + if (std.mem.indexOfDiff(u16, f.conv_input_want, got[0..f.conv_input_want.len]) != null) return error.ConvInputMismatch; + } + + pub fn checkOutput(f: Conv, bits: []const u16) !Report { + var r: Report = .{}; + const C: usize = @intCast(Model.conv_channels); + for (0..@intCast(tokens)) |t| { + for (0..C) |ch| { + var acc: f64 = 0; + for (0..4) |k| { + acc += c.fromBf16(f.conv_input_want[(t + k) * C + ch]) * c.fromBf16(f.weight_bits[ch * 4 + k]); + } + const got = c.fromBf16(bits[t * C + ch]); + if (!std.math.isFinite(got)) return error.NonFiniteOutput; + const d = @abs(got - acc); + const bound = budget(4, 0, 0, acc, @abs(acc)); + if (d > bound) return error.OverBound; + const ratio = d / bound; + if (ratio > r.worst) r = .{ .worst = ratio, .where = t * C + ch }; + } + } + return r; + } +}; diff --git a/zig/tests/vulkan/stage4_fixture_state.zig b/zig/tests/vulkan/stage4_fixture_state.zig new file mode 100644 index 00000000000..a7c60be6d51 --- /dev/null +++ b/zig/tests/vulkan/stage4_fixture_state.zig @@ -0,0 +1,391 @@ +// Stage 4 fixture: the prefill recurrence walker, the recurrence, the decode and the gated norm fixtures. + +const std = @import("std"); +const core = @import("stage4_fixture_core.zig"); + +const c = core.c; +const Model = core.Model; +const tokens = core.tokens; +const budget = core.budget; +const gateRefOracle = core.gateRefOracle; +const f32Of = core.f32Of; +const poisonedShell = core.poisonedShell; +const violated = core.violated; +const Report = core.Report; +const Output = core.Output; +const Violation = core.Violation; +const sampleDyadic = core.sampleDyadic; +const sampleDyadicF32 = core.sampleDyadicF32; + +// ---- Prefill recurrence and decode: one serial f64 walker ---- + +pub const RecurWalk = struct { + // state[head][row][dk]; y[token][head*128+row]; want_snap = FULL state after snap_after-1 (scan's snapshot slot) + pub fn run( + soa_bits: []const u16, + g_pairs: []const u16, // f32 as u16 pairs [T*16*2] + beta_bits: []const u16, + h0_bits: []const u16, // f32 pairs [262144*2] + want_y: []f64, + want_state: []f64, // [262144] + max_state: *f64, + snap_after: ?u64, + want_snap: ?[]f64, + ) void { + const dk: usize = @intCast(Model.head_dim); + const heads: usize = @intCast(Model.heads); + const state_elems: usize = @intCast(Model.state_elements); + const section: usize = @intCast(tokens * Model.linear_dim); + var st: [16 * 128 * 128]f64 = @splat(0); + for (0..heads) |h| { + for (0..dk) |row| { + for (0..dk) |i| { + st[h * dk * dk + row * dk + i] = f32Of(h0_bits, h * dk * dk + row * dk + i); + } + } + } + for (0..@intCast(tokens)) |t| { + for (0..heads) |h| { + const g: f64 = f32Of(g_pairs, t * 16 + h); + const beta: f64 = c.fromBf16(beta_bits[t * 16 + h]); + const qsec = 0 * section + (t * 16 + h) * dk; + const ksec = 1 * section + (t * 16 + h) * dk; + const vsec = 2 * section + (t * 16 + h) * dk; + for (0..dk) |row| { + var kv: f64 = 0; + for (0..dk) |i| { + const s = st[h * dk * dk + row * dk + i]; + max_state.* = @max(max_state.*, @abs(s)); + kv += (s * g) * c.fromBf16(soa_bits[ksec + i]); + } + const delta = (c.fromBf16(soa_bits[vsec + row]) - kv) * beta; + var o: f64 = 0; + for (0..dk) |i| { + const idx = h * dk * dk + row * dk + i; + const ns = (st[idx] * g) + delta * c.fromBf16(soa_bits[ksec + i]); + st[idx] = ns; + o += ns * c.fromBf16(soa_bits[qsec + i]); + } + want_y[(t * 16 + h) * dk + row] = o; + } + } + // the interior chunk boundary's snapshot slot + if (snap_after) |after| if (want_snap) |ws| { + if (@as(u64, t) + 1 == after) @memcpy(ws[0..state_elems], st[0..state_elems]); + }; + } + @memcpy(want_state[0..state_elems], st[0..state_elems]); + } +}; + +// Prefill recurrence gate: reference inputs are the ACTUAL soa/g/beta bits the GPU produced (downloaded first). +pub const Recur = struct { + a: std.mem.Allocator, + h0_bits: []u16, // f32 pairs + want_y: []f64, + want_state: []f64, + max_state: f64, + // Chunk boundary reference snapshot: non-null iff chunked (0 < chunk < tokens, chunks >= 2; boundary at `chunk`). + want_snap: ?[]f64, + + pub fn init(a: std.mem.Allocator, soa_bits: []const u16, g_pairs: []const u16, beta_bits: []const u16, seed: u64, chunk: u64) !Recur { + const state_elems: usize = @intCast(Model.state_elements); + const y_elems: usize = @intCast(tokens * Model.linear_dim); + const h0_bits = try poisonedShell(a, 2 * state_elems); + errdefer a.free(h0_bits); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (0..state_elems) |i| { + const bits: u32 = @bitCast(sampleDyadicF32(rng, -6, -2)); // small state + h0_bits[2 * i] = @truncate(bits); + h0_bits[2 * i + 1] = @truncate(bits >> 16); + } + const want_y = try a.alloc(f64, y_elems); + errdefer a.free(want_y); + const want_state = try a.alloc(f64, state_elems); + errdefer a.free(want_state); + var max_state: f64 = 0; + var want_snap: ?[]f64 = null; + errdefer if (want_snap) |ws| a.free(ws); + if (chunk > 0 and chunk < tokens) { + // chunks >= 2: exactly ONE interior boundary (after `chunk` tokens) + const ws = try a.alloc(f64, state_elems); + want_snap = ws; + RecurWalk.run(soa_bits, g_pairs, beta_bits, h0_bits, want_y, want_state, &max_state, chunk, ws); + } else { + RecurWalk.run(soa_bits, g_pairs, beta_bits, h0_bits, want_y, want_state, &max_state, null, null); + } + return .{ .a = a, .h0_bits = h0_bits, .want_y = want_y, .want_state = want_state, .max_state = max_state, .want_snap = want_snap }; + } + + pub fn deinit(f: Recur) void { + f.a.free(f.h0_bits); + f.a.free(f.want_y); + f.a.free(f.want_state); + if (f.want_snap) |ws| f.a.free(ws); + } + + // Snapshot slot vs boundary reference; poison and partial writes are refused by the same finite/budget gates as hf. + pub fn checkSnap(f: Recur, snap_pairs: []const u16, violation: ?*Violation) !Report { + const ws = f.want_snap orelse return error.NoSnapshotExpected; + var r: Report = .{}; + for (ws, 0..) |w, i| { + const got = f32Of(snap_pairs, i); + if (!std.math.isFinite(got)) { + violated(violation, .snapshot, i, w, got, 0); + return error.NonFiniteOutput; + } + const d = @abs(got - w); + const bound = budget(130 * 4 + 2, 2, 0, w, f.max_state); // chunk tokens deep + if (d > bound) { + violated(violation, .snapshot, i, w, got, bound); + return error.OverBound; + } + const ratio = d / bound; + if (ratio > r.worst) r = .{ .worst = ratio, .where = i }; + } + return r; + } + + pub fn check(f: Recur, y_bits: []const u16, hf_pairs: []const u16, violation: ?*Violation) !Report { + var r: Report = .{}; + for (f.want_y, 0..) |w, i| { + const got = c.fromBf16(y_bits[i]); + if (!std.math.isFinite(got)) { + violated(violation, .y, i, w, got, 0); + return error.NonFiniteOutput; + } + const d = @abs(got - w); + const bound = budget(130 * @as(f64, tokens) + 130, 4, 0, w, f.max_state); + if (d > bound) { + violated(violation, .y, i, w, got, bound); + return error.OverBound; + } + const ratio = d / bound; + if (ratio > r.worst) r = .{ .worst = ratio, .where = i }; + } + for (f.want_state, 0..) |w, i| { + const got = f32Of(hf_pairs, i); + if (!std.math.isFinite(got)) { + violated(violation, .state, i, w, got, 0); + return error.NonFiniteOutput; + } + const d = @abs(got - w); + const bound = budget(130 * @as(f64, tokens) + 2, 2, 0, w, f.max_state); + if (d > bound) { + violated(violation, .state, i, w, got, bound); + return error.OverBound; + } + const ratio = d / bound; + if (ratio > r.worst) r = .{ .worst = ratio, .where = i }; + } + return r; + } +}; + +// Decode gate: one raw-gates token from state_in; g/beta reference comes from a/b/A_log/dt (kernel prologue inputs). +pub const Decode = struct { + a: std.mem.Allocator, + soa_bits: []u16, // [3*2048] + a_bits: []u16, + b_bits: []u16, + a_log_bits: []u16, // f32 pairs + dt_bits: []u16, + state_in_bits: []u16, // f32 pairs + want_y: []f64, + want_state: []f64, + max_state: f64, + + pub fn init(a: std.mem.Allocator, seed: u64) !Decode { + const state_elems: usize = @intCast(Model.state_elements); + const soa_bits = try poisonedShell(a, @intCast(3 * Model.linear_dim)); + errdefer a.free(soa_bits); + const a_bits = try poisonedShell(a, @intCast(Model.heads)); + errdefer a.free(a_bits); + const b_bits = try poisonedShell(a, @intCast(Model.heads)); + errdefer a.free(b_bits); + const dt_bits = try poisonedShell(a, @intCast(Model.heads)); + errdefer a.free(dt_bits); + const a_log_bits = try poisonedShell(a, @intCast(2 * Model.heads)); + errdefer a.free(a_log_bits); + const state_in_bits = try poisonedShell(a, 2 * state_elems); + errdefer a.free(state_in_bits); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (soa_bits[0..@intCast(3 * Model.linear_dim)]) |*x| x.* = sampleDyadic(rng, -4, 2); + for (a_bits[0..@intCast(Model.heads)]) |*x| x.* = sampleDyadic(rng, -3, 2); + for (b_bits[0..@intCast(Model.heads)]) |*x| x.* = sampleDyadic(rng, -3, 2); + for (dt_bits[0..@intCast(Model.heads)]) |*x| x.* = sampleDyadic(rng, -4, 1); + for (0..@intCast(Model.heads)) |h| { + const bits: u32 = @bitCast(sampleDyadicF32(rng, -4, 1)); + a_log_bits[2 * h] = @truncate(bits); + a_log_bits[2 * h + 1] = @truncate(bits >> 16); + } + for (0..state_elems) |i| { + const bits: u32 = @bitCast(sampleDyadicF32(rng, -6, -2)); + state_in_bits[2 * i] = @truncate(bits); + state_in_bits[2 * i + 1] = @truncate(bits >> 16); + } + // serial reference + const want_y = try a.alloc(f64, @intCast(Model.linear_dim)); + errdefer a.free(want_y); + const want_state = try a.alloc(f64, state_elems); + errdefer a.free(want_state); + var max_state: f64 = 0; + const section: usize = @intCast(Model.linear_dim); + var st: [16 * 128 * 128]f64 = @splat(0); + for (0..16 * 128 * 128) |i| st[i] = f32Of(state_in_bits, i); + for (0..16) |h| { + // HOST oracle chain: the exact bf16 rounds the decode prologue does + const gref = gateRefOracle( + a_bits[h], + b_bits[h], + f32Of(a_log_bits, h), + dt_bits[h], + ); + const g = gref.g; + const beta = gref.beta; + for (0..128) |row| { + var kv: f64 = 0; + for (0..128) |i| { + const s = st[h * 128 * 128 + row * 128 + i]; + max_state = @max(max_state, @abs(s)); + kv += (s * g) * c.fromBf16(soa_bits[1 * section + h * 128 + i]); + } + const delta = (c.fromBf16(soa_bits[2 * section + h * 128 + row]) - kv) * beta; + var o: f64 = 0; + for (0..128) |i| { + const idx = h * 128 * 128 + row * 128 + i; + const ns = (st[idx] * g) + delta * c.fromBf16(soa_bits[1 * section + h * 128 + i]); + st[idx] = ns; + o += ns * c.fromBf16(soa_bits[0 * section + h * 128 + i]); + } + want_y[h * 128 + row] = o; + } + } + @memcpy(want_state[0..state_elems], st[0..state_elems]); + return .{ + .a = a, + .soa_bits = soa_bits, + .a_bits = a_bits, + .b_bits = b_bits, + .a_log_bits = a_log_bits, + .dt_bits = dt_bits, + .state_in_bits = state_in_bits, + .want_y = want_y, + .want_state = want_state, + .max_state = max_state, + }; + } + + pub fn deinit(f: Decode) void { + f.a.free(f.soa_bits); + f.a.free(f.a_bits); + f.a.free(f.b_bits); + f.a.free(f.dt_bits); + f.a.free(f.a_log_bits); + f.a.free(f.state_in_bits); + f.a.free(f.want_y); + f.a.free(f.want_state); + } + + pub fn check(f: Decode, y_bits: []const u16, state_out_pairs: []const u16, violation: ?*Violation) !Report { + var r: Report = .{}; + for (f.want_y, 0..) |w, i| { + const got = c.fromBf16(y_bits[i]); + if (!std.math.isFinite(got)) { + violated(violation, .y, i, w, got, 0); + return error.NonFiniteOutput; + } + const d = @abs(got - w); + const bound = budget(260, 4, 0, w, f.max_state); + if (d > bound) { + violated(violation, .y, i, w, got, bound); + return error.OverBound; + } + const ratio = d / bound; + if (ratio > r.worst) r = .{ .worst = ratio, .where = i }; + } + for (f.want_state, 0..) |w, i| { + const got = f32Of(state_out_pairs, i); + if (!std.math.isFinite(got)) { + violated(violation, .state, i, w, got, 0); + return error.NonFiniteOutput; + } + const d = @abs(got - w); + const bound = budget(132, 2, 0, w, f.max_state); + if (d > bound) { + violated(violation, .state, i, w, got, bound); + return error.OverBound; + } + const ratio = d / bound; + if (ratio > r.worst) r = .{ .worst = ratio, .where = i }; + } + return r; + } +}; + +// ---- Gated norm ---- + +// gatedNorm tokens=4: PER-HEAD 128-wide RMS (oracle gdn_post), [4,16,128]; out=bf16(silu_div(z)*bf16(w*(x*inv))). +pub const Norm = struct { + a: std.mem.Allocator, + y_bits: []u16, // norm input + z_bits: []u16, // gate + w_bits: []u16, // [128]: per-head weight, repeated for every head + eps: f64, + + pub fn init(a: std.mem.Allocator, seed: u64) !Norm { + const total: usize = @intCast(4 * Model.linear_dim); + const y_bits = try poisonedShell(a, total); + errdefer a.free(y_bits); + const z_bits = try poisonedShell(a, total); + errdefer a.free(z_bits); + const w_bits = try poisonedShell(a, @intCast(Model.head_dim)); + errdefer a.free(w_bits); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (y_bits[0..total]) |*x| x.* = sampleDyadic(rng, -4, 2); + for (z_bits[0..total]) |*x| x.* = sampleDyadic(rng, -4, 2); + for (w_bits[0..@intCast(Model.head_dim)]) |*x| x.* = sampleDyadic(rng, -3, 1); + return .{ .a = a, .y_bits = y_bits, .z_bits = z_bits, .w_bits = w_bits, .eps = 1e-6 }; + } + + pub fn deinit(f: Norm) void { + f.a.free(f.y_bits); + f.a.free(f.z_bits); + f.a.free(f.w_bits); + } + + // Per-head-row ref (tokens=4, 16 heads, width 128): each row of 128 has its OWN mean square; weight = column index. + pub fn check(f: Norm, bits: []const u16) !Report { + var r: Report = .{}; + const width: usize = @intCast(Model.head_dim); + for (0..4 * @as(usize, @intCast(Model.heads))) |row| { + var ss: f64 = 0; + var xmax: f64 = 0; + for (0..width) |i| { + const v = c.fromBf16(f.y_bits[row * width + i]); + ss += v * v; + xmax = @max(xmax, @abs(v)); + } + const inv: f64 = 1.0 / @sqrt(ss / @as(f64, width) + f.eps); + for (0..width) |i| { + const x = c.fromBf16(f.y_bits[row * width + i]); + const z = c.fromBf16(f.z_bits[row * width + i]); + const scaled = x * inv; + const normed = c.fromBf16(c.toBf16(@floatCast(c.fromBf16(f.w_bits[i]) * scaled))); + const silu = z / (1.0 + @exp(-z)); + const want = silu * normed; + const got = c.fromBf16(bits[row * width + i]); + if (!std.math.isFinite(got)) return error.NonFiniteOutput; + const d = @abs(got - want); + const bound = budget(128, 3, 2, want, xmax * xmax * @as(f64, width)); + if (d > bound) return error.OverBound; + const ratio = d / bound; + if (ratio > r.worst) r = .{ .worst = ratio, .where = row * width + i }; + } + } + return r; + } +}; diff --git a/zig/tests/vulkan/stage4_probe.zig b/zig/tests/vulkan/stage4_probe.zig new file mode 100644 index 00000000000..17c1a2f3b6c --- /dev/null +++ b/zig/tests/vulkan/stage4_probe.zig @@ -0,0 +1,329 @@ +// Stage 4: GatedDeltaNet (conv prefill, gate split, recurrence, raw-gates decode, RMSNorm) via Gdn; exits 0/1/2/124. + +const std = @import("std"); +const vk = @import("vulkan"); +const shaders = @import("stage4_shaders"); +const fx = @import("stage4_fixture.zig"); +const shared = @import("stage4_probe_shared.zig"); +const gates = @import("stage4_probe_gates.zig"); + +const Elapsed = shared.Elapsed; +const repeats = shared.repeats; +const uploadShell = shared.uploadShell; +const fetch = shared.fetch; +const unchanged = shared.unchanged; +const printSplitGateInputs = shared.printSplitGateInputs; + +const usage = "usage: tf-stage4-gpu [--vendor ID] [--timeout-ms MS] | --help\n" ++ + " --help show this text and exit without opening a loader or device\n" ++ + " --vendor ID refuse a device whose Vulkan vendorID differs (decimal or 0x hex; no default)\n" ++ + " --timeout-ms MS bound on every fence wait, 1 to 600000 (default 10000); a timeout ends the process\n"; + +pub fn main(init: std.process.Init) !u8 { + const clock: Elapsed = .{ .io = init.io, .start = std.Io.Clock.awake.now(init.io) }; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + for (args[1..]) |arg| if (std.mem.eql(u8, arg, "--help")) { + std.debug.print(usage, .{}); + return 0; + }; + var vendor: ?u32 = null; + var timeout_ms: u64 = 10_000; + var i: usize = 1; + while (i < args.len) : (i += 2) { + const value = if (i + 1 < args.len) args[i + 1] else ""; + if (std.mem.eql(u8, args[i], "--vendor")) { + vendor = std.fmt.parseInt(u32, value, 0) catch null; + if (vendor == null) break; + } else if (std.mem.eql(u8, args[i], "--timeout-ms")) { + timeout_ms = std.fmt.parseInt(u64, value, 10) catch 0; + if (timeout_ms == 0 or timeout_ms > 600_000) break; + } else break; + } else return if (try probe(init.gpa, clock, vendor, timeout_ms * std.time.ns_per_ms)) 0 else 1; + std.debug.print(usage, .{}); + return 2; +} +fn probe(a: std.mem.Allocator, clock: Elapsed, vendor: ?u32, timeout: u64) !bool { + var rt = try vk.Runtime.open(); + defer rt.close(); + var dev = try vk.Device.open(&rt, .{ .features = .{ .storage16 = true, .int16 = true } }); + defer dev.close(); + const pr = dev.properties; + const name = std.mem.sliceTo(&pr.deviceName, 0); + const v = pr.apiVersion; + std.debug.print("MILESTONE device-found elapsed_ms={d} name=\"{s}\" vendor=0x{x:0>4} device=0x{x:0>4} " ++ + "api={d}.{d}.{d} storage16={} int16={} subgroups={d} shuffle-relative={} subgroup-size={d} fence-timeout={d}ms\n", .{ + clock.ms(), name, + pr.vendorID, pr.deviceID, + v >> 22 & 0x7f, v >> 12 & 0x3ff, + v & 0xfff, dev.features.storage16, + dev.features.int16, dev.v11.subgroupSupportedStages, + dev.v11.subgroupSupportedOperations & vk.c.VK_SUBGROUP_FEATURE_SHUFFLE_RELATIVE_BIT != 0, dev.v11.subgroupSize, + timeout / std.time.ns_per_ms, + }); + if (pr.deviceType == vk.c.VK_PHYSICAL_DEVICE_TYPE_CPU) { + std.debug.print("FAIL device: a CPU (software) Vulkan device is never a GPU result\n", .{}); + return false; + } + if (vendor) |want| if (pr.vendorID != want) { + std.debug.print("FAIL device: vendor 0x{x:0>4} is not the requested 0x{x:0>4}\n", .{ pr.vendorID, want }); + return false; + }; + + var cmds = try vk.Commands.init(&dev); + defer cmds.deinit(); + var staging = try vk.Buffer.init(&dev, .staging, 4 << 20); + defer staging.deinit(); + var m = try vk.Gdn.init(&dev, .{ + .conv_dw1d_bf16 = &shaders.conv_dw1d_bf16, + .gdn_conv_decode_bf16 = &shaders.gdn_conv_decode_bf16, + .gdn_split_bf16 = &shaders.gdn_split_bf16, + .prefill_recur_bf16 = &shaders.prefill_recur_bf16, + .gated_delta_decode_bf16 = &shaders.gated_delta_decode_bf16, + .fast_norm_gated_bf16 = &shaders.fast_norm_gated_bf16, + .copy_u16 = &shaders.copy_u16, + .conv_dw1d_f32 = &shaders.conv_dw1d_f32, + .gdn_split_f32 = &shaders.gdn_split_f32, + .prefill_recur_f32 = &shaders.prefill_recur_f32, + .fast_norm_gated_f32 = &shaders.fast_norm_gated_f32, + .copy_u32 = &shaders.gdn_copy_u32, + }); + defer m.deinit(); + + var passed: usize = 0; + var total: usize = 0; + total += 1; + passed += @intFromBool(try convSplitGate(a, clock, &dev, &cmds, staging, &m, timeout)); + total += 1; + passed += @intFromBool(try gates.prefillGate(a, clock, &dev, &cmds, staging, &m, 8, null, timeout)); + total += 1; + passed += @intFromBool(try gates.prefillGate(a, clock, &dev, &cmds, staging, &m, 4, 1, timeout)); + total += 1; + passed += @intFromBool(try gates.decodeGate(a, clock, &dev, &cmds, staging, &m, timeout)); + total += 1; + passed += @intFromBool(try gates.convDecodeGate(a, clock, &dev, &cmds, staging, &m, timeout)); + total += 1; + passed += @intFromBool(try gates.normGate(a, clock, &dev, &cmds, staging, &m, timeout)); + std.debug.print("RESULT stage4 elapsed_ms={d} on \"{s}\": passed={d} total={d}\n", .{ clock.ms(), name, passed, total }); + return passed == total; +} +// convPrefill: conv input + state carry-out BIT-EXACT; conv output and split soa/g/beta vs f64; repeats and refusals. +fn convSplitGate(a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, m: *vk.Gdn, timeout: u64) !bool { + const T = fx.tokens; + const C: u64 = fx.Model.conv_channels; + const soa_elems: u64 = 3 * T * fx.Model.linear_dim; + const gate_elems: u64 = T * fx.Model.heads; + const f = try fx.Split.init(a, 0x57a4); + defer f.deinit(); + const cf = try fx.Conv.init(a, 0x57a4 + 1); + defer cf.deinit(); + + var state = try uploadShell(dev, cmds, staging, cf.state_bits, timeout); + defer state.deinit(); + var qkv = try uploadShell(dev, cmds, staging, cf.qkv_bits, timeout); + defer qkv.deinit(); + var weight = try uploadShell(dev, cmds, staging, cf.weight_bits, timeout); + defer weight.deinit(); + var conv_input = try vk.Buffer.init(dev, .device, @intCast((T + 3) * C * 2)); + defer conv_input.deinit(); + var conv_out = try vk.Buffer.init(dev, .device, @intCast(T * C * 2)); + defer conv_out.deinit(); + var soa = try vk.Buffer.init(dev, .device, @intCast(soa_elems * 2)); + defer soa.deinit(); + var g_out = try vk.Buffer.init(dev, .device, @intCast(gate_elems * 4)); + defer g_out.deinit(); + var beta_out = try vk.Buffer.init(dev, .device, @intCast(gate_elems * 2)); + defer beta_out.deinit(); + var state_out = try vk.Buffer.init(dev, .device, @intCast(3 * C * 2)); + defer state_out.deinit(); + var a_log = try uploadShell(dev, cmds, staging, f.a_log_bits, timeout); + defer a_log.deinit(); + + const conv_input_view = try vk.tensor.View.init(&conv_input, .bf16, 0, (T + 3) * C); + const state_view = try vk.tensor.View.init(&state, .bf16, 0, 3 * C); + const qkv_view = try vk.tensor.View.init(&qkv, .bf16, 0, T * C); + const weight_view = try vk.tensor.View.init(&weight, .bf16, 0, 4 * C); + const conv_out_view = try vk.tensor.View.init(&conv_out, .bf16, 0, T * C); + const soa_view = try vk.tensor.View.init(&soa, .bf16, 0, soa_elems); + const g_view = try vk.tensor.View.init(&g_out, .f32, 0, gate_elems); + const beta_view = try vk.tensor.View.init(&beta_out, .bf16, 0, gate_elems); + var a_buf = try uploadShell(dev, cmds, staging, f.a_bits, timeout); + defer a_buf.deinit(); + var b_buf = try uploadShell(dev, cmds, staging, f.b_bits, timeout); + defer b_buf.deinit(); + var dt_buf = try uploadShell(dev, cmds, staging, f.dt_bits, timeout); + defer dt_buf.deinit(); + const a_log_view = try vk.tensor.View.init(&a_log, .f32, 0, fx.Model.heads); + const a_view2 = try vk.tensor.View.init(&a_buf, .bf16, 0, gate_elems); + const b_view = try vk.tensor.View.init(&b_buf, .bf16, 0, gate_elems); + const dt_view = try vk.tensor.View.init(&dt_buf, .bf16, 0, fx.Model.heads); + const state_out_view = try vk.tensor.View.init(&state_out, .bf16, 0, 3 * C); + + var first: ?[]u16 = null; + defer if (first) |x| a.free(x); + var g_first: ?[]u16 = null; + defer if (g_first) |x| a.free(x); + var beta_first: ?[]u16 = null; + defer if (beta_first) |x| a.free(x); + var conv_first: ?[]u16 = null; + defer if (conv_first) |x| a.free(x); + var state_first: ?[]u16 = null; + defer if (state_first) |x| a.free(x); + // poison shells for EVERY output: soa, g, beta, conv_out, state_out + const init16 = try a.alloc(u16, @intCast(soa_elems)); + defer a.free(init16); + @memset(init16, fx.poison16); + const init32 = try a.alloc(u16, @intCast(2 * gate_elems)); + defer a.free(init32); + @memset(init32, 0xffff); + const init_beta = try a.alloc(u16, @intCast(gate_elems)); + defer a.free(init_beta); + @memset(init_beta, fx.poison16); + const init_conv = try a.alloc(u16, @intCast(T * C)); + defer a.free(init_conv); + @memset(init_conv, fx.poison16); + const init_state = try a.alloc(u16, @intCast(3 * C)); + defer a.free(init_state); + @memset(init_state, fx.poison16); + // per-repeat readbacks of every output, compared byte-for-byte + const soa_repeat = try a.alloc(u16, @intCast(soa_elems)); + defer a.free(soa_repeat); + const g_repeat = try a.alloc(u16, @intCast(2 * gate_elems)); + defer a.free(g_repeat); + const beta_repeat = try a.alloc(u16, @intCast(gate_elems)); + defer a.free(beta_repeat); + const conv_repeat = try a.alloc(u16, @intCast(T * C)); + defer a.free(conv_repeat); + const state_repeat = try a.alloc(u16, @intCast(3 * C)); + defer a.free(state_repeat); + for (0..repeats) |rep| { + try vk.orExit(vk.upload(cmds, soa, 0, staging, std.mem.sliceAsBytes(init16[0..@intCast(soa_elems)]), timeout)); + try vk.orExit(vk.upload(cmds, g_out, 0, staging, std.mem.sliceAsBytes(init32), timeout)); + try vk.orExit(vk.upload(cmds, beta_out, 0, staging, std.mem.sliceAsBytes(init_beta), timeout)); + try vk.orExit(vk.upload(cmds, conv_out, 0, staging, std.mem.sliceAsBytes(init_conv), timeout)); + try vk.orExit(vk.upload(cmds, state_out, 0, staging, std.mem.sliceAsBytes(init_state), timeout)); + try vk.orExit(m.convPrefill(cmds, soa_view, g_view, beta_view, conv_out_view, conv_input_view, state_view, qkv_view, weight_view, a_view2, b_view, a_log_view, dt_view, state_out_view, T, timeout)); + try fetch(cmds, staging, soa, std.mem.sliceAsBytes(soa_repeat), timeout); + try fetch(cmds, staging, g_out, std.mem.sliceAsBytes(g_repeat), timeout); + try fetch(cmds, staging, beta_out, std.mem.sliceAsBytes(beta_repeat), timeout); + try fetch(cmds, staging, conv_out, std.mem.sliceAsBytes(conv_repeat), timeout); + try fetch(cmds, staging, state_out, std.mem.sliceAsBytes(state_repeat), timeout); + if (first) |x| { + if (std.mem.indexOfDiff(u16, x, soa_repeat)) |at| { + std.debug.print("FAIL conv-split: repeat {d} soa differs at element {d}\n", .{ rep, at }); + return false; + } + // compare the remaining outputs against the FIRST repeat's copies + if (g_first) |gf| { + if (std.mem.indexOfDiff(u16, gf, g_repeat)) |at| { + std.debug.print("FAIL conv-split: repeat {d} g differs at element {d}\n", .{ rep, at }); + return false; + } + } else g_first = try a.dupe(u16, g_repeat); + if (beta_first) |bf| { + if (std.mem.indexOfDiff(u16, bf, beta_repeat)) |at| { + std.debug.print("FAIL conv-split: repeat {d} beta differs at element {d}\n", .{ rep, at }); + return false; + } + } else beta_first = try a.dupe(u16, beta_repeat); + if (conv_first) |cf2| { + if (std.mem.indexOfDiff(u16, cf2, conv_repeat)) |at| { + std.debug.print("FAIL conv-split: repeat {d} conv_out differs at element {d}\n", .{ rep, at }); + return false; + } + } else conv_first = try a.dupe(u16, conv_repeat); + if (state_first) |sf2| { + if (std.mem.indexOfDiff(u16, sf2, state_repeat)) |at| { + std.debug.print("FAIL conv-split: repeat {d} state_out differs at element {d}\n", .{ rep, at }); + return false; + } + } else state_first = try a.dupe(u16, state_repeat); + } else { + first = try a.dupe(u16, soa_repeat); + g_first = try a.dupe(u16, g_repeat); + beta_first = try a.dupe(u16, beta_repeat); + conv_first = try a.dupe(u16, conv_repeat); + state_first = try a.dupe(u16, state_repeat); + } + } + // download and check everything + const soa_bits = try a.alloc(u16, @intCast(soa_elems)); + defer a.free(soa_bits); + const g_pairs = try a.alloc(u16, @intCast(2 * gate_elems)); + defer a.free(g_pairs); + const beta_bits = try a.alloc(u16, @intCast(gate_elems)); + defer a.free(beta_bits); + const conv_input_bits = try a.alloc(u16, @intCast((T + 3) * C)); + defer a.free(conv_input_bits); + const conv_out_bits = try a.alloc(u16, @intCast(T * C)); + defer a.free(conv_out_bits); + const state_out_bits = try a.alloc(u16, @intCast(3 * C)); + defer a.free(state_out_bits); + try fetch(cmds, staging, soa, std.mem.sliceAsBytes(soa_bits), timeout); + try fetch(cmds, staging, g_out, std.mem.sliceAsBytes(g_pairs), timeout); + try fetch(cmds, staging, beta_out, std.mem.sliceAsBytes(beta_bits), timeout); + try fetch(cmds, staging, conv_input, std.mem.sliceAsBytes(conv_input_bits), timeout); + try fetch(cmds, staging, conv_out, std.mem.sliceAsBytes(conv_out_bits), timeout); + try fetch(cmds, staging, state_out, std.mem.sliceAsBytes(state_out_bits), timeout); + + // The split's reference reads the ACTUAL conv output the kernel produced, not the synthetic conv_out_bits. + @memcpy(f.conv_out_bits[0..conv_out_bits.len], conv_out_bits); + var sr_diag: fx.Violation = undefined; + const sr = fx.Split.check(f, soa_bits, g_pairs, beta_bits, &sr_diag) catch |e| { + std.debug.print("FAIL conv-split: split {s} first at index {d}: want={e} got={e} bound={e} ratio={e} want_f32_bits={x}\n", .{ @errorName(e), sr_diag.index, sr_diag.want, sr_diag.got, sr_diag.bound, sr_diag.ratio, sr_diag.want_f32_bits }); + printSplitGateInputs(f, sr_diag); + return false; + }; + fx.Conv.checkInput(cf, conv_input_bits) catch |e| { + std.debug.print("FAIL conv-split: conv input {s}\n", .{@errorName(e)}); + return false; + }; + const cr = fx.Conv.checkOutput(cf, conv_out_bits) catch |e| { + std.debug.print("FAIL conv-split: conv {s}\n", .{@errorName(e)}); + return false; + }; + // the state carry-out is conv input rows tokens..tokens+3, bit-exact + if (std.mem.indexOfDiff(u16, cf.conv_input_want[@intCast(T * C)..], state_out_bits) != null) { + std.debug.print("FAIL conv-split: the conv state carry-out is not the trailing window\n", .{}); + return false; + } + if (!try unchanged(a, cmds, staging, cf.state_bits, state, timeout) or + !try unchanged(a, cmds, staging, cf.qkv_bits, qkv, timeout) or + !try unchanged(a, cmds, staging, cf.weight_bits, weight, timeout) or + !try unchanged(a, cmds, staging, f.a_bits, a_buf, timeout) or + !try unchanged(a, cmds, staging, f.b_bits, b_buf, timeout) or + !try unchanged(a, cmds, staging, f.dt_bits, dt_buf, timeout) or + !try unchanged(a, cmds, staging, f.a_log_bits, a_log, timeout)) + { + std.debug.print("FAIL conv-split: an input changed\n", .{}); + return false; + } + { + const mutated = try a.dupe(u16, soa_bits); + defer a.free(mutated); + const gmut = try a.dupe(u16, g_pairs); + defer a.free(gmut); + var first_finite: ?usize = null; + for (soa_bits, 0..) |bits, i| { + if (std.math.isFinite(fx.c.fromBf16(bits))) { + first_finite = i; + break; + } + } + const at = first_finite orelse return error.NoFiniteReference; + mutated[at] = 0x4f00; // 2^31: an absolutely wrong finite the budgets must refuse + if (fx.Split.check(f, mutated, g_pairs, beta_bits, null)) |_| { + std.debug.print("FAIL conv-split: a planted wrong finite split output passed\n", .{}); + return false; + } else |_| {} + for ([_]u16{ 0x7fc0, 0x7f80, 0xff80 }) |pb| { + mutated[at] = pb; + if (fx.Split.check(f, mutated, g_pairs, beta_bits, null)) |_| { + std.debug.print("FAIL conv-split: a planted nonfinite split output passed\n", .{}); + return false; + } else |_| {} + } + } + std.debug.print("MILESTONE conv-split elapsed_ms={d} tokens={d} split worst={e} conv worst={e} of bounds, " ++ + "conv input and state carry-out bit-exact, repeats byte-identical, inputs unchanged, planted outputs refused\n", .{ clock.ms(), T, sr.worst, cr.worst }); + return true; +} diff --git a/zig/tests/vulkan/stage4_probe_gates.zig b/zig/tests/vulkan/stage4_probe_gates.zig new file mode 100644 index 00000000000..3b855137575 --- /dev/null +++ b/zig/tests/vulkan/stage4_probe_gates.zig @@ -0,0 +1,524 @@ +// Stage 4 probe: the conv-split, prefill, decode, conv-decode and gated-norm GPU gates. + +const std = @import("std"); +const vk = @import("vulkan"); +const fx = @import("stage4_fixture.zig"); +const shared = @import("stage4_probe_shared.zig"); + +const Elapsed = shared.Elapsed; +const repeats = shared.repeats; +const state_elems = shared.state_elems; +const uploadShell = shared.uploadShell; +const fetch = shared.fetch; +const unchanged = shared.unchanged; +const plantRefuse = shared.plantRefuse; +const plantRefuse2 = shared.plantRefuse2; +const printGateInputs = shared.printGateInputs; +const printGateInputsOne = shared.printGateInputsOne; +const printSplitGateInputs = shared.printSplitGateInputs; +const printRecurInputs = shared.printRecurInputs; + +const SharedSplit = struct { + a: std.mem.Allocator, + f: fx.Split, + soa_buf: vk.Buffer, + g_buf: vk.Buffer, + beta_buf: vk.Buffer, + soa_bits: []u16, + g_pairs: []u16, + beta_bits: []u16, + soa_view: vk.tensor.View, + g_view: vk.tensor.View, + beta_view: vk.tensor.View, + a_log_buf: vk.Buffer, + a_buf: vk.Buffer, + b_buf: vk.Buffer, + dt_buf: vk.Buffer, + + fn deinit(s: *SharedSplit) void { + s.f.deinit(); + s.soa_buf.deinit(); + s.g_buf.deinit(); + s.beta_buf.deinit(); + s.a_log_buf.deinit(); + s.a_buf.deinit(); + s.b_buf.deinit(); + s.dt_buf.deinit(); + s.a.free(s.soa_bits); + s.a.free(s.g_pairs); + s.a.free(s.beta_bits); + } +}; + +// The split runs standalone into a caller-owned struct (its views must outlive it); the gates read real GPU soa/g/beta. +pub fn runSplit(s: *SharedSplit, a: std.mem.Allocator, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, m: *vk.Gdn, timeout: u64) !void { + const T = fx.tokens; + const soa_elems: u64 = 3 * T * fx.Model.linear_dim; + const gate_elems: u64 = T * fx.Model.heads; + // staged construction: each field owns its errdefer, so a failure never runs deinit on undefined handles + s.a = a; + s.f = try fx.Split.init(a, 0x57a4); + errdefer s.f.deinit(); + s.soa_bits = try a.alloc(u16, @intCast(soa_elems)); + errdefer a.free(s.soa_bits); + s.g_pairs = try a.alloc(u16, @intCast(2 * gate_elems)); + errdefer a.free(s.g_pairs); + s.beta_bits = try a.alloc(u16, @intCast(gate_elems)); + errdefer a.free(s.beta_bits); + s.soa_buf = try vk.Buffer.init(dev, .device, @intCast(soa_elems * 2)); + errdefer s.soa_buf.deinit(); + s.g_buf = try vk.Buffer.init(dev, .device, @intCast(gate_elems * 4)); + errdefer s.g_buf.deinit(); + s.beta_buf = try vk.Buffer.init(dev, .device, @intCast(gate_elems * 2)); + errdefer s.beta_buf.deinit(); + s.soa_view = undefined; + s.g_view = undefined; + s.beta_view = undefined; + s.a_log_buf = try uploadShell(dev, cmds, staging, s.f.a_log_bits, timeout); + errdefer s.a_log_buf.deinit(); + s.a_buf = try uploadShell(dev, cmds, staging, s.f.a_bits, timeout); + errdefer s.a_buf.deinit(); + s.b_buf = try uploadShell(dev, cmds, staging, s.f.b_bits, timeout); + errdefer s.b_buf.deinit(); + s.dt_buf = try uploadShell(dev, cmds, staging, s.f.dt_bits, timeout); + errdefer s.dt_buf.deinit(); + s.soa_view = try vk.tensor.View.init(&s.soa_buf, .bf16, 0, soa_elems); + s.g_view = try vk.tensor.View.init(&s.g_buf, .f32, 0, gate_elems); + s.beta_view = try vk.tensor.View.init(&s.beta_buf, .bf16, 0, gate_elems); + var conv_out_buf = try uploadShell(dev, cmds, staging, s.f.conv_out_bits, timeout); + defer conv_out_buf.deinit(); + // poison ALL THREE split outputs before the dispatch: a dead or partial g/beta write fails the check. + const poison = try a.alloc(u16, @intCast(soa_elems)); + defer a.free(poison); + @memset(poison, fx.poison16); + const poison_g = try a.alloc(u16, @intCast(2 * gate_elems)); + defer a.free(poison_g); + @memset(poison_g, 0xffff); + const poison_beta = try a.alloc(u16, @intCast(gate_elems)); + defer a.free(poison_beta); + @memset(poison_beta, fx.poison16); + try vk.orExit(vk.upload(cmds, s.soa_buf, 0, staging, std.mem.sliceAsBytes(poison), timeout)); + try vk.orExit(vk.upload(cmds, s.g_buf, 0, staging, std.mem.sliceAsBytes(poison_g), timeout)); + try vk.orExit(vk.upload(cmds, s.beta_buf, 0, staging, std.mem.sliceAsBytes(poison_beta), timeout)); + try vk.orExit(m.split(cmds, s.soa_view, s.g_view, s.beta_view, try vk.tensor.View.init(&conv_out_buf, .bf16, 0, T * fx.Model.conv_channels), try vk.tensor.View.init(&s.a_buf, .bf16, 0, gate_elems), try vk.tensor.View.init(&s.b_buf, .bf16, 0, gate_elems), try vk.tensor.View.init(&s.a_log_buf, .f32, 0, fx.Model.heads), try vk.tensor.View.init(&s.dt_buf, .bf16, 0, fx.Model.heads), T, timeout)); + try fetch(cmds, staging, s.soa_buf, std.mem.sliceAsBytes(s.soa_bits), timeout); + try fetch(cmds, staging, s.g_buf, std.mem.sliceAsBytes(s.g_pairs), timeout); + try fetch(cmds, staging, s.beta_buf, std.mem.sliceAsBytes(s.beta_bits), timeout); + var s_diag: fx.Violation = undefined; + const srep = s.f.check(s.soa_bits, s.g_pairs, s.beta_bits, &s_diag) catch |e| { + std.debug.print("FAIL split standalone: {s} first at index {d}: want={e} got={e} bound={e} ratio={e} want_f32_bits={x}\n", .{ @errorName(e), s_diag.index, s_diag.want, s_diag.got, s_diag.bound, s_diag.ratio, s_diag.want_f32_bits }); + printSplitGateInputs(s.f, s_diag); + return error.SplitFailed; + }; + std.debug.print("MILESTONE split-standalone worst={e} of bound\n", .{srep.worst}); +} + +// One prefill recurrence gate at chunk size `chunk` (snap slots when >1) vs the f64 reference over GPU soa/g/beta. +pub fn prefillGate(a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, m: *vk.Gdn, chunk: u64, snap_chunks: ?u64, timeout: u64) !bool { + const T = fx.tokens; + const section: u64 = T * fx.Model.linear_dim; + var split: SharedSplit = undefined; + runSplit(&split, a, dev, cmds, staging, m, timeout) catch return false; + defer split.deinit(); + const rf = try fx.Recur.init(a, split.soa_bits, split.g_pairs, split.beta_bits, 0x57a7, chunk); + defer rf.deinit(); + + var h0 = try uploadShell(dev, cmds, staging, rf.h0_bits, timeout); + defer h0.deinit(); + var y = try vk.Buffer.init(dev, .device, @intCast(section * 2)); + defer y.deinit(); + var hf = try vk.Buffer.init(dev, .device, @intCast(state_elems * 4)); + defer hf.deinit(); + var snap: ?vk.Buffer = null; + defer if (snap) |*s| s.deinit(); + if (snap_chunks) |n| { + snap = try vk.Buffer.init(dev, .device, @intCast((n * state_elems) * 4)); + } + const h0_view = try vk.tensor.View.init(&h0, .f32, 0, state_elems); + const y_view = try vk.tensor.View.init(&y, .bf16, 0, section); + const hf_view = try vk.tensor.View.init(&hf, .f32, 0, state_elems); + const snap_view: ?vk.tensor.View = if (snap) |*s| try vk.tensor.View.init(s, .f32, 0, snap_chunks.? * state_elems) else null; + + var r_diag: fx.Violation = undefined; + // poison y and hf/snap before EVERY repeat: a dead or partial store leaves poison and fails; readback vs f64. + const poison = try a.alloc(u16, @intCast(section)); + defer a.free(poison); + @memset(poison, fx.poison16); + const poison_state = try a.alloc(u16, @intCast(2 * state_elems)); + defer a.free(poison_state); + @memset(poison_state, 0xffff); + var first: ?[]u16 = null; + defer if (first) |x| a.free(x); + var hf_first: ?[]u16 = null; + defer if (hf_first) |x| a.free(x); + var snap_first: ?[]u16 = null; + defer if (snap_first) |x| a.free(x); + const ybits = try a.alloc(u16, @intCast(section)); + defer a.free(ybits); + const hfpairs = try a.alloc(u16, @intCast(2 * state_elems)); + defer a.free(hfpairs); + const snappairs = try a.alloc(u16, @intCast(2 * state_elems)); + defer a.free(snappairs); + for (0..repeats) |rep| { + try vk.orExit(vk.upload(cmds, y, 0, staging, std.mem.sliceAsBytes(poison), timeout)); + try vk.orExit(vk.upload(cmds, hf, 0, staging, std.mem.sliceAsBytes(poison_state), timeout)); + if (snap) |*s| try vk.orExit(vk.upload(cmds, s.*, 0, staging, std.mem.sliceAsBytes(poison_state), timeout)); + try vk.orExit(m.prefillRecur(cmds, split.soa_view, split.g_view, split.beta_view, h0_view, y_view, hf_view, snap_view, T, chunk, timeout)); + try fetch(cmds, staging, y, std.mem.sliceAsBytes(ybits), timeout); + try fetch(cmds, staging, hf, std.mem.sliceAsBytes(hfpairs), timeout); + // EVERY repeat: y byte-equal to the first AND hf re-checked against the f64 reference. + if (first) |x| { + if (std.mem.indexOfDiff(u16, x, ybits)) |at| { + std.debug.print("FAIL prefill chunk={d}: repeat {d} y differs at element {d}\n", .{ chunk, rep, at }); + return false; + } + if (hf_first) |hf0| { + if (std.mem.indexOfDiff(u16, hf0, hfpairs)) |at| { + std.debug.print("FAIL prefill chunk={d}: repeat {d} hf differs at element {d}\n", .{ chunk, rep, at }); + return false; + } + } else hf_first = try a.dupe(u16, hfpairs); + } else { + first = try a.dupe(u16, ybits); + hf_first = try a.dupe(u16, hfpairs); + } + _ = rf.check(ybits, hfpairs, &r_diag) catch |e| { + std.debug.print("FAIL prefill chunk={d}: repeat {d} {s} first at index {d}: want={e} got={e} bound={e} ratio={e} want_f32_bits={x}\n", .{ chunk, rep, @errorName(e), r_diag.index, r_diag.want, r_diag.got, r_diag.bound, r_diag.ratio, r_diag.want_f32_bits }); + printRecurInputs(split.f, r_diag, T - 1); + return false; + }; + // the snapshot slot: downloaded, checked against the walker's boundary state, and byte-compared across repeats + if (snap) |*s| { + try fetch(cmds, staging, s.*, std.mem.sliceAsBytes(snappairs), timeout); + if (snap_first) |sp0| { + if (std.mem.indexOfDiff(u16, sp0, snappairs)) |at| { + std.debug.print("FAIL prefill chunk={d}: repeat {d} snap differs at element {d}\n", .{ chunk, rep, at }); + return false; + } + } else snap_first = try a.dupe(u16, snappairs); + _ = rf.checkSnap(snappairs, &r_diag) catch |e| { + std.debug.print("FAIL prefill chunk={d}: repeat {d} snap {s} first at index {d}: want={e} got={e} bound={e} ratio={e}\n", .{ chunk, rep, @errorName(e), r_diag.index, r_diag.want, r_diag.got, r_diag.bound, r_diag.ratio }); + printRecurInputs(split.f, r_diag, chunk - 1); + return false; + }; + } + } + const rrep = rf.check(ybits, hfpairs, &r_diag) catch |e| { + std.debug.print("FAIL prefill chunk={d}: {s} first at index {d}: want={e} got={e} bound={e} ratio={e} want_f32_bits={x}\n", .{ chunk, @errorName(e), r_diag.index, r_diag.want, r_diag.got, r_diag.bound, r_diag.ratio, r_diag.want_f32_bits }); + printRecurInputs(split.f, r_diag, T - 1); + return false; + }; + if (!try unchanged(a, cmds, staging, rf.h0_bits, h0, timeout) or + !try unchanged(a, cmds, staging, split.soa_bits, split.soa_buf, timeout)) + { + std.debug.print("FAIL prefill chunk={d}: an input changed\n", .{chunk}); + return false; + } + if (!try plantRefuse2(rf, first.?, hfpairs, "y", "prefill")) return false; + std.debug.print("MILESTONE prefill elapsed_ms={d} tokens={d} chunk={d} chunks={?d} worst={e} of bound, " ++ + "hf matches independent f64 reference from poison, repeats byte-identical, h0 and soa unchanged, planted outputs refused\n", .{ clock.ms(), T, chunk, snap_chunks, rrep.worst }); + return true; +} + +// The raw-gates decode step: one token, state in/out and y vs the serial reference from the kernel's raw inputs. +pub fn decodeGate(a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, m: *vk.Gdn, timeout: u64) !bool { + const linear: u64 = fx.Model.linear_dim; + const f = try fx.Decode.init(a, 0x57a8); + defer f.deinit(); + var soa = try uploadShell(dev, cmds, staging, f.soa_bits, timeout); + defer soa.deinit(); + var a_buf = try uploadShell(dev, cmds, staging, f.a_bits, timeout); + defer a_buf.deinit(); + var b_buf = try uploadShell(dev, cmds, staging, f.b_bits, timeout); + defer b_buf.deinit(); + var a_log = try uploadShell(dev, cmds, staging, f.a_log_bits, timeout); + defer a_log.deinit(); + var dt = try uploadShell(dev, cmds, staging, f.dt_bits, timeout); + defer dt.deinit(); + var state_in = try uploadShell(dev, cmds, staging, f.state_in_bits, timeout); + defer state_in.deinit(); + var state_out = try vk.Buffer.init(dev, .device, @intCast(state_elems * 4)); + defer state_out.deinit(); + var y = try vk.Buffer.init(dev, .device, @intCast(linear * 2)); + defer y.deinit(); + + const soa_view = try vk.tensor.View.init(&soa, .bf16, 0, 3 * linear); + const a_view = try vk.tensor.View.init(&a_buf, .bf16, 0, fx.Model.heads); + const b_view = try vk.tensor.View.init(&b_buf, .bf16, 0, fx.Model.heads); + const a_log_view = try vk.tensor.View.init(&a_log, .u16, 0, 2 * fx.Model.heads); + const dt_view = try vk.tensor.View.init(&dt, .bf16, 0, fx.Model.heads); + const state_in_view = try vk.tensor.View.init(&state_in, .f32, 0, state_elems); + const state_out_view = try vk.tensor.View.init(&state_out, .f32, 0, state_elems); + const y_view = try vk.tensor.View.init(&y, .bf16, 0, linear); + + const poison = try a.alloc(u16, @intCast(linear)); + defer a.free(poison); + @memset(poison, fx.poison16); + // poison state_out before EVERY repeat: a dead store leaves poison, which the f64 comparison refuses. + const poison_state = try a.alloc(u16, @intCast(2 * state_elems)); + defer a.free(poison_state); + @memset(poison_state, 0xffff); + var first: ?[]u16 = null; + defer if (first) |x| a.free(x); + const ybits = try a.alloc(u16, @intCast(linear)); + defer a.free(ybits); + const souts = try a.alloc(u16, @intCast(2 * state_elems)); + defer a.free(souts); + for (0..repeats) |rep| { + try vk.orExit(vk.upload(cmds, y, 0, staging, std.mem.sliceAsBytes(poison), timeout)); + try vk.orExit(vk.upload(cmds, state_out, 0, staging, std.mem.sliceAsBytes(poison_state), timeout)); + try vk.orExit(m.decodeRecur(cmds, soa_view, a_view, b_view, a_log_view, dt_view, state_in_view, state_out_view, y_view, timeout)); + try fetch(cmds, staging, y, std.mem.sliceAsBytes(ybits), timeout); + try fetch(cmds, staging, state_out, std.mem.sliceAsBytes(souts), timeout); + if (first) |x| { + if (std.mem.indexOfDiff(u16, x, ybits)) |at| { + std.debug.print("FAIL decode: repeat {d} differs at element {d}\n", .{ rep, at }); + return false; + } + } else first = try a.dupe(u16, ybits); + } + var d_diag: fx.Violation = undefined; + const drep = f.check(ybits, souts, &d_diag) catch |e| { + std.debug.print("FAIL decode: {s} first at index {d}: want={e} got={e} bound={e} ratio={e} want_f32_bits={x}\n", .{ @errorName(e), d_diag.index, d_diag.want, d_diag.got, d_diag.bound, d_diag.ratio, d_diag.want_f32_bits }); + std.debug.print(" DIAG output={s}\n", .{@tagName(d_diag.output)}); + const hv = if (d_diag.output == .state) d_diag.index / (128 * 128) else d_diag.index / 128; + printGateInputs(f.a_bits, f.b_bits, f.a_log_bits, f.dt_bits, hv); + return false; + }; + if (!try unchanged(a, cmds, staging, f.soa_bits, soa, timeout) or + !try unchanged(a, cmds, staging, f.state_in_bits, state_in, timeout) or + !try unchanged(a, cmds, staging, f.a_bits, a_buf, timeout) or + !try unchanged(a, cmds, staging, f.b_bits, b_buf, timeout) or + !try unchanged(a, cmds, staging, f.a_log_bits, a_log, timeout) or + !try unchanged(a, cmds, staging, f.dt_bits, dt, timeout)) + { + std.debug.print("FAIL decode: an input changed\n", .{}); + return false; + } + if (!try plantRefuse2(f, first.?, souts, "y", "decode")) return false; + std.debug.print("MILESTONE decode elapsed_ms={d} worst={e} of bound, repeats byte-identical, inputs unchanged, planted outputs refused\n", .{ clock.ms(), drep.worst }); + return true; +} + +// convDecodeStep: 4-tap conv vs f64 (raw store; split owns SiLU); carry-out bit-exact; split ACTUAL bits vs SplitOne. +pub fn convDecodeGate(a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, m: *vk.Gdn, timeout: u64) !bool { + const C: u64 = fx.Model.conv_channels; + const linear: u64 = fx.Model.linear_dim; + const cf = try fx.Conv.init(a, 0x57ab); + defer cf.deinit(); + const sf = try fx.SplitOne.init(a, 0x57ac); + defer sf.deinit(); + // x is the first conv channel-row of the qkv fixture + const x_bits = cf.qkv_bits[0..@intCast(C)]; + + var state = try uploadShell(dev, cmds, staging, cf.state_bits, timeout); + defer state.deinit(); + var xbuf = try uploadShell(dev, cmds, staging, x_bits, timeout); + defer xbuf.deinit(); + var weight = try uploadShell(dev, cmds, staging, cf.weight_bits, timeout); + defer weight.deinit(); + var a_log = try uploadShell(dev, cmds, staging, sf.a_log_bits, timeout); + defer a_log.deinit(); + var a_buf = try uploadShell(dev, cmds, staging, sf.a_bits, timeout); + defer a_buf.deinit(); + var b_buf = try uploadShell(dev, cmds, staging, sf.b_bits, timeout); + defer b_buf.deinit(); + var dt_buf = try uploadShell(dev, cmds, staging, sf.dt_bits, timeout); + defer dt_buf.deinit(); + var activated = try vk.Buffer.init(dev, .device, @intCast(C * 2)); + defer activated.deinit(); + var state_out = try vk.Buffer.init(dev, .device, @intCast(3 * C * 2)); + defer state_out.deinit(); + var soa = try vk.Buffer.init(dev, .device, @intCast(3 * linear * 2)); + defer soa.deinit(); + var g_out = try vk.Buffer.init(dev, .device, @intCast(fx.Model.heads * 4)); + defer g_out.deinit(); + var beta_out = try vk.Buffer.init(dev, .device, @intCast(fx.Model.heads * 2)); + defer beta_out.deinit(); + + const state_view = try vk.tensor.View.init(&state, .bf16, 0, 3 * C); + const x_view = try vk.tensor.View.init(&xbuf, .bf16, 0, C); + const weight_view = try vk.tensor.View.init(&weight, .bf16, 0, 4 * C); + const activated_view = try vk.tensor.View.init(&activated, .bf16, 0, C); + const state_out_view = try vk.tensor.View.init(&state_out, .bf16, 0, 3 * C); + const soa_view = try vk.tensor.View.init(&soa, .bf16, 0, 3 * linear); + const g_view = try vk.tensor.View.init(&g_out, .f32, 0, fx.Model.heads); + const beta_view = try vk.tensor.View.init(&beta_out, .bf16, 0, fx.Model.heads); + const a_log_view = try vk.tensor.View.init(&a_log, .f32, 0, fx.Model.heads); + const a_view = try vk.tensor.View.init(&a_buf, .bf16, 0, fx.Model.heads); + const b_view = try vk.tensor.View.init(&b_buf, .bf16, 0, fx.Model.heads); + const dt_view = try vk.tensor.View.init(&dt_buf, .bf16, 0, fx.Model.heads); + + const poison = try a.alloc(u16, @intCast(C)); + defer a.free(poison); + @memset(poison, fx.poison16); + const init32 = try a.alloc(u16, @intCast(2 * fx.Model.heads)); + defer a.free(init32); + @memset(init32, 0xffff); + var first: ?[]u16 = null; + defer if (first) |xx| a.free(xx); + const act_bits = try a.alloc(u16, @intCast(C)); + defer a.free(act_bits); + for (0..repeats) |rep| { + try vk.orExit(vk.upload(cmds, activated, 0, staging, std.mem.sliceAsBytes(poison), timeout)); + try vk.orExit(vk.upload(cmds, g_out, 0, staging, std.mem.sliceAsBytes(init32), timeout)); + try vk.orExit(m.convDecodeStep(cmds, soa_view, g_view, beta_view, activated_view, state_view, x_view, weight_view, a_view, b_view, a_log_view, dt_view, state_out_view, timeout)); + try fetch(cmds, staging, activated, std.mem.sliceAsBytes(act_bits), timeout); + if (first) |xx| { + if (std.mem.indexOfDiff(u16, xx, act_bits)) |at| { + std.debug.print("FAIL conv-decode: repeat {d} differs at element {d}\n", .{ rep, at }); + return false; + } + } else first = try a.dupe(u16, act_bits); + } + // raw conv reference: ascending 4 taps on exact dyadic operands + var cr: fx.Report = .{}; + for (0..@intCast(C)) |ch| { + var acc: f64 = 0; + for (0..3) |k| acc += fx.c.fromBf16(cf.state_bits[k * @as(usize, @intCast(C)) + ch]) * fx.c.fromBf16(cf.weight_bits[ch * 4 + k]); + acc += fx.c.fromBf16(x_bits[ch]) * fx.c.fromBf16(cf.weight_bits[ch * 4 + 3]); + const got = fx.c.fromBf16(first.?[ch]); + if (!std.math.isFinite(got)) { + std.debug.print("FAIL conv-decode: nonfinite conv output at {d}\n", .{ch}); + return false; + } + const d = @abs(got - acc); + const bound = fx.budget(4, 0, 0, acc, @abs(acc)); + if (d > bound) { + std.debug.print("FAIL conv-decode: conv output over bound at {d}\n", .{ch}); + return false; + } + const ratio = d / bound; + if (ratio > cr.worst) cr = .{ .worst = ratio, .where = ch }; + } + // negative controls on the conv stage + { + const mutated = try a.dupe(u16, first.?); + defer a.free(mutated); + mutated[0] = 0x4f00; + var refused = false; + for (mutated, 0..) |bits, ch| { + var acc: f64 = 0; + for (0..3) |k| acc += fx.c.fromBf16(cf.state_bits[k * @as(usize, @intCast(C)) + ch]) * fx.c.fromBf16(cf.weight_bits[ch * 4 + k]); + acc += fx.c.fromBf16(x_bits[ch]) * fx.c.fromBf16(cf.weight_bits[ch * 4 + 3]); + if (!std.math.isFinite(fx.c.fromBf16(bits)) or @abs(fx.c.fromBf16(bits) - acc) > fx.budget(4, 0, 0, acc, @abs(acc))) { + refused = true; + break; + } + } + if (!refused) { + std.debug.print("FAIL conv-decode: a planted wrong conv output passed\n", .{}); + return false; + } + } + // carry-out state: concat(state[1..3], x), bit-exact + const souts = try a.alloc(u16, @intCast(3 * C)); + defer a.free(souts); + try fetch(cmds, staging, state_out, std.mem.sliceAsBytes(souts), timeout); + var want_state = try a.alloc(u16, @intCast(3 * C)); + defer a.free(want_state); + @memcpy(want_state[0 .. 2 * @as(usize, @intCast(C))], cf.state_bits[@intCast(C)..][0 .. 2 * @as(usize, @intCast(C))]); + @memcpy(want_state[2 * @as(usize, @intCast(C)) ..], x_bits); + if (std.mem.indexOfDiff(u16, want_state, souts) != null) { + std.debug.print("FAIL conv-decode: the carry-out state is not concat(state[1..3], x)\n", .{}); + return false; + } + // split on the ACTUAL activated bits + @memcpy(sf.conv_out_bits[0..@intCast(C)], first.?); + const soa_bits = try a.alloc(u16, @intCast(3 * linear)); + defer a.free(soa_bits); + const g_pairs = try a.alloc(u16, @intCast(2 * fx.Model.heads)); + defer a.free(g_pairs); + const beta_bits = try a.alloc(u16, @intCast(fx.Model.heads)); + defer a.free(beta_bits); + try fetch(cmds, staging, soa, std.mem.sliceAsBytes(soa_bits), timeout); + try fetch(cmds, staging, g_out, std.mem.sliceAsBytes(g_pairs), timeout); + try fetch(cmds, staging, beta_out, std.mem.sliceAsBytes(beta_bits), timeout); + var s2_diag: fx.Violation = undefined; + const srep = sf.check(soa_bits, g_pairs, beta_bits, &s2_diag) catch |e| { + std.debug.print("FAIL conv-decode: split {s} first at index {d}: want={e} got={e} bound={e} ratio={e} want_f32_bits={x}\n", .{ @errorName(e), s2_diag.index, s2_diag.want, s2_diag.got, s2_diag.bound, s2_diag.ratio, s2_diag.want_f32_bits }); + printGateInputsOne(sf, s2_diag); + return false; + }; + { + const mutated = try a.dupe(u16, soa_bits); + defer a.free(mutated); + mutated[0] = 0x4f00; + if (sf.check(mutated, g_pairs, beta_bits, null)) |_| { + std.debug.print("FAIL conv-decode: a planted wrong split output passed\n", .{}); + return false; + } else |_| {} + mutated[0] = 0x7fc0; + if (sf.check(mutated, g_pairs, beta_bits, null)) |_| { + std.debug.print("FAIL conv-decode: a planted nonfinite split output passed\n", .{}); + return false; + } else |_| {} + } + if (!try unchanged(a, cmds, staging, cf.state_bits, state, timeout) or + !try unchanged(a, cmds, staging, x_bits, xbuf, timeout) or + !try unchanged(a, cmds, staging, cf.weight_bits, weight, timeout) or + !try unchanged(a, cmds, staging, sf.a_bits, a_buf, timeout) or + !try unchanged(a, cmds, staging, sf.b_bits, b_buf, timeout) or + !try unchanged(a, cmds, staging, sf.dt_bits, dt_buf, timeout) or + !try unchanged(a, cmds, staging, sf.a_log_bits, a_log, timeout)) + { + std.debug.print("FAIL conv-decode: an input changed\n", .{}); + return false; + } + std.debug.print("MILESTONE conv-decode elapsed_ms={d} conv worst={e} split worst={e} of bounds, " ++ + "carry-out state bit-exact, split ran on the ACTUAL activated bits, repeats byte-identical, " ++ + "inputs unchanged, planted outputs refused\n", .{ clock.ms(), cr.worst, srep.worst }); + return true; +} + +/// The gated RMSNorm epilogue at rows=4. +pub fn normGate(a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, m: *vk.Gdn, timeout: u64) !bool { + const rows: u64 = 4; + const total: u64 = rows * fx.Model.linear_dim; + const f = try fx.Norm.init(a, 0x57a9); + defer f.deinit(); + var y = try uploadShell(dev, cmds, staging, f.y_bits, timeout); + defer y.deinit(); + var z = try uploadShell(dev, cmds, staging, f.z_bits, timeout); + defer z.deinit(); + var w = try uploadShell(dev, cmds, staging, f.w_bits, timeout); + defer w.deinit(); + var out = try vk.Buffer.init(dev, .device, @intCast(total * 2)); + defer out.deinit(); + const y_view = try vk.tensor.View.init(&y, .bf16, 0, total); + const z_view = try vk.tensor.View.init(&z, .bf16, 0, total); + const w_view = try vk.tensor.View.init(&w, .bf16, 0, fx.Model.head_dim); + const out_view = try vk.tensor.View.init(&out, .bf16, 0, total); + const poison = try a.alloc(u16, @intCast(total)); + defer a.free(poison); + @memset(poison, fx.poison16); + var first: ?[]u16 = null; + defer if (first) |x| a.free(x); + const bits = try a.alloc(u16, @intCast(total)); + defer a.free(bits); + for (0..repeats) |rep| { + try vk.orExit(vk.upload(cmds, out, 0, staging, std.mem.sliceAsBytes(poison), timeout)); + try vk.orExit(m.gatedNorm(cmds, out_view, y_view, z_view, w_view, rows, 1e-6, timeout)); + try fetch(cmds, staging, out, std.mem.sliceAsBytes(bits), timeout); + if (first) |x| { + if (std.mem.indexOfDiff(u16, x, bits)) |at| { + std.debug.print("FAIL norm: repeat {d} differs at element {d}\n", .{ rep, at }); + return false; + } + } else first = try a.dupe(u16, bits); + } + const nrep = f.check(first.?) catch |e| { + std.debug.print("FAIL norm: {s}\n", .{@errorName(e)}); + return false; + }; + if (!try unchanged(a, cmds, staging, f.y_bits, y, timeout) or + !try unchanged(a, cmds, staging, f.z_bits, z, timeout) or + !try unchanged(a, cmds, staging, f.w_bits, w, timeout)) + { + std.debug.print("FAIL norm: an input changed\n", .{}); + return false; + } + if (!try plantRefuse(f, first.?, "norm", "norm")) return false; + std.debug.print("MILESTONE norm elapsed_ms={d} rows={d} worst={e} of bound, repeats byte-identical, inputs unchanged, planted outputs refused\n", .{ clock.ms(), rows, nrep.worst }); + return true; +} diff --git a/zig/tests/vulkan/stage4_probe_shared.zig b/zig/tests/vulkan/stage4_probe_shared.zig new file mode 100644 index 00000000000..5c8b2de318d --- /dev/null +++ b/zig/tests/vulkan/stage4_probe_shared.zig @@ -0,0 +1,134 @@ +// Stage 4 probe: the clock, the buffer helpers, the planted refusals and the input dumps the gates share. + +const std = @import("std"); +const vk = @import("vulkan"); +const fx = @import("stage4_fixture.zig"); + +pub const repeats = 3; +pub const state_elems: u64 = fx.Model.state_elements; + +pub const Elapsed = struct { + io: std.Io, + start: std.Io.Timestamp, + + pub fn ms(e: Elapsed) i64 { + return e.start.durationTo(std.Io.Clock.awake.now(e.io)).toMilliseconds(); + } +}; + +pub fn uploadShell(dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, shell: []const u16, timeout: u64) !vk.Buffer { + var b = try vk.Buffer.init(dev, .device, std.mem.sliceAsBytes(shell).len); + errdefer b.deinit(); + try vk.orExit(vk.upload(cmds, b, 0, staging, std.mem.sliceAsBytes(shell), timeout)); + return b; +} + +pub fn fetch(cmds: *vk.Commands, staging: vk.Buffer, b: vk.Buffer, out: []u8, timeout: u64) !void { + try vk.orExit(vk.download(cmds, b, 0, staging, out, timeout)); +} + +pub fn unchanged(a: std.mem.Allocator, cmds: *vk.Commands, staging: vk.Buffer, shell: []const u16, buf: vk.Buffer, timeout: u64) !bool { + const back = try a.alloc(u16, shell.len); + defer a.free(back); + try fetch(cmds, staging, buf, std.mem.sliceAsBytes(back), timeout); + return std.mem.indexOfDiff(u16, shell, back) == null; +} + +/// Planted refusals on a downloaded GOOD output. +pub fn plantRefuse(f: anytype, good: []const u16, name: []const u8, gate: []const u8) !bool { + const mutated = try f.a.dupe(u16, good); + defer f.a.free(mutated); + var first_finite: ?usize = null; + for (good, 0..) |bits, i| { + if (std.math.isFinite(fx.c.fromBf16(bits))) { + first_finite = i; + break; + } + } + const at = first_finite orelse return error.NoFiniteReference; + mutated[at] = 0x4f00; // 2^31: an absolutely wrong finite the budgets must refuse + if (f.check(mutated)) |_| { + std.debug.print("FAIL {s} case=\"{s}\": a planted wrong finite output passed\n", .{ gate, name }); + return false; + } else |_| {} + for ([_]u16{ 0x7fc0, 0x7f80, 0xff80 }, [_][]const u8{ "nan", "+inf", "-inf" }) |pb, pname| { + mutated[at] = pb; + if (f.check(mutated)) |_| { + std.debug.print("FAIL {s} case=\"{s}\": a planted {s} output passed\n", .{ gate, name, pname }); + return false; + } else |_| {} + } + return true; +} + +// Planted refusals on a downloaded GOOD output; `extra` is passed through to two-slice checkers (state bits beside y). +pub fn plantRefuse2(f: anytype, good: []const u16, extra: []const u16, name: []const u8, gate: []const u8) !bool { + const mutated = try f.a.dupe(u16, good); + defer f.a.free(mutated); + const extra_mut = try f.a.dupe(u16, extra); + defer f.a.free(extra_mut); + var first_finite: ?usize = null; + for (good, 0..) |bits, i| { + if (std.math.isFinite(fx.c.fromBf16(bits))) { + first_finite = i; + break; + } + } + const at = first_finite orelse return error.NoFiniteReference; + mutated[at] = 0x4f00; // 2^31: an absolutely wrong finite the budgets must refuse + if (f.check(mutated, extra_mut, null)) |_| { + std.debug.print("FAIL {s} case=\"{s}\": a planted wrong finite output passed\n", .{ gate, name }); + return false; + } else |_| {} + for ([_]u16{ 0x7fc0, 0x7f80, 0xff80 }) |pb| { + mutated[at] = pb; + if (f.check(mutated, extra_mut, null)) |_| { + std.debug.print("FAIL {s} case=\"{s}\": a planted nonfinite output passed\n", .{ gate, name }); + return false; + } else |_| {} + } + return true; +} + +/// Uploaded host inputs; GPU softplus/exp intermediates are not stored. +pub fn printGateInputs(a_bits: []const u16, b_bits: []const u16, a_log_pairs: []const u16, dt_bits: []const u16, gate_index: usize) void { + const hv = gate_index % 16; + const alog_word = @as(u32, a_log_pairs[2 * hv]) | (@as(u32, a_log_pairs[2 * hv + 1]) << 16); + const alog: f32 = @bitCast(alog_word); + const ref = fx.gateRefOracle(a_bits[gate_index], b_bits[gate_index], alog, dt_bits[hv]); + std.debug.print(" HOST gate token={d} head={d} a=0x{x:0>4} ({e}) b=0x{x:0>4} ({e}) dt=0x{x:0>4} ({e}) A_log=0x{x:0>8} ({e}) oracle_g={e} oracle_beta={e}\n", .{ + gate_index / 16, hv, a_bits[gate_index], fx.c.fromBf16(a_bits[gate_index]), b_bits[gate_index], fx.c.fromBf16(b_bits[gate_index]), dt_bits[hv], fx.c.fromBf16(dt_bits[hv]), alog_word, alog, ref.g, ref.beta, + }); +} + +pub fn printGateInputsOne(f: fx.SplitOne, v: fx.Violation) void { + std.debug.print(" DIAG output={s}\n", .{@tagName(v.output)}); + if (v.output == .g or v.output == .beta) { + printGateInputs(f.a_bits, f.b_bits, f.a_log_bits, f.dt_bits, v.index); + } else { + std.debug.print(" HOST soa index={d} conv=0x{x:0>4} ({e})\n", .{ v.index, f.conv_out_bits[v.index], fx.c.fromBf16(f.conv_out_bits[v.index]) }); + } +} + +pub fn printSplitGateInputs(f: fx.Split, v: fx.Violation) void { + std.debug.print(" DIAG output={s}\n", .{@tagName(v.output)}); + if (v.output == .g or v.output == .beta) { + printGateInputs(f.a_bits, f.b_bits, f.a_log_bits, f.dt_bits, v.index); + } else { + const section: usize = @intCast(fx.tokens * fx.Model.linear_dim); + const kind = v.index / section; + const within = v.index % section; + const token = within / 2048; + const head = (within % 2048) / 128; + const col = within % 128; + const source = token * 6144 + kind * 2048 + head * 128 + col; + std.debug.print(" HOST soa token={d} section={d} head={d} col={d} conv=0x{x:0>4} ({e})\n", .{ token, kind, head, col, f.conv_out_bits[source], fx.c.fromBf16(f.conv_out_bits[source]) }); + } +} + +pub fn printRecurInputs(f: fx.Split, v: fx.Violation, final_token: u64) void { + std.debug.print(" DIAG output={s}\n", .{@tagName(v.output)}); + const head = if (v.output == .y) (v.index / 128) % 16 else v.index / (128 * 128); + const token: usize = if (v.output == .y) v.index / 2048 else @intCast(final_token); + printGateInputs(f.a_bits, f.b_bits, f.a_log_bits, f.dt_bits, token * 16 + head); +} diff --git a/zig/tests/vulkan/stage5_fixture.zig b/zig/tests/vulkan/stage5_fixture.zig new file mode 100644 index 00000000000..1f8bd0ad4e8 --- /dev/null +++ b/zig/tests/vulkan/stage5_fixture.zig @@ -0,0 +1,235 @@ +// Stage 5: reference re-walks widened bits (not production helpers), q.backend.draw rule; align=256B, tail=16x0xffff +const std = @import("std"); +const contract = @import("contract"); +const c = contract.spec; + +/// Poisoned elements after each row's data, so an overrun reads poison. +pub const tail = 16; +pub const poison16: u16 = 0xffff; + +/// The real Qwen3.5 vocab and the argreduce chunk width (greedy.zig block_elements). +pub const vocab: usize = 248320; +pub const chunk: usize = 4096; + +// The offset is 128 bf16 = 256 bytes, Vulkan's largest minStorageBufferOffsetAlignment; stages 2-4 use the same. +pub const at: usize = 128; + +pub const ReferenceError = error{ NonFinite, Invalid }; + +// Reference rejects non-finite values; strict ascending comparison keeps the largest and lowest index (+0/-0 equal). +pub fn referenceBits(data: []const u16) ReferenceError!u32 { + if (data.len == 0) return error.Invalid; + var best = -std.math.inf(f32); + var index: u32 = 0; + for (data, 0..) |b, i| { + const v = c.fromBf16(b); + if (!std.math.isFinite(v)) return error.NonFinite; + if (v > best) { + best = v; + index = @intCast(i); + } + } + return index; +} + +// The first non-finite element of a row — the evidence decode() must report; the caller guarantees one exists. +pub fn firstBadBits(data: []const u16) usize { + for (data, 0..) |b, i| { + if (!std.math.isFinite(c.fromBf16(b))) return i; + } + unreachable; +} + +pub const NonFinite = enum { nan, pinf, ninf }; +pub const ZeroOrder = enum { plus_first, minus_first }; + +pub const Row = struct { + a: std.mem.Allocator, + /// The whole poisoned shell; the row's data is bits[at..][0..count]. + bits: []u16, + count: usize, + + pub fn deinit(r: Row) void { + r.a.free(r.bits); + } + + pub fn data(r: Row) []u16 { + return r.bits[at..][0..r.count]; + } + + pub fn reference(r: Row) ReferenceError!u32 { + return referenceBits(r.data()); + } + + pub fn firstBad(r: Row) usize { + return firstBadBits(r.data()); + } + + // The downloaded shell equals the uploaded shell bit for bit; guards keep poison and logits stay unchanged. + pub fn checkShell(r: Row, got: []const u16) !void { + if (got.len != r.bits.len) return error.LengthMismatch; + if (std.mem.indexOfDiff(u16, r.bits, got) != null) return error.Changed; + } +}; + +fn shell(a: std.mem.Allocator, count: usize) ![]u16 { + const buf = try a.alloc(u16, at + count + tail); + @memset(buf, poison16); + return buf; +} + +// Random finite bf16 magnitudes use [0.25, cap), both signs unless negative-only; rounding never reaches or passes cap. +fn fillRandom(data: []u16, rng: std.Random, cap: f32, neg_only: bool) void { + for (data) |*d| { + const mag = rng.float(f32) * (cap - 0.25) + 0.25; + const v = if (neg_only) -mag else if (rng.boolean()) mag else -mag; + d.* = c.toBf16(v); + } +} + +// A realistic full-vocab row: random finite logits with a unique max of exactly 5.0 at count/2. +pub fn random(a: std.mem.Allocator, count: usize, seed: u64) !Row { + const bits = try shell(a, count); + errdefer a.free(bits); + var prng = std.Random.DefaultPrng.init(seed); + fillRandom(bits[at..][0..count], prng.random(), 4.75, false); + bits[at + count / 2] = c.toBf16(5.0); + return .{ .a = a, .bits = bits, .count = count }; +} + +// Ties at 1, chunk starts, final partial start, and row end test whether lowest index wins across chunk/combine. +pub fn ties(a: std.mem.Allocator, count: usize, seed: u64) !Row { + const bits = try shell(a, count); + errdefer a.free(bits); + var prng = std.Random.DefaultPrng.init(seed); + const data = bits[at..][0..count]; + fillRandom(data, prng.random(), 3.0, false); + const tie_value = c.toBf16(3.5); + const tie_indices = [_]usize{ 1, chunk, 2 * chunk, (count - 1) / chunk * chunk, count - 1 }; + for (tie_indices) |i| data[i] = tie_value; + return .{ .a = a, .bits = bits, .count = count }; +} + +/// All-negative logits with a unique max of exactly -0.125 at count/2. +pub fn negative(a: std.mem.Allocator, count: usize, seed: u64) !Row { + const bits = try shell(a, count); + errdefer a.free(bits); + var prng = std.Random.DefaultPrng.init(seed); + fillRandom(bits[at..][0..count], prng.random(), 4.0, true); + bits[at + count / 2] = c.toBf16(-0.125); + return .{ .a = a, .bits = bits, .count = count }; +} + +// With -1.0 background, +0 at element 5 and -0 at count-7 win in either order because zeros compare equal. +pub fn zero(a: std.mem.Allocator, count: usize, seed: u64, order: ZeroOrder) !Row { + const bits = try shell(a, count); + errdefer a.free(bits); + _ = seed; + @memset(bits[at..][0..count], c.toBf16(-1.0)); + bits[at + 5] = switch (order) { + .plus_first => 0x0000, + .minus_first => 0x8000, + }; + bits[at + count - 7] = switch (order) { + .plus_first => 0x8000, + .minus_first => 0x0000, + }; + return .{ .a = a, .bits = bits, .count = count }; +} + +// A finite row has a unique max at count/2; a non-finite `bad` (nan/+inf/-inf) at that spot must still be refused. +pub fn nonFinite(a: std.mem.Allocator, count: usize, seed: u64, kind: NonFinite, bad: usize) !Row { + const bits = try shell(a, count); + errdefer a.free(bits); + var prng = std.Random.DefaultPrng.init(seed); + const data = bits[at..][0..count]; + fillRandom(data, prng.random(), 4.75, false); + data[count / 2] = c.toBf16(5.0); + data[bad] = switch (kind) { + .nan => 0x7fc0, + .pinf => 0x7f80, + .ninf => 0xff80, + }; + return .{ .a = a, .bits = bits, .count = count }; +} + +/// The smallest legal row: two elements, the second one largest. +pub fn minTwo(a: std.mem.Allocator) !Row { + const bits = try shell(a, 2); + errdefer a.free(bits); + bits[at] = c.toBf16(-2.0); + bits[at + 1] = c.toBf16(3.5); + return .{ .a = a, .bits = bits, .count = 2 }; +} + +const testing = std.testing; + +fn guardsPoisoned(r: Row) !void { + for (r.bits, 0..) |g, i| { + if (i >= at and i < at + r.count) continue; + try testing.expectEqual(poison16, g); + } +} + +test "the reference applies the draw rule on raw bf16 bits" { + const neg = [_]u16{ c.toBf16(-1.5), c.toBf16(-1.5), c.toBf16(-0.5), c.toBf16(-2.0) }; + try testing.expectEqual(@as(u32, 2), try referenceBits(&neg)); + const tie = [_]u16{ c.toBf16(-1.0), c.toBf16(3.5), c.toBf16(3.5), c.toBf16(3.5) }; + try testing.expectEqual(@as(u32, 1), try referenceBits(&tie)); + const zeros_a = [_]u16{ 0x0000, 0x8000, c.toBf16(-1.0) }; + try testing.expectEqual(@as(u32, 0), try referenceBits(&zeros_a)); + const zeros_b = [_]u16{ 0x8000, 0x0000, c.toBf16(-1.0) }; + try testing.expectEqual(@as(u32, 0), try referenceBits(&zeros_b)); + try testing.expectError(error.Invalid, referenceBits(&.{})); + try testing.expectError(error.NonFinite, referenceBits(&.{ c.toBf16(1.0), 0x7fc0 })); + try testing.expectError(error.NonFinite, referenceBits(&.{ c.toBf16(1.0), 0x7f80 })); + try testing.expectError(error.NonFinite, referenceBits(&.{ 0xff80, c.toBf16(-1.0) })); +} + +test "every builder's row agrees with the reference and keeps its guards poisoned" { + const a = testing.allocator; + { + var r = try random(a, 512, 7); + defer r.deinit(); + _ = try r.reference(); + try guardsPoisoned(r); + } + { + var r = try ties(a, vocab, 8); + defer r.deinit(); + try testing.expectEqual(@as(u32, 1), try r.reference()); + try guardsPoisoned(r); + } + { + var r = try negative(a, 4096, 9); + defer r.deinit(); + _ = try r.reference(); + try guardsPoisoned(r); + } +} + +test "zero rows pick element 5 in both orders; non-finite rows refuse and name the offender" { + const a = testing.allocator; + { + var r = try zero(a, 64, 10, .plus_first); + defer r.deinit(); + try testing.expectEqual(@as(u32, 5), try r.reference()); + } + { + var r = try zero(a, 64, 11, .minus_first); + defer r.deinit(); + try testing.expectEqual(@as(u32, 5), try r.reference()); + } + { + var r = try nonFinite(a, 4096, 12, .nan, 100); + defer r.deinit(); + try testing.expectError(error.NonFinite, r.reference()); + try testing.expectEqual(@as(usize, 100), r.firstBad()); + } + { + var m = try minTwo(a); + defer m.deinit(); + try testing.expectEqual(@as(u32, 1), try m.reference()); + try guardsPoisoned(m); + } +} diff --git a/zig/tests/vulkan/stage5_probe.zig b/zig/tests/vulkan/stage5_probe.zig new file mode 100644 index 00000000000..271ebb995b2 --- /dev/null +++ b/zig/tests/vulkan/stage5_probe.zig @@ -0,0 +1,319 @@ +// Greedy.select vs CPU max (248320 BF16; low-index ties, ±0 equal); NonFinite first-bad index; exit codes 0/1/2/124. +const std = @import("std"); +const vk = @import("vulkan"); +const shaders = @import("stage5_shaders"); +const fx = @import("stage5_fixture.zig"); + +const usage = "usage: tf-stage5-gpu [--vendor ID] [--timeout-ms MS] | --help\n" ++ + " --help show this text and exit without opening a loader or device\n" ++ + " --vendor ID refuse a device whose Vulkan vendorID differs (decimal or 0x hex; no default)\n" ++ + " --timeout-ms MS bound on every fence wait, 1 to 600000 (default 10000); a timeout ends the process\n"; + +/// Milliseconds on the awake clock since main started, printed on every milestone. +const Elapsed = struct { + io: std.Io, + start: std.Io.Timestamp, + + fn ms(e: Elapsed) i64 { + return e.start.durationTo(std.Io.Clock.awake.now(e.io)).toMilliseconds(); + } +}; + +const repeats = 3; + +pub fn main(init: std.process.Init) !u8 { + const clock: Elapsed = .{ .io = init.io, .start = std.Io.Clock.awake.now(init.io) }; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + for (args[1..]) |arg| if (std.mem.eql(u8, arg, "--help")) { + std.debug.print(usage, .{}); + return 0; + }; + var vendor: ?u32 = null; + var timeout_ms: u64 = 10_000; + var i: usize = 1; + while (i < args.len) : (i += 2) { + const value = if (i + 1 < args.len) args[i + 1] else ""; + if (std.mem.eql(u8, args[i], "--vendor")) { + vendor = std.fmt.parseInt(u32, value, 0) catch null; + if (vendor == null) break; + } else if (std.mem.eql(u8, args[i], "--timeout-ms")) { + timeout_ms = std.fmt.parseInt(u64, value, 10) catch 0; + if (timeout_ms == 0 or timeout_ms > 600_000) break; + } else break; + } else return if (try probe(init.gpa, clock, vendor, timeout_ms * std.time.ns_per_ms)) 0 else 1; + std.debug.print(usage, .{}); + return 2; +} + +fn shaderSet() vk.greedy.Shaders { + return .{ + .finite_logits = &shaders.finite_logits, + .argreduce_chunk = &shaders.argreduce_chunk, + .argreduce_combine = &shaders.argreduce_combine, + }; +} + +fn probe(a: std.mem.Allocator, clock: Elapsed, vendor: ?u32, timeout: u64) !bool { + var rt = try vk.Runtime.open(); + defer rt.close(); + var dev = try vk.Device.open(&rt, .{ .features = .{ .storage16 = true, .int16 = true } }); + defer dev.close(); + const pr = dev.properties; + const name = std.mem.sliceTo(&pr.deviceName, 0); + const v = pr.apiVersion; + std.debug.print("MILESTONE device-found elapsed_ms={d} name=\"{s}\" vendor=0x{x:0>4} device=0x{x:0>4} " ++ + "api={d}.{d}.{d} storage16={} int16={} subgroups={d} shuffle-relative={} subgroup-size={d} fence-timeout={d}ms\n", .{ + clock.ms(), name, + pr.vendorID, pr.deviceID, + v >> 22 & 0x7f, v >> 12 & 0x3ff, + v & 0xfff, dev.features.storage16, + dev.features.int16, dev.v11.subgroupSupportedStages, + dev.v11.subgroupSupportedOperations & vk.c.VK_SUBGROUP_FEATURE_SHUFFLE_RELATIVE_BIT != 0, dev.v11.subgroupSize, + timeout / std.time.ns_per_ms, + }); + if (pr.deviceType == vk.c.VK_PHYSICAL_DEVICE_TYPE_CPU) { + std.debug.print("FAIL device: a CPU (software) Vulkan device is never a GPU result\n", .{}); + return false; + } + if (vendor) |want| if (pr.vendorID != want) { + std.debug.print("FAIL device: vendor 0x{x:0>4} is not the requested 0x{x:0>4}\n", .{ pr.vendorID, want }); + return false; + }; + + var cmds = try vk.Commands.init(&dev); + defer cmds.deinit(); + // A full 248320-element bf16 row rides staging: 485 KiB. + var staging = try vk.Buffer.init(&dev, .staging, 4 << 20); + defer staging.deinit(); + + var passed: usize = 0; + var total: usize = 0; + total += 1; + passed += @intFromBool(try argmaxGate(a, clock, &dev, &cmds, staging, .vocab, timeout)); + total += 1; + passed += @intFromBool(try argmaxGate(a, clock, &dev, &cmds, staging, .ties, timeout)); + total += 1; + passed += @intFromBool(try argmaxGate(a, clock, &dev, &cmds, staging, .negative, timeout)); + total += 1; + passed += @intFromBool(try zeroGate(a, clock, &dev, &cmds, staging, timeout)); + total += 1; + passed += @intFromBool(try nonFiniteGate(a, clock, &dev, &cmds, staging, timeout)); + total += 1; + passed += @intFromBool(try minCountGate(a, clock, &dev, &cmds, staging, timeout)); + total += 1; + passed += @intFromBool(try oddCountGate(a, clock, &dev, &cmds, staging, timeout)); + std.debug.print("RESULT stage5 elapsed_ms={d} on \"{s}\": passed={d} total={d}\n", .{ clock.ms(), name, passed, total }); + return passed == total; +} + +fn gpuFail(gate: []const u8, e: anyerror) bool { + std.debug.print("FAIL greedy gate=\"{s}\": {s}\n", .{ gate, @errorName(e) }); + return false; +} + +fn refFail(gate: []const u8, e: anyerror) bool { + std.debug.print("FAIL greedy gate=\"{s}\": the reference refused a finite row: {s}\n", .{ gate, @errorName(e) }); + return false; +} + +// Upload the full poisoned row; caller owns the buffer and creates its address-stable view. +fn uploadRow(gate: []const u8, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, row: fx.Row, timeout: u64) ?vk.Buffer { + var b = vk.Buffer.init(dev, .device, std.mem.sliceAsBytes(row.bits).len) catch |e| { + _ = gpuFail(gate, e); + return null; + }; + vk.orExit(vk.upload(cmds, b, 0, staging, std.mem.sliceAsBytes(row.bits), timeout)) catch |e| { + b.deinit(); + _ = gpuFail(gate, e); + return null; + }; + return b; +} + +// Run 3 selects via orExit; tokens must match, mismatches print FAIL, and GPU errors propagate raw. +fn selectRepeat(gate: []const u8, g: *vk.Greedy, cmds: *vk.Commands, staging: vk.Buffer, view: vk.tensor.View, timeout: u64) !u32 { + var first: ?u32 = null; + for (0..repeats) |rep| { + const tok = try vk.orExit(g.select(cmds, staging, view, timeout)); + if (first) |f| { + if (tok != f) { + std.debug.print("FAIL greedy gate=\"{s}\": repeat {d} selected {d}, first repeat selected {d}\n", .{ gate, rep, tok, f }); + return error.RepeatMismatch; + } + } else first = tok; + } + return first.?; +} + +fn repeatFail(gate: []const u8, e: anyerror) bool { + if (e != error.RepeatMismatch) return gpuFail(gate, e); + return false; +} + +// Download the full shell and require bit identity: poisoned guards and logits remain unchanged. +fn shellChecked(gate: []const u8, a: std.mem.Allocator, cmds: *vk.Commands, staging: vk.Buffer, row: fx.Row, buf: vk.Buffer, timeout: u64) bool { + const back = a.alloc(u16, row.bits.len) catch |e| return gpuFail(gate, e); + defer a.free(back); + vk.orExit(vk.download(cmds, buf, 0, staging, std.mem.sliceAsBytes(back), timeout)) catch |e| return gpuFail(gate, e); + row.checkShell(back) catch |e| { + std.debug.print("FAIL greedy gate=\"{s}\": the logits shell is not intact ({s})\n", .{ gate, @errorName(e) }); + return false; + }; + return true; +} + +fn greedyFor(gate: []const u8, dev: *const vk.Device, count: u64) ?vk.Greedy { + return vk.Greedy.init(dev, shaderSet(), count) catch |e| { + _ = gpuFail(gate, e); + return null; + }; +} + +const RowKind = enum { vocab, ties, negative }; + +/// One finite row: repeat-stable token equal to the reference, shell intact. +fn singleRowGate(a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, gate: []const u8, row: fx.Row, timeout: u64) !bool { + var buf = uploadRow(gate, dev, cmds, staging, row, timeout) orelse return false; + defer buf.deinit(); + const view = vk.tensor.View.init(&buf, .bf16, fx.at * 2, row.count) catch |e| return gpuFail(gate, e); + var g = greedyFor(gate, dev, row.count) orelse return false; + defer g.deinit(); + const tok = selectRepeat(gate, &g, cmds, staging, view, timeout) catch |e| return repeatFail(gate, e); + const want = row.reference() catch |e| return refFail(gate, e); + if (tok != want) { + std.debug.print("FAIL greedy gate=\"{s}\": selected {d}, the reference says {d}\n", .{ gate, tok, want }); + return false; + } + if (!shellChecked(gate, a, cmds, staging, row, buf, timeout)) return false; + std.debug.print("MILESTONE greedy elapsed_ms={d} gate=\"{s}\" count={d} token={d} reference={d} repeats={d} identical inputs-guards-unchanged\n", .{ clock.ms(), gate, row.count, tok, want, repeats }); + return true; +} + +fn argmaxGate(a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, kind: RowKind, timeout: u64) !bool { + const gate = switch (kind) { + .vocab => "vocab-248320", + .ties => "ties-across-chunks", + .negative => "all-negative", + }; + var row = switch (kind) { + .vocab => fx.random(a, fx.vocab, 0x57a5) catch |e| return gpuFail(gate, e), + .ties => fx.ties(a, fx.vocab, 0x57a6) catch |e| return gpuFail(gate, e), + .negative => fx.negative(a, fx.vocab, 0x57a7) catch |e| return gpuFail(gate, e), + }; + defer row.deinit(); + return singleRowGate(a, clock, dev, cmds, staging, gate, row, timeout); +} + +/// +0/-0 ties in both orders: element 5 wins either way, against one Greedy. +fn zeroGate(a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, timeout: u64) !bool { + const gate = "signed-zero-both-orders"; + var rows: [2]fx.Row = .{ + fx.zero(a, fx.vocab, 0x57a8, .plus_first) catch |e| return gpuFail(gate, e), + fx.zero(a, fx.vocab, 0x57a8 + 1, .minus_first) catch |e| return gpuFail(gate, e), + }; + defer for (&rows) |*r| r.deinit(); + var g = greedyFor(gate, dev, fx.vocab) orelse return false; + defer g.deinit(); + var tokens: [2]u32 = undefined; + for (&rows, 0..) |*row, i| { + var buf = uploadRow(gate, dev, cmds, staging, row.*, timeout) orelse return false; + defer buf.deinit(); + const view = vk.tensor.View.init(&buf, .bf16, fx.at * 2, row.count) catch |e| return gpuFail(gate, e); + tokens[i] = selectRepeat(gate, &g, cmds, staging, view, timeout) catch |e| return repeatFail(gate, e); + const want = row.reference() catch |e| return refFail(gate, e); + if (tokens[i] != want) { + std.debug.print("FAIL greedy gate=\"{s}\": order {d} selected {d}, the reference says {d}\n", .{ gate, i, tokens[i], want }); + return false; + } + if (!shellChecked(gate, a, cmds, staging, row.*, buf, timeout)) return false; + } + std.debug.print("MILESTONE greedy elapsed_ms={d} gate=\"{s}\" count={d} plus-first-token={d} minus-first-token={d} reference=lowest-index repeats={d} identical inputs-guards-unchanged\n", .{ clock.ms(), gate, fx.vocab, tokens[0], tokens[1], repeats }); + return true; +} + +// Non-winning -inf/+inf/NaN (even at winner) must refuse; decode() reports first-bad index, not GPU failure. +fn nonFiniteGate(a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, timeout: u64) !bool { + const gate = "nonfinite-refused"; + const cases = [_]struct { kind: fx.NonFinite, bad: usize, label: []const u8 }{ + .{ .kind = .nan, .bad = fx.vocab / 3, .label = "nan" }, + .{ .kind = .pinf, .bad = fx.vocab / 2, .label = "+inf" }, + .{ .kind = .ninf, .bad = 2 * fx.vocab / 3, .label = "-inf" }, + }; + var g = greedyFor(gate, dev, fx.vocab) orelse return false; + defer g.deinit(); + var bad: [cases.len]usize = undefined; + for (cases, 0..) |cs, i| { + var row = fx.nonFinite(a, fx.vocab, 0x57a9 + i, cs.kind, cs.bad) catch |e| return gpuFail(gate, e); + defer row.deinit(); + if (row.reference()) |_| { + std.debug.print("FAIL greedy gate=\"{s}\" kind=\"{s}\": the fixture reference accepted a non-finite row\n", .{ gate, cs.label }); + return false; + } else |e| { + if (e != error.NonFinite) return refFail(gate, e); + } + if (row.firstBad() != cs.bad) { + std.debug.print("FAIL greedy gate=\"{s}\" kind=\"{s}\": the fixture names first bad {d}, planted {d}\n", .{ gate, cs.label, row.firstBad(), cs.bad }); + return false; + } + var buf = uploadRow(gate, dev, cmds, staging, row, timeout) orelse return false; + defer buf.deinit(); + const view = vk.tensor.View.init(&buf, .bf16, fx.at * 2, row.count) catch |e| return gpuFail(gate, e); + if (vk.orExit(g.select(cmds, staging, view, timeout))) |tok| { + std.debug.print("FAIL greedy gate=\"{s}\" kind=\"{s}\": a non-finite row selected token {d}\n", .{ gate, cs.label, tok }); + return false; + } else |e| { + if (e != error.NonFinite) { + std.debug.print("FAIL greedy gate=\"{s}\" kind=\"{s}\": expected NonFinite, got {s}\n", .{ gate, cs.label, @errorName(e) }); + return false; + } + } + // Refusal is not GPU failure: read production Result and require the planted first-offender index. + var ev: vk.greedy.Result = undefined; + vk.orExit(vk.download(cmds, g.results, 0, staging, std.mem.asBytes(&ev), timeout)) catch |e2| { + std.debug.print("FAIL greedy gate=\"{s}\" kind=\"{s}\": evidence download: {s}\n", .{ gate, cs.label, @errorName(e2) }); + return false; + }; + if (ev.finite_flag == 0) { + std.debug.print("FAIL greedy gate=\"{s}\" kind=\"{s}\": the evidence claims a finite verdict\n", .{ gate, cs.label }); + return false; + } + if (ev.first_bad != cs.bad) { + std.debug.print("FAIL greedy gate=\"{s}\" kind=\"{s}\": first-bad evidence {d} is not the planted {d}\n", .{ gate, cs.label, ev.first_bad, cs.bad }); + return false; + } + if (vk.greedy.decode(@intCast(row.count), ev)) |_| { + std.debug.print("FAIL greedy gate=\"{s}\" kind=\"{s}\": the evidence decodes to a token, not a refusal\n", .{ gate, cs.label }); + return false; + } else |e2| { + if (e2 != error.NonFinite) return gpuFail(gate, e2); + } + if (!shellChecked(gate, a, cmds, staging, row, buf, timeout)) return false; + bad[i] = ev.first_bad; + } + std.debug.print("MILESTONE greedy elapsed_ms={d} gate=\"{s}\" count={d} nan-first-bad={d} pinf-first-bad={d} ninf-first-bad={d} verdict=NonFinite decode-refuses inputs-guards-unchanged\n", .{ clock.ms(), gate, fx.vocab, bad[0], bad[1], bad[2] }); + return true; +} + +fn minCountGate(a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, timeout: u64) !bool { + const gate = "min-count-2"; + var row = fx.minTwo(a) catch |e| return gpuFail(gate, e); + defer row.deinit(); + return singleRowGate(a, clock, dev, cmds, staging, gate, row, timeout); +} + +// Odd bf16 count uploads but View refuses with Misaligned before dispatch. +fn oddCountGate(a: std.mem.Allocator, clock: Elapsed, dev: *const vk.Device, cmds: *vk.Commands, staging: vk.Buffer, timeout: u64) !bool { + const gate = "odd-count-refused"; + var row = fx.random(a, 7, 0x57ad) catch |e| return gpuFail(gate, e); + defer row.deinit(); + var buf = uploadRow(gate, dev, cmds, staging, row, timeout) orelse return false; + defer buf.deinit(); + if (vk.tensor.View.init(&buf, .bf16, fx.at * 2, row.count)) |_| { + std.debug.print("FAIL greedy gate=\"{s}\": an odd bf16 count produced a usable view\n", .{gate}); + return false; + } else |e| { + if (e != error.Misaligned) return gpuFail(gate, e); + } + std.debug.print("MILESTONE greedy elapsed_ms={d} gate=\"{s}\" count={d} refused=Misaligned before-any-dispatch\n", .{ clock.ms(), gate, row.count }); + return true; +}