diff --git a/README.md b/README.md index 8925e38b8c..59c0d83c56 100644 --- a/README.md +++ b/README.md @@ -173,7 +173,7 @@ Conversions are explicitly marked: their upstream weights need the linked conver | GLM-5.3-Flash, FP4-Local | `tensorfold pull nvidia/GLM-5.3-Flash-FP4-Local@76348ddb2be2e35e39ecfc00e75dc07703f4f3d5` | Two GB10s at TP=2, or three at TP=3 over a ring of direct links: @Bizuayeu loaded it as stored and served it at both sizes. 165.3 GiB of weights, 13% less than NVFP4's 190.4 GiB. | [Native CUDA](docs/recipes/glm-5.3-flash-cuda.md), in 1.1.5 (#1050). NVIDIA's mixed-precision checkpoint: NVFP4 routed experts with compressed block scales, MXFP8 and 128×128-block FP8 for most of the rest. As close to BF16 as NVFP4 at its smaller size: at TP=2 its NLL is 1.11% above BF16's against NVFP4's 1.37% ([contributed receipt](docs/receipts/glm-fp4local-80517fc3/README.md) on #1050's head, not rerun on main), a difference the per-document intervals in [XSFP4's card](https://huggingface.co/Bizuayeu/GLM-5.3-Flash-XSFP4#measured) do not resolve. | | GLM-5.3-Flash, XSFP4 | `tensorfold pull Bizuayeu/GLM-5.3-Flash-XSFP4@88736f916f8939607fde8a4b881dcae749f7551d` | Two GB10s at TP=2, or three at TP=3 over a ring of direct links: on main `8086c205c` @Bizuayeu loaded it in 370.7 s at TP=2 (rank 0's model 96.97 GiB at a 524,288-token window) and in 313.8 s at TP=3 (63.60 GiB at 1,048,576), and both served. 165.7 GiB of weights, 0.4 GiB more than FP4-Local's. | [Native CUDA](docs/recipes/glm-5.3-flash-cuda.md), in 1.1.5 through FP4-Local's path. FP4-Local with its routed experts' NVFP4 blocks fitted again from the BF16 weights, a smoothing folded into them, so it runs on the same kernels. Closer to BF16 than FP4-Local and NVFP4 at TP=2 by all three measures in [its card](https://huggingface.co/Bizuayeu/GLM-5.3-Flash-XSFP4#measured), on #1050's code: top-1 agreement 93.26%, NLL 0.85% above and a KL bound of 0.0431, against FP4-Local's 92.82%, 1.11% and 0.0510. | | GLM-5.3, EXLR8 `home` | `tensorfold pull kindlingai/glm-5.3-exlr8-k2-k3-k4@f459f5d0272a529403dd9b3fd32f771f4250a0d4 --exlr8 home --node 0/2 --dir ~/models/glm-5.3-exlr8-home` | Two M5 Ultra Macs, 256 GB each; about 160 GB resident per Mac at 8k context / four streams. | [Two-Mac recipe](docs/glm-5.3-exlr8.md). Use `--node 1/2` on the other Mac; Thunderbolt 5 and MCDMA required. No CUDA backend. | -| LFM2.5-24B-A2B, MLX 4 / 6 / 8-bit | `tensorfold pull LiquidAI/LFM2.5-24B-A2B-MLX-4bit@667408d494bde191c2cb9dd4f4549827816b9275`
`tensorfold pull LiquidAI/LFM2.5-24B-A2B-MLX-6bit@1deec53eb49571e21cdcbf82fd5839248f07ed88`
`tensorfold pull LiquidAI/LFM2.5-24B-A2B-MLX-8bit@67ab417dd02d0c9545802a726df668620eae049d` | One Mac; M4 Pro checked. | [Native recipe](docs/recipes/lfm2.5-24b-a2b.md), [release history](CHANGELOG.md). | +| LFM2.5-24B-A2B, MLX 4 / 6 / 8-bit | `tensorfold pull LiquidAI/LFM2.5-24B-A2B-MLX-4bit@667408d494bde191c2cb9dd4f4549827816b9275`
`tensorfold pull LiquidAI/LFM2.5-24B-A2B-MLX-6bit@1deec53eb49571e21cdcbf82fd5839248f07ed88`
`tensorfold pull LiquidAI/LFM2.5-24B-A2B-MLX-8bit@67ab417dd02d0c9545802a726df668620eae049d` | One Mac (M4 Pro checked), or one NVIDIA GPU with the 4-bit export (GB10 checked). | [Native recipe](docs/recipes/lfm2.5-24b-a2b.md) ([CUDA](docs/recipes/lfm2.5-24b-a2b.md#cuda)), [release history](CHANGELOG.md). Weights: LFM Open License v1.0. | | Qwen3.6-35B-A3B, MLX with MTP | `tensorfold pull TensorFold/Qwen3.6-35B-A3B-MLX-4bit-MTP@81169a9bc511a27c1b4eedb77a2cd98ced431847` | One GB10; 20.9 GB checkpoint. | [Qwen3.6](docs/recipes/qwen3.6-moe.md). Native CUDA; the recipe's Metal path belongs to Python. | | Gemma 4 31B, MLX 4-bit | `tensorfold pull mlx-community/gemma-4-31B-it-4bit@696d436c404745a59f30e4939a658162b0a9e57f`
`tensorfold pull z-lab/gemma-4-31B-it-DFlash@eabd648301ce28583cc14757912e5e0f84e152e1` | One M5 Mac; M5 Ultra checked. M1-M4 refused. | [Native 31B](docs/recipes/gemma-4-native.md). Name the drafter with `--drafter`; images with `--vision`. | | Gemma 4 12B, converted MLX 4-bit | `hf download google/gemma-4-12B-it --revision 707f0a3b8a3c7ad586ed01e27eafbad8a27dd0f7 --local-dir ~/models/gemma-4-12b-hf`
`hf download google/gemma-4-12B-it-assistant --revision 46d4c6f13f0ac0ad827b915669b8df9b81c64c51 --local-dir ~/models/gemma-4-12b-assistant` | One M5 Mac; converted text weights 6.3 GB, image weights 100 MB, plus assistant/caches. M1-M4 refused. | [Convert and serve](docs/recipes/gemma-4-12b-native.md). Upstream download needs `hf`; conversion needs Python, mlx and mlx-lm. Serving is native. | diff --git a/build.zig b/build.zig index 7f0a9d2403..952913a822 100644 --- a/build.zig +++ b/build.zig @@ -97,6 +97,7 @@ pub fn build(b: *std.Build) void { @import("zig/build/qwen35_cache_host.zig").compileHost(b, step5_mods.metal, step5_mods.engine, step5_mods.lanes); } @import("zig/build/h3_host.zig").tests(b, target, test_step); + @import("zig/build/lfm3_host.zig").steps(b, test_step); @import("zig/build/digest.zig").steps(b, test_step, draft_ids, build_options); @import("zig/build/flashnext_budget.zig").steps(b, test_step); @import("zig/build/flashnext_batch_host.zig").tests(b, test_step); diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 1ed43ac146..972079f809 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -190,6 +190,13 @@ GLM's Metal window projections register narrow and wide rows at 2, 3, 4, 5, 6 an Independent FP64 linear cases accompany byte replacements against frozen pre-window-optimization kernels. The gathered SwiGLU keeps its predecessor's clipped BF16 output and nonfinite rules; gate/up taps qualify its dots separately. Controls dispatch real source mutations that couple rows or select the wrong expert; host tests validate their witnesses. +CUDA also checks LFM2.5-24B-A2B's kernels: the router's fp32 logits (`tf_lfm2m_route`) and its top-4 rule, Liquid's +3-tap short conv with its per-row state snapshots and the state commit (`liquid_conv.cuh`), each against the host's +replay of its fp32 and bf16 steps (three controls: ties to the higher expert, bx unrounded, every stream committing into +the first stream's slot). Their paged attention at head dim 64 runs over shuffled 512-key pages with several streams +a launch: the decode rows' chunks and merge (`tf_paged_attn_chunk_segs_64`, `tf_paged_attn_merge_segs_64`) and the +prompt rows (`pattn_kernel<64, 8, 4, 8>`), head by head within `attention_ref`'s float64 bound (two controls: the +decode reading its values from the neighbouring lane's dims, the prompt from the neighbouring 8 dims). CUDA also checks the affine-3 experts' split decode pair, bit-equal to the pair kernel it replaces (two controls, one dropping a last round of K groups short of the warps, one chaining the groups in reverse). CUDA also checks Flash Next's NVFP4 experts on prompt tiles with the decode form's sums (SwiGLU gate-up and fp32 diff --git a/docs/recipes/lfm2.5-24b-a2b.md b/docs/recipes/lfm2.5-24b-a2b.md index 7b8b2cdf4b..2e770c0864 100644 --- a/docs/recipes/lfm2.5-24b-a2b.md +++ b/docs/recipes/lfm2.5-24b-a2b.md @@ -1,4 +1,4 @@ -# LFM2.5-24B-A2B, native Metal +# LFM2.5-24B-A2B, native Metal and CUDA This recipe covers [`LiquidAI/LFM2.5-24B-A2B-MLX-4bit`](https://huggingface.co/LiquidAI/LFM2.5-24B-A2B-MLX-4bit) at revision `667408d494bde191c2cb9dd4f4549827816b9275`. The numbers below are measured on this 4-bit export only. @@ -41,3 +41,44 @@ the shared-round kernels on it. The M4 Pro is not yet measured on a quiet machin Every reply was token for token equal before and after, and equal to its solo run and to `"draft": false`. Cold prefill at 2k, 8k and 32k tokens and peak memory did not change. + +## CUDA + +`lfm2_moe` also serves the 4-bit export above (same revision, `667408d494bde191c2cb9dd4f4549827816b9275`) on one +NVIDIA GPU, through the lane host. Checked on a GB10 (sm_121). + +```bash +zig build native -Doptimize=ReleaseFast -Dnvcc=/usr/local/cuda/bin/nvcc +zig-out/native/bin/tensorfold-native serve "$HOME/models/lfm2.5-24b-a2b-4bit" --port 8080 --parallel 8 +``` + +The default window is 32,768 tokens (`--context`), and `--parallel` holds up to 128 streams if the memory budget fits +them. Streams lease 512-position pages of keys and values from one pool as they need them +(`TENSORFOLD_LFM2_KV_PAGES=0` gives each stream fixed pages instead). The prompt cache keeps a chat's state at its +history marks, the conv states and the K/V rows before the mark, in pages allocated once inside its budget +(`--prompt-cache-gib`). The family reads MLX affine 4-bit weights in groups of +64, with the embedding at 4, 6 or 8 bits. The 6- and 8-bit exports are refused by name and stay on Metal. Drafts +are the lane core's copy drafts from the prompt's own text; there is no draft head yet. + +Most launches are shared CUDA operators: the affine linears and experts, the norms, `qk_rows`, the MoE combine, the +pick and the draws. The family adds three kernels: the router (fp32 weights, sigmoid plus bias, top-4) in +`lfm2_moe.cu`, and the 3-tap short conv and its state commit in `liquid_conv.cuh`. It also adds two instances of shared +templates at head dim 64: `pattn_kernel<64, 8, 4, 8>` for prompt rows, and `tf_paged_attn_chunk_segs_64` with +`tf_paged_attn_merge_segs_64` (`paged_attn_segs.cuh`) for decode rows. All of them are under the kernel contract +(see [the contract](../kernel-contract.md#coverage)). + +Receipt (GB10, nvcc 13.0, `--parallel 8`, greedy): + +- drafted == plain and concurrent == solo on 8 of 8 prompts, and the token hashes are the same after a restart; +- a float64 forward of the same weights on the CPU (`tools/lfm2_moe/ref64.py`) picks the engine's greedy token on + 421 of 435 teacher-forced prompt rows and 123 of 128 reply rows; `tools/lfm2_moe/REF64.md` lists every miss with + both tokens and margins; +- decode, total tok/s: 87 at 1 stream, 194 at 4, 303 at 8; +- cold TTFT: 0.32 s at 2.8k prompt tokens, 1.32 s at 11.2k. + +## Licence + +The weights are under Liquid AI's LFM Open License v1.0 (`license: other`, `lfm1.0` on the model card; the text is the +`LICENSE` file in the checkpoint). Use, changes and redistribution are allowed; commercial use is licensed only to +an organisation under US$10 million in annual revenue, and a non-profit may use it for research or other +non-commercial purposes regardless. Read the licence before serving it commercially. diff --git a/docs/recipes/lfm3-cuda.md b/docs/recipes/lfm3-cuda.md new file mode 100644 index 0000000000..642ff8b407 --- /dev/null +++ b/docs/recipes/lfm3-cuda.md @@ -0,0 +1,72 @@ +# LFM3-1.7B, native CUDA + +`lfm3` serves three of LiquidAI's LFM3-1.7B text exports on one NVIDIA GPU through the lane host. Checked on a GB10 +(sm_121). + +| export | revision | weights on CUDA | +|---|---|---| +| [`LiquidAI/LFM3-1.7B-NVFP4`](https://huggingface.co/LiquidAI/LFM3-1.7B-NVFP4) | `b30025bbbbcbf62f6c110b40d51d36af6579f6a2` | FP4 FFN, bf16 rest, weight-only | +| [`LiquidAI/LFM3-1.7B-FP8`](https://huggingface.co/LiquidAI/LFM3-1.7B-FP8) | `243955da95b1cd68026f293b29de4b4c13f3638f` | block FP8 FFN and q/k/v/o, bf16 rest, weight-only | +| [`LiquidAI/LFM3-1.7B-MLX-4bit`](https://huggingface.co/LiquidAI/LFM3-1.7B-MLX-4bit) | `d0ca73dcfa70ce18733754a492cdc1fb0123262d` | MLX affine 4 bits in groups of 64, as stored | + +The 6- and 8-bit MLX exports and the VL model are not loaded on CUDA yet. Metal refuses the NVFP4 and FP8 exports by +name (`Lfm3ModeloptNvfp4NotOnMetal`, `Lfm3Fp8BlockNotOnMetal`). + +```bash +hf download LiquidAI/LFM3-1.7B-NVFP4 --revision b30025bbbbcbf62f6c110b40d51d36af6579f6a2 \ + --local-dir "$HOME/models/lfm3-1.7b-nvfp4" +zig build native -Doptimize=ReleaseFast -Dnvcc=/usr/local/cuda/bin/nvcc +zig-out/native/bin/tensorfold-native serve "$HOME/models/lfm3-1.7b-nvfp4" --port 8080 --parallel 8 +``` + +## Weight-only (NVFP4 and FP8) + +Both exports run weight-only. NVFP4: the FFN's FP4 weights (w1, w3, w2, E2M1 in groups of 16 with e4m3 block scales and one fp32 +tensor scale) are used as stored, and every activation stays bf16. The export also carries static FP4 input scales; +they are read, because the admission requires them, but not applied. The mixers (conv, q/k/v/gate, o), the embedding +and the tied head are bf16 in the export and run as bf16. + +FP8: the FFN and q/k/v/o weights (e4m3 with a bf16 scale per 128 x 128 block) run on the block-FP8 lane, and the +activations stay bf16; the export's dynamic FP8 input quantization (groups of 128) is not applied. The conv +projections, the attention gate, the embedding and the tied head are bf16. The gate is a projection of its own here, +as it is stored bf16 beside the FP8 q, k and v. + +## MLX 4-bit + +Every projection, the embedding and the tied head are affine 4-bit in groups of 64, used as stored: prompt chunks on +the affine prompt GEMM, decode rows on the affine verify tiles. The embedding's rows are read from the same codes. + +## What runs + +- Conv layers (13): Liquid's 7-tap short conv with its rolling state (`liquid_conv.cuh`, shared with LFM2.5). +- Full-attention layers (6): keys and values in pages of the core KV pool, partial RoPE through `qk_rows`' frequency + table, `pattn_kernel<128, 8, 4, 8>` for prompt rows and the shared paged chunk and merge kernels for decode rows. +- Sliding-attention layers (5, window 256): a ring a stream through the shared window-attention kernels. +- The lane backend, the kept states of the prompt cache and the native glue are shared with LFM2.5-24B-A2B + (`zig/src/families/liquid/`). Kept states live in pages allocated once inside the prompt-cache budget. +- The default window is 32,768 tokens (`--context`); the model allows 262,144. + +## Receipt + +NVFP4, GB10, greedy, `--parallel 8`: + +- drafted == plain and concurrent == solo on 8 of 8 prompts, prompts of 300 to 2,960 tokens across the window, + prompt rows == one-row decode at 32 cut points, resumed == fresh, and the token hashes after a restart; +- a float64 forward of the same weights on the CPU (`tools/lfm3/ref64.py`) picks the engine's greedy token on 825 of + 829 teacher-forced prompt rows and 127 of 128 reply rows; `tools/lfm3/REF64.md` lists every miss with both tokens + and margins; +- decode, total tok/s: 112.8 at 1 stream, 420.0 at 4, 767.4 at 8; cold TTFT 0.150 s at 2k tokens, 0.615 s at 8k, + 2.84 s at 32k. + +FP8 and MLX-4bit, GB10, greedy, `--parallel 8`, drafts off: + +- concurrent == solo: 4 prompts of 64 tokens give the same token hashes alone, 4 at once and 8 at once; +- the float64 forward picks the engine's token on 824 of 829 prompt rows and 128 of 128 reply rows (FP8) and on 825 of + 829 and 128 of 128 (MLX-4bit); every miss is listed in `tools/lfm3/REF64.md`, with float64 margins of 0.004 to 0.085. + +## Licence + +The weights are under Liquid AI's LFM Open License v1.0 (`license: other`, `lfm1.0` on the model card; the text is the +`LICENSE` file in the checkpoint), as the MLX exports are. Use, changes and redistribution are allowed; commercial +use is licensed only to an organisation under US$10 million in annual revenue, and a non-profit may use it for +research or other non-commercial purposes regardless. TensorFold's code stays Apache-2.0. diff --git a/docs/recipes/lfm3-native.md b/docs/recipes/lfm3-native.md new file mode 100644 index 0000000000..7833381748 --- /dev/null +++ b/docs/recipes/lfm3-native.md @@ -0,0 +1,72 @@ +# LFM3 native Metal bring-up + +Tracking issue: [#1491](https://github.com/ashhart/TensorFold/issues/1491). + +This branch adds host checks for LFM3 checkpoint configuration and text tensor layouts. +It does not register a serving engine. Text generation and image input are not implemented or qualified here. + +## Inspect a checkpoint + +Use Python 3.11 or newer. The inspector uses the standard library. It resolves a full Hugging Face commit, +reads the config and index, and requests only the safetensors headers. It refuses a server that returns a +complete weight file instead of the requested byte range. + +```sh +python3 tools/lfm3/inspect_checkpoint.py LiquidAI/LFM3-1.7B-MLX-4bit --output text4.json +python3 tools/lfm3/inspect_checkpoint.py LiquidAI/LFM3-VL-2B-MLX-4bit --output vl4.json +``` + +Replace `4bit` with `6bit` or `8bit` to inspect those exports. Pass `--revision COMMIT` to repeat a pinned inspection. +Use `--local-dir MODEL_DIR` instead of a repository name to inspect files already on disk. + +## Check the native text layout + +Use the Zig version in `.zig-version`. + +```sh +zig build test-lfm3-host tf-lfm3-check +./zig-out/bin/tf-lfm3-check text4.json +./zig-out/bin/tf-lfm3-check vl4.json +python3 -m unittest discover -s tools/lfm3 -p 'test_*.py' -v +``` + +The native checker reads the saved metadata, validates shard header geometry, and checks the LFM3 text decoder's +required names, dtypes, shapes, and byte counts. It performs no projection, GPU allocation, or model inference. +For VL, it checks the text decoder only. A passing result does not validate vision weights or image processing. + +The config reader preserves full-attention versus sliding-attention layers, distinct rotary settings, sandwich norms, +and each module's affine width and group size. Unsupported geometry and conflicting quantization declarations fail. + +## Checkpoint details + +All six inspected MLX exports store the decoder under `language_model.model.*`, including the text-only model. +Each decoder has 594 tensors. Convolution kernels are BF16 with shape `[2048, 7, 1]`. + +Text 4/6/8-bit and VL 6/8-bit use affine groups of 64. VL 4-bit uses groups of 32 and an 8-bit embedding override. +The 138 additional VL vision/projector tensors are BF16 in all three exports. Global quantization metadata must not +cause these dense tensors to be treated as packed affine weights. + +The header inspections used these revisions: + +- Text 4-bit: `d0ca73dcfa70ce18733754a492cdc1fb0123262d`. +- Text 6-bit: `26be898c36b545d625d84ac0761c6c3eb664aef3`. +- Text 8-bit: `ee7667f3cfccb2a41725f91820531823fcabc6b7`. +- VL 4-bit: `c823fade98e35a3d6942ad4d38b1ea15f675cd12`. +- VL 6-bit: `6d71bfae121abedc75f6d531fc7f6ceb4f5e3501`. +- VL 8-bit: `83336d5be3874a2bdd3297bde77b2114dbd3a816`. + +## Next gate + +Implement the text operators against an independent reference before advertising generation. +Check the 7-tap convolution, per-head attention gate, sandwich normalization, sliding window, and partial-RoPE formula. +Then exercise actual 1.7B 4-bit weights through the shared lane engine, followed by 6/8-bit checkpoints. +The VL encoder and projector need their own native Metal operators and reference checks after text bring-up. + +Agree model-fidelity and performance acceptance criteria on the tracking issue before claiming release readiness. + +## Weight licence and checked hardware + +The official text and VL MLX exports retain the [LFM Open License v1.0](https://huggingface.co/LiquidAI/LFM3-1.7B-MLX-4bit/blob/main/LICENSE). +The VL export carries its own [copy of that licence](https://huggingface.co/LiquidAI/LFM3-VL-2B-MLX-4bit/blob/main/LICENSE). +TensorFold code remains Apache-2.0; admission does not change the weights’ licence. +This slice validates metadata on the host and does not register inference or qualify a GPU. diff --git a/tools/lfm2_moe/REF64.md b/tools/lfm2_moe/REF64.md new file mode 100644 index 0000000000..3f5f4be85b --- /dev/null +++ b/tools/lfm2_moe/REF64.md @@ -0,0 +1,65 @@ +# Float64 reference receipt: LFM2.5-24B-A2B on CUDA + +- Checkpoint: [LiquidAI/LFM2.5-24B-A2B-MLX-4bit](https://huggingface.co/LiquidAI/LFM2.5-24B-A2B-MLX-4bit) at revision + `667408d494bde191c2cb9dd4f4549827816b9275`. +- Engine: this PR's serving commit, before the rebase onto main (`c6645063a`). The rebased head serves the same + tokens: the token hashes of 8 prompts and of a chat conversation are equal. + Served on a GB10 (sm_121), greedy, drafts off, `--parallel 8 --context 16384`. +- Reference: `ref64.py` on the CPU in float64 (Python 3.12, torch 2.14, safetensors 0.8, tokenizers 0.23), 2026-10-28. + +Reproduce (the reference pass takes about 10 minutes and 17 GB of host memory): + +``` +zig-out/native/bin/tensorfold-native serve MODEL_DIR --port 8080 --parallel 8 --context 16384 +python -I tools/lfm2_moe/ref64.py collect http://127.0.0.1:8080 MODEL_DIR engine.json +python -I tools/lfm2_moe/ref64.py MODEL_DIR engine.json --bf16 +``` + +Rows. A prompt row is the engine's one-token greedy pick after the prefix [0, t] of a prompt, one request per prefix. +A reply row is a token of the engine's 32-token greedy reply, teacher-forced on that reply. `t` is the logits row after +token t, so prompt row P - 1 and reply token 0 are the same position. The prompts (`prompts()` in the script) are +capitals (13 tokens), code (19), a chat turn (19) and items (384, a list of 30 items and a question). + +The float64 pass dequantizes the MLX weights exactly. The `--bf16` pass adds the engine's bf16 storage points (see the +script's docstring). It shows which side of a miss rounding alone lands on; it is not bit-exact. + +## Result + +- f64: prompt rows 421/435, reply rows 123/128 +- bf16: prompt rows 424/435, reply rows 123/128 + +Float64 top-1 misses (margin = f64 logit of its top minus f64 logit of the engine's token; bf16 margin = the same difference in the bf16-points pass, < 0 means it prefers the engine's token): + +| case | row | t | engine | float64 | f64 margin | bf16 margin | bf16 picks | +|---|---|---|---|---|---|---|---| +| capitals | prompt | 1 | 589 "'s" | 124901 "" | 0.3389 | -0.1250 (pre-round -0.1026) | engine's | +| capitals | prompt | 3 | 1394 " each" | 278 " the" | 0.0704 | -1.0625 (pre-round -1.0871) | engine's | +| capitals | reply | 39 | 278 " the" | 17755 " Greece" | 0.0095 | +0.0000 (pre-round -0.0276) | engine's | +| code | prompt | 7 | 1288 " return" | 21744 " \"\"\"" | 0.3008 | +0.2500 (pre-round +0.2923) | f64's | +| code | reply | 26 | 4876 "_m" | 71 "_" | 0.3628 | +0.3125 (pre-round +0.2956) | f64's | +| code | reply | 35 | 73 "a" | 2762 "add" | 0.0687 | -0.3750 (pre-round -0.4420) | engine's | +| chat | prompt | 13 | 90978 " Explanation" | 3103 " \n" | 0.7478 | -0.8125 (pre-round -0.7755) | engine's | +| chat | prompt | 15 | 25 "1" | 3103 " \n" | 0.1301 | +0.0625 (pre-round +0.0216) | f64's | +| items | prompt | 3 | 22 "." | 34 ":" | 0.9143 | +0.8750 (pre-round +0.8393) | f64's | +| items | prompt | 5 | 2719 " problem" | 3451 " question" | 0.1379 | +0.0625 (pre-round +0.0963) | f64's | +| items | prompt | 6 | 9402 " stone" | 24062 " brick" | 0.1680 | +1.0000 (pre-round +0.9824) | f64's | +| items | prompt | 8 | 229 " " | 267 " a" | 0.0823 | -0.1250 (pre-round -0.0781) | engine's | +| items | prompt | 314 | 27 "3" | 1035 "14" | 0.1292 | -0.3750 (pre-round -0.3704) | engine's | +| items | prompt | 338 | 25 "1" | 26 "2" | 0.3699 | +0.3750 (pre-round +0.4515) | f64's | +| items | prompt | 371 | 437 " are" | 296 " in" | 0.0565 | -0.6250 (pre-round -0.5979) | engine's | +| items | prompt | 375 | 342 " (" | 8187 " item" | 0.3365 | -1.1250 (pre-round -1.1544) | engine's | +| items | prompt | 377 | 342 " (" | 41501 " Item" | 0.6276 | -2.3750 (pre-round -2.3544) | engine's | +| items | reply | 387 | 41501 " Item" | 342 " (" | 0.7665 | +0.6250 (pre-round +0.5957) | f64's | +| items | reply | 396 | 2959 " red" | 6542 " blue" | 0.2816 | +0.5000 (pre-round +0.5314) | f64's | + +bf16-points misses where float64 agrees with the engine (margin = bf16 logit of its pick minus the engine's token; f64 margin = how far float64 prefers the engine's token): + +| case | row | t | engine | bf16 pick | bf16 margin | f64 margin | +|---|---|---|---|---|---|---| +| capitals | prompt | 10 | 10125 " Italy" | 6398 " Germany" | 0.0000 (pre-round -0.0632) | 0.0079 | +| code | reply | 44 | 1113 " add" | 1619 " main" | 0.6250 (pre-round +0.5908) | 0.1500 | +| code | reply | 47 | 71 "_" | 65296 "_sum" | 0.1250 (pre-round +0.1222) | 0.3751 | +| items | prompt | 7 | 22 "." | 20 "," | 0.0000 (pre-round +0.0882) | 0.0279 | +| items | prompt | 9 | 30 "6" | 27 "3" | 0.5000 (pre-round +0.4381) | 0.1054 | +| items | prompt | 36 | 440 " The" | 41501 " Item" | 0.1250 (pre-round +0.1183) | 0.1051 | +| items | prompt | 379 | 8187 " item" | 229 " " | 0.2500 (pre-round +0.2420) | 0.0420 | diff --git a/tools/lfm2_moe/ref64.py b/tools/lfm2_moe/ref64.py new file mode 100644 index 0000000000..bce94597dc --- /dev/null +++ b/tools/lfm2_moe/ref64.py @@ -0,0 +1,254 @@ +#!/usr/bin/env python3 +"""Float64 CPU reference of LFM2.5-24B-A2B (lfm2_moe) from its MLX affine checkpoint, checked against the CUDA +engine's greedy tokens. + +The stored weights are dequantized exactly (q * scale + bias) and every op runs in float64. The engine's greedy +tokens come from a running server (`collect`): + prompt rows the one-token greedy pick after each prefix [0, t] of a prompt (a fresh request per prefix) + reply rows a 32-token greedy reply, checked teacher-forced on the engine's own reply +Row t is the logits row after token t, so prompt row P - 1 and reply token 0 are the same position. + +--bf16 adds a second pass with the engine's bf16 storage points (float64 inside every op, round-to-nearest-even to +bf16 where the CUDA kernels store bf16): embed, residual, every RMSNorm (bf16(w * bf16(h * inv))), projection outputs, +q/k norm and RoPE, conv B*x / acc / C*y, SwiGLU, expert outputs, MoE combine and the logits. Prompt rows take the +prefill path (weights rounded once to bf16, attention P rounded before PV, expert down stored bf16), decode rows the +exact-weight GEMV path (fp32 attention P, fp32 expert down). The router stays unrounded. Softmax P is taken against +the row's final max (the kernels' running max over 64-key tiles is not replayed). Greedy over bf16 logits breaks ties +to the lower id, as the engine does. This pass shows which side of a miss rounding alone lands on; it is not bit-exact. + +usage: + ref64.py collect URL MODEL_DIR ENGINE_JSON the greedy rows of prompts() from a server at URL + ref64.py MODEL_DIR ENGINE_JSON [--bf16] prints the miss table (markdown) and the agreement totals +""" +import json +import math +import sys +import urllib.request + +import torch +from safetensors import safe_open + +torch.set_default_dtype(torch.float64) +F64 = torch.float64 + + +def complete(url, ids, n): + body = {"prompt": ids, "max_tokens": n, "temperature": 0, "ignore_eos": True, "draft": False, + "return_token_ids": True} + req = urllib.request.Request(url + "/v1/completions", data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}) + with urllib.request.urlopen(req, timeout=900) as r: + out = json.loads(r.read()) + return out.get("tensorfold", out)["token_ids"] + + +def prompts(n_items): + """The receipt's prompt set: texts, encoded with the checkpoint's tokenizer.json after BOS.""" + colors = "red green blue amber violet".split() + items = " ".join(f"Item {i}: the {colors[i % 5]} box holds {7 * i % 23} stones." for i in range(n_items)) + return [("capitals", "The capital of France is Paris. The capital of Germany is"), + ("code", "def add(a, b):\n return a + b\n\ndef mul(a, b):\n"), + ("chat", "<|im_start|>user\nExplain why the sky is blue in two sentences.<|im_end|>\n" + "<|im_start|>assistant\n"), + ("items", items + " Question: how many stones does item 3 hold? Answer: item 3 holds")] + + +def collect(url, model_dir, out_path, reply_len=32): + from tokenizers import Tokenizer + tok = Tokenizer.from_file(f"{model_dir}/tokenizer.json") + bos = json.load(open(f"{model_dir}/config.json"))["bos_token_id"] + cases = [] + for name, text in prompts(30): + ids = [bos] + tok.encode(text).ids + picks = [complete(url, ids[: t + 1], 1)[0] for t in range(len(ids))] + cases.append({"name": name, "prompt": ids, "prompt_picks": picks, "reply": complete(url, ids, reply_len)}) + print(f"{name}: {len(ids)} prompt rows, {len(cases[-1]['reply'])} reply tokens", file=sys.stderr) + json.dump({"cases": cases}, open(out_path, "w")) + + +def bf16(x): + """Nearest bf16 (8 significant bits, ties to even), kept as float64: exact, no float32 double rounding.""" + m, e = torch.frexp(x) + return torch.ldexp(torch.round(m * 256.0), (e - 8).to(torch.int32)).to(F64) + + +def same(x): + return x + + +class Model: + def __init__(self, d): + self.d = d + self.cfg = json.load(open(f"{d}/config.json")) + self.map = json.load(open(f"{d}/model.safetensors.index.json"))["weight_map"] + self.files = {} + q = self.cfg["quantization"] + self.embed_bits = q.get("model.embed_tokens", q)["bits"] + + def raw(self, k): + f = self.map[k] + if f not in self.files: + self.files[f] = safe_open(f"{self.d}/{f}", "pt") + return self.files[f].get_tensor(k) + + def deq(self, prefix, bits=4, rows=None): + """MLX affine: codes are a little-endian bitstream over uint32 words, value i at bits [i*bits, (i+1)*bits).""" + w, s, b = (self.raw(prefix + n) for n in (".weight", ".scales", ".biases")) + if rows is not None: + w, s, b = w[rows], s[rows], b[rows] + by = w.contiguous().view(torch.uint8).to(torch.int64) + if 32 % bits == 0: + words = by.reshape(*by.shape[:-1], -1, 4) << torch.tensor([0, 8, 16, 24]) + words = words.sum(-1, keepdim=True) + q = ((words >> (bits * torch.arange(32 // bits))) & ((1 << bits) - 1)).flatten(-2) + else: + nb = by.shape[-1] * 8 + bit = ((by.unsqueeze(-1) >> torch.arange(8)) & 1).reshape(*by.shape[:-1], nb) + q = (bit[..., : nb // bits * bits].reshape(*by.shape[:-1], -1, bits) << torch.arange(bits)).sum(-1) + q = q.to(F64).reshape(*q.shape[:-1], s.shape[-1], -1) + return (q * s.to(F64)[..., None] + b.to(F64)[..., None]).flatten(-2) + + def forward(self, ids, P, bf): + """Logits [T, V] (float64, not rounded). bf = same: plain float64; bf = bf16: the engine's storage points with + rows [0, P) on the prefill path and the rest on the decode path.""" + cfg, T = self.cfg, len(ids) + H, HK, hd = cfg["num_attention_heads"], cfg["num_key_value_heads"], 64 + silu = torch.nn.functional.silu + + def lin(x, W, p=P): # prefill rows multiply bf16(W), decode rows the exact W; fp64 accumulation + return torch.cat([x[:p] @ bf(W).T, x[p:] @ W.T]) + + def rms(h, w): # the engine's form: bf16(w * bf16(h * inv)), inv from the bf16 row + return bf(w.to(F64) * bf(h * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + cfg["norm_eps"]))) + + inv = 1.0 / cfg["rope_parameters"]["rope_theta"] ** (torch.arange(0, hd, 2, dtype=F64) / hd) + ang = torch.arange(T, dtype=F64)[:, None] * inv[None] + cos, sin = torch.cat([ang.cos()] * 2, -1), torch.cat([ang.sin()] * 2, -1) + + def rot(x): + return x * cos + torch.cat([-x[..., hd // 2:], x[..., : hd // 2]], -1) * sin + + future = torch.triu(torch.ones(T, T, dtype=torch.bool), 1) + h = bf(self.deq("model.embed_tokens", self.embed_bits, rows=torch.tensor(ids))) + for L, kind in enumerate(cfg["layer_types"]): + p = f"model.layers.{L}." + x = rms(h, self.raw(p + "operator_norm.weight")) + if kind == "conv": + B, C, xx = bf(lin(x, self.deq(p + "conv.in_proj"))).chunk(3, -1) + bx = bf(B * xx) + w = self.raw(p + "conv.conv.weight").to(F64).reshape(-1, 3) # tap 2 is the current row + pad = torch.cat([torch.zeros(2, bx.shape[1]), bx]) + y = bf(C * bf(pad[0:T] * w[:, 0] + pad[1:T + 1] * w[:, 1] + pad[2:T + 2] * w[:, 2])) + o = bf(lin(y, self.deq(p + "conv.out_proj"))) + else: + q, k, v = (bf(lin(x, self.deq(p + f"self_attn.{n}_proj"))) for n in "qkv") + q = bf(rot(rms(q.reshape(T, H, hd), self.raw(p + "self_attn.q_layernorm.weight")).transpose(0, 1))) + k = bf(rot(rms(k.reshape(T, HK, hd), self.raw(p + "self_attn.k_layernorm.weight")).transpose(0, 1))) + k = k.repeat_interleave(H // HK, 0) + v = v.reshape(T, HK, hd).transpose(0, 1).repeat_interleave(H // HK, 0) + s = ((q @ k.transpose(1, 2)) / math.sqrt(hd)).masked_fill(future, -math.inf) + e = torch.exp(s - s.max(-1, keepdim=True).values) + pv = torch.cat([bf(e[:, :P]), e[:, P:]], 1) # prefill rounds P before the PV product, decode keeps it + a = bf((pv @ v) / e.sum(-1, keepdim=True)) + o = bf(lin(a.transpose(0, 1).reshape(T, H * hd), self.deq(p + "self_attn.out_proj"))) + h = bf(h + o) + x = rms(h, self.raw(p + "ffn_norm.weight")) + f = p + "feed_forward." + if L < cfg["num_dense_layers"]: + g, u = bf(lin(x, self.deq(f + "gate_proj"))), bf(lin(x, self.deq(f + "up_proj"))) + o = bf(lin(bf(bf(silu(g)) * u), self.deq(f + "down_proj"))) + else: + sc = torch.sigmoid(x @ self.raw(f + "gate.weight").to(F64).T) + pick = torch.topk(sc + self.raw(f + "expert_bias").to(F64), cfg["num_experts_per_tok"], -1).indices + wts = sc.gather(-1, pick) + wts = wts / (wts.sum(-1, keepdim=True) + 1e-6) * cfg["routed_scaling_factor"] + acc = torch.zeros_like(x) + m = f + "switch_mlp." + for ex in pick.unique().tolist(): + sel = pick == ex + rr = sel.any(-1).nonzero().squeeze(-1) # ascending: its first pe rows are prefill rows + pe = int((rr < P).sum()) + g = bf(lin(x[rr], self.deq(m + "gate_proj", rows=ex), pe)) + u = bf(lin(x[rr], self.deq(m + "up_proj", rows=ex), pe)) + y = lin(bf(bf(silu(g)) * u), self.deq(m + "down_proj", rows=ex), pe) + acc[rr] += (wts * sel).sum(-1)[rr, None] * torch.cat([bf(y[:pe]), y[pe:]]) + o = bf(acc) + h = bf(h + o) + print(f" layer {L} {kind}", file=sys.stderr, flush=True) + x = rms(h, self.raw("model.embedding_norm.weight")) + V = cfg["vocab_size"] + return torch.cat([x @ self.deq("model.embed_tokens", self.embed_bits, rows=torch.arange(a, min(a + 8000, V))).T + for a in range(0, V, 8000)], -1) + + +def greedy(lg): + """argmax in (value desc, id asc) order, as the engine picks over its bf16 logits.""" + return int((lg == lg.max()).nonzero()[0]) + + +def main(model_dir, engine_path, with_bf16): + mdl = Model(model_dir) + try: + from tokenizers import Tokenizer + tok = Tokenizer.from_file(f"{model_dir}/tokenizer.json") + + def show(t): + return f"{t} {json.dumps(tok.decode([t], skip_special_tokens=False))}".replace("|", "\\|") + except ImportError: + show = str + cases = json.load(open(engine_path))["cases"] + tot = {(m, k): [0, 0] for m in ("f64", "bf16") for k in ("prompt", "reply")} + f64_rows, bf16_rows = [], [] + for c in cases: + P = len(c["prompt"]) + ids = c["prompt"] + c["reply"][:-1] + print(f"{c['name']}: {len(ids)} rows", file=sys.stderr, flush=True) + lf = mdl.forward(ids, P, same) + lb_pre = mdl.forward(ids, P, bf16) if with_bf16 else None + lb = bf16(lb_pre) if with_bf16 else None + rows = [("prompt", t, tk) for t, tk in enumerate(c["prompt_picks"])] + rows += [("reply", P - 1 + j, tk) for j, tk in enumerate(c["reply"])] + for kind, t, eng in rows: + top = int(lf[t].argmax()) + tot["f64", kind][0] += top == eng + tot["f64", kind][1] += 1 + bcol = "" + if with_bf16: + bp = greedy(lb[t]) + tot["bf16", kind][0] += bp == eng + tot["bf16", kind][1] += 1 + if top != eng: + if with_bf16: + d, dp = lb[t, top] - lb[t, eng], lb_pre[t, top] - lb_pre[t, eng] + side = "engine's" if bp == eng else "f64's" if bp == top else f"other {show(bp)}" + bcol = f" {d:+.4f} (pre-round {dp:+.4f}) | {side} |" + f64_rows.append(f"| {c['name']} | {kind} | {t} | {show(eng)} | {show(top)} | " + f"{lf[t, top] - lf[t, eng]:.4f} |{bcol}") + elif with_bf16 and bp != eng: + bf16_rows.append(f"| {c['name']} | {kind} | {t} | {show(eng)} | {show(bp)} | " + f"{lb[t, bp] - lb[t, eng]:.4f} (pre-round {lb_pre[t, bp] - lb_pre[t, eng]:+.4f}) | " + f"{lf[t, eng] - lf[t, bp]:.4f} |") + mode = [m for m in ("f64", "bf16") if m == "f64" or with_bf16] + for m in mode: + print(f"{m}: prompt rows {tot[m, 'prompt'][0]}/{tot[m, 'prompt'][1]}, " + f"reply rows {tot[m, 'reply'][0]}/{tot[m, 'reply'][1]}") + print("\nFloat64 top-1 misses (margin = f64 logit of its top minus f64 logit of the engine's token" + + ("; bf16 margin = the same difference in the bf16-points pass, < 0 means it prefers the engine's token" + if with_bf16 else "") + "):\n") + print("| case | row | t | engine | float64 | f64 margin |" + (" bf16 margin | bf16 picks |" if with_bf16 else "")) + print("|---|---|---|---|---|---|" + ("---|---|" if with_bf16 else "")) + print("\n".join(f64_rows) or "(none)") + if with_bf16: + print("\nbf16-points misses where float64 agrees with the engine (margin = bf16 logit of its pick minus the " + "engine's token; f64 margin = how far float64 prefers the engine's token):\n") + print("| case | row | t | engine | bf16 pick | bf16 margin | f64 margin |\n|---|---|---|---|---|---|---|") + print("\n".join(bf16_rows) or "(none)") + + +if __name__ == "__main__": + if len(sys.argv) == 5 and sys.argv[1] == "collect": + collect(*sys.argv[2:]) + elif len(sys.argv) in (3, 4) and sys.argv[3:] in ([], ["--bf16"]): + main(sys.argv[1], sys.argv[2], len(sys.argv) == 4) + else: + sys.exit(__doc__) diff --git a/tools/lfm3/REF64.md b/tools/lfm3/REF64.md new file mode 100644 index 0000000000..2241109eaf --- /dev/null +++ b/tools/lfm3/REF64.md @@ -0,0 +1,76 @@ +# Float64 reference receipt: LFM3-1.7B NVFP4 on CUDA + +- Checkpoint: [LiquidAI/LFM3-1.7B-NVFP4](https://huggingface.co/LiquidAI/LFM3-1.7B-NVFP4) at revision + `b30025bbbbcbf62f6c110b40d51d36af6579f6a2`. +- Engine: `9265a0f5f`, which is the PR head `08b338c7c` before its rebase (range-diff: the same patches except one + test-only line in `contract_limits.zig`). Served on a GB10 (sm_121), greedy, drafts off, + `--parallel 8 --context 32768`. The `collect` step was repeated on 2026-10-28 and returned the same 957 tokens. +- Reference: `ref64.py` on the CPU in float64 (Python 3.12, torch 2.14, safetensors 0.8, tokenizers 0.23), 2026-10-28. + +Reproduce (the reference pass takes about a minute and 20 GB of host memory): + +``` +zig-out/native/bin/tensorfold-native serve MODEL_DIR --port 8080 --parallel 8 --context 32768 +python -I tools/lfm3/ref64.py collect http://127.0.0.1:8080 MODEL_DIR engine.json +python -I tools/lfm3/ref64.py MODEL_DIR engine.json --bf16 +``` + +Rows. A prompt row is the engine's one-token greedy pick after the prefix [0, t] of a prompt, one request per prefix. +A reply row is a token of the engine's 32-token greedy reply, teacher-forced on that reply. `t` is the logits row after +token t, so prompt row P - 1 and reply token 0 are the same position. The prompts (`prompts()` in the script) are +capitals (20 tokens), code (40), a chat turn (19) and items (750, a list of 60 items and a question; it crosses the +256-key sliding window). + +The float64 pass dequantizes the NVFP4 weights exactly. The `--bf16` pass adds the engine's bf16 storage points (see +the script's docstring). It shows which side of a miss rounding alone lands on; it is not bit-exact. + +## Result + +- f64: prompt rows 825/829, reply rows 127/128 +- bf16: prompt rows 827/829, reply rows 126/128 + +Float64 top-1 misses (margin = f64 logit of its top minus f64 logit of the engine's token; bf16 margin = the same difference in the bf16-points pass, < 0 means it prefers the engine's token): + +| case | row | t | engine | float64 | f64 margin | bf16 margin | bf16 picks | +|---|---|---|---|---|---|---|---| +| capitals | reply | 24 | 11097 " Spain" | 3888 " Japan" | 0.1176 | +0.1250 (pre-round +0.1802) | f64's | +| chat | prompt | 12 | 593 " per" | 22 "." | 0.1530 | -0.0625 (pre-round -0.0671) | engine's | +| items | prompt | 5 | 1008 " first" | 1579 " number" | 0.0996 | +0.0000 (pre-round +0.1009) | engine's | +| items | prompt | 9 | 29 "5" | 654 "10" | 0.0033 | +0.0000 (pre-round +0.0093) | engine's | +| items | prompt | 180 | 28 "4" | 30 "6" | 0.0674 | +0.0000 (pre-round +0.0076) | engine's | + +bf16-points misses where float64 agrees with the engine (margin = bf16 logit of its pick minus the engine's token; f64 margin = how far float64 prefers the engine's token): + +| case | row | t | engine | bf16 pick | bf16 margin | f64 margin | +|---|---|---|---|---|---|---| +| items | prompt | 4 | 440 " The" | 229 " " | 0.0000 (pre-round -0.0985) | 0.0881 | +| items | prompt | 631 | 32 "8" | 25 "1" | 0.0000 (pre-round -0.0467) | 0.0636 | +| items | reply | 772 | 41501 " Item" | 1672 " So" | 0.0000 (pre-round -0.0814) | 0.1797 | + +## FP8 and MLX-4bit (2026-10-29) + +- Checkpoints: [LiquidAI/LFM3-1.7B-FP8](https://huggingface.co/LiquidAI/LFM3-1.7B-FP8) at + `243955da95b1cd68026f293b29de4b4c13f3638f` (block FP8, dequantized exactly: e4m3 times its block's bf16 scale) and + [LiquidAI/LFM3-1.7B-MLX-4bit](https://huggingface.co/LiquidAI/LFM3-1.7B-MLX-4bit) at + `d0ca73dcfa70ce18733754a492cdc1fb0123262d` (MLX affine, dequantized exactly). +- Served the same way with `--no-drafts`; the same prompts and rows. Only the float64 pass was run: the `--bf16` + pass models the NVFP4 engine's rounding points. + +FP8: prompt rows 824/829, reply rows 128/128. + +| case | row | t | engine | float64 | f64 margin | +|---|---|---|---|---|---| +| items | prompt | 131 | 30 "6" | 1325 "24" | 0.0348 | +| items | prompt | 582 | 654 "10" | 1035 "14" | 0.0756 | +| items | prompt | 619 | 25 "1" | 29 "5" | 0.0066 | +| items | prompt | 631 | 27 "3" | 32 "8" | 0.0290 | +| items | prompt | 693 | 9402 " stone" | 17937 " stones" | 0.0773 | + +MLX-4bit: prompt rows 825/829, reply rows 128/128. + +| case | row | t | engine | float64 | f64 margin | +|---|---|---|---|---|---| +| items | prompt | 10 | 6802 " items" | 93266 " marbles" | 0.0466 | +| items | prompt | 61 | 41501 " Item" | 29155 " Question" | 0.0042 | +| items | prompt | 595 | 27 "3" | 1035 "14" | 0.0300 | +| items | prompt | 748 | 34 ":" | 10494 " holds" | 0.0848 | diff --git a/tools/lfm3/inspect_checkpoint.py b/tools/lfm3/inspect_checkpoint.py new file mode 100644 index 0000000000..2fec63457d --- /dev/null +++ b/tools/lfm3/inspect_checkpoint.py @@ -0,0 +1,237 @@ +#!/usr/bin/env python3 +"""Inspect pinned Hugging Face configs and safetensors headers without downloading tensor data.""" + +import argparse +import collections +import json +import math +import os +import re +import struct +import sys +import urllib.error +import urllib.parse +import urllib.request +from pathlib import Path, PurePosixPath + +MAX_JSON = 16 * 1024 * 1024 +MAX_HEADER = 16 * 1024 * 1024 + + +class InspectionError(Exception): + pass + + +def safe_name(name): + path = PurePosixPath(name) + if not name or path.is_absolute() or '..' in path.parts or '\\' in name: + raise InspectionError(f'Unsafe checkpoint file name: {name!r}') + return name + + +def read_json_bytes(data, description): + try: + return json.loads(data) + except (ValueError, UnicodeError) as exc: + raise InspectionError(f'Invalid JSON in {description}: {exc}') from exc + + +class Remote: + def __init__(self, repository, revision): + if not re.fullmatch(r'[\w.-]+/[\w.-]+', repository): + raise InspectionError('Repository must have owner/name form') + api = 'https://huggingface.co/api/models/' + repository + '/revision/' + metadata = self.json_url(api + urllib.parse.quote(revision, safe='')) + self.revision = metadata.get('sha', '') + if not re.fullmatch(r'[0-9a-f]{40}', self.revision): + raise InspectionError('Hugging Face did not resolve a full commit revision') + self.repository = repository + self.files = {entry['rfilename'] for entry in metadata.get('siblings', [])} + self.base = f'https://huggingface.co/{repository}/resolve/{self.revision}/' + + @staticmethod + def fetch(url, limit, byte_range=None, with_total=False): + headers = {'User-Agent': 'TensorFold-checkpoint-inspector', 'Accept-Encoding': 'identity'} + if byte_range is not None: + start, end = byte_range + headers['Range'] = f'bytes={start}-{end}' + try: + with urllib.request.urlopen(urllib.request.Request(url, headers=headers), timeout=60) as response: + if byte_range is not None: + match = re.fullmatch(r'bytes (\d+)-(\d+)/(\d+)', response.headers.get('Content-Range', '')) + if response.status != 206 or not match: + raise InspectionError('Server refused an exact HTTP byte range; refusing tensor download') + first, last, total = map(int, match.groups()) + if (first, last) != byte_range or total <= last: + raise InspectionError('Server returned an unexpected byte range') + if response.headers.get('Content-Encoding', 'identity') != 'identity': + raise InspectionError('Compressed range response is not supported') + length = response.headers.get('Content-Length') + if length is not None and int(length) > limit: + raise InspectionError(f'Response exceeds {limit}-byte body limit') + data = response.read(limit + 1) + if len(data) > limit: + raise InspectionError(f'Response exceeds {limit}-byte body limit') + if byte_range is not None and len(data) != byte_range[1] - byte_range[0] + 1: + raise InspectionError('Truncated byte-range response') + return (data, total) if with_total else data + except (urllib.error.URLError, TimeoutError, OSError, ValueError) as exc: + raise InspectionError(f'Cannot fetch {url}: {exc}') from exc + + @classmethod + def json_url(cls, url): + return read_json_bytes(cls.fetch(url, MAX_JSON), url) + + def json(self, name): + return self.json_url(self.base + urllib.parse.quote(safe_name(name), safe='/')) + + def header(self, name): + url = self.base + urllib.parse.quote(safe_name(name), safe='/') + prefix, total = self.fetch(url, 8, (0, 7), with_total=True) + size = header_size(prefix) + if total < 8 + size: + raise InspectionError('Safetensors header extends beyond shard size') + data, second_total = self.fetch(url, size, (8, 7 + size), with_total=True) + if total != second_total: + raise InspectionError('Shard size changed between header requests') + return read_json_bytes(data, name), total - 8 - size + + +class Local: + def __init__(self, directory): + self.directory = Path(directory).resolve() + self.repository = str(self.directory) + self.revision = None + self.files = {str(path.relative_to(self.directory)) for path in self.directory.rglob('*.safetensors')} + if (self.directory / 'model.safetensors.index.json').is_file(): + self.files.add('model.safetensors.index.json') + + def path(self, name): + path = (self.directory / safe_name(name)).resolve() + if not path.is_relative_to(self.directory): + raise InspectionError('Checkpoint file escapes local directory') + return path + + def json(self, name): + with self.path(name).open('rb') as stream: + data = stream.read(MAX_JSON + 1) + if len(data) > MAX_JSON: + raise InspectionError('Local JSON exceeds body limit') + return read_json_bytes(data, name) + + def header(self, name): + with self.path(name).open('rb') as stream: + size = header_size(stream.read(8)) + data = stream.read(size) + total = os.fstat(stream.fileno()).st_size + if len(data) != size: + raise InspectionError('Truncated local safetensors header') + return read_json_bytes(data, name), total - 8 - size + + +def header_size(prefix): + if len(prefix) != 8: + raise InspectionError('Truncated safetensors size prefix') + size = struct.unpack('user\nExplain why the sky is blue in two sentences.<|im_end|>\n" + "<|im_start|>assistant\n"), + ("items", items + " Question: how many stones does item 3 hold? Answer: item 3 holds")] + + +def collect(url, model_dir, out_path, reply_len=32): + from tokenizers import Tokenizer + tok = Tokenizer.from_file(f"{model_dir}/tokenizer.json") + bos = json.load(open(f"{model_dir}/config.json"))["bos_token_id"] + cases = [] + for name, text in prompts(60): + ids = [bos] + tok.encode(text).ids + picks = [complete(url, ids[: t + 1], 1)[0] for t in range(len(ids))] + cases.append({"name": name, "prompt": ids, "prompt_picks": picks, "reply": complete(url, ids, reply_len)}) + print(f"{name}: {len(ids)} prompt rows, {len(cases[-1]['reply'])} reply tokens", file=sys.stderr) + json.dump({"cases": cases}, open(out_path, "w")) + + +def bf16(x): + """Nearest bf16 (8 significant bits, ties to even), kept as float64: exact, no float32 double rounding.""" + m, e = torch.frexp(x) + return torch.ldexp(torch.round(m * 256.0), (e - 8).to(torch.int32)).to(F64) + + +def same(x): + return x + + +def rot(t, cos, sin): + hd = t.shape[-1] + return t * cos + torch.cat([-t[..., hd // 2:], t[..., : hd // 2]], -1) * sin + + +def strip(k): + return re.sub(r"^(model\.language_model\.|language_model\.|model\.)+", "", k) + + +class Model: + def __init__(self, d): + self.d = d + root = json.load(open(f"{d}/config.json")) + self.cfg = root.get("text_config", root) + idx = f"{d}/model.safetensors.index.json" + if os.path.exists(idx): + self.map = json.load(open(idx))["weight_map"] + else: + with safe_open(f"{d}/model.safetensors", "pt") as f: + self.map = {k: "model.safetensors" for k in f.keys()} + self.files, self.cache = {}, {} + self.prefix = "language_model.model." if "language_model.model.embed_tokens.weight" in self.map else "model." + self.mlx = root.get("quantization") if root.get("quantization", {}).get("mode", "affine") == "affine" else None + + def raw(self, k): + f = self.map[k] + if f not in self.files: + self.files[f] = safe_open(f"{self.d}/{f}", "pt") + return self.files[f].get_tensor(k) + + def vec(self, name): + return self.raw(self.prefix + name).to(F64) + + def lin(self, name, rows=None): + """A Linear weight [out, in] in float64, dequantized exactly by its stored format.""" + p = self.prefix + name + w = self.raw(p + ".weight") + if rows is not None: + w = w[rows] + if self.mlx and p + ".scales" in self.map: + s, b = self.raw(p + ".scales"), self.raw(p + ".biases") + if rows is not None: + s, b = s[rows], b[rows] + bits = next((v["bits"] for k, v in self.mlx.items() if isinstance(v, dict) and strip(k) == strip(p)), + self.mlx["bits"]) + by = w.contiguous().view(torch.uint8).to(torch.int64) + nb = by.shape[-1] * 8 + bit = ((by.unsqueeze(-1) >> torch.arange(8)) & 1).reshape(*by.shape[:-1], nb) + q = (bit[..., : nb // bits * bits].reshape(*by.shape[:-1], -1, bits) << torch.arange(bits)).sum(-1) + q = q.to(F64).reshape(*q.shape[:-1], s.shape[-1], -1) + return (q * s.to(F64)[..., None] + b.to(F64)[..., None]).flatten(-2) + if p + ".weight_scale_2" in self.map: + sc = self.raw(p + ".weight_scale") + if rows is not None: + sc = sc[rows] + n = w.shape[0] + codes = torch.stack([w & 0xF, w >> 4], dim=-1).reshape(n, -1).long() + v = E2M1[codes].view(n, -1, 16) * sc.to(F64)[..., None] + return (v * self.raw(p + ".weight_scale_2").to(F64)).reshape(n, -1) + if w.dtype == torch.float8_e4m3fn: + sc = self.raw(p + ".weight_scale").to(F64).repeat_interleave(128, 0).repeat_interleave(128, 1) + return w.to(F64) * (sc if rows is None else sc[rows]) + return w.to(F64) + + def W(self, name): + if name not in self.cache: + self.cache[name] = self.lin(name) + return self.cache[name] + + def forward(self, ids, P, bf): + """Logits [T, V] (float64, not rounded). bf = same: plain float64; bf = bf16: the engine's storage points with + rows [0, P) on the prompt path and the rest on the decode path.""" + cfg, T = self.cfg, len(ids) + D, H, HK, hd = cfg["hidden_size"], cfg["num_attention_heads"], cfg["num_key_value_heads"], cfg["head_dim"] + taps, vec, W = cfg["conv_L_cache"], self.vec, self.W + + def rms(h, w): # the engine's form: bf16(w * bf16(h * inv)), inv from the bf16 row + return bf(w * bf(h * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + cfg["norm_eps"]))) + + tabs = {} + for kind, rp in cfg["rope_parameters"].items(): + angles = int(rp.get("partial_rotary_factor", 1.0) * hd) // 2 + inv = 1.0 / rp["rope_theta"] ** (torch.arange(0, 2 * angles, 2, dtype=F64) / hd) # the full head dim's + inv = torch.cat([inv, torch.zeros(hd // 2 - angles)]) + ang = torch.arange(T, dtype=F64)[:, None] * inv[None] + tabs[kind] = torch.cat([ang.cos()] * 2, -1), torch.cat([ang.sin()] * 2, -1) + i, j = torch.arange(T)[:, None], torch.arange(T)[None] + masks = {"full_attention": j > i, "sliding_attention": (j > i) | (i - j >= cfg["sliding_window"])} + h = self.lin("embed_tokens", rows=torch.tensor(ids)) + for L, kind in enumerate(cfg["layer_types"]): + p = f"layers.{L}." + x = rms(h, vec(p + "operator_norm.weight")) + if kind == "conv": + B, C, xx = bf(x @ W(p + "conv.in_proj").T).chunk(3, -1) + bx = bf(B * xx) + w = vec(p + "conv.conv.weight").reshape(D, taps) + pad = torch.cat([torch.zeros(taps - 1, D), bx]) + acc = torch.zeros(T, D) + for k in range(taps): # oldest tap first + acc = acc + pad[k:k + T] * w[:, k] + o = bf(bf(C * bf(acc)) @ W(p + "conv.out_proj").T) + else: + q, k, v = (bf(x @ W(p + f"self_attn.{n}_proj").T) for n in "qkv") + g = bf(x @ W(p + "self_attn.gate_proj").T) # [T, H] + q = rms(q.reshape(T, H, hd), vec(p + "self_attn.q_norm.weight")).transpose(0, 1) # [H, T, hd] + k = rms(k.reshape(T, HK, hd), vec(p + "self_attn.k_norm.weight")).transpose(0, 1) + q, k = bf(rot(q, *tabs[kind])), bf(rot(k, *tabs[kind])) + k = k.repeat_interleave(H // HK, 0) + v = v.reshape(T, HK, hd).transpose(0, 1).repeat_interleave(H // HK, 0) + s = ((q @ k.transpose(1, 2)) / math.sqrt(hd)).masked_fill(masks[kind], -math.inf) + e = torch.exp(s - s.max(-1, keepdim=True).values) + pv = bf(e) + if kind == "full_attention": + pv = torch.cat([pv[:, :P], e[:, P:]], 1) # full-attention decode keeps P unrounded + a = bf((pv @ v) / e.sum(-1, keepdim=True)).transpose(0, 1) # [T, H, hd] + a = bf(a * bf(torch.sigmoid(g))[..., None]) + o = bf(a.reshape(T, H * hd) @ W(p + "self_attn.o_proj").T) + h = bf(h + rms(o, vec(p + "post_operator_norm.weight"))) + x = rms(h, vec(p + "ffn_norm.weight")) + f = p + "feed_forward." + act = bf(bf(torch.nn.functional.silu(bf(x @ W(f + "w1").T))) * bf(x @ W(f + "w3").T)) + h = bf(h + rms(bf(act @ W(f + "w2").T), vec(p + "post_ffn_norm.weight"))) + x = rms(h, vec("embedding_norm.weight")) + V = cfg["vocab_size"] + return torch.cat([x @ self.lin("embed_tokens", rows=torch.arange(a, min(a + 16000, V))).T + for a in range(0, V, 16000)], -1) + + +def greedy(lg): + """argmax in (value desc, id asc) order, as the engine picks over its bf16 logits.""" + return int((lg == lg.max()).nonzero()[0]) + + +def main(model_dir, engine_path, with_bf16): + mdl = Model(model_dir) + try: + from tokenizers import Tokenizer + tok = Tokenizer.from_file(f"{model_dir}/tokenizer.json") + + def show(t): + return f"{t} {json.dumps(tok.decode([t], skip_special_tokens=False))}".replace("|", "\\|") + except ImportError: + show = str + cases = json.load(open(engine_path))["cases"] + tot = {(m, k): [0, 0] for m in ("f64", "bf16") for k in ("prompt", "reply")} + f64_rows, bf16_rows = [], [] + for c in cases: + P = len(c["prompt"]) + ids = c["prompt"] + c["reply"][:-1] + print(f"{c['name']}: {len(ids)} rows", file=sys.stderr, flush=True) + lf = mdl.forward(ids, P, same) + lb_pre = mdl.forward(ids, P, bf16) if with_bf16 else None + lb = bf16(lb_pre) if with_bf16 else None + rows = [("prompt", t, tk) for t, tk in enumerate(c["prompt_picks"])] + rows += [("reply", P - 1 + j, tk) for j, tk in enumerate(c["reply"])] + for kind, t, eng in rows: + top = int(lf[t].argmax()) + tot["f64", kind][0] += top == eng + tot["f64", kind][1] += 1 + bcol = "" + if with_bf16: + bp = greedy(lb[t]) + tot["bf16", kind][0] += bp == eng + tot["bf16", kind][1] += 1 + if top != eng: + if with_bf16: + d, dp = lb[t, top] - lb[t, eng], lb_pre[t, top] - lb_pre[t, eng] + side = "engine's" if bp == eng else "f64's" if bp == top else f"other {show(bp)}" + bcol = f" {d:+.4f} (pre-round {dp:+.4f}) | {side} |" + f64_rows.append(f"| {c['name']} | {kind} | {t} | {show(eng)} | {show(top)} | " + f"{lf[t, top] - lf[t, eng]:.4f} |{bcol}") + elif with_bf16 and bp != eng: + bf16_rows.append(f"| {c['name']} | {kind} | {t} | {show(eng)} | {show(bp)} | " + f"{lb[t, bp] - lb[t, eng]:.4f} (pre-round {lb_pre[t, bp] - lb_pre[t, eng]:+.4f}) | " + f"{lf[t, eng] - lf[t, bp]:.4f} |") + mode = [m for m in ("f64", "bf16") if m == "f64" or with_bf16] + for m in mode: + print(f"{m}: prompt rows {tot[m, 'prompt'][0]}/{tot[m, 'prompt'][1]}, " + f"reply rows {tot[m, 'reply'][0]}/{tot[m, 'reply'][1]}") + print("\nFloat64 top-1 misses (margin = f64 logit of its top minus f64 logit of the engine's token" + + ("; bf16 margin = the same difference in the bf16-points pass, < 0 means it prefers the engine's token" + if with_bf16 else "") + "):\n") + print("| case | row | t | engine | float64 | f64 margin |" + (" bf16 margin | bf16 picks |" if with_bf16 else "")) + print("|---|---|---|---|---|---|" + ("---|---|" if with_bf16 else "")) + print("\n".join(f64_rows) or "(none)") + if with_bf16: + print("\nbf16-points misses where float64 agrees with the engine (margin = bf16 logit of its pick minus the " + "engine's token; f64 margin = how far float64 prefers the engine's token):\n") + print("| case | row | t | engine | bf16 pick | bf16 margin | f64 margin |\n|---|---|---|---|---|---|---|") + print("\n".join(bf16_rows) or "(none)") + + +if __name__ == "__main__": + if len(sys.argv) == 5 and sys.argv[1] == "collect": + collect(*sys.argv[2:]) + elif len(sys.argv) in (3, 4) and sys.argv[3:] in ([], ["--bf16"]): + main(sys.argv[1], sys.argv[2], len(sys.argv) == 4) + else: + sys.exit(__doc__) diff --git a/tools/lfm3/test_inspect_checkpoint.py b/tools/lfm3/test_inspect_checkpoint.py new file mode 100644 index 0000000000..d647542ade --- /dev/null +++ b/tools/lfm3/test_inspect_checkpoint.py @@ -0,0 +1,142 @@ +"""Host checks for bounded checkpoint metadata inspection.""" + +import io +import json +import struct +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +import inspect_checkpoint as inspector + + +class Response(io.BytesIO): + def __init__(self, body, status, headers): + super().__init__(body) + self.status = status + self.headers = headers + self.read_calls = 0 + + def read(self, size=-1): + self.read_calls += 1 + return super().read(size) + + +class InspectorTests(unittest.TestCase): + def test_refuses_full_weight_response_before_reading(self): + response = Response(b'weight data', 200, {}) + with patch.object(inspector.urllib.request, 'urlopen', return_value=response): + with self.assertRaisesRegex(inspector.InspectionError, 'refused'): + inspector.Remote.fetch('https://example.org/model', 8, (0, 7)) + self.assertEqual(response.read_calls, 0) + + def test_rejects_wrong_range(self): + response = Response(b'12345678', 206, {'Content-Range': 'bytes 8-15/100'}) + with patch.object(inspector.urllib.request, 'urlopen', return_value=response): + with self.assertRaisesRegex(inspector.InspectionError, 'unexpected'): + inspector.Remote.fetch('https://example.org/model', 8, (0, 7)) + + def test_rejects_truncated_range(self): + response = Response(b'123', 206, {'Content-Range': 'bytes 0-7/100'}) + with patch.object(inspector.urllib.request, 'urlopen', return_value=response): + with self.assertRaisesRegex(inspector.InspectionError, 'Truncated'): + inspector.Remote.fetch('https://example.org/model', 8, (0, 7)) + + def test_accepts_exact_range(self): + response = Response(b'12345678', 206, {'Content-Range': 'bytes 0-7/100', 'Content-Length': '8'}) + with patch.object(inspector.urllib.request, 'urlopen', return_value=response): + self.assertEqual(inspector.Remote.fetch('https://example.org/model', 8, (0, 7)), b'12345678') + + def test_caps_non_range_metadata(self): + response = Response(b'123456789', 200, {}) + with patch.object(inspector.urllib.request, 'urlopen', return_value=response): + with self.assertRaisesRegex(inspector.InspectionError, 'body limit'): + inspector.Remote.fetch('https://example.org/config', 8) + + def test_rejects_oversize_header_and_unsafe_names(self): + with self.assertRaises(inspector.InspectionError): + inspector.header_size(struct.pack('(P, Wc, STATE, SLOT, OFF, Y, SNAP, SNAP0, D, SS); +} + +extern "C" __global__ void __launch_bounds__(256) tf_ctl_liquid_commit_3(__nv_bfloat16* STATE, const __nv_bfloat16* SNAP, const int* SLOT, + const int* ROW, int D, long long state_stride, long long snap_stride, + long long SS) { + tf_liquid::liquid_commit<3, 1>(STATE, SNAP, SLOT, ROW, D, state_stride, snap_stride, SS); +} + +// LFM3's 7-tap conv and commit, the same planted faults as the 3-tap pair. +extern "C" __global__ void __launch_bounds__(256) tf_ctl_liquid_conv_7(const __nv_bfloat16* P, const __nv_bfloat16* Wc, const __nv_bfloat16* STATE, + const int* SLOT, const int* OFF, __nv_bfloat16* Y, __nv_bfloat16* SNAP, + int SNAP0, int D, long long SS) { + tf_liquid::liquid_conv<7, 1>(P, Wc, STATE, SLOT, OFF, Y, SNAP, SNAP0, D, SS); +} + +extern "C" __global__ void __launch_bounds__(256) tf_ctl_liquid_commit_7(__nv_bfloat16* STATE, const __nv_bfloat16* SNAP, const int* SLOT, + const int* ROW, int D, long long state_stride, long long snap_stride, + long long SS) { + tf_liquid::liquid_commit<7, 1>(STATE, SNAP, SLOT, ROW, D, state_stride, snap_stride, SS); +} + +// LFM2.5-24B-A2B's router, one planted fault: ties go to the higher expert id. +#include "lfm2_moe.cuh" + +extern "C" __global__ void __launch_bounds__(256) tf_ctl_lfm2m_route(const __nv_bfloat16* X, const float* W, const float* BIAS, int* PICK, + float* WTS, float* LOGITS, int rows, int D, float scale) { + tf_lfm2m::lfm2m_route<1>(X, W, BIAS, PICK, WTS, LOGITS, rows, D, scale); +} + +// Paged attention at head 64, one planted fault each: values from the next lane's dims (decode) or next 8 (prompt). +namespace tf_ctl_paged { +using tf_glue::bf; +using tf_glue::tf_exp; +#include "paged_attn.cuh" +#define SEG_HD 64 +#define SEG_KW (64 / 2 + 2) +#define SEG_WIDE 0 +#define SEG_VLANE 1 +#define SEG_CHUNK tf_ctl_paged_attn_chunk_segs_64 +#define SEG_MERGE tf_ctl_paged_attn_merge_segs_64 +#include "paged_attn_segs.cuh" +} // namespace tf_ctl_paged + +#define PATTN_VLANE 1 +#include "prefill_attention.cuh" + +namespace tf_prefill_attention { +extern "C" __global__ void __launch_bounds__(256, 1) tf_ctl_pattn_64(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kc, + const __nv_bfloat16* __restrict__ vc, __nv_bfloat16* __restrict__ out, int p0, + int W, int H, int HK, int G, float scale, const int* __restrict__ PT, long long PS) { + constexpr int D = 64, WARPS = 8, HPC = 4, NS = 8; + constexpr bool causal = true; + constexpr int skip = 0; + constexpr bool HPAGED = false, F8 = false, HM = false; + constexpr int CAP = 0, PTOK = 1; + const float *ks = nullptr, *vs = nullptr; + const int HO = H; + auto first_row = [] { return blockIdx.x * 16 * (WARPS / HPC); }; +#include "prefill_attention_rows.cuh" +} +// The same fault at head dim 128 (LFM3's 16 query heads over 4). +extern "C" __global__ void __launch_bounds__(256, 1) tf_ctl_pattn_128(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kc, + const __nv_bfloat16* __restrict__ vc, __nv_bfloat16* __restrict__ out, int p0, + int W, int H, int HK, int G, float scale, const int* __restrict__ PT, long long PS) { + constexpr int D = 128, WARPS = 8, HPC = 4, NS = 8; + constexpr bool causal = true; + constexpr int skip = 0; + constexpr bool HPAGED = false, F8 = false, HM = false; + constexpr int CAP = 0, PTOK = 1; + const float *ks = nullptr, *vs = nullptr; + const int HO = H; + auto first_row = [] { return blockIdx.x * 16 * (WARPS / HPC); }; +#include "prefill_attention_rows.cuh" +} +} // namespace tf_prefill_attention diff --git a/zig/kernels/cuda/lfm2_moe.cu b/zig/kernels/cuda/lfm2_moe.cu new file mode 100644 index 0000000000..d7b3783388 --- /dev/null +++ b/zig/kernels/cuda/lfm2_moe.cu @@ -0,0 +1,26 @@ +// LFM2.5-24B-A2B's CUDA glue: Liquid's 3-tap short conv and commit (liquid_conv.cuh), the router (lfm2_moe.cuh). + +#include "liquid_conv.cuh" +#include "lfm2_moe.cuh" + +using namespace tf_lfm2m; + +extern "C" __global__ void __launch_bounds__(tf_liquid::NT) tf_liquid_conv_3(const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ Wc, + const __nv_bfloat16* __restrict__ STATE, const int* __restrict__ SLOT, + const int* __restrict__ OFF, __nv_bfloat16* __restrict__ Y, + __nv_bfloat16* __restrict__ SNAP, int SNAP0, int D, long long SS) { + tf_liquid::liquid_conv<3, 0>(P, Wc, STATE, SLOT, OFF, Y, SNAP, SNAP0, D, SS); +} + +extern "C" __global__ void __launch_bounds__(tf_liquid::NT) tf_liquid_commit_3(__nv_bfloat16* __restrict__ STATE, const __nv_bfloat16* __restrict__ SNAP, + const int* __restrict__ SLOT, const int* __restrict__ ROW, int D, + long long state_stride, long long snap_stride, long long SS) { + tf_liquid::liquid_commit<3, 0>(STATE, SNAP, SLOT, ROW, D, state_stride, snap_stride, SS); +} + +extern "C" __global__ void __launch_bounds__(NT) tf_lfm2m_route(const __nv_bfloat16* __restrict__ X, const float* __restrict__ W, + const float* __restrict__ BIAS, int* __restrict__ PICK, + float* __restrict__ WTS, float* __restrict__ LOGITS, int rows, + int D, float scale) { + lfm2m_route<0>(X, W, BIAS, PICK, WTS, LOGITS, rows, D, scale); +} diff --git a/zig/kernels/cuda/lfm2_moe.cuh b/zig/kernels/cuda/lfm2_moe.cuh new file mode 100644 index 0000000000..e6f390f05c --- /dev/null +++ b/zig/kernels/cuda/lfm2_moe.cuh @@ -0,0 +1,91 @@ +// LFM2.5-24B-A2B's router; each row's bits its own; CTL 1 is the planted control (short conv: liquid_conv.cuh). + +#pragma once + +#include +#include +#include + +#include "common/glue_math.cuh" + +namespace tf_lfm2m { + +constexpr int NT = 256; // threads a block +constexpr int E = 64; // experts +constexpr int TOPK = 4; // picks a row +constexpr int RB = 8; // router rows a block: the weights are read once for them +constexpr int DMAX = 2048; // hidden width the router stages + +// (value, id) wins over (other value, other id): the larger value, the lower id among equals. +template +__device__ __forceinline__ bool beats(float v, int i, float w, int j) { return v > w || (v == w && (CTL == 1 ? i > j : i < j)); } + + +// Block of RB rows: fp32 logits (quarters of D in order, (q0 + q1) + (q2 + q3)), top-4 of sigmoid + bias, weights. +template // CTL 1: ties go to the higher expert id. +__device__ __forceinline__ void lfm2m_route(const __nv_bfloat16* __restrict__ X, const float* __restrict__ W, + const float* __restrict__ BIAS, int* __restrict__ PICK, + float* __restrict__ WTS, float* __restrict__ LOGITS, int rows, + int D, float scale) { + __shared__ __align__(16) __nv_bfloat16 xs[RB][DMAX]; + __shared__ float lg[RB][E]; + const int row0 = blockIdx.x * RB, t = threadIdx.x, e = t / 4, q = t % 4, n = min(RB, rows - row0); + for (int i = t; i < n * D; i += NT) xs[i / D][i % D] = X[static_cast(row0 + i / D) * D + i % D]; + __syncthreads(); + float acc[RB]; +#pragma unroll + for (int b = 0; b < RB; ++b) acc[b] = 0.0f; + const int per = D / 4, k0 = q * per; + const float* w = W + static_cast(e) * D + k0; + for (int k = 0; k < per; ++k) { + const float wk = w[k]; +#pragma unroll + for (int b = 0; b < RB; ++b) acc[b] = __fmaf_rn(tf_glue::bf(xs[b][k0 + k]), wk, acc[b]); + } +#pragma unroll + for (int b = 0; b < RB; ++b) { + float v = __fadd_rn(acc[b], __shfl_xor_sync(0xffffffffu, acc[b], 1)); + v = __fadd_rn(v, __shfl_xor_sync(0xffffffffu, v, 2)); + if (q == 0) lg[b][e] = v; + } + __syncthreads(); + const int b = t / 32, lane = t % 32; + if (b >= n) return; + const int64_t r = row0 + b; + float gate[2], sel[2]; +#pragma unroll + for (int h = 0; h < 2; ++h) { + const float l = lg[b][lane + 32 * h]; + if (LOGITS != nullptr) LOGITS[r * E + lane + 32 * h] = l; + gate[h] = tf_glue::tf_sigmoid(l); + const float v = __fadd_rn(gate[h], BIAS[lane + 32 * h]); + sel[h] = v != v ? -INFINITY : v; // a NaN logit is never picked ahead of a number + } + float total = 0.0f, mine = 0.0f; + int pick = 0; + for (int k = 0; k < TOPK; ++k) { + float bv = -INFINITY; + int bi = CTL == 1 ? -1 : 0x7fffffff; +#pragma unroll + for (int h = 0; h < 2; ++h) + if (beats(sel[h], lane + 32 * h, bv, bi)) bv = sel[h], bi = lane + 32 * h; +#pragma unroll + for (int off = 16; off > 0; off >>= 1) { + const float ov = __shfl_xor_sync(0xffffffffu, bv, off); + const int oi = __shfl_xor_sync(0xffffffffu, bi, off); + if (beats(ov, oi, bv, bi)) bv = ov, bi = oi; + } + float p = 0.0f; +#pragma unroll + for (int h = 0; h < 2; ++h) + if (bi == lane + 32 * h) p = gate[h], sel[h] = NAN; // taken: NaN beats nothing, not even a -inf left + p = __shfl_sync(0xffffffffu, p, bi % 32); + total = __fadd_rn(total, p); + if (lane == k) mine = p, pick = bi; + } + if (lane >= TOPK) return; + PICK[r * TOPK + lane] = pick; + WTS[r * TOPK + lane] = __fmul_rn(__fdiv_rn(mine, __fadd_rn(total, 1e-6f)), scale); +} + +} // namespace tf_lfm2m diff --git a/zig/kernels/cuda/lfm3.cu b/zig/kernels/cuda/lfm3.cu new file mode 100644 index 0000000000..d0e9f8934a --- /dev/null +++ b/zig/kernels/cuda/lfm3.cu @@ -0,0 +1,16 @@ +// LFM3's CUDA glue: Liquid's 7-tap short conv and its commit (liquid_conv.cuh); everything else is shared. + +#include "liquid_conv.cuh" + +extern "C" __global__ void __launch_bounds__(tf_liquid::NT) tf_liquid_conv_7(const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ Wc, + const __nv_bfloat16* __restrict__ STATE, const int* __restrict__ SLOT, + const int* __restrict__ OFF, __nv_bfloat16* __restrict__ Y, + __nv_bfloat16* __restrict__ SNAP, int SNAP0, int D, long long SS) { + tf_liquid::liquid_conv<7, 0>(P, Wc, STATE, SLOT, OFF, Y, SNAP, SNAP0, D, SS); +} + +extern "C" __global__ void __launch_bounds__(tf_liquid::NT) tf_liquid_commit_7(__nv_bfloat16* __restrict__ STATE, const __nv_bfloat16* __restrict__ SNAP, + const int* __restrict__ SLOT, const int* __restrict__ ROW, int D, + long long state_stride, long long snap_stride, long long SS) { + tf_liquid::liquid_commit<7, 0>(STATE, SNAP, SLOT, ROW, D, state_stride, snap_stride, SS); +} diff --git a/zig/kernels/cuda/liquid_conv.cuh b/zig/kernels/cuda/liquid_conv.cuh new file mode 100644 index 0000000000..51d1ec0e7e --- /dev/null +++ b/zig/kernels/cuda/liquid_conv.cuh @@ -0,0 +1,65 @@ +// Liquid's gated short conv over stream state and its commit, by tap count (LFM2.5 3, LFM3 7); CTL 1: control. + +#pragma once + +#include +#include + +#include "common/glue_math.cuh" + +namespace tf_liquid { + +constexpr int NT = 256; // threads a block + +// Block (row, channel block): y = bf16(C * bf16(sum of w * bf16(B * x), oldest tap first)); SNAP rows keep state. +template // CTL 1: bx = B * x not rounded to bf16 before its tap. +__device__ __forceinline__ void liquid_conv(const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ Wc, + const __nv_bfloat16* __restrict__ STATE, const int* __restrict__ SLOT, + const int* __restrict__ OFF, __nv_bfloat16* __restrict__ Y, + __nv_bfloat16* __restrict__ SNAP, int SNAP0, int D, long long SS) { + constexpr int S = TAPS - 1; // state rows: the earlier taps come from the round (OFF before the row), else the slot + const int64_t r = blockIdx.x, D3 = 3 * static_cast(D); + const int c = blockIdx.y * NT + threadIdx.x; + if (c >= D) return; + const int o = OFF[r]; + const __nv_bfloat16* st = STATE + SLOT[r] * SS; + __nv_bfloat16 bx[TAPS]; + float wide[TAPS]; +#pragma unroll + for (int j = 0; j < TAPS; ++j) wide[j] = 0.0f; +#pragma unroll + for (int j = 0; j < TAPS; ++j) { + const int back = S - j; // rows before this one + if (o >= back) { + const int64_t src = r - back; + const float v = __fmul_rn(tf_glue::bf(P[src * D3 + c]), tf_glue::bf(P[src * D3 + 2 * D + c])); + bx[j] = __float2bfloat16_rn(v); + if (CTL == 1) wide[j] = v; + } else { + bx[j] = st[static_cast(S - back + o) * D + c]; + if (CTL == 1) wide[j] = tf_glue::bf(bx[j]); + } + } + float acc = 0.0f; +#pragma unroll + for (int j = 0; j < TAPS; ++j) acc = __fadd_rn(acc, __fmul_rn(tf_glue::bf(Wc[c * TAPS + j]), CTL == 1 ? wide[j] : tf_glue::bf(bx[j]))); + Y[r * D + c] = __float2bfloat16_rn(__fmul_rn(tf_glue::bf(P[r * D3 + D + c]), tf_glue::rbf(acc))); + if (SNAP == nullptr || r < SNAP0) return; +#pragma unroll + for (int j = 1; j < TAPS; ++j) SNAP[((r - SNAP0) * S + (j - 1)) * D + c] = bx[j]; +} + +// Block (layer, stream): stream k's [TAPS - 1][D] conv state set to its snapshot row ROW[k]; ROW < 0 keeps the state. +template // CTL 1: every stream's state written to the first stream's slot. +__device__ __forceinline__ void liquid_commit(__nv_bfloat16* __restrict__ STATE, const __nv_bfloat16* __restrict__ SNAP, + const int* __restrict__ SLOT, const int* __restrict__ ROW, int D, + long long state_stride, long long snap_stride, long long SS) { + constexpr int S = TAPS - 1; // state rows + const int layer = blockIdx.x, k = blockIdx.y, row = ROW[k]; + if (row < 0) return; + __nv_bfloat16* dst = STATE + layer * state_stride + SLOT[CTL == 1 ? 0 : k] * SS; + const __nv_bfloat16* src = SNAP + layer * snap_stride + static_cast(row) * S * D; + for (int i = threadIdx.x; i < S * D; i += NT) dst[i] = src[i]; +} + +} // namespace tf_liquid diff --git a/zig/kernels/cuda/nemotron_attention.cu b/zig/kernels/cuda/nemotron_attention.cu index 313b096f28..6c66d56704 100644 --- a/zig/kernels/cuda/nemotron_attention.cu +++ b/zig/kernels/cuda/nemotron_attention.cu @@ -5,45 +5,15 @@ #include #include "common/glue_math.cuh" +#include "paged_attn.cuh" using namespace tf_glue; namespace { constexpr int HD = 128; // head dim the chunk kernels take -constexpr int GMAX = 16; // query heads a KV head -constexpr int TILE = 64; // keys a tile: the online softmax's unit constexpr int KW = HD / 2 + 2; // a key row's words in shared memory, padded so 64-bit reads of 16 rows miss no bank -__device__ __forceinline__ float lo(uint32_t w) { return __uint_as_float(w << 16); } -__device__ __forceinline__ float hi(uint32_t w) { return __uint_as_float(w & 0xFFFF0000u); } - -// Max over the 16 lanes of a half warp (xor 8, 4, 2, 1). -__device__ __forceinline__ float half_max(float v) { -#pragma unroll - for (int off = 8; off > 0; off >>= 1) v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, off)); - return v; -} - -// Sum over the 16 lanes of a half warp (xor 8, 4, 2, 1): each lane ends with the same sum. -__device__ __forceinline__ float half_sum(float v) { -#pragma unroll - for (int off = 8; off > 0; off >>= 1) v = __fadd_rn(v, __shfl_xor_sync(0xffffffffu, v, off)); - return v; -} - -constexpr int PAGE = 512; // rows a KV page holds: one attention chunk - -// Element offset of chunk c's first key row: contiguous rows, or the page the stream's table names (PT null: rows). -__device__ __forceinline__ int64_t chunk_at(const int* PT, int64_t PS, int c, int KVD) { - return PT ? static_cast(PT[c]) * PS : static_cast(c) * PAGE * KVD; -} - -// Element offset of key row k: its chunk's first row, then its row inside the chunk. -__device__ __forceinline__ int64_t key_at(const int* PT, int64_t PS, int64_t k, int KVD) { - return chunk_at(PT, PS, static_cast(k / PAGE), KVD) + (k % PAGE) * KVD; -} - } // namespace // Row r's keys and values written at its position: K and V follow the query columns of each qkv row. @@ -197,23 +167,6 @@ extern "C" __global__ void __launch_bounds__(256) tf_nemo_attn_merge(const float } // Shared-window copies of the kernels above (a shared body slowed a lone stream): stream k's rows from ROW0[k] on. -constexpr int MAXSEG = 8; -struct AttnSegs { - __nv_bfloat16* KC[MAXSEG]; - __nv_bfloat16* VC[MAXSEG]; - const int* META[MAXSEG]; - const int* PT[MAXSEG]; // each stream's page table (null: its rows are contiguous) - long long PS; // elements between pages - int ROW0[MAXSEG]; - int N; -}; - -__device__ __forceinline__ int seg_of(const AttnSegs& sg, int rq) { - int k = 0; - while (k + 1 < sg.N && sg.ROW0[k + 1] <= rq) ++k; - return k; -} - extern "C" __global__ void tf_nemo_kv_write_segs(const AttnSegs sg, const __nv_bfloat16* __restrict__ QKV, int NQKV, int QD, int KVD) { const int rq = blockIdx.x, k = seg_of(sg, rq); __nv_bfloat16* __restrict__ KC = sg.KC[k]; @@ -225,144 +178,19 @@ extern "C" __global__ void tf_nemo_kv_write_segs(const AttnSegs sg, const __nv_b } } -extern "C" __global__ void __launch_bounds__(256) tf_nemo_attn_chunk_segs(const AttnSegs sg, const __nv_bfloat16* __restrict__ QKV, - float* __restrict__ PO, float* __restrict__ PM, float* __restrict__ PL, - int NQKV, int H, int HK, int G, int CH, int NCH, float scale) { - __shared__ __align__(16) float qs[GMAX][HD]; - __shared__ __align__(16) uint32_t kw[TILE][KW]; - __shared__ __align__(16) uint32_t vw[TILE][HD / 2]; - __shared__ float ps[GMAX][TILE]; - const int rq = blockIdx.x, k = seg_of(sg, rq), r = rq - sg.ROW0[k], hk = blockIdx.y, t = threadIdx.x; - const __nv_bfloat16* __restrict__ KC = sg.KC[k]; - const __nv_bfloat16* __restrict__ VC = sg.VC[k]; - const int* __restrict__ META = sg.META[k]; - const int limit = META[0] + r + 1; - if (static_cast(blockIdx.z) * CH >= limit) return; - const int g = t / 16, jj = t % 16; // softmax and values: head g, lanes jj of its 16 - const int key = t % TILE, quad = t / TILE; // scores: one key, heads 4 quad .. 4 quad + 3 - for (int i = t; i < GMAX * HD; i += 256) { - const int gg = i / HD, d = i % HD; - qs[gg][d] = gg < G ? bf(QKV[static_cast(rq) * NQKV + (hk * G + gg) * HD + d]) : 0.0f; - } - for (int c = blockIdx.z; c * CH < limit; c += gridDim.z) { // this block's chunks: every gridDim.z-th - float m = -INFINITY, den = 0.0f, o[8]; -#pragma unroll - for (int e = 0; e < 8; ++e) o[e] = 0.0f; - const int64_t cb = chunk_at(sg.PT[k], sg.PS, c, HK * HD) - static_cast(c) * CH * HK * HD; - for (int key0 = c * CH; key0 < (c + 1) * CH && key0 < limit; key0 += TILE) { - __syncthreads(); - for (int i = t; i < TILE * (HD / 2); i += 256) { - const int j = i / (HD / 2), w = i % (HD / 2); - const int64_t at = cb + (static_cast(key0 + j) * HK + hk) * HD; - const bool in = key0 + j < limit; - kw[j][w] = in ? reinterpret_cast(KC + at)[w] : 0u; - vw[j][w] = in ? reinterpret_cast(VC + at)[w] : 0u; - } - __syncthreads(); - float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; -#pragma unroll 4 - for (int w = 0; w < HD / 2; w += 2) { - const uint2 kk = *reinterpret_cast(&kw[key][w]); - const float k0 = lo(kk.x), k1 = hi(kk.x), k2 = lo(kk.y), k3 = hi(kk.y); -#pragma unroll - for (int h = 0; h < 4; ++h) { - const float4 q = *reinterpret_cast(&qs[4 * quad + h][2 * w]); - acc[h] = __fmaf_rn(q.x, k0, acc[h]); - acc[h] = __fmaf_rn(q.y, k1, acc[h]); - acc[h] = __fmaf_rn(q.z, k2, acc[h]); - acc[h] = __fmaf_rn(q.w, k3, acc[h]); - } - } -#pragma unroll - for (int h = 0; h < 4; ++h) ps[4 * quad + h][key] = key0 + key < limit ? __fmul_rn(acc[h], scale) : -INFINITY; - __syncthreads(); - float sc[4]; -#pragma unroll - for (int u = 0; u < 4; ++u) sc[u] = ps[g][jj + 16 * u]; - const float tile_m = half_max(fmaxf(fmaxf(sc[0], sc[1]), fmaxf(sc[2], sc[3]))); - const bool active = tile_m != -INFINITY; - const float next = active ? fmaxf(m, tile_m) : m; - const float alpha = active ? (m == -INFINITY ? 0.0f : tf_exp(__fsub_rn(m, next))) : 1.0f; - float local = 0.0f; -#pragma unroll - for (int u = 0; u < 4; ++u) { - const float p = active && key0 + jj + 16 * u < limit ? tf_exp(__fsub_rn(sc[u], next)) : 0.0f; - ps[g][jj + 16 * u] = p; - local = __fadd_rn(local, p); - } - den = __fmaf_rn(den, alpha, half_sum(local)); - m = next; - __syncthreads(); - float pv[8]; -#pragma unroll - for (int e = 0; e < 8; ++e) pv[e] = 0.0f; - for (int j = 0; j < TILE; ++j) { - const float p = ps[g][j]; - const uint4 v = *reinterpret_cast(&vw[j][4 * jj]); - const uint32_t vv[4] = {v.x, v.y, v.z, v.w}; -#pragma unroll - for (int e = 0; e < 4; ++e) { - pv[2 * e] = __fmaf_rn(p, lo(vv[e]), pv[2 * e]); - pv[2 * e + 1] = __fmaf_rn(p, hi(vv[e]), pv[2 * e + 1]); - } - } -#pragma unroll - for (int e = 0; e < 8; ++e) o[e] = __fmaf_rn(o[e], alpha, pv[e]); - } - if (g >= G) continue; - const int64_t base = (static_cast(rq) * NCH + c) * H + hk * G + g; - float4* out = reinterpret_cast(PO + base * HD + 8 * jj); - out[0] = make_float4(o[0], o[1], o[2], o[3]); - out[1] = make_float4(o[4], o[5], o[6], o[7]); - if (jj == 0) { - PM[base] = m; - PL[base] = den; - } - } -} - - -extern "C" __global__ void __launch_bounds__(256) tf_nemo_attn_merge_segs(const AttnSegs sg, const float* __restrict__ PO, const float* __restrict__ PM, - const float* __restrict__ PL, __nv_bfloat16* __restrict__ OUT, float* __restrict__ XS, - int H, int HK, int G, int CH, int NCH) { - __shared__ float ys[GMAX][HD]; - const int rq = blockIdx.x, k = seg_of(sg, rq), r = rq - sg.ROW0[k], hk = blockIdx.y, t = threadIdx.x, g = t / 16, jj = t % 16; - const int* __restrict__ META = sg.META[k]; - const int limit = META[0] + r + 1, nch = (limit + CH - 1) / CH, head = hk * G + g; - float m = -INFINITY, den = 0.0f, o[8]; -#pragma unroll - for (int e = 0; e < 8; ++e) o[e] = 0.0f; - if (g < G) { - for (int c = 0; c < nch; ++c) { - const int64_t base = (static_cast(rq) * NCH + c) * H + head; - const float cm = PM[base], cl = PL[base]; - const bool active = cl > 0.0f; - const float next = active ? fmaxf(m, cm) : m; - const float a = active ? (m == -INFINITY ? 0.0f : tf_exp(__fsub_rn(m, next))) : 1.0f; - const float b = active ? tf_exp(__fsub_rn(cm, next)) : 0.0f; - const float4* po = reinterpret_cast(PO + base * HD + 8 * jj); - const float4 p0 = po[0], p1 = po[1]; - const float co[8] = {p0.x, p0.y, p0.z, p0.w, p1.x, p1.y, p1.z, p1.w}; -#pragma unroll - for (int e = 0; e < 8; ++e) o[e] = __fmaf_rn(o[e], a, __fmul_rn(co[e], b)); - den = __fmaf_rn(den, a, __fmul_rn(cl, b)); - m = next; - } -#pragma unroll - for (int e = 0; e < 8; ++e) { - const __nv_bfloat16 y = __float2bfloat16_rn(__fdiv_rn(o[e], den)); - OUT[(static_cast(rq) * H + head) * HD + 8 * jj + e] = y; - ys[g][8 * jj + e] = bf(y); - } - } - __syncthreads(); - if (t >= GMAX * (HD / 64)) return; - const int gg = t / (HD / 64), gi = t % (HD / 64); - if (gg >= G) return; - float s = 0.0f; - for (int j = 0; j < 64; ++j) s = __fadd_rn(s, ys[gg][gi * 64 + j]); - XS[static_cast(rq) * (H * HD / 64) + (hk * G + gg) * (HD / 64) + gi] = s; -} +#define SEG_HD HD +#define SEG_KW KW +#define SEG_WIDE 1 +#define SEG_CHUNK tf_nemo_attn_chunk_segs +#define SEG_MERGE tf_nemo_attn_merge_segs +#include "paged_attn_segs.cuh" +// The head dim 64 forms (LFM2.5-24B-A2B: 32 query heads over 8 KV heads), four dims a lane. +#define SEG_HD 64 +#define SEG_KW (64 / 2 + 2) +#define SEG_WIDE 0 +#define SEG_CHUNK tf_paged_attn_chunk_segs_64 +#define SEG_MERGE tf_paged_attn_merge_segs_64 +#include "paged_attn_segs.cuh" // tf_nemo_attn_chunk with the next key tile loading into registers while this one computes (same values, same order). extern "C" __global__ void __launch_bounds__(256, 2) tf_nemo_attn_chunk_pf(const __nv_bfloat16* __restrict__ QKV, const __nv_bfloat16* __restrict__ KC, diff --git a/zig/kernels/cuda/paged_attn.cuh b/zig/kernels/cuda/paged_attn.cuh new file mode 100644 index 0000000000..f5bba2e249 --- /dev/null +++ b/zig/kernels/cuda/paged_attn.cuh @@ -0,0 +1,62 @@ +// Paged attention's shared pieces (Nemotron's chunk/merge kernels, paged_attn_segs.cuh): tiles, pages, stream segments. +#pragma once + +#include +#include +#include + +#include "common/glue_math.cuh" + +namespace { + +constexpr int GMAX = 16; // query heads a KV head +constexpr int TILE = 64; // keys a tile: the online softmax's unit + +__device__ __forceinline__ float lo(uint32_t w) { return __uint_as_float(w << 16); } +__device__ __forceinline__ float hi(uint32_t w) { return __uint_as_float(w & 0xFFFF0000u); } + +// Max over the 16 lanes of a half warp (xor 8, 4, 2, 1). +__device__ __forceinline__ float half_max(float v) { +#pragma unroll + for (int off = 8; off > 0; off >>= 1) v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, off)); + return v; +} + +// Sum over the 16 lanes of a half warp (xor 8, 4, 2, 1): each lane ends with the same sum. +__device__ __forceinline__ float half_sum(float v) { +#pragma unroll + for (int off = 8; off > 0; off >>= 1) v = __fadd_rn(v, __shfl_xor_sync(0xffffffffu, v, off)); + return v; +} + +constexpr int PAGE = 512; // rows a KV page holds: one attention chunk + +// Element offset of chunk c's first key row: contiguous rows, or the page the stream's table names (PT null: rows). +__device__ __forceinline__ int64_t chunk_at(const int* PT, int64_t PS, int c, int KVD) { + return PT ? static_cast(PT[c]) * PS : static_cast(c) * PAGE * KVD; +} + +// Element offset of key row k: its chunk's first row, then its row inside the chunk. +__device__ __forceinline__ int64_t key_at(const int* PT, int64_t PS, int64_t k, int KVD) { + return chunk_at(PT, PS, static_cast(k / PAGE), KVD) + (k % PAGE) * KVD; +} + +} // namespace + +// Stream k's rows from ROW0[k] on, for the shared-window kernels. +constexpr int MAXSEG = 8; +struct AttnSegs { + __nv_bfloat16* KC[MAXSEG]; + __nv_bfloat16* VC[MAXSEG]; + const int* META[MAXSEG]; + const int* PT[MAXSEG]; // each stream's page table (null: its rows are contiguous) + long long PS; // elements between pages + int ROW0[MAXSEG]; + int N; +}; + +__device__ __forceinline__ int seg_of(const AttnSegs& sg, int rq) { + int k = 0; + while (k + 1 < sg.N && sg.ROW0[k + 1] <= rq) ++k; + return k; +} diff --git a/zig/kernels/cuda/paged_attn_segs.cuh b/zig/kernels/cuda/paged_attn_segs.cuh new file mode 100644 index 0000000000..5a18303cb4 --- /dev/null +++ b/zig/kernels/cuda/paged_attn_segs.cuh @@ -0,0 +1,167 @@ +// Paged shared-window chunk/merge bodies shared by families (Nemotron head 128, LFM2 head 64); SEG_* set the form. + +#define SEG_E (SEG_WIDE ? 8 : 4) +#ifndef SEG_VLANE +#define SEG_VLANE 0 // 1 in contract_controls.cu only: the narrow form reads its values from the neighbouring lane's dims +#endif + +extern "C" __global__ void __launch_bounds__(256) SEG_CHUNK(const AttnSegs sg, const __nv_bfloat16* __restrict__ QKV, + float* __restrict__ PO, float* __restrict__ PM, float* __restrict__ PL, + int NQKV, int H, int HK, int G, int CH, int NCH, float scale) { + __shared__ __align__(16) float qs[GMAX][SEG_HD]; + __shared__ __align__(16) uint32_t kw[TILE][SEG_KW]; + __shared__ __align__(16) uint32_t vw[TILE][SEG_HD / 2]; + __shared__ float ps[GMAX][TILE]; + const int rq = blockIdx.x, k = seg_of(sg, rq), r = rq - sg.ROW0[k], hk = blockIdx.y, t = threadIdx.x; + const __nv_bfloat16* __restrict__ KC = sg.KC[k]; + const __nv_bfloat16* __restrict__ VC = sg.VC[k]; + const int* __restrict__ META = sg.META[k]; + const int limit = META[0] + r + 1; + if (static_cast(blockIdx.z) * CH >= limit) return; + const int g = t / 16, jj = t % 16; // softmax and values: head g, lanes jj of its 16 + const int key = t % TILE, quad = t / TILE; // scores: one key, heads 4 quad .. 4 quad + 3 + for (int i = t; i < GMAX * SEG_HD; i += 256) { + const int gg = i / SEG_HD, d = i % SEG_HD; + qs[gg][d] = gg < G ? bf(QKV[static_cast(rq) * NQKV + (hk * G + gg) * SEG_HD + d]) : 0.0f; + } + for (int c = blockIdx.z; c * CH < limit; c += gridDim.z) { // this block's chunks: every gridDim.z-th + float m = -INFINITY, den = 0.0f, o[SEG_E]; +#pragma unroll + for (int e = 0; e < SEG_E; ++e) o[e] = 0.0f; + const int64_t cb = chunk_at(sg.PT[k], sg.PS, c, HK * SEG_HD) - static_cast(c) * CH * HK * SEG_HD; + for (int key0 = c * CH; key0 < (c + 1) * CH && key0 < limit; key0 += TILE) { + __syncthreads(); + for (int i = t; i < TILE * (SEG_HD / 2); i += 256) { + const int j = i / (SEG_HD / 2), w = i % (SEG_HD / 2); + const int64_t at = cb + (static_cast(key0 + j) * HK + hk) * SEG_HD; + const bool in = key0 + j < limit; + kw[j][w] = in ? reinterpret_cast(KC + at)[w] : 0u; + vw[j][w] = in ? reinterpret_cast(VC + at)[w] : 0u; + } + __syncthreads(); + float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; +#pragma unroll 4 + for (int w = 0; w < SEG_HD / 2; w += 2) { + const uint2 kk = *reinterpret_cast(&kw[key][w]); + const float k0 = lo(kk.x), k1 = hi(kk.x), k2 = lo(kk.y), k3 = hi(kk.y); +#pragma unroll + for (int h = 0; h < 4; ++h) { + const float4 q = *reinterpret_cast(&qs[4 * quad + h][2 * w]); + acc[h] = __fmaf_rn(q.x, k0, acc[h]); + acc[h] = __fmaf_rn(q.y, k1, acc[h]); + acc[h] = __fmaf_rn(q.z, k2, acc[h]); + acc[h] = __fmaf_rn(q.w, k3, acc[h]); + } + } +#pragma unroll + for (int h = 0; h < 4; ++h) ps[4 * quad + h][key] = key0 + key < limit ? __fmul_rn(acc[h], scale) : -INFINITY; + __syncthreads(); + float sc[4]; +#pragma unroll + for (int u = 0; u < 4; ++u) sc[u] = ps[g][jj + 16 * u]; + const float tile_m = half_max(fmaxf(fmaxf(sc[0], sc[1]), fmaxf(sc[2], sc[3]))); + const bool active = tile_m != -INFINITY; + const float next = active ? fmaxf(m, tile_m) : m; + const float alpha = active ? (m == -INFINITY ? 0.0f : tf_exp(__fsub_rn(m, next))) : 1.0f; + float local = 0.0f; +#pragma unroll + for (int u = 0; u < 4; ++u) { + const float p = active && key0 + jj + 16 * u < limit ? tf_exp(__fsub_rn(sc[u], next)) : 0.0f; + ps[g][jj + 16 * u] = p; + local = __fadd_rn(local, p); + } + den = __fmaf_rn(den, alpha, half_sum(local)); + m = next; + __syncthreads(); + float pv[SEG_E]; +#pragma unroll + for (int e = 0; e < SEG_E; ++e) pv[e] = 0.0f; + for (int j = 0; j < TILE; ++j) { + const float p = ps[g][j]; +#if SEG_WIDE + const uint4 v = *reinterpret_cast(&vw[j][4 * jj]); + const uint32_t vv[4] = {v.x, v.y, v.z, v.w}; +#pragma unroll + for (int e = 0; e < 4; ++e) { +#else + const uint2 v = *reinterpret_cast(&vw[j][2 * (jj ^ SEG_VLANE)]); + const uint32_t vv[2] = {v.x, v.y}; +#pragma unroll + for (int e = 0; e < 2; ++e) { +#endif + pv[2 * e] = __fmaf_rn(p, lo(vv[e]), pv[2 * e]); + pv[2 * e + 1] = __fmaf_rn(p, hi(vv[e]), pv[2 * e + 1]); + } + } +#pragma unroll + for (int e = 0; e < SEG_E; ++e) o[e] = __fmaf_rn(o[e], alpha, pv[e]); + } + if (g >= G) continue; + const int64_t base = (static_cast(rq) * NCH + c) * H + hk * G + g; + float4* out = reinterpret_cast(PO + base * SEG_HD + SEG_E * jj); + out[0] = make_float4(o[0], o[1], o[2], o[3]); +#if SEG_WIDE + out[1] = make_float4(o[4], o[5], o[6], o[7]); +#endif + if (jj == 0) { + PM[base] = m; + PL[base] = den; + } + } +} + + +extern "C" __global__ void __launch_bounds__(256) SEG_MERGE(const AttnSegs sg, const float* __restrict__ PO, const float* __restrict__ PM, + const float* __restrict__ PL, __nv_bfloat16* __restrict__ OUT, float* __restrict__ XS, + int H, int HK, int G, int CH, int NCH) { + __shared__ float ys[GMAX][SEG_HD]; + const int rq = blockIdx.x, k = seg_of(sg, rq), r = rq - sg.ROW0[k], hk = blockIdx.y, t = threadIdx.x, g = t / 16, jj = t % 16; + const int* __restrict__ META = sg.META[k]; + const int limit = META[0] + r + 1, nch = (limit + CH - 1) / CH, head = hk * G + g; + float m = -INFINITY, den = 0.0f, o[SEG_E]; +#pragma unroll + for (int e = 0; e < SEG_E; ++e) o[e] = 0.0f; + if (g < G) { + for (int c = 0; c < nch; ++c) { + const int64_t base = (static_cast(rq) * NCH + c) * H + head; + const float cm = PM[base], cl = PL[base]; + const bool active = cl > 0.0f; + const float next = active ? fmaxf(m, cm) : m; + const float a = active ? (m == -INFINITY ? 0.0f : tf_exp(__fsub_rn(m, next))) : 1.0f; + const float b = active ? tf_exp(__fsub_rn(cm, next)) : 0.0f; + const float4* po = reinterpret_cast(PO + base * SEG_HD + SEG_E * jj); +#if SEG_WIDE + const float4 p0 = po[0], p1 = po[1]; + const float co[8] = {p0.x, p0.y, p0.z, p0.w, p1.x, p1.y, p1.z, p1.w}; +#else + const float4 p0 = po[0]; + const float co[4] = {p0.x, p0.y, p0.z, p0.w}; +#endif +#pragma unroll + for (int e = 0; e < SEG_E; ++e) o[e] = __fmaf_rn(o[e], a, __fmul_rn(co[e], b)); + den = __fmaf_rn(den, a, __fmul_rn(cl, b)); + m = next; + } +#pragma unroll + for (int e = 0; e < SEG_E; ++e) { + const __nv_bfloat16 y = __float2bfloat16_rn(__fdiv_rn(o[e], den)); + OUT[(static_cast(rq) * H + head) * SEG_HD + SEG_E * jj + e] = y; + ys[g][SEG_E * jj + e] = bf(y); + } + } + __syncthreads(); + if (t >= GMAX * (SEG_HD / 64)) return; + const int gg = t / (SEG_HD / 64), gi = t % (SEG_HD / 64); + if (gg >= G) return; + float s = 0.0f; + for (int j = 0; j < 64; ++j) s = __fadd_rn(s, ys[gg][gi * 64 + j]); + XS[static_cast(rq) * (H * SEG_HD / 64) + (hk * G + gg) * (SEG_HD / 64) + gi] = s; +} + +#undef SEG_E +#undef SEG_VLANE +#undef SEG_HD +#undef SEG_KW +#undef SEG_WIDE +#undef SEG_CHUNK +#undef SEG_MERGE diff --git a/zig/kernels/cuda/prefill_attention.cu b/zig/kernels/cuda/prefill_attention.cu index 7fffbb9307..0bfe2eecb8 100644 --- a/zig/kernels/cuda/prefill_attention.cu +++ b/zig/kernels/cuda/prefill_attention.cu @@ -1,63 +1,8 @@ // Device code of src/tensorfold/cuda/kernels/prefill_attention.cu (comments, ATen includes dropped), paged keys added. -#include -#include -#include -#include -#include - -#include "common/qmm_frag.cuh" +#include "prefill_attention.cuh" namespace tf_prefill_attention { -using qmm_frag::commit; -using qmm_frag::cp16z; -using qmm_frag::ldmatrix4; -using qmm_frag::mma; -using qmm_frag::mma0; - -constexpr int BN = 64, HALF = 32; // keys a tile (the fold's unit), keys a staging slot -constexpr float LOG2E = 1.4426950408889634f; -constexpr int SMEM_MAX = 96 * 1024; - -template -__device__ __forceinline__ void wait_group() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); } - -__device__ __forceinline__ void ldmatrix4t(uint32_t (&r)[4], const void* p) { - asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n" - : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) : "r"(qmm_frag::smem(p))); -} - -__device__ __forceinline__ float texp(float x) { - float y; - asm("ex2.approx.f32 %0, %1;\n" : "=f"(y) : "f"(__fmul_rn(x, LOG2E))); - return y; -} - -__device__ __forceinline__ float tdiv(float a, float b) { - float y; - asm("div.full.f32 %0, %1, %2;\n" : "=f"(y) : "f"(a), "f"(b)); - return y; -} - -__device__ __forceinline__ uint32_t pack(float lo, float hi) { - __nv_bfloat162 v = __floats2bfloat162_rn(lo, hi); - return *reinterpret_cast(&v); -} - -template -__device__ __forceinline__ int sw(int r, int c) { return r * (D / 8) + (c ^ (r & 7)); } - -__device__ __forceinline__ float rowsum(const float (&c)[8]) { - float w[8]; -#pragma unroll - for (int j = 0; j < 8; ++j) { - const float x = __fadd_rn(c[j], __shfl_xor_sync(0xffffffffu, c[j], 2)); - w[j] = __fadd_rn(x, __shfl_xor_sync(0xffffffffu, x, 1)); - } - return __fadd_rn(__fadd_rn(__fadd_rn(w[0], w[4]), __fadd_rn(w[2], w[6])), - __fadd_rn(__fadd_rn(w[1], w[5]), __fadd_rn(w[3], w[7]))); -} - template __global__ void __launch_bounds__(32 * WARPS, 1) pattn_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kc, @@ -324,6 +269,14 @@ vision_kernel(const __nv_bfloat16* __restrict__ qkv, __nv_bfloat16* __restrict__ template __global__ void tf_prefill_attention::pattn_kernel<128, 8, 8, 8>(const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, __nv_bfloat16* __restrict__, int, int, int, int, int, float, const int* __restrict__, long long); +// Head dim 128, eight warps, four query heads a block, eight staging slots, over paged caches (LFM3: 4 a KV head). +template __global__ void tf_prefill_attention::pattn_kernel<128, 8, 4, 8>(const __nv_bfloat16* __restrict__, + const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, __nv_bfloat16* __restrict__, int, int, int, + int, int, float, const int* __restrict__, long long); +// Head dim 64, eight warps, four query heads a block, eight staging slots, over paged caches (LFM2.5-24B-A2B's chunks). +template __global__ void tf_prefill_attention::pattn_kernel<64, 8, 4, 8>(const __nv_bfloat16* __restrict__, + const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, __nv_bfloat16* __restrict__, int, int, int, + int, int, float, const int* __restrict__, long long); // Head dim 64, eight warps, four query heads a block, eight staging slots (LFM2's 4 heads a KV head, causal or trees). template __global__ void tf_prefill_attention::tree_kernel<64, 8, 4, 8>(const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, __nv_bfloat16* __restrict__, const int* __restrict__, int, int, int, int, int, float); diff --git a/zig/kernels/cuda/prefill_attention.cuh b/zig/kernels/cuda/prefill_attention.cuh new file mode 100644 index 0000000000..2bebef02aa --- /dev/null +++ b/zig/kernels/cuda/prefill_attention.cuh @@ -0,0 +1,67 @@ +// prefill_attention.cu's staging, MMA and fold helpers, shared with its controls (contract_controls.cu). +#pragma once + +#include +#include +#include +#include +#include + +#include "common/qmm_frag.cuh" + +namespace tf_prefill_attention { + +using qmm_frag::commit; +using qmm_frag::cp16z; +using qmm_frag::ldmatrix4; +using qmm_frag::mma; +using qmm_frag::mma0; + +constexpr int BN = 64, HALF = 32; // keys a tile (the fold's unit), keys a staging slot +constexpr float LOG2E = 1.4426950408889634f; +constexpr int SMEM_MAX = 96 * 1024; + +template +__device__ __forceinline__ void wait_group() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); } + +__device__ __forceinline__ void ldmatrix4t(uint32_t (&r)[4], const void* p) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) : "r"(qmm_frag::smem(p))); +} + +__device__ __forceinline__ float texp(float x) { + float y; + asm("ex2.approx.f32 %0, %1;\n" : "=f"(y) : "f"(__fmul_rn(x, LOG2E))); + return y; +} + +__device__ __forceinline__ float tdiv(float a, float b) { + float y; + asm("div.full.f32 %0, %1, %2;\n" : "=f"(y) : "f"(a), "f"(b)); + return y; +} + +__device__ __forceinline__ uint32_t pack(float lo, float hi) { + __nv_bfloat162 v = __floats2bfloat162_rn(lo, hi); + return *reinterpret_cast(&v); +} + +template +__device__ __forceinline__ int sw(int r, int c) { return r * (D / 8) + (c ^ (r & 7)); } + +__device__ __forceinline__ float rowsum(const float (&c)[8]) { + float w[8]; +#pragma unroll + for (int j = 0; j < 8; ++j) { + const float x = __fadd_rn(c[j], __shfl_xor_sync(0xffffffffu, c[j], 2)); + w[j] = __fadd_rn(x, __shfl_xor_sync(0xffffffffu, x, 1)); + } + return __fadd_rn(__fadd_rn(__fadd_rn(w[0], w[4]), __fadd_rn(w[2], w[6])), + __fadd_rn(__fadd_rn(w[1], w[5]), __fadd_rn(w[3], w[7]))); +} + +#ifndef PATTN_VLANE +#define PATTN_VLANE 0 // 1 in contract_controls.cu only: values read from the neighbouring 8 dims +#endif + +} // namespace tf_prefill_attention diff --git a/zig/kernels/cuda/prefill_attention_rows.cuh b/zig/kernels/cuda/prefill_attention_rows.cuh index 1986522ed2..0f4a5e9953 100644 --- a/zig/kernels/cuda/prefill_attention_rows.cuh +++ b/zig/kernels/cuda/prefill_attention_rows.cuh @@ -165,14 +165,14 @@ #pragma unroll for (int i = 0; i < NB; i += 2) { uint32_t b[4]; - ldmatrix4t(b, sl + sw<8 * CS>(8 * (mt % 2) + lane % 8, i + mt / 2)); + ldmatrix4t(b, sl + sw<8 * CS>(8 * (mt % 2) + lane % 8, (i + mt / 2) ^ PATTN_VLANE)); mma(o[i], pa0, b[0], b[1]); mma(o[i + 1], pa0, b[2], b[3]); } #pragma unroll for (int i = 0; i < NB; i += 2) { uint32_t b[4]; - ldmatrix4t(b, sl + sw<8 * CS>(16 + 8 * (mt % 2) + lane % 8, i + mt / 2)); + ldmatrix4t(b, sl + sw<8 * CS>(16 + 8 * (mt % 2) + lane % 8, (i + mt / 2) ^ PATTN_VLANE)); mma(o[i], pa1, b[0], b[1]); mma(o[i + 1], pa1, b[2], b[3]); } diff --git a/zig/lfm3_check.zig b/zig/lfm3_check.zig new file mode 100644 index 0000000000..577d4910c7 --- /dev/null +++ b/zig/lfm3_check.zig @@ -0,0 +1,55 @@ +//! Check saved header metadata from tools/lfm3/inspect_checkpoint.py without loading model data or running a GPU. +const std = @import("std"); +const lfm3 = @import("lfm3"); +const st = @import("safetensors"); + +pub fn main(init: std.process.Init) !void { + const a = init.arena.allocator(); + const args = try init.minimal.args.toSlice(a); + if (args.len != 2) { + std.debug.print("usage: tf-lfm3-check CHECKPOINT_METADATA.json\n", .{}); + return error.ExpectedMetadata; + } + const bytes = try std.Io.Dir.cwd().readFileAlloc(init.io, args[1], a, .limited(32 << 20)); + const value = try std.json.parseFromSliceLeaky(std.json.Value, a, bytes, .{}); + if (value != .object) return error.BadMetadata; + const root = value.object; + const cfg_bytes = try std.json.Stringify.valueAlloc(a, root.get("config") orelse return error.BadMetadata, .{}); + var cfg = try lfm3.config.parse(a, cfg_bytes); + defer cfg.deinit(); + const tensors = root.get("tensors") orelse return error.BadMetadata; + const shards = root.get("shards") orelse return error.BadMetadata; + const lengths = root.get("shard_data_lengths") orelse return error.BadMetadata; + if (tensors != .object or shards != .array or shards.array.items.len == 0 or lengths != .object) + return error.BadMetadata; + var header: st.Header = .empty; + for (shards.array.items) |shard| { + if (shard != .string) return error.BadMetadata; + var selected: std.json.ObjectMap = .empty; + const length = lengths.object.get(shard.string) orelse return error.BadMetadata; + if (length != .integer or length.integer < 0) return error.BadMetadata; + const data_len = std.math.cast(usize, length.integer) orelse return error.BadMetadata; + var it = tensors.object.iterator(); + while (it.next()) |item| { + const tensor = item.value_ptr.*; + if (tensor != .object) return error.BadMetadata; + const file = tensor.object.get("shard") orelse return error.BadMetadata; + if (file != .string) return error.BadMetadata; + if (!std.mem.eql(u8, file.string, shard.string)) continue; + try selected.put(a, item.key_ptr.*, tensor); + } + if (selected.count() == 0) return error.BadMetadata; + const text = try std.json.Stringify.valueAlloc(a, std.json.Value{ .object = selected }, .{}); + const parsed = try st.parseHeader(a, text, data_len); + for (parsed.keys(), parsed.values()) |key, entry| { + if (header.contains(key)) return error.DuplicateTensor; + try header.put(a, key, entry); + } + } + if (header.count() != tensors.object.count()) return error.BadMetadata; + const manifest = try lfm3.weights.validate(a, &cfg, &header); + std.debug.print("LFM3 text admission passed: {s}, {d} text tensors, {d} stored bytes; ", .{ + @tagName(cfg.kind), manifest.tensors, manifest.bytes, + }); + std.debug.print("{d} other tensors not qualified; GPU inference not tested\n", .{header.count() - manifest.tensors}); +} diff --git a/zig/src/core/contract/lfm2_moe_cuda_records.zig b/zig/src/core/contract/lfm2_moe_cuda_records.zig new file mode 100644 index 0000000000..e517e0f33f --- /dev/null +++ b/zig/src/core/contract/lfm2_moe_cuda_records.zig @@ -0,0 +1,24 @@ +//! LFM2.5-24B-A2B's CUDA records: the 3-tap liquid conv and commit, its route, head-64 paged attention, every SM. +const std = @import("std"); +const registry = @import("registry.zig"); +const R = registry.CudaRecord; + +const pattn64 = "_ZN20tf_prefill_attention12pattn_kernelILi64ELi8ELi4ELi8EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifPKix"; + +pub const cuda = [_]R{ + .{ .launch = .{ .name = "liquid_conv tf_liquid_conv_3", .image = "lfm2_moe", .symbol = "tf_liquid_conv_3", .threads = 256 }, .usage = .{ .{ .regs = 28, .static = 0, .bound = 256 }, .{ .regs = 26, .static = 0, .bound = 256 }, .{ .regs = 26, .static = 0, .bound = 256 }, .{ .regs = 28, .static = 0, .bound = 256 }, .{ .regs = 28, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "liquid_conv tf_liquid_commit_3", .image = "lfm2_moe", .symbol = "tf_liquid_commit_3", .threads = 256 }, .usage = .{ .{ .regs = 24, .static = 0, .bound = 256 }, .{ .regs = 24, .static = 0, .bound = 256 }, .{ .regs = 24, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "lfm2_moe tf_lfm2m_route", .image = "lfm2_moe", .symbol = "tf_lfm2m_route", .threads = 256 }, .usage = .{ .{ .regs = 55, .static = 34816, .bound = 256 }, .{ .regs = 50, .static = 34816, .bound = 256 }, .{ .regs = 50, .static = 34816, .bound = 256 }, .{ .regs = 58, .static = 34816, .bound = 256 }, .{ .regs = 59, .static = 34816, .bound = 256 } } }, + .{ .launch = .{ .name = "paging tf_paged_attn_chunk_segs_64", .image = "nemotron_attention", .symbol = "tf_paged_attn_chunk_segs_64", .threads = 256 }, .usage = .{ .{ .regs = 48, .static = 25088, .bound = 256 }, .{ .regs = 78, .static = 25088, .bound = 256 }, .{ .regs = 78, .static = 25088, .bound = 256 }, .{ .regs = 80, .static = 25088, .bound = 256 }, .{ .regs = 80, .static = 25088, .bound = 256 } } }, + .{ .launch = .{ .name = "paging tf_paged_attn_merge_segs_64", .image = "nemotron_attention", .symbol = "tf_paged_attn_merge_segs_64", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 4096, .bound = 256 }, .{ .regs = 40, .static = 4096, .bound = 256 }, .{ .regs = 40, .static = 4096, .bound = 256 }, .{ .regs = 40, .static = 4096, .bound = 256 }, .{ .regs = 40, .static = 4096, .bound = 256 } } }, + .{ .launch = .{ .name = "paging pattn64", .image = "prefill_attention", .symbol = pattn64, .threads = 256, .dynamic = 32768, .opt_in = true }, .usage = .{ .{ .regs = 148, .static = 0, .bound = 256 }, .{ .regs = 148, .static = 0, .bound = 256 }, .{ .regs = 148, .static = 0, .bound = 256 }, .{ .regs = 164, .static = 0, .bound = 256 }, .{ .regs = 164, .static = 0, .bound = 256 } } }, +}; + +test "LFM2.5-24B-A2B records: every launch fits every SM by its measured record" { + for (cuda) |r| try std.testing.expectEqual(@as(usize, 0), try registry.checkCuda(std.testing.allocator, &.{r}, true)); +} + +/// The launch nvcc 13.2 compiles differently: the router on sm_120 (one register more). +pub const cuda_132 = [_]registry.Override{ + .{ .name = "lfm2_moe tf_lfm2m_route", .usage = .{ .{ .regs = 55, .static = 34816, .bound = 256 }, .{ .regs = 50, .static = 34816, .bound = 256 }, .{ .regs = 50, .static = 34816, .bound = 256 }, .{ .regs = 59, .static = 34816, .bound = 256 }, .{ .regs = 59, .static = 34816, .bound = 256 } } }, +}; diff --git a/zig/src/core/contract/lfm3_cuda_records.zig b/zig/src/core/contract/lfm3_cuda_records.zig new file mode 100644 index 0000000000..c6e04690c5 --- /dev/null +++ b/zig/src/core/contract/lfm3_cuda_records.zig @@ -0,0 +1,16 @@ +//! LFM3's CUDA records: the head-128 paged prompt attention at four query heads a KV head, the 7-tap Liquid conv. +const std = @import("std"); +const registry = @import("registry.zig"); +const R = registry.CudaRecord; + +const pattn128x4 = "_ZN20tf_prefill_attention12pattn_kernelILi128ELi8ELi4ELi8EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifPKix"; + +pub const cuda = [_]R{ + .{ .launch = .{ .name = "liquid_conv tf_liquid_conv_7", .image = "lfm3", .symbol = "tf_liquid_conv_7", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "liquid_conv tf_liquid_commit_7", .image = "lfm3", .symbol = "tf_liquid_commit_7", .threads = 256 }, .usage = .{ .{ .regs = 24, .static = 0, .bound = 256 }, .{ .regs = 24, .static = 0, .bound = 256 }, .{ .regs = 24, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "paging pattn128x4", .image = "prefill_attention", .symbol = pattn128x4, .threads = 256, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 206, .static = 0, .bound = 256 }, .{ .regs = 206, .static = 0, .bound = 256 }, .{ .regs = 206, .static = 0, .bound = 256 }, .{ .regs = 231, .static = 0, .bound = 256 }, .{ .regs = 231, .static = 0, .bound = 256 } } }, +}; + +test "LFM3 records: every launch fits every SM by its measured record" { + for (cuda) |r| try std.testing.expectEqual(@as(usize, 0), try registry.checkCuda(std.testing.allocator, &.{r}, true)); +} diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 7d83d71249..d1dce2f17d 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -207,7 +207,7 @@ pub const cuda = @import("flashnext_cuda_records.zig").cuda ++ [_]CudaRecord{ .{ .launch = .{ .name = "experts3_prefill down", .symbol = "_ZN19tf_experts3_prefill14prefill_kernelILi1ELi3ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_PvifS6_i", .image = "experts3_prefill", .threads = 256, .dynamic = 47104, .opt_in = true }, .usage = .{ .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 98, .static = 0, .bound = 256 }, .{ .regs = 98, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "experts3 tf_experts3_widen", .symbol = "tf_experts3_widen", .image = "experts3", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 32, .static = 0, .bound = 0 } } }, .{ .launch = .{ .name = "experts3 tf_experts3_planes", .symbol = "tf_experts3_planes", .image = "experts3", .threads = 256 }, .usage = .{ .{ .regs = 26, .static = 0, .bound = 0 }, .{ .regs = 28, .static = 0, .bound = 0 }, .{ .regs = 28, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 } } }, -} ++ @import("registry_media.zig").cuda ++ @import("nemotron_prompt_records.zig").cuda ++ @import("glm_cuda_records.zig").cuda ++ @import("registry_shared_attention.zig").records(CudaRecord) ++ @import("gemma_cuda_records.zig").cuda ++ @import("spark_cuda_records.zig").cuda ++ @import("row_cuda_records.zig").cuda ++ @import("flashnext_tile_records.zig").cuda ++ @import("window_rows_records.zig").cuda ++ @import("modelopt_lane_records.zig").cuda ++ @import("nvfp4_fact_records.zig").cuda ++ @import("lane_half_records.zig").cuda; +} ++ @import("registry_media.zig").cuda ++ @import("nemotron_prompt_records.zig").cuda ++ @import("glm_cuda_records.zig").cuda ++ @import("registry_shared_attention.zig").records(CudaRecord) ++ @import("gemma_cuda_records.zig").cuda ++ @import("spark_cuda_records.zig").cuda ++ @import("row_cuda_records.zig").cuda ++ @import("flashnext_tile_records.zig").cuda ++ @import("window_rows_records.zig").cuda ++ @import("modelopt_lane_records.zig").cuda ++ @import("nvfp4_fact_records.zig").cuda ++ @import("lane_half_records.zig").cuda ++ @import("lfm2_moe_cuda_records.zig").cuda ++ @import("lfm3_cuda_records.zig").cuda; /// The nvcc release `cuda`'s usage records come from: the compiler that builds the release fatbins. pub const cuda_release: limits.Nvcc = .{ .major = 13, .minor = 0 }; @@ -319,7 +319,7 @@ pub const cuda_sets = [_]CudaSet{ .{ .name = "latent tf_super_topk_few", .usage = .{ .{ .regs = 32, .static = 4352, .bound = 1024 }, .{ .regs = 38, .static = 4352, .bound = 1024 }, .{ .regs = 38, .static = 4352, .bound = 1024 }, .{ .regs = 32, .static = 4352, .bound = 1024 }, .{ .regs = 32, .static = 4352, .bound = 1024 } } }, .{ .name = "ck tf_ck4_up_relu2_few", .usage = .{ .{ .regs = 19, .static = 0, .bound = 128 }, .{ .regs = 19, .static = 0, .bound = 128 }, .{ .regs = 19, .static = 0, .bound = 128 }, .{ .regs = 124, .static = 0, .bound = 128 }, .{ .regs = 124, .static = 0, .bound = 128 } } }, .{ .name = "mamba tf_nemo_scan", .usage = .{ .{ .regs = 138, .static = 41088, .bound = 128 }, .{ .regs = 144, .static = 41088, .bound = 128 }, .{ .regs = 144, .static = 41088, .bound = 128 }, .{ .regs = 118, .static = 41088, .bound = 128 }, .{ .regs = 118, .static = 41088, .bound = 128 } } }, - } ++ learning.cuda_132 ++ @import("paging_pending.zig").cuda_132 ++ @import("nemotron_prompt_records.zig").cuda_132 ++ @import("glm_cuda_records.zig").cuda_132 ++ @import("qwen_served_records.zig").cuda_132 ++ @import("flashnext_tile_records.zig").cuda_132 ++ @import("nvfp4_fact_records.zig").cuda_132) }, + } ++ learning.cuda_132 ++ @import("paging_pending.zig").cuda_132 ++ @import("lfm2_moe_cuda_records.zig").cuda_132 ++ @import("nemotron_prompt_records.zig").cuda_132 ++ @import("glm_cuda_records.zig").cuda_132 ++ @import("qwen_served_records.zig").cuda_132 ++ @import("flashnext_tile_records.zig").cuda_132 ++ @import("nvfp4_fact_records.zig").cuda_132) }, .{ .release = .{ .major = 13, .minor = 3 }, .usage = &([_]Override{ .{ .name = "router tf_nemo_topk", .usage = .{ .{ .regs = 94, .static = 0, .bound = 32 }, .{ .regs = 95, .static = 0, .bound = 32 }, .{ .regs = 95, .static = 0, .bound = 32 }, .{ .regs = 96, .static = 0, .bound = 32 }, .{ .regs = 96, .static = 0, .bound = 32 } } }, .{ .name = "flashnext_batch fn_qsa_select", .usage = .{ .{ .regs = 32, .static = 1168, .bound = 1024 }, .{ .regs = 50, .static = 1168, .bound = 1024 }, .{ .regs = 50, .static = 1168, .bound = 1024 }, .{ .regs = 44, .static = 1168, .bound = 1024 }, .{ .regs = 44, .static = 1168, .bound = 1024 } } }, diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 33a3bbfe55..12c7d3edf2 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -34,6 +34,8 @@ pub const experts_pack: []const u8 = if (available) &Blob("fatbin_experts_pack") pub const experts_rows: []const u8 = if (available) &Blob("fatbin_experts_rows").bytes else &.{}; pub const prefill_attention: []const u8 = if (available) &Blob("fatbin_prefill_attention").bytes else &.{}; pub const lfm2: []const u8 = if (available) &Blob("fatbin_lfm2").bytes else &.{}; +pub const lfm2_moe: []const u8 = if (available) &Blob("fatbin_lfm2_moe").bytes else &.{}; +pub const lfm3: []const u8 = if (available) &Blob("fatbin_lfm3").bytes else &.{}; pub const vit: []const u8 = if (available) &Blob("fatbin_vit").bytes else &.{}; pub const scan_rows: []const u8 = if (available) &Blob("fatbin_scan_rows").bytes else &.{}; pub const nemotron_ops: []const u8 = if (available) &Blob("fatbin_nemotron_ops").bytes else &.{}; diff --git a/zig/src/families/lfm2_moe/cuda_backend.zig b/zig/src/families/lfm2_moe/cuda_backend.zig new file mode 100644 index 0000000000..94390645ba --- /dev/null +++ b/zig/src/families/lfm2_moe/cuda_backend.zig @@ -0,0 +1,17 @@ +//! LFM2.5-24B-A2B on CUDA behind lanes.Backend: liquid's lane backend over this family's model and forward. +const liquid = @import("liquid_cuda"); + +const B = liquid.lanes.Backend(struct { + pub const cfg = @import("config.zig"); + pub const fwd = @import("cuda_forward.zig"); + pub const mdl = @import("cuda_model.zig"); + pub const kern = @import("cuda_kernels.zig"); + pub const snap = @import("cuda_snap.zig"); + pub const text_only = error.LfmTextOnly; +}); + +pub const Liquid = B.Liquid; +pub const prompt_rows = B.prompt_rows; +pub const round_rows = B.round_rows; +pub const max_window = B.max_window; +pub const pagesFor = B.pagesFor; diff --git a/zig/src/families/lfm2_moe/cuda_family.zig b/zig/src/families/lfm2_moe/cuda_family.zig new file mode 100644 index 0000000000..50263f9d8a --- /dev/null +++ b/zig/src/families/lfm2_moe/cuda_family.zig @@ -0,0 +1,20 @@ +//! LFM2.5-24B-A2B (lfm2_moe) on CUDA: the Metal family's config admission, the shared CUDA operators, lfm2_moe.cu. +pub const config = @import("config.zig"); +pub const weights = @import("cuda_weights.zig"); +pub const kernels = @import("cuda_kernels.zig"); +pub const model = @import("cuda_model.zig"); +pub const forward = @import("cuda_forward.zig"); +pub const backend = @import("cuda_backend.zig"); +pub const native = @import("cuda_native.zig"); +pub const snap = @import("cuda_snap.zig"); + +test { + _ = config; + _ = weights; + _ = kernels; + _ = model; + _ = forward; + _ = backend; + _ = native; + _ = snap; +} diff --git a/zig/src/families/lfm2_moe/cuda_forward.zig b/zig/src/families/lfm2_moe/cuda_forward.zig new file mode 100644 index 0000000000..f6723889f9 --- /dev/null +++ b/zig/src/families/lfm2_moe/cuda_forward.zig @@ -0,0 +1,235 @@ +//! LFM2.5-24B-A2B's forward on CUDA: a prompt chunk or a round of windows, each row's bits its own (solo == batched). +const std = @import("std"); +const cuda = @import("cuda"); +const c = @import("config.zig"); +const wts = @import("cuda_weights.zig"); +const kern = @import("cuda_kernels.zig"); +const mdl = @import("cuda_model.zig"); + +const Model = mdl.Model; +const Kernels = kern.Kernels; +const Stream = cuda.Stream; +const tail: cuda.norms.Tail = .{ .round = true }; // the norms' torch points: bf16(w * bf16(h * inv)) + +/// A stream's rows in one forward: its slot, first position, tokens; a prompt chunk's tap rows keep their conv states. +pub const Chain = struct { slot: usize, p0: usize, tokens: []const u32, taps: []const usize = &.{} }; + +const Pass = struct { + m: *const Model, + k: *const Kernels, + s: Stream, + chains: []const Chain, + rows: usize, + prompt: bool, + + /// x [rows, k] through `l` into out [rows, n]: a prompt chunk's GEMM, or decode rows on their group sums. + fn project(p: Pass, l: wts.Proj, x: u64, xs: u64, out: u64, rows: usize, prompt: bool) !void { + const k = p.k; + if (prompt) return switch (l) { + .a4 => |q| k.affine.prompt(p.s, x, q, out, rows), + .an => |q| k.wide.prompt(p.s, x, q, out, rows), + }; + const tile: usize = switch (l) { + .a4 => k.affine.verifyRows(), + .an => 16, + }; + var at: usize = 0; + while (at < rows) : (at += tile) { + const n = @min(tile, rows - at); + const xa = x + at * l.k() * 2; + const sa = xs + at * (l.k() / 64) * 4; + const oa = out + at * l.n() * 2; + switch (l) { + .a4 => |q| try k.affine.verify(p.s, xa, sa, q, oa, n), + .an => |q| try k.wide.decode(p.s, xa, sa, q, oa, n), + } + } + } + + /// The 64-value sums of x [rows, width] a decode projection reads (prompt GEMMs take none). + fn sums(p: Pass, x: u64, width: usize) !void { + if (!p.prompt) try p.k.act.groupSums(p.s, x, width, p.m.s.xs, p.rows, width); + } + + /// h = bf16(h + r) (r null: as is), x = bf16(w * bf16(h * inv_rms)) and its sums. + fn norm(p: Pass, r: ?u64, w: u64) !void { + const x = p.m.s; + try p.k.norms.addRmsTail(p.s, x.h, r, w, x.h, x.x, x.xs, p.rows, c.hidden, c.eps, tail); + } + + fn conv(p: Pass, cv: wts.Conv, index: usize) !void { + const x = p.m.s; + try p.project(cv.in, x.x, x.xs, x.wide, p.rows, p.prompt); + const state = p.m.states + index * 2 * c.hidden * 2; + const snaps = x.snaps + index * mdl.round_rows * 2 * c.hidden * 2; + // a prompt chunk keeps its last row's state, a round each row's (a window keeps any prefix of its rows) + const snap0: usize = if (p.prompt) p.rows - 1 else 0; + try p.k.convRows(p.s, x.wide, cv.weight, state, mdl.slot_elems, x.slot, x.off, x.y, snaps, snap0, p.rows); + if (p.prompt) for (p.chains[0].taps, 0..) |r, j| { // the kernel again on row r: the same y, its states in tap j + const tap = x.tap + (j * mdl.slot_elems + index * 2 * c.hidden) * 2; + try p.k.convRows(p.s, x.wide + r * 3 * c.hidden * 2, cv.weight, state, mdl.slot_elems, x.slot + r * 4, x.off + r * 4, x.y + r * c.hidden * 2, tap, 0, 1); + }; + try p.sums(x.y, c.hidden); + try p.project(cv.out, x.y, x.xs, x.r, p.rows, p.prompt); + } + + fn attention(p: Pass, a: wts.Attention, index: usize) !void { + const x = p.m.s; + const m = p.m; + const qkv_w = c.hidden + 2 * mdl.kv_width; + try p.project(a.qkv, x.x, x.xs, x.wide, p.rows, p.prompt); + const ld: i64 = @intCast(qkv_w); + const dh: i64 = c.head_dim; + const src: cuda.qk_rows.Src = .{ .q = x.wide, .q_ld = ld, .q_hs = dh, .k = x.wide + c.hidden * 2, .k_ld = ld, .k_hs = dh, .v = x.wide + (c.hidden + mdl.kv_width) * 2, .v_ld = ld, .v_hs = dh }; + const shape: cuda.qk_rows.Shape = .{ .norms = a.norms, .norms_bf16 = true, .inv = m.inv, .rot = c.head_dim, .eps = c.eps, .round = true }; + const out: cuda.qk_rows.Out = .{ .q = x.q, .k = x.kv, .v = x.kv + mdl.kv_width * 2, .kv_ld = 2 * mdl.kv_width }; + try p.k.qk.rowsFrom(p.s, src, .{ .h = c.query_heads, .hk = c.kv_heads, .dim = c.head_dim }, shape, x.pos, null, out, p.rows); + const kc = m.cache(index, false); + const vc = m.cache(index, true); + var row0: usize = 0; + var at: usize = 0; + while (at < p.chains.len) : (at += kern.max_segs) { + const group = p.chains[at..@min(p.chains.len, at + kern.max_segs)]; + var sg = std.mem.zeroes(kern.AttnSegs); + var rows: usize = 0; + for (group, 0..) |ch, j| { + sg.kc[j] = kc; + sg.vc[j] = vc; + sg.meta[j] = x.meta + (at + j) * 4; + sg.pt[j] = m.tableOf(ch.slot); + sg.row0[j] = @intCast(rows); + rows += ch.tokens.len; + } + sg.ps = mdl.page_elems; + sg.n = @intCast(group.len); + try p.k.kvWrite(p.s, sg, x.kv + row0 * 2 * mdl.kv_width * 2, rows); + const qa = x.q + row0 * c.hidden * 2; + const ya = x.y + row0 * c.hidden * 2; + if (p.prompt) { + try p.k.prompt(p.s, qa, kc, vc, ya, group[0].p0, rows, m.tableOf(group[0].slot), mdl.page_elems); + } else { + const xa = x.xs + row0 * (c.hidden / 64) * 4; + try p.k.attend(p.s, sg, qa, x.po, x.pm, x.pl, ya, xa, rows, m.chunks(), m.chunkBlocks()); + } + row0 += rows; + } + try p.project(a.out, x.y, x.xs, x.r, p.rows, p.prompt); + } + + fn dense(p: Pass, dn: wts.Dense) !void { + const x = p.m.s; + try p.project(dn.gate_up, x.x, x.xs, x.wide, p.rows, p.prompt); + try p.k.siluMul(p.s, x.wide, x.y, p.rows, c.intermediate); + try p.sums(x.y, c.intermediate); + try p.project(dn.down, x.y, x.xs, x.r, p.rows, p.prompt); + } + + fn moe(p: Pass, e: wts.Moe) !void { + const x = p.m.s; + const k = p.k; + const pairs = p.rows * c.topk; + try k.routeRows(p.s, x.x, e.router, e.bias, x.picks, x.wts, null, p.rows, p.m.w.config.scale); + const tile: usize = if (p.prompt) cuda.experts.promptTile(.affine4) else 16; + const items = cuda.grouped.maxItems(pairs, c.experts, tile); + const plan: cuda.grouped.Plan = .{ .members = x.members, .items = x.items, .counts = x.counts, .rank = x.rank, .hist = x.hist }; + try k.plan.route(p.s, x.picks, pairs, c.experts, tile, plan); + const grouped: cuda.experts.Grouped = .{ .affine4 = &k.experts }; + const rows_in: cuda.experts.Rows = .{ .x = x.x, .stride = c.hidden, .slots = c.topk }; + const act_in: cuda.experts.Rows = .{ .x = x.eact, .stride = c.moe_intermediate }; + if (p.prompt) { + try grouped.prompt(p.s, .up, rows_in, e.experts, plan, x.eact, items, -1); + try grouped.prompt(p.s, .down_bf16, act_in, e.experts, plan, x.ey, items, -1); + } else { + try grouped.decode(p.s, .up, rows_in, e.experts, plan, x.eact, items, -1); + try grouped.decode(p.s, .down_f32, act_in, e.experts, plan, x.ey, items, -1); + } + try k.moe.combine(p.s, x.ey, !p.prompt, x.wts, x.r, p.rows, c.topk, c.hidden, 0); + } +}; + +/// Every chain's rows (a prompt chunk, or windows); logits rows [head0, head0 + heads); conv states end at last rows. +pub fn forward(m: *const Model, k: *const Kernels, s: Stream, chains: []const Chain, prompt: bool, head0: usize, heads: usize) !void { + var rows: usize = 0; + for (chains) |ch| { + if (ch.tokens.len == 0 or ch.slot >= m.slots or ch.p0 + ch.tokens.len > m.context) return error.BadChain; + if (ch.p0 + ch.tokens.len > m.caps[ch.slot]) return error.PastPages; + if (ch.taps.len > 0 and (!prompt or ch.taps.len > mdl.max_taps)) return error.BadChain; + for (ch.taps) |r| if (r >= ch.tokens.len) return error.BadChain; + rows += ch.tokens.len; + } + if (rows > @as(usize, if (prompt) mdl.prompt_rows else mdl.round_rows) or (prompt and chains.len != 1)) return error.BadWindow; + if (chains.len > mdl.round_rows or heads > mdl.round_rows or head0 + heads > rows) return error.BadWindow; + try upload(m, s, chains, rows); + const p: Pass = .{ .m = m, .k = k, .s = s, .chains = chains, .rows = rows, .prompt = prompt }; + const x = m.s; + try k.embed.affine(s, x.ids, m.w.embed_w, m.w.embed_s, m.w.embed_b, x.h, rows, c.hidden, .{ .bits = @intCast(m.w.config.embed_bits), .group = c.group }); + var conv_at: usize = 0; + var attn_at: usize = 0; + for (m.w.layers, 0..) |l, i| { + try p.norm(if (i == 0) null else x.r, l.operator_norm); + switch (l.mixer) { + .conv => |cv| { + try p.conv(cv, conv_at); + conv_at += 1; + }, + .attention => |a| { + try p.attention(a, attn_at); + attn_at += 1; + }, + } + try p.norm(x.r, l.ffn_norm); + switch (l.ffn) { + .dense => |dn| try p.dense(dn), + .moe => |e| try p.moe(e), + } + } + try p.norm(x.r, m.w.final_norm); + try commitAll(m, k, s, chains, prompt); + if (heads == 0) return; + try p.project(m.w.head, x.x + head0 * c.hidden * 2, x.xs + head0 * (c.hidden / 64) * 4, x.logits, heads, false); +} + +/// Each stream's conv states (slots[i]) set to its state after snapshot row rows[i]; a row < 0 keeps them. +pub fn commit(m: *const Model, k: *const Kernels, s: Stream, slots: []const i32, rows: []const i32) !void { + if (slots.len != rows.len or slots.len > mdl.round_rows) return error.BadWindow; + try put(m, s, m.s.commit_slot, std.mem.sliceAsBytes(slots)); + try put(m, s, m.s.commit_row, std.mem.sliceAsBytes(rows)); + try k.commitRows(s, m.states, m.s.snaps, m.s.commit_slot, m.s.commit_row, slots.len, mdl.conv_layers, 2 * c.hidden, mdl.round_rows * 2 * c.hidden, mdl.slot_elems); +} + +/// Each chain's conv states at its last row. +fn commitAll(m: *const Model, k: *const Kernels, s: Stream, chains: []const Chain, prompt: bool) !void { + var slots: [mdl.round_rows]i32 = undefined; + var rows: [mdl.round_rows]i32 = undefined; + var base: usize = 0; + for (chains, 0..) |ch, i| { + base += ch.tokens.len; + slots[i] = @intCast(ch.slot); + rows[i] = if (prompt) 0 else @intCast(base - 1); + } + try commit(m, k, s, slots[0..chains.len], rows[0..chains.len]); +} + +fn put(m: *const Model, s: Stream, dst: u64, bytes: []const u8) !void { + try m.d.check(m.d.api.cuMemcpyHtoDAsync_v2(dst, bytes.ptr, bytes.len, s.handle), "cuMemcpyHtoDAsync"); +} + +/// Tokens, positions, slots and offsets a row, each chain's first position, staged for the launches. +fn upload(m: *const Model, s: Stream, chains: []const Chain, rows: usize) !void { + const ids = try m.gpa.alloc(i32, 4 * rows + chains.len); + defer m.gpa.free(ids); + var at: usize = 0; + for (chains, 0..) |ch, j| { + for (ch.tokens, 0..) |t, i| { + ids[at + i] = @intCast(t); + ids[rows + at + i] = @intCast(ch.p0 + i); + ids[2 * rows + at + i] = @intCast(ch.slot); + ids[3 * rows + at + i] = @intCast(i); + } + ids[4 * rows + j] = @intCast(ch.p0); + at += ch.tokens.len; + } + const x = m.s; + for ([_]u64{ x.ids, x.pos, x.slot, x.off }, 0..) |dst, i| try put(m, s, dst, std.mem.sliceAsBytes(ids[i * rows ..][0..rows])); + try put(m, s, x.meta, std.mem.sliceAsBytes(ids[4 * rows ..])); +} diff --git a/zig/src/families/lfm2_moe/cuda_kernels.zig b/zig/src/families/lfm2_moe/cuda_kernels.zig new file mode 100644 index 0000000000..2978245e6f --- /dev/null +++ b/zig/src/families/lfm2_moe/cuda_kernels.zig @@ -0,0 +1,193 @@ +//! LFM2.5-24B-A2B's launches on CUDA: the shared operators, liquid_conv.cuh's 3-tap conv and commit, the router. +const std = @import("std"); +const cuda = @import("cuda"); +const c = @import("config.zig"); + +const Args = cuda.Args; +const Function = cuda.Function; +const Stream = cuda.Stream; + +/// prefill_attention.cu's pattn_kernel<64, 8, 4, 8> over paged caches: eight 32-key slots of 64 dims, 32 rows a block. +const pattn_sym = "_ZN20tf_prefill_attention12pattn_kernelILi64ELi8ELi4ELi8EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifPKix"; +const pattn_smem: u32 = 8 * 32 * 8 * 16; +const pattn_rows = 32; +/// Positions a KV page holds (nemotron_attention.cu's PAGE and prefill_attention_rows.cuh's 512). +pub const page = 512; +/// Streams one shared-window attention launch takes (nemotron_attention.cu's MAXSEG). +pub const max_segs = 8; + +/// nemotron_attention.cu's AttnSegs: each stream's caches, position, page table and first row. +pub const AttnSegs = extern struct { kc: [max_segs]u64, vc: [max_segs]u64, meta: [max_segs]u64, pt: [max_segs]u64, ps: i64, row0: [max_segs]c_int, n: c_int }; + +const images = .{ "qmm_group", "qmm_prefill", "lane_gemv", "affine4_pack", "sample", "tree_round", "affine8", "affine_prefill", "experts", "experts_prefill", "experts_pack", "moe", "prefill_attention", "nemotron_attention", "lfm2", "lfm2_moe" }; + +fn int(x: usize) c_int { + return @intCast(x); +} + +pub const Kernels = struct { + mods: [images.len]cuda.Module, + affine: cuda.qlinear.Affine4Kernels, + wide: cuda.qlinear.AffineNKernels, + draws: cuda.sample.Draws, + round: cuda.tree_round.Kernels, + embed: cuda.embed_rows.EmbedRows, + norms: cuda.norms.Norms, + qk: cuda.qk_rows.QkRows, + act: cuda.act_rows.ActRows, + experts: cuda.experts.Affine4Experts, + plan: cuda.grouped.Router, + moe: cuda.moe.Fns, + pattn: Function, + kv_write: Function, + chunk: Function, + merge: Function, + silu_mul: Function, + conv: Function, + commit: Function, + route: Function, + + pub fn load(ctx: *const cuda.Context) !Kernels { + const d = ctx.d; + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var k: Kernels = undefined; + var loaded: usize = 0; + errdefer for (k.mods[0..loaded]) |*m| m.unload(); + inline for (images, 0..) |name, i| { + k.mods[i] = try cuda.Module.load(d, @field(cuda.kernels, name)); + loaded += 1; + } + const m = &k.mods; + const sms: usize = @intCast(try ctx.attribute(.multiprocessor_count)); + const gb10 = try ctx.capability() == 121; + k.affine = try cuda.qlinear.Affine4Kernels.resolve(d, m[0], m[1], m[2], m[3], sms, gb10); + errdefer k.affine.deinit(); + k.wide = try cuda.qlinear.AffineNKernels.resolve(d, m[6], m[7], m[3], sms, gb10); + errdefer k.wide.deinit(); + k.draws = try cuda.sample.Draws.resolve(m[4]); + k.round = try cuda.tree_round.Kernels.resolve(m[5]); + k.experts = try cuda.experts.Affine4Experts.resolve(m[8], m[9], m[10], sms); + k.plan = try cuda.grouped.Router.resolve(m[8]); + k.moe = try cuda.moe.Fns.resolve(m[11]); + k.pattn = try m[12].function(pattn_sym); + try k.pattn.allowDynamicShared(pattn_smem); + k.kv_write = try m[13].function("tf_nemo_kv_write_segs"); + k.chunk = try m[13].function("tf_paged_attn_chunk_segs_64"); + k.merge = try m[13].function("tf_paged_attn_merge_segs_64"); + k.silu_mul = try m[14].function("tf_lfm2_silu_mul"); + k.conv = try m[15].function("tf_liquid_conv_3"); + k.commit = try m[15].function("tf_liquid_commit_3"); + k.route = try m[15].function("tf_lfm2m_route"); + k.embed = try cuda.embed_rows.EmbedRows.load(d); + errdefer k.embed.unload(); + k.norms = try cuda.norms.Norms.load(d); + errdefer k.norms.unload(); + k.qk = try cuda.qk_rows.QkRows.load(d); + errdefer k.qk.unload(); + k.act = try cuda.act_rows.ActRows.load(d); + return k; + } + + pub fn deinit(k: *Kernels) void { + k.act.unload(); + k.qk.unload(); + k.norms.unload(); + k.embed.unload(); + k.wide.deinit(); + k.affine.deinit(); + for (&k.mods) |*m| m.unload(); + } + + fn go(f: Function, s: Stream, grid: [3]usize, block: usize, smem: u32, a: *Args) !void { + const g: cuda.Dim3 = .{ .x = @intCast(grid[0]), .y = @intCast(grid[1]), .z = @intCast(grid[2]) }; + try cuda.launch.launch(f, .{ .grid = g, .block = .{ .x = @intCast(block) }, .shared = smem }, s, a); + } + + /// Short conv over bcx [rows, 3 hidden]: each row's slot and offset; a slot's [2, hidden] state every `ss`. + pub fn convRows(k: *const Kernels, s: Stream, bcx: u64, w: u64, state: u64, ss: usize, slot: u64, off: u64, y: u64, snap: ?u64, snap0: usize, rows: usize) !void { + var a: Args = .{}; + for ([_]u64{ bcx, w, state, slot, off, y, snap orelse 0 }) |v| a.add(v); + a.add(int(snap0)); + a.add(int(c.hidden)); + a.add(@as(i64, @intCast(ss))); + try go(k.conv, s, .{ rows, c.hidden / 256, 1 }, 256, 0, &a); + } + + /// Each stream k's conv states (every conv layer) set to its snapshot row rows[k] (int32; < 0: kept). + pub fn commitRows(k: *const Kernels, s: Stream, state: u64, snap: u64, slots: u64, rows: u64, streams: usize, layers: usize, state_stride: usize, snap_stride: usize, ss: usize) !void { + if (streams == 0) return; + var a: Args = .{}; + for ([_]u64{ state, snap, slots, rows }) |v| a.add(v); + a.add(int(c.hidden)); + a.add(@as(i64, @intCast(state_stride))); + a.add(@as(i64, @intCast(snap_stride))); + a.add(@as(i64, @intCast(ss))); + try go(k.commit, s, .{ layers, streams, 1 }, 256, 0, &a); + } + + /// Router picks and weights ([rows, 4] int32, fp32) of x [rows, hidden] against fp32 rows [experts, hidden]. + pub fn routeRows(k: *const Kernels, s: Stream, x: u64, w: u64, bias: u64, picks: u64, wts: u64, logits: ?u64, rows: usize, scale: f32) !void { + var a: Args = .{}; + for ([_]u64{ x, w, bias, picks, wts, logits orelse 0 }) |v| a.add(v); + a.add(int(rows)); + a.add(int(c.hidden)); + a.add(scale); + try go(k.route, s, .{ (rows + 7) / 8, 1, 1 }, 256, 0, &a); + } + + /// out [rows, f] = bf16(bf16(silu(gate)) * up) from rows gate | up [rows, 2 f] (lfm2.cu). + pub fn siluMul(k: *const Kernels, s: Stream, g: u64, out: u64, rows: usize, f: usize) !void { + var a: Args = .{}; + a.add(g); + a.add(out); + a.add(int(f)); + try go(k.silu_mul, s, .{ rows, (f + 255) / 256, 1 }, 256, 0, &a); + } + + /// Rows k, v (kv [rows, 2 kv width]) into each stream's pages at its position on. + pub fn kvWrite(k: *const Kernels, s: Stream, sg: AttnSegs, kv: u64, rows: usize) !void { + var a: Args = .{}; + a.add(sg); + a.add(kv); + a.add(int(2 * c.kv_heads * c.head_dim)); + a.add(@as(c_int, 0)); + a.add(int(c.kv_heads * c.head_dim)); + try go(k.kv_write, s, .{ rows, 1, 1 }, 256, 0, &a); + } + + /// Decode rows: q [rows, heads, 64] over each stream's pages in 512-key chunks, merged in order into out and xs. + pub fn attend(k: *const Kernels, s: Stream, sg: AttnSegs, q: u64, po: u64, pm: u64, pl: u64, out: u64, xs: u64, rows: usize, nch: usize, z: usize) !void { + const h = c.query_heads; + const hk = c.kv_heads; + const g = h / hk; + var a: Args = .{}; + a.add(sg); + a.add(q); + for ([_]u64{ po, pm, pl }) |v| a.add(v); + for ([_]usize{ h * c.head_dim, h, hk, g, page, nch }) |v| a.add(int(v)); + a.add(1.0 / @sqrt(@as(f32, c.head_dim))); + try go(k.chunk, s, .{ rows, hk, z }, 256, 0, &a); + var b: Args = .{}; + b.add(sg); + for ([_]u64{ po, pm, pl, out, xs }) |v| b.add(v); + for ([_]usize{ h, hk, g, page, nch }) |v| b.add(int(v)); + try go(k.merge, s, .{ rows, hk, 1 }, 256, 0, &b); + } + + /// A prompt chunk's rows at positions p0.. against its stream's pages (filled through them): out [rows, heads, 64]. + pub fn prompt(k: *const Kernels, s: Stream, q: u64, kc: u64, vc: u64, out: u64, p0: usize, rows: usize, pt: u64, ps: usize) !void { + const g = c.query_heads / c.kv_heads; + var a: Args = .{}; + for ([_]u64{ q, kc, vc, out }) |v| a.add(v); + for ([_]usize{ p0, rows, c.query_heads, c.kv_heads, g }) |v| a.add(int(v)); + a.add(1.0 / @sqrt(@as(f32, c.head_dim))); + a.add(pt); + a.add(@as(i64, @intCast(ps))); + try go(k.pattn, s, .{ (rows + pattn_rows - 1) / pattn_rows, c.kv_heads * (g / 4), 1 }, 256, pattn_smem, &a); + } +}; + +test "the attention launches' geometry: 4 query heads a KV head, a pattn block's 4 heads" { + try std.testing.expectEqual(@as(usize, 4), c.query_heads / c.kv_heads); + try std.testing.expectEqual(@as(usize, 304), @sizeOf(AttnSegs)); // nemotron_attention.cu's struct, padded to 8 +} diff --git a/zig/src/families/lfm2_moe/cuda_model.zig b/zig/src/families/lfm2_moe/cuda_model.zig new file mode 100644 index 0000000000..33b0f56755 --- /dev/null +++ b/zig/src/families/lfm2_moe/cuda_model.zig @@ -0,0 +1,350 @@ +//! LFM2.5-24B-A2B's device state on CUDA: KV in 512-position pages, each slot's conv states, the scratch. +const std = @import("std"); +const cuda = @import("cuda"); +const kv = @import("core").kv_pool; +const c = @import("config.zig"); +const wts = @import("cuda_weights.zig"); +const kern = @import("cuda_kernels.zig"); + +const DeviceBuffer = cuda.DeviceBuffer; + +/// Rows a prompt chunk takes. +pub const prompt_rows = 1024; +/// Rows a round of windows takes over every stream. +pub const round_rows = 128; +/// Conv layers (every layer that is not attention) and attention layers of the admitted geometry. +pub const conv_layers = 30; +pub const attention_layers = c.layers - conv_layers; +/// The layers whose keys and values live in pages (liquid's kept-state walk). +pub const paged_layers = attention_layers; +/// Elements a KV row holds (8 heads of 64) and a page holds over every attention layer's keys and values. +pub const kv_width = c.kv_heads * c.head_dim; +pub const page_elems = attention_layers * 2 * kern.page * kv_width; +/// Elements a slot's conv states hold: two bx rows a conv layer. +pub const slot_elems = conv_layers * 2 * c.hidden; +/// Marks a prompt chunk taps (their conv states kept as the pass goes by; no rings to lose); one more ends it. +pub const max_taps = 4; +/// Key chunks one attention launch walks in parallel at most (nemotron's chunk_blocks). +const chunk_blocks = 64; + +pub const Scratch = struct { + ids: u64, // int32 [R]: tokens, positions, slots, offsets in their stream's rows, R each + pos: u64, + slot: u64, + off: u64, + meta: u64, // int32 [round_rows]: each chain's first position (AttnSegs' META), its commit slot and row + commit_slot: u64, + commit_row: u64, + h: u64, // bf16 [R, hidden]: the residual + x: u64, // bf16 [R, hidden]: a norm's output + xs: u64, // fp32 [R, intermediate / 64]: a projection input's 64-value sums + r: u64, // bf16 [R, hidden]: a block's output + wide: u64, // bf16 [R, 2 intermediate]: conv in (3 hidden), q|k|v, dense gate|up + y: u64, // bf16 [R, intermediate]: conv out, attention out, dense act + q: u64, // bf16 [R, hidden]: normed and rotated queries + kv: u64, // bf16 [R, 2 kv_width]: keys then values a row + picks: u64, // int32 [R, 4] + wts: u64, // fp32 [R, 4] + members: u64, + items: u64, + counts: u64, + rank: u64, + hist: u64, + eact: u64, // bf16 [R * 4, moe_intermediate] + ey: u64, // fp32 [R * 4, hidden] (bf16 at prompt rows) + po: u64, // decode attention partials: o, m, l + pm: u64, + pl: u64, + snaps: u64, // bf16 [conv_layers][round_rows][2][hidden]: each window row's conv state after it + tap: u64, // bf16 [max_taps][conv_layers][2][hidden]: the conv states after a chunk's tapped rows (a slot's layout) + logits: u64, // bf16 [round_rows, vocab] +}; + +pub const Model = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + w: *const wts.Weights, + context: usize, // positions a slot holds: whole pages + slots: usize, + paged: bool = false, // slots lease pages of the core pool (attach makes the arena), else each owns a window + pages: usize = 0, // pages in the arena + arena: u64 = 0, // [pages][attention layer][K, V][512][kv_width] bf16 + table: u64 = 0, // int32 [slots][pages a slot]: slot s's page ids in position order + states: u64 = 0, // bf16 [slots][conv layer][2][hidden] + ids: []u32 = &.{}, // the tables on the host + caps: []usize = &.{}, // positions slot s's pages cover + inv: u64 = 0, // fp32 [head_dim / 2]: RoPE's inverse frequencies + s: Scratch = undefined, + buffers: std.ArrayList(DeviceBuffer) = .empty, + planned: ?usize = null, // a plan: the bytes counted, nothing allocated + + pub fn perSlot(m: *const Model) usize { + return m.context / kern.page; + } + + /// A slot's regions outside its pages: its conv states (liquid's kept-state walk). + pub const fixed_bytes = slot_elems * 2; + + pub fn fixed(m: *const Model, slot: usize, i: usize) ?[2]u64 { + return if (i == 0) .{ m.states + slot * fixed_bytes, fixed_bytes } else null; + } + + /// Positions a slot holds for a window of `context`: whole pages. + pub fn positions(context: usize) usize { + return (context + kern.page - 1) / kern.page * kern.page; + } + + /// Bytes one stream's caches take at `context` positions: its pages and its conv states. + pub fn streamBytes(context: usize) usize { + return positions(context) / kern.page * page_elems * 2 + slot_elems * 2; + } + + /// What the host's pool plans: a stream's conv states and table row, a token its rows of the ten layers. + pub fn layout(m: *const Model) kv.Layout { + return .{ .fixed_bytes = slot_elems * 2 + m.perSlot() * 4, .token_bytes = page_elems * 2 / kern.page, .page_tokens = kern.page, .snap_bytes = 0 }; + } + + /// Slots a backend holds: the model's, or (paged) as many as a round takes. + pub fn slotCap(m: *const Model) usize { + return if (m.paged) round_rows else m.slots; + } + + /// Decode attention's key chunks a row walks at most, and the blocks that walk them. + pub fn chunks(m: *const Model) usize { + return m.perSlot(); + } + + pub fn chunkBlocks(m: *const Model) usize { + return @min(m.chunks(), chunk_blocks); + } + + fn alloc(m: *Model, bytes: usize) !u64 { + if (m.planned) |*p| { + p.* += (@max(bytes, 256) + 255) / 256 * 256; + return 0; + } + const b = try DeviceBuffer.alloc(m.d, @max(bytes, 256)); + errdefer { + var x = b; + x.free(); + } + try m.buffers.append(m.gpa, b); + return b.ptr; + } + + /// `slots` streams, each owning a window of pages; paged: one slot's states now, `attach` makes the pages. + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, s: cuda.Stream, w: *const wts.Weights, k: *const kern.Kernels, context: usize, slots: usize, paged: bool) !Model { + _ = k; + var m: Model = .{ .gpa = gpa, .d = d, .w = w, .context = positions(context), .slots = if (paged) 1 else slots, .paged = paged }; + errdefer m.deinit(); + try m.build(s); + return m; + } + + /// The device bytes `init` allocates for these sizes, counted by the same path without allocating. + pub fn openBytes(gpa: std.mem.Allocator, d: *const cuda.Driver, s: cuda.Stream, w: *const wts.Weights, k: *const kern.Kernels, context: usize, slots: usize, paged: bool) !usize { + _ = k; + var m: Model = .{ .gpa = gpa, .d = d, .w = w, .context = positions(context), .slots = if (paged) 1 else slots, .paged = paged, .planned = 0 }; + defer m.deinit(); + try m.build(s); + return m.planned.?; + } + + /// Makes the arena of `pool` pages to lease and the states and tables of `extra + 1` slots (no pages mapped). + pub fn attach(m: *Model, s: cuda.Stream, extra: usize, pool: usize) !void { + if (!m.paged or m.pages != 0 or extra >= round_rows) return error.BadAttach; + m.forget(m.table); + m.forget(m.states); + m.gpa.free(m.ids); + m.gpa.free(m.caps); + m.ids = &.{}; + m.caps = &.{}; + m.slots = extra + 1; + m.pages = pool; + m.arena = try m.alloc(pool * page_elems * 2); + try m.d.check(m.d.api.cuMemsetD8Async(m.arena, 0, pool * page_elems * 2, s.handle), "cuMemsetD8Async"); + try m.makeSlots(s); + try s.synchronize(); + } + + /// Slot `slot`'s pages from table entry `from` on are `ids` (arena page numbers in position order). + pub fn mapPages(m: *Model, s: cuda.Stream, slot: usize, ids: []const u32, from: usize) !void { + const at = try m.note(slot, ids, from); + try m.d.check(m.d.api.cuMemcpyHtoDAsync_v2(m.table + at * 4, ids.ptr, ids.len * 4, s.handle), "cuMemcpyHtoDAsync"); + } + + /// mapPages on the host: the table's copy and the slot's reach; the entry its device copy starts at. + fn note(m: *Model, slot: usize, ids: []const u32, from: usize) !usize { + const per = m.perSlot(); + if (slot >= m.slots or from + ids.len > per) return error.BadChain; + for (ids) |id| if (id >= m.pages) return error.BadChain; + const at = slot * per + from; + @memcpy(m.ids[at..][0..ids.len], ids); + m.caps[slot] = (from + ids.len) * kern.page; + return at; + } + + /// Slot `slot` holds no pages: its table entries stay, no forward may reach them. + pub fn unmapPages(m: *Model, slot: usize) void { + m.caps[slot] = 0; + } + + /// Frees one buffer `alloc` made. + fn forget(m: *Model, ptr: u64) void { + for (m.buffers.items, 0..) |*b, i| if (b.ptr == ptr) { + b.free(); + _ = m.buffers.swapRemove(i); + return; + }; + } + + /// The slots' tables and conv states; a window of pages a slot unless paged (their pages come from leases). + fn makeSlots(m: *Model, s: cuda.Stream) !void { + const per = m.perSlot(); + m.table = try m.alloc(m.slots * per * 4); + m.states = try m.alloc(m.slots * slot_elems * 2); + if (m.planned != null) return; + m.ids = try m.gpa.alloc(u32, m.slots * per); + m.caps = try m.gpa.alloc(usize, m.slots); + for (m.ids, 0..) |*v, i| v.* = if (m.paged) 0 else @intCast(i); + @memset(m.caps, if (m.paged) 0 else m.context); + try m.d.check(m.d.api.cuMemcpyHtoDAsync_v2(m.table, m.ids.ptr, m.ids.len * 4, s.handle), "cuMemcpyHtoDAsync"); + try (DeviceBuffer{ .d = m.d, .ptr = m.states, .len = m.slots * slot_elems * 2, .borrowed = true }).fill8(0, s.handle); + try s.synchronize(); + } + + fn build(m: *Model, s: cuda.Stream) !void { + if (!m.paged) m.pages = m.slots * m.perSlot(); + if (m.pages > 0) m.arena = try m.alloc(m.pages * page_elems * 2); + try m.makeSlots(s); + m.inv = try m.alloc(c.head_dim / 2 * 4); + const r = prompt_rows; + const h = c.hidden; + const pairs = r * c.topk; + const items = cuda.grouped.maxItems(pairs, c.experts, 16); + const parts = round_rows * m.chunks() * c.query_heads; + m.s = .{ + .ids = try m.alloc(r * 4), + .pos = try m.alloc(r * 4), + .slot = try m.alloc(r * 4), + .off = try m.alloc(r * 4), + .meta = try m.alloc(round_rows * 4), + .commit_slot = try m.alloc(round_rows * 4), + .commit_row = try m.alloc(round_rows * 4), + .h = try m.alloc(r * h * 2), + .x = try m.alloc(r * h * 2), + .xs = try m.alloc(r * (c.intermediate / 64) * 4), + .r = try m.alloc(r * h * 2), + .wide = try m.alloc(r * 2 * c.intermediate * 2), + .y = try m.alloc(r * c.intermediate * 2), + .q = try m.alloc(r * h * 2), + .kv = try m.alloc(r * 2 * kv_width * 2), + .picks = try m.alloc(pairs * 4), + .wts = try m.alloc(pairs * 4), + .members = try m.alloc(pairs * 4), + .items = try m.alloc(3 * items * 4), + .counts = try m.alloc(2 * 4), + .rank = try m.alloc(pairs * 4), + .hist = try m.alloc(((pairs + 1023) / 1024) * c.experts * 4), + .eact = try m.alloc(pairs * c.moe_intermediate * 2), + .ey = try m.alloc(pairs * h * 4), + .po = try m.alloc(parts * c.head_dim * 4), + .pm = try m.alloc(parts * 4), + .pl = try m.alloc(parts * 4), + .snaps = try m.alloc(conv_layers * round_rows * 2 * h * 2), + .tap = try m.alloc(max_taps * slot_elems * 2), + .logits = try m.alloc(round_rows * c.vocab * 2), + }; + if (m.planned != null) return; + const inv = inverseFrequencies(m.w.config.theta); + try (DeviceBuffer{ .d = m.d, .ptr = m.inv, .len = inv.len * 4, .borrowed = true }).upload(0, std.mem.sliceAsBytes(&inv)); + } + + pub fn deinit(m: *Model) void { + for (m.buffers.items) |*b| b.free(); + m.buffers.deinit(m.gpa); + m.gpa.free(m.ids); + m.gpa.free(m.caps); + } + + /// Attention layer a's key (or value) cache as the kernels take it: its part of page 0. + pub fn cache(m: *const Model, a: usize, value: bool) u64 { + return m.arena + ((2 * a + @intFromBool(value)) * kern.page * kv_width) * 2; + } + + /// Slot s's page table row. + pub fn tableOf(m: *const Model, slot: usize) u64 { + return m.table + slot * m.perSlot() * 4; + } + + /// Zero slot s's conv states, queued on `s` (a new stream starts from no history). + pub fn clearSlot(m: *const Model, s: cuda.Stream, slot: usize) !void { + try (DeviceBuffer{ .d = m.d, .ptr = m.states + slot * slot_elems * 2, .len = slot_elems * 2, .borrowed = true }).fill8(0, s.handle); + } +}; + +/// RoPE's inverse frequencies as torch's float32 rows: 1 / theta^(2i / 64), the power in float64 rounded once. +pub fn inverseFrequencies(theta: f32) [c.head_dim / 2]f32 { + var out: [c.head_dim / 2]f32 = undefined; + for (&out, 0..) |*v, i| { + const e: f32 = @as(f32, @floatFromInt(2 * i)) / @as(f32, c.head_dim); + const pw: f32 = @floatCast(std.math.pow(f64, theta, e)); + v.* = 1.0 / pw; + } + return out; +} + +test "a page holds 512 positions of the ten attention layers' keys and values: 10 MiB" { + try std.testing.expectEqual(@as(usize, 10 << 20), page_elems * 2); + try std.testing.expectEqual(@as(usize, 4 * (10 << 20) + slot_elems * 2), Model.streamBytes(1537)); + try std.testing.expectEqual(@as(usize, 2048), Model.positions(1537)); + const inv = inverseFrequencies(1e7); + try std.testing.expectEqual(@as(f32, 1), inv[0]); + try std.testing.expect(@abs(inv[31] - 1.0 / std.math.pow(f64, 1e7, 62.0 / 64.0)) < 1e-12); +} + +test "the pool's layout: 20 KiB a token, a stream's conv states and table row fixed, no snapshot slots" { + var m: Model = undefined; + m.context = Model.positions(32768); + const l = m.layout(); + try std.testing.expectEqual(@as(u64, 10 << 20), l.pageBytes()); + try std.testing.expectEqual(@as(u64, 20 << 10), l.token_bytes); + try std.testing.expectEqual(@as(u64, slot_elems * 2 + 64 * 4), l.fixed_bytes); + // 40 GiB past the weights, 8 streams at 32k, the default 16 GiB cache: a window for the first, the rest pages + const p = try kv.plan(l, .{ .room = 40 << 30, .engine = 0, .lanes = 8, .cache = 16 << 30, .window = 32768 + 16, .engine_paged = true }); + try std.testing.expectEqual(@as(u32, 8), p.lanes); + try std.testing.expectEqual(@as(usize, 0), p.snaps); + try std.testing.expect(p.pages * l.pageBytes() + p.copied + 7 * l.fixed_bytes <= 40 << 30); + try std.testing.expect(p.pages >= l.pages(32768 + 16)); +} + +test "leases map a slot's pages in position order; its reach ends at them; an id past the arena is refused" { + var ids: [8]u32 = @splat(0); + var caps = [_]usize{ 0, 0 }; + var m: Model = undefined; + m.context = 4 * kern.page; + m.slots = 2; + m.pages = 16; + m.ids = &ids; + m.caps = ∩︀ + var pool = try kv.Pool.init(std.testing.allocator, 16, page_elems * 2); + defer pool.deinit(); + var lease: kv.Lease = .{}; + try pool.promise(&lease, 3); + const Rec = struct { + m: *Model, + fn write(ctx: *anyopaque, got: []const u32, from: usize) anyerror!void { + const r: *@This() = @ptrCast(@alignCast(ctx)); + _ = try r.m.note(1, got[from..], from); + } + }; + var rec: Rec = .{ .m = &m }; + try pool.cover(&lease, 3, .{ .ctx = &rec, .write = Rec.write }); + try std.testing.expectEqualSlices(u32, lease.ids.items, ids[4..7]); + try std.testing.expectEqual(@as(usize, 3 * kern.page), caps[1]); + try std.testing.expectEqual(@as(usize, 13), pool.available()); + pool.finish(&lease); + try std.testing.expect(pool.idle() and pool.available() == 16); + try std.testing.expectError(error.BadChain, m.note(0, &.{16}, 0)); + try std.testing.expectError(error.BadChain, m.note(0, &.{ 1, 2, 3, 4, 5 }, 0)); +} diff --git a/zig/src/families/lfm2_moe/cuda_native.zig b/zig/src/families/lfm2_moe/cuda_native.zig new file mode 100644 index 0000000000..e13f5f54f9 --- /dev/null +++ b/zig/src/families/lfm2_moe/cuda_native.zig @@ -0,0 +1,85 @@ +//! LFM2.5-24B-A2B (MLX affine 4-bit) behind the native CUDA registry: liquid's glue, this family's checks. +const std = @import("std"); +const liquid = @import("liquid_cuda"); +const cfg = @import("config.zig"); +const mdl = @import("cuda_model.zig"); +const be = @import("cuda_backend.zig"); +const snap = @import("cuda_snap.zig"); + +const N = liquid.native.Native(struct { + pub const wts = @import("cuda_weights.zig"); + pub const kern = @import("cuda_kernels.zig"); + pub const mdl = @import("cuda_model.zig"); + pub const be = @import("cuda_backend.zig"); + pub const snap = @import("cuda_snap.zig"); + pub const pages_env = "TENSORFOLD_LFM2_KV_PAGES"; + pub const check = checkLayers; +}); + +pub const model_type = "lfm2_moe"; +pub const formats: []const []const u8 = &.{"mlx-q4g64"}; +pub const default_context: i64 = 32768; +pub const max_segments: u32 = 1; +pub const explains_open = true; +pub const prompt_rows = N.prompt_rows; +pub const snap_lookahead = N.snap_lookahead; +pub const window_rows = N.window_rows; +pub const Room = N.Room; +pub const Options = N.Options; +pub const Loaded = N.Loaded; +pub const open = N.open; +pub const attach = N.attach; +pub const fits = N.fits; +pub const snap_bytes = N.snap_bytes; +pub const snap_ready = N.snap_ready; +pub const snap_charged = N.snap_charged; +pub const snap_spare = N.snap_spare; +pub const snap_reuses = N.snap_reuses; +pub const snap_save = N.snap_save; +pub const snap_restore = N.snap_restore; +pub const snap_drop = N.snap_drop; + +/// The attention layers the kernels' page arena holds: the admitted 10 of 40 (conv the rest). +fn checkLayers(config: *const cfg.Config) !void { + var n: usize = 0; + for (config.attention) |a| n += @intFromBool(a); + if (n != mdl.attention_layers) return error.UnsupportedLfmGeometry; +} + +pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { + return switch (err) { + error.PastContext => "the prompt and its reply exceed this server's context window: shorten it or lower max_tokens", + error.NoFreeSlot => "every stream's caches are in use: retry once another request ends", + error.NoDraftHead => "LFM2.5-24B-A2B on CUDA has no draft head: it drafts from the prompt's own text", + error.LfmTextOnly => "LFM2.5-24B-A2B is a text model: this request has images", + error.StructuresUnsupported => "structured output is not built for LFM2.5-24B-A2B on CUDA yet", + error.NonFiniteLogits => "the model's logits went non-finite at this row: the request is stopped rather than answered from them", + error.CapturedKernels => "LFM2.5-24B-A2B runs on the engine's own kernels: unset TENSORFOLD_CUDA_KERNELS", + error.UnsupportedLfmQuantization => "LFM2.5-24B-A2B on CUDA reads MLX affine 4-bit checkpoints in groups of 64 (the embedding at 4, 6 or 8 bits): the 6- and 8-bit exports are served on Metal", + error.UnsupportedLfmGeometry, error.UnsupportedLfmConfig, error.UnsupportedLfmRotary => "this lfm2_moe checkpoint's geometry is not LFM2.5-24B-A2B's, which this engine serves", + error.ParallelPastRound => std.fmt.comptimePrint("--parallel: LFM2.5-24B-A2B on CUDA holds at most {d} streams (a round's rows): serve --parallel {d} or fewer, or auto", .{ be.round_rows, be.round_rows }), + error.ParallelPastMemory => "--parallel: the CUDA memory budget fits fewer LFM2.5-24B-A2B streams than asked (each stream's caches are allocated at load): lower --parallel or --context, serve --parallel auto, or adjust TENSORFOLD_MEMORY_RESERVE_GIB / TENSORFOLD_CUDA_MEMORY_LIMIT_GB", + error.NoStreamFits => "the CUDA memory budget fits no LFM2.5-24B-A2B stream after the weights: lower --context, free device memory, or adjust TENSORFOLD_MEMORY_RESERVE_GIB / TENSORFOLD_CUDA_MEMORY_LIMIT_GB", + else => null, + }; +} + +test "every refusal of its own is worded, the rest left to the server" { + for ([_]anyerror{ error.PastContext, error.NoFreeSlot, error.NoDraftHead, error.LfmTextOnly, error.StructuresUnsupported, error.NonFiniteLogits, error.CapturedKernels, error.UnsupportedLfmQuantization, error.UnsupportedLfmGeometry, error.ParallelPastRound, error.ParallelPastMemory, error.NoStreamFits }) |e| { + try std.testing.expect(explain(null, e) != null); + } + try std.testing.expect(explain(null, error.OutOfMemory) == null); +} + +test "the admitted layer pattern has the arena's ten attention layers" { + var config = try cfg.parse(std.testing.allocator, cfg.fixture); + try checkLayers(&config); + config.attention[0] = true; + try std.testing.expectError(error.UnsupportedLfmGeometry, checkLayers(&config)); +} + +test "the prompt cache's hooks compile against the store's vtable" { + const pc: struct { bytes: *const fn (*anyopaque, u32) u64, save: *const fn (*anyopaque, ?*anyopaque, u32) anyerror!*anyopaque, restore: *const fn (*anyopaque, ?*anyopaque, *anyopaque) anyerror!void, drop: *const fn (*anyopaque, *anyopaque) void } = .{ .bytes = snap_bytes, .save = snap_save, .restore = snap_restore, .drop = snap_drop }; + try std.testing.expect(@intFromPtr(pc.save) != 0); + try std.testing.expectEqual(snap.bytes(100), snap_bytes(undefined, 100)); +} diff --git a/zig/src/families/lfm2_moe/cuda_snap.zig b/zig/src/families/lfm2_moe/cuda_snap.zig new file mode 100644 index 0000000000..3618d6d10c --- /dev/null +++ b/zig/src/families/lfm2_moe/cuda_snap.zig @@ -0,0 +1,70 @@ +//! LFM2.5-24B-A2B's kept states: liquid's over this model (the conv states, then the ten attention layers' pages). +const std = @import("std"); +const liquid = @import("liquid_cuda"); +const mdl = @import("cuda_model.zig"); +const kern = @import("cuda_kernels.zig"); + +const S = liquid.snap.States(mdl); +pub const Saved = S.Saved; +pub const bytes = S.bytes; +pub const save = S.save; +pub const restore = S.restore; +pub const drop = S.drop; +pub const page = liquid.snap.page; + +const page_bytes = mdl.page_elems * 2; +const row_bytes = mdl.kv_width * 2; +const state_bytes = mdl.slot_elems * 2; + +const testing = std.testing; + +const Count = struct { + runs: *std.ArrayList([3]usize), + fn add(c: Count, dev: u64, off: usize, len: usize) anyerror!void { + try c.runs.append(testing.allocator, .{ dev, off, len }); + } +}; + +fn fake(ids: []u32, caps: []usize) mdl.Model { + var m: mdl.Model = undefined; + m.context = 4 * kern.page; + m.slots = caps.len; + m.arena = 1 << 40; + m.states = 1 << 36; + m.ids = ids; + m.caps = caps; + return m; +} + +test "a state: the conv states, consecutive whole pages as one copy, the last page's rows a layer's K and V" { + // slot 1 leases pages 7, 8, 3, 9: positions 0..1023 one run, 1024..1535 a page, 1536.. a partial page + var ids = [_]u32{ 0, 0, 0, 0, 7, 8, 3, 9 }; + var caps = [_]usize{ 0, 4 * kern.page }; + const m = fake(&ids, &caps); + const at = 3 * kern.page + 100; + var runs: std.ArrayList([3]usize) = .empty; + defer runs.deinit(testing.allocator); + try S.each(&m, 1, at, null, Count{ .runs = &runs }, Count.add); + try testing.expectEqual(@as(usize, 1 + 2 + 20), runs.items.len); + try testing.expectEqual([3]usize{ (1 << 36) + state_bytes, 0, state_bytes }, runs.items[0]); + try testing.expectEqual([3]usize{ (1 << 40) + 7 * page_bytes, state_bytes, 2 * page_bytes }, runs.items[1]); + try testing.expectEqual([3]usize{ (1 << 40) + 3 * page_bytes, state_bytes + 2 * page_bytes, page_bytes }, runs.items[2]); + const v9 = (1 << 40) + 9 * page_bytes + kern.page * row_bytes; // page 9, layer 0's values + try testing.expectEqual([3]usize{ v9, state_bytes + 3 * page_bytes + 100 * row_bytes, 100 * row_bytes }, runs.items[4]); + var total: usize = 0; + for (runs.items) |r| total += r[2]; + try testing.expectEqual(state_bytes + at * 20 * 1024, total); + try testing.expectEqual((total + page - 1) / page * page, S.bytes(at)); +} + +test "a state on whole pages has no partial copies; one past a slot's pages is refused" { + var ids = [_]u32{ 4, 5, 6, 7 }; + var caps = [_]usize{2 * kern.page}; + const m = fake(&ids, &caps); + var runs: std.ArrayList([3]usize) = .empty; + defer runs.deinit(testing.allocator); + try S.each(&m, 0, 2 * kern.page, null, Count{ .runs = &runs }, Count.add); + try testing.expectEqual(@as(usize, 2), runs.items.len); + try testing.expectEqual(@as(usize, 2 * page_bytes), runs.items[1][2]); + try testing.expectError(error.PastPages, S.each(&m, 0, 2 * kern.page + 1, null, Count{ .runs = &runs }, Count.add)); +} diff --git a/zig/src/families/lfm2_moe/cuda_weights.zig b/zig/src/families/lfm2_moe/cuda_weights.zig new file mode 100644 index 0000000000..05f00a9aff --- /dev/null +++ b/zig/src/families/lfm2_moe/cuda_weights.zig @@ -0,0 +1,239 @@ +//! LFM2.5-24B-A2B's weights on CUDA: affine-4 projections, expert pair/down tables, the tied head, bf16 norms. +const std = @import("std"); +const core = @import("core"); +const cuda = @import("cuda"); +const c = @import("config.zig"); +const kern = @import("cuda_kernels.zig"); + +const DeviceBuffer = cuda.DeviceBuffer; +const Affine4 = cuda.qlinear.Affine4; +const AffineN = cuda.qlinear.AffineN; +const A4 = cuda.experts.Affine4Experts; + +/// A projection at its MLX width: affine-4 tiles, or the tiles of every other width (the head). +pub const Proj = union(enum) { + a4: Affine4, + an: AffineN, + + pub fn n(p: Proj) usize { + return switch (p) { + inline else => |q| q.n, + }; + } + + pub fn k(p: Proj) usize { + return switch (p) { + inline else => |q| q.k, + }; + } +}; + +pub const Conv = struct { in: Proj, out: Proj, weight: u64 }; +/// q, k and v stacked into one projection; the q and k norms bf16 [64] each, q's then k's. +pub const Attention = struct { qkv: Proj, out: Proj, norms: u64 }; +pub const Dense = struct { gate_up: Proj, down: Proj }; +pub const Moe = struct { router: u64, bias: u64, experts: cuda.experts.Layer }; + +pub const Layer = struct { + operator_norm: u64, + ffn_norm: u64, + mixer: union(enum) { conv: Conv, attention: Attention }, + ffn: union(enum) { dense: Dense, moe: Moe }, +}; + +pub const Weights = struct { + gpa: std.mem.Allocator, + config: c.Config, + /// The embedding as MLX stores it: words u32 [V, H bits / 32], scales and biases [V, H / 64]; also the head. + embed_w: u64 = 0, + embed_s: u64 = 0, + embed_b: u64 = 0, + head: Proj = undefined, + final_norm: u64 = 0, + layers: []Layer = &.{}, + buffers: std.ArrayList(DeviceBuffer) = .empty, + + pub fn deinit(w: *Weights) void { + for (w.buffers.items) |*b| b.free(); + w.buffers.deinit(w.gpa); + w.gpa.free(w.layers); + w.* = undefined; + } + + fn alloc(w: *Weights, d: *const cuda.Driver, bytes: usize) !u64 { + const b = try DeviceBuffer.alloc(d, @max(bytes, 256)); + errdefer { + var x = b; + x.free(); + } + try w.buffers.append(w.gpa, b); + return b.ptr; + } + + fn upload(w: *Weights, d: *const cuda.Driver, bytes: []const u8) !u64 { + const b = try DeviceBuffer.fromHost(d, bytes); + errdefer { + var x = b; + x.free(); + } + try w.buffers.append(w.gpa, b); + return b.ptr; + } +}; + +const Ctx = struct { + gpa: std.mem.Allocator, + w: *Weights, + d: *const cuda.Driver, + s: cuda.Stream, + k: *const kern.Kernels, + ck: *core.checkpoint.Checkpoint, + stage: DeviceBuffer, + + fn tensor(x: *Ctx, name: []const u8, dtype: core.checkpoint.DType, shape: []const usize) !u64 { + return x.w.upload(x.d, (try x.ck.expect(name, dtype, shape)).bytes); + } + + /// Parts' bytes one after another into the stage; their device addresses. + fn staged(x: *Ctx, parts: []const []const u8, out: []u64) !void { + var at: usize = 0; + for (parts, out) |p, *o| { + if (at + p.len > x.stage.len) return error.StageTooSmall; + try x.stage.upload(at, p); + o.* = x.stage.ptr + at; + at += (p.len + 255) / 256 * 256; + } + } + + /// MLX affine projections stacked along n (each [n_i, k] at `bits`) packed as one into qlinear's tiles. + fn linear(x: *Ctx, parts: []const []const u8, ns: []const usize, k: usize, bits: u8) !Proj { + var n: usize = 0; + for (ns) |v| n += v; + var buf: [192]u8 = undefined; + const specs = [_]struct { []const u8, core.checkpoint.DType, usize }{ .{ "weight", .u32, k * bits / 32 }, .{ "scales", .bf16, k / 64 }, .{ "biases", .bf16, k / 64 } }; + var raw: [3]u64 = undefined; + var at: usize = 0; + for (specs, &raw) |sp, *r| { + r.* = x.stage.ptr + at; + for (parts, ns) |prefix, rows| { + const t = try x.ck.expect(try std.fmt.bufPrint(&buf, "{s}.{s}", .{ prefix, sp[0] }), sp[1], &.{ rows, sp[2] }); + if (at + t.bytes.len > x.stage.len) return error.StageTooSmall; + try x.stage.upload(at, t.bytes); + at += t.bytes.len; + } + at = (at + 255) / 256 * 256; + } + const served = cuda.qlinear.servedBits(bits); + const out: Proj = if (served == 4) blk: { + const l = Affine4.layout(n, k); + const q: Affine4 = .{ .w = try x.w.alloc(x.d, l.words), .s = try x.w.alloc(x.d, l.scales), .b = try x.w.alloc(x.d, l.scales), .n = n, .k = k, .npad = l.npad }; + try x.k.affine.pack(x.s, raw[0], raw[1], raw[2], q); + break :blk .{ .a4 = q }; + } else blk: { + const l = AffineN.layout(n, k, served); + const q: AffineN = .{ .w = try x.w.alloc(x.d, l.words), .s = try x.w.alloc(x.d, l.scales), .b = try x.w.alloc(x.d, l.scales), .n = n, .k = k, .npad = l.npad, .bits = served }; + try x.k.wide.packWords(x.s, raw[0], bits, n, k, q.w); + try x.k.wide.packTable(x.s, raw[1], n, k, q.npad, q.s); + try x.k.wide.packTable(x.s, raw[2], n, k, q.npad, q.b); + break :blk .{ .an = q }; + }; + try x.s.synchronize(); // the next projection's staging overwrites this one's + return out; + } + + /// The 64 experts' `name` ([64, n, k] at 4 bits) packed as `out`'s mi-th of m interleaved tables. + fn experts(x: *Ctx, prefix: []const u8, name: []const u8, n: usize, k: usize, out: u64, m: usize, mi: usize) !void { + var buf: [192]u8 = undefined; + var parts: [3][]const u8 = undefined; + const specs = [_]struct { []const u8, core.checkpoint.DType, usize }{ .{ "weight", .u32, k / 8 }, .{ "scales", .bf16, k / 64 }, .{ "biases", .bf16, k / 64 } }; + for (specs, &parts) |sp, *p| p.* = (try x.ck.expect(try std.fmt.bufPrint(&buf, "{s}{s}.{s}", .{ prefix, name, sp[0] }), sp[1], &.{ c.experts, n, sp[2] })).bytes; + var raw: [3]u64 = undefined; + try x.staged(&parts, &raw); + try x.k.experts.packInto(x.s, raw[0], raw[1], raw[2], out, c.experts, n, k, m, mi); + try x.s.synchronize(); + } +}; + +/// Bytes the stage needs: the largest of the head's table as stored and one expert table (words, scales, biases). +pub fn stageBytes(embed_bits: u8) usize { + const head = c.vocab * c.hidden * @as(usize, embed_bits) / 8 + 2 * c.vocab * (c.hidden / 64) * 2; + const table = c.experts * c.moe_intermediate * (c.hidden / 2 + 2 * (c.hidden / 64) * 2); + return @max(head, table) + 4 * 256; +} + +/// Every tensor checked against the admitted geometry; projections and experts at 4 bits. `layers` caps them. +pub fn load(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, s: cuda.Stream, k: *const kern.Kernels, dir: []const u8, layers: ?usize) !Weights { + const config = try c.Config.read(gpa, io, dir); + if (config.bits != 4) return error.UnsupportedLfmQuantization; // the experts' kernels read 4-bit tables + var w: Weights = .{ .gpa = gpa, .config = config }; + errdefer w.deinit(); + var ck = try core.checkpoint.Checkpoint.openModel(gpa, io, dir); + defer ck.close(); + const h = c.hidden; + const eb = config.embed_bits; + var stage = try DeviceBuffer.alloc(d, stageBytes(eb)); + defer stage.free(); + var x: Ctx = .{ .gpa = gpa, .w = &w, .d = d, .s = s, .k = k, .ck = &ck, .stage = stage }; + w.embed_w = try x.tensor("model.embed_tokens.weight", .u32, &.{ c.vocab, h * @as(usize, eb) / 32 }); + w.embed_s = try x.tensor("model.embed_tokens.scales", .bf16, &.{ c.vocab, h / 64 }); + w.embed_b = try x.tensor("model.embed_tokens.biases", .bf16, &.{ c.vocab, h / 64 }); + w.head = try x.linear(&.{"model.embed_tokens"}, &.{c.vocab}, h, eb); + w.final_norm = try x.tensor("model.embedding_norm.weight", .bf16, &.{h}); + const n = @min(layers orelse c.layers, c.layers); + w.layers = try gpa.alloc(Layer, n); + var b1: [128]u8 = undefined; + var b2: [128]u8 = undefined; + var b3: [128]u8 = undefined; + for (w.layers, 0..) |*l, i| { + var pb: [64]u8 = undefined; + const p = try std.fmt.bufPrint(&pb, "model.layers.{d}.", .{i}); + const name = struct { + fn of(b: []u8, prefix: []const u8, suffix: []const u8) ![]const u8 { + return std.fmt.bufPrint(b, "{s}{s}", .{ prefix, suffix }); + } + }.of; + l.operator_norm = try x.tensor(try name(&b1, p, "operator_norm.weight"), .bf16, &.{h}); + l.ffn_norm = try x.tensor(try name(&b1, p, "ffn_norm.weight"), .bf16, &.{h}); + if (config.attention[i]) { + const qw = c.query_heads * c.head_dim; + const kw = c.kv_heads * c.head_dim; + const norms = try w.alloc(d, 2 * c.head_dim * 2); + const qn = try ck.expect(try name(&b1, p, "self_attn.q_layernorm.weight"), .bf16, &.{c.head_dim}); + const kn = try ck.expect(try name(&b1, p, "self_attn.k_layernorm.weight"), .bf16, &.{c.head_dim}); + try (DeviceBuffer{ .d = d, .ptr = norms, .len = 4 * c.head_dim, .borrowed = true }).upload(0, qn.bytes); + try (DeviceBuffer{ .d = d, .ptr = norms, .len = 4 * c.head_dim, .borrowed = true }).upload(2 * c.head_dim, kn.bytes); + l.mixer = .{ .attention = .{ + .qkv = try x.linear(&.{ try name(&b1, p, "self_attn.q_proj"), try name(&b2, p, "self_attn.k_proj"), try name(&b3, p, "self_attn.v_proj") }, &.{ qw, kw, kw }, h, 4), + .out = try x.linear(&.{try name(&b1, p, "self_attn.out_proj")}, &.{h}, qw, 4), + .norms = norms, + } }; + } else l.mixer = .{ .conv = .{ + .in = try x.linear(&.{try name(&b1, p, "conv.in_proj")}, &.{3 * h}, h, 4), + .out = try x.linear(&.{try name(&b1, p, "conv.out_proj")}, &.{h}, h, 4), + .weight = try x.tensor(try name(&b1, p, "conv.conv.weight"), .bf16, &.{ h, c.conv_taps, 1 }), + } }; + if (config.moe(i)) { + const ni = c.moe_intermediate; + const pair = try w.alloc(d, A4.pairBytes(c.experts, ni, h)); + const down = try w.alloc(d, A4.bytes(c.experts, h, ni)); + const ff = try name(&b2, p, "feed_forward.switch_mlp."); + try x.experts(ff, "gate_proj", ni, h, pair, 2, 0); + try x.experts(ff, "up_proj", ni, h, pair, 2, 1); + try x.experts(ff, "down_proj", h, ni, down, 1, 0); + l.ffn = .{ .moe = .{ + .router = try x.tensor(try name(&b1, p, "feed_forward.gate.weight"), .f32, &.{ c.experts, h }), + .bias = try x.tensor(try name(&b1, p, "feed_forward.expert_bias"), .f32, &.{c.experts}), + .experts = .{ .format = .affine4, .up = pair, .down = down, .width = ni, .dims = h, .experts = c.experts, .swiglu = true }, + } }; + } else l.ffn = .{ .dense = .{ + .gate_up = try x.linear(&.{ try name(&b1, p, "feed_forward.gate_proj"), try name(&b2, p, "feed_forward.up_proj") }, &.{ c.intermediate, c.intermediate }, h, 4), + .down = try x.linear(&.{try name(&b1, p, "feed_forward.down_proj")}, &.{h}, c.intermediate, 4), + } }; + } + return w; +} + +test "the stage holds the 6-bit head's table and one expert table" { + try std.testing.expect(stageBytes(6) >= c.vocab * c.hidden * 6 / 8); + try std.testing.expect(stageBytes(4) >= c.experts * c.moe_intermediate * c.hidden / 2); +} diff --git a/zig/src/families/lfm3/config.zig b/zig/src/families/lfm3/config.zig new file mode 100644 index 0000000000..baf2b5dc7b --- /dev/null +++ b/zig/src/families/lfm3/config.zig @@ -0,0 +1,401 @@ +//! Admission for LFM3-1.7B and LFM3-VL-2B MLX checkpoints; this parser does not advertise an engine. +const std = @import("std"); + +pub const hidden = 2048; +pub const vocab = 128000; +pub const layers = 24; +pub const intermediate = 7680; +pub const conv_taps = 7; +pub const query_heads = 16; +pub const kv_heads = 4; +pub const head_dim = 128; +pub const sliding_window = 256; +pub const eps: f32 = 1e-5; +pub const Kind = enum { text, vision_language }; +pub const LayerType = enum { conv, full_attention, sliding_attention }; +pub const Rotary = struct { theta: f32, partial: f32 }; +pub const schedule: [layers]LayerType = .{ + .conv, .conv, .full_attention, .conv, .sliding_attention, .conv, + .full_attention, .conv, .sliding_attention, .conv, .full_attention, .conv, + .sliding_attention, .conv, .full_attention, .conv, .sliding_attention, .conv, + .full_attention, .conv, .sliding_attention, .conv, .full_attention, .conv, +}; + +pub const Width = struct { bits: u8, group: u16 }; +/// The projections' storage: MLX affine at `quant`'s widths, ModelOpt NVFP4 (FFN) or block FP8 (FFN, q/k/v/o). +pub const Scheme = enum { mlx_affine, modelopt_nvfp4, fp8_block }; +/// The width `quant` reports for a ModelOpt NVFP4 export (its FFN): 4 bits in groups of 16. +pub const nvfp4_width: Width = .{ .bits = 4, .group = 16 }; +/// The width `quant` reports for a block FP8 export: 8 bits, one scale a 128 x 128 block. +pub const fp8_width: Width = .{ .bits = 8, .group = 128 }; +/// The rows and inputs a block FP8 scale covers. +pub const fp8_block = 128; +pub const Quant = struct { + global: Width, + modules: std.StringHashMapUnmanaged(Width) = .empty, + + pub fn width(self: *const Quant, key: []const u8) Width { + return self.modules.get(canonical(key)) orelse self.global; + } +}; + +fn canonical(key: []const u8) []const u8 { + var k = key; + if (std.mem.endsWith(u8, k, ".weight")) k = k[0 .. k.len - 7]; + for ([_][]const u8{ "model.language_model.", "language_model.", "text_model.", "model." }) |prefix| { + if (std.mem.startsWith(u8, k, prefix)) k = k[prefix.len..]; + } + return k; +} + +fn projection(key: []const u8) bool { + const k = canonical(key); + if (std.mem.eql(u8, k, "embed_tokens")) return true; + if (!std.mem.startsWith(u8, k, "layers.")) return false; + const tail = k[7..]; + const dot = std.mem.indexOfScalar(u8, tail, '.') orelse return false; + if (dot == 0 or (dot > 1 and tail[0] == '0')) return false; + for (tail[0..dot]) |digit| if (digit < '0' or digit > '9') return false; + const index = std.fmt.parseInt(usize, tail[0..dot], 10) catch return false; + if (index >= layers) return false; + const name = tail[dot + 1 ..]; + const mixer = schedule[index]; + const names: []const []const u8 = if (mixer == .conv) + &.{ "conv.in_proj", "conv.out_proj", "feed_forward.w1", "feed_forward.w2", "feed_forward.w3" } + else + &.{ "self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.o_proj", "self_attn.gate_proj", "feed_forward.w1", "feed_forward.w2", "feed_forward.w3" }; + for (names) |name_| if (std.mem.eql(u8, name, name_)) return true; + return false; +} + +fn quantWidth(o: std.json.ObjectMap) !Width { + const bits = try integer(o, "bits"); + const group = try integer(o, "group_size"); + if (o.get("mode")) |mode| { + if (mode != .string or !std.mem.eql(u8, mode.string, "affine")) return error.UnsupportedLfm3Quantization; + } + const w: Width = .{ + .bits = std.math.cast(u8, bits) orelse return error.UnsupportedLfm3Quantization, + .group = std.math.cast(u16, group) orelse return error.UnsupportedLfm3Quantization, + }; + try width(w); + return w; +} + +pub const Config = struct { + arena: std.heap.ArenaAllocator, + kind: Kind, + context: usize, + layer_types: [layers]LayerType = schedule, + full_rotary: Rotary = .{ .theta = 1e6, .partial = 0.5 }, + sliding_rotary: Rotary = .{ .theta = 1e4, .partial = 1 }, + scheme: Scheme = .mlx_affine, + quant: Quant, + + pub fn deinit(self: *Config) void { + self.arena.deinit(); + self.* = undefined; + } + + pub fn width(self: *const Config, key: []const u8) !Width { + return self.quant.width(key); + } + + pub fn read(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Config { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 22)); + defer gpa.free(bytes); + return parse(gpa, bytes); + } +}; + +const Object = std.json.ObjectMap; + +fn object(v: std.json.Value) !Object { + return if (v == .object) v.object else error.BadLfm3Config; +} + +fn field(o: Object, key: []const u8) !std.json.Value { + return o.get(key) orelse error.BadLfm3Config; +} + +fn integer(o: Object, key: []const u8) !usize { + const v = try field(o, key); + if (v != .integer or v.integer < 1) return error.BadLfm3Config; + return std.math.cast(usize, v.integer) orelse error.BadLfm3Config; +} + +fn string(o: Object, key: []const u8, expected: []const u8) !void { + const v = try field(o, key); + if (v != .string or !std.mem.eql(u8, v.string, expected)) return error.UnsupportedLfm3Config; +} + +fn boolean(o: Object, key: []const u8, expected: bool) !void { + const v = try field(o, key); + if (v != .bool or v.bool != expected) return error.UnsupportedLfm3Config; +} + +fn number(o: Object, key: []const u8, expected: f64) !void { + const v = try field(o, key); + const actual: f64 = switch (v) { + .integer => |i| @floatFromInt(i), + .float => |f| f, + else => return error.BadLfm3Config, + }; + if (actual != expected) return error.UnsupportedLfm3Config; +} + +fn geometry(o: Object, dims: anytype) !void { + inline for (dims) |d| if (try integer(o, d[0]) != d[1]) return error.UnsupportedLfm3Geometry; +} + +fn layerTypes(o: Object, expected: []const LayerType) !void { + const v = try field(o, "layer_types"); + if (v != .array or v.array.items.len != expected.len) return error.UnsupportedLfm3LayerSchedule; + for (v.array.items, expected) |actual, wanted| { + if (actual != .string or !std.mem.eql(u8, actual.string, @tagName(wanted))) + return error.UnsupportedLfm3LayerSchedule; + } +} + +fn text(o: Object, kind: Kind) !usize { + try string(o, "model_type", "lfm3"); + try geometry(o, .{ + .{ "hidden_size", hidden }, .{ "vocab_size", vocab }, .{ "num_hidden_layers", layers }, + .{ "intermediate_size", intermediate }, .{ "conv_L_cache", conv_taps }, .{ "num_attention_heads", query_heads }, + .{ "num_key_value_heads", kv_heads }, .{ "head_dim", head_dim }, .{ "sliding_window", sliding_window }, + }); + try boolean(o, "conv_bias", false); + try boolean(o, "tie_word_embeddings", true); + try boolean(o, "use_sandwich_norm", true); + try string(o, "dtype", "bfloat16"); + try number(o, "norm_eps", 1e-5); + try layerTypes(o, &schedule); + const rope = try object(try field(o, "rope_parameters")); + if (rope.count() != 2) return error.UnsupportedLfm3Rotary; + const full = try object(try field(rope, "full_attention")); + const sliding = try object(try field(rope, "sliding_attention")); + try rotary(full, 1e6, 0.5); + try rotary(sliding, 1e4, 1); + const context = try integer(o, "max_position_embeddings"); + if (context != @as(usize, if (kind == .text) 262144 else 128000)) return error.UnsupportedLfm3Context; + return context; +} + +fn rotary(o: Object, theta: f64, partial: f64) !void { + try string(o, "rope_type", "default"); + try number(o, "rope_theta", theta); + if (o.contains("partial_rotary_factor")) try number(o, "partial_rotary_factor", partial) else if (partial != 1) + return error.UnsupportedLfm3Rotary; + var it = o.iterator(); + while (it.next()) |e| { + const k = e.key_ptr.*; + if (!std.mem.eql(u8, k, "rope_type") and !std.mem.eql(u8, k, "rope_theta") and + !std.mem.eql(u8, k, "partial_rotary_factor")) return error.UnsupportedLfm3Rotary; + } +} + +fn vision(root: Object) !void { + try boolean(root, "tie_word_embeddings", true); + try boolean(root, "projector_bias", true); + try string(root, "projector_hidden_act", "gelu"); + try number(root, "projector_layer_norm_eps", 1e-6); + try geometry(root, .{ + .{ "downsample_factor", 2 }, .{ "projector_hidden_size", 4096 }, .{ "image_token_id", 125030 }, + .{ "video_token_id", 125029 }, .{ "vision_start_token_id", 125009 }, .{ "vision_end_token_id", 125010 }, + }); + const o = try object(try field(root, "vision_config")); + try string(o, "model_type", "lfm3_vl_vision"); + try string(o, "dtype", "bfloat16"); + try number(o, "norm_eps", 1e-5); + try geometry(o, .{ + .{ "hidden_size", 1024 }, .{ "intermediate_size", 2560 }, .{ "num_hidden_layers", 14 }, + .{ "num_attention_heads", 16 }, .{ "num_key_value_heads", 8 }, .{ "in_channels", 3 }, + .{ "conv_kernel_size", 3 }, .{ "patch_size", 16 }, .{ "spatial_merge_size", 2 }, + .{ "temporal_patch_size", 2 }, + }); + try layerTypes(o, &.{ .conv, .conv, .full_attention, .conv, .full_attention, .conv, .full_attention, .conv, .full_attention, .conv, .full_attention, .conv, .full_attention, .conv }); + const rope = try object(try field(o, "rope_parameters")); + try string(rope, "rope_type", "axial"); + try number(rope, "rope_theta", 1e6); + if (rope.count() != 2) return error.UnsupportedLfm3Rotary; +} + +fn quantization(a: std.mem.Allocator, value: std.json.Value) !Quant { + const o = try object(value); + var it = o.iterator(); + while (it.next()) |e| { + const key = e.key_ptr.*; + if (std.mem.eql(u8, key, "bits") or std.mem.eql(u8, key, "group_size") or std.mem.eql(u8, key, "mode")) continue; + if (e.value_ptr.* != .object) return error.UnsupportedLfm3Quantization; + const entry = e.value_ptr.object; + if (entry.count() < 2 or entry.count() > 3) return error.UnsupportedLfm3Quantization; + var fields = entry.iterator(); + while (fields.next()) |f| { + const k = f.key_ptr.*; + if (!std.mem.eql(u8, k, "bits") and !std.mem.eql(u8, k, "group_size") and !std.mem.eql(u8, k, "mode")) + return error.UnsupportedLfm3Quantization; + } + if (!projection(key)) return error.UnsupportedLfm3Quantization; + } + var q: Quant = .{ .global = try quantWidth(o) }; + it = o.iterator(); + while (it.next()) |e| { + if (e.value_ptr.* != .object) continue; + const w = try quantWidth(e.value_ptr.object); + const slot = try q.modules.getOrPut(a, canonical(e.key_ptr.*)); + if (slot.found_existing and !std.meta.eql(slot.value_ptr.*, w)) return error.UnsupportedLfm3Quantization; + slot.value_ptr.* = w; + } + return q; +} + +/// ModelOpt's NVFP4 text export: true; another quant_method: false; any other recipe or ignore list is refused. +fn modelopt(value: std.json.Value, kind: Kind) !bool { + const o = try object(value); + const method = o.get("quant_method") orelse return false; + if (method != .string or !std.mem.eql(u8, method.string, "modelopt")) return false; + if (kind != .text or !named(o, "quant_algo", "NVFP4")) return error.UnsupportedLfm3Quantization; + const groups = try object(try field(o, "config_groups")); + if (groups.count() != 1) return error.UnsupportedLfm3Quantization; + const group = try object(groups.values()[0]); + try fp4(try object(try field(group, "weights"))); + try fp4(try object(try field(group, "input_activations"))); + const targets = try field(group, "targets"); + if (targets != .array or targets.array.items.len != 1 or targets.array.items[0] != .string or + !std.mem.eql(u8, targets.array.items[0].string, "Linear")) return error.UnsupportedLfm3Quantization; + try excluded(try field(o, "ignore")); + return true; +} + +/// The compressed-tensors block FP8 text export: true; another quant_method: false; any other recipe is refused. +fn compressed(value: std.json.Value, kind: Kind) !bool { + const o = try object(value); + const method = o.get("quant_method") orelse return false; + if (method != .string or !std.mem.eql(u8, method.string, "compressed-tensors")) return false; + if (kind != .text or !named(o, "format", "float-quantized")) return error.UnsupportedLfm3Quantization; + if (o.get("kv_cache_scheme")) |v| if (v != .null) return error.UnsupportedLfm3Quantization; + const groups = try object(try field(o, "config_groups")); + if (groups.count() != 1) return error.UnsupportedLfm3Quantization; + const group = try object(groups.values()[0]); + const weights = try object(try field(group, "weights")); + const inputs = try object(try field(group, "input_activations")); + try fp8(weights, false, "block"); + try fp8(inputs, true, "group"); + const blocks = try field(weights, "block_structure"); + if (blocks != .array or blocks.array.items.len != 2) return error.UnsupportedLfm3Quantization; + for (blocks.array.items) |b| if (b != .integer or b.integer != fp8_block) return error.UnsupportedLfm3Quantization; + const in_group = try field(inputs, "group_size"); + if (in_group != .integer or in_group.integer != fp8_block) return error.UnsupportedLfm3Quantization; + if (group.get("output_activations")) |v| if (v != .null) return error.UnsupportedLfm3Quantization; + const targets = try field(group, "targets"); + if (targets != .array or targets.array.items.len != 1 or targets.array.items[0] != .string or + !std.mem.eql(u8, targets.array.items[0].string, "Linear")) return error.UnsupportedLfm3Quantization; + const ignore = try field(o, "ignore"); + if (ignore != .array or ignore.array.items.len != fp8_ignore.len) return error.UnsupportedLfm3Quantization; + for (ignore.array.items, fp8_ignore) |item, want| { + if (item != .string or !std.mem.eql(u8, item.string, want)) return error.UnsupportedLfm3Quantization; + } + return true; +} + +/// The FP8 export's ignore list as published: the head, the embedding, the convs and the attention gates stay bf16. +const fp8_ignore = [_][]const u8{ "lm_head", "model.embed_tokens", "re:.*\\.gate$", "re:.*\\.(?:conv|short_conv)\\..*", "re:.*\\.self_attn\\.gate_proj$" }; + +/// A weights or input_activations entry: symmetric e4m3 under `strategy`, dynamic for the inputs only. +fn fp8(o: Object, dynamic: bool, strategy: []const u8) !void { + const d = o.get("dynamic") orelse return error.UnsupportedLfm3Quantization; + const bits = o.get("num_bits") orelse return error.UnsupportedLfm3Quantization; + const sym = o.get("symmetric") orelse return error.UnsupportedLfm3Quantization; + if (d != .bool or d.bool != dynamic or bits != .integer or bits.integer != 8 or sym != .bool or !sym.bool or + !named(o, "type", "float") or !named(o, "strategy", strategy)) return error.UnsupportedLfm3Quantization; +} + +fn named(o: Object, key: []const u8, expected: []const u8) bool { + const v = o.get(key) orelse return false; + return v == .string and std.mem.eql(u8, v.string, expected); +} + +/// A weights or input_activations entry: static FP4 (e2m1) in groups of 16. +fn fp4(o: Object) !void { + const dynamic = o.get("dynamic") orelse return error.UnsupportedLfm3Quantization; + const bits = o.get("num_bits") orelse return error.UnsupportedLfm3Quantization; + const group = o.get("group_size") orelse return error.UnsupportedLfm3Quantization; + if (dynamic != .bool or dynamic.bool or bits != .integer or bits.integer != 4 or group != .integer or + group.integer != 16 or !named(o, "type", "float")) return error.UnsupportedLfm3Quantization; +} + +/// The ignore list as the export writes it: the head, the embedding and every layer's mixer (conv* or self_attn*). +fn excluded(value: std.json.Value) !void { + if (value != .array or value.array.items.len != layers + 2) return error.UnsupportedLfm3Quantization; + var seen: [layers + 2]bool = @splat(false); + for (value.array.items) |item| { + if (item != .string) return error.UnsupportedLfm3Quantization; + const at = ignoredAt(item.string) orelse return error.UnsupportedLfm3Quantization; + if (seen[at]) return error.UnsupportedLfm3Quantization; + seen[at] = true; + } +} + +/// Where an ignore entry sits: 0 the head, 1 the embedding, 2 + i layer i's mixer; null for any other name. +fn ignoredAt(name: []const u8) ?usize { + if (std.mem.eql(u8, name, "lm_head")) return 0; + if (std.mem.eql(u8, name, "model.embed_tokens")) return 1; + const prefix = "model.layers."; + if (!std.mem.startsWith(u8, name, prefix)) return null; + const tail = name[prefix.len..]; + const dot = std.mem.indexOfScalar(u8, tail, '.') orelse return null; + if (dot == 0 or (dot > 1 and tail[0] == '0')) return null; + const index = std.fmt.parseInt(usize, tail[0..dot], 10) catch return null; + if (index >= layers) return null; + const mixer: []const u8 = if (schedule[index] == .conv) "conv*" else "self_attn*"; + return if (std.mem.eql(u8, tail[dot + 1 ..], mixer)) 2 + index else null; +} + +fn width(w: Width) !void { + if ((w.bits != 4 and w.bits != 6 and w.bits != 8) or (w.group != 32 and w.group != 64)) + return error.UnsupportedLfm3Quantization; +} + +fn equalQuant(a: Quant, b: Quant) bool { + if (!std.meta.eql(a.global, b.global) or a.modules.count() != b.modules.count()) return false; + var it = a.modules.iterator(); + while (it.next()) |e| { + const other = b.modules.get(e.key_ptr.*) orelse return false; + if (!std.meta.eql(e.value_ptr.*, other)) return false; + } + return true; +} + +/// The backends a decoder of this admission runs on. +pub const Backend = enum { metal, cuda }; + +/// Whether `backend` serves this scheme: NVFP4 and block FP8 are CUDA's only, so Metal refuses them by name. +pub fn servedOn(config: *const Config, backend: Backend) !void { + if (config.scheme == .modelopt_nvfp4 and backend == .metal) return error.Lfm3ModeloptNvfp4NotOnMetal; + if (config.scheme == .fp8_block and backend == .metal) return error.Lfm3Fp8BlockNotOnMetal; +} + +pub fn parse(gpa: std.mem.Allocator, bytes: []const u8) !Config { + var arena = std.heap.ArenaAllocator.init(gpa); + errdefer arena.deinit(); + const a = arena.allocator(); + const root = try object(try std.json.parseFromSliceLeaky(std.json.Value, a, bytes, .{ .allocate = .alloc_always })); + const tag = try field(root, "model_type"); + if (tag != .string) return error.BadLfm3Config; + const kind: Kind = if (std.mem.eql(u8, tag.string, "lfm3")) .text else if (std.mem.eql(u8, tag.string, "lfm3_vl")) .vision_language else return error.UnsupportedLfm3Config; + const language = if (kind == .text) root else try object(try field(root, "text_config")); + const context = try text(language, kind); + if (kind == .vision_language) try vision(root); + if (root.get("quantization") == null) if (root.get("quantization_config")) |qc| if (try modelopt(qc, kind)) + return .{ .arena = arena, .kind = kind, .context = context, .scheme = .modelopt_nvfp4, .quant = .{ .global = nvfp4_width } }; + if (root.get("quantization") == null) if (root.get("quantization_config")) |qc| if (try compressed(qc, kind)) + return .{ .arena = arena, .kind = kind, .context = context, .scheme = .fp8_block, .quant = .{ .global = fp8_width } }; + const q = try quantization(a, root.get("quantization") orelse root.get("quantization_config") orelse + return error.UnsupportedLfm3Quantization); + if (root.get("quantization_config")) |other| { + if (!equalQuant(q, try quantization(a, other))) return error.UnsupportedLfm3Quantization; + } + return .{ .arena = arena, .kind = kind, .context = context, .quant = q }; +} diff --git a/zig/src/families/lfm3/config_fp8_test.zig b/zig/src/families/lfm3/config_fp8_test.zig new file mode 100644 index 0000000000..56398cae74 --- /dev/null +++ b/zig/src/families/lfm3/config_fp8_test.zig @@ -0,0 +1,160 @@ +//! The block FP8 export's admission, against its published config.json. +const std = @import("std"); +const c = @import("config.zig"); + +test "the block FP8 export parses as its own scheme: FP8 FFN and q/k/v/o in 128 x 128 blocks" { + var config = try c.parse(std.testing.allocator, fp8); + defer config.deinit(); + try std.testing.expectEqual(c.Scheme.fp8_block, config.scheme); + try std.testing.expectEqual(c.fp8_width, try config.width("model.layers.2.self_attn.q_proj")); + try std.testing.expectEqual(@as(usize, 262144), config.context); + try std.testing.expectError(error.Lfm3Fp8BlockNotOnMetal, c.servedOn(&config, .metal)); + try c.servedOn(&config, .cuda); +} + +test "an FP8 export with another block, static inputs, an extra ignore or a kv-cache scheme is refused" { + const edits = [_][2][]const u8{ + .{ "128,\n 128\n", "64,\n 128\n" }, + .{ "\"dynamic\": true", "\"dynamic\": false" }, + .{ "\"strategy\": \"block\"", "\"strategy\": \"channel\"" }, + .{ "\"group_size\": 128", "\"group_size\": 64" }, + .{ "\"lm_head\",", "\"lm_head\",\n \"re:.*mlp.*\"," }, + .{ "\"kv_cache_scheme\": null", "\"kv_cache_scheme\": {}" }, + .{ "\"format\": \"float-quantized\",\n \"global", "\"format\": \"pack-quantized\",\n \"global" }, + .{ "\"Linear\"", "\"Conv1d\"" }, + }; + for (edits) |e| { + const bytes = try std.mem.replaceOwned(u8, std.testing.allocator, fp8, e[0], e[1]); + defer std.testing.allocator.free(bytes); + try std.testing.expect(!std.mem.eql(u8, bytes, fp8)); + try std.testing.expectError(error.UnsupportedLfm3Quantization, c.parse(std.testing.allocator, bytes)); + } +} + +/// LiquidAI/LFM3-1.7B-FP8 (243955da)'s config.json as published. +pub const fp8 = + \\{ + \\ "architectures": [ + \\ "Lfm3ForCausalLM" + \\ ], + \\ "auto_map": { + \\ "AutoConfig": "configuration_lfm3.Lfm3Config", + \\ "AutoModel": "modeling_lfm3.Lfm3Model", + \\ "AutoModelForCausalLM": "modeling_lfm3.Lfm3ForCausalLM" + \\ }, + \\ "bos_token_id": 124894, + \\ "conv_L_cache": 7, + \\ "conv_bias": false, + \\ "dtype": "bfloat16", + \\ "eos_token_id": 124900, + \\ "head_dim": 128, + \\ "hidden_size": 2048, + \\ "intermediate_size": 7680, + \\ "layer_types": [ + \\ "conv", + \\ "conv", + \\ "full_attention", + \\ "conv", + \\ "sliding_attention", + \\ "conv", + \\ "full_attention", + \\ "conv", + \\ "sliding_attention", + \\ "conv", + \\ "full_attention", + \\ "conv", + \\ "sliding_attention", + \\ "conv", + \\ "full_attention", + \\ "conv", + \\ "sliding_attention", + \\ "conv", + \\ "full_attention", + \\ "conv", + \\ "sliding_attention", + \\ "conv", + \\ "full_attention", + \\ "conv" + \\ ], + \\ "max_position_embeddings": 262144, + \\ "model_type": "lfm3", + \\ "norm_eps": 1e-05, + \\ "num_attention_heads": 16, + \\ "num_hidden_layers": 24, + \\ "num_key_value_heads": 4, + \\ "pad_token_id": 124893, + \\ "quantization_config": { + \\ "config_groups": { + \\ "FP8_BLOCK": { + \\ "format": "float-quantized", + \\ "input_activations": { + \\ "actorder": null, + \\ "block_structure": null, + \\ "dynamic": true, + \\ "group_size": 128, + \\ "num_bits": 8, + \\ "observer": null, + \\ "observer_kwargs": {}, + \\ "scale_dtype": null, + \\ "strategy": "group", + \\ "symmetric": true, + \\ "type": "float", + \\ "zp_dtype": null + \\ }, + \\ "output_activations": null, + \\ "targets": [ + \\ "Linear" + \\ ], + \\ "weights": { + \\ "actorder": null, + \\ "block_structure": [ + \\ 128, + \\ 128 + \\ ], + \\ "dynamic": false, + \\ "group_size": null, + \\ "num_bits": 8, + \\ "observer": null, + \\ "observer_kwargs": {}, + \\ "scale_dtype": null, + \\ "strategy": "block", + \\ "symmetric": true, + \\ "type": "float", + \\ "zp_dtype": null + \\ } + \\ } + \\ }, + \\ "format": "float-quantized", + \\ "global_compression_ratio": null, + \\ "ignore": [ + \\ "lm_head", + \\ "model.embed_tokens", + \\ "re:.*\\.gate$", + \\ "re:.*\\.(?:conv|short_conv)\\..*", + \\ "re:.*\\.self_attn\\.gate_proj$" + \\ ], + \\ "kv_cache_scheme": null, + \\ "quant_method": "compressed-tensors", + \\ "quantization_status": "compressed", + \\ "transform_config": {}, + \\ "version": "0.19.0" + \\ }, + \\ "rope_parameters": { + \\ "full_attention": { + \\ "partial_rotary_factor": 0.5, + \\ "rope_theta": 1000000.0, + \\ "rope_type": "default" + \\ }, + \\ "sliding_attention": { + \\ "rope_theta": 10000, + \\ "rope_type": "default" + \\ } + \\ }, + \\ "sliding_window": 256, + \\ "tie_word_embeddings": true, + \\ "transformers_version": "5.12.1", + \\ "use_cache": true, + \\ "use_sandwich_norm": true, + \\ "vocab_size": 128000 + \\} +; diff --git a/zig/src/families/lfm3/config_test.zig b/zig/src/families/lfm3/config_test.zig new file mode 100644 index 0000000000..8507d94a0f --- /dev/null +++ b/zig/src/families/lfm3/config_test.zig @@ -0,0 +1,477 @@ +//! Checkpoint-derived config fixtures and refusal cases for the first LFM3 admission slice. +const std = @import("std"); +const c = @import("config.zig"); + +test { + _ = @import("config_fp8_test.zig"); +} + +test "text and VL preserve their distinct context and actual mixed embedding width" { + var t = try c.parse(std.testing.allocator, text); + defer t.deinit(); + var v = try c.parse(std.testing.allocator, vl); + defer v.deinit(); + try std.testing.expectEqual(c.Kind.text, t.kind); + try std.testing.expectEqual(c.Kind.vision_language, v.kind); + try std.testing.expectEqual(@as(usize, 262144), t.context); + try std.testing.expectEqual(@as(usize, 128000), v.context); + try std.testing.expectEqual(c.Width{ .bits = 4, .group = 64 }, try t.width("model.embed_tokens")); + try std.testing.expectEqual(c.Width{ .bits = 8, .group = 32 }, try v.width("model.embed_tokens.weight")); + try std.testing.expectEqual(c.Width{ .bits = 4, .group = 32 }, try v.width("layers.0.conv.in_proj")); + try std.testing.expectEqual(c.LayerType.sliding_attention, t.layer_types[4]); + try std.testing.expectEqual(@as(f32, 0.5), t.full_rotary.partial); + try std.testing.expectEqual(@as(f32, 1), t.sliding_rotary.partial); +} + +test "6 and 8 bit affine variants use the same config parser" { + inline for (.{ "6", "8" }) |bits| { + const replacement = "\"bits\": " ++ bits; + const bytes = try std.mem.replaceOwned(u8, std.testing.allocator, text, "\"bits\": 4", replacement); + defer std.testing.allocator.free(bytes); + var config = try c.parse(std.testing.allocator, bytes); + defer config.deinit(); + try std.testing.expectEqual(try std.fmt.parseInt(u8, bits, 10), config.quant.global.bits); + } +} + +test "attention output and gate projections honor module widths and reject out_proj" { + inline for (.{ "o_proj", "gate_proj" }) |name| { + const replacement = "\"language_model.model.layers.2.self_attn." ++ name ++ "\": {"; + const bytes = try std.mem.replaceOwned(u8, std.testing.allocator, vl, "\"language_model.model.embed_tokens\": {", replacement); + defer std.testing.allocator.free(bytes); + var config = try c.parse(std.testing.allocator, bytes); + defer config.deinit(); + try std.testing.expectEqual(c.Width{ .bits = 8, .group = 32 }, try config.width("model.layers.2.self_attn." ++ name)); + try std.testing.expectEqual(c.Width{ .bits = 4, .group = 32 }, try config.width("model.layers.6.self_attn." ++ name)); + } + const bad = try std.mem.replaceOwned(u8, std.testing.allocator, vl, "\"language_model.model.embed_tokens\": {", "\"language_model.model.layers.2.self_attn.out_proj\": {"); + defer std.testing.allocator.free(bad); + try std.testing.expectError(error.UnsupportedLfm3Quantization, c.parse(std.testing.allocator, bad)); +} + +test "layer override indices must use canonical decimal spelling" { + inline for (.{ "02", "+2", "00", "-2", "" }) |index| { + const bytes = try std.mem.replaceOwned(u8, std.testing.allocator, vl, "\"language_model.model.embed_tokens\": {", "\"language_model.model.layers." ++ index ++ ".self_attn.gate_proj\": {"); + defer std.testing.allocator.free(bytes); + try std.testing.expectError(error.UnsupportedLfm3Quantization, c.parse(std.testing.allocator, bytes)); + } + const bytes = try std.mem.replaceOwned(u8, std.testing.allocator, vl, "\"language_model.model.embed_tokens\": {", "\"language_model.model.layers.0.conv.in_proj\": {"); + defer std.testing.allocator.free(bytes); + var config = try c.parse(std.testing.allocator, bytes); + defer config.deinit(); + try std.testing.expectEqual(c.Width{ .bits = 8, .group = 32 }, try config.width("layers.0.conv.in_proj")); +} + +test "geometry, mixer schedule, rotary, sandwich and unsupported quantization fail admission" { + const cases = .{ + .{ "\"hidden_size\": 2048", "\"hidden_size\": 1024", error.UnsupportedLfm3Geometry }, + .{ "\"conv_L_cache\": 7", "\"conv_L_cache\": 3", error.UnsupportedLfm3Geometry }, + .{ "\"sliding_attention\",", "\"full_attention\",", error.UnsupportedLfm3LayerSchedule }, + .{ "\"partial_rotary_factor\": 0.5", "\"partial_rotary_factor\": 1.0", error.UnsupportedLfm3Config }, + .{ "\"use_sandwich_norm\": true", "\"use_sandwich_norm\": false", error.UnsupportedLfm3Config }, + .{ "\"bits\": 4", "\"bits\": 5", error.UnsupportedLfm3Quantization }, + .{ "\"group_size\": 64", "\"group_size\": 128", error.UnsupportedLfm3Quantization }, + .{ "\"max_position_embeddings\": 262144", "\"max_position_embeddings\": 128000", error.UnsupportedLfm3Context }, + }; + inline for (cases) |case| { + const bytes = try std.mem.replaceOwned(u8, std.testing.allocator, text, case[0], case[1]); + defer std.testing.allocator.free(bytes); + try std.testing.expectError(case[2], c.parse(std.testing.allocator, bytes)); + } +} + +test "quantization declarations must agree and reject malformed or unknown overrides" { + const mismatch = try std.mem.replaceOwned(u8, std.testing.allocator, vl, "\"quantization_config\": {\n \"group_size\": 32", "\"quantization_config\": {\n \"group_size\": 64"); + defer std.testing.allocator.free(mismatch); + try std.testing.expectError(error.UnsupportedLfm3Quantization, c.parse(std.testing.allocator, mismatch)); + const cases = .{ + .{ "\"language_model.model.embed_tokens\": {", "\"unknown_module\": {" }, + .{ "\"language_model.model.embed_tokens\": {", "\"language_model.model.layers.99.conv.in_proj\": {" }, + .{ "\"mode\": \"affine\"", "\"mode\": \"mxfp4\"" }, + }; + inline for (cases) |case| { + const bytes = try std.mem.replaceOwned(u8, std.testing.allocator, vl, case[0], case[1]); + defer std.testing.allocator.free(bytes); + const result = c.parse(std.testing.allocator, bytes); + if (result) |value| { + var config = value; + config.deinit(); + return error.TestExpectedError; + } else |_| {} + } +} + +test "VL projector and vision geometry are not silently accepted as the known model" { + const bytes = try std.mem.replaceOwned(u8, std.testing.allocator, vl, "\"patch_size\": 16", "\"patch_size\": 14"); + defer std.testing.allocator.free(bytes); + try std.testing.expectError(error.UnsupportedLfm3Geometry, c.parse(std.testing.allocator, bytes)); +} + +test "malformed JSON fails without retaining allocations" { + try std.testing.expectError(error.UnexpectedEndOfInput, c.parse(std.testing.allocator, "{")); +} + +test "the ModelOpt NVFP4 export parses as its own scheme: FP4 g16 FFN, bf16 mixers" { + var config = try c.parse(std.testing.allocator, nvfp4); + defer config.deinit(); + try std.testing.expectEqual(c.Scheme.modelopt_nvfp4, config.scheme); + try std.testing.expectEqual(c.nvfp4_width, try config.width("model.layers.0.feed_forward.w1")); + try std.testing.expectEqual(@as(usize, 262144), config.context); + var mlx = try c.parse(std.testing.allocator, text); + defer mlx.deinit(); + try std.testing.expectEqual(c.Scheme.mlx_affine, mlx.scheme); +} + +test "Metal refuses the ModelOpt NVFP4 export by name; CUDA serves it, and both serve the MLX layouts" { + var config = try c.parse(std.testing.allocator, nvfp4); + defer config.deinit(); + try std.testing.expectError(error.Lfm3ModeloptNvfp4NotOnMetal, c.servedOn(&config, .metal)); + try c.servedOn(&config, .cuda); + var mlx = try c.parse(std.testing.allocator, text); + defer mlx.deinit(); + try c.servedOn(&mlx, .metal); + try c.servedOn(&mlx, .cuda); +} + +test "an NVFP4 export with another recipe, a missing or extra exclusion, or dynamic inputs is refused" { + const edits = [_][2][]const u8{ + .{ "\"quant_algo\": \"NVFP4\"", "\"quant_algo\": \"FP8\"" }, + .{ "\"model.layers.9.conv*\"", "\"model.layers.9.self_attn*\"" }, + .{ " \"model.layers.9.conv*\"\n", " \"model.layers.9.conv*\",\n \"model.layers.9.feed_forward*\"\n" }, + .{ " \"model.layers.8.self_attn*\",\n", "" }, + .{ "\"model.layers.10.self_attn*\"", "\"model.layers.2.self_attn*\"" }, + .{ "\"dynamic\": false", "\"dynamic\": true" }, + .{ "\"group_size\": 16", "\"group_size\": 32" }, + .{ "\"Linear\"", "\"Conv1d\"" }, + }; + for (edits) |e| { + const bytes = try std.mem.replaceOwned(u8, std.testing.allocator, nvfp4, e[0], e[1]); + defer std.testing.allocator.free(bytes); + try std.testing.expect(!std.mem.eql(u8, bytes, nvfp4)); + try std.testing.expectError(error.UnsupportedLfm3Quantization, c.parse(std.testing.allocator, bytes)); + } +} + +pub const text = + \\{ + \\ "model_type": "lfm3", + \\ "conv_L_cache": 7, + \\ "conv_bias": false, + \\ "dtype": "bfloat16", + \\ "head_dim": 128, + \\ "hidden_size": 2048, + \\ "intermediate_size": 7680, + \\ "layer_types": [ + \\ "conv", + \\ "conv", + \\ "full_attention", + \\ "conv", + \\ "sliding_attention", + \\ "conv", + \\ "full_attention", + \\ "conv", + \\ "sliding_attention", + \\ "conv", + \\ "full_attention", + \\ "conv", + \\ "sliding_attention", + \\ "conv", + \\ "full_attention", + \\ "conv", + \\ "sliding_attention", + \\ "conv", + \\ "full_attention", + \\ "conv", + \\ "sliding_attention", + \\ "conv", + \\ "full_attention", + \\ "conv" + \\ ], + \\ "max_position_embeddings": 262144, + \\ "norm_eps": 1e-05, + \\ "num_attention_heads": 16, + \\ "num_hidden_layers": 24, + \\ "num_key_value_heads": 4, + \\ "quantization": { + \\ "group_size": 64, + \\ "bits": 4, + \\ "mode": "affine" + \\ }, + \\ "quantization_config": { + \\ "group_size": 64, + \\ "bits": 4, + \\ "mode": "affine" + \\ }, + \\ "rope_parameters": { + \\ "full_attention": { + \\ "partial_rotary_factor": 0.5, + \\ "rope_theta": 1000000.0, + \\ "rope_type": "default" + \\ }, + \\ "sliding_attention": { + \\ "rope_theta": 10000, + \\ "rope_type": "default" + \\ } + \\ }, + \\ "sliding_window": 256, + \\ "tie_word_embeddings": true, + \\ "use_sandwich_norm": true, + \\ "vocab_size": 128000 + \\} +; + +pub const vl = + \\{ + \\ "downsample_factor": 2, + \\ "dtype": "bfloat16", + \\ "image_token_id": 125030, + \\ "model_type": "lfm3_vl", + \\ "projector_bias": true, + \\ "projector_hidden_act": "gelu", + \\ "projector_hidden_size": 4096, + \\ "projector_layer_norm_eps": 1e-06, + \\ "text_config": { + \\ "conv_L_cache": 7, + \\ "conv_bias": false, + \\ "dtype": "bfloat16", + \\ "head_dim": 128, + \\ "hidden_size": 2048, + \\ "intermediate_size": 7680, + \\ "layer_types": [ + \\ "conv", + \\ "conv", + \\ "full_attention", + \\ "conv", + \\ "sliding_attention", + \\ "conv", + \\ "full_attention", + \\ "conv", + \\ "sliding_attention", + \\ "conv", + \\ "full_attention", + \\ "conv", + \\ "sliding_attention", + \\ "conv", + \\ "full_attention", + \\ "conv", + \\ "sliding_attention", + \\ "conv", + \\ "full_attention", + \\ "conv", + \\ "sliding_attention", + \\ "conv", + \\ "full_attention", + \\ "conv" + \\ ], + \\ "max_position_embeddings": 128000, + \\ "model_type": "lfm3", + \\ "norm_eps": 1e-05, + \\ "num_attention_heads": 16, + \\ "num_hidden_layers": 24, + \\ "num_key_value_heads": 4, + \\ "rope_parameters": { + \\ "full_attention": { + \\ "partial_rotary_factor": 0.5, + \\ "rope_theta": 1000000.0, + \\ "rope_type": "default" + \\ }, + \\ "sliding_attention": { + \\ "rope_theta": 10000, + \\ "rope_type": "default" + \\ } + \\ }, + \\ "sliding_window": 256, + \\ "tie_word_embeddings": true, + \\ "use_sandwich_norm": true, + \\ "vocab_size": 128000 + \\ }, + \\ "tie_word_embeddings": true, + \\ "video_token_id": 125029, + \\ "vision_config": { + \\ "conv_kernel_size": 3, + \\ "dtype": "bfloat16", + \\ "hidden_size": 1024, + \\ "in_channels": 3, + \\ "initializer_range": 0.02, + \\ "intermediate_size": 2560, + \\ "layer_types": [ + \\ "conv", + \\ "conv", + \\ "full_attention", + \\ "conv", + \\ "full_attention", + \\ "conv", + \\ "full_attention", + \\ "conv", + \\ "full_attention", + \\ "conv", + \\ "full_attention", + \\ "conv", + \\ "full_attention", + \\ "conv" + \\ ], + \\ "model_type": "lfm3_vl_vision", + \\ "norm_eps": 1e-05, + \\ "num_attention_heads": 16, + \\ "num_hidden_layers": 14, + \\ "num_key_value_heads": 8, + \\ "patch_size": 16, + \\ "rope_parameters": { + \\ "rope_theta": 1000000.0, + \\ "rope_type": "axial" + \\ }, + \\ "spatial_merge_size": 2, + \\ "temporal_patch_size": 2 + \\ }, + \\ "vision_end_token_id": 125010, + \\ "vision_start_token_id": 125009, + \\ "quantization": { + \\ "group_size": 32, + \\ "bits": 4, + \\ "mode": "affine", + \\ "language_model.model.embed_tokens": { + \\ "group_size": 32, + \\ "bits": 8, + \\ "mode": "affine" + \\ } + \\ }, + \\ "quantization_config": { + \\ "group_size": 32, + \\ "bits": 4, + \\ "mode": "affine", + \\ "language_model.model.embed_tokens": { + \\ "group_size": 32, + \\ "bits": 8, + \\ "mode": "affine" + \\ } + \\ } + \\} +; + +/// LiquidAI/LFM3-1.7B-NVFP4 (b30025bb)'s config.json as published. +pub const nvfp4 = + \\{ + \\ "architectures": [ + \\ "Lfm3ForCausalLM" + \\ ], + \\ "auto_map": { + \\ "AutoConfig": "configuration_lfm3.Lfm3Config", + \\ "AutoModel": "modeling_lfm3.Lfm3Model", + \\ "AutoModelForCausalLM": "modeling_lfm3.Lfm3ForCausalLM" + \\ }, + \\ "bos_token_id": 124894, + \\ "conv_L_cache": 7, + \\ "conv_bias": false, + \\ "dtype": "bfloat16", + \\ "eos_token_id": 124900, + \\ "head_dim": 128, + \\ "hidden_size": 2048, + \\ "initializer_range": 0.02, + \\ "intermediate_size": 7680, + \\ "layer_types": [ + \\ "conv", + \\ "conv", + \\ "full_attention", + \\ "conv", + \\ "sliding_attention", + \\ "conv", + \\ "full_attention", + \\ "conv", + \\ "sliding_attention", + \\ "conv", + \\ "full_attention", + \\ "conv", + \\ "sliding_attention", + \\ "conv", + \\ "full_attention", + \\ "conv", + \\ "sliding_attention", + \\ "conv", + \\ "full_attention", + \\ "conv", + \\ "sliding_attention", + \\ "conv", + \\ "full_attention", + \\ "conv" + \\ ], + \\ "max_position_embeddings": 262144, + \\ "model_type": "lfm3", + \\ "norm_eps": 1e-05, + \\ "num_attention_heads": 16, + \\ "num_hidden_layers": 24, + \\ "num_key_value_heads": 4, + \\ "pad_token_id": 124893, + \\ "rope_parameters": { + \\ "full_attention": { + \\ "partial_rotary_factor": 0.5, + \\ "rope_theta": 1000000.0, + \\ "rope_type": "default" + \\ }, + \\ "sliding_attention": { + \\ "rope_theta": 10000, + \\ "rope_type": "default" + \\ } + \\ }, + \\ "sliding_window": 256, + \\ "tie_word_embeddings": true, + \\ "transformers_version": "5.12.1", + \\ "use_cache": true, + \\ "use_sandwich_norm": true, + \\ "vocab_size": 128000, + \\ "quantization_config": { + \\ "config_groups": { + \\ "group_0": { + \\ "input_activations": { + \\ "dynamic": false, + \\ "num_bits": 4, + \\ "type": "float", + \\ "group_size": 16 + \\ }, + \\ "weights": { + \\ "dynamic": false, + \\ "num_bits": 4, + \\ "type": "float", + \\ "group_size": 16 + \\ }, + \\ "targets": [ + \\ "Linear" + \\ ] + \\ } + \\ }, + \\ "ignore": [ + \\ "lm_head", + \\ "model.embed_tokens", + \\ "model.layers.0.conv*", + \\ "model.layers.1.conv*", + \\ "model.layers.10.self_attn*", + \\ "model.layers.11.conv*", + \\ "model.layers.12.self_attn*", + \\ "model.layers.13.conv*", + \\ "model.layers.14.self_attn*", + \\ "model.layers.15.conv*", + \\ "model.layers.16.self_attn*", + \\ "model.layers.17.conv*", + \\ "model.layers.18.self_attn*", + \\ "model.layers.19.conv*", + \\ "model.layers.2.self_attn*", + \\ "model.layers.20.self_attn*", + \\ "model.layers.21.conv*", + \\ "model.layers.22.self_attn*", + \\ "model.layers.23.conv*", + \\ "model.layers.3.conv*", + \\ "model.layers.4.self_attn*", + \\ "model.layers.5.conv*", + \\ "model.layers.6.self_attn*", + \\ "model.layers.7.conv*", + \\ "model.layers.8.self_attn*", + \\ "model.layers.9.conv*" + \\ ], + \\ "quant_algo": "NVFP4", + \\ "producer": { + \\ "name": "modelopt", + \\ "version": "0.0.1.dev1+g4c1f81382" + \\ }, + \\ "quant_method": "modelopt" + \\ } + \\} +; diff --git a/zig/src/families/lfm3/cuda_backend.zig b/zig/src/families/lfm3/cuda_backend.zig new file mode 100644 index 0000000000..4fab52a7ab --- /dev/null +++ b/zig/src/families/lfm3/cuda_backend.zig @@ -0,0 +1,17 @@ +//! LFM3 on CUDA behind lanes.Backend: liquid's lane backend over this family's model and forward. +const liquid = @import("liquid_cuda"); + +const B = liquid.lanes.Backend(struct { + pub const cfg = @import("config.zig"); + pub const fwd = @import("cuda_forward.zig"); + pub const mdl = @import("cuda_model.zig"); + pub const kern = @import("cuda_kernels.zig"); + pub const snap = @import("cuda_snap.zig"); + pub const text_only = error.Lfm3TextOnly; +}); + +pub const Liquid = B.Liquid; +pub const prompt_rows = B.prompt_rows; +pub const round_rows = B.round_rows; +pub const max_window = B.max_window; +pub const pagesFor = B.pagesFor; diff --git a/zig/src/families/lfm3/cuda_family.zig b/zig/src/families/lfm3/cuda_family.zig new file mode 100644 index 0000000000..0c68d857fb --- /dev/null +++ b/zig/src/families/lfm3/cuda_family.zig @@ -0,0 +1,23 @@ +//! LFM3 text on CUDA: the shared admission (config.zig, weights.zig), the shared CUDA operators, lfm3.cu's 7-tap conv. +pub const config = @import("config.zig"); +pub const admission = @import("weights.zig"); +pub const weights = @import("cuda_weights.zig"); +pub const kernels = @import("cuda_kernels.zig"); +pub const model = @import("cuda_model.zig"); +pub const forward = @import("cuda_forward.zig"); +pub const backend = @import("cuda_backend.zig"); +pub const native = @import("cuda_native.zig"); +pub const snap = @import("cuda_snap.zig"); + +test { + _ = config; + _ = admission; + _ = weights; + _ = kernels; + _ = model; + _ = forward; + _ = backend; + _ = native; + _ = snap; + _ = @import("config_test.zig"); +} diff --git a/zig/src/families/lfm3/cuda_forward.zig b/zig/src/families/lfm3/cuda_forward.zig new file mode 100644 index 0000000000..b6ffc2f47f --- /dev/null +++ b/zig/src/families/lfm3/cuda_forward.zig @@ -0,0 +1,249 @@ +//! LFM3's forward on CUDA: a prompt chunk or a round of windows, each row's bits its own (solo == batched). +const std = @import("std"); +const cuda = @import("cuda"); +const c = @import("config.zig"); +const wts = @import("cuda_weights.zig"); +const kern = @import("cuda_kernels.zig"); +const mdl = @import("cuda_model.zig"); + +const Model = mdl.Model; +const Kernels = kern.Kernels; +const Stream = cuda.Stream; +/// HF's RMSNorm points: bf16(w * bf16(h * inv)), the residual rounded at each add. +const round = true; + +/// A stream's rows in one forward: its slot, first position, tokens. +pub const Chain = struct { slot: usize, p0: usize, tokens: []const u32 }; + +const Pass = struct { + m: *const Model, + k: *const Kernels, + s: Stream, + chains: []const Chain, + rows: usize, + prompt: bool, + + /// out = x w: FP4 prompt rows on the prompt GEMM, the rest on the lane a round at a time (bits free of chunk cuts). + fn project(p: Pass, w: wts.Proj, x: u64, out: u64, rows: usize) !void { + if (w == .affine4) return affine(p.m, p.k, p.s, w.affine4, x, out, rows, p.prompt); + const lin = p.k.linear(); + const in: cuda.qlinear.Rows = .{ .x = x, .part = p.m.s.part }; + if (p.prompt and w == .nvfp4) return lin.prompt(p.s, w, in, out, rows); + var at: usize = 0; + while (at < rows) : (at += mdl.round_rows) { + const n = @min(mdl.round_rows, rows - at); + const slice: cuda.qlinear.Rows = .{ .x = x + at * w.inputs() * 2, .part = in.part }; + try lin.decode(p.s, w, slice, out + at * w.outputs() * 2, n); + } + } + + /// h += post(block) (null: h is the embedding), x = norm(h). + fn norm(p: Pass, post: ?u64, w: u64) !void { + const x = p.m.s; + const tail: cuda.norms.Tail = .{ .round = round, .branch = post }; + if (post == null) return p.k.norms.addRmsTail(p.s, x.h, null, w, x.h, x.x, null, p.rows, c.hidden, c.eps, tail); + try p.k.norms.addRmsTail(p.s, x.r, x.h, w, x.h, x.x, null, p.rows, c.hidden, c.eps, tail); + } + + fn conv(p: Pass, cv: wts.Conv, index: usize) !void { + const x = p.m.s; + try p.project(cv.in, x.x, x.wide, p.rows); + const state = p.m.states + index * mdl.state_rows * c.hidden * 2; + const snaps = x.snaps + index * mdl.round_rows * mdl.state_rows * c.hidden * 2; + // a prompt chunk keeps its last row's state, a round each row's (a window keeps any prefix of its rows) + const snap0: usize = if (p.prompt) p.rows - 1 else 0; + try p.k.convRows(p.s, x.wide, cv.weight, state, mdl.slot_elems, x.slot, x.off, x.y, snaps, snap0, p.rows); + try p.project(cv.out, x.y, x.r, p.rows); + } + + fn attention(p: Pass, a: wts.Attention, full_at: usize, sliding_at: usize) !void { + const x = p.m.s; + const m = p.m; + const qw = c.query_heads * c.head_dim; + try p.project(a.qkvg, x.x, x.wide, p.rows); + const ld: i64 = @intCast(a.qkvg.outputs()); + const dh: i64 = c.head_dim; + const src: cuda.qk_rows.Src = .{ .q = x.wide, .q_ld = ld, .q_hs = dh, .k = x.wide + qw * 2, .k_ld = ld, .k_hs = dh, .v = x.wide + (qw + mdl.kv_width) * 2, .v_ld = ld, .v_hs = dh }; + const shape: cuda.qk_rows.Shape = .{ .norms = a.norms, .norms_bf16 = true, .inv = m.inv[@intFromBool(a.sliding)], .rot = c.head_dim, .eps = c.eps, .round = true }; + var out: cuda.qk_rows.Out = .{ .q = x.q, .k = x.kv, .v = x.kv + mdl.kv_width * 2, .kv_ld = 2 * mdl.kv_width }; + if (a.sliding) { + out.kc = m.ringOf(sliding_at, false); + out.vc = m.ringOf(sliding_at, true); + out.len = mdl.ring; + } + try p.k.qk.rowsFrom(p.s, src, .{ .h = c.query_heads, .hk = c.kv_heads, .dim = c.head_dim }, shape, x.pos, x.slot, out, p.rows); + if (a.sliding) try p.sliding(sliding_at) else try p.full(full_at); + // HF's gate: bf16(o * bf16(sigmoid(gate_proj(x)))), one scalar a query head; on its own, into the merge's sums + var gate = x.wide + (qw + 2 * mdl.kv_width) * 2; + var gate_ld = a.qkvg.outputs(); + if (a.gate) |g| { + try p.project(g, x.x, x.xs, p.rows); + gate = x.xs; + gate_ld = c.query_heads; + } + try p.k.act.headGate(p.s, x.att, gate, gate_ld, x.y, null, p.rows, c.query_heads, c.head_dim); + try p.project(a.out, x.y, x.r, p.rows); + } + + /// Keys (pos - 256, pos] of each row's slot's ring. + fn sliding(p: Pass, l: usize) !void { + const x = p.m.s; + const m = p.m; + const w: cuda.window_attention.Window = .{ .h = c.query_heads, .hk = c.kv_heads, .d = c.head_dim, .win = c.sliding_window, .ring = mdl.ring, .scale = kern.scale }; + if (p.prompt) { + const off = p.chains[0].slot * mdl.ring * mdl.kv_width * 2; + return p.k.win.prompt(p.s, x.q, m.ringOf(l, false) + off, m.ringOf(l, true) + off, x.att, p.chains[0].p0, p.rows, w); + } + try p.k.win.decode(p.s, x.q, m.ringOf(l, false), m.ringOf(l, true), x.att, x.pos, x.slot, w, 0, p.rows); + } + + /// Every key [0, pos] of the stream's pages. + fn full(p: Pass, index: usize) !void { + const x = p.m.s; + const m = p.m; + const kc = m.cache(index, false); + const vc = m.cache(index, true); + var row0: usize = 0; + var at: usize = 0; + while (at < p.chains.len) : (at += kern.max_segs) { + const group = p.chains[at..@min(p.chains.len, at + kern.max_segs)]; + var sg = std.mem.zeroes(kern.AttnSegs); + var rows: usize = 0; + for (group, 0..) |ch, j| { + sg.kc[j] = kc; + sg.vc[j] = vc; + sg.meta[j] = x.meta + (at + j) * 4; + sg.pt[j] = m.tableOf(ch.slot); + sg.row0[j] = @intCast(rows); + rows += ch.tokens.len; + } + sg.ps = mdl.page_elems; + sg.n = @intCast(group.len); + try p.k.kvWrite(p.s, sg, x.kv + row0 * 2 * mdl.kv_width * 2, rows); + const qa = x.q + row0 * c.hidden * 2; + const ya = x.att + row0 * c.hidden * 2; + if (p.prompt) { + try p.k.prompt(p.s, qa, kc, vc, ya, group[0].p0, rows, m.tableOf(group[0].slot), mdl.page_elems); + } else { + const xa = x.xs + row0 * (c.hidden / 64) * 4; + try p.k.attend(p.s, sg, qa, x.po, x.pm, x.pl, ya, xa, rows, m.chunks(), m.chunkBlocks()); + } + row0 += rows; + } + } + + /// w2(bf16(silu(w1 x)) * w3 x): w1 | w3 one FP4 projection. + fn ffn(p: Pass, l: wts.Layer) !void { + const x = p.m.s; + try p.project(l.gate_up, x.x, x.wide, p.rows); + try p.k.siluMul(p.s, x.wide, x.y, p.rows, c.intermediate); + try p.project(l.down, x.y, x.r, p.rows); + } +}; + +/// Every chain's rows (a prompt chunk, or windows); logits rows [head0, head0 + heads); conv states end at last rows. +pub fn forward(m: *const Model, k: *const Kernels, s: Stream, chains: []const Chain, prompt: bool, head0: usize, heads: usize) !void { + var rows: usize = 0; + for (chains) |ch| { + if (ch.tokens.len == 0 or ch.slot >= m.slots or ch.p0 + ch.tokens.len > m.context) return error.BadChain; + if (ch.p0 + ch.tokens.len > m.caps[ch.slot]) return error.PastPages; + if (!prompt and ch.tokens.len + c.sliding_window > mdl.ring) return error.BadWindow; + rows += ch.tokens.len; + } + if (rows > @as(usize, if (prompt) mdl.prompt_rows else mdl.round_rows) or (prompt and chains.len != 1)) return error.BadWindow; + if (chains.len > mdl.round_rows or heads > mdl.round_rows or head0 + heads > rows) return error.BadWindow; + try upload(m, s, chains, rows); + const p: Pass = .{ .m = m, .k = k, .s = s, .chains = chains, .rows = rows, .prompt = prompt }; + const x = m.s; + if (m.w.embed_affine) |e| { + try k.embed.affine4(s, x.ids, e[0], e[1], e[2], x.h, rows, c.hidden); + } else try k.embed.bf16(s, x.ids, m.w.embed, x.h, rows, c.hidden); + var conv_at: usize = 0; + var full_at: usize = 0; + var sliding_at: usize = 0; + var post: ?u64 = null; // the previous FFN's post norm: its output is added before this layer's operator norm + for (m.w.layers) |l| { + try p.norm(post, l.operator_norm); + switch (l.mixer) { + .conv => |cv| { + try p.conv(cv, conv_at); + conv_at += 1; + }, + .attention => |a| { + try p.attention(a, full_at, sliding_at); + if (a.sliding) sliding_at += 1 else full_at += 1; + }, + } + try p.norm(l.post_operator_norm, l.ffn_norm); + try p.ffn(l); + post = l.post_ffn_norm; + } + try p.norm(post, m.w.final_norm); + try commitAll(m, k, s, chains, prompt); + if (heads == 0) return; + const hx = x.x + head0 * c.hidden * 2; + if (m.w.head == .affine4) return affine(m, k, s, m.w.head.affine4, hx, x.logits, heads, false); + try k.linear().decode(s, m.w.head, .{ .x = hx, .part = x.part }, x.logits, heads); +} + +/// out = x q for MLX affine-4: a prompt chunk's GEMM, or decode rows on their 64-value sums, a verify tile at a time. +fn affine(m: *const Model, k: *const Kernels, s: Stream, q: cuda.qlinear.Affine4, x: u64, out: u64, rows: usize, prompt: bool) !void { + if (prompt) return k.affine.prompt(s, x, q, out, rows); + try k.act.groupSums(s, x, q.k, m.s.sums, rows, q.k); + const tile = k.affine.verifyRows(); + var at: usize = 0; + while (at < rows) : (at += tile) { + const n = @min(tile, rows - at); + try k.affine.verify(s, x + at * q.k * 2, m.s.sums + at * (q.k / 64) * 4, q, out + at * q.n * 2, n); + } +} + +/// Each stream's conv states (slots[i]) set to its state after snapshot row rows[i]; a row < 0 keeps them. +pub fn commit(m: *const Model, k: *const Kernels, s: Stream, slots: []const i32, rows: []const i32) !void { + if (slots.len != rows.len or slots.len > mdl.round_rows) return error.BadWindow; + try put(m, s, m.s.commit_slot, std.mem.sliceAsBytes(slots)); + try put(m, s, m.s.commit_row, std.mem.sliceAsBytes(rows)); + const sw = mdl.state_rows * c.hidden; + try k.commitRows(s, m.states, m.s.snaps, m.s.commit_slot, m.s.commit_row, slots.len, mdl.conv_layers, sw, mdl.round_rows * sw, mdl.slot_elems); +} + +/// Each chain's conv states at its last row. +fn commitAll(m: *const Model, k: *const Kernels, s: Stream, chains: []const Chain, prompt: bool) !void { + var slots: [mdl.round_rows]i32 = undefined; + var rows: [mdl.round_rows]i32 = undefined; + var base: usize = 0; + for (chains, 0..) |ch, i| { + base += ch.tokens.len; + slots[i] = @intCast(ch.slot); + rows[i] = if (prompt) 0 else @intCast(base - 1); + } + try commit(m, k, s, slots[0..chains.len], rows[0..chains.len]); +} + +fn put(m: *const Model, s: Stream, dst: u64, bytes: []const u8) !void { + try m.d.check(m.d.api.cuMemcpyHtoDAsync_v2(dst, bytes.ptr, bytes.len, s.handle), "cuMemcpyHtoDAsync"); +} + +/// Tokens, positions, slots and offsets a row, each chain's first position, staged for the launches. +fn upload(m: *const Model, s: Stream, chains: []const Chain, rows: usize) !void { + const ids = try m.gpa.alloc(i32, 4 * rows + chains.len); + defer m.gpa.free(ids); + var at: usize = 0; + for (chains, 0..) |ch, j| { + for (ch.tokens, 0..) |t, i| { + ids[at + i] = @intCast(t); + ids[rows + at + i] = @intCast(ch.p0 + i); + ids[2 * rows + at + i] = @intCast(ch.slot); + ids[3 * rows + at + i] = @intCast(i); + } + ids[4 * rows + j] = @intCast(ch.p0); + at += ch.tokens.len; + } + const x = m.s; + for ([_]u64{ x.ids, x.pos, x.slot, x.off }, 0..) |dst, i| try put(m, s, dst, std.mem.sliceAsBytes(ids[i * rows ..][0..rows])); + try put(m, s, x.meta, std.mem.sliceAsBytes(ids[4 * rows ..])); +} + +test "the forward's launches type-check against the shared operators" { + std.testing.refAllDecls(Pass); +} diff --git a/zig/src/families/lfm3/cuda_kernels.zig b/zig/src/families/lfm3/cuda_kernels.zig new file mode 100644 index 0000000000..0c530287ce --- /dev/null +++ b/zig/src/families/lfm3/cuda_kernels.zig @@ -0,0 +1,195 @@ +//! LFM3's launches on CUDA: the shared operators (linears, norms, QK rows, attention, gates), lfm3.cu's 7-tap conv. +const std = @import("std"); +const cuda = @import("cuda"); +const c = @import("config.zig"); + +const Args = cuda.Args; +const Function = cuda.Function; +const Stream = cuda.Stream; + +/// prefill_attention.cu's pattn_kernel<128, 8, 4, 8> over paged caches: eight 32-key slots, 32 rows a block. +const pattn_sym = "_ZN20tf_prefill_attention12pattn_kernelILi128ELi8ELi4ELi8EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifPKix"; +const pattn_smem: u32 = 8 * 32 * 16 * 16; +const pattn_rows = 32; +/// Positions a KV page holds (nemotron_attention.cu's PAGE and prefill_attention_rows.cuh's 512). +pub const page = 512; +/// Streams one shared-window attention launch takes (nemotron_attention.cu's MAXSEG). +pub const max_segs = 8; +/// The attention's softmax scale: head_dim^-0.5. +pub const scale: f32 = 1.0 / @sqrt(@as(f32, c.head_dim)); + +/// nemotron_attention.cu's AttnSegs: each stream's caches, position, page table and first row. +pub const AttnSegs = extern struct { kc: [max_segs]u64, vc: [max_segs]u64, meta: [max_segs]u64, pt: [max_segs]u64, ps: i64, row0: [max_segs]c_int, n: c_int }; + +const images = .{ "sample", "tree_round", "prefill_attention", "nemotron_attention", "lfm2", "lfm3", "qmm_group", "qmm_prefill", "lane_gemv", "affine4_pack" }; + +fn int(x: usize) c_int { + return @intCast(x); +} + +pub const Kernels = struct { + mods: [images.len]cuda.Module, + lane: cuda.qmmf.Lane, + prompt16: cuda.nvfp4.Prompt, + affine: cuda.qlinear.Affine4Kernels, // MLX 4-bit projections and the head + pack: cuda.modelopt.Packer, + draws: cuda.sample.Draws, + round: cuda.tree_round.Kernels, + embed: cuda.embed_rows.EmbedRows, + norms: cuda.norms.Norms, + qk: cuda.qk_rows.QkRows, + act: cuda.act_rows.ActRows, + win: cuda.window_attention.WindowAttention, + pattn: Function, + kv_write: Function, + chunk: Function, + merge: Function, + silu_mul: Function, + conv: Function, + commit: Function, + + pub fn load(ctx: *const cuda.Context) !Kernels { + const d = ctx.d; + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var k: Kernels = undefined; + var loaded: usize = 0; + errdefer for (k.mods[0..loaded]) |*m| m.unload(); + inline for (images, 0..) |name, i| { + k.mods[i] = try cuda.Module.load(d, @field(cuda.kernels, name)); + loaded += 1; + } + const m = &k.mods; + k.draws = try cuda.sample.Draws.resolve(m[0]); + k.round = try cuda.tree_round.Kernels.resolve(m[1]); + k.pattn = try m[2].function(pattn_sym); + try k.pattn.allowDynamicShared(pattn_smem); + k.kv_write = try m[3].function("tf_nemo_kv_write_segs"); + k.chunk = try m[3].function("tf_nemo_attn_chunk_segs"); + k.merge = try m[3].function("tf_nemo_attn_merge_segs"); + k.silu_mul = try m[4].function("tf_lfm2_silu_mul"); + k.conv = try m[5].function("tf_liquid_conv_7"); + k.commit = try m[5].function("tf_liquid_commit_7"); + const sms: usize = @intCast(try ctx.attribute(.multiprocessor_count)); + k.affine = try cuda.qlinear.Affine4Kernels.resolve(d, m[6], m[7], m[8], m[9], sms, try ctx.capability() == 121); + errdefer k.affine.deinit(); + k.lane = try cuda.qmmf.Lane.load(d, (try ctx.capability()) / 10); + errdefer k.lane.unload(); + k.prompt16 = try cuda.nvfp4.Prompt.load(d); + errdefer k.prompt16.unload(); + k.pack = try cuda.modelopt.Packer.load(d); + errdefer k.pack.unload(); + k.embed = try cuda.embed_rows.EmbedRows.load(d); + errdefer k.embed.unload(); + k.norms = try cuda.norms.Norms.load(d); + errdefer k.norms.unload(); + k.qk = try cuda.qk_rows.QkRows.load(d); + errdefer k.qk.unload(); + k.act = try cuda.act_rows.ActRows.load(d); + errdefer k.act.unload(); + k.win = try cuda.window_attention.WindowAttention.load(d); + return k; + } + + pub fn deinit(k: *Kernels) void { + k.win.unload(); + k.act.unload(); + k.qk.unload(); + k.norms.unload(); + k.embed.unload(); + k.pack.unload(); + k.prompt16.unload(); + k.lane.unload(); + k.affine.deinit(); + for (&k.mods) |*m| m.unload(); + } + + /// The lane for every row of a decode forward and the bf16 projections, prompt16 for an FP4 prompt chunk. + pub fn linear(k: *const Kernels) cuda.qlinear.Linear { + return .{ .lane = &k.lane, .nvfp4_prompt = &k.prompt16, .affine4 = &k.affine }; + } + + fn go(f: Function, s: Stream, grid: [3]usize, block: usize, smem: u32, a: *Args) !void { + const g: cuda.Dim3 = .{ .x = @intCast(grid[0]), .y = @intCast(grid[1]), .z = @intCast(grid[2]) }; + try cuda.launch.launch(f, .{ .grid = g, .block = .{ .x = @intCast(block) }, .shared = smem }, s, a); + } + + /// Short conv over bcx [rows, 3 hidden]: each row's slot and offset; a slot's [6, hidden] state every `ss`. + pub fn convRows(k: *const Kernels, s: Stream, bcx: u64, w: u64, state: u64, ss: usize, slot: u64, off: u64, y: u64, snap: ?u64, snap0: usize, rows: usize) !void { + var a: Args = .{}; + for ([_]u64{ bcx, w, state, slot, off, y, snap orelse 0 }) |v| a.add(v); + a.add(int(snap0)); + a.add(int(c.hidden)); + a.add(@as(i64, @intCast(ss))); + try go(k.conv, s, .{ rows, c.hidden / 256, 1 }, 256, 0, &a); + } + + /// Each stream k's conv states (every conv layer) set to its snapshot row rows[k] (int32; < 0: kept). + pub fn commitRows(k: *const Kernels, s: Stream, state: u64, snap: u64, slots: u64, rows: u64, streams: usize, layers: usize, state_stride: usize, snap_stride: usize, ss: usize) !void { + if (streams == 0) return; + var a: Args = .{}; + for ([_]u64{ state, snap, slots, rows }) |v| a.add(v); + a.add(int(c.hidden)); + a.add(@as(i64, @intCast(state_stride))); + a.add(@as(i64, @intCast(snap_stride))); + a.add(@as(i64, @intCast(ss))); + try go(k.commit, s, .{ layers, streams, 1 }, 256, 0, &a); + } + + /// out [rows, f] = bf16(bf16(silu(gate)) * up) from rows gate | up [rows, 2 f] (lfm2.cu, torch's points). + pub fn siluMul(k: *const Kernels, s: Stream, g: u64, out: u64, rows: usize, f: usize) !void { + var a: Args = .{}; + a.add(g); + a.add(out); + a.add(int(f)); + try go(k.silu_mul, s, .{ rows, (f + 255) / 256, 1 }, 256, 0, &a); + } + + /// Rows k, v (kv [rows, 2 kv width]) into each stream's pages at its position on. + pub fn kvWrite(k: *const Kernels, s: Stream, sg: AttnSegs, kv: u64, rows: usize) !void { + var a: Args = .{}; + a.add(sg); + a.add(kv); + a.add(int(2 * c.kv_heads * c.head_dim)); + a.add(@as(c_int, 0)); + a.add(int(c.kv_heads * c.head_dim)); + try go(k.kv_write, s, .{ rows, 1, 1 }, 256, 0, &a); + } + + /// Decode rows: q [rows, heads, 128] over each stream's pages in 512-key chunks, merged in order (xs: sums). + pub fn attend(k: *const Kernels, s: Stream, sg: AttnSegs, q: u64, po: u64, pm: u64, pl: u64, out: u64, xs: u64, rows: usize, nch: usize, z: usize) !void { + const h = c.query_heads; + const hk = c.kv_heads; + const g = h / hk; + var a: Args = .{}; + a.add(sg); + a.add(q); + for ([_]u64{ po, pm, pl }) |v| a.add(v); + for ([_]usize{ h * c.head_dim, h, hk, g, page, nch }) |v| a.add(int(v)); + a.add(scale); + try go(k.chunk, s, .{ rows, hk, z }, 256, 0, &a); + var b: Args = .{}; + b.add(sg); + for ([_]u64{ po, pm, pl, out, xs }) |v| b.add(v); + for ([_]usize{ h, hk, g, page, nch }) |v| b.add(int(v)); + try go(k.merge, s, .{ rows, hk, 1 }, 256, 0, &b); + } + + /// A prompt chunk's rows at positions p0.. against its stream's pages (filled through them): out [rows, 16, 128]. + pub fn prompt(k: *const Kernels, s: Stream, q: u64, kc: u64, vc: u64, out: u64, p0: usize, rows: usize, pt: u64, ps: usize) !void { + const g = c.query_heads / c.kv_heads; + var a: Args = .{}; + for ([_]u64{ q, kc, vc, out }) |v| a.add(v); + for ([_]usize{ p0, rows, c.query_heads, c.kv_heads, g }) |v| a.add(int(v)); + a.add(scale); + a.add(pt); + a.add(@as(i64, @intCast(ps))); + try go(k.pattn, s, .{ (rows + pattn_rows - 1) / pattn_rows, c.kv_heads * (g / 4), 1 }, 256, pattn_smem, &a); + } +}; + +test "the attention launches' geometry: 4 query heads a KV head, a pattn block's 4 heads, head 128" { + try std.testing.expectEqual(@as(usize, 4), c.query_heads / c.kv_heads); + try std.testing.expectEqual(@as(usize, 304), @sizeOf(AttnSegs)); // nemotron_attention.cu's struct, padded to 8 + try std.testing.expectEqual(@as(u32, 65536), pattn_smem); // lfm3_cuda_records.zig's dynamic bytes + try std.testing.expectEqual(@as(usize, 0), c.hidden % 256); // the conv's channel blocks +} diff --git a/zig/src/families/lfm3/cuda_model.zig b/zig/src/families/lfm3/cuda_model.zig new file mode 100644 index 0000000000..1ab343e31a --- /dev/null +++ b/zig/src/families/lfm3/cuda_model.zig @@ -0,0 +1,387 @@ +//! LFM3 on CUDA: full-attention KV in 512-position pages, each slot's rings and conv states, the scratch. +const std = @import("std"); +const cuda = @import("cuda"); +const kv = @import("core").kv_pool; +const c = @import("config.zig"); +const wts = @import("cuda_weights.zig"); +const kern = @import("cuda_kernels.zig"); + +const DeviceBuffer = cuda.DeviceBuffer; + +/// Rows a prompt chunk takes. +pub const prompt_rows = 1024; +/// Rows a round of windows takes over every stream. +pub const round_rows = 128; +/// Layers of each mixer kind in the admitted schedule. +pub const conv_layers = count(.conv); +pub const full_layers = count(.full_attention); +/// The layers whose keys and values live in pages (liquid's kept-state walk). +pub const paged_layers = full_layers; +pub const sliding_layers = count(.sliding_attention); +/// Elements a KV row holds (4 heads of 128) and a page holds over every full layer's keys and values. +pub const kv_width = c.kv_heads * c.head_dim; +pub const page_elems = full_layers * 2 * kern.page * kv_width; +/// A conv state's rows: the taps before the current row. +pub const state_rows = c.conv_taps - 1; +/// Elements a slot's conv states hold: six bx rows a conv layer. +pub const slot_elems = conv_layers * state_rows * c.hidden; +/// Keys a sliding ring holds: a prompt chunk's rows are written before its first row reads its window back. +pub const ring = (prompt_rows + c.sliding_window + 63) / 64 * 64; +/// Elements a slot's rings hold: every sliding layer's keys and values. +pub const ring_elems = sliding_layers * 2 * ring * kv_width; +/// Key chunks one attention launch walks in parallel at most (nemotron's chunk_blocks). +const chunk_blocks = 64; + +fn count(kind: c.LayerType) usize { + var n: usize = 0; + for (c.schedule) |t| n += @intFromBool(t == kind); + return n; +} + +pub const Scratch = struct { + ids: u64, // int32 [R]: tokens, positions, slots, offsets in their stream's rows, R each + pos: u64, + slot: u64, + off: u64, + meta: u64, // int32 [round_rows]: each chain's first position (AttnSegs' META), its commit slot and row + commit_slot: u64, + commit_row: u64, + h: u64, // bf16 [R, hidden]: the residual + x: u64, // bf16 [R, hidden]: a norm's output + xs: u64, // fp32 [R, hidden / 64]: the decode attention merge's sums (read by nothing) + sums: u64, // fp32 [round_rows, intermediate / 64]: an affine projection's decode rows' 64-value sums + r: u64, // bf16 [R, hidden]: a block's output + wide: u64, // bf16 [R, 2 intermediate]: conv in (3 hidden), q|k|v|gate, w1|w3 + y: u64, // bf16 [R, intermediate]: conv out, the gated attention, the SwiGLU + q: u64, // bf16 [R, hidden]: normed and rotated queries + att: u64, // bf16 [R, hidden]: attention before its head gate + kv: u64, // bf16 [R, 2 kv_width]: keys then values a row + po: u64, // decode attention partials: o, m, l + pm: u64, + pl: u64, + part: u64, // the lane's slice partials (where clusters do not meet them) + snaps: u64, // bf16 [conv_layers][round_rows][6][hidden]: each window row's conv state after it + logits: u64, // bf16 [round_rows, vocab] +}; + +pub const Model = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + w: *const wts.Weights, + context: usize, // positions a slot holds: whole pages + slots: usize, + paged: bool = false, // slots lease pages of the core pool (attach makes the arena), else each owns a window + pages: usize = 0, // pages in the arena + arena: u64 = 0, // [pages][full layer][K, V][512][kv_width] bf16 + table: u64 = 0, // int32 [slots][pages a slot]: slot s's page ids in position order + states: u64 = 0, // bf16 [slots][conv layer][6][hidden] + rings: u64 = 0, // bf16 [sliding layer][K, V][slots][ring][kv_width] + ids: []u32 = &.{}, // the tables on the host + caps: []usize = &.{}, // positions slot s's pages cover + inv: [2]u64 = .{ 0, 0 }, // fp32 [head_dim / 2]: RoPE's inverse frequencies, full then sliding + s: Scratch = undefined, + buffers: std.ArrayList(DeviceBuffer) = .empty, + planned: ?usize = null, // a plan: the bytes counted, nothing allocated + + pub fn perSlot(m: *const Model) usize { + return m.context / kern.page; + } + + /// Positions a slot holds for a window of `context`: whole pages. + pub fn positions(context: usize) usize { + return (context + kern.page - 1) / kern.page * kern.page; + } + + /// Bytes a stream's fixed part takes: its conv states and sliding rings. + pub const fixed_bytes = (slot_elems + ring_elems) * 2; + + /// A slot's regions outside its pages (liquid's walk): conv states, then each sliding layer's K and V rings. + pub fn fixed(m: *const Model, slot: usize, i: usize) ?[2]u64 { + const ring_bytes = ring * kv_width * 2; + if (i == 0) return .{ m.states + slot * slot_elems * 2, slot_elems * 2 }; + if (i > 2 * sliding_layers) return null; + return .{ m.ringOf((i - 1) / 2, (i - 1) % 2 == 1) + slot * ring_bytes, ring_bytes }; + } + + /// Bytes one stream's caches take at `context` positions: its pages, conv states and rings. + pub fn streamBytes(context: usize) usize { + return positions(context) / kern.page * page_elems * 2 + fixed_bytes; + } + + /// What the host's pool plans: a stream's conv states, rings and table row, a token its rows of the six layers. + pub fn layout(m: *const Model) kv.Layout { + return .{ .fixed_bytes = fixed_bytes + m.perSlot() * 4, .token_bytes = page_elems * 2 / kern.page, .page_tokens = kern.page, .snap_bytes = 0 }; + } + + /// Slots a backend holds: the model's, or (paged) as many as a round takes. + pub fn slotCap(m: *const Model) usize { + return if (m.paged) round_rows else m.slots; + } + + /// Decode attention's key chunks a row walks at most, and the blocks that walk them. + pub fn chunks(m: *const Model) usize { + return m.perSlot(); + } + + pub fn chunkBlocks(m: *const Model) usize { + return @min(m.chunks(), chunk_blocks); + } + + fn alloc(m: *Model, bytes: usize) !u64 { + if (m.planned) |*p| { + p.* += (@max(bytes, 256) + 255) / 256 * 256; + return 0; + } + const b = try DeviceBuffer.alloc(m.d, @max(bytes, 256)); + errdefer { + var x = b; + x.free(); + } + try m.buffers.append(m.gpa, b); + return b.ptr; + } + + /// `slots` streams, each owning a window of pages; paged: one slot's states now, `attach` makes the pages. + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, s: cuda.Stream, w: *const wts.Weights, k: *const kern.Kernels, context: usize, slots: usize, paged: bool) !Model { + var m: Model = .{ .gpa = gpa, .d = d, .w = w, .context = positions(context), .slots = if (paged) 1 else slots, .paged = paged }; + errdefer m.deinit(); + try m.build(s, k); + return m; + } + + /// The device bytes `init` allocates for these sizes, counted by the same path without allocating. + pub fn openBytes(gpa: std.mem.Allocator, d: *const cuda.Driver, s: cuda.Stream, w: *const wts.Weights, k: *const kern.Kernels, context: usize, slots: usize, paged: bool) !usize { + var m: Model = .{ .gpa = gpa, .d = d, .w = w, .context = positions(context), .slots = if (paged) 1 else slots, .paged = paged, .planned = 0 }; + defer m.deinit(); + try m.build(s, k); + return m.planned.?; + } + + /// Makes the arena of `pool` pages to lease and the states, rings and tables of `extra + 1` slots (none mapped). + pub fn attach(m: *Model, s: cuda.Stream, extra: usize, pool: usize) !void { + if (!m.paged or m.pages != 0 or extra >= round_rows) return error.BadAttach; + m.forget(m.table); + m.forget(m.states); + m.forget(m.rings); + m.gpa.free(m.ids); + m.gpa.free(m.caps); + m.ids = &.{}; + m.caps = &.{}; + m.slots = extra + 1; + m.pages = pool; + m.arena = try m.alloc(pool * page_elems * 2); + try m.d.check(m.d.api.cuMemsetD8Async(m.arena, 0, pool * page_elems * 2, s.handle), "cuMemsetD8Async"); + try m.makeSlots(s); + try s.synchronize(); + } + + /// Slot `slot`'s pages from table entry `from` on are `ids` (arena page numbers in position order). + pub fn mapPages(m: *Model, s: cuda.Stream, slot: usize, ids: []const u32, from: usize) !void { + const at = try m.note(slot, ids, from); + try m.d.check(m.d.api.cuMemcpyHtoDAsync_v2(m.table + at * 4, ids.ptr, ids.len * 4, s.handle), "cuMemcpyHtoDAsync"); + } + + /// mapPages on the host: the table's copy and the slot's reach; the entry its device copy starts at. + fn note(m: *Model, slot: usize, ids: []const u32, from: usize) !usize { + const per = m.perSlot(); + if (slot >= m.slots or from + ids.len > per) return error.BadChain; + for (ids) |id| if (id >= m.pages) return error.BadChain; + const at = slot * per + from; + @memcpy(m.ids[at..][0..ids.len], ids); + m.caps[slot] = (from + ids.len) * kern.page; + return at; + } + + /// Slot `slot` holds no pages: its table entries stay, no forward may reach them. + pub fn unmapPages(m: *Model, slot: usize) void { + m.caps[slot] = 0; + } + + /// Frees one buffer `alloc` made. + fn forget(m: *Model, ptr: u64) void { + for (m.buffers.items, 0..) |*b, i| if (b.ptr == ptr) { + b.free(); + _ = m.buffers.swapRemove(i); + return; + }; + } + + /// The slots' tables, conv states and rings; a window of pages a slot unless paged (their pages come from leases). + fn makeSlots(m: *Model, s: cuda.Stream) !void { + const per = m.perSlot(); + m.table = try m.alloc(m.slots * per * 4); + m.states = try m.alloc(m.slots * slot_elems * 2); + m.rings = try m.alloc(m.slots * ring_elems * 2); + if (m.planned != null) return; + m.ids = try m.gpa.alloc(u32, m.slots * per); + m.caps = try m.gpa.alloc(usize, m.slots); + for (m.ids, 0..) |*v, i| v.* = if (m.paged) 0 else @intCast(i); + @memset(m.caps, if (m.paged) 0 else m.context); + try m.d.check(m.d.api.cuMemcpyHtoDAsync_v2(m.table, m.ids.ptr, m.ids.len * 4, s.handle), "cuMemcpyHtoDAsync"); + try (DeviceBuffer{ .d = m.d, .ptr = m.states, .len = m.slots * slot_elems * 2, .borrowed = true }).fill8(0, s.handle); + try (DeviceBuffer{ .d = m.d, .ptr = m.rings, .len = m.slots * ring_elems * 2, .borrowed = true }).fill8(0, s.handle); + try s.synchronize(); + } + + /// The lane's slice partials any projection of up to a round's rows needs (0 where clusters meet them). + fn partBytes(m: *const Model, k: *const kern.Kernels) usize { + var most: usize = 0; + const w = m.w; + for (1..round_rows + 1) |rows| { + most = @max(most, partOf(k, rows, w.head)); + for (w.layers) |l| { + for ([_]wts.Proj{ l.gate_up, l.down }) |p| most = @max(most, partOf(k, rows, p)); + switch (l.mixer) { + .conv => |cv| for ([_]wts.Proj{ cv.in, cv.out }) |p| { + most = @max(most, partOf(k, rows, p)); + }, + .attention => |a| for ([_]wts.Proj{ a.qkvg, a.gate orelse a.out, a.out }) |p| { + most = @max(most, partOf(k, rows, p)); + }, + } + } + } + return most; + } + + fn build(m: *Model, s: cuda.Stream, k: *const kern.Kernels) !void { + if (!m.paged) m.pages = m.slots * m.perSlot(); + if (m.pages > 0) m.arena = try m.alloc(m.pages * page_elems * 2); + try m.makeSlots(s); + for (&m.inv) |*p| p.* = try m.alloc(c.head_dim / 2 * 4); + const r = prompt_rows; + const h = c.hidden; + const parts = round_rows * m.chunks() * c.query_heads; + m.s = .{ + .ids = try m.alloc(r * 4), + .pos = try m.alloc(r * 4), + .slot = try m.alloc(r * 4), + .off = try m.alloc(r * 4), + .meta = try m.alloc(round_rows * 4), + .commit_slot = try m.alloc(round_rows * 4), + .commit_row = try m.alloc(round_rows * 4), + .h = try m.alloc(r * h * 2), + .x = try m.alloc(r * h * 2), + .xs = try m.alloc(r * (h / 64) * 4), + .sums = try m.alloc(round_rows * (c.intermediate / 64) * 4), + .r = try m.alloc(r * h * 2), + .wide = try m.alloc(r * 2 * c.intermediate * 2), + .y = try m.alloc(r * c.intermediate * 2), + .q = try m.alloc(r * h * 2), + .att = try m.alloc(r * h * 2), + .kv = try m.alloc(r * 2 * kv_width * 2), + .po = try m.alloc(parts * c.head_dim * 4), + .pm = try m.alloc(parts * 4), + .pl = try m.alloc(parts * 4), + .part = try m.alloc(m.partBytes(k)), + .snaps = try m.alloc(conv_layers * round_rows * state_rows * h * 2), + .logits = try m.alloc(round_rows * c.vocab * 2), + }; + if (m.planned != null) return; + const cfg = &m.w.config; + for ([_]c.Rotary{ cfg.full_rotary, cfg.sliding_rotary }, m.inv) |rot, p| { + const inv = inverseFrequencies(rot); + try (DeviceBuffer{ .d = m.d, .ptr = p, .len = inv.len * 4, .borrowed = true }).upload(0, std.mem.sliceAsBytes(&inv)); + } + } + + pub fn deinit(m: *Model) void { + for (m.buffers.items) |*b| b.free(); + m.buffers.deinit(m.gpa); + m.gpa.free(m.ids); + m.gpa.free(m.caps); + } + + /// Full layer a's key (or value) cache as the kernels take it: its part of page 0. + pub fn cache(m: *const Model, a: usize, value: bool) u64 { + return m.arena + ((2 * a + @intFromBool(value)) * kern.page * kv_width) * 2; + } + + /// Sliding layer l's key (or value) rings: [slots][ring][kv_width], slot s's at s * ring rows. + pub fn ringOf(m: *const Model, l: usize, value: bool) u64 { + return m.rings + (2 * l + @intFromBool(value)) * m.slots * ring * kv_width * 2; + } + + /// Slot s's page table row. + pub fn tableOf(m: *const Model, slot: usize) u64 { + return m.table + slot * m.perSlot() * 4; + } + + /// Zero slot s's conv states, queued on `s` (rings and pages are read only where the new stream wrote). + pub fn clearSlot(m: *const Model, s: cuda.Stream, slot: usize) !void { + try (DeviceBuffer{ .d = m.d, .ptr = m.states + slot * slot_elems * 2, .len = slot_elems * 2, .borrowed = true }).fill8(0, s.handle); + } +}; + +/// A lane weight's slice partials at `rows` (affine-4 projections take none). +fn partOf(k: *const kern.Kernels, rows: usize, p: wts.Proj) usize { + return switch (p) { + .nvfp4, .bf16, .fp8g, .fp8, .mxfp8 => |q| k.lane.partBytes(rows, q), + else => 0, + }; +} + +/// RoPE's inverse frequencies as torch's float32 rows: 1 / theta^(2i / 128) rotated, zero past them (cos 1, sin 0). +pub fn inverseFrequencies(rot: c.Rotary) [c.head_dim / 2]f32 { + var out: [c.head_dim / 2]f32 = undefined; + const angles: usize = @intFromFloat(rot.partial * @as(f32, c.head_dim) / 2); + for (&out, 0..) |*v, i| { + if (i >= angles) { + v.* = 0; + continue; + } + const e: f32 = @as(f32, @floatFromInt(2 * i)) / @as(f32, c.head_dim); + const pw: f32 = @floatCast(std.math.pow(f64, rot.theta, e)); + v.* = 1.0 / pw; + } + return out; +} + +test "the admitted schedule: 13 conv, 6 full and 5 sliding layers; a page 6 MiB, a ring a chunk and a window" { + try std.testing.expectEqual(@as(usize, 13), conv_layers); + try std.testing.expectEqual(@as(usize, 6), full_layers); + try std.testing.expectEqual(@as(usize, 5), sliding_layers); + try std.testing.expectEqual(@as(usize, 6 << 20), page_elems * 2); + try std.testing.expectEqual(@as(usize, 1280), ring); + try std.testing.expectEqual(@as(usize, 2 * (6 << 20) + Model.fixed_bytes), Model.streamBytes(1000)); + try std.testing.expectEqual(@as(usize, 13 * 6 * 2048 * 2), slot_elems * 2); +} + +test "RoPE tables: full layers rotate 32 pairs at theta 1e6 over the full head dim, sliding ones all 64 at 1e4" { + const full = inverseFrequencies(.{ .theta = 1e6, .partial = 0.5 }); + try std.testing.expectEqual(@as(f32, 1), full[0]); + const want31 = 1.0 / std.math.pow(f64, 1e6, 62.0 / 128.0); + try std.testing.expect(@abs(full[31] - want31) <= want31 * 0x1p-23); // one float32 rounding + for (full[32..]) |v| try std.testing.expectEqual(@as(f32, 0), v); + const sliding = inverseFrequencies(.{ .theta = 1e4, .partial = 1 }); + const want63 = 1.0 / std.math.pow(f64, 1e4, 126.0 / 128.0); + try std.testing.expect(@abs(sliding[63] - want63) <= want63 * 0x1p-23); +} + +test "the pool's layout: 12 KiB a token, a stream's conv states, rings and table row fixed" { + var m: Model = undefined; + m.context = Model.positions(32768); + const l = m.layout(); + try std.testing.expectEqual(@as(u64, 6 << 20), l.pageBytes()); + try std.testing.expectEqual(@as(u64, 12 << 10), l.token_bytes); + try std.testing.expectEqual(@as(u64, Model.fixed_bytes + 64 * 4), l.fixed_bytes); + const p = try kv.plan(l, .{ .room = 20 << 30, .engine = 0, .lanes = 8, .cache = 4 << 30, .window = 32768 + 16, .engine_paged = true }); + try std.testing.expectEqual(@as(u32, 8), p.lanes); + try std.testing.expect(p.pages >= l.pages(32768 + 16)); +} + +test "leases map a slot's pages in position order; its reach ends at them; an id past the arena is refused" { + var ids: [8]u32 = @splat(0); + var caps = [_]usize{ 0, 0 }; + var m: Model = undefined; + m.context = 4 * kern.page; + m.slots = 2; + m.pages = 16; + m.ids = &ids; + m.caps = ∩︀ + _ = try m.note(1, &.{ 9, 3, 4 }, 0); + try std.testing.expectEqualSlices(u32, &.{ 9, 3, 4 }, ids[4..7]); + try std.testing.expectEqual(@as(usize, 3 * kern.page), caps[1]); + try std.testing.expectError(error.BadChain, m.note(0, &.{16}, 0)); + try std.testing.expectError(error.BadChain, m.note(0, &.{ 1, 2, 3, 4, 5 }, 0)); +} diff --git a/zig/src/families/lfm3/cuda_native.zig b/zig/src/families/lfm3/cuda_native.zig new file mode 100644 index 0000000000..1ca6394733 --- /dev/null +++ b/zig/src/families/lfm3/cuda_native.zig @@ -0,0 +1,69 @@ +//! LFM3-1.7B (MLX 4-bit; NVFP4 and block FP8 weight-only) behind the native CUDA registry: drafts, prompt reuse. +const std = @import("std"); +const liquid = @import("liquid_cuda"); +const be = @import("cuda_backend.zig"); +const snap = @import("cuda_snap.zig"); + +const N = liquid.native.Native(struct { + pub const wts = @import("cuda_weights.zig"); + pub const kern = @import("cuda_kernels.zig"); + pub const mdl = @import("cuda_model.zig"); + pub const be = @import("cuda_backend.zig"); + pub const snap = @import("cuda_snap.zig"); + pub const pages_env = "TENSORFOLD_LFM3_KV_PAGES"; + pub fn check(_: anytype) !void {} // the admission already fixed the layer schedule +}); + +pub const model_type = "lfm3"; +pub const formats: []const []const u8 = &.{ "mlx-q4g64", "modelopt-nvfp4", "fp8-block128" }; +pub const default_context: i64 = 32768; +pub const max_segments: u32 = 1; +pub const explains_open = true; +pub const prompt_rows = N.prompt_rows; +pub const snap_lookahead = N.snap_lookahead; +pub const window_rows = N.window_rows; +pub const Room = N.Room; +pub const Options = N.Options; +pub const Loaded = N.Loaded; +pub const open = N.open; +pub const attach = N.attach; +pub const fits = N.fits; +pub const snap_bytes = N.snap_bytes; +pub const snap_ready = N.snap_ready; +pub const snap_charged = N.snap_charged; +pub const snap_spare = N.snap_spare; +pub const snap_reuses = N.snap_reuses; +pub const snap_save = N.snap_save; +pub const snap_restore = N.snap_restore; +pub const snap_drop = N.snap_drop; + +pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { + return switch (err) { + error.PastContext => "the prompt and its reply exceed this server's context window: shorten it or lower max_tokens", + error.NoFreeSlot => "every stream's caches are in use: retry once another request ends", + error.NoDraftHead => "LFM3 on CUDA has no draft head: it drafts from the prompt's own text", + error.Lfm3TextOnly => "LFM3-1.7B is a text model, and LFM3-VL's image tower is not built on CUDA yet: this request has images", + error.StructuresUnsupported => "structured output is not built for LFM3 on CUDA yet", + error.NonFiniteLogits => "the model's logits went non-finite at this row: the request is stopped rather than answered from them", + error.CapturedKernels => "LFM3 runs on the engine's own kernels: unset TENSORFOLD_CUDA_KERNELS", + error.UnsupportedLfm3Quantization => "LFM3 on CUDA reads the text exports at MLX 4 bits in groups of 64, ModelOpt NVFP4 and block FP8: its 6- and 8-bit MLX exports and LFM3-VL are not loaded yet", + error.UnsupportedLfm3Config, error.UnsupportedLfm3Geometry, error.UnsupportedLfm3LayerSchedule, error.UnsupportedLfm3Rotary, error.UnsupportedLfm3Context, error.BadLfm3Config => "this lfm3 checkpoint's config is not LFM3-1.7B's, which this engine serves", + error.ParallelPastRound => std.fmt.comptimePrint("--parallel: LFM3 on CUDA holds at most {d} streams (a round's rows): serve --parallel {d} or fewer, or auto", .{ be.round_rows, be.round_rows }), + error.ParallelPastMemory => "--parallel: the CUDA memory budget fits fewer LFM3 streams than asked (each stream's caches are allocated at load): lower --parallel or --context, serve --parallel auto, or adjust TENSORFOLD_MEMORY_RESERVE_GIB / TENSORFOLD_CUDA_MEMORY_LIMIT_GB", + error.NoStreamFits => "the CUDA memory budget fits no LFM3 stream after the weights: lower --context, free device memory, or adjust TENSORFOLD_MEMORY_RESERVE_GIB / TENSORFOLD_CUDA_MEMORY_LIMIT_GB", + else => null, + }; +} + +test "every refusal of its own is worded, the rest left to the server" { + for ([_]anyerror{ error.PastContext, error.NoFreeSlot, error.NoDraftHead, error.Lfm3TextOnly, error.StructuresUnsupported, error.NonFiniteLogits, error.CapturedKernels, error.UnsupportedLfm3Quantization, error.UnsupportedLfm3Geometry, error.ParallelPastRound, error.ParallelPastMemory, error.NoStreamFits }) |e| { + try std.testing.expect(explain(null, e) != null); + } + try std.testing.expect(explain(null, error.OutOfMemory) == null); +} + +test "the prompt cache's hooks compile against the store's vtable" { + const pc: struct { bytes: *const fn (*anyopaque, u32) u64, save: *const fn (*anyopaque, ?*anyopaque, u32) anyerror!*anyopaque, restore: *const fn (*anyopaque, ?*anyopaque, *anyopaque) anyerror!void, drop: *const fn (*anyopaque, *anyopaque) void } = .{ .bytes = snap_bytes, .save = snap_save, .restore = snap_restore, .drop = snap_drop }; + try std.testing.expect(@intFromPtr(pc.save) != 0); + try std.testing.expectEqual(snap.bytes(100), snap_bytes(undefined, 100)); +} diff --git a/zig/src/families/lfm3/cuda_safetensors.zig b/zig/src/families/lfm3/cuda_safetensors.zig new file mode 100644 index 0000000000..4a974cb397 --- /dev/null +++ b/zig/src/families/lfm3/cuda_safetensors.zig @@ -0,0 +1,7 @@ +//! core's safetensors types under the name the shared LFM3 admission (weights.zig) imports, for the CUDA family. +const st = @import("core").safetensors; + +pub const DType = st.DType; +pub const Entry = st.Entry; +pub const Header = st.Header; +pub const max_rank = st.max_rank; diff --git a/zig/src/families/lfm3/cuda_snap.zig b/zig/src/families/lfm3/cuda_snap.zig new file mode 100644 index 0000000000..997c00af20 --- /dev/null +++ b/zig/src/families/lfm3/cuda_snap.zig @@ -0,0 +1,76 @@ +//! LFM3's kept states: liquid's over this model (conv states, the sliding layers' rings, the full layers' pages). +const std = @import("std"); +const liquid = @import("liquid_cuda"); +const mdl = @import("cuda_model.zig"); +const kern = @import("cuda_kernels.zig"); + +const S = liquid.snap.States(mdl); +pub const Saved = S.Saved; +pub const bytes = S.bytes; +pub const save = S.save; +pub const restore = S.restore; +pub const drop = S.drop; +pub const page = liquid.snap.page; + +const page_bytes = mdl.page_elems * 2; +const row_bytes = mdl.kv_width * 2; +const state_bytes = mdl.slot_elems * 2; +const ring_bytes = mdl.ring * row_bytes; // one sliding layer's keys (or values) of a slot + +const testing = std.testing; + +const Count = struct { + runs: *std.ArrayList([3]usize), + fn add(c: Count, dev: u64, off: usize, len: usize) anyerror!void { + try c.runs.append(testing.allocator, .{ dev, off, len }); + } +}; + +fn fake(ids: []u32, caps: []usize) mdl.Model { + var m: mdl.Model = undefined; + m.context = 4 * kern.page; + m.slots = caps.len; + m.arena = 1 << 40; + m.states = 1 << 36; + m.rings = 1 << 38; + m.ids = ids; + m.caps = caps; + return m; +} + +test "a state: conv states, every ring, consecutive whole pages as one copy, the last page's rows a layer's K and V" { + // slot 1 leases pages 7, 8, 3, 9: positions 0..1023 one run, 1024..1535 a page, 1536.. a partial page + var ids = [_]u32{ 0, 0, 0, 0, 7, 8, 3, 9 }; + var caps = [_]usize{ 0, 4 * kern.page }; + const m = fake(&ids, &caps); + const at = 3 * kern.page + 100; + var runs: std.ArrayList([3]usize) = .empty; + defer runs.deinit(testing.allocator); + try S.each(&m, 1, at, null, Count{ .runs = &runs }, Count.add); + const rings = 2 * mdl.sliding_layers; + try testing.expectEqual(@as(usize, 1 + rings + 2 + 12), runs.items.len); + try testing.expectEqual([3]usize{ (1 << 36) + state_bytes, 0, state_bytes }, runs.items[0]); + // sliding layer 0's values: the second ring block, slot 1's ring in it + try testing.expectEqual([3]usize{ (1 << 38) + 2 * mdl.ring * row_bytes + ring_bytes, state_bytes + ring_bytes, ring_bytes }, runs.items[2]); + const fixed = mdl.Model.fixed_bytes; + try testing.expectEqual([3]usize{ (1 << 40) + 7 * page_bytes, fixed, 2 * page_bytes }, runs.items[1 + rings]); + try testing.expectEqual([3]usize{ (1 << 40) + 3 * page_bytes, fixed + 2 * page_bytes, page_bytes }, runs.items[2 + rings]); + const v9 = (1 << 40) + 9 * page_bytes + kern.page * row_bytes; // page 9, full layer 0's values + try testing.expectEqual([3]usize{ v9, fixed + 3 * page_bytes + 100 * row_bytes, 100 * row_bytes }, runs.items[4 + rings]); + var total: usize = 0; + for (runs.items) |r| total += r[2]; + try testing.expectEqual(fixed + at * 12 * 1024, total); + try testing.expectEqual((total + page - 1) / page * page, S.bytes(at)); +} + +test "a state on whole pages has no partial copies; one past a slot's pages is refused" { + var ids = [_]u32{ 4, 5, 6, 7 }; + var caps = [_]usize{2 * kern.page}; + const m = fake(&ids, &caps); + var runs: std.ArrayList([3]usize) = .empty; + defer runs.deinit(testing.allocator); + try S.each(&m, 0, 2 * kern.page, null, Count{ .runs = &runs }, Count.add); + try testing.expectEqual(@as(usize, 2 + 2 * mdl.sliding_layers), runs.items.len); + try testing.expectEqual(@as(usize, 2 * page_bytes), runs.items[runs.items.len - 1][2]); + try testing.expectError(error.PastPages, S.each(&m, 0, 2 * kern.page + 1, null, Count{ .runs = &runs }, Count.add)); +} diff --git a/zig/src/families/lfm3/cuda_weights.zig b/zig/src/families/lfm3/cuda_weights.zig new file mode 100644 index 0000000000..1d1dcadc1c --- /dev/null +++ b/zig/src/families/lfm3/cuda_weights.zig @@ -0,0 +1,331 @@ +//! LFM3's weights on CUDA: MLX affine-4, or weight-only NVFP4 (FFN) or block FP8 (FFN, q/k/v/o) with bf16 the rest. +const std = @import("std"); +const core = @import("core"); +const cuda = @import("cuda"); +const c = @import("config.zig"); +const admit = @import("weights.zig"); +const kern = @import("cuda_kernels.zig"); + +const DeviceBuffer = cuda.DeviceBuffer; +pub const Proj = cuda.qlinear.Weight; + +/// Rows of the stacked attention projection: q (16 x 128), k and v (4 x 128 each), the 16 head gates. +pub const qkvg_rows = c.query_heads * c.head_dim + 2 * c.kv_heads * c.head_dim + c.query_heads; + +pub const Conv = struct { in: Proj, out: Proj, weight: u64 }; +/// q, k, v and the head gates as one projection (or q|k|v and `gate`); the q and k norms bf16 [128] each, q's first. +pub const Attention = struct { qkvg: Proj, gate: ?Proj = null, out: Proj, norms: u64, sliding: bool }; + +pub const Layer = struct { + operator_norm: u64, + post_operator_norm: u64, + ffn_norm: u64, + post_ffn_norm: u64, + mixer: union(enum) { conv: Conv, attention: Attention }, + /// w1 | w3 stacked (gate rows, then up rows), and w2. + gate_up: Proj, + down: Proj, +}; + +pub const Weights = struct { + gpa: std.mem.Allocator, + config: c.Config, + /// The embedding as stored, bf16 [V, H] (the embed rows' table); the head is its lane copy. + embed: u64 = 0, + /// MLX: the embedding's words, scales and biases as stored (the embed rows' table); the head is their tiles. + embed_affine: ?[3]u64 = null, + head: Proj = undefined, + final_norm: u64 = 0, + layers: []Layer = &.{}, + buffers: std.ArrayList(DeviceBuffer) = .empty, + + pub fn deinit(w: *Weights) void { + for (w.buffers.items) |*b| b.free(); + w.buffers.deinit(w.gpa); + w.gpa.free(w.layers); + w.config.deinit(); + w.* = undefined; + } + + fn alloc(w: *Weights, d: *const cuda.Driver, bytes: usize) !u64 { + const b = try DeviceBuffer.alloc(d, @max(bytes, 256)); + errdefer { + var x = b; + x.free(); + } + try w.buffers.append(w.gpa, b); + return b.ptr; + } + + fn upload(w: *Weights, d: *const cuda.Driver, bytes: []const u8) !u64 { + const b = try DeviceBuffer.fromHost(d, bytes); + errdefer { + var x = b; + x.free(); + } + try w.buffers.append(w.gpa, b); + return b.ptr; + } +}; + +/// The largest staged projection: conv in (3 x 2048 rows) in bf16, or the FFN's w1 and w3 as FP8 codes. +pub const stage_bytes = @max(3 * c.hidden * c.hidden * 2, 2 * c.intermediate * c.hidden) + 4 * 256; + +const Ctx = struct { + w: *Weights, + d: *const cuda.Driver, + s: cuda.Stream, + k: *const kern.Kernels, + ck: *core.checkpoint.Checkpoint, + stage: DeviceBuffer, + + fn tensor(x: *Ctx, name: []const u8, dtype: core.checkpoint.DType, shape: []const usize) !u64 { + return x.w.upload(x.d, (try x.ck.expect(name, dtype, shape)).bytes); + } + + /// bf16 projections [n_i, k] stacked along n, packed as one `.bf16` lane weight. + fn bf16(x: *Ctx, names: []const []const u8, ns: []const usize, k: usize) !Proj { + var n: usize = 0; + var buf: [192]u8 = undefined; + for (names, ns) |name, rows| { + const t = try x.ck.expect(try std.fmt.bufPrint(&buf, "{s}.weight", .{name}), .bf16, &.{ rows, k }); + if ((n + rows) * k * 2 > x.stage.len) return error.StageTooSmall; + try x.stage.upload(n * k * 2, t.bytes); + n += rows; + } + return x.pack(x.stage.ptr, n, k); + } + + /// bf16 rows [n, k] on the device as a `.bf16` lane weight. + fn pack(x: *Ctx, src: u64, n: usize, k: usize) !Proj { + const out = try x.w.alloc(x.d, cuda.modelopt.padded(n) * k * 2); + try x.k.pack.bf16Lane(x.s, src, out, n, k); + try x.s.synchronize(); // the next projection's staging overwrites this one's + return .{ .bf16 = cuda.modelopt.bf16Weight(out, n, k) }; + } + + /// MLX affine-4 projections (g64) [n_i, k] stacked along n, packed as one into the affine tiles. + fn affine(x: *Ctx, names: []const []const u8, ns: []const usize, k: usize) !Proj { + var buf: [192]u8 = undefined; + var n: usize = 0; + for (ns) |rows| n += rows; + const specs = [_]struct { []const u8, core.checkpoint.DType, usize }{ .{ "weight", .u32, k / 8 }, .{ "scales", .bf16, k / 64 }, .{ "biases", .bf16, k / 64 } }; + var raw: [3]u64 = undefined; + var at: usize = 0; + for (specs, &raw) |sp, *r| { + r.* = x.stage.ptr + at; + for (names, ns) |name, rows| { + const t = try x.ck.expect(try std.fmt.bufPrint(&buf, "{s}.{s}", .{ name, sp[0] }), sp[1], &.{ rows, sp[2] }); + if (at + t.bytes.len > x.stage.len) return error.StageTooSmall; + try x.stage.upload(at, t.bytes); + at += t.bytes.len; + } + at = (at + 255) / 256 * 256; + } + const q = try x.affineFrom(raw, n, k); + try x.s.synchronize(); // the next projection's staging overwrites this one's + return q; + } + + /// Affine-4 tiles of words, scales and biases [n, k] already on the device. + fn affineFrom(x: *Ctx, raw: [3]u64, n: usize, k: usize) !Proj { + const l = cuda.qlinear.Affine4.layout(n, k); + const q: cuda.qlinear.Affine4 = .{ .w = try x.w.alloc(x.d, l.words), .s = try x.w.alloc(x.d, l.scales), .b = try x.w.alloc(x.d, l.scales), .n = n, .k = k, .npad = l.npad }; + try x.k.affine.pack(x.s, raw[0], raw[1], raw[2], q); + return .{ .affine4 = q }; + } + + /// NVFP4 projections [n_i, k] sharing one tensor scale, stacked along n as one `.nvfp4` lane weight. + fn nvfp4(x: *Ctx, names: []const []const u8, ns: []const usize, k: usize) !Proj { + var buf: [192]u8 = undefined; + var n: usize = 0; + for (ns) |rows| n += rows; + const codes_bytes = n * k / 2; + if (codes_bytes + n * k / 16 > x.stage.len) return error.StageTooSmall; + var global: ?f32 = null; + var at: usize = 0; + for (names, ns) |name, rows| { + const codes = try x.ck.expect(try std.fmt.bufPrint(&buf, "{s}.weight", .{name}), .u8, &.{ rows, k / 2 }); + const blocks = try x.ck.expect(try std.fmt.bufPrint(&buf, "{s}.weight_scale", .{name}), .f8_e4m3, &.{ rows, k / 16 }); + const g = try scalar(x.ck, try std.fmt.bufPrint(&buf, "{s}.weight_scale_2", .{name})); + // weight-only: the static FP4 input scale is read (admission requires it), not applied; rows stay bf16 + _ = try scalar(x.ck, try std.fmt.bufPrint(&buf, "{s}.input_scale", .{name})); + if (global) |v| if (v != g) return error.UnsupportedLfm3Quantization; // one scale a stacked weight + global = g; + try x.stage.upload(at * k / 2, codes.bytes); + try x.stage.upload(codes_bytes + at * k / 16, blocks.bytes); + at += rows; + } + const words = try x.w.alloc(x.d, cuda.modelopt.padded(n) * k / 2); + const tiles = try x.w.alloc(x.d, cuda.modelopt.padded(n) * k / 16); + try x.k.pack.fp4Lane(x.s, x.stage.ptr, x.stage.ptr + codes_bytes, words, tiles, n, k); + try x.s.synchronize(); + return .{ .nvfp4 = cuda.nvfp4.weight(words, tiles, global.?, n, k) }; + } + + /// Block FP8 projections [n_i, k] stacked along n as one `.fp8g` lane weight, a block's scale each 64 inputs. + fn fp8(x: *Ctx, names: []const []const u8, ns: []const usize, k: usize) !Proj { + const b = c.fp8_block; + var buf: [192]u8 = undefined; + var n: usize = 0; + for (ns) |rows| n += rows; + if (n * k > x.stage.len) return error.StageTooSmall; + const blocks = try x.w.gpa.alloc(f32, n / b * (k / b)); + defer x.w.gpa.free(blocks); + var at: usize = 0; + for (names, ns) |name, rows| { + const codes = try x.ck.expect(try std.fmt.bufPrint(&buf, "{s}.weight", .{name}), .f8_e4m3, &.{ rows, k }); + const scales = try x.ck.expect(try std.fmt.bufPrint(&buf, "{s}.weight_scale", .{name}), .bf16, &.{ rows / b, k / b }); + try x.stage.upload(at * k, codes.bytes); + for (std.mem.bytesAsSlice(u16, scales.bytes), blocks[at / b * (k / b) ..][0 .. rows / b * (k / b)]) |v, *o| + o.* = @bitCast(@as(u32, v) << 16); + at += rows; + } + const words = try x.w.alloc(x.d, cuda.modelopt.padded(n) * k); + try x.k.pack.fp8Lane(x.s, x.stage.ptr, words, n, k); + const tiles = try x.w.gpa.alloc(f32, cuda.fp8.padded(n) * (k / cuda.fp8.group)); + defer x.w.gpa.free(tiles); + cuda.fp8.packScales(tiles, blocks, n, k); + const sd = try x.w.upload(x.d, std.mem.sliceAsBytes(tiles)); + try x.s.synchronize(); + return .{ .fp8g = cuda.fp8.weight(words, sd, n, k) }; + } +}; + +/// An fp32 tensor of one element. +fn scalar(ck: *core.checkpoint.Checkpoint, name: []const u8) !f32 { + const t = try ck.get(name); + if (t.dtype != .f32 or t.bytes.len != 4) return error.UnexpectedTensor; + return @bitCast(std.mem.readInt(u32, t.bytes[0..4], .little)); +} + +/// The text decoder's tensors against the shared admission (weights.zig): names, dtypes, shapes, bytes, nothing extra. +fn admitted(gpa: std.mem.Allocator, config: *const c.Config, ck: *const core.checkpoint.Checkpoint) !admit.Namespace { + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + var all: core.safetensors.Header = .empty; + for (ck.files.items) |*f| for (f.names.keys(), f.names.values()) |name, e| try all.put(arena.allocator(), name, e); + return (try admit.validate(gpa, config, &all)).namespace; +} + +/// The MLX widths CUDA serves: affine 4 bits in groups of 64 everywhere (6 and 8 bits, g32 and VL are refused). +fn mlxServed(config: *const c.Config) bool { + return config.quant.global.bits == 4 and config.quant.global.group == 64 and config.quant.modules.count() == 0; +} + +fn named(b: []u8, prefix: []const u8, suffix: []const u8) ![]const u8 { + return std.fmt.bufPrint(b, "{s}{s}", .{ prefix, suffix }); +} + +/// Every tensor checked against the admitted geometry and packed for the lane; `layers` caps them (tests). +pub fn load(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, s: cuda.Stream, k: *const kern.Kernels, dir: []const u8, layers: ?usize) !Weights { + var config = try c.Config.read(gpa, io, dir); + if (config.kind != .text or (config.scheme == .mlx_affine and !mlxServed(&config))) { + config.deinit(); + return error.UnsupportedLfm3Quantization; + } + var w: Weights = .{ .gpa = gpa, .config = config }; + errdefer w.deinit(); + var ck = try core.checkpoint.Checkpoint.openModel(gpa, io, dir); + defer ck.close(); + const pre = (try admitted(gpa, &w.config, &ck)).prefix(); + const fp8 = w.config.scheme == .fp8_block; + const mlx = w.config.scheme == .mlx_affine; + if (!@import("builtin").is_test) { + if (fp8) std.log.info("LFM3 FP8 on CUDA serves weight-only: FP8 weights, bf16 activations; the export's dynamic FP8 input quantization is not applied", .{}); + if (!fp8) std.log.info("LFM3 NVFP4 on CUDA serves weight-only: FP4 weights, bf16 activations; the export's static FP4 input scales are read and not applied", .{}); + } + const h = c.hidden; + var stage = try DeviceBuffer.alloc(d, stage_bytes); + defer stage.free(); + var x: Ctx = .{ .w = &w, .d = d, .s = s, .k = k, .ck = &ck, .stage = stage }; + var b1: [128]u8 = undefined; + var b2: [128]u8 = undefined; + var b3: [128]u8 = undefined; + var b4: [128]u8 = undefined; + if (mlx) { + var e: [3]u64 = undefined; + e[0] = try x.tensor(try named(&b1, pre, "embed_tokens.weight"), .u32, &.{ c.vocab, h / 8 }); + e[1] = try x.tensor(try named(&b1, pre, "embed_tokens.scales"), .bf16, &.{ c.vocab, h / 64 }); + e[2] = try x.tensor(try named(&b1, pre, "embed_tokens.biases"), .bf16, &.{ c.vocab, h / 64 }); + w.embed_affine = e; + w.head = try x.affineFrom(e, c.vocab, h); // tied: the embedding's own codes as tiles + try s.synchronize(); + } else { + w.embed = try x.tensor(try named(&b1, pre, "embed_tokens.weight"), .bf16, &.{ c.vocab, h }); + w.head = try x.pack(w.embed, c.vocab, h); // tied: the lane's copy of the embedding + } + w.final_norm = try x.tensor(try named(&b1, pre, "embedding_norm.weight"), .bf16, &.{h}); + const n = @min(layers orelse c.layers, c.layers); + w.layers = try gpa.alloc(Layer, n); + for (w.layers, w.config.layer_types[0..n], 0..) |*l, kind, i| { + var pb: [64]u8 = undefined; + const p = try std.fmt.bufPrint(&pb, "{s}layers.{d}.", .{ pre, i }); + l.operator_norm = try x.tensor(try named(&b1, p, "operator_norm.weight"), .bf16, &.{h}); + l.post_operator_norm = try x.tensor(try named(&b1, p, "post_operator_norm.weight"), .bf16, &.{h}); + l.ffn_norm = try x.tensor(try named(&b1, p, "ffn_norm.weight"), .bf16, &.{h}); + l.post_ffn_norm = try x.tensor(try named(&b1, p, "post_ffn_norm.weight"), .bf16, &.{h}); + if (kind == .conv) { + const in = try named(&b1, p, "conv.in_proj"); + const out = try named(&b2, p, "conv.out_proj"); + // MLX keeps the depthwise kernel as [channels, taps, 1], torch as [channels, 1, taps]: the same values + const taps: []const usize = if (mlx) &.{ h, c.conv_taps, 1 } else &.{ h, 1, c.conv_taps }; + l.mixer = .{ .conv = .{ + .in = if (mlx) try x.affine(&.{in}, &.{3 * h}, h) else try x.bf16(&.{in}, &.{3 * h}, h), + .out = if (mlx) try x.affine(&.{out}, &.{h}, h) else try x.bf16(&.{out}, &.{h}, h), + .weight = try x.tensor(try named(&b3, p, "conv.conv.weight"), .bf16, taps), + } }; + } else { + const qw = c.query_heads * c.head_dim; + const kw = c.kv_heads * c.head_dim; + const norms = try w.alloc(d, 2 * c.head_dim * 2); + const qn = try ck.expect(try named(&b1, p, "self_attn.q_norm.weight"), .bf16, &.{c.head_dim}); + const kn = try ck.expect(try named(&b1, p, "self_attn.k_norm.weight"), .bf16, &.{c.head_dim}); + const nb: DeviceBuffer = .{ .d = d, .ptr = norms, .len = 4 * c.head_dim, .borrowed = true }; + try nb.upload(0, qn.bytes); + try nb.upload(2 * c.head_dim, kn.bytes); + const names = [_][]const u8{ try named(&b1, p, "self_attn.q_proj"), try named(&b2, p, "self_attn.k_proj"), try named(&b3, p, "self_attn.v_proj"), try named(&b4, p, "self_attn.gate_proj") }; + const sliding = kind == .sliding_attention; + var b5: [128]u8 = undefined; // names holds b1 to b4 + const o_proj = try named(&b5, p, "self_attn.o_proj"); + l.mixer = .{ .attention = if (mlx) .{ + .qkvg = try x.affine(&names, &.{ qw, kw, kw, c.query_heads }, h), + .out = try x.affine(&.{o_proj}, &.{h}, qw), + .norms = norms, + .sliding = sliding, + } else if (fp8) .{ + .qkvg = try x.fp8(names[0..3], &.{ qw, kw, kw }, h), + .gate = try x.bf16(names[3..], &.{c.query_heads}, h), + .out = try x.fp8(&.{o_proj}, &.{h}, qw), + .norms = norms, + .sliding = sliding, + } else .{ + .qkvg = try x.bf16(&names, &.{ qw, kw, kw, c.query_heads }, h), + .out = try x.bf16(&.{o_proj}, &.{h}, qw), + .norms = norms, + .sliding = sliding, + } }; + } + const ffn = [_][]const u8{ try named(&b1, p, "feed_forward.w1"), try named(&b2, p, "feed_forward.w3"), try named(&b3, p, "feed_forward.w2") }; + const ii = [_]usize{ c.intermediate, c.intermediate }; + l.gate_up = switch (w.config.scheme) { + .mlx_affine => try x.affine(ffn[0..2], &ii, h), + .fp8_block => try x.fp8(ffn[0..2], &ii, h), + .modelopt_nvfp4 => try x.nvfp4(ffn[0..2], &ii, h), + }; + l.down = switch (w.config.scheme) { + .mlx_affine => try x.affine(ffn[2..], &.{h}, c.intermediate), + .fp8_block => try x.fp8(ffn[2..], &.{h}, c.intermediate), + .modelopt_nvfp4 => try x.nvfp4(ffn[2..], &.{h}, c.intermediate), + }; + } + return w; +} + +test "the stage holds conv in's bf16 rows and the FFN's stacked codes and scales" { + try std.testing.expect(stage_bytes >= 3 * c.hidden * c.hidden * 2); + try std.testing.expect(stage_bytes >= 2 * c.intermediate * c.hidden); + try std.testing.expect(stage_bytes >= 2 * c.intermediate * c.hidden / 2 + 2 * c.intermediate * c.hidden / 16); + try std.testing.expect(stage_bytes >= qkvg_rows * c.hidden * 2); + try std.testing.expectEqual(@as(usize, 3088), qkvg_rows); +} diff --git a/zig/src/families/lfm3/lfm3.zig b/zig/src/families/lfm3/lfm3.zig new file mode 100644 index 0000000000..7372c2676c --- /dev/null +++ b/zig/src/families/lfm3/lfm3.zig @@ -0,0 +1,3 @@ +//! Host checkpoint admission for LFM3; no runtime family is registered by this module. +pub const config = @import("config.zig"); +pub const weights = @import("weights.zig"); diff --git a/zig/src/families/lfm3/weights.zig b/zig/src/families/lfm3/weights.zig new file mode 100644 index 0000000000..e524f88581 --- /dev/null +++ b/zig/src/families/lfm3/weights.zig @@ -0,0 +1,137 @@ +//! Host-only LFM3 text tensor admission; no weights are uploaded and a vision tower is not qualified here. +const std = @import("std"); +const st = @import("safetensors"); +const c = @import("config.zig"); +const Allocator = std.mem.Allocator; + +pub const Namespace = enum { + bare, + wrapped, + + pub fn prefix(n: Namespace) []const u8 { + return if (n == .bare) "model." else "language_model.model."; + } +}; + +pub const Expected = struct { + name: []const u8, + dtype: st.DType, + rank: u8, + shape: [st.max_rank]usize, +}; + +pub const Manifest = struct { namespace: Namespace, tensors: usize, bytes: usize }; + +fn dense(a: Allocator, out: *std.ArrayList(Expected), name: []const u8, shape: []const usize, dtype: st.DType) !void { + var e: Expected = .{ .name = try a.dupe(u8, name), .dtype = dtype, .rank = @intCast(shape.len), .shape = @splat(1) }; + @memcpy(e.shape[0..shape.len], shape); + try out.append(a, e); +} + +fn named(a: Allocator, out: *std.ArrayList(Expected), prefix: []const u8, suffix: []const u8, shape: []const usize) !void { + var buf: [256]u8 = undefined; + try dense(a, out, try std.fmt.bufPrint(&buf, "{s}{s}", .{ prefix, suffix }), shape, .bf16); +} + +fn affine(a: Allocator, out: *std.ArrayList(Expected), cfg: *const c.Config, prefix: []const u8, suffix: []const u8, n: usize, k: usize) !void { + if (cfg.scheme == .modelopt_nvfp4) return modelopt(a, out, prefix, suffix, n, k); + if (cfg.scheme == .fp8_block) return fp8Block(a, out, prefix, suffix, n, k); + var buf: [256]u8 = undefined; + const module = try std.fmt.bufPrint(&buf, "{s}{s}", .{ prefix, suffix }); + const q = try cfg.width(module); + if (k % q.group != 0 or k * q.bits % 32 != 0) return error.UnsupportedQuantization; + var name: [272]u8 = undefined; + try dense(a, out, try std.fmt.bufPrint(&name, "{s}.weight", .{module}), &.{ n, k * q.bits / 32 }, .u32); + try dense(a, out, try std.fmt.bufPrint(&name, "{s}.scales", .{module}), &.{ n, k / q.group }, .bf16); + try dense(a, out, try std.fmt.bufPrint(&name, "{s}.biases", .{module}), &.{ n, k / q.group }, .bf16); +} + +/// ModelOpt NVFP4: the FFN's codes [n, k / 2], e4m3 scales [n, k / 16], fp32 scalar scales; the rest bf16 [n, k]. +fn modelopt(a: Allocator, out: *std.ArrayList(Expected), prefix: []const u8, suffix: []const u8, n: usize, k: usize) !void { + var name: [272]u8 = undefined; + if (!std.mem.startsWith(u8, suffix, "feed_forward.")) return dense(a, out, try std.fmt.bufPrint(&name, "{s}{s}.weight", .{ prefix, suffix }), &.{ n, k }, .bf16); + if (k % 16 != 0) return error.UnsupportedQuantization; + try dense(a, out, try std.fmt.bufPrint(&name, "{s}{s}.weight", .{ prefix, suffix }), &.{ n, k / 2 }, .u8); + try dense(a, out, try std.fmt.bufPrint(&name, "{s}{s}.weight_scale", .{ prefix, suffix }), &.{ n, k / 16 }, .f8_e4m3); + try dense(a, out, try std.fmt.bufPrint(&name, "{s}{s}.weight_scale_2", .{ prefix, suffix }), &.{}, .f32); + try dense(a, out, try std.fmt.bufPrint(&name, "{s}{s}.input_scale", .{ prefix, suffix }), &.{}, .f32); +} + +/// Block FP8: q/k/v/o and the FFN as e4m3 [n, k] with bf16 scales [n / 128, k / 128]; the rest bf16 [n, k]. +fn fp8Block(a: Allocator, out: *std.ArrayList(Expected), prefix: []const u8, suffix: []const u8, n: usize, k: usize) !void { + var name: [272]u8 = undefined; + const quantized = std.mem.startsWith(u8, suffix, "feed_forward.") or + (std.mem.startsWith(u8, suffix, "self_attn.") and !std.mem.eql(u8, suffix, "self_attn.gate_proj")); + if (!quantized) return dense(a, out, try std.fmt.bufPrint(&name, "{s}{s}.weight", .{ prefix, suffix }), &.{ n, k }, .bf16); + const b = c.fp8_block; + if (n % b != 0 or k % b != 0) return error.UnsupportedQuantization; + try dense(a, out, try std.fmt.bufPrint(&name, "{s}{s}.weight", .{ prefix, suffix }), &.{ n, k }, .f8_e4m3); + try dense(a, out, try std.fmt.bufPrint(&name, "{s}{s}.weight_scale", .{ prefix, suffix }), &.{ n / b, k / b }, .bf16); +} + +/// Expected metadata for one text decoder; allocate the returned names and list in a caller-owned arena. +pub fn schema(a: Allocator, cfg: *const c.Config, namespace: Namespace) ![]Expected { + var out: std.ArrayList(Expected) = .empty; + const prefix = namespace.prefix(); + try affine(a, &out, cfg, prefix, "embed_tokens", c.vocab, c.hidden); + try named(a, &out, prefix, "embedding_norm.weight", &.{c.hidden}); + for (cfg.layer_types, 0..) |kind, i| { + var buf: [128]u8 = undefined; + const p = try std.fmt.bufPrint(&buf, "{s}layers.{d}.", .{ prefix, i }); + for ([_][]const u8{ "operator_norm.weight", "ffn_norm.weight", "post_operator_norm.weight", "post_ffn_norm.weight" }) |n| + try named(a, &out, p, n, &.{c.hidden}); + try affine(a, &out, cfg, p, "feed_forward.w1", c.intermediate, c.hidden); + try affine(a, &out, cfg, p, "feed_forward.w3", c.intermediate, c.hidden); + try affine(a, &out, cfg, p, "feed_forward.w2", c.hidden, c.intermediate); + if (kind == .conv) { + try affine(a, &out, cfg, p, "conv.in_proj", 3 * c.hidden, c.hidden); + try affine(a, &out, cfg, p, "conv.out_proj", c.hidden, c.hidden); + // MLX keeps the depthwise kernel as [channels, taps, 1], torch as [channels, 1, taps]: same values + const taps: []const usize = if (cfg.scheme != .mlx_affine) &.{ c.hidden, 1, c.conv_taps } else &.{ c.hidden, c.conv_taps, 1 }; + try named(a, &out, p, "conv.conv.weight", taps); + } else { + try affine(a, &out, cfg, p, "self_attn.q_proj", c.query_heads * c.head_dim, c.hidden); + try affine(a, &out, cfg, p, "self_attn.k_proj", c.kv_heads * c.head_dim, c.hidden); + try affine(a, &out, cfg, p, "self_attn.v_proj", c.kv_heads * c.head_dim, c.hidden); + try affine(a, &out, cfg, p, "self_attn.o_proj", c.hidden, c.query_heads * c.head_dim); + try affine(a, &out, cfg, p, "self_attn.gate_proj", c.query_heads, c.hidden); + try named(a, &out, p, "self_attn.q_norm.weight", &.{c.head_dim}); + try named(a, &out, p, "self_attn.k_norm.weight", &.{c.head_dim}); + } + } + return out.toOwnedSlice(a); +} + +/// Validate a merged safetensors header; the shared parser must have checked shard byte ranges first. +pub fn validate(a: Allocator, cfg: *const c.Config, header: *const st.Header) !Manifest { + const bare = header.contains("model.embed_tokens.weight"); + const wrapped = header.contains("language_model.model.embed_tokens.weight"); + if (bare and wrapped) return error.MixedNamespaces; + if (!bare and !wrapped) return error.MissingTensor; + const namespace: Namespace = if (bare) .bare else .wrapped; + var arena = std.heap.ArenaAllocator.init(a); + defer arena.deinit(); + const expected = try schema(arena.allocator(), cfg, namespace); + var names: std.StringHashMapUnmanaged(void) = .empty; + var bytes: usize = 0; + for (expected) |e| { + const actual = header.get(e.name) orelse return error.MissingTensor; + if (actual.dtype != e.dtype) return error.TensorDType; + if (actual.rank != e.rank or !std.mem.eql(usize, actual.shape[0..actual.rank], e.shape[0..e.rank])) return error.TensorShape; + var size: usize = e.dtype.size(); + for (e.shape[0..e.rank]) |dim| size = std.math.mul(usize, size, dim) catch return error.TensorBytes; + if (actual.end < actual.begin or actual.end - actual.begin != size) return error.TensorBytes; + bytes = std.math.add(usize, bytes, size) catch return error.TensorBytes; + try names.put(arena.allocator(), e.name, {}); + } + for (header.keys()) |name| { + if (std.mem.startsWith(u8, name, "model.") or std.mem.startsWith(u8, name, "language_model.")) { + if (!names.contains(name)) return error.UnexpectedTensor; + } + } + return .{ .namespace = namespace, .tensors = expected.len, .bytes = bytes }; +} + +test { + _ = @import("weights_test.zig"); +} diff --git a/zig/src/families/lfm3/weights_test.zig b/zig/src/families/lfm3/weights_test.zig new file mode 100644 index 0000000000..e2c8fa486b --- /dev/null +++ b/zig/src/families/lfm3/weights_test.zig @@ -0,0 +1,138 @@ +//! Metadata fixtures exercise admission without allocating tensor payloads or a GPU device. +const std = @import("std"); +const w = @import("weights.zig"); +const c = @import("config.zig"); +const st = @import("safetensors"); + +fn config(bits: u8, group: u16) c.Config { + return .{ .arena = std.heap.ArenaAllocator.init(std.testing.allocator), .kind = .text, .context = 262144, .quant = .{ .global = .{ .bits = bits, .group = group } } }; +} + +fn header(a: std.mem.Allocator, cfg: *const c.Config, namespace: w.Namespace) !st.Header { + var h: st.Header = .empty; + var offset: usize = 0; + for (try w.schema(a, cfg, namespace)) |e| { + var bytes: usize = e.dtype.size(); + for (e.shape[0..e.rank]) |d| bytes *= d; + try h.put(a, e.name, .{ .dtype = e.dtype, .rank = e.rank, .shape = e.shape, .begin = offset, .end = offset + bytes }); + offset += bytes; + } + return h; +} + +test "official text4 tensor count and bytes; both namespaces and every target bit width" { + for ([_]u8{ 4, 6, 8 }) |bits| for ([_]w.Namespace{ .bare, .wrapped }) |namespace| { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + var cfg = config(bits, 64); + defer cfg.deinit(); + const h = try header(arena.allocator(), &cfg, namespace); + const m = try w.validate(std.testing.allocator, &cfg, &h); + try std.testing.expectEqual(@as(usize, 594), m.tensors); + try std.testing.expectEqual(namespace, m.namespace); + if (bits == 4) try std.testing.expectEqual(@as(usize, 973008384), m.bytes); + }; +} + +test "VL text uses mixed8 embedding with group32 without admitting its image tower" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var cfg = config(4, 32); + defer cfg.deinit(); + cfg.kind = .vision_language; + try cfg.quant.modules.put(cfg.arena.allocator(), "embed_tokens", .{ .bits = 8, .group = 32 }); + var h = try header(a, &cfg, .wrapped); + const embed = h.get("language_model.model.embed_tokens.weight").?; + try std.testing.expectEqual(@as(usize, 512), embed.shape[1]); + const gate = h.get("language_model.model.layers.2.self_attn.gate_proj.scales").?; + try std.testing.expectEqualSlices(usize, &.{ 16, 64 }, gate.shape[0..2]); + try h.put(a, "vision_tower.unqualified.weight", .{ .dtype = .bf16, .rank = 1, .shape = .{ 1, 1, 1, 1, 1 }, .begin = 0, .end = 2 }); + const m = try w.validate(std.testing.allocator, &cfg, &h); + try std.testing.expectEqual(@as(usize, 594), m.tensors); + try std.testing.expectEqual(@as(usize, 1212106240), m.bytes); +} + +test "missing head gate, wrong packed width and malformed convolution are rejected" { + const keys = [_][]const u8{ + "language_model.model.layers.2.self_attn.gate_proj.weight", + "language_model.model.layers.0.conv.in_proj.weight", + "language_model.model.layers.0.conv.conv.weight", + "language_model.model.layers.0.post_operator_norm.weight", + }; + for (keys, 0..) |key, i| { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + var cfg = config(4, 64); + defer cfg.deinit(); + var h = try header(arena.allocator(), &cfg, .wrapped); + if (i == 0) { + _ = h.swapRemove(key); + try std.testing.expectError(error.MissingTensor, w.validate(std.testing.allocator, &cfg, &h)); + } else { + h.getPtr(key).?.shape[if (i == 3) 0 else 1] += 1; + try std.testing.expectError(error.TensorShape, w.validate(std.testing.allocator, &cfg, &h)); + } + } +} + +test "bad dtype, byte range, namespace and unexpected text tensors are rejected" { + for (0..4) |mode| { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + var cfg = config(4, 64); + defer cfg.deinit(); + var h = try header(arena.allocator(), &cfg, .wrapped); + const key = "language_model.model.embed_tokens.weight"; + switch (mode) { + 0 => h.getPtr(key).?.dtype = .bf16, + 1 => h.getPtr(key).?.end -= 1, + 2 => try h.put(arena.allocator(), "model.embed_tokens.weight", h.get(key).?), + 3 => try h.put(arena.allocator(), "language_model.lm_head.weight", h.get(key).?), + else => unreachable, + } + const err = switch (mode) { + 0 => error.TensorDType, + 1 => error.TensorBytes, + 2 => error.MixedNamespaces, + 3 => error.UnexpectedTensor, + else => unreachable, + }; + try std.testing.expectError(err, w.validate(std.testing.allocator, &cfg, &h)); + } +} + +test "the ModelOpt NVFP4 export: 502 bare tensors, FP4 FFN with its scales, everything else bf16" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + var cfg = try c.parse(std.testing.allocator, @import("config_test.zig").nvfp4); + defer cfg.deinit(); + const h = try header(arena.allocator(), &cfg, .bare); + const m = try w.validate(std.testing.allocator, &cfg, &h); + try std.testing.expectEqual(@as(usize, 502), m.tensors); + try std.testing.expectEqual(@as(usize, 1829689408), m.bytes); + const codes = h.get("model.layers.0.feed_forward.w2.weight").?; + try std.testing.expect(codes.dtype == .u8 and codes.shape[0] == 2048 and codes.shape[1] == 3840); + try std.testing.expectEqual(@as(u8, 0), h.get("model.layers.0.feed_forward.w2.weight_scale_2").?.rank); + try std.testing.expect(h.get("model.layers.2.self_attn.gate_proj.weight").?.dtype == .bf16); + try std.testing.expect(h.get("model.embed_tokens.weight").?.dtype == .bf16); + try std.testing.expect(h.get("model.layers.0.conv.in_proj.scales") == null); + try std.testing.expectEqualSlices(usize, &.{ 2048, 1, 7 }, h.get("model.layers.0.conv.conv.weight").?.shape[0..3]); +} + +test "the block FP8 export: 402 bare tensors, FP8 FFN and q/k/v/o with block scales, convs and gates bf16" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + var cfg = try c.parse(std.testing.allocator, @import("config_fp8_test.zig").fp8); + defer cfg.deinit(); + const h = try header(arena.allocator(), &cfg, .bare); + const m = try w.validate(std.testing.allocator, &cfg, &h); + try std.testing.expectEqual(@as(usize, 402), m.tensors); + try std.testing.expectEqual(@as(usize, 2209949952), m.bytes); + const down = h.get("model.layers.0.feed_forward.w2.weight_scale").?; + try std.testing.expect(down.dtype == .bf16 and down.shape[0] == 16 and down.shape[1] == 60); + try std.testing.expect(h.get("model.layers.2.self_attn.k_proj.weight").?.dtype == .f8_e4m3); + try std.testing.expect(h.get("model.layers.2.self_attn.gate_proj.weight").?.dtype == .bf16); + try std.testing.expect(h.get("model.layers.0.conv.in_proj.weight").?.dtype == .bf16); + try std.testing.expectEqualSlices(usize, &.{ 2048, 1, 7 }, h.get("model.layers.0.conv.conv.weight").?.shape[0..3]); +} diff --git a/zig/src/families/liquid/cuda.zig b/zig/src/families/liquid/cuda.zig new file mode 100644 index 0000000000..f7e7344b19 --- /dev/null +++ b/zig/src/families/liquid/cuda.zig @@ -0,0 +1,10 @@ +//! What the Liquid families (LFM2.5-24B-A2B, LFM3) share on CUDA: the lane backend, the kept states, the native glue. +pub const lanes = @import("cuda_lanes.zig"); +pub const snap = @import("cuda_snap.zig"); +pub const native = @import("cuda_native.zig"); + +test { + _ = lanes; + _ = snap; + _ = native; +} diff --git a/zig/src/families/liquid/cuda_lanes.zig b/zig/src/families/liquid/cuda_lanes.zig new file mode 100644 index 0000000000..23b03c9755 --- /dev/null +++ b/zig/src/families/liquid/cuda_lanes.zig @@ -0,0 +1,521 @@ +//! The Liquid families' lane backend on CUDA (LFM2.5, LFM3), generic over the family's model and forward. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const kv = @import("core").kv_pool; + +const be = lanes.backend; +const Stream = lanes.Stream; +const Pick = cuda.tree_round.Pick; +const Pages = @import("cuda_snap.zig").Pages; + +/// Positions a page of keys and values holds, in every Liquid family. +pub const page = 512; +/// The widest window a stream verifies: one-row steps and every wider window give a row the same bits. +const window_rows = 16; + +/// Pages a request holds: its prompt, its reply and one last verify window. +pub fn pagesOf(prompt: usize, max_new: usize) usize { + return (prompt + max_new + window_rows + page - 1) / page; +} + +/// Most marks a chunk taps, in any family. +pub const tap_cap = 4; + +/// A chunk's end (the grid line, the prompt's end or a mark) and the marks inside it it taps, in order. +pub const Cut = struct { + end: usize, + at: [tap_cap]usize = undefined, + n: usize = 0, + pub fn taps(c: *const Cut) []const usize { + return c.at[0..c.n]; + } +}; + +/// The chunk from `at` on a `rows` grid: the first `most` marks inside it are tapped, the next one ends it. +pub fn cut(at: usize, len: usize, marks: []const u32, rows: usize, most: usize) Cut { + var c: Cut = .{ .end = @min(len, (at / rows + 1) * rows) }; + var from = at; + while (true) { + var next: ?usize = null; + for (marks) |mk| if (mk > from and mk < c.end and (next == null or mk < next.?)) { + next = mk; + }; + const mk = next orelse break; + if (c.n == most) { + c.end = mk; + break; + } + c.at[c.n] = mk; + c.n += 1; + from = mk; + } + return c; +} + +/// `F` names the family's `cfg`, `mdl`, `kern`, `fwd` (Chain, forward, commit), `snap` and `text_only` (images' error). +pub fn Backend(comptime F: type) type { + return struct { + const cfg = F.cfg; + const fwd = F.fwd; + const mdl = F.mdl; + const kern = F.kern; + const snap = F.snap; + comptime { + std.debug.assert(kern.page == page); + } + /// Marks a chunk taps instead of ending there (`mdl.max_taps`; LFM3 has none, sliding rings would lose them). + const max_taps = if (@hasDecl(mdl, "max_taps")) mdl.max_taps else 0; + comptime { + std.debug.assert(max_taps <= tap_cap); + } + + /// Prompt rows a step takes: the forward's chunk grid, from position 0. + pub const prompt_rows = mdl.prompt_rows; + /// Window rows a round takes over every stream. + pub const round_rows = mdl.round_rows; + /// The widest window a stream verifies: one-row steps and every wider window give a row the same bits. + pub const max_window = window_rows; + /// Drawn tokens a handle names (a ring; read before the round loop could wrap it). + const ring = 1024; + + const Lane = struct { + slot: usize, + len: u64 = 0, // the cache length: positions [0, len) hold this stream's keys + base: u64 = 0, // the cache length the last verify started from + row0: usize = 0, // the stream's first row in the last verify's forward (its conv snapshots) + rows: std.ArrayList(u32) = .empty, // the last verify's row tokens + last: u64, // the prompt's last row's bf16 logits [V] (the first draw reads them) + rule: u64, // sample.cu's Rule, when the stream samples + sampled: bool = false, + fill: ?usize = null, // a prompt pass going a chunk at a time: where its next chunk starts + lease: kv.Lease = .{}, // a paged slot's pages + }; + + /// Pages a request holds: its prompt, its reply and one last verify window. + pub const pagesFor = pagesOf; + + pub const Liquid = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + m: *mdl.Model, + k: *const kern.Kernels, + s: cuda.Stream, + lanes: std.AutoHashMapUnmanaged(*const Stream, Lane) = .empty, + free: []bool, // slot i is free + tok: cuda.DeviceBuffer, // drawn tokens (u32): the handle ring, then a round's rows + picks: cuda.DeviceBuffer, // argmax picks (Pick), the same entries + greedy: []bool, // entry e of the ring (and the round's rows) is a pick, not a draw + meta: cuda.DeviceBuffer, // a draw's first position minus one, a window each + slot_bufs: []cuda.DeviceBuffer, // a slot's last-row logits, then its rule: two a slot + next: u64 = 0, + pool: ?kv.Pool = null, // the pages every slot leases (attach) + snaps: ?Pages = null, // the prompt cache's pages, allocated once inside its budget (snap_ready) + tapped: ?struct { stream: *const Stream, cut: Cut } = null, // the chunk's tapped marks while its hooks run + + /// The device bytes `init` allocates: the token and pick rings, metas, each slot's last row and rule. + pub fn openBytes(slots: usize) usize { + return (ring + round_rows) * (4 + @sizeOf(Pick)) + round_rows * 16 + slots * (cfg.vocab * 2 + @sizeOf(cuda.sample.Rule)); + } + + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, m: *mdl.Model, k: *const kern.Kernels, s: cuda.Stream) !Liquid { + const free = try gpa.alloc(bool, m.slotCap()); + errdefer gpa.free(free); + @memset(free, true); + const greedy = try gpa.alloc(bool, ring + round_rows); + errdefer gpa.free(greedy); + @memset(greedy, true); + var tok = try cuda.DeviceBuffer.alloc(d, (ring + round_rows) * 4); + errdefer tok.free(); + var picks = try cuda.DeviceBuffer.alloc(d, (ring + round_rows) * @sizeOf(Pick)); + errdefer picks.free(); + var meta = try cuda.DeviceBuffer.alloc(d, round_rows * 16); + errdefer meta.free(); + const bufs = try gpa.alloc(cuda.DeviceBuffer, 2 * m.slotCap()); + errdefer gpa.free(bufs); + var made: usize = 0; + errdefer for (bufs[0..made]) |*b| b.free(); + for (bufs, 0..) |*b, i| { + b.* = try cuda.DeviceBuffer.alloc(d, if (i % 2 == 0) cfg.vocab * 2 else @sizeOf(cuda.sample.Rule)); + made += 1; + } + return .{ .gpa = gpa, .d = d, .m = m, .k = k, .s = s, .free = free, .tok = tok, .picks = picks, .greedy = greedy, .meta = meta, .slot_bufs = bufs }; + } + + /// Makes a paged model's arena of `pages` pages for `extra + 1` slots, and the pool that leases them. + pub fn attach(x: *Liquid, extra: usize, pages: usize) !void { + try x.m.attach(x.s, extra, pages); + x.pool = try kv.Pool.init(x.gpa, pages, mdl.page_elems * 2); + } + + pub fn deinit(x: *Liquid) void { + var it = x.lanes.valueIterator(); + while (it.next()) |l| { + if (x.pool) |*p| p.finish(&l.lease); + l.rows.deinit(x.gpa); + } + if (x.pool) |*p| p.deinit(); + if (x.snaps) |*p| p.deinit(); + x.lanes.deinit(x.gpa); + for (x.slot_bufs) |*b| b.free(); + x.gpa.free(x.slot_bufs); + x.gpa.free(x.free); + x.gpa.free(x.greedy); + x.tok.free(); + x.picks.free(); + x.meta.free(); + } + + pub fn backend(x: *Liquid) be.Backend { + return .{ .ptr = x, .vtable = &vtable }; + } + + const vtable: be.Backend.VTable = .{ .prefill = prefill, .prefill_step = prefillStep, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .release = release, .reports_fill = true }; + + /// The round loop's facts: no draft head; host copies verified in windows up to 16 rows, one forward. + pub fn facts(x: *const Liquid) lanes.Model { + return .{ .exact_width = max_window, .gpu_tokens = false, .mtp = false, .speculate = false, .hidden_rows = true, .batch_rows = round_rows, .max_streams = @intCast(@min(x.m.slotCap(), round_rows)) }; + } + + fn self(ptr: *anyopaque) *Liquid { + return @ptrCast(@alignCast(ptr)); + } + + fn upload(x: *Liquid, dst: u64, bytes: []const u8) !void { + try x.d.check(x.d.api.cuMemcpyHtoDAsync_v2(dst, bytes.ptr, bytes.len, x.s.handle), "cuMemcpyHtoDAsync"); + } + + /// Rows [row0, row0 + rows) of bf16 `logits` into entries [at, at + rows), row r at `position0 + r`. + fn drawRows(x: *Liquid, l: *const Lane, logits: u64, row0: usize, rows: usize, position0: u64, at: usize, meta_at: usize) !void { + const v = cfg.vocab; + const src = logits + row0 * v * 2; + @memset(x.greedy[at..][0..rows], !l.sampled); + if (!l.sampled) return x.k.round.argmax(x.s, src, x.picks.ptr + at * @sizeOf(Pick), rows, v, v); + const meta = [4]i32{ @intCast(position0 - 1), 0, 0, 0 }; + const mp = x.meta.ptr + meta_at * 16; + try x.upload(mp, std.mem.asBytes(&meta)); + try x.k.draws.draw(x.s, src, v, l.rule, mp, 0, x.tok.ptr + at * 4, rows, null, null); + } + + /// Entries [at, at + out.len) as tokens, after the stream's work; a non-finite pick is an error. + fn tokens(x: *Liquid, at: usize, out: []u32) !void { + try x.s.synchronize(); + var picks: [round_rows]Pick = undefined; + try x.tok.download(at * 4, std.mem.sliceAsBytes(out)); + try x.picks.download(at * @sizeOf(Pick), std.mem.sliceAsBytes(picks[0..out.len])); + for (out, picks[0..out.len], x.greedy[at..][0..out.len]) |*t, p, g| if (g) { + if (p.nonfinite != 0) return error.NonFiniteLogits; + t.* = p.token; + }; + } + + /// The stream's state at `at` for the prompt cache, copied behind the chunk that ends there. + pub fn saveState(x: *Liquid, owner: ?*anyopaque, at: u32) !*snap.Saved { + const s: *const Stream = @ptrCast(@alignCast(owner orelse return error.NoOwner)); + const l = x.lanes.getPtr(s) orelse return error.UnknownStream; + const j: ?usize = if (x.tapped) |t| if (t.stream == s) std.mem.indexOfScalar(usize, t.cut.taps(), at) else null else null; + if (j == null and (l.fill orelse l.len) != at) return error.NotAtMark; // where the pass stands, or a tap + const tap: ?u64 = if (comptime max_taps > 0) (if (j) |i| x.m.s.tap + i * mdl.slot_elems * 2 else null) else null; + return snap.save(x.gpa, try x.snapPages(), x.m, x.s, l.slot, at, tap); + } + + /// The stream's slot made to hold `k`'s state (a restore outside its prompt pass). + pub fn restoreState(x: *Liquid, owner: ?*anyopaque, k: *const snap.Saved) !void { + const s: *const Stream = @ptrCast(@alignCast(owner orelse return error.NoOwner)); + const l = x.lanes.getPtr(s) orelse return error.UnknownStream; + try snap.restore(try x.snapPages(), x.m, x.s, l.slot, k); + l.len = k.at; + } + + pub fn dropState(x: *Liquid, k: *snap.Saved) void { + snap.drop(x.gpa, &x.snaps.?, k); + } + + fn snapPages(x: *Liquid) !*Pages { + return if (x.snaps) |*p| p else error.NoSnapshotCache; + } + + fn freeSlot(x: *const Liquid) ?usize { + return std.mem.indexOfScalar(bool, x.free[0..x.m.slots], true); + } + + fn fitsContext(x: *const Liquid, prompt: usize, max_new: usize) bool { + return prompt + max_new + max_window <= x.m.context; + } + + /// Whether a request could start now: a free slot and (paged) its pages; one that never fits passes. + pub fn admits(x: *const Liquid, prompt: usize, max_new: usize) bool { + if (!x.fitsContext(prompt, max_new)) return true; // its prompt pass refuses it by name + const p = x.pool orelse return x.freeSlot() != null; + return x.freeSlot() != null and p.fits(pagesFor(prompt, max_new)); + } + + const Map = struct { + x: *Liquid, + slot: usize, + fn write(ctx: *anyopaque, ids: []const u32, from: usize) anyerror!void { + const m: *const Map = @ptrCast(@alignCast(ctx)); + try m.x.m.mapPages(m.x.s, m.slot, ids[from..], from); + } + }; + + /// A free slot for the request, its pages leased now when paged; LaneMemory: not yet. + fn takeSlot(x: *Liquid, lease: *kv.Lease, prompt: usize, max_new: usize) !usize { + const slot = x.freeSlot() orelse return error.NoFreeSlot; + const p = if (x.pool) |*p| p else return slot; + const n = pagesFor(prompt, max_new); + try p.promise(lease, n); + errdefer p.finish(lease); + var map: Map = .{ .x = x, .slot = slot }; + try p.cover(lease, n, .{ .ctx = &map, .write = Map.write }); + return slot; + } + + /// A prompt pass begins: a slot, the stream's rule, a kept state restored; where its first chunk starts. + fn opening(x: *Liquid, s: *Stream) !*Lane { + const p = s.prompt(); + if (p.len == 0) return error.EmptyPrompt; + if (s.images.len > 0) return F.text_only; + if (!x.fitsContext(p.len, s.max_new)) return error.PastContext; + const got = try x.lanes.getOrPut(x.gpa, s); + if (!got.found_existing) { + var lease: kv.Lease = .{}; + const slot = x.takeSlot(&lease, p.len, s.max_new) catch |e| { + x.lanes.removeByPtr(got.key_ptr); + return e; + }; + x.free[slot] = false; + got.value_ptr.* = .{ .slot = slot, .last = x.slot_bufs[2 * slot].ptr, .rule = x.slot_bufs[2 * slot + 1].ptr, .lease = lease }; + } + const l = got.value_ptr; + l.sampled = false; + if (s.sampling) |sp| if (sp.temperature > 0) { + const r = cuda.sample.Rule.from(lanes.gpu_rule.rule(sp)); + try x.upload(l.rule, std.mem.asBytes(&r)); + l.sampled = true; + }; + s.cached = 0; + try x.m.clearSlot(x.s, l.slot); + l.fill = x.resumeAt(s, l, p.len); + s.fill_at.store(@intCast(l.fill.?), .monotonic); + return l; + } + + /// The kept state's position, restored into the slot (0: none, or not this prompt's). + fn resumeAt(x: *Liquid, s: *Stream, l: *const Lane, len: usize) usize { + const saved = s.reuse.saved orelse return 0; + const k: *const snap.Saved = @ptrCast(@alignCast(saved)); + if (k.at != s.reuse.at or k.at == 0 or k.at >= len) { + s.reuse_failed = true; + return 0; + } + const p = x.snapPages() catch { + s.reuse_failed = true; + return 0; + }; + snap.restore(p, x.m, x.s, l.slot, k) catch { + s.reuse_failed = true; + return 0; + }; + s.cached = k.at; + return k.at; + } + + /// One chunk from `l.fill` to the next grid line or mark; true once the prompt's last logits are written. + fn chunk(x: *Liquid, s: *Stream, l: *Lane) !bool { + const p = s.prompt(); + const at = l.fill.?; + if (s.isCancelled()) return error.Cancelled; + const c = cut(at, p.len, s.reuse.marks, prompt_rows, max_taps); + const end = c.end; + const done = end == p.len; + var rows: [tap_cap]usize = undefined; + for (c.taps(), 0..) |t, j| rows[j] = t - at - 1; + const ch: fwd.Chain = if (comptime max_taps > 0) .{ .slot = l.slot, .p0 = at, .tokens = p[at..end], .taps = rows[0..c.n] } else .{ .slot = l.slot, .p0 = at, .tokens = p[at..end] }; + try fwd.forward(x.m, x.k, x.s, &.{ch}, true, end - at - 1, @intFromBool(done)); + if (done) try x.d.check(x.d.api.cuMemcpyDtoDAsync_v2(l.last, x.m.s.logits, cfg.vocab * 2, x.s.handle), "cuMemcpyDtoDAsync"); + s.fill_at.store(@intCast(end), .monotonic); + l.fill = if (done) null else end; + if (done) l.len = p.len; + if (s.reuse.hook) |hk| { + x.tapped = .{ .stream = s, .cut = c }; + defer x.tapped = null; + for (c.taps()) |t| hk.at(hk.ptr, s, @intCast(t)); + if (std.mem.indexOfScalar(u32, s.reuse.marks, @intCast(end)) != null) hk.at(hk.ptr, s, @intCast(end)); + } + return done; + } + + fn prefill(ptr: *anyopaque, s: *Stream) anyerror!void { + const x = self(ptr); + const l = try x.opening(s); + errdefer x.dropLane(s); + while (!try x.chunk(s, l)) {} + } + + fn prefillStep(ptr: *anyopaque, s: *Stream) anyerror!bool { + const x = self(ptr); + const going = if (x.lanes.getPtr(s)) |l| l.fill != null else false; + const l = if (going) x.lanes.getPtr(s).? else try x.opening(s); + errdefer x.dropLane(s); + return x.chunk(s, l); + } + + fn handle(x: *Liquid) usize { + const h = x.next; + x.next += 1; + return @intCast(h % ring); + } + + fn first(ptr: *anyopaque, s: *Stream, position: u64) anyerror!u64 { + const x = self(ptr); + const l = x.lanes.getPtr(s) orelse return error.UnknownStream; + if (position != l.len or l.fill != null) return error.PositionMismatch; + const h = x.next; + try x.drawRows(l, l.last, 0, 1, position, x.handle(), 0); + return h; + } + + fn value(x: *Liquid, feed: be.Feed) !u32 { + return switch (feed) { + .value => |v| v, + .handle => |h| read(x, h), + }; + } + + fn queue(ptr: *anyopaque, s: *Stream, feed: be.Feed, position: u64) anyerror!u64 { + const x = self(ptr); + const l = x.lanes.getPtr(s) orelse return error.UnknownStream; + const t = try x.value(feed); + if (position != l.len + 1) return error.PositionMismatch; + try fwd.forward(x.m, x.k, x.s, &.{.{ .slot = l.slot, .p0 = l.len, .tokens = &.{t} }}, false, 0, 1); + l.len += 1; + const h = x.next; + try x.drawRows(l, x.m.s.logits, 0, 1, position, x.handle(), 0); + return h; + } + + fn read(ptr: *anyopaque, h: u64) anyerror!u32 { + const x = self(ptr); + if (h >= x.next or x.next - h > ring) return error.NoSuchToken; + var out: [1]u32 = undefined; + try x.tokens(@intCast(h % ring), &out); + return out[0]; + } + + fn verify(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { + const x = self(ptr); + const chains = try x.gpa.alloc(fwd.Chain, windows.len); + defer x.gpa.free(chains); + var total: usize = 0; + for (windows, chains) |w, *c| { + if (w.held != 0) return error.NoHeldDrafts; // no draft head: every draft comes from the host + if (w.parents != null) return error.ChainsOnly; + if (w.masks.len != 0) return error.StructuresUnsupported; + const l = x.lanes.getPtr(w.stream) orelse return error.UnknownStream; + if (w.rows() > max_window) return error.WindowTooWide; + l.base = l.len; + l.rows.clearRetainingCapacity(); + try l.rows.append(x.gpa, w.pending); + try l.rows.appendSlice(x.gpa, w.tokens); + for (w.positions, 0..) |pos, r| if (pos != l.base + r + 1) return error.PositionMismatch; + c.* = .{ .slot = l.slot, .p0 = l.base, .tokens = l.rows.items }; + l.row0 = total; + total += l.rows.items.len; + } + if (total > round_rows) return error.RoundTooWide; + try fwd.forward(x.m, x.k, x.s, chains, false, 0, total); + var at: usize = 0; + for (windows, 0..) |w, i| { + const l = x.lanes.getPtr(w.stream).?; + try x.drawRows(l, x.m.s.logits, at, l.rows.items.len, w.positions[0], ring + at, i); + at += l.rows.items.len; + } + var got: [round_rows]u32 = undefined; + try x.tokens(ring, got[0..total]); + at = 0; + for (windows, out) |w, o| { + const l = x.lanes.getPtr(w.stream).?; + const n = l.rows.items.len; + @memcpy(o.sampled[0..n], got[at..][0..n]); + @memcpy(o.drafts[0 .. n - 1], l.rows.items[1..]); + l.len = l.base + n; // a chain keeps every row unless rolled back + at += n; + } + } + + fn keep(ptr: *anyopaque, windows: []const be.Window, paths: []const []const u32) anyerror!void { + const x = self(ptr); + if (windows.len > round_rows) return error.RoundTooWide; + var slots: [round_rows]i32 = undefined; + var rows: [round_rows]i32 = undefined; + for (windows, paths, 0..) |w, path, i| { + const l = x.lanes.getPtr(w.stream) orelse return error.UnknownStream; + if (path.len == 0 or path.len > l.rows.items.len) return error.EmptyPath; + for (path, 0..) |r, j| if (r != j) return error.ChainsOnly; // a chain keeps a prefix of its rows + l.len = l.base + path.len; // caches are positional: later rows overwrite what the rejected ones wrote + slots[i] = @intCast(l.slot); + rows[i] = @intCast(l.row0 + path.len - 1); // the conv states after the last kept row + } + try fwd.commit(x.m, x.k, x.s, slots[0..windows.len], rows[0..windows.len]); + } + + /// No draft head: absorbing is a no-op, drafts are refused by name (copy drafts come from the host). + fn draft(_: *anyopaque, requests: []const be.DraftRequest) anyerror!void { + for (requests) |r| if (r.depth != 0 or r.lanes != null) return error.NoDraftHead; + } + + fn dropLane(x: *Liquid, s: *const Stream) void { + if (x.lanes.fetchRemove(s)) |entry| { + var l = entry.value; + if (x.pool) |*p| { + p.finish(&l.lease); + x.m.unmapPages(l.slot); + } + x.free[l.slot] = true; + l.rows.deinit(x.gpa); + if (x.lanes.count() == 0) x.report(); + } + } + + /// The idle server's line: nothing promised, every page free (kept states are copies, not pages). + fn report(x: *const Liquid) void { + if (@import("builtin").is_test) return; + const p = x.pool orelse return; + const c = p.census(); + const ok = p.idle() and c.free == c.total; + std.log.info("kv pool idle{s}: promised {d}, free {d} + store {d} == total {d}", .{ if (ok) "" else " (UNEVEN)", c.promised, c.free, c.kept, c.total }); + } + + fn release(ptr: *anyopaque, s: *Stream) void { + self(ptr).dropLane(s); + } + }; + }; +} + +test "a request holds its prompt, its reply and one verify window in whole pages" { + try std.testing.expectEqual(@as(usize, 1), pagesOf(1, 1)); + try std.testing.expectEqual(@as(usize, 1), pagesOf(400, 96)); + try std.testing.expectEqual(@as(usize, 2), pagesOf(400, 97)); + try std.testing.expectEqual(@as(usize, 65), pagesOf(32768, 0)); +} + +test "a chunk taps the marks inside it, up to its family's most, and ends at the next one, the grid line or the end" { + const cases = [_]struct { at: usize, len: usize, marks: []const u32, most: usize = 4, end: usize, taps: []const usize }{ + .{ .at = 0, .len = 3000, .marks = &.{}, .end = 1024, .taps = &.{} }, + .{ .at = 0, .len = 3000, .marks = &.{700}, .end = 1024, .taps = &.{700} }, + .{ .at = 0, .len = 3000, .marks = &.{700}, .most = 0, .end = 700, .taps = &.{} }, + .{ .at = 0, .len = 3000, .marks = &.{ 900, 300, 2900 }, .end = 1024, .taps = &.{ 300, 900 } }, + .{ .at = 0, .len = 3000, .marks = &.{ 900, 300, 2900 }, .most = 0, .end = 300, .taps = &.{} }, + .{ .at = 0, .len = 3000, .marks = &.{ 5, 4, 3, 2, 1 }, .end = 5, .taps = &.{ 1, 2, 3, 4 } }, + .{ .at = 9216, .len = 10144, .marks = &.{ 10050, 10140 }, .end = 10144, .taps = &.{ 10050, 10140 } }, + .{ .at = 2048, .len = 2500, .marks = &.{ 2100, 2500 }, .end = 2500, .taps = &.{2100} }, + .{ .at = 700, .len = 3000, .marks = &.{700}, .end = 1024, .taps = &.{} }, + }; + for (cases) |k| { + const c = cut(k.at, k.len, k.marks, 1024, k.most); + try std.testing.expectEqual(k.end, c.end); + try std.testing.expectEqualSlices(usize, k.taps, c.taps()); + } +} diff --git a/zig/src/families/liquid/cuda_native.zig b/zig/src/families/liquid/cuda_native.zig new file mode 100644 index 0000000000..e2b9552071 --- /dev/null +++ b/zig/src/families/liquid/cuda_native.zig @@ -0,0 +1,191 @@ +//! The Liquid families behind the native CUDA registry: open, admission, the pool and the prompt cache's hooks. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const kv = @import("core").kv_pool; + +/// `F` names the family's `wts`, `kern`, `mdl`, `be`, `snap`, `pages_env` and `check` (beyond the admission). +pub fn Native(comptime F: type) type { + return struct { + const wts = F.wts; + const kern = F.kern; + const mdl = F.mdl; + const be = F.be; + const snap = F.snap; + const Pages = @import("cuda_snap.zig").Pages; + + pub const prompt_rows: u32 = be.prompt_rows; + /// A kept state is the caches through its position and nothing past it (no draft head reads a token on). + pub const snap_lookahead: u32 = 0; + /// A request's last verify window: its pages hold that many rows past its prompt and reply. + pub const window_rows: u32 = be.max_window; + + /// The server's reading of the room left after the weights; null if unreadable. + pub const Room = struct { ctx: *anyopaque, read: *const fn (*anyopaque) ?u64 }; + /// No `precision` field: the server refuses --precision for this family by name. + pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false, lanes: u32 = 8, lanes_fixed: bool = false, room: ?Room = null }; + + pub const Loaded = struct { + backend: lanes.backend.Backend, + facts: lanes.Model, + rows: u32, + stream_bytes: usize, + free_streams: u32, + ctx: *anyopaque, + deinit: *const fn (*anyopaque) void, + target: *be.Liquid, + /// The attention caches as pages of the core pool (null: every stream owns a window, made at open). + layout: ?kv.Layout = null, + /// With a layout: the first stream's rows are pages too, a window of them set aside by the plan. + engine_paged: bool = true, + }; + + const Owned = struct { gpa: std.mem.Allocator, w: wts.Weights, k: kern.Kernels, s: cuda.Stream, m: mdl.Model, lb: be.Liquid }; + + /// Streams to allocate in `room`: as asked, else the most that fit; a fixed count past them refuses. + fn fit(room: u64, asked: u32, fixed: bool, plan: anytype) !usize { + if (fixed and asked > be.round_rows) return error.ParallelPastRound; + var n = std.math.clamp(@as(usize, asked), 1, be.round_rows); + while (try plan.bytes(n) > room) : (n -= 1) { + if (fixed) return error.ParallelPastMemory; + if (n == 1) return error.NoStreamFits; + } + return n; + } + + /// The pages switch (TENSORFOLD_LFM*_KV_PAGES) at 0: each stream owns its pages, else it leases the pool's. + fn pagesOn() bool { + const v = std.c.getenv(F.pages_env) orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); + } + + /// The device bytes `n` streams take (paged: the first's states and a round's slot buffers, no pages). + const Plan = struct { + own: *const Owned, + d: *const cuda.Driver, + context: usize, + paged: bool, + + fn bytes(p: Plan, n: usize) !u64 { + const o = p.own; + return try mdl.Model.openBytes(o.gpa, p.d, o.s, &o.w, &o.k, p.context, n, p.paged) + be.Liquid.openBytes(if (p.paged) be.round_rows else n); + } + }; + + /// The pool: `streams` more slots beside the first and an arena of `n` pages. + pub fn attach(p: *anyopaque, streams: u32, n: usize) !void { + const own: *Owned = @ptrCast(@alignCast(p)); + try own.lb.attach(streams, n); + } + + /// Whether a request's caches could start now (the lane host's admission, on its engine thread). + pub fn fits(p: *anyopaque, prompt: u32, max_tokens: u32) bool { + const own: *const Owned = @ptrCast(@alignCast(p)); + return own.lb.admits(prompt, max_tokens); + } + + pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { + if (kernels != null) return error.CapturedKernels; + if (o.lanes_fixed and o.lanes > be.round_rows) return error.ParallelPastRound; + const own = try gpa.create(Owned); + errdefer gpa.destroy(own); + own.* = .{ .gpa = gpa, .w = undefined, .k = undefined, .s = undefined, .m = undefined, .lb = undefined }; + own.k = try kern.Kernels.load(ctx); + errdefer own.k.deinit(); + own.s = try cuda.Stream.init(ctx.d, false); + errdefer own.s.deinit(); + own.w = try wts.load(gpa, io, ctx.d, own.s, &own.k, dir, null); + errdefer own.w.deinit(); + try F.check(&own.w.config); + const context = @min(mdl.Model.positions(o.context), own.w.config.context); + const paged = pagesOn(); + const plan: Plan = .{ .own = own, .d = ctx.d, .context = context, .paged = paged }; + const room: u64 = if (o.room) |r| r.read(r.ctx) orelse 0 else std.math.maxInt(u64); + const n = if (paged) try fitFirst(room, plan) else fit(room, o.lanes, o.lanes_fixed, plan) catch |e| { + if (e == error.ParallelPastMemory) { + const most = fit(room, o.lanes, false, plan) catch 0; + std.log.err("--parallel {d} needs {d:.2} GiB at a {d}-token window and the CUDA memory budget leaves {d:.2} GiB after the weights: {d} stream{s} fit", .{ o.lanes, gib(try plan.bytes(o.lanes)), context, gib(room), most, if (most == 1) "" else "s" }); + } + return e; + }; + own.m = try mdl.Model.init(gpa, ctx.d, own.s, &own.w, &own.k, context, n, paged); + errdefer own.m.deinit(); + own.lb = try be.Liquid.init(gpa, ctx.d, &own.m, &own.k, own.s); + errdefer own.lb.deinit(); + return .{ + .backend = own.lb.backend(), + .facts = own.lb.facts(), + .rows = be.max_window, + .stream_bytes = mdl.Model.streamBytes(context), + .free_streams = @intCast(n), + .ctx = own, + .deinit = release, + .target = &own.lb, + .layout = if (paged) own.m.layout() else null, + }; + } + + /// Paged: the first stream's states and the scratch must fit; the host plans the pool after. + fn fitFirst(room: u64, plan: anytype) !usize { + if (try plan.bytes(1) > room) return error.NoStreamFits; + return 1; + } + + fn gib(bytes: u64) f64 { + return @as(f64, @floatFromInt(bytes)) / (1 << 30); + } + + fn lfm(ptr: *anyopaque) *be.Liquid { + return @ptrCast(@alignCast(ptr)); + } + + /// Bytes a state at `at` prompt tokens takes: whole snapshot pages. + pub fn snap_bytes(_: *anyopaque, at: u32) u64 { + return snap.bytes(at); + } + + /// The cache's pages, allocated once before any request; the store counts held and spare pages via the hooks. + pub fn snap_ready(ptr: *anyopaque, budget: u64) !u64 { + const x = lfm(ptr); + if (x.snaps != null) return error.SnapshotCacheMade; + x.snaps = try Pages.init(x.gpa, x.d, budget); + return 0; + } + + pub fn snap_charged(_: *anyopaque, saved: *anyopaque) u64 { + const k: *const snap.Saved = @ptrCast(@alignCast(saved)); + return k.lease.ids.items.len * @import("cuda_snap.zig").page; + } + + pub fn snap_spare(ptr: *anyopaque) u64 { + return if (lfm(ptr).snaps) |*p| p.spare() else 0; + } + + pub fn snap_reuses(ptr: *anyopaque, at: u32) bool { + return snap_spare(ptr) >= snap.bytes(at); + } + + /// The prompt pass stands at `at` (a chunk end or a mark): copy the stream's caches there. + pub fn snap_save(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!*anyopaque { + return @ptrCast(try lfm(ptr).saveState(owner, at)); + } + + pub fn snap_restore(ptr: *anyopaque, owner: ?*anyopaque, saved: *anyopaque) anyerror!void { + return lfm(ptr).restoreState(owner, @ptrCast(@alignCast(saved))); + } + + pub fn snap_drop(ptr: *anyopaque, saved: *anyopaque) void { + lfm(ptr).dropState(@ptrCast(@alignCast(saved))); + } + + fn release(p: *anyopaque) void { + const own: *Owned = @ptrCast(@alignCast(p)); + own.lb.deinit(); + own.m.deinit(); + own.w.deinit(); + own.s.deinit(); + own.k.deinit(); + own.gpa.destroy(own); + } + }; +} diff --git a/zig/src/families/liquid/cuda_snap.zig b/zig/src/families/liquid/cuda_snap.zig new file mode 100644 index 0000000000..20e235d609 --- /dev/null +++ b/zig/src/families/liquid/cuda_snap.zig @@ -0,0 +1,160 @@ +//! Liquid's kept states in pages allocated once inside the prompt-cache budget (as GLM's): keeps allocate nothing. +const std = @import("std"); +const cuda = @import("cuda"); +const kv = @import("core").kv_pool; +const rows_page = @import("cuda_lanes.zig").page; + +/// A snapshot page: the unit the budget is cut into. +pub const page: usize = 2 << 20; + +/// The budget's pages, allocated once; states lease them and give them back. +pub const Pages = struct { + pages: kv.Pool, + arena: ?cuda.DeviceBuffer, + + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, budget: u64) !Pages { + const count = std.math.cast(u32, budget / page) orelse return error.SnapshotBudget; + var pages = try kv.Pool.init(gpa, count, page); + errdefer pages.deinit(); + return .{ .pages = pages, .arena = if (count > 0) try cuda.DeviceBuffer.alloc(d, @as(usize, count) * page) else null }; + } + + pub fn deinit(p: *Pages) void { + if (p.arena) |*a| a.free(); + p.pages.deinit(); + } + + /// Bytes of pages no state holds. + pub fn spare(p: *const Pages) u64 { + return p.pages.available() * page; + } + + fn take(p: *Pages, n: usize) !kv.Lease { + if (n == 0 or n > p.pages.available() or p.arena == null) return error.NoSnapshotRoom; + var lease: kv.Lease = .{}; + errdefer p.pages.finish(&lease); + try p.pages.promise(&lease, n); + try p.pages.cover(&lease, n, null); + return lease; + } +}; + +/// Region bytes [off, off + len) over a state's pages: (arena offset, region offset, length), a run of consecutive ids. +fn pageRuns(ids: []const u32, off: usize, len: usize, ctx: anytype, comptime f: fn (@TypeOf(ctx), held: u64, done: usize, n: usize) anyerror!void) !void { + var done: usize = 0; + while (done < len) { + const at = off + done; + var q = at / page; + var n = @min(len - done, page - at % page); + while (done + n < len and ids[q + 1] == ids[q] + 1) : (q += 1) n = @min(len - done, n + page); + try f(ctx, @as(u64, ids[at / page]) * page + at % page, done, n); + done += n; + } +} + +/// Copies between a slot's region at `dev` and a lease's pages. +const Copy = struct { d: *const cuda.Driver, s: cuda.Stream, arena: u64, dev: u64, out: bool }; + +fn copyRun(c: Copy, held: u64, done: usize, n: usize) anyerror!void { + const src = if (c.out) c.dev + done else c.arena + held; + const dst = if (c.out) c.arena + held else c.dev + done; + try c.d.check(c.d.api.cuMemcpyDtoDAsync_v2(dst, src, n, c.s.handle), "cuMemcpyDtoDAsync"); +} + +const Leased = struct { d: *const cuda.Driver, s: cuda.Stream, arena: u64, ids: []const u32, out: bool }; + +fn copy(c: Leased, dev: u64, off: usize, len: usize) anyerror!void { + try pageRuns(c.ids, off, len, Copy{ .d = c.d, .s = c.s, .arena = c.arena, .dev = dev, .out = c.out }, copyRun); +} + +/// `mdl` names `page_elems`, `kv_width`, `paged_layers`, a Model (`fixed_bytes`, `fixed(slot, i)`, arena, ids, caps). +pub fn States(comptime mdl: type) type { + return struct { + /// A kept state: its position and the snapshot pages that hold it, in order. + pub const Saved = struct { at: u32, lease: kv.Lease }; + + const page_bytes = mdl.page_elems * 2; + const row_bytes = mdl.kv_width * 2; + + /// Bytes a state at `at` holds: whole snapshot pages. + pub fn bytes(at: usize) u64 { + const n = mdl.Model.fixed_bytes + at * (page_bytes / rows_page); + return (n + page - 1) / page * page; + } + + /// `slot`'s state at `at`: fixed regions (the first from `tap` if set), a copy a run of whole pages, then rows. + pub fn each(m: *const mdl.Model, slot: usize, at: usize, tap: ?u64, ctx: anytype, comptime f: fn (@TypeOf(ctx), dev: u64, off: usize, len: usize) anyerror!void) !void { + if (at > m.caps[slot]) return error.PastPages; + var off: usize = 0; + var i: usize = 0; + while (m.fixed(slot, i)) |r| : (i += 1) { // a ring holds positions, not rows: whole + try f(ctx, if (i == 0) tap orelse r[0] else r[0], off, @intCast(r[1])); + off += @intCast(r[1]); + } + const ids = m.ids[slot * m.perSlot() ..]; + const whole = at / rows_page; + var q: usize = 0; + while (q < whole) { + var n: usize = 1; + while (q + n < whole and ids[q + n] == ids[q] + n) n += 1; + try f(ctx, m.arena + @as(u64, ids[q]) * page_bytes, off, n * page_bytes); + off += n * page_bytes; + q += n; + } + const rows = at % rows_page; + if (rows == 0) return; + const last = m.arena + @as(u64, ids[whole]) * page_bytes; + for (0..2 * mdl.paged_layers) |l| { // layer a's keys at block 2a, its values at 2a + 1 + try f(ctx, last + l * rows_page * row_bytes, off, rows * row_bytes); + off += rows * row_bytes; + } + } + + /// `slot`'s state at `at` copied into leased pages, queued behind the chunk that ends there or taps it (`tap`). + pub fn save(gpa: std.mem.Allocator, p: *Pages, m: *const mdl.Model, s: cuda.Stream, slot: usize, at: u32, tap: ?u64) !*Saved { + if (at == 0 or at > m.context) return error.NotAtMark; + const k = try gpa.create(Saved); + errdefer gpa.destroy(k); + k.* = .{ .at = at, .lease = try p.take(bytes(at) / page) }; + errdefer p.pages.finish(&k.lease); + try each(m, slot, at, tap, Leased{ .d = m.d, .s = s, .arena = p.arena.?.ptr, .ids = k.lease.ids.items, .out = true }, copy); + return k; + } + + /// Make `slot`'s pages and fixed regions hold `k`'s state, queued on `s` before the next chunk. + pub fn restore(p: *const Pages, m: *const mdl.Model, s: cuda.Stream, slot: usize, k: *const Saved) !void { + if (k.at > m.context) return error.PastContext; + try each(m, slot, k.at, null, Leased{ .d = m.d, .s = s, .arena = p.arena.?.ptr, .ids = k.lease.ids.items, .out = false }, copy); + } + + /// Give a state's pages back: every copy of them is queued on the one stream, ahead of any later one. + pub fn drop(gpa: std.mem.Allocator, p: *Pages, k: *Saved) void { + p.pages.finish(&k.lease); + gpa.destroy(k); + } + }; +} + +const testing = std.testing; + +const Run = struct { + out: *std.ArrayList([3]usize), + fn add(r: Run, held: u64, done: usize, n: usize) anyerror!void { + try r.out.append(testing.allocator, .{ held, done, n }); + } +}; + +test "a region over snapshot pages: one copy a run of consecutive ids, split where the ids jump" { + const ids = [_]u32{ 4, 5, 9, 10, 11, 2 }; + var out: std.ArrayList([3]usize) = .empty; + defer out.deinit(testing.allocator); + // bytes [page / 2, 5 page + 10): pages 4-5 (from mid-page), 9-11, then 2 + try pageRuns(&ids, page / 2, 4 * page + page / 2 + 10, Run{ .out = &out }, Run.add); + try testing.expectEqual(@as(usize, 3), out.items.len); + try testing.expectEqual([3]usize{ 4 * page + page / 2, 0, page + page / 2 }, out.items[0]); + try testing.expectEqual([3]usize{ 9 * page, page + page / 2, 3 * page }, out.items[1]); + try testing.expectEqual([3]usize{ 2 * page, 4 * page + page / 2, 10 }, out.items[2]); + out.clearRetainingCapacity(); + try pageRuns(&ids, 3 * page + 7, 100, Run{ .out = &out }, Run.add); // inside one page + try testing.expectEqual([3]usize{ 10 * page + 7, 0, 100 }, out.items[0]); +} diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index ab6a6e9443..7c464d8d78 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -18,6 +18,8 @@ const kolibri1 = @import("kolibri1"); const gemma4 = @import("gemma4_cuda"); const spark25 = @import("spark25"); const dory = @import("dory"); +const lfm2_moe = @import("lfm2_moe_cuda"); +const lfm3 = @import("lfm3_cuda"); const Allocator = std.mem.Allocator; pub const budget = @import("cuda_memory.zig"); const lane_backend = @import("cuda_bound.zig"); @@ -28,7 +30,7 @@ const vision_offer = @import("cuda_vision_offer.zig"); const drafted = @import("cuda_drafted.zig"); const kv_pages = @import("cuda_kv_pool.zig"); /// CUDA families provide metadata, open their lane backend and explain their refusals. -const registry = .{ nemotron.native, nemotron.omni, glm.native, qwen35.native, qwen35.native_moe, flashnext.native, kolibri1.native, gemma4.native, gemma4.native31, spark25.native, dory.native }; +const registry = .{ nemotron.native, nemotron.omni, glm.native, qwen35.native, qwen35.native_moe, flashnext.native, kolibri1.native, gemma4.native, gemma4.native31, spark25.native, dory.native, lfm2_moe.native, lfm3.native }; /// --record opens here (engines.zig refuses it for a backend without this). pub const records = true; pub const backends: []const []const u8 = &.{"cuda"}; diff --git a/zig/tests/cuda/contract_elementwise.zig b/zig/tests/cuda/contract_elementwise.zig index 0b6d8c46da..554dc6258e 100644 --- a/zig/tests/cuda/contract_elementwise.zig +++ b/zig/tests/cuda/contract_elementwise.zig @@ -5,6 +5,9 @@ const core = @import("core"); const check = @import("check.zig"); const embed_k = @import("contract_embed.zig"); const penalize_k = @import("contract_penalize.zig"); +const lfm2m_k = @import("contract_lfm2_moe.zig"); +const paged_k = @import("contract_paged_attn.zig"); +const lfm3_k = @import("contract_lfm3.zig"); const g = core.contract.gates; @@ -15,6 +18,9 @@ pub const Set = struct { sample: ?cuda.Module = null, draws: cuda.sample.Draws = undefined, penalize: ?penalize_k.Penalize = null, + lfm2m: lfm2m_k.Set = .{}, // LFM2.5-24B-A2B's router, Liquid's 3-tap conv and commit (rows with sums, beside these) + paged: paged_k.Set = .{}, // the Liquid families' paged attention: head 64 decode and prompt rows, head 128 prompt rows + lfm3: lfm3_k.Set = .{}, // LFM3's 7-tap conv and commit /// Every registration onto `kernels`; the embed controls' symbols come from the contract_controls image `ctl_mod`. pub fn add(self: *Set, gpu: check.Gpu, s: cuda.Stream, ctl_mod: cuda.Module, kernels: *std.ArrayList(g.Kernel)) !void { @@ -29,9 +35,15 @@ pub const Set = struct { self.draws = try cuda.sample.Draws.resolve(self.sample.?); self.penalize = .{ .rig = .{ .gpu = gpu, .s = s }, .draws = &self.draws }; try kernels.append(a, self.penalize.?.kernel()); + try self.lfm2m.add(gpu, s, ctl_mod, kernels); + try self.paged.add(gpu, s, ctl_mod, kernels); + try self.lfm3.add(gpu, s, ctl_mod, kernels); } pub fn deinit(self: *Set) void { + self.lfm3.deinit(); + self.paged.deinit(); + self.lfm2m.deinit(); if (self.penalize) |*p| p.deinit(); if (self.sample) |*m| m.unload(); for (self.embeds[0..self.n_embeds]) |*x| x.deinit(); diff --git a/zig/tests/cuda/contract_lfm2_moe.zig b/zig/tests/cuda/contract_lfm2_moe.zig new file mode 100644 index 0000000000..8e6dca3256 --- /dev/null +++ b/zig/tests/cuda/contract_lfm2_moe.zig @@ -0,0 +1,505 @@ +//! lfm2_moe.cu under the contract: router logits and top-4 rule, 3-tap liquid conv and snapshots, commit; controlled. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); +const rig_ = @import("contract_rig.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const rows_ = contract.rows; +const Dense = contract.dense.Dense; + +const experts = 64; +const topk = 4; + +fn noSums(_: c.Shape) usize { + return 0; +} + +/// A quarter of the row an fma chain, then the quarters' two adds. +fn quarterChain(s: c.Shape) usize { + return s.k / 4 + 2; +} + +fn f32At(bytes: []const u8, i: usize) f32 { + return @bitCast(std.mem.readInt(u32, bytes[4 * i ..][0..4], .little)); +} + +fn routeArgs(x: u64, w: u64, bias: u64, picks: u64, wts: u64, logits: u64, rows: usize, d: usize) cuda.Args { + var a: cuda.Args = .{}; + for ([_]u64{ x, w, bias, picks, wts, logits }) |p| a.add(p); + a.add(@as(c_int, @intCast(rows))); + a.add(@as(c_int, @intCast(d))); + a.add(@as(f32, 1.0)); + return a; +} + +fn launch(f: cuda.Function, s: cuda.Stream, grid: [3]usize, a: *cuda.Args) !void { + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(grid[0]), .y = @intCast(grid[1]), .z = @intCast(grid[2]) }, .block = .{ .x = 256 } }, s, a); + try s.synchronize(); +} + +pub fn logitsSpec() c.Spec { + return .{ .name = "cuda lfm2_moe router logits tf_lfm2m_route (fp32 weights, quarters of the row)", .accum = .f32, .depth = quarterChain, .out = .f32, .tiles = .{ .{ .size = 8 }, .{}, .{ .size = 2048, .step = 4, .min = 4 } }, .served = &.{ .{ .rows = 1, .n = experts, .k = 2048 }, .{ .rows = 16, .n = experts, .k = 2048 }, .{ .rows = 1024, .n = experts, .k = 2048 } } }; +} + +/// x [rows, D] bf16 times fp32 router rows [64, D]: the logits (the picks and weights go to scratch). +pub const Logits = struct { + rig: rig_.Rig, + f: cuda.Function, + d: ?Dense = null, + dev: [4]u64 = @splat(0), + + pub fn deinit(self: *Logits) void { + if (self.d) |d| d.deinit(); + self.rig.deinit(); + } + + pub fn kernel(self: *Logits) g.Kernel { + return .{ .spec = logitsSpec(), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Logits); + + pub fn prepare(self: *Logits, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.n != experts or shape.k % 4 != 0 or shape.k == 0 or shape.k > 2048 or shape.rows == 0) return false; + if (self.d) |d| d.deinit(); + self.d = null; + self.rig.reset(); + const a = self.rig.gpu.gpa; + const d = try Dense.init(a, shape, 1); + self.d = d; + d.fillX(fill, seed, 16); + const w = try a.alloc(f32, d.w.len); + defer a.free(w); + var prng = std.Random.DefaultPrng.init(seed ^ 0x51); + const rng = prng.random(); + const k = shape.k; + // stress: positive integers, exact in fp32 in any order; mirrored signs would cancel inside a 32-lane fp16 sum + for (w, 0..) |*v, i| v.* = if (fill != .stress) rng.float(f32) * 2 - 1 else @floatFromInt(1 + (i % k * 7 + i / k) % 6); + for (d.w, w) |*o, v| o.* = v; + const bias: [experts]f32 = @splat(0); + self.dev[0] = try self.rig.input(std.mem.sliceAsBytes(d.x)); + self.dev[1] = try self.rig.input(std.mem.sliceAsBytes(w)); + self.dev[2] = try self.rig.input(std.mem.sliceAsBytes(&bias)); + self.dev[3] = try self.rig.output(shape.rows * experts * 4); + return true; + } + + pub fn run(self: *Logits) !void { + const d = self.d.?; + const picks = try self.rig.scratch(d.shape.rows * topk * 4); + const wts = try self.rig.scratch(d.shape.rows * topk * 4); + var a = routeArgs(self.dev[0], self.dev[1], self.dev[2], picks, wts, self.dev[3], d.shape.rows, d.shape.k); + try launch(self.f, self.rig.s, .{ (d.shape.rows + 7) / 8, 1, 1 }, &a); + } + + pub fn count(self: *Logits) usize { + return self.d.?.outputs(); + } + + pub fn reference(self: *Logits, want: []f64, mag: []f64) !void { + return self.d.?.reference(want, mag); + } + + pub fn simulate(self: *Logits, acc: c.Accum, order: g.Order, out: []f64) !void { + return self.d.?.simulate(acc, order, out); + } +}; + +const pick_rule = "top 4 of sigmoid(logit) + bias, ties to the lower id, a NaN ranked as -inf; weights sigmoid / (sum + 1e-6)"; + +pub fn selectSpec(control: bool) c.Spec { + return .{ .name = if (control) "control: lfm2_moe router with ties to the higher expert" else "cuda lfm2_moe router top-4 tf_lfm2m_route (its rule on fp32 logits)", .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = pick_rule }, .tiles = .{ .{ .size = 8 }, .{}, .{} }, .served = &.{ .{ .rows = 16, .n = experts, .k = experts }, .{ .rows = 300, .n = experts, .k = experts } } }; +} + +/// The rule's picks on logits the identity router passes through (logit e = x[e]): sixteenths, so ties are common. +pub const Select = struct { + rig: rig_.Rig, + f: cuda.Function, + control: bool = false, + shape: c.Shape = undefined, + x: []f32 = &.{}, + bias: [experts]f32 = undefined, + dev: [5]u64 = @splat(0), + msg: [200]u8 = undefined, + + pub fn deinit(self: *Select) void { + self.rig.gpu.gpa.free(self.x); + self.rig.deinit(); + } + + pub fn kernel(self: *Select) g.Kernel { + return .{ .spec = selectSpec(self.control), .ctx = self, .vt = &vt, .control = if (self.control) .shapes else null }; + } + + const vt = g.hooks(Select); + + pub fn prepare(self: *Select, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.n != experts or shape.k != experts or shape.rows == 0) return false; + self.rig.reset(); + const a = self.rig.gpu.gpa; + self.shape = shape; + a.free(self.x); + self.x = &.{}; + self.x = try a.alloc(f32, shape.rows * experts); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (self.x) |*v| v.* = @as(f32, @floatFromInt(rng.uintLessThan(u32, 64))) / 16 - 2; + for (&self.bias) |*b| b.* = if (rng.boolean()) 0 else 0.0625; + if (fill == .nonfinite) { + const at = [_][2]usize{ .{ 0, 3 }, .{ 1, 5 }, .{ 2, 7 }, .{ 3, 2 }, .{ 3, experts - 1 } }; + const put = [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32), std.math.inf(f32), -std.math.inf(f32) }; + for (at, put) |p, v| if (p[0] < shape.rows) { + self.x[p[0] * experts + p[1]] = v; + }; + } + const xb = try a.alloc(u16, self.x.len); + defer a.free(xb); + for (xb, self.x) |*b, v| b.* = c.toBf16(v); + var w: [experts * experts]f32 = @splat(0); + for (0..experts) |e| w[e * experts + e] = 1; + self.dev[0] = try self.rig.input(std.mem.sliceAsBytes(xb)); + self.dev[1] = try self.rig.input(std.mem.sliceAsBytes(&w)); + self.dev[2] = try self.rig.input(std.mem.sliceAsBytes(&self.bias)); + self.dev[3] = try self.rig.output(shape.rows * topk * 4); + self.dev[4] = try self.rig.output(shape.rows * topk * 4); + return true; + } + + pub fn run(self: *Select) !void { + const logits = try self.rig.scratch(self.shape.rows * experts * 4); + var a = routeArgs(self.dev[0], self.dev[1], self.dev[2], self.dev[3], self.dev[4], logits, self.shape.rows, experts); + try launch(self.f, self.rig.s, .{ (self.shape.rows + 7) / 8, 1, 1 }, &a); + } + + pub fn count(self: *Select) usize { + return self.shape.rows * topk * 8; + } + + pub fn reference(_: *Select, want: []f64, mag: []f64) !void { + @memset(want, 0); + @memset(mag, 0); + } + + /// Picks by the rule exactly, weights within 2^-20 of float64 (NaN where a picked logit is NaN). + pub fn check(self: *Select, _: g.Gate) !?[]const u8 { + const out = try self.rig.outputBytes(); + const n = self.shape.rows * topk; + for (0..self.shape.rows) |r| { + var lg: [experts]f64 = undefined; // the identity's logits as a dot: one non-finite x makes the whole row NaN + for (&lg, 0..) |*l, e| { + l.* = 0; + for (self.x[r * experts ..][0..experts], 0..) |v, i| l.* += @as(f64, v) * @as(f64, if (i == e) 1 else 0); + } + var taken: [experts]bool = @splat(false); + var gate: [topk]f64 = undefined; + var sum: f64 = 0; + for (0..topk) |k| { + var best: ?usize = null; + var bv: f64 = 0; + for (lg, 0..) |l, e| { + const s = 1 / (1 + @exp(-l)); + const sv: f64 = @as(f32, @floatCast(s)) + @as(f64, self.bias[e]); + const v = if (std.math.isNan(sv)) -std.math.inf(f64) else sv; + if (taken[e]) continue; + if (best == null or v > bv) { + best = e; + bv = v; + } + } + const id = best orelse return try std.fmt.bufPrint(&self.msg, "row {d}: fewer than 4 numbers", .{r}); + taken[id] = true; + const got = std.mem.readInt(u32, out[(r * topk + k) * 4 ..][0..4], .little); + if (got != id) return try std.fmt.bufPrint(&self.msg, "row {d} pick {d} is {d} where the rule gives {d}", .{ r, k, got, id }); + gate[k] = 1 / (1 + @exp(-lg[id])); + sum += gate[k]; + } + for (0..topk) |k| { + const want = gate[k] / (sum + 1e-6); + const got: f64 = f32At(out, n + r * topk + k); + if (std.math.isNan(want) and std.math.isNan(got)) continue; + if (!(@abs(got - want) <= @abs(want) * 0x1p-20 + 1e-30)) return try std.fmt.bufPrint(&self.msg, "row {d} slot {d}: weight {d:.8} where float64 gives {d:.8}", .{ r, k, got, want }); + } + } + return null; + } +}; + +const conv_rule = "the host's replay of each row's fp32 and bf16 steps, bit for bit (NaN as NaN)"; + +pub fn convSpec(taps: usize, control: bool) c.Spec { + const tiles: [3]c.Tile = .{ .{}, .{ .size = 256, .step = 64, .min = 64 }, .{} }; + if (taps == 7) return .{ .name = if (control) "control: liquid 7-tap conv with bx unrounded" else "cuda liquid 7-tap short conv tf_liquid_conv_7 (stream state, every row's snapshot)", .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = conv_rule }, .tiles = tiles, .served = &.{ .{ .rows = 1, .n = 2048, .k = 7 }, .{ .rows = 16, .n = 2048, .k = 7 }, .{ .rows = 128, .n = 2048, .k = 7 } } }; + return .{ .name = if (control) "control: liquid 3-tap conv with bx unrounded" else "cuda liquid 3-tap short conv tf_liquid_conv_3 (stream state, every row's snapshot)", .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = conv_rule }, .tiles = tiles, .served = &.{ .{ .rows = 1, .n = 2048, .k = 3 }, .{ .rows = 16, .n = 2048, .k = 3 }, .{ .rows = 128, .n = 2048, .k = 3 } } }; +} + +/// Rows of two streams (the first rows/2 slot 1, the rest slot 0) over their states: y and every row's state after it. +pub const Conv = struct { + rig: rig_.Rig, + f: cuda.Function, + control: bool = false, + taps: usize = 3, + shape: c.Shape = undefined, + p: []u16 = &.{}, + w: []u16 = &.{}, + st: []u16 = &.{}, + slot: []i32 = &.{}, + off: []i32 = &.{}, + dev: [7]u64 = @splat(0), + msg: [200]u8 = undefined, + + pub fn deinit(self: *Conv) void { + self.free(); + self.rig.deinit(); + } + + fn free(self: *Conv) void { + const a = self.rig.gpu.gpa; + inline for (.{ "p", "w", "st", "slot", "off" }) |f| { + a.free(@field(self, f)); + @field(self, f) = &.{}; + } + } + + pub fn kernel(self: *Conv) g.Kernel { + return .{ .spec = convSpec(self.taps, self.control), .ctx = self, .vt = &vt, .control = if (self.control) .shapes else null }; + } + + const vt = g.hooks(Conv); + + pub fn prepare(self: *Conv, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.n % 64 != 0 or shape.k != self.taps) return false; + self.rig.reset(); + self.free(); + const a = self.rig.gpu.gpa; + self.shape = shape; + const r = shape.rows; + const d = shape.n; + const sr = self.taps - 1; // state rows + self.p = try a.alloc(u16, r * 3 * d); + self.w = try a.alloc(u16, d * self.taps); + self.st = try a.alloc(u16, 2 * sr * d); + self.slot = try a.alloc(i32, r); + self.off = try a.alloc(i32, r); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (self.p) |*v| v.* = c.toBf16(rng.float(f32) * 4 - 2); + for (self.w) |*v| v.* = c.toBf16(rng.float(f32) - 0.5); + for (self.st) |*v| v.* = c.toBf16(rng.float(f32) * 4 - 2); + const half = r / 2; + for (self.slot, self.off, 0..) |*s, *o, i| { + s.* = if (i < half) 1 else 0; + o.* = @intCast(if (i < half) i else i - half); + } + if (fill == .nonfinite) { + const at = [_][2]usize{ .{ 0, d / 3 }, .{ 1, d + 1 }, .{ 2, 2 * d + 5 }, .{ 3, 0 }, .{ 3, 3 * d - 1 } }; + const put = [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32), std.math.inf(f32), -std.math.inf(f32) }; + for (at, put) |q, v| if (q[0] < r) { + self.p[q[0] * 3 * d + q[1]] = c.toBf16(v); + }; + } + self.dev[0] = try self.rig.input(std.mem.sliceAsBytes(self.p)); + self.dev[1] = try self.rig.input(std.mem.sliceAsBytes(self.w)); + self.dev[2] = try self.rig.input(std.mem.sliceAsBytes(self.st)); + self.dev[3] = try self.rig.input(std.mem.sliceAsBytes(self.slot)); + self.dev[4] = try self.rig.input(std.mem.sliceAsBytes(self.off)); + self.dev[5] = try self.rig.output(r * d * 2); + self.dev[6] = try self.rig.output(r * sr * d * 2); + return true; + } + + pub fn run(self: *Conv) !void { + const d = self.shape.n; + var a: cuda.Args = .{}; + for (self.dev) |p| a.add(p); + a.add(@as(c_int, 0)); + a.add(@as(c_int, @intCast(d))); + a.add(@as(i64, @intCast((self.taps - 1) * d))); + try launch(self.f, self.rig.s, .{ self.shape.rows, (d + 255) / 256, 1 }, &a); + } + + pub fn count(self: *Conv) usize { + return self.shape.rows * self.shape.n * 2 * self.taps; + } + + pub fn reference(_: *Conv, want: []f64, mag: []f64) !void { + @memset(want, 0); + @memset(mag, 0); + } + + fn val(h: u16) f32 { + return c.fromBf16(h); + } + + fn same(got: u16, want: u16) bool { + return got == want or (std.math.isNan(val(got)) and std.math.isNan(val(want))); + } + + pub fn check(self: *Conv, _: g.Gate) !?[]const u8 { + const out = try self.rig.outputBytes(); + const d = self.shape.n; + const r = self.shape.rows; + const taps = self.taps; + const sr = taps - 1; + for (0..r) |row| for (0..d) |ch| { + const o: usize = @intCast(self.off[row]); + var bx: [7]u16 = undefined; + for (0..taps) |j| { + const back = sr - j; + bx[j] = if (o >= back) blk: { + const src = (row - back) * 3 * d; + break :blk c.toBf16(val(self.p[src + ch]) * val(self.p[src + 2 * d + ch])); + } else self.st[@as(usize, @intCast(self.slot[row])) * sr * d + (sr - back + o) * d + ch]; + } + var acc: f32 = 0; + for (0..taps) |j| acc += val(self.w[ch * taps + j]) * val(bx[j]); + const y = c.toBf16(val(self.p[row * 3 * d + d + ch]) * val(c.toBf16(acc))); + for (0..taps) |which| { // y, then the row's state rows: its taps 1.. (oldest first) + const at = if (which == 0) row * d + ch else r * d + (row * sr + which - 1) * d + ch; + const want: u16 = if (which == 0) y else bx[which]; + const got = std.mem.readInt(u16, out[2 * at ..][0..2], .little); + if (!same(got, want)) return try std.fmt.bufPrint(&self.msg, "row {d} channel {d} output {d}: 0x{x:0>4} where the replay gives 0x{x:0>4}", .{ row, ch, which, got, want }); + } + }; + return null; + } +}; + +pub fn commitSpec(taps: usize, control: bool) c.Spec { + if (taps == 7) return .{ .name = if (control) "control: liquid 7-tap commit into the first stream's slot" else "cuda liquid 7-tap conv-state commit tf_liquid_commit_7", .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = "a copy: every bit as stored" }, .served = &.{ .{ .rows = 1, .n = 2048, .k = 13 }, .{ .rows = 16, .n = 2048, .k = 13 } } }; + return .{ .name = if (control) "control: liquid 3-tap commit into the first stream's slot" else "cuda liquid 3-tap conv-state commit tf_liquid_commit_3", .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = "a copy: every bit as stored" }, .served = &.{ .{ .rows = 1, .n = 2048, .k = 30 }, .{ .rows = 16, .n = 2048, .k = 30 } } }; +} + +/// Streams k = 0.. into slots rows - k of rows + 1, each its snapshot row 2k (stream 1: none, its slot as it was). +pub const Commit = struct { + rig: rig_.Rig, + f: cuda.Function, + control: bool = false, + taps: usize = 3, + shape: c.Shape = undefined, + snaps: []u16 = &.{}, + dev: [4]u64 = @splat(0), + msg: [200]u8 = undefined, + + pub fn deinit(self: *Commit) void { + self.rig.gpu.gpa.free(self.snaps); + self.rig.deinit(); + } + + pub fn kernel(self: *Commit) g.Kernel { + return .{ .spec = commitSpec(self.taps, self.control), .ctx = self, .vt = &vt, .control = if (self.control) .shapes else null }; + } + + const vt = g.hooks(Commit); + + fn rowOf(k: usize) i32 { + return if (k == 1) -1 else @intCast(2 * k); + } + + pub fn prepare(self: *Commit, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.k == 0) return false; + self.rig.reset(); + const a = self.rig.gpu.gpa; + self.shape = shape; + a.free(self.snaps); + self.snaps = &.{}; + const r = shape.rows; + const sw = (self.taps - 1) * shape.n; // a state's elements + self.snaps = try a.alloc(u16, shape.k * 2 * r * sw); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (self.snaps) |*v| v.* = if (fill == .nonfinite and rng.uintLessThan(u32, 97) == 0) 0x7fc1 else rng.int(u16); + const slots = try a.alloc(i32, r); + defer a.free(slots); + const rows = try a.alloc(i32, r); + defer a.free(rows); + for (slots, rows, 0..) |*s, *w, k| { + s.* = @intCast(r - k); + w.* = rowOf(k); + } + self.dev[0] = try self.rig.output(shape.k * (r + 1) * sw * 2); + self.dev[1] = try self.rig.input(std.mem.sliceAsBytes(self.snaps)); + self.dev[2] = try self.rig.input(std.mem.sliceAsBytes(slots)); + self.dev[3] = try self.rig.input(std.mem.sliceAsBytes(rows)); + return true; + } + + pub fn run(self: *Commit) !void { + const d = self.shape.n; + const r = self.shape.rows; + var a: cuda.Args = .{}; + for (self.dev) |p| a.add(p); + const sw = (self.taps - 1) * d; + a.add(@as(c_int, @intCast(d))); + a.add(@as(i64, @intCast((r + 1) * sw))); + a.add(@as(i64, @intCast(2 * r * sw))); + a.add(@as(i64, @intCast(sw))); + try launch(self.f, self.rig.s, .{ self.shape.k, r, 1 }, &a); + } + + pub fn count(self: *Commit) usize { + return self.shape.k * (self.shape.rows + 1) * (self.taps - 1) * self.shape.n * 2; + } + + pub fn reference(_: *Commit, want: []f64, mag: []f64) !void { + @memset(want, 0); + @memset(mag, 0); + } + + pub fn check(self: *Commit, _: g.Gate) !?[]const u8 { + const out = try self.rig.outputBytes(); + const d = self.shape.n; + const r = self.shape.rows; + const sw = (self.taps - 1) * d; + for (0..self.shape.k) |layer| for (0..r + 1) |slot| for (0..sw) |i| { + const k = r - slot; // the stream that owns this slot (slot 0: none) + const want: u16 = if (slot == 0 or rowOf(k) < 0) 0xffff else self.snaps[layer * 2 * r * sw + @as(usize, @intCast(rowOf(k))) * sw + i]; + const got = std.mem.readInt(u16, out[2 * ((layer * (r + 1) + slot) * sw + i) ..][0..2], .little); + if (got != want) return try std.fmt.bufPrint(&self.msg, "layer {d} slot {d} element {d}: 0x{x:0>4} where 0x{x:0>4}", .{ layer, slot, i, got, want }); + }; + return null; + } +}; + +/// The kernels and their controls, registered onto `kernels`; the controls from the contract_controls image `ctl`. +pub const Set = struct { + mod: ?cuda.Module = null, + logits: Logits = undefined, + select: [2]Select = undefined, + conv: [2]Conv = undefined, + commit: [2]Commit = undefined, + + pub fn add(self: *Set, gpu: check.Gpu, s: cuda.Stream, ctl: cuda.Module, kernels: *std.ArrayList(g.Kernel)) !void { + const a = gpu.gpa; + self.mod = try cuda.Module.load(gpu.d, cuda.kernels.lfm2_moe); + const m = self.mod.?; + const rig: rig_.Rig = .{ .gpu = gpu, .s = s }; + self.logits = .{ .rig = rig, .f = try m.function("tf_lfm2m_route") }; + for (&self.select, [_]bool{ false, true }) |*x, ctrl| x.* = .{ .rig = rig, .f = try (if (ctrl) ctl else m).function(if (ctrl) "tf_ctl_lfm2m_route" else "tf_lfm2m_route"), .control = ctrl }; + for (&self.conv, [_]bool{ false, true }) |*x, ctrl| x.* = .{ .rig = rig, .f = try (if (ctrl) ctl else m).function(if (ctrl) "tf_ctl_liquid_conv_3" else "tf_liquid_conv_3"), .control = ctrl }; + for (&self.commit, [_]bool{ false, true }) |*x, ctrl| x.* = .{ .rig = rig, .f = try (if (ctrl) ctl else m).function(if (ctrl) "tf_ctl_liquid_commit_3" else "tf_liquid_commit_3"), .control = ctrl }; + try kernels.append(a, self.logits.kernel()); + for (&self.select) |*x| try kernels.append(a, x.kernel()); + for (&self.conv) |*x| try kernels.append(a, x.kernel()); + for (&self.commit) |*x| try kernels.append(a, x.kernel()); + } + + pub fn deinit(self: *Set) void { + const m = self.mod orelse return; + self.logits.deinit(); + for (&self.select) |*x| x.deinit(); + for (&self.conv) |*x| x.deinit(); + for (&self.commit) |*x| x.deinit(); + var mm = m; + mm.unload(); + } +}; + +test "the commit's streams: stream 1 keeps its slot, the rest take row 2k" { + try std.testing.expectEqual(@as(i32, -1), Commit.rowOf(1)); + try std.testing.expectEqual(@as(i32, 6), Commit.rowOf(3)); + try std.testing.expectEqual(@as(usize, 514), quarterChain(.{ .rows = 1, .n = 64, .k = 2048 })); +} diff --git a/zig/tests/cuda/contract_lfm3.zig b/zig/tests/cuda/contract_lfm3.zig new file mode 100644 index 0000000000..0afd275517 --- /dev/null +++ b/zig/tests/cuda/contract_lfm3.zig @@ -0,0 +1,40 @@ +//! lfm3.cu under the contract: Liquid's 7-tap conv with its snapshots and its commit, each controlled (the 3-tap rigs). +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); +const rig_ = @import("contract_rig.zig"); +const lfm2m = @import("contract_lfm2_moe.zig"); + +const g = core.contract.gates; + +/// The kernels and their controls, registered onto `kernels`; the controls from the contract_controls image `ctl`. +pub const Set = struct { + mod: ?cuda.Module = null, + conv: [2]lfm2m.Conv = undefined, + commit: [2]lfm2m.Commit = undefined, + + pub fn add(self: *Set, gpu: check.Gpu, s: cuda.Stream, ctl: cuda.Module, kernels: *std.ArrayList(g.Kernel)) !void { + self.mod = try cuda.Module.load(gpu.d, cuda.kernels.lfm3); + const m = self.mod.?; + const rig: rig_.Rig = .{ .gpu = gpu, .s = s }; + for (&self.conv, [_]bool{ false, true }) |*x, ctrl| x.* = .{ .rig = rig, .f = try (if (ctrl) ctl else m).function(if (ctrl) "tf_ctl_liquid_conv_7" else "tf_liquid_conv_7"), .control = ctrl, .taps = 7 }; + for (&self.commit, [_]bool{ false, true }) |*x, ctrl| x.* = .{ .rig = rig, .f = try (if (ctrl) ctl else m).function(if (ctrl) "tf_ctl_liquid_commit_7" else "tf_liquid_commit_7"), .control = ctrl, .taps = 7 }; + for (&self.conv) |*x| try kernels.append(gpu.gpa, x.kernel()); + for (&self.commit) |*x| try kernels.append(gpu.gpa, x.kernel()); + } + + pub fn deinit(self: *Set) void { + const m = self.mod orelse return; + for (&self.conv) |*x| x.deinit(); + for (&self.commit) |*x| x.deinit(); + var mm = m; + mm.unload(); + } +}; + +test "the 7-tap rigs name their own kernels" { + try std.testing.expect(std.mem.indexOf(u8, lfm2m.convSpec(7, false).name, "tf_liquid_conv_7") != null); + try std.testing.expect(std.mem.indexOf(u8, lfm2m.commitSpec(7, true).name, "7-tap") != null); + try std.testing.expectEqual(@as(usize, 13), lfm2m.commitSpec(7, false).served[0].k); +} diff --git a/zig/tests/cuda/contract_limits.zig b/zig/tests/cuda/contract_limits.zig index 2a16afa5b6..83a6ffc468 100644 --- a/zig/tests/cuda/contract_limits.zig +++ b/zig/tests/cuda/contract_limits.zig @@ -25,7 +25,7 @@ pub fn image(name: []const u8) error{UnknownContractImage}![]const u8 { if (std.mem.eql(u8, name, "h3")) return &h3_image.bytes; inline for (.{ "qwen4_projection", "qwen4_hc", "qwen4_gdn", "qwen4_ple" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); - inline for (.{ "tree_round", "act_rows", "gdn_prep", "qk_rows", "norms", "moe", "experts", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "qmm_prefill", "lane_gemv", "prompt16", "qmmf", "ck_experts", "nvfp4_experts", "nvfp4_experts_tile", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "glm_glue", "glm_ops", "glm_index", "nemotron_mamba", "sample", "qwen35", "qwen35_bf16_rows", "qwen35_keep", "qwen35_kv8", "w4a8_prefill", "qwen35_draft", "ck_lane", "gdn_train", "train_q4", "train_attn", "flashnext_train", "train", "attn_train", "qwen27_train", "tree_attention", "block_attention", "window_attention", "window_pages", "moe_glue", "dense_rows", "qmm", "flashnext_fact", "flashnext_hcf", "flashnext_fexp", "experts_rows" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); + inline for (.{ "tree_round", "act_rows", "gdn_prep", "qk_rows", "norms", "moe", "experts", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "qmm_prefill", "lane_gemv", "prompt16", "qmmf", "ck_experts", "nvfp4_experts", "nvfp4_experts_tile", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "glm_glue", "glm_ops", "glm_index", "nemotron_mamba", "sample", "qwen35", "qwen35_bf16_rows", "qwen35_keep", "qwen35_kv8", "w4a8_prefill", "qwen35_draft", "ck_lane", "gdn_train", "train_q4", "train_attn", "flashnext_train", "train", "attn_train", "qwen27_train", "tree_attention", "block_attention", "window_attention", "window_pages", "moe_glue", "dense_rows", "qmm", "flashnext_fact", "flashnext_hcf", "flashnext_fexp", "experts_rows", "lfm2_moe", "lfm3" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); if (std.mem.eql(u8, name, "contract_controls")) return if (options.controls) &controls.bytes else &.{}; std.debug.print("FAIL kernel-contract: unknown image name: {s}\n", .{name}); return error.UnknownContractImage; diff --git a/zig/tests/cuda/contract_paged_attn.zig b/zig/tests/cuda/contract_paged_attn.zig new file mode 100644 index 0000000000..26774c317d --- /dev/null +++ b/zig/tests/cuda/contract_paged_attn.zig @@ -0,0 +1,349 @@ +//! Liquid's paged attention (head 64 LFM2.5, 128 LFM3), decode and prompt rows, over shuffled pages against float64. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); +const rig_ = @import("contract_rig.zig"); +const attn = @import("contract_attn.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; + +const page = 512; // positions a page holds (paged_attn.cuh's PAGE) +const max_segs = 8; // streams a decode launch takes (paged_attn.cuh's MAXSEG) +const chunk_blocks = 64; // key chunks a decode launch walks in parallel at most (the families' chunk_blocks) +const pad = 64; // elements between a page's values and the next page: a stride that ignores PS reads them + +/// Decode: the chunk and merge kernels' symbols; prompt: pattn_kernel's symbol, rows a block, dynamic shared bytes. +pub const Form = union(enum) { + decode: struct { chunk: [:0]const u8, merge: [:0]const u8 }, + prompt: struct { sym: [:0]const u8, rows: usize, smem: u32 }, +}; + +/// A launch's fixed heads and head dim, its kernels, and whether they come from the controls image. +pub const Cfg = struct { name: []const u8, h: usize, hk: usize, d: usize, form: Form, control: bool = false, served: []const c.Shape }; + +const pattn64 = "_ZN20tf_prefill_attention12pattn_kernelILi64ELi8ELi4ELi8EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifPKix"; +const pattn128 = "_ZN20tf_prefill_attention12pattn_kernelILi128ELi8ELi4ELi8EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifPKix"; +const prompt128_served = [_]c.Shape{ .{ .rows = 256, .n = 16, .k = 0 }, .{ .rows = 64, .n = 16, .k = 3000 }, .{ .rows = 100, .n = 16, .k = 500 } }; +const decode64: Form = .{ .decode = .{ .chunk = "tf_paged_attn_chunk_segs_64", .merge = "tf_paged_attn_merge_segs_64" } }; +const prompt64: Form = .{ .prompt = .{ .sym = pattn64, .rows = 32, .smem = 8 * 32 * 8 * 16 } }; +const decode_served = [_]c.Shape{ .{ .rows = 1, .n = 32, .k = 3000 }, .{ .rows = 16, .n = 32, .k = 1500 }, .{ .rows = 128, .n = 32, .k = 700 }, .{ .rows = 8, .n = 32, .k = 4000 } }; +const prompt_served = [_]c.Shape{ .{ .rows = 256, .n = 32, .k = 0 }, .{ .rows = 64, .n = 32, .k = 3000 }, .{ .rows = 100, .n = 32, .k = 500 } }; + +pub const cfgs = [_]Cfg{ + // LFM2.5-24B-A2B: 32 query heads over 8 KV heads of 64 + .{ .name = "cuda paged decode tf_paged_attn_chunk_segs_64 + tf_paged_attn_merge_segs_64, 32/8 heads of 64, 512-key pages", .h = 32, .hk = 8, .d = 64, .form = decode64, .served = &decode_served }, + .{ .name = "control: paged decode at head 64 reading its values from the neighbouring lane's dims", .h = 32, .hk = 8, .d = 64, .form = .{ .decode = .{ .chunk = "tf_ctl_paged_attn_chunk_segs_64", .merge = "tf_paged_attn_merge_segs_64" } }, .control = true, .served = &decode_served }, + .{ .name = "cuda paged prompt pattn_kernel<64, 8, 4, 8>, 32/8 heads of 64, 512-key pages", .h = 32, .hk = 8, .d = 64, .form = prompt64, .served = &prompt_served }, + .{ .name = "control: paged prompt rows at head 64 reading their values from the neighbouring 8 dims", .h = 32, .hk = 8, .d = 64, .form = .{ .prompt = .{ .sym = "tf_ctl_pattn_64", .rows = 32, .smem = 8 * 32 * 8 * 16 } }, .control = true, .served = &prompt_served }, + // LFM3's full-attention layers: 16 query heads over 4 KV heads of 128 + .{ .name = "cuda paged prompt pattn_kernel<128, 8, 4, 8>, 16/4 heads of 128, 512-key pages", .h = 16, .hk = 4, .d = 128, .form = .{ .prompt = .{ .sym = pattn128, .rows = 32, .smem = 8 * 32 * 16 * 16 } }, .served = &prompt128_served }, + .{ .name = "control: paged prompt rows at head 128 reading their values from the neighbouring 8 dims", .h = 16, .hk = 4, .d = 128, .form = .{ .prompt = .{ .sym = "tf_ctl_pattn_128", .rows = 32, .smem = 8 * 32 * 16 * 16 } }, .control = true, .served = &prompt128_served }, +}; + +/// Gates 1 and 3 run attention_ref's bound: decode weights are fp32 through fast exp, prompt P enters the mma as bf16. +const decode_why = "a score's absolute rounding error is its weight's relative error: gates 1 and 3 hold outputs to attention_ref's bound"; + +/// The P x V chain of a row: every key it sees (positions below k, then the rows before it), a tile, then the fold. +fn keyChain(s: c.Shape) usize { + return s.k + s.rows + 64 + 64; +} + +pub fn spec(f: Cfg) c.Spec { + return switch (f.form) { + .decode => .{ .name = f.name, .accum = .f32, .depth = keyChain, .ops = 1, .out = .bf16, .tiles = .{ .{}, .{}, .{ .size = 64 } }, .served = f.served, .sums = .{ .why = decode_why } }, + .prompt => |p| .{ .name = f.name, .accum = .f32, .depth = keyChain, .ops = 1, .out = .bf16, .tiles = .{ .{ .size = p.rows }, .{}, .{ .size = 64, .min = 0 } }, .served = f.served, .sums = .{ .why = attn.why } }, + }; +} + +/// nemotron_attention.cu's AttnSegs (paged_attn.cuh): each stream's caches, position, page table and first row. +const AttnSegs = extern struct { kc: [max_segs]u64, vc: [max_segs]u64, meta: [max_segs]u64, pt: [max_segs]u64, ps: i64, row0: [max_segs]c_int, n: c_int }; + +/// One stream: its rows from row0 on at positions pos0.., and its pages (the first of them at `first` in `ids`). +const Seg = struct { row0: usize, rows: usize, pos0: usize, first: usize, pages: usize }; + +pub const Paged = struct { + rig: rig_.Rig, + f: Cfg, + fns: [2]cuda.Function, // decode: chunk, merge; prompt: the kernel twice + r: ?attn.Ref = null, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + segs: std.ArrayList(Seg) = .empty, + ids: std.ArrayList(i32) = .empty, // every stream's page table, one after another + arena: std.ArrayList(u16) = .empty, // pages of PS elements: page rows of keys, then of values, then pad + q: std.ArrayList(u16) = .empty, + dev: struct { q: u64 = 0, arena: u64 = 0, meta: u64 = 0, ids: u64 = 0, out: u64 = 0, xs: u64 = 0, po: u64 = 0, pm: u64 = 0, pl: u64 = 0 } = .{}, + nch: usize = 0, + + pub fn deinit(self: *Paged) void { + const a = self.rig.gpu.gpa; + if (self.r) |*r| r.deinit(); + inline for (.{ "segs", "ids", "arena", "q" }) |n| @field(self, n).deinit(a); + self.rig.deinit(); + } + + pub fn kernel(self: *Paged) g.Kernel { + return .{ .spec = spec(self.f), .ctx = self, .vt = &vt, .control = if (self.f.control) .shapes else null }; + } + + const vt = g.hooks(Paged); + + fn kvd(self: *const Paged) usize { + return self.f.hk * self.f.d; + } + + fn ps(self: *const Paged) usize { + return 2 * page * self.kvd() + pad; + } + + fn scale(self: *const Paged) f64 { + return @floatCast(@as(f32, 1) / @sqrt(@as(f32, @floatFromInt(self.f.d)))); + } + + /// Row r's stream and position. + fn at(self: *const Paged, r: usize) struct { seg: Seg, p: usize } { + for (self.segs.items) |s| if (r < s.row0 + s.rows) return .{ .seg = s, .p = s.pos0 + r - s.row0 }; + unreachable; + } + + /// Element of key `key`'s row (key head kh) in stream s's pages; values sit page * kvd further on. + fn keyAt(self: *const Paged, s: Seg, key: usize, kh: usize) usize { + const id: usize = @intCast(self.ids.items[s.first + key / page]); + return id * self.ps() + (key % page) * self.kvd() + kh * self.f.d; + } + + fn value(self: *const Paged, s: Seg, key: usize, kh: usize) usize { + return self.keyAt(s, key, kh) + page * self.kvd(); + } + + pub fn prepare(self: *Paged, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + self.rig.reset(); + const f = self.f; + const n = shape.rows; + if (shape.n != f.h or n == 0) return false; + const a = self.rig.gpu.gpa; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + self.shape = shape; + // streams: decode rows split over up to 8 (one under stress), the first ending at k; a prompt chunk is one at k + const nseg: usize = if (f.form == .prompt or fill == .stress) 1 else @min(n, max_segs); + self.segs.clearRetainingCapacity(); + var row0: usize = 0; + var pages: usize = 0; + for (0..nseg) |j| { + const rows = n / nseg + @intFromBool(j < n % nseg); + const pos0 = switch (f.form) { + .prompt => shape.k, + .decode => blk: { + if (shape.k < rows) return false; + const room = shape.k - rows + 1; + break :blk if (j == 0) room - 1 else (j * 7919 + seed % 97) % room; + }, + }; + const need = (pos0 + rows + page - 1) / page; + try self.segs.append(a, .{ .row0 = row0, .rows = rows, .pos0 = pos0, .first = pages, .pages = need }); + row0 += rows; + pages += need; + } + // page ids: every stream's pages shuffled over the arena, two spare pages left random + const total = pages + 2; + try self.ids.resize(a, total); + for (self.ids.items, 0..) |*x, i| x.* = @intCast(i); + rng.shuffle(i32, self.ids.items); + try self.arena.resize(a, total * self.ps()); + for (self.arena.items) |*e| e.* = attn.normal(rng); + try self.q.resize(a, n * f.h * f.d); + for (self.q.items) |*e| e.* = attn.normal(rng); + var most: usize = 1; + var top: usize = 0; + for (self.segs.items) |s| { + most = @max(most, s.pos0 + s.rows); + top = @max(top, s.pos0 + s.rows); + } + self.nch = (top + page - 1) / page + 1; + if (self.r) |*r| r.deinit(); + self.r = null; + self.r = try attn.Ref.init(a, f.d, most); + switch (fill) { + .random => {}, + .nonfinite => self.plant(), + .stress => try self.stress(rng), + } + const metas = try a.alloc(i32, nseg); + defer a.free(metas); + for (metas, self.segs.items) |*m, s| m.* = @intCast(s.pos0); + self.dev.q = try self.rig.input(std.mem.sliceAsBytes(self.q.items)); + self.dev.arena = try self.rig.input(std.mem.sliceAsBytes(self.arena.items)); + self.dev.meta = try self.rig.input(std.mem.sliceAsBytes(metas)); + self.dev.ids = try self.rig.input(std.mem.sliceAsBytes(self.ids.items)); + self.dev.out = try self.rig.output(n * f.h * f.d * 2); + if (f.form == .decode) { + self.dev.xs = try self.rig.output(n * f.h * f.d / 64 * 4); + const parts = n * self.nch * f.h; + self.dev.po = try self.rig.scratch(parts * f.d * 4); + self.dev.pm = try self.rig.scratch(parts * 4); + self.dev.pl = try self.rig.scratch(parts * 4); + } + return true; + } + + /// NaN query (row 0, head 0); +inf and -inf values at row 0's key, NaN where they meet +inf a key earlier; NaN key. + fn plant(self: *Paged) void { + const f = self.f; + const s = self.segs.items[0]; + const p = s.pos0; // row 0's own key: every row of its stream sees it + self.q.items[3] = c.toBf16(std.math.nan(f32)); + self.arena.items[self.value(s, p, 0) + 7] = c.toBf16(std.math.inf(f32)); + self.arena.items[self.value(s, p, 1 % f.hk) + 9] = c.toBf16(-std.math.inf(f32)); + if (p > 0) self.arena.items[self.value(s, p - 1, 1 % f.hk) + 9] = c.toBf16(std.math.inf(f32)); + self.arena.items[self.keyAt(s, p, 2 % f.hk)] = c.toBf16(std.math.nan(f32)); + } + + /// Gate 4: every key of the stream alike (each p is 1); values mirrored over the keys row 0 sees, 0 after them. + fn stress(self: *Paged, rng: std.Random) !void { + const a = self.rig.gpu.gpa; + const f = self.f; + const s = self.segs.items[0]; + const key = try a.alloc(u16, self.kvd()); + defer a.free(key); + for (key) |*e| e.* = attn.normal(rng); + const last = s.pos0 + s.rows; + for (0..last) |k| { + @memcpy(self.arena.items[self.keyAt(s, k, 0)..][0..self.kvd()], key); + @memset(self.arena.items[self.value(s, k, 0)..][0..self.kvd()], 0); + } + const col = try a.alloc(f32, s.pos0 + 1); + defer a.free(col); + for (0..f.hk) |kh| for (0..f.d) |e| { + attn.mirrored(rng, col); + for (col, 0..) |y, k| self.arena.items[self.value(s, k, kh) + e] = c.toBf16(y); + }; + } + + pub fn run(self: *Paged) !void { + const f = self.f; + const n = self.shape.rows; + const g_: usize = f.h / f.hk; + const s = self.rig.s; + switch (f.form) { + .decode => { + var sg = std.mem.zeroes(AttnSegs); + for (self.segs.items, 0..) |x, j| { + sg.kc[j] = self.dev.arena; + sg.vc[j] = self.dev.arena + page * self.kvd() * 2; + sg.meta[j] = self.dev.meta + j * 4; + sg.pt[j] = self.dev.ids + x.first * 4; + sg.row0[j] = @intCast(x.row0); + } + sg.ps = @intCast(self.ps()); + sg.n = @intCast(self.segs.items.len); + var a: cuda.Args = .{}; + a.add(sg); + inline for (.{ "q", "po", "pm", "pl" }) |name| a.add(@field(self.dev, name)); + for ([_]usize{ f.h * f.d, f.h, f.hk, g_, page, self.nch }) |v| a.add(@as(c_int, @intCast(v))); + a.add(@as(f32, @floatCast(self.scale()))); + try cuda.launch.launch(self.fns[0], .{ .grid = .{ .x = @intCast(n), .y = @intCast(f.hk), .z = @intCast(@min(self.nch, chunk_blocks)) }, .block = .{ .x = 256 } }, s, &a); + var b: cuda.Args = .{}; + b.add(sg); + inline for (.{ "po", "pm", "pl", "out", "xs" }) |name| b.add(@field(self.dev, name)); + for ([_]usize{ f.h, f.hk, g_, page, self.nch }) |v| b.add(@as(c_int, @intCast(v))); + try cuda.launch.launch(self.fns[1], .{ .grid = .{ .x = @intCast(n), .y = @intCast(f.hk), .z = 1 }, .block = .{ .x = 256 } }, s, &b); + }, + .prompt => |p| { + const x = self.segs.items[0]; + var a: cuda.Args = .{}; + a.add(self.dev.q); + a.add(self.dev.arena); + a.add(self.dev.arena + page * self.kvd() * 2); + a.add(self.dev.out); + for ([_]usize{ x.pos0, n, f.h, f.hk, g_ }) |v| a.add(@as(c_int, @intCast(v))); + a.add(@as(f32, @floatCast(self.scale()))); + a.add(self.dev.ids + x.first * 4); + a.add(@as(i64, @intCast(self.ps()))); + const hpc = 16 * 8 / p.rows; // a block's query heads: 16 rows a warp, eight warps + try cuda.launch.launch(self.fns[0], .{ .grid = .{ .x = @intCast((n + p.rows - 1) / p.rows), .y = @intCast(f.hk * (g_ / hpc)) }, .block = .{ .x = 256 }, .shared = p.smem }, s, &a); + }, + } + try s.synchronize(); + } + + pub fn count(self: *Paged) usize { + return self.shape.rows * self.f.h * self.f.d; + } + + /// Output head i (row i / h, head i % h): its stream's keys [0, p] through the page table. + pub fn load(self: *Paged, i: usize, hd: *attn.Head) !void { + const f = self.f; + const d = f.d; + const r = i / f.h; + const head = i % f.h; + const kh = head / (f.h / f.hk); + const w = self.at(r); + for (hd.q, 0..) |*x, e| x.* = c.fromBf16(self.q.items[(r * f.h + head) * d + e]); + for (0..w.p + 1) |key| { + const kat = self.keyAt(w.seg, key, kh); + const vat = self.value(w.seg, key, kh); + for (0..d) |e| { + hd.k[key * d + e] = c.fromBf16(self.arena.items[kat + e]); + hd.v[key * d + e] = c.fromBf16(self.arena.items[vat + e]); + } + } + hd.n = w.p + 1; + } + + fn heads(self: *Paged) usize { + return self.shape.rows * self.f.h; + } + + pub fn check(self: *Paged, gate: g.Gate) !?[]const u8 { + return self.r.?.check(self, self.heads(), self.scale(), try self.rig.outputBytes(), gate); + } + + pub fn reference(self: *Paged, want: []f64, mag: []f64) !void { + return self.r.?.reference(self, self.heads(), self.scale(), want, mag); + } + + pub fn simulate(self: *Paged, acc: c.Accum, order: g.Order, out: []f64) !void { + return self.r.?.simulate(self, self.heads(), self.scale(), acc, order, out); + } +}; + +/// The kernels and their controls, registered onto `kernels`; the controls from the contract_controls image `ctl`. +pub const Set = struct { + mods: [2]?cuda.Module = .{ null, null }, + items: [cfgs.len]Paged = undefined, + n: usize = 0, + + pub fn add(self: *Set, gpu: check.Gpu, s: cuda.Stream, ctl: cuda.Module, kernels: *std.ArrayList(g.Kernel)) !void { + self.mods[0] = try cuda.Module.load(gpu.d, cuda.kernels.nemotron_attention); + self.mods[1] = try cuda.Module.load(gpu.d, cuda.kernels.prefill_attention); + const segs = self.mods[0].?; + const prompt = self.mods[1].?; + for (cfgs, &self.items) |f, *x| { + const fns: [2]cuda.Function = switch (f.form) { + .decode => |k| .{ try (if (f.control) ctl else segs).function(k.chunk), try segs.function(k.merge) }, + .prompt => |k| blk: { + const fn_ = try (if (f.control) ctl else prompt).function(k.sym); + try fn_.allowDynamicShared(k.smem); + break :blk .{ fn_, fn_ }; + }, + }; + x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .f = f, .fns = fns }; + self.n += 1; + } + for (self.items[0..self.n]) |*x| try kernels.append(gpu.gpa, x.kernel()); + } + + pub fn deinit(self: *Set) void { + for (self.items[0..self.n]) |*x| x.deinit(); + for (&self.mods) |*m| if (m.*) |*x| x.unload(); + } +}; + +test "the paged forms' geometry: 4 query heads a KV head, AttnSegs as nemotron_attention.cu pads it" { + try std.testing.expectEqual(@as(usize, 304), @sizeOf(AttnSegs)); + for (cfgs) |f| try std.testing.expectEqual(@as(usize, 4), f.h / f.hk); + try std.testing.expectEqual(@as(usize, 32), prompt64.prompt.rows); +}