diff --git a/docs/recipes/xpu-arc.md b/docs/recipes/xpu-arc.md index f73fcec648..bc6be8c855 100644 --- a/docs/recipes/xpu-arc.md +++ b/docs/recipes/xpu-arc.md @@ -50,6 +50,23 @@ Intel Arc Pro B70 Graphics (32.5 GB maximum single allocation), Linux 7.0.0 (`xe Nemotron MLX 4-bit checkpoint (`TensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit`, revision `d9d758fb83953437f7263256b0d96157e2a348b8`, 729 tensors, 17.27 GiB of weights). The commit under test is the head of the pull request. +LFM3-1.7B (`LiquidAI/LFM3-1.7B-MLX-4bit`, MLX affine 4-bit, group 64) on an Intel Iris Xe laptop GPU (Gen12LP, `tgllp`, 96 EUs, no +DPAS) under WSL2 (Ubuntu 24.04, 8 GB given to WSL), Ubuntu's compute runtime 23.43.27642 and Level Zero loader 1.16, Zig 0.17.0, +`zig build xpu-native -Dxpu-device=tgllp` (DPAS off, so the server is built without the Qwen3.8-27B runner): + +| Setup | Decode, c=1, `--no-drafts` | Decode, c=1, copy drafts | Decode, c=4, aggregate | TTFT, 1,024 tokens | TTFT, 8,192 tokens | Memory (`serve --context 16384`) | +| --- | ---: | ---: | ---: | ---: | ---: | --- | +| Iris Xe, LFM3-1.7B MLX 4-bit | 15.6 to 16.8 tok/s | 15.9 (chat) to 26.3 (prose) tok/s | 25 to 32 tok/s | 11.7 to 12.2 s | 116.3 to 117.7 s | 1.33 GB device allocation, 448 MB host resident | + +Decode, c=1: two short chats (160 tokens), a 1,001-token prose prompt and a 601-token code prompt (128 tokens), greedy, three runs +each, with copy drafts (the default) and with `--no-drafts`; c=4: four concurrent 128-token completions of distinct 32-token prompts, +six runs a mode, with or without drafts; TTFT: prompts of numbered item sentences, one token generated; three runs each, through +`/v1/completions`. Copy-draft gains depend on how repetitive the text is: prose and code that repeat their own context decode at 26 +and 22 tok/s, the chats, with few copies to draft, at the speed without drafts. A `/v1/chat/completions` request with the +checkpoint's chat template answers (thinking on, `finish_reason` stop). +With `--parallel N` each stream has its own sequence (about 196 MiB of KV at 16,384 tokens); a second chat turn over an +8,037-token document resumes from the prompt cache in 0.80 s instead of 122.6 s, with the same reply. + ## Measurements Decode, 128 greedy tokens from a 5-token prompt: 5.72 ms/token, 175 tok/s (`tensorfold-xpu run`, five runs on different diff --git a/tools/lfm3/REF64.md b/tools/lfm3/REF64.md new file mode 100644 index 0000000000..1f6d4730ed --- /dev/null +++ b/tools/lfm3/REF64.md @@ -0,0 +1,94 @@ +# Float64 reference receipt: LFM3-1.7B NVFP4 on CUDA + +- Checkpoint: [LiquidAI/LFM3-1.7B-NVFP4](https://huggingface.co/LiquidAI/LFM3-1.7B-NVFP4) at revision + `b30025bbbbcbf62f6c110b40d51d36af6579f6a2`. +- Engine: `9265a0f5f`, which is the PR head `08b338c7c` before its rebase (range-diff: the same patches except one + test-only line in `contract_limits.zig`). Served on a GB10 (sm_121), greedy, drafts off, + `--parallel 8 --context 32768`. The `collect` step was repeated on 2026-10-28 and returned the same 957 tokens. +- Reference: `ref64.py` on the CPU in float64 (Python 3.12, torch 2.14, safetensors 0.8, tokenizers 0.23), 2026-10-28. + +Reproduce (the reference pass takes about a minute and 20 GB of host memory): + +``` +zig-out/native/bin/tensorfold-native serve MODEL_DIR --port 8080 --parallel 8 --context 32768 +python -I tools/lfm3/ref64.py collect http://127.0.0.1:8080 MODEL_DIR engine.json +python -I tools/lfm3/ref64.py MODEL_DIR engine.json --bf16 +``` + +Rows. A prompt row is the engine's one-token greedy pick after the prefix [0, t] of a prompt, one request per prefix. +A reply row is a token of the engine's 32-token greedy reply, teacher-forced on that reply. `t` is the logits row after +token t, so prompt row P - 1 and reply token 0 are the same position. The prompts (`prompts()` in the script) are +capitals (20 tokens), code (40), a chat turn (19) and items (750, a list of 60 items and a question; it crosses the +256-key sliding window). + +The float64 pass dequantizes the NVFP4 weights exactly. The `--bf16` pass adds the engine's bf16 storage points (see +the script's docstring). It shows which side of a miss rounding alone lands on; it is not bit-exact. + +## Result + +- f64: prompt rows 825/829, reply rows 127/128 +- bf16: prompt rows 827/829, reply rows 126/128 + +Float64 top-1 misses (margin = f64 logit of its top minus f64 logit of the engine's token; bf16 margin = the same difference in the bf16-points pass, < 0 means it prefers the engine's token): + +| case | row | t | engine | float64 | f64 margin | bf16 margin | bf16 picks | +|---|---|---|---|---|---|---|---| +| capitals | reply | 24 | 11097 " Spain" | 3888 " Japan" | 0.1176 | +0.1250 (pre-round +0.1802) | f64's | +| chat | prompt | 12 | 593 " per" | 22 "." | 0.1530 | -0.0625 (pre-round -0.0671) | engine's | +| items | prompt | 5 | 1008 " first" | 1579 " number" | 0.0996 | +0.0000 (pre-round +0.1009) | engine's | +| items | prompt | 9 | 29 "5" | 654 "10" | 0.0033 | +0.0000 (pre-round +0.0093) | engine's | +| items | prompt | 180 | 28 "4" | 30 "6" | 0.0674 | +0.0000 (pre-round +0.0076) | engine's | + +bf16-points misses where float64 agrees with the engine (margin = bf16 logit of its pick minus the engine's token; f64 margin = how far float64 prefers the engine's token): + +| case | row | t | engine | bf16 pick | bf16 margin | f64 margin | +|---|---|---|---|---|---|---| +| items | prompt | 4 | 440 " The" | 229 " " | 0.0000 (pre-round -0.0985) | 0.0881 | +| items | prompt | 631 | 32 "8" | 25 "1" | 0.0000 (pre-round -0.0467) | 0.0636 | +| items | reply | 772 | 41501 " Item" | 1672 " So" | 0.0000 (pre-round -0.0814) | 0.1797 | + +## FP8 and MLX-4bit (2026-10-29) + +- Checkpoints: [LiquidAI/LFM3-1.7B-FP8](https://huggingface.co/LiquidAI/LFM3-1.7B-FP8) at + `243955da95b1cd68026f293b29de4b4c13f3638f` (block FP8, dequantized exactly: e4m3 times its block's bf16 scale) and + [LiquidAI/LFM3-1.7B-MLX-4bit](https://huggingface.co/LiquidAI/LFM3-1.7B-MLX-4bit) at + `d0ca73dcfa70ce18733754a492cdc1fb0123262d` (MLX affine, dequantized exactly). +- Served the same way with `--no-drafts`; the same prompts and rows. Only the float64 pass was run: the `--bf16` + pass models the NVFP4 engine's rounding points. + +FP8: prompt rows 824/829, reply rows 128/128. + +| case | row | t | engine | float64 | f64 margin | +|---|---|---|---|---|---| +| items | prompt | 131 | 30 "6" | 1325 "24" | 0.0348 | +| items | prompt | 582 | 654 "10" | 1035 "14" | 0.0756 | +| items | prompt | 619 | 25 "1" | 29 "5" | 0.0066 | +| items | prompt | 631 | 27 "3" | 32 "8" | 0.0290 | +| items | prompt | 693 | 9402 " stone" | 17937 " stones" | 0.0773 | + +MLX-4bit: prompt rows 825/829, reply rows 128/128. + +| case | row | t | engine | float64 | f64 margin | +|---|---|---|---|---|---| +| items | prompt | 10 | 6802 " items" | 93266 " marbles" | 0.0466 | +| items | prompt | 61 | 41501 " Item" | 29155 " Question" | 0.0042 | +| items | prompt | 595 | 27 "3" | 1035 "14" | 0.0300 | +| items | prompt | 748 | 34 ":" | 10494 " holds" | 0.0848 | + +## Without torch: the numpy path and the XPU's bf16 points + +`ref64.py` runs on numpy when torch is not installed (or with `--numpy`): MLX affine and bf16 weights are dequantized +exactly as on the torch path; NVFP4 and FP8 exports are refused by name and need torch. `--bf16-points` replaces the +`--bf16` pass with the Intel GPU (XPU) engine's storage points: the embedding rows and RoPE cos/sin rounded to bf16 too, +attention P and the logits never. `dump` writes the float64 hidden state after every layer for the XPU forward's +layer-by-layer check, and `check` replays an engine's greedy replies teacher-forced: + +``` +python -I tools/lfm3/ref64.py dump MODEL_DIR OUT_DIR +python -I tools/lfm3/ref64.py check MODEL_DIR tokens.json # {"name": {"prompt": [...], "reply": [...]}} +python -I tools/lfm3/ref64.py MODEL_DIR engine.json --bf16-points +``` + +On LiquidAI/LFM3-1.7B-MLX-4bit (Python 3.12, numpy 2.5, no torch; 3 min 14 s and 4.1 GB for `dump`): the three +prompts' bf16-points gap reaches 0.041 (capitals), 0.025 (chat) and 0.111 (items, worst row of any layer), and the +XPU engine's 64-token greedy replies agree with the float64 top-1 at 192 of 192 positions. diff --git a/tools/lfm3/ref64.py b/tools/lfm3/ref64.py new file mode 100644 index 0000000000..9f78aa46b2 --- /dev/null +++ b/tools/lfm3/ref64.py @@ -0,0 +1,521 @@ +#!/usr/bin/env python3 +"""Float64 CPU reference of the LFM3 text decoder (LFM3-1.7B), checked against the CUDA engine's greedy tokens. + +The stored weights are dequantized exactly and every op runs in float64: + nvfp4 ModelOpt NVFP4 (U8 two E2M1 a byte, low nibble first; e4m3 scale per 16 values; f32 weight_scale_2) + mlx MLX affine (bits and group size from config.json "quantization", per-module overrides) + fp8 compressed-tensors block FP8 (e4m3; a bf16 scale per 128 x 128 block) + bf16 plain bf16 weights +The math follows modeling_lfm3.py: sandwich norms (h += post_norm(mixer(norm(h)))), RMSNorm w * x * rsqrt(ms + eps); +conv y = C * causal_depthwise_conv(B * x) with the last tap on the current row; attention with per-head q/k RMSNorm, +RoPE (rotate_half; full layers partial 0.5 in the zero-tail form: inv_freq[i] = theta^(-2i/128) for i < 32, 0 after; +sliding layers all 64 frequencies), GQA, a 256-key sliding window, a per-head output gate sigmoid(gate_proj(x)); the +head is tied to embed_tokens. + +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 NVFP4 weight-only CUDA engine's bf16 storage points (float64 inside every op, +round-to-nearest-even to bf16 where the kernels store bf16): every RMSNorm (bf16(w * bf16(h * inv))), residual, +projection outputs, conv B*x / acc / C*y, q/k norm and RoPE, the gate, attention output, SwiGLU and the logits. Weights +are exact on both paths. Attention P is rounded to bf16 before PV on prompt rows and sliding-window decode rows and kept +unrounded on full-attention decode rows, taken against the row's final max (the kernels' running max 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. + +Without torch (or with --numpy) the forward runs on numpy: MLX affine and bf16 weights are dequantized exactly; NVFP4 +and FP8 stay torch-only and are refused by name. --bf16-points rounds where the XPU kernels store bf16 instead of the +CUDA engine: the embedding rows and RoPE cos and sin as well, attention P and the logits never (numpy path). + +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 | --bf16-points] [--numpy] + prints the miss table (markdown) and the agreement totals + ref64.py dump MODEL_DIR OUT_DIR capitals, chat, items: ids, the hidden state after every layer and + the final norm (f32), float64 top-1 per row, and the bf16-points + pass's per-layer gap (meta.json); numpy path + ref64.py check MODEL_DIR TOKENS_JSON teacher-forced float64 top-1 of {name: {prompt, reply}}; numpy path +""" +import json +import math +import os +import re +import sys +import urllib.request + +import struct + +import numpy as np + +try: + import torch + from safetensors import safe_open +except ImportError: # the numpy path: MLX affine and bf16 weights + torch = None + +if torch is not None: + torch.set_default_dtype(torch.float64) + F64 = torch.float64 + E2M1 = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6, -0.0, -0.5, -1, -1.5, -2, -3, -4, -6], dtype=F64) + + +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 Berlin. " + "The capital of Italy is"), + ("code", "def fibonacci(n):\n if n < 2:\n return n\n" + " return fibonacci(n - 1) + fibonacci(n - 2)\n\n# Rewrite iteratively:\ndef fibonacci(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(60): + ids = [bos] + tok.encode(text).ids + picks = [complete(url, ids[: t + 1], 1)[0] for t in range(len(ids))] + cases.append({"name": name, "prompt": ids, "prompt_picks": picks, "reply": complete(url, ids, reply_len)}) + print(f"{name}: {len(ids)} prompt rows, {len(cases[-1]['reply'])} reply tokens", file=sys.stderr) + json.dump({"cases": cases}, open(out_path, "w")) + + +def bf16(x): + """Nearest bf16 (8 significant bits, ties to even), kept as float64: exact, no float32 double rounding.""" + m, e = torch.frexp(x) + return torch.ldexp(torch.round(m * 256.0), (e - 8).to(torch.int32)).to(F64) + + +def same(x): + return x + + +def rot(t, cos, sin): + hd = t.shape[-1] + return t * cos + torch.cat([-t[..., hd // 2:], t[..., : hd // 2]], -1) * sin + + +def strip(k): + return re.sub(r"^(model\.language_model\.|language_model\.|model\.)+", "", k) + + +class Model: + def __init__(self, d): + self.d = d + root = json.load(open(f"{d}/config.json")) + self.cfg = root.get("text_config", root) + idx = f"{d}/model.safetensors.index.json" + if os.path.exists(idx): + self.map = json.load(open(idx))["weight_map"] + else: + with safe_open(f"{d}/model.safetensors", "pt") as f: + self.map = {k: "model.safetensors" for k in f.keys()} + self.files, self.cache = {}, {} + self.prefix = "language_model.model." if "language_model.model.embed_tokens.weight" in self.map else "model." + self.mlx = root.get("quantization") if root.get("quantization", {}).get("mode", "affine") == "affine" else None + + def raw(self, k): + f = self.map[k] + if f not in self.files: + self.files[f] = safe_open(f"{self.d}/{f}", "pt") + return self.files[f].get_tensor(k) + + def vec(self, name): + return self.raw(self.prefix + name).to(F64) + + def lin(self, name, rows=None): + """A Linear weight [out, in] in float64, dequantized exactly by its stored format.""" + p = self.prefix + name + w = self.raw(p + ".weight") + if rows is not None: + w = w[rows] + if self.mlx and p + ".scales" in self.map: + s, b = self.raw(p + ".scales"), self.raw(p + ".biases") + if rows is not None: + s, b = s[rows], b[rows] + bits = next((v["bits"] for k, v in self.mlx.items() if isinstance(v, dict) and strip(k) == strip(p)), + self.mlx["bits"]) + by = w.contiguous().view(torch.uint8).to(torch.int64) + nb = by.shape[-1] * 8 + bit = ((by.unsqueeze(-1) >> torch.arange(8)) & 1).reshape(*by.shape[:-1], nb) + q = (bit[..., : nb // bits * bits].reshape(*by.shape[:-1], -1, bits) << torch.arange(bits)).sum(-1) + q = q.to(F64).reshape(*q.shape[:-1], s.shape[-1], -1) + return (q * s.to(F64)[..., None] + b.to(F64)[..., None]).flatten(-2) + if p + ".weight_scale_2" in self.map: + sc = self.raw(p + ".weight_scale") + if rows is not None: + sc = sc[rows] + n = w.shape[0] + codes = torch.stack([w & 0xF, w >> 4], dim=-1).reshape(n, -1).long() + v = E2M1[codes].view(n, -1, 16) * sc.to(F64)[..., None] + return (v * self.raw(p + ".weight_scale_2").to(F64)).reshape(n, -1) + if w.dtype == torch.float8_e4m3fn: + sc = self.raw(p + ".weight_scale").to(F64).repeat_interleave(128, 0).repeat_interleave(128, 1) + return w.to(F64) * (sc if rows is None else sc[rows]) + return w.to(F64) + + def W(self, name): + if name not in self.cache: + self.cache[name] = self.lin(name) + return self.cache[name] + + def forward(self, ids, P, bf): + """Logits [T, V] (float64, not rounded). bf = same: plain float64; bf = bf16: the engine's storage points with + rows [0, P) on the prompt path and the rest on the decode path.""" + cfg, T = self.cfg, len(ids) + D, H, HK, hd = cfg["hidden_size"], cfg["num_attention_heads"], cfg["num_key_value_heads"], cfg["head_dim"] + taps, vec, W = cfg["conv_L_cache"], self.vec, self.W + + def rms(h, w): # the engine's form: bf16(w * bf16(h * inv)), inv from the bf16 row + return bf(w * bf(h * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + cfg["norm_eps"]))) + + tabs = {} + for kind, rp in cfg["rope_parameters"].items(): + angles = int(rp.get("partial_rotary_factor", 1.0) * hd) // 2 + inv = 1.0 / rp["rope_theta"] ** (torch.arange(0, 2 * angles, 2, dtype=F64) / hd) # the full head dim's + inv = torch.cat([inv, torch.zeros(hd // 2 - angles)]) + ang = torch.arange(T, dtype=F64)[:, None] * inv[None] + tabs[kind] = torch.cat([ang.cos()] * 2, -1), torch.cat([ang.sin()] * 2, -1) + i, j = torch.arange(T)[:, None], torch.arange(T)[None] + masks = {"full_attention": j > i, "sliding_attention": (j > i) | (i - j >= cfg["sliding_window"])} + h = self.lin("embed_tokens", rows=torch.tensor(ids)) + for L, kind in enumerate(cfg["layer_types"]): + p = f"layers.{L}." + x = rms(h, vec(p + "operator_norm.weight")) + if kind == "conv": + B, C, xx = bf(x @ W(p + "conv.in_proj").T).chunk(3, -1) + bx = bf(B * xx) + w = vec(p + "conv.conv.weight").reshape(D, taps) + pad = torch.cat([torch.zeros(taps - 1, D), bx]) + acc = torch.zeros(T, D) + for k in range(taps): # oldest tap first + acc = acc + pad[k:k + T] * w[:, k] + o = bf(bf(C * bf(acc)) @ W(p + "conv.out_proj").T) + else: + q, k, v = (bf(x @ W(p + f"self_attn.{n}_proj").T) for n in "qkv") + g = bf(x @ W(p + "self_attn.gate_proj").T) # [T, H] + q = rms(q.reshape(T, H, hd), vec(p + "self_attn.q_norm.weight")).transpose(0, 1) # [H, T, hd] + k = rms(k.reshape(T, HK, hd), vec(p + "self_attn.k_norm.weight")).transpose(0, 1) + q, k = bf(rot(q, *tabs[kind])), bf(rot(k, *tabs[kind])) + k = k.repeat_interleave(H // HK, 0) + v = v.reshape(T, HK, hd).transpose(0, 1).repeat_interleave(H // HK, 0) + s = ((q @ k.transpose(1, 2)) / math.sqrt(hd)).masked_fill(masks[kind], -math.inf) + e = torch.exp(s - s.max(-1, keepdim=True).values) + pv = bf(e) + if kind == "full_attention": + pv = torch.cat([pv[:, :P], e[:, P:]], 1) # full-attention decode keeps P unrounded + a = bf((pv @ v) / e.sum(-1, keepdim=True)).transpose(0, 1) # [T, H, hd] + a = bf(a * bf(torch.sigmoid(g))[..., None]) + o = bf(a.reshape(T, H * hd) @ W(p + "self_attn.o_proj").T) + h = bf(h + rms(o, vec(p + "post_operator_norm.weight"))) + x = rms(h, vec(p + "ffn_norm.weight")) + f = p + "feed_forward." + act = bf(bf(torch.nn.functional.silu(bf(x @ W(f + "w1").T))) * bf(x @ W(f + "w3").T)) + h = bf(h + rms(bf(act @ W(f + "w2").T), vec(p + "post_ffn_norm.weight"))) + x = rms(h, vec("embedding_norm.weight")) + V = cfg["vocab_size"] + return torch.cat([x @ self.lin("embed_tokens", rows=torch.arange(a, min(a + 16000, V))).T + for a in range(0, V, 16000)], -1) + + +def torch_logits(self, ids, P, points=None): + """Logits [T, V] as a numpy array: float64, or the CUDA engine's bf16 points (points == "engine").""" + return self.forward(ids, P, bf16 if points else same).numpy() + + +Model.logits = torch_logits + + +def greedy(lg): + """argmax in (value desc, id asc) order, as the engine picks over its bf16 logits.""" + return int(np.flatnonzero(lg == lg.max())[0]) + + +def npbf16(x): + """Nearest bf16 (ties to even) of float64 values, kept as float64.""" + m, e = np.frexp(x) + return np.ldexp(np.round(m * 256.0), e - 8) + + +class NpModel: + """The float64 forward on numpy: MLX affine and bf16 weights read from the safetensors bytes, exactly.""" + + def __init__(self, d): + root = json.load(open(f"{d}/config.json")) + self.cfg = root.get("text_config", root) + idx = f"{d}/model.safetensors.index.json" + shards = sorted(set(json.load(open(idx))["weight_map"].values())) if os.path.exists(idx) else ["model.safetensors"] + self.hdr = {} + for s in shards: + with open(f"{d}/{s}", "rb") as f: + n = struct.unpack("> (np.arange(8, dtype=np.uint32) * 4)) & 15 + else: + bit = np.unpackbits(np.ascontiguousarray(w).view(np.uint8), axis=-1, bitorder="little") + nb = bit.shape[-1] // bits * bits + q = bit[..., :nb].reshape(n, -1, bits).astype(np.int64) @ (1 << np.arange(bits)) + q = q.reshape(n, s.shape[-1], -1).astype(np.float64) + return (q * s[..., None] + b[..., None]).reshape(n, -1) + + def logits(self, ids, P, points=None, dump=None): + """Logits [T, V] in float64; points None, "engine" (--bf16) or "xpu" (--bf16-points); dump(L, h) per layer.""" + cfg, T = self.cfg, len(ids) + D, H, HK, hd = cfg["hidden_size"], cfg["num_attention_heads"], cfg["num_key_value_heads"], cfg["head_dim"] + bf = npbf16 if points else same + xpu = points == "xpu" + + def rms(h, w): + return bf(w * bf(h / np.sqrt((h * h).mean(-1, keepdims=True) + cfg["norm_eps"]))) + + def mm(x, name): + return x @ self.lin(name).T + + tabs = {} + for kind, rp in cfg["rope_parameters"].items(): + angles = int(rp.get("partial_rotary_factor", 1.0) * hd) // 2 + inv = 1.0 / rp["rope_theta"] ** (np.arange(0, 2 * angles, 2, dtype=np.float64) / hd) + inv = np.concatenate([inv, np.zeros(hd // 2 - angles)]) + ang = np.arange(T, dtype=np.float64)[:, None] * inv[None] + cos, sin = np.concatenate([np.cos(ang)] * 2, -1), np.concatenate([np.sin(ang)] * 2, -1) + tabs[kind] = (bf(cos), bf(sin)) if xpu else (cos, sin) + i, j = np.arange(T)[:, None], np.arange(T)[None] + masks = {"full_attention": j > i, "sliding_attention": (j > i) | (i - j >= cfg["sliding_window"])} + h = self.lin("embed_tokens", rows=np.array(ids)) + if xpu: + h = bf(h) + for L, kind in enumerate(cfg["layer_types"]): + p = f"layers.{L}." + x = rms(h, self.vec(p + "operator_norm.weight")) + if kind == "conv": + B, C, xx = np.split(bf(mm(x, p + "conv.in_proj")), 3, -1) + bx = bf(B * xx) + w = self.vec(p + "conv.conv.weight").reshape(D, -1) + taps = w.shape[1] + pad = np.concatenate([np.zeros((taps - 1, D)), bx]) + acc = np.zeros((T, D)) + for k in range(taps): + acc = acc + pad[k:k + T] * w[:, k] + o = bf(mm(bf(C * bf(acc)), p + "conv.out_proj")) + else: + q, k, v = (bf(mm(x, p + f"self_attn.{n}_proj")) for n in "qkv") + g = bf(mm(x, p + "self_attn.gate_proj")) + q = rms(q.reshape(T, H, hd), self.vec(p + "self_attn.q_norm.weight")).transpose(1, 0, 2) + k = rms(k.reshape(T, HK, hd), self.vec(p + "self_attn.k_norm.weight")).transpose(1, 0, 2) + cos, sin = tabs[kind] + q = bf(q * cos + np.concatenate([-q[..., hd // 2:], q[..., : hd // 2]], -1) * sin) + k = bf(k * cos + np.concatenate([-k[..., hd // 2:], k[..., : hd // 2]], -1) * sin) + k = np.repeat(k, H // HK, 0) + v = np.repeat(v.reshape(T, HK, hd).transpose(1, 0, 2), H // HK, 0) + s = (q @ k.transpose(0, 2, 1)) / math.sqrt(hd) + s[:, masks[kind]] = -np.inf + e = np.exp(s - s.max(-1, keepdims=True)) + pv = e + if points == "engine": + pv = bf(e) + if kind == "full_attention": + pv[:, P:] = e[:, P:] # full-attention decode keeps P unrounded + a = bf((pv @ v) / e.sum(-1, keepdims=True)).transpose(1, 0, 2) + a = bf(a * bf(1.0 / (1.0 + np.exp(-g)))[..., None]) + o = bf(mm(a.reshape(T, H * hd), p + "self_attn.o_proj")) + h = bf(h + rms(o, self.vec(p + "post_operator_norm.weight"))) + x = rms(h, self.vec(p + "ffn_norm.weight")) + f = p + "feed_forward." + a1 = bf(mm(x, f + "w1")) + act = bf(bf(a1 / (1.0 + np.exp(-a1))) * bf(mm(x, f + "w3"))) + h = bf(h + rms(bf(mm(act, f + "w2")), self.vec(p + "post_ffn_norm.weight"))) + if dump: + dump(L, h) + x = rms(h, self.vec("embedding_norm.weight")) + if dump: + dump(-1, x) + V = cfg["vocab_size"] + return np.concatenate([x @ self.lin("embed_tokens", rows=np.arange(a, min(a + 8000, V))).T + for a in range(0, V, 8000)], -1) + + +def rel_rows(a, b): + """Per-row ||a - b|| / ||b||: mean and max.""" + r = np.linalg.norm(a - b, axis=-1) / np.linalg.norm(b, axis=-1) + return float(r.mean()), float(r.max()) + + +def dump(model_dir, out): + """capitals, chat, items (BOS, no added specials): float64 hidden states per layer, top-1, bf16-points gaps.""" + from tokenizers import Tokenizer + tok = Tokenizer.from_file(f"{model_dir}/tokenizer.json") + m = NpModel(model_dir) + bos = m.cfg["bos_token_id"] + summary = {} + for name, text in [p for p in prompts(60) if p[0] != "code"]: + ids = [bos] + tok.encode(text, add_special_tokens=False).ids + d = f"{out}/{name}" + os.makedirs(d, exist_ok=True) + f64h, xph = {}, {} + + def keep(L, h): + f64h[L] = h.copy() + h.astype(np.float32).tofile(f"{d}/h{L:02d}.f32" if L >= 0 else f"{d}/final.f32") + + lf = m.logits(ids, len(ids), None, keep) + lb = m.logits(ids, len(ids), "xpu", lambda L, h: xph.__setitem__(L, h.copy())) + gap = {L: rel_rows(xph[L], f64h[L]) for L in f64h} + top = lf.argmax(-1) + np.asarray(top, dtype=np.int32).tofile(f"{d}/top1.i32") + lf[-1].astype(np.float32).tofile(f"{d}/logits_last.f32") + agree = int((lb.argmax(-1) == top).sum()) + json.dump({"ids": ids, "T": len(ids), "top1": top.tolist(), "bf16_gap": {str(L): v for L, v in gap.items()}, + "bf16_top1_agree": agree}, open(f"{d}/meta.json", "w")) + summary[name] = {"T": len(ids), "bf16_gap_max": max(v[1] for v in gap.values()), "next": int(top[-1]), + "bf16_top1_agree": f"{agree}/{len(ids)}"} + print(name, json.dumps(summary[name]), file=sys.stderr, flush=True) + json.dump(summary, open(f"{out}/summary.json", "w"), indent=1) + + +def check(model_dir, path): + """Teacher-forced: the float64 top-1 after each prefix of prompt + reply against the reply's own next token.""" + m = NpModel(model_dir) + res = {} + for name, c in json.load(open(path)).items(): + P, reply = len(c["prompt"]), c["reply"] + lf = m.logits(c["prompt"] + reply[:-1], P)[P - 1:] + picks = lf.argmax(-1) + miss = [{"j": j, "engine": int(reply[j]), "f64": int(picks[j]), + "margin": float(lf[j].max() - lf[j, reply[j]])} for j in range(len(reply)) if picks[j] != reply[j]] + res[name] = {"agree": f"{len(reply) - len(miss)}/{len(reply)}", "misses": miss} + print(name, json.dumps(res[name]), file=sys.stderr, flush=True) + return res + + +def main(model_dir, engine_path, points, numpy_path): + mdl = NpModel(model_dir) if numpy_path or torch is None or points == "xpu" else Model(model_dir) + with_bf16 = points is not None + 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.logits(ids, P) + lb_pre = mdl.logits(ids, P, points) if with_bf16 else None + lb = (npbf16(lb_pre) if points == "engine" else 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__": + a = sys.argv[1:] + flags = a[2:] + if len(a) == 4 and a[0] == "collect": + collect(*a[1:]) + elif len(a) == 3 and a[0] == "dump": + dump(a[1], a[2]) + elif len(a) == 3 and a[0] == "check": + json.dump(check(a[1], a[2]), sys.stdout, indent=1) + elif len(a) >= 2 and set(flags) <= {"--bf16", "--bf16-points", "--numpy"} and len(set(flags)) == len(flags) \ + and not {"--bf16", "--bf16-points"} <= set(flags): + pts = "engine" if "--bf16" in flags else "xpu" if "--bf16-points" in flags else None + main(a[0], a[1], pts, "--numpy" in flags) + else: + sys.exit(__doc__) diff --git a/zig/build/xpu.zig b/zig/build/xpu.zig index e39bea9b48..10bf54badc 100644 --- a/zig/build/xpu.zig +++ b/zig/build/xpu.zig @@ -139,11 +139,20 @@ pub fn steps(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.buil const lanes = b.createModule(.{ .root_source_file = b.path("zig/src/core/lanes/lanes.zig"), .target = target, .optimize = optimize, .link_libc = true }); const nem = nemotronModules(b, target, optimize, ocloc, device, dpas, ze_mod, lanes); const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = lanes }} }); - const gemma4 = b.createModule(.{ .root_source_file = b.path("zig/src/families/gemma4/xpu_engine.zig"), .target = target, .optimize = optimize, .link_libc = true }); - gemma4.addImport("xpu_gemma4_kernels", embedded_g4); - gemma4.addImport("core", core); - gemma4.addImport("ze", ze_mod); - const engines = b.createModule(.{ .root_source_file = b.path("zig/src/native/xpu.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "engine_api", .module = api }, .{ .name = "xpu_qwen27", .module = qwen27 }, .{ .name = "lanes", .module = lanes }, .{ .name = "xpu", .module = nem.xpu }, .{ .name = "nemotron_xpu", .module = nem.nemotron }, .{ .name = "gemma4_xpu", .module = gemma4 } } }); + const engines = b.createModule(.{ .root_source_file = b.path("zig/src/native/xpu.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "engine_api", .module = api }, .{ .name = "lanes", .module = lanes }, .{ .name = "xpu", .module = nem.xpu }, .{ .name = "nemotron_xpu", .module = nem.nemotron }, .{ .name = "lfm3_xpu", .module = lfm3Module(b, target, optimize, nem.xpu) } } }); + // the Qwen3.8-27B runner and Gemma 4 E4B (DPAS, nGEN) only with DPAS: a Gen12LP server serves Nemotron and LFM3 + const engine_options = b.addOptions(); + engine_options.addOption(bool, "qwen27", dpas); + engine_options.addOption(bool, "gemma4", dpas); + engines.addOptions("xpu_engine_options", engine_options); + if (dpas) { + engines.addImport("xpu_qwen27", qwen27); + const gemma4 = b.createModule(.{ .root_source_file = b.path("zig/src/families/gemma4/xpu_engine.zig"), .target = target, .optimize = optimize, .link_libc = true }); + gemma4.addImport("xpu_gemma4_kernels", embedded_g4); + gemma4.addImport("core", core); + gemma4.addImport("ze", ze_mod); + engines.addImport("gemma4_xpu", gemma4); + } const tokenizer = b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true }); const template = b.createModule(.{ .root_source_file = b.path("zig/src/core/template/template.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true }); const checkpoint_cli = b.createModule(.{ .root_source_file = b.path("zig/src/cli/cli.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true, .imports = &.{ .{ .name = "native_engines", .module = engines }, exlr8_build.cliImport(b, target, .ReleaseSafe, core), .{ .name = "core", .module = core } } }); @@ -181,7 +190,10 @@ const kernels = [_]Kernel{ .{ .name = "nem_rows_r4", .src = "nem_rows", .options = cl3 ++ " -DMAXR=4" }, .{ .name = "nem_rows_r8", .src = "nem_rows", .options = cl3 ++ " -DMAXR=8" }, .{ .name = "nem_prow", .options = cl3 }, + .{ .name = "nem_prow_gen12", .src = "nem_prow", .options = cl3 ++ " -DPO=4" }, .{ .name = "nem_attn_dec", .options = cl3 }, + .{ .name = "lfm3", .options = cl3 }, + .{ .name = "lfm3_mv", .options = cl3 }, }; /// The runtime module; `with_kernels` false builds it host-only (empty images). @@ -245,9 +257,16 @@ fn hasDpas(device: []const u8) bool { return true; } +/// LFM3's forward and engine on the runtime module `xpu`. +fn lfm3Module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, xpu: *std.Build.Module) *std.Build.Module { + const m = b.createModule(.{ .root_source_file = b.path("zig/src/families/lfm3/xpu.zig"), .target = target, .optimize = optimize, .link_libc = true }); + m.addImport("xpu", xpu); + return m; +} + /// The standalone op and model tests, each its own program `tf-xpu--test` (step `xpu-tests`). const test_programs = [_][]const u8{ - "panic", "nem_rows", "nem_prefill", "nem_gen", "nem_mtp", "nem_kl", "nem_norm_bench", "nem_mv_bench", "nem_slot", "nem_format", "nem_seam", "nem_attn_dec", "nem_fp64", "memory", + "panic", "nem_rows", "nem_prefill", "nem_gen", "nem_mtp", "nem_kl", "nem_norm_bench", "nem_mv_bench", "nem_slot", "nem_format", "nem_seam", "nem_attn_dec", "nem_fp64", "memory", "lfm3_conv", "lfm3_rope", "lfm3_attn", "lfm3_norm", "lfm3_fwd", "lfm3_bench", "lfm3_rows", "lfm3_streams", "lfm3_chunks", "lfm3_pgemm", }; /// `-Dxpu` (Linux): Nemotron's SPIR-V images, `tensorfold-xpu` and `tf-xpu-test`; needs ocloc, not nvcc or CUDA. @@ -264,11 +283,13 @@ fn nemotronSteps(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std. runner.addImport("nemotron_xpu", mods.nemotron); b.installArtifact(b.addExecutable(.{ .name = "tf-xpu-test", .root_module = runner })); const tests_step = b.step("xpu-tests", "Build the standalone Intel GPU op and model tests (tf-xpu--test)"); + const lfm3 = lfm3Module(b, target, optimize, xpu); for (test_programs) |name| { const mod = b.createModule(.{ .root_source_file = b.path(b.fmt("zig/tests/xpu/{s}.zig", .{name})), .target = target, .optimize = optimize, .link_libc = true }); mod.addImport("xpu", xpu); mod.addImport("core", mods.core); mod.addImport("nemotron_xpu", mods.nemotron); + mod.addImport("lfm3_xpu", lfm3); tests_step.dependOn(&b.addInstallArtifact(b.addExecutable(.{ .name = b.fmt("tf-xpu-{s}-test", .{name}), .root_module = mod }), .{}).step); } } @@ -289,6 +310,11 @@ pub fn hostTests(b: *std.Build, step: *std.Build.Step) void { gemma_native.addImport("ze", ze); gemma_native.addImport("xpu_gemma4_kernels", empty); sources.addImport("gemma4_xpu", gemma_native); + // the host test checks every family's admission, so it builds with both DPAS engines + const source_options = b.addOptions(); + source_options.addOption(bool, "qwen27", true); + source_options.addOption(bool, "gemma4", true); + sources.addOptions("xpu_engine_options", source_options); const source_tests = b.addRunArtifact(b.addTest(.{ .root_module = sources, .filters = &.{"draft sources:"} })); b.step("test-slide-xpu-host", "Unsupported XPU learning refuses before device or checkpoint access").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = sources, .filters = &.{"slide admission:"} })).step); b.step("test-draft-xpu-host", "XPU draft source metadata without devices or kernel compilation").dependOn(&source_tests.step); diff --git a/zig/kernels/xpu/lfm3.cl b/zig/kernels/xpu/lfm3.cl new file mode 100644 index 0000000000..e5144072ea --- /dev/null +++ b/zig/kernels/xpu/lfm3.cl @@ -0,0 +1,257 @@ +// LFM3 ops in plain OpenCL C (no DPAS): bf16 rows, fp32 math, rounded to bf16 where modeling_lfm3.py rounds. +#pragma OPENCL FP_CONTRACT OFF +#define HD 128 // head dim +#ifndef GQ +#define GQ 4 // query heads per KV head (16 / 4) +#endif +#define NPI 16 // values an item of 256 holds in the row ops: n <= 4096 + +inline float bf(ushort v) { return as_float((uint)v << 16); } +inline ushort to_bf(float f) { + uint u = as_uint(f); + u += 0x7fffu + ((u >> 16) & 1u); + return (ushort)(u >> 16); +} +inline float bfr(float f) { return bf(to_bf(f)); } + +// Sum over a work-group of 256 items (16 sub-groups of 16) in a fixed order. +inline float wg_sum256(float v, __local float *part) { + v = sub_group_reduce_add(v); + if (get_sub_group_local_id() == 0) part[get_sub_group_id()] = v; + barrier(CLK_LOCAL_MEM_FENCE); + float t = 0.0f; + for (uint i = 0; i < 16; i++) t += part[i]; + barrier(CLK_LOCAL_MEM_FENCE); + return t; +} + +// RMSNorm of row get_group_id(0): y = bf16(w * bf16(x * rsqrt(mean(x^2) + eps))); 256 items, n % 256 == 0, n <= 4096. +__attribute__((intel_reqd_sub_group_size(16))) +__kernel void lfm3_rmsnorm(__global const ushort *x, __global const ushort *w, __global ushort *y, uint n, float eps) { + __local float part[16]; + const uint row = get_group_id(0); + const uint lid = get_local_id(0); + const uint per = n / 256; + float v[NPI]; + float ss = 0.0f; + for (uint k = 0; k < per; k++) { + v[k] = bf(x[row * n + lid + 256 * k]); + ss += v[k] * v[k]; + } + const float rs = rsqrt(wg_sum256(ss, part) / (float)n + eps); + for (uint k = 0; k < per; k++) y[row * n + lid + 256 * k] = to_bf(bf(w[lid + 256 * k]) * bfr(v[k] * rs)); +} + +// Sandwich norm step: x = bf16(x + rmsnorm(d, wp)) in place, then y = rmsnorm(x, wn); 256 items a row. +__attribute__((intel_reqd_sub_group_size(16))) +__kernel void lfm3_post_add_norm(__global const ushort *d, __global ushort *x, __global const ushort *wp, __global const ushort *wn, __global ushort *y, + uint n, float eps) { + __local float part[16]; + const uint row = get_group_id(0); + const uint lid = get_local_id(0); + const uint per = n / 256; + float v[NPI]; + float ss = 0.0f; + for (uint k = 0; k < per; k++) { + v[k] = bf(d[row * n + lid + 256 * k]); + ss += v[k] * v[k]; + } + const float rd = rsqrt(wg_sum256(ss, part) / (float)n + eps); + ss = 0.0f; + for (uint k = 0; k < per; k++) { + const uint i = row * n + lid + 256 * k; + const float dn = bfr(bf(wp[lid + 256 * k]) * bfr(v[k] * rd)); + v[k] = bfr(bf(x[i]) + dn); + x[i] = to_bf(v[k]); + ss += v[k] * v[k]; + } + const float rx = rsqrt(wg_sum256(ss, part) / (float)n + eps); + for (uint k = 0; k < per; k++) y[row * n + lid + 256 * k] = to_bf(bf(wn[lid + 256 * k]) * bfr(v[k] * rx)); +} + +// Gated short conv: rows [B | C | x], y = bf16(C * bf16(conv7(bf16(B * x)))); state: last six B * x, oldest first. +__kernel void lfm3_conv_rows(__global const ushort *proj, uint pd, __global const ushort *sin, __global ushort *sout, __global const ushort *cw, + __global ushort *y, uint n, uint D) { + const uint c = get_global_id(0); + if (c >= D) return; + float h[6], w[7]; + for (uint k = 0; k < 6; k++) h[k] = bf(sin[k * D + c]); + for (uint k = 0; k < 7; k++) w[k] = bf(cw[c * 7 + k]); + for (uint r = 0; r < n; r++) { + __global const ushort *p = proj + r * pd; + const float bx = bfr(bf(p[c]) * bf(p[2 * D + c])); + float acc = 0.0f; + for (uint k = 0; k < 6; k++) acc += h[k] * w[k]; + acc += bx * w[6]; + y[r * D + c] = to_bf(bf(p[D + c]) * bfr(acc)); + for (uint k = 0; k < 5; k++) h[k] = h[k + 1]; + h[5] = bx; + } + for (uint k = 0; k < 6; k++) sout[k * D + c] = to_bf(h[k]); +} + +// lfm3_conv_rows for a verified window, also ext rows 0..5 = the state before it, 6 + r = row r's B * x (rollback). +__kernel void lfm3_conv_hist(__global const ushort *proj, uint pd, __global const ushort *sin, __global ushort *sout, __global const ushort *cw, + __global ushort *y, uint n, uint D, __global ushort *ext) { + const uint c = get_global_id(0); + if (c >= D) return; + float h[6], w[7]; + for (uint k = 0; k < 6; k++) { + h[k] = bf(sin[k * D + c]); + ext[k * D + c] = to_bf(h[k]); + } + for (uint k = 0; k < 7; k++) w[k] = bf(cw[c * 7 + k]); + for (uint r = 0; r < n; r++) { + __global const ushort *p = proj + r * pd; + const float bx = bfr(bf(p[c]) * bf(p[2 * D + c])); + float acc = 0.0f; + for (uint k = 0; k < 6; k++) acc += h[k] * w[k]; + acc += bx * w[6]; + y[r * D + c] = to_bf(bf(p[D + c]) * bfr(acc)); + ext[(6 + r) * D + c] = to_bf(bx); + for (uint k = 0; k < 5; k++) h[k] = h[k + 1]; + h[5] = bx; + } + for (uint k = 0; k < 6; k++) sout[k * D + c] = to_bf(h[k]); +} + +// One head of row r at position pos: RMSNorm over 128 dims, RoPE on pairs (d, d + 64) with bf16 cos/sin; 64 items. +inline void norm_rope_head(__global const ushort *s, __global ushort *o, __global const ushort *w, __global const float *inv_freq, uint pos, + float eps, __local float *part) { + const uint i = get_local_id(0); // 64 items in 4 sub-groups of 16 (the kernels require 16): part[4] + const float a = bf(s[i]), b = bf(s[i + 64]); + float ss = sub_group_reduce_add(a * a + b * b); + if (get_sub_group_local_id() == 0) part[get_sub_group_id()] = ss; + barrier(CLK_LOCAL_MEM_FENCE); + const uint nsg = get_num_sub_groups(); + float t = 0.0f; + for (uint k = 0; k < nsg; k++) t += part[k]; + const float rs = rsqrt(t / (float)HD + eps); + const float na = bfr(bf(w[i]) * bfr(a * rs)); + const float nb = bfr(bf(w[i + 64]) * bfr(b * rs)); + const float f = (float)pos * inv_freq[i]; + const float co = bfr(cos(f)), si = bfr(sin(f)); + o[i] = to_bf(na * co - nb * si); + o[i + 64] = to_bf(nb * co + na * si); +} + +// Per-head RMSNorm, RoPE on pairs (d, d + 64), bf16 cos/sin; grid (heads, rows); to row r or slot (pos0 + r) % ring. +__attribute__((intel_reqd_sub_group_size(16))) +__kernel void lfm3_qk_norm_rope(__global const ushort *src, uint src_stride, uint src_off, __global ushort *dst, uint dst_stride, + __global const ushort *w, __global const float *inv_freq, uint pos0, uint ring, float eps) { + __local float part[4]; + const uint head = get_group_id(0); + const uint r = get_group_id(1); + const uint row = ring == 0 ? r : (pos0 + r) % ring; + norm_rope_head(src + r * src_stride + src_off + head * HD, dst + row * dst_stride + head * HD, w, inv_freq, pos0 + r, eps, part); +} + +// q heads, then k heads and v rows into slot (pos0 + r) % ring, in one launch; grid (nq + nkv + nkv * 2, rows). +__attribute__((intel_reqd_sub_group_size(16))) +__kernel void lfm3_rope_store(__global const ushort *qkv, uint stride, __global ushort *q, __global ushort *kc, __global ushort *vc, + __global const ushort *qw, __global const ushort *kw, __global const float *inv_freq, uint nq, uint nkv, + uint pos0, uint ring, float eps) { + __local float part[4]; + const uint g = get_group_id(0); + const uint r = get_group_id(1); + const uint slot = (pos0 + r) % ring; + __global const ushort *src = qkv + r * stride; + if (g < nq) { + norm_rope_head(src + g * HD, q + (r * nq + g) * HD, qw, inv_freq, pos0 + r, eps, part); + } else if (g < nq + nkv) { + const uint h = g - nq; + norm_rope_head(src + (nq + h) * HD, kc + (slot * nkv + h) * HD, kw, inv_freq, pos0 + r, eps, part); + } else { + const uint e = (g - nq - nkv) * 64 + get_local_id(0); + vc[slot * nkv * HD + e] = src[(nq + nkv) * HD + e]; + } +} + +// Copies n_elem bf16 of row r to KV slot (pos0 + r) % ring; grid (ceil(n_elem / 64), rows). +__kernel void lfm3_v_store(__global const ushort *src, uint src_stride, uint src_off, __global ushort *dst, uint n_elem, uint pos0, uint ring) { + const uint e = get_global_id(0); + const uint r = get_group_id(1); + if (e < n_elem) dst[((pos0 + r) % ring) * n_elem + e] = src[r * src_stride + src_off + e]; +} + +// First key of row r's attention: window == 0 sees every earlier position, else the last `window` positions. +inline uint key_lo(uint p, uint window) { return (window == 0 || p + 1 <= window) ? 0 : p + 1 - window; } + +// Split-K GQA pass 1: group (kv head, chunk, row), a sub-group per q head, keys at slot k % ring; unnormalised o, m, l. +__attribute__((intel_reqd_sub_group_size(16))) +__kernel void lfm3_attn_partial(__global const ushort *q, __global const ushort *kc, __global const ushort *vc, __global float *po, __global float *pm, + __global float *pl, uint pos0, uint chunk, uint ring, uint window, uint nkv, uint maxc, float scale) { + const uint kvh = get_group_id(0); + const uint c = get_group_id(1); + const uint z = get_group_id(2); + const uint g = get_sub_group_id(); + const uint lane = get_sub_group_local_id(); + const uint nq = nkv * GQ; + const uint h = kvh * GQ + g; + const uint p = pos0 + z; + const uint lo = key_lo(p, window) + c * chunk; + const uint hi = min(p + 1, lo + chunk); + const float8 qv = as_float8(convert_uint8(vload8(0, q + (z * nq + h) * HD + lane * 8)) << 16); + float m = -INFINITY, l = 0.0f; + float8 o = (float8)(0.0f); + for (uint k = lo; k < hi; k++) { + const uint base = ((k % ring) * nkv + kvh) * HD + lane * 8; + const float8 kv = as_float8(convert_uint8(vload8(0, kc + base)) << 16); + const float8 pr = qv * kv; + const float s = sub_group_reduce_add(pr.s0 + pr.s1 + pr.s2 + pr.s3 + pr.s4 + pr.s5 + pr.s6 + pr.s7) * scale; + const float mn = fmax(m, s); + const float corr = exp(m - mn); + const float e = exp(s - mn); + const float8 vv = as_float8(convert_uint8(vload8(0, vc + base)) << 16); + l = l * corr + e; + o = o * corr + e * vv; + m = mn; + } + const uint slot = (z * nq + h) * maxc + c; + vstore8(o, 0, po + slot * HD + lane * 8); + if (lane == 0) { + pm[slot] = m; + pl[slot] = l; + } +} + +// Pass 2: row z's chunks of head h merged in order, bf16(O / L) times bf16(sigmoid(gate)); grid (nq, rows), 16 items. +__attribute__((intel_reqd_sub_group_size(16))) +__kernel void lfm3_attn_merge_gate(__global const float *po, __global const float *pm, __global const float *pl, __global const ushort *g, + __global ushort *out, uint pos0, uint chunk, uint window, uint nq, uint maxc) { + const uint h = get_group_id(0); + const uint z = get_group_id(1); + const uint lane = get_sub_group_local_id(); + const uint p = pos0 + z; + const uint len = p + 1 - key_lo(p, window); + const uint nch = (len + chunk - 1) / chunk; + const uint base = (z * nq + h) * maxc; + float mx = -INFINITY; + for (uint c = 0; c < nch; c++) mx = fmax(mx, pm[base + c]); + float L = 0.0f; + float8 O = (float8)(0.0f); + for (uint c = 0; c < nch; c++) { + const float e = exp(pm[base + c] - mx); + L += pl[base + c] * e; + O += vload8(0, po + (base + c) * HD + lane * 8) * e; + } + const float gate = bfr(1.0f / (1.0f + exp(-bf(g[z * nq + h])))); + const float8 r = O / L; + __global ushort *dst = out + (z * nq + h) * HD + lane * 8; + dst[0] = to_bf(bfr(r.s0) * gate); + dst[1] = to_bf(bfr(r.s1) * gate); + dst[2] = to_bf(bfr(r.s2) * gate); + dst[3] = to_bf(bfr(r.s3) * gate); + dst[4] = to_bf(bfr(r.s4) * gate); + dst[5] = to_bf(bfr(r.s5) * gate); + dst[6] = to_bf(bfr(r.s6) * gate); + dst[7] = to_bf(bfr(r.s7) * gate); +} + +// SwiGLU over n bf16 values: out = bf16(bf16(silu(a)) * b), a = w1 x (gate), b = w3 x (up). +__kernel void lfm3_swiglu(__global const ushort *a, __global const ushort *b, __global ushort *out, uint n) { + const uint i = get_global_id(0); + if (i >= n) return; + const float x = bf(a[i]); + out[i] = to_bf(bfr(x / (1.0f + exp(-x))) * bf(b[i])); +} diff --git a/zig/kernels/xpu/lfm3_mv.cl b/zig/kernels/xpu/lfm3_mv.cl new file mode 100644 index 0000000000..05db2de4cf --- /dev/null +++ b/zig/kernels/xpu/lfm3_mv.cl @@ -0,0 +1,154 @@ +// LFM3 decode matvecs fused by launch: attn.cl's qmv_row unchanged (same contraction), more per launch. +inline float bf(ushort v) { return as_float((uint)v << 16); } +inline ushort to_bf(float f) { + uint u = as_uint(f); + u += 0x7fffu + ((u >> 16) & 1u); + return (ushort)(u >> 16); +} + +// attn.cl's qmv_row, verbatim: one 16-lane sub-group a row, the sum valid on every lane. +inline float qmv_row(__global const uint *w, __global const ushort *scales, __global const ushort *biases, + __global const ushort *x, uint row, uint in_dim) { + const uint lane = get_sub_group_local_id(); + const uint words = in_dim / 8; + const uint groups = in_dim / 64; + float acc = 0.0f; + for (uint wi = lane; wi < words; wi += 16) { + const uint g = wi / 8; + const float sc = bf(scales[row * groups + g]); + const float bi = bf(biases[row * groups + g]); + const uint pk = w[row * words + wi]; + const float8 xv = as_float8(convert_uint8(vload8(wi, x)) << 16); + float dot = 0.0f, sx = 0.0f; + dot += xv.s0 * (float)(pk & 15u); sx += xv.s0; + dot += xv.s1 * (float)((pk >> 4) & 15u); sx += xv.s1; + dot += xv.s2 * (float)((pk >> 8) & 15u); sx += xv.s2; + dot += xv.s3 * (float)((pk >> 12) & 15u); sx += xv.s3; + dot += xv.s4 * (float)((pk >> 16) & 15u); sx += xv.s4; + dot += xv.s5 * (float)((pk >> 20) & 15u); sx += xv.s5; + dot += xv.s6 * (float)((pk >> 24) & 15u); sx += xv.s6; + dot += xv.s7 * (float)((pk >> 28) & 15u); sx += xv.s7; + acc += sc * dot + bi * sx; + } + return sub_group_reduce_add(acc); +} + +// Four tables in one launch (q, k, v, gate): rows r0 .. r3 into y0 + o0 .. y3 + o3; 4 rows a work-group. +__attribute__((intel_reqd_sub_group_size(16))) +__kernel void qmv4_bf_x4(__global const uint *w0, __global const ushort *s0, __global const ushort *b0, __global ushort *y0, uint r0, uint o0, + __global const uint *w1, __global const ushort *s1, __global const ushort *b1, __global ushort *y1, uint r1, uint o1, + __global const uint *w2, __global const ushort *s2, __global const ushort *b2, __global ushort *y2, uint r2, uint o2, + __global const uint *w3, __global const ushort *s3, __global const ushort *b3, __global ushort *y3, uint r3, uint o3, + __global const ushort *x, uint in_dim) { + uint row = get_group_id(0) * 4 + get_sub_group_id(); + float acc; + __global ushort *y; + if (row < r0) { + acc = qmv_row(w0, s0, b0, x, row, in_dim); + y = y0 + o0; + } else if ((row -= r0) < r1) { + acc = qmv_row(w1, s1, b1, x, row, in_dim); + y = y1 + o1; + } else if ((row -= r1) < r2) { + acc = qmv_row(w2, s2, b2, x, row, in_dim); + y = y2 + o2; + } else if ((row -= r2) < r3) { + acc = qmv_row(w3, s3, b3, x, row, in_dim); + y = y3 + o3; + } else return; + if (get_sub_group_local_id() == 0) y[row] = to_bf(acc); +} + +// w1 and w3 rows and their SwiGLU: act = bf16(bf16(silu(bf16(w1 x))) * bf16(w3 x)), as qmv4_bf + lfm3_swiglu. +__attribute__((intel_reqd_sub_group_size(16))) +__kernel void qmv4_swiglu(__global const uint *w1, __global const ushort *s1, __global const ushort *b1, __global const uint *w3, __global const ushort *s3, + __global const ushort *b3, __global const ushort *x, __global ushort *act, uint in_dim, uint rows) { + const uint row = get_group_id(0) * 4 + get_sub_group_id(); + if (row >= rows) return; + const float a = bf(to_bf(qmv_row(w1, s1, b1, x, row, in_dim))); + const float u = bf(to_bf(qmv_row(w3, s3, b3, x, row, in_dim))); + if (get_sub_group_local_id() == 0) act[row] = to_bf(bf(to_bf(a / (1.0f + exp(-a)))) * u); +} + +// Prompt windows: bf16(x . (q * scale + bias)), one fma chain a 64-wide k block, blocks in order: any cut, same rows. +__attribute__((reqd_work_group_size(16, 16, 1))) +__kernel void pgemm(__global const uint *w, __global const ushort *scales, __global const ushort *biases, __global const ushort *x, + __global ushort *y, uint k, uint n, uint m, uint ldy, uint y_off, uint qg) { + __local float xs[16][64]; + __local float ws[16][64]; + const uint lx = get_local_id(0), ly = get_local_id(1), t = ly * 16 + lx; + const ulong m0 = (ulong)get_group_id(1) * 64, n0 = (ulong)get_group_id(0) * 64; + const uint lr = t / 4, lk = (t % 4) * 4, words = k / 8, groups = k / qg; + float acc[4][4]; +#pragma unroll + for (uint i = 0; i < 4; i++) +#pragma unroll + for (uint j = 0; j < 4; j++) acc[i][j] = 0.0f; + float part[4][4]; + for (uint k0 = 0; k0 < k; k0 += 16) { + const ulong xr = m0 + lr, wr = n0 + lr; + const uint kk = k0 + lk; + ushort4 xv = (ushort4)(0); + if (xr < m) xv = vload4(0, x + xr * k + kk); + float4 wv = (float4)(0.0f); + if (wr < n) { + const uint p = w[wr * words + kk / 8] >> (4 * (kk % 8)); + const float sc = bf(scales[wr * groups + kk / qg]), bi = bf(biases[wr * groups + kk / qg]); + wv = (float4)((float)(p & 15u) * sc + bi, (float)((p >> 4) & 15u) * sc + bi, (float)((p >> 8) & 15u) * sc + bi, + (float)((p >> 12) & 15u) * sc + bi); + } + xs[lk][lr] = bf(xv.s0); + xs[lk + 1][lr] = bf(xv.s1); + xs[lk + 2][lr] = bf(xv.s2); + xs[lk + 3][lr] = bf(xv.s3); + ws[lk][lr] = wv.s0; + ws[lk + 1][lr] = wv.s1; + ws[lk + 2][lr] = wv.s2; + ws[lk + 3][lr] = wv.s3; + barrier(CLK_LOCAL_MEM_FENCE); + if (k0 % 64 == 0) { +#pragma unroll + for (uint i = 0; i < 4; i++) +#pragma unroll + for (uint j = 0; j < 4; j++) part[i][j] = 0.0f; + } +#pragma unroll + for (uint q = 0; q < 16; q++) { + const float4 a = vload4(ly, &xs[q][0]); + const float4 b = vload4(lx, &ws[q][0]); + part[0][0] = fma(a.s0, b.s0, part[0][0]); + part[0][1] = fma(a.s0, b.s1, part[0][1]); + part[0][2] = fma(a.s0, b.s2, part[0][2]); + part[0][3] = fma(a.s0, b.s3, part[0][3]); + part[1][0] = fma(a.s1, b.s0, part[1][0]); + part[1][1] = fma(a.s1, b.s1, part[1][1]); + part[1][2] = fma(a.s1, b.s2, part[1][2]); + part[1][3] = fma(a.s1, b.s3, part[1][3]); + part[2][0] = fma(a.s2, b.s0, part[2][0]); + part[2][1] = fma(a.s2, b.s1, part[2][1]); + part[2][2] = fma(a.s2, b.s2, part[2][2]); + part[2][3] = fma(a.s2, b.s3, part[2][3]); + part[3][0] = fma(a.s3, b.s0, part[3][0]); + part[3][1] = fma(a.s3, b.s1, part[3][1]); + part[3][2] = fma(a.s3, b.s2, part[3][2]); + part[3][3] = fma(a.s3, b.s3, part[3][3]); + } + if ((k0 + 16) % 64 == 0 || k0 + 16 >= k) { +#pragma unroll + for (uint i = 0; i < 4; i++) +#pragma unroll + for (uint j = 0; j < 4; j++) acc[i][j] += part[i][j]; + } + barrier(CLK_LOCAL_MEM_FENCE); + } +#pragma unroll + for (uint i = 0; i < 4; i++) { + const ulong r = m0 + ly * 4 + i; + if (r >= m) continue; +#pragma unroll + for (uint j = 0; j < 4; j++) { + const ulong o = n0 + lx * 4 + j; + if (o < n) y[r * ldy + y_off + o] = to_bf(acc[i][j]); + } + } +} diff --git a/zig/kernels/xpu/nem_prow.cl b/zig/kernels/xpu/nem_prow.cl index 7cc36c1a92..3a88dcecfa 100644 --- a/zig/kernels/xpu/nem_prow.cl +++ b/zig/kernels/xpu/nem_prow.cl @@ -1,7 +1,11 @@ // Nemotron-H prompt rows (any count) with the decode step's arithmetic: each output sums as the one-token path does. #define SG 16 +#ifndef PR #define PR 8 // prompt rows a sub-group holds -#define PO 8 // output rows a sub-group holds: one load and one sum of a word's inputs serve all of them +#endif +#ifndef PO +#define PO 8 // output rows a sub-group holds: one load and one sum of a word's inputs serve all (nem_prow_gen12: 4) +#endif #define PSG (16 / PO) // sub-groups a work-group: 16 output rows inline float bf(ushort v) { return as_float((uint)v << 16); } inline ushort to_bf(float f) { diff --git a/zig/src/families/lfm3/config.zig b/zig/src/families/lfm3/config.zig index 84ac2ae5b5..3cd647ec7e 100644 --- a/zig/src/families/lfm3/config.zig +++ b/zig/src/families/lfm3/config.zig @@ -12,8 +12,38 @@ pub const head_dim = 128; pub const sliding_window = 256; pub const eps: f32 = 1e-5; pub const Kind = enum { text, vision_language }; -pub const LayerType = enum { conv, full_attention, sliding_attention }; -pub const Rotary = struct { theta: f32, partial: f32 }; +pub const LayerType = enum { + conv, + full_attention, + sliding_attention, + + /// Keys before a row its attention reads: the sliding window, or 0 for every one (a conv layer reads none). + pub fn window(t: LayerType) usize { + return if (t == .sliding_attention) sliding_window else 0; + } +}; + +pub const Rotary = struct { + theta: f32, + partial: f32, + + /// RoPE's inverse frequencies as torch's float32 rows: 1 / theta^(2i / 128) rotated, zero past them (cos 1, sin 0). + pub fn inverseFrequencies(rot: Rotary) [head_dim / 2]f32 { + var out: [head_dim / 2]f32 = undefined; + const angles: usize = @intFromFloat(rot.partial * @as(f32, head_dim) / 2); + for (&out, 0..) |*v, i| { + if (i >= angles) { + v.* = 0; + continue; + } + const e: f32 = @as(f32, @floatFromInt(2 * i)) / @as(f32, head_dim); + const pw: f32 = @floatCast(std.math.pow(f64, rot.theta, e)); + v.* = 1.0 / pw; + } + return out; + } +}; + pub const schedule: [layers]LayerType = .{ .conv, .conv, .full_attention, .conv, .sliding_attention, .conv, .full_attention, .conv, .sliding_attention, .conv, .full_attention, .conv, @@ -21,6 +51,18 @@ pub const schedule: [layers]LayerType = .{ .full_attention, .conv, .sliding_attention, .conv, .full_attention, .conv, }; +/// Layers of `kind` in the admitted schedule. +pub fn count(kind: LayerType) usize { + var n: usize = 0; + for (schedule) |t| n += @intFromBool(t == kind); + return n; +} + +/// Keys a sliding layer's ring holds when `rows` rows are written before the first of them reads its window back. +pub fn ring(rows: usize) usize { + return sliding_window + rows; +} + pub const Width = struct { bits: u8, group: u16 }; pub const Quant = struct { global: Width, diff --git a/zig/src/families/lfm3/config_test.zig b/zig/src/families/lfm3/config_test.zig index f1a1df00ce..136223dbd5 100644 --- a/zig/src/families/lfm3/config_test.zig +++ b/zig/src/families/lfm3/config_test.zig @@ -19,6 +19,41 @@ test "text and VL preserve their distinct context and actual mixed embedding wid try std.testing.expectEqual(@as(f32, 1), t.sliding_rotary.partial); } +test "RoPE: both settings' inverse frequencies bit for bit torch 2.14's float32 rows, zero past the rotated pairs" { + const full = [32]u32{ + 0x3f800000, 0x3f4e4bad, 0x3f263de0, 0x3f05f6ef, 0x3ed7e89b, 0x3eadfcff, 0x3e8c3504, 0x3e61f835, + 0x3e361887, 0x3e12bd91, 0x3dec7fd6, 0x3dbe94c6, 0x3d99940d, 0x3d778513, 0x3d47763f, 0x3d20bc1d, + 0x3d0186e3, 0x3cd0c1a8, 0x3ca8398b, 0x3c879008, 0x3c5a7bf2, 0x3c301052, 0x3c0de12d, 0x3be4aa46, + 0x3bb8449c, 0x3b947dae, 0x3b6f520e, 0x3b40dac5, 0x3b1b690d, 0x3afa78f0, 0x3ac9d75c, 0x3aa2a6f7, + }; + const sliding = [64]u32{ + 0x3f800000, 0x3f5dafd7, 0x3f3ff911, 0x3f263de0, 0x3f0ff59a, 0x3ef953cf, 0x3ed7e89b, 0x3ebaf81b, + 0x3ea1e89b, 0x3e8c3504, 0x3e72d423, 0x3e5247ed, 0x3e361887, 0x3e1db040, 0x3e088d77, 0x3dec7fd6, + 0x3dcccccd, 0x3db15978, 0x3d99940d, 0x3d84fe4d, 0x3d6655c2, 0x3d47763f, 0x3d2cba15, 0x3d159348, + 0x3d0186e3, 0x3ce054d2, 0x3cc2434f, 0x3ca8398b, 0x3c91ad39, 0x3c7c4d33, 0x3c5a7bf2, 0x3c3d3311, + 0x3c23d70a, 0x3c0de12d, 0x3bf5b9b0, 0x3bd4ca15, 0x3bb8449c, 0x3b9f91cc, 0x3b8a2e77, 0x3b6f520e, + 0x3b4f3e38, 0x3b33770f, 0x3b1b690d, 0x3b06946f, 0x3ae91528, 0x3ac9d75c, 0x3aaec98e, 0x3a975c0e, + 0x3a83126f, 0x3a6301e2, 0x3a44948c, 0x3a2a3b44, 0x3a136a16, 0x39ff4fac, 0x39dd1725, 0x39bf74d7, + 0x39a5cb60, 0x398f9272, 0x3978a815, 0x395753e4, 0x393a7753, 0x39217916, 0x390bd472, 0x38f22ce2, + }; + const f = (c.Rotary{ .theta = 1e6, .partial = 0.5 }).inverseFrequencies(); + const s = (c.Rotary{ .theta = 1e4, .partial = 1 }).inverseFrequencies(); + for (full, f[0..32]) |want, v| try std.testing.expectEqual(want, @as(u32, @bitCast(v))); + for (f[32..]) |v| try std.testing.expectEqual(@as(f32, 0), v); + for (sliding, s) |want, v| try std.testing.expectEqual(want, @as(u32, @bitCast(v))); +} + +test "the schedule: 13 conv, 6 full and 5 sliding layers; only a sliding layer has a window, its ring the window and the rows" { + try std.testing.expectEqual(@as(usize, 13), c.count(.conv)); + try std.testing.expectEqual(@as(usize, 6), c.count(.full_attention)); + try std.testing.expectEqual(@as(usize, 5), c.count(.sliding_attention)); + try std.testing.expectEqual(@as(usize, 256), c.LayerType.sliding_attention.window()); + try std.testing.expectEqual(@as(usize, 0), c.LayerType.full_attention.window()); + try std.testing.expectEqual(@as(usize, 0), c.LayerType.conv.window()); + try std.testing.expectEqual(@as(usize, 256 + 16), c.ring(16)); + try std.testing.expectEqual(@as(usize, 1280), c.ring(1024)); +} + test "6 and 8 bit affine variants use the same config parser" { inline for (.{ "6", "8" }) |bits| { const replacement = "\"bits\": " ++ bits; diff --git a/zig/src/families/lfm3/xpu.zig b/zig/src/families/lfm3/xpu.zig new file mode 100644 index 0000000000..3cefba81da --- /dev/null +++ b/zig/src/families/lfm3/xpu.zig @@ -0,0 +1,4 @@ +//! LFM3 on the Intel GPU (XPU backend): the forward and the config it reads. +pub const model = @import("xpu_model.zig"); +pub const engine = @import("xpu_engine.zig"); +pub const config = @import("config.zig"); diff --git a/zig/src/families/lfm3/xpu_engine.zig b/zig/src/families/lfm3/xpu_engine.zig new file mode 100644 index 0000000000..f71d117870 --- /dev/null +++ b/zig/src/families/lfm3/xpu_engine.zig @@ -0,0 +1,234 @@ +//! LFM3 behind the Intel GPU lane host: a sequence a stream, windows verified bit for bit the steps. +const std = @import("std"); +const xpu = @import("xpu"); +const cfg = @import("config.zig"); +const model = @import("xpu_model.zig"); + +pub const Options = struct { context: usize = 4096, window_rows: u32 = 64 }; + +pub const Engine = struct { + pub const Sequence = model.Seq; + pub const Snapshot = model.Snap; + /// The lane core drafts by copying the context (no head): windows of up to 16 rows, bit for bit the steps. + pub const copy_only = true; + head: ?u8 = null, + gpa: std.mem.Allocator, + stream: xpu.Stream, + r: xpu.rt.Runtime, + loader: xpu.loader.Loader, + c: cfg.Config, + m: model.Model, + max_len: usize, + window_rows: u32, + load_seconds: f64, + /// The first stream's state, held here so the lane host can tell it from the streams it adds (newSeq). + own: Sequence, + /// Each slot's sequence (0 is `own`); verifyMulti names streams by slot. + slots: [model.max_streams]?*Sequence, + + pub fn init(gpa: std.mem.Allocator, io: std.Io, ctx: *const xpu.Context, dir: []const u8, o: Options) !*Engine { + const t0 = std.Io.Clock.awake.now(io); + const e = try gpa.create(Engine); + errdefer gpa.destroy(e); + e.gpa = gpa; + e.c = try cfg.Config.read(gpa, io, dir); + errdefer e.c.deinit(); + if (e.c.kind != .text) return error.UnsupportedFormat; + const q = e.c.quant.global; + if (q.bits != 4 or q.group != 64 or e.c.quant.modules.count() != 0) return error.UnsupportedFormat; + e.stream = try xpu.Stream.init(ctx); + errdefer e.stream.deinit(); + e.r = try xpu.rt.Runtime.init(ctx, e.stream); + errdefer e.r.deinit(); + e.loader = try xpu.loader.Loader.init(gpa, &e.r, dir); + e.window_rows = o.window_rows; + e.m = try model.Model.load(gpa, &e.r, &e.loader, &e.c, @intCast(o.context), o.window_rows); + e.own = e.m.own.*; + gpa.destroy(e.m.own); + e.m.own = &e.own; + e.m.seq = &e.own; + e.head = null; + e.slots = @splat(null); + e.slots[0] = &e.own; + e.max_len = o.context; + e.load_seconds = xpu.decode.seconds(io, t0); + return e; + } + + /// The process exits right after: device memory goes with it once the queue is drained. + pub fn deinit(e: *Engine) void { + e.r.deinit(); + e.stream.deinit(); + e.c.deinit(); + e.gpa.destroy(e); + } + + pub fn deviceBytes(e: *const Engine) u64 { + return e.loader.total; + } + + pub fn vocab(_: *const Engine) usize { + return cfg.vocab; + } + + pub fn reset(e: *Engine) !void { + try e.m.reset(); + } + + pub fn feed(e: *Engine, token: u32, logits: bool) !void { + try e.m.forward(&.{token}, logits, null); + } + + pub fn sync(e: *Engine) !void { + try e.r.sync(); + } + + pub fn argmax(e: *Engine) !u32 { + return e.m.argmax(); + } + + pub fn fetchLogits(e: *Engine, out: []f32) !void { + try e.m.fetchLogits(out); + } + + pub fn newSeq(e: *Engine) !*Sequence { + const slot = for (e.slots[1..], 1..) |q, i| { + if (q == null) break i; + } else return error.NoFreeSlot; + const q = try e.m.newSeq(); + q.slot = @intCast(slot); + e.slots[slot] = q; + return q; + } + + pub fn freeSeq(e: *Engine, q: *Sequence) void { + e.slots[q.slot] = null; + e.m.freeSeq(q); + } + + pub fn slotOf(_: *const Engine, q: *const Sequence) u32 { + return q.slot; + } + + pub fn bind(e: *Engine, q: *Sequence) void { + e.m.bind(q); + } + + /// Prompt tokens [from, to) of the bound stream in windows of at most `rows`; the last window keeps logits. + pub fn prefillSpan(e: *Engine, ids: []const u32, from: usize, to: usize, rows: u32) !void { + if (e.m.pos != from) return error.PositionMismatch; + const w = @min(rows, e.window_rows); + var t = from; + while (t < to) { + const n = @min(w, to - t); + try e.m.forwardPrompt(ids[t .. t + n], t + n == ids.len, null); + t += n; + } + } + + /// The pending token and drafts of the bound stream as one window; every row keeps its logits and argmax. + pub fn verifyRows(e: *Engine, ids: []const u32) !void { + try e.m.forwardSegs(&.{.{ .q = e.m.seq, .ids = ids }}, .rows, true, null); + } + + /// Windows of several streams (by slot) in one forward over the weights, each row at its stream's position. + pub fn verifyMulti(e: *Engine, slots: []const u32, windows: []const []const u32) !void { + var segs: [model.max_streams]model.Seg = undefined; + if (slots.len > segs.len) return error.WindowTooWide; + for (slots, windows, segs[0..slots.len]) |sl, ids, *g| g.* = .{ .q = e.slots[sl] orelse return error.UnknownStream, .ids = ids }; + try e.m.forwardSegs(segs[0..slots.len], .rows, true, null); + } + + /// Keeps the first c rows of the bound stream's last window. + pub fn keepRows(e: *Engine, c: u32) !void { + try e.m.keepRows(c); + } + + /// The last window's greedy tokens, a row each (all streams' rows in order after verifyMulti). + pub fn rowTokens(e: *Engine, out: []i32) !void { + try e.m.rowTokens(out); + } + + pub fn rowLogits(e: *Engine, out: []f32) !void { + try e.m.rowLogits(out); + } + + /// No head: the lane core's copies are the drafts. + pub fn draftKept(_: *Engine, _: []const u32, _: u32, _: u32, _: []u32) !usize { + return 0; + } + + pub fn draftFirst(_: *Engine, _: u32, _: u32, _: []u32) !usize { + return 0; + } + + /// Window ms by rows on a scratch sequence: 1, 2, 4, 8, 16 timed (best of 3), the rest interpolated; no head: 0. + pub fn timeRounds(e: *Engine, ms: []f64) !f64 { + const prev = e.m.seq; + const q = try e.m.newSeq(); + e.m.bind(q); + defer { + e.m.bind(prev); + e.m.freeSeq(q); + } + const ids: [model.window_rows]u32 = @splat(1); + const timed = [_]usize{ 1, 2, 4, 8, model.window_rows }; + std.debug.assert(ms.len == model.window_rows); + for (timed) |w| { + var best: f64 = std.math.inf(f64); + for (0..3) |_| { + try e.m.reset(); + try e.r.sync(); + var ts0: std.c.timespec = undefined; + _ = std.c.clock_gettime(.MONOTONIC, &ts0); + try e.m.forwardSegs(&.{.{ .q = q, .ids = ids[0..w] }}, .rows, true, null); + try e.r.sync(); + var ts1: std.c.timespec = undefined; + _ = std.c.clock_gettime(.MONOTONIC, &ts1); + best = @min(best, @as(f64, @floatFromInt(ts1.sec - ts0.sec)) * 1e3 + @as(f64, @floatFromInt(ts1.nsec - ts0.nsec)) / 1e6); + } + ms[w - 1] = best; + } + for (timed[0 .. timed.len - 1], timed[1..]) |a, b| for (a + 1..b) |w| { + const f = @as(f64, @floatFromInt(w - a)) / @as(f64, @floatFromInt(b - a)); + ms[w - 1] = ms[a - 1] + f * (ms[b - 1] - ms[a - 1]); + }; + try e.m.reset(); + return 0; + } + + /// One decode step of the bound stream, with its logits. + pub fn stepRow(e: *Engine, token: u32) !void { + try e.m.forward(&.{token}, true, null); + } + + pub fn snapBytes(e: *const Engine, at: u32) u64 { + return e.m.snapBytes(at); + } + + /// The bound stream's state at prompt position `at`, which it has just reached. + pub fn snapSave(e: *Engine, at: u32) !*Snapshot { + if (at != e.m.pos) return error.PositionMismatch; + return e.m.snapSave(); + } + + pub fn snapRestore(e: *Engine, sn: *const Snapshot) !void { + try e.m.snapRestore(sn); + } + + pub fn snapDrop(e: *Engine, sn: *Snapshot) void { + e.m.snapDrop(sn); + } + + /// The prompt head in windows of at most `rows` tokens, without logits. + pub fn prefillWindows(e: *Engine, head: []const u32, rows: u32) !void { + var t: usize = 0; + const w = @min(rows, e.window_rows); + while (t < head.len) { + const n = @min(w, head.len - t); + try e.m.forwardPrompt(head[t .. t + n], false, null); + t += n; + } + try e.r.sync(); + } +}; diff --git a/zig/src/families/lfm3/xpu_model.zig b/zig/src/families/lfm3/xpu_model.zig new file mode 100644 index 0000000000..3fda63bdd4 --- /dev/null +++ b/zig/src/families/lfm3/xpu_model.zig @@ -0,0 +1,578 @@ +//! LFM3 forward on the Intel GPU: MLX 4-bit weights on the device, a conv state and KV rings, windows of rows, no DPAS. + +const std = @import("std"); +const xpu = @import("xpu"); +const cfg = @import("config.zig"); +const names = @import("names.zig"); +const rt = xpu.rt; +const ld = xpu.loader; + +const Buf = rt.Buffer; +const H: u32 = cfg.hidden; +const FF: u32 = cfg.intermediate; +const V: u32 = cfg.vocab; +const NQ: u32 = cfg.query_heads; +const NKV: u32 = cfg.kv_heads; +const HD: u32 = cfg.head_dim; +const QKV: u32 = (NQ + 2 * NKV) * HD; +const chunk: u32 = 128; // keys a split-K work-group scans +const argmax_parts: u32 = 64; +/// Without DPAS (Gen12LP) prompt projections take nem_prow_gen12: 4 outputs a sub-group, as 8 x 8 spills on Gen12LP. +const gen12 = !xpu.kernels.dpas; +const tile_rows: u32 = 8; +const tile_outs: u32 = if (gen12) 4 else 8; + +pub const Table = struct { w: Buf, s: Buf, b: Buf }; +const Conv = struct { in: Table, out: Table, cw: Buf }; +const Attn = struct { q: Table, k: Table, v: Table, g: Table, o: Table, qn: Buf, kn: Buf, ring: u32, window: u32 }; +const Mixer = union(enum) { conv: Conv, attn: Attn }; +const Layer = struct { op_norm: Buf, post_op: Buf, ffn_norm: Buf, post_ffn: Buf, w1: Table, w3: Table, w2: Table, mixer: Mixer }; + +const Kernels = struct { + embed: rt.Kernel, + embed_tok: rt.Kernel, + conv_hist: rt.Kernel, + qmv: rt.Kernel, + head: rt.Kernel, + am_part: rt.Kernel, + am_fin: rt.Kernel, + pdense: rt.Kernel, + rms: rt.Kernel, + pan: rt.Kernel, + conv: rt.Kernel, + rope: rt.Kernel, + vst: rt.Kernel, + part: rt.Kernel, + merge: rt.Kernel, + swiglu: rt.Kernel, + qkvg: rt.Kernel, + w13: rt.Kernel, + rope_store: rt.Kernel, + pgemm: rt.Kernel, +}; + +/// One stream's state: each conv layer's last six B*x rows, each attention layer's KV slots, and its position. +pub const Seq = struct { + conv: [cfg.layers]Buf = undefined, + kc: [cfg.layers]Buf = undefined, + vc: [cfg.layers]Buf = undefined, + pos: u32 = 0, + /// Each conv layer's state before the last verified window, then its rows' B * x: keepRows' source. + ext: [cfg.layers]Buf = undefined, + /// The last verified window still to keep: its first position and rows (0: none). + win_start: u32 = 0, + win_rows: u32 = 0, + /// The engine's slot for this stream (0 is the engine's own). + slot: u32 = 0, +}; + +/// Rows of a verified window at most (all streams' rows of a shared forward), and streams a shared forward takes. +pub const window_rows: u32 = 16; +pub const max_streams: u32 = 8; + +/// A stream's rows in one forward: `ids` at the stream's own position (its window, or one row of several streams). +pub const Seg = struct { q: *Seq, ids: []const u32 }; +/// What a forward leaves: nothing, the last row's logits and argmax (a prompt), or every row's (verified windows). +pub const Out = enum { none, last, rows }; + +/// A kept prompt state at `at`: conv rows, the full layers' first `at` KV slots, the sliding rings, in one buffer. +pub const Snap = struct { at: u32, buf: Buf }; + +const conv_bytes: usize = 6 * H * 2; +const ext_bytes: usize = (6 + window_rows) * H * 2; +const slot_bytes: usize = NKV * HD * 2; + +/// `b` from byte `off` on. +fn from(b: Buf, off: usize) Buf { + return .{ .rt = b.rt, .ptr = @ptrFromInt(@intFromPtr(b.ptr.?) + off), .len = b.len - off }; +} + +/// Device scratch for a window of up to `max_rows` rows. +const Scratch = struct { ids: Buf, x: Buf, xn: Buf, d: Buf, proj: Buf, qkv: Buf, g: Buf, q: Buf, att: Buf, a: Buf, b: Buf, act: Buf, po: Buf, pm: Buf, pl: Buf, logits: Buf, am_scratch: Buf, am_out: Buf }; + +/// Called after each layer (index) and after the final norm (-1) with the window's rows, bf16 [n][2048]. +pub const Hook = struct { ctx: *anyopaque, f: *const fn (ctx: *anyopaque, layer: i32, rows: []const u16) void }; + +fn run(k: *rt.Kernel, groups: [3]u32, a: anytype) !void { + inline for (a, 0..) |v, i| { + const T = @TypeOf(v); + if (T == Buf) { + try k.setBuffer(i, v); + } else if (T == f32) { + try k.setF32(i, v); + } else try k.setU32(i, @as(u32, v)); + } + try k.launch(groups); +} + +pub const Model = struct { + gpa: std.mem.Allocator, + r: *rt.Runtime, + layers: []Layer, + emb: Table, + norm_f: Buf, + inv_full: Buf, + inv_slide: Buf, + k: Kernels, + s: Scratch, + host_rows: []u16, + cap: u32, + max_rows: u32, + maxc: u32, + pos: u32 = 0, + device_bytes: u64 = 0, + /// The bound stream's state (`own` until another is bound); `pos` is the bound stream's position. + seq: *Seq = undefined, + own: *Seq = undefined, + /// Fused launches by bit (1 w1+w3+SwiGLU, 2 q/k/v/gate, 4 rope and KV store); LFM3_FUSE sets them for A/B. + fuse: u32 = 7, + /// In forwardPrompt: projections take pgemm, whose rows are the same however the prompt is cut. + prompt: bool = false, + /// The MLX group size of the projections. + qg: u32 = 64, + + /// The affine tensors of module `m` (in layer `layer`, or a model tensor) by names.zig in namespace `ns`. + fn table(l: *ld.Loader, ns: names.Namespace, layer: ?usize, m: names.Module) !Table { + var a: [160]u8 = undefined; + var b: [176]u8 = undefined; + const base = try names.moduleName(&a, ns, layer, m); + return .{ + .w = try l.load(try std.fmt.bufPrint(&b, "{s}.weight", .{base})), + .s = try l.load(try std.fmt.bufPrint(&b, "{s}.scales", .{base})), + .b = try l.load(try std.fmt.bufPrint(&b, "{s}.biases", .{base})), + }; + } + + fn vec(l: *ld.Loader, ns: names.Namespace, layer: ?usize, m: names.Module) !Buf { + var a: [160]u8 = undefined; + return l.load(try names.moduleName(&a, ns, layer, m)); + } + + /// The checkpoint's namespace (bare `model.` or wrapped `language_model.model.`) from the embedding it holds. + fn namespace(l: *const ld.Loader) !names.Namespace { + var a: [64]u8 = undefined; + var b: [64]u8 = undefined; + return names.detect(l.map.contains(try names.embedName(&a, .bare)), l.map.contains(try names.embedName(&b, .wrapped))); + } + + /// Uploads the checkpoint: `cap` positions on full-attention layers, windows of at most `max_rows` rows. + pub fn load(gpa: std.mem.Allocator, r: *rt.Runtime, l: *ld.Loader, c: *const cfg.Config, cap: u32, max_rows: u32) !Model { + if (max_rows == 0 or max_rows > 256 or cap < max_rows) return error.Invalid; + var bm = try r.module(xpu.kernels.basic); + var am = try r.module(xpu.kernels.attn); + var pm = try r.module(if (gen12) xpu.kernels.nem_prow_gen12 else xpu.kernels.nem_prow); + var fm = try r.module(xpu.kernels.lfm3); + var vm = try r.module(xpu.kernels.lfm3_mv); + const k: Kernels = .{ + .embed = try bm.kernel("embed4", .{ 64, 1, 1 }), + .embed_tok = try bm.kernel("embed4_tok", .{ 64, 1, 1 }), + .conv_hist = try fm.kernel("lfm3_conv_hist", .{ 64, 1, 1 }), + .qmv = try am.kernel("qmv4_bf", .{ 64, 1, 1 }), + .head = try am.kernel("qmv4_f32", .{ 64, 1, 1 }), + .am_part = try am.kernel("argmax_partial", .{ 256, 1, 1 }), + .am_fin = try am.kernel("argmax_final", .{ 256, 1, 1 }), + .pdense = try pm.kernel("pdense", .{ 16 * (16 / tile_outs), 1, 1 }), + .rms = try fm.kernel("lfm3_rmsnorm", .{ 256, 1, 1 }), + .pan = try fm.kernel("lfm3_post_add_norm", .{ 256, 1, 1 }), + .conv = try fm.kernel("lfm3_conv_rows", .{ 64, 1, 1 }), + .rope = try fm.kernel("lfm3_qk_norm_rope", .{ 64, 1, 1 }), + .vst = try fm.kernel("lfm3_v_store", .{ 64, 1, 1 }), + .part = try fm.kernel("lfm3_attn_partial", .{ 64, 1, 1 }), + .merge = try fm.kernel("lfm3_attn_merge_gate", .{ 16, 1, 1 }), + .swiglu = try fm.kernel("lfm3_swiglu", .{ 64, 1, 1 }), + .qkvg = try vm.kernel("qmv4_bf_x4", .{ 64, 1, 1 }), + .w13 = try vm.kernel("qmv4_swiglu", .{ 64, 1, 1 }), + .rope_store = try fm.kernel("lfm3_rope_store", .{ 64, 1, 1 }), + .pgemm = try vm.kernel("pgemm", .{ 16, 16, 1 }), + }; + const maxc = (cap + chunk - 1) / chunk; + const w: usize = max_rows; + const s: Scratch = .{ + .ids = try l.empty(w * 4), + .x = try l.empty(w * H * 2), + .xn = try l.empty(w * H * 2), + .d = try l.empty(w * H * 2), + .proj = try l.empty(w * 3 * H * 2), + .qkv = try l.empty(w * QKV * 2), + .g = try l.empty(w * NQ * 2), + .q = try l.empty(w * NQ * HD * 2), + .att = try l.empty(w * NQ * HD * 2), + .a = try l.empty(w * FF * 2), + .b = try l.empty(w * FF * 2), + .act = try l.empty(w * FF * 2), + .po = try l.empty(w * NQ * maxc * HD * 4), + .pm = try l.empty(w * NQ * maxc * 4), + .pl = try l.empty(w * NQ * maxc * 4), + .logits = try l.empty(@as(usize, window_rows) * V * 4), + .am_scratch = try l.empty(window_rows * argmax_parts * 8), + .am_out = try l.empty(window_rows * 4), + }; + const ns = try namespace(l); + const layers = try gpa.alloc(Layer, cfg.layers); + for (layers, 0..) |*ly, i| { + const kind = c.layer_types[i]; + ly.* = .{ + .op_norm = try vec(l, ns, i, .operator_norm), + .post_op = try vec(l, ns, i, .post_operator_norm), + .ffn_norm = try vec(l, ns, i, .ffn_norm), + .post_ffn = try vec(l, ns, i, .post_ffn_norm), + .w1 = try table(l, ns, i, .w1), + .w3 = try table(l, ns, i, .w3), + .w2 = try table(l, ns, i, .w2), + .mixer = if (kind == .conv) .{ .conv = .{ + .in = try table(l, ns, i, .conv_in), + .out = try table(l, ns, i, .conv_out), + .cw = try vec(l, ns, i, .conv), + } } else blk: { + const ring: u32 = if (kind == .sliding_attention) @intCast(cfg.ring(max_rows)) else cap; + break :blk .{ .attn = .{ + .q = try table(l, ns, i, .q_proj), + .k = try table(l, ns, i, .k_proj), + .v = try table(l, ns, i, .v_proj), + .g = try table(l, ns, i, .gate_proj), + .o = try table(l, ns, i, .o_proj), + .qn = try vec(l, ns, i, .q_norm), + .kn = try vec(l, ns, i, .k_norm), + .ring = ring, + .window = @intCast(kind.window()), + } }; + }, + }; + } + // the family's RoPE frequencies (torch's float32 values), as every backend's + const inv_full = try l.empty(HD / 2 * 4); + try r.upload(inv_full, std.mem.sliceAsBytes(&c.full_rotary.inverseFrequencies())); + const inv_slide = try l.empty(HD / 2 * 4); + try r.upload(inv_slide, std.mem.sliceAsBytes(&c.sliding_rotary.inverseFrequencies())); + try r.sync(); + var m: Model = .{ + .gpa = gpa, + .r = r, + .layers = layers, + .emb = try table(l, ns, null, .embed_tokens), + .norm_f = try vec(l, ns, null, .embedding_norm), + .inv_full = inv_full, + .inv_slide = inv_slide, + .k = k, + .s = s, + .host_rows = try gpa.alloc(u16, w * H), + .cap = cap, + .max_rows = max_rows, + .maxc = maxc, + .device_bytes = l.total, + .qg = c.quant.global.group, + }; + if (std.c.getenv("LFM3_FUSE")) |v| m.fuse = std.fmt.parseInt(u32, std.mem.span(v), 10) catch 7; + m.own = try m.newSeq(); + m.seq = m.own; + return m; + } + + /// A new stream state at position 0 (zeroed conv rows; KV slots are only read once written). + pub fn newSeq(m: *Model) !*Seq { + const q = try m.gpa.create(Seq); + errdefer m.gpa.destroy(q); + q.* = .{}; + for (m.layers, 0..) |ly, i| switch (ly.mixer) { + .conv => { + q.conv[i] = try m.r.alloc(conv_bytes); + q.ext[i] = try m.r.alloc(ext_bytes); + }, + .attn => |a| { + q.kc[i] = try m.r.alloc(@as(usize, a.ring) * slot_bytes); + q.vc[i] = try m.r.alloc(@as(usize, a.ring) * slot_bytes); + }, + }; + try m.zeroConv(q); + return q; + } + + pub fn freeSeq(m: *Model, q: *Seq) void { + if (q == m.seq) { // `own` bound again with its own position, not the freed stream's + m.seq = m.own; + m.pos = m.own.pos; + } + for (m.layers, 0..) |ly, i| switch (ly.mixer) { + .conv => { + q.conv[i].free(); + q.ext[i].free(); + }, + .attn => { + q.kc[i].free(); + q.vc[i].free(); + }, + }; + m.gpa.destroy(q); + } + + /// Makes `q` the stream the next forwards read and write. + pub fn bind(m: *Model, q: *Seq) void { + m.seq.pos = m.pos; + m.seq = q; + m.pos = q.pos; + } + + fn zeroConv(m: *Model, q: *Seq) !void { + const zero: [6 * 256]u16 = @splat(0); + for (m.layers, 0..) |ly, i| if (ly.mixer == .conv) { + var off: usize = 0; + while (off < conv_bytes) : (off += zero.len * 2) try m.r.upload(from(q.conv[i], off), std.mem.sliceAsBytes(&zero)); + }; + try m.r.sync(); + } + + /// The bound stream back to position 0. + pub fn reset(m: *Model) !void { + try m.zeroConv(m.seq); + m.pos = 0; + } + + /// Bytes of a kept state at position `at`: conv rows, `at` slots of each full layer, each sliding ring whole. + pub fn snapBytes(m: *const Model, at: u32) u64 { + var n: u64 = 0; + for (m.layers) |ly| n += switch (ly.mixer) { + .conv => conv_bytes, + .attn => |a| 2 * @as(u64, if (a.window == 0) at else a.ring) * slot_bytes, + }; + return n; + } + + /// One piece of a kept state at byte `off` of its buffer, to (save) or from it; the next piece's offset. + fn piece(m: *Model, sn: *const Snap, save: bool, b: Buf, bytes: usize, off: usize) !usize { + if (bytes > 0) { + if (save) try m.r.copy(from(sn.buf, off), b, bytes) else try m.r.copy(b, from(sn.buf, off), bytes); + } + return off + bytes; + } + + /// Copies the bound stream's state to (save) or from a kept state's buffer, on the device. + fn snapCopy(m: *Model, sn: *const Snap, save: bool) !void { + var off: usize = 0; + for (m.layers, 0..) |ly, i| switch (ly.mixer) { + .conv => off = try m.piece(sn, save, m.seq.conv[i], conv_bytes, off), + .attn => |a| { + const bytes = @as(usize, if (a.window == 0) sn.at else a.ring) * slot_bytes; + off = try m.piece(sn, save, m.seq.kc[i], bytes, off); + off = try m.piece(sn, save, m.seq.vc[i], bytes, off); + }, + }; + try m.r.sync(); + } + + /// The bound stream's state at its position, kept. + pub fn snapSave(m: *Model) !*Snap { + const sn = try m.gpa.create(Snap); + errdefer m.gpa.destroy(sn); + sn.* = .{ .at = m.pos, .buf = try m.r.alloc(@intCast(@max(m.snapBytes(m.pos), 1))) }; + try m.snapCopy(sn, true); + return sn; + } + + /// The bound stream continues from a kept state. + pub fn snapRestore(m: *Model, sn: *const Snap) !void { + m.pos = sn.at; + try m.snapCopy(sn, false); + } + + pub fn snapDrop(m: *Model, sn: *Snap) void { + sn.buf.free(); + m.gpa.destroy(sn); + } + + /// y[r * y_stride + y_off + o] = x[r] . W[o] for n rows: pgemm in a prompt window, else qmv (one row) or pdense. + fn proj(m: *Model, t: Table, x: Buf, y: Buf, in_dim: u32, rows: u32, n: u32, y_stride: u32, y_off: u32) !void { + if (m.prompt) return run(&m.k.pgemm, .{ (rows + 63) / 64, (n + 63) / 64, 1 }, .{ t.w, t.s, t.b, x, y, in_dim, rows, n, y_stride, y_off, m.qg }); + if (n == 1) return run(&m.k.qmv, .{ (rows + 3) / 4, 1, 1 }, .{ t.w, t.s, t.b, x, y, in_dim, @as(u32, 0), y_off, rows }); + try run(&m.k.pdense, .{ (rows + 15) / 16, (n + tile_rows - 1) / tile_rows, 1 }, .{ t.w, t.s, t.b, x, y, in_dim, rows, n, y_stride, y_off, @as(u32, 0) }); + } + + /// The token mixer of every segment's rows: matvecs over all rows, conv and attention per stream. + fn mixer(m: *Model, ly: *const Layer, li: usize, segs: []const Seg, n: u32, hist: bool) !void { + const s = &m.s; + switch (ly.mixer) { + .conv => |c| { + try m.proj(c.in, s.xn, s.proj, H, 3 * H, n, 3 * H, 0); + var r0: usize = 0; + for (segs) |g| { + const k: u32 = @intCast(g.ids.len); + const pr = from(s.proj, r0 * 3 * H * 2); + const y = from(s.att, r0 * H * 2); + if (hist) { + try run(&m.k.conv_hist, .{ H / 64, 1, 1 }, .{ pr, 3 * H, g.q.conv[li], g.q.conv[li], c.cw, y, k, H, g.q.ext[li] }); + } else try run(&m.k.conv, .{ H / 64, 1, 1 }, .{ pr, 3 * H, g.q.conv[li], g.q.conv[li], c.cw, y, k, H }); + r0 += k; + } + try m.proj(c.out, s.att, s.d, H, H, n, H, 0); + }, + .attn => |a| { + if (n == 1 and m.fuse & 2 != 0) { // the four decode matvecs in one launch, each row as qmv4_bf sums it + try run(&m.k.qkvg, .{ (QKV + NQ + 3) / 4, 1, 1 }, .{ a.q.w, a.q.s, a.q.b, s.qkv, NQ * HD, @as(u32, 0), a.k.w, a.k.s, a.k.b, s.qkv, NKV * HD, NQ * HD, a.v.w, a.v.s, a.v.b, s.qkv, NKV * HD, (NQ + NKV) * HD, a.g.w, a.g.s, a.g.b, s.g, NQ, @as(u32, 0), s.xn, H }); + } else { + try m.proj(a.q, s.xn, s.qkv, H, NQ * HD, n, QKV, 0); + try m.proj(a.k, s.xn, s.qkv, H, NKV * HD, n, QKV, NQ * HD); + try m.proj(a.v, s.xn, s.qkv, H, NKV * HD, n, QKV, (NQ + NKV) * HD); + try m.proj(a.g, s.xn, s.g, H, NQ, n, NQ, 0); + } + const inv = if (a.window == 0) m.inv_full else m.inv_slide; + const scale: f32 = 1.0 / @sqrt(@as(f32, HD)); + var r0: usize = 0; + for (segs) |g| { + const k: u32 = @intCast(g.ids.len); + const pos = g.q.pos; + const qkv = from(s.qkv, r0 * QKV * 2); + const q = from(s.q, r0 * NQ * HD * 2); + if (m.fuse & 4 != 0) { // q, k and v of every row in one launch + try run(&m.k.rope_store, .{ NQ + NKV + NKV * HD / 64, k, 1 }, .{ qkv, QKV, q, g.q.kc[li], g.q.vc[li], a.qn, a.kn, inv, NQ, NKV, pos, a.ring, cfg.eps }); + } else { + try run(&m.k.rope, .{ NQ, k, 1 }, .{ qkv, QKV, @as(u32, 0), q, NQ * HD, a.qn, inv, pos, @as(u32, 0), cfg.eps }); + try run(&m.k.rope, .{ NKV, k, 1 }, .{ qkv, QKV, NQ * HD, g.q.kc[li], NKV * HD, a.kn, inv, pos, a.ring, cfg.eps }); + try run(&m.k.vst, .{ NKV * HD / 64, k, 1 }, .{ qkv, QKV, (NQ + NKV) * HD, g.q.vc[li], NKV * HD, pos, a.ring }); + } + const span = if (a.window == 0) pos + k else @min(pos + k, a.window); + try run(&m.k.part, .{ NKV, (span + chunk - 1) / chunk, k }, .{ q, g.q.kc[li], g.q.vc[li], s.po, s.pm, s.pl, pos, chunk, a.ring, a.window, NKV, m.maxc, scale }); + try run(&m.k.merge, .{ NQ, k, 1 }, .{ s.po, s.pm, s.pl, from(s.g, r0 * NQ * 2), from(s.att, r0 * NQ * HD * 2), pos, chunk, a.window, NQ, m.maxc }); + r0 += k; + } + try m.proj(a.o, s.att, s.d, NQ * HD, H, n, H, 0); + }, + } + } + + fn hook(m: *Model, h: ?Hook, layer: i32, b: Buf, n: u32) !void { + const k = h orelse return; + const rows = m.host_rows[0 .. n * H]; + try m.r.sync(); + try m.r.download(std.mem.sliceAsBytes(rows), b); + try m.r.sync(); + k.f(k.ctx, layer, rows); + } + + /// The tokens as one window of the bound stream (1 .. max_rows rows); `logits`: the last row's logits and argmax. + pub fn forward(m: *Model, ids: []const u32, logits: bool, h: ?Hook) !void { + try m.forwardSegs(&.{.{ .q = m.seq, .ids = ids }}, if (logits) .last else .none, false, h); + } + + /// A prompt window of the bound stream, every projection in pgemm (verified windows and decode keep theirs). + pub fn forwardPrompt(m: *Model, ids: []const u32, logits: bool, h: ?Hook) !void { + const fuse = m.fuse; + m.prompt = true; + m.fuse = fuse & 4; // a one-row window too: not the decode matvecs + defer { + m.prompt = false; + m.fuse = fuse; + } + try m.forward(ids, logits, h); + } + + /// One forward over every segment's rows, each at its stream's position; `hist` keeps a window's rollback source. + pub fn forwardSegs(m: *Model, segs: []const Seg, out: Out, hist: bool, h: ?Hook) !void { + var n: u32 = 0; + m.seq.pos = m.pos; + for (segs) |g| { + if (g.ids.len == 0 or (hist and g.ids.len > window_rows)) return error.Invalid; + if (g.q.pos + g.ids.len > m.cap) return error.ContextFull; + n += @intCast(g.ids.len); + } + if (n > m.max_rows or (out == .rows and n > window_rows)) return error.Invalid; + const s = &m.s; + if (n == 1 or out == .rows) { // ids as kernel arguments: the queue reads no host memory the caller may reuse + var r: usize = 0; + for (segs) |g| for (g.ids) |id| { + try run(&m.k.embed_tok, .{ 1, 1, 1 }, .{ m.emb.w, m.emb.s, m.emb.b, id, from(s.x, r * H * 2), H }); + r += 1; + }; + } else { // a prompt window: its ids live as long as its stream + try m.r.upload(s.ids, std.mem.sliceAsBytes(segs[0].ids)); + try run(&m.k.embed, .{ n, 1, 1 }, .{ m.emb.w, m.emb.s, m.emb.b, s.ids, s.x, H }); + } + try run(&m.k.rms, .{ n, 1, 1 }, .{ s.x, m.layers[0].op_norm, s.xn, H, cfg.eps }); + for (m.layers, 0..) |*ly, i| { + try m.mixer(ly, i, segs, n, hist); + try run(&m.k.pan, .{ n, 1, 1 }, .{ s.d, s.x, ly.post_op, ly.ffn_norm, s.xn, H, cfg.eps }); + if (n == 1 and m.fuse & 1 != 0) { // w1, w3 and their SwiGLU in one launch, each value as the three launches round it + try run(&m.k.w13, .{ FF / 4, 1, 1 }, .{ ly.w1.w, ly.w1.s, ly.w1.b, ly.w3.w, ly.w3.s, ly.w3.b, s.xn, s.act, H, FF }); + } else { + try m.proj(ly.w1, s.xn, s.a, H, FF, n, FF, 0); + try m.proj(ly.w3, s.xn, s.b, H, FF, n, FF, 0); + try run(&m.k.swiglu, .{ (n * FF + 63) / 64, 1, 1 }, .{ s.a, s.b, s.act, n * FF }); + } + try m.proj(ly.w2, s.act, s.d, FF, H, n, H, 0); + const next = if (i + 1 < m.layers.len) m.layers[i + 1].op_norm else m.norm_f; + try run(&m.k.pan, .{ n, 1, 1 }, .{ s.d, s.x, ly.post_ffn, next, s.xn, H, cfg.eps }); + try m.hook(h, @intCast(i), s.x, n); + } + try m.hook(h, -1, s.xn, n); + for (segs) |g| { + if (hist) { + g.q.win_start = g.q.pos; + g.q.win_rows = @intCast(g.ids.len); + } + g.q.pos += @intCast(g.ids.len); + } + m.pos = m.seq.pos; + const per = (V + argmax_parts - 1) / argmax_parts; + switch (out) { + .none => {}, + .last => { + try run(&m.k.head, .{ V / 4, 1, 1 }, .{ m.emb.w, m.emb.s, m.emb.b, s.xn, s.logits, H, (n - 1) * H, @as(u32, 0), V }); + try run(&m.k.am_part, .{ argmax_parts, 1, 1 }, .{ s.logits, s.am_scratch, V, per }); + try run(&m.k.am_fin, .{ 1, 1, 1 }, .{ s.am_scratch, s.am_out, argmax_parts }); + }, + .rows => { + // every row's logits: pdense's fp32 epilogue sums each as qmv4_f32 does (tf-xpu-lfm3_rows-test) + try run(&m.k.pdense, .{ V / 16, (n + tile_rows - 1) / tile_rows, 1 }, .{ m.emb.w, m.emb.s, m.emb.b, s.xn, s.logits, H, V, n, V, @as(u32, 0), @as(u32, 1) }); + try run(&m.k.am_part, .{ argmax_parts, n, 1 }, .{ s.logits, s.am_scratch, V, per }); + try run(&m.k.am_fin, .{ n, 1, 1 }, .{ s.am_scratch, s.am_out, argmax_parts }); + }, + } + } + + /// Keeps the first c rows of the bound stream's last verified window: its position, and its conv rows from ext. + pub fn keepRows(m: *Model, c: u32) !void { + const q = m.seq; + if (q.win_rows == 0) return; + if (c > q.win_rows) return error.Invalid; + if (c < q.win_rows) for (m.layers, 0..) |ly, i| if (ly.mixer == .conv) try m.r.copy(q.conv[i], from(q.ext[i], c * H * 2), conv_bytes); + m.pos = q.win_start + c; + q.win_rows = 0; + } + + /// The greedy tokens of the last forward's rows (Out.rows), in row order. + pub fn rowTokens(m: *Model, out: []i32) !void { + try m.r.sync(); + try m.r.download(std.mem.sliceAsBytes(out), m.s.am_out); + try m.r.sync(); + } + + /// The logits of the last forward's rows (Out.rows), [rows][vocab]. + pub fn rowLogits(m: *Model, out: []f32) !void { + try m.r.sync(); + try m.r.download(std.mem.sliceAsBytes(out), m.s.logits); + try m.r.sync(); + } + + /// Prefill of a prompt in windows of max_rows; the last window leaves its logits and argmax. + pub fn prefill(m: *Model, ids: []const u32, h: ?Hook) !void { + var i: usize = 0; + while (i < ids.len) { + const n = @min(ids.len - i, m.max_rows); + try m.forwardPrompt(ids[i .. i + n], i + n == ids.len, h); + i += n; + } + } + + /// The greedy token of the last forward with logits. + pub fn argmax(m: *Model) !u32 { + var out: [1]i32 = undefined; + try m.r.sync(); + try m.r.download(std.mem.sliceAsBytes(&out), m.s.am_out); + try m.r.sync(); + return @intCast(out[0]); + } + + pub fn fetchLogits(m: *Model, out: []f32) !void { + try m.r.sync(); + try m.r.download(std.mem.sliceAsBytes(out[0..V]), m.s.logits); + try m.r.sync(); + } +}; diff --git a/zig/src/native/xpu.zig b/zig/src/native/xpu.zig index 935793b2a8..cce64198d1 100644 --- a/zig/src/native/xpu.zig +++ b/zig/src/native/xpu.zig @@ -1,7 +1,9 @@ //! XPU engines (Intel Xe2 over Level Zero): Qwen3.8-27B behind the lane host; a lone stream uses the backend's driver. const std = @import("std"); const api = @import("engine_api"); -const xq = @import("xpu_qwen27"); +/// The Qwen3.8-27B runner needs DPAS and nGEN: a build without DPAS (-Dxpu-dpas=false, Gen12LP) leaves it out. +const has_qwen27 = @import("xpu_engine_options").qwen27; +const xq = if (has_qwen27) @import("xpu_qwen27") else struct {}; const Drafter = xq.drafter.Drafter; const lanes = @import("lanes"); const sampling = lanes.sampling; @@ -11,7 +13,9 @@ const admit = @import("xpu_admit.zig"); const Allocator = std.mem.Allocator; pub const backends: []const []const u8 = &.{"xpu"}; -pub const families: []const api.Family = &.{ .{ .model_type = "qwen3_5", .formats = &.{ "mlx-q4g64", "gptq-int4-g128" } }, .{ .model_type = "nemotron_h", .formats = &.{"mlx-q4g64"} }, .{ .model_type = "gemma4", .formats = &.{"mlx-q4g64"} } }; +const qwen27_family: []const api.Family = if (has_qwen27) &.{.{ .model_type = "qwen3_5", .formats = &.{ "mlx-q4g64", "gptq-int4-g128" } }} else &.{}; +const gemma4_family: []const api.Family = if (nem.has_gemma4) &.{.{ .model_type = "gemma4", .formats = &.{"mlx-q4g64"} }} else &.{}; +pub const families: []const api.Family = qwen27_family ++ &[_]api.Family{ .{ .model_type = "nemotron_h", .formats = &.{"mlx-q4g64"} }, .{ .model_type = "lfm3", .formats = &.{"mlx-q4g64"} } } ++ gemma4_family; pub fn chip(_: Allocator) ?[]const u8 { return "Intel Xe2 (Level Zero)"; @@ -175,10 +179,14 @@ pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]c return null; } if (nem.serves(o.model_type)) return nem.open(a, gpa, io, o, problem); - if (!std.mem.eql(u8, o.model_type, "qwen3_5")) { - problem.* = try std.fmt.allocPrint(a, "the XPU engine serves Qwen3.8-27B (qwen3_5, MLX affine 4/64 or symmetric GPTQ 4/128) and Nemotron-H (nemotron_h, MLX 4/64), not {s}", .{o.model_type}); - return null; - } + if (comptime has_qwen27) { + if (std.mem.eql(u8, o.model_type, "qwen3_5")) return openQwen27(a, gpa, io, o, problem); + problem.* = try std.fmt.allocPrint(a, "the XPU engine serves Qwen3.8-27B (qwen3_5, MLX affine 4/64 or symmetric GPTQ 4/128), Nemotron-H (nemotron_h, MLX 4/64) and LFM3 (lfm3, MLX 4/64), not {s}", .{o.model_type}); + } else problem.* = try std.fmt.allocPrint(a, "this XPU build (no DPAS) serves Nemotron-H (nemotron_h, MLX 4/64) and LFM3 (lfm3, MLX 4/64), not {s}: Qwen3.8-27B and Gemma 4 E4B need DPAS (-Dxpu-dpas=true)", .{o.model_type}); + return null; +} + +fn openQwen27(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !?api.Opened { const context = xq.limits.context(o.context orelse 32768) catch { problem.* = try std.fmt.allocPrint(a, "the XPU engine holds 1 to {d} tokens a stream, not {d}", .{ xq.limits.MAX_CONTEXT, o.context.? }); return null; diff --git a/zig/src/native/xpu_lfm3.zig b/zig/src/native/xpu_lfm3.zig new file mode 100644 index 0000000000..f0fb4ec536 --- /dev/null +++ b/zig/src/native/xpu_lfm3.zig @@ -0,0 +1,40 @@ +//! The LFM3 entry of the Intel GPU lane host: its family table row, loader and prompt window. +const std = @import("std"); +const xpu = @import("xpu"); +const api = @import("engine_api"); +const lfm3 = @import("lfm3_xpu"); +const Allocator = std.mem.Allocator; + +pub const model_type = "lfm3"; +pub const formats: []const []const u8 = &.{"mlx-q4g64"}; +pub const name = "xpu-lfm3"; +pub const Engine = lfm3.engine.Engine; + +/// Windows cut at any row give the same bits: states are kept and resumed anywhere, short prompts too. +pub const cache_rules: api.prompt_cache.Rules = .{ .min_prompt = 0, .min_gap = 64 }; + +/// Copies drafted after an 8-token match (the core's 4 lost 6% on a short chat on Iris Xe; XPU_COPY_MATCH overrides). +pub const copy_match: i64 = 8; + +/// Rows of one prompt window (pdense above one row; the scratch is sized for it): 128 beat 32 and 64 on Iris Xe. +const prompt_rows: u32 = 128; + +pub fn promptRows(_: *const Engine) u32 { + return prompt_rows; +} + +/// What the startup line adds about the engine's settings. +pub fn note(a: Allocator, e: *const Engine) ![]const u8 { + return std.fmt.allocPrint(a, ", prompt windows of {d} rows, copy drafts of the lane core (no head)", .{e.window_rows}); +} + +/// The engine for `o.dir`, or null with `problem` set (the caller releases the device). +pub fn init(a: Allocator, io: std.Io, ctx: *const xpu.Context, o: api.Open, window: usize, problem: *[]const u8) !?*Engine { + return Engine.init(std.heap.page_allocator, io, ctx, o.dir, .{ .context = window, .window_rows = prompt_rows }) catch |e| { + problem.* = if (e == error.UnsupportedFormat) + try std.fmt.allocPrint(a, "{s}: the Intel GPU engine reads LFM3 text checkpoints in MLX affine 4-bit, group 64", .{o.dir}) + else + try std.fmt.allocPrint(a, "the Intel GPU engine cannot load {s} ({s})", .{ o.dir, @errorName(e) }); + return null; + }; +} diff --git a/zig/src/native/xpu_nem.zig b/zig/src/native/xpu_nem.zig index c14f29fcb5..3d5ae12fe8 100644 --- a/zig/src/native/xpu_nem.zig +++ b/zig/src/native/xpu_nem.zig @@ -1,4 +1,4 @@ -//! Nemotron on the XPU server: its lane host (one family behind the lane core, one stream, host sampling). +//! Nemotron and LFM3 on the XPU server: their lane host (a family behind the lane core, one stream, host sampling). const std = @import("std"); const xpu = @import("xpu"); const api = @import("engine_api"); @@ -6,11 +6,11 @@ const lanes = @import("lanes"); const lane = @import("xpu_lane.zig"); const Allocator = std.mem.Allocator; +/// Gemma 4 E4B's GEMMs are nGEN DPAS kernels: a build without DPAS (-Dxpu-dpas=false, Gen12LP) leaves it out. +pub const has_gemma4 = @import("xpu_engine_options").gemma4; + /// The families this host serves, each a file with its table row, loader and prompt window. -const specs = .{ - @import("xpu_nemotron.zig"), - @import("xpu_gemma4.zig"), -}; +const specs = if (has_gemma4) .{ @import("xpu_nemotron.zig"), @import("xpu_gemma4.zig"), @import("xpu_lfm3.zig") } else .{ @import("xpu_nemotron.zig"), @import("xpu_lfm3.zig") }; /// Prompt plus reply tokens when --context names none (and the model's window is larger). const default_context: i64 = 32768; @@ -146,6 +146,7 @@ fn openFamily(comptime S: type, a: Allocator, gpa: Allocator, io: std.Io, o: api h.backend.started = chainStopHandlers; try h.backend.measure(); if (@hasDecl(S.Engine, "copy_only")) { + if (@hasDecl(S, "copy_match")) h.backend.copy_match = S.copy_match; if (getenv("XPU_COPY_MATCH")) |v| h.backend.copy_match = std.fmt.parseInt(i64, v, 10) catch 4; h.backend.streams = @min(@max(1, o.lanes), 8); // one sequence each (the engine holds up to 8) } else if (h.backend.facts().exact_width > 1) h.backend.streams = @max(1, o.lanes); // drafted windows: up to --parallel streams a round diff --git a/zig/src/xpu/kernels.zig b/zig/src/xpu/kernels.zig index 14b3df7017..93178cba5e 100644 --- a/zig/src/xpu/kernels.zig +++ b/zig/src/xpu/kernels.zig @@ -24,4 +24,8 @@ pub const nem_rows: []const u8 = if (available) &Blob("spv_nem_rows").bytes else /// nem_rows built for windows of at most 1, 2, 4 and 8 rows (fewer registers a lane). pub const nem_rows_r: [4][]const u8 = if (available) .{ &Blob("spv_nem_rows_r1").bytes, &Blob("spv_nem_rows_r2").bytes, &Blob("spv_nem_rows_r4").bytes, &Blob("spv_nem_rows_r8").bytes } else .{ &.{}, &.{}, &.{}, &.{} }; pub const nem_prow: []const u8 = if (available) &Blob("spv_nem_prow").bytes else &.{}; +/// nem_prow with 4 outputs a sub-group (8 x 4 tiles): fastest of PR, PO in {2, 4, 8} on Gen12LP, where 8 x 8 spills. +pub const nem_prow_gen12: []const u8 = if (available) &Blob("spv_nem_prow_gen12").bytes else &.{}; pub const nem_attn_dec: []const u8 = if (available) &Blob("spv_nem_attn_dec").bytes else &.{}; +pub const lfm3: []const u8 = if (available) &Blob("spv_lfm3").bytes else &.{}; +pub const lfm3_mv: []const u8 = if (available) &Blob("spv_lfm3_mv").bytes else &.{}; diff --git a/zig/tests/xpu/lfm3_attn.zig b/zig/tests/xpu/lfm3_attn.zig new file mode 100644 index 0000000000..1bfad3d6ee --- /dev/null +++ b/zig/tests/xpu/lfm3_attn.zig @@ -0,0 +1,124 @@ +//! LFM3 GQA attention (partial + merge with gate) over a KV ring, full and sliding, against FP64 on random data. + +const std = @import("std"); +const o = @import("lfm3_ops.zig"); +const stop = @import("xpu").stop; + +const HD: usize = 128; +const NQ: u32 = 16; +const NKV: u32 = 4; +const CHUNK: u32 = 64; + +/// FP64 attention of row z (position pos0 + z) over keys by absolute position, window 0 = all, times sigmoid(gate). +fn refRows(q: []const f64, k: []const f64, v: []const f64, g: []const f64, n: usize, pos0: usize, window: usize, out: []f64) !void { + const scale = 1.0 / @sqrt(@as(f64, HD)); + const s = try o.gpa.alloc(f64, pos0 + n); + defer o.gpa.free(s); + for (0..n) |z| { + const p = pos0 + z; + const lo = if (window == 0 or p + 1 <= window) 0 else p + 1 - window; + for (0..NQ) |h| { + const kvh = h / (NQ / NKV); + const qh = q[(z * NQ + h) * HD ..][0..HD]; + var mx: f64 = -std.math.inf(f64); + for (lo..p + 1) |key| { + var a: f64 = 0; + for (0..HD) |d| a += qh[d] * k[(key * NKV + kvh) * HD + d]; + s[key] = a * scale; + mx = @max(mx, s[key]); + } + var den: f64 = 0; + const oh = out[(z * NQ + h) * HD ..][0..HD]; + @memset(oh, 0); + for (lo..p + 1) |key| { + const e = @exp(s[key] - mx); + den += e; + for (0..HD) |d| oh[d] += e * v[(key * NKV + kvh) * HD + d]; + } + const gate = 1.0 / (1.0 + @exp(-g[z * NQ + h])); + for (oh) |*x| x.* = x.* / den * gate; + } + } +} + +const Case = struct { what: []const u8, pos0: u32, n: u32, window: u32, ring: u32 }; + +pub const panic = std.debug.FullPanic(@import("rig.zig").panicDrain); +pub fn main() !void { + stop.install(); + var r = try o.rt.open(); + defer r.deinit(); + var m = try r.module(@import("xpu").kernels.lfm3); + var kp = try m.kernel("lfm3_attn_partial", .{ 64, 1, 1 }); + var km = try m.kernel("lfm3_attn_merge_gate", .{ 16, 1, 1 }); + var t: o.Tally = .{ .name = "lfm3_attn" }; + var prng = std.Random.DefaultPrng.init(13); + const rnd = prng.random(); + const maxc: u32 = 64; + const max_rows: u32 = 8; + const po = try r.alloc(@as(usize, max_rows) * NQ * maxc * HD * 4); + const pm = try r.alloc(@as(usize, max_rows) * NQ * maxc * 4); + const pl = try r.alloc(@as(usize, max_rows) * NQ * maxc * 4); + std.debug.print("GQA attention, 16 q / 4 kv heads, chunks of {d} keys, sigmoid gate per head, against FP64\n", .{CHUNK}); + + const cases = [_]Case{ + .{ .what = "full layer, 6 rows at 700 (12 chunks)", .pos0 = 700, .n = 6, .window = 0, .ring = 4096 }, + .{ .what = "full layer, decode row at 0 (one key)", .pos0 = 0, .n = 1, .window = 0, .ring = 4096 }, + .{ .what = "sliding 256, 8 rows at 900 in a 512-slot ring (wraps)", .pos0 = 900, .n = 8, .window = 256, .ring = 512 }, + .{ .what = "sliding 256, 3 rows at 100 (shorter than the window)", .pos0 = 100, .n = 3, .window = 256, .ring = 512 }, + }; + for (cases, 0..) |c, ci| { + const npos: usize = c.pos0 + c.n; + const kh = try o.gpa.alloc(u16, npos * NKV * HD); + const vh = try o.gpa.alloc(u16, kh.len); + o.randBf(kh, rnd, 1.0); + o.randBf(vh, rnd, 1.0); + const qh = try o.gpa.alloc(u16, c.n * NQ * HD); + o.randBf(qh, rnd, 1.0); + const gh = try o.gpa.alloc(u16, c.n * NQ); + o.randBf(gh, rnd, 2.0); + const planted = c.window != 0 and c.pos0 >= c.window; + if (planted) { // a key just outside each row's window, aligned with its query: any row that reads it is dominated by it + for (0..c.n) |z| { + const key = c.pos0 + z - c.window; + for (0..NKV) |kvh| for (0..HD) |d| { + kh[(key * NKV + kvh) * HD + d] = o.toBf(3.0 * o.bf(qh[(z * NQ + kvh * (NQ / NKV)) * HD + d])); + vh[(key * NKV + kvh) * HD + d] = o.toBf(8.0); + }; + } + } + // the device ring holds position p at slot p % ring; only the positions a row may read must survive + const ring_k = try o.gpa.alloc(u16, @as(usize, c.ring) * NKV * HD); + const ring_v = try o.gpa.alloc(u16, ring_k.len); + @memset(ring_k, 0); + @memset(ring_v, 0); + const first = if (npos > c.ring) npos - c.ring else 0; + for (first..npos) |p| { + const slot = p % c.ring; + @memcpy(ring_k[slot * NKV * HD ..][0 .. NKV * HD], kh[p * NKV * HD ..][0 .. NKV * HD]); + @memcpy(ring_v[slot * NKV * HD ..][0 .. NKV * HD], vh[p * NKV * HD ..][0 .. NKV * HD]); + } + const kc = try o.upload(&r, ring_k); + const vc = try o.upload(&r, ring_v); + const q = try o.upload(&r, qh); + const g = try o.upload(&r, gh); + const out = try r.alloc(c.n * NQ * HD * 2); + const span = if (c.window == 0) c.pos0 + c.n else @min(c.pos0 + c.n, c.window); + const nch = (span + CHUNK - 1) / CHUNK; + try o.run(&kp, .{ NKV, nch, c.n }, .{ q, kc, vc, po, pm, pl, c.pos0, CHUNK, c.ring, c.window, NKV, maxc, @as(f32, 1.0 / @sqrt(128.0)) }); + try o.run(&km, .{ NQ, c.n, 1 }, .{ po, pm, pl, g, out, c.pos0, CHUNK, c.window, NQ, maxc }); + const oh = try o.fetch(&r, out, c.n * NQ * HD); + const qf = try o.widen(o.gpa, qh); + const kf = try o.widen(o.gpa, kh); + const vf = try o.widen(o.gpa, vh); + const gf = try o.widen(o.gpa, gh); + const want = try o.gpa.alloc(f64, oh.len); + try refRows(qf, kf, vf, gf, c.n, c.pos0, c.window, want); + t.check(c.what, o.relErr(oh, want)); + if (ci == 2) { // planted control: a window of 257 reads the planted key and must not match + try refRows(qf, kf, vf, gf, c.n, c.pos0, c.window + 1, want); + t.control("sliding layer replayed with a window of 257", o.relErr(oh, want)); + } + } + try t.finish(); +} diff --git a/zig/tests/xpu/lfm3_bench.zig b/zig/tests/xpu/lfm3_bench.zig new file mode 100644 index 0000000000..8f40ebbcea --- /dev/null +++ b/zig/tests/xpu/lfm3_bench.zig @@ -0,0 +1,55 @@ +//! LFM3 prefill and decode timing on the device: a prompt of N random ids in windows of W rows, then 64 greedy tokens. + +const std = @import("std"); +const rt = @import("rig.zig"); +const lfm3 = @import("lfm3_xpu"); +const stop = @import("xpu").stop; + +fn nowNs() u64 { + var ts: std.c.timespec = undefined; + _ = std.c.clock_gettime(.MONOTONIC, &ts); + return @as(u64, @intCast(ts.sec)) * 1_000_000_000 + @as(u64, @intCast(ts.nsec)); +} + +pub const panic = std.debug.FullPanic(@import("rig.zig").panicDrain); +pub fn main(init: std.process.Init) !void { + stop.install(); + const gpa = init.gpa; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len < 2) { + std.debug.print("skipped: usage tf-xpu-lfm3_bench-test MODEL_DIR [PROMPT_TOKENS] [WINDOW_ROWS]\n", .{}); + return; + } + const n: usize = if (args.len > 2) try std.fmt.parseInt(usize, args[2], 10) else 1024; + const w: u32 = if (args.len > 3) try std.fmt.parseInt(u32, args[3], 10) else 64; + var r = try rt.open(); + defer r.deinit(); + var c = try lfm3.config.Config.read(gpa, init.io, args[1]); + defer c.deinit(); + var l = try @import("xpu").loader.Loader.init(gpa, &r, args[1]); + var m = try lfm3.model.Model.load(gpa, &r, &l, &c, @intCast(n + 128), w); + var prng = std.Random.DefaultPrng.init(5); + const ids = try gpa.alloc(u32, n); + for (ids) |*v| v.* = prng.random().uintLessThan(u32, 120000); + try m.reset(); + var t0 = nowNs(); + try m.prefill(ids, null); + var tok = try m.argmax(); + const pre = @as(f64, @floatFromInt(nowNs() - t0)) / 1e9; + t0 = nowNs(); + for (0..64) |_| { + try m.forward(&.{tok}, true, null); + tok = try m.argmax(); + } + const dec = @as(f64, @floatFromInt(nowNs() - t0)) / 1e9; + if (std.c.getenv("LFM3_PROF") != null) { // 16 more steps with a sync around every launch: device time by kernel + @import("xpu").rt.prof_on = true; + for (0..16) |_| { + try m.forward(&.{tok}, true, null); + tok = try m.argmax(); + } + @import("xpu").rt.prof_on = false; + @import("xpu").rt.profDump("16 decode steps"); + } + std.debug.print("prompt {d} tokens in windows of {d}: prefill {d:.2} s ({d:.1} tok/s); decode after it {d:.1} tok/s; device {d:.0} MiB\n", .{ n, w, pre, @as(f64, @floatFromInt(n)) / pre, 64.0 / dec, @as(f64, @floatFromInt(m.device_bytes)) / (1 << 20) }); +} diff --git a/zig/tests/xpu/lfm3_chunks.zig b/zig/tests/xpu/lfm3_chunks.zig new file mode 100644 index 0000000000..8c27c736c0 --- /dev/null +++ b/zig/tests/xpu/lfm3_chunks.zig @@ -0,0 +1,93 @@ +//! LFM3 prompt windows of 64, 128, 256 (the most) rows and a ragged cut: rows, logits and the next step, bit for bit. + +const std = @import("std"); +const rt = @import("rig.zig"); +const lfm3 = @import("lfm3_xpu"); +const stop = @import("xpu").stop; +const Model = lfm3.model.Model; +const H = lfm3.config.hidden; +const V = lfm3.config.vocab; + +const Rows = struct { + out: []u16, + at: usize = 0, + fn take(ctx: *anyopaque, layer: i32, rows: []const u16) void { + if (layer != -1) return; + const r: *Rows = @ptrCast(@alignCast(ctx)); + @memcpy(r.out[r.at..][0..rows.len], rows); + r.at += rows.len; + } +}; + +/// The prompt in windows of `cut` (cycled): each row's final norm, the last logits, then one decode step's logits. +fn run(m: *Model, ids: []const u32, cut: []const usize, prompt: bool, rows: []u16, last: []f32, step: []f32) !u32 { + var cap: Rows = .{ .out = rows }; + try m.reset(); + var i: usize = 0; + var c: usize = 0; + while (i < ids.len) : (c += 1) { + const n = @min(cut[c % cut.len], ids.len - i); + const hk: lfm3.model.Hook = .{ .ctx = &cap, .f = Rows.take }; + if (prompt) try m.forwardPrompt(ids[i .. i + n], i + n == ids.len, hk) else try m.forward(ids[i .. i + n], i + n == ids.len, hk); + i += n; + } + try m.fetchLogits(last); + const tok = try m.argmax(); + try m.forward(&.{tok}, true, null); + try m.fetchLogits(step); + return tok; +} + +pub const panic = std.debug.FullPanic(@import("rig.zig").panicDrain); +pub fn main(init: std.process.Init) !void { + stop.install(); + const gpa = init.gpa; + const io = init.io; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len < 3) { + std.debug.print("skipped: usage tf-xpu-lfm3_chunks-test MODEL_DIR REF_DIR (its items prompt, repeated to 2,100 tokens)\n", .{}); + return; + } + var r = try rt.open(); + defer r.deinit(); + var c = try lfm3.config.Config.read(gpa, io, args[1]); + defer c.deinit(); + var l = try @import("xpu").loader.Loader.init(gpa, &r, args[1]); + var m = try Model.load(gpa, &r, &l, &c, 2304, 256); + const path = try std.fmt.allocPrint(gpa, "{s}/items/meta.json", .{args[2]}); + const meta = try std.json.parseFromSlice(std.json.Value, gpa, try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .unlimited), .{}); + const vals = meta.value.object.get("ids").?.array.items; + const ids = try gpa.alloc(u32, 2100); + for (ids, 0..) |*o, i| o.* = @intCast(vals[if (i < vals.len) i else 1 + (i - 1) % (vals.len - 1)].integer); + const cuts = [_]struct { name: []const u8, cut: []const usize, prompt: bool = true }{ + .{ .name = "windows of 128", .cut = &.{128} }, + .{ .name = "windows of 64", .cut = &.{64} }, + .{ .name = "windows of 256", .cut = &.{256} }, + .{ .name = "ragged (1, 7, 100, 3, 250, 64, 1)", .cut = &.{ 1, 7, 100, 3, 250, 64, 1 } }, + .{ .name = "planted control: windows of 128 in pdense's order (must differ)", .cut = &.{128}, .prompt = false }, + }; + var ref_rows: []u16 = &.{}; + var ref_last: []f32 = &.{}; + var ref_step: []f32 = &.{}; + var ref_tok: u32 = 0; + var bad = false; + for (cuts, 0..) |ct, k| { + const rows = try gpa.alloc(u16, ids.len * H); + const last = try gpa.alloc(f32, V); + const step = try gpa.alloc(f32, V); + const tok = try run(&m, ids, ct.cut, ct.prompt, rows, last, step); + if (k == 0) { + ref_rows = rows; + ref_last = last; + ref_step = step; + ref_tok = tok; + std.debug.print(" {s}: the reference ({d} rows, next token {d})\n", .{ ct.name, ids.len, tok }); + continue; + } + const same = std.mem.eql(u16, rows, ref_rows) and std.mem.eql(f32, last, ref_last) and std.mem.eql(f32, step, ref_step) and tok == ref_tok; + if (same != ct.prompt) bad = true; + std.debug.print(" {s}: every row, the logits and the decode step after {s}\n", .{ ct.name, if (same) "bit for bit equal" else "differ" }); + } + std.debug.print("{s}\n", .{if (bad) "lfm3_chunks: FAILED" else "lfm3_chunks: ok"}); + if (bad) return error.TestFailed; +} diff --git a/zig/tests/xpu/lfm3_conv.zig b/zig/tests/xpu/lfm3_conv.zig new file mode 100644 index 0000000000..bbfe076e9f --- /dev/null +++ b/zig/tests/xpu/lfm3_conv.zig @@ -0,0 +1,92 @@ +//! LFM3 gated short conv (lfm3_conv_rows) against an FP64 replay on random data: a prompt window, then decode steps. + +const std = @import("std"); +const o = @import("lfm3_ops.zig"); +const stop = @import("xpu").stop; + +const D: u32 = 2048; +const PD: u32 = 3 * D; + +/// FP64 replay of rows of [B | C | x]: y = C * sum_k hist[k] * w[k] over the last 7 B * x; taps reversed when `flip`. +fn refConv(proj: []const f64, w: []const f64, hist: []f64, n: usize, flip: bool, y: []f64) void { + for (0..D) |c| { + for (0..n) |r| { + const p = proj[r * PD ..]; + const bx = p[c] * p[2 * D + c]; + var acc: f64 = 0; + for (0..6) |k| acc += hist[k * D + c] * w[c * 7 + (if (flip) 6 - k else k)]; + acc += bx * w[c * 7 + @as(usize, if (flip) 0 else 6)]; + y[r * D + c] = p[D + c] * acc; + for (0..5) |k| hist[k * D + c] = hist[(k + 1) * D + c]; + hist[5 * D + c] = bx; + } + } +} + +pub const panic = std.debug.FullPanic(@import("rig.zig").panicDrain); +pub fn main() !void { + stop.install(); + var r = try o.rt.open(); + defer r.deinit(); + var m = try r.module(@import("xpu").kernels.lfm3); + var k = try m.kernel("lfm3_conv_rows", .{ 64, 1, 1 }); + var t: o.Tally = .{ .name = "lfm3_conv" }; + var prng = std.Random.DefaultPrng.init(11); + const rnd = prng.random(); + + const n_pre: u32 = 37; + const n_dec: u32 = 5; + const rows = n_pre + n_dec; + const ph = try o.gpa.alloc(u16, rows * PD); + o.randBf(ph, rnd, 1.0); + const wh = try o.gpa.alloc(u16, D * 7); + o.randBf(wh, rnd, 0.3); + const proj = try o.upload(&r, ph); + const cw = try o.upload(&r, wh); + const zero = try o.gpa.alloc(u16, 6 * D); + @memset(zero, 0); + const st = try o.upload(&r, zero); + const y = try r.alloc(rows * D * 2); + + // the prompt window from an empty state, then single decode steps on the state it leaves (state in place) + try o.run(&k, .{ D / 64, 1, 1 }, .{ proj, PD, st, st, cw, y, n_pre, D }); + for (0..n_dec) |i| { + const off = (n_pre + i) * PD * 2; + const pr: o.Buf = .{ .rt = proj.rt, .ptr = @ptrFromInt(@intFromPtr(proj.ptr.?) + off), .len = proj.len - off }; + const yo = (n_pre + i) * D * 2; + const yr: o.Buf = .{ .rt = y.rt, .ptr = @ptrFromInt(@intFromPtr(y.ptr.?) + yo), .len = y.len - yo }; + try o.run(&k, .{ D / 64, 1, 1 }, .{ pr, PD, st, st, cw, yr, @as(u32, 1), D }); + } + const yh = try o.fetch(&r, y, rows * D); + const sh = try o.fetch(&r, st, 6 * D); + + const pf = try o.widen(o.gpa, ph); + const wf = try o.widen(o.gpa, wh); + const hist = try o.gpa.alloc(f64, 6 * D); + const yf = try o.gpa.alloc(f64, rows * D); + @memset(hist, 0); + refConv(pf, wf, hist, rows, false, yf); + std.debug.print("gated short conv, D {d}: {d} prompt rows then {d} decode steps, against FP64\n", .{ D, n_pre, n_dec }); + t.check("prompt rows y", o.relErr(yh[0 .. n_pre * D], yf[0 .. n_pre * D])); + t.check("decode steps y", o.relErr(yh[n_pre * D ..], yf[n_pre * D ..])); + t.check("state after the last step (last six B*x)", o.relErr(sh, hist)); + + // the same rows one at a time: the state handoff gives the window's bits + const st1 = try o.upload(&r, zero); + const y1 = try r.alloc(rows * D * 2); + for (0..rows) |i| { + const off = i * PD * 2; + const pr: o.Buf = .{ .rt = proj.rt, .ptr = @ptrFromInt(@intFromPtr(proj.ptr.?) + off), .len = proj.len - off }; + const yo = i * D * 2; + const yr: o.Buf = .{ .rt = y1.rt, .ptr = @ptrFromInt(@intFromPtr(y1.ptr.?) + yo), .len = y1.len - yo }; + try o.run(&k, .{ D / 64, 1, 1 }, .{ pr, PD, st1, st1, cw, yr, @as(u32, 1), D }); + } + const y1h = try o.fetch(&r, y1, rows * D); + t.same("all rows stepped one at a time against the window", std.mem.eql(u16, y1h, yh)); + + // planted control: the taps in reverse order (oldest input on the newest tap) must not match + @memset(hist, 0); + refConv(pf, wf, hist, rows, true, yf); + t.control("reversed taps", o.relErr(yh, yf)); + try t.finish(); +} diff --git a/zig/tests/xpu/lfm3_fwd.zig b/zig/tests/xpu/lfm3_fwd.zig new file mode 100644 index 0000000000..4aa1fdfc11 --- /dev/null +++ b/zig/tests/xpu/lfm3_fwd.zig @@ -0,0 +1,163 @@ +//! LFM3 forward on the checkpoint: every layer against ref64.py's float64 dumps, then 64 greedy tokens, twice. + +const std = @import("std"); +const rt = @import("rig.zig"); +const lfm3 = @import("lfm3_xpu"); +const stop = @import("xpu").stop; + +const H = 2048; +const prompts = [_][]const u8{ "capitals", "chat", "items" }; +const reply_len = 64; + +fn nowNs() u64 { + var ts: std.c.timespec = undefined; + _ = std.c.clock_gettime(.MONOTONIC, &ts); + return @as(u64, @intCast(ts.sec)) * 1_000_000_000 + @as(u64, @intCast(ts.nsec)); +} + +fn bf(v: u16) f64 { + return @as(f32, @bitCast(@as(u32, v) << 16)); +} + +/// The hidden rows of every layer (index 24: the final norm) as the windows pass, bf16. +const Rows = struct { + layers: [25][]u16, + row0: usize = 0, + + fn keep(ctx: *anyopaque, layer: i32, rows: []const u16) void { + const s: *Rows = @ptrCast(@alignCast(ctx)); + const li: usize = if (layer < 0) 24 else @intCast(layer); + @memcpy(s.layers[li][s.row0 * H ..][0..rows.len], rows); + } +}; + +/// Per-row ||xpu - ref|| / ||ref||: mean and max over the rows. +fn relRows(got: []const u16, want: []const f32) [2]f64 { + const n = want.len / H; + var sum: f64 = 0; + var mx: f64 = 0; + for (0..n) |r| { + var d: f64 = 0; + var w: f64 = 0; + for (0..H) |i| { + const a = bf(got[r * H + i]); + const b: f64 = want[r * H + i]; + d += (a - b) * (a - b); + w += b * b; + } + const e = @sqrt(d / w); + sum += e; + mx = @max(mx, e); + } + return .{ sum / @as(f64, @floatFromInt(n)), mx }; +} + +fn readF32(gpa: std.mem.Allocator, io: std.Io, path: []const u8) ![]f32 { + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .unlimited); + const out = try gpa.alloc(f32, bytes.len / 4); + @memcpy(std.mem.sliceAsBytes(out), bytes); + gpa.free(bytes); + return out; +} + +pub const panic = std.debug.FullPanic(@import("rig.zig").panicDrain); +pub fn main(init: std.process.Init) !void { + stop.install(); + const gpa = init.gpa; + const io = init.io; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len < 3) { + std.debug.print("skipped: usage tf-xpu-lfm3_fwd-test MODEL_DIR REF_DIR [WINDOW_ROWS] [OUT_JSON]\n", .{}); + return; + } + const dir = args[1]; + const ref = args[2]; + const rows: u32 = if (args.len > 3) try std.fmt.parseInt(u32, args[3], 10) else 64; + var r = try rt.open(); + defer r.deinit(); + var c = try lfm3.config.Config.read(gpa, io, dir); + defer c.deinit(); + var l = try @import("xpu").loader.Loader.init(gpa, &r, dir); + var t0 = nowNs(); + var m = try lfm3.model.Model.load(gpa, &r, &l, &c, 4096, rows); + std.debug.print("loaded in {d:.1} s, {d:.0} MiB on the device, windows of {d} rows\n", .{ @as(f64, @floatFromInt(nowNs() - t0)) / 1e9, @as(f64, @floatFromInt(m.device_bytes)) / (1 << 20), rows }); + + var out: std.Io.Writer.Allocating = .init(gpa); + defer out.deinit(); + try out.writer.writeAll("{"); + var bad = false; + for (prompts, 0..) |name, pi| { + const meta_path = try std.fmt.allocPrint(gpa, "{s}/{s}/meta.json", .{ ref, name }); + const meta_bytes = try std.Io.Dir.cwd().readFileAlloc(io, meta_path, gpa, .unlimited); + const meta = try std.json.parseFromSlice(std.json.Value, gpa, meta_bytes, .{}); + const id_vals = meta.value.object.get("ids").?.array.items; + const ids = try gpa.alloc(u32, id_vals.len); + for (id_vals, ids) |v, *o| o.* = @intCast(v.integer); + const T = ids.len; + var keep: Rows = .{ .layers = undefined }; + for (&keep.layers) |*ly| ly.* = try gpa.alloc(u16, T * H); + const hk: lfm3.model.Hook = .{ .ctx = &keep, .f = Rows.keep }; + try m.reset(); + var i: usize = 0; + while (i < T) { + const n = @min(T - i, rows); + keep.row0 = i; + try m.forwardPrompt(ids[i .. i + n], i + n == T, hk); + i += n; + } + const next = try m.argmax(); + const gaps = meta.value.object.get("bf16_gap").?.object; + std.debug.print("{s}: {d} prompt tokens; per-row relative L2 error against float64 (mean / max), and the bf16-points pass's own gap\n", .{ name, T }); + var worst: f64 = 0; + for (0..25) |li| { + const layer: i32 = if (li == 24) -1 else @intCast(li); + const path = if (li == 24) try std.fmt.allocPrint(gpa, "{s}/{s}/final.f32", .{ ref, name }) else try std.fmt.allocPrint(gpa, "{s}/{s}/h{d:0>2}.f32", .{ ref, name, li }); + const want = try readF32(gpa, io, path); + defer gpa.free(want); + const e = relRows(keep.layers[li], want); + var key: [8]u8 = undefined; + const g = gaps.get(try std.fmt.bufPrint(&key, "{d}", .{layer})).?.array.items; + const gm: f64 = switch (g[0]) { + .float => |f| f, + else => 0, + }; + const gx: f64 = switch (g[1]) { + .float => |f| f, + else => 0, + }; + const ratio = e[0] / gm; + worst = @max(worst, ratio); + std.debug.print(" {s} {d:>2}: xpu {e:.2} / {e:.2} bf16 gap {e:.2} / {e:.2} mean ratio {d:.2}\n", .{ if (li == 24) "final" else "layer", if (li == 24) 0 else li, e[0], e[1], gm, gx, ratio }); + } + const top = meta.value.object.get("top1").?.array.items; + const want_next: u32 = @intCast(top[T - 1].integer); + std.debug.print(" next token after the prompt: xpu {d}, float64 {d}; worst layer mean ratio {d:.2}\n", .{ next, want_next, worst }); + if (worst > 4.0) bad = true; + // greedy reply, twice from a fresh prefill: the same tokens + var reply: [2][reply_len]u32 = undefined; + var dec_ns: u64 = 0; + for (0..2) |run_i| { + try m.reset(); + try m.prefill(ids, null); + reply[run_i][0] = try m.argmax(); + t0 = nowNs(); + for (1..reply_len) |j| { + try m.forward(reply[run_i][j - 1 .. j], true, null); + reply[run_i][j] = try m.argmax(); + } + dec_ns = nowNs() - t0; + } + const same = std.mem.eql(u32, &reply[0], &reply[1]); + if (!same) bad = true; + std.debug.print(" {d} greedy tokens, second run {s}; decode {d:.1} tok/s (synchronous argmax per token)\n", .{ reply_len, if (same) "identical" else "DIFFERENT", @as(f64, reply_len - 1) / (@as(f64, @floatFromInt(dec_ns)) / 1e9) }); + try out.writer.print("{s}\"{s}\": {{\"prompt\": [", .{ if (pi == 0) "" else ", ", name }); + for (ids, 0..) |v, k| try out.writer.print("{s}{d}", .{ if (k == 0) "" else ", ", v }); + try out.writer.writeAll("], \"reply\": ["); + for (reply[0], 0..) |v, k| try out.writer.print("{s}{d}", .{ if (k == 0) "" else ", ", v }); + try out.writer.writeAll("]}"); + } + try out.writer.writeAll("}\n"); + if (args.len > 4) try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = args[4], .data = out.written() }); + std.debug.print("{s}\n", .{if (bad) "lfm3_fwd: FAILED" else "lfm3_fwd: ok"}); + if (bad) return error.TestFailed; +} diff --git a/zig/tests/xpu/lfm3_norm.zig b/zig/tests/xpu/lfm3_norm.zig new file mode 100644 index 0000000000..460412ab92 --- /dev/null +++ b/zig/tests/xpu/lfm3_norm.zig @@ -0,0 +1,112 @@ +//! LFM3 row ops (lfm3_rmsnorm, the sandwich step lfm3_post_add_norm, lfm3_swiglu) against FP64 on random data. + +const std = @import("std"); +const o = @import("lfm3_ops.zig"); +const stop = @import("xpu").stop; + +const N: u32 = 2048; +const ROWS: u32 = 5; +const FF: u32 = 7680; +const eps: f64 = 1e-5; + +fn rms(x: []const f64, w: []const f64, out: []f64) void { + var ss: f64 = 0; + for (x) |v| ss += v * v; + const rs = 1.0 / @sqrt(ss / @as(f64, @floatFromInt(x.len)) + eps); + for (x, w, out) |v, wv, *y| y.* = wv * v * rs; +} + +pub const panic = std.debug.FullPanic(@import("rig.zig").panicDrain); +pub fn main() !void { + stop.install(); + var r = try o.rt.open(); + defer r.deinit(); + var m = try r.module(@import("xpu").kernels.lfm3); + var k_rms = try m.kernel("lfm3_rmsnorm", .{ 256, 1, 1 }); + var k_pan = try m.kernel("lfm3_post_add_norm", .{ 256, 1, 1 }); + var k_swi = try m.kernel("lfm3_swiglu", .{ 64, 1, 1 }); + var t: o.Tally = .{ .name = "lfm3_norm" }; + var prng = std.Random.DefaultPrng.init(14); + const rnd = prng.random(); + std.debug.print("row ops over {d} rows of {d}, against FP64\n", .{ ROWS, N }); + + const xh = try o.gpa.alloc(u16, ROWS * N); + const dh = try o.gpa.alloc(u16, ROWS * N); + const wph = try o.gpa.alloc(u16, N); + const wnh = try o.gpa.alloc(u16, N); + o.randBf(xh, rnd, 3.0); + o.randBf(dh, rnd, 0.5); + for (wph) |*v| v.* = o.toBf(0.2 + 0.1 * rnd.floatNorm(f64)); + for (wnh) |*v| v.* = o.toBf(1.0 + 0.3 * rnd.floatNorm(f64)); + const xf = try o.widen(o.gpa, xh); + const df = try o.widen(o.gpa, dh); + const wpf = try o.widen(o.gpa, wph); + const wnf = try o.widen(o.gpa, wnh); + + // RMSNorm (the final embedding_norm and every pre-norm) + { + const x = try o.upload(&r, xh); + const w = try o.upload(&r, wnh); + const y = try r.alloc(ROWS * N * 2); + try o.run(&k_rms, .{ ROWS, 1, 1 }, .{ x, w, y, N, @as(f32, 1e-5) }); + const yh = try o.fetch(&r, y, ROWS * N); + const want = try o.gpa.alloc(f64, ROWS * N); + for (0..ROWS) |z| rms(xf[z * N ..][0..N], wnf, want[z * N ..][0..N]); + t.check("rmsnorm", o.relErr(yh, want)); + } + + // sandwich step: x += post_norm(d); y = next_norm(x) + { + const x = try o.upload(&r, xh); + const d = try o.upload(&r, dh); + const wp = try o.upload(&r, wph); + const wn = try o.upload(&r, wnh); + const y = try r.alloc(ROWS * N * 2); + try o.run(&k_pan, .{ ROWS, 1, 1 }, .{ d, x, wp, wn, y, N, @as(f32, 1e-5) }); + const xo = try o.fetch(&r, x, ROWS * N); + const yh = try o.fetch(&r, y, ROWS * N); + const hw = try o.gpa.alloc(f64, ROWS * N); + const yw = try o.gpa.alloc(f64, ROWS * N); + const tmp = try o.gpa.alloc(f64, N); + for (0..ROWS) |z| { + rms(df[z * N ..][0..N], wpf, tmp); + for (0..N) |i| hw[z * N + i] = xf[z * N + i] + tmp[i]; + rms(hw[z * N ..][0..N], wnf, yw[z * N ..][0..N]); + } + t.check("post_add_norm residual h = x + post_norm(d)", o.relErr(xo, hw)); + t.check("post_add_norm output next_norm(h)", o.relErr(yh, yw)); + // planted control: the residual added before the post norm (h = post_norm(x + d)) must not match + for (0..ROWS) |z| { + for (0..N) |i| tmp[i] = xf[z * N + i] + df[z * N + i]; + rms(tmp, wpf, hw[z * N ..][0..N]); + } + t.control("residual added before the post norm", o.relErr(xo, hw)); + } + + // SwiGLU: silu(w1 x) * (w3 x) + { + const n = 3 * FF; + const ah = try o.gpa.alloc(u16, n); + const bh = try o.gpa.alloc(u16, n); + o.randBf(ah, rnd, 2.0); + o.randBf(bh, rnd, 1.0); + const a = try o.upload(&r, ah); + const b = try o.upload(&r, bh); + const y = try r.alloc(n * 2); + try o.run(&k_swi, .{ (n + 63) / 64, 1, 1 }, .{ a, b, y, n }); + const yh = try o.fetch(&r, y, n); + const want = try o.gpa.alloc(f64, n); + for (0..n) |i| { + const x = o.bf(ah[i]); + want[i] = x / (1.0 + @exp(-x)) * o.bf(bh[i]); + } + t.check("swiglu", o.relErr(yh, want)); + // planted control: gate and up swapped (silu on w3 x) must not match + for (0..n) |i| { + const x = o.bf(bh[i]); + want[i] = x / (1.0 + @exp(-x)) * o.bf(ah[i]); + } + t.control("gate and up swapped", o.relErr(yh, want)); + } + try t.finish(); +} diff --git a/zig/tests/xpu/lfm3_ops.zig b/zig/tests/xpu/lfm3_ops.zig new file mode 100644 index 0000000000..993557ff33 --- /dev/null +++ b/zig/tests/xpu/lfm3_ops.zig @@ -0,0 +1,95 @@ +//! Shared pieces of the LFM3 op tests: bf16 data, launches, the FP64 comparison and its control. + +const std = @import("std"); +pub const rt = @import("rig.zig"); + +pub const gpa = std.heap.page_allocator; +pub const Buf = rt.Buffer; +/// Largest |device - FP64| over the largest |FP64| a bf16 output may show; a control must land above it. +pub const tol: f64 = 1e-2; + +pub fn bf(v: u16) f64 { + return @as(f32, @bitCast(@as(u32, v) << 16)); +} + +/// Round to nearest even, as the kernels' to_bf. +pub fn toBf(f: f64) u16 { + var u: u32 = @bitCast(@as(f32, @floatCast(f))); + u +%= 0x7fff + ((u >> 16) & 1); + return @intCast(u >> 16); +} + +pub fn randBf(buf: []u16, rnd: std.Random, sd: f64) void { + for (buf) |*v| v.* = toBf(rnd.floatNorm(f64) * sd); +} + +pub fn widen(gpa_: std.mem.Allocator, h: []const u16) ![]f64 { + const out = try gpa_.alloc(f64, h.len); + for (h, out) |v, *o| o.* = bf(v); + return out; +} + +pub fn run(k: *rt.Kernel, groups: [3]u32, a: anytype) !void { + inline for (a, 0..) |v, i| { + const T = @TypeOf(v); + if (T == Buf) { + try k.setBuffer(i, v); + } else if (T == f32 or T == comptime_float) { + try k.setF32(i, v); + } else try k.setU32(i, @as(u32, v)); + } + try k.launch(groups); +} + +pub fn upload(r: *rt.Runtime, h: []const u16) !Buf { + const b = try r.alloc(h.len * 2); + try r.upload(b, std.mem.sliceAsBytes(h)); + return b; +} + +pub fn fetch(r: *rt.Runtime, b: Buf, n: usize) ![]u16 { + const h = try gpa.alloc(u16, n); + try r.sync(); + try r.download(std.mem.sliceAsBytes(h), b); + try r.sync(); + return h; +} + +/// max|got - want| / max|want|. +pub fn relErr(got: []const u16, want: []const f64) f64 { + var e: f64 = 0; + var m: f64 = 0; + for (got, want) |g, w| { + e = @max(e, @abs(bf(g) - w)); + m = @max(m, @abs(w)); + } + return if (m == 0) e else e / m; +} + +/// Tallies the checks of one test program: real checks must stay within tol, planted controls must exceed it. +pub const Tally = struct { + name: []const u8, + bad: bool = false, + + pub fn check(t: *Tally, what: []const u8, err: f64) void { + const ok = err <= tol; + if (!ok) t.bad = true; + std.debug.print(" {s}: rel err {e:.2} (tol {e:.0}) {s}\n", .{ what, err, tol, if (ok) "pass" else "FAIL" }); + } + + pub fn control(t: *Tally, what: []const u8, err: f64) void { + const fails = err > tol; + if (!fails) t.bad = true; + std.debug.print(" control {s}: rel err {e:.2} {s}\n", .{ what, err, if (fails) "fails as planted" else "PASSES (control broken)" }); + } + + pub fn same(t: *Tally, what: []const u8, eq: bool) void { + if (!eq) t.bad = true; + std.debug.print(" {s}: {s}\n", .{ what, if (eq) "bitwise identical" else "DIFFERENT" }); + } + + pub fn finish(t: *Tally) !void { + std.debug.print("{s}: {s}\n", .{ t.name, if (t.bad) "FAILED" else "ok" }); + if (t.bad) return error.TestFailed; + } +}; diff --git a/zig/tests/xpu/lfm3_pgemm.zig b/zig/tests/xpu/lfm3_pgemm.zig new file mode 100644 index 0000000000..7b289e52b4 --- /dev/null +++ b/zig/tests/xpu/lfm3_pgemm.zig @@ -0,0 +1,78 @@ +//! pgemm, the prompt windows' GEMM: MLX 4-bit in groups of 64 and 32 against FP64; each row alone = in a window. + +const std = @import("std"); +const o = @import("lfm3_ops.zig"); +const stop = @import("xpu").stop; +const kernels = @import("xpu").kernels; + +/// W[out][i] = q * scale + bias of an MLX 4-bit table in groups of `group`, FP64. +fn dequant(w: []const u32, s: []const f64, b: []const f64, outs: usize, in_dim: usize, group: usize) ![]f64 { + const d = try o.gpa.alloc(f64, outs * in_dim); + for (0..outs) |r| for (0..in_dim) |i| { + const q: f64 = @floatFromInt((w[r * (in_dim / 8) + i / 8] >> @intCast(4 * (i % 8))) & 15); + const g = r * (in_dim / group) + i / group; + d[r * in_dim + i] = q * s[g] + b[g]; + }; + return d; +} + +pub const panic = std.debug.FullPanic(@import("rig.zig").panicDrain); +pub fn main() !void { + stop.install(); + var r = try o.rt.open(); + defer r.deinit(); + var vm = try r.module(kernels.lfm3_mv); + var k = try vm.kernel("pgemm", .{ 16, 16, 1 }); + var t: o.Tally = .{ .name = "lfm3_pgemm" }; + var prng = std.Random.DefaultPrng.init(64); + const rnd = prng.random(); + std.debug.print("pgemm: 4-bit tables in groups of 64 and 32, against FP64; rows alone and in a window\n", .{}); + var buf: [96]u8 = undefined; + const outs: u32 = 1000; // not a multiple of the 64-wide tile + const ld: u32 = outs + 24; + const off: u32 = 16; + for ([_]u32{ 2048, 7680 }) |in_dim| for ([_]u32{ 64, 32 }) |group| { + const wh = try o.gpa.alloc(u32, outs * in_dim / 8); + for (wh) |*v| v.* = rnd.int(u32); + const w = try r.alloc(wh.len * 4); + try r.upload(w, std.mem.sliceAsBytes(wh)); + const sh = try o.gpa.alloc(u16, outs * in_dim / group); + const bh = try o.gpa.alloc(u16, sh.len); + o.randBf(sh, rnd, 0.02); + o.randBf(bh, rnd, 0.1); + const s = try o.upload(&r, sh); + const b = try o.upload(&r, bh); + const wd = try dequant(wh, try o.widen(o.gpa, sh), try o.widen(o.gpa, bh), outs, in_dim, group); + const rows: u32 = 130; + const xh = try o.gpa.alloc(u16, rows * in_dim); + o.randBf(xh, rnd, 1.0); + const x = try o.upload(&r, xh); + const xf = try o.widen(o.gpa, xh); + const want = try o.gpa.alloc(f64, rows * outs); + for (0..rows) |i| for (0..outs) |j| { + var acc: f64 = 0; + for (0..in_dim) |c| acc += xf[i * in_dim + c] * wd[j * in_dim + c]; + want[i * outs + j] = acc; + }; + const y = try r.alloc(rows * ld * 2); + const one = try r.alloc(rows * outs * 2); + for (0..rows) |i| try o.run(&k, .{ (outs + 63) / 64, 1, 1 }, .{ w, s, b, sub(x, i * in_dim * 2), sub(one, i * outs * 2), in_dim, outs, @as(u32, 1), outs, @as(u32, 0), group }); + const alone = try o.fetch(&r, one, rows * outs); + t.check(try std.fmt.bufPrint(&buf, "groups of {d}, in {d}, rows one at a time", .{ group, in_dim }), o.relErr(alone, want)); + for ([_]u32{ 5, 64, 130 }) |m| { + try o.run(&k, .{ (outs + 63) / 64, (m + 63) / 64, 1 }, .{ w, s, b, x, y, in_dim, outs, m, ld, off, group }); + const got = try o.fetch(&r, y, m * ld); + var eq = true; + for (0..m) |i| eq = eq and std.mem.eql(u16, got[i * ld + off ..][0..outs], alone[i * outs ..][0..outs]); + t.same(try std.fmt.bufPrint(&buf, "groups of {d}, in {d}, a window of {d:>3} rows (stride, offset) = alone", .{ group, in_dim, m }), eq); + } + // planted control: the table read in the other group size must not match + try o.run(&k, .{ (outs + 63) / 64, 1, 1 }, .{ w, s, b, x, one, in_dim, outs, @as(u32, 1), outs, @as(u32, 0), 96 - group }); + t.control(try std.fmt.bufPrint(&buf, "groups of {d} read as {d}, in {d}", .{ group, 96 - group, in_dim }), o.relErr(try o.fetch(&r, one, outs), want[0..outs])); + }; + try t.finish(); +} + +fn sub(b: o.Buf, off: usize) o.Buf { + return .{ .rt = b.rt, .ptr = @ptrFromInt(@intFromPtr(b.ptr.?) + off), .len = b.len - off }; +} diff --git a/zig/tests/xpu/lfm3_rope.zig b/zig/tests/xpu/lfm3_rope.zig new file mode 100644 index 0000000000..d5417b321b --- /dev/null +++ b/zig/tests/xpu/lfm3_rope.zig @@ -0,0 +1,108 @@ +//! LFM3 per-head QK RMSNorm with zero-tail partial RoPE, and the V store, against FP64 on random data. + +const std = @import("std"); +const o = @import("lfm3_ops.zig"); +const stop = @import("xpu").stop; + +const HD = 128; +const eps: f64 = 1e-5; + +/// Inverse frequencies as upstream builds them in fp32: base^(-2i/128) for i < angles, then zeros. +fn invFreq(base: f32, angles: usize, out: []f32) void { + for (out, 0..) |*f, i| f.* = if (i < angles) 1.0 / std.math.pow(f32, base, @as(f32, @floatFromInt(2 * i)) / 128.0) else 0; +} + +/// FP64 replay of one row's heads: RMSNorm over 128 dims, then rotate_half RoPE at `pos` with `angles` rotated pairs. +fn refRow(src: []const f64, w: []const f64, heads: usize, base: f64, angles: usize, pos: f64, out: []f64) void { + for (0..heads) |h| { + const s = src[h * HD ..][0..HD]; + var ss: f64 = 0; + for (s) |v| ss += v * v; + const rs = 1.0 / @sqrt(ss / HD + eps); + for (0..64) |i| { + const a = w[i] * s[i] * rs; + const b = w[i + 64] * s[i + 64] * rs; + const f = if (i < angles) pos / std.math.pow(f64, base, @as(f64, @floatFromInt(2 * i)) / 128.0) else 0; + out[h * HD + i] = a * @cos(f) - b * @sin(f); + out[h * HD + i + 64] = b * @cos(f) + a * @sin(f); + } + } +} + +const Case = struct { what: []const u8, base: f32, angles: usize, heads: u32, pos0: u32, ring: u32 }; + +pub const panic = std.debug.FullPanic(@import("rig.zig").panicDrain); +pub fn main() !void { + stop.install(); + var r = try o.rt.open(); + defer r.deinit(); + var m = try r.module(@import("xpu").kernels.lfm3); + var kr = try m.kernel("lfm3_qk_norm_rope", .{ 64, 1, 1 }); + var kv = try m.kernel("lfm3_v_store", .{ 64, 1, 1 }); + var t: o.Tally = .{ .name = "lfm3_rope" }; + var prng = std.Random.DefaultPrng.init(12); + const rnd = prng.random(); + + const n: u32 = 9; + const stride: u32 = 2048 + 2 * 512; // a fused q | k | v projection row + const sh = try o.gpa.alloc(u16, n * stride); + o.randBf(sh, rnd, 1.5); + const src = try o.upload(&r, sh); + const sf = try o.widen(o.gpa, sh); + const wh = try o.gpa.alloc(u16, HD); + for (wh) |*v| v.* = o.toBf(1.0 + 0.3 * rnd.floatNorm(f64)); + const w = try o.upload(&r, wh); + const wf = try o.widen(o.gpa, wh); + const inv = try r.alloc(64 * 4); + var invh: [64]f32 = undefined; + std.debug.print("QK RMSNorm + RoPE over 128-dim heads, {d} rows, against FP64\n", .{n}); + + const cases = [_]Case{ + .{ .what = "q, full layer (theta 1e6, 32 of 64 pairs)", .base = 1e6, .angles = 32, .heads = 16, .pos0 = 7000, .ring = 0 }, + .{ .what = "q, sliding layer (theta 1e4, all 64 pairs)", .base = 1e4, .angles = 64, .heads = 16, .pos0 = 5, .ring = 0 }, + .{ .what = "k into a 256-slot ring across its wrap", .base = 1e4, .angles = 64, .heads = 4, .pos0 = 250, .ring = 256 }, + }; + var full_out: []u16 = &.{}; + for (cases) |c| { + invFreq(c.base, c.angles, &invh); + try r.upload(inv, std.mem.sliceAsBytes(&invh)); + const slots: u32 = if (c.ring == 0) n else c.ring; + const dst_stride: u32 = c.heads * HD; + const dst = try r.alloc(slots * dst_stride * 2); + const src_off: u32 = if (c.heads == 16) 0 else 2048; + try o.run(&kr, .{ c.heads, n, 1 }, .{ src, stride, src_off, dst, dst_stride, w, inv, c.pos0, c.ring, @as(f32, 1e-5) }); + const dh = try o.fetch(&r, dst, slots * dst_stride); + const got = try o.gpa.alloc(u16, n * dst_stride); + const want = try o.gpa.alloc(f64, n * dst_stride); + for (0..n) |z| { + const slot = if (c.ring == 0) z else (c.pos0 + z) % c.ring; + @memcpy(got[z * dst_stride ..][0..dst_stride], dh[slot * dst_stride ..][0..dst_stride]); + refRow(sf[z * stride + src_off ..], wf, c.heads, c.base, c.angles, @floatFromInt(c.pos0 + z), want[z * dst_stride ..]); + } + t.check(c.what, o.relErr(got, want)); + if (c.angles == 32) full_out = got; + } + + // V rows into the ring: a plain copy + { + const ring: u32 = 256; + const pos0: u32 = 250; + const dst = try r.alloc(ring * 512 * 2); + try o.run(&kv, .{ 512 / 64, n, 1 }, .{ src, stride, @as(u32, 2048 + 512), dst, @as(u32, 512), pos0, ring }); + const dh = try o.fetch(&r, dst, ring * 512); + var eq = true; + for (0..n) |z| { + const slot = (pos0 + z) % ring; + if (!std.mem.eql(u16, dh[slot * 512 ..][0..512], sh[z * stride + 2048 + 512 ..][0..512])) eq = false; + } + t.same("v rows stored at slot (pos0 + r) % 256", eq); + } + + // planted control: the full layer's q replayed with every pair rotated (partial rotary ignored) must not match + { + const want = try o.gpa.alloc(f64, n * 16 * HD); + for (0..n) |z| refRow(sf[z * stride ..], wf, 16, 1e6, 64, @floatFromInt(7000 + z), want[z * 16 * HD ..]); + t.control("full layer with all 64 pairs rotated", o.relErr(full_out, want)); + } + try t.finish(); +} diff --git a/zig/tests/xpu/lfm3_rows.zig b/zig/tests/xpu/lfm3_rows.zig new file mode 100644 index 0000000000..36bec81272 --- /dev/null +++ b/zig/tests/xpu/lfm3_rows.zig @@ -0,0 +1,65 @@ +//! LFM3's window matvecs (pdense, bf16 and fp32 out) against the one-row kernels, bit for bit, on random data. + +const std = @import("std"); +const o = @import("lfm3_ops.zig"); +const stop = @import("xpu").stop; +const kernels = @import("xpu").kernels; + +const ROWS: u32 = 1024; + +fn randWords(r: *o.rt.Runtime, n: usize, rnd: std.Random) !o.Buf { + const h = try o.gpa.alloc(u32, n); + for (h) |*v| v.* = rnd.int(u32); + const b = try r.alloc(n * 4); + try r.upload(b, std.mem.sliceAsBytes(h)); + return b; +} + +fn fetch32(r: *o.rt.Runtime, b: o.Buf, n: usize) ![]u32 { + const h = try o.gpa.alloc(u32, n); + try r.sync(); + try r.download(std.mem.sliceAsBytes(h), b); + try r.sync(); + return h; +} + +pub const panic = std.debug.FullPanic(@import("rig.zig").panicDrain); +pub fn main() !void { + stop.install(); + var r = try o.rt.open(); + defer r.deinit(); + var am = try r.module(kernels.attn); + var pm = try r.module(if (kernels.dpas) kernels.nem_prow else kernels.nem_prow_gen12); + const tile_outs: u32 = if (kernels.dpas) 8 else 4; + var k_bf = try am.kernel("qmv4_bf", .{ 64, 1, 1 }); + var k_f32 = try am.kernel("qmv4_f32", .{ 64, 1, 1 }); + var k_pd = try pm.kernel("pdense", .{ 16 * (16 / tile_outs), 1, 1 }); + var t: o.Tally = .{ .name = "lfm3_rows" }; + var prng = std.Random.DefaultPrng.init(21); + const rnd = prng.random(); + std.debug.print("window matvecs (pdense) against the one-row kernels, {d} output rows, bit for bit\n", .{ROWS}); + for ([_]u32{ 2048, 7680 }) |in_dim| { + const sh = try o.gpa.alloc(u16, ROWS * in_dim / 64); + const bh = try o.gpa.alloc(u16, sh.len); + o.randBf(sh, rnd, 0.02); + o.randBf(bh, rnd, 0.1); + const w = try randWords(&r, ROWS * in_dim / 8, rnd); + const s = try o.upload(&r, sh); + const b = try o.upload(&r, bh); + const xh = try o.gpa.alloc(u16, 16 * in_dim); + o.randBf(xh, rnd, 1.0); + const x = try o.upload(&r, xh); + const y1 = try r.alloc(16 * ROWS * 4); + const yn = try r.alloc(16 * ROWS * 4); + for ([_]u32{ 1, 2, 3, 4, 8, 13, 16 }) |n| { + var buf: [96]u8 = undefined; + for (0..n) |z| try o.run(&k_bf, .{ ROWS / 4, 1, 1 }, .{ w, s, b, x, y1, in_dim, @as(u32, @intCast(z)) * in_dim, @as(u32, @intCast(z)) * ROWS, ROWS }); + try o.run(&k_pd, .{ ROWS / 16, (n + 7) / 8, 1 }, .{ w, s, b, x, yn, in_dim, ROWS, n, ROWS, @as(u32, 0), @as(u32, 0) }); + t.same(try std.fmt.bufPrint(&buf, "pdense bf16 in {d} n {d:>2} vs qmv4_bf", .{ in_dim, n }), std.mem.eql(u16, try o.fetch(&r, y1, n * ROWS), try o.fetch(&r, yn, n * ROWS))); + for (0..n) |z| try o.run(&k_f32, .{ ROWS / 4, 1, 1 }, .{ w, s, b, x, y1, in_dim, @as(u32, @intCast(z)) * in_dim, @as(u32, @intCast(z)) * ROWS, ROWS }); + try o.run(&k_pd, .{ ROWS / 16, (n + 7) / 8, 1 }, .{ w, s, b, x, yn, in_dim, ROWS, n, ROWS, @as(u32, 0), @as(u32, 1) }); + t.same(try std.fmt.bufPrint(&buf, "pdense fp32 in {d} n {d:>2} vs qmv4_f32", .{ in_dim, n }), std.mem.eql(u32, try fetch32(&r, y1, n * ROWS), try fetch32(&r, yn, n * ROWS))); + } + } + try t.finish(); +} diff --git a/zig/tests/xpu/lfm3_streams.zig b/zig/tests/xpu/lfm3_streams.zig new file mode 100644 index 0000000000..ee6682d296 --- /dev/null +++ b/zig/tests/xpu/lfm3_streams.zig @@ -0,0 +1,204 @@ +//! LFM3 streams in shared forwards and verified windows with rollback, against each stream alone. + +const std = @import("std"); +const rt = @import("rig.zig"); +const lfm3 = @import("lfm3_xpu"); +const stop = @import("xpu").stop; +const Model = lfm3.model.Model; +const Seq = lfm3.model.Seq; + +const steps = 64; + +fn readIds(gpa: std.mem.Allocator, io: std.Io, ref: []const u8, name: []const u8) ![]u32 { + const path = try std.fmt.allocPrint(gpa, "{s}/{s}/meta.json", .{ ref, name }); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .unlimited); + const meta = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); + const vals = meta.value.object.get("ids").?.array.items; + const ids = try gpa.alloc(u32, vals.len); + for (vals, ids) |v, *o| o.* = @intCast(v.integer); + return ids; +} + +/// The stream's prompt on its own sequence (bound), then its first token. +fn prefill(m: *Model, q: *Seq, ids: []const u32) !u32 { + m.bind(q); + try m.reset(); + try m.prefill(ids, null); + return m.argmax(); +} + +/// One row of each stream in one forward (verifyMulti's one-row windows); the rows' greedy tokens in `out`. +fn shared(m: *Model, qs: []const *Seq, toks: []const u32, out: []i32) !void { + var segs: [8]lfm3.model.Seg = undefined; + for (qs, toks, segs[0..qs.len]) |q, *t, *g| g.* = .{ .q = q, .ids = t[0..1] }; + try m.forwardSegs(segs[0..qs.len], .rows, false, null); + try m.rowTokens(out[0..qs.len]); +} + +/// Stream q alone: the prompt, then `steps` greedy tokens. +fn solo(m: *Model, q: *Seq, ids: []const u32, out: []u32) !void { + out[0] = try prefill(m, q, ids); + for (1..out.len) |j| { + try m.forward(out[j - 1 .. j], true, null); + out[j] = try m.argmax(); + } +} + +pub const panic = std.debug.FullPanic(@import("rig.zig").panicDrain); +pub fn main(init: std.process.Init) !void { + stop.install(); + const gpa = init.gpa; + const io = init.io; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len < 3) { + std.debug.print("skipped: usage tf-xpu-lfm3_streams-test MODEL_DIR REF_DIR\n", .{}); + return; + } + var r = try rt.open(); + defer r.deinit(); + var c = try lfm3.config.Config.read(gpa, io, args[1]); + defer c.deinit(); + var l = try @import("xpu").loader.Loader.init(gpa, &r, args[1]); + var m = try Model.load(gpa, &r, &l, &c, 2048, 128); + const names = [_][]const u8{ "chat", "capitals", "items" }; + var ids: [3][]u32 = undefined; + for (names, &ids) |n, *v| v.* = try readIds(gpa, io, args[2], n); + var want: [3][steps]u32 = undefined; + const q0 = try m.newSeq(); + for (0..3) |k| try solo(&m, q0, ids[k], &want[k]); + // shared: chat and items decode together; items is freed at step 20 and capitals starts on a new sequence + const qa = try m.newSeq(); + const qb = try m.newSeq(); + var got: [3][steps]u32 = undefined; + got[0][0] = try prefill(&m, qa, ids[0]); + got[2][0] = try prefill(&m, qb, ids[2]); + var qs: [2]*Seq = .{ qa, qb }; + var toks: [2]u32 = undefined; + var row: [2]i32 = undefined; + var ja: usize = 1; + var jb: usize = 1; + while (ja < 21) : ({ + ja += 1; + jb += 1; + }) { + toks = .{ got[0][ja - 1], got[2][jb - 1] }; + try shared(&m, &qs, &toks, &row); + got[0][ja] = @intCast(row[0]); + got[2][jb] = @intCast(row[1]); + } + m.freeSeq(qb); + const qc = try m.newSeq(); + got[1][0] = try prefill(&m, qc, ids[1]); + qs = .{ qa, qc }; + var jc: usize = 1; + while (jc < steps - 21) : ({ + ja += 1; + jc += 1; + }) { + toks = .{ got[0][ja - 1], got[1][jc - 1] }; + try shared(&m, &qs, &toks, &row); + got[0][ja] = @intCast(row[0]); + got[1][jc] = @intCast(row[1]); + } + var bad = false; + // four streams in every step: chat, capitals, items, chat again (a duplicate must equal its twin) + { + const order = [_]usize{ 0, 1, 2, 0 }; + var q4: [4]*Seq = undefined; + var g4: [4][steps]u32 = undefined; + for (&q4, order, 0..) |*q, k, z| { + q.* = try m.newSeq(); + g4[z][0] = try prefill(&m, q.*, ids[k]); + } + for (1..steps) |j| { + var t4: [4]u32 = undefined; + for (0..4) |z| t4[z] = g4[z][j - 1]; + var r4: [4]i32 = undefined; + try shared(&m, &q4, &t4, &r4); + for (0..4) |z| g4[z][j] = @intCast(r4[z]); + } + for (order, 0..) |k, z| { + const eq = std.mem.eql(u32, &g4[z], &want[k]); + if (!eq) bad = true; + std.debug.print(" four streams, row {d} ({s}): {s}\n", .{ z, names[k], if (eq) "equal to alone" else "DIFFERS from alone" }); + } + } + // the server's round: haiku (a chat prompt from CHAT_JSON), chat, capitals and items as rows 0..3, 66 steps + if (args.len > 3) { + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, args[3], gpa, .unlimited); + const doc = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); + const hv = doc.value.object.get("chat_haiku").?.array.items; + const hids = try gpa.alloc(u32, hv.len); + for (hv, hids) |v, *d| d.* = @intCast(v.integer); + var hwant: [70]u32 = undefined; + try solo(&m, q0, hids, &hwant); + const prompts = [_][]const u32{ hids, ids[0], ids[1], ids[2] }; + var q4: [4]*Seq = undefined; + var g4: [4][70]u32 = undefined; + for (&q4, prompts, 0..) |*q, pr, z| { + q.* = try m.newSeq(); + g4[z][0] = try prefill(&m, q.*, pr); + } + for (1..70) |j| { + var t4: [4]u32 = undefined; + for (0..4) |z| t4[z] = g4[z][j - 1]; + var r4: [4]i32 = undefined; + try shared(&m, &q4, &t4, &r4); + for (0..4) |z| g4[z][j] = @intCast(r4[z]); + } + std.debug.print(" haiku alone, tokens 58..66: {any}\n", .{hwant[58..66]}); + const eq = std.mem.eql(u32, &g4[0], &hwant); + if (!eq) bad = true; + const first = std.mem.indexOfDiff(u32, &g4[0], &hwant); + std.debug.print(" haiku in the server's four-stream round: {s} {any}\n", .{ if (eq) "equal to alone" else "DIFFERS from alone at", first }); + for (1..4) |z| { + const k = z - 1; + const e2 = std.mem.eql(u32, g4[z][0..steps], &want[k]); + if (!e2) bad = true; + std.debug.print(" row {d} ({s}) in that round: {s}\n", .{ z, names[k], if (e2) "equal to alone" else "DIFFERS from alone" }); + } + } + // chat in verified windows (the pending token, up to 4 copied drafts, every third wrong), kept rows only + { + const q = try m.newSeq(); + var got_w: [steps]u32 = undefined; + got_w[0] = try prefill(&m, q, ids[0]); + var j: usize = 1; + var round: usize = 0; + var kept_total: usize = 0; + while (j < steps) : (round += 1) { + const room = @min(5, steps - j + 1); + var win: [5]u32 = undefined; + win[0] = got_w[j - 1]; + for (1..room) |d| win[d] = if ((round + d) % 3 == 0) 7 else want[0][j - 1 + d]; + try m.forwardSegs(&.{.{ .q = q, .ids = win[0..room] }}, .rows, true, null); + var g: [5]i32 = undefined; + try m.rowTokens(g[0..room]); + var keep: usize = 1; + while (keep < room and win[keep] == @as(u32, @intCast(g[keep - 1]))) keep += 1; + m.bind(q); + try m.keepRows(@intCast(keep)); + for (0..keep) |k| if (j + k < steps) { + got_w[j + k] = @intCast(g[k]); + }; + j += keep; + kept_total += keep; + } + const eq = std.mem.eql(u32, &got_w, &want[0]); + if (!eq) bad = true; + std.debug.print(" chat in verified windows with rollback ({d} rounds for {d} tokens): {s} {any}\n", .{ round, kept_total, if (eq) "equal to alone" else "DIFFERS from alone at", std.mem.indexOfDiff(u32, &got_w, &want[0]) }); + } + const checks = [_]struct { k: usize, n: usize, what: []const u8 }{ + .{ .k = 0, .n = ja, .what = "chat, shared all the way (a partner freed and a new one prefilled at step 20)" }, + .{ .k = 2, .n = 21, .what = "items, shared until freed at step 20" }, + .{ .k = 1, .n = jc, .what = "capitals, prefilled mid-way, then shared" }, + }; + for (checks) |ck| { + const eq = std.mem.eql(u32, got[ck.k][0..ck.n], want[ck.k][0..ck.n]); + if (!eq) bad = true; + const first = std.mem.indexOfDiff(u32, got[ck.k][0..ck.n], want[ck.k][0..ck.n]); + std.debug.print(" {s}: {d} tokens {s}{s}\n", .{ ck.what, ck.n, if (eq) "equal to alone" else "DIFFER from alone at ", if (first) |f| try std.fmt.allocPrint(gpa, "{d}", .{f}) else "" }); + } + std.debug.print("{s}\n", .{if (bad) "lfm3_streams: FAILED" else "lfm3_streams: ok"}); + if (bad) return error.TestFailed; +}