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..3515cba5d3 --- /dev/null +++ b/docs/recipes/minimax-h3-metal.md @@ -0,0 +1,98 @@ +# 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 | 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. + +```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, 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 + 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`, `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 | +| `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. + +## 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. +- 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..28753f69a4 --- /dev/null +++ b/zig/src/families/h3/h3.zig @@ -0,0 +1,380 @@ +//! 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 `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, + 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, + condition_latents: 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; + + // 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]; + 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.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; + 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.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); + 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; + // 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"); + 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..520f6a5c25 --- /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}, 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); + // 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 }); +}