From 2e66326a99bf0ae552698d95e6454512d9b3baa4 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Tue, 27 Oct 2026 01:38:57 -0500 Subject: [PATCH 1/6] MOSS-Image text to image: its transformer on the Metal runtime zig/src/families/moss_image runs MOSS-Image-Base's 5B joint transformer from the released bf16 diffusers checkpoint. A prompt's rows go through the input projection and the two refiner blocks once. At each step the image and prompt rows go through every joint block together, under one modulation from the timestep. The card's guidance runs two passes a step, the second with the empty negative prompt and block 22 skipped. Activations are bf16, rounded once per elementwise op as the publisher's torch code rounds them. Every sum is fp32, and there is no int8. Biased projections take fp32 sums from the prefill split-K GEMM (one part), then add the bias and round once, as a bf16 GEMM's bias epilogue does. The bf16 projections reuse the Qwen-VL tower's tfp::gemm_nax, and attention reuses the prefill tfp::attention_nax template in its full form. The time MLP and the modulation run on the host once a timestep, with the sinusoid in float64. tf-moss-image checks the transformer against a case from tools/zig/moss_image_case.py, which runs the publisher's own pipeline_moss_image.py (prompt encoder, transformer, scheduler, decoder) in torch. The case holds both branches' velocities at every step on the pipeline's own input; then the whole guided schedule runs and is timed. A case from `start` holds only the prompts, the schedule and the noise, and makes a picture. Gate 5 holds the family's 11 launches. --- tools/zig/moss_image_case.py | 264 ++++++++++++++++ zig/build/metal_targets.zig | 1 + zig/kernels/metal/moss_image.metal | 139 +++++++++ zig/kernels/metal/sources.zig | 2 + zig/src/core/contract.zig | 1 + zig/src/core/contract/moss_image_records.zig | 20 ++ zig/src/core/contract/registry.zig | 3 +- zig/src/families/moss_image/config.zig | 90 ++++++ zig/src/families/moss_image/host.zig | 162 ++++++++++ zig/src/families/moss_image/kernels.zig | 133 ++++++++ zig/src/families/moss_image/model.zig | 309 +++++++++++++++++++ zig/src/families/moss_image/moss_image.zig | 17 + zig/src/families/moss_image/weights.zig | 118 +++++++ zig/src/tensorfold.zig | 2 + zig/tests/contract/metal_limits.zig | 6 + zig/tests/moss_image.zig | 178 +++++++++++ 16 files changed, 1444 insertions(+), 1 deletion(-) create mode 100644 tools/zig/moss_image_case.py create mode 100644 zig/kernels/metal/moss_image.metal create mode 100644 zig/src/core/contract/moss_image_records.zig create mode 100644 zig/src/families/moss_image/config.zig create mode 100644 zig/src/families/moss_image/host.zig create mode 100644 zig/src/families/moss_image/kernels.zig create mode 100644 zig/src/families/moss_image/model.zig create mode 100644 zig/src/families/moss_image/moss_image.zig create mode 100644 zig/src/families/moss_image/weights.zig create mode 100644 zig/tests/moss_image.zig diff --git a/tools/zig/moss_image_case.py b/tools/zig/moss_image_case.py new file mode 100644 index 0000000000..c3b9eb6dc7 --- /dev/null +++ b/tools/zig/moss_image_case.py @@ -0,0 +1,264 @@ +#!/usr/bin/env python3 +"""MOSS-Image reference cases for tf-moss-image from the publisher's own pipeline (pipeline_moss_image.py, torch).""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import math +import sys +import time +from pathlib import Path + +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 + decode MODEL_DIR LATENTS.f32 PNG raw float32 latents (rows x columns x 128) from tf-moss-image to a picture + 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. +""" + + +def module(model_dir: str): + """The checkpoint's pipeline_moss_image.py, imported from its folder.""" + spec = importlib.util.spec_from_file_location("pipeline_moss_image", Path(model_dir) / "pipeline_moss_image.py") + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +def pipeline(model_dir: str, device: str, text: bool = True, dit: bool = True): + """The publisher's MOSSImagePipeline from the folder's components (the encoder and transformer optional).""" + import torch + from diffusers import AutoencoderKLFlux2 + from transformers import AutoModel, AutoTokenizer + + m = module(model_dir) + root = Path(model_dir) + transformer = m.MOSSImageTransformer.from_pretrained(root / "transformer", torch_dtype=torch.bfloat16).eval() if dit else None + if dit and device == "mps": + # MPS has no float64: the publisher's float64 time sinusoid runs on the host; all else runs as published. + class HostTimes: + def __getattr__(self, name): + return getattr(torch, name) + + @staticmethod + def as_tensor(data, device=None, dtype=None): + return torch.as_tensor(data, dtype=dtype).cpu() + + m.torch = HostTimes() + transformer.time_embedding.register_forward_pre_hook(lambda mod, a: tuple(x.to(mod[0].weight.device) for x in a)) + scheduler = m.MOSSImageScheduler.from_pretrained(root / "scheduler") + tokenizer = AutoTokenizer.from_pretrained(root / "tokenizer") if text else None + encoder = AutoModel.from_pretrained(root / "text_encoder", dtype=torch.bfloat16).eval() if text else None + vae = AutoencoderKLFlux2.from_pretrained(root / "vae", torch_dtype=torch.float32).eval() + return m.MOSSImagePipeline(transformer, scheduler, encoder, tokenizer, vae).to(device), m + + +def rows(x): + """[1, C, H, W] to (H, W, C) float32 on the host: tf-moss-image's row order.""" + return x[0].float().permute(1, 2, 0).contiguous().cpu() + + +def to_png(pipe, latents, path: str) -> None: + """Latents [1, 128, H, W] through the pipeline's own decode path (its batch-norm undo, unpack, VAE) to a PNG.""" + import torch + + with torch.no_grad(): + z = latents.to(device=pipe.vae.device, dtype=pipe.vae.dtype) + bn = pipe.vae.bn + eps = pipe.vae.config.batch_norm_eps + z = z * (bn.running_var + eps).sqrt().to(z).view(1, -1, 1, 1) + bn.running_mean.to(z).view(1, -1, 1, 1) + b, c, h, w = z.shape + z = z.reshape(b, c // 4, 2, 2, h, w).permute(0, 1, 4, 2, 5, 3).reshape(b, c // 4, 2 * h, 2 * w) + image = pipe.vae.decode(z, return_dict=False)[0].clamp(-1, 1) + pipe.image_processor.postprocess(image, output_type="pil")[0].save(path) + + +def case(args) -> int: + 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) + mark = time.perf_counter() + positive, _ = pipe.encode_prompt(Path(args.prompt_file).read_text().strip()) + negative, _ = pipe.encode_prompt("") + encode_s = time.perf_counter() - mark + # the encoder is done; the transformer and decoder stay + pipe.text_encoder = None + if args.device == "mps": + torch.mps.empty_cache() + 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) + calls = [] + forward = pipe.transformer.forward + + def recorded(hidden_states, timestep, encoder_hidden_states, encoder_attention_mask=None, return_dict=True, + skip_layers=None): + mark = time.perf_counter() + result = forward(hidden_states, timestep, encoder_hidden_states, encoder_attention_mask, return_dict, + skip_layers) + sample = result.sample if return_dict else result[0] + sample.float().sum().item() + calls.append({"input": rows(hidden_states), "t": float(timestep), "skip": sorted(skip_layers or ()), + "velocity": rows(sample), "s": time.perf_counter() - mark}) + return result + + pipe.transformer.forward = recorded + image = pipe(prompt_embeds=positive, negative_prompt_embeds=negative, height=args.height, width=args.width, + num_inference_steps=args.steps, latents=noise.to(args.device), output_type="latent").images + pipe.transformer.forward = forward + cond, uncond = calls[0::2], calls[1::2] + if not all(c["skip"] == [] for c in cond) or not all(u["skip"] == [22] for u in uncond): + sys.exit("the pipeline's calls are not (prompt, negative skipping 22) pairs") + sched = pipe.scheduler + tensors = {"prompt": positive[0].cpu().contiguous(), "negative": negative[0].cpu().contiguous(), + "timesteps": sched.timesteps.float().cpu().contiguous(), "sigmas": sched.sigmas.float().cpu().contiguous(), + "latents": noise[0].permute(1, 2, 0).contiguous(), "inputs": torch.stack([c["input"] for c in cond]), + "velocity": torch.stack([c["velocity"] for c in cond]), + "velocity_negative": torch.stack([u["velocity"] for u in uncond])} + if args.fp32_steps: + # the same transformer and inputs in float32 throughout: the bf16 reference's own distance from exact arithmetic + pipe.transformer.float() + exact, exact_negative = [], [] + with torch.no_grad(): + for s in range(min(args.fp32_steps, len(cond))): + x = cond[s]["input"].permute(2, 0, 1)[None].to(args.device) + t = sched.timesteps[s] + exact.append(rows(pipe.transformer(x, t, positive.float()).sample)) + exact_negative.append(rows(pipe.transformer(x, t, negative.float(), skip_layers=[22]).sample)) + tensors["velocity_f32"] = torch.stack(exact) + tensors["velocity_negative_f32"] = torch.stack(exact_negative) + pipe.transformer.to(torch.bfloat16) + save_file(tensors, str(out / "case.safetensors")) + image.float().cpu().permute(0, 2, 3, 1).contiguous().numpy().tofile(out / "reference.latents.f32") + to_png(pipe, image, str(out / "reference.png")) + meta = {"width": args.width, "height": args.height, "seed": args.seed, "steps": len(cond), + "prompt_rows": int(positive.shape[1]), "negative_rows": int(negative.shape[1]), "device": args.device, + "encode_s": round(encode_s, 2), "torch_bf16_forward_s": [round(c["s"], 3) for c in calls], + "torch": torch.__version__} + (out / "reference.json").write_text(json.dumps(meta, indent=1) + "\n") + print(json.dumps({k: v for k, v in meta.items() if k != "torch_bf16_forward_s"})) + return 0 + + +def start(args) -> int: + """What tf-moss-image needs to make a picture: both prompts' encoder rows, the schedule and the seeded 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) + 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(), + "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])})) + return 0 + + +def exact(args) -> int: + """The case's whole guided schedule with the transformer in float32: where exact arithmetic lands.""" + import torch + from safetensors.torch import load_file + + case_dir = Path(args.case_dir) + case = load_file(str(case_dir / "case.safetensors")) + 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, + 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")) + return 0 + + +def decode(args) -> int: + import numpy as np + import torch + + raw = np.fromfile(args.latents, dtype=np.float32) + h, w = args.height // 16, args.width // 16 + if raw.size != h * w * 128: + sys.exit(f"{args.latents} holds {raw.size} values; {args.width}x{args.height} needs {h * w * 128}") + pipe, _ = pipeline(args.model_dir, args.device, text=False, dit=False) + to_png(pipe, torch.from_numpy(raw).reshape(1, h, w, 128).permute(0, 3, 1, 2), args.png) + return 0 + + +def compare(args) -> int: + import numpy as np + from PIL import Image + + a = np.asarray(Image.open(args.a).convert("RGB"), dtype=np.float64) + b = np.asarray(Image.open(args.b).convert("RGB"), dtype=np.float64) + mse = float(np.mean((a - b) ** 2)) + psnr = float("inf") if mse == 0 else 10 * math.log10(255.0 ** 2 / mse) + print(json.dumps({"psnr_db": round(psnr, 2), "mean_abs": round(float(np.mean(np.abs(a - b))), 3), + "max_abs": int(np.max(np.abs(a - b)))})) + if args.side: + side = np.concatenate([a, b], axis=1).astype(np.uint8) + Image.fromarray(side).save(args.side) + return 0 + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__, epilog=USAGE, formatter_class=argparse.RawDescriptionHelpFormatter) + sub = parser.add_subparsers(dest="mode", required=True) + c = sub.add_parser("case") + c.add_argument("model_dir") + c.add_argument("out_dir") + c.add_argument("--prompt-file", required=True) + c.add_argument("--width", type=int, default=1024) + c.add_argument("--height", type=int, default=1024) + c.add_argument("--seed", type=int, default=42) + c.add_argument("--steps", type=int, default=50) + c.add_argument("--fp32-steps", type=int, default=1, help="steps also run in float32 (0: none)") + c.add_argument("--device", default="mps") + 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("--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) + s.add_argument("--seed", type=int, default=42) + s.add_argument("--steps", type=int, default=50) + s.add_argument("--device", default="mps") + x = sub.add_parser("exact") + x.add_argument("model_dir") + x.add_argument("case_dir") + x.add_argument("--steps", type=int, default=50) + x.add_argument("--device", default="mps") + d = sub.add_parser("decode") + d.add_argument("model_dir") + d.add_argument("latents") + d.add_argument("png") + d.add_argument("--width", type=int, default=1024) + d.add_argument("--height", type=int, default=1024) + d.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) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/zig/build/metal_targets.zig b/zig/build/metal_targets.zig index d5ab369dd7..2eb6464016 100644 --- a/zig/build/metal_targets.zig +++ b/zig/build/metal_targets.zig @@ -314,6 +314,7 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s .{ .name = "tf-glm-video", .path = "zig/tests/glm_video.zig", .about = "GLM-5.3-Flash's video preparation on a clip (libtfvideo), written for tools/zig/glm_video_fixtures.py" }, .{ .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-glm-attn-probe", .path = "zig/tests/glm_attn_probe.zig", .about = "GLM-5.3-Flash's dense latent attention kernels on given inputs (scores, probabilities, outputs)" }, .{ .name = "tf-moe-bench", .path = "zig/tests/moe_bench.zig", .about = "Routed-expert prompt kernels from .metal files on a skewed routing: chained speed, bits against the first" }, .{ .name = "tf-mm-bench", .path = "zig/tests/mm_bench.zig", .about = "Prompt matmul kernels from .metal files: chained speed at one shape, bits against the first, error against fp64" }, diff --git a/zig/kernels/metal/moss_image.metal b/zig/kernels/metal/moss_image.metal new file mode 100644 index 0000000000..75ebb8c0ba --- /dev/null +++ b/zig/kernels/metal/moss_image.metal @@ -0,0 +1,139 @@ +// MOSS-Image's steps after attention_nax.metal, each rounded to bf16 where the publisher's torch code rounds. + +inline float mi_bf(float x) { return float(bfloat(x)); } + +// One threadgroup's sum of every thread's `v` (at most 1,024 threads), returned to all of them. +inline float mi_sum(float v, threadgroup float* part, uint n, uint lane, uint sg) { + v = simd_sum(v); + if (lane == 0) part[sg] = v; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + float s = lane < (n + 31) / 32 ? part[lane] : 0.0f; + s = simd_sum(s); + if (lane == 0) part[32] = s; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const float out = part[32]; + threadgroup_barrier(mem_flags::mem_threadgroup); + return out; +} + +struct MiRows { + uint width; + float eps; + uint modulated; +}; + +// y = bf16(bf16(x * rsqrt(mean(x^2) + eps)) * w), then if modulated bf16(bf16(y * s) + b); a threadgroup a row. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void mi_norm( + const device bfloat* x [[buffer(0)]], const device bfloat* w [[buffer(1)]], const device bfloat* s [[buffer(2)]], + const device bfloat* b [[buffer(3)]], device bfloat* y [[buffer(4)]], constant MiRows& p [[buffer(5)]], + 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.width; + device bfloat* yr = y + ulong(row) * p.width; + 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) { + float v = mi_bf(mi_bf(float(xr[i]) * inv) * float(w[i])); + if (p.modulated) v = mi_bf(v * float(s[i])) + float(b[i]); + yr[i] = bfloat(v); + } +} + +struct MiHeads { + uint rows; + uint heads; + uint stride; + float eps; +}; + +// q, k: head RMSNorm, then halves (c, c + 64) rotated by cs, sn; q, k, v to [heads, rows, 128]. A simdgroup a head. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void mi_qk( + const device bfloat* qr [[buffer(0)]], const device bfloat* kr [[buffer(1)]], const device bfloat* vr [[buffer(2)]], + const device bfloat* nq [[buffer(3)]], const device bfloat* nk [[buffer(4)]], const device float* cs [[buffer(5)]], + const device float* sn [[buffer(6)]], device bfloat* Q [[buffer(7)]], device bfloat* K [[buffer(8)]], + device bfloat* V [[buffer(9)]], constant MiHeads& p [[buffer(10)]], 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.stride + h * 128 + lane * 4; + float q[4], k[4]; + float qq = 0.0f, kk = 0.0f; + for (uint i = 0; i < 4; i++) { + q[i] = float(qr[from + i]); + k[i] = float(kr[from + i]); + qq += q[i] * q[i]; + kk += k[i] * k[i]; + } + const float qn = metal::precise::rsqrt(simd_sum(qq) / 128.0f + p.eps); + const float kn = metal::precise::rsqrt(simd_sum(kk) / 128.0f + p.eps); + for (uint i = 0; i < 4; i++) { + q[i] = mi_bf(mi_bf(q[i] * qn) * float(nq[lane * 4 + i])); + k[i] = mi_bf(mi_bf(k[i] * kn) * float(nk[lane * 4 + i])); + } + // lane l holds channels 4l..4l+3; its partner half sits 16 lanes away. The first half turns by -partner. + const float sign = lane < 16 ? -1.0f : 1.0f; + const ulong to = (ulong(h) * p.rows + r) * 128 + lane * 4; + for (uint i = 0; i < 4; i++) { + const float qp = sign * simd_shuffle_xor(q[i], ushort(16)); + const float kp = sign * simd_shuffle_xor(k[i], ushort(16)); + const uint j = (lane % 16) * 4 + i; + const float c = cs[ulong(r) * 64 + j], s = sn[ulong(r) * 64 + j]; + Q[to + i] = bfloat(mi_bf(q[i] * c) + mi_bf(qp * s)); + K[to + i] = bfloat(mi_bf(k[i] * c) + mi_bf(kp * s)); + V[to + i] = vr[from + i]; + } +} + +struct MiSitu { + uint width; + uint stride; +}; + +// o = situ(h, g) = (4 tanh(g / 4) sigmoid(g)) (25 tanh(h / 25)), each op rounded to bf16; h and g rows of `stride`. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void mi_situ( + const device bfloat* h [[buffer(0)]], const device bfloat* g [[buffer(1)]], device bfloat* o [[buffer(2)]], + constant MiSitu& p [[buffer(3)]], uint2 at [[thread_position_in_grid]]) { + const ulong i = ulong(at.y) * p.stride + at.x; + const float gv = float(g[i]), hv = float(h[i]); + const float sig = mi_bf(1.0f / (1.0f + metal::precise::exp(-gv))); + const float left = mi_bf(mi_bf(4.0f * mi_bf(metal::precise::tanh(mi_bf(gv * 0.25f)))) * sig); + // torch divides by a host scalar as a product with its fp32 reciprocal + const float right = mi_bf(25.0f * mi_bf(metal::precise::tanh(mi_bf(hv * (1.0f / 25.0f))))); + o[ulong(at.y) * p.width + at.x] = bfloat(left * right); +} + +// x = bf16(x + bf16(g * y)) by column; g is a modulation gate, or ones for an ungated residual. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void mi_gate_add( + device bfloat* x [[buffer(0)]], const device bfloat* y [[buffer(1)]], const device bfloat* g [[buffer(2)]], + constant uint& width [[buffer(3)]], uint2 at [[thread_position_in_grid]]) { + const ulong i = ulong(at.y) * width + at.x; + x[i] = bfloat(float(x[i]) + mi_bf(float(g[at.x]) * float(y[i]))); +} + +// y = bf16(c + b) by column: a projection's fp32 sums and its bias, rounded once (the bias epilogue of a bf16 GEMM). +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void mi_bias( + const device float* c [[buffer(0)]], const device bfloat* b [[buffer(1)]], device bfloat* y [[buffer(2)]], + constant uint& width [[buffer(3)]], uint2 at [[thread_position_in_grid]]) { + const ulong i = ulong(at.y) * width + at.x; + y[i] = bfloat(c[i] + float(b[at.x])); +} + +// tfp::attention_nax (MLX's steel attention on the tensor units), full: every row sees every row. +#define MI_ATTENTION(NAME, AQ, AK) \ + [[max_total_threads_per_threadgroup(128)]] [[kernel]] void NAME( \ + const device bfloat16_t* Q [[buffer(0)]], const device bfloat16_t* K [[buffer(1)]], \ + const device bfloat16_t* V [[buffer(2)]], const device int32_t* P [[buffer(3)]], \ + const device float* F [[buffer(4)]], device bfloat16_t* O [[buffer(5)]], \ + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], \ + uint3 tg [[threadgroup_position_in_grid]]) { \ + tfp::attention_nax(Q, K, V, O, P, F, tg, sg, lane); \ + } +MI_ATTENTION(mi_attention_00, false, false) +MI_ATTENTION(mi_attention_01, false, true) +MI_ATTENTION(mi_attention_10, true, false) +MI_ATTENTION(mi_attention_11, true, true) diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig index 51b14b5a40..b8b0f37435 100644 --- a/zig/kernels/metal/sources.zig +++ b/zig/kernels/metal/sources.zig @@ -127,6 +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"); /// 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.zig b/zig/src/core/contract.zig index b38e46eccd..09e8d931a3 100644 --- a/zig/src/core/contract.zig +++ b/zig/src/core/contract.zig @@ -19,6 +19,7 @@ test { _ = @import("contract/vit_records.zig"); _ = @import("contract/simd_pairs_records.zig"); _ = @import("contract/glm_rows_records.zig"); + _ = @import("contract/moss_image_records.zig"); _ = spec; _ = gates; _ = rows; diff --git a/zig/src/core/contract/moss_image_records.zig b/zig/src/core/contract/moss_image_records.zig new file mode 100644 index 0000000000..32b7ec030a --- /dev/null +++ b/zig/src/core/contract/moss_image_records.zig @@ -0,0 +1,20 @@ +//! MOSS-Image's transformer launches: its own kernels, and the prefill split-K GEMM it reads fp32 sums from. +const std = @import("std"); +const M = @import("limits.zig").MetalLaunch; +pub const records = [_]M{ + .{ .name = "moss image gemm custom_kernel_tf_gemm_splitk_nax_bf16_n_t_n_t_64_64_256_2_2_bfloat16_t_bfloat16_t_int32_t_float", .threads = 128, .tg_memory = 0, .reserved = 128 }, + .{ .name = "moss image gemm custom_kernel_tf_gemm_splitk_nax_bf16_n_t_t_t_64_64_256_2_2_bfloat16_t_bfloat16_t_int32_t_float", .threads = 128, .tg_memory = 0, .reserved = 128 }, + .{ .name = "moss image mi_norm", .threads = 256, .tg_memory = 144, .reserved = 256 }, + .{ .name = "moss image mi_qk", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "moss image mi_situ", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "moss image mi_gate_add", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "moss image mi_bias", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .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_10", .threads = 128, .tg_memory = 0, .reserved = 128 }, + .{ .name = "moss image mi_attention_11", .threads = 128, .tg_memory = 0, .reserved = 128 }, +}; + +test "the MOSS-Image launches fit every Apple family" { + try std.testing.expectEqual(@as(usize, 0), try @import("registry.zig").checkMetal(std.testing.allocator, &records, true)); +} diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 8d5e1ee270..211e6c9ff5 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -570,7 +570,8 @@ test { _ = paging; } /// The Metal collector keeps measured slots and the restored vision declarations. -pub const all_metal = metal ++ @import("delta_slots_records.zig").records ++ @import("flashnext_pending.zig").records ++ @import("simd_pairs_records.zig").records ++ @import("glm_rows_records.zig").records ++ @import("weight_only_records.zig").metal ++ @import("hc_activation_records.zig").metal ++ @import("qwen4_records.zig").metal; +pub const all_metal = metal ++ @import("delta_slots_records.zig").records ++ @import("flashnext_pending.zig").records ++ @import("simd_pairs_records.zig").records ++ @import("glm_rows_records.zig").records ++ @import("weight_only_records.zig").metal ++ @import("hc_activation_records.zig").metal ++ @import("qwen4_records.zig").metal ++ + @import("moss_image_records.zig").records; const learning = @import("registry_learning.zig"); pub const gradient_pending = learning.gradient_pending; pub const qwen_learning_pending = learning.qwen_learning_pending; diff --git a/zig/src/families/moss_image/config.zig b/zig/src/families/moss_image/config.zig new file mode 100644 index 0000000000..9153293f18 --- /dev/null +++ b/zig/src/families/moss_image/config.zig @@ -0,0 +1,90 @@ +//! MOSS-Image's transformer geometry, and the check that a checkpoint's transformer/config.json is that geometry. +const std = @import("std"); + +pub const hidden = 3072; +pub const heads = 24; +pub const head_dim = hidden / heads; +pub const ffn = 8192; +pub const layers = 42; +/// The prompt refiner's blocks ahead of the joint ones. +pub const text_layers = 2; +pub const context = 4096; +pub const time_dim = 256; +pub const channels = 128; +pub const eps: f32 = 1e-6; +/// Rotary: four axes (unused, frame, height, width) of 16 frequencies each, the rotation of halves. +pub const bands = head_dim / 8; +pub const theta: f64 = 10000; +/// The pipeline's longest prompt (max_sequence_length). +pub const max_text = 4096; +/// One latent row a 16 x 16 pixel patch: the image decoder's 8x reduction and the 2 x 2 pack. +pub const pixels_per_row = 16; + +/// The latent grid of a width x height image: rows (height / 16) by columns (width / 16). +pub fn grid(width: usize, height: usize) !struct { rows: usize, columns: usize } { + if (width == 0 or height == 0 or width % pixels_per_row != 0 or height % pixels_per_row != 0) return error.BadImageSize; + const rows = height / pixels_per_row; + const columns = width / pixels_per_row; + const count = std.math.mul(usize, rows, columns) catch return error.ImageTooLarge; + // Metal row offsets use uint elements in the widest buffer, the MLP's (rows, 2 * ffn), prompt rows included. + if (rows > std.math.maxInt(u32) or columns > std.math.maxInt(u32) or count > std.math.maxInt(u32) / (2 * ffn) - max_text) + return error.ImageTooLarge; + return .{ .rows = rows, .columns = columns }; +} + +/// Accepts the transformer config this family runs (MOSSImageTransformer at its released size), else an error. +pub fn check(gpa: std.mem.Allocator, text: []const u8) !void { + const doc = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer doc.deinit(); + if (doc.value != .object) return error.BadConfig; + const o = doc.value.object; + const name = o.get("_class_name") orelse return error.BadConfig; + if (name != .string or !std.mem.eql(u8, name.string, "MOSSImageTransformer")) return error.NotMossImage; + const want = .{ .{ "dim", hidden }, .{ "ffn_dim", ffn }, .{ "num_heads", heads }, .{ "num_layers", layers }, .{ "text_dim", context }, .{ "freq_dim", time_dim }, .{ "latent_channels", channels } }; + 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: {s} is not {d}; this family runs the released geometry only", .{ kv[0], kv[1] }); + return error.UnsupportedGeometry; + } + } + const base = o.get("rope_base") orelse return error.BadConfig; + const value: f64 = switch (base) { + .float => |v| v, + .integer => |v| @floatFromInt(v), + else => return error.BadConfig, + }; + if (value != theta) return error.UnsupportedGeometry; +} + +const sample = + \\{"_class_name": "MOSSImageTransformer", "_diffusers_version": "0.38.0", "dim": 3072, "ffn_dim": 8192, + \\ "freq_dim": 256, "latent_channels": 128, "num_heads": 24, "num_layers": 42, "rope_base": 10000.0, + \\ "text_dim": 4096} +; + +test "the released transformer config passes; another depth, rotary base or class does not" { + const a = std.testing.allocator; + try check(a, sample); + const deeper = try std.mem.replaceOwned(u8, a, sample, "\"num_layers\": 42", "\"num_layers\": 40"); + defer a.free(deeper); + try std.testing.expectError(error.UnsupportedGeometry, check(a, deeper)); + const rope = try std.mem.replaceOwned(u8, a, sample, "10000.0", "1000000.0"); + defer a.free(rope); + try std.testing.expectError(error.UnsupportedGeometry, check(a, rope)); + const missing = try std.mem.replaceOwned(u8, a, sample, "\"rope_base\": 10000.0,", ""); + defer a.free(missing); + try std.testing.expectError(error.BadConfig, check(a, missing)); + const other = try std.mem.replaceOwned(u8, a, sample, "MOSSImageTransformer", "MOSSVideoTransformer"); + defer a.free(other); + try std.testing.expectError(error.NotMossImage, check(a, other)); + try std.testing.expectEqual(@as(usize, 64), (try grid(1024, 1024)).rows); + try std.testing.expectEqual(@as(usize, 84), (try grid(1344, 768)).columns); + try std.testing.expectError(error.BadImageSize, grid(1000, 768)); +} + +test "aligned huge image dimensions are refused without overflowing their latent grid" { + const huge = std.math.maxInt(usize) - (pixels_per_row - 1); + try std.testing.expectError(error.ImageTooLarge, grid(huge, huge)); + try std.testing.expectError(error.ImageTooLarge, grid(huge, pixels_per_row)); +} diff --git a/zig/src/families/moss_image/host.zig b/zig/src/families/moss_image/host.zig new file mode 100644 index 0000000000..b442df7759 --- /dev/null +++ b/zig/src/families/moss_image/host.zig @@ -0,0 +1,162 @@ +//! The host's part of a MOSS-Image pass: the rotary tables, and the time embedding's modulation all blocks share. +const std = @import("std"); +const cfg = @import("config.zig"); + +/// fp32 to bf16, round to nearest even (torch's cast); NaN stays NaN. +pub fn bf16(x: f32) u16 { + const b: u32 = @bitCast(x); + if (b & 0x7fffffff > 0x7f800000) return @truncate((b >> 16) | 0x40); + return @truncate((b + 0x7fff + ((b >> 16) & 1)) >> 16); +} + +pub fn widen(h: u16) f32 { + return @bitCast(@as(u32, h) << 16); +} + +pub fn round(x: f32) f32 { + return widen(bf16(x)); +} + +/// Rotary cos, sin (image + text rows, 64) in bf16: image rows at frame `text`, centred; then prompt rows by frame. +pub fn rotary(cos: []f32, sin: []f32, bands: *const [cfg.bands]f32, text: usize, rows: usize, columns: usize) void { + const half = cfg.head_dim / 2; + const image = rows * columns; + std.debug.assert(cos.len == (image + text) * half and sin.len == cos.len); + const top = @as(f32, @floatFromInt(rows - 1)) / 2; + const left = @as(f32, @floatFromInt(columns - 1)) / 2; + for (0..image + text) |r| { + const at: [4]f32 = if (r < image) + .{ 0, @floatFromInt(text), @as(f32, @floatFromInt(r / columns)) - top, @as(f32, @floatFromInt(r % columns)) - left } + else + .{ 0, @floatFromInt(r - image), 0, 0 }; + for (0..half) |j| { + const angle = at[j / cfg.bands] * bands[j % cfg.bands]; + cos[r * half + j] = round(@cos(angle)); + sin[r * half + j] = round(@sin(angle)); + } + } +} + +/// y[o] = bf16(sum over k of x[k] w[o, k] + b[o]), w stored (outputs, inputs): fp32 sums, one rounding after the bias. +pub fn project(w: []const u16, b: []const u16, x: []const f32, y: []f32) void { + std.debug.assert(w.len == x.len * y.len and b.len == y.len); + const V = @Vector(8, f32); + for (y, 0..) |*o, row| { + const ws = w[row * x.len ..][0..x.len]; + var acc: V = @splat(0); + var k: usize = 0; + while (k + 8 <= x.len) : (k += 8) { + const bits: @Vector(8, u32) = @as(@Vector(8, u16), ws[k..][0..8].*); + const wv: V = @bitCast(bits << @splat(16)); + acc += wv * @as(V, x[k..][0..8].*); + } + var sum = @reduce(.Add, acc); + while (k < x.len) : (k += 1) sum += widen(ws[k]) * x[k]; + o.* = round(sum + widen(b[row])); + } +} + +/// The bf16 projections the modulation reads, each stored (outputs, inputs) with its bias. +pub const Weights = struct { + time_1: []const u16, + time_1_b: []const u16, + time_2: []const u16, + time_2_b: []const u16, + blocks: []const u16, + blocks_b: []const u16, + head: []const u16, + head_b: []const u16, +}; + +/// The shared vectors a pass at one timestep gives the kernels, each `hidden` bf16 values. +pub const Vector = enum(usize) { scale_a, shift_a, gate_a, scale_f, shift_f, gate_f, scale_out, shift_out, ones }; +pub const vectors = @typeInfo(Vector).@"enum".field_names.len; + +/// The blocks' and head's modulation at `timestep` (sigma x 1,000); scales as bf16 `1 + scale`, as torch makes them. +pub fn modulation(out: *[vectors][cfg.hidden]u16, gpa: std.mem.Allocator, w: Weights, timestep: f32) !void { + const h = cfg.hidden; + const emb = try gpa.alloc(f32, cfg.time_dim); + defer gpa.free(emb); + const mid = try gpa.alloc(f32, h); + defer gpa.free(mid); + const time = try gpa.alloc(f32, h); + defer gpa.free(time); + const table = try gpa.alloc(f32, 6 * h); + defer gpa.free(table); + const head = try gpa.alloc(f32, 2 * h); + defer gpa.free(head); + const half = cfg.time_dim / 2; + for (0..half) |i| { + const freq = std.math.pow(f64, 10000, -@as(f64, @floatFromInt(i)) / half); + const phase = @as(f64, timestep) * freq; + emb[i] = round(@floatCast(@cos(phase))); + emb[half + i] = round(@floatCast(@sin(phase))); + } + project(w.time_1, w.time_1_b, emb, mid); + for (mid) |*x| x.* = round(x.* / (1.0 + @exp(-x.*))); + project(w.time_2, w.time_2_b, mid, time); + project(w.blocks, w.blocks_b, time, table); + project(w.head, w.head_b, time, head); + // the blocks' table is shift, scale, gate for attention then the MLP; the head's is shift, scale + const from = [_]struct { Vector, []const f32, bool }{ + .{ .shift_a, table[0..h], false }, .{ .scale_a, table[h .. 2 * h], true }, + .{ .gate_a, table[2 * h .. 3 * h], false }, .{ .shift_f, table[3 * h .. 4 * h], false }, + .{ .scale_f, table[4 * h .. 5 * h], true }, .{ .gate_f, table[5 * h .. 6 * h], false }, + .{ .shift_out, head[0..h], false }, .{ .scale_out, head[h .. 2 * h], true }, + }; + for (from) |f| for (&out[@backingInt(f[0])], f[1]) |*o, v| { + o.* = bf16(if (f[2]) 1.0 + v else v); + }; + @memset(&out[@backingInt(Vector.ones)], bf16(1.0)); +} + +test "bf16 rounds to nearest even and keeps NaN" { + try std.testing.expectEqual(@as(u16, 0x3f80), bf16(1.0)); + try std.testing.expectEqual(@as(u16, 0x3f80), bf16(@bitCast(@as(u32, 0x3f808000)))); + try std.testing.expectEqual(@as(u16, 0x3f82), bf16(@bitCast(@as(u32, 0x3f818000)))); + try std.testing.expect(std.math.isNan(widen(bf16(std.math.nan(f32))))); +} + +test "rotary positions: the image centred at the prompt's length, then the prompt along the frame axis" { + const a = std.testing.allocator; + const text = 3; + const rows = 4; + const columns = 6; + const n = (text + rows * columns) * 64; + const cos = try a.alloc(f32, n); + defer a.free(cos); + const sin = try a.alloc(f32, n); + defer a.free(sin); + var bands: [cfg.bands]f32 = undefined; + for (&bands, 0..) |*b, i| b.* = 1.0 / std.math.pow(f32, 10000, @as(f32, @floatFromInt(i)) / cfg.bands); + rotary(cos, sin, &bands, text, rows, columns); + // band 0 of each axis turns by its coordinate: axis 0 never, frame at 16, height at 32, width at 48 + try std.testing.expectEqual(@as(f32, 0), sin[0]); + try std.testing.expectEqual(@as(f32, 1), cos[0]); + try std.testing.expectEqual(round(@sin(@as(f32, 3))), sin[16]); + try std.testing.expectEqual(round(@sin(@as(f32, -1.5))), sin[32]); + try std.testing.expectEqual(round(@sin(@as(f32, -2.5))), sin[48]); + const last_image = (rows * columns - 1) * 64; + try std.testing.expectEqual(round(@sin(@as(f32, 1.5))), sin[last_image + 32]); + try std.testing.expectEqual(round(@sin(@as(f32, 2.5))), sin[last_image + 48]); + const second_text = (rows * columns + 1) * 64; + try std.testing.expectEqual(round(@sin(@as(f32, 1))), sin[second_text + 16]); + try std.testing.expectEqual(@as(f32, 0), sin[second_text + 32]); +} + +test "project sums in fp32 and adds the bias before its one rounding" { + const row = [_]u16{ bf16(1), bf16(2), bf16(3), bf16(4), bf16(5), bf16(6), bf16(7), bf16(8), bf16(9), bf16(0.5) }; + const w = row ++ row; + const x = [_]f32{ 1, 1, 1, 1, 1, 1, 1, 1, 1, 1 }; + const b = [_]u16{ bf16(0.25), bf16(0.0) }; + var y: [2]f32 = undefined; + project(&w, &b, &x, &y); + try std.testing.expectEqual(@as(f32, 45.75), y[0]); + try std.testing.expectEqual(@as(f32, 45.5), y[1]); + // 45.6 + 0.125 rounds to 45.75; rounding 45.6 first (45.5) would leave a tie at 45.625 that goes to 45.5 + const one = [_]u16{bf16(1)}; + const eighth = [_]u16{bf16(0.125)}; + var z: [1]f32 = undefined; + project(&one, &eighth, &[_]f32{45.6}, &z); + try std.testing.expectEqual(@as(f32, 45.75), z[0]); +} diff --git a/zig/src/families/moss_image/kernels.zig b/zig/src/families/moss_image/kernels.zig new file mode 100644 index 0000000000..dd15caee7b --- /dev/null +++ b/zig/src/families/moss_image/kernels.zig @@ -0,0 +1,133 @@ +//! MOSS-Image's Metal pipelines: core_qwen_vision's bf16 and fp32-sum projections, then its own source. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +const frags = @import("../../core/frags.zig"); +const Ref = @import("../../core/qwen_vl/weights.zig").Ref; + +/// Each pipeline by its function: the tower's four bf16 projections, the two fp32-sum ones, then this family's. +pub const names = [_][]const u8{ + "qwv_gemm_bf16_nt_00", + "qwv_gemm_bf16_nt_01", + "qwv_gemm_bf16_nt_10", + "qwv_gemm_bf16_nt_11", + "custom_kernel_tf_gemm_splitk_nax_bf16_n_t_n_t_64_64_256_2_2_bfloat16_t_bfloat16_t_int32_t_float", + "custom_kernel_tf_gemm_splitk_nax_bf16_n_t_t_t_64_64_256_2_2_bfloat16_t_bfloat16_t_int32_t_float", + "mi_norm", + "mi_qk", + "mi_situ", + "mi_gate_add", + "mi_bias", + "mi_attention_00", + "mi_attention_01", + "mi_attention_10", + "mi_attention_11", +}; +const tower_count = 6; +const sums = 4; +const attention_first = 11; + +pub const Kernels = struct { + tower: mtl.Library, + own: mtl.Library, + p: [names.len]mtl.Pipeline, + + pub fn load(gpa: std.mem.Allocator, device: mtl.Device) !Kernels { + const src = try frags.source(device, gpa, sources.core_qwen_vision); + defer gpa.free(src); + var k: Kernels = undefined; + k.tower = try mtl.Library.fromSource(device, src, mtl.CompileOptions.mlx()); + errdefer k.tower.deinit(); + const own = try frags.source(device, gpa, sources.moss_image); + defer gpa.free(own); + k.own = try mtl.Library.fromSource(device, own, mtl.CompileOptions.mlx()); + errdefer k.own.deinit(); + var made: usize = 0; + errdefer for (k.p[0..made]) |p| p.deinit(); + for (names, 0..) |name, i| { + k.p[i] = try mtl.Pipeline.init(device, if (i < tower_count) k.tower else k.own, name, false); + made += 1; + } + return k; + } + + pub fn deinit(k: *Kernels) void { + for (k.p) |p| p.deinit(); + k.own.deinit(); + k.tower.deinit(); + } + + pub fn get(k: *const Kernels, comptime name: []const u8) mtl.Pipeline { + inline for (names, 0..) |n, i| if (comptime std.mem.eql(u8, n, name)) return k.p[i]; + @compileError("no pipeline " ++ name); + } + + /// y [M, N] = x [M, K] w^T in bf16 with fp32 sums (tfp::gemm_nax), w as stored [N, K]. + pub fn linear(k: *const Kernels, e: mtl.ComputeEncoder, x: Ref, w: Ref, y: Ref, M: usize, N: usize, K: usize) void { + const tn = (N + 127) / 128; + const tm = (M + 63) / 64; + e.setPipeline(k.p[@as(usize, @intFromBool(M % 64 == 0)) * 2 + @intFromBool(N % 128 == 0)]); + e.setBuffer(x.buf, x.off, 0); + e.setBuffer(w.buf, w.off, 1); + var p: [16]i32 = @splat(0); + const vals = [_]usize{ M, N, K, K, K, N, tn, tm, 2, K / 256 }; + for (vals, 0..) |v, i| p[i] = @intCast(v); + e.setBytes(std.mem.asBytes(&p), 2); + e.setBuffer(y.buf, y.off, 3); + e.dispatchThreads(mtl.Size.of(tn * 4 * 32, (tm + 3) / 4 * 4, 2), mtl.Size.of(32, 4, 2)); + } + + /// y [M, N] = bf16(x w^T + b): one split-K part's fp32 sums into `sum`, the bias, one rounding (N % 64 = 0). + pub fn linearBias(k: *const Kernels, e: mtl.ComputeEncoder, x: Ref, w: Ref, b: Ref, sum: Ref, y: Ref, M: usize, N: usize, K: usize) void { + std.debug.assert(N % 64 == 0); + const tn = N / 64; + const tm = (M + 63) / 64; + const swz: u6 = if (tm <= 3) 0 else 1; + e.setPipeline(k.p[sums + @as(usize, @intFromBool(M % 64 == 0))]); + e.setBuffer(x.buf, x.off, 0); + e.setBuffer(w.buf, w.off, 1); + var p: [16]i32 = @splat(0); + const vals = [_]usize{ M, N, K, K, K, N, tn, tm, 1, M * N, K, swz, 0, 0 }; + for (vals, 0..) |v, i| p[i] = @intCast(v); + e.setBytes(std.mem.asBytes(&p), 2); + e.setBuffer(sum.buf, sum.off, 3); + const tiles = (tn << swz) * ((tm + (@as(usize, 1) << swz) - 1) >> swz); + e.dispatchThreads(mtl.Size.of(tiles * 32, 2, 2), mtl.Size.of(32, 2, 2)); + bind(e, k.get("mi_bias"), &.{ sum, b, y }, @as(u32, @intCast(N))); + e.dispatchThreads(mtl.Size.of(N, M, 1), mtl.Size.of(256, 1, 1)); + } + + /// o = softmax(q k^T / sqrt(128)) v, every row to every row: q, k, v [heads, L, 128], o rows [L, heads * 128]. + pub fn attention(k: *const Kernels, e: mtl.ComputeEncoder, q: Ref, kb: Ref, vb: Ref, o: Ref, L: usize, heads: usize) void { + const nq = (L + 63) / 64; + const nk = (L + 31) / 32; + e.setPipeline(k.p[attention_first + @as(usize, @intFromBool(L % 64 == 0)) * 2 + @intFromBool(L % 32 == 0)]); + for ([_]Ref{ q, kb, vb }, 0..) |r, i| e.setBuffer(r.buf, r.off, i); + const d = 128; + const vals = [_]usize{ 1, heads, L, L, 1, nq, nk, L / 64, L / 32, L % 64, L % 32, 0, 0, L * d, d, 0, L * d, d, 0, L * d, d, 0, d, heads * d }; + var p: [32]i32 = @splat(0); + for (vals, 0..) |v, i| p[i] = @intCast(v); + e.setBytes(std.mem.asBytes(&p), 3); + var f: [16]f32 = @splat(0); + f[0] = 1.0 / @sqrt(@as(f32, d)); + e.setBytes(std.mem.asBytes(&f), 4); + e.setBuffer(o.buf, o.off, 5); + e.dispatchThreads(mtl.Size.of(nq * 32, heads * 4, 1), mtl.Size.of(32, 4, 1)); + } + + /// One kernel's buffers in order, then its parameters at the next index. + pub fn bind(e: mtl.ComputeEncoder, p: mtl.Pipeline, refs: []const Ref, params: anytype) void { + e.setPipeline(p); + for (refs, 0..) |r, i| e.setBuffer(r.buf, r.off, i); + e.setValue(params, refs.len); + } +}; + +test "every MOSS-Image pipeline compiles and loads on this GPU" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + var k = try Kernels.load(std.testing.allocator, device); + k.deinit(); +} diff --git a/zig/src/families/moss_image/model.zig b/zig/src/families/moss_image/model.zig new file mode 100644 index 0000000000..e553b0a968 --- /dev/null +++ b/zig/src/families/moss_image/model.zig @@ -0,0 +1,309 @@ +//! MOSS-Image's transformer on Metal: a prompt refined once, then image and prompt rows through every joint block. +const std = @import("std"); +const mtl = @import("metal"); +const cfg = @import("config.zig"); +const host = @import("host.zig"); +const wts = @import("weights.zig"); +const ks = @import("kernels.zig"); +const Ref = @import("../../core/qwen_vl/weights.zig").Ref; + +const H = cfg.hidden; +const F = cfg.ffn; +const D = cfg.head_dim; +const half = D / 2; + +/// One buffer's parts for a `text`-row prompt at most and `image` rows, each 256-aligned. +pub const Scratch = struct { + parts: [17]usize, + + fn mul(a: usize, b: usize) error{ScratchOverflow}!usize { + return std.math.mul(usize, a, b) catch error.ScratchOverflow; + } + + fn add(a: usize, b: usize) error{ScratchOverflow}!usize { + return std.math.add(usize, a, b) catch error.ScratchOverflow; + } + + pub fn of(text: usize, image: usize) !Scratch { + if (text == 0) return error.EmptyPrompt; + if (text > cfg.max_text) return error.PromptTooLong; + if (image == 0) return error.BadImageSize; + const rows = try add(text, image); + const wide = try mul(rows, H * 2); + const sums = try mul(@max(try mul(image, H), try mul(text, 3 * H)), 4); + const lat = try mul(image, cfg.channels * 2); + if (rows > std.math.maxInt(u32) / (2 * F)) return error.ImageTooLarge; + return .{ .parts = .{ + wide, wide, wide, wide, wide, wide, wide, try mul(rows, 2 * F * 2), + try mul(rows, F * 2), wide, wide, wide, sums, lat, lat, try mul(text, cfg.context * 2), + host.vectors * H * 2, + } }; + } + + pub fn bytes(s: Scratch) !usize { + var total: usize = 0; + for (s.parts) |n| total = try add(total, (try add(n, 255)) & ~@as(usize, 255)); + return total; + } +}; + +pub const Timing = struct { wall: f64, gpu: f64, host: f64 }; + +/// A refined prompt: its rows after the refiner [rows, 3072] bf16, and the rotary tables of the image and its rows. +pub const Prompt = struct { + rows: usize, + text: mtl.Buffer, + /// cos then sin, each (image rows + rows, 64) fp32 holding bf16 values. + rope: mtl.Buffer, + + pub fn deinit(p: *Prompt) void { + p.rope.deinit(); + p.text.deinit(); + } + + fn cos(p: *const Prompt) Ref { + return .{ .buf = p.rope, .off = 0 }; + } + + fn sin(p: *const Prompt, image: usize) Ref { + return .{ .buf = p.rope, .off = (image + p.rows) * half * 4 }; + } +}; + +pub const Model = struct { + gpa: std.mem.Allocator, + device: mtl.Device, + queue: mtl.Queue, + k: ks.Kernels, + w: wts.Weights, + max_text: usize, + image: usize, + rows: usize, + columns: usize, + scratch: mtl.Buffer, + // bf16 rows: the stream, a norm's output, q k v, the heads side by side, a projection's output + x: Ref, + xn: Ref, + qr: Ref, + kr: Ref, + vr: Ref, + att: Ref, + y: Ref, + /// the MLP's (hidden | gate) rows, and the refiner's packed q k v or its up and gate rows + ug: Ref, + /// situ's output rows + s: Ref, + // bf16 heads [heads, rows, 128] + q: Ref, + k_: Ref, + v: Ref, + /// a biased projection's fp32 sums + sums: Ref, + lat: Ref, + vel: Ref, + ctx: Ref, + mods: Ref, + /// The timestep `mods` holds, if any. + timestep: ?f32 = null, + + /// The transformer of folder `dir`, with scratch for prompts of up to `max_text` rows and a width x height image. + pub fn load(gpa: std.mem.Allocator, io: std.Io, device: mtl.Device, dir: []const u8, max_text: usize, width: usize, height: usize) !*Model { + const g = try cfg.grid(width, height); + const image = try Scratch.mul(g.rows, g.columns); + const plan = try Scratch.of(max_text, image); + const scratch_bytes = try plan.bytes(); + const t = try gpa.create(Model); + errdefer gpa.destroy(t); + t.* = .{ .gpa = gpa, .device = device, .queue = undefined, .k = undefined, .w = undefined, .max_text = max_text, .image = image, .rows = g.rows, .columns = g.columns, .scratch = undefined, .x = undefined, .xn = undefined, .qr = undefined, .kr = undefined, .vr = undefined, .att = undefined, .y = undefined, .ug = undefined, .s = undefined, .q = undefined, .k_ = undefined, .v = undefined, .sums = undefined, .lat = undefined, .vel = undefined, .ctx = undefined, .mods = undefined }; + t.queue = try device.queue(); + errdefer t.queue.deinit(); + t.k = try ks.Kernels.load(gpa, device); + errdefer t.k.deinit(); + t.w = try wts.load(gpa, io, device, dir); + errdefer t.w.deinit(); + t.scratch = try device.buffer(scratch_bytes, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + var at: usize = 0; + const views = [_]*Ref{ &t.x, &t.xn, &t.qr, &t.kr, &t.vr, &t.att, &t.y, &t.ug, &t.s, &t.q, &t.k_, &t.v, &t.sums, &t.lat, &t.vel, &t.ctx, &t.mods }; + for (views, plan.parts) |view, n| { + view.* = .{ .buf = t.scratch, .off = at }; + at += std.mem.alignForward(usize, n, 256); + } + return t; + } + + pub fn deinit(t: *Model) void { + t.scratch.deinit(); + t.w.deinit(); + t.k.deinit(); + t.queue.deinit(); + t.gpa.destroy(t); + } + + fn floats(r: Ref, n: usize) []f32 { + return @as([*]f32, @ptrCast(@alignCast(r.addr())))[0..n]; + } + + fn halves(r: Ref, n: usize) []u16 { + return @as([*]u16, @ptrCast(@alignCast(r.addr())))[0..n]; + } + + fn vec(t: *const Model, which: host.Vector) Ref { + return t.mods.at(@backingInt(which) * H * 2); + } + + /// A prompt's rows (rows, 4096) bf16, the text encoder's last hidden states, through the refiner; with its tables. + pub fn prompt(t: *Model, embeds: []const u16) !Prompt { + if (embeds.len == 0 or embeds.len % cfg.context != 0) return error.PromptRows; + const T = embeds.len / cfg.context; + if (T > t.max_text) return error.PromptTooLong; + const opts = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; + var p: Prompt = .{ .rows = T, .text = try t.device.buffer(T * H * 2, opts), .rope = undefined }; + errdefer p.text.deinit(); + const table = (t.image + T) * half; + p.rope = try t.device.buffer(2 * table * 4, opts); + errdefer p.rope.deinit(); + host.rotary(floats(p.cos(), table), floats(p.sin(t.image), table), &t.w.bands, T, t.rows, t.columns); + @memcpy(halves(t.ctx, embeds.len), embeds); + @memset(halves(t.vec(.ones), H), host.bf16(1.0)); + const started = mtl.clock.seconds(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const cb = t.queue.commandBuffer(); + const e = cb.compute(.serial); + t.k.linearBias(e, t.ctx, t.w.text_in, t.w.text_in_b, t.sums, t.x, T, H, cfg.context); + const rope = t.image * half * 4; + for (t.w.text) |b| t.textBlock(e, b, T, p.cos().at(rope), p.sin(t.image).at(rope)); + _ = try finish(cb, e, started, 0); + @memcpy(halves(.{ .buf = p.text }, T * H), halves(t.x, T * H)); + return p; + } + + /// The velocity of `latents` (image rows, 128) at `timestep` under prompt `p`, skipping block `skip` if given. + pub fn forward(t: *Model, p: *const Prompt, latents: []const f32, timestep: f32, skip: ?usize, velocity: []f32) !Timing { + const R = t.image; + if (latents.len != R * cfg.channels or velocity.len != latents.len) return error.LatentRows; + if (skip) |s| if (s >= cfg.layers) return error.BadSkipLayer; + const started = mtl.clock.seconds(); + var spent: f64 = 0; + if (t.timestep != timestep) { + const vectors: *[host.vectors][H]u16 = @ptrCast(@alignCast(t.mods.addr())); + t.timestep = null; + try host.modulation(vectors, t.gpa, t.w.modulation, timestep); + t.timestep = timestep; + spent = mtl.clock.seconds() - started; + } + for (halves(t.lat, latents.len), latents) |*o, v| o.* = host.bf16(v); + const N = R + p.rows; + @memcpy(halves(t.x.at(R * H * 2), p.rows * H), halves(.{ .buf = p.text }, p.rows * H)); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const cb = t.queue.commandBuffer(); + const e = cb.compute(.serial); + t.k.linearBias(e, t.lat, t.w.video_in, t.w.video_in_b, t.sums, t.x, R, H, cfg.channels); + for (t.w.blocks, 0..) |b, i| if (skip != i) t.block(e, b, N, p.cos(), p.sin(R)); + t.norm(e, R, t.w.final_norm, .{ .scale = .scale_out, .shift = .shift_out }); + t.k.linear(e, t.xn, t.w.final_out, t.vel, R, cfg.channels, H); + const timing = try finish(cb, e, started, spent); + for (velocity, halves(t.vel, velocity.len)) |*o, v| o.* = host.widen(v); + return timing; + } + + fn finish(cb: mtl.CommandBuffer, e: mtl.ComputeEncoder, started: f64, spent: f64) !Timing { + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.log.err("moss image: a pass failed on the GPU: {s}", .{msg}); + return error.GpuFailure; + } + return .{ .wall = mtl.clock.seconds() - started, .gpu = cb.gpuSeconds(), .host = spent }; + } + + const Modulated = struct { scale: host.Vector, shift: host.Vector }; + + /// xn = the RMSNorm of the stream's `rows` rows with weight `w`, then modulated when given. + fn norm(t: *const Model, e: mtl.ComputeEncoder, rows: usize, w: Ref, m: ?Modulated) void { + const s = if (m) |v| t.vec(v.scale) else t.vec(.ones); + const b = if (m) |v| t.vec(v.shift) else t.vec(.ones); + const p = extern struct { width: u32, eps: f32, modulated: u32 }{ .width = H, .eps = cfg.eps, .modulated = @intFromBool(m != null) }; + ks.Kernels.bind(e, t.k.get("mi_norm"), &.{ t.x, w, s, b, t.xn }, p); + e.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(256, 1, 1)); + } + + /// x += y gated by shared vector `gate` (ones: plain). + fn gateAdd(t: *const Model, e: mtl.ComputeEncoder, rows: usize, gate: host.Vector) void { + ks.Kernels.bind(e, t.k.get("mi_gate_add"), &.{ t.x, t.y, t.vec(gate) }, @as(u32, H)); + e.dispatchThreads(mtl.Size.of(H, rows, 1), mtl.Size.of(256, 1, 1)); + } + + /// q k v rows of `stride` to normed, rotated heads in q, k_, v. + fn heads(t: *const Model, e: mtl.ComputeEncoder, qr: Ref, kr: Ref, vr: Ref, stride: usize, nq: Ref, nk: Ref, cos: Ref, sin: Ref, rows: usize) void { + const p = extern struct { rows: u32, heads: u32, stride: u32, eps: f32 }{ .rows = @intCast(rows), .heads = cfg.heads, .stride = @intCast(stride), .eps = cfg.eps }; + ks.Kernels.bind(e, t.k.get("mi_qk"), &.{ qr, kr, vr, nq, nk, cos, sin, t.q, t.k_, t.v }, p); + e.dispatchGroups(mtl.Size.of(rows, cfg.heads / 8, 1), mtl.Size.of(256, 1, 1)); + } + + /// s = situ(hidden, gate) of rows `stride` apart. + fn situ(t: *const Model, e: mtl.ComputeEncoder, hidden: Ref, gate: Ref, rows: usize, stride: usize) void { + ks.Kernels.bind(e, t.k.get("mi_situ"), &.{ hidden, gate, t.s }, [2]u32{ F, @intCast(stride) }); + e.dispatchThreads(mtl.Size.of(F, rows, 1), mtl.Size.of(256, 1, 1)); + } + + /// One joint block over the image and prompt rows (`rows` in all), under the shared modulation. + fn block(t: *const Model, e: mtl.ComputeEncoder, b: wts.Block, rows: usize, cos: Ref, sin: Ref) void { + t.norm(e, rows, b.norm1, .{ .scale = .scale_a, .shift = .shift_a }); + t.k.linear(e, t.xn, b.q, t.qr, rows, H, H); + t.k.linear(e, t.xn, b.k, t.kr, rows, H, H); + t.k.linear(e, t.xn, b.v, t.vr, rows, H, H); + t.heads(e, t.qr, t.kr, t.vr, H, b.q_norm, b.k_norm, cos, sin, rows); + t.k.attention(e, t.q, t.k_, t.v, t.att, rows, cfg.heads); + t.k.linear(e, t.att, b.o, t.y, rows, H, H); + t.gateAdd(e, rows, .gate_a); + t.norm(e, rows, b.norm2, .{ .scale = .scale_f, .shift = .shift_f }); + t.k.linear(e, t.xn, b.up_gate, t.ug, rows, 2 * F, H); + t.situ(e, t.ug, t.ug.at(F * 2), rows, 2 * F); + t.k.linear(e, t.s, b.down, t.y, rows, H, F); + t.gateAdd(e, rows, .gate_f); + } + + /// One refiner block over the prompt's `rows` rows: plain norms and residuals, biased projections. + fn textBlock(t: *const Model, e: mtl.ComputeEncoder, b: wts.TextBlock, rows: usize, cos: Ref, sin: Ref) void { + t.norm(e, rows, b.norm1, null); + t.k.linearBias(e, t.xn, b.qkv, b.qkv_b, t.sums, t.ug, rows, 3 * H, H); + t.heads(e, t.ug, t.ug.at(H * 2), t.ug.at(2 * H * 2), 3 * H, b.q_norm, b.k_norm, cos, sin, rows); + t.k.attention(e, t.q, t.k_, t.v, t.att, rows, cfg.heads); + t.k.linearBias(e, t.att, b.proj, b.proj_b, t.sums, t.y, rows, H, H); + t.gateAdd(e, rows, .ones); + t.norm(e, rows, b.norm2, null); + const gate = t.ug.at(rows * F * 2); + t.k.linearBias(e, t.xn, b.up, b.up_b, t.sums, t.ug, rows, F, H); + t.k.linearBias(e, t.xn, b.gate, b.gate_b, t.sums, gate, rows, F, H); + t.situ(e, t.ug, gate, rows, F); + t.k.linearBias(e, t.s, b.down, b.down_b, t.sums, t.y, rows, H, F); + t.gateAdd(e, rows, .ones); + } +}; + +test "scratch for a 1024 x 1024 image after prompts of up to 4,096 rows" { + const s = try Scratch.of(4096, 4096); + try std.testing.expectEqual(@as(usize, 1_092_671_488), try s.bytes()); +} + +test "scratch rejects row products, alignment and total overflow without allocating" { + const huge = std.math.maxInt(usize); + try std.testing.expectError(error.ScratchOverflow, Scratch.of(1, huge)); + try std.testing.expectError(error.ScratchOverflow, Scratch.of(1, huge / 2)); + const alignment: Scratch = .{ .parts = @splat(huge) }; + try std.testing.expectError(error.ScratchOverflow, alignment.bytes()); + const total: Scratch = .{ .parts = @splat(huge / 4 + 1) }; + try std.testing.expectError(error.ScratchOverflow, total.bytes()); + try std.testing.expectError(error.EmptyPrompt, Scratch.of(0, 1)); + try std.testing.expectError(error.PromptTooLong, Scratch.of(cfg.max_text + 1, 1)); + try std.testing.expectError(error.BadImageSize, Scratch.of(1, 0)); +} + +test "model load validates huge dimensions before touching a Metal device or checkpoint" { + const huge = std.math.maxInt(usize) - (cfg.pixels_per_row - 1); + try std.testing.expectError(error.ImageTooLarge, Model.load(std.testing.allocator, std.testing.io, undefined, "unused", 1, huge, huge)); + try std.testing.expectError(error.PromptTooLong, Model.load(std.testing.allocator, std.testing.io, undefined, "unused", cfg.max_text + 1, 16, 16)); +} diff --git a/zig/src/families/moss_image/moss_image.zig b/zig/src/families/moss_image/moss_image.zig new file mode 100644 index 0000000000..739c2dd08a --- /dev/null +++ b/zig/src/families/moss_image/moss_image.zig @@ -0,0 +1,17 @@ +//! MOSS-Image text to image: its transformer on Metal (a denoiser beside the lane engine, not a token family). +pub const config = @import("config.zig"); +pub const host = @import("host.zig"); +pub const weights = @import("weights.zig"); +pub const kernels = @import("kernels.zig"); +pub const model = @import("model.zig"); +pub const Model = model.Model; +pub const Prompt = model.Prompt; +pub const Timing = model.Timing; + +test { + _ = config; + _ = host; + _ = weights; + _ = kernels; + _ = model; +} diff --git a/zig/src/families/moss_image/weights.zig b/zig/src/families/moss_image/weights.zig new file mode 100644 index 0000000000..1758e22bf0 --- /dev/null +++ b/zig/src/families/moss_image/weights.zig @@ -0,0 +1,118 @@ +//! MOSS-Image's transformer tensors from a diffusers checkpoint folder, bf16 as stored, into shared Metal buffers. +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("config.zig"); +const host = @import("host.zig"); +const Ref = @import("../../core/qwen_vl/weights.zig").Ref; + +/// A joint block: attention projections, the MLP's (hidden | gate) and its down projection, then the norms. +pub const Block = struct { q: Ref, k: Ref, v: Ref, o: Ref, up_gate: Ref, down: Ref, norm1: Ref, norm2: Ref, q_norm: Ref, k_norm: Ref }; + +/// A prompt refiner block: q, k, v packed in one projection, every projection with a bias, no modulation. +pub const TextBlock = struct { norm1: Ref, norm2: Ref, q_norm: Ref, k_norm: Ref, qkv: Ref, qkv_b: Ref, proj: Ref, proj_b: Ref, gate: Ref, gate_b: Ref, up: Ref, up_b: Ref, down: Ref, down_b: Ref }; + +pub const Weights = struct { + ck: ck.Checkpoint, + blocks: [cfg.layers]Block, + text: [cfg.text_layers]TextBlock, + text_in: Ref, + text_in_b: Ref, + video_in: Ref, + video_in_b: Ref, + final_norm: Ref, + final_out: Ref, + /// The rotary's frequencies as the checkpoint stores them (fp32). + bands: [cfg.bands]f32, + modulation: host.Weights, + + pub fn deinit(w: *Weights) void { + w.ck.deinit(); + } +}; + +/// The transformer's shard files in `dir`: the index's weight map, else the one diffusion_pytorch_model.safetensors. +fn shards(arena: std.mem.Allocator, io: std.Io, dir: []const u8) ![][:0]const u8 { + var files: std.StringArrayHashMapUnmanaged(void) = .empty; + const index = try std.fs.path.join(arena, &.{ dir, "diffusion_pytorch_model.safetensors.index.json" }); + if (std.Io.Dir.cwd().readFileAlloc(io, index, arena, .limited(16 << 20))) |text| { + const doc = try std.json.parseFromSliceLeaky(std.json.Value, arena, text, .{}); + var it = (try index_format.weightMap(doc)).iterator(); + while (it.next()) |kv| try files.put(arena, kv.value_ptr.string, {}); + } else |err| { + if (err != error.FileNotFound) return err; + try files.put(arena, "diffusion_pytorch_model.safetensors", {}); + } + const out = try arena.alloc([:0]const u8, files.count()); + for (files.keys(), out) |name, *path| path.* = try std.fmt.allocPrintSentinel(arena, "{s}/{s}", .{ dir, name }, 0); + return out; +} + +/// The tensor `name` of `dtype` and exactly `shape`. +fn typed(c: *const ck.Checkpoint, name: []const u8, dtype: ck.DType, shape: []const usize) !Ref { + const t = try c.get(name); + if (t.dtype != dtype or t.rank != shape.len or !std.mem.eql(usize, t.shape[0..t.rank], shape)) { + std.log.err("moss image: {s} is {t} {any}; the transformer reads {t} {any}", .{ name, t.dtype, t.shape[0..t.rank], dtype, shape }); + return error.UnexpectedTensor; + } + if (t.offset % 16 != 0) return error.UnalignedTensor; + return .{ .buf = t.buffer, .off = t.offset }; +} + +fn take(c: *const ck.Checkpoint, name: []const u8, shape: []const usize) !Ref { + return typed(c, name, .bf16, shape); +} + +fn values(r: Ref, n: usize) []const u16 { + return @as([*]const u16, @ptrCast(@alignCast(r.addr())))[0..n]; +} + +/// The transformer folder `dir` (a checkpoint's transformer/): config checked, then every tensor this family reads. +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_path = try std.fs.path.join(arena, &.{ dir, "config.json" }); + const config = try std.Io.Dir.cwd().readFileAlloc(io, config_path, arena, .limited(1 << 20)); + try cfg.check(gpa, config); + var w: Weights = undefined; + w.ck = ck.Checkpoint.init(gpa); + errdefer w.ck.deinit(); + for (try shards(arena, io, dir)) |path| try w.ck.addFile(device, path, ""); + const h = cfg.hidden; + const f = cfg.ffn; + const d = cfg.head_dim; + var name: [96]u8 = undefined; + for (&w.blocks, 0..) |*b, i| { + const fields = .{ .{ "proj.q.weight", &b.q, &[_]usize{ h, h } }, .{ "proj.k.weight", &b.k, &[_]usize{ h, h } }, .{ "proj.v.weight", &b.v, &[_]usize{ h, h } }, .{ "proj.o.weight", &b.o, &[_]usize{ h, h } }, .{ "proj.ffn_up_gate.weight", &b.up_gate, &[_]usize{ 2 * f, h } }, .{ "proj.ffn_down.weight", &b.down, &[_]usize{ h, f } }, .{ "norm1.weight", &b.norm1, &[_]usize{h} }, .{ "norm2.weight", &b.norm2, &[_]usize{h} }, .{ "q_norm.weight", &b.q_norm, &[_]usize{d} }, .{ "k_norm.weight", &b.k_norm, &[_]usize{d} } }; + inline for (fields) |x| x[1].* = try take(&w.ck, try std.fmt.bufPrint(&name, "layers.{d}.{s}", .{ i, x[0] }), x[2]); + } + for (&w.text, 0..) |*b, i| { + const fields = .{ .{ "norm1.weight", &b.norm1, &[_]usize{h} }, .{ "norm2.weight", &b.norm2, &[_]usize{h} }, .{ "q_norm.weight", &b.q_norm, &[_]usize{d} }, .{ "k_norm.weight", &b.k_norm, &[_]usize{d} }, .{ "qkv.weight", &b.qkv, &[_]usize{ 3 * h, h } }, .{ "qkv.bias", &b.qkv_b, &[_]usize{3 * h} }, .{ "proj.weight", &b.proj, &[_]usize{ h, h } }, .{ "proj.bias", &b.proj_b, &[_]usize{h} }, .{ "ffn_gate.weight", &b.gate, &[_]usize{ f, h } }, .{ "ffn_gate.bias", &b.gate_b, &[_]usize{f} }, .{ "ffn_up.weight", &b.up, &[_]usize{ f, h } }, .{ "ffn_up.bias", &b.up_b, &[_]usize{f} }, .{ "ffn_down.weight", &b.down, &[_]usize{ h, f } }, .{ "ffn_down.bias", &b.down_b, &[_]usize{h} } }; + inline for (fields) |x| x[1].* = try take(&w.ck, try std.fmt.bufPrint(&name, "text_embedder.blocks.{d}.{s}", .{ i, x[0] }), x[2]); + } + w.text_in = try take(&w.ck, "text_embedder.proj.weight", &.{ h, cfg.context }); + w.text_in_b = try take(&w.ck, "text_embedder.proj.bias", &.{h}); + w.video_in = try take(&w.ck, "video_embedder.weight", &.{ h, cfg.channels }); + w.video_in_b = try take(&w.ck, "video_embedder.bias", &.{h}); + w.final_norm = try take(&w.ck, "final_norm_video.weight", &.{h}); + w.final_out = try take(&w.ck, "final_linear_video.weight", &.{ cfg.channels, h }); + const bands = try typed(&w.ck, "rope.bands", .f32, &.{cfg.bands}); + w.bands = @as(*const [cfg.bands]f32, @ptrCast(@alignCast(bands.addr()))).*; + const parts = .{ .{ "time_embedding.0.weight", [_]usize{ h, cfg.time_dim } }, .{ "time_embedding.0.bias", [_]usize{h} }, .{ "time_embedding.2.weight", [_]usize{ h, h } }, .{ "time_embedding.2.bias", [_]usize{h} }, .{ "shared_block_modulation.weight", [_]usize{ 6 * h, h } }, .{ "shared_block_modulation.bias", [_]usize{6 * h} }, .{ "video_head_modulation.weight", [_]usize{ 2 * h, h } }, .{ "video_head_modulation.bias", [_]usize{2 * h} } }; + var host_parts: [parts.len][]const u16 = undefined; + inline for (parts, 0..) |p, i| { + const shape: [p[1].len]usize = p[1]; + var n: usize = 1; + for (shape) |s| n *= s; + host_parts[i] = values(try take(&w.ck, p[0], &shape), n); + } + w.modulation = .{ .time_1 = host_parts[0], .time_1_b = host_parts[1], .time_2 = host_parts[2], .time_2_b = host_parts[3], .blocks = host_parts[4], .blocks_b = host_parts[5], .head = host_parts[6], .head_b = host_parts[7] }; + const read = 10 * cfg.layers + 14 * cfg.text_layers + 2 + 2 + 2 + 1 + parts.len; + if (w.ck.tensors.count() != read) { + std.log.err("moss image: the checkpoint has {d} transformer tensors; this family reads {d}", .{ w.ck.tensors.count(), read }); + return error.UnexpectedTensors; + } + return w; +} diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index 47b6a354c4..4fa4fbe147 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -11,6 +11,8 @@ pub const qwen_vl = @import("core/qwen_vl/qwen_vl.zig"); pub const qwen_vl_tower = @import("core/qwen_vl/metal.zig"); /// Qwen-Image-2.1's transformer on Metal: a denoiser driven by tf-qwen-image, outside the lane engine. pub const qwen_image = @import("families/qwen_image/qwen_image.zig"); +/// MOSS-Image's transformer on Metal: a denoiser driven by tf-moss-image, outside the lane engine. +pub const moss_image = @import("families/moss_image/moss_image.zig"); pub const checkpoint = @import("core/checkpoint_metal.zig"); pub const shard_edit = @import("core/shard_edit.zig"); pub const lora_merge = @import("core/lora_merge_dir.zig"); diff --git a/zig/tests/contract/metal_limits.zig b/zig/tests/contract/metal_limits.zig index ab64f33ed6..bf7109eca4 100644 --- a/zig/tests/contract/metal_limits.zig +++ b/zig/tests/contract/metal_limits.zig @@ -85,6 +85,12 @@ fn launchOf(a: std.mem.Allocator, device: mtl.Device, d: limits.MetalLaunch) !La } else if (std.mem.startsWith(u8, d.name, "qwen image ")) { text = sources.qwen_image; frags = true; + } else if (std.mem.startsWith(u8, d.name, "moss image gemm ")) { + text = sources.core_qwen_vision; + frags = true; + } else if (std.mem.startsWith(u8, d.name, "moss image ")) { + text = sources.moss_image; + frags = true; } else if (std.mem.startsWith(u8, d.name, "core deltanet ")) { text = sources.core_deltanet; } else if (std.mem.startsWith(u8, d.name, "qwen27 ")) { diff --git a/zig/tests/moss_image.zig b/zig/tests/moss_image.zig new file mode 100644 index 0000000000..35908917e0 --- /dev/null +++ b/zig/tests/moss_image.zig @@ -0,0 +1,178 @@ +//! tf-moss-image: MOSS-Image's transformer on a case from tools/zig/moss_image_case.py, compared and timed. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const mi = tf.moss_image; + +const usage = + \\tf-moss-image TRANSFORMER_DIR CASE OUT_DIR + \\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. +; + +/// The card's sampling defaults: the guidance scale, and the block the negative branch skips. +const guidance: f32 = 3.0; +const negative_skip: usize = 22; + +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 f32, 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]) |a, b| { + const d = @as(f64, a) - b; + dot += @as(f64, a) * b; + aa += @as(f64, a) * a; + bb += @as(f64, b) * b; + rd += d * d; + rb += @as(f64, b) * b; + most = @max(most, @abs(d)); + } + worst = @max(worst, @sqrt(rd / @max(rb, 1e-30))); + } + return .{ .max_abs = most, .cosine = dot / @sqrt(aa * bb), .worst_row = worst }; +} + +/// The case tensor `name` of `dtype`, copied out as T values. +fn take(comptime T: type, a: std.mem.Allocator, f: *const tf.safetensors.File, name: []const u8, dtype: tf.safetensors.DType) ![]T { + const t = f.get(name) orelse { + std.debug.print("the case has no {s}\n", .{name}); + return error.MissingTensor; + }; + if (t.dtype != dtype) return error.BadCaseTensor; + const out = try a.alloc(T, t.numel()); + @memcpy(std.mem.sliceAsBytes(out), t.bytes); + return out; +} + +fn dims(f: *const tf.safetensors.File, name: []const u8) ![]const usize { + const e = f.names.getPtr(name) orelse return error.MissingTensor; + return e.shape[0..e.rank]; +} + +fn print(list: *std.ArrayList(u8), a: std.mem.Allocator, comptime fmt: []const u8, args: anytype) !void { + try list.appendSlice(a, try std.fmt.allocPrint(a, fmt, args)); +} + +fn json(list: *std.ArrayList(u8), a: std.mem.Allocator, key: []const u8, m: Match) !void { + try print(list, a, ", \"{s}\": {{\"max_abs\": {d:.6}, \"cosine\": {d:.8}, \"worst_row\": {d:.6}}}", .{ key, m.max_abs, m.cosine, m.worst_row }); +} + +/// Each step's two branches on the reference's input against its velocities; true when a repeated forward repeats. +fn teacher(a: std.mem.Allocator, case: *const tf.safetensors.File, model: *mi.Model, branches: *[2]mi.Prompt, timesteps: []const f32, count: usize, out: *std.ArrayList(u8)) !bool { + const steps = timesteps.len; + const inputs = try take(f32, a, case, "inputs", .f32); + const reference = [2][]f32{ try take(f32, a, case, "velocity", .f32), try take(f32, a, case, "velocity_negative", .f32) }; + const exact = [2]?[]f32{ + if (case.get("velocity_f32") != null) try take(f32, a, case, "velocity_f32", .f32) else null, + if (case.get("velocity_negative_f32") != null) try take(f32, a, case, "velocity_negative_f32", .f32) else null, + }; + if (inputs.len != steps * count) return error.BadCaseTensor; + for (reference) |r| if (r.len != steps * count) return error.BadCaseTensor; + try print(out, a, ", \"teacher_forced\": [", .{}); + const skips = [2]?usize{ null, negative_skip }; + const labels = [2][]const u8{ "prompt", "negative" }; + const velocity = try a.alloc(f32, count); + const again = try a.alloc(f32, count); + for (0..steps) |s| { + try print(out, a, "{s}{{\"step\": {d}, \"timestep\": {d:.4}", .{ if (s == 0) "" else ", ", s + 1, timesteps[s] }); + for (0..2) |b| { + _ = try model.forward(&branches[b], inputs[s * count ..][0..count], timesteps[s], skips[b], velocity); + const theirs = reference[b][s * count ..][0..count]; + const m = match(velocity, theirs, mi.config.channels); + std.debug.print("step {d} {s}: against the bf16 reference max |diff| {d:.5}, cosine {d:.7}, worst row {d:.5}\n", .{ s + 1, labels[b], m.max_abs, m.cosine, m.worst_row }); + try json(out, a, try std.fmt.allocPrint(a, "{s}_bf16", .{labels[b]}), m); + if (exact[b]) |f| if (f.len >= (s + 1) * count) { + const want = f[s * count ..][0..count]; + const e = match(velocity, want, mi.config.channels); + const r = match(theirs, want, mi.config.channels); + std.debug.print(" against the fp32 reference max |diff| {d:.5}, cosine {d:.7}; the bf16 reference's own: {d:.5}, {d:.7}\n", .{ e.max_abs, e.cosine, r.max_abs, r.cosine }); + try json(out, a, try std.fmt.allocPrint(a, "{s}_fp32", .{labels[b]}), e); + try json(out, a, try std.fmt.allocPrint(a, "{s}_bf16_ref_vs_fp32", .{labels[b]}), r); + }; + } + try out.append(a, '}'); + } + _ = try model.forward(&branches[0], inputs[0..count], timesteps[0], null, again); + _ = try model.forward(&branches[1], inputs[0..count], timesteps[1], negative_skip, velocity); + _ = try model.forward(&branches[0], inputs[0..count], timesteps[0], null, velocity); + const repeat = std.mem.eql(u8, std.mem.sliceAsBytes(velocity), std.mem.sliceAsBytes(again)); + std.debug.print("the same step twice, another between: {s}\n", .{if (repeat) "identical bytes" else "DIFFERENT bytes"}); + try print(out, a, "], \"repeat_identical\": {}", .{repeat}); + return repeat; +} + +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) { + std.debug.print("{s}\n", .{usage}); + return 2; + } + var case = try tf.safetensors.File.open(a, io, args[2]); + defer case.close(io); + const grid = try dims(&case, "latents"); + 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); + 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); + const steps = timesteps.len; + if (sigmas.len != steps + 1 or start.len != count) return error.BadCaseTensor; + + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + const loading = mtl.clock.seconds(); + const model = try mi.Model.load(std.heap.page_allocator, io, device, args[1], @max(text, negative_rows), grid[1] * 16, grid[0] * 16); + defer model.deinit(); + const load_s = mtl.clock.seconds() - loading; + std.debug.print("{s}: {d} image rows, prompts of {d} and {d} rows, {d} steps; loaded in {d:.1} s\n", .{ device.name(), rows, text, negative_rows, steps, load_s }); + + var out: std.ArrayList(u8) = .empty; + const refining = mtl.clock.seconds(); + var branches = [2]mi.Prompt{ try model.prompt(prompt), try model.prompt(negative) }; + defer for (&branches) |*b| b.deinit(); + 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 }); + const repeat = if (case.get("inputs") == null) true else try teacher(a, &case, model, &branches, timesteps, count, &out); + try print(&out, a, ", \"run\": [", .{}); + + const x = try a.dupe(f32, start); + const cond = try a.alloc(f32, count); + const uncond = try a.alloc(f32, count); + var total: f64 = 0; + for (0..steps) |s| { + const c = try model.forward(&branches[0], x, timesteps[s], null, cond); + const u = try model.forward(&branches[1], x, timesteps[s], negative_skip, uncond); + const by = sigmas[s + 1] - sigmas[s]; + for (x, cond, uncond) |*o, pc, pu| o.* += (pu + guidance * (pc - pu)) * by; + const wall = c.wall + u.wall; + total += wall; + std.debug.print("step {d}/{d}: {d:.3} s ({d:.3} s on the GPU, {d:.3} s modulation on the host)\n", .{ s + 1, steps, wall, c.gpu + u.gpu, c.host + u.host }); + try print(&out, a, "{s}{{\"wall_s\": {d:.4}, \"gpu_s\": {d:.4}, \"host_s\": {d:.4}}}", .{ if (s == 0) "" else ", ", wall, c.gpu + u.gpu, c.host + u.host }); + } + try print(&out, a, "], \"denoise_s\": {d:.3}}}\n", .{total}); + std.debug.print("denoise: {d:.2} s for {d} steps ({d} forwards)\n", .{ total, steps, 2 * steps }); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = try std.fmt.allocPrint(a, "{s}/latents.f32", .{args[3]}), .data = std.mem.sliceAsBytes(x) }); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = try std.fmt.allocPrint(a, "{s}/run.json", .{args[3]}), .data = out.items }); + return if (repeat) 0 else 1; +} From 508e59c7c0971491cef5258d181f4a81c82ffed3 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Tue, 27 Oct 2026 05:32:50 -0500 Subject: [PATCH 2/6] MOSS-Image: its kernels through the contract's gates 1-4, and the M5's tensor units required mi_norm (plain and modulated), mi_qk (both row layouts), mi_situ (both), mi_gate_add, mi_bias and the three attention forms, each output checked at the bf16 points the publisher's torch code rounds at; eight controls each drop or add one rounding point and must go red at gate 3. The launches move to kernels.zig so the tests run the engine's own dispatch. mi_attention_10 goes: no L picks it. The family is qualified on the M5 alone and refuses a GPU without tensor units. --- zig/kernels/metal/moss_image.metal | 3 +- zig/src/core/contract/moss_image_records.zig | 1 - zig/src/families/moss_image/config.zig | 1 + zig/src/families/moss_image/kernels.zig | 80 +++- zig/src/families/moss_image/model.zig | 16 +- zig/tests/contract/metal_main.zig | 2 + zig/tests/contract/metal_moss_image.zig | 397 ++++++++++++++++++ .../contract/metal_moss_image_common.zig | 144 +++++++ zig/tests/contract/metal_moss_image_norms.zig | 363 ++++++++++++++++ zig/tests/contract/metal_moss_image_steps.zig | 340 +++++++++++++++ 10 files changed, 1314 insertions(+), 33 deletions(-) create mode 100644 zig/tests/contract/metal_moss_image.zig create mode 100644 zig/tests/contract/metal_moss_image_common.zig create mode 100644 zig/tests/contract/metal_moss_image_norms.zig create mode 100644 zig/tests/contract/metal_moss_image_steps.zig diff --git a/zig/kernels/metal/moss_image.metal b/zig/kernels/metal/moss_image.metal index 75ebb8c0ba..786744562c 100644 --- a/zig/kernels/metal/moss_image.metal +++ b/zig/kernels/metal/moss_image.metal @@ -123,7 +123,7 @@ struct MiSitu { y[i] = bfloat(c[i] + float(b[at.x])); } -// tfp::attention_nax (MLX's steel attention on the tensor units), full: every row sees every row. +// tfp::attention_nax, full (every row sees every row): forms for L % 32 > 0, L % 32 = 0, and L % 64 = 0. #define MI_ATTENTION(NAME, AQ, AK) \ [[max_total_threads_per_threadgroup(128)]] [[kernel]] void NAME( \ const device bfloat16_t* Q [[buffer(0)]], const device bfloat16_t* K [[buffer(1)]], \ @@ -135,5 +135,4 @@ struct MiSitu { } MI_ATTENTION(mi_attention_00, false, false) MI_ATTENTION(mi_attention_01, false, true) -MI_ATTENTION(mi_attention_10, true, false) MI_ATTENTION(mi_attention_11, true, true) diff --git a/zig/src/core/contract/moss_image_records.zig b/zig/src/core/contract/moss_image_records.zig index 32b7ec030a..2e1a572220 100644 --- a/zig/src/core/contract/moss_image_records.zig +++ b/zig/src/core/contract/moss_image_records.zig @@ -11,7 +11,6 @@ pub const records = [_]M{ .{ .name = "moss image mi_bias", .threads = 256, .tg_memory = 0, .reserved = 256 }, .{ .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_10", .threads = 128, .tg_memory = 0, .reserved = 128 }, .{ .name = "moss image mi_attention_11", .threads = 128, .tg_memory = 0, .reserved = 128 }, }; diff --git a/zig/src/families/moss_image/config.zig b/zig/src/families/moss_image/config.zig index 9153293f18..981165c227 100644 --- a/zig/src/families/moss_image/config.zig +++ b/zig/src/families/moss_image/config.zig @@ -11,6 +11,7 @@ pub const text_layers = 2; pub const context = 4096; pub const time_dim = 256; pub const channels = 128; +/// Every RMSNorm's, fixed in the publisher's _RMSNorm (pipeline_moss_image.py), not a config setting. pub const eps: f32 = 1e-6; /// Rotary: four axes (unused, frame, height, width) of 16 frequencies each, the rotation of halves. pub const bands = head_dim / 8; diff --git a/zig/src/families/moss_image/kernels.zig b/zig/src/families/moss_image/kernels.zig index dd15caee7b..33aa044f01 100644 --- a/zig/src/families/moss_image/kernels.zig +++ b/zig/src/families/moss_image/kernels.zig @@ -3,7 +3,7 @@ const std = @import("std"); const mtl = @import("metal"); const sources = @import("kernel_sources"); const frags = @import("../../core/frags.zig"); -const Ref = @import("../../core/qwen_vl/weights.zig").Ref; +pub const Ref = @import("../../core/qwen_vl/weights.zig").Ref; /// Each pipeline by its function: the tower's four bf16 projections, the two fp32-sum ones, then this family's. pub const names = [_][]const u8{ @@ -20,7 +20,6 @@ pub const names = [_][]const u8{ "mi_bias", "mi_attention_00", "mi_attention_01", - "mi_attention_10", "mi_attention_11", }; const tower_count = 6; @@ -93,26 +92,12 @@ pub const Kernels = struct { e.setBuffer(sum.buf, sum.off, 3); const tiles = (tn << swz) * ((tm + (@as(usize, 1) << swz) - 1) >> swz); e.dispatchThreads(mtl.Size.of(tiles * 32, 2, 2), mtl.Size.of(32, 2, 2)); - bind(e, k.get("mi_bias"), &.{ sum, b, y }, @as(u32, @intCast(N))); - e.dispatchThreads(mtl.Size.of(N, M, 1), mtl.Size.of(256, 1, 1)); + bias(e, k.get("mi_bias"), sum, b, y, M, N); } - /// o = softmax(q k^T / sqrt(128)) v, every row to every row: q, k, v [heads, L, 128], o rows [L, heads * 128]. - pub fn attention(k: *const Kernels, e: mtl.ComputeEncoder, q: Ref, kb: Ref, vb: Ref, o: Ref, L: usize, heads: usize) void { - const nq = (L + 63) / 64; - const nk = (L + 31) / 32; - e.setPipeline(k.p[attention_first + @as(usize, @intFromBool(L % 64 == 0)) * 2 + @intFromBool(L % 32 == 0)]); - for ([_]Ref{ q, kb, vb }, 0..) |r, i| e.setBuffer(r.buf, r.off, i); - const d = 128; - const vals = [_]usize{ 1, heads, L, L, 1, nq, nk, L / 64, L / 32, L % 64, L % 32, 0, 0, L * d, d, 0, L * d, d, 0, L * d, d, 0, d, heads * d }; - var p: [32]i32 = @splat(0); - for (vals, 0..) |v, i| p[i] = @intCast(v); - e.setBytes(std.mem.asBytes(&p), 3); - var f: [16]f32 = @splat(0); - f[0] = 1.0 / @sqrt(@as(f32, d)); - e.setBytes(std.mem.asBytes(&f), 4); - e.setBuffer(o.buf, o.off, 5); - e.dispatchThreads(mtl.Size.of(nq * 32, heads * 4, 1), mtl.Size.of(32, 4, 1)); + /// o = softmax(q k^T / sqrt(128)) v, every row to every row: q, k, v [count, L, 128], o rows [L, count * 128]. + pub fn attention(k: *const Kernels, e: mtl.ComputeEncoder, q: Ref, kb: Ref, vb: Ref, o: Ref, L: usize, count: usize) void { + attend(e, k.p[attention_first + attentionForm(L)], q, kb, vb, o, L, count); } /// One kernel's buffers in order, then its parameters at the next index. @@ -123,6 +108,61 @@ pub const Kernels = struct { } }; +/// The attention form for L rows: 0 for L % 32 > 0, 1 for 32-aligned, 2 for 64-aligned (its query blocks whole too). +pub fn attentionForm(L: usize) usize { + return if (L % 64 == 0) 2 else @intFromBool(L % 32 == 0); +} + +/// mi_attention_* `p` (the form for L) over `count` heads of L rows. +pub fn attend(e: mtl.ComputeEncoder, p: mtl.Pipeline, q: Ref, kb: Ref, vb: Ref, o: Ref, L: usize, count: usize) void { + const nq = (L + 63) / 64; + const nk = (L + 31) / 32; + e.setPipeline(p); + for ([_]Ref{ q, kb, vb }, 0..) |r, i| e.setBuffer(r.buf, r.off, i); + const d = 128; + const vals = [_]usize{ 1, count, L, L, 1, nq, nk, L / 64, L / 32, L % 64, L % 32, 0, 0, L * d, d, 0, L * d, d, 0, L * d, d, 0, d, count * d }; + var params: [32]i32 = @splat(0); + for (vals, 0..) |v, i| params[i] = @intCast(v); + e.setBytes(std.mem.asBytes(¶ms), 3); + var f: [16]f32 = @splat(0); + f[0] = 1.0 / @sqrt(@as(f32, d)); + e.setBytes(std.mem.asBytes(&f), 4); + e.setBuffer(o.buf, o.off, 5); + e.dispatchThreads(mtl.Size.of(nq * 32, count * 4, 1), mtl.Size.of(32, 4, 1)); +} + +/// y = mi_norm of x's `rows` rows of `width` with weight w, then by s and b when `modulated`; a threadgroup a row. +pub fn norm(e: mtl.ComputeEncoder, p: mtl.Pipeline, x: Ref, w: Ref, s: Ref, b: Ref, y: Ref, rows: usize, width: usize, eps: f32, modulated: bool) void { + const params = extern struct { width: u32, eps: f32, modulated: u32 }{ .width = @intCast(width), .eps = eps, .modulated = @intFromBool(modulated) }; + Kernels.bind(e, p, &.{ x, w, s, b, y }, params); + e.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(256, 1, 1)); +} + +/// q, k, v rows `stride` apart to normed, rotated heads [count, rows, 128] in q, k and v (copied). +pub fn heads(e: mtl.ComputeEncoder, p: mtl.Pipeline, qr: Ref, kr: Ref, vr: Ref, nq: Ref, nk: Ref, cos: Ref, sin: Ref, q: Ref, k: Ref, v: Ref, rows: usize, count: usize, stride: usize, eps: f32) void { + const params = extern struct { rows: u32, heads: u32, stride: u32, eps: f32 }{ .rows = @intCast(rows), .heads = @intCast(count), .stride = @intCast(stride), .eps = eps }; + Kernels.bind(e, p, &.{ qr, kr, vr, nq, nk, cos, sin, q, k, v }, params); + e.dispatchGroups(mtl.Size.of(rows, (count + 7) / 8, 1), mtl.Size.of(256, 1, 1)); +} + +/// o [rows, width] = situ(h, g) of rows `stride` apart. +pub fn situ(e: mtl.ComputeEncoder, p: mtl.Pipeline, h: Ref, g: Ref, o: Ref, rows: usize, width: usize, stride: usize) void { + Kernels.bind(e, p, &.{ h, g, o }, [2]u32{ @intCast(width), @intCast(stride) }); + e.dispatchThreads(mtl.Size.of(width, rows, 1), mtl.Size.of(256, 1, 1)); +} + +/// x [rows, width] += y gated by the vector g, in place. +pub fn gateAdd(e: mtl.ComputeEncoder, p: mtl.Pipeline, x: Ref, y: Ref, g: Ref, rows: usize, width: usize) void { + Kernels.bind(e, p, &.{ x, y, g }, @as(u32, @intCast(width))); + e.dispatchThreads(mtl.Size.of(width, rows, 1), mtl.Size.of(256, 1, 1)); +} + +/// y [rows, width] = bf16(c + b): fp32 sums and a bias vector, rounded once. +pub fn bias(e: mtl.ComputeEncoder, p: mtl.Pipeline, c: Ref, b: Ref, y: Ref, rows: usize, width: usize) void { + Kernels.bind(e, p, &.{ c, b, y }, @as(u32, @intCast(width))); + e.dispatchThreads(mtl.Size.of(width, rows, 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/model.zig b/zig/src/families/moss_image/model.zig index e553b0a968..2655c66508 100644 --- a/zig/src/families/moss_image/model.zig +++ b/zig/src/families/moss_image/model.zig @@ -112,6 +112,8 @@ pub const Model = struct { const image = try Scratch.mul(g.rows, g.columns); const plan = try Scratch.of(max_text, image); const scratch_bytes = try plan.bytes(); + // qualified on the M5's tensor units alone: the projections and attention take another path before them + if (!device.tensorUnits()) return error.TensorUnitsRequired; const t = try gpa.create(Model); errdefer gpa.destroy(t); t.* = .{ .gpa = gpa, .device = device, .queue = undefined, .k = undefined, .w = undefined, .max_text = max_text, .image = image, .rows = g.rows, .columns = g.columns, .scratch = undefined, .x = undefined, .xn = undefined, .qr = undefined, .kr = undefined, .vr = undefined, .att = undefined, .y = undefined, .ug = undefined, .s = undefined, .q = undefined, .k_ = undefined, .v = undefined, .sums = undefined, .lat = undefined, .vel = undefined, .ctx = undefined, .mods = undefined }; @@ -225,28 +227,22 @@ pub const Model = struct { fn norm(t: *const Model, e: mtl.ComputeEncoder, rows: usize, w: Ref, m: ?Modulated) void { const s = if (m) |v| t.vec(v.scale) else t.vec(.ones); const b = if (m) |v| t.vec(v.shift) else t.vec(.ones); - const p = extern struct { width: u32, eps: f32, modulated: u32 }{ .width = H, .eps = cfg.eps, .modulated = @intFromBool(m != null) }; - ks.Kernels.bind(e, t.k.get("mi_norm"), &.{ t.x, w, s, b, t.xn }, p); - e.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(256, 1, 1)); + ks.norm(e, t.k.get("mi_norm"), t.x, w, s, b, t.xn, rows, H, cfg.eps, m != null); } /// x += y gated by shared vector `gate` (ones: plain). fn gateAdd(t: *const Model, e: mtl.ComputeEncoder, rows: usize, gate: host.Vector) void { - ks.Kernels.bind(e, t.k.get("mi_gate_add"), &.{ t.x, t.y, t.vec(gate) }, @as(u32, H)); - e.dispatchThreads(mtl.Size.of(H, rows, 1), mtl.Size.of(256, 1, 1)); + ks.gateAdd(e, t.k.get("mi_gate_add"), t.x, t.y, t.vec(gate), rows, H); } /// q k v rows of `stride` to normed, rotated heads in q, k_, v. fn heads(t: *const Model, e: mtl.ComputeEncoder, qr: Ref, kr: Ref, vr: Ref, stride: usize, nq: Ref, nk: Ref, cos: Ref, sin: Ref, rows: usize) void { - const p = extern struct { rows: u32, heads: u32, stride: u32, eps: f32 }{ .rows = @intCast(rows), .heads = cfg.heads, .stride = @intCast(stride), .eps = cfg.eps }; - ks.Kernels.bind(e, t.k.get("mi_qk"), &.{ qr, kr, vr, nq, nk, cos, sin, t.q, t.k_, t.v }, p); - e.dispatchGroups(mtl.Size.of(rows, cfg.heads / 8, 1), mtl.Size.of(256, 1, 1)); + ks.heads(e, t.k.get("mi_qk"), qr, kr, vr, nq, nk, cos, sin, t.q, t.k_, t.v, rows, cfg.heads, stride, cfg.eps); } /// s = situ(hidden, gate) of rows `stride` apart. fn situ(t: *const Model, e: mtl.ComputeEncoder, hidden: Ref, gate: Ref, rows: usize, stride: usize) void { - ks.Kernels.bind(e, t.k.get("mi_situ"), &.{ hidden, gate, t.s }, [2]u32{ F, @intCast(stride) }); - e.dispatchThreads(mtl.Size.of(F, rows, 1), mtl.Size.of(256, 1, 1)); + ks.situ(e, t.k.get("mi_situ"), hidden, gate, t.s, rows, F, stride); } /// One joint block over the image and prompt rows (`rows` in all), under the shared modulation. diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index f2a067de87..8ee1cf3b14 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -325,6 +325,8 @@ pub fn main(init: std.process.Init) !u8 { defer native_ple.destroy(); const native_gdn = try @import("metal_qwen4_gdn.zig").Suite.create(a, device, queue, &kernels, o.only); defer native_gdn.destroy(); + const moss_image = try @import("metal_moss_image.zig").Suite.create(a, device, queue, &kernels, o.only); + defer moss_image.destroy(); var t: g.Tally = .{}; const t0 = std.Io.Clock.awake.now(init.io); if (o.gates) { diff --git a/zig/tests/contract/metal_moss_image.zig b/zig/tests/contract/metal_moss_image.zig new file mode 100644 index 0000000000..30b158ac65 --- /dev/null +++ b/zig/tests/contract/metal_moss_image.zig @@ -0,0 +1,397 @@ +//! MOSS-Image's Metal kernels under the contract, each at its torch bf16 points; controls drop or add one such point. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const contract = @import("contract"); +const sources = @import("kernel_sources"); +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 c = mi.c; +const g = mi.g; +const ks = mi.ks; +const D = mi.D; +const kref = mi.kref; +const Ends = mi.Ends; +const Norm = norms.Norm; +const Qk = norms.Qk; +const Situ = steps.Situ; +const GateAdd = steps.GateAdd; +const Bias = steps.Bias; + +// --- mi_attention_*: softmax(q k^T / sqrt(128)) v over every row, against attention_ref's bf16-P bound ------------ + +/// Each form's name, served lengths and edges: L % 32 > 0, L a multiple of 32 alone, and of 64. +const forms = [3]struct { name: []const u8, served: []const c.Shape, tile: c.Tile }{ + .{ + .name = "mi_attention_00", + .served = &.{ .{ .rows = 2, .n = D, .k = 63 }, .{ .rows = 2, .n = D, .k = 4188 }, .{ .rows = 2, .n = D, .k = 4130 }, .{ .rows = 2, .n = D, .k = 34 }, .{ .rows = 2, .n = D, .k = 10 } }, + .tile = .{ .size = 32 }, + }, + .{ + .name = "mi_attention_01", + .served = &.{ .{ .rows = 2, .n = D, .k = 32 }, .{ .rows = 2, .n = D, .k = 4128 } }, + .tile = .{ .size = 64, .step = 32, .min = 32 }, + }, + .{ + .name = "mi_attention_11", + .served = &.{ .{ .rows = 2, .n = D, .k = 64 }, .{ .rows = 2, .n = D, .k = 4160 } }, + .tile = .{ .size = 128, .step = 64, .min = 64 }, + }, +}; + +/// The scores' 128-deep sums, then the keys' sums. +fn attentionDepth(s: c.Shape) usize { + return D + s.k; +} + +pub const Attention = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + form: usize, + 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: heads (each the launch's own y, so two stand for the 24); n: the head width; k: the rows L. + pub fn spec(name: []const u8, form: usize) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = attentionDepth, + .ops = 8, + .out = .bf16, + .tiles = .{ .{}, .{}, forms[form].tile }, + .served = forms[form].served, + }; + } + + pub fn deinit(t: *Attention) void { + t.free(); + t.rig.deinit(); + } + fn free(t: *Attention) void { + t.a.free(t.q); + t.a.free(t.k); + t.a.free(t.v); + t.q = &.{}; + t.k = &.{}; + t.v = &.{}; + } + pub fn kernel(t: *Attention) g.Kernel { + return .{ .spec = spec(t.name, t.form), .ctx = t, .vt = &vt }; + } + const vt = g.hooks(Attention); + + pub fn prepare(t: *Attention, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n != D or shape.k == 0 or ks.attentionForm(shape.k) != t.form) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + t.stress = fill == .stress; + const L = shape.k; + const n = shape.rows * L * D; + t.q = try t.a.alloc(u16, n); + t.k = try t.a.alloc(u16, n); + t.v = try t.a.alloc(u16, n); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (t.stress) { + // equal scores (q = k = 0), each column's values integers whose partial sums reach ~2^22 and cancel to 4 + @memset(t.q, 0); + @memset(t.k, 0); + const half = L / 2; + for (0..shape.rows) |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) { + // head 0: values NaN, +Inf, -Inf, and both infinities in one column; head 1: a NaN query, an infinite key + for ([_][2]usize{ .{ 0, 0 }, .{ 1, 1 }, .{ 2, 2 }, .{ 3, 3 }, .{ 4, 3 } }, [_]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 (shape.rows > 1) { + t.q[(L + 0) * D + 5] = mi.nan16; + if (L > 7) t.k[(L + 7) * D + 3] = mi.inf16; + } + } + 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(n * 2); + return true; + } + + pub fn run(t: *Attention) !void { + const p = t.rig.begin(); + ks.attend(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); + try t.rig.finish(p); + } + + pub fn count(t: *Attention) usize { + return t.shape.rows * t.shape.k * D; + } + + /// attention_ref.head's formula and bound, one pass over the keys a query (its scores kept), O in [L, heads, 128]. + pub fn reference(t: *Attention, want: []f64, mag: []f64) !void { + const L = t.shape.k; + const heads = t.shape.rows; + const scale: f64 = @as(f32, 1.0 / @sqrt(@as(f32, D))); + const kf = try t.a.alloc(f64, L * D); + defer t.a.free(kf); + const vf = try t.a.alloc(f64, L * D); + defer t.a.free(vf); + const s = try t.a.alloc(f64, L); + defer t.a.free(s); + var qf: [D]f64 = undefined; + var out: [D]f64 = undefined; + var tol: [D]f64 = undefined; + for (0..heads) |h| { + for (kf, vf, t.k[h * L * D ..][0 .. L * D], t.v[h * L * D ..][0 .. L * D]) |*kd, *vd, kb, vb| { + kd.* = c.fromBf16(kb); + vd.* = c.fromBf16(vb); + } + for (0..L) |i| { + var zero = true; + for (&qf, t.q[(h * L + i) * D ..][0..D]) |*x, b| { + x.* = c.fromBf16(b); + zero = zero and x.* == 0; + } + var m: f64 = -std.math.inf(f64); + var nan = false; + for (s, 0..) |*sj, j| { + var dot: f64 = 0; + if (!zero) { + const Vec = @Vector(8, f64); + var acc: Vec = @splat(0); + var e: usize = 0; + while (e < D) : (e += 8) acc += @as(Vec, qf[e..][0..8].*) * @as(Vec, kf[j * D + e ..][0..8].*); + dot = @reduce(.Add, acc); + } + sj.* = dot * scale; + nan = nan or std.math.isNan(sj.*); + m = @max(m, sj.*); + } + if (nan) m = std.math.nan(f64); + var l: f64 = 0; + @memset(&out, 0); + @memset(&tol, 0); + for (s, 0..) |sj, j| { + const p = @exp(sj - m); + l += p; + for (&out, &tol, vf[j * D ..][0..D]) |*o, *tl, vv| { + o.* += p * vv; + tl.* += p * @abs(vv); + } + } + const at = (i * heads + h) * D; + for (out, tol, 0..) |o, tl, e| { + want[at + e] = o / l; + mag[at + e] = tl / l * 0x1p-7 + @abs(o / l) * 0x1p-8 + 1e-6; + } + } + } + } + + pub fn bounds(t: *Attention, want: []const f64, tolerance: []const f64, allow: []f64) !void { + for (want, tolerance, allow) |w, tl, *l| l.* = if (t.stress) c.boundExact(spec(t.name, t.form), w) else if (std.math.isFinite(w)) tl else 0; + } + + pub fn simulate(t: *Attention, acc: c.Accum, order: g.Order, out: []f64) !void { + if (!t.stress) return error.AttentionStressNeedsEqualScores; + const L = t.shape.k; + const terms = try t.a.alloc(f64, L); + defer t.a.free(terms); + const ones = try t.a.alloc(f64, L); + defer t.a.free(ones); + @memset(ones, 1); + const l = contract.rows.sumIn(acc, order, ones); + for (0..t.shape.rows) |h| for (0..D) |e| { + for (terms, 0..) |*v, j| v.* = c.fromBf16(t.v[(h * L + j) * D + e]); + const o = contract.rows.sumIn(acc, order, terms) / l; + for (0..L) |i| out[(i * t.shape.rows + h) * D + e] = o; + }; + } +}; + +// --- the suite: the served source's pipelines, and each control's own copy of the source ------------------------ + +/// A control: the kernel it breaks and the one text change that drops (or adds) a rounding point. +const Control = struct { kernel: []const u8, what: []const u8, from: []const u8, to: []const u8 }; + +const controls = [_]?Control{ + .{ .kernel = "mi_norm", .what = "drops bf16(x inv)", .from = "mi_bf(mi_bf(float(xr[i]) * inv) * float(w[i]))", .to = "mi_bf((float(xr[i]) * inv) * float(w[i]))" }, + .{ .kernel = "mi_norm modulated", .what = "drops bf16(y s)", .from = "v = mi_bf(v * float(s[i])) + float(b[i]);", .to = "v = (v * float(s[i])) + float(b[i]);" }, + .{ .kernel = "mi_qk", .what = "drops bf16(q inv)", .from = "q[i] = mi_bf(mi_bf(q[i] * qn)", .to = "q[i] = mi_bf((q[i] * qn)" }, + .{ .kernel = "mi_qk", .what = "drops bf16(q cos)", .from = "Q[to + i] = bfloat(mi_bf(q[i] * c) + mi_bf(qp * s));", .to = "Q[to + i] = bfloat((q[i] * c) + mi_bf(qp * s));" }, + .{ .kernel = "mi_situ", .what = "drops bf16(sigmoid(g))", .from = "const float sig = mi_bf(1.0f / (1.0f + metal::precise::exp(-gv)));", .to = "const float sig = (1.0f / (1.0f + metal::precise::exp(-gv)));" }, + .{ .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]));" }, +}; + +/// The family's source as the engine compiles it, with `ctl`'s change made where it occurs, exactly once. +fn compile(a: std.mem.Allocator, device: mtl.Device, ctl: ?Control) !mtl.Library { + const served = try tf.frags.source(device, a, sources.moss_image); + defer a.free(served); + const text = if (ctl) |x| blk: { + if (std.mem.count(u8, served, x.from) != 1) return error.ControlTextNotFoundOnce; + break :blk try std.mem.replaceOwned(u8, a, served, x.from, x.to); + } else try a.dupe(u8, served); + defer a.free(text); + return mtl.Library.fromSource(device, text, mtl.CompileOptions.mlx()); +} + +/// A source compiled the first time a kernel registered on it asks for a pipeline. +const Lazy = struct { + a: std.mem.Allocator, + device: mtl.Device, + ctl: ?Control, + lib: ?mtl.Library = null, + + fn get(l: *Lazy) !mtl.Library { + if (l.lib == null) l.lib = try compile(l.a, l.device, l.ctl); + return l.lib.?; + } + fn deinit(l: *Lazy) void { + if (l.lib) |x| x.deinit(); + } +}; + +pub const Suite = struct { + a: std.mem.Allocator, + pipes: std.ArrayList(mtl.Pipeline) = .empty, + names: std.ArrayList([]u8) = .empty, + norms: std.ArrayList(*Norm) = .empty, + qks: std.ArrayList(*Qk) = .empty, + situs: std.ArrayList(*Situ) = .empty, + adds: std.ArrayList(*GateAdd) = .empty, + biases: std.ArrayList(*Bias) = .empty, + attns: std.ArrayList(*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); + errdefer p.deinit(); + try s.pipes.append(s.a, p); + return p; + } + + fn label(s: *Suite, comptime fmt: []const u8, args: anytype) ![]const u8 { + const name = try std.fmt.allocPrint(s.a, fmt, args); + errdefer s.a.free(name); + try s.names.append(s.a, name); + return name; + } + + fn make(s: *Suite, comptime T: type, list: *std.ArrayList(*T), value: T) !*T { + const t = try s.a.create(T); + errdefer s.a.destroy(t); + t.* = value; + try list.append(s.a, t); + return t; + } + + /// Each step (or, with `ctl`, the one it breaks) on `lib`'s pipelines. + fn add(s: *Suite, device: mtl.Device, queue: mtl.Queue, lib: *Lazy, ctl: ?Control, out: *std.ArrayList(g.Kernel), only: ?[]const u8) !void { + const suffix = if (ctl) |x| try s.label(" control: {s}", .{x.what}) else ""; + const gate: ?g.Gate = if (ctl != null) .shapes else null; + const rig: rig_.Rig = .{ .a = s.a, .device = device, .queue = queue }; + const ends: Ends = .{ .a = s.a }; + const want = struct { + fn is(x: ?Control, kernel: []const u8) bool { + return if (x) |y| std.mem.eql(u8, y.kernel, kernel) else true; + } + fn keep(o: ?[]const u8, name: []const u8) bool { + 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"; + 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 }); + try out.append(s.a, t.kernel(gate)); + } + for ([_]bool{ false, true }) |packed_rows| { + const name = try s.label("metal moss image mi_qk {s}{s}", .{ if (packed_rows) "packed qkv rows" else "q k v rows", suffix }); + if (!want.is(ctl, "mi_qk") or !want.keep(only, name) or (ctl != null and packed_rows)) continue; + const t = try s.make(Qk, &s.qks, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "mi_qk"), .packed_rows = packed_rows, .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + for ([_]bool{ true, false }) |halves| { + const name = try s.label("metal moss image mi_situ {s}{s}", .{ if (halves) "hidden | gate rows" else "two inputs", suffix }); + if (!want.is(ctl, "mi_situ") or !want.keep(only, name) or (ctl != null and !halves)) continue; + const t = try s.make(Situ, &s.situs, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "mi_situ"), .halves = halves, .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + { + const name = try s.label("metal moss image mi_gate_add{s}", .{suffix}); + if (want.is(ctl, "mi_gate_add") and want.keep(only, name)) { + const t = try s.make(GateAdd, &s.adds, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "mi_gate_add"), .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + } + { + const name = try s.label("metal moss image mi_bias{s}", .{suffix}); + if (want.is(ctl, "mi_bias") and want.keep(only, name)) { + const t = try s.make(Bias, &s.biases, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "mi_bias"), .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + } + if (ctl != null) return; + for (forms, 0..) |f, form| { + const name = try s.label("metal moss image {s}", .{f.name}); + if (!want.keep(only, name)) continue; + const t = try s.make(Attention, &s.attns, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, f.name), .form = form, .name = name }); + try out.append(s.a, t.kernel()); + } + } + + pub fn create(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, out: *std.ArrayList(g.Kernel), only: ?[]const u8) !*Suite { + const s = try a.create(Suite); + s.* = .{ .a = a }; + errdefer s.destroy(); + if (!device.tensorUnits()) { + std.debug.print("UNSUPPORTED metal moss image: the family runs on tensor units (M5) alone; its kernels unchecked here\n", .{}); + return s; + } + for ([_]?Control{null} ++ controls) |x| { + var lib: Lazy = .{ .a = a, .device = device, .ctl = x }; + defer lib.deinit(); + try s.add(device, queue, &lib, x, out, only); + } + return s; + } + + 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| { + for (list.items) |t| { + t.deinit(); + a.destroy(t); + } + list.deinit(a); + } + for (s.pipes.items) |p| p.deinit(); + s.pipes.deinit(a); + for (s.names.items) |n| a.free(n); + s.names.deinit(a); + a.destroy(s); + } +}; diff --git a/zig/tests/contract/metal_moss_image_common.zig b/zig/tests/contract/metal_moss_image_common.zig new file mode 100644 index 0000000000..dd5ce450b0 --- /dev/null +++ b/zig/tests/contract/metal_moss_image_common.zig @@ -0,0 +1,144 @@ +//! MOSS-Image's kernel tests' shared parts: float64 values at bf16 points, their ranges, the launches' buffers. +const std = @import("std"); +const mtl = @import("metal"); +pub const tf = @import("tensorfold"); +const contract = @import("contract"); +const sources = @import("kernel_sources"); +const rig_ = @import("rig.zig"); + +pub const c = contract.spec; +pub const g = contract.gates; +pub const ks = tf.moss_image.kernels; +pub const cfg = tf.moss_image.config; +pub const eps: f64 = cfg.eps; +pub const D = cfg.head_dim; + +pub fn noSums(_: c.Shape) usize { + return 0; +} + +/// A value as fp32 then bf16 round it. +pub fn bf(v: f64) f64 { + return c.fromBf16(c.toBf16(@floatCast(v))); +} + +pub fn kref(r: rig_.Ref) ks.Ref { + return .{ .buf = r.buf, .off = r.off }; +} + +pub fn kat(r: rig_.Ref, bytes: usize) ks.Ref { + return .{ .buf = r.buf, .off = r.off + bytes }; +} + +/// Every value the kernel may hold at a step: its fp32 steps widen both ends, its bf16 points round them. +pub const Range = struct { + lo: f64, + hi: f64, + pub fn of(v: f64) Range { + return .{ .lo = v, .hi = v }; + } + /// An fp32 step good to `ulps` units in the last place (a correctly rounded one: 0.5). + pub fn fp32(a: Range, ulps: f64) Range { + const u = ulps * 0x1p-23; + return .{ .lo = a.lo - u * @abs(a.lo) - 0x1p-149, .hi = a.hi + u * @abs(a.hi) + 0x1p-149 }; + } + pub fn bf16(a: Range) Range { + return .{ .lo = bf(a.lo), .hi = bf(a.hi) }; + } + pub fn add(a: Range, b: Range) Range { + return .{ .lo = a.lo + b.lo, .hi = a.hi + b.hi }; + } + pub fn mul(a: Range, b: Range) Range { + const p = [4]f64{ a.lo * b.lo, a.lo * b.hi, a.hi * b.lo, a.hi * b.hi }; + return .{ .lo = @min(@min(p[0], p[1]), @min(p[2], p[3])), .hi = @max(@max(p[0], p[1]), @max(p[2], p[3])) }; + } + /// An increasing function at both ends. + pub fn up(a: Range, comptime f: fn (f64) f64) Range { + return .{ .lo = f(a.lo), .hi = f(a.hi) }; + } + /// 1 / sqrt(a), a positive. + pub fn rsqrt(a: Range) Range { + return .{ .lo = 1 / @sqrt(a.hi), .hi = 1 / @sqrt(a.lo) }; + } + /// 1 / a, a positive. + pub fn recip(a: Range) Range { + return .{ .lo = 1 / a.hi, .hi = 1 / a.lo }; + } +}; + +pub fn tanh(x: f64) f64 { + return std.math.tanh(x); +} + +// Metal's error bounds with the engine's options (fast math on, the kernels' own `precise::` calls). +pub const div_ulps = 2.5; +pub const rsqrt_ulps = 2; +pub const exp_ulps = 4; +pub const tanh_ulps = 5; + +/// Each output's float64 value at the kernel's rounding points, and its range, kept for the bounds and intervals. +pub const Ends = struct { + a: std.mem.Allocator, + lo: []f64 = &.{}, + hi: []f64 = &.{}, + + pub fn deinit(e: *Ends) void { + e.a.free(e.lo); + e.a.free(e.hi); + e.lo = &.{}; + e.hi = &.{}; + } + + pub fn size(e: *Ends, n: usize) !void { + e.deinit(); + e.lo = try e.a.alloc(f64, n); + e.hi = try e.a.alloc(f64, n); + } + + pub fn put(e: *Ends, want: []f64, mag: []f64, i: usize, w: f64, r: Range) void { + want[i] = w; + mag[i] = if (std.math.isFinite(w)) @abs(w) else 0; + e.lo[i] = r.lo; + e.hi[i] = r.hi; + } + + /// The allowance about the float64 value: its range's farther end (a non-finite value is checked by class). + pub fn bounds(e: *const Ends, want: []const f64, allow: []f64) void { + for (allow, want, e.lo, e.hi) |*l, w, lo, hi| l.* = if (std.math.isFinite(w)) @max(@abs(w - lo), @abs(hi - w)) else 0; + } + + pub fn intervals(e: *const Ends, low: []f64, high: []f64) void { + @memcpy(low, e.lo); + @memcpy(high, e.hi); + } +}; + +pub fn randomBf16(rng: std.Random, out: []u16, lo: f32, hi: f32) void { + for (out) |*v| v.* = c.toBf16(lo + rng.float(f32) * (hi - lo)); +} + +pub fn integers(rng: std.Random, out: []u16, bound: u32) void { + for (out) |*v| v.* = c.toBf16(@as(f32, @floatFromInt(rng.uintLessThan(u32, 2 * bound + 1))) - @as(f32, @floatFromInt(bound))); +} + +pub const nan16 = 0x7fc0; +pub const inf16 = 0x7f80; +pub const ninf16 = 0xff80; + +/// A row's sum of squares in float64 and its range from fp32 sums `depth` deep (exact for gate 4's integers). +pub fn squares(x: []const u16, depth: usize, exact: bool) struct { ss: f64, r: Range } { + var ss: f64 = 0; + for (x) |v| { + const f: f64 = c.fromBf16(v); + ss += f * f; + } + if (exact) return .{ .ss = ss, .r = Range.of(ss) }; + const e = c.gamma(@floatFromInt(depth), c.unit(.f32)) * ss; + return .{ .ss = ss, .r = .{ .lo = ss - e, .hi = ss + e } }; +} + +/// rsqrt(ss / n + eps) in float64 and its range through the kernel's division, add and precise rsqrt. +pub fn invNorm(ss: f64, r: Range, n: f64) struct { inv: f64, r: Range } { + const mean = (Range{ .lo = r.lo / n, .hi = r.hi / n }).fp32(div_ulps); + return .{ .inv = 1 / @sqrt(ss / n + eps), .r = mean.add(Range.of(eps)).fp32(0.5).rsqrt().fp32(rsqrt_ulps) }; +} diff --git a/zig/tests/contract/metal_moss_image_norms.zig b/zig/tests/contract/metal_moss_image_norms.zig new file mode 100644 index 0000000000..6f64b519bc --- /dev/null +++ b/zig/tests/contract/metal_moss_image_norms.zig @@ -0,0 +1,363 @@ +//! mi_norm and mi_qk under the contract: RMS norms at the publisher's bf16 points, their fp32 sums proven in gate 4. +const std = @import("std"); +const mtl = @import("metal"); +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 cfg = mi.cfg; +const D = mi.D; +const Range = mi.Range; +const Ends = mi.Ends; +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 -- + +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"; + +/// A thread's share of the row in turn, then the simdgroup tree and the tree of the threadgroup's 8 simdgroups. +fn normDepth(s: c.Shape) usize { + return (s.k + 255) / 256 + 10; +} + +pub const Norm = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + modulated: bool, + name: []const u8, + shape: c.Shape = undefined, + fill: g.Fill = .random, + x: []u16 = &.{}, + /// w, then s and b + v: []u16 = &.{}, + ins: [4]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + pub fn spec(name: []const u8, modulated: bool) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = normDepth, + .ops = 2 * (mi.div_ulps + mi.rsqrt_ulps) + 1, + .mids = if (modulated) 3 else 1, + .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 } }, + .sums = .{ .bf16 = 0, .why = norm_why }, + }; + } + + pub fn deinit(t: *Norm) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *Norm) void { + t.a.free(t.x); + t.a.free(t.v); + t.x = &.{}; + 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 }; + } + const vt = g.hooks(Norm); + + pub fn prepare(t: *Norm, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.k == 0) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + t.fill = fill; + const k = shape.k; + t.x = try t.a.alloc(u16, shape.rows * k); + t.v = try t.a.alloc(u16, 3 * k); + 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); + // 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); + if (fill == .nonfinite) { + const at = [_][2]usize{ .{ 0, k / 3 }, .{ 1, k / 2 }, .{ 2, k - 1 }, .{ 3, 0 }, .{ 3, k - 1 } }; + for (at, [_]u16{ mi.nan16, mi.inf16, mi.ninf16, mi.inf16, mi.ninf16 }) |p, v| if (p[0] < shape.rows) { + t.x[p[0] * k + p[1]] = v; + }; + if (k > 7) { + t.v[5] = mi.nan16; + t.v[k + 6] = mi.inf16; + t.v[2 * k + 7] = mi.ninf16; + } + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.x)); + for (0..3) |i| t.ins[1 + i] = try t.rig.input(std.mem.sliceAsBytes(t.v[i * k ..][0..k])); + t.out = try t.rig.output(shape.rows * k * 2); + return true; + } + + 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); + try t.rig.finish(p); + } + + pub fn count(t: *Norm) usize { + return t.shape.rows * t.shape.k; + } + + fn row(t: *Norm, ends: *Ends, r: usize, ss: f64, ssr: Range, want: []f64, mag: []f64) void { + const k = t.shape.k; + const n = mi.invNorm(ss, ssr, @floatFromInt(k)); + for (0..k) |j| { + const x: f64 = c.fromBf16(t.x[r * k + j]); + const w: f64 = c.fromBf16(t.v[j]); + 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) { + 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); + yr = yr.mul(Range.of(s)).fp32(0.5).bf16().add(Range.of(b)).fp32(0.5).bf16(); + } + ends.put(want, mag, r * k + j, y, yr); + } + } + + pub fn reference(t: *Norm, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + const k = t.shape.k; + for (0..t.shape.rows) |r| { + const s = mi.squares(t.x[r * k ..][0..k], normDepth(t.shape), t.fill == .stress); + t.row(&t.ends, r, s.ss, s.r, want, mag); + } + } + pub fn bounds(t: *Norm, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *Norm, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } + + pub fn simulate(t: *Norm, acc: c.Accum, order: g.Order, out: []f64) !void { + const k = t.shape.k; + const terms = try t.a.alloc(f64, k); + defer t.a.free(terms); + 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()); + for (0..t.shape.rows) |r| { + for (terms, t.x[r * k ..][0..k]) |*v, x| { + const f: f64 = c.fromBf16(x); + v.* = f * f; + } + const ss = contract.rows.sumIn(acc, order, terms); + t.row(&scratch, r, ss, Range.of(ss), out, mag); + } + } +}; + +// --- mi_qk: q and k heads normed, (c, c + 64) rotated by cos and sin; v copied -------------------------------------- + +/// Four channels a lane in turn, then the simdgroup tree. +fn qkDepth(_: c.Shape) usize { + return 4 + 5; +} + +pub const Qk = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + /// The refiner's form: q, k and v side by side in one row of 3 x 3072; else three buffers of 3072-wide rows. + packed_rows: bool, + name: []const u8, + shape: c.Shape = undefined, + fill: g.Fill = .random, + /// q, k, v as [rows, heads * 128] each + qkv: [3][]u16 = .{ &.{}, &.{}, &.{} }, + /// the q and k head norms' weights + nw: []u16 = &.{}, + /// cos then sin [rows, 64], bf16 values + rot: []f32 = &.{}, + ins: [6]rig_.Ref = undefined, + outs: [3]rig_.Ref = undefined, + ends: Ends, + + /// Rows: positions; n: heads (24); k: the head width (128). + pub fn spec(name: []const u8, packed_rows: bool) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = qkDepth, + .ops = 2 * (mi.div_ulps + mi.rsqrt_ulps) + 1, + .mids = 3, + .out = .bf16, + .served = if (packed_rows) + &.{ .{ .rows = 63, .n = 24, .k = D }, .{ .rows = 34, .n = 24, .k = D }, .{ .rows = 10, .n = 24, .k = D } } + else + &.{ .{ .rows = 4188, .n = 24, .k = D }, .{ .rows = 4130, .n = 24, .k = D }, .{ .rows = 4160, .n = 24, .k = D } }, + .sums = .{ .bf16 = 0, .why = norm_why }, + }; + } + + pub fn deinit(t: *Qk) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *Qk) void { + for (&t.qkv) |*x| { + t.a.free(x.*); + x.* = &.{}; + } + t.a.free(t.nw); + t.a.free(t.rot); + t.nw = &.{}; + t.rot = &.{}; + } + pub fn kernel(t: *Qk, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name, t.packed_rows), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Qk); + + fn width(t: *const Qk) usize { + return t.shape.n * D; + } + + pub fn prepare(t: *Qk, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.k != D) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + t.fill = fill; + const w = t.width(); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (&t.qkv, 0..) |*x, i| { + x.* = try t.a.alloc(u16, shape.rows * w); + if (fill == .stress and i < 2) mi.integers(rng, x.*, 128) else mi.randomBf16(rng, x.*, -4, 4); + } + t.nw = try t.a.alloc(u16, 2 * D); + mi.randomBf16(rng, t.nw, 0.25, 2); + t.rot = try t.a.alloc(f32, 2 * shape.rows * D / 2); + const half = shape.rows * D / 2; + 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 = [_][4]usize{ .{ 0, 0, 0, 5 }, .{ 1, 1, 2, 70 }, .{ 0, 2, 3, 0 }, .{ 2, 3, 1, 9 }, .{ 1, 4, 0, 127 } }; + for (at, [_]u16{ mi.nan16, mi.inf16, mi.ninf16, mi.nan16, mi.ninf16 }) |p, v| if (p[1] < shape.rows and p[2] < shape.n) { + t.qkv[p[0]][p[1] * w + p[2] * D + p[3]] = v; + }; + } + if (t.packed_rows) { + const both = try t.a.alloc(u16, shape.rows * 3 * w); + defer t.a.free(both); + for (0..shape.rows) |r| for (t.qkv, 0..) |x, i| @memcpy(both[(r * 3 + i) * w ..][0..w], x[r * w ..][0..w]); + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(both)); + } else for (t.qkv, 0..) |x, i| { + t.ins[i] = try t.rig.input(std.mem.sliceAsBytes(x)); + } + t.ins[3] = try t.rig.input(std.mem.sliceAsBytes(t.nw)); + t.ins[4] = try t.rig.input(std.mem.sliceAsBytes(t.rot[0..half])); + t.ins[5] = try t.rig.input(std.mem.sliceAsBytes(t.rot[half..])); + for (&t.outs) |*o| o.* = try t.rig.output(shape.rows * w * 2); + return true; + } + + pub fn run(t: *Qk) !void { + const p = t.rig.begin(); + const w = t.width(); + const qkv: [3]ks.Ref = if (t.packed_rows) .{ kref(t.ins[0]), kat(t.ins[0], w * 2), kat(t.ins[0], 2 * w * 2) } else .{ kref(t.ins[0]), kref(t.ins[1]), kref(t.ins[2]) }; + ks.heads(p.e, t.pipe, qkv[0], qkv[1], qkv[2], kref(t.ins[3]), kat(t.ins[3], D * 2), kref(t.ins[4]), kref(t.ins[5]), kref(t.outs[0]), kref(t.outs[1]), kref(t.outs[2]), t.shape.rows, t.shape.n, if (t.packed_rows) 3 * w else w, cfg.eps); + try t.rig.finish(p); + } + + pub fn count(t: *Qk) usize { + return 3 * t.shape.rows * t.width(); + } + + /// Head h of row r of q (0) or k (1) with its sum of squares `ss` (range `ssr`), into its outputs. + fn head(t: *Qk, ends: *Ends, which: usize, r: usize, h: usize, ss: f64, ssr: Range, want: []f64, mag: []f64) void { + const rows = t.shape.rows; + const x = t.qkv[which][r * t.width() + h * D ..][0..D]; + const w = t.nw[which * D ..][0..D]; + const n = mi.invNorm(ss, ssr, D); + var normed: [D]f64 = undefined; + var nr: [D]Range = undefined; + for (x, w, &normed, &nr) |xv, wv, *o, *orr| { + const f: f64 = c.fromBf16(xv); + const wf: f64 = c.fromBf16(wv); + o.* = bf(bf(f * n.inv) * wf); + orr.* = Range.of(f).mul(n.r).fp32(0.5).bf16().mul(Range.of(wf)).fp32(0.5).bf16(); + } + const half = rows * D / 2; + const base = which * rows * t.width() + (h * rows + r) * D; + for (0..D) |ch| { + const j = ch % (D / 2); + const partner = if (ch < D / 2) ch + D / 2 else ch - D / 2; + const sign: f64 = if (ch < D / 2) -1 else 1; + const cs: f64 = t.rot[r * D / 2 + j]; + const sn: f64 = t.rot[half + r * D / 2 + 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); + } + } + + /// v's heads, copied bit for bit. + fn values(t: *Qk, ends: *Ends, want: []f64, mag: []f64) void { + const rows = t.shape.rows; + for (0..rows) |r| for (0..t.shape.n) |h| for (0..D) |ch| { + const f: f64 = c.fromBf16(t.qkv[2][r * t.width() + h * D + ch]); + ends.put(want, mag, 2 * rows * t.width() + (h * rows + r) * D + ch, f, Range.of(f)); + }; + } + + pub fn reference(t: *Qk, 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.qkv[which][r * t.width() + h * D ..][0..D], qkDepth(t.shape), t.fill == .stress); + t.head(&t.ends, which, r, h, s.ss, s.r, want, mag); + }; + t.values(&t.ends, want, mag); + } + pub fn bounds(t: *Qk, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *Qk, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } + + pub fn simulate(t: *Qk, acc: c.Accum, order: g.Order, out: []f64) !void { + const mag = try t.a.alloc(f64, out.len); + defer t.a.free(mag); + var scratch: Ends = .{ .a = t.a }; + defer scratch.deinit(); + try scratch.size(t.count()); + var terms: [D]f64 = undefined; + for (0..2) |which| for (0..t.shape.rows) |r| for (0..t.shape.n) |h| { + for (&terms, t.qkv[which][r * t.width() + h * D ..][0..D]) |*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.values(&scratch, out, mag); + } +}; diff --git a/zig/tests/contract/metal_moss_image_steps.zig b/zig/tests/contract/metal_moss_image_steps.zig new file mode 100644 index 0000000000..3d5a0eb585 --- /dev/null +++ b/zig/tests/contract/metal_moss_image_steps.zig @@ -0,0 +1,340 @@ +//! mi_situ, mi_gate_add and mi_bias under the contract: elementwise steps at the publisher's bf16 points. +const std = @import("std"); +const mtl = @import("metal"); +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 cfg = mi.cfg; +const D = mi.D; +const Range = mi.Range; +const Ends = mi.Ends; +const bf = mi.bf; +const kref = mi.kref; +const kat = mi.kat; + +// --- mi_situ: (4 tanh(g / 4) sigmoid(g)) (25 tanh(h / 25)), each op rounded to bf16 --------------------------------- + +pub const Situ = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + /// The joint blocks' form: hidden and gate halves of one (rows, 2 x width) projection; else two buffers. + halves: bool, + name: []const u8, + shape: c.Shape = undefined, + h: []u16 = &.{}, + gv: []u16 = &.{}, + ins: [2]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + /// Rows: positions; n: the MLP's width (k unused). + pub fn spec(name: []const u8, halves: bool) c.Spec { + return .{ + .name = name, + .accum = .none, + .depth = mi.noSums, + .ops = 2 * (mi.exp_ulps + 2 * mi.tanh_ulps + mi.div_ulps) + 4, + .mids = 8, + .out = .bf16, + .tiles = .{ .{}, .{ .size = 256 }, .{} }, + .served = if (halves) + &.{ .{ .rows = 4188, .n = cfg.ffn, .k = 1 }, .{ .rows = 4130, .n = cfg.ffn, .k = 1 } } + else + &.{ .{ .rows = 63, .n = cfg.ffn, .k = 1 }, .{ .rows = 10, .n = cfg.ffn, .k = 1 } }, + }; + } + + pub fn deinit(t: *Situ) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *Situ) void { + t.a.free(t.h); + t.a.free(t.gv); + t.h = &.{}; + t.gv = &.{}; + } + pub fn kernel(t: *Situ, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name, t.halves), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Situ); + + pub fn prepare(t: *Situ, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + const n = shape.rows * shape.n; + t.h = try t.a.alloc(u16, n); + t.gv = try t.a.alloc(u16, n); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + // gates past +-20 reach the sigmoid's flat ends, hidden values past +-50 the clamp's + mi.randomBf16(rng, t.h, -80, 80); + mi.randomBf16(rng, t.gv, -24, 24); + if (fill == .nonfinite) { + const at = [_]usize{ 0, 1, 2 }; + for (at, [_]u16{ mi.nan16, mi.inf16, mi.ninf16 }) |p, v| { + t.h[(p * 7) % n] = v; + t.gv[(p * 11 + 3) % n] = v; + } + } + if (t.halves) { + const both = try t.a.alloc(u16, 2 * n); + defer t.a.free(both); + for (0..shape.rows) |r| { + @memcpy(both[2 * r * shape.n ..][0..shape.n], t.h[r * shape.n ..][0..shape.n]); + @memcpy(both[(2 * r + 1) * shape.n ..][0..shape.n], t.gv[r * shape.n ..][0..shape.n]); + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(both)); + } else { + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.h)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.gv)); + } + t.out = try t.rig.output(n * 2); + return true; + } + + pub fn run(t: *Situ) !void { + const p = t.rig.begin(); + const w = t.shape.n; + const gate = if (t.halves) kat(t.ins[0], w * 2) else kref(t.ins[1]); + ks.situ(p.e, t.pipe, kref(t.ins[0]), gate, kref(t.out), t.shape.rows, w, if (t.halves) 2 * w else w); + try t.rig.finish(p); + } + + pub fn count(t: *Situ) usize { + return t.shape.rows * t.shape.n; + } + + pub fn reference(t: *Situ, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + const fifth: f64 = @as(f32, 1.0 / 25.0); + for (t.h, t.gv, 0..) |hb, gb, i| { + const h: f64 = c.fromBf16(hb); + const gv: f64 = c.fromBf16(gb); + const sig = bf(1 / (1 + @exp(-gv))); + const left = bf(bf(4 * bf(mi.tanh(bf(gv * 0.25)))) * sig); + const right = bf(25 * bf(mi.tanh(bf(h * fifth)))); + const sr = Range.of(@exp(-gv)).fp32(mi.exp_ulps).add(Range.of(1)).fp32(0.5).recip().fp32(mi.div_ulps).bf16(); + const lr = Range.of(gv * 0.25).fp32(0.5).bf16().up(mi.tanh).fp32(mi.tanh_ulps).bf16().mul(Range.of(4)).fp32(0.5).bf16().mul(sr).fp32(0.5).bf16(); + const rr = Range.of(h * fifth).fp32(0.5).bf16().up(mi.tanh).fp32(mi.tanh_ulps).bf16().mul(Range.of(25)).fp32(0.5).bf16(); + t.ends.put(want, mag, i, bf(left * right), lr.mul(rr).fp32(0.5).bf16()); + } + } + pub fn bounds(t: *Situ, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *Situ, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } +}; + +// --- mi_gate_add: x = bf16(x + bf16(g y)), in place -------------------------------------------------------------- + +pub const GateAdd = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + name: []const u8, + shape: c.Shape = undefined, + old: []u16 = &.{}, + y: []u16 = &.{}, + gv: []u16 = &.{}, + ins: [2]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + /// Rows: positions; n: the stream's width (k unused). + pub fn spec(name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .none, + .depth = mi.noSums, + .ops = 1, + .mids = 1, + .out = .bf16, + .in_place = true, + .tiles = .{ .{}, .{ .size = 256 }, .{} }, + .served = &.{ .{ .rows = 4188, .n = cfg.hidden, .k = 1 }, .{ .rows = 63, .n = cfg.hidden, .k = 1 } }, + }; + } + + pub fn deinit(t: *GateAdd) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *GateAdd) void { + t.a.free(t.old); + t.a.free(t.y); + t.a.free(t.gv); + t.old = &.{}; + t.y = &.{}; + t.gv = &.{}; + } + pub fn kernel(t: *GateAdd, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = blk: { + var hooks = g.hooks(GateAdd); + hooks.overlap = overlap; + break :blk hooks; + }; + /// The stream is its own output. + fn overlap(_: *anyopaque) bool { + return true; + } + + pub fn prepare(t: *GateAdd, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + const n = shape.rows * shape.n; + t.old = try t.a.alloc(u16, n); + t.y = try t.a.alloc(u16, n); + t.gv = try t.a.alloc(u16, shape.n); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + mi.randomBf16(rng, t.old, -8, 8); + mi.randomBf16(rng, t.y, -4, 4); + mi.randomBf16(rng, t.gv, -2, 2); + if (fill == .nonfinite) { + for ([_]u16{ mi.nan16, mi.inf16, mi.ninf16 }, 0..) |v, i| { + t.old[(i * 13) % n] = v; + t.y[(i * 17 + 1) % n] = v; + } + if (shape.n > 6) t.gv[5] = mi.nan16; + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.y)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.gv)); + t.out = try t.rig.output(n * 2); + return true; + } + + pub fn run(t: *GateAdd) !void { + // in place: the stream's rows go back before each independent run + @memcpy(t.rig.bytes(t.out, t.count() * 2), std.mem.sliceAsBytes(t.old)); + const p = t.rig.begin(); + ks.gateAdd(p.e, t.pipe, kref(t.out), kref(t.ins[0]), kref(t.ins[1]), t.shape.rows, t.shape.n); + try t.rig.finish(p); + } + + pub fn count(t: *GateAdd) usize { + return t.shape.rows * t.shape.n; + } + + pub fn reference(t: *GateAdd, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + for (t.old, t.y, 0..) |xb, yb, i| { + const x: f64 = c.fromBf16(xb); + const y: f64 = c.fromBf16(yb); + const gv: f64 = c.fromBf16(t.gv[i % t.shape.n]); + const r = Range.of(gv * y).fp32(0.5).bf16().add(Range.of(x)).fp32(0.5).bf16(); + t.ends.put(want, mag, i, bf(x + bf(gv * y)), r); + } + } + pub fn bounds(t: *GateAdd, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *GateAdd, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } +}; + +// --- mi_bias: y = bf16(c + b), a projection's fp32 sums and its bias rounded once -------------------------------- + +pub const Bias = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + name: []const u8, + shape: c.Shape = undefined, + sums: []f32 = &.{}, + b: []u16 = &.{}, + ins: [2]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + /// Rows: positions; n: the projection's outputs (k unused). + pub fn spec(name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .none, + .depth = mi.noSums, + .ops = 1, + .out = .bf16, + .tiles = .{ .{}, .{ .size = 256 }, .{} }, + .served = &.{ .{ .rows = 4125, .n = cfg.hidden, .k = 1 }, .{ .rows = 63, .n = 3 * cfg.hidden, .k = 1 }, .{ .rows = 63, .n = cfg.ffn, .k = 1 }, .{ .rows = 63, .n = cfg.hidden, .k = 1 } }, + }; + } + + pub fn deinit(t: *Bias) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *Bias) void { + t.a.free(t.sums); + t.a.free(t.b); + t.sums = &.{}; + t.b = &.{}; + } + pub fn kernel(t: *Bias, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Bias); + + pub fn prepare(t: *Bias, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + const n = shape.rows * shape.n; + t.sums = try t.a.alloc(f32, n); + t.b = try t.a.alloc(u16, shape.n); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (t.sums) |*v| v.* = (rng.float(f32) * 2 - 1) * 16; + mi.randomBf16(rng, t.b, -2, 2); + if (fill == .nonfinite) { + for ([_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32) }, 0..) |v, i| t.sums[(i * 13) % n] = v; + if (shape.n > 6) t.b[5] = mi.inf16; + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.sums)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.b)); + t.out = try t.rig.output(n * 2); + return true; + } + + pub fn run(t: *Bias) !void { + const p = t.rig.begin(); + ks.bias(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: *Bias) usize { + return t.shape.rows * t.shape.n; + } + + pub fn reference(t: *Bias, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + for (t.sums, 0..) |s, i| { + const b: f64 = c.fromBf16(t.b[i % t.shape.n]); + t.ends.put(want, mag, i, bf(s + b), Range.of(s).add(Range.of(b)).fp32(0.5).bf16()); + } + } + pub fn bounds(t: *Bias, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *Bias, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } +}; From 7740f9dc9b7f98a3ce08495368b521a67cd5619e Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Tue, 27 Oct 2026 06:43:39 -0500 Subject: [PATCH 3/6] docs: MOSS-Image's transformer on Metal, how it was checked and measured --- docs/recipes/moss-image.md | 118 +++++++++++++++++++++++++++++++++++++ 1 file changed, 118 insertions(+) create mode 100644 docs/recipes/moss-image.md diff --git a/docs/recipes/moss-image.md b/docs/recipes/moss-image.md new file mode 100644 index 0000000000..3ca4c46115 --- /dev/null +++ b/docs/recipes/moss-image.md @@ -0,0 +1,118 @@ +# MOSS-Image on Metal + +MOSS-Image is a text-to-image model. A 5B joint transformer (42 blocks, hidden 3,072, 24 heads of 128, an MLP of +8,192 with the "situ" gate) denoises 128-channel latents by flow matching. `zig/src/families/moss_image` runs that +transformer on the Metal runtime from the released bf16 diffusers checkpoint: + +- a prompt's rows once a prompt: the input projection and the two refiner blocks (bidirectional attention, biased + projections, no modulation); +- the image and prompt rows together at each step: every joint block, one modulation from the timestep shared by all of + them, full attention over both. + +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. + +It is a denoiser, not a token family. It does not register with the lane engine, the server or `families.detect`. + +It runs on the M5's tensor units alone, the only chip it was checked on: `Model.load` refuses a GPU without them. + +## Checkpoint + +`OpenMOSS-Team/MOSS-Image-Base` at `85975793d8b3d432f1b6e27f352abaa4c2a4c0d3`, Apache 2.0. `transformer/config.json` +must be `MOSSImageTransformer` at the released geometry; any other config is refused, and so is a checkpoint with +tensors this family does not read. + +## Precision + +- bf16 activations, rounded where the publisher's torch code rounds them: once per elementwise op, as torch's bf16 + tensors do. +- fp32 sums in every projection and in attention. Biased projections add the bias to the fp32 sums and round once, as + a bf16 GEMM's bias epilogue does: their sums come from the prefill split-K GEMM (one part), then `mi_bias`. +- fp32 for norms, softmax and the time sinusoid's input (the sinusoid itself in float64, as the publisher computes it). +- The rotary tables in bf16, as the pipeline casts them; its frequencies are the checkpoint's `rope.bands`. +- No int8. + +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. + +## Kernel contract + +`zig/tests/contract/metal_moss_image*.zig` registers every kernel of `moss_image.metal`: + +- `mi_norm`, plain and modulated; +- `mi_qk`, in both row layouts: q, k and v in three buffers, and packed in one row; +- `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. + +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. + +`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 +``` + +## Run a case + +```sh +python3 tools/zig/moss_image_case.py case MODEL_DIR case/ --prompt-file prompt.txt # needs torch, diffusers, transformers +zig build tf-moss-image -Dcpu=apple_m1 +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 +``` + +The case holds: + +- both prompts' encoder rows; +- the timesteps and noise levels; +- the starting latents; +- the publisher's pipeline's input and both velocities at every step; +- with `--fp32-steps N`, a float32 run of the same transformer for the first N steps. + +`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. + +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. + +## Measured + +Setup: + +- Mac Studio M5 Ultra 256 GB, macOS 27; this commit's `tf-moss-image`. +- The card's defaults: 50 Euler steps, guidance 3, the empty negative prompt skipping block 22. +- Reference: the checkpoint's `pipeline_moss_image.py` on torch 2.14.1 (MPS, bf16), diffusers 0.42.0, transformers + 5.21.0. +- Two cases: a fox at 1024x1024 (4,096 image rows, 34 prompt rows, seed 42), and the card's Chinese sticky-note prompt + at 1200x880 (an odd 55 x 75 grid, 4,125 image rows, 63 prompt rows, seed 7). The negative prompt has 10 rows. + +Results: + +- **Velocity**, both branches of each of the 50 steps on the pipeline's own input: + - Against torch bf16: cosine 0.99973 to 0.99998, worst at the last steps. + - Against a float32 run of the same transformer on the same inputs (the second case, every step): mean 1 - cosine + 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), our picture is the closer to + it 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. +- **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). + - Ours: 0.489 s and 0.486 s a forward, 1.52x and 1.56x torch's (the median of three runs of 50 guided steps; + 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. From 62fc3120293d594f478b4ef27e1bb31b284cfa49 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Tue, 27 Oct 2026 09:17:40 -0500 Subject: [PATCH 4/6] MOSS-Image: review fixes tf-moss-image no longer leaks the first prompt when the second fails, Model.load says why it refuses a GPU without tensor units, and moss_image_case.py imports the checkpoint's pipeline_moss_image.py only when its sha256 is the pinned revision's (the docs say so, with the download line). --- docs/recipes/moss-image.md | 4 ++++ tools/zig/moss_image_case.py | 13 +++++++++++-- zig/src/families/moss_image/model.zig | 5 ++++- zig/tests/moss_image.zig | 7 +++++-- 4 files changed, 24 insertions(+), 5 deletions(-) diff --git a/docs/recipes/moss-image.md b/docs/recipes/moss-image.md index 3ca4c46115..608a6430c1 100644 --- a/docs/recipes/moss-image.md +++ b/docs/recipes/moss-image.md @@ -62,7 +62,11 @@ SUMMARY 11 kernels: 11 pass gates 1-4, 0 fail, 0 known findings; 8 controls: 8 r ## Run a case +`moss_image_case.py` imports and runs the checkpoint's own `pipeline_moss_image.py`, the publisher's Python. It does so +only when the file's sha256 is the pinned revision's (`1dac1aca…8b31`) and refuses any other file. + ```sh +hf download OpenMOSS-Team/MOSS-Image-Base --revision 85975793d8b3d432f1b6e27f352abaa4c2a4c0d3 --local-dir MODEL_DIR python3 tools/zig/moss_image_case.py case MODEL_DIR case/ --prompt-file prompt.txt # needs torch, diffusers, transformers zig build tf-moss-image -Dcpu=apple_m1 zig-out/bin/tf-moss-image MODEL_DIR/transformer case/case.safetensors case/ diff --git a/tools/zig/moss_image_case.py b/tools/zig/moss_image_case.py index c3b9eb6dc7..30f3df5217 100644 --- a/tools/zig/moss_image_case.py +++ b/tools/zig/moss_image_case.py @@ -4,6 +4,7 @@ from __future__ import annotations import argparse +import hashlib import importlib.util import json import math @@ -22,9 +23,17 @@ """ +# pipeline_moss_image.py at OpenMOSS-Team/MOSS-Image-Base revision 85975793d8b3d432f1b6e27f352abaa4c2a4c0d3 +PIPELINE_SHA256 = "1dac1aca9feb1bfdfc727a401e4b67d61a505d3190706d1b7d2d784f88678b31" + + def module(model_dir: str): - """The checkpoint's pipeline_moss_image.py, imported from its folder.""" - spec = importlib.util.spec_from_file_location("pipeline_moss_image", Path(model_dir) / "pipeline_moss_image.py") + """The checkpoint's pipeline_moss_image.py, imported from its folder once its sha256 is the pinned revision's.""" + path = Path(model_dir) / "pipeline_moss_image.py" + digest = hashlib.sha256(path.read_bytes()).hexdigest() + if digest != PIPELINE_SHA256: + sys.exit(f"{path} has sha256 {digest}; this tool runs only the pinned revision's ({PIPELINE_SHA256})") + spec = importlib.util.spec_from_file_location("pipeline_moss_image", path) mod = importlib.util.module_from_spec(spec) spec.loader.exec_module(mod) return mod diff --git a/zig/src/families/moss_image/model.zig b/zig/src/families/moss_image/model.zig index 2655c66508..374d8cb0c3 100644 --- a/zig/src/families/moss_image/model.zig +++ b/zig/src/families/moss_image/model.zig @@ -113,7 +113,10 @@ pub const Model = struct { const plan = try Scratch.of(max_text, image); const scratch_bytes = try plan.bytes(); // qualified on the M5's tensor units alone: the projections and attention take another path before them - if (!device.tensorUnits()) return error.TensorUnitsRequired; + if (!device.tensorUnits()) { + std.log.err("moss image: {s} has no tensor units; this family runs on the M5's alone, the only chip it was checked on", .{device.name()}); + return error.TensorUnitsRequired; + } const t = try gpa.create(Model); errdefer gpa.destroy(t); t.* = .{ .gpa = gpa, .device = device, .queue = undefined, .k = undefined, .w = undefined, .max_text = max_text, .image = image, .rows = g.rows, .columns = g.columns, .scratch = undefined, .x = undefined, .xn = undefined, .qr = undefined, .kr = undefined, .vr = undefined, .att = undefined, .y = undefined, .ug = undefined, .s = undefined, .q = undefined, .k_ = undefined, .v = undefined, .sums = undefined, .lat = undefined, .vel = undefined, .ctx = undefined, .mods = undefined }; diff --git a/zig/tests/moss_image.zig b/zig/tests/moss_image.zig index 35908917e0..af5cbccf07 100644 --- a/zig/tests/moss_image.zig +++ b/zig/tests/moss_image.zig @@ -148,8 +148,11 @@ pub fn main(init: std.process.Init) !u8 { var out: std.ArrayList(u8) = .empty; const refining = mtl.clock.seconds(); - var branches = [2]mi.Prompt{ try model.prompt(prompt), try model.prompt(negative) }; - defer for (&branches) |*b| b.deinit(); + var branches: [2]mi.Prompt = undefined; + branches[0] = try model.prompt(prompt); + defer branches[0].deinit(); + branches[1] = try model.prompt(negative); + defer branches[1].deinit(); 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 }); From 6aa59372087ae022c36d0a870ea594ca567376ac Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Tue, 27 Oct 2026 15:17:50 -0500 Subject: [PATCH 5/6] docs: MOSS-Image's float32 ruler, both with torch's bf16 encoder rows and with the encoder in float32 --- docs/recipes/moss-image.md | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/docs/recipes/moss-image.md b/docs/recipes/moss-image.md index 608a6430c1..137390227f 100644 --- a/docs/recipes/moss-image.md +++ b/docs/recipes/moss-image.md @@ -110,9 +110,11 @@ Results: 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), our picture is the closer to - it 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. + 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. 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. - **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). From 0ebed8b8bc3248c2d5fe558ed8cfd7e5afce0cb7 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Tue, 27 Oct 2026 16:24:00 -0500 Subject: [PATCH 6/6] docs: MOSS-Image's float32-encoder ruler came from the prompt encoder's follow-up tool --- docs/recipes/moss-image.md | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/docs/recipes/moss-image.md b/docs/recipes/moss-image.md index 137390227f..5680d3e0e1 100644 --- a/docs/recipes/moss-image.md +++ b/docs/recipes/moss-image.md @@ -113,8 +113,10 @@ Results: 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. 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. + 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. - **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).