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/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/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/zig/build/cuda.zig b/zig/build/cuda.zig index 054d85922b..61e81bab26 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -446,7 +446,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st @import("glm_family_host.zig").steps(b, cuda, glm, mods, step); @import("kolibri.zig").host(b, kolibri1Module(b, host, .debug, cuda, mods.core, mods.lanes), native_modules.api, step); @import("gemma4_cuda.zig").host(b, @import("gemma4_cuda.zig").module(b, host, .debug, cuda, mods.core, mods.lanes), step); - @import("spark25.zig").host(b, @import("spark25.zig").module(b, host, .debug, cuda, mods.core, mods.lanes), step); + inline for (.{ @import("spark25.zig"), @import("lfm2_moe_cuda.zig") }) |m| m.host(b, m.module(b, host, .debug, cuda, mods.core, mods.lanes), step); @import("glm_contracts.zig").tests(b, glm, mods.lanes, cuda, mods.core); // the shares a GLM-5.3-Flash rank takes (cluster/canon.zig against the Python family's) const shares = b.createModule(.{ .root_source_file = b.path("zig/src/cluster/canon_glm_test.zig"), .target = host, .optimize = .debug }); diff --git a/zig/build/cuda_engines.zig b/zig/build/cuda_engines.zig index bb929f7a23..e2f2c40b1a 100644 --- a/zig/build/cuda_engines.zig +++ b/zig/build/cuda_engines.zig @@ -52,6 +52,7 @@ pub fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu .{ .name = "kolibri1", .module = @import("kolibri.zig").module(b, target, optimize, cuda, core, lanes) }, .{ .name = "gemma4_cuda", .module = @import("gemma4_cuda.zig").module(b, target, optimize, cuda, core, lanes) }, .{ .name = "spark25", .module = @import("spark25.zig").module(b, target, optimize, cuda, core, lanes) }, + .{ .name = "lfm2_moe_cuda", .module = @import("lfm2_moe_cuda.zig").module(b, target, optimize, cuda, core, lanes) }, .{ .name = "dory", .module = @import("dory.zig").module(b, target, optimize, cuda, core, lanes) }, }, }); diff --git a/zig/build/cuda_kernels.zig b/zig/build/cuda_kernels.zig index 887b2d84ba..38a3d44ae3 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -83,6 +83,7 @@ pub const kernels = [_]Kernel{ .{ .name = "norms", .flags = &.{"-O3"} }, // ours: RMS norm rows for every family (cuda/norms.zig) .{ .name = "embed_rows", .flags = &.{"-O3"} }, // ours: MLX affine and bf16 token rows (cuda/embed_rows.zig) .{ .name = "lfm2", .flags = glue }, // ours: LFM2's short convolution, QK norm and rotary, gates and label rows + .{ .name = "lfm2_moe", .flags = glue }, // ours: LFM2.5-24B-A2B's router, Liquid's 3-tap conv and commit (liquid_conv.cuh) .{ .name = "vit", .flags = glue }, // ours: vision towers' LayerNorm, Linear epilogues, activations, unshuffle .{ .name = "tree_round", .flags = &.{"-O3"} }, // ours: verified tree rounds for every family (cuda/tree_round.zig) .{ .name = "act_rows", .flags = &.{"-O3"} }, // ours: SwiGLU, output gates, group sums (cuda/act_rows.zig) diff --git a/zig/build/lfm2_moe_cuda.zig b/zig/build/lfm2_moe_cuda.zig new file mode 100644 index 0000000000..e02fb69494 --- /dev/null +++ b/zig/build/lfm2_moe_cuda.zig @@ -0,0 +1,12 @@ +//! LFM2.5-24B-A2B's CUDA family over the maintained runtime, the core and the lane core; its host tests. +const std = @import("std"); + +pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module) *std.Build.Module { + return b.createModule(.{ .root_source_file = b.path("zig/src/families/lfm2_moe/cuda_family.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core }, .{ .name = "lanes", .module = lanes } } }); +} + +pub fn host(b: *std.Build, m: *std.Build.Module, all: *std.Build.Step) void { + const run = b.addRunArtifact(b.addTest(.{ .root_module = m })); + b.step("test-lfm2-moe-cuda-host", "LFM2.5-24B-A2B on CUDA: config admission, page and state sizes, RoPE table, refusals without a GPU").dependOn(&run.step); + all.dependOn(&run.step); +} diff --git a/zig/kernels/cuda/contract_controls.cu b/zig/kernels/cuda/contract_controls.cu index 1cc1fd9c84..c3e58c102c 100644 --- a/zig/kernels/cuda/contract_controls.cu +++ b/zig/kernels/cuda/contract_controls.cu @@ -342,3 +342,59 @@ TF_CTL_SWIGLU_L2(tf_ctl_swiglu_l2_drop_gs, tf_act_l2::US | tf_act_l2::PROD) TF_CTL_SWIGLU_L2(tf_ctl_swiglu_l2_drop_us, tf_act_l2::GS | tf_act_l2::PROD) TF_CTL_SWIGLU_L2(tf_ctl_swiglu_l2_drop_prod, tf_act_l2::GS | tf_act_l2::US) #undef TF_CTL_SWIGLU_L2 + +// Liquid's 3-tap short conv and commit, one planted fault each: conv bx unrounded, commits to stream 0's slot. +#include "liquid_conv.cuh" + +extern "C" __global__ void __launch_bounds__(256) tf_ctl_liquid_conv_3(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<3, 1>(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); +} + +// 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" +} +} // 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/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..db33c34253 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,10 @@ 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 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/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/registry.zig b/zig/src/core/contract/registry.zig index 7d83d71249..efa15e00ef 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; /// 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..5e059a1b7e 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -34,6 +34,7 @@ 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 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..681c7e64c2 --- /dev/null +++ b/zig/src/families/lfm2_moe/cuda_backend.zig @@ -0,0 +1,491 @@ +//! LFM2.5-24B-A2B on CUDA behind lanes.Backend: a slot a stream, a chunk a prompt step, one forward a round of windows. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const cfg = @import("config.zig"); +const fwd = @import("cuda_forward.zig"); +const mdl = @import("cuda_model.zig"); +const kern = @import("cuda_kernels.zig"); +const kv = @import("core").kv_pool; +const snap = @import("cuda_snap.zig"); + +const be = lanes.backend; +const Stream = lanes.Stream; +const Pick = cuda.tree_round.Pick; + +/// 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 = 16; +/// 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 fn pagesFor(prompt: usize, max_new: usize) usize { + return (prompt + max_new + max_window + kern.page - 1) / kern.page; +} + +pub const Lfm = 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) + tapped: ?struct { stream: *const Stream, cut: Cut } = null, // the chunk's tapped marks while its hooks run + snaps: ?snap.Pages = null, // the prompt cache's pages, allocated once inside its budget (snap_ready) + + /// 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) !Lfm { + 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: *Lfm, 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: *Lfm) 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: *Lfm) 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, streams sharing a forward. + pub fn facts(x: *const Lfm) 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) *Lfm { + return @ptrCast(@alignCast(ptr)); + } + + fn upload(x: *Lfm, 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: *Lfm, 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 pick over non-finite logits is an error. + fn tokens(x: *Lfm, 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: *Lfm, 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 tap: ?usize = if (x.tapped) |t| if (t.stream == s) std.mem.indexOfScalar(usize, t.cut.taps(), at) else null else null; + if (tap == null and (l.fill orelse l.len) != at) return error.NotAtMark; // where the pass stands, or a tap + const conv = if (tap) |j| x.m.s.tap + j * mdl.slot_elems * 2 else x.m.states + l.slot * mdl.slot_elems * 2; + return snap.save(x.gpa, try x.snapPages(), x.m, x.s, l.slot, at, conv); + } + + /// The stream's slot made to hold `k`'s state (a restore outside its prompt pass). + pub fn restoreState(x: *Lfm, 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: *Lfm, k: *snap.Saved) void { + snap.drop(x.gpa, &x.snaps.?, k); + } + + fn snapPages(x: *Lfm) !*snap.Pages { + return if (x.snaps) |*p| p else error.NoSnapshotCache; + } + + fn freeSlot(x: *const Lfm) ?usize { + return std.mem.indexOfScalar(bool, x.free[0..x.m.slots], true); + } + + fn fitsContext(x: *const Lfm, 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 Lfm, 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: *Lfm, + 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: *Lfm, 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: *Lfm, s: *Stream) !*Lane { + const p = s.prompt(); + if (p.len == 0) return error.EmptyPrompt; + if (s.images.len > 0) return error.LfmTextOnly; + 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: *Lfm, 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 row's logits are written. + fn chunk(x: *Lfm, 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); + const done = c.end == p.len; + var rows: [mdl.max_taps]usize = undefined; + for (c.taps(), 0..) |t, j| rows[j] = t - at - 1; + try fwd.forward(x.m, x.k, x.s, &.{.{ .slot = l.slot, .p0 = at, .tokens = p[at..c.end], .taps = rows[0..c.n] }}, true, c.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(c.end), .monotonic); + l.fill = if (done) null else c.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(c.end)) != null) hk.at(hk.ptr, s, @intCast(c.end)); + } + return done; + } + + /// A chunk's end (the grid line, the prompt's end or a mark) and the marks inside it, tapped in order. + const Cut = struct { + end: usize, + at: [mdl.max_taps]usize = undefined, + n: usize = 0, + fn taps(c: *const Cut) []const usize { + return c.at[0..c.n]; + } + }; + + /// The chunk from `at`: up to max_taps marks inside it are tapped, the next one ends it. + fn cut(at: usize, len: usize, marks: []const u32) Cut { + var c: Cut = .{ .end = @min(len, (at / prompt_rows + 1) * prompt_rows) }; + var from = at; + while (true) { // the marks inside in order, the first max_taps tapped + 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 == mdl.max_taps) { + c.end = mk; + break; + } + c.at[c.n] = mk; + c.n += 1; + from = mk; + } + return c; + } + + 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: *Lfm) 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: *Lfm, 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, a request for drafts is 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: *Lfm, 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 Lfm) 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), pagesFor(1, 1)); + try std.testing.expectEqual(@as(usize, 1), pagesFor(400, 96)); + try std.testing.expectEqual(@as(usize, 2), pagesFor(400, 97)); + try std.testing.expectEqual(@as(usize, 65), pagesFor(32768, 0)); +} + +test "a chunk taps the marks inside it, up to four, and ends at the next one, the grid line or the prompt's end" { + const cases = [_]struct { at: usize, len: usize, marks: []const u32, end: usize, taps: []const usize }{ + .{ .at = 0, .len = 3000, .marks = &.{}, .end = prompt_rows, .taps = &.{} }, + .{ .at = 0, .len = 3000, .marks = &.{700}, .end = prompt_rows, .taps = &.{700} }, + .{ .at = 0, .len = 3000, .marks = &.{ 900, 300, 2900 }, .end = prompt_rows, .taps = &.{ 300, 900 } }, + .{ .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 = prompt_rows, .taps = &.{} }, + }; + for (cases) |k| { + const c = Lfm.cut(k.at, k.len, k.marks); + try std.testing.expectEqual(k.end, c.end); + try std.testing.expectEqualSlices(usize, k.taps, c.taps()); + } +} 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..6a70bf38b6 --- /dev/null +++ b/zig/src/families/lfm2_moe/cuda_model.zig @@ -0,0 +1,339 @@ +//! 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; +/// 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); one more ends the chunk. +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; + } + + /// 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, 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); + 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, 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); + 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..b377869309 --- /dev/null +++ b/zig/src/families/lfm2_moe/cuda_native.zig @@ -0,0 +1,235 @@ +//! LFM2.5-24B-A2B (MLX affine 4-bit) behind the native CUDA registry: plain and copy-drafted decoding, prompt reuse. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const cfg = @import("config.zig"); +const wts = @import("cuda_weights.zig"); +const kern = @import("cuda_kernels.zig"); +const mdl = @import("cuda_model.zig"); +const be = @import("cuda_backend.zig"); +const snap = @import("cuda_snap.zig"); +const kv = @import("core").kv_pool; + +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 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; +pub const explains_open = true; + +/// 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.Lfm, + /// 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.Lfm }; + +/// Streams to allocate in `room`: as asked, else (auto) the most that fit; a fixed count past them or a round refuses. +fn fit(room: u64, asked: u32, fixed: bool, plan: anytype) !usize { + 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; +} + +/// TENSORFOLD_LFM2_KV_PAGES=0: every stream owns a window of pages, made at open; else they lease the pool's. +fn pagesOn() bool { + const v = std.c.getenv("TENSORFOLD_LFM2_KV_PAGES") 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, p.context, n, p.paged) + be.Lfm.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); +} + +/// 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 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 checkLayers(&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, context, n, paged); + errdefer own.m.deinit(); + own.lb = try be.Lfm.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.Lfm { + 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 through the hooks below. +pub fn snap_ready(ptr: *anyopaque, budget: u64) !u64 { + const x = lfm(ptr); + if (x.snaps != null) return error.SnapshotCacheMade; + x.snaps = try snap.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 * snap.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))); +} + +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, + }; +} + +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); +} + +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..965546d10d --- /dev/null +++ b/zig/src/families/lfm2_moe/cuda_snap.zig @@ -0,0 +1,209 @@ +//! A stream's kept state (conv states, K/V [0, at)) in pages allocated once in the prompt-cache budget, as GLM's. +const std = @import("std"); +const cuda = @import("cuda"); +const kv = @import("core").kv_pool; +const mdl = @import("cuda_model.zig"); +const kern = @import("cuda_kernels.zig"); + +/// 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; +const state_bytes = mdl.slot_elems * 2; + +/// A snapshot page: the unit the budget is cut into. +pub const page: usize = 2 << 20; + +/// Bytes a state at `at` holds: whole snapshot pages. +pub fn bytes(at: usize) u64 { + const n = state_bytes + at * (page_bytes / kern.page); + return (n + page - 1) / page * page; +} + +/// 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; + } +}; + +/// `slot`'s state at `at`: conv states (from `conv`), a copy a run of consecutive whole pages, the last's rows a layer. +fn each(m: *const mdl.Model, slot: usize, conv: u64, at: usize, ctx: anytype, comptime f: fn (@TypeOf(ctx), dev: u64, off: usize, len: usize) anyerror!void) !void { + if (at > m.caps[slot]) return error.PastPages; + try f(ctx, conv, 0, state_bytes); + var off: usize = state_bytes; + const ids = m.ids[slot * m.perSlot() ..]; + const whole = at / kern.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 % kern.page; + if (rows == 0) return; + const last = m.arena + @as(u64, ids[whole]) * page_bytes; + for (0..2 * mdl.attention_layers) |i| { // layer a's keys at block 2a, its values at 2a + 1 + try f(ctx, last + i * kern.page * row_bytes, off, rows * row_bytes); + off += rows * row_bytes; + } +} + +/// 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); +} + +/// `slot`'s state at `at` copied into leased pages behind the chunk that ends there or taps it (conv states at `conv`). +pub fn save(gpa: std.mem.Allocator, p: *Pages, m: *const mdl.Model, s: cuda.Stream, slot: usize, at: u32, conv: 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, conv, at, 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 conv states 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, m.states + slot * state_bytes, k.at, 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 that reads or writes 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 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 each(&m, 1, m.states + state_bytes, at, 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, 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 each(&m, 0, m.states, 2 * kern.page, 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, each(&m, 0, m.states, 2 * kern.page + 1, Count{ .runs = &runs }, Count.add)); +} + +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]); +} + +test "a state's pages: whole snapshot pages, the budget's pages leased and given back" { + try testing.expectEqual(@as(u64, page), bytes(1)); + try testing.expect(bytes(10_000) >= state_bytes + 10_000 * 20 * 1024 and bytes(10_000) % page == 0); +} 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/native/cuda.zig b/zig/src/native/cuda.zig index ab6a6e9443..c808cb8d84 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -18,6 +18,7 @@ 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 Allocator = std.mem.Allocator; pub const budget = @import("cuda_memory.zig"); const lane_backend = @import("cuda_bound.zig"); @@ -28,7 +29,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 }; /// --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..fc949f0032 100644 --- a/zig/tests/cuda/contract_elementwise.zig +++ b/zig/tests/cuda/contract_elementwise.zig @@ -5,6 +5,8 @@ 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 g = core.contract.gates; @@ -15,6 +17,8 @@ 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 at head 64, decode and prompt rows /// 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 +33,13 @@ 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); } pub fn deinit(self: *Set) void { + 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..7b30139e39 --- /dev/null +++ b/zig/tests/cuda/contract_lfm2_moe.zig @@ -0,0 +1,493 @@ +//! 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(control: bool) c.Spec { + 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 = .{ .{}, .{ .size = 256, .step = 64, .min = 64 }, .{} }, .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, + 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.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 != 3) return false; + self.rig.reset(); + self.free(); + const a = self.rig.gpu.gpa; + self.shape = shape; + const r = shape.rows; + const d = shape.n; + self.p = try a.alloc(u16, r * 3 * d); + self.w = try a.alloc(u16, d * 3); + self.st = try a.alloc(u16, 2 * 2 * 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 * 2 * 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(2 * 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 * 6; + } + + 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; + for (0..r) |row| for (0..d) |ch| { + const o: usize = @intCast(self.off[row]); + var bx: [3]u16 = undefined; + for (0..3) |j| { + const back = 2 - 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])) * 2 * d + (2 - back + o) * d + ch]; + } + var acc: f32 = 0; + for (0..3) |j| acc += val(self.w[ch * 3 + j]) * val(bx[j]); + const y = c.toBf16(val(self.p[row * 3 * d + d + ch]) * val(c.toBf16(acc))); + const wants = [_][2]usize{ .{ row * d + ch, y }, .{ r * d + row * 2 * d + ch, bx[1] }, .{ r * d + row * 2 * d + d + ch, bx[2] } }; + for (wants, 0..) |wv, which| { + const got = std.mem.readInt(u16, out[2 * wv[0] ..][0..2], .little); + if (!same(got, @intCast(wv[1]))) 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, wv[1] }); + } + }; + return null; + } +}; + +pub fn commitSpec(control: bool) c.Spec { + 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, + 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.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; + self.snaps = try a.alloc(u16, shape.k * 2 * r * 2 * shape.n); + 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) * 2 * shape.n * 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); + a.add(@as(c_int, @intCast(d))); + a.add(@as(i64, @intCast((r + 1) * 2 * d))); + a.add(@as(i64, @intCast(2 * r * 2 * d))); + a.add(@as(i64, @intCast(2 * d))); + 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) * 2 * 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; + for (0..self.shape.k) |layer| for (0..r + 1) |slot| for (0..2 * d) |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 * 2 * d + @as(usize, @intCast(rowOf(k))) * 2 * d + i]; + const got = std.mem.readInt(u16, out[2 * ((layer * (r + 1) + slot) * 2 * d + 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_limits.zig b/zig/tests/cuda/contract_limits.zig index 2a16afa5b6..1fd1a1cc3f 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" }) |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..ff5fceaa19 --- /dev/null +++ b/zig/tests/cuda/contract_paged_attn.zig @@ -0,0 +1,344 @@ +//! Liquid's paged attention, decode (segs chunk and merge) 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 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 }, +}; + +/// 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); +}