From f9651de3032d8bc840b5898f388f549d86d77f61 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Tue, 27 Oct 2026 08:07:21 -0500 Subject: [PATCH 1/3] MOSS-Image: the prompt encoder on Metal encoder.zig runs the checkpoint's text_encoder (Qwen3.5's text model: Gated DeltaNet and gated full attention, hidden 4096) as the pipeline's encode_prompt does: the prompt in its chat format, one causal forward, the final norm's rows. bf16 weights as published; fp32 sums, DeltaNet state and attention; the other steps rounded where the torch code rounds. tf-moss-image --prompt encodes both prompts natively, tf-moss-image-encoder checks a case from moss_image_case.py encode, and start --no-encoder makes a case without torch's encoder. The tokenizer stays outside the engine module: the test programs get it. Nine kernels in moss_image_encoder.metal, each through the contract's gates 1-5 with eight controls. --- tools/zig/moss_image_case.py | 118 +++- zig/build/metal_targets.zig | 5 +- zig/kernels/metal/moss_image_encoder.metal | 225 +++++++ zig/kernels/metal/sources.zig | 4 +- zig/src/core/contract/moss_image_records.zig | 11 +- zig/src/families/moss_image/encoder.zig | 221 +++++++ .../families/moss_image/encoder_config.zig | 112 ++++ .../families/moss_image/encoder_weights.zig | 115 ++++ zig/src/families/moss_image/kernels.zig | 58 ++ zig/src/families/moss_image/moss_image.zig | 7 + zig/tests/contract/metal_moss_image.zig | 70 ++- .../contract/metal_moss_image_common.zig | 22 + .../contract/metal_moss_image_encoder.zig | 496 +++++++++++++++ .../metal_moss_image_encoder_state.zig | 581 ++++++++++++++++++ zig/tests/contract/metal_moss_image_norms.zig | 44 +- zig/tests/moss_image.zig | 47 +- zig/tests/moss_image_encoder.zig | 115 ++++ 17 files changed, 2207 insertions(+), 44 deletions(-) create mode 100644 zig/kernels/metal/moss_image_encoder.metal create mode 100644 zig/src/families/moss_image/encoder.zig create mode 100644 zig/src/families/moss_image/encoder_config.zig create mode 100644 zig/src/families/moss_image/encoder_weights.zig create mode 100644 zig/tests/contract/metal_moss_image_encoder.zig create mode 100644 zig/tests/contract/metal_moss_image_encoder_state.zig create mode 100644 zig/tests/moss_image_encoder.zig diff --git a/tools/zig/moss_image_case.py b/tools/zig/moss_image_case.py index 30f3df5217..1f3d0abfd6 100644 --- a/tools/zig/moss_image_case.py +++ b/tools/zig/moss_image_case.py @@ -14,9 +14,14 @@ USAGE = """\ case MODEL_DIR OUT_DIR --prompt-file F prompt rows, timesteps, start latents, each step's inputs and both velocities - start MODEL_DIR OUT_DIR --prompt-file F only what a picture needs: both prompts' rows, the schedule, the noise - exact MODEL_DIR CASE_DIR the case's schedule with a float32 transformer: CASE_DIR/exact.png + start MODEL_DIR OUT_DIR --prompt-file F only what a picture needs: both prompts' rows, the schedule, the noise; + with --no-encoder, the schedule and the noise alone + exact MODEL_DIR CASE_DIR the case's schedule with a float32 transformer: CASE_DIR/exact.png; + with --fp32-encoder F, F's rows from a float32 encoder too + (CASE_DIR/exact_fp32_encoder.png) decode MODEL_DIR LATENTS.f32 PNG raw float32 latents (rows x columns x 128) from tf-moss-image to a picture + encode MODEL_DIR OUT_DIR [--prompt-file F] the prompt encoder alone: token ids, embedding rows, every layer's + output and the final rows in bf16, and the final rows in float32 compare A.png B.png [SIDE.png] PSNR and mean absolute difference of two pictures, optionally side by side MODEL_DIR is the OpenMOSS-Team/MOSS-Image-Base folder. The case runs the card's defaults: 50 Euler steps, CFG 3, an empty negative prompt whose branch skips block 22, the bf16 transformer and text encoder, the float32 decoder. @@ -157,24 +162,29 @@ def recorded(hidden_states, timestep, encoder_hidden_states, encoder_attention_m def start(args) -> int: - """What tf-moss-image needs to make a picture: both prompts' encoder rows, the schedule and the seeded noise.""" + """What tf-moss-image needs for a picture: both prompts' rows (unless --no-encoder), the schedule, the noise.""" import torch from safetensors.torch import save_file out = Path(args.out_dir) out.mkdir(parents=True, exist_ok=True) - pipe, _ = pipeline(args.model_dir, args.device, dit=False) - positive, _ = pipe.encode_prompt(Path(args.prompt_file).read_text().strip()) - negative, _ = pipe.encode_prompt(args.negative) + if not args.no_encoder and not args.prompt_file: + sys.exit("start needs --prompt-file unless --no-encoder") + pipe, _ = pipeline(args.model_dir, args.device, text=not args.no_encoder, dit=False) h, w = args.height // 16, args.width // 16 noise = torch.randn((1, 128, h, w), generator=torch.Generator("cpu").manual_seed(args.seed), dtype=torch.float32) pipe.scheduler.set_timesteps(args.steps) - save_file({"prompt": positive[0].cpu().contiguous(), "negative": negative[0].cpu().contiguous(), - "timesteps": pipe.scheduler.timesteps.float().cpu().contiguous(), + tensors = {"timesteps": pipe.scheduler.timesteps.float().cpu().contiguous(), "sigmas": pipe.scheduler.sigmas.float().cpu().contiguous(), - "latents": noise[0].permute(1, 2, 0).contiguous()}, str(out / "case.safetensors")) - print(json.dumps({"width": args.width, "height": args.height, "seed": args.seed, "steps": args.steps, - "prompt_rows": int(positive.shape[1]), "negative_rows": int(negative.shape[1])})) + "latents": noise[0].permute(1, 2, 0).contiguous()} + meta = {"width": args.width, "height": args.height, "seed": args.seed, "steps": args.steps} + if not args.no_encoder: + positive, _ = pipe.encode_prompt(Path(args.prompt_file).read_text().strip()) + negative, _ = pipe.encode_prompt(args.negative) + tensors.update(prompt=positive[0].cpu().contiguous(), negative=negative[0].cpu().contiguous()) + meta.update(prompt_rows=int(positive.shape[1]), negative_rows=int(negative.shape[1])) + save_file(tensors, str(out / "case.safetensors")) + print(json.dumps(meta)) return 0 @@ -185,15 +195,21 @@ def exact(args) -> int: case_dir = Path(args.case_dir) case = load_file(str(case_dir / "case.safetensors")) + rows = {"prompt": case["prompt"], "negative": case["negative"]} + if args.fp32_encoder: + # the prompt encoder in float32 too: both prompts' rows as exact arithmetic gives them + prompt = Path(args.fp32_encoder).read_text().strip() + rows = {key: encoder_rows(args.model_dir, text, torch.float32, args.device) for key, text in (("prompt", prompt), ("negative", ""))} pipe, _ = pipeline(args.model_dir, args.device, text=False) pipe.transformer.float() h, w, _ = case["latents"].shape - image = pipe(prompt_embeds=case["prompt"][None].float().to(args.device), - negative_prompt_embeds=case["negative"][None].float().to(args.device), height=16 * h, width=16 * w, + image = pipe(prompt_embeds=rows["prompt"][None].float().to(args.device), + negative_prompt_embeds=rows["negative"][None].float().to(args.device), height=16 * h, width=16 * w, num_inference_steps=args.steps, latents=case["latents"].permute(2, 0, 1)[None].to(args.device), output_type="latent").images - image.float().cpu().permute(0, 2, 3, 1).contiguous().numpy().tofile(case_dir / "exact.latents.f32") - to_png(pipe, image, str(case_dir / "exact.png")) + name = "exact_fp32_encoder" if args.fp32_encoder else "exact" + image.float().cpu().permute(0, 2, 3, 1).contiguous().numpy().tofile(case_dir / f"{name}.latents.f32") + to_png(pipe, image, str(case_dir / f"{name}.png")) return 0 @@ -226,6 +242,67 @@ def compare(args) -> int: return 0 +def encoder_rows(model_dir: str, prompt: str, dtype, device: str): + """One prompt's final rows from the text encoder in `dtype`, formatted and tokenized as the pipeline does.""" + import torch + from transformers import AutoModel, AutoTokenizer + + root = Path(model_dir) + tokenizer = AutoTokenizer.from_pretrained(root / "tokenizer") + formatted = tokenizer.apply_chat_template([{"role": "user", "content": prompt}], tokenize=False, + add_generation_prompt=True, enable_thinking=True) + inputs = tokenizer(formatted, return_tensors="pt", truncation=True, max_length=4096).to(device) + encoder = AutoModel.from_pretrained(root / "text_encoder", dtype=dtype).eval().to(device) + with torch.no_grad(): + rows = encoder(**inputs, use_cache=False, return_dict=True).last_hidden_state[0].cpu() + del encoder + if device == "mps": + torch.mps.empty_cache() + return rows + + +def encode(args) -> int: + """The prompt encoder as the pipeline runs it, every layer's rows kept; then the same prompt in float32.""" + import torch + from safetensors.torch import save_file + from transformers import AutoModel, AutoTokenizer + + root = Path(args.model_dir) + out = Path(args.out_dir) + out.mkdir(parents=True, exist_ok=True) + tokenizer = AutoTokenizer.from_pretrained(root / "tokenizer") + prompt = Path(args.prompt_file).read_text().strip() if args.prompt_file else "" + # as MOSSImagePipeline.encode_prompt formats and tokenizes it + formatted = tokenizer.apply_chat_template([{"role": "user", "content": prompt}], tokenize=False, + add_generation_prompt=True, enable_thinking=True) + inputs = tokenizer(formatted, return_tensors="pt", truncation=True, max_length=4096) + tensors = {"ids": inputs["input_ids"][0].to(torch.int32).contiguous()} + for dtype in (torch.bfloat16, torch.float32): + encoder = AutoModel.from_pretrained(root / "text_encoder", dtype=dtype).eval().to(args.device) + text = encoder.language_model + kept = [] + hooks = [text.embed_tokens.register_forward_hook(lambda mod, a, o: kept.append(o))] + hooks += [layer.register_forward_hook(lambda mod, a, o: kept.append(o[0] if isinstance(o, tuple) else o)) + for layer in text.layers] + with torch.no_grad(): + final = encoder(**inputs.to(args.device), use_cache=False, return_dict=True).last_hidden_state[0] + for h in hooks: + h.remove() + if dtype is torch.bfloat16: + tensors["embed"] = kept[0][0].cpu().contiguous() + for i, rows in enumerate(kept[1:]): + tensors[f"layer.{i}"] = rows[0].cpu().contiguous() + tensors["final"] = final.cpu().contiguous() + else: + tensors["final_fp32"] = final.cpu().contiguous() + del encoder, text, kept + if args.device == "mps": + torch.mps.empty_cache() + save_file(tensors, str(out / "encoder.safetensors")) + print(json.dumps({"rows": int(tensors["ids"].numel()), "prompt": prompt[:60]})) + return 0 + + def main() -> int: parser = argparse.ArgumentParser(description=__doc__, epilog=USAGE, formatter_class=argparse.RawDescriptionHelpFormatter) sub = parser.add_subparsers(dest="mode", required=True) @@ -242,7 +319,8 @@ def main() -> int: s = sub.add_parser("start") s.add_argument("model_dir") s.add_argument("out_dir") - s.add_argument("--prompt-file", required=True) + s.add_argument("--prompt-file") + s.add_argument("--no-encoder", action="store_true", help="no prompt rows: tf-moss-image --prompt encodes them") s.add_argument("--negative", default="", help="the negative prompt (the card's default: empty)") s.add_argument("--width", type=int, default=1024) s.add_argument("--height", type=int, default=1024) @@ -253,6 +331,7 @@ def main() -> int: x.add_argument("model_dir") x.add_argument("case_dir") x.add_argument("--steps", type=int, default=50) + x.add_argument("--fp32-encoder", metavar="PROMPT_FILE", help="the prompt encoder in float32 too, on this prompt") x.add_argument("--device", default="mps") d = sub.add_parser("decode") d.add_argument("model_dir") @@ -261,12 +340,17 @@ def main() -> int: d.add_argument("--width", type=int, default=1024) d.add_argument("--height", type=int, default=1024) d.add_argument("--device", default="mps") + n = sub.add_parser("encode") + n.add_argument("model_dir") + n.add_argument("out_dir") + n.add_argument("--prompt-file", help="the prompt (default: the empty negative prompt)") + n.add_argument("--device", default="mps") p = sub.add_parser("compare") p.add_argument("a") p.add_argument("b") p.add_argument("side", nargs="?") args = parser.parse_args() - return {"case": case, "start": start, "exact": exact, "decode": decode, "compare": compare}[args.mode](args) + return {"case": case, "start": start, "exact": exact, "decode": decode, "encode": encode, "compare": compare}[args.mode](args) if __name__ == "__main__": diff --git a/zig/build/metal_targets.zig b/zig/build/metal_targets.zig index 8a16d7432e..69684d3b8c 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-moss-image-encoder", .path = "zig/tests/moss_image_encoder.zig", .about = "MOSS-Image's prompt encoder on a reference case: the ids checked, each layer's rows compared, 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" }, @@ -331,13 +332,15 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s .{ .name = "tf-spark-run", .path = "zig/tests/spark25/run.zig", .about = "Spark2.5 on Metal: greedy tokens, every position's logits or a prompt's last logits, for the reference's check" }, .{ .name = "tf-qwen36-row-seam", .path = "zig/tests/qwen36_row_seam.zig", .about = "Model-free fixture: the row projection's bias sums stay fp32 at every width" }, }; + // the tokenizer for the programs that tokenize prompts themselves (the engine module has none) + const tokenizer = b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = target, .optimize = optimize, .link_libc = true }); for (engine_programs) |p| { const mod = b.createModule(.{ .root_source_file = b.path(p.path), .target = target, .optimize = optimize, .link_libc = true, - .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "tensorfold", .module = engine }, .{ .name = "core", .module = engine }, .{ .name = "cluster", .module = cluster }, .{ .name = "kernel_sources", .module = sources }, .{ .name = "native_engines", .module = engines }, .{ .name = "engine_api", .module = api }, exlr8_build.cliImport(b, target, optimize, engine) }, + .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "tensorfold", .module = engine }, .{ .name = "core", .module = engine }, .{ .name = "cluster", .module = cluster }, .{ .name = "kernel_sources", .module = sources }, .{ .name = "native_engines", .module = engines }, .{ .name = "engine_api", .module = api }, .{ .name = "tokenizer", .module = tokenizer }, exlr8_build.cliImport(b, target, optimize, engine) }, }); if (p.c_source) |file| { mod.addCSourceFile(.{ .file = b.path(file), .flags = &.{"-std=c11"} }); diff --git a/zig/kernels/metal/moss_image_encoder.metal b/zig/kernels/metal/moss_image_encoder.metal new file mode 100644 index 0000000000..c302b174a9 --- /dev/null +++ b/zig/kernels/metal/moss_image_encoder.metal @@ -0,0 +1,225 @@ +// MOSS-Image's prompt encoder (Qwen3.5) after moss_image.metal, each step rounded where the publisher's torch rounds. + +struct MeRows { + uint width; + uint in_stride; + uint out_stride; + float eps; +}; + +// y = bf16(x rsqrt(mean x^2 + eps) (1 + w)) in fp32 (Qwen3.5's zero-centered RMSNorm); a threadgroup a row. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void me_norm( + const device bfloat* x [[buffer(0)]], const device bfloat* w [[buffer(1)]], device bfloat* y [[buffer(2)]], + constant MeRows& p [[buffer(3)]], 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 bfloat* xr = x + ulong(row) * p.in_stride; + device bfloat* yr = y + ulong(row) * p.out_stride; + float a = 0.0f; + for (uint i = t; i < p.width; i += n) a += float(xr[i]) * float(xr[i]); + const float inv = metal::precise::rsqrt(mi_sum(a, part, n, lane, sg) / float(p.width) + p.eps); + for (uint i = t; i < p.width; i += n) yr[i] = bfloat(float(xr[i]) * inv * (1.0f + float(w[i]))); +} + +struct MeHeads { + uint rows; + uint heads; + uint in_stride; + uint head_stride; + float eps; + uint normed; +}; + +// 256-wide heads to [heads, rows, 256], normed and the first 64 channels turned (else copied); a simdgroup each. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void me_head( + const device bfloat* x [[buffer(0)]], const device bfloat* w [[buffer(1)]], const device float* cs [[buffer(2)]], + const device float* sn [[buffer(3)]], device bfloat* y [[buffer(4)]], constant MeHeads& p [[buffer(5)]], + 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 >= p.heads) return; + const ulong from = ulong(r) * p.in_stride + h * p.head_stride + lane * 8; + const ulong to = (ulong(h) * p.rows + r) * 256 + lane * 8; + if (p.normed == 0) { + for (uint i = 0; i < 8; i++) y[to + i] = x[from + i]; + return; + } + float v[8]; + float ss = 0.0f; + for (uint i = 0; i < 8; i++) { + v[i] = float(x[from + i]); + ss += v[i] * v[i]; + } + const float inv = metal::precise::rsqrt(simd_sum(ss) / 256.0f + p.eps); + for (uint i = 0; i < 8; i++) v[i] = mi_bf(v[i] * inv * (1.0f + float(w[lane * 8 + i]))); + // lanes 0-7 hold the rotated 64 channels; a channel's partner is 32 channels, four lanes, away + const float sign = (lane & 4) != 0 ? 1.0f : -1.0f; + for (uint i = 0; i < 8; i++) { + const float pv = sign * simd_shuffle_xor(v[i], ushort(4)); + float o = v[i]; + if (lane < 8) { + const uint j = (lane & 3) * 8 + i; + o = mi_bf(v[i] * cs[ulong(r) * 32 + j]) + mi_bf(pv * sn[ulong(r) * 32 + j]); + } + y[to + i] = bfloat(o); + } +} + +struct MeAttn { + uint rows; + uint heads; + uint kv_heads; + float scale; +}; + +// o [rows, heads * 256] = causal softmax(q k^T scale) v in fp32, keys in order; a threadgroup a query row and head. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void me_attention( + const device bfloat* Q [[buffer(0)]], const device bfloat* K [[buffer(1)]], const device bfloat* V [[buffer(2)]], + device bfloat* O [[buffer(3)]], constant MeAttn& p [[buffer(4)]], uint2 tg [[threadgroup_position_in_grid]], + uint2 dt [[thread_position_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint sg [[simdgroup_index_in_threadgroup]]) { + threadgroup float part[33]; + const uint t = tg.x, h = tg.y, d = dt.x, kvh = h / (p.heads / p.kv_heads); + const float q = float(Q[(ulong(h) * p.rows + t) * 256 + d]); + float m = -INFINITY, l = 0.0f, acc = 0.0f; + for (uint j = 0; j <= t; j++) { + const ulong at = (ulong(kvh) * p.rows + j) * 256 + d; + const float s = mi_sum(q * float(K[at]), part, 256, lane, sg) * p.scale; + const float top = max(m, s); + const float old = metal::precise::exp(m - top), e = metal::precise::exp(s - top); + l = l * old + e; + acc = acc * old + e * float(V[at]); + m = top; + } + O[(ulong(t) * p.heads + h) * 256 + d] = bfloat(acc / l); +} + +// o = bf16(a bf16(sigmoid(g))): the attention rows gated by the second half of each 512-wide q_proj head. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void me_gate( + const device bfloat* a [[buffer(0)]], const device bfloat* qg [[buffer(1)]], device bfloat* o [[buffer(2)]], + constant uint& width [[buffer(3)]], uint2 at [[thread_position_in_grid]]) { + const ulong i = ulong(at.y) * width + at.x; + const float g = float(qg[ulong(at.y) * 2 * width + (at.x / 256) * 512 + 256 + at.x % 256]); + o[i] = bfloat(float(a[i]) * mi_bf(1.0f / (1.0f + metal::precise::exp(-g)))); +} + +// o = bf16(bf16(silu(g)) u): the MLP's product. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void me_silu_mul( + const device bfloat* g [[buffer(0)]], const device bfloat* u [[buffer(1)]], device bfloat* o [[buffer(2)]], + constant uint& width [[buffer(3)]], uint2 at [[thread_position_in_grid]]) { + const ulong i = ulong(at.y) * width + at.x; + const float gv = float(g[i]); + o[i] = bfloat(mi_bf(gv / (1.0f + metal::precise::exp(-gv))) * float(u[i])); +} + +// y = bf16(silu(bf16(the 4-tap causal depthwise conv of x))) by channel, the taps in order, rows before 0 zero. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void me_conv( + const device bfloat* x [[buffer(0)]], const device bfloat* w [[buffer(1)]], device bfloat* y [[buffer(2)]], + constant uint& width [[buffer(3)]], uint2 at [[thread_position_in_grid]]) { + float s = 0.0f; + for (uint k = 0; k < 4; k++) { + if (at.y + k >= 3) s += float(w[at.x * 4 + k]) * float(x[ulong(at.y + k - 3) * width + at.x]); + } + const float c = mi_bf(s); + y[ulong(at.y) * width + at.x] = bfloat(c / (1.0f + metal::precise::exp(-c))); +} + +struct MeGdn { + uint rows; + uint kheads; + uint vheads; + uint mix_stride; + float scale; +}; + +// The DeltaNet's fp32 inputs: q, k L2-normed (q scaled), beta = bf16(sigmoid(b)), g = -exp(A_log) softplus(a + dt). +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void me_gdn_prep( + const device bfloat* mix [[buffer(0)]], const device bfloat* a [[buffer(1)]], const device bfloat* b [[buffer(2)]], + const device bfloat* a_log [[buffer(3)]], const device bfloat* dt_bias [[buffer(4)]], device float* qn [[buffer(5)]], + device float* kn [[buffer(6)]], device float* g [[buffer(7)]], device float* beta [[buffer(8)]], + constant MeGdn& p [[buffer(9)]], 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 >= p.kheads) return; + const ulong row = ulong(r) * p.mix_stride; + const uint kdim = p.kheads * 128; + float q[4], k[4]; + float qq = 0.0f, kk = 0.0f; + for (uint i = 0; i < 4; i++) { + q[i] = float(mix[row + h * 128 + lane * 4 + i]); + k[i] = float(mix[row + kdim + h * 128 + lane * 4 + i]); + qq += q[i] * q[i]; + kk += k[i] * k[i]; + } + const float qi = metal::precise::rsqrt(simd_sum(qq) + 1e-6f), ki = metal::precise::rsqrt(simd_sum(kk) + 1e-6f); + const ulong to = (ulong(r) * p.kheads + h) * 128 + lane * 4; + for (uint i = 0; i < 4; i++) { + qn[to + i] = q[i] * qi * p.scale; + kn[to + i] = k[i] * ki; + } + const uint per = p.vheads / p.kheads; + if (lane < per) { + const uint v = h * per + lane; + const ulong at = ulong(r) * p.vheads + v; + beta[at] = mi_bf(1.0f / (1.0f + metal::precise::exp(-float(b[at])))); + const float x = float(a[at]) + float(dt_bias[v]); + const float soft = x > 20.0f ? x : metal::precise::log(1.0f + metal::precise::exp(x)); + g[at] = -metal::precise::exp(float(a_log[v])) * soft; + } +} + +// S = S exp(g); d = (v - S^T k) beta; S += k d^T; o = S^T q over the rows in order, a thread a value column in fp32. +[[max_total_threads_per_threadgroup(128)]] [[kernel]] void me_gdn( + const device float* qn [[buffer(0)]], const device float* kn [[buffer(1)]], const device bfloat* mix [[buffer(2)]], + const device float* g [[buffer(3)]], const device float* beta [[buffer(4)]], device bfloat* out [[buffer(5)]], + constant MeGdn& p [[buffer(6)]], uint h [[threadgroup_position_in_grid]], uint j [[thread_position_in_threadgroup]]) { + threadgroup float kq[256]; + float s[128]; + for (uint i = 0; i < 128; i++) s[i] = 0.0f; + const uint kh = h / (p.vheads / p.kheads); + const uint voff = 2 * p.kheads * 128 + h * 128 + j; + for (uint t = 0; t < p.rows; t++) { + const ulong at = (ulong(t) * p.kheads + kh) * 128 + j; + kq[j] = kn[at]; + kq[128 + j] = qn[at]; + threadgroup_barrier(mem_flags::mem_threadgroup); + const float decay = metal::precise::exp(g[ulong(t) * p.vheads + h]); + const float v = float(mix[ulong(t) * p.mix_stride + voff]); + float kv = 0.0f; + for (uint i = 0; i < 128; i++) { + s[i] *= decay; + kv += s[i] * kq[i]; + } + const float delta = (v - kv) * beta[ulong(t) * p.vheads + h]; + float o = 0.0f; + for (uint i = 0; i < 128; i++) { + s[i] += kq[i] * delta; + o += s[i] * kq[128 + i]; + } + out[(ulong(t) * p.vheads + h) * 128 + j] = bfloat(o); + threadgroup_barrier(mem_flags::mem_threadgroup); + } +} + +// y = bf16(bf16(w bf16(x rsqrt(mean x^2 + eps))) silu(z)) over each 128-wide head; a simdgroup a row and head. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void me_gated_norm( + const device bfloat* x [[buffer(0)]], const device bfloat* z [[buffer(1)]], const device bfloat* w [[buffer(2)]], + device bfloat* y [[buffer(3)]], constant MeRows& p [[buffer(4)]], 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 * 128 >= p.width) return; + const ulong at = ulong(r) * p.width + h * 128 + lane * 4; + float v[4]; + float ss = 0.0f; + for (uint i = 0; i < 4; i++) { + v[i] = float(x[at + i]); + ss += v[i] * v[i]; + } + const float inv = metal::precise::rsqrt(simd_sum(ss) / 128.0f + p.eps); + for (uint i = 0; i < 4; i++) { + const float n = mi_bf(float(w[lane * 4 + i]) * mi_bf(v[i] * inv)); + const float zv = float(z[at + i]); + y[at + i] = bfloat(n * (zv / (1.0f + metal::precise::exp(-zv)))); + } +} diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig index 2506fb4905..912441bf68 100644 --- a/zig/kernels/metal/sources.zig +++ b/zig/kernels/metal/sources.zig @@ -127,8 +127,8 @@ pub const glm_vision = @embedFile("prefill/gemm_nax.metal") ++ @embedFile("glm_v pub const core_qwen_vision = @embedFile("prefill/gemm_nax.metal") ++ @embedFile("core/qwen_vision.metal"); /// Qwen-Image-2.1's transformer: the prefill attention template, then its own steps (projections: core_qwen_vision). 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"); +/// MOSS-Image: the prefill attention template, then its transformer's and encoder's steps (gemms: core_qwen_vision). +pub const moss_image = @embedFile("prefill/attention_nax.metal") ++ @embedFile("moss_image.metal") ++ @embedFile("moss_image_encoder.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/moss_image_records.zig b/zig/src/core/contract/moss_image_records.zig index 2e1a572220..5995c26fee 100644 --- a/zig/src/core/contract/moss_image_records.zig +++ b/zig/src/core/contract/moss_image_records.zig @@ -1,4 +1,4 @@ -//! MOSS-Image's transformer launches: its own kernels, and the prefill split-K GEMM it reads fp32 sums from. +//! MOSS-Image's launches: its transformer's and prompt encoder's own kernels, and the split-K GEMM it reads sums from. const std = @import("std"); const M = @import("limits.zig").MetalLaunch; pub const records = [_]M{ @@ -12,6 +12,15 @@ pub const records = [_]M{ .{ .name = "moss image mi_attention_00", .threads = 128, .tg_memory = 0, .reserved = 128 }, .{ .name = "moss image mi_attention_01", .threads = 128, .tg_memory = 0, .reserved = 128 }, .{ .name = "moss image mi_attention_11", .threads = 128, .tg_memory = 0, .reserved = 128 }, + .{ .name = "moss image me_norm", .threads = 256, .tg_memory = 144, .reserved = 256 }, + .{ .name = "moss image me_head", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "moss image me_attention", .threads = 256, .tg_memory = 144, .reserved = 256 }, + .{ .name = "moss image me_gate", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "moss image me_silu_mul", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "moss image me_conv", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "moss image me_gdn_prep", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "moss image me_gdn", .threads = 128, .tg_memory = 1024, .reserved = 128 }, + .{ .name = "moss image me_gated_norm", .threads = 256, .tg_memory = 0, .reserved = 256 }, }; test "the MOSS-Image launches fit every Apple family" { diff --git a/zig/src/families/moss_image/encoder.zig b/zig/src/families/moss_image/encoder.zig new file mode 100644 index 0000000000..94069976b3 --- /dev/null +++ b/zig/src/families/moss_image/encoder.zig @@ -0,0 +1,221 @@ +//! MOSS-Image's prompt encoder on Metal: the chat-formatted prompt through Qwen3.5's text model to its final-norm rows. +const std = @import("std"); +const mtl = @import("metal"); +const cfg = @import("encoder_config.zig"); +const wts = @import("encoder_weights.zig"); +const ks = @import("kernels.zig"); +const Ref = ks.Ref; + +const H = cfg.hidden; +const F = cfg.ffn; +const D = cfg.head_dim; + +/// The pipeline's chat-formatted prompt: one user turn, then the assistant's turn opened with thinking on. +pub fn format(a: std.mem.Allocator, prompt: []const u8) ![]u8 { + return std.mem.concat(a, u8, &.{ "<|im_start|>user\n", prompt, "<|im_end|>\n<|im_start|>assistant\n\n" }); +} + +/// The prompt's ids by `tok` (the checkpoint's tokenizer.json) as the pipeline makes them, cut at 4,096; caller frees. +pub fn ids(tok: anytype, a: std.mem.Allocator, prompt: []const u8) ![]u32 { + const text = try format(a, prompt); + defer a.free(text); + const all = try tok.encode(a, text); + if (all.len <= cfg.max_rows) return all; + defer a.free(all); + return a.dupe(u32, all[0..cfg.max_rows]); +} + +/// The rotary tables [rows, 32] as the text model computes them: fp32 angles, their cos and sin, rounded to bf16. +pub fn rotary(cos: []f32, sin: []f32, rows: usize) void { + const half = cfg.rotary / 2; + for (0..half) |i| { + const exponent: f32 = @as(f32, @floatFromInt(2 * i)) / @as(f32, cfg.rotary); + const freq: f32 = 1.0 / std.math.pow(f32, @floatCast(cfg.theta), exponent); + for (0..rows) |r| { + const angle: f32 = @as(f32, @floatFromInt(r)) * freq; + cos[r * half + i] = bf(@floatCast(@cos(@as(f64, angle)))); + sin[r * half + i] = bf(@floatCast(@sin(@as(f64, angle)))); + } + } +} + +fn bf(v: f32) f32 { + const u: u32 = @bitCast(v); + return @bitCast(((u +% 0x7fff +% ((u >> 16) & 1)) >> 16) << 16); +} + +pub const Encoder = struct { + gpa: std.mem.Allocator, + device: mtl.Device, + queue: mtl.Queue, + k: ks.Kernels, + w: wts.Weights, + + /// The checkpoint folder `dir`: its text_encoder/. + pub fn load(gpa: std.mem.Allocator, io: std.Io, device: mtl.Device, dir: []const u8) !*Encoder { + if (!device.tensorUnits()) { + std.log.err("moss image encoder: {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(Encoder); + errdefer gpa.destroy(t); + t.gpa = gpa; + t.device = device; + t.queue = try device.queue(); + errdefer t.queue.deinit(); + t.k = try ks.Kernels.load(gpa, device); + errdefer t.k.deinit(); + const text = try std.fs.path.join(gpa, &.{ dir, "text_encoder" }); + defer gpa.free(text); + t.w = try wts.load(gpa, io, device, text); + return t; + } + + pub fn deinit(t: *Encoder) void { + t.w.deinit(); + t.k.deinit(); + t.queue.deinit(); + t.gpa.destroy(t); + } + + /// The final-norm rows [tokens, 4096] bf16 into `out`; with `kept`, each layer's rows [layers, tokens, 4096]. + pub fn rows(t: *Encoder, tokens: []const u32, out: []u16, kept: ?[]u16) !void { + const T = tokens.len; + if (T == 0 or T > cfg.max_rows) return error.PromptRows; + if (out.len != T * H) return error.PromptRows; + if (kept) |k| if (k.len != cfg.layers * T * H) return error.PromptRows; + const s = try Scratch.make(t.device, T); + defer s.buf.deinit(); + const x = s.at(.x); + for (tokens, 0..) |id, r| { + if (id >= cfg.vocab) return error.TokenOutOfRange; + @memcpy(halves(x.at(r * H * 2), H), t.w.embed[id * H ..][0..H]); + } + rotary(floats(s.at(.cos), T * 32), floats(s.at(.sin), T * 32), T); + @memset(halves(s.at(.ones), H), 0x3f80); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + var cb = t.queue.commandBuffer(); + var e = cb.compute(.serial); + for (t.w.layers, 0..) |l, i| { + t.layer(e, s, l, T); + if (kept) |k| { + try finish(cb, e); + @memcpy(k[i * T * H ..][0 .. T * H], halves(x, T * H)); + cb = t.queue.commandBuffer(); + e = cb.compute(.serial); + } + } + ks.encoderNorm(e, t.k.get("me_norm"), x, t.w.norm, s.at(.out), T, H, H, H, cfg.eps); + try finish(cb, e); + @memcpy(out, halves(s.at(.out), T * H)); + } + + fn layer(t: *const Encoder, e: mtl.ComputeEncoder, s: Scratch, l: wts.Layer, T: usize) void { + const x = s.at(.x); + const xn = s.at(.xn); + const y = s.at(.y); + ks.encoderNorm(e, t.k.get("me_norm"), x, l.norm1, xn, T, H, H, H, cfg.eps); + switch (l.mixer) { + .delta => |m| t.delta(e, s, m, T), + .attention => |m| t.attention(e, s, m, T), + } + ks.gateAdd(e, t.k.get("mi_gate_add"), x, y, s.at(.ones), T, H); + ks.encoderNorm(e, t.k.get("me_norm"), x, l.norm2, xn, T, H, H, H, cfg.eps); + t.k.linear(e, xn, l.gate, s.at(.gate), T, F, H); + t.k.linear(e, xn, l.up, s.at(.up), T, F, H); + ks.encoderRows(e, t.k.get("me_silu_mul"), s.at(.gate), s.at(.up), s.at(.act), T, F); + t.k.linear(e, s.at(.act), l.down, y, T, H, F); + ks.gateAdd(e, t.k.get("mi_gate_add"), x, y, s.at(.ones), T, H); + } + + fn delta(t: *const Encoder, e: mtl.ComputeEncoder, s: Scratch, m: wts.Delta, T: usize) void { + const xn = s.at(.xn); + t.k.linear(e, xn, m.qkv, s.at(.mix), T, cfg.mix, H); + t.k.linear(e, xn, m.z, s.at(.z), T, H, H); + t.k.linear(e, xn, m.b, s.at(.b), T, cfg.v_heads, H); + t.k.linear(e, xn, m.a, s.at(.a), T, cfg.v_heads, H); + ks.encoderRows(e, t.k.get("me_conv"), s.at(.mix), m.conv, s.at(.conv), T, cfg.mix); + const g: ks.Gdn = .{ .rows = @intCast(T), .kheads = cfg.k_heads, .vheads = cfg.v_heads, .mix_stride = cfg.mix, .scale = 1.0 / @sqrt(@as(f32, cfg.k_dim)) }; + ks.encoderGdnPrep(e, t.k.get("me_gdn_prep"), s.at(.conv), s.at(.a), s.at(.b), m.a_log, m.dt_bias, s.at(.qn), s.at(.kn), s.at(.g), s.at(.beta), g); + ks.encoderGdn(e, t.k.get("me_gdn"), s.at(.qn), s.at(.kn), s.at(.conv), s.at(.g), s.at(.beta), s.at(.core), g); + ks.encoderGatedNorm(e, t.k.get("me_gated_norm"), s.at(.core), s.at(.z), m.norm, s.at(.mixed), T, H, cfg.eps); + t.k.linear(e, s.at(.mixed), m.out, s.at(.y), T, H, H); + } + + fn attention(t: *const Encoder, e: mtl.ComputeEncoder, s: Scratch, m: wts.Attention, T: usize) void { + const xn = s.at(.xn); + const kv = cfg.kv_heads * D; + t.k.linear(e, xn, m.q, s.at(.mix), T, 2 * cfg.heads * D, H); + t.k.linear(e, xn, m.k, s.at(.kp), T, kv, H); + t.k.linear(e, xn, m.v, s.at(.vp), T, kv, H); + const heads = t.k.get("me_head"); + ks.encoderHeads(e, heads, s.at(.mix), m.q_norm, s.at(.cos), s.at(.sin), s.at(.q), T, cfg.heads, 2 * cfg.heads * D, 2 * D, cfg.eps, true); + ks.encoderHeads(e, heads, s.at(.kp), m.k_norm, s.at(.cos), s.at(.sin), s.at(.k), T, cfg.kv_heads, kv, D, cfg.eps, true); + ks.encoderHeads(e, heads, s.at(.vp), m.k_norm, s.at(.cos), s.at(.sin), s.at(.v), T, cfg.kv_heads, kv, D, cfg.eps, false); + ks.encoderAttention(e, t.k.get("me_attention"), s.at(.q), s.at(.k), s.at(.v), s.at(.att), T, cfg.heads, cfg.kv_heads, 1.0 / @sqrt(@as(f32, D))); + ks.encoderRows(e, t.k.get("me_gate"), s.at(.att), s.at(.mix), s.at(.mixed), T, H); + t.k.linear(e, s.at(.mixed), m.o, s.at(.y), T, H, H); + } +}; + +fn finish(cb: mtl.CommandBuffer, e: mtl.ComputeEncoder) !void { + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.log.err("moss image encoder: a pass failed on the GPU: {s}", .{msg}); + return error.GpuFailure; + } +} + +fn floats(r: Ref, n: usize) []f32 { + return @as([*]f32, @ptrCast(@alignCast(r.addr())))[0..n]; +} + +fn halves(r: Ref, n: usize) []u16 { + return @as([*]u16, @ptrCast(@alignCast(r.addr())))[0..n]; +} + +/// One buffer's parts for a prompt of `rows` rows, each 256-aligned. +const Scratch = struct { + const Part = enum { x, xn, y, mix, conv, z, a, b, qn, kn, g, beta, core, mixed, kp, vp, q, k, v, att, gate, up, act, cos, sin, ones, out }; + buf: mtl.Buffer, + offs: [std.enums.values(Part).len]usize, + + fn bytes(p: Part, T: usize) usize { + return switch (p) { + .x, .xn, .y, .z, .core, .mixed, .q, .att, .out => T * H * 2, + .mix, .conv => T * cfg.mix * 2, + .a, .b => T * cfg.v_heads * 2, + .qn, .kn => T * cfg.k_heads * cfg.k_dim * 4, + .g, .beta => T * cfg.v_heads * 4, + .kp, .vp, .k, .v => T * cfg.kv_heads * D * 2, + .gate, .up, .act => T * F * 2, + .cos, .sin => T * 32 * 4, + .ones => H * 2, + }; + } + + fn make(device: mtl.Device, T: usize) !Scratch { + var s: Scratch = undefined; + var total: usize = 0; + for (std.enums.values(Part), 0..) |p, i| { + s.offs[i] = total; + total += std.mem.alignForward(usize, bytes(p, T), 256); + } + s.buf = try device.buffer(total, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + return s; + } + + fn at(s: Scratch, p: Part) Ref { + return .{ .buf = s.buf, .off = s.offs[@backingInt(p)] }; + } +}; + +test "the pipeline's chat format" { + const a = std.testing.allocator; + const text = try format(a, "A red fox"); + defer a.free(text); + try std.testing.expectEqualStrings("<|im_start|>user\nA red fox<|im_end|>\n<|im_start|>assistant\n\n", text); +} diff --git a/zig/src/families/moss_image/encoder_config.zig b/zig/src/families/moss_image/encoder_config.zig new file mode 100644 index 0000000000..977ac93aac --- /dev/null +++ b/zig/src/families/moss_image/encoder_config.zig @@ -0,0 +1,112 @@ +//! The prompt encoder's geometry (the checkpoint's text_encoder: Qwen3.5's text model), and its config check. +const std = @import("std"); + +pub const hidden = 4096; +pub const layers = 32; +pub const vocab = 248320; +pub const ffn = 12288; +/// Every fourth layer full attention, the others Gated DeltaNet. +pub const full_every = 4; +pub const heads = 16; +pub const kv_heads = 4; +pub const head_dim = 256; +/// The rotated channels of each head: a quarter, halves 32 apart. +pub const rotary = 64; +pub const theta: f64 = 10_000_000; +pub const k_heads = 16; +pub const v_heads = 32; +pub const k_dim = 128; +pub const v_dim = 128; +pub const conv_taps = 4; +/// The DeltaNet's mixed row: q and k (k_heads x 128 each), then v (v_heads x 128). +pub const mix = 2 * k_heads * k_dim + v_heads * v_dim; +pub const eps: f32 = 1e-6; +/// The pipeline's longest prompt (max_sequence_length): longer ones are cut there, as its tokenizer call truncates. +pub const max_rows = 4096; + +pub fn isFull(layer: usize) bool { + return layer % full_every == full_every - 1; +} + +/// Accepts text_encoder/config.json at the released Qwen3.5 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 kind = doc.value.object.get("model_type") orelse return error.BadConfig; + if (kind != .string or !std.mem.eql(u8, kind.string, "qwen3_5")) return error.NotQwen35; + const tc = doc.value.object.get("text_config") orelse return error.BadConfig; + if (tc != .object) return error.BadConfig; + const o = tc.object; + const want = .{ + .{ "hidden_size", hidden }, .{ "num_hidden_layers", layers }, .{ "vocab_size", vocab }, + .{ "intermediate_size", ffn }, .{ "full_attention_interval", full_every }, .{ "num_attention_heads", heads }, + .{ "num_key_value_heads", kv_heads }, .{ "head_dim", head_dim }, .{ "linear_num_key_heads", k_heads }, + .{ "linear_num_value_heads", v_heads }, .{ "linear_key_head_dim", k_dim }, .{ "linear_value_head_dim", v_dim }, + .{ "linear_conv_kernel_dim", conv_taps }, + }; + 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("moss image encoder: {s} is not {d}; the encoder runs the released geometry only", .{ kv[0], kv[1] }); + return error.UnsupportedGeometry; + } + } + const flags = .{ .{ "attn_output_gate", true }, .{ "attention_bias", false } }; + inline for (flags) |kv| { + const v = o.get(kv[0]) orelse return error.BadConfig; + if (v != .bool or v.bool != kv[1]) return error.UnsupportedGeometry; + } + if (@as(f32, @floatCast(try number(o, "rms_norm_eps"))) != eps) return error.UnsupportedGeometry; + if (try number(o, "partial_rotary_factor") != 0.25) return error.UnsupportedGeometry; + const rope = o.get("rope_parameters") orelse return error.BadConfig; + if (rope != .object) return error.BadConfig; + if (try number(rope.object, "rope_theta") != theta) return error.UnsupportedGeometry; + const types = o.get("layer_types") orelse return error.BadConfig; + if (types != .array or types.array.items.len != layers) return error.BadConfig; + for (types.array.items, 0..) |t, i| { + if (t != .string) return error.BadConfig; + const full = std.mem.eql(u8, t.string, "full_attention"); + if (!full and !std.mem.eql(u8, t.string, "linear_attention")) return error.BadConfig; + if (full != isFull(i)) return error.UnsupportedGeometry; + } +} + +fn number(o: std.json.ObjectMap, name: []const u8) !f64 { + const v = o.get(name) orelse return error.BadConfig; + return switch (v) { + .float => |f| f, + .integer => |n| @floatFromInt(n), + else => error.BadConfig, + }; +} + +test "only the released text encoder's geometry passes" { + const a = std.testing.allocator; + var types: std.ArrayList(u8) = .empty; + defer types.deinit(a); + for (0..layers) |i| try types.appendSlice(a, if (isFull(i)) "\"full_attention\"," else "\"linear_attention\","); + _ = types.pop(); + const sample = try std.fmt.allocPrint(a, + \\{{"model_type": "qwen3_5", "text_config": {{"hidden_size": 4096, "num_hidden_layers": 32, "vocab_size": 248320, + \\ "intermediate_size": 12288, "full_attention_interval": 4, "num_attention_heads": 16, "num_key_value_heads": 4, + \\ "head_dim": 256, "linear_num_key_heads": 16, "linear_num_value_heads": 32, "linear_key_head_dim": 128, + \\ "linear_value_head_dim": 128, "linear_conv_kernel_dim": 4, "attn_output_gate": true, "attention_bias": false, + \\ "rms_norm_eps": 1e-06, "partial_rotary_factor": 0.25, "rope_parameters": {{"rope_theta": 10000000}}, + \\ "layer_types": [{s}]}}}} + , .{types.items}); + defer a.free(sample); + try check(a, sample); + const deeper = try std.mem.replaceOwned(u8, a, sample, "\"intermediate_size\": 12288", "\"intermediate_size\": 9728"); + defer a.free(deeper); + try std.testing.expectError(error.UnsupportedGeometry, check(a, deeper)); + const theta_off = try std.mem.replaceOwned(u8, a, sample, "10000000", "1000000"); + defer a.free(theta_off); + try std.testing.expectError(error.UnsupportedGeometry, check(a, theta_off)); + const swapped = try std.mem.replaceOwned(u8, a, sample, "\"full_attention\",\"linear_attention\"", "\"linear_attention\",\"full_attention\""); + defer a.free(swapped); + try std.testing.expectError(error.UnsupportedGeometry, check(a, swapped)); + const other = try std.mem.replaceOwned(u8, a, sample, "\"qwen3_5\"", "\"qwen3\""); + defer a.free(other); + try std.testing.expectError(error.NotQwen35, check(a, other)); +} diff --git a/zig/src/families/moss_image/encoder_weights.zig b/zig/src/families/moss_image/encoder_weights.zig new file mode 100644 index 0000000000..826107dcf5 --- /dev/null +++ b/zig/src/families/moss_image/encoder_weights.zig @@ -0,0 +1,115 @@ +//! The prompt encoder's tensors from the checkpoint's text_encoder folder, bf16 as stored, its vision tower unread. +const std = @import("std"); +const mtl = @import("metal"); +const ck = @import("../../core/checkpoint_metal.zig"); +const index_format = @import("../../core/checkpoint_index.zig"); +const cfg = @import("encoder_config.zig"); +const Ref = @import("../../core/qwen_vl/weights.zig").Ref; + +/// A Gated DeltaNet layer's mixer: the q k v projection and its conv, z, b and a, the decay's terms, its norm, out. +pub const Delta = struct { qkv: Ref, z: Ref, b: Ref, a: Ref, conv: Ref, a_log: Ref, dt_bias: Ref, norm: Ref, out: Ref }; + +/// A full-attention layer's mixer: q (with its gate), k, v, o and the head norms. +pub const Attention = struct { q: Ref, k: Ref, v: Ref, o: Ref, q_norm: Ref, k_norm: Ref }; + +pub const Layer = struct { + norm1: Ref, + norm2: Ref, + gate: Ref, + up: Ref, + down: Ref, + mixer: union(enum) { delta: Delta, attention: Attention }, +}; + +pub const Weights = struct { + ck: ck.Checkpoint, + layers: [cfg.layers]Layer, + /// The embedding table [vocab, hidden] on the host (shared storage). + embed: []const u16, + norm: Ref, + + pub fn deinit(w: *Weights) void { + w.ck.deinit(); + } +}; + +const prefix = "language_model."; + +fn take(c: *const ck.Checkpoint, name: []const u8, shape: []const usize) !Ref { + const t = try c.get(name); + if (t.dtype != .bf16 or t.rank != shape.len or !std.mem.eql(usize, t.shape[0..t.rank], shape)) { + std.log.err("moss image encoder: {s} is {t} {any}; the encoder reads bf16 {any}", .{ name, t.dtype, t.shape[0..t.rank], shape }); + return error.UnexpectedTensor; + } + if (t.offset % 16 != 0) return error.UnalignedTensor; + return .{ .buf = t.buffer, .off = t.offset }; +} + +/// The folder `dir` (a checkpoint's text_encoder/): config checked, then every language-model tensor. +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); + const index = try std.Io.Dir.cwd().readFileAlloc(io, try std.fs.path.join(arena, &.{ dir, "model.safetensors.index.json" }), arena, .limited(16 << 20)); + const doc = try std.json.parseFromSliceLeaky(std.json.Value, arena, index, .{}); + var files: std.StringArrayHashMapUnmanaged(void) = .empty; + var it = (try index_format.weightMap(doc)).iterator(); + while (it.next()) |kv| if (std.mem.startsWith(u8, kv.key_ptr.*, prefix)) try files.put(arena, kv.value_ptr.string, {}); + var w: Weights = undefined; + w.ck = ck.Checkpoint.init(gpa); + errdefer w.ck.deinit(); + for (files.keys()) |file| { + const path = try std.fmt.allocPrintSentinel(arena, "{s}/{s}", .{ dir, file }, 0); + try w.ck.addFileSelected(device, path, "", prefix); + } + const h = cfg.hidden; + const vd = cfg.v_heads * cfg.v_dim; + var name: [96]u8 = undefined; + for (&w.layers, 0..) |*l, i| { + const at = struct { + fn n(buf: []u8, layer: usize, field: []const u8) ![]const u8 { + return std.fmt.bufPrint(buf, prefix ++ "layers.{d}.{s}", .{ layer, field }); + } + }; + l.norm1 = try take(&w.ck, try at.n(&name, i, "input_layernorm.weight"), &.{h}); + l.norm2 = try take(&w.ck, try at.n(&name, i, "post_attention_layernorm.weight"), &.{h}); + l.gate = try take(&w.ck, try at.n(&name, i, "mlp.gate_proj.weight"), &.{ cfg.ffn, h }); + l.up = try take(&w.ck, try at.n(&name, i, "mlp.up_proj.weight"), &.{ cfg.ffn, h }); + l.down = try take(&w.ck, try at.n(&name, i, "mlp.down_proj.weight"), &.{ h, cfg.ffn }); + if (cfg.isFull(i)) { + const d = cfg.head_dim; + l.mixer = .{ .attention = .{ + .q = try take(&w.ck, try at.n(&name, i, "self_attn.q_proj.weight"), &.{ 2 * cfg.heads * d, h }), + .k = try take(&w.ck, try at.n(&name, i, "self_attn.k_proj.weight"), &.{ cfg.kv_heads * d, h }), + .v = try take(&w.ck, try at.n(&name, i, "self_attn.v_proj.weight"), &.{ cfg.kv_heads * d, h }), + .o = try take(&w.ck, try at.n(&name, i, "self_attn.o_proj.weight"), &.{ h, cfg.heads * d }), + .q_norm = try take(&w.ck, try at.n(&name, i, "self_attn.q_norm.weight"), &.{d}), + .k_norm = try take(&w.ck, try at.n(&name, i, "self_attn.k_norm.weight"), &.{d}), + } }; + } else { + l.mixer = .{ .delta = .{ + .qkv = try take(&w.ck, try at.n(&name, i, "linear_attn.in_proj_qkv.weight"), &.{ cfg.mix, h }), + .z = try take(&w.ck, try at.n(&name, i, "linear_attn.in_proj_z.weight"), &.{ vd, h }), + .b = try take(&w.ck, try at.n(&name, i, "linear_attn.in_proj_b.weight"), &.{ cfg.v_heads, h }), + .a = try take(&w.ck, try at.n(&name, i, "linear_attn.in_proj_a.weight"), &.{ cfg.v_heads, h }), + .conv = try take(&w.ck, try at.n(&name, i, "linear_attn.conv1d.weight"), &.{ cfg.mix, 1, cfg.conv_taps }), + .a_log = try take(&w.ck, try at.n(&name, i, "linear_attn.A_log"), &.{cfg.v_heads}), + .dt_bias = try take(&w.ck, try at.n(&name, i, "linear_attn.dt_bias"), &.{cfg.v_heads}), + .norm = try take(&w.ck, try at.n(&name, i, "linear_attn.norm.weight"), &.{cfg.v_dim}), + .out = try take(&w.ck, try at.n(&name, i, "linear_attn.out_proj.weight"), &.{ h, vd }), + } }; + } + } + const embed = try take(&w.ck, prefix ++ "embed_tokens.weight", &.{ cfg.vocab, h }); + w.embed = @as([*]const u16, @ptrCast(@alignCast(embed.addr())))[0 .. cfg.vocab * h]; + w.norm = try take(&w.ck, prefix ++ "norm.weight", &.{h}); + const full = cfg.layers / cfg.full_every; + const read = 5 * cfg.layers + 6 * full + 9 * (cfg.layers - full) + 2; + if (w.ck.tensors.count() != read) { + std.log.err("moss image encoder: the checkpoint has {d} language-model tensors; the encoder reads {d}", .{ w.ck.tensors.count(), read }); + return error.UnexpectedTensors; + } + return w; +} diff --git a/zig/src/families/moss_image/kernels.zig b/zig/src/families/moss_image/kernels.zig index 33aa044f01..388e2c092e 100644 --- a/zig/src/families/moss_image/kernels.zig +++ b/zig/src/families/moss_image/kernels.zig @@ -21,6 +21,15 @@ pub const names = [_][]const u8{ "mi_attention_00", "mi_attention_01", "mi_attention_11", + "me_norm", + "me_head", + "me_attention", + "me_gate", + "me_silu_mul", + "me_conv", + "me_gdn_prep", + "me_gdn", + "me_gated_norm", }; const tower_count = 6; const sums = 4; @@ -163,6 +172,55 @@ pub fn bias(e: mtl.ComputeEncoder, p: mtl.Pipeline, c: Ref, b: Ref, y: Ref, rows e.dispatchThreads(mtl.Size.of(width, rows, 1), mtl.Size.of(256, 1, 1)); } +/// y = Qwen3.5's zero-centered RMSNorm of `rows` rows of `width` (rows `in_stride` and `out_stride` apart). +pub fn encoderNorm(e: mtl.ComputeEncoder, p: mtl.Pipeline, x: Ref, w: Ref, y: Ref, rows: usize, width: usize, in_stride: usize, out_stride: usize, eps: f32) void { + const params = extern struct { width: u32, in_stride: u32, out_stride: u32, eps: f32 }{ .width = @intCast(width), .in_stride = @intCast(in_stride), .out_stride = @intCast(out_stride), .eps = eps }; + Kernels.bind(e, p, &.{ x, w, y }, params); + e.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(256, 1, 1)); +} + +/// 256-wide heads to [count, rows, 256]: normed and turned by cs, sn [rows, 32] when `normed`, else copied. +pub fn encoderHeads(e: mtl.ComputeEncoder, p: mtl.Pipeline, x: Ref, w: Ref, cs: Ref, sn: Ref, y: Ref, rows: usize, count: usize, in_stride: usize, head_stride: usize, eps: f32, normed: bool) void { + const params = extern struct { rows: u32, heads: u32, in_stride: u32, head_stride: u32, eps: f32, normed: u32 }{ .rows = @intCast(rows), .heads = @intCast(count), .in_stride = @intCast(in_stride), .head_stride = @intCast(head_stride), .eps = eps, .normed = @intFromBool(normed) }; + Kernels.bind(e, p, &.{ x, w, cs, sn, y }, params); + e.dispatchGroups(mtl.Size.of(rows, (count + 7) / 8, 1), mtl.Size.of(256, 1, 1)); +} + +/// o [rows, count * 256] = causal attention of q [count, rows, 256] over k, v [kv_count, rows, 256]. +pub fn encoderAttention(e: mtl.ComputeEncoder, p: mtl.Pipeline, q: Ref, k: Ref, v: Ref, o: Ref, rows: usize, count: usize, kv_count: usize, scale: f32) void { + const params = extern struct { rows: u32, heads: u32, kv_heads: u32, scale: f32 }{ .rows = @intCast(rows), .heads = @intCast(count), .kv_heads = @intCast(kv_count), .scale = scale }; + Kernels.bind(e, p, &.{ q, k, v, o }, params); + e.dispatchGroups(mtl.Size.of(rows, count, 1), mtl.Size.of(256, 1, 1)); +} + +/// One of the elementwise steps (me_gate, me_silu_mul, me_conv) over `rows` rows of `width`. +pub fn encoderRows(e: mtl.ComputeEncoder, p: mtl.Pipeline, a: Ref, b: Ref, y: Ref, rows: usize, width: usize) void { + Kernels.bind(e, p, &.{ a, b, y }, @as(u32, @intCast(width))); + e.dispatchThreads(mtl.Size.of(width, rows, 1), mtl.Size.of(256, 1, 1)); +} + +/// The DeltaNet's geometry for me_gdn_prep and me_gdn. +pub const Gdn = extern struct { rows: u32, kheads: u32, vheads: u32, mix_stride: u32, scale: f32 }; + +/// qn, kn [rows, kheads, 128] and g, beta [rows, vheads] in fp32 from the conv's rows and the a, b projections. +pub fn encoderGdnPrep(e: mtl.ComputeEncoder, p: mtl.Pipeline, mix: Ref, a: Ref, b: Ref, a_log: Ref, dt_bias: Ref, qn: Ref, kn: Ref, g: Ref, beta: Ref, params: Gdn) void { + Kernels.bind(e, p, &.{ mix, a, b, a_log, dt_bias, qn, kn, g, beta }, params); + e.dispatchGroups(mtl.Size.of(params.rows, (params.kheads + 7) / 8, 1), mtl.Size.of(256, 1, 1)); +} + +/// The gated delta rule's out [rows, vheads, 128] bf16: a threadgroup a value head. +pub fn encoderGdn(e: mtl.ComputeEncoder, p: mtl.Pipeline, qn: Ref, kn: Ref, mix: Ref, g: Ref, beta: Ref, out: Ref, params: Gdn) void { + Kernels.bind(e, p, &.{ qn, kn, mix, g, beta, out }, params); + e.dispatchGroups(mtl.Size.of(params.vheads, 1, 1), mtl.Size.of(128, 1, 1)); +} + +/// y [rows, width] = the gated norm of x by its 128-wide heads with gate z and weight w. +pub fn encoderGatedNorm(e: mtl.ComputeEncoder, p: mtl.Pipeline, x: Ref, z: Ref, w: Ref, y: Ref, rows: usize, width: usize, eps: f32) void { + const params = extern struct { width: u32, in_stride: u32, out_stride: u32, eps: f32 }{ .width = @intCast(width), .in_stride = @intCast(width), .out_stride = @intCast(width), .eps = eps }; + Kernels.bind(e, p, &.{ x, z, w, y }, params); + e.dispatchGroups(mtl.Size.of(rows, (width / 128 + 7) / 8, 1), mtl.Size.of(256, 1, 1)); +} + test "every MOSS-Image pipeline compiles and loads on this GPU" { const pool = mtl.objc.Pool.push(); defer pool.pop(); diff --git a/zig/src/families/moss_image/moss_image.zig b/zig/src/families/moss_image/moss_image.zig index 739c2dd08a..d8fa5640c0 100644 --- a/zig/src/families/moss_image/moss_image.zig +++ b/zig/src/families/moss_image/moss_image.zig @@ -4,6 +4,10 @@ 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 encoder_config = @import("encoder_config.zig"); +pub const encoder_weights = @import("encoder_weights.zig"); +pub const encoder = @import("encoder.zig"); +pub const Encoder = encoder.Encoder; pub const Model = model.Model; pub const Prompt = model.Prompt; pub const Timing = model.Timing; @@ -14,4 +18,7 @@ test { _ = weights; _ = kernels; _ = model; + _ = encoder_config; + _ = encoder_weights; + _ = encoder; } diff --git a/zig/tests/contract/metal_moss_image.zig b/zig/tests/contract/metal_moss_image.zig index 30b158ac65..aebab59ab8 100644 --- a/zig/tests/contract/metal_moss_image.zig +++ b/zig/tests/contract/metal_moss_image.zig @@ -8,6 +8,8 @@ const rig_ = @import("rig.zig"); const mi = @import("metal_moss_image_common.zig"); const norms = @import("metal_moss_image_norms.zig"); const steps = @import("metal_moss_image_steps.zig"); +const enc = @import("metal_moss_image_encoder.zig"); +const enc_state = @import("metal_moss_image_encoder_state.zig"); const c = mi.c; const g = mi.g; @@ -244,6 +246,14 @@ const controls = [_]?Control{ .{ .kernel = "mi_situ", .what = "drops bf16(tanh(h / 25))", .from = "mi_bf(25.0f * mi_bf(metal::precise::tanh(", .to = "mi_bf(25.0f * (metal::precise::tanh(" }, .{ .kernel = "mi_gate_add", .what = "drops bf16(g y)", .from = "x[i] = bfloat(float(x[i]) + mi_bf(float(g[at.x]) * float(y[i])));", .to = "x[i] = bfloat(float(x[i]) + (float(g[at.x]) * float(y[i])));" }, .{ .kernel = "mi_bias", .what = "adds bf16(c) before the bias", .from = "y[i] = bfloat(c[i] + float(b[at.x]));", .to = "y[i] = bfloat(mi_bf(c[i]) + float(b[at.x]));" }, + .{ .kernel = "me_norm", .what = "adds bf16(x inv)", .from = "yr[i] = bfloat(float(xr[i]) * inv * (1.0f + float(w[i])));", .to = "yr[i] = bfloat(mi_bf(float(xr[i]) * inv) * (1.0f + float(w[i])));" }, + .{ .kernel = "me_head q", .what = "drops bf16(q cos)", .from = "o = mi_bf(v[i] * cs[ulong(r) * 32 + j]) + mi_bf(pv * sn[ulong(r) * 32 + j]);", .to = "o = (v[i] * cs[ulong(r) * 32 + j]) + mi_bf(pv * sn[ulong(r) * 32 + j]);" }, + .{ .kernel = "me_gated_norm", .what = "drops bf16(x inv)", .from = "mi_bf(float(w[lane * 4 + i]) * mi_bf(v[i] * inv))", .to = "mi_bf(float(w[lane * 4 + i]) * (v[i] * inv))" }, + .{ .kernel = "me_gate", .what = "drops bf16(sigmoid(g))", .from = "o[i] = bfloat(float(a[i]) * mi_bf(1.0f / (1.0f + metal::precise::exp(-g))));", .to = "o[i] = bfloat(float(a[i]) * (1.0f / (1.0f + metal::precise::exp(-g))));" }, + .{ .kernel = "me_silu_mul", .what = "drops bf16(silu(g))", .from = "o[i] = bfloat(mi_bf(gv / (1.0f + metal::precise::exp(-gv))) * float(u[i]));", .to = "o[i] = bfloat((gv / (1.0f + metal::precise::exp(-gv))) * float(u[i]));" }, + .{ .kernel = "me_conv", .what = "drops bf16(the taps' sum)", .from = "const float c = mi_bf(s);", .to = "const float c = s;" }, + .{ .kernel = "me_gdn_prep", .what = "drops bf16(beta)", .from = "beta[at] = mi_bf(1.0f / (1.0f + metal::precise::exp(-float(b[at]))));", .to = "beta[at] = (1.0f / (1.0f + metal::precise::exp(-float(b[at]))));" }, + .{ .kernel = "me_gdn", .what = "adds bf16(the state)", .from = "s[i] += kq[i] * delta;", .to = "s[i] = mi_bf(s[i] + kq[i] * delta);" }, }; /// The family's source as the engine compiles it, with `ctl`'s change made where it occurs, exactly once. @@ -284,6 +294,12 @@ pub const Suite = struct { adds: std.ArrayList(*GateAdd) = .empty, biases: std.ArrayList(*Bias) = .empty, attns: std.ArrayList(*Attention) = .empty, + heads: std.ArrayList(*enc.Head) = .empty, + gated: std.ArrayList(*enc.GatedNorm) = .empty, + rows: std.ArrayList(*enc.Step) = .empty, + preps: std.ArrayList(*enc_state.Prep) = .empty, + gdns: std.ArrayList(*enc_state.Gdn) = .empty, + encoder_attns: std.ArrayList(*enc_state.Attention) = .empty, fn pipeline(s: *Suite, device: mtl.Device, lib: *Lazy, name: []const u8) !mtl.Pipeline { const p = try mtl.Pipeline.init(device, try lib.get(), name, false); @@ -321,11 +337,15 @@ pub const Suite = struct { return if (o) |needle| std.mem.indexOf(u8, name, needle) != null else true; } }; - for ([_]bool{ false, true }) |modulated| { - const kind = if (modulated) "mi_norm modulated" else "mi_norm"; + for (std.enums.values(norms.NormKind)) |nk| { + const kind = switch (nk) { + .plain => "mi_norm", + .modulated => "mi_norm modulated", + .zero_centered => "me_norm", + }; const name = try s.label("metal moss image {s}{s}", .{ kind, suffix }); if (!want.is(ctl, kind) or !want.keep(only, name)) continue; - const t = try s.make(Norm, &s.norms, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "mi_norm"), .modulated = modulated, .name = name, .ends = ends }); + const t = try s.make(Norm, &s.norms, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, if (nk == .zero_centered) "me_norm" else "mi_norm"), .kind = nk, .name = name, .ends = ends }); try out.append(s.a, t.kernel(gate)); } for ([_]bool{ false, true }) |packed_rows| { @@ -354,7 +374,49 @@ pub const Suite = struct { try out.append(s.a, t.kernel(gate)); } } + for (std.enums.values(enc.HeadForm)) |form| { + const kind = try s.label("me_head {t}", .{form}); + const name = try s.label("metal moss image {s}{s}", .{ kind, suffix }); + if (!want.is(ctl, kind) or !want.keep(only, name)) continue; + const t = try s.make(enc.Head, &s.heads, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "me_head"), .form = form, .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + { + const name = try s.label("metal moss image me_gated_norm{s}", .{suffix}); + if (want.is(ctl, "me_gated_norm") and want.keep(only, name)) { + const t = try s.make(enc.GatedNorm, &s.gated, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "me_gated_norm"), .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + } + for (std.enums.values(enc.StepKind)) |kind| { + const fn_name = try s.label("me_{t}", .{kind}); + const name = try s.label("metal moss image {s}{s}", .{ fn_name, suffix }); + if (!want.is(ctl, fn_name) or !want.keep(only, name)) continue; + const t = try s.make(enc.Step, &s.rows, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, fn_name), .kind = kind, .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + { + const name = try s.label("metal moss image me_gdn_prep{s}", .{suffix}); + if (want.is(ctl, "me_gdn_prep") and want.keep(only, name)) { + const t = try s.make(enc_state.Prep, &s.preps, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "me_gdn_prep"), .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + } + { + const name = try s.label("metal moss image me_gdn{s}", .{suffix}); + if (want.is(ctl, "me_gdn") and want.keep(only, name)) { + const t = try s.make(enc_state.Gdn, &s.gdns, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "me_gdn"), .name = name }); + try out.append(s.a, t.kernel(gate)); + } + } if (ctl != null) return; + { + const name = try s.label("metal moss image me_attention", .{}); + if (want.keep(only, name)) { + const t = try s.make(enc_state.Attention, &s.encoder_attns, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "me_attention"), .name = name }); + try out.append(s.a, t.kernel(null)); + } + } for (forms, 0..) |f, form| { const name = try s.label("metal moss image {s}", .{f.name}); if (!want.keep(only, name)) continue; @@ -381,7 +443,7 @@ pub const Suite = struct { pub fn destroy(s: *Suite) void { const a = s.a; - inline for (.{ &s.norms, &s.qks, &s.situs, &s.adds, &s.biases, &s.attns }) |list| { + inline for (.{ &s.norms, &s.qks, &s.situs, &s.adds, &s.biases, &s.attns, &s.heads, &s.gated, &s.rows, &s.preps, &s.gdns, &s.encoder_attns }) |list| { for (list.items) |t| { t.deinit(); a.destroy(t); diff --git a/zig/tests/contract/metal_moss_image_common.zig b/zig/tests/contract/metal_moss_image_common.zig index dd5ce450b0..f594589b19 100644 --- a/zig/tests/contract/metal_moss_image_common.zig +++ b/zig/tests/contract/metal_moss_image_common.zig @@ -64,6 +64,10 @@ pub const Range = struct { pub fn recip(a: Range) Range { return .{ .lo = 1 / a.hi, .hi = 1 / a.lo }; } + /// Both ranges' span. + pub fn join(a: Range, b: Range) Range { + return .{ .lo = @min(a.lo, b.lo), .hi = @max(a.hi, b.hi) }; + } }; pub fn tanh(x: f64) f64 { @@ -75,6 +79,24 @@ pub const div_ulps = 2.5; pub const rsqrt_ulps = 2; pub const exp_ulps = 4; pub const tanh_ulps = 5; +pub const log_ulps = 4; + +/// 1 + exp(-x) at a point x through the kernel's precise exp and add; past fp32's range, exp's fp32 result is +Inf. +fn denominator(x: f64) Range { + var e = Range.of(@exp(-x)).fp32(exp_ulps); + if (e.hi > std.math.floatMax(f32)) e.hi = std.math.inf(f64); + return e.add(Range.of(1)).fp32(0.5); +} + +/// 1 / (1 + exp(-x)) at a point x through the kernel's precise exp, add and fast division (before any rounding). +pub fn sigmoid(x: f64) Range { + return denominator(x).recip().fp32(div_ulps); +} + +/// x / (1 + exp(-x)) at a point x, as the kernels compute SiLU. +pub fn silu(x: f64) Range { + return Range.of(x).mul(denominator(x).recip()).fp32(div_ulps); +} /// Each output's float64 value at the kernel's rounding points, and its range, kept for the bounds and intervals. pub const Ends = struct { diff --git a/zig/tests/contract/metal_moss_image_encoder.zig b/zig/tests/contract/metal_moss_image_encoder.zig new file mode 100644 index 0000000000..d06f3adfc0 --- /dev/null +++ b/zig/tests/contract/metal_moss_image_encoder.zig @@ -0,0 +1,496 @@ +//! The encoder's me_head, me_gated_norm, me_gate, me_silu_mul and me_conv under the contract, at torch's bf16 points. +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_moss_image_common.zig"); + +const c = mi.c; +const g = mi.g; +const ks = mi.ks; +const ec = tf.moss_image.encoder_config; +const Range = mi.Range; +const Ends = mi.Ends; +const bf = mi.bf; +const kref = mi.kref; +const W = ec.head_dim; + +const norm_why = "a bf16 output follows the sum of squares as ss^-1/2: a bf16 sum's error is a few output ulps"; + +// --- me_head: 256-wide heads normed (1 + w), the first 64 channels turned, laid out [heads, rows, 256]; v copied --- + +/// Eight channels a lane in turn, then the simdgroup tree. +fn headDepth(_: c.Shape) usize { + return 8 + 5; +} + +/// q: heads 512 apart in q_proj's rows (each followed by its gate); k: packed heads; v: packed heads, copied. +pub const HeadForm = enum { q, k, v }; + +pub const Head = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + form: HeadForm, + name: []const u8, + shape: c.Shape = undefined, + fill: g.Fill = .random, + x: []u16 = &.{}, + w: []u16 = &.{}, + rot: []f32 = &.{}, + ins: [4]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + /// Rows: positions; n: heads (16 q, 4 k and v); k: the head width (256). + pub fn spec(name: []const u8, form: HeadForm) c.Spec { + const served: []const c.Shape = switch (form) { + .q => &.{ .{ .rows = 63, .n = ec.heads, .k = W }, .{ .rows = 10, .n = ec.heads, .k = W }, .{ .rows = 4096, .n = ec.heads, .k = W } }, + .k, .v => &.{ .{ .rows = 63, .n = ec.kv_heads, .k = W }, .{ .rows = 10, .n = ec.kv_heads, .k = W }, .{ .rows = 4096, .n = ec.kv_heads, .k = W } }, + }; + return .{ + .name = name, + .accum = if (form == .v) .none else .f32, + .depth = if (form == .v) mi.noSums else headDepth, + .ops = 2 * (mi.div_ulps + mi.rsqrt_ulps) + 3, + .mids = 3, + .out = .bf16, + .served = served, + .sums = .{ .bf16 = 0, .why = norm_why }, + }; + } + + pub fn deinit(t: *Head) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *Head) void { + t.a.free(t.x); + t.a.free(t.w); + t.a.free(t.rot); + t.x = &.{}; + t.w = &.{}; + t.rot = &.{}; + } + pub fn kernel(t: *Head, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name, t.form), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Head); + + fn stride(t: *const Head) usize { + return if (t.form == .q) 2 * t.shape.n * W else t.shape.n * W; + } + fn headStride(t: *const Head) usize { + return if (t.form == .q) 2 * W else W; + } + + pub fn prepare(t: *Head, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.k != W) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + t.fill = fill; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + t.x = try t.a.alloc(u16, shape.rows * t.stride()); + if (fill == .stress) mi.integers(rng, t.x, 128) else mi.randomBf16(rng, t.x, -4, 4); + t.w = try t.a.alloc(u16, W); + mi.randomBf16(rng, t.w, -0.75, 1); + const half = shape.rows * 32; + t.rot = try t.a.alloc(f32, 2 * half); + for (0..half) |i| { + const angle = rng.float(f32) * 2 * std.math.pi; + t.rot[i] = c.fromBf16(c.toBf16(@cos(angle))); + t.rot[half + i] = c.fromBf16(c.toBf16(@sin(angle))); + } + if (fill == .nonfinite) { + const at = [_][3]usize{ .{ 0, 0, 5 }, .{ 1, 1, 40 }, .{ 2, 2, 200 }, .{ 3, 0, 0 }, .{ 4, 3, 63 } }; + for (at, [_]u16{ mi.nan16, mi.inf16, mi.ninf16, mi.nan16, mi.ninf16 }) |p, v| if (p[0] < shape.rows and p[1] < shape.n) { + t.x[p[0] * t.stride() + p[1] * t.headStride() + p[2]] = v; + }; + } + 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.rot[0..half])); + t.ins[3] = try t.rig.input(std.mem.sliceAsBytes(t.rot[half..])); + t.out = try t.rig.output(shape.rows * shape.n * W * 2); + return true; + } + + pub fn run(t: *Head) !void { + const p = t.rig.begin(); + ks.encoderHeads(p.e, t.pipe, kref(t.ins[0]), kref(t.ins[1]), kref(t.ins[2]), kref(t.ins[3]), kref(t.out), t.shape.rows, t.shape.n, t.stride(), t.headStride(), ec.eps, t.form != .v); + try t.rig.finish(p); + } + + pub fn count(t: *Head) usize { + return t.shape.rows * t.shape.n * W; + } + + fn head(t: *Head, ends: *Ends, r: usize, h: usize, ss: f64, ssr: Range, want: []f64, mag: []f64) void { + const rows = t.shape.rows; + const x = t.x[r * t.stride() + h * t.headStride() ..][0..W]; + const base = (h * rows + r) * W; + if (t.form == .v) { + for (x, 0..) |v, ch| ends.put(want, mag, base + ch, c.fromBf16(v), Range.of(c.fromBf16(v))); + return; + } + const n = mi.invNorm(ss, ssr, W); + var normed: [W]f64 = undefined; + var nr: [W]Range = undefined; + for (x, t.w, &normed, &nr) |xv, wv, *o, *orr| { + const f: f64 = c.fromBf16(xv); + const one_w = 1 + @as(f64, c.fromBf16(wv)); + o.* = bf(f * n.inv * one_w); + orr.* = Range.of(f).mul(n.r).fp32(0.5).mul(Range.of(one_w).fp32(0.5)).fp32(0.5).bf16(); + } + const half = rows * 32; + for (0..W) |ch| { + if (ch >= ec.rotary) { + ends.put(want, mag, base + ch, normed[ch], nr[ch]); + continue; + } + const j = ch % 32; + const partner = if (ch < 32) ch + 32 else ch - 32; + const sign: f64 = if (ch < 32) -1 else 1; + const cs: f64 = t.rot[r * 32 + j]; + const sn: f64 = t.rot[half + r * 32 + j]; + const v = bf(bf(normed[ch] * cs) + bf(sign * normed[partner] * sn)); + const vr = nr[ch].mul(Range.of(cs)).fp32(0.5).bf16().add(nr[partner].mul(Range.of(sign * sn)).fp32(0.5).bf16()).fp32(0.5).bf16(); + ends.put(want, mag, base + ch, v, vr); + } + } + + pub fn reference(t: *Head, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + for (0..t.shape.rows) |r| for (0..t.shape.n) |h| { + const s = mi.squares(t.x[r * t.stride() + h * t.headStride() ..][0..W], headDepth(t.shape), t.fill == .stress); + t.head(&t.ends, r, h, s.ss, s.r, want, mag); + }; + } + pub fn bounds(t: *Head, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *Head, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } + + pub fn simulate(t: *Head, 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: [W]f64 = undefined; + for (0..t.shape.rows) |r| for (0..t.shape.n) |h| { + for (&terms, t.x[r * t.stride() + h * t.headStride() ..][0..W]) |*v, x| { + const f: f64 = c.fromBf16(x); + v.* = f * f; + } + const ss = contract.rows.sumIn(acc, order, &terms); + t.head(&scratch, r, h, ss, Range.of(ss), out, mag); + }; + } +}; + +// --- me_gated_norm: bf16(bf16(w bf16(x rsqrt(mean x^2 + eps))) silu(z)) by 128-wide head ----------------------------- + +/// Four channels a lane in turn, then the simdgroup tree. +fn gatedDepth(_: c.Shape) usize { + return 4 + 5; +} + +pub const GatedNorm = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + name: []const u8, + shape: c.Shape = undefined, + fill: g.Fill = .random, + x: []u16 = &.{}, + z: []u16 = &.{}, + w: []u16 = &.{}, + ins: [3]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + /// Rows: positions; n: the value width (32 heads of 128); k: the head width. + pub fn spec(name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = gatedDepth, + .ops = 2 * (mi.div_ulps + mi.rsqrt_ulps + mi.exp_ulps + mi.div_ulps) + 2, + .mids = 2, + .out = .bf16, + .served = &.{ .{ .rows = 63, .n = ec.hidden, .k = ec.v_dim }, .{ .rows = 10, .n = ec.hidden, .k = ec.v_dim }, .{ .rows = 4096, .n = ec.hidden, .k = ec.v_dim } }, + .sums = .{ .bf16 = 0, .why = norm_why }, + }; + } + + pub fn deinit(t: *GatedNorm) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *GatedNorm) void { + for ([_]*[]u16{ &t.x, &t.z, &t.w }) |p| { + t.a.free(p.*); + p.* = &.{}; + } + } + pub fn kernel(t: *GatedNorm, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(GatedNorm); + + pub fn prepare(t: *GatedNorm, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.k != ec.v_dim or shape.n == 0 or shape.n % ec.v_dim != 0) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + t.fill = fill; + const n = shape.rows * shape.n; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + t.x = try t.a.alloc(u16, n); + t.z = try t.a.alloc(u16, n); + t.w = try t.a.alloc(u16, ec.v_dim); + if (fill == .stress) mi.integers(rng, t.x, 128) else mi.randomBf16(rng, t.x, -2, 2); + mi.randomBf16(rng, t.z, -12, 12); + mi.randomBf16(rng, t.w, 0.25, 2); + if (fill == .nonfinite) { + for ([_]u16{ mi.nan16, mi.inf16, mi.ninf16 }, 0..) |v, i| { + t.x[(i * 131 + 7) % n] = v; + t.z[(i * 257 + 3) % n] = v; + } + t.w[9] = mi.nan16; + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.x)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.z)); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.w)); + t.out = try t.rig.output(n * 2); + return true; + } + + pub fn run(t: *GatedNorm) !void { + const p = t.rig.begin(); + ks.encoderGatedNorm(p.e, t.pipe, kref(t.ins[0]), kref(t.ins[1]), kref(t.ins[2]), kref(t.out), t.shape.rows, t.shape.n, ec.eps); + try t.rig.finish(p); + } + + pub fn count(t: *GatedNorm) usize { + return t.shape.rows * t.shape.n; + } + + fn head(t: *GatedNorm, ends: *Ends, at: usize, ss: f64, ssr: Range, want: []f64, mag: []f64) void { + const n = mi.invNorm(ss, ssr, ec.v_dim); + for (0..ec.v_dim) |i| { + const x: f64 = c.fromBf16(t.x[at + i]); + const z: f64 = c.fromBf16(t.z[at + i]); + const w: f64 = c.fromBf16(t.w[i]); + const y = bf(bf(w * bf(x * n.inv)) * (z / (1 + @exp(-z)))); + const yr = Range.of(w).mul(Range.of(x).mul(n.r).fp32(0.5).bf16()).fp32(0.5).bf16().mul(mi.silu(z)).fp32(0.5).bf16(); + ends.put(want, mag, at + i, y, yr); + } + } + + pub fn reference(t: *GatedNorm, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + var at: usize = 0; + while (at < t.count()) : (at += ec.v_dim) { + const s = mi.squares(t.x[at..][0..ec.v_dim], gatedDepth(t.shape), t.fill == .stress); + t.head(&t.ends, at, s.ss, s.r, want, mag); + } + } + pub fn bounds(t: *GatedNorm, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *GatedNorm, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } + + pub fn simulate(t: *GatedNorm, 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: [ec.v_dim]f64 = undefined; + var at: usize = 0; + while (at < t.count()) : (at += ec.v_dim) { + for (&terms, t.x[at..][0..ec.v_dim]) |*v, x| { + const f: f64 = c.fromBf16(x); + v.* = f * f; + } + const ss = contract.rows.sumIn(acc, order, &terms); + t.head(&scratch, at, ss, Range.of(ss), out, mag); + } + } +}; + +// --- me_gate, me_silu_mul and me_conv: elementwise steps at the publisher's bf16 points ---------------------------- + +pub const StepKind = enum { gate, silu_mul, conv }; + +/// The conv's four taps, summed in turn. +fn convDepth(_: c.Shape) usize { + return ec.conv_taps; +} + +pub const Step = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + kind: StepKind, + name: []const u8, + shape: c.Shape = undefined, + fill: g.Fill = .random, + /// gate: the attention rows; silu_mul: the gate; conv: the rows + x: []u16 = &.{}, + /// gate: q_proj's rows (twice as wide); silu_mul: up; conv: the taps [width, 4] + y: []u16 = &.{}, + ins: [2]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + /// Rows: positions; n: the row width; k: the conv's taps (else 1). + pub fn spec(name: []const u8, kind: StepKind) c.Spec { + return switch (kind) { + .gate => .{ .name = name, .accum = .none, .depth = mi.noSums, .ops = 2 * (mi.exp_ulps + mi.div_ulps) + 2, .mids = 1, .out = .bf16, .tiles = .{ .{}, .{ .size = 512, .step = 256, .min = 256 }, .{} }, .served = &.{ .{ .rows = 63, .n = ec.hidden, .k = 1 }, .{ .rows = 4096, .n = ec.hidden, .k = 1 } } }, + .silu_mul => .{ .name = name, .accum = .none, .depth = mi.noSums, .ops = 2 * (mi.exp_ulps + mi.div_ulps) + 2, .mids = 1, .out = .bf16, .tiles = .{ .{}, .{ .size = 256 }, .{} }, .served = &.{ .{ .rows = 63, .n = ec.ffn, .k = 1 }, .{ .rows = 4096, .n = ec.ffn, .k = 1 } } }, + .conv => .{ .name = name, .accum = .f32, .depth = convDepth, .ops = 2 * (mi.exp_ulps + mi.div_ulps) + 2, .mids = 1, .out = .bf16, .tiles = .{ .{ .size = 4 }, .{ .size = 256 }, .{} }, .served = &.{ .{ .rows = 63, .n = ec.mix, .k = ec.conv_taps }, .{ .rows = 10, .n = ec.mix, .k = ec.conv_taps }, .{ .rows = 4096, .n = ec.mix, .k = ec.conv_taps } } }, + }; + } + + pub fn deinit(t: *Step) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *Step) void { + t.a.free(t.x); + t.a.free(t.y); + t.x = &.{}; + t.y = &.{}; + } + pub fn kernel(t: *Step, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name, t.kind), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Step); + + pub fn prepare(t: *Step, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0) return false; + if (t.kind == .gate and shape.n % W != 0) return false; + if (t.kind == .conv and shape.k != ec.conv_taps) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + t.fill = fill; + const n = shape.rows * shape.n; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + t.x = try t.a.alloc(u16, n); + t.y = try t.a.alloc(u16, switch (t.kind) { + .gate => 2 * n, + .silu_mul => n, + .conv => shape.n * ec.conv_taps, + }); + mi.randomBf16(rng, t.x, -8, 8); + mi.randomBf16(rng, t.y, -8, 8); + if (fill == .stress) { + // the taps 256, 1, 1, -256 over rows repeating 256, 1, 1, 256: 2 at every fourth row, exact only in fp32 + for (0..shape.n) |ch| for ([_]f32{ 256, 1, 1, -256 }, 0..) |v, k| { + t.y[ch * ec.conv_taps + k] = c.toBf16(v); + }; + for (0..shape.rows) |r| @memset(t.x[r * shape.n ..][0..shape.n], c.toBf16(if (r % 4 == 0 or r % 4 == 3) 256 else 1)); + } + if (fill == .nonfinite) { + for ([_]u16{ mi.nan16, mi.inf16, mi.ninf16 }, 0..) |v, i| { + t.x[(i * 131 + 7) % n] = v; + t.y[(i * 257 + 3) % t.y.len] = v; + } + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.x)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.y)); + t.out = try t.rig.output(n * 2); + return true; + } + + pub fn run(t: *Step) !void { + const p = t.rig.begin(); + ks.encoderRows(p.e, t.pipe, kref(t.ins[0]), kref(t.ins[1]), kref(t.out), t.shape.rows, t.shape.n); + try t.rig.finish(p); + } + + pub fn count(t: *Step) usize { + return t.shape.rows * t.shape.n; + } + + /// The conv's terms for output (r, ch): the taps over rows r - 3 .. r, zero before row 0. + fn taps(t: *const Step, r: usize, ch: usize, terms: *[ec.conv_taps]f64) []f64 { + var n: usize = 0; + for (0..ec.conv_taps) |k| if (r + k >= 3) { + terms[n] = @as(f64, c.fromBf16(t.y[ch * ec.conv_taps + k])) * c.fromBf16(t.x[(r + k - 3) * t.shape.n + ch]); + n += 1; + }; + return terms[0..n]; + } + + fn convOut(sum: Range) Range { + const lo = bf(sum.lo); + const hi = bf(sum.hi); + return mi.silu(lo).bf16().join(mi.silu(hi).bf16()); + } + + pub fn reference(t: *Step, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + const w = t.shape.n; + for (0..t.count()) |i| { + const r = i / w; + const ch = i % w; + const x: f64 = c.fromBf16(t.x[i]); + switch (t.kind) { + .gate => { + const gv: f64 = c.fromBf16(t.y[r * 2 * w + (ch / W) * 2 * W + W + ch % W]); + t.ends.put(want, mag, i, bf(x * bf(1 / (1 + @exp(-gv)))), Range.of(x).mul(mi.sigmoid(gv).bf16()).fp32(0.5).bf16()); + }, + .silu_mul => { + const u: f64 = c.fromBf16(t.y[i]); + t.ends.put(want, mag, i, bf(bf(x / (1 + @exp(-x))) * u), mi.silu(x).bf16().mul(Range.of(u)).fp32(0.5).bf16()); + }, + .conv => { + var buf: [ec.conv_taps]f64 = undefined; + const terms = t.taps(r, ch, &buf); + var sum: f64 = 0; + var m: f64 = 0; + for (terms) |v| { + sum += v; + m += @abs(v); + } + const e = if (t.fill == .stress) 0 else c.gamma(@floatFromInt(ec.conv_taps), c.unit(.f32)) * m; + const cv = bf(sum); + t.ends.put(want, mag, i, bf(cv / (1 + @exp(-cv))), convOut(.{ .lo = sum - e, .hi = sum + e })); + }, + } + } + } + pub fn bounds(t: *Step, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *Step, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } + + pub fn simulate(t: *Step, acc: c.Accum, order: g.Order, out: []f64) !void { + if (t.kind != .conv) return error.NoSums; + const w = t.shape.n; + for (out, 0..) |*o, i| { + var buf: [ec.conv_taps]f64 = undefined; + const cv = bf(contract.rows.sumIn(acc, order, t.taps(i / w, i % w, &buf))); + o.* = bf(cv / (1 + @exp(-cv))); + } + } +}; diff --git a/zig/tests/contract/metal_moss_image_encoder_state.zig b/zig/tests/contract/metal_moss_image_encoder_state.zig new file mode 100644 index 0000000000..dc7bfcc5bd --- /dev/null +++ b/zig/tests/contract/metal_moss_image_encoder_state.zig @@ -0,0 +1,581 @@ +//! The encoder's me_gdn_prep, me_gdn and me_attention under the contract, against float64 within their fp32 error. +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_moss_image_common.zig"); + +const c = mi.c; +const g = mi.g; +const ks = mi.ks; +const ec = tf.moss_image.encoder_config; +const Range = mi.Range; +const Ends = mi.Ends; +const bf = mi.bf; +const kref = mi.kref; +const K = ec.k_dim; +const u = c.unit(.f32); + +fn exp(x: f64) f64 { + return @exp(x); +} +fn log(x: f64) f64 { + return @log(x); +} + +// --- me_gdn_prep: q, k L2-normed (q scaled), beta = bf16(sigmoid(b)), g = -exp(A_log) softplus(a + dt_bias) --------- + +fn prepDepth(_: c.Shape) usize { + return 4 + 5; +} + +pub const Prep = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + name: []const u8, + shape: c.Shape = undefined, + fill: g.Fill = .random, + mix: []u16 = &.{}, + /// a, then b [rows, vheads] + ab: []u16 = &.{}, + /// A_log, then dt_bias [vheads] + decay: []u16 = &.{}, + ins: [5]rig_.Ref = undefined, + outs: [4]rig_.Ref = undefined, + ends: Ends, + + /// Rows: positions; n: key heads (16, two value heads each); k: the key width. + pub fn spec(name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = prepDepth, + .ops = 2 * (mi.rsqrt_ulps + mi.exp_ulps + mi.log_ulps) + 4, + .out = .f32, + .served = &.{ .{ .rows = 63, .n = ec.k_heads, .k = K }, .{ .rows = 10, .n = ec.k_heads, .k = K }, .{ .rows = 4096, .n = ec.k_heads, .k = K } }, + }; + } + + pub fn deinit(t: *Prep) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *Prep) void { + for ([_]*[]u16{ &t.mix, &t.ab, &t.decay }) |p| { + t.a.free(p.*); + p.* = &.{}; + } + } + pub fn kernel(t: *Prep, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Prep); + + fn vheads(t: *const Prep) usize { + return 2 * t.shape.n; + } + fn stride(t: *const Prep) usize { + return 2 * t.shape.n * K + t.vheads() * K; + } + fn params(t: *const Prep) ks.Gdn { + return .{ .rows = @intCast(t.shape.rows), .kheads = @intCast(t.shape.n), .vheads = @intCast(t.vheads()), .mix_stride = @intCast(t.stride()), .scale = 1.0 / @sqrt(@as(f32, K)) }; + } + + pub fn prepare(t: *Prep, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.k != K) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + t.fill = fill; + const vh = t.vheads(); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + t.mix = try t.a.alloc(u16, shape.rows * t.stride()); + if (fill == .stress) mi.integers(rng, t.mix, 128) else mi.randomBf16(rng, t.mix, -3, 3); + t.ab = try t.a.alloc(u16, 2 * shape.rows * vh); + mi.randomBf16(rng, t.ab, -6, 6); + t.decay = try t.a.alloc(u16, 2 * vh); + mi.randomBf16(rng, t.decay[0..vh], -4, 3); + mi.randomBf16(rng, t.decay[vh..], -2, 2); + if (fill == .nonfinite) { + for ([_]u16{ mi.nan16, mi.inf16, mi.ninf16 }, 0..) |v, i| { + t.mix[i * t.stride() + i * 37] = v; + t.mix[(i + 1) * t.stride() + shape.n * K + i * 53] = v; + t.ab[(i + 2) * vh + i] = v; + t.ab[shape.rows * vh + (i + 3) * vh + i] = v; + } + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.mix)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.ab[0 .. shape.rows * vh])); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.ab[shape.rows * vh ..])); + t.ins[3] = try t.rig.input(std.mem.sliceAsBytes(t.decay[0..vh])); + t.ins[4] = try t.rig.input(std.mem.sliceAsBytes(t.decay[vh..])); + const qk = shape.rows * shape.n * K * 4; + for ([_]usize{ qk, qk, shape.rows * vh * 4, shape.rows * vh * 4 }, &t.outs) |n, *o| o.* = try t.rig.output(n); + return true; + } + + pub fn run(t: *Prep) !void { + const p = t.rig.begin(); + ks.encoderGdnPrep(p.e, t.pipe, kref(t.ins[0]), kref(t.ins[1]), kref(t.ins[2]), kref(t.ins[3]), kref(t.ins[4]), kref(t.outs[0]), kref(t.outs[1]), kref(t.outs[2]), kref(t.outs[3]), t.params()); + try t.rig.finish(p); + } + + pub fn count(t: *Prep) usize { + return 2 * t.shape.rows * t.shape.n * K + 2 * t.shape.rows * t.vheads(); + } + + /// Row r's key head h of q (0) or k (1), normed from its sum of squares `ss` (range `ssr`). + fn head(t: *Prep, ends: *Ends, which: usize, r: usize, h: usize, ss: f64, ssr: Range, want: []f64, mag: []f64) void { + const x = t.mix[r * t.stride() + which * t.shape.n * K + h * K ..][0..K]; + const inv = 1 / @sqrt(ss + 1e-6); + const invr = ssr.add(Range.of(@as(f32, 1e-6))).fp32(0.5).rsqrt().fp32(mi.rsqrt_ulps); + const scale: f64 = @as(f32, 1.0 / @sqrt(@as(f32, K))); + const base = which * t.shape.rows * t.shape.n * K + (r * t.shape.n + h) * K; + for (x, 0..) |v, i| { + const f: f64 = c.fromBf16(v); + var y = f * inv; + var yr = Range.of(f).mul(invr).fp32(0.5); + if (which == 0) { + y *= scale; + yr = yr.mul(Range.of(scale)).fp32(0.5); + } + ends.put(want, mag, base + i, y, yr); + } + } + + /// beta and g of every row and value head. + fn gates(t: *Prep, ends: *Ends, want: []f64, mag: []f64) void { + const vh = t.vheads(); + const n = t.shape.rows * vh; + const at = 2 * t.shape.rows * t.shape.n * K; + for (0..n) |i| { + const v = i % vh; + const b: f64 = c.fromBf16(t.ab[n + i]); + ends.put(want, mag, at + n + i, bf(1 / (1 + @exp(-b))), mi.sigmoid(b).bf16()); + const x: f64 = @as(f32, @floatCast(@as(f64, c.fromBf16(t.ab[i])) + c.fromBf16(t.decay[vh + v]))); + const a_log: f64 = c.fromBf16(t.decay[v]); + const soft = if (x > 20) x else @log(1 + @exp(x)); + const xr = Range.of(@as(f64, c.fromBf16(t.ab[i])) + c.fromBf16(t.decay[vh + v])).fp32(0.5); + const sr = if (x > 20) xr else xr.up(exp).fp32(mi.exp_ulps).add(Range.of(1)).fp32(0.5).up(log).fp32(mi.log_ulps); + const er = Range.of(@exp(a_log)).fp32(mi.exp_ulps); + ends.put(want, mag, at + i, -@exp(a_log) * soft, (Range{ .lo = -er.hi, .hi = -er.lo }).mul(sr).fp32(0.5)); + } + } + + pub fn reference(t: *Prep, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + for (0..2) |which| for (0..t.shape.rows) |r| for (0..t.shape.n) |h| { + const s = mi.squares(t.mix[r * t.stride() + which * t.shape.n * K + h * K ..][0..K], prepDepth(t.shape), t.fill == .stress); + t.head(&t.ends, which, r, h, s.ss, s.r, want, mag); + }; + t.gates(&t.ends, want, mag); + } + pub fn bounds(t: *Prep, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *Prep, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } + + pub fn simulate(t: *Prep, 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: [K]f64 = undefined; + for (0..2) |which| for (0..t.shape.rows) |r| for (0..t.shape.n) |h| { + for (&terms, t.mix[r * t.stride() + which * t.shape.n * K + h * K ..][0..K]) |*v, x| { + const f: f64 = c.fromBf16(x); + v.* = f * f; + } + const ss = contract.rows.sumIn(acc, order, &terms); + t.head(&scratch, which, r, h, ss, Range.of(ss), out, mag); + }; + t.gates(&scratch, out, mag); + } +}; + +// --- me_gdn: S = S exp(g); d = (v - S^T k) beta; S += k d^T; o = S^T q, over the rows in order ----------------------- + +/// A step's two 128-term sums, the update and the decay, every row before it. +fn gdnDepth(s: c.Shape) usize { + return s.rows * (2 * K + 8); +} + +pub const Gdn = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + name: []const u8, + shape: c.Shape = undefined, + stress: bool = false, + /// q then k, fp32 [rows, kheads, 128] + qk: []f32 = &.{}, + mix: []u16 = &.{}, + /// g then beta, fp32 [rows, vheads] + gb: []f32 = &.{}, + ins: [5]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + /// Rows: positions; n: value heads (two a key head; each its own threadgroup, so 4 stand for 32); k: the width. + pub fn spec(name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = gdnDepth, + .ops = 2 * mi.exp_ulps, + .out = .bf16, + .served = &.{ .{ .rows = 63, .n = ec.v_heads, .k = K }, .{ .rows = 10, .n = ec.v_heads, .k = K }, .{ .rows = 4096, .n = 4, .k = K } }, + }; + } + + pub fn deinit(t: *Gdn) void { + t.free(); + t.rig.deinit(); + } + fn free(t: *Gdn) void { + t.a.free(t.qk); + t.a.free(t.mix); + t.a.free(t.gb); + t.qk = &.{}; + t.mix = &.{}; + t.gb = &.{}; + } + pub fn kernel(t: *Gdn, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Gdn); + + fn kheads(t: *const Gdn) usize { + return t.shape.n / 2; + } + fn stride(t: *const Gdn) usize { + return 2 * t.kheads() * K + t.shape.n * K; + } + fn params(t: *const Gdn) ks.Gdn { + return .{ .rows = @intCast(t.shape.rows), .kheads = @intCast(t.kheads()), .vheads = @intCast(t.shape.n), .mix_stride = @intCast(t.stride()), .scale = 1 }; + } + + pub fn prepare(t: *Gdn, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.n % 2 != 0 or shape.k != K) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + t.stress = fill == .stress; + const T = shape.rows; + const kh = t.kheads(); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + t.qk = try t.a.alloc(f32, 2 * T * kh * K); + t.mix = try t.a.alloc(u16, T * t.stride()); + t.gb = try t.a.alloc(f32, 2 * T * shape.n); + mi.randomBf16(rng, t.mix, -2, 2); + if (t.stress) { + // no decay, beta 1, one-hot keys: columns 0, 1, 2 hold j, 2^20, -2^20; q sums them back to j, exact in fp32 + @memset(t.qk, 0); + @memset(t.gb[0 .. T * shape.n], 0); + @memset(t.gb[T * shape.n ..], 1); + for (0..T) |r| for (0..kh) |h| { + const at = (r * kh + h) * K; + for (0..3) |i| t.qk[at + i] = 1; + t.qk[T * kh * K + at + r % 3] = 1; + }; + for (0..T) |r| for (0..shape.n) |h| for (0..K) |j| { + const v: f32 = switch (r % 3) { + 0 => @floatFromInt(j), + 1 => 0x1p20, + else => -0x1p20, + }; + t.mix[r * t.stride() + 2 * kh * K + h * K + j] = c.toBf16(v); + }; + } else { + // unit keys and queries (the prep's output), decays in (0, 1), beta in (0, 1) + var row: usize = 0; + while (row < 2 * T * kh) : (row += 1) { + const x = t.qk[row * K ..][0..K]; + var ss: f32 = 0; + for (x) |*v| { + v.* = rng.float(f32) * 2 - 1; + ss += v.* * v.*; + } + for (x) |*v| v.* /= @sqrt(ss); + } + for (t.gb[0 .. T * shape.n]) |*v| v.* = -rng.float(f32) * 2; + for (t.gb[T * shape.n ..]) |*v| v.* = c.fromBf16(c.toBf16(rng.float(f32))); + } + if (fill == .nonfinite) { + t.mix[2 * t.stride() + 2 * kh * K + 5] = mi.nan16; + t.mix[3 * t.stride() + 2 * kh * K + K + 9] = mi.inf16; + t.qk[(4 * kh + 0) * K + 7] = std.math.inf(f32); + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.qk[0 .. T * kh * K])); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.qk[T * kh * K ..])); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.mix)); + t.ins[3] = try t.rig.input(std.mem.sliceAsBytes(t.gb[0 .. T * shape.n])); + t.ins[4] = try t.rig.input(std.mem.sliceAsBytes(t.gb[T * shape.n ..])); + t.out = try t.rig.output(T * shape.n * K * 2); + return true; + } + + pub fn run(t: *Gdn) !void { + const p = t.rig.begin(); + ks.encoderGdn(p.e, t.pipe, kref(t.ins[0]), kref(t.ins[1]), kref(t.ins[2]), kref(t.ins[3]), kref(t.ins[4]), kref(t.out), t.params()); + try t.rig.finish(p); + } + + pub fn count(t: *Gdn) usize { + return t.shape.rows * t.shape.n * K; + } + + /// The recurrence in float64 (sums in `acc`, `order` if given); each allowance carries every earlier step's error. + fn recur(t: *Gdn, want: []f64, allow: ?[]f64, sums: ?struct { acc: c.Accum, order: g.Order }) !void { + const T = t.shape.rows; + const vh = t.shape.n; + const kh = t.kheads(); + const s = try t.a.alloc(f64, 2 * K * K); + defer t.a.free(s); + const st = s[0 .. K * K]; + const mag = s[K * K ..]; + var err: [K]f64 = undefined; + var terms: [K]f64 = undefined; + for (0..vh) |h| { + @memset(s, 0); + @memset(&err, 0); + for (0..T) |r| { + const q = t.qk[(r * kh + h / 2) * K ..][0..K]; + const k = t.qk[T * kh * K + (r * kh + h / 2) * K ..][0..K]; + const decay = @exp(@as(f64, t.gb[r * vh + h])); + const beta: f64 = t.gb[T * vh + r * vh + h]; + var qn: f64 = 0; + var kn: f64 = 0; + for (q, k) |a, b| { + qn += @as(f64, a) * a; + kn += @as(f64, b) * b; + } + for (0..K) |j| { + const v: f64 = c.fromBf16(t.mix[r * t.stride() + 2 * kh * K + h * K + j]); + var kv: f64 = 0; + var kv_mag: f64 = 0; + for (0..K) |i| { + st[i * K + j] *= decay; + mag[i * K + j] *= decay; + terms[i] = st[i * K + j] * k[i]; + kv_mag += mag[i * K + j] * @abs(k[i]); + } + kv = if (sums) |x| contract.rows.sumIn(x.acc, x.order, &terms) else blk: { + var sum: f64 = 0; + for (terms) |v2| sum += v2; + break :blk sum; + }; + const delta = (v - kv) * beta; + const delta_mag = (@abs(v) + kv_mag) * beta; + var col: f64 = 0; + var o_mag: f64 = 0; + for (0..K) |i| { + st[i * K + j] += k[i] * delta; + mag[i * K + j] += @abs(k[i]) * delta_mag; + col += mag[i * K + j] * mag[i * K + j]; + terms[i] = st[i * K + j] * q[i]; + o_mag += mag[i * K + j] * @abs(q[i]); + } + const o = if (sums) |x| contract.rows.sumIn(x.acc, x.order, &terms) else blk: { + var sum: f64 = 0; + for (terms) |v2| sum += v2; + break :blk sum; + }; + const at = (r * vh + h) * K + j; + want[at] = o; + if (allow) |al| { + // this step's error in the column (2-norm): the decay, the key sum, the delta and the update + err[j] += (2 * mi.exp_ulps * u + 3 * u) * @sqrt(col) + c.gamma(K, u) * kv_mag * beta * @sqrt(kn) + u * delta_mag * @sqrt(kn); + const work = err[j] * @sqrt(qn) + c.gamma(K, u) * o_mag; + al[at] = if (std.math.isFinite(o)) work + 0x1p-8 * (@abs(o) + work) + 0x1p-133 else 0; + } + } + } + } + } + + pub fn reference(t: *Gdn, want: []f64, mag: []f64) !void { + try t.recur(want, mag, null); + } + pub fn bounds(t: *Gdn, 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), w) else tl; + } + pub fn simulate(t: *Gdn, acc: c.Accum, order: g.Order, out: []f64) !void { + if (!t.stress) return error.GdnStressNeedsExactSums; + try t.recur(out, null, .{ .acc = acc, .order = order }); + } +}; + +// --- me_attention: causal softmax(q k^T / 16) v in fp32 over each row's keys, against attention_ref ----------------- + +fn attentionDepth(s: c.Shape) usize { + return ec.head_dim + s.k; +} + +pub const Attention = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + name: []const u8, + shape: c.Shape = undefined, + stress: bool = false, + q: []u16 = &.{}, + k: []u16 = &.{}, + v: []u16 = &.{}, + ins: [3]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + /// Rows: query heads (four a key head; each its own threadgroups, so 8 stand for 16); n: head width; k: rows L. + pub fn spec(name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = attentionDepth, + .ops = 8, + .out = .bf16, + .served = &.{ .{ .rows = 8, .n = ec.head_dim, .k = 63 }, .{ .rows = 8, .n = ec.head_dim, .k = 10 }, .{ .rows = 8, .n = ec.head_dim, .k = 34 }, .{ .rows = 8, .n = ec.head_dim, .k = 4096 } }, + }; + } + + 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), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Attention); + + fn kvHeads(t: *const Attention) usize { + return t.shape.rows / 4; + } + + pub fn prepare(t: *Attention, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.rows % 4 != 0 or shape.n != ec.head_dim or shape.k == 0) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + t.stress = fill == .stress; + const L = shape.k; + const D = ec.head_dim; + t.q = try t.a.alloc(u16, shape.rows * L * D); + t.k = try t.a.alloc(u16, t.kvHeads() * L * D); + t.v = try t.a.alloc(u16, t.kvHeads() * L * D); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (t.stress) { + // equal scores; each column's values partial sums reaching ~2^22 that cancel to 4 over the whole row + @memset(t.q, 0); + @memset(t.k, 0); + const half = L / 2; + for (0..t.kvHeads()) |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], c.toBf16(v)); + }; + } else { + mi.randomBf16(rng, t.q, -1, 1); + mi.randomBf16(rng, t.k, -1, 1); + mi.randomBf16(rng, t.v, -1, 1); + } + if (fill == .nonfinite) { + for ([_][2]usize{ .{ 0, 0 }, .{ 1, 1 }, .{ 2, 2 }, .{ 3, 3 }, .{ 4, 3 } }, [_]u16{ mi.nan16, mi.inf16, mi.ninf16, mi.inf16, mi.ninf16 }) |p, v| if (p[0] < L) { + t.v[p[0] * D + p[1]] = v; + }; + if (L > 2) t.q[(L + 2) * D + 5] = mi.nan16; + } + 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(L * shape.rows * D * 2); + return true; + } + + pub fn run(t: *Attention) !void { + const p = t.rig.begin(); + ks.encoderAttention(p.e, t.pipe, kref(t.ins[0]), kref(t.ins[1]), kref(t.ins[2]), kref(t.out), t.shape.k, t.shape.rows, t.kvHeads(), 1.0 / @sqrt(@as(f32, ec.head_dim))); + try t.rig.finish(p); + } + + pub fn count(t: *Attention) usize { + return t.shape.k * t.shape.rows * ec.head_dim; + } + + /// attention_ref.head over each row's keys 0..i; O in [L, heads, 256]. + pub fn reference(t: *Attention, want: []f64, mag: []f64) !void { + const L = t.shape.k; + const D = ec.head_dim; + const heads = t.shape.rows; + const kf = try t.a.alloc(f64, 2 * L * D); + defer t.a.free(kf); + const vf = kf[L * D ..]; + const keys = try t.a.alloc(usize, L); + defer t.a.free(keys); + for (keys, 0..) |*k, i| k.* = i; + var qf: [D]f64 = undefined; + const scale: f64 = @as(f32, 1.0 / @sqrt(@as(f32, D))); + for (0..t.kvHeads()) |kh| { + for (kf[0 .. L * D], vf, t.k[kh * L * D ..][0 .. L * D], t.v[kh * L * D ..][0 .. L * D]) |*a, *b, x, y| { + a.* = c.fromBf16(x); + b.* = c.fromBf16(y); + } + if (t.stress) { + // q = k = 0: every score 0, so each row is the mean of the values so far (the same in every column) + var sum: f64 = 0; + for (0..L) |i| { + sum += vf[i * D]; + for (kh * 4..kh * 4 + 4) |h| { + @memset(want[(i * heads + h) * D ..][0..D], sum / @as(f64, @floatFromInt(i + 1))); + @memset(mag[(i * heads + h) * D ..][0..D], 0); + } + } + continue; + } + for (kh * 4..kh * 4 + 4) |h| for (0..L) |i| { + for (&qf, t.q[(h * L + i) * D ..][0..D]) |*x, b| x.* = c.fromBf16(b); + const at = (i * heads + h) * D; + tf.attention_ref.head(&qf, kf[0 .. L * D], vf, keys[0 .. i + 1], D, scale, want[at..][0..D], mag[at..][0..D]); + }; + } + } + + 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), 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.shape.k; + const D = ec.head_dim; + 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); + // every column of a key head holds the same values: one column's sums stand for all, and for its query heads + for (0..t.kvHeads()) |kh| { + for (terms, 0..) |*v, j| v.* = c.fromBf16(t.v[(kh * L + j) * D]); + for (0..L) |i| { + const o = contract.rows.sumIn(acc, order, terms[0 .. i + 1]) / contract.rows.sumIn(acc, order, ones[0 .. i + 1]); + for (kh * 4..kh * 4 + 4) |h| @memset(out[(i * t.shape.rows + h) * D ..][0..D], o); + } + } + } +}; diff --git a/zig/tests/contract/metal_moss_image_norms.zig b/zig/tests/contract/metal_moss_image_norms.zig index 6f64b519bc..475d6f4e92 100644 --- a/zig/tests/contract/metal_moss_image_norms.zig +++ b/zig/tests/contract/metal_moss_image_norms.zig @@ -1,6 +1,7 @@ -//! mi_norm and mi_qk under the contract: RMS norms at the publisher's bf16 points, their fp32 sums proven in gate 4. +//! mi_norm, me_norm and mi_qk under the contract: RMS norms at the publisher's bf16 points, fp32 sums proven in gate 4. 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_moss_image_common.zig"); @@ -16,7 +17,7 @@ const bf = mi.bf; const kref = mi.kref; const kat = mi.kat; -// --- mi_norm: bf16(bf16(x rsqrt(mean x^2 + eps)) w), then bf16(bf16(y s) + b) when modulated -- +// --- mi_norm: bf16(bf16(x rsqrt(mean x^2 + eps)) w), then bf16(bf16(y s) + b) when modulated; me_norm: one rounding -- const norm_why = "a bf16 output follows the sum of squares as ss^-1/2: a bf16 sum's error is a few output ulps"; @@ -25,11 +26,14 @@ fn normDepth(s: c.Shape) usize { return (s.k + 255) / 256 + 10; } +/// mi_norm plain or modulated, or the encoder's me_norm: bf16(x rsqrt(mean x^2 + eps) (1 + w)). +pub const NormKind = enum { plain, modulated, zero_centered }; + pub const Norm = struct { a: std.mem.Allocator, rig: rig_.Rig, pipe: mtl.Pipeline, - modulated: bool, + kind: NormKind, name: []const u8, shape: c.Shape = undefined, fill: g.Fill = .random, @@ -40,19 +44,24 @@ pub const Norm = struct { out: rig_.Ref = undefined, ends: Ends, - pub fn spec(name: []const u8, modulated: bool) c.Spec { + pub fn spec(name: []const u8, kind: NormKind) c.Spec { return .{ .name = name, .accum = .f32, .depth = normDepth, .ops = 2 * (mi.div_ulps + mi.rsqrt_ulps) + 1, - .mids = if (modulated) 3 else 1, + .mids = switch (kind) { + .plain => 1, + .modulated => 3, + .zero_centered => 0, + }, .out = .bf16, .tiles = .{ .{}, .{}, .{ .size = 256 } }, - .served = if (modulated) - &.{ .{ .rows = 4188, .n = 3072, .k = 3072 }, .{ .rows = 4130, .n = 3072, .k = 3072 }, .{ .rows = 4125, .n = 3072, .k = 3072 } } - else - &.{ .{ .rows = 63, .n = 3072, .k = 3072 }, .{ .rows = 34, .n = 3072, .k = 3072 }, .{ .rows = 10, .n = 3072, .k = 3072 } }, + .served = switch (kind) { + .modulated => &.{ .{ .rows = 4188, .n = 3072, .k = 3072 }, .{ .rows = 4130, .n = 3072, .k = 3072 }, .{ .rows = 4125, .n = 3072, .k = 3072 } }, + .plain => &.{ .{ .rows = 63, .n = 3072, .k = 3072 }, .{ .rows = 34, .n = 3072, .k = 3072 }, .{ .rows = 10, .n = 3072, .k = 3072 } }, + .zero_centered => &.{ .{ .rows = 63, .n = 4096, .k = 4096 }, .{ .rows = 34, .n = 4096, .k = 4096 }, .{ .rows = 10, .n = 4096, .k = 4096 }, .{ .rows = 4096, .n = 4096, .k = 4096 } }, + }, .sums = .{ .bf16 = 0, .why = norm_why }, }; } @@ -69,7 +78,7 @@ pub const Norm = struct { t.v = &.{}; } pub fn kernel(t: *Norm, control: ?g.Gate) g.Kernel { - return .{ .spec = spec(t.name, t.modulated), .ctx = t, .vt = &vt, .control = control }; + return .{ .spec = spec(t.name, t.kind), .ctx = t, .vt = &vt, .control = control }; } const vt = g.hooks(Norm); @@ -85,7 +94,7 @@ pub const Norm = struct { var prng = std.Random.DefaultPrng.init(seed); const rng = prng.random(); if (fill == .stress) mi.integers(rng, t.x, 32) else mi.randomBf16(rng, t.x, -3, 3); - mi.randomBf16(rng, t.v[0..k], 0.25, 2); + if (t.kind == .zero_centered) mi.randomBf16(rng, t.v[0..k], -0.75, 1) else mi.randomBf16(rng, t.v[0..k], 0.25, 2); // the modulation's 1 + scale, and its shift mi.randomBf16(rng, t.v[k .. 2 * k], 0.5, 1.5); mi.randomBf16(rng, t.v[2 * k ..], -1, 1); @@ -108,7 +117,11 @@ pub const Norm = struct { pub fn run(t: *Norm) !void { const p = t.rig.begin(); - ks.norm(p.e, t.pipe, kref(t.ins[0]), kref(t.ins[1]), kref(t.ins[2]), kref(t.ins[3]), kref(t.out), t.shape.rows, t.shape.k, cfg.eps, t.modulated); + const k = t.shape.k; + if (t.kind == .zero_centered) + ks.encoderNorm(p.e, t.pipe, kref(t.ins[0]), kref(t.ins[1]), kref(t.out), t.shape.rows, k, k, k, tf.moss_image.encoder_config.eps) + else + ks.norm(p.e, t.pipe, kref(t.ins[0]), kref(t.ins[1]), kref(t.ins[2]), kref(t.ins[3]), kref(t.out), t.shape.rows, k, cfg.eps, t.kind == .modulated); try t.rig.finish(p); } @@ -122,9 +135,14 @@ pub const Norm = struct { for (0..k) |j| { const x: f64 = c.fromBf16(t.x[r * k + j]); const w: f64 = c.fromBf16(t.v[j]); + if (t.kind == .zero_centered) { + const r1 = Range.of(x).mul(n.r).fp32(0.5).mul(Range.of(1 + w).fp32(0.5)).fp32(0.5).bf16(); + ends.put(want, mag, r * k + j, bf(x * n.inv * (1 + w)), r1); + continue; + } var y = bf(bf(x * n.inv) * w); var yr = Range.of(x).mul(n.r).fp32(0.5).bf16().mul(Range.of(w)).fp32(0.5).bf16(); - if (t.modulated) { + if (t.kind == .modulated) { const s: f64 = c.fromBf16(t.v[k + j]); const b: f64 = c.fromBf16(t.v[2 * k + j]); y = bf(bf(y * s) + b); diff --git a/zig/tests/moss_image.zig b/zig/tests/moss_image.zig index af5cbccf07..c9eb95985b 100644 --- a/zig/tests/moss_image.zig +++ b/zig/tests/moss_image.zig @@ -3,14 +3,17 @@ const std = @import("std"); const mtl = @import("metal"); const tf = @import("tensorfold"); const mi = tf.moss_image; +const tokenizer = @import("tokenizer"); const usage = - \\tf-moss-image TRANSFORMER_DIR CASE OUT_DIR + \\tf-moss-image TRANSFORMER_DIR CASE OUT_DIR [--prompt FILE [--negative FILE]] \\CASE (tools/zig/moss_image_case.py) holds the prompt's and the empty negative prompt's encoder rows, the timesteps \\and noise levels and the starting latents. A `case` also holds a reference run (each step's input and both \\branches' velocities): each step then runs both branches on the reference's input and compares them (max |diff|, \\cosine, worst row). Then the whole schedule runs from the starting latents with the card's guidance (scale 3, the \\negative branch skipping block 22), timed, and OUT_DIR gets latents.f32 (rows x columns x 128) and run.json. + \\With --prompt, the encoder in TRANSFORMER_DIR's checkpoint folder makes both prompts' rows (the negative from + \\--negative's file, else empty) in place of the case's. ; /// The card's sampling defaults: the guidance scale, and the block the negative branch skips. @@ -113,11 +116,22 @@ fn teacher(a: std.mem.Allocator, case: *const tf.safetensors.File, model: *mi.Mo return repeat; } +/// A prompt file's text, trimmed as the case tool trims it. +fn read(a: std.mem.Allocator, io: std.Io, path: []const u8) ![]const u8 { + return std.mem.trim(u8, try std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(1 << 20)), " \t\r\n"); +} + 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 != 4) { + var prompt_file: ?[]const u8 = null; + var negative_file: ?[]const u8 = null; + var i: usize = 4; + while (i + 1 < args.len) : (i += 2) { + if (std.mem.eql(u8, args[i], "--prompt")) prompt_file = args[i + 1] else if (std.mem.eql(u8, args[i], "--negative")) negative_file = args[i + 1] else break; + } + if (args.len < 4 or i != args.len or (negative_file != null and prompt_file == null)) { std.debug.print("{s}\n", .{usage}); return 2; } @@ -127,10 +141,30 @@ pub fn main(init: std.process.Init) !u8 { if (grid.len != 3 or grid[2] != mi.config.channels) return error.BadCaseTensor; const rows = grid[0] * grid[1]; const count = rows * mi.config.channels; - const text = (try dims(&case, "prompt"))[0]; - const negative_rows = (try dims(&case, "negative"))[0]; - const prompt = try take(u16, a, &case, "prompt", .bf16); - const negative = try take(u16, a, &case, "negative", .bf16); + var encode_s: ?f64 = null; + const both: [2][]u16 = if (prompt_file) |file| blk: { + const texts = [2][]const u8{ try read(a, io, file), if (negative_file) |n| try read(a, io, n) else "" }; + const device = try mtl.Device.init(); + const dir = std.fs.path.dirname(args[1]) orelse "."; + const enc = try mi.Encoder.load(std.heap.page_allocator, io, device, dir); + defer enc.deinit(); + var tok = try tokenizer.loadTokenizer(io, a, try std.fs.path.join(a, &.{ dir, "tokenizer", "tokenizer.json" })); + defer tok.deinit(); + const encoding = mtl.clock.seconds(); + var made: [2][]u16 = undefined; + for (texts, &made) |t, *m| { + const ids = try mi.encoder.ids(&tok, a, t); + m.* = try a.alloc(u16, ids.len * mi.encoder_config.hidden); + try enc.rows(ids, m.*, null); + } + encode_s = mtl.clock.seconds() - encoding; + std.debug.print("both prompts encoded natively: {d:.3} s\n", .{encode_s.?}); + break :blk made; + } else .{ try take(u16, a, &case, "prompt", .bf16), try take(u16, a, &case, "negative", .bf16) }; + const prompt = both[0]; + const negative = both[1]; + const text = prompt.len / mi.config.context; + const negative_rows = negative.len / mi.config.context; const timesteps = try take(f32, a, &case, "timesteps", .f32); const sigmas = try take(f32, a, &case, "sigmas", .f32); const start = try take(f32, a, &case, "latents", .f32); @@ -156,6 +190,7 @@ pub fn main(init: std.process.Init) !u8 { const refine_s = mtl.clock.seconds() - refining; std.debug.print("both prompts refined: {d:.3} s\n", .{refine_s}); try print(&out, a, "{{\"device\": \"{s}\", \"image_rows\": {d}, \"prompt_rows\": {d}, \"negative_rows\": {d}, \"load_s\": {d:.2}, \"refine_s\": {d:.4}", .{ device.name(), rows, text, negative_rows, load_s, refine_s }); + if (encode_s) |s| try print(&out, a, ", \"native_encoder\": true, \"encode_s\": {d:.4}", .{s}); const repeat = if (case.get("inputs") == null) true else try teacher(a, &case, model, &branches, timesteps, count, &out); try print(&out, a, ", \"run\": [", .{}); diff --git a/zig/tests/moss_image_encoder.zig b/zig/tests/moss_image_encoder.zig new file mode 100644 index 0000000000..42ba5d3692 --- /dev/null +++ b/zig/tests/moss_image_encoder.zig @@ -0,0 +1,115 @@ +//! tf-moss-image-encoder: MOSS-Image's prompt encoder on a case from tools/zig/moss_image_case.py `encode`, compared. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const mi = tf.moss_image; +const tokenizer = @import("tokenizer"); + +const usage = + \\tf-moss-image-encoder MODEL_DIR CASE [PROMPT_FILE] + \\MODEL_DIR is the OpenMOSS-Team/MOSS-Image-Base folder. CASE (moss_image_case.py encode) holds the prompt's token + \\ids, every layer's rows and the final rows from the publisher's torch path in bf16, and the final rows in float32. + \\The prompt (PROMPT_FILE, else the empty negative prompt) is tokenized and checked against the ids, then encoded: + \\each layer's rows and the final ones are compared (max |diff|, cosine, worst row) and the encoding is timed. +; + +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`. +fn match(ours: []const u16, theirs: []const f32, width: usize) Match { + 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]) |h, b| { + const a: f64 = @as(f32, @bitCast(@as(u32, h) << 16)); + const d = a - b; + dot += a * b; + aa += 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(aa * bb), .worst_row = worst }; +} + +/// The case tensor `name` as f32 (from bf16 or f32). +fn floats(a: std.mem.Allocator, f: *const tf.safetensors.File, name: []const u8) ![]f32 { + const t = f.get(name) orelse { + std.debug.print("the case has no {s}\n", .{name}); + return error.MissingTensor; + }; + const out = try a.alloc(f32, t.numel()); + switch (t.dtype) { + .f32 => @memcpy(std.mem.sliceAsBytes(out), t.bytes), + .bf16 => for (out, std.mem.bytesAsSlice(u16, t.bytes)) |*o, h| { + o.* = @bitCast(@as(u32, h) << 16); + }, + else => return error.BadCaseTensor, + } + return out; +} + +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, args[2]); + defer case.close(io); + const prompt = if (args.len == 4) std.mem.trim(u8, try std.Io.Dir.cwd().readFileAlloc(io, args[3], a, .limited(1 << 20)), " \t\r\n") else ""; + const want_ids = case.get("ids") orelse return error.MissingTensor; + if (want_ids.dtype != .i32) return error.BadCaseTensor; + + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + const loading = mtl.clock.seconds(); + const enc = try mi.Encoder.load(std.heap.page_allocator, io, device, args[1]); + defer enc.deinit(); + const load_s = mtl.clock.seconds() - loading; + var tok = try tokenizer.loadTokenizer(io, a, try std.fs.path.join(a, &.{ args[1], "tokenizer", "tokenizer.json" })); + defer tok.deinit(); + const ids = try mi.encoder.ids(&tok, a, prompt); + const same_ids = std.mem.eql(u8, std.mem.sliceAsBytes(ids), want_ids.bytes); + std.debug.print("{s}: {d} tokens, the same ids as the tokenizer's: {}; loaded in {d:.1} s\n", .{ device.name(), ids.len, same_ids, load_s }); + if (!same_ids) return 1; + + const H = mi.encoder_config.hidden; + const L = mi.encoder_config.layers; + const T = ids.len; + const out = try a.alloc(u16, T * H); + const kept = try a.alloc(u16, L * T * H); + try enc.rows(ids, out, kept); + var worst: f64 = 1; + for (0..L) |i| { + var name: [32]u8 = undefined; + const ref = try floats(a, &case, try std.fmt.bufPrint(&name, "layer.{d}", .{i})); + const m = match(kept[i * T * H ..][0 .. T * H], ref, H); + worst = @min(worst, m.cosine); + std.debug.print("layer {d:>2} ({s}): max |diff| {d:.4}, cosine {d:.6}, worst row {d:.5}\n", .{ i, if (mi.encoder_config.isFull(i)) "attention" else "delta ", m.max_abs, m.cosine, m.worst_row }); + } + const bf = match(out, try floats(a, &case, "final"), H); + const fp = match(out, try floats(a, &case, "final_fp32"), H); + std.debug.print("final rows against torch bf16: max |diff| {d:.4}, cosine {d:.6}, worst row {d:.5}\n", .{ bf.max_abs, bf.cosine, bf.worst_row }); + std.debug.print("final rows against float32: max |diff| {d:.4}, cosine {d:.6}, worst row {d:.5}\n", .{ fp.max_abs, fp.cosine, fp.worst_row }); + + const again = try a.alloc(u16, T * H); + const timing = mtl.clock.seconds(); + try enc.rows(ids, again, null); + const encode_s = mtl.clock.seconds() - timing; + const repeat = std.mem.eql(u16, out, again); + std.debug.print("encoded again in {d:.3} s, the same bytes: {}\n", .{ encode_s, repeat }); + return if (repeat) 0 else 1; +} From 6dc591433512d359131d4c07a3d9aebfae7ff93e Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Tue, 27 Oct 2026 08:57:04 -0500 Subject: [PATCH 2/3] docs: MOSS-Image's prompt encoder on Metal, how it was checked and measured --- docs/recipes/moss-image.md | 120 ++++++++++++++++++++++++++++++------- 1 file changed, 99 insertions(+), 21 deletions(-) diff --git a/docs/recipes/moss-image.md b/docs/recipes/moss-image.md index 5680d3e0e1..eb0326dbb8 100644 --- a/docs/recipes/moss-image.md +++ b/docs/recipes/moss-image.md @@ -11,9 +11,9 @@ transformer on the Metal runtime from the released bf16 diffusers checkpoint: The card's guidance runs two passes a step: the prompt, and the empty negative prompt with block 22 skipped. -The prompt encoder (a Qwen3.5 model) and the image decoder (FLUX.2's small decoder) are not ported. -`tools/zig/moss_image_case.py` runs them through the publisher's own `pipeline_moss_image.py` in torch, a development -reference. +The prompt encoder, the checkpoint's Qwen3.5 text model, runs on Metal too (below): `tf-moss-image --prompt` makes both +prompts' rows itself. The image decoder (FLUX.2's small decoder) is not ported. `tools/zig/moss_image_case.py` runs it, +and the reference pipeline, through the publisher's own `pipeline_moss_image.py` in torch. It is a denoiser, not a token family. It does not register with the lane engine, the server or `families.detect`. @@ -38,26 +38,81 @@ tensors this family does not read. The bf16 projections are the Qwen-VL tower's `tfp::gemm_nax`. Attention is the prefill `tfp::attention_nax` template's full (non-causal) form. The time MLP and both modulation projections run on the host, once a timestep. +## Prompt encoder + +`encoder.zig` runs the checkpoint's `text_encoder/` (Qwen3.5's text model: 32 layers, three Gated DeltaNet layers to +each full-attention one, hidden 4,096) as the pipeline's `encode_prompt` does: + +- the prompt in the pipeline's chat format (one user turn, then the assistant's turn with thinking on), cut at 4,096 + tokens; +- one causal forward; +- the final norm's rows. + +The tokenizer is the native one on the checkpoint's `tokenizer/tokenizer.json`, and its ids equal the publisher's +tokenizer's on every prompt checked. The vision tower is not read. + +Precision: + +- bf16 weights as published and bf16 activations, rounded where the torch code rounds them; +- fp32 sums in every projection (the tower's `gemm_nax`), norm and attention; +- the DeltaNet's q, k, decay and state in fp32, its state a column a thread over the rows in order (the publisher's + chunked algorithm computes the same recurrence); +- attention in fp32 throughout, the keys in order. + +Its own kernels are in `moss_image_encoder.metal`: + +- `me_norm`: the zero-centered RMSNorm; +- `me_head`: the head norms, the 64 rotated channels and the heads' layout; +- `me_attention`; +- `me_gate`; +- `me_silu_mul`; +- `me_conv`: the DeltaNet's 4-tap conv; +- `me_gdn_prep`; +- `me_gdn`: the recurrence; +- `me_gated_norm`. + +The residual adds reuse `mi_gate_add`. + +`tf-moss-image-encoder` checks a case from `moss_image_case.py encode`: + +- the token ids; +- every layer's rows against the publisher's torch path in bf16; +- the final rows against bf16 and float32. + ## Kernel contract -`zig/tests/contract/metal_moss_image*.zig` registers every kernel of `moss_image.metal`: +`zig/tests/contract/metal_moss_image*.zig` registers every kernel of `moss_image.metal` and +`moss_image_encoder.metal`: -- `mi_norm`, plain and modulated; +- `mi_norm`, plain and modulated, and `me_norm`; - `mi_qk`, in both row layouts: q, k and v in three buffers, and packed in one row; +- `me_head` for q, k and v; - `mi_situ`, in both layouts: hidden and gate halves of one row, and two buffers; -- `mi_gate_add` (in place) and `mi_bias`; -- the attention's three forms. +- `mi_gate_add` (in place), `mi_bias`, `me_gate`, `me_silu_mul`, `me_conv` and `me_gated_norm`; +- `me_gdn_prep` and `me_gdn`; +- the attentions: the transformer's three forms and `me_attention`. The steps are checked output by output at the bf16 points the publisher's torch code rounds at. Each output must -lie in the range its fp32 steps allow, with Metal's error bounds for the `precise::` calls and fast-math division. -Eight controls each drop or add one rounding point, and gate 3 must catch each. The attention is checked against -`attention_ref`'s float64 formula and bf16-P bound. Its gate 4 uses equal scores over values whose sums reach about -2^22 and cancel. Gate 5 holds the source's eight launches and the two split-K projections. +lie in the range its fp32 steps allow, with Metal's error bounds for the `precise::` calls and fast-math division. Each +of 16 controls drops or adds one rounding point (the last adds bf16 rounding to the DeltaNet's state), and gate 3 must +catch each. + +`me_gdn` is checked against the recurrence in float64. Each output's allowance is its two sums' rounding over the +state's magnitude, plus every earlier step's error: a decay and a projection by a unit key never grow it. The +attentions are checked against `attention_ref`'s float64 formula and bf16-P bound. + +Gate 4: + +- the attentions: equal scores over values whose sums reach about 2^22 and cancel; +- `me_gdn`: one-hot keys writing j, 2^20 and -2^20 to three state rows that its query sums back to j; +- `me_conv`: taps of 256, 1, 1 and -256 over rows that cancel to 2. + +Gate 5 holds the 17 launches of the family's source and the two split-K projections. `tf-kernel-contract --only "moss image"` on the M5 Ultra: ``` -SUMMARY 11 kernels: 11 pass gates 1-4, 0 fail, 0 known findings; 8 controls: 8 red as planted, 0 missed; 0 limit problems +SUMMARY 22 kernels: 22 pass gates 1-4, 0 fail, 0 known findings; 16 controls: 16 red as planted, 0 missed; 0 limit problems ``` ## Run a case @@ -73,6 +128,13 @@ zig-out/bin/tf-moss-image MODEL_DIR/transformer case/case.safetensors case/ python3 tools/zig/moss_image_case.py decode MODEL_DIR case/latents.f32 case/zig.png python3 tools/zig/moss_image_case.py compare case/reference.png case/zig.png case/side.png python3 tools/zig/moss_image_case.py exact MODEL_DIR case/ # the same schedule with a float32 transformer +python3 tools/zig/moss_image_case.py exact MODEL_DIR case/ --fp32-encoder prompt.txt # and a float32 encoder + +# the prompt encoder alone, and a picture with it +python3 tools/zig/moss_image_case.py encode MODEL_DIR enc/ --prompt-file prompt.txt +zig-out/bin/tf-moss-image-encoder MODEL_DIR enc/encoder.safetensors prompt.txt +python3 tools/zig/moss_image_case.py start MODEL_DIR pic/ --no-encoder --seed 42 +zig-out/bin/tf-moss-image MODEL_DIR/transformer pic/case.safetensors pic/ --prompt prompt.txt ``` The case holds: @@ -85,7 +147,9 @@ The case holds: `tf-moss-image` first runs both branches of each step on the pipeline's own input and compares the velocities. It then runs the whole guided schedule from the starting latents, times it, and writes the final latents. A case made by -`start` holds only the prompts, the schedule and the noise; `tf-moss-image` then just makes the picture. +`start` holds only the prompts, the schedule and the noise; `tf-moss-image` then just makes the picture. With +`--prompt` (and `--negative`, else the empty negative prompt), it encodes both prompts natively in place of the case's +rows; `start --no-encoder` leaves them out of the case, and torch never loads the encoder. On MPS, which has no float64, the case tool computes the time sinusoid on the host in float64 and moves its float32 result to the GPU; every other step of the publisher's code runs unchanged. @@ -109,14 +173,14 @@ Results: 1.602e-4 for the prompt branch, as for torch bf16's own (1.602e-4); 1.735e-4 against torch's 1.715e-4 for the negative branch. Ours is the closer of the two at 25 and 27 of the 50 steps. - **Image:** our latents through the pipeline's decoder are 48.52 dB PSNR from the pipeline's own picture for the fox - and 28.93 dB for the sticky note. Fifty guided steps magnify bf16 rounding into different fine detail, so the - ruler is a float32 run of the whole schedule. Over 8 cases (both prompts, 4 seeds each), with the transformer in - float32 on torch's bf16 encoder rows, our picture is the closer to it in 6, at a mean of 32.3 dB PSNR against torch - bf16's 31.4 dB. With the encoder in float32 too, torch bf16 averages 32.5 dB and ours 31.2 dB: six cases are within - 0.5 dB or favor ours, and two sticky-note seeds favor torch's by 4.6 and 8.7 dB. That ruler came from a follow-up - tool, `exact --fp32-encoder`, which this tree's case tool does not have; it comes with the prompt encoder (#1459). - Those 8 ran on an earlier build of this family (0b69e127): this commit's latents are byte for byte its on the two - cases above. + and 28.93 dB for the sticky note. Fifty guided steps magnify bf16 rounding into different fine detail, so the ruler + is a float32 run of the whole schedule. There are two rulers, over the same 8 cases (both prompts, 4 seeds each): + - The transformer in float32 on torch's bf16 encoder rows: our picture is the closer in 6, at a mean of 32.3 dB PSNR + against torch bf16's 31.4 dB. Those 8 ran on an earlier build of this family (0b69e127); this commit's latents are + byte for byte its on the two cases above. + - Everything in float32, the encoder too (`exact --fp32-encoder`): torch bf16's picture averages 32.5 dB, ours 31.2 + dB. Six cases are within 0.5 dB of each other or favor ours (fox seed 1 by 2.9 dB). Two sticky-note seeds favor + torch's by 4.6 and 8.7 dB, where the schedule settled on different fine detail. - **Repeat:** the same step run twice, another between, gives identical bytes. - **Speed**, one session with no other GPU work: - torch bf16 on MPS: 0.744 s a forward for the fox, 0.760 s for the sticky note (the median of its 100). @@ -124,3 +188,17 @@ Results: 0.472 to 0.544 s over all six). 0.08 s for both prompts' refiner once. - **Memory:** 11.3 GB peak process footprint (10.6 GiB), with the transformer's bf16 weights (9.5 GB) in shared Metal buffers. +- **Prompt encoder**, both the prompt and the empty negative prompt (10, 34 and 63 tokens): + - Token ids: equal to the publisher's tokenizer's. + - Each layer's rows against torch bf16's, given the same token ids: cosine 0.99974 or more, the worst at the last + layers. + - The final rows against a float32 run of the encoder, ours beside torch bf16's own: + - cosine 0.999826 and worst row 0.0261 against 0.999806 and 0.0294 (empty); + - 0.999931 and 0.0261 against 0.999934 and 0.0259 (fox); + - 0.999926 and 0.0454 against 0.999938 and 0.0318 (sticky note). + - Repeat: identical bytes. + - Pictures with our encoder, against the all-float32 ruler over the same 8 cases: a mean of 31.13 dB. Our + transformer on torch's encoder rows gives 31.20 dB, and each is the closer in 4 of 8. The fox, against the + pipeline's own picture: 47.07 dB. + - Speed, one session with no other GPU work: 0.054 to 0.067 s a prompt against torch MPS's 0.158 to 0.165 s + (warm). Loading takes 0.6 s against 2.60 s. From 92c829887ed66c685ab800e36376a18bc1ddf44d Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Wed, 28 Oct 2026 00:16:30 -0500 Subject: [PATCH 3/3] MOSS-Image: tf-moss-image-encoder gates each step against float32 The check passed on repeatability alone. Now each layer and the final norm run alone (Encoder.steps) on the float32 run's input for that step, rounded to bf16, and each step's rows must be within twice torch bf16's distance from float32 on the same input, or plain bf16 rounding's when larger, in cosine error and in its worst row. moss_image_case.py encode keeps every layer's float32 rows and torch bf16's rows for each step alone. The whole prompt's rows are still compared with torch bf16's and float32 but not gated: in bf16 the layers magnify each implementation's rounding differently, so two correct runs part by more than any one step's error. The second encoding must still repeat the first byte for byte. reference.zig holds the comparison and the gate, and its tests plant a row 5% off that must fail with an identical repeat. --- docs/recipes/moss-image.md | 19 +++- tools/zig/moss_image_case.py | 43 ++++++-- zig/build/metal_targets.zig | 2 +- zig/src/families/moss_image/encoder.zig | 33 +++++- zig/src/families/moss_image/moss_image.zig | 2 + zig/src/families/moss_image/reference.zig | 118 +++++++++++++++++++++ zig/tests/moss_image_encoder.zig | 85 ++++++++------- 7 files changed, 249 insertions(+), 53 deletions(-) create mode 100644 zig/src/families/moss_image/reference.zig diff --git a/docs/recipes/moss-image.md b/docs/recipes/moss-image.md index eb0326dbb8..104638587f 100644 --- a/docs/recipes/moss-image.md +++ b/docs/recipes/moss-image.md @@ -76,8 +76,15 @@ The residual adds reuse `mi_gate_add`. `tf-moss-image-encoder` checks a case from `moss_image_case.py encode`: - the token ids; -- every layer's rows against the publisher's torch path in bf16; -- the final rows against bf16 and float32. +- each step alone, every layer and the final norm, on the float32 run's input for that step rounded to bf16: its rows + against float32 must be within twice torch bf16's distance on the same input (or plain bf16 rounding's, when + larger) in cosine error and in its worst row; +- the whole prompt, every layer's rows against torch bf16's and the final rows against float32, reported but not + gated: in bf16 the 32 layers magnify each implementation's rounding differently, so two correct runs part by more + than any one step's error; +- a second encoding must repeat the first byte for byte. + +It fails otherwise, and `reference.zig`'s tests check that a planted mismatch fails even when its repeat is identical. ## Kernel contract @@ -197,6 +204,14 @@ Results: - 0.999931 and 0.0261 against 0.999934 and 0.0259 (fox); - 0.999926 and 0.0454 against 0.999938 and 0.0318 (sticky note). - Repeat: identical bytes. + - `tf-moss-image-encoder`'s gates, all 33 steps on these three prompts and six more (31 to 75 tokens): each step's + distance from float32 is a mean of 0.50 of its limit, torch bf16's own, and at most 0.65 (the fox's layer 13, + worst row). A row planted 5% off at layer 20, in the float32 rows and torch's, fails there and at layer 21, its + repeat identical. + - The whole prompt in bf16 spreads by more than any step: on the 75-token prompt, torch's own runs (MPS or CPU, + chunks of 8 to 64, or the recurrent form) end 0.86 to 4.97 times its default's cosine error from float32, and + ours 1.82 times. Against float64, which float32 matches to four places, the sticky note's row 58 is 0.0153 to + 0.0320 across eight of those runs, and 0.0454 in ours. - Pictures with our encoder, against the all-float32 ruler over the same 8 cases: a mean of 31.13 dB. Our transformer on torch's encoder rows gives 31.20 dB, and each is the closer in 4 of 8. The fox, against the pipeline's own picture: 47.07 dB. diff --git a/tools/zig/moss_image_case.py b/tools/zig/moss_image_case.py index 1f3d0abfd6..ddba6b2931 100644 --- a/tools/zig/moss_image_case.py +++ b/tools/zig/moss_image_case.py @@ -20,8 +20,9 @@ with --fp32-encoder F, F's rows from a float32 encoder too (CASE_DIR/exact_fp32_encoder.png) decode MODEL_DIR LATENTS.f32 PNG raw float32 latents (rows x columns x 128) from tf-moss-image to a picture - encode MODEL_DIR OUT_DIR [--prompt-file F] the prompt encoder alone: token ids, embedding rows, every layer's - output and the final rows in bf16, and the final rows in float32 + encode MODEL_DIR OUT_DIR [--prompt-file F] the prompt encoder alone: token ids, embedding rows, and every layer's + output and the final rows in bf16 and in float32, and + each step alone in bf16 on the float32 run's input compare A.png B.png [SIDE.png] PSNR and mean absolute difference of two pictures, optionally side by side MODEL_DIR is the OpenMOSS-Team/MOSS-Image-Base folder. The case runs the card's defaults: 50 Euler steps, CFG 3, an empty negative prompt whose branch skips block 22, the bf16 transformer and text encoder, the float32 decoder. @@ -261,8 +262,19 @@ def encoder_rows(model_dir: str, prompt: str, dtype, device: str): return rows +def given(rows, device: str): + """A forward pre-hook that hands a layer `rows` in place of its input.""" + def hook(module, args, kwargs): + x = rows.to(device).unsqueeze(0) + if "hidden_states" in kwargs: + kwargs["hidden_states"] = x + return args, kwargs + return (x,) + tuple(args[1:]), kwargs + return hook + + def encode(args) -> int: - """The prompt encoder as the pipeline runs it, every layer's rows kept; then the same prompt in float32.""" + """The prompt encoder as the pipeline runs it, in float32 and bf16, every layer's rows kept; each step alone too.""" import torch from safetensors.torch import save_file from transformers import AutoModel, AutoTokenizer @@ -277,7 +289,7 @@ def encode(args) -> int: add_generation_prompt=True, enable_thinking=True) inputs = tokenizer(formatted, return_tensors="pt", truncation=True, max_length=4096) tensors = {"ids": inputs["input_ids"][0].to(torch.int32).contiguous()} - for dtype in (torch.bfloat16, torch.float32): + for dtype in (torch.float32, torch.bfloat16): encoder = AutoModel.from_pretrained(root / "text_encoder", dtype=dtype).eval().to(args.device) text = encoder.language_model kept = [] @@ -288,13 +300,30 @@ def encode(args) -> int: final = encoder(**inputs.to(args.device), use_cache=False, return_dict=True).last_hidden_state[0] for h in hooks: h.remove() - if dtype is torch.bfloat16: + if dtype is torch.float32: + for i, rows in enumerate(kept[1:]): + tensors[f"layer_fp32.{i}"] = rows[0].cpu().contiguous() + tensors["final_fp32"] = final.cpu().contiguous() + else: tensors["embed"] = kept[0][0].cpu().contiguous() for i, rows in enumerate(kept[1:]): tensors[f"layer.{i}"] = rows[0].cpu().contiguous() tensors["final"] = final.cpu().contiguous() - else: - tensors["final_fp32"] = final.cpu().contiguous() + # each step alone on the float32 run's input in bf16: torch bf16's own error at that step + steps = [tensors["embed"]] + [tensors[f"layer_fp32.{i}"].to(dtype) for i in range(len(text.layers))] + alone = [] + hooks = [layer.register_forward_pre_hook(given(steps[i], args.device), with_kwargs=True) + for i, layer in enumerate(text.layers)] + hooks += [layer.register_forward_hook(lambda mod, a, o: alone.append(o[0] if isinstance(o, tuple) else o)) + for layer in text.layers] + with torch.no_grad(): + encoder(**inputs.to(args.device), use_cache=False, return_dict=True) + alone.append(text.norm(steps[-1].to(args.device).unsqueeze(0))) + for h in hooks: + h.remove() + for i, rows in enumerate(alone[:-1]): + tensors[f"alone.{i}"] = rows[0].cpu().contiguous() + tensors["alone.final"] = alone[-1][0].cpu().contiguous() del encoder, text, kept if args.device == "mps": torch.mps.empty_cache() diff --git a/zig/build/metal_targets.zig b/zig/build/metal_targets.zig index 69684d3b8c..61105a3079 100644 --- a/zig/build/metal_targets.zig +++ b/zig/build/metal_targets.zig @@ -317,7 +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-moss-image-encoder", .path = "zig/tests/moss_image_encoder.zig", .about = "MOSS-Image's prompt encoder on a reference case: the ids checked, each layer's rows compared, then timed" }, + .{ .name = "tf-moss-image-encoder", .path = "zig/tests/moss_image_encoder.zig", .about = "MOSS-Image's prompt encoder on a reference case: the ids checked, each step gated against float32, 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/src/families/moss_image/encoder.zig b/zig/src/families/moss_image/encoder.zig index 94069976b3..f11b23093f 100644 --- a/zig/src/families/moss_image/encoder.zig +++ b/zig/src/families/moss_image/encoder.zig @@ -84,15 +84,13 @@ pub const Encoder = struct { if (T == 0 or T > cfg.max_rows) return error.PromptRows; if (out.len != T * H) return error.PromptRows; if (kept) |k| if (k.len != cfg.layers * T * H) return error.PromptRows; - const s = try Scratch.make(t.device, T); + const s = try scratch(t.device, T); defer s.buf.deinit(); const x = s.at(.x); for (tokens, 0..) |id, r| { if (id >= cfg.vocab) return error.TokenOutOfRange; @memcpy(halves(x.at(r * H * 2), H), t.w.embed[id * H ..][0..H]); } - rotary(floats(s.at(.cos), T * 32), floats(s.at(.sin), T * 32), T); - @memset(halves(s.at(.ones), H), 0x3f80); const pool = mtl.objc.Pool.push(); defer pool.pop(); var cb = t.queue.commandBuffer(); @@ -111,6 +109,27 @@ pub const Encoder = struct { @memcpy(out, halves(s.at(.out), T * H)); } + /// For checks, layer i on block i of `inputs` [layers + 1, rows, 4096], the final norm on the last. + pub fn steps(t: *Encoder, T: usize, inputs: []const u16, out: []u16) !void { + const n = (cfg.layers + 1) * T * H; + if (T == 0 or T > cfg.max_rows) return error.PromptRows; + if (inputs.len != n or out.len != n) return error.PromptRows; + const s = try scratch(t.device, T); + defer s.buf.deinit(); + const x = s.at(.x); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + for (0..cfg.layers + 1) |i| { + @memcpy(halves(x, T * H), inputs[i * T * H ..][0 .. T * H]); + const cb = t.queue.commandBuffer(); + const e = cb.compute(.serial); + const last = i == cfg.layers; + if (last) ks.encoderNorm(e, t.k.get("me_norm"), x, t.w.norm, s.at(.out), T, H, H, H, cfg.eps) else t.layer(e, s, t.w.layers[i], T); + try finish(cb, e); + @memcpy(out[i * T * H ..][0 .. T * H], halves(if (last) s.at(.out) else x, T * H)); + } + } + fn layer(t: *const Encoder, e: mtl.ComputeEncoder, s: Scratch, l: wts.Layer, T: usize) void { const x = s.at(.x); const xn = s.at(.xn); @@ -169,6 +188,14 @@ fn finish(cb: mtl.CommandBuffer, e: mtl.ComputeEncoder) !void { } } +/// A prompt of `T` rows' scratch with its rotary tables and unit gate filled. +fn scratch(device: mtl.Device, T: usize) !Scratch { + const s = try Scratch.make(device, T); + rotary(floats(s.at(.cos), T * 32), floats(s.at(.sin), T * 32), T); + @memset(halves(s.at(.ones), H), 0x3f80); + return s; +} + fn floats(r: Ref, n: usize) []f32 { return @as([*]f32, @ptrCast(@alignCast(r.addr())))[0..n]; } diff --git a/zig/src/families/moss_image/moss_image.zig b/zig/src/families/moss_image/moss_image.zig index d8fa5640c0..ba7c388e9f 100644 --- a/zig/src/families/moss_image/moss_image.zig +++ b/zig/src/families/moss_image/moss_image.zig @@ -8,6 +8,7 @@ pub const encoder_config = @import("encoder_config.zig"); pub const encoder_weights = @import("encoder_weights.zig"); pub const encoder = @import("encoder.zig"); pub const Encoder = encoder.Encoder; +pub const reference = @import("reference.zig"); pub const Model = model.Model; pub const Prompt = model.Prompt; pub const Timing = model.Timing; @@ -21,4 +22,5 @@ test { _ = encoder_config; _ = encoder_weights; _ = encoder; + _ = reference; } diff --git a/zig/src/families/moss_image/reference.zig b/zig/src/families/moss_image/reference.zig new file mode 100644 index 0000000000..c2e8e81ed8 --- /dev/null +++ b/zig/src/families/moss_image/reference.zig @@ -0,0 +1,118 @@ +//! A run's rows against float32 ones, gated by the bf16 reference's own distance from them. +const std = @import("std"); +const host = @import("host.zig"); + +pub const Match = struct { max_abs: f64, cosine: f64, worst_row: f64, worst_at: usize }; + +/// Largest absolute difference, cosine, and the largest relative L2 error of one row of `width` (and its row). +pub fn match(comptime T: type, ours: []const T, theirs: []const f32, width: usize) Match { + return measure(T, ours, theirs, width, false); +} + +/// `theirs` rounded to bf16 against itself: the closest a bf16 result can be. +pub fn rounding(theirs: []const f32, width: usize) Match { + return measure(f32, theirs, theirs, width, true); +} + +fn measure(comptime T: type, ours: []const T, theirs: []const f32, width: usize, comptime rounded: bool) Match { + var dot: f64 = 0; + var aa: f64 = 0; + var bb: f64 = 0; + var m: Match = .{ .max_abs = 0, .cosine = 0, .worst_row = 0, .worst_at = 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]) |h, b| { + const a: f64 = if (rounded) host.round(h) else if (T == u16) host.widen(h) else h; + const d = a - b; + dot += a * b; + aa += a * a; + bb += @as(f64, b) * b; + rd += d * d; + rb += @as(f64, b) * b; + m.max_abs = @max(m.max_abs, @abs(d)); + } + const rel = @sqrt(rd / @max(rb, 1e-30)); + if (rel > m.worst_row) { + m.worst_row = rel; + m.worst_at = row / width; + } + } + m.cosine = dot / @sqrt(aa * bb); + return m; +} + +/// How far a result may be from float32: this many times the bf16 reference's distance, or bf16 rounding's. +pub const factor = 2.0; + +pub const Verdict = struct { pass: bool, cosine_err: f64, cosine_limit: f64, row: f64, row_at: usize, row_limit: f64 }; + +/// Ours against float32: its cosine error and worst row each within `factor` of the bf16 reference's or rounding's. +pub fn gate(ours: Match, reference: Match, rounded: Match) Verdict { + const cosine_limit = factor * @max(1 - reference.cosine, 1 - rounded.cosine); + const row_limit = factor * @max(reference.worst_row, rounded.worst_row); + const cosine_err = 1 - ours.cosine; + return .{ + .pass = cosine_err <= cosine_limit and ours.worst_row <= row_limit, + .cosine_err = cosine_err, + .cosine_limit = cosine_limit, + .row = ours.worst_row, + .row_at = ours.worst_at, + .row_limit = row_limit, + }; +} + +/// A check passes only when every gate passes and the repeat gave the same bytes. +pub fn accept(verdicts: []const Verdict, repeat: bool) bool { + for (verdicts) |v| if (!v.pass) return false; + return repeat; +} + +const cols = 64; +const n = 8; + +fn synthetic(out: []f32, seed: u64) void { + var prng = std.Random.DefaultPrng.init(seed); + for (out) |*x| x.* = prng.random().floatNorm(f32) * 3; +} + +fn bf16Of(out: []u16, from: []const f32) void { + for (out, from) |*o, x| o.* = host.bf16(x); +} + +test "a result as close as the bf16 reference passes" { + var exact: [cols * n]f32 = undefined; + synthetic(&exact, 1); + var ref: [cols * n]f32 = undefined; + for (&ref, exact, 0..) |*r, x, i| r.* = host.round(x * (1 + 0.004 * @as(f32, if (i % 3 == 0) 1 else -1))); + var ours: [cols * n]u16 = undefined; + bf16Of(&ours, &exact); + const v = gate(match(u16, &ours, &exact, cols), match(f32, &ref, &exact, cols), rounding(&exact, cols)); + try std.testing.expect(v.pass); + try std.testing.expect(accept(&.{v}, true)); +} + +test "a planted mismatch fails even when its repeat is identical" { + var exact: [cols * n]f32 = undefined; + synthetic(&exact, 2); + var ref: [cols * n]f32 = undefined; + for (&ref, exact, 0..) |*r, x, i| r.* = host.round(x * (1 + 0.004 * @as(f32, if (i % 3 == 0) 1 else -1))); + var ours: [cols * n]u16 = undefined; + bf16Of(&ours, &exact); + // one row off by 5%, the rest exact: the cosine stays near 1, the worst row does not + for (ours[5 * cols ..][0..cols], exact[5 * cols ..][0..cols]) |*o, x| o.* = host.bf16(x * 1.05); + var again: [cols * n]u16 = ours; + const repeat = std.mem.eql(u16, &ours, &again); + const m = match(u16, &ours, &exact, cols); + const v = gate(m, match(f32, &ref, &exact, cols), rounding(&exact, cols)); + try std.testing.expect(repeat); + try std.testing.expect(!v.pass); + try std.testing.expectEqual(@as(usize, 5), m.worst_at); + try std.testing.expect(!accept(&.{v}, repeat)); +} + +test "a different repeat fails even when every gate passes" { + const v: Verdict = .{ .pass = true, .cosine_err = 0, .cosine_limit = 1, .row = 0, .row_at = 0, .row_limit = 1 }; + try std.testing.expect(!accept(&.{ v, v }, false)); +} diff --git a/zig/tests/moss_image_encoder.zig b/zig/tests/moss_image_encoder.zig index 42ba5d3692..4708bbdc64 100644 --- a/zig/tests/moss_image_encoder.zig +++ b/zig/tests/moss_image_encoder.zig @@ -8,37 +8,28 @@ const tokenizer = @import("tokenizer"); const usage = \\tf-moss-image-encoder MODEL_DIR CASE [PROMPT_FILE] \\MODEL_DIR is the OpenMOSS-Team/MOSS-Image-Base folder. CASE (moss_image_case.py encode) holds the prompt's token - \\ids, every layer's rows and the final rows from the publisher's torch path in bf16, and the final rows in float32. - \\The prompt (PROMPT_FILE, else the empty negative prompt) is tokenized and checked against the ids, then encoded: - \\each layer's rows and the final ones are compared (max |diff|, cosine, worst row) and the encoding is timed. + \\ids, every layer's rows and the final rows from the publisher's torch path in bf16 and in float32, and torch + \\bf16's rows for each step alone on the float32 run's input. The prompt (PROMPT_FILE, else the empty negative + \\prompt) is tokenized and checked against the ids. Each layer and the final norm, run alone on the same input, must + \\be as close to float32 as twice torch bf16's (or bf16 rounding's) cosine error and worst row there. The whole + \\prompt is then encoded twice, compared with both torch runs, and must repeat byte for byte; else it fails. ; -const Match = struct { max_abs: f64, cosine: f64, worst_row: f64 }; +const ref = mi.reference; +const H = mi.encoder_config.hidden; +const L = mi.encoder_config.layers; -/// Largest absolute difference, cosine, and the largest relative L2 error of one row of `width`. -fn match(ours: []const u16, theirs: []const f32, width: usize) Match { - 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]) |h, b| { - const a: f64 = @as(f32, @bitCast(@as(u32, h) << 16)); - const d = a - b; - dot += a * b; - aa += 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(aa * bb), .worst_row = worst }; +/// One comparison's line: against float32 within its gate, beside torch bf16's own distance. +fn report(label: []const u8, bf: ref.Match, v: ref.Verdict) void { + std.debug.print("{s}: float32 1-cosine {e:.2} of {e:.2}, worst row {d:.5} (row {d}) of {d:.5}, torch bf16 {d:.5} (row {d}): {s}\n", .{ label, v.cosine_err, v.cosine_limit, v.row, v.row_at, v.row_limit, bf.worst_row, bf.worst_at, if (v.pass) "pass" else "FAIL" }); +} + +/// Ours against float32 rows `fp32_name`, gated by torch bf16's rows `name` on the same input. +fn check(a: std.mem.Allocator, case: *const tf.safetensors.File, ours: []const u16, name: []const u8, fp32_name: []const u8) !struct { ref.Match, ref.Verdict } { + const bf16 = try floats(a, case, name); + const exact = try floats(a, case, fp32_name); + const theirs = ref.match(f32, bf16, exact, H); + return .{ theirs, ref.gate(ref.match(u16, ours, exact, H), theirs, ref.rounding(exact, H)) }; } /// The case tensor `name` as f32 (from bf16 or f32). @@ -86,24 +77,36 @@ pub fn main(init: std.process.Init) !u8 { std.debug.print("{s}: {d} tokens, the same ids as the tokenizer's: {}; loaded in {d:.1} s\n", .{ device.name(), ids.len, same_ids, load_s }); if (!same_ids) return 1; - const H = mi.encoder_config.hidden; - const L = mi.encoder_config.layers; const T = ids.len; + const inputs = try a.alloc(u16, (L + 1) * T * H); + for (0..L + 1) |i| { + var name: [32]u8 = undefined; + const rows = try floats(a, &case, if (i == 0) "embed" else try std.fmt.bufPrint(&name, "layer_fp32.{d}", .{i - 1})); + for (inputs[i * T * H ..][0 .. T * H], rows) |*o, x| o.* = mi.host.bf16(x); + } + const alone = try a.alloc(u16, (L + 1) * T * H); + try enc.steps(T, inputs, alone); + var verdicts: [L + 1]ref.Verdict = undefined; + for (0..L + 1) |i| { + var name: [32]u8 = undefined; + var fp32_name: [32]u8 = undefined; + var label: [32]u8 = undefined; + const last = i == L; + const bf, const v = try check(a, &case, alone[i * T * H ..][0 .. T * H], if (last) "alone.final" else try std.fmt.bufPrint(&name, "alone.{d}", .{i}), if (last) "final_fp32" else try std.fmt.bufPrint(&fp32_name, "layer_fp32.{d}", .{i})); + report(if (last) "final norm alone" else try std.fmt.bufPrint(&label, "layer {d:>2} alone ({s})", .{ i, if (mi.encoder_config.isFull(i)) "attention" else "delta " }), bf, v); + verdicts[i] = v; + } + const out = try a.alloc(u16, T * H); const kept = try a.alloc(u16, L * T * H); try enc.rows(ids, out, kept); - var worst: f64 = 1; + var lowest: f64 = 1; for (0..L) |i| { var name: [32]u8 = undefined; - const ref = try floats(a, &case, try std.fmt.bufPrint(&name, "layer.{d}", .{i})); - const m = match(kept[i * T * H ..][0 .. T * H], ref, H); - worst = @min(worst, m.cosine); - std.debug.print("layer {d:>2} ({s}): max |diff| {d:.4}, cosine {d:.6}, worst row {d:.5}\n", .{ i, if (mi.encoder_config.isFull(i)) "attention" else "delta ", m.max_abs, m.cosine, m.worst_row }); + lowest = @min(lowest, ref.match(u16, kept[i * T * H ..][0 .. T * H], try floats(a, &case, try std.fmt.bufPrint(&name, "layer.{d}", .{i})), H).cosine); } - const bf = match(out, try floats(a, &case, "final"), H); - const fp = match(out, try floats(a, &case, "final_fp32"), H); - std.debug.print("final rows against torch bf16: max |diff| {d:.4}, cosine {d:.6}, worst row {d:.5}\n", .{ bf.max_abs, bf.cosine, bf.worst_row }); - std.debug.print("final rows against float32: max |diff| {d:.4}, cosine {d:.6}, worst row {d:.5}\n", .{ fp.max_abs, fp.cosine, fp.worst_row }); + const bf, const v = try check(a, &case, out, "final", "final_fp32"); + std.debug.print("whole prompt, not gated: each layer against torch bf16's, cosine {d:.6} or more; the final rows against float32, 1-cosine {e:.2}, worst row {d:.5} (row {d}), torch bf16's {e:.2}, {d:.5} (row {d})\n", .{ lowest, v.cosine_err, v.row, v.row_at, 1 - bf.cosine, bf.worst_row, bf.worst_at }); const again = try a.alloc(u16, T * H); const timing = mtl.clock.seconds(); @@ -111,5 +114,7 @@ pub fn main(init: std.process.Init) !u8 { const encode_s = mtl.clock.seconds() - timing; const repeat = std.mem.eql(u16, out, again); std.debug.print("encoded again in {d:.3} s, the same bytes: {}\n", .{ encode_s, repeat }); - return if (repeat) 0 else 1; + const pass = ref.accept(&verdicts, repeat); + std.debug.print("{s}\n", .{if (pass) "PASS every step within its gate, repeated byte for byte" else "FAIL"}); + return if (pass) 0 else 1; }