MiniMax H3 / FastH3 on Metal: clips that start from an image

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(.{

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 <transformer_dir> <shards> 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.

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, .{});

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<device bfloat, dextents<int32_t, 2>, tensor_inline> q((device bfloat*)Q + hb * D, dextents<int32_t, 2>(D, TS));
+  tensor<device bfloat, dextents<int32_t, 2>, tensor_inline> k((device bfloat*)K + hb * D, dextents<int32_t, 2>(D, TS));
+  tensor<device bfloat, dextents<int32_t, 2>, tensor_inline> v((device bfloat*)V + hb * D, dextents<int32_t, 2>(D, TS));
+  threadgroup float st[TQ * KEYS];
+  threadgroup bfloat pt[TQ * KEYS];
+  threadgroup float tops[256];
+  threadgroup float fac[TQ];
+  tensor<threadgroup bfloat, dextents<int32_t, 2>, tensor_inline> p((threadgroup bfloat*)pt, dextents<int32_t, 2>(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<d1, execution_simdgroups<8>> op1;
+  matmul2d<d2, execution_simdgroups<8>> op2;
+  auto a = q.slice<D, TQ>(0, r0);
+  auto b0 = k.slice<D, KEYS>(0, 0);
+  auto v0 = v.slice<D, KEYS>(0, 0);
+  auto acc = op1.template get_destination_cooperative_tensor<decltype(a), decltype(b0), float>();
+  auto out = op2.template get_destination_cooperative_tensor<decltype(p), decltype(v0), float>();
+  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<D, KEYS>(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<D, KEYS>(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]]);
+  }
+}

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;
+}