From eb56483683a60591c77a58614f7ac8bddcbf4eb2 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 00:51:50 +0800 Subject: [PATCH 1/9] idol: Idol's denoiser as a native Metal family, with tf-idol and its reference case tool --- tools/zig/idol_case.py | 243 ++++++++++++++++++ zig/build/metal_targets.zig | 1 + zig/kernels/metal/idol.metal | 329 +++++++++++++++++++++++++ zig/kernels/metal/idol_attention.metal | 149 +++++++++++ zig/kernels/metal/sources.zig | 2 + zig/src/families/idol/config.zig | 71 ++++++ zig/src/families/idol/host.zig | 147 +++++++++++ zig/src/families/idol/idol.zig | 19 ++ zig/src/families/idol/kernels.zig | 209 ++++++++++++++++ zig/src/families/idol/model.zig | 244 ++++++++++++++++++ zig/src/families/idol/reference.zig | 79 ++++++ zig/src/families/idol/weights.zig | 310 +++++++++++++++++++++++ zig/src/tensorfold.zig | 1 + zig/tests/idol.zig | 242 ++++++++++++++++++ 14 files changed, 2046 insertions(+) create mode 100644 tools/zig/idol_case.py create mode 100644 zig/kernels/metal/idol.metal create mode 100644 zig/kernels/metal/idol_attention.metal create mode 100644 zig/src/families/idol/config.zig create mode 100644 zig/src/families/idol/host.zig create mode 100644 zig/src/families/idol/idol.zig create mode 100644 zig/src/families/idol/kernels.zig create mode 100644 zig/src/families/idol/model.zig create mode 100644 zig/src/families/idol/reference.zig create mode 100644 zig/src/families/idol/weights.zig create mode 100644 zig/tests/idol.zig diff --git a/tools/zig/idol_case.py b/tools/zig/idol_case.py new file mode 100644 index 0000000000..83453a8e54 --- /dev/null +++ b/tools/zig/idol_case.py @@ -0,0 +1,243 @@ +#!/usr/bin/env python3 +"""Idol reference cases for tf-idol from the publisher's idol-sdk (torch): the denoiser's schedule and its blocks.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import sys +import time +from pathlib import Path + +USAGE = """\ + case WEIGHTS SDK OUT --features F [F] noise, each step's input and both velocities, the final sample + blocks WEIGHTS SDK CASE [--step S] every block's rows at step S, and each block in float32 on its input + exact WEIGHTS SDK CASE [--steps N] both velocities in float32 on the case's first N inputs (the ruler) +WEIGHTS is a speridlabs/idol snapshot; SDK an idol-sdk checkout at the pinned commit. F holds feature1 (DINOv3 rows +[1, 4101, 1280]) and feature2 (FLUX.2 VAE rows [1, 4101, 128]), in one file or two, as the SDK's ImageEncoder returns. +""" + +SDK_COMMIT = "4702a7f00850e34b6e83a2170cc123c3e9570744" +SDK_SHA256 = { + "idol/models/dit.py": "fbe725cc478b0acd245d9404ac8d9df327ebfbf79fe51c954d94b005c5bd6af4", + "idol/models/layers.py": "2dbf24107396194636c9a2e89e7acee1bcf91b14dba1b4cde0318cb85b3f482f", + "idol/sampler.py": "532987b5599be678361c2fcf566335bb3e1b0171696f985818d10b7460bfc8b0", +} +WEIGHTS_REVISION = "2731c2291fc65a797c5cb627a5e5526fc1490902" +DENOISER_ARGS = {"q_token_length": 8192, "in_channels": 16, "cam_channels": 5, "model_channels": 1024, + "cond_channels": 1280, "cond2_channels": 128, "num_blocks": 24, "num_refiner_blocks": 2, + "num_heads": 16, "mlp_ratio": 4} + + +def sdk(path: str): + """The SDK's denoiser and sampler, imported once each file used has the pinned commit's sha256.""" + root = Path(path) + for name, want in SDK_SHA256.items(): + digest = hashlib.sha256((root / name).read_bytes()).hexdigest() + if digest != want: + sys.exit(f"{root / name} has sha256 {digest}; this tool runs only idol-sdk {SDK_COMMIT[:12]}'s ({want})") + sys.path.insert(0, str(root)) + from idol.models.dit import LatentSeqMMFlowModel + from idol.sampler import sample_flow + return LatentSeqMMFlowModel, sample_flow + + +def denoiser(weights: str, sdk_dir: str, device: str): + """The SDK's denoiser with the snapshot's weights, strictly loaded, as IdolPipeline builds it.""" + import torch + from safetensors.torch import load_file + + model_cls, _ = sdk(sdk_dir) + config = json.loads((Path(weights) / "config.json").read_text()) + args = config["denoiser"]["args"] + if config["denoiser"]["name"] != "LatentSeqMMFlowModel" or args != DENOISER_ARGS: + sys.exit(f"{weights}/config.json is not speridlabs/idol {WEIGHTS_REVISION[:12]}'s denoiser") + model = model_cls(**args) + model.load_state_dict(load_file(str(Path(weights) / "denoiser.safetensors")), strict=True) + return model.eval().requires_grad_(False).to(torch.device(device)) + + +def features(paths: list[str]): + """feature1 and feature2 [1, 4101, C] float32 from one or two safetensors files.""" + from safetensors.torch import load_file + + found = {} + for path in paths: + found.update(load_file(path)) + for key, width in (("feature1", 1280), ("feature2", 128)): + if key not in found or tuple(found[key].shape) != (1, 4101, width): + sys.exit(f"the features need {key} [1, 4101, {width}]") + return found["feature1"].float(), found["feature2"].float() + + +def condition(f1, f2, device: str): + """The sampler's cond and neg_cond (both streams zeroed), as IdolPipeline passes them.""" + from tensordict import TensorDict + + cond = TensorDict({"feature1": f1, "feature2": f2}, batch_size=1, device=device) + return cond, cond.apply(lambda x: x * 0) + + +def sample(latent, camera, device: str): + """The sampler's TensorDict for one latent [8192, 16] and camera [5].""" + from tensordict import TensorDict + + return TensorDict({"latent": latent[None].to(device), "camera": camera.view(1, 1, 5).to(device)}, batch_size=1, + device=device) + + +def flush(device: str) -> None: + import torch + + if device == "mps": + torch.mps.synchronize() + + +def case(args) -> int: + import torch + from safetensors.torch import save_file + + out = Path(args.out_dir) + out.mkdir(parents=True, exist_ok=True) + model = denoiser(args.weights, args.sdk, args.device) + _, sample_flow = sdk(args.sdk) + f1, f2 = features(args.features) + cond, neg = condition(f1, f2, args.device) + gen = torch.Generator("cpu").manual_seed(args.seed) + latent = torch.randn((8192, 16), generator=gen, dtype=torch.float32) + camera = torch.randn((5,), generator=gen, dtype=torch.float32) + calls = [] + forward = model.forward + + def recorded(x, t, c): + mark = time.perf_counter() + v = forward(x, t, c) + flush(args.device) + calls.append({"latent": x["latent"][0].float().cpu(), "camera": x["camera"][0, 0].float().cpu(), + "t": float(t[0]), "negative": bool(c["feature1"].abs().max() == 0), + "v": v["latent"][0].float().cpu(), "v_camera": v["camera"][0, 0].float().cpu(), + "s": time.perf_counter() - mark}) + return v + + model.forward = recorded + with torch.no_grad(): + final = sample_flow(model, sample(latent, camera, args.device), cond=cond, neg_cond=neg, steps=args.steps, + rescale_t=args.rescale_t, cfg_strength=args.cfg, verbose=False, batch_cfg=False) + model.forward = forward + pos, negs = calls[0::2], calls[1::2] + if args.cfg <= 0 or any(c["negative"] for c in pos) or not all(c["negative"] for c in negs): + sys.exit("the sampler's calls are not (conditional, negative) pairs") + stack = lambda rows, key: torch.stack([r[key] for r in rows]).contiguous() + tensors = {"feature1": f1[0].contiguous(), "feature2": f2[0].contiguous(), "latent": latent, "camera": camera, + "timesteps": torch.tensor([c["t"] for c in pos], dtype=torch.float32), + "inputs": stack(pos, "latent"), "inputs_camera": stack(pos, "camera"), + "velocity": stack(pos, "v"), "velocity_camera": stack(pos, "v_camera"), + "velocity_negative": stack(negs, "v"), "velocity_negative_camera": stack(negs, "v_camera"), + "sample": final["latent"][0].float().cpu().contiguous(), + "sample_camera": final["camera"][0, 0].float().cpu().contiguous()} + save_file(tensors, str(out / "case.safetensors")) + meta = {"seed": args.seed, "steps": args.steps, "cfg": args.cfg, "rescale_t": args.rescale_t, + "device": args.device, "torch": torch.__version__, "sdk": SDK_COMMIT, "weights": WEIGHTS_REVISION, + "torch_forward_s": [round(c["s"], 3) for c in calls]} + (out / "reference.json").write_text(json.dumps(meta, indent=1) + "\n") + times = sorted(meta["torch_forward_s"]) + print(json.dumps({k: v for k, v in meta.items() if k != "torch_forward_s"} | {"median_forward_s": times[len(times) // 2]})) + return 0 + + +def blocks(args) -> int: + """Each block's rows in the conditional forward at one step, and each block again in float32 on the same input.""" + import copy + + import torch + from safetensors.torch import load_file, save_file + + case_dir = Path(args.case_dir) + data = load_file(str(case_dir / "case.safetensors")) + model = denoiser(args.weights, args.sdk, args.device) + cond, _ = condition(data["feature1"][None], data["feature2"][None], args.device) + seen = {} + + def keep(name): + def hook(_module, _inputs, output): + seen[name] = output.detach().contiguous() + return hook + + stages = [("noise_refiner", model.noise_refiner), ("context_refiner", model.context_refiner), ("block", model.blocks)] + hooks = [m.register_forward_hook(keep(f"{name}.{i}")) for name, layers in stages for i, m in enumerate(layers)] + x = sample(data["inputs"][args.step], data["inputs_camera"][args.step], args.device) + t = data["timesteps"][args.step : args.step + 1].to(args.device) + with torch.no_grad(): + v = model(x, t, cond) + for h in hooks: + h.remove() + # the ruler: each joint block in float32 on the fp16 rows the SDK fed it, with an unrounded modulation + t_emb = model.t_embedder(t) + t_mod = model.adaLN_modulation(t_emb).float() + exact = {} + for i in range(1, len(model.blocks)): + block = copy.deepcopy(model.blocks[i]).float() + h = seen[f"block.{i - 1}"].float() + exact[f"block.{i}"] = block(h, mod=t_mod, rotary_emb=model.repo_layers[i](h))[0].cpu().contiguous() + del block + tensors = {k: x[0].cpu().contiguous() for k, x in seen.items()} + tensors["velocity"] = v["latent"][0].float().cpu().contiguous() + tensors["velocity_camera"] = v["camera"][0, 0].float().cpu().contiguous() + save_file(tensors, str(case_dir / f"blocks{args.step}.safetensors")) + save_file(exact, str(case_dir / f"blocks{args.step}_f32.safetensors")) + drift = (tensors["velocity"] - data["velocity"][args.step]).abs().max().item() + print(json.dumps({"step": args.step, "tensors": len(tensors), "velocity_max_abs_vs_case": drift})) + return 0 + + +def exact(args) -> int: + """Both branches in float32 throughout on the case's own inputs: the fp16 reference's distance from exact.""" + import torch + from safetensors.torch import load_file, save_file + + case_dir = Path(args.case_dir) + data = load_file(str(case_dir / "case.safetensors")) + model = denoiser(args.weights, args.sdk, args.device).float() + model.dtype = torch.float32 + cond, neg = condition(data["feature1"][None], data["feature2"][None], args.device) + keys = ("velocity", "velocity_camera", "velocity_negative", "velocity_negative_camera") + rows = {k: [] for k in keys} + with torch.no_grad(): + for s in range(min(args.steps, data["inputs"].shape[0])): + x = sample(data["inputs"][s], data["inputs_camera"][s], args.device) + t = data["timesteps"][s : s + 1].to(args.device) + for prefix, c in (("velocity", cond), ("velocity_negative", neg)): + v = model(x, t, c) + rows[prefix].append(v["latent"][0].float().cpu()) + rows[prefix + "_camera"].append(v["camera"][0, 0].float().cpu()) + save_file({k + "_f32": torch.stack(v).contiguous() for k, v in rows.items()}, str(case_dir / "exact.safetensors")) + print(json.dumps({"steps": len(rows["velocity"])})) + return 0 + + +def main(argv: list[str]) -> int: + ap = argparse.ArgumentParser(description=__doc__, epilog=USAGE, formatter_class=argparse.RawDescriptionHelpFormatter) + sub = ap.add_subparsers(dest="cmd", required=True) + for name in ("case", "blocks", "exact"): + p = sub.add_parser(name) + p.add_argument("weights") + p.add_argument("sdk") + p.add_argument("out_dir" if name == "case" else "case_dir") + p.add_argument("--device", default="mps", choices=("mps", "cpu", "cuda")) + if name == "case": + p.add_argument("--features", nargs="+", required=True) + p.add_argument("--seed", type=int, default=42) + p.add_argument("--steps", type=int, default=50) + p.add_argument("--cfg", type=float, default=3.0) + p.add_argument("--rescale-t", type=float, default=3.0) + elif name == "blocks": + p.add_argument("--step", type=int, default=0) + else: + p.add_argument("--steps", type=int, default=50) + args = ap.parse_args(argv) + return {"case": case, "blocks": blocks, "exact": exact}[args.cmd](args) + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1:])) diff --git a/zig/build/metal_targets.zig b/zig/build/metal_targets.zig index 5bac04a9b0..f08c980d2f 100644 --- a/zig/build/metal_targets.zig +++ b/zig/build/metal_targets.zig @@ -317,6 +317,7 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s .{ .name = "tf-qwen-vision", .path = "zig/tests/qwen_vision.zig", .about = "The Qwen-VL tower on picture files: patches, grids and merged rows written for a float32 reference" }, .{ .name = "tf-qwen-image", .path = "zig/tests/qwen_image.zig", .about = "Qwen-Image-2.1's transformer on a reference case: each step compared, then the schedule timed" }, .{ .name = "tf-moss-image", .path = "zig/tests/moss_image.zig", .about = "MOSS-Image's transformer on a reference case: each step's two branches compared, then the guided schedule timed" }, + .{ .name = "tf-idol", .path = "zig/tests/idol.zig", .about = "Idol's denoiser on a reference case: each step's two branches compared, then the guided schedule timed" }, .{ .name = "tf-glm-attn-probe", .path = "zig/tests/glm_attn_probe.zig", .about = "GLM-5.3-Flash's dense latent attention kernels on given inputs (scores, probabilities, outputs)" }, .{ .name = "tf-moe-bench", .path = "zig/tests/moe_bench.zig", .about = "Routed-expert prompt kernels from .metal files on a skewed routing: chained speed, bits against the first" }, .{ .name = "tf-mm-bench", .path = "zig/tests/mm_bench.zig", .about = "Prompt matmul kernels from .metal files: chained speed at one shape, bits against the first, error against fp64" }, diff --git a/zig/kernels/metal/idol.metal b/zig/kernels/metal/idol.metal new file mode 100644 index 0000000000..1a8048d1b2 --- /dev/null +++ b/zig/kernels/metal/idol.metal @@ -0,0 +1,329 @@ +// Idol's denoiser steps after gemm_nax.metal: fp16 rows rounded where idol-sdk's torch code rounds them. + +inline float id_h(float x) { return float(half(x)); } + +// One threadgroup's sum of every thread's `v` (at most 1,024 threads), returned to all of them. +inline float id_sum(float v, threadgroup float* part, uint n, uint lane, uint sg) { + v = simd_sum(v); + if (lane == 0) part[sg] = v; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + float s = lane < (n + 31) / 32 ? part[lane] : 0.0f; + s = simd_sum(s); + if (lane == 0) part[32] = s; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const float out = part[32]; + threadgroup_barrier(mem_flags::mem_threadgroup); + return out; +} + +// torch's tanh GELU of an fp16 value, computed in fp32. +inline float id_gelu(float x) { + const float inner = 0.7978845608028654f * (x + 0.044715f * x * x * x); + return 0.5f * x * (1.0f + metal::precise::tanh(inner)); +} + +// ---- y [M, N] = fp16(x w^T + b) (EPI 1: then GELU, rounded again): tfp::gemm_nax's block, fp32 sums ---- +template +inline void id_linear(const device half* A, const device half* B, const device half* bias, device half* D, + const device int* P, uint3 tg, uint sg, uint lane) { + constexpr int BM = 64, BN = 128, BK = 256, WN = 4; + const int tx = int(tg.x), ty = int(tg.y); + if (tx >= P[6] || ty >= P[7]) return; + const int lda = P[3], ldb = P[4], ldd = P[5]; + const int row = ty * BM + 32 * int(sg / WN), col = tx * BN + 32 * int(sg % WN); + A += long(row) * lda; + B += long(col) * ldb; + D += long(row) * ldd + col; + const int rows = AM ? 32 : min(32, P[0] - row), cols = AN ? 32 : min(32, P[1] - col); + const bool full = (AM || rows == 32) && (AN || cols == 32); + const short2 home = tfp::frag_home(ushort(lane)); + tfp::frag acc[2][2]; + if (full) { + tfp::gemm_block(acc, A, B, lda, ldb, P[2], P[9], rows, cols, home); + } else { + tfp::gemm_block(acc, A, B, lda, ldb, P[2], P[9], rows, cols, home); + } + TF_UNROLL + for (short j = 0; j < 2; j++) { + TF_UNROLL + for (short e = 0; e < 8; e++) { + const int c = col + 16 * j + home.x + TF_COL(e); + const float b = c < P[1] ? float(bias[c]) : 0.0f; + TF_UNROLL + for (short i = 0; i < 2; i++) { + const float v = id_h(acc[i][j][e] + b); + acc[i][j][e] = EPI == 1 ? id_gelu(v) : v; + } + } + } + if (full) { + tfp::gemm_store(acc, D, ldd, rows, cols, home); + } else if (rows > 0 && cols > 0) { + tfp::gemm_store(acc, D, ldd, rows, cols, home); + } +} + +#define ID_LINEAR(NAME, EPI, AM, AN) \ + [[max_total_threads_per_threadgroup(256)]] [[kernel]] void NAME( \ + const device half* A [[buffer(0)]], const device half* B [[buffer(1)]], const device int32_t* P [[buffer(2)]], \ + device half* D [[buffer(3)]], const device half* bias [[buffer(4)]], uint sg [[simdgroup_index_in_threadgroup]], \ + uint lane [[thread_index_in_simdgroup]], uint3 tg [[threadgroup_position_in_grid]]) { \ + id_linear(A, B, bias, D, P, tg, sg, lane); \ + } +ID_LINEAR(id_linear_00, 0, false, false) +ID_LINEAR(id_linear_01, 0, false, true) +ID_LINEAR(id_linear_10, 0, true, false) +ID_LINEAR(id_linear_11, 0, true, true) +ID_LINEAR(id_linear_gelu_00, 1, false, false) +ID_LINEAR(id_linear_gelu_01, 1, false, true) +ID_LINEAR(id_linear_gelu_10, 1, true, false) +ID_LINEAR(id_linear_gelu_11, 1, true, true) + +struct IdGemm32 { + uint M, N, K, ldy; +}; + +// y [M, N] (rows ldy apart) = x [M, K] w^T, fp32 FMA on 8x8 simdgroup matrices; K % 32, N % 64 (or N < 64, % 8) = 0. +[[max_total_threads_per_threadgroup(128)]] [[kernel]] void id_gemm32( + const device float* x [[buffer(0)]], const device float* w [[buffer(1)]], device float* y [[buffer(2)]], + constant IdGemm32& p [[buffer(3)]], uint2 tg [[threadgroup_position_in_grid]], uint t [[thread_index_in_threadgroup]], + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]]) { + constexpr uint BK = 32, LD = BK + 4; + threadgroup float ta[64 * LD]; + threadgroup float tb[64 * LD]; + threadgroup float stage[4][64]; + const uint row0 = tg.y * 64, col0 = tg.x * 64; + const uint r0 = (sg / 2) * 32, c0 = (sg % 2) * 32; + const uint cols = min(64u, p.N - col0); + simdgroup_float8x8 acc[4][4]; + for (uint i = 0; i < 4; i++) + for (uint j = 0; j < 4; j++) acc[i][j] = make_filled_simdgroup_matrix(0.0f); + for (uint k = 0; k < p.K; k += BK) { + // 64 x 32 of x and of w: 2,048 values each, 16 a thread as float4s; rows past M or N read as zeros + for (uint e = t; e < 64 * BK / 4; e += 128) { + const uint r = e / (BK / 4), c = (e % (BK / 4)) * 4; + const float4 va = row0 + r < p.M ? *(const device float4*)(x + ulong(row0 + r) * p.K + k + c) : float4(0.0f); + const float4 vb = r < cols ? *(const device float4*)(w + ulong(col0 + r) * p.K + k + c) : float4(0.0f); + *(threadgroup float4*)(ta + r * LD + c) = va; + *(threadgroup float4*)(tb + r * LD + c) = vb; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint kk = 0; kk < BK; kk += 8) { + simdgroup_float8x8 a[4], b[4]; + for (uint i = 0; i < 4; i++) simdgroup_load(a[i], ta + (r0 + 8 * i) * LD + kk, LD); + for (uint j = 0; j < 4; j++) simdgroup_load(b[j], tb + (c0 + 8 * j) * LD + kk, LD, ulong2(0, 0), true); + for (uint i = 0; i < 4; i++) + for (uint j = 0; j < 4; j++) simdgroup_multiply_accumulate(acc[i][j], a[i], b[j], acc[i][j]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + for (uint i = 0; i < 4; i++) { + for (uint j = 0; j < 4; j++) { + const uint r = row0 + r0 + 8 * i, c = c0 + 8 * j; + if (r >= p.M || c >= cols) continue; + if (r + 8 <= p.M) { + simdgroup_store(acc[i][j], y + ulong(r) * p.ldy + col0 + c, p.ldy); + } else { + simdgroup_store(acc[i][j], stage[sg], 8); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (uint e = lane; e < 64; e += 32) + if (r + e / 8 < p.M) y[ulong(r + e / 8) * p.ldy + col0 + c + e % 8] = stage[sg][e]; + simdgroup_barrier(mem_flags::mem_threadgroup); + } + } + } +} + +// c [M, N] fp32 = a [M, K] b^T over fp16 operands, tfp::gemm_splitk_nax's one partition (fp32 sums). +#define ID_SPLIT(NAME, AM) \ + [[max_total_threads_per_threadgroup(128)]] [[kernel]] void NAME( \ + const device half* A [[buffer(0)]], const device half* B [[buffer(1)]], const device int32_t* P [[buffer(2)]], \ + device float* C [[buffer(3)]], uint sg [[simdgroup_index_in_threadgroup]], \ + uint lane [[thread_index_in_simdgroup]], uint3 tg [[threadgroup_position_in_grid]]) { \ + tfp::gemm_splitk_nax(A, B, C, P, tg, sg, lane); \ + } +ID_SPLIT(id_split_0, false) +ID_SPLIT(id_split_1, true) + +struct IdNorm { + uint width; + float eps; + uint mode; // 0 plain, 1 affine (fp32 w, b), 2 modulated (fp16 1 + scale, shift) +}; + +// LayerNorm32 of an fp16 row in fp32, rounded to fp16; then affine in fp32, or modulated in fp16. A threadgroup a row. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void id_norm( + const device half* x [[buffer(0)]], const device float* w [[buffer(1)]], const device float* b [[buffer(2)]], + const device half* s1 [[buffer(3)]], const device half* sh [[buffer(4)]], device half* y [[buffer(5)]], + constant IdNorm& p [[buffer(6)]], uint row [[threadgroup_position_in_grid]], uint t [[thread_position_in_threadgroup]], + uint n [[threads_per_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint sg [[simdgroup_index_in_threadgroup]]) { + threadgroup float part[33]; + const device half* xr = x + ulong(row) * p.width; + device half* yr = y + ulong(row) * p.width; + float a = 0.0f; + for (uint i = t; i < p.width; i += n) a += float(xr[i]); + const float mean = id_sum(a, part, n, lane, sg) / float(p.width); + float q = 0.0f; + for (uint i = t; i < p.width; i += n) q += (float(xr[i]) - mean) * (float(xr[i]) - mean); + const float inv = metal::precise::rsqrt(id_sum(q, part, n, lane, sg) / float(p.width) + p.eps); + for (uint i = t; i < p.width; i += n) { + float v = (float(xr[i]) - mean) * inv; + if (p.mode == 1) v = v * w[i] + b[i]; + v = id_h(v); + if (p.mode == 2) v = id_h(id_h(v * float(s1[i])) + float(sh[i])); + yr[i] = half(v); + } +} + +// Affine LayerNorm32 of an fp16 row as fp16 [hi | hi | lo] of 3 * width: times [w_hi | w_lo | w_hi], near-fp32 sums. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void id_norm_split( + const device half* x [[buffer(0)]], const device float* w [[buffer(1)]], const device float* b [[buffer(2)]], + device half* y [[buffer(3)]], constant IdNorm& p [[buffer(4)]], uint row [[threadgroup_position_in_grid]], + uint t [[thread_position_in_threadgroup]], uint n [[threads_per_threadgroup]], + uint lane [[thread_index_in_simdgroup]], uint sg [[simdgroup_index_in_threadgroup]]) { + threadgroup float part[33]; + const device half* xr = x + ulong(row) * p.width; + device half* yr = y + ulong(row) * 3 * p.width; + float a = 0.0f; + for (uint i = t; i < p.width; i += n) a += float(xr[i]); + const float mean = id_sum(a, part, n, lane, sg) / float(p.width); + float q = 0.0f; + for (uint i = t; i < p.width; i += n) q += (float(xr[i]) - mean) * (float(xr[i]) - mean); + const float inv = metal::precise::rsqrt(id_sum(q, part, n, lane, sg) / float(p.width) + p.eps); + for (uint i = t; i < p.width; i += n) { + const float v = (float(xr[i]) - mean) * inv * w[i] + b[i]; + const half hi = half(v); + yr[i] = hi; + yr[p.width + i] = hi; + yr[2 * p.width + i] = half(v - float(hi)); + } +} + +// LayerNorm of an fp16 row in fp32 kept in fp32: affine (w, b) when given, else times (1 + scale) plus shift (mode 2). +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void id_norm32( + const device half* x [[buffer(0)]], const device float* w [[buffer(1)]], const device float* b [[buffer(2)]], + device float* y [[buffer(3)]], constant IdNorm& p [[buffer(4)]], uint row [[threadgroup_position_in_grid]], + uint t [[thread_position_in_threadgroup]], uint n [[threads_per_threadgroup]], + uint lane [[thread_index_in_simdgroup]], uint sg [[simdgroup_index_in_threadgroup]]) { + threadgroup float part[33]; + const device half* xr = x + ulong(row) * p.width; + device float* yr = y + ulong(row) * p.width; + float a = 0.0f; + for (uint i = t; i < p.width; i += n) a += float(xr[i]); + const float mean = id_sum(a, part, n, lane, sg) / float(p.width); + float q = 0.0f; + for (uint i = t; i < p.width; i += n) q += (float(xr[i]) - mean) * (float(xr[i]) - mean); + const float inv = metal::precise::rsqrt(id_sum(q, part, n, lane, sg) / float(p.width) + p.eps); + for (uint i = t; i < p.width; i += n) { + const float v = (float(xr[i]) - mean) * inv; + yr[i] = p.mode == 2 ? v * (1.0f + w[i]) + b[i] : v * w[i] + b[i]; + } +} + +struct IdRepo { + uint rows; + uint stride; // fp32 values from one row's gate sums to the next; content sums follow gate's 128 + float inv; // the maps' power-of-two scale undone, exactly +}; + +// RePo: d = final_map(silu(gate) * content), 3 per head; angles clamp_mul(d, freqs) * pi as cos, sin [rows, 16, 32]. +[[max_total_threads_per_threadgroup(64)]] [[kernel]] void id_repo( + const device float* gc [[buffer(0)]], const device float* wf [[buffer(1)]], const device float* f [[buffer(2)]], + device float* cs [[buffer(3)]], device float* sn [[buffer(4)]], constant IdRepo& p [[buffer(5)]], + uint row [[threadgroup_position_in_grid]], uint t [[thread_position_in_threadgroup]]) { + threadgroup float u[128]; + threadgroup float d[48]; + const device float* g = gc + ulong(row) * p.stride; + for (uint i = t; i < 128; i += 64) { + const float a = g[i] * p.inv; + u[i] = (a / (1.0f + metal::precise::exp(-a))) * (g[128 + i] * p.inv); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (t < 48) { + float s = 0.0f; + for (uint k = 0; k < 128; k++) s = fma(u[k], wf[t * 128 + k], s); + d[t] = s; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // freqs: 10 x, 10 y, 12 z; each pair of the 32 takes its axis's offset + for (uint i = t; i < 512; i += 64) { + const uint h = i / 32, j = i % 32; + const uint axis = j < 10 ? 0 : (j < 20 ? 1 : 2); + const float x = d[h * 3 + axis], fr = f[j]; + const float ft = metal::precise::tanh(fr); + const float ang = (x * ft + x * (fr - ft)) * 3.14159265358979323846f; + cs[ulong(row) * 512 + i] = metal::precise::cos(ang); + sn[ulong(row) * 512 + i] = metal::precise::sin(ang); + } +} + +struct IdHeads { + uint rows; + float scale; // sqrt(head width): MultiHeadRMSNorm's +}; + +// q, k of qkv rows: pairs turned by the row and head's cos, sin; F.normalize * gamma * scale; out [16, rows, 64]. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void id_qk( + const device half* qkv [[buffer(0)]], const device float* gq [[buffer(1)]], const device float* gk [[buffer(2)]], + const device float* cs [[buffer(3)]], const device float* sn [[buffer(4)]], device half* Q [[buffer(5)]], + device half* K [[buffer(6)]], device half* V [[buffer(7)]], constant IdHeads& p [[buffer(8)]], + uint2 tg [[threadgroup_position_in_grid]], uint lane [[thread_index_in_simdgroup]], + uint sg [[simdgroup_index_in_threadgroup]]) { + const uint r = tg.x, h = tg.y * 8 + sg; + if (h >= 16) return; + const ulong from = ulong(r) * 3072 + h * 64 + lane * 2; + const ulong at = ulong(r) * 512 + h * 32 + lane; + const float c = cs[at], s = sn[at]; + const float q0 = float(qkv[from]), q1 = float(qkv[from + 1]); + const float k0 = float(qkv[from + 1024]), k1 = float(qkv[from + 1025]); + const float rq0 = id_h(q0 * c - q1 * s), rq1 = id_h(q0 * s + q1 * c); + const float rk0 = id_h(k0 * c - k1 * s), rk1 = id_h(k0 * s + k1 * c); + const float qn = max(metal::precise::sqrt(simd_sum(rq0 * rq0 + rq1 * rq1)), 1e-12f); + const float kn = max(metal::precise::sqrt(simd_sum(rk0 * rk0 + rk1 * rk1)), 1e-12f); + const uint g = h * 64 + lane * 2; + const ulong to = (ulong(h) * p.rows + r) * 64 + lane * 2; + Q[to] = half(rq0 / qn * gq[g] * p.scale); + Q[to + 1] = half(rq1 / qn * gq[g + 1] * p.scale); + K[to] = half(rk0 / kn * gk[g] * p.scale); + K[to + 1] = half(rk1 / kn * gk[g + 1] * p.scale); + V[to] = qkv[from + 2048]; + V[to + 1] = qkv[from + 2049]; +} + +// x = fp16(x + fp16(g * y)) by column; g is a modulation gate, or ones for an ungated residual. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void id_gate_add( + device half* x [[buffer(0)]], const device half* y [[buffer(1)]], const device half* g [[buffer(2)]], + constant uint& width [[buffer(3)]], uint2 at [[thread_position_in_grid]]) { + const ulong i = ulong(at.y) * width + at.x; + x[i] = half(float(x[i]) + id_h(float(g[at.x]) * float(y[i]))); +} + +// x = fp16(fp16(z w^T + b) + pe): the latent's input projection (16 inputs, fp32) and its fp16 position embedding. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void id_embed( + const device float* z [[buffer(0)]], const device float* w [[buffer(1)]], const device float* b [[buffer(2)]], + const device half* pe [[buffer(3)]], device half* x [[buffer(4)]], uint2 at [[thread_position_in_grid]]) { + const uint c = at.x, r = at.y; + float s = 0.0f; + for (uint k = 0; k < 16; k++) s = fma(z[r * 16 + k], w[c * 16 + k], s); + const ulong i = ulong(r) * 1024 + c; + x[i] = half(float(half(s + b[c])) + float(pe[i])); +} + +// y = fp16(a + ba + c + bc) by column: the two condition projections' fp32 sums and biases (each torch Linear). +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void id_cond( + const device float* a [[buffer(0)]], const device float* ba [[buffer(1)]], const device float* c [[buffer(2)]], + const device float* bc [[buffer(3)]], device half* y [[buffer(4)]], constant uint& width [[buffer(5)]], + uint2 at [[thread_position_in_grid]]) { + const ulong i = ulong(at.y) * width + at.x; + y[i] = half((a[i] + ba[at.x]) + (c[i] + bc[at.x])); +} + +// y = x + b by column in fp32 (an fp32 Linear's bias), in place. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void id_bias32( + device float* y [[buffer(0)]], const device float* b [[buffer(1)]], constant uint& width [[buffer(2)]], + uint2 at [[thread_position_in_grid]]) { + y[ulong(at.y) * width + at.x] += b[at.x]; +} diff --git a/zig/kernels/metal/idol_attention.metal b/zig/kernels/metal/idol_attention.metal new file mode 100644 index 0000000000..f676630fc5 --- /dev/null +++ b/zig/kernels/metal/idol_attention.metal @@ -0,0 +1,149 @@ +// Idol's attention: tfp::attention_nax's online softmax at head width 64, every query row to every key row. + +constant constexpr float kIdMasked = -3.402823466e+38f; +constant constexpr int kIdAnyCols = 1 << 30; + +// Fold one fragment's rows into r[2] (rows home.y, home.y + 8): pairs, then lanes 1 and 8 apart, then into r. +template +inline void id_fold(thread const tfp::frag& f, thread float (&r)[2]) { + TF_UNROLL + for (short h = 0; h < 2; h++) { + const short b = 4 * h; + float t = MAX ? max(max(f[b], f[b + 1]), max(f[b + 2], f[b + 3])) : (f[b] + f[b + 1]) + (f[b + 2] + f[b + 3]); + const float u = simd_shuffle_xor(t, ushort(1)); + t = MAX ? max(t, u) : t + u; + const float w = simd_shuffle_xor(t, ushort(8)); + t = MAX ? max(t, w) : t + w; + r[h] = MAX ? max(r[h], t) : r[h] + t; + } +} + +// One 64-query block of one head over 32-key blocks; P and F as Kernels.attention sets them (kernels.zig). +template +inline void id_attention(const device T* Q, const device T* K, const device T* V, device T* O, const device int* P, + const device float* F, uint3 tg, uint sg, uint lane) { + constexpr int BQ = 64, BK = 32, D = 64; + const int h = int(tg.y), lq = P[1], lk = P[2]; + Q += (long(h) * lq + long(tg.x) * BQ) * D; + K += long(h) * lk * D; + V += long(h) * lk * D; + const int os = P[0] * D; + O += long(tg.x) * BQ * os + h * D; + const float scale2 = F[0] * 1.44269504089f; + const short tm = 16 * short(sg); + Q += tm * D; + const short2 home = tfp::frag_home(ushort(lane)); + + tfp::frag acc[4]; + TF_UNROLL + for (short i = 0; i < 4; i++) acc[i] = tfp::frag(0); + float top[2] = {kIdMasked, kIdMasked}, total[2] = {0.0f, 0.0f}; + const bool last_q = int(tg.x) == P[4]; + const short q_rows = short(P[6] - tm), k_rows = short(P[7]); + + for (int kb = 0; kb < P[3]; kb++) { + const bool last_k = kb == P[5]; + tfp::frag s[2] = {tfp::frag(0), tfp::frag(0)}; + TF_UNROLL + for (short d = 0; d < 4; d++) { + tfp::frag q, k0, k1; + if (!AQ && last_q) { + tfp::frag_get_in(q, Q, D, 0, 16 * d, home, q_rows, kIdAnyCols); + } else { + tfp::frag_get(q, Q, D, 0, 16 * d, home); + } + if (!AK && last_k) { + tfp::frag_get_t_in(k0, K, D, 0, 16 * d, home, k_rows, kIdAnyCols); + tfp::frag_get_t_in(k1, K, D, 16, 16 * d, home, k_rows, kIdAnyCols); + } else { + tfp::frag_get_t(k0, K, D, 0, 16 * d, home); + tfp::frag_get_t(k1, K, D, 16, 16 * d, home); + } + tfp::mma_16x32(s[0], s[1], q, k0, k1); + } + TF_UNROLL + for (short f = 0; f < 2; f++) { + TF_UNROLL + for (short e = 0; e < 8; e++) s[f][e] *= scale2; + } + if (!AK && last_k) { + TF_UNROLL + for (short f = 0; f < 2; f++) { + TF_UNROLL + for (short e = 0; e < 8; e++) s[f][e] = (16 * f + home.x + TF_COL(e)) < k_rows ? s[f][e] : kIdMasked; + } + } + float top_new[2] = {top[0], top[1]}; + id_fold(s[0], top_new); + id_fold(s[1], top_new); + TF_UNROLL + for (short f = 0; f < 2; f++) { + TF_UNROLL + for (short e = 0; e < 8; e++) s[f][e] = fast::exp2(s[f][e] - top_new[e >> 2]); + } + float scale_old[2]; + TF_UNROLL + for (short r = 0; r < 2; r++) { + scale_old[r] = fast::exp2(top[r] - top_new[r]); + top[r] = top_new[r]; + total[r] = total[r] * scale_old[r]; + } + id_fold(s[0], total); + id_fold(s[1], total); + TF_UNROLL + for (short i = 0; i < 4; i++) { + TF_UNROLL + for (short e = 0; e < 8; e++) acc[i][e] = acc[i][e] * scale_old[e >> 2]; + } + simdgroup_barrier(mem_flags::mem_none); + TF_UNROLL + for (short d = 0; d < 4; d += 2) { + TF_UNROLL + for (short k = 0; k < 2; k++) { + tfp::frag v0, v1; + if (!AK && last_k) { + tfp::frag_get_in(v0, V, D, 16 * k, 16 * d, home, k_rows, kIdAnyCols); + tfp::frag_get_in(v1, V, D, 16 * k, 16 * d + 16, home, k_rows, kIdAnyCols); + } else { + tfp::frag_get(v0, V, D, 16 * k, 16 * d, home); + tfp::frag_get(v1, V, D, 16 * k, 16 * d + 16, home); + } + tfp::mma_16x32(acc[d], acc[d + 1], s[k], v0, v1); + } + } + K += BK * D; + V += BK * D; + } + + threadgroup_barrier(mem_flags::mem_none); + float inv[2]; + TF_UNROLL + for (short r = 0; r < 2; r++) inv[r] = 1.f / total[r]; + TF_UNROLL + for (short i = 0; i < 4; i++) { + TF_UNROLL + for (short e = 0; e < 8; e++) acc[i][e] = acc[i][e] * inv[e >> 2]; + } + O += tm * os; + if (!AQ && last_q) { + if (q_rows <= 0) return; + TF_UNROLL + for (short i = 0; i < 4; i++) tfp::frag_put_in(acc[i], O, os, 0, 16 * i, home, q_rows, kIdAnyCols); + } else { + TF_UNROLL + for (short i = 0; i < 4; i++) tfp::frag_put(acc[i], O, os, 0, 16 * i, home); + } +} + +#define ID_ATTENTION(NAME, T, AQ, AK) \ + [[max_total_threads_per_threadgroup(128)]] [[kernel]] void NAME( \ + const device T* Q [[buffer(0)]], const device T* K [[buffer(1)]], const device T* V [[buffer(2)]], \ + const device int32_t* P [[buffer(3)]], const device float* F [[buffer(4)]], device T* O [[buffer(5)]], \ + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], \ + uint3 tg [[threadgroup_position_in_grid]]) { \ + id_attention(Q, K, V, O, P, F, tg, sg, lane); \ + } +ID_ATTENTION(id_attention_00, half, false, false) +ID_ATTENTION(id_attention_01, half, false, true) +ID_ATTENTION(id_attention_10, half, true, false) +ID_ATTENTION(id_attention_11, half, true, true) diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig index f5d5861ad8..ef3221a88b 100644 --- a/zig/kernels/metal/sources.zig +++ b/zig/kernels/metal/sources.zig @@ -133,6 +133,8 @@ pub const core_qwen_vision = @embedFile("prefill/gemm_nax.metal") ++ @embedFile( pub const qwen_image = @embedFile("prefill/attention_nax.metal") ++ @embedFile("qwen_image.metal"); /// MOSS-Image's transformer: the prefill attention template, then its own steps (projections: core_qwen_vision). pub const moss_image = @embedFile("prefill/attention_nax.metal") ++ @embedFile("moss_image.metal"); +/// Idol's denoiser: the prefill GEMM template, then its own steps and head-64 attention. +pub const idol = @embedFile("prefill/gemm_nax.metal") ++ @embedFile("idol.metal") ++ @embedFile("idol_attention.metal"); /// A prompt chunk's KDA layer in three passes, appended to the generated kda_rows source (its helpers). pub const glm_kda_prompt = @embedFile("glm_kda_prompt.metal"); /// GLM-5.3-Flash speed-up mode's exchange kernels (families/glm/ep.zig). diff --git a/zig/src/families/idol/config.zig b/zig/src/families/idol/config.zig new file mode 100644 index 0000000000..d399755886 --- /dev/null +++ b/zig/src/families/idol/config.zig @@ -0,0 +1,71 @@ +//! Idol's denoiser geometry, and the check that a snapshot's config.json names that denoiser. +const std = @import("std"); + +pub const hidden = 1024; +pub const heads = 16; +pub const head_dim = hidden / heads; +pub const ffn = 4 * hidden; +pub const layers = 24; +/// The noise and condition refiners' blocks, each. +pub const refiners = 2; +pub const latent_rows = 8192; +pub const latent_channels = 16; +pub const camera_channels = 5; +/// DINOv3 ViT-H+/16 rows at 1024 x 1024: 64 x 64 patches, the class token and four registers. +pub const cond_rows = 4101; +pub const cond_channels = 1280; +pub const cond2_channels = 128; +/// Latent, condition and camera rows through the joint blocks. +pub const joint_rows = latent_rows + cond_rows + 1; +/// RePo's hidden width (model_channels / 8) and its frequencies per head: 10, 10 and 12 for x, y and z. +pub const repo_hidden = hidden / 8; +pub const repo_pairs = head_dim / 2; +pub const time_dim = 256; +/// Every transformer block's LayerNorm32, fixed in idol-sdk's layers.py, not a config setting. +pub const block_eps: f32 = 1e-6; +/// RePo's LayerNorm32 and the output's F.layer_norm take torch's default. +pub const default_eps: f32 = 1e-5; +/// Modulated blocks: the noise refiner's, then the joint blocks. +pub const modulated = refiners + layers; + +/// Accepts the snapshot config whose denoiser this family runs (LatentSeqMMFlowModel at its released size). +pub fn check(gpa: std.mem.Allocator, text: []const u8) !void { + const doc = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer doc.deinit(); + if (doc.value != .object) return error.BadConfig; + const den = doc.value.object.get("denoiser") orelse return error.BadConfig; + if (den != .object) return error.BadConfig; + const name = den.object.get("name") orelse return error.BadConfig; + if (name != .string or !std.mem.eql(u8, name.string, "LatentSeqMMFlowModel")) return error.NotIdol; + const args = den.object.get("args") orelse return error.BadConfig; + if (args != .object) return error.BadConfig; + const want = .{ .{ "q_token_length", latent_rows }, .{ "in_channels", latent_channels }, .{ "cam_channels", camera_channels }, .{ "model_channels", hidden }, .{ "cond_channels", cond_channels }, .{ "cond2_channels", cond2_channels }, .{ "num_blocks", layers }, .{ "num_refiner_blocks", refiners }, .{ "num_heads", heads }, .{ "mlp_ratio", 4 } }; + inline for (want) |kv| { + const v = args.object.get(kv[0]) orelse return error.BadConfig; + if (v != .integer or v.integer != kv[1]) { + std.log.warn("idol: denoiser {s} is not {d}; this family runs the released geometry only", .{ kv[0], kv[1] }); + return error.UnsupportedGeometry; + } + } +} + +const sample = + \\{"denoiser": {"name": "LatentSeqMMFlowModel", "args": {"q_token_length": 8192, "in_channels": 16, + \\ "cam_channels": 5, "model_channels": 1024, "cond_channels": 1280, "cond2_channels": 128, "num_blocks": 24, + \\ "num_refiner_blocks": 2, "num_heads": 16, "mlp_ratio": 4}}, "image_size": 1024, "max_voxel_level": 8} +; + +test "the released denoiser config passes; another depth, width or class does not" { + const a = std.testing.allocator; + try check(a, sample); + const deeper = try std.mem.replaceOwned(u8, a, sample, "\"num_blocks\": 24", "\"num_blocks\": 12"); + defer a.free(deeper); + try std.testing.expectError(error.UnsupportedGeometry, check(a, deeper)); + const missing = try std.mem.replaceOwned(u8, a, sample, "\"mlp_ratio\": 4", "\"mlp\": 4"); + defer a.free(missing); + try std.testing.expectError(error.BadConfig, check(a, missing)); + const other = try std.mem.replaceOwned(u8, a, sample, "LatentSeqMMFlowModel", "LatentFlowModel"); + defer a.free(other); + try std.testing.expectError(error.NotIdol, check(a, other)); + try std.testing.expectError(error.BadConfig, check(a, "{\"decoder\": {}}")); +} diff --git a/zig/src/families/idol/host.zig b/zig/src/families/idol/host.zig new file mode 100644 index 0000000000..33dd008649 --- /dev/null +++ b/zig/src/families/idol/host.zig @@ -0,0 +1,147 @@ +//! The host's part of an Idol denoiser pass: the schedule, the time modulation, the camera token and its velocity. +const std = @import("std"); +const cfg = @import("config.zig"); +const Host = @import("weights.zig").Host; + +const H = cfg.hidden; + +pub fn half(x: f32) f32 { + return @as(f16, @floatCast(x)); +} + +/// y[o] = sum over k of x[k] w[o, k] + b[o] in fp32, w stored (outputs, inputs). +pub fn project(w: []const f32, b: []const f32, x: []const f32, y: []f32) void { + std.debug.assert(w.len == x.len * y.len and b.len == y.len); + const V = @Vector(8, f32); + for (y, 0..) |*o, row| { + const ws = w[row * x.len ..][0..x.len]; + var acc: V = @splat(0); + var k: usize = 0; + while (k + 8 <= x.len) : (k += 8) acc += @as(V, ws[k..][0..8].*) * @as(V, x[k..][0..8].*); + var sum = @reduce(.Add, acc); + while (k < x.len) : (k += 1) sum += ws[k] * x[k]; + o.* = sum + b[row]; + } +} + +fn silu(x: f32) f32 { + return x / (1.0 + @exp(-x)); +} + +fn gelu(x: f32) f32 { + return 0.5 * x * (1.0 + std.math.tanh(0.7978845608028654 * (x + 0.044715 * x * x * x))); +} + +/// The sampler's times: t = s u / (1 + (s - 1) u) for u from 1 to 0 in `steps` (steps + 1 values), in float64. +pub fn schedule(out: []f64, rescale: f64) void { + const steps = out.len - 1; + for (out, 0..) |*t, i| { + const u = 1.0 - @as(f64, @floatFromInt(i)) / @as(f64, @floatFromInt(steps)); + t.* = rescale * u / (1.0 + (rescale - 1.0) * u); + } +} + +/// The modulated blocks' six fp16 vectors each (shift, 1 + scale, gate for attention, then for the MLP). +pub const vectors = 6; +pub const Mods = [cfg.modulated][vectors][H]f16; + +/// t_embedder's output [1024] at model time `t` (1000 times the flow time), in fp32. +pub fn timeEmbedding(w: Host, t: f32, out: *[H]f32) void { + const halfdim = cfg.time_dim / 2; + var freq: [cfg.time_dim]f32 = undefined; + const lg: f32 = @floatCast(-@log(@as(f64, 10000))); + for (0..halfdim) |i| { + const a = t * @exp(lg * @as(f32, @floatFromInt(i)) / @as(f32, halfdim)); + freq[i] = @cos(a); + freq[halfdim + i] = @sin(a); + } + var mid: [H]f32 = undefined; + project(w.t1, w.t1_b, &freq, &mid); + for (&mid) |*v| v.* = silu(v.*); + project(w.t2, w.t2_b, &mid, out); +} + +/// Every modulated block's vectors and the output's fp32 shift and scale at model time `t`. +pub fn modulation(w: Host, t: f32, mods: *Mods, head: *[2][H]f32) void { + var emb: [H]f32 = undefined; + timeEmbedding(w, t, &emb); + var act: [H]f32 = undefined; + for (&act, emb) |*o, v| o.* = silu(v); + var tmod: [6 * H]f32 = undefined; + project(w.ada, w.ada_b, &act, &tmod); + for (&tmod) |*v| v.* = half(v.*); + for (mods, w.shifts) |*m, shift| { + for (0..6 * H) |i| { + const v = half(tmod[i] + half(shift[i])); + const part = i / H; + // the scales (parts 1 and 4) enter as fp16(1 + scale) + m[part][i % H] = @floatCast(if (part == 1 or part == 4) half(1.0 + v) else v); + } + } + for (0..2) |p| for (0..H) |i| { + head[p][i] = w.head[p * H + i] + emb[i]; + }; +} + +/// The camera token: cam_refiner's MLP of the five camera values, in fp32, then fp16. +pub fn cameraToken(w: Host, camera: *const [cfg.camera_channels]f32, out: []f16) void { + var mid: [H]f32 = undefined; + project(w.cam1, w.cam1_b, camera, &mid); + for (&mid) |*v| v.* = gelu(v.*); + var tok: [H]f32 = undefined; + project(w.cam2, w.cam2_b, &mid, &tok); + for (out, tok) |*o, v| o.* = @floatCast(v); +} + +/// The camera velocity from the last joint row: F.layer_norm (eps 1e-5), modulated by `head`, then cam_out_layer. +pub fn cameraVelocity(w: Host, row: []const f16, head: *const [2][H]f32, out: *[cfg.camera_channels]f32) void { + var x: [H]f32 = undefined; + var mean: f32 = 0; + for (&x, row) |*o, v| { + o.* = v; + mean += o.*; + } + mean /= H; + var q: f32 = 0; + for (x) |v| q += (v - mean) * (v - mean); + const inv = 1.0 / @sqrt(q / H + cfg.default_eps); + // shift is the table's first half, scale its second + for (&x, 0..) |*v, i| v.* = ((v.* - mean) * inv) * (1.0 + head[1][i]) + head[0][i]; + project(w.cam_out, w.cam_out_b, &x, out); +} + +/// The latents' fixed position embedding [8192, 1024] in fp16: sin then cos of each axis at 2^k frequencies. +pub fn positions(pos: []const f32, out: []f16) void { + std.debug.assert(pos.len == cfg.latent_rows * 3 and out.len == cfg.latent_rows * H); + const fd = H / 3 / 2; + const max_res = 16; + var freqs: [fd]f32 = undefined; + const rest = fd - max_res; + for (&freqs, 0..) |*f, i| { + const e: f32 = if (i < max_res) @floatFromInt(i) else @as(f32, @floatFromInt(i - max_res)) / @as(f32, rest) * max_res; + f.* = std.math.pow(f32, 2, e); + } + for (0..cfg.latent_rows) |r| { + const o = out[r * H ..][0..H]; + // outer(x.reshape(-1), freqs) then cat(sin, cos) per value: each axis's fd sines, then its fd cosines + for (0..3) |a| for (0..fd) |j| { + const ang = pos[r * 3 + a] * freqs[j] * 2 * std.math.pi; + o[a * 2 * fd + j] = @floatCast(@sin(ang)); + o[a * 2 * fd + fd + j] = @floatCast(@cos(ang)); + }; + for (o[6 * fd ..]) |*v| v.* = 0; + } +} + +test "the schedule runs from 1 to 0, denser at high noise for rescale 3" { + var t: [51]f64 = undefined; + schedule(&t, 3.0); + try std.testing.expectEqual(@as(f64, 1), t[0]); + try std.testing.expectEqual(@as(f64, 0), t[50]); + try std.testing.expectApproxEqAbs(@as(f64, 0.75), t[25], 1e-12); +} + +test "fp16 rounding is to nearest, ties to even" { + try std.testing.expectEqual(@as(f32, 1.0), half(1.0 + 1.0 / 4096.0)); + try std.testing.expectEqual(@as(f32, 1.0 + 1.0 / 512.0), half(1.0 + 3.0 / 2048.0)); +} diff --git a/zig/src/families/idol/idol.zig b/zig/src/families/idol/idol.zig new file mode 100644 index 0000000000..a2aff60944 --- /dev/null +++ b/zig/src/families/idol/idol.zig @@ -0,0 +1,19 @@ +//! Idol image to 3D: its denoiser on Metal (a denoiser beside the lane engine, not a token family). +pub const config = @import("config.zig"); +pub const host = @import("host.zig"); +pub const weights = @import("weights.zig"); +pub const kernels = @import("kernels.zig"); +pub const model = @import("model.zig"); +pub const reference = @import("reference.zig"); +pub const Model = model.Model; +pub const Condition = model.Condition; +pub const Timing = model.Timing; + +test { + _ = config; + _ = host; + _ = weights; + _ = kernels; + _ = model; + _ = reference; +} diff --git a/zig/src/families/idol/kernels.zig b/zig/src/families/idol/kernels.zig new file mode 100644 index 0000000000..2efe840cdc --- /dev/null +++ b/zig/src/families/idol/kernels.zig @@ -0,0 +1,209 @@ +//! Idol's Metal pipelines: its fp16 projections and steps, fp32 projections and head-64 attention. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +const frags = @import("../../core/frags.zig"); +const cfg = @import("config.zig"); +pub const Ref = @import("../../core/qwen_vl/weights.zig").Ref; + +/// Each pipeline by its function. +pub const names = [_][]const u8{ + "id_linear_00", + "id_linear_01", + "id_linear_10", + "id_linear_11", + "id_linear_gelu_00", + "id_linear_gelu_01", + "id_linear_gelu_10", + "id_linear_gelu_11", + "id_norm", + "id_norm32", + "id_repo", + "id_qk", + "id_gate_add", + "id_embed", + "id_cond", + "id_bias32", + "id_gemm32", + "id_split_0", + "id_split_1", + "id_norm_split", + "id_attention_00", + "id_attention_01", + "id_attention_10", + "id_attention_11", +}; +const linear_first = 0; +const attention_first = blk: { + for (names, 0..) |n, i| if (std.mem.eql(u8, n, "id_attention_00")) break :blk i; + unreachable; +}; + +pub const Norm = enum(u32) { plain = 0, affine = 1, modulated = 2 }; + +pub const Kernels = struct { + own: mtl.Library, + p: [names.len]mtl.Pipeline, + + pub fn load(gpa: std.mem.Allocator, device: mtl.Device) !Kernels { + var k: Kernels = undefined; + const own = try frags.source(device, gpa, sources.idol); + defer gpa.free(own); + k.own = try mtl.Library.fromSource(device, own, mtl.CompileOptions.mlx()); + errdefer k.own.deinit(); + var made: usize = 0; + errdefer for (k.p[0..made]) |p| p.deinit(); + for (names, 0..) |name, i| { + k.p[i] = try mtl.Pipeline.init(device, k.own, name, false); + made += 1; + } + return k; + } + + pub fn deinit(k: *Kernels) void { + for (k.p) |p| p.deinit(); + k.own.deinit(); + } + + pub fn get(k: *const Kernels, comptime name: []const u8) mtl.Pipeline { + inline for (names, 0..) |n, i| if (comptime std.mem.eql(u8, n, name)) return k.p[i]; + @compileError("no pipeline " ++ name); + } + + /// y [M, N] (rows `ldd` apart) = x [M, K] w^T in fp32 FMA, w as stored [N, K]; K % 32, N % 64 (or N < 64, % 8) = 0. + pub fn linear32(k: *const Kernels, e: mtl.ComputeEncoder, x: Ref, w: Ref, y: Ref, M: usize, N: usize, K: usize, ldd: usize) void { + std.debug.assert(K % 32 == 0 and (N % 64 == 0 or (N < 64 and N % 8 == 0))); + const params = [4]u32{ @intCast(M), @intCast(N), @intCast(K), @intCast(ldd) }; + bind(e, k.get("id_gemm32"), &.{ x, w, y }, params); + e.dispatchGroups(mtl.Size.of((N + 63) / 64, (M + 63) / 64, 1), mtl.Size.of(128, 1, 1)); + } + + /// c [M, 256] fp32 = x3 [M, 3072] w3^T over fp16: a split LayerNorm's rows and a RePo layer's split maps. + pub fn split(k: *const Kernels, e: mtl.ComputeEncoder, x3: Ref, w3: Ref, c: Ref, M: usize) void { + const N = 2 * cfg.repo_hidden; + const K = 3 * cfg.hidden; + const tn: usize = N / 64; + const tm = (M + 63) / 64; + const swz: u6 = if (tm <= 3) 0 else 1; + e.setPipeline(if (M % 64 == 0) k.get("id_split_1") else k.get("id_split_0")); + e.setBuffer(x3.buf, x3.off, 0); + e.setBuffer(w3.buf, w3.off, 1); + var p: [16]i32 = @splat(0); + const vals = [_]usize{ M, N, K, K, K, N, tn, tm, 1, M * N, K, swz, 0, 0 }; + for (vals, 0..) |v, i| p[i] = @intCast(v); + e.setBytes(std.mem.asBytes(&p), 2); + e.setBuffer(c.buf, c.off, 3); + const tiles = (tn << swz) * ((tm + (@as(usize, 1) << swz) - 1) >> swz); + e.dispatchThreads(mtl.Size.of(tiles * 32, 2, 2), mtl.Size.of(32, 2, 2)); + } + + /// x3 rows [hi | hi | lo] of the affine LayerNorm32 of x's rows, for `split`. + pub fn normSplit(k: *const Kernels, e: mtl.ComputeEncoder, x: Ref, w: Ref, b: Ref, x3: Ref, rows: usize, eps: f32) void { + const params = extern struct { width: u32, eps: f32, mode: u32 }{ .width = cfg.hidden, .eps = eps, .mode = 1 }; + bind(e, k.get("id_norm_split"), &.{ x, w, b, x3 }, params); + e.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(256, 1, 1)); + } + + /// y [M, N] = fp16(x w^T + b) in fp16 with fp32 sums, then GELU (rounded again) when `gelu`. + pub fn linear(k: *const Kernels, e: mtl.ComputeEncoder, x: Ref, w: Ref, b: Ref, y: Ref, M: usize, N: usize, K: usize, gelu: bool) void { + const tn = (N + 127) / 128; + const tm = (M + 63) / 64; + const form = @as(usize, @intFromBool(M % 64 == 0)) * 2 + @intFromBool(N % 128 == 0); + e.setPipeline(k.p[linear_first + @as(usize, @intFromBool(gelu)) * 4 + form]); + e.setBuffer(x.buf, x.off, 0); + e.setBuffer(w.buf, w.off, 1); + var p: [16]i32 = @splat(0); + const vals = [_]usize{ M, N, K, K, K, N, tn, tm, 0, K / 256 }; + for (vals, 0..) |v, i| p[i] = @intCast(v); + e.setBytes(std.mem.asBytes(&p), 2); + e.setBuffer(y.buf, y.off, 3); + e.setBuffer(b.buf, b.off, 4); + e.dispatchGroups(mtl.Size.of(tn, tm, 1), mtl.Size.of(256, 1, 1)); + } + + /// o rows [Lq, heads * 64] = softmax(q k^T / 8) v for q [heads, Lq, 64] and k, v [heads, Lk, 64]. + pub fn attention(k: *const Kernels, e: mtl.ComputeEncoder, q: Ref, kb: Ref, vb: Ref, o: Ref, lq: usize, lk: usize) void { + const form = @as(usize, @intFromBool(lq % 64 == 0)) * 2 + @intFromBool(lk % 32 == 0); + const nq = (lq + 63) / 64; + const nk = (lk + 31) / 32; + e.setPipeline(k.p[attention_first + form]); + for ([_]Ref{ q, kb, vb }, 0..) |r, i| e.setBuffer(r.buf, r.off, i); + // heads, q rows, k rows, k blocks, last q block, last k block, rows in the last q block and in the last k block + var params: [8]i32 = .{ cfg.heads, @intCast(lq), @intCast(lk), @intCast(nk), @intCast(nq - 1), @intCast(nk - 1), @intCast(lq - (nq - 1) * 64), @intCast(lk - (nk - 1) * 32) }; + e.setBytes(std.mem.asBytes(¶ms), 3); + var f: [4]f32 = .{ 1.0 / @sqrt(@as(f32, cfg.head_dim)), 0, 0, 0 }; + e.setBytes(std.mem.asBytes(&f), 4); + e.setBuffer(o.buf, o.off, 5); + e.dispatchGroups(mtl.Size.of(nq, cfg.heads, 1), mtl.Size.of(32, 4, 1)); + } + + /// y = LayerNorm32 of x's rows in fp32 rounded to fp16, then affine (w, b fp32) or modulated (s1, sh fp16). + pub fn norm(k: *const Kernels, e: mtl.ComputeEncoder, x: Ref, w: Ref, b: Ref, s1: Ref, sh: Ref, y: Ref, rows: usize, eps: f32, mode: Norm) void { + const params = extern struct { width: u32, eps: f32, mode: u32 }{ .width = cfg.hidden, .eps = eps, .mode = @backingInt(mode) }; + bind(e, k.get("id_norm"), &.{ x, w, b, s1, sh, y }, params); + e.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(256, 1, 1)); + } + + /// y fp32 = LayerNorm of x's rows times w plus b (affine), or times (1 + w) plus b (modulated). + pub fn norm32(k: *const Kernels, e: mtl.ComputeEncoder, x: Ref, w: Ref, b: Ref, y: Ref, rows: usize, eps: f32, mode: Norm) void { + const params = extern struct { width: u32, eps: f32, mode: u32 }{ .width = cfg.hidden, .eps = eps, .mode = @backingInt(mode) }; + bind(e, k.get("id_norm32"), &.{ x, w, b, y }, params); + e.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(256, 1, 1)); + } + + /// cos, sin [rows, 16, 32] of RePo from its gate and content sums (fp32 rows of 256) and final map. + pub fn repo(k: *const Kernels, e: mtl.ComputeEncoder, gc: Ref, inv: f32, final: Ref, freqs: Ref, cos: Ref, sin: Ref, rows: usize) void { + const params = extern struct { rows: u32, stride: u32, inv: f32 }{ .rows = @intCast(rows), .stride = 2 * cfg.repo_hidden, .inv = inv }; + bind(e, k.get("id_repo"), &.{ gc, final, freqs, cos, sin }, params); + e.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(64, 1, 1)); + } + + /// q, k, v heads [16, rows, 64] from packed qkv rows: rotated by cos, sin, then normed by gq, gk. + pub fn heads(k: *const Kernels, e: mtl.ComputeEncoder, qkv: Ref, gq: Ref, gk: Ref, cos: Ref, sin: Ref, q: Ref, kk: Ref, v: Ref, rows: usize) void { + const params = extern struct { rows: u32, scale: f32 }{ .rows = @intCast(rows), .scale = @sqrt(@as(f32, cfg.head_dim)) }; + bind(e, k.get("id_qk"), &.{ qkv, gq, gk, cos, sin, q, kk, v }, params); + e.dispatchGroups(mtl.Size.of(rows, 2, 1), mtl.Size.of(256, 1, 1)); + } + + /// x [rows, 1024] += y gated by the fp16 vector g, in place. + pub fn gateAdd(k: *const Kernels, e: mtl.ComputeEncoder, x: Ref, y: Ref, g: Ref, rows: usize) void { + bind(e, k.get("id_gate_add"), &.{ x, y, g }, @as(u32, cfg.hidden)); + e.dispatchThreads(mtl.Size.of(cfg.hidden, rows, 1), mtl.Size.of(256, 1, 1)); + } + + /// x [8192, 1024] = fp16(fp16(z w^T + b) + pe). + pub fn embed(k: *const Kernels, e: mtl.ComputeEncoder, z: Ref, w: Ref, b: Ref, pe: Ref, x: Ref) void { + e.setPipeline(k.get("id_embed")); + for ([_]Ref{ z, w, b, pe, x }, 0..) |r, i| e.setBuffer(r.buf, r.off, i); + e.dispatchThreads(mtl.Size.of(cfg.hidden, cfg.latent_rows, 1), mtl.Size.of(256, 1, 1)); + } + + /// y [rows, 1024] = fp16((a + ba) + (c + bc)). + pub fn cond(k: *const Kernels, e: mtl.ComputeEncoder, a: Ref, ba: Ref, c: Ref, bc: Ref, y: Ref, rows: usize) void { + bind(e, k.get("id_cond"), &.{ a, ba, c, bc, y }, @as(u32, cfg.hidden)); + e.dispatchThreads(mtl.Size.of(cfg.hidden, rows, 1), mtl.Size.of(256, 1, 1)); + } + + /// y [rows, width] += b by column, fp32. + pub fn bias32(k: *const Kernels, e: mtl.ComputeEncoder, y: Ref, b: Ref, rows: usize, width: usize) void { + bind(e, k.get("id_bias32"), &.{ y, b }, @as(u32, @intCast(width))); + e.dispatchThreads(mtl.Size.of(width, rows, 1), mtl.Size.of(@min(width, 256), 1, 1)); + } +}; + +/// One kernel's buffers in order, then its parameters at the next index. +pub fn bind(e: mtl.ComputeEncoder, p: mtl.Pipeline, refs: []const Ref, params: anytype) void { + e.setPipeline(p); + for (refs, 0..) |r, i| e.setBuffer(r.buf, r.off, i); + e.setValue(params, refs.len); +} + +test "every Idol pipeline compiles and loads on this GPU" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + if (!device.tensorUnits()) return error.SkipZigTest; + var k = try Kernels.load(std.testing.allocator, device); + k.deinit(); +} diff --git a/zig/src/families/idol/model.zig b/zig/src/families/idol/model.zig new file mode 100644 index 0000000000..27c6e12a13 --- /dev/null +++ b/zig/src/families/idol/model.zig @@ -0,0 +1,244 @@ +//! Idol's denoiser on Metal: the image condition refined once, then latent, condition and camera rows each step. +const std = @import("std"); +const mtl = @import("metal"); +const cfg = @import("config.zig"); +const host = @import("host.zig"); +const wts = @import("weights.zig"); +const ks = @import("kernels.zig"); +const Ref = ks.Ref; + +const H = cfg.hidden; +const F = cfg.ffn; +const L = cfg.joint_rows; +const R = cfg.latent_rows; +const C = cfg.cond_rows; + +pub const Timing = struct { wall: f64, gpu: f64, host: f64 }; + +/// A refined image condition: the condition refiner's rows [4101, 1024] fp16. +pub const Condition = struct { + rows: mtl.Buffer, + + pub fn deinit(c: *Condition) void { + c.rows.deinit(); + } +}; + +/// The scratch buffer's parts in bytes, in the order of Model's views. +const parts = [_]usize{ + L * H * 2, L * H * 2, L * 3 * H * 2, L * H * 2, L * H * 2, L * H * 2, L * H * 2, L * H * 2, + L * F * 2, L * H * 4, L * 2 * cfg.repo_hidden * 4, L * cfg.heads * cfg.repo_pairs * 4, L * cfg.heads * cfg.repo_pairs * 4, R * cfg.latent_channels * 4, R * cfg.latent_channels * 4, cfg.modulated * host.vectors * H * 2, + 2 * H * 4, C * cfg.cond_channels * 4, C * cfg.cond2_channels * 4, H * 2, R * H * 2, +}; + +pub fn scratchBytes() usize { + var total: usize = 0; + for (parts) |n| total += std.mem.alignForward(usize, n, 256); + return total; +} + +pub const Model = struct { + gpa: std.mem.Allocator, + device: mtl.Device, + queue: mtl.Queue, + k: ks.Kernels, + w: wts.Weights, + scratch: mtl.Buffer, + // fp16 rows: the stream, a norm's output, packed qkv, heads [16, rows, 64], attention rows, a projection's output + x: Ref, + xn: Ref, + qkv: Ref, + q: Ref, + k_: Ref, + v: Ref, + att: Ref, + y: Ref, + /// the MLP's GELU rows + u: Ref, + /// fp32 rows: RePo's norm (and the condition's two projections), its gate | content sums, cos and sin + ln: Ref, + gc: Ref, + cos: Ref, + sin: Ref, + lat: Ref, + vel: Ref, + mods: Ref, + head: Ref, + f1: Ref, + f2: Ref, + ones: Ref, + pe: Ref, + /// The model time `mods` and `head` hold, if any. + time: ?f32 = null, + + /// The denoiser of snapshot folder `dir`. + pub fn load(gpa: std.mem.Allocator, io: std.Io, device: mtl.Device, dir: []const u8) !*Model { + // qualified on the M5's tensor units alone: the projections and attention take another path before them + if (!device.tensorUnits()) { + std.log.err("idol: {s} has no tensor units; this family runs on the M5's alone, the only chip it was checked on", .{device.name()}); + return error.TensorUnitsRequired; + } + const t = try gpa.create(Model); + errdefer gpa.destroy(t); + t.* = undefined; + t.gpa = gpa; + t.device = device; + t.time = null; + t.queue = try device.queue(); + errdefer t.queue.deinit(); + t.k = try ks.Kernels.load(gpa, device); + errdefer t.k.deinit(); + t.w = try wts.load(gpa, io, device, dir); + errdefer t.w.deinit(); + t.scratch = try device.buffer(scratchBytes(), mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + var at: usize = 0; + const views = [_]*Ref{ &t.x, &t.xn, &t.qkv, &t.q, &t.k_, &t.v, &t.att, &t.y, &t.u, &t.ln, &t.gc, &t.cos, &t.sin, &t.lat, &t.vel, &t.mods, &t.head, &t.f1, &t.f2, &t.ones, &t.pe }; + for (views, parts) |view, n| { + view.* = .{ .buf = t.scratch, .off = at }; + at += std.mem.alignForward(usize, n, 256); + } + @memset(halves(t.ones, H), 1.0); + host.positions(t.w.host.pos, halves(t.pe, R * H)); + return t; + } + + pub fn deinit(t: *Model) void { + t.scratch.deinit(); + t.w.deinit(); + t.k.deinit(); + t.queue.deinit(); + t.gpa.destroy(t); + } + + fn floats(r: Ref, n: usize) []f32 { + return @as([*]f32, @ptrCast(@alignCast(r.addr())))[0..n]; + } + + fn halves(r: Ref, n: usize) []f16 { + return @as([*]f16, @ptrCast(@alignCast(r.addr())))[0..n]; + } + + /// The modulated block `i`'s vector `part` (0 shift, 1 one plus scale, 2 gate; 3 to 5 the MLP's). + fn vec(t: *const Model, i: usize, part: usize) Ref { + return t.mods.at(((i * host.vectors) + part) * H * 2); + } + + /// The image condition: DINOv3 rows [4101, 1280] and FLUX.2 VAE rows [4101, 128] (zeros: the negative branch). + pub fn condition(t: *Model, feature1: []const f32, feature2: []const f32) !Condition { + if (feature1.len != C * cfg.cond_channels or feature2.len != C * cfg.cond2_channels) return error.ConditionRows; + var c: Condition = .{ .rows = try t.device.buffer(C * H * 2, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked) }; + errdefer c.rows.deinit(); + @memcpy(floats(t.f1, feature1.len), feature1); + @memcpy(floats(t.f2, feature2.len), feature2); + const started = mtl.clock.seconds(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const cb = t.queue.commandBuffer(); + const e = cb.compute(.serial); + const a = t.ln; + const b = t.ln.at(C * H * 4); + t.k.linear32(e, t.f1, t.w.cond, a, C, H, cfg.cond_channels, H); + t.k.linear32(e, t.f2, t.w.cond2, b, C, H, cfg.cond2_channels, H); + t.k.cond(e, a, t.w.cond_b, b, t.w.cond2_b, t.x, C); + for (t.w.context, t.w.context_repo) |blk, r| t.block(e, blk, r, C, null); + _ = try finish(cb, e, started, 0); + @memcpy(halves(.{ .buf = c.rows }, C * H), halves(t.x, C * H)); + return c; + } + + /// The velocity of latent [8192, 16] and camera [5] at model time `time` (1000 times t) under condition `c`. + pub fn forward(t: *Model, c: *const Condition, latent: []const f32, camera: *const [cfg.camera_channels]f32, time: f32, velocity: []f32, camera_velocity: *[cfg.camera_channels]f32) !Timing { + if (latent.len != R * cfg.latent_channels or velocity.len != latent.len) return error.LatentRows; + const started = mtl.clock.seconds(); + if (t.time != time) { + t.time = null; + host.modulation(t.w.host, time, @ptrCast(@alignCast(t.mods.addr())), @ptrCast(@alignCast(t.head.addr()))); + t.time = time; + } + @memcpy(floats(t.lat, latent.len), latent); + @memcpy(halves(t.x.at(R * H * 2), C * H), halves(.{ .buf = c.rows }, C * H)); + host.cameraToken(t.w.host, camera, halves(t.x.at((L - 1) * H * 2), H)); + const spent = mtl.clock.seconds() - started; + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const cb = t.queue.commandBuffer(); + const e = cb.compute(.serial); + t.k.embed(e, t.lat, t.w.input, t.w.input_b, t.pe, t.x); + for (t.w.noise, t.w.noise_repo, 0..) |blk, r, i| t.block(e, blk, r, R, i); + for (t.w.blocks, t.w.repo, 0..) |blk, r, i| t.block(e, blk, r, L, cfg.refiners + i); + const head: Ref = t.head; + t.k.norm32(e, t.x, head.at(H * 4), head, t.ln, R, cfg.default_eps, .modulated); + t.k.linear32(e, t.ln, t.w.out, t.vel, R, cfg.latent_channels, H, cfg.latent_channels); + t.k.bias32(e, t.vel, t.w.out_b, R, cfg.latent_channels); + const timing = try finish(cb, e, started, spent); + @memcpy(velocity, floats(t.vel, velocity.len)); + host.cameraVelocity(t.w.host, halves(t.x.at((L - 1) * H * 2), H), @ptrCast(@alignCast(t.head.addr())), camera_velocity); + return timing; + } + + /// Modulated block `index` alone at model time `time` on `rows` (the noise refiner's 8192, else 12294) given rows. + pub fn blockAt(t: *Model, index: usize, rows: []const f16, time: f32) !void { + const n: usize = if (index < cfg.refiners) R else L; + if (index >= cfg.modulated or rows.len != n * H) return error.BlockRows; + if (t.time != time) { + t.time = null; + host.modulation(t.w.host, time, @ptrCast(@alignCast(t.mods.addr())), @ptrCast(@alignCast(t.head.addr()))); + t.time = time; + } + @memcpy(halves(t.x, rows.len), rows); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const cb = t.queue.commandBuffer(); + const e = cb.compute(.serial); + if (index < cfg.refiners) t.block(e, t.w.noise[index], t.w.noise_repo[index], n, index) else t.block(e, t.w.blocks[index - cfg.refiners], t.w.repo[index - cfg.refiners], n, index); + _ = try finish(cb, e, mtl.clock.seconds(), 0); + } + + /// The stream's first `rows` rows as the last pass left them, fp16. + pub fn stream(t: *const Model, rows: usize) []const f16 { + return halves(t.x, rows * H); + } + + fn finish(cb: mtl.CommandBuffer, e: mtl.ComputeEncoder, started: f64, spent: f64) !Timing { + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.log.err("idol: a pass failed on the GPU: {s}", .{msg}); + return error.GpuFailure; + } + return .{ .wall = mtl.clock.seconds() - started, .gpu = cb.gpuSeconds(), .host = spent }; + } + + /// One block over the stream's first `rows` rows: modulated block `mod`, or the condition refiner's (null). + fn block(t: *const Model, e: mtl.ComputeEncoder, b: wts.Block, r: wts.Repo, rows: usize, mod: ?usize) void { + const k = &t.k; + // RePo: angles from the block's input rows + k.normSplit(e, t.x, r.norm, r.norm_b, t.qkv, rows, cfg.default_eps); + k.split(e, t.qkv, r.maps, t.gc, rows); + k.repo(e, t.gc, r.inv, r.final, r.freqs, t.cos, t.sin, rows); + if (mod) |m| { + k.norm(e, t.x, t.ln, t.ln, t.vec(m, 1), t.vec(m, 0), t.xn, rows, cfg.block_eps, .modulated); + } else { + k.norm(e, t.x, b.norm1, b.norm1_b, t.ones, t.ones, t.xn, rows, cfg.block_eps, .affine); + } + k.linear(e, t.xn, b.qkv, b.qkv_b, t.qkv, rows, 3 * H, H, false); + k.heads(e, t.qkv, b.q_norm, b.k_norm, t.cos, t.sin, t.q, t.k_, t.v, rows); + k.attention(e, t.q, t.k_, t.v, t.att, rows, rows); + k.linear(e, t.att, b.out, b.out_b, t.y, rows, H, H, false); + k.gateAdd(e, t.x, t.y, if (mod) |m| t.vec(m, 2) else t.ones, rows); + if (mod) |m| { + k.norm(e, t.x, t.ln, t.ln, t.vec(m, 4), t.vec(m, 3), t.xn, rows, cfg.block_eps, .modulated); + } else { + k.norm(e, t.x, b.norm2, b.norm2_b, t.ones, t.ones, t.xn, rows, cfg.block_eps, .affine); + } + k.linear(e, t.xn, b.up, b.up_b, t.u, rows, F, H, true); + k.linear(e, t.u, b.down, b.down_b, t.y, rows, H, F, false); + k.gateAdd(e, t.x, t.y, if (mod) |m| t.vec(m, 5) else t.ones, rows); + } +}; + +test "scratch for the joint rows fits a GPU buffer, each part 256-aligned" { + try std.testing.expect(scratchBytes() < 1 << 31); + try std.testing.expectEqual(@as(usize, 12294), L); +} diff --git a/zig/src/families/idol/reference.zig b/zig/src/families/idol/reference.zig new file mode 100644 index 0000000000..549b92fa99 --- /dev/null +++ b/zig/src/families/idol/reference.zig @@ -0,0 +1,79 @@ +//! Comparing a velocity with idol-sdk's: cosine, largest difference and worst row, and the gate against fp32. +const std = @import("std"); + +pub const Match = struct { max_abs: f64, cosine: f64, worst_row: f64 }; + +/// Largest absolute difference, cosine, and the largest relative L2 error of one row of `width`. +pub fn match(ours: []const f32, theirs: []const f32, width: usize) Match { + std.debug.assert(ours.len == theirs.len and ours.len % width == 0); + var dot: f64 = 0; + var aa: f64 = 0; + var bb: f64 = 0; + var most: f64 = 0; + var worst: f64 = 0; + var row: usize = 0; + while (row < ours.len) : (row += width) { + var rd: f64 = 0; + var rb: f64 = 0; + for (ours[row..][0..width], theirs[row..][0..width]) |a, b| { + const d = @as(f64, a) - b; + dot += @as(f64, a) * b; + aa += @as(f64, a) * a; + bb += @as(f64, b) * b; + rd += d * d; + rb += @as(f64, b) * b; + most = @max(most, @abs(d)); + } + worst = @max(worst, @sqrt(rd / @max(rb, 1e-30))); + } + return .{ .max_abs = most, .cosine = dot / @sqrt(@max(aa * bb, 1e-300)), .worst_row = worst }; +} + +/// How far ours may sit from the fp32 ruler: this many times the torch fp16 reference's own distance. +pub const factor = 2.0; + +/// Against the fp32 ruler: worst row and 1 - cosine within `factor` times the fp16 reference's own; non-finite fails. +pub fn within(ours: Match, reference: Match) bool { + if (!std.math.isFinite(ours.cosine) or !std.math.isFinite(ours.worst_row)) return false; + return ours.worst_row <= factor * reference.worst_row and (1 - ours.cosine) <= factor * (1 - reference.cosine); +} + +/// The cosine half of `within`: for a velocity, whose worst row the gate takes over the whole run instead. +pub fn withinCosine(ours: Match, reference: Match) bool { + return std.math.isFinite(ours.cosine) and (1 - ours.cosine) <= factor * (1 - reference.cosine); +} + +test "a velocity matches itself exactly" { + const v = [_]f32{ 1, -2, 3, 0.5, 0.25, -1 }; + const m = match(&v, &v, 3); + try std.testing.expectEqual(@as(f64, 0), m.max_abs); + try std.testing.expectEqual(@as(f64, 0), m.worst_row); + try std.testing.expectApproxEqAbs(@as(f64, 1), m.cosine, 1e-15); +} + +test "a planted 5% error in one row fails the gate though everything else repeats exactly" { + const a = std.testing.allocator; + const n = 64 * 16; + const ruler = try a.alloc(f32, n); + defer a.free(ruler); + const fp16 = try a.alloc(f32, n); + defer a.free(fp16); + const ours = try a.alloc(f32, n); + defer a.free(ours); + var rng = std.Random.DefaultPrng.init(7); + for (ruler, fp16, ours) |*r, *f, *o| { + r.* = rng.random().floatNorm(f32); + f.* = r.* * (1 + 1e-3 * rng.random().floatNorm(f32)); + o.* = r.* * (1 + 1e-3 * rng.random().floatNorm(f32)); + } + const reference = match(fp16, ruler, 16); + try std.testing.expect(within(match(ours, ruler, 16), reference)); + for (ours[5 * 16 ..][0..16]) |*o| o.* *= 1.05; + try std.testing.expect(!within(match(ours, ruler, 16), reference)); +} + +test "non-finite values fail the gate" { + const ruler = [_]f32{ 1, 2, 3, 4 }; + const bad = [_]f32{ 1, std.math.nan(f32), 3, 4 }; + try std.testing.expect(!within(match(&bad, &ruler, 2), match(&ruler, &ruler, 2))); +} diff --git a/zig/src/families/idol/weights.zig b/zig/src/families/idol/weights.zig new file mode 100644 index 0000000000..27d872168c --- /dev/null +++ b/zig/src/families/idol/weights.zig @@ -0,0 +1,310 @@ +//! Idol's denoiser tensors (fp32 as stored): the blocks' linears to fp16 at load as idol-sdk does, the rest fp32. +const std = @import("std"); +const mtl = @import("metal"); +const ck = @import("../../core/checkpoint_metal.zig"); +const cfg = @import("config.zig"); +const Ref = @import("../../core/qwen_vl/weights.zig").Ref; + +const H = cfg.hidden; + +/// A transformer block: fp16 linears and biases, fp32 head RMSNorm gammas [16, 64]; norms only in the refiner. +pub const Block = struct { qkv: Ref, qkv_b: Ref, out: Ref, out_b: Ref, up: Ref, up_b: Ref, down: Ref, down_b: Ref, q_norm: Ref, k_norm: Ref, norm1: Ref = undefined, norm1_b: Ref = undefined, norm2: Ref = undefined, norm2_b: Ref = undefined }; + +/// A RePo layer: fp32 norm; gate | content maps over `inv` (a power of two) as fp16 [hi | lo | hi]; final map, freqs. +pub const Repo = struct { norm: Ref, norm_b: Ref, maps: Ref, inv: f32, final: Ref, freqs: Ref }; + +const repos = 2 * cfg.refiners + cfg.layers; +const map_bytes = 2 * cfg.repo_hidden * 3 * H * 2; + +/// The power of two that puts a map's largest magnitude in [2^13, 2^14), so its fp16 halves keep their bits. +fn mapScale(gate: []const f32, content: []const f32) f32 { + var most: f32 = 0; + for (gate) |v| most = @max(most, @abs(v)); + for (content) |v| most = @max(most, @abs(v)); + if (most == 0 or !std.math.isFinite(most)) return 1; + return std.math.ldexp(@as(f32, 1), 13 - std.math.ilogb(most)); +} + +/// Row `n` of a map as [hi | lo | hi] fp16 of `scale` times its values. +fn split(src: []const f32, scale: f32, dst: []f16) void { + for (src, 0..) |v, i| { + const x = v * scale; + const hi: f16 = @floatCast(x); + dst[i] = hi; + dst[H + i] = @floatCast(x - @as(f32, hi)); + dst[2 * H + i] = hi; + } +} + +/// The fp32 vectors and small matrices the host reads each timestep. +pub const Host = struct { + t1: []const f32, + t1_b: []const f32, + t2: []const f32, + t2_b: []const f32, + ada: []const f32, + ada_b: []const f32, + cam1: []const f32, + cam1_b: []const f32, + cam2: []const f32, + cam2_b: []const f32, + cam_out: []const f32, + cam_out_b: []const f32, + /// Each modulated block's own shift table [6144]: the noise refiner's, then the joint blocks'. + shifts: [cfg.modulated][]const f32, + /// The output's shift table [2, 1024]. + head: []const f32, + pos: []const f32, +}; + +pub const Weights = struct { + ck: ck.Checkpoint, + noise: [cfg.refiners]Block, + context: [cfg.refiners]Block, + blocks: [cfg.layers]Block, + noise_repo: [cfg.refiners]Repo, + context_repo: [cfg.refiners]Repo, + repo: [cfg.layers]Repo, + input: Ref, + input_b: Ref, + cond: Ref, + cond_b: Ref, + cond2: Ref, + cond2_b: Ref, + out: Ref, + out_b: Ref, + host: Host, + + pub fn deinit(w: *Weights) void { + w.ck.deinit(); + } +}; + +const linears = [_][]const u8{ "attn.qkv", "attn.out", "mlp.mlp.0", "mlp.mlp.2" }; + +/// The blocks' linear weights stay in the file: load reads each once into its fp16 copy. +fn later(name: []const u8, rank: usize) bool { + if (rank != 2) return false; + const owned = std.mem.startsWith(u8, name, "blocks.") or std.mem.startsWith(u8, name, "noise_refiner.") or std.mem.startsWith(u8, name, "context_refiner."); + return owned and std.mem.endsWith(u8, name, ".weight"); +} + +fn check(t: ck.Tensor, name: []const u8, shape: []const usize) !void { + if (t.dtype != .f32 or t.rank != shape.len or !std.mem.eql(usize, t.shape[0..t.rank], shape)) { + std.log.err("idol: {s} is {t} {any}; the denoiser reads f32 {any}", .{ name, t.dtype, t.shape[0..t.rank], shape }); + return error.UnexpectedTensor; + } +} + +fn take(c: *const ck.Checkpoint, name: []const u8, shape: []const usize) !Ref { + const t = try c.get(name); + try check(t, name, shape); + if (t.offset % 16 != 0) return error.UnalignedTensor; + return .{ .buf = t.buffer, .off = t.offset }; +} + +fn floats(c: *const ck.Checkpoint, name: []const u8, shape: []const usize) ![]const f32 { + const t = try c.get(name); + try check(t, name, shape); + return t.host(f32); +} + +/// fp16 values written into an owned buffer, bump-allocated 256-aligned. +const Halves = struct { + buf: mtl.Buffer, + at: usize = 0, + stage: []f32, + + fn put(h: *Halves, src: []const f32) Ref { + const dst = @as([*]f16, @ptrCast(@alignCast(h.buf.contents() + h.at)))[0..src.len]; + for (dst, src) |*o, v| o.* = @floatCast(v); + const r: Ref = .{ .buf = h.buf, .off = h.at }; + h.at += std.mem.alignForward(usize, src.len * 2, ck.alignment); + return r; + } + + /// A deferred fp32 weight read in staged pieces, rounded to fp16. + fn read(h: *Halves, c: *const ck.Checkpoint, name: []const u8, shape: [2]usize) !Ref { + const d = c.deferred.get(name) orelse return error.MissingTensor; + if (d.dtype != .f32 or d.rank != 2 or d.shape[0] != shape[0] or d.shape[1] != shape[1]) { + std.log.err("idol: {s} is {t} {any}; the denoiser reads f32 {any}", .{ name, d.dtype, d.shape[0..d.rank], shape }); + return error.UnexpectedTensor; + } + const n = shape[0] * shape[1]; + const r: Ref = .{ .buf = h.buf, .off = h.at }; + const dst = @as([*]f16, @ptrCast(@alignCast(h.buf.contents() + h.at)))[0..n]; + var done: usize = 0; + while (done < n) { + const part = @min(h.stage.len, n - done); + try c.read(name, done * 4, std.mem.sliceAsBytes(h.stage[0..part])); + for (dst[done..][0..part], h.stage[0..part]) |*o, v| o.* = @floatCast(v); + done += part; + } + h.at += std.mem.alignForward(usize, n * 2, ck.alignment); + return r; + } +}; + +const per_block = blk: { + var n: usize = 0; + for ([_]usize{ 3 * H * H, 3 * H, H * H, H, cfg.ffn * H, cfg.ffn, H * cfg.ffn, H }) |v| n += std.mem.alignForward(usize, v * 2, ck.alignment); + break :blk n; +}; + +fn block(c: *const ck.Checkpoint, h: *Halves, prefix: []const u8, i: usize, affine: bool) !Block { + var name: [96]u8 = undefined; + const P = struct { + fn n(buf: []u8, p: []const u8, j: usize, s: []const u8) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}.{d}.{s}", .{ p, j, s }); + } + }; + var b: Block = undefined; + b.qkv = try h.read(c, try P.n(&name, prefix, i, "attn.qkv.weight"), .{ 3 * H, H }); + b.qkv_b = h.put(try floats(c, try P.n(&name, prefix, i, "attn.qkv.bias"), &.{3 * H})); + b.out = try h.read(c, try P.n(&name, prefix, i, "attn.out.weight"), .{ H, H }); + b.out_b = h.put(try floats(c, try P.n(&name, prefix, i, "attn.out.bias"), &.{H})); + b.up = try h.read(c, try P.n(&name, prefix, i, "mlp.mlp.0.weight"), .{ cfg.ffn, H }); + b.up_b = h.put(try floats(c, try P.n(&name, prefix, i, "mlp.mlp.0.bias"), &.{cfg.ffn})); + b.down = try h.read(c, try P.n(&name, prefix, i, "mlp.mlp.2.weight"), .{ H, cfg.ffn }); + b.down_b = h.put(try floats(c, try P.n(&name, prefix, i, "mlp.mlp.2.bias"), &.{H})); + b.q_norm = try take(c, try P.n(&name, prefix, i, "attn.q_norm.gamma"), &.{ cfg.heads, cfg.head_dim }); + b.k_norm = try take(c, try P.n(&name, prefix, i, "attn.k_norm.gamma"), &.{ cfg.heads, cfg.head_dim }); + if (affine) { + b.norm1 = try take(c, try P.n(&name, prefix, i, "norm1.weight"), &.{H}); + b.norm1_b = try take(c, try P.n(&name, prefix, i, "norm1.bias"), &.{H}); + b.norm2 = try take(c, try P.n(&name, prefix, i, "norm2.weight"), &.{H}); + b.norm2_b = try take(c, try P.n(&name, prefix, i, "norm2.bias"), &.{H}); + } + return b; +} + +/// A RePo layer; its two maps into `pack` at `slot`, and its three frequency vectors as 32 values after every map. +fn repo(c: *const ck.Checkpoint, pack: mtl.Buffer, slot: usize, prefix: []const u8, i: usize) !Repo { + var name: [96]u8 = undefined; + var r: Repo = undefined; + r.norm = try take(c, try std.fmt.bufPrint(&name, "{s}.{d}.norm.weight", .{ prefix, i }), &.{H}); + r.norm_b = try take(c, try std.fmt.bufPrint(&name, "{s}.{d}.norm.bias", .{ prefix, i }), &.{H}); + const gate = try floats(c, try std.fmt.bufPrint(&name, "{s}.{d}.gate_map.weight", .{ prefix, i }), &.{ cfg.repo_hidden, H }); + const content = try floats(c, try std.fmt.bufPrint(&name, "{s}.{d}.content_map.weight", .{ prefix, i }), &.{ cfg.repo_hidden, H }); + const scale = mapScale(gate, content); + const maps = @as([*]f16, @ptrCast(@alignCast(pack.contents() + slot * map_bytes)))[0 .. 2 * cfg.repo_hidden * 3 * H]; + for (0..2 * cfg.repo_hidden) |n| { + const src = if (n < cfg.repo_hidden) gate[n * H ..][0..H] else content[(n - cfg.repo_hidden) * H ..][0..H]; + split(src, scale, maps[n * 3 * H ..][0 .. 3 * H]); + } + r.maps = .{ .buf = pack, .off = slot * map_bytes }; + r.inv = 1 / scale; + r.final = try take(c, try std.fmt.bufPrint(&name, "{s}.{d}.final_map.weight", .{ prefix, i }), &.{ 3 * cfg.heads, cfg.repo_hidden }); + const freqs = repos * map_bytes; + const dst = @as([*]f32, @ptrCast(@alignCast(pack.contents() + freqs)))[slot * cfg.repo_pairs ..][0..cfg.repo_pairs]; + var at: usize = 0; + inline for (.{ .{ "freqs_0", 10 }, .{ "freqs_1", 10 }, .{ "freqs_2", 12 } }) |f| { + const src = try floats(c, try std.fmt.bufPrint(&name, "{s}.{d}.{s}", .{ prefix, i, f[0] }), &.{f[1]}); + @memcpy(dst[at..][0..f[1]], src); + at += f[1]; + } + r.freqs = .{ .buf = pack, .off = freqs + slot * cfg.repo_pairs * 4 }; + return r; +} + +/// The snapshot folder `dir`: config.json checked, then denoiser.safetensors's every tensor, each read once. +pub fn load(gpa: std.mem.Allocator, io: std.Io, device: mtl.Device, dir: []const u8) !Weights { + var arena_state = std.heap.ArenaAllocator.init(gpa); + defer arena_state.deinit(); + const arena = arena_state.allocator(); + const config = try std.Io.Dir.cwd().readFileAlloc(io, try std.fs.path.join(arena, &.{ dir, "config.json" }), arena, .limited(1 << 20)); + try cfg.check(gpa, config); + var w: Weights = undefined; + w.ck = ck.Checkpoint.init(gpa); + errdefer w.ck.deinit(); + const path = try std.fmt.allocPrintSentinel(arena, "{s}/denoiser.safetensors", .{dir}, 0); + try w.ck.addFileSplit(device, path, "", null, later); + const opts = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; + const blocks = 2 * cfg.refiners + cfg.layers; + const fp16 = try device.buffer(blocks * per_block, opts); + w.ck.own(fp16, blocks * per_block) catch |err| { + fp16.deinit(); + return err; + }; + const pack_bytes = repos * (map_bytes + cfg.repo_pairs * 4); + const pack = try device.buffer(pack_bytes, opts); + w.ck.own(pack, pack_bytes) catch |err| { + pack.deinit(); + return err; + }; + var h: Halves = .{ .buf = fp16, .stage = try arena.alloc(f32, 4 << 20) }; + for (&w.noise, &w.context, &w.noise_repo, &w.context_repo, 0..) |*n, *x, *nr, *xr, i| { + n.* = try block(&w.ck, &h, "noise_refiner", i, false); + x.* = try block(&w.ck, &h, "context_refiner", i, true); + nr.* = try repo(&w.ck, pack, i, "noise_repo_layers", i); + xr.* = try repo(&w.ck, pack, cfg.refiners + i, "context_repo_layers", i); + } + for (&w.blocks, &w.repo, 0..) |*b, *r, i| { + b.* = try block(&w.ck, &h, "blocks", i, false); + r.* = try repo(&w.ck, pack, 2 * cfg.refiners + i, "repo_layers", i); + } + w.ck.closeFiles(); + w.input = try take(&w.ck, "input_layer.weight", &.{ H, cfg.latent_channels }); + w.input_b = try take(&w.ck, "input_layer.bias", &.{H}); + w.cond = try take(&w.ck, "cond_embedder.weight", &.{ H, cfg.cond_channels }); + w.cond_b = try take(&w.ck, "cond_embedder.bias", &.{H}); + w.cond2 = try take(&w.ck, "cond_embedder2.weight", &.{ H, cfg.cond2_channels }); + w.cond2_b = try take(&w.ck, "cond_embedder2.bias", &.{H}); + w.out = try take(&w.ck, "out_layer.weight", &.{ cfg.latent_channels, H }); + w.out_b = try take(&w.ck, "out_layer.bias", &.{cfg.latent_channels}); + var hw: Host = undefined; + hw.t1 = try floats(&w.ck, "t_embedder.mlp.0.weight", &.{ H, cfg.time_dim }); + hw.t1_b = try floats(&w.ck, "t_embedder.mlp.0.bias", &.{H}); + hw.t2 = try floats(&w.ck, "t_embedder.mlp.2.weight", &.{ H, H }); + hw.t2_b = try floats(&w.ck, "t_embedder.mlp.2.bias", &.{H}); + hw.ada = try floats(&w.ck, "adaLN_modulation.1.weight", &.{ 6 * H, H }); + hw.ada_b = try floats(&w.ck, "adaLN_modulation.1.bias", &.{6 * H}); + hw.cam1 = try floats(&w.ck, "cam_refiner.mlp.0.weight", &.{ H, cfg.camera_channels }); + hw.cam1_b = try floats(&w.ck, "cam_refiner.mlp.0.bias", &.{H}); + hw.cam2 = try floats(&w.ck, "cam_refiner.mlp.2.weight", &.{ H, H }); + hw.cam2_b = try floats(&w.ck, "cam_refiner.mlp.2.bias", &.{H}); + hw.cam_out = try floats(&w.ck, "cam_out_layer.weight", &.{ cfg.camera_channels, H }); + hw.cam_out_b = try floats(&w.ck, "cam_out_layer.bias", &.{cfg.camera_channels}); + var name: [64]u8 = undefined; + for (&hw.shifts, 0..) |*s, i| { + const n = if (i < cfg.refiners) try std.fmt.bufPrint(&name, "noise_refiner.{d}.shift_table", .{i}) else try std.fmt.bufPrint(&name, "blocks.{d}.shift_table", .{i - cfg.refiners}); + s.* = try floats(&w.ck, n, &.{ 1, 6 * H }); + } + hw.head = try floats(&w.ck, "shift_table", &.{ 1, 2, H }); + hw.pos = try floats(&w.ck, "pos_pe", &.{ 1, cfg.latent_rows, 3 }); + w.host = hw; + // a block's 4 linears, 4 biases and 2 gammas, the context refiner's 4 norms, the shift tables; 8 a RePo layer + const all = 2 * cfg.refiners + cfg.layers; + const read = 10 * all + 4 * cfg.refiners + cfg.modulated + 8 * all + 8 + 12 + 2; + const total = w.ck.tensors.count() + w.ck.deferred.count(); + if (total != read) { + std.log.err("idol: denoiser.safetensors has {d} tensors; this family reads {d}", .{ total, read }); + return error.UnexpectedTensors; + } + return w; +} + +test "a map's split halves sum back to its scaled values within fp32's rounding" { + var src: [H]f32 = undefined; + var rng = std.Random.DefaultPrng.init(3); + for (&src) |*v| v.* = 0.05 * rng.random().floatNorm(f32); + const scale = mapScale(&src, &src); + try std.testing.expect(scale >= 1); + var dst: [3 * H]f16 = undefined; + split(&src, scale, &dst); + for (src, 0..) |v, i| { + const back = (@as(f32, dst[i]) + @as(f32, dst[H + i])) / scale; + try std.testing.expect(@abs(back - v) <= @abs(v) * 1e-6 + 1e-12); + try std.testing.expectEqual(dst[i], dst[2 * H + i]); + } +} + +test "the blocks' linear weights, and only they, stay in the file for their fp16 copies" { + try std.testing.expect(later("blocks.3.attn.qkv.weight", 2)); + try std.testing.expect(later("context_refiner.1.mlp.mlp.2.weight", 2)); + try std.testing.expect(!later("blocks.3.attn.qkv.bias", 1)); + try std.testing.expect(!later("blocks.3.attn.q_norm.gamma", 2)); + try std.testing.expect(!later("context_refiner.0.norm1.weight", 1)); + try std.testing.expect(!later("repo_layers.0.gate_map.weight", 2)); + try std.testing.expect(!later("cond_embedder.weight", 2)); +} diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index 23655fbcbe..89eb6e678c 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -13,6 +13,7 @@ pub const qwen_vl_tower = @import("core/qwen_vl/metal.zig"); pub const qwen_image = @import("families/qwen_image/qwen_image.zig"); /// MOSS-Image's transformer on Metal: a denoiser driven by tf-moss-image, outside the lane engine. pub const moss_image = @import("families/moss_image/moss_image.zig"); +pub const idol = @import("families/idol/idol.zig"); pub const checkpoint = @import("core/checkpoint_metal.zig"); pub const shard_edit = @import("core/shard_edit.zig"); pub const lora_merge = @import("core/lora_merge_dir.zig"); diff --git a/zig/tests/idol.zig b/zig/tests/idol.zig new file mode 100644 index 0000000000..4359d38e45 --- /dev/null +++ b/zig/tests/idol.zig @@ -0,0 +1,242 @@ +//! tf-idol: Idol's denoiser on a case from tools/zig/idol_case.py, compared with idol-sdk's, gated and timed. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const id = tf.idol; +const cfg = id.config; +const ref = id.reference; + +const usage = + \\tf-idol SNAPSHOT_DIR CASE_DIR [blocks STEP] + \\SNAPSHOT_DIR is a speridlabs/idol snapshot (config.json, denoiser.safetensors). CASE_DIR (tools/zig/idol_case.py) + \\holds case.safetensors: both image conditions, the starting noise and the SDK's run (each step's input and both + \\branches' velocities), and optionally exact.safetensors (the same steps in float32). Each step runs both branches on + \\the SDK's input; with exact.safetensors each 1 - cosine against float32 must be within twice the SDK's own, and + \\so must the worst row over the run. Then the whole guided schedule runs from the noise, timed; CASE_DIR gets tf-idol.json and + \\sample.f32 (8192 x 16, then the 5 camera values). Exit 1 on a failed gate or a forward that does not repeat. + \\With `blocks STEP`, runs each joint block alone on the SDK's input rows at that step (blocksSTEP.safetensors) + \\and gates it against the same block in float32 (blocksSTEP_f32.safetensors), worst row and 1 - cosine. +; + +/// The card's sampling defaults. +const guidance: f32 = 3.0; +const rescale: f64 = 3.0; + +fn take(a: std.mem.Allocator, f: *const tf.safetensors.File, name: []const u8, n: usize) ![]f32 { + const t = f.get(name) orelse { + std.debug.print("the case has no {s}\n", .{name}); + return error.MissingTensor; + }; + if (t.dtype != .f32 or t.numel() != n) { + std.debug.print("{s} is not {d} float32 values\n", .{ name, n }); + return error.BadCaseTensor; + } + const out = try a.alloc(f32, n); + @memcpy(std.mem.sliceAsBytes(out), t.bytes); + return out; +} + +fn print(list: *std.ArrayList(u8), a: std.mem.Allocator, comptime fmt: []const u8, args: anytype) !void { + try list.appendSlice(a, try std.fmt.allocPrint(a, fmt, args)); +} + +fn json(list: *std.ArrayList(u8), a: std.mem.Allocator, key: []const u8, m: ref.Match) !void { + try print(list, a, ", \"{s}\": {{\"max_abs\": {d:.6}, \"cosine\": {d:.9}, \"worst_row\": {d:.6}}}", .{ key, m.max_abs, m.cosine, m.worst_row }); +} + +const Case = struct { + steps: usize, + timesteps: []f32, + inputs: []f32, + inputs_camera: []f32, + velocity: [2][]f32, + velocity_camera: [2][]f32, + exact: ?[2][]f32, +}; + +const N = cfg.latent_rows * cfg.latent_channels; +const labels = [2][]const u8{ "conditional", "negative" }; + +/// Each step's two branches on the SDK's input against its velocities; returns false when a gate fails. +fn teacher(a: std.mem.Allocator, c: Case, model: *id.Model, conds: *[2]id.Condition, out: *std.ArrayList(u8)) !bool { + var pass = true; + var worst = [2]f64{ 0, 0 }; + const v = try a.alloc(f32, N); + var vc: [cfg.camera_channels]f32 = undefined; + try print(out, a, ", \"teacher_forced\": [", .{}); + for (0..c.steps) |s| { + try print(out, a, "{s}{{\"step\": {d}, \"time\": {d:.4}", .{ if (s == 0) "" else ", ", s + 1, c.timesteps[s] }); + const cam: *const [cfg.camera_channels]f32 = c.inputs_camera[s * cfg.camera_channels ..][0..cfg.camera_channels]; + for (0..2) |b| { + _ = try model.forward(&conds[b], c.inputs[s * N ..][0..N], cam, c.timesteps[s], v, &vc); + const theirs = c.velocity[b][s * N ..][0..N]; + const m = ref.match(v, theirs, cfg.latent_channels); + const mc = ref.match(&vc, c.velocity_camera[b][s * cfg.camera_channels ..][0..cfg.camera_channels], cfg.camera_channels); + std.debug.print("step {d} {s}: against the SDK max |diff| {d:.5}, cosine {d:.8}, worst row {d:.5}; camera cosine {d:.7}\n", .{ s + 1, labels[b], m.max_abs, m.cosine, m.worst_row, mc.cosine }); + try json(out, a, try std.fmt.allocPrint(a, "{s}_sdk", .{labels[b]}), m); + try json(out, a, try std.fmt.allocPrint(a, "{s}_camera_sdk", .{labels[b]}), mc); + if (c.exact) |ex| if (ex[b].len >= (s + 1) * N) { + const want = ex[b][s * N ..][0..N]; + const ours = ref.match(v, want, cfg.latent_channels); + const sdk = ref.match(theirs, want, cfg.latent_channels); + const ok = ref.withinCosine(ours, sdk); + pass = pass and ok; + worst[0] = @max(worst[0], ours.worst_row); + worst[1] = @max(worst[1], sdk.worst_row); + std.debug.print(" against float32: worst row {d:.5}, cosine {d:.8}; the SDK's own {d:.5}, {d:.8}: {s}\n", .{ ours.worst_row, ours.cosine, sdk.worst_row, sdk.cosine, if (ok) "pass" else "FAIL" }); + try json(out, a, try std.fmt.allocPrint(a, "{s}_fp32", .{labels[b]}), ours); + try json(out, a, try std.fmt.allocPrint(a, "{s}_sdk_vs_fp32", .{labels[b]}), sdk); + }; + } + try out.append(a, '}'); + } + if (c.exact != null) { + const ok = worst[0] <= ref.factor * worst[1]; + std.debug.print("worst row over the run against float32: {d:.5}; the SDK's own {d:.5}: {s}\n", .{ worst[0], worst[1], if (ok) "pass" else "FAIL" }); + pass = pass and ok; + } + try print(out, a, "], \"worst_row_fp32\": {d:.6}, \"sdk_worst_row_fp32\": {d:.6}", .{ worst[0], worst[1] }); + return pass; +} + +/// The same step twice with the other branch between: identical bytes or not. +fn repeats(a: std.mem.Allocator, c: Case, model: *id.Model, conds: *[2]id.Condition) !bool { + const first = try a.alloc(f32, N); + const again = try a.alloc(f32, N); + var vc: [cfg.camera_channels]f32 = undefined; + const cam: *const [cfg.camera_channels]f32 = c.inputs_camera[0..cfg.camera_channels]; + _ = try model.forward(&conds[0], c.inputs[0..N], cam, c.timesteps[0], first, &vc); + _ = try model.forward(&conds[1], c.inputs[0..N], cam, c.timesteps[@min(1, c.steps - 1)], again, &vc); + _ = try model.forward(&conds[0], c.inputs[0..N], cam, c.timesteps[0], again, &vc); + return std.mem.eql(u8, std.mem.sliceAsBytes(first), std.mem.sliceAsBytes(again)); +} + +/// The case tensor `name`, fp16 or fp32 rows, as f32. +fn rows(a: std.mem.Allocator, f: *const tf.safetensors.File, name: []const u8) ![]f32 { + const t = f.get(name) orelse return error.MissingTensor; + const out = try a.alloc(f32, t.numel()); + switch (t.dtype) { + .f16 => for (out, @as([*]align(1) const f16, @ptrCast(t.bytes.ptr))[0..out.len]) |*o, v| { + o.* = v; + }, + .f32 => @memcpy(std.mem.sliceAsBytes(out), t.bytes), + else => return error.BadCaseTensor, + } + return out; +} + +/// Each joint block alone on the SDK's own input rows at step `step`, against that block in float32; false on a fail. +fn blockGates(a: std.mem.Allocator, io: std.Io, dir: []const u8, step: usize, c: Case, model: *id.Model) !bool { + var f = try tf.safetensors.File.open(a, io, try std.fmt.allocPrint(a, "{s}/blocks{d}.safetensors", .{ dir, step })); + defer f.close(io); + var g = try tf.safetensors.File.open(a, io, try std.fmt.allocPrint(a, "{s}/blocks{d}_f32.safetensors", .{ dir, step })); + defer g.close(io); + const H = cfg.hidden; + var pass = true; + var name: [32]u8 = undefined; + for (1..cfg.layers) |i| { + const input = try rows(a, &f, try std.fmt.bufPrint(&name, "block.{d}", .{i - 1})); + const in16 = try a.alloc(f16, input.len); + for (in16, input) |*o, v| o.* = @floatCast(v); + try model.blockAt(cfg.refiners + i, in16, c.timesteps[step]); + const ours = try a.alloc(f32, input.len); + for (ours, model.stream(cfg.joint_rows)) |*o, v| o.* = v; + const label = try std.fmt.bufPrint(&name, "block.{d}", .{i}); + const want = try rows(a, &g, label); + const m = ref.match(ours, want, H); + const r = ref.match(try rows(a, &f, label), want, H); + const ok = ref.within(m, r); + pass = pass and ok; + std.debug.print("step {d} block {d} against float32: worst row {d:.5}, cosine {d:.9}; the SDK's own {d:.5}, {d:.9}: {s}\n", .{ step + 1, i, m.worst_row, m.cosine, r.worst_row, r.cosine, if (ok) "pass" else "FAIL" }); + } + return pass; +} + +pub fn main(init: std.process.Init) !u8 { + const a = init.arena.allocator(); + const io = init.io; + const args = try init.minimal.args.toSlice(a); + if (args.len != 3 and args.len != 5) { + std.debug.print("{s}\n", .{usage}); + return 2; + } + var case = try tf.safetensors.File.open(a, io, try std.fmt.allocPrint(a, "{s}/case.safetensors", .{args[2]})); + defer case.close(io); + const steps = (case.get("timesteps") orelse return error.MissingTensor).numel(); + const S = steps; + const cam_n = cfg.camera_channels; + var c: Case = .{ .steps = S, .timesteps = try take(a, &case, "timesteps", S), .inputs = try take(a, &case, "inputs", S * N), .inputs_camera = try take(a, &case, "inputs_camera", S * cam_n), .velocity = .{ try take(a, &case, "velocity", S * N), try take(a, &case, "velocity_negative", S * N) }, .velocity_camera = .{ try take(a, &case, "velocity_camera", S * cam_n), try take(a, &case, "velocity_negative_camera", S * cam_n) }, .exact = null }; + const exact_path = try std.fmt.allocPrint(a, "{s}/exact.safetensors", .{args[2]}); + var exact_file: ?tf.safetensors.File = tf.safetensors.File.open(a, io, exact_path) catch null; + defer if (exact_file) |*f| f.close(io); + if (exact_file) |*f| { + const n = (f.get("velocity_f32") orelse return error.MissingTensor).numel(); + c.exact = .{ try take(a, f, "velocity_f32", n), try take(a, f, "velocity_negative_f32", n) }; + } + const f1 = try take(a, &case, "feature1", cfg.cond_rows * cfg.cond_channels); + const f2 = try take(a, &case, "feature2", cfg.cond_rows * cfg.cond2_channels); + const noise = try take(a, &case, "latent", N); + const noise_cam = try take(a, &case, "camera", cam_n); + + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + const loading = mtl.clock.seconds(); + const model = try id.Model.load(std.heap.page_allocator, io, device, args[1]); + defer model.deinit(); + const load_s = mtl.clock.seconds() - loading; + std.debug.print("{s}: {d} steps; loaded in {d:.1} s\n", .{ device.name(), S, load_s }); + + var out: std.ArrayList(u8) = .empty; + const conditioning = mtl.clock.seconds(); + var conds: [2]id.Condition = undefined; + conds[0] = try model.condition(f1, f2); + defer conds[0].deinit(); + @memset(f1, 0); + @memset(f2, 0); + conds[1] = try model.condition(f1, f2); + defer conds[1].deinit(); + const condition_s = mtl.clock.seconds() - conditioning; + if (args.len == 5) return if (try blockGates(a, io, args[2], try std.fmt.parseInt(usize, args[4], 10), c, model)) 0 else 1; + try print(&out, a, "{{\"device\": \"{s}\", \"steps\": {d}, \"load_s\": {d:.2}, \"condition_s\": {d:.4}", .{ device.name(), S, load_s, condition_s }); + const pass = try teacher(a, c, model, &conds, &out); + const repeat = try repeats(a, c, model, &conds); + std.debug.print("the same step twice, another between: {s}\n", .{if (repeat) "identical bytes" else "DIFFERENT bytes"}); + try print(&out, a, ", \"gates_pass\": {}, \"repeat_identical\": {}, \"run\": [", .{ pass and c.exact != null, repeat }); + + const ts = try a.alloc(f64, S + 1); + id.host.schedule(ts, rescale); + const x = try a.alloc(f32, N + cam_n); + @memcpy(x[0..N], noise); + @memcpy(x[N..], noise_cam); + const vc = try a.alloc(f32, N); + const vu = try a.alloc(f32, N); + var cc: [cam_n]f32 = undefined; + var cu: [cam_n]f32 = undefined; + var total: f64 = 0; + for (0..S) |s| { + const time: f32 = @floatCast(1000 * ts[s]); + const cam: *const [cam_n]f32 = x[N..][0..cam_n]; + const p = try model.forward(&conds[0], x[0..N], cam, time, vc, &cc); + const q = try model.forward(&conds[1], x[0..N], cam, time, vu, &cu); + const by: f32 = @floatCast(ts[s] - ts[s + 1]); + for (x[0..N], vc, vu) |*o, pc, pu| o.* -= ((1 + guidance) * pc - guidance * pu) * by; + for (x[N..], cc, cu) |*o, pc, pu| o.* -= ((1 + guidance) * pc - guidance * pu) * by; + const wall = p.wall + q.wall; + total += wall; + std.debug.print("step {d}/{d}: {d:.3} s ({d:.3} s on the GPU)\n", .{ s + 1, S, wall, p.gpu + q.gpu }); + try print(&out, a, "{s}{{\"wall_s\": {d:.4}, \"gpu_s\": {d:.4}}}", .{ if (s == 0) "" else ", ", wall, p.gpu + q.gpu }); + } + if (case.get("sample") != null) { + const sample = try take(a, &case, "sample", N); + const m = ref.match(x[0..N], sample, cfg.latent_channels); + std.debug.print("final sample against the SDK's: cosine {d:.6}, worst row {d:.4}\n", .{ m.cosine, m.worst_row }); + try print(&out, a, "]", .{}); + try json(&out, a, "sample_sdk", m); + try print(&out, a, ", \"denoise_s\": {d:.3}}}\n", .{total}); + } else try print(&out, a, "], \"denoise_s\": {d:.3}}}\n", .{total}); + std.debug.print("denoise: {d:.2} s for {d} steps ({d} forwards)\n", .{ total, S, 2 * S }); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = try std.fmt.allocPrint(a, "{s}/sample.f32", .{args[2]}), .data = std.mem.sliceAsBytes(x) }); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = try std.fmt.allocPrint(a, "{s}/tf-idol.json", .{args[2]}), .data = out.items }); + return if (repeat and (c.exact == null or pass)) 0 else 1; +} From 6b1b074e29a5719b2ab997ee58082e5e5948b09a Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 02:39:20 +0800 Subject: [PATCH 2/9] contract: Idol's denoiser kernels through gates 1-5 Every pipeline of idol.metal and idol_attention.metal is registered: id_linear's eight forms (plain and GELU), id_split's two, id_gemm32, id_norm (plain, affine, modulated), id_norm32 (affine, modulated), id_norm_split, id_repo, id_qk, id_gate_add, id_embed, id_cond, id_bias32 and the attention's four forms. Each output is checked at idol-sdk's fp16 points as the range its fp32 steps allow; the NAX projections against the tensor unit's 16-deep step bound, the attention against attention_ref's formula with P at the tensor unit's fp32. Inputs at the large served shapes repeat at periods prime to every tile, so float64 stays cheap and every output is checked. 19 controls: rounding points dropped or added (gate 3), fp16 sums in id_norm32, id_repo, id_qk, id_embed and id_gemm32, and the NAX fp32 path in id_gemm32's place (gate 4), and the attention's tail keys unmasked. Gate 5 holds the 24 launches, each reserved by its source. Known finding: id_qk gives a head holding a NaN +-inf in its other outputs where F.normalize gives NaN, as Metal's max drops the NaN norm that clamp_min keeps. --- docs/kernel-contract.md | 6 + zig/src/core/contract.zig | 1 + zig/src/core/contract/idol_records.zig | 34 ++ zig/src/core/contract/registry.zig | 2 +- zig/tests/contract/metal_idol.zig | 190 +++++++++ zig/tests/contract/metal_idol_attention.zig | 233 +++++++++++ zig/tests/contract/metal_idol_common.zig | 245 +++++++++++ zig/tests/contract/metal_idol_gemm.zig | 430 +++++++++++++++++++ zig/tests/contract/metal_idol_norms.zig | 251 +++++++++++ zig/tests/contract/metal_idol_rows.zig | 442 ++++++++++++++++++++ zig/tests/contract/metal_idol_steps.zig | 426 +++++++++++++++++++ zig/tests/contract/metal_limits.zig | 3 + zig/tests/contract/metal_main.zig | 2 + 13 files changed, 2264 insertions(+), 1 deletion(-) create mode 100644 zig/src/core/contract/idol_records.zig create mode 100644 zig/tests/contract/metal_idol.zig create mode 100644 zig/tests/contract/metal_idol_attention.zig create mode 100644 zig/tests/contract/metal_idol_common.zig create mode 100644 zig/tests/contract/metal_idol_gemm.zig create mode 100644 zig/tests/contract/metal_idol_norms.zig create mode 100644 zig/tests/contract/metal_idol_rows.zig create mode 100644 zig/tests/contract/metal_idol_steps.zig diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index b690ec8f9e..472881c3b5 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -117,6 +117,8 @@ MLX's formula keeps the infinity, because its codes enter the MMAs as q + 128 an `tf_glm_scores` scores a NaN dot as 0 where float64 keeps the NaN: its relu is `fmaxf(x, 0)`, as `tl.maximum` is in the Python family's Triton glue it was ported from. Nemotron's Metal `tf_xup_w3` gives 0 for a NaN dot where float64 keeps the NaN: its relu2 is `metal::max(bfloat(acc), 0.0f)`; the other up pipelines give its bytes on the same inputs. +Idol's `id_qk` gives a head holding a NaN +-inf in its other outputs where torch's `F.normalize` gives NaN: its +norm is `max(precise::sqrt(ss), 1e-12f)`, and Metal's `max` drops the NaN that `clamp_min` keeps. HIP’s MLX stream tile keeps its finite dot2 sums unchanged. A non-finite result with an actual non-finite operand is recomputed by full quantization groups using the factored float64 formula. The plain tile is @@ -185,6 +187,10 @@ CUDA also checks `dense_rows`' scaled epilogue (Dory's logits, fp32(bf16(x . w)) unscaled tile's outputs rounded and scaled on the host. HIP: the RMS norm rows, a wave a row (fp32, fp16 and bf16 out) and a block a row (fp16 and bf16 out). +Metal also checks Idol's denoiser (`metal_idol*.zig`): its fp16 projections and GELU form, the split-K fp32 rows, +`id_gemm32` (true fp32 sums: the NAX fp32 path in its place is a control gate 4 must catch), the LayerNorms, RePo's +angles, the heads, the elementwise steps and the attention's four forms, each at idol-sdk's fp16 points (19 controls). + RADIO's fused bias, activation and residual (`tf_vit_bias_act`) replaces `tf_vit_bias` then `tf_vit_pointwise`, byte for byte. CUDA also checks GLM-5.3-Flash's KV page kernels on a page table out of order: the latent write and the pooled keys diff --git a/zig/src/core/contract.zig b/zig/src/core/contract.zig index 6828c844cc..4153341694 100644 --- a/zig/src/core/contract.zig +++ b/zig/src/core/contract.zig @@ -25,6 +25,7 @@ test { _ = @import("contract/glm_rows_records.zig"); _ = @import("contract/nemotron_experts_records.zig"); _ = @import("contract/moss_image_records.zig"); + _ = @import("contract/idol_records.zig"); _ = @import("contract/metal_pair_records.zig"); _ = spec; _ = gates; diff --git a/zig/src/core/contract/idol_records.zig b/zig/src/core/contract/idol_records.zig new file mode 100644 index 0000000000..ae39549164 --- /dev/null +++ b/zig/src/core/contract/idol_records.zig @@ -0,0 +1,34 @@ +//! Idol's denoiser launches (idol.metal and idol_attention.metal after gemm_nax.metal), each source-reserved. +const std = @import("std"); +const M = @import("limits.zig").MetalLaunch; +pub const records = [_]M{ + .{ .name = "idol id_linear_00", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol id_linear_01", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol id_linear_10", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol id_linear_11", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol id_linear_gelu_00", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol id_linear_gelu_01", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol id_linear_gelu_10", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol id_linear_gelu_11", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol id_split_0", .threads = 128, .tg_memory = 0, .reserved = 128 }, + .{ .name = "idol id_split_1", .threads = 128, .tg_memory = 0, .reserved = 128 }, + .{ .name = "idol id_gemm32", .threads = 128, .tg_memory = 19456, .reserved = 128 }, + .{ .name = "idol id_norm", .threads = 256, .tg_memory = 144, .reserved = 256 }, + .{ .name = "idol id_norm32", .threads = 256, .tg_memory = 144, .reserved = 256 }, + .{ .name = "idol id_norm_split", .threads = 256, .tg_memory = 144, .reserved = 256 }, + .{ .name = "idol id_repo", .threads = 64, .tg_memory = 704, .reserved = 64 }, + .{ .name = "idol id_qk", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol id_gate_add", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol id_embed", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol id_cond", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol id_bias32", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol id_attention_00", .threads = 128, .tg_memory = 0, .reserved = 128 }, + .{ .name = "idol id_attention_01", .threads = 128, .tg_memory = 0, .reserved = 128 }, + .{ .name = "idol id_attention_10", .threads = 128, .tg_memory = 0, .reserved = 128 }, + .{ .name = "idol id_attention_11", .threads = 128, .tg_memory = 0, .reserved = 128 }, +}; + +test "the Idol launches fit every Apple family, one record a pipeline" { + try std.testing.expectEqual(@as(usize, 0), try @import("registry.zig").checkMetal(std.testing.allocator, &records, true)); + try std.testing.expectEqual(@as(usize, 24), records.len); +} diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index d48c6541bc..f55d146bf7 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -569,7 +569,7 @@ test { } /// The Metal collector keeps measured slots and the restored vision declarations. pub const all_metal = metal ++ @import("delta_slots_records.zig").records ++ @import("flashnext_pending.zig").records ++ @import("simd_pairs_records.zig").records ++ @import("metal_pair_records.zig").records ++ @import("glm_rows_records.zig").records ++ @import("weight_only_records.zig").metal ++ @import("hc_activation_records.zig").metal ++ @import("qwen4_records.zig").metal ++ @import("lfm2_records.zig").records ++ @import("nemotron_experts_records.zig").records ++ @import("nemotron_pages_records.zig").records ++ @import("nemotron_attention_records.zig").records ++ - @import("moss_image_records.zig").records ++ @import("minimax_records.zig").records; + @import("moss_image_records.zig").records ++ @import("minimax_records.zig").records ++ @import("idol_records.zig").records; const learning = @import("registry_learning.zig"); pub const gradient_pending = learning.gradient_pending; pub const qwen_learning_pending = learning.qwen_learning_pending; diff --git a/zig/tests/contract/metal_idol.zig b/zig/tests/contract/metal_idol.zig new file mode 100644 index 0000000000..50d1284ae4 --- /dev/null +++ b/zig/tests/contract/metal_idol.zig @@ -0,0 +1,190 @@ +//! Idol's Metal kernels under the contract, each at idol-sdk's fp16 points; controls drop or add one rounding point. +const std = @import("std"); +const mtl = @import("metal"); +const rig_ = @import("rig.zig"); +const id = @import("metal_idol_common.zig"); +const gemm = @import("metal_idol_gemm.zig"); +const norms = @import("metal_idol_norms.zig"); +const steps = @import("metal_idol_steps.zig"); +const elems = @import("metal_idol_rows.zig"); +const attn = @import("metal_idol_attention.zig"); + +const g = id.g; +const ks = id.ks; +const Control = id.Control; +const Lib = id.Lib; + +/// The NAX fp32 GEMM (the tensor unit's fp32 path, not fp32-accurate on M5) in id_gemm32's place. +const nax_gemm32 = + \\[[max_total_threads_per_threadgroup(128)]] [[kernel]] void id_gemm32( + \\ const device float* x [[buffer(0)]], const device float* w [[buffer(1)]], device float* y [[buffer(2)]], + \\ constant IdGemm32& p [[buffer(3)]], uint2 tg [[threadgroup_position_in_grid]], + \\ uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]]) { + \\ const int row = int(tg.y) * 64 + 32 * int(sg / 2), col = int(tg.x) * 64 + 32 * int(sg % 2); + \\ const int rows = min(32, int(p.M) - row), cols = min(32, int(p.N) - col); + \\ const short2 home = tfp::frag_home(ushort(lane)); + \\ tfp::frag acc[2][2]; + \\ tfp::gemm_block(acc, x + long(row) * p.K, w + long(col) * p.K, int(p.K), int(p.K), + \\ int(p.K), int(p.K) / 32, rows, cols, home); + \\ if (rows > 0 && cols > 0) tfp::gemm_store(acc, y + long(row) * p.ldy + col, int(p.ldy), rows, cols, home); + \\} +; + +const mma_line = "for (uint j = 0; j < 4; j++) simdgroup_multiply_accumulate(acc[i][j], a[i], b[j], acc[i][j]);"; +const mma_half = "for (uint j = 0; j < 4; j++) { simdgroup_multiply_accumulate(acc[i][j], a[i], b[j], acc[i][j]); " ++ + "acc[i][j].thread_elements()[0] = id_h(acc[i][j].thread_elements()[0]); " ++ + "acc[i][j].thread_elements()[1] = id_h(acc[i][j].thread_elements()[1]); }"; + +const controls = [_]Control{ + .{ .kernel = "id_linear_gelu_01", .what = "drops fp16(x w + b) before the GELU", .edits = &.{.{ .from = "const float v = id_h(acc[i][j][e] + b);", .to = "const float v = (acc[i][j][e] + b);" }} }, + .{ .kernel = "id_linear_01", .what = "adds fp16(x w) before the bias", .edits = &.{.{ .from = "const float v = id_h(acc[i][j][e] + b);", .to = "const float v = id_h(id_h(acc[i][j][e]) + b);" }} }, + .{ .kernel = "id_norm modulated", .what = "drops fp16(y s1)", .edits = &.{.{ .from = "v = id_h(id_h(v * float(s1[i])) + float(sh[i]));", .to = "v = id_h((v * float(s1[i])) + float(sh[i]));" }} }, + .{ .kernel = "id_norm affine", .what = "adds fp16(y) before the affine", .edits = &.{.{ .from = "if (p.mode == 1) v = v * w[i] + b[i];", .to = "if (p.mode == 1) v = id_h(v) * w[i] + b[i];" }} }, + .{ .kernel = "id_norm32 modulated", .what = "adds fp16(y) before the modulation", .edits = &.{.{ .from = "yr[i] = p.mode == 2 ? v * (1.0f + w[i])", .to = "yr[i] = p.mode == 2 ? id_h(v) * (1.0f + w[i])" }} }, + .{ .kernel = "id_norm32 modulated", .what = "sums the simdgroups in fp16", .gate = .sums, .edits = &.{.{ .from = " v = simd_sum(v);\n if (lane == 0) part[sg] = v;", .to = " v = float(simd_sum(half(v)));\n if (lane == 0) part[sg] = v;" }} }, + .{ .kernel = "id_norm_split", .what = "rounds the row to fp16 before its split", .edits = &.{.{ .from = "const float v = (float(xr[i]) - mean) * inv * w[i] + b[i];", .to = "const float v = id_h((float(xr[i]) - mean) * inv * w[i] + b[i]);" }} }, + .{ .kernel = "id_repo", .what = "adds fp16(u)", .edits = &.{.{ .from = "u[i] = (a / (1.0f + metal::precise::exp(-a))) * (g[128 + i] * p.inv);", .to = "u[i] = id_h((a / (1.0f + metal::precise::exp(-a))) * (g[128 + i] * p.inv));" }} }, + .{ .kernel = "id_repo", .what = "sums d in fp16", .gate = .sums, .edits = &.{.{ .from = "s = fma(u[k], wf[t * 128 + k], s);", .to = "s = id_h(fma(u[k], wf[t * 128 + k], s));" }} }, + .{ .kernel = "id_qk", .what = "drops fp16(q0 c - q1 s)", .edits = &.{.{ .from = "const float rq0 = id_h(q0 * c - q1 * s), rq1", .to = "const float rq0 = (q0 * c - q1 * s), rq1" }} }, + .{ .kernel = "id_qk", .what = "sums q's squares in fp16", .gate = .sums, .edits = &.{.{ .from = "simd_sum(rq0 * rq0 + rq1 * rq1)", .to = "float(simd_sum(half(rq0 * rq0 + rq1 * rq1)))" }} }, + .{ .kernel = "id_gate_add", .what = "drops fp16(g y)", .edits = &.{.{ .from = "half(float(x[i]) + id_h(float(g[at.x]) * float(y[i])));", .to = "half(float(x[i]) + (float(g[at.x]) * float(y[i])));" }} }, + .{ .kernel = "id_embed", .what = "drops fp16(z w + b)", .edits = &.{.{ .from = "x[i] = half(float(half(s + b[c])) + float(pe[i]));", .to = "x[i] = half((s + b[c]) + float(pe[i]));" }} }, + .{ .kernel = "id_embed", .what = "sums z w in fp16", .gate = .sums, .edits = &.{.{ .from = "s = fma(z[r * 16 + k], w[c * 16 + k], s);", .to = "s = id_h(fma(z[r * 16 + k], w[c * 16 + k], s));" }} }, + .{ .kernel = "id_cond", .what = "adds fp16(a + ba)", .edits = &.{.{ .from = "y[i] = half((a[i] + ba[at.x]) + (c[i] + bc[at.x]));", .to = "y[i] = half(id_h(a[i] + ba[at.x]) + (c[i] + bc[at.x]));" }} }, + .{ .kernel = "id_bias32", .what = "adds fp16(b)", .edits = &.{.{ .from = "y[ulong(at.y) * width + at.x] += b[at.x];", .to = "y[ulong(at.y) * width + at.x] += id_h(b[at.x]);" }} }, + .{ .kernel = "id_gemm32", .what = "sums in fp16", .gate = .sums, .edits = &.{.{ .from = mma_line, .to = mma_half }} }, + .{ .kernel = "id_gemm32", .what = "takes the NAX fp32 path", .gate = .sums, .edits = &.{.{ .from = "void id_gemm32(", .to = "void id_gemm32_simd(" }}, .append = nax_gemm32 }, + .{ .kernel = "id_attention_00", .what = "keys past Lk unmasked", .edits = &.{.{ .from = "s[f][e] = (16 * f + home.x + TF_COL(e)) < k_rows ? s[f][e] : kIdMasked;", .to = "s[f][e] = s[f][e];" }} }, +}; + +/// A test the suite made, and how it goes. +const Made = struct { ptr: *anyopaque, free: *const fn (std.mem.Allocator, *anyopaque) void }; + +pub const Suite = struct { + a: std.mem.Allocator, + libs: std.ArrayList(*Lib) = .empty, + names: std.ArrayList([]u8) = .empty, + made: std.ArrayList(Made) = .empty, + + fn label(s: *Suite, comptime fmt: []const u8, args: anytype) ![]const u8 { + const name = try std.fmt.allocPrint(s.a, fmt, args); + errdefer s.a.free(name); + try s.names.append(s.a, name); + return name; + } + + fn make(s: *Suite, comptime T: type, value: T) !*T { + const t = try s.a.create(T); + errdefer s.a.destroy(t); + t.* = value; + const f = struct { + fn free(a: std.mem.Allocator, p: *anyopaque) void { + const x: *T = @ptrCast(@alignCast(p)); + x.deinit(); + a.destroy(x); + } + }; + try s.made.append(s.a, .{ .ptr = t, .free = f.free }); + return t; + } + + /// Each kernel (or, with `ctl`, the one it breaks) on `lib`'s pipelines. + fn add(s: *Suite, device: mtl.Device, queue: mtl.Queue, lib: *Lib, ctl: ?Control, out: *std.ArrayList(g.Kernel), only: ?[]const u8) !void { + const suffix = if (ctl) |x| try s.label(" control: {s}", .{x.what}) else ""; + const gate: ?g.Gate = if (ctl) |x| x.gate else null; + const rig: rig_.Rig = .{ .a = s.a, .device = device, .queue = queue }; + const ends: id.Ends = .{ .a = s.a }; + const W = struct { + s: *Suite, + ctl: ?Control, + suffix: []const u8, + only: ?[]const u8, + /// The test's name when this pass registers `kernel`, else null. + fn name(w: @This(), kernel: []const u8) !?[]const u8 { + if (w.ctl) |x| if (!std.mem.eql(u8, x.kernel, kernel)) return null; + const full = try w.s.label("metal idol {s}{s}", .{ kernel, w.suffix }); + if (w.only) |needle| if (std.mem.indexOf(u8, full, needle) == null) return null; + return full; + } + }; + const w = W{ .s = s, .ctl = ctl, .suffix = suffix, .only = only }; + for ([_]gemm.Kind{ .linear, .gelu, .split }) |kind| for (0..4) |f| { + const form: u2 = @intCast(f); + if (kind == .split and form & 1 == 0) continue; + const fun = gemm.pipeline(kind, form); + const name = (try w.name(fun)) orelse continue; + const t = try s.make(gemm.Gemm16, .{ .a = s.a, .rig = rig, .lib = lib, .kind = kind, .form = form, .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + }; + if (try w.name("id_gemm32")) |name| { + const t = try s.make(gemm.Gemm32, .{ .a = s.a, .rig = rig, .lib = lib, .name = name }); + try out.append(s.a, t.kernel(gate)); + } + const norm_kinds = [_]struct { kind: norms.Kind, mode: ks.Norm, fun: []const u8 }{ + .{ .kind = .norm, .mode = .plain, .fun = "id_norm plain" }, + .{ .kind = .norm, .mode = .affine, .fun = "id_norm affine" }, + .{ .kind = .norm, .mode = .modulated, .fun = "id_norm modulated" }, + .{ .kind = .norm32, .mode = .affine, .fun = "id_norm32 affine" }, + .{ .kind = .norm32, .mode = .modulated, .fun = "id_norm32 modulated" }, + .{ .kind = .split, .mode = .affine, .fun = "id_norm_split" }, + }; + for (norm_kinds) |nk| { + const name = (try w.name(nk.fun)) orelse continue; + const t = try s.make(norms.Norm, .{ .a = s.a, .rig = rig, .lib = lib, .kind = nk.kind, .mode = nk.mode, .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + if (try w.name("id_repo")) |name| { + const t = try s.make(steps.Repo, .{ .a = s.a, .rig = rig, .lib = lib, .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + if (try w.name("id_qk")) |name| { + const t = try s.make(steps.Qk, .{ .a = s.a, .rig = rig, .lib = lib, .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + inline for (.{ .{ elems.GateAdd, "id_gate_add" }, .{ elems.Embed, "id_embed" }, .{ elems.Cond, "id_cond" }, .{ elems.Bias32, "id_bias32" } }) |e| { + if (try w.name(e[1])) |name| { + const t = try s.make(e[0], .{ .a = s.a, .rig = rig, .lib = lib, .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + } + for (attn.forms, 0..) |f, i| { + const name = (try w.name(f.name)) orelse continue; + const t = try s.make(attn.Attention, .{ .a = s.a, .rig = rig, .lib = lib, .form = i, .name = name }); + try out.append(s.a, t.kernel(gate)); + } + } + + pub fn create(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, out: *std.ArrayList(g.Kernel), only: ?[]const u8) !*Suite { + const s = try a.create(Suite); + s.* = .{ .a = a }; + errdefer s.destroy(); + if (!device.tensorUnits()) { + std.debug.print("UNSUPPORTED metal idol: the family runs on tensor units (M5) alone; its kernels unchecked here\n", .{}); + return s; + } + for ([_]?Control{null} ++ controls) |x| { + const lib = try a.create(Lib); + lib.* = .{ .a = a, .device = device, .ctl = x }; + s.libs.append(a, lib) catch |err| { + a.destroy(lib); + return err; + }; + try s.add(device, queue, lib, x, out, only); + } + return s; + } + + pub fn destroy(s: *Suite) void { + const a = s.a; + for (s.made.items) |m| m.free(a, m.ptr); + s.made.deinit(a); + for (s.libs.items) |l| { + l.deinit(); + a.destroy(l); + } + s.libs.deinit(a); + for (s.names.items) |n| a.free(n); + s.names.deinit(a); + a.destroy(s); + } +}; diff --git a/zig/tests/contract/metal_idol_attention.zig b/zig/tests/contract/metal_idol_attention.zig new file mode 100644 index 0000000000..fc75fac747 --- /dev/null +++ b/zig/tests/contract/metal_idol_attention.zig @@ -0,0 +1,233 @@ +//! id_attention under the contract: softmax(q k^T / 8) v at head width 64, against attention_ref's float64 formula. +const std = @import("std"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const id = @import("metal_idol_common.zig"); + +const c = id.c; +const g = id.g; +const D = id.D; +const kref = id.kref; + +const heads = id.cfg.heads; + +/// Query rows repeat at this period (prime to the 64-row blocks): float64 costs 89 queries a head over every key. +const query_period = 89; + +/// Each form's name, served lengths and edges by (Lq % 64 == 0)(Lk % 32 == 0); Lq is n, Lk is k. +pub const forms = [4]struct { name: []const u8, served: []const c.Shape, tiles: [3]c.Tile }{ + .{ + .name = "id_attention_00", + .served = &.{ .{ .rows = heads, .n = 12294, .k = 12294 }, .{ .rows = heads, .n = 4101, .k = 4101 } }, + .tiles = .{ .{}, .{ .size = 64 }, .{ .size = 32 } }, + }, + // no served call has Lq != Lk: the cross-length forms are checked at their tails + .{ + .name = "id_attention_01", + .served = &.{ .{ .rows = heads, .n = 4101, .k = 4128 }, .{ .rows = heads, .n = 100, .k = 64 } }, + .tiles = .{ .{}, .{ .size = 64 }, .{ .size = 64, .step = 32, .min = 32 } }, + }, + .{ + .name = "id_attention_10", + .served = &.{ .{ .rows = heads, .n = 4096, .k = 4101 }, .{ .rows = heads, .n = 128, .k = 70 } }, + .tiles = .{ .{}, .{ .size = 128, .step = 64, .min = 64 }, .{ .size = 32 } }, + }, + .{ + .name = "id_attention_11", + .served = &.{.{ .rows = heads, .n = 8192, .k = 8192 }}, + .tiles = .{ .{}, .{ .size = 128, .step = 64, .min = 64 }, .{ .size = 64, .step = 32, .min = 32 } }, + }, +}; + +pub fn form(lq: usize, lk: usize) usize { + return @as(usize, @intFromBool(lq % 64 == 0)) * 2 + @intFromBool(lk % 32 == 0); +} + +/// The scores' 64-deep sums, then the keys' sums. +fn attentionDepth(s: c.Shape) usize { + return D + s.k; +} + +pub const Attention = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *id.Lib, + form: usize, + name: []const u8, + lq: usize = 0, + lk: usize = 0, + stress: bool = false, + q: []f16 = &.{}, + k: []f16 = &.{}, + v: []f16 = &.{}, + ins: [3]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + /// Rows: the launch's 16 heads (fixed; a gate's fewer rows run all 16); n: query rows Lq; k: key rows Lk. + pub fn spec(name: []const u8, f: usize) c.Spec { + return .{ .name = name, .accum = .f32, .depth = attentionDepth, .ops = 8, .out = .f16, .tiles = forms[f].tiles, .served = forms[f].served }; + } + + pub fn deinit(t: *Attention) void { + t.free(); + t.rig.deinit(); + } + fn free(t: *Attention) void { + t.a.free(t.q); + t.a.free(t.k); + t.a.free(t.v); + t.q = &.{}; + t.k = &.{}; + t.v = &.{}; + } + pub fn kernel(t: *Attention, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name, t.form), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Attention); + + pub fn prepare(t: *Attention, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.n == 0 or shape.k == 0 or form(shape.n, shape.k) != t.form) return false; + t.rig.reset(); + t.free(); + t.lq = shape.n; + t.lk = shape.k; + t.stress = fill == .stress; + const L = t.lk; + t.q = try t.a.alloc(f16, heads * t.lq * D); + t.k = try t.a.alloc(f16, heads * L * D); + t.v = try t.a.alloc(f16, heads * L * D); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (t.stress) { + // equal scores (q = k = 0), each column's values integers whose partial sums reach ~2^23 and cancel to 4 + @memset(t.q, 0); + @memset(t.k, 0); + const half = L / 2; + for (0..heads) |h| for (0..L) |j| { + const at = if (j < half) j else j - half; + const positive: f32 = if (at % 2 == 0) 127 else 4064; + const residual: f32 = if (at < 8 and at % 2 == 0) 1 else 0; + const v: f32 = if (j < half) positive else if (j < 2 * half) residual - positive else 0; + @memset(t.v[(h * L + j) * D ..][0..D], @floatCast(v)); + }; + } else { + id.randomF16(rng, t.q, -2, 2); + id.randomF16(rng, t.k, -2, 2); + id.randomF16(rng, t.v, -1, 1); + } + if (fill == .nonfinite) { + // head 0: values NaN, +Inf, -Inf, and both infinities in one column; head 1: a NaN query, an infinite key + for ([_][2]usize{ .{ 0, 0 }, .{ 1, 1 }, .{ 2, 2 }, .{ 3, 3 }, .{ 4, 3 } }, [_]f16{ id.nan16, id.inf16, -id.inf16, id.inf16, -id.inf16 }) |p, v| if (p[0] < L) { + t.v[p[0] * D + p[1]] = v; + }; + t.q[(t.lq + 0) * D + 5] = id.nan16; + if (L > 7) t.k[(L + 7) * D + 3] = id.inf16; + } + // every head's queries repeat at the period, so float64 computes each distinct one once + const p = @min(t.lq, query_period); + for (0..heads) |h| for (p..t.lq) |i| { + @memcpy(t.q[(h * t.lq + i) * D ..][0..D], t.q[(h * t.lq + i % p) * D ..][0..D]); + }; + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.q)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.k)); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.v)); + t.out = try t.rig.output(heads * t.lq * D * 2); + return true; + } + + pub fn run(t: *Attention) !void { + try t.lib.need(forms[t.form].name); + const p = t.rig.begin(); + t.lib.k.attention(p.e, kref(t.ins[0]), kref(t.ins[1]), kref(t.ins[2]), kref(t.out), t.lq, t.lk); + try t.rig.finish(p); + } + + pub fn count(t: *Attention) usize { + return heads * t.lq * D; + } + + /// attention_ref.head's formula, one pass over the keys a distinct query; O rows [Lq, heads * 64]. + pub fn reference(t: *Attention, want: []f64, mag: []f64) !void { + const L = t.lk; + const scale: f64 = @as(f32, 1.0 / @sqrt(@as(f32, D))); + const kf = try t.a.alloc(f64, L * D); + defer t.a.free(kf); + const vf = try t.a.alloc(f64, L * D); + defer t.a.free(vf); + const s = try t.a.alloc(f64, L); + defer t.a.free(s); + var qf: [D]f64 = undefined; + var out: [D]f64 = undefined; + var tol: [D]f64 = undefined; + const p = @min(t.lq, query_period); + for (0..heads) |h| { + for (kf, vf, t.k[h * L * D ..][0 .. L * D], t.v[h * L * D ..][0 .. L * D]) |*kd, *vd, kb, vb| { + kd.* = kb; + vd.* = vb; + } + for (0..p) |i| { + var zero = true; + for (&qf, t.q[(h * t.lq + i) * D ..][0..D]) |*x, b| { + x.* = b; + zero = zero and x.* == 0; + } + var m: f64 = -std.math.inf(f64); + var nan = false; + for (s, 0..) |*sj, j| { + var dot: f64 = 0; + if (!zero) { + const Vec = @Vector(8, f64); + var acc: Vec = @splat(0); + var e: usize = 0; + while (e < D) : (e += 8) acc += @as(Vec, qf[e..][0..8].*) * @as(Vec, kf[j * D + e ..][0..8].*); + dot = @reduce(.Add, acc); + } + sj.* = dot * scale; + nan = nan or std.math.isNan(sj.*); + m = @max(m, sj.*); + } + if (nan) m = std.math.nan(f64); + var l: f64 = 0; + @memset(&out, 0); + @memset(&tol, 0); + for (s, 0..) |sj, j| { + const pj = @exp(sj - m); + l += pj; + for (&out, &tol, vf[j * D ..][0..D]) |*o, *tl, vv| { + o.* += pj * vv; + tl.* += pj * @abs(vv); + } + } + var r = i; + while (r < t.lq) : (r += p) { + const at = (r * heads + h) * D; + for (out, tol, 0..) |o, tl, e| { + want[at + e] = o / l; + // P enters P x V at the tensor unit's fp32 (~2^-10 relative); O rounds to fp16 + mag[at + e] = tl / l * 0x1p-9 + @abs(o / l) * 0x1p-10 + 1e-6; + } + } + } + } + } + + pub fn bounds(t: *Attention, want: []const f64, tolerance: []const f64, allow: []f64) !void { + for (want, tolerance, allow) |w, tl, *l| l.* = if (t.stress) c.boundExact(spec(t.name, t.form), w) else if (std.math.isFinite(w)) tl else 0; + } + + pub fn simulate(t: *Attention, acc: c.Accum, order: g.Order, out: []f64) !void { + if (!t.stress) return error.AttentionStressNeedsEqualScores; + const L = t.lk; + const terms = try t.a.alloc(f64, L); + defer t.a.free(terms); + const ones = try t.a.alloc(f64, L); + defer t.a.free(ones); + @memset(ones, 1); + const l = contract.rows.sumIn(acc, order, ones); + for (0..heads) |h| for (0..D) |e| { + for (terms, 0..) |*v, j| v.* = t.v[(h * L + j) * D + e]; + const o = contract.rows.sumIn(acc, order, terms) / l; + for (0..t.lq) |i| out[(i * heads + h) * D + e] = o; + }; + } +}; diff --git a/zig/tests/contract/metal_idol_common.zig b/zig/tests/contract/metal_idol_common.zig new file mode 100644 index 0000000000..222b6b22b4 --- /dev/null +++ b/zig/tests/contract/metal_idol_common.zig @@ -0,0 +1,245 @@ +//! Idol's kernel tests' shared parts: float64 values at fp16 points, their ranges, inputs and the family's library. +const std = @import("std"); +const mtl = @import("metal"); +pub const tf = @import("tensorfold"); +const contract = @import("contract"); +const sources = @import("kernel_sources"); +const rig_ = @import("rig.zig"); + +pub const c = contract.spec; +pub const g = contract.gates; +pub const rows = contract.rows; +pub const ks = tf.idol.kernels; +pub const cfg = tf.idol.config; +pub const H = cfg.hidden; +pub const D = cfg.head_dim; + +pub fn noSums(_: c.Shape) usize { + return 0; +} + +/// A value rounded once to fp16 (the kernels' half(fp32) of a value float64 holds exactly, or float64's own). +pub fn hf(v: f64) f64 { + return @as(f16, @floatCast(v)); +} + +pub fn kref(r: rig_.Ref) ks.Ref { + return .{ .buf = r.buf, .off = r.off }; +} + +/// Every value the kernel may hold at a step: its fp32 steps widen both ends, its fp16 points round them. +pub const Range = struct { + lo: f64, + hi: f64, + pub fn of(v: f64) Range { + return .{ .lo = v, .hi = v }; + } + /// An fp32 step good to `ulps` units in the last place (a correctly rounded one: 0.5). + pub fn fp32(a: Range, ulps: f64) Range { + const u = ulps * 0x1p-23; + return .{ .lo = a.lo - u * @abs(a.lo) - 0x1p-149, .hi = a.hi + u * @abs(a.hi) + 0x1p-149 }; + } + pub fn fp16(a: Range) Range { + return .{ .lo = hf(a.lo), .hi = hf(a.hi) }; + } + pub fn add(a: Range, b: Range) Range { + return .{ .lo = a.lo + b.lo, .hi = a.hi + b.hi }; + } + pub fn neg(a: Range) Range { + return .{ .lo = -a.hi, .hi = -a.lo }; + } + pub fn sub(a: Range, b: Range) Range { + return a.add(b.neg()); + } + pub fn widen(a: Range, e: f64) Range { + return .{ .lo = a.lo - e, .hi = a.hi + e }; + } + pub fn mul(a: Range, b: Range) Range { + const p = [4]f64{ a.lo * b.lo, a.lo * b.hi, a.hi * b.lo, a.hi * b.hi }; + return .{ .lo = @min(@min(p[0], p[1]), @min(p[2], p[3])), .hi = @max(@max(p[0], p[1]), @max(p[2], p[3])) }; + } + /// An increasing function at both ends. + pub fn up(a: Range, comptime f: fn (f64) f64) Range { + return .{ .lo = f(a.lo), .hi = f(a.hi) }; + } + /// 1 / sqrt(a), a positive. + pub fn rsqrt(a: Range) Range { + return .{ .lo = 1 / @sqrt(a.hi), .hi = 1 / @sqrt(a.lo) }; + } + /// 1 / a, a positive. + pub fn recip(a: Range) Range { + return .{ .lo = 1 / a.hi, .hi = 1 / a.lo }; + } + /// a^2 over the range. + pub fn square(a: Range) Range { + const m = @max(a.lo * a.lo, a.hi * a.hi); + return .{ .lo = if (a.lo <= 0 and a.hi >= 0) 0 else @min(a.lo * a.lo, a.hi * a.hi), .hi = m }; + } +}; + +pub fn tanh(x: f64) f64 { + return std.math.tanh(x); +} + +pub fn sqrt(x: f64) f64 { + return @sqrt(x); +} + +// Metal's error bounds with the engine's options (fast math on, the kernels' own `precise::` calls). +pub const div_ulps = 2.5; +pub const rsqrt_ulps = 2; +pub const sqrt_ulps = 3; +pub const exp_ulps = 4; +pub const tanh_ulps = 5; +pub const trig_ulps = 4; + +/// f over [lo, hi] (cos or sin): its ends, and 1 or -1 where an extreme (a multiple of pi past `at`) falls inside. +fn trigOver(lo: f64, hi: f64, comptime f: fn (f64) f64, at: f64) Range { + const pi = std.math.pi; + if (!(hi - lo < 2 * pi)) return .{ .lo = -1, .hi = 1 }; + var out = Range{ .lo = @min(f(lo), f(hi)), .hi = @max(f(lo), f(hi)) }; + var k = @ceil((lo - at) / pi); + while (k * pi + at <= hi) : (k += 1) { + if (@mod(k, 2) == 0) out.hi = 1 else out.lo = -1; + } + return out; +} + +fn cos(x: f64) f64 { + return @cos(x); +} + +fn sin(x: f64) f64 { + return @sin(x); +} + +pub fn cosOver(lo: f64, hi: f64) Range { + return trigOver(lo, hi, cos, 0); +} + +pub fn sinOver(lo: f64, hi: f64) Range { + return trigOver(lo, hi, sin, std.math.pi / 2.0); +} + +/// Each output's float64 value at the kernel's rounding points, and its range, kept for the bounds and intervals. +pub const Ends = struct { + a: std.mem.Allocator, + lo: []f64 = &.{}, + hi: []f64 = &.{}, + + pub fn deinit(e: *Ends) void { + e.a.free(e.lo); + e.a.free(e.hi); + e.lo = &.{}; + e.hi = &.{}; + } + + pub fn size(e: *Ends, n: usize) !void { + e.deinit(); + e.lo = try e.a.alloc(f64, n); + e.hi = try e.a.alloc(f64, n); + } + + pub fn put(e: *Ends, want: []f64, mag: []f64, i: usize, w: f64, r: Range) void { + want[i] = w; + mag[i] = if (std.math.isFinite(w)) @abs(w) else 0; + e.lo[i] = r.lo; + e.hi[i] = r.hi; + } + + /// The allowance about the float64 value: its range's farther end (a non-finite value is checked by class). + pub fn bounds(e: *const Ends, want: []const f64, allow: []f64) void { + for (allow, want, e.lo, e.hi) |*l, w, lo, hi| l.* = if (std.math.isFinite(w)) @max(@abs(w - lo), @abs(hi - w)) else 0; + } + + pub fn intervals(e: *const Ends, low: []f64, high: []f64) void { + @memcpy(low, e.lo); + @memcpy(high, e.hi); + } +}; + +pub fn randomF16(rng: std.Random, out: []f16, lo: f32, hi: f32) void { + for (out) |*v| v.* = @floatCast(lo + rng.float(f32) * (hi - lo)); +} + +pub fn randomF32(rng: std.Random, out: []f32, lo: f32, hi: f32) void { + for (out) |*v| v.* = lo + rng.float(f32) * (hi - lo); +} + +/// An integer in [lo, hi]. +pub fn int(rng: std.Random, lo: i32, hi: i32) f32 { + return @floatFromInt(lo + @as(i32, @intCast(rng.uintLessThan(u32, @intCast(hi - lo + 1))))); +} + +/// +1 or -1. +pub fn sign(rng: std.Random) f32 { + return if (rng.boolean()) 1 else -1; +} + +pub const nan16: f16 = std.math.nan(f16); +pub const inf16: f16 = std.math.inf(f16); +pub const nan32: f32 = std.math.nan(f32); +pub const inf32: f32 = std.math.inf(f32); + +/// [count, width] rows of which row r is `distinct`'s row r % (its rows): a period prime to every tile. +pub fn periodic(comptime T: type, a: std.mem.Allocator, distinct: []const T, width: usize, count: usize) ![]T { + const p = distinct.len / width; + const out = try a.alloc(T, count * width); + for (0..count) |r| @memcpy(out[r * width ..][0..width], distinct[(r % p) * width ..][0..width]); + return out; +} + +/// One text change a control makes to the family's source. +pub const Edit = struct { from: []const u8, to: []const u8 }; + +/// A control: the kernel it breaks, the gate that must go red, its source edits and any text it appends. +pub const Control = struct { kernel: []const u8, what: []const u8, gate: g.Gate = .shapes, edits: []const Edit = &.{}, append: []const u8 = "" }; + +/// The family's source as the engine compiles it (or a control's copy), its pipelines made as tests ask for them. +pub const Lib = struct { + a: std.mem.Allocator, + device: mtl.Device, + ctl: ?Control, + k: ks.Kernels = undefined, + built: bool = false, + made: [ks.names.len]bool = @splat(false), + + fn text(l: *Lib) ![]u8 { + var out = try tf.frags.source(l.device, l.a, sources.idol); + const x = l.ctl orelse return out; + for (x.edits) |e| { + errdefer l.a.free(out); + if (std.mem.count(u8, out, e.from) != 1) { + std.debug.print("control {s}: text not found once: {s}\n", .{ x.what, e.from }); + return error.ControlTextNotFoundOnce; + } + const next = try std.mem.replaceOwned(u8, l.a, out, e.from, e.to); + l.a.free(out); + out = next; + } + if (x.append.len == 0) return out; + defer l.a.free(out); + return std.mem.concat(l.a, u8, &.{ out, "\n", x.append }); + } + + /// The pipeline `name` made on this library (compiled on first use). + pub fn need(l: *Lib, name: []const u8) !void { + if (!l.built) { + const src = try l.text(); + defer l.a.free(src); + l.k.own = try mtl.Library.fromSource(l.device, src, mtl.CompileOptions.mlx()); + l.built = true; + } + for (ks.names, 0..) |n, i| if (std.mem.eql(u8, n, name)) { + if (!l.made[i]) l.k.p[i] = try mtl.Pipeline.init(l.device, l.k.own, n, false); + l.made[i] = true; + return; + }; + return error.NoSuchPipeline; + } + + pub fn deinit(l: *Lib) void { + for (l.k.p, l.made) |p, m| if (m) p.deinit(); + if (l.built) l.k.own.deinit(); + } +}; diff --git a/zig/tests/contract/metal_idol_gemm.zig b/zig/tests/contract/metal_idol_gemm.zig new file mode 100644 index 0000000000..7cb64720f4 --- /dev/null +++ b/zig/tests/contract/metal_idol_gemm.zig @@ -0,0 +1,430 @@ +//! id_linear (plain and GELU), id_split and id_gemm32 under the contract: products against float64, sums proven fp32. +const std = @import("std"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const id = @import("metal_idol_common.zig"); + +const c = id.c; +const g = id.g; +const H = id.H; +const Range = id.Range; +const Ends = id.Ends; +const hf = id.hf; +const kref = id.kref; + +/// Rows of x and of w (with the bias) repeat at these periods, prime to every tile: 97 x 131 dots a shape. +const row_period = 97; +const col_period = 131; + +pub const Kind = enum { linear, gelu, split }; + +/// An output's float64 value and its range. +const Out = struct { w: f64, r: Range }; + +/// The M5 tensor unit's fp32 sums: 36 roundings of 2^-24 of |acc| + sum |a b| a 16-deep step (attention_range's). +fn naxDepth(s: c.Shape) usize { + return s.k / 16 * contract.attention_range.mma_step + 8; +} + +/// float64's dot of an fp16 row pair and the tensor unit's allowance on its fp32 sums, step by 16-deep step. +fn naxDot(x: []const f16, w: []const f16) [3]f64 { + var s: f64 = 0; + var spread: f64 = 0; + var mag: f64 = 0; + var i: usize = 0; + while (i < x.len) : (i += 16) { + var t: f64 = 0; + var abs: f64 = 0; + for (x[i..@min(i + 16, x.len)], w[i..@min(i + 16, x.len)]) |a, b| { + const p = @as(f64, a) * @as(f64, b); + t += p; + abs += @abs(p); + } + spread += @abs(s) + abs; + s += t; + mag += abs; + } + return .{ s, spread * contract.attention_range.mma_step * 0x1p-24, mag }; +} + +const gelu_c: f64 = @as(f32, 0.7978845608028654); +const gelu_k: f64 = @as(f32, 0.044715); + +fn inner(x: f64) f64 { + return gelu_c * (x + gelu_k * x * x * x); +} + +/// torch's tanh GELU in float64 at the kernel's fp32 constants. +fn gelu(x: f64) f64 { + return 0.5 * x * (1 + std.math.tanh(inner(x))); +} + +/// The GELU's fp32 steps over a range of fp16 inputs: its inner polynomial rises, tanh rises, the rest as products. +fn geluRange(v: Range) Range { + const t = (Range{ .lo = inner(v.lo), .hi = inner(v.hi) }).fp32(3).up(id.tanh).fp32(id.tanh_ulps); + return (Range{ .lo = 0.5 * v.lo, .hi = 0.5 * v.hi }).mul(t.add(Range.of(1)).fp32(0.5)).fp32(0.5); +} + +const Served = struct { served: []const c.Shape, tiles: [3]c.Tile }; + +/// Forms by (M % 64 == 0)(N % 128 == 0); rows: 64-row tiles in whole-row forms, a last tile's rows first in the rest. +fn envelope(kind: Kind, form: u2) Served { + const row_tail: c.Tile = .{ .size = 64 }; + const row_whole: c.Tile = .{ .size = 2, .step = 1, .min = 1 }; + const col_tail: c.Tile = .{ .size = 128 }; + const col_whole: c.Tile = .{ .size = 256, .step = 128, .min = 128 }; + const k_tile: c.Tile = .{ .size = 256, .step = 32, .min = 32 }; + const rt = if (form & 2 != 0) row_whole else row_tail; + if (kind == .split) return .{ + .served = if (form & 2 != 0) &.{.{ .rows = 8192 / 64, .n = 256, .k = 3 * H }} else &.{ .{ .rows = 6, .n = 256, .k = 3 * H }, .{ .rows = 12294, .n = 256, .k = 3 * H }, .{ .rows = 4101, .n = 256, .k = 3 * H } }, + .tiles = .{ rt, .{}, .{} }, + }; + const ct = if (form & 1 != 0) col_whole else col_tail; + const tiles: [3]c.Tile = .{ rt, ct, k_tile }; + if (kind == .gelu) return .{ + .tiles = tiles, + .served = switch (form) { + 3 => &.{.{ .rows = 8192 / 64, .n = 4 * H, .k = H }}, + 1 => &.{ .{ .rows = 6, .n = 4 * H, .k = H }, .{ .rows = 12294, .n = 4 * H, .k = H }, .{ .rows = 4101, .n = 4 * H, .k = H } }, + // no served call takes a width short of 128 columns: these forms are checked at one + 2 => &.{.{ .rows = 4096 / 64, .n = 1000, .k = H }}, + 0 => &.{ .{ .rows = 5, .n = 1000, .k = H }, .{ .rows = 4101, .n = 1000, .k = H } }, + }, + }; + return .{ .tiles = tiles, .served = switch (form) { + 3 => &.{ .{ .rows = 8192 / 64, .n = 3 * H, .k = H }, .{ .rows = 8192 / 64, .n = H, .k = H }, .{ .rows = 8192 / 64, .n = H, .k = 4 * H } }, + 1 => &.{ .{ .rows = 6, .n = 3 * H, .k = H }, .{ .rows = 12294, .n = 3 * H, .k = H }, .{ .rows = 12294, .n = H, .k = H }, .{ .rows = 12294, .n = H, .k = 4 * H }, .{ .rows = 4101, .n = 3 * H, .k = H }, .{ .rows = 4101, .n = H, .k = H }, .{ .rows = 4101, .n = H, .k = 4 * H } }, + 2 => &.{.{ .rows = 4096 / 64, .n = 1000, .k = H }}, + 0 => &.{ .{ .rows = 5, .n = 1000, .k = H }, .{ .rows = 4101, .n = 1000, .k = H } }, + } }; +} + +/// The pipeline a kind and form run. +pub fn pipeline(kind: Kind, form: u2) []const u8 { + const names = [3][4][]const u8{ + .{ "id_linear_00", "id_linear_01", "id_linear_10", "id_linear_11" }, + .{ "id_linear_gelu_00", "id_linear_gelu_01", "id_linear_gelu_10", "id_linear_gelu_11" }, + .{ "id_split_0", "id_split_0", "id_split_1", "id_split_1" }, + }; + return names[@backingInt(kind)][form]; +} + +pub const Gemm16 = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *id.Lib, + kind: Kind, + /// split's forms are 1 and 3: its N, 256, is whole 128-column tiles. + form: u2, + name: []const u8, + m: usize = 0, + n: usize = 0, + k: usize = 0, + stress: bool = false, + x: []f16 = &.{}, + w: []f16 = &.{}, + b: []f16 = &.{}, + out: rig_.Ref = undefined, + ins: [3]rig_.Ref = undefined, + ends: Ends, + + pub fn spec(name: []const u8, kind: Kind, form: u2) c.Spec { + const e = envelope(kind, form); + return .{ + .name = name, + .accum = .f32, + .depth = naxDepth, + .ops = 1, + .out = if (kind == .split) .f32 else .f16, + .tiles = e.tiles, + .served = e.served, + }; + } + + pub fn deinit(t: *Gemm16) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *Gemm16) void { + t.a.free(t.x); + t.a.free(t.w); + t.a.free(t.b); + t.x = &.{}; + t.w = &.{}; + t.b = &.{}; + } + pub fn kernel(t: *Gemm16, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name, t.kind, t.form), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Gemm16); + + pub fn prepare(t: *Gemm16, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + const m = if (t.form & 2 != 0) shape.rows * 64 else shape.rows; + if (m == 0 or shape.n == 0 or shape.k == 0 or shape.k % 2 != 0) return false; + if (t.kind == .split and (shape.n != 256 or shape.k != 3 * H)) return false; + const form = @as(u2, @intFromBool(m % 64 == 0)) * 2 + @intFromBool(shape.n % 128 == 0); + if (form != t.form) return false; + t.rig.reset(); + t.free(); + t.m = m; + t.n = shape.n; + t.k = shape.k; + t.stress = fill == .stress; + const K = shape.k; + const px = @min(m, row_period); + const pw = @min(shape.n, col_period); + t.x = try t.a.alloc(f16, px * K); + t.w = try t.a.alloc(f16, pw * K); + t.b = try t.a.alloc(f16, pw); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (t.stress) { + // integers in [hi / 2, hi), hi^2 K <= 2^24: fp32 sums exact in any order, a 32nd of a half past fp16's + const hi: i32 = @intFromFloat(@sqrt(@as(f64, 0x1p24) / @as(f64, @floatFromInt(K)))); + for (0..px) |r| for (0..K / 2) |i| { + const v = id.int(rng, @divTrunc(hi, 2), hi - 1); + t.x[r * K + i] = @floatCast(v); + // x's halves differ by 1 in their first 8: the sums cancel to a few hundred + t.x[r * K + K / 2 + i] = @floatCast(v - @as(f32, if (i < 8) 1 else 0)); + }; + for (0..pw) |j| for (0..K / 2) |i| { + const v = id.int(rng, @divTrunc(hi, 2), hi - 1); + t.w[j * K + i] = @floatCast(v); + t.w[j * K + K / 2 + i] = @floatCast(-v); + }; + for (t.b) |*v| v.* = @floatCast(id.int(rng, -4, 4)); + } else { + id.randomF16(rng, t.x, -2, 2); + id.randomF16(rng, t.w, -0.25, 0.25); + id.randomF16(rng, t.b, -1, 1); + } + if (fill == .nonfinite) { + const at = [_][2]usize{ .{ 0, K / 3 }, .{ 1, K / 2 }, .{ 2, K - 1 }, .{ 3, 0 }, .{ 3, K - 1 } }; + for (at, [_]f16{ id.nan16, id.inf16, -id.inf16, id.inf16, -id.inf16 }) |p, v| if (p[0] < px) { + t.x[p[0] * K + p[1]] = v; + }; + if (pw > 2) t.b[2] = id.inf16; + } + const xs = try id.periodic(f16, t.a, t.x, K, m); + defer t.a.free(xs); + const ws = try id.periodic(f16, t.a, t.w, K, shape.n); + defer t.a.free(ws); + const bs = try id.periodic(f16, t.a, t.b, 1, shape.n); + defer t.a.free(bs); + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(xs)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(ws)); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(bs)); + t.out = try t.rig.output(m * shape.n * @as(usize, if (t.kind == .split) 4 else 2)); + return true; + } + + pub fn run(t: *Gemm16) !void { + try t.lib.need(pipeline(t.kind, t.form)); + const p = t.rig.begin(); + const k = &t.lib.k; + switch (t.kind) { + .split => k.split(p.e, kref(t.ins[0]), kref(t.ins[1]), kref(t.out), t.m), + else => k.linear(p.e, kref(t.ins[0]), kref(t.ins[1]), kref(t.ins[2]), kref(t.out), t.m, t.n, t.k, t.kind == .gelu), + } + try t.rig.finish(p); + } + + pub fn count(t: *Gemm16) usize { + return t.m * t.n; + } + + /// The output of a sum `s` (range `sr`) of column j: its bias, fp16 rounding and GELU; split's fp32 sum as is. + fn finish(t: *const Gemm16, s: f64, sr: Range, j: usize) Out { + if (t.kind == .split) return .{ .w = s, .r = sr.fp32(0.5) }; + const b: f64 = t.b[j]; + const v = hf(s + b); + const vr = sr.add(Range.of(b)).fp32(0.5).fp16(); + if (t.kind == .linear) return .{ .w = v, .r = vr }; + return .{ .w = hf(gelu(v)), .r = geluRange(vr).fp16() }; + } + + /// The distinct outputs' values and ranges by `value(r, j)`, then every output from its row's and column's. + fn table(t: *Gemm16, ends: *Ends, want: []f64, mag: []f64, ctx: anytype, comptime value: anytype) !void { + const px = t.x.len / t.k; + const pw = t.w.len / t.k; + const tw = try t.a.alloc(f64, 3 * px * pw); + defer t.a.free(tw); + for (0..px) |r| for (0..pw) |j| { + const o = value(ctx, t, r, j); + tw[3 * (r * pw + j) ..][0..3].* = .{ o.w, o.r.lo, o.r.hi }; + }; + for (0..t.m) |r| for (0..t.n) |j| { + const e = tw[3 * ((r % px) * pw + j % pw) ..][0..3]; + ends.put(want, mag, r * t.n + j, e[0], .{ .lo = e[1], .hi = e[2] }); + }; + } + + pub fn reference(t: *Gemm16, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + const f = struct { + fn value(_: void, s: *const Gemm16, r: usize, j: usize) Out { + const d = naxDot(s.x[r * s.k ..][0..s.k], s.w[j * s.k ..][0..s.k]); + return s.finish(d[0], if (s.stress) Range.of(d[0]) else Range.of(d[0]).widen(d[1]), j); + } + }; + try t.table(&t.ends, want, mag, {}, f.value); + } + pub fn bounds(t: *Gemm16, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *Gemm16, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } + + pub fn simulate(t: *Gemm16, acc: c.Accum, order: g.Order, out: []f64) !void { + const mag = try t.a.alloc(f64, out.len); + defer t.a.free(mag); + var scratch: Ends = .{ .a = t.a }; + defer scratch.deinit(); + try scratch.size(t.count()); + const terms = try t.a.alloc(f64, t.k); + defer t.a.free(terms); + const Ctx = struct { acc: c.Accum, order: g.Order, terms: []f64 }; + const f = struct { + fn value(x: Ctx, s: *const Gemm16, r: usize, j: usize) Out { + for (x.terms, s.x[r * s.k ..][0..s.k], s.w[j * s.k ..][0..s.k]) |*v, a, b| v.* = @as(f64, a) * @as(f64, b); + const sum = id.rows.sumIn(x.acc, x.order, x.terms); + return s.finish(sum, Range.of(sum), j); + } + }; + try t.table(&scratch, out, mag, Ctx{ .acc = acc, .order = order, .terms = terms }, f.value); + } +}; + +// --- id_gemm32: fp32 x [M, K] times w [N, K]^T, sums that must be fp32's (the NAX fp32 path's are not) ----------- + +/// Each 8x8 step's eight FMAs in one chain over K. +fn gemm32Depth(s: c.Shape) usize { + return s.k + 8; +} + +pub const Gemm32 = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *id.Lib, + name: []const u8, + shape: c.Shape = undefined, + x: []f32 = &.{}, + w: []f32 = &.{}, + ins: [2]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + pub fn spec(name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = gemm32Depth, + .out = .f32, + .tiles = .{ .{ .size = 64 }, .{ .size = 64, .step = 8, .min = 8 }, .{ .size = 32, .step = 32, .min = 32 } }, + .served = &.{ .{ .rows = 4101, .n = H, .k = 1280 }, .{ .rows = 4101, .n = H, .k = 128 }, .{ .rows = 8192, .n = 16, .k = H } }, + }; + } + + pub fn deinit(t: *Gemm32) void { + t.free(); + t.rig.deinit(); + } + fn free(t: *Gemm32) void { + t.a.free(t.x); + t.a.free(t.w); + t.x = &.{}; + t.w = &.{}; + } + pub fn kernel(t: *Gemm32, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Gemm32); + + pub fn prepare(t: *Gemm32, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + const n_ok = shape.n % 64 == 0 or (shape.n < 64 and shape.n % 8 == 0); + if (shape.rows == 0 or shape.n == 0 or shape.k == 0 or shape.k % 32 != 0 or !n_ok) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + const K = shape.k; + const px = @min(shape.rows, row_period); + const pw = @min(shape.n, col_period); + t.x = try t.a.alloc(f32, px * K); + t.w = try t.a.alloc(f32, pw * K); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (fill == .stress) { + // 13-bit integers against +-1, in turn on either side: exact fp32 products and sums below 2^24 + for (0..px) |r| for (0..K) |i| { + t.x[r * K + i] = id.sign(rng) * if (i % 2 == 0) id.int(rng, 4097, 8191) else 1; + }; + for (0..pw) |j| for (0..K) |i| { + t.w[j * K + i] = id.sign(rng) * if (i % 2 == 1) id.int(rng, 4097, 8191) else 1; + }; + } else { + id.randomF32(rng, t.x, -2, 2); + id.randomF32(rng, t.w, -0.25, 0.25); + } + if (fill == .nonfinite) { + const at = [_][2]usize{ .{ 0, K / 3 }, .{ 1, K / 2 }, .{ 2, K - 1 }, .{ 3, 0 }, .{ 3, K - 1 } }; + for (at, [_]f32{ id.nan32, id.inf32, -id.inf32, id.inf32, -id.inf32 }) |p, v| if (p[0] < px) { + t.x[p[0] * K + p[1]] = v; + }; + } + const xs = try id.periodic(f32, t.a, t.x, K, shape.rows); + defer t.a.free(xs); + const ws = try id.periodic(f32, t.a, t.w, K, shape.n); + defer t.a.free(ws); + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(xs)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(ws)); + t.out = try t.rig.output(shape.rows * shape.n * 4); + return true; + } + + pub fn run(t: *Gemm32) !void { + try t.lib.need("id_gemm32"); + const p = t.rig.begin(); + t.lib.k.linear32(p.e, kref(t.ins[0]), kref(t.ins[1]), kref(t.out), t.shape.rows, t.shape.n, t.shape.k, t.shape.n); + try t.rig.finish(p); + } + + pub fn count(t: *Gemm32) usize { + return t.shape.rows * t.shape.n; + } + + fn each(t: *Gemm32, want: []f64, mag: []f64, acc: ?c.Accum, order: g.Order) !void { + const K = t.shape.k; + const px = t.x.len / K; + const pw = t.w.len / K; + const tw = try t.a.alloc(f64, 2 * px * pw); + defer t.a.free(tw); + const terms = try t.a.alloc(f64, K); + defer t.a.free(terms); + for (0..px) |r| for (0..pw) |j| { + var s: f64 = 0; + var m: f64 = 0; + for (terms, t.x[r * K ..][0..K], t.w[j * K ..][0..K]) |*v, a, b| { + v.* = @as(f64, a) * @as(f64, b); + s += v.*; + m += @abs(v.*); + } + tw[2 * (r * pw + j) ..][0..2].* = .{ if (acc) |x| id.rows.sumIn(x, order, terms) else s, m }; + }; + for (0..t.shape.rows) |r| for (0..t.shape.n) |j| { + const e = tw[2 * ((r % px) * pw + j % pw) ..][0..2]; + want[r * t.shape.n + j] = e[0]; + mag[r * t.shape.n + j] = e[1]; + }; + } + + pub fn reference(t: *Gemm32, want: []f64, mag: []f64) !void { + try t.each(want, mag, null, .sequential); + } + + pub fn simulate(t: *Gemm32, acc: c.Accum, order: g.Order, out: []f64) !void { + const mag = try t.a.alloc(f64, out.len); + defer t.a.free(mag); + try t.each(out, mag, acc, order); + } +}; diff --git a/zig/tests/contract/metal_idol_norms.zig b/zig/tests/contract/metal_idol_norms.zig new file mode 100644 index 0000000000..c9ea523659 --- /dev/null +++ b/zig/tests/contract/metal_idol_norms.zig @@ -0,0 +1,251 @@ +//! id_norm, id_norm32 and id_norm_split under the contract: LayerNorm32 in fp32, at idol-sdk's fp16 points. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const id = @import("metal_idol_common.zig"); + +const c = id.c; +const g = id.g; +const ks = id.ks; +const cfg = id.cfg; +const H = id.H; +const Range = id.Range; +const Ends = id.Ends; +const hf = id.hf; +const kref = id.kref; + +pub const Kind = enum { norm, norm32, split }; + +const norm_why = "an fp16 output follows the sums as var^-1/2: a bf16 sum's error is a few output ulps"; + +/// A thread's 4 values of the 1,024 in turn, the simdgroup tree, then the tree over the threadgroup's 8 simdgroups. +const sum_depth = 1024 / 256 + 5 + 5; + +fn normDepth(_: c.Shape) usize { + return sum_depth + 3; +} + +pub const Norm = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *id.Lib, + kind: Kind, + mode: ks.Norm, + name: []const u8, + shape: c.Shape = undefined, + stress: bool = false, + x: []f16 = &.{}, + /// w then b, fp32 + wb: []f32 = &.{}, + /// the modulation's 1 + scale, then its shift, fp16 + mod: []f16 = &.{}, + ins: [5]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + /// Rows: positions; n and k: the width, 1,024 (the launcher's). + pub fn spec(name: []const u8, kind: Kind, mode: ks.Norm) c.Spec { + const all: []const c.Shape = &.{ .{ .rows = 12294, .n = H, .k = H }, .{ .rows = 8192, .n = H, .k = H }, .{ .rows = 4101, .n = H, .k = H } }; + return .{ + .name = name, + .accum = .f32, + .depth = normDepth, + .ops = 2 * (id.div_ulps + id.rsqrt_ulps) + 3, + .out = if (kind == .norm32) .f32 else .f16, + .served = switch (kind) { + .norm => switch (mode) { + .modulated => all[0..2], + .affine => all[2..], + // no block calls the plain form: checked at every stream's length + .plain => all, + }, + // the output head's modulated LayerNorm; the affine form at the same rows + .norm32 => all[1..2], + .split => all, + }, + .sums = if (kind == .norm) .{ .bf16 = 0, .why = norm_why } else .{}, + }; + } + + pub fn deinit(t: *Norm) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *Norm) void { + t.a.free(t.x); + t.a.free(t.wb); + t.a.free(t.mod); + t.x = &.{}; + t.wb = &.{}; + t.mod = &.{}; + } + pub fn kernel(t: *Norm, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name, t.kind, t.mode), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Norm); + + fn eps(t: *const Norm) f32 { + return if (t.kind == .norm) cfg.block_eps else cfg.default_eps; + } + + pub fn prepare(t: *Norm, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n != H or shape.k != H) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + t.stress = fill == .stress; + t.x = try t.a.alloc(f16, shape.rows * H); + t.wb = try t.a.alloc(f32, 2 * H); + t.mod = try t.a.alloc(f16, 2 * H); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (t.stress) { + // integers mirrored in each row's halves: a mean of exactly 0 and squares summing to ~2^22, exact in fp32 + for (0..shape.rows) |r| for (0..H / 2) |i| { + const v = id.int(rng, -64, 64); + t.x[r * H + i] = @floatCast(v); + t.x[r * H + H / 2 + i] = @floatCast(-v); + }; + } else id.randomF16(rng, t.x, -3, 3); + id.randomF32(rng, t.wb[0..H], 0.25, 2); + id.randomF32(rng, t.wb[H..], -1, 1); + id.randomF16(rng, t.mod[0..H], 0.5, 1.5); + id.randomF16(rng, t.mod[H..], -1, 1); + if (fill == .nonfinite) { + const at = [_][2]usize{ .{ 0, H / 3 }, .{ 1, H / 2 }, .{ 2, H - 1 }, .{ 3, 0 }, .{ 3, H - 1 } }; + for (at, [_]f16{ id.nan16, id.inf16, -id.inf16, id.inf16, -id.inf16 }) |p, v| if (p[0] < shape.rows) { + t.x[p[0] * H + p[1]] = v; + }; + t.wb[5] = id.nan32; + t.wb[H + 6] = id.inf32; + t.mod[7] = -id.inf16; + t.mod[H + 8] = id.nan16; + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.x)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.wb[0..H])); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.wb[H..])); + t.ins[3] = try t.rig.input(std.mem.sliceAsBytes(t.mod[0..H])); + t.ins[4] = try t.rig.input(std.mem.sliceAsBytes(t.mod[H..])); + t.out = try t.rig.output(t.count() * @as(usize, if (t.kind == .norm32) 4 else 2)); + return true; + } + + pub fn run(t: *Norm) !void { + const k = &t.lib.k; + try t.lib.need(switch (t.kind) { + .norm => "id_norm", + .norm32 => "id_norm32", + .split => "id_norm_split", + }); + const p = t.rig.begin(); + const in = t.ins; + switch (t.kind) { + .norm => k.norm(p.e, kref(in[0]), kref(in[1]), kref(in[2]), kref(in[3]), kref(in[4]), kref(t.out), t.shape.rows, t.eps(), t.mode), + .norm32 => k.norm32(p.e, kref(in[0]), kref(in[1]), kref(in[2]), kref(t.out), t.shape.rows, t.eps(), t.mode), + .split => k.normSplit(p.e, kref(in[0]), kref(in[1]), kref(in[2]), kref(t.out), t.shape.rows, t.eps()), + } + try t.rig.finish(p); + } + + pub fn count(t: *Norm) usize { + return t.shape.rows * H * @as(usize, if (t.kind == .split) 3 else 1); + } + + /// Row r's outputs from the kernel's mean (range `mr`) and sum of squared deviations (range `qr`). + fn row(t: *Norm, ends: *Ends, r: usize, mean: f64, mr: Range, q: f64, qr: Range, want: []f64, mag: []f64) void { + const n: f64 = H; + const e: f64 = t.eps(); + const inv = 1 / @sqrt(q / n + e); + const ir = (Range{ .lo = qr.lo / n, .hi = qr.hi / n }).add(Range.of(e)).fp32(0.5).rsqrt().fp32(id.rsqrt_ulps); + for (0..H) |i| { + const x: f64 = t.x[r * H + i]; + const v = (x - mean) * inv; + const vr = Range.of(x).sub(mr).fp32(0.5).mul(ir).fp32(0.5); + const w: f64 = t.wb[i]; + const b: f64 = t.wb[H + i]; + const s1: f64 = t.mod[i]; + const sh: f64 = t.mod[H + i]; + const at = r * H + i; + switch (t.kind) { + .norm => switch (t.mode) { + .plain => ends.put(want, mag, at, hf(v), vr.fp16()), + .affine => ends.put(want, mag, at, hf(v * w + b), vr.mul(Range.of(w)).fp32(0.5).add(Range.of(b)).fp32(0.5).fp16()), + .modulated => { + const y = vr.fp16().mul(Range.of(s1)).fp32(0.5).fp16().add(Range.of(sh)).fp32(0.5).fp16(); + ends.put(want, mag, at, hf(hf(hf(v) * s1) + sh), y); + }, + }, + .norm32 => switch (t.mode) { + .modulated => ends.put(want, mag, at, v * (1 + w) + b, vr.mul(Range.of(1 + w).fp32(0.5)).fp32(0.5).add(Range.of(b)).fp32(0.5)), + else => ends.put(want, mag, at, v * w + b, vr.mul(Range.of(w)).fp32(0.5).add(Range.of(b)).fp32(0.5)), + }, + .split => { + const y = v * w + b; + const yr = vr.mul(Range.of(w)).fp32(0.5).add(Range.of(b)).fp32(0.5); + const hi = hf(y); + const hr = yr.fp16(); + // hi + lo: lo = fp16(y - hi), y - hi exact in fp32 for whichever fp16 hi the range rounds to + const lr = (Range{ .lo = yr.lo - hr.hi, .hi = yr.hi - hr.lo }).fp16(); + const base = r * 3 * H + i; + ends.put(want, mag, base, hi, hr); + ends.put(want, mag, base + H, hi, hr); + ends.put(want, mag, base + 2 * H, hf(y - hi), lr); + }, + } + } + } + + pub fn reference(t: *Norm, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + const n: f64 = H; + const gs = c.gamma(sum_depth, c.unit(.f32)); + const gq = c.gamma(sum_depth + 3, c.unit(.f32)); + for (0..t.shape.rows) |r| { + const xs = t.x[r * H ..][0..H]; + var s: f64 = 0; + var abs: f64 = 0; + for (xs) |v| { + s += v; + abs += @abs(@as(f64, v)); + } + const mean = s / n; + var q: f64 = 0; + for (xs) |v| q += (v - mean) * (v - mean); + if (t.stress) { + // a mean of 0 exactly, the deviations the integers and their squares' sums exact in fp32 + t.row(&t.ends, r, mean, Range.of(mean), q, Range.of(q), want, mag); + continue; + } + const mr = Range{ .lo = (s - gs * abs) / n, .hi = (s + gs * abs) / n }; + // the squares about the kernel's mean exceed float64's by n (m' - mean)^2, then fp32 rounds them + const dm = @max(@abs(mr.lo - mean), @abs(mr.hi - mean)); + const qr = Range{ .lo = q * (1 - gq), .hi = (q + n * dm * dm) * (1 + gq) }; + t.row(&t.ends, r, mean, mr, q, qr, want, mag); + } + } + pub fn bounds(t: *Norm, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *Norm, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } + + pub fn simulate(t: *Norm, acc: c.Accum, order: g.Order, out: []f64) !void { + const mag = try t.a.alloc(f64, out.len); + defer t.a.free(mag); + var scratch: Ends = .{ .a = t.a }; + defer scratch.deinit(); + try scratch.size(t.count()); + var terms: [H]f64 = undefined; + for (0..t.shape.rows) |r| { + const xs = t.x[r * H ..][0..H]; + for (&terms, xs) |*v, x| v.* = x; + const mean = contract.rows.sumIn(acc, order, &terms) / H; + for (&terms, xs) |*v, x| v.* = (x - mean) * (x - mean); + const q = contract.rows.sumIn(acc, order, &terms); + t.row(&scratch, r, mean, Range.of(mean), q, Range.of(q), out, mag); + } + } +}; diff --git a/zig/tests/contract/metal_idol_rows.zig b/zig/tests/contract/metal_idol_rows.zig new file mode 100644 index 0000000000..63fa5b1043 --- /dev/null +++ b/zig/tests/contract/metal_idol_rows.zig @@ -0,0 +1,442 @@ +//! id_gate_add, id_embed, id_cond and id_bias32 under the contract: elementwise steps at idol-sdk's fp16 points. +const std = @import("std"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const id = @import("metal_idol_common.zig"); + +const c = id.c; +const g = id.g; +const cfg = id.cfg; +const H = id.H; +const Range = id.Range; +const Ends = id.Ends; +const hf = id.hf; +const kref = id.kref; + +/// The steps' common parts: the outputs' ends, and its bounds and intervals from them. +fn Common(comptime T: type) type { + return struct { + pub fn bounds(t: *T, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *T, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } + }; +} + +// --- id_gate_add: x = fp16(x + fp16(g y)), in place ------------------------------------------------------------ + +pub const GateAdd = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *id.Lib, + name: []const u8, + rows: usize = 0, + old: []f16 = &.{}, + y: []f16 = &.{}, + gv: []f16 = &.{}, + ins: [2]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + /// Rows: positions; n: the stream's width, 1,024 (the launcher's). + pub fn spec(name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .none, + .depth = id.noSums, + .ops = 1, + .mids = 1, + .out = .f16, + .in_place = true, + .served = &.{ .{ .rows = 12294, .n = H, .k = 1 }, .{ .rows = 8192, .n = H, .k = 1 }, .{ .rows = 4101, .n = H, .k = 1 } }, + }; + } + + pub fn deinit(t: *GateAdd) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *GateAdd) void { + t.a.free(t.old); + t.a.free(t.y); + t.a.free(t.gv); + t.old = &.{}; + t.y = &.{}; + t.gv = &.{}; + } + pub fn kernel(t: *GateAdd, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = blk: { + var hooks = g.hooks(GateAdd); + hooks.overlap = overlap; + break :blk hooks; + }; + /// The stream is its own output. + fn overlap(_: *anyopaque) bool { + return true; + } + pub const bounds = Common(GateAdd).bounds; + pub const intervals = Common(GateAdd).intervals; + + pub fn prepare(t: *GateAdd, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n != H) return false; + t.rig.reset(); + t.free(); + t.rows = shape.rows; + const n = shape.rows * H; + t.old = try t.a.alloc(f16, n); + t.y = try t.a.alloc(f16, n); + t.gv = try t.a.alloc(f16, H); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + id.randomF16(rng, t.old, -8, 8); + id.randomF16(rng, t.y, -4, 4); + id.randomF16(rng, t.gv, -2, 2); + if (fill == .nonfinite) { + for ([_]f16{ id.nan16, id.inf16, -id.inf16 }, 0..) |v, i| { + t.old[(i * 13) % n] = v; + t.y[(i * 17 + 1) % n] = v; + } + t.gv[5] = id.nan16; + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.y)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.gv)); + t.out = try t.rig.output(n * 2); + return true; + } + + pub fn run(t: *GateAdd) !void { + try t.lib.need("id_gate_add"); + // in place: the stream's rows go back before each independent run + @memcpy(t.rig.bytes(t.out, t.count() * 2), std.mem.sliceAsBytes(t.old)); + const p = t.rig.begin(); + t.lib.k.gateAdd(p.e, kref(t.out), kref(t.ins[0]), kref(t.ins[1]), t.rows); + try t.rig.finish(p); + } + + pub fn count(t: *GateAdd) usize { + return t.rows * H; + } + + pub fn reference(t: *GateAdd, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + for (t.old, t.y, 0..) |xb, yb, i| { + const x: f64 = xb; + const y: f64 = yb; + const gv: f64 = t.gv[i % H]; + const r = Range.of(gv * y).fp32(0.5).fp16().add(Range.of(x)).fp32(0.5).fp16(); + t.ends.put(want, mag, i, hf(x + hf(gv * y)), r); + } + } +}; + +// --- id_embed: x = fp16(fp16(z w^T + b) + pe), z [8192, 16] fp32 -------------------------------------------------- + +/// The 16 inputs' FMA chain. +fn embedDepth(_: c.Shape) usize { + return cfg.latent_channels + 1; +} + +pub const Embed = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *id.Lib, + name: []const u8, + stress: bool = false, + z: []f32 = &.{}, + /// w [1024, 16] then b [1024] + wb: []f32 = &.{}, + pe: []f16 = &.{}, + ins: [4]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + const R = cfg.latent_rows; + const K = cfg.latent_channels; + + /// Rows: the latent's 8,192, which the launcher fixes (a gate's fewer rows run all of them); k: its 16 channels. + pub fn spec(name: []const u8) c.Spec { + return .{ .name = name, .accum = .f32, .depth = embedDepth, .ops = 2, .mids = 1, .out = .f16, .served = &.{.{ .rows = R, .n = H, .k = K }} }; + } + + pub fn deinit(t: *Embed) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *Embed) void { + t.a.free(t.z); + t.a.free(t.wb); + t.a.free(t.pe); + t.z = &.{}; + t.wb = &.{}; + t.pe = &.{}; + } + pub fn kernel(t: *Embed, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Embed); + pub const bounds = Common(Embed).bounds; + pub const intervals = Common(Embed).intervals; + + pub fn prepare(t: *Embed, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n != H or shape.k != K) return false; + t.rig.reset(); + t.free(); + t.stress = fill == .stress; + t.z = try t.a.alloc(f32, R * K); + t.wb = try t.a.alloc(f32, H * K + H); + t.pe = try t.a.alloc(f16, R * H); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (t.stress) { + // products to 2^16 whose halves cancel but for the first: partial sums near 2^19, exact in fp32 + for (0..R) |r| for (0..K / 2) |k| { + t.z[r * K + k] = id.int(rng, 1, 255); + t.z[r * K + K / 2 + k] = t.z[r * K + k] + @as(f32, if (k == 0) 1 else 0); + }; + for (0..H) |o| for (0..K / 2) |k| { + t.wb[o * K + k] = id.sign(rng) * id.int(rng, 1, 255); + t.wb[o * K + K / 2 + k] = -t.wb[o * K + k]; + }; + for (t.wb[H * K ..]) |*v| v.* = id.int(rng, -8, 8); + for (t.pe) |*v| v.* = @floatCast(id.int(rng, -8, 8)); + } else { + id.randomF32(rng, t.z, -3, 3); + id.randomF32(rng, t.wb, -0.5, 0.5); + id.randomF16(rng, t.pe, -1, 1); + } + if (fill == .nonfinite) { + for ([_]f32{ id.nan32, id.inf32, -id.inf32 }, 0..) |v, i| t.z[i * K + 3 * i] = v; + t.wb[H * K + 9] = id.nan32; + t.pe[5 * H + 11] = id.inf16; + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.z)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.wb[0 .. H * K])); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.wb[H * K ..])); + t.ins[3] = try t.rig.input(std.mem.sliceAsBytes(t.pe)); + t.out = try t.rig.output(R * H * 2); + return true; + } + + pub fn run(t: *Embed) !void { + try t.lib.need("id_embed"); + const p = t.rig.begin(); + t.lib.k.embed(p.e, kref(t.ins[0]), kref(t.ins[1]), kref(t.ins[2]), kref(t.ins[3]), kref(t.out)); + try t.rig.finish(p); + } + + pub fn count(_: *Embed) usize { + return R * H; + } + + fn each(t: *Embed, ends: *Ends, want: []f64, mag: []f64, acc: ?c.Accum, order: g.Order) void { + const gd = c.gamma(@floatFromInt(embedDepth(.{ .rows = 1, .n = 1, .k = 1 })), c.unit(.f32)); + var terms: [K]f64 = undefined; + for (0..R) |r| for (0..H) |o| { + var s: f64 = 0; + var abs: f64 = 0; + for (&terms, t.z[r * K ..][0..K], t.wb[o * K ..][0..K]) |*v, z, w| { + v.* = @as(f64, z) * @as(f64, w); + s += v.*; + abs += @abs(v.*); + } + var sr = Range.of(s).widen(gd * abs); + if (acc) |x| { + s = contract.rows.sumIn(x, order, &terms); + sr = Range.of(s); + } else if (t.stress) sr = Range.of(s); + const b: f64 = t.wb[H * K + o]; + const pe: f64 = t.pe[r * H + o]; + const y = sr.add(Range.of(b)).fp32(0.5).fp16().add(Range.of(pe)).fp32(0.5).fp16(); + ends.put(want, mag, r * H + o, hf(hf(s + b) + pe), y); + }; + } + + pub fn reference(t: *Embed, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + t.each(&t.ends, want, mag, null, .sequential); + } + + pub fn simulate(t: *Embed, acc: c.Accum, order: g.Order, out: []f64) !void { + const mag = try t.a.alloc(f64, out.len); + defer t.a.free(mag); + var scratch: Ends = .{ .a = t.a }; + defer scratch.deinit(); + try scratch.size(t.count()); + t.each(&scratch, out, mag, acc, order); + } +}; + +// --- id_cond: y = fp16((a + ba) + (c + bc)), the two condition projections' fp32 sums and biases ----------------- + +pub const Cond = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *id.Lib, + name: []const u8, + rows: usize = 0, + /// a, c [rows, 1024], then ba, bc [1024] + v: []f32 = &.{}, + ins: [4]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + /// Rows: the condition's 4,101; n: its width, 1,024 (the launcher's). + pub fn spec(name: []const u8) c.Spec { + return .{ .name = name, .accum = .none, .depth = id.noSums, .ops = 3, .out = .f16, .served = &.{.{ .rows = cfg.cond_rows, .n = H, .k = 1 }} }; + } + + pub fn deinit(t: *Cond) void { + t.a.free(t.v); + t.v = &.{}; + t.ends.deinit(); + t.rig.deinit(); + } + pub fn kernel(t: *Cond, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Cond); + pub const bounds = Common(Cond).bounds; + pub const intervals = Common(Cond).intervals; + + pub fn prepare(t: *Cond, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n != H) return false; + t.rig.reset(); + t.a.free(t.v); + t.rows = shape.rows; + const n = shape.rows * H; + t.v = try t.a.alloc(f32, 2 * n + 2 * H); + var prng = std.Random.DefaultPrng.init(seed); + id.randomF32(prng.random(), t.v, -8, 8); + if (fill == .nonfinite) { + for ([_]f32{ id.nan32, id.inf32, -id.inf32 }, 0..) |v, i| { + t.v[(i * 13) % n] = v; + t.v[n + (i * 17 + 1) % n] = v; + } + t.v[2 * n + 5] = id.inf32; + t.v[2 * n + H + 6] = id.nan32; + } + for (0..4) |i| { + const part = if (i < 2) t.v[i * n ..][0..n] else t.v[2 * n + (i - 2) * H ..][0..H]; + t.ins[i] = try t.rig.input(std.mem.sliceAsBytes(part)); + } + t.out = try t.rig.output(n * 2); + return true; + } + + pub fn run(t: *Cond) !void { + try t.lib.need("id_cond"); + const p = t.rig.begin(); + t.lib.k.cond(p.e, kref(t.ins[0]), kref(t.ins[2]), kref(t.ins[1]), kref(t.ins[3]), kref(t.out), t.rows); + try t.rig.finish(p); + } + + pub fn count(t: *Cond) usize { + return t.rows * H; + } + + pub fn reference(t: *Cond, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + const n = t.count(); + for (0..n) |i| { + const av: f64 = t.v[i]; + const cv: f64 = t.v[n + i]; + const ba: f64 = t.v[2 * n + i % H]; + const bc: f64 = t.v[2 * n + H + i % H]; + const r = Range.of(av + ba).fp32(0.5).add(Range.of(cv + bc).fp32(0.5)).fp32(0.5).fp16(); + t.ends.put(want, mag, i, hf((av + ba) + (cv + bc)), r); + } + } +}; + +// --- id_bias32: y += b by column, fp32, in place -------------------------------------------------------------- + +pub const Bias32 = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *id.Lib, + name: []const u8, + shape: c.Shape = undefined, + old: []f32 = &.{}, + b: []f32 = &.{}, + ins: [1]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + /// Rows: the latent's; n: the velocity's 16 channels. + pub fn spec(name: []const u8) c.Spec { + return .{ .name = name, .accum = .none, .depth = id.noSums, .ops = 1, .out = .f32, .in_place = true, .served = &.{.{ .rows = cfg.latent_rows, .n = cfg.latent_channels, .k = 1 }} }; + } + + pub fn deinit(t: *Bias32) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *Bias32) void { + t.a.free(t.old); + t.a.free(t.b); + t.old = &.{}; + t.b = &.{}; + } + pub fn kernel(t: *Bias32, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = blk: { + var hooks = g.hooks(Bias32); + hooks.overlap = overlap; + break :blk hooks; + }; + /// The sums are their own output. + fn overlap(_: *anyopaque) bool { + return true; + } + pub const bounds = Common(Bias32).bounds; + pub const intervals = Common(Bias32).intervals; + + pub fn prepare(t: *Bias32, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + const n = shape.rows * shape.n; + t.old = try t.a.alloc(f32, n); + t.b = try t.a.alloc(f32, shape.n); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + id.randomF32(rng, t.old, -16, 16); + id.randomF32(rng, t.b, -2, 2); + if (fill == .nonfinite) { + for ([_]f32{ id.nan32, id.inf32, -id.inf32 }, 0..) |v, i| t.old[(i * 13) % n] = v; + if (shape.n > 6) t.b[5] = id.inf32; + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.b)); + t.out = try t.rig.output(n * 4); + return true; + } + + pub fn run(t: *Bias32) !void { + try t.lib.need("id_bias32"); + @memcpy(t.rig.bytes(t.out, t.count() * 4), std.mem.sliceAsBytes(t.old)); + const p = t.rig.begin(); + t.lib.k.bias32(p.e, kref(t.out), kref(t.ins[0]), t.shape.rows, t.shape.n); + try t.rig.finish(p); + } + + pub fn count(t: *Bias32) usize { + return t.shape.rows * t.shape.n; + } + + pub fn reference(t: *Bias32, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + for (t.old, 0..) |y, i| { + const b: f64 = t.b[i % t.shape.n]; + t.ends.put(want, mag, i, y + b, Range.of(y + b).fp32(0.5)); + } + } +}; diff --git a/zig/tests/contract/metal_idol_steps.zig b/zig/tests/contract/metal_idol_steps.zig new file mode 100644 index 0000000000..617729535b --- /dev/null +++ b/zig/tests/contract/metal_idol_steps.zig @@ -0,0 +1,426 @@ +//! id_repo and id_qk under the contract: RePo's angles in fp32, and the heads rotated, normed and rounded to fp16. +const std = @import("std"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const id = @import("metal_idol_common.zig"); + +const c = id.c; +const g = id.g; +const H = id.H; +const D = id.D; +const Range = id.Range; +const Ends = id.Ends; +const hf = id.hf; +const kref = id.kref; + +const streams: []const c.Shape = &.{ .{ .rows = 12294, .n = 512, .k = 128 }, .{ .rows = 8192, .n = 512, .k = 128 }, .{ .rows = 4101, .n = 512, .k = 128 } }; + +// --- id_repo: u = silu(g) c, d = u final^T (48), angles x f pi as cos and sin [rows, 16, 32] -------------------- + +const pi32: f64 = @as(f32, std.math.pi); + +/// d's FMA chain over 128. +fn repoDepth(_: c.Shape) usize { + return 128 + 1; +} + +/// A pair's axis: 10 x, 10 y, then 12 z. +fn axis(j: usize) usize { + return if (j < 10) 0 else if (j < 20) 1 else 2; +} + +pub const Repo = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *id.Lib, + name: []const u8, + rows: usize = 0, + stress: bool = false, + /// gate sums then content sums, 128 each a row + gc: []f32 = &.{}, + /// the final map [48, 128] + wf: []f32 = &.{}, + freqs: []f32 = &.{}, + inv: f32 = 0.5, + ins: [3]rig_.Ref = undefined, + outs: [2]rig_.Ref = undefined, + ends: Ends, + + /// Rows: positions; n: 16 heads x 32 pairs; k: RePo's hidden width. + pub fn spec(name: []const u8) c.Spec { + return .{ .name = name, .accum = .f32, .depth = repoDepth, .ops = id.exp_ulps + id.div_ulps + 2, .out = .f32, .served = streams }; + } + + pub fn deinit(t: *Repo) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *Repo) void { + t.a.free(t.gc); + t.a.free(t.wf); + t.a.free(t.freqs); + t.gc = &.{}; + t.wf = &.{}; + t.freqs = &.{}; + } + pub fn kernel(t: *Repo, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Repo); + + pub fn prepare(t: *Repo, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n != 512 or shape.k != 128) return false; + t.rig.reset(); + t.free(); + t.rows = shape.rows; + t.stress = fill == .stress; + t.gc = try t.a.alloc(f32, shape.rows * 256); + t.wf = try t.a.alloc(f32, 48 * 128); + t.freqs = try t.a.alloc(f32, 32); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (t.stress) { + // gates of 20 to 40 (sigmoid 1 in fp32), integer contents: u's halves differ by a gate twice, d small + for (0..shape.rows) |r| for (0..64) |i| { + const gr = t.gc[r * 256 ..][0..256]; + gr[i] = id.int(rng, 20, 40) / t.inv; + gr[i + 64] = gr[i]; + gr[128 + i] = id.sign(rng) * id.int(rng, 1, 8) / t.inv; + gr[128 + i + 64] = gr[128 + i] + @as(f32, if (i < 2) 1 / t.inv else 0); + }; + for (0..48) |o| for (0..64) |i| { + t.wf[o * 128 + i] = id.sign(rng) * id.int(rng, 1, 32); + t.wf[o * 128 + 64 + i] = -t.wf[o * 128 + i]; + }; + id.randomF32(rng, t.freqs, 0.001, 0.01); + } else { + id.randomF32(rng, t.gc, -4, 4); + id.randomF32(rng, t.wf, -0.2, 0.2); + id.randomF32(rng, t.freqs, 0.05, 1.5); + } + if (fill == .nonfinite) { + const at = [_]usize{ 0 * 256 + 5, 1 * 256 + 7, 2 * 256 + 128 + 3 }; + for (at, [_]f32{ id.nan32, id.inf32, -id.inf32 }) |p, v| if (p < t.gc.len) { + t.gc[p] = v; + }; + t.freqs[4] = id.nan32; + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.gc)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.wf)); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.freqs)); + for (&t.outs) |*o| o.* = try t.rig.output(shape.rows * 512 * 4); + return true; + } + + pub fn run(t: *Repo) !void { + try t.lib.need("id_repo"); + const p = t.rig.begin(); + t.lib.k.repo(p.e, kref(t.ins[0]), t.inv, kref(t.ins[1]), kref(t.ins[2]), kref(t.outs[0]), kref(t.outs[1]), t.rows); + try t.rig.finish(p); + } + + pub fn count(t: *Repo) usize { + return 2 * t.rows * 512; + } + + /// Row r's u values in float64 and their ranges through the kernel's exp, add, division and product. + fn us(t: *const Repo, r: usize, u: *[128]f64, ur: *[128]Range) void { + const inv: f64 = t.inv; + for (u, ur, 0..) |*v, *vr, i| { + const a = @as(f64, t.gc[r * 256 + i]) * inv; + const cc = @as(f64, t.gc[r * 256 + 128 + i]) * inv; + v.* = a / (1 + @exp(-a)) * cc; + // a >= 20: 1 + exp(-a) is 1 in fp32 and a / 1 is a, so u = a c exactly + if (t.stress) { + v.* = a * cc; + vr.* = Range.of(v.*); + continue; + } + const den = Range.of(@exp(-a)).fp32(id.exp_ulps).add(Range.of(1)).fp32(0.5); + vr.* = Range.of(a).mul(den.recip()).fp32(id.div_ulps).mul(Range.of(cc)).fp32(0.5); + } + } + + /// Row r's outputs from its 48 offsets d (ranges dr). + fn row(t: *Repo, ends: *Ends, r: usize, d: *const [48]f64, dr: *const [48]Range, want: []f64, mag: []f64) void { + for (0..512) |i| { + const h = i / 32; + const j = i % 32; + const x = d[h * 3 + axis(j)]; + const xr = dr[h * 3 + axis(j)]; + const f: f64 = t.freqs[j]; + const th = std.math.tanh(f); + const ang = (x * th + x * (f - th)) * pi32; + // tanh's error cancels between the two products: fp32 rounds each product, the sum and the scale + const xm = @max(@abs(xr.lo), @abs(xr.hi)); + const e = 2 * 0x1p-24 * pi32 * xm * (@abs(th) + 2 * @abs(f - th) + 2 * @abs(f)) + 0x1p-140; + const lo = @min(xr.lo * f, xr.hi * f) * pi32 - e; + const hi = @max(xr.lo * f, xr.hi * f) * pi32 + e; + ends.put(want, mag, r * 512 + i, @cos(ang), id.cosOver(lo, hi).fp32(id.trig_ulps)); + ends.put(want, mag, (t.rows + r) * 512 + i, @sin(ang), id.sinOver(lo, hi).fp32(id.trig_ulps)); + } + } + + pub fn reference(t: *Repo, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + const gd = c.gamma(@floatFromInt(repoDepth(.{ .rows = 1, .n = 1, .k = 1 })), c.unit(.f32)); + var u: [128]f64 = undefined; + var ur: [128]Range = undefined; + var d: [48]f64 = undefined; + var dr: [48]Range = undefined; + for (0..t.rows) |r| { + t.us(r, &u, &ur); + for (&d, &dr, 0..) |*dv, *dv_r, o| { + var s: f64 = 0; + var rr = Range.of(0); + var abs: f64 = 0; + for (u, ur, t.wf[o * 128 ..][0..128]) |uv, uvr, wv| { + s += uv * wv; + rr = rr.add(uvr.mul(Range.of(wv))); + abs += @max(@abs(uvr.lo), @abs(uvr.hi)) * @abs(@as(f64, wv)); + } + dv.* = s; + dv_r.* = if (t.stress) rr else rr.widen(gd * abs); + } + t.row(&t.ends, r, &d, &dr, want, mag); + } + } + pub fn bounds(t: *Repo, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *Repo, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } + + pub fn simulate(t: *Repo, acc: c.Accum, order: g.Order, out: []f64) !void { + const mag = try t.a.alloc(f64, out.len); + defer t.a.free(mag); + var scratch: Ends = .{ .a = t.a }; + defer scratch.deinit(); + try scratch.size(t.count()); + var u: [128]f64 = undefined; + var ur: [128]Range = undefined; + var d: [48]f64 = undefined; + var dr: [48]Range = undefined; + var terms: [128]f64 = undefined; + for (0..t.rows) |r| { + t.us(r, &u, &ur); + for (&d, &dr, 0..) |*dv, *dv_r, o| { + for (&terms, u, t.wf[o * 128 ..][0..128]) |*v, uv, wv| v.* = uv * wv; + dv.* = contract.rows.sumIn(acc, order, &terms); + dv_r.* = Range.of(dv.*); + } + t.row(&scratch, r, &d, &dr, out, mag); + } + } +}; + +// --- id_qk: (q, k) pairs turned by cos and sin and rounded, F.normalize, times gamma and 8, rounded; v copied ----- + +const qk_why = "an fp16 output follows the head's sum of squares as its square root: a bf16 sum is off a few ulps"; + +const qk_known = "a NaN in a head gives its other outputs +-inf where F.normalize gives NaN: max(sqrt(ss), 1e-12f) drops it"; + +/// A lane's two squares, then the simdgroup tree. +fn qkDepth(_: c.Shape) usize { + return 1 + 5; +} + +pub const Qk = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *id.Lib, + name: []const u8, + rows: usize = 0, + stress: bool = false, + qkv: []f16 = &.{}, + /// gq then gk [1024] each + gam: []f32 = &.{}, + /// cos then sin [rows, 512] each + rot: []f32 = &.{}, + ins: [5]rig_.Ref = undefined, + outs: [3]rig_.Ref = undefined, + ends: Ends, + + /// Rows: positions; n: heads; k: the head width. + pub fn spec(name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = qkDepth, + .ops = id.sqrt_ulps + id.div_ulps + 2, + .out = .f16, + .served = &.{ .{ .rows = 12294, .n = 16, .k = D }, .{ .rows = 8192, .n = 16, .k = D }, .{ .rows = 4101, .n = 16, .k = D } }, + .sums = .{ .bf16 = 0, .why = qk_why }, + .known = .{ .gate = .nonfinite, .finding = qk_known }, + }; + } + + pub fn deinit(t: *Qk) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *Qk) void { + t.a.free(t.qkv); + t.a.free(t.gam); + t.a.free(t.rot); + t.qkv = &.{}; + t.gam = &.{}; + t.rot = &.{}; + } + pub fn kernel(t: *Qk, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Qk); + + pub fn prepare(t: *Qk, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n != 16 or shape.k != D) return false; + t.rig.reset(); + t.free(); + t.rows = shape.rows; + t.stress = fill == .stress; + const n = shape.rows; + t.qkv = try t.a.alloc(f16, n * 3 * H); + t.gam = try t.a.alloc(f32, 2 * H); + t.rot = try t.a.alloc(f32, 2 * n * 512); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + id.randomF16(rng, t.qkv, -4, 4); + id.randomF32(rng, t.gam, 0.25, 2); + for (0..n * 512) |i| { + if (t.stress) { + // quarter turns: exact rotations of integer heads, whose squares sum to ~2^20 exactly + const turns = [4][2]f32{ .{ 1, 0 }, .{ 0, 1 }, .{ -1, 0 }, .{ 0, -1 } }; + const turn = turns[rng.uintLessThan(usize, 4)]; + t.rot[i] = turn[0]; + t.rot[n * 512 + i] = turn[1]; + continue; + } + const angle = rng.float(f32) * 2 * std.math.pi; + t.rot[i] = @cos(angle); + t.rot[n * 512 + i] = @sin(angle); + } + if (t.stress) for (0..n) |r| for (0..2 * H) |i| { + t.qkv[r * 3 * H + i] = @floatCast(id.int(rng, -128, 128)); + }; + if (fill == .nonfinite) { + // q NaN, k +inf, v -inf, in turn; then a NaN cos + const at = [_]usize{ 0 * 3 * H + 0 * D + 5, 1 * 3 * H + H + 2 * D + 9, 2 * 3 * H + 2 * H + 3 * D + 1 }; + for (at, [_]f16{ id.nan16, id.inf16, -id.inf16 }) |p, v| if (p < t.qkv.len) { + t.qkv[p] = v; + }; + if (n > 3) t.rot[3 * 512 + 4 * 32 + 6] = id.nan32; + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.qkv)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.gam[0..H])); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.gam[H..])); + t.ins[3] = try t.rig.input(std.mem.sliceAsBytes(t.rot[0 .. n * 512])); + t.ins[4] = try t.rig.input(std.mem.sliceAsBytes(t.rot[n * 512 ..])); + for (&t.outs) |*o| o.* = try t.rig.output(n * H * 2); + return true; + } + + pub fn run(t: *Qk) !void { + try t.lib.need("id_qk"); + const p = t.rig.begin(); + const in = t.ins; + t.lib.k.heads(p.e, kref(in[0]), kref(in[1]), kref(in[2]), kref(in[3]), kref(in[4]), kref(t.outs[0]), kref(t.outs[1]), kref(t.outs[2]), t.rows); + try t.rig.finish(p); + } + + pub fn count(t: *Qk) usize { + return 3 * t.rows * H; + } + + /// Head h of row r of q (0) or k (1): its pairs turned and rounded, as float64 values and ranges. + fn turned(t: *const Qk, which: usize, r: usize, h: usize, v: *[D]f64, vr: *[D]Range) void { + const x = t.qkv[r * 3 * H + which * H + h * D ..][0..D]; + for (0..D / 2) |p| { + const cs: f64 = t.rot[r * 512 + h * 32 + p]; + const sn: f64 = t.rot[(t.rows + r) * 512 + h * 32 + p]; + const a: f64 = x[2 * p]; + const b: f64 = x[2 * p + 1]; + v[2 * p] = hf(a * cs - b * sn); + v[2 * p + 1] = hf(a * sn + b * cs); + const ac = Range.of(a * cs).fp32(0.5); + const bs = Range.of(b * sn).fp32(0.5); + const as_ = Range.of(a * sn).fp32(0.5); + const bc = Range.of(b * cs).fp32(0.5); + vr[2 * p] = ac.sub(bs).fp32(0.5).fp16(); + vr[2 * p + 1] = as_.add(bc).fp32(0.5).fp16(); + } + } + + /// The head's outputs from its sum of squares `ss` (range `ssr`). + fn head(t: *Qk, ends: *Ends, which: usize, r: usize, h: usize, v: *const [D]f64, vr: *const [D]Range, ss: f64, ssr: Range, want: []f64, mag: []f64) void { + const e12: f64 = @as(f32, 1e-12); + // torch's clamp_min keeps a NaN norm (@max would drop it) + const norm = if (std.math.isNan(ss)) ss else @max(@sqrt(ss), e12); + const nr0 = ssr.up(id.sqrt).fp32(id.sqrt_ulps); + const nr = Range{ .lo = @max(nr0.lo, e12), .hi = @max(nr0.hi, e12) }; + const gam = t.gam[which * H + h * D ..][0..D]; + const base = which * t.rows * H + (h * t.rows + r) * D; + for (0..D) |i| { + const gv: f64 = gam[i]; + const w = hf(v[i] / norm * gv * 8); + const rr = vr[i].mul(nr.recip()).fp32(id.div_ulps).mul(Range.of(gv)).fp32(0.5).mul(Range.of(8)).fp32(0.5).fp16(); + ends.put(want, mag, base + i, w, rr); + } + } + + /// v's heads, copied bit for bit. + fn values(t: *Qk, ends: *Ends, want: []f64, mag: []f64) void { + for (0..t.rows) |r| for (0..16) |h| for (0..D) |i| { + const f: f64 = t.qkv[r * 3 * H + 2 * H + h * D + i]; + ends.put(want, mag, 2 * t.rows * H + (h * t.rows + r) * D + i, f, Range.of(f)); + }; + } + + pub fn reference(t: *Qk, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + const gd = c.gamma(@floatFromInt(qkDepth(.{ .rows = 1, .n = 1, .k = 1 })), c.unit(.f32)); + var v: [D]f64 = undefined; + var vr: [D]Range = undefined; + for (0..2) |which| for (0..t.rows) |r| for (0..16) |h| { + t.turned(which, r, h, &v, &vr); + var ss: f64 = 0; + var lo: f64 = 0; + var hi: f64 = 0; + for (v, vr) |x, xr| { + ss += x * x; + const q = xr.square(); + lo += q.lo; + hi += q.hi; + } + const ssr = if (t.stress) Range.of(ss) else Range{ .lo = lo * (1 - gd), .hi = hi * (1 + gd) }; + t.head(&t.ends, which, r, h, &v, &vr, ss, ssr, want, mag); + }; + t.values(&t.ends, want, mag); + } + pub fn bounds(t: *Qk, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *Qk, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } + + pub fn simulate(t: *Qk, acc: c.Accum, order: g.Order, out: []f64) !void { + const mag = try t.a.alloc(f64, out.len); + defer t.a.free(mag); + var scratch: Ends = .{ .a = t.a }; + defer scratch.deinit(); + try scratch.size(t.count()); + var v: [D]f64 = undefined; + var vr: [D]Range = undefined; + var terms: [D]f64 = undefined; + for (0..2) |which| for (0..t.rows) |r| for (0..16) |h| { + t.turned(which, r, h, &v, &vr); + for (&terms, v) |*s, x| s.* = x * x; + const ss = contract.rows.sumIn(acc, order, &terms); + t.head(&scratch, which, r, h, &v, &vr, ss, Range.of(ss), out, mag); + }; + t.values(&scratch, out, mag); + } +}; diff --git a/zig/tests/contract/metal_limits.zig b/zig/tests/contract/metal_limits.zig index 796ed74d6b..ece8073ecc 100644 --- a/zig/tests/contract/metal_limits.zig +++ b/zig/tests/contract/metal_limits.zig @@ -102,6 +102,9 @@ fn launchOf(a: std.mem.Allocator, device: mtl.Device, d: limits.MetalLaunch) !La } else if (std.mem.startsWith(u8, d.name, "moss image ")) { text = sources.moss_image; frags = true; + } else if (std.mem.startsWith(u8, d.name, "idol ")) { + text = sources.idol; + frags = true; } else if (std.mem.startsWith(u8, d.name, "core deltanet ")) { text = sources.core_deltanet; } else if (std.mem.startsWith(u8, d.name, "qwen27 ")) { diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index b32e2baa9e..9799191182 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -427,6 +427,8 @@ pub fn main(init: std.process.Init) !u8 { defer native_gdn.destroy(); const moss_image = try @import("metal_moss_image.zig").Suite.create(a, device, queue, &kernels, o.only); defer moss_image.destroy(); + const idol = try @import("metal_idol.zig").Suite.create(a, device, queue, &kernels, o.only); + defer idol.destroy(); var t: g.Tally = .{}; const t0 = std.Io.Clock.awake.now(init.io); if (o.gates) { From e99604cca4ac2615c04dfd70fb19c44df80b0464 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 02:48:17 +0800 Subject: [PATCH 3/9] idol: id_qk keeps a NaN norm as F.normalize does; the recipe and changelog line --- CHANGELOG.md | 2 + docs/kernel-contract.md | 2 - docs/recipes/idol.md | 129 ++++++++++++++++++++++++ zig/kernels/metal/idol.metal | 6 +- zig/tests/contract/metal_idol_steps.zig | 3 - 5 files changed, 135 insertions(+), 7 deletions(-) create mode 100644 docs/recipes/idol.md diff --git a/CHANGELOG.md b/CHANGELOG.md index 5f676abf2a..cad527edba 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Idol (speridlabs/idol, image to 3D Gaussian splat), first step: its 0.37B denoiser runs natively on M5 Metal from the released checkpoint, the condition refiner once an image and all 50 guided steps on the GPU. On an M5 Ultra a forward takes 0.273 s against the publisher's SDK on torch MPS at 0.939 s, with a 1.84 GB peak footprint. Each step's velocities on the SDK's own input sit as close to an fp32 run as the SDK's own, and each joint block alone passes the same gate; `tf-idol` runs those gates from a case `tools/zig/idol_case.py` records. RePo's fp32 maps run as an fp16 hi/lo split on the tensor units, whose fp32 GEMM is about fp16-accurate. The image encoders and the decoders are not ported yet (fixes #1757). Thanks to @jayleaton. + - DeepSeek-V4.x: `tool_choice: "required"` and named tool calls are now constrained to the model's own DSML call markup, so those requests are held to a valid call on DeepSeek too. - Qwen3.6-35B-A3B on CUDA: `--drafter` also takes a DFlash2 draft model, chosen by the `dflash_config` in its `config.json`; the checkpoint's own directory still drafts with its MTP layer. With `incoai/Qwen3.6-35B-A3B-DFlash2` over `nvidia/Qwen3.6-35B-A3B-NVFP4`, one GB10 decodes 130 tok/s at one stream on a 4.2K-token prompt (MTP: 99), and every drafted reply equals its plain run. diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 472881c3b5..b2b465117a 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -117,8 +117,6 @@ MLX's formula keeps the infinity, because its codes enter the MMAs as q + 128 an `tf_glm_scores` scores a NaN dot as 0 where float64 keeps the NaN: its relu is `fmaxf(x, 0)`, as `tl.maximum` is in the Python family's Triton glue it was ported from. Nemotron's Metal `tf_xup_w3` gives 0 for a NaN dot where float64 keeps the NaN: its relu2 is `metal::max(bfloat(acc), 0.0f)`; the other up pipelines give its bytes on the same inputs. -Idol's `id_qk` gives a head holding a NaN +-inf in its other outputs where torch's `F.normalize` gives NaN: its -norm is `max(precise::sqrt(ss), 1e-12f)`, and Metal's `max` drops the NaN that `clamp_min` keeps. HIP’s MLX stream tile keeps its finite dot2 sums unchanged. A non-finite result with an actual non-finite operand is recomputed by full quantization groups using the factored float64 formula. The plain tile is diff --git a/docs/recipes/idol.md b/docs/recipes/idol.md new file mode 100644 index 0000000000..27995b8120 --- /dev/null +++ b/docs/recipes/idol.md @@ -0,0 +1,129 @@ +# Idol's denoiser on Metal + +Idol (speridlabs) turns one image into a 3D Gaussian splat. A 0.37B flow transformer denoises 8,192 latent rows of 16 +channels, conditioned on the image's DINOv3 and FLUX.2 VAE rows; two decoders then grow an octree from that latent and +emit the Gaussians. `zig/src/families/idol` runs the denoiser on the Metal runtime from the released checkpoint: + +- once an image: the condition embedding (DINOv3's 1,280 channels and the VAE's 128, each projected to 1,024 and + summed) and the two condition-refiner blocks (no modulation, affine LayerNorms); +- at each step: the two noise-refiner blocks over the latent rows, then the 24 joint blocks over 12,294 rows (8,192 + latent, 4,101 condition, one camera token), then the output head and the camera velocity. + +Every block is 16 heads of 64 at width 1,024 with a tanh-GELU MLP of 4,096. Its rotary is RePo: a learned map from +each row's LayerNorm to per-head 3D positions, angles from 10, 10 and 12 frequencies per axis. + +The sampler's defaults: 50 Euler steps, `t = 3u / (1 + 2u)`, guidance `v = 4 v_cond - 3 v_uncond` (two passes a step: +with the condition, and with zeros). + +The image encoders, the decoders and the PLY writer are not ported yet. `tools/zig/idol_case.py` runs the publisher's +SDK in torch, a development reference. It is a denoiser, not a token family: it does not register with the lane +engine, the server or `families.detect`. + +It runs on the M5's tensor units alone, the only chip it was checked on: `Model.load` refuses a GPU without them. + +## Checkpoint + +`speridlabs/idol` at `2731c2291fc65a797c5cb627a5e5526fc1490902`, Apache 2.0, ungated. `config.json` must name +`LatentSeqMMFlowModel` at the released geometry, and `denoiser.safetensors` must hold its 560 tensors; anything else is +refused. The reference is idol-sdk at `4702a7f00850e34b6e83a2170cc123c3e9570744` (Apache 2.0). + +## Precision + +The SDK loads the checkpoint's fp32 weights and converts the blocks' linear layers to fp16; everything else stays fp32. +This family keeps that split: + +- fp16 block projections with fp32 sums; the bias joins the sums before one fp16 rounding, GELU after it. +- LayerNorms in fp32, rounded to fp16; the modulation's shift and scale in fp16 ops, as torch's fp16 tensors run them. +- RePo in fp32: its LayerNorm, both 1,024 to 128 maps, the 128 to 48 map and the angles. +- Attention: fp16 q, k and v, fp32 scores and sums. q and k are rotated, then L2-normalized and scaled by the + checkpoint's gamma times 8. +- The condition embedding and the output projection in fp32; the time MLP, modulation and camera MLPs on the host. + +### The tensor units' fp32 matmul is not fp32 + +On the M5, a GEMM whose inputs are fp32 on the tensor units (`tfp::gemm_nax`, the Qwen-VL tower's +`qwv_gemm_f32`, MLX's fp32 matmul) rounds its inputs to about fp16: its error against float64 is about 7.5e-4 of the +result, against about 1e-7 for an fp32 FMA. RePo's maps at that precision move block 8's attention out of its gate. + +`id_norm_split` therefore writes each RePo input row as fp16 `[hi | hi | lo]` (`lo = x - hi`), and the map's weights are +stored `[hi | lo | hi]`, scaled by a power of two so `lo` stays normal. One fp16 split-K GEMM with K = 3,072 then sums +`hi*hi + hi*lo + lo*hi` in fp32; `id_repo` undoes the scale. With it every block passes its gate. `id_gemm32`, +an fp32 FMA GEMM on simdgroup matrices, is exact but takes 0.47 s a forward if it runs RePo; it serves the condition +embedding (once an image) and the output projection. + +## Kernel contract + +`zig/tests/contract/metal_idol*.zig` registers every kernel of `idol.metal` and `idol_attention.metal`: + +- `id_linear`'s eight forms (plain and GELU, each tile tail) and the split-K halves `id_split_0` and `id_split_1`; +- `id_norm` (plain, affine, modulated), `id_norm32` (affine, modulated) and `id_norm_split`; +- `id_repo`, `id_qk`, `id_gate_add`, `id_embed`, `id_cond`, `id_bias32` and `id_gemm32`; +- the attention's four forms. + +Each output must lie in the range its fp32 steps allow at the SDK's fp16 rounding points. Nineteen controls drop or +add one rounding point or sum in fp16, and each must go red at its gate. One runs the tensor units' fp32 GEMM in +`id_gemm32`'s place: gate 4 misses it by 13 million times the exact-sum bound. `id_qk` keeps a NaN norm as +`F.normalize` does. Forms no engine call reaches (`id_linear_00` and `_10`, attention `01` and `10`, plain `id_norm`, +affine `id_norm32`) are checked at made-up shapes. The tensor-unit kernels do not build for M1 or M2 offline, as +MOSS-Image's do not; the rest build at their reserved thread counts. + +`tf-kernel-contract --only "metal idol"` on the M5 Ultra: + +``` +SUMMARY 27 kernels: 27 pass gates 1-4, 0 fail, 0 known findings; 19 controls: 19 red as planted, 0 missed; 0 limit problems +``` + +## Run a case + +`idol_case.py` imports the SDK's `dit.py`, `layers.py` and `sampler.py` only when each file's sha256 is the pinned +commit's, and the snapshot only when it is the pinned revision. + +```sh +hf download speridlabs/idol --revision 2731c2291fc65a797c5cb627a5e5526fc1490902 --local-dir WEIGHTS +git clone https://github.com/speridlabs/idol-sdk SDK && git -C SDK checkout 4702a7f00850e34b6e83a2170cc123c3e9570744 +python3 tools/zig/idol_case.py case WEIGHTS SDK case/ --features dino.safetensors vae.safetensors +python3 tools/zig/idol_case.py exact WEIGHTS SDK case/ # the same steps with an fp32 denoiser +python3 tools/zig/idol_case.py blocks WEIGHTS SDK case/ --step 0 # each block's input and output, fp16 and fp32 +zig build tf-idol -Dcpu=apple_m1 +zig-out/bin/tf-idol WEIGHTS case/ +zig-out/bin/tf-idol WEIGHTS case/ blocks 0 +``` + +The features are the SDK's own `feature1` [1, 4101, 1280] and `feature2` [1, 4101, 128] for an image. The case holds the +seeded noise (torch's CPU generator), the schedule, the SDK's input, both velocities and camera velocities at every +step, and its final sample. `exact` adds fp32 velocities on the same inputs; `blocks` adds every block's input and +output at one step, and each joint block run alone in fp32 on the SDK's fp16 input. + +`tf-idol` gates: + +- **Each step**, both branches on the SDK's own input: cosine against the fp32 velocity within 2x the SDK's own + distance (`1 - cos`). Worst row over the run within 2x the SDK's worst row. Fifty fp16 steps make single rows chaotic, + so the per-row gate runs per block instead: +- **Each block** alone on the SDK's fp16 input, against that block in fp32: worst row and cosine within 2x the SDK's. +- **Repeat**: the same step twice, another between, byte for byte. + +It then runs the whole guided schedule from the noise, times it, and writes `sample.f32`. `reference.zig`'s tests plant +a 5% error in one row and require the gate to fail it. + +## Measured + +Setup: + +- Mac Studio M5 Ultra 256 GB, macOS 27; this commit's `tf-idol`. +- The SDK on torch 2.13.0 (MPS, fp16 blocks, `ATTN_BACKEND=sdpa`). +- The SDK's avocado example: its own DINOv3 and VAE features, seed 42, 50 steps, guidance 3. + +Results: + +- **Velocity**, both branches of all 50 steps on the SDK's own input (100 of 100 pass): + - against the SDK: cosine 0.999997 or better; + - against fp32: our worst row over the run 0.0456, the SDK's 0.0398; per step our `1 - cos` is the SDK's within + the 2x bound at every step, and the smaller at 40 of the 100. +- **Blocks**, at steps 1 and 41, each of the 23 joint blocks after the first on the SDK's input: all pass, worst rows + 0.0007 to 0.0017 against fp32, where the SDK's are 0.0007 to 0.0025. +- **Sample:** our 50 guided steps from the same noise end at cosine 0.99999 to the SDK's sample. +- **Repeat:** identical bytes. +- **Speed**, one session: 0.273 s a forward (27.3 s for 50 guided steps), 3.4x the SDK's 0.939 s (the median of its + 100). Load 0.2 to 1.6 s; the condition, once an image, 0.07 s. +- **Memory:** 1.84 GB peak process footprint (`/usr/bin/time -l`, the whole `tf-idol` run with its case), the + denoiser's fp16 weights in shared Metal buffers. diff --git a/zig/kernels/metal/idol.metal b/zig/kernels/metal/idol.metal index 1a8048d1b2..f687b50e5a 100644 --- a/zig/kernels/metal/idol.metal +++ b/zig/kernels/metal/idol.metal @@ -281,8 +281,10 @@ struct IdHeads { const float k0 = float(qkv[from + 1024]), k1 = float(qkv[from + 1025]); const float rq0 = id_h(q0 * c - q1 * s), rq1 = id_h(q0 * s + q1 * c); const float rk0 = id_h(k0 * c - k1 * s), rk1 = id_h(k0 * s + k1 * c); - const float qn = max(metal::precise::sqrt(simd_sum(rq0 * rq0 + rq1 * rq1)), 1e-12f); - const float kn = max(metal::precise::sqrt(simd_sum(rk0 * rk0 + rk1 * rk1)), 1e-12f); + // F.normalize's clamp_min keeps a NaN norm, where Metal's max would drop it + const float sq = metal::precise::sqrt(simd_sum(rq0 * rq0 + rq1 * rq1)); + const float sk = metal::precise::sqrt(simd_sum(rk0 * rk0 + rk1 * rk1)); + const float qn = isnan(sq) ? sq : max(sq, 1e-12f), kn = isnan(sk) ? sk : max(sk, 1e-12f); const uint g = h * 64 + lane * 2; const ulong to = (ulong(h) * p.rows + r) * 64 + lane * 2; Q[to] = half(rq0 / qn * gq[g] * p.scale); diff --git a/zig/tests/contract/metal_idol_steps.zig b/zig/tests/contract/metal_idol_steps.zig index 617729535b..0202b284b2 100644 --- a/zig/tests/contract/metal_idol_steps.zig +++ b/zig/tests/contract/metal_idol_steps.zig @@ -220,8 +220,6 @@ pub const Repo = struct { const qk_why = "an fp16 output follows the head's sum of squares as its square root: a bf16 sum is off a few ulps"; -const qk_known = "a NaN in a head gives its other outputs +-inf where F.normalize gives NaN: max(sqrt(ss), 1e-12f) drops it"; - /// A lane's two squares, then the simdgroup tree. fn qkDepth(_: c.Shape) usize { return 1 + 5; @@ -253,7 +251,6 @@ pub const Qk = struct { .out = .f16, .served = &.{ .{ .rows = 12294, .n = 16, .k = D }, .{ .rows = 8192, .n = 16, .k = D }, .{ .rows = 4101, .n = 16, .k = D } }, .sums = .{ .bf16 = 0, .why = qk_why }, - .known = .{ .gate = .nonfinite, .finding = qk_known }, }; } From 3114e80268272a005bc925f2bedd53d332b658fe Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 02:54:43 +0800 Subject: [PATCH 4/9] changelog: #1766 for Idol's denoiser --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cad527edba..38fd52543d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased -- Idol (speridlabs/idol, image to 3D Gaussian splat), first step: its 0.37B denoiser runs natively on M5 Metal from the released checkpoint, the condition refiner once an image and all 50 guided steps on the GPU. On an M5 Ultra a forward takes 0.273 s against the publisher's SDK on torch MPS at 0.939 s, with a 1.84 GB peak footprint. Each step's velocities on the SDK's own input sit as close to an fp32 run as the SDK's own, and each joint block alone passes the same gate; `tf-idol` runs those gates from a case `tools/zig/idol_case.py` records. RePo's fp32 maps run as an fp16 hi/lo split on the tensor units, whose fp32 GEMM is about fp16-accurate. The image encoders and the decoders are not ported yet (fixes #1757). Thanks to @jayleaton. +- Idol (speridlabs/idol, image to 3D Gaussian splat), first step: its 0.37B denoiser runs natively on M5 Metal from the released checkpoint, the condition refiner once an image and all 50 guided steps on the GPU. On an M5 Ultra a forward takes 0.273 s against the publisher's SDK on torch MPS at 0.939 s, with a 1.84 GB peak footprint. Each step's velocities on the SDK's own input sit as close to an fp32 run as the SDK's own, and each joint block alone passes the same gate; `tf-idol` runs those gates from a case `tools/zig/idol_case.py` records. RePo's fp32 maps run as an fp16 hi/lo split on the tensor units, whose fp32 GEMM is about fp16-accurate. The image encoders and the decoders are not ported yet (#1766, fixes #1757). Thanks to @jayleaton. - DeepSeek-V4.x: `tool_choice: "required"` and named tool calls are now constrained to the model's own DSML call markup, so those requests are held to a valid call on DeepSeek too. From 427c7465eea4dc617afb871114586b1de95ac464 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 03:09:00 +0800 Subject: [PATCH 5/9] idol: attention rounds P to fp16 before P V, so its denominator sums the values the units multiply --- CHANGELOG.md | 2 +- docs/recipes/idol.md | 15 ++++++++------- zig/kernels/metal/idol_attention.metal | 3 ++- 3 files changed, 11 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 38fd52543d..a04a4c04eb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased -- Idol (speridlabs/idol, image to 3D Gaussian splat), first step: its 0.37B denoiser runs natively on M5 Metal from the released checkpoint, the condition refiner once an image and all 50 guided steps on the GPU. On an M5 Ultra a forward takes 0.273 s against the publisher's SDK on torch MPS at 0.939 s, with a 1.84 GB peak footprint. Each step's velocities on the SDK's own input sit as close to an fp32 run as the SDK's own, and each joint block alone passes the same gate; `tf-idol` runs those gates from a case `tools/zig/idol_case.py` records. RePo's fp32 maps run as an fp16 hi/lo split on the tensor units, whose fp32 GEMM is about fp16-accurate. The image encoders and the decoders are not ported yet (#1766, fixes #1757). Thanks to @jayleaton. +- Idol (speridlabs/idol, image to 3D Gaussian splat), first step: its 0.37B denoiser runs natively on M5 Metal from the released checkpoint, the condition refiner once an image and all 50 guided steps on the GPU. On an M5 Ultra a forward takes 0.275 s against the publisher's SDK on torch MPS at 0.939 s, with a 1.84 GB peak footprint. Each step's velocities on the SDK's own input sit as close to an fp32 run as the SDK's own, and each joint block alone passes the same gate; `tf-idol` runs those gates from a case `tools/zig/idol_case.py` records. RePo's fp32 maps run as an fp16 hi/lo split on the tensor units, whose fp32 GEMM is about fp16-accurate. The image encoders and the decoders are not ported yet (#1766, fixes #1757). Thanks to @jayleaton. - DeepSeek-V4.x: `tool_choice: "required"` and named tool calls are now constrained to the model's own DSML call markup, so those requests are held to a valid call on DeepSeek too. diff --git a/docs/recipes/idol.md b/docs/recipes/idol.md index 27995b8120..007f76fe6b 100644 --- a/docs/recipes/idol.md +++ b/docs/recipes/idol.md @@ -36,7 +36,8 @@ This family keeps that split: - LayerNorms in fp32, rounded to fp16; the modulation's shift and scale in fp16 ops, as torch's fp16 tensors run them. - RePo in fp32: its LayerNorm, both 1,024 to 128 maps, the 128 to 48 map and the angles. - Attention: fp16 q, k and v, fp32 scores and sums. q and k are rotated, then L2-normalized and scaled by the - checkpoint's gamma times 8. + checkpoint's gamma times 8. The tensor units take an fp32 P short, which left outputs 3.5e-4 low on average, so + `id_attention` rounds P to fp16 itself and sums the softmax's denominator over those same values. - The condition embedding and the output projection in fp32; the time MLP, modulation and camera MLPs on the host. ### The tensor units' fp32 matmul is not fp32 @@ -116,14 +117,14 @@ Setup: Results: - **Velocity**, both branches of all 50 steps on the SDK's own input (100 of 100 pass): - - against the SDK: cosine 0.999997 or better; - - against fp32: our worst row over the run 0.0456, the SDK's 0.0398; per step our `1 - cos` is the SDK's within - the 2x bound at every step, and the smaller at 40 of the 100. + - against the SDK: cosine 0.999998 or better; + - against fp32: our worst row over the run 0.0474, the SDK's 0.0398; per step our `1 - cos` is the SDK's within + the 2x bound at every step, and the smaller at 75 of the 100. - **Blocks**, at steps 1 and 41, each of the 23 joint blocks after the first on the SDK's input: all pass, worst rows - 0.0007 to 0.0017 against fp32, where the SDK's are 0.0007 to 0.0025. -- **Sample:** our 50 guided steps from the same noise end at cosine 0.99999 to the SDK's sample. + 0.0008 to 0.0017 against fp32, where the SDK's are 0.0007 to 0.0025. +- **Sample:** our 50 guided steps from the same noise end at cosine 0.999993 to the SDK's sample. - **Repeat:** identical bytes. -- **Speed**, one session: 0.273 s a forward (27.3 s for 50 guided steps), 3.4x the SDK's 0.939 s (the median of its +- **Speed**, one session: 0.275 s a forward (27.5 s for 50 guided steps), 3.4x the SDK's 0.939 s (the median of its 100). Load 0.2 to 1.6 s; the condition, once an image, 0.07 s. - **Memory:** 1.84 GB peak process footprint (`/usr/bin/time -l`, the whole `tf-idol` run with its case), the denoiser's fp16 weights in shared Metal buffers. diff --git a/zig/kernels/metal/idol_attention.metal b/zig/kernels/metal/idol_attention.metal index f676630fc5..8cc6919b33 100644 --- a/zig/kernels/metal/idol_attention.metal +++ b/zig/kernels/metal/idol_attention.metal @@ -79,7 +79,8 @@ inline void id_attention(const device T* Q, const device T* K, const device T* V TF_UNROLL for (short f = 0; f < 2; f++) { TF_UNROLL - for (short e = 0; e < 8; e++) s[f][e] = fast::exp2(s[f][e] - top_new[e >> 2]); + // the units take an fp32 P short: round it to fp16 here, so the denominator sums the values P V multiplies + for (short e = 0; e < 8; e++) s[f][e] = float(half(fast::exp2(s[f][e] - top_new[e >> 2]))); } float scale_old[2]; TF_UNROLL From 94a0392497f8dad77c0fd237e1880fcf007c7b70 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 04:03:17 +0800 Subject: [PATCH 6/9] idol: the recipe names the tensor units' fp32 operand truncation (#1772) --- CHANGELOG.md | 2 +- docs/recipes/idol.md | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a04a4c04eb..2abaa9af8a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased -- Idol (speridlabs/idol, image to 3D Gaussian splat), first step: its 0.37B denoiser runs natively on M5 Metal from the released checkpoint, the condition refiner once an image and all 50 guided steps on the GPU. On an M5 Ultra a forward takes 0.275 s against the publisher's SDK on torch MPS at 0.939 s, with a 1.84 GB peak footprint. Each step's velocities on the SDK's own input sit as close to an fp32 run as the SDK's own, and each joint block alone passes the same gate; `tf-idol` runs those gates from a case `tools/zig/idol_case.py` records. RePo's fp32 maps run as an fp16 hi/lo split on the tensor units, whose fp32 GEMM is about fp16-accurate. The image encoders and the decoders are not ported yet (#1766, fixes #1757). Thanks to @jayleaton. +- Idol (speridlabs/idol, image to 3D Gaussian splat), first step: its 0.37B denoiser runs natively on M5 Metal from the released checkpoint, the condition refiner once an image and all 50 guided steps on the GPU. On an M5 Ultra a forward takes 0.275 s against the publisher's SDK on torch MPS at 0.939 s, with a 1.84 GB peak footprint. Each step's velocities on the SDK's own input sit as close to an fp32 run as the SDK's own, and each joint block alone passes the same gate; `tf-idol` runs those gates from a case `tools/zig/idol_case.py` records. RePo's fp32 maps run as an fp16 hi/lo split on the tensor units, whose fp32 GEMM truncates its inputs to 10 mantissa bits (#1772). The image encoders and the decoders are not ported yet (#1766, fixes #1757). Thanks to @jayleaton. - DeepSeek-V4.x: `tool_choice: "required"` and named tool calls are now constrained to the model's own DSML call markup, so those requests are held to a valid call on DeepSeek too. diff --git a/docs/recipes/idol.md b/docs/recipes/idol.md index 007f76fe6b..5b506262a9 100644 --- a/docs/recipes/idol.md +++ b/docs/recipes/idol.md @@ -43,8 +43,9 @@ This family keeps that split: ### The tensor units' fp32 matmul is not fp32 On the M5, a GEMM whose inputs are fp32 on the tensor units (`tfp::gemm_nax`, the Qwen-VL tower's -`qwv_gemm_f32`, MLX's fp32 matmul) rounds its inputs to about fp16: its error against float64 is about 7.5e-4 of the -result, against about 1e-7 for an fp32 FMA. RePo's maps at that precision move block 8's attention out of its gate. +`qwv_gemm_f32`, MLX's fp32 matmul) truncates each input to 10 mantissa bits, toward zero (#1772): its error against +float64 is about 7.5e-4 of the result, against about 1e-7 for an fp32 FMA. RePo's maps at that precision move +block 8's attention out of its gate. `id_norm_split` therefore writes each RePo input row as fp16 `[hi | hi | lo]` (`lo = x - hi`), and the map's weights are stored `[hi | lo | hi]`, scaled by a power of two so `lo` stays normal. One fp16 split-K GEMM with K = 3,072 then sums From 5d5e9fa86db874457603dfa9d619f9828250d38e Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 03:52:06 +0800 Subject: [PATCH 7/9] idol: the FLUX.2 VAE image encoder natively on Metal, in fp32, with tf-idol-vae and its case Builds the denoiser's feature2 [4101, 128] from the preprocessed picture as idol-sdk's ImageEncoder does: the encoder, quant_conv, the posterior's sample with the SDK's noise as an input, the 2 x 2 patches, the batch norm's statistics and five zero rows. fp32 throughout on the SDK's bf16 weights (FLUX.2-klein-4B, or FLUX.2-dev rounded): Winograd 3x3s, implicit-GEMM downsamples and 1x1s, two-pass GroupNorm, single-head attention. idol_case.py vae records the SDK's MPS rows, its noise and a float64 ruler; tf-idol-vae gates worst row and 1 - cosine within twice the SDK's distance from float64 (ours 5.1e-5 and 5.0e-11 against 2.7e-2 and 5.1e-5), and repeat bytes. All eight kernels register with the contract: 13 forms pass gates 1-4, 12 controls red. --- CHANGELOG.md | 2 + docs/kernel-contract.md | 2 + docs/recipes/idol.md | 99 +++- tools/zig/idol_case.py | 148 ++++- zig/build/metal_targets.zig | 1 + zig/kernels/metal/idol_vae.metal | 446 +++++++++++++++ zig/kernels/metal/sources.zig | 2 + zig/src/core/contract/idol_records.zig | 12 +- zig/src/families/idol/idol.zig | 8 +- zig/src/families/idol/vae.zig | 233 ++++++++ zig/src/families/idol/vae_kernels.zig | 154 ++++++ zig/src/families/idol/vae_weights.zig | 308 +++++++++++ zig/tests/contract/metal_idol_vae.zig | 176 ++++++ zig/tests/contract/metal_idol_vae_rows.zig | 437 +++++++++++++++ zig/tests/contract/metal_idol_vae_steps.zig | 578 ++++++++++++++++++++ zig/tests/contract/metal_idol_vae_wino.zig | 261 +++++++++ zig/tests/contract/metal_limits.zig | 2 + zig/tests/contract/metal_main.zig | 2 + zig/tests/idol_vae.zig | 132 +++++ 19 files changed, 2998 insertions(+), 5 deletions(-) create mode 100644 zig/kernels/metal/idol_vae.metal create mode 100644 zig/src/families/idol/vae.zig create mode 100644 zig/src/families/idol/vae_kernels.zig create mode 100644 zig/src/families/idol/vae_weights.zig create mode 100644 zig/tests/contract/metal_idol_vae.zig create mode 100644 zig/tests/contract/metal_idol_vae_rows.zig create mode 100644 zig/tests/contract/metal_idol_vae_steps.zig create mode 100644 zig/tests/contract/metal_idol_vae_wino.zig create mode 100644 zig/tests/idol_vae.zig diff --git a/CHANGELOG.md b/CHANGELOG.md index 2abaa9af8a..810a7ef6f2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Idol (speridlabs/idol, image to 3D Gaussian splat), first step: its 0.37B denoiser runs natively on M5 Metal from the released checkpoint, the condition refiner once an image and all 50 guided steps on the GPU. On an M5 Ultra a forward takes 0.275 s against the publisher's SDK on torch MPS at 0.939 s, with a 1.84 GB peak footprint. Each step's velocities on the SDK's own input sit as close to an fp32 run as the SDK's own, and each joint block alone passes the same gate; `tf-idol` runs those gates from a case `tools/zig/idol_case.py` records. RePo's fp32 maps run as an fp16 hi/lo split on the tensor units, whose fp32 GEMM truncates its inputs to 10 mantissa bits (#1772). The image encoders and the decoders are not ported yet (#1766, fixes #1757). Thanks to @jayleaton. +- Idol's FLUX.2 VAE image encoder runs natively on Metal in fp32: one 1,024 x 1,024 picture to the denoiser's 4,101 condition rows, with the posterior's noise as an input. It is 0.185 s an encode on an M5 Ultra against 0.095 s for the SDK's bf16 encoder on torch MPS, and lies far closer to a float64 run: worst row 5.1e-5 against the SDK's 2.7e-2. `tf-idol-vae` gates it from a case `tools/zig/idol_case.py vae` records, with the ungated FLUX.2-klein-4B weights (after #1766, for #1757). Thanks to @jayleaton. + - DeepSeek-V4.x: `tool_choice: "required"` and named tool calls are now constrained to the model's own DSML call markup, so those requests are held to a valid call on DeepSeek too. - Qwen3.6-35B-A3B on CUDA: `--drafter` also takes a DFlash2 draft model, chosen by the `dflash_config` in its `config.json`; the checkpoint's own directory still drafts with its MTP layer. With `incoai/Qwen3.6-35B-A3B-DFlash2` over `nvidia/Qwen3.6-35B-A3B-NVFP4`, one GB10 decodes 130 tok/s at one stream on a 4.2K-token prompt (MTP: 99), and every drafted reply equals its plain run. diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index b2b465117a..1d920643c7 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -188,6 +188,8 @@ HIP: the RMS norm rows, a wave a row (fp32, fp16 and bf16 out) and a block a row Metal also checks Idol's denoiser (`metal_idol*.zig`): its fp16 projections and GELU form, the split-K fp32 rows, `id_gemm32` (true fp32 sums: the NAX fp32 path in its place is a control gate 4 must catch), the LayerNorms, RePo's angles, the heads, the elementwise steps and the attention's four forms, each at idol-sdk's fp16 points (19 controls). +Its image encoder (`metal_idol_vae*.zig`, FLUX.2's VAE encoder in fp32) registers its convolutions (direct, implicit +GEMM and Winograd), GroupNorm's passes, the 512-wide attention and the condition rows against float64 (12 controls). RADIO's fused bias, activation and residual (`tf_vit_bias_act`) replaces `tf_vit_bias` then `tf_vit_pointwise`, byte for byte. diff --git a/docs/recipes/idol.md b/docs/recipes/idol.md index 5b506262a9..a449405e51 100644 --- a/docs/recipes/idol.md +++ b/docs/recipes/idol.md @@ -15,7 +15,8 @@ each row's LayerNorm to per-head 3D positions, angles from 10, 10 and 12 frequen The sampler's defaults: 50 Euler steps, `t = 3u / (1 + 2u)`, guidance `v = 4 v_cond - 3 v_uncond` (two passes a step: with the condition, and with zeros). -The image encoders, the decoders and the PLY writer are not ported yet. `tools/zig/idol_case.py` runs the publisher's +The FLUX.2 VAE image encoder runs natively too ([below](#vae-encoder)); DINOv3, the decoders and the PLY writer are +not ported yet. `tools/zig/idol_case.py` runs the publisher's SDK in torch, a development reference. It is a denoiser, not a token family: it does not register with the lane engine, the server or `families.detect`. @@ -129,3 +130,99 @@ Results: 100). Load 0.2 to 1.6 s; the condition, once an image, 0.07 s. - **Memory:** 1.84 GB peak process footprint (`/usr/bin/time -l`, the whole `tf-idol` run with its case), the denoiser's fp16 weights in shared Metal buffers. + +## VAE encoder + +`zig/src/families/idol/vae*.zig` (with `idol_vae.metal`) builds the denoiser's `feature2` [4101, 128] from the 1,024 x +1,024 picture `IdolPipeline.preprocess` returns, the way the SDK's `ImageEncoder._flux2_tokens` does: + +- FLUX.2's VAE encoder: `conv_in`, four `DownEncoderBlock2D` stages of two ResNet blocks (128, 256, 512, 512 wide; + GroupNorm 32 and SiLU) with a stride-2 3x3 downsample padded right and below after the first three, the mid block + (ResNet, single-head attention of 512 over the 16,384 positions, ResNet), `conv_norm_out`, SiLU, `conv_out` and + `quant_conv`, on `2 x - 1` of the picture; +- the posterior's **sample**, not its mean: `mean + exp(logvar / 2) * noise` with logvar clamped to [-30, 20]. Some + channels' standard deviations reach 1, so the noise matters. The SDK draws it from torch's MPS generator in bf16 after + `torch.manual_seed(seed)`, which `encode` cannot reproduce; it takes the noise [128, 128, 32] as an input, and the + case records the SDK's; +- each 2 x 2 patch of the 32 latent channels as channel `c * 4 + dy * 2 + dx` (128 channels on a 64 x 64 grid), less + the batch norm's `running_mean` over `sqrt(running_var + 1e-4)`; +- five zero rows, then the 4,096 patches in row order: 4,101 rows. + +### Weights and precision + +`black-forest-labs/FLUX.2-klein-4B` `vae/` at `e7b7dc27f91deacad38e78976d1f2b499d76a294` (Apache 2.0, ungated, bf16). +The SDK loads the VAE with `torch_dtype=torch.bfloat16`. Klein's file is FLUX.2-dev's fp32 VAE rounded to bf16, so both +give the same weights. The loader takes either: it rounds an fp32 file to bf16 as torch does. It reads the 106 encoder +tensors, `quant_conv` and the two batch-norm statistics, and refuses any other geometry. + +The SDK then runs the whole encoder in bf16. This port computes in fp32 on those bf16 weights: 3x3 convolutions by +Winograd F(2x2, 3x3) with `G g G^T` made in float64 and rounded once, the stride-2 downsamples and 1x1s as implicit GEMMs, +GroupNorm statistics in two passes (the mean, then the squared deviations), fp32 attention, `precise::exp`. None of it +uses the tensor units, whose fp32 GEMM is about fp16-accurate (above). + +### Gates + +```sh +hf download black-forest-labs/FLUX.2-klein-4B --revision e7b7dc27f91deacad38e78976d1f2b499d76a294 --include 'vae/*' --local-dir KLEIN +python3 tools/zig/idol_case.py vae KLEIN/vae SDK picture.png case/ --stages +zig build tf-idol-vae -Dcpu=apple_m1 +zig-out/bin/tf-idol-vae KLEIN/vae case/ [features.safetensors] +``` + +`idol_case.py vae` imports `idol/conditioning.py` and diffusers 0.42.1's VAE files only when their sha256 values are +the pinned ones, and the weights only when their sha256 is klein's or FLUX.2-dev's. It runs the SDK's own +`_flux2_tokens` on MPS, seeded as `IdolPipeline` seeds it, and records the noise it draws. Then it runs the same steps +in float64 on the CPU with that noise: the ruler, from the exact picture rather than its bf16 rounding. `--stages` adds +float64's output of each down block, the mid block and `quant_conv`. + +`tf-idol-vae` gates our rows against float64. Worst row (the largest relative L2 error of a row) and `1 - cosine` +must both be within twice the SDK's own distance, and the padding rows exactly zero. Five more encodes must repeat the +first byte for byte. `vae.zig`'s test plants a 2% error in one row, then a non-zero padding value, and requires the +gate to fail each. + +### Kernel contract + +`metal_idol_vae*.zig` registers all eight kernels in 13 forms, each against float64: `iv_conv_in`, `iv_conv` (1x1, +1x1 plus residual, stride-2 3x3), `iv_wino` (plain, plus residual), `iv_norm` (SiLU, plain), `iv_group_part` (sums, +squared deviations), `iv_group_finish`, `iv_attention` and `iv_tokens`. + +There are twelve controls: + +- inputs staged in half (`iv_conv_in`, the 1x1, Winograd's patch); +- symmetric padding in the stride-2 3x3; +- a dropped term of Winograd's output transform; +- a half SiLU; +- half sums (Winograd, both GroupNorm passes), caught at gate 4; +- attention without its rescale; +- swapped patch axes; +- a NaN logvar lost to Metal's `clamp`. `torch.clamp` keeps NaN, so `iv_tokens` does too. + +Winograd's float64 reference is its own transform, so a non-finite input reaches its whole 2 x 2 tile. The pictures +and 1x1 row counts are smaller than the engine's, and rows only tile. These kernels use no tensor ops; their records +reserve their thread counts in the source. + +``` +SUMMARY 13 kernels: 13 pass gates 1-4, 0 fail, 0 known findings; 12 controls: 12 red as planted, 0 missed; 0 limit problems +``` + +### Measured + +On the M5 Ultra, with the avocado example the SDK preprocessed, seed 42: + +| against float64 | worst row | 1 - cosine | +| --- | --- | --- | +| ours | 5.1e-5 | 5.0e-11 | +| the SDK (bf16, MPS) | 2.73e-2 | 5.05e-5 | + +- **Gate:** passes, by about 500 times on worst row. +- **Against the SDK:** cosine 0.999949, worst row 0.0275. The case's SDK rows equal the SDK's own recorded + `feature2` for this picture bit for bit, so its noise is the pipeline's. +- **Stages:** each down block, the mid block and the moments sit within 4e-5 (worst pixel) of float64. +- **Repeat:** identical bytes. +- **Speed:** 0.185 s an encode (median of five, 0.185 s of it on the GPU), where the SDK takes 0.095 s on torch MPS + (median of five, after a first call at 0.19 s). The SDK's bf16 runs on the tensor units; this fp32 path runs at 9 + to 14 TFLOP/s, and the attention, which reads K and V once every 8 queries, takes about 40 ms of it. Load: 0.3 to + 0.4 s. +- **Memory:** 3.28 GB peak footprint (`/usr/bin/time -l`, the whole `tf-idol-vae` run with its case and stage + copies), 2 GB of it four 512 MB picture tensors for the 1,024 x 1,024 x 128 stage. +- **FLUX.2-dev's fp32 file** in klein's place gives the same rows bit for bit, as rounding it to bf16 should. diff --git a/tools/zig/idol_case.py b/tools/zig/idol_case.py index 83453a8e54..3bb1d378cd 100644 --- a/tools/zig/idol_case.py +++ b/tools/zig/idol_case.py @@ -14,8 +14,11 @@ case WEIGHTS SDK OUT --features F [F] noise, each step's input and both velocities, the final sample blocks WEIGHTS SDK CASE [--step S] every block's rows at step S, and each block in float32 on its input exact WEIGHTS SDK CASE [--steps N] both velocities in float32 on the case's first N inputs (the ruler) + vae VAE SDK IMAGE OUT [--stages] the FLUX.2 VAE rows of a preprocessed image: the SDK's, then float64's WEIGHTS is a speridlabs/idol snapshot; SDK an idol-sdk checkout at the pinned commit. F holds feature1 (DINOv3 rows [1, 4101, 1280]) and feature2 (FLUX.2 VAE rows [1, 4101, 128]), in one file or two, as the SDK's ImageEncoder returns. +VAE is the vae/ folder of black-forest-labs/FLUX.2-klein-4B at the pinned revision (FLUX.2-dev's gives the same bf16 +weights); IMAGE the 1024 x 1024 RGB picture IdolPipeline.preprocess returns. """ SDK_COMMIT = "4702a7f00850e34b6e83a2170cc123c3e9570744" @@ -25,6 +28,21 @@ "idol/sampler.py": "532987b5599be678361c2fcf566335bb3e1b0171696f985818d10b7460bfc8b0", } WEIGHTS_REVISION = "2731c2291fc65a797c5cb627a5e5526fc1490902" +CONDITIONING_SHA256 = "8c9cc50d327df3a5269cda392ef46cf4ed9c1e603f42f46050beda83f355f5ab" +DIFFUSERS_VERSION = "0.42.1" +DIFFUSERS_SHA256 = { + "models/autoencoders/autoencoder_kl_flux2.py": "1afe08d76a93e994a453bd9036b57d173516211c2e6b9794ea64f485bea57f57", + "models/autoencoders/vae.py": "59b51fdf2d951da45d697569271e0c99adf4a959ec0f91885f3d9314e925a12c", + "models/unets/unet_2d_blocks.py": "545d158c9d98fb8971c501303117fd3e32d27ae28506b89005c65548404f7437", + "models/resnet.py": "77ad799719fa51bf0fc703f0486915b791c9b69673983abeb3e32c6edcc76650", + "models/downsampling.py": "bce174aa8007b2c40787331e7f46b31eef1eab7785f16b4e54733553b3d14724", + "models/attention_processor.py": "59871fd4e66fc14a26f0e0464706ab24680e13dcd93161c55123678a5e671772", + "utils/torch_utils.py": "bce0a1a4488557e8f80dc02de5782c2b510e91ba2e07f7ae8d26168cee9afc2a", +} +# FLUX.2-klein-4B's vae/ (Apache 2.0, ungated) at its pinned revision, and FLUX.2-dev's: the same weights once in bf16 +VAE_REVISION = "e7b7dc27f91deacad38e78976d1f2b499d76a294" +VAE_SHA256 = {"ca70d2202afe6415bdbcb8793ba8cd99fd159cfe6192381504d6c4d3036e0f04": "FLUX.2-klein-4B", + "d64f3a68e1cc4f9f4e29b6e0da38a0204fe9a49f2d4053f0ec1fa1ca02f9c4b5": "FLUX.2-dev"} DENOISER_ARGS = {"q_token_length": 8192, "in_channels": 16, "cam_channels": 5, "model_channels": 1024, "cond_channels": 1280, "cond2_channels": 128, "num_blocks": 24, "num_refiner_blocks": 2, "num_heads": 16, "mlp_ratio": 4} @@ -216,9 +234,137 @@ def exact(args) -> int: return 0 +def flux2_vae(vae_dir: str, sdk_dir: str): + """The SDK's ImageEncoder and diffusers' FLUX.2 VAE (bf16, as the SDK loads it), each file used checked first.""" + import diffusers + import torch + + def pinned(path: Path, want: str, what: str) -> None: + digest = hashlib.sha256(path.read_bytes()).hexdigest() + if digest != want: + sys.exit(f"{path} has sha256 {digest}; this tool runs only {what}'s ({want})") + + pinned(Path(sdk_dir) / "idol/conditioning.py", CONDITIONING_SHA256, f"idol-sdk {SDK_COMMIT[:12]}") + if diffusers.__version__ != DIFFUSERS_VERSION: + sys.exit(f"diffusers is {diffusers.__version__}; this tool runs only {DIFFUSERS_VERSION}") + root = Path(diffusers.__file__).parent + for name, want in DIFFUSERS_SHA256.items(): + pinned(root / name, want, f"diffusers {DIFFUSERS_VERSION}") + weights = Path(vae_dir) / "diffusion_pytorch_model.safetensors" + digest = hashlib.sha256(weights.read_bytes()).hexdigest() + if digest not in VAE_SHA256: + sys.exit(f"{weights} is neither FLUX.2-klein-4B {VAE_REVISION[:12]}'s vae/ nor FLUX.2-dev's") + sys.path.insert(0, str(Path(sdk_dir))) + from diffusers.models import AutoencoderKLFlux2 + from idol.conditioning import ImageEncoder + + vae = AutoencoderKLFlux2.from_pretrained(vae_dir, torch_dtype=torch.bfloat16) + return ImageEncoder, vae.eval().requires_grad_(False), VAE_SHA256[digest] + + +def tokens64(vae, pixels, noise): + """ImageEncoder._flux2_tokens's steps in float64 with the SDK's noise: the ruler, and each stage when asked.""" + import torch + + latents = vae.encode(pixels * 2 - 1).latent_dist + latents = latents.mean + latents.std * noise + batch, channels, height, width = latents.shape + latents = latents.view(batch, channels, height // 2, 2, width // 2, 2).permute(0, 1, 3, 5, 2, 4) + latents = latents.reshape(batch, channels * 4, height // 2, width // 2) + mean = vae.bn.running_mean.view(1, -1, 1, 1) + std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps) + rows = ((latents - mean) / std).flatten(2).transpose(1, 2) + return torch.cat([torch.zeros(batch, 5, rows.shape[2], dtype=rows.dtype), rows], dim=1) + + +def distance(ours, ruler) -> dict: + """1 - cosine and the worst row's relative L2 error against the ruler, the rows of 128 as tf-idol-vae reads them.""" + a, b = ours.double().reshape(-1, 128), ruler.double().reshape(-1, 128) + cos = float((a * b).sum() / (a.norm() * b.norm())) + rows = ((a - b).norm(dim=1) / b.norm(dim=1).clamp_min(1e-30)).max() + return {"one_minus_cos": 1 - cos, "worst_row": float(rows), "max_abs": float((a - b).abs().max())} + + +def vae_case(args) -> int: + """The SDK's feature2 for one image on MPS (seeded as IdolPipeline seeds it), its noise, and the float64 ruler.""" + import types + + import diffusers.models.autoencoders.vae as vae_module + import numpy as np + import torch + from PIL import Image + from safetensors.torch import save_file + + out = Path(args.out_dir) + out.mkdir(parents=True, exist_ok=True) + image_encoder, vae, source = flux2_vae(args.vae, args.sdk) + picture = Image.open(args.image).convert("RGB") + if picture.size != (1024, 1024): + sys.exit(f"{args.image} is {picture.size}; the encoders see the 1024 x 1024 picture preprocess returns") + pixels = torch.from_numpy(np.array(picture)).permute(2, 0, 1).float().div(255.0)[None] + encoder = types.SimpleNamespace(flux2_vae=vae.to(args.device)) + drawn = [] + randn = vae_module.randn_tensor + + def recorded(*a, **k): + drawn.append(randn(*a, **k)) + return drawn[-1] + + vae_module.randn_tensor = recorded + times = [] + with torch.no_grad(): + for i in range(args.repeats + 1): + torch.manual_seed(args.seed) + mark = time.perf_counter() + rows = image_encoder._flux2_tokens(encoder, pixels.to(args.device)) + flush(args.device) + times.append(time.perf_counter() - mark) + if i == 0: + sdk_rows, noise = rows[0].cpu(), drawn[0].float().cpu() + vae_module.randn_tensor = randn + if any(not torch.equal(d.float().cpu(), noise) for d in drawn): + sys.exit("the seeded noise differs between encodes") + stages = {} + hooks = [] + if args.stages: + named = [(f"down.{i}", b) for i, b in enumerate(vae.encoder.down_blocks)] + [("mid", vae.encoder.mid_block)] + named.append(("moments", vae.quant_conv)) + for name, module in named: + def keep(_m, _i, output, name=name): + stages[name] = output[0].permute(1, 2, 0).float().contiguous() + hooks.append(module.register_forward_hook(keep)) + vae64 = vae.to("cpu").double() + with torch.no_grad(): + mark = time.perf_counter() + ruler = tokens64(vae64, pixels.double(), noise.double())[0] + ruler_s = time.perf_counter() - mark + for h in hooks: + h.remove() + tensors = {"image": pixels[0].permute(1, 2, 0).contiguous(), "noise": noise[0].permute(1, 2, 0).contiguous(), + "sdk": sdk_rows.float().contiguous(), "ruler": ruler.float().contiguous()} + save_file(tensors, str(out / "vae.safetensors")) + if stages: + save_file(stages, str(out / "vae_stages.safetensors")) + sdk = distance(sdk_rows, ruler) + meta = {"seed": args.seed, "device": args.device, "torch": torch.__version__, "sdk": SDK_COMMIT, + "diffusers": DIFFUSERS_VERSION, "vae": source, "vae_revision": VAE_REVISION, + "sdk_vs_float64": sdk, "float64_s": round(ruler_s, 2), "encode_s": [round(t, 4) for t in times]} + (out / "vae.json").write_text(json.dumps(meta, indent=1) + "\n") + warm = sorted(times[1:]) or times + print(json.dumps(meta | {"median_encode_s": warm[len(warm) // 2]})) + return 0 + + def main(argv: list[str]) -> int: ap = argparse.ArgumentParser(description=__doc__, epilog=USAGE, formatter_class=argparse.RawDescriptionHelpFormatter) sub = ap.add_subparsers(dest="cmd", required=True) + p = sub.add_parser("vae") + for name in ("vae", "sdk", "image", "out_dir"): + p.add_argument(name) + p.add_argument("--device", default="mps", choices=("mps", "cpu", "cuda")) + p.add_argument("--seed", type=int, default=42) + p.add_argument("--repeats", type=int, default=5, help="timed encodes after the first") + p.add_argument("--stages", action="store_true", help="float64's down blocks, mid block and moments too") for name in ("case", "blocks", "exact"): p = sub.add_parser(name) p.add_argument("weights") @@ -236,7 +382,7 @@ def main(argv: list[str]) -> int: else: p.add_argument("--steps", type=int, default=50) args = ap.parse_args(argv) - return {"case": case, "blocks": blocks, "exact": exact}[args.cmd](args) + return {"case": case, "blocks": blocks, "exact": exact, "vae": vae_case}[args.cmd](args) if __name__ == "__main__": diff --git a/zig/build/metal_targets.zig b/zig/build/metal_targets.zig index f08c980d2f..b2ed5b4efa 100644 --- a/zig/build/metal_targets.zig +++ b/zig/build/metal_targets.zig @@ -318,6 +318,7 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s .{ .name = "tf-qwen-image", .path = "zig/tests/qwen_image.zig", .about = "Qwen-Image-2.1's transformer on a reference case: each step compared, then the schedule timed" }, .{ .name = "tf-moss-image", .path = "zig/tests/moss_image.zig", .about = "MOSS-Image's transformer on a reference case: each step's two branches compared, then the guided schedule timed" }, .{ .name = "tf-idol", .path = "zig/tests/idol.zig", .about = "Idol's denoiser on a reference case: each step's two branches compared, then the guided schedule timed" }, + .{ .name = "tf-idol-vae", .path = "zig/tests/idol_vae.zig", .about = "Idol's image encoder on a reference case: its condition rows gated against float64 beside the SDK's, then timed" }, .{ .name = "tf-glm-attn-probe", .path = "zig/tests/glm_attn_probe.zig", .about = "GLM-5.3-Flash's dense latent attention kernels on given inputs (scores, probabilities, outputs)" }, .{ .name = "tf-moe-bench", .path = "zig/tests/moe_bench.zig", .about = "Routed-expert prompt kernels from .metal files on a skewed routing: chained speed, bits against the first" }, .{ .name = "tf-mm-bench", .path = "zig/tests/mm_bench.zig", .about = "Prompt matmul kernels from .metal files: chained speed at one shape, bits against the first, error against fp64" }, diff --git a/zig/kernels/metal/idol_vae.metal b/zig/kernels/metal/idol_vae.metal new file mode 100644 index 0000000000..de517018b1 --- /dev/null +++ b/zig/kernels/metal/idol_vae.metal @@ -0,0 +1,446 @@ +// Idol's image encoder (FLUX.2's VAE encoder) in fp32, each picture tensor [rows, columns, channels]. + +#include +using namespace metal; + +#define TF_UNROLL_IV _Pragma("clang loop unroll(full)") + +struct IvConvIn { + uint rows; + uint cols; +}; + +// y [pixels, 128] = conv3x3(2 x - 1) + b of the picture x [pixels, 3] in [0, 1], zero padded; w [27, 128] by tap. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void iv_conv_in( + const device float* x [[buffer(0)]], const device float* w [[buffer(1)]], const device float* b [[buffer(2)]], + device float* y [[buffer(3)]], constant IvConvIn& p [[buffer(4)]], uint2 at [[thread_position_in_grid]]) { + if (at.y >= p.rows * p.cols) return; + const int py = int(at.y / p.cols); + const int px = int(at.y % p.cols); + float s = 0.0f; + TF_UNROLL_IV + for (int ky = 0; ky < 3; ky++) { + TF_UNROLL_IV + for (int kx = 0; kx < 3; kx++) { + const int iy = py + ky - 1; + const int ix = px + kx - 1; + if (iy < 0 || ix < 0 || iy >= int(p.rows) || ix >= int(p.cols)) continue; + const device float* v = x + ulong(uint(iy) * p.cols + uint(ix)) * 3; + TF_UNROLL_IV + for (int c = 0; c < 3; c++) s = fma(v[c] * 2.0f - 1.0f, w[((ky * 3 + kx) * 3 + c) * 128 + at.x], s); + } + } + y[ulong(at.y) * 128 + at.x] = s + b[at.x]; +} + +struct IvConv { + uint rows; + uint cols; + uint in_rows; + uint in_cols; + uint cin; + uint cout; + uint ldw; + uint taps; + uint add; +}; + +// y = conv(x) + b (+ r): a 1x1 (taps 1) or a 3x3 of stride 2 padded right and below (taps 9), w [taps * cin, ldw]. +[[max_total_threads_per_threadgroup(128)]] [[kernel]] void iv_conv( + const device float* x [[buffer(0)]], const device float* w [[buffer(1)]], const device float* b [[buffer(2)]], + const device float* r [[buffer(3)]], device float* y [[buffer(4)]], constant IvConv& p [[buffer(5)]], + uint2 tg [[threadgroup_position_in_grid]], uint t [[thread_index_in_threadgroup]], + uint sg [[simdgroup_index_in_threadgroup]]) { + constexpr uint BK = 32; + constexpr uint lda = BK + 4; + // an implicit GEMM, x [in_rows, in_cols, cin] read a tap at a time: a 64 x 64 tile a threadgroup, fp32 sums + threadgroup float As[64 * lda]; + threadgroup float Cs[64 * 68]; + const uint pixels = p.rows * p.cols; + const uint m0 = tg.y * 64; + const uint n0 = tg.x * 64; + const uint arow = t / 2; + const uint acol = (t % 2) * 16; + const uint pix = m0 + arow; + const uint oy = pix / p.cols; + const uint ox = pix % p.cols; + const uint stride = p.taps == 9 ? 2 : 1; + const uint sr = (sg / 2) * 32; + const uint sc = (sg % 2) * 32; + simdgroup_float8x8 acc[4][4]; + TF_UNROLL_IV + for (uint i = 0; i < 4; i++) { + TF_UNROLL_IV + for (uint j = 0; j < 4; j++) acc[i][j] = simdgroup_float8x8(0.0f); + } + const uint K = p.taps * p.cin; + for (uint k0 = 0; k0 < K; k0 += BK) { + const uint tap = k0 / p.cin; + const uint iy = oy * stride + tap / 3; + const uint ix = ox * stride + tap % 3; + threadgroup float4* dst = (threadgroup float4*)(As + arow * lda + acol); + if (pix < pixels && iy < p.in_rows && ix < p.in_cols) { + const device float4* src = (const device float4*)(x + ulong(iy * p.in_cols + ix) * p.cin + k0 % p.cin + acol); + TF_UNROLL_IV + for (uint i = 0; i < 4; i++) dst[i] = src[i]; + } else { + TF_UNROLL_IV + for (uint i = 0; i < 4; i++) dst[i] = float4(0.0f); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const device float* wk = w + ulong(k0) * p.ldw + n0 + sc; + TF_UNROLL_IV + for (uint kk = 0; kk < BK; kk += 8) { + simdgroup_float8x8 am[4]; + simdgroup_float8x8 bm[4]; + TF_UNROLL_IV + for (uint i = 0; i < 4; i++) simdgroup_load(am[i], As + (sr + i * 8) * lda + kk, lda); + TF_UNROLL_IV + for (uint j = 0; j < 4; j++) simdgroup_load(bm[j], wk + ulong(kk) * p.ldw + j * 8, p.ldw); + TF_UNROLL_IV + for (uint i = 0; i < 4; i++) { + TF_UNROLL_IV + for (uint j = 0; j < 4; j++) simdgroup_multiply_accumulate(acc[i][j], am[i], bm[j], acc[i][j]); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + TF_UNROLL_IV + for (uint i = 0; i < 4; i++) { + TF_UNROLL_IV + for (uint j = 0; j < 4; j++) simdgroup_store(acc[i][j], Cs + (sr + i * 8) * 68 + sc + j * 8, 68); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint e = t; e < 64 * 64; e += 128) { + const uint row = m0 + e / 64; + const uint col = n0 + e % 64; + if (row >= pixels || col >= p.cout) continue; + const ulong at = ulong(row) * p.cout + col; + float v = Cs[(e / 64) * 68 + e % 64] + b[col]; + if (p.add) v = v + r[at]; + y[at] = v; + } +} + +// The patch rows (or columns) a, b and b's sign that B^T's row i combines: (0, 2, -), (1, 2, +), (2, 1, -), (1, 3, -). +inline uint iv_wa(uint i) { return i == 2 ? 2 : (i == 0 ? 0 : 1); } +inline uint iv_wb(uint i) { return i == 0 ? 2 : (i == 3 ? 3 : (i == 1 ? 2 : 1)); } +inline float iv_ws(uint i) { return i == 1 ? 1.0f : -1.0f; } + +// y = conv3x3(x) + b (+ r) by Winograd F(2x2, 3x3), u = G g G^T [16, cin, ldw]; 16 tiles by 64 channels a group. +[[max_total_threads_per_threadgroup(512)]] [[kernel]] void iv_wino( + const device float* x [[buffer(0)]], const device float* u [[buffer(1)]], const device float* b [[buffer(2)]], + const device float* r [[buffer(3)]], device float* y [[buffer(4)]], constant IvConv& p [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], uint t [[thread_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], uint sg [[simdgroup_index_in_threadgroup]]) { + constexpr uint NF = 8; + constexpr uint BK = 16; + constexpr uint TILES = 16; + constexpr uint PR = 6; + constexpr uint PC = 18; + // the input patch [6, 18, 16]; the epilogue's [16 points, 16 tiles, 32 channels] after + threadgroup float buf[16 * 16 * 32]; + const uint n0 = tg.x * NF * 8; + const int y0 = int(tg.y) * 4; + const int x0 = int(tg.z) * 16; + // this lane's two elements of every 8 x 8 fragment: row fm, columns fn and fn + 1 + const uint fm = ((lane / 4) & 4) + (lane / 2) % 4; + const uint fn = ((lane / 4) & 2) * 2 + (lane % 2) * 2; + const uint pi = sg / 4; + const uint pj = sg % 4; + simdgroup_float8x8 acc[2][NF]; + TF_UNROLL_IV + for (uint i = 0; i < 2; i++) { + TF_UNROLL_IV + for (uint j = 0; j < NF; j++) acc[i][j] = simdgroup_float8x8(0.0f); + } + for (uint c0 = 0; c0 < p.cin; c0 += BK) { + for (uint e = t; e < PR * PC * (BK / 4); e += 512) { + const uint px = e / (BK / 4); + const uint q = e % (BK / 4); + const int iy = y0 - 1 + int(px / PC); + const int ix = x0 - 1 + int(px % PC); + float4 v = float4(0.0f); + if (iy >= 0 && ix >= 0 && iy < int(p.rows) && ix < int(p.cols)) + v = *(const device float4*)(x + ulong(uint(iy) * p.cols + uint(ix)) * p.cin + c0 + q * 4); + *(threadgroup float4*)(buf + px * BK + q * 4) = v; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + TF_UNROLL_IV + for (uint kk = 0; kk < BK; kk += 8) { + simdgroup_float8x8 vm[2]; + TF_UNROLL_IV + for (uint i = 0; i < 2; i++) { + // B^T d B at this point, for tile i * 8 + fm and channels fn, fn + 1 + const uint tile = i * 8 + fm; + const threadgroup float* d = buf + (((tile / 8) * 2) * PC + (tile % 8) * 2) * BK + kk + fn; + TF_UNROLL_IV + for (uint c = 0; c < 2; c++) { + const threadgroup float* q = d + c; + const float wa = q[(iv_wa(pi) * PC + iv_wa(pj)) * BK] + iv_ws(pi) * q[(iv_wb(pi) * PC + iv_wa(pj)) * BK]; + const float wb = q[(iv_wa(pi) * PC + iv_wb(pj)) * BK] + iv_ws(pi) * q[(iv_wb(pi) * PC + iv_wb(pj)) * BK]; + vm[i].thread_elements()[c] = wa + iv_ws(pj) * wb; + } + } + const device float* up_ = u + (ulong(sg) * p.cin + c0 + kk) * p.ldw + n0; + TF_UNROLL_IV + for (uint j = 0; j < NF; j++) { + simdgroup_float8x8 um; + simdgroup_load(um, up_ + j * 8, p.ldw); + TF_UNROLL_IV + for (uint i = 0; i < 2; i++) simdgroup_multiply_accumulate(acc[i][j], vm[i], um, acc[i][j]); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + // the output transform A^T M A, 32 channels at a time + for (uint h = 0; h < NF / 4; h++) { + TF_UNROLL_IV + for (uint i = 0; i < 2; i++) { + TF_UNROLL_IV + for (uint j = 0; j < 4; j++) simdgroup_store(acc[i][h * 4 + j], buf + (sg * TILES + i * 8) * 32 + j * 8, 32); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + { + // one tile and channel a thread: its 16 points, then the four outputs A^T M A, rows then columns + const uint co = t % 32; + const uint tile = t / 32; + const uint col = n0 + h * 32 + co; + const threadgroup float* m = buf + tile * 32 + co; + float s[4][2]; + TF_UNROLL_IV + for (uint j = 0; j < 4; j++) { + const float m0 = m[(0 * 4 + j) * TILES * 32]; + const float m1 = m[(1 * 4 + j) * TILES * 32]; + const float m2 = m[(2 * 4 + j) * TILES * 32]; + const float m3 = m[(3 * 4 + j) * TILES * 32]; + s[j][0] = (m0 + m1) + m2; + s[j][1] = (m1 - m2) - m3; + } + const int ty = y0 + int(tile / 8) * 2; + const int tx = x0 + int(tile % 8) * 2; + TF_UNROLL_IV + for (uint rr = 0; rr < 2; rr++) { + const float out[2] = {(s[0][rr] + s[1][rr]) + s[2][rr], (s[1][rr] - s[2][rr]) - s[3][rr]}; + TF_UNROLL_IV + for (uint cc = 0; cc < 2; cc++) { + const int oy = ty + int(rr); + const int ox = tx + int(cc); + if (col >= p.cout || oy >= int(p.rows) || ox >= int(p.cols)) continue; + const ulong at = ulong(uint(oy) * p.cols + uint(ox)) * p.cout + col; + float v = out[cc] + b[col]; + if (p.add) v = v + r[at]; + y[at] = v; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } +} + +struct IvNorm { + uint pixels; + uint channels; + uint silu; +}; + +// y = GroupNorm(x) by stats (mean, rstd per group) with g, beta, then SiLU when silu; x, y [pixels, channels]. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void iv_norm( + const device float* x [[buffer(0)]], const device float* stats [[buffer(1)]], const device float* g [[buffer(2)]], + const device float* beta [[buffer(3)]], device float* y [[buffer(4)]], constant IvNorm& p [[buffer(5)]], + uint2 at [[thread_position_in_grid]]) { + if (at.y >= p.pixels || at.x >= p.channels) return; + const uint grp = at.x / (p.channels / 32); + const ulong i = ulong(at.y) * p.channels + at.x; + float v = (x[i] - stats[grp]) * stats[32 + grp]; + v = v * g[at.x] + beta[at.x]; + y[i] = p.silu ? v / (1.0f + precise::exp(-v)) : v; +} + +struct IvGroups { + uint pixels; + uint channels; + uint chunk; + uint chunks; + uint centered; + float eps; +}; + +// One chunk of pixels' sums of x a group, or of (x - mean)^2 when centered, into part [32, chunks]: a thread a channel. +[[max_total_threads_per_threadgroup(512)]] [[kernel]] void iv_group_part( + const device float* x [[buffer(0)]], const device float* stats [[buffer(1)]], device float* part [[buffer(2)]], + constant IvGroups& p [[buffer(3)]], uint tg [[threadgroup_position_in_grid]], + uint t [[thread_index_in_threadgroup]]) { + threadgroup float sums[512]; + const uint size = p.channels / 32; + const uint first = tg * p.chunk; + const uint last = min(first + p.chunk, p.pixels); + const float mean = p.centered ? stats[t / size] : 0.0f; + float s = 0.0f; + for (uint px = first; px < last; px++) { + const float v = x[ulong(px) * p.channels + t]; + const float d = v - mean; + s += p.centered ? d * d : v; + } + sums[t] = s; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (t < 32) { + float all = 0.0f; + for (uint c = 0; c < size; c++) all += sums[t * size + c]; + part[t * p.chunks + tg] = all; + } +} + +// stats from part's chunk sums, a simdgroup a group: mean[32] when not centered, else rstd[32] = rsqrt(var + eps). +[[max_total_threads_per_threadgroup(1024)]] [[kernel]] void iv_group_finish( + const device float* part [[buffer(0)]], device float* stats [[buffer(1)]], constant IvGroups& p [[buffer(2)]], + uint lane [[thread_index_in_simdgroup]], uint sg [[simdgroup_index_in_threadgroup]]) { + float s = 0.0f; + for (uint i = lane; i < p.chunks; i += 32) s += part[sg * p.chunks + i]; + s = simd_sum(s); + if (lane != 0) return; + const float n = float(p.pixels) * float(p.channels / 32); + if (p.centered) + stats[32 + sg] = metal::precise::rsqrt(s / n + p.eps); + else + stats[sg] = s / n; +} + +struct IvAttention { + uint rows; + float scale; +}; + +// o [rows, 512] = softmax(q k^T * scale) v of qkv rows [rows, 1536] (both padded to 32 rows): 8 queries a group. +[[max_total_threads_per_threadgroup(128)]] [[kernel]] void iv_attention( + const device float* qkv [[buffer(0)]], device float* o [[buffer(1)]], constant IvAttention& p [[buffer(2)]], + uint tg [[threadgroup_position_in_grid]], uint t [[thread_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], uint sg [[simdgroup_index_in_threadgroup]]) { + constexpr uint D = 512; + constexpr uint Q = D / 4; + constexpr uint ld = 3 * D; + constexpr uint BK = 32; + threadgroup float Sp[4][8 * BK]; + threadgroup float Ps[8 * BK]; + threadgroup float As[8]; + threadgroup uint rescale; + // the rows of this lane's two elements of every 8 x 8 fragment + const uint fm = ((lane / 4) & 4) + (lane / 2) % 4; + const uint q0 = tg * 8; + const uint d0 = sg * Q; + simdgroup_float8x8 qm[Q / 8]; + simdgroup_float8x8 acc[Q / 8]; + TF_UNROLL_IV + for (uint d = 0; d < Q / 8; d++) { + simdgroup_load(qm[d], qkv + ulong(q0) * ld + d0 + d * 8, ld); + acc[d] = simdgroup_float8x8(0.0f); + } + const uint row = t / 16; + const uint col = (t % 16) * 2; + float best = -INFINITY; + float total = 0.0f; + for (uint k0 = 0; k0 < p.rows; k0 += BK) { + const device float* kp = qkv + ulong(k0) * ld + D + d0; + TF_UNROLL_IV + for (uint j = 0; j < BK / 8; j++) { + simdgroup_float8x8 s = simdgroup_float8x8(0.0f); + TF_UNROLL_IV + for (uint d = 0; d < Q / 8; d++) { + simdgroup_float8x8 km; + simdgroup_load(km, kp + ulong(j * 8) * ld + d * 8, ld, ulong2(0, 0), true); + simdgroup_multiply_accumulate(s, qm[d], km, s); + } + simdgroup_store(s, Sp[sg] + j * 8, BK); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float v[2]; + float most = best; + TF_UNROLL_IV + for (uint c = 0; c < 2; c++) { + const uint i = row * BK + col + c; + const float sum = (Sp[0][i] + Sp[1][i]) + (Sp[2][i] + Sp[3][i]); + v[c] = k0 + col + c < p.rows ? sum * p.scale : -INFINITY; + most = max(most, v[c]); + } + TF_UNROLL_IV + for (ushort m = 1; m < 16; m <<= 1) most = max(most, simd_shuffle_xor(most, m)); + const float alpha = precise::exp(best - most); + float add = 0.0f; + TF_UNROLL_IV + for (uint c = 0; c < 2; c++) { + const float e = precise::exp(v[c] - most); + Ps[row * BK + col + c] = e; + add += e; + } + TF_UNROLL_IV + for (ushort m = 1; m < 16; m <<= 1) add += simd_shuffle_xor(add, m); + total = total * alpha + add; + best = most; + if (t % 16 == 0) As[row] = alpha; + if (t == 0) rescale = 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (t % 16 == 0 && alpha != 1.0f) rescale = 1; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (rescale) { + const float f = As[fm]; + TF_UNROLL_IV + for (uint d = 0; d < Q / 8; d++) { + acc[d].thread_elements()[0] *= f; + acc[d].thread_elements()[1] *= f; + } + } + const device float* vp = qkv + ulong(k0) * ld + 2 * D + d0; + TF_UNROLL_IV + for (uint j = 0; j < BK / 8; j++) { + simdgroup_float8x8 pm; + simdgroup_load(pm, Ps + j * 8, BK); + TF_UNROLL_IV + for (uint d = 0; d < Q / 8; d++) { + simdgroup_float8x8 vm; + simdgroup_load(vm, vp + ulong(j * 8) * ld + d * 8, ld); + simdgroup_multiply_accumulate(acc[d], pm, vm, acc[d]); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (t % 16 == 0) As[row] = 1.0f / total; + threadgroup_barrier(mem_flags::mem_threadgroup); + const float inv = As[fm]; + TF_UNROLL_IV + for (uint d = 0; d < Q / 8; d++) { + acc[d].thread_elements()[0] *= inv; + acc[d].thread_elements()[1] *= inv; + simdgroup_store(acc[d], o + ulong(q0) * D + d0 + d * 8, D); + } +} + +struct IvTokens { + uint rows; + uint cols; + uint pad; +}; + +// Rows [pad + patches, 128]: pad zeros, then each 2 x 2 patch of z = mean + std * noise as (z - m) / s. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void iv_tokens( + const device float* moments [[buffer(0)]], const device float* noise [[buffer(1)]], + const device float* m [[buffer(2)]], const device float* s [[buffer(3)]], device float* out [[buffer(4)]], + constant IvTokens& p [[buffer(5)]], uint2 at [[thread_position_in_grid]]) { + const uint w2 = p.cols / 2; + if (at.x >= 128 || at.y >= p.pad + (p.rows / 2) * w2) return; + const ulong o = ulong(at.y) * 128 + at.x; + if (at.y < p.pad) { + out[o] = 0.0f; + return; + } + // moments [rows, cols, 64] as mean | logvar, noise [rows, cols, 32]; a patch's channel c * 4 + dy * 2 + dx + const uint patch = at.y - p.pad; + const uint c = at.x / 4; + const uint py = (patch / w2) * 2 + (at.x % 4) / 2; + const uint px = (patch % w2) * 2 + at.x % 2; + const uint pixel = py * p.cols + px; + // torch.clamp keeps a NaN, where Metal's clamp (fmin, fmax) would drop it + const float raw = moments[ulong(pixel) * 64 + 32 + c]; + const float logvar = isnan(raw) ? raw : clamp(raw, -30.0f, 20.0f); + const float z = moments[ulong(pixel) * 64 + c] + precise::exp(0.5f * logvar) * noise[ulong(pixel) * 32 + c]; + out[o] = (z - m[at.x]) / s[at.x]; +} diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig index ef3221a88b..51aee74e04 100644 --- a/zig/kernels/metal/sources.zig +++ b/zig/kernels/metal/sources.zig @@ -135,6 +135,8 @@ pub const qwen_image = @embedFile("prefill/attention_nax.metal") ++ @embedFile(" pub const moss_image = @embedFile("prefill/attention_nax.metal") ++ @embedFile("moss_image.metal"); /// Idol's denoiser: the prefill GEMM template, then its own steps and head-64 attention. pub const idol = @embedFile("prefill/gemm_nax.metal") ++ @embedFile("idol.metal") ++ @embedFile("idol_attention.metal"); +/// Idol's image encoder (FLUX.2's VAE encoder) in fp32. +pub const idol_vae = @embedFile("idol_vae.metal"); /// A prompt chunk's KDA layer in three passes, appended to the generated kda_rows source (its helpers). pub const glm_kda_prompt = @embedFile("glm_kda_prompt.metal"); /// GLM-5.3-Flash speed-up mode's exchange kernels (families/glm/ep.zig). diff --git a/zig/src/core/contract/idol_records.zig b/zig/src/core/contract/idol_records.zig index ae39549164..1644357cef 100644 --- a/zig/src/core/contract/idol_records.zig +++ b/zig/src/core/contract/idol_records.zig @@ -1,4 +1,4 @@ -//! Idol's denoiser launches (idol.metal and idol_attention.metal after gemm_nax.metal), each source-reserved. +//! Idol's launches: the denoiser's (idol.metal, idol_attention.metal after gemm_nax.metal), the image encoder's. const std = @import("std"); const M = @import("limits.zig").MetalLaunch; pub const records = [_]M{ @@ -26,9 +26,17 @@ pub const records = [_]M{ .{ .name = "idol id_attention_01", .threads = 128, .tg_memory = 0, .reserved = 128 }, .{ .name = "idol id_attention_10", .threads = 128, .tg_memory = 0, .reserved = 128 }, .{ .name = "idol id_attention_11", .threads = 128, .tg_memory = 0, .reserved = 128 }, + .{ .name = "idol vae iv_conv_in", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol vae iv_conv", .threads = 128, .tg_memory = 26624, .reserved = 128 }, + .{ .name = "idol vae iv_wino", .threads = 512, .tg_memory = 32768, .reserved = 512 }, + .{ .name = "idol vae iv_norm", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol vae iv_group_part", .threads = 512, .tg_memory = 2048, .reserved = 512 }, + .{ .name = "idol vae iv_group_finish", .threads = 1024, .tg_memory = 0, .reserved = 1024 }, + .{ .name = "idol vae iv_attention", .threads = 128, .tg_memory = 5168, .reserved = 128 }, + .{ .name = "idol vae iv_tokens", .threads = 256, .tg_memory = 0, .reserved = 256 }, }; test "the Idol launches fit every Apple family, one record a pipeline" { try std.testing.expectEqual(@as(usize, 0), try @import("registry.zig").checkMetal(std.testing.allocator, &records, true)); - try std.testing.expectEqual(@as(usize, 24), records.len); + try std.testing.expectEqual(@as(usize, 32), records.len); } diff --git a/zig/src/families/idol/idol.zig b/zig/src/families/idol/idol.zig index a2aff60944..70fa00e6e6 100644 --- a/zig/src/families/idol/idol.zig +++ b/zig/src/families/idol/idol.zig @@ -1,10 +1,13 @@ -//! Idol image to 3D: its denoiser on Metal (a denoiser beside the lane engine, not a token family). +//! Idol image to 3D: its denoiser and image encoder on Metal (beside the lane engine, not a token family). pub const config = @import("config.zig"); pub const host = @import("host.zig"); pub const weights = @import("weights.zig"); pub const kernels = @import("kernels.zig"); pub const model = @import("model.zig"); pub const reference = @import("reference.zig"); +pub const vae = @import("vae.zig"); +pub const vae_kernels = @import("vae_kernels.zig"); +pub const vae_weights = @import("vae_weights.zig"); pub const Model = model.Model; pub const Condition = model.Condition; pub const Timing = model.Timing; @@ -16,4 +19,7 @@ test { _ = kernels; _ = model; _ = reference; + _ = vae; + _ = vae_kernels; + _ = vae_weights; } diff --git a/zig/src/families/idol/vae.zig b/zig/src/families/idol/vae.zig new file mode 100644 index 0000000000..890cda28e4 --- /dev/null +++ b/zig/src/families/idol/vae.zig @@ -0,0 +1,233 @@ +//! Idol's image encoder on Metal: FLUX.2's VAE encoder in fp32, a 1024 x 1024 picture to the condition's 4,101 rows. +const std = @import("std"); +const mtl = @import("metal"); +const ks = @import("vae_kernels.zig"); +const wts = @import("vae_weights.zig"); +const ref = @import("reference.zig"); +const Ref = ks.Ref; + +/// The picture the encoder sees, its latent grid, and the condition rows: 5 zero rows, then a row a 2 x 2 patch. +pub const size = 1024; +pub const latent = size / 8; +pub const pad = 5; +pub const rows = pad + (latent / 2) * (latent / 2); +pub const channels = 4 * wts.latent_channels; + +pub const Timing = struct { wall: f64, gpu: f64 }; + +/// Where `encode` copies each stage's output when asked, [rows, columns, channels] fp32 (null: not wanted). +pub const Stages = struct { down: [4]?[]f32 = @splat(null), mid: ?[]f32 = null, moments: ?[]f32 = null }; + +/// The widest picture tensor: the first stage's 1024 x 1024 x 128. +const most = size * size * wts.widths[0]; + +pub const Encoder = struct { + gpa: std.mem.Allocator, + device: mtl.Device, + queue: mtl.Queue, + k: ks.Kernels, + w: wts.Weights, + scratch: mtl.Buffer, + image: Ref, + noise: Ref, + t: [4]Ref, + moments: Ref, + out: Ref, + stats: Ref, + part: Ref, + + /// The vae/ folder `dir` (FLUX.2-klein-4B's, or FLUX.2-dev's rounded to bf16 as the SDK loads it). + pub fn load(gpa: std.mem.Allocator, io: std.Io, device: mtl.Device, dir: []const u8) !*Encoder { + const v = try gpa.create(Encoder); + errdefer gpa.destroy(v); + v.* = undefined; + v.gpa = gpa; + v.device = device; + v.queue = try device.queue(); + errdefer v.queue.deinit(); + v.k = try ks.Kernels.load(device); + errdefer v.k.deinit(); + v.w = try wts.load(gpa, io, device, dir); + errdefer v.w.deinit(); + const L = latent * latent; + const parts = [_]usize{ size * size * 3, L * wts.latent_channels, most, most, most, most, L * 2 * wts.latent_channels, rows * channels, 2 * ks.groups, ks.groups * ks.chunks(size * size) }; + var total: usize = 0; + for (parts) |n| total += std.mem.alignForward(usize, n * 4, 256); + v.scratch = try device.buffer(total, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + var at: usize = 0; + const views = [_]*Ref{ &v.image, &v.noise, &v.t[0], &v.t[1], &v.t[2], &v.t[3], &v.moments, &v.out, &v.stats, &v.part }; + for (views, parts) |view, n| { + view.* = .{ .buf = v.scratch, .off = at }; + at += std.mem.alignForward(usize, n * 4, 256); + } + return v; + } + + pub fn deinit(v: *Encoder) void { + v.scratch.deinit(); + v.w.deinit(); + v.k.deinit(); + v.queue.deinit(); + v.gpa.destroy(v); + } + + /// The condition rows [4101, 128] of `image` [1024, 1024, 3] in [0, 1] with the posterior's `noise` [128, 128, 32]. + pub fn encode(v: *Encoder, image: []const f32, noise: []const f32, out: []f32, stages: Stages) !Timing { + if (image.len != size * size * 3) return error.PictureSize; + if (noise.len != latent * latent * wts.latent_channels or out.len != rows * channels) return error.LatentRows; + @memcpy(floats(v.image, image.len), image); + @memcpy(floats(v.noise, noise.len), noise); + const started = mtl.clock.seconds(); + var gpu: f64 = 0; + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + var cb = v.queue.commandBuffer(); + var e = cb.compute(.serial); + var R: usize = size; + ks.convIn(e, v.k.get("iv_conv_in"), v.image, v.w.conv_in.w, v.w.conv_in.b, v.t[0], R, R); + var x: usize = 0; + for (v.w.stages, 0..) |s, i| { + for (s.resnets) |r| x = v.resnet(e, r, x, R); + if (s.down) |d| { + const y = (x + 1) % v.t.len; + ks.conv(e, v.k.get("iv_conv"), v.t[x], d.w, d.b, null, v.t[y], .{ .rows = R / 2, .cols = R / 2, .in_rows = R, .in_cols = R, .cin = d.cin, .cout = d.cout, .ldw = d.ldw, .taps = 9 }); + x = y; + R /= 2; + } + if (stages.down[i]) |dst| { + gpu += try finish(cb, e); + @memcpy(dst, floats(v.t[x], dst.len)); + cb = v.queue.commandBuffer(); + e = cb.compute(.serial); + } + } + x = v.resnet(e, v.w.mid[0], x, R); + x = v.attention(e, x, R * R); + x = v.resnet(e, v.w.mid[1], x, R); + if (stages.mid) |dst| { + gpu += try finish(cb, e); + @memcpy(dst, floats(v.t[x], dst.len)); + cb = v.queue.commandBuffer(); + e = cb.compute(.serial); + } + const n = (x + 1) % v.t.len; + const o = (x + 2) % v.t.len; + v.norm(e, v.t[x], v.t[n], R * R, wts.widths[3], v.w.norm_out, true); + v.conv(e, v.w.conv_out, v.t[n], v.t[o], R, null); + v.conv(e, v.w.quant, v.t[o], v.moments, R, null); + ks.tokens(e, v.k.get("iv_tokens"), v.moments, v.noise, v.w.bn_mean, v.w.bn_std, v.out, R, R, pad); + gpu += try finish(cb, e); + if (stages.moments) |dst| @memcpy(dst, floats(v.moments, dst.len)); + @memcpy(out, floats(v.out, out.len)); + return .{ .wall = mtl.clock.seconds() - started, .gpu = gpu }; + } + + fn finish(cb: mtl.CommandBuffer, e: mtl.ComputeEncoder) !f64 { + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.log.err("idol vae: a pass failed on the GPU: {s}", .{msg}); + return error.GpuFailure; + } + return cb.gpuSeconds(); + } + + /// One residual block on t[x] at R x R, its output in another of t; returns that one's index. + fn resnet(v: *const Encoder, e: mtl.ComputeEncoder, r: wts.Resnet, x: usize, R: usize) usize { + const n = (x + 1) % v.t.len; + const h = (x + 2) % v.t.len; + const s = (x + 3) % v.t.len; + v.norm(e, v.t[x], v.t[n], R * R, r.conv1.cin, r.norm1, true); + v.conv(e, r.conv1, v.t[n], v.t[h], R, null); + v.norm(e, v.t[h], v.t[n], R * R, r.conv1.cout, r.norm2, true); + const skip = if (r.shortcut) |sc| blk: { + v.conv(e, sc, v.t[x], v.t[s], R, null); + break :blk v.t[s]; + } else v.t[x]; + v.conv(e, r.conv2, v.t[n], v.t[h], R, skip); + return h; + } + + /// The mid block's self-attention on t[x] (`n` positions): GroupNorm, q | k | v, one head of 512, out + t[x]. + fn attention(v: *const Encoder, e: mtl.ComputeEncoder, x: usize, n: usize) usize { + const xn = (x + 1) % v.t.len; + const qkv = v.t[(x + 2) % v.t.len]; + const att = v.t[(x + 3) % v.t.len]; + const W = ks.width; + v.norm(e, v.t[x], v.t[xn], n, W, v.w.attn_norm, false); + ks.conv(e, v.k.get("iv_conv"), v.t[xn], v.w.qkv.w, v.w.qkv.b, null, qkv, .{ .rows = n, .cols = 1, .in_rows = n, .in_cols = 1, .cin = W, .cout = 3 * W, .ldw = 3 * W }); + ks.attend(e, v.k.get("iv_attention"), qkv, att, n); + ks.conv(e, v.k.get("iv_conv"), att, v.w.attn_out.w, v.w.attn_out.b, v.t[x], v.t[xn], .{ .rows = n, .cols = 1, .in_rows = n, .in_cols = 1, .cin = W, .cout = W, .ldw = W }); + return xn; + } + + /// y = GroupNorm(x) with `nm` (then SiLU when `silu`), x and y [pixels, channels], statistics in two passes. + fn norm(v: *const Encoder, e: mtl.ComputeEncoder, x: Ref, y: Ref, n: usize, ch: usize, nm: wts.Norm, silu: bool) void { + for ([_]bool{ false, true }) |centered| { + ks.groupPart(e, v.k.get("iv_group_part"), x, v.stats, v.part, n, ch, centered); + ks.groupFinish(e, v.k.get("iv_group_finish"), v.part, v.stats, n, ch, centered); + } + ks.normed(e, v.k.get("iv_norm"), x, v.stats, nm.g, nm.b, y, n, ch, silu); + } + + /// y = conv(x) + b (+ r) at R x R, a same-size 3x3 (Winograd) or a 1x1. + fn conv(v: *const Encoder, e: mtl.ComputeEncoder, c: wts.Conv, x: Ref, y: Ref, R: usize, r: ?Ref) void { + const s: ks.ConvShape = .{ .rows = R, .cols = R, .in_rows = R, .in_cols = R, .cin = c.cin, .cout = c.cout, .ldw = c.ldw }; + switch (c.kind) { + .wino => ks.winograd(e, v.k.get("iv_wino"), x, c.w, c.b, r, y, s), + .one => ks.conv(e, v.k.get("iv_conv"), x, c.w, c.b, r, y, s), + .strided, .input => unreachable, + } + } +}; + +fn floats(r: Ref, n: usize) []f32 { + return @as([*]f32, @ptrCast(@alignCast(r.addr())))[0..n]; +} + +/// The encoder's gate against the float64 ruler: worst row and 1 - cosine within twice the SDK's own, rows of 128. +pub const Verdict = struct { ours: ref.Match, sdk: ref.Match, pass: bool }; + +pub fn gate(ours: []const f32, sdk: []const f32, ruler: []const f32) Verdict { + const o = ref.match(ours[pad * channels ..], ruler[pad * channels ..], channels); + const s = ref.match(sdk[pad * channels ..], ruler[pad * channels ..], channels); + const zeros = for (ours[0 .. pad * channels]) |z| { + if (z != 0) break false; + } else true; + return .{ .ours = o, .sdk = s, .pass = zeros and ref.within(o, s) }; +} + +test "the encoder's scratch fits a GPU buffer and the rows are the denoiser's condition rows" { + try std.testing.expectEqual(@as(usize, 4101), rows); + try std.testing.expectEqual(@as(usize, 128), channels); + try std.testing.expect(4 * most * 4 < 1 << 33); +} + +test "a planted mismatch in one condition row fails the encoder's gate, as does a padding row not zero" { + const a = std.testing.allocator; + const n = rows * channels; + const ruler = try a.alloc(f32, n); + defer a.free(ruler); + const sdk = try a.alloc(f32, n); + defer a.free(sdk); + const ours = try a.alloc(f32, n); + defer a.free(ours); + var rng = std.Random.DefaultPrng.init(11); + for (ruler, sdk, ours, 0..) |*r, *s, *o, i| { + const z = i < pad * channels; + r.* = if (z) 0 else rng.random().floatNorm(f32); + s.* = r.* * (1 + 4e-3 * rng.random().floatNorm(f32)); + o.* = r.* * (1 + 1e-6 * rng.random().floatNorm(f32)); + } + try std.testing.expect(gate(ours, sdk, ruler).pass); + const row = 2000; + for (ours[row * channels ..][0..channels]) |*o| o.* *= 1.02; + const v = gate(ours, sdk, ruler); + try std.testing.expect(!v.pass); + try std.testing.expect(v.ours.worst_row > 0.019); + for (ours[row * channels ..][0..channels], ruler[row * channels ..][0..channels]) |*o, r| o.* = r; + try std.testing.expect(gate(ours, sdk, ruler).pass); + ours[3 * channels + 7] = 1e-3; + try std.testing.expect(!gate(ours, sdk, ruler).pass); +} diff --git a/zig/src/families/idol/vae_kernels.zig b/zig/src/families/idol/vae_kernels.zig new file mode 100644 index 0000000000..091005fee8 --- /dev/null +++ b/zig/src/families/idol/vae_kernels.zig @@ -0,0 +1,154 @@ +//! Idol's image encoder pipelines (idol_vae.metal, fp32): convolutions, GroupNorm, attention and the condition rows. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +pub const Ref = @import("../../core/qwen_vl/weights.zig").Ref; + +/// Each pipeline by its function. +pub const names = [_][]const u8{ "iv_conv_in", "iv_conv", "iv_wino", "iv_norm", "iv_group_part", "iv_group_finish", "iv_attention", "iv_tokens" }; + +/// GroupNorm's groups and epsilon, the pixels a statistics part sums, and the mid block's attention width. +pub const groups = 32; +pub const eps: f32 = 1e-6; +pub const chunk = 256; +pub const width = 512; + +pub const Kernels = struct { + lib: mtl.Library, + p: [names.len]mtl.Pipeline, + + pub fn load(device: mtl.Device) !Kernels { + return loadSource(device, sources.idol_vae); + } + + /// The pipelines of `text` (the served source, or a contract control's copy of it). + pub fn loadSource(device: mtl.Device, text: []const u8) !Kernels { + var k: Kernels = undefined; + k.lib = try mtl.Library.fromSource(device, text, mtl.CompileOptions.mlx()); + errdefer k.lib.deinit(); + var made: usize = 0; + errdefer for (k.p[0..made]) |p| p.deinit(); + for (names, 0..) |name, i| { + k.p[i] = try mtl.Pipeline.init(device, k.lib, name, false); + made += 1; + } + return k; + } + + pub fn deinit(k: *Kernels) void { + for (k.p) |p| p.deinit(); + k.lib.deinit(); + } + + pub fn get(k: *const Kernels, comptime name: []const u8) mtl.Pipeline { + inline for (names, 0..) |n, i| if (comptime std.mem.eql(u8, n, name)) return k.p[i]; + @compileError("no pipeline " ++ name); + } +}; + +/// One kernel's buffers in order, then its parameters at the next index. +pub fn bind(e: mtl.ComputeEncoder, p: mtl.Pipeline, refs: []const Ref, params: anytype) void { + e.setPipeline(p); + for (refs, 0..) |r, i| e.setBuffer(r.buf, r.off, i); + e.setValue(params, refs.len); +} + +/// iv_conv_in: y [rows * cols, 128] = conv3x3(2 x - 1) + b of the picture x [rows * cols, 3], w [27, 128]. +pub fn convIn(e: mtl.ComputeEncoder, p: mtl.Pipeline, x: Ref, w: Ref, b: Ref, y: Ref, rows: usize, cols: usize) void { + bind(e, p, &.{ x, w, b, y }, [2]u32{ @intCast(rows), @intCast(cols) }); + e.dispatchThreads(mtl.Size.of(128, rows * cols, 1), mtl.Size.of(128, 2, 1)); +} + +/// A convolution's size: the output picture, its input picture, channels, the weights' row length and taps (1 or 9). +pub const ConvShape = struct { rows: usize, cols: usize, in_rows: usize, in_cols: usize, cin: usize, cout: usize, ldw: usize, taps: usize = 1 }; + +fn convParams(s: ConvShape, add: bool) extern struct { rows: u32, cols: u32, in_rows: u32, in_cols: u32, cin: u32, cout: u32, ldw: u32, taps: u32, add: u32 } { + return .{ .rows = @intCast(s.rows), .cols = @intCast(s.cols), .in_rows = @intCast(s.in_rows), .in_cols = @intCast(s.in_cols), .cin = @intCast(s.cin), .cout = @intCast(s.cout), .ldw = @intCast(s.ldw), .taps = @intCast(s.taps), .add = @intFromBool(add) }; +} + +/// iv_conv: y = conv(x) + b (+ r when given), a 1x1 or a stride-2 3x3; cin % 32 = 0, ldw % 64 = 0. +pub fn conv(e: mtl.ComputeEncoder, p: mtl.Pipeline, x: Ref, w: Ref, b: Ref, r: ?Ref, y: Ref, s: ConvShape) void { + std.debug.assert(s.cin % 32 == 0 and s.ldw % 64 == 0 and (s.taps == 1 or s.taps == 9)); + bind(e, p, &.{ x, w, b, r orelse y, y }, convParams(s, r != null)); + e.dispatchGroups(mtl.Size.of(s.ldw / 64, (s.rows * s.cols + 63) / 64, 1), mtl.Size.of(128, 1, 1)); +} + +/// iv_wino: y = conv3x3(x) + b (+ r when given) by Winograd F(2x2, 3x3), same size; cin % 16 = 0, ldw % 64 = 0. +pub fn winograd(e: mtl.ComputeEncoder, p: mtl.Pipeline, x: Ref, u: Ref, b: Ref, r: ?Ref, y: Ref, s: ConvShape) void { + std.debug.assert(s.cin % 16 == 0 and s.ldw % 64 == 0); + bind(e, p, &.{ x, u, b, r orelse y, y }, convParams(s, r != null)); + e.dispatchGroups(mtl.Size.of(s.ldw / 64, (s.rows + 3) / 4, (s.cols + 15) / 16), mtl.Size.of(512, 1, 1)); +} + +/// The Winograd F(2x2, 3x3) form of one 3x3 kernel g: G g G^T in float64, each of its 16 values rounded once. +pub fn transform(g: *const [9]f32) [16]f32 { + const G = [4][3]f64{ .{ 1, 0, 0 }, .{ 0.5, 0.5, 0.5 }, .{ 0.5, -0.5, 0.5 }, .{ 0, 0, 1 } }; + var out: [16]f32 = undefined; + for (0..4) |i| for (0..4) |j| { + var v: f64 = 0; + for (0..3) |a| for (0..3) |b| { + v += G[i][a] * @as(f64, g[a * 3 + b]) * G[j][b]; + }; + out[i * 4 + j] = @floatCast(v); + }; + return out; +} + +/// iv_norm: y = GroupNorm(x) by stats (32 means, then 32 rstd) with g and beta, then SiLU when `silu`. +pub fn normed(e: mtl.ComputeEncoder, p: mtl.Pipeline, x: Ref, stats: Ref, g: Ref, beta: Ref, y: Ref, n: usize, channels: usize, silu: bool) void { + const params = extern struct { pixels: u32, channels: u32, silu: u32 }{ .pixels = @intCast(n), .channels = @intCast(channels), .silu = @intFromBool(silu) }; + bind(e, p, &.{ x, stats, g, beta, y }, params); + e.dispatchThreads(mtl.Size.of(channels, n, 1), mtl.Size.of(32, 8, 1)); +} + +/// GroupNorm's parts over `n` pixels: one sum a chunk of `chunk` pixels and a group. +pub fn chunks(n: usize) usize { + return (n + chunk - 1) / chunk; +} + +fn groupParams(n: usize, channels: usize, centered: bool) extern struct { pixels: u32, channels: u32, chunk: u32, chunks: u32, centered: u32, eps: f32 } { + return .{ .pixels = @intCast(n), .channels = @intCast(channels), .chunk = chunk, .chunks = @intCast(chunks(n)), .centered = @intFromBool(centered), .eps = eps }; +} + +/// iv_group_part: part [32, chunks] = each chunk's sum of a group's values, or (centered) of (x - mean)^2. +pub fn groupPart(e: mtl.ComputeEncoder, p: mtl.Pipeline, x: Ref, stats: Ref, part: Ref, n: usize, channels: usize, centered: bool) void { + std.debug.assert(channels % groups == 0 and channels <= 512); + bind(e, p, &.{ x, stats, part }, groupParams(n, channels, centered)); + e.dispatchGroups(mtl.Size.of(chunks(n), 1, 1), mtl.Size.of(channels, 1, 1)); +} + +/// iv_group_finish: stats' 32 means from part, or (centered) its 32 rstd. +pub fn groupFinish(e: mtl.ComputeEncoder, p: mtl.Pipeline, part: Ref, stats: Ref, n: usize, channels: usize, centered: bool) void { + bind(e, p, &.{ part, stats }, groupParams(n, channels, centered)); + e.dispatchGroups(mtl.Size.of(1, 1, 1), mtl.Size.of(32 * groups, 1, 1)); +} + +/// iv_attention: o [rows, 512] = softmax(q k^T / sqrt(512)) v of qkv rows [rows, 1536], both padded to 32 rows. +pub fn attend(e: mtl.ComputeEncoder, p: mtl.Pipeline, qkv: Ref, o: Ref, rows: usize) void { + const params = extern struct { rows: u32, scale: f32 }{ .rows = @intCast(rows), .scale = 1.0 / @sqrt(@as(f32, width)) }; + bind(e, p, &.{ qkv, o }, params); + e.dispatchGroups(mtl.Size.of((rows + 7) / 8, 1, 1), mtl.Size.of(128, 1, 1)); +} + +/// iv_tokens: out [pad + rows * cols / 4, 128] from moments [rows, cols, 64], noise [rows, cols, 32], m and s [128]. +pub fn tokens(e: mtl.ComputeEncoder, p: mtl.Pipeline, moments: Ref, noise: Ref, m: Ref, s: Ref, out: Ref, rows: usize, cols: usize, pad: usize) void { + bind(e, p, &.{ moments, noise, m, s, out }, [3]u32{ @intCast(rows), @intCast(cols), @intCast(pad) }); + e.dispatchThreads(mtl.Size.of(128, pad + rows * cols / 4, 1), mtl.Size.of(128, 2, 1)); +} + +test "every image encoder pipeline compiles and loads on this GPU" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + var k = try Kernels.load(device); + k.deinit(); +} + +test "the Winograd transform of a delta kernel is G's outer products" { + var g: [9]f32 = @splat(0); + g[4] = 1; + const u = transform(&g); + const col = [4]f32{ 0, 0.5, -0.5, 0 }; + for (0..4) |i| for (0..4) |j| try std.testing.expectEqual(col[i] * col[j], u[i * 4 + j]); +} diff --git a/zig/src/families/idol/vae_weights.zig b/zig/src/families/idol/vae_weights.zig new file mode 100644 index 0000000000..817794aa39 --- /dev/null +++ b/zig/src/families/idol/vae_weights.zig @@ -0,0 +1,308 @@ +//! Idol's image encoder weights: FLUX.2's VAE encoder, quant_conv and batch norm statistics, as the SDK holds them. +const std = @import("std"); +const mtl = @import("metal"); +const ck = @import("../../core/checkpoint_metal.zig"); +const ks = @import("vae_kernels.zig"); +const Ref = ks.Ref; + +/// The encoder's widths by stage, and the rest of the released geometry. +pub const widths = [4]usize{ 128, 256, 512, 512 }; +pub const latent_channels = 32; +pub const resnets = 2; +pub const batch_norm_eps: f32 = 1e-4; + +/// Accepts vae/config.json at FLUX.2's released encoder geometry, else an error. +pub fn check(gpa: std.mem.Allocator, text: []const u8) !void { + const doc = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer doc.deinit(); + if (doc.value != .object) return error.BadConfig; + const o = doc.value.object; + const name = o.get("_class_name") orelse return error.BadConfig; + if (name != .string or !std.mem.eql(u8, name.string, "AutoencoderKLFlux2")) return error.NotFlux2Vae; + const want = .{ .{ "in_channels", 3 }, .{ "latent_channels", latent_channels }, .{ "layers_per_block", resnets }, .{ "norm_num_groups", ks.groups } }; + inline for (want) |kv| { + const v = o.get(kv[0]) orelse return error.BadConfig; + if (v != .integer or v.integer != kv[1]) { + std.log.warn("idol vae: {s} is not {d}; the encoder runs the released geometry only", .{ kv[0], kv[1] }); + return error.UnsupportedGeometry; + } + } + const lists = .{ .{ "block_out_channels", [_]i64{ 128, 256, 512, 512 } }, .{ "patch_size", [_]i64{ 2, 2 } } }; + inline for (lists) |kv| { + const v = o.get(kv[0]) orelse return error.BadConfig; + if (v != .array or v.array.items.len != kv[1].len) return error.UnsupportedGeometry; + for (v.array.items, kv[1]) |item, n| if (item != .integer or item.integer != n) return error.UnsupportedGeometry; + } + inline for (.{ "mid_block_add_attention", "use_quant_conv" }) |f| { + const v = o.get(f) orelse return error.BadConfig; + if (v != .bool or !v.bool) return error.UnsupportedGeometry; + } + const blocks = o.get("down_block_types") orelse return error.BadConfig; + if (blocks != .array or blocks.array.items.len != widths.len) return error.UnsupportedGeometry; + for (blocks.array.items) |b| if (b != .string or !std.mem.eql(u8, b.string, "DownEncoderBlock2D")) return error.UnsupportedGeometry; + const act = o.get("act_fn") orelse return error.BadConfig; + if (act != .string or !std.mem.eql(u8, act.string, "silu")) return error.UnsupportedGeometry; + const bn = o.get("batch_norm_eps") orelse return error.BadConfig; + const bn_eps: f64 = switch (bn) { + .float => |f| f, + .integer => |n| @floatFromInt(n), + else => return error.BadConfig, + }; + if (@as(f32, @floatCast(bn_eps)) != batch_norm_eps) return error.UnsupportedGeometry; +} + +/// A convolution's weights (1x1 [cin, ldw], stride-2 3x3 [9 * cin, ldw], 3x3 G g G^T [16, cin, ldw]) and bias. +pub const Conv = struct { w: Ref, b: Ref, cin: usize, cout: usize, ldw: usize, kind: Kind }; +pub const Kind = enum { one, wino, strided, input }; +pub const Norm = struct { g: Ref, b: Ref }; +pub const Resnet = struct { norm1: Norm, conv1: Conv, norm2: Norm, conv2: Conv, shortcut: ?Conv }; +pub const Stage = struct { resnets: [resnets]Resnet, down: ?Conv }; + +pub const Weights = struct { + buf: mtl.Buffer, + conv_in: Conv, + stages: [4]Stage, + mid: [2]Resnet, + attn_norm: Norm, + qkv: Conv, + attn_out: Conv, + norm_out: Norm, + conv_out: Conv, + quant: Conv, + /// The batch norm's running mean and sqrt(running_var + eps), each of 128. + bn_mean: Ref, + bn_std: Ref, + + pub fn deinit(w: *Weights) void { + w.buf.deinit(); + } +}; + +/// The bf16 value an fp32 weight rounds to (to nearest, ties to even), as torch's cast to bf16 makes it. +pub fn bf16(v: f32) f32 { + if (std.math.isNan(v)) return v; + const bits: u32 = @bitCast(v); + const rounded = bits +% (0x7fff + ((bits >> 16) & 1)); + return @bitCast(rounded & 0xffff0000); +} + +/// The folder `dir` (config.json, diffusion_pytorch_model.safetensors in bf16 or f32): every encoder tensor, packed. +pub fn load(gpa: std.mem.Allocator, io: std.Io, device: mtl.Device, dir: []const u8) !Weights { + var arena_state = std.heap.ArenaAllocator.init(gpa); + defer arena_state.deinit(); + const arena = arena_state.allocator(); + const config = try std.Io.Dir.cwd().readFileAlloc(io, try std.fs.path.join(arena, &.{ dir, "config.json" }), arena, .limited(1 << 20)); + try check(gpa, config); + var c = ck.Checkpoint.init(gpa); + defer c.deinit(); + try c.addFile(device, try std.fmt.allocPrintSentinel(arena, "{s}/diffusion_pytorch_model.safetensors", .{dir}, 0), ""); + var w: Weights = undefined; + var p: Packer = .{ .c = &c, .arena = arena }; + try layout(&w, &p); + if (p.read != expected(&c)) { + std.log.err("idol vae: the checkpoint has {d} encoder tensors; the encoder reads {d}", .{ expected(&c), p.read }); + return error.UnexpectedTensors; + } + w.buf = try device.buffer(p.at * 4, mtl.ResourceOptions.shared); + errdefer w.buf.deinit(); + p.out = @as([*]f32, @ptrCast(@alignCast(w.buf.contents())))[0..p.at]; + p.at = 0; + p.read = 0; + p.buf = w.buf; + try layout(&w, &p); + return w; +} + +/// Walks every tensor in the order it is packed: sizes it on the first pass, copies it on the second. +fn layout(w: *Weights, p: *Packer) !void { + var name: [96]u8 = undefined; + w.conv_in = try p.conv("encoder.conv_in", 3, widths[0], .input); + var prev = widths[0]; + for (&w.stages, widths, 0..) |*s, out, i| { + for (&s.resnets, 0..) |*r, j| { + r.* = try p.resnet(try std.fmt.bufPrint(&name, "encoder.down_blocks.{d}.resnets.{d}", .{ i, j }), if (j == 0) prev else out, out); + } + s.down = if (i + 1 < widths.len) try p.conv(try std.fmt.bufPrint(&name, "encoder.down_blocks.{d}.downsamplers.0.conv", .{i}), out, out, .strided) else null; + prev = out; + } + const W = widths[3]; + for (&w.mid, 0..) |*r, i| r.* = try p.resnet(try std.fmt.bufPrint(&name, "encoder.mid_block.resnets.{d}", .{i}), W, W); + const at = "encoder.mid_block.attentions.0"; + w.attn_norm = try p.norm(at ++ ".group_norm", W); + w.qkv = try p.linears(&.{ at ++ ".to_q", at ++ ".to_k", at ++ ".to_v" }, W); + w.attn_out = try p.linears(&.{at ++ ".to_out.0"}, W); + w.norm_out = try p.norm("encoder.conv_norm_out", W); + w.conv_out = try p.conv("encoder.conv_out", W, 2 * latent_channels, .wino); + w.quant = try p.conv("quant_conv", 2 * latent_channels, 2 * latent_channels, .one); + const n = 4 * latent_channels; + const mean = try p.vector("bn.running_mean", n); + const var_ = try p.vector("bn.running_var", n); + w.bn_mean = p.alloc(n); + w.bn_std = p.alloc(n); + if (p.out) |out| { + // the SDK's sqrt(var + eps) in bf16, here in float64 rounded once + for (out[w.bn_mean.off / 4 ..][0..n], out[w.bn_std.off / 4 ..][0..n], mean, var_) |*mo, *so, mv, v| { + mo.* = mv; + so.* = @floatCast(@sqrt(@as(f64, v) + 1e-4)); + } + } +} + +/// The checkpoint's encoder-side tensors: the encoder's, quant_conv's and the batch norm's two statistics. +fn expected(c: *const ck.Checkpoint) usize { + var n: usize = 0; + var it = c.tensors.keyIterator(); + while (it.next()) |k| { + if (std.mem.startsWith(u8, k.*, "encoder.") or std.mem.startsWith(u8, k.*, "quant_conv.") or + std.mem.eql(u8, k.*, "bn.running_mean") or std.mem.eql(u8, k.*, "bn.running_var")) n += 1; + } + return n; +} + +/// Lays the encoder's tensors out in one fp32 buffer: sizes on the first pass (out null), copies on the second. +const Packer = struct { + c: *const ck.Checkpoint, + arena: std.mem.Allocator, + buf: mtl.Buffer = undefined, + out: ?[]f32 = null, + at: usize = 0, + read: usize = 0, + + fn alloc(p: *Packer, n: usize) Ref { + const r: Ref = .{ .buf = p.buf, .off = p.at * 4 }; + p.at = std.mem.alignForward(usize, p.at + n, 64); + return r; + } + + /// Tensor `name` of `shape` as the bf16 values the SDK computes with (a bf16 file's own, an f32 file's rounded). + fn tensor(p: *Packer, name: []const u8, shape: []const usize) ![]const f32 { + const t = try p.c.get(name); + if ((t.dtype != .f32 and t.dtype != .bf16) or t.rank != shape.len or !std.mem.eql(usize, t.shape[0..t.rank], shape)) { + std.log.err("idol vae: {s} is {t} {any}; the encoder reads bf16 or f32 {any}", .{ name, t.dtype, t.shape[0..t.rank], shape }); + return error.UnexpectedTensor; + } + p.read += 1; + if (p.out == null) return &.{}; + const n = t.count(); + const v = try p.arena.alloc(f32, n); + switch (t.dtype) { + .bf16 => for (v, t.host(u16)) |*o, h| { + o.* = @bitCast(@as(u32, h) << 16); + }, + else => for (v, t.host(f32)) |*o, f| { + o.* = bf16(f); + }, + } + return v; + } + + fn vector(p: *Packer, name: []const u8, n: usize) ![]const f32 { + return p.tensor(name, &.{n}); + } + + /// A copied vector `name` of n. + fn copy(p: *Packer, name: []const u8, n: usize) !Ref { + const v = try p.vector(name, n); + const r = p.alloc(n); + if (p.out) |out| @memcpy(out[r.off / 4 ..][0..n], v); + return r; + } + + /// Convolution `name` (.weight [cout, cin, k, k], .bias) in `kind`'s layout, its columns padded to ldw. + fn conv(p: *Packer, name: []const u8, cin: usize, cout: usize, kind: Kind) !Conv { + var buf: [96]u8 = undefined; + const taps: usize = if (kind == .one) 1 else 9; + const k: usize = if (kind == .one) 1 else 3; + const ldw = if (kind == .input) cout else std.mem.alignForward(usize, cout, 64); + const rows: usize = switch (kind) { + .one => cin, + .wino => 16 * cin, + .strided, .input => 9 * cin, + }; + const w = try p.tensor(try std.fmt.bufPrint(&buf, "{s}.weight", .{name}), &.{ cout, cin, k, k }); + const r = p.alloc(rows * ldw); + if (p.out) |out| { + const dst = out[r.off / 4 ..][0 .. rows * ldw]; + @memset(dst, 0); + for (0..cout) |o| for (0..cin) |i| { + const g = w[(o * cin + i) * taps ..][0..taps]; + switch (kind) { + .one => dst[i * ldw + o] = g[0], + .wino => for (ks.transform(g[0..9]), 0..) |v, pt| { + dst[(pt * cin + i) * ldw + o] = v; + }, + .strided, .input => for (g, 0..) |v, tap| { + dst[(tap * cin + i) * ldw + o] = v; + }, + } + }; + } + const b = try p.copy(try std.fmt.bufPrint(&buf, "{s}.bias", .{name}), cout); + return .{ .w = r, .b = b, .cin = cin, .cout = cout, .ldw = ldw, .kind = kind }; + } + + /// Linear layers `names` (.weight [width, width], .bias) stacked as one 1x1 convolution. + fn linears(p: *Packer, names_: []const []const u8, width: usize) !Conv { + var buf: [96]u8 = undefined; + const n = names_.len * width; + const r = p.alloc(width * n); + for (names_, 0..) |name, part| { + const w = try p.tensor(try std.fmt.bufPrint(&buf, "{s}.weight", .{name}), &.{ width, width }); + if (p.out) |out| { + const dst = out[r.off / 4 ..][0 .. width * n]; + for (0..width) |o| for (0..width) |i| { + dst[i * n + part * width + o] = w[o * width + i]; + }; + } + } + const b = p.alloc(n); + for (names_, 0..) |name, i| { + const v = try p.vector(try std.fmt.bufPrint(&buf, "{s}.bias", .{name}), width); + if (p.out) |out| @memcpy(out[b.off / 4 + i * width ..][0..width], v); + } + return .{ .w = r, .b = b, .cin = width, .cout = n, .ldw = n, .kind = .one }; + } + + fn norm(p: *Packer, name: []const u8, width: usize) !Norm { + var buf: [96]u8 = undefined; + return .{ .g = try p.copy(try std.fmt.bufPrint(&buf, "{s}.weight", .{name}), width), .b = try p.copy(try std.fmt.bufPrint(&buf, "{s}.bias", .{name}), width) }; + } + + fn resnet(p: *Packer, name: []const u8, cin: usize, cout: usize) !Resnet { + var buf: [96]u8 = undefined; + return .{ + .norm1 = try p.norm(try std.fmt.bufPrint(&buf, "{s}.norm1", .{name}), cin), + .conv1 = try p.conv(try std.fmt.bufPrint(&buf, "{s}.conv1", .{name}), cin, cout, .wino), + .norm2 = try p.norm(try std.fmt.bufPrint(&buf, "{s}.norm2", .{name}), cout), + .conv2 = try p.conv(try std.fmt.bufPrint(&buf, "{s}.conv2", .{name}), cout, cout, .wino), + .shortcut = if (cin != cout) try p.conv(try std.fmt.bufPrint(&buf, "{s}.conv_shortcut", .{name}), cin, cout, .one) else null, + }; + } +}; + +test "only FLUX.2's released encoder geometry passes" { + const a = std.testing.allocator; + const sample = + \\{"_class_name": "AutoencoderKLFlux2", "act_fn": "silu", "batch_norm_eps": 0.0001, "block_out_channels": [128, 256, 512, 512], + \\ "down_block_types": ["DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D"], "in_channels": 3, + \\ "latent_channels": 32, "layers_per_block": 2, "mid_block_add_attention": true, "norm_num_groups": 32, + \\ "patch_size": [2, 2], "use_quant_conv": true} + ; + try check(a, sample); + const narrow = try std.mem.replaceOwned(u8, a, sample, "[128, 256, 512, 512]", "[96, 192, 384, 384]"); + defer a.free(narrow); + try std.testing.expectError(error.UnsupportedGeometry, check(a, narrow)); + const no_quant = try std.mem.replaceOwned(u8, a, sample, "\"use_quant_conv\": true", "\"use_quant_conv\": false"); + defer a.free(no_quant); + try std.testing.expectError(error.UnsupportedGeometry, check(a, no_quant)); + const other = try std.mem.replaceOwned(u8, a, sample, "AutoencoderKLFlux2", "AutoencoderKL"); + defer a.free(other); + try std.testing.expectError(error.NotFlux2Vae, check(a, other)); +} + +test "an f32 weight rounds to bf16 as torch's cast does: to nearest, ties to even" { + try std.testing.expectEqual(@as(f32, 1.0), bf16(1.0 + 0x1p-9)); + try std.testing.expectEqual(@as(f32, 1.0 + 0x1p-7), bf16(1.0 + 0x1p-7 + 0x1p-9)); + try std.testing.expectEqual(@as(f32, 1.0 + 0x1p-6), bf16(1.0 + 0x1p-7 + 0x1p-8)); + try std.testing.expectEqual(@as(f32, 1.0 + 0x1p-7), bf16(1.0 + 0x1p-7 + 0x1p-8 - 0x1p-20)); +} diff --git a/zig/tests/contract/metal_idol_vae.zig b/zig/tests/contract/metal_idol_vae.zig new file mode 100644 index 0000000000..5ec14ca3b2 --- /dev/null +++ b/zig/tests/contract/metal_idol_vae.zig @@ -0,0 +1,176 @@ +//! Idol's image encoder kernels (idol_vae.metal, fp32) under the contract; each control breaks one gate. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +const rig_ = @import("rig.zig"); +const id = @import("metal_idol_common.zig"); +const st = @import("metal_idol_vae_steps.zig"); +const wn = @import("metal_idol_vae_wino.zig"); +const rw = @import("metal_idol_vae_rows.zig"); + +const g = id.g; +const Ends = id.Ends; + +/// A control: the kernel it breaks, the gate that must catch it, and the one text change that breaks it. +const Control = struct { kernel: []const u8, gate: g.Gate, what: []const u8, from: []const u8, to: []const u8 }; + +const wino_mma = "for (uint i = 0; i < 2; i++) simdgroup_multiply_accumulate(acc[i][j], vm[i], um, acc[i][j]);"; +const wino_mma_half = "for (uint i = 0; i < 2; i++) { simdgroup_multiply_accumulate(acc[i][j], vm[i], um, acc[i][j]); " ++ + "acc[i][j].thread_elements()[0] = float(half(acc[i][j].thread_elements()[0])); " ++ + "acc[i][j].thread_elements()[1] = float(half(acc[i][j].thread_elements()[1])); }"; +const patch_axes = "const uint py = (patch / w2) * 2 + (at.x % 4) / 2;\n const uint px = (patch % w2) * 2 + at.x % 2;"; +const patch_swapped = "const uint py = (patch / w2) * 2 + at.x % 2;\n const uint px = (patch % w2) * 2 + (at.x % 4) / 2;"; + +const controls = [_]Control{ + .{ .kernel = "iv_conv_in", .gate = .shapes, .what = "stages 2 x - 1 in half", .from = "s = fma(v[c] * 2.0f - 1.0f,", .to = "s = fma(float(half(v[c] * 2.0f - 1.0f))," }, + .{ .kernel = "iv_conv 1x1", .gate = .shapes, .what = "stages its inputs in half", .from = "for (uint i = 0; i < 4; i++) dst[i] = src[i];", .to = "for (uint i = 0; i < 4; i++) dst[i] = float4(half4(src[i]));" }, + .{ .kernel = "iv_conv 3x3 stride 2", .gate = .shapes, .what = "pads above as padding=1 would", .from = "const uint iy = oy * stride + tap / 3;", .to = "const uint iy = oy * stride + tap / 3 - (p.taps == 9 ? 1 : 0);" }, + .{ .kernel = "iv_wino 3x3", .gate = .shapes, .what = "stages its patch in half", .from = "*(threadgroup float4*)(buf + px * BK + q * 4) = v;", .to = "*(threadgroup float4*)(buf + px * BK + q * 4) = float4(half4(v));" }, + .{ .kernel = "iv_wino 3x3 + residual", .gate = .shapes, .what = "drops a term of the output transform", .from = "s[j][1] = (m1 - m2) - m3;", .to = "s[j][1] = (m1 - m2);" }, + .{ .kernel = "iv_wino 3x3", .gate = .sums, .what = "sums in half", .from = wino_mma, .to = wino_mma_half }, + .{ .kernel = "iv_norm silu", .gate = .shapes, .what = "its SiLU in half", .from = "y[i] = p.silu ? v / (1.0f + precise::exp(-v)) : v;", .to = "y[i] = p.silu ? float(half(v) / (1.0h + exp(-half(v)))) : v;" }, + .{ .kernel = "iv_group_part", .gate = .sums, .what = "sums in half", .from = "s += p.centered ? d * d : v;", .to = "s = float(half(s) + half(p.centered ? d * d : v));" }, + .{ .kernel = "iv_group_finish", .gate = .sums, .what = "sums in half", .from = "for (uint i = lane; i < p.chunks; i += 32) s += part[sg * p.chunks + i];", .to = "for (uint i = lane; i < p.chunks; i += 32) s = float(half(s) + half(part[sg * p.chunks + i]));" }, + .{ .kernel = "iv_attention", .gate = .shapes, .what = "keeps stale maxima (no rescale)", .from = " if (rescale) {", .to = " if (false) {" }, + .{ .kernel = "iv_tokens", .gate = .shapes, .what = "swaps the patch's axes", .from = patch_axes, .to = patch_swapped }, + .{ .kernel = "iv_tokens", .gate = .nonfinite, .what = "clamps a NaN logvar (Metal's clamp)", .from = "const float logvar = isnan(raw) ? raw : clamp(raw, -30.0f, 20.0f);", .to = "const float logvar = clamp(raw, -30.0f, 20.0f);" }, +}; + +/// The encoder's source as the engine compiles it, with `ctl`'s change made where it occurs, exactly once. +fn compile(a: std.mem.Allocator, device: mtl.Device, ctl: ?Control) !mtl.Library { + const text = if (ctl) |x| blk: { + if (std.mem.count(u8, sources.idol_vae, x.from) != 1) { + std.debug.print("control {s}: text not found once: {s}\n", .{ x.what, x.from }); + return error.ControlTextNotFoundOnce; + } + break :blk try std.mem.replaceOwned(u8, a, sources.idol_vae, x.from, x.to); + } else try a.dupe(u8, sources.idol_vae); + defer a.free(text); + return mtl.Library.fromSource(device, text, mtl.CompileOptions.mlx()); +} + +/// A test the suite made, and how it goes. +const Made = struct { ptr: *anyopaque, free: *const fn (std.mem.Allocator, *anyopaque) void }; + +pub const Suite = struct { + a: std.mem.Allocator, + pipes: std.ArrayList(mtl.Pipeline) = .empty, + names: std.ArrayList([]u8) = .empty, + libs: std.ArrayList(mtl.Library) = .empty, + made: std.ArrayList(Made) = .empty, + + fn pipeline(s: *Suite, device: mtl.Device, lib: mtl.Library, name: []const u8) !mtl.Pipeline { + const p = try mtl.Pipeline.init(device, lib, name, false); + errdefer p.deinit(); + try s.pipes.append(s.a, p); + return p; + } + + fn label(s: *Suite, comptime fmt: []const u8, args: anytype) ![]const u8 { + const name = try std.fmt.allocPrint(s.a, fmt, args); + errdefer s.a.free(name); + try s.names.append(s.a, name); + return name; + } + + fn make(s: *Suite, comptime T: type, value: T) !*T { + const t = try s.a.create(T); + errdefer s.a.destroy(t); + t.* = value; + const f = struct { + fn free(a: std.mem.Allocator, p: *anyopaque) void { + const x: *T = @ptrCast(@alignCast(p)); + x.deinit(); + a.destroy(x); + } + }; + try s.made.append(s.a, .{ .ptr = t, .free = f.free }); + return t; + } + + /// Each kernel (or, with `ctl`, the one it breaks) on `lib`'s pipelines. + fn add(s: *Suite, device: mtl.Device, queue: mtl.Queue, lib: mtl.Library, ctl: ?Control, out: *std.ArrayList(g.Kernel), only: ?[]const u8) !void { + const suffix = if (ctl) |x| try s.label(" control: {s}", .{x.what}) else ""; + const gate: ?g.Gate = if (ctl) |x| x.gate else null; + const rig: rig_.Rig = .{ .a = s.a, .device = device, .queue = queue }; + const ends: Ends = .{ .a = s.a }; + const W = struct { + s: *Suite, + ctl: ?Control, + suffix: []const u8, + only: ?[]const u8, + /// The test's name when this pass registers `kernel`, else null. + fn name(w: @This(), kernel: []const u8) !?[]const u8 { + if (w.ctl) |x| if (!std.mem.eql(u8, x.kernel, kernel)) return null; + const full = try w.s.label("metal idol vae {s}{s}", .{ kernel, w.suffix }); + if (w.only) |needle| if (std.mem.indexOf(u8, full, needle) == null) return null; + return full; + } + }; + const w = W{ .s = s, .ctl = ctl, .suffix = suffix, .only = only }; + if (try w.name("iv_conv_in")) |name| { + const t = try s.make(st.ConvIn, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "iv_conv_in"), .name = name }); + try out.append(s.a, t.kernel(gate)); + } + for (st.conv_forms) |f| { + const name = (try w.name(f.label)) orelse continue; + const t = try s.make(st.Conv, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "iv_conv"), .form = f, .name = name }); + try out.append(s.a, t.kernel(gate)); + } + for (wn.forms) |f| { + const name = (try w.name(f.label)) orelse continue; + const t = try s.make(wn.Wino, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "iv_wino"), .form = f, .name = name }); + try out.append(s.a, t.kernel(gate)); + } + for ([_]bool{ true, false }) |silu| { + const name = (try w.name(if (silu) "iv_norm silu" else "iv_norm")) orelse continue; + const t = try s.make(st.Norm, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "iv_norm"), .silu = silu, .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + for ([_]bool{ false, true }) |centered| { + const name = (try w.name(if (centered) "iv_group_part centered" else "iv_group_part")) orelse continue; + const t = try s.make(rw.GroupPart, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "iv_group_part"), .centered = centered, .name = name }); + try out.append(s.a, t.kernel(gate)); + } + if (try w.name("iv_group_finish")) |name| { + const t = try s.make(rw.GroupFinish, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "iv_group_finish"), .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + if (try w.name("iv_attention")) |name| { + const t = try s.make(rw.Attention, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "iv_attention"), .name = name }); + try out.append(s.a, t.kernel(gate)); + } + if (try w.name("iv_tokens")) |name| { + const t = try s.make(st.Tokens, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "iv_tokens"), .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + } + + pub fn create(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, out: *std.ArrayList(g.Kernel), only: ?[]const u8) !*Suite { + const s = try a.create(Suite); + s.* = .{ .a = a }; + errdefer s.destroy(); + for ([_]?Control{null} ++ controls) |x| { + const lib = try compile(a, device, x); + s.libs.append(a, lib) catch |err| { + lib.deinit(); + return err; + }; + try s.add(device, queue, lib, x, out, only); + } + return s; + } + + pub fn destroy(s: *Suite) void { + const a = s.a; + for (s.made.items) |m| m.free(a, m.ptr); + s.made.deinit(a); + for (s.pipes.items) |p| p.deinit(); + s.pipes.deinit(a); + for (s.libs.items) |l| l.deinit(); + s.libs.deinit(a); + for (s.names.items) |n| a.free(n); + s.names.deinit(a); + a.destroy(s); + } +}; diff --git a/zig/tests/contract/metal_idol_vae_rows.zig b/zig/tests/contract/metal_idol_vae_rows.zig new file mode 100644 index 0000000000..10824594ab --- /dev/null +++ b/zig/tests/contract/metal_idol_vae_rows.zig @@ -0,0 +1,437 @@ +//! iv_group_part, iv_group_finish and iv_attention under the contract: fp32 sums against float64. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const mi = @import("metal_idol_common.zig"); +const st = @import("metal_idol_vae_steps.zig"); + +const c = mi.c; +const g = mi.g; +const Range = mi.Range; +const Ends = mi.Ends; +const dec = st.ks; +const ref = st.ref; +const uniform = st.uniform; +const plant = st.plant; +const nan = st.nan; +const inf = st.inf; + +// --- iv_group_part: GroupNorm's chunk sums, of the values or of their squared deviations --------------------------- + +pub const GroupPart = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + centered: bool, + name: []const u8, + shape: c.Shape = undefined, + x: []f32 = &.{}, + stats: [64]f32 = undefined, + ins: [2]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + /// A thread's chunk of pixels, then its group's channels (up to 16) summed in turn. + fn depth(s: c.Shape) usize { + return @min(s.k, dec.chunk) + s.n / 32 + 1; + } + + /// Rows unused (1); n: channels (32 groups); k: pixels, summed in chunks. + pub fn spec(name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = depth, + .out = .f32, + .tiles = .{ .{}, .{ .size = 128, .step = 32, .min = 32 }, .{ .size = dec.chunk } }, + .served = &.{ .{ .rows = 1, .n = 512, .k = 16384 }, .{ .rows = 1, .n = 512, .k = 16500 }, .{ .rows = 1, .n = 512, .k = 65536 }, .{ .rows = 1, .n = 256, .k = 262144 }, .{ .rows = 1, .n = 128, .k = 1048576 } }, + }; + } + + pub fn deinit(t: *GroupPart) void { + t.a.free(t.x); + t.rig.deinit(); + } + pub fn kernel(t: *GroupPart, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(GroupPart); + + pub fn prepare(t: *GroupPart, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.n == 0 or shape.n % 32 != 0 or shape.k == 0) return false; + t.rig.reset(); + t.a.free(t.x); + t.shape = shape; + t.x = try t.a.alloc(f32, shape.k * shape.n); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (fill == .stress) { + // integers: plain sums reach ~2^21 and cancel; deviations within 18 keep squares' sums under 2^24 + for (0..32) |grp| t.stats[grp] = @floatFromInt(rng.intRangeAtMost(i32, -4, 4)); + for (0..shape.k) |p| for (0..shape.n) |ch| { + const grp = ch / (shape.n / 32); + const in = p % dec.chunk; + const half = dec.chunk / 2; + const v: f32 = if (t.centered) + t.stats[grp] + @as(f32, @floatFromInt(rng.intRangeAtMost(i32, -18, 18))) + else if (in < half) @as(f32, if (in % 2 == 0) 127 else 511) else -@as(f32, if ((in - half) % 2 == 0) 127 else 511) + @as(f32, if (in - half < 4) 1 else 0); + t.x[p * shape.n + ch] = v; + }; + } else { + uniform(rng, t.x, -3, 3); + uniform(rng, t.stats[0..32], -0.5, 0.5); + if (fill == .nonfinite) plant(t.x, shape.n); + } + uniform(rng, t.stats[32..], 0.5, 2); + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.x)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(&t.stats)); + t.out = try t.rig.output(t.count() * 4); + return true; + } + + pub fn run(t: *GroupPart) !void { + const p = t.rig.begin(); + dec.groupPart(p.e, t.pipe, ref(t.ins[0]), ref(t.ins[1]), ref(t.out), t.shape.k, t.shape.n, t.centered); + try t.rig.finish(p); + } + + pub fn count(t: *GroupPart) usize { + return 32 * dec.chunks(t.shape.k); + } + + /// Output (group, chunk)'s terms: each pixel's channels of the group, or their squared deviations. + fn terms(t: *const GroupPart, grp: usize, ck: usize, out: []f64) usize { + const size = t.shape.n / 32; + var n: usize = 0; + for (ck * dec.chunk..@min((ck + 1) * dec.chunk, t.shape.k)) |p| for (0..size) |i| { + const v: f64 = t.x[p * t.shape.n + grp * size + i]; + const d = v - t.stats[grp]; + out[n] = if (t.centered) d * d else v; + n += 1; + }; + return n; + } + + pub fn reference(t: *GroupPart, want: []f64, mag: []f64) !void { + const buf = try t.a.alloc(f64, dec.chunk * (t.shape.n / 32)); + defer t.a.free(buf); + const chunks = dec.chunks(t.shape.k); + for (0..32) |grp| for (0..chunks) |ck| { + const n = t.terms(grp, ck, buf); + var s: f64 = 0; + var m: f64 = 0; + for (buf[0..n]) |v| { + s += v; + m += @abs(v); + } + want[grp * chunks + ck] = s; + mag[grp * chunks + ck] = if (std.math.isFinite(m)) m else 0; + }; + } + + pub fn simulate(t: *GroupPart, acc: c.Accum, order: g.Order, out: []f64) !void { + const buf = try t.a.alloc(f64, dec.chunk * (t.shape.n / 32)); + defer t.a.free(buf); + const chunks = dec.chunks(t.shape.k); + for (0..32) |grp| for (0..chunks) |ck| { + const n = t.terms(grp, ck, buf); + out[grp * chunks + ck] = contract.rows.sumIn(acc, order, buf[0..n]); + }; + } +}; + +// --- iv_group_finish: the 32 means, then the 32 rstd, from the chunk sums ------------------------------------------- + +pub const GroupFinish = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + name: []const u8, + shape: c.Shape = undefined, + part: []f32 = &.{}, + ins: [1]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + /// Each lane's chunks / 32 sums in turn, then the simdgroup tree. + fn depth(s: c.Shape) usize { + return (s.k + 31) / 32 + 5; + } + + /// Rows unused (1); n: the 32 groups; k: the chunks. The group size is the widest stage's, 512 / 32. + pub fn spec(name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = depth, + .out = .f32, + .tiles = .{ .{}, .{}, .{ .size = 32 } }, + .served = &.{ .{ .rows = 1, .n = 32, .k = 64 }, .{ .rows = 1, .n = 32, .k = 65 }, .{ .rows = 1, .n = 32, .k = 256 }, .{ .rows = 1, .n = 32, .k = 1024 }, .{ .rows = 1, .n = 32, .k = 4096 } }, + }; + } + + pub fn deinit(t: *GroupFinish) void { + t.a.free(t.part); + t.ends.deinit(); + t.rig.deinit(); + } + pub fn kernel(t: *GroupFinish, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(GroupFinish); + + fn pixels(t: *const GroupFinish) usize { + return t.shape.k * dec.chunk; + } + + pub fn prepare(t: *GroupFinish, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.n != 32 or shape.k == 0) return false; + t.rig.reset(); + t.a.free(t.part); + t.shape = shape; + t.part = try t.a.alloc(f32, 32 * shape.k); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (fill == .stress) { + // integer sums near 2^16 (past fp16's range), the second half of chunks their negatives plus 64 + for (0..32) |grp| for (0..shape.k) |ck| { + const half = shape.k / 2; + const v: f32 = if (ck < half) 65533 + @as(f32, @floatFromInt(ck % 7)) else if (ck < 2 * half) -65533 - @as(f32, @floatFromInt((ck - half) % 7)) + 64 else 64; + t.part[grp * shape.k + ck] = v; + }; + } else { + uniform(rng, t.part, 0, 5000); + if (fill == .nonfinite) plant(t.part, shape.k); + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.part)); + t.out = try t.rig.output(64 * 4); + return true; + } + + pub fn run(t: *GroupFinish) !void { + const p = t.rig.begin(); + for ([_]bool{ false, true }) |centered| dec.groupFinish(p.e, t.pipe, ref(t.ins[0]), ref(t.out), t.pixels(), 512, centered); + try t.rig.finish(p); + } + + pub fn count(_: *GroupFinish) usize { + return 64; + } + + fn outputs(t: *GroupFinish, sums: [32]f64, mags: [32]f64, want: []f64, mag: []f64, ends: *Ends) void { + const n: f64 = @floatFromInt(t.pixels() * 16); + const nf: f64 = @as(f32, @floatCast(n)); + const eps: f64 = dec.eps; + for (0..32) |grp| { + const s = sums[grp]; + const e = c.gamma(@floatFromInt(depth(t.shape)), c.unit(.f32)) * mags[grp]; + const sr = Range{ .lo = s - e, .hi = s + e }; + const mean = (Range{ .lo = sr.lo / nf, .hi = sr.hi / nf }).fp32(mi.div_ulps); + ends.put(want, mag, grp, s / n, mean); + const pos = mean.add(Range.of(eps)).fp32(0.5); + const inv = if (pos.lo > 0) pos.rsqrt().fp32(mi.rsqrt_ulps) else Range.of(1 / @sqrt(s / n + eps)); + ends.put(want, mag, 32 + grp, 1 / @sqrt(s / n + eps), inv); + } + } + + pub fn reference(t: *GroupFinish, want: []f64, mag: []f64) !void { + try t.ends.size(64); + var sums: [32]f64 = @splat(0); + var mags: [32]f64 = @splat(0); + for (0..32) |grp| for (t.part[grp * t.shape.k ..][0..t.shape.k]) |v| { + sums[grp] += v; + mags[grp] += @abs(@as(f64, v)); + }; + t.outputs(sums, mags, want, mag, &t.ends); + } + pub fn bounds(t: *GroupFinish, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *GroupFinish, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } + + pub fn simulate(t: *GroupFinish, acc: c.Accum, order: g.Order, out: []f64) !void { + const terms = try t.a.alloc(f64, t.shape.k); + defer t.a.free(terms); + const n: f64 = @floatFromInt(t.pixels() * 16); + for (0..32) |grp| { + for (terms, t.part[grp * t.shape.k ..][0..t.shape.k]) |*v, p| v.* = p; + const s = contract.rows.sumIn(acc, order, terms); + out[grp] = s / n; + out[32 + grp] = 1 / @sqrt(s / n + dec.eps); + } + } +}; + +// --- iv_attention: softmax(q k^T / sqrt(512)) v, one head of 512, fp32 throughout ---------------------------------- + +pub const Attention = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + name: []const u8, + shape: c.Shape = undefined, + stress: bool = false, + qkv: []f32 = &.{}, + tol: []f64 = &.{}, + ins: [1]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + const D = 512; + const ld = 3 * D; + + /// The scores' 512-deep sums, the keys' sums and one rescale a block of 32 keys. + fn depth(s: c.Shape) usize { + return D + s.k + s.k / 32 + 8; + } + + /// Rows: the queries compared (the first `rows` of all k); n: the head width; k: the positions L. + pub fn spec(name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = depth, + .ops = mi.div_ulps + 2, + .out = .f32, + .tiles = .{ .{ .size = 8 }, .{}, .{ .size = 32 } }, + .served = &.{ .{ .rows = 24, .n = D, .k = 16384 }, .{ .rows = 24, .n = D, .k = 16500 }, .{ .rows = 24, .n = D, .k = 16128 }, .{ .rows = 40, .n = D, .k = 40 } }, + }; + } + + pub fn deinit(t: *Attention) void { + t.a.free(t.qkv); + t.a.free(t.tol); + t.rig.deinit(); + } + pub fn kernel(t: *Attention, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Attention); + + pub fn prepare(t: *Attention, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.n != D or shape.k == 0 or shape.rows == 0 or shape.rows > shape.k) return false; + t.rig.reset(); + t.a.free(t.qkv); + t.shape = shape; + t.stress = fill == .stress; + const L = shape.k; + const padded = std.mem.alignForward(usize, L, 32); + t.qkv = try t.a.alloc(f32, padded * ld); + @memset(t.qkv, 0); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (t.stress) { + // equal scores (q = k = 0); each column's values integers whose partial sums reach ~2^23 and cancel to 4 + const half = L / 2; + for (0..L) |j| { + const at = if (j < half) j else j - half; + const positive: f32 = if (at % 2 == 0) 127 else 2041; + const residual: f32 = if (at < 8 and at % 2 == 0) 1 else 0; + const v: f32 = if (j < half) positive else if (j < 2 * half) residual - positive else 0; + @memset(t.qkv[j * ld + 2 * D ..][0..D], v); + } + } else { + for (0..L) |j| uniform(rng, t.qkv[j * ld ..][0..ld], -1, 1); + if (fill == .nonfinite) { + // values NaN, +Inf, -Inf, and both infinities in one column; a NaN query, an infinite key + for ([_][2]usize{ .{ 0, 0 }, .{ 1, 1 }, .{ 2, 2 }, .{ 3, 3 }, .{ 4, 3 } }, [_]f32{ nan, inf, -inf, inf, -inf }) |p, v| if (p[0] < L) { + t.qkv[p[0] * ld + 2 * D + p[1]] = v; + }; + if (L > 6) t.qkv[1 * ld + 5] = nan; + if (L > 7) t.qkv[7 * ld + D + 3] = inf; + } + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.qkv)); + t.out = try t.rig.output(padded * D * 4); + return true; + } + + pub fn run(t: *Attention) !void { + const p = t.rig.begin(); + dec.attend(p.e, t.pipe, ref(t.ins[0]), ref(t.out), t.shape.k); + try t.rig.finish(p); + } + + pub fn count(t: *Attention) usize { + return t.shape.rows * D; + } + + /// Each compared query's output in float64, and its bound: the scores' error moved through exp, then the sums'. + pub fn reference(t: *Attention, want: []f64, mag: []f64) !void { + const L = t.shape.k; + const scale: f64 = @as(f32, 1.0 / @sqrt(@as(f32, D))); + const u = c.unit(.f32); + t.a.free(t.tol); + t.tol = try t.a.alloc(f64, t.count()); + const s = try t.a.alloc(f64, L); + defer t.a.free(s); + const sabs = try t.a.alloc(f64, L); + defer t.a.free(sabs); + var out: [D]f64 = undefined; + var spread: [D]f64 = undefined; + for (0..t.shape.rows) |i| { + const q = t.qkv[i * ld ..][0..D]; + var m: f64 = -std.math.inf(f64); + var nan_row = false; + var worst: f64 = 0; + for (s, sabs, 0..) |*sj, *aj, j| { + const k = t.qkv[j * ld + D ..][0..D]; + var dot: f64 = 0; + var mag_: f64 = 0; + for (q, k) |a, b| { + dot += @as(f64, a) * b; + mag_ += @abs(@as(f64, a) * b); + } + sj.* = dot * scale; + aj.* = mag_ * scale; + nan_row = nan_row or std.math.isNan(sj.*); + m = @max(m, sj.*); + worst = @max(worst, aj.*); + } + if (nan_row) m = std.math.nan(f64); + // a score's error: its sums and the scale, then the subtraction of the maximum + const es = c.gamma(D + 4, u) * worst + 2 * u * (@abs(m) + worst); + const rel = 2 * es + (mi.exp_ulps + 1) * 0x1p-24; + var l: f64 = 0; + @memset(&out, 0); + @memset(&spread, 0); + for (s, 0..) |sj, j| { + const p = @exp(sj - m); + l += p; + for (&out, &spread, t.qkv[j * ld + 2 * D ..][0..D]) |*o, *sp, v| { + o.* += p * v; + sp.* += p * @abs(@as(f64, v)); + } + } + const sums = c.gamma(@floatFromInt(depth(t.shape)), u); + for (out, spread, 0..) |o, sp, e| { + const at = i * D + e; + want[at] = o / l; + mag[at] = sp / l; + t.tol[at] = (2 * rel + 2 * sums) * sp / l; + } + } + } + + pub fn bounds(t: *Attention, want: []const f64, _: []const f64, allow: []f64) !void { + for (allow, want, t.tol) |*l, w, tl| l.* = if (t.stress) c.boundExact(spec(t.name), w) else if (std.math.isFinite(w)) tl + c.boundExact(spec(t.name), w) else 0; + } + + pub fn simulate(t: *Attention, acc: c.Accum, order: g.Order, out: []f64) !void { + if (!t.stress) return error.AttentionStressNeedsEqualScores; + const L = t.shape.k; + const terms = try t.a.alloc(f64, L); + defer t.a.free(terms); + const ones = try t.a.alloc(f64, L); + defer t.a.free(ones); + @memset(ones, 1); + const l = contract.rows.sumIn(acc, order, ones); + for (0..D) |e| { + for (terms, 0..) |*v, j| v.* = t.qkv[j * ld + 2 * D + e]; + const o = contract.rows.sumIn(acc, order, terms) / l; + for (0..t.shape.rows) |i| out[i * D + e] = o; + } + } +}; diff --git a/zig/tests/contract/metal_idol_vae_steps.zig b/zig/tests/contract/metal_idol_vae_steps.zig new file mode 100644 index 0000000000..5fb097959f --- /dev/null +++ b/zig/tests/contract/metal_idol_vae_steps.zig @@ -0,0 +1,578 @@ +//! iv_conv_in, iv_conv, iv_norm and iv_tokens under the contract: Idol's image encoder's fp32 steps against float64. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const id = @import("metal_idol_common.zig"); + +const c = id.c; +const g = id.g; +const Range = id.Range; +const Ends = id.Ends; +pub const ks = tf.idol.vae_kernels; + +pub fn ref(r: rig_.Ref) ks.Ref { + return .{ .buf = r.buf, .off = r.off }; +} + +pub fn uniform(rng: std.Random, out: []f32, lo: f32, hi: f32) void { + for (out) |*v| v.* = lo + rng.float(f32) * (hi - lo); +} + +pub const nan = std.math.nan(f32); +pub const inf = std.math.inf(f32); + +/// NaN, +Inf, -Inf at three places, and +Inf with -Inf at two places `apart` further on (one row holding both). +pub fn plant(x: []f32, apart: usize) void { + for ([_]usize{ 0, 1, 2 }, [_]f32{ nan, inf, -inf }) |at, v| x[(at * 7) % x.len] = v; + if (x.len > 3 * apart + 1) { + x[3 * apart] = inf; + x[3 * apart + 1] = -inf; + } +} + +// --- iv_conv_in: conv3x3(2 x - 1) + b of the picture, 3 channels to 128 ---------------------------------------------- + +pub const ConvIn = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + name: []const u8, + shape: c.Shape = undefined, + x: []f32 = &.{}, + w: []f32 = &.{}, + b: []f32 = &.{}, + ins: [3]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + /// The 27 taps' products, 2 x - 1's rounding and the bias. + fn depth(_: c.Shape) usize { + return 27 + 2; + } + + /// Rows: the picture's rows; n: the 128 output channels; k: its columns (the encoder's picture is 1024 x 1024). + pub fn spec(name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = depth, + .out = .f32, + .tiles = .{ .{ .size = 2 }, .{}, .{ .size = 2 } }, + .served = &.{ .{ .rows = 12, .n = 128, .k = 1024 }, .{ .rows = 1024, .n = 128, .k = 12 }, .{ .rows = 37, .n = 128, .k = 41 } }, + }; + } + + pub fn deinit(t: *ConvIn) void { + t.free(); + t.rig.deinit(); + } + fn free(t: *ConvIn) void { + inline for (.{ &t.x, &t.w, &t.b }) |p| { + t.a.free(p.*); + p.* = &.{}; + } + } + pub fn kernel(t: *ConvIn, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(ConvIn); + + pub fn prepare(t: *ConvIn, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.k == 0 or shape.n != 128) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + t.x = try t.a.alloc(f32, shape.rows * shape.k * 3); + t.w = try t.a.alloc(f32, 27 * 128); + t.b = try t.a.alloc(f32, 128); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (fill == .stress) { + // 2 x - 1 an odd integer (127 or 2041), channel 1 two more, channel 2 zero; weights 509, -509, 1 cancel + for (0..shape.rows * shape.k) |p| { + const v: f32 = if (p % 2 == 0) 2041 else 127; + t.x[p * 3] = (v + 1) / 2; + t.x[p * 3 + 1] = (v + 3) / 2; + t.x[p * 3 + 2] = 0.5; + } + for (0..27) |r| @memset(t.w[r * 128 ..][0..128], switch (r % 3) { + 0 => 509, + 1 => -509, + else => 1, + }); + @memset(t.b, 0); + } else { + uniform(rng, t.x, 0, 1); + uniform(rng, t.w, -0.5, 0.5); + uniform(rng, t.b, -1, 1); + if (fill == .nonfinite) plant(t.x, shape.k * 3); + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.x)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.w)); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.b)); + t.out = try t.rig.output(t.count() * 4); + return true; + } + + pub fn run(t: *ConvIn) !void { + const p = t.rig.begin(); + ks.convIn(p.e, t.pipe, ref(t.ins[0]), ref(t.ins[1]), ref(t.ins[2]), ref(t.out), t.shape.rows, t.shape.k); + try t.rig.finish(p); + } + + pub fn count(t: *ConvIn) usize { + return t.shape.rows * t.shape.k * 128; + } + + /// Output (pixel, o)'s 27 products in the kernel's order (row, column, channel), zero past the picture's edge. + fn terms(t: *const ConvIn, px: usize, o: usize, out: *[27]f64) void { + const py: isize = @intCast(px / t.shape.k); + const pxx: isize = @intCast(px % t.shape.k); + for (0..3) |ky| for (0..3) |kx| for (0..3) |ch| { + const iy = py + @as(isize, @intCast(ky)) - 1; + const ix = pxx + @as(isize, @intCast(kx)) - 1; + const at = (ky * 3 + kx) * 3 + ch; + out[at] = if (iy < 0 or ix < 0 or iy >= t.shape.rows or ix >= t.shape.k) 0 else blk: { + const v: f64 = t.x[(@as(usize, @intCast(iy)) * t.shape.k + @as(usize, @intCast(ix))) * 3 + ch]; + break :blk (2 * v - 1) * t.w[at * 128 + o]; + }; + }; + } + + pub fn reference(t: *ConvIn, want: []f64, mag: []f64) !void { + var buf: [27]f64 = undefined; + for (0..t.shape.rows * t.shape.k) |px| for (0..128) |o| { + t.terms(px, o, &buf); + var s: f64 = 0; + var m: f64 = 0; + for (buf) |v| { + s += v; + m += @abs(v); + } + want[px * 128 + o] = s + t.b[o]; + m += @abs(@as(f64, t.b[o])); + mag[px * 128 + o] = if (std.math.isFinite(m)) m else 0; + }; + } + + pub fn simulate(t: *ConvIn, acc: c.Accum, order: g.Order, out: []f64) !void { + var buf: [27]f64 = undefined; + for (0..t.shape.rows * t.shape.k) |px| for (0..128) |o| { + t.terms(px, o, &buf); + out[px * 128 + o] = contract.rows.sumIn(acc, order, &buf) + t.b[o]; + }; + } +}; + +// --- iv_conv: a 1x1 or stride-2 3x3 convolution as an implicit GEMM, fp32 sums, + bias (+ a residual) --------------- + +/// A form the encoder launches: its taps (1, or 9 at stride 2), a residual, and its served shapes. +pub const Form = struct { label: []const u8, taps: usize = 1, add: bool = false, served: []const c.Shape }; + +/// A 1x1's rows are pixels (fewer than the engine's: rows only tile); a 3x3's are output rows of 9 columns. +pub const conv_forms = [_]Form{ + .{ .label = "iv_conv 1x1", .served = &.{ .{ .rows = 1100, .n = 1536, .k = 512 }, .{ .rows = 1100, .n = 256, .k = 128 }, .{ .rows = 1100, .n = 512, .k = 256 }, .{ .rows = 1100, .n = 64, .k = 64 } } }, + .{ .label = "iv_conv 1x1 + residual", .add = true, .served = &.{.{ .rows = 1100, .n = 512, .k = 512 }} }, + .{ .label = "iv_conv 3x3 stride 2", .taps = 9, .served = &.{ .{ .rows = 8, .n = 128, .k = 128 }, .{ .rows = 8, .n = 256, .k = 256 }, .{ .rows = 8, .n = 512, .k = 512 } } }, +}; + +const out_cols = 9; + +pub const Conv = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + form: Form, + name: []const u8, + shape: c.Shape = undefined, + x: []f32 = &.{}, + w: []f32 = &.{}, + b: []f32 = &.{}, + r: []f32 = &.{}, + ins: [4]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + /// Each output sums taps * cin products, then the bias and the residual. + fn depth1(s: c.Shape) usize { + return s.k + 2; + } + fn depth9(s: c.Shape) usize { + return 9 * s.k + 2; + } + + /// Rows: pixels (1x1) or output rows (3x3); n: output channels; k: input channels. + pub fn spec(name: []const u8, f: Form) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = if (f.taps == 9) depth9 else depth1, + .out = .f32, + .tiles = .{ if (f.taps == 9) .{ .size = 8 } else .{ .size = 64 }, .{ .size = 64 }, .{ .size = 32, .step = 32, .min = 32 } }, + .served = f.served, + }; + } + + pub fn deinit(t: *Conv) void { + t.free(); + t.rig.deinit(); + } + fn free(t: *Conv) void { + inline for (.{ &t.x, &t.w, &t.b, &t.r }) |p| { + t.a.free(p.*); + p.* = &.{}; + } + } + pub fn kernel(t: *Conv, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name, t.form), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Conv); + + fn ldw(t: *const Conv) usize { + return std.mem.alignForward(usize, t.shape.n, 64); + } + fn cols(t: *const Conv) usize { + return if (t.form.taps == 9) out_cols else 1; + } + fn pixels(t: *const Conv) usize { + return t.shape.rows * t.cols(); + } + fn inRows(t: *const Conv) usize { + return if (t.form.taps == 9) 2 * t.shape.rows else t.shape.rows; + } + fn inCols(t: *const Conv) usize { + return if (t.form.taps == 9) 2 * out_cols else 1; + } + + pub fn prepare(t: *Conv, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.k == 0 or shape.k % 32 != 0) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + const L = t.ldw(); + const K = t.form.taps * shape.k; + const in_px = t.inRows() * t.inCols(); + t.x = try t.a.alloc(f32, in_px * shape.k); + t.w = try t.a.alloc(f32, K * L); + t.b = try t.a.alloc(f32, shape.n); + t.r = try t.a.alloc(f32, if (t.form.add) t.pixels() * shape.n else 0); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + @memset(t.w, 0); + if (fill == .stress) { + // the channels' second half mirroring the first, weights 1 then -1: partial sums reach ~2^18 and cancel + const half = shape.k / 2; + for (0..in_px) |p| for (0..shape.k) |ci| { + const at = if (ci < half) ci else ci - half; + const v: f32 = if (at % 2 == 0) 127 else 2041; + t.x[p * shape.k + ci] = if (ci < half) v else v + @as(f32, if (at < 4) 1 else 0); + }; + for (0..K) |k| for (0..shape.n) |o| { + t.w[k * L + o] = if (k % shape.k < half) 1 else -1; + }; + @memset(t.b, 0); + for (t.r) |*v| v.* = @floatFromInt(rng.intRangeAtMost(i32, -3, 3)); + } else { + uniform(rng, t.x, -2, 2); + for (0..K) |k| uniform(rng, t.w[k * L ..][0..shape.n], -0.1, 0.1); + uniform(rng, t.b, -1, 1); + uniform(rng, t.r, -3, 3); + if (fill == .nonfinite) plant(t.x, shape.k); + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.x)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.w)); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.b)); + if (t.form.add) t.ins[3] = try t.rig.input(std.mem.sliceAsBytes(t.r)); + t.out = try t.rig.output(t.count() * 4); + return true; + } + + pub fn run(t: *Conv) !void { + const p = t.rig.begin(); + const s: ks.ConvShape = .{ .rows = t.shape.rows, .cols = t.cols(), .in_rows = t.inRows(), .in_cols = t.inCols(), .cin = t.shape.k, .cout = t.shape.n, .ldw = t.ldw(), .taps = t.form.taps }; + ks.conv(p.e, t.pipe, ref(t.ins[0]), ref(t.ins[1]), ref(t.ins[2]), if (t.form.add) ref(t.ins[3]) else null, ref(t.out), s); + try t.rig.finish(p); + } + + pub fn count(t: *Conv) usize { + return t.pixels() * t.shape.n; + } + + /// Output (p, o)'s products in the kernel's order, tap then input channel; a tap past the right or lower edge is 0. + fn terms(t: *const Conv, p: usize, o: usize, out: []f64) void { + const K = t.shape.k; + const oy = p / t.cols(); + const ox = p % t.cols(); + const stride: usize = if (t.form.taps == 9) 2 else 1; + for (0..t.form.taps) |tap| { + const iy = oy * stride + tap / 3; + const ix = ox * stride + tap % 3; + for (0..K) |ci| { + const inside = iy < t.inRows() and ix < t.inCols(); + out[tap * K + ci] = if (inside) @as(f64, t.x[(iy * t.inCols() + ix) * K + ci]) * t.w[(tap * K + ci) * t.ldw() + o] else 0; + } + } + } + + pub fn reference(t: *Conv, want: []f64, mag: []f64) !void { + const buf = try t.a.alloc(f64, t.form.taps * t.shape.k); + defer t.a.free(buf); + for (0..t.pixels()) |p| for (0..t.shape.n) |o| { + t.terms(p, o, buf); + var sum: f64 = 0; + var m: f64 = 0; + for (buf) |v| { + sum += v; + m += @abs(v); + } + const i = p * t.shape.n + o; + const r: f64 = if (t.form.add) t.r[i] else 0; + want[i] = sum + t.b[o] + r; + m += @abs(@as(f64, t.b[o])) + @abs(r); + mag[i] = if (std.math.isFinite(m)) m else 0; + }; + } + + pub fn simulate(t: *Conv, acc: c.Accum, order: g.Order, out: []f64) !void { + const buf = try t.a.alloc(f64, t.form.taps * t.shape.k); + defer t.a.free(buf); + for (0..t.pixels()) |p| for (0..t.shape.n) |o| { + t.terms(p, o, buf); + const i = p * t.shape.n + o; + out[i] = contract.rows.sumIn(acc, order, buf) + t.b[o] + (if (t.form.add) t.r[i] else 0); + }; + } +}; + +// --- iv_norm: GroupNorm by given statistics, then SiLU or not ------------------------------------------------------- + +pub const Norm = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + silu: bool, + name: []const u8, + shape: c.Shape = undefined, + x: []f32 = &.{}, + stats: [64]f32 = undefined, + gamma: []f32 = &.{}, + beta: []f32 = &.{}, + ins: [4]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + /// Rows: pixels; n: channels (32 groups of n / 32). + pub fn spec(name: []const u8, silu: bool) c.Spec { + return .{ + .name = name, + .accum = .none, + .depth = id.noSums, + .ops = 4 + id.exp_ulps + id.div_ulps + 1, + .out = .f32, + .tiles = .{ .{ .size = 8 }, .{ .size = 128, .step = 32, .min = 32 }, .{} }, + .served = if (silu) &.{ .{ .rows = 4100, .n = 128, .k = 1 }, .{ .rows = 4100, .n = 256, .k = 1 }, .{ .rows = 4100, .n = 512, .k = 1 } } else &.{.{ .rows = 4100, .n = 512, .k = 1 }}, + }; + } + + pub fn deinit(t: *Norm) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *Norm) void { + inline for (.{ &t.x, &t.gamma, &t.beta }) |p| { + t.a.free(p.*); + p.* = &.{}; + } + } + pub fn kernel(t: *Norm, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name, t.silu), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Norm); + + pub fn prepare(t: *Norm, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.n % 32 != 0) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + t.x = try t.a.alloc(f32, shape.rows * shape.n); + t.gamma = try t.a.alloc(f32, shape.n); + t.beta = try t.a.alloc(f32, shape.n); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + // values stay within +-60 after the norm, short of fp32 exp's overflow in the SiLU + uniform(rng, t.x, -4, 4); + uniform(rng, t.stats[0..32], -1, 1); + uniform(rng, t.stats[32..], 0.2, 3); + uniform(rng, t.gamma, -2, 2); + uniform(rng, t.beta, -1, 1); + if (fill == .nonfinite) plant(t.x, shape.n); + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.x)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(&t.stats)); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.gamma)); + t.ins[3] = try t.rig.input(std.mem.sliceAsBytes(t.beta)); + t.out = try t.rig.output(t.count() * 4); + return true; + } + + pub fn run(t: *Norm) !void { + const p = t.rig.begin(); + ks.normed(p.e, t.pipe, ref(t.ins[0]), ref(t.ins[1]), ref(t.ins[2]), ref(t.ins[3]), ref(t.out), t.shape.rows, t.shape.n, t.silu); + try t.rig.finish(p); + } + + pub fn count(t: *Norm) usize { + return t.shape.rows * t.shape.n; + } + + pub fn reference(t: *Norm, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + const size = t.shape.n / 32; + for (t.x, 0..) |xv, i| { + const ch = i % t.shape.n; + const grp = ch / size; + const x: f64 = xv; + const m: f64 = t.stats[grp]; + const r: f64 = t.stats[32 + grp]; + const gm: f64 = t.gamma[ch]; + const b: f64 = t.beta[ch]; + const v = (x - m) * r * gm + b; + var vr = Range.of(x - m).fp32(0.5).mul(Range.of(r)).fp32(0.5).mul(Range.of(gm)).fp32(0.5).add(Range.of(b)).fp32(0.5); + var w = v; + if (t.silu) { + w = v / (1 + @exp(-v)); + const e = (Range{ .lo = @exp(-vr.hi), .hi = @exp(-vr.lo) }).fp32(id.exp_ulps); + vr = vr.mul(e.add(Range.of(1)).fp32(0.5).recip()).fp32(id.div_ulps); + } + t.ends.put(want, mag, i, w, vr); + } + } + pub fn bounds(t: *Norm, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *Norm, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } +}; + +// --- iv_tokens: the posterior's sample, patchified and normalized by the batch norm's statistics -------------------- + +pub const Tokens = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + name: []const u8, + shape: c.Shape = undefined, + moments: []f32 = &.{}, + noise: []f32 = &.{}, + m: [128]f32 = undefined, + s: [128]f32 = undefined, + ins: [4]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + const pad = 5; + + /// Rows: the latent grid's rows; n: the 128 channels of a row; k: the grid's columns. + pub fn spec(name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .none, + .depth = id.noSums, + .ops = id.exp_ulps + id.div_ulps + 3, + .out = .f32, + .tiles = .{ .{ .size = 2, .step = 2, .min = 2 }, .{}, .{ .size = 2, .step = 2, .min = 2 } }, + .served = &.{ .{ .rows = 128, .n = 128, .k = 128 }, .{ .rows = 6, .n = 128, .k = 10 } }, + }; + } + + pub fn deinit(t: *Tokens) void { + t.a.free(t.moments); + t.a.free(t.noise); + t.ends.deinit(); + t.rig.deinit(); + } + pub fn kernel(t: *Tokens, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Tokens); + + fn latents(t: *const Tokens) usize { + return t.shape.rows * t.shape.k; + } + + pub fn prepare(t: *Tokens, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.n != 128 or shape.rows == 0 or shape.k == 0 or shape.rows % 2 != 0 or shape.k % 2 != 0) return false; + t.rig.reset(); + t.a.free(t.moments); + t.a.free(t.noise); + t.shape = shape; + const px = t.latents(); + t.moments = try t.a.alloc(f32, px * 64); + t.noise = try t.a.alloc(f32, px * 32); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (0..px) |p| { + uniform(rng, t.moments[p * 64 ..][0..32], -3, 3); + // logvars past both of the clamp's ends (-30, 20) as well as inside them + uniform(rng, t.moments[p * 64 + 32 ..][0..32], -40, 30); + } + uniform(rng, t.noise, -3, 3); + uniform(rng, &t.m, -1, 1); + uniform(rng, &t.s, 0.5, 4); + if (fill == .nonfinite) { + plant(t.moments, 64); + plant(t.noise, 32); + t.moments[5 * 64 + 40] = nan; + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.moments)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.noise)); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(&t.m)); + t.ins[3] = try t.rig.input(std.mem.sliceAsBytes(&t.s)); + t.out = try t.rig.output(t.count() * 4); + return true; + } + + pub fn run(t: *Tokens) !void { + const p = t.rig.begin(); + ks.tokens(p.e, t.pipe, ref(t.ins[0]), ref(t.ins[1]), ref(t.ins[2]), ref(t.ins[3]), ref(t.out), t.shape.rows, t.shape.k, pad); + try t.rig.finish(p); + } + + pub fn count(t: *Tokens) usize { + return (pad + t.latents() / 4) * 128; + } + + /// torch.clamp(x, -30, 20): NaN stays NaN. + fn clamp(x: f64) f64 { + return if (std.math.isNan(x)) x else std.math.clamp(x, -30, 20); + } + + pub fn reference(t: *Tokens, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + for (0..pad * 128) |i| t.ends.put(want, mag, i, 0, Range.of(0)); + const w2 = t.shape.k / 2; + for (0..t.latents() / 4) |patch| for (0..128) |k| { + const ch = k / 4; + const py = (patch / w2) * 2 + (k % 4) / 2; + const px = (patch % w2) * 2 + k % 2; + const p = py * t.shape.k + px; + const mean: f64 = t.moments[p * 64 + ch]; + const lv = clamp(t.moments[p * 64 + 32 + ch]); + const n: f64 = t.noise[p * 32 + ch]; + const sd = @exp(0.5 * lv); + const z = mean + sd * n; + const v = (z - t.m[k]) / t.s[k]; + const er = (Range{ .lo = @exp(0.5 * lv), .hi = @exp(0.5 * lv) }).fp32(id.exp_ulps); + const zr = er.mul(Range.of(n)).fp32(0.5).add(Range.of(mean)).fp32(0.5); + const vr = zr.sub(Range.of(t.m[k])).fp32(0.5).mul(Range.of(1 / @as(f64, t.s[k]))).fp32(id.div_ulps); + t.ends.put(want, mag, (pad + patch) * 128 + k, v, vr); + }; + } + pub fn bounds(t: *Tokens, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *Tokens, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } +}; diff --git a/zig/tests/contract/metal_idol_vae_wino.zig b/zig/tests/contract/metal_idol_vae_wino.zig new file mode 100644 index 0000000000..9e55840dae --- /dev/null +++ b/zig/tests/contract/metal_idol_vae_wino.zig @@ -0,0 +1,261 @@ +//! iv_wino under the contract: Winograd F(2x2, 3x3) against the float64 convolution, in Winograd's own float64 form. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const id = @import("metal_idol_common.zig"); +const st = @import("metal_idol_vae_steps.zig"); + +const c = id.c; +const g = id.g; +const ks = st.ks; +const ref = st.ref; +const uniform = st.uniform; +const plant = st.plant; + +/// B^T, G and A^T of F(2x2, 3x3). +const BT = [4][4]f64{ .{ 1, 0, -1, 0 }, .{ 0, 1, 1, 0 }, .{ 0, -1, 1, 0 }, .{ 0, 1, 0, -1 } }; +const G = [4][3]f64{ .{ 1, 0, 0 }, .{ 0.5, 0.5, 0.5 }, .{ 0.5, -0.5, 0.5 }, .{ 0, 0, 1 } }; +const AT = [2][4]f64{ .{ 1, 1, 1, 0 }, .{ 0, 1, -1, -1 } }; + +/// A 3x3 form the encoder launches: with a residual or not. +pub const Form = struct { label: []const u8, add: bool = false, served: []const c.Shape }; + +/// The pictures here are 8 x 18 (the encoder's up to 1024 x 1024): the sums' depth is the input channels alone. +pub const forms = [_]Form{ + .{ .label = "iv_wino 3x3", .served = &.{ .{ .rows = 8, .n = 128, .k = 128 }, .{ .rows = 8, .n = 256, .k = 128 }, .{ .rows = 8, .n = 256, .k = 256 }, .{ .rows = 8, .n = 512, .k = 256 }, .{ .rows = 8, .n = 512, .k = 512 }, .{ .rows = 8, .n = 64, .k = 512 } } }, + .{ .label = "iv_wino 3x3 + residual", .add = true, .served = &.{ .{ .rows = 8, .n = 128, .k = 128 }, .{ .rows = 8, .n = 256, .k = 256 }, .{ .rows = 8, .n = 512, .k = 512 } } }, +}; + +const cols = 18; + +pub const Wino = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + form: Form, + name: []const u8, + shape: c.Shape = undefined, + x: []f32 = &.{}, + /// the weights as released [cout, cin, 9], and their transform as the decoder packs it [16, cin, ldw] + w: []f32 = &.{}, + u: []f32 = &.{}, + b: []f32 = &.{}, + r: []f32 = &.{}, + ins: [4]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + /// The transforms' roundings (2 in B^T d B, 1 in G g G^T), cin products, A^T M A's 4, the bias and the residual. + fn depth(s: c.Shape) usize { + return s.k + 3 + 4 + 2; + } + + /// Rows: the output picture's rows (18 columns); n: output channels; k: input channels. + pub fn spec(name: []const u8, f: Form) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = depth, + .out = .f32, + .tiles = .{ .{ .size = 4, .step = 2, .min = 2 }, .{ .size = 64 }, .{ .size = 16, .step = 16, .min = 16 } }, + .served = f.served, + }; + } + + pub fn deinit(t: *Wino) void { + t.free(); + t.rig.deinit(); + } + fn free(t: *Wino) void { + inline for (.{ &t.x, &t.w, &t.u, &t.b, &t.r }) |p| { + t.a.free(p.*); + p.* = &.{}; + } + } + pub fn kernel(t: *Wino, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name, t.form), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Wino); + + fn ldw(t: *const Wino) usize { + return std.mem.alignForward(usize, t.shape.n, 64); + } + fn pixels(t: *const Wino) usize { + return t.shape.rows * cols; + } + + pub fn prepare(t: *Wino, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.rows % 2 != 0 or shape.n == 0 or shape.k == 0 or shape.k % 16 != 0) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + const K = shape.k; + const L = t.ldw(); + const in_px = t.pixels(); + t.x = try t.a.alloc(f32, in_px * K); + t.w = try t.a.alloc(f32, shape.n * K * 9); + t.u = try t.a.alloc(f32, 16 * K * L); + t.b = try t.a.alloc(f32, shape.n); + t.r = try t.a.alloc(f32, if (t.form.add) t.pixels() * shape.n else 0); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (fill == .stress) { + // equal pixels, mirrored channel halves, weights 1 then -1: quarters throughout, sums cancel + const half = K / 2; + for (0..in_px) |p| for (0..K) |ci| { + const pos = if (ci < half) ci else ci - half; + const v: f32 = if (pos % 2 == 0) 127 else 1021; + t.x[p * K + ci] = if (ci < half) v else v + @as(f32, if (pos < 4) 1 else 0); + }; + for (0..shape.n) |o| for (0..K) |ci| @memset(t.w[(o * K + ci) * 9 ..][0..9], if (ci < half) 1 else -1); + @memset(t.b, 0); + for (t.r) |*v| v.* = @floatFromInt(rng.intRangeAtMost(i32, -3, 3)); + } else { + uniform(rng, t.x, -2, 2); + uniform(rng, t.w, -0.1, 0.1); + uniform(rng, t.b, -1, 1); + uniform(rng, t.r, -3, 3); + if (fill == .nonfinite) plant(t.x, K); + } + @memset(t.u, 0); + for (0..shape.n) |o| for (0..K) |ci| { + for (ks.transform(t.w[(o * K + ci) * 9 ..][0..9]), 0..) |v, pt| t.u[(pt * K + ci) * L + o] = v; + }; + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.x)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.u)); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.b)); + if (t.form.add) t.ins[3] = try t.rig.input(std.mem.sliceAsBytes(t.r)); + t.out = try t.rig.output(t.count() * 4); + return true; + } + + pub fn run(t: *Wino) !void { + const p = t.rig.begin(); + const s: ks.ConvShape = .{ .rows = t.shape.rows, .cols = cols, .in_rows = t.shape.rows, .in_cols = cols, .cin = t.shape.k, .cout = t.shape.n, .ldw = t.ldw() }; + ks.winograd(p.e, t.pipe, ref(t.ins[0]), ref(t.ins[1]), ref(t.ins[2]), if (t.form.add) ref(t.ins[3]) else null, ref(t.out), s); + try t.rig.finish(p); + } + + pub fn count(t: *Wino) usize { + return t.pixels() * t.shape.n; + } + + /// The input at picture position (y, x), zero in the padding. + fn at(t: *const Wino, y: isize, x: isize, ci: usize) f64 { + if (y < 0 or x < 0 or y >= t.shape.rows or x >= cols) return 0; + const uy: usize = @intCast(y); + const ux: usize = @intCast(x); + return t.x[(uy * cols + ux) * t.shape.k + ci]; + } + + /// B^T d B of tile (ty, tx)'s patch for channel ci, and the same of |d| (the magnitude its roundings act on). + fn tileV(t: *const Wino, ty: usize, tx: usize, ci: usize, v: *[16]f64, va: *[16]f64) void { + var d: [4][4]f64 = undefined; + for (0..4) |i| for (0..4) |j| { + d[i][j] = t.at(@as(isize, @intCast(ty * 2 + i)) - 1, @as(isize, @intCast(tx * 2 + j)) - 1, ci); + }; + for (0..4) |i| for (0..4) |j| { + var s: f64 = 0; + var m: f64 = 0; + for (0..4) |a| for (0..4) |b| { + const k = BT[i][a] * BT[j][b]; + if (k == 0) continue; + s += k * d[a][b]; + m += @abs(k * d[a][b]); + }; + v[i * 4 + j] = s; + va[i * 4 + j] = m; + }; + } + + /// G g G^T exactly, and the same of |g|. + fn exactU(t: *const Wino, o: usize, ci: usize, u: *[16]f64, ua: *[16]f64) void { + const w = t.w[(o * t.shape.k + ci) * 9 ..][0..9]; + for (0..4) |i| for (0..4) |j| { + var s: f64 = 0; + var m: f64 = 0; + for (0..3) |a| for (0..3) |b| { + const k = G[i][a] * G[j][b]; + s += k * w[a * 3 + b]; + m += @abs(k * w[a * 3 + b]); + }; + u[i * 4 + j] = s; + ua[i * 4 + j] = m; + }; + } + + /// Each output as A^T M A, M each point's products V U summed by `sum` (U exact or as packed); `mag` its bound's. + fn each(t: *const Wino, packed_u: bool, sum: anytype, out: []f64, mag: ?[]f64) !void { + const K = t.shape.k; + const vs = try t.a.alloc([16]f64, K); + defer t.a.free(vs); + const vas = try t.a.alloc([16]f64, K); + defer t.a.free(vas); + const terms = try t.a.alloc(f64, K); + defer t.a.free(terms); + const us = try t.a.alloc([16]f64, K); + defer t.a.free(us); + const uas = try t.a.alloc([16]f64, K); + defer t.a.free(uas); + for (0..t.shape.rows / 2) |ty| for (0..cols / 2) |tx| { + for (0..K) |ci| t.tileV(ty, tx, ci, &vs[ci], &vas[ci]); + for (0..t.shape.n) |o| { + for (0..K) |ci| t.exactU(o, ci, &us[ci], &uas[ci]); + var m: [16]f64 = undefined; + var ma: [16]f64 = undefined; + for (0..16) |pt| { + var mm: f64 = 0; + for (0..K) |ci| { + const uv = if (packed_u) @as(f64, t.u[(pt * K + ci) * t.ldw() + o]) else us[ci][pt]; + terms[ci] = vs[ci][pt] * uv; + mm += vas[ci][pt] * uas[ci][pt]; + } + m[pt] = sum.of(terms); + ma[pt] = mm; + } + for (0..2) |rr| for (0..2) |cc| { + var y: f64 = 0; + var ya: f64 = 0; + for (0..4) |i| for (0..4) |j| { + const k = AT[rr][i] * AT[cc][j]; + if (k == 0) continue; + y += k * m[i * 4 + j]; + ya += @abs(k) * ma[i * 4 + j]; + }; + const px = (ty * 2 + rr) * cols + tx * 2 + cc; + const idx = px * t.shape.n + o; + const res: f64 = if (t.form.add) t.r[idx] else 0; + out[idx] = y + t.b[o] + res; + if (mag) |mg| { + const total = ya + @abs(@as(f64, t.b[o])) + @abs(res); + mg[idx] = if (std.math.isFinite(total)) total else 0; + } + }; + } + }; + } + + const Exact = struct { + fn of(_: Exact, terms: []const f64) f64 { + var s: f64 = 0; + for (terms) |v| s += v; + return s; + } + }; + const Replay = struct { + acc: c.Accum, + order: g.Order, + fn of(r: Replay, terms: []const f64) f64 { + return contract.rows.sumIn(r.acc, r.order, terms); + } + }; + + pub fn reference(t: *Wino, want: []f64, mag: []f64) !void { + try t.each(false, Exact{}, want, mag); + } + + pub fn simulate(t: *Wino, acc: c.Accum, order: g.Order, out: []f64) !void { + try t.each(true, Replay{ .acc = acc, .order = order }, out, null); + } +}; diff --git a/zig/tests/contract/metal_limits.zig b/zig/tests/contract/metal_limits.zig index ece8073ecc..ff3088ccc0 100644 --- a/zig/tests/contract/metal_limits.zig +++ b/zig/tests/contract/metal_limits.zig @@ -102,6 +102,8 @@ fn launchOf(a: std.mem.Allocator, device: mtl.Device, d: limits.MetalLaunch) !La } else if (std.mem.startsWith(u8, d.name, "moss image ")) { text = sources.moss_image; frags = true; + } else if (std.mem.startsWith(u8, d.name, "idol vae ")) { + text = sources.idol_vae; } else if (std.mem.startsWith(u8, d.name, "idol ")) { text = sources.idol; frags = true; diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index 9799191182..4de4d69394 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -429,6 +429,8 @@ pub fn main(init: std.process.Init) !u8 { defer moss_image.destroy(); const idol = try @import("metal_idol.zig").Suite.create(a, device, queue, &kernels, o.only); defer idol.destroy(); + const idol_vae = try @import("metal_idol_vae.zig").Suite.create(a, device, queue, &kernels, o.only); + defer idol_vae.destroy(); var t: g.Tally = .{}; const t0 = std.Io.Clock.awake.now(init.io); if (o.gates) { diff --git a/zig/tests/idol_vae.zig b/zig/tests/idol_vae.zig new file mode 100644 index 0000000000..3a0e38b5d5 --- /dev/null +++ b/zig/tests/idol_vae.zig @@ -0,0 +1,132 @@ +//! tf-idol-vae: Idol's image encoder on a case from tools/zig/idol_case.py's vae, gated against float64 and timed. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const vae = tf.idol.vae; +const ref = tf.idol.reference; + +const usage = + \\tf-idol-vae VAE_DIR CASE_DIR [FEATURES] + \\VAE_DIR is FLUX.2-klein-4B's vae/ folder (config.json, diffusion_pytorch_model.safetensors). CASE_DIR (idol_case.py vae) + \\holds vae.safetensors: the picture [1024, 1024, 3] in [0, 1], the SDK's posterior noise [128, 128, 32], the SDK's + \\condition rows (`sdk`, [4101, 128]) and float64's on the same picture and noise (`ruler`). Our rows' worst row and + \\1 - cosine against float64 must be within twice the SDK's own, and two more encodes must repeat them byte for byte. + \\With vae_stages.safetensors, each stage is compared with float64's too. FEATURES, a safetensors file holding + \\`feature2` [1, 4101, 128] (the SDK's own run), is compared as well. CASE_DIR gets tf-idol-vae.json and rows.f32. +; + +/// Encodes timed after the gated one. +const timed = 5; + +fn take(a: std.mem.Allocator, f: *const tf.safetensors.File, name: []const u8, n: usize) ![]f32 { + const t = f.get(name) orelse { + std.debug.print("the case has no {s}\n", .{name}); + return error.MissingTensor; + }; + if (t.dtype != .f32 or t.numel() != n) { + std.debug.print("{s} is not {d} float32 values\n", .{ name, n }); + return error.BadCaseTensor; + } + const out = try a.alloc(f32, n); + @memcpy(std.mem.sliceAsBytes(out), t.bytes); + return out; +} + +fn print(list: *std.ArrayList(u8), a: std.mem.Allocator, comptime fmt: []const u8, args: anytype) !void { + try list.appendSlice(a, try std.fmt.allocPrint(a, fmt, args)); +} + +fn json(list: *std.ArrayList(u8), a: std.mem.Allocator, key: []const u8, m: ref.Match) !void { + try print(list, a, ", \"{s}\": {{\"max_abs\": {d:.6}, \"cosine\": {d:.12}, \"worst_row\": {d:.8}}}", .{ key, m.max_abs, m.cosine, m.worst_row }); +} + +const stage_names = [_][]const u8{ "down.0", "down.1", "down.2", "down.3", "mid", "moments" }; +const stage_sizes = [_]usize{ 512 * 512 * 128, 256 * 256 * 256, 128 * 128 * 512, 128 * 128 * 512, 128 * 128 * 512, 128 * 128 * 64 }; +const stage_widths = [_]usize{ 128, 256, 512, 512, 512, 64 }; + +pub fn main(init: std.process.Init) !u8 { + const a = init.arena.allocator(); + const io = init.io; + const args = try init.minimal.args.toSlice(a); + if (args.len != 3 and args.len != 4) { + std.debug.print("{s}\n", .{usage}); + return 2; + } + var case = try tf.safetensors.File.open(a, io, try std.fmt.allocPrint(a, "{s}/vae.safetensors", .{args[2]})); + defer case.close(io); + const n = vae.rows * vae.channels; + const image = try take(a, &case, "image", vae.size * vae.size * 3); + const noise = try take(a, &case, "noise", vae.latent * vae.latent * 32); + const sdk = try take(a, &case, "sdk", n); + const ruler = try take(a, &case, "ruler", n); + var stages_file: ?tf.safetensors.File = tf.safetensors.File.open(a, io, try std.fmt.allocPrint(a, "{s}/vae_stages.safetensors", .{args[2]})) catch null; + defer if (stages_file) |*f| f.close(io); + + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + const loading = mtl.clock.seconds(); + const enc = try vae.Encoder.load(std.heap.page_allocator, io, device, args[1]); + defer enc.deinit(); + const load_s = mtl.clock.seconds() - loading; + std.debug.print("{s}: encoder loaded in {d:.2} s\n", .{ device.name(), load_s }); + var out: std.ArrayList(u8) = .empty; + try print(&out, a, "{{\"device\": \"{s}\", \"load_s\": {d:.3}", .{ device.name(), load_s }); + + if (stages_file) |*f| { + var stages: vae.Stages = .{}; + var mine: [stage_names.len][]f32 = undefined; + for (&mine, stage_sizes) |*m, s| m.* = try a.alloc(f32, s); + for (0..4) |i| stages.down[i] = mine[i]; + stages.mid = mine[4]; + stages.moments = mine[5]; + const scratch = try a.alloc(f32, n); + _ = try enc.encode(image, noise, scratch, stages); + try print(&out, a, ", \"stages\": {{\"_\": 0", .{}); + for (stage_names, mine, stage_widths) |name, m, w| { + const want = try take(a, f, name, m.len); + const r = ref.match(m, want, w); + std.debug.print("stage {s} against float64: worst pixel {e:.3}, 1 - cosine {e:.3}, max |diff| {e:.3}\n", .{ name, r.worst_row, 1 - r.cosine, r.max_abs }); + try json(&out, a, name, r); + } + try out.append(a, '}'); + } + + const rows = try a.alloc(f32, n); + const first = try enc.encode(image, noise, rows, .{}); + const v = vae.gate(rows, sdk, ruler); + std.debug.print("against float64: ours worst row {e:.3}, 1 - cosine {e:.3}; the SDK's {e:.3}, {e:.3}: {s}\n", .{ v.ours.worst_row, 1 - v.ours.cosine, v.sdk.worst_row, 1 - v.sdk.cosine, if (v.pass) "pass" else "FAIL" }); + const vs_sdk = ref.match(rows, sdk, vae.channels); + std.debug.print("against the SDK: worst row {d:.4}, cosine {d:.8}, max |diff| {d:.4}\n", .{ vs_sdk.worst_row, vs_sdk.cosine, vs_sdk.max_abs }); + try json(&out, a, "ours_vs_float64", v.ours); + try json(&out, a, "sdk_vs_float64", v.sdk); + try json(&out, a, "ours_vs_sdk", vs_sdk); + if (args.len == 4) { + var feats = try tf.safetensors.File.open(a, io, try std.fmt.allocPrint(a, "{s}", .{args[3]})); + defer feats.close(io); + const theirs = try take(a, &feats, "feature2", n); + const m = ref.match(rows, theirs, vae.channels); + const same = std.mem.eql(u8, std.mem.sliceAsBytes(theirs), std.mem.sliceAsBytes(sdk)); + std.debug.print("against {s}: worst row {d:.4}, cosine {d:.8} (the case's SDK rows {s} it)\n", .{ args[3], m.worst_row, m.cosine, if (same) "equal" else "differ from" }); + try json(&out, a, "ours_vs_features", m); + try print(&out, a, ", \"case_sdk_equals_features\": {}", .{same}); + } + + const again = try a.alloc(f32, n); + var repeat = true; + var times: [timed]f64 = undefined; + var gpus: [timed]f64 = undefined; + for (×, &gpus) |*t, *gp| { + const tm = try enc.encode(image, noise, again, .{}); + t.* = tm.wall; + gp.* = tm.gpu; + repeat = repeat and std.mem.eql(u8, std.mem.sliceAsBytes(rows), std.mem.sliceAsBytes(again)); + } + std.mem.sort(f64, ×, {}, std.sort.asc(f64)); + std.mem.sort(f64, &gpus, {}, std.sort.asc(f64)); + std.debug.print("{d} more encodes: {s}; first {d:.3} s, median {d:.4} s ({d:.4} s on the GPU)\n", .{ timed, if (repeat) "identical bytes" else "DIFFERENT bytes", first.wall, times[timed / 2], gpus[timed / 2] }); + try print(&out, a, ", \"gate_pass\": {}, \"repeat_identical\": {}, \"first_s\": {d:.4}, \"median_s\": {d:.4}, \"median_gpu_s\": {d:.4}}}\n", .{ v.pass, repeat, first.wall, times[timed / 2], gpus[timed / 2] }); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = try std.fmt.allocPrint(a, "{s}/rows.f32", .{args[2]}), .data = std.mem.sliceAsBytes(rows) }); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = try std.fmt.allocPrint(a, "{s}/tf-idol-vae.json", .{args[2]}), .data = out.items }); + return if (v.pass and repeat) 0 else 1; +} From 7747ed59dae2ac0d7559d45a12f738a0ecae65fe Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 04:10:09 +0800 Subject: [PATCH 8/9] idol: the VAE recipe names the tensor units' fp32 operand truncation (#1772) --- docs/recipes/idol.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/recipes/idol.md b/docs/recipes/idol.md index a449405e51..a3b4b734db 100644 --- a/docs/recipes/idol.md +++ b/docs/recipes/idol.md @@ -156,9 +156,9 @@ give the same weights. The loader takes either: it rounds an fp32 file to bf16 a tensors, `quant_conv` and the two batch-norm statistics, and refuses any other geometry. The SDK then runs the whole encoder in bf16. This port computes in fp32 on those bf16 weights: 3x3 convolutions by -Winograd F(2x2, 3x3) with `G g G^T` made in float64 and rounded once, the stride-2 downsamples and 1x1s as implicit GEMMs, -GroupNorm statistics in two passes (the mean, then the squared deviations), fp32 attention, `precise::exp`. None of it -uses the tensor units, whose fp32 GEMM is about fp16-accurate (above). +Winograd F(2x2, 3x3) with `G g G^T` made in float64 and rounded once, the stride-2 downsamples and 1x1s as implicit +GEMMs, GroupNorm statistics in two passes (the mean, then the squared deviations), fp32 attention, `precise::exp`. None +of it uses the tensor units, whose fp32 GEMM truncates its inputs to 10 mantissa bits (above). ### Gates From b5a74eb1e6593652fb240d28f4816b30f3e3cabd Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 04:10:49 +0800 Subject: [PATCH 9/9] changelog: #1774 for Idol's VAE encoder --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 810a7ef6f2..2c5ff309ae 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,7 +6,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Idol (speridlabs/idol, image to 3D Gaussian splat), first step: its 0.37B denoiser runs natively on M5 Metal from the released checkpoint, the condition refiner once an image and all 50 guided steps on the GPU. On an M5 Ultra a forward takes 0.275 s against the publisher's SDK on torch MPS at 0.939 s, with a 1.84 GB peak footprint. Each step's velocities on the SDK's own input sit as close to an fp32 run as the SDK's own, and each joint block alone passes the same gate; `tf-idol` runs those gates from a case `tools/zig/idol_case.py` records. RePo's fp32 maps run as an fp16 hi/lo split on the tensor units, whose fp32 GEMM truncates its inputs to 10 mantissa bits (#1772). The image encoders and the decoders are not ported yet (#1766, fixes #1757). Thanks to @jayleaton. -- Idol's FLUX.2 VAE image encoder runs natively on Metal in fp32: one 1,024 x 1,024 picture to the denoiser's 4,101 condition rows, with the posterior's noise as an input. It is 0.185 s an encode on an M5 Ultra against 0.095 s for the SDK's bf16 encoder on torch MPS, and lies far closer to a float64 run: worst row 5.1e-5 against the SDK's 2.7e-2. `tf-idol-vae` gates it from a case `tools/zig/idol_case.py vae` records, with the ungated FLUX.2-klein-4B weights (after #1766, for #1757). Thanks to @jayleaton. +- Idol's FLUX.2 VAE image encoder runs natively on Metal in fp32: one 1,024 x 1,024 picture to the denoiser's 4,101 condition rows, with the posterior's noise as an input. It is 0.185 s an encode on an M5 Ultra against 0.095 s for the SDK's bf16 encoder on torch MPS, and lies far closer to a float64 run: worst row 5.1e-5 against the SDK's 2.7e-2. `tf-idol-vae` gates it from a case `tools/zig/idol_case.py vae` records, with the ungated FLUX.2-klein-4B weights (#1774, after #1766, for #1757). Thanks to @jayleaton. - DeepSeek-V4.x: `tool_choice: "required"` and named tool calls are now constrained to the model's own DSML call markup, so those requests are held to a valid call on DeepSeek too.