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 01/14] 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 02/14] 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 03/14] 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 04/14] 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 05/14] 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 06/14] 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 f04ce61162996dbc6c7ef04e66f551b9c32df62b Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 02:03:31 +0800 Subject: [PATCH 07/14] idol: CUDA denoiser for sm_121/sm_120 (untested on GB10) Port Idol's denoiser from Metal to CUDA under the same precision contract: fp16 block linears via cuBLASLt with fp32 sums and one rounding after the bias (tf_idol_bias), LayerNorm32 to fp16 with fp16 modulation, true-fp32 RePo/condition/output projections (CUBLAS_COMPUTE_32F), and a head-64 fp16 attention with fp32 softmax and accumulation in its own image (idol.cu, idol_attention.cuh), leaving prefill_attention.cu and its SASS untouched. - family modules zig/src/families/idol/cuda*.zig; the model refuses any GPU but sm_120/sm_121 by name with a log line; timing via the sys shim - tf-idol on CUDA (idol-cuda step) sharing the case driver with Metal (zig/tests/idol_run.zig) - cuBLASLt Linear.initTyped for fp16 and fp32 inputs; bf16 path unchanged - kernel contract: idol-contract (gates 1-4 over every idol.cu kernel and the attention, three controls dropping an fp16 rounding), joined to the family sets; nvcc 13.0 records for sm_80/86/89/120/121; image in contract_limits - host tests for the family and the contract's range references --- zig/build/cuda.zig | 1 + zig/build/cuda_family_host.zig | 1 + zig/build/cuda_kernels.zig | 1 + zig/build/idol_cuda.zig | 34 +++ zig/kernels/cuda/contract_controls.cu | 56 ++++ zig/kernels/cuda/idol.cu | 205 +++++++++++++ zig/kernels/cuda/idol_attention.cuh | 188 ++++++++++++ zig/src/core/contract/idol_cuda_records.zig | 24 ++ zig/src/core/contract/registry.zig | 3 +- zig/src/cuda/cublaslt.zig | 11 +- zig/src/cuda/kernels.zig | 1 + zig/src/families/idol/cuda.zig | 20 ++ zig/src/families/idol/cuda_gemm.zig | 58 ++++ zig/src/families/idol/cuda_kernels.zig | 153 ++++++++++ zig/src/families/idol/cuda_model.zig | 288 ++++++++++++++++++ zig/src/families/idol/cuda_weights.zig | 191 ++++++++++++ zig/src/families/idol/host.zig | 22 +- zig/src/families/idol/weights.zig | 22 +- zig/tests/cuda/contract_attn.zig | 4 +- zig/tests/cuda/contract_family_sets.zig | 8 +- zig/tests/cuda/contract_idol.zig | 123 ++++++++ zig/tests/cuda/contract_idol_attn.zig | 153 ++++++++++ zig/tests/cuda/contract_idol_case.zig | 135 +++++++++ zig/tests/cuda/contract_idol_host.zig | 8 + zig/tests/cuda/contract_idol_norms.zig | 208 +++++++++++++ zig/tests/cuda/contract_idol_ref.zig | 213 +++++++++++++ zig/tests/cuda/contract_idol_rotary.zig | 312 ++++++++++++++++++++ zig/tests/cuda/contract_idol_steps.zig | 217 ++++++++++++++ zig/tests/cuda/contract_limits.zig | 2 +- zig/tests/cuda/main.zig | 2 + zig/tests/idol.zig | 253 ++-------------- zig/tests/idol_cuda.zig | 43 +++ zig/tests/idol_run.zig | 248 ++++++++++++++++ 33 files changed, 2952 insertions(+), 256 deletions(-) create mode 100644 zig/build/idol_cuda.zig create mode 100644 zig/kernels/cuda/idol.cu create mode 100644 zig/kernels/cuda/idol_attention.cuh create mode 100644 zig/src/core/contract/idol_cuda_records.zig create mode 100644 zig/src/families/idol/cuda.zig create mode 100644 zig/src/families/idol/cuda_gemm.zig create mode 100644 zig/src/families/idol/cuda_kernels.zig create mode 100644 zig/src/families/idol/cuda_model.zig create mode 100644 zig/src/families/idol/cuda_weights.zig create mode 100644 zig/tests/cuda/contract_idol.zig create mode 100644 zig/tests/cuda/contract_idol_attn.zig create mode 100644 zig/tests/cuda/contract_idol_case.zig create mode 100644 zig/tests/cuda/contract_idol_host.zig create mode 100644 zig/tests/cuda/contract_idol_norms.zig create mode 100644 zig/tests/cuda/contract_idol_ref.zig create mode 100644 zig/tests/cuda/contract_idol_rotary.zig create mode 100644 zig/tests/cuda/contract_idol_steps.zig create mode 100644 zig/tests/idol_cuda.zig create mode 100644 zig/tests/idol_run.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index ca14e5a728..076f903640 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -110,6 +110,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu runner.addImport("gemma4_cuda", @import("gemma4_cuda.zig").module(b, target, optimize, cuda, mods.core, mods.lanes)); runner.addImport("spark25", @import("spark25.zig").module(b, target, optimize, cuda, mods.core, mods.lanes)); runner.addImport("moss_image_cuda", b.createModule(.{ .root_source_file = b.path("zig/src/families/moss_image/cuda_kernels.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "cuda", .module = cuda }} })); + @import("idol_cuda.zig").targets(b, target, optimize, strip, cuda, mods.core, runner); const runner_exe = b.addExecutable(.{ .name = "tf-cuda-test", .root_module = runner }); b.installArtifact(runner_exe); // tf-tp-watch-check: the ranks' lines with no GPU or NCCL; a failing or dying rank stops every rank diff --git a/zig/build/cuda_family_host.zig b/zig/build/cuda_family_host.zig index d425fe8541..d6ce31728d 100644 --- a/zig/build/cuda_family_host.zig +++ b/zig/build/cuda_family_host.zig @@ -4,6 +4,7 @@ const std = @import("std"); pub fn steps(b: *std.Build, cuda: *std.Build.Module, mods: anytype, step: *std.Build.Step) void { const host = b.graph.host; @import("minimax_cuda.zig").hostTests(b, host, cuda, mods.core, mods.lanes, step); + @import("idol_cuda.zig").hostTests(b, cuda, mods.core, step); const prefill = b.addRunArtifact(b.addTest(.{ .root_module = mods.qwen35, .filters = &.{"Qwen stepped prefill:"} })); b.step("test-qwen35-prefill-host", "Interrupted stepped prompt lifecycle without devices").dependOn(&prefill.step); step.dependOn(&prefill.step); diff --git a/zig/build/cuda_kernels.zig b/zig/build/cuda_kernels.zig index 04bded45c9..78c291fc69 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -104,6 +104,7 @@ pub const kernels = [_]Kernel{ .{ .name = "glm_index", .flags = &.{"-O3"} }, // a prompt chunk's DSA scores and top-512 pools .{ .name = "minimax_m3", .flags = glue }, // ours: MiniMax-M3.1-Flash's glue (norms, RoPE prep, routing, DSpark, tower RoPE) .{ .name = "minimax_m3_attn", .flags = glue }, // ours: MiniMax-M3.1-Flash's indexer and block-sparse attention + .{ .name = "idol", .flags = glue }, // ours: Idol's denoiser glue and its fp16 head-64 attention .{ .name = "glm_quant", .flags = torch_ops }, // ours: GLM's draft-only weight quantizers, torch at load in the Python family .{ .name = "exl3_linear", .flags = &.{ "-O3", "--expt-relaxed-constexpr" } }, // exl3/linear.cu's device code, tensorfold_exl3_linear_v5 .{ .name = "exl3_experts", .flags = &.{ "-O3", "--expt-relaxed-constexpr" } }, // exl3/experts.cu's device code, tensorfold_exl3_experts diff --git a/zig/build/idol_cuda.zig b/zig/build/idol_cuda.zig new file mode 100644 index 0000000000..c9d32101ce --- /dev/null +++ b/zig/build/idol_cuda.zig @@ -0,0 +1,34 @@ +//! Idol's CUDA denoiser in the build: its kernels for tf-cuda-test's contract, tf-idol on CUDA, its host tests. +const std = @import("std"); + +/// The family's module over the runtime and the core. +pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module) *std.Build.Module { + return b.createModule(.{ .root_source_file = b.path("zig/src/families/idol/cuda.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core } } }); +} + +/// The launchers alone (idol.cu's kernels), as tf-cuda-test's contract imports them. +fn launchers(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module) *std.Build.Module { + return b.createModule(.{ .root_source_file = b.path("zig/src/families/idol/cuda_kernels.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "cuda", .module = cuda }} }); +} + +/// The launchers in tf-cuda-test's `runner` (idol-contract), and tf-idol: the case driver on the first GPU. +pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, strip: ?bool, cuda: *std.Build.Module, core: *std.Build.Module, runner: *std.Build.Module) void { + runner.addImport("idol_cuda", launchers(b, target, optimize, cuda)); + const tool = b.createModule(.{ .root_source_file = b.path("zig/tests/idol_cuda.zig"), .target = target, .optimize = optimize, .link_libc = true, .strip = strip, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core }, .{ .name = "idol", .module = module(b, target, optimize, cuda, core) } } }); + const exe = b.addExecutable(.{ .name = "tf-idol", .root_module = tool }); + b.installArtifact(exe); + b.step("idol-cuda", "Build tf-idol: Idol's denoiser on a reference case on CUDA").dependOn(&b.addInstallArtifact(exe, .{}).step); +} + +/// The family's host tests (config, schedule, scratch, admission, launch geometry, contract ranges) in `all`. +pub fn hostTests(b: *std.Build, cuda: *std.Build.Module, core: *std.Build.Module, all: *std.Build.Step) void { + const host = b.graph.host; + const run = b.addRunArtifact(b.addTest(.{ .root_module = module(b, host, .debug, cuda, core) })); + const refs = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/contract_idol_host.zig"), .target = host, .optimize = .debug, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core }, .{ .name = "idol_cuda", .module = launchers(b, host, .debug, cuda) } } }); + const ref_run = b.addRunArtifact(b.addTest(.{ .root_module = refs })); + const step = b.step("test-idol-cuda-host", "Idol's CUDA denoiser host tests and contract references (no GPU work)"); + for ([_]*std.Build.Step{ &run.step, &ref_run.step }) |s| { + step.dependOn(s); + all.dependOn(s); + } +} diff --git a/zig/kernels/cuda/contract_controls.cu b/zig/kernels/cuda/contract_controls.cu index d047c7353d..bbd8db2602 100644 --- a/zig/kernels/cuda/contract_controls.cu +++ b/zig/kernels/cuda/contract_controls.cu @@ -537,3 +537,59 @@ TF_CTL_EXL3_GROUPED(tf_ctl_exl3_grouped_old_2, 2, tf_exl3::PLAIN) TF_CTL_EXL3_GROUPED(tf_ctl_exl3_grouped_skip, 2, tf_exl3::SKIP) TF_CTL_EXL3_GROUPED(tf_ctl_exl3_grouped_stale, 2, tf_exl3::STALE) #undef TF_CTL_EXL3_GROUPED + +// Gate 3: Idol's glue (idol.cu) with one fp16 rounding dropped each; the intervals hold idol-sdk's rounding points. +namespace tf_ctl_idol { +__device__ __forceinline__ float h(float x) { return __half2float(__float2half_rn(x)); } +__device__ __forceinline__ float ld(const __half* p) { return __half2float(*p); } +__device__ __forceinline__ float block_sum(float v, float* part) { + for (int o = 16; o > 0; o >>= 1) v += __shfl_xor_sync(0xffffffffu, v, o); + __syncthreads(); + if ((threadIdx.x & 31) == 0) part[threadIdx.x >> 5] = v; + __syncthreads(); + float total = 0.0f; + for (int w = 0; w < 8; ++w) total += part[w]; + return total; +} +} // namespace tf_ctl_idol + +// tf_idol_norm modulating the unrounded LayerNorm: fp16(fp16(y s1) + sh) where idol-sdk rounds y first. +extern "C" __global__ void __launch_bounds__(256) tf_ctl_idol_norm_late(const __half* X, const float* W, const float* B, + const __half* S1, const __half* SH, __half* Y, + int width, float eps, int mode) { + using namespace tf_ctl_idol; + __shared__ float part[8]; + const __half* x = X + (size_t)blockIdx.x * width; + float a = 0.0f, q = 0.0f; + for (int i = threadIdx.x; i < width; i += 256) a += ld(x + i); + const float mean = block_sum(a, part) / (float)width; + for (int i = threadIdx.x; i < width; i += 256) q += (ld(x + i) - mean) * (ld(x + i) - mean); + const float inv = 1.0f / sqrtf(block_sum(q, part) / (float)width + eps); + for (int i = threadIdx.x; i < width; i += 256) { + float v = (ld(x + i) - mean) * inv; + v = mode == 1 ? h(v * W[i] + B[i]) : h(h(v * ld(S1 + i)) + ld(SH + i)); + Y[(size_t)blockIdx.x * width + i] = __float2half_rn(v); + } +} + +// tf_idol_gate_add with g y left in fp32: fp16(x + g y) where idol-sdk rounds g y to fp16 first. +extern "C" __global__ void __launch_bounds__(256) tf_ctl_idol_gate_add_once(__half* X, const __half* Y, const __half* G, + int width, int rows) { + using namespace tf_ctl_idol; + const int c = blockIdx.y * 256 + threadIdx.x, r = blockIdx.x; + if (c >= width || r >= rows) return; + const size_t i = (size_t)r * width + c; + X[i] = __float2half_rn(ld(X + i) + ld(G + c) * ld(Y + i)); +} + +// tf_idol_bias taking GELU of the fp32 sum: fp16(GELU(acc + b)) where idol-sdk rounds the linear's output first. +extern "C" __global__ void __launch_bounds__(256) tf_ctl_idol_bias_late(const float* ACC, const __half* B, __half* Y, + int width, int rows, int act) { + using namespace tf_ctl_idol; + const int c = blockIdx.y * 256 + threadIdx.x, r = blockIdx.x; + if (c >= width || r >= rows) return; + const size_t i = (size_t)r * width + c; + float v = ACC[i] + ld(B + c); + if (act) v = 0.5f * v * (1.0f + tanhf(0.7978845608028654f * (v + 0.044715f * v * v * v))); + Y[i] = __float2half_rn(act ? v : h(v)); +} diff --git a/zig/kernels/cuda/idol.cu b/zig/kernels/cuda/idol.cu new file mode 100644 index 0000000000..e0131e1848 --- /dev/null +++ b/zig/kernels/cuda/idol.cu @@ -0,0 +1,205 @@ +// Idol's denoiser glue on CUDA: idol.metal's steps op for op, fp16 rows rounded where idol-sdk's torch code rounds. +#include +#include +#include + +#include "idol_attention.cuh" + +namespace tf_idol { + +typedef __half f16; + +__device__ __forceinline__ float h(float x) { return __half2float(__float2half_rn(x)); } +__device__ __forceinline__ float ld(const f16* p) { return __half2float(*p); } + +__device__ __forceinline__ float warp_sum(float v) { +#pragma unroll + for (int o = 16; o > 0; o >>= 1) v += __shfl_xor_sync(0xffffffffu, v, o); + return v; +} + +// A block of 256's sum of every thread's `v`, returned to all of them: lanes, then the 8 warps in order. +__device__ __forceinline__ float block_sum(float v, float* part) { + v = warp_sum(v); + __syncthreads(); + if ((threadIdx.x & 31) == 0) part[threadIdx.x >> 5] = v; + __syncthreads(); + float total = 0.0f; +#pragma unroll + for (int w = 0; w < 8; ++w) total += part[w]; + return total; +} + +// The row's mean and 1 / sqrt(var + eps), each sum in fp32: LayerNorm32's statistics. +__device__ __forceinline__ void stats(const f16* x, int width, float eps, float* part, float& mean, float& inv) { + float a = 0.0f; + for (int i = threadIdx.x; i < width; i += 256) a += ld(x + i); + mean = block_sum(a, part) / (float)width; + float q = 0.0f; + for (int i = threadIdx.x; i < width; i += 256) { + const float d = ld(x + i) - mean; + q += d * d; + } + inv = 1.0f / sqrtf(block_sum(q, part) / (float)width + eps); +} + +// F.normalize's clamp_min(1e-12) as torch does it: a NaN norm stays NaN (fmaxf would drop it). +__device__ __forceinline__ float clamp_min(float n) { return n < 1e-12f ? 1e-12f : n; } + +// torch's tanh GELU of an fp16 value, in fp32. +__device__ __forceinline__ float gelu(float x) { + const float inner = 0.7978845608028654f * (x + 0.044715f * x * x * x); + return 0.5f * x * (1.0f + tanhf(inner)); +} + +} // namespace tf_idol + +using namespace tf_idol; + +// LayerNorm32 of an fp16 row to fp16: affine (mode 1) in fp32 before the rounding, modulated (mode 2) in fp16 after. +extern "C" __global__ void __launch_bounds__(256) tf_idol_norm(const f16* __restrict__ X, const float* __restrict__ W, + const float* __restrict__ B, const f16* __restrict__ S1, + const f16* __restrict__ SH, f16* __restrict__ Y, int width, + float eps, int mode) { + __shared__ float part[8]; + const f16* x = X + (size_t)blockIdx.x * width; + f16* y = Y + (size_t)blockIdx.x * width; + float mean, inv; + stats(x, width, eps, part, mean, inv); + for (int i = threadIdx.x; i < width; i += 256) { + float v = (ld(x + i) - mean) * inv; + if (mode == 1) v = v * W[i] + B[i]; + v = h(v); + if (mode == 2) v = h(h(v * ld(S1 + i)) + ld(SH + i)); + y[i] = __float2half_rn(v); + } +} + +// LayerNorm of an fp16 row kept in fp32: affine (mode 1) y w + b, or modulated (mode 2) y (1 + w) + b. +extern "C" __global__ void __launch_bounds__(256) tf_idol_norm32(const f16* __restrict__ X, const float* __restrict__ W, + const float* __restrict__ B, float* __restrict__ Y, + int width, float eps, int mode) { + __shared__ float part[8]; + const f16* x = X + (size_t)blockIdx.x * width; + float* y = Y + (size_t)blockIdx.x * width; + float mean, inv; + stats(x, width, eps, part, mean, inv); + for (int i = threadIdx.x; i < width; i += 256) { + const float v = (ld(x + i) - mean) * inv; + y[i] = mode == 2 ? v * (1.0f + W[i]) + B[i] : v * W[i] + B[i]; + } +} + +// RePo: d = final_map(silu(gate) content), 3 a head; angles clamp_mul(d, freqs) pi as cos, sin [rows, 16, 32]. +extern "C" __global__ void __launch_bounds__(64) tf_idol_repo(const float* __restrict__ GC, const float* __restrict__ WF, + const float* __restrict__ F, float* __restrict__ CS, + float* __restrict__ SN, int rows) { + __shared__ float u[128]; + __shared__ float d[48]; + const int row = blockIdx.x, t = threadIdx.x; + if (row >= rows) return; + const float* g = GC + (size_t)row * 256; + for (int i = t; i < 128; i += 64) { + const float a = g[i]; + u[i] = (a / (1.0f + expf(-a))) * g[128 + i]; + } + __syncthreads(); + if (t < 48) { + float s = 0.0f; + for (int k = 0; k < 128; ++k) s = __fmaf_rn(u[k], WF[t * 128 + k], s); + d[t] = s; + } + __syncthreads(); + for (int i = t; i < 512; i += 64) { + const int hd = i / 32, j = i % 32; + const int axis = j < 10 ? 0 : (j < 20 ? 1 : 2); + const float x = d[hd * 3 + axis], fr = F[j]; + const float ft = tanhf(fr); + const float ang = (x * ft + x * (fr - ft)) * 3.14159265358979323846f; + CS[(size_t)row * 512 + i] = cosf(ang); + SN[(size_t)row * 512 + i] = sinf(ang); + } +} + +// q, k of qkv rows turned by cos, sin (to fp16), F.normalize times gamma times `scale`; v copied. A warp a head. +extern "C" __global__ void __launch_bounds__(256) tf_idol_qk(const f16* __restrict__ QKV, const float* __restrict__ GQ, + const float* __restrict__ GK, const float* __restrict__ CS, + const float* __restrict__ SN, f16* __restrict__ OUT, int rows, + float scale) { + const int r = blockIdx.x, hd = blockIdx.y * 8 + (threadIdx.x >> 5), lane = threadIdx.x & 31; + if (hd >= 16 || r >= rows) return; + const size_t from = (size_t)r * 3072 + hd * 64 + lane * 2; + const size_t at = (size_t)r * 512 + hd * 32 + lane; + const float c = CS[at], s = SN[at]; + const float q0 = ld(QKV + from), q1 = ld(QKV + from + 1); + const float k0 = ld(QKV + from + 1024), k1 = ld(QKV + from + 1025); + const float rq0 = h(q0 * c - q1 * s), rq1 = h(q0 * s + q1 * c); + const float rk0 = h(k0 * c - k1 * s), rk1 = h(k0 * s + k1 * c); + const float qn = clamp_min(sqrtf(warp_sum(rq0 * rq0 + rq1 * rq1))); + const float kn = clamp_min(sqrtf(warp_sum(rk0 * rk0 + rk1 * rk1))); + const int gi = hd * 64 + lane * 2; + f16* o = OUT + from; + o[0] = __float2half_rn(rq0 / qn * GQ[gi] * scale); + o[1] = __float2half_rn(rq1 / qn * GQ[gi + 1] * scale); + o[1024] = __float2half_rn(rk0 / kn * GK[gi] * scale); + o[1025] = __float2half_rn(rk1 / kn * GK[gi + 1] * scale); + o[2048] = QKV[from + 2048]; + o[2049] = QKV[from + 2049]; +} + +// x = fp16(x + fp16(g y)) by column; g is a modulation gate, or ones for an ungated residual. In place. +extern "C" __global__ void __launch_bounds__(256) tf_idol_gate_add(f16* __restrict__ X, const f16* __restrict__ Y, + const f16* __restrict__ G, int width, int rows) { + const int c = blockIdx.y * 256 + threadIdx.x, r = blockIdx.x; + if (c >= width || r >= rows) return; + const size_t i = (size_t)r * width + c; + X[i] = __float2half_rn(ld(X + i) + h(ld(G + c) * ld(Y + i))); +} + +// y = fp16(acc + b): a block linear's fp32 sums and its fp16 bias, one rounding; then fp16(GELU(y)) when `act`. +extern "C" __global__ void __launch_bounds__(256) tf_idol_bias(const float* __restrict__ ACC, const f16* __restrict__ B, + f16* __restrict__ Y, int width, int rows, int act) { + const int c = blockIdx.y * 256 + threadIdx.x, r = blockIdx.x; + if (c >= width || r >= rows) return; + const size_t i = (size_t)r * width + c; + float v = h(ACC[i] + ld(B + c)); + if (act) v = gelu(v); + Y[i] = __float2half_rn(v); +} + +// x = fp16(fp16(z w^T + b) + pe): the latent's input projection (16 inputs, fp32) and its fp16 position embedding. +extern "C" __global__ void __launch_bounds__(256) tf_idol_embed(const float* __restrict__ Z, const float* __restrict__ W, + const float* __restrict__ B, const f16* __restrict__ PE, + f16* __restrict__ X, int rows) { + const int c = blockIdx.y * 256 + threadIdx.x, r = blockIdx.x; + if (c >= 1024 || r >= rows) return; + float s = 0.0f; +#pragma unroll + for (int k = 0; k < 16; ++k) s = __fmaf_rn(Z[(size_t)r * 16 + k], W[c * 16 + k], s); + const size_t i = (size_t)r * 1024 + c; + X[i] = __float2half_rn(h(s + B[c]) + ld(PE + i)); +} + +// y = fp16((a + ba) + (c + bc)) by column: the two condition projections' fp32 sums and biases. +extern "C" __global__ void __launch_bounds__(256) tf_idol_cond(const float* __restrict__ A, const float* __restrict__ BA, + const float* __restrict__ C, const float* __restrict__ BC, + f16* __restrict__ Y, int width, int rows) { + const int c = blockIdx.y * 256 + threadIdx.x, r = blockIdx.x; + if (c >= width || r >= rows) return; + const size_t i = (size_t)r * width + c; + Y[i] = __float2half_rn((A[i] + BA[c]) + (C[i] + BC[c])); +} + +// y += b by column in fp32 (an fp32 Linear's bias), in place. +extern "C" __global__ void __launch_bounds__(256) tf_idol_bias32(float* __restrict__ Y, const float* __restrict__ B, + int width, int rows) { + const int c = blockIdx.y * 256 + threadIdx.x, r = blockIdx.x; + if (c >= width || r >= rows) return; + Y[(size_t)r * width + c] += B[c]; +} + +// Full attention at head width 64 over fp16 packed rows; idol_attention.cuh has the body. +extern "C" __global__ void __launch_bounds__(256) tf_idol_attention(const f16* __restrict__ qkv, f16* __restrict__ out, + int W, int HO, float scale) { + tf_idol_attn::full<64, 8>(qkv, out, W, HO, scale); +} diff --git a/zig/kernels/cuda/idol_attention.cuh b/zig/kernels/cuda/idol_attention.cuh new file mode 100644 index 0000000000..5a1aec1bbc --- /dev/null +++ b/zig/kernels/cuda/idol_attention.cuh @@ -0,0 +1,188 @@ +// Idol's attention: vision_kernel's row arithmetic over fp16 q, k, v and P (fp32 scores, softmax, sums), its own copy. +#pragma once + +#include +#include +#include + +#include "common/qmm_frag.cuh" + +namespace tf_idol_attn { + +using qmm_frag::commit; +using qmm_frag::cp16z; +using qmm_frag::ldmatrix4; +using qmm_frag::ldmatrix4t; + +constexpr int BN = 64; // keys a tile +constexpr float LOG2E = 1.4426950408889634f; + +__device__ __forceinline__ void mma(float (&d)[4], const uint32_t (&a)[4], uint32_t b0, uint32_t b1) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1)); +} + +__device__ __forceinline__ void mma0(float (&d)[4], const uint32_t (&a)[4], uint32_t b0, uint32_t b1) { + d[0] = d[1] = d[2] = d[3] = 0.0f; + mma(d, a, b0, b1); +} + +__device__ __forceinline__ float texp(float x) { + float y; + asm("ex2.approx.f32 %0, %1;\n" : "=f"(y) : "f"(__fmul_rn(x, LOG2E))); + return y; +} + +__device__ __forceinline__ float tdiv(float a, float b) { + float y; + asm("div.full.f32 %0, %1, %2;\n" : "=f"(y) : "f"(a), "f"(b)); + return y; +} + +__device__ __forceinline__ uint32_t pack(float lo, float hi) { + __half2 v = __floats2half2_rn(lo, hi); + return *reinterpret_cast(&v); +} + +__device__ __forceinline__ float rowsum(const float (&c)[8]) { + float w[8]; +#pragma unroll + for (int j = 0; j < 8; ++j) { + const float x = __fadd_rn(c[j], __shfl_xor_sync(0xffffffffu, c[j], 2)); + w[j] = __fadd_rn(x, __shfl_xor_sync(0xffffffffu, x, 1)); + } + return __fadd_rn(__fadd_rn(__fadd_rn(w[0], w[4]), __fadd_rn(w[2], w[6])), + __fadd_rn(__fadd_rn(w[1], w[5]), __fadd_rn(w[3], w[7]))); +} + +// Every row over every key, a block a head and 16 WARPS rows: qkv [W, 3 HO D] fp16 -> out [W, HO D] fp16. +template +__device__ __forceinline__ void full(const __half* __restrict__ qkv, __half* __restrict__ out, int W, int HO, + float scale) { + static_assert(D % 16 == 0, "whole 16-dim steps"); + constexpr int C = D / 8, RS = C + 1, K16 = D / 16, NT = BN / 8, THREADS = 32 * WARPS; + extern __shared__ uint4 smem[]; + const int warp = threadIdx.x / 32, lane = threadIdx.x % 32, g = lane / 4, qd = lane % 4, mt = lane / 8; + const int head = blockIdx.y, wrow = (blockIdx.x * WARPS + warp) * 16, last = W - 1, tiles = (W + BN - 1) / BN; + const int64_t stride = 3LL * HO * D; + const __half* __restrict__ q = qkv + head * D; + const __half* __restrict__ kc = qkv + (HO + head) * D; + const __half* __restrict__ vc = qkv + (2 * HO + head) * D; + auto stage = [&](int t) { // tile t's K rows then V rows into buffer t % 2 + uint4* dst = smem + (t & 1) * 2 * BN * RS; + for (int i = threadIdx.x; i < 2 * BN * C; i += THREADS) { + const int kv = i / (BN * C), r = (i / C) % BN, c = i % C, key = t * BN + r; + const bool in = key < W; + cp16z(dst + (kv * BN + r) * RS + c, (kv ? vc : kc) + (int64_t)(in ? key : 0) * stride + c * 8, in); + } + commit(); + }; + uint32_t qa[K16][4]; + const int r0 = wrow + g, r1 = r0 + 8; + { + const uint32_t* q0 = reinterpret_cast(q + (int64_t)min(r0, last) * stride); + const uint32_t* q1 = reinterpret_cast(q + (int64_t)min(r1, last) * stride); +#pragma unroll + for (int k = 0; k < K16; ++k) { + qa[k][0] = r0 < W ? q0[8 * k + qd] : 0u; + qa[k][1] = r1 < W ? q1[8 * k + qd] : 0u; + qa[k][2] = r0 < W ? q0[8 * k + qd + 4] : 0u; + qa[k][3] = r1 < W ? q1[8 * k + qd + 4] : 0u; + } + } + float o[C][4]; +#pragma unroll + for (int i = 0; i < C; ++i) o[i][0] = o[i][1] = o[i][2] = o[i][3] = 0.0f; + float m0 = -INFINITY, m1 = -INFINITY, l0 = 0.0f, l1 = 0.0f; + stage(0); + for (int t = 0; t < tiles; ++t) { + qmm_frag::wait<0>(); + __syncthreads(); // tile t landed, and every warp is done with tile t - 1 + if (t + 1 < tiles) stage(t + 1); + const uint4* ks = smem + (t & 1) * 2 * BN * RS; + const uint4* vs = ks + BN * RS; + float sc[NT][4]; +#pragma unroll + for (int k = 0; k < K16; ++k) { +#pragma unroll + for (int n = 0; n < NT; n += 2) { + uint32_t b[4]; + ldmatrix4(b, ks + (8 * (n + mt / 2) + lane % 8) * RS + 2 * k + mt % 2); + if (k == 0) { + mma0(sc[n], qa[k], b[0], b[1]); + mma0(sc[n + 1], qa[k], b[2], b[3]); + } else { + mma(sc[n], qa[k], b[0], b[1]); + mma(sc[n + 1], qa[k], b[2], b[3]); + } + } + } + const int kb = t * BN + 2 * qd; + float tm0 = -INFINITY, tm1 = -INFINITY; +#pragma unroll + for (int n = 0; n < NT; ++n) { + const int c = kb + 8 * n; + sc[n][0] = c <= last ? __fmul_rn(sc[n][0], scale) : -INFINITY; + sc[n][1] = c + 1 <= last ? __fmul_rn(sc[n][1], scale) : -INFINITY; + sc[n][2] = c <= last ? __fmul_rn(sc[n][2], scale) : -INFINITY; + sc[n][3] = c + 1 <= last ? __fmul_rn(sc[n][3], scale) : -INFINITY; + tm0 = fmaxf(tm0, fmaxf(sc[n][0], sc[n][1])); + tm1 = fmaxf(tm1, fmaxf(sc[n][2], sc[n][3])); + } + for (int x = 1; x <= 2; x *= 2) { + tm0 = fmaxf(tm0, __shfl_xor_sync(0xffffffffu, tm0, x)); + tm1 = fmaxf(tm1, __shfl_xor_sync(0xffffffffu, tm1, x)); + } + const bool act0 = tm0 != -INFINITY, act1 = tm1 != -INFINITY; + const float n0 = act0 ? fmaxf(m0, tm0) : m0, n1 = act1 ? fmaxf(m1, tm1) : m1; + const float a0 = act0 ? (m0 == -INFINITY ? 0.0f : texp(__fsub_rn(m0, n0))) : 1.0f; + const float a1 = act1 ? (m1 == -INFINITY ? 0.0f : texp(__fsub_rn(m1, n1))) : 1.0f; + float c0[NT], c1[NT]; +#pragma unroll + for (int n = 0; n < NT; ++n) { + const int c = kb + 8 * n; + sc[n][0] = act0 && c <= last ? texp(__fsub_rn(sc[n][0], n0)) : 0.0f; + sc[n][1] = act0 && c + 1 <= last ? texp(__fsub_rn(sc[n][1], n0)) : 0.0f; + sc[n][2] = act1 && c <= last ? texp(__fsub_rn(sc[n][2], n1)) : 0.0f; + sc[n][3] = act1 && c + 1 <= last ? texp(__fsub_rn(sc[n][3], n1)) : 0.0f; + c0[n] = __fadd_rn(sc[n][0], sc[n][1]); + c1[n] = __fadd_rn(sc[n][2], sc[n][3]); + } + l0 = __fmaf_rn(l0, a0, rowsum(c0)); + l1 = __fmaf_rn(l1, a1, rowsum(c1)); + m0 = n0; + m1 = n1; +#pragma unroll + for (int i = 0; i < C; ++i) { + o[i][0] = __fmul_rn(o[i][0], a0); + o[i][1] = __fmul_rn(o[i][1], a0); + o[i][2] = __fmul_rn(o[i][2], a1); + o[i][3] = __fmul_rn(o[i][3], a1); + } +#pragma unroll + for (int j = 0; j < NT / 2; ++j) { // 16-key chunks, in key order for every output tile + const uint32_t pa[4] = {pack(sc[2 * j][0], sc[2 * j][1]), pack(sc[2 * j][2], sc[2 * j][3]), + pack(sc[2 * j + 1][0], sc[2 * j + 1][1]), pack(sc[2 * j + 1][2], sc[2 * j + 1][3])}; +#pragma unroll + for (int i = 0; i < C; i += 2) { + uint32_t b[4]; + ldmatrix4t(b, vs + (16 * j + 8 * (mt % 2) + lane % 8) * RS + i + mt / 2); + mma(o[i], pa, b[0], b[1]); + mma(o[i + 1], pa, b[2], b[3]); + } + } + } +#pragma unroll + for (int i = 0; i < C; ++i) { + const int d = 8 * i + 2 * qd; + if (r0 < W) + *reinterpret_cast(out + ((int64_t)r0 * HO + head) * D + d) = pack(tdiv(o[i][0], l0), tdiv(o[i][1], l0)); + if (r1 < W) + *reinterpret_cast(out + ((int64_t)r1 * HO + head) * D + d) = pack(tdiv(o[i][2], l1), tdiv(o[i][3], l1)); + } +} + +} // namespace tf_idol_attn diff --git a/zig/src/core/contract/idol_cuda_records.zig b/zig/src/core/contract/idol_cuda_records.zig new file mode 100644 index 0000000000..1ff9a8ee00 --- /dev/null +++ b/zig/src/core/contract/idol_cuda_records.zig @@ -0,0 +1,24 @@ +//! Idol's denoiser kernels on CUDA (idol.cu): every SM's compile records from nvcc 13.0, the release compiler. +const registry = @import("registry.zig"); + +/// nvcc 13.0's usage (cuobjdump -res-usage of the glue-flag image; sm_120's 1024 reserved shared bytes left out). +pub const cuda = [_]registry.CudaRecord{ + .{ .launch = .{ .name = "idol tf_idol_norm", .symbol = "tf_idol_norm", .image = "idol", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 32, .static = 32, .bound = 256 }, .{ .regs = 40, .static = 32, .bound = 256 }, .{ .regs = 40, .static = 32, .bound = 256 }, .{ .regs = 36, .static = 32, .bound = 256 }, .{ .regs = 38, .static = 32, .bound = 256 } } }, + .{ .launch = .{ .name = "idol tf_idol_norm32", .symbol = "tf_idol_norm32", .image = "idol", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 32, .static = 32, .bound = 256 }, .{ .regs = 38, .static = 32, .bound = 256 }, .{ .regs = 38, .static = 32, .bound = 256 }, .{ .regs = 38, .static = 32, .bound = 256 }, .{ .regs = 38, .static = 32, .bound = 256 } } }, + .{ .launch = .{ .name = "idol tf_idol_repo", .symbol = "tf_idol_repo", .image = "idol", .threads = 64, .min_sm = 80 }, .usage = .{ .{ .regs = 38, .static = 704, .bound = 64 }, .{ .regs = 66, .static = 704, .bound = 64 }, .{ .regs = 44, .static = 704, .bound = 64 }, .{ .regs = 37, .static = 704, .bound = 64 }, .{ .regs = 37, .static = 704, .bound = 64 } } }, + .{ .launch = .{ .name = "idol tf_idol_qk", .symbol = "tf_idol_qk", .image = "idol", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 25, .static = 0, .bound = 256 }, .{ .regs = 25, .static = 0, .bound = 256 }, .{ .regs = 25, .static = 0, .bound = 256 }, .{ .regs = 27, .static = 0, .bound = 256 }, .{ .regs = 27, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "idol tf_idol_gate_add", .symbol = "tf_idol_gate_add", .image = "idol", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "idol tf_idol_bias", .symbol = "tf_idol_bias", .image = "idol", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 15, .static = 0, .bound = 256 }, .{ .regs = 15, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "idol tf_idol_embed", .symbol = "tf_idol_embed", .image = "idol", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "idol tf_idol_cond", .symbol = "tf_idol_cond", .image = "idol", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 16, .static = 0, .bound = 256 }, .{ .regs = 16, .static = 0, .bound = 256 }, .{ .regs = 16, .static = 0, .bound = 256 }, .{ .regs = 18, .static = 0, .bound = 256 }, .{ .regs = 18, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "idol tf_idol_bias32", .symbol = "tf_idol_bias32", .image = "idol", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 10, .static = 0, .bound = 256 }, .{ .regs = 10, .static = 0, .bound = 256 }, .{ .regs = 10, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "idol tf_idol_attention", .symbol = "tf_idol_attention", .image = "idol", .threads = 256, .dynamic = 36864, .min_sm = 80 }, .usage = .{ .{ .regs = 132, .static = 0, .bound = 256 }, .{ .regs = 132, .static = 0, .bound = 256 }, .{ .regs = 132, .static = 0, .bound = 256 }, .{ .regs = 120, .static = 0, .bound = 256 }, .{ .regs = 120, .static = 0, .bound = 256 } } }, +}; + +test "Idol CUDA: measured launches fit their physical resource limits" { + const t = @import("std").testing; + for (cuda) |r| { + for (r.usage) |u| try t.expect(u.regs > 0); + try t.expectEqual(@as(usize, 0), try registry.checkCuda(t.allocator, &.{r}, true)); + } +} diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index f55d146bf7..22e8338f01 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -205,7 +205,7 @@ pub const cuda = @import("flashnext_cuda_records.zig").cuda ++ [_]CudaRecord{ .{ .launch = .{ .name = "experts3_prefill down", .symbol = "_ZN19tf_experts3_prefill14prefill_kernelILi1ELi3ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_PvifS6_i", .image = "experts3_prefill", .threads = 256, .dynamic = 47104, .opt_in = true }, .usage = .{ .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 98, .static = 0, .bound = 256 }, .{ .regs = 98, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "experts3 tf_experts3_widen", .symbol = "tf_experts3_widen", .image = "experts3", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 32, .static = 0, .bound = 0 } } }, .{ .launch = .{ .name = "experts3 tf_experts3_planes", .symbol = "tf_experts3_planes", .image = "experts3", .threads = 256 }, .usage = .{ .{ .regs = 26, .static = 0, .bound = 0 }, .{ .regs = 28, .static = 0, .bound = 0 }, .{ .regs = 28, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 } } }, -} ++ @import("registry_media.zig").cuda ++ @import("nemotron_prompt_records.zig").cuda ++ @import("glm_cuda_records.zig").cuda ++ @import("registry_shared_attention.zig").records(CudaRecord) ++ @import("gemma_cuda_records.zig").cuda ++ @import("spark_cuda_records.zig").cuda ++ @import("row_cuda_records.zig").cuda ++ @import("flashnext_tile_records.zig").cuda ++ @import("window_rows_records.zig").cuda ++ @import("modelopt_lane_records.zig").cuda ++ @import("nvfp4_fact_records.zig").cuda ++ @import("lane_half_records.zig").cuda ++ @import("pocket_records.zig").cuda ++ @import("minimax_cuda_records.zig").cuda ++ @import("moss_image_cuda_records.zig").cuda; +} ++ @import("registry_media.zig").cuda ++ @import("nemotron_prompt_records.zig").cuda ++ @import("glm_cuda_records.zig").cuda ++ @import("registry_shared_attention.zig").records(CudaRecord) ++ @import("gemma_cuda_records.zig").cuda ++ @import("spark_cuda_records.zig").cuda ++ @import("row_cuda_records.zig").cuda ++ @import("flashnext_tile_records.zig").cuda ++ @import("window_rows_records.zig").cuda ++ @import("modelopt_lane_records.zig").cuda ++ @import("nvfp4_fact_records.zig").cuda ++ @import("lane_half_records.zig").cuda ++ @import("pocket_records.zig").cuda ++ @import("minimax_cuda_records.zig").cuda ++ @import("moss_image_cuda_records.zig").cuda ++ @import("idol_cuda_records.zig").cuda; /// The nvcc release `cuda`'s usage records come from: the compiler that builds the release fatbins. pub const cuda_release: limits.Nvcc = .{ .major = 13, .minor = 0 }; @@ -585,6 +585,7 @@ test { _ = @import("moss_image_cuda_records.zig"); _ = @import("glm_kv_pending.zig"); _ = @import("minimax_cuda_records.zig"); + _ = @import("idol_cuda_records.zig"); _ = @import("glm_f32w_pending.zig"); _ = @import("registry_pending_test.zig"); _ = @import("flashnext_pages_pending.zig"); diff --git a/zig/src/cuda/cublaslt.zig b/zig/src/cuda/cublaslt.zig index c620786cb5..823544d90d 100644 --- a/zig/src/cuda/cublaslt.zig +++ b/zig/src/cuda/cublaslt.zig @@ -110,7 +110,7 @@ pub const Library = struct { } }; -/// One row-major bf16 product D[m,n] = X[m,k] . W[n,k]^T (a torch Linear), fp32 accumulate, bf16 or fp32 out. +/// One row-major product D[m,n] = X[m,k] . W[n,k]^T (a torch Linear), fp32 accumulate; bf16 inputs unless typed. pub const Linear = struct { lt: *const Library, handle: Handle, @@ -130,6 +130,11 @@ pub const Linear = struct { /// `init`'s plan, its algorithm also running `epilogue` when it supports it (`epilogue` says if); see `runBias`. pub fn initEpilogue(lt: *const Library, m: u64, n: u64, k: u64, out: DataType, workspace_limit: usize, epilogue: Epilogue) Error!Linear { + return initTyped(lt, m, n, k, .bf16, out, workspace_limit, epilogue); + } + + /// `initEpilogue` with X and W of type `in` (bf16, fp16, or fp32 for true fp32 products), fp32 compute. + pub fn initTyped(lt: *const Library, m: u64, n: u64, k: u64, in: DataType, out: DataType, workspace_limit: usize, epilogue: Epilogue) Error!Linear { var self: Linear = undefined; self.lt = lt; try lt.check(lt.api.cublasLtCreate(&self.handle), "cublasLtCreate"); @@ -141,9 +146,9 @@ pub const Linear = struct { try lt.check(lt.api.cublasLtMatmulDescSetAttribute(self.desc, desc_transa, &ta, @sizeOf(c_int)), "set transa"); try lt.check(lt.api.cublasLtMatmulDescSetAttribute(self.desc, desc_transb, &tb, @sizeOf(c_int)), "set transb"); // column-major view: D^T[n,m] = W[n,k] (stored k x n, transposed) . X^T[k,m] - try lt.check(lt.api.cublasLtMatrixLayoutCreate(&self.a, .bf16, k, n, @intCast(k)), "layout W"); + try lt.check(lt.api.cublasLtMatrixLayoutCreate(&self.a, in, k, n, @intCast(k)), "layout W"); errdefer _ = lt.api.cublasLtMatrixLayoutDestroy(self.a); - try lt.check(lt.api.cublasLtMatrixLayoutCreate(&self.b, .bf16, k, m, @intCast(k)), "layout X"); + try lt.check(lt.api.cublasLtMatrixLayoutCreate(&self.b, in, k, m, @intCast(k)), "layout X"); errdefer _ = lt.api.cublasLtMatrixLayoutDestroy(self.b); try lt.check(lt.api.cublasLtMatrixLayoutCreate(&self.c, out, n, m, @intCast(n)), "layout D"); errdefer _ = lt.api.cublasLtMatrixLayoutDestroy(self.c); diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index db13a74f50..ab6626047c 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -72,6 +72,7 @@ pub const glm_ops: []const u8 = if (available) &Blob("fatbin_glm_ops").bytes els pub const glm_index: []const u8 = if (available) &Blob("fatbin_glm_index").bytes else &.{}; pub const minimax_m3: []const u8 = if (available) &Blob("fatbin_minimax_m3").bytes else &.{}; pub const minimax_m3_attn: []const u8 = if (available) &Blob("fatbin_minimax_m3_attn").bytes else &.{}; +pub const idol: []const u8 = if (available) &Blob("fatbin_idol").bytes else &.{}; pub const glm_quant: []const u8 = if (available) &Blob("fatbin_glm_quant").bytes else &.{}; pub const glm_vision: []const u8 = if (available) &Blob("fatbin_glm_vision").bytes else &.{}; pub const glm_mm: []const u8 = if (available) &Blob("fatbin_glm_mm").bytes else &.{}; diff --git a/zig/src/families/idol/cuda.zig b/zig/src/families/idol/cuda.zig new file mode 100644 index 0000000000..ea600075b4 --- /dev/null +++ b/zig/src/families/idol/cuda.zig @@ -0,0 +1,20 @@ +//! Idol's denoiser on CUDA: driven by tf-idol-cuda beside the lane engine, as the Metal family is by tf-idol. +pub const config = @import("config.zig"); +pub const host = @import("host.zig"); +pub const reference = @import("reference.zig"); +pub const kernels = @import("cuda_kernels.zig"); +pub const Model = @import("cuda_model.zig").Model; +pub const Condition = @import("cuda_model.zig").Condition; +pub const Timing = @import("cuda_model.zig").Timing; +pub const seconds = @import("cuda_model.zig").seconds; +pub const admitted = @import("cuda_model.zig").admitted; + +test { + _ = @import("cuda_kernels.zig"); + _ = @import("cuda_gemm.zig"); + _ = @import("cuda_model.zig"); + _ = @import("cuda_weights.zig"); + _ = config; + _ = host; + _ = reference; +} diff --git a/zig/src/families/idol/cuda_gemm.zig b/zig/src/families/idol/cuda_gemm.zig new file mode 100644 index 0000000000..a39a36f347 --- /dev/null +++ b/zig/src/families/idol/cuda_gemm.zig @@ -0,0 +1,58 @@ +//! Idol's products on cuBLASLt: fp16 block linears into fp32 sums, and true fp32 products (no TF32), a plan a shape. +const std = @import("std"); +const cuda = @import("cuda"); +const lt_ = cuda.cublaslt; + +pub const workspace_bytes = 32 << 20; + +/// The operands' type: fp16 for the blocks' linears, fp32 for RePo's maps and the condition and output projections. +pub const In = enum { f16, f32 }; + +pub const Gemms = struct { + gpa: std.mem.Allocator, + lt: lt_.Library, + workspace: cuda.DeviceBuffer, + plans: std.AutoHashMapUnmanaged([4]u64, lt_.Linear) = .empty, + + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver) !Gemms { + var lt = try lt_.Library.open(); + errdefer lt.close(); + return .{ .gpa = gpa, .lt = lt, .workspace = try cuda.DeviceBuffer.alloc(d, workspace_bytes) }; + } + + pub fn deinit(g: *Gemms) void { + var it = g.plans.valueIterator(); + while (it.next()) |l| l.deinit(); + g.plans.deinit(g.gpa); + g.workspace.free(); + g.lt.close(); + } + + /// y fp32 [rows, n] (rows `n` apart) = x [rows, k] . w [n, k]^T with fp32 compute and sums, on `stream`. + pub fn run(g: *Gemms, s: cuda.Stream, in: In, x: u64, w: u64, y: u64, rows: usize, n: usize, k: usize) !void { + if (rows == 0) return; + const p = try g.plan(in, rows, n, k); + try p.run(x, w, y, g.workspace.ptr, workspace_bytes, s.handle); + } + + fn plan(g: *Gemms, in: In, rows: usize, n: usize, k: usize) !*const lt_.Linear { + const got = try g.plans.getOrPut(g.gpa, .{ @backingInt(in), rows, n, k }); + if (got.found_existing) return got.value_ptr; + const t: lt_.DataType = switch (in) { + .f16 => .f16, + .f32 => .f32, + }; + got.value_ptr.* = lt_.Linear.initTyped(&g.lt, rows, n, k, t, .f32, workspace_bytes, .none) catch |e| { + g.plans.removeByPtr(got.key_ptr); + return e; + }; + return got.value_ptr; + } +}; + +test "cuBLASLt's fp32 compute type is CUBLAS_COMPUTE_32F, not a TF32 or fast mode" { + // cublasComputeType_t: 68 is CUBLAS_COMPUTE_32F; 74 and 77 are its FAST_TF32 and FAST_16F forms + try std.testing.expectEqual(@as(c_int, 68), lt_.compute_32f); + try std.testing.expectEqual(@as(c_int, 2), @backingInt(lt_.DataType.f16)); + try std.testing.expectEqual(@as(c_int, 0), @backingInt(lt_.DataType.f32)); +} diff --git a/zig/src/families/idol/cuda_kernels.zig b/zig/src/families/idol/cuda_kernels.zig new file mode 100644 index 0000000000..02aa16fc79 --- /dev/null +++ b/zig/src/families/idol/cuda_kernels.zig @@ -0,0 +1,153 @@ +//! Idol's own CUDA kernels (idol.cu): LayerNorms, RePo's angles, q/k turned and normed, epilogues, head-64 attention. +const std = @import("std"); +const cuda = @import("cuda"); +const cfg = @import("config.zig"); + +fn int(x: usize) c_int { + return @intCast(x); +} + +fn ceil(n: usize, d: usize) u32 { + return @intCast((n + d - 1) / d); +} + +/// The LayerNorm's form: plain, affine (fp32 weight and bias), or modulated (after the fp16 rounding). +pub const Norm = enum(c_int) { plain = 0, affine = 1, modulated = 2 }; + +/// The attention's tile: 8 warps of 16 query rows a block, two buffers of 64 K and V rows padded to 9 chunks of 16 B. +pub const attention_warps = 8; +pub const attention_smem: u32 = 2 * 2 * 64 * (cfg.head_dim / 8 + 1) * 16; + +pub const Kernels = struct { + mod: cuda.Module, + norm: cuda.Function, + norm32: cuda.Function, + repo: cuda.Function, + qk: cuda.Function, + gate_add: cuda.Function, + bias: cuda.Function, + embed: cuda.Function, + cond: cuda.Function, + bias32: cuda.Function, + attention: cuda.Function, + + pub fn load(d: *const cuda.Driver) !Kernels { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var mod = try cuda.Module.load(d, cuda.kernels.idol); + errdefer mod.unload(); + var k: Kernels = undefined; + k.mod = mod; + inline for (.{ "norm", "norm32", "repo", "qk", "gate_add", "bias", "embed", "cond", "bias32", "attention" }) |n| + @field(k, n) = try mod.function("tf_idol_" ++ n); + return k; + } + + pub fn unload(k: *Kernels) void { + k.mod.unload(); + } +}; + +fn go(f: cuda.Function, s: cuda.Stream, grid: cuda.Dim3, threads: u32, a: *cuda.Args) !void { + try cuda.launch.launch(f, .{ .grid = grid, .block = .{ .x = threads } }, s, a); +} + +fn addAll(a: *cuda.Args, ptrs: []const u64) void { + for (ptrs) |p| a.add(p); +} + +/// y fp16 = the LayerNorm32 of x's `rows` rows, rounded to fp16: plain, affine (w, b fp32) or modulated (s1, sh fp16). +pub fn norm(k: *const Kernels, s: cuda.Stream, x: u64, w: u64, b: u64, s1: u64, sh: u64, y: u64, rows: usize, eps: f32, mode: Norm) !void { + var a: cuda.Args = .{}; + addAll(&a, &.{ x, w, b, s1, sh, y }); + a.add(int(cfg.hidden)); + a.add(eps); + a.add(@backingInt(mode)); + try go(k.norm, s, .{ .x = @intCast(rows) }, 256, &a); +} + +/// y fp32 = the LayerNorm of x's rows times w plus b (affine), or times (1 + w) plus b (modulated). +pub fn norm32(k: *const Kernels, s: cuda.Stream, x: u64, w: u64, b: u64, y: u64, rows: usize, eps: f32, mode: Norm) !void { + var a: cuda.Args = .{}; + addAll(&a, &.{ x, w, b, y }); + a.add(int(cfg.hidden)); + a.add(eps); + a.add(@backingInt(mode)); + try go(k.norm32, s, .{ .x = @intCast(rows) }, 256, &a); +} + +/// cos, sin [rows, 16, 32] fp32 of RePo from its gate | content sums [rows, 256] fp32, final map and frequencies. +pub fn repo(k: *const Kernels, s: cuda.Stream, gc: u64, final: u64, freqs: u64, cos: u64, sin: u64, rows: usize) !void { + var a: cuda.Args = .{}; + addAll(&a, &.{ gc, final, freqs, cos, sin }); + a.add(int(rows)); + try go(k.repo, s, .{ .x = @intCast(rows) }, 64, &a); +} + +/// Attention's packed rows [rows, q | k | v] from qkv rows: q, k turned by cos, sin, normed by gq, gk times 8. +pub fn qk(k: *const Kernels, s: cuda.Stream, qkv: u64, gq: u64, gk: u64, cos: u64, sin: u64, out: u64, rows: usize) !void { + var a: cuda.Args = .{}; + addAll(&a, &.{ qkv, gq, gk, cos, sin, out }); + a.add(int(rows)); + a.add(@as(f32, @sqrt(@as(f32, cfg.head_dim)))); + try go(k.qk, s, .{ .x = @intCast(rows), .y = ceil(cfg.heads, 8) }, 256, &a); +} + +/// x [rows, width] += y gated by the fp16 vector g, in place. +pub fn gateAdd(k: *const Kernels, s: cuda.Stream, x: u64, y: u64, g: u64, rows: usize, width: usize) !void { + var a: cuda.Args = .{}; + addAll(&a, &.{ x, y, g }); + for ([_]usize{ width, rows }) |v| a.add(int(v)); + try go(k.gate_add, s, .{ .x = @intCast(rows), .y = ceil(width, 256) }, 256, &a); +} + +/// y fp16 = fp16(acc + b) from fp32 sums, then fp16(GELU) when `gelu`. +pub fn bias(k: *const Kernels, s: cuda.Stream, acc: u64, b: u64, y: u64, rows: usize, width: usize, gelu: bool) !void { + var a: cuda.Args = .{}; + addAll(&a, &.{ acc, b, y }); + for ([_]usize{ width, rows }) |v| a.add(int(v)); + a.add(@as(c_int, @intFromBool(gelu))); + try go(k.bias, s, .{ .x = @intCast(rows), .y = ceil(width, 256) }, 256, &a); +} + +/// x [rows, 1024] = fp16(fp16(z w^T + b) + pe), z [rows, 16] fp32. +pub fn embed(k: *const Kernels, s: cuda.Stream, z: u64, w: u64, b: u64, pe: u64, x: u64, rows: usize) !void { + var a: cuda.Args = .{}; + addAll(&a, &.{ z, w, b, pe, x }); + a.add(int(rows)); + try go(k.embed, s, .{ .x = @intCast(rows), .y = ceil(cfg.hidden, 256) }, 256, &a); +} + +/// y [rows, width] = fp16((a + ba) + (c + bc)). +pub fn cond(k: *const Kernels, s: cuda.Stream, a_: u64, ba: u64, c: u64, bc: u64, y: u64, rows: usize, width: usize) !void { + var a: cuda.Args = .{}; + addAll(&a, &.{ a_, ba, c, bc, y }); + for ([_]usize{ width, rows }) |v| a.add(int(v)); + try go(k.cond, s, .{ .x = @intCast(rows), .y = ceil(width, 256) }, 256, &a); +} + +/// y [rows, width] += b by column, fp32. +pub fn bias32(k: *const Kernels, s: cuda.Stream, y: u64, b: u64, rows: usize, width: usize) !void { + var a: cuda.Args = .{}; + addAll(&a, &.{ y, b }); + for ([_]usize{ width, rows }) |v| a.add(int(v)); + try go(k.bias32, s, .{ .x = @intCast(rows), .y = ceil(width, 256) }, 256, &a); +} + +/// out [rows, heads 64] = softmax(q k^T scale) v of packed rows [rows, q | k | v heads 64], fp16, every key. +pub fn attention(k: *const Kernels, s: cuda.Stream, qkv: u64, out: u64, rows: usize, heads: usize, scale: f32) !void { + if (rows == 0) return; + var a: cuda.Args = .{}; + addAll(&a, &.{ qkv, out }); + for ([_]usize{ rows, heads }) |v| a.add(int(v)); + a.add(scale); + const grid: cuda.Dim3 = .{ .x = ceil(rows, 16 * attention_warps), .y = @intCast(heads) }; + try cuda.launch.launch(k.attention, .{ .grid = grid, .block = .{ .x = 32 * attention_warps }, .shared = attention_smem }, s, &a); +} + +test "launch geometry: rows on grid x, 256-column strips on y, attention's tile under 48 KiB" { + try std.testing.expectEqual(@as(u32, 4), ceil(cfg.hidden, 256)); + try std.testing.expectEqual(@as(u32, 2), ceil(cfg.heads, 8)); + try std.testing.expectEqual(@as(u32, 97), ceil(cfg.joint_rows, 16 * attention_warps)); + try std.testing.expectEqual(@as(u32, 36864), attention_smem); + try std.testing.expect(attention_smem <= 48 * 1024); +} diff --git a/zig/src/families/idol/cuda_model.zig b/zig/src/families/idol/cuda_model.zig new file mode 100644 index 0000000000..b4dd9792f7 --- /dev/null +++ b/zig/src/families/idol/cuda_model.zig @@ -0,0 +1,288 @@ +//! Idol's denoiser on CUDA: the image condition refined once, then latent, condition and camera rows each step. +const std = @import("std"); +const cuda = @import("cuda"); +const cfg = @import("config.zig"); +const host = @import("host.zig"); +const wts = @import("cuda_weights.zig"); +const ks = @import("cuda_kernels.zig"); +const gm = @import("cuda_gemm.zig"); + +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 }; + +pub fn seconds() f64 { + const c = @import("sys").c; + var ts: c.timespec = undefined; + _ = c.clock_gettime(c.CLOCK.MONOTONIC, &ts); + return @as(f64, @floatFromInt(ts.sec)) + @as(f64, @floatFromInt(ts.nsec)) * 1e-9; +} + +/// The GPUs this family admits: GB10 (sm_121) and the RTX 50 / PRO Blackwell cards (sm_120), refused others by name. +pub fn admitted(sm: u32) bool { + return sm == 120 or sm == 121; +} + +/// A refined image condition: the condition refiner's rows [4101, 1024] fp16 on the device. +pub const Condition = struct { + rows: cuda.DeviceBuffer, + + pub fn deinit(c: *Condition) void { + c.rows.free(); + } +}; + +/// 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 * 3 * H * 2, L * H * 2, + L * H * 2, L * F * 2, L * F * 4, L * H * 4, L * 2 * cfg.repo_hidden * 4, + L * cfg.heads * 32 * 4, L * cfg.heads * 32 * 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, + d: *const cuda.Driver, + stream: cuda.Stream, + k: ks.Kernels, + gemms: gm.Gemms, + w: *wts.Weights, + scratch: cuda.DeviceBuffer, + /// A pass's start and end on the stream, for its GPU time. + marks: [2]cuda.Event, + // fp16 rows: the stream, a norm's output, a qkv projection, attention's packed heads, its output, a projection's + x: u64 = 0, + xn: u64 = 0, + qkv: u64 = 0, + qkh: u64 = 0, + att: u64 = 0, + y: u64 = 0, + /// the MLP's GELU rows, and a block linear's fp32 sums before its bias + u: u64 = 0, + acc: u64 = 0, + /// fp32 rows: RePo's norm (and the condition's two projections), its gate | content sums, cos and sin + ln: u64 = 0, + gc: u64 = 0, + cos: u64 = 0, + sin: u64 = 0, + lat: u64 = 0, + vel: u64 = 0, + mods: u64 = 0, + head: u64 = 0, + f1: u64 = 0, + f2: u64 = 0, + ones: u64 = 0, + pe: u64 = 0, + /// The model time `mods` and `head` hold, if any, and their host copies. + time: ?f32 = null, + mods_host: host.Mods = undefined, + head_host: [2][H]f32 = undefined, + /// The stream's rows downloaded for `stream`, and the camera row each forward reads back. + rows_host: []f16 = &.{}, + + /// The denoiser of snapshot folder `dir` on `ctx`'s GPU, if it is one this family admits. + pub fn load(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8) !*Model { + const sm = try ctx.capability(); + if (!admitted(sm)) { + var buf: [256]u8 = undefined; + const name = ctx.name(&buf) catch "this GPU"; + std.log.err("idol: {s} (sm_{d}) is not a GPU this family admits; it runs on sm_120 and sm_121 alone", .{ name, sm }); + return error.UncheckedGpu; + } + const d = ctx.d; + const t = try gpa.create(Model); + errdefer gpa.destroy(t); + t.* = .{ .gpa = gpa, .d = d, .stream = undefined, .k = undefined, .gemms = undefined, .w = undefined, .scratch = undefined, .marks = undefined }; + t.rows_host = try gpa.alloc(f16, L * H); + errdefer gpa.free(t.rows_host); + t.stream = try cuda.Stream.init(d, false); + errdefer t.stream.deinit(); + t.k = try ks.Kernels.load(d); + errdefer t.k.unload(); + t.gemms = try gm.Gemms.init(gpa, d); + errdefer t.gemms.deinit(); + t.w = try wts.load(gpa, io, d, dir); + errdefer { + t.w.deinit(); + gpa.destroy(t.w); + } + t.scratch = try cuda.DeviceBuffer.alloc(d, scratchBytes()); + errdefer t.scratch.free(); + t.marks[0] = try cuda.Event.init(d, true); + errdefer t.marks[0].deinit(); + t.marks[1] = try cuda.Event.init(d, true); + errdefer t.marks[1].deinit(); + var at: usize = t.scratch.ptr; + const views = [_]*u64{ &t.x, &t.xn, &t.qkv, &t.qkh, &t.att, &t.y, &t.u, &t.acc, &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.* = at; + at += std.mem.alignForward(usize, n, 256); + } + const ones = try gpa.alloc(f16, H); + defer gpa.free(ones); + @memset(ones, 1.0); + try t.upload(t.ones, std.mem.sliceAsBytes(ones)); + const pe = try gpa.alloc(f16, R * H); + defer gpa.free(pe); + host.positions(t.w.host.pos, pe); + try t.upload(t.pe, std.mem.sliceAsBytes(pe)); + return t; + } + + pub fn deinit(t: *Model) void { + for (&t.marks) |*m| m.deinit(); + t.scratch.free(); + t.w.deinit(); + t.gpa.destroy(t.w); + t.gemms.deinit(); + t.k.unload(); + t.stream.deinit(); + t.gpa.free(t.rows_host); + t.gpa.destroy(t); + } + + fn upload(t: *const Model, at: u64, bytes: []const u8) !void { + try t.d.check(t.d.api.cuMemcpyHtoDAsync_v2(at, bytes.ptr, bytes.len, t.stream.handle), "cuMemcpyHtoDAsync"); + try t.stream.synchronize(); + } + + fn download(t: *const Model, out: []u8, at: u64) !void { + try t.d.check(t.d.api.cuMemcpyDtoHAsync_v2(out.ptr, at, out.len, t.stream.handle), "cuMemcpyDtoHAsync"); + try t.stream.synchronize(); + } + + fn copy(t: *const Model, to: u64, from: u64, n: usize) !void { + try t.d.check(t.d.api.cuMemcpyDtoDAsync_v2(to, from, n, t.stream.handle), "cuMemcpyDtoDAsync"); + } + + /// 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) u64 { + return t.mods + ((i * host.vectors) + part) * H * 2; + } + + /// Every modulated block's vectors and the output's at model time `time`, unless they are already there. + fn modulate(t: *Model, time: f32) !void { + if (t.time == time) return; + t.time = null; + host.modulation(t.w.host, time, &t.mods_host, &t.head_host); + try t.upload(t.mods, std.mem.sliceAsBytes(&t.mods_host)); + try t.upload(t.head, std.mem.sliceAsBytes(&t.head_host)); + t.time = time; + } + + /// 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 cuda.DeviceBuffer.alloc(t.d, C * H * 2) }; + errdefer c.rows.free(); + try t.upload(t.f1, std.mem.sliceAsBytes(feature1)); + try t.upload(t.f2, std.mem.sliceAsBytes(feature2)); + const a = t.ln; + const b = t.ln + C * H * 4; + try t.gemms.run(t.stream, .f32, t.f1, t.w.cond, a, C, H, cfg.cond_channels); + try t.gemms.run(t.stream, .f32, t.f2, t.w.cond2, b, C, H, cfg.cond2_channels); + try ks.cond(&t.k, t.stream, a, t.w.cond_b, b, t.w.cond2_b, t.x, C, H); + for (t.w.context, t.w.context_repo) |blk, r| try t.block(blk, r, C, null); + try t.copy(c.rows.ptr, t.x, C * H * 2); + try t.stream.synchronize(); + 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 = seconds(); + try t.modulate(time); + try t.upload(t.lat, std.mem.sliceAsBytes(latent)); + const camera_row = t.rows_host[0..H]; + host.cameraToken(t.w.host, camera, camera_row); + try t.upload(t.x + (L - 1) * H * 2, std.mem.sliceAsBytes(camera_row)); + const spent = seconds() - started; + try t.marks[0].record(t.stream); + try t.copy(t.x + R * H * 2, c.rows.ptr, C * H * 2); + try ks.embed(&t.k, t.stream, t.lat, t.w.input, t.w.input_b, t.pe, t.x, R); + for (t.w.noise, t.w.noise_repo, 0..) |blk, r, i| try t.block(blk, r, R, i); + for (t.w.blocks, t.w.repo, 0..) |blk, r, i| try t.block(blk, r, L, cfg.refiners + i); + try ks.norm32(&t.k, t.stream, t.x, t.head + H * 4, t.head, t.ln, R, cfg.default_eps, .modulated); + try t.gemms.run(t.stream, .f32, t.ln, t.w.out, t.vel, R, cfg.latent_channels, H); + try ks.bias32(&t.k, t.stream, t.vel, t.w.out_b, R, cfg.latent_channels); + try t.marks[1].record(t.stream); + try t.download(std.mem.sliceAsBytes(velocity), t.vel); + try t.download(std.mem.sliceAsBytes(camera_row), t.x + (L - 1) * H * 2); + host.cameraVelocity(t.w.host, camera_row, &t.head_host, camera_velocity); + const gpu = @as(f64, try cuda.Event.elapsedMs(t.marks[0], t.marks[1])) / 1000; + return .{ .wall = seconds() - started, .gpu = gpu, .host = spent }; + } + + /// 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; + try t.modulate(time); + try t.upload(t.x, std.mem.sliceAsBytes(rows)); + if (index < cfg.refiners) { + try t.block(t.w.noise[index], t.w.noise_repo[index], n, index); + } else try t.block(t.w.blocks[index - cfg.refiners], t.w.repo[index - cfg.refiners], n, index); + try t.stream.synchronize(); + } + + /// The residual stream's first `rows` rows as the last pass left them, fp16 (valid until the next call). + pub fn residual(t: *Model, rows: usize) ![]const f16 { + const out = t.rows_host[0 .. rows * H]; + try t.download(std.mem.sliceAsBytes(out), t.x); + return out; + } + + /// y fp16 [rows, n] = fp16(x w^T + b) from fp32 sums, then fp16(GELU) when `gelu`. + fn linear(t: *Model, x: u64, w: u64, b: u64, y: u64, rows: usize, n: usize, k: usize, gelu: bool) !void { + try t.gemms.run(t.stream, .f16, x, w, t.acc, rows, n, k); + try ks.bias(&t.k, t.stream, t.acc, b, y, rows, n, gelu); + } + + /// One block over the stream's first `rows` rows: modulated block `mod`, or the condition refiner's (null). + fn block(t: *Model, b: wts.Block, r: wts.Repo, rows: usize, mod: ?usize) !void { + const k = &t.k; + const s = t.stream; + // RePo: angles from the block's input rows, its projections in fp32 + try ks.norm32(k, s, t.x, r.norm, r.norm_b, t.ln, rows, cfg.default_eps, .affine); + try t.gemms.run(s, .f32, t.ln, r.maps, t.gc, rows, 2 * cfg.repo_hidden, H); + try ks.repo(k, s, t.gc, r.final, r.freqs, t.cos, t.sin, rows); + if (mod) |m| { + try ks.norm(k, s, t.x, 0, 0, t.vec(m, 1), t.vec(m, 0), t.xn, rows, cfg.block_eps, .modulated); + } else try ks.norm(k, s, t.x, b.norm1, b.norm1_b, t.ones, t.ones, t.xn, rows, cfg.block_eps, .affine); + try t.linear(t.xn, b.qkv, b.qkv_b, t.qkv, rows, 3 * H, H, false); + try ks.qk(k, s, t.qkv, b.q_norm, b.k_norm, t.cos, t.sin, t.qkh, rows); + try ks.attention(k, s, t.qkh, t.att, rows, cfg.heads, 1.0 / @sqrt(@as(f32, cfg.head_dim))); + try t.linear(t.att, b.out, b.out_b, t.y, rows, H, H, false); + try ks.gateAdd(k, s, t.x, t.y, if (mod) |m| t.vec(m, 2) else t.ones, rows, H); + if (mod) |m| { + try ks.norm(k, s, t.x, 0, 0, t.vec(m, 4), t.vec(m, 3), t.xn, rows, cfg.block_eps, .modulated); + } else try ks.norm(k, s, t.x, b.norm2, b.norm2_b, t.ones, t.ones, t.xn, rows, cfg.block_eps, .affine); + try t.linear(t.xn, b.up, b.up_b, t.u, rows, F, H, true); + try t.linear(t.u, b.down, b.down_b, t.y, rows, H, F, false); + try ks.gateAdd(k, s, t.x, t.y, if (mod) |m| t.vec(m, 5) else t.ones, rows, H); + } +}; + +test "scratch for the joint rows: every part 256-aligned, the condition's two projections fit RePo's norm rows" { + try std.testing.expectEqual(@as(usize, 12294), L); + try std.testing.expect(2 * C * H * 4 <= parts[8]); + try std.testing.expect(scratchBytes() < 1 << 30); + // the kernels index rows and offsets as c_int + try std.testing.expect(L * F < std.math.maxInt(c_int)); +} + +test "admission: sm_120 and sm_121 alone" { + try std.testing.expect(admitted(121) and admitted(120)); + for ([_]u32{ 80, 86, 89, 90, 100, 103, 122 }) |sm| try std.testing.expect(!admitted(sm)); +} diff --git a/zig/src/families/idol/cuda_weights.zig b/zig/src/families/idol/cuda_weights.zig new file mode 100644 index 0000000000..82fb1e97da --- /dev/null +++ b/zig/src/families/idol/cuda_weights.zig @@ -0,0 +1,191 @@ +//! Idol's denoiser tensors on CUDA: blocks' linears and biases to fp16 as idol-sdk does, the rest fp32. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const cfg = @import("config.zig"); +const host = @import("host.zig"); + +const H = cfg.hidden; + +/// A block's device addresses: fp16 linears and biases, fp32 head RMSNorm gammas; fp32 norms in the refiner alone. +pub const Block = struct { qkv: u64, qkv_b: u64, out: u64, out_b: u64, up: u64, up_b: u64, down: u64, down_b: u64, q_norm: u64, k_norm: u64, norm1: u64 = 0, norm1_b: u64 = 0, norm2: u64 = 0, norm2_b: u64 = 0 }; + +/// A RePo layer: fp32 norm, gate | content maps [256, 1024], final map [48, 128] and its 32 frequencies. +pub const Repo = struct { norm: u64, norm_b: u64, maps: u64, final: u64, freqs: u64 }; + +pub const Weights = struct { + gpa: std.mem.Allocator, + ck: core.Checkpoint, + buffers: std.ArrayList(cuda.DeviceBuffer) = .empty, + noise: [cfg.refiners]Block = undefined, + context: [cfg.refiners]Block = undefined, + blocks: [cfg.layers]Block = undefined, + noise_repo: [cfg.refiners]Repo = undefined, + context_repo: [cfg.refiners]Repo = undefined, + repo: [cfg.layers]Repo = undefined, + input: u64 = 0, + input_b: u64 = 0, + cond: u64 = 0, + cond_b: u64 = 0, + cond2: u64 = 0, + cond2_b: u64 = 0, + out: u64 = 0, + out_b: u64 = 0, + /// The host's tensors, mapped from the checkpoint (host.modulation, cameraToken, cameraVelocity, positions). + host: host.Host = undefined, + bytes: u64 = 0, + + pub fn deinit(w: *Weights) void { + for (w.buffers.items) |*b| b.free(); + w.buffers.deinit(w.gpa); + w.ck.close(); + } + + fn typed(w: *Weights, name: []const u8, shape: []const usize) !core.checkpoint.Tensor { + const t = try w.ck.get(name); + 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; + } + if (@intFromPtr(t.bytes.ptr) % 4 != 0) return error.UnalignedTensor; + return t; + } + + /// The fp32 tensor `name` of `shape`, its values on the host (mapped from the checkpoint). + fn floats(w: *Weights, name: []const u8, shape: []const usize) ![]const f32 { + const t = try w.typed(name, shape); + return @alignCast(std.mem.bytesAsSlice(f32, t.bytes)); + } + + fn upload(w: *Weights, d: *const cuda.Driver, bytes: []const u8) !u64 { + var b = try cuda.DeviceBuffer.fromHost(d, bytes); + errdefer b.free(); + try w.buffers.append(w.gpa, b); + w.bytes += b.len; + return b.ptr; + } + + /// The fp32 tensor `name` of `shape` on the device as stored. + fn take(w: *Weights, d: *const cuda.Driver, name: []const u8, shape: []const usize) !u64 { + return w.upload(d, std.mem.sliceAsBytes(try w.floats(name, shape))); + } + + /// The fp32 tensor `name` of `shape` on the device rounded to fp16 (`stage` holds the halves). + fn half(w: *Weights, d: *const cuda.Driver, stage: []f16, name: []const u8, shape: []const usize) !u64 { + const src = try w.floats(name, shape); + for (stage[0..src.len], src) |*o, v| o.* = @floatCast(v); + return w.upload(d, std.mem.sliceAsBytes(stage[0..src.len])); + } +}; + +fn block(w: *Weights, d: *const cuda.Driver, stage: []f16, prefix: []const u8, i: usize, affine: bool) !Block { + var name: [96]u8 = undefined; + const N = 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 w.half(d, stage, try N.n(&name, prefix, i, "attn.qkv.weight"), &.{ 3 * H, H }); + b.qkv_b = try w.half(d, stage, try N.n(&name, prefix, i, "attn.qkv.bias"), &.{3 * H}); + b.out = try w.half(d, stage, try N.n(&name, prefix, i, "attn.out.weight"), &.{ H, H }); + b.out_b = try w.half(d, stage, try N.n(&name, prefix, i, "attn.out.bias"), &.{H}); + b.up = try w.half(d, stage, try N.n(&name, prefix, i, "mlp.mlp.0.weight"), &.{ cfg.ffn, H }); + b.up_b = try w.half(d, stage, try N.n(&name, prefix, i, "mlp.mlp.0.bias"), &.{cfg.ffn}); + b.down = try w.half(d, stage, try N.n(&name, prefix, i, "mlp.mlp.2.weight"), &.{ H, cfg.ffn }); + b.down_b = try w.half(d, stage, try N.n(&name, prefix, i, "mlp.mlp.2.bias"), &.{H}); + b.q_norm = try w.take(d, try N.n(&name, prefix, i, "attn.q_norm.gamma"), &.{ cfg.heads, cfg.head_dim }); + b.k_norm = try w.take(d, try N.n(&name, prefix, i, "attn.k_norm.gamma"), &.{ cfg.heads, cfg.head_dim }); + b.norm1 = 0; + b.norm1_b = 0; + b.norm2 = 0; + b.norm2_b = 0; + if (affine) { + b.norm1 = try w.take(d, try N.n(&name, prefix, i, "norm1.weight"), &.{H}); + b.norm1_b = try w.take(d, try N.n(&name, prefix, i, "norm1.bias"), &.{H}); + b.norm2 = try w.take(d, try N.n(&name, prefix, i, "norm2.weight"), &.{H}); + b.norm2_b = try w.take(d, try N.n(&name, prefix, i, "norm2.bias"), &.{H}); + } + return b; +} + +/// A RePo layer: its gate and content maps stacked [256, 1024] fp32, its three frequency vectors as 32 values. +fn repo(w: *Weights, d: *const cuda.Driver, stage: []f32, prefix: []const u8, i: usize) !Repo { + var name: [96]u8 = undefined; + var r: Repo = undefined; + r.norm = try w.take(d, try std.fmt.bufPrint(&name, "{s}.{d}.norm.weight", .{ prefix, i }), &.{H}); + r.norm_b = try w.take(d, try std.fmt.bufPrint(&name, "{s}.{d}.norm.bias", .{ prefix, i }), &.{H}); + const n = cfg.repo_hidden * H; + @memcpy(stage[0..n], try w.floats(try std.fmt.bufPrint(&name, "{s}.{d}.gate_map.weight", .{ prefix, i }), &.{ cfg.repo_hidden, H })); + @memcpy(stage[n .. 2 * n], try w.floats(try std.fmt.bufPrint(&name, "{s}.{d}.content_map.weight", .{ prefix, i }), &.{ cfg.repo_hidden, H })); + r.maps = try w.upload(d, std.mem.sliceAsBytes(stage[0 .. 2 * n])); + r.final = try w.take(d, try std.fmt.bufPrint(&name, "{s}.{d}.final_map.weight", .{ prefix, i }), &.{ 3 * cfg.heads, cfg.repo_hidden }); + var freqs: [cfg.repo_pairs]f32 = undefined; + var at: usize = 0; + inline for (.{ .{ "freqs_0", 10 }, .{ "freqs_1", 10 }, .{ "freqs_2", 12 } }) |f| { + @memcpy(freqs[at..][0..f[1]], try w.floats(try std.fmt.bufPrint(&name, "{s}.{d}.{s}", .{ prefix, i, f[0] }), &.{f[1]})); + at += f[1]; + } + r.freqs = try w.upload(d, std.mem.sliceAsBytes(&freqs)); + return r; +} + +/// The snapshot folder `dir`: config.json checked, then denoiser.safetensors's every tensor, none left unread. +pub fn load(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, 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); + const w = try gpa.create(Weights); + errdefer gpa.destroy(w); + w.* = .{ .gpa = gpa, .ck = .{ .gpa = gpa, .io = io } }; + errdefer w.deinit(); + try w.ck.add(dir, "denoiser.safetensors"); + const halves = try arena.alloc(f16, cfg.ffn * H); + const stage = try arena.alloc(f32, 2 * cfg.repo_hidden * H); + for (&w.noise, &w.context, &w.noise_repo, &w.context_repo, 0..) |*n, *x, *nr, *xr, i| { + n.* = try block(w, d, halves, "noise_refiner", i, false); + x.* = try block(w, d, halves, "context_refiner", i, true); + nr.* = try repo(w, d, stage, "noise_repo_layers", i); + xr.* = try repo(w, d, stage, "context_repo_layers", i); + } + for (&w.blocks, &w.repo, 0..) |*b, *r, i| { + b.* = try block(w, d, halves, "blocks", i, false); + r.* = try repo(w, d, stage, "repo_layers", i); + } + w.input = try w.take(d, "input_layer.weight", &.{ H, cfg.latent_channels }); + w.input_b = try w.take(d, "input_layer.bias", &.{H}); + w.cond = try w.take(d, "cond_embedder.weight", &.{ H, cfg.cond_channels }); + w.cond_b = try w.take(d, "cond_embedder.bias", &.{H}); + w.cond2 = try w.take(d, "cond_embedder2.weight", &.{ H, cfg.cond2_channels }); + w.cond2_b = try w.take(d, "cond_embedder2.bias", &.{H}); + w.out = try w.take(d, "out_layer.weight", &.{ cfg.latent_channels, H }); + w.out_b = try w.take(d, "out_layer.bias", &.{cfg.latent_channels}); + var hw: host.Host = undefined; + hw.t1 = try w.floats("t_embedder.mlp.0.weight", &.{ H, cfg.time_dim }); + hw.t1_b = try w.floats("t_embedder.mlp.0.bias", &.{H}); + hw.t2 = try w.floats("t_embedder.mlp.2.weight", &.{ H, H }); + hw.t2_b = try w.floats("t_embedder.mlp.2.bias", &.{H}); + hw.ada = try w.floats("adaLN_modulation.1.weight", &.{ 6 * H, H }); + hw.ada_b = try w.floats("adaLN_modulation.1.bias", &.{6 * H}); + hw.cam1 = try w.floats("cam_refiner.mlp.0.weight", &.{ H, cfg.camera_channels }); + hw.cam1_b = try w.floats("cam_refiner.mlp.0.bias", &.{H}); + hw.cam2 = try w.floats("cam_refiner.mlp.2.weight", &.{ H, H }); + hw.cam2_b = try w.floats("cam_refiner.mlp.2.bias", &.{H}); + hw.cam_out = try w.floats("cam_out_layer.weight", &.{ cfg.camera_channels, H }); + hw.cam_out_b = try w.floats("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 w.floats(n, &.{ 1, 6 * H }); + } + hw.head = try w.floats("shift_table", &.{ 1, 2, H }); + hw.pos = try w.floats("pos_pe", &.{ 1, cfg.latent_rows, 3 }); + w.host = hw; + if (w.ck.unused() != 0) { + std.log.err("idol: denoiser.safetensors has {d} tensors this family does not read", .{w.ck.unused()}); + return error.UnexpectedTensors; + } + return w; +} diff --git a/zig/src/families/idol/host.zig b/zig/src/families/idol/host.zig index 33dd008649..a6d7de560b 100644 --- a/zig/src/families/idol/host.zig +++ b/zig/src/families/idol/host.zig @@ -1,10 +1,30 @@ //! 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; +/// 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 fn half(x: f32) f32 { return @as(f16, @floatCast(x)); } diff --git a/zig/src/families/idol/weights.zig b/zig/src/families/idol/weights.zig index 27d872168c..4dcee77a12 100644 --- a/zig/src/families/idol/weights.zig +++ b/zig/src/families/idol/weights.zig @@ -36,26 +36,8 @@ fn split(src: []const f32, scale: f32, dst: []f16) void { } } -/// 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, -}; +/// The fp32 vectors and small matrices the host reads each timestep (host.zig, shared with CUDA). +pub const Host = @import("host.zig").Host; pub const Weights = struct { ck: ck.Checkpoint, diff --git a/zig/tests/cuda/contract_attn.zig b/zig/tests/cuda/contract_attn.zig index afa284d55c..11ad3705b9 100644 --- a/zig/tests/cuda/contract_attn.zig +++ b/zig/tests/cuda/contract_attn.zig @@ -24,6 +24,8 @@ pub const Ref = struct { want: []f64, tol: []f64, msg: [240]u8 = undefined, + /// The outputs' type as the kernel writes them (Idol's attention writes fp16). + out: c.Elem = .bf16, pub fn init(a: std.mem.Allocator, d: usize, most: usize) !Ref { const keys = try a.alloc(usize, most); @@ -54,7 +56,7 @@ pub const Ref = struct { try t.load(i, &r.head); ref.head(r.head.q, r.head.k, r.head.v, r.keys[0..r.head.n], r.d, scale, r.want, r.tol); for (r.want, r.tol, 0..) |w, tl, e| { - const got = c.read(.bf16, out, i * r.d + e); + const got = c.read(r.out, out, i * r.d + e); if (c.classOf(w) != .finite) planted += 1; if (c.classOf(got) != c.classOf(w)) return try std.fmt.bufPrint(&r.msg, "head {d} element {d}: {t} where float64 gives {t}", .{ i, e, c.classOf(got), c.classOf(w) }); if (at != .nonfinite and std.math.isFinite(w) and !(@abs(got - w) <= tl)) return try std.fmt.bufPrint(&r.msg, "head {d} element {d} is {e:.6} where float64 gives {e:.6}: {d:.2} of attention_ref's bound", .{ i, e, got, w, @abs(got - w) / tl }); diff --git a/zig/tests/cuda/contract_family_sets.zig b/zig/tests/cuda/contract_family_sets.zig index 1e8d024f7f..9635cba38a 100644 --- a/zig/tests/cuda/contract_family_sets.zig +++ b/zig/tests/cuda/contract_family_sets.zig @@ -1,4 +1,4 @@ -//! The family sets in the full kernel contract: Nemotron's pocket, MiniMax-M3's, MOSS-Image's. +//! The family sets in the full kernel contract: Nemotron's pocket, MiniMax-M3's, MOSS-Image's, Idol's. const std = @import("std"); const cuda = @import("cuda"); const check = @import("check.zig"); @@ -8,6 +8,7 @@ pub const Sets = struct { pocket: *@import("contract_pocket.zig").Set, minimax: *@import("contract_minimax.zig").Set, moss: *@import("contract_moss_image.zig").Set, + idol: *@import("contract_idol.zig").Set, /// Each set's kernels appended to `list`, in this order. pub fn init(gpu: check.Gpu, s: cuda.Stream, a: std.mem.Allocator, list: *std.ArrayList(g.Kernel)) !Sets { @@ -16,10 +17,13 @@ pub const Sets = struct { const minimax = try @import("contract_minimax.zig").Set.init(gpu, s, a, list); errdefer minimax.deinit(); const moss = try @import("contract_moss_image.zig").Set.init(gpu, s, a, list); - return .{ .pocket = pocket, .minimax = minimax, .moss = moss }; + errdefer moss.deinit(); + const idol = try @import("contract_idol.zig").Set.init(gpu, s, a, list); + return .{ .pocket = pocket, .minimax = minimax, .moss = moss, .idol = idol }; } pub fn deinit(sets: Sets) void { + sets.idol.deinit(); sets.moss.deinit(); sets.minimax.deinit(); sets.pocket.deinit(); diff --git a/zig/tests/cuda/contract_idol.zig b/zig/tests/cuda/contract_idol.zig new file mode 100644 index 0000000000..f8de4c6298 --- /dev/null +++ b/zig/tests/cuda/contract_idol.zig @@ -0,0 +1,123 @@ +//! Idol's denoiser kernels (idol.cu) under the contract, with controls that drop one fp16 rounding (idol-contract). +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const idol = @import("idol_cuda"); +const check = @import("check.zig"); +const case = @import("contract_idol_case.zig"); +const norms = @import("contract_idol_norms.zig"); +const steps = @import("contract_idol_steps.zig"); +const rotary = @import("contract_idol_rotary.zig"); +const attn = @import("contract_idol_attn.zig"); + +const g = core.contract.gates; + +const controls = struct { + pub const bytes align(16) = @embedFile("fatbin_contract_controls").*; +}; + +/// Every registration, on the heap: the gates hold pointers to them. +pub const Set = struct { + gpa: std.mem.Allocator, + k: idol.Kernels, + ctl: cuda.Module, + /// The kernels with one function swapped for its control. + late_norm: idol.Kernels = undefined, + once_add: idol.Kernels = undefined, + late_bias: idol.Kernels = undefined, + norms: [2]case.Case(norms.Norm) = undefined, + norm32s: [2]case.Case(norms.Norm32) = undefined, + adds: case.Case(steps.GateAdd) = undefined, + biases: [2]case.Case(steps.Bias) = undefined, + cond: case.Case(steps.Cond) = undefined, + bias32: case.Case(steps.Bias32) = undefined, + qk: case.Case(rotary.Qk) = undefined, + repo: case.Case(rotary.Repo) = undefined, + embed: case.Case(rotary.Embed) = undefined, + attn: attn.Attn = undefined, + ctl_norm: case.Case(norms.Norm) = undefined, + ctl_add: case.Case(steps.GateAdd) = undefined, + ctl_bias: case.Case(steps.Bias) = undefined, + + /// The set, its kernels appended to `list` (the controls last). + pub fn init(gpu: check.Gpu, s: cuda.Stream, a: std.mem.Allocator, list: *std.ArrayList(g.Kernel)) !*Set { + const set = try gpu.gpa.create(Set); + errdefer gpu.gpa.destroy(set); + var k = try idol.Kernels.load(gpu.d); + errdefer k.unload(); + var ctl = try cuda.Module.load(gpu.d, &controls.bytes); + errdefer ctl.unload(); + set.* = .{ .gpa = gpu.gpa, .k = k, .ctl = ctl }; + set.late_norm = k; + set.late_norm.norm = try ctl.function("tf_ctl_idol_norm_late"); + set.once_add = k; + set.once_add.gate_add = try ctl.function("tf_ctl_idol_gate_add_once"); + set.late_bias = k; + set.late_bias.bias = try ctl.function("tf_ctl_idol_bias_late"); + const rig: @import("contract_rig.zig").Rig = .{ .gpu = gpu, .s = s }; + const norm_names = [_][]const u8{ "cuda idol tf_idol_norm, LayerNorm32 affine to fp16", "cuda idol tf_idol_norm, LayerNorm32 to fp16, modulated" }; + const norm32_names = [_][]const u8{ "cuda idol tf_idol_norm32, LayerNorm affine in fp32", "cuda idol tf_idol_norm32, LayerNorm times 1 + w in fp32" }; + for (&set.norms, norm_names, 1..) |*x, n, v| x.* = case.make(norms.Norm, rig, k, n, @intCast(v), null); + for (&set.norm32s, norm32_names, 1..) |*x, n, v| x.* = case.make(norms.Norm32, rig, k, n, @intCast(v), null); + set.adds = case.make(steps.GateAdd, rig, k, "cuda idol tf_idol_gate_add, x + fp16(g y)", 0, null); + set.biases[0] = case.make(steps.Bias, rig, k, "cuda idol tf_idol_bias, fp16(acc + b)", 0, null); + set.biases[1] = case.make(steps.Bias, rig, k, "cuda idol tf_idol_bias, fp16(GELU(fp16(acc + b)))", 1, null); + set.cond = case.make(steps.Cond, rig, k, "cuda idol tf_idol_cond, the condition projections' sum", 0, null); + set.bias32 = case.make(steps.Bias32, rig, k, "cuda idol tf_idol_bias32, fp32 bias in place", 0, null); + set.qk = case.make(rotary.Qk, rig, k, "cuda idol tf_idol_qk, RePo turn, F.normalize, gamma", 0, null); + set.repo = case.make(rotary.Repo, rig, k, "cuda idol tf_idol_repo, final map and angles", 0, null); + set.embed = case.make(rotary.Embed, rig, k, "cuda idol tf_idol_embed, input projection and position", 0, null); + set.attn = .{ .rig = rig, .k = &set.k, .name = "cuda idol tf_idol_attention, full<64, 8>, fp16" }; + set.ctl_norm = case.make(norms.Norm, rig, set.late_norm, "control: idol modulated norm unrounded before scale", 2, .shapes); + set.ctl_add = case.make(steps.GateAdd, rig, set.once_add, "control: idol gate_add without fp16(g y)", 0, .shapes); + set.ctl_bias = case.make(steps.Bias, rig, set.late_bias, "control: idol GELU of the unrounded linear", 1, .shapes); + for (&set.norms) |*x| try list.append(a, x.kernel()); + for (&set.norm32s) |*x| try list.append(a, x.kernel()); + try list.append(a, set.adds.kernel()); + for (&set.biases) |*x| try list.append(a, x.kernel()); + try list.append(a, set.cond.kernel()); + try list.append(a, set.bias32.kernel()); + try list.append(a, set.qk.kernel()); + try list.append(a, set.repo.kernel()); + try list.append(a, set.embed.kernel()); + try list.append(a, set.attn.kernel()); + try list.append(a, set.ctl_norm.kernel()); + try list.append(a, set.ctl_add.kernel()); + try list.append(a, set.ctl_bias.kernel()); + return set; + } + + pub fn deinit(set: *Set) void { + for (&set.norms) |*x| x.deinit(); + for (&set.norm32s) |*x| x.deinit(); + set.adds.deinit(); + for (&set.biases) |*x| x.deinit(); + set.cond.deinit(); + set.bias32.deinit(); + set.qk.deinit(); + set.repo.deinit(); + set.embed.deinit(); + set.attn.deinit(); + set.ctl_norm.deinit(); + set.ctl_add.deinit(); + set.ctl_bias.deinit(); + set.ctl.unload(); + set.k.unload(); + set.gpa.destroy(set); + } +}; + +/// Idol's kernels' gates 1-4 and their controls alone. +pub fn run(gpu: check.Gpu) !void { + const a = gpu.gpa; + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var list: std.ArrayList(g.Kernel) = .empty; + defer list.deinit(a); + const set = try Set.init(gpu, s, a, &list); + defer set.deinit(); + var t: g.Tally = .{}; + try g.runAll(a, list.items, 0x1d01, &t); + std.debug.print("SUMMARY idol: {d} kernels, {d} pass gates 1-4, {d} fail; {d} controls: {d} red as planted, {d} missed\n", .{ t.kernels, t.passed, t.failed, t.controls, t.caught, t.missed }); + try check.expect(t.problems() == 0, "idol-contract: {d} gate failures, {d} controls missed", .{ t.failed, t.missed }); +} diff --git a/zig/tests/cuda/contract_idol_attn.zig b/zig/tests/cuda/contract_idol_attn.zig new file mode 100644 index 0000000000..a964278f95 --- /dev/null +++ b/zig/tests/cuda/contract_idol_attn.zig @@ -0,0 +1,153 @@ +//! tf_idol_attention under the contract: every row over every key, head width 64, fp16 in and out, head by head. +const std = @import("std"); +const cuda = @import("cuda"); +const idol = @import("idol_cuda"); +const rig_ = @import("contract_rig.zig"); +const attn = @import("contract_attn.zig"); +const r = @import("contract_idol_ref.zig"); + +const c = r.c; +const g = r.g; + +const why = "P enters the P x V mma as fp16: gates 1 and 3 hold each output to attention_ref's float64 bound (bf16 weights)"; + +/// A row's P x V chain: every key, a 64-key tile at a time, then the 4 lanes' fold. +fn keyChain(s: c.Shape) usize { + return s.k + 64 + 2; +} + +/// Heads (rows) over tokens (n, each a query and key): served token counts at one head, the head strides at 16. +const served = [_]c.Shape{ .{ .rows = 1, .n = r.cond, .k = r.cond }, .{ .rows = 1, .n = r.latent, .k = r.latent }, .{ .rows = 1, .n = r.joint, .k = r.joint }, .{ .rows = r.heads, .n = 300, .k = 300 } }; + +pub fn spec(name: []const u8) c.Spec { + return .{ .name = name, .accum = .f32, .depth = keyChain, .ops = 6, .out = .f16, .tiles = .{ .{}, .{ .size = 16 * idol.attention_warps }, .{} }, .served = &served, .sums = .{ .why = why } }; +} + +fn half(x: f64) u16 { + return @bitCast(@as(f16, @floatCast(@as(f32, @floatCast(x))))); +} + +fn value(b: u16) f64 { + return @as(f16, @bitCast(b)); +} + +pub const Attn = struct { + rig: rig_.Rig, + k: *const idol.Kernels, + name: []const u8, + ref: ?attn.Ref = null, + heads: usize = 0, + tokens: usize = 0, + qkv: std.ArrayList(u16) = .empty, + /// The key head whose K and V rows the reference holds now. + held: ?usize = null, + dev: u64 = 0, + out: u64 = 0, + + pub fn deinit(t: *Attn) void { + if (t.ref) |*x| x.deinit(); + t.qkv.deinit(t.rig.gpu.gpa); + t.rig.deinit(); + } + + pub fn kernel(t: *Attn) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt }; + } + const vt = g.hooks(Attn); + + fn stride(t: *const Attn) usize { + return 3 * t.heads * r.D; + } + + pub fn prepare(t: *Attn, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0) return false; + t.rig.reset(); + const a = t.rig.gpu.gpa; + t.heads = shape.rows; + t.tokens = shape.n; + t.held = null; + try t.qkv.resize(a, t.tokens * t.stride()); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (t.qkv.items) |*e| e.* = half(rng.floatNorm(f64)); + if (t.ref) |*x| x.deinit(); + t.ref = null; + t.ref = try attn.Ref.init(a, r.D, t.tokens); + t.ref.?.out = .f16; + switch (fill) { + .random => {}, + .nonfinite => t.plant(), + .stress => t.stress(rng), + } + t.dev = try t.rig.input(std.mem.sliceAsBytes(t.qkv.items)); + t.out = try t.rig.output(t.count() * 2); + return true; + } + + fn at(t: *Attn, row: usize, part: usize, head: usize, e: usize) *u16 { + return &t.qkv.items[row * t.stride() + (part * t.heads + head) * r.D + e]; + } + + /// NaN in row 0's query (head 0); +inf and -inf in one value column of the last head, at two keys. + fn plant(t: *Attn) void { + const last = t.heads - 1; + t.at(0, 0, 0, 3).* = half(r.nan); + t.at(t.tokens / 2, 2, last, 7).* = half(r.inf); + t.at(t.tokens - 1, 2, last, 7).* = half(-r.inf); + t.at(t.tokens / 3, 2, last, 9).* = half(r.inf); + } + + /// Gate 4: every key of a head alike (each weight 1), values small odd integers mirrored over the keys. + fn stress(t: *Attn, rng: std.Random) void { + const n = t.tokens; + for (0..t.heads) |hd| for (0..r.D) |e| { + const key = t.at(0, 1, hd, e).*; + for (0..n) |row| t.at(row, 1, hd, e).* = key; + for (0..n / 2) |row| { + const v: f64 = @floatFromInt(1 + 2 * rng.uintLessThan(u32, 8)); + t.at(row, 2, hd, e).* = half(v); + t.at(n / 2 + row, 2, hd, e).* = half((if (row < 8) @as(f64, 1) else 0) - v); + } + if (n % 2 == 1) t.at(n - 1, 2, hd, e).* = half(0); + }; + } + + pub fn run(t: *Attn) !void { + try idol.attention(t.k, t.rig.s, t.dev, t.out, t.tokens, t.heads, 1.0 / 8.0); + try t.rig.s.synchronize(); + } + + pub fn count(t: *Attn) usize { + return t.tokens * t.heads * r.D; + } + + /// Output head i (token i / heads, head i % heads): its query and the head's every key and value. + pub fn load(t: *Attn, i: usize, hd: *attn.Head) !void { + const row = i / t.heads; + const head = i % t.heads; + for (hd.q, 0..) |*x, e| x.* = value(t.at(row, 0, head, e).*); + hd.n = t.tokens; + if (t.held == head) return; + for (0..t.tokens) |j| for (0..r.D) |e| { + hd.k[j * r.D + e] = value(t.at(j, 1, head, e).*); + hd.v[j * r.D + e] = value(t.at(j, 2, head, e).*); + }; + t.held = head; + } + + fn outputs(t: *Attn) usize { + return t.tokens * t.heads; + } + + pub fn check(t: *Attn, gate: g.Gate) !?[]const u8 { + return t.ref.?.check(t, t.outputs(), 1.0 / 8.0, try t.rig.outputBytes(), gate); + } + + pub fn reference(t: *Attn, want: []f64, mag: []f64) !void { + return t.ref.?.reference(t, t.outputs(), 1.0 / 8.0, want, mag); + } + + pub fn simulate(t: *Attn, acc: c.Accum, order: g.Order, out: []f64) !void { + return t.ref.?.simulate(t, t.outputs(), 1.0 / 8.0, acc, order, out); + } +}; diff --git a/zig/tests/cuda/contract_idol_case.zig b/zig/tests/cuda/contract_idol_case.zig new file mode 100644 index 0000000000..46c83434fd --- /dev/null +++ b/zig/tests/cuda/contract_idol_case.zig @@ -0,0 +1,135 @@ +//! A test of one Idol kernel under the contract, from its `Op`: inputs, launch and float64 reference with ranges. +const std = @import("std"); +const cuda = @import("cuda"); +const idol = @import("idol_cuda"); +const rig_ = @import("contract_rig.zig"); +const r = @import("contract_idol_ref.zig"); + +const c = r.c; +const g = r.g; + +/// The gates' hooks over an Op: its kinds, out, in_place (input the output starts as), spec, sizes, fill, launch, ref. +pub fn Case(comptime Op: type) type { + const n = Op.kinds.len; + return struct { + const Self = @This(); + rig: rig_.Rig, + k: idol.Kernels, + name: []const u8, + /// The Op's own form (a norm's mode, the bias's GELU), passed to each of its functions. + variant: u32 = 0, + control: ?g.Gate = null, + shape: c.Shape = undefined, + fill: g.Fill = .random, + vals: [n][]f64 = @splat(&.{}), + dev: [n]u64 = @splat(0), + out: u64 = 0, + ends: r.Ends, + + pub fn deinit(t: *Self) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + + fn free(t: *Self) void { + for (&t.vals) |*v| { + t.rig.gpu.gpa.free(v.*); + v.* = &.{}; + } + } + + pub fn kernel(t: *Self) g.Kernel { + return .{ .spec = Op.spec(t.name, t.variant), .ctx = t, .vt = &vt, .control = t.control }; + } + const vt = g.hooks(Self); + + pub fn prepare(t: *Self, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (!Op.accepts(shape)) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + t.fill = fill; + const a = t.rig.gpu.gpa; + for (&t.vals, Op.sizes(shape)) |*v, len| { + v.* = try a.alloc(f64, len); + @memset(v.*, 0); + } + var prng = std.Random.DefaultPrng.init(seed); + Op.fill(prng.random(), fill, shape, t.variant, &t.vals); + for (&t.vals, &t.dev, Op.kinds) |v, *d, kind| { + const bytes = try encode(a, kind, v); + defer a.free(bytes); + d.* = try t.rig.input(bytes); + } + t.out = try t.rig.output(Op.count(shape) * @as(usize, if (Op.out == .f16) 2 else 4)); + return true; + } + + /// The values as the device holds them (each rounded to its type), and their bytes. + fn encode(a: std.mem.Allocator, kind: r.Kind, v: []f64) ![]u8 { + const size: usize = if (kind == .f16) 2 else 4; + const out = try a.alloc(u8, v.len * size); + for (v, 0..) |*x, i| switch (kind) { + .f16 => { + const hv: f16 = @floatCast(@as(f32, @floatCast(x.*))); + std.mem.writeInt(u16, out[2 * i ..][0..2], @bitCast(hv), .little); + x.* = hv; + }, + .f32 => { + const fv: f32 = @floatCast(x.*); + std.mem.writeInt(u32, out[4 * i ..][0..4], @bitCast(fv), .little); + x.* = fv; + }, + }; + return out; + } + + pub fn run(t: *Self) !void { + if (Op.in_place) |i| { + const bytes = t.count() * @as(usize, if (Op.out == .f16) 2 else 4); + try (cuda.DeviceBuffer{ .d = t.rig.gpu.d, .ptr = t.out, .len = bytes }).copyFrom(0, t.dev[i], bytes, t.rig.s.handle); + } + try Op.launch(&t.k, t.rig.s, t.dev, t.out, t.shape, t.variant); + try t.rig.s.synchronize(); + } + + pub fn count(t: *Self) usize { + return Op.count(t.shape); + } + + fn constVals(t: *Self) [n][]const f64 { + var v: [n][]const f64 = undefined; + for (&v, t.vals) |*o, x| o.* = x; + return v; + } + + pub fn reference(t: *Self, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + try Op.ref(t.rig.gpu.gpa, t.shape, t.variant, t.constVals(), .{ .range = t.fill == .stress }, &t.ends, want, mag); + } + + pub fn bounds(t: *Self, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + + pub fn intervals(t: *Self, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } + + pub fn simulate(t: *Self, acc: c.Accum, order: g.Order, out: []f64) !void { + const a = t.rig.gpu.gpa; + const mag = try a.alloc(f64, out.len); + defer a.free(mag); + var scratch: r.Ends = .{ .a = a }; + defer scratch.deinit(); + try scratch.size(t.count()); + try Op.ref(a, t.shape, t.variant, t.constVals(), .{ .sim = .{ .acc = acc, .order = order } }, &scratch, out, mag); + } + }; +} + +/// The usual registration of a test: its rig on the stream, the kernels (or a control's copy), name and form. +pub fn make(comptime Op: type, rig: rig_.Rig, k: idol.Kernels, name: []const u8, variant: u32, control: ?g.Gate) Case(Op) { + return .{ .rig = rig, .k = k, .name = name, .variant = variant, .control = control, .ends = .{ .a = rig.gpu.gpa } }; +} diff --git a/zig/tests/cuda/contract_idol_host.zig b/zig/tests/cuda/contract_idol_host.zig new file mode 100644 index 0000000000..ee61edd549 --- /dev/null +++ b/zig/tests/cuda/contract_idol_host.zig @@ -0,0 +1,8 @@ +//! Idol's kernel-contract references on the host: their ranges hold the float64 values, no GPU work. +test { + _ = @import("contract_idol_ref.zig"); + _ = @import("contract_idol_norms.zig"); + _ = @import("contract_idol_steps.zig"); + _ = @import("contract_idol_rotary.zig"); + _ = @import("contract_idol_case.zig"); +} diff --git a/zig/tests/cuda/contract_idol_norms.zig b/zig/tests/cuda/contract_idol_norms.zig new file mode 100644 index 0000000000..010dd0e0a1 --- /dev/null +++ b/zig/tests/cuda/contract_idol_norms.zig @@ -0,0 +1,208 @@ +//! tf_idol_norm and tf_idol_norm32 under the contract: LayerNorm32 at idol-sdk's fp16 points, fp32 sums in gate 4. +const std = @import("std"); +const cuda = @import("cuda"); +const idol = @import("idol_cuda"); +const r = @import("contract_idol_ref.zig"); + +const c = r.c; +const Range = r.Range; +const Sum = r.Sum; + +const why = "an output follows the sums as (x - mean) / sqrt(var): an fp16 or bf16 sum's error is many output ulps"; + +/// A thread's share of the row in turn, the warp's 5 shuffles, then the 8 warps' partials in order. +fn depth(s: c.Shape) usize { + return (s.k + 255) / 256 + 5 + 8; +} + +const served = [_]c.Shape{ .{ .rows = r.joint, .n = r.hidden, .k = r.hidden }, .{ .rows = r.latent, .n = r.hidden, .k = r.hidden }, .{ .rows = r.cond, .n = r.hidden, .k = r.hidden } }; + +fn rowSpec(name: []const u8, out: c.Elem) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = depth, + .ops = 2 * (r.div_ulps + r.sqrt_ulps) + 3, + .mids = 2, + .out = out, + .served = &served, + .sums = .{ .f16 = 0, .bf16 = 0, .why = why }, + }; +} + +fn rowOk(s: c.Shape) bool { + return s.rows > 0 and s.k == r.hidden and s.n == r.hidden; +} + +/// One row's mean and 1 / sqrt(var + eps) as float64 and in range, each sum formed as `sum` says. +const Stats = struct { + mean: f64, + mr: Range, + inv: f64, + ir: Range, + + fn of(a: std.mem.Allocator, x: []const f64, eps: f64, sum: Sum) !Stats { + const k: f64 = @floatFromInt(x.len); + const dp = (x.len + 255) / 256 + 5 + 8; + const m = sum.total(x, x, x, dp); + const mean = m.v / k; + const mr = (Range{ .lo = m.r.lo / k, .hi = m.r.hi / k }).fp32(r.div_ulps); + const buf = try a.alloc(f64, 3 * x.len); + defer a.free(buf); + const sq = buf[0..x.len]; + const lo = buf[x.len .. 2 * x.len]; + const hi = buf[2 * x.len ..]; + for (x, sq, lo, hi) |v, *q, *l, *u| { + q.* = (v - mean) * (v - mean); + const qr = Range.of(v).sub(mr).fp32(0.5).square().fp32(0.5); + l.* = qr.lo; + u.* = qr.hi; + } + const s = sum.total(sq, lo, hi, dp); + const vr = (Range{ .lo = s.r.lo / k, .hi = s.r.hi / k }).fp32(r.div_ulps).add(Range.of(eps)).fp32(0.5); + const root = vr.up(sqrt).fp32(r.sqrt_ulps); + return .{ .mean = mean, .mr = mr, .inv = 1 / @sqrt(s.v / k + eps), .ir = Range.of(1).div(root).fp32(r.div_ulps) }; + } + + /// (x - mean) inv: its float64 value and range. + fn at(st: Stats, x: f64) struct { v: f64, r: Range } { + return .{ .v = (x - st.mean) * st.inv, .r = Range.of(x).sub(st.mr).fp32(0.5).mul(st.ir).fp32(0.5) }; + } +}; + +fn sqrt(x: f64) f64 { + return @sqrt(@max(x, 0)); +} + +/// Gate 4's rows: integers whose second half mirrors the first, so the mean is 0 and every fp32 sum is exact. +fn rowsFill(rng: std.Random, fill: c.Shape, mode: r.g.Fill, x: []f64) void { + if (mode == .stress) { + const k = fill.k; + for (0..fill.rows) |row| { + const v = x[row * k ..][0..k]; + r.integers(rng, v[0 .. k / 2], 8); + for (v[k / 2 ..], v[0 .. k / 2]) |*o, a| o.* = -a; + } + } else r.uniform(rng, x, -3, 3); + if (mode == .nonfinite) { + const k = fill.k; + r.plant(x, &.{ k / 3, k + k / 2, 2 * k + k - 1 }); + } +} + +/// tf_idol_norm: fp16(y w + b) (variant 1, affine) or fp16(fp16(fp16(y) s1) + sh) (variant 2, modulated). +pub const Norm = struct { + pub const kinds = [_]r.Kind{ .f16, .f32, .f32, .f16, .f16 }; + pub const out: r.Kind = .f16; + pub const in_place: ?usize = null; + pub const accepts = rowOk; + + pub fn spec(name: []const u8, _: u32) c.Spec { + return rowSpec(name, .f16); + } + + pub fn sizes(s: c.Shape) [5]usize { + return .{ s.rows * s.k, s.k, s.k, s.k, s.k }; + } + + pub fn count(s: c.Shape) usize { + return s.rows * s.k; + } + + pub fn fill(rng: std.Random, mode: r.g.Fill, s: c.Shape, _: u32, v: *[5][]f64) void { + rowsFill(rng, s, mode, v[0]); + r.uniform(rng, v[1], 0.25, 2); + r.uniform(rng, v[2], -1, 1); + // the modulation's 1 + scale, and its shift + r.uniform(rng, v[3], 0.5, 1.5); + r.uniform(rng, v[4], -1, 1); + if (mode == .nonfinite) for (v[1..]) |col| r.plant(col, &.{7}); + } + + pub fn launch(k: *const idol.Kernels, s: cuda.Stream, dev: [5]u64, o: u64, shape: c.Shape, variant: u32) !void { + try idol.norm(k, s, dev[0], dev[1], dev[2], dev[3], dev[4], o, shape.rows, @as(f32, 1e-6), if (variant == 2) .modulated else .affine); + } + + pub fn ref(a: std.mem.Allocator, shape: c.Shape, variant: u32, v: [5][]const f64, sum: Sum, e: *r.Ends, want: []f64, mag: []f64) !void { + const k = shape.k; + for (0..shape.rows) |row| { + const x = v[0][row * k ..][0..k]; + const st = try Stats.of(a, x, r.f(1e-6), sum); + for (x, 0..) |xv, i| { + const y = st.at(xv); + const at = row * k + i; + if (variant == 2) { + const s1 = v[3][i]; + const sh = v[4][i]; + e.put(want, mag, at, r.h(r.h(r.h(y.v) * s1) + sh), y.r.half().mul(Range.of(s1)).fp32(0.5).half().add(Range.of(sh)).fp32(0.5).half()); + } else { + const w = v[1][i]; + const b = v[2][i]; + e.put(want, mag, at, r.h(y.v * w + b), y.r.mul(Range.of(w)).fp32(0.5).add(Range.of(b)).fp32(0.5).half()); + } + } + } + } +}; + +/// tf_idol_norm32: y w + b (variant 1, RePo's norm) or y (1 + w) + b (variant 2, the output's), in fp32. +pub const Norm32 = struct { + pub const kinds = [_]r.Kind{ .f16, .f32, .f32 }; + pub const out: r.Kind = .f32; + pub const in_place: ?usize = null; + pub const accepts = rowOk; + + pub fn spec(name: []const u8, _: u32) c.Spec { + return rowSpec(name, .f32); + } + + pub fn sizes(s: c.Shape) [3]usize { + return .{ s.rows * s.k, s.k, s.k }; + } + + pub fn count(s: c.Shape) usize { + return s.rows * s.k; + } + + pub fn fill(rng: std.Random, mode: r.g.Fill, s: c.Shape, _: u32, v: *[3][]f64) void { + rowsFill(rng, s, mode, v[0]); + r.uniform(rng, v[1], -1, 2); + r.uniform(rng, v[2], -1, 1); + if (mode == .nonfinite) for (v[1..]) |col| r.plant(col, &.{9}); + } + + pub fn launch(k: *const idol.Kernels, s: cuda.Stream, dev: [3]u64, o: u64, shape: c.Shape, variant: u32) !void { + try idol.norm32(k, s, dev[0], dev[1], dev[2], o, shape.rows, @as(f32, 1e-5), if (variant == 2) .modulated else .affine); + } + + pub fn ref(a: std.mem.Allocator, shape: c.Shape, variant: u32, v: [3][]const f64, sum: Sum, e: *r.Ends, want: []f64, mag: []f64) !void { + const k = shape.k; + for (0..shape.rows) |row| { + const x = v[0][row * k ..][0..k]; + const st = try Stats.of(a, x, r.f(1e-5), sum); + for (x, 0..) |xv, i| { + const y = st.at(xv); + const w = v[1][i]; + const b = v[2][i]; + const at = row * k + i; + if (variant == 2) { + const one = Range.of(1).add(Range.of(w)).fp32(0.5); + e.put(want, mag, at, y.v * (1 + w) + b, y.r.mul(one).fp32(0.5).add(Range.of(b)).fp32(0.5)); + } else e.put(want, mag, at, y.v * w + b, y.r.mul(Range.of(w)).fp32(0.5).add(Range.of(b)).fp32(0.5)); + } + } + } +}; + +test "a LayerNorm row's statistics lie in their ranges, and gate 4's mirrored integers sum to a zero mean" { + const a = std.testing.allocator; + var x: [1024]f64 = undefined; + var prng = std.Random.DefaultPrng.init(5); + rowsFill(prng.random(), .{ .rows = 1, .n = 1024, .k = 1024 }, .stress, &x); + const st = try Stats.of(a, &x, r.f(1e-6), .{ .range = true }); + try std.testing.expectEqual(@as(f64, 0), st.mean); + try std.testing.expect(st.ir.lo <= st.inv and st.inv <= st.ir.hi); + rowsFill(prng.random(), .{ .rows = 1, .n = 1024, .k = 1024 }, .random, &x); + const rs = try Stats.of(a, &x, r.f(1e-6), .{ .range = false }); + try std.testing.expect(rs.mr.lo <= rs.mean and rs.mean <= rs.mr.hi and rs.ir.lo <= rs.inv and rs.inv <= rs.ir.hi); +} diff --git a/zig/tests/cuda/contract_idol_ref.zig b/zig/tests/cuda/contract_idol_ref.zig new file mode 100644 index 0000000000..fe004a1d2e --- /dev/null +++ b/zig/tests/cuda/contract_idol_ref.zig @@ -0,0 +1,213 @@ +//! Idol's CUDA kernel tests' shared parts: float64 values at the kernels' fp16 points, and every value they may take. +const std = @import("std"); +const core = @import("core"); + +pub const c = core.contract.spec; +pub const g = core.contract.gates; +pub const rows = core.contract.rows; + +pub const hidden = 1024; +pub const heads = 16; +pub const D = 64; +pub const ffn = 4096; +/// The served row counts: the joint rows, the noise refiner's latent rows, the condition refiner's rows. +pub const joint = 12294; +pub const latent = 8192; +pub const cond = 4101; + +// nvcc's error bounds under the glue flags: IEEE division and square root; expf, tanhf, cosf and sinf 2 ulps +pub const div_ulps = 0.5; +pub const sqrt_ulps = 0.5; +pub const fn_ulps = 2; + +/// An input's element type on the device. +pub const Kind = enum { f16, f32 }; + +/// A value as fp32, then fp16, round to nearest even (overflow to infinity). +pub fn h(v: f64) f64 { + return @as(f16, @floatCast(@as(f32, @floatCast(v)))); +} + +/// A value as fp32. +pub fn f(v: f64) f64 { + return @as(f32, @floatCast(v)); +} + +/// fmaxf: a NaN operand gives the other. +pub fn fmax(a: f64, b: f64) f64 { + if (std.math.isNan(a)) return b; + if (std.math.isNan(b)) return a; + return @max(a, b); +} + +/// 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); infinities stay. + pub fn fp32(a: Range, ulps: f64) Range { + const u = ulps * 0x1p-23; + return .{ .lo = if (std.math.isFinite(a.lo)) a.lo - u * @abs(a.lo) - 0x1p-149 else a.lo, .hi = if (std.math.isFinite(a.hi)) a.hi + u * @abs(a.hi) + 0x1p-149 else a.hi }; + } + pub fn half(a: Range) Range { + return .{ .lo = h(a.lo), .hi = h(a.hi) }; + } + pub fn add(a: Range, b: Range) Range { + return .{ .lo = a.lo + b.lo, .hi = a.hi + b.hi }; + } + pub fn sub(a: Range, b: Range) Range { + return .{ .lo = a.lo - b.hi, .hi = a.hi - b.lo }; + } + 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])) }; + } + /// a / b for b of one sign. + pub fn div(a: Range, b: Range) Range { + return a.mul(.{ .lo = 1 / b.hi, .hi = 1 / b.lo }); + } + /// An increasing function at both ends. + pub fn up(a: Range, comptime fun: fn (f64) f64) Range { + return .{ .lo = fun(a.lo), .hi = fun(a.hi) }; + } + /// fmaxf against a constant. + pub fn atLeast(a: Range, v: f64) Range { + return .{ .lo = @max(a.lo, v), .hi = @max(a.hi, v) }; + } + /// x^2 over the range. + pub fn square(a: Range) Range { + const lo = if (a.lo <= 0 and a.hi >= 0) 0 else @min(a.lo * a.lo, a.hi * a.hi); + return .{ .lo = lo, .hi = @max(a.lo * a.lo, a.hi * a.hi) }; + } + /// cos or sin (`sine`) over the range, its extrema included when one lies inside. + pub fn trig(a: Range, sine: bool) Range { + if (!(std.math.isFinite(a.lo) and std.math.isFinite(a.hi)) or a.hi - a.lo >= 2 * std.math.pi) return .{ .lo = -1, .hi = 1 }; + const fun: *const fn (f64) f64 = if (sine) sinf64 else cosf64; + var lo = @min(fun(a.lo), fun(a.hi)); + var hi = @max(fun(a.lo), fun(a.hi)); + // extrema at k pi (cos) or pi / 2 + k pi (sin) + const shift: f64 = if (sine) std.math.pi / 2.0 else 0; + var k = @ceil((a.lo - shift) / std.math.pi); + while (k * std.math.pi + shift <= a.hi) : (k += 1) { + const v = fun(k * std.math.pi + shift); + lo = @min(lo, v); + hi = @max(hi, v); + } + return .{ .lo = lo, .hi = hi }; + } +}; + +fn sinf64(x: f64) f64 { + return @sin(x); +} +fn cosf64(x: f64) f64 { + return @cos(x); +} + +/// How a reference forms its sums: in float64 with their fp32 range (exact for gate 4's inputs), or in an accumulator. +pub const Sum = union(enum) { + range: bool, + sim: struct { acc: c.Accum, order: g.Order }, + + /// The terms' sum and its range, `lo` and `hi` each term's own ends, `depth` fp32 roundings deep. + pub fn total(s: Sum, terms: []const f64, lo: []const f64, hi: []const f64, depth: usize) struct { v: f64, r: Range } { + switch (s) { + .sim => |x| { + const v = rows.sumIn(x.acc, x.order, terms); + return .{ .v = v, .r = Range.of(v) }; + }, + .range => |exact| { + var v: f64 = 0; + var a: f64 = 0; + var b: f64 = 0; + var m: f64 = 0; + for (terms, lo, hi) |t, l, u| { + v += t; + a += l; + b += u; + m += @max(@abs(l), @abs(u)); + } + if (exact) return .{ .v = v, .r = Range.of(v) }; + const e = c.gamma(@floatFromInt(depth), c.unit(.f32)) * m; + return .{ .v = v, .r = .{ .lo = a - e, .hi = b + e } }; + }, + } + } +}; + +/// 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) and std.math.isFinite(lo) and std.math.isFinite(hi)) @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); + } +}; + +/// Values uniform in [lo, hi]. +pub fn uniform(rng: std.Random, out: []f64, lo: f64, hi: f64) void { + for (out) |*v| v.* = lo + rng.float(f64) * (hi - lo); +} + +/// Integers in [-bound, bound]. +pub fn integers(rng: std.Random, out: []f64, bound: u32) void { + for (out) |*v| v.* = @as(f64, @floatFromInt(rng.uintLessThan(u32, 2 * bound + 1))) - @as(f64, @floatFromInt(bound)); +} + +pub const nan = std.math.nan(f64); +pub const inf = std.math.inf(f64); + +/// Plants NaN, +inf and -inf at `at` (indices taken modulo the length). +pub fn plant(out: []f64, at: []const usize) void { + const vals = [_]f64{ nan, inf, -inf }; + for (at, 0..) |i, j| out[i % out.len] = vals[j % 3]; +} + +test "fp16 points round to nearest even, ranges keep the exact value inside" { + try std.testing.expectEqual(@as(f64, 1.0), h(1.0 + 1.0 / 4096.0)); + try std.testing.expectEqual(inf, h(70000)); + const r = Range.of(0.3).fp32(0.5).mul(Range.of(1.7)).fp32(0.5).half(); + try std.testing.expect(r.lo <= h(0.3 * 1.7) and h(0.3 * 1.7) <= r.hi); + const t = (Range{ .lo = 3.1, .hi = 3.2 }).trig(false); + try std.testing.expectEqual(@as(f64, -1), t.lo); +} + +test "a sum's range holds the exact sum, and is the sum itself for gate 4's exact inputs" { + const t = [_]f64{ 1.5, -2.25, 3 }; + const s = (Sum{ .range = false }).total(&t, &t, &t, 3); + try std.testing.expect(s.r.lo < 2.25 and 2.25 < s.r.hi); + const e = (Sum{ .range = true }).total(&t, &t, &t, 3); + try std.testing.expectEqual(e.r.lo, e.r.hi); +} diff --git a/zig/tests/cuda/contract_idol_rotary.zig b/zig/tests/cuda/contract_idol_rotary.zig new file mode 100644 index 0000000000..6d23482ba9 --- /dev/null +++ b/zig/tests/cuda/contract_idol_rotary.zig @@ -0,0 +1,312 @@ +//! Idol's summing glue under the contract: RePo's angles, q and k turned and normed, the latent's input embed. +const std = @import("std"); +const cuda = @import("cuda"); +const idol = @import("idol_cuda"); +const r = @import("contract_idol_ref.zig"); + +const c = r.c; +const Range = r.Range; +const Sum = r.Sum; +const Fill = r.g.Fill; + +const why = "gate 4's integers keep fp32 exact; fp16 and bf16 sums miss by a margin it reports but cannot bound ahead"; + +fn summed(name: []const u8, out: c.Elem, served: []const c.Shape, depth: *const fn (c.Shape) usize) c.Spec { + return .{ .name = name, .accum = .f32, .depth = depth, .ops = 8, .mids = 2, .out = out, .served = served, .sums = .{ .f16 = 0, .bf16 = 0, .why = why } }; +} + +fn rowsOf(n: usize, k: usize) [3]c.Shape { + return .{ .{ .rows = r.joint, .n = n, .k = k }, .{ .rows = r.latent, .n = n, .k = k }, .{ .rows = r.cond, .n = n, .k = k } }; +} + +/// tf_idol_qk: q and k of [rows, 3072] turned by cos, sin [rows, 512], F.normalize times gamma times 8; v copied. +pub const Qk = struct { + pub const kinds = [_]r.Kind{ .f16, .f32, .f32, .f32, .f32 }; + pub const out: r.Kind = .f16; + pub const in_place: ?usize = null; + const served = rowsOf(3 * r.hidden, r.D); + + pub fn accepts(s: c.Shape) bool { + return s.rows > 0 and s.n == 3 * r.hidden and s.k == r.D; + } + + /// A lane's two squares and their sum, then the warp's 5 shuffles. + fn depth(_: c.Shape) usize { + return 7; + } + + pub fn spec(name: []const u8, _: u32) c.Spec { + return summed(name, .f16, &served, depth); + } + + pub fn sizes(s: c.Shape) [5]usize { + return .{ s.rows * 3 * r.hidden, r.hidden, r.hidden, s.rows * 512, s.rows * 512 }; + } + + pub fn count(s: c.Shape) usize { + return s.rows * 3 * r.hidden; + } + + pub fn fill(rng: std.Random, mode: Fill, _: c.Shape, _: u32, v: *[5][]f64) void { + if (mode == .stress) { + r.integers(rng, v[0], 64); + @memset(v[3], 1); + @memset(v[4], 0); + } else { + r.uniform(rng, v[0], -4, 4); + for (v[3], v[4]) |*co, *si| { + const t = rng.float(f64) * 2 * std.math.pi; + co.* = @cos(t); + si.* = @sin(t); + } + } + r.uniform(rng, v[1], 0.5, 1.5); + r.uniform(rng, v[2], 0.5, 1.5); + if (mode == .nonfinite) { + r.plant(v[0], &.{ 5, 3072 + 1024 + 77, 2 * 3072 + 2048 + 9 }); + r.plant(v[1], &.{300}); + } + } + + pub fn launch(k: *const idol.Kernels, s: cuda.Stream, dev: [5]u64, o: u64, shape: c.Shape, _: u32) !void { + try idol.qk(k, s, dev[0], dev[1], dev[2], dev[3], dev[4], o, shape.rows); + } + + pub fn ref(_: std.mem.Allocator, shape: c.Shape, _: u32, v: [5][]const f64, sum: Sum, e: *r.Ends, want: []f64, mag: []f64) !void { + var rot: [64]f64 = undefined; + var rr: [64]Range = undefined; + var sq: [32]f64 = undefined; + var lo: [32]f64 = undefined; + var hi: [32]f64 = undefined; + for (0..shape.rows) |row| for (0..r.heads) |hd| { + const base = row * 3 * r.hidden + hd * r.D; + for (0..r.D) |j| e.put(want, mag, base + 2048 + j, v[0][base + 2048 + j], Range.of(v[0][base + 2048 + j])); + for ([_]usize{ 0, 1024 }, [_][]const f64{ v[1], v[2] }) |off, gamma| { + for (0..32) |l| { + const x0 = v[0][base + off + 2 * l]; + const x1 = v[0][base + off + 2 * l + 1]; + const co = v[3][row * 512 + hd * 32 + l]; + const si = v[4][row * 512 + hd * 32 + l]; + rot[2 * l] = r.h(x0 * co - x1 * si); + rot[2 * l + 1] = r.h(x0 * si + x1 * co); + const a = Range.of(x0).mul(Range.of(co)).fp32(0.5); + const b = Range.of(x1).mul(Range.of(si)).fp32(0.5); + const p = Range.of(x0).mul(Range.of(si)).fp32(0.5); + const q = Range.of(x1).mul(Range.of(co)).fp32(0.5); + rr[2 * l] = a.sub(b).fp32(0.5).half(); + rr[2 * l + 1] = p.add(q).fp32(0.5).half(); + sq[l] = rot[2 * l] * rot[2 * l] + rot[2 * l + 1] * rot[2 * l + 1]; + const t = rr[2 * l].square().fp32(0.5).add(rr[2 * l + 1].square().fp32(0.5)).fp32(0.5); + lo[l] = t.lo; + hi[l] = t.hi; + } + const s = sum.total(&sq, &lo, &hi, 5); + const n = clampMin(@sqrt(s.v)); + const nr = s.r.up(sqrt).fp32(r.sqrt_ulps).up(clampMin); + for (0..r.D) |j| { + const g = gamma[hd * r.D + j]; + const w = r.h(rot[j] / n * g * 8); + const wr = rr[j].div(nr).fp32(r.div_ulps).mul(Range.of(g)).fp32(0.5).mul(Range.of(8)).fp32(0.5).half(); + e.put(want, mag, base + off + j, w, wr); + } + } + }; + } +}; + +fn sqrt(x: f64) f64 { + return @sqrt(@max(x, 0)); +} + +/// F.normalize's clamp_min(1e-12) in fp32, NaN kept. +fn clampMin(x: f64) f64 { + return if (x < r.f(1e-12)) r.f(1e-12) else x; +} + +/// tf_idol_repo: from gate | content [rows, 256], d = final (silu(gate) content), angles; out cos then sin [rows, 512]. +pub const Repo = struct { + pub const kinds = [_]r.Kind{ .f32, .f32, .f32 }; + pub const out: r.Kind = .f32; + pub const in_place: ?usize = null; + const served = rowsOf(r.hidden, 128); + const pi = r.f(std.math.pi); + + pub fn accepts(s: c.Shape) bool { + return s.rows > 0 and s.n == r.hidden and s.k == 128; + } + + /// The final map's 128 fused steps in turn. + fn depth(_: c.Shape) usize { + return 128; + } + + pub fn spec(name: []const u8, _: u32) c.Spec { + return summed(name, .f32, &served, depth); + } + + pub fn sizes(s: c.Shape) [3]usize { + return .{ s.rows * 256, 48 * 128, 32 }; + } + + pub fn count(s: c.Shape) usize { + return s.rows * 1024; + } + + pub fn fill(rng: std.Random, mode: Fill, s: c.Shape, _: u32, v: *[3][]f64) void { + if (mode == .stress) { + // gate 64..80: silu is the gate itself in fp32, so every term is an integer and every fp32 sum exact + for (0..s.rows) |row| { + for (v[0][row * 256 ..][0..128]) |*x| x.* = @floatFromInt(64 + rng.uintLessThan(u32, 17)); + r.integers(rng, v[0][row * 256 + 128 ..][0..128], 8); + } + r.integers(rng, v[1], 8); + @memset(v[2], 0x1p-19); + return; + } + r.uniform(rng, v[0], -4, 4); + r.uniform(rng, v[1], -0.2, 0.2); + r.uniform(rng, v[2], 0.5, 3); + if (mode == .nonfinite) { + r.plant(v[0], &.{ 7, 256 + 130, 2 * 256 + 50 }); + r.plant(v[1], &.{ 128 * 5 + 3, 128 * 20 }); + r.plant(v[2], &.{31}); + } + } + + pub fn launch(k: *const idol.Kernels, s: cuda.Stream, dev: [3]u64, o: u64, shape: c.Shape, _: u32) !void { + try idol.repo(k, s, dev[0], dev[1], dev[2], o, o + shape.rows * 512 * 4, shape.rows); + } + + pub fn ref(_: std.mem.Allocator, shape: c.Shape, _: u32, v: [3][]const f64, sum: Sum, e: *r.Ends, want: []f64, mag: []f64) !void { + var u: [128]f64 = undefined; + var ur: [128]Range = undefined; + var t: [128]f64 = undefined; + var lo: [128]f64 = undefined; + var hi: [128]f64 = undefined; + var d: [48]f64 = undefined; + var dr: [48]Range = undefined; + const half = shape.rows * 512; + for (0..shape.rows) |row| { + const g = v[0][row * 256 ..][0..256]; + for (0..128) |k| { + const a = g[k]; + const cv = g[128 + k]; + u[k] = a / (1 + @exp(-a)) * cv; + const den = Range.of(-a).up(exp).fp32(r.fn_ulps).add(Range.of(1)).fp32(0.5); + ur[k] = Range.of(a).div(den).fp32(r.div_ulps).mul(Range.of(cv)).fp32(0.5); + } + for (0..48) |o| { + for (0..128) |k| { + const w = v[1][o * 128 + k]; + t[k] = u[k] * w; + const p = ur[k].mul(Range.of(w)); + lo[k] = p.lo; + hi[k] = p.hi; + } + const s = sum.total(&t, &lo, &hi, 128); + d[o] = s.v; + dr[o] = s.r; + } + for (0..512) |i| { + const hd = i / 32; + const j = i % 32; + const axis: usize = if (j < 10) 0 else if (j < 20) 1 else 2; + const x = d[hd * 3 + axis]; + const xr = dr[hd * 3 + axis]; + const fr = v[2][j]; + const ft = std.math.tanh(fr); + const ang = (x * ft + x * (fr - ft)) * pi; + const ftr = Range.of(fr).up(tanh).fp32(r.fn_ulps); + const a1 = xr.mul(ftr).fp32(0.5); + const a2 = xr.mul(Range.of(fr).sub(ftr).fp32(0.5)).fp32(0.5); + const ar = a1.add(a2).fp32(0.5).mul(Range.of(pi)).fp32(0.5); + const at = row * 512 + i; + e.put(want, mag, at, @cos(ang), ar.trig(false).fp32(r.fn_ulps)); + e.put(want, mag, half + at, @sin(ang), ar.trig(true).fp32(r.fn_ulps)); + } + } + } +}; + +fn exp(x: f64) f64 { + return @exp(x); +} + +fn tanh(x: f64) f64 { + return std.math.tanh(x); +} + +/// tf_idol_embed: x = fp16(fp16(z w^T + b) + pe), 16 fp32 inputs a latent row, its fp16 position embedding. +pub const Embed = struct { + pub const kinds = [_]r.Kind{ .f32, .f32, .f32, .f16 }; + pub const out: r.Kind = .f16; + pub const in_place: ?usize = null; + const served = [_]c.Shape{.{ .rows = r.latent, .n = r.hidden, .k = 16 }}; + + pub fn accepts(s: c.Shape) bool { + return s.rows > 0 and s.n == r.hidden and s.k == 16; + } + + fn depth(_: c.Shape) usize { + return 16; + } + + pub fn spec(name: []const u8, _: u32) c.Spec { + return summed(name, .f16, &served, depth); + } + + pub fn sizes(s: c.Shape) [4]usize { + return .{ s.rows * 16, r.hidden * 16, r.hidden, s.rows * r.hidden }; + } + + pub fn count(s: c.Shape) usize { + return s.rows * r.hidden; + } + + pub fn fill(rng: std.Random, mode: Fill, _: c.Shape, _: u32, v: *[4][]f64) void { + if (mode == .stress) { + for (v) |col| r.integers(rng, col, 8); + return; + } + r.uniform(rng, v[0], -3, 3); + r.uniform(rng, v[1], -0.5, 0.5); + r.uniform(rng, v[2], -0.5, 0.5); + r.uniform(rng, v[3], -1, 1); + if (mode == .nonfinite) for (v, 0..) |col, i| r.plant(col, &.{ 3 + 40 * i, 900 + i }); + } + + pub fn launch(k: *const idol.Kernels, s: cuda.Stream, dev: [4]u64, o: u64, shape: c.Shape, _: u32) !void { + try idol.embed(k, s, dev[0], dev[1], dev[2], dev[3], o, shape.rows); + } + + pub fn ref(_: std.mem.Allocator, shape: c.Shape, _: u32, v: [4][]const f64, sum: Sum, e: *r.Ends, want: []f64, mag: []f64) !void { + var t: [16]f64 = undefined; + for (0..shape.rows) |row| for (0..r.hidden) |col| { + for (&t, 0..) |*x, k| x.* = v[0][row * 16 + k] * v[1][col * 16 + k]; + const s = sum.total(&t, &t, &t, 16); + const b = v[2][col]; + const at = row * r.hidden + col; + const pe = v[3][at]; + e.put(want, mag, at, r.h(r.h(s.v + b) + pe), s.r.add(Range.of(b)).fp32(0.5).half().add(Range.of(pe)).fp32(0.5).half()); + }; + } +}; + +test "q and k turned and normed by their warp's sum keep the float64 value inside its range" { + const a = std.testing.allocator; + const shape: c.Shape = .{ .rows = 1, .n = 3 * r.hidden, .k = r.D }; + var bufs: [5][]f64 = undefined; + for (&bufs, Qk.sizes(shape)) |*b, n| b.* = try a.alloc(f64, n); + defer for (bufs) |b| a.free(b); + var prng = std.Random.DefaultPrng.init(3); + Qk.fill(prng.random(), .random, shape, 0, &bufs); + var e: r.Ends = .{ .a = a }; + defer e.deinit(); + try e.size(Qk.count(shape)); + const want = try a.alloc(f64, Qk.count(shape)); + defer a.free(want); + const mag = try a.alloc(f64, Qk.count(shape)); + defer a.free(mag); + try Qk.ref(a, shape, 0, .{ bufs[0], bufs[1], bufs[2], bufs[3], bufs[4] }, .{ .range = false }, &e, want, mag); + for (want, e.lo, e.hi) |w, lo, hi| try std.testing.expect(lo <= w and w <= hi); +} diff --git a/zig/tests/cuda/contract_idol_steps.zig b/zig/tests/cuda/contract_idol_steps.zig new file mode 100644 index 0000000000..aa4d1f9fa8 --- /dev/null +++ b/zig/tests/cuda/contract_idol_steps.zig @@ -0,0 +1,217 @@ +//! Idol's epilogues under the contract: gated residual, linear biases (GELU), condition sum, fp32 bias and the embed. +const std = @import("std"); +const cuda = @import("cuda"); +const idol = @import("idol_cuda"); +const r = @import("contract_idol_ref.zig"); + +const c = r.c; +const Range = r.Range; +const Sum = r.Sum; +const Fill = r.g.Fill; + +fn none(_: c.Shape) usize { + return 0; +} + +const why = "no sums: each output is a few fp32 steps and its fp16 rounding, every one inside its interval"; + +/// A kernel with no sums: its outputs are held to their pointwise intervals alone. +fn plain(name: []const u8, out: c.Elem, served: []const c.Shape) c.Spec { + return .{ .name = name, .accum = .none, .depth = none, .ops = 4, .mids = 2, .out = out, .served = served, .sums = .{ .f16 = 0, .bf16 = 0, .why = why } }; +} + +fn wide(rows: usize, n: usize) c.Shape { + return .{ .rows = rows, .n = n, .k = n }; +} + +fn any(s: c.Shape) bool { + return s.rows > 0 and s.n > 0 and s.n == s.k; +} + +/// The 1024-wide residual stream's rows at every refiner. +const stream = [_]c.Shape{ wide(r.joint, r.hidden), wide(r.latent, r.hidden), wide(r.cond, r.hidden) }; + +/// tf_idol_gate_add: x = fp16(x + fp16(g y)), in place over x (input 0). +pub const GateAdd = struct { + pub const kinds = [_]r.Kind{ .f16, .f16, .f16 }; + pub const out: r.Kind = .f16; + pub const in_place: ?usize = 0; + pub const accepts = any; + + pub fn spec(name: []const u8, _: u32) c.Spec { + return plain(name, .f16, &stream); + } + + pub fn sizes(s: c.Shape) [3]usize { + return .{ s.rows * s.n, s.rows * s.n, s.n }; + } + + pub fn count(s: c.Shape) usize { + return s.rows * s.n; + } + + pub fn fill(rng: std.Random, mode: Fill, _: c.Shape, _: u32, v: *[3][]f64) void { + r.uniform(rng, v[0], -8, 8); + r.uniform(rng, v[1], -4, 4); + r.uniform(rng, v[2], -2, 2); + if (mode == .nonfinite) for (v, 0..) |col, i| r.plant(col, &.{ 3 + 5 * i, 1000 + i }); + } + + pub fn launch(k: *const idol.Kernels, s: cuda.Stream, dev: [3]u64, o: u64, shape: c.Shape, _: u32) !void { + try idol.gateAdd(k, s, o, dev[1], dev[2], shape.rows, shape.n); + } + + pub fn ref(_: std.mem.Allocator, shape: c.Shape, _: u32, v: [3][]const f64, _: Sum, e: *r.Ends, want: []f64, mag: []f64) !void { + for (0..shape.rows * shape.n) |i| { + const x = v[0][i]; + const y = v[1][i]; + const g = v[2][i % shape.n]; + const gy = Range.of(g).mul(Range.of(y)).fp32(0.5).half(); + e.put(want, mag, i, r.h(x + r.h(g * y)), gy.add(Range.of(x)).fp32(0.5).half()); + } + } +}; + +/// The block linears' widths: attention's 3072 and 1024, the MLP's 4096 (GELU) and 1024, at every refiner. +const linears = [_]c.Shape{ wide(r.joint, 3 * r.hidden), wide(r.joint, r.ffn), wide(r.joint, r.hidden), wide(r.latent, r.ffn), wide(r.cond, r.ffn) }; + +/// torch's tanh GELU of an fp16 value as the kernel forms it in fp32 (tanhf within fn_ulps): value and range. +pub fn gelu(x: f64) struct { v: f64, r: Range } { + const c1 = r.f(0.7978845608028654); + const c2 = r.f(0.044715); + const v = 0.5 * x * (1 + std.math.tanh(c1 * (x + c2 * x * x * x))); + const xr = Range.of(x); + const t = Range.of(c2).mul(xr).fp32(0.5).mul(xr).fp32(0.5).mul(xr).fp32(0.5); + const inner = Range.of(c1).mul(xr.add(t).fp32(0.5)).fp32(0.5); + const th = inner.up(tanh).fp32(r.fn_ulps); + return .{ .v = v, .r = Range.of(0.5).mul(xr).mul(Range.of(1).add(th).fp32(0.5)).fp32(0.5) }; +} + +fn tanh(x: f64) f64 { + return std.math.tanh(x); +} + +/// tf_idol_bias: y = fp16(acc + b) from fp32 sums and an fp16 bias; variant 1 then takes fp16(GELU(y)). +pub const Bias = struct { + pub const kinds = [_]r.Kind{ .f32, .f16 }; + pub const out: r.Kind = .f16; + pub const in_place: ?usize = null; + pub const accepts = any; + + pub fn spec(name: []const u8, _: u32) c.Spec { + return plain(name, .f16, &linears); + } + + pub fn sizes(s: c.Shape) [2]usize { + return .{ s.rows * s.n, s.n }; + } + + pub fn count(s: c.Shape) usize { + return s.rows * s.n; + } + + pub fn fill(rng: std.Random, mode: Fill, _: c.Shape, _: u32, v: *[2][]f64) void { + r.uniform(rng, v[0], -6, 6); + r.uniform(rng, v[1], -1, 1); + if (mode == .nonfinite) for (v, 0..) |col, i| r.plant(col, &.{ 11 + 7 * i, 2048 + i }); + } + + pub fn launch(k: *const idol.Kernels, s: cuda.Stream, dev: [2]u64, o: u64, shape: c.Shape, variant: u32) !void { + try idol.bias(k, s, dev[0], dev[1], o, shape.rows, shape.n, variant == 1); + } + + pub fn ref(_: std.mem.Allocator, shape: c.Shape, variant: u32, v: [2][]const f64, _: Sum, e: *r.Ends, want: []f64, mag: []f64) !void { + for (0..shape.rows * shape.n) |i| { + const y = r.h(v[0][i] + v[1][i % shape.n]); + const yr = Range.of(v[0][i]).add(Range.of(v[1][i % shape.n])).fp32(0.5).half(); + if (variant != 1) { + e.put(want, mag, i, y, yr); + continue; + } + // GELU's ranges at both fp16 ends, and its minimum (-0.16997 at -0.7518) when the range holds it + const lo = gelu(yr.lo); + const hi = gelu(yr.hi); + const g = gelu(y); + var gr: Range = .{ .lo = @min(@min(lo.r.lo, hi.r.lo), g.r.lo), .hi = @max(@max(lo.r.hi, hi.r.hi), g.r.hi) }; + if (yr.lo <= -0.751 and yr.hi >= -0.753) gr.lo = @min(gr.lo, -0.1701); + e.put(want, mag, i, r.h(g.v), gr.half()); + } + } +}; + +/// tf_idol_cond: y = fp16((a + ba) + (c + bc)), the two condition projections' fp32 sums and biases. +pub const Cond = struct { + pub const kinds = [_]r.Kind{ .f32, .f32, .f32, .f32 }; + pub const out: r.Kind = .f16; + pub const in_place: ?usize = null; + pub const accepts = any; + const served = [_]c.Shape{wide(r.cond, r.hidden)}; + + pub fn spec(name: []const u8, _: u32) c.Spec { + return plain(name, .f16, &served); + } + + pub fn sizes(s: c.Shape) [4]usize { + return .{ s.rows * s.n, s.n, s.rows * s.n, s.n }; + } + + pub fn count(s: c.Shape) usize { + return s.rows * s.n; + } + + pub fn fill(rng: std.Random, mode: Fill, _: c.Shape, _: u32, v: *[4][]f64) void { + for (v) |col| r.uniform(rng, col, -4, 4); + if (mode == .nonfinite) for (v, 0..) |col, i| r.plant(col, &.{ 17 + 3 * i, 3000 + i }); + } + + pub fn launch(k: *const idol.Kernels, s: cuda.Stream, dev: [4]u64, o: u64, shape: c.Shape, _: u32) !void { + try idol.cond(k, s, dev[0], dev[1], dev[2], dev[3], o, shape.rows, shape.n); + } + + pub fn ref(_: std.mem.Allocator, shape: c.Shape, _: u32, v: [4][]const f64, _: Sum, e: *r.Ends, want: []f64, mag: []f64) !void { + for (0..shape.rows * shape.n) |i| { + const j = i % shape.n; + const a = Range.of(v[0][i]).add(Range.of(v[1][j])).fp32(0.5); + const b = Range.of(v[2][i]).add(Range.of(v[3][j])).fp32(0.5); + e.put(want, mag, i, r.h((v[0][i] + v[1][j]) + (v[2][i] + v[3][j])), a.add(b).fp32(0.5).half()); + } + } +}; + +/// tf_idol_bias32: y += b in fp32, in place over y (input 0): RePo's 256 gate | content sums, the output's 16. +pub const Bias32 = struct { + pub const kinds = [_]r.Kind{ .f32, .f32 }; + pub const out: r.Kind = .f32; + pub const in_place: ?usize = 0; + pub const accepts = any; + const served = [_]c.Shape{ wide(r.joint, 256), wide(r.latent, 256), wide(r.cond, 256), wide(r.latent, 16) }; + + pub fn spec(name: []const u8, _: u32) c.Spec { + return plain(name, .f32, &served); + } + + pub fn sizes(s: c.Shape) [2]usize { + return .{ s.rows * s.n, s.n }; + } + + pub fn count(s: c.Shape) usize { + return s.rows * s.n; + } + + pub fn fill(rng: std.Random, mode: Fill, _: c.Shape, _: u32, v: *[2][]f64) void { + r.uniform(rng, v[0], -10, 10); + r.uniform(rng, v[1], -1, 1); + if (mode == .nonfinite) for (v, 0..) |col, i| r.plant(col, &.{ 5 + 9 * i, 777 + i }); + } + + pub fn launch(k: *const idol.Kernels, s: cuda.Stream, dev: [2]u64, o: u64, shape: c.Shape, _: u32) !void { + try idol.bias32(k, s, o, dev[1], shape.rows, shape.n); + } + + pub fn ref(_: std.mem.Allocator, shape: c.Shape, _: u32, v: [2][]const f64, _: Sum, e: *r.Ends, want: []f64, mag: []f64) !void { + for (0..shape.rows * shape.n) |i| { + const b = v[1][i % shape.n]; + e.put(want, mag, i, r.f(v[0][i] + b), Range.of(v[0][i]).add(Range.of(b)).fp32(0.5)); + } + } +}; diff --git a/zig/tests/cuda/contract_limits.zig b/zig/tests/cuda/contract_limits.zig index ca9e662d17..419cc615ef 100644 --- a/zig/tests/cuda/contract_limits.zig +++ b/zig/tests/cuda/contract_limits.zig @@ -26,7 +26,7 @@ pub fn image(name: []const u8) error{UnknownContractImage}![]const u8 { if (std.mem.eql(u8, name, "h3")) return &h3_image.bytes; inline for (.{ "qwen4_projection", "qwen4_hc", "qwen4_gdn", "qwen4_ple" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); - inline for (.{ "tree_round", "act_rows", "gdn_prep", "qk_rows", "norms", "moe", "experts", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "qmm_prefill", "lane_gemv", "prompt16", "qmmf", "ck_experts", "nvfp4_experts", "nvfp4_experts_tile", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "glm_glue", "glm_ops", "glm_index", "nemotron_mamba", "sample", "qwen35", "qwen35_bf16_rows", "qwen35_keep", "qwen35_kv8", "w4a8_prefill", "qwen35_draft", "ck_lane", "gdn_train", "train_q4", "train_attn", "flashnext_train", "train", "attn_train", "qwen27_train", "tree_attention", "block_attention", "window_attention", "window_pages", "moe_glue", "dense_rows", "qmm", "flashnext_fact", "flashnext_hcf", "flashnext_fexp", "experts_rows", "nemotron_ops", "minimax_m3", "minimax_m3_attn", "moss_image" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); + inline for (.{ "tree_round", "act_rows", "gdn_prep", "qk_rows", "norms", "moe", "experts", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "qmm_prefill", "lane_gemv", "prompt16", "qmmf", "ck_experts", "nvfp4_experts", "nvfp4_experts_tile", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "glm_glue", "glm_ops", "glm_index", "nemotron_mamba", "sample", "qwen35", "qwen35_bf16_rows", "qwen35_keep", "qwen35_kv8", "w4a8_prefill", "qwen35_draft", "ck_lane", "gdn_train", "train_q4", "train_attn", "flashnext_train", "train", "attn_train", "qwen27_train", "tree_attention", "block_attention", "window_attention", "window_pages", "moe_glue", "dense_rows", "qmm", "flashnext_fact", "flashnext_hcf", "flashnext_fexp", "experts_rows", "nemotron_ops", "minimax_m3", "minimax_m3_attn", "moss_image", "idol" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); if (std.mem.eql(u8, name, "contract_controls")) return if (options.controls) &controls.bytes else &.{}; std.debug.print("FAIL kernel-contract: unknown image name: {s}\n", .{name}); return error.UnknownContractImage; diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 3521a83762..e4c0d80f93 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -237,6 +237,7 @@ const usage = \\ pocket-contract Nemotron's pocket kernels (learned neurons) through gates 1-4, controls red \\ minimax-contract MiniMax-M3.1-Flash's indexer and block-sparse attention through gates 1-4, controls red \\ moss-contract MOSS-Image's own kernels and its full attention through gates 1-4, controls red + \\ idol-contract Idol's denoiser glue and head-64 attention through gates 1-4, controls red \\ glm-exl3 GLM's EXL3 routed experts, decode and prompt, against float64 (synthetic) \\ glm-vision MODEL F32 BF16 IMG GLM's image tower and preparation against vision_ref.py's torch runs \\ glm-mm-fixed Prompt K policies: FP64, every row alone/batched, offsets and guarded tails @@ -458,6 +459,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "mtp-rows")) return @import("contract_qwen_mtp_rows.zig").alone(gpu); if (std.mem.eql(u8, cmd, "qwen-pages")) return @import("contract_qwen_pages.zig").run(gpu); if (std.mem.eql(u8, cmd, "minimax-contract")) return @import("contract_minimax.zig").run(gpu); + if (std.mem.eql(u8, cmd, "idol-contract")) return @import("contract_idol.zig").run(gpu); if (std.mem.eql(u8, cmd, "kernel-contract-limits")) return contract_tests.fatbins(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "pocket-contract")) return @import("contract_pocket.zig").run(gpu); if (std.mem.eql(u8, cmd, "moss-contract")) return @import("contract_moss_image.zig").run(gpu); diff --git a/zig/tests/idol.zig b/zig/tests/idol.zig index 4359d38e45..330d091952 100644 --- a/zig/tests/idol.zig +++ b/zig/tests/idol.zig @@ -1,242 +1,41 @@ -//! tf-idol: Idol's denoiser on a case from tools/zig/idol_case.py, compared with idol-sdk's, gated and timed. +//! tf-idol on Metal: idol_run.zig's case driver on the system GPU. 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. -; +const Metal = struct { + pub const id = tf.idol; + pub const st = tf.safetensors; + pub const seconds = mtl.clock.seconds; -/// 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; + pub fn stream(m: *id.Model, rows: usize) ![]const f16 { + return m.stream(rows); } - 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" }; + pub const Device = struct { + pool: mtl.objc.Pool, + device: mtl.Device, -/// 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); - }; + pub fn init(g: *Device) !void { + g.pool = mtl.objc.Pool.push(); + errdefer g.pool.pop(); + g.device = try mtl.Device.init(); } - 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)); -} + pub fn deinit(g: *Device) void { + g.pool.pop(); + } -/// 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; -} + pub fn name(g: *Device) []const u8 { + return std.mem.span(g.device.name()); + } -/// 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 load(g: *Device, gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !*id.Model { + return id.Model.load(gpa, io, g.device, dir); + } + }; +}; 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; + return @import("idol_run.zig").run(Metal, init); } diff --git a/zig/tests/idol_cuda.zig b/zig/tests/idol_cuda.zig new file mode 100644 index 0000000000..45d2328607 --- /dev/null +++ b/zig/tests/idol_cuda.zig @@ -0,0 +1,43 @@ +//! tf-idol on CUDA: idol_run.zig's case driver on the first GPU (sm_120 or sm_121; others are refused by name). +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); + +const Cuda = struct { + pub const id = @import("idol"); + pub const st = core.safetensors; + pub const seconds = id.seconds; + + pub fn stream(m: *id.Model, rows: usize) ![]const f16 { + return m.residual(rows); + } + + pub const Device = struct { + driver: cuda.Driver, + ctx: cuda.Context, + buf: [256]u8, + + pub fn init(g: *Device) !void { + g.driver = try cuda.Driver.open(); + errdefer g.driver.close(); + g.ctx = try cuda.Context.init(&g.driver, 0); + } + + pub fn deinit(g: *Device) void { + g.ctx.deinit(); + g.driver.close(); + } + + pub fn name(g: *Device) []const u8 { + return g.ctx.name(&g.buf) catch "a CUDA GPU"; + } + + pub fn load(g: *Device, gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !*id.Model { + return id.Model.load(gpa, io, &g.ctx, dir); + } + }; +}; + +pub fn main(init: std.process.Init) !u8 { + return @import("idol_run.zig").run(Cuda, init); +} diff --git a/zig/tests/idol_run.zig b/zig/tests/idol_run.zig new file mode 100644 index 0000000000..095ca78d21 --- /dev/null +++ b/zig/tests/idol_run.zig @@ -0,0 +1,248 @@ +//! tf-idol's case driver for a backend `B`: its family `id`, safetensors `st`, `seconds`, `Device` and `stream`. +const std = @import("std"); + +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: anytype, 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: anytype) !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 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(comptime B: type, a: std.mem.Allocator, c: Case, model: *B.id.Model, conds: *[2]B.id.Condition, out: *std.ArrayList(u8)) !bool { + const cfg = B.id.config; + const ref = B.id.reference; + const N = cfg.latent_rows * cfg.latent_channels; + 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(comptime B: type, a: std.mem.Allocator, c: Case, model: *B.id.Model, conds: *[2]B.id.Condition) !bool { + const cfg = B.id.config; + const N = cfg.latent_rows * cfg.latent_channels; + 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: anytype, 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(comptime B: type, a: std.mem.Allocator, io: std.Io, dir: []const u8, step: usize, c: Case, model: *B.id.Model) !bool { + const cfg = B.id.config; + const ref = B.id.reference; + var f = try B.st.File.open(a, io, try std.fmt.allocPrint(a, "{s}/blocks{d}.safetensors", .{ dir, step })); + defer f.close(io); + var g = try B.st.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, try B.stream(model, 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 run(comptime B: type, init: std.process.Init) !u8 { + const id = B.id; + const cfg = id.config; + const ref = id.reference; + const N = cfg.latent_rows * cfg.latent_channels; + 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 B.st.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: ?B.st.File = B.st.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); + + var gpu: B.Device = undefined; + try gpu.init(); + defer gpu.deinit(); + const device = gpu.name(); + const loading = B.seconds(); + const model = try gpu.load(std.heap.page_allocator, io, args[1]); + defer model.deinit(); + const load_s = B.seconds() - loading; + std.debug.print("{s}: {d} steps; loaded in {d:.1} s\n", .{ device, S, load_s }); + + var out: std.ArrayList(u8) = .empty; + const conditioning = B.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 = B.seconds() - conditioning; + if (args.len == 5) return if (try blockGates(B, 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, S, load_s, condition_s }); + const pass = try teacher(B, a, c, model, &conds, &out); + const repeat = try repeats(B, 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 8cdab33631f3bdc460a0383c6571897f879f6452 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 03:11:26 +0800 Subject: [PATCH 08/14] idol: CUDA recipe and changelog line --- CHANGELOG.md | 2 ++ docs/recipes/idol-cuda.md | 41 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 43 insertions(+) create mode 100644 docs/recipes/idol-cuda.md diff --git a/CHANGELOG.md b/CHANGELOG.md index 2abaa9af8a..6d9be636a3 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's denoiser on CUDA, beside the Metal family and driven by the same `tf-idol` case driver and gates: cuBLASLt fp16 projections with fp32 sums, true fp32 (no TF32) for RePo, the condition and the output, and its own head-64 attention. It admits sm_121 (GB10) and sm_120 and refuses other GPUs. It compiles for both and its host tests pass, but it has not yet run on a GPU (after #1766, for #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-cuda.md b/docs/recipes/idol-cuda.md new file mode 100644 index 0000000000..6a9810f25c --- /dev/null +++ b/docs/recipes/idol-cuda.md @@ -0,0 +1,41 @@ +# Idol's denoiser on CUDA + +Idol's denoiser (see `idol.md`) on the CUDA runtime, beside the Metal family and driven by the same case driver +(`zig/tests/idol_run.zig`): `tf-idol` on Metal and on CUDA reads the same case and runs the same gates. The image +encoders, the decoders and the PLY writer are not ported, and nothing registers with the lane engine or the server. + +**Not run on a GPU yet.** It compiles for sm_120 and sm_121 and its host tests pass; no CUDA kernel of this family has +run on a device. Until a GB10 or an RTX 50 card has run the kernel contract and a case, treat it as unqualified. + +`Model.load` admits sm_121 (GB10) and sm_120 (RTX 50, RTX PRO Blackwell) and refuses any other GPU by name. + +## Precision + +The same contract as on Metal, in CUDA's own terms: + +- fp16 block projections on cuBLASLt (`CUBLAS_COMPUTE_32F`, fp16 inputs, fp32 output); `idol.cu` adds the bias to + those fp32 sums and rounds once to fp16, GELU after it. +- RePo's maps, the condition embedding and the output projection in fp32 on cuBLASLt with `CUBLAS_COMPUTE_32F` and + fp32 inputs: true fp32 (no TF32), so no hi/lo split is needed here. +- LayerNorms in fp32, rounded to fp16; the modulation in fp16 ops. +- Attention: `idol_attention.cuh`, head 64, fp16 q, k and v, fp32 scores, softmax and sums, scale 1/8, every query row + to every key row. It is Idol's own kernel: no shared attention template gains an instance. +- `idol.cu` builds with the glue flags (`--fmad=false`, IEEE division and square root). `F.normalize`'s clamp keeps a + NaN norm, as on Metal. + +## Kernel contract + +`tf-cuda-test idol-contract` runs gates 1-4 over every `idol.cu` kernel and the attention, with three controls that +each drop one fp16 rounding point; the full contract runs it among its family sets. Gate-5 records +(`zig/src/core/contract/idol_cuda_records.zig`) are nvcc 13.0's for sm_80, 86, 89, 120 and 121. No SUMMARY line yet: +it needs a GPU. + +## Run a case + +```sh +zig build idol-cuda +zig-out/bin/tf-idol WEIGHTS case/ +zig-out/bin/tf-idol WEIGHTS case/ blocks 0 +``` + +The case comes from `tools/zig/idol_case.py`, as in `idol.md`. From 95da5ab57cde8e6d3de37db6e20eb0d0bd2830f9 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 03:24:13 +0800 Subject: [PATCH 09/14] changelog: #1767 for Idol's CUDA denoiser --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6d9be636a3..65bea78b74 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's denoiser on CUDA, beside the Metal family and driven by the same `tf-idol` case driver and gates: cuBLASLt fp16 projections with fp32 sums, true fp32 (no TF32) for RePo, the condition and the output, and its own head-64 attention. It admits sm_121 (GB10) and sm_120 and refuses other GPUs. It compiles for both and its host tests pass, but it has not yet run on a GPU (after #1766, for #1757). Thanks to @jayleaton. +- Idol's denoiser on CUDA, beside the Metal family and driven by the same `tf-idol` case driver and gates: cuBLASLt fp16 projections with fp32 sums, true fp32 (no TF32) for RePo, the condition and the output, and its own head-64 attention. It admits sm_121 (GB10) and sm_120 and refuses other GPUs. It compiles for both and its host tests pass, but it has not yet run on a GPU (#1767, after #1766, for #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. 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 10/14] 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 11/14] 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 12/14] 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. From f4a07e29e43ad9cd86c18ab65d47d6ff27037fd4 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 18:57:08 +0800 Subject: [PATCH 13/14] cuda contract controls: close tf_ctl_pocket_in_fmax before the MOSS-Image controls --- zig/kernels/cuda/contract_controls.cu | 1 + 1 file changed, 1 insertion(+) diff --git a/zig/kernels/cuda/contract_controls.cu b/zig/kernels/cuda/contract_controls.cu index bbd8db2602..b866fd560c 100644 --- a/zig/kernels/cuda/contract_controls.cu +++ b/zig/kernels/cuda/contract_controls.cu @@ -364,6 +364,7 @@ extern "C" __global__ void __launch_bounds__(128) tf_ctl_pocket_in_fmax(const __ const float z = fmaxf(s, 0.0f); act[static_cast(r) * 64 + j] = z * z; } +} // Gate 3: moss_image.cu's steps, each with one of the publisher's bf16 points dropped; their range checks must go red. namespace tf_ctl_moss { typedef __nv_bfloat16 bf16; From 7f68bb85fa1f610d8c0da55d31d7bf2d9bdbea86 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 18:57:23 +0800 Subject: [PATCH 14/14] idol: the FLUX.2 VAE image encoder on CUDA for sm_121/sm_120 (untested on GB10) Port Idol's image encoder (#1774) to CUDA under the same float64 gate, fp32 throughout on the SDK's bf16 weights. Compiles and passes host tests; never run on a GPU; GB10 qualification pending. - idol_vae.cu / idol_vae.cuh: every convolution as an fp32 implicit GEMM of FMAs (same-size 3x3, stride-2 3x3 padded right and below, 1x1; no Winograd, no tensor cores, no TF32), conv_in, two-pass GroupNorm, GroupNorm + SiLU, bias epilogues, the attention's softmax and the condition rows (NaN-keeping clamp); glue flags. The mid block's q/k/v/out and P.V run on cuBLASLt in true fp32 (cuda_gemm.zig, CUBLAS_COMPUTE_32F), inside free picture tensors. - cuda_vae*.zig: encoder, launchers, weights; admits sm_120/sm_121 alone and refuses others by name; posterior noise stays an input. - vae_config.zig: geometry, config check, bf16 rounding and the gate shared with Metal; tf-idol-vae's driver shared too (idol_vae_run.zig), the Metal tool's output unchanged; tf-idol-vae on CUDA (idol-vae-cuda step). - kernel contract: idol-vae-contract, 16 forms through gates 1-4 and eleven controls built from idol_vae.cuh with one planted flaw each, joined to the family sets; nvcc 13.0 records for sm_80/86/89/120/121; image in contract_limits; host tests for geometry, scratch, admission and references. - contract_controls.cu: restore the closing brace of tf_ctl_pocket_in_fmax that the MOSS-Image CUDA commit dropped (as on main), so the controls build. --- CHANGELOG.md | 2 + docs/kernel-contract.md | 3 + docs/recipes/idol-cuda.md | 62 ++- zig/build/cuda_kernels.zig | 1 + zig/build/idol_cuda.zig | 29 +- zig/kernels/cuda/contract_controls.cu | 4 + zig/kernels/cuda/idol_vae.cu | 55 ++ zig/kernels/cuda/idol_vae.cuh | 318 +++++++++++ zig/kernels/cuda/idol_vae_controls.cuh | 56 ++ .../core/contract/idol_vae_cuda_records.zig | 22 + zig/src/core/contract/registry.zig | 3 +- zig/src/cuda/kernels.zig | 1 + zig/src/families/idol/cuda.zig | 8 +- zig/src/families/idol/cuda_vae.zig | 226 ++++++++ zig/src/families/idol/cuda_vae_kernels.zig | 182 +++++++ zig/src/families/idol/cuda_vae_weights.zig | 208 ++++++++ zig/src/families/idol/idol.zig | 2 + zig/src/families/idol/vae.zig | 63 +-- zig/src/families/idol/vae_config.zig | 155 ++++++ zig/src/families/idol/vae_kernels.zig | 7 +- zig/src/families/idol/vae_weights.zig | 93 +--- zig/tests/cuda/contract_family_sets.zig | 8 +- zig/tests/cuda/contract_idol_case.zig | 9 +- zig/tests/cuda/contract_idol_host.zig | 4 +- zig/tests/cuda/contract_idol_vae.zig | 151 ++++++ zig/tests/cuda/contract_idol_vae_convs.zig | 275 ++++++++++ zig/tests/cuda/contract_idol_vae_rows.zig | 501 ++++++++++++++++++ zig/tests/cuda/contract_limits.zig | 2 +- zig/tests/cuda/main.zig | 2 + zig/tests/idol_vae.zig | 146 +---- zig/tests/idol_vae_cuda.zig | 41 ++ zig/tests/idol_vae_run.zig | 130 +++++ 32 files changed, 2494 insertions(+), 275 deletions(-) create mode 100644 zig/kernels/cuda/idol_vae.cu create mode 100644 zig/kernels/cuda/idol_vae.cuh create mode 100644 zig/kernels/cuda/idol_vae_controls.cuh create mode 100644 zig/src/core/contract/idol_vae_cuda_records.zig create mode 100644 zig/src/families/idol/cuda_vae.zig create mode 100644 zig/src/families/idol/cuda_vae_kernels.zig create mode 100644 zig/src/families/idol/cuda_vae_weights.zig create mode 100644 zig/src/families/idol/vae_config.zig create mode 100644 zig/tests/cuda/contract_idol_vae.zig create mode 100644 zig/tests/cuda/contract_idol_vae_convs.zig create mode 100644 zig/tests/cuda/contract_idol_vae_rows.zig create mode 100644 zig/tests/idol_vae_cuda.zig create mode 100644 zig/tests/idol_vae_run.zig diff --git a/CHANGELOG.md b/CHANGELOG.md index 8cb23ae9e3..d29a0d3df1 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's FLUX.2 VAE image encoder on CUDA, driven by the same `tf-idol-vae` case driver and float64 gate as on Metal: fp32 throughout on the SDK's bf16 weights, every convolution an fp32 implicit GEMM on the FMA units (no Winograd, no tensor cores, no TF32), two-pass GroupNorm, and the mid block's attention on cuBLASLt in true fp32 with its own softmax. The posterior's noise stays an input. It admits sm_121 (GB10) and sm_120 and refuses other GPUs. It compiles for both and its host tests pass, but it has never run on a GPU; GB10 qualification is pending (after #1774 and #1767, for #1757). Thanks to @jayleaton. + - Idol's denoiser on CUDA, beside the Metal family and driven by the same `tf-idol` case driver and gates: cuBLASLt fp16 projections with fp32 sums, true fp32 (no TF32) for RePo, the condition and the output, and its own head-64 attention. It admits sm_121 (GB10) and sm_120 and refuses other GPUs. It compiles for both and its host tests pass, but it has not yet run on a GPU (#1767, after #1766, for #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. diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 1d920643c7..1056fbf326 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -227,6 +227,9 @@ Older SMs build main's loop, and their records are main's. CUDA also checks MOSS-Image's glue (`moss_image.cu`: the norm, plain and modulated; q/k norm and rotary into packed rows; situ; the gated add), each output within float64's range at its torch bf16 points (five controls, each dropping one point), and `vision_kernel` at head dim 128 against `contract_attn.zig`'s float64 reference (`moss-contract`). +Idol's CUDA image encoder (`idol_vae.cu`, fp32) registers its implicit-GEMM convolutions, GroupNorm's passes, the +attention's biases and softmax, and the condition rows against float64; eleven controls each build one body of +`idol_vae.cuh` with a planted flaw (`idol-vae-contract`). ## CUDA records per nvcc release diff --git a/docs/recipes/idol-cuda.md b/docs/recipes/idol-cuda.md index 6a9810f25c..318d605025 100644 --- a/docs/recipes/idol-cuda.md +++ b/docs/recipes/idol-cuda.md @@ -1,8 +1,9 @@ # Idol's denoiser on CUDA Idol's denoiser (see `idol.md`) on the CUDA runtime, beside the Metal family and driven by the same case driver -(`zig/tests/idol_run.zig`): `tf-idol` on Metal and on CUDA reads the same case and runs the same gates. The image -encoders, the decoders and the PLY writer are not ported, and nothing registers with the lane engine or the server. +(`zig/tests/idol_run.zig`): `tf-idol` on Metal and on CUDA reads the same case and runs the same gates. The FLUX.2 VAE +image encoder is ported too ([below](#vae-encoder)); DINOv3, the decoders and the PLY writer are not, and nothing +registers with the lane engine or the server. **Not run on a GPU yet.** It compiles for sm_120 and sm_121 and its host tests pass; no CUDA kernel of this family has run on a device. Until a GB10 or an RTX 50 card has run the kernel contract and a case, treat it as unqualified. @@ -39,3 +40,60 @@ zig-out/bin/tf-idol WEIGHTS case/ blocks 0 ``` The case comes from `tools/zig/idol_case.py`, as in `idol.md`. + +## VAE encoder + +**Not run on a GPU yet.** It compiles and passes host tests; it has never run on a GPU, and GB10 qualification is +pending. + +`cuda_vae.zig` runs the encoder of `idol.md`'s [VAE encoder](idol.md#vae-encoder) section on CUDA: the picture +`IdolPipeline.preprocess` returns to the denoiser's `feature2` [4101, 128], with the posterior's noise as an input. +`tf-idol-vae` shares its case driver with Metal (`zig/tests/idol_vae_run.zig`), so it reads the same case and runs the +same float64 gate. `Encoder.load` admits sm_121 and sm_120 alone, as the denoiser does. The loader takes klein's bf16 +file or FLUX.2-dev's fp32 one rounded to bf16, through `vae_config.zig`, which Metal shares. + +### Precision + +fp32 throughout on the SDK's bf16 weights, as on Metal, against the same float64 ruler: + +- Every convolution is `tf_iv_conv`, an implicit GEMM of fp32 FMAs (`__fmaf_rn`) on the FMA units: same-size 3x3s, + the stride-2 downsamples padded right and below, and the 1x1s. Each output is one FMA chain over its taps and input + channels in order, then the bias and the residual. There is no Winograd: with no tensor-unit workaround to pay for, + the direct form keeps each output an fp32 dot product, with no `G g G^T` rounding or transform sums. +- `conv_in` (`2 x - 1`, three channels) is its own kernel, as on Metal. +- GroupNorm in two passes (the mean, then the squared deviations), then the affine and SiLU, `expf` and IEEE + division. +- The mid block's attention runs on cuBLASLt in true fp32 (`CUBLAS_COMPUTE_32F`, fp32 operands, never TF32): q and k + as torch's Linear, v as `W_v xn^T` so that `P V` is cuBLASLt's `X W^T` too, and the scores 4,096 query rows at a + time. `tf_iv_softmax` takes the row's maximum, `expf` and the fp32 sum, and divides. The buffers fit inside the + free picture tensors: no memory beyond the four 512 MB ones. +- `tf_iv_tokens` keeps a NaN logvar through the clamp, as `torch.clamp` does. + +`idol_vae.cu` builds with the glue flags (`--fmad=false`, IEEE division and square root). + +### Kernel contract + +`tf-cuda-test idol-vae-contract` runs gates 1-4 over all eight `idol_vae.cu` kernels in 16 forms, each against +float64: `tf_iv_conv_in`, `tf_iv_conv` (3x3, 3x3 plus residual, stride-2 3x3, 1x1), `tf_iv_group_part` (sums, +squared deviations), `tf_iv_group_finish` (means, rstd), `tf_iv_norm` (plain, SiLU), `tf_iv_bias` (by column, by +column plus residual, by row), `tf_iv_softmax` and `tf_iv_tokens`. Gate 4's softmax rows are equal scores, so its +denominator is exact in fp32 and an fp16 sum stalls at 2,048. + +Eleven controls each build one body of `idol_vae.cuh` with a planted flaw, and each must go red at its gate: half +inputs (`conv_in`, the conv's tile), symmetric padding in the stride-2 3x3, a half SiLU, half sums (the conv, both +GroupNorm passes, the softmax), a by-row bias read by column, swapped patch axes, and a NaN logvar lost to +`fminf`/`fmaxf`. The engine's kernels are the flaw-free instances. The full contract runs this set among its family +sets. Gate 5's records (`idol_vae_cuda_records.zig`) are nvcc 13.0's for sm_80, 86, 89, 120 and 121. No SUMMARY line +yet: it needs a GPU. + +### Run a case + +```sh +zig build idol-vae-cuda -Doptimize=ReleaseFast -Dnvcc=/usr/local/cuda/bin/nvcc -Dsm=121 # -Dsm=120 for an RTX 50 +zig-out/bin/tf-cuda-test idol-vae-contract +zig-out/bin/tf-idol-vae KLEIN/vae case/ +``` + +The case comes from `tools/zig/idol_case.py vae`, as in `idol.md`. A pass is the line `against float64: ... pass` +(worst row and `1 - cosine` within twice the SDK's), five more encodes with identical bytes, and exit 0. + diff --git a/zig/build/cuda_kernels.zig b/zig/build/cuda_kernels.zig index 78c291fc69..24cd106a82 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -105,6 +105,7 @@ pub const kernels = [_]Kernel{ .{ .name = "minimax_m3", .flags = glue }, // ours: MiniMax-M3.1-Flash's glue (norms, RoPE prep, routing, DSpark, tower RoPE) .{ .name = "minimax_m3_attn", .flags = glue }, // ours: MiniMax-M3.1-Flash's indexer and block-sparse attention .{ .name = "idol", .flags = glue }, // ours: Idol's denoiser glue and its fp16 head-64 attention + .{ .name = "idol_vae", .flags = glue }, // ours: Idol's fp32 image encoder (implicit-GEMM convs, GroupNorm, rows) .{ .name = "glm_quant", .flags = torch_ops }, // ours: GLM's draft-only weight quantizers, torch at load in the Python family .{ .name = "exl3_linear", .flags = &.{ "-O3", "--expt-relaxed-constexpr" } }, // exl3/linear.cu's device code, tensorfold_exl3_linear_v5 .{ .name = "exl3_experts", .flags = &.{ "-O3", "--expt-relaxed-constexpr" } }, // exl3/experts.cu's device code, tensorfold_exl3_experts diff --git a/zig/build/idol_cuda.zig b/zig/build/idol_cuda.zig index c9d32101ce..8bc10de233 100644 --- a/zig/build/idol_cuda.zig +++ b/zig/build/idol_cuda.zig @@ -1,4 +1,4 @@ -//! Idol's CUDA denoiser in the build: its kernels for tf-cuda-test's contract, tf-idol on CUDA, its host tests. +//! Idol on CUDA in the build: its kernels for tf-cuda-test's contract, tf-idol and tf-idol-vae, its host tests. const std = @import("std"); /// The family's module over the runtime and the core. @@ -11,22 +11,35 @@ fn launchers(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.buil return b.createModule(.{ .root_source_file = b.path("zig/src/families/idol/cuda_kernels.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "cuda", .module = cuda }} }); } -/// The launchers in tf-cuda-test's `runner` (idol-contract), and tf-idol: the case driver on the first GPU. +/// The image encoder's launchers alone (idol_vae.cu's kernels), as tf-cuda-test's contract imports them. +fn vaeLaunchers(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module) *std.Build.Module { + return b.createModule(.{ .root_source_file = b.path("zig/src/families/idol/cuda_vae_kernels.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "cuda", .module = cuda }} }); +} + +/// The launchers in tf-cuda-test's `runner` (idol-contract, idol-vae-contract); tf-idol and tf-idol-vae on CUDA. pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, strip: ?bool, cuda: *std.Build.Module, core: *std.Build.Module, runner: *std.Build.Module) void { runner.addImport("idol_cuda", launchers(b, target, optimize, cuda)); - const tool = b.createModule(.{ .root_source_file = b.path("zig/tests/idol_cuda.zig"), .target = target, .optimize = optimize, .link_libc = true, .strip = strip, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core }, .{ .name = "idol", .module = module(b, target, optimize, cuda, core) } } }); - const exe = b.addExecutable(.{ .name = "tf-idol", .root_module = tool }); - b.installArtifact(exe); - b.step("idol-cuda", "Build tf-idol: Idol's denoiser on a reference case on CUDA").dependOn(&b.addInstallArtifact(exe, .{}).step); + runner.addImport("idol_vae_cuda", vaeLaunchers(b, target, optimize, cuda)); + const family = module(b, target, optimize, cuda, core); + const tools = [_]struct { []const u8, []const u8, []const u8, []const u8 }{ + .{ "tf-idol", "zig/tests/idol_cuda.zig", "idol-cuda", "Build tf-idol: Idol's denoiser on a reference case on CUDA" }, + .{ "tf-idol-vae", "zig/tests/idol_vae_cuda.zig", "idol-vae-cuda", "Build tf-idol-vae: Idol's image encoder on a reference case on CUDA" }, + }; + for (tools) |t| { + const tool = b.createModule(.{ .root_source_file = b.path(t[1]), .target = target, .optimize = optimize, .link_libc = true, .strip = strip, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core }, .{ .name = "idol", .module = family } } }); + const exe = b.addExecutable(.{ .name = t[0], .root_module = tool }); + b.installArtifact(exe); + b.step(t[2], t[3]).dependOn(&b.addInstallArtifact(exe, .{}).step); + } } /// The family's host tests (config, schedule, scratch, admission, launch geometry, contract ranges) in `all`. pub fn hostTests(b: *std.Build, cuda: *std.Build.Module, core: *std.Build.Module, all: *std.Build.Step) void { const host = b.graph.host; const run = b.addRunArtifact(b.addTest(.{ .root_module = module(b, host, .debug, cuda, core) })); - const refs = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/contract_idol_host.zig"), .target = host, .optimize = .debug, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core }, .{ .name = "idol_cuda", .module = launchers(b, host, .debug, cuda) } } }); + const refs = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/contract_idol_host.zig"), .target = host, .optimize = .debug, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core }, .{ .name = "idol_cuda", .module = launchers(b, host, .debug, cuda) }, .{ .name = "idol_vae_cuda", .module = vaeLaunchers(b, host, .debug, cuda) } } }); const ref_run = b.addRunArtifact(b.addTest(.{ .root_module = refs })); - const step = b.step("test-idol-cuda-host", "Idol's CUDA denoiser host tests and contract references (no GPU work)"); + const step = b.step("test-idol-cuda-host", "Idol's CUDA denoiser and image encoder host tests and contract references (no GPU work)"); for ([_]*std.Build.Step{ &run.step, &ref_run.step }) |s| { step.dependOn(s); all.dependOn(s); diff --git a/zig/kernels/cuda/contract_controls.cu b/zig/kernels/cuda/contract_controls.cu index bbd8db2602..dc2b9185ce 100644 --- a/zig/kernels/cuda/contract_controls.cu +++ b/zig/kernels/cuda/contract_controls.cu @@ -364,6 +364,7 @@ extern "C" __global__ void __launch_bounds__(128) tf_ctl_pocket_in_fmax(const __ const float z = fmaxf(s, 0.0f); act[static_cast(r) * 64 + j] = z * z; } +} // Gate 3: moss_image.cu's steps, each with one of the publisher's bf16 points dropped; their range checks must go red. namespace tf_ctl_moss { typedef __nv_bfloat16 bf16; @@ -593,3 +594,6 @@ extern "C" __global__ void __launch_bounds__(256) tf_ctl_idol_bias_late(const fl if (act) v = 0.5f * v * (1.0f + tanhf(0.7978845608028654f * (v + 0.044715f * v * v * v))); Y[i] = __float2half_rn(act ? v : h(v)); } + +// Idol's image encoder (idol_vae.cu), each control one planted flaw of idol_vae.cuh's bodies. +#include "idol_vae_controls.cuh" diff --git a/zig/kernels/cuda/idol_vae.cu b/zig/kernels/cuda/idol_vae.cu new file mode 100644 index 0000000000..f3e91b2f0f --- /dev/null +++ b/zig/kernels/cuda/idol_vae.cu @@ -0,0 +1,55 @@ +// Idol's image encoder on CUDA: idol_vae.cuh's bodies as the engine runs them (flaw none), fp32 throughout. +#include "idol_vae.cuh" + +using namespace tf_iv; + +extern "C" __global__ void __launch_bounds__(256) tf_iv_conv_in(const float* __restrict__ x, const float* __restrict__ w, + const float* __restrict__ b, float* __restrict__ y, + int rows, int cols) { + conv_in(x, w, b, y, rows, cols); +} + +extern "C" __global__ void __launch_bounds__(256) tf_iv_conv(const float* __restrict__ X, const float* __restrict__ W, + const float* __restrict__ B, const float* __restrict__ R, + float* __restrict__ Y, int rows, int cols, int in_rows, + int in_cols, int cin, int cout, int ldw, int taps, + int stride, int pad, int add) { + conv(X, W, B, R, Y, rows, cols, in_rows, in_cols, cin, cout, ldw, taps, stride, pad, add); +} + +extern "C" __global__ void __launch_bounds__(512) tf_iv_group_part(const float* __restrict__ x, + const float* __restrict__ stats, + float* __restrict__ part, int pixels, int channels, + int chunk, int chunks, int centered) { + group_part(x, stats, part, pixels, channels, chunk, chunks, centered); +} + +extern "C" __global__ void __launch_bounds__(1024) tf_iv_group_finish(const float* __restrict__ part, + float* __restrict__ stats, int pixels, + int channels, int chunks, int centered, + float eps) { + group_finish(part, stats, pixels, channels, chunks, centered, eps); +} + +extern "C" __global__ void __launch_bounds__(256) tf_iv_norm(const float* __restrict__ x, const float* __restrict__ stats, + const float* __restrict__ g, const float* __restrict__ beta, + float* __restrict__ y, int pixels, int channels, int silu) { + norm(x, stats, g, beta, y, pixels, channels, silu); +} + +extern "C" __global__ void __launch_bounds__(256) tf_iv_bias(const float* acc, const float* __restrict__ b, + const float* r, float* y, int rows, int cols, int by_row, + int add) { + bias(acc, b, r, y, rows, cols, by_row, add); +} + +extern "C" __global__ void __launch_bounds__(256) tf_iv_softmax(float* s, int n, float scale) { + softmax(s, n, scale); +} + +extern "C" __global__ void __launch_bounds__(256) tf_iv_tokens(const float* __restrict__ moments, + const float* __restrict__ noise, + const float* __restrict__ m, const float* __restrict__ s, + float* __restrict__ out, int rows, int cols, int pad) { + tokens(moments, noise, m, s, out, rows, cols, pad); +} diff --git a/zig/kernels/cuda/idol_vae.cuh b/zig/kernels/cuda/idol_vae.cuh new file mode 100644 index 0000000000..8a7c27bc71 --- /dev/null +++ b/zig/kernels/cuda/idol_vae.cuh @@ -0,0 +1,318 @@ +// Idol's image encoder (FLUX.2's VAE encoder) in fp32 on CUDA, each picture tensor [rows, columns, channels]. +#pragma once +#include +#include +#include + +namespace tf_iv { + +// Each body takes a flaw F: idol_vae.cu builds F = none alone; contract_controls.cu builds the controls' flaws. +enum Flaw { + none = 0, + conv_in_half = 1, // conv_in stages 2 x - 1 in half + conv_half_in = 2, // conv stages its input tile in half + conv_sym_pad = 3, // conv's stride-2 3x3 pads above and left as padding=1 would + conv_half_sum = 4, // conv sums in half + norm_half = 5, // norm's SiLU in half + part_half = 6, // group_part sums in half + finish_half = 7, // group_finish sums in half + softmax_half = 8, // softmax's denominator summed in half + bias_swap = 9, // bias by row reads the column's bias + tokens_axes = 10, // tokens swaps a patch's axes + tokens_clamp = 11, // tokens clamps a NaN logvar (fminf, fmaxf) +}; + +__device__ __forceinline__ float h(float x) { return __half2float(__float2half_rn(x)); } + +__device__ __forceinline__ float warp_sum(float v) { +#pragma unroll + for (int o = 16; o > 0; o >>= 1) v += __shfl_xor_sync(0xffffffffu, v, o); + return v; +} + +__device__ __forceinline__ float warp_max(float v) { +#pragma unroll + for (int o = 16; o > 0; o >>= 1) v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, o)); + return v; +} + +// y [pixels, 128] = conv3x3(2 x - 1) + b of the picture x [pixels, 3] in [0, 1], zero padded; w [27, 128] by tap. +template +__device__ __forceinline__ void conv_in(const float* __restrict__ x, const float* __restrict__ w, + const float* __restrict__ b, float* __restrict__ y, int rows, int cols) { + const int o = threadIdx.x; + const long long p = (long long)blockIdx.x * 2 + threadIdx.y; + if (p >= (long long)rows * cols) return; + const int py = (int)(p / cols), px = (int)(p % cols); + float s = 0.0f; +#pragma unroll + for (int ky = 0; ky < 3; ++ky) { +#pragma unroll + for (int kx = 0; kx < 3; ++kx) { + const int iy = py + ky - 1, ix = px + kx - 1; + if (iy < 0 || ix < 0 || iy >= rows || ix >= cols) continue; + const float* v = x + ((size_t)iy * cols + ix) * 3; +#pragma unroll + for (int c = 0; c < 3; ++c) { + float a = v[c] * 2.0f - 1.0f; + if (F == conv_in_half) a = h(a); + s = __fmaf_rn(a, w[((ky * 3 + kx) * 3 + c) * 128 + o], s); + } + } + } + y[(size_t)p * 128 + o] = s + b[o]; +} + +// The implicit GEMM's tile: 128 pixels by 128 output channels a block of 256, 8 of K a step, 8 x 8 outputs a thread. +constexpr int BM = 128; +constexpr int BN = 128; +constexpr int BK = 8; +constexpr int LDA = BM + 4; + +// y = conv(x) + b (+ r when add), w [taps * cin, ldw]: one fp32 FMA chain over the taps' input channels, then b, r. +template +__device__ __forceinline__ void conv(const float* __restrict__ X, const float* __restrict__ W, + const float* __restrict__ B, const float* __restrict__ R, float* __restrict__ Y, + int rows, int cols, int in_rows, int in_cols, int cin, int cout, int ldw, int taps, + int stride, int pad, int add) { + __shared__ __align__(16) float As[2][BK][LDA]; + __shared__ __align__(16) float Bs[2][BK][BN]; + const int t = threadIdx.x; + const long long pixels = (long long)rows * cols; + const long long m0 = (long long)blockIdx.x * BM; + const int n0 = blockIdx.y * BN; + // pad: rows and columns above and left (a same-size 3x3: 1; the stride-2 downsample, padded right and below: 0) + if (F == conv_sym_pad && stride == 2) pad = 1; + // this thread's loads: four input channels of one pixel, four weight columns of one K row + const int am = t >> 1, aq = (t & 1) * 4; + const long long apix = m0 + am; + const bool alive = apix < pixels; + const int oy = alive ? (int)(apix / cols) : 0; + const int ox = alive ? (int)(apix % cols) : 0; + const int bk = t >> 5, bc = (t & 31) * 4; + const int K = taps * cin; + const int tx = t & 15, ty = t >> 4; + float acc[8][8]; +#pragma unroll + for (int i = 0; i < 8; ++i) { +#pragma unroll + for (int j = 0; j < 8; ++j) acc[i][j] = 0.0f; + } + float4 a, bq; + // K row k0's tap and channel, the input pixel it reads (zero outside the picture), and the weights + auto load = [&](int k0) { + const int tap = k0 / cin; + const int c0 = k0 - tap * cin; + const int ky = taps == 9 ? tap / 3 : 0; + const int kx = taps == 9 ? tap % 3 : 0; + const int iy = oy * stride + ky - pad, ix = ox * stride + kx - pad; + a = make_float4(0.0f, 0.0f, 0.0f, 0.0f); + if (alive && iy >= 0 && ix >= 0 && iy < in_rows && ix < in_cols) + a = *reinterpret_cast(X + ((size_t)iy * in_cols + ix) * cin + c0 + aq); + bq = *reinterpret_cast(W + (size_t)(k0 + bk) * ldw + n0 + bc); + }; + auto store = [&](int s) { + if (F == conv_half_in) a = make_float4(h(a.x), h(a.y), h(a.z), h(a.w)); + As[s][aq + 0][am] = a.x; + As[s][aq + 1][am] = a.y; + As[s][aq + 2][am] = a.z; + As[s][aq + 3][am] = a.w; + *reinterpret_cast(&Bs[s][bk][bc]) = bq; + }; + load(0); + store(0); + __syncthreads(); + int s = 0; + for (int k0 = 0; k0 < K; k0 += BK) { + const bool more = k0 + BK < K; + if (more) load(k0 + BK); +#pragma unroll + for (int kk = 0; kk < BK; ++kk) { + const float4 a0 = *reinterpret_cast(&As[s][kk][ty * 4]); + const float4 a1 = *reinterpret_cast(&As[s][kk][64 + ty * 4]); + const float4 b0 = *reinterpret_cast(&Bs[s][kk][tx * 4]); + const float4 b1 = *reinterpret_cast(&Bs[s][kk][64 + tx * 4]); + const float av[8] = {a0.x, a0.y, a0.z, a0.w, a1.x, a1.y, a1.z, a1.w}; + const float bv[8] = {b0.x, b0.y, b0.z, b0.w, b1.x, b1.y, b1.z, b1.w}; +#pragma unroll + for (int i = 0; i < 8; ++i) { +#pragma unroll + for (int j = 0; j < 8; ++j) { + acc[i][j] = __fmaf_rn(av[i], bv[j], acc[i][j]); + if (F == conv_half_sum) acc[i][j] = h(acc[i][j]); + } + } + } + // the other buffer was last read before the previous step's barrier + if (more) store(s ^ 1); + __syncthreads(); + s ^= 1; + } +#pragma unroll + for (int i = 0; i < 8; ++i) { + const long long row = m0 + (i < 4 ? ty * 4 + i : 64 + ty * 4 + i - 4); + if (row >= pixels) continue; +#pragma unroll + for (int jh = 0; jh < 2; ++jh) { + const int col = n0 + jh * 64 + tx * 4; + if (col >= cout) continue; + const size_t at = (size_t)row * cout + col; + float4 v = make_float4(acc[i][jh * 4] + B[col], acc[i][jh * 4 + 1] + B[col + 1], acc[i][jh * 4 + 2] + B[col + 2], + acc[i][jh * 4 + 3] + B[col + 3]); + if (add) { + const float4 r = *reinterpret_cast(R + at); + v = make_float4(v.x + r.x, v.y + r.y, v.z + r.z, v.w + r.w); + } + *reinterpret_cast(Y + at) = v; + } + } +} + +// One chunk of pixels' sums of x a group, or of (x - mean)^2 when centered, into part [32, chunks]: a thread a channel. +template +__device__ __forceinline__ void group_part(const float* __restrict__ x, const float* __restrict__ stats, + float* __restrict__ part, int pixels, int channels, int chunk, int chunks, + int centered) { + __shared__ float sums[512]; + const int t = threadIdx.x; + const int size = channels / 32; + const long long first = (long long)blockIdx.x * chunk; + const long long last = min(first + chunk, (long long)pixels); + const float mean = centered ? stats[t / size] : 0.0f; + float s = 0.0f; + for (long long px = first; px < last; ++px) { + const float v = x[(size_t)px * channels + t]; + const float d = v - mean; + const float term = centered ? d * d : v; + s = F == part_half ? h(h(s) + h(term)) : s + term; + } + sums[t] = s; + __syncthreads(); + if (t < 32) { + float all = 0.0f; + for (int c = 0; c < size; ++c) all = F == part_half ? h(h(all) + h(sums[t * size + c])) : all + sums[t * size + c]; + part[t * chunks + blockIdx.x] = all; + } +} + +// stats from part's chunk sums, a warp a group: mean[32] when not centered, else rstd[32] = 1 / sqrt(var + eps). +template +__device__ __forceinline__ void group_finish(const float* __restrict__ part, float* __restrict__ stats, int pixels, + int channels, int chunks, int centered, float eps) { + const int lane = threadIdx.x & 31, g = threadIdx.x >> 5; + float s = 0.0f; + for (int i = lane; i < chunks; i += 32) s = F == finish_half ? h(h(s) + h(part[g * chunks + i])) : s + part[g * chunks + i]; + s = warp_sum(s); + if (lane != 0) return; + const float n = (float)pixels * (float)(channels / 32); + if (centered) + stats[32 + g] = 1.0f / sqrtf(s / n + eps); + else + stats[g] = s / n; +} + +// y = GroupNorm(x) by stats (32 means, then 32 rstd) with g and beta, then SiLU when silu; x, y [pixels, channels]. +template +__device__ __forceinline__ void norm(const float* __restrict__ x, const float* __restrict__ stats, + const float* __restrict__ g, const float* __restrict__ beta, float* __restrict__ y, + int pixels, int channels, int silu) { + const size_t i = (size_t)blockIdx.x * 256 + threadIdx.x; + if (i >= (size_t)pixels * channels) return; + const int c = (int)(i % channels); + const int grp = c / (channels / 32); + float v = (x[i] - stats[grp]) * stats[32 + grp]; + v = v * g[c] + beta[c]; + if (silu && F == norm_half) { + v = h(h(v) / h(1.0f + h(expf(-h(v))))); + } else if (silu) { + v = v / (1.0f + expf(-v)); + } + y[i] = v; +} + +// y = acc + b (+ r when add), b by column, or by row when by_row; y may be acc. +template +__device__ __forceinline__ void bias(const float* acc, const float* __restrict__ b, const float* r, float* y, int rows, + int cols, int by_row, int add) { + const size_t i = (size_t)blockIdx.x * 256 + threadIdx.x; + if (i >= (size_t)rows * cols) return; + const int col = (int)(i % cols); + const int row = (int)(i / cols); + float v = acc[i] + b[by_row && F != bias_swap ? row : col]; + if (add) v = v + r[i]; + y[i] = v; +} + +// The softmax of each row of s [rows, n] times scale, in place, n at most 256 * SOFTMAX_PER: a block of 256 a row. +constexpr int SOFTMAX_PER = 64; + +template +__device__ __forceinline__ void softmax(float* s, int n, float scale) { + __shared__ float part[8]; + float* row = s + (size_t)blockIdx.x * n; + const int t = threadIdx.x; + float v[SOFTMAX_PER]; + float m = -INFINITY; +#pragma unroll + for (int j = 0; j < SOFTMAX_PER; ++j) { + const int i = t + j * 256; + v[j] = i < n ? row[i] * scale : -INFINITY; + m = fmaxf(m, v[j]); + } + m = warp_max(m); + if ((t & 31) == 0) part[t >> 5] = m; + __syncthreads(); +#pragma unroll + for (int w = 0; w < 8; ++w) m = fmaxf(m, part[w]); + __syncthreads(); + float sum = 0.0f; +#pragma unroll + for (int j = 0; j < SOFTMAX_PER; ++j) { + const int i = t + j * 256; + if (i < n) { + v[j] = expf(v[j] - m); + sum = F == softmax_half ? h(h(sum) + h(v[j])) : sum + v[j]; + } + } + sum = warp_sum(sum); + if ((t & 31) == 0) part[t >> 5] = sum; + __syncthreads(); + float total = 0.0f; +#pragma unroll + for (int w = 0; w < 8; ++w) total += part[w]; +#pragma unroll + for (int j = 0; j < SOFTMAX_PER; ++j) { + const int i = t + j * 256; + if (i < n) row[i] = v[j] / total; + } +} + +// Rows [pad + patches, 128]: pad zeros, then each 2 x 2 patch of z = mean + std * noise as (z - m) / s. +template +__device__ __forceinline__ void tokens(const float* __restrict__ moments, const float* __restrict__ noise, + const float* __restrict__ m, const float* __restrict__ s, float* __restrict__ out, + int rows, int cols, int pad) { + const int k = threadIdx.x; + const int r = blockIdx.x * 2 + threadIdx.y; + const int w2 = cols / 2; + if (r >= pad + (rows / 2) * w2) return; + const size_t o = (size_t)r * 128 + k; + if (r < 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 int patch = r - pad; + const int c = k / 4; + const int dy = F == tokens_axes ? k % 2 : (k % 4) / 2; + const int dx = F == tokens_axes ? (k % 4) / 2 : k % 2; + const size_t pixel = (size_t)((patch / w2) * 2 + dy) * cols + (patch % w2) * 2 + dx; + // torch.clamp keeps a NaN, where fminf and fmaxf would drop it + const float raw = moments[pixel * 64 + 32 + c]; + const float clamped = fminf(fmaxf(raw, -30.0f), 20.0f); + const float logvar = F == tokens_clamp ? clamped : (isnan(raw) ? raw : clamped); + const float z = moments[pixel * 64 + c] + expf(0.5f * logvar) * noise[pixel * 32 + c]; + out[o] = (z - m[k]) / s[k]; +} + +} // namespace tf_iv diff --git a/zig/kernels/cuda/idol_vae_controls.cuh b/zig/kernels/cuda/idol_vae_controls.cuh new file mode 100644 index 0000000000..9d5ef6693e --- /dev/null +++ b/zig/kernels/cuda/idol_vae_controls.cuh @@ -0,0 +1,56 @@ +// Idol's image encoder controls (contract_controls.cu): idol_vae.cuh's bodies each built with one planted flaw. +#pragma once +#include "idol_vae.cuh" + +extern "C" __global__ void __launch_bounds__(256) tf_ctl_iv_conv_in_half(const float* x, const float* w, const float* b, + float* y, int rows, int cols) { + tf_iv::conv_in(x, w, b, y, rows, cols); +} + +#define TF_CTL_IV_CONV(name, flaw) \ + extern "C" __global__ void __launch_bounds__(256) \ + name(const float* X, const float* W, const float* B, const float* R, float* Y, int rows, int cols, \ + int in_rows, int in_cols, int cin, int cout, int ldw, int taps, int stride, int pad, int add) { \ + tf_iv::conv(X, W, B, R, Y, rows, cols, in_rows, in_cols, cin, cout, ldw, taps, stride, pad, add); \ + } +TF_CTL_IV_CONV(tf_ctl_iv_conv_half_in, tf_iv::conv_half_in) +TF_CTL_IV_CONV(tf_ctl_iv_conv_sym_pad, tf_iv::conv_sym_pad) +TF_CTL_IV_CONV(tf_ctl_iv_conv_half_sum, tf_iv::conv_half_sum) +#undef TF_CTL_IV_CONV + +extern "C" __global__ void __launch_bounds__(256) tf_ctl_iv_norm_half(const float* x, const float* stats, const float* g, + const float* beta, float* y, int pixels, + int channels, int silu) { + tf_iv::norm(x, stats, g, beta, y, pixels, channels, silu); +} + +extern "C" __global__ void __launch_bounds__(512) tf_ctl_iv_group_part_half(const float* x, const float* stats, + float* part, int pixels, int channels, + int chunk, int chunks, int centered) { + tf_iv::group_part(x, stats, part, pixels, channels, chunk, chunks, centered); +} + +extern "C" __global__ void __launch_bounds__(1024) tf_ctl_iv_group_finish_half(const float* part, float* stats, + int pixels, int channels, int chunks, + int centered, float eps) { + tf_iv::group_finish(part, stats, pixels, channels, chunks, centered, eps); +} + +extern "C" __global__ void __launch_bounds__(256) tf_ctl_iv_softmax_half(float* s, int n, float scale) { + tf_iv::softmax(s, n, scale); +} + +extern "C" __global__ void __launch_bounds__(256) tf_ctl_iv_bias_swap(const float* acc, const float* b, const float* r, + float* y, int rows, int cols, int by_row, + int add) { + tf_iv::bias(acc, b, r, y, rows, cols, by_row, add); +} + +#define TF_CTL_IV_TOKENS(name, flaw) \ + extern "C" __global__ void __launch_bounds__(256) name(const float* moments, const float* noise, const float* m, \ + const float* s, float* out, int rows, int cols, int pad) { \ + tf_iv::tokens(moments, noise, m, s, out, rows, cols, pad); \ + } +TF_CTL_IV_TOKENS(tf_ctl_iv_tokens_axes, tf_iv::tokens_axes) +TF_CTL_IV_TOKENS(tf_ctl_iv_tokens_clamp, tf_iv::tokens_clamp) +#undef TF_CTL_IV_TOKENS diff --git a/zig/src/core/contract/idol_vae_cuda_records.zig b/zig/src/core/contract/idol_vae_cuda_records.zig new file mode 100644 index 0000000000..be50dbb21a --- /dev/null +++ b/zig/src/core/contract/idol_vae_cuda_records.zig @@ -0,0 +1,22 @@ +//! Idol's image encoder kernels on CUDA (idol_vae.cu): every SM's compile records from nvcc 13.0, the release compiler. +const registry = @import("registry.zig"); + +/// nvcc 13.0's usage (cuobjdump -res-usage of the glue-flag image; sm_120's 1024 reserved shared bytes left out). +pub const cuda = [_]registry.CudaRecord{ + .{ .launch = .{ .name = "idol vae tf_iv_conv_in", .symbol = "tf_iv_conv_in", .image = "idol_vae", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 24, .static = 0, .bound = 256 }, .{ .regs = 25, .static = 0, .bound = 256 }, .{ .regs = 25, .static = 0, .bound = 256 }, .{ .regs = 37, .static = 0, .bound = 256 }, .{ .regs = 37, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "idol vae tf_iv_conv", .symbol = "tf_iv_conv", .image = "idol_vae", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 128, .static = 16640, .bound = 256 }, .{ .regs = 128, .static = 16640, .bound = 256 }, .{ .regs = 128, .static = 16640, .bound = 256 }, .{ .regs = 145, .static = 16640, .bound = 256 }, .{ .regs = 145, .static = 16640, .bound = 256 } } }, + .{ .launch = .{ .name = "idol vae tf_iv_group_part", .symbol = "tf_iv_group_part", .image = "idol_vae", .threads = 512, .min_sm = 80 }, .usage = .{ .{ .regs = 32, .static = 2048, .bound = 512 }, .{ .regs = 40, .static = 2048, .bound = 512 }, .{ .regs = 40, .static = 2048, .bound = 512 }, .{ .regs = 34, .static = 2048, .bound = 512 }, .{ .regs = 40, .static = 2048, .bound = 512 } } }, + .{ .launch = .{ .name = "idol vae tf_iv_group_finish", .symbol = "tf_iv_group_finish", .image = "idol_vae", .threads = 1024, .min_sm = 80 }, .usage = .{ .{ .regs = 31, .static = 0, .bound = 1024 }, .{ .regs = 40, .static = 0, .bound = 1024 }, .{ .regs = 40, .static = 0, .bound = 1024 }, .{ .regs = 44, .static = 0, .bound = 1024 }, .{ .regs = 44, .static = 0, .bound = 1024 } } }, + .{ .launch = .{ .name = "idol vae tf_iv_norm", .symbol = "tf_iv_norm", .image = "idol_vae", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 17, .static = 0, .bound = 256 }, .{ .regs = 16, .static = 0, .bound = 256 }, .{ .regs = 16, .static = 0, .bound = 256 }, .{ .regs = 16, .static = 0, .bound = 256 }, .{ .regs = 16, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "idol vae tf_iv_bias", .symbol = "tf_iv_bias", .image = "idol_vae", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 16, .static = 0, .bound = 256 }, .{ .regs = 16, .static = 0, .bound = 256 }, .{ .regs = 16, .static = 0, .bound = 256 }, .{ .regs = 16, .static = 0, .bound = 256 }, .{ .regs = 16, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "idol vae tf_iv_softmax", .symbol = "tf_iv_softmax", .image = "idol_vae", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 91, .static = 32, .bound = 256 }, .{ .regs = 109, .static = 32, .bound = 256 }, .{ .regs = 109, .static = 32, .bound = 256 }, .{ .regs = 196, .static = 32, .bound = 256 }, .{ .regs = 196, .static = 32, .bound = 256 } } }, + .{ .launch = .{ .name = "idol vae tf_iv_tokens", .symbol = "tf_iv_tokens", .image = "idol_vae", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 19, .static = 0, .bound = 256 }, .{ .regs = 20, .static = 0, .bound = 256 }, .{ .regs = 20, .static = 0, .bound = 256 }, .{ .regs = 18, .static = 0, .bound = 256 }, .{ .regs = 18, .static = 0, .bound = 256 } } }, +}; + +test "Idol VAE CUDA: measured launches fit their physical resource limits" { + const t = @import("std").testing; + for (cuda) |r| { + for (r.usage) |u| try t.expect(u.regs > 0); + try t.expectEqual(@as(usize, 0), try registry.checkCuda(t.allocator, &.{r}, true)); + } +} diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 22e8338f01..a8c15337b5 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -205,7 +205,7 @@ pub const cuda = @import("flashnext_cuda_records.zig").cuda ++ [_]CudaRecord{ .{ .launch = .{ .name = "experts3_prefill down", .symbol = "_ZN19tf_experts3_prefill14prefill_kernelILi1ELi3ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_PvifS6_i", .image = "experts3_prefill", .threads = 256, .dynamic = 47104, .opt_in = true }, .usage = .{ .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 98, .static = 0, .bound = 256 }, .{ .regs = 98, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "experts3 tf_experts3_widen", .symbol = "tf_experts3_widen", .image = "experts3", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 32, .static = 0, .bound = 0 } } }, .{ .launch = .{ .name = "experts3 tf_experts3_planes", .symbol = "tf_experts3_planes", .image = "experts3", .threads = 256 }, .usage = .{ .{ .regs = 26, .static = 0, .bound = 0 }, .{ .regs = 28, .static = 0, .bound = 0 }, .{ .regs = 28, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 } } }, -} ++ @import("registry_media.zig").cuda ++ @import("nemotron_prompt_records.zig").cuda ++ @import("glm_cuda_records.zig").cuda ++ @import("registry_shared_attention.zig").records(CudaRecord) ++ @import("gemma_cuda_records.zig").cuda ++ @import("spark_cuda_records.zig").cuda ++ @import("row_cuda_records.zig").cuda ++ @import("flashnext_tile_records.zig").cuda ++ @import("window_rows_records.zig").cuda ++ @import("modelopt_lane_records.zig").cuda ++ @import("nvfp4_fact_records.zig").cuda ++ @import("lane_half_records.zig").cuda ++ @import("pocket_records.zig").cuda ++ @import("minimax_cuda_records.zig").cuda ++ @import("moss_image_cuda_records.zig").cuda ++ @import("idol_cuda_records.zig").cuda; +} ++ @import("registry_media.zig").cuda ++ @import("nemotron_prompt_records.zig").cuda ++ @import("glm_cuda_records.zig").cuda ++ @import("registry_shared_attention.zig").records(CudaRecord) ++ @import("gemma_cuda_records.zig").cuda ++ @import("spark_cuda_records.zig").cuda ++ @import("row_cuda_records.zig").cuda ++ @import("flashnext_tile_records.zig").cuda ++ @import("window_rows_records.zig").cuda ++ @import("modelopt_lane_records.zig").cuda ++ @import("nvfp4_fact_records.zig").cuda ++ @import("lane_half_records.zig").cuda ++ @import("pocket_records.zig").cuda ++ @import("minimax_cuda_records.zig").cuda ++ @import("moss_image_cuda_records.zig").cuda ++ @import("idol_cuda_records.zig").cuda ++ @import("idol_vae_cuda_records.zig").cuda; /// The nvcc release `cuda`'s usage records come from: the compiler that builds the release fatbins. pub const cuda_release: limits.Nvcc = .{ .major = 13, .minor = 0 }; @@ -586,6 +586,7 @@ test { _ = @import("glm_kv_pending.zig"); _ = @import("minimax_cuda_records.zig"); _ = @import("idol_cuda_records.zig"); + _ = @import("idol_vae_cuda_records.zig"); _ = @import("glm_f32w_pending.zig"); _ = @import("registry_pending_test.zig"); _ = @import("flashnext_pages_pending.zig"); diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index ab6626047c..f4a1e025ec 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -73,6 +73,7 @@ pub const glm_index: []const u8 = if (available) &Blob("fatbin_glm_index").bytes pub const minimax_m3: []const u8 = if (available) &Blob("fatbin_minimax_m3").bytes else &.{}; pub const minimax_m3_attn: []const u8 = if (available) &Blob("fatbin_minimax_m3_attn").bytes else &.{}; pub const idol: []const u8 = if (available) &Blob("fatbin_idol").bytes else &.{}; +pub const idol_vae: []const u8 = if (available) &Blob("fatbin_idol_vae").bytes else &.{}; pub const glm_quant: []const u8 = if (available) &Blob("fatbin_glm_quant").bytes else &.{}; pub const glm_vision: []const u8 = if (available) &Blob("fatbin_glm_vision").bytes else &.{}; pub const glm_mm: []const u8 = if (available) &Blob("fatbin_glm_mm").bytes else &.{}; diff --git a/zig/src/families/idol/cuda.zig b/zig/src/families/idol/cuda.zig index ea600075b4..eaeede5f29 100644 --- a/zig/src/families/idol/cuda.zig +++ b/zig/src/families/idol/cuda.zig @@ -1,4 +1,4 @@ -//! Idol's denoiser on CUDA: driven by tf-idol-cuda beside the lane engine, as the Metal family is by tf-idol. +//! Idol on CUDA: the denoiser (tf-idol) and the image encoder (tf-idol-vae), beside the lane engine as on Metal. pub const config = @import("config.zig"); pub const host = @import("host.zig"); pub const reference = @import("reference.zig"); @@ -8,12 +8,18 @@ pub const Condition = @import("cuda_model.zig").Condition; pub const Timing = @import("cuda_model.zig").Timing; pub const seconds = @import("cuda_model.zig").seconds; pub const admitted = @import("cuda_model.zig").admitted; +pub const vae = @import("cuda_vae.zig"); +pub const vae_kernels = @import("cuda_vae_kernels.zig"); test { _ = @import("cuda_kernels.zig"); _ = @import("cuda_gemm.zig"); _ = @import("cuda_model.zig"); _ = @import("cuda_weights.zig"); + _ = @import("cuda_vae.zig"); + _ = @import("cuda_vae_kernels.zig"); + _ = @import("cuda_vae_weights.zig"); + _ = @import("vae_config.zig"); _ = config; _ = host; _ = reference; diff --git a/zig/src/families/idol/cuda_vae.zig b/zig/src/families/idol/cuda_vae.zig new file mode 100644 index 0000000000..84fc176692 --- /dev/null +++ b/zig/src/families/idol/cuda_vae.zig @@ -0,0 +1,226 @@ +//! Idol's image encoder on CUDA: FLUX.2's VAE encoder in fp32, a 1024 x 1024 picture to the condition's 4,101 rows. +const std = @import("std"); +const cuda = @import("cuda"); +const cfg = @import("vae_config.zig"); +const ks = @import("cuda_vae_kernels.zig"); +const wts = @import("cuda_vae_weights.zig"); +const gm = @import("cuda_gemm.zig"); +const model = @import("cuda_model.zig"); + +pub const size = cfg.size; +pub const latent = cfg.latent; +pub const pad = cfg.pad; +pub const rows = cfg.rows; +pub const channels = cfg.channels; +pub const Stages = cfg.Stages; +pub const Verdict = cfg.Verdict; +pub const gate = cfg.gate; + +pub const Timing = struct { wall: f64, gpu: f64 }; + +/// The widest picture tensor: the first stage's 1024 x 1024 x 128. +const most = size * size * cfg.widths[0]; +/// The mid block's positions, and the query rows a softmax pass takes (its scores fill one picture tensor at most). +const mid = latent * latent; +pub const query_chunk = 4096; + +/// The scratch's parts in bytes, as Encoder's views: picture, noise, four picture tensors, moments, rows, GroupNorm's. +const parts = [_]usize{ size * size * 3 * 4, mid * cfg.latent_channels * 4, most * 4, most * 4, most * 4, most * 4, mid * 2 * cfg.latent_channels * 4, rows * channels * 4, 2 * ks.groups * 4, ks.groups * ks.chunks(size * size) * 4 }; + +pub fn scratchBytes() usize { + var total: usize = 0; + for (parts) |n| total += std.mem.alignForward(usize, n, 256); + return total; +} + +pub const Encoder = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + stream: cuda.Stream, + k: ks.Kernels, + gemms: gm.Gemms, + w: wts.Weights, + scratch: cuda.DeviceBuffer, + marks: [2]cuda.Event, + image: u64 = 0, + noise: u64 = 0, + t: [4]u64 = @splat(0), + moments: u64 = 0, + out: u64 = 0, + stats: u64 = 0, + part: u64 = 0, + + /// The vae/ folder `dir` (FLUX.2-klein-4B's, or FLUX.2-dev's rounded to bf16) on `ctx`'s GPU, if admitted. + pub fn load(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8) !*Encoder { + const sm = try ctx.capability(); + if (!model.admitted(sm)) { + var buf: [256]u8 = undefined; + const name = ctx.name(&buf) catch "this GPU"; + std.log.err("idol vae: {s} (sm_{d}) is not a GPU this family admits; it runs on sm_120 and sm_121 alone", .{ name, sm }); + return error.UncheckedGpu; + } + const d = ctx.d; + const v = try gpa.create(Encoder); + errdefer gpa.destroy(v); + v.* = .{ .gpa = gpa, .d = d, .stream = undefined, .k = undefined, .gemms = undefined, .w = undefined, .scratch = undefined, .marks = undefined }; + v.stream = try cuda.Stream.init(d, false); + errdefer v.stream.deinit(); + v.k = try ks.Kernels.load(d); + errdefer v.k.unload(); + v.gemms = try gm.Gemms.init(gpa, d); + errdefer v.gemms.deinit(); + v.w = try wts.load(gpa, io, d, dir); + errdefer v.w.deinit(); + v.scratch = try cuda.DeviceBuffer.alloc(d, scratchBytes()); + errdefer v.scratch.free(); + v.marks[0] = try cuda.Event.init(d, true); + errdefer v.marks[0].deinit(); + v.marks[1] = try cuda.Event.init(d, true); + errdefer v.marks[1].deinit(); + var at: u64 = v.scratch.ptr; + const views = [_]*u64{ &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.* = at; + at += std.mem.alignForward(usize, n, 256); + } + return v; + } + + pub fn deinit(v: *Encoder) void { + for (&v.marks) |*m| m.deinit(); + v.scratch.free(); + v.w.deinit(); + v.gemms.deinit(); + v.k.unload(); + v.stream.deinit(); + v.gpa.destroy(v); + } + + fn upload(v: *const Encoder, at: u64, bytes: []const u8) !void { + try v.d.check(v.d.api.cuMemcpyHtoDAsync_v2(at, bytes.ptr, bytes.len, v.stream.handle), "cuMemcpyHtoDAsync"); + } + + fn download(v: *const Encoder, out: []u8, at: u64) !void { + try v.d.check(v.d.api.cuMemcpyDtoHAsync_v2(out.ptr, at, out.len, v.stream.handle), "cuMemcpyDtoHAsync"); + try v.stream.synchronize(); + } + + /// 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 * cfg.latent_channels or out.len != rows * channels) return error.LatentRows; + const started = model.seconds(); + try v.upload(v.image, std.mem.sliceAsBytes(image)); + try v.upload(v.noise, std.mem.sliceAsBytes(noise)); + try v.marks[0].record(v.stream); + const s = v.stream; + var R: usize = size; + try ks.convIn(&v.k, s, 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..) |st, i| { + for (st.resnets) |r| x = try v.resnet(r, x, R); + if (st.down) |d| { + const y = (x + 1) % v.t.len; + try ks.conv(&v.k, s, v.t[x], d.w, d.b, null, v.t[y], ks.down(R, d.cin, d.cout, d.ldw)); + x = y; + R /= 2; + } + if (stages.down[i]) |dst| try v.download(std.mem.sliceAsBytes(dst), v.t[x]); + } + x = try v.resnet(v.w.mid[0], x, R); + x = try v.attention(x, R * R); + x = try v.resnet(v.w.mid[1], x, R); + if (stages.mid) |dst| try v.download(std.mem.sliceAsBytes(dst), v.t[x]); + const n = (x + 1) % v.t.len; + const o = (x + 2) % v.t.len; + try v.norm(v.t[x], v.t[n], R * R, cfg.widths[3], v.w.norm_out, true); + try v.conv(v.w.conv_out, v.t[n], v.t[o], R, null); + try v.conv(v.w.quant, v.t[o], v.moments, R, null); + try ks.tokens(&v.k, s, v.moments, v.noise, v.w.bn_mean, v.w.bn_std, v.out, R, R, pad); + try v.marks[1].record(s); + try v.download(std.mem.sliceAsBytes(out), v.out); + if (stages.moments) |dst| try v.download(std.mem.sliceAsBytes(dst), v.moments); + const gpu = @as(f64, try cuda.Event.elapsedMs(v.marks[0], v.marks[1])) / 1000; + return .{ .wall = model.seconds() - started, .gpu = gpu }; + } + + /// 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, r: wts.Resnet, x: usize, R: usize) !usize { + const n = (x + 1) % v.t.len; + const h = (x + 2) % v.t.len; + const sc = (x + 3) % v.t.len; + try v.norm(v.t[x], v.t[n], R * R, r.conv1.cin, r.norm1, true); + try v.conv(r.conv1, v.t[n], v.t[h], R, null); + try v.norm(v.t[h], v.t[n], R * R, r.conv1.cout, r.norm2, true); + const skip = if (r.shortcut) |c| blk: { + try v.conv(c, v.t[x], v.t[sc], R, null); + break :blk v.t[sc]; + } else v.t[x]; + try v.conv(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 and v, one head of 512, out + t[x]. + fn attention(v: *Encoder, x: usize, n: usize) !usize { + // cuBLASLt in true fp32 (CUBLAS_COMPUTE_32F, fp32 operands); v as W_v xn^T, so P V is cuBLASLt's X W^T too + const W = cfg.width; + const s = v.stream; + const xn = v.t[(x + 1) % v.t.len]; + const qkv = v.t[(x + 2) % v.t.len]; + const scores = v.t[(x + 3) % v.t.len]; + const q = qkv; + const k = qkv + n * W * 4; + const vt = qkv + 2 * n * W * 4; + const att = qkv + 3 * n * W * 4; + try v.norm(v.t[x], xn, n, W, v.w.attn_norm, false); + try v.gemms.run(s, .f32, xn, v.w.q.w, q, n, W, W); + try ks.bias(&v.k, s, q, v.w.q.b, null, q, n, W, false); + try v.gemms.run(s, .f32, xn, v.w.k.w, k, n, W, W); + try ks.bias(&v.k, s, k, v.w.k.b, null, k, n, W, false); + try v.gemms.run(s, .f32, v.w.v.w, xn, vt, W, n, W); + try ks.bias(&v.k, s, vt, v.w.v.b, null, vt, W, n, true); + const scale = 1.0 / @sqrt(@as(f32, W)); + var r0: usize = 0; + while (r0 < n) : (r0 += query_chunk) { + const m = @min(query_chunk, n - r0); + try v.gemms.run(s, .f32, q + r0 * W * 4, k, scores, m, n, W); + try ks.softmax(&v.k, s, scores, m, n, scale); + try v.gemms.run(s, .f32, scores, vt, att + r0 * W * 4, m, W, n); + } + // the output projection's sums into q's rows (q is spent), then its bias and the residual into xn's + try v.gemms.run(s, .f32, att, v.w.out.w, q, n, W, W); + try ks.bias(&v.k, s, q, v.w.out.b, v.t[x], xn, n, W, false); + return (x + 1) % v.t.len; + } + + /// y = GroupNorm(x) with `nm` (then SiLU when `silu`), x and y [pixels, channels], statistics in two passes. + fn norm(v: *const Encoder, x: u64, y: u64, n: usize, ch: usize, nm: wts.Norm, silu: bool) !void { + for ([_]bool{ false, true }) |centered| { + try ks.groupPart(&v.k, v.stream, x, v.stats, v.part, n, ch, centered); + try ks.groupFinish(&v.k, v.stream, v.part, v.stats, n, ch, centered, cfg.eps); + } + try ks.norm(&v.k, v.stream, x, v.stats, nm.g, nm.b, y, n, ch, silu); + } + + /// y = conv(x) + b (+ r) at R x R, a same-size 3x3 or a 1x1. + fn conv(v: *const Encoder, c: wts.Conv, x: u64, y: u64, R: usize, r: ?u64) !void { + const shape = if (c.taps == 9) ks.same(R, c.cin, c.cout, c.ldw) else ks.one(R, c.cin, c.cout, c.ldw); + try ks.conv(&v.k, v.stream, x, c.w, c.b, r, y, shape); + } +}; + +test "scratch: four picture tensors, and the mid block's attention fits two of them" { + try std.testing.expect(scratchBytes() < 2400 << 20); + // q, k, v^T and the attention's output in one picture tensor; a chunk's scores in another + try std.testing.expect(4 * mid * cfg.width <= most); + try std.testing.expect(query_chunk * mid <= most); + try std.testing.expect(mid % query_chunk == 0 and mid <= ks.softmax_max); + // GroupNorm's chunk sums at the widest stage, and every kernel index as c_int + try std.testing.expectEqual(@as(usize, 4096), ks.chunks(size * size)); + try std.testing.expect(most < std.math.maxInt(c_int)); +} + +test "admission: the encoder takes the denoiser's GPUs, sm_120 and sm_121 alone" { + try std.testing.expect(model.admitted(121) and model.admitted(120)); + for ([_]u32{ 80, 86, 89, 90, 100 }) |sm| try std.testing.expect(!model.admitted(sm)); +} diff --git a/zig/src/families/idol/cuda_vae_kernels.zig b/zig/src/families/idol/cuda_vae_kernels.zig new file mode 100644 index 0000000000..cef5cbc835 --- /dev/null +++ b/zig/src/families/idol/cuda_vae_kernels.zig @@ -0,0 +1,182 @@ +//! Idol's image encoder kernels on CUDA (idol_vae.cu, fp32): convolutions, GroupNorm, attention's softmax and the rows. +const std = @import("std"); +const cuda = @import("cuda"); + +/// GroupNorm's groups, the pixels a statistics part sums, the convolution's tile, and the softmax's widest row. +pub const groups = 32; +pub const chunk = 256; +pub const tile_m = 128; +pub const tile_n = 128; +pub const tile_k = 8; +pub const softmax_max = 256 * 64; + +/// Each function by its symbol, `tf_iv_` and the field's name. +pub const names = [_][]const u8{ "conv_in", "conv", "group_part", "group_finish", "norm", "bias", "softmax", "tokens" }; + +pub const Kernels = struct { + mod: cuda.Module, + conv_in: cuda.Function, + conv: cuda.Function, + group_part: cuda.Function, + group_finish: cuda.Function, + norm: cuda.Function, + bias: cuda.Function, + softmax: cuda.Function, + tokens: cuda.Function, + + pub fn load(d: *const cuda.Driver) !Kernels { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var mod = try cuda.Module.load(d, cuda.kernels.idol_vae); + errdefer mod.unload(); + var k: Kernels = undefined; + k.mod = mod; + inline for (names) |n| @field(k, n) = try mod.function("tf_iv_" ++ n); + return k; + } + + pub fn unload(k: *Kernels) void { + k.mod.unload(); + } +}; + +fn int(x: usize) c_int { + return @intCast(x); +} + +fn ceil(n: usize, d: usize) u32 { + return @intCast((n + d - 1) / d); +} + +fn go(f: cuda.Function, s: cuda.Stream, grid: cuda.Dim3, block: cuda.Dim3, a: *cuda.Args) !void { + try cuda.launch.launch(f, .{ .grid = grid, .block = block }, s, a); +} + +fn addAll(a: *cuda.Args, ptrs: []const u64) void { + for (ptrs) |p| a.add(p); +} + +fn addInts(a: *cuda.Args, vals: []const usize) void { + for (vals) |v| a.add(int(v)); +} + +/// tf_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(k: *const Kernels, s: cuda.Stream, x: u64, w: u64, b: u64, y: u64, rows: usize, cols: usize) !void { + var a: cuda.Args = .{}; + addAll(&a, &.{ x, w, b, y }); + addInts(&a, &.{ rows, cols }); + try go(k.conv_in, s, .{ .x = ceil(rows * cols, 2) }, .{ .x = 128, .y = 2 }, &a); +} + +/// A convolution's output and input pictures, channels, weights' row, taps (1 or 9), stride and padding above and left. +pub const ConvShape = struct { rows: usize, cols: usize, in_rows: usize, in_cols: usize, cin: usize, cout: usize, ldw: usize, taps: usize = 1, stride: usize = 1, pad: usize = 0 }; + +/// The convolutions the encoder runs: a same-size 3x3, the stride-2 downsample, a 1x1 at `R` x `R`. +pub fn same(R: usize, cin: usize, cout: usize, ldw: usize) ConvShape { + return .{ .rows = R, .cols = R, .in_rows = R, .in_cols = R, .cin = cin, .cout = cout, .ldw = ldw, .taps = 9, .pad = 1 }; +} +pub fn down(R: usize, cin: usize, cout: usize, ldw: usize) ConvShape { + return .{ .rows = R / 2, .cols = R / 2, .in_rows = R, .in_cols = R, .cin = cin, .cout = cout, .ldw = ldw, .taps = 9, .stride = 2 }; +} +pub fn one(R: usize, cin: usize, cout: usize, ldw: usize) ConvShape { + return .{ .rows = R, .cols = R, .in_rows = R, .in_cols = R, .cin = cin, .cout = cout, .ldw = ldw }; +} + +/// The weights' row length for `cout` output channels: whole column tiles, so no weight load needs a bound. +pub fn ldwOf(cout: usize) usize { + return std.mem.alignForward(usize, cout, tile_n); +} + +/// What tf_iv_conv requires of a shape: K in whole steps of one tap, float4 rows, whole weight tiles. +pub fn convOk(c: ConvShape) bool { + return c.cin % tile_k == 0 and c.cout % 4 == 0 and c.ldw % tile_n == 0 and c.ldw >= c.cout and (c.taps == 1 or c.taps == 9) and + (c.stride == 1 or c.stride == 2) and c.rows * c.cols > 0; +} + +/// The grid of a convolution: 128-pixel tiles on x, 128-channel tiles of the weights' row on y. +pub fn convGrid(c: ConvShape) cuda.Dim3 { + return .{ .x = ceil(c.rows * c.cols, tile_m), .y = @intCast(c.ldw / tile_n) }; +} + +/// tf_iv_conv: y = conv(x) + b (+ r when given), an implicit GEMM in fp32 FMAs (no tensor cores, no TF32). +pub fn conv(k: *const Kernels, s: cuda.Stream, x: u64, w: u64, b: u64, r: ?u64, y: u64, c: ConvShape) !void { + std.debug.assert(convOk(c)); + var a: cuda.Args = .{}; + addAll(&a, &.{ x, w, b, r orelse y, y }); + addInts(&a, &.{ c.rows, c.cols, c.in_rows, c.in_cols, c.cin, c.cout, c.ldw, c.taps, c.stride, c.pad, @intFromBool(r != null) }); + try go(k.conv, s, convGrid(c), .{ .x = 256 }, &a); +} + +/// 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; +} + +/// tf_iv_group_part: part [32, chunks] = each chunk's sum of a group's values, or (centered) of (x - mean)^2. +pub fn groupPart(k: *const Kernels, s: cuda.Stream, x: u64, stats: u64, part: u64, n: usize, channels: usize, centered: bool) !void { + std.debug.assert(channels % groups == 0 and channels <= 512); + var a: cuda.Args = .{}; + addAll(&a, &.{ x, stats, part }); + addInts(&a, &.{ n, channels, chunk, chunks(n), @intFromBool(centered) }); + try go(k.group_part, s, .{ .x = @intCast(chunks(n)) }, .{ .x = @intCast(channels) }, &a); +} + +/// tf_iv_group_finish: stats' 32 means from part, or (centered) its 32 rstd. +pub fn groupFinish(k: *const Kernels, s: cuda.Stream, part: u64, stats: u64, n: usize, channels: usize, centered: bool, eps: f32) !void { + var a: cuda.Args = .{}; + addAll(&a, &.{ part, stats }); + addInts(&a, &.{ n, channels, chunks(n), @intFromBool(centered) }); + a.add(eps); + try go(k.group_finish, s, .{ .x = 1 }, .{ .x = 32 * groups }, &a); +} + +/// tf_iv_norm: y = GroupNorm(x) by stats (32 means, then 32 rstd) with g and beta, then SiLU when `silu`. +pub fn norm(k: *const Kernels, s: cuda.Stream, x: u64, stats: u64, g: u64, beta: u64, y: u64, n: usize, channels: usize, silu: bool) !void { + var a: cuda.Args = .{}; + addAll(&a, &.{ x, stats, g, beta, y }); + addInts(&a, &.{ n, channels, @intFromBool(silu) }); + try go(k.norm, s, .{ .x = ceil(n * channels, 256) }, .{ .x = 256 }, &a); +} + +/// tf_iv_bias: y [rows, cols] = acc + b by column (or by row when `by_row`) (+ r when given); y may be acc. +pub fn bias(k: *const Kernels, s: cuda.Stream, acc: u64, b: u64, r: ?u64, y: u64, rows: usize, cols: usize, by_row: bool) !void { + var a: cuda.Args = .{}; + addAll(&a, &.{ acc, b, r orelse y, y }); + addInts(&a, &.{ rows, cols, @intFromBool(by_row), @intFromBool(r != null) }); + try go(k.bias, s, .{ .x = ceil(rows * cols, 256) }, .{ .x = 256 }, &a); +} + +/// tf_iv_softmax: each of `rows` rows of s [rows, n] becomes softmax(s scale), in place; n <= softmax_max. +pub fn softmax(k: *const Kernels, s: cuda.Stream, scores: u64, rows: usize, n: usize, scale: f32) !void { + std.debug.assert(n <= softmax_max); + var a: cuda.Args = .{}; + a.add(scores); + a.add(int(n)); + a.add(scale); + try go(k.softmax, s, .{ .x = @intCast(rows) }, .{ .x = 256 }, &a); +} + +/// tf_iv_tokens: out [pad + rows * cols / 4, 128] from moments [rows, cols, 64], noise [rows, cols, 32], m and s [128]. +pub fn tokens(k: *const Kernels, s: cuda.Stream, moments: u64, noise: u64, m: u64, sd: u64, out: u64, rows: usize, cols: usize, pad: usize) !void { + var a: cuda.Args = .{}; + addAll(&a, &.{ moments, noise, m, sd, out }); + addInts(&a, &.{ rows, cols, pad }); + try go(k.tokens, s, .{ .x = ceil(pad + rows * cols / 4, 2) }, .{ .x = 128, .y = 2 }, &a); +} + +test "launch geometry: 128 x 128 conv tiles over the weights' padded row, a group-part block a chunk" { + const c = same(1024, 128, 128, ldwOf(128)); + try std.testing.expect(convOk(c)); + try std.testing.expectEqual(@as(u32, 8192), convGrid(c).x); + try std.testing.expectEqual(@as(u32, 1), convGrid(c).y); + const o = same(128, 512, 64, ldwOf(64)); + try std.testing.expectEqual(@as(usize, 128), o.ldw); + try std.testing.expect(convOk(o)); + const d = down(1024, 128, 128, ldwOf(128)); + try std.testing.expectEqual(@as(usize, 512), d.rows); + try std.testing.expectEqual(@as(u32, 2048), convGrid(d).x); + try std.testing.expect(!convOk(.{ .rows = 4, .cols = 4, .in_rows = 4, .in_cols = 4, .cin = 3, .cout = 128, .ldw = 128 })); + try std.testing.expectEqual(@as(usize, 4096), chunks(1024 * 1024)); + try std.testing.expectEqual(@as(usize, 16384), softmax_max); + // every grid dimension and kernel index fits its type at the widest stage + try std.testing.expect(1024 * 1024 * 128 < std.math.maxInt(c_int)); +} diff --git a/zig/src/families/idol/cuda_vae_weights.zig b/zig/src/families/idol/cuda_vae_weights.zig new file mode 100644 index 0000000000..d31a614264 --- /dev/null +++ b/zig/src/families/idol/cuda_vae_weights.zig @@ -0,0 +1,208 @@ +//! Idol's image encoder weights on CUDA: FLUX.2's VAE encoder, quant_conv and batch norm, fp32 of the SDK's bf16. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const cfg = @import("vae_config.zig"); +const ks = @import("cuda_vae_kernels.zig"); + +/// A conv's weights [taps * cin, ldw] by tap, then input channel (ldw: cout in whole column tiles), and its bias. +pub const Conv = struct { w: u64, b: u64, cin: usize, cout: usize, ldw: usize, taps: usize }; +pub const Norm = struct { g: u64, b: u64 }; +pub const Resnet = struct { norm1: Norm, conv1: Conv, norm2: Norm, conv2: Conv, shortcut: ?Conv }; +pub const Stage = struct { resnets: [cfg.resnets]Resnet, down: ?Conv }; +/// One of the mid block's attention linears as torch holds it, w [512 out, 512 in] (cuBLASLt's X . W^T), and bias. +pub const Linear = struct { w: u64, b: u64 }; + +pub const Weights = struct { + buf: cuda.DeviceBuffer, + conv_in: Conv, + stages: [4]Stage, + mid: [2]Resnet, + attn_norm: Norm, + q: Linear, + k: Linear, + v: Linear, + out: Linear, + norm_out: Norm, + conv_out: Conv, + quant: Conv, + /// The batch norm's running mean and sqrt(running_var + eps), each of 128. + bn_mean: u64, + bn_std: u64, + + pub fn deinit(w: *Weights) void { + w.buf.free(); + } +}; + +/// 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, d: *const cuda.Driver, 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 c: core.Checkpoint = .{ .gpa = gpa, .io = io }; + defer c.close(); + try c.add(dir, "diffusion_pytorch_model.safetensors"); + var w: Weights = undefined; + var p: Packer = .{ .c = &c, .arena = arena }; + try layout(&w, &p); + const left = unread(&c); + if (left != 0) { + std.log.err("idol vae: the checkpoint has {d} encoder tensors the encoder does not read", .{left}); + return error.UnexpectedTensors; + } + const host = try arena.alloc(f32, p.at); + w.buf = try cuda.DeviceBuffer.alloc(d, p.at * 4); + errdefer w.buf.free(); + p = .{ .c = &c, .arena = arena, .out = host, .base = w.buf.ptr }; + try layout(&w, &p); + try w.buf.upload(0, std.mem.sliceAsBytes(host)); + 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 { + const widths = cfg.widths; + var name: [96]u8 = undefined; + w.conv_in = try p.conv("encoder.conv_in", 3, widths[0], 9, false); + 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, 9, true) 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.q = try p.linear(at ++ ".to_q", W); + w.k = try p.linear(at ++ ".to_k", W); + w.v = try p.linear(at ++ ".to_v", W); + w.out = try p.linear(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 * cfg.latent_channels, 9, true); + w.quant = try p.conv("quant_conv", 2 * cfg.latent_channels, 2 * cfg.latent_channels, 1, true); + const n = 4 * cfg.latent_channels; + const mean = try p.tensor("bn.running_mean", &.{n}); + const var_ = try p.tensor("bn.running_var", &.{n}); + const m_at = p.alloc(n); + const s_at = p.alloc(n); + w.bn_mean = p.addr(m_at); + w.bn_std = p.addr(s_at); + if (p.out) |out| for (out[m_at..][0..n], out[s_at..][0..n], mean, var_) |*mo, *so, mv, v| { + mo.* = mv; + so.* = cfg.bnStd(v); + }; +} + +/// Encoder-side tensors (vae_config.encoderSide) no `get` took. +fn unread(c: *const core.Checkpoint) usize { + var n: usize = 0; + for (c.files.items) |*f| for (f.names.keys()) |k| { + if (cfg.encoderSide(k) and !c.used.contains(k)) 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: *core.Checkpoint, + arena: std.mem.Allocator, + out: ?[]f32 = null, + base: u64 = 0, + at: usize = 0, + + /// n floats' offset, the next one 256 bytes on. + fn alloc(p: *Packer, n: usize) usize { + const r = p.at; + p.at = std.mem.alignForward(usize, p.at + n, 64); + return r; + } + + fn addr(p: *const Packer, off: usize) u64 { + return p.base + off * 4; + } + + /// 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; + } + if (p.out == null) return &.{}; + const n = t.numel(); + const v = try p.arena.alloc(f32, n); + switch (t.dtype) { + .bf16 => for (v, 0..) |*o, i| { + o.* = @bitCast(@as(u32, std.mem.readInt(u16, t.bytes[2 * i ..][0..2], .little)) << 16); + }, + else => for (v, 0..) |*o, i| { + o.* = cfg.bf16(@bitCast(std.mem.readInt(u32, t.bytes[4 * i ..][0..4], .little))); + }, + } + return v; + } + + /// A copied vector `name` of n. + fn copy(p: *Packer, name: []const u8, n: usize) !u64 { + const v = try p.tensor(name, &.{n}); + const r = p.alloc(n); + if (p.out) |out| @memcpy(out[r..][0..n], v); + return p.addr(r); + } + + /// Convolution `name` (.weight [cout, cin, k, k], .bias) in tf_iv_conv's layout (conv_in's [27, 128] unpadded). + fn conv(p: *Packer, name: []const u8, cin: usize, cout: usize, taps: usize, padded: bool) !Conv { + var buf: [96]u8 = undefined; + const k: usize = if (taps == 1) 1 else 3; + const ldw = if (padded) ks.ldwOf(cout) else cout; + const w = try p.tensor(try std.fmt.bufPrint(&buf, "{s}.weight", .{name}), &.{ cout, cin, k, k }); + const r = p.alloc(taps * cin * ldw); + if (p.out) |out| { + const dst = out[r..][0 .. taps * cin * ldw]; + @memset(dst, 0); + for (0..cout) |o| for (0..cin) |i| for (0..taps) |tap| { + dst[(tap * cin + i) * ldw + o] = w[(o * cin + i) * taps + tap]; + }; + } + const b = try p.copy(try std.fmt.bufPrint(&buf, "{s}.bias", .{name}), cout); + return .{ .w = p.addr(r), .b = b, .cin = cin, .cout = cout, .ldw = ldw, .taps = taps }; + } + + /// Linear `name` (.weight [width, width], .bias) as stored. + fn linear(p: *Packer, name: []const u8, width: usize) !Linear { + var buf: [96]u8 = undefined; + const w = try p.tensor(try std.fmt.bufPrint(&buf, "{s}.weight", .{name}), &.{ width, width }); + const r = p.alloc(width * width); + if (p.out) |out| @memcpy(out[r..][0 .. width * width], w); + return .{ .w = p.addr(r), .b = try p.copy(try std.fmt.bufPrint(&buf, "{s}.bias", .{name}), width) }; + } + + 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, 9, true), + .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, 9, true), + .shortcut = if (cin != cout) try p.conv(try std.fmt.bufPrint(&buf, "{s}.conv_shortcut", .{name}), cin, cout, 1, true) else null, + }; + } +}; + +test "conv weights go tap-major, then input channel, into rows of whole column tiles" { + try std.testing.expectEqual(@as(usize, 128), ks.ldwOf(64)); + try std.testing.expectEqual(@as(usize, 512), ks.ldwOf(512)); + try std.testing.expect(cfg.encoderSide("encoder.conv_in.weight") and cfg.encoderSide("bn.running_var")); + try std.testing.expect(!cfg.encoderSide("decoder.conv_in.weight") and !cfg.encoderSide("post_quant_conv.bias")); +} diff --git a/zig/src/families/idol/idol.zig b/zig/src/families/idol/idol.zig index 70fa00e6e6..9f0a84dbb7 100644 --- a/zig/src/families/idol/idol.zig +++ b/zig/src/families/idol/idol.zig @@ -6,6 +6,7 @@ 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_config = @import("vae_config.zig"); pub const vae_kernels = @import("vae_kernels.zig"); pub const vae_weights = @import("vae_weights.zig"); pub const Model = model.Model; @@ -20,6 +21,7 @@ test { _ = model; _ = reference; _ = vae; + _ = vae_config; _ = vae_kernels; _ = vae_weights; } diff --git a/zig/src/families/idol/vae.zig b/zig/src/families/idol/vae.zig index 890cda28e4..14dcfaeaf6 100644 --- a/zig/src/families/idol/vae.zig +++ b/zig/src/families/idol/vae.zig @@ -3,20 +3,19 @@ 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 cfg = @import("vae_config.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; +/// The picture the encoder sees, its latent grid, and the condition rows (vae_config.zig). +pub const size = cfg.size; +pub const latent = cfg.latent; +pub const pad = cfg.pad; +pub const rows = cfg.rows; +pub const channels = cfg.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 }; +pub const Stages = cfg.Stages; /// The widest picture tensor: the first stage's 1024 x 1024 x 128. const most = size * size * wts.widths[0]; @@ -186,48 +185,10 @@ 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 }; +/// The encoder's gate against the float64 ruler (vae_config.zig). +pub const Verdict = cfg.Verdict; +pub const gate = cfg.gate; -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); +test "the encoder's scratch fits a GPU buffer" { 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_config.zig b/zig/src/families/idol/vae_config.zig new file mode 100644 index 0000000000..f13223a7b4 --- /dev/null +++ b/zig/src/families/idol/vae_config.zig @@ -0,0 +1,155 @@ +//! Idol's image encoder, its backend-free parts: FLUX.2's VAE geometry, the condition rows and the float64 gate. +const std = @import("std"); +const ref = @import("reference.zig"); + +/// 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; + +/// GroupNorm's groups and epsilon, and the mid block's attention width. +pub const groups = 32; +pub const eps: f32 = 1e-6; +pub const width = 512; + +/// 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 * latent_channels; + +/// 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 }; + +/// 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", 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; +} + +/// 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 checkpoint's encoder-side tensor names: the encoder's, quant_conv's and the batch norm's two statistics. +pub fn encoderSide(name: []const u8) bool { + return std.mem.startsWith(u8, name, "encoder.") or std.mem.startsWith(u8, name, "quant_conv.") or + std.mem.eql(u8, name, "bn.running_mean") or std.mem.eql(u8, name, "bn.running_var"); +} + +/// The batch norm's sqrt(running_var + eps), as the SDK forms it in bf16, here in float64 rounded once. +pub fn bnStd(v: f32) f32 { + return @floatCast(@sqrt(@as(f64, v) + 1e-4)); +} + +/// 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 rows are the denoiser's condition rows" { + try std.testing.expectEqual(@as(usize, 4101), rows); + try std.testing.expectEqual(@as(usize, 128), channels); +} + +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); +} + +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/src/families/idol/vae_kernels.zig b/zig/src/families/idol/vae_kernels.zig index 091005fee8..4a435487bb 100644 --- a/zig/src/families/idol/vae_kernels.zig +++ b/zig/src/families/idol/vae_kernels.zig @@ -3,15 +3,16 @@ const std = @import("std"); const mtl = @import("metal"); const sources = @import("kernel_sources"); pub const Ref = @import("../../core/qwen_vl/weights.zig").Ref; +const cfg = @import("vae_config.zig"); /// 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 groups = cfg.groups; +pub const eps = cfg.eps; pub const chunk = 256; -pub const width = 512; +pub const width = cfg.width; pub const Kernels = struct { lib: mtl.Library, diff --git a/zig/src/families/idol/vae_weights.zig b/zig/src/families/idol/vae_weights.zig index 817794aa39..49b0a24e77 100644 --- a/zig/src/families/idol/vae_weights.zig +++ b/zig/src/families/idol/vae_weights.zig @@ -3,53 +3,16 @@ const std = @import("std"); const mtl = @import("metal"); const ck = @import("../../core/checkpoint_metal.zig"); const ks = @import("vae_kernels.zig"); +const cfg = @import("vae_config.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; -} +/// The encoder's widths by stage, and the rest of the released geometry (vae_config.zig). +pub const widths = cfg.widths; +pub const latent_channels = cfg.latent_channels; +pub const resnets = cfg.resnets; +pub const batch_norm_eps = cfg.batch_norm_eps; +pub const check = cfg.check; +pub const bf16 = cfg.bf16; /// 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 }; @@ -78,14 +41,6 @@ pub const Weights = struct { } }; -/// 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); @@ -143,7 +98,7 @@ fn layout(w: *Weights, p: *Packer) !void { // 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)); + so.* = cfg.bnStd(v); } } } @@ -153,8 +108,7 @@ 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; + if (cfg.encoderSide(k.*)) n += 1; } return n; } @@ -279,30 +233,3 @@ const Packer = struct { }; } }; - -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/cuda/contract_family_sets.zig b/zig/tests/cuda/contract_family_sets.zig index 9635cba38a..94c6881034 100644 --- a/zig/tests/cuda/contract_family_sets.zig +++ b/zig/tests/cuda/contract_family_sets.zig @@ -1,4 +1,4 @@ -//! The family sets in the full kernel contract: Nemotron's pocket, MiniMax-M3's, MOSS-Image's, Idol's. +//! The family sets in the full kernel contract: Nemotron's pocket, MiniMax-M3's, MOSS-Image's, Idol's two. const std = @import("std"); const cuda = @import("cuda"); const check = @import("check.zig"); @@ -9,6 +9,7 @@ pub const Sets = struct { minimax: *@import("contract_minimax.zig").Set, moss: *@import("contract_moss_image.zig").Set, idol: *@import("contract_idol.zig").Set, + idol_vae: *@import("contract_idol_vae.zig").Set, /// Each set's kernels appended to `list`, in this order. pub fn init(gpu: check.Gpu, s: cuda.Stream, a: std.mem.Allocator, list: *std.ArrayList(g.Kernel)) !Sets { @@ -19,10 +20,13 @@ pub const Sets = struct { const moss = try @import("contract_moss_image.zig").Set.init(gpu, s, a, list); errdefer moss.deinit(); const idol = try @import("contract_idol.zig").Set.init(gpu, s, a, list); - return .{ .pocket = pocket, .minimax = minimax, .moss = moss, .idol = idol }; + errdefer idol.deinit(); + const idol_vae = try @import("contract_idol_vae.zig").Set.init(gpu, s, a, list); + return .{ .pocket = pocket, .minimax = minimax, .moss = moss, .idol = idol, .idol_vae = idol_vae }; } pub fn deinit(sets: Sets) void { + sets.idol_vae.deinit(); sets.idol.deinit(); sets.moss.deinit(); sets.minimax.deinit(); diff --git a/zig/tests/cuda/contract_idol_case.zig b/zig/tests/cuda/contract_idol_case.zig index 46c83434fd..258c41558e 100644 --- a/zig/tests/cuda/contract_idol_case.zig +++ b/zig/tests/cuda/contract_idol_case.zig @@ -8,13 +8,18 @@ const r = @import("contract_idol_ref.zig"); const c = r.c; const g = r.g; +/// The kernels an Op launches: its own `Kernels` (the image encoder's), else the denoiser's. +pub fn KernelsOf(comptime Op: type) type { + return if (@hasDecl(Op, "Kernels")) Op.Kernels else idol.Kernels; +} + /// The gates' hooks over an Op: its kinds, out, in_place (input the output starts as), spec, sizes, fill, launch, ref. pub fn Case(comptime Op: type) type { const n = Op.kinds.len; return struct { const Self = @This(); rig: rig_.Rig, - k: idol.Kernels, + k: KernelsOf(Op), name: []const u8, /// The Op's own form (a norm's mode, the bias's GELU), passed to each of its functions. variant: u32 = 0, @@ -130,6 +135,6 @@ pub fn Case(comptime Op: type) type { } /// The usual registration of a test: its rig on the stream, the kernels (or a control's copy), name and form. -pub fn make(comptime Op: type, rig: rig_.Rig, k: idol.Kernels, name: []const u8, variant: u32, control: ?g.Gate) Case(Op) { +pub fn make(comptime Op: type, rig: rig_.Rig, k: KernelsOf(Op), name: []const u8, variant: u32, control: ?g.Gate) Case(Op) { return .{ .rig = rig, .k = k, .name = name, .variant = variant, .control = control, .ends = .{ .a = rig.gpu.gpa } }; } diff --git a/zig/tests/cuda/contract_idol_host.zig b/zig/tests/cuda/contract_idol_host.zig index ee61edd549..c64dc8b230 100644 --- a/zig/tests/cuda/contract_idol_host.zig +++ b/zig/tests/cuda/contract_idol_host.zig @@ -1,8 +1,10 @@ -//! Idol's kernel-contract references on the host: their ranges hold the float64 values, no GPU work. +//! Idol's kernel-contract references on the host (denoiser and image encoder): ranges hold float64, no GPU work. test { _ = @import("contract_idol_ref.zig"); _ = @import("contract_idol_norms.zig"); _ = @import("contract_idol_steps.zig"); _ = @import("contract_idol_rotary.zig"); _ = @import("contract_idol_case.zig"); + _ = @import("contract_idol_vae_convs.zig"); + _ = @import("contract_idol_vae_rows.zig"); } diff --git a/zig/tests/cuda/contract_idol_vae.zig b/zig/tests/cuda/contract_idol_vae.zig new file mode 100644 index 0000000000..b6b85108ec --- /dev/null +++ b/zig/tests/cuda/contract_idol_vae.zig @@ -0,0 +1,151 @@ +//! Idol's image encoder kernels (idol_vae.cu, fp32) under the contract, each control breaking one gate. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const vae = @import("idol_vae_cuda"); +const check = @import("check.zig"); +const case = @import("contract_idol_case.zig"); +const cv = @import("contract_idol_vae_convs.zig"); +const rw = @import("contract_idol_vae_rows.zig"); + +const g = core.contract.gates; + +const controls = struct { + pub const bytes align(16) = @embedFile("fatbin_contract_controls").*; +}; + +const Same = cv.Conv(cv.forms.same); +const Residual = cv.Conv(cv.forms.residual); +const Down = cv.Conv(cv.forms.down); +const One = cv.Conv(cv.forms.one); +const ByCol = rw.Bias(false, false); +const ByColAdd = rw.Bias(false, true); +const ByRow = rw.Bias(true, false); + +/// A control: the field of Kernels it swaps for the control's symbol (Set.init names the gate that must catch it). +const Swap = struct { field: []const u8, symbol: [:0]const u8 }; +const swaps = [_]Swap{ + .{ .field = "conv_in", .symbol = "tf_ctl_iv_conv_in_half" }, + .{ .field = "conv", .symbol = "tf_ctl_iv_conv_half_in" }, + .{ .field = "conv", .symbol = "tf_ctl_iv_conv_sym_pad" }, + .{ .field = "conv", .symbol = "tf_ctl_iv_conv_half_sum" }, + .{ .field = "norm", .symbol = "tf_ctl_iv_norm_half" }, + .{ .field = "group_part", .symbol = "tf_ctl_iv_group_part_half" }, + .{ .field = "group_finish", .symbol = "tf_ctl_iv_group_finish_half" }, + .{ .field = "softmax", .symbol = "tf_ctl_iv_softmax_half" }, + .{ .field = "bias", .symbol = "tf_ctl_iv_bias_swap" }, + .{ .field = "tokens", .symbol = "tf_ctl_iv_tokens_axes" }, + .{ .field = "tokens", .symbol = "tf_ctl_iv_tokens_clamp" }, +}; + +/// Every registration, on the heap: the gates hold pointers to them. +pub const Set = struct { + gpa: std.mem.Allocator, + k: vae.Kernels, + ctl: cuda.Module, + /// The kernels with one function swapped for a control, in `swaps`' order. + bad: [swaps.len]vae.Kernels = undefined, + conv_in: case.Case(cv.ConvIn) = undefined, + same: case.Case(Same) = undefined, + residual: case.Case(Residual) = undefined, + down: case.Case(Down) = undefined, + one: case.Case(One) = undefined, + parts: [2]case.Case(rw.GroupPart) = undefined, + finishes: [2]case.Case(rw.GroupFinish) = undefined, + norms: [2]case.Case(rw.Norm) = undefined, + by_col: case.Case(ByCol) = undefined, + by_col_add: case.Case(ByColAdd) = undefined, + by_row: case.Case(ByRow) = undefined, + softmax: case.Case(rw.Softmax) = undefined, + tokens: case.Case(rw.Tokens) = undefined, + ctl_conv_in: case.Case(cv.ConvIn) = undefined, + ctl_half_in: case.Case(Same) = undefined, + ctl_sym_pad: case.Case(Down) = undefined, + ctl_half_sum: case.Case(Same) = undefined, + ctl_norm: case.Case(rw.Norm) = undefined, + ctl_part: case.Case(rw.GroupPart) = undefined, + ctl_finish: case.Case(rw.GroupFinish) = undefined, + ctl_softmax: case.Case(rw.Softmax) = undefined, + ctl_bias: case.Case(ByRow) = undefined, + ctl_axes: case.Case(rw.Tokens) = undefined, + ctl_clamp: case.Case(rw.Tokens) = undefined, + + /// The set, its kernels appended to `list` (the controls last). + pub fn init(gpu: check.Gpu, s: cuda.Stream, a: std.mem.Allocator, list: *std.ArrayList(g.Kernel)) !*Set { + const set = try gpu.gpa.create(Set); + errdefer gpu.gpa.destroy(set); + var k = try vae.Kernels.load(gpu.d); + errdefer k.unload(); + var ctl = try cuda.Module.load(gpu.d, &controls.bytes); + errdefer ctl.unload(); + set.* = .{ .gpa = gpu.gpa, .k = k, .ctl = ctl }; + inline for (swaps, 0..) |sw, i| { + set.bad[i] = k; + @field(set.bad[i], sw.field) = try ctl.function(sw.symbol); + } + const rig: @import("contract_rig.zig").Rig = .{ .gpu = gpu, .s = s }; + const n = "cuda idol vae "; + set.conv_in = case.make(cv.ConvIn, rig, k, n ++ "tf_iv_conv_in, conv3x3(2 x - 1) + b", 0, null); + set.same = case.make(Same, rig, k, n ++ "tf_iv_conv, 3x3 same size", 0, null); + set.residual = case.make(Residual, rig, k, n ++ "tf_iv_conv, 3x3 same size + residual", 0, null); + set.down = case.make(Down, rig, k, n ++ "tf_iv_conv, 3x3 stride 2 padded right and below", 0, null); + set.one = case.make(One, rig, k, n ++ "tf_iv_conv, 1x1", 0, null); + const parts = [_][]const u8{ n ++ "tf_iv_group_part, sums", n ++ "tf_iv_group_part, squared deviations" }; + const finishes = [_][]const u8{ n ++ "tf_iv_group_finish, means", n ++ "tf_iv_group_finish, rstd" }; + const norms = [_][]const u8{ n ++ "tf_iv_norm, GroupNorm affine", n ++ "tf_iv_norm, GroupNorm affine then SiLU" }; + for (&set.parts, parts, 0..) |*x, m, v| x.* = case.make(rw.GroupPart, rig, k, m, @intCast(v), null); + for (&set.finishes, finishes, 0..) |*x, m, v| x.* = case.make(rw.GroupFinish, rig, k, m, @intCast(v), null); + for (&set.norms, norms, 0..) |*x, m, v| x.* = case.make(rw.Norm, rig, k, m, @intCast(v), null); + set.by_col = case.make(ByCol, rig, k, n ++ "tf_iv_bias, by column", 0, null); + set.by_col_add = case.make(ByColAdd, rig, k, n ++ "tf_iv_bias, by column + residual", 0, null); + set.by_row = case.make(ByRow, rig, k, n ++ "tf_iv_bias, by row (v transposed)", 0, null); + set.softmax = case.make(rw.Softmax, rig, k, n ++ "tf_iv_softmax, attention rows of 16384", 0, null); + set.tokens = case.make(rw.Tokens, rig, k, n ++ "tf_iv_tokens, posterior sample to condition rows", 0, null); + const c = "control: idol vae "; + set.ctl_conv_in = case.make(cv.ConvIn, rig, set.bad[0], c ++ "conv_in stages 2 x - 1 in half", 0, .shapes); + set.ctl_half_in = case.make(Same, rig, set.bad[1], c ++ "conv stages its input in half", 0, .shapes); + set.ctl_sym_pad = case.make(Down, rig, set.bad[2], c ++ "stride-2 conv pads above as padding=1 would", 0, .shapes); + set.ctl_half_sum = case.make(Same, rig, set.bad[3], c ++ "conv sums in half", 0, .sums); + set.ctl_norm = case.make(rw.Norm, rig, set.bad[4], c ++ "norm's SiLU in half", 1, .shapes); + set.ctl_part = case.make(rw.GroupPart, rig, set.bad[5], c ++ "group_part sums in half", 0, .sums); + set.ctl_finish = case.make(rw.GroupFinish, rig, set.bad[6], c ++ "group_finish sums in half", 0, .sums); + set.ctl_softmax = case.make(rw.Softmax, rig, set.bad[7], c ++ "softmax's denominator in half", 0, .sums); + set.ctl_bias = case.make(ByRow, rig, set.bad[8], c ++ "bias by row reads the column's", 0, .shapes); + set.ctl_axes = case.make(rw.Tokens, rig, set.bad[9], c ++ "tokens swaps a patch's axes", 0, .shapes); + set.ctl_clamp = case.make(rw.Tokens, rig, set.bad[10], c ++ "tokens clamps a NaN logvar (fminf, fmaxf)", 0, .nonfinite); + inline for (.{ "conv_in", "same", "residual", "down", "one" }) |f| try list.append(a, @field(set, f).kernel()); + for (&set.parts) |*x| try list.append(a, x.kernel()); + for (&set.finishes) |*x| try list.append(a, x.kernel()); + for (&set.norms) |*x| try list.append(a, x.kernel()); + inline for (.{ "by_col", "by_col_add", "by_row", "softmax", "tokens" }) |f| try list.append(a, @field(set, f).kernel()); + inline for (ctl_fields) |f| try list.append(a, @field(set, f).kernel()); + return set; + } + + const ctl_fields = .{ "ctl_conv_in", "ctl_half_in", "ctl_sym_pad", "ctl_half_sum", "ctl_norm", "ctl_part", "ctl_finish", "ctl_softmax", "ctl_bias", "ctl_axes", "ctl_clamp" }; + + pub fn deinit(set: *Set) void { + inline for (.{ "conv_in", "same", "residual", "down", "one", "by_col", "by_col_add", "by_row", "softmax", "tokens" } ++ ctl_fields) |f| @field(set, f).deinit(); + for (&set.parts) |*x| x.deinit(); + for (&set.finishes) |*x| x.deinit(); + for (&set.norms) |*x| x.deinit(); + set.ctl.unload(); + set.k.unload(); + set.gpa.destroy(set); + } +}; + +/// The image encoder's kernels' gates 1-4 and their controls alone. +pub fn run(gpu: check.Gpu) !void { + const a = gpu.gpa; + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var list: std.ArrayList(g.Kernel) = .empty; + defer list.deinit(a); + const set = try Set.init(gpu, s, a, &list); + defer set.deinit(); + var t: g.Tally = .{}; + try g.runAll(a, list.items, 0x1d02, &t); + std.debug.print("SUMMARY idol vae: {d} kernels, {d} pass gates 1-4, {d} fail; {d} controls: {d} red as planted, {d} missed\n", .{ t.kernels, t.passed, t.failed, t.controls, t.caught, t.missed }); + try check.expect(t.problems() == 0, "idol-vae-contract: {d} gate failures, {d} controls missed", .{ t.failed, t.missed }); +} diff --git a/zig/tests/cuda/contract_idol_vae_convs.zig b/zig/tests/cuda/contract_idol_vae_convs.zig new file mode 100644 index 0000000000..8874d9973e --- /dev/null +++ b/zig/tests/cuda/contract_idol_vae_convs.zig @@ -0,0 +1,275 @@ +//! tf_iv_conv_in and tf_iv_conv under the contract: Idol's image encoder's fp32 convolutions against float64. +const std = @import("std"); +const cuda = @import("cuda"); +const vae = @import("idol_vae_cuda"); +const r = @import("contract_idol_ref.zig"); + +const c = r.c; +const Range = r.Range; +const Sum = r.Sum; +const Fill = r.g.Fill; + +/// NaN, +Inf and -Inf at `x`'s first values, and +Inf with -Inf side by side `apart` further on (one row holding both). +fn plant(x: []f64, apart: usize) void { + r.plant(x, &.{ 0, 7, 14 }); + if (x.len > 3 * apart + 1) { + x[3 * apart] = r.inf; + x[3 * apart + 1] = -r.inf; + } +} + +// --- tf_iv_conv_in: conv3x3(2 x - 1) + b of the picture, 3 channels to 128 ------------------------------------------ + +pub const ConvIn = struct { + pub const Kernels = vae.Kernels; + pub const kinds = [_]r.Kind{ .f32, .f32, .f32 }; + pub const out: r.Kind = .f32; + pub const in_place: ?usize = null; + + /// The 27 taps' products in one FMA chain (2 x - 1 rounded once before each), then 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, _: u32) 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 accepts(s: c.Shape) bool { + return s.rows > 0 and s.k > 0 and s.n == 128; + } + + pub fn sizes(s: c.Shape) [3]usize { + return .{ s.rows * s.k * 3, 27 * 128, 128 }; + } + + pub fn count(s: c.Shape) usize { + return s.rows * s.k * 128; + } + + pub fn fill(rng: std.Random, mode: Fill, s: c.Shape, _: u32, v: *[3][]f64) void { + if (mode == .stress) { + // 2 x - 1 an odd integer (127 or 2041), channel 1 two more, channel 2 zero; weights 509, -509, 1 cancel + for (0..s.rows * s.k) |p| { + const x: f64 = if (p % 2 == 0) 2041 else 127; + v[0][p * 3] = (x + 1) / 2; + v[0][p * 3 + 1] = (x + 3) / 2; + v[0][p * 3 + 2] = 0.5; + } + for (0..27) |row| @memset(v[1][row * 128 ..][0..128], switch (row % 3) { + 0 => 509, + 1 => -509, + else => 1, + }); + @memset(v[2], 0); + return; + } + r.uniform(rng, v[0], 0, 1); + r.uniform(rng, v[1], -0.5, 0.5); + r.uniform(rng, v[2], -1, 1); + if (mode == .nonfinite) plant(v[0], s.k * 3); + } + + pub fn launch(k: *const Kernels, s: cuda.Stream, dev: [3]u64, o: u64, shape: c.Shape, _: u32) !void { + try vae.convIn(k, s, dev[0], dev[1], dev[2], o, shape.rows, shape.k); + } + + pub fn ref(a: std.mem.Allocator, shape: c.Shape, _: u32, v: [3][]const f64, sum: Sum, e: *r.Ends, want: []f64, mag: []f64) !void { + const buf = try a.alloc(f64, 3 * 27); + defer a.free(buf); + const terms = buf[0..27]; + const lo = buf[27..54]; + const hi = buf[54..]; + for (0..shape.rows * shape.k) |px| for (0..128) |o| { + const py: isize = @intCast(px / shape.k); + const pxx: isize = @intCast(px % 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; + if (iy < 0 or ix < 0 or iy >= shape.rows or ix >= shape.k) { + terms[at] = 0; + lo[at] = 0; + hi[at] = 0; + continue; + } + const x = v[0][(@as(usize, @intCast(iy)) * shape.k + @as(usize, @intCast(ix))) * 3 + ch]; + const w = v[1][at * 128 + o]; + const t = Range.of(2 * x - 1).fp32(0.5).mul(Range.of(w)); + terms[at] = (2 * x - 1) * w; + lo[at] = t.lo; + hi[at] = t.hi; + }; + const s = sum.total(terms, lo, hi, 27); + const b = v[2][o]; + e.put(want, mag, px * 128 + o, s.v + b, s.r.add(Range.of(b)).fp32(0.5)); + }; + } +}; + +// --- tf_iv_conv: same-size and stride-2 3x3s and 1x1s as one implicit GEMM, fp32 sums, + bias (+ a residual) --------- + +/// A form the encoder launches: taps (9 or 1), stride, padding above and left, a residual, and its served shapes. +pub const Form = struct { taps: usize, stride: usize = 1, pad: usize = 0, add: bool = false, served: []const c.Shape }; + +/// A 3x3's rows are output rows of 16 columns (8: a 128-pixel tile), a 1x1's pixels; fewer than the engine's. +pub const forms = struct { + pub const same: Form = .{ .taps = 9, .pad = 1, .served = &.{ .{ .rows = 9, .n = 128, .k = 128 }, .{ .rows = 9, .n = 256, .k = 128 }, .{ .rows = 9, .n = 256, .k = 256 }, .{ .rows = 9, .n = 512, .k = 256 }, .{ .rows = 9, .n = 512, .k = 512 }, .{ .rows = 9, .n = 64, .k = 512 } } }; + pub const residual: Form = .{ .taps = 9, .pad = 1, .add = true, .served = &.{ .{ .rows = 9, .n = 128, .k = 128 }, .{ .rows = 9, .n = 256, .k = 256 }, .{ .rows = 9, .n = 512, .k = 512 } } }; + pub const down: Form = .{ .taps = 9, .stride = 2, .served = &.{ .{ .rows = 9, .n = 128, .k = 128 }, .{ .rows = 9, .n = 256, .k = 256 }, .{ .rows = 9, .n = 512, .k = 512 } } }; + pub const one: Form = .{ .taps = 1, .served = &.{ .{ .rows = 1100, .n = 256, .k = 128 }, .{ .rows = 1100, .n = 512, .k = 256 }, .{ .rows = 1100, .n = 64, .k = 64 } } }; +}; + +const out_cols = 16; + +pub fn Conv(comptime f: Form) type { + return struct { + pub const Kernels = vae.Kernels; + pub const kinds = [_]r.Kind{ .f32, .f32, .f32, .f32 }; + pub const out: r.Kind = .f32; + pub const in_place: ?usize = null; + + /// Each output sums taps * cin products in one FMA chain, then the bias and the residual. + fn depth(s: c.Shape) usize { + return f.taps * s.k + 2; + } + + /// Rows: output rows (3x3) or pixels (1x1); n: output channels; k: input channels. + pub fn spec(name: []const u8, _: u32) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = depth, + .out = .f32, + .tiles = .{ if (f.taps == 9) .{ .size = vae.tile_m / out_cols } else .{ .size = vae.tile_m }, .{ .size = vae.tile_n, .step = 64, .min = 64 }, .{ .size = vae.tile_k, .step = vae.tile_k, .min = vae.tile_k } }, + .served = f.served, + }; + } + + fn cols() usize { + return if (f.taps == 9) out_cols else 1; + } + fn shapeOf(s: c.Shape) vae.ConvShape { + return .{ .rows = s.rows, .cols = cols(), .in_rows = s.rows * f.stride, .in_cols = cols() * f.stride, .cin = s.k, .cout = s.n, .ldw = vae.ldwOf(s.n), .taps = f.taps, .stride = f.stride, .pad = f.pad }; + } + + pub fn accepts(s: c.Shape) bool { + return s.rows > 0 and s.n > 0 and s.k > 0 and vae.convOk(shapeOf(s)); + } + + pub fn sizes(s: c.Shape) [4]usize { + const g = shapeOf(s); + return .{ g.in_rows * g.in_cols * s.k, f.taps * s.k * g.ldw, s.n, if (f.add) count(s) else 0 }; + } + + pub fn count(s: c.Shape) usize { + return s.rows * cols() * s.n; + } + + pub fn fill(rng: std.Random, mode: Fill, s: c.Shape, _: u32, v: *[4][]f64) void { + const ldw = vae.ldwOf(s.n); + if (mode == .stress) { + // the channels' second half mirroring the first, weights 1 then -1: partial sums reach ~2^18 and cancel + const half = s.k / 2; + for (0..v[0].len / s.k) |p| for (0..s.k) |ci| { + const at = if (ci < half) ci else ci - half; + const x: f64 = if (at % 2 == 0) 127 else 2041; + v[0][p * s.k + ci] = if (ci < half) x else x + @as(f64, if (at < 4) 1 else 0); + }; + for (0..f.taps * s.k) |k| for (0..ldw) |o| { + v[1][k * ldw + o] = if (o >= s.n) 0 else if (k % s.k < half) 1 else -1; + }; + @memset(v[2], 0); + r.integers(rng, v[3], 3); + return; + } + r.uniform(rng, v[0], -2, 2); + r.uniform(rng, v[1], -0.1, 0.1); + r.uniform(rng, v[2], -1, 1); + r.uniform(rng, v[3], -3, 3); + if (mode == .nonfinite) plant(v[0], s.k); + } + + pub fn launch(k: *const Kernels, s: cuda.Stream, dev: [4]u64, o: u64, shape: c.Shape, _: u32) !void { + try vae.conv(k, s, dev[0], dev[1], dev[2], if (f.add) dev[3] else null, o, shapeOf(shape)); + } + + /// Output (p, o)'s products in the kernel's order, tap then input channel; a tap outside the input is 0. + fn terms(shape: c.Shape, v: [4][]const f64, p: usize, o: usize, out_: []f64) void { + const g = shapeOf(shape); + const K = shape.k; + const oy: isize = @intCast(p / g.cols); + const ox: isize = @intCast(p % g.cols); + const st: isize = @intCast(f.stride); + const pd: isize = @intCast(f.pad); + for (0..f.taps) |tap| { + const iy = oy * st + @as(isize, @intCast(if (f.taps == 9) tap / 3 else 0)) - pd; + const ix = ox * st + @as(isize, @intCast(if (f.taps == 9) tap % 3 else 0)) - pd; + const inside = iy >= 0 and ix >= 0 and iy < g.in_rows and ix < g.in_cols; + for (0..K) |ci| { + out_[tap * K + ci] = if (inside) blk: { + const at = (@as(usize, @intCast(iy)) * g.in_cols + @as(usize, @intCast(ix))) * K + ci; + break :blk v[0][at] * v[1][(tap * K + ci) * g.ldw + o]; + } else 0; + } + } + } + + pub fn ref(a: std.mem.Allocator, shape: c.Shape, _: u32, v: [4][]const f64, sum: Sum, e: *r.Ends, want: []f64, mag: []f64) !void { + const K = f.taps * shape.k; + const buf = try a.alloc(f64, K); + defer a.free(buf); + for (0..shape.rows * cols()) |p| for (0..shape.n) |o| { + terms(shape, v, p, o, buf); + const s = sum.total(buf, buf, buf, K); + const i = p * shape.n + o; + const b = v[2][o]; + var val = s.v + b; + var vr = s.r.add(Range.of(b)).fp32(0.5); + if (f.add) { + val += v[3][i]; + vr = vr.add(Range.of(v[3][i])).fp32(0.5); + } + e.put(want, mag, i, val, vr); + }; + } + }; +} + +test "conv references: a delta kernel copies its input, the stride-2 3x3 pads right and below alone" { + const a = std.testing.allocator; + const C = Conv(forms.down); + const shape: c.Shape = .{ .rows = 2, .n = 64, .k = 8 }; + var vals: [4][]f64 = undefined; + for (&vals, C.sizes(shape)) |*v, n| { + v.* = try a.alloc(f64, n); + @memset(v.*, 0); + } + defer for (vals) |v| a.free(v); + for (vals[0], 0..) |*x, i| x.* = @floatFromInt(i + 1); + // tap (2, 2), channel 3 to output 5: output pixel (oy, ox) reads input (2 oy + 2, 2 ox + 2) + const ldw = vae.ldwOf(64); + vals[1][(8 * 8 + 3) * ldw + 5] = 1; + var ends: r.Ends = .{ .a = a }; + defer ends.deinit(); + try ends.size(C.count(shape)); + const want = try a.alloc(f64, C.count(shape)); + defer a.free(want); + const mag = try a.alloc(f64, C.count(shape)); + defer a.free(mag); + try C.ref(a, shape, 0, .{ vals[0], vals[1], vals[2], vals[3] }, .{ .range = true }, &ends, want, mag); + try std.testing.expectEqual(vals[0][(2 * 32 + 2) * 8 + 3], want[5]); + try std.testing.expectEqual(vals[0][(2 * 32 + 30) * 8 + 3], want[14 * 64 + 5]); + // past the right edge (input column 32) and below (input row 4): the padding's zeros + try std.testing.expectEqual(@as(f64, 0), want[15 * 64 + 5]); + for (want[16 * 64 ..]) |w| try std.testing.expectEqual(@as(f64, 0), w); + try std.testing.expect(C.accepts(shape) and !C.accepts(.{ .rows = 1, .n = 64, .k = 12 })); +} diff --git a/zig/tests/cuda/contract_idol_vae_rows.zig b/zig/tests/cuda/contract_idol_vae_rows.zig new file mode 100644 index 0000000000..eda38cfe21 --- /dev/null +++ b/zig/tests/cuda/contract_idol_vae_rows.zig @@ -0,0 +1,501 @@ +//! Idol's image encoder's row kernels under the contract: GroupNorm, attention's biases and softmax, the rows. +const std = @import("std"); +const cuda = @import("cuda"); +const vae = @import("idol_vae_cuda"); +const r = @import("contract_idol_ref.zig"); + +const c = r.c; +const Range = r.Range; +const Sum = r.Sum; +const Fill = r.g.Fill; + +fn exp(x: f64) f64 { + return @exp(x); +} + +fn sqrt(x: f64) f64 { + return @sqrt(@max(x, 0)); +} + +const why = "no sums: each output is a few fp32 steps, every one inside its interval"; + +fn noSums(_: c.Shape) usize { + return 0; +} + +// --- tf_iv_group_part: GroupNorm's chunk sums, of the values (variant 0) or their squared deviations (1) ------------- + +pub const GroupPart = struct { + pub const Kernels = vae.Kernels; + pub const kinds = [_]r.Kind{ .f32, .f32 }; + pub const out: r.Kind = .f32; + pub const in_place: ?usize = null; + + /// A thread's chunk of pixels in turn, then its group's channels (up to 16) in turn. + fn depth(s: c.Shape) usize { + return @min(s.k, vae.chunk) + s.n / 32; + } + + /// Rows unused (1); n: channels (32 groups); k: pixels, summed in chunks. + pub fn spec(name: []const u8, _: u32) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = depth, + .out = .f32, + .tiles = .{ .{}, .{ .size = 128, .step = 32, .min = 32 }, .{ .size = vae.chunk } }, + .served = &.{ .{ .rows = 1, .n = 512, .k = 16384 }, .{ .rows = 1, .n = 512, .k = 65536 }, .{ .rows = 1, .n = 256, .k = 262144 }, .{ .rows = 1, .n = 128, .k = 1048576 } }, + }; + } + + pub fn accepts(s: c.Shape) bool { + return s.n > 0 and s.n % 32 == 0 and s.n <= 512 and s.k > 0; + } + + pub fn sizes(s: c.Shape) [2]usize { + return .{ s.k * s.n, 64 }; + } + + pub fn count(s: c.Shape) usize { + return 32 * vae.chunks(s.k); + } + + pub fn fill(rng: std.Random, mode: Fill, s: c.Shape, centered: u32, v: *[2][]f64) void { + if (mode == .stress) { + // integers: plain sums reach ~2^15 a channel and cancel; deviations within 18 keep squares' sums under 2^24 + for (0..32) |grp| v[1][grp] = @floatFromInt(rng.intRangeAtMost(i32, -4, 4)); + const half = vae.chunk / 2; + for (0..s.k) |p| for (0..s.n) |ch| { + const grp = ch / (s.n / 32); + const in = p % vae.chunk; + v[0][p * s.n + ch] = if (centered == 1) + v[1][grp] + @as(f64, @floatFromInt(rng.intRangeAtMost(i32, -18, 18))) + else if (in < half) @as(f64, if (in % 2 == 0) 127 else 511) else -@as(f64, if ((in - half) % 2 == 0) 127 else 511) + @as(f64, if (in - half < 4) 1 else 0); + }; + } else { + r.uniform(rng, v[0], -3, 3); + r.uniform(rng, v[1][0..32], -0.5, 0.5); + if (mode == .nonfinite) r.plant(v[0], &.{ 3, s.n + 9, 5 * s.n + 2 }); + } + r.uniform(rng, v[1][32..], 0.5, 2); + } + + pub fn launch(k: *const Kernels, s: cuda.Stream, dev: [2]u64, o: u64, shape: c.Shape, centered: u32) !void { + try vae.groupPart(k, s, dev[0], dev[1], o, shape.k, shape.n, centered == 1); + } + + pub fn ref(a: std.mem.Allocator, shape: c.Shape, centered: u32, v: [2][]const f64, sum: Sum, e: *r.Ends, want: []f64, mag: []f64) !void { + const size = shape.n / 32; + const most = vae.chunk * size; + const buf = try a.alloc(f64, 3 * most); + defer a.free(buf); + const chunks = vae.chunks(shape.k); + for (0..32) |grp| for (0..chunks) |ck| { + var n: usize = 0; + const mean = v[1][grp]; + for (ck * vae.chunk..@min((ck + 1) * vae.chunk, shape.k)) |p| for (0..size) |i| { + const x = v[0][p * shape.n + grp * size + i]; + const t = if (centered == 1) Range.of(x).sub(Range.of(mean)).fp32(0.5).square().fp32(0.5) else Range.of(x); + buf[n] = if (centered == 1) (x - mean) * (x - mean) else x; + buf[most + n] = t.lo; + buf[2 * most + n] = t.hi; + n += 1; + }; + const s = sum.total(buf[0..n], buf[most..][0..n], buf[2 * most ..][0..n], depth(shape)); + e.put(want, mag, grp * chunks + ck, s.v, s.r); + }; + } +}; + +// --- tf_iv_group_finish: the 32 means (variant 0) or 32 rstd (1) from the chunk sums -------------------------------- + +pub const GroupFinish = struct { + pub const Kernels = vae.Kernels; + pub const kinds = [_]r.Kind{ .f32, .f32 }; + pub const out: r.Kind = .f32; + /// stats [64] starts as input 1: the kernel writes the means' half or the rstd's, the other half stays. + pub const in_place: ?usize = 1; + + /// A lane's share of the chunks in turn, then the warp's 5 shuffles. + fn depth(s: c.Shape) usize { + return (vae.chunks(s.k) + 31) / 32 + 5; + } + + /// Rows unused (1); n: channels; k: pixels (their chunk sums are the input). + pub fn spec(name: []const u8, centered: u32) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = depth, + .ops = if (centered == 1) 2 * r.div_ulps + r.sqrt_ulps + 1 else r.div_ulps, + .out = .f32, + .tiles = .{ .{}, .{}, .{ .size = 32 * vae.chunk } }, + .served = &.{ .{ .rows = 1, .n = 512, .k = 16384 }, .{ .rows = 1, .n = 256, .k = 262144 }, .{ .rows = 1, .n = 128, .k = 1048576 } }, + }; + } + + pub fn accepts(s: c.Shape) bool { + return s.n > 0 and s.n % 32 == 0 and s.n <= 512 and s.k > 0; + } + + pub fn sizes(s: c.Shape) [2]usize { + return .{ 32 * vae.chunks(s.k), 64 }; + } + + pub fn count(_: c.Shape) usize { + return 64; + } + + pub fn fill(rng: std.Random, mode: Fill, _: c.Shape, centered: u32, v: *[2][]f64) void { + if (mode == .stress) { + // integers whose sums (at most 4096 chunks) stay under 2^23: plain ones mirror, squared ones are positive + const chunks = v[0].len / 32; + for (0..32) |grp| for (0..chunks) |i| { + const at = if (i < chunks / 2) i else i - chunks / 2; + const x: f64 = @floatFromInt(1 + (at * 37 + grp) % 2000); + v[0][grp * chunks + i] = if (centered == 1) x else if (i < chunks / 2) x else -x + @as(f64, if (at < 3) 1 else 0); + }; + } else if (centered == 1) r.uniform(rng, v[0], 0, 1000) else r.uniform(rng, v[0], -500, 500); + r.uniform(rng, v[1], -1, 1); + if (mode == .nonfinite) r.plant(v[0], &.{ 1, v[0].len / 32 + 3, 2 * (v[0].len / 32) + 5 }); + } + + pub fn launch(k: *const Kernels, s: cuda.Stream, dev: [2]u64, o: u64, shape: c.Shape, centered: u32) !void { + try vae.groupFinish(k, s, dev[0], o, shape.k, shape.n, centered == 1, 1e-6); + } + + pub fn ref(_: std.mem.Allocator, shape: c.Shape, centered: u32, v: [2][]const f64, sum: Sum, e: *r.Ends, want: []f64, mag: []f64) !void { + const chunks = vae.chunks(shape.k); + const n = r.f(@as(f64, @floatFromInt(shape.k)) * @as(f64, @floatFromInt(shape.n / 32))); + const eps = r.f(1e-6); + for (0..64) |i| e.put(want, mag, i, v[1][i], Range.of(v[1][i])); + for (0..32) |g| { + const part = v[0][g * chunks ..][0..chunks]; + const s = sum.total(part, part, part, depth(shape)); + const mr = (Range{ .lo = s.r.lo / n, .hi = s.r.hi / n }).fp32(r.div_ulps); + if (centered == 0) { + e.put(want, mag, g, s.v / n, mr); + continue; + } + const root = mr.add(Range.of(eps)).fp32(0.5).up(sqrt).fp32(r.sqrt_ulps); + e.put(want, mag, 32 + g, 1 / @sqrt(s.v / n + eps), Range.of(1).div(root).fp32(r.div_ulps)); + } + } +}; + +// --- tf_iv_norm: GroupNorm by given statistics, then SiLU (variant 1) or not ----------------------------------------- + +pub const Norm = struct { + pub const Kernels = vae.Kernels; + pub const kinds = [_]r.Kind{ .f32, .f32, .f32, .f32 }; + pub const out: r.Kind = .f32; + pub const in_place: ?usize = null; + + /// Rows: pixels; n: channels (32 groups of n / 32). + pub fn spec(name: []const u8, silu: u32) c.Spec { + return .{ + .name = name, + .accum = .none, + .depth = noSums, + .ops = 4 + r.fn_ulps + r.div_ulps + 1, + .out = .f32, + .tiles = .{ .{ .size = 8 }, .{ .size = 128, .step = 32, .min = 32 }, .{} }, + .served = if (silu == 1) &.{ .{ .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 }}, + .sums = .{ .f16 = 0, .bf16 = 0, .why = why }, + }; + } + + pub fn accepts(s: c.Shape) bool { + return s.rows > 0 and s.n > 0 and s.n % 32 == 0; + } + + pub fn sizes(s: c.Shape) [4]usize { + return .{ s.rows * s.n, 64, s.n, s.n }; + } + + pub fn count(s: c.Shape) usize { + return s.rows * s.n; + } + + pub fn fill(rng: std.Random, mode: Fill, s: c.Shape, _: u32, v: *[4][]f64) void { + // values stay within +-60 after the norm, short of fp32 exp's overflow in the SiLU + r.uniform(rng, v[0], -4, 4); + r.uniform(rng, v[1][0..32], -1, 1); + r.uniform(rng, v[1][32..], 0.2, 3); + r.uniform(rng, v[2], -2, 2); + r.uniform(rng, v[3], -1, 1); + if (mode == .nonfinite) r.plant(v[0], &.{ 3, s.n + 7, 2 * s.n + 1 }); + } + + pub fn launch(k: *const Kernels, s: cuda.Stream, dev: [4]u64, o: u64, shape: c.Shape, silu: u32) !void { + try vae.norm(k, s, dev[0], dev[1], dev[2], dev[3], o, shape.rows, shape.n, silu == 1); + } + + pub fn ref(_: std.mem.Allocator, shape: c.Shape, silu: u32, v: [4][]const f64, _: Sum, e: *r.Ends, want: []f64, mag: []f64) !void { + const size = shape.n / 32; + for (v[0], 0..) |x, i| { + const ch = i % shape.n; + const grp = ch / size; + const m = v[1][grp]; + const rs = v[1][32 + grp]; + const g = v[2][ch]; + const b = v[3][ch]; + const y = (x - m) * rs * g + b; + var yr = Range.of(x).sub(Range.of(m)).fp32(0.5).mul(Range.of(rs)).fp32(0.5).mul(Range.of(g)).fp32(0.5).add(Range.of(b)).fp32(0.5); + var w = y; + if (silu == 1) { + w = y / (1 + @exp(-y)); + const ex = (Range{ .lo = @exp(-yr.hi), .hi = @exp(-yr.lo) }).fp32(r.fn_ulps); + yr = yr.div(ex.add(Range.of(1)).fp32(0.5)).fp32(r.div_ulps); + } + e.put(want, mag, i, w, yr); + } + } +}; + +// --- tf_iv_bias: an attention linear's bias by column or by row (v^T), and the residual ----------------------------- + +pub fn Bias(comptime by_row: bool, comptime add: bool) type { + return struct { + pub const Kernels = vae.Kernels; + pub const kinds = [_]r.Kind{ .f32, .f32, .f32 }; + pub const out: r.Kind = .f32; + pub const in_place: ?usize = null; + + /// Rows by n: q, k and the output's [16384, 512], v^T's [512, 16384] (fewer rows: rows only tile). + pub fn spec(name: []const u8, _: u32) c.Spec { + return .{ + .name = name, + .accum = .none, + .depth = noSums, + .ops = 2, + .out = .f32, + .tiles = .{ .{ .size = 1 }, .{ .size = 256 }, .{} }, + .served = if (by_row) &.{.{ .rows = 512, .n = 4100, .k = 1 }} else &.{.{ .rows = 4100, .n = 512, .k = 1 }}, + .sums = .{ .f16 = 0, .bf16 = 0, .why = why }, + }; + } + + pub fn accepts(s: c.Shape) bool { + return s.rows > 0 and s.n > 0; + } + + pub fn sizes(s: c.Shape) [3]usize { + return .{ s.rows * s.n, if (by_row) s.rows else s.n, if (add) s.rows * s.n else 0 }; + } + + pub fn count(s: c.Shape) usize { + return s.rows * s.n; + } + + pub fn fill(rng: std.Random, mode: Fill, _: c.Shape, _: u32, v: *[3][]f64) void { + r.uniform(rng, v[0], -6, 6); + r.uniform(rng, v[1], -1, 1); + r.uniform(rng, v[2], -3, 3); + if (mode == .nonfinite) for (v, 0..) |col, i| if (col.len > 0) r.plant(col, &.{ 11 + 7 * i, 2048 + i }); + } + + pub fn launch(k: *const Kernels, s: cuda.Stream, dev: [3]u64, o: u64, shape: c.Shape, _: u32) !void { + try vae.bias(k, s, dev[0], dev[1], if (add) dev[2] else null, o, shape.rows, shape.n, by_row); + } + + pub fn ref(_: std.mem.Allocator, shape: c.Shape, _: u32, v: [3][]const f64, _: Sum, e: *r.Ends, want: []f64, mag: []f64) !void { + for (0..shape.rows * shape.n) |i| { + const b = v[1][if (by_row) i / shape.n else i % shape.n]; + var y = v[0][i] + b; + var yr = Range.of(v[0][i]).add(Range.of(b)).fp32(0.5); + if (add) { + y += v[2][i]; + yr = yr.add(Range.of(v[2][i])).fp32(0.5); + } + e.put(want, mag, i, y, yr); + } + } + }; +} + +// --- tf_iv_softmax: the attention's softmax(s / sqrt(512)) of a row of scores, in place ----------------------------- + +pub const Softmax = struct { + pub const Kernels = vae.Kernels; + pub const kinds = [_]r.Kind{.f32}; + pub const out: r.Kind = .f32; + pub const in_place: ?usize = 0; + const scale: f32 = 1.0 / @sqrt(@as(f32, 512)); + + /// A thread's share of the row in turn, the warp's 5 shuffles, then the 8 warps' partials in order. + fn depth(s: c.Shape) usize { + return (s.n + 255) / 256 + 5 + 8; + } + + /// Rows: query rows (the engine's chunk is 4,096; rows only tile); n = k: the keys, 16,384. + pub fn spec(name: []const u8, _: u32) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = depth, + .ops = 2 + r.fn_ulps + r.div_ulps, + .out = .f32, + .tiles = .{ .{}, .{ .size = 256 }, .{} }, + .served = &.{ .{ .rows = 6, .n = 16384, .k = 16384 }, .{ .rows = 6, .n = 4096, .k = 4096 } }, + }; + } + + pub fn accepts(s: c.Shape) bool { + return s.rows > 0 and s.n > 0 and s.n <= vae.softmax_max; + } + + pub fn sizes(s: c.Shape) [1]usize { + return .{s.rows * s.n}; + } + + pub fn count(s: c.Shape) usize { + return s.rows * s.n; + } + + pub fn fill(rng: std.Random, mode: Fill, s: c.Shape, _: u32, v: *[1][]f64) void { + // gate 4's rows are equal scores: every exp is 1 and the denominator n, exact in fp32 in any order + if (mode == .stress) return @memset(v[0], 3.0); + r.uniform(rng, v[0], -200, 200); + if (mode == .nonfinite) r.plant(v[0], &.{ 3, s.n + 5, 2 * s.n + 7 }); + } + + pub fn launch(k: *const Kernels, s: cuda.Stream, _: [1]u64, o: u64, shape: c.Shape, _: u32) !void { + try vae.softmax(k, s, o, shape.rows, shape.n, scale); + } + + pub fn ref(a: std.mem.Allocator, shape: c.Shape, _: u32, v: [1][]const f64, sum: Sum, e: *r.Ends, want: []f64, mag: []f64) !void { + const n = shape.n; + const buf = try a.alloc(f64, 3 * n); + defer a.free(buf); + const ex = buf[0..n]; + const lo = buf[n .. 2 * n]; + const hi = buf[2 * n ..]; + for (0..shape.rows) |row| { + const x = v[0][row * n ..][0..n]; + // s scale and its difference from the row's maximum are single correctly rounded fp32 steps + var m = -r.inf; + for (x) |s| m = r.fmax(m, r.f(s * scale)); + for (x, ex, lo, hi) |s, *t, *l, *u| { + const d = r.f(r.f(s * scale) - m); + t.* = @exp(d); + const er = Range.of(d).up(exp).fp32(r.fn_ulps); + l.* = er.lo; + u.* = er.hi; + } + const total = sum.total(ex, lo, hi, depth(shape)); + for (0..n) |i| { + const er: Range = .{ .lo = lo[i], .hi = hi[i] }; + e.put(want, mag, row * n + i, ex[i] / total.v, er.div(total.r).fp32(r.div_ulps)); + } + } + } +}; + +// --- tf_iv_tokens: the posterior's sample, patchified and normalized by the batch norm's statistics ----------------- + +pub const Tokens = struct { + pub const Kernels = vae.Kernels; + pub const kinds = [_]r.Kind{ .f32, .f32, .f32, .f32 }; + pub const out: r.Kind = .f32; + pub const in_place: ?usize = null; + 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, _: u32) c.Spec { + return .{ + .name = name, + .accum = .none, + .depth = noSums, + .ops = r.fn_ulps + r.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 } }, + .sums = .{ .f16 = 0, .bf16 = 0, .why = why }, + }; + } + + pub fn accepts(s: c.Shape) bool { + return s.n == 128 and s.rows > 0 and s.k > 0 and s.rows % 2 == 0 and s.k % 2 == 0; + } + + pub fn sizes(s: c.Shape) [4]usize { + return .{ s.rows * s.k * 64, s.rows * s.k * 32, 128, 128 }; + } + + pub fn count(s: c.Shape) usize { + return (pad + s.rows * s.k / 4) * 128; + } + + pub fn fill(rng: std.Random, mode: Fill, _: c.Shape, _: u32, v: *[4][]f64) void { + for (0..v[0].len / 64) |p| { + r.uniform(rng, v[0][p * 64 ..][0..32], -3, 3); + // logvars past both of the clamp's ends (-30, 20) as well as inside them + r.uniform(rng, v[0][p * 64 + 32 ..][0..32], -40, 30); + } + r.uniform(rng, v[1], -3, 3); + r.uniform(rng, v[2], -1, 1); + r.uniform(rng, v[3], 0.5, 4); + if (mode == .nonfinite) { + r.plant(v[0], &.{ 0, 64 + 1, 2 * 64 + 2 }); + r.plant(v[1], &.{ 7, 32 + 8, 2 * 32 + 9 }); + v[0][5 * 64 + 40] = r.nan; + } + } + + pub fn launch(k: *const Kernels, s: cuda.Stream, dev: [4]u64, o: u64, shape: c.Shape, _: u32) !void { + try vae.tokens(k, s, dev[0], dev[1], dev[2], dev[3], o, shape.rows, shape.k, pad); + } + + /// 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 ref(_: std.mem.Allocator, shape: c.Shape, _: u32, v: [4][]const f64, _: Sum, e: *r.Ends, want: []f64, mag: []f64) !void { + for (0..pad * 128) |i| e.put(want, mag, i, 0, Range.of(0)); + const w2 = shape.k / 2; + for (0..shape.rows * shape.k / 4) |patch| for (0..128) |k| { + const ch = k / 4; + const p = ((patch / w2) * 2 + (k % 4) / 2) * shape.k + (patch % w2) * 2 + k % 2; + const mean = v[0][p * 64 + ch]; + const lv = clamp(v[0][p * 64 + 32 + ch]); + const nz = v[1][p * 32 + ch]; + const sd = @exp(0.5 * lv); + const val = (mean + sd * nz - v[2][k]) / v[3][k]; + const er = Range.of(0.5 * lv).up(exp).fp32(r.fn_ulps); + const zr = er.mul(Range.of(nz)).fp32(0.5).add(Range.of(mean)).fp32(0.5); + const vr = zr.sub(Range.of(v[2][k])).fp32(0.5).div(Range.of(v[3][k])).fp32(r.div_ulps); + e.put(want, mag, (pad + patch) * 128 + k, val, vr); + }; + } +}; + +test "the softmax reference's ranges hold its float64 values, and gate 4's equal rows give 1 / n exactly" { + const a = std.testing.allocator; + const shape: c.Shape = .{ .rows = 2, .n = 300, .k = 300 }; + const x = try a.alloc(f64, 600); + defer a.free(x); + var prng = std.Random.DefaultPrng.init(5); + var vals = [1][]f64{x}; + Softmax.fill(prng.random(), .random, shape, 0, &vals); + var ends: r.Ends = .{ .a = a }; + defer ends.deinit(); + try ends.size(600); + const want = try a.alloc(f64, 600); + defer a.free(want); + const mag = try a.alloc(f64, 600); + defer a.free(mag); + try Softmax.ref(a, shape, 0, .{x}, .{ .range = false }, &ends, want, mag); + for (want, ends.lo, ends.hi) |w, lo, hi| try std.testing.expect(lo <= w and w <= hi); + Softmax.fill(prng.random(), .stress, shape, 0, &vals); + try Softmax.ref(a, shape, 0, .{x}, .{ .range = true }, &ends, want, mag); + for (want) |w| try std.testing.expectEqual(@as(f64, 1.0 / 300.0), w); + // fp16 sums of 4096 ones stall at 2048: gate 4's margin + const ones = try a.alloc(f64, 3 * 4096); + defer a.free(ones); + @memset(ones[0..4096], 3.0); + try ends.size(4096); + try Softmax.ref(a, .{ .rows = 1, .n = 4096, .k = 4096 }, 0, .{ones[0..4096]}, .{ .sim = .{ .acc = .f16, .order = .sequential } }, &ends, ones[4096..8192], ones[8192..]); + try std.testing.expectEqual(@as(f64, 1.0 / 2048.0), ones[4096]); +} diff --git a/zig/tests/cuda/contract_limits.zig b/zig/tests/cuda/contract_limits.zig index 419cc615ef..e0f64620e9 100644 --- a/zig/tests/cuda/contract_limits.zig +++ b/zig/tests/cuda/contract_limits.zig @@ -26,7 +26,7 @@ pub fn image(name: []const u8) error{UnknownContractImage}![]const u8 { if (std.mem.eql(u8, name, "h3")) return &h3_image.bytes; inline for (.{ "qwen4_projection", "qwen4_hc", "qwen4_gdn", "qwen4_ple" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); - inline for (.{ "tree_round", "act_rows", "gdn_prep", "qk_rows", "norms", "moe", "experts", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "qmm_prefill", "lane_gemv", "prompt16", "qmmf", "ck_experts", "nvfp4_experts", "nvfp4_experts_tile", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "glm_glue", "glm_ops", "glm_index", "nemotron_mamba", "sample", "qwen35", "qwen35_bf16_rows", "qwen35_keep", "qwen35_kv8", "w4a8_prefill", "qwen35_draft", "ck_lane", "gdn_train", "train_q4", "train_attn", "flashnext_train", "train", "attn_train", "qwen27_train", "tree_attention", "block_attention", "window_attention", "window_pages", "moe_glue", "dense_rows", "qmm", "flashnext_fact", "flashnext_hcf", "flashnext_fexp", "experts_rows", "nemotron_ops", "minimax_m3", "minimax_m3_attn", "moss_image", "idol" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); + inline for (.{ "tree_round", "act_rows", "gdn_prep", "qk_rows", "norms", "moe", "experts", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "qmm_prefill", "lane_gemv", "prompt16", "qmmf", "ck_experts", "nvfp4_experts", "nvfp4_experts_tile", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "glm_glue", "glm_ops", "glm_index", "nemotron_mamba", "sample", "qwen35", "qwen35_bf16_rows", "qwen35_keep", "qwen35_kv8", "w4a8_prefill", "qwen35_draft", "ck_lane", "gdn_train", "train_q4", "train_attn", "flashnext_train", "train", "attn_train", "qwen27_train", "tree_attention", "block_attention", "window_attention", "window_pages", "moe_glue", "dense_rows", "qmm", "flashnext_fact", "flashnext_hcf", "flashnext_fexp", "experts_rows", "nemotron_ops", "minimax_m3", "minimax_m3_attn", "moss_image", "idol", "idol_vae" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); if (std.mem.eql(u8, name, "contract_controls")) return if (options.controls) &controls.bytes else &.{}; std.debug.print("FAIL kernel-contract: unknown image name: {s}\n", .{name}); return error.UnknownContractImage; diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index e4c0d80f93..c2e2c6e4d9 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -238,6 +238,7 @@ const usage = \\ minimax-contract MiniMax-M3.1-Flash's indexer and block-sparse attention through gates 1-4, controls red \\ moss-contract MOSS-Image's own kernels and its full attention through gates 1-4, controls red \\ idol-contract Idol's denoiser glue and head-64 attention through gates 1-4, controls red + \\ idol-vae-contract Idol's fp32 image encoder kernels through gates 1-4, controls red \\ glm-exl3 GLM's EXL3 routed experts, decode and prompt, against float64 (synthetic) \\ glm-vision MODEL F32 BF16 IMG GLM's image tower and preparation against vision_ref.py's torch runs \\ glm-mm-fixed Prompt K policies: FP64, every row alone/batched, offsets and guarded tails @@ -460,6 +461,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "qwen-pages")) return @import("contract_qwen_pages.zig").run(gpu); if (std.mem.eql(u8, cmd, "minimax-contract")) return @import("contract_minimax.zig").run(gpu); if (std.mem.eql(u8, cmd, "idol-contract")) return @import("contract_idol.zig").run(gpu); + if (std.mem.eql(u8, cmd, "idol-vae-contract")) return @import("contract_idol_vae.zig").run(gpu); if (std.mem.eql(u8, cmd, "kernel-contract-limits")) return contract_tests.fatbins(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "pocket-contract")) return @import("contract_pocket.zig").run(gpu); if (std.mem.eql(u8, cmd, "moss-contract")) return @import("contract_moss_image.zig").run(gpu); diff --git a/zig/tests/idol_vae.zig b/zig/tests/idol_vae.zig index 3a0e38b5d5..b6cd72f873 100644 --- a/zig/tests/idol_vae.zig +++ b/zig/tests/idol_vae.zig @@ -1,132 +1,38 @@ -//! tf-idol-vae: Idol's image encoder on a case from tools/zig/idol_case.py's vae, gated against float64 and timed. +//! tf-idol-vae on Metal: idol_vae_run.zig's case driver on the system GPU. 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. -; +const Metal = struct { + pub const vae = tf.idol.vae; + pub const ref = tf.idol.reference; + pub const st = tf.safetensors; + pub const seconds = mtl.clock.seconds; -/// Encodes timed after the gated one. -const timed = 5; + pub const Device = struct { + pool: mtl.objc.Pool, + device: mtl.Device, -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); + pub fn init(g: *Device) !void { + g.pool = mtl.objc.Pool.push(); + errdefer g.pool.pop(); + g.device = try mtl.Device.init(); + } - 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 }); + pub fn deinit(g: *Device) void { + g.pool.pop(); + } - 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); + pub fn name(g: *Device) []const u8 { + return std.mem.span(g.device.name()); } - 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}); - } + pub fn load(g: *Device, gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !*vae.Encoder { + return vae.Encoder.load(gpa, io, g.device, dir); + } + }; +}; - 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; +pub fn main(init: std.process.Init) !u8 { + return @import("idol_vae_run.zig").run(Metal, init); } diff --git a/zig/tests/idol_vae_cuda.zig b/zig/tests/idol_vae_cuda.zig new file mode 100644 index 0000000000..c733698f33 --- /dev/null +++ b/zig/tests/idol_vae_cuda.zig @@ -0,0 +1,41 @@ +//! tf-idol-vae on CUDA: idol_vae_run.zig's case driver on the first GPU (sm_120 or sm_121; others are refused by name). +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); + +const Cuda = struct { + const id = @import("idol"); + pub const vae = id.vae; + pub const ref = id.reference; + pub const st = core.safetensors; + pub const seconds = id.seconds; + + pub const Device = struct { + driver: cuda.Driver, + ctx: cuda.Context, + buf: [256]u8, + + pub fn init(g: *Device) !void { + g.driver = try cuda.Driver.open(); + errdefer g.driver.close(); + g.ctx = try cuda.Context.init(&g.driver, 0); + } + + pub fn deinit(g: *Device) void { + g.ctx.deinit(); + g.driver.close(); + } + + pub fn name(g: *Device) []const u8 { + return g.ctx.name(&g.buf) catch "a CUDA GPU"; + } + + pub fn load(g: *Device, gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !*vae.Encoder { + return vae.Encoder.load(gpa, io, &g.ctx, dir); + } + }; +}; + +pub fn main(init: std.process.Init) !u8 { + return @import("idol_vae_run.zig").run(Cuda, init); +} diff --git a/zig/tests/idol_vae_run.zig b/zig/tests/idol_vae_run.zig new file mode 100644 index 0000000000..16b1c74f57 --- /dev/null +++ b/zig/tests/idol_vae_run.zig @@ -0,0 +1,130 @@ +//! tf-idol-vae's case driver for a backend `B`: its encoder `vae`, `ref`, safetensors `st`, `seconds` and `Device`. +const std = @import("std"); + +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: anytype, 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: anytype) !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 run(comptime B: type, init: std.process.Init) !u8 { + const vae = B.vae; + const ref = B.ref; + 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 B.st.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: ?B.st.File = B.st.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); + + var gpu: B.Device = undefined; + try gpu.init(); + defer gpu.deinit(); + const loading = B.seconds(); + const enc = try gpu.load(std.heap.page_allocator, io, args[1]); + defer enc.deinit(); + const load_s = B.seconds() - loading; + std.debug.print("{s}: encoder loaded in {d:.2} s\n", .{ gpu.name(), load_s }); + var out: std.ArrayList(u8) = .empty; + try print(&out, a, "{{\"device\": \"{s}\", \"load_s\": {d:.3}", .{ gpu.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 B.st.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; +}