From a428b8b784b39771d44f97d04bfa41a640205ec1 Mon Sep 17 00:00:00 2001 From: drowzeys <264866615+drowzeys@users.noreply.github.com> Date: Fri, 16 Oct 2026 08:57:08 -0400 Subject: [PATCH 1/2] MiniMax H3 / FastH3 transformer on Metal in bf16, as a test tool beside the engine A video and audio model's denoising transformer on the M5 tensor units, in the checkpoint's precision: bf16 weights read in place, products on the prefill kernels' tfp::gemm_nax (MLX's GEMM, fp32 sums), FastH3's tile-routed attention in bf16 with fp32 scores and an online softmax. zig build tf-h3-dit runs every denoising forward for a case file. The lane core, the server and the other families are untouched. Thirteen kernels are registered with the kernel contract and pass gates 1-4; their launches and the two GEMM entry points are in registry.zig. The recipe holds the receipt. --- build.zig | 12 + docs/recipes/minimax-h3-metal.md | 84 ++++++ zig/build/kernel_contract.zig | 1 + zig/kernels/metal/h3/h3_attention.metal | 113 ++++++++ zig/kernels/metal/h3/h3_routing.metal | 153 ++++++++++ zig/kernels/metal/h3/h3_rows.metal | 153 ++++++++++ zig/kernels/metal/h3/kernels.zig | 6 + zig/src/core/contract/registry.zig | 15 + zig/src/families/h3/block.zig | 129 +++++++++ zig/src/families/h3/h3.zig | 369 ++++++++++++++++++++++++ zig/src/families/h3/tensors.zig | 97 +++++++ zig/tests/contract/metal_h3.zig | 211 ++++++++++++++ zig/tests/contract/metal_h3_prep.zig | 311 ++++++++++++++++++++ zig/tests/contract/metal_h3_ref.zig | 266 +++++++++++++++++ zig/tests/contract/metal_h3_topk.zig | 136 +++++++++ zig/tests/contract/metal_limits.zig | 7 + zig/tests/contract/metal_main.zig | 17 ++ zig/tests/h3_dit.zig | 116 ++++++++ 18 files changed, 2196 insertions(+) create mode 100644 docs/recipes/minimax-h3-metal.md create mode 100644 zig/kernels/metal/h3/h3_attention.metal create mode 100644 zig/kernels/metal/h3/h3_routing.metal create mode 100644 zig/kernels/metal/h3/h3_rows.metal create mode 100644 zig/kernels/metal/h3/kernels.zig create mode 100644 zig/src/families/h3/block.zig create mode 100644 zig/src/families/h3/h3.zig create mode 100644 zig/src/families/h3/tensors.zig create mode 100644 zig/tests/contract/metal_h3.zig create mode 100644 zig/tests/contract/metal_h3_prep.zig create mode 100644 zig/tests/contract/metal_h3_ref.zig create mode 100644 zig/tests/contract/metal_h3_topk.zig create mode 100644 zig/tests/h3_dit.zig diff --git a/build.zig b/build.zig index 380a38b3de..eb8ce13219 100644 --- a/build.zig +++ b/build.zig @@ -460,6 +460,18 @@ fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b glm_checks.sampler(b, target, optimize, metal, lanes, sources); + // MiniMax H3 / FastH3's transformer: the packed rows' denoising forwards against a case file (M5 tensor units). + const h3_kernels = b.createModule(.{ .root_source_file = b.path("zig/kernels/metal/h3/kernels.zig") }); + const h3_safetensors = b.createModule(.{ .root_source_file = b.path("zig/src/core/safetensors.zig"), .target = target, .optimize = optimize }); + const h3_mod = b.createModule(.{ .root_source_file = b.path("zig/src/families/h3/h3.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "safetensors", .module = h3_safetensors }, .{ .name = "h3_kernels", .module = h3_kernels }, .{ .name = "kernel_sources", .module = sources } } }); + const h3_dit = b.addExecutable(.{ .name = "tf-h3-dit", .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/h3_dit.zig"), + .target = target, + .optimize = optimize, + .link_libc = true, + .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "h3", .module = h3_mod } }, + }) }); + b.step("tf-h3-dit", "FastH3's Metal transformer against a case file: velocities, every step, timings").dependOn(&b.addInstallArtifact(h3_dit, .{}).step); // The core row projection (chips without tensor units) on synthetic matrices; runs on any Mac's GPU const row_mod = b.createModule(.{ .root_source_file = b.path("zig/src/core/row_projection.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "kernel_sources", .module = sources } } }); const row_check = b.addExecutable(.{ .name = "tf-row-check", .root_module = b.createModule(.{ diff --git a/docs/recipes/minimax-h3-metal.md b/docs/recipes/minimax-h3-metal.md new file mode 100644 index 0000000000..288e1029da --- /dev/null +++ b/docs/recipes/minimax-h3-metal.md @@ -0,0 +1,84 @@ +# MiniMax H3 / FastH3 transformer on Metal (M5 tensor units) + +This is a video model's denoising transformer, not a token model. It has no lane round and no drafted tokens. It +keeps the checkpoint's precision: bf16 weights read in place from the checkpoint, bf16 activations, fp32 sums. It is +a test tool beside the engine; nothing in the lane core, the server or another family changes. + +## What it is + +`zig build tf-h3-dit -Dcpu=apple_m1 -Doptimize=ReleaseFast` builds `zig-out/bin/tf-h3-dit`. It reads the FastH3 +checkpoint's shards and a case file, runs every denoising forward for the packed [text | audio | video] +rows, prints the first forward's velocities against the reference velocities the case holds, and writes the final +video and audio rows (and the first forward's velocities) as raw float32. + +```sh +zig-out/bin/tf-h3-dit case.safetensors out_prefix [dense | profile | DEFINE, ...] +``` + +It needs the M5's tensor units and refuses to load without them. `dense` makes every query tile attend to every tile +and leaves the pooled branch out (not how the checkpoint was trained; for kernel work only). `profile` times one more +forward stage by stage. + +## Arithmetic + +- Block projections: the prefill kernels' `tfp::gemm_nax` (MLX 0.32.3's NAX GEMM), bf16 x w^T as the checkpoint + stores w, fp32 sums, bf16 out. The SwiGLU's first projection is one product over its stored [value; gate] rows. +- Norms, rotary, the gated adds and the SwiGLU round to bf16 where the reference's bf16 MLX arithmetic rounds. +- Attention is FastH3's video sparse attention over 64-row tiles: a text or audio tile attends every + tile; a video tile attends the prefix tiles and its top 20% of video tiles by pooled score; the pooled branch is + gated in. Tile means, scores, the top-k and the pooled branch are fp32. The routed attention is bf16 q, k and v + on the tensor units with fp32 scores and an online softmax across key tiles; its weights are rounded to bf16 + before the value product, as MLX's attention rounds them. + +## Kernel contract + +Thirteen kernels are registered (`zig/tests/contract/metal_h3*.zig`), each against its formula in float64, and their +launches in `registry.zig`. The two GEMM entry points are the prefill source's own and are registered as launches. +Kernels whose outputs pass through intermediate bf16 roundings declare them as `mids`; their test inputs are of one +sign where a sum could cancel, so no rounding hides behind a small result. + +```sh +zig build tf-kernel-contract -Dcpu=apple_m1 -Doptimize=ReleaseFast +zig-out/bin/tf-kernel-contract --only h3 +``` + +## The case file + +The prompt's rows, the schedule and the tile map come in one safetensors file. The exporter that writes one is not in +this change: it uses the Python 0.6 H3 family, which lives on a fork (`drowzeys/TensorFold`, `tools/zig/h3_case.py` on +branch `h3-1.0.4`). + +| Tensor | Holds | +|---|---| +| `text` | The prompt's rows after the condition projection and the token refiner, bf16 | +| `video`, `audio` | Starting noise rows | +| `cos`, `sin` | Rotary tables for every packed row | +| `adaln`, `times`, `tables.N`, `final` | Each row's line in a block's six modulation tables per step, and the tables | +| `video_step`, `audio_step` | Per step: the sigma the model sees and the Euler ratio | +| `tile_slot`, `tile_sizes`, `geometry` | FastVideo's tile map: each row's padded slot, each tile's rows, the counts | +| `video_in.*`, `audio_in.*`, `video_out.*`, `audio_out.*`, `final_norm.weight` | The small float projections | +| `video_velocity`, `audio_velocity` | The reference's first forward, for the parity line the tool prints | + +## Receipt + +Mac Studio, Apple M5 Ultra, 256 GB, macOS 27.0.1, Zig 0.17.0. Checkpoint `FastVideo/FastVideo-FastH3-8-Step-V2`, +864x480, 124 frames, 8 steps, 15,435 packed rows (36 text, 414 audio, 14,985 video), 288 tiles. Reference: the +Python 0.6 family in bf16 on MLX 0.32.3 with FastVideo's reference routing; with FastVideo's fp Metal sparse attention +it takes 27.9 s a forward on the same Mac. Peak process footprint 4.3 GB, the checkpoint mapped read-only beside it. +The pull request has the full tables. + +| Run | A forward (median) | Video velocity against the reference | Audio velocity against the reference | +|---|---|---|---| +| Step 1 of 8 | 10.1 s | cosine 0.98981, relative error 0.143 | cosine 0.99987, relative error 0.017 | +| Step 5 of 8 | 10.1 s | cosine 0.99838, relative error 0.057 | cosine 0.99980, relative error 0.020 | + +At step 1 every video tile is noise, so FastH3's top-k among them is decided by differences far below bf16. The same +binary with only the fp32 summation order of the tile scores reversed differs from itself by video cosine 0.99021 at +step 1 and 0.99834 at step 5: the comparison with the reference is at that floor. + +## Not done + +- No exporter in the tree, so the tool cannot be exercised from this repository alone. +- One machine (M5 Ultra). Chips without tensor units are refused. +- Text to video only; a clip that starts from an image is a follow-up change. +- No `tensorfold generate`, no text encoder, decoders or sampler: the forward only. diff --git a/zig/build/kernel_contract.zig b/zig/build/kernel_contract.zig index efcdef5194..7b2882edae 100644 --- a/zig/build/kernel_contract.zig +++ b/zig/build/kernel_contract.zig @@ -24,6 +24,7 @@ pub fn metal(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.buil .{ .name = "tensorfold", .module = engine }, .{ .name = "kernel_sources", .module = sources }, .{ .name = "contract", .module = module(b, target, optimize) }, + .{ .name = "h3_kernels", .module = b.createModule(.{ .root_source_file = b.path("zig/kernels/metal/h3/kernels.zig") }) }, }, }) }); const install = b.addInstallArtifact(exe, .{}); diff --git a/zig/kernels/metal/h3/h3_attention.metal b/zig/kernels/metal/h3/h3_attention.metal new file mode 100644 index 0000000000..138b26f6aa --- /dev/null +++ b/zig/kernels/metal/h3/h3_attention.metal @@ -0,0 +1,113 @@ +// MiniMax H3 / FastH3 on Metal: tile-routed attention in bf16 on the tensor units, fp32 sums and an online softmax. + +// softmax(q k^T) v over a query tile's key tiles in bf16, fp32 sums and an online softmax. [query tiles, H] of 256. +[[kernel]] void h3_attention( + const device bfloat* Q [[buffer(0)]], + const device bfloat* K [[buffer(1)]], + const device bfloat* V [[buffer(2)]], + const device int32_t* IDX [[buffer(3)]], + const device int32_t* SIZES [[buffer(4)]], + const device int32_t* ROWOF [[buffer(5)]], + const constant int32_t* P [[buffer(6)]], + const constant float* SC [[buffer(7)]], + device bfloat* Y [[buffer(8)]], + uint tid [[thread_index_in_threadgroup]], + uint3 tg [[threadgroup_position_in_grid]]) { + using namespace mpp::tensor_ops; + constexpr int D = 128, TQ = 64, KEYS = 64, TS = H3_SLOTS, H = H3_HEADS; + constexpr int TPR = 256 / TQ; + constexpr int C1 = TQ * KEYS / 256, C2 = TQ * D / 256; + const int NQ = P[0], KSEL = P[1]; + const int head = int(tg.y), r0 = (P[2] + int(tg.x)) * TQ; + const long hb = long(head) * TS; + const device int32_t* list = IDX + (P[3] ? (long(head) * NQ + tg.x) * KSEL : 0); + tensor, tensor_inline> q((device bfloat*)Q + hb * D, dextents(D, TS)); + tensor, tensor_inline> k((device bfloat*)K + hb * D, dextents(D, TS)); + tensor, tensor_inline> v((device bfloat*)V + hb * D, dextents(D, TS)); + threadgroup float st[TQ * KEYS]; + threadgroup bfloat pt[TQ * KEYS]; + threadgroup float tops[256]; + threadgroup float fac[TQ]; + tensor, tensor_inline> p((threadgroup bfloat*)pt, dextents(KEYS, TQ)); + constexpr auto d1 = matmul2d_descriptor(TQ, KEYS, D, false, true, true, matmul2d_descriptor::mode::multiply_accumulate); + constexpr auto d2 = matmul2d_descriptor(TQ, D, KEYS, false, false, true, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d> op1; + matmul2d> op2; + auto a = q.slice(0, r0); + auto b0 = k.slice(0, 0); + auto v0 = v.slice(0, 0); + auto acc = op1.template get_destination_cooperative_tensor(); + auto out = op2.template get_destination_cooperative_tensor(); + short srow[C1], scol[C1], orow[C2], ocol[C2]; + H3_UNROLL + for (ushort i = 0; i < C1; i++) { + auto ids = acc.get_multidimensional_index(i); + scol[i] = ids[0]; + srow[i] = ids[1]; + } + H3_UNROLL + for (ushort i = 0; i < C2; i++) { + auto ids = out.get_multidimensional_index(i); + ocol[i] = ids[0]; + orow[i] = ids[1]; + out[i] = 0.0f; + } + const float scale = SC[0]; + const int row = int(tid) / TPR, first = (int(tid) % TPR) * C1; + const int mine = row * KEYS + first; + float top = -INFINITY, total = 0.0f; + float s[C1]; + for (int sel = 0; sel < KSEL; sel++) { + const int tile = list[sel]; + const int k0 = tile * KEYS, size = SIZES[tile]; + H3_UNROLL + for (ushort i = 0; i < C1; i++) acc[i] = 0.0f; + auto b = k.slice(0, k0); + op1.run(a, b, acc); + H3_UNROLL + for (ushort i = 0; i < C1; i++) st[srow[i] * KEYS + scol[i]] = scol[i] < size ? acc[i] * scale : -INFINITY; + threadgroup_barrier(mem_flags::mem_threadgroup); + float local = -INFINITY; + H3_UNROLL + for (ushort j = 0; j < C1; j++) { + s[j] = st[mine + j]; + local = max(local, s[j]); + } + tops[tid] = local; + threadgroup_barrier(mem_flags::mem_threadgroup); + float next = top; + H3_UNROLL + for (ushort j = 0; j < TPR; j++) next = max(next, tops[row * TPR + j]); + // a tile whose keys are all padding leaves next at -inf; its weights are zero and nothing shrinks + const float shrink = next == -INFINITY ? 1.0f : exp(top - next); + if (tid % TPR == 0) fac[row] = shrink; + total *= shrink; + top = next; + H3_UNROLL + for (ushort j = 0; j < C1; j++) { + const float weight = next == -INFINITY ? 0.0f : exp(s[j] - next); + total += weight; + pt[mine + j] = bfloat(weight); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + H3_UNROLL + for (ushort i = 0; i < C2; i++) out[i] *= fac[orow[i]]; + auto vb = v.slice(0, k0); + op2.run(p, vb, out); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + tops[tid] = total; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid % TPR == 0) { + float sum = 0.0f; + H3_UNROLL + for (ushort j = 0; j < TPR; j++) sum += tops[tid + j]; + fac[row] = 1.0f / sum; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + H3_UNROLL + for (ushort i = 0; i < C2; i++) { + const int to = ROWOF[r0 + orow[i]]; + if (to >= 0) Y[(long(to) * H + head) * D + ocol[i]] = bfloat(out[i] * fac[orow[i]]); + } +} diff --git a/zig/kernels/metal/h3/h3_routing.metal b/zig/kernels/metal/h3/h3_routing.metal new file mode 100644 index 0000000000..96dbfd30a1 --- /dev/null +++ b/zig/kernels/metal/h3/h3_routing.metal @@ -0,0 +1,153 @@ +// MiniMax H3 / FastH3 on Metal: q, k and v in tile order, and FastH3's routing (tile means, scores, top-k, pooling). + +// q, k, v to bf16 in tile order: head RMSNorm, then rotate-half on q and k in bf16 as the reference rounds it. +[[kernel]] void h3_heads( + 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* COS [[buffer(5)]], + const device float* SIN [[buffer(6)]], + const device int32_t* SLOT [[buffer(7)]], + const constant int32_t* P [[buffer(8)]], + const constant float* E [[buffer(9)]], + device bfloat* TQ [[buffer(10)]], + device bfloat* TK [[buffer(11)]], + device bfloat* TV [[buffer(12)]], + uint lane [[thread_index_in_simdgroup]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int H = P[1], ROT = P[2], HALF = ROT / 2; + const int row = int(tg.x), head = int(tg.y); + const int c0 = 4 * int(lane); + const long base = (long(row) * H + head) * 128; + float4 q, k; + for (int j = 0; j < 4; j++) { + q[j] = float(QR[base + c0 + j]); + k[j] = float(KR[base + c0 + j]); + } + const float qi = rsqrt(simd_sum(dot(q, q)) / 128.0f + E[0]), ki = rsqrt(simd_sum(dot(k, k)) / 128.0f + E[0]); + const device float* cs = COS + long(row) * ROT; + const device float* sn = SIN + long(row) * ROT; + const long at = (long(head) * H3_SLOTS + SLOT[row]) * 128; + for (int j = 0; j < 4; j++) { + const int c = c0 + j; + float a = float(bfloat(q[j] * qi * float(NQ[c]))), b = float(bfloat(k[j] * ki * float(NK[c]))); + if (c < ROT) { + // rotate-half: channel c pairs with c + HALF below the half and c - HALF above it; tables rounded to bf16 + const int p = c < HALF ? c + HALF : c - HALF; + const float sign = c < HALF ? -1.0f : 1.0f; + const float qp = float(bfloat(float(QR[base + p]) * qi * float(NQ[p]))); + const float kp = float(bfloat(float(KR[base + p]) * ki * float(NK[p]))); + const float co = float(bfloat(cs[c])), si = float(bfloat(sn[c])); + a = float(bfloat(a * co)) + float(bfloat(sign * qp * si)); + b = float(bfloat(b * co)) + float(bfloat(sign * kp * si)); + } + TQ[at + c] = bfloat(a); + TK[at + c] = bfloat(b); + TV[at + c] = VR[base + c]; + } +} + +// Each tile's mean q, k and v over its real rows, in float. QP, KP, VP: (H, tiles, 128). +[[kernel]] void h3_pool( + const device bfloat* TQ [[buffer(0)]], + const device bfloat* TK [[buffer(1)]], + const device bfloat* TV [[buffer(2)]], + const device int32_t* SIZES [[buffer(3)]], + device float* QP [[buffer(4)]], + device float* KP [[buffer(5)]], + device float* VP [[buffer(6)]], + uint lane [[thread_index_in_simdgroup]], + uint3 tg [[threadgroup_position_in_grid]], + uint3 grid [[threadgroups_per_grid]]) { + const int tile = int(tg.x), head = int(tg.y), size = SIZES[tile], c0 = 4 * int(lane); + const long first = long(head) * H3_SLOTS + long(tile) * 64; + float4 q = 0.0f, k = 0.0f, v = 0.0f; + for (int s = 0; s < size; s++) { + const long at = (first + s) * 128 + c0; + q += float4(float(TQ[at]), float(TQ[at + 1]), float(TQ[at + 2]), float(TQ[at + 3])); + k += float4(float(TK[at]), float(TK[at + 1]), float(TK[at + 2]), float(TK[at + 3])); + v += float4(float(TV[at]), float(TV[at + 1]), float(TV[at + 2]), float(TV[at + 3])); + } + const float n = float(size); + const long out = (long(head) * grid.x + tile) * 128 + c0; + for (int j = 0; j < 4; j++) { + QP[out + j] = q[j] / n; + KP[out + j] = k[j] / n; + VP[out + j] = v[j] / n; + } +} + +// Tile against tile: S[h, tq, tk] = QP[h, tq] . KP[h, tk] / sqrt(128). Threads [tiles, tiles, H]. +[[kernel]] void h3_tile_scores( + const device float* QP [[buffer(0)]], + const device float* KP [[buffer(1)]], + const constant int32_t* P [[buffer(2)]], + device float* S [[buffer(3)]], + uint3 at [[thread_position_in_grid]]) { + const int NT = P[0], H = P[1]; + if (int(at.x) >= NT || int(at.y) >= NT || int(at.z) >= H) return; + const device float* q = QP + (long(at.z) * NT + at.y) * 128; + const device float* k = KP + (long(at.z) * NT + at.x) * 128; + float sum = 0.0f; + for (int c = 0; c < 128; c++) sum += q[c] * k[c]; + S[(long(at.z) * NT + at.y) * NT + at.x] = sum * 0.08838834764831845f; +} + +// A video tile's key tiles: the prefix, then its KEEP best video tiles by score, in tile order. +[[kernel]] void h3_topk( + const device float* S [[buffer(0)]], + const constant int32_t* P [[buffer(1)]], + device int32_t* IDX [[buffer(2)]], + uint lane [[thread_index_in_simdgroup]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int NT = P[0], P0 = P[1], KEEP = P[2], NV = NT - P0; + const int tq = int(tg.x), head = int(tg.y); + const device float* s = S + (long(head) * NT + P0 + tq) * NT + P0; + device int32_t* out = IDX + (long(head) * NV + tq) * (P0 + KEEP); + for (int i = int(lane); i < P0; i += 32) out[i] = i; + // a tile is kept when fewer than KEEP score above it (a NaN ranks as -inf); its place counts the kept tiles before it + threadgroup uchar kept[H3_TILES]; + for (int j = int(lane); j < NV; j += 32) { + const float mine = isnan(s[j]) ? -INFINITY : s[j]; + int above = 0; + for (int i = 0; i < NV; i++) { + const float other = isnan(s[i]) ? -INFINITY : s[i]; + above += (other > mine || (other == mine && i < j)) ? 1 : 0; + } + kept[j] = above < KEEP ? 1 : 0; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int j = int(lane); j < NV; j += 32) { + if (!kept[j]) continue; + int before = 0; + for (int i = 0; i < j; i++) before += kept[i]; + out[P0 + before] = P0 + j; + } +} + +// The pooled branch: C[h, tq] = softmax over tiles of S[h, tq] times VP[h]. C: (H, tiles, 128) float. +[[kernel]] void h3_coarse( + const device float* S [[buffer(0)]], + const device float* VP [[buffer(1)]], + const constant int32_t* P [[buffer(2)]], + device float* C [[buffer(3)]], + uint lane [[thread_index_in_simdgroup]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int NT = P[0], tq = int(tg.x), head = int(tg.y), c0 = 4 * int(lane); + const device float* s = S + (long(head) * NT + tq) * NT; + float top = -3.0e38f; + for (int i = int(lane); i < NT; i += 32) top = max(top, s[i]); + top = simd_max(top); + float mass = 0.0f; + for (int i = int(lane); i < NT; i += 32) mass += exp(s[i] - top); + const float inv = 1.0f / simd_sum(mass); + float4 acc = 0.0f; + for (int i = 0; i < NT; i++) { + const device float* v = VP + (long(head) * NT + i) * 128 + c0; + acc += exp(s[i] - top) * float4(v[0], v[1], v[2], v[3]); + } + device float* out = C + (long(head) * NT + tq) * 128 + c0; + for (int j = 0; j < 4; j++) out[j] = acc[j] * inv; +} diff --git a/zig/kernels/metal/h3/h3_rows.metal b/zig/kernels/metal/h3/h3_rows.metal new file mode 100644 index 0000000000..6442da5ec8 --- /dev/null +++ b/zig/kernels/metal/h3/h3_rows.metal @@ -0,0 +1,153 @@ +// MiniMax H3 / FastH3 on Metal in the checkpoint's bf16: the row kernels around the products and the attention. +#include +#include +using namespace metal; + +#define H3_UNROLL _Pragma("clang loop unroll(full)") + +// The model's sizes come ahead of the source at load; these are FastH3's at 864x480 and 124 frames. +#ifndef H3_HEADS +#define H3_HEADS 56 +#define H3_SLOTS 18432 +#define H3_TILES 288 +#endif + +// The modulated norm one value at a time, rounded to bf16 where the reference's bf16 arithmetic rounds. +inline float h3_mod(float x, float inv, float w, float scale, float shift) { + return float(bfloat(float(bfloat(float(bfloat(x * inv * w)) * float(bfloat(1.0f + scale)))) + shift)); +} + +// Y[row0 + r, n] = B[n] + sum over k of X[r, k] W[n, k]: the latent rows into the stream. Threads [N, R, 1]. +[[kernel]] void h3_rows_in( + const device float* X [[buffer(0)]], + const device float* W [[buffer(1)]], + const device float* B [[buffer(2)]], + const constant int32_t* P [[buffer(3)]], + device bfloat* Y [[buffer(4)]], + uint3 at [[thread_position_in_grid]]) { + const int R = P[0], K = P[1], N = P[2], row0 = P[3]; + if (int(at.x) >= N || int(at.y) >= R) return; + const device float* x = X + long(at.y) * K; + const device float* w = W + long(at.x) * K; + float sum = 0.0f; + for (int k = 0; k < K; k++) sum += x[k] * w[k]; + Y[long(row0 + int(at.y)) * N + at.x] = bfloat(sum + B[at.x]); +} + +// A block's modulated RMSNorm to bf16 rows, after an optional gated add of Y into X. Threadgroups [M, 1, 1] of 32. +[[kernel]] void h3_norm( + device bfloat* X [[buffer(0)]], + const device bfloat* Y [[buffer(1)]], + const device bfloat* W [[buffer(2)]], + const device bfloat* TG [[buffer(3)]], + const device bfloat* TAB [[buffer(4)]], + const device int32_t* LINE [[buffer(5)]], + const constant int32_t* P [[buffer(6)]], + const constant float* E [[buffer(7)]], + device bfloat* O [[buffer(8)]], + uint lane [[thread_index_in_simdgroup]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int C = P[1], L = P[2], row = int(tg.x); + const int line = LINE[row]; + device bfloat* x = X + long(row) * C; + float sq = 0.0f; + if (P[3] >= 0) { + const device bfloat* g = TG + (long(P[3]) * L + line) * C; + const device bfloat* y = Y + long(row) * C; + for (int c = int(lane); c < C; c += 32) { + const bfloat v = bfloat(float(x[c]) + float(bfloat(float(g[c]) * float(y[c])))); + x[c] = v; + sq += float(v) * float(v); + } + } else { + for (int c = int(lane); c < C; c += 32) sq += float(x[c]) * float(x[c]); + } + const float inv = rsqrt(simd_sum(sq) / float(C) + E[0]); + const device bfloat* s = TAB + (long(P[4]) * L + line) * C; + const device bfloat* h = TAB + (long(P[5]) * L + line) * C; + device bfloat* o = O + long(row) * C; + for (int c = int(lane); c < C; c += 32) o[c] = bfloat(h3_mod(float(x[c]), inv, float(W[c]), float(s[c]), float(h[c]))); +} + +// X += TG[part, line] * Y, the last block's MLP branch. P: M, C, L, part. Threadgroups [M, 1, 1] of [32, 1, 1]. +[[kernel]] void h3_gate_add( + device bfloat* X [[buffer(0)]], + const device bfloat* Y [[buffer(1)]], + const device bfloat* TG [[buffer(2)]], + const device int32_t* LINE [[buffer(3)]], + const constant int32_t* P [[buffer(4)]], + uint lane [[thread_index_in_simdgroup]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int C = P[1], L = P[2], row = int(tg.x); + const device bfloat* g = TG + (long(P[3]) * L + LINE[row]) * C; + device bfloat* x = X + long(row) * C; + const device bfloat* y = Y + long(row) * C; + for (int c = int(lane); c < C; c += 32) x[c] = bfloat(float(x[c]) + float(bfloat(float(g[c]) * float(y[c])))); +} + +// The pooled branch gated into the attention output: Y[row, h] += C[h, tile of row] * G[row, h]. Threads [R, H, 1]. +[[kernel]] void h3_gate_mix( + device bfloat* Y [[buffer(0)]], + const device bfloat* G [[buffer(1)]], + const device float* C [[buffer(2)]], + const device int32_t* SLOT [[buffer(3)]], + const constant int32_t* P [[buffer(4)]], + uint3 at [[thread_position_in_grid]]) { + const int R = P[0], H = P[1], NT = P[2]; + if (int(at.x) >= R || int(at.y) >= H) return; + const long base = (long(at.x) * H + at.y) * 128; + const device float* c = C + (long(at.y) * NT + SLOT[at.x] / 64) * 128; + for (int j = 0; j < 128; j++) Y[base + j] = bfloat(float(Y[base + j]) + float(bfloat(c[j] * float(G[base + j])))); +} + +// silu(gate) * value in bf16 from the fused first projection's rows [value | gate]. Threads [N, M, 1]. +[[kernel]] void h3_swiglu( + const device bfloat* F [[buffer(0)]], + const constant int32_t* P [[buffer(1)]], + device bfloat* A [[buffer(2)]], + uint3 at [[thread_position_in_grid]]) { + const int M = P[0], N = P[1]; + if (int(at.x) >= N || int(at.y) >= M) return; + const device bfloat* f = F + long(at.y) * 2 * N; + const float g = float(f[N + at.x]), v = float(f[at.x]); + const float silu = float(bfloat(g * float(bfloat(1.0f / (1.0f + exp(-g)))))); + A[long(at.y) * N + at.x] = bfloat(silu * v); +} + +// The final layer's modulated RMSNorm. TAB: (T, 2 C) bf16, shift then scale. Y: (M, C) bf16. +[[kernel]] void h3_final_norm( + const device bfloat* X [[buffer(0)]], + const device bfloat* W [[buffer(1)]], + const device bfloat* TAB [[buffer(2)]], + const device int32_t* LINE [[buffer(3)]], + const constant int32_t* P [[buffer(4)]], + const constant float* E [[buffer(5)]], + device bfloat* Y [[buffer(6)]], + uint lane [[thread_index_in_simdgroup]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int C = P[1], row = int(tg.x); + const device bfloat* x = X + long(row) * C; + const device bfloat* t = TAB + long(LINE[row]) * 2 * C; + float sq = 0.0f; + for (int c = int(lane); c < C; c += 32) sq += float(x[c]) * float(x[c]); + const float inv = rsqrt(simd_sum(sq) / float(C) + E[0]); + for (int c = int(lane); c < C; c += 32) Y[long(row) * C + c] = bfloat(h3_mod(float(x[c]), inv, float(W[c]), float(t[C + c]), float(t[c]))); +} + +// Y[r, n] = B[n] + sum over k of X[row0 + r, k] W[n, k]: an output head over its bf16 rows, in float. +[[kernel]] void h3_rows_out( + const device bfloat* X [[buffer(0)]], + const device float* W [[buffer(1)]], + const device float* B [[buffer(2)]], + const constant int32_t* P [[buffer(3)]], + device float* Y [[buffer(4)]], + uint lane [[thread_index_in_simdgroup]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int K = P[1], N = P[2], row0 = P[3]; + const device bfloat* x = X + long(row0 + int(tg.y)) * K; + const device float* w = W + long(tg.x) * K; + float sum = 0.0f; + for (int k = int(lane); k < K; k += 32) sum += float(x[k]) * w[k]; + sum = simd_sum(sum); + if (lane == 0) Y[long(tg.y) * N + tg.x] = sum + B[tg.x]; +} diff --git a/zig/kernels/metal/h3/kernels.zig b/zig/kernels/metal/h3/kernels.zig new file mode 100644 index 0000000000..a7be157f50 --- /dev/null +++ b/zig/kernels/metal/h3/kernels.zig @@ -0,0 +1,6 @@ +//! MiniMax H3's Metal kernels: three sources joined and compiled at load, after the prefill kernels' gemm_nax. +pub const source = @embedFile("h3_rows.metal") ++ @embedFile("h3_routing.metal") ++ @embedFile("h3_attention.metal"); +/// MLX's bf16 GEMM, x w^T with N a multiple of 128 and M a multiple of 64 (whole) or not (ragged). +pub const gemm_whole = "custom_kernel_tf_gemm_nax_bf16_n_t_t_t_n_64_128_256_2_4_bfloat16_t_bfloat16_t_int32_t_bfloat16_t"; +pub const gemm_ragged = "custom_kernel_tf_gemm_nax_bf16_n_t_n_t_n_64_128_256_2_4_bfloat16_t_bfloat16_t_int32_t_bfloat16_t"; +pub const names = [_][:0]const u8{ "h3_rows_in", "h3_norm", "h3_gate_add", "h3_gate_mix", "h3_swiglu", "h3_final_norm", "h3_rows_out", "h3_heads", "h3_pool", "h3_tile_scores", "h3_topk", "h3_coarse", "h3_attention", gemm_whole, gemm_ragged }; diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index b79337b5c5..5fda07b576 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -126,6 +126,21 @@ pub const metal = [_]M{ .{ .name = "tile_qmm twin picks by expert tf_tile_qmv_group", .threads = 128, .tg_memory = 1024, .m1 = 1024, .m2 = 1024 }, .{ .name = "tile_qmm group sums tf_tile_xsum", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, .{ .name = "glm hc_expand_10 custom_kernel_tf_glm5_fused_hc_expand_cc77af98f2__4096_1_0_0_0_bfloat16_t_bfloat16_t_float_float_floatc_bfloat16_t_float_float", .threads = 1024, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "h3 rows in h3_rows_in", .threads = 1024, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "h3 modulated norm h3_norm", .threads = 32, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "h3 gated add h3_gate_add", .threads = 32, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "h3 pooled branch mix h3_gate_mix", .threads = 256, .tg_memory = 0, .m1 = 928, .m2 = 928 }, + .{ .name = "h3 SwiGLU h3_swiglu", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "h3 final norm h3_final_norm", .threads = 32, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "h3 rows out h3_rows_out", .threads = 32, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "h3 head layout h3_heads", .threads = 32, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "h3 tile means h3_pool", .threads = 32, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "h3 tile scores h3_tile_scores", .threads = 256, .tg_memory = 0, .m1 = 928, .m2 = 928 }, + .{ .name = "h3 tile top-k h3_topk", .threads = 32, .tg_memory = 288, .m1 = 1024, .m2 = 1024 }, + .{ .name = "h3 pooled branch h3_coarse", .threads = 32, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "h3 routed attention h3_attention", .threads = 256, .tg_memory = 25856, .offline = false }, + .{ .name = "h3 gemm whole custom_kernel_tf_gemm_nax_bf16_n_t_t_t_n_64_128_256_2_4_bfloat16_t_bfloat16_t_int32_t_bfloat16_t", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "h3 gemm ragged custom_kernel_tf_gemm_nax_bf16_n_t_n_t_n_64_128_256_2_4_bfloat16_t_bfloat16_t_int32_t_bfloat16_t", .threads = 256, .tg_memory = 0, .reserved = 256 }, .{ .name = "qwen vision qwv_gemm_bf16_nt_00", .threads = 256, .tg_memory = 0, .reserved = 256 }, .{ .name = "qwen vision qwv_gemm_bf16_nt_01", .threads = 256, .tg_memory = 0, .reserved = 256 }, .{ .name = "qwen vision qwv_gemm_bf16_nt_10", .threads = 256, .tg_memory = 0, .reserved = 256 }, diff --git a/zig/src/families/h3/block.zig b/zig/src/families/h3/block.zig new file mode 100644 index 0000000000..4884120569 --- /dev/null +++ b/zig/src/families/h3/block.zig @@ -0,0 +1,129 @@ +//! One FastH3 block in bf16: products on tfp::gemm_nax (MLX's GEMM, fp32 sums), FastH3's routed attention in bf16. +const std = @import("std"); +const mtl = @import("metal"); +const h3 = @import("h3.zig"); + +const kernels = @import("h3_kernels"); + +const Model = h3.Model; +const Ref = h3.Ref; + +/// d (rows, n) = x (rows, k) w^T for w (n, k) in bf16. +fn gemm(m: *const Model, enc: mtl.ComputeEncoder, x: mtl.Buffer, w: Ref, d: mtl.Buffer, n: usize, k: usize) void { + const tn = n / 128; + const tm = (m.rows + 63) / 64; + enc.setPipeline(if (m.rows % 64 == 0) m.pipeline(kernels.gemm_whole) else m.pipeline(kernels.gemm_ragged)); + enc.setBuffer(x, 0, 0); + enc.setBuffer(w.buffer, w.offset, 1); + var p: [16]i32 = @splat(0); + const values = [_]usize{ m.rows, n, k, k, k, n, tn, tm, 2, k / 256 }; + for (values, 0..) |v, i| p[i] = @intCast(v); + enc.setBytes(std.mem.asBytes(&p), 2); + enc.setBuffer(d, 0, 3); + enc.dispatchThreads(mtl.Size.of(tn * 4 * 32, (tm + 3) / 4 * 4, 2), mtl.Size.of(32, 4, 2)); +} + +/// The stream's modulated norm to bf16 rows in xn; with `gates` it first adds the branch in `y` gated by that part. +fn norm(m: *const Model, enc: mtl.ComputeEncoder, w: Ref, gates: ?Ref, gate: i32, tables: Ref, scale: i32, shift: i32) void { + const g = gates orelse tables; + enc.setPipeline(m.pipeline("h3_norm")); + enc.setBuffer(m.x, 0, 0); + enc.setBuffer(m.y, 0, 1); + enc.setBuffer(w.buffer, w.offset, 2); + enc.setBuffer(g.buffer, g.offset, 3); + enc.setBuffer(tables.buffer, tables.offset, 4); + enc.setBuffer(m.line, 0, 5); + enc.setValue([6]i32{ @intCast(m.rows), @intCast(m.hidden), @intCast(m.lines), if (gates == null) -1 else gate, scale, shift }, 6); + enc.setValue([1]f32{h3.eps}, 7); + enc.setBuffer(m.xn, 0, 8); + enc.dispatchGroups(mtl.Size.of(m.rows, 1, 1), mtl.Size.of(32, 1, 1)); +} + +fn attend(m: *const Model, enc: mtl.ComputeEncoder, list: mtl.Buffer, queries: usize, keys: usize, first: usize, per_query: bool) void { + enc.setPipeline(m.pipeline("h3_attention")); + enc.setBuffer(m.tq, 0, 0); + enc.setBuffer(m.tk, 0, 1); + enc.setBuffer(m.tv, 0, 2); + enc.setBuffer(list, 0, 3); + enc.setBuffer(m.sizes.buffer, m.sizes.offset, 4); + enc.setBuffer(m.row_of, 0, 5); + enc.setValue([4]i32{ @intCast(queries), @intCast(keys), @intCast(first), @intFromBool(per_query) }, 6); + enc.setValue([1]f32{1.0 / @sqrt(@as(f32, h3.head_dim))}, 7); + enc.setBuffer(m.qr, 0, 8); + enc.dispatchGroups(mtl.Size.of(queries, m.heads, 1), mtl.Size.of(256, 1, 1)); +} + +/// Block `index` of a forward; the block before left its MLP's rows in y, added here with this block's first norm. +pub fn run(m: *Model, pass: *h3.Pass, index: usize) void { + const b = &m.blocks[index]; + const enc = &pass.enc; + const heads = m.heads; + const gates: ?Ref = if (index == 0) null else m.blocks[index - 1].tables; + norm(m, enc.*, b.norm1, gates, h3.gate_m, b.tables, h3.scale_a, h3.shift_a); + m.lap(pass, "norms"); + gemm(m, enc.*, m.xn, b.q, m.qr, m.inner, m.hidden); + gemm(m, enc.*, m.xn, b.k, m.kr, m.inner, m.hidden); + gemm(m, enc.*, m.xn, b.v, m.vr, m.inner, m.hidden); + if (!m.dense) gemm(m, enc.*, m.xn, b.compress, m.cr, m.inner, m.hidden); + m.lap(pass, "q k v gate"); + enc.setPipeline(m.pipeline("h3_heads")); + enc.setBuffer(m.qr, 0, 0); + enc.setBuffer(m.kr, 0, 1); + enc.setBuffer(m.vr, 0, 2); + enc.setBuffer(b.norm_q.buffer, b.norm_q.offset, 3); + enc.setBuffer(b.norm_k.buffer, b.norm_k.offset, 4); + enc.setBuffer(m.cos.buffer, m.cos.offset, 5); + enc.setBuffer(m.sin.buffer, m.sin.offset, 6); + enc.setBuffer(m.slot.buffer, m.slot.offset, 7); + enc.setValue([4]i32{ @intCast(m.rows), @intCast(heads), @intCast(m.rot), 0 }, 8); + enc.setValue([1]f32{h3.eps}, 9); + enc.setBuffer(m.tq, 0, 10); + enc.setBuffer(m.tk, 0, 11); + enc.setBuffer(m.tv, 0, 12); + enc.dispatchGroups(mtl.Size.of(m.rows, heads, 1), mtl.Size.of(32, 1, 1)); + m.lap(pass, "heads"); + if (m.dense) { + attend(m, enc.*, m.all_tiles, m.tiles, m.tiles, 0, false); + m.lap(pass, "attention"); + } else { + enc.setPipeline(m.pipeline("h3_pool")); + enc.setBuffer(m.tq, 0, 0); + enc.setBuffer(m.tk, 0, 1); + enc.setBuffer(m.tv, 0, 2); + enc.setBuffer(m.sizes.buffer, m.sizes.offset, 3); + enc.setBuffer(m.qp, 0, 4); + enc.setBuffer(m.kp, 0, 5); + enc.setBuffer(m.vp, 0, 6); + enc.dispatchGroups(mtl.Size.of(m.tiles, heads, 1), mtl.Size.of(32, 1, 1)); + m.route(enc.*); + m.lap(pass, "routing"); + // the prefix's query tiles see every tile; a video tile sees the prefix and its chosen tiles + attend(m, enc.*, m.all_tiles, m.prefix_tiles, m.tiles, 0, false); + attend(m, enc.*, m.chosen, m.tiles - m.prefix_tiles, m.prefix_tiles + m.keep, m.prefix_tiles, true); + m.lap(pass, "attention"); + // the pooled branch gated into the attention rows in place + enc.setPipeline(m.pipeline("h3_gate_mix")); + enc.setBuffer(m.qr, 0, 0); + enc.setBuffer(m.cr, 0, 1); + enc.setBuffer(m.coarse, 0, 2); + enc.setBuffer(m.slot.buffer, m.slot.offset, 3); + enc.setValue([3]i32{ @intCast(m.rows), @intCast(heads), @intCast(m.tiles) }, 4); + enc.dispatchThreads(mtl.Size.of(m.rows, heads, 1), mtl.Size.of(64, 4, 1)); + m.lap(pass, "gate mix"); + } + gemm(m, enc.*, m.qr, b.out, m.y, m.hidden, m.inner); + m.lap(pass, "attention out"); + norm(m, enc.*, b.norm2, b.tables, h3.gate_a, b.tables, h3.scale_m, h3.shift_m); + m.lap(pass, "norms"); + // the first projection's rows are [value | gate], one product of twice the MLP's width + gemm(m, enc.*, m.xn, b.fc1, m.wide, 2 * m.mlp, m.hidden); + enc.barrier(); + enc.setPipeline(m.pipeline("h3_swiglu")); + enc.setBuffer(m.wide, 0, 0); + enc.setValue([2]i32{ @intCast(m.rows), @intCast(m.mlp) }, 1); + enc.setBuffer(m.act, 0, 2); + enc.dispatchThreads(mtl.Size.of(m.mlp, m.rows, 1), mtl.Size.of(256, 1, 1)); + m.lap(pass, "mlp in"); + gemm(m, enc.*, m.act, b.fc2, m.y, m.hidden, m.mlp); + m.lap(pass, "mlp out"); +} diff --git a/zig/src/families/h3/h3.zig b/zig/src/families/h3/h3.zig new file mode 100644 index 0000000000..20e211a236 --- /dev/null +++ b/zig/src/families/h3/h3.zig @@ -0,0 +1,369 @@ +//! MiniMax H3 / FastH3's transformer on Metal in the checkpoint's bf16: one denoising forward over the packed rows. +const std = @import("std"); +const mtl = @import("metal"); +const st = @import("safetensors"); +const sources = @import("kernel_sources"); +const tensors = @import("tensors.zig"); +const kernels = @import("h3_kernels"); +const block = @import("block.zig"); + +pub const Tensors = tensors.Tensors; +pub const Ref = tensors.Ref; +pub const Checkpoint = tensors.Checkpoint; +pub const head_dim = 128; +pub const eps: f32 = 1e-5; +/// Rows of one routing tile, fixed by the checkpoint's training. +const slot_tile = 64; +/// The six modulation tables of a block, in the order the reference splits them. +pub const shift_a = 0; +pub const scale_a = 1; +pub const gate_a = 2; +pub const shift_m = 3; +pub const scale_m = 4; +pub const gate_m = 5; + +/// A block's projections as the checkpoint stores them, (N, K) bf16 read in place; fc1 is [value; gate]. +pub const Block = struct { q: Ref, k: Ref, v: Ref, compress: Ref, out: Ref, fc1: Ref, fc2: Ref, norm1: Ref, norm2: Ref, norm_q: Ref, norm_k: Ref, tables: Ref }; + +/// `defines` goes ahead of the kernels' source (a probe or a knock-out under test). +pub const Options = struct { defines: []const u8 = "" }; + +pub const Timing = struct { wall: f64, gpu: f64 }; + +/// Where a profiled forward's GPU time went. +pub const stages = [_][]const u8{ "rows in", "norms", "q k v gate", "heads", "routing", "attention", "gate mix", "quantize", "attention out", "mlp in", "mlp out", "final" }; + +pub const Pass = struct { cb: mtl.CommandBuffer, enc: mtl.ComputeEncoder }; + +pub const Model = struct { + gpa: std.mem.Allocator, + device: mtl.Device, + queue: mtl.Queue, + pipelines: [kernels.names.len]mtl.Pipeline, + checkpoint: Checkpoint, + case: *Tensors, + hidden: usize, + inner: usize, + heads: usize, + mlp: usize, + video_in: usize, + audio_in: usize, + /// Packed rows: `text`, then `audio`, then `video`. + rows: usize, + text: usize, + audio: usize, + video: usize, + tiles: usize, + prefix_tiles: usize, + keep: usize, + lines: usize, + rot: usize, + blocks: []Block, + prompt: []const u16, + adaln: []const i32, + times: []const i32, + final_table: Ref, + final_norm: Ref, + w_video_in: Ref, + b_video_in: Ref, + w_audio_in: Ref, + b_audio_in: Ref, + w_video_out: Ref, + b_video_out: Ref, + w_audio_out: Ref, + b_audio_out: Ref, + cos: Ref, + sin: Ref, + slot: Ref, + sizes: Ref, + row_of: mtl.Buffer, + all_tiles: mtl.Buffer, + line: mtl.Buffer, + time: mtl.Buffer, + // the stream and a branch's rows, the normed rows, the projections at the attention width, the MLP's rows + x: mtl.Buffer, + y: mtl.Buffer, + xn: mtl.Buffer, + qr: mtl.Buffer, + kr: mtl.Buffer, + vr: mtl.Buffer, + cr: mtl.Buffer, + wide: mtl.Buffer, + act: mtl.Buffer, + // queries, keys and values in tile order, the tiles' means, scores, chosen tiles, pooled branch + tq: mtl.Buffer, + tk: mtl.Buffer, + tv: mtl.Buffer, + qp: mtl.Buffer, + kp: mtl.Buffer, + vp: mtl.Buffer, + tile_scores: mtl.Buffer, + chosen: mtl.Buffer, + coarse: mtl.Buffer, + video_latents: mtl.Buffer, + audio_latents: mtl.Buffer, + normed: mtl.Buffer, + video_velocity: mtl.Buffer, + audio_velocity: mtl.Buffer, + /// Every query tile attends to every tile and the pooled branch is left out (not how the checkpoint was trained). + dense: bool = false, + profile: bool = false, + spent: [stages.len]f64 = @splat(0), + + pub fn pipeline(self: *const Model, comptime name: []const u8) mtl.Pipeline { + inline for (kernels.names, 0..) |n, i| if (comptime std.mem.eql(u8, n, name)) return self.pipelines[i]; + @compileError("no kernel " ++ name); + } + + fn slots(self: *const Model) usize { + return self.tiles * slot_tile; + } + + /// `checkpoint_dir`: the FastH3 transformer folder of `shards` files. `case`: tools/zig/h3_case.py's output. + pub fn load(gpa: std.mem.Allocator, checkpoint_dir: []const u8, shards: usize, case: *Tensors, options: Options) !*Model { + const device = try mtl.Device.init(); + if (!device.tensorUnits()) return error.NeedsTensorUnits; + const self = try gpa.create(Model); + self.gpa = gpa; + self.device = device; + self.queue = try device.queue(); + self.case = case; + self.dense = false; + self.profile = false; + self.spent = @splat(0); + self.checkpoint = try Checkpoint.open(gpa, device, checkpoint_dir, shards); + const ck = &self.checkpoint; + const to_q = ck.entry("transformer_blocks.0.attn.to_q.weight") orelse return error.MissingTensor; + const fc2 = ck.entry("transformer_blocks.0.ff.net.2.weight") orelse return error.MissingTensor; + self.inner = to_q.shape[0]; + self.hidden = to_q.shape[1]; + self.heads = self.inner / head_dim; + self.mlp = fc2.shape[1]; + var layers: usize = 0; + var name: [128]u8 = undefined; + while (ck.entry(try std.fmt.bufPrint(&name, "transformer_blocks.{d}.norm1.weight", .{layers})) != null) layers += 1; + + const counts = (case.names.get("geometry") orelse return error.MissingTensor).shape[0]; + // a case made for a first frame has a seventh count, the keyframe's rows: not read here + if (counts != 6) return error.UnsupportedCase; + const geometry = try case.copy(i32, gpa, try case.at("geometry", .i32, counts), counts); + defer gpa.free(geometry); + self.prefix_tiles = @intCast(geometry[0]); + self.tiles = @intCast(geometry[0] + geometry[1]); + self.keep = @intCast(geometry[2]); + self.text = @intCast(geometry[3]); + self.audio = @intCast(geometry[4]); + self.video = @intCast(geometry[5]); + self.rows = self.text + self.audio + self.video; + const rows = self.rows; + const hidden = self.hidden; + const inner = self.inner; + const heads = self.heads; + self.video_in = (case.names.get("video_in.weight") orelse return error.MissingTensor).shape[1]; + self.audio_in = (case.names.get("audio_in.weight") orelse return error.MissingTensor).shape[1]; + self.lines = (case.names.get("tables.0") orelse return error.MissingTensor).shape[1]; + self.rot = (case.names.get("cos") orelse return error.MissingTensor).shape[1]; + const times = (case.names.get("final") orelse return error.MissingTensor).shape[0]; + + // the kernels take the model's sizes as compile-time values + const source = try std.fmt.allocPrint(gpa, "#define H3_HEADS {d}\n#define H3_SLOTS {d}\n#define H3_TILES {d}\n{s}\n{s}", .{ heads, self.slots(), self.tiles, options.defines, kernels.source }); + defer gpa.free(source); + // MLX's bf16 GEMM (the prefill kernels' gemm_nax) ahead of the family's own kernels; nax.h inlined + const gemm = try std.mem.replaceOwned(u8, gpa, gemmSource(), "#include \"../nax.h\"", sources.nax); + defer gpa.free(gemm); + const all = try std.mem.concat(gpa, u8, &.{ gemm, source }); + defer gpa.free(all); + const lib = try mtl.Library.fromSource(device, all, mtl.CompileOptions.mlx()); + defer lib.deinit(); + for (kernels.names, 0..) |kernel, i| self.pipelines[i] = try mtl.Pipeline.init(device, lib, kernel, false); + + self.prompt = try case.copy(u16, gpa, try case.at("text", .bf16, self.text * hidden), self.text * hidden); + self.final_table = try self.held("final", .bf16, times * 2 * hidden); + self.final_norm = try self.held("final_norm.weight", .bf16, hidden); + self.w_video_in = try self.held("video_in.weight", .f32, hidden * self.video_in); + self.b_video_in = try self.held("video_in.bias", .f32, hidden); + self.w_audio_in = try self.held("audio_in.weight", .f32, hidden * self.audio_in); + self.b_audio_in = try self.held("audio_in.bias", .f32, hidden); + self.w_video_out = try self.held("video_out.weight", .f32, hidden * self.video_in); + self.b_video_out = try self.held("video_out.bias", .f32, self.video_in); + self.w_audio_out = try self.held("audio_out.weight", .f32, hidden * self.audio_in); + self.b_audio_out = try self.held("audio_out.bias", .f32, self.audio_in); + self.cos = try self.held("cos", .f32, rows * self.rot); + self.sin = try self.held("sin", .f32, rows * self.rot); + self.slot = try self.held("tile_slot", .i32, rows); + self.sizes = try self.held("tile_sizes", .i32, self.tiles); + + const steps = (case.names.get("adaln") orelse return error.MissingTensor).shape[0]; + self.adaln = try case.copy(i32, gpa, try case.at("adaln", .i32, steps * rows), steps * rows); + self.times = try case.copy(i32, gpa, try case.at("times", .i32, steps * rows), steps * rows); + const total = self.slots(); + const shared = mtl.ResourceOptions.shared; + self.row_of = try device.buffer(total * 4, shared); + @memset(self.row_of.slice(i32, total), -1); + for (self.slot.buffer.slice(i32, rows), 0..) |at, row| self.row_of.slice(i32, total)[@intCast(at)] = @intCast(row); + self.all_tiles = try device.buffer(self.tiles * 4, shared); + for (self.all_tiles.slice(i32, self.tiles), 0..) |*t, i| t.* = @intCast(i); + self.line = try device.buffer(rows * 4, shared); + self.time = try device.buffer(rows * 4, shared); + inline for (.{ "x", "y", "xn" }) |f| @field(self, f) = try device.buffer(rows * hidden * 2, shared); + inline for (.{ "qr", "kr", "vr", "cr" }) |f| @field(self, f) = try device.buffer(rows * inner * 2, shared); + self.wide = try device.buffer(rows * 2 * self.mlp * 2, shared); + self.act = try device.buffer(rows * self.mlp * 2, shared); + // padding slots are never written and must read as zeros + const slot_bytes = heads * total * head_dim * 2; + inline for (.{ "tq", "tk", "tv" }) |f| { + @field(self, f) = try device.buffer(slot_bytes, shared); + @memset(@field(self, f).slice(u8, slot_bytes), 0); + } + inline for (.{ "qp", "kp", "vp", "coarse" }) |f| @field(self, f) = try device.buffer(heads * self.tiles * head_dim * 4, shared); + self.tile_scores = try device.buffer(heads * self.tiles * self.tiles * 4, shared); + self.chosen = try device.buffer(heads * (self.tiles - self.prefix_tiles) * (self.prefix_tiles + self.keep) * 4, shared); + self.video_latents = try device.buffer(self.video * self.video_in * 4, shared); + self.audio_latents = try device.buffer(self.audio * self.audio_in * 4, shared); + self.normed = try device.buffer(rows * hidden * 2, shared); + self.video_velocity = try device.buffer(self.video * self.video_in * 4, shared); + self.audio_velocity = try device.buffer(self.audio * self.audio_in * 4, shared); + + self.blocks = try gpa.alloc(Block, layers); + for (self.blocks, 0..) |*b, i| { + const parts = .{ .{ "q", "attn.to_q", inner, hidden }, .{ "k", "attn.to_k", inner, hidden }, .{ "v", "attn.to_v", inner, hidden }, .{ "compress", "attn.to_gate_compress", inner, hidden }, .{ "out", "attn.to_out.0", hidden, inner }, .{ "fc1", "ff.net.0.proj", 2 * self.mlp, hidden }, .{ "fc2", "ff.net.2", hidden, self.mlp } }; + inline for (parts) |p| @field(b, p[0]) = try ck.ref(try std.fmt.bufPrint(&name, "transformer_blocks.{d}." ++ p[1] ++ ".weight", .{i}), p[2] * p[3]); + b.norm1 = try ck.ref(try std.fmt.bufPrint(&name, "transformer_blocks.{d}.norm1.weight", .{i}), hidden); + b.norm2 = try ck.ref(try std.fmt.bufPrint(&name, "transformer_blocks.{d}.norm2.weight", .{i}), hidden); + b.norm_q = try ck.ref(try std.fmt.bufPrint(&name, "transformer_blocks.{d}.attn.norm_q.weight", .{i}), head_dim); + b.norm_k = try ck.ref(try std.fmt.bufPrint(&name, "transformer_blocks.{d}.attn.norm_k.weight", .{i}), head_dim); + b.tables = try self.held(try std.fmt.bufPrint(&name, "tables.{d}", .{i}), .bf16, 6 * self.lines * hidden); + } + return self; + } + + /// The case's tensor `name` in a buffer of its own (the case file's tensors start wherever its header ends). + fn held(self: *const Model, name: []const u8, dtype: st.DType, count: usize) !Ref { + const from = try self.case.at(name, dtype, count); + const bytes = count * dtype.size(); + const buffer = try self.device.buffer(bytes, mtl.ResourceOptions.shared); + @memcpy(buffer.slice(u8, bytes), self.case.map.bytes[from.offset..][0..bytes]); + return .{ .buffer = buffer }; + } + + pub fn lap(self: *Model, pass: *Pass, comptime stage: []const u8) void { + pass.enc.barrier(); + if (!self.profile) return; + pass.enc.end(); + pass.cb.commit(); + pass.cb.wait(); + inline for (stages, 0..) |s, i| if (comptime std.mem.eql(u8, s, stage)) { + self.spent[i] += pass.cb.gpuSeconds(); + }; + pass.cb = self.queue.commandBuffer(); + pass.enc = pass.cb.compute(.concurrent); + } + + /// After the tiles' means: tile scores, each video tile's chosen key tiles and the pooled branch. + pub fn route(self: *const Model, enc: mtl.ComputeEncoder) void { + const heads = self.heads; + enc.barrier(); + enc.setPipeline(self.pipeline("h3_tile_scores")); + enc.setBuffer(self.qp, 0, 0); + enc.setBuffer(self.kp, 0, 1); + enc.setValue([2]i32{ @intCast(self.tiles), @intCast(heads) }, 2); + enc.setBuffer(self.tile_scores, 0, 3); + enc.dispatchThreads(mtl.Size.of(self.tiles, self.tiles, heads), mtl.Size.of(32, 8, 1)); + enc.barrier(); + enc.setPipeline(self.pipeline("h3_topk")); + enc.setBuffer(self.tile_scores, 0, 0); + enc.setValue([3]i32{ @intCast(self.tiles), @intCast(self.prefix_tiles), @intCast(self.keep) }, 1); + enc.setBuffer(self.chosen, 0, 2); + enc.dispatchGroups(mtl.Size.of(self.tiles - self.prefix_tiles, heads, 1), mtl.Size.of(32, 1, 1)); + enc.setPipeline(self.pipeline("h3_coarse")); + enc.setBuffer(self.tile_scores, 0, 0); + enc.setBuffer(self.vp, 0, 1); + enc.setValue([1]i32{@intCast(self.tiles)}, 2); + enc.setBuffer(self.coarse, 0, 3); + enc.dispatchGroups(mtl.Size.of(self.tiles, heads, 1), mtl.Size.of(32, 1, 1)); + } + + fn rowsIn(self: *const Model, enc: mtl.ComputeEncoder, from: mtl.Buffer, w: Ref, b: Ref, count: usize, width: usize, first: usize) void { + enc.setPipeline(self.pipeline("h3_rows_in")); + enc.setBuffer(from, 0, 0); + enc.setBuffer(w.buffer, w.offset, 1); + enc.setBuffer(b.buffer, b.offset, 2); + enc.setValue([4]i32{ @intCast(count), @intCast(width), @intCast(self.hidden), @intCast(first) }, 3); + enc.setBuffer(self.x, 0, 4); + enc.dispatchThreads(mtl.Size.of(self.hidden, count, 1), mtl.Size.of(64, 16, 1)); + } + + fn rowsOut(self: *const Model, enc: mtl.ComputeEncoder, w: Ref, b: Ref, to: mtl.Buffer, count: usize, width: usize, first: usize) void { + enc.setPipeline(self.pipeline("h3_rows_out")); + enc.setBuffer(self.normed, 0, 0); + enc.setBuffer(w.buffer, w.offset, 1); + enc.setBuffer(b.buffer, b.offset, 2); + enc.setValue([4]i32{ @intCast(count), @intCast(self.hidden), @intCast(width), @intCast(first) }, 3); + enc.setBuffer(to, 0, 4); + enc.dispatchGroups(mtl.Size.of(width, count, 1), mtl.Size.of(32, 1, 1)); + } + + /// The video and audio velocities for the latent rows at denoising step `step` of the case's schedule. + pub fn forward(self: *Model, video: []const f32, audio: []const f32, step: usize, video_velocity: []f32, audio_velocity: []f32) !Timing { + const started = mtl.clock.seconds(); + const rows = self.rows; + const hidden = self.hidden; + // each row's line in the blocks' tables and in the final layer's, for this step + @memcpy(self.line.slice(i32, rows), self.adaln[step * rows ..][0..rows]); + @memcpy(self.time.slice(i32, rows), self.times[step * rows ..][0..rows]); + @memcpy(self.x.slice(u16, self.prompt.len), self.prompt); + @memcpy(self.video_latents.slice(f32, video.len), video); + @memcpy(self.audio_latents.slice(f32, audio.len), audio); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + var pass: Pass = undefined; + pass.cb = self.queue.commandBuffer(); + pass.enc = pass.cb.compute(.concurrent); + var gpu: f64 = 0; + const enc = &pass.enc; + const media = self.text; + self.rowsIn(enc.*, self.audio_latents, self.w_audio_in, self.b_audio_in, self.audio, self.audio_in, media); + self.rowsIn(enc.*, self.video_latents, self.w_video_in, self.b_video_in, self.video, self.video_in, media + self.audio); + self.lap(&pass, "rows in"); + for (0..self.blocks.len) |index| block.run(self, &pass, index); + const last = self.blocks[self.blocks.len - 1].tables; + enc.setPipeline(self.pipeline("h3_gate_add")); + enc.setBuffer(self.x, 0, 0); + enc.setBuffer(self.y, 0, 1); + enc.setBuffer(last.buffer, last.offset, 2); + enc.setBuffer(self.line, 0, 3); + enc.setValue([4]i32{ @intCast(rows), @intCast(hidden), @intCast(self.lines), gate_m }, 4); + enc.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(32, 1, 1)); + enc.barrier(); + enc.setPipeline(self.pipeline("h3_final_norm")); + enc.setBuffer(self.x, 0, 0); + enc.setBuffer(self.final_norm.buffer, self.final_norm.offset, 1); + enc.setBuffer(self.final_table.buffer, self.final_table.offset, 2); + enc.setBuffer(self.time, 0, 3); + enc.setValue([2]i32{ @intCast(rows), @intCast(hidden) }, 4); + enc.setValue([1]f32{eps}, 5); + enc.setBuffer(self.normed, 0, 6); + enc.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(32, 1, 1)); + enc.barrier(); + self.rowsOut(enc.*, self.w_audio_out, self.b_audio_out, self.audio_velocity, self.audio, self.audio_in, media); + self.rowsOut(enc.*, self.w_video_out, self.b_video_out, self.video_velocity, self.video, self.video_in, media + self.audio); + pass.enc.end(); + pass.cb.commit(); + pass.cb.wait(); + gpu += pass.cb.gpuSeconds(); + if (self.profile) self.spent[stages.len - 1] += pass.cb.gpuSeconds(); + if (pass.cb.failure()) |text| { + std.debug.print("command buffer failed: {s}\n", .{text}); + return error.GpuFailed; + } + @memcpy(video_velocity, self.video_velocity.slice(f32, video_velocity.len)); + @memcpy(audio_velocity, self.audio_velocity.slice(f32, audio_velocity.len)); + return .{ .wall = mtl.clock.seconds() - started, .gpu = gpu }; + } +}; + +/// The prefill kernels' gemm_nax source, which the block's products use. +fn gemmSource() []const u8 { + inline for (sources.prefill) |f| if (comptime std.mem.eql(u8, f.name, "gemm_nax")) return f.text; + @compileError("no gemm_nax source"); +} diff --git a/zig/src/families/h3/tensors.zig b/zig/src/families/h3/tensors.zig new file mode 100644 index 0000000000..2763949de0 --- /dev/null +++ b/zig/src/families/h3/tensors.zig @@ -0,0 +1,97 @@ +//! Safetensors files mapped once and handed to Metal whole, and a checkpoint's shards searched by tensor name. +const std = @import("std"); +const mtl = @import("metal"); +const st = @import("safetensors"); + +/// A safetensors file mapped once and handed to Metal whole; a tensor is an offset into that one buffer. +pub const Tensors = struct { + map: mtl.MappedFile, + buffer: mtl.Buffer, + data: usize, + names: st.Header, + arena: std.heap.ArenaAllocator, + + pub fn open(gpa: std.mem.Allocator, device: mtl.Device, path: [*:0]const u8) !Tensors { + const map = try mtl.MappedFile.open(path); + errdefer map.deinit(); + if (map.size < 8) return error.BadSafetensors; + const header_len: usize = @intCast(std.mem.readInt(u64, map.bytes[0..8], .little)); + if (header_len > map.size - 8) return error.BadSafetensors; + var arena = std.heap.ArenaAllocator.init(gpa); + errdefer arena.deinit(); + const names = try st.parseHeader(arena.allocator(), map.bytes[8..][0..header_len], map.size - 8 - header_len); + const buffer = try device.bufferNoCopy(map.bytes.ptr, map.bytes.len, mtl.ResourceOptions.shared); + return .{ .map = map, .buffer = buffer, .data = 8 + header_len, .names = names, .arena = arena }; + } + + pub fn deinit(self: *Tensors) void { + self.buffer.deinit(); + self.arena.deinit(); + self.map.deinit(); + } + + /// The tensor `name` of `dtype` and `count` values: where it starts in the buffer. + pub fn at(self: *const Tensors, name: []const u8, dtype: st.DType, count: usize) !Ref { + const e = self.names.get(name) orelse { + std.debug.print("no tensor {s}\n", .{name}); + return error.MissingTensor; + }; + const offset = self.data + e.begin; + if (e.dtype != dtype or e.end - e.begin != count * dtype.size()) { + std.debug.print("tensor {s}: unexpected type or size\n", .{name}); + return error.BadTensor; + } + return .{ .buffer = self.buffer, .offset = offset }; + } + + /// `at` for a tensor a kernel reads in place, which must start on a multiple of its value size. + pub fn inPlace(self: *const Tensors, name: []const u8, dtype: st.DType, count: usize) !Ref { + const ref = try self.at(name, dtype, count); + if (ref.offset % dtype.size() != 0) { + std.debug.print("tensor {s} is not aligned for the GPU\n", .{name}); + return error.BadTensor; + } + return ref; + } + + pub fn values(self: *const Tensors, comptime T: type, ref: Ref, count: usize) []const T { + return @as([*]const T, @ptrCast(@alignCast(self.map.bytes.ptr + ref.offset)))[0..count]; + } + + /// The tensor's values copied out, wherever it starts in the file. + pub fn copy(self: *const Tensors, comptime T: type, gpa: std.mem.Allocator, ref: Ref, count: usize) ![]T { + const out = try gpa.alloc(T, count); + @memcpy(std.mem.sliceAsBytes(out), self.map.bytes[ref.offset..][0 .. count * @sizeOf(T)]); + return out; + } +}; + +pub const Ref = struct { buffer: mtl.Buffer, offset: usize = 0 }; + +/// The checkpoint's shards, each mapped once; a tensor is found in whichever holds it. +pub const Checkpoint = struct { + shards: []Tensors, + + pub fn open(gpa: std.mem.Allocator, device: mtl.Device, dir: []const u8, count: usize) !Checkpoint { + const shards = try gpa.alloc(Tensors, count); + var path: [1024]u8 = undefined; + for (shards, 1..) |*shard, i| { + const text = try std.fmt.bufPrint(path[0 .. path.len - 1], "{s}/diffusion_pytorch_model-{d:0>5}-of-{d:0>5}.safetensors", .{ dir, i, count }); + path[text.len] = 0; + shard.* = try Tensors.open(gpa, device, @ptrCast(text.ptr)); + } + return .{ .shards = shards }; + } + + pub fn entry(self: *const Checkpoint, name: []const u8) ?st.Entry { + for (self.shards) |*shard| if (shard.names.get(name)) |e| return e; + return null; + } + + /// The bf16 tensor `name` of `count` values, read in place. + pub fn ref(self: *const Checkpoint, name: []const u8, count: usize) !Ref { + for (self.shards) |*shard| if (shard.names.get(name) != null) return shard.inPlace(name, .bf16, count); + std.debug.print("no tensor {s}\n", .{name}); + return error.MissingTensor; + } +}; diff --git a/zig/tests/contract/metal_h3.zig b/zig/tests/contract/metal_h3.zig new file mode 100644 index 0000000000..4d7e1144cf --- /dev/null +++ b/zig/tests/contract/metal_h3.zig @@ -0,0 +1,211 @@ +//! MiniMax H3's bf16 Metal kernels under the contract, each against its own formula in float64 (metal_h3_ref.zig). +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const h3k = @import("h3_kernels"); +const rig_ = @import("rig.zig"); +const ref = @import("metal_h3_ref.zig"); + +const c = contract.spec; +const g = contract.gates; + +pub const Which = enum { rows_in, rows_out, norm, final_norm, gate_add, gate_mix, swiglu, heads, pool, tile_scores, coarse, attention }; + +/// The geometry the test kernels are compiled for: two heads, up to 16 tiles of 64 slots. +const heads = 2; +const tiles = 16; +const slots = tiles * 64; +const defines = std.fmt.comptimePrint("#define H3_HEADS {d}\n#define H3_SLOTS {d}\n#define H3_TILES {d}\n", .{ heads, slots, tiles }); + +fn kChain(s: c.Shape) usize { + return s.k + 2; +} +/// A lane's chain over every 32nd input, then one simd_sum tree. +fn laneChain(s: c.Shape) usize { + return (s.k + 31) / 32 + 6; +} +fn headChain(_: c.Shape) usize { + return 4 + 6; +} +fn poolChain(_: c.Shape) usize { + return 64; +} +fn tileChain(s: c.Shape) usize { + return s.rows + 2; +} +/// The tensor op's dot over 128 channels, then the value product's sum over every selected key. +fn attnChain(s: c.Shape) usize { + return 128 + s.k + 4; +} +fn none(_: c.Shape) usize { + return 0; +} + +const soft_why = "softmax weights in [0, 1], rounded to bf16 before the value product: the bf16 output hides the sums' error"; +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"; + +pub fn spec(w: Which) c.Spec { + const norm_sums = c.Sums{ .bf16 = 0, .f16 = 0, .exact = false, .why = norm_why }; + return switch (w) { + .rows_in => .{ .name = "metal h3 rows in h3_rows_in", .accum = .f32, .depth = kChain, .ops = 1, .out = .bf16, .tiles = .{ .{ .size = 16 }, .{ .size = 64 }, .{} }, .served = &.{ .{ .rows = 64, .n = 5376, .k = 96 }, .{ .rows = 16, .n = 5376, .k = 32 } } }, + .rows_out => .{ .name = "metal h3 rows out h3_rows_out", .accum = .f32, .depth = laneChain, .ops = 1, .out = .f32, .tiles = .{ .{}, .{}, .{ .size = 32 } }, .served = &.{ .{ .rows = 32, .n = 96, .k = 5376 }, .{ .rows = 16, .n = 32, .k = 5376 } } }, + .norm => .{ .name = "metal h3 modulated norm h3_norm", .accum = .f32, .depth = laneChain, .ops = 8, .mids = 5, .out = .bf16, .in_place = true, .tiles = .{ .{}, .{}, .{ .size = 32 } }, .served = &.{.{ .rows = 16, .n = 5376, .k = 5376 }}, .sums = norm_sums }, + .final_norm => .{ .name = "metal h3 final norm h3_final_norm", .accum = .f32, .depth = laneChain, .ops = 8, .mids = 3, .out = .bf16, .tiles = .{ .{}, .{}, .{ .size = 32 } }, .served = &.{.{ .rows = 16, .n = 5376, .k = 5376 }}, .sums = norm_sums }, + .gate_add => .{ .name = "metal h3 gated add h3_gate_add", .accum = .none, .depth = none, .mids = 1, .out = .bf16, .in_place = true, .served = &.{.{ .rows = 16, .n = 5376, .k = 1 }} }, + .gate_mix => .{ .name = "metal h3 pooled branch mix h3_gate_mix", .accum = .none, .depth = none, .mids = 1, .out = .bf16, .in_place = true, .tiles = .{ .{ .size = 64 }, .{}, .{} }, .served = &.{.{ .rows = slots, .n = heads * 128, .k = 1 }} }, + .swiglu => .{ .name = "metal h3 SwiGLU h3_swiglu", .accum = .none, .depth = none, .ops = 4, .mids = 2, .out = .bf16, .tiles = .{ .{}, .{ .size = 256 }, .{} }, .served = &.{.{ .rows = 16, .n = 14336, .k = 1 }} }, + .heads => .{ .name = "metal h3 head layout h3_heads", .accum = .f32, .depth = headChain, .ops = 8, .mids = 2, .out = .bf16, .served = &.{ .{ .rows = slots, .n = heads * 128, .k = 64 }, .{ .rows = slots, .n = heads * 128, .k = 128 } }, .sums = norm_sums }, + .pool => .{ .name = "metal h3 tile means h3_pool", .accum = .f32, .depth = poolChain, .ops = 1, .out = .f32, .served = &.{.{ .rows = tiles, .n = 128, .k = 64 }}, .sums = .{ .bf16 = 0, .f16 = 0, .exact = false, .why = "a mean of at most 64 bf16 values: no sum reaches past fp16's range" } }, + .tile_scores => .{ .name = "metal h3 tile scores h3_tile_scores", .accum = .f32, .depth = kChain, .ops = 1, .out = .f32, .served = &.{.{ .rows = tiles, .n = tiles, .k = 128 }} }, + .coarse => .{ .name = "metal h3 pooled branch h3_coarse", .accum = .f32, .depth = tileChain, .ops = 6, .out = .f32, .tiles = .{ .{ .size = 4 }, .{}, .{} }, .served = &.{.{ .rows = tiles, .n = 128, .k = tiles }}, .sums = .{ .bf16 = 0, .f16 = 0, .exact = false, .why = soft_why } }, + .attention => .{ .name = "metal h3 routed attention h3_attention", .accum = .f32, .depth = attnChain, .ops = 8, .mids = 1, .out = .bf16, .tiles = .{ .{}, .{}, .{ .size = 256, .step = 64, .min = 64 } }, .served = &.{ .{ .rows = tiles, .n = 128, .k = tiles * 64 }, .{ .rows = tiles, .n = 128, .k = 6 * 64 } }, .sums = .{ .bf16 = 0, .f16 = 0, .exact = false, .why = soft_why } }, + }; +} + +/// The kernel `w` launches: its function name in h3_kernels' sources. +fn function(w: Which) []const u8 { + return switch (w) { + .rows_in => "h3_rows_in", + .rows_out => "h3_rows_out", + .norm => "h3_norm", + .final_norm => "h3_final_norm", + .gate_add => "h3_gate_add", + .gate_mix => "h3_gate_mix", + .swiglu => "h3_swiglu", + .heads => "h3_heads", + .pool => "h3_pool", + .tile_scores => "h3_tile_scores", + .coarse => "h3_coarse", + .attention => "h3_attention", + }; +} + +const Slot = struct { ref: rig_.Ref, len: usize }; + +pub const H3 = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + which: Which, + pipe: mtl.Pipeline, + shape: c.Shape = undefined, + arena: std.heap.ArenaAllocator, + in: ref.In = .{ .rows = 0, .n = 0, .k = 0 }, + bufs: [12]rig_.Ref = undefined, + nbufs: usize = 0, + /// The in-place buffer's starting bytes, copied in before every launch. + fresh: []const u8 = &.{}, + outs: usize = 0, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, which: Which) !H3 { + const src = try std.mem.concat(a, u8, &.{ defines, h3k.source }); + defer a.free(src); + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .which = which, .pipe = try rig_.pipeline(device, src, function(which)), .arena = std.heap.ArenaAllocator.init(a) }; + } + + pub fn deinit(self: *H3) void { + self.rig.deinit(); + self.pipe.deinit(); + self.arena.deinit(); + } + + pub fn kernel(self: *H3) g.Kernel { + return .{ .spec = spec(self.which), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(H3); + + /// `n` values uniform in [lo, hi), bf16-valued when `half`, some planted non-finite when `plant`. + pub fn values(self: *H3, n: usize, lo: f64, hi: f64, fill: g.Fill, rng: std.Random, half: bool, plant: bool) ![]f64 { + const v = try self.arena.allocator().alloc(f64, n); + for (v, 0..) |*x, i| { + const r = lo + (hi - lo) * rng.float(f64); + x.* = switch (fill) { + .stress => @floatFromInt(@as(i64, @intFromFloat(@round(r * 8)))), + else => r, + }; + if (i % 61 == 7 and fill != .stress) x.* *= 6; + x.* = if (half) ref.bf(x.*) else @as(f32, @floatCast(x.*)); + } + if (fill == .nonfinite and plant) { + const at = [_]usize{ n / 3, n / 2, n - 1 }; + const put = [_]f64{ std.math.nan(f64), std.math.inf(f64), -std.math.inf(f64) }; + for (at, put) |p, x| v[p] = x; + } + return v; + } + + /// Upload `v` as bf16 (half) or fp32, as an input. + pub fn upload(self: *H3, v: []const f64, half: bool) !void { + const al = self.arena.allocator(); + const bytes = if (half) blk: { + const h = try al.alloc(u16, v.len); + for (h, v) |*o, x| o.* = c.toBf16(@floatCast(x)); + break :blk std.mem.sliceAsBytes(h); + } else blk: { + const f = try al.alloc(f32, v.len); + for (f, v) |*o, x| o.* = @floatCast(x); + break :blk std.mem.sliceAsBytes(f); + }; + self.bufs[self.nbufs] = try self.rig.input(bytes); + self.nbufs += 1; + } + + pub fn ints(self: *H3, v: []const i32) !void { + self.bufs[self.nbufs] = try self.rig.input(std.mem.sliceAsBytes(v)); + self.nbufs += 1; + } + + pub fn output(self: *H3, bytes: usize) !void { + self.bufs[self.nbufs] = try self.rig.output(bytes); + self.nbufs += 1; + } + + pub fn prepare(self: *H3, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + self.rig.reset(); + _ = self.arena.reset(.retain_capacity); + self.nbufs = 0; + self.fresh = &.{}; + self.shape = shape; + var prng = std.Random.DefaultPrng.init(seed); + return @import("metal_h3_prep.zig").prepare(self, shape, fill, prng.random()); + } + + pub fn run(self: *H3) !void { + if (self.fresh.len > 0) @memcpy(self.rig.bytes(self.bufs[0], self.fresh.len), self.fresh); + const p = self.rig.begin(); + try @import("metal_h3_prep.zig").launch(self, p.e); + try self.rig.finish(p); + } + + pub fn count(self: *H3) usize { + return self.outs; + } + + fn compute(self: *H3, s: ref.Sum, want: []f64, mag: []f64) !void { + const scratch = try self.a.alloc(f64, @max(2 * self.in.n, 1024) + self.in.k); + defer self.a.free(scratch); + switch (self.which) { + .rows_in, .rows_out => ref.dot(self.in, s, scratch, want, mag), + .norm => ref.norm(self.in, s, scratch, want, mag), + .final_norm => ref.finalNorm(self.in, s, scratch, want, mag), + .gate_add => ref.gated(self.in, false, want, mag), + .gate_mix => ref.gated(self.in, true, want, mag), + .swiglu => ref.swiglu(self.in, want, mag), + .heads => ref.heads(self.in, s, scratch, want, mag), + .pool => ref.pool(self.in, s, scratch, want, mag), + .tile_scores => ref.scores(self.in, s, scratch, want, mag), + .coarse => ref.coarse(self.in, s, scratch, want, mag), + .attention => ref.attention(self.in, s, scratch, want, mag), + } + } + + pub fn reference(self: *H3, want: []f64, mag: []f64) !void { + return self.compute(.{}, want, mag); + } + + pub fn simulate(self: *H3, acc: c.Accum, order: g.Order, out: []f64) !void { + const mag = try self.a.alloc(f64, out.len); + defer self.a.free(mag); + return self.compute(.{ .acc = acc, .order = order }, out, mag); + } +}; diff --git a/zig/tests/contract/metal_h3_prep.zig b/zig/tests/contract/metal_h3_prep.zig new file mode 100644 index 0000000000..aa094ae02a --- /dev/null +++ b/zig/tests/contract/metal_h3_prep.zig @@ -0,0 +1,311 @@ +//! The H3 contract test's inputs for each kernel and its launch, bound as the family binds it (families/h3). +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const t = @import("metal_h3.zig"); +const ref = @import("metal_h3_ref.zig"); + +const c = contract.spec; +const g = contract.gates; +const H3 = t.H3; +const heads = 2; +const slots = 1024; + +/// Gate 4's dot rows: integer terms whose second half cancels the first, so fp32 sums are exact and the result small. +fn mirror(x: []f64, w: []f64, k: usize, rows: usize, n: usize) void { + const h = k / 2; + for (0..rows) |r| for (0..h) |i| { + x[r * k + h + i] = x[r * k + i]; + }; + for (0..n) |o| for (0..h) |i| { + w[o * k + h + i] = -w[o * k + i] + @as(f64, if (i < 4) 1 else 0); + }; +} + +/// Tile sizes: most whole, some short (the last rows of a segment), never empty. +fn sizes(self: *H3, n: usize, rng: std.Random) ![]i32 { + const v = try self.arena.allocator().alloc(i32, n); + for (v) |*s| s.* = if (rng.uintLessThan(u32, 4) == 0) @intCast(1 + rng.uintLessThan(u32, 64)) else 64; + return v; +} + +pub fn prepare(self: *H3, s: c.Shape, fill: g.Fill, rng: std.Random) !bool { + const al = self.arena.allocator(); + const stress = fill == .stress; + self.in = .{ .rows = s.rows, .n = s.n, .k = s.k, .heads = heads, .slots = slots }; + switch (self.which) { + .rows_in, .rows_out => { + const lo: f64 = if (stress) -6 else -1; + const xr = s.rows; + const half = self.which == .rows_out; + const x = try self.values(xr * s.k, lo, -lo, fill, rng, half, true); + const w = try self.values(s.n * s.k, lo, -lo, fill, rng, false, false); + if (stress) mirror(x, w, s.k, xr, s.n); + const b = try self.values(s.n, -0.5, 0.5, if (stress) .random else fill, rng, false, false); + if (stress) @memset(b, 0); + self.in.a = x; + self.in.b = w; + self.in.d = b; + try self.upload(x, half); + try self.upload(w, false); + try self.upload(b, false); + try self.output(s.rows * s.n * @as(usize, if (half) 4 else 2)); + }, + .norm, .final_norm, .gate_add => { + const C = s.n; + if (self.which != .gate_add and s.k != C) return false; + // terms of one sign: no cancellation hides an intermediate rounding from the bound + const x = try self.values(s.rows * C, 0, 2, fill, rng, true, true); + const y = try self.values(s.rows * C, 0, 2, .random, rng, true, false); + const w = try self.values(C, 0.5, 1.5, .random, rng, true, false); + // the six modulation tables of one line: gate, scale and shift rows (the kernel reads them by part) + const tab = try self.values(6 * C, 0, 0.5, .random, rng, true, false); + self.in.a = x; + self.in.b = y; + self.in.d = w; + self.in.e = tab[2 * C .. 3 * C]; + self.in.f = tab[C .. 2 * C]; + self.in.g = tab[0..C]; + if (self.which == .final_norm) { + // the final table is (shift, scale) a line + self.in.g = tab[0..C]; + self.in.f = tab[C .. 2 * C]; + } + const line = try al.alloc(i32, s.rows); + @memset(line, 0); + const fresh = try al.alloc(u16, x.len); + for (fresh, x) |*o, v| o.* = c.toBf16(@floatCast(v)); + if (self.which == .final_norm) try self.upload(x, true) else { + try self.output(x.len * 2); + self.fresh = std.mem.sliceAsBytes(fresh); + } + try self.upload(y, true); + try self.upload(w, true); + try self.upload(tab, true); + try self.ints(line); + if (self.which == .norm) try self.output(x.len * 2); + if (self.which == .final_norm) try self.output(x.len * 2); + }, + .gate_mix => { + const R = s.rows; + const y = try self.values(R * s.n, 0, 2, fill, rng, true, true); + const gte = try self.values(R * s.n, 0, 2, .random, rng, true, false); + const cc = try self.values(heads * (slots / 64) * 128, 0, 1, .random, rng, false, false); + const slot = try al.alloc(i32, R); + for (slot, 0..) |*o, i| o.* = @intCast(i); + rng.shuffle(i32, slot); + self.in.a = y; + self.in.b = gte; + self.in.d = cc; + self.in.ints = slot; + self.in.slots = slots / 64; + const fresh = try al.alloc(u16, y.len); + for (fresh, y) |*o, v| o.* = c.toBf16(@floatCast(v)); + try self.output(y.len * 2); + self.fresh = std.mem.sliceAsBytes(fresh); + try self.upload(gte, true); + try self.upload(cc, false); + try self.ints(slot); + }, + .swiglu => { + const f = try self.values(s.rows * 2 * s.n, -4, 4, fill, rng, true, true); + self.in.a = f; + try self.upload(f, true); + try self.output(s.rows * s.n * 2); + }, + .heads => { + // every slot written, so the outputs hold no poison: all the slots' rows whatever the shape asks + const R: usize = slots; + const n = R * heads * 128; + const q = try self.values(n, 0, 2, fill, rng, true, true); + const k = try self.values(n, 0, 2, .random, rng, true, false); + const v = try self.values(n, -2, 2, .random, rng, true, false); + const w = try self.values(256, 0.5, 1.5, .random, rng, true, false); + const cs = try self.values(R * s.k, 0, 1, .random, rng, false, false); + const sn = try self.values(R * s.k, 0, 1, .random, rng, false, false); + // rotate-half subtracts below the half: a negative sine there keeps every rotated term of one sign + for (0..R) |r| for (0..s.k / 2) |i| { + sn[r * s.k + i] = -sn[r * s.k + i]; + }; + const slot = try al.alloc(i32, R); + for (slot, 0..) |*o, i| o.* = @intCast(i); + rng.shuffle(i32, slot); + self.in = .{ .a = q, .b = k, .d = v, .e = cs, .f = sn, .g = w, .ints = slot, .rows = R, .n = s.n, .k = s.k, .heads = heads, .slots = slots }; + for ([_][]const f64{ q, k, v }) |x| try self.upload(x, true); + try self.upload(w[0..128], true); + try self.upload(w[128..], true); + try self.upload(cs, false); + try self.upload(sn, false); + try self.ints(slot); + for (0..3) |_| try self.output(heads * slots * 128 * 2); + }, + .pool => { + const n = heads * slots * 128; + const lo: f64 = if (stress) -16 else -2; + const q = try self.values(n, lo, -lo, fill, rng, true, true); + const k = try self.values(n, lo, -lo, .random, rng, true, false); + const v = try self.values(n, lo, -lo, .random, rng, true, false); + const sz = try sizes(self, s.rows, rng); + self.in = .{ .a = q, .b = k, .d = v, .ints = sz, .rows = s.rows, .n = s.n, .k = s.k, .heads = heads, .slots = slots }; + for ([_][]const f64{ q, k, v }) |x| try self.upload(x, true); + try self.ints(sz); + for (0..3) |_| try self.output(heads * s.rows * 128 * 4); + }, + .tile_scores => { + const lo: f64 = if (stress) -6 else -1; + const qp = try self.values(heads * s.rows * 128, lo, -lo, fill, rng, false, true); + const kp = try self.values(heads * s.rows * 128, lo, -lo, fill, rng, false, false); + if (stress) mirror(qp, kp, 128, heads * s.rows, heads * s.rows); + self.in.a = qp; + self.in.b = kp; + try self.upload(qp, false); + try self.upload(kp, false); + try self.output(heads * s.rows * s.rows * 4); + }, + .coarse => { + const sc = try self.values(heads * s.rows * s.rows, -4, 4, fill, rng, false, true); + const vp = try self.values(heads * s.rows * 128, -2, 2, .random, rng, false, false); + self.in.a = sc; + self.in.b = vp; + try self.upload(sc, false); + try self.upload(vp, false); + try self.output(heads * s.rows * 128 * 4); + }, + .attention => { + // `rows` query tiles over all 16 tiles of keys; each query tile attends `k` / 64 of them + const T = s.rows; + const ksel = s.k / 64; + if (s.k % 64 != 0 or ksel > 16 or T > 16) return false; + const n = heads * slots * 128; + const q = try self.values(n, -2, 2, fill, rng, true, true); + const k = try self.values(n, -2, 2, .random, rng, true, false); + // values of one sign: a weighted mean, so a weight's bf16 rounding is relative to the output + const v = try self.values(n, 0, 2, .random, rng, true, false); + const sz = try sizes(self, 16, rng); + // each (head, query tile)'s key tiles: ksel distinct ones in tile order + const list = try al.alloc(i32, heads * T * ksel); + const pick = try al.alloc(i32, 16); + for (0..heads * T) |row| { + for (pick, 0..) |*p, i| p.* = @intCast(i); + rng.shuffle(i32, pick); + std.mem.sort(i32, pick[0..ksel], {}, std.sort.asc(i32)); + @memcpy(list[row * ksel ..][0..ksel], pick[0..ksel]); + } + const row_of = try al.alloc(i32, slots); + @memset(row_of, -1); + var rows: usize = 0; + for (sz[0..T], 0..) |size, tile| for (0..@intCast(size)) |r| { + row_of[tile * 64 + r] = @intCast(rows); + rows += 1; + }; + self.in = .{ .a = q, .b = k, .d = v, .ints = list, .more = sz, .rows = T, .n = s.n, .k = s.k, .heads = heads, .slots = slots, .scale = 1.0 / @sqrt(128.0) }; + for ([_][]const f64{ q, k, v }) |x| try self.upload(x, true); + try self.ints(list); + try self.ints(sz); + try self.ints(row_of); + try self.output(rows * heads * 128 * 2); + self.outs = rows * heads * 128; + return true; + }, + } + self.outs = switch (self.which) { + .rows_in, .rows_out, .swiglu => s.rows * s.n, + .norm => 2 * s.rows * s.n, + .final_norm, .gate_add, .gate_mix => s.rows * s.n, + .heads => 3 * heads * slots * 128, + .pool => 3 * heads * s.rows * 128, + .tile_scores => heads * s.rows * s.rows, + .coarse => heads * s.rows * 128, + .attention => unreachable, + }; + return true; +} + +fn bind(self: *H3, e: mtl.ComputeEncoder, i: usize, at: usize) void { + e.setBuffer(self.bufs[i].buf, self.bufs[i].off, at); +} + +pub fn launch(self: *H3, e: mtl.ComputeEncoder) !void { + const s = self.shape; + const eps = [1]f32{1e-5}; + e.setPipeline(self.pipe); + switch (self.which) { + .rows_in, .rows_out => { + for (0..3) |i| bind(self, e, i, i); + e.setValue([4]i32{ @intCast(s.rows), @intCast(s.k), @intCast(s.n), 0 }, 3); + bind(self, e, 3, 4); + if (self.which == .rows_in) e.dispatchThreads(mtl.Size.of(s.n, s.rows, 1), mtl.Size.of(64, 16, 1)) else e.dispatchGroups(mtl.Size.of(s.n, s.rows, 1), mtl.Size.of(32, 1, 1)); + }, + .norm => { + // X (in place), Y, W, the gate table, the tables, LINE; then the normed rows + for (0..3) |i| bind(self, e, i, i); + bind(self, e, 3, 3); + bind(self, e, 3, 4); + bind(self, e, 4, 5); + e.setValue([6]i32{ @intCast(s.rows), @intCast(s.n), 1, 2, 1, 0 }, 6); + e.setValue(eps, 7); + bind(self, e, 5, 8); + e.dispatchGroups(mtl.Size.of(s.rows, 1, 1), mtl.Size.of(32, 1, 1)); + }, + .gate_add => { + for (0..2) |i| bind(self, e, i, i); + bind(self, e, 3, 2); + bind(self, e, 4, 3); + e.setValue([4]i32{ @intCast(s.rows), @intCast(s.n), 1, 2 }, 4); + e.dispatchGroups(mtl.Size.of(s.rows, 1, 1), mtl.Size.of(32, 1, 1)); + }, + .final_norm => { + bind(self, e, 0, 0); + bind(self, e, 2, 1); + bind(self, e, 3, 2); + bind(self, e, 4, 3); + e.setValue([2]i32{ @intCast(s.rows), @intCast(s.n) }, 4); + e.setValue(eps, 5); + bind(self, e, 5, 6); + e.dispatchGroups(mtl.Size.of(s.rows, 1, 1), mtl.Size.of(32, 1, 1)); + }, + .gate_mix => { + for (0..4) |i| bind(self, e, i, i); + e.setValue([3]i32{ @intCast(s.rows), heads, slots / 64 }, 4); + e.dispatchThreads(mtl.Size.of(s.rows, heads, 1), mtl.Size.of(64, 2, 1)); + }, + .swiglu => { + bind(self, e, 0, 0); + e.setValue([2]i32{ @intCast(s.rows), @intCast(s.n) }, 1); + bind(self, e, 1, 2); + e.dispatchThreads(mtl.Size.of(s.n, s.rows, 1), mtl.Size.of(256, 1, 1)); + }, + .heads => { + for (0..8) |i| bind(self, e, i, i); + e.setValue([4]i32{ slots, heads, @intCast(s.k), 0 }, 8); + e.setValue(eps, 9); + for (0..3) |i| bind(self, e, 8 + i, 10 + i); + e.dispatchGroups(mtl.Size.of(slots, heads, 1), mtl.Size.of(32, 1, 1)); + }, + .pool => { + for (0..7) |i| bind(self, e, i, i); + e.dispatchGroups(mtl.Size.of(s.rows, heads, 1), mtl.Size.of(32, 1, 1)); + }, + .tile_scores => { + bind(self, e, 0, 0); + bind(self, e, 1, 1); + e.setValue([2]i32{ @intCast(s.rows), heads }, 2); + bind(self, e, 2, 3); + e.dispatchThreads(mtl.Size.of(s.rows, s.rows, heads), mtl.Size.of(32, 8, 1)); + }, + .coarse => { + bind(self, e, 0, 0); + bind(self, e, 1, 1); + e.setValue([1]i32{@intCast(s.rows)}, 2); + bind(self, e, 2, 3); + e.dispatchGroups(mtl.Size.of(s.rows, heads, 1), mtl.Size.of(32, 1, 1)); + }, + .attention => { + for (0..6) |i| bind(self, e, i, i); + e.setValue([4]i32{ @intCast(s.rows), @intCast(s.k / 64), 0, 1 }, 6); + e.setValue([1]f32{@floatCast(self.in.scale)}, 7); + bind(self, e, 6, 8); + e.dispatchGroups(mtl.Size.of(s.rows, heads, 1), mtl.Size.of(256, 1, 1)); + }, + } +} diff --git a/zig/tests/contract/metal_h3_ref.zig b/zig/tests/contract/metal_h3_ref.zig new file mode 100644 index 0000000000..fa72f5bcfd --- /dev/null +++ b/zig/tests/contract/metal_h3_ref.zig @@ -0,0 +1,266 @@ +//! MiniMax H3's bf16 Metal kernels in float64: each kernel's formula unrounded; its bf16 roundings are its spec's mids. +const std = @import("std"); +const contract = @import("contract"); + +const c = contract.spec; +const Order = contract.gates.Order; +const sumIn = contract.rows.sumIn; + +/// A value rounded to bf16 as the kernels round an input table (to nearest even); NaN and the infinities kept. +pub fn bf(v: f64) f64 { + return c.fromBf16(c.toBf16(@floatCast(v))); +} + +/// What the references read: every input as float64 values (bf16 inputs already hold bf16 values) and the sizes. +pub const In = struct { + a: []const f64 = &.{}, + b: []const f64 = &.{}, + d: []const f64 = &.{}, + e: []const f64 = &.{}, + f: []const f64 = &.{}, + g: []const f64 = &.{}, + ints: []const i32 = &.{}, + more: []const i32 = &.{}, + rows: usize, + n: usize, + k: usize, + heads: usize = 1, + slots: usize = 0, + scale: f64 = 1, + eps: f64 = 1e-5, +}; + +/// How the sums are taken: in float64 (the reference) or in an accumulator and order (gate 4's simulation). +pub const Sum = struct { + acc: c.Accum = .none, + order: Order = .sequential, + + pub fn of(s: Sum, t: []const f64) f64 { + if (s.acc == .none) { + var x: f64 = 0; + for (t) |v| x += v; + return x; + } + return sumIn(s.acc, s.order, t); + } +}; + +fn absSum(t: []const f64) f64 { + var x: f64 = 0; + for (t) |v| x += @abs(v); + return x; +} + +/// Y[r, n] = sum over k of X[r, k] W[n, k] + B[n]: the latent rows in (bf16 out) and the heads out (fp32 out). +pub fn dot(in: In, s: Sum, scratch: []f64, want: []f64, mag: []f64) void { + for (0..in.rows) |r| for (0..in.n) |n| { + const t = scratch[0..in.k]; + for (t, 0..) |*v, k| v.* = in.a[r * in.k + k] * in.b[n * in.k + k]; + const y = s.of(t) + in.d[n]; + want[r * in.n + n] = y; + mag[r * in.n + n] = absSum(t) + @abs(in.d[n]); + }; +} + +/// The modulated norm of one value (h3_rows.metal's h3_mod, which rounds the product and its scale to bf16). +pub fn modulated(x: f64, inv: f64, w: f64, scale: f64, shift: f64) f64 { + return x * inv * w * (1 + scale) + shift; +} + +/// h3_norm: x += gate y, then the modulated norm (a x, b y, d w, e gate, f scale, g shift); stream out, then normed. +pub fn norm(in: In, s: Sum, scratch: []f64, want: []f64, mag: []f64) void { + const C = in.n; + const half = in.rows * C; + for (0..in.rows) |r| { + const x = scratch[0..C]; + for (x, 0..) |*v, i| v.* = in.a[r * C + i] + in.e[i] * in.b[r * C + i]; + const sq = scratch[C .. 2 * C]; + for (sq, x) |*q, v| q.* = v * v; + const inv = 1 / @sqrt(s.of(sq) / @as(f64, @floatFromInt(C)) + in.eps); + for (0..C) |i| { + want[r * C + i] = x[i]; + mag[r * C + i] = @abs(x[i]); + want[half + r * C + i] = modulated(x[i], inv, in.d[i], in.f[i], in.g[i]); + mag[half + r * C + i] = @abs(x[i] * inv * in.d[i] * (1 + in.f[i])) + @abs(in.g[i]); + } + } +} + +/// h3_final_norm: the final layer's modulated RMS norm in float. a x, d weight, f scale row, g shift row. +pub fn finalNorm(in: In, s: Sum, scratch: []f64, want: []f64, mag: []f64) void { + const C = in.n; + for (0..in.rows) |r| { + const sq = scratch[0..C]; + for (sq, 0..) |*q, i| q.* = in.a[r * C + i] * in.a[r * C + i]; + const inv = 1 / @sqrt(s.of(sq) / @as(f64, @floatFromInt(C)) + in.eps); + for (0..C) |i| { + want[r * C + i] = modulated(in.a[r * C + i], inv, in.d[i], in.f[i], in.g[i]); + mag[r * C + i] = @abs(in.a[r * C + i] * inv * in.d[i] * (1 + in.f[i])) + @abs(in.g[i]); + } + } +} + +/// h3_heads: q, k norms and rotate-half on `k` channels, all in tile order (a q, b k, d v, e cos, f sin, g norms). +pub fn heads(in: In, s: Sum, scratch: []f64, want: []f64, mag: []f64) void { + const H = in.heads; + const rot = in.k; + const block = H * in.slots * 128; + for (0..in.rows) |r| for (0..H) |h| { + const base = (r * H + h) * 128; + const at = (h * in.slots + @as(usize, @intCast(in.ints[r]))) * 128; + for ([_][]const f64{ in.a, in.b }, 0..) |x, which| { + const sq = scratch[0..128]; + for (sq, 0..) |*q, i| q.* = x[base + i] * x[base + i]; + const inv = 1 / @sqrt(s.of(sq) / 128 + in.eps); + const w = in.g[which * 128 ..][0..128]; + for (0..128) |i| { + var y = x[base + i] * inv * w[i]; + if (i < rot) { + const half = rot / 2; + const p = if (i < half) i + half else i - half; + const sign: f64 = if (i < half) -1 else 1; + const pair = x[base + p] * inv * w[p]; + y = y * bf(in.e[r * rot + i]) + sign * pair * bf(in.f[r * rot + i]); + } + want[which * block + at + i] = y; + mag[which * block + at + i] = 4 * @abs(x[base + i] * inv * w[i]) + 4 * @abs(x[base + (i + rot / 2) % 128] * inv); + } + } + for (0..128) |i| { + want[2 * block + at + i] = in.d[base + i]; + mag[2 * block + at + i] = @abs(in.d[base + i]); + } + }; +} + +/// h3_pool: each tile's mean q, k and v over its real rows. a, b, d: q, k, v in tile order; ints: the tile sizes. +pub fn pool(in: In, s: Sum, scratch: []f64, want: []f64, mag: []f64) void { + const T = in.rows; + const block = in.heads * T * 128; + for ([_][]const f64{ in.a, in.b, in.d }, 0..) |x, which| for (0..in.heads) |h| for (0..T) |t| for (0..128) |i| { + const size: usize = @intCast(in.ints[t]); + const terms = scratch[0..size]; + for (terms, 0..) |*v, j| v.* = x[((h * in.slots) + t * 64 + j) * 128 + i]; + const o = which * block + (h * T + t) * 128 + i; + want[o] = s.of(terms) / @as(f64, @floatFromInt(size)); + mag[o] = absSum(terms) / @as(f64, @floatFromInt(size)); + }; +} + +/// h3_tile_scores: S[h, tq, tk] = QP[h, tq] . KP[h, tk] / sqrt(128). a QP, b KP. +pub fn scores(in: In, s: Sum, scratch: []f64, want: []f64, mag: []f64) void { + const T = in.rows; + for (0..in.heads) |h| for (0..T) |q| for (0..T) |k| { + const t = scratch[0..128]; + for (t, 0..) |*v, i| v.* = in.a[(h * T + q) * 128 + i] * in.b[(h * T + k) * 128 + i]; + const o = (h * T + q) * T + k; + want[o] = s.of(t) * 0.08838834764831845; + mag[o] = absSum(t) * 0.08838834764831845; + }; +} + +/// h3_coarse: C[h, tq] = softmax over tiles of S[h, tq] times VP[h]. a S, b VP. +pub fn coarse(in: In, s: Sum, scratch: []f64, want: []f64, mag: []f64) void { + const T = in.rows; + for (0..in.heads) |h| for (0..T) |q| { + const row = in.a[(h * T + q) * T ..][0..T]; + var top = -std.math.inf(f64); + for (row) |v| top = @max(top, v); + const e = scratch[0..T]; + for (e, row) |*w, v| w.* = @exp(v - top); + const total = s.of(e); + const t = scratch[T .. 2 * T]; + for (0..128) |i| { + for (t, 0..) |*v, j| v.* = e[j] * in.b[(h * T + j) * 128 + i]; + const o = (h * T + q) * 128 + i; + want[o] = s.of(t) / total; + mag[o] = absSum(t) / total; + } + }; +} + +/// h3_attention: softmax(q k^T scale) v over each query tile's key tiles (ints), the online softmax's order kept. +pub fn attention(in: In, s: Sum, scratch: []f64, want: []f64, mag: []f64) void { + const T = in.rows; + const ksel = in.k / 64; + const H = in.heads; + var row: usize = 0; + for (0..T) |tq| { + const size: usize = @intCast(in.more[tq]); + for (0..size) |r| { + for (0..H) |h| attendRow(in, s, scratch, h, tq, r, ksel, want[(row * H + h) * 128 ..][0..128], mag[(row * H + h) * 128 ..][0..128]); + row += 1; + } + } +} + +fn attendRow(in: In, s: Sum, scratch: []f64, h: usize, tq: usize, r: usize, ksel: usize, want: []f64, mag: []f64) void { + const q = in.a[((h * in.slots) + tq * 64 + r) * 128 ..][0..128]; + const list = in.ints[(h * in.rows + tq) * ksel ..][0..ksel]; + var top = -std.math.inf(f64); + var total: f64 = 0; + const acc = scratch[0..128]; + const accm = scratch[128..256]; + @memset(acc, 0); + @memset(accm, 0); + const t = scratch[256..384]; + const w = scratch[384..448]; + const terms = scratch[448 .. 448 + 64]; + for (list) |tile_i| { + const tile: usize = @intCast(tile_i); + const size: usize = @intCast(in.more[tile]); + var next = top; + var sc: [64]f64 = undefined; + for (0..64) |j| { + if (j >= size) { + sc[j] = -std.math.inf(f64); + continue; + } + const key = in.b[((h * in.slots) + tile * 64 + j) * 128 ..][0..128]; + for (t, q, key) |*v, a, b| v.* = a * b; + sc[j] = s.of(t) * in.scale; + next = @max(next, sc[j]); + } + const shrink = if (next == -std.math.inf(f64)) 1 else @exp(top - next); + total *= shrink; + for (acc, accm) |*o, *m| { + o.* *= shrink; + m.* *= shrink; + } + top = next; + for (w, sc) |*x, v| x.* = if (next == -std.math.inf(f64)) 0 else @exp(v - next); + total += s.of(w); + for (0..128) |i| { + for (terms, w, 0..) |*v, x, j| v.* = x * in.d[((h * in.slots) + tile * 64 + j) * 128 + i]; + acc[i] += s.of(terms); + accm[i] += absSum(terms); + } + } + for (want, mag, acc, accm) |*o, *m, a, b| { + o.* = a / total; + m.* = b / total; + } +} + +/// h3_swiglu: silu(gate) * value in bf16 from rows [value | gate]. a the fused rows. +pub fn swiglu(in: In, want: []f64, mag: []f64) void { + const N = in.n; + for (0..in.rows) |r| for (0..N) |n| { + const v = in.a[r * 2 * N + n]; + const g = in.a[r * 2 * N + N + n]; + const silu = g / (1 + @exp(-g)); + want[r * N + n] = silu * v; + mag[r * N + n] = @abs(silu * v); + }; +} + +/// h3_gate_add: X += bf16(g y) per row. a x, b y, e the gate row. h3_gate_mix: Y += bf16(C[tile] G). a y, b g, d C. +pub fn gated(in: In, mix: bool, want: []f64, mag: []f64) void { + for (0..in.rows) |r| for (0..in.n) |i| { + const o = r * in.n + i; + const g = if (mix) in.d[((i / 128) * in.slots + @as(usize, @intCast(in.ints[r])) / 64) * 128 + i % 128] else in.e[i]; + const add = g * in.b[o]; + want[o] = in.a[o] + add; + mag[o] = @abs(in.a[o]) + @abs(add); + }; +} diff --git a/zig/tests/contract/metal_h3_topk.zig b/zig/tests/contract/metal_h3_topk.zig new file mode 100644 index 0000000000..0c5aaf9dc0 --- /dev/null +++ b/zig/tests/contract/metal_h3_topk.zig @@ -0,0 +1,136 @@ +//! FastH3's tile top-k (h3_topk) under the contract: its bytes against a CPU oracle with the same comparisons. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const h3k = @import("h3_kernels"); +const rig_ = @import("rig.zig"); + +const c = contract.spec; +const g = contract.gates; + +fn noSums(_: c.Shape) usize { + return 0; +} + +const rule = "a NaN score ranks as -inf, ties go to the lower tile, so every row keeps exactly n tiles"; + +/// Rows: query video tiles; n: tiles kept a row; k: video tiles. Compiled for 288 tiles, FastH3's at 864x480. +pub fn spec() c.Spec { + return .{ + .name = "metal h3 tile top-k h3_topk", + .accum = .none, + .depth = noSums, + .out = .raw, + .nan = .{ .stated = rule }, + .tiles = .{ .{}, .{}, .{ .size = 32 } }, + .served = &.{ .{ .rows = 280, .n = 56, .k = 280 }, .{ .rows = 8, .n = 3, .k = 12 } }, + }; +} + +const heads = 2; +const prefix = 8; + +fn rank(v: f32) f32 { + return if (std.math.isNan(v)) -std.math.inf(f32) else v; +} + +pub const TopK = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + shape: c.Shape = undefined, + scores: []f32 = &.{}, + s: rig_.Ref = undefined, + out: rig_.Ref = undefined, + msg: [200]u8 = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue) !TopK { + const src = try std.mem.concat(a, u8, &.{ "#define H3_HEADS 2\n#define H3_SLOTS 18432\n#define H3_TILES 288\n", h3k.source }); + defer a.free(src); + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .pipe = try rig_.pipeline(device, src, "h3_topk") }; + } + + pub fn deinit(t: *TopK) void { + t.a.free(t.scores); + t.pipe.deinit(); + t.rig.deinit(); + } + + pub fn kernel(t: *TopK) g.Kernel { + return .{ .spec = spec(), .ctx = t, .vt = &vt }; + } + + const vt = g.hooks(TopK); + + fn tiles(t: *const TopK) usize { + return prefix + t.shape.k; + } + + pub fn prepare(t: *TopK, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.n > shape.k or shape.rows > shape.k or prefix + shape.k > 288) return false; + t.rig.reset(); + t.shape = shape; + const nt = t.tiles(); + t.a.free(t.scores); + t.scores = &.{}; + t.scores = try t.a.alloc(f32, heads * nt * nt); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + // eighths: many ties + for (t.scores) |*v| v.* = @as(f32, @floatFromInt(rng.uintLessThan(u32, 64))) / 8 - 4; + if (fill == .nonfinite) { + const at = [_]usize{ (prefix + 1) * nt + prefix + 2, (prefix + 2) * nt + nt - 1, (prefix + 3) * nt + prefix }; + const put = [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32) }; + for (at, put) |p, v| t.scores[p] = v; + } + t.s = try t.rig.input(std.mem.sliceAsBytes(t.scores)); + t.out = try t.rig.output(t.count()); + return true; + } + + pub fn run(t: *TopK) !void { + const p = t.rig.begin(); + p.e.setPipeline(t.pipe); + p.e.setBuffer(t.s.buf, t.s.off, 0); + p.e.setValue([3]i32{ @intCast(t.tiles()), prefix, @intCast(t.shape.n) }, 1); + p.e.setBuffer(t.out.buf, t.out.off, 2); + p.e.dispatchGroups(mtl.Size.of(t.shape.rows, heads, 1), mtl.Size.of(32, 1, 1)); + try t.rig.finish(p); + } + + pub fn count(t: *TopK) usize { + return heads * t.shape.k * (prefix + t.shape.n) * 4; + } + + pub fn reference(_: *TopK, want: []f64, mag: []f64) !void { + @memset(want, 0); + @memset(mag, 0); + } + + /// Each launched row's chosen tiles against the oracle: the prefix, then each tile fewer than `n` score above. + pub fn check(t: *TopK, at: g.Gate) !?[]const u8 { + const nt = t.tiles(); + const nv = t.shape.k; + const keep = t.shape.n; + const got = std.mem.bytesAsSlice(i32, t.rig.bytes(t.out, t.count())); + var planted: usize = 0; + for (0..heads) |h| for (0..t.shape.rows) |q| { + const row = t.scores[(h * nt + prefix + q) * nt + prefix ..][0..nv]; + const out = got[(h * nv + q) * (prefix + keep) ..][0 .. prefix + keep]; + for (0..prefix) |i| if (out[i] != i) return try std.fmt.bufPrint(&t.msg, "head {d} row {d}: prefix slot {d} holds {d}", .{ h, q, i, out[i] }); + var placed: usize = prefix; + for (row, 0..) |raw, j| { + if (!std.math.isFinite(raw)) planted += 1; + const mine = rank(raw); + var above: usize = 0; + for (row, 0..) |other, i| above += @intFromBool(rank(other) > mine or (rank(other) == mine and i < j)); + if (above >= keep) continue; + if (placed >= out.len or out[placed] != prefix + j) return try std.fmt.bufPrint(&t.msg, "head {d} row {d}: tile {d} kept, slot {d} differs", .{ h, q, j, placed }); + placed += 1; + } + if (placed != out.len) return try std.fmt.bufPrint(&t.msg, "head {d} row {d}: {d} tiles kept, not {d}", .{ h, q, placed - prefix, keep }); + }; + if (at == .nonfinite and planted == 0) return try std.fmt.bufPrint(&t.msg, "no planted score in a launched row: no proof", .{}); + return null; + } +}; diff --git a/zig/tests/contract/metal_limits.zig b/zig/tests/contract/metal_limits.zig index 234800b800..8af56d980f 100644 --- a/zig/tests/contract/metal_limits.zig +++ b/zig/tests/contract/metal_limits.zig @@ -34,6 +34,13 @@ fn launchOf(a: std.mem.Allocator, device: mtl.Device, d: limits.MetalLaunch) !La if (std.mem.startsWith(u8, fun, "tf_ctl_")) { text = sources.core_contract_controls; if (has(fun, "big_tg")) head = "#define TF_CTL_BIG_TG 1\n"; + } else if (std.mem.startsWith(u8, d.name, "h3 gemm ")) { + text = for (sources.prefill) |f| { + if (std.mem.eql(u8, f.name, "gemm_nax")) break f.text; + } else return error.UnknownLaunch; + frags = true; + } else if (std.mem.startsWith(u8, d.name, "h3 ")) { + text = @import("h3_kernels").source; } else if (std.mem.startsWith(u8, d.name, "glm ")) { const key = d.name[4..std.mem.indexOfScalarPos(u8, d.name, 4, ' ').?]; text = for (sources.glm.all) |k| { diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index 8426690ae3..b4801f4074 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -8,6 +8,8 @@ const affine_k = @import("metal_affine.zig"); const topk_k = @import("metal_topk.zig"); const tile_k = @import("metal_tile.zig"); const limits_k = @import("metal_limits.zig"); +const h3_k = @import("metal_h3.zig"); +const h3_topk = @import("metal_h3_topk.zig"); const g = contract.gates; const usage = "usage: tf-kernel-contract [--only TEXT] [--no-gates] [--no-limits] [--offline DIR] [--seed N] [--repro]"; @@ -108,6 +110,21 @@ pub fn main(init: std.process.Init) !u8 { try proj_list.append(a, p); try kernels.append(a, p.kernel()); }; + var h3_list: std.ArrayList(*h3_k.H3) = .empty; + defer for (h3_list.items) |x| x.deinit(); + if (tensor) for (std.enums.values(h3_k.Which)) |w| { + if (!wanted(o, h3_k.spec(w).name)) continue; + const x = try a.create(h3_k.H3); + x.* = try h3_k.H3.init(a, device, queue, w); + try h3_list.append(a, x); + try kernels.append(a, x.kernel()); + }; + var h3_top: ?h3_topk.TopK = null; + defer if (h3_top) |*x| x.deinit(); + if (wanted(o, h3_topk.spec().name)) { + h3_top = try h3_topk.TopK.init(a, device, queue); + try kernels.append(a, h3_top.?.kernel()); + } var tile_p = try tf.tile_qmm.Pipelines.load(device); defer tile_p.deinit(); var tiles: [2 * tf.tile_qmm.bit_widths.len]tile_k.Tile = undefined; diff --git a/zig/tests/h3_dit.zig b/zig/tests/h3_dit.zig new file mode 100644 index 0000000000..3839d9da1d --- /dev/null +++ b/zig/tests/h3_dit.zig @@ -0,0 +1,116 @@ +//! FastH3's Metal transformer against a case file: the first forward's velocities, then every step, timed. +const std = @import("std"); +const mtl = @import("metal"); +const h3 = @import("h3"); + +extern "c" fn fopen(path: [*:0]const u8, mode: [*:0]const u8) ?*anyopaque; +extern "c" fn fwrite(ptr: *const anyopaque, size: usize, count: usize, file: *anyopaque) usize; +extern "c" fn fclose(file: *anyopaque) c_int; + +const Match = struct { cosine: f64, relative: f64, worst_row: f64 }; + +/// Cosine and relative error over everything, and the largest relative 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 dd: 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| { + dot += @as(f64, a) * b; + aa += @as(f64, a) * a; + bb += @as(f64, b) * b; + rd += (@as(f64, a) - b) * (@as(f64, a) - b); + rb += @as(f64, b) * b; + } + dd += rd; + worst = @max(worst, @sqrt(rd / @max(rb, 1e-30))); + } + return .{ .cosine = dot / @sqrt(aa * bb), .relative = @sqrt(dd / bb), .worst_row = worst }; +} + +fn write(gpa: std.mem.Allocator, prefix: []const u8, comptime suffix: []const u8, values: []const f32) !void { + const path = try std.fmt.allocPrint(gpa, "{s}" ++ suffix ++ "\x00", .{prefix}); + defer gpa.free(path); + const file = fopen(@ptrCast(path.ptr), "wb") orelse return error.CannotWrite; + defer _ = fclose(file); + if (fwrite(values.ptr, 4, values.len, file) != values.len) return error.CannotWrite; +} + +/// One Euler step as the reference takes it: x0 from the sigma the model saw, then the grid's ratio. +fn advance(x: []f32, velocity: []const f32, seen: f32, ratio: f32) void { + for (x, velocity) |*o, v| o.* = ratio * o.* + (1.0 - ratio) * (o.* + seen * v); +} + +pub fn main(init: std.process.Init) !void { + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len != 5 and args.len != 6) { + std.debug.print("usage: tf-h3-dit transformer_dir shards case.safetensors out_prefix [dense | profile | DEFINE, ...]\n", .{}); + return error.BadArguments; + } + const gpa = init.gpa; + const device = try mtl.Device.init(); + var case = try h3.Tensors.open(gpa, device, args[3]); + defer case.deinit(); + const loading = mtl.clock.seconds(); + // The optional last argument: dense, profile, or a define for the kernels. + var options = h3.Options{}; + var defines: [256]u8 = undefined; + var used: usize = 0; + var dense = false; + var profile = false; + if (args.len == 6) { + var it = std.mem.splitScalar(u8, args[5], ','); + while (it.next()) |word| { + if (std.mem.eql(u8, word, "dense")) dense = true else if (std.mem.eql(u8, word, "profile")) profile = true else { + const text = try std.fmt.bufPrint(defines[used..], "#define {s}\n", .{word}); + used += text.len; + } + } + } + options.defines = defines[0..used]; + const model = try h3.Model.load(gpa, args[1], try std.fmt.parseInt(usize, args[2], 10), &case, options); + model.dense = dense; + const steps = (case.names.get("video_step") orelse return error.MissingTensor).shape[0]; + const nv = model.video * model.video_in; + const na = model.audio * model.audio_in; + const video = try case.copy(f32, gpa, try case.at("video", .f32, nv), nv); + const audio = try case.copy(f32, gpa, try case.at("audio", .f32, na), na); + const video_start = try case.copy(f32, gpa, try case.at("video", .f32, nv), nv); + const audio_start = try case.copy(f32, gpa, try case.at("audio", .f32, na), na); + const expect_video = try case.copy(f32, gpa, try case.at("video_velocity", .f32, nv), nv); + const expect_audio = try case.copy(f32, gpa, try case.at("audio_velocity", .f32, na), na); + const video_step = try case.copy(f32, gpa, try case.at("video_step", .f32, 2 * steps), 2 * steps); + const audio_step = try case.copy(f32, gpa, try case.at("audio_step", .f32, 2 * steps), 2 * steps); + std.debug.print("{s}: {d} rows (text {d}, audio {d}, video {d}), {d} tiles ({d} prefix, {d} video kept a query), {d} blocks of {d}/{d}/{d}, {d} steps, loaded in {d:.1} s\n", .{ device.name(), model.rows, model.text, model.audio, model.video, model.tiles, model.prefix_tiles, model.keep, model.blocks.len, model.hidden, model.inner, model.mlp, steps, mtl.clock.seconds() - loading }); + + const vv = try gpa.alloc(f32, nv); + const av = try gpa.alloc(f32, na); + // a knock-out or a probe only wants the profile + const quick = args.len == 6 and (std.mem.indexOf(u8, args[5], "H3_KO") != null or std.mem.indexOf(u8, args[5], "H3_PROBE") != null); + for (0..if (quick) 0 else steps) |step| { + const took = try model.forward(video, audio, step, vv, av); + if (step == 0) { + const m = match(vv, expect_video, model.video_in); + const n = match(av, expect_audio, model.audio_in); + std.debug.print("first video velocity against the reference: cosine {d:.6}, relative error {d:.5}, worst row {d:.5}\n", .{ m.cosine, m.relative, m.worst_row }); + std.debug.print("first audio velocity against the reference: cosine {d:.6}, relative error {d:.5}, worst row {d:.5}\n", .{ n.cosine, n.relative, n.worst_row }); + try write(gpa, args[4], ".first-video.f32", vv); + try write(gpa, args[4], ".first-audio.f32", av); + } + advance(video, vv, video_step[2 * step], video_step[2 * step + 1]); + advance(audio, av, audio_step[2 * step], audio_step[2 * step + 1]); + std.debug.print("step {d}/{d}: {d:.3} s wall, {d:.3} s on the GPU\n", .{ step + 1, steps, took.wall, took.gpu }); + } + try write(gpa, args[4], ".video.f32", video); + try write(gpa, args[4], ".audio.f32", audio); + if (!profile and !quick) return; + model.profile = true; + _ = try model.forward(video_start, audio_start, 0, vv, av); + std.debug.print("one forward by stage, each in its own command buffer:\n", .{}); + for (h3.stages, model.spent) |stage, seconds| std.debug.print(" {s}: {d:.3} s\n", .{ stage, seconds }); +} From 551dff9db8a021bfdf2f0deb3322c5736636e1d6 Mon Sep 17 00:00:00 2001 From: drowzeys <264866615+drowzeys@users.noreply.github.com> Date: Fri, 16 Oct 2026 09:09:39 -0400 Subject: [PATCH 2/2] MiniMax H3 / FastH3 on Metal: clips that start from an image A case exported with a first frame holds the keyframe's rows between the text and the audio. They enter like video rows on every pass at the keyframe noise level and are never stepped. Text to video is bit-identical to before. --- docs/recipes/minimax-h3-metal.md | 22 ++++++++++++++++++---- zig/src/families/h3/h3.zig | 21 ++++++++++++++++----- zig/tests/h3_dit.zig | 2 +- 3 files changed, 35 insertions(+), 10 deletions(-) diff --git a/docs/recipes/minimax-h3-metal.md b/docs/recipes/minimax-h3-metal.md index 288e1029da..3515cba5d3 100644 --- a/docs/recipes/minimax-h3-metal.md +++ b/docs/recipes/minimax-h3-metal.md @@ -7,7 +7,7 @@ a test tool beside the engine; nothing in the lane core, the server or another f ## What it is `zig build tf-h3-dit -Dcpu=apple_m1 -Doptimize=ReleaseFast` builds `zig-out/bin/tf-h3-dit`. It reads the FastH3 -checkpoint's shards and a case file, runs every denoising forward for the packed [text | audio | video] +checkpoint's shards and a case file, runs every denoising forward for the packed [text | keyframe | audio | video] rows, prints the first forward's velocities against the reference velocities the case holds, and writes the final video and audio rows (and the first forward's velocities) as raw float32. @@ -24,7 +24,7 @@ forward stage by stage. - Block projections: the prefill kernels' `tfp::gemm_nax` (MLX 0.32.3's NAX GEMM), bf16 x w^T as the checkpoint stores w, fp32 sums, bf16 out. The SwiGLU's first projection is one product over its stored [value; gate] rows. - Norms, rotary, the gated adds and the SwiGLU round to bf16 where the reference's bf16 MLX arithmetic rounds. -- Attention is FastH3's video sparse attention over 64-row tiles: a text or audio tile attends every +- Attention is FastH3's video sparse attention over 64-row tiles: a text, keyframe or audio tile attends every tile; a video tile attends the prefix tiles and its top 20% of video tiles by pooled score; the pooled branch is gated in. Tile means, scores, the top-k and the pooled branch are fp32. The routed attention is bf16 q, k and v on the tensor units with fp32 scores and an online softmax across key tiles; its weights are rounded to bf16 @@ -51,7 +51,7 @@ branch `h3-1.0.4`). | Tensor | Holds | |---|---| | `text` | The prompt's rows after the condition projection and the token refiner, bf16 | -| `video`, `audio` | Starting noise rows | +| `video`, `audio`, `condition` | Starting noise rows; a first frame's rows when the clip starts from an image | | `cos`, `sin` | Rotary tables for every packed row | | `adaln`, `times`, `tables.N`, `final` | Each row's line in a block's six modulation tables per step, and the tables | | `video_step`, `audio_step` | Per step: the sigma the model sees and the Euler ratio | @@ -76,9 +76,23 @@ At step 1 every video tile is noise, so FastH3's top-k among them is decided by binary with only the fp32 summation order of the tile scores reversed differs from itself by video cosine 0.99021 at step 1 and 0.99834 at step 5: the comparison with the reference is at that floor. +## A clip that starts from an image + +A case exported with a first frame (`h3_case.py --first-frame IMAGE`) has a seventh geometry count: the keyframe's +rows, between the text (which then also holds the image's vision tokens) and the audio. They enter like video rows +on every pass, at the keyframe noise level the case's tables give them, and are never stepped. + +Same Mac, checkpoint and size, a 864x480 first frame: 16,253 packed rows (449 text, 405 keyframe, 414 audio, +14,985 video), 302 tiles (22 prefix). + +| Run | Video velocity against the reference, step 1 | Audio velocity against the reference, step 1 | +|---|---|---| +| From a first frame | cosine 0.99711, relative error 0.076 | cosine 0.99986, relative error 0.017 | + +Text to video is bit-identical with and without this change (final video and audio rows, 8 steps). + ## Not done - No exporter in the tree, so the tool cannot be exercised from this repository alone. - One machine (M5 Ultra). Chips without tensor units are refused. -- Text to video only; a clip that starts from an image is a follow-up change. - No `tensorfold generate`, no text encoder, decoders or sampler: the forward only. diff --git a/zig/src/families/h3/h3.zig b/zig/src/families/h3/h3.zig index 20e211a236..28753f69a4 100644 --- a/zig/src/families/h3/h3.zig +++ b/zig/src/families/h3/h3.zig @@ -48,9 +48,10 @@ pub const Model = struct { mlp: usize, video_in: usize, audio_in: usize, - /// Packed rows: `text`, then `audio`, then `video`. + /// Packed rows: `text`, then `condition` (a first frame's rows, held through the run), then `audio`, then `video`. rows: usize, text: usize, + condition: usize, audio: usize, video: usize, tiles: usize, @@ -100,6 +101,7 @@ pub const Model = struct { tile_scores: mtl.Buffer, chosen: mtl.Buffer, coarse: mtl.Buffer, + condition_latents: mtl.Buffer, video_latents: mtl.Buffer, audio_latents: mtl.Buffer, normed: mtl.Buffer, @@ -143,9 +145,8 @@ pub const Model = struct { var name: [128]u8 = undefined; while (ck.entry(try std.fmt.bufPrint(&name, "transformer_blocks.{d}.norm1.weight", .{layers})) != null) layers += 1; + // a case made for a first frame has a seventh count, the keyframe's rows const counts = (case.names.get("geometry") orelse return error.MissingTensor).shape[0]; - // a case made for a first frame has a seventh count, the keyframe's rows: not read here - if (counts != 6) return error.UnsupportedCase; const geometry = try case.copy(i32, gpa, try case.at("geometry", .i32, counts), counts); defer gpa.free(geometry); self.prefix_tiles = @intCast(geometry[0]); @@ -154,7 +155,8 @@ pub const Model = struct { self.text = @intCast(geometry[3]); self.audio = @intCast(geometry[4]); self.video = @intCast(geometry[5]); - self.rows = self.text + self.audio + self.video; + self.condition = if (counts > 6) @intCast(geometry[6]) else 0; + self.rows = self.text + self.condition + self.audio + self.video; const rows = self.rows; const hidden = self.hidden; const inner = self.inner; @@ -218,6 +220,13 @@ pub const Model = struct { inline for (.{ "qp", "kp", "vp", "coarse" }) |f| @field(self, f) = try device.buffer(heads * self.tiles * head_dim * 4, shared); self.tile_scores = try device.buffer(heads * self.tiles * self.tiles * 4, shared); self.chosen = try device.buffer(heads * (self.tiles - self.prefix_tiles) * (self.prefix_tiles + self.keep) * 4, shared); + self.condition_latents = try device.buffer(@max(self.condition, 1) * self.video_in * 4, shared); + if (self.condition > 0) { + const n = self.condition * self.video_in; + const held_rows = try case.copy(f32, gpa, try case.at("condition", .f32, n), n); + defer gpa.free(held_rows); + @memcpy(self.condition_latents.slice(f32, n), held_rows); + } self.video_latents = try device.buffer(self.video * self.video_in * 4, shared); self.audio_latents = try device.buffer(self.audio * self.audio_in * 4, shared); self.normed = try device.buffer(rows * hidden * 2, shared); @@ -321,7 +330,9 @@ pub const Model = struct { pass.enc = pass.cb.compute(.concurrent); var gpu: f64 = 0; const enc = &pass.enc; - const media = self.text; + // a first frame's rows enter like video rows on every pass and are never stepped + const media = self.text + self.condition; + if (self.condition > 0) self.rowsIn(enc.*, self.condition_latents, self.w_video_in, self.b_video_in, self.condition, self.video_in, self.text); self.rowsIn(enc.*, self.audio_latents, self.w_audio_in, self.b_audio_in, self.audio, self.audio_in, media); self.rowsIn(enc.*, self.video_latents, self.w_video_in, self.b_video_in, self.video, self.video_in, media + self.audio); self.lap(&pass, "rows in"); diff --git a/zig/tests/h3_dit.zig b/zig/tests/h3_dit.zig index 3839d9da1d..520f6a5c25 100644 --- a/zig/tests/h3_dit.zig +++ b/zig/tests/h3_dit.zig @@ -86,7 +86,7 @@ pub fn main(init: std.process.Init) !void { const expect_audio = try case.copy(f32, gpa, try case.at("audio_velocity", .f32, na), na); const video_step = try case.copy(f32, gpa, try case.at("video_step", .f32, 2 * steps), 2 * steps); const audio_step = try case.copy(f32, gpa, try case.at("audio_step", .f32, 2 * steps), 2 * steps); - std.debug.print("{s}: {d} rows (text {d}, audio {d}, video {d}), {d} tiles ({d} prefix, {d} video kept a query), {d} blocks of {d}/{d}/{d}, {d} steps, loaded in {d:.1} s\n", .{ device.name(), model.rows, model.text, model.audio, model.video, model.tiles, model.prefix_tiles, model.keep, model.blocks.len, model.hidden, model.inner, model.mlp, steps, mtl.clock.seconds() - loading }); + std.debug.print("{s}: {d} rows (text {d}, keyframe {d}, audio {d}, video {d}), {d} tiles ({d} prefix, {d} video kept a query), {d} blocks of {d}/{d}/{d}, {d} steps, loaded in {d:.1} s\n", .{ device.name(), model.rows, model.text, model.condition, model.audio, model.video, model.tiles, model.prefix_tiles, model.keep, model.blocks.len, model.hidden, model.inner, model.mlp, steps, mtl.clock.seconds() - loading }); const vv = try gpa.alloc(f32, nv); const av = try gpa.alloc(f32, na);